From 2d8c08476863968f2b3f30164a915be00d520cde Mon Sep 17 00:00:00 2001 From: Kushaan Gupta Date: Fri, 9 Oct 2026 14:44:47 -0400 Subject: [PATCH 1/4] Allow 3D distances for local common reference --- src/spikeinterface/core/recording_tools.py | 15 ++++++++---- .../core/tests/test_recording_tools.py | 19 +++++++++++++++ .../preprocessing/common_reference.py | 8 ++++++- .../tests/test_common_reference.py | 23 +++++++++++++++++++ 4 files changed, 59 insertions(+), 6 deletions(-) diff --git a/src/spikeinterface/core/recording_tools.py b/src/spikeinterface/core/recording_tools.py index 09d0ce0deb..1bcf6d39f4 100644 --- a/src/spikeinterface/core/recording_tools.py +++ b/src/spikeinterface/core/recording_tools.py @@ -419,17 +419,20 @@ def get_random_data_chunks(recording, return_in_uV=False, concatenated=True, **r ) -def get_channel_distances(recording): +def get_channel_distances(recording, axes="xy"): """ - Distance between channel pairs + Distance between channel pairs using the selected spatial axes. + + ``axes`` defaults to ``"xy"`` for backward compatibility. Use ``"xyz"`` + to measure physical distances on a 3-D probe. """ - locations = recording.get_channel_locations() + locations = recording.get_channel_locations(axes=axes) channel_distances = np.linalg.norm(locations[:, np.newaxis] - locations[np.newaxis, :], axis=2) return channel_distances -def get_closest_channels(recording, channel_ids=None, num_channels=None): +def get_closest_channels(recording, channel_ids=None, num_channels=None, axes="xy"): """Get closest channels + distances Parameters @@ -440,6 +443,8 @@ def get_closest_channels(recording, channel_ids=None, num_channels=None): List of channels ids to compute there near neighborhood num_channels : int, default: None Maximum number of neighborhood channels to return + axes : str, default: "xy" + Spatial axes used to calculate distances. Use "xyz" for a 3-D probe. Returns ------- @@ -453,7 +458,7 @@ def get_closest_channels(recording, channel_ids=None, num_channels=None): if num_channels is None: num_channels = len(channel_ids) - 1 - locations = recording.get_channel_locations(channel_ids=channel_ids) + locations = recording.get_channel_locations(channel_ids=channel_ids, axes=axes) closest_channels_inds = [] dists = [] diff --git a/src/spikeinterface/core/tests/test_recording_tools.py b/src/spikeinterface/core/tests/test_recording_tools.py index 13d68e3b45..c3d4f98156 100644 --- a/src/spikeinterface/core/tests/test_recording_tools.py +++ b/src/spikeinterface/core/tests/test_recording_tools.py @@ -3,6 +3,7 @@ import numpy as np from spikeinterface.core import NumpyRecording, generate_recording, MockRecording +from probeinterface import Probe from spikeinterface.core.binaryrecordingextractor import BinaryRecordingExtractor @@ -179,6 +180,24 @@ def test_get_closest_channels(): dist = get_channel_distances(rec) +def test_channel_distances_can_use_3d_probe_geometry(): + recording = NumpyRecording(np.zeros((10, 3), dtype="float32"), sampling_frequency=30_000) + probe = Probe(ndim=3, si_units="um") + probe.set_contacts( + positions=np.array([[0, 0, 0], [50, 0, 50], [100, 0, 0]]), + plane_axes=np.broadcast_to(np.eye(3)[:2], (3, 2, 3)), + ) + probe.set_device_channel_indices(np.arange(3)) + recording.set_probe(probe) + + assert get_channel_distances(recording)[0, 1] == 50 + assert np.isclose(get_channel_distances(recording, axes="xyz")[0, 1], np.sqrt(5000)) + _, distances_xy = get_closest_channels(recording) + _, distances_xyz = get_closest_channels(recording, axes="xyz") + assert distances_xy[0, 0] == 50 + assert np.isclose(distances_xyz[0, 0], np.sqrt(5000)) + + def test_get_noise_levels(): job_kwargs = dict(n_jobs=1, progress_bar=True) rec = generate_recording(num_channels=2, sampling_frequency=1000.0, durations=[60.0]) diff --git a/src/spikeinterface/preprocessing/common_reference.py b/src/spikeinterface/preprocessing/common_reference.py index 6da823d0e4..7069eab90a 100644 --- a/src/spikeinterface/preprocessing/common_reference.py +++ b/src/spikeinterface/preprocessing/common_reference.py @@ -82,6 +82,10 @@ class CommonReferenceRecording(BasePreprocessor): warning is raised at init. dtype : None or dtype, default: None If None the parent dtype is kept. + local_axes : str, default: "xy" + Spatial axes used to select neighbors for local referencing. Set to + "xyz" to use physical distances on a 3-D probe. The default retains + the existing 2-D behavior. Returns ------- @@ -100,6 +104,7 @@ def __init__( local_radius: tuple[float, float] = (30.0, 55.0), min_local_neighbors: int = 5, dtype: str | np.dtype | None = None, + local_axes: str = "xy", ): num_chans = recording.get_num_channels() local_kernel = None @@ -148,7 +153,7 @@ def __init__( ), "Some 'ref_channel_ids' are wrong!" elif reference == "local": assert groups is None, "With 'local' CAR, the group option should not be used." - closest_inds, dist = get_closest_channels(recording) + closest_inds, dist = get_closest_channels(recording, axes=local_axes) # The neighbor kernel is a matrix that will be used to calculate the local reference. # It has shape (num_chans, num_chans) and is filled with zeros except for the columns corresponding to the # neighbors of each channel, which are filled with 1 / number of neighbors. This way, when we do a dot @@ -220,6 +225,7 @@ def __init__( ref_channel_ids=ref_channel_ids, local_radius=local_radius, min_local_neighbors=min_local_neighbors, + local_axes=local_axes, dtype=dtype_.str, ) diff --git a/src/spikeinterface/preprocessing/tests/test_common_reference.py b/src/spikeinterface/preprocessing/tests/test_common_reference.py index 139213056f..00fe490500 100644 --- a/src/spikeinterface/preprocessing/tests/test_common_reference.py +++ b/src/spikeinterface/preprocessing/tests/test_common_reference.py @@ -2,6 +2,7 @@ import pytest from spikeinterface.core import generate_recording, NumpyRecording +from probeinterface import Probe from spikeinterface.preprocessing import common_reference @@ -249,6 +250,28 @@ def test_min_local_radius(): ) +def test_local_reference_can_use_3d_probe_geometry(): + traces = np.tile(np.array([[10, 2, 1]], dtype="float32"), (10, 1)) + recording = NumpyRecording(traces, sampling_frequency=30_000) + probe = Probe(ndim=3, si_units="um") + probe.set_contacts( + positions=np.array([[0, 0, 0], [50, 0, 50], [100, 0, 0]]), + plane_axes=np.broadcast_to(np.eye(3)[:2], (3, 2, 3)), + ) + probe.set_device_channel_indices(np.arange(3)) + recording.set_probe(probe) + + with pytest.warns(UserWarning, match="not have enough neighbors"): + reference_xy = common_reference( + recording, reference="local", local_radius=(60, 80), min_local_neighbors=1 + ) + reference_xyz = common_reference( + recording, reference="local", local_radius=(60, 80), min_local_neighbors=1, local_axes="xyz" + ) + assert np.all(reference_xy.get_traces()[:, 0] == 9) + assert np.all(reference_xyz.get_traces()[:, 0] == 8) + + @pytest.mark.parametrize("operator", ["median", "average"]) def test_local_reference_no_channel_beyond_exclude_radius(operator): # When no channel at all lies beyond the exclude (inner) radius, there is nothing to build a From 1a6e7eefff2c3ff8fa595504c3427d42bdb72934 Mon Sep 17 00:00:00 2001 From: Kushaan Gupta Date: Fri, 9 Oct 2026 14:44:47 -0400 Subject: [PATCH 2/4] Preserve 3D geometry when estimating sparsity --- src/spikeinterface/core/sparsity.py | 9 +++-- .../core/tests/test_sparsity.py | 37 +++++++++++++++++++ 2 files changed, 43 insertions(+), 3 deletions(-) diff --git a/src/spikeinterface/core/sparsity.py b/src/spikeinterface/core/sparsity.py index 9863dfbb77..a7a0efed9f 100644 --- a/src/spikeinterface/core/sparsity.py +++ b/src/spikeinterface/core/sparsity.py @@ -833,14 +833,17 @@ def estimate_sparsity( if method == "radius" and main_channel_indices is not None: assert len(main_channel_indices) == len(sorting.unit_ids) - chan_locs = recording.get_channel_locations() + # The recording accessor defaults to xy, even when the attached probe + # has physical 3-D contact positions. + axes = "xyz" if len(recording.get_probes()) == 1 and recording.get_probe().ndim == 3 else "xy" + chan_locs = recording.get_channel_locations(axes=axes) sparsity = ChannelSparsity.from_radius_and_main_channel( sorting.unit_ids, recording.channel_ids, main_channel_indices, chan_locs, radius_um ) elif method != "by_property": - if recording.get_probes() == 1: - # standard case + if len(recording.get_probes()) == 1: + # Preserve the original probe geometry, including z coordinates. probe = recording.get_probe() else: all_locations = recording.get_channel_locations() diff --git a/src/spikeinterface/core/tests/test_sparsity.py b/src/spikeinterface/core/tests/test_sparsity.py index cc9da07387..a2fb5d7eef 100644 --- a/src/spikeinterface/core/tests/test_sparsity.py +++ b/src/spikeinterface/core/tests/test_sparsity.py @@ -9,6 +9,7 @@ from spikeinterface.core.core_tools import check_json from spikeinterface.core import generate_ground_truth_recording from spikeinterface.core import create_sorting_analyzer +from spikeinterface.core import NumpyRecording, NumpySorting def test_ChannelSparsity(): @@ -162,6 +163,42 @@ def get_dataset(): return recording, sorting +def test_estimate_sparsity_uses_single_probe_3d_geometry(): + positions = np.array([[0, 0, 0], [50, 0, 50], [100, 0, 0], [0, 0, 200]]) + probe = pi.Probe(ndim=3, si_units="um") + probe.set_contacts( + positions=positions, + plane_axes=np.broadcast_to(np.eye(3)[:2], (len(positions), 2, 3)), + ) + probe.set_device_channel_indices(np.arange(len(positions))) + + rng = np.random.default_rng(0) + traces = rng.normal(0, 0.1, (4000, len(positions))).astype("float32") + spike_samples = np.array([400, 900, 1400, 1900, 2400, 2900, 3400]) + traces[spike_samples, 0] -= 20 + recording = NumpyRecording(traces, sampling_frequency=20_000) + recording.set_probe(probe) + sorting = NumpySorting.from_unit_dict({1: spike_samples}, sampling_frequency=20_000) + expected_mask = [[True, True, False, False]] + + direct = estimate_sparsity( + sorting, recording, method="radius", radius_um=80, main_channel_indices=[0] + ) + np.testing.assert_array_equal(direct.mask, expected_mask) + + estimated = estimate_sparsity( + sorting, recording, method="radius", radius_um=80, + num_spikes_for_sparsity=len(spike_samples), n_jobs=1, + ) + np.testing.assert_array_equal(estimated.mask, expected_mask) + + closest = estimate_sparsity( + sorting, recording, method="closest_channels", num_channels=2, + num_spikes_for_sparsity=len(spike_samples), n_jobs=1, + ) + np.testing.assert_array_equal(closest.mask, expected_mask) + + def test_estimate_sparsity(): recording, sorting = get_dataset() num_units = sorting.unit_ids.size From b3d6ab179e481a81afe5ecac1c207dc9eee2ba84 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Mon, 12 Oct 2026 19:24:23 +0000 Subject: [PATCH 3/4] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- .../preprocessing/tests/test_common_reference.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/src/spikeinterface/preprocessing/tests/test_common_reference.py b/src/spikeinterface/preprocessing/tests/test_common_reference.py index 00fe490500..3931bad99a 100644 --- a/src/spikeinterface/preprocessing/tests/test_common_reference.py +++ b/src/spikeinterface/preprocessing/tests/test_common_reference.py @@ -262,9 +262,7 @@ def test_local_reference_can_use_3d_probe_geometry(): recording.set_probe(probe) with pytest.warns(UserWarning, match="not have enough neighbors"): - reference_xy = common_reference( - recording, reference="local", local_radius=(60, 80), min_local_neighbors=1 - ) + reference_xy = common_reference(recording, reference="local", local_radius=(60, 80), min_local_neighbors=1) reference_xyz = common_reference( recording, reference="local", local_radius=(60, 80), min_local_neighbors=1, local_axes="xyz" ) From a33c7e761c623a0c73571ba20b933fea2ce9af09 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Mon, 12 Oct 2026 19:37:40 +0000 Subject: [PATCH 4/4] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- .../core/tests/test_sparsity.py | 20 ++++++++++++------- 1 file changed, 13 insertions(+), 7 deletions(-) diff --git a/src/spikeinterface/core/tests/test_sparsity.py b/src/spikeinterface/core/tests/test_sparsity.py index a2fb5d7eef..bb741160be 100644 --- a/src/spikeinterface/core/tests/test_sparsity.py +++ b/src/spikeinterface/core/tests/test_sparsity.py @@ -181,20 +181,26 @@ def test_estimate_sparsity_uses_single_probe_3d_geometry(): sorting = NumpySorting.from_unit_dict({1: spike_samples}, sampling_frequency=20_000) expected_mask = [[True, True, False, False]] - direct = estimate_sparsity( - sorting, recording, method="radius", radius_um=80, main_channel_indices=[0] - ) + direct = estimate_sparsity(sorting, recording, method="radius", radius_um=80, main_channel_indices=[0]) np.testing.assert_array_equal(direct.mask, expected_mask) estimated = estimate_sparsity( - sorting, recording, method="radius", radius_um=80, - num_spikes_for_sparsity=len(spike_samples), n_jobs=1, + sorting, + recording, + method="radius", + radius_um=80, + num_spikes_for_sparsity=len(spike_samples), + n_jobs=1, ) np.testing.assert_array_equal(estimated.mask, expected_mask) closest = estimate_sparsity( - sorting, recording, method="closest_channels", num_channels=2, - num_spikes_for_sparsity=len(spike_samples), n_jobs=1, + sorting, + recording, + method="closest_channels", + num_channels=2, + num_spikes_for_sparsity=len(spike_samples), + n_jobs=1, ) np.testing.assert_array_equal(closest.mask, expected_mask)