abfactor 0.5.1__py3-none-any.whl

abfactor/design.py

@@ -0,0 +1,168 @@
+"""The linear Gaussian reduction every quadrature analysis runs on.
+
+Per experiment, with estimates d[a, t] and sampling variances v[a, t] on the model scale:
+
+    d[a, t] = level[t] + (Q o*)[a] + (Q e*)[a] * on[t] + noise,   noise ~ N(0, v[a, t])
+
+    level[t] ~ N(m0, level_sd^2)          baseline per period
+    o*       ~ N(0, offset_sd^2 I)        persistent arm offset (optional)
+    e* | g   ~ N(0, g I)                  treatment effect
+
+with Q the orthonormal sum-to-zero basis. Given g everything is Gaussian, so
+
+    S(g) = N(d; mean, B + g K) / N(d; mean, B),   B = nuisance + noise covariance,
+                                                  K = effect covariance
+
+and the posterior of (level, o*, e*) given g is Gaussian in closed form.
+"""
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Optional
+
+import numpy as np
+from scipy.special import logsumexp
+from scipy.stats import norm
+
+
+def sum_to_zero_basis(a: int) -> np.ndarray:
+    """a x (a-1) orthonormal basis of the sum-to-zero subspace: the eigenvectors of
+    I - J/a with eigenvalue 1 (Q_a in Rouder et al. 2012, section 7.2)."""
+    w, v = np.linalg.eigh(np.eye(a) - np.ones((a, a)) / a)
+    return v[:, w > 0.5]
+
+
+@dataclass
+class Design:
+    arms: int
+    periods: int
+    m0: float
+    level_sd: float = 1.0
+    treatment_start: int = 0
+    offset_sd: Optional[float] = None          # None = no offsets
+
+    def __post_init__(self):
+        A, P = self.arms, self.periods
+        if A < 2:
+            raise ValueError("Need at least two arms")
+        if not 0 <= self.treatment_start < P:
+            raise ValueError("treatment_start must be one of the periods")
+        if self.offset_sd is not None and self.treatment_start == 0:
+            raise ValueError("Arm offsets need at least one pre-treatment period")
+        self.Q = sum_to_zero_basis(A)
+        self.on = (np.arange(P) >= self.treatment_start).astype(float)
+        # observation order: arm-major, (a, t) -> a * P + t
+        Xl = np.kron(np.ones((A, 1)), np.eye(P))
+        Xo = np.kron(self.Q, np.ones((P, 1)))
+        Xe = np.kron(self.Q, self.on[:, None])
+        self.Xl, self.Xo, self.Xe = Xl, Xo, Xe
+        blocks = [Xl] + ([Xo] if self.offset_sd is not None else []) + [Xe]
+        self.X = np.hstack(blocks)
+        self.n_level, self.n_off, self.n_eff = P, (A - 1 if self.offset_sd is not None else 0), A - 1
+        self.prior_mean = np.concatenate([np.full(P, self.m0), np.zeros(self.n_off + self.n_eff)])
+        nuis = self.level_sd ** 2 * Xl @ Xl.T
+        if self.offset_sd is not None:
+            nuis = nuis + self.offset_sd ** 2 * Xo @ Xo.T
+        self.nuisance_cov = nuis
+        self.K = Xe @ Xe.T
+        self.report_period = self.treatment_start
+
+    # ---- noise ---------------------------------------------------------------
+
+    def noise(self, v: np.ndarray, cov: Optional[np.ndarray] = None) -> np.ndarray:
+        """Sampling covariance of the estimates, shape (E, A*P, A*P) in arm-major order.
+        v: variances (E, A, P); cov: optional covariance across arms within each period,
+        (E, P, A, A), for estimates that share a term (e.g. a pooled pre-period mean)."""
+        E, A, P = v.shape
+        N = np.zeros((E, A * P, A * P))
+        idx = np.arange(A * P)
+        N[:, idx, idx] = v.reshape(E, -1)
+        if cov is not None:
+            for t in range(P):
+                at = np.arange(A) * P + t
+                N[:, at[:, None], at[None, :]] = cov[:, t]
+        return N
+
+    # ---- integrand ---------------------------------------------------------
+
+    def log_S(self, d: np.ndarray, N: np.ndarray, u: np.ndarray) -> np.ndarray:
+        """log S_e(g) on g = exp(u) for every experiment. d: (E, A, P); N: noise covariance
+        (E, A*P, A*P) from noise(). Returns (E, K)."""
+        E = d.shape[0]
+        g = np.exp(u)
+        out = np.empty((E, len(u)))
+        mean = self.Xl @ np.full(self.periods, self.m0)
+        for e in range(E):
+            B = self.nuisance_cov + N[e]
+            Lb = np.linalg.cholesky(np.linalg.inv(B))
+            lam, U = np.linalg.eigh(Lb.T @ self.K @ Lb)
+            lam = np.clip(lam, 0, None)
+            proj = (U.T @ (Lb.T @ (d[e].ravel() - mean))) ** 2
+            gl = g[:, None] * lam[None, :]
+            out[e] = 0.5 * np.sum(proj * gl / (1 + gl), axis=1) - 0.5 * np.sum(np.log1p(gl), axis=1)
+        return out
+
+    # ---- conditional posterior ----------------------------------------------
+
+    def posterior_given_g(self, d_e: np.ndarray, N_e: np.ndarray, g: np.ndarray):
+        """Gaussian posterior of theta = (level, offset*, effect*) for each g in a 1-D array,
+        given estimates d_e (A, P) with noise covariance N_e (A*P, A*P).
+        Returns means (G, p) and covariances (G, p, p)."""
+        g = np.atleast_1d(g)
+        Ninv = np.linalg.inv(N_e)
+        XtDX = self.X.T @ Ninv @ self.X
+        XtDd = self.X.T @ Ninv @ d_e.ravel()
+        base_prec = np.concatenate([np.full(self.n_level, 1 / self.level_sd ** 2),
+                                    np.full(self.n_off, 1 / self.offset_sd ** 2) if self.n_off else np.zeros(0)])
+        prec = np.empty((len(g), self.X.shape[1]))
+        prec[:, :len(base_prec)] = base_prec
+        prec[:, len(base_prec):] = 1 / np.maximum(g, 1e-300)[:, None]
+        P = XtDX[None] + np.einsum("gi,ij->gij", prec, np.eye(self.X.shape[1]))
+        cov = np.linalg.inv(P)
+        rhs = XtDd[None] + prec * self.prior_mean[None]
+        mean = np.einsum("gij,gj->gi", cov, rhs)
+        return mean, cov
+
+    def effect_rows(self) -> slice:
+        return slice(self.n_level + self.n_off, self.n_level + self.n_off + self.n_eff)
+
+    def contrast_vector(self, k: int, control: int = 0) -> np.ndarray:
+        return self.Q[k] - self.Q[control]
+
+    def natural_eta(self, theta: np.ndarray) -> np.ndarray:
+        """Model-scale value of each arm in the report period, offsets removed.
+        theta: (S, p) -> (S, A)."""
+        level = theta[:, self.report_period]
+        eff = theta[:, self.effect_rows()] @ self.Q.T
+        return level[:, None] + eff
+
+    # ---- exact pairwise Savage-Dickey ------------------------------------------
+
+    def pairwise(self, d_e, N_e, g, log_w_post, log_w_prior, k, control=0, tol=1e-12):
+        """Exact Savage-Dickey BF10 for psi = effect[k] - effect[control]: prior and posterior
+        densities of psi at 0 are mixtures of normals over the g grid. Returns
+        (log BF10, P(psi > 0 | data), posterior mean, posterior sd)."""
+        c = self.contrast_vector(k, control)
+        cc = float(c @ c)
+        keep = log_w_post > log_w_post.max() + np.log(tol)
+        wp = np.exp(log_w_post[keep] - logsumexp(log_w_post[keep]))
+        mean, cov = self.posterior_given_g(d_e, N_e, g[keep])
+        r = self.effect_rows()
+        m = mean[:, r] @ c
+        s = np.sqrt(np.einsum("i,gij,j->g", c, cov[:, r, r], c))
+        log_post0 = logsumexp(np.log(wp) + norm.logpdf(0, m, s))
+        log_prior0 = logsumexp(log_w_prior + norm.logpdf(0, 0, np.sqrt(cc * g)))
+        p_pos = float(np.sum(wp * norm.sf(0, m, s)))
+        mu = float(np.sum(wp * m))
+        sd = float(np.sqrt(np.sum(wp * (s ** 2 + m ** 2)) - mu ** 2))
+        return float(log_prior0 - log_post0), p_pos, mu, sd
+
+    def sample(self, d_e, N_e, g, log_w, size, rng):
+        """Exact posterior draws of theta: g from its grid posterior, then the conditional Gaussian."""
+        w = np.exp(log_w - logsumexp(log_w))
+        idx = rng.choice(len(g), size=size, p=w)
+        uniq, inv = np.unique(idx, return_inverse=True)
+        mean, cov = self.posterior_given_g(d_e, N_e, g[uniq])
+        chol = np.linalg.cholesky(cov + 1e-15 * np.eye(cov.shape[-1]))
+        z = rng.standard_normal((size, cov.shape[-1]))
+        return mean[inv] + np.einsum("sij,sj->si", chol[inv], z)

abfactor/metrics.py

@@ -0,0 +1,319 @@
+"""Metrics: what is measured, on which scale it is modelled, and its likelihood.
+
+Every metric reduces Stats to estimates d[e, a, t] with sampling variances
+v[e, a, t] on its model scale (used by the quadrature engine), and also supplies
+its exact log-likelihood as a function of the linear predictor (used by the NUTS
+and SMC engines).
+"""
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from dataclasses import dataclass
+from typing import Optional
+
+import numpy as np
+from scipy.special import expit, logit
+
+from .stats import Stats, square
+
+
+class Metric(ABC):
+    """warn_below / note_below: smallest-cell sizes below which the quadrature engine's
+    normal approximation is flagged (see docs/accuracy.md for how they were set)."""
+    name: str
+    scale: str = "raw"
+    warn_below: int = 5
+    note_below: int = 50
+
+    def min_cell(self, stats: Stats) -> np.ndarray:
+        """Smallest count driving the approximation in each experiment, shape (experiment,)."""
+        raise NotImplementedError
+
+    @abstractmethod
+    def estimates(self, stats: Stats) -> tuple[np.ndarray, np.ndarray]:
+        """(estimate, variance), each of shape (experiment, arm, period), on the model scale."""
+
+    def covariance(self, stats: Stats) -> Optional[np.ndarray]:
+        """Covariance of the estimates across arms within each period, (experiment, period,
+        arm, arm), when they share a term; None when the arms are independent."""
+        return None
+
+    @abstractmethod
+    def to_model(self, natural: float) -> float:
+        """A natural-scale value (e.g. a baseline rate) on the model scale."""
+
+    @abstractmethod
+    def to_natural(self, eta: np.ndarray) -> np.ndarray:
+        """Model-scale values back on the natural scale (numpy)."""
+
+    @abstractmethod
+    def pymc_loglik(self, eta, stats: Stats):
+        """Exact log-likelihood per experiment (shape (experiment,)) given the linear
+        predictor eta of shape (experiment, arm, period), as a PyTensor expression."""
+
+    def check(self, stats: Stats) -> None:
+        """Raise ValueError if `stats` lacks what this metric needs."""
+        pass
+
+
+# ---------------------------------------------------------------------------
+# Binary outcomes
+# ---------------------------------------------------------------------------
+
+@dataclass
+class Proportion(Metric):
+    """A binary outcome per user (converted or not), modelled on the log-odds scale.
+    Uses n and the statistic `name` (the number of successes). With several periods,
+    each period's users are assumed distinct; for one-off events among the same users
+    use TimeToEvent."""
+    name: str
+    scale: str = "log odds"
+
+    def estimates(self, stats):
+        y, n = stats.get(self.name), stats.get("n")
+        if np.any(y > n) or np.any(y < 0):
+            raise ValueError(f"{self.name}: successes must be between 0 and n")
+        return np.log((y + 0.5) / (n - y + 0.5)), 1 / (y + 0.5) + 1 / (n - y + 0.5)
+
+    def min_cell(self, stats):
+        y, n = stats.get(self.name), stats.get("n")
+        return np.minimum(y, n - y).min(axis=(1, 2))
+
+    def to_model(self, natural):
+        return float(logit(natural))
+
+    def to_natural(self, eta):
+        return expit(eta)
+
+    def pymc_loglik(self, eta, stats):
+        import pymc as pm
+        y, n = stats.get(self.name), stats.get("n")
+        return pm.logp(pm.Binomial.dist(n=n, logit_p=eta), y).sum(axis=(1, 2))
+
+
+@dataclass
+class TimeToEvent(Metric):
+    """A one-off event (a first conversion) over periods: a discrete-time hazard.
+
+    In each period only users who haven't had the event are at risk:
+        events_t ~ Binomial(at_risk_t, h_t),  at_risk_t = n_t - events before t
+    where n_t is the users assigned by the end of period t. link="cloglog" makes
+    effects log hazard ratios (proportional hazards), the natural choice when
+    periods differ in length."""
+    name: str
+    link: str = "logit"
+
+    def __post_init__(self):
+        if self.link not in ("logit", "cloglog"):
+            raise ValueError("link must be 'logit' or 'cloglog'")
+        self.scale = f"{self.link} hazard"
+
+    def at_risk(self, stats):
+        """Users still at risk at the start of each period (those not yet converted)."""
+        c, n = stats.get(self.name), stats.get("n")
+        before = np.concatenate([np.zeros_like(c[..., :1]), np.cumsum(c, axis=-1)[..., :-1]], axis=-1)
+        r = n - before
+        if np.any(c < 0) or np.any(c > r):
+            raise ValueError(f"{self.name}: events must be between 0 and the users still at risk")
+        return c, r
+
+    def min_cell(self, stats):
+        c, r = self.at_risk(stats)
+        return np.minimum(c, r - c).min(axis=(1, 2))
+
+    def estimates(self, stats):
+        c, r = self.at_risk(stats)
+        if self.link == "logit":
+            return np.log((c + 0.5) / (r - c + 0.5)), 1 / (c + 0.5) + 1 / (r - c + 0.5)
+        h = (c + 0.5) / (r + 1)
+        l1 = np.log1p(-h)
+        return np.log(-l1), h / ((1 - h) * (r + 1) * l1 ** 2)
+
+    def to_model(self, natural):
+        return float(logit(natural)) if self.link == "logit" else float(np.log(-np.log1p(-natural)))
+
+    def to_natural(self, eta):
+        return expit(eta) if self.link == "logit" else -np.expm1(-np.exp(eta))
+
+    def pymc_loglik(self, eta, stats):
+        import pymc as pm
+        import pytensor.tensor as pt
+        c, r = self.at_risk(stats)
+        if self.link == "logit":
+            dist = pm.Binomial.dist(n=r, logit_p=eta)
+        else:
+            dist = pm.Binomial.dist(n=r, p=pt.clip(1 - pt.exp(-pt.exp(eta)), 1e-12, 1 - 1e-12))
+        return pm.logp(dist, c).sum(axis=(1, 2))
+
+
+# ---------------------------------------------------------------------------
+# Counts
+# ---------------------------------------------------------------------------
+
+@dataclass
+class Count(Metric):
+    """Events per user (orders, boxes, visits) on the log-rate scale.
+
+    dispersion: "auto" uses `name__sq` when present to estimate phi = Var/mean of the
+    per-user count and inflates the variance by phi (quasi-Poisson); "none" assumes
+    Poisson; "require" errors without `name__sq`. Users differ, so the Poisson variance
+    is almost always too small: provide `name__sq`."""
+    name: str
+    dispersion: str = "auto"
+    scale: str = "log rate"
+
+    def phi(self, stats):
+        """Quasi-Poisson dispersion per experiment, arm and period (1 when dispersion='none')."""
+        y, n = stats.get(self.name), stats.get("n")
+        if self.dispersion == "none":
+            return np.ones_like(y)
+        if not stats.has(square(self.name)):
+            if self.dispersion == "require":
+                raise KeyError(f"dispersion='require' needs {square(self.name)!r}")
+            return np.ones_like(y)
+        m = y / n
+        var = (stats.get(square(self.name)) - n * m ** 2) / np.maximum(n - 1, 1)
+        return np.maximum(1.0, var / np.maximum(m, 1e-12))
+
+    def min_cell(self, stats):
+        return stats.get(self.name).min(axis=(1, 2))
+
+    def estimates(self, stats):
+        y, n = stats.get(self.name), stats.get("n")
+        return np.log((y + 0.5) / n), self.phi(stats) / (y + 0.5)
+
+    def to_model(self, natural):
+        return float(np.log(natural))
+
+    def to_natural(self, eta):
+        return np.exp(eta)
+
+    def pymc_loglik(self, eta, stats):
+        import pymc as pm
+        import pytensor.tensor as pt
+        y, n = stats.get(self.name), stats.get("n")
+        ll = pm.logp(pm.Poisson.dist(mu=n * pt.exp(eta)), y)
+        return (ll / self.phi(stats)).sum(axis=(1, 2))
+
+
+# ---------------------------------------------------------------------------
+# Continuous metrics
+# ---------------------------------------------------------------------------
+
+@dataclass
+class PrePeriod:
+    """Pre-period adjustment for a Mean, when the same users were observed before
+    the intervention. method="cuped" (ANCOVA; most precise under randomisation) or
+    "did" (difference in differences; removes a persistent arm difference exactly).
+    `pre` is the statistic name of the pre-period metric; needs pre, pre__sq and the
+    cross product with the post-period metric."""
+    method: str
+    pre: str
+
+    def __post_init__(self):
+        if self.method not in ("cuped", "did"):
+            raise ValueError("method must be 'cuped' or 'did'")
+
+
+class _Gaussian(Metric):
+    scale = "raw"
+    warn_below = 30          # users per arm: a CLT rule of thumb, not measured (docs/accuracy.md)
+    note_below = 0
+
+    def min_cell(self, stats):
+        return stats.get("n").min(axis=(1, 2))
+
+    def to_model(self, natural):
+        return float(natural)
+
+    def to_natural(self, eta):
+        return eta
+
+    def pymc_loglik(self, eta, stats):
+        import pymc as pm
+        est, var = self.estimates(stats)
+        cov = self.covariance(stats)
+        if cov is None:
+            return pm.logp(pm.Normal.dist(mu=eta, sigma=np.sqrt(var)), est).sum(axis=(1, 2))
+        return sum(pm.logp(pm.MvNormal.dist(mu=eta[:, :, t], cov=cov[:, t]), est[:, :, t])
+                   for t in range(est.shape[2]))
+
+
+def _moments(stats, name):
+    n = stats.get("n")
+    if np.any(n < 2):
+        raise ValueError("Need at least 2 users per arm")
+    s = stats.get(name)
+    m = s / n
+    var = (stats.get(square(name)) - n * m ** 2) / (n - 1)
+    return n, m, var
+
+
+@dataclass
+class Mean(_Gaussian):
+    """A per-user mean (revenue per user, orders per user), normal by the CLT with the
+    variance from per-user squares. Needs n, name, name__sq (and pre-period terms for
+    `adjust`). Adjusted means stay on the post-period scale."""
+    name: str
+    adjust: Optional[PrePeriod] = None
+
+    def _adjusted(self, stats):
+        """Adjusted means est_a = r_a + k * xbar, with r_a = m_a - k * mx_a and xbar the
+        pre-period mean pooled over arms (k = theta for CUPED, 1 for DiD), and their full
+        covariance across arms. xbar is shared, so the arms' estimates are correlated:
+            Cov(est_a, est_b) = [a == b] Var(r_a) + k (C_a + C_b) + k^2 Var(xbar),
+            C_a = Cov(r_a, xbar) = (cov_a - k varx_a) / N,  Var(xbar) = sum_b n_b varx_b / N^2
+        with k treated as known. The correlation is entirely in the level direction, so
+        contrasts between arms are unaffected; levels get the uncertainty in xbar."""
+        n, m, var = _moments(stats, self.name)
+        _, mx, varx = _moments(stats, self.adjust.pre)
+        cxy = (stats.get_cross(self.name, self.adjust.pre) - n * m * mx) / (n - 1)
+        N = n.sum(1, keepdims=True)
+        xbar = stats.get(self.adjust.pre).sum(1, keepdims=True) / N
+        if self.adjust.method == "cuped":
+            k = ((n - 1) * cxy).sum(1, keepdims=True) / np.maximum(((n - 1) * varx).sum(1, keepdims=True), 1e-300)
+        else:
+            k = np.ones_like(xbar)
+        var_r = np.maximum(var - 2 * k * cxy + k ** 2 * varx, 1e-300) / n          # (E, A, P)
+        C = (cxy - k * varx) / N                                                    # (E, A, P)
+        V = (n * varx).sum(1, keepdims=True) / N ** 2                               # (E, 1, P)
+        est = m - k * mx + k * xbar
+        Ct = np.moveaxis(C, 1, 2)                                                   # (E, P, A)
+        kt, Vt = np.moveaxis(k, 1, 2), np.moveaxis(V, 1, 2)                        # (E, P, 1)
+        cov = (kt[..., None] * (Ct[..., :, None] + Ct[..., None, :]) + (kt ** 2 * Vt)[..., None]
+               + np.einsum("epa,ab->epab", np.moveaxis(var_r, 1, 2), np.eye(var.shape[1])))
+        return est, cov
+
+    def estimates(self, stats):
+        if self.adjust is None:
+            n, m, var = _moments(stats, self.name)
+            return m, var / n
+        est, cov = self._adjusted(stats)
+        return est, np.moveaxis(np.diagonal(cov, axis1=2, axis2=3), 1, 2).copy()
+
+    def covariance(self, stats):
+        return None if self.adjust is None else self._adjusted(stats)[1]
+
+
+@dataclass
+class Ratio(_Gaussian):
+    """A ratio of per-user totals, sum(numerator) / sum(denominator), e.g. AOV =
+    revenue / orders. The user is the unit of randomisation, so the variance comes
+    from the delta method over users. Needs n, both metrics, their squares and their
+    cross product."""
+    numerator: str
+    denominator: str
+
+    @property
+    def name(self):
+        """The metric's label in results, "numerator/denominator"."""
+        return f"{self.numerator}/{self.denominator}"
+
+    def estimates(self, stats):
+        n, mx, varx = _moments(stats, self.numerator)
+        _, my, vary = _moments(stats, self.denominator)
+        if np.any(my <= 0):
+            raise ValueError(f"Sum of {self.denominator!r} must be positive in every arm")
+        cov = (stats.get_cross(self.numerator, self.denominator) - n * mx * my) / (n - 1)
+        R = mx / my
+        return R, np.maximum((varx - 2 * R * cov + R ** 2 * vary) / (n * my ** 2), 1e-300)

abfactor/pooling.py

@@ -0,0 +1,87 @@
+"""How experiments in a batch inform each other.
+
+Each experiment contributes its integrand L[e] = log S_e(g) on a common grid over g
+(the log likelihood ratio of "effects at scale g" against "no effects", with all
+other parameters integrated out). Pooling combines them:
+
+  Independent   each experiment has its own g: BF_e = int S_e(g) pi(g) dg
+  SharedScale   one g for the batch
+  SpikeAndSlab  one g, and each experiment has effects with probability w ~ Beta(a_w, b_w)
+
+For every experiment the result gives log p(g | the other experiments), the prior for
+its own Bayes factor, and log p(g | everything, this experiment has effects), used to
+draw its effects.
+"""
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+import numpy as np
+from scipy.special import betaln, logsumexp
+
+
+@dataclass
+class Independent:
+    pass
+
+
+@dataclass
+class SharedScale:
+    pass
+
+
+@dataclass
+class SpikeAndSlab:
+    a_w: float = 1.0
+    b_w: float = 1.0
+    w_grid_size: int = 201
+
+
+def _normalise(x, axis=-1):
+    return x - logsumexp(x, axis=axis, keepdims=True)
+
+
+def pool(L: np.ndarray, log_pg: np.ndarray, pooling) -> dict:
+    """L: (E, K) log S_e(g); log_pg: (K,) log prior mass per grid cell."""
+    E, K = L.shape
+    if isinstance(pooling, Independent):
+        rest = np.broadcast_to(_normalise(log_pg), (E, K))
+        post = _normalise(L + log_pg)
+        return dict(log_bf=logsumexp(L + log_pg, axis=1), log_g_rest=rest, log_g_post=post,
+                    p_effect=None, w_mean=None, log_bf_any=None, log_g_batch=None)
+
+    if isinstance(pooling, SharedScale):
+        total = L.sum(axis=0)
+        rest = _normalise(log_pg[None, :] + total[None, :] - L)
+        post = np.broadcast_to(_normalise(log_pg + total), (E, K))
+        return dict(log_bf=logsumexp(L + rest, axis=1), log_g_rest=rest, log_g_post=post,
+                    p_effect=None, w_mean=None, log_bf_any=float(logsumexp(log_pg + total)),
+                    log_g_batch=_normalise(log_pg + total))
+
+    if isinstance(pooling, SpikeAndSlab):
+        W = pooling.w_grid_size
+        w = (np.arange(W) + 0.5) / W
+        log_pw = ((pooling.a_w - 1) * np.log(w) + (pooling.b_w - 1) * np.log1p(-w)
+                  - betaln(pooling.a_w, pooling.b_w) - np.log(W))
+        lw, l1w = np.log(w)[:, None], np.log1p(-w)[:, None]
+        term = lambda e: np.logaddexp(l1w, lw + L[e][None, :])          # (W, K)
+        total = np.zeros((W, K))
+        for e in range(E):
+            total += term(e)
+        joint = total + log_pw[:, None] + log_pg[None, :]
+        log_bf_any = float(logsumexp(joint))
+        post_wg = np.exp(joint - log_bf_any)
+        log_bf, p_eff = np.empty(E), np.empty(E)
+        rest, post = np.empty((E, K)), np.empty((E, K))
+        for e in range(E):
+            t_e = term(e)
+            r_wg = total - t_e + log_pw[:, None] + log_pg[None, :]
+            rest[e] = _normalise(logsumexp(r_wg, axis=0))                 # p(g | others), w integrated
+            log_bf[e] = logsumexp(L[e] + rest[e])
+            # slab posterior for g: p(g | all data, this experiment has effects)
+            post[e] = _normalise(logsumexp(r_wg + lw, axis=0) + L[e])
+            p_eff[e] = float(np.sum(post_wg * np.exp(lw + L[e][None, :] - t_e)))
+        return dict(log_bf=log_bf, log_g_rest=rest, log_g_post=post, p_effect=p_eff,
+                    w_mean=float(np.sum(post_wg.sum(axis=1) * w)), log_bf_any=log_bf_any,
+                    log_g_batch=_normalise(logsumexp(joint + lw, axis=0)))
+    raise TypeError(f"Unknown pooling {pooling!r}")

abfactor/priors.py

@@ -0,0 +1,88 @@
+"""Priors on the treatment effect.
+
+Effects are a* | g ~ N(0, g I) on the sum-to-zero projection (Rouder et al. 2012,
+section 7.2), so the difference between any two arms is N(0, 2g) given g. A prior
+is a distribution over g, discretised on an even grid in u = log g: `grid()`
+returns u and the log prior mass of each cell.
+"""
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+
+import numpy as np
+
+from .brease import BreasePrior
+
+
+def log_inv_chi2_density(g: np.ndarray, r: float) -> np.ndarray:
+    """log density of g ~ r^2 Inv-chi2(1)."""
+    return np.log(r) - 0.5 * np.log(2 * np.pi) - 1.5 * np.log(g) - r ** 2 / (2 * g)
+
+
+class EffectPrior:
+    """Base class for priors on the effect scale g (the variance of each projected effect)."""
+
+    def grid(self, size: int) -> tuple[np.ndarray, np.ndarray]:
+        """(u, log_mass): grid points u = log g and the prior's log probability mass on each."""
+        raise NotImplementedError
+
+    def pymc_scale(self, dims):
+        """PyMC expression for sqrt(g), for the NUTS engine."""
+        raise NotImplementedError
+
+
+@dataclass
+class DefaultFixedEffects(EffectPrior):
+    """g ~ r^2 Inv-chi2(1): every pairwise contrast is Cauchy(0, sqrt(2) r) on the model scale.
+    For conversion rates on log odds, r around 0.1 suits online experiments; Rouder et al.'s
+    defaults (0.5 to 1) are for standardised effects and make the null win by default here."""
+    r: float = 0.1
+
+    def grid(self, size=3001):
+        u = np.linspace(np.log(self.r ** 2) - 25, np.log(self.r ** 2) + 12, size)
+        g = np.exp(u)
+        return u, log_inv_chi2_density(g, self.r) + u + np.log(u[1] - u[0])
+
+    def pymc_scale(self, dims=None):
+        import pymc as pm
+        g = pm.InverseGamma("g", alpha=0.5, beta=self.r ** 2 / 2, dims=dims)
+        return pm.math.sqrt(g)
+
+
+@dataclass
+class NormalEffects(EffectPrior):
+    """g fixed: every pairwise contrast is N(0, contrast_sd^2). Equivalent to an exchangeable
+    multivariate normal prior on the arms with that contrast SD."""
+    contrast_sd: float
+
+    def grid(self, size=None):
+        return np.array([np.log(self.contrast_sd ** 2 / 2)]), np.array([0.0])
+
+    def pymc_scale(self, dims=None):
+        return self.contrast_sd / np.sqrt(2)
+
+
+@dataclass
+class LearnedScale(EffectPrior):
+    """A distribution over g learned from completed experiments (Analysis.learn_prior).
+    Fixed before a new experiment starts, so it keeps sequential guarantees."""
+    u: np.ndarray
+    log_mass: np.ndarray
+    source: str = ""
+
+    def grid(self, size=None):
+        lm = self.log_mass - np.logaddexp.reduce(self.log_mass)
+        return self.u, lm
+
+    def pymc_scale(self, dims=None):
+        raise NotImplementedError("LearnedScale is supported by the quadrature engine only")
+
+
+@dataclass
+class Brease(EffectPrior):
+    """The BREASE prior of Irons & Cinelli (2023) for a binary two-arm experiment:
+    baseline risk, efficacy and side effects. `outcome="success"` treats the counted
+    statistic as a success (converted), so efficacy means preventing non-conversion;
+    "adverse" treats it as the adverse outcome directly."""
+    prior: BreasePrior = field(default_factory=BreasePrior)
+    outcome: str = "success"