Merge pull request #3062 from samuelgarcia/dredge_lfp

Dredge lfp and dredge ap
SpikeInterface · Jul 15, 2024 · c9fc8e1 · c9fc8e1
2 parents ad3e924 + abb6a7c
commit c9fc8e1
Show file tree

Hide file tree

Showing 39 changed files with 3,822 additions and 1,880 deletions.
diff --git a/doc/api.rst b/doc/api.rst
@@ -407,12 +407,6 @@ Peak Detection
 
     .. autofunction:: detect_peaks
 
-Motion Correction
-~~~~~~~~~~~~~~~~~
-.. automodule:: spikeinterface.sortingcomponents.motion_interpolation
-
-    .. autoclass:: InterpolateMotionRecording
-
 Clustering
 ~~~~~~~~~~
 .. automodule:: spikeinterface.sortingcomponents.clustering
@@ -424,3 +418,15 @@ Template Matching
 .. automodule:: spikeinterface.sortingcomponents.matching
 
     .. autofunction:: find_spikes_from_templates
+
+Motion Correction
+~~~~~~~~~~~~~~~~~
+.. automodule:: spikeinterface.sortingcomponents.motion
+
+    .. autoclass:: Motion
+    .. autofunction:: estimate_motion
+    .. autofunction:: interpolate_motion
+    .. autofunction:: correct_motion_on_peaks
+    .. autofunction:: interpolate_motion_on_traces
+    .. autofunction:: clean_motion_vector
+    .. autoclass:: InterpolateMotionRecording
diff --git a/doc/how_to/benchmark_with_hybrid_recordings.rst b/doc/how_to/benchmark_with_hybrid_recordings.rst
@@ -24,7 +24,7 @@ order to smoothly inject spikes into the recording.
     import spikeinterface.generation as sgen
     import spikeinterface.widgets as sw
 
-    from spikeinterface.sortingcomponents.motion_estimation import estimate_motion
+    from spikeinterface.sortingcomponents.motion import estimate_motion
 
     import numpy as np
     import matplotlib.pyplot as plt
@@ -1202,63 +1202,63 @@ drifts when injecting hybrid spikes.
      0.23076923 0.23076923 0.30769231 0.30769231 0.38461538 0.38461538
      0.46153846 0.46153846 0.53846154 0.53846154 0.61538462 0.61538462
      0.69230769 0.69230769 0.76923077 0.76923077 0.84615385 0.84615385
-     0.         0.         0.07692308 0.07692308 0.15384615 0.15384615
+     1.         0.         0.07692308 0.07692308 0.15384615 0.15384615
      0.23076923 0.23076923 0.30769231 0.30769231 0.38461538 0.38461538
      0.46153846 0.46153846 0.53846154 0.53846154 0.61538462 0.61538462
      0.69230769 0.69230769 0.76923077 0.76923077 0.84615385 0.84615385
-     0.         0.         0.07692308 0.07692308 0.15384615 0.15384615
+     2.         0.         0.07692308 0.07692308 0.15384615 0.15384615
      0.23076923 0.23076923 0.30769231 0.30769231 0.38461538 0.38461538
      0.46153846 0.46153846 0.53846154 0.53846154 0.61538462 0.61538462
      0.69230769 0.69230769 0.76923077 0.76923077 0.84615385 0.84615385
-     0.         0.         0.07692308 0.07692308 0.15384615 0.15384615
+     3.         0.         0.07692308 0.07692308 0.15384615 0.15384615
      0.23076923 0.23076923 0.30769231 0.30769231 0.38461538 0.38461538
      0.46153846 0.46153846 0.53846154 0.53846154 0.61538462 0.61538462
      0.69230769 0.69230769 0.76923077 0.76923077 0.84615385 0.84615385
-     0.         0.         0.07692308 0.07692308 0.15384615 0.15384615
+     4.         0.         0.07692308 0.07692308 0.15384615 0.15384615
      0.23076923 0.23076923 0.30769231 0.30769231 0.38461538 0.38461538
      0.46153846 0.46153846 0.53846154 0.53846154 0.61538462 0.61538462
      0.69230769 0.69230769 0.76923077 0.76923077 0.84615385 0.84615385
-     0.         0.         0.07692308 0.07692308 0.15384615 0.15384615
+     5.         0.         0.07692308 0.07692308 0.15384615 0.15384615
      0.23076923 0.23076923 0.30769231 0.30769231 0.38461538 0.38461538
      0.46153846 0.46153846 0.53846154 0.53846154 0.61538462 0.61538462
      0.69230769 0.69230769 0.76923077 0.76923077 0.84615385 0.84615385
-     0.         0.         0.07692308 0.07692308 0.15384615 0.15384615
+     6.         0.         0.07692308 0.07692308 0.15384615 0.15384615
      0.23076923 0.23076923 0.30769231 0.30769231 0.38461538 0.38461538
      0.46153846 0.46153846 0.53846154 0.53846154 0.61538462 0.61538462
      0.69230769 0.69230769 0.76923077 0.76923077 0.84615385 0.84615385
-     0.         0.         0.07692308 0.07692308 0.15384615 0.15384615
+     7.         0.         0.07692308 0.07692308 0.15384615 0.15384615
      0.23076923 0.23076923 0.30769231 0.30769231 0.38461538 0.38461538
      0.46153846 0.46153846 0.53846154 0.53846154 0.61538462 0.61538462
      0.69230769 0.69230769 0.76923077 0.76923077 0.84615385 0.84615385
-     0.         0.         0.07692308 0.07692308 0.15384615 0.15384615
+     8.         0.         0.07692308 0.07692308 0.15384615 0.15384615
      0.23076923 0.23076923 0.30769231 0.30769231 0.38461538 0.38461538
      0.46153846 0.46153846 0.53846154 0.53846154 0.61538462 0.61538462
      0.69230769 0.69230769 0.76923077 0.76923077 0.84615385 0.84615385
-     0.         0.         0.07692308 0.07692308 0.15384615 0.15384615
+     9.         0.         0.07692308 0.07692308 0.15384615 0.15384615
      0.23076923 0.23076923 0.30769231 0.30769231 0.38461538 0.38461538
      0.46153846 0.46153846 0.53846154 0.53846154 0.61538462 0.61538462
      0.69230769 0.69230769 0.76923077 0.76923077 0.84615385 0.84615385
-     0.         0.         0.07692308 0.07692308 0.15384615 0.15384615
+     10.        0.         0.07692308 0.07692308 0.15384615 0.15384615
      0.23076923 0.23076923 0.30769231 0.30769231 0.38461538 0.38461538
      0.46153846 0.46153846 0.53846154 0.53846154 0.61538462 0.61538462
      0.69230769 0.69230769 0.76923077 0.76923077 0.84615385 0.84615385
-     0.         0.         0.07692308 0.07692308 0.15384615 0.15384615
+     11.        0.         0.07692308 0.07692308 0.15384615 0.15384615
      0.23076923 0.23076923 0.30769231 0.30769231 0.38461538 0.38461538
      0.46153846 0.46153846 0.53846154 0.53846154 0.61538462 0.61538462
      0.69230769 0.69230769 0.76923077 0.76923077 0.84615385 0.84615385
-     0.         0.         0.07692308 0.07692308 0.15384615 0.15384615
+     12.        0.         0.07692308 0.07692308 0.15384615 0.15384615
      0.23076923 0.23076923 0.30769231 0.30769231 0.38461538 0.38461538
      0.46153846 0.46153846 0.53846154 0.53846154 0.61538462 0.61538462
      0.69230769 0.69230769 0.76923077 0.76923077 0.84615385 0.84615385
-     0.         0.         0.07692308 0.07692308 0.15384615 0.15384615
+     13.        0.         0.07692308 0.07692308 0.15384615 0.15384615
      0.23076923 0.23076923 0.30769231 0.30769231 0.38461538 0.38461538
      0.46153846 0.46153846 0.53846154 0.53846154 0.61538462 0.61538462
      0.69230769 0.69230769 0.76923077 0.76923077 0.84615385 0.84615385
-     0.         0.         0.07692308 0.07692308 0.15384615 0.15384615
+     14.        0.         0.07692308 0.07692308 0.15384615 0.15384615
      0.23076923 0.23076923 0.30769231 0.30769231 0.38461538 0.38461538
      0.46153846 0.46153846 0.53846154 0.53846154 0.61538462 0.61538462
      0.69230769 0.69230769 0.76923077 0.76923077 0.84615385 0.84615385
-     0.         0.         0.07692308 0.07692308 0.15384615 0.15384615
+     15.        0.         0.07692308 0.07692308 0.15384615 0.15384615
      0.23076923 0.23076923 0.30769231 0.30769231 0.38461538 0.38461538
      0.46153846 0.46153846 0.53846154 0.53846154 0.61538462 0.61538462
      0.69230769 0.69230769 0.76923077 0.76923077 0.84615385 0.84615385]</details></ul></details>

diff --git a/doc/how_to/drift_with_lfp.rst b/doc/how_to/drift_with_lfp.rst
@@ -0,0 +1,163 @@
+Estimate drift using the LFP traces
+===================================
+
+Drift is a well known issue for long shank probes. Some datasets, especially from primates and humans,
+can experience very fast motion due to breathing and heart beats. In these cases, the standard motion
+estimation methods that use detected spikes as a basis for motion inference will fail, because there
+are not enough spikes to "follow" such fast drifts.
+
+Charlie Windolf and colleagues from the Paninski Lab at Columbia have developed a method to estimate
+the motion using the LFP signal: **DREDge**. (more details about the method in the paper
+`DREDge: robust motion correction for high-density extracellular recordings across species <https://doi.org/10.1101/2023.10.24.563768>`_).
+
+This method is particularly suited for the open dataset recorded at Massachusetts General Hospital by Angelique Paulk and colleagues in humans (more details in the [paper](https://doi.org/10.1038/s41593-021-00997-0)). The dataset can be dowloaed from [datadryad](https://datadryad.org/stash/dataset/doi:10.5061/dryad.d2547d840) and it contains recordings on human patients with a Neuropixels probe, some of which with very high and fast motion on the probe, which prevents accurate spike sorting without a proper and adequate motion correction
+
+The **DREDge** method has two options: **dredge_lfp** and **dredge_ap**, which have both been ported inside `SpikeInterface`.
+
+Here we will demonstrate the **dredge_lfp** method to estimate the fast and high drift on this recording.
+
+For each patient, the dataset contains two streams:
+
+* a highpass "action potential" (AP), sampled at 30kHz
+* a lowpass "local field" (LF) sampled at 2.5kHz
+
+For this demonstration, we will use the LF stream.
+
+.. code:: ipython3
+
+    %matplotlib inline
+    %load_ext autoreload
+    %autoreload 2
+
+.. code:: ipython3
+
+    from pathlib import Path
+    import matplotlib.pyplot as plt
+
+    import spikeinterface.full as si
+    from spikeinterface.sortingcomponents.motion import estimate_motion
+
+.. code:: ipython3
+
+    # the dataset has been locally downloaded
+    base_folder = Path("/mnt/data/sam/DataSpikeSorting/")
+    np_data_drift = base_folder / 'human_neuropixel/Pt02/'
+
+Read the spikeglx file
+~~~~~~~~~~~~~~~~~~~~~~
+
+.. code:: ipython3
+
+    raw_rec = si.read_spikeglx(np_data_drift)
+    print(raw_rec)
+
+
+.. parsed-literal::
+
+    SpikeGLXRecordingExtractor: 384 channels - 2.5kHz - 1 segments - 2,183,292 samples
+                                873.32s (14.56 minutes) - int16 dtype - 1.56 GiB
+
+
+Preprocessing
+~~~~~~~~~~~~~
+
+Contrary to the **dredge_ap** approach, which needs detected peaks and peak locations, the **dredge_lfp**
+method is estimating the motion directly on traces.
+Importantly, the method requires some additional pre-processing steps:
+  * ``bandpass_filter``: to "focus" the signal on a particular band
+  * ``phase_shift``: to compensate for the sampling misalignement
+  * ``resample``: to further reduce the sampling fequency of the signal and speed up the computation. The sampling frequency of the estimated motion will be the same as the resampling frequency. Here we choose 250Hz, which corresponds to a sampling interval of 4ms.
+  * ``directional_derivative``: this optional step applies a second order derivative in the spatial dimension to enhance edges on the traces.
+    This is not a general rules and need to be tested case by case.
+  * ``average_across_direction``: Neuropixels 1.0 probes have two contacts per depth. This steps averages them to obtain a unique virtual signal along the probe depth ("y" in ``spikeinterface``).
+
+After appying this preprocessing chain, the motion can be estimated almost by eyes ont the traces plotted with the map mode.
+
+.. code:: ipython3
+
+    lfprec = si.bandpass_filter(
+        raw_rec,
+        freq_min=0.5,
+        freq_max=250,
+
+        margin_ms=1500.,
+        filter_order=3,
+        dtype="float32",
+        add_reflect_padding=True,
+    )
+    lfprec = si.phase_shift(lfprec)
+    lfprec = si.resample(lfprec, resample_rate=250, margin_ms=1000)
+
+    lfprec = si.directional_derivative(lfprec, order=2, edge_order=1)
+    lfprec = si.average_across_direction(lfprec)
+
+    print(lfprec)
+
+
+.. parsed-literal::
+
+    AverageAcrossDirectionRecording: 192 channels - 0.2kHz - 1 segments - 218,329 samples
+                                     873.32s (14.56 minutes) - float32 dtype - 159.91 MiB
+
+
+.. code:: ipython3
+
+    %matplotlib inline
+    si.plot_traces(lfprec, backend="matplotlib", mode="map", clim=(-0.05, 0.05), time_range=(400, 420))
+
+
+
+.. image:: drift_with_lfp_files/drift_with_lfp_8_1.png
+
+
+Run the method
+~~~~~~~~~~~~~~
+
+``estimate_motion()`` is the generic function to estimate motion with multiple
+methods in ``spikeinterface``.
+
+This function returns a ``Motion`` object and we can notice that the interval is exactly
+the same as downsampled signal.
+
+Here we use ``rigid=True``, which means that we have one unqiue signal to
+describe the motion across the entire probe depth.
+
+.. code:: ipython3
+
+    motion = estimate_motion(lfprec, method='dredge_lfp', rigid=True, progress_bar=True)
+    motion
+
+
+.. parsed-literal::
+
+    Online chunks [10.0s each]:   0%|          | 0/87 [00:00<?, ?it/s]
+
+
+.. parsed-literal::
+
+    Motion rigid - interval 0.004s - 1 segments
+
+
+
+Plot the drift
+~~~~~~~~~~~~~~
+
+When plotting the drift, we can notice a very fast drift which corresponds to the heart rate.
+The slower oscillations can be attributed to the breathing signal.
+
+We can appreciate how the estimated motion signal matches the processed LFP traces plotted above.
+
+.. code:: ipython3
+
+    fig, ax = plt.subplots()
+    si.plot_motion(motion, mode='line', ax=ax)
+    ax.set_xlim(400, 420)
+    ax.set_ylim(800, 1300)
+
+
+.. parsed-literal::
+
+    (800.0, 1300.0)
+
+
+.. image:: drift_with_lfp_files/drift_with_lfp_12_1.png
diff --git a/doc/how_to/drift_with_lfp_files/drift_with_lfp_12_1.png b/doc/how_to/drift_with_lfp_files/drift_with_lfp_12_1.png
diff --git a/doc/how_to/drift_with_lfp_files/drift_with_lfp_8_1.png b/doc/how_to/drift_with_lfp_files/drift_with_lfp_8_1.png
diff --git a/doc/how_to/handle_drift.rst b/doc/how_to/handle_drift.rst
@@ -245,7 +245,7 @@ to display the results.
 
 .. code:: ipython
 
-    from spikeinterface.sortingcomponents.motion_interpolation import correct_motion_on_peaks
+    from spikeinterface.sortingcomponents.motion import correct_motion_on_peaks
 
     for preset in some_presets:
         folder = base_folder / 'motion_folder_dataset1' / preset

diff --git a/doc/how_to/index.rst b/doc/how_to/index.rst
@@ -14,3 +14,4 @@ Guides on how to solve specific, short problems in SpikeInterface. Learn how to.
     process_by_channel_group
     load_your_data_into_sorting
     benchmark_with_hybrid_recordings
+    drift_with_lfp
diff --git a/doc/modules/motion_correction.rst b/doc/modules/motion_correction.rst
@@ -151,8 +151,7 @@ The high-level :py:func:`~spikeinterface.preprocessing.correct_motion()` is inte
     from spikeinterface.sortingcomponents.peak_detection import detect_peaks
     from spikeinterface.sortingcomponents.peak_selection import select_peaks
     from spikeinterface.sortingcomponents.peak_localization import localize_peaks
-    from spikeinterface.sortingcomponents.motion_estimation import estimate_motion
-    from spikeinterface.sortingcomponents.motion_interpolation import interpolate_motion
+    from spikeinterface.sortingcomponents.motion import estimate_motion, interpolate_motion
 
     job_kwargs = dict(chunk_duration="1s", n_jobs=20, progress_bar=True)
     # Step 1 : activity profile

diff --git a/doc/modules/sortingcomponents.rst b/doc/modules/sortingcomponents.rst
@@ -190,10 +190,10 @@ Here is an example with non-rigid motion estimation:
     peak_locations = localize_peaks(recording=recording, peaks=peaks, ...) # as above
 
 
-    from spikeinterface.sortingcomponents.motion_estimation import estimate_motion
+    from spikeinterface.sortingcomponents.motion import estimate_motion
     motion, temporal_bins, spatial_bins,
                 extra_check = estimate_motion(recording=recording, peaks=peaks, peak_locations=peak_locations,
-                                              direction='y', bin_duration_s=10., bin_um=10., margin_um=0.,
+                                              direction='y', bin_s=10., bin_um=10., margin_um=0.,
                                               method='decentralized_registration',
                                               rigid=False, win_shape='gaussian', win_step_um=50., win_sigma_um=150.,
                                               progress_bar=True, verbose=True)
@@ -206,7 +206,7 @@ Motion interpolation
 
 The estimated motion can be used to interpolate traces, in other words, for drift correction.
 One possible way is to make an interpolation sample-by-sample to compensate for the motion.
-The :py:class:`~spikeinterface.sortingcomponents.motion_interpolation.InterpolateMotionRecording` is a preprocessing
+The :py:class:`~spikeinterface.sortingcomponents.motion.InterpolateMotionRecording` is a preprocessing
 step doing this. This preprocessing is *lazy*, so that interpolation is done on-the-fly. However, the class needs the
 "motion vector" as input, which requires a relatively long computation (peak detection, localization and motion
 estimation).
@@ -216,7 +216,7 @@ Here is a short example that depends on the output of "Motion interpolation":
 
 .. code-block:: python
 
-  from spikeinterface.sortingcomponents.motion_interpolation import InterpolateMotionRecording
+  from spikeinterface.sortingcomponents.motion import InterpolateMotionRecording
 
   recording_corrected = InterpolateMotionRecording(recording=recording_with_drift, motion=motion, temporal_bins=temporal_bins, spatial_bins=spatial_bins
                                                    spatial_interpolation_method='kriging',

diff --git a/examples/how_to/benchmark_with_hybrid_recordings.py b/examples/how_to/benchmark_with_hybrid_recordings.py
@@ -32,7 +32,7 @@
 import spikeinterface.generation as sgen
 import spikeinterface.widgets as sw
 
-from spikeinterface.sortingcomponents.motion_estimation import estimate_motion
+from spikeinterface.sortingcomponents.motion import estimate_motion
 
 import numpy as np
 import matplotlib.pyplot as plt