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))
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 : rx,ry=np.meshgrid(np.arange(3),np.arange(4)) In : R=np.dstack((rx,ry)) In : R.shape Out: (4, 3, 2) In : kx,ky=1,2 In : np.einsum('i,jki->jk',[kx,ky],R) Out: array([[0, 1, 2], [2, 3, 4], [4, 5, 6], [6, 7, 8]])
There are other versions of
einsum is the one I like to use. I've worked with it enough to quickly set up a multidimensional
Now just apply the
exp to each element:
In : np.exp(np.einsum('i,jki->jk',[kx,ky],R)*1j) Out: 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]])