spectrs 0.2.0

Fast spectrogram creation library
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193

//use clap::ValueEnum;
use rayon::prelude::*;
use rustfft::{FftPlanner, num_complex::Complex};
use std::f32::consts::PI;

// Different spectrogram types
#[derive(Debug, Clone, Copy)]
#[cfg_attr(feature = "cli", derive(clap::ValueEnum))]
pub enum SpectrogramType {
    Magnitude,
    Power,
}

/// Create Hann window, see e.g. https://en.wikipedia.org/wiki/Hann_function
fn create_hann_window(length: usize) -> Vec<f32> {
    (0..length)
        .map(|i| 0.5 * (1.0 - (2.0 * PI * i as f32 / (length - 1) as f32).cos()))
        .collect()
}

/// Compute the spectrogram (single-threaded)
/// n_samples: number of samples in each Fast Fourier Transform (FFT) window
/// hop_length: stride between windows, i.e. number of samples between successive FFT frames
/// win_length: number of samples in the window function applied before FFT
/// Pad with zeros if needed. This is because usually win_length < n_samples
/// and the missing are just zeros (in this case complex zeros)
pub fn compute_spectrogram(
    audio: &[f32],
    n_samples: usize,
    hop_length: usize,
    win_length: usize,
    center: bool,
    spectrogram_type: SpectrogramType,
) -> Vec<Vec<f32>> {
    // Set-up FFT
    let mut planner = FftPlanner::<f32>::new();
    let fft = planner.plan_fft_forward(n_samples);

    // Choose the transformation function to create the spectrogram
    let transform_fn: fn(&Complex<f32>) -> f32 = match spectrogram_type {
        SpectrogramType::Magnitude => |c| c.norm(),
        SpectrogramType::Power => |c| c.norm_sqr(),
    };

    // Create (Hann) window
    let window = create_hann_window(win_length);

    // Determine the number of frames
    let n_frames = (audio.len().saturating_sub(win_length)) / hop_length + 1;

    // Determine number of frequency bins
    let n_freq_bins = n_samples / 2 + 1;

    // Directly create spectrogram in [freq][time] format (no transpose needed)
    let mut spectrogram = vec![vec![0.0f32; n_frames]; n_freq_bins];

    // Sequential loop over frames
    for frame_idx in 0..n_frames {
        // Determine start and end sample for each frame
        let start = frame_idx * hop_length;
        let end = (start + win_length).clamp(0, audio.len());

        // Skip if start is beyond the end of the file
        if start > audio.len() {
            continue;
        }

        // Init buffer to be filled with windowed audio
        let mut frame = vec![Complex::<f32>::new(0.0, 0.0); n_samples];

        // Add an offset if the window needs to be centered
        let centering_offset = if center {
            (n_samples - win_length) / 2_usize
        } else {
            0_usize
        };

        // Window & copy into complex buffer
        let src = &audio[start..end];
        let win = &window[..src.len()];
        for (dst, (&s, &w)) in frame
            .iter_mut()
            .skip(centering_offset)
            .zip(src.iter().zip(win.iter()))
        {
            dst.re = s * w; // Convolve audio and window
            dst.im = 0.0; // No imaginary part
        }

        // Run FFT
        fft.process(&mut frame);

        // Store positive freqs only and apply transformation fn
        for (freq_idx, c) in frame.iter().take(n_freq_bins).enumerate() {
            spectrogram[freq_idx][frame_idx] = transform_fn(c);
        }
    }

    spectrogram
}

/// Compute the spectrogram (parallelized with rayon)
/// n_samples: number of samples in each Fast Fourier Transform (FFT) window
/// hop_length: stride between windows, i.e. number of samples between successive FFT frames
/// win_length: number of samples in the window function applied before FFT
/// Pad with zeros if needed. This is because usually win_length < n_samples
/// and the missing are just zeros (in this case complex zeros)
pub fn par_compute_spectrogram(
    audio: &[f32],
    n_samples: usize,
    hop_length: usize,
    win_length: usize,
    center: bool,
    spectrogram_type: SpectrogramType,
) -> Vec<Vec<f32>> {
    // Set-up FFT
    let mut planner = FftPlanner::<f32>::new();
    let fft = planner.plan_fft_forward(n_samples);

    // Choose the transformation function to create the spectrogram
    let transform_fn: fn(&Complex<f32>) -> f32 = match spectrogram_type {
        SpectrogramType::Magnitude => |c| c.norm(),
        SpectrogramType::Power => |c| c.norm_sqr(),
    };

    // Create (Hann) window
    let window = create_hann_window(win_length);

    // Determine the number of frames
    let n_frames = (audio.len().saturating_sub(win_length)) / hop_length + 1;

    // Determine number of frequency bins
    let n_freq_bins = n_samples / 2 + 1;

    // Frame-major spectrogram for safe parallel writes: spectrogram[frame][freq]
    // Eventually to be transposed
    let mut transposed_spectrogram = vec![vec![0.0f32; n_freq_bins]; n_frames];

    // Parallel loop over frames
    transposed_spectrogram
        .par_iter_mut() // Auto-parallelize with rayon
        .enumerate() // Extract frame idx
        .for_each(|(frame_idx, out_row)| {
            // Determine start and end sample for each frame, recalling that hop_length is a stride
            // If the end is after the end of the audio it might still be good (depending on start, see after)
            let start = frame_idx * hop_length;
            let end = (start + win_length).clamp(0, audio.len());

            // Start is beyond the end of the file
            if start > audio.len() {
                return;
            }

            // Init thread-local buffers to be filled with windowed audio
            let mut frame = vec![Complex::<f32>::new(0.0, 0.0); n_samples];

            // Add an offset if the window needs to be centered
            let centering_offset = if center {
                (n_samples - win_length) / 2_usize
            } else {
                0_usize
            };

            // Window & copy into complex buffer
            let src = &audio[start..end];
            let win = &window[..src.len()];
            for (dst, (&s, &w)) in frame
                .iter_mut()
                .skip(centering_offset)
                .zip(src.iter().zip(win.iter()))
            {
                dst.re = s * w; // Convolve audio and window
                dst.im = 0.0; // No imaginary part
            }

            // Run FFT
            fft.process(&mut frame);

            // Store positive freqs only and apply transformation fn depending on request
            for (k, c) in frame.iter().take(n_freq_bins).enumerate() {
                out_row[k] = transform_fn(c);
            }
        });

    // If your downstream expects [freq][frame], transpose once (cache-friendly)
    let mut spectrogram = vec![vec![0.0f32; n_frames]; n_freq_bins];
    for (t, row) in transposed_spectrogram.into_iter().enumerate() {
        for (f, v) in row.into_iter().enumerate() {
            spectrogram[f][t] = v;
        }
    }
    spectrogram
}