如何用SciPy solve_ivp批量求解大规模耦合ODE方程组

问题背景 #

我已经掌握了用Python SciPy的solve_ivp求解2-3个耦合常微分方程（ODE）的方法，现在需要扩展到包含K个n类方程与(N-K)个p类方程的大规模耦合场景（方程数量可达数百个，无法手动逐个定义）。尝试用数组批量实现时，出现如下ValueError报错：

ValueError: setting an array element with a sequence. The requested array has an inhomogeneous shape after 1 dimensions. The detected shape was (2,) + inhomogeneous part.

原小规模求解代码示例 #

def sol_fun():
    def dndt(t,V):
        
        GTotUp = GUp + np.sum(GUpM*V[0:2]) #GUpM是长度为2的向量
        GTotDown = GDown+np.sum(GDownM*(V[0:2]+1)) #GDownM是长度为2的向量
        n1=V[0]
        n2=V[1]
        p=V[2]

        n_dot = -kappa*n1 + N*GDownM[0]*(n1+1)*p - N*GUpM[0] * n1*(1-p)
        n_dot2 = -kappa*n2 + N*GDownM[1]*(n2+1)*p - N*GUpM[1] * n2*(1-p)
        pdot = GTotUp*(1-p)-GTotDown*p

        return [n_dot,n_dot2, pdot]

    sol = solve_ivp(dndt, [t[0], t[-1]], [0,0,0], method='LSODA', t_eval=t)
    return sol

尝试的批量实现代码（报错） #

def sol_fun():
    def dndt(t,V):

        # 以下定义引发错误
        n=V[0:K] # 取V的前K个元素作为n向量
        p=V[K:N] # 取V的第K到N个元素作为p向量
    
        GTotUp = GUp + np.sum(GUpM*V[0:K]) #GUpM是长度为K的向量
        GTotDown = GDown+np.sum(GDownM*(V[0:K]+1)) #GDownM是长度为K的向量
    
        pm=np.sum(p) # pm是标量
    
        ndot=-kappa*n + rho*GDownM*(n + 1)*pm - rho*GUpM*n*(1-pm) 
        pdot = GTotUp*(1-p) - GTotDown*p
    
        return [ndot, pdot]

问题原因 #

报错的核心是：solve_ivp要求导数函数dndt必须返回一维numpy数组，而你返回的[ndot, pdot]是两个数组组成的列表，形状不匹配（ndot是K维数组，pdot是(N-K)维数组，列表无法被识别为统一的一维数组）。

修正方案 #

把ndot和pdot拼接成一个一维numpy数组返回即可，具体步骤：

• 使用np.concatenate([ndot, pdot])或np.hstack([ndot, pdot])将两个数组合并为一维数组
• 确保所有运算都是numpy元素级运算（你的代码中这部分已经正确，比如-kappa*n是数组元素级乘法）

修正后的完整代码 #

import numpy as np
from scipy.integrate import solve_ivp

def sol_fun(K, N, t, GUp, GDown, GUpM, GDownM, kappa, rho):
    def dndt(t, V):
        n = V[:K]  # 前K个元素为n类变量
        p = V[K:N] # 从K到N的元素为p类变量
    
        GTotUp = GUp + np.sum(GUpM * n)
        GTotDown = GDown + np.sum(GDownM * (n + 1))
    
        pm = np.sum(p)
    
        # 计算n类变量的导数（K维数组）
        ndot = -kappa * n + rho * GDownM * (n + 1) * pm - rho * GUpM * n * (1 - pm)
        # 计算p类变量的导数（N-K维数组）
        pdot = GTotUp * (1 - p) - GTotDown * p
    
        # 拼接成一维数组返回，符合solve_ivp要求
        return np.concatenate([ndot, pdot])

    # 初始条件：全0的一维数组，长度为N
    initial_conditions = np.zeros(N)
    sol = solve_ivp(dndt, [t[0], t[-1]], initial_conditions, method='LSODA', t_eval=t)
    return sol

关键说明 #

1. V是solve_ivp传入的一维状态数组，长度为N（K个n类变量 + N-K个p类变量）
2. ndot和pdot分别是K维和(N-K)维的numpy数组，运算均为元素级，保证批量计算正确
3. np.concatenate将两个数组合并为长度为N的一维数组，完美匹配solve_ivp对返回值的要求
4. 初始条件也需要是长度为N的一维数组，不能是嵌套列表

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