Repository: SpikeInterface/spikeinterface
Pull request #3476 (2024-10-15 11:49:56)

Unit localization
 1. Use existing sparsity for unit location
 2. Add location with max channel


Patch:
diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
index 4c36d6fb8..1e133694b 100644
--- a/.pre-commit-config.yaml
+++ b/.pre-commit-config.yaml
@@ -6,7 +6,7 @@ repos:
     -   id: end-of-file-fixer
     -   id: trailing-whitespace
 -   repo: https://github.com/psf/black
-    rev: 24.10.0
+    rev: 24.8.0
     hooks:
     -   id: black
         files: ^src/
diff --git a/src/spikeinterface/benchmark/benchmark_matching.py b/src/spikeinterface/benchmark/benchmark_matching.py
index 1934b65ef..3799fa19b 100644
--- a/src/spikeinterface/benchmark/benchmark_matching.py
+++ b/src/spikeinterface/benchmark/benchmark_matching.py
@@ -80,7 +80,6 @@ class MatchingStudy(BenchmarkStudy):
     def plot_collisions(self, case_keys=None, figsize=None):
         if case_keys is None:
             case_keys = list(self.cases.keys())
-        import matplotlib.pyplot as plt
 
         fig, axs = plt.subplots(ncols=len(case_keys), nrows=1, figsize=figsize, squeeze=False)
 
diff --git a/src/spikeinterface/postprocessing/localization_tools.py b/src/spikeinterface/postprocessing/localization_tools.py
index e6278fc59..837b98305 100644
--- a/src/spikeinterface/postprocessing/localization_tools.py
+++ b/src/spikeinterface/postprocessing/localization_tools.py
@@ -77,7 +77,11 @@ def compute_monopolar_triangulation(
 
     contact_locations = sorting_analyzer_or_templates.get_channel_locations()
 
-    sparsity = compute_sparsity(sorting_analyzer_or_templates, method="radius", radius_um=radius_um)
+    if sorting_analyzer_or_templates.sparsity is None:
+        sparsity = compute_sparsity(sorting_analyzer_or_templates, method="radius", radius_um=radius_um)
+    else:
+        sparsity = sorting_analyzer_or_templates.sparsity
+
     templates = get_dense_templates_array(
         sorting_analyzer_or_templates, return_scaled=get_return_scaled(sorting_analyzer_or_templates)
     )
@@ -157,9 +161,13 @@ def compute_center_of_mass(
 
     assert feature in ["ptp", "mean", "energy", "peak_voltage"], f"{feature} is not a valid feature"
 
-    sparsity = compute_sparsity(
-        sorting_analyzer_or_templates, peak_sign=peak_sign, method="radius", radius_um=radius_um
-    )
+    if sorting_analyzer_or_templates.sparsity is None:
+        sparsity = compute_sparsity(
+            sorting_analyzer_or_templates, peak_sign=peak_sign, method="radius", radius_um=radius_um
+        )
+    else:
+        sparsity = sorting_analyzer_or_templates.sparsity
+
     templates = get_dense_templates_array(
         sorting_analyzer_or_templates, return_scaled=get_return_scaled(sorting_analyzer_or_templates)
     )
@@ -650,8 +658,55 @@ if HAVE_NUMBA:
     enforce_decrease_shells = numba.jit(enforce_decrease_shells_data, nopython=True)
 
 
+def compute_location_max_channel(
+    templates_or_sorting_analyzer: SortingAnalyzer | Templates,
+    unit_ids=None,
+    peak_sign: "neg" | "pos" | "both" = "neg",
+    mode: "extremum" | "at_index" | "peak_to_peak" = "extremum",
+) -> np.ndarray:
+    """
+    Localize a unit using max channel.
+
+    This uses internally `get_template_extremum_channel()`
+
+
+    Parameters
+    ----------
+    templates_or_sorting_analyzer : SortingAnalyzer | Templates
+        A SortingAnalyzer or Templates object
+    unit_ids: list[str] | list[int] | None
+        A list of unit_id to restrict the computation
+    peak_sign :  "neg" | "pos" | "both"
+        Sign of the template to find extremum channels
+    mode : "extremum" | "at_index" | "peak_to_peak", default: "at_index"
+        Where the amplitude is computed
+        * "extremum" : take the peak value (max or min depending on `peak_sign`)
+        * "at_index" : take value at `nbefore` index
+        * "peak_to_peak" : take the peak-to-peak amplitude
+
+    Returns
+    -------
+    unit_locations: np.ndarray
+        2d
+    """
+    extremum_channels_index = get_template_extremum_channel(
+        templates_or_sorting_analyzer, peak_sign=peak_sign, mode=mode, outputs="index"
+    )
+    contact_locations = templates_or_sorting_analyzer.get_channel_locations()
+    if unit_ids is None:
+        unit_ids = templates_or_sorting_analyzer.unit_ids
+    else:
+        unit_ids = np.asarray(unit_ids)
+    unit_locations = np.zeros((unit_ids.size, 2), dtype="float32")
+    for i, unit_id in enumerate(unit_ids):
+        unit_locations[i, :] = contact_locations[extremum_channels_index[unit_id]]
+
+    return unit_locations
+
+
 _unit_location_methods = {
     "center_of_mass": compute_center_of_mass,
     "grid_convolution": compute_grid_convolution,
     "monopolar_triangulation": compute_monopolar_triangulation,
+    "max_channel": compute_location_max_channel,
 }
diff --git a/src/spikeinterface/sorters/internal/spyking_circus2.py b/src/spikeinterface/sorters/internal/spyking_circus2.py
index eed693b34..211adba99 100644
--- a/src/spikeinterface/sorters/internal/spyking_circus2.py
+++ b/src/spikeinterface/sorters/internal/spyking_circus2.py
@@ -27,7 +27,7 @@ class Spykingcircus2Sorter(ComponentsBasedSorter):
         "general": {"ms_before": 2, "ms_after": 2, "radius_um": 75},
         "sparsity": {"method": "snr", "amplitude_mode": "peak_to_peak", "threshold": 0.25},
         "filtering": {"freq_min": 150, "freq_max": 7000, "ftype": "bessel", "filter_order": 2},
-        "whitening": {"mode": "local", "regularize": False},
+        "whitening": {"mode": "local", "regularize": True},
         "detection": {"peak_sign": "neg", "detect_threshold": 4},
         "selection": {
             "method": "uniform",
@@ -100,12 +100,6 @@ class Spykingcircus2Sorter(ComponentsBasedSorter):
         except:
             HAVE_HDBSCAN = False
 
-        try:
-            import torch
-        except ImportError:
-            HAVE_TORCH = False
-            print("spykingcircus2 could benefit from using torch. Consider installing it")
-
         assert HAVE_HDBSCAN, "spykingcircus2 needs hdbscan to be installed"
 
         # this is importanted only on demand because numba import are too heavy
diff --git a/src/spikeinterface/sortingcomponents/clustering/circus.py b/src/spikeinterface/sortingcomponents/clustering/circus.py
index 99c59f493..b7e71d3b4 100644
--- a/src/spikeinterface/sortingcomponents/clustering/circus.py
+++ b/src/spikeinterface/sortingcomponents/clustering/circus.py
@@ -92,7 +92,7 @@ class CircusClustering:
         # SVD for time compression
         few_peaks = select_peaks(peaks, recording=recording, method="uniform", n_peaks=10000, margin=(nbefore, nafter))
         few_wfs = extract_waveform_at_max_channel(
-            recording, few_peaks, ms_before=ms_before, ms_after=ms_after, **job_kwargs
+            recording, few_peaks, ms_before=ms_before, ms_after=ms_after, **params["job_kwargs"]
         )
 
         wfs = few_wfs[:, :, 0]
@@ -141,7 +141,7 @@ class CircusClustering:
             all_pc_data = run_node_pipeline(
                 recording,
                 pipeline_nodes,
-                job_kwargs,
+                params["job_kwargs"],
                 job_name="extracting features",
             )
 
@@ -176,7 +176,7 @@ class CircusClustering:
             _ = run_node_pipeline(
                 recording,
                 pipeline_nodes,
-                job_kwargs,
+                params["job_kwargs"],
                 job_name="extracting features",
                 gather_mode="npy",
                 gather_kwargs=dict(exist_ok=True),
diff --git a/src/spikeinterface/sortingcomponents/matching/circus.py b/src/spikeinterface/sortingcomponents/matching/circus.py
index 3b97f2dc6..d1b2139c5 100644
--- a/src/spikeinterface/sortingcomponents/matching/circus.py
+++ b/src/spikeinterface/sortingcomponents/matching/circus.py
@@ -18,15 +18,6 @@ spike_dtype = [
     ("segment_index", "int64"),
 ]
 
-try:
-    import torch
-    import torch.nn.functional as F
-
-    HAVE_TORCH = True
-    from torch.nn.functional import conv1d
-except ImportError:
-    HAVE_TORCH = False
-
 from .base import BaseTemplateMatching
 
 
@@ -52,9 +43,9 @@ def compress_templates(
 
     temporal, singular, spatial = np.linalg.svd(templates_array, full_matrices=False)
     # Keep only the strongest components
-    temporal = temporal[:, :, :approx_rank].astype(np.float32)
-    singular = singular[:, :approx_rank].astype(np.float32)
-    spatial = spatial[:, :approx_rank, :].astype(np.float32)
+    temporal = temporal[:, :, :approx_rank]
+    singular = singular[:, :approx_rank]
+    spatial = spatial[:, :approx_rank, :]
 
     if return_new_templates:
         templates_array = np.matmul(temporal * singular[:, np.newaxis, :], spatial)
@@ -116,22 +107,18 @@ class CircusOMPSVDPeeler(BaseTemplateMatching):
 
     Parameters
     ----------
-    amplitude : tuple
+    amplitude: tuple
         (Minimal, Maximal) amplitudes allowed for every template
-    max_failures : int
+    max_failures: int
         Stopping criteria of the OMP algorithm, as number of retry while updating amplitudes
-    sparse_kwargs : dict
+    sparse_kwargs: dict
         Parameters to extract a sparsity mask from the waveform_extractor, if not
         already sparse.
-    rank : int, default: 5
+    rank: int, default: 5
         Number of components used internally by the SVD
-    vicinity : int
+    vicinity: int
         Size of the area surrounding a spike to perform modification (expressed in terms
         of template temporal width)
-    engine : string in ["numpy", "torch", "auto"]. Default "auto"
-        The engine to use for the convolutions
-    torch_device : string in ["cpu", "cuda", None]. Default "cpu"
-        Controls torch device if the torch engine is selected
     -----
     """
 
@@ -161,8 +148,6 @@ class CircusOMPSVDPeeler(BaseTemplateMatching):
         ignore_inds=[],
         vicinity=2,
         precomputed=None,
-        engine="numpy",
-        torch_device="cpu",
     ):
 
         BaseTemplateMatching.__init__(self, recording, templates, return_output=True, parents=None)
@@ -173,19 +158,6 @@ class CircusOMPSVDPeeler(BaseTemplateMatching):
         self.nafter = templates.nafter
         self.sampling_frequency = recording.get_sampling_frequency()
         self.vicinity = vicinity * self.num_samples
-        assert engine in ["numpy", "torch", "auto"], "engine should be numpy, torch or auto"
-        if engine == "auto":
-            if HAVE_TORCH:
-                self.engine = "torch"
-            else:
-                self.engine = "numpy"
-        else:
-            if engine == "torch":
-                assert HAVE_TORCH, "please install torch to use the torch engine"
-            self.engine = engine
-
-        assert torch_device in ["cuda", "cpu", None]
-        self.torch_device = torch_device
 
         self.amplitudes = amplitudes
         self.stop_criteria = stop_criteria
@@ -211,7 +183,6 @@ class CircusOMPSVDPeeler(BaseTemplateMatching):
             self.unit_overlaps_tables[i][self.unit_overlaps_indices[i]] = np.arange(len(self.unit_overlaps_indices[i]))
 
         self.margin = 2 * self.num_samples
-        self.is_pushed = False
 
     def _prepare_templates(self):
 
@@ -283,14 +254,6 @@ class CircusOMPSVDPeeler(BaseTemplateMatching):
         self.temporal = np.moveaxis(self.temporal, [0, 1, 2], [1, 2, 0])
         self.singular = self.singular.T[:, :, np.newaxis]
 
-    def _push_to_torch(self):
-        if self.engine == "torch":
-            self.spatial = torch.as_tensor(self.spatial, device=self.torch_device)
-            self.singular = torch.as_tensor(self.singular, device=self.torch_device)
-            self.temporal = torch.as_tensor(self.temporal.copy(), device=self.torch_device).swapaxes(0, 1)
-            self.temporal = torch.flip(self.temporal, (2,))
-        self.is_pushed = True
-
     def get_extra_outputs(self):
         output = {}
         for key in self._more_output_keys:
@@ -305,15 +268,15 @@ class CircusOMPSVDPeeler(BaseTemplateMatching):
         import scipy
         from scipy import ndimage
 
-        if not self.is_pushed:
-            self._push_to_torch()
-
         (potrs,) = scipy.linalg.get_lapack_funcs(("potrs",), dtype=np.float32)
+
         (nrm2,) = scipy.linalg.get_blas_funcs(("nrm2",), dtype=np.float32)
 
-        omp_tol = np.finfo(np.float32).eps
-        neighbor_window = self.num_samples - 1
+        overlaps_array = self.overlaps
 
+        omp_tol = np.finfo(np.float32).eps
+        num_samples = self.nafter + self.nbefore
+        neighbor_window = num_samples - 1
         if isinstance(self.amplitudes, list):
             min_amplitude, max_amplitude = self.amplitudes
         else:
@@ -321,36 +284,27 @@ class CircusOMPSVDPeeler(BaseTemplateMatching):
             min_amplitude = min_amplitude[:, np.newaxis]
             max_amplitude = max_amplitude[:, np.newaxis]
 
-        if self.engine == "torch":
-            blank = np.zeros((neighbor_window, self.num_channels), dtype=np.float32)
-            traces = np.vstack((blank, traces, blank))
-            num_timesteps = traces.shape[0]
-            torch_traces = torch.as_tensor(traces.T[np.newaxis, :, :], device=self.torch_device)
-            num_templates, num_channels = self.temporal.shape[0], self.temporal.shape[1]
-            spatially_filtered_data = torch.matmul(self.spatial, torch_traces)
-            scaled_filtered_data = (spatially_filtered_data * self.singular).swapaxes(0, 1)
-            scaled_filtered_data_ = scaled_filtered_data.reshape(1, num_templates * num_channels, num_timesteps)
-            scalar_products = conv1d(scaled_filtered_data_, self.temporal, groups=num_templates, padding="valid")
-            scalar_products = scalar_products.cpu().numpy()[0, :, self.num_samples - 1 : -neighbor_window]
-        else:
-            num_timesteps = traces.shape[0]
-            num_peaks = num_timesteps - neighbor_window
-            conv_shape = (self.num_templates, num_peaks)
-            scalar_products = np.zeros(conv_shape, dtype=np.float32)
-            # Filter using overlap-and-add convolution
-            spatially_filtered_data = np.matmul(self.spatial, traces.T[np.newaxis, :, :])
-            scaled_filtered_data = spatially_filtered_data * self.singular
-            from scipy import signal
+        num_timesteps = len(traces)
 
-            objective_by_rank = signal.oaconvolve(scaled_filtered_data, self.temporal, axes=2, mode="valid")
-            scalar_products += np.sum(objective_by_rank, axis=0)
-
-        num_peaks = scalar_products.shape[1]
+        num_peaks = num_timesteps - num_samples + 1
+        conv_shape = (self.num_templates, num_peaks)
+        scalar_products = np.zeros(conv_shape, dtype=np.float32)
 
         # Filter using overlap-and-add convolution
         if len(self.ignore_inds) > 0:
-            scalar_products[self.ignore_inds] = -np.inf
             not_ignored = ~np.isin(np.arange(self.num_templates), self.ignore_inds)
+            spatially_filtered_data = np.matmul(self.spatial[:, not_ignored, :], traces.T[np.newaxis, :, :])
+            scaled_filtered_data = spatially_filtered_data * self.singular[:, not_ignored, :]
+            objective_by_rank = scipy.signal.oaconvolve(
+                scaled_filtered_data, self.temporal[:, not_ignored, :], axes=2, mode="valid"
+            )
+            scalar_products[not_ignored] += np.sum(objective_by_rank, axis=0)
+            scalar_products[self.ignore_inds] = -np.inf
+        else:
+            spatially_filtered_data = np.matmul(self.spatial, traces.T[np.newaxis, :, :])
+            scaled_filtered_data = spatially_filtered_data * self.singular
+            objective_by_rank = scipy.signal.oaconvolve(scaled_filtered_data, self.temporal, axes=2, mode="valid")
+            scalar_products += np.sum(objective_by_rank, axis=0)
 
         num_spikes = 0
 
@@ -368,7 +322,7 @@ class CircusOMPSVDPeeler(BaseTemplateMatching):
         is_in_vicinity = np.zeros(0, dtype=np.int32)
 
         if self.stop_criteria == "omp_min_sps":
-            stop_criteria = self.omp_min_sps * np.maximum(self.norms, np.sqrt(self.num_channels * self.num_samples))
+            stop_criteria = self.omp_min_sps * np.maximum(self.norms, np.sqrt(self.num_channels * num_samples))
         elif self.stop_criteria == "max_failures":
             num_valids = 0
             nb_failures = self.max_failures
@@ -400,11 +354,11 @@ class CircusOMPSVDPeeler(BaseTemplateMatching):
 
                 if num_selection > 0:
                     delta_t = selection[1] - peak_index
-                    idx = np.flatnonzero((delta_t < self.num_samples) & (delta_t > -self.num_samples))
+                    idx = np.flatnonzero((delta_t < num_samples) & (delta_t > -num_samples))
                     myline = neighbor_window + delta_t[idx]
                     myindices = selection[0, idx]
 
-                    local_overlaps = self.overlaps[best_cluster_ind]
+                    local_overlaps = overlaps_array[best_cluster_ind]
                     overlapping_templates = self.unit_overlaps_indices[best_cluster_ind]
                     table = self.unit_overlaps_tables[best_cluster_ind]
 
@@ -482,10 +436,10 @@ class CircusOMPSVDPeeler(BaseTemplateMatching):
                 for i in modified:
                     tmp_best, tmp_peak = sub_selection[:, i]
                     diff_amp = diff_amplitudes[i] * self.norms[tmp_best]
-                    local_overlaps = self.overlaps[tmp_best]
+                    local_overlaps = overlaps_array[tmp_best]
                     overlapping_templates = self.units_overlaps[tmp_best]
                     tmp = tmp_peak - neighbor_window
-                    idx = [max(0, tmp), min(num_peaks, tmp_peak + self.num_samples)]
+                    idx = [max(0, tmp), min(num_peaks, tmp_peak + num_samples)]
                     tdx = [idx[0] - tmp, idx[1] - tmp]
                     to_add = diff_amp * local_overlaps[:, tdx[0] : tdx[1]]
                     scalar_products[overlapping_templates, idx[0] : idx[1]] -= to_add
@@ -546,27 +500,27 @@ class CircusPeeler(BaseTemplateMatching):
 
     Parameters
     ----------
-    peak_sign : str
+    peak_sign: str
         Sign of the peak (neg, pos, or both)
-    exclude_sweep_ms : float
+    exclude_sweep_ms: float
         The number of samples before/after to classify a peak (should be low)
-    jitter : int
+    jitter: int
         The number of samples considered before/after every peak to search for
         matches
-    detect_threshold : int
+    detect_threshold: int
         The detection threshold
-    noise_levels : array
+    noise_levels: array
         The noise levels, for every channels
-    random_chunk_kwargs : dict
+    random_chunk_kwargs: dict
         Parameters for computing noise levels, if not provided (sub optimal)
-    max_amplitude : float
+    max_amplitude: float
         Maximal amplitude allowed for every template
-    min_amplitude : float
+    min_amplitude: float
         Minimal amplitude allowed for every template
-    use_sparse_matrix_threshold : float
+    use_sparse_matrix_threshold: float
         If density of the templates is below a given threshold, sparse matrix
         are used (memory efficient)
-    sparse_kwargs : dict
+    sparse_kwargs: dict
         Parameters to extract a sparsity mask from the waveform_extractor, if not
         already sparse.
     -----
diff --git a/src/spikeinterface/sortingcomponents/matching/wobble.py b/src/spikeinterface/sortingcomponents/matching/wobble.py
index 59e171fe5..3099448b1 100644
--- a/src/spikeinterface/sortingcomponents/matching/wobble.py
+++ b/src/spikeinterface/sortingcomponents/matching/wobble.py
@@ -8,15 +8,6 @@ from typing import List, Tuple, Optional
 from .base import BaseTemplateMatching, _base_matching_dtype
 from spikeinterface.core.template import Templates
 
-try:
-    import torch
-    import torch.nn.functional as F
-
-    HAVE_TORCH = True
-    from torch.nn.functional import conv1d
-except ImportError:
-    HAVE_TORCH = False
-
 
 @dataclass
 class WobbleParameters:
@@ -50,10 +41,6 @@ class WobbleParameters:
         Maximum value for ampltiude scaling of templates.
     scale_amplitudes : bool
         If True, scale amplitudes of templates to match spikes.
-    engine : string in ["numpy", "torch", "auto"]. Default "auto"
-        The engine to use for the convolutions
-    torch_device : string in ["cpu", "cuda", None]. Default "cpu"
-        Controls torch device if the torch engine is selected
 
     Notes
     -----
@@ -75,8 +62,6 @@ class WobbleParameters:
     scale_min: float = 0
     scale_max: float = np.inf
     scale_amplitudes: bool = False
-    engine: str = "numpy"
-    torch_device: str = "cpu"
 
     def __post_init__(self):
         assert self.amplitude_variance >= 0, "amplitude_variance must be a non-negative scalar"
@@ -359,8 +344,6 @@ class WobbleMatch(BaseTemplateMatching):
         parents=None,
         templates=None,
         parameters={},
-        engine="numpy",
-        torch_device="cpu",
     ):
 
         BaseTemplateMatching.__init__(self, recording, templates, return_output=True, parents=None)
@@ -369,21 +352,6 @@ class WobbleMatch(BaseTemplateMatching):
 
         # Aggregate useful parameters/variables for handy access in downstream functions
         params = WobbleParameters(**parameters)
-
-        assert engine in ["numpy", "torch", "auto"], "engine should be numpy, torch or auto"
-        if engine == "auto":
-            if HAVE_TORCH:
-                self.engine = "torch"
-            else:
-                self.engine = "numpy"
-        else:
-            if engine == "torch":
-                assert HAVE_TORCH, "please install torch to use the torch engine"
-            self.engine = engine
-
-        assert torch_device in ["cuda", "cpu", None]
-        self.torch_device = torch_device
-
         template_meta = TemplateMetadata.from_parameters_and_templates(params, templates_array)
         if not templates.are_templates_sparse():
             sparsity = WobbleSparsity.from_parameters_and_templates(params, templates_array)
@@ -398,21 +366,13 @@ class WobbleMatch(BaseTemplateMatching):
         pairwise_convolution = convolve_templates(
             compressed_templates, params.jitter_factor, params.approx_rank, template_meta.jittered_indices, sparsity
         )
-
         norm_squared = compute_template_norm(sparsity.visible_channels, templates_array)
-
-        spatial = np.moveaxis(spatial, [0, 1, 2], [1, 0, 2])
-        temporal = np.moveaxis(temporal, [0, 1, 2], [1, 2, 0])
-        singular = singular.T[:, :, np.newaxis]
-
-        compressed_templates = (temporal, singular, spatial, temporal_jittered)
         template_data = TemplateData(
             compressed_templates=compressed_templates,
             pairwise_convolution=pairwise_convolution,
             norm_squared=norm_squared,
         )
 
-        self.is_pushed = False
         self.params = params
         self.template_meta = template_meta
         self.sparsity = sparsity
@@ -424,24 +384,11 @@ class WobbleMatch(BaseTemplateMatching):
         # self.margin = int(buffer_ms*1e-3 * recording.sampling_frequency)
         self.margin = 300  # To ensure equivalence with spike-psvae version of the algorithm
 
-    def _push_to_torch(self):
-        if self.engine == "torch":
-            temporal, singular, spatial, temporal_jittered = self.template_data.compressed_templates
-            spatial = torch.as_tensor(spatial, device=self.torch_device)
-            singular = torch.as_tensor(singular, device=self.torch_device)
-            temporal = torch.as_tensor(temporal.copy(), device=self.torch_device).swapaxes(0, 1)
-            temporal = torch.flip(temporal, (2,))
-            self.template_data.compressed_templates = (temporal, singular, spatial, temporal_jittered)
-        self.is_pushed = True
-
     def get_trace_margin(self):
         return self.margin
 
     def compute_matching(self, traces, start_frame, end_frame, segment_index):
 
-        if not self.is_pushed:
-            self._push_to_torch()
-
         # Unpack method_kwargs
         # nbefore, nafter = method_kwargs["nbefore"], method_kwargs["nafter"]
         # template_meta = method_kwargs["template_meta"]
@@ -453,9 +400,7 @@ class WobbleMatch(BaseTemplateMatching):
         assert traces.dtype == np.float32, "traces must be specified as np.float32"
 
         # Compute objective
-        objective = compute_objective(
-            traces, self.template_data, self.params.approx_rank, self.engine, self.torch_device
-        )
+        objective = compute_objective(traces, self.template_data, self.params.approx_rank)
         objective_normalized = 2 * objective - self.template_data.norm_squared[:, np.newaxis]
 
         # Compute spike train
@@ -841,11 +786,10 @@ def compress_templates(templates, approx_rank) -> tuple[np.ndarray, np.ndarray,
     temporal, singular, spatial = np.linalg.svd(templates, full_matrices=False)
 
     # Keep only the strongest components
-    temporal = temporal[:, :, :approx_rank].astype(np.float32)
+    temporal = temporal[:, :, :approx_rank]
     temporal = np.flip(temporal, axis=1)
-    singular = singular[:, :approx_rank].astype(np.float32)
-    spatial = spatial[:, :approx_rank, :].astype(np.float32)
-
+    singular = singular[:, :approx_rank]
+    spatial = spatial[:, :approx_rank, :]
     return temporal, singular, spatial
 
 
@@ -883,6 +827,7 @@ def upsample_and_jitter(temporal, jitter_factor, num_samples):
 
     shape_temporal_jittered = (-1, num_samples, approx_rank)
     temporal_jittered = np.reshape(temporal_jittered[:, shifted_index, :], shape_temporal_jittered)
+
     temporal_jittered = np.flip(temporal_jittered, axis=1)
     return temporal_jittered
 
@@ -944,7 +889,7 @@ def convolve_templates(compressed_templates, jitter_factor, approx_rank, jittere
     return pairwise_convolution
 
 
-def compute_objective(traces, template_data, approx_rank, engine="numpy", torch_device=None) -> np.ndarray:
+def compute_objective(traces, template_data, approx_rank) -> np.ndarray:
     """Compute objective by convolving templates with voltage traces.
 
     Parameters
@@ -953,39 +898,31 @@ def compute_objective(traces, template_data, approx_rank, engine="numpy", torch_
         Voltage traces for a chunk of the recording.
     template_data : TemplateData
         Dataclass object for aggregating template data together.
+    approx_rank : int
+        Rank of the compressed template matrices.
 
     Returns
     -------
     objective : ndarray (template_meta.num_templates, traces.shape[0]+template_meta.num_samples-1)
             Template matching objective for each template.
     """
-    temporal, singular, spatial, _ = template_data.compressed_templates
-    if engine == "torch":
-        nt = temporal.shape[2] - 1
-        num_channels = traces.shape[1]
-        blank = np.zeros((nt, num_channels), dtype=np.float32)
-        traces = np.vstack((blank, traces, blank))
-        torch_traces = torch.as_tensor(traces.T[None, :, :], device=torch_device)
-        num_templates, num_channels = temporal.shape[0], temporal.shape[1]
-        num_timesteps = torch_traces.shape[2]
-        spatially_filtered_data = torch.matmul(spatial, torch_traces)
-        scaled_filtered_data = (spatially_filtered_data * singular).swapaxes(0, 1)
-        scaled_filtered_data_ = scaled_filtered_data.reshape(1, num_templates * num_channels, num_timesteps)
-        objective = conv1d(scaled_filtered_data_, temporal, groups=num_templates, padding="valid")
-        objective = objective.cpu().numpy()[0, :, :]
-    elif engine == "numpy":
-        num_channels, num_templates = temporal.shape[0], temporal.shape[1]
-        num_timesteps = temporal.shape[2]
-        objective_len = get_convolution_len(traces.shape[0], num_timesteps)
-        conv_shape = (num_templates, objective_len)
-        objective = np.zeros(conv_shape, dtype=np.float32)
-        # Filter using overlap-and-add convolution
-        spatially_filtered_data = np.matmul(spatial, traces.T[np.newaxis, :, :])
-        scaled_filtered_data = spatially_filtered_data * singular
-        from scipy import signal
+    temporal, singular, spatial, temporal_jittered = template_data.compressed_templates
+    num_templates = temporal.shape[0]
+    num_samples = temporal.shape[1]
+    objective_len = get_convolution_len(traces.shape[0], num_samples)
+    conv_shape = (num_templates, objective_len)
+    objective = np.zeros(conv_shape, dtype=np.float32)
+    spatial_filters = np.moveaxis(spatial[:, :approx_rank, :], [0, 1, 2], [1, 0, 2])
+    temporal_filters = np.moveaxis(temporal[:, :, :approx_rank], [0, 1, 2], [1, 2, 0])
+    singular_filters = singular.T[:, :, np.newaxis]
+
+    # Filter using overlap-and-add convolution
+    spatially_filtered_data = np.matmul(spatial_filters, traces.T[np.newaxis, :, :])
+    scaled_filtered_data = spatially_filtered_data * singular_filters
+    from scipy import signal
 
-        objective_by_rank = signal.oaconvolve(scaled_filtered_data, temporal, axes=2, mode="full")
-        objective += np.sum(objective_by_rank, axis=0)
+    objective_by_rank = signal.oaconvolve(scaled_filtered_data, temporal_filters, axes=2, mode="full")
+    objective += np.sum(objective_by_rank, axis=0)
     return objective
 
 
diff --git a/src/spikeinterface/sortingcomponents/peak_detection.py b/src/spikeinterface/sortingcomponents/peak_detection.py
index 5b1d33b33..d608c5d10 100644
--- a/src/spikeinterface/sortingcomponents/peak_detection.py
+++ b/src/spikeinterface/sortingcomponents/peak_detection.py
@@ -603,13 +603,13 @@ class DetectPeakMatchedFiltering(PeakDetector):
     params_doc = (
         DetectPeakByChannel.params_doc
         + """
-    radius_um : float
+    radius_um: float
         The radius to use to select neighbour channels for locally exclusive detection.
-    prototype : array
+    prototype: array
         The canonical waveform of action potentials
-    ms_before : float
-        The time in ms before the maximial value of the absolute prototype
-    weight_method : dict
+    rank : int (default 1)
+        The rank for SVD convolution of spatiotemporal templates with the traces
+    weight_method: dict
         Parameter that should be provided to the get_convolution_weights() function
         in order to know how to estimate the positions. One argument is mode that could
         be either gaussian_2d (KS like) or exponential_3d (default)
@@ -625,12 +625,12 @@ class DetectPeakMatchedFiltering(PeakDetector):
         detect_threshold=5,
         exclude_sweep_ms=0.1,
         radius_um=50,
+        rank=1,
         noise_levels=None,
         random_chunk_kwargs={"num_chunks_per_segment": 5},
         weight_method={},
     ):
         PeakDetector.__init__(self, recording, return_output=True)
-        from scipy.sparse import csr_matrix
 
         if not HAVE_NUMBA:
             raise ModuleNotFoundError('matched_filtering" needs numba which is not installed')
@@ -664,7 +664,7 @@ class DetectPeakMatchedFiltering(PeakDetector):
             self.num_templates *= 2
 
         self.weights = self.weights.reshape(self.num_templates * self.num_z_factors, -1)
-        self.weights = csr_matrix(self.weights)
+
         random_data = get_random_data_chunks(recording, return_scaled=False, **random_chunk_kwargs)
         conv_random_data = self.get_convolved_traces(random_data)
         medians = np.median(conv_random_data, axis=1)
@@ -734,10 +734,10 @@ class DetectPeakMatchedFiltering(PeakDetector):
         return (local_peaks,)
 
     def get_convolved_traces(self, traces):
-        from scipy.signal import oaconvolve
+        import scipy.signal
 
-        tmp = oaconvolve(self.prototype[None, :], traces.T, axes=1, mode="valid")
-        scalar_products = self.weights.dot(tmp)
+        tmp = scipy.signal.oaconvolve(self.prototype[None, :], traces.T, axes=1, mode="valid")
+        scalar_products = np.dot(self.weights, tmp)
         return scalar_products
 
 


Test patch:
diff --git a/src/spikeinterface/postprocessing/tests/test_unit_locations.py b/src/spikeinterface/postprocessing/tests/test_unit_locations.py
index c40a917a2..545edb349 100644
--- a/src/spikeinterface/postprocessing/tests/test_unit_locations.py
+++ b/src/spikeinterface/postprocessing/tests/test_unit_locations.py
@@ -13,6 +13,7 @@ class TestUnitLocationsExtension(AnalyzerExtensionCommonTestSuite):
             dict(method="grid_convolution", radius_um=150, weight_method={"mode": "gaussian_2d"}),
             dict(method="monopolar_triangulation", radius_um=150),
             dict(method="monopolar_triangulation", radius_um=150, optimizer="minimize_with_log_penality"),
+            dict(method="max_channel"),
         ],
     )
     def test_extension(self, params):
diff --git a/src/spikeinterface/sortingcomponents/tests/test_wobble.py b/src/spikeinterface/sortingcomponents/tests/test_wobble.py
index 0d46b790a..d6d1e1e0b 100644
--- a/src/spikeinterface/sortingcomponents/tests/test_wobble.py
+++ b/src/spikeinterface/sortingcomponents/tests/test_wobble.py
@@ -44,7 +44,7 @@ def test_compress_templates():
         elif test_case == "num_channels == num_samples":
             num_channels = rng.integers(1, 100)
             num_samples = num_channels
-        templates = rng.random((num_templates, num_samples, num_channels), dtype=np.float32)
+        templates = rng.random((num_templates, num_samples, num_channels))
         full_rank = np.minimum(num_samples, num_channels)
         approx_rank = rng.integers(1, full_rank)
 
@@ -66,31 +66,15 @@ def test_compress_templates():
         assert np.all(singular_full >= 0)
         # check that svd matrices are orthonormal if applicable
         if num_channels > num_samples:
-            assert np.allclose(
-                np.matmul(temporal_full, temporal_full.transpose(0, 2, 1)),
-                np.eye(num_samples, dtype=np.float32),
-                atol=1e-3,
-            )
+            assert np.allclose(np.matmul(temporal_full, temporal_full.transpose(0, 2, 1)), np.eye(num_samples))
         elif num_samples > num_channels:
-            assert np.allclose(
-                np.matmul(spatial_full, spatial_full.transpose(0, 2, 1)),
-                np.eye(num_channels, dtype=np.float32),
-                atol=1e-3,
-            )
+            assert np.allclose(np.matmul(spatial_full, spatial_full.transpose(0, 2, 1)), np.eye(num_channels))
         elif num_channels == num_samples:
-            assert np.allclose(
-                np.matmul(temporal_full, temporal_full.transpose(0, 2, 1)),
-                np.eye(num_samples, dtype=np.float32),
-                atol=1e-3,
-            )
-            assert np.allclose(
-                np.matmul(spatial_full, spatial_full.transpose(0, 2, 1)),
-                np.eye(num_channels, dtype=np.float32),
-                atol=1e-3,
-            )
+            assert np.allclose(np.matmul(temporal_full, temporal_full.transpose(0, 2, 1)), np.eye(num_samples))
+            assert np.allclose(np.matmul(spatial_full, spatial_full.transpose(0, 2, 1)), np.eye(num_channels))
             # check that the full rank svd matrices reconstruct the original templates
             reconstructed_templates = np.matmul(temporal_full * singular_full[:, np.newaxis, :], spatial_full)
-            assert np.allclose(reconstructed_templates, templates, atol=1e-3)
+            assert np.allclose(reconstructed_templates, templates)
 
 
 def test_upsample_and_jitter():
@@ -227,33 +211,18 @@ def test_compute_objective():
     approx_rank = rng.integers(1, num_samples)
     num_channels = rng.integers(1, 100)
     chunk_len = rng.integers(num_samples * 2, num_samples * 10)
-    traces = rng.random((chunk_len, num_channels), dtype=np.float32)
+    traces = rng.random((chunk_len, num_channels))
     temporal = rng.random((num_templates, num_samples, approx_rank))
     singular = rng.random((num_templates, approx_rank))
     spatial = rng.random((num_templates, approx_rank, num_channels))
-
-    spatial_transformed = np.moveaxis(spatial, [0, 1, 2], [1, 0, 2])
-    temporal_transformed = np.moveaxis(temporal, [0, 1, 2], [1, 2, 0])
-    singular_transformed = singular.T[:, :, np.newaxis]
-
-    compressed_templates_transformed = (
-        temporal_transformed,
-        singular_transformed,
-        spatial_transformed,
-        temporal_transformed,
-    )
-    norm_squared = np.random.rand(num_templates)
-
-    template_data_transformed = wobble.TemplateData(
-        compressed_templates=compressed_templates_transformed, pairwise_convolution=[], norm_squared=norm_squared
-    )
-    # Act: run compute_objective
-    objective = wobble.compute_objective(traces, template_data_transformed, approx_rank, engine="numpy")
-
     compressed_templates = (temporal, singular, spatial, temporal)
+    norm_squared = np.random.rand(num_templates)
     template_data = wobble.TemplateData(
         compressed_templates=compressed_templates, pairwise_convolution=[], norm_squared=norm_squared
     )
+
+    # Act: run compute_objective
+    objective = wobble.compute_objective(traces, template_data, approx_rank)
     expected_objective = compute_objective_loopy(traces, template_data, approx_rank)
 
     # Assert: check shape and equivalence to expected_objective