abfactor 0.5.1__py3-none-any.whl

abfactor/anova.py

@@ -0,0 +1,197 @@
+"""
+Default (g-prior) ANOVA Bayes factors, Rouder, Morey, Speckman & Province (2012),
+computed from per-group sufficient statistics and extended to many experiments.
+
+Within one experiment with a variations (control + variants), Section 6 of the paper
+gives, with Jeffreys priors on the grand mean and error variance,
+
+    y = mu 1 + sigma X theta + eps,   theta | g ~ N(0, g I),   g ~ r^2 Inv-chi2(1)
+    BF10 = integral over g of S(g) pi(g),
+    S(g) = |G|^{-1/2} |V_g|^{-1/2} (SST / (SST - c' V_g^{-1} c))^{(N-1)/2}
+
+theta is a vector of standardised effects (effect / sigma), so the prior on each
+is Cauchy(0, r). Everything in S(g) depends on the data only through the group
+counts, sums and sums of squares, so it can be computed from aggregated data:
+for a per-user metric those are population, sum over users and sum of squares
+over users (the same columns as GaussianMeanModel in models.py).
+
+Multiple experiments
+--------------------
+Stacking experiments into one design matrix doesn't work: the model has one mu
+and one sigma for all the data, so it would need experiment as a factor too,
+assume every experiment has the same noise level, and treat "variant 1" in
+different experiments as the same level of a crossed factor when it's nested.
+Instead, each experiment keeps its own mu_e and sigma_e (with Jeffreys priors),
+and so its own S_e(g). Experiments are then connected through what's shared:
+
+* independent_bfs: each experiment has its own g. BF_e = int S_e(g) pi(g) dg.
+  This is exactly the paper's one-way ANOVA BF, once per experiment.
+* SharedScaleModel: one g for all experiments (standardised effects in your
+  experiments are on a common scale), plus a spike-and-slab across
+  experiments: each experiment has no effect with probability 1 - w and
+  effects of scale g with probability w, w ~ Beta(a_w, b_w). The model learns
+  from the batch how big effects usually are and how often there are any, and
+  gives, for every experiment, its Bayes factor given all the others and its
+  posterior probability of having an effect.
+
+Fixed vs random variation effects
+---------------------------------
+Variations are fixed effects (they're specific treatments, not a sample from a
+population), so by default the paper's sum-to-zero projection (Section 7.2) is
+used: a - 1 effects with Q the eigenvectors of I - J/a. For two groups this
+gives the JZS two-sample t-test Bayes factor with Cauchy scale sqrt(2) * r on
+the standardised difference.
+
+Choosing r
+----------
+The paper's defaults (r = 1, or 0.5 in the BayesFactor package's "medium")
+assume standardised effects of order 0.5. In online experiments a
+standardised effect is (difference in means) / (SD between users), which is
+usually tiny: a 3% lift in revenue per user with a CV of 4 is 0.0075. With
+r = 0.5 the alternative spends almost all its prior mass on effects
+hundreds of times larger than anything plausible, and the Bayes factor
+favours the null almost regardless of the data (Lindley's paradox). Pick r on
+the standardised scale your experiments actually live on, e.g. from history,
+or let SharedScaleModel learn g from the batch.
+"""
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Optional
+
+import numpy as np
+from scipy.special import logsumexp
+
+from .pooling import SpikeAndSlab, pool
+
+__all__ = ["ExperimentStats", "experiment_stats", "bf10", "independent_bfs", "SharedScaleModel", "stats_from_wide",
+           "log_S", "log_prior_g"]
+
+
+# ---------------------------------------------------------------------------
+# One experiment: log S(g) from sufficient statistics
+# ---------------------------------------------------------------------------
+
+def _sum_to_zero_basis(a: int) -> np.ndarray:
+    """a x (a-1) orthonormal basis of the sum-to-zero subspace (paper's Q_a)."""
+    w, v = np.linalg.eigh(np.eye(a) - np.ones((a, a)) / a)
+    return v[:, w > 0.5]
+
+
+@dataclass
+class ExperimentStats:
+    """Precomputed pieces of S(g) for one experiment."""
+    N: float
+    sst: float
+    lam: np.ndarray      # eigenvalues of X~'X~ (in the effect parameterisation)
+    proj2: np.ndarray    # squared projections of X~'y~ onto the eigenvectors
+    p: int               # number of effect parameters
+
+
+def experiment_stats(n, s, ss, fixed: bool = True) -> ExperimentStats:
+    """n, s, ss: per-group counts, sums and sums of squares (length a)."""
+    n, s, ss = (np.asarray(x, dtype=float) for x in (n, s, ss))
+    N = n.sum()
+    ybar = s.sum() / N
+    sst = ss.sum() - N * ybar ** 2
+    if sst <= 0:
+        raise ValueError("Total sum of squares must be positive")
+    M = np.diag(n) - np.outer(n, n) / N          # X~'X~ for the cell-indicator design
+    c = s - n * ybar                              # X~'y~
+    if fixed:
+        Q = _sum_to_zero_basis(len(n))
+        M, c = Q.T @ M @ Q, Q.T @ c
+    lam, U = np.linalg.eigh(M)
+    return ExperimentStats(N, sst, np.clip(lam, 0, None), (U.T @ c) ** 2, M.shape[0])
+
+
+def log_S(st: ExperimentStats, g: np.ndarray) -> np.ndarray:
+    """log S(g) on an array of g values (paper eq. 8 integrand, without pi(g))."""
+    g = np.asarray(g, dtype=float)[..., None]
+    d = st.lam + 1.0 / g                                   # eigenvalues of V_g
+    quad = np.sum(st.proj2 / d, axis=-1)                   # c' V_g^{-1} c
+    log_det = np.sum(np.log(d), axis=-1)
+    resid = np.maximum(st.sst - quad, st.sst * 1e-300)
+    return (-0.5 * st.p * np.log(g[..., 0]) - 0.5 * log_det
+            + 0.5 * (st.N - 1) * (np.log(st.sst) - np.log(resid)))
+
+
+def log_prior_g(g: np.ndarray, r: float) -> np.ndarray:
+    """log density of g ~ r^2 Inv-chi2(1), i.e. Cauchy(0, r) on each standardised effect."""
+    return np.log(r) - 0.5 * np.log(2 * np.pi) - 1.5 * np.log(g) - r ** 2 / (2 * g)
+
+
+def _log_g_grid(r: float, stats: list, size: int) -> np.ndarray:
+    """Grid on log g covering the prior and every experiment's likelihood peak."""
+    lo, hi = np.log(r ** 2) - 30, np.log(r ** 2) + 12
+    probe = np.linspace(-60, 20, 2001)
+    for st in stats:
+        peak = probe[np.argmax(log_S(st, np.exp(probe)))]
+        lo, hi = min(lo, peak - 8), max(hi, peak + 8)
+    return np.linspace(lo, hi, size)
+
+
+def bf10(n, s, ss, r: float = 0.5, fixed: bool = True, grid_size: int = 20001) -> float:
+    """log BF10 for 'some variation differs' within one experiment."""
+    st = experiment_stats(n, s, ss, fixed)
+    u = _log_g_grid(r, [st], grid_size)
+    g = np.exp(u)
+    integrand = log_S(st, g) + log_prior_g(g, r) + u          # dg = g du
+    return float(logsumexp(integrand) + np.log(u[1] - u[0]))
+
+
+def independent_bfs(stats_by_expt, r: float = 0.5, fixed: bool = True) -> np.ndarray:
+    """log BF10 per experiment, each with its own g. stats_by_expt: list of (n, s, ss)."""
+    return np.array([bf10(n, s, ss, r, fixed) for n, s, ss in stats_by_expt])
+
+
+# ---------------------------------------------------------------------------
+# Many experiments: shared effect scale + spike-and-slab across experiments
+# ---------------------------------------------------------------------------
+
+class SharedScaleModel:
+    """One g shared by all experiments, each experiment null with probability 1 - w.
+
+        p(data | g, w) / p(data | all null) = prod_e [ (1 - w) + w S_e(g) ]
+        g ~ r^2 Inv-chi2(1),  w ~ Beta(a_w, b_w)
+
+    fit() returns, per experiment:
+
+    - `log_BF10_given_rest`: log of int S_e(g) p(g | other experiments) dg, the Bayes
+        factor for this experiment's effect, with the effect scale informed by the rest of
+        the batch
+    - `P_effect`: posterior probability this experiment has an effect
+
+    and for the batch: posterior mean of w, posterior of g on the grid, and
+    log BF of "effects exist somewhere" vs "no experiment has an effect".
+    """
+
+    def __init__(self, r: float = 0.5, a_w: float = 1.0, b_w: float = 1.0,
+                 fixed: bool = True, g_grid_size: int = 3001, w_grid_size: int = 201):
+        self.r, self.a_w, self.b_w, self.fixed = r, a_w, b_w, fixed
+        self.g_grid_size, self.w_grid_size = g_grid_size, w_grid_size
+
+    def fit(self, stats_by_expt) -> dict:
+        """Fit the shared-scale model to a list of (n, s, ss) per experiment; returns per-experiment and batch results."""
+        stats = [experiment_stats(n, s, ss, self.fixed) for n, s, ss in stats_by_expt]
+        u = _log_g_grid(self.r, stats, self.g_grid_size)
+        L = np.stack([log_S(st, np.exp(u)) for st in stats])            # (E, K)
+        g = np.exp(u)
+        log_pg = log_prior_g(g, self.r) + u + np.log(u[1] - u[0])
+        r = pool(L, log_pg, SpikeAndSlab(self.a_w, self.b_w, self.w_grid_size))
+        return dict(log_BF10_given_rest=r["log_bf"], P_effect=r["p_effect"],
+                    log_BF_any_effect=r["log_bf_any"], w_mean=r["w_mean"], g_grid=g,
+                    g_posterior=np.exp(r["log_g_batch"]))
+
+
+def stats_from_wide(data, numerator: str = "money", population: str = "population",
+                    n_groups: Optional[int] = None):
+    """(n, s, ss) per experiment from the wide frame used by models.py:
+    population_<g>, <numerator>_<g>, <numerator>_sq_<g> (per-user sums)."""
+    import re
+    def cols(prefix):
+        pat = re.compile(rf"^{re.escape(prefix)}_(\d+)$")
+        found = sorted((int(m.group(1)), c) for c in data.columns if (m := pat.match(str(c))))
+        return data[[c for _, c in found]].to_numpy(dtype=float)
+    n, s, ss = cols(population), cols(numerator), cols(f"{numerator}_sq")
+    return [(n[e], s[e], ss[e]) for e in range(len(data))]

abfactor/baselines.py

@@ -0,0 +1,142 @@
+"""
+Baselines from Gronau, Raj & Wagenmakers (2021), for comparison with BREASE:
+the two approaches in Gronau, Raj & Wagenmakers (2021),
+"Informed Bayesian Inference for the A/B Test":
+
+  * Independent beta ("separate" priors per arm, as in bayesAB / BayesianFirstAid)
+  * abtest's logit-transform model (Kass & Vaidyanathan 1992):
+        logit p1 = beta - psi/2,  logit p2 = beta + psi/2,
+        beta ~ N(mu_b, s_b^2),  psi ~ N(mu_psi, s_psi^2)  (truncated for H+/H-)
+
+LT marginal likelihoods are computed by brute-force 2-D quadrature on a grid
+centred on the posterior mode (abtest uses Laplace + importance sampling).
+All functions take data in abtest's success coding: y1/n1 control, y2/n2 treatment.
+"""
+import numpy as np
+from scipy.special import betaln, gammaln, expit, log_expit, logsumexp
+from scipy.stats import norm
+from scipy.optimize import minimize, least_squares
+
+from .brease import Data, BreasePrior, log_bayes_factor_10, one_sided_bayes_factors
+
+__all__ = ["lt_log_ml", "lt_log_ml_h0", "lt_bfs", "ib_bf10", "ib_one_sided", "brease_bfs", "elicit_lt",
+           "elicit_brease"]
+
+
+def lchoose(n, k):
+    return gammaln(n + 1) - gammaln(k + 1) - gammaln(n - k + 1)
+
+
+# ---------------------------------------------------------------- LT / abtest
+def _lt_loglik(b, psi, y1, n1, y2, n2):
+    e1, e2 = b - psi / 2, b + psi / 2
+    return (lchoose(n1, y1) + y1 * log_expit(e1) + (n1 - y1) * log_expit(-e1)
+            + lchoose(n2, y2) + y2 * log_expit(e2) + (n2 - y2) * log_expit(-e2))
+
+
+def lt_log_ml(y1, n1, y2, n2, mu_b=0, s_b=1, mu_psi=0, s_psi=1,
+              side=None, n_grid=801, width=10):
+    """log p(data | H1), or H+ (side=+1: psi>0) / H- (side=-1: psi<0)."""
+    f = lambda x: -(_lt_loglik(x[0], x[1], y1, n1, y2, n2)
+                    + norm.logpdf(x[0], mu_b, s_b) + norm.logpdf(x[1], mu_psi, s_psi))
+    opt = minimize(f, [0.0, 0.0], method="BFGS")
+    sd = np.sqrt(np.diag(opt.hess_inv)) * width
+    bg = np.linspace(opt.x[0] - sd[0], opt.x[0] + sd[0], n_grid)
+    pg = np.linspace(opt.x[1] - sd[1], opt.x[1] + sd[1], n_grid)
+    if side is not None:   # make sure the grid covers the constrained region
+        pg = np.linspace(0, max(pg.max(), 4 * s_psi + abs(mu_psi)), n_grid) * side
+        pg = np.sort(pg)
+    B, P = np.meshgrid(bg, pg, indexing="ij")
+    lp = (_lt_loglik(B, P, y1, n1, y2, n2)
+          + norm.logpdf(B, mu_b, s_b) + norm.logpdf(P, mu_psi, s_psi))
+    if side is not None:   # truncated-normal normaliser
+        mass = norm.sf(0, mu_psi, s_psi) if side > 0 else norm.cdf(0, mu_psi, s_psi)
+        lp = lp - np.log(mass)
+    return logsumexp(lp) + np.log(bg[1] - bg[0]) + np.log(abs(pg[1] - pg[0]))
+
+
+def lt_log_ml_h0(y1, n1, y2, n2, mu_b=0, s_b=1):
+    """Log marginal likelihood of the logit-transform model under H0 (psi = 0), by quadrature."""
+    b = np.linspace(-15, 15, 200_001)
+    lp = _lt_loglik(b, 0.0, y1, n1, y2, n2) + norm.logpdf(b, mu_b, s_b)
+    return logsumexp(lp) + np.log(b[1] - b[0])
+
+
+def lt_bfs(y1, n1, y2, n2, **prior):
+    """BF10, BF+0 and BF-0 for the logit-transform (abtest) model."""
+    mu_b, s_b = prior.get("mu_b", 0), prior.get("s_b", 1)
+    l0 = lt_log_ml_h0(y1, n1, y2, n2, mu_b, s_b)
+    return {k: float(np.exp(lt_log_ml(y1, n1, y2, n2, side=s, **prior) - l0))
+            for k, s in (("BF10", None), ("BF+0", +1), ("BF-0", -1))}
+
+
+# ------------------------------------------------------- independent beta (IB)
+def ib_bf10(y1, n1, y2, n2, a=1.0):
+    """Two-sided IB Bayes factor via Savage-Dickey (Dablander et al. 2022; BREASE App. F)."""
+    lbf = (betaln(2 * a - 1, 2 * a - 1) + betaln(a + y1, a + n1 - y1)
+           + betaln(a + y2, a + n2 - y2)
+           - betaln(2 * a + y1 + y2 - 1, 2 * a + n1 - y1 + n2 - y2 - 1)
+           - 2 * betaln(a, a))
+    return float(np.exp(lbf))
+
+
+def ib_one_sided(y1, n1, y2, n2, a=1.0, size=2_000_000, rng=0):
+    """IB one-sided BFs: H+ is the IB prior truncated to p2>p1 (prior mass 1/2)."""
+    rng = np.random.default_rng(rng)
+    p1 = rng.beta(a + y1, a + n1 - y1, size)
+    p2 = rng.beta(a + y2, a + n2 - y2, size)
+    post_plus = np.mean(p2 > p1)
+    bf10 = ib_bf10(y1, n1, y2, n2, a)
+    return {"BF10": bf10, "BF+0": bf10 * post_plus / 0.5, "BF-0": bf10 * (1 - post_plus) / 0.5}
+
+
+# ------------------------------------------------------------------- BREASE
+def to_brease(y1, n1, y2, n2):
+    """abtest success coding -> BREASE adverse-outcome coding (y = failures)."""
+    return Data(y0=n1 - y1, N0=n1, y1=n2 - y2, N1=n2)
+
+
+def brease_bfs(y1, n1, y2, n2, prior=BreasePrior(), rng=0):
+    """BF10 and one-sided Bayes factors under a BREASE prior, in the baselines' format."""
+    d = to_brease(y1, n1, y2, n2)
+    os_ = one_sided_bayes_factors(d, prior, size=400_000, rng=rng)
+    # BREASE "benefit" = fewer adverse outcomes = abtest H+
+    return {"BF10": float(np.exp(log_bayes_factor_10(d, prior))),
+            "BF+0": os_["BF_minus_0"], "BF-0": os_["BF_plus_0"]}
+
+
+# ------------------------------------------------- prior elicitation helpers
+def elicit_lt(q, prob, mu_b=0, s_b=1, size=200_000, rng=0):
+    """Mimic abtest::elicit_prior(what='arisk'): fit (mu_psi, s_psi) so the
+    implied prior on p2 - p1 matches the given quantiles (eq. 2)."""
+    z = np.random.default_rng(rng).standard_normal((2, size))
+    b = mu_b + s_b * z[0]
+
+    def resid(th):
+        psi = th[0] + np.exp(th[1]) * z[1]
+        ar = expit(b + psi / 2) - expit(b - psi / 2)
+        return [np.mean(ar <= qi) - pi for qi, pi in zip(q, prob)]
+    th = least_squares(resid, [0.5, np.log(0.3)], diff_step=1e-2).x
+    return th[0], float(np.exp(th[1]))
+
+
+def elicit_brease(q, prob, mu0=0.5, n0=2, mus=0.02, ns=1, size=200_000, rng=0):
+    """Same quantiles, but fit the *efficacy* prior (mu_e, n_e) of BREASE,
+    with a small side-effect prior fixed. In abtest coding the absolute risk
+    p2 - p1 = theta0 - theta1 (control failure rate minus treated failure rate)."""
+    u = np.random.default_rng(rng).random((3, size))
+    from scipy.stats import beta as B
+    t0 = B.ppf(u[0], mu0 * n0, (1 - mu0) * n0)
+    es = B.ppf(u[2], mus * ns, (1 - mus) * ns)
+
+    def resid(th):
+        mue, ne = expit(th[0]), np.exp(th[1])
+        ee = B.ppf(u[1], mue * ne, (1 - mue) * ne)
+        t1 = (1 - ee) * t0 + es * (1 - t0)
+        return [np.mean(t0 - t1 <= qi) - pi for qi, pi in zip(q, prob)]
+    th = least_squares(resid, [-1.0, np.log(5.0)], diff_step=1e-2).x
+    return BreasePrior(mu0, float(expit(th[0])), mus, n0, float(np.exp(th[1])), ns)
+
+
+def fmt(d):
+    return "  ".join(f"{k}={v:8.4f}" for k, v in d.items())

abfactor/brease.py

@@ -0,0 +1,349 @@
+"""
+BREASE: Bayesian analysis of binary experiments with causally sound priors.
+
+Implementation of Irons & Cinelli (2023), "Causally Sound Priors for Binary
+Experiments". The likelihood is parameterised by
+
+    theta0 : baseline risk          P(Y(0) = 1)
+    eta_e  : efficacy               P(Y(1) = 0 | Y(0) = 1)
+    eta_s  : risk of side effects   P(Y(1) = 1 | Y(0) = 0)
+
+with theta1 = (1 - eta_e) * theta0 + eta_s * (1 - theta0)       (eq. 3.3)
+
+and independent priors  theta0 ~ Beta*(mu0, n0), eta_e ~ Beta*(mue, ne),
+eta_s ~ Beta*(mus, ns), where Beta*(mu, n) = Beta(mu * n, (1 - mu) * n).
+
+Everything is computed in log space. The posterior is a finite mixture of
+independent betas indexed by the latent counterfactual counts
+
+    j = y1(1)  treated subjects with the outcome who'd have had it anyway
+    k = x1(1)  treated subjects without the outcome who'd have had it untreated
+
+so sampling is exact (no MCMC) and the marginal likelihood is a finite sum.
+The mixture grid has (y1 + 1) x (N1 - y1 + 1) cells, which is cheap even for
+trials with tens of thousands of participants.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+import numpy as np
+from scipy.special import betaln, gammaln, logsumexp
+
+__all__ = ["Data", "BreasePrior", "bayes_factor_10", "log_bayes_factor_10", "one_sided_bayes_factors",
+           "sample_posterior", "sample_posterior_no_harm", "sample_posterior_no_benefit",
+           "log_marginal_likelihood_M1", "log_marginal_likelihood_M0", "log_marginal_likelihood_no_harm",
+           "log_marginal_likelihood_no_benefit", "partial_id_bounds", "theta1_from", "sensitivity_grid",
+           "prior_moments", "summarise"]
+
+
+# ---------------------------------------------------------------------------
+# Data containers
+# ---------------------------------------------------------------------------
+
+@dataclass(frozen=True)
+class Data:
+    """Observed trial: y0 events of N0 in control, y1 events of N1 in treatment."""
+    y0: int
+    N0: int
+    y1: int
+    N1: int
+
+    def __post_init__(self):
+        if not (0 <= self.y0 <= self.N0 and 0 <= self.y1 <= self.N1):
+            raise ValueError("Need 0 <= y <= N in each arm.")
+
+    @property
+    def N(self) -> int:
+        """Total participants."""
+        return self.N0 + self.N1
+
+
+@dataclass(frozen=True)
+class BreasePrior:
+    """BREASE(mu0, mue, mus; n0, ne, ns). Defaults are the paper's default prior."""
+    mu0: float = 0.5
+    mue: float = 0.3
+    mus: float = 0.3
+    n0: float = 2.0
+    ne: float = 1.0
+    ns: float = 1.0
+
+    # Beta(a, b) shape parameters
+    @property
+    def a0(self):
+        """Beta parameters of the three components: control rate (a0, b0), effect (ae, be),
+        side effect (as_, bs), from each mean mu and prior sample size n."""
+        return self.mu0 * self.n0
+    @property
+    def b0(self): return (1 - self.mu0) * self.n0
+    @property
+    def ae(self): return self.mue * self.ne
+    @property
+    def be(self): return (1 - self.mue) * self.ne
+    @property
+    def as_(self): return self.mus * self.ns
+    @property
+    def bs(self): return (1 - self.mus) * self.ns
+
+    @classmethod
+    def dirichlet(cls, mu0, mue, mus, n0):
+        """Dirichlet prior on response types = BREASE with 'equal confidence' (eq. 3.6)."""
+        return cls(mu0, mue, mus, n0, mu0 * n0, (1 - mu0) * n0)
+
+    @classmethod
+    def empirical_bayes(cls, data: Data, n: float = 1.0):
+        """BREASE-EB(n) from Appendix E.3: prior means at midpoints of the
+        partially identified intervals, evaluated at (y+1)/(N+2)."""
+        t0 = (data.y0 + 1) / (data.N0 + 2)
+        t1 = (data.y1 + 1) / (data.N1 + 2)
+        (le, ue), (ls, us) = partial_id_bounds(t0, t1)
+        return cls(0.5, (le + ue) / 2, (ls + us) / 2, 2.0, n, n)
+
+    def sample(self, size, rng=None):
+        """Draws from the prior: a tuple of arrays (theta0, eta_e, eta_s)."""
+        rng = np.random.default_rng(rng)
+        t0 = rng.beta(self.a0, self.b0, size)
+        ee = rng.beta(self.ae, self.be, size)
+        es = rng.beta(self.as_, self.bs, size)
+        return t0, ee, es
+
+
+# ---------------------------------------------------------------------------
+# Helpers
+# ---------------------------------------------------------------------------
+
+def _lchoose(n, k):
+    return gammaln(n + 1) - gammaln(k + 1) - gammaln(n - k + 1)
+
+
+def theta1_from(theta0, eta_e, eta_s):
+    """Eq. (3.3)."""
+    return (1 - eta_e) * theta0 + eta_s * (1 - theta0)
+
+
+def partial_id_bounds(theta0, theta1):
+    """Sharp bounds on (eta_e, eta_s) given (theta0, theta1), eqs. (E.3)-(E.4)."""
+    theta0 = np.asarray(theta0, float)
+    theta1 = np.asarray(theta1, float)
+    le = np.maximum(0, (theta0 - theta1) / theta0)
+    ue = np.minimum(1, (1 - theta1) / theta0)
+    ls = np.maximum(0, (theta1 - theta0) / (1 - theta0))
+    us = np.minimum(1, theta1 / (1 - theta0))
+    return (le, ue), (ls, us)
+
+
+def _mixture_log_weights(d: Data, p: BreasePrior):
+    """Unnormalised log mixture weights over the (j, k) grid, plus the
+    posterior beta parameters for each cell (eqs. 3.11-3.12, with j = y1 - y1(0)).
+
+    The returned weights already include the prior normalising constants, so
+    logsumexp(weights) + binomial coefficients = log marginal likelihood (3.16).
+    """
+    j = np.arange(d.y1 + 1)[:, None]          # y1(1)
+    k = np.arange(d.N1 - d.y1 + 1)[None, :]   # x1(1)
+
+    # posterior beta parameters for each mixture component
+    A0 = d.y0 + j + k + p.a0
+    B0 = d.N - (d.y0 + j + k) + p.b0
+    Ae = k + p.ae
+    Be = j + p.be
+    As = d.y1 - j + p.as_
+    Bs = d.N1 - d.y1 - k + p.bs
+
+    logw = (
+        _lchoose(d.y1, j) + _lchoose(d.N1 - d.y1, k)
+        + betaln(A0, B0) - betaln(p.a0, p.b0)
+        + betaln(Ae, Be) - betaln(p.ae, p.be)
+        + betaln(As, Bs) - betaln(p.as_, p.bs)
+    )
+    params = np.broadcast_arrays(A0, B0, Ae, Be, As, Bs)
+    return logw, params
+
+
+def _log_binom_coef(d: Data):
+    return _lchoose(d.N0, d.y0) + _lchoose(d.N1, d.y1)
+
+
+# ---------------------------------------------------------------------------
+# Marginal likelihoods (Theorem 3 and Appendix E.2)
+# ---------------------------------------------------------------------------
+
+def log_marginal_likelihood_M1(d: Data, p: BreasePrior) -> float:
+    """Unconstrained BREASE model M1, eq. (3.16)."""
+    logw, _ = _mixture_log_weights(d, p)
+    return float(_log_binom_coef(d) + logsumexp(logw))
+
+
+def log_marginal_likelihood_M0(d: Data, p: BreasePrior) -> float:
+    """Null model M0: theta0 ~ Beta*(mu0, n0), theta1 = theta0 (beta-binomial)."""
+    s = d.y0 + d.y1
+    return float(_log_binom_coef(d)
+                 + betaln(s + p.a0, d.N - s + p.b0) - betaln(p.a0, p.b0))
+
+
+def log_marginal_likelihood_no_harm(d: Data, p: BreasePrior) -> float:
+    """M'_-: eta_s = 0, so theta1 = (1 - eta_e) theta0 (Appendix E.2)."""
+    k = np.arange(d.N1 - d.y1 + 1)
+    logw = (_lchoose(d.N1 - d.y1, k)
+            + betaln(d.y0 + d.y1 + k + p.a0, d.N - (d.y0 + d.y1 + k) + p.b0)
+            - betaln(p.a0, p.b0)
+            + betaln(k + p.ae, d.y1 + p.be) - betaln(p.ae, p.be))
+    return float(_log_binom_coef(d) + logsumexp(logw))
+
+
+def log_marginal_likelihood_no_benefit(d: Data, p: BreasePrior) -> float:
+    """M'_+: eta_e = 0, so theta1 = theta0 + eta_s (1 - theta0) (Appendix E.2)."""
+    j = np.arange(d.y1 + 1)
+    logw = (_lchoose(d.y1, j)
+            + betaln(d.y0 + j + p.a0, d.N - (d.y0 + j) + p.b0)
+            - betaln(p.a0, p.b0)
+            + betaln(d.y1 - j + p.as_, d.N1 - d.y1 + p.bs) - betaln(p.as_, p.bs))
+    return float(_log_binom_coef(d) + logsumexp(logw))
+
+
+# ---------------------------------------------------------------------------
+# Exact posterior sampling (Algorithm 1, Theorems 4-5)
+# ---------------------------------------------------------------------------
+
+def sample_posterior(d: Data, p: BreasePrior, size: int = 10_000, rng=None):
+    """Exact draws from the BREASE posterior (Algorithm 1).
+
+    Returns a dict of arrays: theta0, theta1, eta_e, eta_s, risk_ratio,
+    risk_diff, and the latent counts j = y1(1), k = x1(1).
+    """
+    rng = np.random.default_rng(rng)
+    logw, (A0, B0, Ae, Be, As, Bs) = _mixture_log_weights(d, p)
+    w = np.exp(logw - logsumexp(logw)).ravel()
+    idx = rng.choice(w.size, size=size, p=w)  # joint draw of (y1(0), x1(1))
+
+    t0 = rng.beta(A0.ravel()[idx], B0.ravel()[idx])
+    ee = rng.beta(Ae.ravel()[idx], Be.ravel()[idx])
+    es = rng.beta(As.ravel()[idx], Bs.ravel()[idx])
+    return _pack(t0, ee, es, *np.unravel_index(idx, logw.shape))
+
+
+def sample_posterior_no_harm(d: Data, p: BreasePrior, size=10_000, rng=None):
+    """Exact draws under eta_s = 0 (Algorithm 2)."""
+    rng = np.random.default_rng(rng)
+    k = np.arange(d.N1 - d.y1 + 1)
+    A0 = d.y0 + d.y1 + k + p.a0
+    B0 = d.N - (d.y0 + d.y1 + k) + p.b0
+    Ae, Be = k + p.ae, d.y1 + p.be + 0 * k
+    logw = _lchoose(d.N1 - d.y1, k) + betaln(A0, B0) + betaln(Ae, Be)
+    idx = rng.choice(k.size, size=size, p=np.exp(logw - logsumexp(logw)))
+    t0 = rng.beta(A0[idx], B0[idx])
+    ee = rng.beta(Ae[idx], Be[idx])
+    return _pack(t0, ee, np.zeros(size), None, k[idx])
+
+
+def sample_posterior_no_benefit(d: Data, p: BreasePrior, size=10_000, rng=None):
+    """Exact draws under eta_e = 0 (Algorithm 3). Here j indexes y1(1)."""
+    rng = np.random.default_rng(rng)
+    j = np.arange(d.y1 + 1)
+    A0 = d.y0 + j + p.a0
+    B0 = d.N - (d.y0 + j) + p.b0
+    As, Bs = d.y1 - j + p.as_, d.N1 - d.y1 + p.bs + 0 * j
+    logw = _lchoose(d.y1, j) + betaln(A0, B0) + betaln(As, Bs)
+    idx = rng.choice(j.size, size=size, p=np.exp(logw - logsumexp(logw)))
+    t0 = rng.beta(A0[idx], B0[idx])
+    es = rng.beta(As[idx], Bs[idx])
+    return _pack(t0, np.zeros(size), es, j[idx], None)
+
+
+def _pack(t0, ee, es, j, k):
+    t1 = theta1_from(t0, ee, es)
+    out = dict(theta0=t0, theta1=t1, eta_e=ee, eta_s=es,
+               risk_ratio=t1 / t0, risk_diff=t1 - t0)
+    if j is not None:
+        out["j"] = j
+    if k is not None:
+        out["k"] = k
+    return out
+
+
+# ---------------------------------------------------------------------------
+# Bayes factors
+# ---------------------------------------------------------------------------
+
+def bayes_factor_10(d: Data, p: BreasePrior = BreasePrior()) -> float:
+    """BF10 for H1: theta1 != theta0 (M1) vs H0: theta1 = theta0 (M0)."""
+    return float(np.exp(log_marginal_likelihood_M1(d, p)
+                        - log_marginal_likelihood_M0(d, p)))
+
+
+def log_bayes_factor_10(d: Data, p: BreasePrior = BreasePrior()) -> float:
+    """Use this when BF10 overflows (e.g. the COVID vaccine trial)."""
+    return log_marginal_likelihood_M1(d, p) - log_marginal_likelihood_M0(d, p)
+
+
+def one_sided_bayes_factors(d: Data, p: BreasePrior = BreasePrior(),
+                            size: int = 200_000, rng=None):
+    """Bayes factors against H0 for H-: theta1 < theta0 (benefit) and
+    H+: theta1 > theta0 (harm), using the truncated-M1 priors of eq. (3.15)
+    and the identity (3.17). Prior/posterior tail probabilities are Monte Carlo
+    estimates from exact draws.
+    """
+    rng = np.random.default_rng(rng)
+    t0, ee, es = p.sample(size, rng)
+    prior_minus = np.mean(theta1_from(t0, ee, es) < t0)
+    post = sample_posterior(d, p, size, rng)
+    post_minus = np.mean(post["theta1"] < post["theta0"])
+
+    lbf10 = log_bayes_factor_10(d, p)
+    with np.errstate(divide="ignore"):
+        return dict(
+            BF_minus_0=float(np.exp(lbf10 + np.log(post_minus) - np.log(prior_minus))),
+            BF_plus_0=float(np.exp(lbf10 + np.log1p(-post_minus) - np.log1p(-prior_minus))),
+            prior_P_benefit=float(prior_minus),
+            post_P_benefit=float(post_minus),
+        )
+
+
+def sensitivity_grid(d: Data, mue_grid, mus_grid, base: BreasePrior = BreasePrior()):
+    """log BF10 over a (mu_e, mu_s) grid, holding the other hyperparameters at
+    `base` (as in Figures 3b / 4b). Rows index mu_e, columns mu_s."""
+    out = np.empty((len(mue_grid), len(mus_grid)))
+    for a, me in enumerate(mue_grid):
+        for b, ms in enumerate(mus_grid):
+            pr = BreasePrior(base.mu0, me, ms, base.n0, base.ne, base.ns)
+            out[a, b] = log_bayes_factor_10(d, pr)
+    return out
+
+
+# ---------------------------------------------------------------------------
+# Prior moments (Appendix A)
+# ---------------------------------------------------------------------------
+
+def prior_moments(p: BreasePrior):
+    """Prior mean, variance and correlation of (theta0, theta1)."""
+    v0 = p.mu0 * (1 - p.mu0) / (p.n0 + 1)
+    ve = p.mue * (1 - p.mue) / (p.ne + 1)
+    vs = p.mus * (1 - p.mus) / (p.ns + 1)
+    cov = v0 * (1 - p.mue - p.mus)
+    var1 = (v0 * (1 - p.mue - p.mus) ** 2
+            + ve * (v0 + p.mu0 ** 2)
+            + vs * (v0 + (1 - p.mu0) ** 2))
+    mean1 = (1 - p.mue) * p.mu0 + p.mus * (1 - p.mu0)
+    return dict(mean_theta0=p.mu0, mean_theta1=mean1, var_theta0=v0,
+                var_theta1=var1, cov=cov, cor=cov / np.sqrt(v0 * var1))
+
+
+# ---------------------------------------------------------------------------
+# Convenience summary
+# ---------------------------------------------------------------------------
+
+def summarise(d: Data, p: BreasePrior = BreasePrior(), size=50_000, rng=None):
+    """Posterior summaries and Bayes factors for one trial, as a dict."""
+    post = sample_posterior(d, p, size, rng)
+    q = lambda x: np.quantile(x, [0.5, 0.025, 0.975])
+    lbf = log_bayes_factor_10(d, p)
+    return {
+        "risk_ratio (median, 2.5%, 97.5%)": q(post["risk_ratio"]),
+        "risk_diff (median, 2.5%, 97.5%)": q(post["risk_diff"]),
+        "eta_e (median, 2.5%, 97.5%)": q(post["eta_e"]),
+        "eta_s (median, 2.5%, 97.5%)": q(post["eta_s"]),
+        "log BF10": lbf,
+        "P(H1 | D) at 50:50 prior odds": 1 / (1 + np.exp(-lbf)),
+    }