Repository: SpikeInterface/spikeinterface
Pull request #2623 (2024-03-25 15:21:22)

Fast templates: compute `std` with and extend to widgets
With the same approach of computing the template means with an accumulator, this PR adds the computation of std "online".

The idea is the following:

$$
template_{std}  = \sqrt{\frac{\sum_i (w_i - \mu)^2}{N}} = \sqrt{\frac{\sum_i w_i^2 + N \cdot \mu - 2 \cdot \mu \cdot \sum_i w_i}{N}}
$$

Since the `estimate_template_average` (now renamed `extimate_template_online` -- another possible name is `estimate_template_one_pass`/`from_traces`) already accumulates the waveforms $\sum_i w_i$, we just need to extent the mechanism to $\sum_i w_i^2$ and the trick is quickly done.

Added a test to assess it gives the same result as `ComputeTemplates` ;)


Patch:
diff --git a/src/spikeinterface/core/__init__.py b/src/spikeinterface/core/__init__.py
index 48c1f0b80..de1e9d9a4 100644
--- a/src/spikeinterface/core/__init__.py
+++ b/src/spikeinterface/core/__init__.py
@@ -103,7 +103,7 @@ from .recording_tools import (
 )
 from .sorting_tools import spike_vector_to_spike_trains, random_spikes_selection
 
-from .waveform_tools import extract_waveforms_to_buffers, estimate_templates, estimate_templates_average
+from .waveform_tools import extract_waveforms_to_buffers, estimate_templates, estimate_templates_with_accumulator
 from .snippets_tools import snippets_from_sorting
 
 # waveform extractor
diff --git a/src/spikeinterface/core/analyzer_extension_core.py b/src/spikeinterface/core/analyzer_extension_core.py
index e60712cd3..73ed97aae 100644
--- a/src/spikeinterface/core/analyzer_extension_core.py
+++ b/src/spikeinterface/core/analyzer_extension_core.py
@@ -13,7 +13,7 @@ It also implement:
 import numpy as np
 
 from .sortinganalyzer import AnalyzerExtension, register_result_extension
-from .waveform_tools import extract_waveforms_to_single_buffer, estimate_templates_average
+from .waveform_tools import extract_waveforms_to_single_buffer, estimate_templates_with_accumulator
 from .recording_tools import get_noise_levels
 from .template import Templates
 from .sorting_tools import random_spikes_selection
@@ -252,7 +252,7 @@ class ComputeWaveforms(AnalyzerExtension):
             chan_inds = self.sorting_analyzer.sparsity.unit_id_to_channel_indices[unit_id]
             wfs = wfs[:, :, : chan_inds.size]
             if force_dense:
-                num_channels = self.get_num_channels()
+                num_channels = self.sorting_analyzer.get_num_channels()
                 dense_wfs = np.zeros((wfs.shape[0], wfs.shape[1], num_channels), dtype=wfs.dtype)
                 dense_wfs[:, :, chan_inds] = wfs
                 wfs = dense_wfs
@@ -410,9 +410,9 @@ class ComputeTemplates(AnalyzerExtension):
 
     def get_templates(self, unit_ids=None, operator="average", percentile=None, save=True):
         """
-        Return templates (average, std, median or percentil) for multiple units.
+        Return templates (average, std, median or percentiles) for multiple units.
 
-        I not computed yet then this is computed on demand and optionally saved.
+        If not computed yet then this is computed on demand and optionally saved.
 
         Parameters
         ----------
@@ -486,7 +486,7 @@ register_result_extension(ComputeTemplates)
 class ComputeFastTemplates(AnalyzerExtension):
     """
     AnalyzerExtension which is similar to the extension "templates" (ComputeTemplates)
-    **but only for average**.
+    **but only for average and standard deviation**.
     This is way faster because it does not need "waveforms" to be computed first.
 
     Parameters
@@ -531,8 +531,15 @@ class ComputeFastTemplates(AnalyzerExtension):
         return_scaled = self.params["return_scaled"]
 
         # TODO jobw_kwargs
-        self.data["average"] = estimate_templates_average(
-            recording, some_spikes, unit_ids, self.nbefore, self.nafter, return_scaled=return_scaled, **job_kwargs
+        self.data["average"], self.data["std"] = estimate_templates_with_accumulator(
+            recording,
+            some_spikes,
+            unit_ids,
+            self.nbefore,
+            self.nafter,
+            return_scaled=return_scaled,
+            return_std=True,
+            **job_kwargs,
         )
 
     def _set_params(
@@ -573,28 +580,30 @@ class ComputeFastTemplates(AnalyzerExtension):
 
         return new_data
 
-    def get_templates(self, unit_ids=None, operator="average"):
+    def get_templates(self, unit_ids=None, operator="average", save=True):
         """
-        Return average templates for multiple units.
+        Return templates (average, std) for multiple units.
 
         Parameters
         ----------
-        unit_ids: list or None, default: None
+        unit_ids: list or None
             Unit ids to retrieve waveforms for
-        operator: str
-            MUST be "average" (only one supported by fast_templates)
-            The argument exist to have the same signature as ComputeTemplates.get_templates
+        operator: "average" | "std", default: "average"
+            The mode to compute the templates
+        save: bool, default True
+            In case, the operator is not computed yet it can be saved to folder or zarr.
 
         Returns
         -------
         templates: np.array
             The returned templates (num_units, num_samples, num_channels)
         """
+        assert operator in ("average", "std"), "ComputeFastTemplates only support average and std"
 
-        assert (
-            operator == "average"
-        ), f"Analyzer extension `fast_templates` only works with 'average' templates. Given operator = {operator}"
-        templates = self.data["average"]
+        templates = self.data[operator]
+
+        if save:
+            self.save()
 
         if unit_ids is not None:
             unit_indices = self.sorting_analyzer.sorting.ids_to_indices(unit_ids)
diff --git a/src/spikeinterface/core/sortinganalyzer.py b/src/spikeinterface/core/sortinganalyzer.py
index 06e4d032c..b1e3adaad 100644
--- a/src/spikeinterface/core/sortinganalyzer.py
+++ b/src/spikeinterface/core/sortinganalyzer.py
@@ -32,7 +32,7 @@ from .node_pipeline import run_node_pipeline
 
 # high level function
 def create_sorting_analyzer(
-    sorting, recording, format="memory", folder=None, sparse=True, sparsity=None, **sparsity_kwargs
+    sorting, recording, format="memory", folder=None, sparse=True, sparsity=None, overwrite=False, **sparsity_kwargs
 ):
     """
     Create a SortingAnalyzer by pairing a Sorting and the corresponding Recording.
@@ -96,6 +96,12 @@ def create_sorting_analyzer(
     In some situation, sparsity is not needed, so to make it fast creation, you need to turn
     sparsity off (or give external sparsity) like this.
     """
+    if format != "memory":
+        if Path(folder).is_dir():
+            if not overwrite:
+                raise ValueError(f"Folder already exists {folder}! Use overwrite=True to overwrite it.")
+            else:
+                shutil.rmtree(folder)
 
     # handle sparsity
     if sparsity is not None:
diff --git a/src/spikeinterface/core/sparsity.py b/src/spikeinterface/core/sparsity.py
index 5d6946456..5e9fa84c6 100644
--- a/src/spikeinterface/core/sparsity.py
+++ b/src/spikeinterface/core/sparsity.py
@@ -7,7 +7,7 @@ from .basesorting import BaseSorting
 from .baserecording import BaseRecording
 from .sorting_tools import random_spikes_selection
 from .job_tools import _shared_job_kwargs_doc
-from .waveform_tools import estimate_templates_average
+from .waveform_tools import estimate_templates_with_accumulator
 
 
 _sparsity_doc = """
@@ -553,7 +553,7 @@ def estimate_sparsity(
       * all units are computed in one read of recording
       * it doesn't require a folder
       * it doesn't consume too much memory
-      * it uses internally the `estimate_templates_average()` which is fast and parallel
+      * it uses internally the `estimate_templates_with_accumulator()` which is fast and parallel
 
     Parameters
     ----------
@@ -613,7 +613,7 @@ def estimate_sparsity(
     spikes = sorting.to_spike_vector()
     spikes = spikes[random_spikes_indices]
 
-    templates_array = estimate_templates_average(
+    templates_array = estimate_templates_with_accumulator(
         recording,
         spikes,
         sorting.unit_ids,
diff --git a/src/spikeinterface/core/waveform_tools.py b/src/spikeinterface/core/waveform_tools.py
index b066107e6..0b51b662e 100644
--- a/src/spikeinterface/core/waveform_tools.py
+++ b/src/spikeinterface/core/waveform_tools.py
@@ -712,7 +712,7 @@ def estimate_templates(
 ):
     """
     Estimate dense templates with "average" or "median".
-    If "average" internaly estimate_templates_average() is used to saved memory/
+    If "average" internally estimate_templates_with_accumulator() is used to saved memory/
 
     Parameters
     ----------
@@ -742,7 +742,7 @@ def estimate_templates(
         job_name = "estimate_templates"
 
     if operator == "average":
-        templates_array = estimate_templates_average(
+        templates_array = estimate_templates_with_accumulator(
             recording, spikes, unit_ids, nbefore, nafter, return_scaled=return_scaled, job_name=job_name, **job_kwargs
         )
     elif operator == "median":
@@ -772,7 +772,7 @@ def estimate_templates(
     return templates_array
 
 
-def estimate_templates_average(
+def estimate_templates_with_accumulator(
     recording: BaseRecording,
     spikes: np.ndarray,
     unit_ids: list | np.ndarray,
@@ -780,13 +780,16 @@ def estimate_templates_average(
     nafter: int,
     return_scaled: bool = True,
     job_name=None,
+    return_std: bool = False,
     **job_kwargs,
 ):
     """
-    This is a fast implementation to compute average templates.
+    This is a fast implementation to compute template averages and standard deviations.
     This is useful to estimate sparsity without the need to allocate large waveform buffers.
-    The mechanism is pretty simple: it accumulates and sums spike waveforms in-place per worker and per unit.
-    Note that std, median and percentiles can't be computed with this method.
+    The mechanism is pretty simple: it accumulates and sums spike waveforms (and their squared)
+    in-place per worker and per unit.
+    Note that median and percentiles can't be computed with this method, because they don't support
+    the accumulator implementation.
 
     Parameters
     ----------
@@ -803,6 +806,8 @@ def estimate_templates_average(
         Number of samples to cut out after a spike
     return_scaled: bool, default: True
         If True, the traces are scaled before averaging
+    return_std: bool, default: False
+        If True, the standard deviation is also computed.
 
     Returns
     -------
@@ -820,8 +825,14 @@ def estimate_templates_average(
 
     shape = (num_worker, num_units, nbefore + nafter, num_chans)
     dtype = np.dtype("float32")
-    waveforms_per_worker, shm = make_shared_array(shape, dtype)
+    waveform_accumulator_per_worker, shm = make_shared_array(shape, dtype)
     shm_name = shm.name
+    if return_std:
+        waveform_squared_accumulator_per_worker, shm_squared = make_shared_array(shape, dtype)
+        shm_squared_name = shm_squared.name
+    else:
+        waveform_squared_accumulator_per_worker = None
+        shm_squared_name = None
 
     # trick to get the work_index given pid arrays
     lock = multiprocessing.Lock()
@@ -836,6 +847,7 @@ def estimate_templates_average(
         recording,
         spikes,
         shm_name,
+        shm_squared_name,
         shape,
         dtype,
         nbefore,
@@ -846,27 +858,46 @@ def estimate_templates_average(
     )
 
     if job_name is None:
-        job_name = "estimate_templates_average"
+        job_name = "estimate_templates_with_accumulator"
     processor = ChunkRecordingExecutor(recording, func, init_func, init_args, job_name=job_name, **job_kwargs)
     processor.run()
 
     # average
-    templates_array = np.sum(waveforms_per_worker, axis=0)
+    waveforms_sum = np.sum(waveform_accumulator_per_worker, axis=0)
+    template_means = waveforms_sum
     unit_indices, spike_count = np.unique(spikes["unit_index"], return_counts=True)
-    templates_array[unit_indices, :, :] /= spike_count[:, np.newaxis, np.newaxis]
-
-    # important : release the sharemem
-    del waveforms_per_worker
+    template_means[unit_indices, :, :] /= spike_count[:, np.newaxis, np.newaxis]
+
+    if return_std:
+        waveforms_squared_sum = np.sum(waveform_squared_accumulator_per_worker, axis=0)
+        # standard deviation
+        template_stds = np.zeros_like(template_means)
+        for i, (unit_index, count) in enumerate(zip(unit_indices, spike_count)):
+            residuals = (
+                waveforms_squared_sum[unit_index] - 2 * count * template_means[unit_index] * waveforms_sum[unit_index]
+            ) + count * template_means[unit_index] ** 2
+            template_stds[unit_index] = np.sqrt(residuals / count)
+        assert np.all(template_stds >= 0)
+        del waveform_squared_accumulator_per_worker
+        shm_squared.unlink()
+        shm_squared.close()
+
+    # important : release the sharedmem
+    del waveform_accumulator_per_worker
     shm.unlink()
     shm.close()
 
-    return templates_array
+    if return_std:
+        return template_means, template_stds
+    else:
+        return template_means
 
 
 def _init_worker_estimate_templates(
     recording,
     spikes,
     shm_name,
+    shm_squared_name,
     shape,
     dtype,
     nbefore,
@@ -886,9 +917,15 @@ def _init_worker_estimate_templates(
     import multiprocessing
 
     shm = SharedMemory(shm_name)
-    waveforms_per_worker = np.ndarray(shape=shape, dtype=dtype, buffer=shm.buf)
+    waveform_accumulator_per_worker = np.ndarray(shape=shape, dtype=dtype, buffer=shm.buf)
+
     worker_ctx["shm"] = shm
-    worker_ctx["waveforms_per_worker"] = waveforms_per_worker
+    worker_ctx["waveform_accumulator_per_worker"] = waveform_accumulator_per_worker
+    if shm_squared_name is not None:
+        shm_squared = SharedMemory(shm_squared_name)
+        waveform_squared_accumulator_per_worker = np.ndarray(shape=shape, dtype=dtype, buffer=shm_squared.buf)
+        worker_ctx["shm_squared"] = shm_squared
+        worker_ctx["waveform_squared_accumulator_per_worker"] = waveform_squared_accumulator_per_worker
 
     # prepare segment slices
     segment_slices = []
@@ -920,7 +957,8 @@ def _worker_estimate_templates(segment_index, start_frame, end_frame, worker_ctx
     spikes = worker_ctx["spikes"]
     nbefore = worker_ctx["nbefore"]
     nafter = worker_ctx["nafter"]
-    waveforms_per_worker = worker_ctx["waveforms_per_worker"]
+    waveform_accumulator_per_worker = worker_ctx["waveform_accumulator_per_worker"]
+    waveform_squared_accumulator_per_worker = worker_ctx.get("waveform_squared_accumulator_per_worker", None)
     worker_index = worker_ctx["worker_index"]
     return_scaled = worker_ctx["return_scaled"]
 
@@ -954,4 +992,6 @@ def _worker_estimate_templates(segment_index, start_frame, end_frame, worker_ctx
             unit_index = spikes[spike_index]["unit_index"]
             wf = traces[sample_index - start - nbefore : sample_index - start + nafter, :]
 
-            waveforms_per_worker[worker_index, unit_index, :, :] += wf
+            waveform_accumulator_per_worker[worker_index, unit_index, :, :] += wf
+            if waveform_squared_accumulator_per_worker is not None:
+                waveform_squared_accumulator_per_worker[worker_index, unit_index, :, :] += wf**2
diff --git a/src/spikeinterface/sorters/internal/tridesclous2.py b/src/spikeinterface/sorters/internal/tridesclous2.py
index 8a9dfc1ce..48c3c4cc1 100644
--- a/src/spikeinterface/sorters/internal/tridesclous2.py
+++ b/src/spikeinterface/sorters/internal/tridesclous2.py
@@ -8,7 +8,7 @@ from spikeinterface.core import (
     get_noise_levels,
     NumpySorting,
     get_channel_distances,
-    estimate_templates_average,
+    estimate_templates_with_accumulator,
     Templates,
     compute_sparsity,
 )
@@ -302,7 +302,7 @@ class Tridesclous2Sorter(ComponentsBasedSorter):
 
         nbefore = int(params["templates"]["ms_before"] * sampling_frequency / 1000.0)
         nafter = int(params["templates"]["ms_after"] * sampling_frequency / 1000.0)
-        templates_array = estimate_templates_average(
+        templates_array = estimate_templates_with_accumulator(
             recording,
             sorting_pre_peeler.to_spike_vector(),
             sorting_pre_peeler.unit_ids,
diff --git a/src/spikeinterface/sortingcomponents/clustering/position_and_features.py b/src/spikeinterface/sortingcomponents/clustering/position_and_features.py
index 3c58b5edb..d24af3c17 100644
--- a/src/spikeinterface/sortingcomponents/clustering/position_and_features.py
+++ b/src/spikeinterface/sortingcomponents/clustering/position_and_features.py
@@ -18,7 +18,7 @@ from spikeinterface.core import get_global_tmp_folder, get_noise_levels
 from spikeinterface.core.waveform_tools import extract_waveforms_to_buffers
 from .clustering_tools import remove_duplicates, remove_duplicates_via_matching, remove_duplicates_via_dip
 from spikeinterface.core import NumpySorting
-from spikeinterface.core import estimate_templates_average, Templates
+from spikeinterface.core import estimate_templates_with_accumulator, Templates
 from spikeinterface.sortingcomponents.features_from_peaks import compute_features_from_peaks
 
 
@@ -174,7 +174,7 @@ class PositionAndFeaturesClustering:
 
             nbefore = int(params["ms_before"] * fs / 1000.0)
             nafter = int(params["ms_after"] * fs / 1000.0)
-            templates_array = estimate_templates_average(
+            templates_array = estimate_templates_with_accumulator(
                 recording,
                 sorting.to_spike_vector(),
                 sorting.unit_ids,
diff --git a/src/spikeinterface/widgets/unit_waveforms.py b/src/spikeinterface/widgets/unit_waveforms.py
index ab415ae2f..a52a5847a 100644
--- a/src/spikeinterface/widgets/unit_waveforms.py
+++ b/src/spikeinterface/widgets/unit_waveforms.py
@@ -1,13 +1,13 @@
 from __future__ import annotations
 
 import numpy as np
+from warnings import warn
 
 from .base import BaseWidget, to_attr
 from .utils import get_unit_colors
 
 from ..core import ChannelSparsity, SortingAnalyzer
 from ..core.basesorting import BaseSorting
-from ..core.template_tools import _get_dense_templates_array
 
 
 class UnitWaveformsWidget(BaseWidget):
@@ -144,10 +144,18 @@ class UnitWaveformsWidget(BaseWidget):
                 assert isinstance(sparsity, ChannelSparsity), "'sparsity' should be a ChannelSparsity object!"
 
         # get templates
-        ext = sorting_analyzer.get_extension("templates")
-        assert ext is not None, "plot_waveforms() need extension 'templates'"
-        templates = ext.get_templates(unit_ids=unit_ids, operator="average")
+        self.templates_ext = sorting_analyzer.get_extension("templates") or sorting_analyzer.get_extension(
+            "fast_templates"
+        )
+        assert self.templates_ext is not None, "plot_waveforms() need extension 'templates' or 'fast_templates'"
+        templates = self.templates_ext.get_templates(unit_ids=unit_ids, operator="average")
 
+        if templates_percentile_shading is not None and sorting_analyzer.get_extension("templates") is None:
+            warn(
+                "templates_percentile_shading can only be used if the 'templates' extension is available. "
+                "Settimg templates_percentile_shading to None."
+            )
+            templates_percentile_shading = None
         templates_shading = self._get_template_shadings(sorting_analyzer, unit_ids, templates_percentile_shading)
 
         xvectors, y_scale, y_offset, delta_x = get_waveforms_scales(
@@ -418,11 +426,10 @@ class UnitWaveformsWidget(BaseWidget):
             display(self.widget)
 
     def _get_template_shadings(self, sorting_analyzer, unit_ids, templates_percentile_shading):
-        ext = sorting_analyzer.get_extension("templates")
-        templates = ext.get_templates(unit_ids=unit_ids, operator="average")
+        templates = self.templates_ext.get_templates(unit_ids=unit_ids, operator="average")
 
         if templates_percentile_shading is None:
-            templates_std = ext.get_templates(unit_ids=unit_ids, operator="std")
+            templates_std = self.templates_ext.get_templates(unit_ids=unit_ids, operator="std")
             templates_shading = [templates - templates_std, templates + templates_std]
         else:
             if isinstance(templates_percentile_shading, (int, float)):
@@ -436,7 +443,9 @@ class UnitWaveformsWidget(BaseWidget):
                 ), "'templates_percentile_shading' should be a have an even number of elements."
             templates_shading = []
             for percentile in templates_percentile_shading:
-                template_percentile = ext.get_templates(unit_ids=unit_ids, operator="percentile", percentile=percentile)
+                template_percentile = self.templates_ext.get_templates(
+                    unit_ids=unit_ids, operator="percentile", percentile=percentile
+                )
 
                 templates_shading.append(template_percentile)
         return templates_shading
@@ -453,8 +462,7 @@ class UnitWaveformsWidget(BaseWidget):
         do_shading = self.template_shading_button.value
 
         wf_ext = self.sorting_analyzer.get_extension("waveforms")
-        templates_ext = self.sorting_analyzer.get_extension("templates")
-        templates = templates_ext.get_templates(unit_ids=unit_ids, operator="average")
+        templates = self.templates_ext.get_templates(unit_ids=unit_ids, operator="average")
 
         # matplotlib next_data_plot dict update at each call
         data_plot = self.next_data_plot
diff --git a/src/spikeinterface/widgets/unit_waveforms_density_map.py b/src/spikeinterface/widgets/unit_waveforms_density_map.py
index 41c77f59f..a4704d14a 100644
--- a/src/spikeinterface/widgets/unit_waveforms_density_map.py
+++ b/src/spikeinterface/widgets/unit_waveforms_density_map.py
@@ -79,8 +79,8 @@ class UnitWaveformDensityMapWidget(BaseWidget):
         channel_inds = used_sparsity.unit_id_to_channel_indices
 
         # bins
-        # templates = we.get_all_templates(unit_ids=unit_ids)
-        templates = sorting_analyzer.get_extension("templates").get_templates(unit_ids=unit_ids)
+        ext_templates = sorting_analyzer.get_extension("templates") or sorting_analyzer.get_extension("fast_templates")
+        templates = ext_templates.get_templates(unit_ids=unit_ids)
         bin_min = np.min(templates) * 1.3
         bin_max = np.max(templates) * 1.3
         bin_size = (bin_max - bin_min) / 100


Test patch:
diff --git a/src/spikeinterface/core/tests/test_analyzer_extension_core.py b/src/spikeinterface/core/tests/test_analyzer_extension_core.py
index ba7e435c6..4b87f8240 100644
--- a/src/spikeinterface/core/tests/test_analyzer_extension_core.py
+++ b/src/spikeinterface/core/tests/test_analyzer_extension_core.py
@@ -211,6 +211,22 @@ def test_ComputeNoiseLevels(format, sparse):
     assert noise_levels.shape[0] == sorting_analyzer.channel_ids.size
 
 
+def test_ComputeFastTemplates_versus_ComputeTemplates():
+    sorting_analyzer = get_sorting_analyzer(format="memory", sparse=False)
+    sorting_analyzer.compute("random_spikes", max_spikes_per_unit=20, seed=2205)
+
+    # compute fast_templates and templates
+    sorting_analyzer.compute(["fast_templates", "waveforms", "templates"])
+
+    templates_fast = sorting_analyzer.get_extension("fast_templates").get_templates()
+    templates_fast_std = sorting_analyzer.get_extension("fast_templates").get_templates(operator="std")
+    templates = sorting_analyzer.get_extension("templates").get_templates()
+    templates_std = sorting_analyzer.get_extension("fast_templates").get_templates(operator="std")
+
+    np.testing.assert_almost_equal(templates_fast, templates)
+    np.testing.assert_almost_equal(templates_fast_std, templates_std)
+
+
 def test_get_children_dependencies():
     assert "fast_templates" in _extension_children["random_spikes"]
     assert "waveforms" in _extension_children["random_spikes"]
@@ -242,10 +258,9 @@ if __name__ == "__main__":
     # test_ComputeWaveforms(format="zarr", sparse=True)
     # test_ComputeWaveforms(format="zarr", sparse=False)
     # test_ComputeRandomSpikes(format="memory", sparse=True)
-
     # test_ComputeFastTemplates(format="memory", sparse=True)
-
     # test_ComputeNoiseLevels(format="memory", sparse=False)
 
+    test_ComputeFastTemplates_versus_ComputeTemplates()
     test_get_children_dependencies()
     test_delete_on_recompute()
diff --git a/src/spikeinterface/core/tests/test_waveform_tools.py b/src/spikeinterface/core/tests/test_waveform_tools.py
index a5473ae89..5c496fd97 100644
--- a/src/spikeinterface/core/tests/test_waveform_tools.py
+++ b/src/spikeinterface/core/tests/test_waveform_tools.py
@@ -11,7 +11,7 @@ from spikeinterface.core.waveform_tools import (
     extract_waveforms_to_single_buffer,
     split_waveforms_by_units,
     estimate_templates,
-    estimate_templates_average,
+    estimate_templates_with_accumulator,
 )
 
 
@@ -163,7 +163,7 @@ def test_waveform_tools():
     _check_all_wf_equal(list_wfs_sparse)
 
 
-def test_estimate_templates_average():
+def test_estimate_templates_with_accumulator():
     recording, sorting = get_dataset()
 
     ms_before = 1.0
@@ -178,7 +178,7 @@ def test_estimate_templates_average():
 
     job_kwargs = dict(n_jobs=2, progress_bar=True, chunk_duration="1s")
 
-    templates = estimate_templates_average(
+    templates = estimate_templates_with_accumulator(
         recording, spikes, sorting.unit_ids, nbefore, nafter, return_scaled=True, **job_kwargs
     )
     print(templates.shape)
@@ -231,5 +231,5 @@ def test_estimate_templates():
 
 if __name__ == "__main__":
     test_waveform_tools()
-    test_estimate_templates_average()
+    test_estimate_templates_with_accumulator()
     test_estimate_templates()