Experimentation & A/B Testing

"""
NorthMart's experimentation lab (simulated). Provided for Module 3: read the docstrings, not the code,
until the capstone is done (the code contains the true effects).

    from causalkit.lab import population, run, truth

    pop = population()                      # 40,000 customers with 8 weeks of pre-period history
    out = run(pop, treated, "email_v2")     # outcomes over the 4-week test window
    truth("email_v2")                       # the true average effects (for checking only)

A real company has the first two and never the third.
"""
from functools import lru_cache

import numpy as np
import polars as pl

N_CUSTOMERS = 40_000
PRE_WEEKS = 8
TEST_WEEKS = 4
SHIPPING_FEE = 7.0
OLD_THRESHOLD, NEW_THRESHOLD = 50.0, 35.0
GROSS_MARGIN = 0.30

EXPERIMENTS = {
    # lift in order rate, by loyalty status (non-members, members)
    "email_v2":      {"lift": (0.05, 0.05)},
    "free_shipping": {"lift": (0.10, 0.03)},
    "aa_test":       {"lift": (0.0, 0.0)},
}


@lru_cache(maxsize=1)
def _population(seed: int) -> pl.DataFrame:
    rng = np.random.default_rng(seed)
    n = N_CUSTOMERS
    loyalty = rng.random(n) < 0.30
    rate = rng.lognormal(-1.6, 0.6, n) * np.where(loyalty, 1.6, 1.0)        # orders per week
    basket = rng.lognormal(np.log(55), 0.45, n)                             # typical order value
    household = rng.integers(0, int(n / 2.2), n)                            # ~2.2 customers per household
    pre_orders = rng.poisson(rate * PRE_WEEKS)
    pre_revenue = _order_values(rng, pre_orders, basket).round(2)
    return pl.DataFrame({
        "customer_id": [f"L{i:05d}" for i in range(1, n + 1)],
        "household_id": [f"H{h:05d}" for h in household],
        "region": rng.choice(["ON", "QC", "BC", "AB", "ATL"], n, p=[0.38, 0.23, 0.15, 0.12, 0.12]),
        "device": rng.choice(["mobile", "desktop"], n, p=[0.62, 0.38]),
        "loyalty": loyalty.astype(np.int8),
        "pre_orders": pre_orders,
        "pre_revenue": pre_revenue,
        "_rate": rate,
        "_basket": basket,
    })


def _order_values(rng, n_orders: np.ndarray, basket: np.ndarray) -> np.ndarray:
    """Total value of n_orders[i] orders for each customer i."""
    owner = np.repeat(np.arange(len(n_orders)), n_orders)
    values = basket[owner] * rng.lognormal(-0.045, 0.3, owner.size)
    return np.bincount(owner, weights=values, minlength=len(n_orders))


def population(seed: int = 2030) -> pl.DataFrame:
    """40,000 customers: id, household, region, device, loyalty, and 8 weeks of pre-period orders/revenue.
    Columns starting with '_' are the simulator's hidden truth; don't use them in an analysis."""
    return _population(seed).clone()


def run(pop: pl.DataFrame, treated, experiment: str, seed: int = 0,
        spillover: float = 0.0) -> pl.DataFrame:
    """Run an experiment for 4 weeks. `treated` is a 0/1 array aligned with `pop`.
    Returns customer_id, treated, orders, revenue, converted (any order), and margin
    (30% of revenue, minus shipping NorthMart absorbs under free shipping).
    `spillover` (Day 28): share of a treated household member's lift that reaches the others."""
    spec = EXPERIMENTS[experiment]
    rng = np.random.default_rng(seed)
    t = np.asarray(treated).astype(bool)
    loyal = pop["loyalty"].to_numpy().astype(bool)
    lift = np.where(loyal, spec["lift"][1], spec["lift"][0])
    exposure = t.astype(float)
    if spillover:
        hh = pop["household_id"].to_numpy()
        _, inv = np.unique(hh, return_inverse=True)
        treated_in_hh = np.bincount(inv, weights=t)[inv] - t
        exposure = np.where(t, 1.0, np.minimum(treated_in_hh, 1) * spillover)
    rate = pop["_rate"].to_numpy() * (1 + lift * exposure)
    orders = rng.poisson(rate * TEST_WEEKS)

    owner = np.repeat(np.arange(len(orders)), orders)
    values = pop["_basket"].to_numpy()[owner] * rng.lognormal(-0.045, 0.3, owner.size)
    revenue = np.bincount(owner, weights=values, minlength=len(orders))
    absorbed = np.zeros(len(orders))
    if experiment == "free_shipping":
        free_now = t[owner] & (values >= NEW_THRESHOLD) & (values < OLD_THRESHOLD)
        absorbed = np.bincount(owner, weights=free_now * SHIPPING_FEE, minlength=len(orders))
    return pl.DataFrame({
        "customer_id": pop["customer_id"],
        "treated": t.astype(np.int8),
        "orders": orders,
        "revenue": revenue.round(2),
        "converted": (orders > 0).astype(np.int8),
        "margin": (GROSS_MARGIN * revenue - absorbed).round(2),
    })


def truth(experiment: str, seed: int = 2030) -> dict:
    """True average effects per customer over the 4-week window (ATE over the population)."""
    pop = _population(seed)
    rate = pop["_rate"].to_numpy()
    basket = pop["_basket"].to_numpy()
    loyal = pop["loyalty"].to_numpy().astype(bool)
    spec = EXPERIMENTS[experiment]
    lift = np.where(loyal, spec["lift"][1], spec["lift"][0])
    mu0 = rate * TEST_WEEKS
    mu1 = mu0 * (1 + lift)
    d_orders = mu1 - mu0
    d_revenue = d_orders * basket
    d_conv = np.exp(-mu0) - np.exp(-mu1)
    d_margin = GROSS_MARGIN * d_revenue
    if experiment == "free_shipping":
        # share of orders with value in [35, 50): order value = basket * lognormal(-0.045, 0.3)
        from scipy.stats import norm
        lo = (np.log(NEW_THRESHOLD / basket) + 0.045) / 0.3
        hi = (np.log(OLD_THRESHOLD / basket) + 0.045) / 0.3
        share = norm.cdf(hi) - norm.cdf(lo)
        d_margin = d_margin - SHIPPING_FEE * share * mu1
    return {
        "orders": float(d_orders.mean()),
        "revenue": float(d_revenue.mean()),
        "converted": float(d_conv.mean()),
        "margin": float(d_margin.mean()),
        "orders_relative": float(d_orders.sum() / mu0.sum()),
        "revenue_relative": float(d_revenue.sum() / (mu0 * basket).sum()),
    }


def run_weekly(pop: pl.DataFrame, treated, experiment: str, seed: int = 0, weeks: int = TEST_WEEKS) -> pl.DataFrame:
    """Like run(), but one row per customer per week (week 1..weeks), so results can be
    looked at as they accumulate (Days 26-27). Columns: customer_id, week, treated, orders, revenue."""
    frames = []
    for w in range(1, weeks + 1):
        one = run(pop, treated, experiment, seed=seed * 1000 + w)
        # run() simulates 4 weeks; scale to one week by thinning orders (each kept with prob 1/4)
        rng = np.random.default_rng(seed * 1000 + w + 500)
        orders = one["orders"].to_numpy()
        kept = rng.binomial(orders, 1 / TEST_WEEKS)
        share = np.divide(kept, orders, out=np.zeros(len(orders)), where=orders > 0)
        frames.append(pl.DataFrame({
            "customer_id": one["customer_id"], "week": np.full(len(orders), w, dtype=np.int8),
            "treated": one["treated"], "orders": kept,
            "revenue": (one["revenue"].to_numpy() * share).round(2),
        }))
    return pl.concat(frames)