diff --git a/src/spikeinterface/postprocessing/spike_locations.py b/src/spikeinterface/postprocessing/spike_locations.py index 4873e39bce..32334e6b3e 100644 --- a/src/spikeinterface/postprocessing/spike_locations.py +++ b/src/spikeinterface/postprocessing/spike_locations.py @@ -20,6 +20,9 @@ class ComputeSpikeLocations(BaseSpikeVectorExtension): The localization method to use method_kwargs : dict, default: dict() Other kwargs depending on the method. + seed : int or None, default: None + Seed for random number generator. Used by the `grid_convolution` method to + reproducibly subsample peaks when computing the prototype waveform. Returns ------- @@ -48,6 +51,7 @@ def _set_params( spike_retriever_kwargs=None, method="center_of_mass", method_kwargs={}, + seed=None, ): if spike_retriever_kwargs is None: spike_retriever_kwargs = {} @@ -57,6 +61,7 @@ def _set_params( spike_retriever_kwargs=spike_retriever_kwargs, method=method, method_kwargs=method_kwargs, + seed=seed, ) def _get_pipeline_nodes(self): @@ -77,6 +82,7 @@ def _get_pipeline_nodes(self): method_kwargs=self.params["method_kwargs"], ms_before=self.params["ms_before"], ms_after=self.params["ms_after"], + seed=self.params.get("seed"), ) return nodes diff --git a/src/spikeinterface/sortingcomponents/peak_localization/main.py b/src/spikeinterface/sortingcomponents/peak_localization/main.py index 71fb810eda..f49298e407 100644 --- a/src/spikeinterface/sortingcomponents/peak_localization/main.py +++ b/src/spikeinterface/sortingcomponents/peak_localization/main.py @@ -25,6 +25,7 @@ def get_localization_pipeline_nodes( ms_before=0.5, ms_after=0.5, job_kwargs=None, + seed=None, ): assert ( @@ -52,7 +53,7 @@ def get_localization_pipeline_nodes( method_kwargs = method_kwargs.copy() method_kwargs["prototype"], _, _ = get_prototype_and_waveforms_from_peaks( - recording, peaks=peak_source.peaks, ms_before=ms_before, ms_after=ms_after, job_kwargs=job_kwargs + recording, peaks=peak_source.peaks, ms_before=ms_before, ms_after=ms_after, job_kwargs=job_kwargs, seed=seed ) localization_nodes = method_class(recording, parents=[peak_source, extract_dense_waveforms], **method_kwargs) @@ -72,6 +73,7 @@ def localize_peaks( pipeline_kwargs=None, verbose=False, job_kwargs=None, + seed=None, **old_kwargs, ) -> np.ndarray: """Localize peak (spike) in 2D or 3D depending the method. @@ -102,6 +104,9 @@ def localize_peaks( If True, output is verbose job_kwargs : dict | None, default None A job kwargs dict. If None or empty dict, then the global one is used. + seed : int or None, default: None + Seed for random number generator. Used by `grid_convolution` to reproducibly + subsample peaks when computing the prototype waveform. {method_doc} @@ -150,6 +155,7 @@ def localize_peaks( ms_before=ms_before, ms_after=ms_after, job_kwargs=job_kwargs, + seed=seed, ) if pipeline_kwargs is None: diff --git a/src/spikeinterface/sortingcomponents/peak_selection.py b/src/spikeinterface/sortingcomponents/peak_selection.py index 1ddf3b53e4..ddec61c1ae 100644 --- a/src/spikeinterface/sortingcomponents/peak_selection.py +++ b/src/spikeinterface/sortingcomponents/peak_selection.py @@ -106,7 +106,6 @@ def select_peak_indices(peaks, method, seed, **method_kwargs): selected_indices = [] - seed = seed if seed else None rng = np.random.default_rng(seed=seed) if method == "uniform":