Add session displacement generation

[JoeZiminski]

This PR adds an 'inter-session displacement' ground-truth recording generator. This is to act as test data for inter-session alignment e.g. #2626 #3126. The reason for this is to support inter-session alignment (e.g. #2626, #3126 ). The idea is to create separate recordings with the same templates but shifted unit locations across recordings. There are options to model:

• rigid shift in (x, y)
• nonrigid shift in (x, y)
• neurons dropping out across recordings, recording_amplitude_scalings allows scaling the injected template amplitudes to different sizes across recordings.
• additional neurons being introduced into the recording. By default, when the templates shift, the space previously occupied on the probe by those recordings is left empty. In reality, new neurons should be shifted into the recording. shift_units_outside_probe=True introduces new neurons into the recording after probe shift.

This PR tries to use the existing motion machinery where possible. This is quite a large PR (sorry) but I split into commits to make it easier to review commit-by-commit. The first few commits perform refactoring on existing code that are required. Then, the next few commits introduce the generation and test functions. A debugging script can be used to test manually, it can be found in the section below. Please also see some examples of how the function is used below.

I'm not sure about the name 'session displacement'. Maybe just 'generate_multi_session_recording()` is easier to understand. Also, please let me know if any other variable names are unclear.

Example Use

Example 1:

Set the amplitude of 2 units to zero between sessions, also with a 100 um shift. The units are shifted across sessions, and the top / bottom units are removed (e.g. simulating these neurons disappearing between sessions)_.

Example 2

Set shift_units_outside_probe to True which introduces new units into the recording due to shift, alongside a 250 um shift. Note the top unit has been shifted out of the probe, the middle units are shifted up, and 2 new units are introduced at the bottom of the probe.

Debugging code

import spikeinterface.full as si
from spikeinterface.generation.session_displacement_generator import generate_session_displacement_recordings
import matplotlib.pyplot as plt
import numpy as np
from spikeinterface.sortingcomponents.peak_detection import detect_peaks
from spikeinterface.sortingcomponents.peak_localization import localize_peaks

# Generate a ground truth recording where every unit is firing a lot,
# with high amplitude, and is close to the spike, so all picked up.
# This just makes it easy to play around with the units, e.g.
# if specifying 5 units, 5 unit peaks are clearly visible, none are lost
# because their position is too far from probe.

default_unit_params_range = dict(
    alpha=(100.0, 500.0),
    depolarization_ms=(0.09, 0.14),
    repolarization_ms=(0.5, 0.8),
    recovery_ms=(1.0, 1.5),
    positive_amplitude=(0.1, 0.25),
    smooth_ms=(0.03, 0.07),
    spatial_decay=(20, 40),
    propagation_speed=(250.0, 350.0),
    b=(0.1, 1),
    c=(0.1, 1),
    x_angle=(0, np.pi),
    y_angle=(0, np.pi),
    z_angle=(0, np.pi),
)

default_unit_params_range["alpha"] = (500, 500)  # do this or change the margin on generate_unit_locations_kwargs
default_unit_params_range["b"] = (0.5, 1)        # and make the units fatter, easier to receive signal!
default_unit_params_range["c"] = (0.5, 1)

scale_ = [np.array([0.25, 0.5, 1, 1, 0])] * 2
scale_ = [np.ones(5)] + scale_

rec_list, _ = generate_session_displacement_recordings(
    non_rigid_gradient=None,  # 0.05, TODO: note this will set nonlinearity to both x and y (the same)
    num_units=5,
    recording_durations=(25, 25, 25),  # TODO: checks on inputs
    recording_shifts=(
        (0, 0),
        (0, 0),
        (0, 0),
    ),
    recording_amplitude_scalings= {
        "method": "by_amplitude_and_firing_rate",
        "scalings": scale_,
    },
    shift_units_outside_probe=False,
    generate_sorting_kwargs=dict(firing_rates=(0, 200), refractory_period_ms=4.0),
    generate_templates_kwargs=dict(unit_params=default_unit_params_range, ms_before=1.5, ms_after=3),
    seed=None,
    generate_unit_locations_kwargs=dict(
        margin_um=0.0,  # if this is say 20, then units go off the edge of the probe and are such low amplitude they are not picked up.
        minimum_z=5.0,
        maximum_z=45.0,
        minimum_distance=18.0,
        max_iteration=100,
        distance_strict=False,
    ),
)

# Iterate through each recording, plotting the raw traces then
# detecting and plotting the peaks.

for rec in rec_list:

    si.plot_traces(rec, time_range=(0, 1))
    plt.show()

    peaks = detect_peaks(rec, method="locally_exclusive")
    peak_locs = localize_peaks(rec, peaks, method="grid_convolution")

    si.plot_drift_raster_map(
        peaks=peaks,
        peak_locations=peak_locs,
        recording=rec,
        clim=(-300, 0)  # fix clim for comparability across plots
    )
    plt.show()

[cwindolf]

Cool! This seems really nice. I think my only question is whether there can be more logic sharing with the single-session generate_drifting_recording, since some of the changes you've made seem like they could be helpful there. But, not sure how much sense that makes.

[cwindolf]

maybe this could be called something like "simulate_linear_gradient_drift"? i was a bit confused reading it, but it seems to be generating drift which is 0 at the top of the probe and something not zero at the bottom?

maybe someone can help explain what exactly

displacement_unit_factor = factors * (1 - f) + f

ends up producing... is it like there is some global drift plus per-unit linear drift?

[cwindolf]

this is a cool feature! in the single-session case, i have previously dealt with this by adding extra units that are off the probe to start. i'm wondering how this fits into the single session case... would this function be useful there too?

[cwindolf]

i guess it would only need max_displacement padding of units? maybe plus a spread factor of 200um?

also, would it be worth adding this to the single-session generator too?