Numpy 根据位置逐元素应用函数答案

【问题标题】：Numpy apply function elementwise based on positionNumpy 根据位置逐元素应用函数
【发布时间】：2018-02-09 16:48:22
【问题描述】：

我有一个形状为 (M,N) 的 numpy 数组 A。我想创建一个形状为 (M,N,3) 的新数组 B，其结果将与以下内容相同：

import numpy as np

def myfunc(A,sx=1.5,sy=3.5):
    M,N=A.shape
    B=np.zeros((M,N,3))

    for i in range(M):
        for j in range(N):
            B[i,j,0]=i*sx
            B[i,j,1]=j*sy
            B[i,j,2]=A[i,j]
    return B

A=np.array([[1,2,3],[9,8,7]])
print(myfunc(A))

给出结果：

[[[0.  0.  1. ]
  [0.  3.5 2. ]
  [0.  7.  3. ]]

 [[1.5 0.  9. ]
  [1.5 3.5 8. ]
  [1.5 7.  7. ]]]

有没有办法在没有循环的情况下做到这一点？我在想 numpy 是否能够使用数组的索引逐元素应用函数。比如：

def myfuncEW(indx,value,out,vars):
    out[0]=indx[0]*vars[0]
    out[1]=indx[1]*vars[1]
    out[2]=value

M,N=A.shape
B=np.zeros((M,N,3))
np.applyfunctionelementwise(myfuncEW,A,B,(sx,sy))

【问题讨论】：

标签： python numpy

【解决方案1】：

您可以使用mgrid 和moveaxis：

>>> M, N = A.shape
>>> I, J = np.mgrid[:M, :N] * np.array((sx, sy))[:, None, None]
>>> np.moveaxis((I, J, A), 0, -1)
array([[[ 0. ,  0. ,  1. ],
        [ 0. ,  3.5,  2. ],
        [ 0. ,  7. ,  3. ]],

       [[ 1.5,  0. ,  9. ],
        [ 1.5,  3.5,  8. ],
        [ 1.5,  7. ,  7. ]]])
>>>

【讨论】：

【解决方案2】：

你可以使用meshgrid和dstack，像这样：

import numpy as np

def myfunc(A,sx=1.5,sy=3.5):
    M, N = A.shape
    J, I = np.meshgrid(range(N), range(M))
    return np.dstack((I*sx, J*sy, A))

A=np.array([[1,2,3],[9,8,7]])
print(myfunc(A))

# array([[[ 0. ,  0. ,  1. ],
#         [ 0. ,  3.5,  2. ],
#         [ 0. ,  7. ,  3. ]],
#
#        [[ 1.5,  0. ,  9. ],
#         [ 1.5,  3.5,  8. ],
#         [ 1.5,  7. ,  7. ]]])

【讨论】：

【解决方案3】：

通过预先分配 3d 数组 B，与堆叠 I、J 和 A 相比，您可以节省大约一半的时间。

def myfunc(A, sx=1.5, sy=3.5):
    M, N = A.shape
    B = np.zeros((M, N, 3))
    B[:, :, 0] = np.arange(M)[:, None]*sx
    B[:, :, 1] = np.arange(N)[None, :]*sy
    B[:, :, 2] = A
    return B

【讨论】：