use ndarray::Array2;
use num_complex::Complex;
use thiserror::Error;
#[derive(Error, Debug)]
pub enum PhaseRecoveryError {
#[error("Computation failed: {0}")]
ComputationFailed(String),
}
pub fn griffinlim(s: &Array2<f32>) -> GriffinLimBuilder<'_> {
GriffinLimBuilder {
s,
n_iter: 32,
hop_length: None, }
}
#[derive(Debug, Clone)]
pub struct GriffinLimBuilder<'a> {
s: &'a Array2<f32>,
n_iter: usize,
hop_length: Option<usize>,
}
impl<'a> GriffinLimBuilder<'a> {
pub fn n_iter(mut self, n_iter: usize) -> Self {
self.n_iter = n_iter;
self
}
pub fn hop_length(mut self, hop_length: usize) -> Self {
self.hop_length = Some(hop_length);
self
}
pub fn compute(self) -> Result<Vec<f32>, PhaseRecoveryError> {
griffinlim_impl(self.s, self.n_iter, self.hop_length)
}
}
fn griffinlim_impl(s: &Array2<f32>, n_iter: usize, hop_length: Option<usize>) -> Result<Vec<f32>, PhaseRecoveryError> {
let n_fft = (s.shape()[0] - 1) * 2;
let hop = hop_length.unwrap_or(n_fft / 4).max(1);
let signal_len = hop * (s.shape()[1] - 1) + n_fft;
let mut y = crate::signal_generation::generators::tone(440.0, 44100)
.duration(signal_len as f32 / 44100.0)
.compute();
for _ in 0..n_iter {
let stft_y = crate::signal_processing::time_frequency::stft(&y)
.n_fft(n_fft)
.hop_length(hop)
.compute()
.map_err(|e| PhaseRecoveryError::ComputationFailed(format!("STFT computation failed: {}", e)))?;
let (mut mag, mut phase) = crate::signal_processing::time_frequency::magphase(&stft_y, None);
for ((i, j), m) in mag.indexed_iter_mut() {
*m = s[[i, j]]; let p = &mut phase[[i, j]];
if m.abs() > 1e-10 {
*p /= p.norm(); }
}
let new_stft = mag.mapv(|x| Complex::new(x, 0.0)) * phase;
y = crate::signal_processing::time_frequency::istft(&new_stft, Some(hop), None, Some(signal_len));
}
Ok(y)
}