Refactor Quality Metrics Logic + blackformatting

Author: MilagrosMarin

element_array_ephys/ephys_no_curation.py

@@ -1277,11 +1277,11 @@ def make(self, key):
             we: si.WaveformExtractor = si.load_waveforms(
                 si_waveform_dir, with_recording=False
             )
-            unit_id_to_peak_channel_map: dict[int, np.ndarray] = (
-                si.ChannelSparsity.from_best_channels(
-                    we, 1, peak_sign="neg"
-                ).unit_id_to_channel_indices
-            )  # {unit: peak_channel_index}
+            unit_id_to_peak_channel_map: dict[
+                int, np.ndarray
+            ] = si.ChannelSparsity.from_best_channels(
+                we, 1, peak_sign="neg"
+            ).unit_id_to_channel_indices  # {unit: peak_channel_index}
 
             # reorder channel2electrode_map according to recording channel ids
             channel_indices = we.channel_ids_to_indices(we.channel_ids).tolist()
@@ -1315,6 +1315,7 @@ def yield_unit_waveforms():
                     ]
 
                     yield unit_peak_waveform, unit_electrode_waveforms
+
         else:  # read from kilosort outputs
             kilosort_dataset = kilosort.Kilosort(output_dir)
 
@@ -1546,9 +1547,14 @@ def make(self, key):
 
         metrics_df.rename(
             columns={
-                "isi_viol": "isi_violation",
-                "num_viol": "number_violation",
-                "contam_rate": "contamination_rate",
+                "isi_violations_ratio": "isi_violation",
+                "isi_violations_count": "number_violation",
+                "silhouette": "silhouette_score",
+                "rp_contamination": "contamination_rate",
+                "drift_ptp": "max_drift",
+                "drift_mad": "cumulative_drift",
+                "half_width": "halfwidth",
+                "peak_trough_ratio": "pt_ratio",
             },
             inplace=True,
         )

element_array_ephys/spike_sorting/si_spike_sorting.py

@@ -134,15 +134,19 @@ def make(self, key):
         # Add probe information to recording object
         electrodes_df = (
             (
-                ephys.EphysRecording.Channel * probe.ElectrodeConfig.Electrode * probe.ProbeType.Electrode
+                ephys.EphysRecording.Channel
+                * probe.ElectrodeConfig.Electrode
+                * probe.ProbeType.Electrode
                 & key
             )
             .fetch(format="frame")
             .reset_index()
         )
 
         # Create SI probe object
-        si_probe = readers.probe_geometry.to_probeinterface(electrodes_df[["electrode", "x_coord", "y_coord", "shank"]])
+        si_probe = readers.probe_geometry.to_probeinterface(
+            electrodes_df[["electrode", "x_coord", "y_coord", "shank"]]
+        )
         si_probe.set_device_channel_indices(electrodes_df["channel_idx"].values)
         si_recording.set_probe(probe=si_probe, in_place=True)
 
@@ -184,7 +188,9 @@ def make(self, key):
         output_dir = find_full_path(ephys.get_ephys_root_data_dir(), output_dir)
         sorter_name = clustering_method.replace(".", "_")
         recording_file = output_dir / sorter_name / "recording" / "si_recording.pkl"
-        si_recording: si.BaseRecording = si.load_extractor(recording_file, base_folder=output_dir)
+        si_recording: si.BaseRecording = si.load_extractor(
+            recording_file, base_folder=output_dir
+        )
 
         # Run sorting
         # Sorting performed in a dedicated docker environment if the sorter is not built in the spikeinterface package.
@@ -241,8 +247,12 @@ def make(self, key):
         recording_file = output_dir / sorter_name / "recording" / "si_recording.pkl"
         sorting_file = output_dir / sorter_name / "spike_sorting" / "si_sorting.pkl"
 
-        si_recording: si.BaseRecording = si.load_extractor(recording_file, base_folder=output_dir)
-        si_sorting: si.sorters.BaseSorter = si.load_extractor(sorting_file, base_folder=output_dir)
+        si_recording: si.BaseRecording = si.load_extractor(
+            recording_file, base_folder=output_dir
+        )
+        si_sorting: si.sorters.BaseSorter = si.load_extractor(
+            sorting_file, base_folder=output_dir
+        )
 
         # Extract waveforms
         we: si.WaveformExtractor = si.extract_waveforms(
@@ -257,37 +267,61 @@ def make(self, key):
             **params.get("SI_JOB_KWARGS", {"n_jobs": -1, "chunk_size": 30000}),
         )
 
-        # Calculate QC Metrics
-        metrics: pd.DataFrame = si.qualitymetrics.compute_quality_metrics(
-            we,
-            metric_names=[
-                "firing_rate",
-                "snr",
-                "presence_ratio",
-                "isi_violation",
-                "num_spikes",
-                "amplitude_cutoff",
-                "amplitude_median",
-                "sliding_rp_violation",
-                "rp_violation",
-                "drift",
-            ],
-        )
-        # Add PCA based metrics. These will be added to the metrics dataframe above.
+        # Calculate Cluster and Waveform Metrics
+
+        # To provide waveform_principal_component
         _ = si.postprocessing.compute_principal_components(
             waveform_extractor=we, **params.get("SI_QUALITY_METRICS_PARAMS", None)
         )
-        metrics = si.qualitymetrics.compute_quality_metrics(waveform_extractor=we)
+
+        # To estimate the location of each spike in the sorting output.
+        # The drift metrics require the `spike_locations` waveform extension.
+        _ = si.postprocessing.compute_spike_locations(waveform_extractor=we)
+
+        # The `sd_ratio` metric requires the `spike_amplitudes` waveform extension.
+        # It is highly recommended before calculating amplitude-based quality metrics.
+        _ = si.postprocessing.compute_spike_amplitudes(waveform_extractor=we)
+
+        # To compute correlograms for spike trains.
+        _ = si.postprocessing.compute_correlograms(we)
+
+        metric_names = si.qualitymetrics.get_quality_metric_list()
+        metric_names.extend(si.qualitymetrics.get_quality_pca_metric_list())
+
+        # To compute commonly used cluster quality metrics.
+        qc_metrics = si.qualitymetrics.compute_quality_metrics(
+            waveform_extractor=we,
+            metric_names=metric_names,
+        )
+
+        # To compute commonly used waveform/template metrics.
+        template_metric_names = si.postprocessing.get_template_metric_names()
+        template_metric_names.extend(["amplitude", "duration"])
+
+        template_metrics = si.postprocessing.compute_template_metrics(
+            waveform_extractor=we,
+            include_multi_channel_metrics=True,
+            metric_names=template_metric_names,
+        )
+
+        # Save the output (metrics.csv to the output dir)
+        metrics = pd.DataFrame()
+        metrics = pd.concat([qc_metrics, template_metrics], axis=1)
 
         # Save the output (metrics.csv to the output dir)
         metrics_output_dir = output_dir / sorter_name / "metrics"
         metrics_output_dir.mkdir(parents=True, exist_ok=True)
         metrics.to_csv(metrics_output_dir / "metrics.csv")
 
         # Save to phy format
-        si.exporters.export_to_phy(waveform_extractor=we, output_folder=output_dir / sorter_name / "phy")
+        si.exporters.export_to_phy(
+            waveform_extractor=we, output_folder=output_dir / sorter_name / "phy"
+        )
         # Generate spike interface report
-        si.exporters.export_report(waveform_extractor=we, output_folder=output_dir / sorter_name / "spikeinterface_report")
+        si.exporters.export_report(
+            waveform_extractor=we,
+            output_folder=output_dir / sorter_name / "spikeinterface_report",
+        )
 
         self.insert1(
             {