Tools for Generation of Hybrid recordings

src/spikeinterface/core/generate.py

@@ -3,7 +3,6 @@
 import warnings
 import numpy as np
 from typing import Union, Optional, List, Literal
-import warnings
 from math import ceil
 
 from .basesorting import SpikeVectorSortingSegment
@@ -1862,7 +1861,7 @@ def get_traces(
             wf = template[start_template:end_template]
             if self.amplitude_vector is not None:
                 wf = wf * self.amplitude_vector[i]
-            traces[start_traces:end_traces] += wf
+            traces[start_traces:end_traces] += wf.astype(traces.dtype, copy=False)
 
         return traces.astype(self.dtype, copy=False)
 

src/spikeinterface/core/node_pipeline.py

@@ -516,7 +516,6 @@ def _init_peak_pipeline(recording, nodes):
     worker_ctx["recording"] = recording
     worker_ctx["nodes"] = nodes
     worker_ctx["max_margin"] = max(node.get_trace_margin() for node in nodes)
-
     return worker_ctx
 
 

src/spikeinterface/core/recording_tools.py

@@ -709,7 +709,7 @@ def get_chunk_with_margin(
     case zero padding is used, in the second case np.pad is called
     with mod="reflect".
     """
-    length = rec_segment.get_num_samples()
+    length = int(rec_segment.get_num_samples())
 
     if channel_indices is None:
         channel_indices = slice(None)

src/spikeinterface/core/template.py

@@ -294,9 +294,9 @@ def from_zarr_group(cls, zarr_group: "zarr.Group") -> "Templates":
         the `add_templates_to_zarr_group` method.
 
         """
-        templates_array = zarr_group["templates_array"]
-        channel_ids = zarr_group["channel_ids"]
-        unit_ids = zarr_group["unit_ids"]
+        templates_array = zarr_group["templates_array"][:]
+        channel_ids = zarr_group["channel_ids"][:]
+        unit_ids = zarr_group["unit_ids"][:]
         sampling_frequency = zarr_group.attrs["sampling_frequency"]
         nbefore = zarr_group.attrs["nbefore"]
 
@@ -305,7 +305,7 @@ def from_zarr_group(cls, zarr_group: "zarr.Group") -> "Templates":
 
         sparsity_mask = None
         if "sparsity_mask" in zarr_group:
-            sparsity_mask = zarr_group["sparsity_mask"]
+            sparsity_mask = zarr_group["sparsity_mask"][:]
 
         probe = None
         if "probe" in zarr_group:
@@ -352,6 +352,55 @@ def to_json(self):
     def from_json(cls, json_str):
         return cls.from_dict(json.loads(json_str))
 
+    def select_units(self, unit_ids):
+        """
+        Return a new Templates object with only the selected units.
+
+        Parameters
+        ----------
+        unit_ids : list
+            List of unit IDs to select.
+        """
+        unit_ids_list = list(self.unit_ids)
+        unit_indices = np.array([unit_ids_list.index(unit_id) for unit_id in unit_ids], dtype=int)
+        sliced_sparsity_mask = None if self.sparsity_mask is None else self.sparsity_mask[unit_indices]
+        return Templates(
+            templates_array=self.templates_array[unit_indices],
+            sampling_frequency=self.sampling_frequency,
+            nbefore=self.nbefore,
+            sparsity_mask=sliced_sparsity_mask,
+            channel_ids=self.channel_ids,
+            unit_ids=unit_ids,
+            probe=self.probe,
+            check_for_consistent_sparsity=False,
+        )
+
+    def select_channels(self, channel_ids):
+        """
+        Return a new Templates object with only the selected channels.
+        This operation can be useful to remove bad channels for hybrid recording
+        generation.
+
+        Parameters
+        ----------
+        channel_ids : list
+            List of channel IDs to select.
+        """
+        assert not self.are_templates_sparse(), "Cannot select channels on sparse templates"
+        channel_ids_list = list(self.channel_ids)
+        channel_indices = np.array([channel_ids_list.index(channel_id) for channel_id in channel_ids])
+        sliced_sparsity_mask = None if self.sparsity_mask is None else self.sparsity_mask[:, channel_indices]
+        return Templates(
+            templates_array=self.templates_array[:, :, channel_indices],
+            sampling_frequency=self.sampling_frequency,
+            nbefore=self.nbefore,
+            sparsity_mask=sliced_sparsity_mask,
+            channel_ids=channel_ids,
+            unit_ids=self.unit_ids,
+            probe=self.probe,
+            check_for_consistent_sparsity=False,
+        )
+
     def __eq__(self, other):
         """
         Necessary to compare templates because they naturally compare objects by equality of their fields
@@ -390,7 +439,7 @@ def __eq__(self, other):
 
         return True
 
-    def get_channel_locations(self):
+    def get_channel_locations(self) -> np.ndarray:
         assert self.probe is not None, "Templates.get_channel_locations() needs a probe to be set"
         channel_locations = self.probe.contact_positions
         return channel_locations

src/spikeinterface/core/template_tools.py

@@ -50,7 +50,7 @@ def _get_nbefore(one_object):
             raise ValueError("SortingAnalyzer need extension 'templates' to be computed")
         return ext.nbefore
     else:
-        raise ValueError("Input should be Templates or SortingAnalyzer or SortingAnalyzer")
+        raise ValueError("Input should be Templates or SortingAnalyzer")
 
 
 def get_template_amplitudes(

src/spikeinterface/generation/__init__.py

@@ -5,6 +5,14 @@
     InjectDriftingTemplatesRecording,
     make_linear_displacement,
 )
+
+from .hybrid_tools import (
+    generate_hybrid_recording,
+    estimate_templates_from_recording,
+    select_templates,
+    scale_templates,
+    shift_templates,
+)
 from .noise_tools import generate_noise
 from .drifting_generator import (
     make_one_displacement_vector,

src/spikeinterface/generation/drift_tools.py

@@ -125,13 +125,13 @@ def __init__(self, **kwargs):
 
     @classmethod
     def from_static(cls, templates):
-        drifting_teplates = cls(
+        drifting_templates = cls(
             templates_array=templates.templates_array,
             sampling_frequency=templates.sampling_frequency,
             nbefore=templates.nbefore,
             probe=templates.probe,
         )
-        return drifting_teplates
+        return drifting_templates
 
     def move_one_template(self, unit_index, displacement, **interpolation_kwargs):
         """
@@ -264,9 +264,9 @@ def __init__(
     ):
         import scipy.spatial
 
-        assert isinstance(
-            drifting_templates, DriftingTemplates
-        ), "drifting_templates must be a DriftingTemplates object"
+        # assert isinstance(
+        #     drifting_templates, DriftingTemplates
+        # ), "drifting_templates must be a DriftingTemplates object"
         self.drifting_templates = drifting_templates
 
         if parent_recording is None:
@@ -442,7 +442,8 @@ def __init__(
         # TODO: self.upsample_vector = upsample_vector
         self.upsample_vector = None
         self.parent_recording = parent_recording_segment
-        self.num_samples = parent_recording_segment.get_num_frames() if num_samples is None else num_samples
+        self.num_samples = parent_recording_segment.get_num_samples() if num_samples is None else num_samples
+        self.num_samples = int(num_samples)
 
         self.displacement_indices = displacement_indices
         self.templates_array_moved = templates_array_moved
@@ -507,7 +508,7 @@ def get_traces(
             wf = template[start_template:end_template]
             if self.amplitude_vector is not None:
                 wf *= self.amplitude_vector[i]
-            traces[start_traces:end_traces] += wf
+            traces[start_traces:end_traces] += wf.astype(self.dtype, copy=False)
 
         return traces.astype(self.dtype)