diff --git a/distreqx/distributions/__init__.py b/distreqx/distributions/__init__.py index 06b8aa3..ff8efff 100644 --- a/distreqx/distributions/__init__.py +++ b/distreqx/distributions/__init__.py @@ -12,6 +12,7 @@ from ._gamma import Gamma as Gamma from ._independent import Independent as Independent from ._logistic import Logistic as Logistic +from ._lognormal import LogNormal as LogNormal from ._mixture_same_family import MixtureSameFamily as MixtureSameFamily from ._mvn_diag import MultivariateNormalDiag as MultivariateNormalDiag from ._mvn_from_bijector import ( diff --git a/distreqx/distributions/_lognormal.py b/distreqx/distributions/_lognormal.py new file mode 100644 index 0000000..50335d1 --- /dev/null +++ b/distreqx/distributions/_lognormal.py @@ -0,0 +1,149 @@ +"""LogNormal 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) + + +class LogNormal(AbstractProbDistribution): + """ + LogNormal distribution parameterized by + `loc` and `scale` of the underlying Normal. + """ + + loc: Array + scale: Array + + def __init__(self, loc: Array, scale: Array): + """Initializes a LogNormal distribution. + + **Arguments:** + + - `loc`: Mean of the underlying Normal distribution. + - `scale`: Standard deviation of the underlying Normal distribution. + """ + self.loc = jnp.array(loc) + self.scale = jnp.array(scale) + + @property + def event_shape(self) -> tuple[int, ...]: + """Shape of event of distribution samples.""" + return self.loc.shape + + @property + def support(self) -> tuple[Array, Array]: + """See `Distribution.support`.""" + dtype = jnp.result_type(self.loc, self.scale) + return (jnp.array(0.0, dtype=dtype), jnp.array(jnp.inf, dtype=dtype)) + + def _sample_from_std_normal(self, key: Key[Array, ""]) -> Array: + dtype = jnp.result_type(self.loc, self.scale) + return jax.random.normal(key, shape=self.event_shape, dtype=dtype) + + def sample(self, key: Key[Array, ""]) -> Array: + """See `Distribution.sample`.""" + rnd = self._sample_from_std_normal(key) + normal_samples = self.scale * rnd + self.loc + return jnp.exp(normal_samples) + + def sample_and_log_prob(self, key: Key[Array, ""]) -> tuple[Array, Array]: + """See `Distribution.sample_and_log_prob`.""" + rnd = self._sample_from_std_normal(key) + normal_samples = self.scale * rnd + self.loc + samples = jnp.exp(normal_samples) + + # Change of variables: log(P(Y=y)) = log(P(X=x)) - log(y), where x = log(y) + normal_log_prob = -0.5 * jnp.square(rnd) - _half_log2pi - jnp.log(self.scale) + log_prob = normal_log_prob - normal_samples + return samples, log_prob + + def log_prob(self, value: Array) -> Array: + """See `Distribution.log_prob`.""" + log_value = jnp.log(value) + log_unnormalized = -0.5 * jnp.square(self._standardize(log_value)) + log_normalization = _half_log2pi + jnp.log(self.scale) + log_value + return log_unnormalized - log_normalization + + def icdf(self, value: Array) -> Array: + """See `Distribution.icdf`.""" + return jnp.exp(jax.scipy.special.ndtri(value) * self.scale + self.loc) + + def cdf(self, value: Array) -> Array: + """See `Distribution.cdf`.""" + return jax.scipy.special.ndtr(self._standardize(jnp.log(value))) + + def log_cdf(self, value: Array) -> Array: + """See `Distribution.log_cdf`.""" + return jax.scipy.special.log_ndtr( + self._standardize(jnp.log(value)) + ) # pyright: ignore[reportGeneralTypeIssues] + + def survival_function(self, value: Array) -> Array: + """See `Distribution.survival_function`.""" + return jax.scipy.special.ndtr(-self._standardize(jnp.log(value))) + + def log_survival_function(self, value: Array) -> Array: + """See `Distribution.log_survival_function`.""" + return jax.scipy.special.log_ndtr( + -self._standardize(jnp.log(value)) + ) # pyright: ignore[reportGeneralTypeIssues] + + def _standardize(self, value: Array) -> Array: + return (value - self.loc) / self.scale + + def entropy(self) -> Array: + """Calculates the Shannon entropy (in nats).""" + log_normalization = _half_log2pi + jnp.log(self.scale) + entropy = 0.5 + log_normalization + self.loc + return entropy + + def mean(self) -> Array: + """Calculates the mean.""" + return jnp.exp(self.loc + 0.5 * jnp.square(self.scale)) + + def variance(self) -> Array: + """Calculates the variance.""" + var_scale = jnp.square(self.scale) + return jnp.expm1(var_scale) * jnp.exp(2.0 * self.loc + var_scale) + + def stddev(self) -> Array: + """Calculates the standard deviation.""" + return jnp.sqrt(self.variance()) + + def mode(self) -> Array: + """Calculates the mode.""" + return jnp.exp(self.loc - jnp.square(self.scale)) + + def median(self) -> Array: + """Calculates the median.""" + return jnp.exp(self.loc) + + def kl_divergence(self, other_dist, **kwargs) -> Array: + """Calculates the KL divergence to another distribution. + + **Arguments:** + + - `other_dist`: A compatible disteqx LogNormal distribution. + - `kwargs`: Additional kwargs. + + **Returns:** + + The KL divergence `KL(self || other_dist)`. + """ + # KL divergence is invariant under strictly monotonic transformations. + # Thus, KL(LogNormal(m1, s1) || LogNormal(m2, s2)) is identical to + # KL(Normal(m1, s1) || Normal(m2, s2)). + dist1 = self + dist2 = other_dist + diff_log_scale = jnp.log(dist1.scale) - jnp.log(dist2.scale) + return ( + 0.5 * jnp.square(dist1.loc / dist2.scale - dist2.loc / dist2.scale) + + 0.5 * jnp.expm1(2.0 * diff_log_scale) + - diff_log_scale + ) diff --git a/docs/api/distributions/lognormal.md b/docs/api/distributions/lognormal.md new file mode 100644 index 0000000..5bd1d61 --- /dev/null +++ b/docs/api/distributions/lognormal.md @@ -0,0 +1,8 @@ +# LogNormal + +::: distreqx.distributions.LogNormal + options: + members: + - __init__ + - entropy +--- \ No newline at end of file diff --git a/mkdocs.yml b/mkdocs.yml index aa5cf4e..c9a7c50 100644 --- a/mkdocs.yml +++ b/mkdocs.yml @@ -112,6 +112,7 @@ nav: - 'api/distributions/gamma.md' - 'api/distributions/logistic.md' - 'api/distributions/independent.md' + - 'api/distributions/lognormal.md' - 'api/distributions/mixture_same_family.md' - 'api/distributions/uniform.md' - 'api/distributions/transformed.md' diff --git a/tests/lognormal_test.py b/tests/lognormal_test.py new file mode 100644 index 0000000..916c3b1 --- /dev/null +++ b/tests/lognormal_test.py @@ -0,0 +1,228 @@ +"""Tests for `log_normal.py`.""" + +from unittest import TestCase + +import jax +import jax.numpy as jnp +import numpy as np +from parameterized import parameterized # type: ignore + +from distreqx.distributions import LogNormal + + +class LogNormalTest(TestCase): + @parameterized.expand( + [ + ("1d std normal params", (jnp.array(0), jnp.array(1))), + ("2d std normal params", (jnp.zeros(2), jnp.ones(2))), + ("rank 2 std normal params", (jnp.zeros((3, 2)), jnp.ones((3, 2)))), + ] + ) + def test_event_shape(self, name, distr_params): + loc, scale = distr_params + self.assertEqual(loc.shape, LogNormal(loc, scale).event_shape) + + @parameterized.expand( + [ + ("1d std normal params", (jnp.array(0), jnp.array(1))), + ("2d std normal params", (jnp.zeros(2), jnp.ones(2))), + ("rank 2 std normal params", (jnp.zeros((3, 2)), jnp.ones((3, 2)))), + ] + ) + def test_sample_shape(self, name, distr_params): + distr_params = ( + jnp.asarray(distr_params[0], dtype=jnp.float32), + jnp.asarray(distr_params[1], dtype=jnp.float32), + ) + key = jax.random.key(0) + self.assertEqual( + distr_params[0].shape, + LogNormal(distr_params[0], distr_params[1]).sample(key).shape, + ) + + @parameterized.expand( + [ + ("1d std normal params", (jnp.array(0), jnp.array(1))), + ("2d std normal params", (jnp.zeros(2), jnp.ones(2))), + ("rank 2 std normal params", (jnp.zeros((3, 2)), jnp.ones((3, 2)))), + ] + ) + @jax.numpy_rank_promotion("raise") + def test_sample_and_log_prob(self, name, distr_params): + distr_params = ( + jnp.asarray(distr_params[0], dtype=jnp.float32), + jnp.asarray(distr_params[1], dtype=jnp.float32), + ) + key = jax.random.key(0) + dist = LogNormal(distr_params[0], distr_params[1]) + 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 normal params", (jnp.array(0), jnp.array(1))), + ("2d std normal params", (jnp.zeros(2), jnp.ones(2))), + ("rank 2 std normal params", (jnp.zeros((3, 2)), jnp.ones((3, 2)))), + ] + ) + def test_method_with_input(self, name, distr_params): + distr_params = ( + jnp.asarray(distr_params[0], dtype=jnp.float32), + jnp.asarray(distr_params[1], dtype=jnp.float32), + ) + # LogNormal domain is x > 0, so we use exp(loc) to ensure positive test values + value = jnp.exp(jnp.asarray(distr_params[0], dtype=jnp.float32)) + dist = LogNormal(distr_params[0], distr_params[1]) + for method in [ + "log_prob", + "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.0, 1.0), "entropy"), + ("mean", (0, 1), "mean"), + ("mean from 1d params", ([-1, 1], [1, 2]), "mean"), + ("variance", (0, 1), "variance"), + ("variance from np params", (np.ones(2), np.ones(2)), "variance"), + ("stddev", (0, 1), "stddev"), + ("stddev from rank 2 params", (np.ones((2, 3)), np.ones((2, 3))), "stddev"), + ("mode", (0, 1), "mode"), + ] + ) + def test_method(self, name, distr_params, function_string): + distr_params = ( + jnp.asarray(distr_params[0], dtype=jnp.float32), + jnp.asarray(distr_params[1], dtype=jnp.float32), + ) + dist = LogNormal(distr_params[0], distr_params[1]) + result = getattr(dist, function_string)() + self.assertEqual(distr_params[0].shape, result.shape) + + @parameterized.expand( + [ + ("no broadcast", ([0.0, 1.0, -0.5], [0.5, 1.0, 1.5])), + ] + ) + def test_median(self, name, distr_params): + distr_params = ( + jnp.asarray(distr_params[0], dtype=jnp.float32), + jnp.asarray(distr_params[1], dtype=jnp.float32), + ) + dist = LogNormal(distr_params[0], distr_params[1]) + # Median of LogNormal is exp(loc) + expected_median = jnp.exp(distr_params[0]) + np.testing.assert_allclose(dist.median(), expected_median, rtol=1e-3) + + @parameterized.expand( + [ + ("kl", "kl_divergence"), + ("cross-ent", "cross_entropy"), + ] + ) + def test_with_two_distributions(self, name, function_string): + dist1_kwargs = { + "loc": jnp.array(np.random.randn(3, 2)), + "scale": jnp.asarray([[0.8, 0.2], [0.1, 1.2], [1.4, 3.1]]), + } + dist2_kwargs = { + "loc": jnp.array(np.random.randn(3, 2)), + "scale": jnp.array(0.1 + np.random.rand(3, 2)), + } + dist1 = LogNormal(**dist1_kwargs) + dist2 = LogNormal(**dist2_kwargs) + + result = getattr(dist1, function_string)(dist2) + self.assertEqual(dist1_kwargs["loc"].shape, result.shape) + result = getattr(dist1, function_string)(dist1) + self.assertEqual(dist1_kwargs["loc"].shape, result.shape) + if name == "kl": + np.testing.assert_allclose( + jnp.zeros_like(dist1_kwargs["loc"]), result, atol=1e-6 + ) + elif name == "cross-ent": + np.testing.assert_allclose(dist1.entropy(), result, rtol=1e-5) + + @parameterized.expand( + [ + ("1d std normal params", (jnp.array(0.0), jnp.array(1.0))), + ("2d std normal params", (jnp.zeros(2), jnp.ones(2))), + ("rank 2 std normal params", (jnp.zeros((3, 2)), jnp.ones((3, 2)))), + ] + ) + def test_icdf_shape(self, name, distr_params): + distr_params = ( + jnp.asarray(distr_params[0], dtype=jnp.float32), + jnp.asarray(distr_params[1], dtype=jnp.float32), + ) + value = 0.5 * jnp.ones_like(distr_params[0]) + dist = LogNormal(distr_params[0], distr_params[1]) + result = dist.icdf(value) + self.assertEqual(value.shape, result.shape) + + def test_icdf_values(self): + loc = jnp.array([0.0, 1.0, -0.5]) + scale = jnp.array([1.0, 2.0, 0.5]) + dist = LogNormal(loc, scale) + + # icdf(cdf(x)) should be x. For LogNormal, x must be strictly positive. + x = jnp.array([0.5, 1.0, 0.8]) + np.testing.assert_allclose(dist.icdf(dist.cdf(x)), x, rtol=1e-5) + + # cdf(icdf(u)) should be u + u = jnp.array([0.1, 0.5, 0.9]) + np.testing.assert_allclose(dist.cdf(dist.icdf(u)), u, rtol=1e-5) + + def test_vmap_inputs(self): + def log_prob_sum(dist, x): + return dist.log_prob(x).sum() + + dist = LogNormal(jnp.arange(3 * 4 * 5).reshape((3, 4, 5)), jnp.ones((3, 4, 5))) + # For LogNormal, inputs must be strictly positive + x = jnp.ones((3, 4, 5)) + + 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): + return LogNormal(loc.sum(keepdims=True), scale.sum(keepdims=True)) + + loc = jnp.arange((3 * 4 * 5), dtype=jnp.float32).reshape((3, 4, 5)) + scale = jnp.ones((3, 4, 5), dtype=jnp.float32) + + actual = jax.vmap(summed_dist)(loc, scale) + expected = LogNormal( + loc.sum(axis=(1, 2), keepdims=True), scale.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_allclose(actual.log_prob(x), expected.log_prob(x), rtol=1e-6)