Python:使用预先计算的元素加速大双倍和

Question

>我需要计算以下形式的双倍和：

wignersum{ell} = sum_{ell1} sum_{ell2} (2*ell1+1((2*ell2+1( * W{ell，ell1，ell2}^2 * C1(ell1( * C2(ell2(

其中 Wignersum 是由 ell 索引的数组，ell、ell1 和 ell2 都从 0 运行到 ellmax。W{ell，ell1，ell2}^2 是我已经计算过的一组已知系数(称为 w3j(，存储在形状数组(ellmax、ellmax、ellmax(中，作为此函数调用的全局变量。(这些系数的计算非常耗时，我发现从 numpy 文件加载它们更快(。C1 和 C2 是形状系数 (ellmax( 的数组。

我已经通过使用双精度 for 循环并从每个现有数组中获取适当的元素并在每次迭代中更新 wignersum 数组来成功计算出这个总和。我认为有一种更好的方法来矢量化这个问题以加快计算速度。我考虑过将 C1 和 C2 数组制作成与 w3j 数组形状相同的数组，然后在 ell1 和 ell2 轴上使用 np.sum 之前将这些数组按元素相乘。我不确定这是否真的是一种很好的vecotrizing方法，如果是，如何实际做到这一点。

目前的代码是这样的

import numpy as np
ell_max = 400
w3j = np.ones((ell_max, ell_max, ell_max))
C1 = np.arange(ell_max)
C2 = np.arange(ell_max)
def function(ell_max)
ells = np.arange(ell_max)
wignersum = np.zeros(ell_max)
factor = np.array([2*i+1 for i in range(384)])
for ell1 in ells:
A = factor[ell1]
B = C1[ell1]
for ell2 in ells:
D = factor[ell2] * C2[ell2] * w3j[:,ell1,ell2]
wignersum += A * B * D
return wignersum

(请注意，实际C1和C2不是全局变量，而是必须根据提供给function的一组参数计算的局部变量。然而，这不是代码速度的限制因素(

使用双 for 循环，运行 ell_max~400 需要 ~1.5 秒，这对于我使用它的目的来说太长了。我想尽可能多地矢量化以提高速度。

Answer 1

您可以使用 einsum 或矩阵乘法来实现 ~20 倍的加速：

import numpy as np
ell_max = 400
w3j = np.random.randint(1,10,(ell_max, ell_max, ell_max))
C1 = np.random.randint(1,10,ell_max)
C2 = np.random.randint(1,10,ell_max)
def function(ell_max):
ells = np.arange(ell_max)
wignersum = np.zeros(ell_max)
factor = np.array([2*i+1 for i in range(ell_max)])
for ell1 in ells:
A = factor[ell1]
B = C1[ell1]
for ell2 in ells:
D = factor[ell2] * C2[ell2] * w3j[:,ell1,ell2]
wignersum += A * B * D
return wignersum
def pp_es(l_mx):
l = np.arange(l_mx)
f = 2*l+1
return np.einsum("i,i,j,j,kij",f,C1,f,C2,w3j,optimize=True)
def pp_mm(l_mx):
l = np.arange(l_mx)
f = 2*l+1
return w3j.reshape(l_mx,-1)@np.outer(f*C1,f*C2).ravel()
from timeit import timeit
print(timeit(lambda:pp_es(400),number=10))
print(timeit(lambda:pp_mm(400),number=10))
print(timeit(lambda:function(400),number=10))
print((pp_mm(400)==pp_es(400)).all())
print((function(400)==pp_mm(400)).all())

示例运行：

0.6061844169162214 # einsum
0.6111843499820679 # matrix x vector
12.233918005018495 # OP
True # einsum == matrix x vector
True # OP == matrix x vector