Implement Symmetric Multivariate Laplace distribution

Author: ricardoV94

pymc_experimental/distributions/__init__.py

@@ -24,6 +24,7 @@
     Skellam,
 )
 from pymc_experimental.distributions.histogram_utils import histogram_approximation
+from pymc_experimental.distributions.multivariate.laplace import MvLaplace
 from pymc_experimental.distributions.multivariate.r2d2m2cp import R2D2M2CP
 from pymc_experimental.distributions.timeseries import DiscreteMarkovChain
 
@@ -32,6 +33,7 @@
     "DiscreteMarkovChain",
     "GeneralizedPoisson",
     "GenExtreme",
+    "MvLaplace",
     "R2D2M2CP",
     "Skellam",
     "histogram_approximation",

pymc_experimental/distributions/multivariate/laplace.py

@@ -0,0 +1,126 @@
+import numpy as np
+import pytensor.tensor as pt
+import scipy
+
+from pymc.distributions.dist_math import check_parameters
+from pymc.distributions.distribution import Continuous, SymbolicRandomVariable
+from pymc.distributions.multivariate import quaddist_chol, quaddist_matrix
+from pymc.distributions.shape_utils import implicit_size_from_params, rv_size_is_none
+from pymc.pytensorf import normalize_rng_param
+from pytensor.gradient import grad_not_implemented
+from pytensor.scalar import BinaryScalarOp, upgrade_to_float
+from pytensor.tensor.elemwise import Elemwise
+from pytensor.tensor.random.utils import normalize_size_param
+
+
+class Kv(BinaryScalarOp):
+    """
+    Modified Bessel function of the second kind of real order v.
+    """
+
+    nfunc_spec = ("scipy.special.kv", 2, 1)
+
+    @staticmethod
+    def st_impl(v, x):
+        return scipy.special.kv(v, x)
+
+    def impl(self, v, x):
+        return self.st_impl(v, x)
+
+    def grad(self, inputs, grads):
+        v, x = inputs
+        (gz,) = grads
+        return [grad_not_implemented(self, 0, v), gz * scalar_kvp(v, x)]
+
+    def c_code(self, *args, **kwargs):
+        raise NotImplementedError()
+
+
+kv = Elemwise(Kv(upgrade_to_float, name="kv"))
+
+
+class Kvp(BinaryScalarOp):
+    """
+    First-order derivative of real-order Modified Bessel function of the second kind Kv(z)
+    """
+
+    nfunc_spec = ("scipy.special.kvp", 2, 1)
+
+    @staticmethod
+    def st_impl(v, x):
+        return scipy.special.kvp(v, x)
+
+    def impl(self, v, x):
+        return self.st_impl(v, x)
+
+    def c_code(self, *args, **kwargs):
+        raise NotImplementedError()
+
+
+scalar_kvp = Kvp(upgrade_to_float, name="kvp")
+
+
+class MultivariateLaplaceRV(SymbolicRandomVariable):
+    name = "multivariate_laplace"
+    extended_signature = "[rng],[size],(m),(m,m)->[rng],(m)"
+    _print_name = ("MultivariateLaplace", "\\operatorname{MultivariateLaplace}")
+
+    @classmethod
+    def rv_op(cls, mu, cov, *, size=None, rng=None):
+        mu = pt.as_tensor(mu)
+        cov = pt.as_tensor(cov)
+        rng = normalize_rng_param(rng)
+        size = normalize_size_param(size)
+
+        assert mu.type.ndim >= 1
+        assert cov.type.ndim >= 2
+
+        if rv_size_is_none(size):
+            size = implicit_size_from_params(mu, cov, ndims_params=(1, 2))
+
+        next_rng, e = pt.random.exponential(size=size, rng=rng).owner.outputs
+        next_rng, z = pt.random.multivariate_normal(
+            mean=pt.zeros(mu.shape[-1]), cov=cov, size=size, rng=next_rng
+        ).owner.outputs
+        rv = mu + pt.sqrt(e)[..., None] * z
+
+        return cls(
+            inputs=[rng, size, mu, cov],
+            outputs=[next_rng, rv],
+        )(rng, size, mu, cov)
+
+
+class MvLaplace(Continuous):
+    r"""Multivariate (Symmetric) Laplace distribution."""
+
+    rv_type = MultivariateLaplaceRV
+    rv_op = MultivariateLaplaceRV.rv_op
+
+    @classmethod
+    def dist(cls, mu=0, cov=None, *, tau=None, chol=None, lower=True, **kwargs):
+        cov = quaddist_matrix(cov, chol, tau, lower)
+
+        mu = pt.as_tensor_variable(mu)
+        if mu.type.broadcastable[-1] != cov.type.broadcastable[-1]:
+            mu, _ = pt.broadcast_arrays(mu, cov[..., -1])
+        return super().dist([mu, cov], **kwargs)
+
+    def support_point(rv, size, mu, cov):
+        if rv_size_is_none(size):
+            broadcasted_mu, _ = pt.random.utils.broadcast_params([mu, cov], ndims_params=[1, 2])
+        else:
+            broadcast_shape = pt.concatenate([size, [mu.shape[-1]]])
+            broadcasted_mu = pt.broadcast_to(mu, broadcast_shape)
+        return broadcasted_mu
+
+    def logp(value, mu, cov):
+        quaddist, logdet, posdef = quaddist_chol(value, mu, cov)
+
+        k = value.shape[-1].astype("floatX")
+        norm = np.log(2) - 0.5 * k * np.log(2 * np.pi) - logdet
+
+        v = 1 - (k / 2)
+        kernel = ((v / 2) * pt.log(quaddist / 2)) + pt.log(kv(v, pt.sqrt(2 * quaddist)))
+
+        logp_val = norm + kernel
+        return check_parameters(logp_val, posdef, msg="posdef scale")

tests/distributions/multivariate/test_laplace.py

@@ -0,0 +1,72 @@
+import numpy as np
+import pymc as pm
+import scipy
+
+from pymc_experimental.distributions.multivariate.laplace import MvLaplace
+
+
+class TestMvLaplace:
+    def test_mvlaplace_support_point(self):
+        raise NotImplementedError()
+
+    def test_mvlaplace_mean(self):
+        raise NotImplementedError()
+
+    def test_mvlaplace_random(self):
+        mu = [-1, np.pi, 1]
+        cov = [[1, 0.5, 0.25], [0.5, 2, 0.5], [0.25, 0.5, 3]]
+        rv = MvLaplace.dist(mu=mu, cov=cov, size=10_000)
+
+        samples = pm.draw(rv, random_seed=13)
+        assert samples.shape == (10_000, 3)
+        np.testing.assert_allclose(np.mean(samples, axis=0), mu, rtol=0.05)
+        np.testing.assert_allclose(np.cov(samples, rowvar=False), cov, rtol=0.1)
+
+    def test_laplace_logp(self):
+        # Testing against special bivariate cases described in:
+        # https://en.wikipedia.org/wiki/Multivariate_Laplace_distribution#Probability_density_function
+
+        # Zero mean, non-identity covariance case
+        mu = np.zeros(2)
+        s1 = 0.5
+        s2 = 2.0
+        r = -0.25
+        cov = np.array(
+            [
+                [s1**2, r * s1 * s2],
+                [r * s1 * s2, s2**2],
+            ]
+        )
+        rv = MvLaplace.dist(mu=mu, cov=cov)
+        rv_val = np.random.normal(size=(2,))
+        logp_eval = pm.logp(rv, rv_val).eval()
+
+        x1, x2 = rv_val
+        logp_expected = np.log(
+            (1 / (np.pi * s1 * s2 * np.sqrt(1 - r**2)))
+            * scipy.special.kv(
+                0,
+                np.sqrt(
+                    (2 * ((x1**2 / s1**2) - (2 * r * x1 * x2 / (s1 * s2)) + (x2**2 / s2**2)))
+                    / (1 - r**2)
+                ),
+            )
+        )
+        np.testing.assert_allclose(
+            logp_eval,
+            logp_expected,
+        )
+
+        # Non zero mean, identity covariance case
+        mu = np.array([1, 3])
+        rv = MvLaplace.dist(mu=mu, cov=np.eye(2))
+        rv_val = np.random.normal(size=(2,))
+        logp_eval = pm.logp(rv, rv_val).eval()
+
+        logp_expected = np.log(
+            1 / np.pi * scipy.special.kv(0, np.sqrt(2 * np.sum((rv_val - mu) ** 2)))
+        )
+        np.testing.assert_allclose(
+            logp_eval,
+            logp_expected,
+        )