0

我有以下代码：

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，我找到第 th 层内的数组元素的索引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变成X2一维数组，我不确定其中的索引部分E或如何适当地重塑数组。

那么有什么方法可以加快原始代码的速度呢？

4

0 回答 0