rustytorch/crates/specialized/rtx-neuro-signal/src/resample.rs

//! Resampling functions for neuroimaging signals.

use crate::filter::{FilterMethod, lowpass};
use crate::{SignalError, SignalResult};
use std::f64::consts::PI;

/// Resample a signal to a new sampling frequency.
///
/// Applies anti-aliasing lowpass filter when downsampling.
///
/// # Arguments
/// * `signal` - Input signal samples
/// * `sfreq_in` - Original sampling frequency in Hz
/// * `sfreq_out` - Target sampling frequency in Hz
///
/// # Returns
/// Resampled signal
pub fn resample(signal: &[f64], sfreq_in: f64, sfreq_out: f64) -> SignalResult<Vec<f64>> {
    if signal.is_empty() {
        return Ok(Vec::new());
    }

    if sfreq_in <= 0.0 || sfreq_out <= 0.0 {
        return Err(SignalError::InvalidParameters(
            "Sampling frequencies must be positive".to_string(),
        ));
    }

    // Calculate the resampling ratio
    let ratio = sfreq_out / sfreq_in;

    if (ratio - 1.0).abs() < 1e-10 {
        return Ok(signal.to_vec());
    }

    // If downsampling, apply anti-aliasing filter first
    let filtered = if ratio < 1.0 {
        let cutoff = sfreq_out / 2.0 * 0.9; // 90% of new Nyquist
        lowpass(signal, sfreq_in, cutoff, FilterMethod::Fft)?
    } else {
        signal.to_vec()
    };

    // Calculate new length
    let n_out = (filtered.len() as f64 * ratio).round() as usize;

    // Resample using linear interpolation (simple but effective)
    let mut output = Vec::with_capacity(n_out);

    for i in 0..n_out {
        let src_idx = i as f64 / ratio;
        let idx_low = src_idx.floor() as usize;
        let idx_high = (idx_low + 1).min(filtered.len() - 1);
        let frac = src_idx - idx_low as f64;

        let value = filtered[idx_low] * (1.0 - frac) + filtered[idx_high] * frac;
        output.push(value);
    }

    Ok(output)
}

/// Resample using sinc interpolation (higher quality).
///
/// Uses a windowed sinc kernel for interpolation, providing
/// better frequency response than linear interpolation.
pub fn resample_sinc(
    signal: &[f64],
    sfreq_in: f64,
    sfreq_out: f64,
    num_taps: usize,
) -> SignalResult<Vec<f64>> {
    if signal.is_empty() {
        return Ok(Vec::new());
    }

    if sfreq_in <= 0.0 || sfreq_out <= 0.0 {
        return Err(SignalError::InvalidParameters(
            "Sampling frequencies must be positive".to_string(),
        ));
    }

    let ratio = sfreq_out / sfreq_in;

    if (ratio - 1.0).abs() < 1e-10 {
        return Ok(signal.to_vec());
    }

    // Anti-aliasing filter for downsampling
    let filtered = if ratio < 1.0 {
        let cutoff = sfreq_out / 2.0 * 0.9;
        lowpass(signal, sfreq_in, cutoff, FilterMethod::Fft)?
    } else {
        signal.to_vec()
    };

    let n_in = filtered.len();
    let n_out = (n_in as f64 * ratio).round() as usize;
    let half_taps = num_taps / 2;

    let mut output = Vec::with_capacity(n_out);

    for i in 0..n_out {
        let src_idx = i as f64 / ratio;
        let src_int = src_idx.floor() as isize;

        let mut sum = 0.0;
        let mut weight_sum = 0.0;

        for j in -(half_taps as isize)..=(half_taps as isize) {
            let idx = src_int + j;
            if idx < 0 || idx >= n_in as isize {
                continue;
            }

            let t = src_idx - idx as f64;
            let weight = sinc(t) * blackman_window(t, half_taps as f64);

            sum += filtered[idx as usize] * weight;
            weight_sum += weight;
        }

        let value = if weight_sum.abs() > 1e-10 {
            sum / weight_sum
        } else {
            0.0
        };

        output.push(value);
    }

    Ok(output)
}

/// Normalized sinc function: sinc(x) = sin(pi*x) / (pi*x)
fn sinc(x: f64) -> f64 {
    if x.abs() < 1e-10 {
        1.0
    } else {
        let px = PI * x;
        px.sin() / px
    }
}

/// Blackman window function
fn blackman_window(x: f64, half_width: f64) -> f64 {
    if x.abs() > half_width {
        0.0
    } else {
        let t = x / half_width; // Normalize to [-1, 1]
        let theta = PI * (t + 1.0); // Map to [0, 2*pi]
        0.42 - 0.5 * theta.cos() + 0.08 * (2.0 * theta).cos()
    }
}

/// Decimate a signal by an integer factor.
///
/// More efficient than resample for integer factors.
pub fn decimate(signal: &[f64], sfreq: f64, factor: usize) -> SignalResult<Vec<f64>> {
    if factor == 0 {
        return Err(SignalError::InvalidParameters(
            "Decimation factor must be positive".to_string(),
        ));
    }

    if factor == 1 {
        return Ok(signal.to_vec());
    }

    // Apply anti-aliasing filter
    let cutoff = sfreq / factor as f64 / 2.0 * 0.9;
    let filtered = lowpass(signal, sfreq, cutoff, FilterMethod::Fft)?;

    // Decimate
    Ok(filtered.iter().step_by(factor).copied().collect())
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::f64::consts::PI;

    fn generate_sine(sfreq: f64, duration: f64, freq: f64) -> Vec<f64> {
        let n = (sfreq * duration) as usize;
        (0..n)
            .map(|i| (2.0 * PI * freq * i as f64 / sfreq).sin())
            .collect()
    }

    #[test]
    fn test_resample_downsample() {
        let signal = generate_sine(1000.0, 1.0, 10.0);
        let resampled = resample(&signal, 1000.0, 500.0).unwrap();

        assert_eq!(resampled.len(), 500);
    }

    #[test]
    fn test_resample_upsample() {
        let signal = generate_sine(500.0, 1.0, 10.0);
        let resampled = resample(&signal, 500.0, 1000.0).unwrap();

        assert_eq!(resampled.len(), 1000);
    }

    #[test]
    fn test_decimate() {
        let signal: Vec<f64> = (0..100).map(|x| x as f64).collect();
        let decimated = decimate(&signal, 1000.0, 2).unwrap();

        assert_eq!(decimated.len(), 50);
    }

    #[test]
    fn test_sinc() {
        assert!((sinc(0.0) - 1.0).abs() < 1e-10);
        assert!(sinc(1.0).abs() < 1e-10);
        assert!(sinc(2.0).abs() < 1e-10);
    }
}