diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml index 40473a014c..75f7bac323 100644 --- a/.github/workflows/tests.yml +++ b/.github/workflows/tests.yml @@ -76,6 +76,7 @@ jobs: tests/distributions/test_continuous.py tests/distributions/test_multivariate.py tests/distributions/moments/test_means.py + tests/distributions/moments/test_entropy.py - | tests/distributions/test_censored.py diff --git a/pymc/distributions/moments/__init__.py b/pymc/distributions/moments/__init__.py index b61e8b6400..19c3532dbb 100644 --- a/pymc/distributions/moments/__init__.py +++ b/pymc/distributions/moments/__init__.py @@ -14,6 +14,7 @@ """Moments dispatchers for pymc random variables.""" +from pymc.distributions.moments.entropy import entropy from pymc.distributions.moments.means import mean -__all__ = ["mean"] +__all__ = ["entropy", "mean"] diff --git a/pymc/distributions/moments/entropy.py b/pymc/distributions/moments/entropy.py new file mode 100644 index 0000000000..5df8a1c07a --- /dev/null +++ b/pymc/distributions/moments/entropy.py @@ -0,0 +1,261 @@ +# Copyright 2024 - present The PyMC Developers +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Entropy dispatcher for pymc random variables. + +The entropy of a random variable is the differential entropy for continuous +variables and the Shannon entropy (in nats) for discrete variables, matching +the convention of :meth:`scipy.stats.rv_continuous.entropy`. +""" + +from functools import singledispatch + +import numpy as np + +from pytensor import tensor as pt +from pytensor.tensor.random.basic import ( + BernoulliRV, + BetaRV, + CategoricalRV, + CauchyRV, + DirichletRV, + ExponentialRV, + GammaRV, + GeometricRV, + GumbelRV, + HalfNormalRV, + InvGammaRV, + LaplaceRV, + LogisticRV, + LogNormalRV, + MvNormalRV, + NormalRV, + ParetoRV, + StudentTRV, + TriangularRV, + UniformRV, + VonMisesRV, +) +from pytensor.tensor.variable import TensorVariable + +from pymc.distributions.continuous import ( + AsymmetricLaplaceRV, + HalfCauchyRV, + MoyalRV, + WeibullBetaRV, +) +from pymc.distributions.discrete import DiscreteUniformRV +from pymc.distributions.shape_utils import maybe_resize, rv_size_is_none + +__all__ = ["entropy"] + + +def _betaln(a, b): + return pt.gammaln(a) + pt.gammaln(b) - pt.gammaln(a + b) + + +@singledispatch +def _entropy(op, rv, *rv_inputs) -> TensorVariable: + raise NotImplementedError(f"Variable {rv} of type {op} has no entropy implementation.") + + +def entropy(rv: TensorVariable) -> TensorVariable: + """Compute the entropy of a random variable. + + The entropy is the differential entropy for continuous distributions and + the Shannon entropy (in nats) for discrete distributions. This matches the + convention used by ``scipy.stats``'s ``entropy`` method. + + The only parameter to this function is the RandomVariable + for which the entropy is to be derived. + """ + return _entropy(rv.owner.op, rv, *rv.owner.inputs) + + +# --- Continuous univariate --- + + +@_entropy.register(NormalRV) +def normal_entropy(op, rv, rng, size, mu, sigma): + return maybe_resize(0.5 * pt.log(2 * np.pi * np.e) + pt.log(sigma), size) + + +@_entropy.register(UniformRV) +def uniform_entropy(op, rv, rng, size, lower, upper): + return maybe_resize(pt.log(upper - lower), size) + + +@_entropy.register(ExponentialRV) +def exponential_entropy(op, rv, rng, size, mu): + # ``mu`` is the mean (scale) of the distribution + return maybe_resize(1 + pt.log(mu), size) + + +@_entropy.register(LaplaceRV) +def laplace_entropy(op, rv, rng, size, mu, b): + return maybe_resize(1 + pt.log(2 * b), size) + + +@_entropy.register(AsymmetricLaplaceRV) +def asymmetric_laplace_entropy(op, rv, rng, size, b, kappa, mu): + # scale = 1 / b + return maybe_resize(1 + pt.log((kappa + 1 / kappa) / b), size) + + +@_entropy.register(CauchyRV) +def cauchy_entropy(op, rv, rng, size, alpha, beta): + return maybe_resize(pt.log(4 * np.pi * beta), size) + + +@_entropy.register(HalfCauchyRV) +def halfcauchy_entropy(op, rv, rng, size, beta): + return maybe_resize(pt.log(2 * np.pi * beta), size) + + +@_entropy.register(HalfNormalRV) +def halfnormal_entropy(op, rv, rng, size, loc, sigma): + return maybe_resize(0.5 * pt.log(np.pi * sigma**2 / 2) + 0.5, size) + + +@_entropy.register(GammaRV) +def gamma_entropy(op, rv, rng, size, alpha, inv_beta): + # ``inv_beta`` is the scale (1 / rate) + return maybe_resize( + alpha + pt.log(inv_beta) + pt.gammaln(alpha) + (1 - alpha) * pt.digamma(alpha), + size, + ) + + +@_entropy.register(InvGammaRV) +def invgamma_entropy(op, rv, rng, size, alpha, beta): + # ``beta`` is the scale + return maybe_resize( + alpha + pt.log(beta) + pt.gammaln(alpha) - (1 + alpha) * pt.digamma(alpha), + size, + ) + + +@_entropy.register(BetaRV) +def beta_entropy(op, rv, rng, size, alpha, beta): + return maybe_resize( + _betaln(alpha, beta) + - (alpha - 1) * pt.digamma(alpha) + - (beta - 1) * pt.digamma(beta) + + (alpha + beta - 2) * pt.digamma(alpha + beta), + size, + ) + + +@_entropy.register(LogisticRV) +def logistic_entropy(op, rv, rng, size, mu, s): + return maybe_resize(pt.log(s) + 2, size) + + +@_entropy.register(LogNormalRV) +def lognormal_entropy(op, rv, rng, size, mu, sigma): + return maybe_resize(mu + 0.5 * pt.log(2 * np.pi * np.e * sigma**2), size) + + +@_entropy.register(GumbelRV) +def gumbel_entropy(op, rv, rng, size, mu, beta): + return maybe_resize(pt.log(beta) + np.euler_gamma + 1, size) + + +@_entropy.register(ParetoRV) +def pareto_entropy(op, rv, rng, size, alpha, m): + return maybe_resize(pt.log(m / alpha) + 1 / alpha + 1, size) + + +@_entropy.register(WeibullBetaRV) +def weibull_entropy(op, rv, rng, size, alpha, beta): + # ``alpha`` is the shape, ``beta`` is the scale + return maybe_resize(np.euler_gamma * (1 - 1 / alpha) + pt.log(beta / alpha) + 1, size) + + +@_entropy.register(StudentTRV) +def studentt_entropy(op, rv, rng, size, nu, mu, sigma): + return maybe_resize( + pt.log(sigma) + + 0.5 * (nu + 1) * (pt.digamma((nu + 1) / 2) - pt.digamma(nu / 2)) + + pt.log(pt.sqrt(nu)) + + _betaln(nu / 2, 0.5), + size, + ) + + +@_entropy.register(TriangularRV) +def triangular_entropy(op, rv, rng, size, lower, c, upper): + return maybe_resize(0.5 + pt.log((upper - lower) / 2), size) + + +@_entropy.register(MoyalRV) +def moyal_entropy(op, rv, rng, size, mu, sigma): + return maybe_resize( + pt.log(sigma) + 0.5 * pt.log(2 * np.pi) + 0.5 * (np.euler_gamma + np.log(2) + 1), + size, + ) + + +@_entropy.register(VonMisesRV) +def vonmises_entropy(op, rv, rng, size, mu, kappa): + return maybe_resize( + pt.log(2 * np.pi * pt.i0(kappa)) - kappa * pt.i1(kappa) / pt.i0(kappa), size + ) + + +# --- Discrete univariate --- + + +@_entropy.register(BernoulliRV) +def bernoulli_entropy(op, rv, rng, size, p): + return maybe_resize(-p * pt.log(p) - (1 - p) * pt.log1p(-p), size) + + +@_entropy.register(GeometricRV) +def geometric_entropy(op, rv, rng, size, p): + return maybe_resize((-(1 - p) * pt.log1p(-p) - p * pt.log(p)) / p, size) + + +@_entropy.register(DiscreteUniformRV) +def discrete_uniform_entropy(op, rv, rng, size, lower, upper): + return maybe_resize(pt.log(upper - lower + 1), size) + + +@_entropy.register(CategoricalRV) +def categorical_entropy(op, rv, rng, size, p): + return maybe_resize(-pt.sum(p * pt.log(p), axis=-1), size) + + +# --- Multivariate --- + + +@_entropy.register(MvNormalRV) +def mvnormal_entropy(op, rv, rng, size, mu, cov): + k = cov.shape[-1] + _, logdet = pt.linalg.slogdet(cov) + res = 0.5 * k * pt.log(2 * np.pi * np.e) + 0.5 * logdet + if rv_size_is_none(size): + return res + return maybe_resize(res, size) + + +@_entropy.register(DirichletRV) +def dirichlet_entropy(op, rv, rng, size, a): + a0 = pt.sum(a, axis=-1) + k = a.shape[-1] + log_beta = pt.sum(pt.gammaln(a), axis=-1) - pt.gammaln(a0) + res = log_beta + (a0 - k) * pt.digamma(a0) - pt.sum((a - 1) * pt.digamma(a), axis=-1) + if rv_size_is_none(size): + return res + return maybe_resize(res, size) diff --git a/tests/distributions/moments/test_entropy.py b/tests/distributions/moments/test_entropy.py new file mode 100644 index 0000000000..7d42564f23 --- /dev/null +++ b/tests/distributions/moments/test_entropy.py @@ -0,0 +1,228 @@ +# Copyright 2024 - present The PyMC Developers +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import numpy as np +import pytest + +from pytensor import function +from pytensor import tensor as pt +from pytensor.compile.mode import Mode +from scipy import stats +from scipy.stats import ( + bernoulli, + beta, + cauchy, + chi2, + dirichlet, + expon, + gamma, + geom, + gumbel_r, + halfcauchy, + halfnorm, + invgamma, + laplace, + laplace_asymmetric, + logistic, + lognorm, + moyal, + multivariate_normal, + norm, + pareto, + randint, + t, + triang, + uniform, + vonmises, + weibull_min, +) + +from pymc import ( + AsymmetricLaplace, + Bernoulli, + Beta, + Binomial, + Categorical, + Cauchy, + ChiSquared, + Dirichlet, + DiscreteUniform, + ExGaussian, + Exponential, + Flat, + Gamma, + Geometric, + Gumbel, + HalfCauchy, + HalfFlat, + HalfNormal, + HalfStudentT, + InverseGamma, + Kumaraswamy, + Laplace, + Logistic, + LogitNormal, + LogNormal, + Moyal, + MvNormal, + MvStudentT, + Normal, + Pareto, + Poisson, + Rice, + SkewNormal, + StudentT, + Triangular, + Uniform, + VonMises, + Wald, + Weibull, +) +from pymc.distributions.moments.entropy import entropy + +mode = Mode(linker="py", optimizer=None) + + +@pytest.mark.parametrize( + ["dist", "scipy_equiv", "dist_params", "scipy_params"], + [ + [ + AsymmetricLaplace, + laplace_asymmetric, + {"kappa": 2, "mu": 0.2, "b": 1 / 1.2}, + {"kappa": 2, "loc": 0.2, "scale": 1.2}, + ], + [Bernoulli, bernoulli, {"p": 0.6}, {"p": 0.6}], + [Beta, beta, {"alpha": 3, "beta": 2}, {"a": 3, "b": 2}], + [Cauchy, cauchy, {"alpha": 2, "beta": 1.5}, {"loc": 2, "scale": 1.5}], + [ChiSquared, chi2, {"nu": 6}, {"df": 6}], + [DiscreteUniform, randint, {"lower": 2, "upper": 8}, {"low": 2, "high": 9}], + [Exponential, expon, {"lam": 0.5}, {"scale": 2}], + [Gamma, gamma, {"alpha": 4, "beta": 3}, {"a": 4, "scale": 1 / 3}], + [Geometric, geom, {"p": 0.1}, {"p": 0.1}], + [Gumbel, gumbel_r, {"mu": 2, "beta": 1}, {"loc": 2, "scale": 1}], + [HalfCauchy, halfcauchy, {"beta": 1.5}, {"scale": 1.5}], + [HalfNormal, halfnorm, {"sigma": 3}, {"scale": 3}], + [InverseGamma, invgamma, {"alpha": 2, "beta": 2}, {"a": 2, "scale": 2}], + [Laplace, laplace, {"mu": 2, "b": 2}, {"loc": 2, "scale": 2}], + [Logistic, logistic, {"mu": 2, "s": 1.5}, {"loc": 2, "scale": 1.5}], + [LogNormal, lognorm, {"mu": 0.3, "sigma": 0.6}, {"scale": np.exp(0.3), "s": 0.6}], + [Moyal, moyal, {"mu": 2, "sigma": 2}, {"loc": 2, "scale": 2}], + [Normal, norm, {"mu": 2, "sigma": 3}, {"loc": 2, "scale": 3}], + [Pareto, pareto, {"alpha": 5, "m": 2}, {"b": 5, "scale": 2}], + [StudentT, t, {"nu": 6, "mu": 0, "sigma": 2}, {"df": 6, "loc": 0, "scale": 2}], + [Triangular, triang, {"lower": 0, "c": 0.5, "upper": 2}, {"c": 0.25, "loc": 0, "scale": 2}], + [Uniform, uniform, {"lower": -1, "upper": 4}, {"loc": -1, "scale": 5}], + [VonMises, vonmises, {"mu": 1, "kappa": 2.5}, {"loc": 1, "kappa": 2.5}], + [Weibull, weibull_min, {"alpha": 1.5, "beta": 2}, {"c": 1.5, "scale": 2}], + ], +) +def test_entropy_equal_to_scipy(dist, scipy_equiv, dist_params, scipy_params): + rv = dist.dist(**dist_params) + pymc_entropy = entropy(rv).eval(mode=mode) + scipy_entropy = scipy_equiv(**scipy_params).entropy() + assert np.asarray(pymc_entropy).shape == np.asarray(scipy_entropy).shape + np.testing.assert_allclose(pymc_entropy, scipy_entropy, rtol=1e-6) + + # entropy of a batched distribution broadcasts over the batch dimension + pymc_entropy_tiled = entropy(dist.dist(shape=(3,), **dist_params)).eval() + np.testing.assert_allclose(pymc_entropy_tiled, np.tile(pymc_entropy, 3), rtol=1e-6) + + +@pytest.mark.parametrize( + ["dist", "dist_params", "expected"], + [ + [Categorical, {"p": [0.1, 0.2, 0.3, 0.4]}, stats.entropy([0.1, 0.2, 0.3, 0.4])], + ], +) +def test_entropy_equal_expected(dist, dist_params, expected): + rv = dist.dist(**dist_params) + pymc_entropy = entropy(rv).eval(mode=mode) + np.testing.assert_allclose(pymc_entropy, expected, rtol=1e-6) + + +@pytest.mark.parametrize( + ["dist", "dist_params", "scipy_entropy"], + [ + [ + MvNormal, + { + "mu": np.zeros(3), + "cov": np.array([[2.0, 0.3, 0.1], [0.3, 1.5, 0.2], [0.1, 0.2, 1.0]]), + }, + multivariate_normal( + np.zeros(3), np.array([[2.0, 0.3, 0.1], [0.3, 1.5, 0.2], [0.1, 0.2, 1.0]]) + ).entropy(), + ], + [ + Dirichlet, + {"a": np.array([2.0, 3.0, 4.0, 1.0])}, + dirichlet(np.array([2.0, 3.0, 4.0, 1.0])).entropy(), + ], + ], +) +def test_entropy_multivariate(dist, dist_params, scipy_entropy): + rv = dist.dist(**dist_params) + pymc_entropy = entropy(rv).eval(mode=mode) + assert np.asarray(pymc_entropy).shape == () + np.testing.assert_allclose(pymc_entropy, scipy_entropy, rtol=1e-6) + + # a batch of independent multivariate distributions yields one entropy per batch element + pymc_entropy_batched = entropy(dist.dist(size=(4,), **dist_params)).eval() + np.testing.assert_allclose(pymc_entropy_batched, np.full(4, pymc_entropy), rtol=1e-6) + + +@pytest.mark.parametrize( + ["dist", "dist_params"], + [ + [Binomial, {"n": 5, "p": 0.6}], + [ExGaussian, {"mu": 0, "sigma": 1, "nu": 1}], + [Flat, {}], + [HalfFlat, {}], + [HalfStudentT, {"nu": 3, "sigma": 1}], + [Kumaraswamy, {"a": 2, "b": 2}], + [LogitNormal, {"mu": 0, "sigma": 1}], + [MvStudentT, {"mu": np.zeros(3), "scale": np.eye(3), "nu": 4}], + [Poisson, {"mu": 3}], + [Rice, {"nu": 1, "sigma": 1}], + [SkewNormal, {"mu": 0, "sigma": 1, "alpha": 2}], + [Wald, {"mu": 1, "lam": 1}], + ], +) +def test_no_entropy(dist, dist_params): + with pytest.raises(NotImplementedError): + entropy(dist.dist(**dist_params)) + + +@pytest.mark.parametrize( + ["dist", "param", "value"], + [ + # A regularization-by-entropy term (e.g. in RL policies) requires the + # entropy to be differentiable w.r.t. the distribution parameters. + [lambda p: Normal.dist(mu=0.0, sigma=p), "sigma", 2.0], + [lambda p: Gamma.dist(alpha=p, beta=1.0), "alpha", 3.0], + [lambda p: Beta.dist(alpha=p, beta=2.0), "alpha", 3.0], + [lambda p: StudentT.dist(nu=p, mu=0.0, sigma=1.0), "nu", 5.0], + [lambda p: VonMises.dist(mu=0.0, kappa=p), "kappa", 2.5], + ], +) +def test_entropy_is_differentiable(dist, param, value): + p = pt.scalar(param) + grad = pt.grad(entropy(dist(p)).sum(), p) + grad_fn = function([p], grad, mode=mode, on_unused_input="ignore") + analytic = grad_fn(value) + h = 1e-5 + entropy_fn = function([p], entropy(dist(p)), mode=mode, on_unused_input="ignore") + finite_diff = (entropy_fn(value + h) - entropy_fn(value - h)) / (2 * h) + np.testing.assert_allclose(analytic, finite_diff, rtol=1e-4)