dspk_sorter_ms5/
lib.rs

1//! MountainSort 5 (Chung et al., Neuron 2017; `flatironinstitute/mountainsort5`, Apache-2.0),
2//! ported from its source: locally exclusive detection ([`detect`]), masked snippets
3//! ([`snippets`]), PCA and the isosplit6 subdivision method ([`scheme1::subdivision`],
4//! [`isosplit6`]: this crate's port of `magland/isosplit6`), median templates and alignment ([`templates`],
5//! [`scheme1`]), then per-channel classifiers over the whole recording ([`scheme2`]), driven by
6//! [`runner`]. Stages, settings and our choices: the book's MountainSort 5 pages.
7
8pub mod detect;
9pub mod isosplit6;
10pub mod runner;
11pub mod scheme1;
12pub mod scheme2;
13pub mod snippets;
14pub mod templates;
15
16pub use detect::{DetectOptions, Detections, Detector};
17pub use runner::{run, Mountainsort5Result, MOUNTAINSORT5_SORTER};
18pub use scheme1::{cluster_events, cluster_events_with_progress, ClusteringParams, Events, Scheme1Output};
19pub use scheme2::{ClassifierParams, Classifiers};
20pub use snippets::{MaskedSnippets, SnippetRows};
21pub use scheme1::subdivision::{isosplit6_subdivision, isosplit6_subdivision_with_progress, SubdivisionOptions};
22
23use dspk_ops::linalg::TopComponentsOptions;
24use crate::isosplit6::IsosplitOptions;
25
26use dspk_synapse::provenance::{Attributed, Paper, Provenance, ProvenanceKind, UpstreamCode};
27
28/// One pass (scheme 1) or train on a stretch and classify everything (scheme 2, the default).
29#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
30pub enum Scheme {
31    One,
32    #[default]
33    Two,
34}
35
36/// Where scheme 2's training stretch comes from.
37#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
38pub enum TrainingSampling {
39    /// The first `training_duration_sec`.
40    Initial,
41    /// 10 s chunks spread evenly over the recording.
42    #[default]
43    Uniform,
44}
45
46/// MountainSort 5 settings: the package's defaults, and SpikeInterface's wrapper's where the
47/// package leaves them to the caller (book: MountainSort 5 parameters).
48#[derive(Debug, Clone, PartialEq)]
49pub struct Mountainsort5Config {
50    pub scheme: Scheme,
51    /// Common average reference first (not in the wrapper: off).
52    pub do_car: bool,
53    /// Band-pass (a high-pass without an upper edge): SpikeInterface's `bandpass_filter`,
54    /// Butterworth order 5, forward-backward.
55    pub do_bandpass: bool,
56    pub bandpass_low_hz: f64,
57    pub bandpass_high_hz: Option<f64>,
58    pub do_notch: bool,
59    pub notch_hz: f64,
60    pub notch_q: f64,
61    /// Global ZCA whitening (SpikeInterface's `whiten`).
62    pub do_whiten: bool,
63    /// The whitening is fitted on this many chunks of `whitening_chunk_ms`, evenly spaced.
64    pub whitening_chunks: usize,
65    pub whitening_chunk_ms: f64,
66    /// SpikeInterface's `eps` for data in µV (median square ≥ 1); values below
67    /// `dspk_ops::spatial::MIN_WHITENING_EPSILON` act as that floor.
68    pub whitening_epsilon: f32,
69    pub detect_threshold: f32,
70    /// −1 negative peaks, +1 positive, 0 both.
71    pub detect_sign: i8,
72    pub detect_time_radius_ms: f64,
73    /// Scheme 1's detection neighbourhood (µm; `None`: all channels).
74    pub scheme1_detect_channel_radius_um: Option<f32>,
75    /// Scheme 2's phase 1 (training) detection.
76    pub phase1_detect_channel_radius_um: Option<f32>,
77    pub phase1_detect_threshold: f32,
78    pub phase1_detect_time_radius_ms: f64,
79    /// Scheme 2's phase 2 (classification) detection neighbourhood.
80    pub detect_channel_radius_um: Option<f32>,
81    pub snippet_t1: usize,
82    pub snippet_t2: usize,
83    pub snippet_mask_radius_um: Option<f32>,
84    pub npca_per_channel: usize,
85    pub npca_per_subdivision: usize,
86    pub skip_alignment: bool,
87    pub isocut_threshold: f64,
88    pub min_cluster_size: usize,
89    pub k_init: usize,
90    pub max_iterations_per_pass: usize,
91    /// `None`: the whole recording.
92    pub training_duration_sec: Option<f64>,
93    pub training_sampling: TrainingSampling,
94    pub max_num_snippets_per_training_batch: usize,
95    /// `None`: `max(12, 3 · mask channels)`.
96    pub classifier_npca: Option<usize>,
97    /// Samples per window; `None`: `⌈10⁸ / channels⌉` (upstream's classification chunk).
98    pub classification_chunk_sec: Option<f64>,
99    /// Above this many features PCA is randomized (scikit-learn's `covariance_eigh` cap).
100    pub pca_exact_cap: usize,
101    /// The per-channel classifiers' PCA of at most this many snippets (fewer than their features)
102    /// is exact through the snippets' Gram matrix
103    /// (`dspk_ops::linalg::TopComponentsOptions::dual_max_rows`; upstream's is exact too, through
104    /// scikit-learn's covariance), else randomized. Measured on 2 s × 384 ch: classifiers 42 s →
105    /// 13 s. Not used for clustering: there the Gram matrix of thousands of events costs more than
106    /// the randomized solver. 0: never.
107    pub classifier_dual_max_rows: usize,
108    pub seed: u64,
109    /// Independent work on several threads (each with its own device stream): the per-channel
110    /// classifiers and the halves of the isosplit6 subdivision. Same result either way; default
111    /// `true`.
112    pub parallel: bool,
113}
114
115impl Default for Mountainsort5Config {
116    fn default() -> Self {
117        Self {
118            scheme: Scheme::Two,
119            do_car: false,
120            do_bandpass: true,
121            bandpass_low_hz: 300.0,
122            bandpass_high_hz: Some(6000.0),
123            do_notch: false,
124            notch_hz: 60.0,
125            notch_q: 30.0,
126            do_whiten: true,
127            whitening_chunks: 20,
128            whitening_chunk_ms: 500.0,
129            whitening_epsilon: 1e-16,
130            detect_threshold: 5.5,
131            detect_sign: -1,
132            detect_time_radius_ms: 0.5,
133            scheme1_detect_channel_radius_um: Some(150.0),
134            phase1_detect_channel_radius_um: Some(200.0),
135            phase1_detect_threshold: 5.5,
136            phase1_detect_time_radius_ms: 1.5,
137            detect_channel_radius_um: Some(50.0),
138            snippet_t1: 20,
139            snippet_t2: 20,
140            snippet_mask_radius_um: Some(250.0),
141            npca_per_channel: 3,
142            npca_per_subdivision: 10,
143            skip_alignment: false,
144            isocut_threshold: 2.0,
145            min_cluster_size: 10,
146            k_init: 200,
147            max_iterations_per_pass: 500,
148            training_duration_sec: Some(300.0),
149            training_sampling: TrainingSampling::Uniform,
150            max_num_snippets_per_training_batch: 200,
151            classifier_npca: None,
152            classification_chunk_sec: None,
153            pca_exact_cap: 8000,
154            classifier_dual_max_rows: 1024,
155            seed: 0,
156            parallel: true,
157        }
158    }
159}
160
161impl Mountainsort5Config {
162    pub fn pca(&self) -> TopComponentsOptions {
163        TopComponentsOptions { exact_cap: self.pca_exact_cap, seed: self.seed, ..Default::default() }
164    }
165
166    pub fn clustering(&self) -> ClusteringParams {
167        ClusteringParams {
168            npca_per_channel: self.npca_per_channel,
169            subdivision: SubdivisionOptions {
170                npca_per_subdivision: self.npca_per_subdivision,
171                isosplit: IsosplitOptions {
172                    isocut_threshold: self.isocut_threshold,
173                    min_cluster_size: self.min_cluster_size,
174                    k_init: self.k_init,
175                    max_iterations_per_pass: self.max_iterations_per_pass,
176                },
177                pca: self.pca(),
178                parallel: self.parallel,
179            },
180            skip_alignment: self.skip_alignment,
181            detect_sign: self.detect_sign,
182            pca: self.pca(),
183        }
184    }
185
186    pub fn classifier(&self) -> ClassifierParams {
187        ClassifierParams {
188            detect_threshold: self.detect_threshold,
189            detect_sign: self.detect_sign,
190            max_num_snippets_per_training_batch: self.max_num_snippets_per_training_batch,
191            classifier_npca: self.classifier_npca,
192            // Randomized (module docs of `scheme2`): only the subspace matters to the distances
193            pca: TopComponentsOptions { exact_cap: 0, dual_max_rows: self.classifier_dual_max_rows, ..self.pca() },
194            parallel: self.parallel,
195        }
196    }
197}
198
199/// Paper and code of MountainSort 5.
200pub fn mountainsort5_provenance() -> Provenance {
201    Provenance {
202        name: "MountainSort 5".into(),
203        kind: ProvenanceKind::PortedFromCode,
204        paper: Some(Paper {
205            title: "A fully automated approach to spike sorting".into(),
206            authors: vec!["Chung".into(), "Magland".into(), "Barnett".into(), "Tolosa".into(), "Tooker".into(), "Lee".into(), "Shah".into(), "Felix".into(), "Frank".into(), "Greengard".into()],
207            venue: "Neuron".into(),
208            year: 2017,
209            doi: "10.1016/j.neuron.2017.08.030".into(),
210            license: None,
211        }),
212        code: UpstreamCode {
213            url: "https://github.com/flatironinstitute/mountainsort5".into(),
214            license: Some("Apache-2.0".into()),
215            version: "v0.5.9 (isosplit6: magland/isosplit6)".into(),
216        },
217        artifacts: Vec::new(),
218        notes: "Ported stage by stage from the Apache-2.0 source: detection, masked snippets, PCA, isosplit6 subdivision, median templates and alignment (scheme 1), per-channel classifiers (scheme 2). Preprocessing as SpikeInterface's wrapper (band-pass 300-6000 Hz, global whitening). Our choices are listed in the book.".into(),
219    }
220}
221
222/// MountainSort 5 over a whole recording ([`run`]).
223#[derive(Debug, Clone, Default)]
224pub struct Mountainsort5 {
225    pub config: Mountainsort5Config,
226}
227
228impl Mountainsort5 {
229    pub fn new(config: Mountainsort5Config) -> Self {
230        Self { config }
231    }
232
233    /// Runs MountainSort 5 over `source` on `client`'s device runtime.
234    pub fn run(
235        &self,
236        client: &cubecl::prelude::Client,
237        source: &dyn dspk_core::RecordingSource,
238        probe: &dspk_synapse::io::probe::SensorLayout,
239    ) -> dspk_core::DspResult<Mountainsort5Result> {
240        run(client, source, probe, &self.config, &dspk_core::NoProgress)
241    }
242
243    pub fn run_with_progress(
244        &self,
245        client: &cubecl::prelude::Client,
246        source: &dyn dspk_core::RecordingSource,
247        probe: &dspk_synapse::io::probe::SensorLayout,
248        progress: &dyn dspk_core::ProgressSink,
249    ) -> dspk_core::DspResult<Mountainsort5Result> {
250        run(client, source, probe, &self.config, progress)
251    }
252}
253
254impl Attributed for Mountainsort5 {
255    fn provenance(&self) -> Provenance {
256        mountainsort5_provenance()
257    }
258}
dspk_sorter_ms5/lib.rs

dspk_sorter_ms5/
lib.rs
