有可能使这个矩阵运算更快吗?

2024-05-14 23:01:56 发布

您现在位置:Python中文网/ 问答频道 /正文

我想计算

对于固定iN和i0,G是复数的nxn数组,F原则上是复数的(N+1)x(N+1)xnxn数组。然而,F的对称性使得F(k,j)=F(k-j+1,1)如果j>=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更快)


Tags: theinfornprange数组outtemp

热门问题