diffbio 0.1.9__tar.gz → 0.1.10__tar.gz

{diffbio-0.1.9 → diffbio-0.1.10}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.5
 Name: diffbio
-Version: 0.1.9
+Version: 0.1.10
 Summary: End-to-end differentiable bioinformatics for JAX/Flax NNX: alignment, mapping, assembly, variant calling, RNA-seq, single-cell, epigenomics, CRISPR, metabolomics, multi-omics, protein and RNA structure, molecular dynamics and drug-discovery operators composed into trainable pipelines on datarax, artifex, opifex and calibrax
 Project-URL: Bug Tracker, https://github.com/avitai/DiffBio/issues
 Project-URL: Documentation, https://diffbio.readthedocs.io
@@ -45,22 +45,22 @@ Classifier: Topic :: Software Development :: Libraries
 Classifier: Topic :: Software Development :: Libraries :: Python Modules
 Requires-Python: <3.14,>=3.12
 Requires-Dist: anndata>=0.9.1
-Requires-Dist: avitai-artifex>=0.1.10
+Requires-Dist: avitai-artifex>=0.1.15
 Requires-Dist: beartype>=0.14.1
 Requires-Dist: biopython>=1.81
-Requires-Dist: calibrax>=0.1.8
+Requires-Dist: calibrax>=0.1.17
 Requires-Dist: chex>=0.1.7
-Requires-Dist: datarax>=0.1.13
+Requires-Dist: datarax<0.1.18,>=0.1.17
 Requires-Dist: flax>=0.12.9
 Requires-Dist: h5py>=3.7
 Requires-Dist: jax-md>=0.2.27
 Requires-Dist: jax>=0.11.1
 Requires-Dist: jaxtyping>=0.2.20
 Requires-Dist: numpy>=1.24
-Requires-Dist: opifex>=0.2.7
+Requires-Dist: opifex>=0.2.10
 Requires-Dist: optax>=0.2.8
 Requires-Dist: scipy>=1.10
-Requires-Dist: substrax>=0.1.11
+Requires-Dist: substrax>=0.1.25
 Provides-Extra: all
 Requires-Dist: bandit[toml]>=1.8.6; extra == 'all'
 Requires-Dist: beartype>=0.14.1; extra == 'all'
@@ -83,6 +83,7 @@ Requires-Dist: mkdocs-material>=9.6.7; extra == 'all'
 Requires-Dist: mkdocs>=1.6.1; extra == 'all'
 Requires-Dist: mkdocstrings-python>=1.1.2; extra == 'all'
 Requires-Dist: mkdocstrings>=0.28.3; extra == 'all'
+Requires-Dist: mpmath>=1.3; extra == 'all'
 Requires-Dist: optimistix>=0.0.9; extra == 'all'
 Requires-Dist: ott-jax>=0.5.0; extra == 'all'
 Requires-Dist: pandas>=2.0; extra == 'all'
@@ -199,6 +200,7 @@ Requires-Dist: ott-jax>=0.5.0; extra == 'soft-ops-ot'
 Provides-Extra: test
 Requires-Dist: beartype>=0.14.1; extra == 'test'
 Requires-Dist: coverage>=7; extra == 'test'
+Requires-Dist: mpmath>=1.3; extra == 'test'
 Requires-Dist: pytest-asyncio>=0.23; extra == 'test'
 Requires-Dist: pytest-benchmark>=4; extra == 'test'
 Requires-Dist: pytest-cov>=6.1.1; extra == 'test'

{diffbio-0.1.9 → diffbio-0.1.10}/pyproject.toml

@@ -36,17 +36,18 @@ classifiers = [
   "License :: OSI Approved :: MIT License"
 ]
 dependencies = [
-  # Datarax - operator, source, and pipeline contracts
-  "datarax>=0.1.13",
+  # Datarax - operator, source, and pipeline contracts. Capped below 0.1.18, which removes the
+  # eager source helpers diffbio.sources imports; the release that ports the sources raises it.
+  "datarax>=0.1.17,<0.1.18",
   # Artifex - modeling, modality, and transformer substrate
-  "avitai-artifex>=0.1.10",
+  "avitai-artifex>=0.1.15",
   # Opifex - scientific ML, operator learning, and advanced optimization
-  "opifex>=0.2.7",
+  "opifex>=0.2.10",
   # Calibrax - benchmarking, comparison, profiling, and regression control
-  "calibrax>=0.1.8",
+  "calibrax>=0.1.17",
   # Substrax - device identity, meshes, runtime settings and checkpoint store shared by the
-  # siblings; 0.1.6 adds substrax.runtime, substrax.artifacts and substrax.testing.
-  "substrax>=0.1.11",
+  # siblings; calibrax 0.1.17 requires 0.1.25.
+  "substrax>=0.1.25",
   # JAX ecosystem
   "beartype>=0.14.1",
   "chex>=0.1.7",
@@ -71,7 +72,7 @@ license = {file = "LICENSE"}
 name = "diffbio"
 readme = "README.md"
 requires-python = ">=3.12,<3.14"
-version = "0.1.9"
+version = "0.1.10"
 
 [project.optional-dependencies]
 all = ["diffbio[benchmark,chem,cuda12,dev,docs,genomics,metal,soft-ops-advanced,soft-ops-ot,test,torch-io]"]
@@ -148,7 +149,9 @@ test = [
   "pytest-timeout>=2.1",
   "pytest-xdist>=3.6",
   # tests/test_ci_shards.py reads the CI workflow.
-  "pyyaml>=6"
+  "pyyaml>=6",
+  # The exact references of tests/losses/test_count_likelihood_precision.py.
+  "mpmath>=1.3"
 ]
 torch-io = ["torch>=1.13.0"]
 

{diffbio-0.1.9 → diffbio-0.1.10}/src/diffbio/losses/__init__.py

@@ -4,6 +4,7 @@ This module provides loss functions for training differentiable bioinformatics
 pipelines, including biological regularization to prevent adversarial optimization.
 """
 
+from diffbio.losses._negative_binomial import nb_log_prob
 from diffbio.losses.alignment_losses import (
     AlignmentConsistencyLoss,
     AlignmentScoreLoss,
@@ -58,6 +59,7 @@ __all__ = [
     "HMMLikelihoodLoss",
     "NegativeBinomialLoss",
     "VAELoss",
+    "nb_log_prob",
     "zinb_negative_log_likelihood",
     # Self-supervised losses
     "MaskedValueLoss",

diffbio-0.1.10/src/diffbio/losses/_negative_binomial.py

@@ -0,0 +1,388 @@
+"""The negative binomial log-probability by Loader's saddle-point method, with closed-form partials.
+
+Catherine Loader, "Fast and Accurate Computation of Binomial Probabilities" (2002, revising
+2000) writes the binomial mass through the Stirling remainder ``delta`` (eq. 4) and the
+deviance ``bd0`` (eqs. 3, 6), so no two large logarithms are ever subtracted. The negative
+binomial follows from the identity
+
+    NB(x; theta, mu) = theta / (x + theta) * b(theta; x + theta, theta / (theta + mu)),
+
+which with ``n = x + theta`` reduces to
+
+    log NB = delta(n) - delta(theta) - delta(x) - bd0(theta, n P) - bd0(x, n Q)
+             - (log(2 pi) + log(x) + log1p(x / theta)) / 2,
+
+``P = theta / (theta + mu)`` and ``Q = mu / (theta + mu)``; it is Loader's Poisson form (eq. 7)
+as ``theta`` grows. The direct formula subtracts ``lgamma`` values of size ~7e7 at 5e6 counts,
+where float32 resolves 4 (issue #36). A zero count takes R's ``dnbinom_mu`` form,
+``theta log(theta / (theta + mu))`` through ``log1p`` of the small ratio.
+
+Two choices follow from float32 rather than from the paper, which works in float64:
+
+- The Stirling series is used from the argument where its first omitted term falls below the
+  dtype's ``eps`` relative to ``delta``, and ``lgamma`` below it.
+- ``bd0`` is Welinder's ``-x log1pmx((m - x) / x)`` (R bug PR#15628), with the difference
+  ``x - m`` computed in closed form, ``theta (mu - x) / (theta + mu)``. Loader's direct form
+  ``x log(x / m) + m - x`` cancels beyond his series radius and takes the log of a quotient,
+  which XLA's GPU float32 division rounds to about one ULP rather than to the nearest.
+
+The derivatives are not obtained by differentiating these branches: every branch would run in
+the backward pass, and autodiff of the formula cancels. A ``jax.custom_jvp`` gives the exact
+partials instead, each written without a cancelling difference:
+
+    d/dmu    = theta (x - mu) / (mu (theta + mu)),
+    d/dtheta = log1pmx(t) + x / (2 theta (theta + x)) + delta'(theta + x) - delta'(theta),
+    d/dx     = log1p(s) - (2x + theta) / (2x (x + theta)) + delta'(x + theta) - delta'(x),
+
+with ``t = (x - mu) / (theta + mu)``, ``s = theta (mu - x) / (x (theta + mu))`` and
+``log1pmx(t) = log1p(t) - t``. They follow from ``psi(z) = delta'(z) + log(z) - 1 / (2z)``
+applied to the textbook partials ``psi(x + theta) - psi(theta) + log(P) + (mu - x) / (theta +
+mu)`` and ``psi(x + theta) - psi(x + 1) + log(Q)``, whose terms are of size ``1 / theta`` while
+the partial in ``theta`` is of size ``1 / theta**2``. No published source states them; each piece
+is published: ``log1pmx`` by its artanh series (as R's ``log1pmx`` evaluates it), and
+``delta'`` by the asymptotic series of the digamma function (DLMF 5.11.2).
+"""
+
+import functools
+import math
+
+import jax
+import jax.numpy as jnp
+from jax.custom_derivatives import SymbolicZero
+from jax.scipy.special import digamma, gammaln
+from jaxtyping import Array, Float
+
+
+_HALF_LOG_TWO_PI = 0.5 * math.log(2.0 * math.pi)
+
+# Stirling's series for log(z!) - log(sqrt(2 pi z) (z / e)**z): B_{2k} / (2k (2k - 1)) z**(1-2k)
+# (Abramowitz and Stegun 6.1.40; Loader eq. 4 lists the first three).
+_STIRLING_COEFFICIENTS = (
+    1.0 / 12.0,
+    -1.0 / 360.0,
+    1.0 / 1260.0,
+    -1.0 / 1680.0,
+    1.0 / 1188.0,
+    -691.0 / 360360.0,
+    1.0 / 156.0,
+)
+# The first omitted coefficient, -3617 / 122400, bounds the truncation error of the alternating
+# series above (Whittaker and Watson, section 12.33).
+_STIRLING_FIRST_OMITTED = 3617.0 / 122400.0
+
+# Integer counts below this, where theta is below delta'(z)'s series start too, take the partial
+# in theta as an exact harmonic sum: one division per term, so a few at most.
+_HARMONIC_COUNTS = 4
+
+# log1pmx's artanh series runs on this interval of t and log(1 + t) - t outside it, where it
+# loses at most log2(0.5 / 0.193), 1.4 bits, at t = -0.5.
+_LOG1PMX_SERIES_INTERVAL = (-0.5, 1.0)
+
+
+def _stirling_series_start(dtype: jnp.dtype) -> int:
+    """The smallest ``z`` where the series' truncation error is below ``eps * delta(z)``.
+
+    ``delta(z) > 1 / (12 z + 1)`` and the truncation error is below the first omitted term
+    ``c z**-15``, so ``c z**-15 <= eps / (12 z + 1)`` holds from ``z = (13 c / eps)**(1/14)``.
+    """
+    eps = float(jnp.finfo(dtype).eps)
+    return math.ceil((13.0 * _STIRLING_FIRST_OMITTED / eps) ** (1.0 / 14.0))
+
+
+def _exact_stirling_remainders(dtype: jnp.dtype) -> dict[int, float]:
+    """``delta(n)`` at the integers below the series start, evaluated in float64.
+
+    R's ``stirlerr`` keeps such exact values for small arguments: below the series start
+    ``delta`` otherwise comes from ``lgamma`` minus terms many times its size (at ``n = 2``,
+    0.041 left from 0.69 and 1.73). Rounded once to ``dtype``, each value is correctly rounded.
+    """
+    return {
+        n: math.lgamma(n + 1.0) - (n + 0.5) * math.log(n) + n - _HALF_LOG_TWO_PI
+        for n in range(1, _stirling_series_start(dtype))
+    }
+
+
+def _stirling_derivative_series_start(dtype: jnp.dtype) -> int:
+    """The smallest ``z`` where the derivative series' truncation is below ``eps * |delta'(z)|``.
+
+    The series of ``delta'`` is that of ``psi(z) - log(z) + 1 / (2z)`` (DLMF 5.11.2), whose
+    remainder for real ``z > 0`` is below the first omitted term (DLMF 5.11(ii)),
+    ``15 c z**-16``. Its terms alternate and decrease for ``z >= 1``, so ``|delta'(z)| >= (1 -
+    1 / 10) / (12 z**2) > 1 / (14 z**2)``, and ``15 c z**-16 <= eps / (14 z**2)`` holds from
+    ``z = (210 c / eps)**(1/14)``.
+    """
+    eps = float(jnp.finfo(dtype).eps)
+    order = 2 * len(_STIRLING_COEFFICIENTS) + 1
+    return math.ceil((14.0 * order * _STIRLING_FIRST_OMITTED / eps) ** (1.0 / 14.0))
+
+
+def _log1pmx_truncation(t: float, terms: int) -> float:
+    """The bound ``2 |r| y**(K+1) / ((2K + 3)(1 - y))`` on dropping ``sum_{k >= K}``."""
+    r = t / (2.0 + t)
+    y = r * r
+    return 2.0 * abs(r) * y ** (terms + 1) / ((2 * terms + 3) * (1.0 - y))
+
+
+def _log1pmx_series_terms(dtype: jnp.dtype) -> int:
+    """Terms of ``sum_k y**k / (2k + 3)`` that reach ``eps`` relative to ``log1pmx``.
+
+    Checked at the interval's ends, where ``|r|`` and so the truncation are largest.
+    """
+    eps = float(jnp.finfo(dtype).eps)
+    terms = 1
+    while any(
+        _log1pmx_truncation(t, terms) > eps * abs(math.log1p(t) - t)
+        for t in _LOG1PMX_SERIES_INTERVAL
+    ):
+        terms += 1
+    return terms
+
+
+def stirling_remainder(z: Float[Array, "..."]) -> Float[Array, "..."]:
+    """``delta(z) = log(z!) - log(sqrt(2 pi z) (z / e)**z)`` for ``z > 0`` (Loader eq. 4).
+
+    Args:
+        z: Positive arguments.
+
+    Returns:
+        The remainder of Stirling's approximation, elementwise; exact values at the integers
+        below the series start.
+    """
+    start = _stirling_series_start(z.dtype)
+    large = z >= start
+    # Each branch sees an argument in its own domain, so the side the where drops stays finite.
+    z_series = jnp.where(large, z, start)
+    inverse_square = 1.0 / (z_series * z_series)
+    series = jnp.zeros_like(z)
+    for coefficient in reversed(_STIRLING_COEFFICIENTS):
+        series = series * inverse_square + coefficient
+    series = series / z_series
+    z_direct = jnp.where(large, 1.0, z)
+    direct = gammaln(z_direct + 1.0) - (z_direct + 0.5) * jnp.log(z_direct) + z_direct
+    result = jnp.where(large, series, direct - _HALF_LOG_TWO_PI)
+    for n, exact in _exact_stirling_remainders(z.dtype).items():
+        result = jnp.where(z == n, exact, result)
+    return result
+
+
+def stirling_remainder_derivative(z: Float[Array, "..."]) -> Float[Array, "..."]:
+    """``delta'(z) = psi(z + 1) - log(z) - 1 / (2z)`` for ``z > 0``.
+
+    Args:
+        z: Positive arguments.
+
+    Returns:
+        The derivative of the Stirling remainder, by DLMF 5.11.2's series where it reaches the
+        dtype's ``eps`` and by the digamma function below.
+    """
+    start = _stirling_derivative_series_start(z.dtype)
+    large = z >= start
+    # The second derivatives differentiate both sides, so each stays in its own domain.
+    z_series = jnp.where(large, z, start)
+    inverse_square = 1.0 / (z_series * z_series)
+    series = jnp.zeros_like(z)
+    for order, coefficient in reversed(list(enumerate(_STIRLING_COEFFICIENTS, start=1))):
+        series = series * inverse_square - (2 * order - 1) * coefficient
+    series = series * inverse_square
+    z_direct = jnp.where(large, 1.0, z)
+    direct = digamma(z_direct + 1.0) - jnp.log(z_direct) - 0.5 / z_direct
+    return jnp.where(large, series, direct)
+
+
+def log1pmx(t: Float[Array, "..."], one_plus_t: Float[Array, "..."]) -> Float[Array, "..."]:
+    """``log(1 + t) - t`` for ``t > -1``.
+
+    Inside ``[-0.5, 1]`` it is ``r (2 y S(y) - t)`` with ``r = t / (2 + t)``, ``y = r**2`` and
+    ``S(y) = sum_k y**k / (2k + 3)``, from ``log(1 + t) = 2 artanh(r)`` with the cancelling
+    ``2r - t`` formed as ``-r t``; outside, ``log(1 + t) - t`` loses at most 1.4 bits.
+
+    Args:
+        t: Arguments above -1.
+        one_plus_t: ``1 + t`` computed in closed form by the caller, so that it keeps its
+            relative accuracy as ``t`` approaches -1.
+
+    Returns:
+        ``log(1 + t) - t``, elementwise.
+    """
+    lower, upper = _LOG1PMX_SERIES_INTERVAL
+    inside = (t >= lower) & (t <= upper)
+    t_series = jnp.where(inside, t, 0.0)
+    r = t_series / (2.0 + t_series)
+    y = r * r
+    total = jnp.zeros_like(y)
+    for k in reversed(range(_log1pmx_series_terms(t.dtype))):
+        total = total * y + 1.0 / (2 * k + 3)
+    series = r * (2.0 * y * total - t_series)
+    one_plus_t_direct = jnp.where(inside, 1.0, one_plus_t)
+    return jnp.where(inside, series, jnp.log(one_plus_t_direct) - t)
+
+
+def binomial_deviance(
+    x: Float[Array, "..."], m: Float[Array, "..."], difference: Float[Array, "..."]
+) -> Float[Array, "..."]:
+    """``bd0(x, m) = x log(x / m) + m - x`` for ``x, m > 0`` (Loader eq. 6).
+
+    It is ``-x log1pmx(t)`` with ``t = (m - x) / x``, M. Welinder's form (R bug PR#15628; DPQ's
+    ``bd0_l1pm``): from ``m = x (1 + t)``, ``x log(x / m) = -x log1p(t)`` and ``m - x = x t``.
+    Nothing in it cancels, and for small ``|t|`` it is ``log1pmx``'s artanh series, the series of
+    Loader's eq. 6. Loader's direct ``x log(x / m) + (m - x)``, used beyond his series radius,
+    cancels and takes the log of ``x / m``, a quotient XLA's GPU float32 division rounds to
+    about one ULP rather than to the nearest: at ``(x, mu, theta) = (1e5, 5e4, 1e5)`` on an RTX
+    4090 that put 20 ULP of error on ``bd0(theta, n P)``, against 0.4 ULP on a CPU.
+
+    Args:
+        x: Positive first arguments.
+        m: Positive second arguments.
+        difference: ``x - m``, computed without subtracting the two.
+
+    Returns:
+        The deviance term, elementwise.
+    """
+    return -x * log1pmx(-difference / x, m / x)
+
+
+@jax.custom_jvp
+def nb_log_prob(
+    counts: Float[Array, "..."], mu: Float[Array, "..."], theta: Float[Array, "..."]
+) -> Float[Array, "..."]:
+    """``log NB(x | mu, theta)`` in the mean and inverse-dispersion parameterisation.
+
+    Differentiable in all three arguments through the closed-form partials of the module
+    docstring; in ``counts``, of the continuous extension through the gamma function.
+
+    Args:
+        counts: Non-negative counts ``x``.
+        mu: Positive means, broadcastable to ``counts``.
+        theta: Positive inverse dispersions, broadcastable to ``counts``.
+
+    Returns:
+        The log-probability, elementwise.
+    """
+    # No broadcast up front: a term in theta alone, such as delta(theta), runs at theta's own
+    # shape (one value per gene for NegativeBinomialLoss) and broadcasts only when combined.
+    positive = counts > 0
+    # The saddle-point branch sees a count of one at a zero count, where the where drops it.
+    x = jnp.where(positive, counts, 1.0)
+    # n P and n Q are each one product, as R's dnbinom_mu forms them: as x + (theta - n P) the
+    # mean term would subtract nearly equal numbers and lose mu once mu / theta is below eps.
+    # theta - n P = theta (mu - x) / (theta + mu) = x - n Q, so the deviances also see their
+    # arguments' difference in closed form rather than by subtraction.
+    total = theta + mu
+    n = x + theta
+    difference = theta * (mu - x) / total
+    log_positive = (
+        stirling_remainder(n)
+        - stirling_remainder(theta)
+        - stirling_remainder(x)
+        - binomial_deviance(theta, n * (theta / total), difference)
+        - binomial_deviance(x, n * (mu / total), -difference)
+        - 0.5 * (2.0 * _HALF_LOG_TWO_PI + jnp.log(x) + jnp.log1p(x / theta))
+    )
+    return jnp.where(positive, log_positive, _log_prob_of_zero(mu, theta))
+
+
+def _log_prob_of_zero(mu: Float[Array, "..."], theta: Float[Array, "..."]) -> Float[Array, "..."]:
+    """``log NB(0) = theta log(theta / (theta + mu))``, as R's ``dnbinom_mu`` forms it.
+
+    The log takes the small one of ``theta / (theta + mu)`` and ``mu / (theta + mu)``, each a
+    single quotient: ``log`` of the first when ``theta < mu``, ``log1p`` of minus the second
+    otherwise, where the first rounds towards one.
+    """
+    total = theta + mu
+    return theta * jnp.where(theta < mu, jnp.log(theta / total), jnp.log1p(-mu / total))
+
+
+def _log_mean_fraction(mu: Float[Array, "..."], theta: Float[Array, "..."]) -> Float[Array, "..."]:
+    """``log(mu / (theta + mu))``, as ``log1p(-theta / (theta + mu))`` when ``theta < mu``."""
+    total = theta + mu
+    return jnp.where(theta < mu, jnp.log1p(-theta / total), jnp.log(mu / total))
+
+
+def _partial_in_mu(
+    counts: Float[Array, "..."], mu: Float[Array, "..."], theta: Float[Array, "..."]
+) -> Float[Array, "..."]:
+    """``theta (x - mu) / (mu (theta + mu))``: ``x / mu - (x + theta) / (theta + mu)`` combined."""
+    return theta * (counts - mu) / (mu * (theta + mu))
+
+
+def _partial_in_theta(
+    counts: Float[Array, "..."], mu: Float[Array, "..."], theta: Float[Array, "..."]
+) -> Float[Array, "..."]:
+    """``log1pmx(t) + x / (2 theta (theta + x)) + delta'(theta + x) - delta'(theta)``.
+
+    At ``x = 0`` the rational term is 0 and the two ``delta'`` agree exactly, leaving
+    ``log1pmx(-mu / (theta + mu))``, the derivative of ``theta log(theta / (theta + mu))``.
+    The ``delta'`` difference is formed first: added one at a time, each ``delta'`` would absorb
+    a ``log1pmx`` far below its own size before the two cancel.
+
+    Below ``delta'``'s series start in ``theta``, ``delta'(theta)`` rests on digamma, whose
+    float32 error then bounds the result where it is small. There an integer count below
+    ``_HARMONIC_COUNTS`` takes ``psi(theta + x) - psi(theta) = sum_{j < x} 1 / (theta + j)``
+    (DLMF 5.5.2), which jax also uses for its hyp1f1 parameter derivatives: the partial is
+    ``sum_{j < x} (mu - j) / ((theta + j)(theta + mu)) + log1pmx(-mu / (theta + mu))``, exact
+    term by term. From the series start on, the general form is the more accurate: its terms
+    are of the size of the score's ``1 / theta**2``, the harmonic terms of ``1 / theta``.
+    """
+    total = theta + mu
+    stirling_difference = stirling_remainder_derivative(
+        theta + counts
+    ) - stirling_remainder_derivative(theta)
+    general = (
+        log1pmx((counts - mu) / total, (theta + counts) / total)
+        + counts / (2.0 * theta * (theta + counts))
+        + stirling_difference
+    )
+    harmonic = log1pmx(-mu / total, theta / total)
+    for j in range(_HARMONIC_COUNTS - 1):
+        harmonic = harmonic + jnp.where(j < counts, (mu - j) / ((theta + j) * total), 0.0)
+    is_harmonic = (
+        (counts < _HARMONIC_COUNTS)
+        & (counts == jnp.floor(counts))
+        & (theta < _stirling_derivative_series_start(theta.dtype))
+    )
+    return jnp.where(is_harmonic, harmonic, general)
+
+
+def _partial_in_counts(
+    counts: Float[Array, "..."], mu: Float[Array, "..."], theta: Float[Array, "..."]
+) -> Float[Array, "..."]:
+    """``psi(x + theta) - psi(x + 1) + log(mu / (theta + mu))``.
+
+    For ``x >= 1``, ``log1p(s) - (2x + theta) / (2x (x + theta)) + delta'(x + theta) -
+    delta'(x)``, ``s = theta (mu - x) / (x (theta + mu))``, with ``log`` of ``1 + s = mu (x +
+    theta) / (x (theta + mu))`` near ``s = -1`` and the ``delta'`` difference formed first, as in
+    ``_partial_in_theta``: ``psi(x + theta)`` and ``log(theta + mu)`` would otherwise cancel to
+    leave a result of size one when ``theta`` is large. Below one count, where the ``1 / x``
+    terms grow, the digamma functions directly.
+    """
+    large = counts >= 1.0
+    x = jnp.where(large, counts, 1.0)
+    total = theta + mu
+    ratio_minus_one = theta * (mu - x) / (x * total)
+    log_ratio = jnp.where(
+        ratio_minus_one < -0.5,
+        jnp.log(mu * (x + theta) / (x * total)),
+        jnp.log1p(ratio_minus_one),
+    )
+    stirling_difference = stirling_remainder_derivative(x + theta) - stirling_remainder_derivative(
+        x
+    )
+    identity = log_ratio - (2.0 * x + theta) / (2.0 * x * (x + theta)) + stirling_difference
+    x_direct = jnp.where(large, 1.0, counts)
+    direct = digamma(x_direct + theta) - digamma(x_direct + 1.0) + _log_mean_fraction(mu, theta)
+    return jnp.where(large, identity, direct)
+
+
+@functools.partial(nb_log_prob.defjvp, symbolic_zeros=True)
+def _nb_log_prob_jvp(
+    primals: tuple[Array, Array, Array], tangents: tuple[Array, Array, Array]
+) -> tuple[Array, Array]:
+    counts, mu, theta = primals
+    value = nb_log_prob(counts, mu, theta)
+    tangent = jnp.zeros_like(value)
+    partials = (_partial_in_counts, _partial_in_mu, _partial_in_theta)
+    for partial, tangent_in in zip(partials, tangents, strict=True):
+        # An argument nobody differentiates arrives as a symbolic zero: its partial is skipped.
+        if not isinstance(tangent_in, SymbolicZero):
+            tangent = tangent + partial(counts, mu, theta) * tangent_in
+    return value, tangent

{diffbio-0.1.9 → diffbio-0.1.10}/src/diffbio/losses/statistical_losses.py

@@ -15,38 +15,7 @@ import jax.numpy as jnp
 from flax import nnx
 from jaxtyping import Array, Float, Int
 
-from diffbio.constants import EPSILON
-from diffbio.core import soft_ops
-
-
-def _nb_log_prob(
-    counts: Float[Array, "..."],
-    mu: Float[Array, "..."],
-    theta: Float[Array, "..."],
-    eps: float = EPSILON,
-) -> Float[Array, "..."]:
-    """Per-element negative binomial log-probability ``log NB(x | mu, theta)``.
-
-    Shared core of :func:`zinb_negative_log_likelihood` and
-    :class:`NegativeBinomialLoss`.
-
-    Args:
-        counts: Observed counts ``x``.
-        mu: Mean parameter (broadcastable to ``counts``).
-        theta: Dispersion parameter (broadcastable to ``counts``).
-        eps: Numerical stability constant.
-
-    Returns:
-        Elementwise log-probability, same shape as ``counts``.
-    """
-    log_theta_mu = jnp.log(theta + mu + eps)
-    return (
-        jax.scipy.special.gammaln(counts + theta)
-        - jax.scipy.special.gammaln(theta)
-        - jax.scipy.special.gammaln(counts + 1.0)
-        + theta * (jnp.log(theta + eps) - log_theta_mu)
-        + counts * (jnp.log(mu + eps) - log_theta_mu)
-    )
+from diffbio.losses._negative_binomial import nb_log_prob
 
 
 def zinb_negative_log_likelihood(
@@ -81,24 +50,26 @@ def zinb_negative_log_likelihood(
     """
     mu = jnp.exp(log_rate)
     theta = jnp.exp(jnp.clip(log_theta, -10.0, 10.0))
-    eps = EPSILON
 
     # Log-space sigmoid computations (numerically stable)
     softplus_pi = jax.nn.softplus(-pi_logit)  # = -log(sigmoid(pi_logit))
-    log_theta_mu = jnp.log(theta + mu + eps)
-
-    # NB(0) in log-space combined with dropout logit
-    pi_theta_log = -pi_logit + theta * (jnp.log(theta + eps) - log_theta_mu)
 
-    # Case x == 0: log[sigmoid(pi) + (1 - sigmoid(pi)) * NB(0)]
-    case_zero = jax.nn.softplus(pi_theta_log) - softplus_pi
+    # log NB(x); at a zero count it is log NB(0) = -theta log1p(mu / theta), which stays
+    # accurate when theta >> mu, where log(theta) - log(theta + mu) cancels.
+    log_nb = nb_log_prob(counts, mu, theta)
 
-    # Case x > 0: log[(1 - sigmoid(pi)) * NB(x)]. The bracketed term equals
-    # ``-pi_logit + log NB(x)`` once ``pi_theta_log`` is expanded.
-    case_nonzero = -softplus_pi - pi_logit + _nb_log_prob(counts, mu, theta, eps)
+    # Case x == 0: log[pi + (1 - pi) NB(0)] = log(1 - y) with y = (1 - pi)(1 - NB(0)). Both
+    # 1 - NB(0) and 1 - y go through Machler's log1mexp ("Accurately Computing log(1 - exp(-|a|))",
+    # 2012; jax.nn.log1mexp), since log(pi + (1 - pi) NB(0)) as softplus(-pi + log NB(0)) -
+    # softplus(-pi) subtracts two nearly equal terms once NB(0) is close to one.
+    log_y = -jax.nn.softplus(pi_logit) + jax.nn.log1mexp(-log_nb)
+    case_zero = jax.nn.log1mexp(-log_y)
+    # Case x > 0: log[(1 - pi) NB(x)]
+    case_nonzero = -pi_logit + log_nb - softplus_pi
 
-    is_zero = soft_ops.less(counts, jnp.array(eps), softness=0.01)
-    log_prob = is_zero * case_zero + (1.0 - is_zero) * case_nonzero
+    # The counts are observed data, so the branch is chosen exactly, as scvi-tools'
+    # log_zinb_positive does; nothing differentiates the selection.
+    log_prob = jnp.where(counts == 0, case_zero, case_nonzero)
 
     return -jnp.sum(log_prob)
 
@@ -159,7 +130,7 @@ class NegativeBinomialLoss(nnx.Module):
         mu = jnp.maximum(mu, self.eps)
         theta = jnp.maximum(theta, self.eps)
 
-        ll = _nb_log_prob(counts, mu, theta[None, :], self.eps)
+        ll = nb_log_prob(counts, mu, theta[None, :])
 
         # Return mean negative log-likelihood
         return -jnp.mean(ll)

{diffbio-0.1.9 → diffbio-0.1.10}/src/diffbio/operators/statistical/nb_glm.py

@@ -19,6 +19,9 @@ from datarax.core.operator import OperatorModule
 from flax import nnx
 from jaxtyping import Array, Float, PyTree
 
+from diffbio.losses._negative_binomial import nb_log_prob
+
+
 logger = logging.getLogger(__name__)
 
 
@@ -160,30 +163,8 @@ class DifferentiableNBGLM(OperatorModule):
             Total log probability (scalar).
         """
         mu = self.predict_mean(design, size_factor)
-        dispersion = self.get_dispersion()
-
-        # Negative binomial log probability
-        # NB(k; mu, r) where r = dispersion
-        # log P(k) = log(Gamma(k + r)) - log(Gamma(k + 1)) - log(Gamma(r))
-        #          + r * log(r / (r + mu)) + k * log(mu / (r + mu))
-
-        r = dispersion
-        k = counts
-
-        # Compute log probability using the NB PMF
-        # Use jax.scipy.special functions for numerical stability
-        log_prob = (
-            jax.scipy.special.gammaln(k + r)
-            - jax.scipy.special.gammaln(k + 1)
-            - jax.scipy.special.gammaln(r)
-            + r * jnp.log(r / (r + mu + 1e-8))
-            + k * jnp.log(mu / (r + mu + 1e-8) + 1e-8)
-        )
-
-        # Sum over features
-        total_log_prob = jnp.sum(log_prob)
-
-        return total_log_prob
+        # NB(k; mu, r) with r = dispersion, the inverse dispersion: variance mu + mu**2 / r.
+        return jnp.sum(nb_log_prob(counts, mu, self.get_dispersion()))
 
     def batch_log_likelihood(
         self,