Add correct_inter_session_displacementfunction

debugging/playing.py

@@ -0,0 +1,99 @@
+from spikeinterface.generation.session_displacement_generator import generate_session_displacement_recordings
+from spikeinterface.sortingcomponents.peak_detection import detect_peaks
+from spikeinterface.sortingcomponents.peak_localization import localize_peaks
+
+from spikeinterface.preprocessing.inter_session_alignment import (
+    session_alignment,
+)
+from spikeinterface.widgets import plot_session_alignment, plot_activity_histogram_2d
+import matplotlib.pyplot as plt
+
+import spikeinterface.full as si
+import numpy as np
+
+
+si.set_global_job_kwargs(n_jobs=10)
+
+
+if __name__ == '__main__':
+
+ #   recordings_list, _ = generate_session_displacement_recordings(
+ #       num_units=25,
+ #       recording_durations=[800, 800, 800],
+ #       recording_shifts=((0, 0), (0, -400), (0, 200)),  # TODO: can see how well this is recaptured by comparing the displacements to the known displacement + gradient
+ #       non_rigid_gradient=0.3,
+ #       seed=57,  # 52
+ #   )
+
+    # --------------------------------------------------------------------------------------
+    # Load / generate some recordings
+    # --------------------------------------------------------------------------------------
+
+    # try num units 5 and 65
+    recordings_list, _ = generate_session_displacement_recordings(
+        num_units=5,
+        recording_durations=[25, 25, 25],
+        recording_shifts=((0, 0), (0, -200), (0, 150)),
+        non_rigid_gradient=None,
+        seed=55,
+    )
+
+    peaks_list, peak_locations_list = session_alignment.compute_peaks_locations_for_session_alignment(
+        recordings_list,
+        detect_kwargs={"method": "locally_exclusive"},
+        localize_peaks_kwargs={"method": "grid_convolution"},
+    )
+
+    # --------------------------------------------------------------------------------------
+    # Do the estimation
+    # --------------------------------------------------------------------------------------
+    # For each session, an 'activity histogram' is generated. This can be `entire_session`
+    # or the session can be chunked into segments and some summary generated taken over then.
+    # This might be useful if periods of the recording have weird kinetics or noise.
+    # See `session_alignment.py` for docs on these settings.
+
+    non_rigid_window_kwargs = session_alignment.get_non_rigid_window_kwargs()
+    non_rigid_window_kwargs["rigid"] = False
+#    non_rigid_window_kwargs["num_shifts_global"] = 500
+#    non_rigid_window_kwargs["num_shifts_block"] = 24  # TODO: it makes no sense for this to be larger than the windo
+    non_rigid_window_kwargs["win_step_um"] = 125
+    non_rigid_window_kwargs["win_scale_um"] = 60
+
+    estimate_histogram_kwargs = session_alignment.get_estimate_histogram_kwargs()
+    estimate_histogram_kwargs["method"] = "chunked_median"
+    estimate_histogram_kwargs["histogram_type"] = "activity_2d"
+    estimate_histogram_kwargs["bin_um"] = 2
+    estimate_histogram_kwargs["log_scale"] = True
+    estimate_histogram_kwargs["weight_with_amplitude"] = False
+
+    compute_alignment_kwargs = session_alignment.get_compute_alignment_kwargs()
+    compute_alignment_kwargs["interpolate"] = False
+
+    corrected_recordings_list, extra_info = session_alignment.align_sessions(
+        recordings_list,
+        peaks_list,
+        peak_locations_list,
+        alignment_order="to_session_1",  # "to_session_X" or "to_middle"
+        non_rigid_window_kwargs=non_rigid_window_kwargs,
+        estimate_histogram_kwargs=estimate_histogram_kwargs,
+        compute_alignment_kwargs=compute_alignment_kwargs,
+    )
+
+    plot_session_alignment(
+        recordings_list,
+        peaks_list,
+        peak_locations_list,
+        extra_info["session_histogram_list"],
+        **extra_info["corrected"],
+        spatial_bin_centers=extra_info["spatial_bin_centers"],
+        drift_raster_map_kwargs={"clim":(-250, 0), "scatter_decimate": 10}
+    )
+    plt.show()
+
+    if estimate_histogram_kwargs["histogram_type"]  == "activity_2d":
+        plot_activity_histogram_2d(
+            extra_info["session_histogram_list"],
+            extra_info["spatial_bin_centers"],
+            extra_info["corrected"]["corrected_session_histogram_list"]
+        )
+        plt.show()

doc/how_to/inter_session_alignment.rst

@@ -0,0 +1,194 @@
+How to perform inter-session alignment
+======================================
+
+In this how-to, we will assess and correct changes in probe position across multiple experimental sessions
+using `inter-session alignment`.
+
+This is often valuable for chronic-recording experiments, where the goal is to track units across sessions
+
+A full tutorial, including details on the many settings for this procedure, can be found `here <TODO: ADD LINK (INTERNAL)>`_.
+
+Running inter-session alignment
+-------------------------------
+
+In SpikeInterface, it is recommended to perform inter-session alignment
+following within-session motion correction (if used) and before whitening.
+
+Preprocessed recordings should first be stored in a list:
+
+.. code:: python
+
+    recordings_list = [prepro_session_1, prepro_session_2, ...]
+
+Here, we will simulate such an experiment by generating a pair of sessions in
+which the probe is displaced 200 micrometers (μm) along its y-axis (depth).
+
+First, we will import all required packages and functions:
+
+.. code:: python
+
+    # TODO: should add all of the below to spikeinterface.full? (CHECK)
+    import spikeinterface.full as si
+    from spikeinterface.generation.session_displacement_generator import generate_session_displacement_recordings
+    from spikeinterface.preprocessing.inter_session_alignment import session_alignment
+    from spikeinterface.widgets import plot_session_alignment, plot_activity_histogram_2d
+    import matplotlib.pyplot as plt
+
+
+and then generating the test recordings:
+
+.. code:: python
+
+    recordings_list, _ = generate_session_displacement_recordings(  # TODO: add to spikeinterface.full ?
+        num_units=5,
+        recording_durations=[10, 10],
+        recording_shifts=((0, 0), (0, 200)),  # (x offset, y offset) pairs
+    )
+
+We won't explicitly preprocess these recordings in this how-to, but you can imagine
+preprocessing steps have already been run (e.g. filtering, common reference etc.).
+
+To run inter-session alignment, peaks must be detected and localised
+as the locations of firing neurons are used to anchor the sessions alignment.
+
+If you are **running inter-session alignment following motion correction**, the peaks will
+already be detected and localised. In this case, please jump to
+:ref:`inter-session alignment after motion correction <with_motion_correction>`.
+
+In this section we will assume motion correction was not run, so we need to compute the peaks:
+
+.. code:: python
+
+    peaks_list, peak_locations_list = session_alignment.compute_peaks_locations_for_session_alignment(
+                recordings_list,
+                detect_kwargs={"method": "locally_exclusive"},
+                localize_peaks_kwargs={"method": "grid_convolution"},
+    )
+
+The peak locations (before correction) can be visualised with the plotting function:
+
+.. code:: python
+
+    plot_session_alignment(
+        recordings_list,
+        peaks_list,
+        peak_locations_list,
+        drift_raster_map_kwargs={"clim":(-250, 0)}  # fix the color limit across plots for easy comparison
+    )
+    plt.show()
+
+Now, we are ready to perform inter-session alignment. There are many options associated
+with this method—the simplest way to edit these is to fetch the default options
+with the getter function and make select changes as required:
+
+.. code:: python
+
+    estimate_histogram_kwargs = session_alignment.get_estimate_histogram_kwargs()
+    estimate_histogram_kwargs["histogram_type"] = "activity_2d"  # TODO: RENAME
+
+    corrected_recordings_list, extra_info = session_alignment.align_sessions(
+        recordings_list,
+        peaks_list,
+        peak_locations_list,
+        estimate_histogram_kwargs=estimate_histogram_kwargs
+    )
+
+To assess the performance of inter-session alignment, ``plot_session_alignment()``
+will plot both the original and corrected recordings:
+
+.. code:: python
+
+    plot_session_alignment(  # TODO: is this signature confusing?
+        recordings_list,
+        peaks_list,
+        peak_locations_list,
+        extra_info["session_histogram_list"],
+        **extra_info["corrected"],
+        spatial_bin_centers=extra_info["spatial_bin_centers"],
+        drift_raster_map_kwargs={"clim":(-250, 0)}
+    )
+    plt.show()
+
+As we have used 2d histograms for alignment, we can also plot these with ``plot_activity_histogram_2d()``.
+
+.. _with_motion_correction:
+
+Inter-session alignment after motion correction
+-----------------------------------------------
+
+If motion correction has already been performed, it is possible to reuse the
+previously computed peaks and peak locations, avoiding the need for re-computation.
+We will use the special function` `align_sessions_after_motion_correction()`` for this case.
+
+Critically, the last preprocessing step prior to inter-session alignment should be motion correction.
+This ensures the correction for inter-session alignment will be **added directly to the motion correction**.
+This is beneficial as it avoids interpolating the data (i.e. shifting the traces) more than once.
+
+.. admonition:: Warning
+   :class: warning
+
+    To ensure that inter-session alignment adds the displacement directly to the motion-corrected recording
+    to avoid repeated interpolation, motion correction must be the final operation applied to the recording
+    prior to inter-session alignment.
+
+    You can verify this by confirming the recording is an ``InterpolateMotionRecording`` with:
+
+    .. code:: python
+
+        type(recording)  # quick check, should print `InterpolateMotionRecording`
+
+        from spikeinterace.sortingcomponents.motion.motion_utils import InterpolateMotionRecording
+
+        assert isinstance(recording, InterpolateMotionRecording)  # error if not true
+
+    ``align_sessions_after_motion_correction()`` will raise an error if the passed recordings
+    are not all ``InterpolateMotionRecordings``.
+
+Let's first create some test data. We can create a recording with motion errors,
+then split it in two to simulate two separate sessions:
+
+.. code:: python
+
+    # Generate the recording with motion artefact
+    motion_recording, _ = si.generate_drifting_recording(duration=100)
+    total_duration = motion_recording.get_duration()
+    split_time = total_duration / 2
+
+    # Split in two to simulate two sessions
+    recording_part1 = motion_recording.time_slice(start_time=0, end_time=split_time)
+    recording_part2 = motion_recording.time_slice(start_time=split_time, end_time=total_duration)
+
+
+Next, motion correction is performed, storing the results in a list:
+
+.. code:: python
+
+    # perform motion correction on each session, storing the outputs in lists
+    recordings_list = []
+    motion_info_list = []
+    for recording in [recording_part1, recording_part2]:
+
+        rec, motion_info = si.correct_motion(recording, output_motion_info=True, preset="rigid_fast")
+
+        recordings_list.append(rec)
+        motion_info_list.append(motion_info)
+
+Now, we are ready to use ``align_sessions_after_motion_correction()``
+to align the motion-corrected sessions. This function should always be used
+for aligning motion-corrected sessions, as it ensures the alignment
+parameters are properly matched.
+
+We can pass any arguments directly to ``align_sessions`` using the ``align_sessions_kwargs`` argument:
+
+.. code:: python
+
+    estimate_histogram_kwargs = session_alignment.get_estimate_histogram_kwargs()
+    estimate_histogram_kwargs["histogram_type"] = "activity_2d"  # TODO: RENAME
+
+    align_sessions_kwargs = {"estimate_histogram_kwargs": estimate_histogram_kwargs}
+
+    corrected_recordings_list, extra_info = session_alignment.align_sessions_after_motion_correction(
+        recordings_list, motion_info_list, align_sessions_kwargs
+    )
+
+As above, the inter-session alignment can be assessed using ``plot_session_alignment()``.

playing.py

@@ -0,0 +1,73 @@
+from spikeinterface.generation import generate_drifting_recording
+from spikeinterface.preprocessing.motion import correct_motion
+from spikeinterface.sortingcomponents.motion.motion_interpolation import InterpolateMotionRecording
+from spikeinterface.sortingcomponents.peak_detection import detect_peaks
+from spikeinterface.generation.session_displacement_generator import generate_session_displacement_recordings
+from spikeinterface.generation import generate_ground_truth_recording
+from spikeinterface.core import get_noise_levels
+from spikeinterface.sortingcomponents.peak_localization import localize_peaks
+
+
+recordings_list, _ = generate_session_displacement_recordings(
+    num_units=5,
+    recording_durations=[1, 1],
+    recording_shifts=((0, 0), (0, 250)),
+    # TODO: can see how well this is recaptured by comparing the displacements to the known displacement + gradient
+    non_rigid_gradient=None,  # 0.1, # 0.1,
+    seed=55,  # 52
+    generate_sorting_kwargs=dict(firing_rates=(100, 250), refractory_period_ms=4.0),
+    generate_unit_locations_kwargs=dict(
+        margin_um=0.0,
+        # if this is say 20, then units go off the edge of the probe and are such low amplitude they are not picked up.
+        minimum_z=0.0,
+        maximum_z=2.0,
+        minimum_distance=18.0,
+        max_iteration=100,
+        distance_strict=False,
+    ),
+    generate_noise_kwargs=dict(noise_levels=(0.0, 1.0), spatial_decay=1.0),
+)
+rec = recordings_list[1]
+
+detect_kwargs = {
+    "method": "locally_exclusive",
+    "peak_sign": "neg",
+    "detect_threshold": 25,
+    "exclude_sweep_ms": 0.1,
+    "radius_um": 75,
+    "noise_levels": None,
+    "random_chunk_kwargs": {},
+}
+localize_peaks_kwargs = {"method": "grid_convolution"}
+
+# noise_levels = get_noise_levels(rec, return_scaled=False)
+rec_0 = recordings_list[0]
+rec_1 = recordings_list[1]
+
+peaks_before_0 = detect_peaks(rec_0, **detect_kwargs)  # noise_levels=noise_levels,
+peaks_before_1 = detect_peaks(rec_1, **detect_kwargs)
+
+proc_rec_0, motion_info_0 = correct_motion(rec_0, preset="rigid_fast", detect_kwargs=detect_kwargs, localize_peaks_kwargs=localize_peaks_kwargs, output_motion_info=True)
+proc_rec_1, motion_info_1 = correct_motion(rec_1, preset="rigid_fast", detect_kwargs=detect_kwargs, localize_peaks_kwargs=localize_peaks_kwargs, output_motion_info=True)
+
+peaks_after_0 = detect_peaks(proc_rec_0, **detect_kwargs)  #  noise_levels=noise_levels
+peaks_after_1 = detect_peaks(proc_rec_1, **detect_kwargs)
+
+
+import spikeinterface.full as si
+import matplotlib.pyplot as plt
+
+# TODO: need to test multi-shank
+plot = si.plot_traces(rec_1, order_channel_by_depth=True) # , time_range=(0, 0.1))
+x = peaks_before_1["sample_index"] * (1/ rec_1.get_sampling_frequency())
+y = rec_1.get_channel_locations()[peaks_before_1["channel_index"], 1]
+plot.ax.scatter(x, y, color="r", s=2)
+plt.show()
+
+plot = si.plot_traces(proc_rec_1, order_channel_by_depth=True)
+x = peaks_after_1["sample_index"] * (1/ proc_rec_1.get_sampling_frequency())
+y = rec_1.get_channel_locations()[peaks_after_1["channel_index"], 1]
+plot.ax.scatter(x, y, color="r", s=2)
+plt.show()
+
+breakpoint()

pyproject.toml

@@ -150,7 +150,7 @@ test = [
     "pytest-dependency",
     "pytest-cov",
     "psutil",
-
+    "pytest-mock",
     "huggingface_hub",
 
     # preprocessing