Python: how to convert a complex array to a 2D array?

As a C++ programmer, I'm used to access vectors in C++ style:

for (i=0; i<max_x; i++) {
  for (j=0; j<max_y; j++) {
    vec[i][j] = real(complex_number(j+i*max_x))
  }
}

Now I have in Python

    x = np.linspace(x1, x2, step)
    y = np.linspace(y1, y2, step)
    X, Y = np.meshgrid(x, y)
    Z = x + 1j*y

    for z in Z:
      FZ = complex_function(z)

How can I accomplish the same thing as the C++ code in a "pythonic" way? Thanks

EDIT: Checking the reshape function and reanalizing my code, I noted a problem with transformations from 2D arrays to 1D arrays and back. The main problem is that I have a function that accepts an complex array z_list and return a complex array. I need to plot that on a 开发者_如何学Cgrid, and I was planning to use matplotlib, but matplotlib needs a 2D array with each point on that array having the value. How can I do this without generate a 2D array, reshape it to a 1D array, and re-reshape the array back to 2D? Thanks.

Use reshape to turn a 1D array to 2D (or any other shapes).

>>> x_max = 12
>>> y_max = 4
>>> vec1d = np.arange(x_max*y_max, dtype=complex)
>>> vec1d.reshape([x_max, y_max])
array([[  0.+0.j,   1.+0.j,   2.+0.j,   3.+0.j],
       [  4.+0.j,   5.+0.j,   6.+0.j,   7.+0.j],
       [  8.+0.j,   9.+0.j,  10.+0.j,  11.+0.j],
       [ 12.+0.j,  13.+0.j,  14.+0.j,  15.+0.j],
       [ 16.+0.j,  17.+0.j,  18.+0.j,  19.+0.j],
       [ 20.+0.j,  21.+0.j,  22.+0.j,  23.+0.j],
       [ 24.+0.j,  25.+0.j,  26.+0.j,  27.+0.j],
       [ 28.+0.j,  29.+0.j,  30.+0.j,  31.+0.j],
       [ 32.+0.j,  33.+0.j,  34.+0.j,  35.+0.j],
       [ 36.+0.j,  37.+0.j,  38.+0.j,  39.+0.j],
       [ 40.+0.j,  41.+0.j,  42.+0.j,  43.+0.j],
       [ 44.+0.j,  45.+0.j,  46.+0.j,  47.+0.j]])

Instead of doing Z = x + 1j*y then reshaping, you could do:

Z = np.zeros((ydim, xdim), dtype=complex)
Z.real, Z.imag = X, Y

which I think might be more efficient (less operations in total).

Use reshape

>>> Z.reshape(5,10)
array([[  0.00000000 +0.j        ,   0.20408163 +0.20408163j,
          0.40816327 +0.40816327j,   0.61224490 +0.6122449j ,
          0.81632653 +0.81632653j,   1.02040816 +1.02040816j,
          1.22448980 +1.2244898j ,   1.42857143 +1.42857143j,
          1.63265306 +1.63265306j,   1.83673469 +1.83673469j],
       [  2.04081633 +2.04081633j,   2.24489796 +2.24489796j,
          2.44897959 +2.44897959j,   2.65306122 +2.65306122j,
          2.85714286 +2.85714286j,   3.06122449 +3.06122449j,
          3.26530612 +3.26530612j,   3.46938776 +3.46938776j,
          3.67346939 +3.67346939j,   3.87755102 +3.87755102j],
       [  4.08163265 +4.08163265j,   4.28571429 +4.28571429j,
          4.48979592 +4.48979592j,   4.69387755 +4.69387755j,
          4.89795918 +4.89795918j,   5.10204082 +5.10204082j,
          5.30612245 +5.30612245j,   5.51020408 +5.51020408j,
          5.71428571 +5.71428571j,   5.91836735 +5.91836735j],
       [  6.12244898 +6.12244898j,   6.32653061 +6.32653061j,
          6.53061224 +6.53061224j,   6.73469388 +6.73469388j,
          6.93877551 +6.93877551j,   7.14285714 +7.14285714j,
          7.34693878 +7.34693878j,   7.55102041 +7.55102041j,
          7.75510204 +7.75510204j,   7.95918367 +7.95918367j],
       [  8.16326531 +8.16326531j,   8.36734694 +8.36734694j,
          8.57142857 +8.57142857j,   8.77551020 +8.7755102j ,
          8.97959184 +8.97959184j,   9.18367347 +9.18367347j,
          9.38775510 +9.3877551j ,   9.59183673 +9.59183673j,
          9.79591837 +9.79591837j,  10.00000000+10.j        ]])