Test Design & Avoiding Confounds Questions

import numpy as np
import pandas as pd

def block_permutation_test(df, day_col='day', variant_col='variant',
                           metric_col='metric', treat_label='treatment',
                           control_label='control', n_perms=1000, seed=None):
    """
    Block-wise permutation test: permute variant within each day block.
    Returns: observed_ate, p_value, permuted_ates (numpy array length n_perms)
    """
    rng = np.random.default_rng(seed)
    # Basic validation
    if {day_col, variant_col, metric_col} - set(df.columns):
        raise ValueError("Missing required columns")
    df = df[[day_col, variant_col, metric_col]].dropna()
    
    # Observed ATE
    mask_t = df[variant_col] == treat_label
    obs_mean_t = df.loc[mask_t, metric_col].mean()
    obs_mean_c = df.loc[~mask_t, metric_col].mean()
    observed_ate = obs_mean_t - obs_mean_c
    
    # Precompute per-day arrays: metrics, n_treat_in_day
    groups = df.groupby(day_col)
    day_info = []
    total_n_treat = mask_t.sum()
    total_n = len(df)
    total_n_control = total_n - total_n_treat
    if total_n_treat == 0 or total_n_control == 0:
        raise ValueError("Both treatment and control must be present")
    
    for _, g in groups:
        metrics = g[metric_col].to_numpy()
        n_treat = (g[variant_col] == treat_label).sum()
        # store metrics and number of treatment slots for that day
        day_info.append((metrics, int(n_treat)))
    
    perm_ates = np.empty(n_perms, dtype=float)
    # For each permutation, assign treatment indices per day and compute ATE
    for i in range(n_perms):
        sum_treat = 0.0
        for metrics, n_treat in day_info:
            if n_treat == 0:
                continue
            # sample indices for treatment slots without replacement
            idx = rng.choice(len(metrics), size=n_treat, replace=False)
            sum_treat += metrics[idx].sum()
        mean_treat = sum_treat / total_n_treat
        # total sum of metrics:
        # compute once per perm? we can precompute global sum outside loop to save time
        # but here we compute control mean via global sum - sum_treat
        # Precompute global_sum before loop:
        # (moved below)
        perm_ates[i] = mean_treat
    # improve by computing global sum once:
    global_sum = df[metric_col].sum()
    # recalc using better loop (to ensure correctness with global_sum)
    for i in range(n_perms):
        # recompute sums using RNG with same seed handling
        # simpler: redo deterministic approach: regenerate per-perm treatment sums
        sum_treat = 0.0
        for metrics, n_treat in day_info:
            if n_treat == 0:
                continue
            idx = rng.choice(len(metrics), size=n_treat, replace=False)
            sum_treat += metrics[idx].sum()
        mean_t = sum_treat / total_n_treat
        mean_c = (global_sum - sum_treat) / total_n_control
        perm_ates[i] = mean_t - mean_c

    # two-sided p-value
    abs_obs = abs(observed_ate)
    p_value = (np.count_nonzero(np.abs(perm_ates) >= abs_obs) + 1) / (n_perms + 1)
    return observed_ate, p_value, perm_ates