dspk-synapse 0.0.6

Electrophysiology for dspk: recording and sorting formats, probes, spike detection, waveforms, features, clustering, localization, quality metrics, and the shared parts of the sorters
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137

//! Grid Convolution Spike Source Localization (`grid_convolution.rs`).
//!
//! Evaluates a 3D grid of synthetic monopole footprints around the primary channel and computes an
//! exponentially weighted soft-argmax over the cosine similarity of each footprint with the spike's
//! peak-to-peak amplitudes. Inspired by SpikeInterface's `GridConvolution`, which convolves
//! prototype waveforms instead; results differ.

use crate::io::probe::SensorLayout;
use crate::extraction::SnippetBatch;
use crate::core::PeakLocalizer;
use dspk_ops::math::peak_to_peak;

/// Localizes a single spike via Grid Convolution soft-argmax over a local 3D grid.
pub fn localize_spike_grid_convolution(
    channel_ids: &[usize],
    ptp_amplitudes: &[f32],
    layout: &SensorLayout,
    xy_radius_um: f32,
    xy_step_um: f32,
    z_values_um: &[f32],
    softmax_temperature: f32,
) -> [f32; 3] {
    if channel_ids.is_empty() || ptp_amplitudes.is_empty() {
        return [0.0, 0.0, 15.0];
    }

    let prim_site = match layout.get_site(channel_ids[0]) {
        Ok(s) => s,
        Err(_) => return [0.0, 0.0, 15.0],
    };
    let anchor_x = prim_site.position.x_um;
    let anchor_y = prim_site.position.y_um;

    let mut sites = Vec::with_capacity(channel_ids.len());
    let mut obs = Vec::with_capacity(channel_ids.len());
    for (&ch, &amp) in channel_ids.iter().zip(ptp_amplitudes.iter()) {
        if let Ok(s) = layout.get_site(ch) {
            sites.push([s.position.x_um, s.position.y_um, s.position.z_um]);
            obs.push(amp.max(0.0));
        }
    }

    let obs_norm = obs.iter().map(|v| v * v).sum::<f32>().sqrt().max(1e-8);
    let step = xy_step_um.max(2.0);
    let n_steps = (xy_radius_um / step).round() as isize;
    let temp = softmax_temperature.max(0.01);

    let mut grid_pts = Vec::new();
    let mut scores = Vec::new();
    let mut max_score = f32::NEG_INFINITY;

    for ix in -n_steps..=n_steps {
        let gx = anchor_x + (ix as f32) * step;
        for iy in -n_steps..=n_steps {
            let gy = anchor_y + (iy as f32) * step;
            for &gz in z_values_um {
                let mut dot = 0.0f32;
                let mut syn_norm_sq = 0.0f32;
                for (s, &v_obs) in sites.iter().zip(obs.iter()) {
                    let dx = gx - s[0];
                    let dy = gy - s[1];
                    let dz = gz - s[2];
                    let w = 1.0 / (dx * dx + dy * dy + dz * dz + 1.0).sqrt();
                    dot += w * v_obs;
                    syn_norm_sq += w * w;
                }
                let cos_sim = dot / (syn_norm_sq.sqrt().max(1e-8) * obs_norm);
                let logit = cos_sim / temp;
                if logit > max_score {
                    max_score = logit;
                }
                grid_pts.push([gx, gy, gz]);
                scores.push(logit);
            }
        }
    }

    let mut sum_w = 0.0f32;
    let mut est = [0.0f32; 3];
    for (pt, &sc) in grid_pts.iter().zip(scores.iter()) {
        let w = (sc - max_score).exp();
        sum_w += w;
        est[0] += w * pt[0];
        est[1] += w * pt[1];
        est[2] += w * pt[2];
    }

    if sum_w > 1e-12 {
        [est[0] / sum_w, est[1] / sum_w, est[2] / sum_w]
    } else {
        [anchor_x, anchor_y, 15.0]
    }
}

/// Grid Convolution localizer implementing [`PeakLocalizer`].
#[derive(Debug, Clone)]
pub struct GridConvolutionLocalizer {
    pub xy_radius_um: f32,
    pub xy_step_um: f32,
    pub z_values_um: Vec<f32>,
    pub softmax_temperature: f32,
}

impl Default for GridConvolutionLocalizer {
    fn default() -> Self {
        Self {
            xy_radius_um: 30.0,
            xy_step_um: 5.0,
            z_values_um: vec![5.0, 10.0, 20.0, 35.0, 50.0],
            softmax_temperature: 0.05,
        }
    }
}

impl PeakLocalizer for GridConvolutionLocalizer {
    fn localize(&self, batch: &SnippetBatch, layout: &SensorLayout) -> dspk_core::DspResult<Vec<[f32; 3]>> {
        let mut out = Vec::with_capacity(batch.num_spikes);
        let mut ptp_buf = vec![0.0f32; batch.num_channels];

        for i in 0..batch.num_spikes {
            for k in 0..batch.num_channels {
                ptp_buf[k] = peak_to_peak(batch.channel_slice(i, k));
            }
            let ch_ids = batch.spike_channel_ids(i);
            out.push(localize_spike_grid_convolution(
                ch_ids,
                &ptp_buf,
                layout,
                self.xy_radius_um,
                self.xy_step_um,
                &self.z_values_um,
                self.softmax_temperature,
            ));
        }
        Ok(out)
    }
}