使用 Scipy 加速许多粒子的路径计算答案

【问题标题】：Speeding up path calculation of many particles using Scipy使用 Scipy 加速许多粒子的路径计算
【发布时间】：2018-03-17 06:06:47
【问题描述】：

我正在尝试使用 Python 计算许多粒子位置的演变。最终，我将在大约 100000 个时间步长上计算许多粒子（大约 10000 个）。由于我不惜一切代价避免使用 Fortran，因此我正在努力加快这一进程。

我要解决的方程是

d X_i/dt = u
d Y_i/dt = v

所以理想情况下，我会使用二维数组：一个在粒子之间切换，另一个在x 和y 之间切换。因此，如果我有 100 个粒子，我将有一个 100x2 数组。

问题的一部分是因为scipy.integrate.odeint只需要一维数组，所以我必须将我的初始条件展平，将其拆分到我的导函数（RHS_im）中，然后在输出时再次展平，这很慢（这约占RHS_im 调用的 20%）。

我可以用丑陋的方式做到这一点。这是我提出的 MWE

import numpy as np
from scipy.interpolate import RegularGridInterpolator
from scipy.integrate import odeint

x=np.arange(2.5, 500, 5)
y=np.arange(2.5, 500, 5)

X, Y = np.meshgrid(x, y, indexing='xy')

U=np.full_like(X, -0.1)
V=0.2*np.sin(2*np.pi*X/200)

tend=100
dt=1
tsteps=np.arange(0, tend, dt)

#------
# Create interpolation
U_int = RegularGridInterpolator((x, y), U.T)
V_int = RegularGridInterpolator((x, y), V.T)
#------

#------
# Initial conditions
x0=np.linspace(200,300,5)
y0=np.linspace(200,300,5)
#------


#------
# Calculation for many
def RHS_im(XY, t):
    X, Y = np.split(XY, 2, axis=0)
    pts=np.array([X,Y]).T
    return np.concatenate([ U_int(pts), V_int(pts) ])
XY0 = np.concatenate([x0, y0])
XY, info = odeint(RHS_im, XY0, tsteps[:-1], args=(), hmax=dt, hmin=dt, atol=0.1, full_output=True)
X, Y = np.split(XY, 2, axis=1)

有没有办法避免拆分过程？

此外，虽然我选择了odeint 来集成这个系统，但我并不承诺这个选择。如果有另一个函数可以更快地执行此操作（即使使用简单的欧拉方案），我会很容易地改变。我只是没有找到。（插值方案也是如此，这大约是RHS_im 所用时间的 80%）。

编辑

稍微改进了拆分过程（使用np.split），但整体程序仍需改进。

【问题讨论】：

澄清一下：您是从 5x5 网格上的初始 x,y 条件开始，然后在导函数中插入 100x100 网格的 u,v 值吗？
@binaryfunt 不，网格已经是100x100。我认为您将arange 与linspace 混淆了。
但是当你设置x0=np.linspace(200,300,5)时，x0的长度为5
@binaryfunt 我现在明白你的意思了。那就是粒子的数量。我给出了 5 个初始条件，每个粒子一个。我不确定这是否能回答您的问题。
需要使用网格插值吗？

标签： python numpy scipy numerical-methods

【解决方案1】：

感谢您对 MWE 的澄清。通过将U_int 和V_int 组合成UV_int 并更改XY 的布局以便可以使用np.reshape 代替np.split，我能够使MWE 的运行速度提高约50%。这两项更改都有助于就地且连续地访问内存。

x_grid=np.arange(2.5, 500, 5)
y_grid=np.arange(2.5, 500, 5)

x_mesh, y_mesh = np.meshgrid(x_grid, y_grid, indexing='xy')

u_mesh=(np.full_like(x_mesh, -0.1))
v_mesh=(0.2*np.sin(2*np.pi*x_mesh/200))
uv_mesh = np.stack([u_mesh.T, v_mesh.T], axis=-1)

tend=100
dt=1
tsteps=np.arange(0, tend, dt)

#------
# Create interpolation
UV_int = RegularGridInterpolator((x_grid, y_grid), uv_mesh)
#------

#------
# Initial conditions
npart = 5
x0=np.linspace(200,300,npart)
y0=np.linspace(200,300,npart)

XY_pair = np.reshape(np.column_stack([x0, y0]), (2*npart))
#-----


def RHS_im(XY, t):        
    return np.reshape(UV_int(np.reshape(XY, (npart, 2))), (npart*2))

XY, info =  odeint(RHS_im, XY_pair, tsteps[:-1], args=(), hmax=dt, hmin=dt, atol=0.1, full_output=True)

【讨论】：