From 85fd146d366c266c62a6af5d9aec654a15a50e9d Mon Sep 17 00:00:00 2001
From: Florian Pfaff <6773539+FlorianPfaff@users.noreply.github.com>
Date: Sat, 9 May 2026 02:32:42 +0200
Subject: [PATCH 1/8] Add generic pairwise covariance feature utilities

---
 .../utils/pairwise_covariance_features.py     | 234 ++++++++++++++++++
 1 file changed, 234 insertions(+)
 create mode 100644 src/pyrecest/utils/pairwise_covariance_features.py

diff --git a/src/pyrecest/utils/pairwise_covariance_features.py b/src/pyrecest/utils/pairwise_covariance_features.py
new file mode 100644
index 000000000..80ab624b6
--- /dev/null
+++ b/src/pyrecest/utils/pairwise_covariance_features.py
@@ -0,0 +1,234 @@
+"""Pairwise Gaussian and covariance-derived association features.
+
+The helpers in this module operate on plain mean and covariance arrays rather
+than on tracker-, sensor-, or application-specific objects.  They are intended
+for building feature tensors for association models and assignment costs from
+Gaussian state summaries.
+
+All stacked mean and covariance inputs use the standard PyRecEst convention:
+
+* means have shape ``(dim, n_items)``;
+* covariance stacks have shape ``(dim, dim, n_items)``.
+"""
+
+from __future__ import annotations
+
+import math
+from typing import Any
+
+# pylint: disable=no-name-in-module,no-member,redefined-builtin
+from pyrecest.backend import (
+    abs,
+    all as backend_all,
+    asarray,
+    einsum,
+    exp,
+    eye,
+    float64,
+    isfinite,
+    log,
+    maximum,
+    moveaxis,
+    reshape,
+    sqrt,
+    sum as backend_sum,
+    trace,
+    transpose,
+    where,
+    zeros,
+)
+from pyrecest.backend import linalg
+
+
+def pairwise_mahalanobis_distances(
+    means_a: Any,
+    covariances_a: Any,
+    means_b: Any,
+    covariances_b: Any,
+    *,
+    regularization: float = 0.0,
+) -> Any:
+    """Return covariance-normalized distances between two Gaussian stacks.
+
+    For a pair of items ``i`` and ``j``, the returned value is
+
+    ``sqrt((mu_i - nu_j)^T (Sigma_i + Lambda_j + regularization * I)^+ (mu_i - nu_j))``,
+
+    where ``^+`` denotes the Moore-Penrose pseudoinverse.  Using the summed
+    covariance makes the feature symmetric in the two uncertain estimates and is
+    the standard normalization for comparing two independent Gaussian position
+    estimates.
+
+    Parameters
+    ----------
+    means_a, means_b:
+        Mean stacks with shape ``(dim, n_items)``.
+    covariances_a, covariances_b:
+        Covariance stacks with shape ``(dim, dim, n_items)``.  The number of
+        covariance matrices must match the corresponding number of means.
+    regularization:
+        Optional non-negative diagonal loading added to every summed covariance
+        before inversion.
+
+    Returns
+    -------
+    array-like
+        Matrix with shape ``(n_a, n_b)``.
+    """
+
+    if not math.isfinite(float(regularization)) or regularization < 0.0:
+        raise ValueError("regularization must be finite and non-negative")
+
+    means_a, covariances_a = _validate_mean_covariance_stack(
+        "means_a", means_a, "covariances_a", covariances_a
+    )
+    means_b, covariances_b = _validate_mean_covariance_stack(
+        "means_b", means_b, "covariances_b", covariances_b
+    )
+
+    dim, n_a = means_a.shape
+    dim_b, n_b = means_b.shape
+    if dim != dim_b:
+        raise ValueError(
+            "means_a and means_b must have the same leading dimension"
+        )
+
+    if n_a == 0 or n_b == 0:
+        return zeros((n_a, n_b), dtype=float64)
+
+    moved_covariances_a = _symmetrized_covariance_batch(covariances_a)
+    moved_covariances_b = _symmetrized_covariance_batch(covariances_b)
+    covariance_sums = (
+        moved_covariances_a[:, None, :, :] + moved_covariances_b[None, :, :, :]
+    )
+    if regularization > 0.0:
+        covariance_sums = covariance_sums + float(regularization) * eye(dim)[
+            None, None, :, :
+        ]
+
+    mean_differences = transpose(means_a)[:, None, :] - transpose(means_b)[None, :, :]
+    flat_differences = reshape(mean_differences, (n_a * n_b, dim))
+    flat_covariances = reshape(covariance_sums, (n_a * n_b, dim, dim))
+    inverse_covariances = linalg.pinv(flat_covariances)
+    squared_distances = einsum(
+        "ni,nij,nj->n", flat_differences, inverse_covariances, flat_differences
+    )
+    squared_distances = reshape(squared_distances, (n_a, n_b))
+    return sqrt(maximum(squared_distances, 0.0))
+
+
+def pairwise_covariance_shape_components(
+    covariances_a: Any,
+    covariances_b: Any,
+    *,
+    epsilon: float = 1.0e-6,
+) -> tuple[Any, Any, Any]:
+    """Return pairwise covariance shape, scale, and similarity components.
+
+    The covariance-shape cost compares trace-normalized covariance matrices with
+    a Frobenius norm.  This makes the feature sensitive to orientation and
+    anisotropy while ignoring overall scale.  The log-determinant cost measures
+    scale mismatch separately.  The shape similarity is ``exp(-shape_cost)``.
+
+    Parameters
+    ----------
+    covariances_a, covariances_b:
+        Covariance stacks with shape ``(dim, dim, n_items)``.
+    epsilon:
+        Strictly positive floor used for traces and determinants.
+
+    Returns
+    -------
+    shape_cost, logdet_cost, shape_similarity:
+        Three matrices with shape ``(n_a, n_b)``.
+    """
+
+    if not math.isfinite(float(epsilon)) or epsilon <= 0.0:
+        raise ValueError("epsilon must be finite and strictly positive")
+
+    covariances_a = _validate_covariance_stack("covariances_a", covariances_a)
+    covariances_b = _validate_covariance_stack("covariances_b", covariances_b)
+    dim = covariances_a.shape[0]
+    if covariances_b.shape[0] != dim:
+        raise ValueError(
+            "covariances_a and covariances_b must have the same matrix dimension"
+        )
+
+    n_a = covariances_a.shape[2]
+    n_b = covariances_b.shape[2]
+    if n_a == 0 or n_b == 0:
+        empty = zeros((n_a, n_b), dtype=float64)
+        return empty, empty, empty
+
+    moved_covariances_a = _symmetrized_covariance_batch(covariances_a)
+    moved_covariances_b = _symmetrized_covariance_batch(covariances_b)
+
+    traces_a = _positive_floor(
+        trace(moved_covariances_a, axis1=1, axis2=2), epsilon
+    )
+    traces_b = _positive_floor(
+        trace(moved_covariances_b, axis1=1, axis2=2), epsilon
+    )
+    normalized_a = moved_covariances_a / traces_a[:, None, None]
+    normalized_b = moved_covariances_b / traces_b[:, None, None]
+
+    shape_differences = normalized_a[:, None, :, :] - normalized_b[None, :, :, :]
+    frobenius_squared = backend_sum(
+        backend_sum(shape_differences * shape_differences, axis=-1), axis=-1
+    )
+    shape_cost = sqrt(maximum(frobenius_squared, 0.0)) / sqrt(2.0)
+    shape_similarity = exp(-shape_cost)
+
+    determinants_a = _positive_floor(linalg.det(moved_covariances_a), epsilon)
+    determinants_b = _positive_floor(linalg.det(moved_covariances_b), epsilon)
+    logdet_cost = abs(log(determinants_a[:, None] / determinants_b[None, :]))
+    return shape_cost, logdet_cost, shape_similarity
+
+
+def _validate_mean_covariance_stack(
+    mean_name: str,
+    means: Any,
+    covariance_name: str,
+    covariances: Any,
+) -> tuple[Any, Any]:
+    means = asarray(means, dtype=float64)
+    if means.ndim != 2:
+        raise ValueError(f"{mean_name} must have shape (dim, n_items)")
+    if not bool(backend_all(isfinite(means))):
+        raise ValueError(f"{mean_name} must contain only finite values")
+
+    covariances = _validate_covariance_stack(covariance_name, covariances)
+    if covariances.shape[0] != means.shape[0]:
+        raise ValueError(
+            f"{covariance_name} matrix dimension must match the leading dimension of {mean_name}"
+        )
+    if covariances.shape[2] != means.shape[1]:
+        raise ValueError(
+            f"{covariance_name} must contain one covariance matrix per column of {mean_name}"
+        )
+    return means, covariances
+
+
+def _validate_covariance_stack(name: str, covariances: Any) -> Any:
+    covariances = asarray(covariances, dtype=float64)
+    if covariances.ndim != 3 or covariances.shape[0] != covariances.shape[1]:
+        raise ValueError(f"{name} must have shape (dim, dim, n_items)")
+    if not bool(backend_all(isfinite(covariances))):
+        raise ValueError(f"{name} must contain only finite values")
+    return covariances
+
+
+def _symmetrized_covariance_batch(covariances: Any) -> Any:
+    moved = moveaxis(covariances, -1, 0)
+    return 0.5 * (moved + transpose(moved, (0, 2, 1)))
+
+
+def _positive_floor(values: Any, epsilon: float) -> Any:
+    values = asarray(values, dtype=float64)
+    return where(values > float(epsilon), values, float(epsilon))
+
+
+__all__ = [
+    "pairwise_covariance_shape_components",
+    "pairwise_mahalanobis_distances",
+]

From b241e7b15f37f1f4c24ef7bb2c56296d19dc4042 Mon Sep 17 00:00:00 2001
From: Florian Pfaff <6773539+FlorianPfaff@users.noreply.github.com>
Date: Sat, 9 May 2026 02:37:14 +0200
Subject: [PATCH 2/8] Export pairwise covariance feature utilities

---
 src/pyrecest/utils/__init__.py | 6 ++++++
 1 file changed, 6 insertions(+)

diff --git a/src/pyrecest/utils/__init__.py b/src/pyrecest/utils/__init__.py
index 4fa7c99e9..620a010f6 100644
--- a/src/pyrecest/utils/__init__.py
+++ b/src/pyrecest/utils/__init__.py
@@ -22,6 +22,10 @@
     estimate_thin_plate_spline,
     joint_tps_registration_assignment,
 )
+from .pairwise_covariance_features import (
+    pairwise_covariance_shape_components,
+    pairwise_mahalanobis_distances,
+)
 from .track_evaluation import (
     complete_track_set,
     normalize_track_matrix,
@@ -56,6 +60,8 @@
     "estimate_thin_plate_spline",
     "joint_tps_registration_assignment",
     "murty_k_best_assignments",
+    "pairwise_covariance_shape_components",
+    "pairwise_mahalanobis_distances",
     "complete_track_set",
     "normalize_track_matrix",
     "pairwise_track_set",

From 4205f652e6e27411876645fd415d7886c6566e5e Mon Sep 17 00:00:00 2001
From: Florian Pfaff <6773539+FlorianPfaff@users.noreply.github.com>
Date: Sat, 9 May 2026 02:39:42 +0200
Subject: [PATCH 3/8] Add tests for pairwise covariance features

---
 tests/test_pairwise_covariance_features.py | 146 +++++++++++++++++++++
 1 file changed, 146 insertions(+)
 create mode 100644 tests/test_pairwise_covariance_features.py

diff --git a/tests/test_pairwise_covariance_features.py b/tests/test_pairwise_covariance_features.py
new file mode 100644
index 000000000..a64af6356
--- /dev/null
+++ b/tests/test_pairwise_covariance_features.py
@@ -0,0 +1,146 @@
+import math
+import unittest
+
+import numpy.testing as npt
+
+# pylint: disable=no-name-in-module,no-member,redefined-builtin
+from pyrecest.backend import array, exp, log, zeros
+from pyrecest.utils import (
+    pairwise_covariance_shape_components,
+    pairwise_mahalanobis_distances,
+)
+
+
+class TestPairwiseCovarianceFeatures(unittest.TestCase):
+    def test_pairwise_mahalanobis_distances_use_summed_covariances(self):
+        means_a = array([[0.0], [0.0]])
+        means_b = array([[2.0, 0.0], [0.0, 2.0]])
+        covariances_a = array([[[1.0], [0.0]], [[0.0], [1.0]]])
+        covariances_b = array(
+            [
+                [[3.0, 1.0], [0.0, 0.0]],
+                [[0.0, 0.0], [1.0, 3.0]],
+            ]
+        )
+
+        distances = pairwise_mahalanobis_distances(
+            means_a,
+            covariances_a,
+            means_b,
+            covariances_b,
+        )
+
+        self.assertEqual(distances.shape, (1, 2))
+        npt.assert_allclose(distances, array([[1.0, 1.0]]))
+
+    def test_pairwise_mahalanobis_distances_support_diagonal_regularization(self):
+        means_a = array([[0.0], [0.0]])
+        means_b = array([[2.0], [0.0]])
+        zero_covariance = zeros((2, 2, 1))
+
+        distances = pairwise_mahalanobis_distances(
+            means_a,
+            zero_covariance,
+            means_b,
+            zero_covariance,
+            regularization=1.0,
+        )
+
+        npt.assert_allclose(distances, array([[2.0]]))
+
+    def test_pairwise_mahalanobis_distances_return_empty_pairwise_shape(self):
+        means_a = zeros((2, 0))
+        means_b = zeros((2, 3))
+        covariances_a = zeros((2, 2, 0))
+        covariances_b = zeros((2, 2, 3))
+
+        distances = pairwise_mahalanobis_distances(
+            means_a,
+            covariances_a,
+            means_b,
+            covariances_b,
+        )
+
+        self.assertEqual(distances.shape, (0, 3))
+
+    def test_pairwise_covariance_shape_components_separate_shape_and_scale(self):
+        covariances_a = array(
+            [
+                [[2.0], [0.0]],
+                [[0.0], [1.0]],
+            ]
+        )
+        covariances_b = array(
+            [
+                [[1.0, 2.0], [0.0, 0.0]],
+                [[0.0, 0.0], [1.0, 2.0]],
+            ]
+        )
+
+        shape_cost, logdet_cost, shape_similarity = (
+            pairwise_covariance_shape_components(covariances_a, covariances_b)
+        )
+
+        npt.assert_allclose(shape_cost, array([[1.0 / 6.0, 1.0 / 6.0]]))
+        npt.assert_allclose(logdet_cost, array([[log(2.0), log(2.0)]]))
+        npt.assert_allclose(shape_similarity, exp(-shape_cost))
+
+    def test_pairwise_covariance_shape_components_return_zero_for_equal_shapes(self):
+        covariances_a = array(
+            [
+                [[1.0], [0.0]],
+                [[0.0], [1.0]],
+            ]
+        )
+        covariances_b = array(
+            [
+                [[2.0], [0.0]],
+                [[0.0], [2.0]],
+            ]
+        )
+
+        shape_cost, logdet_cost, shape_similarity = (
+            pairwise_covariance_shape_components(covariances_a, covariances_b)
+        )
+
+        npt.assert_allclose(shape_cost, array([[0.0]]))
+        npt.assert_allclose(logdet_cost, array([[math.log(4.0)]]))
+        npt.assert_allclose(shape_similarity, array([[1.0]]))
+
+    def test_pairwise_covariance_shape_components_support_empty_stacks(self):
+        shape_cost, logdet_cost, shape_similarity = (
+            pairwise_covariance_shape_components(zeros((2, 2, 0)), zeros((2, 2, 4)))
+        )
+
+        self.assertEqual(shape_cost.shape, (0, 4))
+        self.assertEqual(logdet_cost.shape, (0, 4))
+        self.assertEqual(shape_similarity.shape, (0, 4))
+
+    def test_invalid_inputs_raise(self):
+        with self.assertRaises(ValueError):
+            pairwise_mahalanobis_distances(
+                array([[0.0]]),
+                zeros((1, 1, 1)),
+                array([[0.0]]),
+                zeros((1, 1, 1)),
+                regularization=-1.0,
+            )
+
+        with self.assertRaises(ValueError):
+            pairwise_covariance_shape_components(
+                zeros((2, 2, 1)),
+                zeros((2, 2, 1)),
+                epsilon=0.0,
+            )
+
+        with self.assertRaises(ValueError):
+            pairwise_mahalanobis_distances(
+                array([[0.0, 1.0]]),
+                zeros((1, 1, 1)),
+                array([[0.0]]),
+                zeros((1, 1, 1)),
+            )
+
+
+if __name__ == "__main__":
+    unittest.main()

From 01f8179729e7f6f3bdb3c332c7fa342d174c5852 Mon Sep 17 00:00:00 2001
From: Florian Pfaff <6773539+FlorianPfaff@users.noreply.github.com>
Date: Sat, 9 May 2026 02:47:23 +0200
Subject: [PATCH 4/8] Document pairwise covariance feature utilities

---
 docs/api-overview.md | 6 ++++--
 1 file changed, 4 insertions(+), 2 deletions(-)

diff --git a/docs/api-overview.md b/docs/api-overview.md
index a5037df66..6f0d21560 100644
--- a/docs/api-overview.md
+++ b/docs/api-overview.md
@@ -144,13 +144,15 @@ Common starting points include:
 ### `pyrecest.utils`
 
 Reusable utility helpers for assignment, association models, history recording,
-multi-session assignment, generic track-matrix evaluation, and point-set
-registration.
+multi-session assignment, generic track-matrix evaluation, Gaussian/covariance
+pairwise features, and point-set registration.
 
 Common starting points include:
 
 - `murty_k_best_assignments`;
 - `LogisticPairwiseAssociationModel`;
+- `pairwise_mahalanobis_distances`;
+- `pairwise_covariance_shape_components`;
 - `HistoryRecorder`;
 - `solve_multisession_assignment`;
 - `normalize_track_matrix`;

From 493ee99ea7cdd1b2fc46faca857629895e06bf43 Mon Sep 17 00:00:00 2001
From: Florian Pfaff <pfaff@ias.uni-stuttgart.de>
Date: Sat, 9 May 2026 03:17:31 +0200
Subject: [PATCH 5/8] Fix pairwise covariance feature CI failures

---
 pyproject.toml                                |  2 +
 src/pyrecest/filters/_linear_gaussian.py      |  9 ++-
 .../utils/pairwise_covariance_features.py     | 28 ++++----
 tests/filters/test_relaxed_s3f_so3.py         | 70 +++++++++++++++----
 4 files changed, 75 insertions(+), 34 deletions(-)

diff --git a/pyproject.toml b/pyproject.toml
index 9184ee720..8b440fbdd 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -82,6 +82,8 @@ disable = [
   "too-many-branches",
   "too-many-statements",
   "too-many-arguments",
+  "too-many-positional-arguments",
+  "too-many-instance-attributes",
   "duplicate-code",
 ]
 
diff --git a/src/pyrecest/filters/_linear_gaussian.py b/src/pyrecest/filters/_linear_gaussian.py
index e84c37319..16e519273 100644
--- a/src/pyrecest/filters/_linear_gaussian.py
+++ b/src/pyrecest/filters/_linear_gaussian.py
@@ -1,4 +1,4 @@
-# pylint: disable=no-name-in-module,no-member
+# pylint: disable=no-name-in-module,no-member,redefined-outer-name
 """Backend-native linear-Gaussian predict/update primitives."""
 
 from pyrecest.backend import (
@@ -45,8 +45,7 @@ def normalized_innovation_squared(innovation, innovation_covariance):
     innovation_dim = innovation.shape[0]
     if innovation_covariance.shape != (innovation_dim, innovation_dim):
         raise ValueError(
-            "innovation_covariance must have shape "
-            "(innovation_dim, innovation_dim)"
+            "innovation_covariance must have shape (innovation_dim, innovation_dim)"
         )
     return transpose(innovation) @ linalg.solve(innovation_covariance, innovation)
 
@@ -261,8 +260,8 @@ def linear_gaussian_update(
     identity = eye(state_dim)
     correction = identity - kalman_gain @ measurement_matrix
     updated_covariance = correction @ covariance @ transpose(correction)
-    updated_covariance = updated_covariance + kalman_gain @ scaled_meas_noise @ transpose(
-        kalman_gain
+    updated_covariance = (
+        updated_covariance + kalman_gain @ scaled_meas_noise @ transpose(kalman_gain)
     )
     updated_covariance = 0.5 * (updated_covariance + transpose(updated_covariance))
 
diff --git a/src/pyrecest/utils/pairwise_covariance_features.py b/src/pyrecest/utils/pairwise_covariance_features.py
index 80ab624b6..03f677882 100644
--- a/src/pyrecest/utils/pairwise_covariance_features.py
+++ b/src/pyrecest/utils/pairwise_covariance_features.py
@@ -19,25 +19,29 @@
 # pylint: disable=no-name-in-module,no-member,redefined-builtin
 from pyrecest.backend import (
     abs,
-    all as backend_all,
+)
+from pyrecest.backend import all as backend_all
+from pyrecest.backend import (
     asarray,
     einsum,
     exp,
     eye,
     float64,
     isfinite,
+    linalg,
     log,
     maximum,
     moveaxis,
     reshape,
     sqrt,
-    sum as backend_sum,
+)
+from pyrecest.backend import sum as backend_sum
+from pyrecest.backend import (
     trace,
     transpose,
     where,
     zeros,
 )
-from pyrecest.backend import linalg
 
 
 def pairwise_mahalanobis_distances(
@@ -89,9 +93,7 @@ def pairwise_mahalanobis_distances(
     dim, n_a = means_a.shape
     dim_b, n_b = means_b.shape
     if dim != dim_b:
-        raise ValueError(
-            "means_a and means_b must have the same leading dimension"
-        )
+        raise ValueError("means_a and means_b must have the same leading dimension")
 
     if n_a == 0 or n_b == 0:
         return zeros((n_a, n_b), dtype=float64)
@@ -102,9 +104,9 @@ def pairwise_mahalanobis_distances(
         moved_covariances_a[:, None, :, :] + moved_covariances_b[None, :, :, :]
     )
     if regularization > 0.0:
-        covariance_sums = covariance_sums + float(regularization) * eye(dim)[
-            None, None, :, :
-        ]
+        covariance_sums = (
+            covariance_sums + float(regularization) * eye(dim)[None, None, :, :]
+        )
 
     mean_differences = transpose(means_a)[:, None, :] - transpose(means_b)[None, :, :]
     flat_differences = reshape(mean_differences, (n_a * n_b, dim))
@@ -163,12 +165,8 @@ def pairwise_covariance_shape_components(
     moved_covariances_a = _symmetrized_covariance_batch(covariances_a)
     moved_covariances_b = _symmetrized_covariance_batch(covariances_b)
 
-    traces_a = _positive_floor(
-        trace(moved_covariances_a, axis1=1, axis2=2), epsilon
-    )
-    traces_b = _positive_floor(
-        trace(moved_covariances_b, axis1=1, axis2=2), epsilon
-    )
+    traces_a = _positive_floor(trace(moved_covariances_a), epsilon)
+    traces_b = _positive_floor(trace(moved_covariances_b), epsilon)
     normalized_a = moved_covariances_a / traces_a[:, None, None]
     normalized_b = moved_covariances_b / traces_b[:, None, None]
 
diff --git a/tests/filters/test_relaxed_s3f_so3.py b/tests/filters/test_relaxed_s3f_so3.py
index b6ba7255c..7e445f870 100644
--- a/tests/filters/test_relaxed_s3f_so3.py
+++ b/tests/filters/test_relaxed_s3f_so3.py
@@ -7,7 +7,9 @@
 from pyrecest.distributions.hypersphere_subset.hyperhemispherical_grid_distribution import (
     HyperhemisphericalGridDistribution,
 )
-from pyrecest.distributions.nonperiodic.gaussian_distribution import GaussianDistribution
+from pyrecest.distributions.nonperiodic.gaussian_distribution import (
+    GaussianDistribution,
+)
 from pyrecest.filters.relaxed_s3f_so3 import (
     _cached_s3r3_cell_statistics,
     predict_s3r3_relaxed,
@@ -39,10 +41,15 @@ def test_cell_statistics_reuses_identical_grid_cache(self):
         _cached_s3r3_cell_statistics.cache_clear()
         grid = _small_quaternion_grid()
 
-        stats_a = s3r3_cell_statistics(grid, array([0.4, 0.1, 0.2]), cell_sample_count=27)
-        stats_b = s3r3_cell_statistics(array(grid, dtype=float), array([0.4, 0.1, 0.2]), cell_sample_count=27)
+        stats_a = s3r3_cell_statistics(
+            grid, array([0.4, 0.1, 0.2]), cell_sample_count=27
+        )
+        stats_b = s3r3_cell_statistics(
+            array(grid, dtype=float), array([0.4, 0.1, 0.2]), cell_sample_count=27
+        )
 
         self.assertIs(stats_a, stats_b)
+        # pylint: disable-next=no-value-for-parameter
         cache_info = _cached_s3r3_cell_statistics.cache_info()
         self.assertEqual(cache_info.misses, 1)
         self.assertEqual(cache_info.hits, 1)
@@ -59,9 +66,20 @@ def test_prediction_conserves_grid_mass_and_inflates_covariance(self):
             cell_sample_count=27,
         )
 
-        self.assertTrue(bool(allclose(filter_.filter_state.gd.grid_values, values_before, atol=1e-12)))
-        self.assertTrue(bool(allclose(float(filter_.filter_state.gd.integrate()), 1.0, atol=1e-12)))
-        self.assertTrue(any(float(linalg.norm(covariance)) > 0.0 for covariance in stats.covariance_inflations))
+        self.assertTrue(
+            bool(
+                allclose(filter_.filter_state.gd.grid_values, values_before, atol=1e-12)
+            )
+        )
+        self.assertTrue(
+            bool(allclose(float(filter_.filter_state.gd.integrate()), 1.0, atol=1e-12))
+        )
+        self.assertTrue(
+            any(
+                float(linalg.norm(covariance)) > 0.0
+                for covariance in stats.covariance_inflations
+            )
+        )
 
     def test_baseline_and_r1_variants_use_expected_inputs(self):
         body_increment = array([0.4, 0.0, 0.0])
@@ -83,8 +101,18 @@ def test_baseline_and_r1_variants_use_expected_inputs(self):
 
         baseline_mean0 = baseline_filter.filter_state.linear_distributions[0].mu
         r1_mean0 = r1_filter.filter_state.linear_distributions[0].mu
-        self.assertTrue(bool(allclose(baseline_mean0, baseline_stats.representative_displacements[0], atol=1e-12)))
-        self.assertTrue(bool(allclose(r1_mean0, r1_stats.mean_displacements[0], atol=1e-12)))
+        self.assertTrue(
+            bool(
+                allclose(
+                    baseline_mean0,
+                    baseline_stats.representative_displacements[0],
+                    atol=1e-12,
+                )
+            )
+        )
+        self.assertTrue(
+            bool(allclose(r1_mean0, r1_stats.mean_displacements[0], atol=1e-12))
+        )
         self.assertFalse(bool(allclose(baseline_mean0, r1_mean0, atol=1e-14)))
 
     def test_quaternion_rotation_and_distance_helpers(self):
@@ -94,18 +122,32 @@ def test_quaternion_rotation_and_distance_helpers(self):
         rotated = rotate_quaternion_body_increment(identity, array([0.4, 0.1, 0.2]))
 
         self.assertTrue(bool(allclose(rotated[0], array([0.4, 0.1, 0.2]), atol=1e-12)))
-        self.assertAlmostEqual(s3r3_orientation_distance(identity, -identity), 0.0, places=12)
-        self.assertAlmostEqual(s3r3_orientation_distance(identity, half_turn_z), 3.141592653589793, places=12)
+        self.assertAlmostEqual(
+            s3r3_orientation_distance(identity, -identity), 0.0, places=12
+        )
+        self.assertAlmostEqual(
+            s3r3_orientation_distance(identity, half_turn_z),
+            3.141592653589793,
+            places=12,
+        )
 
     def test_validation_errors_are_explicit(self):
         with self.assertRaisesRegex(ValueError, "method"):
-            s3r3_cell_statistics(_small_quaternion_grid(), array([0.4, 0.1, 0.2]), method="exact_voronoi")
+            s3r3_cell_statistics(
+                _small_quaternion_grid(), array([0.4, 0.1, 0.2]), method="exact_voronoi"
+            )
         with self.assertRaisesRegex(ValueError, "cell_sample_count"):
-            s3r3_cell_statistics(_small_quaternion_grid(), array([0.4, 0.1, 0.2]), cell_sample_count=0)
+            s3r3_cell_statistics(
+                _small_quaternion_grid(), array([0.4, 0.1, 0.2]), cell_sample_count=0
+            )
         with self.assertRaisesRegex(ValueError, "body_increment"):
-            s3r3_cell_statistics(_small_quaternion_grid(), array([0.4, 0.1]), cell_sample_count=27)
+            s3r3_cell_statistics(
+                _small_quaternion_grid(), array([0.4, 0.1]), cell_sample_count=27
+            )
         with self.assertRaisesRegex(ValueError, "variant"):
-            predict_s3r3_relaxed(_make_filter(), array([0.4, 0.1, 0.2]), variant="r2_only")
+            predict_s3r3_relaxed(
+                _make_filter(), array([0.4, 0.1, 0.2]), variant="r2_only"
+            )
 
 
 def _make_filter() -> StateSpaceSubdivisionFilter:

From 54c9bc62910d489f2dc3c755afe3e695c537a543 Mon Sep 17 00:00:00 2001
From: Florian Pfaff <pfaff@ias.uni-stuttgart.de>
Date: Sat, 9 May 2026 03:31:04 +0200
Subject: [PATCH 6/8] Normalize SO3 grid transition columns

---
 .../sd_half_cond_sd_half_grid_distribution.py         |  4 +++-
 src/pyrecest/filters/so3_grid_transition.py           | 11 +++++++----
 2 files changed, 10 insertions(+), 5 deletions(-)

diff --git a/src/pyrecest/distributions/conditional/sd_half_cond_sd_half_grid_distribution.py b/src/pyrecest/distributions/conditional/sd_half_cond_sd_half_grid_distribution.py
index 6d486305b..ae92fb677 100644
--- a/src/pyrecest/distributions/conditional/sd_half_cond_sd_half_grid_distribution.py
+++ b/src/pyrecest/distributions/conditional/sd_half_cond_sd_half_grid_distribution.py
@@ -57,9 +57,11 @@ def __init__(self, grid, grid_values, enforce_pdf_nonnegative=True):
 
     def _check_normalization(self, tol=0.01):
         """Warn if any column is not normalized to 1 over the hemisphere."""
+        # Match HyperhemisphericalGridDistribution.get_manifold_size(), which
+        # uses the grid embedding dimension for the existing quadrature rule.
         hemisphere_surface = 0.5 * (
             AbstractHypersphereSubsetDistribution.compute_unit_hypersphere_surface(
-                self.grid.shape[1] - 1
+                self.grid.shape[1]
             )
         )
         ints = mean(self.grid_values, axis=0) * hemisphere_surface
diff --git a/src/pyrecest/filters/so3_grid_transition.py b/src/pyrecest/filters/so3_grid_transition.py
index d2d200f92..b7ca25c19 100644
--- a/src/pyrecest/filters/so3_grid_transition.py
+++ b/src/pyrecest/filters/so3_grid_transition.py
@@ -67,8 +67,9 @@ def so3_right_multiplication_grid_transition(
 
     ``exp(kappa * |<q_next, q_current * delta_q>|**2)``.
 
-    The columns are normalized by the S3 upper-hemisphere grid quadrature rule,
-    so ``mean(grid_values[:, j]) * surface(S3+) == 1`` for every column.
+    The columns are normalized by the hyperhemispherical grid quadrature rule
+    used by :class:`HyperhemisphericalGridFilter`, so
+    ``mean(grid_values[:, j]) * manifold_size == 1`` for every column.
     """
 
     if kappa <= 0.0:
@@ -88,7 +89,7 @@ def so3_right_multiplication_grid_transition(
     manifold_size = (
         0.5
         * AbstractHypersphereSubsetDistribution.compute_unit_hypersphere_surface(
-            quaternion_grid.shape[1] - 1
+            quaternion_grid.shape[1]
         )
     )
     column_integrals = mean(scores, axis=0, keepdims=True) * manifold_size
@@ -121,7 +122,9 @@ def _as_quaternion_grid(grid):
 
     quaternion_grid = array(grid, dtype=float)
     if ndim(quaternion_grid) != 2 or quaternion_grid.shape[1] != 4:
-        raise ValueError("grid must have shape (n_grid, 4) with scalar-last quaternions.")
+        raise ValueError(
+            "grid must have shape (n_grid, 4) with scalar-last quaternions."
+        )
     if quaternion_grid.shape[0] == 0:
         raise ValueError("grid must contain at least one quaternion.")
     if not all(linalg.norm(quaternion_grid, axis=1) > 0.0):

From 7c648be5a56b84325015113265b84f7e3e6a9abd Mon Sep 17 00:00:00 2001
From: Florian Pfaff <pfaff@ias.uni-stuttgart.de>
Date: Sat, 9 May 2026 09:02:49 +0200
Subject: [PATCH 7/8] Fix backend compatibility for covariance features

---
 src/pyrecest/filters/so3_grid_transition.py     |  9 ++++++---
 .../filters/so3_product_particle_filter.py      | 17 +++++++++++------
 .../utils/pairwise_covariance_features.py       |  9 ++++++---
 3 files changed, 23 insertions(+), 12 deletions(-)

diff --git a/src/pyrecest/filters/so3_grid_transition.py b/src/pyrecest/filters/so3_grid_transition.py
index b7ca25c19..4c588c630 100644
--- a/src/pyrecest/filters/so3_grid_transition.py
+++ b/src/pyrecest/filters/so3_grid_transition.py
@@ -9,8 +9,9 @@
     clip,
     exp,
     linalg,
-    mean,
     ndim,
+    reshape,
+    sum,
     transpose,
 )
 from pyrecest.distributions._so3_helpers import (
@@ -84,7 +85,8 @@ def so3_right_multiplication_grid_transition(
     # Subtracting the per-column maximum keeps the normalization stable for
     # large kappa without changing the normalized conditional density.
     exponents = kappa * inner_products**2
-    scores = exp(exponents - amax(exponents, axis=0, keepdims=True))
+    column_maxima = reshape(amax(exponents, axis=0), (1, exponents.shape[1]))
+    scores = exp(exponents - column_maxima)
 
     manifold_size = (
         0.5
@@ -92,7 +94,8 @@ def so3_right_multiplication_grid_transition(
             quaternion_grid.shape[1]
         )
     )
-    column_integrals = mean(scores, axis=0, keepdims=True) * manifold_size
+    column_integrals = sum(scores, axis=0, keepdims=True) / scores.shape[0]
+    column_integrals = column_integrals * manifold_size
     density_values = scores / column_integrals
 
     return SdHalfCondSdHalfGridDistribution(
diff --git a/src/pyrecest/filters/so3_product_particle_filter.py b/src/pyrecest/filters/so3_product_particle_filter.py
index e38808503..9c800a52b 100644
--- a/src/pyrecest/filters/so3_product_particle_filter.py
+++ b/src/pyrecest/filters/so3_product_particle_filter.py
@@ -1,5 +1,6 @@
 """Particle filter for Cartesian products of SO(3)."""
 
+import math
 from collections.abc import Callable
 
 # pylint: disable=no-name-in-module,no-member
@@ -157,7 +158,7 @@ def _normalize_log_weights(log_weights):
     @staticmethod
     def _validate_probability(name: str, value: float) -> float:
         value = float(value)
-        if not isfinite(value) or value < 0.0 or value > 1.0:
+        if not math.isfinite(value) or value < 0.0 or value > 1.0:
             raise ValueError(f"{name} must be a finite probability in [0, 1].")
         return value
 
@@ -198,7 +199,9 @@ def _as_component_noise_std(noise_std, component_noise_std, num_rotations):
 
         if ndim(sigma) == 0:
             if not isfinite(sigma) or sigma <= 0.0:
-                raise ValueError("noise standard deviations must be positive and finite.")
+                raise ValueError(
+                    "noise standard deviations must be positive and finite."
+                )
             return ones(num_rotations) * sigma
         if sigma.shape != (num_rotations,):
             raise ValueError("component_noise_std must have shape (num_rotations,).")
@@ -235,13 +238,15 @@ def confidence_to_noise_std(
         min_sigma = float(noise_std)
         max_sigma = float(max_noise_std)
         exponent = float(confidence_exponent)
-        if min_sigma <= 0.0 or not isfinite(min_sigma):
+        if min_sigma <= 0.0 or not math.isfinite(min_sigma):
             raise ValueError("noise_std must be positive and finite.")
-        if max_sigma <= 0.0 or not isfinite(max_sigma):
+        if max_sigma <= 0.0 or not math.isfinite(max_sigma):
             raise ValueError("max_noise_std must be positive and finite.")
         if max_sigma < min_sigma:
-            raise ValueError("max_noise_std must be greater than or equal to noise_std.")
-        if exponent <= 0.0 or not isfinite(exponent):
+            raise ValueError(
+                "max_noise_std must be greater than or equal to noise_std."
+            )
+        if exponent <= 0.0 or not math.isfinite(exponent):
             raise ValueError("confidence_exponent must be positive and finite.")
 
         variance = min_sigma * min_sigma + (1.0 - confidence) ** exponent * (
diff --git a/src/pyrecest/utils/pairwise_covariance_features.py b/src/pyrecest/utils/pairwise_covariance_features.py
index 03f677882..306d43b46 100644
--- a/src/pyrecest/utils/pairwise_covariance_features.py
+++ b/src/pyrecest/utils/pairwise_covariance_features.py
@@ -37,7 +37,6 @@
 )
 from pyrecest.backend import sum as backend_sum
 from pyrecest.backend import (
-    trace,
     transpose,
     where,
     zeros,
@@ -165,8 +164,8 @@ def pairwise_covariance_shape_components(
     moved_covariances_a = _symmetrized_covariance_batch(covariances_a)
     moved_covariances_b = _symmetrized_covariance_batch(covariances_b)
 
-    traces_a = _positive_floor(trace(moved_covariances_a), epsilon)
-    traces_b = _positive_floor(trace(moved_covariances_b), epsilon)
+    traces_a = _positive_floor(_batch_trace(moved_covariances_a), epsilon)
+    traces_b = _positive_floor(_batch_trace(moved_covariances_b), epsilon)
     normalized_a = moved_covariances_a / traces_a[:, None, None]
     normalized_b = moved_covariances_b / traces_b[:, None, None]
 
@@ -221,6 +220,10 @@ def _symmetrized_covariance_batch(covariances: Any) -> Any:
     return 0.5 * (moved + transpose(moved, (0, 2, 1)))
 
 
+def _batch_trace(matrices: Any) -> Any:
+    return einsum("nii->n", matrices)
+
+
 def _positive_floor(values: Any, epsilon: float) -> Any:
     values = asarray(values, dtype=float64)
     return where(values > float(epsilon), values, float(epsilon))

From 63299a60930a756110e9efb34bfbe7b619d7fa59 Mon Sep 17 00:00:00 2001
From: Florian Pfaff <pfaff@ias.uni-stuttgart.de>
Date: Sat, 9 May 2026 09:27:54 +0200
Subject: [PATCH 8/8] Relax JAX precision assertions

---
 tests/filters/test_relaxed_s3f_so3.py  | 8 +++++---
 tests/test_affine_transform_algebra.py | 9 ++++++---
 2 files changed, 11 insertions(+), 6 deletions(-)

diff --git a/tests/filters/test_relaxed_s3f_so3.py b/tests/filters/test_relaxed_s3f_so3.py
index 7e445f870..234fe18ea 100644
--- a/tests/filters/test_relaxed_s3f_so3.py
+++ b/tests/filters/test_relaxed_s3f_so3.py
@@ -1,6 +1,6 @@
 import unittest
 
-from pyrecest.backend import allclose, array, eye, linalg, ones, zeros
+from pyrecest.backend import __backend_name__, allclose, array, eye, linalg, ones, zeros
 from pyrecest.distributions.cart_prod.state_space_subdivision_gaussian_distribution import (
     StateSpaceSubdivisionGaussianDistribution,
 )
@@ -19,6 +19,8 @@
 )
 from pyrecest.filters.state_space_subdivision_filter import StateSpaceSubdivisionFilter
 
+_JAX_ATOL = 1e-6 if __backend_name__ == "jax" else 1e-12
+
 
 class RelaxedS3FSO3Test(unittest.TestCase):
     def test_covariance_inflation_is_positive_semidefinite(self):
@@ -123,12 +125,12 @@ def test_quaternion_rotation_and_distance_helpers(self):
 
         self.assertTrue(bool(allclose(rotated[0], array([0.4, 0.1, 0.2]), atol=1e-12)))
         self.assertAlmostEqual(
-            s3r3_orientation_distance(identity, -identity), 0.0, places=12
+            s3r3_orientation_distance(identity, -identity), 0.0, delta=_JAX_ATOL
         )
         self.assertAlmostEqual(
             s3r3_orientation_distance(identity, half_turn_z),
             3.141592653589793,
-            places=12,
+            delta=_JAX_ATOL,
         )
 
     def test_validation_errors_are_explicit(self):
diff --git a/tests/test_affine_transform_algebra.py b/tests/test_affine_transform_algebra.py
index 2ed2e21bc..5e9ccb24f 100644
--- a/tests/test_affine_transform_algebra.py
+++ b/tests/test_affine_transform_algebra.py
@@ -1,9 +1,12 @@
 import unittest
 
 import numpy.testing as npt
-from pyrecest.backend import array
+
+from pyrecest.backend import __backend_name__, array
 from pyrecest.utils.point_set_registration import AffineTransform
 
+_ATOL = 1e-6 if __backend_name__ == "jax" else 1e-10
+
 
 class TestAffineTransformAlgebra(unittest.TestCase):
     def test_inverse_round_trips_points(self):
@@ -15,7 +18,7 @@ def test_inverse_round_trips_points(self):
 
         recovered = transform.inverse().apply(transform.apply(points))
 
-        npt.assert_allclose(recovered, points, atol=1e-10)
+        npt.assert_allclose(recovered, points, atol=_ATOL)
 
     def test_compose_matches_sequential_application(self):
         points = array([[0.0, 0.0], [1.0, -2.0], [3.0, 4.0]])
@@ -33,7 +36,7 @@ def test_compose_matches_sequential_application(self):
         npt.assert_allclose(
             composed.apply(points),
             second.apply(first.apply(points)),
-            atol=1e-10,
+            atol=_ATOL,
         )
 
     def test_compose_rejects_dimension_mismatch(self):