Statistical Foundations for Experimentation Questions

import numpy as np

def ipw_ate(treatment, outcome, propensity, clip=0.01):
    """
    Returns stabilized-IPW ATE and robust SE.
    treatment: binary array (0/1)
    outcome: numeric array
    propensity: estimated P(T=1|X), same length
    clip: min prop after clipping (prop in [clip, 1-clip])
    """
    treatment = np.asarray(treatment).astype(float)
    outcome = np.asarray(outcome).astype(float)
    propensity = np.asarray(propensity).astype(float)
    n = len(treatment)
    if not (len(outcome) == n and len(propensity) == n):
        raise ValueError("All inputs must have same length")
    # clip propensity
    eps = clip
    propensity = np.clip(propensity, eps, 1 - eps)
    p_bar = treatment.mean()
    # stabilized weights
    sw = treatment * (p_bar / propensity) + (1 - treatment) * ((1 - p_bar) / (1 - propensity))
    # weighted means (normalized by sum of weights in each group)
    w_t = sw * treatment
    w_c = sw * (1 - treatment)
    sum_wt = w_t.sum()
    sum_wc = w_c.sum()
    if sum_wt == 0 or sum_wc == 0:
        raise ValueError("One group has zero total weight")
    mu1 = (w_t * outcome).sum() / sum_wt
    mu0 = (w_c * outcome).sum() / sum_wc
    ate = mu1 - mu0
    # influence function for robust SE
    # IF_i = sw_i * [ T_i*(Y_i - mu1)/sum_wt - (1-T_i)*(Y_i - mu0)/sum_wc ]
    IF = sw * ( treatment * (outcome - mu1) / sum_wt - (1 - treatment) * (outcome - mu0) / sum_wc )
    var_if = IF.var(ddof=0)  # population variance
    se = np.sqrt(var_if / n)
    return {'ate': float(ate), 'se': float(se), 'mean_weight': float(sw.mean()), 'max_weight': float(sw.max()),
            'ess': float((sw.sum()**2 / (sw**2).sum())) }

import pandas as pd
import numpy as np
import statsmodels.api as sm
from statsmodels.regression.linear_model import GLSAR
from statsmodels.stats.sandwich_covariance import cov_hac

# assume df_control and df_treat provided with ['date','conversions','users']
df_control = df_control.copy(); df_treat = df_treat.copy()
df_control['group'] = 'control'; df_treat['group'] = 'treatment'
df = pd.concat([df_control, df_treat], ignore_index=True)
df['date'] = pd.to_datetime(df['date'])
df = df.sort_values(['group','date']).reset_index(drop=True)

# outcome = conversion rate
df['rate'] = df['conversions'] / df['users']

# create day index and post indicator (treatment starts on day 11)
start = df['date'].min()
df['day'] = (df['date'] - start).dt.days + 1
df['post'] = (df['day'] >= 11).astype(int)
df['treat'] = (df['group'] == 'treatment').astype(int)
df['treat_post'] = df['treat'] * df['post']

# create lagged outcome by group to avoid leakage across groups
df['rate_lag1'] = df.groupby('group')['rate'].shift(1)
df = df.dropna(subset=['rate_lag1'])  # drop first day per group

# regressors: intercept, lag, treat, post, treat_post (DiD with lag)
X = sm.add_constant(df[['rate_lag1','treat','post','treat_post']])
y = df['rate']

# GLSAR to model AR(1) errors
model = GLSAR(y, X, rho=1)   # start with rho=1 iteration
res = model.iterative_fit(5) # iterate to refine rho estimate

# robust covariance: Newey-West (HAC) with lag=1 or more based on autocorr
max_lag = 1
hac_cov = cov_hac(res, nlags=max_lag)
se_hac = np.sqrt(np.diag(hac_cov))

# coefficient, robust SE, 95% CI for treatment effect (treat_post)
coef = res.params['treat_post']
se = se_hac[list(res.params.index).index('treat_post')]
ci_lower = coef - 1.96 * se
ci_upper = coef + 1.96 * se

print("Estimated treatment effect (treat_post):", coef)
print("HAC SE:", se, "95% CI:", (ci_lower, ci_upper))