Added pairwise association-cost for GNN (#1739)

* Added pairwise association-cost for GNN

* Fix pylint/flake8 issues: remove unused imports, add pylint disables

Agent-Logs-Url: https://github.com/FlorianPfaff/PyRecEst/sessions/3594bc92-b3a4-4c14-b600-7754a7f5398d

Co-authored-by: FlorianPfaff <6773539+FlorianPfaff@users.noreply.github.com>

* Suppress duplicate-code (R0801) in GNN filter and pairwise-cost test

Agent-Logs-Url: https://github.com/FlorianPfaff/PyRecEst/sessions/37f5b6c8-c759-4657-8c8f-3a514ce88de3

Co-authored-by: FlorianPfaff <6773539+FlorianPfaff@users.noreply.github.com>

---------

Co-authored-by: copilot-swe-agent[bot] <198982749+Copilot@users.noreply.github.com>
Co-authored-by: FlorianPfaff <6773539+FlorianPfaff@users.noreply.github.com>

Author: FlorianPfaff

pyrecest/filters/global_nearest_neighbor.py

@@ -1,4 +1,7 @@
-# pylint: disable=redefined-builtin,no-name-in-module,no-member
+# pylint: disable=redefined-builtin,no-name-in-module,no-member,duplicate-code
+import warnings
+
+import pyrecest.backend
 from pyrecest.backend import all, empty, full, repeat, squeeze, stack
 from scipy.optimize import linear_sum_assignment
 from scipy.spatial.distance import cdist
@@ -8,19 +11,36 @@
 
 
 class GlobalNearestNeighbor(AbstractNearestNeighborTracker):
+    """Global nearest-neighbor tracker for linear/Gaussian multitarget tracking.
+
+    Besides the built-in geometric association costs, this implementation can
+    optionally fuse an externally computed ``pairwise_cost_matrix`` of shape
+    ``(n_targets, n_meas)``. This is useful for domains such as longitudinal
+    calcium-imaging cell tracking where association should depend on arbitrary
+    pairwise cues like ROI overlap, footprint correlation, or appearance
+    embeddings in addition to centroid distance.
+    """
+
     def __init__(
         self,
         initial_prior=None,
         association_param=None,
         log_prior_estimates=True,
         log_posterior_estimates=True,
     ):
+        default_association_param = {
+            "distance_metric_pos": "Mahalanobis",
+            "square_dist": True,
+            "max_new_tracks": 10,
+            "gating_distance_threshold": chi2.ppf(0.999, 2) ** 2,
+            "pairwise_cost_weight": 1.0,
+        }
         if association_param is None:
+            association_param = default_association_param
+        else:
             association_param = {
-                "distance_metric_pos": "Mahalanobis",
-                "square_dist": True,
-                "max_new_tracks": 10,
-                "gating_distance_threshold": chi2.ppf(0.999, 2) ** 2,
+                **default_association_param,
+                **association_param,
             }
 
         super().__init__(
@@ -30,14 +50,51 @@ def __init__(
             log_posterior_estimates=log_posterior_estimates,
         )
 
-    # pylint: disable=too-many-locals
+    @staticmethod
+    def _validate_pairwise_cost_matrix(pairwise_cost_matrix, n_targets, n_meas):
+        if pairwise_cost_matrix is None:
+            return None
+        if pairwise_cost_matrix.shape != (n_targets, n_meas):
+            raise ValueError(
+                "pairwise_cost_matrix must have shape "
+                f"({n_targets}, {n_meas}), got {pairwise_cost_matrix.shape}."
+            )
+        return pairwise_cost_matrix
+
+    def _apply_pairwise_cost_matrix(self, dists, pairwise_cost_matrix):
+        if pairwise_cost_matrix is None:
+            return dists
+        pairwise_cost_weight = self.association_param.get("pairwise_cost_weight", 1.0)
+        if pairwise_cost_weight == 0.0:
+            return dists
+        return dists + pairwise_cost_weight * pairwise_cost_matrix
+
+    # pylint: disable=too-many-locals,too-many-positional-arguments
     def find_association(
         self,
         measurements,
         measurement_matrix,
         cov_mats_meas,
         warn_on_no_meas_for_track=True,
+        pairwise_cost_matrix=None,
     ):
+        """Find the minimum-cost measurement-to-track assignment.
+
+        Parameters
+        ----------
+        measurements : array-like, shape (dim_meas, n_meas)
+            Measurements for the current update step.
+        measurement_matrix : array-like
+            Linear measurement model.
+        cov_mats_meas : array-like
+            Measurement covariance matrix or per-measurement covariance tensor.
+        warn_on_no_meas_for_track : bool, optional
+            Whether to emit a warning when a track remains unassigned.
+        pairwise_cost_matrix : array-like, optional
+            Additional target/measurement association costs of shape
+            ``(n_targets, n_meas)``. These costs are added to the geometric cost
+            matrix before running the Hungarian algorithm.
+        """
         n_targets = len(self.filter_bank)
         n_meas = measurements.shape[1]
 
@@ -71,11 +128,14 @@ def find_association(
             )
 
             if all_cov_mat_meas_equal and all_cov_mat_state_equal:
+                shared_cov_mats_meas = (
+                    cov_mats_meas if cov_mats_meas.ndim == 2 else cov_mats_meas[:, :, 0]
+                )
                 curr_cov_mahalanobis = (
                     measurement_matrix
                     @ all_cov_mats_prior[:, :, 0]
                     @ measurement_matrix.T
-                    + cov_mats_meas[:, :, 0]
+                    + shared_cov_mats_meas
                 )
                 dists = cdist(
                     (measurement_matrix @ all_means_prior).T,
@@ -84,6 +144,9 @@ def find_association(
                     VI=curr_cov_mahalanobis,
                 )
             elif all_cov_mat_meas_equal:
+                shared_cov_mats_meas = (
+                    cov_mats_meas if cov_mats_meas.ndim == 2 else cov_mats_meas[:, :, 0]
+                )
                 all_mats_mahalanobis = empty(
                     (
                         measurements.shape[0],
@@ -96,7 +159,7 @@ def find_association(
                         measurement_matrix
                         @ all_cov_mats_prior[:, :, i]
                         @ measurement_matrix.T
-                        + cov_mats_meas
+                        + shared_cov_mats_meas
                     )
                 for i in range(n_targets):
                     dists[i, :] = cdist(
@@ -125,6 +188,11 @@ def find_association(
         else:
             raise ValueError("Association scheme not recognized")
 
+        pairwise_cost_matrix = self._validate_pairwise_cost_matrix(
+            pairwise_cost_matrix, n_targets, n_meas
+        )
+        dists = self._apply_pairwise_cost_matrix(dists, pairwise_cost_matrix)
+
         # Pad to square and add max_new_tracks rows and columns
         pad_to = max(n_targets, n_meas) + self.association_param["max_new_tracks"]
         association_matrix = full(
@@ -137,9 +205,48 @@ def find_association(
 
         association = col_ind[:n_targets]
 
-        if warn_on_no_meas_for_track and any(association > n_meas):
+        if warn_on_no_meas_for_track and any(association >= n_meas):
             print(
                 "GNN: No measurement was within gating threshold for at least one target."
             )
 
         return association
+
+    def update_linear(
+        self,
+        measurements,
+        measurement_matrix,
+        covMatsMeas,
+        pairwise_cost_matrix=None,
+    ):
+        """Update the tracker with an optional additional association cost matrix."""
+        assert (
+            pyrecest.backend.__backend_name__ == "numpy"
+        ), "Only supported for numpy backend"
+        if len(self.filter_bank) == 0:
+            warnings.warn("Currently, there are zero targets")
+            return
+        assert (
+            measurement_matrix.shape[0] == measurements.shape[0]
+            and measurement_matrix.shape[1]
+            == self.filter_bank[0].get_point_estimate().shape[0]
+        ), (
+            "Dimensions of measurement matrix must match state and measurement dimensions."
+        )
+        association = self.find_association(
+            measurements,
+            measurement_matrix,
+            covMatsMeas,
+            pairwise_cost_matrix=pairwise_cost_matrix,
+        )
+        currMeasCov = covMatsMeas
+        for i in range(self.get_number_of_targets()):
+            if association[i] < measurements.shape[1]:
+                if covMatsMeas.ndim != 2:
+                    currMeasCov = covMatsMeas[:, :, association[i]]
+                self.filter_bank[i].update_linear(
+                    measurements[:, association[i]], measurement_matrix, currMeasCov
+                )
+
+        if self.log_posterior_estimates:
+            self.store_posterior_estimates()

pyrecest/tests/filters/test_global_nearest_neighbor_pairwise_costs.py

@@ -0,0 +1,148 @@
+import unittest
+
+import numpy.testing as npt
+
+# pylint: disable=no-member,duplicate-code
+import pyrecest.backend
+
+from pyrecest.backend import array, column_stack, diag, eye, zeros
+from pyrecest.distributions import GaussianDistribution
+from pyrecest.filters import KalmanFilter
+from pyrecest.filters.global_nearest_neighbor import GlobalNearestNeighbor
+
+
+class GlobalNearestNeighborPairwiseCostTest(unittest.TestCase):
+    """Regression tests for pairwise-cost fusion in GNN."""
+
+    def setUp(self):
+        self.kfs_init = [
+            KalmanFilter(
+                GaussianDistribution(zeros(4), diag(array([1.0, 2.0, 3.0, 4.0])))
+            ),
+            KalmanFilter(
+                GaussianDistribution(
+                    array([1.0, 2.0, 3.0, 4.0]), diag(array([2.0, 2.0, 2.0, 2.0]))
+                )
+            ),
+            KalmanFilter(
+                GaussianDistribution(
+                    -array([1.0, 2.0, 3.0, 4.0]), diag(array([4.0, 3.0, 2.0, 1.0]))
+                )
+            ),
+        ]
+        self.meas_mat = array([[1.0, 0.0, 0.0, 0.0], [0.0, 0.0, 1.0, 0.0]])
+
+    @unittest.skipIf(
+        pyrecest.backend.__backend_name__ in ("pytorch", "jax"),
+        reason="Not supported on this backend",
+    )
+    def test_association_supports_shared_2d_measurement_covariance(self):
+        tracker = GlobalNearestNeighbor()
+        tracker.filter_state = [
+            KalmanFilter(
+                GaussianDistribution(zeros(4), diag(array([1.0, 2.0, 3.0, 4.0])))
+            ),
+            KalmanFilter(
+                GaussianDistribution(
+                    array([1.0, 2.0, 3.0, 4.0]), diag(array([1.0, 2.0, 3.0, 4.0]))
+                )
+            ),
+            KalmanFilter(
+                GaussianDistribution(
+                    -array([1.0, 2.0, 3.0, 4.0]), diag(array([1.0, 2.0, 3.0, 4.0]))
+                )
+            ),
+        ]
+        all_gaussians = [kf.filter_state for kf in tracker.filter_bank]
+        perfect_meas_ordered = self.meas_mat @ column_stack(
+            [gaussian.mu for gaussian in all_gaussians]
+        )
+        association = tracker.find_association(
+            perfect_meas_ordered, self.meas_mat, eye(2)
+        )
+        npt.assert_array_equal(association, [0, 1, 2])
+
+    @unittest.skipIf(
+        pyrecest.backend.__backend_name__ in ("pytorch", "jax"),
+        reason="Not supported on this backend",
+    )
+    def test_pairwise_cost_matrix_can_override_geometric_assignment(self):
+        tracker = GlobalNearestNeighbor()
+        tracker.filter_state = self.kfs_init
+        all_gaussians = [kf.filter_state for kf in self.kfs_init]
+        perfect_meas_ordered = self.meas_mat @ column_stack(
+            [gaussian.mu for gaussian in all_gaussians]
+        )
+
+        pairwise_cost_matrix = array(
+            [[10.0, -10.0, 10.0], [10.0, 10.0, -10.0], [-10.0, 10.0, 10.0]]
+        )
+        association = tracker.find_association(
+            perfect_meas_ordered,
+            self.meas_mat,
+            eye(2),
+            pairwise_cost_matrix=pairwise_cost_matrix,
+        )
+        npt.assert_array_equal(association, [1, 2, 0])
+
+    @unittest.skipIf(
+        pyrecest.backend.__backend_name__ in ("pytorch", "jax"),
+        reason="Not supported on this backend",
+    )
+    def test_pairwise_cost_matrix_shape_is_validated(self):
+        tracker = GlobalNearestNeighbor()
+        tracker.filter_state = self.kfs_init
+        all_gaussians = [kf.filter_state for kf in self.kfs_init]
+        perfect_meas_ordered = self.meas_mat @ column_stack(
+            [gaussian.mu for gaussian in all_gaussians]
+        )
+
+        with self.assertRaises(ValueError):
+            tracker.find_association(
+                perfect_meas_ordered,
+                self.meas_mat,
+                eye(2),
+                pairwise_cost_matrix=zeros((2, 3)),
+            )
+
+    @unittest.skipIf(
+        pyrecest.backend.__backend_name__ in ("pytorch", "jax"),
+        reason="Not supported on this backend",
+    )
+    def test_update_linear_accepts_pairwise_cost_matrix(self):
+        tracker_manual = GlobalNearestNeighbor()
+        tracker_pairwise = GlobalNearestNeighbor()
+        tracker_manual.filter_state = self.kfs_init
+        tracker_pairwise.filter_state = self.kfs_init
+
+        all_gaussians = [kf.filter_state for kf in self.kfs_init]
+        perfect_meas_ordered = self.meas_mat @ column_stack(
+            [gaussian.mu for gaussian in all_gaussians]
+        )
+        forced_permutation = [1, 2, 0]
+
+        for track_index, meas_index in enumerate(forced_permutation):
+            tracker_manual.filter_bank[track_index].update_linear(
+                perfect_meas_ordered[:, meas_index],
+                self.meas_mat,
+                eye(2),
+            )
+
+        pairwise_cost_matrix = array(
+            [[10.0, -10.0, 10.0], [10.0, 10.0, -10.0], [-10.0, 10.0, 10.0]]
+        )
+        tracker_pairwise.update_linear(
+            perfect_meas_ordered,
+            self.meas_mat,
+            eye(2),
+            pairwise_cost_matrix=pairwise_cost_matrix,
+        )
+
+        npt.assert_allclose(
+            tracker_pairwise.get_point_estimate(),
+            tracker_manual.get_point_estimate(),
+        )
+
+
+if __name__ == "__main__":
+    unittest.main()