jordancheney89/causality
0
1"""2causal_utils.py3---------------4All causal estimation logic for the Hillstrom dashboard.5Results are cached to .cache/results.pkl on first run and loaded on subsequent runs.6"""7 8import os9import pickle10import warnings11import numpy as np12import pandas as pd13 14# Narrow suppression: PyMC / ArviZ emit deprecation chatter on import. Keep15# everything else visible.16warnings.filterwarnings("ignore", category=DeprecationWarning)17 18CACHE_DIR = ".cache"19CACHE_FILE = os.path.join(CACHE_DIR, "results.pkl")20CACHE_SCHEMA_VERSION = 321RANDOM_SEED = 1022 23# Set to False to force a full recompute (and overwrite the pickle) the next24# time the app starts - use after any change to this file's estimation logic.25# Leave True for production/Plotly Cloud so the pre-computed pickle is loaded26# instantly instead of re-running ~3-5 min of PSM/Bayesian/uplift fits.27USE_CACHE = True28 29def smd(a, b):30 """Standardised mean difference with pooled SD. nan if either group has <2 obs."""31 if len(a) < 2 or len(b) < 2:32 return float("nan")33 pooled_std = np.sqrt((np.var(a, ddof=1) + np.var(b, ddof=1)) / 2)34 if pooled_std == 0:35 return 0.036 return (np.mean(a) - np.mean(b)) / pooled_std37 38# ---------------------------------------------------------------------------39# Data Loading40# ---------------------------------------------------------------------------41 42 43def load_data():44 """Load and preprocess the Hillstrom dataset."""45 from sklift.datasets import fetch_hillstrom46 47 bunch = fetch_hillstrom(target_col="all")48 # bunch["data"] = features, bunch["target"] = visit/conversion/spend, bunch["treatment"] = segment49 df = bunch["data"].copy()50 df["segment"] = bunch["treatment"].values51 target_df = bunch["target"]52 for col in target_df.columns:53 df[col] = target_df[col].values54 55 # One-hot encode categoricals (drop_first=False so all levels are explicit,56 # we drop one reference level manually to avoid perfect multicollinearity)57 # zip_code: reference = Urban58 df["zip_suburban"] = (df["zip_code"] == "Surburban").astype(int)59 df["zip_rural"] = (df["zip_code"] == "Rural").astype(int)60 # channel: reference = Phone61 df["channel_web"] = (df["channel"] == "Web").astype(int)62 df["channel_multichannel"] = (df["channel"] == "Multichannel").astype(int)63 64 # Keep ordinal encodings too, used only for OLS interaction display65 df["zip_code_enc"] = df["zip_code"].map({"Urban": 0, "Surburban": 1, "Rural": 2})66 df["channel_enc"] = df["channel"].map({"Phone": 0, "Web": 1, "Multichannel": 2})67 68 # Binary treatment indicators69 df["is_mens"] = (df["segment"] == "Mens E-Mail").astype(int)70 df["is_womens"] = (df["segment"] == "Womens E-Mail").astype(int)71 df["is_control"] = (df["segment"] == "No E-Mail").astype(int)72 73 return df74 75 76# ---------------------------------------------------------------------------77# Propensity Score Matching (PSM)78# ---------------------------------------------------------------------------79 80COVARIATES = [81 "recency",82 "history",83 "mens",84 "womens",85 "zip_suburban",86 "zip_rural",87 "channel_web",88 "channel_multichannel",89 "newbie"90]91 92 93def _logit_propensity(propensity):94 eps = np.finfo(np.float64).eps * 295 p = np.asarray(propensity, dtype=float)96 p = np.clip(p, eps, 1.0 - eps)97 return np.log(p / (1.0 - p))98 99 100def _fit_propensity_and_match(sub, arm_col, seed, caliper_sds=0.2):101 """102 Fit logistic propensity scores, run 1:1 nearest neighbours in **logit(PS)**103 space with a ``caliper_sds``-scaled logit pooled-SD band: pairs exceeding the104 caliper on Euclidean logit-distance are discarded.105 """106 from sklearn.linear_model import LogisticRegression107 from sklearn.preprocessing import StandardScaler108 from sklearn.neighbors import NearestNeighbors109 110 X = sub[COVARIATES].values111 y = sub[arm_col].values112 113 scaler = StandardScaler()114 X_scaled = scaler.fit_transform(X)115 116 lr = LogisticRegression(max_iter=1000, random_state=seed)117 lr.fit(X_scaled, y)118 propensity = lr.predict_proba(X_scaled)[:, 1]119 120 sub = sub.copy()121 sub["propensity"] = propensity122 sub["logit_ps"] = _logit_propensity(propensity)123 124 pooled_sd = float(np.std(sub["logit_ps"].values, ddof=1))125 if pooled_sd <= 0 or not np.isfinite(pooled_sd):126 pooled_sd = 1e-6127 caliper_width = float(caliper_sds * pooled_sd)128 129 treated_all = sub[sub[arm_col] == 1].reset_index(drop=True)130 control = sub[sub[arm_col] == 0].reset_index(drop=True)131 meta_base = {"caliper_width": caliper_width, "pooled_logit_sd": pooled_sd}132 133 if len(treated_all) == 0 or len(control) == 0:134 meta_base["n_dropped_no_caliper"] = len(treated_all)135 return (136 treated_all,137 control,138 treated_all.iloc[[]],139 control.iloc[[]],140 np.array([], dtype=float),141 meta_base,142 )143 144 # Nearest neighbours in logit(PS) space so ordering matches caliper thresholds.145 ctl_l = control[["logit_ps"]].values.reshape(-1, 1)146 trt_l = treated_all[["logit_ps"]].values.reshape(-1, 1)147 nn = NearestNeighbors(n_neighbors=1, metric="euclidean")148 nn.fit(ctl_l)149 dist_nn, indices_ps = nn.kneighbors(trt_l)150 151 match_idx = indices_ps.flatten().astype(np.int64)152 dist_logit = dist_nn.flatten()153 matched_ok = dist_logit <= caliper_width154 155 if not matched_ok.any():156 meta_base["n_dropped_no_caliper"] = int(len(treated_all))157 return (158 treated_all,159 control,160 treated_all.iloc[[]],161 control.iloc[[]],162 np.array([], dtype=float),163 meta_base,164 )165 166 matched_control = control.iloc[match_idx[matched_ok]].reset_index(drop=True)167 distances_logit_kept = dist_logit[matched_ok]168 treated_matched = treated_all.loc[matched_ok].reset_index(drop=True)169 170 meta_base["n_dropped_no_caliper"] = int((~matched_ok).sum())171 return (172 treated_all,173 control,174 treated_matched,175 matched_control,176 distances_logit_kept,177 meta_base,178 )179 180 181def _compute_psm_for_arm(df, arm):182 """183 Run PSM for one arm vs control.184 185 ATT uncertainty is approximated via a full re-fit/re-match bootstrap:186 each replicate resamples the combined treated+control pool, re-fits the187 propensity score, re-matches, and recomputes ATT. We report this as a188 practical sensitivity interval rather than an exact finite-sample CI for189 nearest-neighbour matching.190 191 arm: "mens" | "womens"192 Returns dict with all PSM artefacts.193 """194 arm_col = f"is_{arm}"195 mask = (df[arm_col] == 1) | (df["is_control"] == 1)196 sub = df[mask].copy().reset_index(drop=True)197 198 (199 treated_all,200 control,201 treated_matched,202 matched_control,203 distances_logit,204 match_meta,205 ) = _fit_propensity_and_match(sub, arm_col, RANDOM_SEED)206 207 n_treated_total = len(treated_all)208 n_matched = len(treated_matched)209 if n_matched > 0:210 matched_diffs = (211 treated_matched["spend"].values - matched_control["spend"].values212 )213 att_point = float(np.mean(matched_diffs))214 avg_logit_ps_distance = float(np.mean(distances_logit))215 # Simple matched-pair analytical SE for the retained pairs. This avoids216 # presenting the nonparametric bootstrap as an exact NN-matching CI, but217 # it is still pedagogical rather than a full nearest-neighbour matching218 # variance estimator: control re-use and matching uncertainty are not219 # modelled here.220 if n_matched > 1:221 att_se_matched = float(222 np.std(matched_diffs, ddof=1) / np.sqrt(n_matched)223 )224 else:225 att_se_matched = float("nan")226 att_ci_lo_matched = (227 att_point - 1.96 * att_se_matched228 if np.isfinite(att_se_matched)229 else float("nan")230 )231 att_ci_hi_matched = (232 att_point + 1.96 * att_se_matched233 if np.isfinite(att_se_matched)234 else float("nan")235 )236 else:237 att_point = float("nan")238 avg_logit_ps_distance = float("nan")239 att_se_matched = float("nan")240 att_ci_lo_matched = float("nan")241 att_ci_hi_matched = float("nan")242 243 # Common support on point-estimate propensities (full treated pool)244 cs_lower = float(245 max(treated_all["propensity"].min(), control["propensity"].min())246 )247 cs_upper = float(248 min(treated_all["propensity"].max(), control["propensity"].max())249 )250 n_outside_support = int(251 ((treated_all["propensity"] < cs_lower)252 | (treated_all["propensity"] > cs_upper)).sum()253 )254 255 smd_before = {}256 smd_after = {}257 smd_after_caliper_match = {}258 for cov in COVARIATES:259 smd_before[cov] = smd(260 treated_all[cov].values, control[cov].values261 )262 if n_matched > 0:263 s = smd(264 treated_matched[cov].values, matched_control[cov].values265 )266 smd_after[cov] = s267 smd_after_caliper_match[cov] = s268 else:269 smd_after[cov] = float("nan")270 smd_after_caliper_match[cov] = float("nan")271 272 # Causal bootstrap: stratified resample (treated and control drawn separately273 # with replacement at their observed sizes) so each replicate preserves the274 # treated/control ratio. An unstratified pooled resample lets the ratio275 # drift, which inflates the CI for reasons unrelated to matching uncertainty.276 # 200 reps gives stable 95% percentile CIs at reasonable compute cost277 # (~2-4 min/arm).278 rng = np.random.default_rng(RANDOM_SEED)279 n_boot = 200280 min_matched_boot = 15281 treated_pool = sub[sub[arm_col] == 1]282 control_pool = sub[sub[arm_col] == 0]283 n_treated_pool = len(treated_pool)284 n_control_pool = len(control_pool)285 att_boot = []286 if n_treated_pool >= 10 and n_control_pool >= 10:287 for b in range(n_boot):288 t_idx = rng.integers(0, n_treated_pool, size=n_treated_pool)289 c_idx = rng.integers(0, n_control_pool, size=n_control_pool)290 boot_sub = pd.concat(291 [treated_pool.iloc[t_idx], control_pool.iloc[c_idx]],292 ignore_index=True,293 )294 try:295 _, _, t_b_matched, mc_b, _, _ = _fit_propensity_and_match(296 boot_sub, arm_col, RANDOM_SEED + b297 )298 if len(t_b_matched) >= min_matched_boot:299 att_boot.append(300 float(301 np.mean(t_b_matched["spend"].values)302 - np.mean(mc_b["spend"].values)303 )304 )305 except Exception:306 continue307 308 att_boot = np.asarray(att_boot, dtype=float)309 if len(att_boot) > 0:310 att_ci_lo = float(np.percentile(att_boot, 2.5))311 att_ci_hi = float(np.percentile(att_boot, 97.5))312 else:313 att_ci_lo = float("nan")314 att_ci_hi = float("nan")315 316 pct_matched = (317 100.0 * n_matched / n_treated_total if n_treated_total > 0 else 0.0318 )319 320 return {321 "arm": arm,322 "propensity_treated": treated_all["propensity"].values,323 "propensity_control": control["propensity"].values,324 "smd_before": smd_before,325 "smd_after": smd_after,326 "smd_after_caliper_match": smd_after_caliper_match,327 "att_point": att_point,328 "att_se_matched": att_se_matched,329 "att_ci_lo_matched": att_ci_lo_matched,330 "att_ci_hi_matched": att_ci_hi_matched,331 "att_ci_lo": att_ci_lo,332 "att_ci_hi": att_ci_hi,333 "n_matched": n_matched,334 "n_treated_total": n_treated_total,335 "pct_matched": pct_matched,336 "avg_logit_ps_distance": avg_logit_ps_distance,337 # Backward compatibility for old pickles338 "avg_ps_distance": avg_logit_ps_distance,339 "cs_lower": cs_lower,340 "cs_upper": cs_upper,341 "n_outside_support": n_outside_support,342 "n_boot_successful": len(att_boot),343 "caliper_width": match_meta["caliper_width"],344 "pooled_logit_sd": match_meta["pooled_logit_sd"],345 "n_dropped_no_caliper": match_meta["n_dropped_no_caliper"],346 }347 348 349def run_psm(df):350 """Run PSM for both Men's and Women's arms."""351 return {352 "mens": _compute_psm_for_arm(df, "mens"),353 "womens": _compute_psm_for_arm(df, "womens")354 }355 356 357# ---------------------------------------------------------------------------358# Bayesian A/B Test359# ---------------------------------------------------------------------------360 361ARM_PAIRS = {362 "mens_vs_control": ("Mens E-Mail", "No E-Mail"),363 "womens_vs_control": ("Womens E-Mail", "No E-Mail"),364 "mens_vs_womens": ("Mens E-Mail", "Womens E-Mail")365}366 367# Posterior-predictive draws are thinned to this many posterior samples before368# the display payload is built, so the pickle stays small.369_PPC_N_DRAWS = 400370 371 372def _native_ppc_pack(draws_a, draws_b, cap=8000, seed=RANDOM_SEED):373 """374 Build the small PPC display payload from native posterior-predictive draws.375 376 ``draws_a`` / ``draws_b`` are the simulated spend arrays for the two hurdle377 likelihoods, each shaped (n_draws, n_customers). For each arm we keep:378 379 * per-draw conversion rate (mean of spend > 0 at the arm's real n, so the380 dispersion matches binomial Monte Carlo variance at empirical n),381 * a capped sample of the full simulated spend (zeros + positive tail),382 * a capped sample of the positive amounts.383 384 Only these summaries are returned, so the pickle stays small even though the385 simulated arrays span every customer.386 """387 rng = np.random.default_rng(seed)388 out = {}389 390 def _sample(arr, n_cap):391 if len(arr) <= n_cap:392 return arr393 return rng.choice(arr, size=n_cap, replace=False)394 395 for arm, draws in (("a", draws_a), ("b", draws_b)):396 draws = np.asarray(draws).reshape(-1, np.asarray(draws).shape[-1])397 flat = draws.ravel()398 pos = flat[flat > 0.0]399 out[f"ppc_conv_rep_mean_{arm}"] = (draws > 0.0).mean(axis=1)400 out[f"ppc_spend_display_{arm}"] = _sample(flat, cap)401 out[f"ppc_amount_pos_{arm}"] = _sample(pos, cap)402 return out403 404 405def _run_bayesian_pair(df, pair_key):406 """407 Fit a two-part (hurdle) PyMC model comparing spend for one arm pair.408 409 Spend is ~99% zeros with a right-skewed positive tail. A plain Normal410 likelihood is a severe misspecification (it assigns mass to negative411 spend and the sigma term is uninterpretable). Each arm is modelled with a412 native hurdle:413 414 spend ~ HurdleLogNormal(psi, mu, sigma)415 416 where psi = P(spend > 0) and the positive amounts follow LogNormal(mu,417 sigma). The per-customer expected spend is E[spend] = psi * exp(mu +418 sigma**2/2), and `delta` is the difference between the two arms' expected419 spend. This is the target of interest for the A/B comparison and is on the420 same dollar scale as the raw arm means, so it is directly comparable to421 PSM's ATT and OLS's main effect.422 423 Runs on the full arm data (no subsampling) via the nutpie NUTS424 sampler, should be quick.425 """426 import pymc as pm427 import arviz as az428 429 arm_a_label, arm_b_label = ARM_PAIRS[pair_key]430 431 a_spend = df[df["segment"] == arm_a_label]["spend"].values.astype(float)432 b_spend = df[df["segment"] == arm_b_label]["spend"].values.astype(float)433 434 a_converted = (a_spend > 0).astype(int)435 b_converted = (b_spend > 0).astype(int)436 437 a_pos = a_spend[a_spend > 0]438 b_pos = b_spend[b_spend > 0]439 440 # Weakly informative prior on log-spend, centred on pooled log mean441 log_pooled = np.log(np.concatenate([a_pos, b_pos]))442 log_mu_prior = float(np.mean(log_pooled))443 log_sd_prior = float(np.std(log_pooled))444 445 with pm.Model():446 # psi = P(spend > 0). Beta(1,1) = uniform prior on the hurdle.447 psi_a = pm.Beta("psi_a", alpha=1.0, beta=1.0)448 psi_b = pm.Beta("psi_b", alpha=1.0, beta=1.0)449 450 # Log-amount among converters (weakly informative, data-derived).451 mu_log_a = pm.Normal("mu_log_a", mu=log_mu_prior, sigma=log_sd_prior * 2)452 mu_log_b = pm.Normal("mu_log_b", mu=log_mu_prior, sigma=log_sd_prior * 2)453 sigma_log_a = pm.HalfNormal("sigma_log_a", sigma=log_sd_prior)454 sigma_log_b = pm.HalfNormal("sigma_log_b", sigma=log_sd_prior)455 456 # Native hurdle on the full spend vector (zero spike + positive tail).457 pm.HurdleLogNormal(458 "obs_a", psi=psi_a, mu=mu_log_a, sigma=sigma_log_a, observed=a_spend459 )460 pm.HurdleLogNormal(461 "obs_b", psi=psi_b, mu=mu_log_b, sigma=sigma_log_b, observed=b_spend462 )463 464 # Expected per-customer spend: E[spend] = psi * E[amount | spend > 0]465 exp_spend_a = pm.Deterministic(466 "exp_spend_a",467 psi_a * pm.math.exp(mu_log_a + 0.5 * sigma_log_a**2),468 )469 exp_spend_b = pm.Deterministic(470 "exp_spend_b",471 psi_b * pm.math.exp(mu_log_b + 0.5 * sigma_log_b**2),472 )473 pm.Deterministic("delta", exp_spend_a - exp_spend_b)474 475 idata = pm.sample(476 draws=2000,477 tune=1000,478 chains=2,479 nuts_sampler="nutpie",480 random_seed=RANDOM_SEED,481 progressbar=True,482 return_inferencedata=True,483 )484 485 # Native posterior predictive, thinned so the display payload stays small.486 # The two hurdle likelihoods have different arm sizes, so they are487 # sampled one at a time: a single combined call makes PyTensor try to488 # inline-rewrite both diracdelta components together and fail on the489 # shape mismatch.490 ppc_pack = None491 try:492 n_chains = int(idata.posterior.sizes["chain"])493 n_draws = int(idata.posterior.sizes["draw"])494 per_chain_keep = max(1, _PPC_N_DRAWS // max(1, n_chains))495 step = max(1, n_draws // per_chain_keep)496 idata_thin = idata.sel(draw=slice(None, None, step))497 498 def _ppc_draws(var):499 pp = pm.sample_posterior_predictive(500 idata_thin,501 var_names=[var],502 random_seed=RANDOM_SEED + 11,503 progressbar=False,504 )505 return np.asarray(pp.posterior_predictive[var].values)506 507 ppc_pack = _native_ppc_pack(508 _ppc_draws("obs_a"), _ppc_draws("obs_b"), seed=RANDOM_SEED + 11509 )510 except Exception:511 ppc_pack = None512 513 delta_samples = idata.posterior["delta"].values.flatten()514 hdi = az.hdi(idata, var_names=["delta"], prob=0.95)["delta"].values515 516 report_vars = ["delta", "exp_spend_a", "exp_spend_b", "psi_a", "psi_b",517 "mu_log_a", "mu_log_b", "sigma_log_a", "sigma_log_b"]518 diagnostics = az.summary(idata, var_names=report_vars, round_to=3)519 rhat_delta = float(diagnostics.loc["delta", "r_hat"])520 bulk_ess_delta = float(diagnostics.loc["delta", "ess_bulk"])521 tail_ess_delta = float(diagnostics.loc["delta", "ess_tail"])522 523 diag_rows = []524 for var in report_vars:525 if var in diagnostics.index:526 diag_rows.append(527 {528 "parameter": var,529 "r_hat": diagnostics.loc[var, "r_hat"],530 "ess_bulk": diagnostics.loc[var, "ess_bulk"],531 "ess_tail": diagnostics.loc[var, "ess_tail"],532 "mean": diagnostics.loc[var, "mean"],533 "sd": diagnostics.loc[var, "sd"],534 }535 )536 537 delta_chains = idata.posterior["delta"].values # (chains, draws)538 539 return {540 "pair_key": pair_key,541 "arm_a_label": arm_a_label,542 "arm_b_label": arm_b_label,543 "delta_samples": delta_samples,544 "delta_chains": delta_chains,545 "exp_spend_a_samples": idata.posterior["exp_spend_a"].values.flatten(),546 "exp_spend_b_samples": idata.posterior["exp_spend_b"].values.flatten(),547 "hdi_lo": float(hdi[0]),548 "hdi_hi": float(hdi[1]),549 "p_positive": float(np.mean(delta_samples > 0)),550 "mean_a": float(np.mean(idata.posterior["exp_spend_a"].values)),551 "mean_b": float(np.mean(idata.posterior["exp_spend_b"].values)),552 "rhat_delta": rhat_delta,553 "bulk_ess_delta": bulk_ess_delta,554 "tail_ess_delta": tail_ess_delta,555 "diagnostics_table": diag_rows,556 "observed_spend_a": a_spend,557 "observed_spend_b": b_spend,558 "observed_amount_a": a_pos,559 "observed_amount_b": b_pos,560 "obs_conv_rate_a": float(np.mean(a_converted)),561 "obs_conv_rate_b": float(np.mean(b_converted)),562 "ppc_pack": ppc_pack,563 "ppc_amount_a": None564 if ppc_pack is None565 else ppc_pack.get("ppc_amount_pos_a"),566 "ppc_amount_b": None567 if ppc_pack is None568 else ppc_pack.get("ppc_amount_pos_b"),569 }570 571 572def run_bayesian_ab(df):573 """Run Bayesian A/B for all three arm pairs."""574 results = {}575 for key in ARM_PAIRS:576 print(f" [Bayesian] Fitting {key}...")577 results[key] = _run_bayesian_pair(df, key)578 return results579 580 581# ---------------------------------------------------------------------------582# Uplift Modelling / HTE583# ---------------------------------------------------------------------------584 585 586# Shuffles for the AUUC permutation test. The smallest reportable p-value is587# 1 / (AUUC_N_PERM + 1); display code uses this to report floor values as588# "p < ..." rather than an exact p.589AUUC_N_PERM = 500590 591 592def _permutation_p_auuc(593 sub_sorted, arm_col, n_perm=AUUC_N_PERM, seed=RANDOM_SEED, scale_to_audience=True594):595 """596 Permutation p-value for Qini AUUC > 0 with fixed CATE ranking.597 598 H0: the predicted CATE ranking carries no information about the treatment599 response. Under H0, treatment labels are exchangeable across the ranked600 list, so we shuffle them, recompute AUUC and check how often the601 permuted AUUC reaches the observed value. This is a refit-free602 permutation. The model and its ranking are held fixed, only the603 treatment labels are reshuffled. It tests whether the *ranking* picks604 out responders, conditional on the model that produced it.605 606 When `scale_to_audience=True`, AUUC is computed on the same audience-scaled607 curve shown in the dashboard and used by the policy calculator.608 609 Returns (observed_auuc, p_value, null_distribution).610 """611 rng = np.random.default_rng(seed)612 treatment = sub_sorted[arm_col].values.astype(np.float64)613 y = sub_sorted["spend"].values.astype(np.float64)614 n = len(sub_sorted)615 if n == 0:616 return 0.0, 1.0, np.array([], dtype=np.float64)617 xs_full = np.arange(1, n + 1, dtype=np.float64) / n618 619 def _auuc_for_treatment(t_vec):620 cum_t = np.cumsum(y * t_vec)621 cum_c = np.cumsum(y * (1.0 - t_vec))622 n_t = np.cumsum(t_vec)623 n_c = np.cumsum(1.0 - t_vec)624 valid = (n_t > 0) & (n_c > 0)625 if not valid.any():626 return 0.0627 ratio = np.zeros_like(n_t)628 ratio[valid] = n_t[valid] / n_c[valid]629 ys = cum_t - cum_c * ratio630 if scale_to_audience:631 rank_counts = np.arange(1, n + 1, dtype=np.float64)632 ys[valid] = ys[valid] * (rank_counts[valid] / n_t[valid])633 xs_v = xs_full[valid]634 ys_v = ys[valid]635 if len(xs_v) < 2:636 return 0.0637 return float(np.trapezoid(ys_v, xs_v))638 639 obs_auuc = _auuc_for_treatment(treatment)640 null_aucs = np.empty(n_perm, dtype=np.float64)641 for b in range(n_perm):642 null_aucs[b] = _auuc_for_treatment(rng.permutation(treatment))643 644 # One-sided p-value: how often does a random ranking match or beat the645 # observed AUUC? Add +1 numerator/denominator (Phipson & Smyth 2010646 # correction) so the p-value can't be exactly zero.647 p_value = float((np.sum(null_aucs >= obs_auuc) + 1) / (n_perm + 1))648 return obs_auuc, p_value, null_aucs649 650 651def _qini_curve_continuous(sorted_df, arm_col, scale_to_audience=True):652 """653 Radcliffe (2007) Qini-style curve for continuous outcomes.654 655 For a population ranked by predicted uplift (descending), the canonical656 Qini value at rank k is:657 658 Q(k) = R_T(k) - R_C(k) * (N_T(k) / N_C(k))659 660 where R_T(k) and R_C(k) are cumulative outcomes in treated/control and661 N_T(k), N_C(k) are cumulative treated/control counts. At k=N this equals662 R_T_total - R_C_total * (N_T/N_C), i.e. the total incremental revenue663 among the treated sample. For dashboard and policy use, `scale_to_audience`664 converts that to the full ranked-audience scale by multiplying by665 k / N_T(k), so revenue and send cost are expressed for the same number of666 customers. scikit-uplift's `qini_curve` enforces binary outcomes and can't667 be used on `spend` directly. This is the same formula generalised.668 """669 n_rows = len(sorted_df)670 if n_rows == 0:671 return [], []672 is_t = sorted_df[arm_col].values.astype(float)673 y_vals = sorted_df["spend"].values.astype(float)674 675 cum_t = np.cumsum(y_vals * is_t)676 cum_c = np.cumsum(y_vals * (1 - is_t))677 n_t_cum = np.cumsum(is_t)678 n_c_cum = np.cumsum(1 - is_t)679 680 valid = (n_t_cum > 0) & (n_c_cum > 0)681 rank_counts = np.arange(1, n_rows + 1)682 xs_arr = (rank_counts / n_rows)[valid]683 ys_arr = cum_t[valid] - cum_c[valid] * (n_t_cum[valid] / n_c_cum[valid])684 if scale_to_audience:685 ys_arr = ys_arr * (rank_counts[valid] / n_t_cum[valid])686 687 xs = xs_arr.tolist()688 ys = ys_arr.tolist()689 690 # Prepend origin for a clean plot691 if xs and xs[0] > 0:692 xs = [0.0] + xs693 ys = [0.0] + ys694 return xs, ys695 696 697def _make_rf():698 """699 Random forest hyperparameters tuned for ~99%-zero, right-skewed spend.700 701 The defaults (`n_estimators=100`, no depth cap, `min_samples_leaf=1`) memorise702 individual converters because positive spend is rare. We cap depth, force703 leaves to span many customers, and subsample features per split so trees704 decorrelate. This trades some in-sample fit for honest CATE generalisation.705 """706 from sklearn.ensemble import RandomForestRegressor707 708 return RandomForestRegressor(709 n_estimators=200,710 max_depth=8,711 min_samples_leaf=50,712 max_features=0.5,713 random_state=RANDOM_SEED,714 n_jobs=-1,715 )716 717 718def _fit_x_learner(X_train_df, y_train, t_train):719 """720 Fit X-Learner fold models (Künzel et al. 2019).721 722 Stage 1: fit outcome models μ̂₀, μ̂₁ on control and treated arms.723 Stage 2: impute counterfactual treatment effects on training data:724 D̃¹_i = Y_i - μ̂₀(X_i) for treated i725 D̃⁰_i = μ̂₁(X_i) - Y_i for control i726 Stage 3: fit τ̂₁, τ̂₀ regressing the imputed effects on covariates,727 then retain e for the propensity-weighted prediction step.728 """729 treated = t_train == 1730 control = t_train == 0731 732 mu0 = _make_rf()733 mu0.fit(X_train_df.loc[control], y_train[control])734 mu1 = _make_rf()735 mu1.fit(X_train_df.loc[treated], y_train[treated])736 737 d1 = y_train[treated] - mu0.predict(X_train_df.loc[treated])738 d0 = mu1.predict(X_train_df.loc[control]) - y_train[control]739 740 tau1 = _make_rf()741 tau1.fit(X_train_df.loc[treated], d1)742 tau0 = _make_rf()743 tau0.fit(X_train_df.loc[control], d0)744 745 e = float(treated.mean())746 return {"tau0": tau0, "tau1": tau1, "e": e}747 748 749def _predict_x_learner(fit, X_test_df):750 """751 Predict X-Learner CATE from a fitted fold bundle.752 753 τ̂(x) = e * τ̂₀(x) + (1 - e) * τ̂₁(x), where e is the observed treated754 share in the fold. The minority arm's τ estimate is up-weighted, which is755 why X-Learner is useful when treatment/control sizes are uneven.756 """757 return fit["e"] * fit["tau0"].predict(X_test_df) + (1.0 - fit["e"]) * fit[758 "tau1"759 ].predict(X_test_df)760 761 762def _solo_design(X, t):763 """764 Treatment-interaction design for the S-Learner: ``[X, t, t * X]``.765 766 A single learner on this design lets the treatment effect vary with the767 covariates (the interaction block), matching the behaviour of the previous768 scikit-uplift ``SoloModel(method="treatment_interaction")``.769 """770 X = np.asarray(X, dtype=float)771 t = np.asarray(t, dtype=float).reshape(-1, 1)772 return np.hstack([X, t, X * t])773 774 775def _fit_s_learner(X_train_df, y_train, t_train):776 """S-Learner: one model on the treatment-interaction design."""777 model = _make_rf()778 model.fit(_solo_design(X_train_df.values, t_train), y_train)779 return {"model": model}780 781 782def _predict_s_learner(fit, X_test_df):783 """CATE = f(x, t=1) - f(x, t=0) on the interaction design."""784 X = X_test_df.values785 n = len(X)786 pred_t = fit["model"].predict(_solo_design(X, np.ones(n)))787 pred_c = fit["model"].predict(_solo_design(X, np.zeros(n)))788 return pred_t - pred_c789 790 791def _avg_cate_ci(cate, n_boot=1000, seed=RANDOM_SEED):792 """793 Percentile bootstrap CI for the average CATE.794 795 Resamples customers with replacement and recomputes the mean of the796 cross-fit CATE. Like the decile-lift CIs, this captures the sampling797 variability of the average given the fitted CATE surface.798 """799 cate = np.asarray(cate, dtype=float)800 n = len(cate)801 if n == 0:802 return float("nan"), float("nan")803 rng = np.random.default_rng(seed)804 means = np.array(805 [cate[rng.integers(0, n, size=n)].mean() for _ in range(n_boot)]806 )807 return float(np.percentile(means, 2.5)), float(np.percentile(means, 97.5))808 809 810def _run_uplift_arm(df, arm):811 """812 Run S-Learner and X-Learner for one arm vs control.813 814 Uses 5-fold *stratified* cross-fitting on the treatment indicator so every815 fold has both arms represented, which matters for small-sample fold fits.816 CATE estimates are out-of-sample for every observation.817 818 Heterogeneity importance is computed per learner by **permutation819 importance on the predicted CATE**: for each fold and feature, the column820 is shuffled in the held-out X_test, the selected learner re-predicts CATE821 on the permuted inputs, and the mean absolute change relative to the822 unpermuted prediction is recorded. Features that drive heterogeneity show823 large shifts, features irrelevant to the treatment-effect surface show824 small ones. This is a model-agnostic, bias-free alternative to RF impurity825 importance, which is biased toward high-cardinality / continuous features.826 Permutation is repeated `n_perm_repeats` times per (fold, feature) to827 reduce shuffle noise.828 """829 from sklearn.model_selection import StratifiedKFold830 831 arm_col = f"is_{arm}"832 mask = (df[arm_col] == 1) | (df["is_control"] == 1)833 sub = df[mask].copy().reset_index(drop=True)834 835 X = sub[COVARIATES].values836 y = sub["spend"].values837 treatment = sub[arm_col].values838 839 kf = StratifiedKFold(n_splits=5, shuffle=True, random_state=RANDOM_SEED)840 cate_s = np.zeros(len(sub))841 cate_x = np.zeros(len(sub))842 perm_imp_accum = {843 "s": np.zeros(len(COVARIATES)),844 "x": np.zeros(len(COVARIATES)),845 }846 n_perm_repeats = 5847 perm_rng = np.random.default_rng(RANDOM_SEED)848 849 for _, (train_idx, test_idx) in enumerate(kf.split(X, treatment)):850 X_train, X_test = X[train_idx], X[test_idx]851 y_train = y[train_idx]852 t_train = treatment[train_idx]853 854 X_train_df = pd.DataFrame(X_train, columns=COVARIATES)855 X_test_df = pd.DataFrame(X_test, columns=COVARIATES)856 857 def _accumulate_importance(model_key, predict_fn, baseline_pred):858 for j, _col in enumerate(COVARIATES):859 shifts = np.zeros(n_perm_repeats)860 for r in range(n_perm_repeats):861 X_perm = X_test_df.copy()862 X_perm.iloc[:, j] = perm_rng.permutation(X_perm.iloc[:, j].values)863 cate_perm = predict_fn(X_perm)864 shifts[r] = float(np.mean(np.abs(cate_perm - baseline_pred)))865 perm_imp_accum[model_key][j] += shifts.mean()866 867 # S-Learner868 s_fit = _fit_s_learner(X_train_df, y_train, t_train)869 cate_s_fold = _predict_s_learner(s_fit, X_test_df)870 cate_s[test_idx] = cate_s_fold871 _accumulate_importance(872 "s", lambda X_perm: _predict_s_learner(s_fit, X_perm), cate_s_fold873 )874 875 # X-Learner876 x_fit = _fit_x_learner(X_train_df, y_train, t_train)877 cate_x_fold = _predict_x_learner(x_fit, X_test_df)878 cate_x[test_idx] = cate_x_fold879 _accumulate_importance(880 "x", lambda X_perm: _predict_x_learner(x_fit, X_perm), cate_x_fold881 )882 883 feat_imp_diff = {884 key: vals / kf.get_n_splits() for key, vals in perm_imp_accum.items()885 }886 887 def _normalise_importance(vals):888 return vals / vals.sum() if vals.sum() > 0 else vals889 890 feat_imp_norm = {891 key: _normalise_importance(vals) for key, vals in feat_imp_diff.items()892 }893 894 sub["cate_s"] = cate_s895 sub["cate_x"] = cate_x896 897 def _decile_lift(sub_sorted, n_boot=500):898 """899 Actual spend lift per decile for a population sorted by predicted CATE,900 with stratified bootstrap 95% CIs on the difference-in-means within901 each decile. Within-decile treated and control are resampled separately902 with replacement so the CI reflects sampling uncertainty in the lift903 estimate, not the decile boundaries.904 """905 sub_sorted = sub_sorted.reset_index(drop=True)906 sub_sorted["decile"] = pd.qcut(sub_sorted.index, q=10, labels=False)907 rng = np.random.default_rng(RANDOM_SEED)908 rows = []909 for d in range(10):910 dec = sub_sorted[sub_sorted["decile"] == d]911 t_vals = dec[dec[arm_col] == 1]["spend"].values912 c_vals = dec[dec[arm_col] == 0]["spend"].values913 if len(t_vals) == 0 or len(c_vals) == 0:914 rows.append({915 "decile": d + 1, "lift": 0.0,916 "ci_lo": float("nan"), "ci_hi": float("nan"),917 })918 continue919 lift = float(t_vals.mean() - c_vals.mean())920 boot = np.empty(n_boot, dtype=float)921 for b in range(n_boot):922 t_b = rng.choice(t_vals, size=len(t_vals), replace=True).mean()923 c_b = rng.choice(c_vals, size=len(c_vals), replace=True).mean()924 boot[b] = t_b - c_b925 rows.append({926 "decile": d + 1,927 "lift": lift,928 "ci_lo": float(np.percentile(boot, 2.5)),929 "ci_hi": float(np.percentile(boot, 97.5)),930 })931 return rows932 933 sub_sorted_s = sub.sort_values("cate_s", ascending=False)934 sub_sorted_x = sub.sort_values("cate_x", ascending=False)935 936 decile_lift_s = _decile_lift(sub_sorted_s)937 decile_lift_x = _decile_lift(sub_sorted_x)938 939 qini_x_s, qini_y_s = _qini_curve_continuous(sub_sorted_s, arm_col)940 qini_x_x, qini_y_x = _qini_curve_continuous(sub_sorted_x, arm_col)941 942 # Bootstrap CIs on the average CATE (estimation-conditional lower bound).943 avg_cate_s_lo, avg_cate_s_hi = _avg_cate_ci(cate_s)944 avg_cate_x_lo, avg_cate_x_hi = _avg_cate_ci(cate_x)945 946 # Permutation p-values for AUUC. 500 shuffles per method, cheap because947 # we hold the predicted ranking fixed and only relabel treatment.948 _, qini_p_s, _ = _permutation_p_auuc(sub_sorted_s, arm_col)949 _, qini_p_x, _ = _permutation_p_auuc(sub_sorted_x, arm_col)950 951 # AUUC-like area under the cumulative incremental gain curve. With952 # subtract_baseline=True, reports only the excess area above the953 # random-targeting baseline (the line from (0, 0) to (1, final gain)).954 def _qini_auc(xs, ys, subtract_baseline=False):955 if len(xs) < 2:956 return 0.0957 xs_arr = np.asarray(xs, dtype=float)958 ys_arr = np.asarray(ys, dtype=float)959 if subtract_baseline:960 ys_arr = ys_arr - xs_arr * ys_arr[-1]961 return float(np.trapezoid(ys_arr, xs_arr))962 963 return {964 "arm": arm,965 "cate_s": cate_s,966 "cate_x": cate_x,967 "feat_imp_s": dict(zip(COVARIATES, feat_imp_norm["s"])),968 "feat_imp_x": dict(zip(COVARIATES, feat_imp_norm["x"])),969 "feat_imp_label_s": "Heterogeneity importance (S-Learner CATE permutation)",970 "feat_imp_label_x": "Heterogeneity importance (X-Learner CATE permutation)",971 "decile_lift_s": decile_lift_s,972 "decile_lift_x": decile_lift_x,973 "qini_x_s": qini_x_s,974 "qini_y_s": qini_y_s,975 "qini_x_x": qini_x_x,976 "qini_y_x": qini_y_x,977 "qini_auc_s": _qini_auc(qini_x_s, qini_y_s),978 "qini_auc_x": _qini_auc(qini_x_x, qini_y_x),979 "qini_excess_auc_s": _qini_auc(qini_x_s, qini_y_s, subtract_baseline=True),980 "qini_excess_auc_x": _qini_auc(qini_x_x, qini_y_x, subtract_baseline=True),981 "qini_p_s": qini_p_s,982 "qini_p_x": qini_p_x,983 "avg_cate_s": float(np.mean(cate_s)),984 "avg_cate_x": float(np.mean(cate_x)),985 "avg_cate_s_lo": avg_cate_s_lo,986 "avg_cate_s_hi": avg_cate_s_hi,987 "avg_cate_x_lo": avg_cate_x_lo,988 "avg_cate_x_hi": avg_cate_x_hi,989 }990 991 992def run_uplift(df):993 """Run uplift modelling for both arms."""994 return {995 "mens": _run_uplift_arm(df, "mens"),996 "womens": _run_uplift_arm(df, "womens")997 }998 999 1000# ---------------------------------------------------------------------------1001# Multi-Arm OLS1002# ---------------------------------------------------------------------------1003 1004# Covariates that interact with each arm dummy in the OLS model, i.e. the terms1005# that shift an arm's marginal effect away from the reference subgroup.1006OLS_INTERACTION_COVARIATES = [1007 "newbie",1008 "channel_web",1009 "channel_multichannel",1010 "zip_suburban",1011 "zip_rural",1012]1013 1014 1015def _subgroup_marginal_effect(params, arm_prefix, subgroup_vals):1016 """Marginal effect of an arm for one covariate subgroup: main effect plus1017 each interaction coefficient weighted by that subgroup's covariate value."""1018 me = params.get(arm_prefix, 0)1019 for cov in OLS_INTERACTION_COVARIATES:1020 me += params.get(f"{arm_prefix}:{cov}", 0) * subgroup_vals[cov]1021 return me1022 1023 1024def run_ols(df):1025 """1026 OLS regression with treatment dummies, covariates, and interaction terms.1027 Outcome: spend. Categoricals (zip_code, channel) are one-hot encoded.1028 Reference levels are Urban and Phone respectively.1029 """1030 import statsmodels.formula.api as smf1031 1032 model_df = df.copy()1033 model_df["mens_email"] = (model_df["segment"] == "Mens E-Mail").astype(int)1034 model_df["womens_email"] = (model_df["segment"] == "Womens E-Mail").astype(int)1035 1036 # Main effects + interactions with OHE categoricals1037 formula = (1038 "spend ~ mens_email + womens_email "1039 "+ recency + history + newbie "1040 "+ zip_suburban + zip_rural "1041 "+ channel_web + channel_multichannel "1042 "+ mens_email:newbie + womens_email:newbie "1043 "+ mens_email:channel_web + womens_email:channel_web "1044 "+ mens_email:channel_multichannel + womens_email:channel_multichannel "1045 "+ mens_email:zip_suburban + womens_email:zip_suburban "1046 "+ mens_email:zip_rural + womens_email:zip_rural"1047 )1048 1049 # HC3 heteroscedasticity-robust standard errors: spend is right-skewed with1050 # variance that scales with the treatment means, so default OLS SEs would1051 # be biased. HC3 is the recommended small-sample-corrected White estimator.1052 result = smf.ols(formula, data=model_df).fit(cov_type="HC3")1053 1054 # Coefficient table1055 _ci = result.conf_int()1056 coef_df = (1057 pd.DataFrame(1058 {1059 "coef": result.params,1060 "ci_lo": _ci[0],1061 "ci_hi": _ci[1],1062 "pvalue": result.pvalues1063 }1064 )1065 .reset_index()1066 .rename(columns={"index": "term"})1067 )1068 1069 # Marginal effects by subgroup1070 subgroups = []1071 for newbie_val, newbie_label in [(0, "Existing"), (1, "New")]:1072 for channel_web, channel_mc, channel_label in [1073 (0, 0, "Phone"),1074 (1, 0, "Web"),1075 (0, 1, "Multichannel")1076 ]:1077 for zip_sub, zip_rural_val, zip_label in [1078 (0, 0, "Urban"),1079 (1, 0, "Suburban"),1080 (0, 1, "Rural")1081 ]:1082 subgroup_vals = {1083 "newbie": newbie_val,1084 "channel_web": channel_web,1085 "channel_multichannel": channel_mc,1086 "zip_suburban": zip_sub,1087 "zip_rural": zip_rural_val,1088 }1089 subgroups.append(1090 {1091 "newbie": newbie_label,1092 "channel": channel_label,1093 "zip_code": zip_label,1094 "me_mens": _subgroup_marginal_effect(1095 result.params, "mens_email", subgroup_vals1096 ),1097 "me_womens": _subgroup_marginal_effect(1098 result.params, "womens_email", subgroup_vals1099 ),1100 }1101 )1102 1103 subgroup_df = pd.DataFrame(subgroups)1104 1105 # Population-weighted ATE and its HC3 CI via the delta method.1106 # The `mens_email` / `womens_email` coefficients on their own are the1107 # effect for the reference subgroup (Existing + Phone + Urban), they are1108 # not directly comparable to PSM's ATT or the Bayesian delta. The ATE1109 # below is the average of the linear-prediction marginal effects over1110 # the actual sample distribution of the covariates, which IS on the same1111 # scale as the other methods.1112 def _ate_with_ci(arm_prefix):1113 term_main = arm_prefix1114 inter_terms = {1115 f"{arm_prefix}:{cov}": float(model_df[cov].mean())1116 for cov in OLS_INTERACTION_COVARIATES1117 }1118 params = result.params1119 cov = result.cov_params()1120 contrast = np.zeros(len(params))1121 if term_main in params.index:1122 contrast[params.index.get_loc(term_main)] = 1.01123 for t, w in inter_terms.items():1124 if t in params.index:1125 contrast[params.index.get_loc(t)] = w1126 ate = float(contrast @ params.values)1127 se = float(np.sqrt(contrast @ cov.values @ contrast))1128 return ate, ate - 1.96 * se, ate + 1.96 * se1129 1130 ate_mens, ate_mens_lo, ate_mens_hi = _ate_with_ci("mens_email")1131 ate_womens, ate_womens_lo, ate_womens_hi = _ate_with_ci("womens_email")1132 1133 return {1134 "coef_df": coef_df,1135 "subgroup_df": subgroup_df,1136 "r_squared": result.rsquared,1137 "n_obs": int(result.nobs),1138 "summary_text": result.summary().as_text(),1139 "ate_mens": ate_mens,1140 "ate_mens_lo": ate_mens_lo,1141 "ate_mens_hi": ate_mens_hi,1142 "ate_womens": ate_womens,1143 "ate_womens_lo": ate_womens_lo,1144 "ate_womens_hi": ate_womens_hi,1145 }1146 1147 1148# ---------------------------------------------------------------------------1149# Cache Management1150# ---------------------------------------------------------------------------1151 1152 1153def build_cache():1154 """Compute all results and save to disk. Returns the results dict."""1155 os.makedirs(CACHE_DIR, exist_ok=True)1156 1157 print("[Cache] Loading data...")1158 df = load_data()1159 1160 print("[Cache] Running PSM (re-fit/re-match bootstrap 200 reps x 2 arms, ~4-8 min)...")1161 psm = run_psm(df)1162 1163 print("[Cache] Running Bayesian A/B (PyMC, 3 arm pairs)...")1164 bayesian = run_bayesian_ab(df)1165 1166 print("[Cache] Running Uplift models (S/X-Learners w/ 2 arms)...")1167 uplift = run_uplift(df)1168 1169 print("[Cache] Running Multi-Arm OLS...")1170 ols = run_ols(df)1171 1172 results = {1173 "schema_version": CACHE_SCHEMA_VERSION,1174 "df": df,1175 "psm": psm,1176 "bayesian": bayesian,1177 "uplift": uplift,1178 "ols": ols,1179 }1180 1181 with open(CACHE_FILE, "wb") as f:1182 pickle.dump(results, f)1183 1184 print(f"[Cache] Saved to {CACHE_FILE}")1185 return results1186 1187 1188def load_or_build_cache():1189 """Load cached results if USE_CACHE and a pickle exists, otherwise recompute."""1190 if USE_CACHE and os.path.exists(CACHE_FILE):1191 print(f"[Cache] USE_CACHE=True - loading from {CACHE_FILE}...")1192 try:1193 with open(CACHE_FILE, "rb") as f:1194 results = pickle.load(f)1195 except ModuleNotFoundError as exc:1196 print(1197 f"[Cache] Existing cache needs unavailable module ({exc.name}); "1198 "rebuilding..."1199 )1200 else: