abfactor 0.5.1__py3-none-any.whl

abfactor/__init__.py

@@ -0,0 +1,27 @@
+"""Bayesian A/B testing with Bayes factors.
+
+    from abfactor import Analysis, Stats, Proportion, DefaultFixedEffects, SpikeAndSlab
+
+    stats = Stats.from_users(users, metrics={"converted": "converted"}, control="control")
+    result = Analysis(Proportion("converted"), baseline=0.05,
+                      prior=DefaultFixedEffects(r=0.1)).fit(stats)
+    result.summary()
+
+Pieces: Stats (sufficient statistics) x Metric (Proportion, TimeToEvent, Count, Mean,
+Ratio) x prior (DefaultFixedEffects, NormalEffects, LearnedScale, Brease) x pooling
+(Independent, SharedScale, SpikeAndSlab) x engine ("quadrature", "nuts", "smc").
+Monitor runs sequential tests. See docs/design.md.
+"""
+from .analysis import Analysis, SmallSampleWarning
+from .metrics import Count, Mean, PrePeriod, Proportion, Ratio, TimeToEvent
+from .pooling import Independent, SharedScale, SpikeAndSlab
+from .priors import Brease, DefaultFixedEffects, LearnedScale, NormalEffects
+from .results import Result
+from .sequential import Monitor
+from .stats import Stats, cross, square
+
+__version__ = "0.5.1"
+
+__all__ = ["Analysis", "SmallSampleWarning", "Stats", "Result", "Monitor", "Proportion", "TimeToEvent", "Count", "Mean",
+           "Ratio", "PrePeriod", "DefaultFixedEffects", "NormalEffects", "LearnedScale", "Brease",
+           "Independent", "SharedScale", "SpikeAndSlab", "square", "cross"]

abfactor/aa.py

@@ -0,0 +1,75 @@
+"""Learning from A/A tests: how much identical arms really differ.
+
+In real experiments, arms that show users exactly the same thing still differ by more
+than sampling noise allows: traffic shifts during a test, repeat visitors, bots and
+logging all add variation between arms that more data does not average away. In the
+Upworthy Research Archive's exploratory set, identical arms differed by about 12% in
+click-through rate beyond binomial noise (docs/real_data.md). A model that ignores it
+reports effects that are not there.
+
+estimate_arm_sd measures that extra noise from your own A/A tests, as an SD on the
+metric's model scale; pass it to Analysis(arm_sd=...). aa_check reports how often an
+analysis finds evidence of an effect in A/A tests, which should be rare.
+"""
+from __future__ import annotations
+
+import numpy as np
+import pandas as pd
+from scipy.optimize import minimize_scalar
+
+from .analysis import Analysis
+from .design import sum_to_zero_basis
+from .metrics import Metric
+from .stats import Stats
+
+__all__ = ["estimate_arm_sd", "aa_check", "arm_noise_loglik"]
+
+
+def _contrasts(metric: Metric, stats: Stats):
+    d, v = metric.estimates(stats)
+    E, A, P = d.shape
+    Q = sum_to_zero_basis(A)
+    c = np.einsum("eap,ak->ekp", d, Q)                       # within-period contrasts, (E, A-1, P)
+    return c, v, Q
+
+
+def arm_noise_loglik(tau: float, metric: Metric, stats: Stats) -> float:
+    """Restricted log-likelihood of an extra per-arm SD tau: the contrasts Q'd within each
+    experiment and period are N(0, Q'(diag(v) + tau^2 I)Q) = N(0, Q' diag(v) Q + tau^2 I)."""
+    c, v, Q = _contrasts(metric, stats)
+    total = 0.0
+    E, K, P = c.shape
+    for p in range(P):
+        S = np.einsum("ak,ea,al->ekl", Q, v[..., p], Q) + tau ** 2 * np.eye(K)[None]
+        sign, logdet = np.linalg.slogdet(S)
+        quad = np.einsum("ek,ekl,el->e", c[..., p], np.linalg.inv(S), c[..., p])
+        total += float(-0.5 * (logdet + quad + K * np.log(2 * np.pi)).sum())
+    return total
+
+
+def estimate_arm_sd(metric: Metric, aa_stats: Stats, upper: float = 2.0) -> dict:
+    """Estimate the extra per-arm SD from A/A tests (all arms identical), by restricted
+    maximum likelihood on the contrasts, so each test's unknown common level drops out.
+
+    aa_stats may contain experiments with any number of identical arms as separate Stats;
+    pass a list to combine them. Returns the estimate, an approximate 95% interval from
+    the likelihood ratio, and the number of tests and arms used."""
+    parts = aa_stats if isinstance(aa_stats, (list, tuple)) else [aa_stats]
+    f = lambda t: -sum(arm_noise_loglik(t, metric, s) for s in parts)
+    best = minimize_scalar(f, bounds=(0.0, upper), method="bounded", options={"xatol": 1e-4})
+    tau, lmax = float(best.x), -float(best.fun)
+    grid = np.linspace(0, upper, 801)
+    ll = np.array([-f(t) for t in grid])
+    inside = grid[ll >= lmax - 1.92]                               # chi2(1) 95% / 2
+    return {"arm_sd": tau, "interval_95": (float(inside.min()), float(inside.max())),
+            "loglik_gain_vs_zero": lmax - (-f(0.0)),
+            "tests": int(sum(len(s) for s in parts)), "arms": int(sum(len(s) * s.shape[1] for s in parts))}
+
+
+def aa_check(analysis: Analysis, aa_stats: Stats) -> dict:
+    """How often an analysis finds evidence of a difference in A/A tests. For a calibrated
+    analysis, Bayes factors above 20 should be very rare, and evidence for no difference
+    (BF10 below 1/3) common."""
+    lbf = analysis.log_bf10(aa_stats)
+    return {"tests": len(lbf), "BF10>3": float(np.mean(lbf > np.log(3))),
+            "BF10>20": float(np.mean(lbf > np.log(20))), "BF10<1/3": float(np.mean(lbf < -np.log(3)))}

abfactor/analysis.py

@@ -0,0 +1,455 @@
+"""Analysis: metric x prior x pooling x engine."""
+from __future__ import annotations
+
+import logging
+import warnings
+from typing import Optional
+
+import numpy as np
+import pandas as pd
+import xarray as xr
+from scipy.special import logsumexp
+from scipy.stats import norm, skewnorm
+
+from . import brease as _brease
+from .design import Design
+from .metrics import Metric, Proportion
+from .pooling import Independent, SharedScale, SpikeAndSlab, pool
+from .priors import Brease, DefaultFixedEffects, EffectPrior, LearnedScale, NormalEffects
+from .results import Result
+from .stats import Stats
+
+OFFSETS = ("none", "fixed", "free", "shrunk")
+
+
+class SmallSampleWarning(UserWarning):
+    """The quadrature engine's normal approximation is unreliable for some experiments."""
+
+
+def small_sample_check(metric, stats) -> dict:
+    """Flag experiments whose smallest cell is too small for the normal approximation.
+
+    Below metric.warn_below, posterior intervals can be far too narrow: warn, and point to
+    the exact-likelihood engines. Below metric.note_below, intervals are fine but Bayes
+    factors understate the evidence for an effect (never overstate it, in the measurements
+    behind docs/accuracy.md): record a note for borderline calls."""
+    try:
+        m = metric.min_cell(stats)
+    except NotImplementedError:
+        return {}
+    exps = stats.experiments
+    unreliable = [e for e, v in zip(exps, m) if v < metric.warn_below]
+    conservative = [e for e, v in zip(exps, m) if metric.warn_below <= v < metric.note_below]
+    if unreliable:
+        warnings.warn(f"{len(unreliable)} experiment(s) have fewer than {metric.warn_below} in their "
+                      f"smallest cell ({unreliable[:5]}{'...' if len(unreliable) > 5 else ''}): the normal "
+                      "approximation is unreliable there; use engine='smc' or 'nuts'", SmallSampleWarning,
+                      stacklevel=3)
+    out = {"min_cell": dict(zip(map(str, exps), map(float, m)))}
+    if unreliable:
+        out["unreliable"] = unreliable
+    if conservative:
+        out["conservative_bf"] = conservative
+    return out
+ENGINES = ("quadrature", "nuts", "smc")
+
+
+class Analysis:
+    """An analysis of one kind of metric: what is measured, the prior on effects, how
+    experiments in a batch share information, and how the posterior is computed. `fit`
+    returns a `Result`; `log_bf10` returns Bayes factors alone; `learn_prior` turns completed
+    experiments into a prior for new ones.
+
+    - `metric`: what is measured (Proportion, TimeToEvent, Count, Mean, Ratio)
+    - `baseline`: prior centre for the control arm, on the natural scale (a conversion rate,
+      a rate per user, a mean)
+    - `prior`: DefaultFixedEffects(r) | NormalEffects(sd) | LearnedScale | Brease
+    - `pooling`: Independent() | SharedScale() | SpikeAndSlab()
+    - `level_sd`: prior SD of each period's baseline on the model scale. Default: 1.0 on the
+      log-odds, log-rate and hazard scales; for raw-scale metrics (Mean, Ratio) the
+      baseline's magnitude, so a revenue baseline of 50 gets SD 50
+    - `treatment_start`: first treated period; earlier periods are A/A periods
+    - `arm_sd`: SD of extra noise per arm and period on the model scale, beyond sampling
+      noise: identical arms in real A/A tests differ by more than the binomial allows
+      (time-varying traffic, repeat visitors, bots). Estimate it from your own A/A tests
+      with abfactor.aa.estimate_arm_sd. Default 0 (none).
+    - `offsets`: persistent arm offsets, identified by the A/A periods: "none", "fixed"
+      (N(0, offset_sd^2) on each projected offset), "free" (as wide as the level prior) or
+      "shrunk" (scale learned across the batch from prior scale offset_sd; NUTS only)
+    - `engine`: "quadrature" (default: fast, exact draws and pairwise BFs on the Gaussian
+      reduction), "nuts" (exact likelihood), "smc" (exact likelihood, omnibus BFs by
+      marginal likelihood; Independent only)
+    - `offset_sd`: prior SD of the projected arm offsets for offsets="fixed" (prior scale
+      for "shrunk")
+    - `grid_size`: points on the log g grid for quadrature
+    - `draws`: posterior draws per experiment (quadrature)
+    - `seed`: random seed; the same settings and data give the same draws
+    - `nuts`: options for engine="nuts" and "smc": draws, tune, chains, target_accept
+    """
+
+    def __init__(self, metric: Metric, baseline: Optional[float] = None,
+                 prior: EffectPrior = DefaultFixedEffects(0.1), pooling=Independent(),
+                 level_sd: Optional[float] = None, treatment_start: int = 0, offsets: str = "none",
+                 offset_sd: float = 0.05, engine: str = "quadrature", grid_size: int = 3001,
+                 draws: int = 4000, seed: int = 0, nuts: Optional[dict] = None, arm_sd: float = 0.0):
+        if arm_sd < 0:
+            raise ValueError("arm_sd must be non-negative")
+        self.arm_sd = float(arm_sd)
+        if offsets not in OFFSETS:
+            raise ValueError(f"offsets must be one of {OFFSETS}")
+        if engine not in ENGINES:
+            raise ValueError(f"engine must be one of {ENGINES}")
+        if offsets == "shrunk" and engine != "nuts":
+            raise ValueError("offsets='shrunk' needs engine='nuts'; use 'fixed' or 'free' with quadrature")
+        if baseline is None and not isinstance(prior, Brease):
+            raise ValueError("baseline is required")
+        self.metric, self.baseline, self.prior, self.pooling = metric, baseline, prior, pooling
+        if level_sd is None:
+            level_sd = max(abs(float(baseline)), 1e-6) if metric.scale == "raw" and baseline is not None else 1.0
+        self.level_sd, self.treatment_start = level_sd, treatment_start
+        self.offsets, self.offset_sd = offsets, offset_sd
+        self.engine, self.grid_size, self.draws, self.seed = engine, grid_size, draws, seed
+        self.nuts = {"draws": 1000, "tune": 1000, "chains": 4, "target_accept": 0.95, **(nuts or {})}
+
+    # ---- settings ------------------------------------------------------------
+
+    @property
+    def settings(self) -> dict:
+        """Everything needed to recompute a result, as plain JSON-able data."""
+        from .spec import to_spec, versions
+        return {"metric": to_spec(self.metric), "baseline": self.baseline, "prior": to_spec(self.prior),
+                "pooling": to_spec(self.pooling), "level_sd": self.level_sd,
+                "treatment_start": self.treatment_start, "offsets": self.offsets,
+                "offset_sd": self.offset_sd, "arm_sd": self.arm_sd, "engine": self.engine, "grid_size": self.grid_size,
+                "draws": self.draws, "seed": self.seed, "nuts": dict(self.nuts), "versions": versions()}
+
+    @classmethod
+    def from_settings(cls, settings: dict) -> "Analysis":
+        """Rebuild an Analysis from `Analysis.settings` or `Result.settings`."""
+        from .spec import from_spec
+        kw = {k: v for k, v in settings.items() if k != "versions"}
+        for key in ("metric", "prior", "pooling"):
+            kw[key] = from_spec(kw[key])
+        return cls(**kw)
+
+    # ---- public --------------------------------------------------------------
+
+    def fit(self, stats):
+        """Fit the analysis. `stats` is a Stats (experiments with the same number of arms and
+        periods), or a list of them, which returns a list of Results. With a list, pooling
+        spans every experiment in every part: a batch can mix two-arm and three-arm tests."""
+        if isinstance(stats, (list, tuple)):
+            return self._fit_parts(list(stats))
+        if isinstance(self.prior, Brease):
+            return self._fit_brease(stats)
+        design = self._design(stats)
+        if self.engine == "quadrature":
+            return self._fit_quadrature([stats], [design])[0]
+        if self.engine == "nuts":
+            return self._fit_nuts(stats, design)
+        return self._fit_smc(stats, design)
+
+    def _fit_parts(self, parts):
+        if not parts:
+            raise ValueError("fit needs at least one Stats")
+        if self.engine == "quadrature" and not isinstance(self.prior, Brease):
+            return self._fit_quadrature(parts, [self._design(st) for st in parts])
+        if isinstance(self.pooling, Independent):
+            return [self.fit(st) for st in parts]
+        raise NotImplementedError("pooling across a list of Stats needs engine='quadrature'")
+
+    def learn_prior(self, completed, pooling=None) -> LearnedScale:
+        """The effect scale learned from completed experiments, as a prior for new ones.
+        `completed` is a Stats, or a list of them when experiments have different numbers of
+        arms. Uses the slab of a spike-and-slab by default, so experiments with no effect
+        don't drag the scale towards zero. Set arm_sd first if your A/A tests show extra
+        noise between identical arms: otherwise that noise is learned as effects."""
+        if isinstance(self.prior, (Brease, LearnedScale)):
+            raise ValueError("learn_prior needs a DefaultFixedEffects or NormalEffects prior")
+        pooling = pooling or SpikeAndSlab()
+        parts = completed if isinstance(completed, (list, tuple)) else [completed]
+        u, log_pg = self.prior.grid(self.grid_size)
+        L = []
+        for st in parts:
+            design = self._design(st)
+            d, N = self._estimates(st, design)
+            L.append(design.log_S(d, N, u))
+        res = pool(np.vstack(L), log_pg, pooling)
+        n = sum(len(st) for st in parts)
+        return LearnedScale(u=u, log_mass=res["log_g_batch"],
+                            source=f"{n} completed experiments, {type(pooling).__name__}")
+
+    def log_bf10(self, stats) -> np.ndarray:
+        """Omnibus log BF10 per experiment only, without posterior draws (quadrature).
+        `stats` may be a list of Stats; the result is then concatenated in order."""
+        parts = stats if isinstance(stats, (list, tuple)) else [stats]
+        u, log_pg = self.prior.grid(self.grid_size)
+        L = []
+        for st in parts:
+            small_sample_check(self.metric, st)
+            design = self._design(st)
+            d, N = self._estimates(st, design)
+            L.append(design.log_S(d, N, u))
+        return pool(np.vstack(L), log_pg, self.pooling)["log_bf"]
+
+    # ---- shared ---------------------------------------------------------------
+
+    def _design(self, stats: Stats) -> Design:
+        E, A, P = stats.shape
+        offset_sd = {"none": None, "fixed": self.offset_sd, "free": self.level_sd,
+                     "shrunk": None}[self.offsets]
+        if self.offsets != "none" and self.treatment_start == 0:
+            raise ValueError("offsets need at least one pre-treatment period (treatment_start >= 1)")
+        return Design(A, P, m0=self.metric.to_model(self.baseline), level_sd=self.level_sd,
+                      treatment_start=self.treatment_start, offset_sd=offset_sd)
+
+    def _estimates(self, stats: Stats, design: Design):
+        """Estimates (E, A, P) and their full noise covariance (E, A*P, A*P)."""
+        d, v = self.metric.estimates(stats)
+        N = design.noise(v, self.metric.covariance(stats))
+        if self.arm_sd > 0:                                  # independent extra noise per arm and period
+            N = N + self.arm_sd ** 2 * np.eye(N.shape[-1])[None]
+        return d, N
+
+    def _prior_density_at_zero(self, cc: float = 2.0) -> float:
+        """Marginal prior density of a pairwise contrast at 0."""
+        if isinstance(self.prior, DefaultFixedEffects):
+            return 1 / (np.pi * np.sqrt(cc) * self.prior.r)          # Cauchy(0, sqrt(2) r)
+        if isinstance(self.prior, NormalEffects):
+            return float(norm.pdf(0, 0, self.prior.contrast_sd))
+        u, lm = self.prior.grid()
+        return float(np.exp(logsumexp(lm + norm.logpdf(0, 0, np.sqrt(cc * np.exp(u))))))
+
+    def _assemble(self, stats, design, theta_level, effect, omnibus, pairwise, pooling_info, diag):
+        """theta_level: (E, S) report-period level; effect: (E, A, S) arm effects (Q e*)."""
+        arms, ctrl = stats.arms, stats.control
+        eta = theta_level[:, None, :] + effect
+        natural = xr.DataArray(self.metric.to_natural(eta), dims=("experiment", "arm", "sample"),
+                               coords={"experiment": stats.experiments, "arm": arms})
+        contrasts = xr.DataArray(effect[:, 1:, :] - effect[:, :1, :], dims=("experiment", "variant", "sample"),
+                                 coords={"experiment": stats.experiments, "variant": arms[1:]})
+        return Result(metric=getattr(self.metric, "name", type(self.metric).__name__),
+                      scale=self.metric.scale, control=ctrl, arms=arms, experiments=stats.experiments,
+                      contrasts=contrasts, natural=natural, omnibus=omnibus, pairwise=pairwise,
+                      pooling=pooling_info, diagnostics=diag, settings=self.settings)
+
+    # ---- quadrature -------------------------------------------------------------
+
+    def _fit_quadrature(self, parts, designs):
+        """One pooled quadrature over every experiment in `parts`; one Result per part."""
+        rng = np.random.default_rng(self.seed)
+        u, log_pg = self.prior.grid(self.grid_size)
+        g = np.exp(u)
+        est = [self._estimates(st, design) for st, design in zip(parts, designs)]
+        res = pool(np.vstack([design.log_S(d, v, u) for (d, v), design in zip(est, designs)]), log_pg, self.pooling)
+        pooling_info = {"kind": type(self.pooling).__name__, "w_mean": res["w_mean"],
+                        "log_BF_any": res["log_bf_any"], "g_grid": g, "log_g_batch": res["log_g_batch"]}
+        out, start = [], 0
+        for stats, design, (d, v) in zip(parts, designs, est):
+            E, A, P = stats.shape
+            idx = slice(start, start + E)
+            start += E
+            r = {k: (res[k][idx] if res[k] is not None else None)
+                 for k in ("log_bf", "log_g_rest", "log_g_post", "p_effect")}
+            out.append(self._quadrature_result(stats, design, d, v, g, r, pooling_info, rng))
+        return out
+
+    def _quadrature_result(self, stats, design, d, v, g, res, pooling_info, rng):
+        E, A, P = stats.shape
+        S = self.draws
+        level, effect = np.empty((E, S)), np.empty((E, A, S))
+        rows = []
+        for e in range(E):
+            if res["p_effect"] is None:
+                theta = design.sample(d[e], v[e], g, res["log_g_post"][e], S, rng)
+            else:
+                slab = rng.random(S) < res["p_effect"][e]
+                theta = np.empty((S, design.X.shape[1]))
+                if slab.any():
+                    theta[slab] = design.sample(d[e], v[e], g, res["log_g_post"][e], int(slab.sum()), rng)
+                if (~slab).any():
+                    theta[~slab] = design.sample(d[e], v[e], np.array([1e-14]), np.array([0.0]),
+                                                 int((~slab).sum()), rng)
+                    theta[np.ix_(~slab, np.arange(theta.shape[1])[design.effect_rows()])] = 0.0   # no effect: exactly zero
+            level[e] = theta[:, design.report_period]
+            effect[e] = (theta[:, design.effect_rows()] @ design.Q.T).T
+            for k in range(1, A):
+                lbf, p_pos, mu, sd = design.pairwise(d[e], v[e], g, res["log_g_post"][e],
+                                                     res["log_g_rest"][e], k)
+                bf = np.exp(min(lbf, 700.0))                 # very strong evidence: BF capped, log_BF10 exact
+                rows.append(dict(experiment=stats.experiments[e], variant=stats.arms[k], log_BF10=lbf,
+                                 BF10=bf, BFplus0=2 * bf * p_pos, BFminus0=2 * bf * (1 - p_pos),
+                                 P_better=p_pos, contrast_mean=mu, contrast_sd=sd))
+        omnibus = pd.DataFrame({"experiment": stats.experiments, "log_BF10": res["log_bf"],
+                                "P_effect": res["p_effect"] if res["p_effect"] is not None else np.nan})
+        small = small_sample_check(self.metric, stats)
+        notes = ["Gaussian reduction of the likelihood on the model scale"]
+        if small.get("conservative_bf"):
+            notes.append(f"{len(small['conservative_bf'])} experiment(s) with smallest cell under "
+                         f"{self.metric.note_below}: Bayes factors there may understate the evidence for an "
+                         "effect; confirm borderline calls with engine='smc'")
+        diag = {"engine": "quadrature", "small_sample": small, "notes": notes}
+        return self._assemble(stats, design, level, effect, omnibus, pd.DataFrame(rows), dict(pooling_info), diag)
+
+    # ---- NUTS / SMC ----------------------------------------------------------------
+
+    def _pymc_model(self, stats, design, null=False, pooling=None):
+        import pymc as pm
+        import pytensor.tensor as pt
+        pooling = self.pooling if pooling is None else pooling
+        E, A, P = stats.shape
+        Q, on = design.Q, design.on
+        coords = {"experiment": range(E), "arm": range(A), "period": range(P), "contrast": range(A - 1)}
+        with pm.Model(coords=coords) as model:
+            level = pm.Normal("level", design.m0, self.level_sd, dims=("experiment", "period"))
+            eta0 = level[:, None, :]
+            if self.arm_sd > 0:                              # latent extra noise per arm and period
+                eta0 = eta0 + pm.Normal("arm_noise", 0, self.arm_sd, dims=("experiment", "arm", "period"))
+            if self.offsets != "none":
+                zo = pm.Normal("zo", 0, 1, dims=("experiment", "contrast"))
+                if self.offsets == "shrunk":
+                    so = pm.math.sqrt(pm.InverseGamma("g_offset", alpha=0.5, beta=self.offset_sd ** 2 / 2))
+                else:
+                    so = design.offset_sd
+                offset = pm.Deterministic("offset", (so * zo) @ Q.T, dims=("experiment", "arm"))
+                eta0 = eta0 + offset[:, :, None]
+            if null:
+                effect = pt.zeros((E, A))
+            else:
+                if isinstance(pooling, Independent):
+                    sc = self.prior.pymc_scale(dims="experiment")
+                    sc = sc[:, None] if hasattr(sc, "ndim") and getattr(sc, "ndim", 0) == 1 else sc
+                else:
+                    sc = self.prior.pymc_scale(dims=None)
+                ze = pm.Normal("ze", 0, 1, dims=("experiment", "contrast"))
+                effect = (sc * ze) @ Q.T
+            effect = pm.Deterministic("effect", effect, dims=("experiment", "arm"))
+            eta = eta0 + effect[:, :, None] * on[None, None, :]
+            ll = self.metric.pymc_loglik(eta, stats)
+            if isinstance(pooling, SpikeAndSlab) and not null:
+                w = pm.Beta("w", pooling.a_w, pooling.b_w)
+                ll0 = self.metric.pymc_loglik(eta0, stats)
+                slab, spike = pt.log(w) + ll, pt.log1p(-w) + ll0
+                pm.Potential("loglik", pt.logaddexp(slab, spike).sum())
+                pm.Deterministic("resp", pm.math.sigmoid(slab - spike), dims="experiment")
+            else:
+                pm.Potential("loglik", ll.sum())
+        return model
+
+    def _from_draws(self, stats, design, post, rng, omnibus_bf, diag):
+        st = lambda name: post[name].stack(sample=("chain", "draw"))
+        level = st("level").isel(period=design.report_period).transpose("experiment", "sample").values
+        effect = st("effect").transpose("experiment", "arm", "sample").values
+        E, A, S = effect.shape
+        p_eff = None
+        if "resp" in post.data_vars:
+            resp = st("resp").transpose("experiment", "sample").values
+            p_eff = resp.mean(axis=1)
+            effect = np.where((rng.random(resp.shape) < resp)[:, None, :], effect, 0.0)
+        rows = []
+        prior0 = self._prior_density_at_zero()
+        for e in range(E):
+            for k in range(1, A):
+                psi = effect[e, k] - effect[e, 0]
+                p_pos = float(np.clip((psi > 0).mean(), 1 / S, 1 - 1 / S))
+                if p_eff is None:
+                    bf = prior0 / skewnorm.pdf(0.0, *skewnorm.fit(psi))
+                    lbf = float(np.log(bf))
+                else:
+                    bf = lbf = np.nan                    # mixture with a point mass: use quadrature
+                rows.append(dict(experiment=stats.experiments[e], variant=stats.arms[k], log_BF10=lbf,
+                                 BF10=bf, BFplus0=2 * bf * p_pos, BFminus0=2 * bf * (1 - p_pos),
+                                 P_better=p_pos, contrast_mean=float(psi.mean()), contrast_sd=float(psi.std())))
+        omnibus = pd.DataFrame({"experiment": stats.experiments,
+                                "log_BF10": omnibus_bf if omnibus_bf is not None else np.nan,
+                                "P_effect": p_eff if p_eff is not None else np.nan})
+        pooling_info = {"kind": type(self.pooling).__name__,
+                        "w_mean": float(st("w").mean()) if "w" in post.data_vars else None}
+        return self._assemble(stats, design, level, effect, omnibus, pd.DataFrame(rows), pooling_info, diag)
+
+    def _fit_nuts(self, stats, design):
+        import pymc as pm
+        if isinstance(self.prior, LearnedScale):
+            raise NotImplementedError("LearnedScale priors need engine='quadrature'")
+        kw = dict(self.nuts)
+        with self._pymc_model(stats, design):
+            idata = pm.sample(kw.pop("draws"), tune=kw.pop("tune"), chains=kw.pop("chains"),
+                              random_seed=self.seed, progressbar=False, **kw)
+        div = int(idata.sample_stats["diverging"].sum())
+        if div:
+            logging.warning(f"{div} divergent transitions")
+        diag = {"engine": "nuts", "divergences": div,
+                "notes": ["omnibus BFs are not computed by NUTS: use engine='quadrature' or 'smc'"]}
+        self.idata_ = idata
+        return self._from_draws(stats, design, idata.posterior, np.random.default_rng(self.seed), None, diag)
+
+    def _fit_smc(self, stats, design):
+        import pymc as pm
+        if not isinstance(self.pooling, Independent):
+            raise ValueError("engine='smc' supports pooling=Independent() only")
+        if isinstance(self.prior, LearnedScale):
+            raise NotImplementedError("LearnedScale priors need engine='quadrature'")
+        lbf, posts = [], []
+        for e in stats.experiments:
+            sub = stats.sel([e])
+            lml = {}
+            for null in (False, True):
+                with self._pymc_model(sub, design, null=null, pooling=Independent()):
+                    smc = pm.sample_smc(self.nuts["draws"], chains=self.nuts["chains"],
+                                        random_seed=self.seed, progressbar=False)
+                lml[null] = _smc_log_ml(smc)
+                if not null:
+                    posts.append(smc.posterior)
+            lbf.append(lml[False] - lml[True])
+        post = xr.concat([p.to_dataset() if hasattr(p, "to_dataset") else p for p in posts], dim="experiment")
+        post = post.assign_coords(experiment=range(len(stats)))
+        diag = {"engine": "smc", "notes": ["omnibus BFs from SMC marginal likelihoods of the exact likelihood"]}
+        return self._from_draws(stats, design, post, np.random.default_rng(self.seed), np.array(lbf), diag)
+
+    # ---- BREASE ---------------------------------------------------------------------
+
+    def _fit_brease(self, stats):
+        E, A, P = stats.shape
+        if not isinstance(self.metric, Proportion) or A != 2 or P != 1:
+            raise ValueError("The BREASE prior needs a Proportion metric, two arms and one period")
+        y, n = stats.get(self.metric.name)[..., 0], stats.get("n")[..., 0]
+        adverse = (n - y) if self.prior.outcome == "success" else y
+        rng = np.random.default_rng(self.seed)
+        S = self.draws
+        level, effect, rows, om = np.empty((E, S)), np.empty((E, 2, S)), [], []
+        for e in range(E):
+            data = _brease.Data(y0=int(adverse[e, 0]), N0=int(n[e, 0]), y1=int(adverse[e, 1]), N1=int(n[e, 1]))
+            lbf = _brease.log_bayes_factor_10(data, self.prior.prior)
+            post = _brease.sample_posterior(data, self.prior.prior, S, rng)
+            p0, p1 = post["theta0"], post["theta1"]
+            if self.prior.outcome == "success":
+                p0, p1 = 1 - p0, 1 - p1
+            lo0, lo1 = np.log(p0 / (1 - p0)), np.log(p1 / (1 - p1))
+            level[e] = (lo0 + lo1) / 2
+            effect[e] = np.stack([lo0 - level[e], lo1 - level[e]])
+            os_ = _brease.one_sided_bayes_factors(data, self.prior.prior, size=200_000, rng=rng)
+            # "better" = the counted statistic is higher in the variant
+            up, down = ("BF_minus_0", "BF_plus_0") if self.prior.outcome == "success" else ("BF_plus_0", "BF_minus_0")
+            psi = lo1 - lo0
+            rows.append(dict(experiment=stats.experiments[e], variant=stats.arms[1], log_BF10=lbf,
+                             BF10=float(np.exp(lbf)), BFplus0=os_[up], BFminus0=os_[down],
+                             P_better=float((p1 > p0).mean()), contrast_mean=float(psi.mean()),
+                             contrast_sd=float(psi.std())))
+            om.append(lbf)
+        omnibus = pd.DataFrame({"experiment": stats.experiments, "log_BF10": om, "P_effect": np.nan})
+        design = Design(2, 1, m0=0.0)
+        diag = {"engine": "brease", "notes": ["exact BREASE posterior and marginal likelihoods"]}
+        return self._assemble(stats, design, level, effect, omnibus, pd.DataFrame(rows),
+                              {"kind": "Independent"}, diag)
+
+
+def _smc_log_ml(idata) -> float:
+    """Log marginal likelihood from pm.sample_smc output, averaged over chains. Stored per
+    chain as per-stage values, NaN-padded; the last finite value is the total."""
+    v = idata.sample_stats["log_marginal_likelihood"].values
+    per_chain = []
+    for chain in v:
+        flat = np.concatenate([np.atleast_1d(np.asarray(x, dtype=float)).ravel()
+                               for x in np.atleast_1d(chain).ravel()])
+        per_chain.append(flat[np.isfinite(flat)][-1])
+    return float(logsumexp(per_chain) - np.log(len(per_chain)))