我正在实施一种使用 Euler-Maruyama 方法求解随机微分方程的数值方法。
我所拥有的有效,但效率不高。原因是由于问题的随机性,我有很多轨迹。现在我正在一一解决。我感觉我应该能够并行化它们,因为它们是独立的。
工作代码如下所示
%matplotlib inline
import matplotlib.pyplot as plt
import numpy as np
import time
from numba import jit, njit
import os
def A(u):
x=u[0]
y=u[1]
z=u[2]
omega=1/2*np.sqrt((1+8*kappa*z*z))
A=np.array([[-2,omega,0],
[-omega,0,0],
[0,0,-kappa]])
du=A.dot(u)
return du
def B(u,w):
x=u[0]
y=u[1]
z=u[2]
g=np.sqrt(kappa*nth)
B=np.array([[0],
[1],
[1]])*g
return np.reshape(B*w,len(u0))
def SDE(A,B):
u = np.zeros((len(u0),Nmax+1,Mmax),dtype=np.complex64)
for m in range(Mmax):
u[:,0,m]=u0
for n in range(0,Nmax):
u[:,n+1,m] = u[:,n,m]+dt*A(u[:,n,m])+B(u[:,n,m],w[n,m])*np.sqrt(dt)
return u
#Parameters
kappa=0.05
nth=1.
gamma=1
Mmax=100 #number of trajectories
Tmax=10. ##max value for time
dt=0.05
Nmax=int(Tmax/dt) ##number of steps
t_list=np.arange(0,Tmax+dt/2,dt)
w = np.random.randn(Nmax+1,Mmax)
u0 = np.array([1., 0., np.sqrt(nth)/2])
u_t=SDE(A,B)
u_mean=np.mean(u_t,axis=2)
此代码是我的真实代码的简化,其中我有更大的系统维度和更多的轨迹。
请注意,这效率不高,因为当我增加 Mmax 时,我必须循环遍历它们。
理想情况下,我希望我的解算器看起来像这样
def SDE(A,B):
u = np.zeros((len(u0),Nmax+1,Mmax),dtype=np.complex64)
u[:,0,:]=u0
for n in range(0,Nmax):
u[:,n+1,:] = u[:,n,:]+dt*A(u[:,n,:])+B(u[:,n,:],w[n,:])*np.sqrt(dt)
return u
即,能够忽略 m 上的循环并以并行方式进行。然而,天真地这样做是行不通的。
提高效率的另一种理想方法是使用 Numba。然而,经过多次尝试,我一直无法用我定义的SDE求解器来实现njit。
以下是我在提供的代码中发现的性能问题:
A.dot(u)
在这里显然是多余的,因为 A
矩阵是一个只有 4 个非零值的 3x3 矩阵。最好手动计算表达式,尤其是使用 Numba。u
时交换第一个和最后一个轴应该会更快。它的形状是(Mmax,Nmax+1,len(u0))
。有关更多信息,请阅读:AoS 与 SoA。此外,
np.reshape(B*w,len(u0))
令人困惑,因为 B 的大小为 3,因此 u0
的大小也必须为 3,而 reshape
似乎毫无用处。请注意,np.complex64
用于简单精度(如评论中所述)。
这是生成的 Numba 代码:
import numba as nb # new
@nb.njit('(complex128[:], complex128[::1])')
def A(u, res):
x=u[0]
y=u[1]
z=u[2]
omega = 0.5 * np.sqrt((1+8*kappa*z*z))
res[0] = -2 * x + omega * y
res[1] = -omega * x
res[2] = -kappa * z
return res
@nb.njit('(complex128[:], float64, complex128[::1])')
def B(u, w, res):
g = np.sqrt(kappa*nth)
res[0] = 0
res[1] = g * w
res[2] = g * w
return res
# No signature is provided so the first call will be much slower
# But providing a signature here is complicated since A and B are functions
@nb.njit
def SDE(A,B):
u = np.zeros((len(u0),Nmax+1,Mmax), dtype=np.complex128)
sqrt_dt = np.sqrt(dt)
for m in range(Mmax):
u[:,0,m] = u0
tmp1 = np.empty(3, dtype=np.complex128)
tmp2 = np.empty(3, dtype=np.complex128)
for n in range(0,Nmax):
A(u[:,n,m],tmp1)
B(u[:,n,m],w[n,m],tmp2)
for i in range(3):
u[i,n+1,m] = u[i,n,m] + dt * tmp1[i] + tmp2[i] * sqrt_dt
return u
在我的机器上(配备 i5-9600KF CPU),速度快了约 500 倍。
我认为一旦代码优化,就不需要使用多线程,因为计算最终相当快。如果这还不够,您可以添加标志parallel=True
并将
for m in range(Mmax)
替换为
for m in nb.prange(Mmax)
。然而,由于不良的内存布局导致错误共享,这将无法很好地扩展。如上表所述,您应该交换轴 1 和轴 3 来解决此问题。最后,一旦并行化并具有更好的内存布局,最终代码应该如下所示:
# A and B are the same as before
# No signature is provided so the first call will be much slower
# But providing a signature here is complicated since A and B are functions
@nb.njit(parallel=True)
def SDE(A,B):
u = np.zeros((Mmax,Nmax+1,len(u0)), dtype=np.complex128)
sqrt_dt = np.sqrt(dt)
for m in nb.prange(Mmax):
u[m,0,:] = u0
tmp1 = np.empty(3, dtype=np.complex128)
tmp2 = np.empty(3, dtype=np.complex128)
for n in range(0,Nmax):
A(u[m,n,:],tmp1)
B(u[m,n,:],w[m,n],tmp2)
for i in range(3):
u[m,n+1,i] = u[m,n,i] + dt * tmp1[i] + tmp2[i] * sqrt_dt
return u
# [...] same code
w = np.random.randn(Nmax+1,Mmax).T.copy()
u0 = np.array([1., 0., np.sqrt(nth)/2])
u_t=SDE(A,B)
u_mean=np.mean(u_t,axis=0)
此代码快了约 2000 倍。
注意w
是一个全局变量,因此它被 Numba 视为编译时常量(即它在应用程序生命周期内永远不会改变),如果不是这种情况,您应该将其传递到参数中。顺便说一句,请注意,对于此类计算,Julia 可能是一种更好的语言,因为标准实现是 JIT 编译器,我们可以轻松避免创建新的临时数组(尽管即使在 Julia 中,创建新数组仍然相当昂贵) .