dspk_ops/resampler/
decimate.rs

1//! Integer down-sampling behind an anti-aliasing filter (`scipy.signal.decimate`, zero phase).
2
3use cubecl::prelude::*;
4use cubecl::server::Handle;
5use dspk_core::compute::LaunchGeometry;
6
7use super::design::{firwin, FirWindow};
8use super::kernels::downsample_kernel;
9use super::poly::{execute_resample_poly, ResampleFilter, RESAMPLE_POLY_DEFAULT_EDGE};
10use crate::core::{buffer, DspFloat};
11use crate::filter::iir::DeviceFilter;
12use crate::filter::{FilterBand, OpsError, FilterMode, FilterSpec};
13
14/// Anti-aliasing cutoff as a fraction of the new Nyquist frequency (`decimate`'s IIR design).
15pub const DECIMATE_IIR_CUTOFF: f64 = 0.8;
16
17/// FIR taps per unit of the factor (`decimate`'s FIR design: `20 · q + 1` taps).
18pub const DECIMATE_FIR_TAPS_PER_FACTOR: usize = 20;
19
20/// Anti-aliasing filter of [`execute_decimate`].
21#[derive(Debug, Clone, Copy, PartialEq)]
22pub enum DecimateFilter {
23    /// Chebyshev type I low-pass at [`DECIMATE_IIR_CUTOFF`] of the new Nyquist, run forward-backward,
24    /// then every `q`-th sample (`ftype="iir"`).
25    Iir { order: usize, ripple_db: f64 },
26    /// Hamming `firwin` of `taps_per_factor · q + 1` taps through [`execute_resample_poly`] (`ftype="fir"`).
27    Fir { taps_per_factor: usize },
28}
29
30/// `decimate`'s default: an order-8 Chebyshev type I with 0.05 dB ripple.
31pub const DECIMATE_DEFAULT: DecimateFilter = DecimateFilter::Iir { order: 8, ripple_db: 0.05 };
32
33/// Samples out of `samples` decimated by `q` (`ceil(samples / q)`).
34pub fn decimate_len(samples: usize, q: usize) -> usize {
35    samples.div_ceil(q.max(1))
36}
37
38/// Down-samples every channel of a `[channels, samples]` buffer of `F` by `q` into `output`
39/// (`[channels, decimate_len(samples, q)]`) behind `filter`. Returns the output length.
40///
41/// The IIR path pads with odd reflection over the filter's settling length (where `sosfiltfilt`
42/// pads a fixed `3·(2·sections + 1)`), so samples near the ends can differ slightly from scipy.
43///
44/// # Panics
45/// If `q` is zero.
46///
47/// # Errors
48///
49/// [`OpsError`] when the anti-aliasing filter's design is invalid at the input rate.
50#[allow(clippy::too_many_arguments)]
51pub fn execute_decimate<F: DspFloat>(
52    client: &Client,
53    input: &Handle,
54    output: &Handle,
55    channels: usize,
56    samples: usize,
57    q: usize,
58    filter: DecimateFilter,
59) -> Result<usize, OpsError> {
60    assert!(q > 0, "decimation factor must be positive");
61    let out_len = decimate_len(samples, q);
62    if channels == 0 || samples == 0 {
63        return Ok(out_len);
64    }
65    match filter {
66        DecimateFilter::Iir { order, ripple_db } => {
67            // Designed on scipy's normalized grid: fs = 2, so the new Nyquist is 1 / q
68            let spec = FilterSpec::chebyshev1(order, ripple_db, FilterBand::Lowpass(DECIMATE_IIR_CUTOFF / q as f64))
69                .with_mode(FilterMode::ForwardBackward);
70            let device = DeviceFilter::<F>::new(client, &spec, 2.0)?;
71            let filtered = buffer::empty::<F>(client, channels * samples);
72            let scratch = buffer::empty::<F>(client, device.scratch_len(channels, samples));
73            let state = buffer::empty::<F>(client, channels * device.state_len());
74            device.apply(client, input, &filtered, &scratch, &state, channels, samples);
75
76            let geom = LaunchGeometry::channels_samples(client, channels, out_len);
77            unsafe {
78                downsample_kernel::launch::<F>(
79                    client,
80                    geom.cube_count,
81                    geom.cube_dim,
82                    BufferArg::from_raw_parts(filtered, channels * samples),
83                    BufferArg::from_raw_parts(output.clone(), channels * out_len),
84                    channels as u32,
85                    samples as u32,
86                    out_len as u32,
87                    q as u32,
88                    0u32,
89                );
90            }
91            Ok(out_len)
92        }
93        DecimateFilter::Fir { taps_per_factor } => {
94            let taps = firwin(taps_per_factor * q + 1, 1.0 / q as f64, FirWindow::Hamming);
95            let filter = ResampleFilter::Taps(taps);
96            Ok(execute_resample_poly::<F>(client, input, output, channels, samples, 1, q, &filter, RESAMPLE_POLY_DEFAULT_EDGE))
97        }
98    }
99}
100
101#[cfg(test)]
102mod tests {
103    use super::*;
104
105    fn iir_matches_host(client: &Client) {
106        let (samples, q) = (5_000usize, 4usize);
107        let x: Vec<f64> = (0..samples).map(|i| (i as f64 * 0.01).sin() * 50.0 + (i as f64 * 1.9).sin() * 5.0 - 20.0).collect();
108        let input = buffer::upload(client, &x.iter().map(|v| *v as f32).collect::<Vec<_>>());
109        let out_len = decimate_len(samples, q);
110        let output = buffer::empty::<f32>(client, out_len);
111        assert_eq!(execute_decimate::<f32>(client, &input, &output, 1, samples, q, DECIMATE_DEFAULT).unwrap(), out_len);
112        let got = buffer::download::<f32>(client, output);
113
114        let DecimateFilter::Iir { order, ripple_db } = DECIMATE_DEFAULT else { unreachable!() };
115        let sos = crate::filter::design::chebyshev1_sos(order, ripple_db, FilterBand::Lowpass(DECIMATE_IIR_CUTOFF / q as f64), 2.0).unwrap();
116        let filtered = sos.filtfilt(&x, sos.settling_samples(crate::filter::design::DEFAULT_SETTLING_TOLERANCE).min(samples - 1));
117        for (m, g) in got.iter().enumerate() {
118            let w = filtered[m * q];
119            assert!((*g as f64 - w).abs() < 2e-3 * 70.0, "{} sample {m}: {g} vs {w}", client.name());
120        }
121    }
122    runtime_test!(test_decimate_iir_matches_host, iir_matches_host);
123}
dspk_ops/resampler/decimate.rs

dspk_ops/resampler/
decimate.rs
