Vladimir Vargas Vladimir Vargas - 9 months ago 34
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.

Thanks for your help.


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)
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)
array([[ 1.00000000+0.j        ,  0.54030231+0.84147098j,
       [-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 ,