use rosa::matrix::Matrix;
use rosa::sequence::*;
fn load_refs() -> serde_json::Value {
let data = include_str!("references/sequence/sequence.json");
serde_json::from_str(data).expect("Failed to parse sequence.json")
}
fn flatten_matrix(nested: &serde_json::Value) -> Vec<f64> {
nested
.as_array()
.unwrap()
.iter()
.flat_map(|row| row.as_array().unwrap().iter().map(|v| v.as_f64().unwrap()))
.collect()
}
fn assert_matrix_eq(actual: &Matrix, expected_nested: &serde_json::Value, tol: f64, name: &str) {
let expected = flatten_matrix(expected_nested);
let actual_data = actual.as_slice();
assert_eq!(
actual_data.len(),
expected.len(),
"{name}: size mismatch: {} vs {}",
actual_data.len(),
expected.len()
);
for (i, (a, e)) in actual_data.iter().zip(expected.iter()).enumerate() {
let diff = (a - e).abs();
assert!(
diff < tol,
"{name}[{i}]: got {a}, expected {e}, diff {diff}"
);
}
}
#[test]
fn test_transition_uniform_3_matches_librosa() {
let refs = load_refs();
let case = &refs["transition_uniform_3"];
let n = case["n_states"].as_u64().unwrap() as usize;
let t = transition_uniform(n);
assert_matrix_eq(&t, &case["matrix"], 1e-12, "transition_uniform_3");
}
#[test]
fn test_transition_uniform_5_matches_librosa() {
let refs = load_refs();
let case = &refs["transition_uniform_5"];
let n = case["n_states"].as_u64().unwrap() as usize;
let t = transition_uniform(n);
assert_matrix_eq(&t, &case["matrix"], 1e-12, "transition_uniform_5");
}
#[test]
fn test_transition_loop_scalar_matches_librosa() {
let refs = load_refs();
let case = &refs["transition_loop_3_0.5"];
let n = case["n_states"].as_u64().unwrap() as usize;
let prob: Vec<f64> = serde_json::from_value(case["prob"].clone()).unwrap();
let t = transition_loop(n, &prob);
assert_matrix_eq(&t, &case["matrix"], 1e-12, "transition_loop_3_0.5");
}
#[test]
fn test_transition_loop_varied_matches_librosa() {
let refs = load_refs();
let case = &refs["transition_loop_3_varied"];
let n = case["n_states"].as_u64().unwrap() as usize;
let prob: Vec<f64> = serde_json::from_value(case["prob"].clone()).unwrap();
let t = transition_loop(n, &prob);
assert_matrix_eq(&t, &case["matrix"], 1e-12, "transition_loop_3_varied");
}
#[test]
fn test_transition_loop_4_matches_librosa() {
let refs = load_refs();
let case = &refs["transition_loop_4_0.9"];
let n = case["n_states"].as_u64().unwrap() as usize;
let prob: Vec<f64> = serde_json::from_value(case["prob"].clone()).unwrap();
let t = transition_loop(n, &prob);
assert_matrix_eq(&t, &case["matrix"], 1e-12, "transition_loop_4_0.9");
}
#[test]
fn test_transition_local_5_3_triangle_matches_librosa() {
let refs = load_refs();
let case = &refs["transition_local_5_3_triangle"];
let n = case["n_states"].as_u64().unwrap() as usize;
let w = case["width"].as_u64().unwrap() as usize;
let t = transition_local(n, w, "triangle", false);
assert_matrix_eq(&t, &case["matrix"], 1e-12, "transition_local_5_3_triangle");
}
#[test]
fn test_transition_local_5_3_triangle_wrap_matches_librosa() {
let refs = load_refs();
let case = &refs["transition_local_5_3_triangle_wrap"];
let n = case["n_states"].as_u64().unwrap() as usize;
let w = case["width"].as_u64().unwrap() as usize;
let t = transition_local(n, w, "triangle", true);
assert_matrix_eq(
&t,
&case["matrix"],
1e-12,
"transition_local_5_3_triangle_wrap",
);
}
#[test]
fn test_transition_local_5_3_ones_matches_librosa() {
let refs = load_refs();
let case = &refs["transition_local_5_3_ones"];
let n = case["n_states"].as_u64().unwrap() as usize;
let w = case["width"].as_u64().unwrap() as usize;
let t = transition_local(n, w, "ones", false);
assert_matrix_eq(&t, &case["matrix"], 1e-12, "transition_local_5_3_ones");
}
#[test]
fn test_transition_local_8_5_triangle_matches_librosa() {
let refs = load_refs();
let case = &refs["transition_local_8_5_triangle"];
let n = case["n_states"].as_u64().unwrap() as usize;
let w = case["width"].as_u64().unwrap() as usize;
let t = transition_local(n, w, "triangle", false);
assert_matrix_eq(&t, &case["matrix"], 1e-12, "transition_local_8_5_triangle");
}
#[test]
fn test_viterbi_2state_matches_librosa() {
let refs = load_refs();
let case = &refs["viterbi_2state"];
let prob_flat = flatten_matrix(&case["prob"]);
let n_states = case["prob"].as_array().unwrap().len();
let n_steps = case["prob"][0].as_array().unwrap().len();
let prob = Matrix::from_vec(prob_flat, n_states, n_steps);
let trans_flat = flatten_matrix(&case["transition"]);
let trans = Matrix::from_vec(trans_flat, n_states, n_states);
let p_init: Vec<f64> = serde_json::from_value(case["p_init"].clone()).unwrap();
let path = viterbi(&prob, &trans, Some(&p_init));
let expected: Vec<usize> = serde_json::from_value(case["path"].clone()).unwrap();
assert_eq!(path, expected, "viterbi_2state path mismatch");
}
#[test]
fn test_viterbi_3state_matches_librosa() {
let refs = load_refs();
let case = &refs["viterbi_3state"];
let prob_flat = flatten_matrix(&case["prob"]);
let n_states = case["prob"].as_array().unwrap().len();
let n_steps = case["prob"][0].as_array().unwrap().len();
let prob = Matrix::from_vec(prob_flat, n_states, n_steps);
let trans_flat = flatten_matrix(&case["transition"]);
let trans = Matrix::from_vec(trans_flat, n_states, n_states);
let path = viterbi(&prob, &trans, None);
let expected: Vec<usize> = serde_json::from_value(case["path"].clone()).unwrap();
assert_eq!(path, expected, "viterbi_3state path mismatch");
}
#[test]
fn test_viterbi_discriminative_matches_librosa() {
let refs = load_refs();
let case = &refs["viterbi_discriminative_2state"];
let prob_flat = flatten_matrix(&case["prob"]);
let n_states = case["prob"].as_array().unwrap().len();
let n_steps = case["prob"][0].as_array().unwrap().len();
let prob = Matrix::from_vec(prob_flat, n_states, n_steps);
let trans_flat = flatten_matrix(&case["transition"]);
let trans = Matrix::from_vec(trans_flat, n_states, n_states);
let path = viterbi_discriminative(&prob, &trans, None, None);
let expected: Vec<usize> = serde_json::from_value(case["path"].clone()).unwrap();
assert_eq!(path, expected, "viterbi_discriminative path mismatch");
}
#[test]
fn test_viterbi_binary_matches_librosa() {
let refs = load_refs();
let case = &refs["viterbi_binary"];
let prob_flat = flatten_matrix(&case["prob"]);
let n_states = case["prob"].as_array().unwrap().len();
let n_steps = case["prob"][0].as_array().unwrap().len();
let prob = Matrix::from_vec(prob_flat, n_states, n_steps);
let trans_flat = flatten_matrix(&case["transition"]);
let trans = Matrix::from_vec(trans_flat, 2, 2);
let result = viterbi_binary(&prob, &trans, Some(0.5), Some(0.5));
let expected_flat = flatten_matrix(&case["states"]);
let expected_n_states = case["states"].as_array().unwrap().len();
let expected_n_steps = case["states"][0].as_array().unwrap().len();
assert_eq!(result.shape(), (expected_n_states, expected_n_steps));
for (i, (a, e)) in result
.as_slice()
.iter()
.zip(expected_flat.iter())
.enumerate()
{
assert_eq!(
*a as usize, *e as usize,
"viterbi_binary[{i}]: got {a}, expected {e}"
);
}
}