如何用更Pythonic的方式优化基于fsolve的循环求解代码？

问题背景 #

给定如下DataFrame：

需要逐次使用fsolve迭代求解变量x，原实现代码如下：

result = []
for i in range(len(df) - 1):
  def R(x):
    ZDownturn = norm.ppf(df['PD Downturn'])[i] #Exclude 'Default' class
    ZShift = np.sqrt(
        x / (1 - x)
    ) * norm.ppf(0.999)
    ZPortion = np.sqrt(
        1 / (1 - x)
    ) * norm.ppf(df['PD TTC'])[i] #Exclude 'Default' class
    target = ZShift + ZPortion
    error =  np.abs(ZDownturn - target)
    return error

  # Initial guess
  x0 = [0.01]
  
  # Solver
  Rho = fsolve(R, x0)
  result.append(Rho[0])

原代码可得到结果：

[0.19153452995548875,
 0.15906256238706026,
 0.08759684851688349,
 0.1348702069117432]

现希望优化这段代码，写出更符合Python风格的实现方式。

优化方案 #

核心优化思路 #

1. 避免在循环内重复定义函数，改用通用函数+参数绑定的方式
2. 预计算常量与DataFrame的统计转换值，减少重复计算
3. 用更简洁的列表推导式或pandas.apply替代显式循环

优化后代码示例 #

import numpy as np
from scipy.optimize import fsolve
from scipy.stats import norm
import pandas as pd
from functools import partial

# 预计算固定常量
quantile_999 = norm.ppf(0.999)

# 定义通用的误差计算函数
def calculate_error(x, z_downturn, z_ttc):
    sqrt_x_over_1mx = np.sqrt(x / (1 - x))
    sqrt_1_over_1mx = np.sqrt(1 / (1 - x))
    target = sqrt_x_over_1mx * quantile_999 + sqrt_1_over_1mx * z_ttc
    return np.abs(z_downturn - target)

# 提前计算DataFrame列的ppf转换值
df['Z_Downturn'] = norm.ppf(df['PD Downturn'])
df['Z_TTC'] = norm.ppf(df['PD TTC'])

# 用列表推导式结合partial完成求解
result = [
    fsolve(partial(calculate_error, z_downturn=row['Z_Downturn'], z_ttc=row['Z_TTC']), [0.01])[0]
    for _, row in df.iloc[:-1].iterrows()  # 排除最后一行Default类
]

print(result)

优化点说明 #

• 预计算提升效率：把重复调用的norm.ppf(0.999)和DataFrame列的ppf值提前计算，避免循环内重复执行统计计算
• 函数复用性增强：将原本循环内的R函数改为通用的calculate_error，通过functools.partial绑定每行的专属参数，避免循环内重复定义函数
• 代码简洁性提升：用列表推导式遍历DataFrame行，代码紧凑易读，符合Pythonic风格
• 可读性优化：调整变量命名，让逻辑更清晰

可选实现：使用pandas.apply #

如果更习惯pandas的操作风格，也可以用apply方法实现：

def solve_for_row(row):
    return fsolve(partial(calculate_error, z_downturn=row['Z_Downturn'], z_ttc=row['Z_TTC']), [0.01])[0]

result = df.iloc[:-1].apply(solve_for_row, axis=1).tolist()

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