rosa 0.1.0

Audio feature extraction library matching librosa output
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

//! Predominant Local Pulse matching `librosa.beat.plp`.
//!
//! Grosche & Müller (2011): Extracts locally stable tempos from onset strength.

use crate::dsp::util::{Norm, tiny};
use crate::matrix::Matrix;
use crate::stft::{IstftParams, istft};

use super::tempo::fourier_tempo_frequencies;
use super::tempogram::{FourierTempogramParams, fourier_tempogram};

/// Parameters for PLP computation.
pub struct PlpParams {
    pub sr: f64,
    pub hop_length: usize,
    pub win_length: usize,
    pub tempo_min: Option<f64>,
    pub tempo_max: Option<f64>,
}

impl PlpParams {
    pub fn new(sr: f64) -> Self {
        Self {
            sr,
            hop_length: 512,
            win_length: 384,
            tempo_min: Some(30.0),
            tempo_max: Some(300.0),
        }
    }
}

/// Compute predominant local pulse.
///
/// Matches `librosa.beat.plp`.
pub fn plp(onset_envelope: &[f64], params: &PlpParams) -> Vec<f64> {
    let n_frames = onset_envelope.len();
    if n_frames == 0 {
        return Vec::new();
    }

    // Step 1: Compute Fourier tempogram
    let ft_params = FourierTempogramParams {
        sr: params.sr,
        hop_length: params.hop_length,
        win_length: params.win_length,
        center: true,
    };
    let (ft_real, ft_imag) = fourier_tempogram(onset_envelope, &ft_params);

    let n_bins = ft_real.rows();
    let n_cols = ft_real.cols();

    // Step 2: Get tempo frequencies
    let tempo_freqs = fourier_tempo_frequencies(params.sr, params.win_length, params.hop_length);

    // Step 3: Zero out bins outside [tempo_min, tempo_max]
    // Build mask of valid bins
    let mut valid = vec![true; n_bins];
    if let Some(tempo_min) = params.tempo_min {
        for (i, &freq) in tempo_freqs.iter().enumerate() {
            if freq < tempo_min {
                valid[i] = false;
            }
        }
    }
    if let Some(tempo_max) = params.tempo_max {
        for (i, &freq) in tempo_freqs.iter().enumerate() {
            if freq > tempo_max {
                valid[i] = false;
            }
        }
    }

    // Step 4: For each frame, keep only the peak-magnitude bin
    let mut out_real = vec![0.0; n_bins * n_cols];
    let mut out_imag = vec![0.0; n_bins * n_cols];

    for t in 0..n_cols {
        // Compute log-magnitude for valid bins
        let mut peak_val = f64::NEG_INFINITY;
        let mut peak_bin = 0;

        for bin in 0..n_bins {
            if !valid[bin] {
                continue;
            }
            let re = ft_real.get(bin, t);
            let im = ft_imag.get(bin, t);
            let mag = (re * re + im * im).sqrt();
            let log_mag = (1e6 * mag).ln_1p();
            if log_mag > peak_val {
                peak_val = log_mag;
                peak_bin = bin;
            }
        }

        // Keep only the peak bin, normalize to unit magnitude (phase only)
        let re = ft_real.get(peak_bin, t);
        let im = ft_imag.get(peak_bin, t);
        let mag = (re * re + im * im).sqrt();
        let tiny_sqrt = tiny().sqrt();
        let denom = mag.max(tiny_sqrt);

        let idx = peak_bin * n_cols + t;
        out_real[idx] = re / denom;
        out_imag[idx] = im / denom;
    }

    let stft_real = Matrix::from_vec(out_real, n_bins, n_cols);
    let stft_imag = Matrix::from_vec(out_imag, n_bins, n_cols);

    // Step 5: ISTFT with hop_length=1
    let istft_params = IstftParams {
        n_fft: params.win_length,
        hop_length: 1,
        win_length: params.win_length,
        center: true,
        length: Some(n_frames),
    };

    let mut pulse = istft(&stft_real, &stft_imag, &istft_params);

    // Step 6: Clip negatives
    for val in pulse.iter_mut() {
        *val = val.max(0.0);
    }

    // Step 7: Normalize
    crate::dsp::util::normalize(&mut pulse, Norm::Linf);

    pulse
}