使用scipy.newton_krylov优化时残差函数疑似代码错误排查

用户尝试用scipy.newton_krylov求解包含权重矩阵W和两个对偶变量的非线性方程组，但报错提示Hessian非方阵（实际为23890×23892，预期23892×23892）。其中W含23890个元素，加上2个对偶变量，总变量数为23892，但残差函数输出维度与输入维度不匹配。

用户原代码 #

残差函数 #

## This function is what will be used in the newton_krylov optimization process
def residual_function(all_vars):
    """
    Calculates the residual of the weight matrix equation.

    Args:
        all_vars: A combined vector of W elements and lambda_vals.

    Returns:
        A NumPy array representing the element-wise residual of the equation.
    """

    n, m = w_init.shape  # Determine dimensions of W from w_init (this is defined before the residual_function in my code

    W = all_vars[:n*m].reshape(n, m) ## extracts the first n*m elements from all_vars and assigns to W, reshaping to matrix

    lambda_vals = all_vars[n*m:] ## extracts remaining elements and assigns to lambda_vals

    exp_term = np.exp(-d1 - lambda_vals[0] * Owner - lambda_vals[1] * NH_White) - W ## Owner and NH_white are known column vectors of the same length as W

    residual = d1 * exp_term

    return residual

求解代码 #

## set up dual values 
## Combine initial guesses for lambda and gamma
all_lambda = np.array([lambda_init, gamma_init]) ## lambda_init and gamma_init are prespecified as 1 elsewhere in my program

## Combine all necessary inputs, including flattened w_init matrix
all_vars = np.concatenate((w_init.flatten(), all_lambda))

test = residual_function(all_vars = all_vars)

print(test) 
## Solve using newton_krylov with residual function
sol = newton_krylov(residual_function, all_vars)

问题原因 #

newton_krylov用于求解非线性方程组F(x)=0，要求输入变量向量x的维度必须等于残差函数F(x)的输出维度。用户当前的残差仅对应W的23890个元素，未包含两个对偶变量的约束方程，导致输出维度（23890）≠输入维度（23892），Jacobian矩阵为残差维度×变量维度，无法形成方阵，进而报错。

解决方案 #

需要补充对偶变量对应的约束方程，使残差总维度与输入变量维度一致（23892）。具体步骤如下：

1. 将原W对应的残差展平为一维数组（长度23890）；
2. 根据原优化问题的约束（如KKT条件、互补松弛条件等），添加两个对偶变量的残差方程（每个为标量）；
3. 合并所有残差，形成长度为23892的一维向量。

修改后的残差函数示例 #

def residual_function(all_vars):
    n, m = w_init.shape
    total_w_elements = n * m
    W = all_vars[:total_w_elements].reshape(n, m)
    lambda_vals = all_vars[total_w_elements:]  # [λ₀, λ₁]

    # 原W的残差方程，展平为一维
    exp_term = np.exp(-d1 - lambda_vals[0] * Owner - lambda_vals[1] * NH_White) - W
    w_residual = (d1 * exp_term).flatten()

    # --------------------------
    # 关键：补充对偶变量的约束方程
    # 这里需要根据你的原优化问题替换为实际条件，比如KKT中的对偶可行性/互补松弛条件
    # 示例（需替换）：假设原问题要求∑(W * Owner) = target_value，或者梯度为零的条件
    lambda_residual_0 = ...  # 第一个对偶变量的残差，标量
    lambda_residual_1 = ...  # 第二个对偶变量的残差，标量
    # --------------------------

    # 合并所有残差，总长度与输入变量一致
    total_residual = np.concatenate([w_residual, [lambda_residual_0, lambda_residual_1]])

    return total_residual

验证 #

修改后，调用residual_function(all_vars)的输出长度应为23892，与输入all_vars的长度匹配，此时newton_krylov可正常计算Jacobian矩阵并求解，最终从结果sol中提取后两个元素即为λ和γ的最优值。

内容的提问来源于stack exchange，提问作者Kelsey O'Hollaren