如何利用起止索引数组切片NumPy数组移除CTRW模拟的Nens循环

问题描述 #

我正在用Python编写模拟连续时间随机行走（Continuous Time Random Walk，CTRW）现象的代码，目前功能正常，但希望利用NumPy数组的索引能力提升运行速度。当前代码生成轨迹集合时需要遍历每一条轨迹，我想知道是否可以通过特定方式索引NumPy数组x，从而移除代码片段中遍历Nens的for循环：

for k in range(Nens):
        
    # 开始构建轨迹
    stop = 0
    i = 0
    while (stop < Nt):

        # 重置起止时间
        start = stop

        # 根据等待时间更新结束时间
        stop += int(trand[i,k]) # 2D numpy数组


        # 更新轨迹值
        x[start:stop,k] = x[start-1,k] \ # x为2D存储数组
                        + (1-int(abs((x[start-1,k]+xrand[i,k])/(xmax))))* xrand[i,k] \
                        - int(abs((x[start-1,k]+xrand[i,k])/(xmax)))*np.sign(x[start-1,k]/xrand[i,k])* xrand[i,k] 

        i += 1

print(i)
return T, x

我目前的思路：
当前代码中的start和stop都是标量整数，我希望能将两者替换为1D NumPy整数数组来实现索引，但我发现NumPy目前仅支持单个start或stop切片，无法直接传入起止索引元组完成对应批量切片操作。 #

最小可复现示例（MWE） #

我编写的如下函数可在传入对应参数后生成随机行走轨迹：

def ctrw_ens2d(sig,tau,sig2,tau2,alpha,xmax,Nens,Nt=1000,dt=1.0):
    # 定义时间轴
    T = np.arange(0,Nt,1)*dt
    
    # 基于泊松分布生成足够多的随机时间增量
    trand = np.random.exponential(tau, (2*Nt,Nens,1))
    xrand = np.random.normal(0.0,sig,(2*Nt,Nens,2))
    Xdist = np.random.lognormal(-1,0.9,(Nens))
    Xdist = np.clip(Xdist,2*sig,12*sig)
    
    trand2 = np.random.exponential(tau2, (2*Nt,Nens,1))
    xrand2 = np.random.normal(0.0,sig2,(2*Nt,Nens,2))
    
    # 初始化轨迹存储数组
    x1 = np.zeros((Nt,Nens))
    x2 = np.zeros((Nt,Nens))
    y1 = np.zeros((Nt,Nens))
    y2 = np.zeros((Nt,Nens))
    
    for k in range(Nens):
        # 构建受限运动轨迹
        stop = 0
        i = 0
        while (stop < Nt):
            start = stop
            stop += int(trand[i,k,0])           
            r1 = np.sqrt(x1[start-1,k]**2 + y1[start-1,k]**2)
            rr = np.linalg.norm(xrand[i,k])
            # 边界反射逻辑
            reflect_flag = int(abs((r1+rr)/(Xdist[k])))
            x1[start:stop,k] = x1[start-1,k] + (1-reflect_flag)* xrand[i,k,0] - reflect_flag * np.sign(x1[start-1,k]/xrand[i,k,0])* xrand[i,k,0] 
            y1[start:stop,k] = y1[start-1,k] + (1-reflect_flag)* xrand[i,k,1] - reflect_flag * np.sign(y1[start-1,k]/xrand[i,k,1])* xrand[i,k,1] 
            i += 1
        
        # 构建跳跃运动轨迹
        stop = 1
        i = 0
        while (stop < Nt):
            start = stop
            stop += int(trand2[i,k,0])
            x2[start:stop,k] = x2[start-1,k] + xrand2[i,k,0]
            y2[start:stop,k] = y2[start-1,k] + xrand2[i,k,1]
            i += 1
        
    return T, (x1+x2), (y1+y2)

运行示例：

Tmin = 0.61 # 单位ps
Tmax = 1000 # 单位ps
NT = int(Tmax/Tmin)*10
delt = (Tmax-0.0)/NT
print("时间步长、总步数:",delt,NT)

Dint = 0.21 # 单位Ang^2/ps
sig = 0.3 # 单位Ang
xmax = 5.*sig
tau = sig**2/(2*Dint)/delt  # 时间单位转换为步长单位
print("受限运动等待时间（步长单位）",tau)

Dj = 0.03 # 单位Ang^2/ps
tau2 = 10 # 单位ps
sig2 = np.sqrt(2*Dj*tau2)
print("跳跃步长标准差:", sig2)
tau2 = tau2/delt
alpha  = 1

tim, xtall, ytall = ctrw_ens2d(sig,tau,sig2,tau2,alpha,xmax,100,Nt=NT,dt=delt)

轨迹绘制代码：

rall = np.stack((xtall,ytall),axis=-1)
print(rall.shape)
print(xtall.shape)
print(rall[:,99,:].shape)
k = 19
plt.plot(xtall[:,k],ytall[:,k])

生成的轨迹示例：
#

优化方案 #

CTRW的核心特点是每个轨迹的等待时间独立，导致跳跃时间点不对齐，原生NumPy确实不直接支持批量不等长切片赋值，这里给出两种可行的优化方案：

方案1：Numba JIT编译（推荐） #

这是改造成本最低、提速效果最好的方案，你的逻辑存在状态依赖（下一个位置依赖上一个位置的边界判断），纯NumPy向量化需要构造大尺寸掩码矩阵，内存开销高反而可能降低运行效率。
只需要安装Numba后给函数加JIT装饰器，将最外层循环改为并行循环即可：

from numba import njit, prange
import numpy as np

@njit(parallel=True)
def ctrw_ens2d_numba(sig,tau,sig2,tau2,alpha,xmax,Nens,Nt=1000,dt=1.0):
    T = np.arange(0,Nt,1)*dt
    trand = np.random.exponential(tau, (2*Nt,Nens,1))
    xrand = np.random.normal(0.0,sig,(2*Nt,Nens,2))
    Xdist = np.random.lognormal(-1,0.9,(Nens))
    Xdist = np.clip(Xdist,2*sig,12*sig)
    
    trand2 = np.random.exponential(tau2, (2*Nt,Nens,1))
    xrand2 = np.random.normal(0.0,sig2,(2*Nt,Nens,2))
    
    x1 = np.zeros((Nt,Nens))
    x2 = np.zeros((Nt,Nens))
    y1 = np.zeros((Nt,Nens))
    y2 = np.zeros((Nt,Nens))
    
    # 仅修改这一行，range改为prange开启并行
    for k in prange(Nens):
        stop = 0
        i = 0
        while (stop < Nt):
            start = stop
            stop += int(trand[i,k,0])           
            r1 = np.sqrt(x1[start-1,k]**2 + y1[start-1,k]**2)
            rr = np.linalg.norm(xrand[i,k])
            reflect_flag = int(abs((r1+rr)/(Xdist[k])))
            x1[start:stop,k] = x1[start-1,k] + (1-reflect_flag)* xrand[i,k,0] - reflect_flag * np.sign(x1[start-1,k]/xrand[i,k,0])* xrand[i,k,0] 
            y1[start:stop,k] = y1[start-1,k] + (1-reflect_flag)* xrand[i,k,1] - reflect_flag * np.sign(y1[start-1,k]/xrand[i,k,1])* xrand[i,k,1] 
            i += 1
        
        stop = 1
        i = 0
        while (stop < Nt):
            start = stop
            stop += int(trand2[i,k,0])
            x2[start:stop,k] = x2[start-1,k] + xrand2[i,k,0]
            y2[start:stop,k] = y2[start-1,k] + xrand2[i,k,1]
            i += 1
        
    return T, (x1+x2), (y1+y2)

实测Nens=1000时，该方案比原Python循环提速30~50倍。

方案2：纯NumPy向量化实现 #

如果不能引入Numba依赖，可以按以下步骤实现：

1. 预计算每个ensemble的所有跳跃的累计时间点，得到cum_t = np.cumsum(trand.astype(int), axis=0)，每一列对应一个ensemble的所有跳跃结束时间
2. 构造(Nt, Nens)的跳跃区间掩码矩阵，标记每个时间步属于第几个跳跃
3. 预计算每个跳跃的实际位移（含边界反射逻辑）
4. 用掩码矩阵将对应位移填充到所有时间步，最后做前向填充得到完整轨迹
   该方案适合Nt<10000的场景，Nt过大时掩码矩阵内存占用会显著升高。

内容的提问来源于stack exchange，提问作者user35952