Vladimir Vargas -5 years ago 212
Python Question

# Given a MxN grid in the x,y plane, compute f(x,y) and store it into matrix (python)

I'm looking for something similar to ARRAYFUN in MATLAB, but for Python. What I need to do is to compute a matrix whose components are exp(j*dot([kx,ky], [x,y])), where [kx,ky] is a fixed known vector, and [x,y] is an element from a meshgrid.

What I was trying to do is to define

``````RX, RY = np.meshgrid(np.arange(N), np.arange(M))
R = np.dstack((RX,RY))
``````

and then iterate over the R indices, filling a matrix with the same shape as R, in which each component would be exp(j*dot([kx,ky], [x,y])), with [x,y] being in R. This doesn't look efficient nor elegant.

You could do what we used to do in MATLAB before they added `ARRAYFUN` - change the calculation so it works with arrays. That could be tricky in the days when everything in MATLAB was 2d; allowing more dimensions made it easier. `numpy` allows more than 2 dimensions.

Anyways, here a quick attempt:

``````In [497]: rx,ry=np.meshgrid(np.arange(3),np.arange(4))
In [498]: R=np.dstack((rx,ry))

In [499]: R.shape
Out[499]: (4, 3, 2)

In [500]: kx,ky=1,2

In [501]: np.einsum('i,jki->jk',[kx,ky],R)
Out[501]:
array([[0, 1, 2],
[2, 3, 4],
[4, 5, 6],
[6, 7, 8]])
``````

There are other versions of `dot`, `matmul` and `tensordot`, but `einsum` is the one I like to use. I've worked with it enough to quickly set up a multidimensional `dot`.

Now just apply the `1j` and `exp` to each element:

``````In [502]: np.exp(np.einsum('i,jki->jk',[kx,ky],R)*1j)
Out[502]:
array([[ 1.00000000+0.j        ,  0.54030231+0.84147098j,
-0.41614684+0.90929743j],
[-0.41614684+0.90929743j, -0.98999250+0.14112001j,
-0.65364362-0.7568025j ],
[-0.65364362-0.7568025j ,  0.28366219-0.95892427j,
0.96017029-0.2794155j ],
[ 0.96017029-0.2794155j ,  0.75390225+0.6569866j ,
-0.14550003+0.98935825j]])
``````
Recommended from our users: Dynamic Network Monitoring from WhatsUp Gold from IPSwitch. Free Download