Fix PCA metric recomputation when "soft" merging

src/spikeinterface/metrics/quality/quality_metrics.py

@@ -156,11 +156,14 @@ def _prepare_data(self, sorting_analyzer, unit_ids=None):
         if pca_ext is None:
             return tmp_data
 
-        if unit_ids is None:
-            unit_ids = sorting_analyzer.unit_ids
+        target_unit_ids = sorting_analyzer.unit_ids if unit_ids is None else unit_ids
+        context_unit_ids = sorting_analyzer.unit_ids
 
-        # Get dense PCA projections for all requested units
-        dense_projections, spike_unit_indices = pca_ext.get_some_projections(channel_ids=None, unit_ids=unit_ids)
+        # Unit-level PCA metrics require neighboring units as comparison data, even when
+        # only a subset of metric rows is being recomputed after a merge or split.
+        dense_projections, spike_unit_indices = pca_ext.get_some_projections(
+            channel_ids=None, unit_ids=context_unit_ids
+        )
         all_labels = sorting_analyzer.sorting.unit_ids[spike_unit_indices]
 
         # Get extremum channels for neighbor selection in sparse mode
@@ -170,22 +173,22 @@ def _prepare_data(self, sorting_analyzer, unit_ids=None):
         # Pre-compute spike counts and firing rates if advanced NN metrics are requested
         advanced_nn_metrics = ["nn_advanced"]  # Our grouped advanced NN metric
         if any(m in advanced_nn_metrics for m in requested_pca_metrics):
-            tmp_data["n_spikes_all_units"] = compute_num_spikes(sorting_analyzer, unit_ids=unit_ids)
-            tmp_data["fr_all_units"] = compute_firing_rates(sorting_analyzer, unit_ids=unit_ids)
+            tmp_data["n_spikes_all_units"] = compute_num_spikes(sorting_analyzer, unit_ids=context_unit_ids)
+            tmp_data["fr_all_units"] = compute_firing_rates(sorting_analyzer, unit_ids=context_unit_ids)
 
         # Pre-compute per-unit PCA data and neighbor information
         pca_data_per_unit = {}
-        for unit_id in unit_ids:
+        for unit_id in target_unit_ids:
             # Determine neighbor units based on sparsity
             if sorting_analyzer.is_sparse():
                 neighbor_channel_ids = sorting_analyzer.sparsity.unit_id_to_channel_ids[unit_id]
                 neighbor_unit_ids = [
-                    other_unit for other_unit in unit_ids if main_channels[other_unit] in neighbor_channel_ids
+                    other_unit for other_unit in context_unit_ids if main_channels[other_unit] in neighbor_channel_ids
                 ]
                 neighbor_channel_indices = sorting_analyzer.channel_ids_to_indices(neighbor_channel_ids)
             else:
                 neighbor_channel_ids = sorting_analyzer.channel_ids
-                neighbor_unit_ids = unit_ids
+                neighbor_unit_ids = context_unit_ids
                 neighbor_channel_indices = sorting_analyzer.channel_ids_to_indices(neighbor_channel_ids)
 
             # Filter projections to neighbor units

src/spikeinterface/metrics/quality/tests/test_quality_metric_calculator.py

@@ -1,3 +1,5 @@
+import warnings
+
 import pytest
 import numpy as np
 
@@ -74,6 +76,7 @@ def test_compute_quality_metrics(sorting_analyzer_simple):
 def test_merging_quality_metrics(sorting_analyzer_simple):
 
     sorting_analyzer = sorting_analyzer_simple
+    sorting_analyzer.compute("principal_components")
 
     metrics = compute_quality_metrics(
         sorting_analyzer,
@@ -84,7 +87,9 @@ def test_merging_quality_metrics(sorting_analyzer_simple):
     )
 
     # sorting_analyzer_simple has ten units
-    new_sorting_analyzer = sorting_analyzer.merge_units([["0", "1"]])
+    with warnings.catch_warnings():
+        warnings.filterwarnings("error", message="No other units found in the vicinity")
+        new_sorting_analyzer, new_unit_ids = sorting_analyzer.merge_units([["0", "1"]], return_new_unit_ids=True)
     new_metrics = new_sorting_analyzer.get_extension("quality_metrics").get_data()
 
     # we should copy over the metrics after merge
@@ -96,6 +101,9 @@ def test_merging_quality_metrics(sorting_analyzer_simple):
     # 10 units vs 9 units
     assert len(metrics.index) > len(new_metrics.index)
 
+    merged_unit_metrics = new_metrics.loc[new_unit_ids[0]]
+    assert merged_unit_metrics["nn_hit_rate"] != 1 or merged_unit_metrics["nn_miss_rate"] != 0
+
 
 def test_compute_quality_metrics_recordingless(sorting_analyzer_simple):