Julia随机微分方程代码运行过慢：请求排查原因与优化方案

问题背景 #

我用Euler-Murayama方法求解耦合Chemical Langevin方程，模拟野生型（W）与突变型（M）细胞的随机波传播。已尝试向量化操作、@inbounds、列主索引、@floop并行、预分配数组等优化手段，但代码运行仍极慢（即使仅1个独立进程N=1，1e8时间步的计算也耗时惊人），急需排查瓶颈并获取加速方案。

方程说明 #

野生型细胞的CLE方程：
$$
dW_i(t) = \left[ r_w W_i(t) \left(K - W_i(t) - M_i(t)\right) + K\left(W_{i+1}(t) - 2W_i(t) + W_{i-1}(t)\right) \right]dt + \sqrt{6W_i(t)(K-W_i(t)) + (K-2W_i(t))(W_{i+1}(t)-2W_i(t)+W_{i-1}(t)) - r_w W_i(t)(K-W_i(t)-M_i(t))}dN(0,1)
$$

突变型细胞的CLE方程：
$$
dM_i(t) = \left[ r_m M_i(t) \left(K - W_i(t) - M_i(t)\right) + K\left(M_{i+1}(t) - 2M_i(t) + M_{i-1}(t)\right) \right]dt + \sqrt{6M_i(t)(K-M_i(t)) + (K-2M_i(t))(M_{i+1}(t)-2M_i(t)+M_{i-1}(t)) - r_m M_i(t)(K-W_i(t)-M_i(t))}dN(0,1)
$$

参数定义：

• K：环境容纳量（种群最大数量）
• i：种群分布区（deme）
• N(0,1)：标准正态随机变量
• $r_w$/$r_m$：野生型/突变型细胞的生长率

核心代码 #

using Random, Distributions
using StatsBase
using Statistics
using FLoops

# CLE 参数设置
K = 100 # 环境容纳量
M = 100 # 种群分布区数量
T = 100_000_000 # 时间步数
dt = 1e-1 # 时间步长
g = Normal(0.0, sqrt(dt)) # 正态分布（均值0，标准差√dt）
r_w = 0.1 # 野生型生长率
r_m = 0.2 # 突变型生长率
r_wm = [r_w, r_m]' # 生长率向量（转置）
N = 1 # 独立进程数（N=1时仍慢）

# 初始波型：野生型呈阶跃分布，第76个分布区初始有100个突变型
state_init = Matrix(reshape(repeat([K, 0.0]', M+2), (M+2, 2))) 
state_init[M÷2+2:end, 1] .= 0   
state_init[76, 2] = 100.0
state_init[1, :] .= [K, 0]
state_init[end, :] .= [0, 0]

state = deepcopy(state_init)
state_plus = zeros(size(state_init)) # 存储i+1分布区的状态（用于计算导数）
state_minus = zeros(size(state_init)) # 存储i-1分布区的状态（用于计算导数）

function sim!(state_init::Matrix{Float64}, state::Matrix{Float64},
              T::Int64, dt::Float64, N::Int64, M::Int64, K::Int64,
              hist_data::Array{Int64,3}, g::Normal{Float64})

    @inbounds @floop for n in 1:N
        state .= deepcopy(state_init) # 初始化状态
        @inbounds for t in 1:T
            
            state_plus .= circshift(state, -1) # 生成i+1状态
            state_plus[1, :] .= [0, 0] # 固定边界条件
            state_minus .= circshift(state, 1) # 生成i-1状态
            state_minus[end, :] .= [0, 0] # 固定边界条件
            state_shift = circshift(state, (0, 1)) 
            # 转换状态向量顺序：从[野生型, 突变型]变为[突变型, 野生型]

            # Euler-Murayama方法更新状态，用clamp限制细胞数在[0,K]范围内（避免根号下负数）
            state .= clamp.(state .+ 
            dt .* (r_wm .* state .* (K .- state .- state_shift) .+ 
            K .* (state_plus .- 2.0 .* state .+ state_minus)) .+ 
            sqrt.(clamp.( 
            6 .* state .* (K .- state) .+ 
            (K .- 2.0 .* state) .* (state_plus .- 2.0 .* state .+ 
            state_minus) .- r_wm .* state .* (K .- state .- state_shift), 
            0.0, 1.0*K*K)) .* 
            rand(g, M+2), 0.0, 1.0*K) 
        end
    end
end

sim!(state_init, state, T, dt, N, M, K, hist_data, g)

# 后续分析代码非性能瓶颈，此处省略

瓶颈分析 #

• circshift高频调用开销：每步循环调用3次circshift，每次都会复制整个数组，1e8次循环累计开销极大。
• 临时数组泛滥：复杂向量化表达式会生成大量中间临时数组，内存分配与回收的开销在大规模循环中被放大。
• 随机数重复分配：每次rand(g, M+2)都会创建新数组，未复用预分配内存。
• 冗余初始化操作：@floop内的deepcopy(state_init)是冗余操作，可直接内存拷贝替代。
• 边界条件重复赋值：每步都对固定边界重新赋值，无必要重复执行。

加速方案 #

1. 替换circshift为手动索引，消除数组复制 #

直接用索引访问相邻分布区，避免circshift的数组复制开销：

# 计算相邻项时直接用索引，例如：
# 对中间分布区2<=i<=M+1：
# state_plus_i = state[i+1, :]
# state_minus_i = state[i-1, :]

2. 预分配随机数数组，复用内存 #

提前创建随机数数组，每步仅更新值而非重新分配：

# 在函数开头预分配
rand_vals = similar(state)

# 循环内替换rand(g, M+2)为：
rand!(g, rand_vals)

3. 拆分复杂表达式，减少临时数组 #

将大表达式拆分为多步，用预分配的中间数组存储结果，避免临时数组创建：

# 预分配中间数组
growth_term = similar(state)
sqrt_arg = similar(state)
diffusion_term = similar(state)

@inbounds for t in 1:T
    # 计算生长项（手动处理边界）
    @. growth_term[2:end-1,:] = r_wm * state[2:end-1,:] * (K - state[2:end-1,:] - state_shift[2:end-1,:]) + K * (state[3:end,:] - 2state[2:end-1,:] + state[1:end-2,:])
    growth_term[1,:] .= 0.0
    growth_term[end,:] .= 0.0

    # 计算根号内项
    @. sqrt_arg[2:end-1,:] = 6 * state[2:end-1,:] * (K - state[2:end-1,:]) + (K - 2state[2:end-1,:]) * (state[3:end,:] - 2state[2:end-1,:] + state[1:end-2,:]) - r_wm * state[2:end-1,:] * (K - state[2:end-1,:] - state_shift[2:end-1,:])
    @. sqrt_arg = clamp(sqrt_arg, 0.0, 1.0*K*K)
    @. diffusion_term = sqrt(sqrt_arg)

    # 更新随机数
    rand!(g, rand_vals)

    # 更新状态
    @. state = clamp(state + dt * growth_term + diffusion_term * rand_vals, 0.0, 1.0*K)
end

4. 优化初始化逻辑，替换deepcopy #

用直接内存拷贝替代deepcopy，减少初始化开销：

# 替换state .= deepcopy(state_init)为：
state .= state_init

5. GPU加速（可选） #

若有GPU资源，使用CUDA.jl将数组转移到GPU计算，向量化操作在GPU上的并行效率远高于CPU，可大幅缩短大规模循环的耗时。

内容的提问来源于stack exchange，提问作者Jason M Gray