多个轨迹上的高效随机数值积分

Question

我正在实施一种使用 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。

Answer 1

以下是我在提供的代码中发现的性能问题：

• A.dot(u)

  在这里显然是多余的，因为

  A

  矩阵是一个只有 4 个非零值的 3x3 矩阵。最好手动计算表达式，尤其是使用 Numba。
• 创建许多小型数组，例如 3x3 2D 数组，甚至只有 3 个项目的 1D 数组，成本非常昂贵，即使在 Numba 中也是如此（因为数组是动态分配和引用计数的）。您应该预先分配数组，甚至只使用 3 个变量（效率更高）。因此，在 Numba 中，返回新数组通常比写入参数中传递的现有预分配数组效率低（出于上述原因）。
• 在最后一个轴上迭代通常更快（在连续数组上），因为最后一个轴的项目连续存储在内存中，而其他轴的情况则不然。因此，在计算

  u

  时交换第一个和最后一个轴应该会更快。它的形状是

  (Mmax,Nmax+1,len(u0))

  。有关更多信息，请阅读：AoS 与 SoA。
• 您可以预先计算一些常量，特别是当值乘以 0 时。

此外，

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 中，创建新数组仍然相当昂贵） .