Selection on Observables & Instrumental Variables

"""
NorthMart's loyalty programme (simulated). Provided for Module 4: read the docstrings, not the code,
until the Day 40 capstone is done (the code contains the answers).

    from causalkit.loyalty import world, truth

    df = world()                         # 20,000 customers; who joined the programme, and next-year spend
    df = world(hidden=1.0)               # same, but an unobserved trait also drives joining and spending
    truth(df)                            # true ATE, ATT, ATC and LATE (for checking only)

Columns an analyst has:
    customer_id, tenure_years, age, mobile, engaged (opens marketing emails),
    pre_spend (last 12 months, $), pre_orders (last 12 months),
    invited (1 = randomly invited to join with a $20 sign-up bonus),
    joined (1 = joined the loyalty programme: the TREATMENT),
    spend (next 12 months, $: the OUTCOME)
Columns starting with '_' are the simulator's hidden truth.
"""
import numpy as np
import polars as pl


def _sigmoid(x):
    return 1 / (1 + np.exp(-x))


def world(n: int = 20_000, seed: int = 31, hidden: float = 0.0, overlap: float = 1.0,
          invite_effect: float = 0.9, bonus_spend_effect: float = 0.0) -> pl.DataFrame:
    """Simulate one loyalty-programme dataset.
    hidden:  strength of an UNOBSERVED trait (brand affinity) that drives both joining and spending.
             0 = selection on observables holds (Days 31-36); 1 = it fails (Days 37-40).
    overlap: >1 makes joining more predictable from the covariates (worse overlap, Day 35).
    invite_effect: how strongly the random invitation raises joining (weak instruments, Day 38).
    bonus_spend_effect: $ the invitation adds to spend DIRECTLY (breaks the exclusion restriction, Day 38)."""
    rng = np.random.default_rng(seed)
    tenure = rng.uniform(0, 8, n).round(1)
    age = np.clip(rng.normal(45, 14, n), 18, 80).round()
    mobile = (rng.random(n) < 0.6).astype(np.int8)
    engaged = (rng.random(n) < _sigmoid(-0.5 + 0.2 * tenure)).astype(np.int8)
    pre_spend = rng.lognormal(np.log(400) + 0.08 * tenure, 0.7, n).round(2)
    pre_orders = rng.poisson(pre_spend / 80)
    affinity = rng.normal(size=n)                                  # hidden brand affinity
    invited = (rng.random(n) < 0.5).astype(np.int8)                # randomised encouragement

    ls = np.log(pre_spend / 400)
    index = overlap * (-2.0 + 1.1 * ls + 0.9 * engaged + 0.15 * tenure - 0.01 * (age - 45)) + hidden * 1.0 * affinity
    u = rng.random(n)
    joined_if_not_invited = (u < _sigmoid(index)).astype(np.int8)
    joined_if_invited = (u < _sigmoid(index + invite_effect)).astype(np.int8)
    joined = np.where(invited == 1, joined_if_invited, joined_if_not_invited)
    ps = _sigmoid(index + invite_effect * invited)                 # true propensity (given everything)

    y0 = (150 + 0.75 * pre_spend + 0.0001 * pre_spend**2 + 110 * engaged + 60 * np.sqrt(tenure)
          + 15 * mobile + hidden * 90 * affinity + bonus_spend_effect * invited + rng.normal(0, 120, n))
    tau = 50 + 0.12 * pre_spend + 40 * engaged                     # effect of joining, larger for big spenders
    y1 = y0 + tau
    spend = np.where(joined == 1, y1, y0).round(2)

    return pl.DataFrame({
        "customer_id": [f"M{i:05d}" for i in range(1, n + 1)],
        "tenure_years": tenure, "age": age.astype(np.int16), "mobile": mobile, "engaged": engaged,
        "pre_spend": pre_spend, "pre_orders": pre_orders, "invited": invited,
        "joined": joined.astype(np.int8), "spend": spend,
        "_y0": y0, "_y1": y1, "_tau": tau, "_ps": ps, "_affinity": affinity,
        "_complier": ((joined_if_invited == 1) & (joined_if_not_invited == 0)).astype(np.int8),
    })


def truth(df: pl.DataFrame) -> dict:
    """True average effects of joining in THIS dataset (they need the hidden columns)."""
    t = df["_tau"].to_numpy()
    d = df["joined"].to_numpy() == 1
    c = df["_complier"].to_numpy() == 1
    return {"ATE": float(t.mean()), "ATT": float(t[d].mean()), "ATC": float(t[~d].mean()),
            "LATE": float(t[c].mean()), "share_joined": float(d.mean()), "share_compliers": float(c.mean())}


def public(df: pl.DataFrame) -> pl.DataFrame:
    """Drop the hidden columns: what an analyst would actually receive."""
    return df.select([c for c in df.columns if not c.startswith("_")])

"""
Module 4 estimators, built up day by day. Every function takes a pandas DataFrame with
columns `joined` (treatment), `spend` (outcome) and the covariates you name.

    from causalkit.estimators import COVARIATES, g_computation, propensity, ...
"""
import numpy as np
import pandas as pd
import statsmodels.formula.api as smf

COVARIATES = ["pre_spend", "pre_orders", "engaged", "tenure_years", "age", "mobile"]
OUTCOME_FORMULA = ("spend ~ pre_spend + I(pre_spend**2) + pre_orders + engaged"
                   " + tenure_years + I(tenure_years**0.5) + age + mobile")


def g_computation(df: pd.DataFrame, formula: str = OUTCOME_FORMULA) -> dict:
    """Fit one outcome model per group, predict BOTH potential outcomes for everyone, average."""
    m1 = smf.ols(formula, df[df.joined == 1]).fit()
    m0 = smf.ols(formula, df[df.joined == 0]).fit()
    y1_hat, y0_hat = m1.predict(df), m0.predict(df)
    diff = y1_hat - y0_hat
    d = df.joined == 1
    return {"ATE": diff.mean(), "ATT": diff[d].mean(), "ATC": diff[~d].mean()}


def bootstrap(estimator, df: pd.DataFrame, reps: int = 200, seed: int = 0, **kwargs) -> dict:
    """Standard errors by resampling customers with replacement and re-estimating."""
    rng = np.random.default_rng(seed)
    draws = [estimator(df.iloc[rng.integers(0, len(df), len(df))].reset_index(drop=True), **kwargs)
             for _ in range(reps)]
    return {k: float(np.std([d[k] for d in draws], ddof=1)) for k in draws[0]}


PS_FORMULA = "joined ~ np.log(pre_spend) + pre_orders + engaged + tenure_years + age + mobile"


def propensity(df: pd.DataFrame, formula: str = PS_FORMULA) -> np.ndarray:
    """P(joined = 1 | covariates), by logistic regression."""
    return smf.logit(formula, df).fit(disp=False).predict(df).to_numpy()


def smd(a: pd.Series, b: pd.Series, wa=None, wb=None) -> float:
    """Standardised mean difference, optionally weighted (Days 33-34)."""
    ma, mb = np.average(a, weights=wa), np.average(b, weights=wb)
    return (ma - mb) / np.sqrt((a.var() + b.var()) / 2)


MATCH_VARS = ["log_pre_spend", "engaged", "tenure_years", "age"]


def match_att(df: pd.DataFrame, on: str = "ps", caliper: float | None = None,
              bias_correct: bool = False) -> dict:
    """1-to-1 nearest-neighbour matching WITH replacement, for the ATT.
    on="ps": match on the logit of the propensity score; on="mahalanobis": on the covariates.
    caliper (in SDs of the matching distance): drop treated customers with no close match.
    bias_correct: adjust each pair for its remaining covariate gap (Abadie & Imbens 2011)."""
    from sklearn.neighbors import NearestNeighbors

    d = df.assign(log_pre_spend=np.log(df.pre_spend)).reset_index(drop=True)
    if on == "ps":
        p = propensity(d)
        X = np.log(p / (1 - p)).reshape(-1, 1)
    else:
        Z = d[MATCH_VARS].to_numpy(float)
        L = np.linalg.cholesky(np.linalg.inv(np.cov(Z, rowvar=False)))
        X = Z @ L                                            # Mahalanobis distance = Euclidean here
    t, c = np.flatnonzero(d.joined == 1), np.flatnonzero(d.joined == 0)
    dist, idx = NearestNeighbors(n_neighbors=1).fit(X[c]).kneighbors(X[t])
    dist, match = dist[:, 0], c[idx[:, 0]]
    keep = np.ones(len(t), bool) if caliper is None else dist <= caliper * X.std()
    y_t, y_c = d.spend.to_numpy()[t[keep]], d.spend.to_numpy()[match[keep]]
    if bias_correct:
        m0 = smf.ols(OUTCOME_FORMULA, d.iloc[c]).fit()
        y_c = y_c + (m0.predict(d.iloc[t[keep]]).to_numpy() - m0.predict(d.iloc[match[keep]]).to_numpy())
    return {"ATT": float((y_t - y_c).mean()), "matched_treated": int(keep.sum()),
            "distinct_controls_used": int(len(np.unique(match[keep]))),
            "_pairs": (t[keep], match[keep])}


def ipw_weights(d: np.ndarray, ps: np.ndarray, estimand: str = "ATE") -> np.ndarray:
    """ATE: everyone re-weighted to look like the whole population.
    ATT: non-joiners re-weighted to look like the joiners (joiners keep weight 1)."""
    if estimand == "ATE":
        return np.where(d == 1, 1 / ps, 1 / (1 - ps))
    return np.where(d == 1, 1.0, ps / (1 - ps))


def ipw(df: pd.DataFrame, ps: np.ndarray | None = None, normalised: bool = True) -> dict:
    """IPW estimates of the ATE and ATT. normalised=True is the Hajek version (weights sum to 1 per group)."""
    ps = propensity(df) if ps is None else ps
    d, y = df.joined.to_numpy(), df.spend.to_numpy()
    out = {}
    for est in ["ATE", "ATT"]:
        w = ipw_weights(d, ps, est)
        if normalised:
            out[est] = np.average(y[d == 1], weights=w[d == 1]) - np.average(y[d == 0], weights=w[d == 0])
        else:   # Horvitz-Thompson: divide by the number of units, not the sum of weights
            n = len(y) if est == "ATE" else (d == 1).sum()
            out[est] = (w * y * d).sum() / n - (w * y * (1 - d)).sum() / n
    return out


def effective_sample_size(w: np.ndarray) -> float:
    """How many equally weighted observations the weights are 'worth' (Kish)."""
    return w.sum() ** 2 / (w**2).sum()


def _fit_predict_outcome(train: pd.DataFrame, test: pd.DataFrame, method: str, formula: str):
    if method == "ols":
        return {g: smf.ols(formula, train[train.joined == g]).fit().predict(test).to_numpy() for g in [0, 1]}
    from sklearn.ensemble import HistGradientBoostingRegressor
    out = {}
    for g in [0, 1]:
        m = HistGradientBoostingRegressor(max_iter=300, learning_rate=0.05, random_state=0)
        m.fit(train.loc[train.joined == g, COVARIATES], train.spend[train.joined == g])
        out[g] = m.predict(test[COVARIATES])
    return out


def _fit_predict_ps(train: pd.DataFrame, test: pd.DataFrame, method: str, formula: str) -> np.ndarray:
    if method == "logit":
        p = smf.logit(formula, train).fit(disp=False).predict(test).to_numpy()
    else:
        from sklearn.ensemble import HistGradientBoostingClassifier
        m = HistGradientBoostingClassifier(max_iter=200, learning_rate=0.05, random_state=0)
        p = m.fit(train[COVARIATES], train.joined).predict_proba(test[COVARIATES])[:, 1]
    return np.clip(p, 0.01, 0.99)


def aipw(df: pd.DataFrame, outcome: str = "ols", ps_model: str = "logit",
         outcome_formula: str = OUTCOME_FORMULA, ps_formula: str = PS_FORMULA,
         folds: int = 5, seed: int = 0) -> dict:
    """Doubly robust estimates of the ATE and ATT, with standard errors from the influence function.
    Each customer's predictions come from models fitted on the OTHER folds (cross-fitting)."""
    from sklearn.model_selection import KFold

    df = df.reset_index(drop=True)
    n = len(df)
    mu0, mu1, ps = np.empty(n), np.empty(n), np.empty(n)
    splits = KFold(folds, shuffle=True, random_state=seed).split(df) if folds > 1 else [(np.arange(n), np.arange(n))]
    for train_idx, test_idx in splits:
        train, test = df.iloc[train_idx], df.iloc[test_idx]
        mu = _fit_predict_outcome(train, test, outcome, outcome_formula)
        mu0[test_idx], mu1[test_idx] = mu[0], mu[1]
        ps[test_idx] = _fit_predict_ps(train, test, ps_model, ps_formula)

    d, y = df.joined.to_numpy(), df.spend.to_numpy()
    psi_ate = mu1 - mu0 + d * (y - mu1) / ps - (1 - d) * (y - mu0) / (1 - ps)
    p1 = d.mean()
    psi_att = (d * (y - mu0) - (1 - d) * ps / (1 - ps) * (y - mu0)) / p1
    return {"ATE": psi_ate.mean(), "ATE_se": psi_ate.std(ddof=1) / np.sqrt(n),
            "ATT": psi_att.mean(), "ATT_se": psi_att.std(ddof=1) / np.sqrt(n)}


IV_CONTROLS = ["pre_spend", "pre_spend_sq", "pre_orders", "engaged", "tenure_years", "sqrt_tenure", "age", "mobile"]


def two_sls(df: pd.DataFrame, controls: list[str] | None = IV_CONTROLS):
    """2SLS of spend on joined, instrumenting joined with invited. Returns the linearmodels result."""
    from linearmodels.iv import IV2SLS

    d = df.assign(const=1.0, pre_spend_sq=df.pre_spend**2 / 1e4, sqrt_tenure=np.sqrt(df.tenure_years))
    exog = d[["const"] + (controls or [])]
    return IV2SLS(d.spend, exog, d.joined, d.invited).fit(cov_type="robust")