Skip to main content

gpu_fft/
utils.rs

1use std::f32::consts::PI;
2
3/// Generates a sine wave signal based on the specified frequency, sample rate, and duration.
4///
5/// # Parameters
6///
7/// - `frequency`: The frequency of the sine wave in Hertz (Hz).
8/// - `sample_rate`: The number of samples per second.
9/// - `duration`: The duration of the sine wave in seconds.
10///
11/// # Returns
12///
13/// A vector of `f32` samples. Its length equals `(sample_rate * duration) as usize`.
14///
15/// # Example
16///
17/// ```
18/// # use gpu_fft::utils::generate_sine_wave;
19/// let frequency = 440.0f32;    // A4 note
20/// let sample_rate = 44100.0f32; // CD quality
21/// let duration = 1.0f32;        // 1 second
22/// let sine_wave = generate_sine_wave(frequency, sample_rate, duration);
23/// assert_eq!(sine_wave.len(), 44100);
24/// ```
25#[must_use]
26pub fn generate_sine_wave(frequency: f32, sample_rate: f32, duration: f32) -> Vec<f32> {
27    let num_samples = (sample_rate * duration) as usize;
28    (0..num_samples)
29        .map(|n| (2.0 * PI * frequency * n as f32 / sample_rate).sin())
30        .collect()
31}
32
33/// Returns the frequency (in Hz) corresponding to each bin of a full (two-sided) DFT output.
34///
35/// Bin `k` maps to `k * sample_rate / n` Hz.  The upper half of the returned frequencies
36/// (`> sample_rate / 2`) represent negative frequencies — they are the conjugate mirrors
37/// of the lower half and carry no additional information for real-valued signals.
38///
39/// For real signals, prefer [`calculate_one_sided_frequencies`] instead.
40///
41/// # Example
42///
43/// ```
44/// # use gpu_fft::utils::calculate_frequencies;
45/// let frequencies = calculate_frequencies(1024, 44100.0);
46/// assert_eq!(frequencies.len(), 1024);
47/// assert_eq!(frequencies[0], 0.0);
48/// ```
49#[must_use]
50pub fn calculate_frequencies(n: usize, sample_rate: f32) -> Vec<f32> {
51    (0..n).map(|k| k as f32 * sample_rate / n as f32).collect()
52}
53
54/// Returns the `n_total / 2 + 1` unique positive-frequency bins (0 Hz … Nyquist) for a
55/// real-valued DFT of `n_total` samples at `sample_rate` Hz.
56///
57/// For a real input the DFT spectrum is conjugate-symmetric, so only the first half plus
58/// the DC and Nyquist bins are unique.  Use this together with slicing the PSD to the
59/// same length to avoid spurious mirror-image peaks.
60///
61/// # Example
62///
63/// ```
64/// # use gpu_fft::utils::calculate_one_sided_frequencies;
65/// let freqs = calculate_one_sided_frequencies(1000, 200.0);
66/// assert_eq!(freqs.len(), 501);       // n/2 + 1
67/// assert_eq!(freqs[0], 0.0);          // DC
68/// assert!((freqs[500] - 100.0).abs() < 1e-4); // Nyquist = sample_rate / 2
69/// ```
70#[must_use]
71pub fn calculate_one_sided_frequencies(n_total: usize, sample_rate: f32) -> Vec<f32> {
72    (0..=n_total / 2)
73        .map(|k| k as f32 * sample_rate / n_total as f32)
74        .collect()
75}
76
77/// Finds the dominant frequencies in a Power Spectral Density (PSD) by looking for local
78/// peaks above a threshold.
79///
80/// A peak is a bin whose value exceeds both immediate neighbours and the threshold.
81/// The first and last bins are never reported (they cannot be local peaks).
82///
83/// For real-valued signals, pass only the **one-sided** PSD (first `n/2 + 1` bins) and the
84/// matching frequencies from [`calculate_one_sided_frequencies`] to avoid spurious
85/// mirror-image peaks in the upper half of the spectrum.
86///
87/// # Example
88///
89/// ```
90/// # use gpu_fft::utils::find_dominant_frequencies;
91/// let psd = vec![0.1f32, 0.5, 0.3, 0.7, 0.2];
92/// let frequencies = vec![0.0f32, 100.0, 200.0, 300.0, 400.0];
93/// let dominant = find_dominant_frequencies(&psd, &frequencies, 0.4);
94/// // Bins 1 (100 Hz) and 3 (300 Hz) are local peaks above the threshold.
95/// assert_eq!(dominant.len(), 2);
96/// assert_eq!(dominant[0].0, 100.0);
97/// assert_eq!(dominant[1].0, 300.0);
98/// ```
99#[must_use]
100pub fn find_dominant_frequencies(
101    psd: &[f32],
102    frequencies: &[f32],
103    threshold: f32,
104) -> Vec<(f32, f32)> {
105    assert_eq!(psd.len(), frequencies.len(), "psd and frequencies must have the same length");
106    (1..psd.len().saturating_sub(1))
107        .filter(|&i| psd[i] > psd[i - 1] && psd[i] > psd[i + 1] && psd[i] > threshold)
108        .map(|i| (frequencies[i], psd[i]))
109        .collect()
110}