Timing the computations of quality metrics and some optimizations

src/spikeinterface/core/analyzer_extension_core.py

@@ -1105,6 +1105,9 @@ def _compute_metrics(
         -------
         metrics : pd.DataFrame
             DataFrame containing the computed metrics for each unit.
+        run_times : dict
+            Dictionary containing the computation time for each metric.
+
         """
         import pandas as pd
 
@@ -1121,11 +1124,17 @@ def _compute_metrics(
 
         metrics = pd.DataFrame(index=unit_ids, columns=list(column_names_dtypes.keys()))
 
+        run_times = {}
+
         for metric_name in metric_names:
             metric = [m for m in self.metric_list if m.metric_name == metric_name][0]
             column_names = list(metric.metric_columns.keys())
+            import time
+
+            t_start = time.perf_counter()
             try:
                 metric_params = self.params["metric_params"].get(metric_name, {})
+
                 res = metric.compute(
                     sorting_analyzer,
                     unit_ids=unit_ids,
@@ -1139,6 +1148,8 @@ def _compute_metrics(
                     res = {unit_id: np.nan for unit_id in unit_ids}
                 else:
                     res = namedtuple("MetricResult", column_names)(*([np.nan] * len(column_names)))
+            t_end = time.perf_counter()
+            run_times[metric_name] = t_end - t_start
 
             # res is a namedtuple with several dictionary entries (one per column)
             if isinstance(res, dict):
@@ -1150,7 +1161,7 @@ def _compute_metrics(
 
         metrics = self._cast_metrics(metrics)
 
-        return metrics
+        return metrics, run_times
 
     def _run(self, **job_kwargs):
 
@@ -1161,7 +1172,7 @@ def _run(self, **job_kwargs):
         job_kwargs = fix_job_kwargs(job_kwargs)
 
         # compute the metrics which have been specified by the user
-        computed_metrics = self._compute_metrics(
+        computed_metrics, run_times = self._compute_metrics(
             sorting_analyzer=self.sorting_analyzer, unit_ids=None, metric_names=metrics_to_compute, **job_kwargs
         )
 
@@ -1189,6 +1200,7 @@ def _run(self, **job_kwargs):
                 computed_metrics[column_name] = extension.data["metrics"][column_name]
 
         self.data["metrics"] = computed_metrics
+        self.data["runtime_s"] = run_times
 
     def _get_data(self):
         # convert to correct dtype
@@ -1265,7 +1277,7 @@ def _merge_extension_data(
         metrics = pd.DataFrame(index=all_unit_ids, columns=old_metrics.columns)
 
         metrics.loc[not_new_ids, :] = old_metrics.loc[not_new_ids, :]
-        metrics.loc[new_unit_ids, :] = self._compute_metrics(
+        metrics.loc[new_unit_ids, :], _ = self._compute_metrics(
             sorting_analyzer=new_sorting_analyzer, unit_ids=new_unit_ids, metric_names=metric_names, **job_kwargs
         )
         metrics = self._cast_metrics(metrics)
@@ -1309,7 +1321,7 @@ def _split_extension_data(
         metrics = pd.DataFrame(index=all_unit_ids, columns=old_metrics.columns)
 
         metrics.loc[not_new_ids, :] = old_metrics.loc[not_new_ids, :]
-        metrics.loc[new_unit_ids_f, :] = self._compute_metrics(
+        metrics.loc[new_unit_ids_f, :], _ = self._compute_metrics(
             sorting_analyzer=new_sorting_analyzer, unit_ids=new_unit_ids_f, metric_names=metric_names, **job_kwargs
         )
         metrics = self._cast_metrics(metrics)