From 437d8b4252e48149441b0cb637f2968d9a20c149 Mon Sep 17 00:00:00 2001
From: Pierre Yger <pierre.yger@gmail.com>
Date: Wed, 8 Jan 2025 15:56:18 +0100
Subject: [PATCH] Docstrings and cosmetics

---
 .../sorters/internal/spyking_circus2.py       |  8 ++---
 .../sortingcomponents/peak_detection.py       |  6 +++-
 src/spikeinterface/sortingcomponents/tools.py | 29 ++++++++++++++++++-
 3 files changed, 36 insertions(+), 7 deletions(-)

diff --git a/src/spikeinterface/sorters/internal/spyking_circus2.py b/src/spikeinterface/sorters/internal/spyking_circus2.py
index 47e5fc191f..5ccae819c4 100644
--- a/src/spikeinterface/sorters/internal/spyking_circus2.py
+++ b/src/spikeinterface/sorters/internal/spyking_circus2.py
@@ -27,7 +27,7 @@ class Spykingcircus2Sorter(ComponentsBasedSorter):
         "general": {"ms_before": 2, "ms_after": 2, "radius_um": 50},
         "sparsity": {"method": "snr", "amplitude_mode": "peak_to_peak", "threshold": 0.25},
         "filtering": {"freq_min": 150, "freq_max": 7000, "ftype": "bessel", "filter_order": 2},
-        "whitening": {"mode": "local", "regularize": False, "radius_um": 100},
+        "whitening": {"mode": "local", "regularize": False},
         "detection": {"peak_sign": "neg", "detect_threshold": 4},
         "selection": {
             "method": "uniform",
@@ -52,7 +52,7 @@ class Spykingcircus2Sorter(ComponentsBasedSorter):
         "matched_filtering": True,
         "cache_preprocessing": {"mode": "memory", "memory_limit": 0.5, "delete_cache": True},
         "multi_units_only": False,
-        "job_kwargs": {"n_jobs": 0.8},
+        "job_kwargs": {"n_jobs": 0.5},
         "seed": 42,
         "debug": False,
     }
@@ -169,7 +169,6 @@ def _run_from_folder(cls, sorter_output_folder, params, verbose):
         if num_channels == 1:
             whitening_kwargs["regularize"] = False
         if whitening_kwargs["regularize"]:
-            n_jobs = job_kwargs["n_jobs"]
             whitening_kwargs["regularize_kwargs"] = {"method": "LedoitWolf"}
 
         recording_w = whiten(recording_f, **whitening_kwargs)
@@ -204,8 +203,7 @@ def _run_from_folder(cls, sorter_output_folder, params, verbose):
 
         if params["matched_filtering"]:
             prototype, waveforms = get_prototype_and_waveforms(
-                recording_w, 
-                n_peaks=5000, 
+                recording_w,
                 ms_before=ms_before, 
                 ms_after=ms_after, 
                 seed=params["seed"],
diff --git a/src/spikeinterface/sortingcomponents/peak_detection.py b/src/spikeinterface/sortingcomponents/peak_detection.py
index 8a22b35152..9224a1dc6b 100644
--- a/src/spikeinterface/sortingcomponents/peak_detection.py
+++ b/src/spikeinterface/sortingcomponents/peak_detection.py
@@ -118,7 +118,11 @@ def detect_peaks(
         squeeze_output = True
     else:
         squeeze_output = False
-        job_name += f"  + {len(pipeline_nodes)} nodes"
+        if len(pipeline_nodes) == 1:
+            plural = ''
+        else:
+            plural = 's'
+        job_name += f" + {len(pipeline_nodes)} node{plural}"
 
         # because node are modified inplace (insert parent) they need to copy incase
         # the same pipeline is run several times
diff --git a/src/spikeinterface/sortingcomponents/tools.py b/src/spikeinterface/sortingcomponents/tools.py
index a5526871f1..439e71f0fa 100644
--- a/src/spikeinterface/sortingcomponents/tools.py
+++ b/src/spikeinterface/sortingcomponents/tools.py
@@ -70,7 +70,34 @@ def extract_waveform_at_max_channel(rec, peaks, ms_before=0.5, ms_after=1.5, **j
 
 def get_prototype_and_waveforms(recording, n_peaks=5000, peaks=None, ms_before=0.5, ms_after=0.5, seed=None, return_waveforms=False, **all_kwargs):
     """
-    Helper function to extract a prototype waveform from a peak list or from a peak detection
+    Function to extract a prototype waveform from a peak list or from a peak detection. Note that in case
+    of a peak detection, the detection stops as soon as n_peaks are detected.
+
+    Parameters
+    ----------
+    recording : Recording
+        The recording object containing the data.
+    n_peaks : int, optional
+        Number of peaks to consider, by default 5000.
+    peaks : numpy.array, optional
+        Array of peaks, if None, peaks will be detected, by default None.
+    ms_before : float, optional
+        Time in milliseconds before the peak to extract the waveform, by default 0.5.
+    ms_after : float, optional
+        Time in milliseconds after the peak to extract the waveform, by default 0.5.
+    seed : int or None, optional
+        Seed for random number generator, by default None.
+    return_waveforms : bool, optional
+        Whether to return the waveforms along with the prototype, by default False.
+    **all_kwargs : dict
+        Additional keyword arguments for peak detection and job kwargs.
+
+    Returns
+    -------
+    prototype : numpy.array
+        The prototype waveform.
+    waveforms : numpy.array, optional
+        The extracted waveforms, returned if return_waveforms is True.
     """
     
     seed = seed if seed else None