dspk_synapse/sorting/
cbss.rs

1//! Convolutive Blind Source Separation (cBSS) for High-Density Surface EMG (`cbss.rs`).
2//!
3//! Decomposes multi-channel HD-EMG grids (`[channels, samples]`) into constituent Motor Unit
4//! Pulse Trains (IPTs) using the Holobar & Zazula (CKC) / Negro et al. (2016) pipeline:
5//! 1. **Temporal Extension**: Extends $M$ channels by $L$ delayed replicas into an $M \cdot L \times S$
6//!    block-Hankel observation matrix $\tilde{\mathbf{X}}$.
7//! 2. **Spatial-Temporal ZCA Whitening & FastICA**: Orthogonalizes extended observations and extracts
8//!    sparse, super-Gaussian source Innovation Pulse Trains $s_k(t)$ via [`dspk_ops::linalg::FastIcaModel`].
9//! 3. **Peak Picking & Discharge Quality**: Extracts motor unit discharge timestamps from squared IPTs
10//!    $s_k^2(t)$, deduplicates duplicate delayed replicas, and computes the Coefficient of Variation
11//!    of Inter-Spike Intervals ($\text{CoV}_{\text{ISI}} = \sigma_{\text{ISI}} / \mu_{\text{ISI}}$) and
12//!    Pulse-to-Noise Ratio ($\text{PNR}$ in dB).
13
14use cubecl::prelude::Client;
15use dspk_ops::linalg::{FastIcaModel, IcaContrast};
16use serde::{Deserialize, Serialize};
17
18/// FastICA convergence tolerance on the change of each unmixing vector.
19pub const ICA_TOLERANCE: f32 = 1e-4;
20
21/// Decomposed single Motor Unit (MU) pulse train and discharge quality metrics.
22#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
23pub struct MotorUnitPulseTrain {
24    pub unit_id: usize,
25    /// Discharge sample indices (`0..samples`).
26    pub spike_samples: Vec<u64>,
27    /// Innovation Pulse Train (IPT) $s_k(t)$ of length `samples`.
28    pub ipt: Vec<f32>,
29    /// Coefficient of Variation of Inter-Spike Intervals ($\sigma_{\text{ISI}} / \mu_{\text{ISI}}$).
30    /// Physiological motor units typically have $\text{CoV}_{\text{ISI}} < 0.35$.
31    pub cov_isi: f32,
32    /// Pulse-to-Noise Ratio (PNR) in dB ($10 \log_{10}(\mathbb{E}[s^2(t_k)] / \mathbb{E}[s^2(t \notin t_k)])$).
33    pub pnr_db: f32,
34}
35
36/// Convolutive Blind Source Separation (cBSS) decomposer for HD-EMG motor unit identification.
37#[derive(Debug, Clone)]
38pub struct ConvolutiveBssDecomposer {
39    /// Extension factor $L$ (number of delayed replicas per channel, typically $4..16$).
40    pub extension_factor: usize,
41    /// Number of candidate ICA sources to extract.
42    pub num_sources: usize,
43    /// Minimum refractory distance between discharges of the same motor unit (in ms, typically $20\,\text{ms}$ for MUs).
44    pub refractory_ms: f64,
45    /// Threshold factor (in standard deviations of $|s(t)|$) for discharge peak detection.
46    pub peak_threshold_sigma: f32,
47    /// Maximum FastICA iterations.
48    pub max_iterations: usize,
49}
50
51impl Default for ConvolutiveBssDecomposer {
52    fn default() -> Self {
53        Self {
54            extension_factor: 6,
55            num_sources: 8,
56            refractory_ms: 20.0,
57            peak_threshold_sigma: 3.5,
58            max_iterations: 80,
59        }
60    }
61}
62
63impl ConvolutiveBssDecomposer {
64    pub fn new(extension_factor: usize, num_sources: usize, refractory_ms: f64) -> Self {
65        Self {
66            extension_factor: extension_factor.max(1),
67            num_sources: num_sources.max(1),
68            refractory_ms: refractory_ms.max(2.0),
69            ..Self::default()
70        }
71    }
72
73    /// Decomposes a multi-channel HD-EMG recording (`[channels, samples]`) into deduplicated Motor
74    /// Unit Pulse Trains, fitting ICA on `client`.
75    pub fn decompose(
76        &self,
77        client: &Client,
78        data: &[f32],
79        channels: usize,
80        samples: usize,
81        sample_rate_hz: f64,
82    ) -> Vec<MotorUnitPulseTrain> {
83        assert_eq!(data.len(), channels * samples);
84        if channels == 0 || samples < 16 {
85            return Vec::new();
86        }
87
88        let l_ext = self.extension_factor.clamp(1, samples / 4);
89        let ext_channels = channels * l_ext;
90        let k_sources = self.num_sources.clamp(1, ext_channels);
91
92        // 1. Construct temporally extended observation matrix X_tilde of shape [channels * L, samples]
93        let mut extended = vec![0.0f32; ext_channels * samples];
94        for c in 0..channels {
95            let src_row = &data[c * samples..(c + 1) * samples];
96            for lag in 0..l_ext {
97                let dst_row = &mut extended[(c * l_ext + lag) * samples..(c * l_ext + lag + 1) * samples];
98                dst_row[lag..].copy_from_slice(&src_row[..samples - lag]);
99            }
100        }
101
102        // 2. Fit FastICA with Cube (kurtosis) contrast over the extended observations
103        let ica = FastIcaModel::fit::<f32>(
104            client,
105            &extended,
106            ext_channels,
107            samples,
108            k_sources,
109            IcaContrast::Cube,
110            self.max_iterations,
111            ICA_TOLERANCE,
112        );
113        let sources = ica.transform_cpu(&extended, ext_channels, samples);
114
115        let ref_samples = ((sample_rate_hz * self.refractory_ms * 1e-3).round() as usize).max(2);
116        let mut candidates: Vec<MotorUnitPulseTrain> = Vec::new();
117
118        for s_idx in 0..ica.num_components {
119            let mut ipt = sources[s_idx * samples..(s_idx + 1) * samples].to_vec();
120            // Orient sign so dominant skewness/extrema are positive
121            let max_pos = ipt.iter().copied().fold(0.0f32, f32::max);
122            let max_neg = ipt.iter().copied().fold(0.0f32, f32::min).abs();
123            if max_neg > max_pos {
124                for v in &mut ipt {
125                    *v = -*v;
126                }
127            }
128
129            let rms = (ipt.iter().map(|v| v * v).sum::<f32>() / (samples as f32)).sqrt().max(1e-8);
130            let thresh = self.peak_threshold_sigma * rms;
131
132            let mut spikes = Vec::new();
133            let mut last_t: Option<usize> = None;
134            for t in 1..(samples - 1) {
135                let v = ipt[t];
136                if v > thresh && v >= ipt[t - 1] && v > ipt[t + 1] {
137                    if last_t.map_or(true, |prev| t > prev + ref_samples) {
138                        spikes.push(t as u64);
139                        last_t = Some(t);
140                    } else if let Some(prev) = last_t {
141                        if v > ipt[prev] {
142                            *spikes.last_mut().unwrap() = t as u64;
143                            last_t = Some(t);
144                        }
145                    }
146                }
147            }
148
149            if spikes.len() < 3 {
150                continue;
151            }
152
153            let cov_isi = compute_cov_isi(&spikes);
154            let pnr_db = compute_pnr_db(&ipt, &spikes);
155
156            candidates.push(MotorUnitPulseTrain {
157                unit_id: candidates.len(),
158                spike_samples: spikes,
159                ipt,
160                cov_isi,
161                pnr_db,
162            });
163        }
164
165        // 3. Sort by PNR (highest quality first) and deduplicate delayed replica trains (>50% spike coincidence within +/- L samples)
166        candidates.sort_by(|a, b| b.pnr_db.partial_cmp(&a.pnr_db).unwrap_or(std::cmp::Ordering::Equal));
167        let mut unique_units: Vec<MotorUnitPulseTrain> = Vec::new();
168        let tol_samples = (l_ext as u64) + 2;
169
170        for mut cand in candidates {
171            let is_dup = unique_units.iter().any(|u| {
172                spike_train_coincidence(&u.spike_samples, &cand.spike_samples, tol_samples) > 0.5
173            });
174            if !is_dup {
175                cand.unit_id = unique_units.len();
176                unique_units.push(cand);
177            }
178        }
179
180        unique_units
181    }
182}
183
184fn compute_cov_isi(spikes: &[u64]) -> f32 {
185    if spikes.len() < 3 {
186        return f32::NAN;
187    }
188    let isis: Vec<f32> = spikes.windows(2).map(|w| (w[1] - w[0]) as f32).collect();
189    let n = isis.len() as f32;
190    let mean = isis.iter().sum::<f32>() / n;
191    if mean <= 1e-6 {
192        return f32::NAN;
193    }
194    let var = isis.iter().map(|&d| (d - mean) * (d - mean)).sum::<f32>() / n;
195    var.sqrt() / mean
196}
197
198fn compute_pnr_db(ipt: &[f32], spikes: &[u64]) -> f32 {
199    if ipt.is_empty() || spikes.is_empty() {
200        return 0.0;
201    }
202    let mut is_peak = vec![false; ipt.len()];
203    let mut peak_sum = 0.0f32;
204    for &t in spikes {
205        let idx = t as usize;
206        if idx < ipt.len() {
207            peak_sum += ipt[idx] * ipt[idx];
208            for w in idx.saturating_sub(2)..=(idx + 2).min(ipt.len() - 1) {
209                is_peak[w] = true;
210            }
211        }
212    }
213    let mean_peak_sq = peak_sum / (spikes.len() as f32);
214
215    let mut noise_sum = 0.0f32;
216    let mut noise_count = 0usize;
217    for (i, &v) in ipt.iter().enumerate() {
218        if !is_peak[i] {
219            noise_sum += v * v;
220            noise_count += 1;
221        }
222    }
223    let mean_noise_sq = (noise_sum / (noise_count.max(1) as f32)).max(1e-12);
224    10.0 * (mean_peak_sq / mean_noise_sq).max(1e-6).log10()
225}
226
227fn spike_train_coincidence(a: &[u64], b: &[u64], tol: u64) -> f32 {
228    if a.is_empty() || b.is_empty() {
229        return 0.0;
230    }
231    // Check best lag within [-tol, tol]
232    let max_shift = tol as i64;
233    let mut best_matches = 0usize;
234    for shift in -max_shift..=max_shift {
235        let mut matches = 0usize;
236        let mut j = 0usize;
237        for &ta in a {
238            let shifted_a = (ta as i64 + shift).max(0) as u64;
239            while j < b.len() && b[j] + 2 < shifted_a {
240                j += 1;
241            }
242            if j < b.len() && b[j].abs_diff(shifted_a) <= 2 {
243                matches += 1;
244            }
245        }
246        if matches > best_matches {
247            best_matches = matches;
248        }
249    }
250    (best_matches as f32) / (a.len().min(b.len()) as f32)
251}
252
253#[cfg(test)]
254mod tests {
255    use super::*;
256
257    #[test]
258    fn test_cbss_decomposer_recovers_motor_unit_pulse_trains() {
259        let channels = 4;
260        let samples = 2000;
261        let fs = 2000.0;
262
263        // Two motor units firing regularly with different periods (80 samples = 25 Hz, 115 samples = 17.4 Hz)
264        let mu1_times: Vec<usize> = (60..samples - 40).step_by(80).collect();
265        let mu2_times: Vec<usize> = (90..samples - 40).step_by(115).collect();
266
267        let mut data = vec![0.0f32; channels * samples];
268        let w1 = [1.0f32, 0.7, 0.2, 0.05];
269        let w2 = [0.05f32, 0.25, 0.8, 1.0];
270
271        for &t0 in &mu1_times {
272            for c in 0..channels {
273                // Convolutive delayed biphasic MUAP
274                let center = t0 + c;
275                data[c * samples + center] += 25.0 * w1[c];
276                data[c * samples + center + 1] -= 12.0 * w1[c];
277            }
278        }
279        for &t0 in &mu2_times {
280            for c in 0..channels {
281                let center = t0 + (channels - 1 - c);
282                data[c * samples + center] += 25.0 * w2[c];
283                data[c * samples + center + 1] -= 12.0 * w2[c];
284            }
285        }
286
287        struct Decompose<'a>(&'a [f32], usize, usize, f64);
288        impl dspk_core::compute::ComputeTask for Decompose<'_> {
289            type Output = Vec<MotorUnitPulseTrain>;
290            fn run(self, client: Client) -> Self::Output {
291                ConvolutiveBssDecomposer::new(3, 4, 20.0).decompose(&client, self.0, self.1, self.2, self.3)
292            }
293        }
294        for target in dspk_core::compute::ComputeTarget::for_tests() {
295            let units = target.run(Decompose(&data, channels, samples, fs)).unwrap();
296            assert!(units.len() >= 2, "{}: expected at least 2 motor units, got {}", target.name(), units.len());
297            assert!(units[0].cov_isi < 0.15, "cov_isi={}", units[0].cov_isi);
298            assert!(units[0].pnr_db > 15.0, "pnr_db={}", units[0].pnr_db);
299        }
300    }
301}
dspk_synapse/sorting/cbss.rs

dspk_synapse/sorting/
cbss.rs
