ML辅助A/B随机化检验 — 有限样本随机化测试检测异质处理效应与干扰效应

"""
ML辅助A/B随机化检验
基于 arXiv:2501.07722 (2025)
有限样本随机化测试，检测ATE/HTE/干扰效应
"""
import numpy as np
from sklearn.ensemble import GradientBoostingClassifier
from sklearn.model_selection import cross_val_score
import warnings
warnings.filterwarnings('ignore')


def compute_ml_test_statistic(X, D, Y, model=None, n_folds=5):
    """
    计算ML辅助测试统计量
    T_ML = CV_Error(无处理) - CV_Error(有处理)
    
    Args:
        X: 协变量矩阵 (n_samples, n_features)
        D: 处理变量 (n_samples,)，0=控制，1=处理
        Y: 结果变量 (n_samples,)
        model: ML模型（默认GBM）
    Returns:
        T_ML: 测试统计量（正值=处理有效）
    """
    if model is None:
        model = GradientBoostingClassifier(n_estimators=50, max_depth=3, random_state=42)

    n = len(Y)
    D_col = D.reshape(-1, 1)

    # 有处理变量的模型
    X_with_D = np.hstack([X, D_col])
    scores_with = cross_val_score(model, X_with_D, Y, cv=n_folds, scoring='neg_log_loss')

    # 无处理变量的模型
    scores_without = cross_val_score(model, X, Y, cv=n_folds, scoring='neg_log_loss')

    # 测试统计量：有处理时误差更小说明处理有效
    T_ML = np.mean(scores_without) - np.mean(scores_with)
    return T_ML


def randomization_test(X, D, Y, n_permutations=500, alpha=0.05):
    """
    随机化检验（排列检验）
    
    Args:
        X: 协变量
        D: 处理变量
        Y: 结果变量
        n_permutations: 排列次数
        alpha: 显著性水平
    Returns:
        p_value: p值
        T_observed: 观测统计量
        T_null_dist: 零假设分布
        significant: 是否显著
    """
    # 观测统计量