diff --git a/doc/api.rst b/doc/api.rst index ce850d1291..02cb50749d 100755 --- a/doc/api.rst +++ b/doc/api.rst @@ -475,7 +475,7 @@ Drift .. autofunction:: generate_displacement_vector .. autofunction:: make_one_displacement_vector .. autofunction:: make_linear_displacement - .. autofunction:: move_dense_templates + .. autofunction:: move_all_dense_templates_by_displacement .. autofunction:: interpolate_templates .. autoclass:: DriftingTemplates .. autoclass:: InjectDriftingTemplatesRecording diff --git a/src/spikeinterface/core/generate.py b/src/spikeinterface/core/generate.py index 012f47fcea..657f618bf4 100644 --- a/src/spikeinterface/core/generate.py +++ b/src/spikeinterface/core/generate.py @@ -2253,8 +2253,19 @@ def generate_unit_locations( rng = np.random.default_rng(seed=seed) units_locations = np.zeros((num_units, 3), dtype="float32") - minimum_x, maximum_x = np.min(channel_locations[:, 0]) - margin_um, np.max(channel_locations[:, 0]) + margin_um - minimum_y, maximum_y = np.min(channel_locations[:, 1]) - margin_um, np.max(channel_locations[:, 1]) + margin_um + if isinstance(margin_um, (int, float)): + margin_um = (float(margin_um), float(margin_um)) + else: + assert len(margin_um) == 2, "margin_um must be a float or a tuple of two floats" + + minimum_x, maximum_x = ( + np.min(channel_locations[:, 0]) - margin_um[0], + np.max(channel_locations[:, 0]) + margin_um[0], + ) + minimum_y, maximum_y = ( + np.min(channel_locations[:, 1]) - margin_um[1], + np.max(channel_locations[:, 1]) + margin_um[1], + ) units_locations[:, 0] = rng.uniform(minimum_x, maximum_x, size=num_units) if distribution == "uniform": diff --git a/src/spikeinterface/generation/__init__.py b/src/spikeinterface/generation/__init__.py index e0fb98dbbd..d1b11f0f64 100644 --- a/src/spikeinterface/generation/__init__.py +++ b/src/spikeinterface/generation/__init__.py @@ -1,5 +1,5 @@ from .drift_tools import ( - move_dense_templates, + move_all_dense_templates_by_displacement, interpolate_templates, DriftingTemplates, InjectDriftingTemplatesRecording, diff --git a/src/spikeinterface/generation/drift_tools.py b/src/spikeinterface/generation/drift_tools.py index 59565e2303..44cd279183 100644 --- a/src/spikeinterface/generation/drift_tools.py +++ b/src/spikeinterface/generation/drift_tools.py @@ -7,6 +7,9 @@ from probeinterface import Probe +from spikeinterface.core import ms_to_samples +from spikeinterface.core.generate import generate_templates + def interpolate_templates( templates_array: np.ndarray, @@ -86,9 +89,11 @@ def interpolate_templates( return new_templates_array -def move_dense_templates(templates_array, displacements, source_probe, dest_probe=None, interpolation_method="cubic"): +def move_all_dense_templates_by_displacement( + templates_array, displacements, source_probe, dest_probe=None, interpolation_method="cubic" +): """ - Move all templates_array given some displacements using spatial interpolation (cubic or linear). + Move all templates_array at onces given some displacements using spatial interpolation (cubic or linear). Optionally, the displaced templates can be remapped to another probe with a different geometry. This function operates on dense templates only. @@ -96,6 +101,8 @@ def move_dense_templates(templates_array, displacements, source_probe, dest_prob Note: in this function no checks are done to see if templates_array can be interpolatable after displacements. To check if the given displacements are interpolatable use the higher-level function move_templates(). + See also move_templates_by_position(), more flexible. + Parameters ---------- templates_array : np.array @@ -131,6 +138,69 @@ def move_dense_templates(templates_array, displacements, source_probe, dest_prob return templates_array_moved +def move_templates_by_position( + templates_array, + source_templates_locations, + source_probe, + dest_templates_locations, + dest_probe, + displacements, + interpolation_method="cubic", +): + """ + + + Parameters + ---------- + templates_array : np.array + A numpy array with dense templates_array. + shape = (num_templates, num_samples, num_channels) + source_templates_locations : np.array + Positions of templates in the source probe coordinates. + shape : (num_templates, 2) + source_probe : Probe + The Probe object on which templates_array are defined + dest_templates_locations : np.array + Positions of templates in the dest probe coordinates. + shape : (num_templates, 2) + dest_probe : Probe + Destination Probe. Can be different geometry than the original. + displacements : np.array + Displacement vector + shape : (num_displacement, 2) + interpolation_method : "cubic" | "linear", default: "cubic" + The interpolation method. + + Returns + ------- + new_templates_array : np.array + shape = (num_displacement, num_templates, num_samples, num_channels) + """ + num_displacement = displacements.shape[0] + num_templates = templates_array.shape[0] + num_samples = templates_array.shape[1] + num_channels = dest_probe.get_contact_count() + + moved_templates_array = np.zeros( + (num_displacement, num_templates, num_samples, num_channels), dtype=templates_array.dtype + ) + for i in range(num_templates): + + shift = source_templates_locations[i : i + 1, :] - dest_templates_locations[i : i + 1, :] + src_channel_locations = source_probe.contact_positions + dest_channel_locations = dest_probe.contact_positions + shift + moved_locations = dest_channel_locations[np.newaxis, :, :] - displacements.reshape(-1, 1, 2) + + moved_templates_array[:, i : i + 1, :, :] = interpolate_templates( + templates_array[i : i + 1, :, :], + src_channel_locations, + moved_locations, + interpolation_method=interpolation_method, + ) + + return moved_templates_array + + class DriftingTemplates(Templates): """ Templates with drift. @@ -242,7 +312,7 @@ def move_one_template(self, unit_index, displacement, **interpolation_kwargs): displacements : np.array The displacement vector. shape = (1, 2) - **interpolation_kwargs : keyword arguments for `move_dense_templates` function + **interpolation_kwargs : keyword arguments for `move_all_dense_templates_by_displacement` function Returns ------- @@ -255,7 +325,7 @@ def move_one_template(self, unit_index, displacement, **interpolation_kwargs): one_template_array = self.get_one_template_dense(unit_index) one_template_array = one_template_array[np.newaxis, :, :] - template_array_moved = move_dense_templates( + template_array_moved = move_all_dense_templates_by_displacement( one_template_array, displacement, self.probe, **interpolation_kwargs ) # one motion one template keep only (num_samples, num_channels) @@ -272,16 +342,132 @@ def precompute_displacements(self, displacements, **interpolation_kwargs): displacements : np.array The displacement vector. shape = (num_displacements, 2) - **interpolation_kwargs : keyword arguments for `move_dense_templates` function + **interpolation_kwargs : keyword arguments for `move_all_dense_templates_by_displacement` function """ dense_static_templates = self.get_dense_templates() - self.templates_array_moved = move_dense_templates( + self.templates_array_moved = move_all_dense_templates_by_displacement( dense_static_templates, displacements, self.probe, **interpolation_kwargs ) self.displacements = displacements +def generate_drifting_templates_synthetic( + probe, unit_locations, displacements, sampling_frequency, generate_templates_kwargs, seed +): + """ + Generate synthetic drifting template by moving the location of the units. + This avoid interpolation and then can have units on border moving outside with correct shape. + + Parameters + ---------- + probe : Probe + The target probe on which the drifting templates will be defined. + unit_locations : np.array + The locations of the units in the probe coordinates. + displacements : np.array + The displacement vector (num_displacement, 2) + sampling_frequency : float + The sampling frequency of the templates. + generate_templates_kwargs : dict + Keyword arguments for `generate_templates` function. + seed : int + Random seed for reproducibility. + + Returns + ------- + drifting_templates : DriftingTemplates + The drifting templates object. + """ + channel_locations = probe.contact_positions + + templates_array = generate_templates( + channel_locations, unit_locations, sampling_frequency=sampling_frequency, seed=seed, **generate_templates_kwargs + ) + + num_displacement = displacements.shape[0] + templates_array_moved = np.zeros(shape=(num_displacement,) + templates_array.shape, dtype=templates_array.dtype) + for i in range(num_displacement): + unit_locations_moved = unit_locations.copy() + unit_locations_moved[:, :2] += displacements[i, :][np.newaxis, :] + templates_array_moved[i, :, :, :] = generate_templates( + channel_locations, + unit_locations_moved, + sampling_frequency=sampling_frequency, + seed=seed, + **generate_templates_kwargs, + ) + + ms_before = generate_templates_kwargs["ms_before"] + nbefore = ms_to_samples(ms_before, sampling_frequency) + static_templates = Templates( + templates_array=templates_array, + sampling_frequency=sampling_frequency, + nbefore=nbefore, + probe=probe, + is_in_uV=True, + ) + + drifting_templates = DriftingTemplates.from_static_templates(static_templates) + + drifting_templates.templates_array_moved = templates_array_moved + drifting_templates.displacements = displacements + + return drifting_templates, static_templates + + +def generate_drifting_templates_by_interpolation( + templates, probe, unit_locations, displacements, interpolation_method="cubic" +): + """ + Generate drifting templates by interpolation of the static templates. + This is useful to generate drifting templates from real data. + + Parameters + ---------- + templates : Templates + The static templates, potentially higher density than the probe. + probe : Probe + The target probe on which the drifting templates will be defined. + unit_locations : np.array + The locations of the units in the probe coordinates. + If None, the main channel of the templates is used. + displacements : np.array + The displacement vector (num_displacement, 2) + interpolation_method : str, default: "cubic" + The interpolation method. + + Returns + ------- + drifting_templates : DriftingTemplates + The drifting templates object. + """ + main_channel_indices = templates.get_main_channels("both", "extremum", outputs="index") + + # We use the channel locations of the main channels of the templates as source locations for interpolation + source_templates_locations = templates.probe.contact_positions[main_channel_indices] + + templates_array_moved = move_templates_by_position( + templates.templates_array, + source_templates_locations, + templates.probe, + unit_locations[:, :2], + probe, + displacements, + interpolation_method=interpolation_method, + ) + + drifting_templates = DriftingTemplates.from_precomputed_templates( + templates_array_moved, + displacements, + templates.sampling_frequency, + templates.nbefore, + probe, + ) + + return drifting_templates + + def make_linear_displacement(start, stop, num_step=10): """ Generates 2D linear displacements between `start` and `stop` positions (included in returned displacements). @@ -302,6 +488,8 @@ def make_linear_displacement(start, stop, num_step=10): """ if num_step < 1: raise ValueError("make_linear_displacement needs num_step > 0") + start = np.array(start, dtype="float32") + stop = np.array(stop, dtype="float32") if num_step == 1: displacements = ((start + stop) / 2)[np.newaxis, :] else: diff --git a/src/spikeinterface/generation/drifting_generator.py b/src/spikeinterface/generation/drifting_generator.py index 21d25cdd41..77c20d1882 100644 --- a/src/spikeinterface/generation/drifting_generator.py +++ b/src/spikeinterface/generation/drifting_generator.py @@ -22,7 +22,12 @@ _ensure_unit_params, _ensure_seed, ) -from .drift_tools import DriftingTemplates, make_linear_displacement, InjectDriftingTemplatesRecording +from .drift_tools import ( + DriftingTemplates, + generate_drifting_templates_synthetic, + make_linear_displacement, + InjectDriftingTemplatesRecording, +) from .noise_tools import generate_noise @@ -470,11 +475,6 @@ def generate_drifting_recording( probe.set_device_channel_indices(np.arange(num_channels)) channel_locations = probe.contact_positions - # import matplotlib.pyplot as plt - # import probeinterface.plotting - # fig, ax = plt.subplots() - # probeinterface.plotting.plot_probe(probe, ax=ax) - # plt.show() # unit locations if unit_locations is None: @@ -512,47 +512,12 @@ def generate_drifting_recording( generate_templates_kwargs["unit_params"] = unit_params # generate templates - templates_array = generate_templates( - channel_locations, unit_locations, sampling_frequency=sampling_frequency, seed=seed, **generate_templates_kwargs - ) - - num_displacement = displacements_steps.shape[0] - templates_array_moved = np.zeros(shape=(num_displacement,) + templates_array.shape, dtype=templates_array.dtype) - for i in range(num_displacement): - unit_locations_moved = unit_locations.copy() - unit_locations_moved[:, :2] += displacements_steps[i, :][np.newaxis, :] - templates_array_moved[i, :, :, :] = generate_templates( - channel_locations, - unit_locations_moved, - sampling_frequency=sampling_frequency, - seed=seed, - **generate_templates_kwargs, - ) - - ms_before = generate_templates_kwargs["ms_before"] - nbefore = ms_to_samples(ms_before, sampling_frequency) - templates = Templates( - templates_array=templates_array, - sampling_frequency=sampling_frequency, - nbefore=nbefore, - probe=probe, - is_in_uV=True, + drifting_templates, static_templates = generate_drifting_templates_synthetic( + probe, unit_locations, displacements_steps, sampling_frequency, generate_templates_kwargs, seed ) - drifting_templates = DriftingTemplates.from_static_templates(templates) - - sorting.set_property("gt_unit_locations", unit_locations) - distances = np.linalg.norm(unit_locations[:, np.newaxis, :2] - channel_locations[np.newaxis, :, :], axis=2) max_channel_index = np.argmin(distances, axis=1) - sorting.set_property("max_channel_index", max_channel_index) - - ## Important precompute displacement do not work on border and so do not work for tetrode - # here we bypass the interpolation and regenerate templates at several positions. - ## drifting_templates.precompute_displacements(displacements_steps) - # shape (num_displacement, num_templates, num_samples, num_channels) - drifting_templates.templates_array_moved = templates_array_moved - drifting_templates.displacements = displacements_steps if noise is None: noise = generate_noise( @@ -567,6 +532,10 @@ def generate_drifting_recording( assert noise.probe.get_contact_count() == probe.get_contact_count(), "Noise num channels mismatch" assert noise.get_total_duration() == duration, "Noise duration should be the same as the recording duration" + sorting.set_property("gt_unit_locations", unit_locations) + sorting.set_property("max_channel_index", max_channel_index) + sorting.set_property("main_channel_id", noise.channel_ids[max_channel_index]) + amplitude_factor = synthesize_amplitude_factor( num_spikes=sorting.count_total_num_spikes(), amplitude_factor=amplitude_factor, @@ -603,7 +572,8 @@ def generate_drifting_recording( unit_locations=unit_locations, displacement_unit_factor=displacement_unit_factor, unit_displacements=unit_displacements, - templates=templates, + templates=static_templates, + drifting_templates=drifting_templates, generate_templates_kwargs=generate_templates_kwargs, ) return static_recording, drifting_recording, sorting, extra_infos diff --git a/src/spikeinterface/generation/hybrid_tools.py b/src/spikeinterface/generation/hybrid_tools.py index 84e6187ac7..c1986e1100 100644 --- a/src/spikeinterface/generation/hybrid_tools.py +++ b/src/spikeinterface/generation/hybrid_tools.py @@ -5,10 +5,8 @@ from spikeinterface.core import BaseRecording, BaseSorting, Templates, ms_to_samples from spikeinterface.core.generate import ( - generate_templates, generate_unit_locations, generate_sorting, - InjectTemplatesRecording, _ensure_seed, synthesize_amplitude_factor, ) @@ -16,10 +14,10 @@ from spikeinterface.core.motion import Motion from spikeinterface.generation.drift_tools import ( InjectDriftingTemplatesRecording, - DriftingTemplates, + generate_drifting_templates_by_interpolation, + generate_drifting_templates_synthetic, make_linear_displacement, - interpolate_templates, - move_dense_templates, + move_all_dense_templates_by_displacement, ) @@ -295,7 +293,7 @@ def relocate_templates( displacement = -displacement displacement_vector = np.zeros(2) displacement_vector[depth_dimension] = displacement - templates_array_moved[i] = move_dense_templates( + templates_array_moved[i] = move_all_dense_templates_by_displacement( templates.templates_array[i][None], displacements=displacement_vector[None], source_probe=templates.probe, @@ -316,12 +314,10 @@ def generate_hybrid_recording( recording: BaseRecording, sorting: BaseSorting | None = None, templates: Templates | None = None, + relocate_templates: bool = True, motion: Motion | None = None, - are_templates_scaled: bool = True, unit_locations: np.ndarray | None = None, drift_step_um: float = 1.0, - upsample_factor: int | None = None, - upsample_vector: np.ndarray | None = None, amplitude_std: float = 0.05, amplitude_factor: np.ndarray | None = None, generate_sorting_kwargs: dict = dict(num_units=10, firing_rates=15, refractory_period_ms=4.0, seed=2205), @@ -352,18 +348,13 @@ def generate_hybrid_recording( If None they are generated. motion : Motion | None, default: None The motion object to use for the drifting templates. - are_templates_scaled : bool, default: True - If True, the templates are assumed to be in uV, otherwise in the same unit as the recording. - In case the recording has scaling, the templates are "unscaled" before injection. unit_locations : np.array, default: None - The locations at which the templates should be injected. If not provided, generated (see - generate_unit_location_kwargs). + The locations at which the templates should be injected. If not provided, they are randomly generated, + unless relocate_templates is False. (see `generate_unit_locations_kwargs`). + relocate_templates : bool, default: True + If True, the templates are relocated across the target probe of the recording. drift_step_um : float, default: 1.0 The step in um to use for the drifting templates. - upsample_factor : None or int, default: None - A upsampling factor used only when templates are not provided. - upsample_vector : np.array or None - Optional the upsample_vector can given. This has the same shape as spike_vector amplitude_std : float, default: 0.05 The standard deviation of the modulation to apply to the spikes when injecting them into the recording. @@ -380,7 +371,7 @@ def generate_hybrid_recording( Returns ------- - recording: BaseRecording + recording: InjectDriftingTemplatesRecording The generated hybrid recording extractor. sorting: Sorting The generated sorting extractor for the injected units. @@ -393,7 +384,7 @@ def generate_hybrid_recording( sampling_frequency = recording.sampling_frequency probe = recording.get_probe() num_segments = recording.get_num_segments() - dtype = recording.dtype + # dtype = recording.dtype num_samples = np.array([recording.get_num_samples(segment_index) for segment_index in range(num_segments)]) # since the recording can have timestamps with some small gaps, we use the number of samples to compute # the duration used for the sorting generation @@ -421,57 +412,83 @@ def generate_hybrid_recording( else: generate_sorting_kwargs["num_units"] = num_units - if templates is None: - if unit_locations is None: - unit_locations = generate_unit_locations(num_units, channel_locations, **generate_unit_locations_kwargs) - else: - assert len(unit_locations) == num_units, "unit_locations and num_units should have the same length" - templates_array = generate_templates( - channel_locations, - unit_locations, - sampling_frequency, - upsample_factor=upsample_factor, - seed=seed, - dtype=dtype, - **generate_templates_kwargs, + # handle units location + if unit_locations is not None: + assert len(unit_locations) == num_units, "unit_locations and num_units should have the same length" + elif relocate_templates: + # propagate the seed so that the hybrid recording is reproducible (unless a seed is + # explicitly provided in generate_unit_locations_kwargs) + generate_unit_locations_kwargs = {"seed": seed, **generate_unit_locations_kwargs} + unit_locations = generate_unit_locations(num_units, channel_locations, **generate_unit_locations_kwargs) + else: + # We use the channel locations of the main channels of the templates as source locations for interpolation + main_channel_indices = templates.get_main_channels("both", "extremum", outputs="index") + unit_locations = templates.probe.contact_positions[main_channel_indices] + + # handle motion and displacement + if motion is None: + # make a one bin motion with no displacement + displacements = np.zeros((1, 2)) + motion = Motion( + displacement=[np.zeros((2, 1)) for _ in range(num_segments)], + temporal_bins_s=[np.array([0, durations[0]]) for _ in range(num_segments)], + spatial_bins_um=np.array([0]), + direction="y", ) - ms_before = generate_templates_kwargs["ms_before"] - ms_after = generate_templates_kwargs["ms_after"] - nbefore = ms_to_samples(ms_before, sampling_frequency) - nafter = ms_to_samples(ms_after, sampling_frequency) - templates_ = Templates(templates_array, sampling_frequency, nbefore, True, None, None, None, probe) else: - from spikeinterface.postprocessing.localization_tools import compute_monopolar_triangulation + assert num_segments == motion.num_segments, "recording and motion should have the same number of segments" + dim = motion.dim + motion_array_concat = np.concatenate(motion.displacement) + if dim == 0: + start = np.array([np.min(motion_array_concat), 0]) + stop = np.array([np.max(motion_array_concat), 0]) + elif dim == 1: + start = np.array([0, np.min(motion_array_concat)]) + stop = np.array([0, np.max(motion_array_concat)]) + elif dim == 2: + raise NotImplementedError("3D motion not implemented yet") + num_step = int((stop - start)[dim] / drift_step_um) + num_step = max(1, num_step) + displacements = make_linear_displacement(start, stop, num_step=num_step) - assert isinstance(templates, Templates), "templates should be a Templates object" - assert ( - templates.num_channels == recording.get_num_channels() - ), "templates and recording should have the same number of channels" - nbefore = templates.nbefore - nafter = templates.nafter - unit_locations = compute_monopolar_triangulation(templates, peak_sign="both", peak_mode="extremum") - - channel_locations_rel = channel_locations - channel_locations[0] - templates_locations = templates.get_channel_locations() - templates_locations_rel = templates_locations - templates_locations[0] - - if not np.allclose(channel_locations_rel, templates_locations_rel): - warnings.warn("Channel locations are different between recording and templates. Interpolating templates.") - templates_array = np.zeros(templates.templates_array.shape, dtype=dtype) - for i in range(len(templates_array)): - src_template = templates.templates_array[i][np.newaxis, :, :] - templates_array[i] = interpolate_templates(src_template, templates_locations_rel, channel_locations_rel) - else: + # calculate displacement vectors for each segment and unit + # for each unit, we interpolate the motion at its location + displacement_sampling_frequency = 1.0 / np.diff(motion.temporal_bins_s[0])[0] + displacement_vectors = [] + for segment_index in range(motion.num_segments): + temporal_bins_segment = motion.temporal_bins_s[segment_index] + displacement_vector = np.zeros((len(temporal_bins_segment), 2, num_units)) + for unit_index in range(num_units): + motion_for_unit = motion.get_displacement_at_time_and_depth( + times_s=temporal_bins_segment, + locations_um=unit_locations[unit_index], + segment_index=segment_index, + grid=True, + ) + displacement_vector[:, motion.dim, unit_index] = motion_for_unit[motion.dim, :] + displacement_vectors.append(displacement_vector) + # since displacement is estimated by interpolation for each unit, the unit factor is an eye + displacement_unit_factor = np.eye(num_units) + + # generate synthetic drifting templates or interpolate from the input templates + if templates is None: + drifting_templates, _ = generate_drifting_templates_synthetic( + probe, unit_locations, displacements, sampling_frequency, generate_templates_kwargs, seed + ) + else: + if recording.has_scaleable_traces() and templates.is_in_uV: + # In this case we "unscale" the templates, so we don't need to scale the recording to uV during injection. templates_array = templates.templates_array + templates_array = (templates_array - recording.get_channel_offsets()) / recording.get_channel_gains() + # make a copy of the templates and reset templates_array (might have scaled templates) + templates_ = templates.select_units(templates.unit_ids) + templates_.templates_array = templates_array + else: + templates_ = templates - # manage scaling of templates - templates_ = templates - if recording.has_scaleable_traces(): - if are_templates_scaled: - templates_array = (templates_array - recording.get_channel_offsets()) / recording.get_channel_gains() - # make a copy of the templates and reset templates_array (might have scaled templates) - templates_ = templates.select_units(templates.unit_ids) - templates_.templates_array = templates_array + drifting_templates = generate_drifting_templates_by_interpolation( + templates_, probe, unit_locations, displacements, interpolation_method="cubic" + ) if sorting is None: generate_sorting_kwargs = generate_sorting_kwargs.copy() @@ -483,18 +500,11 @@ def generate_hybrid_recording( assert sorting.sampling_frequency == sampling_frequency num_spikes = sorting.to_spike_vector().size + distances = np.linalg.norm(unit_locations[:, np.newaxis, :2] - channel_locations[np.newaxis, :, :], axis=2) + main_channel_indices = np.argmin(distances, axis=1) + main_channel_ids = recording.channel_ids[main_channel_indices] sorting.set_property("gt_unit_locations", unit_locations) - - assert (nbefore + nafter) == templates_array.shape[ - 1 - ], "templates and ms_before, ms_after should have the same length" - - if templates_array.ndim == 3: - upsample_vector = None - else: - if upsample_vector is None: - upsample_factor = templates_array.shape[3] - upsample_vector = rng.integers(0, upsample_factor, size=num_spikes) + sorting.set_property("main_channel_id", main_channel_ids) amplitude_factor = synthesize_amplitude_factor( num_spikes, @@ -503,68 +513,15 @@ def generate_hybrid_recording( seed=rng, ) - if motion is not None: - assert num_segments == motion.num_segments, "recording and motion should have the same number of segments" - dim = motion.dim - motion_array_concat = np.concatenate(motion.displacement) - if dim == 0: - start = np.array([np.min(motion_array_concat), 0]) - stop = np.array([np.max(motion_array_concat), 0]) - elif dim == 1: - start = np.array([0, np.min(motion_array_concat)]) - stop = np.array([0, np.max(motion_array_concat)]) - elif dim == 2: - raise NotImplementedError("3D motion not implemented yet") - num_step = int((stop - start)[dim] / drift_step_um) - num_step = max(1, num_step) - displacements = make_linear_displacement(start, stop, num_step=num_step) - - # use templates_, because templates_array might have been scaled - drifting_templates = DriftingTemplates.from_static_templates(templates_) - drifting_templates.precompute_displacements(displacements) - - # calculate displacement vectors for each segment and unit - # for each unit, we interpolate the motion at its location - displacement_sampling_frequency = 1.0 / np.diff(motion.temporal_bins_s[0])[0] - displacement_vectors = [] - for segment_index in range(motion.num_segments): - temporal_bins_segment = motion.temporal_bins_s[segment_index] - displacement_vector = np.zeros((len(temporal_bins_segment), 2, num_units)) - for unit_index in range(num_units): - motion_for_unit = motion.get_displacement_at_time_and_depth( - times_s=temporal_bins_segment, - locations_um=unit_locations[unit_index], - segment_index=segment_index, - grid=True, - ) - displacement_vector[:, motion.dim, unit_index] = motion_for_unit[motion.dim, :] - displacement_vectors.append(displacement_vector) - # since displacement is estimated by interpolation for each unit, the unit factor is an eye - displacement_unit_factor = np.eye(num_units) - - hybrid_recording = InjectDriftingTemplatesRecording( - sorting=sorting, - parent_recording=recording, - drifting_templates=drifting_templates, - displacement_vectors=displacement_vectors, - displacement_sampling_frequency=displacement_sampling_frequency, - displacement_unit_factor=displacement_unit_factor, - num_samples=num_samples.astype("int64"), - amplitude_factor=amplitude_factor, - ) - - else: - warnings.warn( - "No Motion is provided! Please check that your recording is drift-free, otherwise the hybrid recording " - "will have stationary units over a drifting recording..." - ) - hybrid_recording = InjectTemplatesRecording( - sorting, - templates_array, - nbefore=nbefore, - parent_recording=recording, - upsample_vector=upsample_vector, - amplitude_factor=amplitude_factor, - ) + hybrid_recording = InjectDriftingTemplatesRecording( + sorting=sorting, + parent_recording=recording, + drifting_templates=drifting_templates, + displacement_vectors=displacement_vectors, + displacement_sampling_frequency=displacement_sampling_frequency, + displacement_unit_factor=displacement_unit_factor, + num_samples=num_samples.astype("int64"), + amplitude_factor=amplitude_factor, + ) return hybrid_recording, sorting diff --git a/src/spikeinterface/generation/tests/test_drift_tools.py b/src/spikeinterface/generation/tests/test_drift_tools.py index 4dfcc785dc..d0d6d8114c 100644 --- a/src/spikeinterface/generation/tests/test_drift_tools.py +++ b/src/spikeinterface/generation/tests/test_drift_tools.py @@ -3,7 +3,7 @@ from spikeinterface.generation import ( interpolate_templates, - move_dense_templates, + move_all_dense_templates_by_displacement, make_linear_displacement, DriftingTemplates, InjectDriftingTemplatesRecording, @@ -84,7 +84,7 @@ def test_interpolate_templates(): ) -def test_move_dense_templates(): +def test_move_all_dense_templates_by_displacement(): templates = make_some_templates() num_move = 5 @@ -92,7 +92,9 @@ def test_move_dense_templates(): displacements = np.zeros((num_move, 2)) displacements[:, 1] = np.linspace(-amplitude_motion_um, amplitude_motion_um, num_move) - templates_moved = move_dense_templates(templates.templates_array, displacements, templates.probe) + templates_moved = move_all_dense_templates_by_displacement( + templates.templates_array, displacements, templates.probe + ) assert templates_moved.shape == (num_move,) + templates.templates_array.shape @@ -215,6 +217,6 @@ def test_InjectDriftingTemplatesRecording(create_cache_folder): if __name__ == "__main__": test_interpolate_templates() - test_move_dense_templates() + test_move_all_dense_templates_by_displacement() test_DriftingTemplates() test_InjectDriftingTemplatesRecording() diff --git a/src/spikeinterface/generation/tests/test_hybrid_tools.py b/src/spikeinterface/generation/tests/test_hybrid_tools.py index 239890e938..48e53e707d 100644 --- a/src/spikeinterface/generation/tests/test_hybrid_tools.py +++ b/src/spikeinterface/generation/tests/test_hybrid_tools.py @@ -13,9 +13,27 @@ from spikeinterface.generation.hybrid_tools import ( estimate_templates_from_recording, generate_hybrid_recording, + select_templates, + scale_template_to_range, + relocate_templates, ) +def _make_templates(num_units=10, num_channels=16, ms_before=1.0, ms_after=3.0, seed=0): + """Helper to build a Templates object (with probe) from a generated recording.""" + rec, _ = generate_ground_truth_recording( + sampling_frequency=20000, durations=[5], num_channels=num_channels, seed=seed + ) + channel_locations = rec.get_channel_locations() + unit_locations = generate_unit_locations(num_units, channel_locations=channel_locations, seed=seed) + templates_array = generate_templates( + channel_locations, unit_locations, rec.sampling_frequency, ms_before, ms_after, seed=seed + ) + nbefore = ms_to_samples(ms_before, rec.sampling_frequency) + templates = Templates(templates_array, rec.sampling_frequency, nbefore, True, None, None, None, rec.get_probe()) + return templates + + def test_generate_hybrid_no_motion(): rec, _ = generate_ground_truth_recording(sampling_frequency=20000, seed=0) hybrid, _ = generate_hybrid_recording(rec, seed=0) @@ -33,7 +51,7 @@ def test_generate_hybrid_with_sorting(): assert rec.get_num_frames() == hybrid.get_num_frames() assert rec.get_num_segments() == hybrid.get_num_segments() assert np.array_equal(rec.get_channel_locations(), hybrid.get_channel_locations()) - assert sorting_hybrid.get_num_units() == len(hybrid.templates) + assert sorting_hybrid.get_num_units() == hybrid.drifting_templates.num_units def test_generate_hybrid_motion(): @@ -65,8 +83,10 @@ def test_generate_hybrid_from_templates(): ) nbefore = ms_to_samples(ms_before, rec.sampling_frequency) templates = Templates(templates_array, rec.sampling_frequency, nbefore, True, None, None, None, rec.get_probe()) - hybrid, sorting_hybrid = generate_hybrid_recording(rec, templates=templates, seed=0) - assert np.array_equal(hybrid.templates, templates.templates_array) + hybrid, sorting_hybrid = generate_hybrid_recording(rec, templates=templates, relocate_templates=False, seed=0) + # use allclose (not array_equal): interpolation onto the target probe introduces platform-dependent + # float noise (e.g. macOS Accelerate BLAS), even though the templates are injected at their own locations + assert np.allclose(hybrid.drifting_templates.templates_array, templates.templates_array) assert rec.get_num_channels() == hybrid.get_num_channels() assert rec.get_num_frames() == hybrid.get_num_frames() assert rec.get_num_segments() == hybrid.get_num_segments() @@ -74,6 +94,252 @@ def test_generate_hybrid_from_templates(): assert sorting_hybrid.get_num_units() == num_units +# --------------------------------------------------------------------------- +# select_templates +# --------------------------------------------------------------------------- +def test_select_templates_by_amplitude(): + templates = _make_templates(num_units=10) + selected = select_templates(templates, min_amplitude=0.0) + assert selected.num_units <= templates.num_units + assert set(selected.unit_ids).issubset(set(templates.unit_ids)) + + +def test_select_templates_by_amplitude_min_max(): + templates = _make_templates(num_units=10) + selected = select_templates(templates, min_amplitude=0.0, max_amplitude=np.inf) + assert selected.num_units == templates.num_units + + +def test_select_templates_by_depth(): + templates = _make_templates(num_units=10) + channel_depths = templates.get_channel_locations()[:, 1] + selected = select_templates(templates, min_depth=np.min(channel_depths), max_depth=np.max(channel_depths)) + assert selected.num_units <= templates.num_units + assert set(selected.unit_ids).issubset(set(templates.unit_ids)) + + +@pytest.mark.parametrize("amplitude_function", ["ptp", "min", "max"]) +def test_select_templates_amplitude_functions(amplitude_function): + templates = _make_templates(num_units=10) + selected = select_templates(templates, min_amplitude=-np.inf, amplitude_function=amplitude_function) + assert selected.num_units == templates.num_units + + +@pytest.mark.parametrize("depth_direction", ["x", "y"]) +def test_select_templates_depth_direction(depth_direction): + templates = _make_templates(num_units=10) + dim = ["x", "y"].index(depth_direction) + depths = templates.get_channel_locations()[:, dim] + selected = select_templates( + templates, min_depth=np.min(depths), max_depth=np.max(depths), depth_direction=depth_direction + ) + assert selected.num_units <= templates.num_units + + +def test_select_templates_no_filter_raises(): + templates = _make_templates(num_units=5) + with pytest.raises(AssertionError): + select_templates(templates) + + +def test_select_templates_empty_returns_none(): + templates = _make_templates(num_units=5) + with pytest.warns(UserWarning): + selected = select_templates(templates, min_amplitude=np.inf) + assert selected is None + + +def test_select_templates_depth_without_probe_raises(): + templates = _make_templates(num_units=5) + templates_no_probe = Templates( + templates_array=templates.templates_array, + sampling_frequency=templates.sampling_frequency, + nbefore=templates.nbefore, + is_in_uV=templates.is_in_uV, + ) + with pytest.raises(AssertionError): + select_templates(templates_no_probe, min_depth=0.0) + + +# --------------------------------------------------------------------------- +# scale_template_to_range +# --------------------------------------------------------------------------- +@pytest.mark.parametrize("amplitude_function", ["ptp", "min", "max"]) +def test_scale_template_to_range(amplitude_function): + templates = _make_templates(num_units=10) + scaled = scale_template_to_range( + templates, min_amplitude=50.0, max_amplitude=200.0, amplitude_function=amplitude_function + ) + assert scaled.num_units == templates.num_units + assert scaled.templates_array.shape == templates.templates_array.shape + # metadata preserved + assert scaled.sampling_frequency == templates.sampling_frequency + assert scaled.nbefore == templates.nbefore + assert np.array_equal(scaled.unit_ids, templates.unit_ids) + + +def test_scale_template_to_range_correctness(): + # the smallest and largest templates should map exactly onto the requested amplitude range + templates = _make_templates(num_units=10) + min_amplitude, max_amplitude = 50.0, 200.0 + scaled = scale_template_to_range(templates, min_amplitude=min_amplitude, max_amplitude=max_amplitude) + + main_channel_indices = scaled.get_main_channels(outputs="index", with_dict=False) + scaled_array = scaled.templates_array + amplitudes = np.array([np.ptp(scaled_array[i, :, main_channel_indices[i]]) for i in range(scaled.num_units)]) + assert np.isclose(np.min(amplitudes), min_amplitude) + assert np.isclose(np.max(amplitudes), max_amplitude) + + +# --------------------------------------------------------------------------- +# relocate_templates +# --------------------------------------------------------------------------- +def test_relocate_templates(): + templates = _make_templates(num_units=10) + relocated = relocate_templates(templates, min_displacement=10.0, max_displacement=30.0, seed=0) + assert relocated.num_units == templates.num_units + assert relocated.templates_array.shape == templates.templates_array.shape + assert np.array_equal(relocated.unit_ids, templates.unit_ids) + + +@pytest.mark.parametrize("favor_borders", [True, False]) +def test_relocate_templates_favor_borders(favor_borders): + templates = _make_templates(num_units=10) + relocated = relocate_templates( + templates, min_displacement=10.0, max_displacement=30.0, favor_borders=favor_borders, seed=0 + ) + assert relocated.num_units == templates.num_units + + +@pytest.mark.parametrize("depth_direction", ["x", "y"]) +def test_relocate_templates_depth_direction(depth_direction): + templates = _make_templates(num_units=10) + relocated = relocate_templates( + templates, min_displacement=10.0, max_displacement=30.0, depth_direction=depth_direction, seed=0 + ) + assert relocated.num_units == templates.num_units + + +def test_relocate_templates_negative_margin_raises(): + templates = _make_templates(num_units=5) + with pytest.raises(AssertionError): + relocate_templates(templates, min_displacement=10.0, max_displacement=30.0, margin=-1.0, seed=0) + + +def test_relocate_templates_reproducible(): + templates = _make_templates(num_units=10) + relocated_a = relocate_templates(templates, min_displacement=10.0, max_displacement=30.0, seed=42) + relocated_b = relocate_templates(templates, min_displacement=10.0, max_displacement=30.0, seed=42) + assert np.array_equal(relocated_a.templates_array, relocated_b.templates_array) + + +def test_relocate_templates_correctness(): + # place units near the middle of the probe so a moderate displacement stays on the probe + # (no border clipping), then check the templates are physically moved along the depth axis + # by an amount within [min_displacement, max_displacement]. + rec, _ = generate_ground_truth_recording(sampling_frequency=20000, durations=[5], num_channels=16, seed=0) + channel_locations = rec.get_channel_locations() + unit_locations = np.array( + [[0.0, 60.0, 10.0], [20.0, 80.0, 10.0], [0.0, 70.0, 10.0], [20.0, 50.0, 10.0], [0.0, 90.0, 10.0]] + ) + num_units = len(unit_locations) + templates_array = generate_templates(channel_locations, unit_locations, rec.sampling_frequency, 1.0, 3.0, seed=0) + nbefore = ms_to_samples(1.0, rec.sampling_frequency) + templates = Templates(templates_array, rec.sampling_frequency, nbefore, True, None, None, None, rec.get_probe()) + + min_displacement, max_displacement = 10.0, 20.0 + relocated = relocate_templates( + templates, + min_displacement=min_displacement, + max_displacement=max_displacement, + favor_borders=False, + margin=0.0, + seed=0, + ) + + # energy-weighted center of mass of each template, per axis + locs = templates.get_channel_locations() + + def center_of_mass(tmp): + weights = np.ptp(tmp.templates_array, axis=1) # (num_units, num_channels) + com_x = (weights * locs[:, 0]).sum(axis=1) / weights.sum(axis=1) + com_y = (weights * locs[:, 1]).sum(axis=1) / weights.sum(axis=1) + return com_x, com_y + + com_x_before, com_y_before = center_of_mass(templates) + com_x_after, com_y_after = center_of_mass(relocated) + + depth_shift = np.abs(com_y_after - com_y_before) + lateral_shift = np.abs(com_x_after - com_x_before) + + # allow a small tolerance for interpolation smearing onto discrete channels + tol = 2.0 + assert np.all(depth_shift >= min_displacement - tol) + assert np.all(depth_shift <= max_displacement + tol) + # the move is along the depth ("y") direction only + assert np.all(lateral_shift < 1.0) + + +# --------------------------------------------------------------------------- +# generate_hybrid_recording : additional branches +# --------------------------------------------------------------------------- +def test_generate_hybrid_relocate_templates_true(): + templates = _make_templates(num_units=10) + rec, _ = generate_ground_truth_recording(sampling_frequency=20000, durations=[10], num_channels=16, seed=0) + hybrid, sorting_hybrid = generate_hybrid_recording(rec, templates=templates, relocate_templates=True, seed=0) + assert rec.get_num_channels() == hybrid.get_num_channels() + assert rec.get_num_frames() == hybrid.get_num_frames() + assert sorting_hybrid.get_num_units() == templates.num_units + + +def test_generate_hybrid_with_unit_locations(): + num_units = 10 + templates = _make_templates(num_units=num_units) + rec, _ = generate_ground_truth_recording(sampling_frequency=20000, durations=[10], num_channels=16, seed=0) + unit_locations = generate_unit_locations(num_units, channel_locations=rec.get_channel_locations(), seed=1) + hybrid, sorting_hybrid = generate_hybrid_recording(rec, templates=templates, unit_locations=unit_locations, seed=0) + assert sorting_hybrid.get_num_units() == num_units + assert np.array_equal(sorting_hybrid.get_property("gt_unit_locations"), unit_locations) + + +def test_generate_hybrid_unit_locations_wrong_length_raises(): + templates = _make_templates(num_units=10) + rec, _ = generate_ground_truth_recording(sampling_frequency=20000, durations=[10], num_channels=16, seed=0) + unit_locations = generate_unit_locations(5, channel_locations=rec.get_channel_locations(), seed=1) + with pytest.raises(AssertionError): + generate_hybrid_recording(rec, templates=templates, unit_locations=unit_locations, seed=0) + + +def test_generate_hybrid_with_amplitude_factor(): + gt_sorting = generate_sorting(durations=[10], num_units=10, sampling_frequency=20000, seed=0) + rec, _ = generate_ground_truth_recording(durations=[10], sampling_frequency=20000, sorting=gt_sorting, seed=0) + num_spikes = gt_sorting.to_spike_vector().size + amplitude_factor = np.ones(num_spikes) + hybrid, sorting_hybrid = generate_hybrid_recording( + rec, sorting=gt_sorting, amplitude_factor=amplitude_factor, seed=0 + ) + assert sorting_hybrid.get_num_units() == gt_sorting.get_num_units() + + +def test_generate_hybrid_num_units_mismatch_raises(): + templates = _make_templates(num_units=10) + gt_sorting = generate_sorting(durations=[10], num_units=5, sampling_frequency=20000, seed=0) + rec, _ = generate_ground_truth_recording(sampling_frequency=20000, durations=[10], num_channels=16, seed=0) + with pytest.raises(AssertionError): + generate_hybrid_recording(rec, sorting=gt_sorting, templates=templates, seed=0) + + +def test_generate_hybrid_reproducible(): + # the seed must propagate to the internal generate_unit_locations() call so that two calls + # with the same seed produce identical recordings (default path, templates auto-generated). + rec, _ = generate_ground_truth_recording(sampling_frequency=20000, durations=[5], num_channels=16, seed=0) + hybrid_a, _ = generate_hybrid_recording(rec, seed=0) + hybrid_b, _ = generate_hybrid_recording(rec, seed=0) + traces_a = hybrid_a.get_traces(start_frame=0, end_frame=2000) + traces_b = hybrid_b.get_traces(start_frame=0, end_frame=2000) + assert np.array_equal(traces_a, traces_b) + + @pytest.mark.skip("Spykingcircus2 is not stable enought for estimating templates from recording") def test_estimate_templates_from_recording(create_cache_folder): cache_folder = create_cache_folder