CausalML项目移除pygam依赖的技术方案分析

背景与问题

在CausalML项目中，propensity模块的校准功能当前依赖于pygam库。然而，pygam对scipy版本的限制（要求版本低于1.12）导致了与numpy 2.x和Cython 3.x的兼容性问题，这给项目的依赖管理和升级带来了障碍。

技术评估

经过深入分析，我们发现propensity.calibrate()方法主要实现的是倾向性得分的校准功能。pygam在此处的作用是通过广义加性模型(GAM)来调整预测概率，使其更接近真实观察到的结果分布。

替代方案研究

我们评估了scikit-learn提供的概率校准方法作为替代方案，特别是IsotonicRegression（等张回归）方法。通过对比实验发现：

1. 等张回归在校准性能上表现优异，在某些测试案例中甚至优于pygam
2. 从评估指标来看（如log_loss和brier_score_loss），等张回归能够提供相当或更好的校准效果
3. 等张回归作为scikit-learn的一部分，具有更好的维护性和兼容性

实现方案

新的校准函数实现简洁明了：

from sklearn.isotonic import IsotonicRegression

def calibrate(ps, treatment):
    pm_ir = IsotonicRegression(out_of_bounds="clip")
    ps_ir = pm_ir.fit_transform(ps, treatment)
    return ps_ir

该实现具有以下优势：

• 代码更简洁，依赖更少
• 参数out_of_bounds="clip"确保输出值保持在[0,1]范围内
• 完全兼容现有的输入输出接口

性能对比

实验数据显示，在多个测试场景下：

• 等张回归的校准效果至少不逊于pygam
• 在某些情况下，等张回归能提供更好的校准结果
• 计算效率方面，等张回归通常表现更优

迁移影响评估

这一变更对现有用户的影响极小：

1. 接口保持不变，用户无需修改调用代码
2. 输出结果的质量有保障
3. 解决了依赖冲突问题，使项目能使用更新的科学计算库版本

结论与建议

基于以上分析，建议采用scikit-learn的IsotonicRegression作为pygam的替代方案。这一变更不仅能解决当前的依赖冲突问题，还能带来代码简化、维护性提升等额外好处，同时对现有用户的影响最小化。