#!/usr/bin/env python3 """A/B test planner: fixed split vs Thompson sampling vs top-two Thompson. Simulates the next ``--days`` of traffic for a control plus one or more variants and compares three ways of allocating it: fixed Classic A/B(/n) test. Equal split until a pre-registered sample size (two-sided z-test, Bonferroni over variants, given power and minimum detectable effect) is reached; then a variant is shipped only if it beats control significantly, otherwise control is kept. Clean frequentist inference, highest exploration cost. ts Batched Thompson sampling: each day traffic is split in proportion to the posterior probability that each arm is best. Minimises lost conversions; never 'stops', and naive estimates after adaptive allocation are biased (do not run a plain t-test on its data). ttts Top-two Thompson sampling (Russo, 2016), beta = 1/2: built for *identifying* the best arm quickly; stops as soon as one arm has posterior P(best) >= --confidence and ships it. All arms are Bernoulli with Beta(1,1) priors. Traffic is allocated in daily batches (that is how most teams can actually deploy changes). Results are Monte Carlo averages over ``--sims`` simulated futures. Example: python tools/ab_test_planner.py --baseline 0.05 --lifts 0.05,0.10 \ --traffic 2000 --days 28 --sims 300 --plot out/ab_plan.png If --lifts is omitted, each simulation draws the true relative lifts of --variants treatment arms uniformly from [-max_lift, +max_lift] (a crude 'we don't know' prior), so the report averages over plausible worlds. """ from __future__ import annotations import argparse import math import sys import numpy as np from scipy import stats # -------------------------------------------------------------------------- # helpers # -------------------------------------------------------------------------- def prob_best(alpha, beta, rng, draws: int = 1000): """Monte Carlo P(arm is best) under independent Beta posteriors. alpha, beta: (R, K). Returns (R, K). """ R, K = alpha.shape samples = rng.beta(alpha[:, None, :], beta[:, None, :], size=(R, draws, K)) winners = samples.argmax(axis=2) return np.stack([(winners == k).mean(axis=1) for k in range(K)], axis=1) def top_two_allocation(p, beta: float = 0.5, floor: float = 1e-4): """Allocation of top-two TS given P(best) vector(s) p. Leader I ~ p; with prob beta play I, else re-sample until a different arm J appears: P(J=i | I=j) = p_i / (1 - p_j). Hence psi_i = beta p_i + (1-beta) p_i sum_{j != i} p_j / (1 - p_j). A small floor keeps the formula defined when some p_j -> 1. """ p = np.clip(p, floor, 1 - floor) p = p / p.sum(axis=-1, keepdims=True) ratio = p / (1 - p) other = ratio.sum(axis=-1, keepdims=True) - ratio psi = beta * p + (1 - beta) * p * other return psi / psi.sum(axis=-1, keepdims=True) def fixed_sample_size(p0: float, mde_rel: float, n_variants: int, alpha: float = 0.05, power: float = 0.8) -> int: """Per-arm n for a two-sided two-proportion z-test vs control, Bonferroni-corrected over the number of treatment arms.""" p1 = p0 * (1 + mde_rel) a = alpha / max(n_variants, 1) z_a = stats.norm.ppf(1 - a / 2) z_b = stats.norm.ppf(power) var = p0 * (1 - p0) + p1 * (1 - p1) return int(math.ceil((z_a + z_b) ** 2 * var / (p1 - p0) ** 2)) def allocate(traffic: int, probs, rng): """Split a day's traffic over arms by multinomial draws (R, K).""" return np.stack([rng.multinomial(traffic, pr) for pr in probs]) # -------------------------------------------------------------------------- # strategies # -------------------------------------------------------------------------- def run_strategy(name, rates, args, rng): """Simulate one strategy on R parallel worlds. rates: (R, K) true conversion rates (arm 0 = control). Returns dict: daily lost conversions (R, days), chosen arm (R,), decision day (R,) (np.inf = never confident within horizon). """ R, K = rates.shape best = rates.max(axis=1) S = np.zeros((R, K)) N = np.zeros((R, K)) lost = np.zeros((R, args.days)) decided = np.full(R, np.inf) shipped = np.full(R, -1) # arm deployed to 100% after a decision n_per_arm = fixed_sample_size(args.baseline, args.mde, K - 1, args.alpha, args.power) test_days = math.ceil(n_per_arm * K / args.traffic) for d in range(args.days): if name == "fixed": if d < test_days: probs = np.full((R, K), 1.0 / K) else: if d == test_days: # analyse the pre-registered test once shipped = fixed_decision(S, N, args.alpha) decided[:] = d probs = np.eye(K)[shipped] else: post = prob_best(1 + S, 1 + N - S, rng, args.draws) if name == "ts": probs = post newly = np.isinf(decided) & (post.max(axis=1) >= args.confidence) decided[newly] = d else: # ttts: stop & ship once confident newly = (shipped < 0) & (post.max(axis=1) >= args.confidence) shipped[newly] = post[newly].argmax(axis=1) decided[newly] = d probs = np.where((shipped >= 0)[:, None], np.eye(K)[np.maximum(shipped, 0)], top_two_allocation(post, args.ttts_beta)) n = allocate(args.traffic, probs, rng) conv = rng.binomial(n, rates) S += conv N += n lost[:, d] = (n * (best[:, None] - rates)).sum(axis=1) # expected loss # final recommendation if name == "fixed": if test_days >= args.days: # horizon shorter than the planned test shipped = fixed_decision(S, N, args.alpha) decided[:] = args.days chosen = shipped elif name == "ts": chosen = prob_best(1 + S, 1 + N - S, rng, args.draws).argmax(axis=1) else: final = prob_best(1 + S, 1 + N - S, rng, args.draws).argmax(axis=1) chosen = np.where(shipped >= 0, shipped, final) return {"lost": lost, "chosen": chosen, "decided": decided, "test_days": test_days, "n_per_arm": n_per_arm} def fixed_decision(S, N, alpha): """Ship the best significantly-better variant (one-sided at the Bonferroni level, matching the two-sided design), else control (arm 0).""" R, K = S.shape p = S / np.maximum(N, 1) pooled = (S[:, :1] + S) / np.maximum(N[:, :1] + N, 1) se = np.sqrt(pooled * (1 - pooled) * (1 / np.maximum(N[:, :1], 1) + 1 / np.maximum(N, 1))) z = (p - p[:, :1]) / np.maximum(se, 1e-12) crit = stats.norm.ppf(1 - alpha / max(K - 1, 1) / 2) sig = z > crit sig[:, 0] = False cand = np.where(sig, p, -np.inf) return np.where(sig.any(axis=1), cand.argmax(axis=1), 0) # -------------------------------------------------------------------------- # CLI # -------------------------------------------------------------------------- def parse_args(argv=None): ap = argparse.ArgumentParser(description=__doc__.split("\n\n")[0], formatter_class=argparse.RawDescriptionHelpFormatter) ap.add_argument("--baseline", type=float, default=0.05, help="control conversion rate") ap.add_argument("--lifts", type=str, default=None, help="comma-separated true relative lifts of the variants, e.g. 0.05,0.10") ap.add_argument("--variants", type=int, default=2, help="number of treatment arms when --lifts is not given") ap.add_argument("--max-lift", type=float, default=0.15, help="lifts ~ U[-max, max] per simulation when --lifts is not given") ap.add_argument("--traffic", type=int, default=2000, help="visitors per day (all arms)") ap.add_argument("--days", type=int, default=28, help="horizon in days") ap.add_argument("--mde", type=float, default=None, help="relative minimum detectable effect for the fixed test " "(default: smallest positive lift, or max-lift/2)") ap.add_argument("--alpha", type=float, default=0.05) ap.add_argument("--power", type=float, default=0.8) ap.add_argument("--confidence", type=float, default=0.95, help="posterior P(best) needed to call a winner (TS / TTTS)") ap.add_argument("--ttts-beta", type=float, default=0.5) ap.add_argument("--sims", type=int, default=300) ap.add_argument("--draws", type=int, default=1000, help="posterior draws per day") ap.add_argument("--seed", type=int, default=0) ap.add_argument("--plot", type=str, default=None, help="save a PNG summary here") args = ap.parse_args(argv) if args.lifts is not None: args.lift_list = [float(x) for x in args.lifts.split(",") if x.strip()] else: args.lift_list = None if args.mde is None: if args.lift_list: pos = [x for x in args.lift_list if x > 0] args.mde = min(pos) if pos else 0.05 else: args.mde = args.max_lift / 2 return args def make_worlds(args, rng): if args.lift_list is not None: lifts = np.tile(np.array(args.lift_list), (args.sims, 1)) else: lifts = rng.uniform(-args.max_lift, args.max_lift, size=(args.sims, args.variants)) rates = args.baseline * (1 + np.concatenate([np.zeros((args.sims, 1)), lifts], axis=1)) return np.clip(rates, 0, 1) def plan(args): rng = np.random.default_rng(args.seed) rates = make_worlds(args, rng) best_arm = rates.argmax(axis=1) # 'correct' is ambiguous when two arms tie for best; we count any tied best is_best = np.isclose(rates, rates.max(axis=1, keepdims=True)) out = {} for name in ("fixed", "ts", "ttts"): res = run_strategy(name, rates, args, rng) res["correct"] = is_best[np.arange(len(best_arm)), res["chosen"]] out[name] = res return rates, out def report(rates, out, args): K = rates.shape[1] total = args.traffic * args.days lines = [] lines.append(f"Arms: control + {K - 1} variant(s); baseline {args.baseline:.3%}; " f"{args.traffic}/day for {args.days} days ({total:,} visitors); {args.sims} sims") if args.lift_list is not None: lines.append("True relative lifts: " + ", ".join(f"{x:+.1%}" for x in args.lift_list)) else: lines.append(f"True lifts drawn U[-{args.max_lift:.0%}, +{args.max_lift:.0%}] per simulation") f = out["fixed"] lines.append(f"Fixed test plan: n = {f['n_per_arm']:,}/arm for MDE {args.mde:+.1%} " f"(alpha {args.alpha}, power {args.power}, Bonferroni) " f"-> {f['test_days']} days" + (" [LONGER THAN HORIZON: underpowered]" if f["test_days"] > args.days else "")) lines.append("") hdr = f"{'strategy':<26}{'lost conv. (mean ± sd)':>24}{'P(pick best)':>14}{'decided by end':>16}{'median day*':>12}" lines.append(hdr) lines.append("-" * len(hdr)) labels = {"fixed": "fixed split A/B", "ts": "Thompson sampling", "ttts": "top-two TS (stop+ship)"} for name, res in out.items(): lost = res["lost"].sum(axis=1) dec = res["decided"] if name == "fixed": frac = "planned" if res["test_days"] <= args.days else "forced@end" else: frac = f"{np.isfinite(dec).mean():.0%}" if name == "fixed": med = f"{res['test_days']}" if res["test_days"] <= args.days else f">{args.days}" else: med = f"{np.median(dec[np.isfinite(dec)]):.0f}" if np.isfinite(dec).any() else "-" lines.append(f"{labels[name]:<26}{lost.mean():>14.1f} ± {lost.std():<7.1f}" f"{res['correct'].mean():>14.1%}{frac:>16}{med:>12}") lines.append("* days of data at decision (fixed: planned test length; TS/TTTS: median over decided sims)") lines.append("") lines.append("Reading this honestly:") lines.append(" * 'lost conversions' = expected conversions lost vs always showing the true best arm.") lines.append(" * fixed A/B: 'P(pick best)' counts keeping control when no variant is significant;") lines.append(" its error rates are controlled by design (alpha, power) and its estimates are unbiased.") lines.append(" * TS / TTTS 'decided' uses posterior P(best) >= %.2f. That is a Bayesian statement under a" % args.confidence) lines.append(" flat prior; under optional stopping the frequentist error rate is not guaranteed to be") lines.append(" %.0f%% - the simulated P(pick best) above is the number to trust." % (100 * (1 - args.confidence))) lines.append(" * After adaptive allocation, naive conversion-rate estimates and t-tests are biased;") lines.append(" use the bandit for *allocation*, and a held-out fixed split if you need a clean effect size.") return "\n".join(lines) def make_plot(rates, out, args, path): import matplotlib matplotlib.use("Agg") import matplotlib.pyplot as plt labels = {"fixed": "fixed split A/B", "ts": "Thompson sampling", "ttts": "top-two TS"} colors = {"fixed": "#4C72B0", "ts": "#DD8452", "ttts": "#55A868"} fig, axes = plt.subplots(1, 3, figsize=(14, 4.2)) days = np.arange(1, args.days + 1) for name, res in out.items(): cum = res["lost"].cumsum(axis=1) m = cum.mean(axis=0) lo, hi = np.percentile(cum, [10, 90], axis=0) axes[0].plot(days, m, color=colors[name], label=labels[name]) axes[0].fill_between(days, lo, hi, color=colors[name], alpha=0.15, lw=0) axes[0].set_title("Cumulative lost conversions (mean, 10-90%)") axes[0].set_xlabel("day") axes[0].legend() names = list(out) axes[1].bar([labels[n] for n in names], [out[n]["correct"].mean() for n in names], color=[colors[n] for n in names]) axes[1].set_ylim(0, 1) axes[1].set_title("P(final choice is the true best arm)") axes[1].tick_params(axis="x", labelsize=8) for name in ("ts", "ttts"): dec = out[name]["decided"] x = np.sort(dec) y = np.arange(1, len(x) + 1) / len(x) axes[2].step(np.r_[0, x[np.isfinite(x)]], np.r_[0, y[np.isfinite(x)]], where="post", color=colors[name], label=labels[name]) td = out["fixed"]["test_days"] if td <= args.days: axes[2].axvline(td, color=colors["fixed"], ls="--", label="fixed test ends") axes[2].set_xlim(0, args.days) axes[2].set_ylim(0, 1) axes[2].set_title(f"Time to confident decision (P(best)≥{args.confidence})") axes[2].set_xlabel("day") axes[2].set_ylabel("fraction of simulations decided") axes[2].legend(fontsize=8) for ax in axes: ax.grid(alpha=0.3) fig.tight_layout() fig.savefig(path, dpi=130) plt.close(fig) def main(argv=None): args = parse_args(argv) rates, out = plan(args) print(report(rates, out, args)) if args.plot: make_plot(rates, out, args, args.plot) print(f"\nplot saved to {args.plot}") return 0 if __name__ == "__main__": sys.exit(main())