binference 0.1.0__py3-none-any.whl

binference/__init__.py

@@ -0,0 +1,85 @@
+"""binference: population and hierarchical inference with PDFs.
+
+Infer population distributions -- bin fractions or rates, with honest
+uncertainties -- from noisy, heteroscedastic per-object PDFs or likelihoods
+over bins or categories, together with consistent posteriors for every
+individual object.
+
+Quick start
+-----------
+>>> import binference
+>>> mock = binference.simulate.mock_1d(500, rng=1)
+>>> result = binference.infer(mock.likelihoods, rng=2)
+>>> result.summary()  # doctest: +SKIP
+
+Main entry points
+-----------------
+* :class:`Bins` -- the support (1-D/N-D grids or categories).
+* :class:`Likelihoods` -- per-object likelihoods, with converters from
+  tables, posteriors (+ interim priors), densities, samples, Gaussian
+  measurements and binned measurements.
+* :class:`Response` -- the response of a binned measurement (a detector, or a
+  classifier's scores on a labelled validation set).
+* :func:`infer` -- one-call inference (``method="auto"`` by default: Gibbs, Laplace or NUTS).
+* :mod:`binference.priors` -- Dirichlet family, smoothness priors, rates.
+* :mod:`binference.samplers` -- inference engines.
+* :class:`PopulationResult` -- draws, summaries, per-object outputs, plots.
+* :mod:`binference.simulate` and :mod:`binference.validation` -- mocks, SBC, coverage.
+"""
+
+from . import diagnostics, priors, samplers, simulate, validation
+from ._version import __version__
+from .basis import Basis
+from .bins import Bins
+from .density import Density
+from .estimator import PopulationEstimator
+from .inference import infer
+from .likelihoods import Likelihoods
+from .measure import Measure
+from .model import PopulationModel
+from .response import Response
+from .results import PointEstimate, PopulationResult, Summary
+from .selection import Selection
+from .warnings import (
+    ApproximationWarning,
+    BinferenceWarning,
+    ConvergenceWarning,
+    DiagnosticsError,
+    InputWarning,
+)
+
+__all__ = [
+    "ApproximationWarning",
+    "Basis",
+    "BinferenceWarning",
+    "Bins",
+    "ConvergenceWarning",
+    "Density",
+    "DiagnosticsError",
+    "InputWarning",
+    "Likelihoods",
+    "Measure",
+    "PointEstimate",
+    "PopulationEstimator",
+    "PopulationModel",
+    "PopulationResult",
+    "Response",
+    "Selection",
+    "Summary",
+    "__version__",
+    "diagnostics",
+    "infer",
+    "priors",
+    "samplers",
+    "simulate",
+    "validation",
+]
+
+
+def __getattr__(name: str):
+    # binference.plotting and binference.io load on first use (matplotlib stays optional)
+    if name in ("io", "plotting"):
+        import importlib
+
+        return importlib.import_module(f".{name}", __name__)
+    raise AttributeError(f"module 'binference' has no attribute {name!r}")

binference/_kernels.py

@@ -0,0 +1,252 @@
+"""Low-level numerical kernels shared by models and engines.
+
+Every function accepts the likelihood matrix ``L`` (shape ``(N, K)``) either as
+a dense ``numpy.ndarray`` or as a ``scipy.sparse`` CSR array/matrix. The hot
+paths are two matrix-vector products:
+
+    s     = L @ rho                       (per-object marginal likelihoods)
+    n_hat = rho * (L.T @ (w / s))          (expected counts / EM "stack")
+"""
+
+from __future__ import annotations
+
+import math
+
+import numpy as np
+import scipy.sparse as sp
+
+from ._utils import chunk_slices
+
+__all__ = [
+    "loglike",
+    "loglike_and_counts",
+    "loglike_many",
+    "matvec",
+    "possible",
+    "responsibilities",
+    "rmatvec",
+    "sample_counts",
+    "sample_labels",
+]
+
+# Gibbs labels were drawn as fast in blocks of 4096 as of 65536 objects (N = 2e4 to 2e5,
+# K = 20 to 80)
+#: Objects per block of the chunked likelihood kernels.
+DEFAULT_CHUNK = 4096
+
+# rows are normalised to a maximum of one, so this is a log-likelihood below -575
+#: An object whose normalised evidence is below this is impossible.
+MIN_EVIDENCE = 1e-250
+
+# the largest fill at which sparse beat dense in every measured case (DECISIONS #77)
+#: Likelihoods are stored sparse below this share of non-zero entries.
+SPARSE_FILL = 0.1
+
+
+def matvec(L, x: np.ndarray) -> np.ndarray:
+    """``L @ x`` for dense or sparse ``L``."""
+    return np.asarray(L @ x).ravel()
+
+
+def rmatvec(L, y: np.ndarray) -> np.ndarray:
+    """``L.T @ y`` for dense or sparse ``L``."""
+    return np.asarray(L.T @ y).ravel()
+
+
+def possible(s: np.ndarray) -> bool:
+    """Whether every object's evidence ``s_i = L_i . rho`` is positive.
+
+    Evidences below :data:`MIN_EVIDENCE` count as zero; the one test used by the
+    log-likelihood, EM and the KKT gap (it also keeps ``w / s`` finite).
+    """
+    return bool(s.min() > MIN_EVIDENCE)  # the comparison also catches NaN
+
+
+def loglike(L, w: np.ndarray, rho: np.ndarray) -> float:
+    """``sum_i w_i log(L_i . rho)`` (``-inf`` if some object is impossible), one pass."""
+    return loglike_of_evidences(w, matvec(L, rho))
+
+
+def loglike_of_evidences(w: np.ndarray, s: np.ndarray) -> float:
+    """``sum_i w_i log s_i`` for the evidences ``s = L @ rho`` (``-inf`` if one is zero)."""
+    return float(np.dot(w, np.log(s))) if possible(s) else -np.inf
+
+
+def loglike_many(
+    L, w: np.ndarray, V: np.ndarray, block: int = 256, chunk: int | None = DEFAULT_CHUNK
+) -> np.ndarray:
+    """:func:`loglike` at each row of ``V`` ``(n, K)``: ``chunk`` objects at a time (each
+    sliced once), ``block`` rows of ``V`` per product (so at most ``block * chunk``
+    evidences held)."""
+    N, n = L.shape[0], V.shape[0]
+    out, low = np.zeros(n), np.full(n, np.inf)
+    for sl in chunk_slices(N, chunk):
+        Ls = L if sl.stop - sl.start == N else L[sl]
+        for pts in chunk_slices(n, block):
+            s = np.asarray(Ls @ V[pts].T)
+            low[pts] = np.minimum(low[pts], s.min(axis=0))
+            with np.errstate(divide="ignore", invalid="ignore"):
+                out[pts] += w[sl] @ np.log(s)
+    # the test of possible() per point
+    return np.where((low > MIN_EVIDENCE) & np.isfinite(out), out, -np.inf)
+
+
+def loglike_and_counts(L, w: np.ndarray, rho: np.ndarray) -> tuple[float, np.ndarray, np.ndarray]:
+    """Weighted marginal log-likelihood, expected counts and per-object evidences.
+
+    Parameters
+    ----------
+    L : (N, K) array or sparse matrix
+        Bin-averaged likelihoods (rows may carry arbitrary positive scales).
+    w : (N,) array
+        Object weights.
+    rho : (K,) array
+        Population fractions (or rates); must be non-negative.
+
+    Returns
+    -------
+    loglike : float
+        ``sum_i w_i log(L_i . rho)`` (``-inf`` if any object has zero evidence).
+    n_hat : (K,) array
+        ``rho_k * sum_i w_i L_ik / (L_i . rho)``.
+    s : (N,) array
+        ``L_i . rho``.
+    """
+    s = matvec(L, rho)
+    if not possible(s):
+        return -np.inf, np.zeros_like(rho), s
+    ll = float(np.dot(w, np.log(s)))
+    if not math.isfinite(ll):
+        return -np.inf, np.zeros_like(rho), s
+    n_hat = rho * rmatvec(L, w / s)
+    return ll, n_hat, s
+
+
+def responsibilities(
+    L, rho: np.ndarray, idx=None, chunk: int | None = DEFAULT_CHUNK
+) -> np.ndarray:
+    """Per-object posterior over bins given ``rho``: ``r_ik = L_ik rho_k / (L_i . rho)``.
+
+    Returns a dense ``(n, K)`` array for the selected objects ``idx`` (all by
+    default). Memory scales as ``n * K``; select a subset for large catalogues.
+    """
+    if idx is not None:
+        L = L[np.atleast_1d(idx)]
+    N = L.shape[0]
+    out = np.empty((N, rho.size))
+    for sl in chunk_slices(N, chunk):
+        block = L[sl]
+        block = block.toarray() if sp.issparse(block) else np.asarray(block)
+        P = block * rho
+        tot = P.sum(axis=1, keepdims=True)
+        with np.errstate(invalid="ignore", divide="ignore"):
+            out[sl] = np.where(tot > 0, P / tot, 0.0)
+    return out
+
+
+def _sample_labels_dense(Lc: np.ndarray, rho: np.ndarray, rng: np.random.Generator) -> np.ndarray:
+    c = Lc * rho
+    np.cumsum(c, axis=1, out=c)
+    # u in (0, total]: the first index with c_k >= u always has P_k > 0
+    u = (1.0 - rng.random(Lc.shape[0])) * c[:, -1]
+    idx = np.sum(c < u[:, None], axis=1)
+    return np.minimum(idx, Lc.shape[1] - 1)
+
+
+def _sample_labels_csr(Lc, rho: np.ndarray, rng: np.random.Generator) -> np.ndarray:
+    Lc = sp.csr_array(Lc)
+    indptr, indices = Lc.indptr, Lc.indices
+    data = Lc.data * rho[indices]
+    n = Lc.shape[0]
+    # normalise each row so that cumulative sums stay O(n): avoids cancellation
+    row_tot = np.asarray(Lc @ rho).ravel()
+    counts = np.diff(indptr)
+    data = data / np.repeat(np.where(row_tot > 0, row_tot, 1.0), counts)
+    cs = np.cumsum(data)
+    start = np.concatenate(([0.0], cs))[indptr[:-1]]
+    u = start + (1.0 - rng.random(n))
+    pos = np.searchsorted(cs, u, side="left")
+    pos = np.clip(pos, indptr[:-1], np.maximum(indptr[1:] - 1, indptr[:-1]))
+    return indices[pos]
+
+
+def sample_labels(L, rho: np.ndarray, rng: np.random.Generator, chunk: int | None = DEFAULT_CHUNK):
+    """Draw one bin label per object from ``Cat(r_i(rho))`` (vectorised; dense rows
+    ``chunk`` at a time, sparse rows at once, whose cost and memory follow the nonzeros:
+    slicing CSR rows per chunk made them slower, 57% at N = 2e5, K = 200, 2% filled)."""
+    if sp.issparse(L):
+        return _sample_labels_csr(sp.csr_array(L), rho, rng)
+    N = L.shape[0]
+    out = np.empty(N, dtype=np.int64)
+    for sl in chunk_slices(N, chunk):
+        block = L if sl.stop - sl.start == N else L[sl]  # one chunk: no copy
+        out[sl] = _sample_labels_dense(np.asarray(block), rho, rng)
+    return out
+
+
+def _multinomial_csr(Lc, m: np.ndarray, rho: np.ndarray, rng: np.random.Generator) -> np.ndarray:
+    """Draws of ``Mult(m_i, r_i)`` on the nonzeros of each CSR row, one per stored entry.
+
+    Conditional binomials position by position within rows, rows sorted by
+    their number of nonzeros: the cost follows the nonzeros, not K or ``m_i``.
+    """
+    nnz = np.diff(Lc.indptr)
+    order = np.argsort(-nnz, kind="stable")
+    n_sorted = nnz[order]
+    start = Lc.indptr[:-1][order]
+    p = Lc.data * rho[Lc.indices]
+    mass = np.asarray(Lc @ rho, dtype=float).ravel()[order]
+    left = m[order].astype(np.int64)
+    out = np.zeros(p.size)
+    for j in range(int(n_sorted[0]) if n_sorted.size else 0):
+        a = int(np.searchsorted(-n_sorted, -j, side="left"))  # rows with more than j nonzeros
+        pos = start[:a] + j
+        q = np.divide(p[pos], mass[:a], out=np.ones(a), where=mass[:a] > 0)
+        q = np.where(n_sorted[:a] == j + 1, 1.0, np.clip(q, 0.0, 1.0))  # last entry takes the rest
+        d = rng.binomial(left[:a], q)
+        out[pos] = d
+        left[:a] -= d
+        mass[:a] -= p[pos]
+    return out
+
+
+def sample_counts(
+    L,
+    w_int: np.ndarray,
+    rho: np.ndarray,
+    rng: np.random.Generator,
+    chunk: int | None = DEFAULT_CHUNK,
+    unit_weights: bool | None = None,
+) -> tuple[np.ndarray, np.ndarray | None]:
+    """Draw latent bin counts ``n_k = sum_i w_i 1[z_i = k]``.
+
+    For unit weights a label is drawn per object (and returned). Integer weights
+    ``m_i`` (compressed duplicate rows) draw one ``Mult(m_i, r_i)`` per row and
+    return ``labels=None``; for sparse ``L`` only over the row's nonzeros, so
+    the cost follows the nonzeros whatever K and the ``m_i`` (audit G-2).
+
+    Returns
+    -------
+    counts : (K,) float array
+    labels : (N,) int array or None
+    """
+    K = rho.size
+    if unit_weights is None:
+        unit_weights = bool(np.all(w_int == 1))
+    if unit_weights:
+        labels = sample_labels(L, rho, rng, chunk)
+        return np.bincount(labels, minlength=K).astype(float), labels
+    counts = np.zeros(K)
+    N = L.shape[0]
+    m_all = np.asarray(np.round(w_int), dtype=np.int64)
+    if sp.issparse(L):  # all rows at once, as sample_labels
+        Lc = sp.csr_array(L)
+        n = _multinomial_csr(Lc, m_all, rho, rng)
+        return np.bincount(Lc.indices, weights=n, minlength=K), None
+    for sl in chunk_slices(N, chunk):
+        R = responsibilities(L[sl], rho, chunk=None)
+        # guard tiny negative round-off and renormalise
+        R = np.clip(R, 0.0, None)
+        R /= R.sum(axis=1, keepdims=True)
+        counts += rng.multinomial(m_all[sl], R).sum(axis=0)
+    return counts, None

binference/_quadrature.py

@@ -0,0 +1,202 @@
+"""Adaptive bin averages, in the measure coordinate or in ``z`` (DECISIONS #23)."""
+
+from __future__ import annotations
+
+from collections.abc import Callable
+from itertools import product
+from typing import Literal
+
+import numpy as np
+from scipy import special
+
+from ._utils import chunk_slices, warn
+from .bins import Bins
+from .measure import _identity
+from .warnings import InputWarning
+
+#: Most function values a quadrature evaluates in one call.
+MAX_ELEMENTS = 2**22
+
+
+class _Rule:
+    # a tensor-product Gauss-Legendre rule on the unit cube, kept per axis
+    def __init__(self, n: int, ndim: int) -> None:
+        x, w = np.polynomial.legendre.leggauss(int(n))
+        self.t = 0.5 * (x + 1.0)
+        idx = np.array(list(product(range(int(n)), repeat=ndim)))  # (M, ndim)
+        self.idx = idx.T
+        self.logw = np.log(0.5 * w)[idx].sum(axis=1)  # (M,)
+
+    def points(self, lo: np.ndarray, hi: np.ndarray, inverse) -> np.ndarray:
+        # (C, M, ndim) points in z for cells [lo, hi] given in the integration coordinate
+        axes = [
+            inv(lo[:, d, None] + (hi - lo)[:, d, None] * self.t)[:, i]
+            for d, (inv, i) in enumerate(zip(inverse, self.idx, strict=True))
+        ]
+        return np.stack(axes, axis=-1)
+
+
+def bin_log_average(
+    logf: Callable[..., np.ndarray],
+    bins: Bins,
+    *,
+    rows: int | None = None,
+    coordinate: Literal["measure", "z"] = "measure",
+    tol: float = 1e-3,
+    max_depth: int = 12,
+    n_sub: int | None = None,
+    must_split: Callable[[np.ndarray, np.ndarray, np.ndarray], np.ndarray] | None = None,
+) -> np.ndarray:
+    """Log of the average of ``exp(logf)`` over each active bin, shape ``(N, K)``.
+
+    The average is taken uniformly in the measure coordinate ``u`` (the
+    within-bin population shape) or in ``z``. Two calling conventions:
+
+    * ``rows=None``: ``logf(z)`` receives points of shape ``(M,)`` (1-D) or
+      ``(M, ndim)`` and returns ``(N, M)`` log-values for all ``N`` rows; a
+      cell is refined for every row at once.
+    * ``rows=N``: ``logf(z, r)`` receives points of shape ``(P, M)`` or
+      ``(P, M, ndim)`` and row indices ``r`` of shape ``(P,)``, and returns
+      ``(P, M)``: row ``r[p]`` at the points ``z[p]``. Cells are refined per row.
+
+    Each bin starts as one cell. A cell is estimated with 5 and 3 Gauss-Legendre
+    points per axis; while their difference exceeds ``tol`` relative to the row's
+    largest bin average (in ``u``) or bin mass (in ``z``), it is split in half
+    along every axis, at most ``max_depth`` times, after which an
+    :class:`~binference.InputWarning` is raised. With ``n_sub`` a fixed
+    ``n_sub``-point rule is used instead.
+
+    ``must_split(z_lo, z_hi, r)`` (per-row mode) marks cells, given by their
+    bounds in ``z``, that must be split whatever the error estimate says: a
+    function narrower than a cell can fall between the points of both rules,
+    which then agree on a wrong value.
+    """
+    if not tol > 0 or max_depth < 0:
+        raise ValueError(
+            f"tol must be positive and max_depth non-negative; got {tol}, {max_depth}."
+        )
+    if rows is not None:
+        # refine blocks of rows separately, so the per-row state stays bounded
+        block = max(1, MAX_ELEMENTS // (8 * bins.K))
+        if rows > block:
+            kw = {"coordinate": coordinate, "tol": tol, "max_depth": max_depth, "n_sub": n_sub}
+            return np.concatenate(
+                [
+                    bin_log_average(
+                        lambda z, r, start=sl.start: logf(z, r + start),
+                        bins,
+                        rows=sl.stop - sl.start,
+                        must_split=None
+                        if must_split is None
+                        else (lambda a, b, r, start=sl.start: must_split(a, b, r + start)),
+                        **kw,
+                    )
+                    for sl in chunk_slices(rows, block)
+                ]
+            )
+    nd = bins.ndim
+    if coordinate == "measure":
+        edges, inverse = bins.u_edges, [m.inverse for m in bins.measures]
+    else:
+        edges, inverse = bins.edges, [_identity] * nd
+    j = bins.cell_index
+    lo = np.stack([e[:-1][i] for e, i in zip(edges, j, strict=True)], axis=-1)
+    hi = np.stack([e[1:][i] for e, i in zip(edges, j, strict=True)], axis=-1)
+    owner = np.arange(bins.K)
+    rules = [_Rule(n_sub, nd)] if n_sub is not None else [_Rule(5, nd), _Rule(3, nd)]
+
+    if rows is None:
+
+        def integrate(rule, lo, hi, row):
+            # (N, C) log averages over the cells, chunked to bound memory
+            z = rule.points(lo, hi, inverse)
+            C, M = z.shape[:2]
+            first = _call(logf, z[:1], M)
+            if first.shape[0] == 0:
+                raise ValueError("The function returned no rows.")
+            size = max(1, MAX_ELEMENTS // (first.shape[0] * M))
+            parts = [special.logsumexp(first.reshape(-1, 1, M) + rule.logw, axis=2)]
+            for sl in chunk_slices(C - 1, size):
+                v = _call(logf, z[1:][sl], M)
+                parts.append(special.logsumexp(v.reshape(v.shape[0], -1, M) + rule.logw, axis=2))
+            return np.concatenate(parts, axis=1)
+
+        row = None
+    else:
+
+        def integrate(rule, lo, hi, row):
+            # (P,) log averages, one row per cell
+            out = np.empty(lo.shape[0])
+            for sl in chunk_slices(lo.shape[0], max(1, MAX_ELEMENTS // len(rule.logw))):
+                z = rule.points(lo[sl], hi[sl], inverse)
+                v = np.asarray(logf(z[..., 0] if nd == 1 else z, row[sl]), dtype=float)
+                out[sl] = special.logsumexp(v + rule.logw, axis=1)
+            return out
+
+        K = bins.K
+        row = np.repeat(np.arange(rows), K)
+        owner = np.tile(owner, rows)
+        lo, hi = np.tile(lo, (rows, 1)), np.tile(hi, (rows, 1))
+
+    logfrac = np.zeros(owner.size)
+    fine, *coarse = (integrate(rule, lo, hi, row) for rule in rules)
+    N = fine.shape[0] if rows is None else rows
+    out = np.full((N, bins.K), -np.inf)
+    if n_sub is not None:
+        return fine if rows is None else fine.reshape(N, bins.K)
+    # errors are measured on averages (in u) or on masses (in z), relative to the
+    # row's largest; the reference grows as refinement resolves narrow peaks
+    logvol = np.zeros(bins.K) if coordinate == "measure" else np.log(bins.volumes)
+    rowmax = np.full(N, -np.inf)
+    corners = np.array(list(product((0.0, 0.5), repeat=nd)))
+    unconverged = np.zeros(bins.K, dtype=bool)
+    for depth in range(max_depth + 1):
+        shift = logfrac + logvol[owner]
+        if rows is None:
+            rowmax = np.maximum(rowmax, (fine + shift).max(axis=1))
+        else:
+            np.maximum.at(rowmax, row, fine + shift)
+        rowmax = np.maximum(rowmax, (out + logvol).max(axis=1))
+        scale = (rowmax[:, None] if rows is None else rowmax[row]) - logvol[owner]
+        with np.errstate(invalid="ignore"):
+            diff = np.nan_to_num(np.abs(np.exp(fine - scale) - np.exp(coarse[0] - scale)))
+        err = np.exp(logfrac) * (diff.max(axis=0) if rows is None else diff)
+        ok = err <= tol
+        if must_split is not None:
+            z_lo = np.stack([inv(lo[:, d]) for d, inv in enumerate(inverse)], axis=-1)
+            z_hi = np.stack([inv(hi[:, d]) for d, inv in enumerate(inverse)], axis=-1)
+            ok &= ~must_split(z_lo, z_hi, row)
+        if depth == max_depth:
+            unconverged[owner[~ok]] = True
+            ok[:] = True
+        if rows is None:
+            np.logaddexp.at(out.T, owner[ok], (fine[:, ok] + logfrac[ok]).T)
+        else:
+            np.logaddexp.at(out, (row[ok], owner[ok]), fine[ok] + logfrac[ok])
+        if ok.all():
+            break
+        keep = ~ok
+        half = 0.5 * (hi[keep] - lo[keep])
+        lo = (lo[keep][:, None, :] + 2 * half[:, None, :] * corners[None]).reshape(-1, nd)
+        hi = lo + np.repeat(half, len(corners), axis=0)
+        owner = np.repeat(owner[keep], len(corners))
+        logfrac = np.repeat(logfrac[keep], len(corners)) - nd * np.log(2.0)
+        if row is not None:
+            row = np.repeat(row[keep], len(corners))
+        fine, *coarse = (integrate(rule, lo, hi, row) for rule in rules)
+    if unconverged.any():
+        warn(
+            f"Bin averages did not reach tol={tol:g} in {int(unconverged.sum())} bin(s) after "
+            f"{max_depth} refinements; raise max_depth or tol.",
+            InputWarning,
+        )
+    return out
+
+
+def _call(logf: Callable, z: np.ndarray, M: int) -> np.ndarray:
+    # logf at the points of the cells in z (C, M, ndim), as (N, C * M)
+    pts = z.reshape(-1, z.shape[-1])
+    v = np.asarray(logf(pts[:, 0] if z.shape[-1] == 1 else pts), dtype=float)
+    if v.ndim != 2 or v.shape[1] != pts.shape[0]:
+        raise ValueError("The function must return an array of shape (N, n_points).")
+    return v

binference/_simplex.py

@@ -0,0 +1,54 @@
+"""Simplex helpers: softmax, isometric log-ratio (ILR) and centred log-ratio maps.
+
+The ILR map uses the orthonormal Helmert basis ``V`` (shape ``(K, K-1)``,
+columns orthonormal and summing to zero):
+
+    eta = V.T @ log(rho)          (ILR)
+    rho = softmax(V @ eta)        (inverse ILR)
+
+The Jacobian of ``eta -> rho[:K-1]`` has determinant ``sqrt(K) * prod(rho)``,
+so the log-density of ``eta`` is ``log p(rho) + sum(log rho) + 0.5 log K``.
+"""
+
+from __future__ import annotations
+
+from functools import lru_cache
+
+import numpy as np
+
+from ._utils import frozen
+
+__all__ = ["helmert_basis", "log_softmax", "softmax"]
+
+
+@lru_cache(maxsize=64)
+def _helmert(K: int) -> np.ndarray:
+    V = np.zeros((K, K - 1))
+    for j in range(1, K):
+        V[:j, j - 1] = 1.0 / j
+        V[j, j - 1] = -1.0
+        V[:, j - 1] *= np.sqrt(j / (j + 1.0))
+    return frozen(V)
+
+
+def helmert_basis(K: int) -> np.ndarray:
+    """Orthonormal ``(K, K-1)`` basis of the zero-sum subspace of R^K."""
+    if K < 2:
+        raise ValueError("The simplex needs at least two bins (K >= 2).")
+    return _helmert(int(K))
+
+
+def softmax(x: np.ndarray, axis: int = -1) -> np.ndarray:
+    """Numerically stable softmax along ``axis``."""
+    x = np.asarray(x, dtype=float)
+    m = np.max(x, axis=axis, keepdims=True)
+    e = np.exp(x - m)
+    return e / np.sum(e, axis=axis, keepdims=True)
+
+
+def log_softmax(x: np.ndarray, axis: int = -1) -> np.ndarray:
+    """Numerically stable log-softmax along ``axis``."""
+    x = np.asarray(x, dtype=float)
+    m = np.max(x, axis=axis, keepdims=True)
+    z = x - m
+    return z - np.log(np.sum(np.exp(z), axis=axis, keepdims=True))