Temporary rebase staging for #4290

src/pyrecest/distributions/hypertorus/hypertoroidal_wrapped_normal_distribution.py

@@ -143,20 +143,22 @@ def pdf(self, xs, m: Union[int, int32, int64] = 3):
         """
         m = _validate_series_order(m)
         xs = as_hypertoroidal_points(xs, self.dim)
-        xs = (xs + pi) % (2 * pi) - pi
+        residuals = mod(xs - self.mu + pi, 2.0 * pi) - pi
 
         # Generate all combinations of offsets for each dimension
         offsets = [arange(-m, m + 1) * 2.0 * pi for _ in range(self.dim)]
         offset_combinations = stack(meshgrid(*offsets, indexing="ij"), -1).reshape(
             -1, self.dim
         )
 
-        # Calculate the PDF values by considering all combinations of offsets
+        # Center the finite image expansion on the nearest wrapped residual.
         pdf_values = zeros(xs.shape[0])
         for offset in offset_combinations:
-            shifted_xa = xs + offset[None, :]
+            shifted_residuals = residuals + offset[None, :]
             pdf_values += multivariate_normal.pdf(
-                shifted_xa, mean=self.mu.flatten(), cov=self.C
+                shifted_residuals,
+                mean=zeros(self.dim),
+                cov=self.C,
             )
 
         return pdf_values

tests/distributions/test_hypertoroidal_wrapped_normal_truncation.py

@@ -0,0 +1,29 @@
+import numpy as np
+import numpy.testing as npt
+
+from pyrecest.backend import array, to_numpy
+from pyrecest.distributions.hypertorus.hypertoroidal_wrapped_normal_distribution import (
+    HypertoroidalWrappedNormalDistribution,
+)
+
+
+def test_zero_order_pdf_centers_truncation_on_wrapped_residual():
+    sigma = 0.2
+    mean = 5.5
+    distribution = HypertoroidalWrappedNormalDistribution(
+        array([mean]),
+        array([[sigma**2]]),
+    )
+
+    values = distribution.pdf(
+        array([[mean], [mean - 2.0 * np.pi]]),
+        m=0,
+    )
+
+    expected_peak = 1.0 / (np.sqrt(2.0 * np.pi) * sigma)
+    npt.assert_allclose(
+        to_numpy(values),
+        [expected_peak, expected_peak],
+        rtol=1.0e-6,
+        atol=1.0e-6,
+    )