zenith_float_num/
dsp.rs

1//! Discrete cosine / sine transforms, real-input FFT wrappers, and windows.
2
3use crate::defs::WORD_BIT_SIZE;
4use crate::Consts;
5use crate::ExactNum;
6use crate::ExactNumArray;
7use crate::RoundingMode;
8use alloc::vec::Vec;
9
10/// Maximum real length for [`dct`] / [`idct`] / [`dst`] / [`idst`] and the
11/// window generators. Transforms use a `2N`-point FFT, so this is half of
12/// `FFT_MAX_POINTS`.
13pub const DSP_MAX_POINTS: usize = 2048;
14
15fn work_p(p: usize) -> usize {
16    p.saturating_add(WORD_BIT_SIZE)
17}
18
19fn finite(x: &ExactNum) -> bool {
20    !x.is_nan() && !x.is_inf()
21}
22
23fn copy_p(x: &ExactNum, p: usize, rm: RoundingMode) -> ExactNum {
24    let mut y = x.clone();
25    let _ = y.set_precision(p, rm);
26    y
27}
28
29fn to_row(p: usize, rm: RoundingMode, vals: &[ExactNum]) -> Option<ExactNumArray> {
30    let rounded: Vec<ExactNum> = vals.iter().map(|x| copy_p(x, p, rm)).collect();
31    ExactNumArray::from_shape(p, 1, rounded.len(), &rounded)
32}
33
34/// Real vector from a `(1, n)` or `(n, 1)` array.
35fn as_real(signal: &ExactNumArray, p: usize, rm: RoundingMode) -> Option<Vec<ExactNum>> {
36    let (rows, cols) = signal.shape();
37    if !((rows == 1 && cols > 0) || (cols == 1 && rows > 0)) {
38        return None;
39    }
40    let mut out = Vec::with_capacity(signal.len());
41    for v in signal.as_slice() {
42        if !finite(v) {
43            return None;
44        }
45        out.push(copy_p(v, p, rm));
46    }
47    Some(out)
48}
49
50fn real_len_ok(n: usize) -> bool {
51    n > 0 && n.is_power_of_two() && n <= DSP_MAX_POINTS
52}
53
54fn pad_fft(
55    x: &[ExactNum],
56    n2: usize,
57    p: usize,
58    rm: RoundingMode,
59    cc: &mut Consts,
60) -> Option<ExactNumArray> {
61    let zero = ExactNum::from_u8(0, p);
62    let mut y = x.to_vec();
63    y.resize(n2, zero);
64    ExactNumArray::from_values(p, &y).fft(p, rm, cc)
65}
66
67/// Discrete cosine transform, type II, via a `2N`-point FFT.
68///
69/// `X_k = Σ_n x_n cos(π(n+1/2)k/N)`. Input is a real row or column. `N` must
70/// be a power of two and at most [`DSP_MAX_POINTS`]. Empty, non-finite, or a
71/// bad shape returns `None`.
72pub fn dct(
73    signal: &ExactNumArray,
74    p: usize,
75    rm: RoundingMode,
76    cc: &mut Consts,
77) -> Option<ExactNumArray> {
78    let wrk = work_p(p);
79    let x = as_real(signal, wrk, RoundingMode::None)?;
80    let n = x.len();
81    let n2 = n.checked_mul(2)?;
82    if !real_len_ok(n) {
83        return None;
84    }
85    let spec = pad_fft(&x, n2, wrk, RoundingMode::None, cc)?;
86    let pi = cc.pi(wrk, RoundingMode::None);
87    let two_n = ExactNum::from_u32(n2 as u32, wrk);
88    let mut out = Vec::with_capacity(n);
89    for k in 0..n {
90        let ang = pi
91            .mul(&ExactNum::from_u32(k as u32, wrk), wrk, RoundingMode::None)
92            .div(&two_n, wrk, RoundingMode::None);
93        let (sn, cs) = ang.sin_cos(wrk, RoundingMode::None, cc);
94        let yr = spec.get2(0, k)?;
95        let yi = spec.get2(1, k)?;
96        let re = yr.mul(&cs, wrk, RoundingMode::None).add(
97            &yi.mul(&sn, wrk, RoundingMode::None),
98            wrk,
99            RoundingMode::None,
100        );
101        out.push(re);
102    }
103    to_row(p, rm, &out)
104}
105
106/// Inverse of [`dct`]: type-III DCT scaled by `2/N` so `idct(dct(x)) = x`.
107///
108/// `x_n = X_0/N + (2/N) Σ_{k=1}^{N-1} X_k cos(π k (n+1/2)/N)`.
109pub fn idct(
110    signal: &ExactNumArray,
111    p: usize,
112    rm: RoundingMode,
113    cc: &mut Consts,
114) -> Option<ExactNumArray> {
115    let wrk = work_p(p);
116    let x = as_real(signal, wrk, RoundingMode::None)?;
117    let n = x.len();
118    if !real_len_ok(n) {
119        return None;
120    }
121    let none = RoundingMode::None;
122    let pi = cc.pi(wrk, none);
123    let nn = ExactNum::from_u32(n as u32, wrk);
124    let two = ExactNum::from_u8(2, wrk);
125    let inv_n = ExactNum::from_u8(1, wrk).div(&nn, wrk, none);
126    let two_n = two.div(&nn, wrk, none);
127    let half = ExactNum::from_u8(1, wrk).div(&two, wrk, none);
128    let mut out = Vec::with_capacity(n);
129    for ni in 0..n {
130        let n_half = ExactNum::from_u32(ni as u32, wrk).add(&half, wrk, none);
131        let mut acc = x[0].mul(&inv_n, wrk, none);
132        for k in 1..n {
133            let ang = pi
134                .mul(&ExactNum::from_u32(k as u32, wrk), wrk, none)
135                .mul(&n_half, wrk, none)
136                .div(&nn, wrk, none);
137            let c = ang.cos(wrk, none, cc);
138            acc = acc.add(&x[k].mul(&c, wrk, none).mul(&two_n, wrk, none), wrk, none);
139        }
140        if !finite(&acc) {
141            return None;
142        }
143        out.push(acc);
144    }
145    to_row(p, rm, &out)
146}
147
148/// Discrete sine transform, type II, via a `2N`-point FFT.
149///
150/// `X_k = Σ_n x_n sin(π(n+1/2)(k+1)/N)`. Same length rules as [`dct`].
151pub fn dst(
152    signal: &ExactNumArray,
153    p: usize,
154    rm: RoundingMode,
155    cc: &mut Consts,
156) -> Option<ExactNumArray> {
157    let wrk = work_p(p);
158    let x = as_real(signal, wrk, RoundingMode::None)?;
159    let n = x.len();
160    let n2 = n.checked_mul(2)?;
161    if !real_len_ok(n) {
162        return None;
163    }
164    let spec = pad_fft(&x, n2, wrk, RoundingMode::None, cc)?;
165    let pi = cc.pi(wrk, RoundingMode::None);
166    let two_n = ExactNum::from_u32(n2 as u32, wrk);
167    let mut out = Vec::with_capacity(n);
168    for k in 0..n {
169        let m = k + 1;
170        let ang = pi
171            .mul(&ExactNum::from_u32(m as u32, wrk), wrk, RoundingMode::None)
172            .div(&two_n, wrk, RoundingMode::None);
173        let (sn, cs) = ang.sin_cos(wrk, RoundingMode::None, cc);
174        let yr = spec.get2(0, m)?;
175        let yi = spec.get2(1, m)?;
176        // −Im(Y[m] exp(−iθ)) = yr sinθ − yi cosθ
177        let val = yr.mul(&sn, wrk, RoundingMode::None).sub(
178            &yi.mul(&cs, wrk, RoundingMode::None),
179            wrk,
180            RoundingMode::None,
181        );
182        out.push(val);
183    }
184    to_row(p, rm, &out)
185}
186
187/// Inverse of [`dst`]: type-III DST scaled by `2/N` so `idst(dst(x)) = x`.
188///
189/// `x_n = (2/N)[ (1/2)(−1)^n X_{N−1} + Σ_{k=0}^{N−2} X_k sin(π(k+1)(n+1/2)/N) ]`.
190pub fn idst(
191    signal: &ExactNumArray,
192    p: usize,
193    rm: RoundingMode,
194    cc: &mut Consts,
195) -> Option<ExactNumArray> {
196    let wrk = work_p(p);
197    let x = as_real(signal, wrk, RoundingMode::None)?;
198    let n = x.len();
199    if !real_len_ok(n) {
200        return None;
201    }
202    let none = RoundingMode::None;
203    let pi = cc.pi(wrk, none);
204    let nn = ExactNum::from_u32(n as u32, wrk);
205    let two = ExactNum::from_u8(2, wrk);
206    let two_n = two.div(&nn, wrk, none);
207    let half = ExactNum::from_u8(1, wrk).div(&two, wrk, none);
208    let last = x[n - 1].mul(&half, wrk, none);
209    let mut out = Vec::with_capacity(n);
210    for ni in 0..n {
211        let n_half = ExactNum::from_u32(ni as u32, wrk).add(&half, wrk, none);
212        let mut acc = if ni % 2 == 0 { last.clone() } else { last.neg() };
213        for k in 0..n - 1 {
214            let ang = pi
215                .mul(&ExactNum::from_u32((k + 1) as u32, wrk), wrk, none)
216                .mul(&n_half, wrk, none)
217                .div(&nn, wrk, none);
218            let s = ang.sin(wrk, none, cc);
219            acc = acc.add(&x[k].mul(&s, wrk, none), wrk, none);
220        }
221        acc = acc.mul(&two_n, wrk, none);
222        if !finite(&acc) {
223            return None;
224        }
225        out.push(acc);
226    }
227    to_row(p, rm, &out)
228}
229
230/// Unnormalized DFT of a real row or column. Output is `(2, n)` (row 0 real,
231/// row 1 imaginary). Same power-of-two rule as [`ExactNumArray::fft`].
232pub fn fft_real(
233    signal: &ExactNumArray,
234    p: usize,
235    rm: RoundingMode,
236    cc: &mut Consts,
237) -> Option<ExactNumArray> {
238    let _ = as_real(signal, p, rm)?;
239    signal.fft(p, rm, cc)
240}
241
242fn window_len_ok(n: usize) -> bool {
243    n > 0 && n <= DSP_MAX_POINTS
244}
245
246fn frac(num: u32, den: u32, p: usize, rm: RoundingMode) -> ExactNum {
247    ExactNum::from_u32(num, p).div(&ExactNum::from_u32(den, p), p, rm)
248}
249
250/// Symmetric cosine argument `2π k / (n−1)`. `n ≥ 2`.
251fn two_pi_k_over_nm1(k: usize, n: usize, p: usize, rm: RoundingMode, cc: &mut Consts) -> ExactNum {
252    let two = ExactNum::from_u8(2, p);
253    two.mul(&cc.pi(p, rm), p, rm)
254        .mul(&ExactNum::from_u32(k as u32, p), p, rm)
255        .div(&ExactNum::from_u32((n - 1) as u32, p), p, rm)
256}
257
258fn ones_row(n: usize, p: usize, rm: RoundingMode) -> Option<ExactNumArray> {
259    if !window_len_ok(n) {
260        return None;
261    }
262    let mut one = ExactNum::from_u8(1, p);
263    let _ = one.set_precision(p, rm);
264    Some(ExactNumArray::filled(p, n, &one))
265}
266
267/// Symmetric Hann: `½(1 − cos(2πk/(n−1)))`. `n = 1` is `[1]`.
268///
269/// The plan wrote `2πk/N` (periodic). The locked gold `hann_window(4) =
270/// [0, 3/4, 3/4, 0]` is the symmetric form.
271pub fn hann_window(n: usize, p: usize, rm: RoundingMode, cc: &mut Consts) -> Option<ExactNumArray> {
272    if n == 1 {
273        return ones_row(n, p, rm);
274    }
275    if !window_len_ok(n) {
276        return None;
277    }
278    let wrk = work_p(p);
279    let none = RoundingMode::None;
280    let half = frac(1, 2, wrk, none);
281    let one = ExactNum::from_u8(1, wrk);
282    let mut vals = Vec::with_capacity(n);
283    for k in 0..n {
284        let c = two_pi_k_over_nm1(k, n, wrk, none, cc).cos(wrk, none, cc);
285        let w = half.mul(&one.sub(&c, wrk, none), wrk, none);
286        if !finite(&w) {
287            return None;
288        }
289        vals.push(w);
290    }
291    to_row(p, rm, &vals)
292}
293
294/// Symmetric Hamming: `0.54 − 0.46 cos(2πk/(n−1))`. Endpoints are `0.08`.
295pub fn hamming_window(
296    n: usize,
297    p: usize,
298    rm: RoundingMode,
299    cc: &mut Consts,
300) -> Option<ExactNumArray> {
301    if n == 1 {
302        return ones_row(n, p, rm);
303    }
304    if !window_len_ok(n) {
305        return None;
306    }
307    let wrk = work_p(p);
308    let none = RoundingMode::None;
309    let a0 = frac(27, 50, wrk, none);
310    let a1 = frac(23, 50, wrk, none);
311    let mut vals = Vec::with_capacity(n);
312    for k in 0..n {
313        let c = two_pi_k_over_nm1(k, n, wrk, none, cc).cos(wrk, none, cc);
314        let w = a0.sub(&a1.mul(&c, wrk, none), wrk, none);
315        if !finite(&w) {
316            return None;
317        }
318        vals.push(w);
319    }
320    to_row(p, rm, &vals)
321}
322
323/// Symmetric Blackman: `0.42 − 0.5 cos(θ) + 0.08 cos(2θ)`, `θ = 2πk/(n−1)`.
324pub fn blackman_window(
325    n: usize,
326    p: usize,
327    rm: RoundingMode,
328    cc: &mut Consts,
329) -> Option<ExactNumArray> {
330    if n == 1 {
331        return ones_row(n, p, rm);
332    }
333    if !window_len_ok(n) {
334        return None;
335    }
336    let wrk = work_p(p);
337    let none = RoundingMode::None;
338    let a0 = frac(21, 50, wrk, none);
339    let a1 = frac(1, 2, wrk, none);
340    let a2 = frac(2, 25, wrk, none);
341    let two = ExactNum::from_u8(2, wrk);
342    let mut vals = Vec::with_capacity(n);
343    for k in 0..n {
344        let th = two_pi_k_over_nm1(k, n, wrk, none, cc);
345        let c1 = th.cos(wrk, none, cc);
346        let c2 = th.mul(&two, wrk, none).cos(wrk, none, cc);
347        let w = a0
348            .sub(&a1.mul(&c1, wrk, none), wrk, none)
349            .add(&a2.mul(&c2, wrk, none), wrk, none);
350        if !finite(&w) {
351            return None;
352        }
353        vals.push(w);
354    }
355    to_row(p, rm, &vals)
356}
357
358/// Kaiser–Bessel: `I_0(β √(1−t_k²)) / I_0(β)` with `t_k = (k−(n−1)/2)/((n−1)/2)`.
359///
360/// `beta = 0` is the rectangular window. `beta < 0` or non-finite is `None`.
361pub fn kaiser_window(
362    n: usize,
363    beta: &ExactNum,
364    p: usize,
365    rm: RoundingMode,
366    cc: &mut Consts,
367) -> Option<ExactNumArray> {
368    if !window_len_ok(n) || !finite(beta) || beta.is_negative() {
369        return None;
370    }
371    if n == 1 || beta.is_zero() {
372        return ones_row(n, p, rm);
373    }
374    let wrk = work_p(p);
375    let none = RoundingMode::None;
376    let b = copy_p(beta, wrk, none);
377    let nu0 = ExactNum::from_u8(0, wrk);
378    let i0b = b.bessel_i(&nu0, wrk, none, cc);
379    if !finite(&i0b) || i0b.is_zero() {
380        return None;
381    }
382    let two = ExactNum::from_u8(2, wrk);
383    let mid = ExactNum::from_u32((n - 1) as u32, wrk).div(&two, wrk, none);
384    let one = ExactNum::from_u8(1, wrk);
385    let mut vals = Vec::with_capacity(n);
386    for k in 0..n {
387        let t = ExactNum::from_u32(k as u32, wrk)
388            .sub(&mid, wrk, none)
389            .div(&mid, wrk, none);
390        let rad = one.sub(&t.mul(&t, wrk, none), wrk, none);
391        if rad.is_negative() {
392            return None;
393        }
394        let arg = b.mul(&rad.sqrt(wrk, none), wrk, none);
395        let w = arg.bessel_i(&nu0, wrk, none, cc).div(&i0b, wrk, none);
396        if !finite(&w) {
397            return None;
398        }
399        vals.push(w);
400    }
401    to_row(p, rm, &vals)
402}
403
404/// Rectangular window: `n` ones.
405pub fn rectangular_window(n: usize, p: usize, rm: RoundingMode) -> Option<ExactNumArray> {
406    ones_row(n, p, rm)
407}
408
409/// Inverse of [`fft_real`]: [`ExactNumArray::ifft`] then the real row.
410///
411/// Input must be a `(2, n)` spectrum. Imaginary residuals are dropped.
412pub fn ifft_real(
413    spectrum: &ExactNumArray,
414    p: usize,
415    rm: RoundingMode,
416    cc: &mut Consts,
417) -> Option<ExactNumArray> {
418    let (rows, cols) = spectrum.shape();
419    if rows != 2 || cols == 0 {
420        return None;
421    }
422    let spec = spectrum.ifft(p, rm, cc)?;
423    let mut re = Vec::with_capacity(cols);
424    for j in 0..cols {
425        re.push(spec.get2(0, j)?.clone());
426    }
427    to_row(p, rm, &re)
428}
429
430#[cfg(test)]
431mod tests {
432    use super::*;
433    use crate::Consts;
434
435    fn gold_p() -> (usize, RoundingMode) {
436        (256, RoundingMode::ToEven)
437    }
438
439    fn near(a: &ExactNum, b: &ExactNum, p: usize) -> bool {
440        let d = a.sub(b, p, RoundingMode::None).abs();
441        d.is_zero() || d.exponent().unwrap_or(0) < -((p as i32) - 40)
442    }
443
444    fn n_at(p: usize, k: u8) -> ExactNum {
445        ExactNum::from_u8(k, p)
446    }
447
448    #[test]
449    fn dsp_dct_dst_fft_real() {
450        let (p, rm) = gold_p();
451        let mut cc = Consts::new().expect("consts");
452        let n = |k: u8| n_at(p, k);
453        let zero = n(0);
454
455        let x = ExactNumArray::from_values(p, &[n(1), n(2), n(3), n(4)]);
456        let cx = dct(&x, p, rm, &mut cc).expect("dct");
457        let back = idct(&cx, p, rm, &mut cc).expect("idct");
458        assert_eq!(back.len(), 4);
459        for j in 0..4 {
460            assert!(
461                near(back.get(j).unwrap(), x.get(j).unwrap(), p),
462                "idct(dct(x))[{j}]"
463            );
464        }
465
466        let cst = ExactNumArray::from_values(p, &[n(3), n(3), n(3), n(3)]);
467        let dc = dct(&cst, p, rm, &mut cc).expect("dct const");
468        let twelve = ExactNum::from_u8(12, p);
469        assert!(near(dc.get(0).unwrap(), &twelve, p), "DC bin");
470        for j in 1..4 {
471            assert!(near(dc.get(j).unwrap(), &zero, p), "const bin {j}");
472        }
473
474        let sx = dst(&x, p, rm, &mut cc).expect("dst");
475        let sback = idst(&sx, p, rm, &mut cc).expect("idst");
476        for j in 0..4 {
477            assert!(
478                near(sback.get(j).unwrap(), x.get(j).unwrap(), p),
479                "idst(dst(x))[{j}]"
480            );
481        }
482
483        let n8 = ExactNum::from_u8(8, p);
484        let two_pi = n(2).mul(&cc.pi(p, rm), p, rm);
485        let mut cos_vals = Vec::with_capacity(8);
486        for k in 0..8u8 {
487            let kn = ExactNum::from_u8(k, p);
488            let ang = two_pi.mul(&kn, p, rm).div(&n8, p, rm);
489            cos_vals.push(ang.cos(p, rm, &mut cc));
490        }
491        let cosine = ExactNumArray::from_values(p, &cos_vals);
492        let cspec = fft_real(&cosine, p, rm, &mut cc).expect("fft_real cos");
493        assert_eq!(cspec.shape(), (2, 8));
494        let four = n(4);
495        for j in 0..8 {
496            let re = cspec.get2(0, j).unwrap();
497            let im = cspec.get2(1, j).unwrap();
498            if j == 1 || j == 7 {
499                assert!(near(re, &four, p), "cos bin {j} re");
500            } else {
501                assert!(near(re, &zero, p), "cos bin {j} re");
502            }
503            assert!(near(im, &zero, p), "cos bin {j} im");
504        }
505
506        let rec = ifft_real(&cspec, p, rm, &mut cc).expect("ifft_real");
507        for j in 0..8 {
508            assert!(near(rec.get(j).unwrap(), cosine.get(j).unwrap(), p));
509        }
510
511        let mut e_t = ExactNum::from_u8(0, p);
512        let mut e_f = ExactNum::from_u8(0, p);
513        for j in 0..8 {
514            let xv = cosine.get(j).unwrap();
515            e_t = e_t.add(&xv.mul(xv, p, rm), p, rm);
516            let xr = cspec.get2(0, j).unwrap();
517            let xi = cspec.get2(1, j).unwrap();
518            e_f = e_f
519                .add(&xr.mul(xr, p, rm), p, rm)
520                .add(&xi.mul(xi, p, rm), p, rm);
521        }
522        let parseval = e_f.div(&n8, p, rm);
523        assert!(near(&parseval, &e_t, p));
524
525        assert!(dct(
526            &ExactNumArray::from_values(p, &[n(1), n(2), n(3)]),
527            p,
528            rm,
529            &mut cc
530        )
531        .is_none());
532        assert!(fft_real(&cspec, p, rm, &mut cc).is_none());
533    }
534
535    fn window_sum_pos(w: &ExactNumArray, p: usize) -> bool {
536        let mut s = ExactNum::from_u8(0, p);
537        for i in 0..w.len() {
538            s = s.add(w.get(i).unwrap(), p, RoundingMode::ToEven);
539        }
540        s.is_positive() && !s.is_zero()
541    }
542
543    #[test]
544    fn dsp_windows() {
545        let (p, rm) = gold_p();
546        let mut cc = Consts::new().expect("consts");
547        let zero = ExactNum::from_u8(0, p);
548        let three_fourths = ExactNum::from_u8(3, p).div(&ExactNum::from_u8(4, p), p, rm);
549        let eight_hundredths = ExactNum::from_u8(2, p).div(&ExactNum::from_u8(25, p), p, rm);
550
551        let hann = hann_window(4, p, rm, &mut cc).expect("hann");
552        assert_eq!(hann.len(), 4);
553        assert!(near(hann.get(0).unwrap(), &zero, p));
554        assert!(near(hann.get(1).unwrap(), &three_fourths, p));
555        assert!(near(hann.get(2).unwrap(), &three_fourths, p));
556        assert!(near(hann.get(3).unwrap(), &zero, p));
557
558        let hamm = hamming_window(4, p, rm, &mut cc).expect("hamming");
559        assert!(near(hamm.get(0).unwrap(), &eight_hundredths, p));
560        assert!(near(hamm.get(3).unwrap(), &eight_hundredths, p));
561        assert!(!near(hamm.get(0).unwrap(), &zero, p));
562
563        let rect = rectangular_window(8, p, rm).expect("rect");
564        let k0 = kaiser_window(8, &zero, p, rm, &mut cc).expect("kaiser0");
565        assert_eq!(k0.len(), 8);
566        for j in 0..8 {
567            assert!(near(k0.get(j).unwrap(), rect.get(j).unwrap(), p));
568        }
569
570        let blk = blackman_window(8, p, rm, &mut cc).expect("blackman");
571        assert!(window_sum_pos(&hann, p));
572        assert!(window_sum_pos(&hamm, p));
573        assert!(window_sum_pos(&blk, p));
574        assert!(window_sum_pos(&rect, p));
575        assert!(window_sum_pos(&k0, p));
576
577        assert!(hann_window(0, p, rm, &mut cc).is_none());
578        assert!(rectangular_window(0, p, rm).is_none());
579        assert!(kaiser_window(4, &ExactNum::from_i64(-1, p), p, rm, &mut cc).is_none());
580    }
581}
zenith_float_num/dsp.rs

zenith_float_num/
dsp.rs
