定义向量函数及使用solve_ivp时遇TypeError错误求排查

问题背景 #

尝试定义目标向量函数并使用solve_ivp进行积分时，触发错误：

TypeError: fun() takes 3 positional arguments but 4 were given

原代码如下：

def fun(s, rho_0, rho_1):
    return lambda t, s:np.dot(np.array([0.775416, 0,0, 0.308968]).reshape(2,2), s) + np.array([rho_0,rho_1]).reshape(2,1)

def fun2(t, rho_0, rho_1):
    res = solve_ivp(fun, [0, 5], y0 = [0, 0], t_eval=np.arange(0,5), args = (rho_0, rho_1), vectorized = True)
    return res.y[1]

fun2(t = 0, rho_0 = 0.0099532, rho_1 = 0.001699)

错误原因 #

1. 参数匹配错误：solve_ivp会自动向传入的微分方程函数传递3类参数：当前时间t、当前状态向量y、args中的额外参数。但你定义的fun只接受s, rho_0, rho_1三个参数，实际调用时solve_ivp会传入t, s, rho_0, rho_1四个参数，直接导致参数不匹配。
2. 冗余嵌套设计：fun返回lambda的写法完全没必要，solve_ivp只需要直接接收符合参数要求的微分方程函数。

修正代码 #

1. 修正微分方程函数 #

按照solve_ivp的要求，函数需以t, y作为前两个参数，后续为额外参数：

import numpy as np
from scipy.integrate import solve_ivp

def fun(t, s, rho_0, rho_1):
    # 明确定义系数矩阵，提升可读性
    A = np.array([[0.775416, 0],
                  [0, 0.308968]])
    # 构造常数项向量
    b = np.array([rho_0, rho_1]).reshape(2, 1)
    # 确保状态向量维度匹配后计算
    return np.dot(A, s.reshape(2, 1)) + b

2. 修正积分调用逻辑 #

移除无意义的t参数，调整solve_ivp调用细节：

def fun2(rho_0, rho_1):
    res = solve_ivp(fun, [0, 5], y0=[0, 0], 
                    t_eval=np.arange(0, 5), 
                    args=(rho_0, rho_1), 
                    vectorized=True)
    # 返回第二个状态变量的积分结果
    return res.y[1]

# 测试调用
fun2(rho_0=0.0099532, rho_1=0.001699)

额外注意点 #

• 维度兼容：将状态向量s转为(2,1)列向量后再做矩阵乘法，避免维度不匹配报错；
• vectorized参数：启用vectorized=True时，函数需支持批量输入，当前代码通过reshape已兼容；
• t_eval取值：np.arange(0,5)生成的是[0,1,2,3,4]，如果需要包含5，改为np.arange(0, 5.1, 1)。

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