abkit 0.1.0__py3-none-any.whl

abkit/__init__.py

@@ -0,0 +1,23 @@
+"""abkit: audit an A/B-test readout before you ship the decision."""
+
+from .audit import run_audit
+from .io import dump_design, dump_units, load_design, load_units
+from .result import AuditResult, ProbeResult, SkippedProbe
+from .schema import Design, Look, Unit
+
+__version__ = "0.1.0"
+
+__all__ = [
+    "AuditResult",
+    "Design",
+    "Look",
+    "ProbeResult",
+    "SkippedProbe",
+    "Unit",
+    "__version__",
+    "dump_design",
+    "dump_units",
+    "load_design",
+    "load_units",
+    "run_audit",
+]

abkit/audit.py

@@ -0,0 +1,58 @@
+"""Run every applicable probe against one experiment readout."""
+
+from __future__ import annotations
+
+from collections.abc import Sequence
+
+from . import probes
+from .result import AuditResult, ProbeResult, SkippedProbe
+from .schema import Design, Unit
+from .stats import welch
+
+
+def run_audit(units: Sequence[Unit], design: Design, name: str = "experiment") -> AuditResult:
+    result = AuditResult(experiment=name)
+
+    treat = [
+        u.metrics[design.primary_metric]
+        for u in units
+        if u.arm == design.treatment and design.primary_metric in u.metrics
+    ]
+    ctrl = [
+        u.metrics[design.primary_metric]
+        for u in units
+        if u.arm == design.control and design.primary_metric in u.metrics
+    ]
+    if len(treat) >= 2 and len(ctrl) >= 2:
+        effect, se, z, p = welch(treat, ctrl)
+        result.summary = {
+            "primary_metric": design.primary_metric,
+            "effect": effect,
+            "se": se,
+            "z": z,
+            "p": p,
+            "n_treatment": len(treat),
+            "n_control": len(ctrl),
+            "alpha": design.alpha,
+        }
+
+    outcomes: list[ProbeResult | SkippedProbe | list[ProbeResult]] = [
+        probes.sample_ratio(units, design),
+        probes.contamination(units),
+        probes.covariate_balance(units, design),
+        probes.outlier_fragility(units, design),
+        probes.uncorrected_winners(units, design),
+        probes.peeking(design),
+        probes.novelty(units, design),
+        probes.winners_curse(units, design),
+        probes.variance_ratio(units, design),
+        probes.cuped_headroom(units, design),
+    ]
+    for outcome in outcomes:
+        if isinstance(outcome, SkippedProbe):
+            result.skipped.append(outcome)
+        elif isinstance(outcome, list):
+            result.probes.extend(outcome)
+        else:
+            result.probes.append(outcome)
+    return result

abkit/bootstrap.py

@@ -0,0 +1,46 @@
+"""Stratified bootstrap confidence intervals.
+
+Units are resampled within their arm (or other stratum), never across, so the
+group sizes that determine a statistic's sampling distribution are preserved.
+Percentile intervals, seeded PCG64, no BLAS.
+"""
+
+from __future__ import annotations
+
+import math
+from collections.abc import Callable, Sequence
+
+import numpy as np
+
+MAX_UNDEFINED_SHARE = 0.1
+"""If more than this share of resamples yields nan, the CI is not reported."""
+
+
+def bootstrap_ci(
+    statistic: Callable[..., float],
+    groups: Sequence[Sequence[float]],
+    n_resamples: int = 1000,
+    seed: int = 0,
+) -> tuple[float, tuple[float, float] | None]:
+    """Point estimate and 95% percentile CI, resampling within each group."""
+    point = statistic(*groups)
+    rng = np.random.Generator(np.random.PCG64(seed))
+    arrays = [list(group) for group in groups]
+    values: list[float] = []
+    undefined = 0
+    for _ in range(n_resamples):
+        resampled = []
+        for group in arrays:
+            indices = rng.integers(0, len(group), size=len(group)).tolist()
+            resampled.append([group[i] for i in indices])
+        value = statistic(*resampled)
+        if math.isnan(value):
+            undefined += 1
+        else:
+            values.append(value)
+    if undefined > MAX_UNDEFINED_SHARE * n_resamples or not values:
+        return point, None
+    values.sort()
+    lo = values[int(0.025 * len(values))]
+    hi = values[min(len(values) - 1, int(0.975 * len(values)))]
+    return point, (lo, hi)

abkit/cli.py

@@ -0,0 +1,71 @@
+"""Command line interface."""
+
+from __future__ import annotations
+
+import argparse
+from pathlib import Path
+
+from . import __version__
+from .audit import run_audit
+from .demo import run_demo
+from .io import load_design, load_units
+from .report import inject_readme, write_report
+
+
+def main(argv: list[str] | None = None) -> int:
+    parser = argparse.ArgumentParser(
+        prog="abkit", description="Audit an A/B-test readout before you ship the decision."
+    )
+    parser.add_argument("--version", action="version", version=f"abkit {__version__}")
+    sub = parser.add_subparsers(dest="command", required=True)
+
+    report_cmd = sub.add_parser("report", help="audit a units JSONL against a design JSON")
+    report_cmd.add_argument("units", help="JSONL file, one unit per line")
+    report_cmd.add_argument("--design", required=True, help="design JSON file")
+    report_cmd.add_argument("--out", default="audit", help="output directory")
+    report_cmd.add_argument("--name", default=None, help="experiment name in the report")
+    report_cmd.add_argument(
+        "--fail-on-flags", action="store_true", help="exit 2 if any probe triggers"
+    )
+
+    demo_cmd = sub.add_parser("demo", help="rebuild the implanted-defect demo")
+    demo_cmd.add_argument("--out", default="results", help="output directory")
+
+    inject_cmd = sub.add_parser(
+        "inject-readme", help="refresh the demo table between README markers"
+    )
+    inject_cmd.add_argument("readme")
+    inject_cmd.add_argument("--results", default="results", help="demo output directory")
+
+    args = parser.parse_args(argv)
+
+    if args.command == "report":
+        units = load_units(args.units)
+        design = load_design(args.design)
+        name = args.name or Path(args.units).stem
+        audit = run_audit(units, design, name=name)
+        report_path = write_report(audit, args.out, design=design, units=units)
+        print(f"wrote {report_path}")
+        for probe in audit.flags:
+            print(f"FLAG {probe.name} -- {probe.detail}")
+        if args.fail_on_flags and audit.flags:
+            return 2
+        return 0
+
+    if args.command == "demo":
+        audits = run_demo(args.out)
+        flagged = sum(1 for audit in audits if audit.flags)
+        print(f"wrote {args.out}/report.md ({flagged}/{len(audits)} experiments flagged)")
+        return 0
+
+    if args.command == "inject-readme":
+        table = (Path(args.results) / "table.md").read_text(encoding="utf-8")
+        inject_readme(args.readme, table)
+        print(f"updated {args.readme}")
+        return 0
+
+    raise AssertionError("unreachable")
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

abkit/demo.py

@@ -0,0 +1,177 @@
+"""Eight experiments, seven implanted defects, one table.
+
+Each demo experiment implants exactly one failure mode at a severity chosen to
+be decisive, then runs the full audit. The table is regenerated from fixed
+seeds; CI diffs it against the committed copy.
+
+Two experiments deliberately search over seeds: the peeker keeps trying null
+experiments until one stops early at an interim look, and the underpowered
+winner keeps trying until a true effect of 0.1 comes out significant. That
+selection is not a trick — it is the defect. Peeking and publishing only
+winners are selection processes, and the demo reproduces the selection.
+"""
+
+from __future__ import annotations
+
+from pathlib import Path
+
+from .audit import run_audit
+from .probes import uncorrected_winners
+from .report import (
+    exaggeration_figure,
+    fragility_figure,
+    peeking_figure,
+    render_report,
+)
+from .result import AuditResult, SkippedProbe
+from .schema import Design, Unit
+from .stats import design_analysis, welch
+from .synthetic import simulate, simulate_peeking
+
+DEFECTS = {
+    "clean": "nothing",
+    "traffic-leak": "3% of traffic diverted to treatment",
+    "double-dipper": "25 units assigned to both arms",
+    "whale-driven": "3 extreme units carry the significance",
+    "peeker": "null effect, stopped at the first significant look",
+    "fading-novelty": "early-half effect 0.30, late-half effect 0.00",
+    "metric-fisher": "null effect, 40 secondary metrics tested",
+    "underpowered-winner": "true effect 0.1, powered for 0.14 of that chance",
+}
+
+
+def _confirm(
+    units: list[Unit], design: Design, name: str, target: set[str]
+) -> tuple[list[Unit], Design] | None:
+    audit = run_audit(units, design, name=name)
+    if {probe.name for probe in audit.flags} == target:
+        return units, design
+    return None
+
+
+def _find_peeker() -> tuple[list[Unit], Design]:
+    for seed in range(500):
+        units, design = simulate_peeking(n_max=5000, look_every=250, effect=0.0, seed=seed)
+        looks = design.looks or []
+        if not looks or looks[-1].n < 1500 or looks[-1].n > 3500 or abs(looks[-1].z) < 2.2:
+            continue
+        found = _confirm(units, design, "peeker", {"peeking"})
+        if found:
+            return found
+    raise RuntimeError("no early-stopping null experiment found in 500 seeds")
+
+
+def _find_underpowered() -> tuple[list[Unit], Design]:
+    for seed in range(500):
+        units, design = simulate(n=300, effect=0.1, mde=0.1, seed=seed)
+        treat = [u.metrics["outcome"] for u in units if u.arm == "treatment"]
+        ctrl = [u.metrics["outcome"] for u in units if u.arm == "control"]
+        effect, _, _, p = welch(treat, ctrl)
+        if p >= design.alpha or effect <= 0:
+            continue
+        found = _confirm(units, design, "underpowered-winner", {"winner's curse"})
+        if found:
+            return found
+    raise RuntimeError("no significant underpowered experiment found in 500 seeds")
+
+
+def _find_metric_fisher() -> tuple[list[Unit], Design]:
+    for seed in range(500):
+        units, design = simulate(n=2000, effect=0.0, extra_metrics=40, seed=seed)
+        probe = uncorrected_winners(units, design)
+        if isinstance(probe, SkippedProbe) or not probe.triggered:
+            continue
+        found = _confirm(units, design, "metric-fisher", {"uncorrected winners"})
+        if found:
+            return found
+    raise RuntimeError("no experiment with uncorrected winners found in 500 seeds")
+
+
+def build_experiments() -> list[tuple[str, list[Unit], Design]]:
+    return [
+        ("clean", *simulate(n=8000, effect=0.2, mde=0.2, seed=11)),
+        ("traffic-leak", *simulate(n=8000, effect=0.2, srm=0.03, mde=0.2, seed=12)),
+        ("double-dipper", *simulate(n=4000, effect=0.2, contaminated=25, mde=0.2, seed=13)),
+        (
+            "whale-driven",
+            *simulate(n=1000, effect=0.12, whales=3, whale_value=15.0, mde=0.2, seed=2),
+        ),
+        ("peeker", *_find_peeker()),
+        ("fading-novelty", *simulate(n=6000, effect=0.15, novelty=0.3, mde=0.15, seed=16)),
+        ("metric-fisher", *_find_metric_fisher()),
+        ("underpowered-winner", *_find_underpowered()),
+    ]
+
+
+def _cell(audit: AuditResult, probe_name: str, fmt: str) -> str:
+    for probe in audit.probes:
+        if probe.name == probe_name:
+            text = format(probe.value, fmt)
+            return f"**{text}**" if probe.triggered else text
+    return "-"
+
+
+def demo_table(audits: list[AuditResult]) -> str:
+    lines = [
+        "| experiment | implanted defect | SRM p | contam. | fragility "
+        "| look FPR | novelty gap | uncorrected | exagg. | flags |",
+        "|---|---|---|---|---|---|---|---|---|---|",
+    ]
+    curse = "winner's curse"
+    for audit in audits:
+        flags = ", ".join(p.name for p in audit.flags) or "none"
+        lines.append(
+            f"| {audit.experiment} | {DEFECTS[audit.experiment]} "
+            f"| {_cell(audit, 'sample ratio mismatch', '.2g')} "
+            f"| {_cell(audit, 'assignment contamination', '.0f')} "
+            f"| {_cell(audit, 'outlier fragility', '.0f')} "
+            f"| {_cell(audit, 'peeking', '.2f')} "
+            f"| {_cell(audit, 'novelty', '.2f')} "
+            f"| {_cell(audit, 'uncorrected winners', '.0f')} "
+            f"| {_cell(audit, curse, '.1f')} "
+            f"| {flags} |"
+        )
+    return "\n".join(lines)
+
+
+def run_demo(out_dir: str | Path) -> list[AuditResult]:
+    out = Path(out_dir)
+    figures = out / "figures"
+    figures.mkdir(parents=True, exist_ok=True)
+    experiments = build_experiments()
+    audits = []
+    sections = []
+    marks = []
+    for name, units, design in experiments:
+        audit = run_audit(units, design, name=name)
+        audits.append(audit)
+        sections.append(render_report(audit).replace("# abkit audit:", "## "))
+        if (
+            name in ("clean", "whale-driven", "underpowered-winner")
+            and audit.summary
+            and design.mde is not None
+        ):
+            se = float(audit.summary["se"])  # type: ignore[arg-type]
+            power, _, ratio = design_analysis(design.mde, se, design.alpha)
+            if float(audit.summary["p"]) < design.alpha:  # type: ignore[arg-type]
+                marks.append((name, power, ratio))
+        if name == "peeker" and design.looks:
+            peeking_figure(design.looks, design.alpha, figures / "peeking.png")
+        if name == "whale-driven":
+            treat = [u.metrics["outcome"] for u in units if u.arm == "treatment"]
+            ctrl = [u.metrics["outcome"] for u in units if u.arm == "control"]
+            fragility_figure(treat, ctrl, design.alpha, figures / "fragility.png")
+    exaggeration_figure(marks, 0.05, figures / "exaggeration.png")
+
+    table = demo_table(audits)
+    (out / "table.md").write_text(table + "\n", encoding="utf-8")
+    header = (
+        "# abkit demo\n\n"
+        "Eight synthetic experiments, seven implanted defects. Bold values are\n"
+        "flags; each lands on the row where its defect was implanted.\n\n"
+    )
+    (out / "report.md").write_text(header + table + "\n\n" + "\n".join(sections), encoding="utf-8")
+    return audits
+
+
+__all__ = ["build_experiments", "demo_table", "run_demo", "DEFECTS"]

abkit/io.py

@@ -0,0 +1,38 @@
+"""Read and write experiment data.
+
+Units are one JSON object per line; the design is a single JSON document.
+"""
+
+from __future__ import annotations
+
+import json
+from pathlib import Path
+
+from .schema import Design, Unit
+
+
+def load_units(path: str | Path) -> list[Unit]:
+    units = []
+    with open(path, encoding="utf-8") as handle:
+        for line in handle:
+            line = line.strip()
+            if line:
+                units.append(Unit.model_validate_json(line))
+    return units
+
+
+def dump_units(units: list[Unit], path: str | Path) -> None:
+    with open(path, "w", encoding="utf-8") as handle:
+        for unit in units:
+            handle.write(unit.model_dump_json(exclude_defaults=True) + "\n")
+
+
+def load_design(path: str | Path) -> Design:
+    return Design.model_validate_json(Path(path).read_text(encoding="utf-8"))
+
+
+def dump_design(design: Design, path: str | Path) -> None:
+    Path(path).write_text(
+        json.dumps(design.model_dump(exclude_none=True), indent=2) + "\n",
+        encoding="utf-8",
+    )