Python 使用 numpy 魔法避免三重嵌套 for 循环

我已经编写了以下代码，它可以满足我的需要，但我希望对其进行优化.

import numpy as np
N_rs = 100
N_thetas = 30

Nlm_container = np.ones( (N_thetas, 2*N_thetas-1), dtype=np.complex128 )
lpmv_container = np.ones ( (N_thetas, 2*N_thetas-1, N_thetas), dtype=np.complex128 )
Psi_m_res1 = np.zeros( (2*N_thetas - 1,    N_rs - 1,    N_thetas) , dtype=np.complex128 )
Psi_m_res2 = np.zeros( (2*N_thetas - 1,    N_rs - 1,    N_thetas) , dtype=np.complex128 )
Psi_m_res3 = np.zeros( (2*N_thetas - 1,    N_rs - 1,    N_thetas) , dtype=np.complex128 )
Psi_lm = np.ones ( ( N_rs - 1, 2*N_thetas-1, N_thetas), dtype=np.complex128 )

for m in range(-N_thetas+1,  N_thetas,  1):
    for j in range(N_rs - 1):
        for k in range(N_thetas):
            Psi_m_res1[m + (N_thetas-1),  j,  k] = np.sum(  Psi_lm[j, m+(N_thetas-1), np.abs(m):] * Nlm_container[np.abs(m):, m+(N_thetas-1)]  *  lpmv_container[abs(m):, m+(N_thetas-1), k]  )

我希望避免嵌套的三元组for循环. 我认为使用:会更好，而不是使用j之上的循环和k之上的循环.我这样写(作为第一步):

for m in range(-N_thetas+1,  N_thetas,  1):
    for k in range(N_thetas):
        Psi_m_res2[m + (N_thetas-1),  :,  k] = np.sum(  Psi_lm[:, m+(N_thetas-1), np.abs(m):] * Nlm_container[np.abs(m):, m+(N_thetas-1)]  *  lpmv_container[abs(m):, m+(N_thetas-1), k]  )

尽管如此，这并没有产生我想要的:

print(np.testing.assert_allclose(Psi_m_res1, Psi_m_res2))
Mismatched elements: 175230 / 175230 (100%)
Max absolute difference: 2940.
Max relative difference: 0.98989899

更重要的是，当同时消除k索引for循环时，程序根本不会计算任何形状不匹配.

for m in range(-N_thetas+1,  N_thetas,  1):
    Psi_m_res3[m + (N_thetas-1),  :,  :] = np.sum(  Psi_lm[:, m+(N_thetas-1), np.abs(m):] * Nlm_container[np.abs(m):, m+(N_thetas-1)]  *  lpmv_container[abs(m):, m+(N_thetas-1), :]  )

ValueError: operands could not be broadcast together with shapes (99,2) (2,30)

我可以创建线程，将Psi_lm、Nlm_container和lpmv_container的片段发送给它们，然后请求执行工作，然后重新组合结果，但我认为麻木魔法可能会做得更好.

在实践中，N_rs和N_thetas应该更高--它们可以上升到N_rs = 1600和N_thetas=100，这是在代码开始运行之前设置的，并在代码运行时期间进行修复.

有没有办法用麻木魔法来解决这个问题？

谢谢!