From 5f3a2311c8a0730c77a072419827d82740152cf3 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jes=C3=BAs=20Royeth?= Date: Fri, 18 Sep 2026 20:40:01 -0300 Subject: [PATCH] Precompute unit-pair spatial overlap once in find_collisions --- .../postprocessing/amplitude_scalings.py | 77 +++++++++++-------- .../tests/test_amplitude_scalings.py | 43 ++++++++++- 2 files changed, 85 insertions(+), 35 deletions(-) diff --git a/src/spikeinterface/postprocessing/amplitude_scalings.py b/src/spikeinterface/postprocessing/amplitude_scalings.py index be5967ad1b..d28d2ab1c0 100644 --- a/src/spikeinterface/postprocessing/amplitude_scalings.py +++ b/src/spikeinterface/postprocessing/amplitude_scalings.py @@ -192,17 +192,22 @@ def __init__( assert spike_retriever.include_spikes_in_margin, "Need SpikeRetriever with include_spikes_in_margin=True" if not handle_collisions: self._margin = max(nbefore, nafter) + overlap_matrix = None else: # in this case we extend the margin to be able to get with collisions outside the chunk margin_waveforms = max(nbefore, nafter) max_margin_collisions = delta_collision_samples + margin_waveforms self._margin = max_margin_collisions + # sparsity_mask is fixed for the node's lifetime, so the unit-pair overlap matrix + # is computed once here instead of once per chunk inside compute()/find_collisions. + overlap_matrix = _unit_pair_overlap_matrix(sparsity_mask) # for some edge cases a template can be zero, leading to problems later template_is_zero = [np.all(template == 0) for template in all_templates] self._all_templates = all_templates self._sparsity_mask = sparsity_mask + self._overlap_matrix = overlap_matrix self._nbefore = nbefore self._nafter = nafter self._cut_out_before = cut_out_before @@ -232,6 +237,7 @@ def compute(self, traces, peaks): offsets = self._offsets all_templates = self._all_templates sparsity_mask = self._sparsity_mask + overlap_matrix = self._overlap_matrix nbefore = self._nbefore cut_out_before = self._cut_out_before cut_out_after = self._cut_out_after @@ -255,7 +261,7 @@ def compute(self, traces, peaks): local_spikes, local_spikes_within_margin, delta_collision_samples, - sparsity_mask, + overlap_matrix, local_spike_indices, ) else: @@ -324,30 +330,34 @@ def get_margin(self): ### Collision handling ### -def _are_units_spatially_overlapping(sparsity_mask, i, j): +def _unit_pair_overlap_matrix(sparsity_mask): """ - Returns True if the unit indices i and j are - spatially overlapping, False otherwise + Precompute, once, whether every pair of units shares at least one channel. + + Unit-pair spatial overlap is a fixed fact of `sparsity_mask` alone: there are only + `num_units ** 2` possible (i, j) answers, independent of which spikes are being compared. + `find_collisions` looks this up once per temporally-overlapping spike-pair candidate + (millions of times on a realistic recording), so computing it here with one matrix + multiplication instead of a fresh `np.any(sparsity_mask[i] & sparsity_mask[j])` per lookup + removes an O(num_channels) recomputation of an already-known answer. Parameters ---------- - sparsity_mask: boolean mask + sparsity_mask : boolean mask A num_units x num_channels boolean array indicating whether the unit is represented on the channel. - i: int - The first unit index - j: int - The second unit index Returns ------- - bool - True if the units i and j are spatially overlapping, False otherwise + np.ndarray + A num_units x num_units boolean array where entry (i, j) is True if units i and j + are spatially overlapping, False otherwise. """ - if np.any(sparsity_mask[i] & sparsity_mask[j]): - return True - else: - return False + # int32 avoids any risk of the dot-product overflowing (it accumulates at most num_channels + # per entry, far below the int32 range) while staying far cheaper than a bool broadcast that + # would materialize a full num_units x num_units x num_channels intermediate array. + sparsity_mask_int = np.asarray(sparsity_mask, dtype=np.int32) + return (sparsity_mask_int @ sparsity_mask_int.T) > 0 def _ordinary_scaling_slope(template, local_waveform): @@ -387,7 +397,7 @@ def _ordinary_scaling_slope(template, local_waveform): return covariance / template_variance -def find_collisions(spikes, spikes_within_margin, delta_collision_samples, sparsity_mask, spike_indices): +def find_collisions(spikes, spikes_within_margin, delta_collision_samples, overlap_matrix, spike_indices): """ Finds the collisions between spikes. @@ -415,9 +425,12 @@ def find_collisions(spikes, spikes_within_margin, delta_collision_samples, spars another spike within a given margin delta_collision_samples: int The maximum number of samples between two spikes to consider them as overlapping - sparsity_mask: boolean mask - A num_units x num_channels boolean array indicating whether - the unit is represented on the channel. + overlap_matrix : np.ndarray + A num_units x num_units boolean array, as returned by `_unit_pair_overlap_matrix`, + where entry (i, j) is True if units i and j are spatially overlapping. Callers that + run this once per chunk (like `AmplitudeScalingNode.compute`) should precompute it + once from `sparsity_mask` and reuse it, since spatial overlap is a fixed fact of the + sparsity mask, not of the spikes being compared. spike_indices : np.ndarray The indices of `spikes` in `spikes_within_margin`. Providing these indices avoids searching `spikes_within_margin` once for every spike. @@ -428,7 +441,8 @@ def find_collisions(spikes, spikes_within_margin, delta_collision_samples, spars A dictionary with collisions. The key is the index of the spike with collision, the value is an array of overlapping spikes, including the spike itself at position 0. """ - # TODO: refactor to speed-up + spikes_within_margin_unit_index = spikes_within_margin["unit_index"] + collision_spikes_dict = {} for spike_index, (spike, spike_index_within_margin) in enumerate(zip(spikes, spike_indices, strict=True)): # find the spikes that fall within a temporal window around the spike peak @@ -451,19 +465,16 @@ def find_collisions(spikes, spikes_within_margin, delta_collision_samples, spars (pre_possible_consecutive_spike_indices, post_possible_consecutive_spike_indices) ) - # Build the collusion_spikes_dict including only - # spikes that overlap spatially - collision_spikes = [] - for possible_overlapping_spike_index in possible_overlapping_spike_indices: - - if _are_units_spatially_overlapping( - sparsity_mask, - spike["unit_index"], - spikes_within_margin[possible_overlapping_spike_index]["unit_index"], - ): - collision_spikes.append(spikes_within_margin[possible_overlapping_spike_index]) - if collision_spikes: - collision_spikes_dict[spike_index] = np.array([spike, *collision_spikes], dtype=spikes.dtype) + # Keep only the candidates that overlap spatially, looked up from the precomputed matrix + # instead of recomputed per candidate. + other_unit_indices = spikes_within_margin_unit_index[possible_overlapping_spike_indices] + overlaps = overlap_matrix[spike["unit_index"], other_unit_indices] + collision_spike_indices = possible_overlapping_spike_indices[overlaps] + + if collision_spike_indices.size: + collision_spikes_dict[spike_index] = np.array( + [spike, *spikes_within_margin[collision_spike_indices]], dtype=spikes.dtype + ) return collision_spikes_dict diff --git a/src/spikeinterface/postprocessing/tests/test_amplitude_scalings.py b/src/spikeinterface/postprocessing/tests/test_amplitude_scalings.py index d6faabb616..0d6488e6b4 100644 --- a/src/spikeinterface/postprocessing/tests/test_amplitude_scalings.py +++ b/src/spikeinterface/postprocessing/tests/test_amplitude_scalings.py @@ -5,7 +5,12 @@ from spikeinterface.postprocessing import ComputeAmplitudeScalings from spikeinterface.core.base import spike_peak_dtype -from spikeinterface.postprocessing.amplitude_scalings import _ordinary_scaling_slope, find_collisions, fit_collision +from spikeinterface.postprocessing.amplitude_scalings import ( + _ordinary_scaling_slope, + _unit_pair_overlap_matrix, + find_collisions, + fit_collision, +) def test_ordinary_scaling_slope_float32_precision(): @@ -86,7 +91,7 @@ def test_find_collisions_with_margin_indices(monkeypatch): spikes, spikes_within_margin, delta_collision_samples=4, - sparsity_mask=sparsity_mask, + overlap_matrix=_unit_pair_overlap_matrix(sparsity_mask), spike_indices=spike_indices, ) @@ -96,6 +101,40 @@ def test_find_collisions_with_margin_indices(monkeypatch): np.testing.assert_array_equal(collisions[2], spikes_within_margin[[3, 2]]) +def test_unit_pair_overlap_matrix_matches_naive_reference(): + """ + `_unit_pair_overlap_matrix` replaces a per-pair `np.any(sparsity_mask[i] & sparsity_mask[j])` + lookup (previously recomputed on every temporally-overlapping spike-pair candidate inside + `find_collisions`, millions of times on a realistic recording) with one matrix multiplication + computed once. Verify it agrees, entry by entry, with the direct naive reference on + NON-CONTIGUOUS per-unit channel subsets (a bounded/contiguous slice would pass even a mutant + that only checks a channel range), and on a unit with an entirely empty sparsity row (no + channels at all -- that unit must not spuriously overlap with anything, including itself). + """ + rng = np.random.default_rng(99) + num_units, num_channels = 15, 24 + + sparsity_mask = np.zeros((num_units, num_channels), dtype=bool) + for unit_index in range(num_units): + if unit_index == 0: + continue # unit 0: entirely empty sparsity row + num_active = rng.integers(1, num_channels // 2) + active_channels = rng.choice(num_channels, size=num_active, replace=False) # non-contiguous + sparsity_mask[unit_index, active_channels] = True + + overlap_matrix = _unit_pair_overlap_matrix(sparsity_mask) + + assert overlap_matrix.shape == (num_units, num_units) + assert overlap_matrix.dtype == np.bool_ + assert not overlap_matrix[0, :].any() # the empty-row unit overlaps with nothing + assert not overlap_matrix[:, 0].any() + + for i in range(num_units): + for j in range(num_units): + expected = np.any(sparsity_mask[i] & sparsity_mask[j]) + assert overlap_matrix[i, j] == expected, f"mismatch at unit pair ({i}, {j})" + + class TestAmplitudeScalingsExtension(AnalyzerExtensionCommonTestSuite): @pytest.mark.parametrize("params", [dict(handle_collisions=True), dict(handle_collisions=False)])