1use crate::{Complex64, DspError, invalid};
6
7#[derive(Debug, Clone, Copy)]
9pub enum Fraction<'a> {
10 Scalar(f64),
11 PerBin(&'a [f64]),
12}
13
14pub 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#[derive(Debug)]
40pub struct PowerSmoother {
41 bounds: Vec<(usize, usize)>,
42}
43
44impl PowerSmoother {
45 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 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 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
127pub(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}