From 0d29a422d3adc32d0cc113479ec530c05e6a20a1 Mon Sep 17 00:00:00 2001
From: Zach McKenzie <92116279+zm711@users.noreply.github.com>
Date: Sat, 14 Oct 2023 09:00:34 -0400
Subject: [PATCH] add kwargs and make keep_unit_ids list of strings

---
 .../qualitymetrics/plot_4_curation.py         | 22 +++++++++++--------
 1 file changed, 13 insertions(+), 9 deletions(-)

diff --git a/examples/modules_gallery/qualitymetrics/plot_4_curation.py b/examples/modules_gallery/qualitymetrics/plot_4_curation.py
index c66f55f221..7f33e0bd8f 100644
--- a/examples/modules_gallery/qualitymetrics/plot_4_curation.py
+++ b/examples/modules_gallery/qualitymetrics/plot_4_curation.py
@@ -6,6 +6,8 @@
 quality metrics.
 
 """
+#############################################################################
+# Import the modules and/or functions necessary from spikeinterface
 
 import spikeinterface as si
 import spikeinterface.extractors as se
@@ -15,22 +17,21 @@
 
 
 ##############################################################################
-# First, let's download a simulated dataset
-#  from the repo 'https://gin.g-node.org/NeuralEnsemble/ephy_testing_data'
+# Let's download a simulated dataset
+# from the repo 'https://gin.g-node.org/NeuralEnsemble/ephy_testing_data'
 #
 # Let's imagine that the ground-truth sorting is in fact the output of a sorter.
-#
 
 local_path = si.download_dataset(remote_path='mearec/mearec_test_10s.h5')
-recording, sorting = se.read_mearec(local_path)
+recording, sorting = se.read_mearec(file_path=local_path)
 print(recording)
 print(sorting)
 
 ##############################################################################
-# First, we extract waveforms and compute their PC scores:
+# First, we extract waveforms (to be saved in the folder 'wfs_mearec') and
+# compute their PC scores:
 
-folder = 'wfs_mearec'
-we = si.extract_waveforms(recording, sorting, folder,
+we = si.extract_waveforms(recording, sorting, folder='wfs_mearec',
                           ms_before=1, ms_after=2., max_spikes_per_unit=500,
                           n_jobs=1, chunk_size=30000)
 print(we)
@@ -47,12 +48,15 @@
 ##############################################################################
 # We can now threshold each quality metric and select units based on some rules.
 #
-# The easiest and most intuitive way is to use boolean masking with dataframe:
+# The easiest and most intuitive way is to use boolean masking with a dataframe.
+#
+# Then create a list of unit ids that we want to keep
 
 keep_mask = (metrics['snr'] > 7.5) & (metrics['isi_violations_ratio'] < 0.2) & (metrics['nn_hit_rate'] > 0.90)
 print(keep_mask)
 
 keep_unit_ids = keep_mask[keep_mask].index.values
+keep_unit_ids = [unit_id for unit_id in keep_unit_ids]
 print(keep_unit_ids)
 
 ##############################################################################
@@ -61,4 +65,4 @@
 
 curated_sorting = sorting.select_units(keep_unit_ids)
 print(curated_sorting)
-se.NpzSortingExtractor.write_sorting(curated_sorting, 'curated_sorting.pnz')
+se.NpzSortingExtractor.write_sorting(sorting=curated_sorting, save_path='curated_sorting.npz')