在Python中优化Euler-Maruyama实现

Question

我的任务是解决一个特定的随机微分方程（SDE）和一个普通的DE，它们都相互依赖。我一直在使用Euler-Maruyama method来求解这些方程，并且目前需要解决其中的10万个（它们模拟粒子路径和动量）。代码本身工作正常，但问题是，因为我需要增加我的算法的时间步数，所以计算时间自然也会增加。我最初选择python作为我最习惯的任务，尽管我非常清楚它的声誉并不是HPC（慢循环）的最佳选择。我可能不得不改变使用python作为Fortran子程序的“胶水”，但是我想知道是否仍然有某种方法可以从当前代码中挖掘出一点性能。

我目前正在使用函数EM来解决（模拟）DE，然后收集所有模拟时间，路径和动作并将它们附加到它们各自的数组的辅助函数。代码如下所示：

def EM(tstart, tend, steps, x1, x2, x3, x4, x5): 
    dt = float((tend - tstart))/steps  # Timestep 
    T = np.linspace(tstart, tend, steps) 
    y = [x3*T[0] - 0.01] # Table for the paths 
    p = [x3*x2*x4] # Table for the momentum 
    pos = 0.0 
    mom = 0.0 

    for i in range(1, steps): 
     pos = y[i-1] + dt*(((x1*y[i-1])/(x3*T[i-1])) + (3.0*p[i-1]*(T[0]/T[i-1])/x2)*((2*(x3*T[i-1]-y[i-1])/(y[i-1]+x5))-1)) + (np.sqrt(6*(p[i-1]*(T[0]/T[i-1])/x2)*(x3*T[i-1]-y[i-1])) * np.sqrt(dt) * np.random.normal(0,1)) 

     #Boundary condition 
     if(pos > x3*T[i]): 
      v = (pos-y[i-1])/dt 
      tdot = (y[i-1]-v*T[i-1])/(x3-v) 
      pos = x3*tdot - v*(T[i-1]+dt-tdot) 

     mom = p[i-1] - (1.0/3.0)*p[i-1]*(x1/(x3*T[i-1]))*(1 + (2*y[i-1]/(y[i-1]+x5)))*dt 

     y.append(pos) 
     p.append(mom) 

     #Boundary condition    
     if(pos < 0): 
      break 



    return T[0:i+1], y, p 

其中x1，...，x5是某些常量。截至目前，我需要使用108个000时间步骤，并与内置％timeit测试运行的代码

%timeit EM(1.0, 10.0, 108000, 1.0, 1.0, 2.0, 3.0, 1.0) 

给了我65毫秒之间最好的情况下，结果25毫秒。

我用它来收集所有的这些共同的辅助功能是相当简单：

def helper(tstart, tend, steps, x1, x2, x3, x4, x5, no_of): 
    timespan = [] 
    momentums = [] 
    paths = [] 
    for i in range(0, no_of): 
      t, y, p = EM(tstart, tend, steps, x1, x2, x3, x4, x5) 
      timespan.append(t) 
      paths.append(y) 
      momentums.append(p) 

    return timespan, paths, momentums 

通过timeit以下参数

%timeit multi(1.0, 10.0, 108000, 1.0, 1.0, 2.0, 3.0, 1.0, 1000) 

得到1分钟的最好情况下的结果，并运行此14秒（74秒），其中100000个颗粒将达到7400秒或大约两个小时。我仍然可以用这个工作，但很可能我必须在未来增加模拟或时间步骤。

我最初使用numpy数组，但更改为常规python列表实际上使代码更快。我猜这是因为你必须在np.zeros（）方法使用它们之前声明numpy数组的大小（除非你想使用np.append方法，但在这种情况下非常慢）。因此，虽然使用的步骤数量是108000，但只有一小部分模拟结果很长，所以我最终需要使用np.trim_zeros（）从数组中修剪零点。

我一直在尝试使用Numba library，它的@jit方法，但我无法让它工作。它给了我下面的错误：

NotImplementedError: Failed at nopython (nopython frontend) 
    (<class 'numba.ir.Expr'>, build_list(items=[Var($0.20, <ipython-input- 32-4529989cafc0> (5))])) 

能消除帮助功能，只需用for循环是追加模拟阵列提高了运行时运行的代码？有没有办法运行代码而不使用for-loops，而是使用数组操作？我听说这会加快速度。

还有其他想法吗？非常感谢您的帮助。

Answer 1

由于此问题的迭代性质，无法用数组操作替换循环。

作为替代，我认为Numba确实是一个不错的选择。 Numba不适用于Python列表（因此您收到异常），所以您只能使用Numpy数组。为了处理先前未知的数组大小，ndarray.resize实例方法使用起来很好，因为它释放了未使用的内存（与采用保留对整个数组的引用的切片相反）。该守则将是这个样子：

from numba import jit 

@jit(nopython=True) 
def EM_helper(T, x1, x2, x3, x5, y, p): 
    dt = (T[-1] - T[0])/T.shape[0] 
    for i in range(1, T.shape[0]): 

     # ...big calculation 

     y[i] = pos 
     p[i] = mom 

     #Boundary condition    
     if(pos < 0): 
      break 

    return i+1 

def EM(T, x1, x2, x3, x4, x5): 
    y = np.empty_like(T, dtype=float) # Table for the path 
    p = np.empty_like(T, dtype=float) # Table for the momentum 

    y[0] = x3*T[0] - 0.01 
    p[0] = x3*x2*x4 
    count = EM_helper(T, x1, x2, x3, x5, y, p) 

    y.resize(count) 
    p.resize(count) 
    return T[:count], y, p 

取而代之的是EM_helper功能，你也可以尝试依靠自动“循环升降”，但是这是我的经验不太可靠。

时间阵列T = np.linspace(tstart, tend, steps)我感动的功能之外，因为在快速的测试，我发现它会成为性能瓶颈的产生。