How to Compute CUPED Variance Reduction in Python

import numpy as np

def compute_cuped_reduction(control, variant, covariate):
    """
    Compute variance reduction using CUPED (Controlled Experiment with
    Pre-Experiment Data). Uses pre-experiment covariate values to
    reduce variance of the treatment effect estimate.
    """
    control = np.asarray(control, dtype=float)
    variant = np.asarray(variant, dtype=float)
    covariate = np.asarray(covariate, dtype=float)

    n_control = len(control)
    n_variant = len(variant)

    # Estimate theta (the optimal coefficient) from control group only
    cov_control = np.cov(control, covariate, ddof=1)
    theta = cov_control[0, 1] / cov_control[1, 1]

    # Adjust both groups using the same covariate values (shared pre-period data)
    control_adj = control - theta * (covariate[:n_control] - covariate[:n_control].mean())
    variant_adj = variant - theta * (covariate[n_control:n_control + n_variant] - covariate[n_control:n_control + n_variant].mean())

    # Sample variances of raw and adjusted metrics
    var_control_raw = np.var(control, ddof=1)
    var_variant_raw = np.var(variant, ddof=1)
    var_control_adj = np.var(control_adj, ddof=1)
    var_variant_adj = np.var(variant_adj, ddof=1)

    raw_variance = var_control_raw / n_control + var_variant_raw / n_variant
    adj_variance = var_control_adj / n_control + var_variant_adj / n_variant

    reduction = (1 - adj_variance / raw_variance) * 100
    return reduction, theta, raw_variance, adj_variance


if __name__ == "__main__":
    # Mock experiment data
    rng = np.random.default_rng(42)
    n = 2000
    pre_metric = rng.normal(100, 15, n)
    noise = rng.normal(0, 10, n)
    control = pre_metric + noise
    variant = pre_metric + noise + 5.0  # true lift of 5

    reduction, theta, raw_var, adj_var = compute_cuped_reduction(control, variant, pre_metric)
    print(f"Theta: {theta:.4f}")
    print(f"Raw variance: {raw_var:.4f}")
    print(f"CUPED-adjusted variance: {adj_var:.4f}")
    print(f"Variance reduction: {reduction:.2f}%")