aus 0.1.8

A library of audio processing tools
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
194
195
196
197
198
199
200
201
202
203
204
205
206

// File: fft_tools.rs
// This file contains tools for FFT processing, such as decomposing the spectrum
// into magnitude and phase spectra.

use rustfft::num_complex::Complex;
use super::fft::SpectrumError;

/// This function decomposes a complex spectrum into magnitude and phase spectra. 
/// 
/// # Example
/// 
/// ```
/// use aus::spectrum;
/// let fft_size: usize = 2048;
/// let audio = aus::read("myfile.wav").unwrap();
/// let window = aus::generate_window_hanning(fft_size);
/// // Just choose the first 2048 samples in the audio file. This might be a problem, because those samples might be zeros.
/// let audio_chunk: Vec<f64> = audio.samples[0][..fft_size].iter().zip(window.iter()).map(|(a, b)| a * b).collect();
/// let (mut magnitude_spectrum, mut phase_spectrum) = spectrum::complex_to_polar_rfft(&spectrum::rfft(&audio_chunk, fft_size));
/// ```
pub fn complex_to_polar_rfft(spectrum: &[Complex<f64>]) -> (Vec<f64>, Vec<f64>) {
    let mut magnitude_spectrum = vec![0.0 as f64; spectrum.len()];
    let mut phase_spectrum = vec![0.0 as f64; spectrum.len()];
    for i in 0..spectrum.len() {
        magnitude_spectrum[i] = f64::sqrt(spectrum[i].re.powf(2.0) + spectrum[i].im.powf(2.0));
        phase_spectrum[i] = f64::atan2(spectrum[i].im, spectrum[i].re);
    }
    (magnitude_spectrum, phase_spectrum)
}

/// This function decomposes a complex spectrogram into magnitude and phase spectrograms.
/// 
/// # Example
/// 
/// ```
/// use aus::spectrum;
/// let file = aus::read("myfile.wav").unwrap();
/// let mut imaginary_spectrogram = spectrum::rstft(&file.samples[0], 2048, 1024, aus::WindowType::Hanning);
/// let (mut magnitude_spectrogram, mut phase_spectrogram) = spectrum::complex_to_polar_rstft(&imaginary_spectrogram);
/// ```
pub fn complex_to_polar_rstft(spectrogram: &Vec<Vec<Complex<f64>>>) -> (Vec<Vec<f64>>, Vec<Vec<f64>>) {
    let mut magnitude_spectrogram: Vec<Vec<f64>> = Vec::with_capacity(spectrogram.len());
    let mut phase_spectrogram: Vec<Vec<f64>> = Vec::with_capacity(spectrogram.len());
    for frame_idx in 0..spectrogram.len() {
        let mut frame_magnitude_spectrum = vec![0.0 as f64; spectrogram[frame_idx].len()];
        let mut frame_phase_spectrum = vec![0.0 as f64; spectrogram[frame_idx].len()];
        for fft_bin_idx in 0..spectrogram[frame_idx].len() {
            frame_magnitude_spectrum[fft_bin_idx] = f64::sqrt(spectrogram[frame_idx][fft_bin_idx].re.powf(2.0) + spectrogram[frame_idx][fft_bin_idx].im.powf(2.0));
            frame_phase_spectrum[fft_bin_idx] = f64::atan2(spectrogram[frame_idx][fft_bin_idx].im, spectrogram[frame_idx][fft_bin_idx].re);    
        }
        magnitude_spectrogram.push(frame_magnitude_spectrum);
        phase_spectrogram.push(frame_phase_spectrum);
    }    
    (magnitude_spectrogram, phase_spectrogram)
}

/// An efficient overlap add mechanism.
/// 
/// The formula for overlap add in mathematical definitions of the ISTFT requires checking
/// each of the *M* audio chunks generated by the *M* IRFFT operations. For short audio, this
/// is not particularly costly, but as FFT sizes decrease and audio length increases, it
/// becomes wasteful.
/// 
/// This algorithm keeps track of which audio chunks are relevant to computing the current 
/// sample, so that only those chunks are consulted. It does this by using running indices
/// for the lowest and highest audio chunk indices that are currently relevant.
/// 
/// The algorithm assumes that the hop size is greater than 0.
pub fn overlap_add(audio_chunks: &Vec<Vec<f64>>, fft_size: usize, hop_size: usize) -> Vec<f64> {
    let mut audio: Vec<f64> = Vec::new();

    // Get the global start and end index corresponding to each audio frame
    let mut frame_indices: Vec<(usize, usize)> = Vec::with_capacity(audio_chunks.len());
    let mut current_frame_start_idx: usize = 0;
    for _ in 0..audio_chunks.len() {
        frame_indices.push((current_frame_start_idx, current_frame_start_idx + fft_size));
        current_frame_start_idx += hop_size;
    }

    let mut lower_frame_idx: usize = 0;  // The index of the lowest frame we are adding
    let mut upper_frame_idx: usize = 0;  // The index of the highest frame we are adding
    let mut current_sample_idx: usize = 0;  // The index of the current sample to compute

    // Overlap add
    while lower_frame_idx < audio_chunks.len() {
        // If we've moved beyond the range of the lower frame, we need to move the lower frame index up
        if current_sample_idx >= frame_indices[lower_frame_idx].1 {
            lower_frame_idx += 1;
        }

        // If we've moved into the range of a new upper frame, we need to adjust the upper frame index
        if upper_frame_idx + 1 < audio_chunks.len() {
            if current_sample_idx >= frame_indices[upper_frame_idx + 1].0 {
                upper_frame_idx += 1;
            }
        }

        // Check to make sure the lower frame index is still valid (i.e. we haven't gone beyond the end of the audio)
        if lower_frame_idx < audio_chunks.len() {
            // Build the sample using only the valid frames
            let mut sample: f64 = 0.0;
            for i in lower_frame_idx..upper_frame_idx + 1 {
                let local_frame_idx = current_sample_idx - frame_indices[i].0;
                sample += audio_chunks[i][local_frame_idx];
            }
            audio.push(sample);
            current_sample_idx += 1;
        }
    }

    audio
}

/// This function combines magnitude and phase spectra into a complex spectrum for the inverse real FFT.
/// The magnitude and phase spectra must have the same length. If they do not have the same length,
/// this function will return a SpectrumError.
/// 
/// # Example
/// 
/// ```
/// use aus::spectrum;
/// let fft_size: usize = 2048;
/// let audio = aus::read("myfile.wav").unwrap();
/// let window = aus::generate_window_hanning(fft_size);
/// // Just choose the first 2048 samples in the audio file. This might be a problem, because those samples might be zeros.
/// let audio_chunk: Vec<f64> = audio.samples[0][..fft_size].iter().zip(window.iter()).map(|(a, b)| a * b).collect();
/// let (mut magnitude_spectrum, mut phase_spectrum) = spectrum::complex_to_polar_rfft(&spectrum::rfft(&audio_chunk, fft_size));
/// spectrum::fft_exchange_bins(&mut magnitude_spectrum, &mut phase_spectrum, 20);
/// let new_imag_spectrum = spectrum::polar_to_complex_rfft(&magnitude_spectrum, &phase_spectrum).unwrap();
/// ```
pub fn polar_to_complex_rfft(magnitude_spectrum: &[f64], phase_spectrum: &[f64]) -> Result<Vec<Complex<f64>>, SpectrumError> {
    if magnitude_spectrum.len() != phase_spectrum.len() {
        return Err(SpectrumError{error_msg: String::from(format!("The magnitude spectrum and phase spectrum do not \
            have the same length. The magnitude spectrum has len {} and the phase spectrum has len {}.", 
            magnitude_spectrum.len(), phase_spectrum.len()))});
    }
    let mut spectrum = vec![num::complex::Complex::new(0.0, 0.0); magnitude_spectrum.len()];    
    for i in 0..magnitude_spectrum.len() {
        let real = f64::cos(phase_spectrum[i]) * magnitude_spectrum[i];
        let imag = f64::sin(phase_spectrum[i]) * magnitude_spectrum[i];
        let output: Complex<f64> = Complex::new(real, imag);
        spectrum[i] = output;
    }
    Ok(spectrum)
}

/// This function combines magnitude and phase spectrograms into a complex spectrogram for the inverse real STFT.
/// The magnitude and phase spectrograms must have the same length. If they do not have the same length,
/// this function will return a SpectrumError.
/// 
/// # Example
/// 
/// ```
/// use aus::spectrum;
/// let file = aus::read("myfile.wav").unwrap();
/// let mut imaginary_spectrogram = spectrum::rstft(&file.samples[0], 2048, 1024, aus::WindowType::Hanning);
/// let (mut magnitude_spectrogram, mut phase_spectrogram) = spectrum::complex_to_polar_rstft(&imaginary_spectrogram);
/// spectrum::stft_exchange_frames(&mut magnitude_spectrogram, &mut phase_spectrogram, 10);
/// let new_imaginary_spectrogram = spectrum::polar_to_complex_rstft(&magnitude_spectrogram, &phase_spectrogram).unwrap();
/// let new_audio = spectrum::irstft(&new_imaginary_spectrogram, 2048, 1024, aus::WindowType::Hanning).unwrap();
/// let output_file = aus::AudioFile::new_mono(aus::AudioFormat::S24, file.sample_rate, new_audio);
/// aus::write("myfile2.wav", &output_file);
/// ```
pub fn polar_to_complex_rstft(magnitude_spectrogram: &Vec<Vec<f64>>, phase_spectrogram: &Vec<Vec<f64>>) -> Result<Vec<Vec<Complex<f64>>>, SpectrumError> {
    if magnitude_spectrogram.len() != phase_spectrogram.len() {
        return Err(SpectrumError{error_msg: String::from(format!("The magnitude spectrogram and phase spectrogram do not \
            have the same length. The magnitude spectrogram has len {} and the phase spectrogram has len {}.", 
            magnitude_spectrogram.len(), phase_spectrogram.len()))});
    }
    
    let mut spectrogram: Vec<Vec<Complex<f64>>> = Vec::with_capacity(magnitude_spectrogram.len());
    for i in 0..magnitude_spectrogram.len() {
        if magnitude_spectrogram[i].len() != phase_spectrogram[i].len() {
            return Err(SpectrumError{error_msg: String::from(format!("The {}th magnitude spectrum and phase spectrum do not \
                have the same length. The magnitude spectrum has len {} and the phase spectrum has len {}.", 
                i, magnitude_spectrogram[i].len(), phase_spectrogram[i].len()))});
        }
        let mut frame_spectrum: Vec<Complex<f64>> = Vec::with_capacity(magnitude_spectrogram[i].len()); 
        for j in 0..magnitude_spectrogram[i].len() {
            let real = f64::cos(phase_spectrogram[i][j]) * magnitude_spectrogram[i][j];
            let imag = f64::sin(phase_spectrogram[i][j]) * magnitude_spectrogram[i][j];
            let output: Complex<f64> = Complex::new(real, imag);
            frame_spectrum.push(output);
        }
        spectrogram.push(frame_spectrum);
    }
    Ok(spectrogram)
}

/// Gets the corresponding frequencies for the real FFT bins, given a provided FFT size and sample rate.
/// 
/// # Example
/// 
/// ```
/// use aus::spectrum::rfftfreq;
/// let freqs = rfftfreq(2048, 44100);
/// ```
#[inline(always)]
pub fn rfftfreq(fft_size: usize, sample_rate: u32) -> Vec<f64> {
    let mut freqs = vec![0.0 as f64; fft_size / 2 + 1];
    let f_0 = sample_rate as f64 / fft_size as f64;
    for i in 1..freqs.len() {
        freqs[i] = f_0 * i as f64;
    }
    freqs
}