use crate::io::probe::SensorLayout;
use crate::extraction::SnippetBatch;
use dspk_ops::math::peak_to_peak;
use crate::core::PeakLocalizer;
pub fn localize_spike_center_of_mass(
channel_ids: &[usize],
ptp_amplitudes: &[f32],
layout: &SensorLayout,
power: f32,
) -> [f32; 3] {
let mut sum_w = 0.0f32;
let mut sum_x = 0.0f32;
let mut sum_y = 0.0f32;
let mut sum_z = 0.0f32;
for (&ch_id, &) in channel_ids.iter().zip(ptp_amplitudes.iter()) {
if let Ok(site) = layout.get_site(ch_id) {
let w = amp.max(0.0).powf(power);
sum_w += w;
sum_x += w * site.position.x_um;
sum_y += w * site.position.y_um;
sum_z += w * site.position.z_um;
}
}
if sum_w > 1e-8 {
[sum_x / sum_w, sum_y / sum_w, sum_z / sum_w]
} else if let Some(&first_ch) = channel_ids.first() {
layout
.get_site(first_ch)
.map(|s| [s.position.x_um, s.position.y_um, s.position.z_um])
.unwrap_or([0.0, 0.0, 0.0])
} else {
[0.0, 0.0, 0.0]
}
}
#[derive(Debug, Clone, Copy)]
pub struct CenterOfMassLocalizer {
pub amplitude_power: f32,
}
impl Default for CenterOfMassLocalizer {
fn default() -> Self {
Self {
amplitude_power: 1.0,
}
}
}
impl PeakLocalizer for CenterOfMassLocalizer {
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_center_of_mass(
ch_ids,
&ptp_buf,
layout,
self.amplitude_power,
));
}
Ok(out)
}
}