Allow export_to_phy to work with fast_templates

src/spikeinterface/core/analyzer_extension_core.py

@@ -397,10 +397,10 @@ def get_templates(self, unit_ids=None, operator="average", percentile=None, save
         ----------
         unit_ids: list or None
             Unit ids to retrieve waveforms for
-        mode: "average" | "median" | "std" | "percentile", default: "average"
-            The mode to compute the templates
+        operator: "average" | "median" | "std" | "percentile", default: "average"
+            The operator to compute the templates
         percentile: float, default: None
-            Percentile to use for mode="percentile"
+            Percentile to use for operator="percentile"
         save: bool, default True
             In case, the operator is not computed yet it can be saved to folder or zarr.
 
@@ -520,6 +520,49 @@ def _select_extension_data(self, unit_ids):
 
         return new_data
 
+    def get_templates(self, unit_ids=None):
+        """
+        Return average templates for multiple units.
+
+        Parameters
+        ----------
+        unit_ids: list or None, default: None
+            Unit ids to retrieve waveforms for
+
+        Returns
+        -------
+        templates: np.array
+            The returned templates (num_units, num_samples, num_channels)
+        """
+
+        templates = self.data["average"]
+
+        if unit_ids is not None:
+            unit_indices = self.sorting_analyzer.sorting.ids_to_indices(unit_ids)
+            templates = templates[unit_indices, :, :]
+
+        return np.array(templates)
+
+    def get_unit_template(self, unit_id):
+        """
+        Return average template for a single unit.
+
+        Parameters
+        ----------
+        unit_id: str | int
+            Unit id to retrieve waveforms for
+
+        Returns
+        -------
+        template: np.array
+            The returned template (num_samples, num_channels)
+        """
+
+        templates = self.data["average"]
+        unit_index = self.sorting_analyzer.sorting.id_to_index(unit_id)
+
+        return np.array(templates[unit_index, :, :])
+
 
 compute_fast_templates = ComputeFastTemplates.function_factory()
 register_result_extension(ComputeFastTemplates)

src/spikeinterface/exporters/to_phy.py

@@ -181,9 +181,19 @@ def export_to_phy(
     # export templates/templates_ind/similar_templates
     # shape (num_units, num_samples, max_num_channels)
     templates_ext = sorting_analyzer.get_extension("templates")
-    templates_ext is not None, "export_to_phy need SortingAnalyzer with extension 'templates'"
+    if templates_ext is None:
+        templates_ext = sorting_analyzer.get_extension("fast_templates")
+        if templates_ext is not None and template_mode != "average":
+            raise ValueError(
+                "export_to_phy with SortingAnalyzer with extension 'fast_templates' can only work with template_mode='average'"
+            )
+        dense_templates = templates_ext.get_templates(unit_ids=unit_ids)
+    else:
+        dense_templates = templates_ext.get_templates(unit_ids=unit_ids, operator=template_mode)
+    assert (
+        templates_ext is not None
+    ), "export_to_phy requires SortingAnalyzer with either extension 'templates' or 'fast_templates'"
     max_num_channels = max(len(chan_inds) for chan_inds in sparse_dict.values())
-    dense_templates = templates_ext.get_templates(unit_ids=unit_ids, operator=template_mode)
     num_samples = dense_templates.shape[1]
     templates = np.zeros((len(unit_ids), num_samples, max_num_channels), dtype="float64")
     # here we pad template inds with -1 if len of sparse channels is unequal