多轨迹随机数值积分效率优化：欧拉-马尔可夫方法并行与Numba适配

问题背景 #

我用Euler-Maruyama方法实现随机微分方程（SDE）的数值求解，现有代码能正常运行但效率极低。因为问题本身带随机性，需要计算大量独立轨迹，目前是逐个求解的方式。尝试直接去掉轨迹循环改成向量化形式但未生效，用Numba优化求解器也失败了。

现有可运行但低效的代码：

%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)

实际系统维度更大、轨迹数量更多，期望的高效求解器形式（目前无法正常运行）：

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

解决方案 #

1. 核心问题：A/B函数不支持批量计算 #

原代码中A、B函数仅处理单个轨迹的1维向量输入，无法适配向量化所需的(维度, 轨迹数)批量输入，这是直接去循环失败的根本原因。需先修改这两个函数支持批量计算：

修改后的A函数

def A(u):
    # u形状: (3, Mmax)
    x = u[0, :]
    y = u[1, :]
    z = u[2, :]
    
    omega = 0.5 * np.sqrt(1 + 8 * kappa * z**2)
    
    # 构造批量系数矩阵，形状: (3, 3, Mmax)
    A_mat = np.zeros((3, 3, u.shape[1]), dtype=np.complex64)
    A_mat[0, 0, :] = -2
    A_mat[0, 1, :] = omega
    A_mat[1, 0, :] = -omega
    A_mat[2, 2, :] = -kappa
    
    # 批量矩阵-向量乘法，结果形状: (3, Mmax)
    du = np.einsum('ijk,jk->ik', A_mat, u)
    return du

修改后的B函数

def B(u, w):
    # u形状: (3, Mmax), w形状: (Mmax,)
    g = np.sqrt(kappa * nth)
    
    # 构造批量噪声向量，形状: (3, Mmax)
    B_vec = np.zeros_like(u)
    B_vec[1, :] = g * w
    B_vec[2, :] = g * w
    return B_vec

2. 向量化SDE主函数 #

现在可以直接使用你期望的无轨迹循环形式，注意扩展初始向量维度以匹配批量输入：

def SDE(A, B):
    u = np.zeros((len(u0), Nmax+1, Mmax), dtype=np.complex64)
    u[:, 0, :] = u0[:, np.newaxis]  # 把u0从(3,)扩展为(3,1)，适配轨迹维度
    for n in range(Nmax):
        u[:, n+1, :] = u[:, n, :] + dt * A(u[:, n, :]) + B(u[:, n, :], w[n, :]) * np.sqrt(dt)
    return u

3. Numba优化（针对超大规模场景） #

若向量化后效率仍不足，用Numba的njit装饰器编译优化，注意遵循Numba的数组操作规则：

from numba import njit, prange

@njit(parallel=True)
def A_numba(u, kappa):
    # u形状: (3, Mmax)
    x = u[0, :]
    y = u[1, :]
    z = u[2, :]
    du = np.zeros_like(u)
    
    for m in prange(u.shape[1]):
        omega = 0.5 * np.sqrt(1 + 8 * kappa * z[m]**2)
        du[0, m] = -2 * x[m] + omega * y[m]
        du[1, m] = -omega * x[m]
        du[2, m] = -kappa * z[m]
    return du

@njit(parallel=True)
def B_numba(w, kappa, nth):
    # w形状: (Mmax,)
    g = np.sqrt(kappa * nth)
    B_vec = np.zeros((3, w.shape[0]), dtype=np.complex64)
    B_vec[1, :] = g * w
    B_vec[2, :] = g * w
    return B_vec

@njit(parallel=True)
def SDE_numba(u0, dt, sqrt_dt, Nmax, Mmax, w, kappa, nth):
    u = np.zeros((u0.shape[0], Nmax+1, Mmax), dtype=np.complex64)
    u[:, 0, :] = u0[:, np.newaxis]
    
    for n in range(Nmax):
        current_u = u[:, n, :]
        u[:, n+1, :] = current_u + dt * A_numba(current_u, kappa) + B_numba(w[n, :], kappa, nth) * sqrt_dt
    return u

调用Numba版本代码

# 预计算sqrt_dt避免循环内重复计算
sqrt_dt = np.sqrt(dt)

# 调用优化后的求解器
u_t = SDE_numba(u0, dt, sqrt_dt, Nmax, Mmax, w, kappa, nth)

4. 额外优化建议 #

• 预计算所有常数（如sqrt_dt、g），避免在循环/函数内重复计算
• 超大规模轨迹场景可结合multiprocessing做轨迹分块并行，但Numba的parallel=True已能覆盖多数场景
• 避免在循环内创建新数组，尽量复用内存

内容的提问来源于stack exchange，提问作者J.Agusti