From ec4a69791ccc74f7c35eb75b6178a6c3f8879a14 Mon Sep 17 00:00:00 2001
From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com>
Date: Tue, 31 Mar 2026 12:33:06 +0000
Subject: [PATCH 01/19] Implement TdCondTdGridDistribution with tests

Agent-Logs-Url: https://github.com/FlorianPfaff/PyRecEst/sessions/7d751a84-d968-4c42-864e-1b8382d38b4d

Co-authored-by: FlorianPfaff <6773539+FlorianPfaff@users.noreply.github.com>
---
 .../td_cond_td_grid_distribution.py           | 299 ++++++++++++++++++
 .../test_td_cond_td_grid_distribution.py      | 227 +++++++++++++
 2 files changed, 526 insertions(+)
 create mode 100644 pyrecest/distributions/hypertorus/td_cond_td_grid_distribution.py
 create mode 100644 pyrecest/tests/distributions/test_td_cond_td_grid_distribution.py

diff --git a/pyrecest/distributions/hypertorus/td_cond_td_grid_distribution.py b/pyrecest/distributions/hypertorus/td_cond_td_grid_distribution.py
new file mode 100644
index 000000000..df18c439e
--- /dev/null
+++ b/pyrecest/distributions/hypertorus/td_cond_td_grid_distribution.py
@@ -0,0 +1,299 @@
+"""Conditional grid distribution on the hypertorus (TdCondTdGridDistribution)."""
+
+from __future__ import annotations
+
+import warnings
+
+import numpy as np
+from beartype import beartype
+
+# pylint: disable=redefined-builtin,no-name-in-module,no-member
+from pyrecest.backend import abs, all, any, array, mean, meshgrid, pi, reshape
+
+from ..conditional.abstract_conditional_distribution import (
+    AbstractConditionalDistribution,
+)
+from .hypertoroidal_grid_distribution import HypertoroidalGridDistribution
+
+
+class TdCondTdGridDistribution(AbstractConditionalDistribution):
+    """Conditional distribution on the hypertorus described by grid values.
+
+    Represents ``f(a | b)`` where both ``a`` and ``b`` live on the same
+    ``d``-dimensional hypertorus.  The combined state space is the Cartesian
+    product of two hypertori, and :attr:`grid_values` is a square 2-D matrix
+    where::
+
+        grid_values[i, j]  =  f(grid[i] | grid[j])
+
+    i.e. the *rows* index the conditioned variable ``a`` and the *columns*
+    index the conditioning variable ``b``.
+
+    Parameters
+    ----------
+    grid:
+        Grid points on the individual hypertorus, shape
+        ``(n_grid_points, dim_half)``.
+    grid_values:
+        Conditional probability values, shape
+        ``(n_grid_points, n_grid_points)``.  Must be square; entry ``[i, j]``
+        is ``f(grid[i] | grid[j])``.
+    enforce_pdf_nonnegative:
+        Whether to enforce non-negativity of the pdf values (default ``True``).
+    """
+
+    @beartype
+    def __init__(
+        self,
+        grid,
+        grid_values,
+        enforce_pdf_nonnegative: bool = True,
+    ):
+        super().__init__()
+
+        if grid_values.ndim != 2:
+            raise ValueError("grid_values must be a 2-D array")
+        if grid_values.shape[0] != grid_values.shape[1]:
+            raise ValueError("grid_values must be square")
+        if grid.shape[0] != grid_values.shape[0]:
+            raise ValueError(
+                "grid.shape[0] must equal grid_values.shape[0] (number of grid points)"
+            )
+
+        self.grid = grid  # (n_grid_points, dim_half)
+        self.grid_values = grid_values  # (n_grid_points, n_grid_points)
+        self.enforce_pdf_nonnegative = enforce_pdf_nonnegative
+        # dim is the combined dimension of the two-hypertorus Cartesian product
+        self.dim = 2 * grid.shape[1]
+        # Populated by from_function; used to reshape slices back to the
+        # correct multi-dimensional layout when creating HypertoroidalGridDistribution
+        self._n_grid_points_per_dim: tuple[int, ...] | None = None
+
+        self._check_normalization()
+
+    # ----------------------------------------------------------------- internal
+
+    def _check_normalization(self, tol: float = 0.01) -> None:
+        """Warn if the distribution is not column-normalized.
+
+        For each fixed ``b = grid[j]`` the values ``grid_values[:, j]``
+        should integrate to 1 over the first variable, i.e.
+        ``mean(grid_values[:, j]) * (2π)^(dim/2) ≈ 1``.
+        """
+        dim_half = self.dim // 2
+        manifold_size = (2.0 * float(pi)) ** dim_half
+
+        ints = mean(self.grid_values, axis=0) * manifold_size
+        if any(abs(ints - 1.0) > tol):
+            # Check whether swapping rows/cols would give a normalized distribution
+            ints_swapped = mean(self.grid_values, axis=1) * manifold_size
+            if all(abs(ints_swapped - 1.0) <= tol):
+                raise ValueError(
+                    "Normalization:maybeWrongOrder: The distribution is not "
+                    "normalized, but would be normalized if the order of the "
+                    "tori were swapped.  Check your input."
+                )
+            warnings.warn(
+                "Normalization:notNormalized: When conditioning values for the "
+                "first torus on the second, normalization is not guaranteed.  "
+                "Check your input or increase the tolerance.",
+                RuntimeWarning,
+            )
+
+    # ------------------------------------------------------------------ public
+
+    def normalize(self):
+        """Return *self* unchanged (normalization check only; no rescaling).
+
+        For compatibility with the rest of the API.  Normalization must be
+        ensured by the caller when constructing :class:`TdCondTdGridDistribution`.
+        """
+        self._check_normalization()
+        return self
+
+    @beartype
+    def marginalize_out(self, first_or_second: int):
+        """Marginalize out one of the two variables.
+
+        Parameters
+        ----------
+        first_or_second:
+            ``1`` to sum over the first (conditioned) variable ``a``;
+            ``2`` to sum over the second (conditioning) variable ``b``.
+
+        Returns
+        -------
+        HypertoroidalGridDistribution
+            Marginal distribution on the remaining variable.
+        """
+        if first_or_second == 1:
+            # Sum rows → proportional to the marginal over b
+            grid_values_sgd = np.sum(np.asarray(self.grid_values), axis=0)
+        elif first_or_second == 2:
+            # Sum cols → proportional to the marginal over a
+            grid_values_sgd = np.sum(np.asarray(self.grid_values), axis=1)
+        else:
+            raise ValueError("first_or_second must be 1 or 2")
+
+        return self._make_hypertoroidal_grid_dist(array(grid_values_sgd))
+
+    @beartype
+    def fix_dim(self, first_or_second: int, point):
+        """Fix one variable to a grid point and return the slice distribution.
+
+        Parameters
+        ----------
+        first_or_second:
+            ``1`` to fix the first variable ``a`` to *point*;
+            ``2`` to fix the second variable ``b`` to *point*.
+        point:
+            The value to fix; must be an exact grid point,
+            shape ``(dim_half,)`` or ``(dim_half, 1)``.
+
+        Returns
+        -------
+        HypertoroidalGridDistribution
+            Conditional (or likelihood) distribution on the remaining variable.
+
+        Raises
+        ------
+        ValueError
+            If *point* is not found in the grid.
+        """
+        dim_half = self.dim // 2
+        point_arr = np.asarray(point).ravel()
+        if point_arr.shape != (dim_half,):
+            raise ValueError(
+                f"point must have shape ({dim_half},), got {point_arr.shape}"
+            )
+
+        grid_np = np.asarray(self.grid)
+        diffs = np.all(np.isclose(grid_np, point_arr[None, :]), axis=1)
+        indices = np.where(diffs)[0]
+        if indices.size == 0:
+            raise ValueError(
+                "Cannot fix value at this point because it is not on the grid"
+            )
+        locb = int(indices[0])
+
+        gv = np.asarray(self.grid_values)
+        if first_or_second == 1:
+            grid_values_slice = gv[locb, :]
+        elif first_or_second == 2:
+            grid_values_slice = gv[:, locb]
+        else:
+            raise ValueError("first_or_second must be 1 or 2")
+
+        return self._make_hypertoroidal_grid_dist(array(grid_values_slice))
+
+    # -------------------------------------------------------------- factory
+
+    @staticmethod
+    @beartype
+    def from_function(
+        fun,
+        n_grid_points: int | tuple | list,
+        fun_does_cartesian_product: bool,
+        grid_type: str = "CartesianProd",
+        dim: int = 2,
+    ) -> "TdCondTdGridDistribution":
+        """Construct a :class:`TdCondTdGridDistribution` from a function handle.
+
+        Parameters
+        ----------
+        fun:
+            ``f(a, b)`` where ``a`` and ``b`` each have shape
+            ``(n_eval, dim_half)``.  When *fun_does_cartesian_product* is
+            ``False``, ``n_eval = n_total ** 2`` (all pairs); the function
+            must return shape ``(n_total ** 2,)``.
+            When *fun_does_cartesian_product* is ``True``, ``a`` and ``b``
+            are both the full grid (shape ``(n_total, dim_half)``) and the
+            function must return shape ``(n_total, n_total)``.
+        n_grid_points:
+            Number of grid points per dimension of a *single* hypertorus.
+            Pass an ``int`` to use the same count for every dimension, or a
+            sequence to specify each dimension individually.
+        fun_does_cartesian_product:
+            Whether *fun* handles all grid combinations internally.
+        grid_type:
+            Grid layout; currently only ``'CartesianProd'`` /
+            ``'CartesianProduct'`` is supported.
+        dim:
+            Total dimension of the combined space (must be even; equals
+            ``2 * dim_half``).
+
+        Returns
+        -------
+        TdCondTdGridDistribution
+        """
+        if dim % 2 != 0:
+            raise ValueError(
+                "dim must be even (it represents two copies of a hypertorus)."
+            )
+        # User-facing API mirrors the MATLAB convention ('CartesianProd').
+        # Internally HypertoroidalGridDistribution uses 'cartesian_prod'.
+        if grid_type not in ("CartesianProd", "CartesianProduct"):
+            raise ValueError(
+                "Grid scheme not recognized; only 'CartesianProd' / "
+                "'CartesianProduct' is currently supported."
+            )
+
+        dim_half = dim // 2
+
+        if isinstance(n_grid_points, int):
+            n_grid_points_seq = [n_grid_points] * dim_half
+        else:
+            n_grid_points_seq = list(n_grid_points)
+
+        # Generate the grid on a single hypertorus: (n_total, dim_half)
+        grid = HypertoroidalGridDistribution.generate_cartesian_product_grid(
+            n_grid_points_seq
+        )
+        n_total = grid.shape[0]
+
+        if fun_does_cartesian_product:
+            fvals = fun(grid, grid)
+            fvals = reshape(array(fvals), (n_total, n_total))
+        else:
+            # Build index arrays so that fvals[i, j] = f(grid[i], grid[j]).
+            # Using plain numpy for index generation (pure control flow).
+            r = np.arange(n_total)
+            idx_i_np, idx_j_np = np.meshgrid(r, r, indexing="ij")
+            idx_i = idx_i_np.ravel()
+            idx_j = idx_j_np.ravel()
+
+            fvals_flat = fun(grid[idx_i], grid[idx_j])
+            fvals = reshape(array(fvals_flat), (n_total, n_total))
+
+        td = TdCondTdGridDistribution(grid, fvals)
+        td._n_grid_points_per_dim = tuple(n_grid_points_seq)
+        return td
+
+    # --------------------------------------------------------------- helpers
+
+    def _make_hypertoroidal_grid_dist(
+        self, grid_values_flat
+    ) -> HypertoroidalGridDistribution:
+        """Create a :class:`HypertoroidalGridDistribution` from a 1-D slice.
+
+        When the instance was created via :meth:`from_function` the
+        per-dimension point counts are known and the flat slice is reshaped to
+        the correct multi-dimensional layout required by the ``cartesian_prod``
+        grid type.  Otherwise the ``custom`` grid type is used.
+        """
+        if self._n_grid_points_per_dim is not None:
+            grid_values_shaped = reshape(
+                grid_values_flat, self._n_grid_points_per_dim
+            )
+            return HypertoroidalGridDistribution(
+                grid_values=grid_values_shaped,
+                grid_type="cartesian_prod",
+                grid=self.grid,
+                enforce_pdf_nonnegative=self.enforce_pdf_nonnegative,
+            )
+        return HypertoroidalGridDistribution(
+            grid_values=grid_values_flat,
+            grid_type="custom",
+            grid=self.grid,
+            enforce_pdf_nonnegative=self.enforce_pdf_nonnegative,
+        )
diff --git a/pyrecest/tests/distributions/test_td_cond_td_grid_distribution.py b/pyrecest/tests/distributions/test_td_cond_td_grid_distribution.py
new file mode 100644
index 000000000..4c183ee0b
--- /dev/null
+++ b/pyrecest/tests/distributions/test_td_cond_td_grid_distribution.py
@@ -0,0 +1,227 @@
+import unittest
+import warnings
+
+import numpy as np
+import numpy.testing as npt
+
+from pyrecest.backend import array, pi
+from pyrecest.distributions.hypertorus.hypertoroidal_grid_distribution import (
+    HypertoroidalGridDistribution,
+)
+from pyrecest.distributions.hypertorus.td_cond_td_grid_distribution import (
+    TdCondTdGridDistribution,
+)
+
+
+def _make_normalized_grid_values(n: int) -> np.ndarray:
+    """Return an (n x n) matrix whose columns are normalized (integrate to 1)."""
+    gv = np.random.default_rng(0).uniform(0.5, 1.5, size=(n, n))
+    # Normalize each column so that mean(col) * (2*pi)^1 == 1
+    gv = gv / (gv.mean(axis=0, keepdims=True) * 2.0 * np.pi)
+    return gv
+
+
+class TdCondTdGridDistributionTest(unittest.TestCase):
+    # -------------------------------------------------------------- construction
+
+    def test_construction_t1(self):
+        """Basic construction for T1 x T1."""
+        n = 5
+        grid = np.linspace(0.0, 2.0 * np.pi - 2.0 * np.pi / n, n).reshape(-1, 1)
+        gv = _make_normalized_grid_values(n)
+        td = TdCondTdGridDistribution(array(grid), array(gv))
+        self.assertEqual(td.dim, 2)
+        npt.assert_array_equal(np.asarray(td.grid), grid)
+        npt.assert_array_equal(np.asarray(td.grid_values), gv)
+
+    def test_construction_wrong_shape_raises(self):
+        n = 4
+        grid = np.zeros((n, 1))
+        with self.assertRaises(ValueError):
+            # Non-square grid_values
+            TdCondTdGridDistribution(
+                array(grid), array(np.ones((n, n + 1)) / (n * 2 * np.pi))
+            )
+
+    def test_construction_wrong_order_raises(self):
+        """Transposed (row-normalized) matrix should raise an error."""
+        n = 6
+        grid = np.linspace(0.0, 2.0 * np.pi - 2.0 * np.pi / n, n).reshape(-1, 1)
+        gv = _make_normalized_grid_values(n)
+        # Transpose → rows are normalized, columns are not
+        with self.assertRaises(ValueError):
+            TdCondTdGridDistribution(array(grid), array(gv.T))
+
+    def test_construction_unnormalized_warns(self):
+        """An unnormalized matrix that cannot be fixed by transposing should warn."""
+        n = 5
+        grid = np.linspace(0.0, 2.0 * np.pi - 2.0 * np.pi / n, n).reshape(-1, 1)
+        gv = np.ones((n, n))  # neither rows nor cols sum to 1/(2pi)
+        with self.assertWarns(RuntimeWarning):
+            TdCondTdGridDistribution(array(grid), array(gv))
+
+    # -------------------------------------------------------------- normalize
+
+    def test_normalize_returns_self(self):
+        n = 4
+        grid = np.linspace(0.0, 2.0 * np.pi - 2.0 * np.pi / n, n).reshape(-1, 1)
+        gv = _make_normalized_grid_values(n)
+        td = TdCondTdGridDistribution(array(grid), array(gv))
+        self.assertIs(td.normalize(), td)
+
+    # -------------------------------------------------------------- from_function
+
+    def test_from_function_t1(self):
+        """from_function should recover a wrapped-normal-like conditional."""
+        n = 20
+        dim = 2  # T1 x T1
+
+        def cond_fun(a, b):
+            # Simple Gaussian-like conditional (unnormalized, normalized per column)
+            diff = np.asarray(a)[:, 0] - np.asarray(b)[:, 0]
+            return np.exp(-0.5 * np.minimum(diff**2, (2 * np.pi - np.abs(diff)) ** 2))
+
+        td = TdCondTdGridDistribution.from_function(
+            cond_fun, n, fun_does_cartesian_product=False, grid_type="CartesianProd", dim=dim
+        )
+        self.assertIsInstance(td, TdCondTdGridDistribution)
+        self.assertEqual(td.dim, dim)
+        self.assertEqual(np.asarray(td.grid_values).shape, (n, n))
+        self.assertEqual(np.asarray(td.grid).shape, (n, 1))
+
+    def test_from_function_cartesian_product_flag(self):
+        """from_function with fun_does_cartesian_product=True."""
+        n = 8
+        dim = 2
+
+        def cond_fun_cp(a, b):
+            # a: (n, 1), b: (n, 1) → return (n, n)
+            a_np = np.asarray(a)[:, 0]
+            b_np = np.asarray(b)[:, 0]
+            diff = a_np[:, None] - b_np[None, :]
+            return np.exp(-0.5 * np.minimum(diff**2, (2 * np.pi - np.abs(diff)) ** 2))
+
+        td = TdCondTdGridDistribution.from_function(
+            cond_fun_cp,
+            n,
+            fun_does_cartesian_product=True,
+            grid_type="CartesianProd",
+            dim=dim,
+        )
+        self.assertIsInstance(td, TdCondTdGridDistribution)
+        self.assertEqual(np.asarray(td.grid_values).shape, (n, n))
+
+    def test_from_function_unknown_grid_raises(self):
+        n = 4
+        with self.assertRaises(ValueError):
+            TdCondTdGridDistribution.from_function(
+                lambda a, b: np.ones(len(a)),
+                n,
+                fun_does_cartesian_product=False,
+                grid_type="unknownGrid",
+                dim=2,
+            )
+
+    def test_from_function_odd_dim_raises(self):
+        with self.assertRaises(ValueError):
+            TdCondTdGridDistribution.from_function(
+                lambda a, b: np.ones(len(a)), 4, False, "CartesianProd", dim=3
+            )
+
+    # --------------------------------------------------------- marginalize_out
+
+    def test_marginalize_out_returns_hgd(self):
+        n = 6
+        grid = np.linspace(0.0, 2.0 * np.pi - 2.0 * np.pi / n, n).reshape(-1, 1)
+        gv = _make_normalized_grid_values(n)
+        td = TdCondTdGridDistribution(array(grid), array(gv))
+
+        for first_or_second in (1, 2):
+            with self.subTest(first_or_second=first_or_second):
+                marginal = td.marginalize_out(first_or_second)
+                self.assertIsInstance(marginal, HypertoroidalGridDistribution)
+                self.assertEqual(marginal.dim, 1)
+
+    def test_marginalize_out_sums(self):
+        n = 5
+        grid = np.linspace(0.0, 2.0 * np.pi - 2.0 * np.pi / n, n).reshape(-1, 1)
+        gv = _make_normalized_grid_values(n)
+        td = TdCondTdGridDistribution(array(grid), array(gv))
+
+        # marginalize_out(1) sums rows; HGD normalizes, so check proportionality
+        m1 = td.marginalize_out(1)
+        expected_unnorm = gv.sum(axis=0)
+        actual_unnorm = np.asarray(m1.grid_values) * float(m1.integrate())
+        # Check proportionality (ratio should be constant)
+        ratio = actual_unnorm / expected_unnorm
+        npt.assert_allclose(ratio, ratio[0] * np.ones(n), atol=1e-12)
+
+        # marginalize_out(2) sums cols
+        m2 = td.marginalize_out(2)
+        expected_unnorm2 = gv.sum(axis=1)
+        actual_unnorm2 = np.asarray(m2.grid_values) * float(m2.integrate())
+        ratio2 = actual_unnorm2 / expected_unnorm2
+        npt.assert_allclose(ratio2, ratio2[0] * np.ones(n), atol=1e-12)
+
+    # -------------------------------------------------------------- fix_dim
+
+    def test_fix_dim_returns_hgd(self):
+        n = 5
+        grid_np = np.linspace(0.0, 2.0 * np.pi - 2.0 * np.pi / n, n).reshape(-1, 1)
+        gv = _make_normalized_grid_values(n)
+        td = TdCondTdGridDistribution(array(grid_np), array(gv))
+
+        for first_or_second in (1, 2):
+            with self.subTest(first_or_second=first_or_second):
+                point = grid_np[2]  # third grid point
+                result = td.fix_dim(first_or_second, point)
+                self.assertIsInstance(result, HypertoroidalGridDistribution)
+                self.assertEqual(result.dim, 1)
+
+    def test_fix_dim_off_grid_raises(self):
+        n = 5
+        grid_np = np.linspace(0.0, 2.0 * np.pi - 2.0 * np.pi / n, n).reshape(-1, 1)
+        gv = _make_normalized_grid_values(n)
+        td = TdCondTdGridDistribution(array(grid_np), array(gv))
+        with self.assertRaises(ValueError):
+            td.fix_dim(1, np.array([1.23456789]))
+
+    def test_fix_dim_values_correct(self):
+        """fix_dim(2, grid[j]) should give a distribution proportional to col j."""
+        n = 6
+        grid_np = np.linspace(0.0, 2.0 * np.pi - 2.0 * np.pi / n, n).reshape(-1, 1)
+        gv = _make_normalized_grid_values(n)
+        td = TdCondTdGridDistribution(array(grid_np), array(gv))
+
+        j = 3
+        slice_dist = td.fix_dim(2, grid_np[j])
+        expected = gv[:, j]
+        expected = expected / expected.mean() / (2.0 * np.pi)  # normalize
+        npt.assert_allclose(
+            np.asarray(slice_dist.grid_values),
+            expected,
+            atol=1e-12,
+        )
+
+    # --------------------------------------------------------- from_function + fix_dim round-trip
+
+    def test_from_function_fix_dim_roundtrip(self):
+        """fix_dim on a from_function object should use cartesian_prod layout."""
+        n = 10
+        dim = 2
+
+        def cond_fun(a, b):
+            diff = np.asarray(a)[:, 0] - np.asarray(b)[:, 0]
+            return np.exp(-0.5 * diff**2)
+
+        td = TdCondTdGridDistribution.from_function(
+            cond_fun, n, fun_does_cartesian_product=False, dim=dim
+        )
+        grid_np = np.asarray(td.grid)
+        slice_dist = td.fix_dim(2, grid_np[0])
+        self.assertIsInstance(slice_dist, HypertoroidalGridDistribution)
+        self.assertEqual(slice_dist.grid_type, "cartesian_prod")
+
+
+if __name__ == "__main__":
+    unittest.main()

From 42352bcc9f2bf7605475637da096a4607e4eb0df Mon Sep 17 00:00:00 2001
From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com>
Date: Tue, 31 Mar 2026 12:57:35 +0000
Subject: [PATCH 02/19] Fix MegaLinter errors: remove unused imports, reduce
 locals, fix protected-access

Agent-Logs-Url: https://github.com/FlorianPfaff/PyRecEst/sessions/0221331a-f5da-4059-aae5-86498d14a29d

Co-authored-by: FlorianPfaff <6773539+FlorianPfaff@users.noreply.github.com>
---
 .../td_cond_td_grid_distribution.py           | 32 +++++++++----------
 .../test_td_cond_td_grid_distribution.py      |  3 +-
 2 files changed, 17 insertions(+), 18 deletions(-)

diff --git a/pyrecest/distributions/hypertorus/td_cond_td_grid_distribution.py b/pyrecest/distributions/hypertorus/td_cond_td_grid_distribution.py
index df18c439e..ce9e99217 100644
--- a/pyrecest/distributions/hypertorus/td_cond_td_grid_distribution.py
+++ b/pyrecest/distributions/hypertorus/td_cond_td_grid_distribution.py
@@ -8,7 +8,7 @@
 from beartype import beartype
 
 # pylint: disable=redefined-builtin,no-name-in-module,no-member
-from pyrecest.backend import abs, all, any, array, mean, meshgrid, pi, reshape
+from pyrecest.backend import abs, all, any, array, mean, pi, reshape
 
 from ..conditional.abstract_conditional_distribution import (
     AbstractConditionalDistribution,
@@ -48,6 +48,7 @@ def __init__(
         grid,
         grid_values,
         enforce_pdf_nonnegative: bool = True,
+        n_grid_points_per_dim: tuple | None = None,
     ):
         super().__init__()
 
@@ -65,9 +66,9 @@ def __init__(
         self.enforce_pdf_nonnegative = enforce_pdf_nonnegative
         # dim is the combined dimension of the two-hypertorus Cartesian product
         self.dim = 2 * grid.shape[1]
-        # Populated by from_function; used to reshape slices back to the
-        # correct multi-dimensional layout when creating HypertoroidalGridDistribution
-        self._n_grid_points_per_dim: tuple[int, ...] | None = None
+        # Used to reshape slices back to the correct multi-dimensional layout
+        # when creating HypertoroidalGridDistribution instances.
+        self._n_grid_points_per_dim: tuple[int, ...] | None = n_grid_points_per_dim
 
         self._check_normalization()
 
@@ -252,22 +253,21 @@ def from_function(
         n_total = grid.shape[0]
 
         if fun_does_cartesian_product:
-            fvals = fun(grid, grid)
-            fvals = reshape(array(fvals), (n_total, n_total))
+            fvals = reshape(array(fun(grid, grid)), (n_total, n_total))
         else:
             # Build index arrays so that fvals[i, j] = f(grid[i], grid[j]).
             # Using plain numpy for index generation (pure control flow).
-            r = np.arange(n_total)
-            idx_i_np, idx_j_np = np.meshgrid(r, r, indexing="ij")
-            idx_i = idx_i_np.ravel()
-            idx_j = idx_j_np.ravel()
-
-            fvals_flat = fun(grid[idx_i], grid[idx_j])
-            fvals = reshape(array(fvals_flat), (n_total, n_total))
+            idx_i, idx_j = (
+                m.ravel()
+                for m in np.meshgrid(
+                    np.arange(n_total), np.arange(n_total), indexing="ij"
+                )
+            )
+            fvals = reshape(array(fun(grid[idx_i], grid[idx_j])), (n_total, n_total))
 
-        td = TdCondTdGridDistribution(grid, fvals)
-        td._n_grid_points_per_dim = tuple(n_grid_points_seq)
-        return td
+        return TdCondTdGridDistribution(
+            grid, fvals, n_grid_points_per_dim=tuple(n_grid_points_seq)
+        )
 
     # --------------------------------------------------------------- helpers
 
diff --git a/pyrecest/tests/distributions/test_td_cond_td_grid_distribution.py b/pyrecest/tests/distributions/test_td_cond_td_grid_distribution.py
index 4c183ee0b..1e66b663c 100644
--- a/pyrecest/tests/distributions/test_td_cond_td_grid_distribution.py
+++ b/pyrecest/tests/distributions/test_td_cond_td_grid_distribution.py
@@ -1,10 +1,9 @@
 import unittest
-import warnings
 
 import numpy as np
 import numpy.testing as npt
 
-from pyrecest.backend import array, pi
+from pyrecest.backend import array
 from pyrecest.distributions.hypertorus.hypertoroidal_grid_distribution import (
     HypertoroidalGridDistribution,
 )

From 7998df673f186d8b478c74181a7fa54a0d0628a1 Mon Sep 17 00:00:00 2001
From: Florian Pfaff <pfaff@ias.uni-stuttgart.de>
Date: Tue, 31 Mar 2026 15:39:31 +0200
Subject: [PATCH 03/19] Increase tolerance to fix

---
 .../tests/distributions/test_td_cond_td_grid_distribution.py  | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/pyrecest/tests/distributions/test_td_cond_td_grid_distribution.py b/pyrecest/tests/distributions/test_td_cond_td_grid_distribution.py
index 1e66b663c..dd7df6624 100644
--- a/pyrecest/tests/distributions/test_td_cond_td_grid_distribution.py
+++ b/pyrecest/tests/distributions/test_td_cond_td_grid_distribution.py
@@ -30,8 +30,8 @@ def test_construction_t1(self):
         gv = _make_normalized_grid_values(n)
         td = TdCondTdGridDistribution(array(grid), array(gv))
         self.assertEqual(td.dim, 2)
-        npt.assert_array_equal(np.asarray(td.grid), grid)
-        npt.assert_array_equal(np.asarray(td.grid_values), gv)
+        npt.assert_allclose(td.grid, grid, rtol=1e-6)
+        npt.assert_allclose(td.grid_values, gv, rtol=1e-6)
 
     def test_construction_wrong_shape_raises(self):
         n = 4

From e0710524963e554c2174a7e199217088ac2207ca Mon Sep 17 00:00:00 2001
From: FlorianPfaff <6773539+FlorianPfaff@users.noreply.github.com>
Date: Tue, 31 Mar 2026 12:21:35 +0000
Subject: [PATCH 04/19] [MegaLinter] Apply linters automatic fixes

Co-authored-by: FlorianPfaff <6773539+FlorianPfaff@users.noreply.github.com>
---
 .../test_spherical_harmonics_distribution_complex.py            | 2 +-
 pyrecest/tests/test_evaluation_basic.py                         | 2 +-
 2 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/pyrecest/tests/distributions/test_spherical_harmonics_distribution_complex.py b/pyrecest/tests/distributions/test_spherical_harmonics_distribution_complex.py
index 31241def1..2383b78d2 100644
--- a/pyrecest/tests/distributions/test_spherical_harmonics_distribution_complex.py
+++ b/pyrecest/tests/distributions/test_spherical_harmonics_distribution_complex.py
@@ -1,5 +1,5 @@
-import warnings
 import unittest
+import warnings
 
 import numpy.testing as npt
 import pyrecest.backend
diff --git a/pyrecest/tests/test_evaluation_basic.py b/pyrecest/tests/test_evaluation_basic.py
index 00e7ae643..c21f7e7ba 100644
--- a/pyrecest/tests/test_evaluation_basic.py
+++ b/pyrecest/tests/test_evaluation_basic.py
@@ -1,7 +1,7 @@
-import warnings
 import os
 import tempfile
 import unittest
+import warnings
 from typing import Optional
 
 import matplotlib

From 9673238de58a6a1ec3a414c94977226d10673e0a Mon Sep 17 00:00:00 2001
From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com>
Date: Tue, 31 Mar 2026 12:25:46 +0000
Subject: [PATCH 05/19] Fix 3 bugs in EOT shape database: sample_within return
 type, visibility polygon dimension, and @classmethod on setUp

Agent-Logs-Url: https://github.com/FlorianPfaff/PyRecEst/sessions/27f787c9-722d-4cc0-93e6-0ff0f3f8e4f6

Co-authored-by: FlorianPfaff <6773539+FlorianPfaff@users.noreply.github.com>
---
 pyrecest/evaluation/eot_shape_database.py | 4 ++--
 pyrecest/tests/test_eot_shape_database.py | 1 -
 2 files changed, 2 insertions(+), 3 deletions(-)

diff --git a/pyrecest/evaluation/eot_shape_database.py b/pyrecest/evaluation/eot_shape_database.py
index 55b048605..685d8f295 100644
--- a/pyrecest/evaluation/eot_shape_database.py
+++ b/pyrecest/evaluation/eot_shape_database.py
@@ -55,7 +55,7 @@ def sample_within(self, num_points: int) -> np.ndarray:
 
             points[i] = random_point
 
-        return np.array(points)
+        return np.array([(point.x, point.y) for point in points])
 
 
 class StarShapedPolygon(PolygonWithSampling):  # pylint: disable=abstract-method
@@ -105,7 +105,7 @@ def compute_visibility_polygon(self, point):
                 segments.append(line_of_sight)
 
             # Check for intersections along the direction to the vertex
-            direction = np.array(vertex) - np.array(point.coords)
+            direction = np.array(vertex) - np.array(point.coords[0])
             far_away_point = Point(np.array(vertex) + 1000 * direction)
             ray = LineString([point, far_away_point])
 
diff --git a/pyrecest/tests/test_eot_shape_database.py b/pyrecest/tests/test_eot_shape_database.py
index 1048724e9..eba40be2d 100644
--- a/pyrecest/tests/test_eot_shape_database.py
+++ b/pyrecest/tests/test_eot_shape_database.py
@@ -56,7 +56,6 @@ def test_plotting(self):
 
 
 class TestStarfish(unittest.TestCase):
-    @classmethod
     def setUp(self):
         self.starfish = StarFish()
         # self.starfish_kernel = self.starfish.compute_kernel()

From 32930f7dbd85bac4762fe40bc4b0769332fced8d Mon Sep 17 00:00:00 2001
From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com>
Date: Tue, 31 Mar 2026 12:32:37 +0000
Subject: [PATCH 06/19] Implement SdCondSdGridDistribution with tests

Agent-Logs-Url: https://github.com/FlorianPfaff/PyRecEst/sessions/1c6bef08-40d1-4a9d-9810-72aeaceca573

Co-authored-by: FlorianPfaff <6773539+FlorianPfaff@users.noreply.github.com>
---
 pyrecest/distributions/__init__.py            |   2 +
 .../distributions/conditional/__init__.py     |   1 +
 .../sd_cond_sd_grid_distribution.py           | 300 +++++++++++++++
 .../test_sd_cond_sd_grid_distribution.py      | 357 ++++++++++++++++++
 4 files changed, 660 insertions(+)
 create mode 100644 pyrecest/distributions/conditional/sd_cond_sd_grid_distribution.py
 create mode 100644 pyrecest/tests/distributions/test_sd_cond_sd_grid_distribution.py

diff --git a/pyrecest/distributions/__init__.py b/pyrecest/distributions/__init__.py
index cf98b9104..72efac1ad 100644
--- a/pyrecest/distributions/__init__.py
+++ b/pyrecest/distributions/__init__.py
@@ -91,6 +91,7 @@
 from .conditional.abstract_conditional_distribution import (
     AbstractConditionalDistribution,
 )
+from .conditional.sd_cond_sd_grid_distribution import SdCondSdGridDistribution
 from .custom_hyperrectangular_distribution import CustomHyperrectangularDistribution
 from .disk_uniform_distribution import DiskUniformDistribution
 from .ellipsoidal_ball_uniform_distribution import EllipsoidalBallUniformDistribution
@@ -275,6 +276,7 @@
     "WrappedLaplaceDistribution",
     "WrappedNormalDistribution",
     "AbstractConditionalDistribution",
+    "SdCondSdGridDistribution",
     "CustomHyperrectangularDistribution",
     "DiskUniformDistribution",
     "EllipsoidalBallUniformDistribution",
diff --git a/pyrecest/distributions/conditional/__init__.py b/pyrecest/distributions/conditional/__init__.py
index e69de29bb..e339fdac4 100644
--- a/pyrecest/distributions/conditional/__init__.py
+++ b/pyrecest/distributions/conditional/__init__.py
@@ -0,0 +1 @@
+from .sd_cond_sd_grid_distribution import SdCondSdGridDistribution
diff --git a/pyrecest/distributions/conditional/sd_cond_sd_grid_distribution.py b/pyrecest/distributions/conditional/sd_cond_sd_grid_distribution.py
new file mode 100644
index 000000000..779c785f0
--- /dev/null
+++ b/pyrecest/distributions/conditional/sd_cond_sd_grid_distribution.py
@@ -0,0 +1,300 @@
+import copy
+import warnings
+
+# pylint: disable=redefined-builtin,no-name-in-module,no-member
+from pyrecest.backend import (
+    abs,
+    all,
+    any,
+    arange,
+    argmin,
+    array_equal,
+    linalg,
+    mean,
+    meshgrid,
+    sum,
+)
+
+from pyrecest.distributions.hypersphere_subset.abstract_hypersphere_subset_distribution import (
+    AbstractHypersphereSubsetDistribution,
+)
+
+from .abstract_conditional_distribution import AbstractConditionalDistribution
+
+
+class SdCondSdGridDistribution(AbstractConditionalDistribution):
+    """
+    Conditional distribution on Sd x Sd represented by a grid of values.
+
+    For a conditional distribution f(a|b), ``grid_values[i, j]`` stores
+    the value f(grid[i] | grid[j]).
+
+    Convention
+    ----------
+    - ``grid`` has shape ``(n_points, d)`` where ``d`` is the embedding
+      dimension of the individual sphere (e.g. d=3 for S2).
+    - ``grid_values`` has shape ``(n_points, n_points)``.
+    - ``dim = 2 * d`` is the embedding dimension of the Cartesian product
+      space (convention inherited from libDirectional).
+    """
+
+    def __init__(self, grid, grid_values, enforce_pdf_nonnegative=True):
+        """
+        Parameters
+        ----------
+        grid : array of shape (n_points, d)
+            Grid points on the sphere.  All coordinate values must lie in
+            [-1, 1] (unit sphere).
+        grid_values : array of shape (n_points, n_points)
+            Conditional pdf values: ``grid_values[i, j] = f(grid[i] | grid[j])``.
+            Must be non-negative.
+        enforce_pdf_nonnegative : bool
+            Whether non-negativity of ``grid_values`` is required.
+        """
+        if grid.ndim != 2:
+            raise ValueError("grid must be a 2D array of shape (n_points, d).")
+
+        n_points, d = grid.shape
+
+        if grid_values.ndim != 2 or grid_values.shape != (n_points, n_points):
+            raise ValueError(
+                f"grid_values must be a square 2D array of shape ({n_points}, {n_points})."
+            )
+
+        if any(abs(grid) > 1 + 1e-12):
+            raise ValueError(
+                "Grid points must have coordinates in [-1, 1] (unit sphere)."
+            )
+
+        if enforce_pdf_nonnegative and any(grid_values < 0):
+            raise ValueError("grid_values must be non-negative.")
+
+        self.grid = grid
+        self.grid_values = grid_values
+        self.enforce_pdf_nonnegative = enforce_pdf_nonnegative
+        # Embedding dimension of the Cartesian product space (convention from
+        # libDirectional: dim = 2 * embedding_dim_of_individual_sphere).
+        self.dim = 2 * d
+
+        self._check_normalization()
+
+    # ------------------------------------------------------------------
+    # Normalization
+    # ------------------------------------------------------------------
+
+    def _check_normalization(self, tol=0.01):
+        """Warn if any column is not normalized to 1 over the sphere."""
+        sphere_surface = (
+            AbstractHypersphereSubsetDistribution.compute_unit_hypersphere_surface(
+                self.grid.shape[1] - 1
+            )
+        )
+        # For each fixed second argument j, the mean over i times the sphere
+        # surface area should equal 1.
+        ints = mean(self.grid_values, axis=0) * sphere_surface
+        if any(abs(ints - 1) > tol):
+            # Check whether swapping the two arguments would yield normalisation.
+            ints_swapped = mean(self.grid_values, axis=1) * sphere_surface
+            if all(abs(ints_swapped - 1) <= tol):
+                raise ValueError(
+                    "Normalization:maybeWrongOrder: Not normalized but would be if "
+                    "the order of the two spheres were swapped. Check input."
+                )
+            warnings.warn(
+                "Normalization:notNormalized: When conditioning values for the first "
+                "sphere on the second, normalisation is not ensured. "
+                "Check input or increase tolerance. "
+                "No normalisation is performed; you may want to do this manually.",
+                UserWarning,
+            )
+
+    def normalize(self):
+        """No-op – returns ``self`` for compatibility."""
+        return self
+
+    # ------------------------------------------------------------------
+    # Arithmetic
+    # ------------------------------------------------------------------
+
+    def multiply(self, other):
+        """
+        Element-wise multiply two conditional grid distributions.
+
+        The resulting distribution is *not* normalized.
+
+        Parameters
+        ----------
+        other : SdCondSdGridDistribution
+            Must be defined on the same grid.
+
+        Returns
+        -------
+        SdCondSdGridDistribution
+        """
+        if not array_equal(self.grid, other.grid):
+            raise ValueError(
+                "Multiply:IncompatibleGrid: Can only multiply distributions "
+                "defined on identical grids."
+            )
+        warnings.warn(
+            "Multiply:UnnormalizedResult: Multiplication does not yield a "
+            "normalized result.",
+            UserWarning,
+        )
+        result = copy.deepcopy(self)
+        result.grid_values = result.grid_values * other.grid_values
+        return result
+
+    # ------------------------------------------------------------------
+    # Marginalisation and conditioning
+    # ------------------------------------------------------------------
+
+    def marginalize_out(self, first_or_second):
+        """
+        Marginalize out one of the two spheres.
+
+        Parameters
+        ----------
+        first_or_second : int  (1 or 2)
+            ``1`` marginalizes out the *first* dimension (sums over rows);
+            ``2`` marginalizes out the *second* dimension (sums over columns).
+
+        Returns
+        -------
+        HypersphericalGridDistribution
+        """
+        # Import here to avoid circular imports
+        from pyrecest.distributions.hypersphere_subset.hyperspherical_grid_distribution import (  # pylint: disable=import-outside-toplevel
+            HypersphericalGridDistribution,
+        )
+
+        if first_or_second == 1:
+            grid_values_sgd = sum(self.grid_values, axis=0)
+        elif first_or_second == 2:
+            grid_values_sgd = sum(self.grid_values, axis=1)
+        else:
+            raise ValueError("first_or_second must be 1 or 2.")
+
+        return HypersphericalGridDistribution(self.grid, grid_values_sgd)
+
+    def fix_dim(self, first_or_second, point):
+        """
+        Return the conditional slice for a fixed grid point.
+
+        The supplied ``point`` must be an existing grid point.
+
+        Parameters
+        ----------
+        first_or_second : int  (1 or 2)
+            ``1`` fixes the *first* argument at ``point`` and returns values
+            over the second sphere;
+            ``2`` fixes the *second* argument and returns values over the
+            first sphere.
+        point : array of shape (d,)
+            Grid point at which to evaluate.
+
+        Returns
+        -------
+        HypersphericalGridDistribution
+        """
+        # Import here to avoid circular imports
+        from pyrecest.distributions.hypersphere_subset.hyperspherical_grid_distribution import (  # pylint: disable=import-outside-toplevel
+            HypersphericalGridDistribution,
+        )
+
+        d = self.grid.shape[1]
+        if point.shape[0] != d:
+            raise ValueError(
+                f"point must have length {d} (embedding dimension of the sphere)."
+            )
+
+        diffs = linalg.norm(self.grid - point[None, :], axis=1)
+        locb = argmin(diffs)
+        if diffs[locb] > 1e-10:
+            raise ValueError(
+                "Cannot fix value at this point because it is not on the grid."
+            )
+
+        if first_or_second == 1:
+            grid_values_slice = self.grid_values[locb, :]
+        elif first_or_second == 2:
+            grid_values_slice = self.grid_values[:, locb]
+        else:
+            raise ValueError("first_or_second must be 1 or 2.")
+
+        return HypersphericalGridDistribution(self.grid, grid_values_slice)
+
+    # ------------------------------------------------------------------
+    # Factory
+    # ------------------------------------------------------------------
+
+    @staticmethod
+    def from_function(
+        fun,
+        no_of_grid_points,
+        fun_does_cartesian_product=False,
+        grid_type="leopardi",
+        dim=6,
+    ):
+        """
+        Construct a :class:`SdCondSdGridDistribution` from a callable.
+
+        Parameters
+        ----------
+        fun : callable
+            Conditional pdf function ``f(a, b)``.
+
+            *   When ``fun_does_cartesian_product=False`` (default): ``fun``
+                is called once with two arrays of shape ``(n_pairs, d)``
+                containing all ``n_points²`` pairs, and must return a 1-D
+                array of length ``n_points²``.
+            *   When ``fun_does_cartesian_product=True``: ``fun`` is called
+                with both full grids of shape ``(n_points, d)`` and must
+                return an array of shape ``(n_points, n_points)``.
+        no_of_grid_points : int
+            Number of grid points for each sphere.
+        fun_does_cartesian_product : bool
+            See ``fun`` description above.
+        grid_type : str
+            Grid type passed to :func:`~pyrecest.sampling.hyperspherical_sampler.get_grid_hypersphere`.
+            Defaults to ``'leopardi'``.
+        dim : int
+            Embedding dimension of the Cartesian product space
+            (``2 * embedding_dim_of_individual_sphere``).
+            Defaults to 6 (S2 × S2).
+
+        Returns
+        -------
+        SdCondSdGridDistribution
+        """
+        # Import inside the function to avoid circular imports at module level.
+        from pyrecest.sampling.hyperspherical_sampler import (  # pylint: disable=import-outside-toplevel
+            get_grid_hypersphere,
+        )
+
+        n = no_of_grid_points
+        # Convert from Cartesian-product embedding dim to individual sphere
+        # manifold dim: embedding_dim = dim // 2, manifold_dim = embedding_dim - 1.
+        manifold_dim = dim // 2 - 1
+        grid, _ = get_grid_hypersphere(grid_type, n, manifold_dim)
+        # grid is (n, dim//2)
+
+        if fun_does_cartesian_product:
+            fvals = fun(grid, grid)
+            grid_values = fvals.reshape(n, n)
+        else:
+            # Build index pairs: idx_a[i, j] = i, idx_b[i, j] = j
+            idx_a, idx_b = meshgrid(arange(n), arange(n), indexing="ij")
+            grid_a = grid[idx_a.ravel()]  # (n*n, d)
+            grid_b = grid[idx_b.ravel()]  # (n*n, d)
+            fvals = fun(grid_a, grid_b)  # (n*n,)
+
+            if fvals.shape == (n**2, n**2):
+                raise ValueError(
+                    "Function apparently performs the Cartesian product itself. "
+                    "Set fun_does_cartesian_product=True."
+                )
+
+            grid_values = fvals.reshape(n, n)
+
+        return SdCondSdGridDistribution(grid, grid_values)
diff --git a/pyrecest/tests/distributions/test_sd_cond_sd_grid_distribution.py b/pyrecest/tests/distributions/test_sd_cond_sd_grid_distribution.py
new file mode 100644
index 000000000..84132451f
--- /dev/null
+++ b/pyrecest/tests/distributions/test_sd_cond_sd_grid_distribution.py
@@ -0,0 +1,357 @@
+import unittest
+import warnings
+
+import numpy.testing as npt
+import pyrecest
+from pyrecest.backend import (
+    array,
+    ones,
+    zeros,
+)
+from pyrecest.distributions.conditional.sd_cond_sd_grid_distribution import (
+    SdCondSdGridDistribution,
+)
+from pyrecest.distributions.hypersphere_subset.abstract_hypersphere_subset_distribution import (
+    AbstractHypersphereSubsetDistribution,
+)
+from pyrecest.distributions.hypersphere_subset.hyperspherical_grid_distribution import (
+    HypersphericalGridDistribution,
+)
+
+
+def _make_uniform_conditional(no_of_grid_points=100, dim=6):
+    """Helper: build a properly normalized SdCondSdGridDistribution with uniform values."""
+    from pyrecest.sampling.hyperspherical_sampler import get_grid_hypersphere
+
+    manifold_dim = dim // 2 - 1
+    grid, _ = get_grid_hypersphere("leopardi", no_of_grid_points, manifold_dim)
+    d = grid.shape[1]  # embedding dim
+    surface = AbstractHypersphereSubsetDistribution.compute_unit_hypersphere_surface(
+        d - 1
+    )
+    # Uniform pdf value for each column so that mean * surface == 1
+    uniform_val = 1.0 / surface
+    grid_values = ones((no_of_grid_points, no_of_grid_points)) * uniform_val
+    return SdCondSdGridDistribution(grid, grid_values)
+
+
+class TestSdCondSdGridDistributionInit(unittest.TestCase):
+    """Tests for __init__ validation."""
+
+    @unittest.skipIf(
+        pyrecest.backend.__backend_name__ == "jax",
+        reason="Not supported on JAX backend",
+    )
+    def test_basic_construction_s2(self):
+        """Construct a valid SdCondSdGridDistribution on S2 x S2."""
+        sc = _make_uniform_conditional(no_of_grid_points=50)
+        self.assertEqual(sc.dim, 6)
+        self.assertEqual(sc.grid.shape[1], 3)
+        self.assertEqual(sc.grid_values.shape, (50, 50))
+
+    @unittest.skipIf(
+        pyrecest.backend.__backend_name__ == "jax",
+        reason="Not supported on JAX backend",
+    )
+    def test_wrong_grid_shape_raises(self):
+        """Providing a 1-D grid must raise ValueError."""
+        from pyrecest.sampling.hyperspherical_sampler import get_grid_hypersphere
+
+        grid, _ = get_grid_hypersphere("leopardi", 10, 2)
+        n = grid.shape[0]
+        surface = AbstractHypersphereSubsetDistribution.compute_unit_hypersphere_surface(
+            2
+        )
+        gv = ones((n, n)) / surface
+        with self.assertRaises(ValueError):
+            SdCondSdGridDistribution(grid.ravel(), gv)
+
+    @unittest.skipIf(
+        pyrecest.backend.__backend_name__ == "jax",
+        reason="Not supported on JAX backend",
+    )
+    def test_non_square_grid_values_raises(self):
+        from pyrecest.sampling.hyperspherical_sampler import get_grid_hypersphere
+
+        grid, _ = get_grid_hypersphere("leopardi", 10, 2)
+        n = grid.shape[0]
+        with self.assertRaises(ValueError):
+            SdCondSdGridDistribution(grid, ones((n, n + 1)))
+
+    @unittest.skipIf(
+        pyrecest.backend.__backend_name__ == "jax",
+        reason="Not supported on JAX backend",
+    )
+    def test_negative_grid_values_raises(self):
+        from pyrecest.sampling.hyperspherical_sampler import get_grid_hypersphere
+
+        grid, _ = get_grid_hypersphere("leopardi", 10, 2)
+        n = grid.shape[0]
+        surface = AbstractHypersphereSubsetDistribution.compute_unit_hypersphere_surface(
+            2
+        )
+        gv = ones((n, n)) / surface
+        gv[0, 0] = -1.0
+        with self.assertRaises(ValueError):
+            SdCondSdGridDistribution(grid, gv)
+
+    @unittest.skipIf(
+        pyrecest.backend.__backend_name__ == "jax",
+        reason="Not supported on JAX backend",
+    )
+    def test_unnormalized_warns(self):
+        from pyrecest.sampling.hyperspherical_sampler import get_grid_hypersphere
+
+        grid, _ = get_grid_hypersphere("leopardi", 10, 2)
+        n = grid.shape[0]
+        # Deliberately unnormalized (all ones instead of 1/surface)
+        gv = ones((n, n))
+        with self.assertWarns(UserWarning):
+            SdCondSdGridDistribution(grid, gv)
+
+
+class TestSdCondSdGridDistributionNormalize(unittest.TestCase):
+    """normalize() is a no-op and returns self."""
+
+    @unittest.skipIf(
+        pyrecest.backend.__backend_name__ == "jax",
+        reason="Not supported on JAX backend",
+    )
+    def test_normalize_returns_self(self):
+        sc = _make_uniform_conditional(50)
+        result = sc.normalize()
+        self.assertIs(result, sc)
+
+
+class TestSdCondSdGridDistributionMultiply(unittest.TestCase):
+    """Tests for multiply()."""
+
+    @unittest.skipIf(
+        pyrecest.backend.__backend_name__ == "jax",
+        reason="Not supported on JAX backend",
+    )
+    def test_multiply_same_grid(self):
+        sc1 = _make_uniform_conditional(20)
+        sc2 = _make_uniform_conditional(20)
+        # Force identical grids
+        sc2.grid = sc1.grid
+        with warnings.catch_warnings():
+            warnings.simplefilter("ignore")
+            result = sc1.multiply(sc2)
+        npt.assert_allclose(
+            result.grid_values, sc1.grid_values * sc2.grid_values, rtol=1e-10
+        )
+
+    @unittest.skipIf(
+        pyrecest.backend.__backend_name__ == "jax",
+        reason="Not supported on JAX backend",
+    )
+    def test_multiply_incompatible_grid_raises(self):
+        sc1 = _make_uniform_conditional(20)
+        sc2 = _make_uniform_conditional(30)  # Different number of grid points
+        with self.assertRaises(ValueError) as ctx:
+            sc1.multiply(sc2)
+        self.assertIn("IncompatibleGrid", str(ctx.exception))
+
+
+class TestSdCondSdGridDistributionMarginalizeOut(unittest.TestCase):
+    """Tests for marginalize_out()."""
+
+    @unittest.skipIf(
+        pyrecest.backend.__backend_name__ == "jax",
+        reason="Not supported on JAX backend",
+    )
+    def test_marginalize_out_returns_hgd(self):
+        sc = _make_uniform_conditional(20)
+        sgd1 = sc.marginalize_out(1)
+        sgd2 = sc.marginalize_out(2)
+        self.assertIsInstance(sgd1, HypersphericalGridDistribution)
+        self.assertIsInstance(sgd2, HypersphericalGridDistribution)
+
+    @unittest.skipIf(
+        pyrecest.backend.__backend_name__ == "jax",
+        reason="Not supported on JAX backend",
+    )
+    def test_marginalize_out_shape(self):
+        n = 20
+        sc = _make_uniform_conditional(n)
+        sgd1 = sc.marginalize_out(1)
+        sgd2 = sc.marginalize_out(2)
+        self.assertEqual(sgd1.grid_values.shape, (n,))
+        self.assertEqual(sgd2.grid_values.shape, (n,))
+
+    @unittest.skipIf(
+        pyrecest.backend.__backend_name__ == "jax",
+        reason="Not supported on JAX backend",
+    )
+    def test_marginalize_out_uniform_is_uniform(self):
+        """For a uniform conditional, marginalizing either way gives uniform."""
+        n = 30
+        sc = _make_uniform_conditional(n)
+        sgd1 = sc.marginalize_out(1)
+        sgd2 = sc.marginalize_out(2)
+        # After normalisation in HypersphericalGridDistribution, all values equal
+        npt.assert_allclose(
+            sgd1.grid_values,
+            sgd1.grid_values[0] * ones(n),
+            rtol=1e-10,
+        )
+        npt.assert_allclose(
+            sgd2.grid_values,
+            sgd2.grid_values[0] * ones(n),
+            rtol=1e-10,
+        )
+
+    @unittest.skipIf(
+        pyrecest.backend.__backend_name__ == "jax",
+        reason="Not supported on JAX backend",
+    )
+    def test_marginalize_out_invalid_raises(self):
+        sc = _make_uniform_conditional(10)
+        with self.assertRaises(ValueError):
+            sc.marginalize_out(0)
+        with self.assertRaises(ValueError):
+            sc.marginalize_out(3)
+
+
+class TestSdCondSdGridDistributionFixDim(unittest.TestCase):
+    """Tests for fix_dim()."""
+
+    @unittest.skipIf(
+        pyrecest.backend.__backend_name__ == "jax",
+        reason="Not supported on JAX backend",
+    )
+    def test_fix_dim_returns_hgd(self):
+        n = 20
+        sc = _make_uniform_conditional(n)
+        point = sc.grid[0, :]
+        sgd1 = sc.fix_dim(1, point)
+        sgd2 = sc.fix_dim(2, point)
+        self.assertIsInstance(sgd1, HypersphericalGridDistribution)
+        self.assertIsInstance(sgd2, HypersphericalGridDistribution)
+
+    @unittest.skipIf(
+        pyrecest.backend.__backend_name__ == "jax",
+        reason="Not supported on JAX backend",
+    )
+    def test_fix_dim_values_correct(self):
+        n = 20
+        sc = _make_uniform_conditional(n)
+        # Fix first dim at grid[5]
+        sgd = sc.fix_dim(1, sc.grid[5, :])
+        npt.assert_allclose(
+            sgd.grid_values,
+            sc.fix_dim(1, sc.grid[5, :]).grid_values,
+            rtol=1e-10,
+        )
+
+    @unittest.skipIf(
+        pyrecest.backend.__backend_name__ == "jax",
+        reason="Not supported on JAX backend",
+    )
+    def test_fix_dim_off_grid_raises(self):
+        sc = _make_uniform_conditional(20)
+        # Supply a point that is not on the grid
+        off_grid = array([0.1, 0.2, 0.97])  # arbitrary off-grid S2 point
+        # Normalise to unit sphere
+        from pyrecest.backend import linalg
+
+        off_grid = off_grid / linalg.norm(off_grid)
+        with self.assertRaises(ValueError):
+            sc.fix_dim(1, off_grid)
+
+    @unittest.skipIf(
+        pyrecest.backend.__backend_name__ == "jax",
+        reason="Not supported on JAX backend",
+    )
+    def test_fix_dim_invalid_raises(self):
+        sc = _make_uniform_conditional(10)
+        point = sc.grid[0, :]
+        with self.assertRaises(ValueError):
+            sc.fix_dim(0, point)
+        with self.assertRaises(ValueError):
+            sc.fix_dim(3, point)
+
+
+class TestSdCondSdGridDistributionFromFunction(unittest.TestCase):
+    """Tests for from_function()."""
+
+    @unittest.skipIf(
+        pyrecest.backend.__backend_name__ == "jax",
+        reason="Not supported on JAX backend",
+    )
+    def test_from_function_uniform(self):
+        """Build a uniform conditional from a constant function."""
+        from pyrecest.backend import ones as be_ones
+
+        d = 3  # S2 embedding dim; manifold dim = 2
+        surface = AbstractHypersphereSubsetDistribution.compute_unit_hypersphere_surface(
+            d - 1
+        )
+
+        def uniform_fun(a, _b):
+            return be_ones(a.shape[0]) / surface
+
+        n = 50
+        sc = SdCondSdGridDistribution.from_function(
+            uniform_fun, n, fun_does_cartesian_product=False, grid_type="leopardi", dim=6
+        )
+        self.assertEqual(sc.dim, 6)
+        self.assertEqual(sc.grid_values.shape, (n, n))
+        # All values should be equal to 1/surface
+        npt.assert_allclose(sc.grid_values, ones((n, n)) / surface, rtol=1e-10)
+
+    @unittest.skipIf(
+        pyrecest.backend.__backend_name__ == "jax",
+        reason="Not supported on JAX backend",
+    )
+    def test_from_function_cartesian_product(self):
+        """from_function with fun_does_cartesian_product=True."""
+        from pyrecest.backend import ones as be_ones
+
+        d = 3
+        surface = AbstractHypersphereSubsetDistribution.compute_unit_hypersphere_surface(
+            d - 1
+        )
+        n = 30
+
+        def uniform_fun_cp(grid_a, grid_b):
+            # grid_a and grid_b are both (n, d); return (n, n)
+            return be_ones((grid_a.shape[0], grid_b.shape[0])) / surface
+
+        sc = SdCondSdGridDistribution.from_function(
+            uniform_fun_cp,
+            n,
+            fun_does_cartesian_product=True,
+            grid_type="leopardi",
+            dim=6,
+        )
+        self.assertEqual(sc.grid_values.shape, (n, n))
+        npt.assert_allclose(sc.grid_values, ones((n, n)) / surface, rtol=1e-10)
+
+    @unittest.skipIf(
+        pyrecest.backend.__backend_name__ == "jax",
+        reason="Not supported on JAX backend",
+    )
+    def test_from_function_wrong_shape_raises(self):
+        """Function that returns wrong shape should raise ValueError."""
+        from pyrecest.backend import ones as be_ones
+
+        n = 10
+
+        def bad_fun(a, b):
+            # Return (n*n, n*n) – as if doing Cartesian product
+            return be_ones((a.shape[0], b.shape[0]))
+
+        with self.assertRaises(ValueError):
+            SdCondSdGridDistribution.from_function(
+                bad_fun,
+                n,
+                fun_does_cartesian_product=False,
+                grid_type="leopardi",
+                dim=6,
+            )
+
+
+if __name__ == "__main__":
+    unittest.main()

From f33c532bc4dffc127e3181e7a24d3df92070a201 Mon Sep 17 00:00:00 2001
From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com>
Date: Tue, 31 Mar 2026 12:59:17 +0000
Subject: [PATCH 07/19] Fix MegaLinter errors: remove unused import, add pylint
 disable comments, add __all__

Agent-Logs-Url: https://github.com/FlorianPfaff/PyRecEst/sessions/d724643e-1252-4c50-b114-57aac876ac0f

Co-authored-by: FlorianPfaff <6773539+FlorianPfaff@users.noreply.github.com>
---
 .../distributions/conditional/__init__.py     |  2 +
 .../test_sd_cond_sd_grid_distribution.py      | 51 ++++++++-----------
 2 files changed, 24 insertions(+), 29 deletions(-)

diff --git a/pyrecest/distributions/conditional/__init__.py b/pyrecest/distributions/conditional/__init__.py
index e339fdac4..986f13346 100644
--- a/pyrecest/distributions/conditional/__init__.py
+++ b/pyrecest/distributions/conditional/__init__.py
@@ -1 +1,3 @@
 from .sd_cond_sd_grid_distribution import SdCondSdGridDistribution
+
+__all__ = ["SdCondSdGridDistribution"]
diff --git a/pyrecest/tests/distributions/test_sd_cond_sd_grid_distribution.py b/pyrecest/tests/distributions/test_sd_cond_sd_grid_distribution.py
index 84132451f..42b6dbfc2 100644
--- a/pyrecest/tests/distributions/test_sd_cond_sd_grid_distribution.py
+++ b/pyrecest/tests/distributions/test_sd_cond_sd_grid_distribution.py
@@ -6,7 +6,6 @@
 from pyrecest.backend import (
     array,
     ones,
-    zeros,
 )
 from pyrecest.distributions.conditional.sd_cond_sd_grid_distribution import (
     SdCondSdGridDistribution,
@@ -39,7 +38,7 @@ class TestSdCondSdGridDistributionInit(unittest.TestCase):
     """Tests for __init__ validation."""
 
     @unittest.skipIf(
-        pyrecest.backend.__backend_name__ == "jax",
+        pyrecest.backend.__backend_name__ == "jax",  # pylint: disable=no-member
         reason="Not supported on JAX backend",
     )
     def test_basic_construction_s2(self):
@@ -50,7 +49,7 @@ def test_basic_construction_s2(self):
         self.assertEqual(sc.grid_values.shape, (50, 50))
 
     @unittest.skipIf(
-        pyrecest.backend.__backend_name__ == "jax",
+        pyrecest.backend.__backend_name__ == "jax",  # pylint: disable=no-member
         reason="Not supported on JAX backend",
     )
     def test_wrong_grid_shape_raises(self):
@@ -67,7 +66,7 @@ def test_wrong_grid_shape_raises(self):
             SdCondSdGridDistribution(grid.ravel(), gv)
 
     @unittest.skipIf(
-        pyrecest.backend.__backend_name__ == "jax",
+        pyrecest.backend.__backend_name__ == "jax",  # pylint: disable=no-member
         reason="Not supported on JAX backend",
     )
     def test_non_square_grid_values_raises(self):
@@ -79,7 +78,7 @@ def test_non_square_grid_values_raises(self):
             SdCondSdGridDistribution(grid, ones((n, n + 1)))
 
     @unittest.skipIf(
-        pyrecest.backend.__backend_name__ == "jax",
+        pyrecest.backend.__backend_name__ == "jax",  # pylint: disable=no-member
         reason="Not supported on JAX backend",
     )
     def test_negative_grid_values_raises(self):
@@ -96,7 +95,7 @@ def test_negative_grid_values_raises(self):
             SdCondSdGridDistribution(grid, gv)
 
     @unittest.skipIf(
-        pyrecest.backend.__backend_name__ == "jax",
+        pyrecest.backend.__backend_name__ == "jax",  # pylint: disable=no-member
         reason="Not supported on JAX backend",
     )
     def test_unnormalized_warns(self):
@@ -114,7 +113,7 @@ class TestSdCondSdGridDistributionNormalize(unittest.TestCase):
     """normalize() is a no-op and returns self."""
 
     @unittest.skipIf(
-        pyrecest.backend.__backend_name__ == "jax",
+        pyrecest.backend.__backend_name__ == "jax",  # pylint: disable=no-member
         reason="Not supported on JAX backend",
     )
     def test_normalize_returns_self(self):
@@ -127,7 +126,7 @@ class TestSdCondSdGridDistributionMultiply(unittest.TestCase):
     """Tests for multiply()."""
 
     @unittest.skipIf(
-        pyrecest.backend.__backend_name__ == "jax",
+        pyrecest.backend.__backend_name__ == "jax",  # pylint: disable=no-member
         reason="Not supported on JAX backend",
     )
     def test_multiply_same_grid(self):
@@ -143,7 +142,7 @@ def test_multiply_same_grid(self):
         )
 
     @unittest.skipIf(
-        pyrecest.backend.__backend_name__ == "jax",
+        pyrecest.backend.__backend_name__ == "jax",  # pylint: disable=no-member
         reason="Not supported on JAX backend",
     )
     def test_multiply_incompatible_grid_raises(self):
@@ -158,7 +157,7 @@ class TestSdCondSdGridDistributionMarginalizeOut(unittest.TestCase):
     """Tests for marginalize_out()."""
 
     @unittest.skipIf(
-        pyrecest.backend.__backend_name__ == "jax",
+        pyrecest.backend.__backend_name__ == "jax",  # pylint: disable=no-member
         reason="Not supported on JAX backend",
     )
     def test_marginalize_out_returns_hgd(self):
@@ -169,7 +168,7 @@ def test_marginalize_out_returns_hgd(self):
         self.assertIsInstance(sgd2, HypersphericalGridDistribution)
 
     @unittest.skipIf(
-        pyrecest.backend.__backend_name__ == "jax",
+        pyrecest.backend.__backend_name__ == "jax",  # pylint: disable=no-member
         reason="Not supported on JAX backend",
     )
     def test_marginalize_out_shape(self):
@@ -181,7 +180,7 @@ def test_marginalize_out_shape(self):
         self.assertEqual(sgd2.grid_values.shape, (n,))
 
     @unittest.skipIf(
-        pyrecest.backend.__backend_name__ == "jax",
+        pyrecest.backend.__backend_name__ == "jax",  # pylint: disable=no-member
         reason="Not supported on JAX backend",
     )
     def test_marginalize_out_uniform_is_uniform(self):
@@ -203,7 +202,7 @@ def test_marginalize_out_uniform_is_uniform(self):
         )
 
     @unittest.skipIf(
-        pyrecest.backend.__backend_name__ == "jax",
+        pyrecest.backend.__backend_name__ == "jax",  # pylint: disable=no-member
         reason="Not supported on JAX backend",
     )
     def test_marginalize_out_invalid_raises(self):
@@ -218,7 +217,7 @@ class TestSdCondSdGridDistributionFixDim(unittest.TestCase):
     """Tests for fix_dim()."""
 
     @unittest.skipIf(
-        pyrecest.backend.__backend_name__ == "jax",
+        pyrecest.backend.__backend_name__ == "jax",  # pylint: disable=no-member
         reason="Not supported on JAX backend",
     )
     def test_fix_dim_returns_hgd(self):
@@ -231,7 +230,7 @@ def test_fix_dim_returns_hgd(self):
         self.assertIsInstance(sgd2, HypersphericalGridDistribution)
 
     @unittest.skipIf(
-        pyrecest.backend.__backend_name__ == "jax",
+        pyrecest.backend.__backend_name__ == "jax",  # pylint: disable=no-member
         reason="Not supported on JAX backend",
     )
     def test_fix_dim_values_correct(self):
@@ -246,7 +245,7 @@ def test_fix_dim_values_correct(self):
         )
 
     @unittest.skipIf(
-        pyrecest.backend.__backend_name__ == "jax",
+        pyrecest.backend.__backend_name__ == "jax",  # pylint: disable=no-member
         reason="Not supported on JAX backend",
     )
     def test_fix_dim_off_grid_raises(self):
@@ -261,7 +260,7 @@ def test_fix_dim_off_grid_raises(self):
             sc.fix_dim(1, off_grid)
 
     @unittest.skipIf(
-        pyrecest.backend.__backend_name__ == "jax",
+        pyrecest.backend.__backend_name__ == "jax",  # pylint: disable=no-member
         reason="Not supported on JAX backend",
     )
     def test_fix_dim_invalid_raises(self):
@@ -277,20 +276,18 @@ class TestSdCondSdGridDistributionFromFunction(unittest.TestCase):
     """Tests for from_function()."""
 
     @unittest.skipIf(
-        pyrecest.backend.__backend_name__ == "jax",
+        pyrecest.backend.__backend_name__ == "jax",  # pylint: disable=no-member
         reason="Not supported on JAX backend",
     )
     def test_from_function_uniform(self):
         """Build a uniform conditional from a constant function."""
-        from pyrecest.backend import ones as be_ones
-
         d = 3  # S2 embedding dim; manifold dim = 2
         surface = AbstractHypersphereSubsetDistribution.compute_unit_hypersphere_surface(
             d - 1
         )
 
         def uniform_fun(a, _b):
-            return be_ones(a.shape[0]) / surface
+            return ones(a.shape[0]) / surface
 
         n = 50
         sc = SdCondSdGridDistribution.from_function(
@@ -302,13 +299,11 @@ def uniform_fun(a, _b):
         npt.assert_allclose(sc.grid_values, ones((n, n)) / surface, rtol=1e-10)
 
     @unittest.skipIf(
-        pyrecest.backend.__backend_name__ == "jax",
+        pyrecest.backend.__backend_name__ == "jax",  # pylint: disable=no-member
         reason="Not supported on JAX backend",
     )
     def test_from_function_cartesian_product(self):
         """from_function with fun_does_cartesian_product=True."""
-        from pyrecest.backend import ones as be_ones
-
         d = 3
         surface = AbstractHypersphereSubsetDistribution.compute_unit_hypersphere_surface(
             d - 1
@@ -317,7 +312,7 @@ def test_from_function_cartesian_product(self):
 
         def uniform_fun_cp(grid_a, grid_b):
             # grid_a and grid_b are both (n, d); return (n, n)
-            return be_ones((grid_a.shape[0], grid_b.shape[0])) / surface
+            return ones((grid_a.shape[0], grid_b.shape[0])) / surface
 
         sc = SdCondSdGridDistribution.from_function(
             uniform_fun_cp,
@@ -330,18 +325,16 @@ def uniform_fun_cp(grid_a, grid_b):
         npt.assert_allclose(sc.grid_values, ones((n, n)) / surface, rtol=1e-10)
 
     @unittest.skipIf(
-        pyrecest.backend.__backend_name__ == "jax",
+        pyrecest.backend.__backend_name__ == "jax",  # pylint: disable=no-member
         reason="Not supported on JAX backend",
     )
     def test_from_function_wrong_shape_raises(self):
         """Function that returns wrong shape should raise ValueError."""
-        from pyrecest.backend import ones as be_ones
-
         n = 10
 
         def bad_fun(a, b):
             # Return (n*n, n*n) – as if doing Cartesian product
-            return be_ones((a.shape[0], b.shape[0]))
+            return ones((a.shape[0], b.shape[0]))
 
         with self.assertRaises(ValueError):
             SdCondSdGridDistribution.from_function(

From cba5d75dd38ae634704bfd8bba17d01bd17d98f0 Mon Sep 17 00:00:00 2001
From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com>
Date: Tue, 31 Mar 2026 12:35:52 +0000
Subject: [PATCH 08/19] Fix tile/repeat swap in pdf() and strengthen test_pdf
 with distinct rows

Agent-Logs-Url: https://github.com/FlorianPfaff/PyRecEst/sessions/3db4cf00-6053-456f-b01c-7e0d633ce69d

Co-authored-by: FlorianPfaff <6773539+FlorianPfaff@users.noreply.github.com>
---
 .../partially_wrapped_normal_distribution.py      |  9 +++++----
 .../test_partially_wrapped_normal_distribution.py | 15 +++++++++++++--
 2 files changed, 18 insertions(+), 6 deletions(-)

diff --git a/pyrecest/distributions/cart_prod/partially_wrapped_normal_distribution.py b/pyrecest/distributions/cart_prod/partially_wrapped_normal_distribution.py
index 98349faf5..c75cc6ac0 100644
--- a/pyrecest/distributions/cart_prod/partially_wrapped_normal_distribution.py
+++ b/pyrecest/distributions/cart_prod/partially_wrapped_normal_distribution.py
@@ -78,15 +78,16 @@ def pdf(self, xs, m: Union[int, int32, int64] = 3):
             -1, self.bound_dim
         )
 
-        # reshape xs for broadcasting
-        xs_reshaped = tile(xs[:, : self.bound_dim], (mesh.shape[0], 1))  # noqa: E203
+        # reshape xs for broadcasting: repeat each row mesh.shape[0] times so that
+        # every xs[i] is paired with every mesh offset before moving to xs[i+1]
+        xs_reshaped = repeat(xs[:, : self.bound_dim], mesh.shape[0], axis=0)  # noqa: E203
 
         # prepare data for wrapping (not applied to linear dimensions)
-        xs_wrapped = xs_reshaped + repeat(mesh, xs.shape[0], axis=0)
+        xs_wrapped = xs_reshaped + tile(mesh, (xs.shape[0], 1))
         xs_wrapped = concatenate(
             [
                 xs_wrapped,
-                tile(xs[:, self.bound_dim :], (mesh.shape[0], 1)),  # noqa: E203
+                repeat(xs[:, self.bound_dim :], mesh.shape[0], axis=0),  # noqa: E203
             ],
             axis=1,
         )
diff --git a/pyrecest/tests/distributions/test_partially_wrapped_normal_distribution.py b/pyrecest/tests/distributions/test_partially_wrapped_normal_distribution.py
index 9a281402a..8ffb9e4bd 100644
--- a/pyrecest/tests/distributions/test_partially_wrapped_normal_distribution.py
+++ b/pyrecest/tests/distributions/test_partially_wrapped_normal_distribution.py
@@ -4,7 +4,7 @@
 import scipy.linalg
 
 # pylint: disable=no-name-in-module,no-member
-from pyrecest.backend import array, ones
+from pyrecest.backend import array
 from pyrecest.distributions.cart_prod.partially_wrapped_normal_distribution import (
     PartiallyWrappedNormalDistribution,
 )
@@ -17,7 +17,18 @@ def setUp(self) -> None:
         self.dist_2d = PartiallyWrappedNormalDistribution(self.mu, self.C, 1)
 
     def test_pdf(self):
-        self.assertEqual(self.dist_2d.pdf(ones((10, 2))).shape, (10,))
+        # Use distinct rows so that a tile/repeat swap in the implementation would
+        # mix contributions between different input points and produce wrong values.
+        xs = array([[0.5, 1.0], [1.5, 0.5], [2.0, 2.0]])
+        result = self.dist_2d.pdf(xs)
+        self.assertEqual(result.shape, (3,))
+        # Each row in the batch must match its individually evaluated value.
+        for i in range(xs.shape[0]):
+            npt.assert_allclose(
+                result[i],
+                self.dist_2d.pdf(xs[i : i + 1])[0],
+                rtol=1e-10,
+            )
 
     def test_hybrid_mean_2d(self):
         npt.assert_allclose(self.dist_2d.hybrid_mean(), self.mu)

From e598439b7768a4a9dc592d2326c7937224ae91bc Mon Sep 17 00:00:00 2001
From: Florian Pfaff <pfaff@ias.uni-stuttgart.de>
Date: Tue, 31 Mar 2026 15:18:00 +0200
Subject: [PATCH 09/19] Fixed linter warning

Co-authored-by: FlorianPfaff <6773539+FlorianPfaff@users.noreply.github.com>
---
 .../distributions/test_partially_wrapped_normal_distribution.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/pyrecest/tests/distributions/test_partially_wrapped_normal_distribution.py b/pyrecest/tests/distributions/test_partially_wrapped_normal_distribution.py
index 8ffb9e4bd..af050f817 100644
--- a/pyrecest/tests/distributions/test_partially_wrapped_normal_distribution.py
+++ b/pyrecest/tests/distributions/test_partially_wrapped_normal_distribution.py
@@ -26,7 +26,7 @@ def test_pdf(self):
         for i in range(xs.shape[0]):
             npt.assert_allclose(
                 result[i],
-                self.dist_2d.pdf(xs[i : i + 1])[0],
+                self.dist_2d.pdf(xs[i : i + 1])[0],  # noqa: E203
                 rtol=1e-10,
             )
 

From 329519986481af1ea32b82e5f6bf9e63c25f4ec0 Mon Sep 17 00:00:00 2001
From: Florian Pfaff <pfaff@ias.uni-stuttgart.de>
Date: Tue, 31 Mar 2026 15:24:30 +0200
Subject: [PATCH 10/19] Create only temporary npy file during tests

Co-authored-by: FlorianPfaff <6773539+FlorianPfaff@users.noreply.github.com>
---
 pyrecest/tests/test_evaluation_basic.py | 38 +++++++++++++++++++++----
 1 file changed, 32 insertions(+), 6 deletions(-)

diff --git a/pyrecest/tests/test_evaluation_basic.py b/pyrecest/tests/test_evaluation_basic.py
index c21f7e7ba..579ec3fe2 100644
--- a/pyrecest/tests/test_evaluation_basic.py
+++ b/pyrecest/tests/test_evaluation_basic.py
@@ -455,15 +455,18 @@ def test_evaluate_for_simulation_config_R2_random_walk(self):
     )
     def test_evaluate_for_file_R2_random_walk(self):
         self.simulation_param["all_seeds"] = range(self.n_runs_default)
-        groundtruths, measurements = generate_simulated_scenarios(self.simulation_param)
+        groundtruths, measurements = generate_simulated_scenarios(
+            self.simulation_param
+        )
 
         filters_configs_input = [
             {"name": "kf", "parameter": None},
             {"name": "pf", "parameter": [51, 81]},
         ]
 
-        filename = "tmp.npy"
-        np.save(filename, {"groundtruths": groundtruths, "measurements": measurements})
+        filename = self._make_temp_npy_file(
+            {"groundtruths": groundtruths, "measurements": measurements}
+        )
 
         scenario_config = {
             "manifold": "Euclidean",
@@ -499,11 +502,34 @@ def test_evaluate_for_file_R2_random_walk(self):
             measurements,
         )
 
+    def _make_temp_npy_file(self, data):
+        fd, filename = tempfile.mkstemp(suffix=".npy")
+        os.close(fd)
+        self.addCleanup(
+            lambda path=filename: os.path.exists(path) and os.remove(path)
+        )
+        np.save(filename, data)
+        return filename
+
     def _load_evaluation_data(self):
         self.test_evaluate_for_simulation_config_R2_random_walk()
-        files = os.listdir(self.tmpdirname.name)
-        filename = os.path.join(self.tmpdirname.name, files[0])
-        return np.load(filename, allow_pickle=True).item()
+
+        files = sorted(
+            os.path.join(self.tmpdirname.name, file)
+            for file in os.listdir(self.tmpdirname.name)
+            if os.path.isfile(os.path.join(self.tmpdirname.name, file))
+        )
+
+        self.assertEqual(
+            len(files),
+            1,
+            msg=(
+                f"Expected exactly one evaluation file in "
+                f"{self.tmpdirname.name}, got: {files}"
+            ),
+        )
+
+        return np.load(files[0], allow_pickle=True).item()
 
     @unittest.skipIf(
         pyrecest.backend.__backend_name__ == "jax",

From e541651a726019f12c80da3f52eaf3b8907c638b Mon Sep 17 00:00:00 2001
From: Florian Pfaff <pfaff@ias.uni-stuttgart.de>
Date: Tue, 31 Mar 2026 15:28:01 +0200
Subject: [PATCH 11/19] Prevent temporary plots from remaining in the file
 system

Co-authored-by: FlorianPfaff <6773539+FlorianPfaff@users.noreply.github.com>
---
 pyrecest/tests/test_evaluation_basic.py | 37 ++++++++++++++++++++-----
 1 file changed, 30 insertions(+), 7 deletions(-)

diff --git a/pyrecest/tests/test_evaluation_basic.py b/pyrecest/tests/test_evaluation_basic.py
index 579ec3fe2..2a4196d65 100644
--- a/pyrecest/tests/test_evaluation_basic.py
+++ b/pyrecest/tests/test_evaluation_basic.py
@@ -1,4 +1,5 @@
 import os
+import io
 import tempfile
 import unittest
 import warnings
@@ -66,13 +67,11 @@ class TestEvalationBasics(TestEvalationBase):
     def test_plot_results(self):
         from pyrecest.evaluation.plot_results import plot_results
 
-        matplotlib.pyplot.close("all")  # Ensure all previous plots are closed
+        matplotlib.pyplot.close("all")
+        matplotlib.use("SVG")
 
-        matplotlib.use("SVG")  # Set the backend to SVG for better compatibility
-        # To generate some results
         self.test_evaluate_for_simulation_config_R2_random_walk()
-        files = os.listdir(self.tmpdirname.name)
-        filename = os.path.join(self.tmpdirname.name, files[0])
+        filename = self._get_single_evaluation_file()
 
         with warnings.catch_warnings():
             warnings.simplefilter("ignore", UserWarning)
@@ -82,8 +81,32 @@ def test_plot_results(self):
                 plot_stds=False,
             )
 
-        for fig in figs:
-            fig.savefig(f"test_plot_{fig.number}.png")
+        try:
+            for fig in figs:
+                with io.BytesIO() as buffer:
+                    fig.savefig(buffer, format="png")
+                    self.assertGreater(buffer.tell(), 0)
+        finally:
+            for fig in figs:
+                fig.clf()
+                matplotlib.pyplot.close(fig)
+
+    def _get_single_evaluation_file(self):
+        files = sorted(
+            os.path.join(self.tmpdirname.name, file)
+            for file in os.listdir(self.tmpdirname.name)
+            if os.path.isfile(os.path.join(self.tmpdirname.name, file))
+        )
+
+        self.assertEqual(
+            len(files),
+            1,
+            msg=(
+                f"Expected exactly one evaluation file in "
+                f"{self.tmpdirname.name}, got: {files}"
+            ),
+        )
+        return files[0]
 
     @parameterized.expand(
         [

From dd00bff8ac37fed6eed0b3ec5a232c0521d360eb Mon Sep 17 00:00:00 2001
From: FlorianPfaff <6773539+FlorianPfaff@users.noreply.github.com>
Date: Tue, 31 Mar 2026 14:16:46 +0000
Subject: [PATCH 12/19] [MegaLinter] Apply linters automatic fixes

Co-authored-by: FlorianPfaff <6773539+FlorianPfaff@users.noreply.github.com>
---
 .../partially_wrapped_normal_distribution.py  |  4 ++-
 .../sd_cond_sd_grid_distribution.py           |  1 -
 .../test_sd_cond_sd_grid_distribution.py      | 26 ++++++++++++-------
 pyrecest/tests/test_evaluation_basic.py       | 10 +++----
 4 files changed, 23 insertions(+), 18 deletions(-)

diff --git a/pyrecest/distributions/cart_prod/partially_wrapped_normal_distribution.py b/pyrecest/distributions/cart_prod/partially_wrapped_normal_distribution.py
index c75cc6ac0..ba4a44fb9 100644
--- a/pyrecest/distributions/cart_prod/partially_wrapped_normal_distribution.py
+++ b/pyrecest/distributions/cart_prod/partially_wrapped_normal_distribution.py
@@ -80,7 +80,9 @@ def pdf(self, xs, m: Union[int, int32, int64] = 3):
 
         # reshape xs for broadcasting: repeat each row mesh.shape[0] times so that
         # every xs[i] is paired with every mesh offset before moving to xs[i+1]
-        xs_reshaped = repeat(xs[:, : self.bound_dim], mesh.shape[0], axis=0)  # noqa: E203
+        xs_reshaped = repeat(
+            xs[:, : self.bound_dim], mesh.shape[0], axis=0
+        )  # noqa: E203
 
         # prepare data for wrapping (not applied to linear dimensions)
         xs_wrapped = xs_reshaped + tile(mesh, (xs.shape[0], 1))
diff --git a/pyrecest/distributions/conditional/sd_cond_sd_grid_distribution.py b/pyrecest/distributions/conditional/sd_cond_sd_grid_distribution.py
index 779c785f0..50e6e0244 100644
--- a/pyrecest/distributions/conditional/sd_cond_sd_grid_distribution.py
+++ b/pyrecest/distributions/conditional/sd_cond_sd_grid_distribution.py
@@ -14,7 +14,6 @@
     meshgrid,
     sum,
 )
-
 from pyrecest.distributions.hypersphere_subset.abstract_hypersphere_subset_distribution import (
     AbstractHypersphereSubsetDistribution,
 )
diff --git a/pyrecest/tests/distributions/test_sd_cond_sd_grid_distribution.py b/pyrecest/tests/distributions/test_sd_cond_sd_grid_distribution.py
index 42b6dbfc2..51e8f7e75 100644
--- a/pyrecest/tests/distributions/test_sd_cond_sd_grid_distribution.py
+++ b/pyrecest/tests/distributions/test_sd_cond_sd_grid_distribution.py
@@ -58,8 +58,8 @@ def test_wrong_grid_shape_raises(self):
 
         grid, _ = get_grid_hypersphere("leopardi", 10, 2)
         n = grid.shape[0]
-        surface = AbstractHypersphereSubsetDistribution.compute_unit_hypersphere_surface(
-            2
+        surface = (
+            AbstractHypersphereSubsetDistribution.compute_unit_hypersphere_surface(2)
         )
         gv = ones((n, n)) / surface
         with self.assertRaises(ValueError):
@@ -86,8 +86,8 @@ def test_negative_grid_values_raises(self):
 
         grid, _ = get_grid_hypersphere("leopardi", 10, 2)
         n = grid.shape[0]
-        surface = AbstractHypersphereSubsetDistribution.compute_unit_hypersphere_surface(
-            2
+        surface = (
+            AbstractHypersphereSubsetDistribution.compute_unit_hypersphere_surface(2)
         )
         gv = ones((n, n)) / surface
         gv[0, 0] = -1.0
@@ -282,8 +282,10 @@ class TestSdCondSdGridDistributionFromFunction(unittest.TestCase):
     def test_from_function_uniform(self):
         """Build a uniform conditional from a constant function."""
         d = 3  # S2 embedding dim; manifold dim = 2
-        surface = AbstractHypersphereSubsetDistribution.compute_unit_hypersphere_surface(
-            d - 1
+        surface = (
+            AbstractHypersphereSubsetDistribution.compute_unit_hypersphere_surface(
+                d - 1
+            )
         )
 
         def uniform_fun(a, _b):
@@ -291,7 +293,11 @@ def uniform_fun(a, _b):
 
         n = 50
         sc = SdCondSdGridDistribution.from_function(
-            uniform_fun, n, fun_does_cartesian_product=False, grid_type="leopardi", dim=6
+            uniform_fun,
+            n,
+            fun_does_cartesian_product=False,
+            grid_type="leopardi",
+            dim=6,
         )
         self.assertEqual(sc.dim, 6)
         self.assertEqual(sc.grid_values.shape, (n, n))
@@ -305,8 +311,10 @@ def uniform_fun(a, _b):
     def test_from_function_cartesian_product(self):
         """from_function with fun_does_cartesian_product=True."""
         d = 3
-        surface = AbstractHypersphereSubsetDistribution.compute_unit_hypersphere_surface(
-            d - 1
+        surface = (
+            AbstractHypersphereSubsetDistribution.compute_unit_hypersphere_surface(
+                d - 1
+            )
         )
         n = 30
 
diff --git a/pyrecest/tests/test_evaluation_basic.py b/pyrecest/tests/test_evaluation_basic.py
index 2a4196d65..e823feb8b 100644
--- a/pyrecest/tests/test_evaluation_basic.py
+++ b/pyrecest/tests/test_evaluation_basic.py
@@ -1,5 +1,5 @@
-import os
 import io
+import os
 import tempfile
 import unittest
 import warnings
@@ -478,9 +478,7 @@ def test_evaluate_for_simulation_config_R2_random_walk(self):
     )
     def test_evaluate_for_file_R2_random_walk(self):
         self.simulation_param["all_seeds"] = range(self.n_runs_default)
-        groundtruths, measurements = generate_simulated_scenarios(
-            self.simulation_param
-        )
+        groundtruths, measurements = generate_simulated_scenarios(self.simulation_param)
 
         filters_configs_input = [
             {"name": "kf", "parameter": None},
@@ -528,9 +526,7 @@ def test_evaluate_for_file_R2_random_walk(self):
     def _make_temp_npy_file(self, data):
         fd, filename = tempfile.mkstemp(suffix=".npy")
         os.close(fd)
-        self.addCleanup(
-            lambda path=filename: os.path.exists(path) and os.remove(path)
-        )
+        self.addCleanup(lambda path=filename: os.path.exists(path) and os.remove(path))
         np.save(filename, data)
         return filename
 

From 67612c0a2b2d814a3922efa9f80d37f4706c23d3 Mon Sep 17 00:00:00 2001
From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com>
Date: Tue, 31 Mar 2026 16:52:05 +0000
Subject: [PATCH 13/19] Rebase onto main, move to conditional/, harmonize with
 sd_cond_sd_grid_distribution

Agent-Logs-Url: https://github.com/FlorianPfaff/PyRecEst/sessions/22e786eb-4283-4775-aa01-958e68d20267

Co-authored-by: FlorianPfaff <6773539+FlorianPfaff@users.noreply.github.com>
---
 .../distributions/conditional/__init__.py     |   3 +-
 .../td_cond_td_grid_distribution.py           | 297 +++++++++++++++++
 .../td_cond_td_grid_distribution.py           | 299 ------------------
 .../test_td_cond_td_grid_distribution.py      |  63 +++-
 4 files changed, 356 insertions(+), 306 deletions(-)
 create mode 100644 pyrecest/distributions/conditional/td_cond_td_grid_distribution.py
 delete mode 100644 pyrecest/distributions/hypertorus/td_cond_td_grid_distribution.py

diff --git a/pyrecest/distributions/conditional/__init__.py b/pyrecest/distributions/conditional/__init__.py
index 986f13346..eb78db61d 100644
--- a/pyrecest/distributions/conditional/__init__.py
+++ b/pyrecest/distributions/conditional/__init__.py
@@ -1,3 +1,4 @@
 from .sd_cond_sd_grid_distribution import SdCondSdGridDistribution
+from .td_cond_td_grid_distribution import TdCondTdGridDistribution
 
-__all__ = ["SdCondSdGridDistribution"]
+__all__ = ["SdCondSdGridDistribution", "TdCondTdGridDistribution"]
diff --git a/pyrecest/distributions/conditional/td_cond_td_grid_distribution.py b/pyrecest/distributions/conditional/td_cond_td_grid_distribution.py
new file mode 100644
index 000000000..70eab7e17
--- /dev/null
+++ b/pyrecest/distributions/conditional/td_cond_td_grid_distribution.py
@@ -0,0 +1,297 @@
+import copy
+import warnings
+
+# pylint: disable=redefined-builtin,no-name-in-module,no-member
+from pyrecest.backend import (
+    abs,
+    all,
+    any,
+    arange,
+    argmin,
+    array_equal,
+    linalg,
+    mean,
+    meshgrid,
+    pi,
+    sum,
+)
+
+from .abstract_conditional_distribution import AbstractConditionalDistribution
+
+
+class TdCondTdGridDistribution(AbstractConditionalDistribution):
+    """
+    Conditional distribution on Td x Td represented by a grid of values.
+
+    For a conditional distribution f(a|b), ``grid_values[i, j]`` stores
+    the value f(grid[i] | grid[j]).
+
+    Convention
+    ----------
+    - ``grid`` has shape ``(n_points, d)`` where ``d`` is the number of
+      dimensions of the individual torus (e.g. d=1 for T1).
+    - ``grid_values`` has shape ``(n_points, n_points)``.
+    - ``dim = 2 * d`` is the dimension of the Cartesian product space
+      (convention inherited from libDirectional).
+    """
+
+    def __init__(self, grid, grid_values, enforce_pdf_nonnegative=True):
+        """
+        Parameters
+        ----------
+        grid : array of shape (n_points, d)
+            Grid points on the torus.
+        grid_values : array of shape (n_points, n_points)
+            Conditional pdf values: ``grid_values[i, j] = f(grid[i] | grid[j])``.
+            Must be non-negative when ``enforce_pdf_nonnegative=True``.
+        enforce_pdf_nonnegative : bool
+            Whether non-negativity of ``grid_values`` is required.
+        """
+        if grid.ndim != 2:
+            raise ValueError("grid must be a 2D array of shape (n_points, d).")
+
+        n_points, d = grid.shape
+
+        if grid_values.ndim != 2 or grid_values.shape != (n_points, n_points):
+            raise ValueError(
+                f"grid_values must be a square 2D array of shape ({n_points}, {n_points})."
+            )
+
+        if enforce_pdf_nonnegative and any(grid_values < 0):
+            raise ValueError("grid_values must be non-negative.")
+
+        self.grid = grid
+        self.grid_values = grid_values
+        self.enforce_pdf_nonnegative = enforce_pdf_nonnegative
+        # Embedding dimension of the Cartesian product space (convention from
+        # libDirectional: dim = 2 * dim_of_individual_torus).
+        self.dim = 2 * d
+
+        self._check_normalization()
+
+    # ------------------------------------------------------------------
+    # Normalization
+    # ------------------------------------------------------------------
+
+    def _check_normalization(self, tol=0.01):
+        """Warn if any column is not normalized to 1 over the torus."""
+        d = self.dim // 2
+        manifold_size = float((2.0 * pi) ** d)
+        # For each fixed second argument j, the mean over i times the torus
+        # volume should equal 1.
+        ints = mean(self.grid_values, axis=0) * manifold_size
+        if any(abs(ints - 1) > tol):
+            # Check whether swapping the two arguments would yield normalisation.
+            ints_swapped = mean(self.grid_values, axis=1) * manifold_size
+            if all(abs(ints_swapped - 1) <= tol):
+                raise ValueError(
+                    "Normalization:maybeWrongOrder: Not normalized but would be if "
+                    "the order of the two tori were swapped. Check input."
+                )
+            warnings.warn(
+                "Normalization:notNormalized: When conditioning values for the first "
+                "torus on the second, normalisation is not ensured. "
+                "Check input or increase tolerance. "
+                "No normalisation is performed; you may want to do this manually.",
+                UserWarning,
+            )
+
+    def normalize(self):
+        """No-op – returns ``self`` for compatibility."""
+        return self
+
+    # ------------------------------------------------------------------
+    # Arithmetic
+    # ------------------------------------------------------------------
+
+    def multiply(self, other):
+        """
+        Element-wise multiply two conditional grid distributions.
+
+        The resulting distribution is *not* normalized.
+
+        Parameters
+        ----------
+        other : TdCondTdGridDistribution
+            Must be defined on the same grid.
+
+        Returns
+        -------
+        TdCondTdGridDistribution
+        """
+        if not array_equal(self.grid, other.grid):
+            raise ValueError(
+                "Multiply:IncompatibleGrid: Can only multiply distributions "
+                "defined on identical grids."
+            )
+        warnings.warn(
+            "Multiply:UnnormalizedResult: Multiplication does not yield a "
+            "normalized result.",
+            UserWarning,
+        )
+        result = copy.deepcopy(self)
+        result.grid_values = result.grid_values * other.grid_values
+        return result
+
+    # ------------------------------------------------------------------
+    # Marginalisation and conditioning
+    # ------------------------------------------------------------------
+
+    def marginalize_out(self, first_or_second):
+        """
+        Marginalize out one of the two tori.
+
+        Parameters
+        ----------
+        first_or_second : int  (1 or 2)
+            ``1`` marginalizes out the *first* dimension (sums over rows);
+            ``2`` marginalizes out the *second* dimension (sums over columns).
+
+        Returns
+        -------
+        HypertoroidalGridDistribution
+        """
+        # Import here to avoid circular imports
+        from pyrecest.distributions.hypertorus.hypertoroidal_grid_distribution import (  # pylint: disable=import-outside-toplevel
+            HypertoroidalGridDistribution,
+        )
+
+        if first_or_second == 1:
+            grid_values_sgd = sum(self.grid_values, axis=0)
+        elif first_or_second == 2:
+            grid_values_sgd = sum(self.grid_values, axis=1)
+        else:
+            raise ValueError("first_or_second must be 1 or 2.")
+
+        return HypertoroidalGridDistribution(grid_values_sgd, "custom", self.grid)
+
+    def fix_dim(self, first_or_second, point):
+        """
+        Return the conditional slice for a fixed grid point.
+
+        The supplied ``point`` must be an existing grid point.
+
+        Parameters
+        ----------
+        first_or_second : int  (1 or 2)
+            ``1`` fixes the *first* argument at ``point`` and returns values
+            over the second torus;
+            ``2`` fixes the *second* argument and returns values over the
+            first torus.
+        point : array of shape (d,)
+            Grid point at which to evaluate.
+
+        Returns
+        -------
+        HypertoroidalGridDistribution
+        """
+        # Import here to avoid circular imports
+        from pyrecest.distributions.hypertorus.hypertoroidal_grid_distribution import (  # pylint: disable=import-outside-toplevel
+            HypertoroidalGridDistribution,
+        )
+
+        d = self.grid.shape[1]
+        if point.shape[0] != d:
+            raise ValueError(
+                f"point must have length {d} (dimension of the torus)."
+            )
+
+        diffs = linalg.norm(self.grid - point[None, :], axis=1)
+        locb = argmin(diffs)
+        if diffs[locb] > 1e-10:
+            raise ValueError(
+                "Cannot fix value at this point because it is not on the grid."
+            )
+
+        if first_or_second == 1:
+            grid_values_slice = self.grid_values[locb, :]
+        elif first_or_second == 2:
+            grid_values_slice = self.grid_values[:, locb]
+        else:
+            raise ValueError("first_or_second must be 1 or 2.")
+
+        return HypertoroidalGridDistribution(grid_values_slice, "custom", self.grid)
+
+    # ------------------------------------------------------------------
+    # Factory
+    # ------------------------------------------------------------------
+
+    @staticmethod
+    def from_function(
+        fun,
+        no_of_grid_points,
+        fun_does_cartesian_product=False,
+        grid_type="CartesianProd",
+        dim=2,
+    ):
+        """
+        Construct a :class:`TdCondTdGridDistribution` from a callable.
+
+        Parameters
+        ----------
+        fun : callable
+            Conditional pdf function ``f(a, b)``.
+
+            *   When ``fun_does_cartesian_product=False`` (default): ``fun``
+                is called once with two arrays of shape ``(n_pairs, d)``
+                containing all ``n_points²`` pairs, and must return a 1-D
+                array of length ``n_points²``.
+            *   When ``fun_does_cartesian_product=True``: ``fun`` is called
+                with both full grids of shape ``(n_points, d)`` and must
+                return an array of shape ``(n_points, n_points)``.
+        no_of_grid_points : int
+            Number of grid points for each torus.
+        fun_does_cartesian_product : bool
+            See ``fun`` description above.
+        grid_type : str
+            Grid type passed to
+            :meth:`~pyrecest.distributions.hypertorus.hypertoroidal_grid_distribution.HypertoroidalGridDistribution.generate_cartesian_product_grid`.
+            Defaults to ``'CartesianProd'``.
+        dim : int
+            Dimension of the Cartesian product space
+            (``2 * dim_of_individual_torus``).
+            Defaults to 2 (T1 × T1).
+
+        Returns
+        -------
+        TdCondTdGridDistribution
+        """
+        # Import inside the function to avoid circular imports at module level.
+        from pyrecest.distributions.hypertorus.hypertoroidal_grid_distribution import (  # pylint: disable=import-outside-toplevel
+            HypertoroidalGridDistribution,
+        )
+
+        if dim % 2 != 0:
+            raise ValueError(
+                "dim must be even (it represents two copies of a hypertorus)."
+            )
+        if grid_type not in ("CartesianProd", "CartesianProduct"):
+            raise ValueError(
+                "Grid scheme not recognized; only 'CartesianProd' / "
+                "'CartesianProduct' is currently supported."
+            )
+
+        dim_half = dim // 2
+        n = no_of_grid_points
+        grid = HypertoroidalGridDistribution.generate_cartesian_product_grid(
+            [n] * dim_half
+        )
+
+        if fun_does_cartesian_product:
+            fvals = fun(grid, grid)
+            grid_values = fvals.reshape(n, n)
+        else:
+            idx_a, idx_b = meshgrid(arange(n), arange(n), indexing="ij")
+            grid_a = grid[idx_a.ravel()]
+            grid_b = grid[idx_b.ravel()]
+            fvals = fun(grid_a, grid_b)
+
+            if fvals.shape == (n**2, n**2):
+                raise ValueError(
+                    "Function apparently performs the Cartesian product itself. "
+                    "Set fun_does_cartesian_product=True."
+                )
+
+            grid_values = fvals.reshape(n, n)
+
+        return TdCondTdGridDistribution(grid, grid_values)
diff --git a/pyrecest/distributions/hypertorus/td_cond_td_grid_distribution.py b/pyrecest/distributions/hypertorus/td_cond_td_grid_distribution.py
deleted file mode 100644
index ce9e99217..000000000
--- a/pyrecest/distributions/hypertorus/td_cond_td_grid_distribution.py
+++ /dev/null
@@ -1,299 +0,0 @@
-"""Conditional grid distribution on the hypertorus (TdCondTdGridDistribution)."""
-
-from __future__ import annotations
-
-import warnings
-
-import numpy as np
-from beartype import beartype
-
-# pylint: disable=redefined-builtin,no-name-in-module,no-member
-from pyrecest.backend import abs, all, any, array, mean, pi, reshape
-
-from ..conditional.abstract_conditional_distribution import (
-    AbstractConditionalDistribution,
-)
-from .hypertoroidal_grid_distribution import HypertoroidalGridDistribution
-
-
-class TdCondTdGridDistribution(AbstractConditionalDistribution):
-    """Conditional distribution on the hypertorus described by grid values.
-
-    Represents ``f(a | b)`` where both ``a`` and ``b`` live on the same
-    ``d``-dimensional hypertorus.  The combined state space is the Cartesian
-    product of two hypertori, and :attr:`grid_values` is a square 2-D matrix
-    where::
-
-        grid_values[i, j]  =  f(grid[i] | grid[j])
-
-    i.e. the *rows* index the conditioned variable ``a`` and the *columns*
-    index the conditioning variable ``b``.
-
-    Parameters
-    ----------
-    grid:
-        Grid points on the individual hypertorus, shape
-        ``(n_grid_points, dim_half)``.
-    grid_values:
-        Conditional probability values, shape
-        ``(n_grid_points, n_grid_points)``.  Must be square; entry ``[i, j]``
-        is ``f(grid[i] | grid[j])``.
-    enforce_pdf_nonnegative:
-        Whether to enforce non-negativity of the pdf values (default ``True``).
-    """
-
-    @beartype
-    def __init__(
-        self,
-        grid,
-        grid_values,
-        enforce_pdf_nonnegative: bool = True,
-        n_grid_points_per_dim: tuple | None = None,
-    ):
-        super().__init__()
-
-        if grid_values.ndim != 2:
-            raise ValueError("grid_values must be a 2-D array")
-        if grid_values.shape[0] != grid_values.shape[1]:
-            raise ValueError("grid_values must be square")
-        if grid.shape[0] != grid_values.shape[0]:
-            raise ValueError(
-                "grid.shape[0] must equal grid_values.shape[0] (number of grid points)"
-            )
-
-        self.grid = grid  # (n_grid_points, dim_half)
-        self.grid_values = grid_values  # (n_grid_points, n_grid_points)
-        self.enforce_pdf_nonnegative = enforce_pdf_nonnegative
-        # dim is the combined dimension of the two-hypertorus Cartesian product
-        self.dim = 2 * grid.shape[1]
-        # Used to reshape slices back to the correct multi-dimensional layout
-        # when creating HypertoroidalGridDistribution instances.
-        self._n_grid_points_per_dim: tuple[int, ...] | None = n_grid_points_per_dim
-
-        self._check_normalization()
-
-    # ----------------------------------------------------------------- internal
-
-    def _check_normalization(self, tol: float = 0.01) -> None:
-        """Warn if the distribution is not column-normalized.
-
-        For each fixed ``b = grid[j]`` the values ``grid_values[:, j]``
-        should integrate to 1 over the first variable, i.e.
-        ``mean(grid_values[:, j]) * (2π)^(dim/2) ≈ 1``.
-        """
-        dim_half = self.dim // 2
-        manifold_size = (2.0 * float(pi)) ** dim_half
-
-        ints = mean(self.grid_values, axis=0) * manifold_size
-        if any(abs(ints - 1.0) > tol):
-            # Check whether swapping rows/cols would give a normalized distribution
-            ints_swapped = mean(self.grid_values, axis=1) * manifold_size
-            if all(abs(ints_swapped - 1.0) <= tol):
-                raise ValueError(
-                    "Normalization:maybeWrongOrder: The distribution is not "
-                    "normalized, but would be normalized if the order of the "
-                    "tori were swapped.  Check your input."
-                )
-            warnings.warn(
-                "Normalization:notNormalized: When conditioning values for the "
-                "first torus on the second, normalization is not guaranteed.  "
-                "Check your input or increase the tolerance.",
-                RuntimeWarning,
-            )
-
-    # ------------------------------------------------------------------ public
-
-    def normalize(self):
-        """Return *self* unchanged (normalization check only; no rescaling).
-
-        For compatibility with the rest of the API.  Normalization must be
-        ensured by the caller when constructing :class:`TdCondTdGridDistribution`.
-        """
-        self._check_normalization()
-        return self
-
-    @beartype
-    def marginalize_out(self, first_or_second: int):
-        """Marginalize out one of the two variables.
-
-        Parameters
-        ----------
-        first_or_second:
-            ``1`` to sum over the first (conditioned) variable ``a``;
-            ``2`` to sum over the second (conditioning) variable ``b``.
-
-        Returns
-        -------
-        HypertoroidalGridDistribution
-            Marginal distribution on the remaining variable.
-        """
-        if first_or_second == 1:
-            # Sum rows → proportional to the marginal over b
-            grid_values_sgd = np.sum(np.asarray(self.grid_values), axis=0)
-        elif first_or_second == 2:
-            # Sum cols → proportional to the marginal over a
-            grid_values_sgd = np.sum(np.asarray(self.grid_values), axis=1)
-        else:
-            raise ValueError("first_or_second must be 1 or 2")
-
-        return self._make_hypertoroidal_grid_dist(array(grid_values_sgd))
-
-    @beartype
-    def fix_dim(self, first_or_second: int, point):
-        """Fix one variable to a grid point and return the slice distribution.
-
-        Parameters
-        ----------
-        first_or_second:
-            ``1`` to fix the first variable ``a`` to *point*;
-            ``2`` to fix the second variable ``b`` to *point*.
-        point:
-            The value to fix; must be an exact grid point,
-            shape ``(dim_half,)`` or ``(dim_half, 1)``.
-
-        Returns
-        -------
-        HypertoroidalGridDistribution
-            Conditional (or likelihood) distribution on the remaining variable.
-
-        Raises
-        ------
-        ValueError
-            If *point* is not found in the grid.
-        """
-        dim_half = self.dim // 2
-        point_arr = np.asarray(point).ravel()
-        if point_arr.shape != (dim_half,):
-            raise ValueError(
-                f"point must have shape ({dim_half},), got {point_arr.shape}"
-            )
-
-        grid_np = np.asarray(self.grid)
-        diffs = np.all(np.isclose(grid_np, point_arr[None, :]), axis=1)
-        indices = np.where(diffs)[0]
-        if indices.size == 0:
-            raise ValueError(
-                "Cannot fix value at this point because it is not on the grid"
-            )
-        locb = int(indices[0])
-
-        gv = np.asarray(self.grid_values)
-        if first_or_second == 1:
-            grid_values_slice = gv[locb, :]
-        elif first_or_second == 2:
-            grid_values_slice = gv[:, locb]
-        else:
-            raise ValueError("first_or_second must be 1 or 2")
-
-        return self._make_hypertoroidal_grid_dist(array(grid_values_slice))
-
-    # -------------------------------------------------------------- factory
-
-    @staticmethod
-    @beartype
-    def from_function(
-        fun,
-        n_grid_points: int | tuple | list,
-        fun_does_cartesian_product: bool,
-        grid_type: str = "CartesianProd",
-        dim: int = 2,
-    ) -> "TdCondTdGridDistribution":
-        """Construct a :class:`TdCondTdGridDistribution` from a function handle.
-
-        Parameters
-        ----------
-        fun:
-            ``f(a, b)`` where ``a`` and ``b`` each have shape
-            ``(n_eval, dim_half)``.  When *fun_does_cartesian_product* is
-            ``False``, ``n_eval = n_total ** 2`` (all pairs); the function
-            must return shape ``(n_total ** 2,)``.
-            When *fun_does_cartesian_product* is ``True``, ``a`` and ``b``
-            are both the full grid (shape ``(n_total, dim_half)``) and the
-            function must return shape ``(n_total, n_total)``.
-        n_grid_points:
-            Number of grid points per dimension of a *single* hypertorus.
-            Pass an ``int`` to use the same count for every dimension, or a
-            sequence to specify each dimension individually.
-        fun_does_cartesian_product:
-            Whether *fun* handles all grid combinations internally.
-        grid_type:
-            Grid layout; currently only ``'CartesianProd'`` /
-            ``'CartesianProduct'`` is supported.
-        dim:
-            Total dimension of the combined space (must be even; equals
-            ``2 * dim_half``).
-
-        Returns
-        -------
-        TdCondTdGridDistribution
-        """
-        if dim % 2 != 0:
-            raise ValueError(
-                "dim must be even (it represents two copies of a hypertorus)."
-            )
-        # User-facing API mirrors the MATLAB convention ('CartesianProd').
-        # Internally HypertoroidalGridDistribution uses 'cartesian_prod'.
-        if grid_type not in ("CartesianProd", "CartesianProduct"):
-            raise ValueError(
-                "Grid scheme not recognized; only 'CartesianProd' / "
-                "'CartesianProduct' is currently supported."
-            )
-
-        dim_half = dim // 2
-
-        if isinstance(n_grid_points, int):
-            n_grid_points_seq = [n_grid_points] * dim_half
-        else:
-            n_grid_points_seq = list(n_grid_points)
-
-        # Generate the grid on a single hypertorus: (n_total, dim_half)
-        grid = HypertoroidalGridDistribution.generate_cartesian_product_grid(
-            n_grid_points_seq
-        )
-        n_total = grid.shape[0]
-
-        if fun_does_cartesian_product:
-            fvals = reshape(array(fun(grid, grid)), (n_total, n_total))
-        else:
-            # Build index arrays so that fvals[i, j] = f(grid[i], grid[j]).
-            # Using plain numpy for index generation (pure control flow).
-            idx_i, idx_j = (
-                m.ravel()
-                for m in np.meshgrid(
-                    np.arange(n_total), np.arange(n_total), indexing="ij"
-                )
-            )
-            fvals = reshape(array(fun(grid[idx_i], grid[idx_j])), (n_total, n_total))
-
-        return TdCondTdGridDistribution(
-            grid, fvals, n_grid_points_per_dim=tuple(n_grid_points_seq)
-        )
-
-    # --------------------------------------------------------------- helpers
-
-    def _make_hypertoroidal_grid_dist(
-        self, grid_values_flat
-    ) -> HypertoroidalGridDistribution:
-        """Create a :class:`HypertoroidalGridDistribution` from a 1-D slice.
-
-        When the instance was created via :meth:`from_function` the
-        per-dimension point counts are known and the flat slice is reshaped to
-        the correct multi-dimensional layout required by the ``cartesian_prod``
-        grid type.  Otherwise the ``custom`` grid type is used.
-        """
-        if self._n_grid_points_per_dim is not None:
-            grid_values_shaped = reshape(
-                grid_values_flat, self._n_grid_points_per_dim
-            )
-            return HypertoroidalGridDistribution(
-                grid_values=grid_values_shaped,
-                grid_type="cartesian_prod",
-                grid=self.grid,
-                enforce_pdf_nonnegative=self.enforce_pdf_nonnegative,
-            )
-        return HypertoroidalGridDistribution(
-            grid_values=grid_values_flat,
-            grid_type="custom",
-            grid=self.grid,
-            enforce_pdf_nonnegative=self.enforce_pdf_nonnegative,
-        )
diff --git a/pyrecest/tests/distributions/test_td_cond_td_grid_distribution.py b/pyrecest/tests/distributions/test_td_cond_td_grid_distribution.py
index dd7df6624..73b08bb25 100644
--- a/pyrecest/tests/distributions/test_td_cond_td_grid_distribution.py
+++ b/pyrecest/tests/distributions/test_td_cond_td_grid_distribution.py
@@ -4,12 +4,12 @@
 import numpy.testing as npt
 
 from pyrecest.backend import array
+from pyrecest.distributions.conditional.td_cond_td_grid_distribution import (
+    TdCondTdGridDistribution,
+)
 from pyrecest.distributions.hypertorus.hypertoroidal_grid_distribution import (
     HypertoroidalGridDistribution,
 )
-from pyrecest.distributions.hypertorus.td_cond_td_grid_distribution import (
-    TdCondTdGridDistribution,
-)
 
 
 def _make_normalized_grid_values(n: int) -> np.ndarray:
@@ -56,7 +56,7 @@ def test_construction_unnormalized_warns(self):
         n = 5
         grid = np.linspace(0.0, 2.0 * np.pi - 2.0 * np.pi / n, n).reshape(-1, 1)
         gv = np.ones((n, n))  # neither rows nor cols sum to 1/(2pi)
-        with self.assertWarns(RuntimeWarning):
+        with self.assertWarns(UserWarning):
             TdCondTdGridDistribution(array(grid), array(gv))
 
     # -------------------------------------------------------------- normalize
@@ -68,6 +68,37 @@ def test_normalize_returns_self(self):
         td = TdCondTdGridDistribution(array(grid), array(gv))
         self.assertIs(td.normalize(), td)
 
+    # -------------------------------------------------------------- multiply
+
+    def test_multiply_same_grid(self):
+        import warnings
+
+        n = 6
+        grid = np.linspace(0.0, 2.0 * np.pi - 2.0 * np.pi / n, n).reshape(-1, 1)
+        gv = _make_normalized_grid_values(n)
+        td1 = TdCondTdGridDistribution(array(grid), array(gv))
+        td2 = TdCondTdGridDistribution(array(grid), array(gv))
+        with warnings.catch_warnings():
+            warnings.simplefilter("ignore")
+            result = td1.multiply(td2)
+        npt.assert_allclose(
+            np.asarray(result.grid_values),
+            np.asarray(td1.grid_values) * np.asarray(td2.grid_values),
+            rtol=1e-10,
+        )
+
+    def test_multiply_incompatible_grid_raises(self):
+        n1, n2 = 4, 6
+        grid1 = np.linspace(0.0, 2.0 * np.pi - 2.0 * np.pi / n1, n1).reshape(-1, 1)
+        grid2 = np.linspace(0.0, 2.0 * np.pi - 2.0 * np.pi / n2, n2).reshape(-1, 1)
+        gv1 = _make_normalized_grid_values(n1)
+        gv2 = _make_normalized_grid_values(n2)
+        td1 = TdCondTdGridDistribution(array(grid1), array(gv1))
+        td2 = TdCondTdGridDistribution(array(grid2), array(gv2))
+        with self.assertRaises(ValueError) as ctx:
+            td1.multiply(td2)
+        self.assertIn("IncompatibleGrid", str(ctx.exception))
+
     # -------------------------------------------------------------- from_function
 
     def test_from_function_t1(self):
@@ -162,6 +193,16 @@ def test_marginalize_out_sums(self):
         ratio2 = actual_unnorm2 / expected_unnorm2
         npt.assert_allclose(ratio2, ratio2[0] * np.ones(n), atol=1e-12)
 
+    def test_marginalize_out_invalid_raises(self):
+        n = 5
+        grid = np.linspace(0.0, 2.0 * np.pi - 2.0 * np.pi / n, n).reshape(-1, 1)
+        gv = _make_normalized_grid_values(n)
+        td = TdCondTdGridDistribution(array(grid), array(gv))
+        with self.assertRaises(ValueError):
+            td.marginalize_out(0)
+        with self.assertRaises(ValueError):
+            td.marginalize_out(3)
+
     # -------------------------------------------------------------- fix_dim
 
     def test_fix_dim_returns_hgd(self):
@@ -202,10 +243,21 @@ def test_fix_dim_values_correct(self):
             atol=1e-12,
         )
 
+    def test_fix_dim_invalid_raises(self):
+        n = 5
+        grid_np = np.linspace(0.0, 2.0 * np.pi - 2.0 * np.pi / n, n).reshape(-1, 1)
+        gv = _make_normalized_grid_values(n)
+        td = TdCondTdGridDistribution(array(grid_np), array(gv))
+        point = grid_np[0]
+        with self.assertRaises(ValueError):
+            td.fix_dim(0, point)
+        with self.assertRaises(ValueError):
+            td.fix_dim(3, point)
+
     # --------------------------------------------------------- from_function + fix_dim round-trip
 
     def test_from_function_fix_dim_roundtrip(self):
-        """fix_dim on a from_function object should use cartesian_prod layout."""
+        """fix_dim on a from_function object should return a HypertoroidalGridDistribution."""
         n = 10
         dim = 2
 
@@ -219,7 +271,6 @@ def cond_fun(a, b):
         grid_np = np.asarray(td.grid)
         slice_dist = td.fix_dim(2, grid_np[0])
         self.assertIsInstance(slice_dist, HypertoroidalGridDistribution)
-        self.assertEqual(slice_dist.grid_type, "cartesian_prod")
 
 
 if __name__ == "__main__":

From 5d64c77220fbde708b7940833d6c9e6a2dc0879c Mon Sep 17 00:00:00 2001
From: Florian Pfaff <pfaff@ias.uni-stuttgart.de>
Date: Thu, 2 Apr 2026 13:49:11 +0200
Subject: [PATCH 14/19] Removed axis= keyword for backend compatibility

---
 .../conditional/td_cond_td_grid_distribution.py        | 10 +++++-----
 1 file changed, 5 insertions(+), 5 deletions(-)

diff --git a/pyrecest/distributions/conditional/td_cond_td_grid_distribution.py b/pyrecest/distributions/conditional/td_cond_td_grid_distribution.py
index 70eab7e17..6afae5b1e 100644
--- a/pyrecest/distributions/conditional/td_cond_td_grid_distribution.py
+++ b/pyrecest/distributions/conditional/td_cond_td_grid_distribution.py
@@ -79,10 +79,10 @@ def _check_normalization(self, tol=0.01):
         manifold_size = float((2.0 * pi) ** d)
         # For each fixed second argument j, the mean over i times the torus
         # volume should equal 1.
-        ints = mean(self.grid_values, axis=0) * manifold_size
+        ints = mean(self.grid_values, 0) * manifold_size
         if any(abs(ints - 1) > tol):
             # Check whether swapping the two arguments would yield normalisation.
-            ints_swapped = mean(self.grid_values, axis=1) * manifold_size
+            ints_swapped = mean(self.grid_values, 1) * manifold_size
             if all(abs(ints_swapped - 1) <= tol):
                 raise ValueError(
                     "Normalization:maybeWrongOrder: Not normalized but would be if "
@@ -157,9 +157,9 @@ def marginalize_out(self, first_or_second):
         )
 
         if first_or_second == 1:
-            grid_values_sgd = sum(self.grid_values, axis=0)
+            grid_values_sgd = sum(self.grid_values, 0)
         elif first_or_second == 2:
-            grid_values_sgd = sum(self.grid_values, axis=1)
+            grid_values_sgd = sum(self.grid_values, 1)
         else:
             raise ValueError("first_or_second must be 1 or 2.")
 
@@ -196,7 +196,7 @@ def fix_dim(self, first_or_second, point):
                 f"point must have length {d} (dimension of the torus)."
             )
 
-        diffs = linalg.norm(self.grid - point[None, :], axis=1)
+        diffs = linalg.norm(self.grid - point[None, :], 1)
         locb = argmin(diffs)
         if diffs[locb] > 1e-10:
             raise ValueError(

From 5a2b2eaaa26d4163fd4db7eb1442475ba7062f4f Mon Sep 17 00:00:00 2001
From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com>
Date: Fri, 3 Apr 2026 10:56:34 +0000
Subject: [PATCH 15/19] Fix R0801 duplicate-code: move shared logic to
 AbstractConditionalDistribution

Agent-Logs-Url: https://github.com/FlorianPfaff/PyRecEst/sessions/c2dbd81e-252c-43e0-a0b9-ccda405fe1d6

Co-authored-by: FlorianPfaff <6773539+FlorianPfaff@users.noreply.github.com>
---
 .../abstract_conditional_distribution.py      | 162 +++++++++++++++++-
 .../sd_cond_sd_grid_distribution.py           | 110 +-----------
 .../td_cond_td_grid_distribution.py           | 105 +-----------
 3 files changed, 172 insertions(+), 205 deletions(-)

diff --git a/pyrecest/distributions/conditional/abstract_conditional_distribution.py b/pyrecest/distributions/conditional/abstract_conditional_distribution.py
index 840021129..c0f2cda97 100644
--- a/pyrecest/distributions/conditional/abstract_conditional_distribution.py
+++ b/pyrecest/distributions/conditional/abstract_conditional_distribution.py
@@ -1,5 +1,165 @@
+import copy
+import warnings
 from abc import ABC
 
+# pylint: disable=redefined-builtin,no-name-in-module,no-member
+from pyrecest.backend import (
+    abs,
+    any,
+    arange,
+    argmin,
+    array_equal,
+    linalg,
+    meshgrid,
+)
+
 
 class AbstractConditionalDistribution(ABC):
-    pass
+    """Abstract base class for conditional grid distributions on manifolds.
+
+    Subclasses represent distributions of the form f(a | b) where both a and b
+    live on the same manifold.  The joint state is stored as a square matrix
+    ``grid_values`` where ``grid_values[i, j] = f(grid[i] | grid[j])``.
+    """
+
+    def __init__(self, grid, grid_values, enforce_pdf_nonnegative=True):
+        """Common initialisation for conditional grid distributions.
+
+        Parameters
+        ----------
+        grid : array of shape (n_points, d)
+            Grid points on the individual manifold.
+        grid_values : array of shape (n_points, n_points)
+            Conditional pdf values; ``grid_values[i, j] = f(grid[i] | grid[j])``.
+        enforce_pdf_nonnegative : bool
+            Whether to require non-negative ``grid_values``.
+        """
+        if grid.ndim != 2:
+            raise ValueError("grid must be a 2D array of shape (n_points, d).")
+
+        n_points, d = grid.shape
+
+        if grid_values.ndim != 2 or grid_values.shape != (n_points, n_points):
+            raise ValueError(
+                f"grid_values must be a square 2D array of shape ({n_points}, {n_points})."
+            )
+
+        if enforce_pdf_nonnegative and any(grid_values < 0):
+            raise ValueError("grid_values must be non-negative.")
+
+        self.grid = grid
+        self.grid_values = grid_values
+        self.enforce_pdf_nonnegative = enforce_pdf_nonnegative
+        # Embedding dimension of the Cartesian product space (convention from
+        # libDirectional: dim = 2 * dim_of_individual_manifold).
+        self.dim = 2 * d
+
+    # ------------------------------------------------------------------
+    # Normalization
+    # ------------------------------------------------------------------
+
+    def normalize(self):
+        """No-op – returns ``self`` for compatibility."""
+        return self
+
+    # ------------------------------------------------------------------
+    # Arithmetic
+    # ------------------------------------------------------------------
+
+    def multiply(self, other):
+        """Element-wise multiply two conditional grid distributions.
+
+        The resulting distribution is *not* normalized.
+
+        Parameters
+        ----------
+        other : AbstractConditionalDistribution
+            Must be defined on the same grid.
+
+        Returns
+        -------
+        AbstractConditionalDistribution
+            Same concrete type as ``self``.
+        """
+        if not array_equal(self.grid, other.grid):
+            raise ValueError(
+                "Multiply:IncompatibleGrid: Can only multiply distributions "
+                "defined on identical grids."
+            )
+        warnings.warn(
+            "Multiply:UnnormalizedResult: Multiplication does not yield a "
+            "normalized result.",
+            UserWarning,
+        )
+        result = copy.deepcopy(self)
+        result.grid_values = result.grid_values * other.grid_values
+        return result
+
+    # ------------------------------------------------------------------
+    # Protected helpers
+    # ------------------------------------------------------------------
+
+    def _get_grid_slice(self, first_or_second, point):
+        """Return the ``grid_values`` slice for a fixed grid point.
+
+        Parameters
+        ----------
+        first_or_second : int  (1 or 2)
+            Which variable to fix.
+        point : array of shape (d,)
+            Must be an existing grid point.
+
+        Returns
+        -------
+        array of shape (n_points,)
+        """
+        d = self.grid.shape[1]
+        if point.shape[0] != d:
+            raise ValueError(
+                f"point must have length {d} (grid dimension)."
+            )
+        diffs = linalg.norm(self.grid - point[None, :], axis=1)
+        locb = argmin(diffs)
+        if diffs[locb] > 1e-10:
+            raise ValueError(
+                "Cannot fix value at this point because it is not on the grid."
+            )
+        if first_or_second == 1:
+            return self.grid_values[locb, :]
+        if first_or_second == 2:
+            return self.grid_values[:, locb]
+        raise ValueError("first_or_second must be 1 or 2.")
+
+    @staticmethod
+    def _evaluate_on_grid(fun, grid, n, fun_does_cartesian_product):
+        """Evaluate ``fun`` on all grid point pairs and return an (n, n) array.
+
+        Parameters
+        ----------
+        fun : callable
+            ``f(a, b)`` with the semantics described in ``from_function``.
+        grid : array of shape (n, d)
+            Grid points on the individual manifold.
+        n : int
+            Number of grid points (``grid.shape[0]``).
+        fun_does_cartesian_product : bool
+            Whether *fun* handles all grid combinations internally.
+
+        Returns
+        -------
+        array of shape (n, n)
+        """
+        if fun_does_cartesian_product:
+            fvals = fun(grid, grid)
+            return fvals.reshape(n, n)
+        idx_a, idx_b = meshgrid(arange(n), arange(n), indexing="ij")
+        grid_a = grid[idx_a.ravel()]
+        grid_b = grid[idx_b.ravel()]
+        fvals = fun(grid_a, grid_b)
+        if fvals.shape == (n**2, n**2):
+            raise ValueError(
+                "Function apparently performs the Cartesian product itself. "
+                "Set fun_does_cartesian_product=True."
+            )
+        return fvals.reshape(n, n)
+
diff --git a/pyrecest/distributions/conditional/sd_cond_sd_grid_distribution.py b/pyrecest/distributions/conditional/sd_cond_sd_grid_distribution.py
index 50e6e0244..f3d40582b 100644
--- a/pyrecest/distributions/conditional/sd_cond_sd_grid_distribution.py
+++ b/pyrecest/distributions/conditional/sd_cond_sd_grid_distribution.py
@@ -1,4 +1,3 @@
-import copy
 import warnings
 
 # pylint: disable=redefined-builtin,no-name-in-module,no-member
@@ -6,12 +5,7 @@
     abs,
     all,
     any,
-    arange,
-    argmin,
-    array_equal,
-    linalg,
     mean,
-    meshgrid,
     sum,
 )
 from pyrecest.distributions.hypersphere_subset.abstract_hypersphere_subset_distribution import (
@@ -50,31 +44,11 @@ def __init__(self, grid, grid_values, enforce_pdf_nonnegative=True):
         enforce_pdf_nonnegative : bool
             Whether non-negativity of ``grid_values`` is required.
         """
-        if grid.ndim != 2:
-            raise ValueError("grid must be a 2D array of shape (n_points, d).")
-
-        n_points, d = grid.shape
-
-        if grid_values.ndim != 2 or grid_values.shape != (n_points, n_points):
-            raise ValueError(
-                f"grid_values must be a square 2D array of shape ({n_points}, {n_points})."
-            )
-
-        if any(abs(grid) > 1 + 1e-12):
+        super().__init__(grid, grid_values, enforce_pdf_nonnegative)
+        if any(abs(self.grid) > 1 + 1e-12):
             raise ValueError(
                 "Grid points must have coordinates in [-1, 1] (unit sphere)."
             )
-
-        if enforce_pdf_nonnegative and any(grid_values < 0):
-            raise ValueError("grid_values must be non-negative.")
-
-        self.grid = grid
-        self.grid_values = grid_values
-        self.enforce_pdf_nonnegative = enforce_pdf_nonnegative
-        # Embedding dimension of the Cartesian product space (convention from
-        # libDirectional: dim = 2 * embedding_dim_of_individual_sphere).
-        self.dim = 2 * d
-
         self._check_normalization()
 
     # ------------------------------------------------------------------
@@ -107,43 +81,6 @@ def _check_normalization(self, tol=0.01):
                 UserWarning,
             )
 
-    def normalize(self):
-        """No-op – returns ``self`` for compatibility."""
-        return self
-
-    # ------------------------------------------------------------------
-    # Arithmetic
-    # ------------------------------------------------------------------
-
-    def multiply(self, other):
-        """
-        Element-wise multiply two conditional grid distributions.
-
-        The resulting distribution is *not* normalized.
-
-        Parameters
-        ----------
-        other : SdCondSdGridDistribution
-            Must be defined on the same grid.
-
-        Returns
-        -------
-        SdCondSdGridDistribution
-        """
-        if not array_equal(self.grid, other.grid):
-            raise ValueError(
-                "Multiply:IncompatibleGrid: Can only multiply distributions "
-                "defined on identical grids."
-            )
-        warnings.warn(
-            "Multiply:UnnormalizedResult: Multiplication does not yield a "
-            "normalized result.",
-            UserWarning,
-        )
-        result = copy.deepcopy(self)
-        result.grid_values = result.grid_values * other.grid_values
-        return result
-
     # ------------------------------------------------------------------
     # Marginalisation and conditioning
     # ------------------------------------------------------------------
@@ -201,26 +138,7 @@ def fix_dim(self, first_or_second, point):
             HypersphericalGridDistribution,
         )
 
-        d = self.grid.shape[1]
-        if point.shape[0] != d:
-            raise ValueError(
-                f"point must have length {d} (embedding dimension of the sphere)."
-            )
-
-        diffs = linalg.norm(self.grid - point[None, :], axis=1)
-        locb = argmin(diffs)
-        if diffs[locb] > 1e-10:
-            raise ValueError(
-                "Cannot fix value at this point because it is not on the grid."
-            )
-
-        if first_or_second == 1:
-            grid_values_slice = self.grid_values[locb, :]
-        elif first_or_second == 2:
-            grid_values_slice = self.grid_values[:, locb]
-        else:
-            raise ValueError("first_or_second must be 1 or 2.")
-
+        grid_values_slice = self._get_grid_slice(first_or_second, point)
         return HypersphericalGridDistribution(self.grid, grid_values_slice)
 
     # ------------------------------------------------------------------
@@ -276,24 +194,8 @@ def from_function(
         # manifold dim: embedding_dim = dim // 2, manifold_dim = embedding_dim - 1.
         manifold_dim = dim // 2 - 1
         grid, _ = get_grid_hypersphere(grid_type, n, manifold_dim)
-        # grid is (n, dim//2)
-
-        if fun_does_cartesian_product:
-            fvals = fun(grid, grid)
-            grid_values = fvals.reshape(n, n)
-        else:
-            # Build index pairs: idx_a[i, j] = i, idx_b[i, j] = j
-            idx_a, idx_b = meshgrid(arange(n), arange(n), indexing="ij")
-            grid_a = grid[idx_a.ravel()]  # (n*n, d)
-            grid_b = grid[idx_b.ravel()]  # (n*n, d)
-            fvals = fun(grid_a, grid_b)  # (n*n,)
-
-            if fvals.shape == (n**2, n**2):
-                raise ValueError(
-                    "Function apparently performs the Cartesian product itself. "
-                    "Set fun_does_cartesian_product=True."
-                )
-
-            grid_values = fvals.reshape(n, n)
 
+        grid_values = SdCondSdGridDistribution._evaluate_on_grid(
+            fun, grid, n, fun_does_cartesian_product
+        )
         return SdCondSdGridDistribution(grid, grid_values)
diff --git a/pyrecest/distributions/conditional/td_cond_td_grid_distribution.py b/pyrecest/distributions/conditional/td_cond_td_grid_distribution.py
index 6afae5b1e..86007b153 100644
--- a/pyrecest/distributions/conditional/td_cond_td_grid_distribution.py
+++ b/pyrecest/distributions/conditional/td_cond_td_grid_distribution.py
@@ -1,4 +1,3 @@
-import copy
 import warnings
 
 # pylint: disable=redefined-builtin,no-name-in-module,no-member
@@ -6,12 +5,7 @@
     abs,
     all,
     any,
-    arange,
-    argmin,
-    array_equal,
-    linalg,
     mean,
-    meshgrid,
     pi,
     sum,
 )
@@ -47,26 +41,7 @@ def __init__(self, grid, grid_values, enforce_pdf_nonnegative=True):
         enforce_pdf_nonnegative : bool
             Whether non-negativity of ``grid_values`` is required.
         """
-        if grid.ndim != 2:
-            raise ValueError("grid must be a 2D array of shape (n_points, d).")
-
-        n_points, d = grid.shape
-
-        if grid_values.ndim != 2 or grid_values.shape != (n_points, n_points):
-            raise ValueError(
-                f"grid_values must be a square 2D array of shape ({n_points}, {n_points})."
-            )
-
-        if enforce_pdf_nonnegative and any(grid_values < 0):
-            raise ValueError("grid_values must be non-negative.")
-
-        self.grid = grid
-        self.grid_values = grid_values
-        self.enforce_pdf_nonnegative = enforce_pdf_nonnegative
-        # Embedding dimension of the Cartesian product space (convention from
-        # libDirectional: dim = 2 * dim_of_individual_torus).
-        self.dim = 2 * d
-
+        super().__init__(grid, grid_values, enforce_pdf_nonnegative)
         self._check_normalization()
 
     # ------------------------------------------------------------------
@@ -96,43 +71,6 @@ def _check_normalization(self, tol=0.01):
                 UserWarning,
             )
 
-    def normalize(self):
-        """No-op – returns ``self`` for compatibility."""
-        return self
-
-    # ------------------------------------------------------------------
-    # Arithmetic
-    # ------------------------------------------------------------------
-
-    def multiply(self, other):
-        """
-        Element-wise multiply two conditional grid distributions.
-
-        The resulting distribution is *not* normalized.
-
-        Parameters
-        ----------
-        other : TdCondTdGridDistribution
-            Must be defined on the same grid.
-
-        Returns
-        -------
-        TdCondTdGridDistribution
-        """
-        if not array_equal(self.grid, other.grid):
-            raise ValueError(
-                "Multiply:IncompatibleGrid: Can only multiply distributions "
-                "defined on identical grids."
-            )
-        warnings.warn(
-            "Multiply:UnnormalizedResult: Multiplication does not yield a "
-            "normalized result.",
-            UserWarning,
-        )
-        result = copy.deepcopy(self)
-        result.grid_values = result.grid_values * other.grid_values
-        return result
-
     # ------------------------------------------------------------------
     # Marginalisation and conditioning
     # ------------------------------------------------------------------
@@ -190,26 +128,7 @@ def fix_dim(self, first_or_second, point):
             HypertoroidalGridDistribution,
         )
 
-        d = self.grid.shape[1]
-        if point.shape[0] != d:
-            raise ValueError(
-                f"point must have length {d} (dimension of the torus)."
-            )
-
-        diffs = linalg.norm(self.grid - point[None, :], 1)
-        locb = argmin(diffs)
-        if diffs[locb] > 1e-10:
-            raise ValueError(
-                "Cannot fix value at this point because it is not on the grid."
-            )
-
-        if first_or_second == 1:
-            grid_values_slice = self.grid_values[locb, :]
-        elif first_or_second == 2:
-            grid_values_slice = self.grid_values[:, locb]
-        else:
-            raise ValueError("first_or_second must be 1 or 2.")
-
+        grid_values_slice = self._get_grid_slice(first_or_second, point)
         return HypertoroidalGridDistribution(grid_values_slice, "custom", self.grid)
 
     # ------------------------------------------------------------------
@@ -277,21 +196,7 @@ def from_function(
             [n] * dim_half
         )
 
-        if fun_does_cartesian_product:
-            fvals = fun(grid, grid)
-            grid_values = fvals.reshape(n, n)
-        else:
-            idx_a, idx_b = meshgrid(arange(n), arange(n), indexing="ij")
-            grid_a = grid[idx_a.ravel()]
-            grid_b = grid[idx_b.ravel()]
-            fvals = fun(grid_a, grid_b)
-
-            if fvals.shape == (n**2, n**2):
-                raise ValueError(
-                    "Function apparently performs the Cartesian product itself. "
-                    "Set fun_does_cartesian_product=True."
-                )
-
-            grid_values = fvals.reshape(n, n)
-
+        grid_values = TdCondTdGridDistribution._evaluate_on_grid(
+            fun, grid, n, fun_does_cartesian_product
+        )
         return TdCondTdGridDistribution(grid, grid_values)

From 0c8f19ddf25468cfff488b7371b1b4262d9fada8 Mon Sep 17 00:00:00 2001
From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com>
Date: Fri, 3 Apr 2026 12:56:41 +0000
Subject: [PATCH 16/19] Remove unused abs import from
 abstract_conditional_distribution

Agent-Logs-Url: https://github.com/FlorianPfaff/PyRecEst/sessions/973bac70-4353-4a5f-95b0-ef2eb5b18d82

Co-authored-by: FlorianPfaff <6773539+FlorianPfaff@users.noreply.github.com>
---
 .../conditional/abstract_conditional_distribution.py             | 1 -
 1 file changed, 1 deletion(-)

diff --git a/pyrecest/distributions/conditional/abstract_conditional_distribution.py b/pyrecest/distributions/conditional/abstract_conditional_distribution.py
index c0f2cda97..c5e332177 100644
--- a/pyrecest/distributions/conditional/abstract_conditional_distribution.py
+++ b/pyrecest/distributions/conditional/abstract_conditional_distribution.py
@@ -4,7 +4,6 @@
 
 # pylint: disable=redefined-builtin,no-name-in-module,no-member
 from pyrecest.backend import (
-    abs,
     any,
     arange,
     argmin,

From 108cfa25c37db210c0434a97955a9110c403732f Mon Sep 17 00:00:00 2001
From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com>
Date: Fri, 3 Apr 2026 13:35:40 +0000
Subject: [PATCH 17/19] Replace numpy with pyrecest.backend in
 test_td_cond_td_grid_distribution

Agent-Logs-Url: https://github.com/FlorianPfaff/PyRecEst/sessions/f320cb8b-dfbf-4e21-b822-9b668def8bf2

Co-authored-by: FlorianPfaff <6773539+FlorianPfaff@users.noreply.github.com>
---
 .../test_td_cond_td_grid_distribution.py      | 131 ++++++++++--------
 1 file changed, 71 insertions(+), 60 deletions(-)

diff --git a/pyrecest/tests/distributions/test_td_cond_td_grid_distribution.py b/pyrecest/tests/distributions/test_td_cond_td_grid_distribution.py
index 73b08bb25..5182d35da 100644
--- a/pyrecest/tests/distributions/test_td_cond_td_grid_distribution.py
+++ b/pyrecest/tests/distributions/test_td_cond_td_grid_distribution.py
@@ -1,9 +1,19 @@
 import unittest
 
-import numpy as np
 import numpy.testing as npt
 
-from pyrecest.backend import array
+from pyrecest.backend import (
+    abs,
+    array,
+    asarray,
+    exp,
+    linspace,
+    minimum,
+    ones,
+    pi,
+    random,
+    zeros,
+)
 from pyrecest.distributions.conditional.td_cond_td_grid_distribution import (
     TdCondTdGridDistribution,
 )
@@ -12,11 +22,12 @@
 )
 
 
-def _make_normalized_grid_values(n: int) -> np.ndarray:
+def _make_normalized_grid_values(n: int):
     """Return an (n x n) matrix whose columns are normalized (integrate to 1)."""
-    gv = np.random.default_rng(0).uniform(0.5, 1.5, size=(n, n))
+    random.seed(0)
+    gv = random.uniform(0.5, 1.5, (n, n))
     # Normalize each column so that mean(col) * (2*pi)^1 == 1
-    gv = gv / (gv.mean(axis=0, keepdims=True) * 2.0 * np.pi)
+    gv = gv / (gv.mean(axis=0) * 2.0 * pi)
     return gv
 
 
@@ -26,46 +37,46 @@ class TdCondTdGridDistributionTest(unittest.TestCase):
     def test_construction_t1(self):
         """Basic construction for T1 x T1."""
         n = 5
-        grid = np.linspace(0.0, 2.0 * np.pi - 2.0 * np.pi / n, n).reshape(-1, 1)
+        grid = linspace(0.0, 2.0 * pi - 2.0 * pi / n, n).reshape(-1, 1)
         gv = _make_normalized_grid_values(n)
-        td = TdCondTdGridDistribution(array(grid), array(gv))
+        td = TdCondTdGridDistribution(grid, gv)
         self.assertEqual(td.dim, 2)
         npt.assert_allclose(td.grid, grid, rtol=1e-6)
         npt.assert_allclose(td.grid_values, gv, rtol=1e-6)
 
     def test_construction_wrong_shape_raises(self):
         n = 4
-        grid = np.zeros((n, 1))
+        grid = zeros((n, 1))
         with self.assertRaises(ValueError):
             # Non-square grid_values
             TdCondTdGridDistribution(
-                array(grid), array(np.ones((n, n + 1)) / (n * 2 * np.pi))
+                grid, ones((n, n + 1)) / (n * 2 * pi)
             )
 
     def test_construction_wrong_order_raises(self):
         """Transposed (row-normalized) matrix should raise an error."""
         n = 6
-        grid = np.linspace(0.0, 2.0 * np.pi - 2.0 * np.pi / n, n).reshape(-1, 1)
+        grid = linspace(0.0, 2.0 * pi - 2.0 * pi / n, n).reshape(-1, 1)
         gv = _make_normalized_grid_values(n)
         # Transpose → rows are normalized, columns are not
         with self.assertRaises(ValueError):
-            TdCondTdGridDistribution(array(grid), array(gv.T))
+            TdCondTdGridDistribution(grid, gv.T)
 
     def test_construction_unnormalized_warns(self):
         """An unnormalized matrix that cannot be fixed by transposing should warn."""
         n = 5
-        grid = np.linspace(0.0, 2.0 * np.pi - 2.0 * np.pi / n, n).reshape(-1, 1)
-        gv = np.ones((n, n))  # neither rows nor cols sum to 1/(2pi)
+        grid = linspace(0.0, 2.0 * pi - 2.0 * pi / n, n).reshape(-1, 1)
+        gv = ones((n, n))  # neither rows nor cols sum to 1/(2pi)
         with self.assertWarns(UserWarning):
-            TdCondTdGridDistribution(array(grid), array(gv))
+            TdCondTdGridDistribution(grid, gv)
 
     # -------------------------------------------------------------- normalize
 
     def test_normalize_returns_self(self):
         n = 4
-        grid = np.linspace(0.0, 2.0 * np.pi - 2.0 * np.pi / n, n).reshape(-1, 1)
+        grid = linspace(0.0, 2.0 * pi - 2.0 * pi / n, n).reshape(-1, 1)
         gv = _make_normalized_grid_values(n)
-        td = TdCondTdGridDistribution(array(grid), array(gv))
+        td = TdCondTdGridDistribution(grid, gv)
         self.assertIs(td.normalize(), td)
 
     # -------------------------------------------------------------- multiply
@@ -74,27 +85,27 @@ def test_multiply_same_grid(self):
         import warnings
 
         n = 6
-        grid = np.linspace(0.0, 2.0 * np.pi - 2.0 * np.pi / n, n).reshape(-1, 1)
+        grid = linspace(0.0, 2.0 * pi - 2.0 * pi / n, n).reshape(-1, 1)
         gv = _make_normalized_grid_values(n)
-        td1 = TdCondTdGridDistribution(array(grid), array(gv))
-        td2 = TdCondTdGridDistribution(array(grid), array(gv))
+        td1 = TdCondTdGridDistribution(grid, gv)
+        td2 = TdCondTdGridDistribution(grid, gv)
         with warnings.catch_warnings():
             warnings.simplefilter("ignore")
             result = td1.multiply(td2)
         npt.assert_allclose(
-            np.asarray(result.grid_values),
-            np.asarray(td1.grid_values) * np.asarray(td2.grid_values),
+            asarray(result.grid_values),
+            asarray(td1.grid_values) * asarray(td2.grid_values),
             rtol=1e-10,
         )
 
     def test_multiply_incompatible_grid_raises(self):
         n1, n2 = 4, 6
-        grid1 = np.linspace(0.0, 2.0 * np.pi - 2.0 * np.pi / n1, n1).reshape(-1, 1)
-        grid2 = np.linspace(0.0, 2.0 * np.pi - 2.0 * np.pi / n2, n2).reshape(-1, 1)
+        grid1 = linspace(0.0, 2.0 * pi - 2.0 * pi / n1, n1).reshape(-1, 1)
+        grid2 = linspace(0.0, 2.0 * pi - 2.0 * pi / n2, n2).reshape(-1, 1)
         gv1 = _make_normalized_grid_values(n1)
         gv2 = _make_normalized_grid_values(n2)
-        td1 = TdCondTdGridDistribution(array(grid1), array(gv1))
-        td2 = TdCondTdGridDistribution(array(grid2), array(gv2))
+        td1 = TdCondTdGridDistribution(grid1, gv1)
+        td2 = TdCondTdGridDistribution(grid2, gv2)
         with self.assertRaises(ValueError) as ctx:
             td1.multiply(td2)
         self.assertIn("IncompatibleGrid", str(ctx.exception))
@@ -108,16 +119,16 @@ def test_from_function_t1(self):
 
         def cond_fun(a, b):
             # Simple Gaussian-like conditional (unnormalized, normalized per column)
-            diff = np.asarray(a)[:, 0] - np.asarray(b)[:, 0]
-            return np.exp(-0.5 * np.minimum(diff**2, (2 * np.pi - np.abs(diff)) ** 2))
+            diff = asarray(a)[:, 0] - asarray(b)[:, 0]
+            return exp(-0.5 * minimum(diff**2, (2 * pi - abs(diff)) ** 2))
 
         td = TdCondTdGridDistribution.from_function(
             cond_fun, n, fun_does_cartesian_product=False, grid_type="CartesianProd", dim=dim
         )
         self.assertIsInstance(td, TdCondTdGridDistribution)
         self.assertEqual(td.dim, dim)
-        self.assertEqual(np.asarray(td.grid_values).shape, (n, n))
-        self.assertEqual(np.asarray(td.grid).shape, (n, 1))
+        self.assertEqual(asarray(td.grid_values).shape, (n, n))
+        self.assertEqual(asarray(td.grid).shape, (n, 1))
 
     def test_from_function_cartesian_product_flag(self):
         """from_function with fun_does_cartesian_product=True."""
@@ -126,10 +137,10 @@ def test_from_function_cartesian_product_flag(self):
 
         def cond_fun_cp(a, b):
             # a: (n, 1), b: (n, 1) → return (n, n)
-            a_np = np.asarray(a)[:, 0]
-            b_np = np.asarray(b)[:, 0]
-            diff = a_np[:, None] - b_np[None, :]
-            return np.exp(-0.5 * np.minimum(diff**2, (2 * np.pi - np.abs(diff)) ** 2))
+            a_arr = asarray(a)[:, 0]
+            b_arr = asarray(b)[:, 0]
+            diff = a_arr[:, None] - b_arr[None, :]
+            return exp(-0.5 * minimum(diff**2, (2 * pi - abs(diff)) ** 2))
 
         td = TdCondTdGridDistribution.from_function(
             cond_fun_cp,
@@ -139,13 +150,13 @@ def cond_fun_cp(a, b):
             dim=dim,
         )
         self.assertIsInstance(td, TdCondTdGridDistribution)
-        self.assertEqual(np.asarray(td.grid_values).shape, (n, n))
+        self.assertEqual(asarray(td.grid_values).shape, (n, n))
 
     def test_from_function_unknown_grid_raises(self):
         n = 4
         with self.assertRaises(ValueError):
             TdCondTdGridDistribution.from_function(
-                lambda a, b: np.ones(len(a)),
+                lambda a, b: ones(len(a)),
                 n,
                 fun_does_cartesian_product=False,
                 grid_type="unknownGrid",
@@ -155,16 +166,16 @@ def test_from_function_unknown_grid_raises(self):
     def test_from_function_odd_dim_raises(self):
         with self.assertRaises(ValueError):
             TdCondTdGridDistribution.from_function(
-                lambda a, b: np.ones(len(a)), 4, False, "CartesianProd", dim=3
+                lambda a, b: ones(len(a)), 4, False, "CartesianProd", dim=3
             )
 
     # --------------------------------------------------------- marginalize_out
 
     def test_marginalize_out_returns_hgd(self):
         n = 6
-        grid = np.linspace(0.0, 2.0 * np.pi - 2.0 * np.pi / n, n).reshape(-1, 1)
+        grid = linspace(0.0, 2.0 * pi - 2.0 * pi / n, n).reshape(-1, 1)
         gv = _make_normalized_grid_values(n)
-        td = TdCondTdGridDistribution(array(grid), array(gv))
+        td = TdCondTdGridDistribution(grid, gv)
 
         for first_or_second in (1, 2):
             with self.subTest(first_or_second=first_or_second):
@@ -174,30 +185,30 @@ def test_marginalize_out_returns_hgd(self):
 
     def test_marginalize_out_sums(self):
         n = 5
-        grid = np.linspace(0.0, 2.0 * np.pi - 2.0 * np.pi / n, n).reshape(-1, 1)
+        grid = linspace(0.0, 2.0 * pi - 2.0 * pi / n, n).reshape(-1, 1)
         gv = _make_normalized_grid_values(n)
-        td = TdCondTdGridDistribution(array(grid), array(gv))
+        td = TdCondTdGridDistribution(grid, gv)
 
         # marginalize_out(1) sums rows; HGD normalizes, so check proportionality
         m1 = td.marginalize_out(1)
         expected_unnorm = gv.sum(axis=0)
-        actual_unnorm = np.asarray(m1.grid_values) * float(m1.integrate())
+        actual_unnorm = asarray(m1.grid_values) * float(m1.integrate())
         # Check proportionality (ratio should be constant)
         ratio = actual_unnorm / expected_unnorm
-        npt.assert_allclose(ratio, ratio[0] * np.ones(n), atol=1e-12)
+        npt.assert_allclose(ratio, ratio[0] * ones(n), atol=1e-12)
 
         # marginalize_out(2) sums cols
         m2 = td.marginalize_out(2)
         expected_unnorm2 = gv.sum(axis=1)
-        actual_unnorm2 = np.asarray(m2.grid_values) * float(m2.integrate())
+        actual_unnorm2 = asarray(m2.grid_values) * float(m2.integrate())
         ratio2 = actual_unnorm2 / expected_unnorm2
-        npt.assert_allclose(ratio2, ratio2[0] * np.ones(n), atol=1e-12)
+        npt.assert_allclose(ratio2, ratio2[0] * ones(n), atol=1e-12)
 
     def test_marginalize_out_invalid_raises(self):
         n = 5
-        grid = np.linspace(0.0, 2.0 * np.pi - 2.0 * np.pi / n, n).reshape(-1, 1)
+        grid = linspace(0.0, 2.0 * pi - 2.0 * pi / n, n).reshape(-1, 1)
         gv = _make_normalized_grid_values(n)
-        td = TdCondTdGridDistribution(array(grid), array(gv))
+        td = TdCondTdGridDistribution(grid, gv)
         with self.assertRaises(ValueError):
             td.marginalize_out(0)
         with self.assertRaises(ValueError):
@@ -207,9 +218,9 @@ def test_marginalize_out_invalid_raises(self):
 
     def test_fix_dim_returns_hgd(self):
         n = 5
-        grid_np = np.linspace(0.0, 2.0 * np.pi - 2.0 * np.pi / n, n).reshape(-1, 1)
+        grid_np = linspace(0.0, 2.0 * pi - 2.0 * pi / n, n).reshape(-1, 1)
         gv = _make_normalized_grid_values(n)
-        td = TdCondTdGridDistribution(array(grid_np), array(gv))
+        td = TdCondTdGridDistribution(grid_np, gv)
 
         for first_or_second in (1, 2):
             with self.subTest(first_or_second=first_or_second):
@@ -220,34 +231,34 @@ def test_fix_dim_returns_hgd(self):
 
     def test_fix_dim_off_grid_raises(self):
         n = 5
-        grid_np = np.linspace(0.0, 2.0 * np.pi - 2.0 * np.pi / n, n).reshape(-1, 1)
+        grid_np = linspace(0.0, 2.0 * pi - 2.0 * pi / n, n).reshape(-1, 1)
         gv = _make_normalized_grid_values(n)
-        td = TdCondTdGridDistribution(array(grid_np), array(gv))
+        td = TdCondTdGridDistribution(grid_np, gv)
         with self.assertRaises(ValueError):
-            td.fix_dim(1, np.array([1.23456789]))
+            td.fix_dim(1, array([1.23456789]))
 
     def test_fix_dim_values_correct(self):
         """fix_dim(2, grid[j]) should give a distribution proportional to col j."""
         n = 6
-        grid_np = np.linspace(0.0, 2.0 * np.pi - 2.0 * np.pi / n, n).reshape(-1, 1)
+        grid_np = linspace(0.0, 2.0 * pi - 2.0 * pi / n, n).reshape(-1, 1)
         gv = _make_normalized_grid_values(n)
-        td = TdCondTdGridDistribution(array(grid_np), array(gv))
+        td = TdCondTdGridDistribution(grid_np, gv)
 
         j = 3
         slice_dist = td.fix_dim(2, grid_np[j])
         expected = gv[:, j]
-        expected = expected / expected.mean() / (2.0 * np.pi)  # normalize
+        expected = expected / expected.mean() / (2.0 * pi)  # normalize
         npt.assert_allclose(
-            np.asarray(slice_dist.grid_values),
+            asarray(slice_dist.grid_values),
             expected,
             atol=1e-12,
         )
 
     def test_fix_dim_invalid_raises(self):
         n = 5
-        grid_np = np.linspace(0.0, 2.0 * np.pi - 2.0 * np.pi / n, n).reshape(-1, 1)
+        grid_np = linspace(0.0, 2.0 * pi - 2.0 * pi / n, n).reshape(-1, 1)
         gv = _make_normalized_grid_values(n)
-        td = TdCondTdGridDistribution(array(grid_np), array(gv))
+        td = TdCondTdGridDistribution(grid_np, gv)
         point = grid_np[0]
         with self.assertRaises(ValueError):
             td.fix_dim(0, point)
@@ -262,13 +273,13 @@ def test_from_function_fix_dim_roundtrip(self):
         dim = 2
 
         def cond_fun(a, b):
-            diff = np.asarray(a)[:, 0] - np.asarray(b)[:, 0]
-            return np.exp(-0.5 * diff**2)
+            diff = asarray(a)[:, 0] - asarray(b)[:, 0]
+            return exp(-0.5 * diff**2)
 
         td = TdCondTdGridDistribution.from_function(
             cond_fun, n, fun_does_cartesian_product=False, dim=dim
         )
-        grid_np = np.asarray(td.grid)
+        grid_np = asarray(td.grid)
         slice_dist = td.fix_dim(2, grid_np[0])
         self.assertIsInstance(slice_dist, HypertoroidalGridDistribution)
 

From 7e8dfbfbe508226eb288f80a895022e6b34daf4b Mon Sep 17 00:00:00 2001
From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com>
Date: Fri, 3 Apr 2026 13:50:29 +0000
Subject: [PATCH 18/19] Suppress W0622 redefined-builtin for abs import in test
 file

Agent-Logs-Url: https://github.com/FlorianPfaff/PyRecEst/sessions/45f4bfe8-c75e-4be9-9289-29743706bc26

Co-authored-by: FlorianPfaff <6773539+FlorianPfaff@users.noreply.github.com>
---
 .../tests/distributions/test_td_cond_td_grid_distribution.py    | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/pyrecest/tests/distributions/test_td_cond_td_grid_distribution.py b/pyrecest/tests/distributions/test_td_cond_td_grid_distribution.py
index 5182d35da..f2aa08fce 100644
--- a/pyrecest/tests/distributions/test_td_cond_td_grid_distribution.py
+++ b/pyrecest/tests/distributions/test_td_cond_td_grid_distribution.py
@@ -2,7 +2,7 @@
 
 import numpy.testing as npt
 
-from pyrecest.backend import (
+from pyrecest.backend import (  # pylint: disable=redefined-builtin
     abs,
     array,
     asarray,

From 136db1ca2f06e310ebba5e3e4ebd520f49432cdc Mon Sep 17 00:00:00 2001
From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com>
Date: Fri, 3 Apr 2026 18:07:24 +0000
Subject: [PATCH 19/19] Fix random.uniform call to use cross-backend compatible
 size= kwarg

Agent-Logs-Url: https://github.com/FlorianPfaff/PyRecEst/sessions/cb85a0e9-ba2b-4f41-8016-d692369737b1

Co-authored-by: FlorianPfaff <6773539+FlorianPfaff@users.noreply.github.com>
---
 .../tests/distributions/test_td_cond_td_grid_distribution.py    | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/pyrecest/tests/distributions/test_td_cond_td_grid_distribution.py b/pyrecest/tests/distributions/test_td_cond_td_grid_distribution.py
index f2aa08fce..100007edc 100644
--- a/pyrecest/tests/distributions/test_td_cond_td_grid_distribution.py
+++ b/pyrecest/tests/distributions/test_td_cond_td_grid_distribution.py
@@ -25,7 +25,7 @@
 def _make_normalized_grid_values(n: int):
     """Return an (n x n) matrix whose columns are normalized (integrate to 1)."""
     random.seed(0)
-    gv = random.uniform(0.5, 1.5, (n, n))
+    gv = 0.5 + random.uniform(size=(n, n))
     # Normalize each column so that mean(col) * (2*pi)^1 == 1
     gv = gv / (gv.mean(axis=0) * 2.0 * pi)
     return gv