Causal Thinking, Estimands & DAGs

"""
NorthMart's hidden truth: BOTH potential outcomes for every customer. Provided for Module 2.

Run once from the repository root (after data/make_northmart.py):
    python data/northmart_truth.py

Writes data/truth/potential_outcomes.parquet with one row per customer:
    treated        1 if the customer got the SPRING26 email (what really happened)
    y_observed     completed orders in weeks 9-16 (what the data shows)
    y0, y1         orders in weeks 9-16 WITHOUT and WITH the email
                   (one of them is y_observed; the other is simulated by this script)
    tau            y1 - y0: the individual effect (never observable in real life)
    tau_expected   the expected individual effect, without Poisson noise

No real dataset has these columns. We use them to check our methods against the truth.
"""
from pathlib import Path

import numpy as np
import polars as pl

import make_northmart as nm   # the Module 1 generator, in the same folder

CANCEL_RATE = 0.03            # the generator cancels ~3% of orders at random
POST_WEEKS = range(nm.CAMPAIGN_WEEK, nm.N_WEEKS + 1)
OUT = Path(__file__).parent / "truth"


def main() -> None:
    # 1. Replay the generator's random draws to recover each customer's latent rate and treatment
    rng = np.random.default_rng(nm.SEED)
    customers = nm.make_customers(rng)
    nm.make_products(rng)
    treated = nm.assign_campaign(rng, customers)

    recorded = set(pl.read_csv(nm.OUT / "email_campaign.csv")["customer_id"])
    assert recorded == set(customers.filter(pl.Series(treated))["customer_id"]), "replay mismatch"

    # 2. Observed post-period orders (completed, deduplicated) from the raw file
    orders = (pl.read_csv(nm.OUT / "orders.csv", try_parse_dates=True)
              .unique().filter(pl.col("status") == "completed")
              .filter(pl.col("order_date") >= pl.date(2026, 3, 2)))
    y_obs = orders.group_by("customer_id").agg(y_observed=pl.len())

    # 3. Expected orders over weeks 9-16 without the email (sum of weekly Poisson means)
    base = customers["_base_rate"].to_numpy()
    season = sum(1.0 + 0.10 * np.sin(2 * np.pi * w / 16) for w in POST_WEEKS)
    mu0 = base * season * (1 - CANCEL_RATE)
    mu1 = mu0 * (1 + nm.TRUE_EFFECT)

    # 4. Draw the potential outcome that did NOT happen, coupled to the one that did:
    #    with the email, a customer makes every order they would have made anyway, plus extra ones.
    #    Treated:  y0 keeps each observed order with probability 1/1.15 (so y0 <= y1).
    #    Control:  y1 adds extra orders, Poisson with mean mu1 - mu0 (so y1 >= y0).
    rng_cf = np.random.default_rng(nm.SEED + 1)
    df = (customers.select("customer_id", "is_loyalty", "region")
          .with_columns(treated=pl.Series(treated.astype(np.int8)),
                        mu0=pl.Series(mu0), mu1=pl.Series(mu1))
          .join(y_obs, on="customer_id", how="left")
          .with_columns(pl.col("y_observed").fill_null(0).cast(pl.Int64)))
    y = df["y_observed"].to_numpy()
    kept = rng_cf.binomial(y, 1 / (1 + nm.TRUE_EFFECT))
    extra = rng_cf.poisson(mu1 - mu0)
    t = treated
    df = df.with_columns(
        y0=pl.Series(np.where(t, kept, y)),
        y1=pl.Series(np.where(t, y, y + extra)),
    ).with_columns(
        tau=pl.col("y1") - pl.col("y0"),
        tau_expected=pl.col("mu1") - pl.col("mu0"),
    )

    OUT.mkdir(exist_ok=True)
    df.select("customer_id", "treated", "is_loyalty", "region", "y_observed", "y0", "y1",
              "tau", "tau_expected", "mu0", "mu1").write_parquet(OUT / "potential_outcomes.parquet")
    print(f"Wrote {df.height:,} customers to {OUT / 'potential_outcomes.parquet'}")


if __name__ == "__main__":
    main()