给定x，y平面上的MxN网格，计算f（x，y）并将其存储到矩阵（python）

Question

我正在寻找一些类似于MATLAB中的ARRAYFUN的东西，但是是Python。我需要做的是计算一个矩阵，它的分量是exp（j*dot（[kx，ky]，[x，y]），其中[kx，ky]是一个固定的已知向量，[x，y]是网格中的一个元素。你知道吗

我想做的是定义

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

然后在R索引上迭代，填充一个与R形状相同的矩阵，其中每个分量都是exp（j*dot（[kx，ky]，[x，y]），其中[x，y]在R中。这看起来既不高效也不优雅。你知道吗

谢谢你的帮助。你知道吗

Answer 1

在他们添加ARRAYFUN之前，您可以像我们在MATLAB中所做的那样-更改计算，使其与数组一起工作。在MATLAB中的一切都是2d的时代，这可能会很棘手；允许更多的维度使它变得更容易。numpy允许2个以上的维度。你知道吗

总之，这里有一个快速的尝试：

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

还有其他版本的dot、matmul和tensordot，但是einsum是我喜欢使用的。我已经做了足够多的工作来快速建立一个多维的dot。你知道吗

现在只需对每个元素应用1j和exp：

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