1use ndarray::{Array2, ArrayView2, Axis};
8
9const AUX_DISCRETE_MAX_LEVELS: usize = 64;
15
16const AUX_LEVEL_DEDUP_TOL: f64 = 1.0e-12;
20
21#[derive(Debug, Clone)]
23pub struct AuxRichnessMetrics {
24 pub aux_observed: bool,
26 pub n_nonfinite_aux: usize,
28 pub aux_dim: usize,
30 pub latent_dim: usize,
32 pub n_rows: usize,
34 pub constant_columns: Vec<usize>,
36 pub aux_is_discrete: bool,
38 pub n_distinct_levels: usize,
40 pub jacobian_rank: usize,
43 pub jacobian_rank_estimated: bool,
45}
46
47pub fn aux_richness_metrics(aux: ArrayView2<f64>, latents: ArrayView2<f64>) -> AuxRichnessMetrics {
54 let (n, aux_dim) = aux.dim();
55 let (n_z, latent_dim) = latents.dim();
56 assert_eq!(n, n_z, "aux and latents must share row count");
57
58 let mut n_nonfinite_aux: usize = 0;
60 for &v in aux.iter() {
61 if !v.is_finite() {
62 n_nonfinite_aux += 1;
63 }
64 }
65 let aux_observed = n_nonfinite_aux == 0;
66
67 let mut constant_columns: Vec<usize> = Vec::new();
70 if aux_observed && n >= 1 {
71 for j in 0..aux_dim {
72 let col = aux.column(j);
73 let mean: f64 = col.sum() / n as f64;
75 let mut var = 0.0_f64;
76 for &v in col.iter() {
77 let d = v - mean;
78 var += d * d;
79 }
80 var /= n as f64;
81 if var <= 1.0e-24 {
82 constant_columns.push(j);
83 }
84 }
85 }
86
87 let (aux_is_discrete, n_distinct_levels) = if aux_observed && n >= 1 {
89 let mut discrete = true;
90 for &v in aux.iter() {
91 if (v - v.round()).abs() > 0.0 {
92 discrete = false;
93 break;
94 }
95 }
96 if discrete {
97 for j in 0..aux_dim {
98 let col = aux.column(j);
99 let mut sorted: Vec<f64> = col.iter().copied().collect();
100 sorted.sort_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal));
101 sorted.dedup_by(|a, b| (*a - *b).abs() < AUX_LEVEL_DEDUP_TOL);
102 if sorted.len() > AUX_DISCRETE_MAX_LEVELS {
103 discrete = false;
104 break;
105 }
106 }
107 }
108 if discrete {
109 let mut keys: Vec<Vec<i64>> = Vec::with_capacity(n);
111 for i in 0..n {
112 let mut row = Vec::with_capacity(aux_dim);
113 for j in 0..aux_dim {
114 row.push(aux[[i, j]].round() as i64);
115 }
116 keys.push(row);
117 }
118 keys.sort();
119 keys.dedup();
120 (true, keys.len())
121 } else {
122 (false, 0)
123 }
124 } else {
125 (false, 0)
126 };
127
128 let need_rows = aux_dim.max(latent_dim) + 1;
130 let mut jacobian_rank_estimated = false;
131 let mut jacobian_rank: usize = usize::MAX;
132 let z_finite = latents.iter().all(|v| v.is_finite());
133 if aux_observed && z_finite && n >= need_rows && aux_dim >= 1 && latent_dim >= 1 {
134 let mut a_c = aux.to_owned();
136 let mut z_c = latents.to_owned();
137 let a_mean = a_c
138 .mean_axis(Axis(0))
139 .expect("the n >= need_rows >= 1 guard above rules out an empty axis");
140 let z_mean = z_c
141 .mean_axis(Axis(0))
142 .expect("the n >= need_rows >= 1 guard above rules out an empty axis");
143 for mut row in a_c.rows_mut() {
144 row -= &a_mean;
145 }
146 for mut row in z_c.rows_mut() {
147 row -= &z_mean;
148 }
149 let ata = a_c.t().dot(&a_c);
151 let atz = a_c.t().dot(&z_c);
152 if let Ok(b_hat) = pinv_solve(ata.view(), atz.view())
156 && let Ok(rank) = matrix_rank(b_hat.view(), 1.0e-8)
157 {
158 jacobian_rank = rank;
159 jacobian_rank_estimated = true;
160 }
161 }
162
163 AuxRichnessMetrics {
164 aux_observed,
165 n_nonfinite_aux,
166 aux_dim,
167 latent_dim,
168 n_rows: n,
169 constant_columns,
170 aux_is_discrete,
171 n_distinct_levels,
172 jacobian_rank,
173 jacobian_rank_estimated,
174 }
175}
176
177fn pinv_solve(a: ArrayView2<f64>, b: ArrayView2<f64>) -> Result<Array2<f64>, String> {
181 let (m, n) = a.dim();
182 assert_eq!(m, n, "pinv_solve expects a square normal-equation matrix");
183 let (eigvals, eigvecs) = symmetric_eigen_lower(a)?;
187 let max_abs = eigvals.iter().fold(0.0_f64, |acc, &v| acc.max(v.abs()));
188 let tol = 1.0e-12 * max_abs.max(1.0);
189 let k = eigvals.len();
191 let mut inv_diag = vec![0.0_f64; k];
192 for i in 0..k {
193 if eigvals[i].abs() > tol {
194 inv_diag[i] = 1.0 / eigvals[i];
195 }
196 }
197 let vtb = eigvecs.t().dot(&b);
199 let mut dvtb = vtb.clone();
200 for i in 0..k {
201 let scale = inv_diag[i];
202 for j in 0..dvtb.ncols() {
203 dvtb[[i, j]] *= scale;
204 }
205 }
206 Ok(eigvecs.dot(&dvtb))
207}
208
209fn symmetric_eigen_lower(a: ArrayView2<f64>) -> Result<(Vec<f64>, Array2<f64>), String> {
216 let (values, vectors) = gam_linalg::faer_ndarray::FaerEigh::eigh(&a, faer::Side::Lower)
217 .map_err(|error| format!("identifiability eigendecomposition: {error}"))?;
218 Ok((values.to_vec(), vectors))
219}
220
221fn matrix_rank(m: ArrayView2<f64>, tol: f64) -> Result<usize, String> {
225 let gram = m.t().dot(&m);
226 let (eigvals, _) = symmetric_eigen_lower(gram.view())?;
227 let mut rank = 0usize;
228 for &lam in eigvals.iter() {
229 if lam.max(0.0).sqrt() > tol {
230 rank += 1;
231 }
232 }
233 Ok(rank)
234}
235
236#[derive(Debug, Clone)]
238pub struct JacobianSparsityMetrics {
239 pub n_samples: usize,
241 pub p_features: usize,
242 pub latent_dim: usize,
243 pub mean_sparsity: f64,
246 pub max_abs: f64,
248 pub ranks: Vec<usize>,
250}
251
252pub fn jacobian_sparsity_metrics(
257 jacobians_flat: ArrayView2<f64>,
258 n_samples: usize,
259 zero_threshold: f64,
260) -> Result<JacobianSparsityMetrics, String> {
261 let (np_rows, latent_dim) = jacobians_flat.dim();
262 assert!(np_rows % n_samples == 0, "rows not divisible by n_samples");
263 let p_features = np_rows / n_samples;
264
265 let mut max_abs = 0.0_f64;
267 for &v in jacobians_flat.iter() {
268 let a = v.abs();
269 if a > max_abs {
270 max_abs = a;
271 }
272 }
273 let cutoff = zero_threshold * max_abs;
274
275 let mut total_near_zero: usize = 0;
276 let total_entries = np_rows * latent_dim;
277 if max_abs > 0.0 {
278 for &v in jacobians_flat.iter() {
279 if v.abs() < cutoff {
280 total_near_zero += 1;
281 }
282 }
283 } else {
284 total_near_zero = total_entries;
286 }
287 let mean_sparsity = if total_entries > 0 {
288 total_near_zero as f64 / total_entries as f64
289 } else {
290 0.0
291 };
292
293 let mut ranks = Vec::with_capacity(n_samples);
295 for s in 0..n_samples {
296 let start = s * p_features;
297 let end = start + p_features;
298 let view = jacobians_flat.slice(ndarray::s![start..end, ..]);
299 ranks.push(matrix_rank(view, cutoff)?);
302 }
303
304 Ok(JacobianSparsityMetrics {
305 n_samples,
306 p_features,
307 latent_dim,
308 mean_sparsity,
309 max_abs,
310 ranks,
311 })
312}
313
314#[derive(Debug, Clone)]
316pub struct AnchorConsistencyMetrics {
317 pub n_rows: usize,
319 pub n_atoms: usize,
321 pub n_anchors: usize,
324 pub anchors_per_atom: Vec<usize>,
327}
328
329fn anchor_consistency_metrics(
335 assignments: ArrayView2<f64>,
336 anchor_dominance: f64,
337) -> AnchorConsistencyMetrics {
338 let (n, k) = assignments.dim();
339 let mut anchors_per_atom = vec![0_usize; k];
340 let mut n_anchors = 0_usize;
341 for i in 0..n {
342 let row = assignments.row(i);
343 let mut mass = 0.0_f64;
344 let mut max_val = 0.0_f64;
345 let mut max_j = 0_usize;
346 for j in 0..k {
347 let a = row[j].abs();
348 mass += a;
349 if a > max_val {
350 max_val = a;
351 max_j = j;
352 }
353 }
354 if mass > 0.0 && max_val / mass >= anchor_dominance {
355 n_anchors += 1;
356 anchors_per_atom[max_j] += 1;
357 }
358 }
359 AnchorConsistencyMetrics {
360 n_rows: n,
361 n_atoms: k,
362 n_anchors,
363 anchors_per_atom,
364 }
365}
366
367#[derive(Debug, Clone, PartialEq, Eq)]
385pub struct AnchorConsistencyPreconditions {
386 pub enough_anchors_total: bool,
388 pub anchors_cover_all_atoms: bool,
390}
391
392#[derive(Debug, Clone)]
393pub struct AnchorConsistencyReport {
394 pub metrics: AnchorConsistencyMetrics,
396 pub anchor_dominance: f64,
398 pub anchor_fraction: f64,
400 pub preconditions: AnchorConsistencyPreconditions,
402 pub violations: Vec<String>,
404 pub recommendations: Vec<String>,
406 pub uncovered_atoms: Vec<usize>,
408}
409
410impl AnchorConsistencyReport {
411 pub fn passes(&self) -> bool {
413 self.preconditions.enough_anchors_total && self.preconditions.anchors_cover_all_atoms
414 }
415}
416
417pub const ANCHOR_DOMINANCE_DEFAULT: f64 = f64::from_bits(0.5_f64.to_bits() + 1);
426
427pub fn anchor_consistency_report(
431 assignments: ArrayView2<f64>,
432 anchor_dominance: Option<f64>,
433) -> Result<AnchorConsistencyReport, String> {
434 if let Some(((row, atom), value)) = assignments
435 .indexed_iter()
436 .find(|(_, value)| !value.is_finite())
437 {
438 return Err(format!(
439 "assignments must be finite; entry ({row}, {atom}) is {value}"
440 ));
441 }
442 let anchor_dominance = anchor_dominance.unwrap_or(ANCHOR_DOMINANCE_DEFAULT);
443 if !(anchor_dominance > 0.5 && anchor_dominance <= 1.0) {
444 return Err(format!(
445 "anchor_dominance must be in (0.5, 1]; got {anchor_dominance}"
446 ));
447 }
448 let (_, k) = assignments.dim();
449 if k < 1 {
450 return Err("assignments must have at least one atom column".to_string());
451 }
452 let metrics = anchor_consistency_metrics(assignments, anchor_dominance);
453 let anchor_fraction = metrics.n_anchors as f64 / metrics.n_rows.max(1) as f64;
454
455 let mut violations = Vec::new();
456 let mut recommendations = Vec::new();
457 let mut uncovered_atoms = Vec::new();
458
459 let preconditions = if k == 1 {
460 AnchorConsistencyPreconditions {
461 enough_anchors_total: true,
462 anchors_cover_all_atoms: true,
463 }
464 } else {
465 let enough_anchors = metrics.n_anchors >= k;
466 if !enough_anchors {
467 violations.push(format!(
468 "Only {} anchor row(s) (dominance >= {:.2}) found in a K={}-atom \
469 model; need at least {}. The recovered atoms are identified only \
470 up to a linear transformation in atom space.",
471 metrics.n_anchors, anchor_dominance, k, k
472 ));
473 recommendations.push(format!(
474 "Reduce K to <= {}, sharpen the assignment prior (e.g. lower \
475 temperature / stronger IBP concentration), or collect more \
476 anchor-like rows where a single atom dominates.",
477 metrics.n_anchors.max(1)
478 ));
479 }
480 uncovered_atoms = metrics
481 .anchors_per_atom
482 .iter()
483 .enumerate()
484 .filter_map(|(j, &count)| (count == 0).then_some(j))
485 .collect();
486 let cover_ok = uncovered_atoms.is_empty();
487 if !cover_ok {
488 violations.push(format!(
489 "Atom(s) {:?} have zero anchor rows; they are not individually \
490 identifiable and may be redundant or merged with neighbours.",
491 uncovered_atoms
492 ));
493 recommendations.push(format!(
494 "Prune the {} uncovered atom(s) (refit with K={}) or strengthen \
495 the per-atom sparsity prior so that each atom acquires a \
496 dominant region.",
497 uncovered_atoms.len(),
498 (k - uncovered_atoms.len()).max(1)
499 ));
500 }
501 AnchorConsistencyPreconditions {
502 enough_anchors_total: enough_anchors,
503 anchors_cover_all_atoms: cover_ok,
504 }
505 };
506
507 Ok(AnchorConsistencyReport {
508 metrics,
509 anchor_dominance,
510 anchor_fraction,
511 preconditions,
512 violations,
513 recommendations,
514 uncovered_atoms,
515 })
516}
517
518pub fn concat_decoder_blocks(blocks: &[ArrayView2<f64>]) -> Result<Array2<f64>, String> {
524 if blocks.is_empty() {
525 return Err("concat_decoder_blocks: empty block list".into());
526 }
527 let p = blocks[0].ncols();
528 for (i, b) in blocks.iter().enumerate() {
529 if b.ncols() != p {
530 return Err(format!(
531 "concat_decoder_blocks: block {} has {} cols, expected {}",
532 i,
533 b.ncols(),
534 p
535 ));
536 }
537 }
538 let total_k: usize = blocks.iter().map(|b| b.nrows()).sum();
539 let mut out = Array2::<f64>::zeros((p, total_k));
540 let mut col = 0_usize;
541 for b in blocks {
542 for k in 0..b.nrows() {
544 for row in 0..p {
545 out[[row, col]] = b[[k, row]];
546 }
547 col += 1;
548 }
549 }
550 Ok(out)
551}
552
553#[cfg(test)]
554mod tests {
555 use super::*;
556 use ndarray::array;
557
558 #[test]
559 fn aux_richness_passes_on_rich_2d_aux() {
560 let aux = array![
561 [0.0, 0.0],
562 [0.0, 1.0],
563 [1.0, 0.0],
564 [1.0, 1.0],
565 [2.0, 0.0],
566 [2.0, 1.0],
567 [0.0, 2.0],
568 [1.0, 2.0],
569 [2.0, 2.0],
570 ];
571 let lat = array![
572 [0.10, 0.05],
573 [0.02, 1.01],
574 [1.05, 0.04],
575 [1.01, 1.02],
576 [2.03, 0.07],
577 [2.04, 1.01],
578 [0.05, 2.02],
579 [1.02, 2.01],
580 [2.01, 2.05],
581 ];
582 let m = aux_richness_metrics(aux.view(), lat.view());
583 assert!(m.aux_observed);
584 assert_eq!(m.aux_dim, 2);
585 assert_eq!(m.latent_dim, 2);
586 assert!(m.constant_columns.is_empty());
587 assert!(m.aux_is_discrete);
588 assert!(m.n_distinct_levels >= 3);
589 assert!(m.jacobian_rank_estimated);
590 assert_eq!(m.jacobian_rank, 2);
591 }
592
593 #[test]
594 fn aux_richness_flags_constant_aux() {
595 let aux = Array2::<f64>::zeros((20, 1));
596 let mut lat = Array2::<f64>::zeros((20, 2));
597 for i in 0..20 {
598 lat[[i, 0]] = i as f64;
599 lat[[i, 1]] = (i as f64).cos();
600 }
601 let m = aux_richness_metrics(aux.view(), lat.view());
602 assert_eq!(m.aux_dim, 1);
603 assert_eq!(m.latent_dim, 2);
604 assert_eq!(m.constant_columns, vec![0_usize]);
605 }
606
607 #[test]
608 fn aux_richness_flags_nonfinite_aux() {
609 let mut aux = Array2::<f64>::zeros((10, 1));
610 aux[[3, 0]] = f64::NAN;
611 let lat = Array2::<f64>::zeros((10, 1));
612 let m = aux_richness_metrics(aux.view(), lat.view());
613 assert!(!m.aux_observed);
614 assert_eq!(m.n_nonfinite_aux, 1);
615 }
616
617 #[test]
618 fn jacobian_sparsity_passes_on_diagonal() {
619 let j = array![
621 [1.0_f64, 0.0, 0.0],
622 [0.0, 1.0, 0.0],
623 [0.0, 0.0, 1.0],
624 [0.0, 0.0, 0.0]
625 ];
626 let m = jacobian_sparsity_metrics(j.view(), 1, 1.0e-3).expect("sparsity metrics");
627 assert_eq!(m.p_features, 4);
628 assert_eq!(m.latent_dim, 3);
629 assert!(m.mean_sparsity > 0.5);
630 assert_eq!(m.ranks, vec![3_usize]);
631 }
632
633 #[test]
634 fn jacobian_sparsity_dense_has_low_sparsity() {
635 let mut j = Array2::<f64>::zeros((4, 3));
636 for i in 0..4 {
637 for k in 0..3 {
638 j[[i, k]] = 1.0 + 0.1 * (i + k) as f64;
639 }
640 }
641 let m = jacobian_sparsity_metrics(j.view(), 1, 1.0e-3).expect("sparsity metrics");
642 assert!(m.mean_sparsity < 0.1);
643 }
644
645 #[test]
646 fn anchor_consistency_three_clusters() {
647 let mut a = Array2::<f64>::from_elem((9, 3), 0.01);
648 for i in 0..3 {
649 a[[i, 0]] = 1.0;
650 }
651 for i in 3..6 {
652 a[[i, 1]] = 1.0;
653 }
654 for i in 6..9 {
655 a[[i, 2]] = 1.0;
656 }
657 let m = anchor_consistency_metrics(a.view(), 0.95);
658 assert_eq!(m.n_atoms, 3);
659 assert_eq!(m.n_anchors, 9);
660 assert_eq!(m.anchors_per_atom, vec![3, 3, 3]);
661 }
662
663 #[test]
664 fn anchor_consistency_uniform_has_zero_anchors() {
665 let a = Array2::<f64>::from_elem((10, 4), 0.25);
666 let m = anchor_consistency_metrics(a.view(), 0.95);
667 assert_eq!(m.n_anchors, 0);
668 assert_eq!(m.anchors_per_atom, vec![0, 0, 0, 0]);
669 }
670
671 #[test]
672 fn anchor_consistency_report_owns_the_pass_fail_verdict() {
673 let a = array![[1.0_f64, 0.0, 0.0], [0.0, 1.0, 0.0], [0.0, 0.0, 1.0]];
674 let report = anchor_consistency_report(a.view(), None).unwrap();
675 assert_eq!(report.anchor_dominance, ANCHOR_DOMINANCE_DEFAULT);
676 assert!(report.passes());
677 assert_eq!(
678 report.preconditions,
679 AnchorConsistencyPreconditions {
680 enough_anchors_total: true,
681 anchors_cover_all_atoms: true,
682 }
683 );
684 assert!(report.uncovered_atoms.is_empty());
685 }
686
687 #[test]
688 fn anchor_consistency_report_derives_thresholds_from_atom_count() {
689 let a = Array2::<f64>::from_elem((7, 4), 0.25);
690 let report = anchor_consistency_report(a.view(), Some(0.95)).unwrap();
691 assert!(!report.passes());
692 assert_eq!(
693 report.preconditions,
694 AnchorConsistencyPreconditions {
695 enough_anchors_total: false,
696 anchors_cover_all_atoms: false,
697 }
698 );
699 assert_eq!(report.uncovered_atoms, vec![0, 1, 2, 3]);
700 assert_eq!(report.violations.len(), 2);
701 assert_eq!(report.recommendations.len(), report.violations.len());
702 assert!(report.violations[0].contains("need at least 4"));
703 }
704
705 #[test]
706 fn anchor_consistency_report_rejects_invalid_dominance() {
707 let a = Array2::<f64>::ones((2, 2));
708 let error = anchor_consistency_report(a.view(), Some(0.0)).unwrap_err();
709 assert!(error.contains("anchor_dominance must be in (0.5, 1]"));
710 }
711
712 #[test]
713 fn default_anchor_rule_is_theorem_derived_strict_majority() {
714 let tied = array![[0.5_f64, 0.5], [0.5, 0.5]];
715 let tied_report = anchor_consistency_report(tied.view(), None).unwrap();
716 assert_eq!(tied_report.metrics.n_anchors, 0);
717
718 let majority = array![[0.5_f64.next_up(), 0.5], [0.5, 0.5_f64.next_up()]];
719 let majority_report = anchor_consistency_report(majority.view(), None).unwrap();
720 assert_eq!(majority_report.metrics.n_anchors, 2);
721 assert!(majority_report.passes());
722 }
723
724 #[test]
725 fn anchor_consistency_report_rejects_non_finite_assignments() {
726 for value in [f64::NAN, f64::INFINITY, f64::NEG_INFINITY] {
727 let assignments = array![[1.0_f64, 0.0], [0.0, value]];
728 let error = anchor_consistency_report(assignments.view(), None).unwrap_err();
729 assert!(error.contains("assignments must be finite"));
730 assert!(error.contains("(1, 1)"));
731 }
732 }
733
734 #[test]
735 fn anchor_consistency_report_distinguishes_count_from_atom_coverage() {
736 let a = array![
737 [1.0_f64, 0.0, 0.0],
738 [1.0, 0.0, 0.0],
739 [1.0, 0.0, 0.0],
740 [0.0, 1.0, 0.0],
741 ];
742 let report = anchor_consistency_report(a.view(), None).unwrap();
743 assert!(report.preconditions.enough_anchors_total);
744 assert!(!report.preconditions.anchors_cover_all_atoms);
745 assert_eq!(report.uncovered_atoms, vec![2]);
746 assert!(!report.passes());
747 assert_eq!(report.violations.len(), 1);
748 }
749}