Quasi-Experiments: DiD, Synthetic Control, RDD

"""
NorthMart store panel (simulated). Provided for Module 5: read the docstrings, not the code,
until the Day 50 capstone is done.

    from causalkit.stores import campaign, rollout, truth

    df = campaign()                  # 240 stores x 60 weeks; stores in two regions get a campaign from week 40
    df = campaign(trend_gap=0.15)    # same, but treated stores were already growing faster
    df = rollout()                   # a new store format, adopted in waves (weeks 15, 30, 45) or never
    truth(df)                        # true average effects (for checking only)

Columns an analyst has:
    store_id, region, week, sales ($ thousands per week),
    treated_group (1 = store is ever treated), cohort (first treated week; 0 = never),
    post (1 = week >= start; campaign only), d (1 = treated in this week)
Columns starting with '_' are the simulator's hidden truth: _y0 (sales without treatment), _tau (effect).
"""
import numpy as np
import polars as pl

REGIONS = ["BC", "AB", "SK", "MB", "ON", "QC", "NB", "NS"]


def _base(n_stores: int, weeks: int, rng: np.random.Generator, rho: float = 0.7, region_sd: float = 0.0):
    """Store effects, shared seasonality, regional shocks and AR(1) store noise."""
    region = rng.choice(REGIONS, n_stores)
    store_fe = rng.normal(100, 20, n_stores)
    t = np.arange(1, weeks + 1)
    week_fe = 4 * np.sin(2 * np.pi * t / 52) + 0.05 * t + rng.normal(0, 1.0, weeks)
    region_shock = {r: np.cumsum(rng.normal(0, 1, weeks)) * region_sd for r in REGIONS}
    e = np.zeros((n_stores, weeks))
    u = rng.normal(0, 2.5, (n_stores, weeks))
    e[:, 0] = u[:, 0] / np.sqrt(1 - rho**2)
    for w in range(1, weeks):
        e[:, w] = rho * e[:, w - 1] + u[:, w]
    y0 = store_fe[:, None] + week_fe[None, :] + np.vstack([region_shock[r] for r in region]) + e
    return region, t, y0


def _frame(region, t, y0, tau, d, cohort, extra=None) -> pl.DataFrame:
    n, w = y0.shape
    df = pl.DataFrame({
        "store_id": np.repeat([f"S{i:03d}" for i in range(1, n + 1)], w),
        "region": np.repeat(region, w),
        "week": np.tile(t, n).astype(np.int64),
        "sales": np.round((y0 + tau * d).ravel(), 2),
        "treated_group": np.repeat((cohort > 0).astype(np.int64), w),
        "cohort": np.repeat(cohort, w).astype(np.int64),
        "d": d.ravel().astype(np.int64),
        "_y0": y0.ravel(),
        "_tau": (tau * d).ravel(),
    })
    return df.with_columns(**(extra or {}))


def campaign(n_stores: int = 240, weeks: int = 60, start: int = 40, seed: int = 41,
             effect: float = 3.0, trend_gap: float = 0.0, treated_regions=("ON", "QC"),
             region_sd: float = 0.0) -> pl.DataFrame:
    """A regional campaign: every store in treated_regions is treated from week `start`.
    effect:    true sales lift, $ thousands per store-week (constant over time)
    trend_gap: extra weekly growth of treated stores' sales, before AND after the campaign
               (violates parallel trends when > 0)
    region_sd: size of week-to-week shocks shared by all stores in a region (random walk; 0 = none)"""
    rng = np.random.default_rng(seed)
    region, t, y0 = _base(n_stores, weeks, rng, region_sd=region_sd)
    g = np.isin(region, treated_regions)
    y0 = y0 + trend_gap * g[:, None] * (t[None, :] - start)
    d = g[:, None] & (t[None, :] >= start)
    tau = np.full(y0.shape, effect)
    cohort = np.where(g, start, 0)
    return _frame(region, t, y0, tau, d, cohort,
                  {"post": (pl.col("week") >= start).cast(pl.Int64)}).select(
        "store_id", "region", "week", "sales", "treated_group", "cohort", "post", "d", "_y0", "_tau")


def rollout(n_stores: int = 240, weeks: int = 60, seed: int = 47, cohorts=(15, 30, 45),
            growth: float = 0.4, cap: int = 20, dynamic: bool = True) -> pl.DataFrame:
    """A new store format adopted in waves. A quarter of stores adopt in each cohort week; a quarter never do.
    The effect grows by `growth` per week since adoption up to `cap` weeks (dynamic=True), and is larger
    for earlier adopters (stores that volunteered first were better suited to the format).
    dynamic=False gives the same average effect but constant over time."""
    rng = np.random.default_rng(seed)
    region, t, y0 = _base(n_stores, weeks, rng)
    groups = np.array(list(cohorts) + [0])
    cohort = groups[rng.permutation(np.arange(n_stores) % len(groups))]
    rel = t[None, :] - cohort[:, None]
    d = (cohort[:, None] > 0) & (rel >= 0)
    suit = np.select([cohort == cohorts[0], cohort == cohorts[1]], [1.5, 1.0], 0.6)[:, None]
    if dynamic:
        tau = suit * growth * np.minimum(rel + 1, cap)
    else:
        tau = suit * growth * np.minimum(np.maximum(rel, 0) + 1, cap).mean() * np.ones_like(y0)
    tau = np.where(d, tau, 0.0)
    return _frame(region, t, y0, tau, d, cohort)


def truth(df: pl.DataFrame) -> dict:
    """True ATT over treated store-weeks, and by cohort."""
    tr = df.filter(pl.col("d") == 1)
    out = {"ATT": round(tr["_tau"].mean(), 3)}
    for c in sorted(tr["cohort"].unique().to_list()):
        out[f"ATT_cohort_{c}"] = round(tr.filter(pl.col("cohort") == c)["_tau"].mean(), 3)
    return out


def public(df: pl.DataFrame) -> pl.DataFrame:
    return df.select([c for c in df.columns if not c.startswith("_")])