Experiment Analysis skill for Claude Code

#!/usr/bin/env python3
"""A/B test calculator: sample size, SRM check, proportion and mean tests (stdlib only).

Subcommands:
  size  --baseline 0.04 --mde-rel 0.05 [--alpha 0.05 --power 0.8 --daily 20000 --variants 2]
  size  --sd 42.0 --mde-abs 1.5 [...]                      (continuous metric)
  srm   --counts 50410 49590 [--ratios 1 1]
  prop  --a 1200 30000 --b 1310 30100 [--alpha 0.05]       (conversions visitors)
  mean  --a 52.1 40.3 30000 --b 53.0 41.0 30100            (mean sd n per group, Welch)

All tests are two-sided. Exit codes: 0 ok; 1 when `srm` detects a sample ratio
mismatch (p < 0.001); 2 on bad input. Results are for fixed-horizon tests analyzed
once at the planned sample size; repeated peeking invalidates the p-values.
"""
import argparse
import math
import sys
from statistics import NormalDist

N = NormalDist()


# ---------- special functions ----------
def gammaincc(a, x):
    """Regularized upper incomplete gamma Q(a, x)."""
    if x < 0 or a <= 0:
        raise ValueError("invalid gammaincc arguments")
    if x == 0:
        return 1.0
    gln = math.lgamma(a)
    if x < a + 1:  # series for P, then Q = 1 - P
        ap, s, d = a, 1.0 / a, 1.0 / a
        for _ in range(10_000):
            ap += 1
            d *= x / ap
            s += d
            if abs(d) < abs(s) * 1e-15:
                break
        return max(0.0, 1.0 - s * math.exp(-x + a * math.log(x) - gln))
    b = x + 1 - a  # continued fraction for Q
    c, d = 1 / 1e-300, 1 / b
    h = d
    for i in range(1, 10_000):
        an = -i * (i - a)
        b += 2
        d = an * d + b
        d = 1e-300 if abs(d) < 1e-300 else d
        c = b + an / c
        c = 1e-300 if abs(c) < 1e-300 else c
        d = 1 / d
        delta = d * c
        h *= delta
        if abs(delta - 1) < 1e-15:
            break
    return math.exp(-x + a * math.log(x) - gln) * h


def chi2_sf(x, df):
    return gammaincc(df / 2.0, x / 2.0)


def betacf(a, b, x):
    qab, qap, qam = a + b, a + 1, a - 1
    c, d = 1.0, 1 - qab * x / qap
    d = 1e-300 if abs(d) < 1e-300 else d
    d = 1 / d
    h = d
    for m in range(1, 10_000):
        m2 = 2 * m
        aa = m * (b - m) * x / ((qam + m2) * (a + m2))
        d = 1 + aa * d
        d = 1e-300 if abs(d) < 1e-300 else d
        c = 1 + aa / c
        c = 1e-300 if abs(c) < 1e-300 else c
        d = 1 / d
        h *= d * c
        aa = -(a + m) * (qab + m) * x / ((a + m2) * (qap + m2))
        d = 1 + aa * d
        d = 1e-300 if abs(d) < 1e-300 else d
        c = 1 + aa / c
        c = 1e-300 if abs(c) < 1e-300 else c
        d = 1 / d
        delta = d * c
        h *= delta
        if abs(delta - 1) < 1e-15:
            break
    return h


def betainc(a, b, x):
    """Regularized incomplete beta I_x(a, b)."""
    if x <= 0:
        return 0.0
    if x >= 1:
        return 1.0
    bt = math.exp(math.lgamma(a + b) - math.lgamma(a) - math.lgamma(b) + a * math.log(x) + b * math.log(1 - x))
    if x < (a + 1) / (a + b + 2):
        return bt * betacf(a, b, x) / a
    return 1 - bt * betacf(b, a, 1 - x) / b


def t_sf_two_sided(t, df):
    return betainc(df / 2.0, 0.5, df / (df + t * t))


def t_ppf(q, df):
    """Quantile of Student t via bisection (q in (0.5, 1))."""
    lo, hi = 0.0, 1000.0
    target = 2 * (1 - q)  # two-sided tail
    for _ in range(200):
        mid = (lo + hi) / 2
        if t_sf_two_sided(mid, df) > target:
            lo = mid
        else:
            hi = mid
    return (lo + hi) / 2


# ---------- commands ----------
def cmd_size(a):
    za = N.inv_cdf(1 - a.alpha / 2)
    zb = N.inv_cdf(a.power)
    if a.baseline is not None:
        p1 = a.baseline
        if a.mde_rel is not None:
            p2 = p1 * (1 + a.mde_rel)
        elif a.mde_abs is not None:
            p2 = p1 + a.mde_abs
        else:
            raise SystemExit("error: give --mde-rel or --mde-abs")
        if not (0 < p1 < 1 and 0 < p2 < 1):
            raise SystemExit("error: rates must be within (0, 1)")
        pbar = (p1 + p2) / 2
        n = (za * math.sqrt(2 * pbar * (1 - pbar)) + zb * math.sqrt(p1 * (1 - p1) + p2 * (1 - p2))) ** 2 / (p2 - p1) ** 2
        desc = f"proportion {p1:.4%} -> {p2:.4%} (abs diff {p2 - p1:+.4%}, rel {p2 / p1 - 1:+.2%})"
    else:
        if a.sd is None or a.mde_abs is None:
            raise SystemExit("error: continuous metric needs --sd and --mde-abs")
        n = 2 * (za + zb) ** 2 * a.sd ** 2 / a.mde_abs ** 2
        desc = f"mean difference {a.mde_abs} with sd {a.sd}"
    n = math.ceil(n)
    total = n * a.variants
    print(f"Detect {desc}")
    print(f"  alpha {a.alpha} (two-sided), power {a.power}, {a.variants} groups")
    print(f"  required per group: {n:,}")
    print(f"  required total:     {total:,}")
    if a.daily:
        days = math.ceil(total / a.daily)
        weeks = math.ceil(days / 7)
        print(f"  at {a.daily:,} eligible units/day: {days} days -> run {weeks * 7} days (whole weeks)")
    print("  Note: with more than one treatment vs control, consider alpha / (variants - 1).")
    return 0


def cmd_srm(a):
    counts = a.counts
    ratios = a.ratios or [1] * len(counts)
    if len(ratios) != len(counts) or len(counts) < 2 or any(c < 0 for c in counts):
        raise SystemExit("error: need >= 2 non-negative counts and matching ratios")
    total = sum(counts)
    rsum = sum(ratios)
    expected = [total * r / rsum for r in ratios]
    chi2 = sum((o - e) ** 2 / e for o, e in zip(counts, expected))
    p = chi2_sf(chi2, len(counts) - 1)
    print("Sample ratio mismatch check (chi-square goodness of fit)")
    for i, (o, e) in enumerate(zip(counts, expected)):
        print(f"  group {i}: observed {o:,}  expected {e:,.0f}  share {o / total:.4%}")
    print(f"  chi2 = {chi2:.3f}, df = {len(counts) - 1}, p = {p:.3g}")
    if p < 0.001:
        print("  SRM DETECTED (p < 0.001): do not trust the results; investigate assignment/logging.")
        return 1
    print("  No SRM at p < 0.001.")
    return 0


def cmd_prop(a):
    xa, na = a.a
    xb, nb = a.b
    if not (0 <= xa <= na and 0 <= xb <= nb and na > 0 and nb > 0):
        raise SystemExit("error: need 0 <= conversions <= visitors")
    pa, pb = xa / na, xb / nb
    pooled = (xa + xb) / (na + nb)
    se_pooled = math.sqrt(pooled * (1 - pooled) * (1 / na + 1 / nb))
    z = (pb - pa) / se_pooled if se_pooled else 0.0
    p = 2 * (1 - N.cdf(abs(z)))
    se = math.sqrt(pa * (1 - pa) / na + pb * (1 - pb) / nb)
    zc = N.inv_cdf(1 - a.alpha / 2)
    diff = pb - pa
    lo, hi = diff - zc * se, diff + zc * se
    print("Two-proportion z-test (B vs A)")
    print(f"  A: {xa:,}/{na:,} = {pa:.4%}")
    print(f"  B: {xb:,}/{nb:,} = {pb:.4%}")
    print(f"  absolute diff: {diff:+.4%}  ({1 - a.alpha:.0%} CI {lo:+.4%} to {hi:+.4%})")
    if pa > 0:
        # delta-method CI on relative lift (log ratio)
        if pb > 0:
            se_log = math.sqrt((1 - pa) / xa + (1 - pb) / xb) if xa and xb else float("nan")
            rr = pb / pa
            rlo, rhi = rr * math.exp(-zc * se_log), rr * math.exp(zc * se_log)
            print(f"  relative lift: {rr - 1:+.2%}  (CI {rlo - 1:+.2%} to {rhi - 1:+.2%})")
    print(f"  z = {z:.3f}, p = {p:.4g}  -> {'significant' if p < a.alpha else 'not significant'} at alpha {a.alpha}")
    return 0


def cmd_mean(a):
    ma, sa, na = a.a
    mb, sb, nb = a.b
    na, nb = int(na), int(nb)
    if na < 2 or nb < 2 or sa < 0 or sb < 0:
        raise SystemExit("error: need n >= 2 and sd >= 0")
    va, vb = sa ** 2 / na, sb ** 2 / nb
    se = math.sqrt(va + vb)
    diff = mb - ma
    t = diff / se if se else 0.0
    df = (va + vb) ** 2 / (va ** 2 / (na - 1) + vb ** 2 / (nb - 1)) if (va + vb) else na + nb - 2
    p = t_sf_two_sided(abs(t), df)
    tc = t_ppf(1 - a.alpha / 2, df)
    print("Welch t-test (B vs A)")
    print(f"  A: mean {ma:g}, sd {sa:g}, n {na:,}")
    print(f"  B: mean {mb:g}, sd {sb:g}, n {nb:,}")
    print(f"  diff: {diff:+.4g}  ({1 - a.alpha:.0%} CI {diff - tc * se:+.4g} to {diff + tc * se:+.4g})")
    if ma:
        print(f"  relative: {diff / ma:+.2%}")
    print(f"  t = {t:.3f}, df = {df:.1f}, p = {p:.4g}  -> {'significant' if p < a.alpha else 'not significant'} at alpha {a.alpha}")
    return 0


def main():
    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    sub = ap.add_subparsers(dest="cmd", required=True)

    s = sub.add_parser("size")
    s.add_argument("--baseline", type=float)
    s.add_argument("--mde-rel", type=float)
    s.add_argument("--mde-abs", type=float)
    s.add_argument("--sd", type=float)
    s.add_argument("--alpha", type=float, default=0.05)
    s.add_argument("--power", type=float, default=0.8)
    s.add_argument("--daily", type=int, help="eligible units per day across all groups")
    s.add_argument("--variants", type=int, default=2, help="number of groups including control")

    r = sub.add_parser("srm")
    r.add_argument("--counts", type=int, nargs="+", required=True)
    r.add_argument("--ratios", type=float, nargs="+")

    p = sub.add_parser("prop")
    p.add_argument("--a", type=int, nargs=2, required=True, metavar=("CONV", "N"))
    p.add_argument("--b", type=int, nargs=2, required=True, metavar=("CONV", "N"))
    p.add_argument("--alpha", type=float, default=0.05)

    m = sub.add_parser("mean")
    m.add_argument("--a", type=float, nargs=3, required=True, metavar=("MEAN", "SD", "N"))
    m.add_argument("--b", type=float, nargs=3, required=True, metavar=("MEAN", "SD", "N"))
    m.add_argument("--alpha", type=float, default=0.05)

    a = ap.parse_args()
    try:
        return {"size": cmd_size, "srm": cmd_srm, "prop": cmd_prop, "mean": cmd_mean}[a.cmd](a)
    except SystemExit as exc:
        if isinstance(exc.code, str):
            print(exc.code, file=sys.stderr)
            return 2
        raise


if __name__ == "__main__":
    sys.exit(main())