在 Python 中使用复杂的索引加速 for 循环？

Question

我正在开发一个必须计算（协方差）矩阵的程序，而我正在使用的 for 循环实现非常慢。 Numba 抱怨 4 维 numpy 数组，我真的不知道如何向量化代码。代码如下：

covmat = np.zeros((nbl, nbl, npairs_tot, npairs_tot))

for ell1 in range(nbl):
    for ell2 in range(nbl):
        for p in range (npairs_tot):
            for q in range (npairs_tot):
                
                if ind[p,0] == 0 and ind[p,1] == 0 and ind[q,0] == 0 and ind[q,1] == 0:
                    covmat[ell1,ell2,p,q] = (Rl*Rl*\
                    D_ALL[ell1, ind[p,0], ind[p,1], ind[p,2], ind[p,3]]*\
                    D_ALL[ell2, ind[q,0], ind[q,1], ind[q,2], ind[q,3]]*\
                    matrix[ind[p,2],ind[p,3],ind[q,2],ind[q,3]])
                
                elif ind[p,0] == 0 and ind[p,1] == 0 and ind[q,0] == 1 and ind[q,1] == 0:
                    covmat[ell1,ell2,p,q] = (Rl*Rl*\
                    D_ALL[ell1, ind[p,0], ind[p,1], ind[p,2], ind[p,3]]*\
                    D_ALL[ell2, ind[q,0], ind[q,1], ind[q,2], ind[q,3]]*\
                    matrix[ind[p,2],ind[p,3],ind[q,2] + zbins,ind[q,3]])

                elif ind[p,0] == 0 and ind[p,1] == 0 and ind[q,0] == 1 and ind[q,1] == 1:
                    covmat[ell1,ell2,p,q] = (Rl*Rl*\
                    D_ALL[ell1, ind[p,0], ind[p,1], ind[p,2], ind[p,3]]*\
                    D_ALL[ell2, ind[q,0], ind[q,1], ind[q,2], ind[q,3]]*\
                    matrix[ind[p,2],ind[p,3],ind[q,2] + zbins,ind[q,3] + zbins])
                    
                    # ... similar blocks

正如您所见，if 语句检查我是否应该以及在何处添加 zbins 语句使任务更加复杂 项（它是常数）到矩阵的索引。 ind 只是一个用于存储正确索引的二维数组。
你有什么建议吗？谢谢！！

Answer 1

我们不可能做一个工作示例，因为这是一个不可重现的问题，请遵循建议here。例如，无法理解如何使用ind。

提示

这似乎是一个很难计算的 O(n^4) 问题... numpy 处理数组非常好，所以开始应用 divide and conquer 从内部循环向量化。 条件变为掩码，并且必须与 q 维度具有相同的长度。