Skip to main content

secs_dsp/
smooth.rs

1//! Adapted from Impulcifer-py313 `crates/impulcifer-dsp/src/stages/room_tuning2/primitives.rs`
2//! (MIT, Copyright (c) 2024- 115dkk), corrected against SECS 2/SciPy.
3//! See docs/THIRD_PARTY.md and ../THIRD_PARTY_LICENSES.txt.
4//! Fractional-octave smoothing, including upstream's sequential cumsum.
5use crate::{Complex64, DspError, invalid};
6
7/// Scalar or per-bin fractional-octave denominator (24 means 1/24 octave).
8#[derive(Debug, Clone, Copy)]
9pub enum Fraction<'a> {
10    Scalar(f64),
11    PerBin(&'a [f64]),
12}
13
14/// Linear-grid mean over inclusive floor/ceil bounds. Shape and upstream
15/// out-of-bounds errors return Err (not silently clamped at Nyquist).
16pub fn smooth_power(f: &[f64], power: &[f64], frac: Fraction<'_>) -> Result<Vec<f64>, DspError> {
17    if f.len() < 2 || f.len() != power.len() {
18        return Err(invalid(
19            "smoothing requires matching arrays of at least two bins",
20        ));
21    }
22    if let Fraction::PerBin(x) = frac
23        && x.len() != f.len()
24        && x.len() != 1
25    {
26        return Err(invalid("fraction cannot broadcast to frequency grid"));
27    }
28    let mut c = Vec::with_capacity(power.len() + 1);
29    c.push(0.0);
30    for p in power {
31        c.push(c.last().unwrap() + p);
32    }
33    bounds(f, frac)
34        .map(|bound| bound.map(|(lo, hi)| (c[hi] - c[lo]) / (hi - lo) as f64))
35        .collect()
36}
37
38/// Reusable bin bounds for repeated smoothing on the same frequency grid.
39#[derive(Debug)]
40pub struct PowerSmoother {
41    bounds: Vec<(usize, usize)>,
42}
43
44impl PowerSmoother {
45    /// Records the original inclusive floor/ceil bounds, rejecting the same
46    /// invalid broadcast shapes and out-of-grid upper bounds as `smooth_power`.
47    pub fn new(f: &[f64], frac: Fraction<'_>) -> Result<Self, DspError> {
48        if f.len() < 2 {
49            return Err(invalid("smoothing requires at least two bins"));
50        }
51        if let Fraction::PerBin(x) = frac
52            && x.len() != f.len()
53            && x.len() != 1
54        {
55            return Err(invalid("fraction cannot broadcast to frequency grid"));
56        }
57        Ok(Self {
58            bounds: bounds(f, frac).collect::<Result<_, _>>()?,
59        })
60    }
61
62    /// Smooths a matching power array with the original sequential prefix sum.
63    pub fn apply(&self, power: &[f64]) -> Result<Vec<f64>, DspError> {
64        if power.len() != self.bounds.len() {
65            return Err(invalid("smoothing requires matching arrays"));
66        }
67        let mut c = Vec::with_capacity(power.len() + 1);
68        c.push(0.0);
69        for p in power {
70            c.push(c.last().unwrap() + p);
71        }
72        Ok(self
73            .bounds
74            .iter()
75            .map(|&(lo, hi)| (c[hi] - c[lo]) / (hi - lo) as f64)
76            .collect())
77    }
78}
79
80fn bounds<'a>(
81    f: &'a [f64],
82    frac: Fraction<'a>,
83) -> impl Iterator<Item = Result<(usize, usize), DspError>> + 'a {
84    let df = f[1] - f[0];
85    // NumPy float->int casts NaN/inf/out-of-range to INT64_MIN, then clips.
86    let index = move |x: f64| -> usize {
87        if !x.is_finite() || x >= i64::MAX as f64 || x < i64::MIN as f64 {
88            0
89        } else {
90            (x as i64).clamp(0, (f.len() - 1) as i64) as usize
91        }
92    };
93    let scalar_h = match frac {
94        Fraction::Scalar(b) => Some(2_f64.powf(0.5 / b)),
95        Fraction::PerBin(b) if b.len() == 1 => Some(2_f64.powf(0.5 / b[0])),
96        Fraction::PerBin(_) => None,
97    };
98    let f_len = f.len();
99    f.iter().enumerate().map(move |(i, &f)| {
100        let h = scalar_h.unwrap_or_else(|| match frac {
101            Fraction::PerBin(x) => 2_f64.powf(0.5 / x[i]),
102            Fraction::Scalar(_) => unreachable!(),
103        });
104        let lo = index((f / h / df).floor());
105        let hi = index((f * h / df).ceil()).max(lo + 1);
106        if hi >= f_len {
107            return Err(invalid("smoothing upper bound exceeds cumulative sum"));
108        }
109        Ok((lo, hi + 1))
110    })
111}
112
113pub fn smooth_complex(
114    f: &[f64],
115    h: &[Complex64],
116    frac: Fraction<'_>,
117) -> Result<Vec<Complex64>, DspError> {
118    let re = smooth_power(f, &h.iter().map(|h| h.re).collect::<Vec<_>>(), frac)?;
119    let im = smooth_power(f, &h.iter().map(|h| h.im).collect::<Vec<_>>(), frac)?;
120    Ok(re
121        .into_iter()
122        .zip(im)
123        .map(|(re, im)| re + Complex64::new(0., 1.) * Complex64::new(im, 0.))
124        .collect())
125}
126
127/// Smooths only an output prefix, evaluating input bins lazily. Every bound is
128/// still validated against the original grid. Prefix sums start at zero and
129/// retain exactly the accumulation order of `smooth_complex`, including NaNs.
130pub(crate) fn smooth_complex_prefix(
131    f: &[f64],
132    frac: Fraction<'_>,
133    output_len: usize,
134    mut value: impl FnMut(usize) -> Complex64,
135) -> Result<Vec<Complex64>, DspError> {
136    if f.len() < 2 || output_len > f.len() {
137        return Err(invalid(
138            "smoothing requires matching arrays of at least two bins",
139        ));
140    }
141    if let Fraction::PerBin(x) = frac
142        && x.len() != f.len()
143        && x.len() != 1
144    {
145        return Err(invalid("fraction cannot broadcast to frequency grid"));
146    }
147    let mut selected = Vec::with_capacity(output_len);
148    let mut input_len = 0;
149    for (i, bound) in bounds(f, frac).enumerate() {
150        let (lo, hi) = bound?;
151        if i < output_len {
152            selected.push((lo, hi));
153            input_len = input_len.max(hi);
154        }
155    }
156    let mut re = Vec::with_capacity(input_len + 1);
157    let mut im = Vec::with_capacity(input_len + 1);
158    re.push(0.0);
159    im.push(0.0);
160    for i in 0..input_len {
161        let h = value(i);
162        re.push(re.last().unwrap() + h.re);
163        im.push(im.last().unwrap() + h.im);
164    }
165    Ok(selected
166        .into_iter()
167        .map(|(lo, hi)| {
168            let real = (re[hi] - re[lo]) / (hi - lo) as f64;
169            let imaginary = (im[hi] - im[lo]) / (hi - lo) as f64;
170            real + Complex64::new(0.0, 1.0) * Complex64::new(imaginary, 0.0)
171        })
172        .collect())
173}
174
175#[cfg(test)]
176mod tests {
177    use super::*;
178
179    #[test]
180    fn cached_bounds_match_smoothing_bits_and_errors() {
181        let f: Vec<_> = (0..1024).map(|i| i as f64 * 0.366).collect();
182        let per_bin: Vec<_> = f.iter().map(|f| 6.0 + f / 100.0).collect();
183        for fraction in [
184            Fraction::Scalar(12.0),
185            Fraction::PerBin(&[24.0]),
186            Fraction::PerBin(&per_bin),
187        ] {
188            let cache = PowerSmoother::new(&f, fraction).unwrap();
189            for offset in [0.0, 0.1, -0.1] {
190                let power: Vec<_> = f.iter().map(|f| (f + offset).sin()).collect();
191                let expected = smooth_power(&f, &power, fraction).unwrap();
192                let actual = cache.apply(&power).unwrap();
193                assert!(
194                    actual
195                        .iter()
196                        .zip(expected)
197                        .all(|(a, b)| a.to_bits() == b.to_bits())
198                );
199            }
200            assert!(cache.apply(&[0.0]).is_err());
201        }
202        assert!(PowerSmoother::new(&[], Fraction::Scalar(12.0)).is_err());
203        assert!(PowerSmoother::new(&f, Fraction::PerBin(&[])).is_err());
204        assert!(PowerSmoother::new(&f, Fraction::Scalar(-12.0)).is_err());
205    }
206
207    #[test]
208    fn selected_complex_smoothing_preserves_prefix_bits_and_all_grid_errors() {
209        let f: Vec<_> = (0..4097).map(|i| i as f64 * 0.3662109375).collect();
210        let h: Vec<_> = f
211            .iter()
212            .map(|f| Complex64::new((f * 0.013).sin(), (f * 0.017).cos()))
213            .collect();
214        let per_bin: Vec<_> = f.iter().map(|f| 6.0 + f / 50.0).collect();
215        for frac in [
216            Fraction::Scalar(24.0),
217            Fraction::PerBin(&[24.0]),
218            Fraction::PerBin(&per_bin),
219        ] {
220            let expected = smooth_complex(&f, &h, frac).unwrap();
221            let mut calls = 0;
222            let actual = smooth_complex_prefix(&f, frac, 100, |i| {
223                calls += 1;
224                h[i]
225            })
226            .unwrap();
227            assert!(calls < 150, "unneeded high-frequency bins were evaluated");
228            for (a, b) in actual.iter().zip(&expected) {
229                assert_eq!(a.re.to_bits(), b.re.to_bits());
230                assert_eq!(a.im.to_bits(), b.im.to_bits());
231            }
232        }
233        let mut special = h.clone();
234        special[30] = Complex64::new(f64::INFINITY, f64::NAN);
235        let full = smooth_complex(&f, &special, Fraction::Scalar(24.0)).unwrap();
236        let prefix =
237            smooth_complex_prefix(&f, Fraction::Scalar(24.0), 100, |i| special[i]).unwrap();
238        for (a, b) in prefix.iter().zip(full) {
239            assert!(a.re.to_bits() == b.re.to_bits() || a.re.is_nan() && b.re.is_nan());
240            assert!(a.im.to_bits() == b.im.to_bits() || a.im.is_nan() && b.im.is_nan());
241        }
242        let mut invalid = vec![24.0; f.len()];
243        invalid[f.len() - 1] = -24.0;
244        assert!(smooth_complex_prefix(&f, Fraction::PerBin(&invalid), 10, |i| h[i]).is_err());
245        assert!(smooth_complex_prefix(&f, Fraction::PerBin(&[]), 10, |i| h[i]).is_err());
246    }
247}