From 4a68d55c8ffe93d8c5e263479a02fb6e4842cc1e Mon Sep 17 00:00:00 2001 From: Will Dean Date: Tue, 7 Nov 2023 19:28:57 +0100 Subject: [PATCH 1/5] implement maxwell logp and logcdf --- docs/api_reference.rst | 1 + pymc_experimental/distributions/__init__.py | 3 +- pymc_experimental/distributions/continuous.py | 106 ++++++++++++++++++ .../tests/distributions/test_continuous.py | 25 ++++- 4 files changed, 133 insertions(+), 2 deletions(-) diff --git a/docs/api_reference.rst b/docs/api_reference.rst index 57882776..75839049 100644 --- a/docs/api_reference.rst +++ b/docs/api_reference.rst @@ -29,6 +29,7 @@ Distributions :toctree: generated/ Chi + Maxwell DiscreteMarkovChain GeneralizedPoisson GenExtreme diff --git a/pymc_experimental/distributions/__init__.py b/pymc_experimental/distributions/__init__.py index 2a89e7d0..5cb095b9 100644 --- a/pymc_experimental/distributions/__init__.py +++ b/pymc_experimental/distributions/__init__.py @@ -17,7 +17,7 @@ Experimental probability distributions for stochastic nodes in PyMC. """ -from pymc_experimental.distributions.continuous import Chi, GenExtreme +from pymc_experimental.distributions.continuous import Chi, GenExtreme, Maxwell from pymc_experimental.distributions.discrete import GeneralizedPoisson, Skellam from pymc_experimental.distributions.histogram_utils import histogram_approximation from pymc_experimental.distributions.multivariate import R2D2M2CP @@ -30,4 +30,5 @@ "R2D2M2CP", "histogram_approximation", "Chi", + "Maxwell", ] diff --git a/pymc_experimental/distributions/continuous.py b/pymc_experimental/distributions/continuous.py index 2e957b4f..9929173f 100644 --- a/pymc_experimental/distributions/continuous.py +++ b/pymc_experimental/distributions/continuous.py @@ -280,3 +280,109 @@ def __new__(cls, name, nu, **kwargs): @classmethod def dist(cls, nu, **kwargs): return CustomDist.dist(nu, dist=cls.chi_dist, class_name="Chi", **kwargs) + + +class Maxwell: + R""" + The Maxwell-Boltzmann distribution + + The pdf of this distribution is + + .. math:: + + f(x \mid a) = {\displaystyle {\sqrt {\frac {2}{\pi }}}\,{\frac {x^{2}}{a^{3}}}\,\exp \left({\frac {-x^{2}}{2a^{2}}}\right)} + + Read more about it on `Wikipedia `_ + + .. plot:: + :context: close-figs + + import matplotlib.pyplot as plt + import numpy as np + import scipy.stats as st + import arviz as az + plt.style.use('arviz-darkgrid') + x = np.linspace(0, 20, 200) + for a in [1, 2, 5]: + pdf = st.maxwell.pdf(x, scale=a) + plt.plot(x, pdf, label=r'$a$ = {}'.format(a)) + plt.xlabel('x', fontsize=12) + plt.ylabel('f(x)', fontsize=12) + plt.legend(loc=1) + plt.show() + + ======== ========================================================================= + Support :math:`x \in (0, \infty)` + Mean :math:`2a \sqrt{\frac{2}{\pi}}` + Variance :math:`\frac{a^2(3 \pi - 8)}{\pi}` + ======== ========================================================================= + + Parameters + ---------- + a : tensor_like of float + Scale parameter (a > 0). + + """ + + @staticmethod + def maxwell_dist(a: TensorVariable, size: TensorVariable) -> TensorVariable: + return Chi.dist(nu=3, size=size) * a + + @staticmethod + def maxwell_logp(value, a): + res = ( + pt.log(pt.sqrt(2.0 / pt.pi)) + + 2 * pt.log(value) + - 3 * pt.log(a) + - pt.sqr(value) / (2 * pt.sqr(a)) + ) + res = pt.switch( + value > 0, + res, + -pt.inf, + ) + return check_parameters( + res, + a > 0, + msg="a > 0", + ) + + @staticmethod + def maxwell_logcdf(value, a): + res = pt.erf(value / (pt.sqrt(2) * a)) - pt.sqrt(2.0 / pt.pi) * (value / a) * pt.exp( + -pt.sqr(value) / (2 * pt.sqr(a)) + ) + res = pt.log(res) + + res = pt.switch( + value > 0, + res, + -pt.inf, + ) + return check_parameters( + res, + a > 0, + msg="a > 0", + ) + + def __new__(cls, name, a, **kwargs): + return CustomDist( + name, + a, + dist=cls.maxwell_dist, + logp=cls.maxwell_logp, + logcdf=cls.maxwell_logcdf, + class_name="Maxwell", + **kwargs, + ) + + @classmethod + def dist(cls, a, **kwargs): + return CustomDist.dist( + a, + dist=cls.maxwell_dist, + logp=cls.maxwell_logp, + logcdf=cls.maxwell_logcdf, + class_name="Maxwell", + **kwargs, + ) diff --git a/pymc_experimental/tests/distributions/test_continuous.py b/pymc_experimental/tests/distributions/test_continuous.py index 891e7ab3..5df4aef1 100644 --- a/pymc_experimental/tests/distributions/test_continuous.py +++ b/pymc_experimental/tests/distributions/test_continuous.py @@ -33,7 +33,7 @@ ) # the distributions to be tested -from pymc_experimental.distributions import Chi, GenExtreme +from pymc_experimental.distributions import Chi, GenExtreme, Maxwell class TestGenExtremeClass: @@ -159,3 +159,26 @@ def test_logcdf(self): {"nu": Rplus}, lambda value, nu: sp.chi.logcdf(value, df=nu), ) + + +class TestMaxwell: + """ + Wrapper class so that tests of experimental additions can be dropped into + PyMC directly on adoption. + """ + + def test_logp(self): + check_logp( + Maxwell, + Rplus, + {"a": Rplus}, + lambda value, a: sp.maxwell.logpdf(value, scale=a), + ) + + def test_logcdf(self): + check_logcdf( + Maxwell, + Rplus, + {"a": Rplus}, + lambda value, a: sp.maxwell.logcdf(value, scale=a), + ) From 7e54f50fd1b531d2a393af8fcafc86d60929bfcc Mon Sep 17 00:00:00 2001 From: Will Dean Date: Thu, 9 Nov 2023 18:23:54 +0100 Subject: [PATCH 2/5] remove the logp and logcdf --- pymc_experimental/distributions/continuous.py | 41 ------------------- 1 file changed, 41 deletions(-) diff --git a/pymc_experimental/distributions/continuous.py b/pymc_experimental/distributions/continuous.py index 9929173f..f1747f4e 100644 --- a/pymc_experimental/distributions/continuous.py +++ b/pymc_experimental/distributions/continuous.py @@ -328,50 +328,11 @@ class Maxwell: def maxwell_dist(a: TensorVariable, size: TensorVariable) -> TensorVariable: return Chi.dist(nu=3, size=size) * a - @staticmethod - def maxwell_logp(value, a): - res = ( - pt.log(pt.sqrt(2.0 / pt.pi)) - + 2 * pt.log(value) - - 3 * pt.log(a) - - pt.sqr(value) / (2 * pt.sqr(a)) - ) - res = pt.switch( - value > 0, - res, - -pt.inf, - ) - return check_parameters( - res, - a > 0, - msg="a > 0", - ) - - @staticmethod - def maxwell_logcdf(value, a): - res = pt.erf(value / (pt.sqrt(2) * a)) - pt.sqrt(2.0 / pt.pi) * (value / a) * pt.exp( - -pt.sqr(value) / (2 * pt.sqr(a)) - ) - res = pt.log(res) - - res = pt.switch( - value > 0, - res, - -pt.inf, - ) - return check_parameters( - res, - a > 0, - msg="a > 0", - ) - def __new__(cls, name, a, **kwargs): return CustomDist( name, a, dist=cls.maxwell_dist, - logp=cls.maxwell_logp, - logcdf=cls.maxwell_logcdf, class_name="Maxwell", **kwargs, ) @@ -381,8 +342,6 @@ def dist(cls, a, **kwargs): return CustomDist.dist( a, dist=cls.maxwell_dist, - logp=cls.maxwell_logp, - logcdf=cls.maxwell_logcdf, class_name="Maxwell", **kwargs, ) From 8dae542ef5a3e85090db8dea629ca2be124c9f67 Mon Sep 17 00:00:00 2001 From: Will Dean Date: Fri, 10 Nov 2023 07:40:52 +0100 Subject: [PATCH 3/5] ensure shape --- pymc_experimental/distributions/continuous.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/pymc_experimental/distributions/continuous.py b/pymc_experimental/distributions/continuous.py index f1747f4e..666d97ee 100644 --- a/pymc_experimental/distributions/continuous.py +++ b/pymc_experimental/distributions/continuous.py @@ -326,6 +326,8 @@ class Maxwell: @staticmethod def maxwell_dist(a: TensorVariable, size: TensorVariable) -> TensorVariable: + if rv_size_is_none(size): + size = a.shape return Chi.dist(nu=3, size=size) * a def __new__(cls, name, a, **kwargs): From b1ad63a8aacb821a69ea82f5701f68dba1d2a047 Mon Sep 17 00:00:00 2001 From: Will Dean Date: Fri, 10 Nov 2023 08:28:05 +0100 Subject: [PATCH 4/5] dont check outside paradomain --- pymc_experimental/distributions/continuous.py | 1 + pymc_experimental/tests/distributions/test_continuous.py | 2 ++ 2 files changed, 3 insertions(+) diff --git a/pymc_experimental/distributions/continuous.py b/pymc_experimental/distributions/continuous.py index 666d97ee..61dffed2 100644 --- a/pymc_experimental/distributions/continuous.py +++ b/pymc_experimental/distributions/continuous.py @@ -328,6 +328,7 @@ class Maxwell: def maxwell_dist(a: TensorVariable, size: TensorVariable) -> TensorVariable: if rv_size_is_none(size): size = a.shape + return Chi.dist(nu=3, size=size) * a def __new__(cls, name, a, **kwargs): diff --git a/pymc_experimental/tests/distributions/test_continuous.py b/pymc_experimental/tests/distributions/test_continuous.py index 5df4aef1..a540ddd3 100644 --- a/pymc_experimental/tests/distributions/test_continuous.py +++ b/pymc_experimental/tests/distributions/test_continuous.py @@ -173,6 +173,7 @@ def test_logp(self): Rplus, {"a": Rplus}, lambda value, a: sp.maxwell.logpdf(value, scale=a), + skip_paramdomain_outside_edge_test=True, ) def test_logcdf(self): @@ -181,4 +182,5 @@ def test_logcdf(self): Rplus, {"a": Rplus}, lambda value, a: sp.maxwell.logcdf(value, scale=a), + skip_paramdomain_outside_edge_test=True, ) From 052ff379eb263442d67767b62c2a766a0a804ecd Mon Sep 17 00:00:00 2001 From: Will Dean Date: Fri, 10 Nov 2023 18:12:43 +0100 Subject: [PATCH 5/5] ensure a is positive --- pymc_experimental/distributions/continuous.py | 3 +++ pymc_experimental/tests/distributions/test_continuous.py | 2 -- 2 files changed, 3 insertions(+), 2 deletions(-) diff --git a/pymc_experimental/distributions/continuous.py b/pymc_experimental/distributions/continuous.py index 61dffed2..dcf9b775 100644 --- a/pymc_experimental/distributions/continuous.py +++ b/pymc_experimental/distributions/continuous.py @@ -28,6 +28,7 @@ from pymc.distributions.dist_math import check_parameters from pymc.distributions.distribution import Continuous from pymc.distributions.shape_utils import rv_size_is_none +from pymc.logprob.utils import CheckParameterValue from pymc.pytensorf import floatX from pytensor.tensor.random.op import RandomVariable from pytensor.tensor.variable import TensorVariable @@ -329,6 +330,8 @@ def maxwell_dist(a: TensorVariable, size: TensorVariable) -> TensorVariable: if rv_size_is_none(size): size = a.shape + a = CheckParameterValue("a > 0")(a, pt.all(pt.gt(a, 0))) + return Chi.dist(nu=3, size=size) * a def __new__(cls, name, a, **kwargs): diff --git a/pymc_experimental/tests/distributions/test_continuous.py b/pymc_experimental/tests/distributions/test_continuous.py index a540ddd3..5df4aef1 100644 --- a/pymc_experimental/tests/distributions/test_continuous.py +++ b/pymc_experimental/tests/distributions/test_continuous.py @@ -173,7 +173,6 @@ def test_logp(self): Rplus, {"a": Rplus}, lambda value, a: sp.maxwell.logpdf(value, scale=a), - skip_paramdomain_outside_edge_test=True, ) def test_logcdf(self): @@ -182,5 +181,4 @@ def test_logcdf(self): Rplus, {"a": Rplus}, lambda value, a: sp.maxwell.logcdf(value, scale=a), - skip_paramdomain_outside_edge_test=True, )