diff --git a/pyproject.toml b/pyproject.toml index 9e36d82..b8da1cd 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -9,10 +9,10 @@ readme = "README.md" [tool.poetry.dependencies] scanpy = "^1.11.0" python = "^3.9" -flax = "^0.10.4" -optax = "^0.2.4" -tensorflow_probability = "^0.22.0" -clu = "^0.0.11" +jax = ">=0.4.27" +flax = ">=0.10.4" +optax = ">=0.2.4" +clu = ">=0.0.11" tqdm = "^4.66.1" [build-system] diff --git a/scenvi/ENVI.py b/scenvi/ENVI.py index 1a92ba7..51c7f37 100644 --- a/scenvi/ENVI.py +++ b/scenvi/ENVI.py @@ -7,7 +7,6 @@ import pandas as pd import scanpy as sc import sklearn.neighbors -import tensorflow_probability.substrates.jax as jax_prob # type: ignore from flax import linen as nn from jax import jit, random from tqdm import trange, tqdm @@ -15,6 +14,7 @@ from scenvi._dists import ( KL, AOT_Distance, + fill_triangular, log_nb_pdf, log_normal_pdf, log_pos_pdf, @@ -368,7 +368,7 @@ def grammian_cov(self, dec_cov): :meta private: """ - dec_cov = jax_prob.math.fill_triangular(dec_cov) + dec_cov = fill_triangular(dec_cov) return jnp.matmul(dec_cov, dec_cov.transpose([0, 2, 1])) def create_train_state(self, key=random.key(0), init_lr=3e-4, decay_steps=100): diff --git a/scenvi/_dists.py b/scenvi/_dists.py index 89488bf..328a431 100644 --- a/scenvi/_dists.py +++ b/scenvi/_dists.py @@ -1,5 +1,25 @@ +"""Log-densities and matrix helpers for ENVI's decoders. + +These were tensorflow_probability's jax substrate. tfp's last release is 0.25.0 +(November 2024) and it is unmaintained, while its pinned version here (<0.23) +fails to import against any recent jax: + + AttributeError: module 'jax.interpreters.xla' has no attribute + 'pytype_aval_mappings' + +The surface actually used was five functions, all of them closed-form, so they +are written out here against jax directly. Each one is checked against the tfp +original in tests/test_dists.py. +""" + +import math + import jax.numpy as jnp -import tensorflow_probability.substrates.jax.distributions as jnd +from jax.nn import log_sigmoid +from jax.scipy.special import gammaln, xlogy + +#: ``log(2 * pi)``, the Gaussian normalization constant. +LOG_TWO_PI = math.log(2.0 * math.pi) def KL(mean, log_std): @@ -15,7 +35,9 @@ def log_pos_pdf(sample, l): # noqa: E741 :meta private: """ - log_prob = jnd.Poisson(rate=l).log_prob(sample) + # Poisson(rate=l): k log(l) - l - log(k!). xlogy keeps k = 0, l = 0 at 0 + # rather than nan, which is what tfp does. + log_prob = xlogy(sample, l) - l - gammaln(sample + 1.0) return jnp.mean(log_prob, axis=-1) @@ -24,7 +46,11 @@ def log_nb_pdf(sample, r, p): :meta private: """ - log_prob = jnd.NegativeBinomial(total_count=r, logits=p).log_prob(sample) + # NegativeBinomial(total_count=r, logits=p), i.e. p is the logit of the + # per-trial success probability: + # log C(k + r - 1, k) + k log(sigmoid(p)) + r log(1 - sigmoid(p)) + binomial_coefficient = gammaln(sample + r) - gammaln(sample + 1.0) - gammaln(r) + log_prob = binomial_coefficient + sample * log_sigmoid(p) + r * log_sigmoid(-p) return jnp.mean(log_prob, axis=-1) @@ -33,9 +59,23 @@ def log_zinb_pdf(sample, r, p, d): :meta private: """ - log_prob = jnd.Inflated( - jnd.NegativeBinomial(total_count=r, logits=p), inflated_loc_logits=d - ).log_prob(sample) + # Zero-inflated negative binomial: with probability sigmoid(2d) the sample is + # a structural zero, otherwise it is drawn from the negative binomial. Only + # the k = 0 branch mixes the two, and it is summed in log space. + # + # The factor of two is not a typo. tfp built the mixture's categorical from + # logits [d, -d], whose difference is 2d, so the inflation weight is + # sigmoid(2d) and not sigmoid(d). Getting this wrong rescales the + # zero-inflation parameter silently -- see tests/test_dists.py. + binomial_coefficient = gammaln(sample + r) - gammaln(sample + 1.0) - gammaln(r) + log_nb = binomial_coefficient + sample * log_sigmoid(p) + r * log_sigmoid(-p) + + log_nb_at_zero = r * log_sigmoid(-p) + log_prob = jnp.where( + sample == 0, + jnp.logaddexp(log_sigmoid(2.0 * d), log_sigmoid(-2.0 * d) + log_nb_at_zero), + log_sigmoid(-2.0 * d) + log_nb, + ) return jnp.mean(log_prob, axis=-1) @@ -44,7 +84,7 @@ def log_normal_pdf(sample, mean): :meta private: """ - log_prob = jnd.Normal(loc=mean, scale=1).log_prob(sample) + log_prob = -0.5 * (jnp.square(sample - mean) + LOG_TWO_PI) return jnp.mean(log_prob, axis=-1) @@ -57,3 +97,31 @@ def AOT_Distance(sample, mean): mean = jnp.reshape(mean, [mean.shape[0], -1]) log_prob = -jnp.square(sample - mean) return jnp.mean(log_prob, axis=-1) + + +def fill_triangular(x): + """Pack ``x`` into the lower triangle of a square matrix, tfp's way. + + Reproduces ``tfp.math.fill_triangular``, whose fill order is not the obvious + one -- the last ``m - n`` entries are laid down first, then the whole vector + reversed on top, and the lower triangle taken:: + + [1, 2, 3, 4, 5, 6] -> [[4, 0, 0], + [6, 5, 0], + [3, 2, 1]] + + :param x: (array) ``(..., n * (n + 1) / 2)`` values to pack + + :return: (array) ``(..., n, n)`` lower-triangular matrices + + :meta private: + """ + + m = x.shape[-1] + # m = n (n + 1) / 2, so n is the positive root of n^2 + n - 2m. + n = int(round((math.sqrt(1.0 + 8.0 * m) - 1.0) / 2.0)) + if n * (n + 1) // 2 != m: + raise ValueError(f"last dimension {m} is not a triangular number") + + packed = jnp.concatenate([x[..., n:], jnp.flip(x, axis=-1)], axis=-1) + return jnp.tril(jnp.reshape(packed, (*x.shape[:-1], n, n))) diff --git a/tests/test_dists.py b/tests/test_dists.py new file mode 100644 index 0000000..350e3ba --- /dev/null +++ b/tests/test_dists.py @@ -0,0 +1,121 @@ +"""The tfp-free log-densities must agree with the tensorflow_probability originals. + +scenvi no longer depends on tensorflow_probability, so these run only where it is +still installed and importable -- which, given the <0.23 pin fails against any +recent jax, in practice means a deliberately pinned environment. They are the +record that the replacements in `scenvi/_dists.py` reproduce what they replaced. +""" + +import numpy as np +import pytest + +jnp = pytest.importorskip("jax.numpy") +tfp = pytest.importorskip("tensorflow_probability.substrates.jax") +jnd = tfp.distributions + +from scenvi._dists import ( # noqa: E402 + fill_triangular, + log_nb_pdf, + log_normal_pdf, + log_pos_pdf, + log_zinb_pdf, +) + +SHAPE = (16, 32) + + +@pytest.fixture(autouse=True) +def double_precision(): + """Compare the formulae, not float32 rounding, whatever the ambient config.""" + from jax.experimental import enable_x64 + + with enable_x64(): + yield + + +@pytest.fixture +def rng(): + return np.random.default_rng(0) + + +@pytest.fixture +def counts(rng): + """Non-negative counts, deliberately including exact zeros.""" + sample = rng.poisson(2.0, size=SHAPE).astype(np.float64) + assert (sample == 0).any(), "the zero branch must be exercised" + return jnp.asarray(sample) + + +def test_poisson_matches_tfp(rng, counts): + rate = jnp.asarray(rng.uniform(0.05, 10.0, size=SHAPE)) + + expected = jnp.mean(jnd.Poisson(rate=rate).log_prob(counts), axis=-1) + np.testing.assert_allclose(log_pos_pdf(counts, rate), expected, rtol=1e-12, atol=1e-12) + + +def test_negative_binomial_matches_tfp(rng, counts): + total_count = jnp.asarray(rng.uniform(0.5, 20.0, size=SHAPE)) + logits = jnp.asarray(rng.normal(0.0, 3.0, size=SHAPE)) + + expected = jnp.mean( + jnd.NegativeBinomial(total_count=total_count, logits=logits).log_prob(counts), axis=-1 + ) + np.testing.assert_allclose( + log_nb_pdf(counts, total_count, logits), expected, rtol=1e-11, atol=1e-11 + ) + + +def test_zero_inflated_negative_binomial_matches_tfp(rng, counts): + total_count = jnp.asarray(rng.uniform(0.5, 20.0, size=SHAPE)) + logits = jnp.asarray(rng.normal(0.0, 3.0, size=SHAPE)) + inflation = jnp.asarray(rng.normal(0.0, 3.0, size=SHAPE)) + + expected = jnp.mean( + jnd.Inflated( + jnd.NegativeBinomial(total_count=total_count, logits=logits), + inflated_loc_logits=inflation, + ).log_prob(counts), + axis=-1, + ) + np.testing.assert_allclose( + log_zinb_pdf(counts, total_count, logits, inflation), expected, rtol=1e-11, atol=1e-11 + ) + + +def test_normal_matches_tfp(rng): + sample = jnp.asarray(rng.normal(size=SHAPE)) + mean = jnp.asarray(rng.normal(size=SHAPE)) + + expected = jnp.mean(jnd.Normal(loc=mean, scale=1).log_prob(sample), axis=-1) + np.testing.assert_allclose(log_normal_pdf(sample, mean), expected, rtol=1e-12, atol=1e-12) + + +@pytest.mark.parametrize("n", [1, 2, 3, 4, 8, 16]) +def test_fill_triangular_matches_tfp(rng, n): + x = jnp.asarray(rng.normal(size=(5, n * (n + 1) // 2))) + + np.testing.assert_array_equal(fill_triangular(x), tfp.math.fill_triangular(x)) + + +def test_extreme_parameters_match_tfp(rng): + """Saturated logits and large counts, where a naive log(sigmoid(p)) underflows.""" + sample = jnp.asarray(rng.integers(0, 500, size=SHAPE).astype(np.float64)) + total_count = jnp.asarray(rng.uniform(1.0, 500.0, size=SHAPE)) + logits = jnp.asarray(rng.uniform(-40.0, 40.0, size=SHAPE)) + inflation = jnp.asarray(rng.uniform(-40.0, 40.0, size=SHAPE)) + + nb = jnd.NegativeBinomial(total_count=total_count, logits=logits) + np.testing.assert_allclose( + log_nb_pdf(sample, total_count, logits), + jnp.mean(nb.log_prob(sample), axis=-1), + rtol=1e-11, + atol=1e-11, + ) + np.testing.assert_allclose( + log_zinb_pdf(sample, total_count, logits, inflation), + jnp.mean( + jnd.Inflated(nb, inflated_loc_logits=inflation).log_prob(sample), axis=-1 + ), + rtol=1e-11, + atol=1e-11, + )