use rosa::beat::*;
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() {
if e.is_infinite() || *e >= 1e307 {
assert!(
a.is_infinite() || *a >= 1e307,
"{name}[{i}]: expected inf/very large, got {a}"
);
continue;
}
let diff = (a - e).abs();
let threshold = tol * e.abs().max(1.0);
assert!(
diff < threshold,
"{name}[{i}]: got {a}, expected {e}, diff {diff} > tol {threshold}"
);
}
}
#[test]
fn test_pad_linear_ramp() {
let cases: Vec<serde_json::Value> =
serde_json::from_str(include_str!("references/beat/pad_linear_ramp.json")).unwrap();
for case in &cases {
let name = case["name"].as_str().unwrap();
let input: Vec<f64> = serde_json::from_value(case["input"].clone()).unwrap();
let pad_left = case["pad_left"].as_u64().unwrap() as usize;
let pad_right = case["pad_right"].as_u64().unwrap() as usize;
let expected: Vec<f64> = serde_json::from_value(case["expected"].clone()).unwrap();
let result = rosa::dsp::util::pad_linear_ramp(&input, pad_left, pad_right);
assert_approx_eq(&result, &expected, 1e-12, name);
}
}
#[test]
fn test_tempo_frequencies() {
let data: serde_json::Value =
serde_json::from_str(include_str!("references/beat/tempo_frequencies.json")).unwrap();
for case in data["tempo_frequencies"].as_array().unwrap() {
let name = case["name"].as_str().unwrap();
let n_bins = case["n_bins"].as_u64().unwrap() as usize;
let hop_length = case["hop_length"].as_u64().unwrap() as usize;
let sr = case["sr"].as_f64().unwrap();
let expected: Vec<f64> = serde_json::from_value(case["expected"].clone()).unwrap();
let result = tempo_frequencies(n_bins, hop_length, sr);
assert_approx_eq(
&result,
&expected,
1e-10,
&format!("tempo_frequencies_{name}"),
);
}
}
#[test]
fn test_fourier_tempo_frequencies() {
let data: serde_json::Value =
serde_json::from_str(include_str!("references/beat/tempo_frequencies.json")).unwrap();
for case in data["fourier_tempo_frequencies"].as_array().unwrap() {
let name = case["name"].as_str().unwrap();
let sr = case["sr"].as_f64().unwrap();
let win_length = case["win_length"].as_u64().unwrap() as usize;
let hop_length = case["hop_length"].as_u64().unwrap() as usize;
let expected: Vec<f64> = serde_json::from_value(case["expected"].clone()).unwrap();
let result = fourier_tempo_frequencies(sr, win_length, hop_length);
assert_approx_eq(
&result,
&expected,
1e-10,
&format!("fourier_tempo_frequencies_{name}"),
);
}
}
#[test]
fn test_istft_roundtrip() {
let data: serde_json::Value =
serde_json::from_str(include_str!("references/beat/istft.json")).unwrap();
let original: Vec<f64> = serde_json::from_value(data["original"].clone()).unwrap();
let n_fft = data["n_fft"].as_u64().unwrap() as usize;
let hop_length = data["hop_length"].as_u64().unwrap() as usize;
let stft_params = rosa::stft::StftParams {
n_fft,
hop_length,
win_length: n_fft,
center: true,
window: None,
};
let (stft_real, stft_imag) = rosa::stft::stft_as_real_imag(&original, &stft_params);
let istft_params = rosa::stft::IstftParams {
n_fft,
hop_length,
win_length: n_fft,
center: true,
length: None,
};
let reconstructed = rosa::stft::istft(&stft_real, &stft_imag, &istft_params);
let margin = n_fft;
let end = reconstructed
.len()
.min(original.len())
.saturating_sub(margin);
if end > margin {
assert_approx_eq(
&reconstructed[margin..end],
&original[margin..end],
1e-6,
"istft_roundtrip_core",
);
}
}
#[test]
fn test_tempogram() {
let data: serde_json::Value =
serde_json::from_str(include_str!("references/beat/tempogram.json")).unwrap();
let sr = data["sr"].as_f64().unwrap();
let hop_length = data["hop_length"].as_u64().unwrap() as usize;
let win_length = data["win_length"].as_u64().unwrap() as usize;
let onset_env: Vec<f64> = serde_json::from_value(data["onset_envelope"].clone()).unwrap();
let expected_shape: Vec<usize> = serde_json::from_value(data["shape"].clone()).unwrap();
let params = TempogramParams {
sr,
hop_length,
win_length,
center: true,
norm: Some(rosa::dsp::util::Norm::Linf),
};
let tg = tempogram(&onset_env, ¶ms);
assert_eq!(tg.rows(), expected_shape[0], "tempogram rows mismatch");
assert_eq!(tg.cols(), expected_shape[1], "tempogram cols mismatch");
for col_data in data["sampled_columns"].as_array().unwrap() {
let col_idx = col_data["index"].as_u64().unwrap() as usize;
let expected_col: Vec<f64> = serde_json::from_value(col_data["data"].clone()).unwrap();
if col_idx >= tg.cols() {
continue;
}
let actual_col: Vec<f64> = (0..tg.rows()).map(|r| tg.get(r, col_idx)).collect();
assert_approx_eq(
&actual_col,
&expected_col,
1e-4,
&format!("tempogram_col_{col_idx}"),
);
}
}
#[test]
fn test_tempo() {
let data: serde_json::Value =
serde_json::from_str(include_str!("references/beat/tempo.json")).unwrap();
let sr = data["sr"].as_f64().unwrap();
let hop_length = data["hop_length"].as_u64().unwrap() as usize;
let start_bpm = data["start_bpm"].as_f64().unwrap();
let onset_env: Vec<f64> = serde_json::from_value(data["onset_envelope"].clone()).unwrap();
let expected_tempo = data["tempo"].as_f64().unwrap();
let params = TempoParams {
sr,
hop_length,
start_bpm,
std_bpm: 1.0,
ac_size: 8.0,
max_tempo: Some(320.0),
};
let result = tempo(&onset_env, ¶ms);
let diff = (result - expected_tempo).abs();
let rel_diff = diff / expected_tempo;
assert!(
rel_diff < 0.05 || diff < 5.0,
"tempo: got {result}, expected {expected_tempo}, diff {diff} ({:.1}%)",
rel_diff * 100.0
);
}
#[test]
fn test_beat_track() {
let data: serde_json::Value =
serde_json::from_str(include_str!("references/beat/beat_track.json")).unwrap();
let sr = data["sr"].as_f64().unwrap();
let hop_length = data["hop_length"].as_u64().unwrap() as usize;
let start_bpm = data["start_bpm"].as_f64().unwrap();
let onset_env: Vec<f64> = serde_json::from_value(data["onset_envelope"].clone()).unwrap();
let expected_tempo = data["tempo"].as_f64().unwrap();
let expected_beats: Vec<usize> = serde_json::from_value(data["beats"].clone()).unwrap();
let params = BeatTrackParams {
sr,
hop_length,
start_bpm,
tightness: 100.0,
trim: true,
bpm: None,
};
let (result_tempo, result_beats) = beat_track(&onset_env, ¶ms);
let tempo_diff = (result_tempo - expected_tempo).abs() / expected_tempo;
assert!(
tempo_diff < 0.1,
"beat_track tempo: got {result_tempo}, expected {expected_tempo}"
);
if !expected_beats.is_empty() {
let count_diff = (result_beats.len() as f64 - expected_beats.len() as f64).abs()
/ expected_beats.len() as f64;
assert!(
count_diff < 0.3,
"beat_track beat count: got {}, expected {}",
result_beats.len(),
expected_beats.len()
);
}
}
#[test]
fn test_clicks_beat_track() {
let data: serde_json::Value =
serde_json::from_str(include_str!("references/beat/clicks_beat_track.json")).unwrap();
let sr = data["sr"].as_f64().unwrap();
let hop_length = data["hop_length"].as_u64().unwrap() as usize;
let onset_env: Vec<f64> = serde_json::from_value(data["onset_envelope"].clone()).unwrap();
let known_bpm = data["known_bpm"].as_f64().unwrap();
let params = BeatTrackParams {
sr,
hop_length,
start_bpm: 120.0,
tightness: 100.0,
trim: true,
bpm: None,
};
let (result_tempo, result_beats) = beat_track(&onset_env, ¶ms);
let tempo_diff = (result_tempo - known_bpm).abs();
assert!(
tempo_diff < 10.0,
"clicks tempo: got {result_tempo}, expected ~{known_bpm}"
);
assert!(
!result_beats.is_empty(),
"clicks beat_track: should detect some beats"
);
}