Question

我有一个三维numpy数组A.我想将每个元素A [i，j，k]乘以w *（i / Lx + j / Ly + k / Lz）其中w，Lx，Ly和Lz是实数（浮点数）。在for循环中执行此操作是非常不切实际的，因为我需要能够为大型数组扩展它，并且在O（N ^ 3）中以三个索引进行for循环。

是否有一种有效的方法对关注索引的numpy数组的每个元素执行操作？

Answer 1

您可以使用broadcasting -

M,N,R = A.shape

p1 = np.arange(M)[:,None,None]/Lx
p2 = np.arange(N)[:,None]/Ly
p3 = np.arange(R)/Lz

out = A/(w*(p1 + p2 + p3))

您还可以使用np.ix_获取更优雅的解决方案 -

M,N,R = A.shape
X,Y,Z = np.ix_(np.arange(M),np.arange(N),np.arange(R))
out = A/(w*((X/Lx) + (Y/Ly) + (Z/Lz)))

运行时测试和输出验证 -

功能定义：

def vectorized_app1(A, w, Lx, Ly, Lz ):
    M,N,R = A.shape
    p1 = np.arange(M)[:,None,None]/Lx
    p2 = np.arange(N)[:,None]/Ly
    p3 = np.arange(R)/Lz
    return A/(w*(p1 + p2 + p3))

def vectorized_app2(A, w, Lx, Ly, Lz ):
    M,N,R = A.shape
    X,Y,Z = np.ix_(np.arange(M),np.arange(N),np.arange(R))
    return A/(w*((X/Lx) + (Y/Ly) + (Z/Lz)))

def original_app(A, w, Lx, Ly, Lz ):
    out = np.empty_like(A)
    M,N,R = A.shape
    for i in range(M):
        for j in range(N):
            for k in range(R):
                out[i,j,k] = A[i,j,k]/(w*( (i / Lx) + (j / Ly) + (k / Lz) ))
    return out

时序：

In [197]: # Inputs
     ...: A = np.random.rand(100,100,100)
     ...: w, Lx, Ly, Lz = 2.3, 3.2, 4.2, 5.2
     ...: 

In [198]: np.allclose(original_app(A,w,Lx,Ly,Lz),vectorized_app1(A,w,Lx,Ly,Lz))
Out[198]: True

In [199]: np.allclose(original_app(A,w,Lx,Ly,Lz),vectorized_app2(A,w,Lx,Ly,Lz))
Out[199]: True

In [200]: %timeit original_app(A, w, Lx, Ly, Lz )
1 loops, best of 3: 1.39 s per loop

In [201]: %timeit vectorized_app1(A, w, Lx, Ly, Lz )
10 loops, best of 3: 24.6 ms per loop

In [202]: %timeit vectorized_app2(A, w, Lx, Ly, Lz )
10 loops, best of 3: 24.2 ms per loop

将关心索引的函数应用于numpy数组的每个元素

1 个答案: