Implementation of Maxwell Distribution

docs/source/api/distributions/continuous.rst

@@ -27,6 +27,7 @@ Continuous
    Logistic
    LogitNormal
    LogNormal
+   Maxwell
    Moyal
    Normal
    Pareto

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",

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

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}