Python中含求和循环的多变量Newton Raphson方法实现问题

问题背景 #

当前需要在Python中用MLE方法估计某分布的参数，推导得到的对数似然函数偏导如下：

由于公式形式复杂，计划采用Newton-Raphson方法数值求解方程根，但编写代码后调用jacobian()计算雅可比矩阵时遇到问题：运行目标函数直接返回数值结果，无法正常生成雅可比。

核心问题说明 #

你不需要把函数改写为符号多项式形式。autograd是自动微分框架，不需要符号表达式就能追踪计算图输出雅可比矩阵，你之前的代码无法正常计算雅可比，是三个写法错误导致的：

• 待估参数p0-p4被定义为全局固定值，没有作为函数入参传入，自动微分没有可求导的变量对象
• 混用了标准库math的函数、未经过autograd包装的scipy.special.psi，这类算子无法被autograd追踪，会打断计算图
• 手动写Python原生循环做累加虽然语法上能跑，但效率低，用numpy向量化实现更简洁，也更适配自动微分的运行逻辑

修正方法 #

• 调整函数结构：将所有待估计的参数打包为一个numpy数组，作为目标函数的唯一入参，不要用全局变量固定参数值。
• 统一使用autograd兼容算子：所有数学运算调用autograd.numpy下的函数，不要用标准库math的同名函数；导入autograd.scipy.special下的psi，不要直接使用scipy原生的psi函数。
• 向量化实现去掉循环：先将样本数据转为numpy数组，所有逐元素计算直接用数组广播完成，最后用np.sum求和即可，不需要手动写循环逐行累加。

修正后可运行代码 #

import autograd.numpy as np
from autograd import jacobian
from autograd.scipy.special import psi

# 样本转为numpy数组
x = np.array([0.1, 0.2, 1, 1, 1, 1, 1, 2, 3, 6, 7, 11, 12, 18, 18, 18, 18, 18,
              21, 32, 36, 40, 45, 46, 47, 50, 55, 60, 63, 63, 67, 67, 67, 67,
              72, 75, 79, 82, 82, 83, 84, 84, 84, 85, 85, 85, 85, 85, 86, 86])

def Fs(p):
    p0, p1, p2, p3, p4 = p
    n = len(x)
    # 向量化计算所有中间变量
    z = np.exp(-(p1 / p2) * (np.exp(p2 * x) - 1))
    a = 1 - z
    b = a ** p0
    c = np.exp(p2 * x)
    
    # 直接求和替代循环累加
    sum1 = np.sum(np.log(a))
    sum2 = np.sum(b * np.log(a) / np.log(1 - b))
    sum3 = np.sum(c)
    sum4 = np.sum((c - 1) * z / a)
    sum5 = np.sum((c - 1) * z * (a ** (p0 - 1)) / (1 - b))
    sum6 = np.sum(x)
    sum7 = np.sum(c * (1 - p2 * x))
    sum8 = np.sum((p2 * x * c - c + 1) * z / a)
    sum9 = np.sum((p2 * x * c - c + 1) * z * (a ** (p0 - 1)) / (1 - b))
    sum10 = np.sum(np.log(1 - b))

    f1 = n / p0 + p3 * sum1 - (p4 - 1) * sum2
    f2 = n / p1 + n / p2 - sum3 / p2 + sum4 * (p0 * p3 - 1) / p2 + sum5 * p0 * (p4 - 1) / p2
    f3 = -n * p1 / (p2 ** 2) + sum6 + p1 * sum7 / (p2 ** 2) - p1 * (p3 * p0 - 1) * sum8 / (p2 ** 2) + p0 * p1 * (p4 - 1) * sum9 / (p2 ** 2)
    f4 = -n * (psi(p3) + psi(p3 + p4)) + p0 * sum1
    f5 = -n * (psi(p4) + psi(p3 + p4)) + p0 * sum10

    return np.array([f1, f2, f3, f4, f5]).reshape(-1, 1)

# 生成雅可比计算函数
Fs_jac = jacobian(Fs)

# 初始化参数测试
p_init = np.array([2.5, 0.00003, 0.1, 0.02, 0.1])
print("初始点函数值：\n", Fs(p_init))
print("初始点雅可比矩阵：\n", Fs_jac(p_init).squeeze())

牛顿迭代实现参考 #

拿到雅可比函数后，按如下逻辑迭代即可，求解线性方程组时用np.linalg.solve比直接求矩阵逆数值稳定性更好：

p = p_init.copy()
max_iter = 100
tol = 1e-6
for _ in range(max_iter):
    f_val = Fs(p).flatten()
    jac_val = Fs_jac(p).squeeze()
    delta = np.linalg.solve(jac_val, -f_val)
    p += delta
    if np.max(np.abs(delta)) < tol:
        break
print("MLE参数估计结果：", p)

补充提示 #

• 不需要调用sympy做符号推导，自动微分的计算效率远高于符号微分，完全满足数值迭代的性能需求
• 如果手动实现牛顿迭代收敛性差，可以直接调用scipy.optimize.root，传入Fs函数、初始值即可求解，支持hybr、lm等多种稳健求解算法，传入你计算的雅可比还能进一步提升收敛速度
• 迭代时注意参数的取值范围（比如形状参数、尺度参数通常要求为正），避免出现对数自变量为负、除零等数值错误，必要时可以加边界约束或者对参数做对数变换后再迭代。

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