Scale local reference and sparsity computation to 3D per probe geometry

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 = []

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()

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])

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,48 @@ 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

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,
         )
 

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,26 @@ 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