dspk_synapse/features/
mod.rs

1use cubecl::prelude::Client;
2use dspk_ops::linalg::{PcaModel, PpcaModel};
3use dspk_core::compute::{ComputeTarget, ComputeTask};
4use dspk_core::{DspError, DspResult};
5
6use crate::core::{FeatureEmbedder, SnippetBatch, WaveformSnippet};
7
8pub mod conduction;
9pub(crate) mod kernels;
10
11/// A neighbourhood slot without a channel (`ChannelNeighbourhoods` tables).
12pub use kernels::NO_CHANNEL;
13pub mod local_svd;
14pub mod morphology;
15pub mod wavelet;
16
17pub use local_svd::{ChannelNeighbourhoods, LocalSvd, LocalSvdOptions};
18pub use conduction::{ConductionVelocityEstimate, estimate_hdemg_conduction_velocity};
19pub use morphology::{SpikeMorphology, compute_morphology};
20pub use wavelet::{WaveletFeatureEmbedder, compute_haar_dwt};
21
22/// Components the PCA / PPCA embedders keep unless told otherwise.
23pub const DEFAULT_FEATURE_COMPONENTS: usize = 4;
24
25/// Features × spikes matrix (`matrix[f · spikes + s]`), the layout the linear models fit.
26fn feature_matrix(features: usize, spikes: usize, feature_of: impl Fn(usize) -> Vec<f32>) -> Vec<f32> {
27    let mut matrix = vec![0.0f32; features * spikes];
28    for s in 0..spikes {
29        for (f, v) in feature_of(s).into_iter().enumerate().take(features) {
30            matrix[f * spikes + s] = v;
31        }
32    }
33    matrix
34}
35
36/// Projections `[component][spike]` as one row per spike (`out[s · d + k]`).
37fn per_spike(projected: &[f32], d: usize, spikes: usize) -> Vec<f32> {
38    let mut out = vec![0.0f32; spikes * d];
39    for s in 0..spikes {
40        for k in 0..d {
41            out[s * d + k] = projected[k * spikes + s];
42        }
43    }
44    out
45}
46
47/// Fits a PCA model (`dspk_ops::linalg::PcaModel`) on `client` to spike waveforms (primary trace
48/// of each snippet; snippets of another length are left at zero) and projects each spike.
49pub fn extract_waveform_pca(
50    client: &Client,
51    snippets: &[WaveformSnippet],
52    num_components: usize,
53) -> Option<(PcaModel, Vec<Vec<f32>>)> {
54    let first = snippets.first()?;
55    let (samples, spikes) = (first.num_samples, snippets.len());
56    let matrix = feature_matrix(samples, spikes, |s| {
57        let w = &snippets[s].waveform;
58        if w.len() == samples { w.clone() } else { Vec::new() }
59    });
60    let pca = PcaModel::fit::<f32>(client, &matrix, samples, spikes, num_components);
61    let d = pca.num_components;
62    let flat = per_spike(&pca.project_cpu(&matrix, samples, spikes), d, spikes);
63    let rows = flat.chunks_exact(d.max(1)).map(<[f32]>::to_vec).collect();
64    Some((pca, rows))
65}
66
67/// Which linear model an embedder fits.
68#[derive(Debug, Clone, Copy, PartialEq, Eq)]
69enum LinearModel {
70    Pca,
71    Ppca,
72}
73
74/// Fits `model` to a batch on the task's device and returns `(rows, d)`.
75struct EmbedTask<'a> {
76    model: LinearModel,
77    matrix: &'a [f32],
78    features: usize,
79    spikes: usize,
80    components: usize,
81}
82
83impl ComputeTask for EmbedTask<'_> {
84    type Output = (Vec<f32>, usize);
85    fn run(self, client: Client) -> Self::Output {
86        let (projected, d) = match self.model {
87            LinearModel::Pca => {
88                let m = PcaModel::fit::<f32>(&client, self.matrix, self.features, self.spikes, self.components);
89                (m.project_cpu(self.matrix, self.features, self.spikes), m.num_components)
90            }
91            LinearModel::Ppca => {
92                let m = PpcaModel::fit::<f32>(&client, self.matrix, self.features, self.spikes, self.components);
93                (m.project_cpu(self.matrix, self.features, self.spikes), m.num_components)
94            }
95        };
96        (per_spike(&projected, d, self.spikes), d)
97    }
98}
99
100fn embed_linear(model: LinearModel, target: ComputeTarget, components: usize, batch: &SnippetBatch) -> DspResult<(Vec<f32>, usize)> {
101    let d = components.max(1);
102    if batch.num_spikes == 0 {
103        return Ok((Vec::new(), d));
104    }
105    let features = batch.num_channels * batch.num_samples;
106    let matrix = feature_matrix(features, batch.num_spikes, |s| batch.snippet_slice(s).to_vec());
107    let task = EmbedTask { model, matrix: &matrix, features, spikes: batch.num_spikes, components: d };
108    target.run(task).map_err(|e| DspError::ComputeError(e.to_string()))
109}
110
111/// Classical PCA feature embedder delegating to [`dspk_ops::linalg::PcaModel`], fitted on
112/// `target`.
113#[derive(Debug, Clone, Copy)]
114pub struct PcaFeatureEmbedder {
115    pub num_components: usize,
116    pub target: ComputeTarget,
117}
118
119impl PcaFeatureEmbedder {
120    /// [`DEFAULT_FEATURE_COMPONENTS`] components on `target`.
121    pub fn new(target: ComputeTarget) -> Self {
122        Self { num_components: DEFAULT_FEATURE_COMPONENTS, target }
123    }
124}
125
126impl FeatureEmbedder for PcaFeatureEmbedder {
127    fn embed(&self, batch: &SnippetBatch) -> DspResult<(Vec<f32>, usize)> {
128        embed_linear(LinearModel::Pca, self.target, self.num_components, batch)
129    }
130}
131
132/// Probabilistic PCA (PPCA) embedder delegating to [`dspk_ops::linalg::PpcaModel`], fitted on
133/// `target`.
134#[derive(Debug, Clone, Copy)]
135pub struct PpcaFeatureEmbedder {
136    pub num_components: usize,
137    pub target: ComputeTarget,
138}
139
140impl PpcaFeatureEmbedder {
141    /// [`DEFAULT_FEATURE_COMPONENTS`] components on `target`.
142    pub fn new(target: ComputeTarget) -> Self {
143        Self { num_components: DEFAULT_FEATURE_COMPONENTS, target }
144    }
145}
146
147impl FeatureEmbedder for PpcaFeatureEmbedder {
148    fn embed(&self, batch: &SnippetBatch) -> DspResult<(Vec<f32>, usize)> {
149        embed_linear(LinearModel::Ppca, self.target, self.num_components, batch)
150    }
151}
dspk_synapse/features/mod.rs

dspk_synapse/features/
mod.rs
