因果推断——PSM挑选对照组

问题描述

  • 共 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
    74
    import 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为良好）