segovia 0.4.1

A fast, chunked, memory-bounded Rust engine for electrophysiology (Neuropixels) signal processing, with Python bindings.
Documentation
use std::sync::Arc;

use numpy::ndarray::Array2;
use thiserror::Error;

use crate::core::ChunkSource;
use crate::sim::{mix, Rng};

#[derive(Debug, Error)]
pub enum IfcError {
    #[error("n_channels must be greater than zero")]
    NoChannels,
    #[error("duration_s must be greater than zero")]
    NoDuration,
    #[error("sample_rate must be greater than zero")]
    BadSampleRate,
    #[error("n_populations must be greater than zero")]
    NoPopulations,
    #[error("event_rate must not be negative")]
    BadEventRate,
    #[error("noise_level must not be negative")]
    BadNoise,
    #[error("lsb must be greater than zero")]
    BadLsb,
}

#[derive(Clone, Copy)]
pub struct IfcConfig {
    pub n_channels: usize,
    pub duration_s: f64,
    pub sample_rate: f64,
    pub n_populations: usize,
    pub event_rate: f64,
    pub noise_level: f64,
    pub lsb: f64,
    pub seed: u64,
}

struct Population {
    amplitude: f64,
    gains: Vec<f32>,
    template: Vec<f64>,
}

struct Event {
    sample: i64,
    population: u32,
    amp_scale: f32,
}

struct IfcData {
    n_channels: usize,
    n_samples: usize,
    sample_rate: f64,
    noise_level: f32,
    lsb: f64,
    noise_seed: u64,
    template_len: usize,
    center: usize,
    populations: Vec<Population>,
    events: Vec<Event>,
}

pub struct SyntheticIfcReader {
    data: Arc<IfcData>,
}

impl SyntheticIfcReader {
    pub fn new(config: IfcConfig) -> Result<Self, IfcError> {
        if config.n_channels == 0 {
            return Err(IfcError::NoChannels);
        }
        if config.duration_s <= 0.0 || config.duration_s.is_nan() {
            return Err(IfcError::NoDuration);
        }
        if config.sample_rate <= 0.0 || config.sample_rate.is_nan() {
            return Err(IfcError::BadSampleRate);
        }
        if config.n_populations == 0 {
            return Err(IfcError::NoPopulations);
        }
        if config.event_rate < 0.0 {
            return Err(IfcError::BadEventRate);
        }
        if config.noise_level < 0.0 {
            return Err(IfcError::BadNoise);
        }
        if config.lsb <= 0.0 || config.lsb.is_nan() {
            return Err(IfcError::BadLsb);
        }

        let n_samples = (config.duration_s * config.sample_rate).round() as usize;
        let template_len = ((0.0015 * config.sample_rate).round() as usize).max(3);
        let center = template_len / 2;

        let mut populations = Vec::with_capacity(config.n_populations);
        for p in 0..config.n_populations {
            let mut rng = Rng::seed(mix(config.seed, mix(0x1FC0_0001, p as u64)));
            let amplitude = 0.2 + rng.next_f64() * 1.3;
            let sigma_us = 30.0 + rng.next_f64() * 50.0;
            let separation_us = 150.0 + rng.next_f64() * 200.0;
            let sigma = (sigma_us * 1e-6 * config.sample_rate).max(1.0);
            let separation = separation_us * 1e-6 * config.sample_rate;

            let mut gains = Vec::with_capacity(config.n_channels);
            for c in 0..config.n_channels {
                if c == 0 {
                    gains.push(1.0);
                } else {
                    gains.push((0.6 + rng.next_f64() * 0.4) as f32);
                }
            }

            let c1 = center as f64 - separation * 0.5;
            let c2 = center as f64 + separation * 0.5;
            let template = (0..template_len)
                .map(|k| {
                    let t = k as f64;
                    let a = (-0.5 * ((t - c1) / sigma).powi(2)).exp();
                    let b = (-0.5 * ((t - c2) / sigma).powi(2)).exp();
                    a - b
                })
                .collect();

            populations.push(Population {
                amplitude,
                gains,
                template,
            });
        }

        let mut events = Vec::new();
        if config.event_rate > 0.0 {
            let mut arrivals = Rng::seed(mix(config.seed, 0x1FC0_A551));
            let mut t = 0.0f64;
            let mut idx = 0u64;
            loop {
                let p = arrivals.next_f64().max(f64::MIN_POSITIVE);
                t += -p.ln() / config.event_rate;
                let sample = (t * config.sample_rate).round() as i64;
                if sample < 0 {
                    idx += 1;
                    continue;
                }
                if sample as usize >= n_samples {
                    break;
                }
                let mut ev = Rng::seed(mix(config.seed, mix(0x1FC0_E7E7, idx)));
                let population = (ev.next_f64() * config.n_populations as f64) as usize;
                let population = population.min(config.n_populations - 1);
                let amp_scale = (1.0 + 0.1 * ev.next_gaussian()).max(0.05) as f32;
                events.push(Event {
                    sample,
                    population: population as u32,
                    amp_scale,
                });
                idx += 1;
            }
        }

        events.sort_by_key(|e| e.sample);

        let data = IfcData {
            n_channels: config.n_channels,
            n_samples,
            sample_rate: config.sample_rate,
            noise_level: config.noise_level as f32,
            lsb: config.lsb,
            noise_seed: mix(config.seed, 0x1FC0_0757_D0D0_5EED),
            template_len,
            center,
            populations,
            events,
        };

        Ok(Self {
            data: Arc::new(data),
        })
    }

    pub fn ground_truth(&self) -> (Vec<i64>, Vec<i32>, Vec<i32>) {
        let n = self.data.events.len();
        let mut samples = Vec::with_capacity(n);
        let mut population_ids = Vec::with_capacity(n);
        let mut amplitudes = Vec::with_capacity(n);
        let scale = 1.0 / self.data.lsb;
        for e in &self.data.events {
            samples.push(e.sample);
            population_ids.push(e.population as i32);
            let peak =
                self.data.populations[e.population as usize].amplitude * e.amp_scale as f64 * scale;
            amplitudes.push(peak.round() as i32);
        }
        (samples, population_ids, amplitudes)
    }
}

impl IfcData {
    fn synth_chunk(&self, start: usize, rows: usize) -> Array2<i16> {
        let mut buf = Array2::<f32>::zeros((rows, self.n_channels));

        if self.noise_level > 0.0 {
            for r in 0..rows {
                let s_abs = (start + r) as u64;
                let mut rng = Rng::seed(mix(self.noise_seed, s_abs));
                for c in 0..self.n_channels {
                    buf[[r, c]] = self.noise_level * rng.next_gaussian() as f32;
                }
            }
        }

        let start_i = start as i64;
        let end_i = (start + rows) as i64;
        let tlen = self.template_len as i64;
        let center = self.center as i64;
        let lo = self.events.partition_point(|e| e.sample < start_i - tlen);
        let hi = self.events.partition_point(|e| e.sample < end_i + tlen);

        for event in &self.events[lo..hi] {
            let pop = &self.populations[event.population as usize];
            let t0 = event.sample - center;
            for k in 0..self.template_len {
                let s = t0 + k as i64;
                if !(start_i..end_i).contains(&s) {
                    continue;
                }
                let r = (s - start_i) as usize;
                let shape = (pop.template[k] * pop.amplitude) as f32 * event.amp_scale;
                for c in 0..self.n_channels {
                    buf[[r, c]] += shape * pop.gains[c];
                }
            }
        }

        let scale = 1.0 / self.lsb;
        buf.mapv(|v| {
            let q = (v as f64 * scale).round();
            q.clamp(i16::MIN as f64, i16::MAX as f64) as i16
        })
    }
}

impl ChunkSource for SyntheticIfcReader {
    type Chunks = IfcChunkIter;

    fn n_channels(&self) -> usize {
        self.data.n_channels
    }

    fn n_samples(&self) -> usize {
        self.data.n_samples
    }

    fn sample_rate(&self) -> f64 {
        self.data.sample_rate
    }

    fn chunks(&self, chunk_samples: usize) -> IfcChunkIter {
        IfcChunkIter {
            data: Arc::clone(&self.data),
            chunk_samples,
            pos: 0,
        }
    }
}

pub struct IfcChunkIter {
    data: Arc<IfcData>,
    chunk_samples: usize,
    pos: usize,
}

impl Iterator for IfcChunkIter {
    type Item = Array2<i16>;

    fn next(&mut self) -> Option<Self::Item> {
        if self.pos >= self.data.n_samples {
            return None;
        }
        let rows = self.chunk_samples.min(self.data.n_samples - self.pos);
        let chunk = self.data.synth_chunk(self.pos, rows);
        self.pos += rows;
        Some(chunk)
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    fn config() -> IfcConfig {
        IfcConfig {
            n_channels: 2,
            duration_s: 0.5,
            sample_rate: 100000.0,
            n_populations: 3,
            event_rate: 200.0,
            noise_level: 0.01,
            lsb: 1.0e-4,
            seed: 42,
        }
    }

    fn stack(chunks: impl Iterator<Item = Array2<i16>>) -> Array2<i16> {
        let collected: Vec<Array2<i16>> = chunks.collect();
        let n_cols = collected[0].ncols();
        let n_rows: usize = collected.iter().map(|c| c.nrows()).sum();
        let mut out = Array2::<i16>::zeros((n_rows, n_cols));
        let mut row = 0;
        for chunk in collected {
            for r in 0..chunk.nrows() {
                for c in 0..n_cols {
                    out[[row + r, c]] = chunk[[r, c]];
                }
            }
            row += chunk.nrows();
        }
        out
    }

    #[test]
    fn shape_matches_config() {
        let reader = SyntheticIfcReader::new(config()).unwrap();
        assert_eq!(reader.n_channels(), 2);
        assert_eq!(reader.n_samples(), 50000);
        assert_eq!(reader.sample_rate(), 100000.0);
    }

    #[test]
    fn chunks_reconstruct_full_length() {
        let reader = SyntheticIfcReader::new(config()).unwrap();
        let total: usize = reader.chunks(4096).map(|c| c.nrows()).sum();
        assert_eq!(total, 50000);
        assert!(reader.chunks(4096).all(|c| c.ncols() == 2));
    }

    #[test]
    fn output_is_chunk_size_independent() {
        let reader = SyntheticIfcReader::new(config()).unwrap();
        let big = stack(reader.chunks(50000));
        let small = stack(reader.chunks(1000));
        assert_eq!(big, small);
    }

    #[test]
    fn same_seed_is_deterministic() {
        let a = stack(SyntheticIfcReader::new(config()).unwrap().chunks(4096));
        let b = stack(SyntheticIfcReader::new(config()).unwrap().chunks(4096));
        assert_eq!(a, b);
    }

    #[test]
    fn different_seed_differs() {
        let mut other = config();
        other.seed = 7;
        let a = stack(SyntheticIfcReader::new(config()).unwrap().chunks(4096));
        let b = stack(SyntheticIfcReader::new(other).unwrap().chunks(4096));
        assert_ne!(a, b);
    }

    #[test]
    fn ground_truth_is_sorted_and_in_range() {
        let reader = SyntheticIfcReader::new(config()).unwrap();
        let (samples, pops, amps) = reader.ground_truth();
        assert!(!samples.is_empty());
        assert_eq!(samples.len(), pops.len());
        assert_eq!(samples.len(), amps.len());
        assert!(samples.windows(2).all(|w| w[0] <= w[1]));
        assert!(samples.iter().all(|&s| s >= 0 && (s as usize) < 50000));
        assert!(pops.iter().all(|&p| (0..3).contains(&p)));
    }

    #[test]
    fn pulse_is_bipolar() {
        let cfg = IfcConfig {
            n_populations: 1,
            event_rate: 20.0,
            noise_level: 0.0,
            ..config()
        };
        let reader = SyntheticIfcReader::new(cfg).unwrap();
        let signal = stack(reader.chunks(4096));
        let mut has_pos = false;
        let mut has_neg = false;
        for r in 0..signal.nrows() {
            let v = signal[[r, 0]];
            if v > 1000 {
                has_pos = true;
            }
            if v < -1000 {
                has_neg = true;
            }
        }
        assert!(has_pos && has_neg);
    }
}