如何让scipy的approx_fprime返回数组而非标量以计算偏导数矩阵？

问题 #

我正在尝试让approx_fprime返回多值（数组）而非仅标量，场景为方程组的误差传播：

from scipy.optimize import fsolve, approx_fprime
import numpy as np

def pipeline2(inp):
    sol1,sol1A,sol1B,sol1AB=inp
    def equations(initial):
            x,y,k1,k=initial
            sol1_equation=((k*k1)/(1+k1))-sol1
            sol1A_equation=((k*k1*x**2)/(1+k1*x))-sol1A
            sol1B_equation=((k*k1*y**2)/(1+k1*y))-sol1B
            sol1AB_equation=((k*k1*x**2*y**2)/(1+k1*x*y))-sol1AB
            return(sol1_equation,sol1A_equation,sol1B_equation,sol1AB_equation)
    val1,val2,val3,val4=fsolve(equations,(7,30,0.02,500))
    return np.array([val1,val2,val3,val4])

我手动实现了偏导数矩阵的计算：

h=1e-10
differential=np.array([(pipeline2(sol1_value+h,sol1A_value,sol1B_value,sol1AB_value)-pipeline2(sol1_value,sol1A_value,sol1B_value,sol1AB_value))/h,
                       (pipeline2(sol1_value,sol1A_value+h,sol1B_value,sol1AB_value)-pipeline2(sol1_value,sol1A_value,sol1B_value,sol1AB_value))/h,
                       (pipeline2(sol1_value,sol1A_value,sol1B_value+h,sol1AB_value)-pipeline2(sol1_value,sol1A_value,sol1B_value,sol1AB_value))/h,
                       (pipeline2(sol1_value,sol1A_value,sol1B_value,sol1AB_value+h)-pipeline2(sol1_value,sol1A_value,sol1B_value,sol1AB_value))/h])

该代码能生成4×4的偏导数矩阵，但代码冗长。我希望用approx_fprime替代，但它仅返回标量。请问如何用approx_fprime优雅地计算val1、val2、val3、val4对sol1、sol1A、sol1B、sol1AB的偏导数？

解决方案 #

approx_fprime仅支持标量输出的函数，可通过对pipeline2的每个输出分量单独求导，再组合成偏导数矩阵，以下是两种简洁实现方式：

方法一：循环遍历输出分量 #

对pipeline2返回数组的每个元素，定义仅返回该元素的函数，用approx_fprime计算该分量对输入参数的偏导数，最终堆叠为矩阵：

import numpy as np
from scipy.optimize import approx_fprime

# 输入参数基准值
sol_base = np.array([sol1_value, sol1A_value, sol1B_value, sol1AB_value])
h = 1e-10

# 初始化4×4雅可比矩阵
jacobian = np.zeros((4, 4))

for i in range(4):
    def partial_func(inp):
        return pipeline2(inp)[i]
    jacobian[i] = approx_fprime(sol_base, partial_func, h)

方法二：向量化批量处理 #

利用np.apply_along_axis批量处理所有输出分量，避免显式循环：

import numpy as np
from scipy.optimize import approx_fprime

sol_base = np.array([sol1_value, sol1A_value, sol1B_value, sol1AB_value])
h = 1e-10

jacobian = np.apply_along_axis(
    lambda i: approx_fprime(sol_base, lambda x: pipeline2(x)[i], h),
    axis=0,
    arr=np.arange(4)
)

关键说明 #

• 两种方法生成的4×4雅可比矩阵中，jacobian[i][j]对应第i个输出（val1/val2/val3/val4）对第j个输入（sol1/sol1A/sol1B/sol1AB）的偏导数，与手动实现结果一致。
• approx_fprime默认步长为1e-8，需显式传入h=1e-10以匹配手动实现的精度。

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