pitch_detector/core/
utils.rs

1use fitting::gaussian::fit;
2
3use super::FftPoint;
4
5pub fn sine_wave_signal(num_samples: usize, freq: f64, sample_rate: f64) -> Vec<f64> {
6    (0..num_samples)
7        .map(|r| (2.0 * std::f64::consts::PI * r as f64 * freq / sample_rate).sin())
8        .collect()
9}
10
11pub fn mixed_wave_signal(num_samples: usize, freqs: Vec<f64>, sample_rate: f64) -> Vec<f64> {
12    let mut signal = vec![0.0; num_samples];
13    freqs.iter().for_each(|f| {
14        let s = sine_wave_signal(num_samples, *f, sample_rate);
15        signal.iter_mut().zip(s.iter()).for_each(|(a, b)| *a += *b);
16    });
17    signal
18}
19
20pub fn audio_buffer_to_samples(byte_buffer: &[u8]) -> Box<dyn Iterator<Item = i16> + '_> {
21    Box::new(
22        byte_buffer
23            .chunks_exact(2)
24            .map(|a| i16::from_ne_bytes([a[0], a[1]])),
25    )
26}
27pub fn audio_buffer_to_signal(byte_buffer: &[u8]) -> Box<dyn Iterator<Item = f64> + '_> {
28    Box::new(audio_buffer_to_samples(byte_buffer).map(|x| x as f64))
29}
30
31pub fn interpolated_peak_at(spectrum: &[f64], fft_point_x: usize) -> Option<FftPoint> {
32    let mut idx = fft_point_x;
33    let peak_begin_idx = loop {
34        if idx == 0 {
35            break idx;
36        }
37        if spectrum[idx] < spectrum[idx - 1] || spectrum[idx - 1] <= 0. {
38            break idx;
39        }
40        idx -= 1;
41    };
42    idx = fft_point_x;
43    let peak_end_idx_incl = loop {
44        if idx == spectrum.len() - 1 {
45            break idx;
46        }
47        if spectrum[idx] < spectrum[idx + 1] || spectrum[idx + 1] <= 0. {
48            break idx;
49        }
50        idx += 1;
51    };
52    let y_vals: Vec<f64> = spectrum[peak_begin_idx..=peak_end_idx_incl]
53        .iter()
54        .cloned()
55        .collect();
56    let x_vals: Vec<f64> = (peak_begin_idx..=peak_end_idx_incl)
57        .map(|i| i as f64)
58        .collect();
59
60    assert_eq!(y_vals.len(), x_vals.len(), "Sanity check failed");
61    match x_vals.len() {
62        0 => None,
63        1 => Some(FftPoint {
64            x: x_vals[0],
65            y: y_vals[0],
66        }),
67        2 => {
68            if y_vals[0] > y_vals[1] {
69                Some(FftPoint {
70                    x: x_vals[0],
71                    y: y_vals[0],
72                })
73            } else {
74                Some(FftPoint {
75                    x: x_vals[1],
76                    y: y_vals[1],
77                })
78            }
79        }
80        _ => {
81            let (mu, _, a) = fit(x_vals.into(), y_vals.into()).ok()?;
82            Some(FftPoint { x: mu, y: a })
83        }
84    }
85}