如何用scipy.odeint处理多维初始条件数组，避免循环遍历？

我希望将scipy.integrate.odeint应用于一组初始条件数组，返回与初始条件尺寸一致的结果。若采用循环遍历每个初始条件的方式，当N较大时速度会很慢。以下示例中sol_1可正常运行，但sol_2报错ValueError: Initial condition y0 must be one-dimensional.，求无需循环的解决方案。

原代码示例 #

import numpy as np
from scipy.integrate import odeint

def pend(y, t, b, c):
    theta, omega = y
    dydt = [omega, -b*omega - c*np.sin(theta)]
    return dydt

b = 0.25
c = 5.0
t = np.linspace(0, 10, 101)

# 0D初始值，来自官方文档
y0_1 = [np.pi - 0.1, 0]
sol_1 = odeint(pend, y0_1, t, args=(b, c))

# 1D初始值，来自官方文档
y0_2 = [np.ones((3)) * (np.pi - 0.1), np.zeros((3))]

# 此处报错
sol2 = odeint(pend, y0_2, t, args=(b, c))

无需循环的解决方案 #

核心思路是将批量初始条件扁平化处理，同时修改微分方程函数适配向量化计算，利用numpy的广播特性一次性计算所有初始条件的导数，完全避免循环开销。

步骤1：调整初始条件为一维数组 #

把多组初始条件的状态值（theta、omega）依次拼接成一维数组，符合odeint对y0的格式要求：

# 转换为一维数组：[theta1, theta2, theta3, omega1, omega2, omega3]
y0_2_flat = np.concatenate([np.ones(3)*(np.pi-0.1), np.zeros(3)])

步骤2：修改微分方程函数适配批量计算 #

重构函数，让它能处理扁平化后的批量状态，利用numpy向量化运算同时计算所有样本的导数：

def pend_batch(y, t, b, c):
    # 根据总长度拆分theta和omega数组
    n = len(y) // 2
    theta = y[:n]
    omega = y[n:]
    # 向量化计算所有样本的导数并拼接返回
    dydt = np.concatenate([omega, -b*omega - c*np.sin(theta)])
    return dydt

步骤3：求解并还原结果形状 #

调用odeint后，将扁平化的结果重新调整为(时间步数, 初始条件数, 状态数)的直观形状：

sol2 = odeint(pend_batch, y0_2_flat, t, args=(b, c))
# 还原形状：(101, 3, 2)，对应每个时间点、每个初始条件的theta和omega
sol2_reshaped = sol2.reshape(len(t), 3, 2)

完整可运行代码 #

import numpy as np
from scipy.integrate import odeint

def pend_batch(y, t, b, c):
    n = len(y) // 2
    theta = y[:n]
    omega = y[n:]
    dydt = np.concatenate([omega, -b*omega - c*np.sin(theta)])
    return dydt

b = 0.25
c = 5.0
t = np.linspace(0, 10, 101)

# 单组初始条件验证
y0_1 = [np.pi - 0.1, 0]
sol_1 = odeint(pend_batch, y0_1, t, args=(b, c))

# 批量初始条件求解
y0_2_flat = np.concatenate([np.ones(3)*(np.pi-0.1), np.zeros(3)])
sol2 = odeint(pend_batch, y0_2_flat, t, args=(b, c))
sol2_reshaped = sol2.reshape(len(t), 3, 2)

# 验证：第一组批量初始条件的解与单组解一致
print(np.allclose(sol2_reshaped[:,0,:], sol_1))  # 输出True

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