diff --git a/debugging/_test_session_alignment.py b/debugging/_test_session_alignment.py new file mode 100644 index 0000000000..da04869088 --- /dev/null +++ b/debugging/_test_session_alignment.py @@ -0,0 +1,305 @@ +from __future__ import annotations + +from spikeinterface.generation.session_displacement_generator import generate_session_displacement_recordings +import matplotlib.pyplot as plt +import numpy as np +import pickle +from spikeinterface.preprocessing.inter_session_alignment import ( + session_alignment, + plotting_session_alignment, + alignment_utils +) +from spikeinterface.sortingcomponents.peak_detection import detect_peaks +from spikeinterface.sortingcomponents.peak_localization import localize_peaks +import spikeinterface.full as si + + +# TODO: all of the nonrigid methods (and even rigid) could be having some strange affects on AP +# waveforms. definately needs looking into! + +# TODO: ask about best way to chunk, as ofc the peak detection takes +# recording as inputs so cannot use get traces with chunks function as planned. +# TODO: expose trimmed versions, robust xcorr + +# Note, the cross correlation is intrinsically limited because for large +# shifts the value is too reduced by the reduction in number of points. +# but, of course cannot scale by number of points due to instability at edges +# This is a major problem, e.g. see the strange results for: +""" + scalings = [np.ones(25), np.r_[np.zeros(10), np.ones(15)]] + recordings_list, _ = generate_session_displacement_recordings( + non_rigid_gradient=None, # 0.05, # 0.05, + num_units=55, + recording_durations=(100, 100, 100, 100), + recording_shifts=( + (0, 0), (0, 250), (0, -150), (0, -210), + ), + recording_amplitude_scalings=None, # {"method": "by_amplitude_and_firing_rate", "scalings": scalings}, + generate_unit_locations_kwargs={"margin_um": 0, "minimum_z": 0, "maximum_z": 0}, + seed=42, + ) +""" + +""" +TODO: in this case, it is not necessary to run peak detection across + the entire recording, would probably be sufficient to + take a few chunks, of size determined by firing frequency of + the neurons in the recording (or just take user defined size). + For now, run on the entire recording and discuss the best way to + run on chunked sections with Sam. +""" + +# with nonrigid shift. This is less of a problem when restricting to a small +# windwo for the nonrigid because even if it fails catistrophically the nonrigid +# error will only be max(non rigid shifts). But its still not good. + +# TODO: add different modes (to mean, to nth session...) +# TODO: document that the output is Hz + +# TODO: major check, refactor and tidy up +# list out carefully all notes +# handle the case where the passed recordings are not motion correction recordings. + +# 3) think about and add new neurons that are introduced when shifted + +# 4) add interpolation of the histograms prior to cross correlation +# 5) add robust cross-correlation +# 6) add trimmed methods +# 7) add better way to estimate chunk length. + +# try and interpolate /smooth the xcorr. What about smoothing the activity histograms directly? +# look into te akima spline + +# TODO: think about the nonrigid alignment, it correlates +# over the entire window. is this wise? try cutting it down a bit? + +# TODO: We this interpolate, smooth for xcorr in both rigid and non-rigid case. Think aboout this / check this is ok +# maybe we only want to apply the smoothings etc for nonrigid like KS motion correction + +# TODO: try forcing all unit locations to actually +# be within the probe. Add some notes on this because it is confusing. + +# 1) write argument checks +# 2) go through with a fine tooth comb, fix all outstanding issues, tidy up, +# plot everything to make sure it is working prior to writing tests. + +# 4) to an optimisation shift and scale instead of the current xcorr method. +# 5) finalise estimation of chunk size (skip for now, try a new alignment method) +# and optimal bin size. +# 6) make some presets? should estimate a lot of parameters based on the data, especially for nonrigid. +# these are all basically based on probe geometry. + + +# Note, shifting can move a unit closer to the channel e.g. if separated +# by 20 um which can increase the signal and make shift estimation harder. + +# go through everything and plot to check before writing tests. +# there is a relationship between bin size and nonrigid bins. If bins are +# too small then nonrigid is very unstable. So either choosing a bigger bin +# size or smoothing over the histogram in relation to the number +# of nonrigid bins may make sense. + +# the results for nonrigid are very dependent on chosen parameters, +# in particular the number of nonrigid windows, gaussian scale, +# smoothing of the histgram. An optimaisation method may also +# serve to help reduce the number of parameters to choose. + +# what you really want is for the window size to adapt to how +# busy the histogram is. + +# Suprisingly, the session that is aligned TO can have a +# major affect. + +# problem with current: +# - xcorr is not the best for large shifts due to lower num overlapping samples +# - + +def _prep_recording(recording, plot=False): + """ + :param recording: + :return: + """ + peaks = detect_peaks(recording, method="locally_exclusive") + + peak_locations = localize_peaks(recording, peaks, method="grid_convolution") + + if plot: + si.plot_drift_raster_map( + peaks=peaks, + peak_locations=peak_locations, + recording=recording, + clim=(-300, 0), # fix clim for comparability across plots + ) + plt.show() + + return peaks, peak_locations + +MOTION = True # True +SAVE = True +PLOT = False +BIN_UM = 5 + + +if SAVE: + scalings = [np.ones(25), np.r_[np.zeros(10), np.ones(15)]] + recordings_list, _ = generate_session_displacement_recordings( + non_rigid_gradient=None, # 0.05, # 0.05, # 0.05, + num_units=55, + recording_durations=(50, 50), # , 100), + recording_shifts=( + (0, 0), + (0, 75), + ), + recording_amplitude_scalings=None, # {"method": "by_amplitude_and_firing_rate", "scalings": scalings}, + generate_unit_locations_kwargs={"margin_um": 0, "minimum_z": 0, "maximum_z": 0}, + generate_templates_kwargs=dict( + ms_before=1.5, + ms_after=3.0, + mode="sphere", # this is key to maintaining consistent unit positions with shift + unit_params=dict( + alpha=(75, 125.0), # firing rate + spatial_decay=(10, 45), + ), + ), + seed=42, + ) + + if not MOTION: + peaks_list = [] + peak_locations_list = [] + + for recording in recordings_list: + peaks, peak_locations = _prep_recording( + recording, + plot=PLOT, + ) + peaks_list.append(peaks) + peak_locations_list.append(peak_locations) + + # something relatively easy, only 15 units + with open("all_recordings.pickle", "wb") as handle: + pickle.dump((recordings_list, peaks_list, peak_locations_list), handle, protocol=pickle.HIGHEST_PROTOCOL) + else: + # if False: + # TODO: need to align spatial bin calculation between estimate motion and + # estimate session methods so they are more easily interoperable. OR + # just take spatial bin centers from interpoalte! + recordings_list_new = [] + peaks_list = [] + peak_locations_list = [] + motion_info_list = [] + from spikeinterface.preprocessing.motion import correct_motion + + for i in range(len(recordings_list)): + new_recording, motion_info = correct_motion( + recordings_list[i], + output_motion_info=True, + estimate_motion_kwargs={ + "rigid": False, + "win_shape": "gaussian", + "win_step_um": 50, + "win_margin_um": 0, + }, + ) + recordings_list_new.append(new_recording) + motion_info_list.append(motion_info) + recordings_list = recordings_list_new + + with open("all_recordings_motion.pickle", "wb") as handle: + pickle.dump((recordings_list, motion_info_list), handle, protocol=pickle.HIGHEST_PROTOCOL) + +if MOTION: + with open("all_recordings_motion.pickle", "rb") as handle: + recordings_list, motion_info_list = pickle.load(handle) +else: + with open("all_recordings.pickle", "rb") as handle: + recordings_list, peaks_list, peak_locations_list = pickle.load(handle) + +# TODO: need docs to be super clear from estimate from existing motion, +# as will use motion correction nonrigid bins even if it is suboptimal. + +estimate_histogram_kwargs = { + "bin_um": BIN_UM, + "method": "first_eigenvector", # CHANGE NAME!! # TODO: double check scaling + "chunked_bin_size_s": "estimate", + "log_scale": True, + "depth_smooth_um": 10, + "histogram_type": "activity_1d", # "y_only", "locations_2d", "activity_2d"" TOOD: better names! +} +compute_alignment_kwargs = { + "num_shifts_block": None, # TODO: can be in um so comaprable with window kwargs. + "interpolate": False, + "interp_factor": 10, + "kriging_sigma": 1, + "kriging_p": 2, + "kriging_d": 2, + "smoothing_sigma_bin": False, # 0.5, + "smoothing_sigma_window": False, # 0.5, + "akima_interp_nonrigid": False, +} +non_rigid_window_kwargs = { + "rigid": False, + "win_shape": "gaussian", + "win_step_um": 250, + "win_scale_um": 250, + "win_margin_um": None, + "zero_threshold": None, +} + +if MOTION: + corrected_recordings_list, extra_info = session_alignment.align_sessions_after_motion_correction( + recordings_list, + motion_info_list, + align_sessions_kwargs={ + "alignment_order": "to_middle", + "estimate_histogram_kwargs": estimate_histogram_kwargs, + "compute_alignment_kwargs": compute_alignment_kwargs, + "non_rigid_window_kwargs": non_rigid_window_kwargs, + } + ) + peaks_list = [info["peaks"] for info in motion_info_list] + peak_locations_list = [info["peak_locations"] for info in motion_info_list] +else: + corrected_recordings_list, extra_info = session_alignment.align_sessions( + recordings_list, + peaks_list, + peak_locations_list, + alignment_order="to_session_1", + estimate_histogram_kwargs=estimate_histogram_kwargs, + compute_alignment_kwargs=compute_alignment_kwargs, + ) + + +plotting_session_alignment.SessionAlignmentWidget( + 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)} # TODO: option to fix this across recordings. +) + +plt.show() + +# TODO: estimate chunk size Hz needs to be scaled for time? is it not been done correctly? + +# No, even two sessions is a mess +# TODO: working assumptions, maybe after rigid, make a template for nonrigid alignment +# as at the moment all nonrigid to eachother is a mess +if False: + A = extra_info["histogram_info_list"][2]["chunked_histograms"] + + mean_ = np.mean(A, axis=0) + median_ = alignment_utils.get_chunked_hist_median(A) + supremum_ = alignment_utils.get_chunked_hist_supremum(A) + poisson_ = alignment_utils.get_chunked_hist_poisson_estimate(A) + eigenvector_ = alignment_utils.get_chunked_hist_eigenvector(A) + + plt.plot(mean_) + plt.plot(median_) + plt.plot(supremum_) + plt.plot(poisson_) + plt.plot(eigenvector_) + plt.legend(["mean", "median", "supremum", "poisson", "eigenvector"]) + plt.show() diff --git a/doc/conf.py b/doc/conf.py index 1ed7b26c66..f75fe69530 100644 --- a/doc/conf.py +++ b/doc/conf.py @@ -128,6 +128,7 @@ 'subsection_order': ExplicitOrder([ '../examples/tutorials/core', '../examples/tutorials/extractors', + '../examples/tutorials/preprocessing', '../examples/tutorials/curation', '../examples/tutorials/metrics', '../examples/tutorials/comparison', diff --git a/doc/tutorials_custom_index.rst b/doc/tutorials_custom_index.rst index d50fa56e22..c212f6f250 100755 --- a/doc/tutorials_custom_index.rst +++ b/doc/tutorials_custom_index.rst @@ -92,6 +92,20 @@ These tutorials focus on the :py:mod:`spikeinterface.core` module. :class-card: gallery-card :text-align: center +Preprocessing tutorials +----------------------- + +.. grid:: 1 2 2 3 + :gutter: 2 + + .. grid-item-card:: Inter-session Alignment + :link-type: ref + :link: sphx_glr_tutorials_preprocessing_plot_7_inter_session_alignment.py + :img-top: /tutorials/preprocessing/images/thumb/sphx_glr_plot_7_inter_session_alignment_thumb.png + :img-alt: Inter-session Alignment + :class-card: gallery-card + :text-align: center + Extractors tutorials -------------------- diff --git a/examples/tutorials/preprocessing/README.rst b/examples/tutorials/preprocessing/README.rst new file mode 100644 index 0000000000..50ef0bd6f1 --- /dev/null +++ b/examples/tutorials/preprocessing/README.rst @@ -0,0 +1,2 @@ +Preprocessing Tutorials +----------------------- diff --git a/examples/tutorials/preprocessing/plot_7_inter_session_alignment.py b/examples/tutorials/preprocessing/plot_7_inter_session_alignment.py new file mode 100644 index 0000000000..f719b6d2cf --- /dev/null +++ b/examples/tutorials/preprocessing/plot_7_inter_session_alignment.py @@ -0,0 +1,352 @@ +""" +How to perform inter-session alignment +====================================== + +In this tutorial 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 + + +Running inter-session alignment +------------------------------- + +In SpikeInterface, it is recommended to perform inter-session alignment +following within-session motion correction (if used) and before whitening / sorting. +If you are running inter-session alignment after motion correction, see +:ref:`inter-session alignment after motion correction `. + +Preprocessed recordings should first be stored in a list: + +.. code-block:: python + + recordings_list = [prepro_session_1, prepro_session_2, ...] + +Here, we will simulate an experiment with two sessions 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: +""" + +import spikeinterface.full as si +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 generate the test recordings: + +recordings_list, _ = si.generate_session_displacement_recordings( + num_units=8, + recording_durations=[10, 10], + recording_shifts=((0, 0), (0, 200)), # (x offset, y offset) pairs + seed=42 +) + +# %% +# We won't preprocess the simulated recordings in this tutorial, 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 `. +# +# In this section of the tutorial, we will assume motion correction was not run, so we need to compute the peaks: + +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: + +plot_session_alignment( + recordings_list, + peaks_list, + peak_locations_list, +) +plt.show() + +# %% +# we are now ready to perform inter-session alignment. There are many options associated +# with this method (see sections below). To edit the configurations, fetch the default options +# with the available getters function and make select changes as required: + +estimate_histogram_kwargs = session_alignment.get_estimate_histogram_kwargs() +estimate_histogram_kwargs["histogram_type"] = "2d" + +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: + +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)} +) +plt.show() + +# %% +# As we have used 2d histograms for alignment, we can also plot these with ``plot_activity_histogram_2d()``: + +plot_activity_histogram_2d( + extra_info["session_histogram_list"], + extra_info["spatial_bin_centers"], + extra_info["corrected"]["corrected_session_histogram_list"] +) +plt.show() + +# %% +# .. _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-block:: python +# +# type(recording) # quick check, should print `InterpolateMotionRecording` +# +# from spikeinterface.sortingcomponents.motion.motion_interpolation 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: + +# Generate the recording with motion artefact +motion_recording = si.generate_drifting_recording(duration=100)[0] +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=None) + +# %% +# Next, motion correction is performed, storing the results in a list: +# We override the motion estimator with its NumPy-based alternative so this +# example does not require PyTorch, which the default ``rigid_fast`` estimator uses. + +# perform motion correction on each session, storing the outputs in lists +recordings_list_motion = [] +motion_info_list = [] +for recording in [recording_part1, recording_part2]: + + rec, motion_info = si.correct_motion( + recording, + output_motion_info=True, + preset="rigid_fast", + estimate_motion_kwargs={"method": "decentralized", "conv_engine": "numpy"}, + ) + + recordings_list_motion.append(rec) + motion_info_list.append(motion_info) + +# %% +# Now, we are ready to use ``align_sessions_after_motion_correction()``. +# We can pass any arguments directly to ``align_sessions`` using the ``align_sessions_kwargs`` argument: + +estimate_histogram_kwargs = session_alignment.get_estimate_histogram_kwargs() +estimate_histogram_kwargs["histogram_type"] = "2d" + +align_sessions_kwargs = {"estimate_histogram_kwargs": estimate_histogram_kwargs} + +corrected_recordings_list_motion, _ = session_alignment.align_sessions_after_motion_correction( + recordings_list_motion, motion_info_list, align_sessions_kwargs +) + +# %% +# As above, the inter-session alignment can be assessed using ``plot_session_alignment()``. + +# %% +# Inter-session alignment settings +# -------------------------------- +# +# Below, the settings that control how inter-session alignment is performed +# are explored. These configs can be accessed by the getter functions +# ``get_estimate_histogram_kwargs``, ``get_compute_alignment_kwargs``, +# ``get_non_rigid_window_kwargs``, ``get_interpolate_motion_kwargs``. +# +# Estimate Histogram Kwargs +# ~~~~~~~~~~~~~~~~~~~~~~~~~ +# +# The settings control how the activity histogram (used for alignment) is estimated +# for each session. They can be obtained with ``get_estimate_histogram_kwargs``. +# +# The ``"bin_um"`` parameter controls the bin-size of the activity histogram. +# Along the probe's y-axis, spatial bins will be generated according to this bin size. +# +# The default ``"method"``, ``"entire_session"``, pools all peaks into one histogram +# and divides by the recording duration to obtain rates. Alternatively, +# ``"chunked_median"`` computes rate histograms for time chunks and takes their +# bin-wise median, reducing the influence of unusually active periods. +# For both methods, ``"log_transform"`` applies ``log2(1 + rate)`` after the +# session histogram has been formed, rather than transforming individual chunks. +# ``"chunked_bin_size_s"`` applies only to ``"chunked_median"``. Its default, +# ``"estimate"``, estimates the chunk size based on firing frequency; a float +# specifies the chunk size in seconds. +# +# The ``histogram_type`` can be ``"1d"` or ``"2d"``, +# if 1D the firing rate x spatial bin histogram is generated. Otherwise +# a firing rate x amplitude x spatial bin histogram is generated. +# +# We can visualise the histograms for each time chunk with: + +estimate_histogram_kwargs = session_alignment.get_estimate_histogram_kwargs() +estimate_histogram_kwargs["histogram_type"] = "1d" +estimate_histogram_kwargs["method"] = "chunked_median" +estimate_histogram_kwargs["chunked_bin_size_s"] = 1.0 + +_, extra_info_rigid = session_alignment.align_sessions( + recordings_list, + peaks_list, + peak_locations_list, + estimate_histogram_kwargs=estimate_histogram_kwargs, +) + +plt.plot(extra_info_rigid["histogram_info_list"][0]["chunked_histograms"].T) +plt.xlabel("Spatial bim (um)") +plt.show() + +# %% +# Compute Alignment Kwargs +# ~~~~~~~~~~~~~~~~~~~~~~~~ +# +# Once the histograms have been generated for each session, the displacement +# between sessions is computed. ``get_compute_alignment_kwargs()`` set how this +# displacement is estimated. +# +# The estimation proceeds similar to the `Kilosort motion-correction `_ +# method (see also the "kilosort-like" option in `:func:`correct_motion``.). Briefly, the cross-correlation +# of activity histograms is performed and the peak position used as a linear estimate of the displacement. +# For-non rigid alignment, first linear alignment is performed, then the probe y-axis is split into segments +# and linear estimation performed in each bin. Then, the displacement set at each bin center are interpolated acoss channels. +# +# Most compute-alignment kwargs are similar to those used in motion correction. +# Key arguments and those related to inter-session alignment +# include: +# +# ``"num_shifts_global"``: This is the number of shifts to perform cross-correlation across for linear alignment. +# Put differently, this is the maximum allowed displacement to consider for rigid alignment. +# ``"num_shifts_block"``: The number of shifts to perform cross-correlation across for non-linear alignment (within each spatial bin). +# ``"akima_interp_nonrigid"``: If ``True``, perform akima interpolation across non-rigid spatial bins (rather than linear). +# ``"min_crosscorr_threshold"``: To estimate alignment, normalised cross-correlation is performed. In some cases, particularly +# for non-rigid alignment, there may be little correlation within a bin. To stop aberrant shifts estimated on poor correlations, +# this sets a minimum value for the correlation used to estimate the shift. If less than this value, the shift will be set to zero. +# +# Non-rigid window kwargs +# ~~~~~~~~~~~~~~~~~~~~~~~ +# Non-rigid window kwargs determine how the non-rigid alignment is performed, +# in particular around how the y-axis of the probe is segmented into blocks +# (each which will be aligned using rigid alignment) are found here. +# (and see ``get_non_rigid_window_kwargs()``. +# +# We can see how the ``compute_alignment_kwargs`` control the non-rigid alignment +# by inspecting the output of inter-session alignment. First, we generate a +# pair of recordings with non-rigid displacement and perform rigid alignment: + +recordings_list, _ = si.generate_session_displacement_recordings( + num_units=8, + recording_durations=[10, 10], + recording_shifts=((0, 0), (0, 200)), # (x offset, y offset) pairs + non_rigid_gradient=0.1, + seed=42 +) + +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"}, +) + + +non_rigid_window_kwargs = session_alignment.get_non_rigid_window_kwargs() +non_rigid_window_kwargs["rigid"] = True + +_, extra_info_rigid = session_alignment.align_sessions( + recordings_list, + peaks_list, + peak_locations_list, + estimate_histogram_kwargs=estimate_histogram_kwargs, + non_rigid_window_kwargs=non_rigid_window_kwargs, +) + +plt.plot(extra_info_rigid["corrected"]["corrected_session_histogram_list"][0]) +plt.plot(extra_info_rigid["corrected"]["corrected_session_histogram_list"][1]) +plt.show() + +# %% +# Above, you can see there rigid alignemnt has well-matched one peak but +# the second peak is offset. Next, we can apply non-rigid alignment, +# and visualise the non-rigid segments that the probe is split into. +# Note that by default, Gaussian windows are used: + +# %% + +non_rigid_window_kwargs = session_alignment.get_non_rigid_window_kwargs() +non_rigid_window_kwargs["rigid"] = False +non_rigid_window_kwargs["win_step_um"] = 200 +non_rigid_window_kwargs["win_scale_um"] = 100 + +compute_alignment_kwargs = session_alignment.get_compute_alignment_kwargs() +compute_alignment_kwargs["akima_interp_nonrigid"] = True + +_, extra_info_nonrigid = session_alignment.align_sessions( + recordings_list, + peaks_list, + peak_locations_list, + estimate_histogram_kwargs=estimate_histogram_kwargs, + non_rigid_window_kwargs=non_rigid_window_kwargs, +) + +plt.plot(extra_info_nonrigid["corrected"]["corrected_session_histogram_list"][0]) +plt.plot(extra_info_nonrigid["corrected"]["corrected_session_histogram_list"][1]) +plt.plot(extra_info_nonrigid["non_rigid_windows"].T) +plt.show() + +# %% +# It is notable that the non-rigid alignment is not perfect in this case. This +# is because for each bin, the displacement is computed and we can imagine +# its value being 'positioned' in the center of the bin. Then, the bin center +# values are interpolates across all channels. This leads to non-perfect alignment. +# +# This is in part because this is a simulated test case with only a few peaks and +# the spatial footprint of the APs is small. Nonetheless, the non-rigid window kwargs +# may be adjusting to maximise the performance of non-rigid alignment in the real world case. diff --git a/src/spikeinterface/core/motion.py b/src/spikeinterface/core/motion.py index 668108701b..80e70e57ea 100644 --- a/src/spikeinterface/core/motion.py +++ b/src/spikeinterface/core/motion.py @@ -69,7 +69,13 @@ def __repr__(self): else: rigid_txt = f"non-rigid - {nbins} spatial bins" - interval_s = self.temporal_bins_s[0][1] - self.temporal_bins_s[0][0] + if self.temporal_bins_s[0].size > 1: + interval_s = self.temporal_bins_s[0][1] - self.temporal_bins_s[0][0] + else: + # If there is only one temporal bin (entire session), assume the bin + # left edge is zero, and take twice it for the bin size. + interval_s = self.temporal_bins_s[0][0] * 2 + txt = f"Motion {rigid_txt} - interval {interval_s}s - {self.num_segments} segments" return txt @@ -152,6 +158,13 @@ def get_displacement_at_time_and_depth(self, times_s, locations_um, segment_inde # reshape to grid domain shape if necessary displacement = displacement.reshape(out_shape) + # For the inter-session alignment case + if self.temporal_bins_s[segment_index].size == 1 and self.spatial_bins_um.size == 1: + # TODO: double check, this used to be assert all isnan + # assert np.all(displacement == -0) + assert self.displacement[segment_index].size == 1 + displacement[:] = self.displacement[segment_index] + return displacement def to_dict(self): diff --git a/src/spikeinterface/core/node_pipeline.py b/src/spikeinterface/core/node_pipeline.py index 0e3a209ba5..8fb951e494 100644 --- a/src/spikeinterface/core/node_pipeline.py +++ b/src/spikeinterface/core/node_pipeline.py @@ -660,7 +660,7 @@ def run_node_pipeline( Here a "spike" is a spike with any a label so already sorted. The main idea is to have a graph of nodes. - Every node is doing a computaion of some peaks and related traces. + Every node is doing a computation of some peaks and related traces. The first node is PeakSource so either a peak detector PeakDetector or peak/spike replay (PeakRetriever/SpikeRetriever) Every node can have one or several output that can be directed to other nodes (aka nodes have parents). diff --git a/src/spikeinterface/generation/__init__.py b/src/spikeinterface/generation/__init__.py index ea4cccb911..3f0d96696f 100644 --- a/src/spikeinterface/generation/__init__.py +++ b/src/spikeinterface/generation/__init__.py @@ -22,6 +22,7 @@ generate_displacement_vector, generate_drifting_recording, ) +from .session_displacement_generator import generate_session_displacement_recordings from .template_database import ( fetch_template_object_from_database, diff --git a/src/spikeinterface/preprocessing/__init__.py b/src/spikeinterface/preprocessing/__init__.py index 502c3748da..39614d53ed 100644 --- a/src/spikeinterface/preprocessing/__init__.py +++ b/src/spikeinterface/preprocessing/__init__.py @@ -9,6 +9,15 @@ compute_motion, ) +from .inter_session_alignment.session_alignment import ( + get_estimate_histogram_kwargs, + get_compute_alignment_kwargs, + get_non_rigid_window_kwargs, + get_interpolate_motion_kwargs, + align_sessions, + align_sessions_after_motion_correction, + compute_peaks_locations_for_session_alignment, +) from .preprocessing_tools import get_spatial_interpolation_kernel from .detect_bad_channels import detect_bad_channels from .correct_lsb import correct_lsb diff --git a/src/spikeinterface/preprocessing/inter_session_alignment/__init__.py b/src/spikeinterface/preprocessing/inter_session_alignment/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/src/spikeinterface/preprocessing/inter_session_alignment/alignment_utils.py b/src/spikeinterface/preprocessing/inter_session_alignment/alignment_utils.py new file mode 100644 index 0000000000..bf6e670581 --- /dev/null +++ b/src/spikeinterface/preprocessing/inter_session_alignment/alignment_utils.py @@ -0,0 +1,469 @@ +from spikeinterface import BaseRecording +import numpy as np + +from packaging.version import Version + +# ############################################################################# +# Get Histograms +# ############################################################################# + + +def get_2d_activity_histogram( + recording: BaseRecording, + peaks: np.ndarray, + peak_locations: np.ndarray, + spatial_bin_edges: np.ndarray, + bin_s: float | None, + depth_smooth_um: float | None, + scale_to_hz: bool = False, + weight_with_amplitude: bool = False, + avg_in_bin: bool = True, +): + """ + Generate a 2D activity histogram for the session. Wraps the underlying + spikeinterface function with some adjustments for scaling to time and + log transform. + + Parameters + ---------- + + recording: BaseRecording, + A SpikeInterface recording object. + peaks: np.ndarray, + A SpikeInterface `peaks` array. + peak_locations: np.ndarray, + A SpikeInterface `peak_locations` array. + spatial_bin_edges: np.ndarray, + A (1 x n_bins + 1) array of spatial (probe y dimension) bin edges. + bin_s | None: float, + If `None`, a single histogram will be generated from all session + peaks. Otherwise, multiple histograms will be generated, one for + each time bin. + depth_smooth_um: float | None + If not `None`, smooth the histogram across the spatial + axis. see `make_2d_motion_histogram()` for details. + + """ + # avoid circular imports + from spikeinterface.sortingcomponents.motion.motion_utils import make_2d_motion_histogram + + activity_histogram, temporal_bin_edges, generated_spatial_bin_edges = make_2d_motion_histogram( + recording, + peaks, + peak_locations, + weight_with_amplitude=weight_with_amplitude, + direction="y", + bin_s=(bin_s if bin_s is not None else recording.get_duration(segment_index=0)), + bin_um=None, + hist_margin_um=None, + spatial_bin_edges=spatial_bin_edges, + depth_smooth_um=depth_smooth_um, + avg_in_bin=avg_in_bin, + ) + + if scale_to_hz: + activity_histogram = scale_histogram_to_hz(activity_histogram, temporal_bin_edges, recording.get_duration()) + + temporal_bin_centers = get_bin_centers(temporal_bin_edges) + spatial_bin_centers = get_bin_centers(spatial_bin_edges) + + return activity_histogram, temporal_bin_centers, spatial_bin_centers + + +def scale_histogram_to_hz(histogram, temporal_bin_edges, recording_duration_s): + """Normalize a floating-point histogram in place by each time bin's recorded duration. + + Time is the first axis; remaining axes can represent depth and amplitude. + Bin edges start at the recording start, and the final bin may extend past + the recording end. Only its recorded duration contributes to normalization. + """ + recording_end = temporal_bin_edges[0] + recording_duration_s + chunk_durations = np.diff(np.minimum(temporal_bin_edges, recording_end)) + duration_shape = (-1,) + (1,) * (histogram.ndim - 1) + histogram /= chunk_durations.reshape(duration_shape) + return histogram + + +def get_bin_centers(bin_edges): + return (bin_edges[1:] + bin_edges[:-1]) / 2 + + +def estimate_chunk_size(scaled_activity_histogram): + """ + Estimate a chunk size based on the firing rate. Intuitively, we + want longer chunk size to better estimate low firing rates. The + estimation computes a summary of the firing rates for the session + by taking the value 25% of the max of the activity histogram. + + Then, the chunk size that will accurately estimate this firing rate + within 90% accuracy, 90% of the time based on assumption of Poisson + firing (based on CLT) is computed. + + Parameters + ---------- + + scaled_activity_histogram: np.ndarray + The activity histogram scaled to firing rate in Hz. + """ + print("scaled max", np.max(scaled_activity_histogram)) + + firing_rate = np.max(scaled_activity_histogram) * 0.25 + + lambda_hat_s = firing_rate + range_percent = 0.1 + confidence_z = 1.645 # 90% of samples in the normal distribution + e = lambda_hat_s * range_percent + + t = lambda_hat_s / (e / confidence_z) ** 2 + + print(f"Chunked histogram window size of: {t}s estimated for firing rate (25% of histogram peak) of {lambda_hat_s}") + + return t + + +# ############################################################################# +# Chunked Histogram estimation methods +# ############################################################################# +# Given a set off chunked_session_histograms (num time chunks x num spatial bins) +# take the summary statistic over the time axis. + + +def get_chunked_hist_median(chunked_session_histograms): + """ """ + median_hist = np.median(chunked_session_histograms, axis=0) + + return median_hist + + +def compute_histogram_crosscorrelation( + session_histogram_list: np.ndarray, + non_rigid_windows: np.ndarray, + num_shifts: int, + interpolate: bool, + interp_factor: int, + kriging_sigma: float, + kriging_p: float, + kriging_d: float, + smoothing_sigma_bin: None | float, + smoothing_sigma_window: None | float, + min_crosscorr_threshold: float, +) -> tuple[np.ndarray, np.ndarray]: + """ + Given a list of session activity histograms, cross-correlate + all histograms returning the peak correlation shift (in indices) + in a symmetric (num_session x num_session) matrix. + + Supports non-rigid estimation by windowing the activity histogram + and performing separate cross-correlations on each window separately. + + Parameters + ---------- + + session_histogram_list : np.ndarray + (num_sessions, num_bins) (1d histogram) or (num_sessions, num_bins, 2) (2d histogram) + array of session activity histograms. + non_rigid_windows : np.ndarray + A (num windows x num_bins) binary of weights by which to window + the activity histogram for non-rigid-registration. For example, if + 2 rectangular masks were used, there would be a two row binary mask + the first row with mask of the first half of the probe and the second + row a mask for the second half of the probe. + num_shifts : int + Number of indices by which to shift the histogram to find the maximum + of the cross correlation. If `None`, the entire activity histograms + are cross-correlated. + interpolate : bool + If `True`, the cross-correlation is interpolated before maximum is taken. + interp_factor: + Factor by which to interpolate the cross-correlation. + kriging_sigma : float + sigma parameter for kriging_kernel function. See `kriging_kernel`. + kriging_p : float + p parameter for kriging_kernel function. See `kriging_kernel`. + kriging_d : float + d parameter for kriging_kernel function. See `kriging_kernel`. + smoothing_sigma_bin : float + sigma parameter for the gaussian smoothing kernel over the + spatial bins. + smoothing_sigma_window : float + sigma parameter for the gaussian smoothing kernel over the + non-rigid windows. + + Returns + ------- + + shift_matrix : ndarray + A (num_session x num_session) symmetric matrix of shifts + (indices) between pairs of session activity histograms. + + Notes + ----- + + - This function is very similar to the IterativeTemplateRegistration + function used in motion correct, though slightly difference in scope. + It was not convenient to merge them at this time, but worth looking + into in future. + + - Some obvious performances boosts, not done so because already fast + 1) the cross correlations for each session comparison are performed + twice. They are slightly different due to interpolation, but + still probably better to calculate once and flip. + 2) `num_shifts` is implemented by simply making the full + cross correlation. Would probably be nicer to explicitly calculate + only where needed. However, in general these cross correlations are + only a few thousand datapoints and so are already extremely + fast to cross correlate. + + Notes + ----- + + - The original kilosort method does not work in the inter-session + context because it averages over time bins to form a template to + align too. In this case, averaging over a small number of possibly + quite different session histograms does not work well. + + - In the nonrigid case, this strategy can completely fail when the xcorr + is very bad for a certain window. The smoothing and interpolation + make it much worse, because bad xcorr are merged together. The x-corr + can be bad when the recording is shifted a lot and so there are empty + regions that are correlated with non-empty regions in the nonrigid + approach. A different approach will need to be taken in this case. + + Note that kilosort method does not work because creating a + mean does not make sense over sessions. + """ + # scipy is not a core dependency + from scipy.ndimage import gaussian_filter + + # avoid circular imports + from spikeinterface.sortingcomponents.motion.iterative_template import kriging_kernel + + num_sessions = session_histogram_list.shape[0] + num_bins = session_histogram_list.shape[1] # all hists are same length + num_windows = non_rigid_windows.shape[0] + + shift_matrix = np.zeros((num_sessions, num_sessions, num_windows)) + + center_bin = np.floor((num_bins * 2 - 1) / 2).astype(int) + + # Create the (num windows, num_bins) matrix for this pair of sessions + if num_shifts is None: + num_shifts = num_bins - 1 + shifts_array = np.arange(-(num_shifts), num_shifts + 1) + num_iter = shifts_array.size + + for i in range(num_sessions): + for j in range(i, num_sessions): + xcorr_matrix = np.zeros((non_rigid_windows.shape[0], num_iter)) + + # For each window, window the session histograms (`window` is binary) + # and perform the cross correlations + for win_idx, window in enumerate(non_rigid_windows): + if session_histogram_list.ndim == 3: + # For 2D histogram (spatial, amplitude), manually loop through shifts along + # the spatial axis of the histogram. This is faster than using correlate2d + # because we are not shifting along the amplitude axis. + + windowed_histogram_i = session_histogram_list[i, :] * window[:, np.newaxis] + windowed_histogram_j = session_histogram_list[j, :] * window[:, np.newaxis] + + windowed_histogram_i = (windowed_histogram_i - np.mean(windowed_histogram_i)) / ( + np.std(windowed_histogram_i) + 1e-8 + ) + windowed_histogram_j = (windowed_histogram_j - np.mean(windowed_histogram_j)) / ( + np.std(windowed_histogram_j) + 1e-8 + ) + + xcorr = np.zeros(num_iter) + + for idx, shift in enumerate(shifts_array): + shifted_i = shift_array_fill_zeros(windowed_histogram_i, shift) + flatten_i = shifted_i.flatten() + + xcorr[idx] = np.correlate(flatten_i, windowed_histogram_j.flatten()).item() / flatten_i.size + + else: + # For a 1D histogram, compute the full cross-correlation and + # window the desired shifts ( this is faster than manual looping). + windowed_histogram_i = session_histogram_list[i, :] * window + windowed_histogram_j = session_histogram_list[j, :] * window + + windowed_histogram_i = (windowed_histogram_i - np.mean(windowed_histogram_i)) / ( + np.std(windowed_histogram_i) + 1e-8 + ) + windowed_histogram_j = (windowed_histogram_j - np.mean(windowed_histogram_j)) / ( + np.std(windowed_histogram_j) + 1e-8 + ) + + xcorr = np.correlate( + windowed_histogram_i, + windowed_histogram_j, + mode="full", + ) / (windowed_histogram_i.size) + + if num_shifts: + window_indices = np.arange(center_bin - num_shifts, center_bin + num_shifts + 1) + xcorr = xcorr[window_indices] + + xcorr_matrix[win_idx, :] = xcorr + + # Smooth the cross-correlations across the bins + if smoothing_sigma_bin: + xcorr_matrix = gaussian_filter(xcorr_matrix, sigma=smoothing_sigma_bin, axes=1) + + # Smooth the cross-correlations across the windows + if num_windows > 1 and smoothing_sigma_window: + xcorr_matrix = gaussian_filter(xcorr_matrix, sigma=smoothing_sigma_window, axes=0) + + # Upsample the cross-correlation + if interpolate: + num_points = (shifts_array.size - 1) * interp_factor + 1 + shifts_upsampled = np.linspace(shifts_array[0], shifts_array[-1], num_points) + + K = kriging_kernel( + np.c_[np.ones_like(shifts_array), shifts_array], + np.c_[np.ones_like(shifts_upsampled), shifts_upsampled], + sigma=kriging_sigma, + p=kriging_p, + d=kriging_d, + ) + + xcorr_matrix = np.matmul(xcorr_matrix, K, axes=[(-2, -1), (-2, -1), (-2, -1)]) + xcorr_peak = np.argmax(xcorr_matrix, axis=1) + xcorr_value = np.max(xcorr_matrix, axis=1) + shifts_to_idx = shifts_upsampled + else: + xcorr_peak = np.argmax(xcorr_matrix, axis=1) + xcorr_value = np.max(xcorr_matrix, axis=1) + shifts_to_idx = shifts_array + + shift = shifts_to_idx[xcorr_peak] + + shift[np.where(xcorr_value < min_crosscorr_threshold)] = 0 + + shift_matrix[i, j, :] = shift + + # Reversing a session pair reverses its shift, so fill the lower triangle + # from the transposed upper triangle for every spatial window. + i, j = np.tril_indices(shift_matrix.shape[0], k=-1) + shift_matrix[i, j, :] = -shift_matrix[j, i, :] + + return shift_matrix, xcorr_matrix + + +def shift_array_fill_zeros(array: np.ndarray, shift: int) -> np.ndarray: + """ + Shift an array by `shift` indices, padding with zero. + Samples going out of bounds are dropped i,e, the array is not + extended and samples are not wrapped around to the start of the array. + + Parameters + ---------- + + array : np.ndarray + The array to pad. + shift : int + Number of indices why which to shift the array. If positive, the + zeros are added from the end of the array. If negative, the zeros + are added from the start of the array. + + Returns + ------- + + cut_padded_array : np.ndarray + The `array` padded with zeros and cut down (i.e. out of bounds + samples dropped). + + """ + abs_shift = np.abs(shift) + pad_tuple = (0, abs_shift) if shift > 0 else (abs_shift, 0) + + if array.ndim == 2: + pad_tuple = (pad_tuple, (0, 0)) + + padded_hist = np.pad(array, pad_tuple, mode="constant") + + if padded_hist.ndim == 2: + cut_padded_array = padded_hist[abs_shift:, :] if shift >= 0 else padded_hist[:-abs_shift, :] + else: + cut_padded_array = padded_hist[abs_shift:] if shift >= 0 else padded_hist[:-abs_shift] + + return cut_padded_array + + +def akima_interpolate_nonrigid_shifts( + non_rigid_shifts: np.ndarray, + non_rigid_window_centers: np.ndarray, + spatial_bin_centers: np.ndarray, +): + """ + Perform Akima spline interpolation on a set of non-rigid shifts. + The non-rigid shifts are per segment of the probe, each segment + containing a number of channels. Interpolating these non-rigid + shifts to the spatial bin centers gives a more accurate shift + per channel. + + Parameters + ---------- + non_rigid_shifts : np.ndarray + non_rigid_window_centers : np.ndarray + spatial_bin_centers : np.ndarray + + Returns + ------- + interp_nonrigid_shifts : np.ndarray + An array (length num_spatial_bins) of shifts + interpolated from the non-rigid shifts. + + """ + import scipy + + if Version(scipy.__version__) < Version("1.14.0"): + raise ImportError("Scipy version 14 or higher is required fro Akima interpolation.") + + from scipy.interpolate import Akima1DInterpolator + + x = non_rigid_window_centers + xs = spatial_bin_centers + + num_sessions = non_rigid_shifts.shape[0] + num_bins = spatial_bin_centers.shape[0] + + interp_nonrigid_shifts = np.zeros((num_sessions, num_bins)) + for ses_idx in range(num_sessions): + y = non_rigid_shifts[ses_idx] + y_new = Akima1DInterpolator(x, y, method="akima", extrapolate=True)(xs) + interp_nonrigid_shifts[ses_idx, :] = y_new + + return interp_nonrigid_shifts + + +def get_shifts_from_session_matrix(alignment_order: str, session_offsets_matrix: np.ndarray): + """ + Given a matrix of displacements between all sessions, find the + shifts (one per session) to bring the sessions into alignment. + Assumes `session_offsets_matrix` is skew symmetric. + + Parameters + ---------- + alignment_order : "to_middle" or "to_session_X" where + "N" is the number of the session to align to. + session_offsets_matrix : np.ndarray + The num_sessions x num_sessions symmetric matrix + of displacements between all sessions, generated by + `_compute_session_alignment()`. + + Returns + ------- + optimal_shift_indices : np.ndarray + A 1 x num_sessions array of shifts to apply to + each session in order to bring all sessions into + alignment. + """ + if alignment_order == "to_middle": + optimal_shift_indices = -np.mean(session_offsets_matrix, axis=0) + else: + ses_idx = int(alignment_order.split("_")[-1]) - 1 + optimal_shift_indices = -session_offsets_matrix[ses_idx, :, :] + + return optimal_shift_indices diff --git a/src/spikeinterface/preprocessing/inter_session_alignment/session_alignment.py b/src/spikeinterface/preprocessing/inter_session_alignment/session_alignment.py new file mode 100644 index 0000000000..a3e94307f9 --- /dev/null +++ b/src/spikeinterface/preprocessing/inter_session_alignment/session_alignment.py @@ -0,0 +1,1030 @@ +from __future__ import annotations +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from spikeinterface.core.baserecording import BaseRecording + +import warnings +import numpy as np +from spikeinterface.core import get_noise_levels +from spikeinterface.core.motion import Motion +from spikeinterface.preprocessing.motion import run_peak_detection_pipeline_node +from spikeinterface.preprocessing.inter_session_alignment import alignment_utils + +import copy + + +def get_estimate_histogram_kwargs() -> dict: + """ + A dictionary controlling how the histogram for each session is + computed. Histograms use the entire session or the bin-wise median + of time-chunk histograms, with counts normalized by duration before + aggregation and optional log transformation. + + Returns + ------- + A dictionary with entries: + + "bin_um" : number of spatial histogram bins. As the estimated peak + locations are continuous (i.e. real numbers) this is not constrained + by the number of channels. + "method" : "entire_session" (default) or "chunked_median". + "entire_session" pools all peaks into one histogram and divides by + the recording duration to obtain rates, ignoring "chunked_bin_size_s". + "chunked_median" computes a rate histogram for each time chunk, then + takes the median across chunks for each histogram bin. + See `alignment_utils.py` for details. + "chunked_bin_size_s" : The length in seconds (float) to chunk the recording + for estimating the chunked histograms. Can be set to "estimate" (str), + and the size is estimated from firing frequencies. + "log_transform" : if `True`, apply log2(1 + rate) to the session histogram + after pooling the entire session or taking the median of chunk rates. + "depth_smooth_um" : if `None`, no smoothing is applied. See + `make_2d_motion_histogram`. + """ + return { + "bin_um": 2, + "method": "entire_session", + "chunked_bin_size_s": "estimate", + "log_transform": True, + "depth_smooth_um": None, + "histogram_type": "1d", + "weight_with_amplitude": False, + "avg_in_bin": False, + } + + +def get_compute_alignment_kwargs() -> dict: + """ + A dictionary with settings controlling how inter-session + alignment is estimated and computed given a set of + session activity histograms. + + All keys except for "non_rigid_window_kwargs" determine + how alignment is estimated, based on the kilosort ("kilosort_like" + in spikeinterface) motion correction method. See + `iterative_template_registration` for details. + + "non_rigid_window_kwargs" : if nonrigid alignment + is performed, this determines the nature of the + windows along the probe depth. See `get_spatial_windows`. + """ + return { + "num_shifts_global": None, + "num_shifts_block": 20, + "interpolate": False, + "interp_factor": 10, + "kriging_sigma": 1, + "kriging_p": 2, + "kriging_d": 2, + "smoothing_sigma_bin": 0.5, + "smoothing_sigma_window": 0.5, + "akima_interp_nonrigid": False, + "min_crosscorr_threshold": 0.001, + } + + +def get_non_rigid_window_kwargs(): + """ + see get_spatial_windows() for parameters. + + """ + return { + "rigid": True, + "win_shape": "gaussian", + "win_step_um": 50, + "win_scale_um": 50, + "win_margin_um": None, + "zero_threshold": None, + } + + +def get_interpolate_motion_kwargs(): + """ + Settings to pass to `InterpolateMotionRecording`, + see that class for parameter descriptions. + """ + return { + "border_mode": "force_zeros", # fixed as this until can figure out probe + "spatial_interpolation_method": "kriging", + "sigma_um": 20.0, + "p": 2, + } + + +############################################################################### +# Public Entry Level Functions +############################################################################### + + +def align_sessions( + recordings_list: list[BaseRecording], + peaks_list: list[np.ndarray], + peak_locations_list: list[np.ndarray], + alignment_order: str = "to_middle", + non_rigid_window_kwargs: dict = get_non_rigid_window_kwargs(), + estimate_histogram_kwargs: dict = get_estimate_histogram_kwargs(), + compute_alignment_kwargs: dict = get_compute_alignment_kwargs(), + interpolate_motion_kwargs: None | dict = get_interpolate_motion_kwargs(), +) -> tuple[list[BaseRecording], dict]: + """ + Estimate probe displacement across recording sessions and + return interpolated, displacement-corrected recording. Displacement + is only estimated along the "y" dimension. + + This assumes peaks and peak locations have already been computed. + See `compute_peaks_locations_for_session_alignment` for generating + `peaks_list` and `peak_locations_list` from a `recordings_list`. + + If a recording in `recordings_list` is already an `InterpolateMotionRecording`, + the displacement will be added to the existing shifts to avoid duplicate + interpolations. Note the returned, corrected recording is a copy + (recordings in `recording_list` are not edited in-place). + + Parameters + ---------- + recordings_list : list[BaseRecording] + A list of recordings to be aligned. + peaks_list : list[np.ndarray] + A list of peaks detected from the recordings in `recordings_list`, + as returned from the `detect_peaks` function. Each entry in + `peaks_list` should be from the corresponding entry in `recordings_list`. + peak_locations_list : list[np.ndarray] + A list of peak locations, as computed by `localize_peaks`. Each entry + in `peak_locations_list` should be matched to the corresponding entry + in `peaks_list` and `recordings_list`. + alignment_order : str + "to_middle" will align all sessions to the mean position. + Alternatively, "to_session_N" where "N" is a session number + will align to the Nth session. + non_rigid_window_kwargs : dict + see `get_non_rigid_window_kwargs` + estimate_histogram_kwargs : dict + see `get_estimate_histogram_kwargs()` + compute_alignment_kwargs : dict + see `get_compute_alignment_kwargs()` + interpolate_motion_kwargs : dict + see `get_interpolate_motion_kwargs()` Will not be used if passed + recording is InterpolateMotionRecording (in which case, do not + use this function but use `compute_peaks_locations_for_session_alignment`. + + Returns + ------- + `corrected_recordings_list : list[BaseRecording] + List of displacement-corrected recordings (corresponding + in order to `recordings_list`). If an input recordings is + an InterpolateMotionRecording` recording, the corrected + output recording will be a copy of the input recording with + the additional displacement correction added. + + extra_outputs_dict : dict + Dictionary of features used in the alignment estimation and correction. + + shifts_array : np.ndarray + A (num_sessions x num_rigid_windows) array of shifts. + session_histogram_list : list[np.ndarray] + A list of histograms (one per session) used for the alignment. + spatial_bin_centers : np.ndarray + The spatial bin centers, shared between all recordings. + temporal_bin_centers_list : list[np.ndarray] + List of temporal bin centers. As alignment is based on a single + histogram per session, this contains only 1 value per recording, + which is the mid-timepoint of the recording. + non_rigid_window_centers : np.ndarray + Window centers of the probe segments used for non-rigid alignment. + If rigid alignment is performed, this is a single value (mid-probe). + non_rigid_windows : np.ndarray + A (num nonrigid windows, num spatial_bin_centers) binary array used to mask + the probe segments for non-rigid alignment. If rigid alignment is performed, + this a vector of ones with length (spatial_bin_centers,) + histogram_info_list :list[dict] + see `_get_single_session_activity_histogram()` for details. + motion_objects_list : + List of motion objects containing the shifts and spatial and temporal + bins for each recording. Note this contains only displacement + associated with the inter-session alignment, and so will differ from + the motion on corrected recording objects if the recording is + already an `InterpolateMotionRecording` object containing + within-session motion correction. + corrected : dict + Dictionary containing corrected-histogram + information. + corrected_peak_locations_list : + Displacement-corrected `peak_locations`. + corrected_session_histogram_list : + Corrected activity histogram (computed from the corrected peak locations). + """ + non_rigid_window_kwargs = copy.deepcopy(non_rigid_window_kwargs) + estimate_histogram_kwargs = copy.deepcopy(estimate_histogram_kwargs) + compute_alignment_kwargs = copy.deepcopy(compute_alignment_kwargs) + interpolate_motion_kwargs = copy.deepcopy(interpolate_motion_kwargs) + + # Ensure list lengths match and all channel locations are the same across recordings. + _check_align_sessions_inputs( + recordings_list, + peaks_list, + peak_locations_list, + alignment_order, + estimate_histogram_kwargs, + interpolate_motion_kwargs, + ) + + print("Computing a single activity histogram from each session...") + + session_histogram_list, temporal_bin_centers_list, spatial_bin_centers, spatial_bin_edges, histogram_info_list = ( + _compute_session_histograms(recordings_list, peaks_list, peak_locations_list, **estimate_histogram_kwargs) + ) + + print("Aligning the activity histograms across sessions...") + + contact_depths = recordings_list[0].get_channel_locations()[:, 1] + + shifts_array, non_rigid_windows, non_rigid_window_centers = _compute_session_alignment( + session_histogram_list, + contact_depths, + spatial_bin_centers, + alignment_order, + non_rigid_window_kwargs, + compute_alignment_kwargs, + ) + shifts_array *= estimate_histogram_kwargs["bin_um"] + + print("Creating corrected recordings...") + + corrected_recordings_list, motion_objects_list = _create_motion_recordings( + recordings_list, shifts_array, temporal_bin_centers_list, non_rigid_window_centers, interpolate_motion_kwargs + ) + + print("Creating corrected peak locations and histograms...") + + corrected_peak_locations_list, corrected_session_histogram_list = _correct_session_displacement( + corrected_recordings_list, + peaks_list, + peak_locations_list, + motion_objects_list, + spatial_bin_edges, + estimate_histogram_kwargs, + ) + + extra_outputs_dict = { + "shifts_array": shifts_array, + "session_histogram_list": session_histogram_list, + "spatial_bin_centers": spatial_bin_centers, + "temporal_bin_centers_list": temporal_bin_centers_list, + "non_rigid_window_centers": non_rigid_window_centers, + "non_rigid_windows": non_rigid_windows, + "histogram_info_list": histogram_info_list, + "motion_objects_list": motion_objects_list, + "corrected": { + "corrected_peak_locations_list": corrected_peak_locations_list, + "corrected_session_histogram_list": corrected_session_histogram_list, + }, + } + return corrected_recordings_list, extra_outputs_dict + + +def align_sessions_after_motion_correction( + recordings_list: list[BaseRecording], motion_info_list: list[dict], align_sessions_kwargs: dict | None +) -> tuple[list[BaseRecording], dict]: + """ + Convenience function to run `align_sessions` to correct for + inter-session displacement from the outputs of motion correction. + + The estimated displacement will be added directly to the recording. + + Parameters + ---------- + recordings_list : list[BaseRecording] + A list of motion-corrected (`InterpolateMotionRecording`) recordings. + motion_info_list : list[dict] + A list of `motion_info` objects, as output from `correct_motion`. + Each entry should correspond to a recording in `recording_list`. + align_sessions_kwargs : dict + A dictionary of keyword arguments passed to `align_sessions`. + + """ + # avoid circular imports + from spikeinterface.sortingcomponents.motion import InterpolateMotionRecording + from spikeinterface.sortingcomponents.motion.motion_interpolation import correct_motion_on_peaks + + # Check motion kwargs are the same across all recordings + if not all(isinstance(rec, InterpolateMotionRecording) for rec in recordings_list): + raise ValueError( + "All passed recordings have been run with motion correction as the last step. " + "They must be InterpolateMotionRecording." + ) + + motion_kwargs_list = [info["parameters"]["estimate_motion_kwargs"] for info in motion_info_list] + if not all(kwargs == motion_kwargs_list[0] for kwargs in motion_kwargs_list): + raise ValueError( + "The motion correct settings used on the `recordings_list` must be identical for all recordings" + ) + + motion_window_kwargs = copy.deepcopy(motion_kwargs_list[0]) + + if "direction" in motion_window_kwargs and motion_window_kwargs["direction"] != "y": + raise ValueError("motion correct must have been performed along the 'y' dimension.") + + if align_sessions_kwargs is None: + align_sessions_kwargs = {} + else: + align_sessions_kwargs = copy.deepcopy(align_sessions_kwargs) + # If motion correction was nonrigid, we must use the same settings for + # inter-session alignment, or we will not be able to add the nonrigid + # shifts together. + if ( + "non_rigid_window_kwargs" in align_sessions_kwargs + and not align_sessions_kwargs["non_rigid_window_kwargs"]["rigid"] + ): + if not motion_window_kwargs["rigid"]: + warnings.warn( + "Nonrigid inter-session alignment must use the motion correct " + "nonrigid settings.\n!Now overwriting any passed `non_rigid_window_kwargs` " + "with the motion object's non_rigid_window_kwargs !" + ) + non_rigid_window_kwargs = get_non_rigid_window_kwargs() + + for ( + key, + value, + ) in motion_window_kwargs.items(): + if key in non_rigid_window_kwargs: + non_rigid_window_kwargs[key] = value + + align_sessions_kwargs["non_rigid_window_kwargs"] = non_rigid_window_kwargs + + if "interpolate_motion_kwargs" in align_sessions_kwargs: + raise ValueError( + "Cannot set `interpolate_motion_kwargs` when using this function. " + "The interpolate kwargs from the original motion recording will be used." + ) + # This does not do anything, just makes it explicit these are not used. + align_sessions_kwargs["interpolate_motion_kwargs"] = None + + corrected_peak_locations = [ + correct_motion_on_peaks(info["peaks"], info["peak_locations"], info["motion"], recording) + for info, recording in zip(motion_info_list, recordings_list) + ] + + return align_sessions( + recordings_list, + [info["peaks"] for info in motion_info_list], + corrected_peak_locations, + **align_sessions_kwargs, + ) + + +def compute_peaks_locations_for_session_alignment( + recording_list: list[BaseRecording], + detect_kwargs: dict, + localize_peaks_kwargs: dict, + job_kwargs: dict | None = None, + gather_mode: str = "memory", +): + """ + A convenience function to compute `peaks_list` and `peak_locations_list` + from a list of recordings, for `align_sessions`. + + Parameters + ---------- + recording_list : list[BaseRecording] + A list of recordings to compute `peaks` and + `peak_locations` for. + detect_kwargs : dict + Arguments to be passed to `detect_peaks`. + localize_peaks_kwargs : dict + Arguments to be passed to `localise_peaks`. + job_kwargs : dict | None + `job_kwargs` for `run_node_pipeline()`. + gather_mode : str + The mode for `run_node_pipeline()`. + """ + if job_kwargs is None: + job_kwargs = {} + + peaks_list = [] + peak_locations_list = [] + + for recording in recording_list: + noise_levels = get_noise_levels(recording, return_in_uV=False) + peaks, peak_locations, _ = run_peak_detection_pipeline_node( + recording, noise_levels, gather_mode, detect_kwargs, localize_peaks_kwargs, job_kwargs + ) + peaks_list.append(peaks) + peak_locations_list.append(peak_locations) + + return peaks_list, peak_locations_list + + +############################################################################### +# Private Functions +############################################################################### + + +def _compute_session_histograms( + recordings_list: list[BaseRecording], + peaks_list: list[np.ndarray], + peak_locations_list: list[np.ndarray], + histogram_type, + bin_um: float, + method: str, + chunked_bin_size_s: float | "estimate", + depth_smooth_um: float, + log_transform: bool, + weight_with_amplitude: bool, + avg_in_bin: bool, +) -> tuple[list[np.ndarray], list[np.ndarray], np.ndarray, np.ndarray, list[dict]]: + """ + Compute a 1d activity histogram for the session. As + sessions may be long, the approach taken is to chunk + the recording into time segments and compute separate + histograms for each. Then, a summary statistic is computed + over the histograms. This accounts for periods of noise + in the recording or segments of irregular spiking. + + Parameters + ---------- + see `align_sessions` for `recording_list`, `peaks_list`, + `peak_locations_list`. + + see `get_estimate_histogram_kwargs()` for all other kwargs. + + Returns + ------- + + session_histogram_list : list[np.ndarray] + A list of activity histograms (1 x n_bins), one per session. + This is the histogram which summarises all chunked histograms. + + temporal_bin_centers_list : list[np.ndarray] + A list of temporal bin centers, one per session. We have one + histogram per session, the temporal bin has 1 entry, the + mid-time point of the session. + + spatial_bin_centers : np.ndarray + A list of spatial bin centers corresponding to the session + activity histograms. + + spatial_bin_edges : np.ndarray + The corresponding spatial bin edges + + histogram_info_list : list[dict] + A list of extra information on the histograms generation + (e.g. chunked histograms). One per session. See + `_get_single_session_activity_histogram()` for details. + """ + # avoid circular imports + from spikeinterface.sortingcomponents.motion.motion_utils import get_spatial_bins + + # Get spatial windows (shared across all histograms) + # and estimate the session histograms + temporal_bin_centers_list = [] + + spatial_bin_centers, spatial_bin_edges, _ = get_spatial_bins( + recordings_list[0], direction="y", hist_margin_um=0, bin_um=bin_um + ) + + session_histogram_list = [] + histogram_info_list = [] + + for recording, peaks, peak_locations in zip(recordings_list, peaks_list, peak_locations_list): + session_hist, temporal_bin_centers, histogram_info = _get_single_session_activity_histogram( + recording, + peaks, + peak_locations, + histogram_type=histogram_type, + spatial_bin_edges=spatial_bin_edges, + method=method, + log_transform=log_transform, + chunked_bin_size_s=chunked_bin_size_s, + depth_smooth_um=depth_smooth_um, + weight_with_amplitude=weight_with_amplitude, + avg_in_bin=avg_in_bin, + ) + temporal_bin_centers_list.append(temporal_bin_centers) + session_histogram_list.append(session_hist) + histogram_info_list.append(histogram_info) + + return ( + session_histogram_list, + temporal_bin_centers_list, + spatial_bin_centers, + spatial_bin_edges, + histogram_info_list, + ) + + +def _get_single_session_activity_histogram( + recording: BaseRecording, + peaks: np.ndarray, + peak_locations: np.ndarray, + histogram_type, + spatial_bin_edges: np.ndarray, + method: str, + log_transform: bool, + chunked_bin_size_s: float | "estimate", + depth_smooth_um: float, + weight_with_amplitude: bool, + avg_in_bin: bool, +) -> tuple[np.ndarray, np.ndarray, dict]: + """ + Compute an activity histogram for a single session. + Pool all peaks into a rate histogram or take the bin-wise median of + rate histograms from time chunks. Apply log2(1 + rate), if requested, + only after constructing this session summary. + + For "chunked_median", if `chunked_bin_size_s` is set to `"estimate"` the + histogram for the entire session is first created to get a good + estimate of the firing rates. + The firing rates are used to use a time segment size that will + allow a good estimation of the firing rate. + + Parameters + ---------- + `spatial_bin_edges : np.ndarray + The spatial bin edges for the created histogram. This is + explicitly required as for inter-session alignment, the + session histograms must share bin edges. + + see `_compute_session_histograms()` for all other keyword arguments. + + Returns + ------- + session_histogram : np.ndarray + Summary activity histogram for the session. + temporal_bin_centers : np.ndarray + Temporal bin center (session mid-point as we only have + one time point) for the session. + histogram_info : dict + A dict of additional info including: + "chunked_histograms" : The rate histograms before log transformation + over which the summary histogram was calculated. + "chunked_temporal_bin_centers" : The temporal vin centers + for the chunked histograms, with length num_chunks. + "session_std" : The mean across bin-wise standard deviation + of the chunked histograms. + "chunked_bin_size_s" : time of each chunk used to + calculate the chunked histogram. + """ + times = recording.get_times() + temporal_bin_centers = np.atleast_1d((times[-1] + times[0]) / 2) + + # Entire-session histograms use one temporal bin and need no chunk-size estimation. + if method == "entire_session": + chunked_bin_size_s = recording.get_duration() + elif chunked_bin_size_s == "estimate": + scaled_hist, _, _ = alignment_utils.get_2d_activity_histogram( + recording, + peaks, + peak_locations, + spatial_bin_edges, + bin_s=None, + depth_smooth_um=None, + scale_to_hz=True, + weight_with_amplitude=False, + avg_in_bin=False, + ) + + # It is important that the passed histogram is scaled to firing rate in Hz + chunked_bin_size_s = alignment_utils.estimate_chunk_size(scaled_hist) + chunked_bin_size_s = np.min([chunked_bin_size_s, recording.get_duration()]) + + if histogram_type == "1d": + chunked_histograms, chunked_temporal_bin_centers, _ = alignment_utils.get_2d_activity_histogram( + recording, + peaks, + peak_locations, + spatial_bin_edges, + bin_s=chunked_bin_size_s, + depth_smooth_um=depth_smooth_um, + weight_with_amplitude=weight_with_amplitude, + avg_in_bin=avg_in_bin, + scale_to_hz=True, + ) + + elif histogram_type in ["2d"]: + if histogram_type == "2d": + # avoid circular imports + from spikeinterface.sortingcomponents.motion.motion_utils import make_3d_motion_histograms + + chunked_histograms, chunked_temporal_bin_edges, _ = make_3d_motion_histograms( + recording, + peaks, + peak_locations, + direction="y", + bin_s=chunked_bin_size_s, + bin_um=None, + hist_margin_um=50, + num_amp_bins=20, + log_transform=False, + spatial_bin_edges=spatial_bin_edges, + ) + + chunked_temporal_bin_centers = alignment_utils.get_bin_centers(chunked_temporal_bin_edges) + chunked_histograms = alignment_utils.scale_histogram_to_hz( + chunked_histograms, chunked_temporal_bin_edges, recording.get_duration() + ) + + if method == "entire_session": + session_histogram = chunked_histograms[0] + + elif method == "chunked_median": + session_histogram = alignment_utils.get_chunked_hist_median(chunked_histograms) + + if log_transform: + session_histogram = np.log2(1 + session_histogram) + + histogram_info = { + "chunked_histograms": chunked_histograms, + "chunked_temporal_bin_centers": chunked_temporal_bin_centers, + "chunked_bin_size_s": chunked_bin_size_s, + } + + return session_histogram, temporal_bin_centers, histogram_info + + +def _create_motion_recordings( + recordings_list: list[BaseRecording], + shifts_array: np.ndarray, + temporal_bin_centers_list: list[np.ndarray], + non_rigid_window_centers: np.ndarray, + interpolate_motion_kwargs: dict, +) -> tuple[list[BaseRecording], list[Motion]]: + """ + Given a set of recordings, motion shifts and bin information per-recording, + generate an InterpolateMotionRecording. If the recording is already an + InterpolateMotionRecording, then the shifts will be added to a copy + of it. Copies of the Recordings are made, nothing is changed in-place. + + Parameters + ---------- + shifts_array : num_sessions x num_nonrigid bins + + Returns + ------- + corrected_recordings_list : list[BaseRecording] + A list of InterpolateMotionRecording recordings of shift-corrected + recordings corresponding to `recordings_list`. + + motion_objects_list : list[Motion] + A list of Motion objects. If the recording in `recordings_list` + is already an InterpolateMotionRecording, this will be `None`, as + no motion object is created (the existing motion object is added to) + """ + # avoid circular imports + from spikeinterface.sortingcomponents.motion import InterpolateMotionRecording + + assert all(array.ndim == 1 for array in shifts_array), "time dimension should be 1 for session displacement" + + corrected_recordings_list = [] + motion_objects_list = [] + for ses_idx, recording in enumerate(recordings_list): + session_shift = shifts_array[ses_idx][np.newaxis, :] + + motion = Motion([session_shift], [temporal_bin_centers_list[ses_idx]], non_rigid_window_centers, direction="y") + motion_objects_list.append(motion) + + if isinstance(recording, InterpolateMotionRecording): + print("Recording is already an `InterpolateMotionRecording. Adding shifts directly the recording object.") + + corrected_recording = _add_displacement_to_interpolate_recording(recording, motion) + else: + corrected_recording = InterpolateMotionRecording( + recording, + motion, + interpolation_time_bin_centers_s=motion.temporal_bins_s, + interpolation_time_bin_edges_s=[np.array([recording.get_times()[0], recording.get_times()[-1]])], + **interpolate_motion_kwargs, + ) + corrected_recording.set_probe(recording.get_probe()) + + corrected_recordings_list.append(corrected_recording) + + return corrected_recordings_list, motion_objects_list + + +def _add_displacement_to_interpolate_recording( + original_recording: BaseRecording, + session_displacement_motion: Motion, +): + """ + This function adds a shift to an InterpolateMotionRecording. + + There are four cases: + - The original recording is rigid and new shift is rigid (shifts are added). + - The original recording is rigid and new shifts are non-rigid (sets the + non-rigid shifts onto the recording, then adds back the original shifts). + - The original recording is nonrigid and the new shifts are rigid (rigid + shift added to all nonlinear shifts) + - The original recording is nonrigid and the new shifts are nonrigid + (respective non-rigid shifts are added, must have same number of + non-rigid windows). + + Parameters + ---------- + see `_create_motion_recordings()` + + Returns + ------- + corrected_recording : InterpolateMotionRecording + A copy of the `recording` with new shifts added. + + """ + # Everything is done in place, so keep a short variable + # name reference to the new recordings `motion` object + # and update it.okay + corrected_recording = copy.deepcopy(original_recording) + + shifts_to_add = session_displacement_motion.displacement[0] + new_non_rigid_window_centers = session_displacement_motion.spatial_bins_um + + motion_ref = corrected_recording._recording_segments[0].motion + recording_bins = motion_ref.displacement[0].shape[1] + + # If the new displacement is a scalar (i.e. rigid), + # just add it to the existing displacements + if shifts_to_add.shape[1] == 1: + motion_ref.displacement[0] += shifts_to_add[0, 0] + + else: + if recording_bins == 1: + # If the new displacement is nonrigid (multiple windows) but the motion + # recording is rigid, we update the displacement at all time bins + # with the new, nonrigid displacement added to the old, rigid displacement. + num_time_bins = motion_ref.displacement[0].shape[0] + tiled_nonrigid_displacement = np.repeat(shifts_to_add, num_time_bins, axis=0) + shifts_to_add = tiled_nonrigid_displacement + motion_ref.displacement + + motion_ref.displacement = shifts_to_add + motion_ref.spatial_bins_um = new_non_rigid_window_centers + # Cached interpolators still reference the original rigid array and grid. + motion_ref.interpolators = None + else: + # Otherwise, if both the motion and new displacement are + # nonrigid, we need to make sure the nonrigid windows + # match exactly. + assert np.array_equal(motion_ref.spatial_bins_um, new_non_rigid_window_centers) + assert motion_ref.displacement[0].shape[1] == shifts_to_add.shape[1] + + motion_ref.displacement[0] += shifts_to_add + + return corrected_recording + + +def _correct_session_displacement( + recordings_list: list[BaseRecording], + peaks_list: list[np.ndarray], + peak_locations_list: list[np.ndarray], + motion_objects_list: list[Motion], + spatial_bin_edges: np.ndarray, + estimate_histogram_kwargs: dict, +): + """ + Internal function to apply the correction from `align_sessions` + to build a corrected histogram for comparison. First, create + new shifted peak locations. Then, create a new 'corrected' + activity histogram from the new peak locations. + + Parameters + ---------- + see `align_sessions()` for parameters. + + Returns + ------- + corrected_peak_locations_list : list[np.ndarray] + A list of peak locations corrected by the inter-session + shifts (one entry per session). + corrected_session_histogram_list : list[np.ndarray] + A list of histograms calculated from the corrected peaks (one per session). + """ + # avoid circular imports + from spikeinterface.sortingcomponents.motion.motion_interpolation import correct_motion_on_peaks + + corrected_peak_locations_list = [] + + for recording, peaks, peak_locations, motion in zip( + recordings_list, peaks_list, peak_locations_list, motion_objects_list + ): + # Note this `motion` is not necessarily the same as the motion on the recording. If the recording + # is an `InterpolateMotionRecording`, it will contain correction for both motion and inter-session displacement. + # Here we want to correct only the motion associated with inter-session displacement. + corrected_peak_locs = correct_motion_on_peaks( + peaks, + peak_locations, + motion, + recording, + ) + corrected_peak_locations_list.append(corrected_peak_locs) + + corrected_session_histogram_list = [] + + for recording, peaks, corrected_locations in zip(recordings_list, peaks_list, corrected_peak_locations_list): + session_hist, _, _ = _get_single_session_activity_histogram( + recording, + peaks, + corrected_locations, + estimate_histogram_kwargs["histogram_type"], + spatial_bin_edges, + estimate_histogram_kwargs["method"], + estimate_histogram_kwargs["log_transform"], + estimate_histogram_kwargs["chunked_bin_size_s"], + estimate_histogram_kwargs["depth_smooth_um"], + estimate_histogram_kwargs["weight_with_amplitude"], + estimate_histogram_kwargs["avg_in_bin"], + ) + corrected_session_histogram_list.append(session_hist) + + return corrected_peak_locations_list, corrected_session_histogram_list + + +######################################################################################################################## + + +def _compute_session_alignment( + session_histogram_list: list[np.ndarray], + contact_depths: np.ndarray, + spatial_bin_centers: np.ndarray, + alignment_order: str, + non_rigid_window_kwargs: dict, + compute_alignment_kwargs: dict, +) -> tuple[np.ndarray, ...]: + """ + Given a list of activity histograms (one per session) compute + rigid or non-rigid set of shifts (one per session) that will bring + all sessions into alignment. + + For rigid shifts, a cross-correlation between activity + histograms is performed. For non-rigid shifts, the probe + is split into segments, and linear estimation of shift + performed for each segment. + + Parameters + ---------- + See `align_sessions()` for parameters + + Returns + ------- + shifts : np.ndarray + A (num_sessions x num_rigid_windows) array of shifts to bring + the histograms in `session_histogram_list` into alignment. + non_rigid_windows : np.ndarray + An array (num_non_rigid_windows x num_spatial_bins) of weightings + for each bin in each window. For rect, these are in the range [0, 1], + for Gaussian these are gaussian etc. + non_rigid_window_centers : np.ndarray + The centers (spatial, in um) of each non-rigid window. + """ + # avoid circular imports + from spikeinterface.sortingcomponents.motion.motion_utils import get_spatial_windows + + session_histogram_array = np.array(session_histogram_list) + + akima_interp_nonrigid = compute_alignment_kwargs.pop("akima_interp_nonrigid") + num_shifts_global = compute_alignment_kwargs.pop("num_shifts_global") + num_shifts_block = compute_alignment_kwargs.pop("num_shifts_block") + + non_rigid_windows, non_rigid_window_centers = get_spatial_windows( + contact_depths, spatial_bin_centers, **non_rigid_window_kwargs + ) + + rigid_shifts = _estimate_rigid_alignment( + session_histogram_array, + alignment_order, + num_shifts_global, + compute_alignment_kwargs, + ) + + if non_rigid_window_kwargs["rigid"]: + return rigid_shifts, non_rigid_windows, non_rigid_window_centers + + # For non-rigid, first shift the histograms according to the rigid shift + shifted_histograms = np.zeros_like(session_histogram_array) + for ses_idx, orig_histogram in enumerate(session_histogram_array): + shifted_histogram = alignment_utils.shift_array_fill_zeros( + array=orig_histogram, shift=int(rigid_shifts[ses_idx, 0]) + ) + shifted_histograms[ses_idx, :] = shifted_histogram + + # Then compute the nonrigid shifts + nonrigid_session_offsets_matrix, _ = alignment_utils.compute_histogram_crosscorrelation( + shifted_histograms, non_rigid_windows, num_shifts=num_shifts_block, **compute_alignment_kwargs + ) + non_rigid_shifts = alignment_utils.get_shifts_from_session_matrix(alignment_order, nonrigid_session_offsets_matrix) + + # Akima interpolate the nonrigid bins if required. + if akima_interp_nonrigid: + interp_nonrigid_shifts = alignment_utils.akima_interpolate_nonrigid_shifts( + non_rigid_shifts, non_rigid_window_centers, spatial_bin_centers + ) + shifts = rigid_shifts + interp_nonrigid_shifts + non_rigid_window_centers = spatial_bin_centers + else: + shifts = rigid_shifts + non_rigid_shifts + + return shifts, non_rigid_windows, non_rigid_window_centers + + +def _estimate_rigid_alignment( + session_histogram_array: np.ndarray, + alignment_order: str, + num_shifts: None | int, + compute_alignment_kwargs: dict, +): + """ + Estimate the rigid alignment from a set of activity + histograms, using simple cross-correlation. + + Parameters + ---------- + session_histogram_array : np.ndarray + A (num_sessions x num_spatial_bins) array of activity + histograms to align + alignment_order : str + Align "to_middle" or "to_session_N" (where "N" is the session number) + compute_alignment_kwargs : dict + See `get_compute_alignment_kwargs()`. + + Returns + ------- + optimal_shift_indices : np.ndarray + An array (num_sessions x 1) of shifts to bring all + session histograms into alignment. + """ + compute_alignment_kwargs = copy.deepcopy(compute_alignment_kwargs) + + rigid_window = np.ones(session_histogram_array.shape[1])[np.newaxis, :] + + rigid_session_offsets_matrix, _ = alignment_utils.compute_histogram_crosscorrelation( + session_histogram_array, + rigid_window, + num_shifts=num_shifts, + **compute_alignment_kwargs, # TODO: remove the copy above and pass directly. Consider removing this function... + ) + optimal_shift_indices = alignment_utils.get_shifts_from_session_matrix( + alignment_order, rigid_session_offsets_matrix + ) + + return optimal_shift_indices + + +# ----------------------------------------------------------------------------- +# Checkers +# ----------------------------------------------------------------------------- + + +def _check_align_sessions_inputs( + recordings_list: list[BaseRecording], + peaks_list: list[np.ndarray], + peak_locations_list: list[np.ndarray], + alignment_order: str, + estimate_histogram_kwargs: dict, + interpolate_motion_kwargs: dict, +): + """ + Perform checks on the input of `align_sessions()` + """ + num_sessions = len(recordings_list) + + if len(peaks_list) != num_sessions or len(peak_locations_list) != num_sessions: + raise ValueError( + "`recordings_list`, `peaks_list` and `peak_locations_list` " + "must be the same length. They must contains list of corresponding " + "recordings, peak and peak location objects." + ) + + if not all(rec.get_num_segments() == 1 for rec in recordings_list): + raise ValueError( + "Multi-segment recordings not supported. All recordings in `recordings_list` but have only 1 segment." + ) + + channel_locs = [rec.get_channel_locations() for rec in recordings_list] + if not all([np.array_equal(locs, channel_locs[0]) for locs in channel_locs]): + raise ValueError( + "The recordings in `recordings_list` do not all have " + "the same channel locations. All recordings must be " + "performed using the same probe." + ) + + accepted_hist_methods = [ + "entire_session", + "chunked_median", + ] + method = estimate_histogram_kwargs["method"] + if method not in accepted_hist_methods: + raise ValueError(f"`method` option must be one of: {accepted_hist_methods}") + + if alignment_order != "to_middle": + split_name = alignment_order.split("_") + if not "_".join(split_name[:2]) == "to_session": + raise ValueError( + "`alignment_order` must take be 'to_middle' or take the form 'to_session_X' where X is the session number to align to." + ) + + ses_num = int(split_name[-1]) + if ses_num > num_sessions: + raise ValueError( + f"`alignment_order` session {ses_num} is larger than the number of sessions in `recordings_list`." + ) + + if ses_num == 0: + raise ValueError("`alignment_order` required the session number, not session index.") diff --git a/src/spikeinterface/preprocessing/motion.py b/src/spikeinterface/preprocessing/motion.py index 2629b118c1..616fbf6197 100644 --- a/src/spikeinterface/preprocessing/motion.py +++ b/src/spikeinterface/preprocessing/motion.py @@ -331,10 +331,9 @@ def compute_motion( """ # local import are important because "sortingcomponents" is not important by default - from spikeinterface.sortingcomponents.peak_detection import detect_peaks, detect_peak_methods + from spikeinterface.sortingcomponents.peak_detection import detect_peaks from spikeinterface.sortingcomponents.peak_selection import select_peaks - from spikeinterface.sortingcomponents.peak_localization import localize_peaks, peak_localization_methods - from spikeinterface.core.node_pipeline import ExtractDenseWaveforms, run_node_pipeline + from spikeinterface.sortingcomponents.peak_localization import localize_peaks from spikeinterface.sortingcomponents.motion.motion_estimation import estimate_motion # get preset params and update if necessary @@ -374,47 +373,13 @@ def compute_motion( if no_selection_kwargs: # maybe do this directly in the folder when not None, but might be slow on external storage gather_mode = "memory" - # node detect - detect_peaks_method = detect_kwargs["method"] - method_class = detect_peak_methods[detect_peaks_method] - detect_kwargs_without_method = { - key: detect_kwarg for key, detect_kwarg in detect_kwargs.items() if key != "method" - } - if method_class.need_noise_levels: - detect_kwargs_without_method["noise_levels"] = noise_levels - - node0 = method_class(recording, **detect_kwargs_without_method) - - node1 = ExtractDenseWaveforms(recording, parents=[node0], ms_before=0.1, ms_after=0.3) - # node detect + localize - method = localize_peaks_kwargs["method"] - method_class = peak_localization_methods[method] - localize_peaks_kwargs_without_method = { - key: localize_peaks_kwarg for key, localize_peaks_kwarg in localize_peaks_kwargs.items() if key != "method" - } - node2 = method_class( - recording, parents=[node0, node1], return_output=True, **localize_peaks_kwargs_without_method - ) - pipeline_nodes = [node0, node1, node2] - t0 = time.perf_counter() - peaks, peak_locations = run_node_pipeline( - recording, - pipeline_nodes, - job_kwargs, - job_name="detect and localize", - gather_mode=gather_mode, - gather_kwargs=None, - squeeze_output=False, - dest=None, - names=None, + peaks, peak_locations, run_times = run_peak_detection_pipeline_node( + recording, noise_levels, gather_mode, detect_kwargs, localize_peaks_kwargs, job_kwargs ) - t1 = time.perf_counter() - run_times = dict(detect_and_localize=t1 - t0) + else: # localization is done after select_peaks() - pipeline_nodes = None - t0 = time.perf_counter() method_kwargs = detect_kwargs.copy() method_kwargs["noise_levels"] = noise_levels @@ -461,6 +426,50 @@ def compute_motion( return motion +def run_peak_detection_pipeline_node( + recording, noise_levels, gather_mode, detect_kwargs, localize_peaks_kwargs, job_kwargs +): + + from spikeinterface.sortingcomponents.peak_detection import detect_peak_methods + from spikeinterface.core.node_pipeline import ExtractDenseWaveforms, run_node_pipeline + from spikeinterface.sortingcomponents.peak_localization import peak_localization_methods + + detect_peaks_method = detect_kwargs["method"] + method_class = detect_peak_methods[detect_peaks_method] + detect_kwargs_without_method = {key: detect_kwarg for key, detect_kwarg in detect_kwargs.items() if key != "method"} + if method_class.need_noise_levels: + detect_kwargs_without_method["noise_levels"] = noise_levels + + node0 = method_class(recording, **detect_kwargs_without_method) + + node1 = ExtractDenseWaveforms(recording, parents=[node0], ms_before=0.1, ms_after=0.3) + + # node detect + localize + method = localize_peaks_kwargs["method"] + method_class = peak_localization_methods[method] + localize_peaks_kwargs_without_method = { + key: localize_peaks_kwarg for key, localize_peaks_kwarg in localize_peaks_kwargs.items() if key != "method" + } + node2 = method_class(recording, parents=[node0, node1], return_output=True, **localize_peaks_kwargs_without_method) + pipeline_nodes = [node0, node1, node2] + t0 = time.perf_counter() + peaks, peak_locations = run_node_pipeline( + recording, + pipeline_nodes, + job_kwargs, + job_name="detect and localize", + gather_mode=gather_mode, + gather_kwargs=None, + squeeze_output=False, + dest=None, + names=None, + ) + t1 = time.perf_counter() + run_times = dict(detect_and_localize=t1 - t0) + + return peaks, peak_locations, run_times + + def correct_motion( recording: BaseRecording, preset: Literal[ @@ -532,7 +541,6 @@ def correct_motion( Optional output if `output_motion_info=True`. This dict contains several variable for for plotting. See `plot_motion_info()` """ - from spikeinterface.sortingcomponents.motion import interpolate_motion detect_kwargs, select_kwargs, localize_peaks_kwargs, estimate_motion_kwargs = _update_motion_kwargs( diff --git a/src/spikeinterface/preprocessing/tests/test_inter_session_alignment.py b/src/spikeinterface/preprocessing/tests/test_inter_session_alignment.py new file mode 100644 index 0000000000..82036faafd --- /dev/null +++ b/src/spikeinterface/preprocessing/tests/test_inter_session_alignment.py @@ -0,0 +1,1114 @@ +import numpy as np +import pytest + +from spikeinterface.core import Motion +from spikeinterface.preprocessing.inter_session_alignment import session_alignment, alignment_utils +from spikeinterface.generation.session_displacement_generator import * +import spikeinterface # required for monkeypatching +import spikeinterface.full as si +from spikeinterface.sortingcomponents.motion import InterpolateMotionRecording + +DEBUG = False + + +class TestInterSessionAlignment: + @pytest.fixture(scope="session") + def recording_1(self): + """ + Generate a set of session recordings with displacement. + These parameters are chosen such that simulated AP signal is strong + on the probe to avoid noise in the AP positions. This is important + for checking that the estimated shift matches the known shift. + """ + shifts = ((0, 0), (0, -200), (0, 150)) + + recordings_list, _ = generate_session_displacement_recordings( + num_units=15, + recording_durations=[0.1, 0.2, 0.3], + recording_shifts=shifts, + non_rigid_gradient=None, + seed=55, + generate_sorting_kwargs=dict(firing_rates=(100, 250), refractory_period_ms=4.0), + generate_unit_locations_kwargs=dict( + margin_um=0.0, + 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, 0.0), spatial_decay=1.0), + ) + + 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"}, + ) + return (recordings_list, shifts, peaks_list, peak_locations_list) + + # TODO: need to make this a fixtures somehow, I guess rigid or nonrigid case. + @pytest.fixture(scope="session") + def recording_2(self): + """ + Get a shifted inter-session alignment recording. First, interpolate-motion + within session. The purpose of these tests is then to run inter-session alignment + and check the displacements are properly added. + """ + shifts = ((0, 0), (0, 250)) + + recordings_list, _ = generate_session_displacement_recordings( + num_units=5, + recording_durations=[0.3, 0.3], + recording_shifts=shifts, + non_rigid_gradient=0.2, + 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, + 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, 0.5), spatial_decay=1.0), + # must have some noise, or peak detection becomes completely stoachastic + # because it relies on the std to set the threshold. + ) + + 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"}, + ) + + return (recordings_list, shifts, peaks_list, peak_locations_list) + + def motion_correct_recordings_list(self, recordings_list, rigid_motion): + # Unfortunately this is necessary only in the test environemnt + # because adding offsets means copying the motion object does not work + interpolate_motion_kwargs = {"border_mode": "force_zeros"} + localize_peaks_kwargs = {"method": "grid_convolution"} + + preset = "rigid_fast" if rigid_motion else "kilosort_like" + + # Perform a motion correction, note this is just to make the + # motion correction object with the correct displacment, but + # the displacements should be zero here. These are manulally + # added in the tetsts. + mc_recording_list = [] + mc_motion_info_list = [] + for rec in recordings_list: + corrected_rec, motion_info = si.correct_motion( + rec, + preset=preset, + interpolate_motion_kwargs=interpolate_motion_kwargs, + output_motion_info=True, + localize_peaks_kwargs=localize_peaks_kwargs, + ) + mc_recording_list.append(corrected_rec) + mc_motion_info_list.append(motion_info) + + return mc_recording_list, mc_motion_info_list # , shifts + + ########################################################################### + # Functional Tests + ############################################################################ + + @pytest.mark.parametrize("histogram_type", ["1d", "2d"]) + @pytest.mark.parametrize("num_shifts_global", [None, 200]) + def test_align_sessions_finds_correct_shifts(self, num_shifts_global, recording_1, histogram_type): + """ + Test that `align_sessions` recovers the correct (linear) shifts. + """ + recordings_list, shifts, peaks_list, peak_locations_list = recording_1 + bin_um = 2 + + assert shifts == ( + (0, 0), + (0, -200), + (0, 150), + ), "expected shifts are hard-coded into this test ahould should be set in the fixture.." + + compute_alignment_kwargs = session_alignment.get_compute_alignment_kwargs() + compute_alignment_kwargs["smoothing_sigma_bin"] = None + compute_alignment_kwargs["smoothing_sigma_window"] = None + compute_alignment_kwargs["num_shifts_global"] = num_shifts_global + + estimate_histogram_kwargs = session_alignment.get_estimate_histogram_kwargs() + estimate_histogram_kwargs["bin_um"] = bin_um + estimate_histogram_kwargs["histogram_type"] = histogram_type + estimate_histogram_kwargs["log_transform"] = True + + for mode, expected in zip( + ["to_session_1", "to_session_2", "to_session_3", "to_middle"], + [ + (0, -200, 150), + (200, 0, 350), + (-150, -350, 0), + (16.66, -183.33, 166.66), + ], + ): + corrected_recordings_list, extra_info = session_alignment.align_sessions( + recordings_list, + peaks_list, + peak_locations_list, + alignment_order=mode, + compute_alignment_kwargs=compute_alignment_kwargs, + estimate_histogram_kwargs=estimate_histogram_kwargs, + ) + + if DEBUG: + from spikeinterface.widgets import plot_session_alignment, plot_activity_histogram_2d + import matplotlib.pyplot as plt + + plot = 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() + + assert np.allclose(expected, extra_info["shifts_array"].squeeze(), rtol=0, atol=bin_um) + + corr_peaks_list, corr_peak_loc_list = session_alignment.compute_peaks_locations_for_session_alignment( + corrected_recordings_list, + detect_kwargs={"method": "locally_exclusive"}, + localize_peaks_kwargs={"method": "grid_convolution"}, + ) + + new_histograms = session_alignment._compute_session_histograms( + corrected_recordings_list, corr_peaks_list, corr_peak_loc_list, **estimate_histogram_kwargs + )[0] + + rows, cols = np.triu_indices(len(new_histograms), k=1) + assert np.all( + np.abs(np.corrcoef([hist.flatten() for hist in new_histograms])[rows, cols]) + - np.abs(np.corrcoef([hist.flatten() for hist in extra_info["session_histogram_list"]])[rows, cols]) + >= 0 + ) + + def test_histogram_generation(self, recording_1): + """ """ + recordings_list, _, peaks_list, peak_locations_list = recording_1 + + recording = recordings_list[0] + + channel_locations = recording.get_channel_locations() + loc_start = np.min(channel_locations[:, 1]) + loc_end = np.max(channel_locations[:, 1]) + + # Test some floats as slightly more complex than integer case + bin_s = 1.5 + bin_um = 1.5 + + estimate_histogram_kwargs = session_alignment.get_estimate_histogram_kwargs() + estimate_histogram_kwargs["bin_um"] = bin_um + estimate_histogram_kwargs["method"] = "chunked_median" + estimate_histogram_kwargs["chunked_bin_size_s"] = bin_s + estimate_histogram_kwargs["log_transform"] = False + + ( + session_histogram_list, + temporal_bin_centers_list, + spatial_bin_centers, + spatial_bin_edges, + histogram_info_list, + ) = session_alignment._compute_session_histograms( + recordings_list, peaks_list, peak_locations_list, **estimate_histogram_kwargs + ) + + num_bins = (loc_end - loc_start) / bin_um + + bin_edges = np.linspace(loc_start, loc_end, int(num_bins) + 1) + bin_centers = bin_edges[:-1] + bin_um / 2 + + assert np.array_equal(bin_edges, spatial_bin_edges) + assert np.array_equal(bin_centers, spatial_bin_centers) + + for recording, temporal_bin_center in zip(recordings_list, temporal_bin_centers_list): + times = recording.get_times() + centers = (np.max(times) - np.min(times)) / 2 + assert temporal_bin_center == centers + + for ses_idx, (recording, chunked_histogram_info) in enumerate(zip(recordings_list, histogram_info_list)): + # TODO: this is direct copy from above, can merge + times = recording.get_times() + + chunk_time_window = chunked_histogram_info["chunked_bin_size_s"] + + num_windows = (np.ceil(np.max(times)) - np.min(times)) / chunk_time_window + temp_bin_edges = np.arange(np.ceil(num_windows) + 1) * chunk_time_window + centers = temp_bin_edges[:-1] + chunk_time_window / 2 + + assert chunked_histogram_info["chunked_bin_size_s"] == chunk_time_window + assert np.array_equal(chunked_histogram_info["chunked_temporal_bin_centers"], centers) + + for edge_idx in range(len(temp_bin_edges) - 1): + lower = temp_bin_edges[edge_idx] + upper = temp_bin_edges[edge_idx + 1] + + lower_idx = recording.time_to_sample_index(lower) + upper_idx = recording.time_to_sample_index(upper) + + new_peak_locs = peak_locations_list[ses_idx][ + np.where( + np.logical_and( + peaks_list[ses_idx]["sample_index"] >= lower_idx, + peaks_list[ses_idx]["sample_index"] < upper_idx, + ) + ) + ] + assert np.allclose( + np.histogram(new_peak_locs["y"], bins=bin_edges)[0] + / (min(upper, times[0] + recording.get_duration()) - lower), + chunked_histogram_info["chunked_histograms"][edge_idx, :], + rtol=0, + atol=1e-6, + ) + + @pytest.mark.parametrize("histogram_type", ["1d", "2d"]) + @pytest.mark.parametrize("log_transform", [False, True]) + def test_entire_session_histogram(self, histogram_type, log_transform, monkeypatch): + from spikeinterface.core import NumpyRecording + + recording = NumpyRecording([np.zeros((2500, 2), dtype="float32")], sampling_frequency=1000) + peaks = np.array( + [(100, 1.0), (200, 10.0), (1700, 100.0), (2400, 1000.0)], + dtype=[("sample_index", "int64"), ("amplitude", "float64")], + ) + locations = np.array([(0.5,), (0.5,), (0.5,), (1.5,)], dtype=[("y", "float64")]) + if histogram_type == "1d": + expected = np.array([3.0, 1.0]) + else: + expected = np.zeros((2, 20)) + expected[0, [0, 6, 13]] = 1 + expected[1, 19] = 1 + expected /= recording.get_duration() + + monkeypatch.setattr( + alignment_utils, + "estimate_chunk_size", + lambda _: pytest.fail("Entire-session mode must not estimate chunks"), + ) + for chunk_size in ["estimate", 0.6]: + histogram, _, info = session_alignment._get_single_session_activity_histogram( + recording, + peaks, + locations, + histogram_type=histogram_type, + spatial_bin_edges=np.array([0.0, 1.0, 2.0]), + method="entire_session", + log_transform=log_transform, + chunked_bin_size_s=chunk_size, + depth_smooth_um=None, + weight_with_amplitude=False, + avg_in_bin=False, + ) + np.testing.assert_allclose(histogram, np.log2(1 + expected) if log_transform else expected) + np.testing.assert_allclose(info["chunked_histograms"], expected[None, ...]) + np.testing.assert_allclose(info["chunked_temporal_bin_centers"], [1.25]) + assert info["chunked_bin_size_s"] == 2.5 + + @pytest.mark.parametrize("histogram_type", ["1d", "2d"]) + @pytest.mark.parametrize("method", ["entire_session", "chunked_median"]) + def test_histogram_log_transform(self, histogram_type, method): + """Check rates and apply the log after aggregation, including a partial final chunk.""" + from spikeinterface.core import NumpyRecording + + recording = NumpyRecording([np.zeros((3000, 2), dtype="float32")], sampling_frequency=1000) + peaks = np.array( + [(100, 1.0), (200, 1.0), (300, 1.0), (400, 10.0), (2100, 1.0), (2200, 10.0)], + dtype=[("sample_index", "int64"), ("amplitude", "float64")], + ) + locations = np.array([(0.5,), (0.5,), (0.5,), (1.5,), (0.5,), (1.5,)], dtype=[("y", "float64")]) + + if histogram_type == "1d": + counts = np.array([[3.0, 1.0], [1.0, 1.0]]) + else: + counts = np.zeros((2, 2, 20)) + counts[:, 0, 0] = [3, 1] + counts[:, 1, 19] = [1, 1] + + if method == "entire_session": + expected_chunks = counts.sum(axis=0, keepdims=True) / 3.0 + expected = expected_chunks[0] + else: + # The first chunk lasts two seconds; the final chunk lasts only one. + expected_chunks = counts.copy() + expected_chunks[0] /= 2.0 + expected = np.median(expected_chunks, axis=0) + # With two unequal chunk rates, taking logs before the median would differ. + assert not np.allclose(np.log2(1 + expected), np.median(np.log2(1 + expected_chunks), axis=0)) + + for log_transform in [False, True]: + histogram, _, info = session_alignment._get_single_session_activity_histogram( + recording, + peaks, + locations, + histogram_type=histogram_type, + spatial_bin_edges=np.array([0.0, 1.0, 2.0]), + method=method, + log_transform=log_transform, + chunked_bin_size_s=2.0, + depth_smooth_um=None, + weight_with_amplitude=False, + avg_in_bin=False, + ) + np.testing.assert_allclose(info["chunked_histograms"], expected_chunks) + np.testing.assert_allclose(histogram, np.log2(1 + expected) if log_transform else expected) + + ########################################################################### + # Following Motion Correction + ########################################################################### + # These tests check that the displacement found by the inter-session alignment + # is correctly added to any existing motion-correction results. + + def test_rigid_motion_rigid_intersession(self, recording_1): + """ + Create an inter-session alignment recording and motion correct it so that + it is an InterpolateMotion recording. Add some shifts to the existing displacement + on the InterpolateMotion recordings and check the inter-session alignment shifts + are properly added to this. + """ + recordings_list, shifts, _, _ = recording_1 + + mc_recording_list, mc_motion_info_list = self.motion_correct_recordings_list( + recordings_list, + rigid_motion=True, + ) + first_ses_mc_displacement = mc_recording_list[0]._recording_segments[0].motion.displacement + second_ses_mc_displacement = mc_recording_list[1]._recording_segments[0].motion.displacement + + # Ensure the motion was generated rigid by the test suite + assert first_ses_mc_displacement[0].size == 1 + assert second_ses_mc_displacement[0].size == 1 + + # Add some shifts to represent an existing motion correction + first_ses_mc_displacement[0] += 0.01 + second_ses_mc_displacement[0] += 0.02 + + non_rigid_window_kwargs = session_alignment.get_non_rigid_window_kwargs() + non_rigid_window_kwargs["rigid"] = True + + corrected_recordings, extra_info = session_alignment.align_sessions_after_motion_correction( + mc_recording_list, + mc_motion_info_list, + align_sessions_kwargs={ + "alignment_order": "to_session_1", + "non_rigid_window_kwargs": non_rigid_window_kwargs, + }, + ) + first_ses_total_displacement = corrected_recordings[0]._recording_segments[0].motion.displacement + second_ses_total_displacement = corrected_recordings[1]._recording_segments[0].motion.displacement + + # Check that the shift is the existing motion correction + the inter-session shift + assert first_ses_total_displacement == [np.array([[shifts[0][1] + 0.01]])] + assert second_ses_total_displacement == [np.array([[shifts[1][1] + 0.02]])] + + self.assert_interpolate_recording_not_duplicate(corrected_recordings[0]) + + def test_rigid_motion_nonrigid_intersession(self, recording_2): + """ + Test that non-rigid shifts estimated in inter-session alignment are properly + added to rigid shifts estimated in motion correction. + """ + recordings_list, _, peaks_list, _ = recording_2 + + mc_recording_list, mc_motion_info_list = self.motion_correct_recordings_list( + recordings_list, + rigid_motion=True, + ) + + first_ses_mc_displacement = mc_recording_list[0]._recording_segments[0].motion.displacement + second_ses_mc_displacement = mc_recording_list[1]._recording_segments[0].motion.displacement + + # Ensure the motion was generated rigid by the test suite + assert first_ses_mc_displacement[0].size == 1 + assert second_ses_mc_displacement[0].size == 1 + + # Add some shifts to represent an existing motion correction + first_ses_mc_displacement[0] += 0.01 + second_ses_mc_displacement[0] += 0.02 + + # All of this is direct copy between these tests... + non_rigid_window_kwargs = session_alignment.get_non_rigid_window_kwargs() + non_rigid_window_kwargs["rigid"] = False + + corrected_recordings, extra_info = session_alignment.align_sessions_after_motion_correction( + mc_recording_list, + mc_motion_info_list, + align_sessions_kwargs={ + "alignment_order": "to_session_1", + "non_rigid_window_kwargs": non_rigid_window_kwargs, + }, + ) + + first_ses_total_displacement = corrected_recordings[0]._recording_segments[0].motion.displacement + second_ses_total_displacement = corrected_recordings[1]._recording_segments[0].motion.displacement + + # The shift themselves are not expected to be correct to align this tricky test case + # (see test_interesting_debug_case) but the shift on the motion objects should + # match the estimateed shifts from inter-session alignment + the motion shifts set above + assert np.all(extra_info["shifts_array"][0] + 0.01 == first_ses_total_displacement) + assert np.all(extra_info["shifts_array"][1] + 0.02 == second_ses_total_displacement) + + self.assert_interpolate_recording_not_duplicate(corrected_recordings[0]) + + def test_nonrigid_alignment_matches_combined_motion(self, recording_2): + """Check that nonrigid session alignment produces the expected traces. + + Compare the aligned traces with fresh recordings using the combined + motion, and check that the original recordings remain unchanged. + The value of this test is to check that cached objects such as + interpolators are properly refreshed when align_sessions_after_motion_correction + is applied. + """ + # Start with recordings that already have within-session rigid motion correction. + recordings_list, _, _, _ = recording_2 + mc_recording_list, mc_motion_info_list = self.motion_correct_recordings_list(recordings_list, rigid_motion=True) + # Add small intra-session motion offsets (0.01 and 0.02 um) that session alignment must preserve. + for index, recording in enumerate(mc_recording_list): + recording.segments[0].motion.displacement[0] += 0.01 * (index + 1) + + # Copy the traces and motions which have had the extra displacements added + # This initializes cached interpolators, and saves the pre-alignment state. + original_traces = [recording.get_traces() for recording in mc_recording_list] + original_motions = [recording.segments[0].motion.copy() for recording in mc_recording_list] + non_rigid_window_kwargs = session_alignment.get_non_rigid_window_kwargs() + non_rigid_window_kwargs["rigid"] = False + + # Estimate session alignment and combine it with the existing motion correction. + corrected_recordings, extra_info = session_alignment.align_sessions_after_motion_correction( + mc_recording_list, + mc_motion_info_list, + align_sessions_kwargs={ + "alignment_order": "to_session_1", + "non_rigid_window_kwargs": non_rigid_window_kwargs, + }, + ) + + for index, recording in enumerate(recordings_list): + original_motion = original_motions[index] + + # Independently combine the original rigid motion with the estimated session shifts. + expected_motion = Motion( + original_motion.displacement[0] + extra_info["shifts_array"][index], + original_motion.temporal_bins_s[0], + extra_info["non_rigid_window_centers"], + ) + + # Independently create the motion object and compare. This tests that cached attributes + # on the `corrected_recordings` have been properly refreshed. + expected = InterpolateMotionRecording(recording, expected_motion, border_mode="force_zeros", p=2) + np.testing.assert_allclose(corrected_recordings[index].get_traces(), expected.get_traces()) + np.testing.assert_array_equal(mc_recording_list[index].get_traces(), original_traces[index]) + + @pytest.mark.parametrize("rigid_intersession", [True, False]) + def test_nonrigid_motion(self, rigid_intersession, recording_1, recording_2): + """ + Now test that non-rigid motion estimates are properly combined with the + rigid or non-rigid inter-session alignment estimates. + """ + if rigid_intersession: + recordings_list, _, peaks_list, _ = recording_1 + else: + recordings_list, _, peaks_list, _ = recording_2 + + mc_recording_list, mc_motion_info_list = self.motion_correct_recordings_list( + recordings_list, + rigid_motion=False, + ) + + # Now the motion data has multiple displcements (per probe segment window). + # Add offsets to these different windows. + first_ses_mc_displacement = mc_recording_list[0]._recording_segments[0].motion.displacement + second_ses_mc_displacement = mc_recording_list[1]._recording_segments[0].motion.displacement + + offsets1 = np.linspace(0, 0.1, first_ses_mc_displacement[0].size) + offsets2 = np.linspace(0, 0.1, first_ses_mc_displacement[0].size) + + first_ses_mc_displacement[0] += offsets1 + second_ses_mc_displacement[0] += offsets2 + + # Now run inter-session alignment (either rigid or non-rigid). Check that + # the final displacement on the motion objects is a combination of the + # nonrigid motion estimate and rigid or nonrigid inter-session alignment estimate. + non_rigid_window_kwargs = session_alignment.get_non_rigid_window_kwargs() + non_rigid_window_kwargs["rigid"] = rigid_intersession + + corrected_recordings, extra_info = session_alignment.align_sessions_after_motion_correction( + mc_recording_list, + mc_motion_info_list, + align_sessions_kwargs={ + "alignment_order": "to_session_1", + "non_rigid_window_kwargs": non_rigid_window_kwargs, + }, + ) + + first_ses_total_displacement = corrected_recordings[0]._recording_segments[0].motion.displacement + second_ses_total_displacement = corrected_recordings[1]._recording_segments[0].motion.displacement + + assert np.all(extra_info["shifts_array"][0] + offsets1 == first_ses_total_displacement) + assert np.all(extra_info["shifts_array"][1] + offsets2 == second_ses_total_displacement) + + self.assert_interpolate_recording_not_duplicate(corrected_recordings[0]) + + def assert_interpolate_recording_not_duplicate(self, recording): + """ + Do a quick check that indeed the interpolate recording is not duplicate + (i.e. only one interpolate recording, and the previous is the generated simulation + recording. + """ + assert ( + isinstance(recording, InterpolateMotionRecording) + and recording._parent.name == "InterSessionDisplacementRecording" + ) + + def test_motion_correction_peaks_are_converted(self, mocker, recording_1): + """ + When `align_sessions_after_motion_correction` is run, the peaks locations + used should be those that are already motion corrected, which requires + correcting the peak locations in the function. + + Therefore, check that the final peak locations passed to `align_sessions` + are motion-corrected. + """ + recordings_list, _, peaks_list, peak_locations_list = recording_1 + + # Motion correct recordings, and add a known motion-displacement + mc_recording_list, mc_motion_info_list = self.motion_correct_recordings_list( + recordings_list, + rigid_motion=True, + ) + + first_ses_mc_displacement = mc_recording_list[0]._recording_segments[0].motion.displacement + second_ses_mc_displacement = mc_recording_list[1]._recording_segments[0].motion.displacement + + first_ses_mc_displacement[0] += 0.1 + second_ses_mc_displacement[0] += 0.2 + + # mock the `align_sessions` function to check what was passed + spy_align_sessions = mocker.spy( + spikeinterface.preprocessing.inter_session_alignment.session_alignment, "align_sessions" + ) + + # Call the function, and check that the passed peak-locations are corrected + corrected_recordings, _ = session_alignment.align_sessions_after_motion_correction( + mc_recording_list, mc_motion_info_list, None + ) + + passed_peak_locations_1 = spy_align_sessions.call_args_list[0][0][2][0] + passed_peak_locations_2 = spy_align_sessions.call_args_list[0][0][2][1] + + assert np.allclose( + passed_peak_locations_1["y"], mc_motion_info_list[0]["peak_locations"]["y"] - 0.1, rtol=0, atol=1e-4 + ) + assert np.allclose( + passed_peak_locations_2["y"], mc_motion_info_list[1]["peak_locations"]["y"] - 0.2, rtol=0, atol=1e-4 + ) + + def test_motion_correction_kwargs(self, mocker, recording_1): + """ + For `align_sessions_after_motion_correction`, if the motion-correct is non-rigid + then the non-rigid window kwargs must match for inter-session alignment, + otherwise it will not be possible to add the displacement. + """ + recordings_list, _, _, _ = recording_1 + + mc_recording_list, mc_motion_info_list = self.motion_correct_recordings_list( + recordings_list, + rigid_motion=False, + ) + + spy_align_sessions = mocker.spy( + spikeinterface.preprocessing.inter_session_alignment.session_alignment, "align_sessions" + ) + + # Run `align_sessions_after_motion_correction` with non-rigid window kwargs + # that do not mach those used for motion correction + changed = session_alignment.get_non_rigid_window_kwargs() + changed["rigid"] = False + changed["win_step_um"] = 51 + + session_alignment.align_sessions_after_motion_correction( + mc_recording_list, mc_motion_info_list, {"non_rigid_window_kwargs": changed} + ) + + # Now remove kwargs from the motion-correct and inter-session alignment (passed) non-rigid + # window kwargs that don't match (Some from motion are not relevant for inter-session alignment, + # some for inter-session and not set on motion (they may? predate their introduction). + # Check that these core kwargs match (i.e. align_sessions is using the non-rigid-window + # settings that motion uses. + non_rigid_windows = mc_motion_info_list[0]["parameters"]["estimate_motion_kwargs"] + non_rigid_windows.pop("method") + non_rigid_windows.pop("bin_s") + non_rigid_windows.pop( + "hist_margin_um" + ) # TODO: I think some kwargs are not exposed, this is probably okay, could mention to Sam + + passed_non_rigid_windows = spy_align_sessions.call_args_list[0][1]["non_rigid_window_kwargs"] + passed_non_rigid_windows.pop("zero_threshold") + passed_non_rigid_windows.pop("win_margin_um") + + assert sorted(passed_non_rigid_windows) == sorted(non_rigid_windows) + + ########################################################################### + # Unit Tests + ########################################################################### + + def test_estimate_chunk_size(self): + """ + Test the function that chunks the time before computing histograms. + It estimates chunk size from the firing frequency. + """ + histogram = np.array([1.0, 4.0, 2.0]) + # The representative firing rate is 25% of the maximum: 1 Hz. + # A 10% error tolerance at 90% confidence requires 270.6025 seconds. + chunk_size = alignment_utils.estimate_chunk_size(histogram) + doubled_rate_chunk_size = alignment_utils.estimate_chunk_size(2 * histogram) + + assert chunk_size == pytest.approx(270.6025) + assert doubled_rate_chunk_size == pytest.approx(chunk_size / 2) + + def test_shift_array_fill_zeros(self): + """ + The tested function shifts a 1d array or 2d array (along a certain axis) + and fills space with zero. Check that arrays are shifted as expected. + """ + # Test 1d array + test_1d = np.random.random((10)) + + # shift leftwards + res = alignment_utils.shift_array_fill_zeros(test_1d, 2) + assert np.all(res[8:] == 0) + assert np.array_equal(res[:8], test_1d[2:]) + + # shift rightwards + res = alignment_utils.shift_array_fill_zeros(test_1d, -2) + assert np.all(res[:2] == 0) + assert np.array_equal(res[2:], test_1d[:8]) + + # Test 2d array + test_2d = np.random.random((10, 10)) + + # shift upwards + res = alignment_utils.shift_array_fill_zeros(test_2d, 2) + assert np.all(res[8:, :] == 0) + assert np.array_equal(res[:8, :], test_2d[2:, :]) + + # shift downwards. + res = alignment_utils.shift_array_fill_zeros(test_2d, -2) + assert np.all(res[:2, :] == 0) + assert np.array_equal(res[2:, :], test_2d[:8, :]) + + def test_get_shifts_from_session_matrix(self): + """ + Given a 'session matrix' of shifts (a matrix Mij where each element + is the shift to get from session i to session j). It is skew-symmetric. + """ + matrix = np.random.random((10, 10, 2)) + + res = alignment_utils.get_shifts_from_session_matrix("to_middle", matrix) + assert np.array_equal(res, -np.mean(matrix, axis=0)) + + res = alignment_utils.get_shifts_from_session_matrix("to_session_1", matrix) + assert np.array_equal(res, -matrix[0, :, :]) + + res = alignment_utils.get_shifts_from_session_matrix("to_session_5", matrix) + assert np.array_equal(res, -matrix[4, :, :]) + + res = alignment_utils.get_shifts_from_session_matrix("to_session_10", matrix) + assert np.array_equal(res, -matrix[9, :, :]) + + @pytest.mark.parametrize("interpolate", [True, False]) + @pytest.mark.parametrize("odd_hist_size", [True, False]) + @pytest.mark.parametrize("shifts", [3, -2]) + def test_compute_histogram_crosscorrelation_rigid(self, interpolate, odd_hist_size, shifts): + """ + Create some toy array and shift it, then check that the cross-correlattion + correctly finds the shifts, under a number of conditions. + """ + if odd_hist_size: + hist = np.array([1, 0, 1, 1, 1, 0]) + else: + hist = np.array([0, 0, 1, 1, 0, 1, 0, 1]) + + hist_shift = alignment_utils.shift_array_fill_zeros(hist, shifts) + + session_histogram_list = np.vstack([hist, hist_shift]) + + interp_factor = 50 # not used when interpolate = False + shifts_matrix, xcorr_matrix_unsmoothed = alignment_utils.compute_histogram_crosscorrelation( + session_histogram_list, + non_rigid_windows=np.ones((1, hist.size)), + num_shifts=None, + interpolate=interpolate, + interp_factor=interp_factor, + kriging_sigma=0.2, + kriging_p=2, + kriging_d=2, + smoothing_sigma_bin=None, + smoothing_sigma_window=None, + min_crosscorr_threshold=0.001, + ) + assert np.isclose( + alignment_utils.get_shifts_from_session_matrix("to_session_1", shifts_matrix)[-1], + -shifts, + rtol=0, + atol=0.01, + ) + + num_shifts = hist.size * 2 - 1 + if interpolate: + assert xcorr_matrix_unsmoothed.shape[1] == (num_shifts - 1) * interp_factor + 1 + else: + assert xcorr_matrix_unsmoothed.shape[1] == num_shifts + + def test_compute_histogram_crosscorrelation_multiple_sessions(self): + """Check pairwise shifts and their signs across six sessions. + + Each session contains a single histogram peak at a known offset. + Cross-correlation should recover the relative displacement between + every pair of sessions. + """ + offsets = np.array([0, 1, 4, 9, 16, 25]) + histograms = np.zeros((offsets.size, 64)) + histograms[np.arange(offsets.size), 16 + offsets] = 1 + + shifts_matrix, _ = alignment_utils.compute_histogram_crosscorrelation( + histograms, + non_rigid_windows=np.ones((1, 64)), + num_shifts=30, + interpolate=False, + interp_factor=1, + kriging_sigma=0.2, + kriging_p=2, + kriging_d=2, + smoothing_sigma_bin=None, + smoothing_sigma_window=None, + min_crosscorr_threshold=0.001, + ) + + expected = offsets[:, None] - offsets[None, :] + np.testing.assert_array_equal(shifts_matrix[:, :, 0], expected) + + @pytest.mark.parametrize("histogram_mode", ["1d", "2d"]) + @pytest.mark.filterwarnings( + "error:Conversion of an array with ndim > 0 to a scalar is deprecated:DeprecationWarning" + ) + def test_compute_histogram_crosscorrelation_nonrigid(self, histogram_mode): + """ """ + # fmt: off + # Window 1 | Window 2 | Window 3 + hist_1 = np.array([0.5, 1, 0, 0, 0, 1e-3, 0, 0, 0, 0, 0, 1]) + hist_2 = np.array([0, 0, 0.5, 1, 1e-12, 0, 0, 0, 1, 0, 0, 0]) + + if histogram_mode == "2d": + hist_1 = np.vstack([hist_1, hist_1 * 2]).T + hist_2 = np.vstack([hist_2, hist_2 * 2]).T + + winds = np.array([[ 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0], + [ 0, 0, 0, 0, 1, 1, 1, 1, 0, 0, 0, 0], + [ 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1]]) + + shifts_matrix, xcorr_matrix_unsmoothed = alignment_utils.compute_histogram_crosscorrelation( + np.stack([hist_1, hist_2], axis=0), + non_rigid_windows=winds, + num_shifts=None, + interpolate=None, + interp_factor=1.0, + kriging_sigma=0.2, + kriging_p=2, + kriging_d=2, + smoothing_sigma_bin=None, + smoothing_sigma_window=None, + min_crosscorr_threshold=0.001, + ) + # fmt: on + + wind_1_shift = shifts_matrix[1, 0][0] + wind_2_shift = shifts_matrix[1, 0][1] + wind_3_shift = shifts_matrix[1, 0][2] + + assert wind_1_shift == 2 # the first window is shifted +2 + assert wind_2_shift == 0 # the second window has poor correlation under 0.001, so set to 0 + assert wind_3_shift == -3 # the third window is shifted -3 + + ########################################################################### + # Kwargs Tests + ########################################################################### + + def test_get_estimate_histogram_kwargs(self, mocker, recording_1): + + recordings_list, _, peaks_list, peak_locations_list = recording_1 + + default_kwargs = { + "bin_um": 2, + "method": "entire_session", + "chunked_bin_size_s": "estimate", + "log_transform": True, + "depth_smooth_um": None, + "histogram_type": "1d", + "weight_with_amplitude": False, + "avg_in_bin": False, + } + + assert default_kwargs == session_alignment.get_estimate_histogram_kwargs(), ( + "Default `get_estimate_histogram_kwargs` were changed." + ) + + different_kwargs = session_alignment.get_estimate_histogram_kwargs() + different_kwargs.update( + { + "method": "chunked_median", + "chunked_bin_size_s": 6, + "depth_smooth_um": 5, + "weight_with_amplitude": True, + "avg_in_bin": True, + } + ) + + spy_2d_histogram = mocker.spy(spikeinterface.sortingcomponents.motion.motion_utils, "make_2d_motion_histogram") + session_alignment.align_sessions( + [recordings_list[0]], [peaks_list[0]], [peak_locations_list[0]], estimate_histogram_kwargs=different_kwargs + ) + first_call = spy_2d_histogram.call_args_list[0] + args, kwargs = first_call + + assert kwargs["bin_s"] == different_kwargs["chunked_bin_size_s"] + assert kwargs["bin_um"] is None + assert np.unique(np.diff(kwargs["spatial_bin_edges"])) == different_kwargs["bin_um"] + assert kwargs["depth_smooth_um"] == different_kwargs["depth_smooth_um"] + assert kwargs["weight_with_amplitude"] == different_kwargs["weight_with_amplitude"] + assert kwargs["avg_in_bin"] == different_kwargs["avg_in_bin"] + + def test_compute_alignment_kwargs(self, mocker, recording_1): + + recordings_list, _, peaks_list, peak_locations_list = recording_1 + + default_kwargs = { + "num_shifts_global": None, + "num_shifts_block": 20, + "interpolate": False, + "interp_factor": 10, + "kriging_sigma": 1, + "kriging_p": 2, + "kriging_d": 2, + "smoothing_sigma_bin": 0.5, + "smoothing_sigma_window": 0.5, + "akima_interp_nonrigid": False, + "min_crosscorr_threshold": 0.001, + } + assert session_alignment.get_compute_alignment_kwargs() == default_kwargs, ( + "Default `get_compute_alignment_kwargs` were changed." + ) + + different_kwargs = session_alignment.get_compute_alignment_kwargs() + different_kwargs.update( + { + "interpolate": True, + "kriging_sigma": 5, + "kriging_p": 10, + "kriging_d": 20, + "smoothing_sigma_bin": 1.2, + "smoothing_sigma_window": 1.3, + } + ) + import scipy + + spy_kriging = mocker.spy(spikeinterface.sortingcomponents.motion.iterative_template, "kriging_kernel") + spy_gaussian_filter = mocker.spy(scipy.ndimage, "gaussian_filter") + non_rigid_window_kwargs = session_alignment.get_non_rigid_window_kwargs() + non_rigid_window_kwargs["rigid"] = False + + session_alignment.align_sessions( + [recordings_list[0]], + [peaks_list[0]], + [peak_locations_list[0]], + compute_alignment_kwargs=different_kwargs, + non_rigid_window_kwargs=non_rigid_window_kwargs, + ) + kwargs = spy_kriging.call_args_list[0][1] + assert kwargs["sigma"] == different_kwargs["kriging_sigma"] + assert kwargs["p"] == different_kwargs["kriging_p"] + assert kwargs["d"] == different_kwargs["kriging_d"] + + # First call is overall rigid, then nonrigid smooth bin + kwargs = spy_gaussian_filter.call_args_list[0][1] + assert kwargs["sigma"] == different_kwargs["smoothing_sigma_bin"] + + # then nonrigid smooth window + kwargs = spy_gaussian_filter.call_args_list[2][1] + assert kwargs["sigma"] == different_kwargs["smoothing_sigma_window"] + + def test_non_rigid_window_kwargs(self, mocker, recording_1): + + import spikeinterface # TODO: place up top with a note + + default_kwargs = { + "rigid": True, + "win_shape": "gaussian", + "win_step_um": 50, + "win_scale_um": 50, + "win_margin_um": None, + "zero_threshold": None, + } + assert session_alignment.get_non_rigid_window_kwargs() == default_kwargs, ( + "Default `get_non_rigid_window_kwargs` were changed." + ) + + different_kwargs = { + "rigid": False, + "win_shape": "rect", + "win_step_um": 55, + "win_scale_um": 65, + "win_margin_um": 10, + "zero_threshold": 4, + } + + recordings_list, _, peaks_list, peak_locations_list = recording_1 + + spy_get_spatial_windows = mocker.spy( + spikeinterface.sortingcomponents.motion.motion_utils, "get_spatial_windows" + ) + session_alignment.align_sessions( + [recordings_list[0]], [peaks_list[0]], [peak_locations_list[0]], non_rigid_window_kwargs=different_kwargs + ) + kwargs = spy_get_spatial_windows.call_args_list[0][1] + assert kwargs["rigid"] == different_kwargs["rigid"] + assert kwargs["win_shape"] == different_kwargs["win_shape"] + assert kwargs["win_step_um"] == different_kwargs["win_step_um"] + assert kwargs["win_scale_um"] == different_kwargs["win_scale_um"] + assert kwargs["win_margin_um"] == different_kwargs["win_margin_um"] + assert kwargs["zero_threshold"] == different_kwargs["zero_threshold"] + + def test_interpolate_motion_kwargs(self, mocker, recording_1): + """ """ + + default_kwargs = { + "border_mode": "force_zeros", + "spatial_interpolation_method": "kriging", + "sigma_um": 20.0, + "p": 2, + } + assert session_alignment.get_interpolate_motion_kwargs() == default_kwargs, ( + "Default `get_non_rigid_window_kwargs` were changed." + ) + + different_kwargs = { + "border_mode": "force_zeros", + "spatial_interpolation_method": "nearest", + "sigma_um": 25.0, + "p": 3, + } + + recordings_list, _, peaks_list, peak_locations_list = recording_1 + + spy_get_2d_activity_histogram = mocker.spy( + InterpolateMotionRecording, + "__init__", + ) + session_alignment.align_sessions( + [recordings_list[0]], [peaks_list[0]], [peak_locations_list[0]], interpolate_motion_kwargs=different_kwargs + ) + first_call = spy_get_2d_activity_histogram.call_args_list[0] + args, kwargs = first_call + + assert kwargs["border_mode"] == different_kwargs["border_mode"] + assert kwargs["spatial_interpolation_method"] == different_kwargs["spatial_interpolation_method"] + assert kwargs["sigma_um"] == different_kwargs["sigma_um"] + assert kwargs["p"] == different_kwargs["p"] + + @pytest.mark.parametrize("histogram_type", ["1d", "2d"]) + def test_interesting_debug_case(self, histogram_type, recording_2): + """ + This is an interseting debug case that is included in the tests to act as + both a regression test and highlight how the alignment works and can lead + to imperfect results in the test setting. + + In this case we take a non-rigid alignment, and we see that the right-edge + of the histogram is aligned well, but the middle area (which lies within + the same nonrigid window) is not well aligned. This is in spite of the + shifts being estimated correctly for that segment. + + The problem is that the nonrigid bins are interpolated to get the shifts + for each channel. In this case, the histogram peak being aligned + lies in between two nonrigid window middle points, and so is interpolated. + It is in the window with shift ~144 but next to a window with ~190 + and so ends up around ~170, resulting in the incorrect shift for this segment. + But, in the real world it is necessary to interpolate channels like this + or shifting whole windows would end up with very strange results! + """ + recording_list, _, _, _ = recording_2 + + mc_recording_list, mc_motion_info_list = self.motion_correct_recordings_list(recording_list, rigid_motion=True) + + # Run alignment + non_rigid_window_kwargs = session_alignment.get_non_rigid_window_kwargs() + non_rigid_window_kwargs["rigid"] = False + non_rigid_window_kwargs["win_shape"] = "rect" + non_rigid_window_kwargs["win_step_um"] = 250 + non_rigid_window_kwargs["win_scale_um"] = 250 + + compute_alignment_kwargs = session_alignment.get_compute_alignment_kwargs() + compute_alignment_kwargs["num_shifts_block"] = 75 + compute_alignment_kwargs["smoothing_sigma_bin"] = None + compute_alignment_kwargs["smoothing_sigma_window"] = None + + estimate_histogram_kwargs = session_alignment.get_estimate_histogram_kwargs() + estimate_histogram_kwargs["histogram_type"] = histogram_type + + corrected_recordings, extra_info = session_alignment.align_sessions_after_motion_correction( + mc_recording_list, + mc_motion_info_list, + align_sessions_kwargs={ + "alignment_order": "to_session_1", # to_center + "non_rigid_window_kwargs": non_rigid_window_kwargs, + "compute_alignment_kwargs": compute_alignment_kwargs, + "estimate_histogram_kwargs": estimate_histogram_kwargs, + }, + ) + + if DEBUG: + from spikeinterface.widgets import plot_session_alignment, plot_activity_histogram_2d + import matplotlib.pyplot as plt + + # Plot the results, as well as the shift non-rigid window centers and + # the shifts amount. You can see where the shift is reduced to align + # the peaks in the middle of the histogram, the orange peak is + # not sufficiently moved (because of the interpolation). + peaks_list = [info["peaks"] for info in mc_motion_info_list] + peak_locations_list = [info["peak_locations"] for info in mc_motion_info_list] + plot = plot_session_alignment( + mc_recording_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}, + ) + + window_edges = np.r_[0, np.cumsum(np.diff(extra_info["non_rigid_window_centers"]))] + window_edges[:-1] += np.diff(window_edges) / 2 + + y_window = extra_info["shifts_array"][1] + x_bin = extra_info["non_rigid_window_centers"] + + ax4_twin = plot.figure.axes[4].twinx() + ax5_twin = plot.figure.axes[5].twinx() + ax4_twin.scatter(x_bin, y_window, color="red", s=100, edgecolor="black", zorder=3, label="Points") + ax5_twin.scatter(x_bin, y_window, color="red", s=100, edgecolor="black", zorder=3, label="Points") + + for x_n, y_n in zip(x_bin, y_window): + ax4_twin.plot([x_n, x_n], [0, y_n], color="black", linestyle="--", linewidth=1.5, zorder=2) + ax5_twin.plot([x_n, x_n], [0, y_n], color="black", linestyle="--", linewidth=1.5, zorder=2) + + plt.suptitle("test_interesting_debug_case") + plt.show() diff --git a/src/spikeinterface/sortingcomponents/motion/decentralized.py b/src/spikeinterface/sortingcomponents/motion/decentralized.py index 509841e6f6..05e0f94da1 100644 --- a/src/spikeinterface/sortingcomponents/motion/decentralized.py +++ b/src/spikeinterface/sortingcomponents/motion/decentralized.py @@ -3,7 +3,7 @@ from tqdm.auto import tqdm, trange from spikeinterface.core.motion import Motion -from .motion_utils import get_spatial_windows, get_spatial_bin_edges, make_2d_motion_histogram +from .motion_utils import get_spatial_windows, make_2d_motion_histogram, get_spatial_bins from .dredge import normxcorr1d @@ -135,13 +135,9 @@ def run( lsqr_robust_n_iter=20, weight_with_amplitude=False, ): - - dim = ["x", "y", "z"].index(direction) - contact_depths = recording.get_channel_locations()[:, dim] - - # spatial histogram bins - spatial_bin_edges = get_spatial_bin_edges(recording, direction, hist_margin_um, bin_um) - spatial_bin_centers = 0.5 * (spatial_bin_edges[1:] + spatial_bin_edges[:-1]) + spatial_bin_centers, spatial_bin_edges, contact_depths = get_spatial_bins( + recording, direction, hist_margin_um, bin_um + ) # get spatial windows non_rigid_windows, non_rigid_window_centers = get_spatial_windows( diff --git a/src/spikeinterface/sortingcomponents/motion/iterative_template.py b/src/spikeinterface/sortingcomponents/motion/iterative_template.py index 34d1881137..0c77aa12aa 100644 --- a/src/spikeinterface/sortingcomponents/motion/iterative_template.py +++ b/src/spikeinterface/sortingcomponents/motion/iterative_template.py @@ -289,6 +289,8 @@ def iterative_template_registration( return optimal_shift_indices, target_spikecount_hist, shift_covs_block +# TODO: this is duplicate of get_kriging_kernel_distance() but that +# doesnt expose d parameter, could combine? def kriging_kernel(source_location, target_location, sigma=1, p=2, d=2): from scipy.spatial.distance import cdist diff --git a/src/spikeinterface/sortingcomponents/motion/motion_utils.py b/src/spikeinterface/sortingcomponents/motion/motion_utils.py index ab08326c56..6e52da4c32 100644 --- a/src/spikeinterface/sortingcomponents/motion/motion_utils.py +++ b/src/spikeinterface/sortingcomponents/motion/motion_utils.py @@ -168,6 +168,17 @@ def get_spatial_bin_edges(recording, direction, hist_margin_um, bin_um): return spatial_bins +def get_spatial_bins(recording, direction, hist_margin_um, bin_um): + dim = ["x", "y", "z"].index(direction) + contact_depths = recording.get_channel_locations()[:, dim] + + # spatial histogram bins + spatial_bin_edges = get_spatial_bin_edges(recording, direction, hist_margin_um, bin_um) + spatial_bin_centers = 0.5 * (spatial_bin_edges[1:] + spatial_bin_edges[:-1]) + + return spatial_bin_centers, spatial_bin_edges, contact_depths + + def make_2d_motion_histogram( recording, peaks, diff --git a/src/spikeinterface/widgets/inter_session_alignment.py b/src/spikeinterface/widgets/inter_session_alignment.py new file mode 100644 index 0000000000..adeea37089 --- /dev/null +++ b/src/spikeinterface/widgets/inter_session_alignment.py @@ -0,0 +1,425 @@ +from spikeinterface.core import BaseRecording +import numpy as np +from spikeinterface.widgets.base import BaseWidget +from spikeinterface.widgets.base import to_attr +from spikeinterface.widgets.motion import DriftRasterMapWidget + + +class SessionAlignmentWidget(BaseWidget): + """ + Widget to display the output of inter-session alignment. + In the top section, `DriftRasterMapWidget`s are used to display + the raster maps for each session, before and after alignment. + The order of all lists should correspond to the same recording. + + If histograms are provided, `ActivityHistogram1DWidget` + are used to show the activity histograms, before and after alignment. + See `align_sessions` for context. + + Corrected and uncorrected activity histograms are generated + as part of the `align_sessions` step. + + Parameters + ---------- + + recordings_list : list[BaseRecording] + List of recordings to plot. + peaks_list : list[np.ndarray] + List of detected peaks for each session. + peak_locations_list : list[np.ndarray] + List of detected peak locations for each session. + session_histogram_list : np.ndarray | None + A list of activity histograms as output from `align_sessions`. + If `None`, no histograms will be displayed. + spatial_bin_centers=None : np.ndarray | None + Spatial bin centers for the histogram (each session activity + histogram will have the same spatial bin centers). + corrected_peak_locations_list : list[np.ndarray] | None + A list of corrected peak locations. If provided, the corrected + raster plots will be displayed. + corrected_session_histogram_list : list[np.ndarray] + A list of corrected session activity histograms, as + output from `align_sessions`. + drift_raster_map_kwargs : dict | None + Kwargs to be passed to `DriftRasterMapWidget`. + session_alignment_histogram_kwargs : dict | None + Kwargs to be passed to `ActivityHistogram1DWidget`. + **backend_kwargs + """ + + def __init__( + self, + recordings_list: list[BaseRecording], + peaks_list: list[np.ndarray], + peak_locations_list: list[np.ndarray], + session_histogram_list: list[np.ndarray] | None = None, + spatial_bin_centers: np.ndarray | None = None, + corrected_peak_locations_list: list[np.ndarray] | None = None, + corrected_session_histogram_list: list[np.ndarray] = None, + drift_raster_map_kwargs: dict | None = None, + session_alignment_histogram_kwargs: dict | None = None, + **backend_kwargs, + ): + assert len(recordings_list) <= 8, ( + "At present, this widget supports plotting up to 8 sessions. " + "Please contact SpikeInterface to discuss increasing." + ) + if corrected_session_histogram_list is not None: + if not len(corrected_session_histogram_list) == len(session_histogram_list): + raise ValueError( + "`corrected_session_histogram_list` must be the same length as `session_histogram_list`. " + "Entries should correspond exactly, with the histogram in each position being the corrected" + "version of `session_histogram_list`." + ) + if corrected_peak_locations_list is not None: + if not len(corrected_peak_locations_list) == len(peak_locations_list): + raise ValueError( + "`corrected_peak_locations_list` must be the same length as `peak_locations_list`. " + "Entries should correspond exactly, with the histogram in each position being the corrected" + "version of `peak_locations_list`." + ) + if (corrected_peak_locations_list is None) != (corrected_session_histogram_list is None): + raise ValueError( + "If either `corrected_peak_locations_list` or `corrected_session_histogram_list` " + "is passed, they must both be passed." + ) + + if drift_raster_map_kwargs is None: + drift_raster_map_kwargs = {} + + if session_alignment_histogram_kwargs is None: + session_alignment_histogram_kwargs = {} + + plot_data = dict( + recordings_list=recordings_list, + peaks_list=peaks_list, + peak_locations_list=peak_locations_list, + session_histogram_list=session_histogram_list, + spatial_bin_centers=spatial_bin_centers, + corrected_peak_locations_list=corrected_peak_locations_list, + corrected_session_histogram_list=corrected_session_histogram_list, + drift_raster_map_kwargs=drift_raster_map_kwargs, + session_alignment_histogram_kwargs=session_alignment_histogram_kwargs, + ) + + BaseWidget.__init__(self, plot_data, backend="matplotlib", **backend_kwargs) + + def plot_matplotlib(self, data_plot, **backend_kwargs): + """ + Create the `SessionAlignmentWidget` for matplotlib. + """ + from spikeinterface.widgets.utils_matplotlib import make_mpl_figure + + dp = to_attr(data_plot) + + assert backend_kwargs["axes"] is None, "axes argument is not allowed in SessionAlignmentWidget" + assert backend_kwargs["ax"] is None, "ax argument is not allowed in SessionAlignmentWidget" + + self.figure, self.axes, self.ax = make_mpl_figure(**backend_kwargs) + fig = self.figure + fig.clear() + + # Find the min and max y peak position across all sessions so the + # axis can be set the same across all sessions + min_y = np.min(np.hstack([locs["y"] for locs in dp.peak_locations_list])) + max_y = np.max(np.hstack([locs["y"] for locs in dp.peak_locations_list])) + + # First, plot the peak location raster plots + + if dp.corrected_peak_locations_list is None: + # In this case, we only have uncorrected peak locations. We plot only the + # uncorrected raster maps. If there are more than 4 sessions, move + # onto the second row (usually reserved for the corrected peak raster). + + num_sessions = len(dp.peak_locations_list) + num_cols = min(4, num_sessions) + num_rows = (num_sessions + num_cols - 1) // num_cols + + gs = fig.add_gridspec(num_rows + 1, num_cols, wspace=0.3, hspace=0.5) + + for i in range(num_sessions): + row, col = divmod(i, num_cols) + ax = fig.add_subplot(gs[row, col]) + + DriftRasterMapWidget( + dp.peaks_list[i], + dp.peak_locations_list[i], + recording=dp.recordings_list[i], + ax=ax, + **dp.drift_raster_map_kwargs, + ) + ax.set_ylim((min_y, max_y)) + else: + # In this case, we have corrected and unncorrected peak locations to + # plot in the raster. Uncorrected are on the first row and corrected are + # on the second. Each session is a new column. + + # Own function, then see if can compare + num_cols = len(dp.peak_locations_list) + num_rows = 2 + + gs = fig.add_gridspec(num_rows + 1, num_cols, wspace=0.3, hspace=0.5) + + for i in range(num_cols): + ax_top = fig.add_subplot(gs[0, i]) + ax_bottom = fig.add_subplot(gs[1, i]) + + # Uncorrected session (row 1) + DriftRasterMapWidget( + dp.peaks_list[i], + dp.peak_locations_list[i], + recording=dp.recordings_list[i], + ax=ax_top, + **dp.drift_raster_map_kwargs, + ) + ax_top.set_title(f"Session {i + 1}") + ax_top.set_xlabel(None) + ax_top.set_ylim((min_y, max_y)) + + # Corrected session (row 2) + DriftRasterMapWidget( + dp.peaks_list[i], + dp.corrected_peak_locations_list[i], + recording=dp.recordings_list[i], + ax=ax_bottom, + **dp.drift_raster_map_kwargs, + ) + ax_bottom.set_title(f"Corrected Session {i + 1}") + ax_bottom.set_ylim((min_y, max_y)) + + # Next, plot the activity histograms under the raster plots + # If we only have uncorrected, plot taking up two columns. + # Otherwise, uncorrected histogram on the left column and + # corrected histgoram on the right column + if dp.session_histogram_list: + num_sessions = len(dp.session_histogram_list) + + if "legend" not in dp.session_alignment_histogram_kwargs: + sessions = [f"session {i + 1}" for i in range(num_sessions)] + dp.session_alignment_histogram_kwargs["legend"] = sessions + + if not dp.corrected_session_histogram_list: + ax = fig.add_subplot(gs[num_rows, :]) + + ActivityHistogram1DWidget( + dp.session_histogram_list, + dp.spatial_bin_centers, + ax=ax, + **dp.session_alignment_histogram_kwargs, + ) + ax.legend(loc="upper left") + else: + gs_sub = gs[num_rows, :].subgridspec(1, 2) + + ax_left = fig.add_subplot(gs_sub[0]) + ax_right = fig.add_subplot(gs_sub[1]) + + ActivityHistogram1DWidget( + dp.session_histogram_list, + dp.spatial_bin_centers, + ax=ax_left, + **dp.session_alignment_histogram_kwargs, + ) + ActivityHistogram1DWidget( + dp.corrected_session_histogram_list, + dp.spatial_bin_centers, + ax=ax_right, + **dp.session_alignment_histogram_kwargs, + ) + ax_left.get_legend().set_loc("upper right") + ax_left.set_title("Original Histogram") + ax_right.get_legend().set_loc("upper right") + ax_right.set_title("Corrected Histogram") + + +class ActivityHistogram1DWidget(BaseWidget): + """ + Plot 1D session activity histograms, overlaid on the same plot. + See SessionAlignmentWidget for detail. + + Parameters + ---------- + + session_histogram_list: list[np.ndarray] + List of 1D activity histograms to plot + spatial_bin_centers: list[np.ndarray] | np.ndarray | None + x-axis tick labels (bin centers of the histogram) + legend: None | list[str] = None + List of str to set as plot legend + linewidths: None | float | list[float] = 2, + Linewidths (list of linewidth for different across histograms, + otherwise `None` or specify shared linewidth with `float`. + colors: None | str | list[str] = None, + Colors to set the activity histograms. `None` uses matplotlib defautl colors. + """ + + def __init__( + self, + session_histogram_list: list[np.ndarray], + spatial_bin_centers: list[np.ndarray] | np.ndarray | None, + legend: None | list[str] = None, + linewidths: None | float | list[float] = 2.0, + colors: None | list = None, + **backend_kwargs, + ): + + plot_data = dict( + session_histogram_list=session_histogram_list, + spatial_bin_centers=spatial_bin_centers, + legend=legend, + linewidths=linewidths, + colors=colors, + ) + + BaseWidget.__init__(self, plot_data, backend="matplotlib", **backend_kwargs) + + def plot_matplotlib(self, data_plot, **backend_kwargs): + from spikeinterface.widgets.utils_matplotlib import make_mpl_figure + from matplotlib.colors import is_color_like + + dp = to_attr(data_plot) + + legend = dp.legend + colors = dp.colors + linewidths = dp.linewidths + spatial_bin_centers = dp.spatial_bin_centers + + assert backend_kwargs["axes"] is None, "`axes` argument not supported. Use `ax` to pass an axis to set." + + self.figure, self.axes, self.ax = make_mpl_figure(**backend_kwargs) + + num_histograms = len(dp.session_histogram_list) + + # If passed parameters are not unique across plots, then + # create as lists to set them for all plots + if colors is None or is_color_like(colors): + colors = [colors] * num_histograms + + if not isinstance(linewidths, (list, tuple)): + linewidths = [linewidths] * num_histograms + + spatial_bin_centers = [spatial_bin_centers] * num_histograms + + # If 2D, average across amplitude axis + if dp.session_histogram_list[0].ndim == 2: + histogram_list = [np.sum(hist_, axis=1) for hist_ in dp.session_histogram_list] + print( + "2D histogram passed, will be summed across first (i.e. amplitude) axis.\n" + "Use ActivityHistogram1DWidget to plot the 2D histograms directly." + ) + else: + histogram_list = dp.session_histogram_list + + # Plot the activity histograms + for i in range(num_histograms): + self.ax.plot(spatial_bin_centers[i], histogram_list[i], color=colors[i], linewidth=linewidths[i]) + + if legend is not None: + self.ax.legend(legend) + + self.ax.set_xlabel("Spatial bins (um)") + self.ax.set_ylabel("Activity (p.d.u)") + + +class ActivityHistogram2DWidget(BaseWidget): + """ + Plot 2D (spatial bin, amplitude bin) histograms following inter-session alignment. + The first column is uncorrected histograms, the second (if passed) is the corrected histogram. + + Parameters + ---------- + session_histogram_list : list[np.ndarray] + List of 2D activity histograms (one per sesson) + spatial_bin_centers : np.ndarray + Array of spatial bin centers (shared between all histograms) + corrected_session_histogram_list : None | list[np.ndarray] + A list of 2D corrected activity histograms (one per session, order + corresponding to `session_histogram_list`. + """ + + def __init__( + self, + session_histogram_list: list[np.ndarray], + spatial_bin_centers: np.ndarray, + corrected_session_histogram_list: None | list[np.ndarray] = None, + **backend_kwargs, + ): + if corrected_session_histogram_list: + if not (len(corrected_session_histogram_list) == len(session_histogram_list)): + raise ValueError( + "`corrected_session_histogram_list` must be the same" + "length as `session_histogram_list`, containing a " + "corrected histogram corresponding to each entry." + ) + + plot_data = dict( + session_histogram_list=session_histogram_list, + spatial_bin_centers=spatial_bin_centers, + corrected_session_histogram_list=corrected_session_histogram_list, + ) + + BaseWidget.__init__(self, plot_data, backend="matplotlib", **backend_kwargs) + + def plot_matplotlib(self, data_plot, **backend_kwargs): + """ + Create the `SessionAlignmentWidget` for matplotlib. + """ + from spikeinterface.widgets.utils_matplotlib import make_mpl_figure + + dp = to_attr(data_plot) + + assert backend_kwargs["axes"] is None, "axes argument is not allowed in ActivityHistogram1DWidget" + assert backend_kwargs["ax"] is None, "ax argument is not allowed in ActivityHistogram1DWidget" + + self.figure, self.axes, self.ax = make_mpl_figure(**backend_kwargs) + fig = self.figure + fig.clear() + + num_sessions = len(dp.session_histogram_list) + has_corrected = dp.corrected_session_histogram_list is not None + + num_cols = 2 if has_corrected else 1 + gs = fig.add_gridspec(num_sessions, num_cols, wspace=0.3, hspace=0.5) + + # Show 8 (arbitrary numbers) ticks on the spatial bin axis + bin_centers = dp.spatial_bin_centers + divisor = max(1, bin_centers.size // 8) + xlabels = bin_centers[::divisor] + + for idx in range(num_sessions): + # Plot uncorrected 2d histograms in the first column + ax = fig.add_subplot(gs[idx, 0]) + + num_bins = dp.session_histogram_list[idx].shape[0] + ax.imshow(dp.session_histogram_list[idx].T, aspect="auto") + + ax.set_title(f"Session {idx + 1}") + + self._set_plot_tick_labels(idx, num_sessions, ax, num_bins, xlabels, col=0) + + # If passed, plot corrected 2d histograms in the second column + if has_corrected: + ax = fig.add_subplot(gs[idx, 1]) + + ax.imshow(dp.corrected_session_histogram_list[idx].T, aspect="auto") + ax.set_title(f"Corrected Session {idx + 1}") + + self._set_plot_tick_labels(idx, num_sessions, ax, num_bins, xlabels, col=1) + + def _set_plot_tick_labels(self, idx, num_sessions, ax, num_bins, xlabels, col): + """ + Setup the plot labels. Only the bottom plots should show the x-axis (spatial) + bin ticks. On the left plots should show the y-axis (amplitude) bin label. + The amplitude bins are not specified in units (just bin number: TODO: Q: is this confusing?) + """ + if col == 0: + ax.set_ylabel("Amplitude Bin") + + if idx == num_sessions - 1: + # Set the x-ticks on the bottom plot only + ax.set_xticks(np.linspace(0, num_bins - 1, xlabels.size)) + ax.set_xticklabels([f"{i}" for i in xlabels], rotation=45) + ax.set_xlabel("Spatial Bins (µm)") + else: + ax.set_xticks([]) + ax.set_xticklabels([]) diff --git a/src/spikeinterface/widgets/motion.py b/src/spikeinterface/widgets/motion.py index 4ac8c9b030..f6f3dc17ac 100644 --- a/src/spikeinterface/widgets/motion.py +++ b/src/spikeinterface/widgets/motion.py @@ -371,7 +371,6 @@ def __init__( def plot_matplotlib(self, data_plot, **backend_kwargs): from .utils_matplotlib import make_mpl_figure - from spikeinterface.sortingcomponents.motion import correct_motion_on_peaks dp = to_attr(data_plot) diff --git a/src/spikeinterface/widgets/widget_list.py b/src/spikeinterface/widgets/widget_list.py index a0ef344d12..cb5036667d 100644 --- a/src/spikeinterface/widgets/widget_list.py +++ b/src/spikeinterface/widgets/widget_list.py @@ -8,6 +8,7 @@ from .autocorrelograms import AutoCorrelogramsWidget from .crosscorrelograms import CrossCorrelogramsWidget from .drift_templates import DriftingTemplatesWidget +from .inter_session_alignment import SessionAlignmentWidget, ActivityHistogram1DWidget, ActivityHistogram2DWidget from .isi_distribution import ISIDistributionWidget from .metrics import MetricsHistogramsWidget from .motion import DriftRasterMapWidget, MotionWidget, MotionInfoWidget @@ -42,6 +43,8 @@ from .bombcell_curation import BombcellUpsetPlotWidget widget_list = [ + ActivityHistogram1DWidget, + ActivityHistogram2DWidget, AgreementMatrixWidget, AllAmplitudesDistributionsWidget, AmplitudesWidget, @@ -66,6 +69,7 @@ ProbeMapWidget, QualityMetricsWidget, RasterWidget, + SessionAlignmentWidget, SortingSummaryWidget, SpikeLocationsWidget, SpikesOnTracesWidget, @@ -127,6 +131,8 @@ # make function for all widgets +plot_activity_histogram_1d = ActivityHistogram1DWidget +plot_activity_histogram_2d = ActivityHistogram2DWidget plot_agreement_matrix = AgreementMatrixWidget plot_all_amplitudes_distributions = AllAmplitudesDistributionsWidget plot_amplitudes = AmplitudesWidget @@ -151,6 +157,7 @@ plot_probe_map = ProbeMapWidget plot_quality_metrics = QualityMetricsWidget plot_rasters = RasterWidget +plot_session_alignment = SessionAlignmentWidget plot_sorting_summary = SortingSummaryWidget plot_spike_locations = SpikeLocationsWidget plot_spikes_on_traces = SpikesOnTracesWidget