dspk_synapse/features/
wavelet.rs

1//! Multi-level Discrete Wavelet Transform (DWT) Feature Extraction (`wavelet.rs`, *Wave_Clus* style).
2//!
3//! Decomposes each spike snippet with a multi-level Haar filter bank and selects the $D$ wavelet
4//! coefficients with the largest spread across the spike population, ranked by a normalized
5//! interquartile range (Wave_Clus uses a Lilliefors test; this is a simpler proxy).
6
7use dspk_ops::math::interquartile_range;
8use crate::core::{FeatureEmbedder, SnippetBatch};
9
10/// Computes a multi-level 1D Haar Discrete Wavelet Transform (DWT) of `signal`.
11pub fn compute_haar_dwt(signal: &[f32], levels: usize) -> Vec<f32> {
12    if signal.is_empty() {
13        return Vec::new();
14    }
15    let mut coeffs = signal.to_vec();
16    let mut len = coeffs.len();
17    let inv_sqrt2 = std::f32::consts::FRAC_1_SQRT_2;
18
19    for _ in 0..levels.max(1) {
20        let half = len / 2;
21        if half == 0 {
22            break;
23        }
24        let mut temp = vec![0.0f32; half * 2];
25        for i in 0..half {
26            let a = coeffs[2 * i];
27            let b = coeffs[2 * i + 1];
28            temp[i] = (a + b) * inv_sqrt2;
29            temp[half + i] = (a - b) * inv_sqrt2;
30        }
31        coeffs[..half * 2].copy_from_slice(&temp);
32        len = half;
33    }
34    coeffs
35}
36
37/// Wavelet feature embedder (*Wave_Clus* / Quiroga et al., 2004) implementing [`FeatureEmbedder`].
38#[derive(Debug, Clone, Copy)]
39pub struct WaveletFeatureEmbedder {
40    pub num_components: usize,
41    pub levels: usize,
42}
43
44impl Default for WaveletFeatureEmbedder {
45    fn default() -> Self {
46        Self {
47            num_components: 8,
48            levels: 4,
49        }
50    }
51}
52
53impl FeatureEmbedder for WaveletFeatureEmbedder {
54    fn embed(&self, batch: &SnippetBatch) -> dspk_core::DspResult<(Vec<f32>, usize)> {
55        let d = self.num_components.max(1);
56        if batch.num_spikes == 0 {
57            return Ok((Vec::new(), d));
58        }
59
60        let num_spikes = batch.num_spikes;
61        let ch_count = batch.num_channels;
62        let t_len = batch.num_samples;
63        let total_coeffs = ch_count * t_len;
64        let actual_d = d.min(total_coeffs).max(1);
65
66        // 1. Compute DWT coefficients for every spike across all snippet channels
67        let mut all_coeffs = vec![0.0f32; num_spikes * total_coeffs];
68        for s in 0..num_spikes {
69            for c in 0..ch_count {
70                let wave = batch.channel_slice(s, c);
71                let dwt = compute_haar_dwt(wave, self.levels);
72                let dst = &mut all_coeffs[s * total_coeffs + c * t_len..s * total_coeffs + (c + 1) * t_len];
73                dst.copy_from_slice(&dwt);
74            }
75        }
76
77        // 2. Score each coefficient index by robust IQR across spikes
78        let mut col_buf = vec![0.0f32; num_spikes];
79        let mut scores: Vec<(usize, f32)> = (0..total_coeffs)
80            .map(|f| {
81                for s in 0..num_spikes {
82                    col_buf[s] = all_coeffs[s * total_coeffs + f];
83                }
84                (f, interquartile_range(&col_buf))
85            })
86            .collect();
87        scores.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));
88
89        let selected: Vec<usize> = scores.iter().take(actual_d).map(|&(idx, _)| idx).collect();
90        let mut out = vec![0.0f32; num_spikes * actual_d];
91        for s in 0..num_spikes {
92            for (k, &f_idx) in selected.iter().enumerate() {
93                out[s * actual_d + k] = all_coeffs[s * total_coeffs + f_idx];
94            }
95        }
96
97        Ok((out, actual_d))
98    }
99}
100
101#[cfg(test)]
102mod tests {
103    use super::*;
104
105    #[test]
106    fn test_haar_dwt_preserves_energy() {
107        let sig = vec![1.0f32, -2.0, 4.0, -1.0, 3.0, 0.5, -1.5, 2.0];
108        let e_in: f32 = sig.iter().map(|v| v * v).sum();
109        let coeffs = compute_haar_dwt(&sig, 3);
110        let e_out: f32 = coeffs.iter().map(|v| v * v).sum();
111        assert!((e_in - e_out).abs() < 1e-4);
112    }
113}
dspk_synapse/features/wavelet.rs

dspk_synapse/features/
wavelet.rs
