Add periodic parameter support for 2D plots

- Add convolve2D_periodic() function for 2D periodic convolution
- Support periodic_x, periodic_y, periodic_both convolution modes
- Update get2DDensity() to detect parameter periodicity and select appropriate convolution mode
- Fix edge mask handling to skip boundary corrections on periodic axes
- Add comprehensive test case for periodic 2D density calculations
- Ensure perfect boundary continuity for periodic parameters while preserving edge effects for non-periodic parameters

Author: cmbant

getdist/convolve.py

@@ -203,9 +203,126 @@ def convolve1D(x, y, mode, largest_size=0, cache=None, cache_args=(1, 2)):
 
 
 def convolve2D(x, y, mode, largest_size=0, cache=None, cache_args=(1, 2)):
+    if mode in ("periodic", "periodic_both"):
+        return convolve2D_periodic(x, y, cache, cache_args, periodic_x=True, periodic_y=True)
+    elif mode == "periodic_x":
+        return convolve2D_periodic(x, y, cache, cache_args, periodic_x=True, periodic_y=False)
+    elif mode == "periodic_y":
+        return convolve2D_periodic(x, y, cache, cache_args, periodic_x=False, periodic_y=True)
     return convolveFFTn(x, y, mode, largest_size, cache, cache_args=cache_args)
 
 
+def convolve2D_periodic(x, y, cache=None, cache_args=(1, 2), periodic_x=True, periodic_y=True):
+    """
+    2D convolution with periodic boundary conditions in specified dimensions.
+
+    This implements periodic convolution by extending the 1D approach to 2D.
+    For periodic dimensions, the input is made circular by averaging boundary values.
+
+    :param x: 2D input array (data)
+    :param y: 2D kernel array
+    :param cache: optional cache for FFTs
+    :param cache_args: which arrays to cache (1=x, 2=y)
+    :param periodic_x: whether to use periodic boundaries in x (second) dimension
+    :param periodic_y: whether to use periodic boundaries in y (first) dimension
+    :return: convolved array of same size as x
+    """
+    if x.ndim != 2 or y.ndim != 2:
+        raise ValueError("convolve2D_periodic requires 2D arrays")
+
+    ny, nx = x.shape
+    ky, kx = y.shape
+
+    # Create the circular version of x based on periodicity
+    if periodic_x and periodic_y:
+        # Both dimensions periodic
+        x_circ = x[:-1, :-1].copy()  # Remove last row and column
+        x_circ[0, :] += x[-1, :-1]  # Add last row to first row
+        x_circ[:, 0] += x[:-1, -1]  # Add last column to first column
+        x_circ[0, 0] += x[-1, -1]  # Add corner element
+        N_y, N_x = x_circ.shape
+
+    elif periodic_x and not periodic_y:
+        # Only x dimension periodic
+        x_circ = x[:, :-1].copy()  # Remove last column
+        x_circ[:, 0] += x[:, -1]  # Add last column to first column
+        N_y, N_x = x_circ.shape
+
+    elif periodic_y and not periodic_x:
+        # Only y dimension periodic
+        x_circ = x[:-1, :].copy()  # Remove last row
+        x_circ[0, :] += x[-1, :]  # Add last row to first row
+        N_y, N_x = x_circ.shape
+
+    else:
+        # Neither periodic (shouldn't happen, but handle gracefully)
+        return convolveFFTn(x, y, "same")
+
+    # Prepare kernel for circular convolution
+    hpad = np.zeros((N_y, N_x), dtype=float)
+    hpad[:ky, :kx] = y
+    # Roll kernel to center it properly for convolution
+    hpad = np.roll(hpad, -(ky // 2), axis=0)
+    hpad = np.roll(hpad, -(kx // 2), axis=1)
+
+    # Cache keys
+    if cache is not None:
+        if 1 in cache_args:
+            key_x = ("2D_circ", N_y, N_x, periodic_x, periodic_y, id(x))
+        if 2 in cache_args:
+            key_y = ("2D_kernel", N_y, N_x, id(y))
+
+    # Get or compute FFT of kernel
+    yfft = None
+    if cache is not None and 2 in cache_args:
+        yfft = cache.get(key_y)
+    if yfft is None:
+        yfft = np.fft.rfftn(hpad)
+        if cache is not None and 2 in cache_args:
+            cache[key_y] = yfft
+
+    # Get or compute FFT of data
+    xfft = None
+    if cache is not None and 1 in cache_args:
+        xfft = cache.get(key_x)
+    if xfft is None:
+        xfft = np.fft.rfftn(x_circ)
+        if cache is not None and 1 in cache_args:
+            cache[key_x] = xfft
+
+    # Perform convolution in frequency domain
+    result_fft = xfft * yfft
+    result = np.fft.irfftn(result_fft, (N_y, N_x), axes=(0, 1))
+
+    # Extend result back to original size for periodic dimensions
+    if periodic_x and periodic_y:
+        # Both periodic: append first row/column as last
+        final_result = np.empty((ny, nx))
+        final_result[:-1, :-1] = result
+        final_result[-1, :-1] = result[0, :]  # Last row = first row
+        final_result[:-1, -1] = result[:, 0]  # Last column = first column
+        final_result[-1, -1] = result[0, 0]  # Corner element
+        return final_result
+
+    elif periodic_x and not periodic_y:
+        # Only x periodic: append first column as last
+        final_result = np.empty((ny, nx))
+        final_result[:, :-1] = result
+        final_result[:, -1] = result[:, 0]  # Last column = first column
+        return final_result
+
+    elif periodic_y and not periodic_x:
+        # Only y periodic: append first row as last
+        final_result = np.empty((ny, nx))
+        final_result[:-1, :] = result
+        final_result[-1, :] = result[0, :]  # Last row = first row
+        return final_result
+
+    else:
+        # Neither periodic (fallback)
+        return result
+
+
 def convolve1D_periodic(x, y, cache=None, cache_args=(1, 2)):
     """
     Circular (periodic) 1D convolution of x with kernel y. Returns same-length result.

getdist/mcsamples.py

@@ -1700,25 +1700,31 @@ def get1DDensityGridData(self, j, paramConfid=None, meanlikes=False, **kwargs):
         return density1D
 
     def _setEdgeMask2D(self, parx, pary, prior_mask, winw):
-        if parx.has_limits_bot:
-            prior_mask[:, winw] /= 2
+        # Only apply edge masks on non-periodic axes; periodic axes have no boundaries
+        if not parx.periodic:
+            if parx.has_limits_bot:
+                prior_mask[:, winw] /= 2
+                prior_mask[:, :winw] = 0
+            if parx.has_limits_top:
+                prior_mask[:, -(winw + 1)] /= 2
+                prior_mask[:, -winw:] = 0
+        if not pary.periodic:
+            if pary.has_limits_bot:
+                prior_mask[winw, :] /= 2
+                prior_mask[:winw:] = 0
+            if pary.has_limits_top:
+                prior_mask[-(winw + 1), :] /= 2
+                prior_mask[-winw:, :] = 0
+
+    def _setAllEdgeMask2D(self, prior_mask, winw, periodic_x=False, periodic_y=False):
+        # Zero out margins only along non-periodic axes
+        if not periodic_x:
             prior_mask[:, :winw] = 0
-        if parx.has_limits_top:
-            prior_mask[:, -(winw + 1)] /= 2
             prior_mask[:, -winw:] = 0
-        if pary.has_limits_bot:
-            prior_mask[winw, :] /= 2
+        if not periodic_y:
             prior_mask[:winw:] = 0
-        if pary.has_limits_top:
-            prior_mask[-(winw + 1), :] /= 2
             prior_mask[-winw:, :] = 0
 
-    def _setAllEdgeMask2D(self, prior_mask, winw):
-        prior_mask[:, :winw] = 0
-        prior_mask[:, -winw:] = 0
-        prior_mask[:winw:] = 0
-        prior_mask[-winw:, :] = 0
-
     def _getScaleForParam(self, par):
         # Also ensures that the 1D limits are initialized
         density = self.get1DDensity(par)
@@ -1878,14 +1884,29 @@ def get2DDensityGridData(
         start = time.time()
         cache = {}
         convolvesize = xsize + 2 * winw + Win.shape[0]  # larger than needed for selecting fft pixel count
-        bins2D = convolve2D(histbins, Win, "same", largest_size=convolvesize, cache=cache)
+
+        # Determine convolution mode based on parameter periodicity
+        if parx.periodic and pary.periodic:
+            convolution_mode = "periodic_both"
+        elif parx.periodic:
+            convolution_mode = "periodic_x"
+        elif pary.periodic:
+            convolution_mode = "periodic_y"
+        else:
+            convolution_mode = "same"
+
+        bins2D = convolve2D(histbins, Win, convolution_mode, largest_size=convolvesize, cache=cache)
 
         if meanlikes:
-            bin2Dlikes = convolve2D(finebinlikes, Win, "same", largest_size=convolvesize, cache=cache, cache_args=[2])
+            bin2Dlikes = convolve2D(
+                finebinlikes, Win, convolution_mode, largest_size=convolvesize, cache=cache, cache_args=[2]
+            )
             if mult_bias_correction_order:
                 ix = bin2Dlikes > 0
                 finebinlikes[ix] /= bin2Dlikes[ix]
-                likes2 = convolve2D(finebinlikes, Win, "same", largest_size=convolvesize, cache=cache, cache_args=[2])
+                likes2 = convolve2D(
+                    finebinlikes, Win, convolution_mode, largest_size=convolvesize, cache=cache, cache_args=[2]
+                )
                 likes2[ix] *= bin2Dlikes[ix]
                 bin2Dlikes = likes2
             del finebinlikes
@@ -1896,17 +1917,26 @@ def get2DDensityGridData(
             bin2Dlikes = None
 
         bool_mask = None
+
         if has_prior and boundary_correction_order >= 0 or mult_bias_correction_order or mask_function:
+            # Always pad by winw in both axes so that 'valid' convolution returns (ysize, xsize)
+            # Masks are only applied on non-periodic axes; periodic axes remain as ones.
             prior_mask = np.ones((ysize + 2 * winw, xsize + 2 * winw))
             if mask_function:
                 mask_function(
-                    xbinmin - winw * finewidthx, ybinmin - winw * finewidthy, finewidthx, finewidthy, prior_mask
+                    xbinmin - winw * finewidthx,
+                    ybinmin - winw * finewidthy,
+                    finewidthx,
+                    finewidthy,
+                    prior_mask,
                 )
                 bool_mask = prior_mask[winw:-winw, winw:-winw] < 1e-8
 
-        if has_prior and boundary_correction_order >= 0:
-            # Correct for edge effects
+        if has_prior and boundary_correction_order >= 0 and not (parx.periodic and pary.periodic):
+            # Correct for edge effects; if only one axis is periodic still correct on the other axis
             self._setEdgeMask2D(parx, pary, prior_mask, winw)
+            # Use 'valid' on the padded prior_mask so the result is (ysize, xsize)
+            # For periodic axes prior_mask is ones along that axis, so this is consistent
             a00 = convolve2D(prior_mask, Win, "valid", largest_size=convolvesize, cache=cache)
             ix = a00 * bins2D > np.max(bins2D) * 1e-8
             a00 = a00[ix]
@@ -1933,8 +1963,8 @@ def get2DDensityGridData(
                 a11 = convolve2D(
                     prior_mask, winy * indexes, "valid", largest_size=convolvesize, cache=cache, cache_args=[1]
                 )[ix]
-                xP = convolve2D(histbins, winx, "same", largest_size=convolvesize, cache=cache)[ix]
-                yP = convolve2D(histbins, winy, "same", largest_size=convolvesize, cache=cache)[ix]
+                xP = convolve2D(histbins, winx, convolution_mode, largest_size=convolvesize, cache=cache)[ix]
+                yP = convolve2D(histbins, winy, convolution_mode, largest_size=convolvesize, cache=cache)[ix]
                 denom = a20 * a01**2 + a10**2 * a02 - a00 * a02 * a20 + a11**2 * a00 - 2 * a01 * a10 * a11
                 A = a11**2 - a02 * a20
                 Ax = a10 * a02 - a01 * a11
@@ -1944,14 +1974,16 @@ def get2DDensityGridData(
             else:
                 raise SettingError("unknown boundary_correction_order (expected 0 or 1)")
 
-        if mult_bias_correction_order:
-            self._setAllEdgeMask2D(prior_mask, winw)
+        if mult_bias_correction_order and not (parx.periodic and pary.periodic):
+            # If only one axis is periodic still correct on the non-periodic edges
+            self._setAllEdgeMask2D(prior_mask, winw, periodic_x=parx.periodic, periodic_y=pary.periodic)
+            # Use 'valid' on the padded prior_mask so the result is (ysize, xsize)
             a00 = convolve2D(prior_mask, Win, "valid", largest_size=convolvesize, cache=cache, cache_args=[2])
             for _ in range(mult_bias_correction_order):
                 box = histbins.copy()
                 ix2 = bins2D > np.max(bins2D) * 1e-8
                 box[ix2] /= bins2D[ix2]
-                bins2D *= convolve2D(box, Win, "same", largest_size=convolvesize, cache=cache, cache_args=[2])
+                bins2D *= convolve2D(box, Win, convolution_mode, largest_size=convolvesize, cache=cache, cache_args=[2])
                 if mask_function:
                     bins2D[~bool_mask] /= a00[~bool_mask]
                 else:

getdist/tests/getdist_test.py

@@ -178,6 +178,43 @@ def testNDDensity(self):
 
         self.assertTrue(np.allclose(d2D.P, dND.P, atol=1e-5))
 
+    def testPeriodic2D(self):
+        # Test 2D density calculation with periodic parameters
+        n_samples = 1000
+        np.random.seed(42)
+
+        # Create angle parameter (periodic) and radius parameter (non-periodic)
+        angle = np.random.normal(0, 1, n_samples) % (2 * np.pi)
+        radius = np.abs(np.random.normal(2, 0.5, n_samples))
+
+        samples = np.column_stack([angle, radius])
+        names = ["angle", "radius"]
+        labels = [r"\theta", "r"]
+
+        mcsamples = MCSamples(
+            samples=samples, names=names, labels=labels, ranges={"angle": [0, 2 * np.pi, "periodic"], "radius": [0, 5]}
+        )
+
+        # Test 2D density calculation
+        density = mcsamples.get2DDensity("angle", "radius", fine_bins_2D=32)
+
+        # Basic checks
+        self.assertEqual(density.P.shape, (32, 32))
+        self.assertGreater(np.max(density.P), 0)
+        self.assertGreater(density.norm_integral(), 0)
+
+        # Periodic edges equal checks (merge from former test)
+        d64x = mcsamples.get2DDensity("angle", "radius", fine_bins_2D=64)
+        self.assertTrue(np.allclose(d64x.P[:, 0], d64x.P[:, -1], atol=5e-3, rtol=5e-3))
+        d64y = mcsamples.get2DDensity("radius", "angle", fine_bins_2D=64)
+        self.assertTrue(np.allclose(d64y.P[0, :], d64y.P[-1, :], atol=5e-3, rtol=5e-3))
+
+        # Test that periodic parameter is properly detected
+        angle_param = mcsamples.paramNames.parWithName("angle")
+        radius_param = mcsamples.paramNames.parWithName("radius")
+        self.assertTrue(angle_param.periodic)
+        self.assertFalse(radius_param.periodic)
+
     def testLoads(self):
         # test initiating from multiple chain arrays
         samps = []