Vladimir Vargas - 1 year ago 69

Python Question

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.

Answer Source

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]])
```