rosa 0.1.0

Audio feature extraction library matching librosa output
Documentation
/// Mu-law compression.
///
/// Matches `librosa.mu_compress(x, mu=mu, quantize=quantize)`.
///
/// Input `x` must be in `[-1, +1]`. `mu` must be > 0.
/// Formula: `sign(x) * ln(1 + mu * |x|) / ln(1 + mu)`
pub fn mu_compress(x: &[f64], mu: f64, quantize: bool) -> Vec<f64> {
    let log_1_mu = (1.0 + mu).ln();

    let mut result: Vec<f64> = x
        .iter()
        .map(|&val| val.signum() * (1.0 + mu * val.abs()).ln() / log_1_mu)
        .collect();

    if quantize {
        // librosa: np.digitize(x_mu, np.linspace(-1, 1, int(1+mu)), right=True) - int((1+mu)/2)
        // np.digitize(x, bins, right=True) finds leftmost index i where bins[i] >= x
        let n_edges = (1.0 + mu) as usize;
        let half = ((1.0 + mu) / 2.0) as i64;
        let step = 2.0 / (n_edges as f64 - 1.0);
        for val in result.iter_mut() {
            // Binary search equivalent of np.digitize(val, bins, right=True)
            let mut bin = n_edges as i64; // default: past all edges
            for j in 0..n_edges {
                let edge = -1.0 + j as f64 * step;
                if edge >= *val {
                    bin = j as i64;
                    break;
                }
            }
            *val = (bin - half) as f64;
        }
    }

    result
}

/// Mu-law expansion (inverse of mu_compress).
///
/// Matches `librosa.mu_expand(x, mu=mu, quantize=quantize)`.
///
/// Formula: `sign(x) * (1/mu) * ((1 + mu)^|x| - 1)`
pub fn mu_expand(x: &[f64], mu: f64, quantize: bool) -> Vec<f64> {
    x.iter()
        .map(|&val| {
            let v = if quantize {
                val * 2.0 / (1.0 + mu)
            } else {
                val
            };
            v.signum() / mu * ((1.0 + mu).powf(v.abs()) - 1.0)
        })
        .collect()
}

#[cfg(test)]
mod tests {
    use super::*;
    use serde::Deserialize;

    fn assert_approx_eq(actual: &[f64], expected: &[f64], tol: f64, name: &str) {
        assert_eq!(
            actual.len(),
            expected.len(),
            "{name}: length mismatch: got {} expected {}",
            actual.len(),
            expected.len()
        );
        for (i, (a, e)) in actual.iter().zip(expected.iter()).enumerate() {
            let diff = (a - e).abs();
            assert!(
                diff < tol,
                "{name}[{i}]: got {a}, expected {e}, diff {diff} > tol {tol}"
            );
        }
    }

    #[derive(Deserialize)]
    struct MuLawTestCase {
        name: String,
        function: String,
        input: Vec<f64>,
        mu: f64,
        quantize: bool,
        expected: Vec<f64>,
    }

    fn load_references() -> Vec<MuLawTestCase> {
        let data = include_str!("tests/references/mu_law.json");
        serde_json::from_str(data).expect("Failed to parse mu_law.json")
    }

    #[test]
    fn test_mu_compress_matches_librosa() {
        for case in load_references()
            .into_iter()
            .filter(|c| c.function == "mu_compress")
        {
            let result = mu_compress(&case.input, case.mu, case.quantize);
            let tol = if case.quantize { 0.5 } else { 1e-12 };
            assert_approx_eq(&result, &case.expected, tol, &case.name);
        }
    }

    #[test]
    fn test_mu_expand_matches_librosa() {
        for case in load_references()
            .into_iter()
            .filter(|c| c.function == "mu_expand")
        {
            let result = mu_expand(&case.input, case.mu, case.quantize);
            assert_approx_eq(&result, &case.expected, 1e-10, &case.name);
        }
    }
}