sva-samples 0.14.6

// Concern: resolves each frame into the resonances an all-pole model puts under it | Non-concern: the discrete spectrum (spectrum.rs), naming a peak (pitch.rs) | IO: (planes, sample rate) -> frames

use crate::measure::spectrum::db;

pub const MAX_ORDER: usize = 64;

/// Nearer than this to DC or Nyquist a pole is the overall tilt, not a resonance.
const EDGE_HZ: f64 = 40.0;

#[derive(Clone, Copy, Debug, PartialEq)]
pub struct Formant {
    pub hz: f64,
    pub bandwidth_hz: f64,
    pub db: f64,
}

/// `residual` and `energy` are mean squares, and neither derives from the other.
#[derive(Clone, Debug, PartialEq)]
pub struct FormantFrame {
    pub t_secs: f64,
    pub order: usize,
    pub energy: f64,
    pub residual: f64,
    pub formants: Vec<Formant>,
}

/// Two poles per kilohertz of bandwidth plus two for the tilt — 46 at 44.1 kHz, far more than a
/// voice wants and about right for a full-band instrument.
pub fn default_order(sample_rate: f64) -> usize {
    (2.0 + sample_rate / 1000.0).clamp(2.0, MAX_ORDER as f64) as usize
}

pub fn track(
    planes: &[&[f64]],
    sample_rate: f64,
    start_secs: f64,
    frame_secs: f64,
    order: usize,
    max_formants: usize,
) -> Vec<FormantFrame> {
    let stride = ((frame_secs * sample_rate).round() as usize).max(1);
    let len = crate::measure::held(planes);
    (0..len)
        .step_by(stride)
        .enumerate()
        .map(|(n, from)| {
            let chunk: Vec<&[f64]> = planes
                .iter()
                .map(|p| &p[from..len.min(from + stride)])
                .collect();
            FormantFrame {
                t_secs: start_secs + (n * stride) as f64 / sample_rate,
                ..analyze(&chunk, sample_rate, order, max_formants)
            }
        })
        .collect()
}

/// Hamming, not Hann: a taper reaching zero throws the frame's edges away. Never
/// pre-emphasised; `x - 0.97*x(t - 1sp)` is a composition's own to write.
pub fn analyze(
    planes: &[&[f64]],
    sample_rate: f64,
    order: usize,
    max_formants: usize,
) -> FormantFrame {
    let len = crate::measure::held(planes);
    let order = order.clamp(2, MAX_ORDER).min(len.saturating_sub(1));
    let windowed: Vec<Vec<f64>> = planes
        .iter()
        .map(|frame| {
            frame
                .iter()
                .enumerate()
                .map(|(n, &x)| {
                    let turn = 2.0 * std::f64::consts::PI * n as f64 / len as f64;
                    let w = 0.54 - 0.46 * turn.cos();
                    x * w
                })
                .collect()
        })
        .collect();
    let energy = match len {
        0 => 0.0,
        _ => {
            let each = windowed
                .iter()
                .map(|w| w.iter().map(|x| x * x).sum::<f64>());
            each.sum::<f64>() / (len * planes.len()) as f64
        }
    };
    if order < 2 || energy <= 0.0 {
        return FormantFrame {
            t_secs: 0.0,
            order,
            energy,
            residual: energy,
            formants: Vec::new(),
        };
    }

    let r = crate::measure::averaged(windowed.iter().map(|w| autocorrelation(w, order)).collect());
    let (a, error) = levinson(&r, order);
    let residual = error / len as f64;
    let gain = residual.max(0.0).sqrt();

    let mut formants: Vec<Formant> = roots(&a)
        .into_iter()
        .filter(|z| z.im > 0.0)
        .filter_map(|z| {
            let hz = z.im.atan2(z.re) * sample_rate / (2.0 * std::f64::consts::PI);
            let radius = z.abs();
            if !(EDGE_HZ..sample_rate / 2.0 - EDGE_HZ).contains(&hz) || radius >= 1.0 {
                return None;
            }
            Some(Formant {
                hz,
                bandwidth_hz: -radius.ln() * sample_rate / std::f64::consts::PI,
                db: db(gain / response(&a, hz, sample_rate)),
            })
        })
        .collect();
    formants.sort_by(|x, y| x.hz.total_cmp(&y.hz));
    formants.truncate(max_formants);

    FormantFrame {
        t_secs: 0.0,
        order,
        energy,
        residual,
        formants,
    }
}

fn autocorrelation(x: &[f64], lags: usize) -> Vec<f64> {
    (0..=lags)
        .map(|k| x.iter().skip(k).zip(x).map(|(a, b)| a * b).sum())
        .collect()
}

/// Levinson-Durbin: the order-`p` predictor from `p` autocorrelations, no FFT and no matrix.
fn levinson(r: &[f64], order: usize) -> (Vec<f64>, f64) {
    let mut a = vec![0.0; order + 1];
    a[0] = 1.0;
    let mut error = r[0];

    for i in 1..=order {
        if error <= 0.0 {
            return (a, error.max(0.0));
        }
        let acc: f64 = r[i] + (1..i).map(|j| a[j] * r[i - j]).sum::<f64>();
        let k = -acc / error;
        let held: Vec<f64> = a[1..i].to_vec();
        for (j, prior) in held.iter().enumerate() {
            a[j + 1] = prior + k * held[i - 2 - j];
        }
        a[i] = k;
        error *= 1.0 - k * k;
    }
    (a, error.max(0.0))
}

fn response(a: &[f64], hz: f64, sample_rate: f64) -> f64 {
    let w = 2.0 * std::f64::consts::PI * hz / sample_rate;
    let (mut re, mut im) = (0.0, 0.0);
    for (k, coefficient) in a.iter().enumerate() {
        re += coefficient * (w * k as f64).cos();
        im -= coefficient * (w * k as f64).sin();
    }
    (re * re + im * im).sqrt()
}

#[derive(Clone, Copy, Debug)]
struct C {
    re: f64,
    im: f64,
}

impl C {
    fn mul(self, o: C) -> C {
        C {
            re: self.re * o.re - self.im * o.im,
            im: self.re * o.im + self.im * o.re,
        }
    }

    fn sub(self, o: C) -> C {
        C {
            re: self.re - o.re,
            im: self.im - o.im,
        }
    }

    fn div(self, o: C) -> C {
        let d = o.re * o.re + o.im * o.im;
        match d == 0.0 {
            true => C { re: 0.0, im: 0.0 },
            false => C {
                re: (self.re * o.re + self.im * o.im) / d,
                im: (self.im * o.re - self.re * o.im) / d,
            },
        }
    }

    fn abs(self) -> f64 {
        self.re.hypot(self.im)
    }
}

/// Durand-Kerner: every root moves at once, so no deflation loses the closely spaced pairs an
/// envelope is made of.
fn roots(a: &[f64]) -> Vec<C> {
    const ROUNDS: usize = 500;
    const SETTLED: f64 = 1e-14;

    let degree = a.len() - 1;
    let seed = C { re: 0.4, im: 0.9 };
    let mut z: Vec<C> = Vec::with_capacity(degree);
    let mut power = C { re: 1.0, im: 0.0 };
    for _ in 0..degree {
        z.push(power);
        power = power.mul(seed);
    }

    for _ in 0..ROUNDS {
        let mut moved: f64 = 0.0;
        for i in 0..degree {
            let mut denominator = C { re: 1.0, im: 0.0 };
            for j in 0..degree {
                if j != i {
                    denominator = denominator.mul(z[i].sub(z[j]));
                }
            }
            let step = evaluate(a, z[i]).div(denominator);
            z[i] = z[i].sub(step);
            moved = moved.max(step.abs());
        }
        if moved < SETTLED {
            break;
        }
    }
    z
}

fn evaluate(a: &[f64], at: C) -> C {
    let mut acc = C { re: 0.0, im: 0.0 };
    for coefficient in a {
        acc = acc.mul(at);
        acc.re += coefficient;
    }
    acc
}