python - 给定 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 中的一切都是二维的时代，这可能会很棘手；允许更多的维度使它更容易。 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 是我喜欢的使用。我已经充分使用它来快速设置多维点。

现在只需将 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]])