115 lines
2.5 KiB
Rust
115 lines
2.5 KiB
Rust
// vim: noet
|
|
|
|
use fftw::array::AlignedVec;
|
|
use fftw::plan::*;
|
|
use fftw::types::*;
|
|
use std::f32::consts::PI;
|
|
|
|
pub struct SignalProcessing
|
|
{
|
|
samp_rate: f32,
|
|
|
|
fft_window: Vec<f32>,
|
|
|
|
fft_input: AlignedVec<f32>,
|
|
fft_output: AlignedVec<c32>,
|
|
|
|
fft_plan: R2CPlan32,
|
|
|
|
fft_absolute: Vec<f32>,
|
|
}
|
|
|
|
impl SignalProcessing
|
|
{
|
|
fn hann_window(block_size: usize) -> Vec<f32>
|
|
{
|
|
let mut window = vec![0.0; block_size];
|
|
|
|
for i in 0..block_size {
|
|
window[i] = (PI * (i as f32) / (block_size as f32)).sin().powi(2);
|
|
}
|
|
|
|
window
|
|
}
|
|
|
|
pub fn new(block_size: usize, samp_rate: f32) -> fftw::error::Result<SignalProcessing>
|
|
{
|
|
let freq_domain_size = block_size/2 + 1;
|
|
|
|
let s = SignalProcessing {
|
|
samp_rate: samp_rate,
|
|
fft_window: SignalProcessing::hann_window(block_size),
|
|
fft_input: AlignedVec::new(block_size),
|
|
fft_output: AlignedVec::new(freq_domain_size),
|
|
fft_plan: R2CPlan::aligned(&[block_size], Flag::MEASURE)?,
|
|
|
|
fft_absolute: vec![0.0; freq_domain_size],
|
|
};
|
|
|
|
Ok(s)
|
|
}
|
|
|
|
fn apply_window(&mut self)
|
|
{
|
|
self.fft_input.iter_mut()
|
|
.zip(self.fft_window.iter())
|
|
.for_each(|(s, w)| *s *= w);
|
|
}
|
|
|
|
pub fn import_i16_stereo(&mut self, data: &[i16]) -> std::result::Result<(), &str>
|
|
{
|
|
if data.len() != 2*self.fft_input.len() {
|
|
return Err("Stereo data length does not match 2x the FFT input length.");
|
|
}
|
|
|
|
data.chunks_exact(2)
|
|
.map(|channels| (channels[0] as f32 + channels[1] as f32) / 2.0 / 32768.0)
|
|
.zip(self.fft_input.iter_mut())
|
|
.for_each(|(c, t)| *t = c);
|
|
|
|
self.apply_window();
|
|
|
|
Ok(())
|
|
}
|
|
|
|
pub fn import_i16_mono(&mut self, data: &[i16]) -> std::result::Result<(), &str>
|
|
{
|
|
if data.len() != self.fft_input.len() {
|
|
return Err("Mono data length does not match the FFT input length.");
|
|
}
|
|
|
|
data.iter()
|
|
.map(|&sample| (sample as f32) / 32768.0)
|
|
.zip(self.fft_input.iter_mut())
|
|
.for_each(|(c, t)| *t = c);
|
|
|
|
self.apply_window();
|
|
|
|
Ok(())
|
|
}
|
|
|
|
pub fn update_fft(&mut self) -> fftw::error::Result<()>
|
|
{
|
|
self.fft_plan.r2c(&mut self.fft_input, &mut self.fft_output)?;
|
|
|
|
for (i, abs_sample) in self.fft_absolute.iter_mut().enumerate() {
|
|
*abs_sample = self.fft_output[i].norm();
|
|
}
|
|
|
|
Ok(())
|
|
}
|
|
|
|
fn freq_to_idx(&self, freq: f32) -> usize
|
|
{
|
|
(freq * (self.fft_input.len() as f32) / self.samp_rate) as usize
|
|
}
|
|
|
|
pub fn get_energy_in_band(&self, freq_start: f32, freq_end: f32) -> f32
|
|
{
|
|
let start_bin = self.freq_to_idx(freq_start);
|
|
let end_bin = self.freq_to_idx(freq_end);
|
|
|
|
self.fft_absolute[start_bin ..= end_bin].iter().sum()
|
|
}
|
|
}
|