问题描述
- 共 100 个实验商家，需要从 2000+ 个商家中使用PSM匹配找到相似商家作为 AB 实验
PSM 基本流程
- 共三个集合：实验组、对照组（初始化为空）、候选商家集合
- 特征构造：选择可能影响实验结论的关键特征
- label 构建：标记 100 个实验商家 label 为 1，其他商家 label 为 0
- 模型训练：用构建的样本训练一个分类模型，一般为 LR 模型即可
- PSM 匹配：迭代访问实验商家 A，分别进行如下操作
- 在候选商家集合中选择模型预估值和 A 差异小于一定阈值的商家作为匹配候选集，从匹配候选集中随机选择一个商家 B
- 将商家 B 加入对照组，并在候选商家集合中删除对照组
流程思考
- 如果找不到指定阈值的商家作为匹配候选集
Python 实现 Demo
PSM 匹配 Demo
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74import pandas as pd
import numpy as np
from sklearn.linear_model import LogisticRegression
from sklearn.preprocessing import StandardScaler
# 生成模拟数据（假设有2000个商家）
np.random.seed(42)
n = 2000
# 生成商家特征
data = pd.DataFrame({
"shop_id": range(n),
"monthly_sales": np.random.normal(50000, 15000, n), # 月销售额
"store_size": np.random.choice([50, 100, 150, 200], n), # 店铺面积
"is_chain": np.random.binomial(1, 0.3, n), # 是否连锁店
"city_tier": np.random.choice([1, 2, 3], p=[0.2,0.5,0.3], n), # 城市等级
"category": np.random.choice(["餐饮", "零售", "服务", "其他"], p=[0.4,0.3,0.2,0.1], n) # 类别
})
# 添加实验组标记（随机选择100家作为实验组）
experiment_group = np.random.choice(n, 100, replace=False)
data['is_treated'] = data.shop_id.isin(experiment_group).astype(int)
# 特征预处理
features = data[['monthly_sales', 'store_size', 'is_chain', 'city_tier']]
scaler = StandardScaler()
features_scaled = scaler.fit_transform(features)
# 训练倾向得分模型
model = LogisticRegression(max_iter=1000)
model.fit(features_scaled, data['is_treated'])
# 计算倾向得分
data['propensity_score'] = model.predict_proba(features_scaled)[:, 1]
# PSM匹配函数
def psm_match(data, treated_col='is_treated', score_col='propensity_score', caliper=0.02):
treated = data[data[treated_col] == 1]
control = data[data[treated_col] == 0]
matches = []
for _, treat in treated.iterrows():
# 寻找最近邻匹配
control['distance'] = abs(control[score_col] - treat[score_col])
candidates = control[control.distance <= caliper]
if not candidates.empty:
# 随机选择最近邻
match = candidates.sample(1, random_state=42)
matches.append(match.index[0])
control = control.drop(match.index)
return treated, data.loc[matches]
# 执行匹配
treated_group, control_group = psm_match(data)
# 评估匹配质量
print(f"匹配成功实验商家数：{len(treated_group)}")
print(f"找到对照组商家数：{len(control_group)}")
# 检查协变量平衡性
matched_data = pd.concat([treated_group, control_group])
for col in ['monthly_sales', 'store_size', 'is_chain']:
treat_mean = treated_group[col].mean()
control_mean = control_group[col].mean()
std_diff = (treat_mean - control_mean) / treated_group[col].std()
print(f"{col}标准化差异：{std_diff:.3f}")
# 输出匹配结果
matched_pairs = pd.DataFrame({
'experiment_id': treated_group.shop_id.values,
'control_id': control_group.shop_id.values
})代码说明：
- 数据模拟：生成包含销售额、店铺面积、是否连锁等特征的2000个商家数据
- 实验组选择：随机选择100个商家作为实验组（实际业务中应根据业务逻辑选择）
- 特征标准化：对连续变量进行标准化处理
- 倾向得分模型：使用逻辑回归预测进入实验组的概率
- 卡钳匹配：设置最大允许得分差异（caliper=0.02），确保匹配质量
- 平衡性检查：通过标准化差异评估匹配质量（<0.1为良好）