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 ms到25 ms之间的最佳案例结果。

我用来收集所有这些的辅助函数非常简单：

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秒）的最佳情况结果，其中10万颗粒将达到7400秒或大约2小时。我仍然可以使用它，但我很可能在将来添加模拟或时间步骤。

我最初使用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)我移动到函数之外，因为在快速测试中我发现它会成为性能瓶颈。

在python

1 个答案: