Python 中的倾向得分匹配（Propensity Score Matching）：在大规模数据集上的实现优化
各位听众，大家好！今天我们来深入探讨一个在因果推断领域非常重要的技术——倾向得分匹配（Propensity Score Matching, PSM）。特别是在处理大规模数据集时，如何高效地实现 PSM，以及如何在实践中应对各种挑战。
1. 什么是倾向得分匹配？
在观察性研究中，我们经常遇到组间存在显著差异的情况。例如，研究某种药物对疾病的疗效，接受药物治疗的患者可能本身就比未接受治疗的患者病情更严重，或者有其他的健康问题。这种情况下，简单地比较两组的治疗结果可能会产生偏差，因为我们无法确定观察到的差异是由药物引起的，还是由两组患者本身的差异引起的。
倾向得分匹配就是一种用于减少这种选择偏差的技术。它的核心思想是：尝试创建一个“伪随机”的实验环境，使得接受治疗的组和未接受治疗的组在可观测的协变量上尽可能相似。
倾向得分是指个体接受治疗的概率，它是基于可观测的协变量计算出来的。具体来说，我们可以使用逻辑回归等模型，将个体的协变量作为输入，预测其接受治疗的概率。然后，我们可以使用这个概率来进行匹配，找到接受治疗组和未接受治疗组中倾向得分相似的个体，将它们配对。配对后，我们就可以比较配对组的治疗结果，从而更准确地评估治疗效果。
2. PSM 的基本步骤
PSM 的基本步骤可以概括为以下几点：
- 估计倾向得分： 使用逻辑回归或其他分类模型，将协变量作为输入，预测个体接受治疗的概率。
- 匹配： 根据倾向得分，将接受治疗组和未接受治疗组的个体进行匹配。常见的匹配方法包括最近邻匹配、卡尺匹配、半径匹配和核匹配等。
- 平衡性检验： 检查匹配后，接受治疗组和未接受治疗组在协变量上是否达到平衡。常用的平衡性检验包括 t 检验、卡方检验和标准化均值差异等。
- 估计治疗效果： 比较匹配组的治疗结果，估计治疗效果。
3. Python 实现 PSM
接下来，我们来看一下如何使用 Python 来实现 PSM。我们将使用 scikit-learn 来估计倾向得分，使用 MatchIt (或者自己编写函数) 进行匹配，使用 statsmodels 进行平衡性检验和治疗效果估计。
首先，我们需要安装必要的库：
pip install scikit-learn statsmodels pandas numpy接下来，我们创建一个模拟数据集：
import pandas as pd
import numpy as np
from sklearn.linear_model import LogisticRegression
from sklearn.model_selection import train_test_split
from sklearn.metrics import roc_auc_score
# 设置随机种子，保证结果可重复
np.random.seed(42)
# 生成模拟数据
n_samples = 1000
age = np.random.randint(20, 80, n_samples)
income = np.random.normal(50000, 20000, n_samples)
education = np.random.randint(8, 18, n_samples)
comorbidity = np.random.randint(0, 2, n_samples) # 0: 无合并症, 1: 有合并症
# 计算倾向得分（真实的）
true_propensity = 1 / (1 + np.exp(-(0.05 * age + 0.00001 * income + 0.1 * education - 0.5 * comorbidity - 2)))
# 模拟治疗分配（基于倾向得分）
treatment = np.random.binomial(1, true_propensity, n_samples)
# 模拟结果变量（受治疗和协变量影响）
outcome = 2 * treatment + 0.01 * age - 0.000005 * income + 0.05 * education - 0.2 * comorbidity + np.random.normal(0, 1, n_samples)
# 创建 DataFrame
data = pd.DataFrame({'age': age, 'income': income, 'education': education, 'comorbidity': comorbidity, 'treatment': treatment, 'outcome': outcome})
print(data.head())这段代码生成了一个包含 age、income、education、comorbidity（合并症）、treatment（治疗）和 outcome（结果）的 DataFrame。treatment 是根据一个真实的倾向得分生成的，而 outcome 受到治疗和协变量的影响。
接下来，我们使用逻辑回归来估计倾向得分：
# 准备数据
X = data[['age', 'income', 'education', 'comorbidity']]
y = data['treatment']
# 分割训练集和测试集
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)
# 训练逻辑回归模型
model = LogisticRegression(random_state=42)
model.fit(X_train, y_train)
# 预测倾向得分
propensity_scores = model.predict_proba(X)[:, 1]
data['propensity_score'] = propensity_scores
# 评估模型
auc = roc_auc_score(y, propensity_scores)
print(f"AUC: {auc}")这段代码使用 scikit-learn 中的 LogisticRegression 模型来估计倾向得分。它首先将数据分割成训练集和测试集，然后训练模型，并预测倾向得分。最后，它使用 AUC 来评估模型的性能。
接下来，我们进行匹配。这里我们使用最近邻匹配，并设置卡尺值：
def nearest_neighbor_matching(treatment, propensity_scores, caliper=0.25):
"""
进行倾向得分最近邻匹配.
Args:
treatment: Series, 指示是否接受治疗 (1) 或未接受治疗 (0).
propensity_scores: Series, 每个个体的倾向得分.
caliper: float, 卡尺值，用于限制匹配的倾向得分差异.
Returns:
DataFrame, 包含匹配对的DataFrame，包括治疗组索引、对照组索引和倾向得分差异.
"""
treated_indices = treatment[treatment == 1].index
control_indices = treatment[treatment == 0].index
matched_pairs = []
for treated_index in treated_indices:
treated_score = propensity_scores[treated_index]
# 计算治疗组个体与所有对照组个体的倾向得分差异
score_differences = abs(propensity_scores[control_indices] - treated_score)
# 在卡尺范围内找到最小的倾向得分差异
within_caliper = score_differences[score_differences <= caliper]
if not within_caliper.empty:
matched_control_index = score_differences[score_differences == within_caliper.min()].index[0]
matched_pairs.append({'treated': treated_index,
'control': matched_control_index,
'propensity_score_difference': score_differences[matched_control_index]})
return pd.DataFrame(matched_pairs)
# 应用最近邻匹配
matched_pairs = nearest_neighbor_matching(data['treatment'], data['propensity_score'], caliper=0.25)
print(matched_pairs.head())这个代码定义了一个 nearest_neighbor_matching 函数，它实现了最近邻匹配。它遍历接受治疗组的每个个体，找到未接受治疗组中倾向得分最接近的个体，并将它们配对。卡尺值用于限制匹配的倾向得分差异。
接下来，我们进行平衡性检验。这里我们使用标准化均值差异：
def standardized_mean_difference(group1, group2):
"""
计算两组的标准化均值差异.
Args:
group1: Series, 第一组的数值.
group2: Series, 第二组的数值.
Returns:
float, 标准化均值差异.
"""
mean1 = group1.mean()
mean2 = group2.mean()
std1 = group1.std()
std2 = group2.std()
pooled_std = np.sqrt((std1**2 + std2**2) / 2)
return (mean1 - mean2) / pooled_std
# 创建匹配后的数据集
matched_treated = data.loc[matched_pairs['treated']]
matched_control = data.loc[matched_pairs['control']]
# 计算匹配前的标准化均值差异
for col in ['age', 'income', 'education', 'comorbidity']:
smd_before = standardized_mean_difference(data[data['treatment'] == 1][col], data[data['treatment'] == 0][col])
smd_after = standardized_mean_difference(matched_treated[col], matched_control[col])
print(f"变量 {col}: 匹配前 SMD = {smd_before:.3f}, 匹配后 SMD = {smd_after:.3f}")
# 计算匹配后的样本量
print(f"匹配后的样本量：{len(matched_treated)}")这段代码定义了一个 standardized_mean_difference 函数，它计算两组的标准化均值差异。然后，它计算匹配前和匹配后的标准化均值差异，并输出结果。通常，我们希望匹配后的标准化均值差异小于 0.1，表明协变量在两组之间达到了平衡。
最后，我们估计治疗效果：
# 估计平均治疗效果 (ATE)
ate = matched_treated['outcome'].mean() - matched_control['outcome'].mean()
print(f"平均治疗效果 (ATE): {ate:.3f}")这段代码计算匹配组的治疗结果的平均差异，作为治疗效果的估计。
4. 大规模数据集的优化
当处理大规模数据集时，上述代码可能会变得非常慢。以下是一些优化方法：
- 使用高效的匹配算法： 最近邻匹配的复杂度较高，对于大规模数据集，可以考虑使用基于树的匹配算法，如
KDTree或BallTree，或者使用近似最近邻搜索算法。 - 使用并行计算： 可以使用
multiprocessing或joblib等库，将匹配过程并行化，从而提高匹配速度。 - 减少数据量： 如果数据集非常大，可以考虑使用抽样的方法，减少数据量，从而降低计算复杂度。但是，需要注意抽样可能会引入偏差。
- 使用更高效的数据结构： Pandas DataFrame 在某些情况下可能不是最有效的数据结构。可以考虑使用 NumPy 数组或 Dask DataFrames。
让我们尝试使用 KDTree 来加速最近邻匹配：
from sklearn.neighbors import KDTree
import pandas as pd
import numpy as np
def kdtree_matching(treatment, propensity_scores, caliper=0.25):
"""
使用 KDTree 进行倾向得分最近邻匹配.
Args:
treatment: Series, 指示是否接受治疗 (1) 或未接受治疗 (0).
propensity_scores: Series, 每个个体的倾向得分.
caliper: float, 卡尺值，用于限制匹配的倾向得分差异.
Returns:
DataFrame, 包含匹配对的DataFrame，包括治疗组索引、对照组索引和倾向得分差异.
"""
treated_indices = treatment[treatment == 1].index
control_indices = treatment[treatment == 0].index
# 将倾向得分转换为 NumPy 数组并重塑
treated_scores = propensity_scores[treated_indices].values.reshape(-1, 1)
control_scores = propensity_scores[control_indices].values.reshape(-1, 1)
# 构建 KDTree
kdtree = KDTree(control_scores)
matched_pairs = []
for i, treated_index in enumerate(treated_indices):
treated_score = treated_scores[i]
# 查询最近邻
dist, index = kdtree.query([treated_score], k=1)
nearest_neighbor_index = control_indices[index[0][0]]
nearest_neighbor_score = control_scores[index[0][0]]
# 检查是否在卡尺范围内
score_difference = abs(treated_score[0] - nearest_neighbor_score[0])
if score_difference <= caliper:
matched_pairs.append({'treated': treated_index,
'control': nearest_neighbor_index,
'propensity_score_difference': score_difference})
return pd.DataFrame(matched_pairs)
# 应用 KDTree 匹配
matched_pairs_kdtree = kdtree_matching(data['treatment'], data['propensity_score'], caliper=0.25)
print(matched_pairs_kdtree.head())
这段代码使用 scikit-learn 中的 KDTree 类来加速最近邻搜索。KDTree 是一种基于树的数据结构，可以高效地进行最近邻搜索。
5. 使用MatchIt库进行PSM
虽然我们可以自己编写PSM的代码，但是已经有很多现成的库可以帮助我们更方便地实现PSM。其中一个非常流行的库是MatchIt。
首先，安装MatchIt:
pip install matchit然后，使用MatchIt进行PSM：
import matchit
import pandas as pd
# 使用 MatchIt 进行匹配
m = matchit.MatchIt(formula="treatment ~ age + income + education + comorbidity",
data=data,
method="nearest", # 使用最近邻匹配
distance="logit", # 使用倾向得分的logit值
caliper=0.25) # 设置卡尺值
m.fit()
matched_data = m.get_matches()
print(matched_data.head())
# 计算匹配后的标准化均值差异
for col in ['age', 'income', 'education', 'comorbidity']:
smd_before = standardized_mean_difference(data[data['treatment'] == 1][col], data[data['treatment'] == 0][col])
smd_after = standardized_mean_difference(matched_data[matched_data['treatment'] == 1][col], matched_data[matched_data['treatment'] == 0][col])
print(f"变量 {col}: 匹配前 SMD = {smd_before:.3f}, 匹配后 SMD = {smd_after:.3f}")
# 估计平均治疗效果 (ATE)
treated_matched = matched_data[matched_data['treatment'] == 1]
control_matched = matched_data[matched_data['treatment'] == 0]
ate = treated_matched['outcome'].mean() - control_matched['outcome'].mean()
print(f"平均治疗效果 (ATE): {ate:.3f}")MatchIt 库提供了更简洁的接口，可以方便地进行 PSM，并支持多种匹配方法和距离度量。
6. 其他注意事项
- 选择合适的协变量： 倾向得分模型应该包含所有可能影响治疗分配和结果变量的协变量。遗漏重要的协变量可能会导致偏差。
- 处理缺失值： 在估计倾向得分之前，需要处理缺失值。常用的方法包括删除包含缺失值的行、使用均值或中位数填充缺失值，以及使用插补模型预测缺失值。
- 评估匹配质量： 除了标准化均值差异，还可以使用其他指标来评估匹配质量，如 QQ 图和 KS 检验。
- 进行敏感性分析： 倾向得分匹配只能减少可观测的偏差。为了评估不可观测的偏差的影响，需要进行敏感性分析。
7. 总结
今天，我们深入探讨了倾向得分匹配的原理、步骤和 Python 实现。我们还讨论了如何优化 PSM 在大规模数据集上的性能，以及在实践中需要注意的一些问题。PSM 是一种强大的因果推断工具，可以帮助我们更准确地评估治疗效果。掌握 PSM 技术对于研究者和数据科学家来说非常重要。
8. 匹配算法的选择与平衡性检验
选择合适的匹配算法，并进行充分的平衡性检验，是确保倾向得分匹配结果可靠性的关键步骤。
9. 大规模数据处理的优化方向
对于大规模数据集，优化匹配算法和利用并行计算，是提高倾向得分匹配效率的主要方向。
10. PSM实践中的关键环节
在PSM实践中，正确选择协变量、处理缺失值和评估匹配质量，能够有效降低偏差，提高研究结果的可靠性。
更多IT精英技术系列讲座，到智猿学院