使用scipy.integrate.solve_ivp向量化求解多初值二阶系统出错求助

我尝试用scipy.integrate.solve_ivp并设置vectorized=True，求解包含N个初始状态（x, y, xdot, ydot）的二阶系统，以此避免循环。但得到的结果和常规循环求解的合理结果差异极大。

我的操作是：将传入solve_ivp的扁平初始状态（大小为4N）重塑为(4, N)的形状，第一行是所有x值，第二行是所有y值，第三行是所有xdot值，第四行是所有ydot值。ODE函数返回扁平的[[dxdt], [dydt], [ddxdtt], [ddydtt]]（大小为4N的一维数组），二阶导数ddxdtt使用预先计算好的数组，所有计算通过numpy数组完成。

代码如下：

import numpy as np
import matplotlib.pyplot as plt
from scipy.integrate import solve_ivp

x = np.linspace(-3, 3, 7)
y = np.linspace(-3, 3, 7)
X, Y = np.meshgrid(x, y)
a = 2
inside = np.sqrt(X**2 + Y**2) <= a
outside = np.sqrt(X**2 + Y**2) > a

# 设置背景场
Ux = 2 * X ** 2 + 3 * Y
Ux[outside] = 5 * X[outside] - Y[outside]
dUxdx = 4 * X
dUxdx[outside] = 5
dUxdy = 3 * np.ones_like(Ux)
dUxdy[outside] = -1 
dUxdt = np.zeros_like(Ux)

Uy = -4 * Y + X ** 2
Uy[outside] = X[outside] + 2 * Y[outside]
dUydx = 2 * X
dUydx[outside] = 1
dUydy = -4 * np.ones_like(Uy)
dUydy[outside] = 2
dUydt = np.zeros_like(Uy)

grids = np.row_stack((X.flatten(), Y.flatten()))
IC = np.row_stack((X.flatten(), Y.flatten(), Ux.flatten(), Uy.flatten()))
input = IC.flatten()

def func_vec(t, Z):
    
    Z = np.reshape(Z, (4, -1))

    dZdt = np.zeros_like(Z)
    dZdt[0, :] = Z[2, :]
    dZdt[1, :] = Z[3, :]
    dZdt[2, :] = Ux.flatten() * dUxdx.flatten() + Uy.flatten() * dUxdy.flatten()
    dZdt[3, :] = Ux.flatten() * dUydx.flatten() + Uy.flatten() * dUydy.flatten()

    out = dZdt.flatten()
    return out

sol_vec = solve_ivp(func_vec, [0, 10], input, 'RK45', np.linspace(0, 10, 10), vectorized=True)
sol_vec_state = sol_vec.y.reshape((4, -1, len(sol_vec.t)))
sol_vec_x = sol_vec_state[0, :, :].reshape((X.shape[0], X.shape[1], len(sol_vec.t)))
sol_vec_y = sol_vec_state[1, :, :].reshape((X.shape[0], X.shape[1], len(sol_vec.t)))

你的向量化使用方式结构上是正确的，但核心错误在于：ODE函数中计算二阶导数时，使用的是初始时刻预计算的Ux、Uy及其导数数组，而没有根据每个粒子当前的位置(x,y)重新计算对应的背景场值。

在循环求解中，每个粒子每一步都会根据当前位置计算对应的场量；但你的向量化代码里，一直用初始位置的场值，导致积分过程中粒子的加速度始终是初始位置的加速度，这自然会和循环求解结果产生巨大差异。

修正后的ODE函数 #

需要根据当前Z中的x（Z[0,:]）和y（Z[1,:]），为每个粒子实时计算Ux、dUxdx等值：

def func_vec(t, Z):
    Z = np.reshape(Z, (4, -1))
    x_current = Z[0, :]
    y_current = Z[1, :]
    
    # 计算每个当前位置对应的inside/outside标记
    r_current = np.sqrt(x_current**2 + y_current**2)
    inside_current = r_current <= a
    outside_current = r_current > a
    
    # 计算当前位置的Ux及导数
    Ux_current = np.where(inside_current, 2 * x_current**2 + 3 * y_current, 5 * x_current - y_current)
    dUxdx_current = np.where(inside_current, 4 * x_current, 5)
    dUxdy_current = np.where(inside_current, 3, -1)
    
    # 计算当前位置的Uy及导数
    Uy_current = np.where(inside_current, -4 * y_current + x_current**2, x_current + 2 * y_current)
    dUydx_current = np.where(inside_current, 2 * x_current, 1)
    dUydy_current = np.where(inside_current, -4, 2)
    
    dZdt = np.zeros_like(Z)
    dZdt[0, :] = Z[2, :]
    dZdt[1, :] = Z[3, :]
    dZdt[2, :] = Ux_current * dUxdx_current + Uy_current * dUxdy_current
    dZdt[3, :] = Ux_current * dUydx_current + Uy_current * dUydy_current
    
    return dZdt.flatten()

这样修改后，向量化求解的结果就会和循环求解的结果一致，因为每一步都根据粒子的当前位置计算了正确的加速度。

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