abfactor 0.5.1__py3-none-any.whl

abfactor/results.py

@@ -0,0 +1,163 @@
+"""Result of an analysis: effects first, and saveable."""
+from __future__ import annotations
+
+import json
+from dataclasses import dataclass, field
+from typing import Optional
+
+import numpy as np
+import pandas as pd
+import xarray as xr
+
+FORMAT_VERSION = 1
+PAIRWISE_COLUMNS = ("log_BF10", "BF10", "BFplus0", "BFminus0", "P_better", "contrast_mean", "contrast_sd")
+
+
+def _labels_json(labels) -> str:
+    try:
+        return json.dumps([v.item() if hasattr(v, "item") else v for v in labels])
+    except TypeError:
+        return json.dumps([str(v) for v in labels])
+
+
+def _jsonable(obj):
+    if isinstance(obj, dict):
+        return {k: _jsonable(v) for k, v in obj.items() if not isinstance(v, np.ndarray)}
+    if isinstance(obj, (list, tuple)):
+        return [_jsonable(v) for v in obj]
+    if isinstance(obj, (np.integer, np.floating)):
+        return obj.item()
+    return obj
+
+
+@dataclass
+class Result:
+    """
+    - `contrasts`: draws of each variant's effect minus control on the model scale, dims
+      (experiment, variant, sample)
+    - `natural`: draws of each arm on the natural scale in the first treated period, with
+      any persistent arm offset removed, dims (experiment, arm, sample)
+    - `omnibus`: per experiment: log_BF10 ("some arm differs"; given the rest of the batch
+      when pooled) and P_effect (spike-and-slab only)
+    - `pairwise`: per experiment and variant: log_BF10, BF10, BF+0, BF-0, P(better), the
+      contrast's posterior mean and SD
+    - `pooling`: batch summaries: w_mean, log_BF_any, posterior over g (grid, weights)
+    - `diagnostics`: engine, divergences, warnings, notes
+    - `settings`: everything needed to recompute this result: Analysis.from_settings(settings)
+
+    Save with to_netcdf(path) and load with Result.from_netcdf(path); to_frame() gives
+    the per-variant summary as a flat table.
+    """
+    metric: str
+    scale: str
+    control: object
+    arms: list
+    experiments: list
+    contrasts: xr.DataArray
+    natural: xr.DataArray
+    omnibus: pd.DataFrame
+    pairwise: pd.DataFrame
+    pooling: dict = field(default_factory=dict)
+    diagnostics: dict = field(default_factory=dict)
+    settings: dict = field(default_factory=dict)
+
+    @property
+    def variants(self) -> list:
+        """Arm labels other than control."""
+        return [a for a in self.arms if a != self.control]
+
+    def lifts(self) -> xr.DataArray:
+        """Relative lift of each variant over control on the natural scale, (experiment, variant, sample)."""
+        ctrl = self.natural.sel(arm=self.control)
+        lift = self.natural.sel(arm=self.variants) / ctrl - 1
+        return lift.rename(arm="variant")
+
+    def summary(self, interval: float = 0.9) -> pd.DataFrame:
+        """One row per experiment and variant: evidence, direction and size."""
+        lo, hi = (1 - interval) / 2, 1 - (1 - interval) / 2
+        q = self.lifts().quantile([0.5, lo, hi], dim="sample")
+        om = self.omnibus.set_index("experiment")
+        rows = []
+        for _, r in self.pairwise.iterrows():
+            ql = q.sel(experiment=r["experiment"], variant=r["variant"]).values
+            row = dict(r)
+            row.update(lift_median=float(ql[0]), lift_lo=float(ql[1]), lift_hi=float(ql[2]))
+            row["omnibus_log_BF10"] = float(om.loc[r["experiment"], "log_BF10"])
+            if "P_effect" in om and pd.notna(om.loc[r["experiment"], "P_effect"]):
+                row["P_effect"] = float(om.loc[r["experiment"], "P_effect"])
+            rows.append(row)
+        return pd.DataFrame(rows)
+
+    def to_frame(self, interval: float = 0.9) -> pd.DataFrame:
+        """The summary as a flat table, with the metric and its scale on every row."""
+        out = self.summary(interval)
+        out.insert(0, "metric", self.metric)
+        out.insert(1, "scale", self.scale)
+        return out
+
+    # ---- saving -------------------------------------------------------------------
+
+    def to_dataset(self) -> xr.Dataset:
+        """Everything in the result as one xarray Dataset (labels stored as text,
+        their original types kept in the attributes)."""
+        E, V = [str(e) for e in self.experiments], [str(v) for v in self.variants]
+        ds = xr.Dataset(
+            {"natural": (("experiment", "arm", "sample"), self.natural.values),
+             "contrasts": (("experiment", "variant", "sample"), self.contrasts.values),
+             "omnibus_log_BF10": ("experiment", self.omnibus["log_BF10"].to_numpy(float)),
+             "omnibus_P_effect": ("experiment", self.omnibus["P_effect"].to_numpy(float))},
+            coords={"experiment": E, "arm": [str(a) for a in self.arms], "variant": V})
+        pw = self.pairwise.assign(experiment=self.pairwise["experiment"].map(str),
+                                  variant=self.pairwise["variant"].map(str)).set_index(["experiment", "variant"])
+        for col in PAIRWISE_COLUMNS:
+            grid = pw[col].astype(float).unstack("variant").reindex(index=E, columns=V)
+            ds[f"pairwise_{col}"] = (("experiment", "variant"), grid.to_numpy())
+        for key in ("g_grid", "log_g_batch"):
+            if self.pooling.get(key) is not None:
+                ds[f"pooling_{key}"] = ("g", np.asarray(self.pooling[key], dtype=float))
+        ds.attrs.update(
+            abfactor_format=FORMAT_VERSION, metric=self.metric, scale=self.scale,
+            control=json.dumps(self.control.item() if hasattr(self.control, "item") else self.control),
+            experiment_labels=_labels_json(self.experiments), arm_labels=_labels_json(self.arms),
+            pooling=json.dumps(_jsonable(self.pooling)), diagnostics=json.dumps(_jsonable(self.diagnostics)),
+            settings=json.dumps(self.settings))
+        return ds
+
+    @classmethod
+    def from_dataset(cls, ds: xr.Dataset) -> "Result":
+        """Rebuild a Result from `to_dataset()` output."""
+        if int(ds.attrs.get("abfactor_format", 0)) != FORMAT_VERSION:
+            raise ValueError("Not a abfactor result, or written by an incompatible version")
+        experiments = json.loads(ds.attrs["experiment_labels"])
+        arms = json.loads(ds.attrs["arm_labels"])
+        control = json.loads(ds.attrs["control"])
+        variants = [a for a in arms if a != control]
+        natural = xr.DataArray(ds["natural"].values, dims=("experiment", "arm", "sample"),
+                               coords={"experiment": experiments, "arm": arms})
+        contrasts = xr.DataArray(ds["contrasts"].values, dims=("experiment", "variant", "sample"),
+                                 coords={"experiment": experiments, "variant": variants})
+        omnibus = pd.DataFrame({"experiment": experiments, "log_BF10": ds["omnibus_log_BF10"].values,
+                                "P_effect": ds["omnibus_P_effect"].values})
+        rows = []
+        for i, e in enumerate(experiments):
+            for j, v in enumerate(variants):
+                row = {"experiment": e, "variant": v}
+                row.update({c: float(ds[f"pairwise_{c}"].values[i, j]) for c in PAIRWISE_COLUMNS})
+                rows.append(row)
+        pooling = json.loads(ds.attrs["pooling"])
+        for key in ("g_grid", "log_g_batch"):
+            pooling[key] = ds[f"pooling_{key}"].values if f"pooling_{key}" in ds else None
+        return cls(metric=ds.attrs["metric"], scale=ds.attrs["scale"], control=control, arms=arms,
+                   experiments=experiments, contrasts=contrasts, natural=natural, omnibus=omnibus,
+                   pairwise=pd.DataFrame(rows), pooling=pooling,
+                   diagnostics=json.loads(ds.attrs["diagnostics"]), settings=json.loads(ds.attrs["settings"]))
+
+    def to_netcdf(self, path, **kwargs) -> None:
+        """Save to NetCDF. Works with any xarray backend; install h5netcdf or netCDF4 for
+        compression of large draws."""
+        self.to_dataset().to_netcdf(path, **kwargs)
+
+    @classmethod
+    def from_netcdf(cls, path) -> "Result":
+        """Load a Result saved with `to_netcdf`."""
+        return cls.from_dataset(xr.load_dataset(path))

abfactor/sequential.py

@@ -0,0 +1,259 @@
+"""Sequential monitoring: stop an experiment the first time the evidence crosses a
+threshold, with a false-positive guarantee that survives looking as often as you like.
+
+Two methods:
+
+method="exact" (Proportion only)
+    A conditional e-process. At each look, take the new users in control and in
+    variant k and the total number of them who converted. Under the null (equal rates)
+    the split of those conversions between the two arms is hypergeometric, whatever the
+    common rate is; under a log odds ratio psi it is Fisher's noncentral hypergeometric
+    with odds ratio exp(psi), whatever the baseline is. So
+
+        E_k,t = integral over psi of  prod_{looks s <= t} f(split_s; psi) / f(split_s; 0)  prior(psi)
+
+    is a mixture of likelihood ratios whose denominator is the exact null distribution,
+    a test martingale under every baseline. The prior on psi is the analysis's own prior
+    on a pairwise log odds ratio (Cauchy(0, sqrt(2) r) for DefaultFixedEffects). With
+    several variants the omnibus e-process is the average of the E_k, still a test
+    martingale, and by Ville's inequality
+
+        P(E_t >= threshold at any look) <= 1 / threshold     for every baseline.
+
+    This is exact: no normal approximation and no assumption about the baseline. Each
+    user must be counted once (cumulative counts of distinct users with one outcome
+    each); looks may be scheduled however you like, including on the data so far.
+
+    "Exact" is relative to that model of the data: users who are independent draws with
+    their arm's rate. Real A/A tests can break it (identical arms in the Upworthy archive
+    differed by about 12% beyond binomial noise; docs/real_data.md). Check your own A/A
+    tests with abfactor.aa before relying on it; if identical arms vary, use
+    method="bayes" with the analysis's arm_sd.
+    Conditioning on each look's total discards the little information the totals carry
+    about the effect, so there is almost no loss of power.
+
+method="bayes"
+    The Bayes factor BF10 from the analysis. With a prior fixed in advance and a null
+    that fixes the data distribution it is a test martingale, so the same bound holds.
+    With an unknown baseline the null is composite and the bound is not guaranteed;
+    in simulation it held comfortably (calibrate() checks your own settings). It works
+    for every metric.
+
+Either way the prior must be frozen when the Monitor is created: pool across completed
+experiments with Analysis.learn_prior, never across the running ones.
+"""
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Callable, Optional
+
+import numpy as np
+import pandas as pd
+from scipy.special import gammaln, logsumexp
+from scipy.stats import cauchy, norm
+
+from .analysis import Analysis
+from .metrics import Proportion
+from .pooling import Independent
+from .priors import DefaultFixedEffects, LearnedScale, NormalEffects
+from .stats import Stats
+
+__all__ = ["Monitor", "calibrate", "binomial_aa", "psi_grid", "conditional_log_lr"]
+
+METHODS = ("exact", "bayes")
+
+
+def psi_grid(prior, size: int = 401) -> tuple[np.ndarray, np.ndarray]:
+    """Points and log prior weights for a pairwise log odds ratio psi, equally spaced in
+    prior probability (midpoints of `size` equal-mass cells), so the grid is fine where
+    the prior is concentrated and every cell carries mass 1/size."""
+    q = (np.arange(size) + 0.5) / size
+    if isinstance(prior, DefaultFixedEffects):
+        psi = cauchy.ppf(q, 0, np.sqrt(2) * prior.r)
+    elif isinstance(prior, NormalEffects):
+        psi = norm.ppf(q, 0, prior.contrast_sd)
+    elif isinstance(prior, LearnedScale):
+        u, lm = prior.grid()
+        keep = lm > lm.max() - 40                                       # drop negligible mass
+        w = np.exp(lm[keep] - logsumexp(lm[keep]))
+        sd = np.sqrt(2 * np.exp(u[keep]))
+        x = np.linspace(-12, 12, 48001)
+        cdf = np.zeros_like(x)                                          # mixture of normal CDFs,
+        for c in range(0, len(w), 64):                                  # in chunks to bound memory
+            cdf += w[c:c + 64] @ norm.cdf(x[None, :] / sd[c:c + 64, None])
+        psi = np.interp(q, cdf, x)
+    else:
+        raise ValueError(f"No exact sequential test for prior {type(prior).__name__}")
+    return psi, np.full(size, -np.log(size))
+
+
+def _log_choose(n, k):
+    return gammaln(n + 1) - gammaln(k + 1) - gammaln(n - k + 1)
+
+
+def conditional_log_lr(dy0, dn0, dy1, dn1, psi):
+    """log f(split; psi) - log f(split; 0) for one look's new users, where f is Fisher's
+    noncentral hypergeometric distribution of the variant's conversions given the total.
+    dy0, dn0, dy1, dn1: arrays (E,); psi: (G,). Returns (E, G)."""
+    m = dy0 + dy1
+    lo, hi = np.maximum(0, m - dn0), np.minimum(dn1, m)
+    width = int((hi - lo).max()) + 1 if len(m) else 1
+    u = lo[:, None] + np.arange(width)[None, :]                         # support of the split
+    valid = u <= hi[:, None]
+    lw = np.where(valid, _log_choose(dn0[:, None], m[:, None] - u) + _log_choose(dn1[:, None], u), -np.inf)
+    # log normaliser at each psi: logsumexp_u (lw + u psi), done in chunks of psi
+    out = np.empty((len(m), len(psi)))
+    step = max(1, int(2e7 // max(1, len(m) * width)))                 # bound memory (~160 MB)
+    for c in range(0, len(psi), step):
+        ps = psi[c:c + step]
+        z = logsumexp(lw[:, :, None] + u[:, :, None] * ps[None, None, :], axis=1)      # (E, chunk)
+        out[:, c:c + step] = dy1[:, None] * ps[None, :] - z
+    z0 = logsumexp(lw, axis=1)                                          # psi = 0
+    return out + z0[:, None]                                            # the binomial terms cancel
+
+
+@dataclass
+class _State:
+    y: np.ndarray               # cumulative conversions per arm
+    n: np.ndarray               # cumulative users per arm
+    ll: np.ndarray              # accumulated log likelihood ratio per variant and psi, (A-1, G)
+
+
+class Monitor:
+    """Sequential monitoring: look at running experiments as often as you like, and stop
+    each one the first time the evidence for an effect reaches `threshold`. The
+    false-positive rate stays below 1/threshold (see docs/sequential.md).
+
+    - `analysis`: an Analysis with pooling=Independent() and the quadrature engine; its
+      prior is frozen for the life of the monitor
+    - `threshold`: stop for an effect when the evidence first reaches this; false positives
+      are then at most 1/threshold (exactly for method="exact")
+    - `futility`: optionally stop for no effect when BF01 reaches this (does not affect the
+      false-positive guarantee)
+    - `method`: "exact" (e-process, Proportion only) or "bayes" (Bayes factor); by default
+      "exact" for a Proportion without arm_sd, otherwise "bayes"
+
+    update(stats) takes the cumulative statistics at a look, one period per experiment.
+    """
+
+    def __init__(self, analysis: Analysis, threshold: float = 20.0, futility: Optional[float] = None,
+                 method: Optional[str] = None):
+        if method is None:
+            method = "exact" if isinstance(analysis.metric, Proportion) and analysis.arm_sd == 0 else "bayes"
+        if method not in METHODS:
+            raise ValueError(f"method must be one of {METHODS}")
+        if not isinstance(analysis.pooling, Independent):
+            raise ValueError("Sequential monitoring needs pooling=Independent(): pool through a "
+                             "LearnedScale prior from completed experiments instead")
+        if analysis.engine != "quadrature":
+            raise ValueError("Monitor uses the quadrature engine")
+        if method == "exact" and analysis.arm_sd > 0:
+            raise ValueError("method='exact' assumes each user is an independent draw with the arm's rate; "
+                             "arm_sd > 0 says arms vary beyond that, which the exact test would read as an "
+                             "effect. Use method='bayes' (which models arm_sd) and check it with calibrate()")
+        if method == "exact" and not isinstance(analysis.metric, Proportion):
+            raise ValueError("method='exact' is implemented for Proportion; use method='bayes' "
+                             "and check it with calibrate()")
+        self.analysis, self.threshold, self.futility, self.method = analysis, threshold, futility, method
+        self.history = pd.DataFrame()
+        self._stopped: dict = {}
+        self._state: dict = {}
+        self._psi = self._log_w = None
+        self._last_variant_log_e: dict = {}
+
+    @property
+    def decisions(self) -> dict:
+        """Experiments stopped so far: {experiment: ("effect" | "no effect", look)}."""
+        return dict(self._stopped)
+
+    def _log_e(self, stats: Stats) -> np.ndarray:
+        """Advance each experiment's conditional e-process to this look; returns the
+        omnibus log E_t per experiment (the log of the average over variants)."""
+        if self._psi is None:
+            self._psi, self._log_w = psi_grid(self.analysis.prior)
+        name = self.analysis.metric.name
+        y_all, n_all = stats.get(name)[..., 0], stats.get("n")[..., 0]
+        E, A = y_all.shape
+        prev_y, prev_n, ll = np.empty_like(y_all), np.empty_like(n_all), np.empty((E, A - 1, len(self._psi)))
+        for i, e in enumerate(stats.experiments):
+            st = self._state.get(e)
+            if st is None:
+                st = _State(np.zeros(A), np.zeros(A), np.zeros((A - 1, len(self._psi))))
+            prev_y[i], prev_n[i], ll[i] = st.y, st.n, st.ll
+        dy, dn = y_all - prev_y, n_all - prev_n
+        bad = (dn < 0) | (dy < 0) | (dy > dn)
+        if bad.any():
+            e = stats.experiments[int(np.where(bad.any(axis=1))[0][0])]
+            raise ValueError(f"Experiment {e!r}: statistics must be cumulative and non-decreasing")
+        for k in range(1, A):
+            ll[:, k - 1] += conditional_log_lr(dy[:, 0], dn[:, 0], dy[:, k], dn[:, k], self._psi)
+        log_ek = logsumexp(ll + self._log_w[None, None, :], axis=2)                 # (E, A-1)
+        log_e = logsumexp(log_ek, axis=1) - np.log(A - 1)
+        for i, e in enumerate(stats.experiments):
+            self._state[e] = _State(y_all[i].copy(), n_all[i].copy(), ll[i].copy())
+        self._last_variant_log_e = dict(zip(stats.experiments, log_ek))
+        return log_e
+
+    def update(self, stats: Stats, look=None) -> pd.DataFrame:
+        """Assess each experiment on its cumulative statistics at this look. Experiments
+        already stopped keep their decision (their e-process stops too)."""
+        if stats.shape[2] != 1:
+            raise ValueError("Monitor takes one period per experiment: pass the cumulative totals")
+        look = (self.history["look"].nunique() if len(self.history) else 0) if look is None else look
+        lbf = self.analysis.log_bf10(stats)
+        active = np.array([e not in self._stopped for e in stats.experiments])
+        log_e = np.full(len(stats), np.nan)
+        if self.method == "exact" and active.any():
+            sub = stats.sel([e for e, a in zip(stats.experiments, active) if a])
+            log_e[active] = self._log_e(sub)
+        evidence = log_e if self.method == "exact" else lbf
+        rows = []
+        for i, e in enumerate(stats.experiments):
+            if e in self._stopped:
+                decision = self._stopped[e][0]
+            elif evidence[i] >= np.log(self.threshold):
+                decision = "effect"
+                self._stopped[e] = (decision, look)
+            elif self.futility is not None and -lbf[i] >= np.log(self.futility):
+                decision = "no effect"
+                self._stopped[e] = (decision, look)
+            else:
+                decision = "continue"
+            rows.append(dict(look=look, experiment=e, log_BF10=float(lbf[i]),
+                             log_E=float(log_e[i]), decision=decision))
+        out = pd.DataFrame(rows)
+        self.history = pd.concat([self.history, out], ignore_index=True) if len(self.history) else out
+        return out
+
+
+# ---------------------------------------------------------------------------
+# Calibration by simulation
+# ---------------------------------------------------------------------------
+
+def binomial_aa(baseline: float, users_per_look: int, arms: int = 2) -> Callable:
+    """A/A simulator for calibrate(): every arm converts at `baseline`."""
+    def simulate(rng, n_experiments, looks):
+        y = rng.binomial(users_per_look, baseline, size=(looks, n_experiments, arms)).cumsum(0)
+        n = users_per_look * np.arange(1, looks + 1)
+        return [Stats.from_arrays(n=np.full((n_experiments, arms), float(n[t])), converted=y[t])
+                for t in range(looks)]
+    return simulate
+
+
+def calibrate(make_monitor: Callable[[], Monitor], simulate: Callable, looks: int = 30,
+              n_experiments: int = 1000, seed: int = 0) -> dict:
+    """Share of A/A experiments a monitor stops for an effect, by simulation.
+
+    - `make_monitor`: returns a fresh Monitor (so each run starts with no state)
+    - `simulate`: simulate(rng, n_experiments, looks) -> list of cumulative Stats, one per
+      look; binomial_aa(...) gives one for conversion rates
+
+    Returns the false-positive rate, its 95% interval, and the bound 1/threshold."""
+    rng = np.random.default_rng(seed)
+    mon = make_monitor()
+    for t, st in enumerate(simulate(rng, n_experiments, looks)):
+        mon.update(st, look=t)
+    fp = sum(d == "effect" for d, _ in mon.decisions.values()) / n_experiments
+    se = np.sqrt(max(fp * (1 - fp), 1e-12) / n_experiments)
+    return {"false_positive_rate": fp, "interval_95": (max(0.0, fp - 1.96 * se), fp + 1.96 * se),
+            "bound": 1 / mon.threshold, "method": mon.method, "looks": looks, "experiments": n_experiments}

abfactor/spec.py

@@ -0,0 +1,66 @@
+"""Turn the pieces of an analysis into plain JSON-able data and back.
+
+Every metric, prior and pooling choice is a dataclass; `to_spec` records its type
+and fields, `from_spec` rebuilds it. This is what makes a saved Result reproducible:
+Analysis.from_settings(result.settings).fit(stats) recomputes it exactly.
+"""
+from __future__ import annotations
+
+import dataclasses
+from importlib import metadata
+
+import numpy as np
+
+from .brease import BreasePrior
+from .metrics import Count, Mean, PrePeriod, Proportion, Ratio, TimeToEvent
+from .pooling import Independent, SharedScale, SpikeAndSlab
+from .priors import Brease, DefaultFixedEffects, LearnedScale, NormalEffects
+
+REGISTRY = {cls.__name__: cls for cls in (
+    Proportion, TimeToEvent, Count, Mean, Ratio, PrePeriod,
+    DefaultFixedEffects, NormalEffects, LearnedScale, Brease, BreasePrior,
+    Independent, SharedScale, SpikeAndSlab)}
+
+
+def to_spec(obj):
+    if dataclasses.is_dataclass(obj) and not isinstance(obj, type):
+        name = type(obj).__name__
+        if name not in REGISTRY:
+            raise TypeError(f"Cannot record {name}: not a known abfactor type")
+        out = {"type": name}
+        for f in dataclasses.fields(obj):
+            if f.init:
+                out[f.name] = to_spec(getattr(obj, f.name))
+        return out
+    if isinstance(obj, np.ndarray):
+        return {"type": "ndarray", "values": obj.tolist()}
+    if isinstance(obj, (np.integer, np.floating)):
+        return obj.item()
+    if isinstance(obj, (list, tuple)):
+        return [to_spec(x) for x in obj]
+    if isinstance(obj, dict):
+        return {k: to_spec(v) for k, v in obj.items()}
+    return obj
+
+
+def from_spec(spec):
+    if isinstance(spec, dict) and "type" in spec:
+        if spec["type"] == "ndarray":
+            return np.asarray(spec["values"], dtype=float)
+        cls = REGISTRY[spec["type"]]
+        return cls(**{k: from_spec(v) for k, v in spec.items() if k != "type"})
+    if isinstance(spec, list):
+        return [from_spec(x) for x in spec]
+    if isinstance(spec, dict):
+        return {k: from_spec(v) for k, v in spec.items()}
+    return spec
+
+
+def versions() -> dict:
+    out = {}
+    for pkg in ("abfactor", "numpy", "scipy", "pandas", "xarray", "pymc"):
+        try:
+            out[pkg] = metadata.version(pkg)
+        except metadata.PackageNotFoundError:
+            out[pkg] = None
+    return out