From 9f124f25d00b1bd0dc26669d131f3740b3074f20 Mon Sep 17 00:00:00 2001 From: Gary Allen Date: Thu, 2 Jul 2026 00:00:54 +0200 Subject: [PATCH] Add TruncatedNormal Truncated Normal distribution with loc, scale, low, and high parameters. Every closed-form formula (sample, log_prob, cdf, icdf, survival function, entropy, mean, variance, mode, median) is cross-checked numerically against scipy.stats.truncnorm. KL divergence between arbitrary truncated normals isn't analytically tractable and raises NotImplementedError. --- distreqx/distributions/__init__.py | 1 + distreqx/distributions/_truncated_normal.py | 218 +++++++++++++ docs/api/distributions/truncated_normal.md | 8 + mkdocs.yml | 1 + tests/truncated_normal_test.py | 328 ++++++++++++++++++++ 5 files changed, 556 insertions(+) create mode 100644 distreqx/distributions/_truncated_normal.py create mode 100644 docs/api/distributions/truncated_normal.md create mode 100644 tests/truncated_normal_test.py diff --git a/distreqx/distributions/__init__.py b/distreqx/distributions/__init__.py index 06b8aa3..eeeb0dd 100644 --- a/distreqx/distributions/__init__.py +++ b/distreqx/distributions/__init__.py @@ -28,4 +28,5 @@ AbstractTransformed as AbstractTransformed, Transformed as Transformed, ) +from ._truncated_normal import TruncatedNormal as TruncatedNormal from ._uniform import Uniform as Uniform diff --git a/distreqx/distributions/_truncated_normal.py b/distreqx/distributions/_truncated_normal.py new file mode 100644 index 0000000..cfc2597 --- /dev/null +++ b/distreqx/distributions/_truncated_normal.py @@ -0,0 +1,218 @@ +"""Truncated Normal distribution.""" + +import math + +import jax +import jax.numpy as jnp +from jaxtyping import Array, Key + +from ._distribution import AbstractProbDistribution + +_half_log2pi = 0.5 * math.log(2 * math.pi) +_inv_sqrt_2pi = 1.0 / math.sqrt(2 * math.pi) + + +class TruncatedNormal(AbstractProbDistribution): + """ + Truncated Normal distribution with `loc`, `scale`, + `low`, and `high` parameters. + """ + + loc: Array + scale: Array + low: Array + high: Array + + def __init__(self, loc: Array, scale: Array, low: Array, high: Array): + """Initializes a Truncated Normal distribution. + + **Arguments:** + + - `loc`: Mean of the untruncated distribution. + - `scale`: Standard deviation of the untruncated distribution. + - `low`: Lower bound of the truncation. + - `high`: Upper bound of the truncation. + """ + self.loc = jnp.array(loc) + self.scale = jnp.array(scale) + self.low = jnp.array(low) + self.high = jnp.array(high) + + @property + def event_shape(self) -> tuple[int, ...]: + """Shape of event of distribution samples.""" + return jnp.broadcast_shapes( + self.loc.shape, self.scale.shape, self.low.shape, self.high.shape + ) + + @property + def support(self) -> tuple[Array, Array]: + """See `Distribution.support`.""" + return (self.low, self.high) + + def _standardize(self, value: Array) -> Array: + return (value - self.loc) / self.scale + + @property + def _std_low(self) -> Array: + return self._standardize(self.low) + + @property + def _std_high(self) -> Array: + return self._standardize(self.high) + + def _log_normalizer(self) -> Array: + """Calculates the log of the normalization constant Z.""" + cdf_high = jax.scipy.special.ndtr(self._std_high) + cdf_low = jax.scipy.special.ndtr(self._std_low) + return jnp.log(cdf_high - cdf_low) + + def _sample_from_std_trunc_normal(self, key: Key[Array, ""]) -> Array: + dtype = jnp.result_type(self.loc, self.scale, self.low, self.high) + return jax.random.truncated_normal( + key, + lower=self._std_low, + upper=self._std_high, + shape=self.event_shape, + dtype=dtype, + ) + + def sample(self, key: Key[Array, ""]) -> Array: + """See `Distribution.sample`.""" + rnd = self._sample_from_std_trunc_normal(key) + return self.scale * rnd + self.loc + + def sample_and_log_prob(self, key: Key[Array, ""]) -> tuple[Array, Array]: + """See `Distribution.sample_and_log_prob`.""" + rnd = self._sample_from_std_trunc_normal(key) + samples = self.scale * rnd + self.loc + + log_unnormalized = -0.5 * jnp.square(rnd) - _half_log2pi - jnp.log(self.scale) + log_prob = log_unnormalized - self._log_normalizer() + return samples, log_prob + + def log_prob(self, value: Array) -> Array: + """See `Distribution.log_prob`.""" + z = self._standardize(value) + log_unnormalized = -0.5 * jnp.square(z) - _half_log2pi - jnp.log(self.scale) + log_prob_val = log_unnormalized - self._log_normalizer() + + # Prob is 0 (log_prob is -inf) outside the bounds + is_valid = (value >= self.low) & (value <= self.high) + return jnp.where(is_valid, log_prob_val, -jnp.inf) + + def cdf(self, value: Array) -> Array: + """See `Distribution.cdf`.""" + z = self._standardize(value) + cdf_unnormalized = jax.scipy.special.ndtr(z) - jax.scipy.special.ndtr( + self._std_low + ) + Z = jax.scipy.special.ndtr(self._std_high) - jax.scipy.special.ndtr( + self._std_low + ) + cdf_val = cdf_unnormalized / Z + + # Clamp out-of-bounds CDF values + cdf_val = jnp.where(value < self.low, jnp.zeros_like(cdf_val), cdf_val) + cdf_val = jnp.where(value > self.high, jnp.ones_like(cdf_val), cdf_val) + return cdf_val + + def log_cdf(self, value: Array) -> Array: + """See `Distribution.log_cdf`.""" + return jnp.log(self.cdf(value)) + + def icdf(self, value: Array) -> Array: + """See `Distribution.icdf`.""" + cdf_low = jax.scipy.special.ndtr(self._std_low) + cdf_high = jax.scipy.special.ndtr(self._std_high) + Z = cdf_high - cdf_low + p_unnormalized = value * Z + cdf_low + return self.loc + self.scale * jax.scipy.special.ndtri(p_unnormalized) + + def survival_function(self, value: Array) -> Array: + """See `Distribution.survival_function`.""" + z = self._standardize(value) + Z = jax.scipy.special.ndtr(self._std_high) - jax.scipy.special.ndtr( + self._std_low + ) + surv_unnormalized = jax.scipy.special.ndtr( + self._std_high + ) - jax.scipy.special.ndtr(z) + surv_val = surv_unnormalized / Z + + # Clamp out-of-bounds survival values + surv_val = jnp.where(value < self.low, jnp.ones_like(surv_val), surv_val) + surv_val = jnp.where(value > self.high, jnp.zeros_like(surv_val), surv_val) + return surv_val + + def log_survival_function(self, value: Array) -> Array: + """See `Distribution.log_survival_function`.""" + return jnp.log(self.survival_function(value)) + + def _pdf(self, value: Array) -> Array: + """Standard normal PDF helper for mean/variance/entropy calculations.""" + return _inv_sqrt_2pi * jnp.exp(-0.5 * jnp.square(value)) + + def entropy(self) -> Array: + """Calculates the Shannon entropy (in nats).""" + alpha = self._std_low + beta = self._std_high + Z = jax.scipy.special.ndtr(beta) - jax.scipy.special.ndtr(alpha) + + pdf_alpha = self._pdf(alpha) + pdf_beta = self._pdf(beta) + + term1 = jnp.log(self.scale * Z) + 0.5 * math.log(2 * math.pi * math.e) + term2 = (alpha * pdf_alpha - beta * pdf_beta) / (2.0 * Z) + return term1 + term2 + + def mean(self) -> Array: + """Calculates the mean.""" + alpha = self._std_low + beta = self._std_high + Z = jax.scipy.special.ndtr(beta) - jax.scipy.special.ndtr(alpha) + + pdf_alpha = self._pdf(alpha) + pdf_beta = self._pdf(beta) + + return self.loc + self.scale * ((pdf_alpha - pdf_beta) / Z) + + def variance(self) -> Array: + """Calculates the variance.""" + alpha = self._std_low + beta = self._std_high + Z = jax.scipy.special.ndtr(beta) - jax.scipy.special.ndtr(alpha) + + pdf_alpha = self._pdf(alpha) + pdf_beta = self._pdf(beta) + + term1 = 1.0 + term2 = (alpha * pdf_alpha - beta * pdf_beta) / Z + term3 = jnp.square((pdf_alpha - pdf_beta) / Z) + + return jnp.square(self.scale) * (term1 + term2 - term3) + + def stddev(self) -> Array: + """Calculates the standard deviation.""" + return jnp.sqrt(self.variance()) + + def mode(self) -> Array: + """Calculates the mode.""" + # For a truncated normal, the mode is the mean if the mean is within bounds, + # otherwise it is the bound closest to the mean. + return jnp.clip(self.loc, self.low, self.high) + + def median(self) -> Array: + """Calculates the median.""" + return self.icdf(jnp.array(0.5, dtype=self.loc.dtype)) + + def kl_divergence(self, other_dist, **kwargs) -> Array: + """Calculates the KL divergence to another distribution. + + Raises: + NotImplementedError: The KL divergence between arbitrary truncated normal + distributions lacks a stable general analytical solution. + """ + raise NotImplementedError( + "KL divergence for TruncatedNormal is not analytically tractable." + ) diff --git a/docs/api/distributions/truncated_normal.md b/docs/api/distributions/truncated_normal.md new file mode 100644 index 0000000..1a9fe19 --- /dev/null +++ b/docs/api/distributions/truncated_normal.md @@ -0,0 +1,8 @@ +# Truncated Normal + +::: distreqx.distributions.TruncatedNormal + options: + members: + - __init__ + - entropy +--- \ No newline at end of file diff --git a/mkdocs.yml b/mkdocs.yml index aa5cf4e..557d8e3 100644 --- a/mkdocs.yml +++ b/mkdocs.yml @@ -102,6 +102,7 @@ nav: - Distributions: - Gaussians: - 'api/distributions/normal.md' + - 'api/distributions/truncated_normal.md' - 'api/distributions/mvn_diag.md' - 'api/distributions/mvn_from_bijector.md' - 'api/distributions/mvn_full_covariance.md' diff --git a/tests/truncated_normal_test.py b/tests/truncated_normal_test.py new file mode 100644 index 0000000..725e6ee --- /dev/null +++ b/tests/truncated_normal_test.py @@ -0,0 +1,328 @@ +"""Tests for `truncated_normal.py`.""" + +from unittest import TestCase + +import jax +import jax.numpy as jnp +import numpy as np +from parameterized import parameterized # type: ignore +from scipy import stats + +from distreqx.distributions import TruncatedNormal + + +class TruncatedNormalTest(TestCase): + @parameterized.expand( + [ + ( + "1d std trunc normal", + (jnp.array(0.0), jnp.array(1.0), jnp.array(-1.0), jnp.array(1.0)), + ), + ( + "2d std trunc normal", + (jnp.zeros(2), jnp.ones(2), -jnp.ones(2), jnp.ones(2)), + ), + ( + "rank 2 std trunc normal", + ( + jnp.zeros((3, 2)), + jnp.ones((3, 2)), + -jnp.ones((3, 2)), + jnp.ones((3, 2)), + ), + ), + ] + ) + def test_event_shape(self, name, distr_params): + loc, scale, low, high = distr_params + self.assertEqual(loc.shape, TruncatedNormal(loc, scale, low, high).event_shape) + + @parameterized.expand( + [ + ( + "1d std trunc normal", + (jnp.array(0.0), jnp.array(1.0), jnp.array(-1.0), jnp.array(1.0)), + ), + ( + "2d std trunc normal", + (jnp.zeros(2), jnp.ones(2), -jnp.ones(2), jnp.ones(2)), + ), + ( + "rank 2 std trunc normal", + ( + jnp.zeros((3, 2)), + jnp.ones((3, 2)), + -jnp.ones((3, 2)), + jnp.ones((3, 2)), + ), + ), + ] + ) + def test_sample_shape(self, name, distr_params): + distr_params = tuple(jnp.asarray(p, dtype=jnp.float32) for p in distr_params) + key = jax.random.key(0) + self.assertEqual( + distr_params[0].shape, + TruncatedNormal(*distr_params).sample(key).shape, + ) + + @parameterized.expand( + [ + ( + "1d std trunc normal", + (jnp.array(0.0), jnp.array(1.0), jnp.array(-1.0), jnp.array(1.0)), + ), + ( + "2d std trunc normal", + (jnp.zeros(2), jnp.ones(2), -jnp.ones(2), jnp.ones(2)), + ), + ( + "rank 2 std trunc normal", + ( + jnp.zeros((3, 2)), + jnp.ones((3, 2)), + -jnp.ones((3, 2)), + jnp.ones((3, 2)), + ), + ), + ] + ) + @jax.numpy_rank_promotion("raise") + def test_sample_and_log_prob(self, name, distr_params): + distr_params = tuple(jnp.asarray(p, dtype=jnp.float32) for p in distr_params) + key = jax.random.key(0) + dist = TruncatedNormal(*distr_params) + result = dist.sample_and_log_prob(key) + self.assertEqual(distr_params[0].shape, result[0].shape) + self.assertEqual(distr_params[0].shape, result[1].shape) + + @parameterized.expand( + [ + ( + "1d std trunc normal", + (jnp.array(0.0), jnp.array(1.0), jnp.array(-1.0), jnp.array(1.0)), + ), + ( + "2d std trunc normal", + (jnp.zeros(2), jnp.ones(2), -jnp.ones(2), jnp.ones(2)), + ), + ( + "rank 2 std trunc normal", + ( + jnp.zeros((3, 2)), + jnp.ones((3, 2)), + -jnp.ones((3, 2)), + jnp.ones((3, 2)), + ), + ), + ] + ) + def test_method_with_input(self, name, distr_params): + distr_params = tuple(jnp.asarray(p, dtype=jnp.float32) for p in distr_params) + value = jnp.asarray(distr_params[0], dtype=jnp.float32) + dist = TruncatedNormal(*distr_params) + for method in [ + "log_prob", + "cdf", + "log_cdf", + "survival_function", + "log_survival_function", + ]: + with self.subTest(method): + result = getattr(dist, method)(value) + self.assertEqual(value.shape, result.shape) + + @parameterized.expand( + [ + ("entropy", (0, 1, -1, 1), "entropy"), + ("mean", (0, 1, -1, 1), "mean"), + ("mean from 1d params", ([-1, 1], [1, 2], [-2, 0], [0, 2]), "mean"), + ("variance", (0, 1, -1, 1), "variance"), + ( + "variance from np params", + (np.ones(2), np.ones(2), np.zeros(2), 2 * np.ones(2)), + "variance", + ), + ("stddev", (0, 1, -1, 1), "stddev"), + ("mode", (0, 1, -1, 1), "mode"), + ("median", (0, 1, -1, 1), "median"), + ] + ) + def test_method(self, name, distr_params, function_string): + distr_params = tuple(jnp.asarray(p, dtype=jnp.float32) for p in distr_params) + dist = TruncatedNormal(*distr_params) + result = getattr(dist, function_string)() + self.assertEqual(distr_params[0].shape, result.shape) + + @parameterized.expand( + [ + ("loc within bounds", 0.5, 1.5, -1.0, 2.0), + ("loc outside bounds", 3.0, 1.0, -1.0, 1.0), + ("negative loc", -2.0, 0.5, -3.0, -1.5), + ] + ) + def test_stats_match_scipy(self, name, loc, scale, low, high): + # Cross-check the closed-form moment/entropy formulas against scipy's + # reference implementation, not just their output shapes. + a, b = (low - loc) / scale, (high - loc) / scale + ref = stats.truncnorm(a, b, loc=loc, scale=scale) + dist = TruncatedNormal( + jnp.array(loc), jnp.array(scale), jnp.array(low), jnp.array(high) + ) + + np.testing.assert_allclose(dist.mean(), ref.mean(), rtol=1e-5) + np.testing.assert_allclose(dist.variance(), ref.var(), rtol=1e-5) + np.testing.assert_allclose(dist.stddev(), ref.std(), rtol=1e-5) + np.testing.assert_allclose(dist.entropy(), ref.entropy(), rtol=1e-4, atol=1e-5) + np.testing.assert_allclose(dist.median(), ref.median(), rtol=1e-5) + + xs = np.linspace(low + 1e-3, high - 1e-3, 5) + jxs = jnp.asarray(xs, dtype=jnp.float32) + np.testing.assert_allclose( + dist.log_prob(jxs), + ref.logpdf(xs), # pyright: ignore[reportAttributeAccessIssue] + rtol=1e-4, + atol=1e-5, + ) + np.testing.assert_allclose(dist.cdf(jxs), ref.cdf(xs), rtol=1e-4, atol=1e-6) + + expected_mode = np.clip(loc, low, high) + np.testing.assert_allclose(dist.mode(), expected_mode, rtol=1e-5) + + @parameterized.expand( + [ + ("standard", ([0.0, 1.0], [1.0, 1.0], [-1.0, 0.5], [1.0, 2.0])), + ] + ) + def test_median(self, name, distr_params): + distr_params = tuple(jnp.asarray(p, dtype=jnp.float32) for p in distr_params) + dist = TruncatedNormal(*distr_params) + np.testing.assert_allclose( + dist.median(), dist.icdf(jnp.array(0.5, dtype=jnp.float32)), rtol=1e-5 + ) + + def test_kl_divergence_raises(self): + dist1 = TruncatedNormal( + jnp.array(0.0), jnp.array(1.0), jnp.array(-1.0), jnp.array(1.0) + ) + dist2 = TruncatedNormal( + jnp.array(0.0), jnp.array(1.0), jnp.array(-2.0), jnp.array(2.0) + ) + with self.assertRaises(NotImplementedError): + dist1.kl_divergence(dist2) + + @parameterized.expand( + [ + ( + "1d std trunc normal", + (jnp.array(0.0), jnp.array(1.0), jnp.array(-1.0), jnp.array(1.0)), + ), + ( + "2d std trunc normal", + (jnp.zeros(2), jnp.ones(2), -jnp.ones(2), jnp.ones(2)), + ), + ( + "rank 2 std trunc normal", + ( + jnp.zeros((3, 2)), + jnp.ones((3, 2)), + -jnp.ones((3, 2)), + jnp.ones((3, 2)), + ), + ), + ] + ) + def test_icdf_shape(self, name, distr_params): + distr_params = tuple(jnp.asarray(p, dtype=jnp.float32) for p in distr_params) + value = 0.5 * jnp.ones_like(distr_params[0]) + dist = TruncatedNormal(*distr_params) + result = dist.icdf(value) + self.assertEqual(value.shape, result.shape) + + def test_icdf_values(self): + loc = jnp.array([0.0, 1.0, -2.0]) + scale = jnp.array([1.0, 2.0, 0.5]) + low = jnp.array([-2.0, -1.0, -2.5]) + high = jnp.array([2.0, 3.0, -1.5]) + dist = TruncatedNormal(loc, scale, low, high) + + x = jnp.array([0.5, 1.0, -2.0]) + np.testing.assert_allclose(dist.icdf(dist.cdf(x)), x, rtol=1e-5) + + u = jnp.array([0.1, 0.5, 0.9]) + np.testing.assert_allclose(dist.cdf(dist.icdf(u)), u, rtol=1e-5) + + def test_out_of_bounds(self): + dist = TruncatedNormal( + loc=jnp.array(0.0), + scale=jnp.array(1.0), + low=jnp.array(-1.0), + high=jnp.array(1.0), + ) + + expected_inf = jnp.array(-jnp.inf, dtype=jnp.float32) + np.testing.assert_array_equal(dist.log_prob(jnp.array(1.5)), expected_inf) + np.testing.assert_array_equal(dist.log_prob(jnp.array(-1.5)), expected_inf) + + np.testing.assert_allclose(dist.cdf(jnp.array(-1.5)), 0.0, atol=1e-6) + np.testing.assert_allclose(dist.cdf(jnp.array(1.5)), 1.0, atol=1e-6) + + np.testing.assert_allclose( + dist.survival_function(jnp.array(-1.5)), 1.0, atol=1e-6 + ) + np.testing.assert_allclose( + dist.survival_function(jnp.array(1.5)), 0.0, atol=1e-6 + ) + + def test_vmap_inputs(self): + def log_prob_sum(dist, x): + return dist.log_prob(x).sum() + + loc = jnp.arange(3 * 4 * 5).reshape((3, 4, 5)) + scale = jnp.ones((3, 4, 5)) + low = loc - 1.0 + high = loc + 1.0 + dist = TruncatedNormal(loc, scale, low, high) + x = loc + + with self.subTest("no vmap"): + actual = log_prob_sum(dist, x) + expected = dist.log_prob(x).sum() + np.testing.assert_allclose(actual, expected) + + with self.subTest("axis=0"): + actual = jax.vmap(log_prob_sum, in_axes=0)(dist, x) + expected = dist.log_prob(x).sum(axis=(1, 2)) + np.testing.assert_allclose(actual, expected) + + with self.subTest("axis=1"): + actual = jax.vmap(log_prob_sum, in_axes=1)(dist, x) + expected = dist.log_prob(x).sum(axis=(0, 2)) + np.testing.assert_allclose(actual, expected) + + def test_vmap_outputs(self): + def summed_dist(loc, scale, low, high): + return TruncatedNormal( + loc.sum(keepdims=True), + scale.sum(keepdims=True), + low.sum(keepdims=True), + high.sum(keepdims=True), + ) + + loc = jnp.arange(3 * 4 * 5).reshape((3, 4, 5)).astype(jnp.float32) + scale = jnp.ones((3, 4, 5)) + low = loc - 1.0 + high = loc + 1.0 + + actual = jax.vmap(summed_dist)(loc, scale, low, high) + expected = TruncatedNormal( + loc.sum(axis=(1, 2), keepdims=True), + scale.sum(axis=(1, 2), keepdims=True), + low.sum(axis=(1, 2), keepdims=True), + high.sum(axis=(1, 2), keepdims=True), + ) + + np.testing.assert_equal(actual.event_shape, expected.event_shape) + + x = jnp.array([[[1.0]], [[2.0]], [[3.0]]]) + np.testing.assert_equal(actual.log_prob(x).shape, expected.log_prob(x).shape)