How to Detect Data Drift with PSI in Python

import numpy as np

def calculate_psi(expected, actual, buckets=10):
    """Calculate Population Stability Index (PSI) between two distributions."""
    # Create bucket edges based on expected distribution percentiles
    edges = np.percentile(expected, np.linspace(0, 100, buckets + 1))
    edges[-1] = np.inf  # Ensure last bucket captures all values
    
    # Bucket the data
    expected_buckets, _ = np.histogram(expected, bins=edges)
    actual_buckets, _ = np.histogram(actual, bins=edges)
    
    # Convert to proportions
    expected_prop = expected_buckets / len(expected)
    actual_prop = actual_buckets / len(actual)
    
    # Avoid division by zero and log of zero
    expected_prop = np.where(expected_prop == 0, 0.0001, expected_prop)
    actual_prop = np.where(actual_prop == 0, 0.0001, actual_prop)
    
    # Calculate PSI
    psi = np.sum((actual_prop - expected_prop) * np.log(actual_prop / expected_prop))
    return psi

if __name__ == "__main__":
    # Mock data: expected distribution (training) vs actual distribution (live)
    rng = np.random.default_rng(42)
    expected_data = rng.normal(loc=50, scale=10, size=10000)  # Training distribution
    actual_same = rng.normal(loc=50, scale=10, size=10000)   # Same distribution
    actual_drifted = rng.normal(loc=55, scale=12, size=10000)  # Drifted distribution
    
    psi_same = calculate_psi(expected_data, actual_same)
    psi_drifted = calculate_psi(expected_data, actual_drifted)
    
    print(f"PSI (same distribution): {psi_same:.4f}")
    print(f"PSI (drifted distribution): {psi_drifted:.4f}")
    print("Interpretation: PSI < 0.1 = no drift, 0.1-0.25 = moderate, > 0.25 = significant drift")