Porting ToroidalVMCosineDistribution code

pyrecest/distributions/__init__.py

@@ -182,8 +182,8 @@
 from .hypertorus.toroidal_dirac_distribution import ToroidalDiracDistribution
 from .hypertorus.toroidal_mixture import ToroidalMixture
 from .hypertorus.toroidal_uniform_distribution import ToroidalUniformDistribution
-from .hypertorus.toroidal_vm_rivest_distribution import (
-    ToroidalVMRivestDistribution,
+from .hypertorus.toroidal_von_mises_cosine_distribution import (
+    ToroidalVonMisesCosineDistribution,
 )
 from .hypertorus.toroidal_von_mises_sine_distribution import (
     ToroidalVonMisesSineDistribution,
@@ -321,6 +321,7 @@
     "ToroidalDiracDistribution",
     "ToroidalMixture",
     "ToroidalUniformDistribution",
+    "ToroidalVonMisesCosineDistribution",
     "ToroidalVMRivestDistribution",
     "ToroidalVonMisesSineDistribution",
     "ToroidalWrappedNormalDistribution",

pyrecest/distributions/hypertorus/toroidal_von_mises_cosine_distribution.py

@@ -0,0 +1,96 @@
+# pylint: disable=redefined-builtin,no-name-in-module,no-member
+import numpy as np
+from pyrecest.backend import all, array, cos, exp, mod, pi, sum
+from scipy.special import iv
+
+from .abstract_toroidal_distribution import AbstractToroidalDistribution
+
+_SERIES_TERMS = 10
+
+
+class ToroidalVonMisesCosineDistribution(AbstractToroidalDistribution):
+    """Bivariate von Mises distribution, cosine model.
+
+    Corresponds to A = [-kappa3, 0; 0, -kappa3].
+
+    References:
+        Mardia, K. V.; Taylor, C. C. & Subramaniam, G. K.
+        Protein Bioinformatics and Mixtures of Bivariate von Mises Distributions
+        for Angular Data Biometrics, 2007, 63, 505-512
+
+        Mardia, K. V. & Frellsen, J. in Hamelryck, T.; Mardia, K. &
+        Ferkinghoff-Borg, J. (Eds.)
+        Statistics of Bivariate von Mises Distributions
+        Bayesian Methods in Structural Bioinformatics,
+        Springer Berlin Heidelberg, 2012, 159-178
+    """
+
+    def __init__(self, mu, kappa, kappa3):
+        AbstractToroidalDistribution.__init__(self)
+        assert mu.shape == (2,)
+        assert kappa.shape == (2,)
+        assert kappa3.shape == ()
+        assert all(kappa >= 0.0)
+
+        self.mu = mod(mu, 2.0 * pi)
+        self.kappa = kappa
+        self.kappa3 = kappa3
+
+        self.C = 1.0 / self.norm_const
+
+    @property
+    def norm_const(self):
+        def s(p):
+            return iv(p, self.kappa[0]) * iv(p, self.kappa[1]) * iv(p, -self.kappa3)
+
+        Cinv = 4.0 * pi**2 * (
+            s(0) + 2.0 * sum(array([s(p) for p in range(1, _SERIES_TERMS + 1)]))
+        )
+        return Cinv
+
+    def pdf(self, xs):
+        assert xs.shape[-1] == 2
+        p = self.C * exp(
+            self.kappa[0] * cos(xs[..., 0] - self.mu[0])
+            + self.kappa[1] * cos(xs[..., 1] - self.mu[1])
+            - self.kappa3 * cos(xs[..., 0] - self.mu[0] - xs[..., 1] + self.mu[1])
+        )
+        return p
+
+    def trigonometric_moment(self, n):
+        if n == 1:
+            def s1(m):
+                return (
+                    (iv(m + 1, self.kappa[0]) + iv(m - 1, self.kappa[0]))
+                    * iv(m, self.kappa[1])
+                    * iv(m, -self.kappa3)
+                )
+
+            def s2(m):
+                return (
+                    iv(m, self.kappa[0])
+                    * (iv(m + 1, self.kappa[1]) + iv(m - 1, self.kappa[1]))
+                    * iv(m, -self.kappa3)
+                )
+
+            def s(p):
+                return iv(p, self.kappa[0]) * iv(p, self.kappa[1]) * iv(p, -self.kappa3)
+
+            terms = range(1, _SERIES_TERMS + 1)
+            s1_sum = s1(0) / 2.0 + sum(array([s1(m) for m in terms]))
+            s2_sum = s2(0) / 2.0 + sum(array([s2(m) for m in terms]))
+            s_sum = s(0) + 2.0 * sum(array([s(p) for p in terms]))
+
+            # Use numpy directly here because the result is inherently complex
+            # and pyrecest.backend does not support complex-valued arrays.
+            m1 = float(s1_sum) / float(s_sum) * np.exp(1j * n * float(self.mu[0]))
+            m2 = float(s2_sum) / float(s_sum) * np.exp(1j * n * float(self.mu[1]))
+            return np.array([m1, m2])
+        return self.trigonometric_moment_numerical(n)
+
+    def shift(self, shift_by):
+        assert shift_by.shape == (self.dim,)
+        tvm = ToroidalVonMisesCosineDistribution(
+            mod(self.mu + shift_by, 2.0 * pi), self.kappa, self.kappa3
+        )
+        return tvm

pyrecest/tests/distributions/test_toroidal_von_mises_cosine_distribution.py

@@ -0,0 +1,72 @@
+import unittest
+
+import matplotlib
+import numpy.testing as npt
+
+# pylint: disable=no-name-in-module,no-member
+import pyrecest.backend
+
+# pylint: disable=no-name-in-module,no-member
+from pyrecest.backend import arange, array, column_stack, cos, exp, pi
+from pyrecest.distributions.hypertorus.toroidal_von_mises_cosine_distribution import (
+    ToroidalVonMisesCosineDistribution,
+)
+from pyrecest.tests.distributions.test_toroidal_von_mises_sine_distribution import (
+    ToroidalBivarVMTestMixin,
+)
+
+matplotlib.pyplot.close("all")
+matplotlib.use("Agg")
+
+
+class ToroidalVMCosineDistributionTest(ToroidalBivarVMTestMixin, unittest.TestCase):
+    def setUp(self):
+        self.mu = array([1.0, 2.0])
+        self.kappa = array([0.7, 1.4])
+        self.kappa3 = array(0.5)
+        self.tvm = ToroidalVonMisesCosineDistribution(
+            self.mu, self.kappa, self.kappa3
+        )
+
+    def test_instance(self):
+        self.assertIsInstance(self.tvm, ToroidalVonMisesCosineDistribution)
+
+    def test_mu_kappa_kappa3(self):
+        npt.assert_allclose(self.tvm.mu, self.mu)
+        npt.assert_allclose(self.tvm.kappa, self.kappa)
+        self.assertEqual(self.tvm.kappa3, self.kappa3)
+
+    def _unnormalized_pdf(self, xs):
+        return exp(
+            self.kappa[0] * cos(xs[..., 0] - self.mu[0])
+            + self.kappa[1] * cos(xs[..., 1] - self.mu[1])
+            - self.kappa3
+            * cos(xs[..., 0] - self.mu[0] - xs[..., 1] + self.mu[1])
+        )
+
+    @unittest.skipIf(
+        pyrecest.backend.__backend_name__ == "jax",
+        reason="Not supported on this backend",
+    )
+    def test_trigonometric_moment_analytical(self):
+        m_analytical = self.tvm.trigonometric_moment(1)
+        m_numerical = self.tvm.trigonometric_moment_numerical(1)
+        npt.assert_allclose(m_analytical, m_numerical, rtol=1e-8)
+
+    def test_shift(self):
+        shift_by = array([4.0, 2.0])
+        tvm2 = self.tvm.shift(shift_by)
+        self.assertIsInstance(tvm2, ToroidalVonMisesCosineDistribution)
+        x_test = column_stack(
+            (arange(0.0, 2.0 * pi, 0.3), arange(0.0, 2.0 * pi, 0.3))
+        )
+        npt.assert_allclose(
+            tvm2.pdf(x_test),
+            self.tvm.pdf(x_test - shift_by),
+            atol=1e-10,
+            rtol=1e-6,
+        )
+
+
+if __name__ == "__main__":
+    unittest.main()

pyrecest/tests/distributions/test_toroidal_von_mises_sine_distribution.py

@@ -17,28 +17,18 @@
 matplotlib.use("Agg")
 
 
-class ToroidalVMSineDistributionTest(unittest.TestCase):
-    def setUp(self):
-        self.mu = array([1.0, 2.0])
-        self.kappa = array([0.7, 1.4])
-        self.lambda_ = array(0.5)
-        self.tvm = ToroidalVonMisesSineDistribution(self.mu, self.kappa, self.lambda_)
+class ToroidalBivarVMTestMixin:
+    """Shared tests for bivariate toroidal von Mises distribution test classes.
 
-    def test_instance(self):
-        # sanity check
-        self.assertIsInstance(self.tvm, ToroidalVonMisesSineDistribution)
-
-    def test_mu_kappa_lambda(self):
-        npt.assert_allclose(self.tvm.mu, self.mu)
-        npt.assert_allclose(self.tvm.kappa, self.kappa)
-        self.assertEqual(self.tvm.lambda_, self.lambda_)
+    Subclasses must implement ``setUp`` (setting ``self.mu``, ``self.kappa``,
+    and ``self.tvm``) and ``_unnormalized_pdf``.
+    """
 
     @unittest.skipIf(
         pyrecest.backend.__backend_name__ in ("pytorch", "jax"),
         reason="Not supported on this backend",
     )
     def test_integral(self):
-        # test integral
         self.assertAlmostEqual(self.tvm.integrate(), 1.0, delta=1e-5)
 
     @unittest.skipIf(
@@ -53,17 +43,6 @@ def test_trigonometric_moment_numerical(self):
     def test_plot_2d(self):
         self.tvm.plot()
 
-    # jscpd:ignore-start
-    # pylint: disable=R0801
-    def _unnormalized_pdf(self, xs):
-        return exp(
-            self.kappa[0] * cos(xs[..., 0] - self.mu[0])
-            + self.kappa[1] * cos(xs[..., 1] - self.mu[1])
-            + self.lambda_ * sin(xs[..., 0] - self.mu[0]) * sin(xs[..., 1] - self.mu[1])
-        )
-
-    # jscpd:ignore-end
-
     @parameterized.expand(
         [
             (array([3.0, 2.0]),),
@@ -89,5 +68,29 @@ def pdf(x):
         npt.assert_allclose(self.tvm.pdf(x), expected)
 
 
+class ToroidalVMSineDistributionTest(ToroidalBivarVMTestMixin, unittest.TestCase):
+    def setUp(self):
+        self.mu = array([1.0, 2.0])
+        self.kappa = array([0.7, 1.4])
+        self.lambda_ = array(0.5)
+        self.tvm = ToroidalVonMisesSineDistribution(self.mu, self.kappa, self.lambda_)
+
+    def test_instance(self):
+        # sanity check
+        self.assertIsInstance(self.tvm, ToroidalVonMisesSineDistribution)
+
+    def test_mu_kappa_lambda(self):
+        npt.assert_allclose(self.tvm.mu, self.mu)
+        npt.assert_allclose(self.tvm.kappa, self.kappa)
+        self.assertEqual(self.tvm.lambda_, self.lambda_)
+
+    def _unnormalized_pdf(self, xs):
+        return exp(
+            self.kappa[0] * cos(xs[..., 0] - self.mu[0])
+            + self.kappa[1] * cos(xs[..., 1] - self.mu[1])
+            + self.lambda_ * sin(xs[..., 0] - self.mu[0]) * sin(xs[..., 1] - self.mu[1])
+        )
+
+
 if __name__ == "__main__":
     unittest.main()