#![cfg(test)]
use crate::entity_cluster::{cluster_names, names_plausibly_same_person};
use std::collections::HashMap;
const DRIFT_GROUPS: &[&[&str]] = &[
&["tan-vir", "tanvir", "tanmir", "tanzir"],
&["nadia", "nadya"],
&["aaron", "arron"],
];
const DISTINCT_DISSIMILAR: &[&str] = &[
"deepak", "sarah", "bright", "liam", "monica", "keith", "priya", "tomas", "李雷", "李蕾",
];
const AMBIGUOUS_PAIRS: &[(&str, &str)] = &[("sam", "sami"), ("an", "ann"), ("rena", "rana")];
#[derive(Debug, Default)]
struct ClusterReport {
drift_groups_recovered: usize,
drift_groups_total: usize,
wrong_merges: usize,
ambiguous_suggested: usize,
ambiguous_total: usize,
bcubed_precision: f64,
bcubed_recall: f64,
}
fn build_pool() -> (Vec<String>, HashMap<String, usize>) {
let mut pool: Vec<String> = Vec::new();
let mut truth: HashMap<String, usize> = HashMap::new();
let mut next_id = 0usize;
for group in DRIFT_GROUPS {
let id = next_id;
next_id += 1;
for &name in *group {
pool.push(name.to_string());
truth.insert(name.to_string(), id);
}
}
for &name in DISTINCT_DISSIMILAR {
let id = next_id;
next_id += 1;
pool.push(name.to_string());
truth.insert(name.to_string(), id);
}
(pool, truth)
}
fn run_eval() -> ClusterReport {
let (pool, truth) = build_pool();
let clusters = cluster_names(&pool);
let mut predicted: HashMap<String, usize> = HashMap::new();
for (cid, cluster) in clusters.iter().enumerate() {
for &idx in cluster {
predicted.insert(pool[idx].clone(), cid);
}
}
let mut next_singleton = clusters.len();
for name in &pool {
predicted.entry(name.clone()).or_insert_with(|| {
let id = next_singleton;
next_singleton += 1;
id
});
}
let mut report = ClusterReport {
drift_groups_total: DRIFT_GROUPS.len(),
ambiguous_total: AMBIGUOUS_PAIRS.len(),
..Default::default()
};
for cluster in &clusters {
let mut persons: Vec<usize> = cluster.iter().map(|&i| truth[&pool[i]]).collect();
persons.sort_unstable();
persons.dedup();
if persons.len() > 1 {
report.wrong_merges += 1;
}
}
for group in DRIFT_GROUPS {
let first = predicted[&group[0].to_string()];
if group.iter().all(|&n| predicted[&n.to_string()] == first) {
report.drift_groups_recovered += 1;
}
}
for (a, b) in AMBIGUOUS_PAIRS {
if names_plausibly_same_person(a, b).is_some() {
report.ambiguous_suggested += 1;
}
}
let (mut prec_sum, mut rec_sum) = (0.0f64, 0.0f64);
for name in &pool {
let pc = predicted[name];
let gc = truth[name];
let pred_members: Vec<&String> = pool.iter().filter(|n| predicted[*n] == pc).collect();
let truth_members: Vec<&String> = pool.iter().filter(|n| truth[*n] == gc).collect();
let correct = pred_members.iter().filter(|n| truth[**n] == gc).count() as f64;
prec_sum += correct / pred_members.len() as f64;
rec_sum += correct / truth_members.len() as f64;
}
report.bcubed_precision = prec_sum / pool.len() as f64;
report.bcubed_recall = rec_sum / pool.len() as f64;
report
}
#[test]
fn entity_clustering_meets_gates() {
let report = run_eval();
assert_eq!(
report.wrong_merges, 0,
"a suggested cluster mixed two distinct ground-truth people: {report:?}"
);
assert_eq!(
report.drift_groups_recovered, report.drift_groups_total,
"not all drift groups were co-clustered: {report:?}"
);
assert!(
report.bcubed_precision >= 0.99,
"B-cubed precision regressed: {report:?}"
);
assert!(
report.bcubed_recall >= 0.99,
"B-cubed recall regressed: {report:?}"
);
assert_eq!(
report.ambiguous_suggested, report.ambiguous_total,
"ambiguous short-name pairs should all be suggested (recall-oriented): {report:?}"
);
}
#[test]
fn every_suggested_cluster_is_a_clique() {
let mut pool: Vec<String> = build_pool().0;
pool.extend(["jon", "jan", "jana"].iter().map(|s| s.to_string()));
let clusters = cluster_names(&pool);
for c in &clusters {
for i in 0..c.len() {
for j in (i + 1)..c.len() {
assert!(
names_plausibly_same_person(&pool[c[i]], &pool[c[j]]).is_some(),
"cluster is not a clique (transitive bridge): {:?}",
c.iter().map(|&k| pool[k].as_str()).collect::<Vec<_>>()
);
}
}
}
}