3. PSI 计算 Python 实现：分箱、概率分布与对数求和

本节将使用 pandas 和 numpy 实现 PSI 的核心计算函数。我们将遵循“分箱 → 概率分布 → 对数求和”的标准流程，并提供一个可直接复用的函数模板。

3.1 计算流程总览

PSI 的计算可以拆解为三个连续步骤，下面的流程图展示了从原始数据到最终 PSI 值的完整路径：

3.2 核心函数实现

以下代码实现了完整的 PSI 计算函数，包含三种分箱策略和自动平滑处理：

import pandas as pd
import numpy as np

def calculate_psi(actual, expected, bins=10, bin_type='equal_width'):
    """
    计算 PSI（群体稳定性指数）

    参数:
        actual (array-like): 实际分布数据
        expected (array-like): 预期分布数据
        bins (int): 分箱数量，默认10
        bin_type (str): 分箱方式，'equal_width' | 'equal_freq' | 'quantile'

    返回:
        float: PSI 值
    """
    # 步骤1：分箱
    if bin_type == 'equal_width':
        # 等宽分箱：基于合并数据的分位数
        combined = np.concatenate([actual, expected])
        bin_edges = np.percentile(combined, np.linspace(0, 100, bins + 1))
        # 去重边界
        bin_edges = np.unique(bin_edges)
    elif bin_type == 'equal_freq':
        # 等频分箱：基于实际数据的分位数
        bin_edges = np.percentile(actual, np.linspace(0, 100, bins + 1))
        bin_edges = np.unique(bin_edges)
    elif bin_type == 'quantile':
        # 分位数分箱：使用实际数据的十分位
        bin_edges = np.percentile(actual, np.linspace(0, 100, bins + 1))
        bin_edges = np.unique(bin_edges)
    else:
        raise ValueError("bin_type 必须为 'equal_width', 'equal_freq' 或 'quantile'")

    # 步骤2：计算概率分布
    actual_counts, _ = np.histogram(actual, bins=bin_edges)
    expected_counts, _ = np.histogram(expected, bins=bin_edges)

    # 转换为占比
    actual_pct = actual_counts / len(actual)
    expected_pct = expected_counts / len(expected)

    # 平滑处理：防止除零或取对数时出现无穷大
    epsilon = 1e-6
    actual_pct = np.clip(actual_pct, epsilon, 1 - epsilon)
    expected_pct = np.clip(expected_pct, epsilon, 1 - epsilon)

    # 步骤3：对数求和
    psi_value = np.sum((actual_pct - expected_pct) * np.log(actual_pct / expected_pct))

    return psi_value

3.3 函数使用示例

下面通过一个模拟场景演示函数的使用方法：

# 生成模拟数据
np.random.seed(42)
# 预期分布：正态分布
expected_data = np.random.normal(loc=0, scale=1, size=10000)
# 实际分布1：与预期相同（稳定）
actual_stable = np.random.normal(loc=0, scale=1, size=10000)
# 实际分布2：均值偏移（不稳定）
actual_shifted = np.random.normal(loc=0.5, scale=1.2, size=10000)

# 计算 PSI
psi_stable = calculate_psi(actual_stable, expected_data, bins=10)
psi_shifted = calculate_psi(actual_shifted, expected_data, bins=10)

print(f"稳定场景 PSI: {psi_stable:.4f}")
print(f"偏移场景 PSI: {psi_shifted:.4f}")

3.4 分箱策略对比

不同的分箱策略对 PSI 计算结果有显著影响。下表总结了三种策略的适用场景：

3.5 平滑处理的重要性

在实际计算中，某些分箱可能出现零计数，导致对数计算失败。我们的函数通过 np.clip 将概率限制在 [1e-6, 1-1e-6] 范围内，确保计算稳定性。

3.6 批量计算与监控

在实际监控场景中，通常需要计算多个时间窗口的 PSI。以下代码展示了如何批量计算：

def batch_psi_calculation(reference_data, monitoring_data_list, bins=10):
    """
    批量计算多个监控窗口的 PSI

    参数:
        reference_data: 基准分布（如模型训练集）
        monitoring_data_list: 列表，每个元素为一个时间窗口的数据
        bins: 分箱数量

    返回:
        pd.DataFrame: 包含时间窗口和对应 PSI 值的表格
    """
    results = []
    for i, window_data in enumerate(monitoring_data_list):
        psi = calculate_psi(window_data, reference_data, bins=bins)
        results.append({
            'window_id': f'window_{i+1}',
            'psi_value': round(psi, 4),
            'sample_size': len(window_data)
        })
    return pd.DataFrame(results)

# 使用示例
reference = np.random.normal(0, 1, 50000)
windows = [np.random.normal(0.1 * i, 1, 10000) for i in range(5)]
psi_df = batch_psi_calculation(reference, windows, bins=10)
print(psi_df)

3.7 代码验证与调试

为了确保函数正确性，可以使用以下单元测试进行验证：

def test_psi_identical_distributions():
    """测试相同分布的 PSI 应接近 0"""
    data = np.random.normal(0, 1, 10000)
    psi = calculate_psi(data, data, bins=10)
    assert abs(psi) < 0.01, f"相同分布 PSI 应接近 0，实际为 {psi}"

def test_psi_different_distributions():
    """测试不同分布的 PSI 应显著大于 0"""
    data1 = np.random.normal(0, 1, 10000)
    data2 = np.random.normal(2, 1, 10000)
    psi = calculate_psi(data1, data2, bins=10)
    assert psi > 0.5, f"不同分布 PSI 应较大，实际为 {psi}"

def test_psi_smoothing():
    """测试平滑处理不会导致 NaN"""
    data1 = np.array([1, 2, 3, 4, 5])
    data2 = np.array([10, 20, 30, 40, 50])
    psi = calculate_psi(data1, data2, bins=3)
    assert not np.isnan(psi), "平滑处理后不应出现 NaN"
    assert np.isfinite(psi), "PSI 应为有限值"

# 运行测试
test_psi_identical_distributions()
test_psi_different_distributions()
test_psi_smoothing()
print("所有测试通过！")

通过以上实现，我们获得了一个健壮、可复用的 PSI 计算工具，为后续的监控阈值设定和告警策略奠定了数据基础。