如何基于Scipy最小二乘法为Pandas生成新列并提速？

我有如下结构的Pandas DataFrame：

Race_ID   Date           Student_ID      feature1  
1         1/1/2023       3               0.02167131     
1         1/1/2023       4               0.17349148     
1         1/1/2023       6               0.08438952     
1         1/1/2023       8               0.04143787     
1         1/1/2023       9               0.02589056
1         1/1/2023       1               0.03866752     
1         1/1/2023       10              0.0461553     
1         1/1/2023       45              0.09212758
1         1/1/2023       23              0.10879326      
1         1/1/2023       102             0.186921      
1         1/1/2023       75              0.02990676      
1         1/1/2023       27              0.02731904      
1         1/1/2023       15              0.06020158      
1         1/1/2023       29              0.06302721                          
3         17/4/2022      5               0.2     
3         17/4/2022      2               0.1     
3         17/4/2022      3               0.55     
3         17/4/2022      4               0.15     

希望通过以下方法生成新列：

1. 定义含积分的函数：

import numpy as np
from scipy import integrate
from scipy.stats import norm
import scipy.integrate as integrate
from scipy.optimize import fsolve
from scipy.optimize import least_squares

def integrandforpi_i(xi, ti, *theta):
  prod = 1
  for t in theta:
    prod = prod * (1 - norm.cdf(xi - t))
  return prod * norm.pdf(xi - ti)

def pi_i(ti, *theta):
  return integrate.quad(integrandforpi_i, -np.inf, np.inf, args=(ti, *theta))[0]

2. 针对每个Race_ID，通过scipy.optimize中的least_squares求解非线性方程组，得到每个Student_ID对应新列的值。例如Race_ID为1时，需求解14个非线性方程组，参数范围限制在-2到2之间；Race_ID为3时，需求解4个非线性方程组，最终得到含new_column的目标DataFrame。

目前不清楚如何生成该新列，且实际数据包含大量Race_ID，想请教实现方法及计算提速方案。

1. 明确求解逻辑 #

对每个Race_ID分组，假设该组有n个学生，对应feature1的值为f_1, f_2, ..., f_n，我们需要求解参数数组theta = [θ₁, θ₂, ..., θₙ]，使得每个学生满足pi_i(θ_i, *theta) = f_i。

2. 构建残差函数 #

为适配least_squares，需定义残差函数：输入参数数组theta，输出每个方程的残差（即pi_i(θ_i, *theta) - f_i的数组）：

def residual(theta, f_vals):
    res = []
    n = len(f_vals)
    for i in range(n):
        ti = theta[i]
        res.append(pi_i(ti, *theta) - f_vals[i])
    return np.array(res)

3. 分组求解并合并结果 #

利用Pandas的groupby按Race_ID分组，对每个分组单独求解后合并：

import pandas as pd

# 替换为你的实际DataFrame
df = pd.read_csv("your_data.csv")

def process_group(group):
    f_vals = group['feature1'].values
    n = len(f_vals)
    # 初始化参数：将feature1映射到[-2,2]区间作为初始值
    x0 = np.interp(f_vals, (f_vals.min(), f_vals.max()), (-2, 2))
    # 设置参数边界
    bounds = ([-2]*n, [2]*n)
    # 调用least_squares求解
    result = least_squares(residual, x0, args=(f_vals,), bounds=bounds)
    # 将求解得到的theta赋值给new_column
    group['new_column'] = result.x
    return group

# 分组处理并生成最终DataFrame
final_df = df.groupby('Race_ID').apply(process_group).reset_index(drop=True)

1. 优化积分计算 #

• 替换高精度积分为数值求和：integrate.quad精度高但速度慢，可预定义覆盖正态分布主要概率区间的点（如np.linspace(-5,5,1000)），用数值求和替代积分：

  def pi_i_vec(ti, theta):
      xi = np.linspace(-5, 5, 1000)
      dxi = xi[1] - xi[0]
      # 向量化计算(1 - norm.cdf(xi - t))的乘积
      cdf_terms = 1 - norm.cdf(xi[:, None] - theta)
      prod = np.prod(cdf_terms, axis=1)
      # 计算正态密度项
      pdf_term = norm.pdf(xi - ti)
      # 数值积分求和
      return np.sum(prod * pdf_term) * dxi
  

  之后更新残差函数使用这个向量化版本，能大幅减少计算时间。

• 减少积分点数量：如果业务允许降低精度，可将积分点从1000减少到500甚至300，进一步提速。

2. 并行处理分组 #

每个Race_ID的求解完全独立，可通过多进程并行处理：

from multiprocessing import Pool

# 拆分所有分组
groups = [group for _, group in df.groupby('Race_ID')]
# 根据CPU核心数设置进程数
with Pool(processes=4) as pool:
    processed_groups = pool.map(process_group, groups)
# 合并结果
final_df = pd.concat(processed_groups).reset_index(drop=True)

也可以使用swifter库自动实现apply的并行化，代码更简洁。

3. 优化参数初始化 #

好的初始值能大幅减少least_squares的迭代次数：

• 用feature1的归一化值作为初始值（如示例中的np.interp方法），避免随机初始化带来的无效迭代。
• 对于特征分布相似的Race_ID分组，可复用已求解的参数作为初始值。

4. 放宽收敛条件 #

调整least_squares的收敛参数，放宽精度要求以减少迭代次数：

result = least_squares(residual, x0, args=(f_vals,), bounds=bounds,
                       ftol=1e-4, gtol=1e-4, xtol=1e-4)

默认的精度参数是1e-8，根据业务需求调整到合适的量级即可。

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