计算级数（大和）的最佳迭代技术

Question

我的代码产生了正确的结果，但感觉非常慢。 如何让

Z00

功能更加高效？

我的目标是证明立方体尺寸的数值稳定性（即

n_max

）。

生成许多函数值并将其显示在图表中的最有效方法是什么？

import numpy as np
import matplotlib.pyplot as plt
import matplotlib.ticker as ticker
import itertools as it
import time as t

def Z00(eta,n_max,mu,m,d3):
    precartesian = [range(-n_max,n_max+1),range(-n_max,n_max+1),range(-n_max,n_max+1)] #3d grid
    cartesian = list(it.product(*precartesian))
    gamma = np.sqrt(1+d3**2/(4*m**2+eta**2)) 
    sum=0
    for (n1,n2,n3) in cartesian:
        r2 = gamma**2*(n3-d3/2)**2+n1**2+n2**2 #r²
        sum+=1/((r2-eta**2)*(r2+mu**2)**2)
    return 1/(2*np.sqrt(np.pi))*((mu**2-eta**2)**2*sum+gamma**2*np.pi**2/mu*(eta**2-mu**2))

Z00

函数使用

n_max

来创建3D网格。然后，对于网格上的每个值，计算函数值

r2

。这被进一步操纵并总结为

sum

。

然后我计算值并绘制函数：

#set function parameters
nmax=10
mu=2
m=1

#set plot parameters
xmin=0
xmax=6

#compute function values
x = np.linspace(xmin,xmax,10000)
y = Z00(x,nmax,mu,m,1)
plt.plot(x,y)

一般来说，计算将按

n_max**3

进行缩放，因为我们正在计算立方体内的值的函数。我已经尝试过使用内置函数（例如

np.sum()

）以及列表理解来计算总和：

def Z00(eta,n_max,mu,m,d3):
    precartesian = [range(-n_max,n_max+1),range(-n_max,n_max+1),range(-n_max,n_max+1)] #3d grid
    cartesian = list(it.product(*precartesian))
    gamma = np.sqrt(1+d3**2/(4*m**2+eta**2))
    r2s = [gamma**2*(n3-d3/2)**2+n1**2+n2**2 for (n1,n2,n3) in cartesian]
    sum = np.sum([1/((r2-eta**2)*(r2+mu**2)**2) for r2 in r2s])
    return 1/(2*np.sqrt(np.pi))*((mu**2-eta**2)**2*sum+gamma**2*np.pi**2/mu*(eta**2-mu**2))

此外，我尝试减少分配的局部变量的数量，以及将 3d 网格移动到全局变量，但这也没有帮助。

Answer 1

我首先将您的

Z00

简化为

Z01

，避免重新计算，然后将

Z00

转换为

Z02

，这是一种矢量实现（请参阅下面的详细信息）。

最终我测试了我的实现的正确性，并根据您的原始代码对它们进行了基准测试，

13:46 Fri, Aug 02 ~/ $ ipython-3.12 -i faster.py
Python 3.12.3 (main, Apr 10 2024, 14:51:56) [GCC]
Type 'copyright', 'credits' or 'license' for more information
IPython 8.25.0 -- An enhanced Interactive Python. Type '?' for help.

In [1]: # are the results the same?
   ...: for Z0x in (Z00, Z01, Z02): print(Z0x(x, nmax, mu, m, 1))
[ -2.87639267  -2.87639041  -2.87638363 ... 168.60197891 243.68957108
 493.83389926]
[ -2.87639267  -2.87639041  -2.87638363 ... 168.60197891 243.68957108
 493.83389926]
[ -2.87639267  -2.87639041  -2.87638363 ... 168.60197891 243.68957108
 493.83389926]

In [2]: # which is fastest?
   ...: for Z0x in (Z00, Z01, Z02):
   ...:     %timeit Z0x(x, nmax, mu, m, 1)
320 ms ± 2.54 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
275 ms ± 442 μs per loop (mean ± std. dev. of 7 runs, 1 loop each)
1.16 s ± 706 μs per loop (mean ± std. dev. of 7 runs, 1 loop each)

In [3]: 

正如您所看到的，所有三种实现的结果都一致，简化的实现比您的稍快，而完全矢量化的实现则慢了 3÷4 倍（我想知道我是否犯了某种错误）在此实施中）。

虽然

Z01

更快，但我想说编译器或 Numpy 已经对

Z00

中的重复计算进行了一些优化，并且（至少现在）向量化计算没有帮助。

正如所承诺的，我的代码包含三个不同的函数和数据。

import numpy as np
from itertools import product

nmax = 10
mu = 2
m = 1

#set plot parameters
xmin = 0
xmax = 6

#compute function values
x = np.linspace(xmin, xmax, 10000)

def Z00(eta,n_max,mu,m,d3):
    precartesian = [range(-n_max,n_max+1),range(-n_max,n_max+1),range(-n_max,n_max+1)] #3d grid
    cartesian = list(product(*precartesian))
    gamma = np.sqrt(1+d3**2/(4*m**2+eta**2)) 
    sum=0
    for (n1,n2,n3) in cartesian:
        r2 = gamma**2*(n3-d3/2)**2+n1**2+n2**2 #r²
        sum+=1/((r2-eta**2)*(r2+mu**2)**2)
    return 1/(2*np.sqrt(np.pi))*((mu**2-eta**2)**2*sum+gamma**2*np.pi**2/mu*(eta**2-mu**2))

def Z01(eta, n_max, mu, m, d3):
    cartesian = list(product(*[range(-n_max,n_max+1)]*3))
    eta2 = eta**2
    gamma2 = 1 + d3**2 / (4*m**2+eta2)
    sum=0
    for (n1,n2,n3) in cartesian:
        r2 = gamma2*(n3-d3/2)**2+n1**2+n2**2 #r²
        sum+=1/((r2-eta2)*(r2+mu**2)**2)
    return 1/(2*np.sqrt(np.pi))*((mu**2-eta2)**2*sum+gamma2*np.pi**2/mu*(eta2-mu**2))

def Z02(eta, n_max, mu, m, d3):
    mnmx = np.linspace(-nmax, +nmax, nmax+nmax+1)
    X, Y, DZ = np.meshgrid(mnmx, mnmx, mnmx-d3/2)
    eta2 = eta**2
    gamma2 = 1 + d3**2 / (4*m**2+eta2)
    sum=0
    r2 = (X**2 + Y**2 + np.einsum('ijk,l->lijk', DZ**2, gamma2)).transpose((1,2,3,0))
    sum = np.sum(1/((r2-eta2)*(r2+mu**2)**2), axis=(0,1,2))
    return 1/(2*np.sqrt(np.pi))*((mu**2-eta2)**2*sum+gamma2*np.pi**2/mu*(eta**2-mu**2))