python - 是否有可能使这个矩阵运算更快？

Question

我想计算

对于固定的 iN 和 i0，G 是 nxn 复数数组，F 原则上是 (N+1)x(N+1)xnxn 复数数组。但是，F 具有对称性，如果 j>=1，则 F(k,j)=F(k-j+1,1)，因此 F 可以存储在形状为 (N+1)x2xnxn 的数组中。例如，典型值是 N=20 和 n=4。

目前，我将 Python 与 jit 一起使用，相应的函数如下所示：

@jit(nopython=True)
def calc():
    out = np.complex128(0)
    i = np.empty(N+1,np.int32) # index array
    i[0] = i0
    i[N] = iN
    for j in range(n**(N-1)): # one execution of the inner of this loops corresponds to one term in the large sum
        for k in range(1,N): # this loops calculates the index array
            i[k] = j % n
            j = j // n

        F_temp = np.complex128(1)
        for kpr in range(N+1):
            F_temp = F_temp * F[kpr,0,i[kpr],i[0]]
            for kprpr in range(1,kpr+1):
                F_temp = F_temp * F[kpr-kprpr+1,1,i[kpr],i[kprpr]]

        G_temp = np.complex128(1)
        for k in range(N):
            G_temp = G_temp * G[i[k+1],i[k]]

        out += G_temp * F_temp
    return out

我觉得这个实现远非最佳（可能是一个数量级）。例如，如果我使用 np.complex64 而不是 np.complex128 ，我根本看不到加速（这不应该是有效实现的情况，对吧？）。所以，问题来了：我有没有机会在这方面达到 10 倍的加速（或者至少在那个范围内）？就像用另一种语言实现这个函数并从 python 调用它（甚至会比 numba 更快）？

标签： pythonperformancetensornumba