Synthetic Control

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

    from causalkit.regions import regions, truth, wide

    df = regions()          # 30 regions x 104 weeks; region R07 runs an 8-week TV campaign from week 80
    truth(df)               # true lift: average per post week, cumulative, during the campaign (for checking only)
    Y = wide(df)            # weeks x regions matrix of sales, for synthetic control

Columns an analyst has:
    region, week, sales ($ thousands per week), treated (1 = the campaign region), post (1 = week >= start)
Columns starting with '_' are the simulator's hidden truth: _y0 (sales without the campaign), _tau (lift).

Regions differ in size, seasonality and exposure to two national trends; one of them, urban growth,
accelerates over the two years, and some regions ride it much more than others.
"""
import numpy as np
import polars as pl


def regions(n_regions: int = 30, weeks: int = 104, start: int = 80, seed: int = 45, treated: str = "R07",
            campaign_weeks: int = 8, lift: float = 30.0, decay: float = 0.8, treated_urban: float = 1.6,
            treated_size: float = 0.0,
            noise: float = 4.0) -> pl.DataFrame:
    """treated_size: extra weekly sales of the campaign region ($ thousands) on top of the median region.
    lift:   true weekly sales lift during the campaign ($ thousands); after it ends the lift decays
               by `decay` per week (carry-over).
    treated_urban: the campaign region's exposure to urban growth (other regions: uniform 0 to 2)."""
    rng = np.random.default_rng(seed)
    t = np.arange(1, weeks + 1)
    names = [f"R{j:02d}" for j in range(1, n_regions + 1)]
    level = rng.normal(500, 60, n_regions)
    season_amp = rng.uniform(10, 30, n_regions)
    national = 0.4 * t + np.cumsum(rng.normal(0, 1.5, weeks))
    urban = 30 * (t / weeks) ** 2 + np.cumsum(rng.normal(0, 0.8, weeks))
    load_nat = rng.uniform(0.6, 1.4, n_regions)
    load_urb = rng.uniform(0, 2, n_regions)
    j = names.index(treated)
    load_urb[j] = treated_urban
    level[j] = np.median(level) + treated_size
    season = np.sin(2 * np.pi * (t - 10) / 52)
    y0 = (level[:, None] + season_amp[:, None] * season[None, :] + load_nat[:, None] * national[None, :]
          + load_urb[:, None] * urban[None, :] + rng.normal(0, noise, (n_regions, weeks)))
    k = t - start
    tau_t = np.where(k < 0, 0.0, np.where(k < campaign_weeks, lift, lift * decay ** (k - campaign_weeks + 1)))
    tau = np.zeros_like(y0)
    tau[j] = tau_t
    return pl.DataFrame({
        "region": np.repeat(names, weeks),
        "week": np.tile(t, n_regions).astype(np.int64),
        "sales": np.round((y0 + tau).ravel(), 2),
        "treated": np.repeat((np.arange(n_regions) == j).astype(np.int64), weeks),
        "post": np.tile((t >= start).astype(np.int64), n_regions),
        "_y0": y0.ravel(),
        "_tau": tau.ravel(),
    })


def truth(df: pl.DataFrame, campaign_weeks: int = 8) -> dict:
    tr = df.filter((pl.col("treated") == 1) & (pl.col("post") == 1)).sort("week")
    tau = tr["_tau"].to_numpy()
    return {"avg_post_lift": round(float(tau.mean()), 2), "cumulative_lift": round(float(tau.sum()), 1),
            "campaign_weeks_lift": round(float(tau[:campaign_weeks].sum()), 1)}


def wide(df: pl.DataFrame, col: str = "sales") -> pl.DataFrame:
    """Weeks in rows, one column per region."""
    return df.pivot(on="region", index="week", values=col).sort("week")


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

"""Synthetic control, synthetic DiD and their placebo tests (Days 45-49)."""
import numpy as np
import polars as pl
from scipy.optimize import minimize

from causalkit.regions import wide


def simplex_ls(X: np.ndarray, y: np.ndarray, intercept: bool = False, ridge: float = 0.0):
    """Weights w >= 0 summing to 1 that minimise ||y - (a) - X w||^2 + ridge * ||w||^2.
    With intercept=True, a free constant a absorbs any level difference. Returns (w, a)."""
    k = X.shape[1]
    if intercept:                           # the optimal constant is the mean gap: demean, then solve
        Xc, yc = X - X.mean(0), y - y.mean()
    else:
        Xc, yc = X, y

    def loss(w):
        r = yc - Xc @ w
        return r @ r + ridge * w @ w

    def grad(w):
        return -2 * Xc.T @ (yc - Xc @ w) + 2 * ridge * w

    res = minimize(loss, np.full(k, 1 / k), jac=grad, method="SLSQP", bounds=[(0, 1)] * k,
                   constraints=[{"type": "eq", "fun": lambda w: w.sum() - 1, "jac": lambda w: np.ones(k)}],
                   options={"maxiter": 1000, "ftol": 1e-12})
    w = np.clip(res.x, 0, None)
    w /= w.sum()
    a = float((y - X @ w).mean()) if intercept else 0.0
    return w, a


def _split(df: pl.DataFrame, treated: str, start: int):
    Y = wide(df)
    weeks = Y["week"].to_numpy()
    donors = [c for c in Y.columns if c not in ("week", treated)]
    y1 = Y[treated].to_numpy()
    Y0 = Y.select(donors).to_numpy()
    pre = weeks < start
    return weeks, donors, y1, Y0, pre


def did(df: pl.DataFrame, treated: str, start: int) -> dict:
    """Plain DiD: treated region vs the average of all other regions, post minus pre."""
    weeks, donors, y1, Y0, pre = _split(df, treated, start)
    gap = y1 - Y0.mean(1)
    return {"effect": float(gap[~pre].mean() - gap[pre].mean()), "gap": gap - gap[pre].mean(), "weeks": weeks}


def synthetic_control(df: pl.DataFrame, treated: str, start: int, intercept: bool = False) -> dict:
    """Abadie-style synthetic control: donor weights fitted on pre-period sales only.
    Returns the average post-period gap (effect), the weekly gap, weights and fit statistics."""
    weeks, donors, y1, Y0, pre = _split(df, treated, start)
    w, a = simplex_ls(Y0[pre], y1[pre], intercept=intercept)
    synth = a + Y0 @ w
    gap = y1 - synth
    rmspe_pre = float(np.sqrt((gap[pre] ** 2).mean()))
    rmspe_post = float(np.sqrt((gap[~pre] ** 2).mean()))
    return {"effect": float(gap[~pre].mean()), "gap": gap, "synthetic": synth, "actual": y1, "weeks": weeks,
            "weights": dict(sorted(((d, round(float(x), 3)) for d, x in zip(donors, w) if x > 0.005),
                                   key=lambda kv: -kv[1])),
            "rmspe_pre": rmspe_pre, "rmspe_post": rmspe_post, "ratio": rmspe_post / rmspe_pre}


def sdid(df: pl.DataFrame, treated: str, start: int) -> dict:
    """Synthetic difference-in-differences (Arkhangelsky et al., 2021), one treated unit.
    Unit weights: like SC but with an intercept and a ridge penalty.
    Time weights: pre-period weeks that best predict the donors' post-period average."""
    weeks, donors, y1, Y0, pre = _split(df, treated, start)
    T0, T1 = pre.sum(), (~pre).sum()
    zeta = (1 * T1) ** 0.25 * np.diff(Y0[pre], axis=0).std(ddof=1)
    omega, _ = simplex_ls(Y0[pre], y1[pre], intercept=True, ridge=zeta**2 * T0)
    lam, _ = simplex_ls(Y0[pre].T, Y0[~pre].mean(0), intercept=True, ridge=1e-6 * T0)
    treated_change = y1[~pre].mean() - lam @ y1[pre]
    donor_change = (Y0[~pre].mean(0) - lam @ Y0[pre]) @ omega
    return {"effect": float(treated_change - donor_change),
            "unit_weights": dict(sorted(((d, round(float(x), 3)) for d, x in zip(donors, omega) if x > 0.005),
                                        key=lambda kv: -kv[1])),
            "time_weights": {int(w): round(float(x), 3) for w, x in zip(weeks[pre], lam) if x > 0.005}}


def placebo_in_space(df: pl.DataFrame, treated: str, start: int, intercept: bool = False) -> pl.DataFrame:
    """Re-run synthetic control pretending each region was treated (the real one is dropped from donor pools).
    Inference: where does the real region's post/pre RMSPE ratio rank?"""
    rows = []
    for r in df["region"].unique().sort().to_list():
        sub = df if r == treated else df.filter(pl.col("region") != treated)
        sc = synthetic_control(sub, r, start, intercept=intercept)
        rows.append({"region": r, "is_treated": r == treated, "effect": sc["effect"],
                     "rmspe_pre": sc["rmspe_pre"], "ratio": sc["ratio"]})
    return pl.DataFrame(rows).sort("ratio", descending=True)

# day45/short_pre_period.py   (python -m day45.short_pre_period)
import polars as pl

from causalkit.regions import public, regions, truth
from causalkit.synth import synthetic_control

full = regions()
df = public(full)
for first_week in [1, 60]:
    sc = synthetic_control(df.filter(pl.col("week") >= first_week), "R07", 80)
    print(f"pre-period from week {first_week:>2}: effect {sc['effect']:.1f}, pre RMSPE {sc['rmspe_pre']:.1f}, "
          f"{len(sc['weights'])} donors")
print(f"truth {truth(full)['avg_post_lift']:.1f}")

pre-period from week  1: effect 17.3, pre RMSPE 3.9, 6 donors
pre-period from week 60: effect 27.4, pre RMSPE 2.9, 6 donors
truth 14.3

The short pre-period gives a better fit (RMSPE 2.9) and a worse answer ($27.4k). The weights fitted the noise of 20 weeks, not R07's long-run exposure to urban growth. A good fit over a short window is weak evidence; Abadie (2021) recommends a pre-period much longer than the number of donors that end up with weight.