Python中合并solve_ivp的OdeSolution并实现向量化调用的问题

问题原因 #

1. 原sol_comb函数的if t <= 10判断在输入为数组时，会生成布尔数组，而Python的if语句无法直接处理布尔数组，导致报错。
2. np.vectorize本质是循环处理每个元素，但OdeSolution实例的返回值是多维数组（比如示例中是(2,)的数组，对应x和y两个变量），vectorize会尝试将这些序列塞进一个一维数组，触发ValueError: setting an array element with a sequence错误。

解决方法 #

方法一：直接修改函数支持数组输入 #

通过numpy的掩码索引处理数组输入，同时保持和原OdeSolution一致的输出格式（标量输入返回(2,)数组，数组输入返回(2, N)数组）：

from scipy.integrate import solve_ivp
import numpy as np

def lotkavolterra(t, z, a, b, c, d):
    x, y = z
    return [a*x - b*x*y, -c*y + d*x*y]

sol_1 = solve_ivp(lotkavolterra, [0, 10], [10, 5], args=(1.5, 1, 3, 1), dense_output=True).sol  
sol_2 = solve_ivp(lotkavolterra, [10, 15], [10, 5], args=(1.5, 1, 3, 1), dense_output=True).sol

def sol_comb(t):
    t = np.asarray(t)
    # 划分时间区间的掩码
    mask_left = t <= 10
    mask_right = t > 10
    
    # 初始化结果数组，匹配OdeSolution的输出形状
    result = np.zeros((2, t.size))
    
    # 分别填充两个区间的结果
    if np.any(mask_left):
        result[:, mask_left] = sol_1(t[mask_left])
    if np.any(mask_right):
        result[:, mask_right] = sol_2(t[mask_right])
    
    # 输入为标量时压缩维度，保持和原sol一致的输出格式
    return result.squeeze() if t.ndim == 0 else result

测试验证：

# 标量输入
print(sol_comb(5))  # 等价于sol_1(5)
print(sol_comb(12)) # 等价于sol_2(12)

# 数组输入
t_array = np.linspace(0, 15, 30)
y_array = sol_comb(t_array)
print(y_array.shape)  # 输出(2, 30)，符合预期

方法二：拼接解并重建插值器（可选） #

如果需要更灵活的插值控制，可以将两个解的时间点和数值拼接后，用scipy.interpolate.interp1d重建插值器：

from scipy.interpolate import interp1d

# 从原OdeSolution中采样足够多的点保证精度
t1 = np.linspace(0, 10, 100)
y1 = sol_1(t1)
t2 = np.linspace(10, 15, 50)
y2 = sol_2(t2)

# 拼接时去除重复的t=10点
t_combined = np.concatenate([t1, t2[1:]])
y_combined = np.concatenate([y1, y2[:, 1:]], axis=1)

# 创建插值器，选择和原solve_ivp匹配的插值类型（比如立方插值）
sol_comb_interp = interp1d(t_combined, y_combined, kind='cubic', fill_value="extrapolate")

这种方式适合需要自定义插值策略的场景，但会丢失原OdeSolution的高阶插值精度，推荐优先使用方法一。

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