abkit 0.1.0__py3-none-any.whl

abkit/result.py

@@ -0,0 +1,82 @@
+"""Probe results and the audit container."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+
+
+@dataclass(frozen=True)
+class ProbeResult:
+    """One measured check.
+
+    A probe triggers in one of two ways: band probes carry an ``innocent``
+    interval and trigger when the whole 95% CI lies outside it; exact
+    procedures (a chi-square threshold, a count that must be zero) set
+    ``flagged`` directly and ``innocent`` stays None. Descriptive probes set
+    neither and can never trigger.
+    """
+
+    name: str
+    value: float
+    detail: str
+    n: int
+    unit: str = ""
+    ci: tuple[float, float] | None = None
+    innocent: tuple[float, float] | None = None
+    flagged: bool | None = None
+
+    @property
+    def triggered(self) -> bool:
+        if self.flagged is not None:
+            return self.flagged
+        if self.innocent is None or self.ci is None:
+            return False
+        lo, hi = self.ci
+        band_lo, band_hi = self.innocent
+        return hi < band_lo or lo > band_hi
+
+    def to_dict(self) -> dict[str, object]:
+        return {
+            "name": self.name,
+            "value": self.value,
+            "unit": self.unit,
+            "ci": list(self.ci) if self.ci is not None else None,
+            "innocent": list(self.innocent) if self.innocent is not None else None,
+            "triggered": self.triggered,
+            "detail": self.detail,
+            "n": self.n,
+        }
+
+
+@dataclass(frozen=True)
+class SkippedProbe:
+    """A check that could not run, and exactly what to log to enable it."""
+
+    name: str
+    reason: str
+
+    def to_dict(self) -> dict[str, object]:
+        return {"name": self.name, "skipped": True, "reason": self.reason}
+
+
+@dataclass
+class AuditResult:
+    """Everything the audit measured for one experiment readout."""
+
+    experiment: str
+    summary: dict[str, float | str | None] = field(default_factory=dict)
+    probes: list[ProbeResult] = field(default_factory=list)
+    skipped: list[SkippedProbe] = field(default_factory=list)
+
+    @property
+    def flags(self) -> list[ProbeResult]:
+        return [probe for probe in self.probes if probe.triggered]
+
+    def to_dict(self) -> dict[str, object]:
+        return {
+            "experiment": self.experiment,
+            "summary": self.summary,
+            "probes": [probe.to_dict() for probe in self.probes],
+            "skipped": [skip.to_dict() for skip in self.skipped],
+            "flags": [probe.name for probe in self.flags],
+        }

abkit/schema.py

@@ -0,0 +1,71 @@
+"""Data model for an experiment readout.
+
+The schema accepts the data as it was logged, including the problems an audit
+exists to find: duplicated unit ids, units present in several arms, or arms
+that never appear in the intended split all pass validation and are reported
+by the probes instead of rejected here.
+"""
+
+from __future__ import annotations
+
+from pydantic import BaseModel, ConfigDict, Field, field_validator
+
+
+class Unit(BaseModel):
+    """One randomized unit (user, session, cluster) and what it produced."""
+
+    model_config = ConfigDict(extra="forbid")
+
+    unit_id: str
+    arm: str
+    metrics: dict[str, float] = Field(default_factory=dict)
+    pre: dict[str, float] = Field(default_factory=dict)
+    t: float | None = None
+    """Assignment order or timestamp; any monotone unit works."""
+
+
+class Look(BaseModel):
+    """One interim analysis: cumulative sample size and the z statistic seen."""
+
+    model_config = ConfigDict(extra="forbid")
+
+    n: int = Field(gt=0)
+    z: float
+
+
+class Design(BaseModel):
+    """What the experiment intended, against which the data is audited."""
+
+    model_config = ConfigDict(extra="forbid")
+
+    split: dict[str, float]
+    control: str
+    treatment: str
+    primary_metric: str
+    alpha: float = Field(default=0.05, gt=0.0, lt=1.0)
+    metrics: list[str] | None = None
+    """Metrics that were tested; defaults to every metric observed."""
+    mde: float | None = None
+    """Absolute effect the experiment was designed to detect."""
+    planned_n: int | None = None
+    looks: list[Look] | None = None
+    """Interim analyses actually performed, in order, final look last."""
+
+    @field_validator("split")
+    @classmethod
+    def _positive_weights(cls, value: dict[str, float]) -> dict[str, float]:
+        if len(value) < 2:
+            raise ValueError("split needs at least two arms")
+        if any(weight <= 0 for weight in value.values()):
+            raise ValueError("split weights must be positive")
+        return value
+
+    def model_post_init(self, __context: object) -> None:
+        for arm in (self.control, self.treatment):
+            if arm not in self.split:
+                raise ValueError(f"arm {arm!r} is not in the intended split")
+
+    @property
+    def normalized_split(self) -> dict[str, float]:
+        total = sum(self.split.values())
+        return {arm: weight / total for arm, weight in self.split.items()}

abkit/stats.py

@@ -0,0 +1,222 @@
+"""Statistics used by the probes.
+
+Everything here is closed-form or an explicit seeded simulation, implemented
+on top of the standard library and elementwise numpy — no scipy at runtime and
+no BLAS in any statistic's path, so the same data gives the same digits on any
+machine. The test suite cross-checks each function against scipy or
+statsmodels.
+"""
+
+from __future__ import annotations
+
+import math
+from collections.abc import Mapping, Sequence
+from statistics import NormalDist
+
+import numpy as np
+
+_STANDARD_NORMAL = NormalDist()
+
+
+def normal_cdf(x: float) -> float:
+    return _STANDARD_NORMAL.cdf(x)
+
+
+def normal_pdf(x: float) -> float:
+    return _STANDARD_NORMAL.pdf(x)
+
+
+def normal_ppf(q: float) -> float:
+    return _STANDARD_NORMAL.inv_cdf(q)
+
+
+def _lower_gamma_series(a: float, x: float) -> float:
+    """Regularized lower incomplete gamma P(a, x) by series, for x < a + 1."""
+    term = 1.0 / a
+    total = term
+    denom = a
+    for _ in range(1000):
+        denom += 1.0
+        term *= x / denom
+        total += term
+        if abs(term) < abs(total) * 1e-16:
+            break
+    return total * math.exp(-x + a * math.log(x) - math.lgamma(a))
+
+
+def _upper_gamma_cf(a: float, x: float) -> float:
+    """Regularized upper incomplete gamma Q(a, x) by continued fraction."""
+    tiny = 1e-300
+    b = x + 1.0 - a
+    c = 1.0 / tiny
+    d = 1.0 / b if b != 0 else 1.0 / tiny
+    h = d
+    for i in range(1, 1000):
+        an = -i * (i - a)
+        b += 2.0
+        d = an * d + b
+        if abs(d) < tiny:
+            d = tiny
+        c = b + an / c
+        if abs(c) < tiny:
+            c = tiny
+        d = 1.0 / d
+        delta = d * c
+        h *= delta
+        if abs(delta - 1.0) < 1e-16:
+            break
+    return h * math.exp(-x + a * math.log(x) - math.lgamma(a))
+
+
+def chi2_sf(x: float, df: float) -> float:
+    """Survival function of the chi-square distribution."""
+    if x <= 0:
+        return 1.0
+    a = df / 2.0
+    half = x / 2.0
+    if half < a + 1.0:
+        p = 1.0 - _lower_gamma_series(a, half)
+    else:
+        p = _upper_gamma_cf(a, half)
+    return min(1.0, max(0.0, p))
+
+
+def srm_chi2(counts: Mapping[str, int], split: Mapping[str, float]) -> tuple[float, float]:
+    """Chi-square statistic and p-value for observed arm counts vs the intended split."""
+    total = sum(counts.get(arm, 0) for arm in split)
+    weight_total = sum(split.values())
+    stat = 0.0
+    for arm, weight in split.items():
+        expected = total * weight / weight_total
+        observed = counts.get(arm, 0)
+        stat += (observed - expected) ** 2 / expected
+    return stat, chi2_sf(stat, len(split) - 1)
+
+
+def mean(values: Sequence[float]) -> float:
+    return float(sum(values) / len(values))
+
+
+def variance(values: Sequence[float]) -> float:
+    """Unbiased sample variance."""
+    m = mean(values)
+    return float(sum((v - m) ** 2 for v in values) / (len(values) - 1))
+
+
+def welch(
+    treatment: Sequence[float], control: Sequence[float]
+) -> tuple[float, float, float, float]:
+    """Difference in means with its standard error, z statistic and two-sided p.
+
+    Uses the normal approximation, which is what experimentation platforms
+    report at the sample sizes where an audit matters.
+    """
+    effect = mean(treatment) - mean(control)
+    se = math.sqrt(variance(treatment) / len(treatment) + variance(control) / len(control))
+    if se == 0.0:
+        return effect, 0.0, math.inf if effect else 0.0, 0.0 if effect else 1.0
+    z = effect / se
+    p = 2.0 * (1.0 - normal_cdf(abs(z)))
+    return effect, se, z, p
+
+
+def two_proportion_z(
+    successes_a: int, n_a: int, successes_b: int, n_b: int
+) -> tuple[float, float, float]:
+    """Difference in proportions, pooled z statistic and two-sided p."""
+    p_a = successes_a / n_a
+    p_b = successes_b / n_b
+    pooled = (successes_a + successes_b) / (n_a + n_b)
+    se = math.sqrt(pooled * (1.0 - pooled) * (1.0 / n_a + 1.0 / n_b))
+    if se == 0.0:
+        return p_a - p_b, 0.0, 1.0
+    z = (p_a - p_b) / se
+    return p_a - p_b, z, 2.0 * (1.0 - normal_cdf(abs(z)))
+
+
+def standardized_mean_difference(treatment: Sequence[float], control: Sequence[float]) -> float:
+    """Cohen-style SMD with the pooled (average-variance) denominator."""
+    sd = math.sqrt((variance(treatment) + variance(control)) / 2.0)
+    if sd == 0.0:
+        return 0.0
+    return (mean(treatment) - mean(control)) / sd
+
+
+def benjamini_hochberg(pvalues: Sequence[float], alpha: float) -> list[bool]:
+    """Step-up FDR control; returns a rejection flag per input p-value."""
+    m = len(pvalues)
+    order = sorted(range(m), key=lambda i: pvalues[i])
+    threshold_rank = 0
+    for rank, idx in enumerate(order, start=1):
+        if pvalues[idx] <= rank * alpha / m:
+            threshold_rank = rank
+    rejected = [False] * m
+    for rank, idx in enumerate(order, start=1):
+        if rank <= threshold_rank:
+            rejected[idx] = True
+    return rejected
+
+
+def pearson(x: Sequence[float], y: Sequence[float]) -> float:
+    mx, my = mean(x), mean(y)
+    sxy = sum((a - mx) * (b - my) for a, b in zip(x, y, strict=True))
+    sxx = sum((a - mx) ** 2 for a in x)
+    syy = sum((b - my) ** 2 for b in y)
+    if sxx == 0.0 or syy == 0.0:
+        return math.nan
+    return sxy / math.sqrt(sxx * syy)
+
+
+def cuped_reduction(post: Sequence[float], pre: Sequence[float]) -> float:
+    """Share of metric variance CUPED would have removed, rho squared."""
+    rho = pearson(post, pre)
+    if math.isnan(rho):
+        return 0.0
+    return rho * rho
+
+
+def design_analysis(true_effect: float, se: float, alpha: float) -> tuple[float, float, float]:
+    """Power, Type-S rate and Type-M exaggeration ratio (Gelman and Carlin).
+
+    For an estimator distributed N(true_effect, se) and a two-sided test at
+    ``alpha``: the probability of significance, the probability a significant
+    result has the wrong sign, and the expected |estimate| among significant
+    results divided by |true_effect|.
+    """
+    if true_effect == 0.0 or se <= 0.0:
+        raise ValueError("design_analysis needs a nonzero effect and positive se")
+    crit = normal_ppf(1.0 - alpha / 2.0) * se
+    mu = abs(true_effect)
+    hi = (crit - mu) / se
+    lo = (-crit - mu) / se
+    p_hi = 1.0 - normal_cdf(hi)
+    p_lo = normal_cdf(lo)
+    power = p_hi + p_lo
+    upper_mass = mu * p_hi + se * normal_pdf(hi)
+    lower_mass = mu * p_lo - se * normal_pdf(lo)
+    exaggeration = (upper_mass - lower_mass) / power / mu
+    type_s = p_lo / power
+    return power, type_s, exaggeration
+
+
+def schedule_false_positive_rate(
+    look_ns: Sequence[int], alpha: float, seed: int = 0, n_sim: int = 20_000
+) -> float:
+    """False-positive rate of "stop when any look is significant" under H0.
+
+    Simulates the z-statistic path at the given cumulative sample sizes as a
+    Brownian motion observed at those information fractions, and counts null
+    experiments in which any look crosses the fixed-horizon critical value.
+    Seeded PCG64, elementwise numpy only.
+    """
+    if any(b <= a for a, b in zip(look_ns, look_ns[1:], strict=False)):
+        raise ValueError("look sample sizes must be strictly increasing")
+    z_crit = normal_ppf(1.0 - alpha / 2.0)
+    fractions = np.asarray(look_ns, dtype=np.float64) / float(look_ns[-1])
+    dt = np.diff(np.concatenate([np.zeros(1), fractions]))
+    rng = np.random.Generator(np.random.PCG64(seed))
+    increments = rng.standard_normal((n_sim, len(look_ns))) * np.sqrt(dt)
+    paths = np.cumsum(increments, axis=1)
+    z_paths = paths / np.sqrt(fractions)
+    crossed = (np.abs(z_paths) >= z_crit).any(axis=1)
+    return float(crossed.sum()) / float(n_sim)

abkit/synthetic.py

@@ -0,0 +1,144 @@
+"""Synthetic experiments with defect dials.
+
+Every failure mode abkit probes for can be implanted here at a chosen
+severity, which is how the probes are validated: each must recover its own
+dial, stay monotone in it, and stay silent on the clean experiment.
+"""
+
+from __future__ import annotations
+
+import math
+
+import numpy as np
+
+from .schema import Design, Look, Unit
+from .stats import normal_ppf
+
+PRE_CORRELATION = 0.6
+"""Correlation between the pre-experiment covariate and the outcome."""
+
+
+def simulate(
+    n: int = 8000,
+    effect: float = 0.2,
+    seed: int = 0,
+    srm: float = 0.0,
+    imbalance: float = 0.0,
+    whales: int = 0,
+    whale_value: float = 15.0,
+    novelty: float = 0.0,
+    contaminated: int = 0,
+    extra_metrics: int = 0,
+    mde: float | None = None,
+) -> tuple[list[Unit], Design]:
+    """Simulate a two-arm experiment on a unit outcome with SD 1.
+
+    Dials, all zero-off:
+
+    - ``srm``: extra assignment share diverted to treatment (0.03 gives 53/47).
+    - ``imbalance``: SMD implanted on the pre-experiment covariate in treatment.
+    - ``whales``: units appended to treatment with outcome ``whale_value``.
+    - ``novelty``: early-minus-late effect gap in SD units; the effect decays
+      linearly from ``effect + novelty`` to ``effect - novelty``, so the
+      first-half mean effect exceeds the second-half mean effect by exactly
+      ``novelty``.
+    - ``contaminated``: unit ids duplicated into the other arm.
+    - ``extra_metrics``: independent null secondary metrics per unit.
+    """
+    rng = np.random.Generator(np.random.PCG64(seed))
+    units: list[Unit] = []
+    load = math.sqrt(1.0 - PRE_CORRELATION**2)
+    for i in range(n):
+        in_treatment = rng.random() < 0.5 + srm
+        arm = "treatment" if in_treatment else "control"
+        pre = float(rng.standard_normal())
+        if in_treatment:
+            pre += imbalance
+        noise = float(rng.standard_normal())
+        outcome = PRE_CORRELATION * pre + load * noise
+        if in_treatment:
+            outcome += effect + novelty * (1.0 - 2.0 * i / n)
+        metrics = {"outcome": outcome}
+        for j in range(extra_metrics):
+            metrics[f"m{j + 1}"] = float(rng.standard_normal())
+        units.append(
+            Unit(unit_id=f"u{i}", arm=arm, metrics=metrics, pre={"pre_outcome": pre}, t=float(i))
+        )
+    for w in range(whales):
+        units.append(
+            Unit(
+                unit_id=f"whale{w}",
+                arm="treatment",
+                metrics={"outcome": whale_value},
+                pre={"pre_outcome": float(rng.standard_normal())},
+                t=float(rng.integers(0, n)),
+            )
+        )
+    for c, unit in enumerate(units[:contaminated]):
+        other = "control" if unit.arm == "treatment" else "treatment"
+        units.append(
+            Unit(
+                unit_id=unit.unit_id,
+                arm=other,
+                metrics=dict(unit.metrics),
+                pre=dict(unit.pre),
+                t=float(n + whales + c),
+            )
+        )
+    design = Design(
+        split={"control": 0.5, "treatment": 0.5},
+        control="control",
+        treatment="treatment",
+        primary_metric="outcome",
+        mde=mde,
+    )
+    return units, design
+
+
+def simulate_peeking(
+    n_max: int = 5000,
+    look_every: int = 250,
+    effect: float = 0.0,
+    seed: int = 0,
+    alpha: float = 0.05,
+) -> tuple[list[Unit], Design]:
+    """A sequential experiment stopped at the first significant interim look.
+
+    Returns the units accrued up to the stop together with a design whose
+    ``looks`` record every interim analysis performed — exactly what a peeking
+    analyst would have shipped.
+    """
+    rng = np.random.Generator(np.random.PCG64(seed))
+    z_crit = normal_ppf(1.0 - alpha / 2.0)
+    units: list[Unit] = []
+    looks: list[Look] = []
+    sum_t = sum_c = sumsq_t = sumsq_c = 0.0
+    n_t = n_c = 0
+    for i in range(n_max):
+        in_treatment = bool(rng.random() < 0.5)
+        value = float(rng.standard_normal()) + (effect if in_treatment else 0.0)
+        arm = "treatment" if in_treatment else "control"
+        units.append(Unit(unit_id=f"u{i}", arm=arm, metrics={"outcome": value}, t=float(i)))
+        if in_treatment:
+            n_t, sum_t, sumsq_t = n_t + 1, sum_t + value, sumsq_t + value * value
+        else:
+            n_c, sum_c, sumsq_c = n_c + 1, sum_c + value, sumsq_c + value * value
+        if (i + 1) % look_every == 0 and n_t > 1 and n_c > 1:
+            mean_t, mean_c = sum_t / n_t, sum_c / n_c
+            var_t = (sumsq_t - n_t * mean_t**2) / (n_t - 1)
+            var_c = (sumsq_c - n_c * mean_c**2) / (n_c - 1)
+            se = math.sqrt(var_t / n_t + var_c / n_c)
+            z = (mean_t - mean_c) / se if se > 0 else 0.0
+            looks.append(Look(n=i + 1, z=z))
+            if abs(z) >= z_crit:
+                break
+    design = Design(
+        split={"control": 0.5, "treatment": 0.5},
+        control="control",
+        treatment="treatment",
+        primary_metric="outcome",
+        alpha=alpha,
+        planned_n=n_max,
+        looks=looks,
+    )
+    return units, design