dspk_synapse/spatial/
grid_convolution.rs

1//! Grid Convolution Spike Source Localization (`grid_convolution.rs`).
2//!
3//! Evaluates a 3D grid of synthetic monopole footprints around the primary channel and computes an
4//! exponentially weighted soft-argmax over the cosine similarity of each footprint with the spike's
5//! peak-to-peak amplitudes. Inspired by SpikeInterface's `GridConvolution`, which convolves
6//! prototype waveforms instead; results differ.
7
8use crate::io::probe::SensorLayout;
9use crate::extraction::SnippetBatch;
10use crate::core::PeakLocalizer;
11use dspk_ops::math::peak_to_peak;
12
13/// Localizes a single spike via Grid Convolution soft-argmax over a local 3D grid.
14pub fn localize_spike_grid_convolution(channel_ids: &[usize], ptp_amplitudes: &[f32], layout: &SensorLayout, settings: &GridConvolutionLocalizer) -> [f32; 3] {
15    let (xy_radius_um, xy_step_um, z_values_um, softmax_temperature) = (settings.xy_radius_um, settings.xy_step_um, settings.z_values_um.as_slice(), settings.softmax_temperature);
16    if channel_ids.is_empty() || ptp_amplitudes.is_empty() {
17        return [0.0, 0.0, 15.0];
18    }
19
20    let prim_site = match layout.get_site(channel_ids[0]) {
21        Ok(s) => s,
22        Err(_) => return [0.0, 0.0, 15.0],
23    };
24    let anchor_x = prim_site.position.x_um;
25    let anchor_y = prim_site.position.y_um;
26
27    let mut sites = Vec::with_capacity(channel_ids.len());
28    let mut obs = Vec::with_capacity(channel_ids.len());
29    for (&ch, &amp) in channel_ids.iter().zip(ptp_amplitudes.iter()) {
30        if let Ok(s) = layout.get_site(ch) {
31            sites.push([s.position.x_um, s.position.y_um, s.position.z_um]);
32            obs.push(amp.max(0.0));
33        }
34    }
35
36    let obs_norm = obs.iter().map(|v| v * v).sum::<f32>().sqrt().max(1e-8);
37    let step = xy_step_um.max(2.0);
38    let n_steps = (xy_radius_um / step).round() as isize;
39    let temp = softmax_temperature.max(0.01);
40
41    let mut grid_pts = Vec::new();
42    let mut scores = Vec::new();
43    let mut max_score = f32::NEG_INFINITY;
44
45    for ix in -n_steps..=n_steps {
46        let gx = anchor_x + (ix as f32) * step;
47        for iy in -n_steps..=n_steps {
48            let gy = anchor_y + (iy as f32) * step;
49            for &gz in z_values_um {
50                let mut dot = 0.0f32;
51                let mut syn_norm_sq = 0.0f32;
52                for (s, &v_obs) in sites.iter().zip(obs.iter()) {
53                    let dx = gx - s[0];
54                    let dy = gy - s[1];
55                    let dz = gz - s[2];
56                    let w = 1.0 / (dx * dx + dy * dy + dz * dz + 1.0).sqrt();
57                    dot += w * v_obs;
58                    syn_norm_sq += w * w;
59                }
60                let cos_sim = dot / (syn_norm_sq.sqrt().max(1e-8) * obs_norm);
61                let logit = cos_sim / temp;
62                if logit > max_score {
63                    max_score = logit;
64                }
65                grid_pts.push([gx, gy, gz]);
66                scores.push(logit);
67            }
68        }
69    }
70
71    let mut sum_w = 0.0f32;
72    let mut est = [0.0f32; 3];
73    for (pt, &sc) in grid_pts.iter().zip(scores.iter()) {
74        let w = (sc - max_score).exp();
75        sum_w += w;
76        est[0] += w * pt[0];
77        est[1] += w * pt[1];
78        est[2] += w * pt[2];
79    }
80
81    if sum_w > 1e-12 {
82        [est[0] / sum_w, est[1] / sum_w, est[2] / sum_w]
83    } else {
84        [anchor_x, anchor_y, 15.0]
85    }
86}
87
88/// Grid Convolution localizer implementing [`PeakLocalizer`].
89#[derive(Debug, Clone)]
90pub struct GridConvolutionLocalizer {
91    pub xy_radius_um: f32,
92    pub xy_step_um: f32,
93    pub z_values_um: Vec<f32>,
94    pub softmax_temperature: f32,
95}
96
97impl Default for GridConvolutionLocalizer {
98    fn default() -> Self {
99        Self {
100            xy_radius_um: 30.0,
101            xy_step_um: 5.0,
102            z_values_um: vec![5.0, 10.0, 20.0, 35.0, 50.0],
103            softmax_temperature: 0.05,
104        }
105    }
106}
107
108impl PeakLocalizer for GridConvolutionLocalizer {
109    fn localize(&self, batch: &SnippetBatch, layout: &SensorLayout) -> dspk_core::DspResult<Vec<[f32; 3]>> {
110        let mut out = Vec::with_capacity(batch.num_spikes);
111        let mut ptp_buf = vec![0.0f32; batch.num_channels];
112
113        for i in 0..batch.num_spikes {
114            for k in 0..batch.num_channels {
115                ptp_buf[k] = peak_to_peak(batch.channel_slice(i, k));
116            }
117            let ch_ids = batch.spike_channel_ids(i);
118            out.push(localize_spike_grid_convolution(ch_ids, &ptp_buf, layout, self));
119        }
120        Ok(out)
121    }
122}
dspk_synapse/spatial/grid_convolution.rs

dspk_synapse/spatial/
grid_convolution.rs
