利用粒子粒子加速路径计算

Question

我试图用Python计算许多粒子位置的演化。最终，我将有许多粒子(10000左右)在大约100000步的时间内被计算出来。由于我正不惜一切代价避免使用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来集成这个系统，但我并没有致力于这个选择。如果有另一个更快的函数(即使使用简单的Euler方案)，我会很容易地进行更改。我只是没找到。(插补方案也是如此，大约是RHS_im所需时间的80% )。

编辑

(通过使用np.split)稍微改进了分裂过程，但是整个程序仍然需要改进。

Answer 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)