diff --git a/docs/source/api/distributions/continuous.rst b/docs/source/api/distributions/continuous.rst index 98f1c97ca1..1cf6eb71ce 100644 --- a/docs/source/api/distributions/continuous.rst +++ b/docs/source/api/distributions/continuous.rst @@ -27,6 +27,7 @@ Continuous Logistic LogitNormal LogNormal + Maxwell Moyal Normal Pareto diff --git a/pymc/distributions/__init__.py b/pymc/distributions/__init__.py index bc3d9c7863..e0c072a11d 100644 --- a/pymc/distributions/__init__.py +++ b/pymc/distributions/__init__.py @@ -36,6 +36,7 @@ LogitNormal, LogNormal, Lognormal, + Maxwell, Moyal, Normal, Pareto, @@ -132,6 +133,7 @@ "Bound", "LogNormal", "Lognormal", + "Maxwell", "HalfStudentT", "ChiSquared", "HalfNormal", diff --git a/pymc/distributions/continuous.py b/pymc/distributions/continuous.py index 1df2bd6efe..88e562566f 100644 --- a/pymc/distributions/continuous.py +++ b/pymc/distributions/continuous.py @@ -47,6 +47,7 @@ laplace, logistic, lognormal, + maxwell, normal, pareto, triangular, @@ -114,6 +115,7 @@ def polyagamma_cdf(*args, **kwargs): "Weibull", "HalfStudentT", "LogNormal", + "Maxwell", "ChiSquared", "HalfNormal", "Wald", @@ -1735,6 +1737,94 @@ def icdf(value, mu, sigma): Lognormal = LogNormal +class Maxwell(PositiveContinuous): + r""" + Univariate Maxwell (or Maxwell-Boltzmann) log-likelihood. + + The pdf of this distribution is + + .. math:: + + f(x \mid \mu, \sigma) = + \sqrt{\frac{2}{\pi}} + \frac{(x-\mu)^2 \exp\left\{-(x-\mu)^2/(2\sigma^2)\right\}}{\sigma^3} + + .. 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, 10, 1000) + mus = [0., 0., 0., 2.] + sigmas = [0.4, 1., 2., 0.4] + for mu, sigma in zip(mus, sigmas): + pdf = st.maxwell.pdf(x, mu, sigma) + plt.plot(x, pdf, label=r'$\mu$ = {}, $\sigma$ = {}'.format(mu, sigma)) + 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:`2\sigma\sqrt{\frac{2}{\pi}}` + Variance :math:`\sigma^2(3\pi - 8)/\pi` + ======== ========================================================================= + + Parameters + ---------- + mu : tensor_like of float, default 0 + Location parameter. + sigma : tensor_like of float, default 1 + Scale parameter. (sigma > 0). + + Examples + -------- + .. code-block:: python + + with pm.Model(): + x = pm.Maxwell('x', mu=2, sigma=5) + """ + rv_op = maxwell + + @classmethod + def dist(cls, mu=0, sigma=1, **kwargs): + sigma = pt.as_tensor_variable(sigma) + return super().dist([mu, sigma], **kwargs) + + def moment(rv, size, mu, sigma): + mu, _ = pt.broadcast_arrays(mu, sigma) + if not rv_size_is_none(size): + mu = pt.full(size, mu) + return mu + + def logp(value, mu, sigma): + res = ( + -0.5 * pt.pow((value - mu) / sigma, 2) + + pt.log(pt.sqrt(2.0 / np.pi)) + + 2.0 * pt.log((value - mu) / sigma) + - pt.log(sigma) + ) + return check_parameters( + res, + sigma > 0, + msg="sigma > 0", + ) + + def logcdf(value, mu, sigma): + res = pt.log(pt.gammainc(1.5, 0.5 * pt.pow(value / sigma, 2))) + pt.log( + 2.0 / pt.sqrt(np.pi) + ) + return check_parameters( + res, + sigma > 0, + msg="sigma > 0", + ) + + class StudentTRV(RandomVariable): name = "studentt" ndim_supp = 0 diff --git a/tests/distributions/test_continuous.py b/tests/distributions/test_continuous.py index bab4e281c1..8e03f1a6af 100644 --- a/tests/distributions/test_continuous.py +++ b/tests/distributions/test_continuous.py @@ -529,6 +529,20 @@ def test_lognormal(self): decimal=select_by_precision(float64=4, float32=3), ) + def test_maxwell(self): + check_logp( + pm.Maxwell, + R, + {"mu": R, "sigma": Rplusbig}, + lambda value, mu, sigma: st.maxwell.logpdf(value, loc=mu, scale=sigma), + ) + check_logcdf( + pm.Maxwell, + R, + {"mu": R, "sigma": Rplusbig}, + lambda value, mu, sigma: st.maxwell.logcdf(value, loc=mu, scale=sigma), + ) + def test_studentt_logp(self): check_logp( pm.StudentT, @@ -1241,6 +1255,20 @@ def test_lognormal_moment(self, mu, sigma, size, expected): pm.LogNormal("x", mu=mu, sigma=sigma, size=size) assert_moment_is_expected(model, expected) + @pytest.mark.parametrize( + "mu, sigma, size, expected", + [ + (0, 1, None, 0), + (0, np.ones(5), None, np.zeros(5)), + (np.arange(5), 1, None, np.arange(5)), + (np.arange(5), np.arange(1, 6), (2, 5), np.full((2, 5))), + ], + ) + def test_maxwell_moment(self, mu, sigma, size, expected): + with pm.Model() as model: + pm.Maxwell("x", mu=mu, sigma=sigma, size=size) + assert_moment_is_expected(model, expected) + @pytest.mark.parametrize( "beta, size, expected", [ @@ -1810,6 +1838,19 @@ class TestGumbel(BaseTestDistributionRandom): ] +class TestMaxwell(BaseTestDistributionRandom): + pymc_dist = pm.Maxwell + pymc_dist_params = {"loc": 1.5, "sigma": 3.0} + expected_rv_op_params = {"loc": 1.5, "sigma": 3.0} + reference_dist_params = {"loc": 1.5, "sigma": 3.0} + reference_dist = seeded_scipy_distribution_builder("maxwell") + checks_to_run = [ + "check_pymc_params_match_rv_op", + "check_pymc_draws_match_reference", + "check_rv_size", + ] + + class TestStudentT(BaseTestDistributionRandom): pymc_dist = pm.StudentT pymc_dist_params = {"nu": 5.0, "mu": -1.0, "sigma": 2.0}