numpy - 配列のスライスとブロードキャストを含む numpy コードの高速化

Question

次のコードがあります。

x = sp.linspace(-2,2,1000)
z = sp.linspace(-1,3,2000)
X,Z = sp.meshgrid(x,z) 
X = X[:,:,sp.meshgrid]
Z = Z[:,:,sp.meshgrid]
E = sp.zeros((len(z),len(x),3), dtype=complex)

# e_uvect.shape = (2,N,2,3)
# En.shape = (2,N,2)
# d_cum.shape = (N,)
# pol is either 0 or 1

for n in range(N):
    idx = sp.logical_and(Z<d_cum[n], Z>=d_cum[n-1])        
    E += e_uvect[pol,n,0,:]*En[pol,n,0]*sp.exp(+1j*self.kz[n]*(Z-d_cum[n-1])+1j*self.kx*X)*idx

N-基本的に、上記は層構造の電界を計算するためのコードの一部です。ループ内の反復ごとforに、層内にある配列要素のインデックスを見つけ、N電界を計算した後、すべてを掛けて、idxを満たす正しい部分を「フィルタリング」しますsp.logical_and(Z<d_cum[n], Z>=d_cum[n-1])。

それは正常に動作しますが、numpy 配列スライシングまたは他の方法を使用してこれを行うより効率的な方法があるかどうか疑問に思います。各乗算には、各反復で受け入れられない配列要素の大部分が含まれるためです。Z座標配列の関連部分でのみ動作するように、次のようなことを試しましたX

idx = sp.logical_and(Z<d_cum[n], Z>=d_cum[n-1]) 
Z2 = Z[idx]
X2 = X[idx]
E[???] += e_uvect[pol,n,0,:]*En[pol,n,0]*sp.exp(+1j*self.kz[n]*(Z2-d_cum[n-1])+1j*self.kx*X2)

しかし、その後Z2、 1次元配列になり、その中のインデックス部分や、配列を適切に再形成する方法についてはわかりX2ません。E

元のコードを高速化する方法はありますか?