use faer::{Accum, Mat, linalg::matmul::matmul};
use rayon::prelude::*;
use crate::prelude::*;
use crate::utils::faer_parallelism;
use super::{LdaCorpus, LdaResult, LdaSweepEntry};
const COHERENCE_TOP_N: usize = 20;
const COHERENCE_EPS: f64 = 1e-12;
const ARUN_FLOOR: f64 = 1e-12;
#[derive(Clone, Debug)]
pub struct LdaMetrics<F: BixverseFloat> {
pub arun_2010: F,
pub cao_juan_2009: F,
pub mimno_2011: F,
pub coherence_per_topic: Vec<F>,
pub bound: F,
pub perplexity: F,
}
fn normalise_to_sum(v: &mut [f64]) {
let total: f64 = v.iter().sum();
if total > 0.0 {
v.iter_mut().for_each(|x| *x /= total);
}
}
fn sorted_intersection_len(a: &[u32], b: &[u32]) -> usize {
let (mut i, mut j, mut count) = (0, 0, 0);
while i < a.len() && j < b.len() {
match a[i].cmp(&b[j]) {
std::cmp::Ordering::Less => i += 1,
std::cmp::Ordering::Greater => j += 1,
std::cmp::Ordering::Equal => {
count += 1;
i += 1;
j += 1;
}
}
}
count
}
fn rescale(values: &[f64]) -> Vec<f64> {
let min = values.iter().cloned().fold(f64::INFINITY, f64::min);
let max = values.iter().cloned().fold(f64::NEG_INFINITY, f64::max);
let range = max - min;
if !range.is_finite() || range <= 0.0 {
return vec![0.0; values.len()];
}
values.iter().map(|v| (v - min) / range).collect()
}
pub fn cao_juan_2009<F>(topic_region: &Mat<F>) -> F
where
F: BixverseFloat + Send + Sync,
{
let k = topic_region.ncols();
if k < 2 {
return F::zero();
}
let mut gram = Mat::<F>::zeros(k, k);
matmul(
gram.as_mut(),
Accum::Replace,
topic_region.as_ref().transpose(),
topic_region.as_ref(),
F::one(),
faer_parallelism(),
);
let mut total = F::zero();
for i in 0..k {
for j in (i + 1)..k {
let denom = (gram[(i, i)] * gram[(j, j)]).sqrt();
if denom > F::zero() {
total += gram[(i, j)] / denom;
}
}
}
total / F::from_usize(k * (k - 1) / 2).unwrap()
}
pub fn arun_2010<F>(
topic_region: &Mat<F>,
cell_topic: &Mat<F>,
doc_lengths: &[F],
) -> Result<F, BixverseErrors>
where
F: BixverseFloat + BixverseSimd + Send + Sync,
{
let k = topic_region.ncols();
let n_docs = cell_topic.ncols();
if cell_topic.nrows() != k {
return Err(BixverseErrors::LdaDimensionMismatch {
expected: k,
got: cell_topic.nrows(),
});
}
if doc_lengths.len() != n_docs {
return Err(BixverseErrors::LdaDimensionMismatch {
expected: n_docs,
got: doc_lengths.len(),
});
}
let svd = topic_region
.thin_svd()
.map_err(|e| BixverseErrors::FaerSvdError(format!("{e:?}")))?;
let mut cm1: Vec<f64> = svd
.S()
.column_vector()
.iter()
.map(|v| v.to_f64().unwrap_or(0.0).max(ARUN_FLOOR))
.collect();
let mut cm2_f = vec![F::zero(); k];
for d in 0..n_docs {
let col = cell_topic.col_as_slice(d);
F::bxv_axpy_simd(&mut cm2_f, doc_lengths[d], col);
}
let mut cm2: Vec<f64> = cm2_f
.iter()
.map(|v| v.to_f64().unwrap_or(0.0).max(ARUN_FLOOR))
.collect();
cm2.sort_by(|a, b| b.total_cmp(a));
normalise_to_sum(&mut cm1);
normalise_to_sum(&mut cm2);
let divergence: f64 = cm1
.iter()
.zip(&cm2)
.map(|(a, b)| a * (a / b).ln() + b * (b / a).ln())
.sum();
Ok(F::from_f64(divergence).unwrap_or(F::infinity()))
}
pub fn mimno_2011_coherence<F>(
corpus: &LdaCorpus<F>,
topic_region: &Mat<F>,
top_n: usize,
) -> Result<Vec<F>, BixverseErrors>
where
F: BixverseFloat + BixverseNumeric + Send + Sync,
{
let n_terms = corpus.n_terms;
if top_n > n_terms {
return Err(BixverseErrors::LdaTopNTooLarge {
requested: top_n,
vocab_size: n_terms,
});
}
if topic_region.nrows() != n_terms {
return Err(BixverseErrors::LdaDimensionMismatch {
expected: n_terms,
got: topic_region.nrows(),
});
}
let csc = &corpus.csc;
let pair_norm = 2.0 / (top_n * (top_n - 1)) as f64;
Ok((0..topic_region.ncols())
.into_par_iter()
.map(|topic| {
let col = topic_region.col_as_slice(topic);
let mut ranked: Vec<usize> = (0..n_terms).collect();
ranked.sort_by(|&a, &b| col[b].total_cmp(&col[a]));
let top: Vec<usize> = ranked.into_iter().take(top_n).collect();
let docs: Vec<&[u32]> = top
.iter()
.map(|&w| {
let start = csc.indptr[w] as usize;
let end = csc.indptr[w + 1] as usize;
&csc.indices[start..end]
})
.collect();
let mut total = 0.0_f64;
for m in 1..top_n {
for l in 0..m {
let co = sorted_intersection_len(docs[m], docs[l]) as f64;
let df = docs[l].len() as f64;
if df > 0.0 {
total += ((co + COHERENCE_EPS) / df).ln();
}
}
}
F::from_f64(total * pair_norm).unwrap_or(F::zero())
})
.collect())
}
pub fn lda_metrics<F>(
corpus: &LdaCorpus<F>,
model: &LdaResult<F>,
top_topics_coh: usize,
) -> Result<LdaMetrics<F>, BixverseErrors>
where
F: BixverseFloat + BixverseNumeric + BixverseSimd + Send + Sync,
{
let top_n = COHERENCE_TOP_N.min(corpus.n_terms);
let coherence_per_topic = mimno_2011_coherence(corpus, &model.topic_region, top_n)?;
let mut sorted = coherence_per_topic.clone();
sorted.sort_by(|a, b| b.total_cmp(a));
let take = top_topics_coh.clamp(1, sorted.len());
let mimno_2011 =
sorted[..take].iter().fold(F::zero(), |a, b| a + *b) / F::from_usize(take).unwrap();
Ok(LdaMetrics {
arun_2010: arun_2010(&model.topic_region, &model.cell_topic, &corpus.doc_lengths)?,
cao_juan_2009: cao_juan_2009(&model.topic_region),
mimno_2011,
coherence_per_topic,
bound: model.bound,
perplexity: model.perplexity,
})
}
pub fn select_best_k<F>(entries: &[LdaSweepEntry<F>], min_topics_coh: usize) -> (Vec<F>, usize)
where
F: BixverseFloat,
{
let eligible: Vec<usize> = (0..entries.len())
.filter(|&i| entries[i].k >= min_topics_coh)
.collect();
let eligible = if eligible.is_empty() {
(0..entries.len()).collect()
} else {
eligible
};
let pull = |f: &dyn Fn(&LdaSweepEntry<F>) -> f64| -> Vec<f64> {
eligible.iter().map(|&i| f(&entries[i])).collect()
};
let arun = rescale(&pull(&|e| -e.metrics.arun_2010.to_f64().unwrap_or(0.0)));
let cao = rescale(&pull(&|e| -e.metrics.cao_juan_2009.to_f64().unwrap_or(0.0)));
let coh = rescale(&pull(&|e| e.metrics.mimno_2011.to_f64().unwrap_or(0.0)));
let bound = rescale(&pull(&|e| e.metrics.bound.to_f64().unwrap_or(0.0)));
let mut combined = vec![F::nan(); entries.len()];
let mut best = (f64::NEG_INFINITY, entries[0].k);
for (slot, &i) in eligible.iter().enumerate() {
let score = arun[slot] + cao[slot] + coh[slot] + bound[slot];
combined[i] = F::from_f64(score).unwrap_or(F::nan());
if score > best.0 {
best = (score, entries[i].k);
}
}
(combined, best.1)
}
#[cfg(test)]
mod tests {
use super::*;
use crate::methods::lda::vb::tests::{block_corpus, csr_from_rows};
use crate::methods::lda::{DEFAULT_MIN_TOPICS_COH, LdaParams, lda_fit, lda_k_sweep};
use approx::assert_relative_eq;
fn topic_matrix(cols: &[Vec<f64>]) -> Mat<f64> {
Mat::from_fn(cols[0].len(), cols.len(), |w, t| cols[t][w])
}
fn test_params() -> LdaParams<f64> {
LdaParams {
alpha: 0.1,
alpha_by_topic: false,
max_iter: 40,
check_every: 1,
tol: 1e-9,
seed: 42,
..Default::default()
}
}
#[test]
fn test_cao_juan_orthogonal_topics() {
let m = topic_matrix(&[vec![0.5, 0.5, 0.0, 0.0], vec![0.0, 0.0, 0.5, 0.5]]);
assert_relative_eq!(cao_juan_2009(&m), 0.0, epsilon = 1e-12);
}
#[test]
fn test_cao_juan_identical_topics() {
let col = vec![0.4, 0.3, 0.2, 0.1];
let m = topic_matrix(&[col.clone(), col.clone(), col]);
assert_relative_eq!(cao_juan_2009(&m), 1.0, epsilon = 1e-12);
}
#[test]
fn test_cao_juan_single_topic() {
let m = topic_matrix(&[vec![0.6, 0.4]]);
assert_eq!(cao_juan_2009(&m), 0.0);
}
#[test]
fn test_cao_juan_partial_overlap() {
let m = topic_matrix(&[vec![0.5, 0.5, 0.0, 0.0], vec![0.0, 0.5, 0.5, 0.0]]);
assert_relative_eq!(cao_juan_2009(&m), 0.5, epsilon = 1e-12);
}
#[test]
fn test_arun_is_zero_when_spectra_agree() {
let topic_region = topic_matrix(&[vec![0.5, 0.5, 0.0, 0.0], vec![0.0, 0.0, 0.5, 0.5]]);
let cell_topic = Mat::from_fn(2, 2, |t, d| if t == d { 1.0 } else { 0.0 });
let doc_lengths = vec![1.0, 1.0];
let arun = arun_2010(&topic_region, &cell_topic, &doc_lengths).unwrap();
assert_relative_eq!(arun, 0.0, epsilon = 1e-10);
}
#[test]
fn test_arun_grows_with_imbalance() {
let topic_region = topic_matrix(&[vec![0.5, 0.5, 0.0, 0.0], vec![0.0, 0.0, 0.5, 0.5]]);
let balanced = Mat::from_fn(2, 2, |t, d| if t == d { 1.0 } else { 0.0 });
let lengths_even = vec![1.0, 1.0];
let lengths_skewed = vec![50.0, 1.0];
let even = arun_2010(&topic_region, &balanced, &lengths_even).unwrap();
let skewed = arun_2010(&topic_region, &balanced, &lengths_skewed).unwrap();
assert!(skewed > even, "skewed {skewed} did not exceed even {even}");
assert!(even >= 0.0);
}
#[test]
fn test_arun_dimension_mismatch() {
let topic_region = topic_matrix(&[vec![0.5, 0.5], vec![0.5, 0.5]]);
let cell_topic = Mat::from_fn(3, 2, |_, _| 0.5);
assert!(matches!(
arun_2010(&topic_region, &cell_topic, &[1.0, 1.0]),
Err(BixverseErrors::LdaDimensionMismatch { .. })
));
let cell_topic = Mat::from_fn(2, 2, |_, _| 0.5);
assert!(matches!(
arun_2010(&topic_region, &cell_topic, &[1.0]),
Err(BixverseErrors::LdaDimensionMismatch { .. })
));
}
#[test]
fn test_coherence_perfect_co_occurrence() {
let matrix = csr_from_rows(&[vec![1.0, 1.0], vec![1.0, 1.0]]);
let corpus = LdaCorpus::<f64>::new(&matrix).unwrap();
let topic_region = topic_matrix(&[vec![0.5, 0.5]]);
let coh = mimno_2011_coherence(&corpus, &topic_region, 2).unwrap();
assert_relative_eq!(coh[0], 0.0, epsilon = 1e-9);
}
#[test]
fn test_coherence_never_co_occurring() {
let matrix = csr_from_rows(&[vec![1.0, 0.0], vec![0.0, 1.0]]);
let corpus = LdaCorpus::<f64>::new(&matrix).unwrap();
let topic_region = topic_matrix(&[vec![0.5, 0.5]]);
let coh = mimno_2011_coherence(&corpus, &topic_region, 2).unwrap();
assert_relative_eq!(coh[0], (COHERENCE_EPS / 1.0).ln(), epsilon = 1e-9);
}
#[test]
fn test_coherence_orders_topics() {
let matrix = csr_from_rows(&[
vec![1.0, 1.0, 1.0, 0.0],
vec![1.0, 1.0, 0.0, 1.0],
vec![1.0, 1.0, 1.0, 0.0],
]);
let corpus = LdaCorpus::<f64>::new(&matrix).unwrap();
let topic_region = topic_matrix(&[vec![0.5, 0.5, 0.0, 0.0], vec![0.0, 0.0, 0.5, 0.5]]);
let coh = mimno_2011_coherence(&corpus, &topic_region, 2).unwrap();
assert!(
coh[0] > coh[1],
"co-occurring topic {} scored below the disjoint one {}",
coh[0],
coh[1]
);
}
#[test]
fn test_coherence_top_n_too_large() {
let matrix = csr_from_rows(&[vec![1.0, 1.0]]);
let corpus = LdaCorpus::<f64>::new(&matrix).unwrap();
let topic_region = topic_matrix(&[vec![0.5, 0.5]]);
assert!(matches!(
mimno_2011_coherence(&corpus, &topic_region, 5),
Err(BixverseErrors::LdaTopNTooLarge { .. })
));
}
#[test]
fn test_rescale() {
let r = rescale(&[2.0, 4.0, 6.0]);
assert_relative_eq!(r[0], 0.0, epsilon = 1e-12);
assert_relative_eq!(r[1], 0.5, epsilon = 1e-12);
assert_relative_eq!(r[2], 1.0, epsilon = 1e-12);
assert_eq!(rescale(&[3.0, 3.0, 3.0]), vec![0.0, 0.0, 0.0]);
}
#[test]
fn test_metrics_average_top_topics_only() {
let matrix = block_corpus(6, 5);
let model = lda_fit(&matrix, 3, Some(test_params()), 0).unwrap();
let corpus = LdaCorpus::<f64>::new(&matrix).unwrap();
let all = lda_metrics(&corpus, &model, 3).unwrap();
let best_one = lda_metrics(&corpus, &model, 1).unwrap();
assert!(
best_one.mimno_2011 >= all.mimno_2011,
"top-1 coherence {} below the all-topic mean {}",
best_one.mimno_2011,
all.mimno_2011
);
assert_eq!(all.coherence_per_topic.len(), 3);
}
#[test]
fn test_k_sweep_scores_every_entry() {
let matrix = block_corpus(8, 6);
let ks = [2, 3, 4, 5];
let sweep = lda_k_sweep(&matrix, &ks, Some(test_params()), Some(2), 0).unwrap();
assert_eq!(sweep.entries.len(), ks.len());
assert_eq!(sweep.combined_score.len(), ks.len());
for (entry, &k) in sweep.entries.iter().zip(&ks) {
assert_eq!(entry.k, k);
assert!(entry.metrics.arun_2010.is_finite());
assert!(entry.metrics.cao_juan_2009.is_finite());
assert!(entry.metrics.bound.is_finite());
}
assert!(ks.contains(&sweep.best_k));
}
#[test]
fn test_selection_respects_topic_floor() {
let matrix = block_corpus(8, 6);
let ks = [2, 3, 5, 6];
let sweep = lda_k_sweep(&matrix, &ks, Some(test_params()), Some(2), 0).unwrap();
assert!(sweep.combined_score[0].is_nan(), "k = 2 was scored");
assert!(sweep.combined_score[1].is_nan(), "k = 3 was scored");
assert!(sweep.combined_score[2].is_finite());
assert!(sweep.best_k >= DEFAULT_MIN_TOPICS_COH);
let small = lda_k_sweep(&matrix, &[2, 3], Some(test_params()), Some(2), 0).unwrap();
assert!(small.combined_score.iter().all(|v| v.is_finite()));
assert!([2, 3].contains(&small.best_k));
}
#[test]
fn test_k_sweep_rejects_empty_ks() {
let matrix = block_corpus(4, 4);
assert!(matches!(
lda_k_sweep(&matrix, &[], Some(test_params()), None, 0),
Err(BixverseErrors::LdaInvalidTopicCount { .. })
));
}
}