"""Selection bias in ad measurement: RCT vs exposed-vs-unexposed vs propensity matching.

A small simulation in the spirit of Gordon, Zettelmeyer, Bhargava and Chapsky
(2019, Marketing Science) and Gordon, Moakler and Zettelmeyer (2023, Marketing
Science). From the Vexpower MMM library:
vexpower.com/mmm/facebook-experiments-vs-observational-methods

Each simulated user has a latent purchase propensity. Users are randomized into
a test group (eligible for ads) and a control group (held out). In the test
group, the ad platform delivers ads preferentially to high-propensity users.
We then estimate the ad's lift on the people who saw it in three ways:

  1. RCT: compare everyone in test with everyone in control, then scale up by
     the share of the test group that was exposed (the ITT-to-ATT adjustment
     used in conversion lift studies).
  2. Naive: compare exposed with unexposed users inside the test group.
  3. Stratified propensity score matching (SPSM): fit a logistic propensity
     model on an observable feature, bin users into propensity strata and
     compare exposed with unexposed users within each stratum.

The observable only captures part of the latent propensity (OBS_SHARE), which
is the situation every real advertiser is in. All numbers are invented.

Run:  python selection-bias-ad-measurement.py
Needs numpy (and pandas only if you want the CSV written; it is optional).
"""
import numpy as np

SEED = 7
N_TEST = 400_000          # users in the test group
N_CONTROL = 400_000       # users in the randomized control group
BASE_RATE = 0.02          # average conversion rate without ads
TRUE_LIFT = 0.20          # ads multiply an exposed user's conversion probability by 1.2
REACH = 0.5               # share of the test group that ends up exposed
SPREAD = 1.5              # SD of the latent propensity on the log-odds scale
OBS_SHARE = 0.5           # share of the propensity's variance you can observe
N_STRATA = 20
TARGETING = [0.0, 0.5, 1.0, 1.5, 2.0, 3.0]


def sigmoid(z):
    return 1.0 / (1.0 + np.exp(-z))


def solve_intercept(z, target):
    """Intercept c so that mean(sigmoid(c + z)) == target (bisection)."""
    lo, hi = -20.0, 20.0
    for _ in range(60):
        mid = (lo + hi) / 2
        if sigmoid(mid + z).mean() < target:
            lo = mid
        else:
            hi = mid
    return (lo + hi) / 2


def logistic_fit(x, w, iters=25):
    """Logistic regression of w on [1, x] by Newton-Raphson."""
    X = np.column_stack([np.ones_like(x), x])
    beta = np.zeros(2)
    for _ in range(iters):
        p = sigmoid(X @ beta)
        grad = X.T @ (w - p)
        hess = X.T @ (X * (p * (1 - p))[:, None])
        beta += np.linalg.solve(hess, grad)
    return sigmoid(X @ beta)


def simulate(targeting, rng):
    # Latent propensity for test and control users (same distribution: randomized).
    u_t = rng.normal(size=N_TEST)
    u_c = rng.normal(size=N_CONTROL)
    a = solve_intercept(SPREAD * np.concatenate([u_t, u_c]), BASE_RATE)
    p0_t = sigmoid(a + SPREAD * u_t)
    p0_c = sigmoid(a + SPREAD * u_c)

    # Ad delivery in the test group favours high-propensity users.
    b = solve_intercept(targeting * u_t, REACH)
    exposed = rng.random(N_TEST) < sigmoid(b + targeting * u_t)

    # Outcomes.
    p_t = np.where(exposed, np.minimum(1, p0_t * (1 + TRUE_LIFT)), p0_t)
    y_t = rng.random(N_TEST) < p_t
    y_c = rng.random(N_CONTROL) < p0_c

    # 1. RCT: ITT difference scaled to the exposed (no ads reach the control group).
    itt = y_t.mean() - y_c.mean()
    share_exposed = exposed.mean()
    att = itt / share_exposed
    # Counterfactual conversion rate of the exposed = their observed rate minus the effect.
    exposed_rate = y_t[exposed].mean()
    rct_lift = att / (exposed_rate - att)

    # 2. Naive exposed vs unexposed inside the test group.
    naive_lift = exposed_rate / y_t[~exposed].mean() - 1

    # 3. Stratified propensity score matching on a partial observable.
    x = np.sqrt(OBS_SHARE) * u_t + np.sqrt(1 - OBS_SHARE) * rng.normal(size=N_TEST)
    ps = logistic_fit(x, exposed.astype(float))
    edges = np.quantile(ps, np.linspace(0, 1, N_STRATA + 1))
    strata = np.clip(np.searchsorted(edges, ps, side="right") - 1, 0, N_STRATA - 1)
    num, cf, w_total = 0.0, 0.0, 0.0
    for s in range(N_STRATA):
        m = strata == s
        e_m, u_m = m & exposed, m & ~exposed
        if e_m.sum() == 0 or u_m.sum() == 0:
            continue
        w = e_m.sum()
        num += y_t[e_m].sum()
        cf += w * y_t[u_m].mean()
        w_total += w
    spsm_lift = (num / w_total) / (cf / w_total) - 1

    return {
        "targeting": targeting,
        "share_exposed": share_exposed,
        "exposed_rate": exposed_rate,
        "unexposed_rate": y_t[~exposed].mean(),
        "rct_lift": rct_lift,
        "naive_lift": naive_lift,
        "spsm_lift": spsm_lift,
    }


def main():
    rng = np.random.default_rng(SEED)
    rows = [simulate(t, rng) for t in TARGETING]
    print(f"True lift on the exposed: {TRUE_LIFT:.0%}; observable captures {OBS_SHARE:.0%} of propensity\n")
    print(f"{'targeting':>9}  {'exposed':>8}  {'conv exp':>8}  {'conv unexp':>10}  {'RCT':>7}  {'naive':>7}  {'SPSM':>7}")
    for r in rows:
        print(f"{r['targeting']:>9.1f}  {r['share_exposed']:>8.0%}  {r['exposed_rate']:>8.2%}  {r['unexposed_rate']:>10.2%}  "
              f"{r['rct_lift']:>7.0%}  {r['naive_lift']:>7.0%}  {r['spsm_lift']:>7.0%}")
    print("\nThe RCT stays near the truth (up to sampling noise) whatever the targeting.")
    print("The naive comparison grows with targeting; matching on a partial observable removes only part of the bias.")
    try:
        import pandas as pd
        pd.DataFrame(rows).to_csv("selection-bias-results.csv", index=False)
        print("Wrote selection-bias-results.csv")
    except ImportError:
        pass


if __name__ == "__main__":
    main()