1use faer::Side;
45use gam_runtime::warm_start::{Fingerprint, Fingerprinter};
46use ndarray::{Array1, Array2, ArrayView1, ArrayView2};
47use serde::{Deserialize, Serialize};
48
49use crate::arrow_schur::{ArrowFactorCache, ArrowSchurSystem};
50use crate::priority_selection::{PriorityCandidate, rank_priority_candidates};
51use gam_linalg::faer_ndarray::FaerEigh;
52use gam_linalg::lanczos::{
53 SymmetricLanczosOptions, symmetric_lanczos_eigenpairs, symmetric_lanczos_log_quadrature,
54};
55use gam_linalg::pairwise_reduce::{BASE_CHUNK, pairwise_sum};
56use gam_linalg::triangular::cholesky_solve_vector;
57use gam_math::special::bessel_i0_log_minus_abs_and_ratio;
58
59pub const ANALYTIC_LOGDET_DENSE_DIM_THRESHOLD: usize = 1024;
60const EVIDENCE_LOGDET_SLQ_PROBES: usize = 16;
61const EVIDENCE_LOGDET_LANCZOS_STEPS: usize = 32;
62const EVIDENCE_HVP_SYMMETRY_REL_TOL: f64 = 1e-8;
63const EVIDENCE_HVP_SYMMETRY_PROBES: usize = 4;
64
65#[derive(Clone, Copy)]
69pub struct EvidenceHvpLogDet<'a> {
70 pub dim: usize,
71 pub apply: &'a dyn Fn(&[f64]) -> Vec<f64>,
72}
73
74#[derive(Clone, Copy)]
76pub enum EvidenceLogDetSource<'a> {
77 FactoredArrow {
80 cache: &'a ArrowFactorCache,
81 fallback_hvp: Option<EvidenceHvpLogDet<'a>>,
82 },
83 Hvp(EvidenceHvpLogDet<'a>),
88}
89
90#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
100pub enum TopologyKind {
101 Periodic,
103 Flat,
105 Sphere,
107 Torus,
109}
110
111impl TopologyKind {
112 pub fn complexity_rank(self) -> u8 {
115 match self {
116 TopologyKind::Flat => 0,
117 TopologyKind::Periodic => 1,
118 TopologyKind::Sphere => 2,
119 TopologyKind::Torus => 3,
120 }
121 }
122}
123
124#[derive(Debug, Clone)]
127pub struct TopologyCandidate {
128 pub kind: TopologyKind,
129 pub negative_log_evidence: f64,
132 pub effective_dim: f64,
135 pub n_obs: usize,
138 pub converged: bool,
142 pub exclusion_reason: Option<String>,
145}
146
147#[derive(Debug, Clone)]
149pub struct SelectedTopology {
150 pub winner: TopologyKind,
151 pub ranking: Vec<TopologyCandidate>,
154 pub tie: bool,
158}
159
160#[derive(Debug, Clone, Copy)]
162pub struct TopologySelectOptions {
163 pub tie_tolerance: f64,
167 pub score_scale: TopologyScoreScale,
171}
172
173#[derive(Debug, Clone, Copy, PartialEq, Eq)]
175pub enum TopologyScoreScale {
176 PerObservation,
178 PerEffectiveDim,
180}
181
182#[derive(Debug, Clone, Copy, Serialize, Deserialize)]
184pub struct StackingConfig {
185 pub max_iter: usize,
189 pub kkt_tol: f64,
193}
194
195impl Default for StackingConfig {
196 fn default() -> Self {
197 Self {
198 max_iter: 256,
199 kkt_tol: f64::EPSILON.sqrt(),
200 }
201 }
202}
203
204#[derive(Debug, Clone, Copy, Serialize, Deserialize)]
206pub struct StackingCertificate {
207 pub mean_log_score: f64,
209 pub duality_gap: f64,
212 pub simplex_residual: f64,
214 pub multiplier_residual: f64,
216 pub complementarity_residual: f64,
218}
219
220impl StackingCertificate {
221 pub fn residual(&self) -> f64 {
222 self.duality_gap
223 .max(self.simplex_residual)
224 .max(self.multiplier_residual)
225 .max(self.complementarity_residual)
226 }
227}
228
229#[derive(Debug, Clone, Serialize, Deserialize)]
233pub struct StackingCheckpoint {
234 pub weights: Array1<f64>,
235 pub completed_iterations: usize,
236 density_fingerprint: Fingerprint,
237}
238
239#[derive(Debug, Clone)]
242pub enum StackingError {
243 InvalidInput {
244 message: String,
245 },
246 NumericalFailure {
247 message: String,
248 certificate: Option<StackingCertificate>,
249 checkpoint: Option<StackingCheckpoint>,
250 },
251 DidNotConverge {
252 max_iterations: usize,
253 tolerance: f64,
254 certificate: StackingCertificate,
255 checkpoint: StackingCheckpoint,
256 },
257}
258
259impl std::fmt::Display for StackingError {
260 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
261 match self {
262 Self::InvalidInput { message } => write!(f, "invalid stacking problem: {message}"),
263 Self::NumericalFailure {
264 message,
265 certificate,
266 checkpoint,
267 } => write!(
268 f,
269 "stacking numerical failure: {message} (certificate residual {}, checkpoint iterations {})",
270 certificate.map_or(f64::NAN, |value| value.residual()),
271 checkpoint
272 .as_ref()
273 .map_or(0, |value| value.completed_iterations)
274 ),
275 Self::DidNotConverge {
276 max_iterations,
277 tolerance,
278 certificate,
279 checkpoint,
280 } => write!(
281 f,
282 "stacking did not certify after {max_iterations} additional iterations (total {}): KKT residual {:.6e} exceeds tolerance {:.3e}; resume from the carried weights checkpoint",
283 checkpoint.completed_iterations,
284 certificate.residual(),
285 tolerance
286 ),
287 }
288 }
289}
290
291impl std::error::Error for StackingError {}
292
293#[derive(Debug, Clone)]
296pub struct StackingWeights {
297 pub weights: Array1<f64>,
298 pub iterations: usize,
299 pub certificate: StackingCertificate,
300}
301
302impl StackingWeights {
303 pub fn mean_log_score(&self) -> f64 {
304 self.certificate.mean_log_score
305 }
306}
307
308struct StackingProblem {
309 scaled_density: Array2<f64>,
310 row_log_scale: Array1<f64>,
311}
312
313impl StackingProblem {
314 fn from_log_density(log_density: ArrayView2<'_, f64>) -> Result<Self, StackingError> {
315 let n_obs = log_density.nrows();
316 let n_cand = log_density.ncols();
317 if n_cand == 0 || n_obs == 0 {
318 return Err(StackingError::InvalidInput {
319 message: "at least one candidate and one held-out row are required".to_string(),
320 });
321 }
322 if let Some(((row, col), value)) = log_density
323 .indexed_iter()
324 .find(|(_, value)| value.is_nan() || **value == f64::INFINITY)
325 {
326 return Err(StackingError::InvalidInput {
327 message: format!(
328 "log density at row {row}, candidate {col} is {value}; NaN and +infinity are not predictive densities"
329 ),
330 });
331 }
332 let mut scaled_density = Array2::<f64>::zeros((n_obs, n_cand));
333 let mut row_log_scale = Array1::<f64>::zeros(n_obs);
334 for row in 0..n_obs {
335 let row_max = (0..n_cand)
336 .map(|col| log_density[[row, col]])
337 .fold(f64::NEG_INFINITY, f64::max);
338 if !row_max.is_finite() {
339 return Err(StackingError::InvalidInput {
340 message: format!(
341 "held-out row {row} has zero density under every candidate; deleting it would change the stacking target"
342 ),
343 });
344 }
345 row_log_scale[row] = row_max;
346 for col in 0..n_cand {
347 let value = log_density[[row, col]];
348 if value.is_finite() {
349 scaled_density[[row, col]] = (value - row_max).exp();
350 }
351 }
352 }
353 Ok(Self {
354 scaled_density,
355 row_log_scale,
356 })
357 }
358
359 fn evaluate(
360 &self,
361 weights: ArrayView1<'_, f64>,
362 ) -> Result<(Array1<f64>, StackingCertificate, f64), String> {
363 let n = self.scaled_density.nrows();
364 let k = self.scaled_density.ncols();
365 let mass = weights.sum();
366 if weights.len() != k
367 || weights
368 .iter()
369 .any(|value| !value.is_finite() || *value < 0.0)
370 || !(mass.is_finite() && mass > 0.0)
371 {
372 return Err(
373 "checkpoint weights are not a finite nonnegative simplex vector".to_string(),
374 );
375 }
376 let mut gradient = Array1::<f64>::zeros(k);
377 let mut centered_objective = 0.0_f64;
378 let mut mean_log_score = 0.0_f64;
379 for row in 0..n {
380 let mut mixture = 0.0_f64;
381 for col in 0..k {
382 mixture += weights[col] * self.scaled_density[[row, col]];
383 }
384 if !(mixture.is_finite() && mixture > 0.0) {
385 return Err(format!(
386 "candidate mixture lost held-out row {row} (scaled density {mixture})"
387 ));
388 }
389 let log_mixture = mixture.ln();
390 centered_objective += log_mixture / n as f64;
391 let log_score = self.row_log_scale[row] + log_mixture;
392 let count = (row + 1) as f64;
393 mean_log_score = mean_log_score * ((count - 1.0) / count) + log_score / count;
394 for col in 0..k {
395 gradient[col] += self.scaled_density[[row, col]] / mixture / n as f64;
396 }
397 }
398 if !centered_objective.is_finite()
399 || !mean_log_score.is_finite()
400 || gradient.iter().any(|value| !value.is_finite())
401 {
402 return Err("objective or analytic gradient became non-finite".to_string());
403 }
404 let multiplier = weights.dot(&gradient);
405 let max_gradient = gradient.iter().copied().fold(f64::NEG_INFINITY, f64::max);
406 let certificate = StackingCertificate {
407 mean_log_score,
408 duality_gap: (max_gradient - multiplier).max(0.0),
409 simplex_residual: (mass - 1.0).abs(),
410 multiplier_residual: (multiplier - 1.0).abs(),
411 complementarity_residual: weights
412 .iter()
413 .zip(gradient.iter())
414 .map(|(&weight, &gain)| weight * (gain - multiplier).abs())
415 .fold(0.0_f64, f64::max),
416 };
417 Ok((gradient, certificate, centered_objective))
418 }
419
420 fn centered_objective(&self, weights: ArrayView1<'_, f64>) -> Option<f64> {
421 let n = self.scaled_density.nrows();
422 let mut objective = 0.0_f64;
423 for row in 0..n {
424 let mixture = self.scaled_density.row(row).dot(&weights);
425 if !(mixture.is_finite() && mixture > 0.0) {
426 return None;
427 }
428 objective += mixture.ln() / n as f64;
429 }
430 objective.is_finite().then_some(objective)
431 }
432}
433
434pub fn solve_stacking_weights(
455 log_density: ArrayView2<'_, f64>,
456 config: StackingConfig,
457) -> Result<StackingWeights, StackingError> {
458 solve_stacking_weights_impl(log_density, config, None)
459}
460
461pub fn resume_stacking_weights(
464 log_density: ArrayView2<'_, f64>,
465 config: StackingConfig,
466 checkpoint: &StackingCheckpoint,
467) -> Result<StackingWeights, StackingError> {
468 solve_stacking_weights_impl(log_density, config, Some(checkpoint))
469}
470
471fn solve_stacking_weights_impl(
472 log_density: ArrayView2<'_, f64>,
473 config: StackingConfig,
474 checkpoint: Option<&StackingCheckpoint>,
475) -> Result<StackingWeights, StackingError> {
476 if config.max_iter == 0 {
477 return Err(StackingError::InvalidInput {
478 message: "max_iter must be positive".to_string(),
479 });
480 }
481 let numerical_floor = f64::EPSILON.sqrt();
482 if !config.kkt_tol.is_finite() || config.kkt_tol < numerical_floor {
483 return Err(StackingError::InvalidInput {
484 message: format!(
485 "kkt_tol must be finite and at least the floating-point resolution floor {numerical_floor:.3e}"
486 ),
487 });
488 }
489 let density_fingerprint = evidence_matrix_fingerprint("stacking-log-density-v1", log_density);
490 let problem = StackingProblem::from_log_density(log_density)?;
491 let k = problem.scaled_density.ncols();
492 let (mut weights, completed_before) = if let Some(checkpoint) = checkpoint {
493 if checkpoint.density_fingerprint != density_fingerprint {
494 return Err(StackingError::InvalidInput {
495 message: "checkpoint belongs to a different held-out density table".to_string(),
496 });
497 }
498 if checkpoint.weights.len() != k {
499 return Err(StackingError::InvalidInput {
500 message: format!(
501 "checkpoint has {} weights but the density table has {k} candidates",
502 checkpoint.weights.len()
503 ),
504 });
505 }
506 let mut weights = checkpoint.weights.clone();
507 let mass = weights.sum();
508 if weights
509 .iter()
510 .any(|value| !value.is_finite() || *value < 0.0)
511 || !mass.is_finite()
512 || (mass - 1.0).abs() > config.kkt_tol
513 {
514 return Err(StackingError::InvalidInput {
515 message: "checkpoint weights must be a finite nonnegative simplex vector"
516 .to_string(),
517 });
518 }
519 weights.mapv_inplace(|value| value / mass);
520 (weights, checkpoint.completed_iterations)
521 } else {
522 (Array1::<f64>::from_elem(k, 1.0 / k as f64), 0)
523 };
524
525 for additional_iterations in 0..=config.max_iter {
526 let completed_iterations = completed_before + additional_iterations;
527 let checkpoint = StackingCheckpoint {
528 weights: weights.clone(),
529 completed_iterations,
530 density_fingerprint,
531 };
532 let (gradient, certificate, objective) =
533 problem.evaluate(weights.view()).map_err(|message| {
534 StackingError::NumericalFailure {
535 message,
536 certificate: None,
537 checkpoint: Some(checkpoint.clone()),
538 }
539 })?;
540 if certificate.residual() <= config.kkt_tol {
541 return Ok(StackingWeights {
542 weights,
543 iterations: completed_iterations,
544 certificate,
545 });
546 }
547 if additional_iterations == config.max_iter {
548 return Err(StackingError::DidNotConverge {
549 max_iterations: config.max_iter,
550 tolerance: config.kkt_tol,
551 certificate,
552 checkpoint,
553 });
554 }
555
556 let max_gradient_col = gradient
557 .iter()
558 .enumerate()
559 .max_by(|left, right| left.1.total_cmp(right.1))
560 .map(|(index, _)| index)
561 .expect("stacking has at least one candidate");
562 let candidate = stacking_newton_step(&problem, weights.view(), gradient.view(), objective)
563 .or_else(|| {
564 stacking_vertex_step(&problem, weights.view(), max_gradient_col, objective)
565 })
566 .ok_or_else(|| StackingError::NumericalFailure {
567 message: "positive KKT gap remained but neither the analytic Newton direction nor the exact vertex line solve produced a representable ascent step".to_string(),
568 certificate: Some(certificate),
569 checkpoint: Some(checkpoint),
570 })?;
571 weights = candidate;
572 }
573 Err(StackingError::NumericalFailure {
574 message: format!(
575 "stacking solver exhausted its inclusive iteration budget ({}) without producing a \
576 terminal verdict",
577 config.max_iter
578 ),
579 certificate: None,
580 checkpoint: None,
581 })
582}
583
584fn stacking_newton_step(
585 problem: &StackingProblem,
586 weights: ArrayView1<'_, f64>,
587 gradient: ArrayView1<'_, f64>,
588 objective: f64,
589) -> Option<Array1<f64>> {
590 let active: Vec<usize> = weights
591 .iter()
592 .enumerate()
593 .filter_map(|(index, &weight)| (weight > 0.0).then_some(index))
594 .collect();
595 if active.len() < 2 {
596 return None;
597 }
598 let reference_position = active
599 .iter()
600 .enumerate()
601 .max_by(|left, right| weights[*left.1].total_cmp(&weights[*right.1]))
602 .map(|(position, _)| position)?;
603 let reference = active[reference_position];
604 let free: Vec<usize> = active
605 .iter()
606 .copied()
607 .filter(|&index| index != reference)
608 .collect();
609 let dimension = free.len();
610 let n = problem.scaled_density.nrows();
611 let mut information = Array2::<f64>::zeros((dimension, dimension));
612 for row in 0..n {
613 let mixture = problem.scaled_density.row(row).dot(&weights);
614 if !(mixture.is_finite() && mixture > 0.0) {
615 return None;
616 }
617 let reference_density = problem.scaled_density[[row, reference]];
618 let contrasts: Vec<f64> = free
619 .iter()
620 .map(|&col| (problem.scaled_density[[row, col]] - reference_density) / mixture)
621 .collect();
622 for left in 0..dimension {
623 for right in 0..=left {
624 information[[left, right]] += contrasts[left] * contrasts[right] / n as f64;
625 information[[right, left]] = information[[left, right]];
626 }
627 }
628 }
629 let reduced_gradient =
630 Array1::from_iter(free.iter().map(|&col| gradient[col] - gradient[reference]));
631 let (eigenvalues, eigenvectors) = information.eigh(Side::Lower).ok()?;
632 let spectral_scale = eigenvalues.iter().copied().fold(0.0_f64, f64::max);
633 if !(spectral_scale.is_finite() && spectral_scale > 0.0) {
634 return None;
635 }
636 let rank_tolerance = f64::EPSILON * (dimension as f64) * spectral_scale.max(f64::MIN_POSITIVE);
637 let projected = eigenvectors.t().dot(&reduced_gradient);
638 let mut spectral_step = Array1::<f64>::zeros(dimension);
639 for index in 0..dimension {
640 if eigenvalues[index] > rank_tolerance {
641 spectral_step[index] = projected[index] / eigenvalues[index];
642 }
643 }
644 let reduced_step = eigenvectors.dot(&spectral_step);
645 let ascent = reduced_gradient.dot(&reduced_step);
646 if !(ascent.is_finite() && ascent > 0.0) {
647 return None;
648 }
649 let mut direction = Array1::<f64>::zeros(weights.len());
650 for (position, &col) in free.iter().enumerate() {
651 direction[col] = reduced_step[position];
652 }
653 direction[reference] = -reduced_step.sum();
654 let mut step = 1.0_f64;
655 let mut boundary = None;
656 for col in 0..weights.len() {
657 if direction[col] < 0.0 {
658 let candidate = -weights[col] / direction[col];
659 if candidate < step {
660 step = candidate;
661 boundary = Some(col);
662 }
663 }
664 }
665 loop {
666 let mut candidate = &weights + &(direction.mapv(|value| step * value));
667 if let Some(col) = boundary {
668 if step == -weights[col] / direction[col] {
669 candidate[col] = 0.0;
670 }
671 }
672 for value in candidate.iter_mut() {
673 if *value < 0.0 && *value >= -f64::EPSILON {
674 *value = 0.0;
675 }
676 }
677 let mass = candidate.sum();
678 if mass.is_finite() && mass > 0.0 {
679 candidate.mapv_inplace(|value| value / mass);
680 if problem
681 .centered_objective(candidate.view())
682 .is_some_and(|value| value > objective)
683 {
684 return Some(candidate);
685 }
686 }
687 let next_step = 0.5 * step;
688 if next_step == step || next_step == 0.0 {
689 return None;
690 }
691 step = next_step;
692 boundary = None;
693 }
694}
695
696fn stacking_vertex_step(
697 problem: &StackingProblem,
698 weights: ArrayView1<'_, f64>,
699 vertex: usize,
700 objective: f64,
701) -> Option<Array1<f64>> {
702 let derivative = |step: f64| -> f64 {
703 let mut value = 0.0_f64;
704 let n = problem.scaled_density.nrows();
705 for row in 0..n {
706 let current = problem.scaled_density.row(row).dot(&weights);
707 let target = problem.scaled_density[[row, vertex]];
708 let mixture = (1.0 - step) * current + step * target;
709 if mixture <= 0.0 {
710 return f64::NEG_INFINITY;
711 }
712 value += (target - current) / mixture / n as f64;
713 }
714 value
715 };
716 if derivative(0.0) <= 0.0 {
717 return None;
718 }
719 let mut step = if derivative(1.0) >= 0.0 {
720 1.0
721 } else {
722 let mut lower = 0.0_f64;
723 let mut upper = 1.0_f64;
724 while upper - lower > f64::EPSILON.sqrt() {
725 let middle = 0.5 * (lower + upper);
726 if derivative(middle) > 0.0 {
727 lower = middle;
728 } else {
729 upper = middle;
730 }
731 }
732 0.5 * (lower + upper)
733 };
734 loop {
735 let mut candidate = weights.mapv(|weight| (1.0 - step) * weight);
736 candidate[vertex] += step;
737 if problem
738 .centered_objective(candidate.view())
739 .is_some_and(|value| value > objective)
740 {
741 return Some(candidate);
742 }
743 let next_step = 0.5 * step;
744 if next_step == step || next_step == 0.0 {
745 return None;
746 }
747 step = next_step;
748 }
749}
750
751pub fn stacked_predictive_mean(
753 weights: &Array1<f64>,
754 candidate_means: &[Array1<f64>],
755) -> Result<Array1<f64>, String> {
756 if candidate_means.len() != weights.len() {
757 return Err(format!(
758 "stacked_predictive_mean: {} weights but {} candidate mean vectors",
759 weights.len(),
760 candidate_means.len()
761 ));
762 }
763 let Some(first) = candidate_means.first() else {
764 return Err("stacked_predictive_mean requires at least one candidate".to_string());
765 };
766 let n_rows = first.len();
767 if candidate_means.iter().any(|means| means.len() != n_rows) {
768 return Err(
769 "stacked_predictive_mean: candidate mean vectors disagree on row count".to_string(),
770 );
771 }
772 let mut out = Array1::<f64>::zeros(n_rows);
773 for (weight, means) in weights.iter().zip(candidate_means) {
774 if *weight != 0.0 {
775 out.scaled_add(*weight, means);
776 }
777 }
778 Ok(out)
779}
780
781#[derive(Debug, Clone, Copy, Serialize, Deserialize)]
806pub struct GaussianMixtureConfig {
807 pub max_iter: usize,
812 pub loglik_tol: f64,
814 pub parameter_tol: f64,
820 pub covariance_floor: f64,
824 pub kmeans_max_iter: usize,
826}
827
828impl Default for GaussianMixtureConfig {
829 fn default() -> Self {
830 Self {
831 max_iter: 1000,
832 loglik_tol: f64::EPSILON.sqrt(),
833 parameter_tol: f64::EPSILON.sqrt(),
834 covariance_floor: 1e-6,
835 kmeans_max_iter: 25,
836 }
837 }
838}
839
840#[derive(Debug, Clone, Copy, Serialize, Deserialize)]
843pub struct GaussianMixtureCertificate {
844 pub mean_log_likelihood: f64,
846 pub mean_log_likelihood_gain: f64,
848 pub monotonicity_uncertainty: f64,
852 pub objective_residual: f64,
853 pub objective_tolerance: f64,
854 pub parameter_residual: f64,
855 pub parameter_tolerance: f64,
856 pub contraction_rate: Option<f64>,
861 pub projected_iterations_to_tolerance: Option<usize>,
867}
868
869const EM_RATE_WINDOW: usize = 64;
881
882fn em_contraction_rate(window: &std::collections::VecDeque<f64>) -> Option<f64> {
889 if window.len() < EM_RATE_WINDOW + 1 {
890 return None;
891 }
892 let first = *window.front()?;
893 let last = *window.back()?;
894 if !(first.is_finite() && last.is_finite() && first > 0.0 && last > 0.0) {
895 return None;
896 }
897 let steps = (window.len() - 1) as f64;
898 let rate = (last / first).powf(1.0 / steps);
899 rate.is_finite().then_some(rate)
900}
901
902fn em_projected_iterations(residual: f64, tolerance: f64, rate: f64) -> Option<usize> {
910 if !(residual.is_finite() && tolerance.is_finite() && rate.is_finite()) {
911 return None;
912 }
913 if !(rate > 0.0 && rate < 1.0) || !(residual > tolerance) || tolerance <= 0.0 {
914 return None;
915 }
916 let steps = (tolerance / residual).ln() / rate.ln();
917 (steps.is_finite() && steps >= 0.0).then(|| steps.ceil() as usize)
918}
919
920#[derive(Debug, Clone, Serialize, Deserialize)]
922pub struct GaussianMixtureCheckpoint {
923 pub weights: Array1<f64>,
924 pub means: Array2<f64>,
925 pub covariances: Vec<Array2<f64>>,
926 pub mean_log_likelihood: f64,
927 pub completed_iterations: usize,
928 data_fingerprint: Fingerprint,
929 covariance_floor: f64,
930}
931
932#[derive(Debug, Clone)]
935pub enum GaussianMixtureError {
936 InvalidInput {
937 message: String,
938 },
939 NumericalFailure {
940 message: String,
941 checkpoint: Option<GaussianMixtureCheckpoint>,
942 },
943 MonotonicityViolation {
944 previous_mean_log_likelihood: f64,
945 next_mean_log_likelihood: f64,
946 numerical_uncertainty: f64,
947 checkpoint: GaussianMixtureCheckpoint,
948 },
949 DidNotConverge {
950 max_iterations: usize,
951 certificate: GaussianMixtureCertificate,
952 checkpoint: GaussianMixtureCheckpoint,
953 },
954}
955
956impl std::fmt::Display for GaussianMixtureError {
957 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
958 match self {
959 Self::InvalidInput { message } => write!(f, "invalid Gaussian mixture: {message}"),
960 Self::NumericalFailure {
961 message,
962 checkpoint,
963 } => write!(
964 f,
965 "Gaussian-mixture numerical failure: {message} (checkpoint iterations {})",
966 checkpoint
967 .as_ref()
968 .map_or(0, |value| value.completed_iterations)
969 ),
970 Self::MonotonicityViolation {
971 previous_mean_log_likelihood,
972 next_mean_log_likelihood,
973 numerical_uncertainty,
974 checkpoint,
975 } => write!(
976 f,
977 "Gaussian-mixture EM violated monotone ascent at iteration {}: mean log likelihood {previous_mean_log_likelihood:.12e} -> {next_mean_log_likelihood:.12e} (comparison uncertainty {numerical_uncertainty:.3e}); resume from the carried checkpoint only after diagnosing the numerical failure",
978 checkpoint.completed_iterations
979 ),
980 Self::DidNotConverge {
981 max_iterations,
982 certificate,
983 checkpoint,
984 } => write!(
985 f,
986 "Gaussian-mixture EM did not certify after {max_iterations} additional iterations (total {}): signed mean-log-likelihood gain {:.6e} (numerical uncertainty {:.3e}), objective residual {:.6e}/{:.3e}, parameter-map residual {:.6e}/{:.3e}, contraction rate {} per iteration, projected iterations to tolerance {}; resume from the carried checkpoint, which is not comparable evidence",
987 checkpoint.completed_iterations,
988 certificate.mean_log_likelihood_gain,
989 certificate.monotonicity_uncertainty,
990 certificate.objective_residual,
991 certificate.objective_tolerance,
992 certificate.parameter_residual,
993 certificate.parameter_tolerance,
994 match certificate.contraction_rate {
995 Some(rate) => format!("{rate:.6}"),
996 None => "unmeasured".to_string(),
997 },
998 match certificate.projected_iterations_to_tolerance {
999 Some(steps) => steps.to_string(),
1000 None => "none (not contracting)".to_string(),
1001 }
1002 ),
1003 }
1004 }
1005}
1006
1007impl std::error::Error for GaussianMixtureError {}
1008
1009#[derive(Debug, Clone)]
1011pub struct GaussianMixtureFit {
1012 weights: Array1<f64>,
1014 means: Array2<f64>,
1016 covariances: Vec<Array2<f64>>,
1018 k: usize,
1020 d: usize,
1022 n_obs: usize,
1024 loglik: f64,
1026 iterations: usize,
1028 certificate: GaussianMixtureCertificate,
1029}
1030
1031impl GaussianMixtureFit {
1032 pub fn weights(&self) -> ArrayView1<'_, f64> {
1033 self.weights.view()
1034 }
1035
1036 pub fn means(&self) -> ArrayView2<'_, f64> {
1037 self.means.view()
1038 }
1039
1040 pub fn covariances(&self) -> &[Array2<f64>] {
1041 &self.covariances
1042 }
1043
1044 pub fn iterations(&self) -> usize {
1045 self.iterations
1046 }
1047
1048 pub fn certificate(&self) -> GaussianMixtureCertificate {
1049 self.certificate
1050 }
1051
1052 pub fn num_free_parameters(&self) -> usize {
1057 let cov_per = self.d * (self.d + 1) / 2;
1058 (self.k - 1) + self.k * self.d + self.k * cov_per
1059 }
1060
1061 pub fn per_point_log_density(&self, data: ArrayView2<'_, f64>) -> Result<Array1<f64>, String> {
1065 if data.ncols() != self.d {
1066 return Err(format!(
1067 "mixture log-density expects {} columns, got {}",
1068 self.d,
1069 data.ncols()
1070 ));
1071 }
1072 let n = data.nrows();
1073 let mut comp = Vec::with_capacity(self.k);
1074 for j in 0..self.k {
1075 comp.push(GaussianComponentEval::factor(
1076 self.means.row(j),
1077 &self.covariances[j],
1078 )?);
1079 }
1080 let mut out = Array1::<f64>::zeros(n);
1081 let log_w: Vec<f64> = self.weights.iter().map(|w| w.ln()).collect();
1082 for i in 0..n {
1083 let row = data.row(i);
1084 let mut log_terms = vec![f64::NEG_INFINITY; self.k];
1085 let mut max_term = f64::NEG_INFINITY;
1086 for j in 0..self.k {
1087 let lt = log_w[j] + comp[j].log_density(row);
1088 log_terms[j] = lt;
1089 if lt > max_term {
1090 max_term = lt;
1091 }
1092 }
1093 out[i] = log_sum_exp(&log_terms, max_term);
1094 }
1095 Ok(out)
1096 }
1097
1098 pub fn bic(&self) -> f64 {
1102 -self.loglik + 0.5 * self.num_free_parameters() as f64 * (self.n_obs as f64).ln()
1103 }
1104}
1105
1106#[derive(Debug, Clone)]
1109struct GaussianComponentEval {
1110 residual_origin: Array1<f64>,
1111 residual_scale: Array1<f64>,
1112 residual_normalized_offset: Array1<f64>,
1113 precision: Array2<f64>,
1114 log_norm: f64,
1115 d: usize,
1116}
1117
1118impl GaussianComponentEval {
1119 fn factor(mean: ArrayView1<'_, f64>, cov: &Array2<f64>) -> Result<Self, String> {
1120 let d = mean.len();
1121 if mean.iter().any(|value| !value.is_finite()) {
1122 return Err("mixture component mean must be finite".to_string());
1123 }
1124 if cov.nrows() != d || cov.ncols() != d {
1125 return Err(format!(
1126 "mixture component covariance must be {d}x{d}, got {}x{}",
1127 cov.nrows(),
1128 cov.ncols()
1129 ));
1130 }
1131 let (evals, evecs) = cov
1132 .eigh(Side::Lower)
1133 .map_err(|e| format!("mixture component covariance eigendecomposition failed: {e}"))?;
1134 let mut log_det = 0.0_f64;
1135 let mut inv_evals = Array1::<f64>::zeros(d);
1136 for (idx, &ev) in evals.iter().enumerate() {
1137 if !ev.is_finite() || ev <= 0.0 {
1138 return Err(format!(
1139 "mixture component covariance is not SPD: eigenvalue {idx} is {ev:.3e}"
1140 ));
1141 }
1142 log_det += ev.ln();
1143 let inverse = ev.recip();
1144 if !inverse.is_finite() {
1145 return Err(format!(
1146 "mixture component precision is not representable: eigenvalue {idx} is {ev:.3e}"
1147 ));
1148 }
1149 inv_evals[idx] = inverse;
1150 }
1151 let mut precision = Array2::<f64>::zeros((d, d));
1153 for a in 0..d {
1154 for b in 0..d {
1155 let mut acc = 0.0_f64;
1156 for m in 0..d {
1157 acc += evecs[[a, m]] * inv_evals[m] * evecs[[b, m]];
1158 }
1159 precision[[a, b]] = acc;
1160 }
1161 }
1162 let log_norm = -0.5 * (d as f64 * (2.0 * std::f64::consts::PI).ln() + log_det);
1163 if precision.iter().any(|value| !value.is_finite()) || !log_norm.is_finite() {
1164 return Err(
1165 "mixture component factorization produced non-finite precision or log normalizer"
1166 .to_string(),
1167 );
1168 }
1169 Ok(Self {
1170 residual_origin: mean.to_owned(),
1171 residual_scale: Array1::zeros(d),
1172 residual_normalized_offset: Array1::zeros(d),
1173 precision,
1174 log_norm,
1175 d,
1176 })
1177 }
1178
1179 fn isotropic(charts: &[StableScalarMeanChart], variance: f64) -> Result<Self, String> {
1180 let d = charts.len();
1181 if d == 0 {
1182 return Err("isotropic Gaussian density requires positive dimension".to_string());
1183 }
1184 if !(variance.is_finite() && variance > 0.0) {
1185 return Err(format!(
1186 "isotropic Gaussian variance must be finite and positive, got {variance}"
1187 ));
1188 }
1189 let inverse_variance = variance.recip();
1190 if !inverse_variance.is_finite() {
1191 return Err(format!(
1192 "isotropic Gaussian precision is non-finite for variance {variance}"
1193 ));
1194 }
1195 let mut precision = Array2::<f64>::zeros((d, d));
1196 for axis in 0..d {
1197 precision[[axis, axis]] = inverse_variance;
1198 }
1199 let log_norm = -0.5 * d as f64 * ((2.0 * std::f64::consts::PI).ln() + variance.ln());
1200 if !log_norm.is_finite() {
1201 return Err("isotropic Gaussian log normalizer is non-finite".to_string());
1202 }
1203 Ok(Self {
1204 residual_origin: Array1::from_iter(charts.iter().map(|chart| chart.origin)),
1205 residual_scale: Array1::from_iter(charts.iter().map(|chart| chart.scale)),
1206 residual_normalized_offset: Array1::from_iter(
1207 charts.iter().map(|chart| chart.normalized_offset),
1208 ),
1209 precision,
1210 log_norm,
1211 d,
1212 })
1213 }
1214
1215 #[inline]
1216 fn log_density(&self, y: ArrayView1<'_, f64>) -> f64 {
1217 let residual = self.residual(y);
1218 let pv = self.precision_times_residual(&residual);
1219 let mut quad = 0.0_f64;
1220 for c in 0..self.d {
1221 quad += residual[c] * pv[c];
1222 }
1223 self.log_norm - 0.5 * quad
1224 }
1225
1226 #[inline]
1227 fn residual(&self, y: ArrayView1<'_, f64>) -> Vec<f64> {
1228 let mut residual = vec![0.0_f64; self.d];
1229 for axis in 0..self.d {
1230 residual[axis] = (-self.residual_normalized_offset[axis]).mul_add(
1231 self.residual_scale[axis],
1232 y[axis] - self.residual_origin[axis],
1233 );
1234 }
1235 residual
1236 }
1237
1238 #[inline]
1240 fn precision_times_residual(&self, residual: &[f64]) -> Vec<f64> {
1241 let mut out = vec![0.0_f64; self.d];
1242 for a in 0..self.d {
1243 let mut acc = 0.0_f64;
1244 for b in 0..self.d {
1245 acc += self.precision[[a, b]] * residual[b];
1246 }
1247 out[a] = acc;
1248 }
1249 out
1250 }
1251}
1252
1253#[inline]
1254fn log_sum_exp(terms: &[f64], max_term: f64) -> f64 {
1255 if !max_term.is_finite() {
1256 return f64::NEG_INFINITY;
1257 }
1258 let mut acc = 0.0_f64;
1259 for &t in terms {
1260 acc += (t - max_term).exp();
1261 }
1262 max_term + acc.ln()
1263}
1264
1265fn evidence_matrix_fingerprint(namespace: &str, values: ArrayView2<'_, f64>) -> Fingerprint {
1266 let mut hasher = Fingerprinter::new();
1267 hasher.write_str(namespace);
1268 hasher.write_usize(values.nrows());
1269 hasher.write_usize(values.ncols());
1270 for &value in values {
1273 hasher.write_f64(value);
1274 }
1275 hasher.finalize()
1276}
1277
1278fn mixture_data_fingerprint(data: ArrayView2<'_, f64>) -> Fingerprint {
1279 evidence_matrix_fingerprint("gaussian-mixture-em-v1", data)
1280}
1281
1282pub fn fit_gaussian_mixture(
1291 data: ArrayView2<'_, f64>,
1292 k: usize,
1293 config: GaussianMixtureConfig,
1294) -> Result<GaussianMixtureFit, GaussianMixtureError> {
1295 validate_gaussian_mixture_problem(data, k, config)?;
1296 let means = gam_terms::basis::select_centers_by_strategy(
1299 data,
1300 &gam_terms::basis::CenterStrategy::KMeans {
1301 num_centers: k,
1302 max_iter: config.kmeans_max_iter,
1303 },
1304 )
1305 .map_err(|error| GaussianMixtureError::NumericalFailure {
1306 message: format!("deterministic k-means seeding failed: {error}"),
1307 checkpoint: None,
1308 })?;
1309 if means.nrows() != k || means.ncols() != data.ncols() {
1310 return Err(GaussianMixtureError::NumericalFailure {
1311 message: format!(
1312 "seeding returned {}x{} centers, expected {k}x{}",
1313 means.nrows(),
1314 means.ncols(),
1315 data.ncols()
1316 ),
1317 checkpoint: None,
1318 });
1319 }
1320 let global_covariance =
1321 constrained_data_covariance(data, config.covariance_floor).map_err(|message| {
1322 GaussianMixtureError::NumericalFailure {
1323 message,
1324 checkpoint: None,
1325 }
1326 })?;
1327 let weights = Array1::<f64>::from_elem(k, 1.0 / k as f64);
1328 let covariances = vec![global_covariance; k];
1329 let initial_e_step =
1330 mixture_e_step(data, &weights, &means, &covariances).map_err(|message| {
1331 GaussianMixtureError::NumericalFailure {
1332 message,
1333 checkpoint: None,
1334 }
1335 })?;
1336 let data_fingerprint = mixture_data_fingerprint(data);
1337 let checkpoint = GaussianMixtureCheckpoint {
1338 weights,
1339 means,
1340 covariances,
1341 mean_log_likelihood: initial_e_step.mean_log_likelihood,
1342 completed_iterations: 0,
1343 data_fingerprint,
1344 covariance_floor: config.covariance_floor,
1345 };
1346 run_gaussian_mixture_em(data, config, checkpoint)
1347}
1348
1349pub fn resume_gaussian_mixture(
1351 data: ArrayView2<'_, f64>,
1352 config: GaussianMixtureConfig,
1353 checkpoint: GaussianMixtureCheckpoint,
1354) -> Result<GaussianMixtureFit, GaussianMixtureError> {
1355 let k = checkpoint.weights.len();
1356 validate_gaussian_mixture_problem(data, k, config)?;
1357 validate_gaussian_mixture_checkpoint(data, config.covariance_floor, &checkpoint)?;
1358 run_gaussian_mixture_em(data, config, checkpoint)
1359}
1360
1361fn validate_gaussian_mixture_problem(
1362 data: ArrayView2<'_, f64>,
1363 k: usize,
1364 config: GaussianMixtureConfig,
1365) -> Result<(), GaussianMixtureError> {
1366 let n = data.nrows();
1367 let d = data.ncols();
1368 if k == 0 {
1369 return Err(GaussianMixtureError::InvalidInput {
1370 message: "k must be positive".to_string(),
1371 });
1372 }
1373 if d == 0 {
1374 return Err(GaussianMixtureError::InvalidInput {
1375 message: "at least one data column is required".to_string(),
1376 });
1377 }
1378 if k > n {
1379 return Err(GaussianMixtureError::InvalidInput {
1380 message: format!("requested {k} components but data has {n} rows"),
1381 });
1382 }
1383 if data.iter().any(|value| !value.is_finite()) {
1384 return Err(GaussianMixtureError::InvalidInput {
1385 message: "data must be finite".to_string(),
1386 });
1387 }
1388 if config.max_iter == 0 || config.kmeans_max_iter == 0 {
1389 return Err(GaussianMixtureError::InvalidInput {
1390 message: "max_iter and kmeans_max_iter must be positive".to_string(),
1391 });
1392 }
1393 let numerical_floor = f64::EPSILON.sqrt();
1394 if !config.loglik_tol.is_finite()
1395 || config.loglik_tol < numerical_floor
1396 || !config.parameter_tol.is_finite()
1397 || config.parameter_tol < numerical_floor
1398 || !config.covariance_floor.is_finite()
1399 || config.covariance_floor <= 0.0
1400 {
1401 return Err(GaussianMixtureError::InvalidInput {
1402 message: format!(
1403 "loglik_tol and parameter_tol must be finite and >= {numerical_floor:.3e}, and covariance_floor must be finite and positive"
1404 ),
1405 });
1406 }
1407 Ok(())
1408}
1409
1410fn validate_gaussian_mixture_checkpoint(
1411 data: ArrayView2<'_, f64>,
1412 covariance_floor: f64,
1413 checkpoint: &GaussianMixtureCheckpoint,
1414) -> Result<(), GaussianMixtureError> {
1415 let d = data.ncols();
1416 let k = checkpoint.weights.len();
1417 let mass = checkpoint.weights.sum();
1418 if k == 0
1419 || checkpoint.data_fingerprint != mixture_data_fingerprint(data)
1420 || checkpoint.covariance_floor.to_bits() != covariance_floor.to_bits()
1421 || checkpoint.means.dim() != (k, d)
1422 || checkpoint.covariances.len() != k
1423 || checkpoint
1424 .covariances
1425 .iter()
1426 .any(|covariance| covariance.dim() != (d, d))
1427 || checkpoint
1428 .weights
1429 .iter()
1430 .chain(checkpoint.means.iter())
1431 .chain(checkpoint.covariances.iter().flat_map(|value| value.iter()))
1432 .any(|value| !value.is_finite())
1433 || checkpoint.weights.iter().any(|value| *value <= 0.0)
1434 || !mass.is_finite()
1435 || (mass - 1.0).abs() > f64::EPSILON.sqrt()
1436 || !checkpoint.mean_log_likelihood.is_finite()
1437 {
1438 return Err(GaussianMixtureError::InvalidInput {
1439 message: "checkpoint problem identity, dimensions, interior parameters, likelihood, or simplex mass are invalid".to_string(),
1440 });
1441 }
1442 Ok(())
1443}
1444
1445fn run_gaussian_mixture_em(
1446 data: ArrayView2<'_, f64>,
1447 config: GaussianMixtureConfig,
1448 mut checkpoint: GaussianMixtureCheckpoint,
1449) -> Result<GaussianMixtureFit, GaussianMixtureError> {
1450 validate_gaussian_mixture_checkpoint(data, config.covariance_floor, &checkpoint)?;
1451 let k = checkpoint.weights.len();
1452 let d = data.ncols();
1453 let data_fingerprint = mixture_data_fingerprint(data);
1454
1455 let mut budget = config.max_iter;
1466 let mut extension: Option<usize> = None;
1467 let mut residual_window: std::collections::VecDeque<f64> =
1468 std::collections::VecDeque::with_capacity(EM_RATE_WINDOW + 1);
1469 let mut additional_updates = 0usize;
1470 loop {
1471 let current = mixture_e_step(
1472 data,
1473 &checkpoint.weights,
1474 &checkpoint.means,
1475 &checkpoint.covariances,
1476 )
1477 .map_err(|message| GaussianMixtureError::NumericalFailure {
1478 message,
1479 checkpoint: Some(checkpoint.clone()),
1480 })?;
1481 if (checkpoint.mean_log_likelihood - current.mean_log_likelihood).abs()
1482 > current.mean_log_likelihood_roundoff
1483 {
1484 return Err(GaussianMixtureError::InvalidInput {
1485 message: format!(
1486 "checkpoint mean log likelihood {:.12e} disagrees with its parameters ({:.12e} +/- {:.3e})",
1487 checkpoint.mean_log_likelihood,
1488 current.mean_log_likelihood,
1489 current.mean_log_likelihood_roundoff
1490 ),
1491 });
1492 }
1493 checkpoint.mean_log_likelihood = current.mean_log_likelihood;
1494
1495 let (next_weights, next_means, next_covariances) = mixture_m_step(
1496 data,
1497 current.responsibilities.view(),
1498 config.covariance_floor,
1499 )
1500 .map_err(|message| GaussianMixtureError::NumericalFailure {
1501 message,
1502 checkpoint: Some(checkpoint.clone()),
1503 })?;
1504 let next = mixture_e_step(data, &next_weights, &next_means, &next_covariances).map_err(
1505 |message| GaussianMixtureError::NumericalFailure {
1506 message,
1507 checkpoint: Some(checkpoint.clone()),
1508 },
1509 )?;
1510 let objective_scale = current
1511 .mean_log_likelihood
1512 .abs()
1513 .max(next.mean_log_likelihood.abs())
1514 .max(1.0);
1515 let objective_step = next.mean_log_likelihood - current.mean_log_likelihood;
1516 let objective_residual = objective_step.abs() / objective_scale;
1517 let parameter_residual = empirical_predictive_density_residual(
1518 ¤t.row_log_likelihoods,
1519 &next.row_log_likelihoods,
1520 )
1521 .map_err(|message| GaussianMixtureError::NumericalFailure {
1522 message,
1523 checkpoint: Some(checkpoint.clone()),
1524 })?;
1525 let monotonicity_uncertainty = gaussian_mixture_monotonicity_uncertainty(
1526 objective_scale,
1527 current.mean_log_likelihood_roundoff,
1528 next.mean_log_likelihood_roundoff,
1529 );
1530 residual_window.push_back(parameter_residual);
1531 if residual_window.len() > EM_RATE_WINDOW + 1 {
1532 residual_window.pop_front();
1533 }
1534 let contraction_rate = em_contraction_rate(&residual_window);
1535 let projected_iterations_to_tolerance = contraction_rate
1536 .and_then(|rate| em_projected_iterations(parameter_residual, config.parameter_tol, rate));
1537 let certificate = GaussianMixtureCertificate {
1538 mean_log_likelihood: current.mean_log_likelihood,
1539 mean_log_likelihood_gain: objective_step,
1540 monotonicity_uncertainty,
1541 objective_residual,
1542 objective_tolerance: config.loglik_tol,
1543 parameter_residual,
1544 parameter_tolerance: config.parameter_tol,
1545 contraction_rate,
1546 projected_iterations_to_tolerance,
1547 };
1548 if objective_step < -monotonicity_uncertainty {
1549 return Err(GaussianMixtureError::MonotonicityViolation {
1550 previous_mean_log_likelihood: current.mean_log_likelihood,
1551 next_mean_log_likelihood: next.mean_log_likelihood,
1552 numerical_uncertainty: monotonicity_uncertainty,
1553 checkpoint,
1554 });
1555 }
1556 if objective_residual <= config.loglik_tol && parameter_residual <= config.parameter_tol {
1557 let loglik = current.mean_log_likelihood * data.nrows() as f64;
1558 if !loglik.is_finite() {
1559 return Err(GaussianMixtureError::NumericalFailure {
1560 message: "certified mean log likelihood overflows as a total likelihood"
1561 .to_string(),
1562 checkpoint: Some(checkpoint),
1563 });
1564 }
1565 return Ok(GaussianMixtureFit {
1566 weights: checkpoint.weights,
1567 means: checkpoint.means,
1568 covariances: checkpoint.covariances,
1569 k,
1570 d,
1571 n_obs: data.nrows(),
1572 loglik,
1573 iterations: checkpoint.completed_iterations,
1574 certificate,
1575 });
1576 }
1577 if additional_updates >= budget {
1578 let extend = match (contraction_rate, projected_iterations_to_tolerance) {
1587 (Some(rate), Some(steps)) if rate < 1.0 && extension.is_none() => Some(steps),
1588 _ => None,
1589 };
1590 match extend {
1591 Some(steps) => {
1592 let steps = steps.min(config.max_iter);
1603 budget = budget.saturating_add(steps);
1604 extension = Some(steps);
1605 }
1606 None => {
1607 return Err(GaussianMixtureError::DidNotConverge {
1608 max_iterations: budget,
1609 certificate,
1610 checkpoint,
1611 });
1612 }
1613 }
1614 } else if extension.is_some() && additional_updates.is_multiple_of(EM_RATE_WINDOW) {
1615 if !matches!(contraction_rate, Some(rate) if rate < 1.0) {
1619 return Err(GaussianMixtureError::DidNotConverge {
1620 max_iterations: budget,
1621 certificate,
1622 checkpoint,
1623 });
1624 }
1625 }
1626 checkpoint = GaussianMixtureCheckpoint {
1627 weights: next_weights,
1628 means: next_means,
1629 covariances: next_covariances,
1630 mean_log_likelihood: next.mean_log_likelihood,
1631 completed_iterations: checkpoint.completed_iterations + 1,
1632 data_fingerprint,
1633 covariance_floor: config.covariance_floor,
1634 };
1635 additional_updates += 1;
1636 }
1637}
1638
1639struct GaussianMixtureEStep {
1640 responsibilities: Array2<f64>,
1641 row_log_likelihoods: Vec<f64>,
1642 mean_log_likelihood: f64,
1643 mean_log_likelihood_roundoff: f64,
1644}
1645
1646fn gaussian_mixture_monotonicity_uncertainty(
1658 objective_scale: f64,
1659 current_reduction_roundoff: f64,
1660 next_reduction_roundoff: f64,
1661) -> f64 {
1662 let reduction_roundoff = current_reduction_roundoff + next_reduction_roundoff;
1663 let composite_map_resolution = f64::EPSILON.sqrt() * objective_scale;
1664 reduction_roundoff.max(composite_map_resolution)
1665}
1666
1667fn pairwise_sum_max_depth(term_count: usize) -> usize {
1668 if term_count <= 1 {
1669 return 0;
1670 }
1671 let within_block = term_count.min(BASE_CHUNK) - 1;
1672 let blocks = term_count.div_ceil(BASE_CHUNK);
1673 let tree_levels = if blocks <= 1 {
1674 0
1675 } else {
1676 (usize::BITS - (blocks - 1).leading_zeros()) as usize
1677 };
1678 within_block.saturating_add(tree_levels)
1679}
1680
1681fn pairwise_mean_with_roundoff(values: &[f64]) -> Result<(f64, f64), String> {
1682 if values.is_empty() || values.iter().any(|value| !value.is_finite()) {
1683 return Err("mean log-likelihood terms must be nonempty and finite".to_string());
1684 }
1685 let sum = pairwise_sum(values);
1686 let magnitudes: Vec<f64> = values.iter().map(|value| value.abs()).collect();
1687 let magnitude_sum = pairwise_sum(&magnitudes);
1688 let unit_roundoff = 0.5 * f64::EPSILON;
1689 let accumulated = pairwise_sum_max_depth(values.len()) as f64 * unit_roundoff;
1690 let addition_bound = if accumulated < 1.0 {
1691 accumulated / (1.0 - accumulated) * magnitude_sum
1692 } else {
1693 f64::INFINITY
1694 };
1695 let count = values.len() as f64;
1696 let mean = sum / count;
1697 let roundoff = addition_bound / count + unit_roundoff * mean.abs();
1701 if !(mean.is_finite() && roundoff.is_finite()) {
1702 return Err("mean mixture log likelihood or its rounding bound is non-finite".to_string());
1703 }
1704 Ok((mean, roundoff))
1705}
1706
1707fn mixture_e_step(
1708 data: ArrayView2<'_, f64>,
1709 weights: &Array1<f64>,
1710 means: &Array2<f64>,
1711 covariances: &[Array2<f64>],
1712) -> Result<GaussianMixtureEStep, String> {
1713 let n = data.nrows();
1714 let k = weights.len();
1715 if weights
1716 .iter()
1717 .any(|weight| !weight.is_finite() || *weight <= 0.0)
1718 {
1719 return Err("mixture E-step requires strictly positive finite weights".to_string());
1720 }
1721 let mut components = Vec::with_capacity(k);
1722 for component in 0..k {
1723 components.push(GaussianComponentEval::factor(
1724 means.row(component),
1725 &covariances[component],
1726 )?);
1727 }
1728 let log_weights: Vec<f64> = weights.iter().map(|weight| weight.ln()).collect();
1729 let mut responsibilities = Array2::<f64>::zeros((n, k));
1730 let mut row_log_likelihoods = Vec::with_capacity(n);
1731 for row in 0..n {
1732 let observation = data.row(row);
1733 let mut log_terms = vec![f64::NEG_INFINITY; k];
1734 let mut max_term = f64::NEG_INFINITY;
1735 for component in 0..k {
1736 let term = log_weights[component] + components[component].log_density(observation);
1737 log_terms[component] = term;
1738 max_term = max_term.max(term);
1739 }
1740 let log_mixture = log_sum_exp(&log_terms, max_term);
1741 if !log_mixture.is_finite() {
1742 return Err(format!(
1743 "mixture density is non-finite at training row {row}"
1744 ));
1745 }
1746 row_log_likelihoods.push(log_mixture);
1747 for component in 0..k {
1748 responsibilities[[row, component]] = (log_terms[component] - log_mixture).exp();
1749 }
1750 }
1751 let (mean_log_likelihood, mean_log_likelihood_roundoff) =
1752 pairwise_mean_with_roundoff(&row_log_likelihoods)?;
1753 Ok(GaussianMixtureEStep {
1754 responsibilities,
1755 row_log_likelihoods,
1756 mean_log_likelihood,
1757 mean_log_likelihood_roundoff,
1758 })
1759}
1760
1761fn mixture_m_step(
1762 data: ArrayView2<'_, f64>,
1763 responsibilities: ArrayView2<'_, f64>,
1764 covariance_floor: f64,
1765) -> Result<(Array1<f64>, Array2<f64>, Vec<Array2<f64>>), String> {
1766 let n = data.nrows();
1767 let d = data.ncols();
1768 let k = responsibilities.ncols();
1769 let mut component_mass = Array1::<f64>::zeros(k);
1770 for component in 0..k {
1771 component_mass[component] = responsibilities.column(component).sum();
1772 }
1773 if component_mass
1774 .iter()
1775 .any(|mass| !mass.is_finite() || *mass <= 0.0)
1776 {
1777 return Err(
1778 "M-step reached a zero-mass component; the requested mixture order has no interior fitted density"
1779 .to_string(),
1780 );
1781 }
1782 let mut weights = component_mass.mapv(|mass| mass / n as f64);
1783 let total_weight = weights.sum();
1784 if !(total_weight.is_finite() && total_weight > 0.0) {
1785 return Err("M-step produced invalid mixture-weight mass".to_string());
1786 }
1787 weights.mapv_inplace(|weight| weight / total_weight);
1788 let mut means = Array2::<f64>::zeros((k, d));
1789 let mut covariances = Vec::with_capacity(k);
1790 for component in 0..k {
1791 let mass = component_mass[component];
1792 let mut mean = Array1::<f64>::zeros(d);
1793 for row in 0..n {
1794 let responsibility = responsibilities[[row, component]];
1795 for col in 0..d {
1796 mean[col] += responsibility * data[[row, col]];
1797 }
1798 }
1799 mean.mapv_inplace(|value| value / mass);
1800 means.row_mut(component).assign(&mean);
1801 let mut covariance = Array2::<f64>::zeros((d, d));
1802 for row in 0..n {
1803 let responsibility = responsibilities[[row, component]];
1804 for left in 0..d {
1805 let left_residual = data[[row, left]] - mean[left];
1806 for right in 0..d {
1807 covariance[[left, right]] +=
1808 responsibility * left_residual * (data[[row, right]] - mean[right]);
1809 }
1810 }
1811 }
1812 covariance.mapv_inplace(|value| value / mass);
1813 covariances.push(constrain_covariance(covariance, covariance_floor)?);
1814 }
1815 Ok((weights, means, covariances))
1816}
1817
1818fn relative_parameter_step(previous: f64, next: f64) -> f64 {
1819 (next - previous).abs() / previous.abs().max(next.abs()).max(1.0)
1820}
1821
1822fn empirical_predictive_density_residual(
1835 previous_row_log_density: &[f64],
1836 next_row_log_density: &[f64],
1837) -> Result<f64, String> {
1838 if previous_row_log_density.is_empty()
1839 || previous_row_log_density.len() != next_row_log_density.len()
1840 || previous_row_log_density
1841 .iter()
1842 .chain(next_row_log_density)
1843 .any(|value| !value.is_finite())
1844 {
1845 return Err(
1846 "predictive-density residual requires equal, nonempty, finite log-density vectors"
1847 .to_string(),
1848 );
1849 }
1850 Ok(previous_row_log_density
1851 .iter()
1852 .zip(next_row_log_density)
1853 .map(|(&previous, &next)| (next - previous).abs())
1854 .fold(0.0_f64, f64::max))
1855}
1856
1857fn constrain_covariance(covariance: Array2<f64>, floor: f64) -> Result<Array2<f64>, String> {
1858 let (eigenvalues, eigenvectors) = covariance
1859 .eigh(Side::Lower)
1860 .map_err(|error| format!("covariance eigendecomposition failed: {error}"))?;
1861 let d = covariance.nrows();
1862 let mut constrained = Array2::<f64>::zeros((d, d));
1863 for row in 0..d {
1864 for col in 0..d {
1865 let mut value = 0.0_f64;
1866 for index in 0..d {
1867 value += eigenvectors[[row, index]]
1868 * eigenvalues[index].max(floor)
1869 * eigenvectors[[col, index]];
1870 }
1871 constrained[[row, col]] = value;
1872 }
1873 }
1874 if constrained.iter().any(|value| !value.is_finite()) {
1875 return Err("constrained covariance became non-finite".to_string());
1876 }
1877 Ok(constrained)
1878}
1879
1880fn constrained_data_covariance(
1882 data: ArrayView2<'_, f64>,
1883 floor: f64,
1884) -> Result<Array2<f64>, String> {
1885 let n = data.nrows();
1886 let d = data.ncols();
1887 let mut mean = Array1::<f64>::zeros(d);
1888 for i in 0..n {
1889 for c in 0..d {
1890 mean[c] += data[[i, c]];
1891 }
1892 }
1893 mean.mapv_inplace(|v| v / n.max(1) as f64);
1894 let mut cov = Array2::<f64>::zeros((d, d));
1895 for i in 0..n {
1896 for a in 0..d {
1897 let da = data[[i, a]] - mean[a];
1898 for b in 0..d {
1899 cov[[a, b]] += da * (data[[i, b]] - mean[b]);
1900 }
1901 }
1902 }
1903 let inv = 1.0 / n as f64;
1904 cov.mapv_inplace(|v| v * inv);
1905 constrain_covariance(cov, floor)
1906}
1907
1908#[derive(Debug, Clone)]
1926pub struct RingGaussianMixtureFit {
1927 weights: Array1<f64>,
1928 center: Array1<f64>,
1929 radius: f64,
1930 directions: Array2<f64>,
1931 variance: f64,
1932 k: usize,
1933 n_obs: usize,
1934 loglik: f64,
1935 iterations: usize,
1936 certificate: GaussianMixtureCertificate,
1937}
1938
1939impl RingGaussianMixtureFit {
1940 pub fn weights(&self) -> ArrayView1<'_, f64> {
1941 self.weights.view()
1942 }
1943
1944 pub fn center(&self) -> ArrayView1<'_, f64> {
1945 self.center.view()
1946 }
1947
1948 pub fn radius(&self) -> f64 {
1949 self.radius
1950 }
1951
1952 pub fn directions(&self) -> ArrayView2<'_, f64> {
1953 self.directions.view()
1954 }
1955
1956 pub fn variance(&self) -> f64 {
1957 self.variance
1958 }
1959
1960 pub fn iterations(&self) -> usize {
1961 self.iterations
1962 }
1963
1964 pub fn certificate(&self) -> GaussianMixtureCertificate {
1965 self.certificate
1966 }
1967
1968 pub fn num_free_parameters(&self) -> usize {
1971 2 * self.k + 3
1972 }
1973
1974 pub fn per_point_log_density(&self, data: ArrayView2<'_, f64>) -> Result<Array1<f64>, String> {
1975 if data.ncols() != 2 {
1976 return Err(format!(
1977 "ring-of-clusters density expects two columns, got {}",
1978 data.ncols()
1979 ));
1980 }
1981 ring_mixture_log_density(
1982 data,
1983 &self.weights,
1984 &self.center,
1985 self.radius,
1986 &self.directions,
1987 self.variance,
1988 )
1989 }
1990
1991 pub fn bic(&self) -> f64 {
1994 -self.loglik + 0.5 * self.num_free_parameters() as f64 * (self.n_obs as f64).ln()
1995 }
1996}
1997
1998#[derive(Debug, Clone)]
1999struct RingMixtureState {
2000 weights: Array1<f64>,
2001 center: Array1<f64>,
2002 radius: f64,
2003 directions: Array2<f64>,
2004 variance: f64,
2005 mean_log_likelihood: f64,
2006 completed_iterations: usize,
2007}
2008
2009fn ring_component_means(
2010 center: &Array1<f64>,
2011 radius: f64,
2012 directions: &Array2<f64>,
2013) -> Array2<f64> {
2014 let mut means = Array2::<f64>::zeros((directions.nrows(), 2));
2015 for component in 0..directions.nrows() {
2016 means[[component, 0]] = center[0] + radius * directions[[component, 0]];
2017 means[[component, 1]] = center[1] + radius * directions[[component, 1]];
2018 }
2019 means
2020}
2021
2022fn ring_mixture_log_terms(
2023 data: ArrayView2<'_, f64>,
2024 weights: &Array1<f64>,
2025 center: &Array1<f64>,
2026 radius: f64,
2027 directions: &Array2<f64>,
2028 variance: f64,
2029) -> Result<(Array2<f64>, Vec<f64>), String> {
2030 if data.ncols() != 2
2031 || center.len() != 2
2032 || directions.ncols() != 2
2033 || directions.nrows() != weights.len()
2034 || weights
2035 .iter()
2036 .any(|weight| !weight.is_finite() || *weight <= 0.0)
2037 || !(radius.is_finite() && radius > 0.0)
2038 || !(variance.is_finite() && variance > 0.0)
2039 {
2040 return Err("invalid ring-of-clusters parameter state".to_string());
2041 }
2042 let means = ring_component_means(center, radius, directions);
2043 let log_normalizer = -(std::f64::consts::TAU).ln() - variance.ln();
2044 let mut terms = Array2::<f64>::zeros((data.nrows(), weights.len()));
2045 let mut row_log_likelihoods = Vec::with_capacity(data.nrows());
2046 for row in 0..data.nrows() {
2047 let mut max_term = f64::NEG_INFINITY;
2048 for component in 0..weights.len() {
2049 let dx = data[[row, 0]] - means[[component, 0]];
2050 let dy = data[[row, 1]] - means[[component, 1]];
2051 let term =
2052 weights[component].ln() + log_normalizer - 0.5 * (dx * dx + dy * dy) / variance;
2053 terms[[row, component]] = term;
2054 max_term = max_term.max(term);
2055 }
2056 let values = terms.row(row).to_vec();
2057 let log_likelihood = log_sum_exp(&values, max_term);
2058 if !log_likelihood.is_finite() {
2059 return Err(format!(
2060 "ring-of-clusters density is non-finite at training row {row}"
2061 ));
2062 }
2063 row_log_likelihoods.push(log_likelihood);
2064 }
2065 Ok((terms, row_log_likelihoods))
2066}
2067
2068fn ring_mixture_e_step(
2069 data: ArrayView2<'_, f64>,
2070 state: &RingMixtureState,
2071) -> Result<GaussianMixtureEStep, String> {
2072 let (terms, row_log_likelihoods) = ring_mixture_log_terms(
2073 data,
2074 &state.weights,
2075 &state.center,
2076 state.radius,
2077 &state.directions,
2078 state.variance,
2079 )?;
2080 let mut responsibilities = Array2::<f64>::zeros(terms.raw_dim());
2081 for row in 0..terms.nrows() {
2082 for component in 0..terms.ncols() {
2083 responsibilities[[row, component]] =
2084 (terms[[row, component]] - row_log_likelihoods[row]).exp();
2085 }
2086 }
2087 let (mean_log_likelihood, mean_log_likelihood_roundoff) =
2088 pairwise_mean_with_roundoff(&row_log_likelihoods)?;
2089 Ok(GaussianMixtureEStep {
2090 responsibilities,
2091 row_log_likelihoods,
2092 mean_log_likelihood,
2093 mean_log_likelihood_roundoff,
2094 })
2095}
2096
2097fn ring_mixture_log_density(
2098 data: ArrayView2<'_, f64>,
2099 weights: &Array1<f64>,
2100 center: &Array1<f64>,
2101 radius: f64,
2102 directions: &Array2<f64>,
2103 variance: f64,
2104) -> Result<Array1<f64>, String> {
2105 let (_, row_log_likelihoods) =
2106 ring_mixture_log_terms(data, weights, center, radius, directions, variance)?;
2107 Ok(Array1::from_vec(row_log_likelihoods))
2108}
2109
2110fn fit_weighted_component_circle(
2111 component_means: &Array2<f64>,
2112 component_mass: &Array1<f64>,
2113 initial_center: &Array1<f64>,
2114 initial_radius: f64,
2115 parameter_tol: f64,
2116 max_iter: usize,
2117) -> Result<(Array1<f64>, f64, Array2<f64>), String> {
2118 let k = component_means.nrows();
2119 let total_mass = component_mass.sum();
2120 if component_means.ncols() != 2
2121 || component_mass.len() != k
2122 || component_mass
2123 .iter()
2124 .any(|mass| !mass.is_finite() || *mass <= 0.0)
2125 || !(total_mass.is_finite() && total_mass > 0.0)
2126 {
2127 return Err("ring M-step requires positive component masses and 2-D means".to_string());
2128 }
2129 let mut center = initial_center.clone();
2130 let mut radius = initial_radius;
2131 let mut directions = Array2::<f64>::zeros((k, 2));
2132 for _ in 0..max_iter {
2133 for component in 0..k {
2134 let dx = component_means[[component, 0]] - center[0];
2135 let dy = component_means[[component, 1]] - center[1];
2136 let norm = dx.hypot(dy);
2137 if !(norm.is_finite() && norm > 0.0) {
2138 return Err(
2139 "ring M-step reached a component centroid at the circle center; its angle is unidentified"
2140 .to_string(),
2141 );
2142 }
2143 directions[[component, 0]] = dx / norm;
2144 directions[[component, 1]] = dy / norm;
2145 }
2146
2147 let mut mean_point = Array1::<f64>::zeros(2);
2148 let mut mean_direction = Array1::<f64>::zeros(2);
2149 for component in 0..k {
2150 let weight = component_mass[component] / total_mass;
2151 for axis in 0..2 {
2152 mean_point[axis] += weight * component_means[[component, axis]];
2153 mean_direction[axis] += weight * directions[[component, axis]];
2154 }
2155 }
2156 let mut numerator = 0.0;
2157 let mut denominator = 0.0;
2158 for component in 0..k {
2159 let mass = component_mass[component];
2160 let dux = directions[[component, 0]] - mean_direction[0];
2161 let duy = directions[[component, 1]] - mean_direction[1];
2162 numerator += mass
2163 * (dux * (component_means[[component, 0]] - mean_point[0])
2164 + duy * (component_means[[component, 1]] - mean_point[1]));
2165 denominator += mass * (dux * dux + duy * duy);
2166 }
2167 if !(denominator.is_finite() && denominator > 0.0) {
2168 return Err(
2169 "ring M-step component directions are identical; radius and center are unidentified"
2170 .to_string(),
2171 );
2172 }
2173 let mut next_radius = numerator / denominator;
2174 if !next_radius.is_finite() || next_radius == 0.0 {
2175 return Err("ring M-step produced an unidentified zero radius".to_string());
2176 }
2177 if next_radius < 0.0 {
2178 next_radius = -next_radius;
2179 directions.mapv_inplace(|value| -value);
2180 }
2181 let next_center = Array1::from_vec(vec![
2182 mean_point[0] - next_radius * mean_direction[0],
2183 mean_point[1] - next_radius * mean_direction[1],
2184 ]);
2185 let residual = center
2186 .iter()
2187 .zip(next_center.iter())
2188 .map(|(&left, &right)| relative_parameter_step(left, right))
2189 .chain(std::iter::once(relative_parameter_step(
2190 radius,
2191 next_radius,
2192 )))
2193 .fold(0.0, f64::max);
2194 center = next_center;
2195 radius = next_radius;
2196 if residual <= parameter_tol {
2197 for component in 0..k {
2200 let dx = component_means[[component, 0]] - center[0];
2201 let dy = component_means[[component, 1]] - center[1];
2202 let norm = dx.hypot(dy);
2203 if !(norm.is_finite() && norm > 0.0) {
2204 return Err("ring M-step terminal component angle is unidentified".to_string());
2205 }
2206 directions[[component, 0]] = dx / norm;
2207 directions[[component, 1]] = dy / norm;
2208 }
2209 return Ok((center, radius, directions));
2210 }
2211 }
2212 Err(format!(
2213 "ring M-step did not certify its constrained center/radius fixed point after {max_iter} iterations"
2214 ))
2215}
2216
2217fn ring_mixture_m_step(
2218 data: ArrayView2<'_, f64>,
2219 responsibilities: ArrayView2<'_, f64>,
2220 previous: &RingMixtureState,
2221 config: GaussianMixtureConfig,
2222) -> Result<RingMixtureState, String> {
2223 let n = data.nrows();
2224 let k = responsibilities.ncols();
2225 let mut component_mass = Array1::<f64>::zeros(k);
2226 let mut component_means = Array2::<f64>::zeros((k, 2));
2227 for component in 0..k {
2228 let mass = responsibilities.column(component).sum();
2229 if !(mass.is_finite() && mass > 0.0) {
2230 return Err(
2231 "ring M-step reached a zero-mass component; the requested order is singular"
2232 .to_string(),
2233 );
2234 }
2235 component_mass[component] = mass;
2236 for row in 0..n {
2237 for axis in 0..2 {
2238 component_means[[component, axis]] +=
2239 responsibilities[[row, component]] * data[[row, axis]];
2240 }
2241 }
2242 for axis in 0..2 {
2243 component_means[[component, axis]] /= mass;
2244 }
2245 }
2246 let mut weights = component_mass.mapv(|mass| mass / n as f64);
2247 let weight_sum = weights.sum();
2248 weights.mapv_inplace(|weight| weight / weight_sum);
2249 let (center, radius, directions) = fit_weighted_component_circle(
2250 &component_means,
2251 &component_mass,
2252 &previous.center,
2253 previous.radius,
2254 config.parameter_tol,
2255 config.max_iter,
2256 )?;
2257 let means = ring_component_means(¢er, radius, &directions);
2258 let mut expected_squared_error = 0.0;
2259 for row in 0..n {
2260 for component in 0..k {
2261 let dx = data[[row, 0]] - means[[component, 0]];
2262 let dy = data[[row, 1]] - means[[component, 1]];
2263 expected_squared_error += responsibilities[[row, component]] * (dx * dx + dy * dy);
2264 }
2265 }
2266 let variance = (expected_squared_error / (2 * n) as f64).max(config.covariance_floor);
2267 if !variance.is_finite() {
2268 return Err("ring M-step produced non-finite shared variance".to_string());
2269 }
2270 Ok(RingMixtureState {
2271 weights,
2272 center,
2273 radius,
2274 directions,
2275 variance,
2276 mean_log_likelihood: f64::NAN,
2277 completed_iterations: previous.completed_iterations + 1,
2278 })
2279}
2280
2281pub fn fit_ring_gaussian_mixture(
2284 data: ArrayView2<'_, f64>,
2285 k: usize,
2286 config: GaussianMixtureConfig,
2287) -> Result<RingGaussianMixtureFit, String> {
2288 validate_gaussian_mixture_problem(data, k, config).map_err(|error| error.to_string())?;
2289 if data.ncols() != 2 {
2290 return Err(format!(
2291 "ring-of-clusters fitting requires exactly two columns, got {}",
2292 data.ncols()
2293 ));
2294 }
2295 if k < 3 {
2296 return Err(format!(
2297 "ring-of-clusters fitting requires at least three component centers, got {k}"
2298 ));
2299 }
2300 let seeded_means = gam_terms::basis::select_centers_by_strategy(
2301 data,
2302 &gam_terms::basis::CenterStrategy::KMeans {
2303 num_centers: k,
2304 max_iter: config.kmeans_max_iter,
2305 },
2306 )
2307 .map_err(|error| format!("ring-of-clusters deterministic seeding failed: {error}"))?;
2308 let component_mass = Array1::<f64>::ones(k);
2309 let mut initial_center = Array1::<f64>::zeros(2);
2310 for component in 0..k {
2311 initial_center[0] += seeded_means[[component, 0]] / k as f64;
2312 initial_center[1] += seeded_means[[component, 1]] / k as f64;
2313 }
2314 let mut initial_radius = 0.0;
2315 for component in 0..k {
2316 initial_radius += (seeded_means[[component, 0]] - initial_center[0])
2317 .hypot(seeded_means[[component, 1]] - initial_center[1])
2318 / k as f64;
2319 }
2320 if !(initial_radius.is_finite() && initial_radius > 0.0) {
2321 return Err("ring-of-clusters seed has an unidentified zero radius".to_string());
2322 }
2323 let (center, radius, directions) = fit_weighted_component_circle(
2324 &seeded_means,
2325 &component_mass,
2326 &initial_center,
2327 initial_radius,
2328 config.parameter_tol,
2329 config.max_iter,
2330 )?;
2331 let means = ring_component_means(¢er, radius, &directions);
2332 let mut squared_error = 0.0;
2333 for row in 0..data.nrows() {
2334 let mut nearest = f64::INFINITY;
2335 for component in 0..k {
2336 let dx = data[[row, 0]] - means[[component, 0]];
2337 let dy = data[[row, 1]] - means[[component, 1]];
2338 nearest = nearest.min(dx * dx + dy * dy);
2339 }
2340 squared_error += nearest;
2341 }
2342 let variance = (squared_error / (2 * data.nrows()) as f64).max(config.covariance_floor);
2343 let mut state = RingMixtureState {
2344 weights: Array1::from_elem(k, 1.0 / k as f64),
2345 center,
2346 radius,
2347 directions,
2348 variance,
2349 mean_log_likelihood: f64::NAN,
2350 completed_iterations: 0,
2351 };
2352 for additional_updates in 0..=config.max_iter {
2353 let current = ring_mixture_e_step(data, &state)?;
2354 state.mean_log_likelihood = current.mean_log_likelihood;
2355 let mut next =
2356 ring_mixture_m_step(data, current.responsibilities.view(), &state, config)?;
2357 let next_e_step = ring_mixture_e_step(data, &next)?;
2358 next.mean_log_likelihood = next_e_step.mean_log_likelihood;
2359 let current_mean = current.mean_log_likelihood;
2360 let next_mean = next_e_step.mean_log_likelihood;
2361 let objective_scale = current_mean.abs().max(next_mean.abs()).max(1.0);
2362 let objective_step = next_mean - current_mean;
2363 let objective_residual = objective_step.abs() / objective_scale;
2364 let parameter_residual = empirical_predictive_density_residual(
2365 ¤t.row_log_likelihoods,
2366 &next_e_step.row_log_likelihoods,
2367 )?;
2368 let monotonicity_uncertainty = gaussian_mixture_monotonicity_uncertainty(
2369 objective_scale,
2370 current.mean_log_likelihood_roundoff,
2371 next_e_step.mean_log_likelihood_roundoff,
2372 );
2373 let certificate = GaussianMixtureCertificate {
2374 mean_log_likelihood: current_mean,
2375 mean_log_likelihood_gain: objective_step,
2376 monotonicity_uncertainty,
2377 objective_residual,
2378 objective_tolerance: config.loglik_tol,
2379 parameter_residual,
2380 parameter_tolerance: config.parameter_tol,
2381 contraction_rate: None,
2386 projected_iterations_to_tolerance: None,
2387 };
2388 if objective_step < -monotonicity_uncertainty {
2389 return Err(format!(
2390 "ring-of-clusters generalized EM violated monotone ascent at iteration {}: {current_mean:.12e} -> {next_mean:.12e} (comparison uncertainty {monotonicity_uncertainty:.3e})",
2391 state.completed_iterations
2392 ));
2393 }
2394 if objective_residual <= config.loglik_tol && parameter_residual <= config.parameter_tol {
2395 let loglik = current_mean * data.nrows() as f64;
2396 if !loglik.is_finite() {
2397 return Err("ring-of-clusters total log likelihood overflowed".to_string());
2398 }
2399 return Ok(RingGaussianMixtureFit {
2400 weights: state.weights,
2401 center: state.center,
2402 radius: state.radius,
2403 directions: state.directions,
2404 variance: state.variance,
2405 k,
2406 n_obs: data.nrows(),
2407 loglik,
2408 iterations: state.completed_iterations,
2409 certificate,
2410 });
2411 }
2412 if additional_updates == config.max_iter {
2413 return Err(format!(
2414 "ring-of-clusters generalized EM did not certify after {} iterations: objective residual {:.6e}/{:.3e}, parameter-map residual {:.6e}/{:.3e}",
2415 config.max_iter,
2416 objective_residual,
2417 config.loglik_tol,
2418 parameter_residual,
2419 config.parameter_tol,
2420 ));
2421 }
2422 state = next;
2423 }
2424 Err("ring-of-clusters generalized EM exhausted without a terminal certificate".to_string())
2425}
2426
2427#[derive(Debug, Clone, Copy)]
2448pub struct CircularGaussianFit2d {
2449 center: [f64; 2],
2450 radius: f64,
2451 noise_variance: f64,
2452}
2453
2454impl CircularGaussianFit2d {
2455 pub const NUM_FREE_PARAMETERS: usize = 4;
2457
2458 pub fn from_parameters(
2460 center: [f64; 2],
2461 radius: f64,
2462 noise_variance: f64,
2463 ) -> Result<Self, String> {
2464 if !center.iter().all(|value| value.is_finite()) {
2465 return Err("circular Gaussian center must be finite".to_string());
2466 }
2467 if !(radius.is_finite() && radius >= 0.0) {
2468 return Err("circular Gaussian radius must be finite and nonnegative".to_string());
2469 }
2470 if !(noise_variance.is_finite() && noise_variance > 0.0) {
2471 return Err("circular Gaussian noise variance must be finite and positive".to_string());
2472 }
2473 Ok(Self {
2474 center,
2475 radius,
2476 noise_variance,
2477 })
2478 }
2479
2480 pub fn fit(coords: ArrayView2<'_, f64>, rows: &[usize]) -> Result<Self, String> {
2482 if coords.ncols() != 2 {
2483 return Err(format!(
2484 "circular Gaussian requires 2-D data, got {} columns",
2485 coords.ncols()
2486 ));
2487 }
2488 if rows.is_empty() {
2489 return Err("circular Gaussian requires a nonempty training set".to_string());
2490 }
2491 if rows.iter().any(|&row| row >= coords.nrows()) {
2492 return Err("circular Gaussian row index is out of bounds".to_string());
2493 }
2494 if rows
2495 .iter()
2496 .any(|&row| !coords[[row, 0]].is_finite() || !coords[[row, 1]].is_finite())
2497 {
2498 return Err("circular Gaussian requires finite training coordinates".to_string());
2499 }
2500
2501 let anchor_row = rows[0];
2505 let anchor = [coords[[anchor_row, 0]], coords[[anchor_row, 1]]];
2506 let mut scale = 0.0_f64;
2507 for &row in rows {
2508 let dx = coords[[row, 0]] - anchor[0];
2509 let dy = coords[[row, 1]] - anchor[1];
2510 if !(dx.is_finite() && dy.is_finite()) {
2511 return Err("circular Gaussian coordinate range exceeds f64".to_string());
2512 }
2513 scale = scale.max(dx.hypot(dy));
2514 }
2515 if !(scale.is_finite() && scale > 0.0) {
2516 return Err("circular Gaussian requires nonzero spatial extent".to_string());
2517 }
2518
2519 let mut points = Vec::with_capacity(rows.len());
2520 let mut mean = [0.0_f64; 2];
2521 for &row in rows {
2522 let point = [
2523 (coords[[row, 0]] - anchor[0]) / scale,
2524 (coords[[row, 1]] - anchor[1]) / scale,
2525 ];
2526 points.push(point);
2527 mean[0] += point[0];
2528 mean[1] += point[1];
2529 }
2530 let count = rows.len() as f64;
2531 mean[0] /= count;
2532 mean[1] /= count;
2533
2534 let mut squared_radii = Vec::with_capacity(rows.len());
2539 let mut mean_squared_radius = 0.0_f64;
2540 for point in &points {
2541 let dx = point[0] - mean[0];
2542 let dy = point[1] - mean[1];
2543 let squared_radius = dx * dx + dy * dy;
2544 squared_radii.push(squared_radius);
2545 mean_squared_radius += squared_radius;
2546 }
2547 mean_squared_radius /= count;
2548 let mut squared_radius_variance = 0.0_f64;
2549 for squared_radius in squared_radii {
2550 squared_radius_variance += (squared_radius - mean_squared_radius).powi(2);
2551 }
2552 squared_radius_variance /= count;
2553
2554 let variance_floor = (64.0 * f64::EPSILON * mean_squared_radius).max(f64::MIN_POSITIVE);
2558 let radius_squared = (mean_squared_radius * mean_squared_radius - squared_radius_variance)
2559 .max(0.0)
2560 .sqrt();
2561 let mut radius = radius_squared.sqrt();
2562 let mut noise_variance = (0.5 * (mean_squared_radius - radius_squared)).max(variance_floor);
2563 let mut center = mean;
2564
2565 const MAX_EM_ITERATIONS: usize = 4096;
2570 const EM_TOLERANCE: f64 = 2.0e-12;
2571 let mut posterior_means = vec![[0.0_f64; 2]; points.len()];
2572 let mut converged = false;
2573 for _ in 0..MAX_EM_ITERATIONS {
2574 let mut posterior_mean = [0.0_f64; 2];
2575 for (point, latent_mean) in points.iter().zip(&mut posterior_means) {
2576 let dx = point[0] - center[0];
2577 let dy = point[1] - center[1];
2578 let observed_radius = dx.hypot(dy);
2579 if observed_radius == 0.0 || radius == 0.0 {
2580 *latent_mean = [0.0, 0.0];
2581 } else {
2582 let (_, bessel_ratio) =
2583 circular_gaussian_bessel_terms(radius, observed_radius, noise_variance);
2584 if !(bessel_ratio.is_finite() && (0.0..=1.0).contains(&bessel_ratio)) {
2585 return Err("circular Gaussian Bessel ratio left [0, 1]".to_string());
2586 }
2587 let multiplier = bessel_ratio / observed_radius;
2588 *latent_mean = [multiplier * dx, multiplier * dy];
2589 }
2590 posterior_mean[0] += latent_mean[0];
2591 posterior_mean[1] += latent_mean[1];
2592 }
2593 posterior_mean[0] /= count;
2594 posterior_mean[1] /= count;
2595
2596 let denominator =
2597 1.0 - posterior_mean[0] * posterior_mean[0] - posterior_mean[1] * posterior_mean[1];
2598 if !(denominator.is_finite() && denominator > 0.0) {
2599 return Err("circular Gaussian EM radius update is singular".to_string());
2600 }
2601 let mut radius_numerator = 0.0_f64;
2602 for (point, latent_mean) in points.iter().zip(&posterior_means) {
2603 radius_numerator +=
2604 latent_mean[0] * (point[0] - mean[0]) + latent_mean[1] * (point[1] - mean[1]);
2605 }
2606 let next_radius = (radius_numerator / (count * denominator)).max(0.0);
2607 let next_center = [
2608 mean[0] - next_radius * posterior_mean[0],
2609 mean[1] - next_radius * posterior_mean[1],
2610 ];
2611
2612 let mut residual_sum = 0.0_f64;
2615 for (point, latent_mean) in points.iter().zip(&posterior_means) {
2616 let dx = point[0] - next_center[0];
2617 let dy = point[1] - next_center[1];
2618 let ex = dx - next_radius * latent_mean[0];
2619 let ey = dy - next_radius * latent_mean[1];
2620 let latent_norm_squared =
2621 latent_mean[0] * latent_mean[0] + latent_mean[1] * latent_mean[1];
2622 residual_sum += ex * ex
2623 + ey * ey
2624 + next_radius * next_radius * (1.0 - latent_norm_squared).max(0.0);
2625 }
2626 let next_noise_variance = (residual_sum / (2.0 * count)).max(variance_floor);
2627
2628 let parameter_change = (next_center[0] - center[0])
2629 .hypot(next_center[1] - center[1])
2630 .max((next_radius - radius).abs())
2631 .max(
2632 (next_noise_variance - noise_variance).abs()
2633 / (next_noise_variance + noise_variance),
2634 );
2635 center = next_center;
2636 radius = next_radius;
2637 noise_variance = next_noise_variance;
2638 if parameter_change <= EM_TOLERANCE {
2639 converged = true;
2640 break;
2641 }
2642 }
2643 if !converged {
2644 return Err("circular Gaussian maximum-likelihood fit did not converge".to_string());
2645 }
2646
2647 let fitted_noise_sd = scale * noise_variance.sqrt();
2648 Self::from_parameters(
2649 [anchor[0] + scale * center[0], anchor[1] + scale * center[1]],
2650 scale * radius,
2651 fitted_noise_sd * fitted_noise_sd,
2652 )
2653 .map_err(|error| format!("circular Gaussian fit produced invalid parameters: {error}"))
2654 }
2655
2656 pub const fn center(self) -> [f64; 2] {
2658 self.center
2659 }
2660
2661 pub const fn radius(self) -> f64 {
2663 self.radius
2664 }
2665
2666 pub const fn noise_variance(self) -> f64 {
2668 self.noise_variance
2669 }
2670
2671 pub fn log_density(self, x: f64, y: f64) -> f64 {
2673 let observed_radius = (x - self.center[0]).hypot(y - self.center[1]);
2674 let (log_i0_minus_kappa, _) =
2675 circular_gaussian_bessel_terms(self.radius, observed_radius, self.noise_variance);
2676 let standardized_radial_residual =
2677 (observed_radius - self.radius) / self.noise_variance.sqrt();
2678 -std::f64::consts::TAU.ln()
2681 - self.noise_variance.ln()
2682 - 0.5 * standardized_radial_residual.powi(2)
2683 + log_i0_minus_kappa
2684 }
2685
2686 pub fn log_likelihood(
2688 self,
2689 coords: ArrayView2<'_, f64>,
2690 rows: &[usize],
2691 ) -> Result<f64, String> {
2692 if coords.ncols() != 2 || rows.iter().any(|&row| row >= coords.nrows()) {
2693 return Err(
2694 "circular Gaussian likelihood received invalid coordinates or rows".to_string(),
2695 );
2696 }
2697 let mut log_densities = Vec::with_capacity(rows.len());
2698 for &row in rows {
2699 let value = self.log_density(coords[[row, 0]], coords[[row, 1]]);
2700 if !value.is_finite() {
2701 return Err("circular Gaussian likelihood is not finite".to_string());
2702 }
2703 log_densities.push(value);
2704 }
2705 let log_likelihood = pairwise_sum(&log_densities);
2706 if !log_likelihood.is_finite() {
2707 return Err("circular Gaussian likelihood sum is not finite".to_string());
2708 }
2709 Ok(log_likelihood)
2710 }
2711
2712 pub fn fit_with_bic(
2716 coords: ArrayView2<'_, f64>,
2717 rows: &[usize],
2718 ) -> Result<(Self, f64), String> {
2719 let fit = Self::fit(coords, rows)?;
2720 let log_likelihood = fit.log_likelihood(coords, rows)?;
2721 let bic =
2722 -log_likelihood + 0.5 * Self::NUM_FREE_PARAMETERS as f64 * (rows.len() as f64).ln();
2723 if !bic.is_finite() {
2724 return Err("circular Gaussian BIC is not finite".to_string());
2725 }
2726 Ok((fit, bic))
2727 }
2728}
2729
2730fn circular_gaussian_bessel_terms(
2736 radius: f64,
2737 observed_radius: f64,
2738 noise_variance: f64,
2739) -> (f64, f64) {
2740 if radius == 0.0 || observed_radius == 0.0 {
2741 return (0.0, 0.0);
2742 }
2743 let kappa = radius * observed_radius / noise_variance;
2744 if kappa.is_finite() {
2745 return bessel_i0_log_minus_abs_and_ratio(kappa);
2746 }
2747 let log_kappa = radius.ln() + observed_radius.ln() - noise_variance.ln();
2748 if log_kappa <= f64::MAX.ln() {
2749 return bessel_i0_log_minus_abs_and_ratio(log_kappa.exp());
2752 }
2753 (-0.5 * (std::f64::consts::TAU.ln() + log_kappa), 1.0)
2754}
2755#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
2776pub enum UnionStructure {
2777 CircleCircle,
2779 CirclePointCluster,
2781 LineCluster,
2783}
2784
2785pub const UNION_STRUCTURE_LADDER: &[UnionStructure] = &[
2787 UnionStructure::CircleCircle,
2788 UnionStructure::CirclePointCluster,
2789 UnionStructure::LineCluster,
2790];
2791
2792#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
2800pub enum UnionComponentKind {
2801 Circle,
2802 Line,
2803 PointCluster,
2804}
2805
2806impl UnionStructure {
2807 pub const fn as_str(self) -> &'static str {
2809 match self {
2810 UnionStructure::CircleCircle => "union_circle+circle",
2811 UnionStructure::CirclePointCluster => "union_circle+cluster",
2812 UnionStructure::LineCluster => "union_line+cluster",
2813 }
2814 }
2815
2816 pub const fn components(self) -> &'static [UnionComponentKind] {
2818 match self {
2819 UnionStructure::CircleCircle => {
2820 &[UnionComponentKind::Circle, UnionComponentKind::Circle]
2821 }
2822 UnionStructure::CirclePointCluster => {
2823 &[UnionComponentKind::Circle, UnionComponentKind::PointCluster]
2824 }
2825 UnionStructure::LineCluster => {
2826 &[UnionComponentKind::Line, UnionComponentKind::PointCluster]
2827 }
2828 }
2829 }
2830
2831 pub const fn num_components(self) -> usize {
2833 self.components().len()
2834 }
2835}
2836
2837#[derive(Debug, Clone)]
2843pub struct UnionComponentFit {
2844 pub kind: UnionComponentKind,
2845 pub row_count: usize,
2846 pub num_parameters: usize,
2847 pub mixing_weight: f64,
2848}
2849
2850#[derive(Debug, Clone)]
2854pub struct UnionStructureFit {
2855 pub structure: UnionStructure,
2856 pub components: Vec<UnionComponentFit>,
2857 pub log_likelihood: f64,
2859 pub bic: f64,
2861 pub total_parameters: usize,
2863}
2864
2865pub fn union_responsibility_split(
2870 data: ArrayView2<'_, f64>,
2871 m: usize,
2872 config: GaussianMixtureConfig,
2873) -> Result<Vec<Vec<usize>>, String> {
2874 let n = data.nrows();
2875 if m == 0 {
2876 return Err("union split requires at least one component".to_string());
2877 }
2878 if m > n {
2879 return Err(format!(
2880 "union split requested {m} groups but data has {n} rows"
2881 ));
2882 }
2883 if m == 1 {
2884 return Ok(vec![(0..n).collect()]);
2885 }
2886 let fit = fit_gaussian_mixture(data, m, config).map_err(|error| error.to_string())?;
2887 let mut groups: Vec<Vec<usize>> = vec![Vec::new(); m];
2888 let mut comp = Vec::with_capacity(m);
2890 for j in 0..m {
2891 comp.push(GaussianComponentEval::factor(
2892 fit.means.row(j),
2893 &fit.covariances[j],
2894 )?);
2895 }
2896 let log_w = fit
2897 .weights
2898 .iter()
2899 .enumerate()
2900 .map(|(component, &weight)| {
2901 if weight.is_finite() && weight > 0.0 {
2902 Ok(weight.ln())
2903 } else {
2904 Err(format!(
2905 "union split received invalid fitted weight {weight} for component {component}"
2906 ))
2907 }
2908 })
2909 .collect::<Result<Vec<_>, _>>()?;
2910 for i in 0..n {
2911 let row = data.row(i);
2912 let mut best_j = 0usize;
2913 let mut best_lt = f64::NEG_INFINITY;
2914 for j in 0..m {
2915 let lt = log_w[j] + comp[j].log_density(row);
2916 if lt > best_lt {
2917 best_lt = lt;
2918 best_j = j;
2919 }
2920 }
2921 if !best_lt.is_finite() {
2922 return Err(format!(
2923 "union split produced no finite component score at row {i}"
2924 ));
2925 }
2926 groups[best_j].push(i);
2927 }
2928 Ok(groups)
2929}
2930
2931pub fn fit_union_structure(
2940 data: ArrayView2<'_, f64>,
2941 structure: UnionStructure,
2942 config: GaussianMixtureConfig,
2943) -> Result<UnionStructureFit, String> {
2944 let fitted = fit_union_density(data, structure, config)?;
2945 Ok(UnionStructureFit {
2946 structure,
2947 components: fitted
2948 .components
2949 .iter()
2950 .map(UnionComponentDensity::summary)
2951 .collect(),
2952 log_likelihood: fitted.log_likelihood,
2953 bic: fitted.bic,
2954 total_parameters: fitted.total_parameters,
2955 })
2956}
2957
2958pub fn fit_union_ladder(
2963 data: ArrayView2<'_, f64>,
2964 config: GaussianMixtureConfig,
2965) -> Result<Vec<UnionStructureFit>, String> {
2966 let mut fits = Vec::new();
2967 let mut errors = Vec::new();
2968 for &structure in UNION_STRUCTURE_LADDER {
2969 match fit_union_structure(data, structure, config) {
2970 Ok(fit) => fits.push(fit),
2971 Err(e) => errors.push(format!("{}: {e}", structure.as_str())),
2972 }
2973 }
2974 if !errors.is_empty() {
2975 return Err(format!(
2976 "union ladder comparison failed; every declared structure must fit ({})",
2977 errors.join("; ")
2978 ));
2979 }
2980 if fits.is_empty() {
2981 return Err("union ladder is empty".to_string());
2982 }
2983 let ranked = rank_priority_candidates(
2984 fits.into_iter()
2985 .enumerate()
2986 .map(|(idx, row)| {
2987 let score = row.bic;
2988 let tie = row.total_parameters; PriorityCandidate::new(row, idx, score, tie)
2990 })
2991 .collect(),
2992 )
2993 .into_iter()
2994 .map(|row| row.item)
2995 .collect::<Vec<_>>();
2996 Ok(ranked)
2997}
2998
2999fn gather_union_rows(data: ArrayView2<'_, f64>, idx: &[usize]) -> Array2<f64> {
3000 let d = data.ncols();
3001 let mut out = Array2::<f64>::zeros((idx.len(), d));
3002 for (r, &i) in idx.iter().enumerate() {
3003 for c in 0..d {
3004 out[[r, c]] = data[[i, c]];
3005 }
3006 }
3007 out
3008}
3009
3010fn union_circle_rows(group: ArrayView2<'_, f64>) -> Result<Vec<usize>, String> {
3015 let minimum_rows = CircularGaussianFit2d::NUM_FREE_PARAMETERS + 1;
3016 if group.nrows() < minimum_rows {
3017 return Err(format!(
3018 "union circle component needs at least {minimum_rows} rows, got {}",
3019 group.nrows()
3020 ));
3021 }
3022 Ok((0..group.nrows()).collect())
3023}
3024
3025#[derive(Debug, Clone)]
3030enum UnionDensityModel {
3031 Gaussian(GaussianComponentEval),
3032 Circle(CircularGaussianFit2d),
3033}
3034
3035#[derive(Debug, Clone)]
3036struct UnionComponentDensity {
3037 kind: UnionComponentKind,
3038 row_count: usize,
3039 num_parameters: usize,
3040 mixing_weight: f64,
3041 log_weight: f64,
3042 model: UnionDensityModel,
3043}
3044
3045impl UnionComponentDensity {
3046 fn summary(&self) -> UnionComponentFit {
3047 UnionComponentFit {
3048 kind: self.kind,
3049 row_count: self.row_count,
3050 num_parameters: self.num_parameters,
3051 mixing_weight: self.mixing_weight,
3052 }
3053 }
3054
3055 fn dimension(&self) -> usize {
3056 match &self.model {
3057 UnionDensityModel::Gaussian(eval) => eval.d,
3058 UnionDensityModel::Circle(_) => 2,
3059 }
3060 }
3061
3062 fn weighted_log_density(&self, y: ArrayView1<'_, f64>) -> f64 {
3064 let component_log_density = match &self.model {
3065 UnionDensityModel::Gaussian(eval) => eval.log_density(y),
3066 UnionDensityModel::Circle(fit) => fit.log_density(y[0], y[1]),
3067 };
3068 self.log_weight + component_log_density
3069 }
3070}
3071
3072#[derive(Debug, Clone)]
3073struct FittedUnionDensity {
3074 components: Vec<UnionComponentDensity>,
3075 log_likelihood: f64,
3076 bic: f64,
3077 total_parameters: usize,
3078}
3079
3080fn fit_union_density(
3081 train: ArrayView2<'_, f64>,
3082 structure: UnionStructure,
3083 config: GaussianMixtureConfig,
3084) -> Result<FittedUnionDensity, String> {
3085 let groups = union_responsibility_split(train, structure.num_components(), config)?;
3086 fit_union_density_from_groups(train, structure, &groups, config)
3087}
3088
3089fn fit_union_density_from_groups(
3093 train: ArrayView2<'_, f64>,
3094 structure: UnionStructure,
3095 groups: &[Vec<usize>],
3096 config: GaussianMixtureConfig,
3097) -> Result<FittedUnionDensity, String> {
3098 validate_union_partition(train.nrows(), structure.num_components(), groups)?;
3099 let assignments = unique_union_role_assignments(structure.components());
3100 let mut best: Option<FittedUnionDensity> = None;
3101 let mut errors = Vec::new();
3102
3103 for roles in assignments {
3104 let candidate = (|| {
3105 let mut components = Vec::with_capacity(groups.len());
3106 let n_train = train.nrows() as f64;
3107 for (&kind, rows) in roles.iter().zip(groups) {
3108 let group = gather_union_rows(train, rows);
3109 let mixing_weight = rows.len() as f64 / n_train;
3110 components.push(fit_union_component_density(
3111 group.view(),
3112 kind,
3113 mixing_weight,
3114 config,
3115 )?);
3116 }
3117
3118 let component_parameters = components.iter().try_fold(0usize, |sum, component| {
3119 sum.checked_add(component.num_parameters)
3120 .ok_or_else(|| "union component parameter count overflowed usize".to_string())
3121 })?;
3122 let mixing_parameters = components.len() - 1;
3123 let total_parameters = component_parameters
3124 .checked_add(mixing_parameters)
3125 .ok_or_else(|| "union total parameter count overflowed usize".to_string())?;
3126 let per_point = score_union_components(&components, train)?;
3127 let log_likelihood = pairwise_sum(
3128 per_point
3129 .as_slice()
3130 .expect("owned union score vector must be contiguous"),
3131 );
3132 if !log_likelihood.is_finite() {
3133 return Err("union training log likelihood is non-finite".to_string());
3134 }
3135 let bic = -log_likelihood + 0.5 * total_parameters as f64 * (train.nrows() as f64).ln();
3136 if !bic.is_finite() {
3137 return Err("union normalized soft-mixture BIC is non-finite".to_string());
3138 }
3139 Ok(FittedUnionDensity {
3140 components,
3141 log_likelihood,
3142 bic,
3143 total_parameters,
3144 })
3145 })();
3146
3147 match candidate {
3148 Ok(candidate) => {
3149 let replace = match &best {
3150 Some(current) => candidate.bic.total_cmp(¤t.bic).is_lt(),
3151 None => true,
3152 };
3153 if replace {
3156 best = Some(candidate);
3157 }
3158 }
3159 Err(error) => errors.push(format!("{roles:?}: {error}")),
3160 }
3161 }
3162
3163 best.ok_or_else(|| {
3164 format!(
3165 "union {} has no finite role assignment ({})",
3166 structure.as_str(),
3167 errors.join("; ")
3168 )
3169 })
3170}
3171
3172fn validate_union_partition(
3173 n_rows: usize,
3174 expected_groups: usize,
3175 groups: &[Vec<usize>],
3176) -> Result<(), String> {
3177 if n_rows == 0 {
3178 return Err("union fitting requires at least one training row".to_string());
3179 }
3180 if groups.len() != expected_groups {
3181 return Err(format!(
3182 "union partition has {} groups, expected {expected_groups}",
3183 groups.len()
3184 ));
3185 }
3186 let mut seen = vec![false; n_rows];
3187 for (group_index, rows) in groups.iter().enumerate() {
3188 if rows.is_empty() {
3189 return Err(format!("union partition group {group_index} is empty"));
3190 }
3191 for &row in rows {
3192 if row >= n_rows {
3193 return Err(format!(
3194 "union partition group {group_index} contains out-of-range row {row} for {n_rows} rows"
3195 ));
3196 }
3197 if std::mem::replace(&mut seen[row], true) {
3198 return Err(format!("union partition contains duplicate row {row}"));
3199 }
3200 }
3201 }
3202 if let Some(missing) = seen.iter().position(|included| !included) {
3203 return Err(format!("union partition omits row {missing}"));
3204 }
3205 Ok(())
3206}
3207
3208fn unique_union_role_assignments(roles: &[UnionComponentKind]) -> Vec<Vec<UnionComponentKind>> {
3209 fn visit(
3210 roles: &[UnionComponentKind],
3211 used: &mut [bool],
3212 assignment: &mut Vec<UnionComponentKind>,
3213 out: &mut Vec<Vec<UnionComponentKind>>,
3214 ) {
3215 if assignment.len() == roles.len() {
3216 out.push(assignment.clone());
3217 return;
3218 }
3219 let mut used_at_depth = Vec::new();
3220 for (index, &role) in roles.iter().enumerate() {
3221 if used[index] || used_at_depth.contains(&role) {
3222 continue;
3223 }
3224 used_at_depth.push(role);
3225 used[index] = true;
3226 assignment.push(role);
3227 visit(roles, used, assignment, out);
3228 assignment.pop();
3229 used[index] = false;
3230 }
3231 }
3232
3233 let mut out = Vec::new();
3234 visit(
3235 roles,
3236 &mut vec![false; roles.len()],
3237 &mut Vec::with_capacity(roles.len()),
3238 &mut out,
3239 );
3240 out
3241}
3242
3243fn fit_union_component_density(
3244 group: ArrayView2<'_, f64>,
3245 kind: UnionComponentKind,
3246 mixing_weight: f64,
3247 config: GaussianMixtureConfig,
3248) -> Result<UnionComponentDensity, String> {
3249 if !(mixing_weight.is_finite() && mixing_weight > 0.0 && mixing_weight <= 1.0) {
3250 return Err(format!(
3251 "union component mixing weight must be finite and in (0, 1], got {mixing_weight}"
3252 ));
3253 }
3254 let row_count = group.nrows();
3255 let (model, num_parameters) = match kind {
3256 UnionComponentKind::Line => {
3257 if group.nrows() < group.ncols() + 1 {
3258 return Err(format!(
3259 "union line component needs >= {} rows, got {}",
3260 group.ncols() + 1,
3261 group.nrows()
3262 ));
3263 }
3264 let fit = fit_gaussian_mixture(group, 1, config).map_err(|error| error.to_string())?;
3265 let num_parameters = fit.num_free_parameters();
3266 let eval = GaussianComponentEval::factor(fit.means.row(0), &fit.covariances[0])?;
3267 (UnionDensityModel::Gaussian(eval), num_parameters)
3268 }
3269 UnionComponentKind::PointCluster => {
3270 if group.nrows() < group.ncols() + 1 {
3271 return Err(format!(
3272 "union isotropic point component needs >= {} rows, got {}",
3273 group.ncols() + 1,
3274 group.nrows()
3275 ));
3276 }
3277 let eval = fit_isotropic_gaussian_component(group, config.covariance_floor)?;
3278 (
3279 UnionDensityModel::Gaussian(eval),
3280 group
3281 .ncols()
3282 .checked_add(1)
3283 .ok_or_else(|| "union point parameter count overflowed usize".to_string())?,
3284 )
3285 }
3286 UnionComponentKind::Circle => {
3287 let rows = union_circle_rows(group)?;
3288 let fit = CircularGaussianFit2d::fit(group, &rows)?;
3289 (
3290 UnionDensityModel::Circle(fit),
3291 CircularGaussianFit2d::NUM_FREE_PARAMETERS,
3292 )
3293 }
3294 };
3295 Ok(UnionComponentDensity {
3296 kind,
3297 row_count,
3298 num_parameters,
3299 mixing_weight,
3300 log_weight: mixing_weight.ln(),
3301 model,
3302 })
3303}
3304
3305#[derive(Debug, Clone, Copy)]
3306struct StableScalarMeanChart {
3307 origin: f64,
3308 scale: f64,
3309 normalized_offset: f64,
3310}
3311
3312impl StableScalarMeanChart {
3313 #[inline]
3314 fn centered(self, value: f64) -> Result<f64, String> {
3315 let relative = value - self.origin;
3316 let centered = (-self.normalized_offset).mul_add(self.scale, relative);
3317 if centered.is_finite() {
3318 Ok(centered)
3319 } else {
3320 Err("union isotropic point residual is not representable".to_string())
3321 }
3322 }
3323}
3324
3325fn stable_scalar_mean_chart(values: ArrayView1<'_, f64>) -> Result<StableScalarMeanChart, String> {
3331 if values.is_empty() || values.iter().any(|value| !value.is_finite()) {
3332 return Err("stable scalar mean requires finite nonempty values".to_string());
3333 }
3334 let anchor = values[0];
3335 let anchor_chart_is_representable = values.iter().all(|&value| (value - anchor).is_finite());
3336 let origin = if anchor_chart_is_representable {
3337 anchor
3338 } else {
3339 0.0
3340 };
3341 let scale = values
3342 .iter()
3343 .map(|&value| (value - origin).abs())
3344 .fold(0.0_f64, f64::max);
3345 if scale == 0.0 {
3346 return Ok(StableScalarMeanChart {
3347 origin,
3348 scale: 0.0,
3349 normalized_offset: 0.0,
3350 });
3351 }
3352 let normalized = values
3353 .iter()
3354 .map(|&value| (value - origin) / scale)
3355 .collect::<Vec<_>>();
3356 let normalized_offset = pairwise_sum(&normalized) / values.len() as f64;
3357 let mean = normalized_offset.mul_add(scale, origin);
3358 if !(normalized_offset.is_finite() && mean.is_finite()) {
3359 return Err("union isotropic point mean is not representable".to_string());
3360 }
3361 Ok(StableScalarMeanChart {
3362 origin,
3363 scale,
3364 normalized_offset,
3365 })
3366}
3367
3368fn fit_isotropic_gaussian_component(
3370 group: ArrayView2<'_, f64>,
3371 covariance_floor: f64,
3372) -> Result<GaussianComponentEval, String> {
3373 let n = group.nrows();
3374 let d = group.ncols();
3375 if n == 0 || d == 0 {
3376 return Err("union isotropic point component requires a non-empty matrix".to_string());
3377 }
3378 if !(covariance_floor.is_finite() && covariance_floor > 0.0) {
3379 return Err(format!(
3380 "union isotropic covariance floor must be finite and positive, got {covariance_floor}"
3381 ));
3382 }
3383
3384 for row in group.rows() {
3385 for axis in 0..d {
3386 let value = row[axis];
3387 if !value.is_finite() {
3388 return Err(format!(
3389 "union isotropic point data contains non-finite coordinate {value}"
3390 ));
3391 }
3392 }
3393 }
3394
3395 let mut charts = Vec::with_capacity(d);
3396 for axis in 0..d {
3397 let chart = stable_scalar_mean_chart(group.column(axis))?;
3398 charts.push(chart);
3399 }
3400
3401 let scalar_count = n
3402 .checked_mul(d)
3403 .ok_or_else(|| "union isotropic residual count overflowed usize".to_string())?;
3404 let mut residuals = Vec::with_capacity(scalar_count);
3405 let mut residual_scale = 0.0_f64;
3406 for row in group.rows() {
3407 for axis in 0..d {
3408 let residual = charts[axis].centered(row[axis])?;
3409 residual_scale = residual_scale.max(residual.abs());
3410 residuals.push(residual);
3411 }
3412 }
3413 let variance = if residual_scale == 0.0 {
3414 covariance_floor
3415 } else {
3416 for residual in &mut residuals {
3417 *residual = (*residual / residual_scale).powi(2);
3418 }
3419 let normalized_mean_square = pairwise_sum(&residuals) / scalar_count as f64;
3420 let rms = residual_scale * normalized_mean_square.sqrt();
3421 let unconstrained = rms * rms;
3422 if !unconstrained.is_finite() {
3423 return Err("union isotropic point variance is non-finite".to_string());
3424 }
3425 unconstrained.max(covariance_floor)
3426 };
3427 GaussianComponentEval::isotropic(&charts, variance)
3428}
3429
3430fn score_union_components(
3431 components: &[UnionComponentDensity],
3432 eval: ArrayView2<'_, f64>,
3433) -> Result<Array1<f64>, String> {
3434 if components.is_empty() {
3435 return Err("union density requires at least one component".to_string());
3436 }
3437 if eval.iter().any(|coordinate| !coordinate.is_finite()) {
3438 return Err("union eval coordinates must be finite".to_string());
3439 }
3440 for component in components {
3441 if component.dimension() != eval.ncols() {
3442 return Err(format!(
3443 "union component {:?} has dimension {}, eval has {} columns",
3444 component.kind,
3445 component.dimension(),
3446 eval.ncols()
3447 ));
3448 }
3449 }
3450 let mut out = Array1::<f64>::zeros(eval.nrows());
3451 let mut terms = vec![f64::NEG_INFINITY; components.len()];
3452 for i in 0..eval.nrows() {
3453 let row = eval.row(i);
3454 let mut max_term = f64::NEG_INFINITY;
3455 for (component_index, component) in components.iter().enumerate() {
3456 let term = component.weighted_log_density(row);
3457 terms[component_index] = term;
3458 if term > max_term {
3459 max_term = term;
3460 }
3461 }
3462 let value = log_sum_exp(&terms, max_term);
3463 if !value.is_finite() {
3464 return Err(format!(
3465 "union density produced non-finite log density at eval row {i}"
3466 ));
3467 }
3468 out[i] = value;
3469 }
3470 Ok(out)
3471}
3472
3473pub fn union_per_point_log_density(
3478 train: ArrayView2<'_, f64>,
3479 eval: ArrayView2<'_, f64>,
3480 structure: UnionStructure,
3481 config: GaussianMixtureConfig,
3482) -> Result<Array1<f64>, String> {
3483 if train.ncols() != eval.ncols() {
3484 return Err(format!(
3485 "union held-out density: train has {} columns, eval has {}",
3486 train.ncols(),
3487 eval.ncols()
3488 ));
3489 }
3490 let fitted = fit_union_density(train, structure, config)?;
3491 score_union_components(&fitted.components, eval)
3492}
3493
3494#[derive(Clone, Debug)]
3496pub struct RemlCandidate {
3497 pub index: usize,
3498 pub name: String,
3499 pub score: f64,
3502 pub edf: Option<f64>,
3503 pub log_lik: Option<f64>,
3507 pub family: Option<String>,
3513 pub n_obs: Option<usize>,
3521}
3522
3523impl RemlCandidate {
3524 pub fn ranking_score(&self) -> f64 {
3543 match (self.log_lik, self.edf) {
3544 (Some(log_lik), Some(edf)) if log_lik.is_finite() && edf.is_finite() => {
3545 -2.0 * log_lik + 2.0 * edf
3546 }
3547 _ => self.score,
3548 }
3549 }
3550}
3551
3552#[derive(Clone, Debug)]
3553pub struct RemlComparison {
3554 pub ranking: Vec<RankedRow>,
3555 pub winner: String,
3556 pub evidence_summary: String,
3557 pub score_table: Vec<ScoreRow>,
3558}
3559
3560#[derive(Clone, Debug)]
3561pub struct RankedRow {
3562 pub name: String,
3563 pub score: f64,
3564 pub delta: f64,
3571 pub bayes_factor: f64,
3576 pub edf: Option<f64>,
3577}
3578
3579#[derive(Clone, Debug)]
3580pub struct ScoreRow {
3581 pub name: String,
3582 pub reml_score: f64,
3583 pub delta_reml: f64,
3584 pub bayes_factor_best_over_model: f64,
3585 pub effective_dof: Option<f64>,
3586}
3587
3588#[inline]
3590pub fn log_bayes_factor(reml_score_a: f64, reml_score_b: f64) -> f64 {
3591 reml_score_b - reml_score_a
3592}
3593
3594pub fn compare_reml_fits(mut candidates: Vec<RemlCandidate>) -> Result<RemlComparison, String> {
3598 if candidates.is_empty() {
3599 return Err("compare_models requires at least one fit".to_string());
3600 }
3601 {
3609 let mut seen_family: Option<&str> = None;
3610 for cand in &candidates {
3611 if let Some(fam) = cand.family.as_deref() {
3612 match seen_family {
3613 None => seen_family = Some(fam),
3614 Some(prev) if prev != fam => {
3615 return Err(format!(
3616 "compare_models: cannot compare fits of different response families ('{prev}' vs '{fam}'); their REML/LAML evidence scores are on incomparable base measures. Compare models fit to the same response under the same family."
3617 ));
3618 }
3619 Some(_) => {}
3620 }
3621 }
3622 }
3623 }
3624 {
3635 let mut seen_n: Option<usize> = None;
3636 for cand in &candidates {
3637 if let Some(n) = cand.n_obs {
3638 match seen_n {
3639 None => seen_n = Some(n),
3640 Some(prev) if prev != n => {
3641 return Err(format!(
3642 "compare_models: cannot compare fits made on a different number of \
3643 observations (n={prev} vs n={n}); AIC / REML-LAML evidence scales \
3644 with the sample size, so their score difference is not a Bayes \
3645 factor. Compare models fit to the same response on the same data."
3646 ));
3647 }
3648 Some(_) => {}
3649 }
3650 }
3651 }
3652 }
3653 candidates = rank_priority_candidates(
3654 candidates
3655 .into_iter()
3656 .enumerate()
3657 .map(|(idx, row)| {
3658 let ranking = row.ranking_score();
3661 PriorityCandidate::new(row, idx, ranking, 0)
3662 })
3663 .collect(),
3664 )
3665 .into_iter()
3666 .map(|row| row.item)
3667 .collect();
3668
3669 let winner = candidates[0].name.clone();
3670 let best_ranking_score = candidates[0].ranking_score();
3680 let best_raw_score = candidates
3685 .iter()
3686 .map(|c| c.score)
3687 .fold(f64::INFINITY, f64::min);
3688 let mut ranking = Vec::with_capacity(candidates.len());
3689 let mut score_table = Vec::with_capacity(candidates.len());
3690 for row in &candidates {
3691 let delta = log_bayes_factor(best_ranking_score, row.ranking_score());
3692 let bayes_factor = (0.5 * delta).exp();
3699 let delta_reml = log_bayes_factor(best_raw_score, row.score);
3700 ranking.push(RankedRow {
3701 name: row.name.clone(),
3702 score: row.score,
3703 delta,
3704 bayes_factor,
3705 edf: row.edf,
3706 });
3707 score_table.push(ScoreRow {
3708 name: row.name.clone(),
3709 reml_score: row.score,
3710 delta_reml,
3711 bayes_factor_best_over_model: delta_reml.exp(),
3712 effective_dof: row.edf,
3713 });
3714 }
3715 let evidence_summary = if let Some(runner_up) = candidates.get(1) {
3720 let margin = runner_up.ranking_score() - candidates[0].ranking_score();
3721 format!(
3726 "{} wins by Bayes factor {} over {}",
3727 winner,
3728 format_bayes_factor(0.5 * margin),
3729 runner_up.name
3730 )
3731 } else {
3732 format!("{winner} (single fit; no comparison)")
3733 };
3734 Ok(RemlComparison {
3735 ranking,
3736 winner,
3737 evidence_summary,
3738 score_table,
3739 })
3740}
3741
3742pub fn format_bayes_factor(log_bf: f64) -> String {
3743 if !log_bf.is_finite() {
3744 return "inf".to_string();
3745 }
3746 if log_bf.abs() >= std::f64::consts::LN_10 * 3.0 {
3747 return format!("1e{:+.1}", log_bf / std::f64::consts::LN_10);
3748 }
3749 format_three_significant(log_bf.exp())
3750}
3751
3752pub fn format_three_significant(value: f64) -> String {
3753 if value == 0.0 {
3754 return "0".to_string();
3755 }
3756 if !value.is_finite() {
3757 return format!("{value}");
3758 }
3759 let exponent = value.abs().log10().floor() as i32;
3760 if exponent >= 3 {
3761 return format!("{value:.2e}");
3762 }
3763 let decimals = (2 - exponent).max(0) as usize;
3764 let scale = 10f64.powi(decimals as i32);
3765 let rounded = (value * scale).abs().round() / scale * value.signum();
3766 format!("{rounded:.decimals$}")
3767}
3768
3769impl Default for TopologySelectOptions {
3770 fn default() -> Self {
3771 Self {
3772 tie_tolerance: 1e-3,
3773 score_scale: TopologyScoreScale::PerObservation,
3774 }
3775 }
3776}
3777
3778pub fn laplace_evidence(
3823 logdet_source: EvidenceLogDetSource<'_>,
3824 penalty_log_det: f64,
3825 residual_objective: f64,
3826 effective_dim: f64,
3827 penalty_rank: f64,
3828) -> f64 {
3829 if !(effective_dim.is_finite() && penalty_rank.is_finite()) {
3830 return f64::NAN;
3831 }
3832 let log_det_h = match evidence_hessian_log_det(logdet_source) {
3833 Ok(v) => v,
3834 Err(_) => return f64::NAN,
3835 };
3836 let null_dim = effective_dim - penalty_rank;
3837 if !null_dim.is_finite() || null_dim < -1e-9 {
3838 return f64::NAN;
3839 }
3840 residual_objective + 0.5 * log_det_h
3841 - 0.5 * penalty_log_det
3842 - 0.5 * null_dim.max(0.0) * (2.0 * std::f64::consts::PI).ln()
3843}
3844
3845pub fn evidence_hessian_log_det(source: EvidenceLogDetSource<'_>) -> Result<f64, String> {
3847 match source {
3848 EvidenceLogDetSource::FactoredArrow {
3849 cache,
3850 fallback_hvp,
3851 } => match arrow_log_det_from_cache(cache) {
3852 Some(v) => Ok(v),
3853 None => match fallback_hvp {
3854 Some(hvp) => hessian_log_det_from_hvp(hvp),
3855 None => {
3856 Err("evidence Hessian logdet requires exact factors or HVP fallback".into())
3857 }
3858 },
3859 },
3860 EvidenceLogDetSource::Hvp(hvp) => hessian_log_det_from_hvp(hvp),
3861 }
3862}
3863
3864pub fn hessian_log_det_from_hvp(hvp: EvidenceHvpLogDet<'_>) -> Result<f64, String> {
3871 if hvp.dim == 0 {
3872 return Ok(0.0);
3873 }
3874 if hvp.dim <= ANALYTIC_LOGDET_DENSE_DIM_THRESHOLD {
3875 let mut dense = Array2::<f64>::zeros((hvp.dim, hvp.dim));
3876 let mut basis = vec![0.0_f64; hvp.dim];
3877 for j in 0..hvp.dim {
3878 basis[j] = 1.0;
3879 let col = (hvp.apply)(&basis);
3880 basis[j] = 0.0;
3881 if col.len() != hvp.dim || col.iter().any(|v| !v.is_finite()) {
3882 return Err(format!(
3883 "evidence HVP logdet expected finite column of length {}, got {}",
3884 hvp.dim,
3885 col.len()
3886 ));
3887 }
3888 for i in 0..hvp.dim {
3889 dense[[i, j]] = col[i];
3890 }
3891 }
3892 validate_dense_hvp_symmetry(&dense)?;
3893 for i in 0..hvp.dim {
3894 for j in (i + 1)..hvp.dim {
3895 let avg = 0.5 * (dense[[i, j]] + dense[[j, i]]);
3896 dense[[i, j]] = avg;
3897 dense[[j, i]] = avg;
3898 }
3899 }
3900 dense_spd_log_det(&dense)
3901 } else {
3902 stochastic_hvp_log_det(hvp)
3903 }
3904}
3905
3906fn dense_spd_log_det(matrix: &Array2<f64>) -> Result<f64, String> {
3907 if matrix.nrows() != matrix.ncols() {
3908 return Err(format!(
3909 "evidence dense logdet requires square matrix, got {}x{}",
3910 matrix.nrows(),
3911 matrix.ncols()
3912 ));
3913 }
3914 if gam_gpu::cuda_selected().map_err(|error| error.to_string())? {
3915 return crate::gpu::reml_gpu::evidence_derivatives_gpu(
3916 crate::gpu::reml_gpu::RemlGpuInput {
3917 penalized_hessian: matrix.view(),
3918 derivative_hessians: Vec::new(),
3919 },
3920 )
3921 .map(|evidence| evidence.logdet_hessian);
3922 }
3923 let (evals, _) = matrix
3924 .eigh(Side::Lower)
3925 .map_err(|e| format!("evidence dense logdet eigendecomposition failed: {e}"))?;
3926 let mut logdet = 0.0_f64;
3927 for (idx, &ev) in evals.iter().enumerate() {
3928 if !ev.is_finite() || ev <= 0.0 {
3929 return Err(format!(
3930 "evidence dense logdet expected SPD Hessian, eigenvalue {idx} is {ev:.3e}"
3931 ));
3932 }
3933 logdet += ev.ln();
3934 }
3935 Ok(logdet)
3936}
3937
3938fn validate_dense_hvp_symmetry(matrix: &Array2<f64>) -> Result<(), String> {
3939 let n = matrix.nrows();
3940 let mut norm_sq = 0.0_f64;
3941 for &value in matrix.iter() {
3942 norm_sq += value * value;
3943 }
3944
3945 let mut skew_sq = 0.0_f64;
3946 for i in 0..n {
3947 for j in (i + 1)..n {
3948 let skew = matrix[[i, j]] - matrix[[j, i]];
3949 skew_sq += 2.0 * skew * skew;
3950 }
3951 }
3952
3953 let rel_skew = skew_sq.sqrt() / norm_sq.sqrt().max(1.0);
3954 if !rel_skew.is_finite() || rel_skew > EVIDENCE_HVP_SYMMETRY_REL_TOL {
3955 return Err(format!(
3956 "evidence HVP logdet requires symmetric operator, relative skew norm is {rel_skew:.3e}"
3957 ));
3958 }
3959 Ok(())
3960}
3961
3962fn validate_hvp_randomized_symmetry(hvp: EvidenceHvpLogDet<'_>) -> Result<(), String> {
3963 let inv_norm = 1.0 / (hvp.dim as f64).sqrt();
3964 for probe in 0..EVIDENCE_HVP_SYMMETRY_PROBES.max(1) {
3965 let mut x = vec![0.0_f64; hvp.dim];
3966 let mut y = vec![0.0_f64; hvp.dim];
3967 rademacher_unit_probe_into_slice(&mut x, (2 * probe) as u64, inv_norm);
3968 rademacher_unit_probe_into_slice(&mut y, (2 * probe + 1) as u64, inv_norm);
3969
3970 let hx = (hvp.apply)(&x);
3971 let hy = (hvp.apply)(&y);
3972 if hx.len() != hvp.dim || hx.iter().any(|v| !v.is_finite()) {
3973 return Err(format!(
3974 "evidence HVP symmetry check expected finite vector of length {}, got {}",
3975 hvp.dim,
3976 hx.len()
3977 ));
3978 }
3979 if hy.len() != hvp.dim || hy.iter().any(|v| !v.is_finite()) {
3980 return Err(format!(
3981 "evidence HVP symmetry check expected finite vector of length {}, got {}",
3982 hvp.dim,
3983 hy.len()
3984 ));
3985 }
3986
3987 let lhs = dot_slice(&x, &hy);
3988 let rhs = dot_slice(&hx, &y);
3989 let scale = (norm2_slice(&hx) * norm2_slice(&y))
3990 .max(norm2_slice(&hy) * norm2_slice(&x))
3991 .max(lhs.abs())
3992 .max(rhs.abs())
3993 .max(1.0);
3994 let rel = (lhs - rhs).abs() / scale;
3995 if !rel.is_finite() || rel > EVIDENCE_HVP_SYMMETRY_REL_TOL {
3996 return Err(format!(
3997 "evidence HVP logdet requires symmetric operator, randomized symmetry probe {probe} has relative bilinear mismatch {rel:.3e}"
3998 ));
3999 }
4000 }
4001 Ok(())
4002}
4003
4004fn stochastic_hvp_log_det(hvp: EvidenceHvpLogDet<'_>) -> Result<f64, String> {
4005 validate_hvp_randomized_symmetry(hvp)?;
4006 let probes = EVIDENCE_LOGDET_SLQ_PROBES.max(1);
4007 let steps = EVIDENCE_LOGDET_LANCZOS_STEPS.min(hvp.dim).max(1);
4008 let inv_norm = 1.0 / (hvp.dim as f64).sqrt();
4009 let mut estimate = 0.0_f64;
4010 for probe in 0..probes {
4011 let mut q0 = vec![0.0_f64; hvp.dim];
4012 rademacher_unit_probe_into_slice(&mut q0, probe as u64, inv_norm);
4013 let quad = lanczos_log_quadrature_hvp(hvp, q0, steps)?;
4014 estimate += hvp.dim as f64 * quad;
4015 }
4016 Ok(estimate / probes as f64)
4017}
4018
4019fn lanczos_log_quadrature_hvp(
4020 hvp: EvidenceHvpLogDet<'_>,
4021 q: Vec<f64>,
4022 max_steps: usize,
4023) -> Result<f64, String> {
4024 let n = hvp.dim;
4025 let eigen = symmetric_lanczos_eigenpairs(
4026 n,
4027 &q,
4028 SymmetricLanczosOptions {
4029 max_steps,
4030 residual_tol: 1e-12,
4031 local_reorthogonalize: false,
4032 full_reorthogonalize: false,
4033 },
4034 |q, out| {
4035 let applied = (hvp.apply)(q);
4036 if applied.len() != n || applied.iter().any(|v| !v.is_finite()) {
4037 return Err(format!(
4038 "evidence HVP SLQ expected finite vector of length {n}, got {}",
4039 applied.len()
4040 ));
4041 }
4042 out.copy_from_slice(&applied);
4043 Ok(())
4044 },
4045 )
4046 .map_err(|e| format!("evidence HVP SLQ Lanczos failed: {e}"))?;
4047 symmetric_lanczos_log_quadrature(&eigen, "evidence HVP SLQ expected SPD Hessian")
4048}
4049
4050#[inline]
4051fn dot_slice(a: &[f64], b: &[f64]) -> f64 {
4052 assert_eq!(a.len(), b.len());
4053 let mut s = 0.0_f64;
4054 for i in 0..a.len() {
4055 s += a[i] * b[i];
4056 }
4057 s
4058}
4059
4060#[inline]
4061fn norm2_slice(a: &[f64]) -> f64 {
4062 dot_slice(a, a).sqrt()
4063}
4064
4065fn rademacher_unit_probe_into_slice(z: &mut [f64], probe: u64, scale: f64) {
4066 let mut state = 0x6A09E667F3BCC909_u64 ^ probe.wrapping_mul(0xD1B54A32D192ED03);
4067 let mut bits = 0_u64;
4068 let mut remaining_bits = 0_u32;
4069 for value in z.iter_mut() {
4070 if remaining_bits == 0 {
4071 bits = splitmix64(&mut state);
4072 remaining_bits = 64;
4073 }
4074 *value = if bits & 1 == 0 { scale } else { -scale };
4075 bits >>= 1;
4076 remaining_bits -= 1;
4077 }
4078}
4079
4080#[inline]
4081const fn splitmix64(state: &mut u64) -> u64 {
4082 gam_linalg::utils::splitmix64(state)
4083}
4084
4085pub fn arrow_log_det_from_cache(cache: &ArrowFactorCache) -> Option<f64> {
4093 if let Some(log_det) = cache.joint_hessian_log_det {
4094 return log_det.is_finite().then_some(log_det);
4095 }
4096 if cache.ridge_t != 0.0 || cache.ridge_beta != 0.0 {
4097 return None;
4098 }
4099 if cache.k > 0 && !cache.schur_factor_is_undamped {
4100 return None;
4101 }
4102 cache.compute_undamped_arrow_log_det()
4103}
4104
4105pub fn ift_du_dbeta(cache: &ArrowFactorCache) -> Array2<f64> {
4115 let n = cache.undamped_factor_count();
4116 let total_len = cache.delta_t_len();
4117 let k = cache.k;
4118 if !cache.htbeta_available() {
4119 return Array2::<f64>::from_elem((total_len, k), f64::NAN);
4120 }
4121 let mut out = Array2::<f64>::zeros((total_len, k));
4122 let mut beta_basis = Array1::<f64>::zeros(k);
4123 let mut rhs = Array1::<f64>::zeros(cache.d);
4125 for i in 0..n {
4126 let di = cache.row_dims[i];
4127 let row_base = cache.row_offsets[i];
4128 let factor = cache.undamped_factor(i);
4129 for col in 0..k {
4131 beta_basis.fill(0.0);
4132 beta_basis[col] = 1.0;
4133 let mut rhs_i = rhs.slice_mut(ndarray::s![..di]).to_owned();
4134 if !cache.apply_htbeta_row(i, beta_basis.view(), &mut rhs_i) {
4137 return Array2::<f64>::from_elem((total_len, k), f64::NAN);
4140 }
4141 let y = cholesky_solve_vector(factor, &rhs_i);
4142 for c in 0..di {
4143 out[[row_base + c, col]] = -y[c];
4144 }
4145 }
4146 }
4147 out
4148}
4149
4150pub fn coupling_components(hessian: ArrayView2<'_, f64>) -> Vec<usize> {
4171 let p = hessian.nrows();
4172 if p == 0 || hessian.ncols() != p {
4173 return Vec::new();
4174 }
4175 let mut parent: Vec<usize> = (0..p).collect();
4177 let mut size: Vec<usize> = vec![1; p];
4178
4179 fn find(parent: &mut [usize], mut x: usize) -> usize {
4180 while parent[x] != x {
4181 parent[x] = parent[parent[x]];
4182 x = parent[x];
4183 }
4184 x
4185 }
4186
4187 for i in 0..p {
4188 for j in (i + 1)..p {
4189 if hessian[[i, j]] != 0.0 || hessian[[j, i]] != 0.0 {
4192 let (ri, rj) = (find(&mut parent, i), find(&mut parent, j));
4193 if ri != rj {
4194 let (small, large) = if size[ri] < size[rj] {
4195 (ri, rj)
4196 } else {
4197 (rj, ri)
4198 };
4199 parent[small] = large;
4200 size[large] += size[small];
4201 }
4202 }
4203 }
4204 }
4205
4206 let mut label_of_root: Vec<Option<usize>> = vec![None; p];
4209 let mut next_label = 0usize;
4210 let mut labels = vec![0usize; p];
4211 for idx in 0..p {
4212 let root = find(&mut parent, idx);
4213 let label = match label_of_root[root] {
4214 Some(l) => l,
4215 None => {
4216 let l = next_label;
4217 label_of_root[root] = Some(l);
4218 next_label += 1;
4219 l
4220 }
4221 };
4222 labels[idx] = label;
4223 }
4224 labels
4225}
4226
4227pub fn cone_of_influence(labels: &[usize], support: &[usize]) -> Vec<usize> {
4238 if support.is_empty() {
4239 return Vec::new();
4240 }
4241 let mut in_cone_labels: Vec<usize> = support
4242 .iter()
4243 .filter_map(|&idx| labels.get(idx).copied())
4244 .collect();
4245 in_cone_labels.sort_unstable();
4246 in_cone_labels.dedup();
4247 if in_cone_labels.is_empty() {
4248 return Vec::new();
4249 }
4250 (0..labels.len())
4251 .filter(|idx| in_cone_labels.binary_search(&labels[*idx]).is_ok())
4252 .collect()
4253}
4254
4255pub fn ift_dbeta_drho(
4267 cache: &ArrowFactorCache,
4268 dg_red_drho: ArrayView2<'_, f64>,
4269) -> Option<Array2<f64>> {
4270 if !cache.schur_factor_is_undamped {
4271 return None;
4272 }
4273 let schur = cache.schur_factor.as_ref()?;
4274 if dg_red_drho.nrows() != cache.k || schur.nrows() != cache.k {
4275 return None;
4276 }
4277 crate::sensitivity::FitSensitivity::from_lower_triangular(schur).mode_response(dg_red_drho)
4278}
4279
4280#[derive(Clone)]
4298pub struct EvidenceIftGradientTerms<'a> {
4299 pub dbeta_drho: ArrayView2<'a, f64>,
4300 pub du_drho: ArrayView2<'a, f64>,
4301 pub value_beta: ArrayView1<'a, f64>,
4302 pub value_u: ArrayView1<'a, f64>,
4303 pub logdet_h_beta: ArrayView1<'a, f64>,
4304 pub logdet_h_u: ArrayView1<'a, f64>,
4305}
4306
4307pub fn evidence_ift_gradient_correction(terms: EvidenceIftGradientTerms<'_>) -> Array1<f64> {
4310 let k = terms.dbeta_drho.nrows();
4311 let nd = terms.du_drho.nrows();
4312 let r = terms.dbeta_drho.ncols();
4313 if terms.du_drho.ncols() != r
4314 || terms.value_beta.len() != k
4315 || terms.logdet_h_beta.len() != k
4316 || terms.value_u.len() != nd
4317 || terms.logdet_h_u.len() != nd
4318 {
4319 return Array1::<f64>::from_elem(r, f64::NAN);
4320 }
4321
4322 let mut out = Array1::<f64>::zeros(r);
4323 for a in 0..r {
4324 let mut acc = 0.0_f64;
4325 for j in 0..k {
4326 let mode = terms.dbeta_drho[[j, a]];
4327 acc += terms.value_beta[j] * mode;
4328 acc += 0.5 * terms.logdet_h_beta[j] * mode;
4329 }
4330 for j in 0..nd {
4331 let mode = terms.du_drho[[j, a]];
4332 acc += terms.value_u[j] * mode;
4333 acc += 0.5 * terms.logdet_h_u[j] * mode;
4334 }
4335 out[a] = acc;
4336 }
4337 out
4338}
4339
4340pub fn evidence_grad_rho(
4370 cache: &ArrowFactorCache,
4371 value_rho: ArrayView1<'_, f64>,
4372 huu_drho: &[Vec<Array2<f64>>],
4373 htbeta_drho: &[Vec<Array2<f64>>],
4374 hbb_drho: &[Array2<f64>],
4375 pen_logdet_drho: ArrayView1<'_, f64>,
4376 ift_terms: EvidenceIftGradientTerms<'_>,
4377) -> Array1<f64> {
4378 let r = value_rho.len();
4379 let n = cache.undamped_factor_count();
4380 let k = cache.k;
4381 let mut out = Array1::<f64>::zeros(r);
4382 if !cache.htbeta_available()
4383 || pen_logdet_drho.len() != r
4384 || huu_drho.len() != n
4385 || htbeta_drho.len() != n
4386 || hbb_drho.len() != r
4387 || huu_drho.iter().any(|row| row.len() != r)
4388 || htbeta_drho.iter().any(|row| row.len() != r)
4389 || hbb_drho.iter().any(|m| m.nrows() != k || m.ncols() != k)
4390 || huu_drho.iter().enumerate().any(|(i, row)| {
4391 let di = cache.row_dims[i];
4392 row.iter().any(|m| m.nrows() != di || m.ncols() != di)
4393 })
4394 || htbeta_drho.iter().enumerate().any(|(i, row)| {
4395 let di = cache.row_dims[i];
4396 row.iter().any(|m| m.nrows() != di || m.ncols() != k)
4397 })
4398 {
4399 out.fill(f64::NAN);
4400 return out;
4401 }
4402 let ift_correction = evidence_ift_gradient_correction(ift_terms);
4403 if ift_correction.len() != r || ift_correction.iter().any(|v| v.is_nan()) {
4404 out.fill(f64::NAN);
4405 return out;
4406 }
4407
4408 let schur = match cache.schur_factor.as_ref() {
4409 Some(s) => s,
4410 None => {
4411 for a in 0..r {
4412 out[a] = f64::NAN;
4413 }
4414 return out;
4415 }
4416 };
4417 if !cache.schur_factor_is_undamped {
4418 for a in 0..r {
4419 out[a] = f64::NAN;
4420 }
4421 return out;
4422 }
4423
4424 let mut y_blocks: Vec<Array2<f64>> = Vec::with_capacity(n);
4427 let mut beta_basis = Array1::<f64>::zeros(k);
4428 let mut rhs = Array1::<f64>::zeros(cache.d);
4430 for i in 0..n {
4431 let di = cache.row_dims[i];
4432 let factor = cache.undamped_factor(i);
4433 let mut yi = Array2::<f64>::zeros((di, k));
4434 for col in 0..k {
4435 beta_basis.fill(0.0);
4436 beta_basis[col] = 1.0;
4437 let mut rhs_i = rhs.slice_mut(ndarray::s![..di]).to_owned();
4438 if !cache.apply_htbeta_row(i, beta_basis.view(), &mut rhs_i) {
4440 out.fill(f64::NAN);
4443 return out;
4444 }
4445 let v = cholesky_solve_vector(factor, &rhs_i);
4446 for c in 0..di {
4447 yi[[c, col]] = v[c];
4448 }
4449 }
4450 y_blocks.push(yi);
4451 }
4452
4453 let mut trace_rhs = Array1::<f64>::zeros(cache.d);
4456 let mut da_tmp = Array2::<f64>::zeros((cache.d, k));
4457 let mut col_scratch = Array1::<f64>::zeros(k);
4458 for a in 0..r {
4459 let mut grad = value_rho[a];
4461
4462 let mut row_trace_acc = 0.0_f64;
4469 for i in 0..n {
4470 let di = cache.row_dims[i];
4471 let m_i = &huu_drho[i][a];
4472 assert_eq!(m_i.shape(), &[di, di]);
4473 for col in 0..di {
4474 let mut tr_rhs_i = trace_rhs.slice_mut(ndarray::s![..di]).to_owned();
4475 for r0 in 0..di {
4476 tr_rhs_i[r0] = m_i[[r0, col]];
4477 }
4478 let v = cholesky_solve_vector(cache.undamped_factor(i), &tr_rhs_i);
4479 row_trace_acc += v[col];
4480 }
4481 }
4482
4483 let mut da = hbb_drho[a].clone();
4492 assert_eq!(da.shape(), &[k, k]);
4493 for i in 0..n {
4494 let di = cache.row_dims[i];
4495 let dhtb = &htbeta_drho[i][a]; let yi = &y_blocks[i]; for r0 in 0..k {
4499 for c0 in 0..k {
4500 let mut acc = 0.0;
4501 for cc in 0..di {
4502 acc += dhtb[[cc, r0]] * yi[[cc, c0]];
4503 }
4504 da[[r0, c0]] -= acc;
4505 }
4506 }
4507 for r0 in 0..k {
4509 for c0 in 0..k {
4510 let mut acc = 0.0;
4511 for cc in 0..di {
4512 acc += yi[[cc, r0]] * dhtb[[cc, c0]];
4513 }
4514 da[[r0, c0]] -= acc;
4515 }
4516 }
4517 let dhuu = &huu_drho[i][a];
4519 let mut da_tmp_i = da_tmp.slice_mut(ndarray::s![..di, ..]).to_owned();
4521 for r0 in 0..di {
4522 for c0 in 0..k {
4523 let mut acc = 0.0;
4524 for cc in 0..di {
4525 acc += dhuu[[r0, cc]] * yi[[cc, c0]];
4526 }
4527 da_tmp_i[[r0, c0]] = acc;
4528 }
4529 }
4530 for r0 in 0..k {
4532 for c0 in 0..k {
4533 let mut acc = 0.0;
4534 for cc in 0..di {
4535 acc += yi[[cc, r0]] * da_tmp_i[[cc, c0]];
4536 }
4537 da[[r0, c0]] += acc;
4538 }
4539 }
4540 }
4541
4542 let mut schur_trace_acc = 0.0_f64;
4544 for j in 0..k {
4545 for r0 in 0..k {
4546 col_scratch[r0] = da[[r0, j]];
4547 }
4548 let v = cholesky_solve_vector(schur, &col_scratch);
4549 schur_trace_acc += v[j];
4550 }
4551
4552 grad += 0.5 * (row_trace_acc + schur_trace_acc);
4553 grad += ift_correction[a];
4554
4555 grad -= 0.5 * pen_logdet_drho[a];
4557
4558 out[a] = grad;
4559 }
4560 out
4561}
4562
4563pub fn select_topology(
4589 candidates: &[TopologyCandidate],
4590 options: TopologySelectOptions,
4591) -> SelectedTopology {
4592 let mut valid: Vec<TopologyCandidate> = candidates
4594 .iter()
4595 .filter(|c| {
4596 c.converged
4597 && c.exclusion_reason.is_none()
4598 && c.negative_log_evidence.is_finite()
4599 && topology_selection_score(c, options.score_scale).is_finite()
4600 })
4601 .cloned()
4602 .collect();
4603 let mut excluded: Vec<TopologyCandidate> = candidates
4604 .iter()
4605 .filter(|c| {
4606 !(c.converged && c.exclusion_reason.is_none() && c.negative_log_evidence.is_finite())
4607 || !topology_selection_score(c, options.score_scale).is_finite()
4608 })
4609 .cloned()
4610 .collect();
4611
4612 assert!(
4613 !valid.is_empty(),
4614 "select_topology: no finite valid candidates; proposal §6.11 forbids silent fallback"
4615 );
4616
4617 valid = rank_priority_candidates(
4622 valid
4623 .into_iter()
4624 .enumerate()
4625 .map(|(idx, row)| {
4626 let score = topology_selection_score(&row, options.score_scale);
4627 let tie_break = usize::from(row.kind.complexity_rank());
4628 PriorityCandidate::new(row, idx, score, tie_break)
4629 })
4630 .collect(),
4631 )
4632 .into_iter()
4633 .map(|row| row.item)
4634 .collect();
4635
4636 let tie = if valid.len() >= 2 {
4638 let top = topology_selection_score(&valid[0], options.score_scale);
4639 let next = topology_selection_score(&valid[1], options.score_scale);
4640 (next - top).abs() <= options.tie_tolerance
4641 } else {
4642 false
4643 };
4644
4645 if tie {
4647 let top_score = topology_selection_score(&valid[0], options.score_scale);
4648 let tied_end = valid
4650 .iter()
4651 .position(|c| {
4652 (topology_selection_score(c, options.score_scale) - top_score).abs()
4653 > options.tie_tolerance
4654 })
4655 .unwrap_or(valid.len());
4656 valid[..tied_end].sort_by_key(|c| c.kind.complexity_rank());
4658 }
4659
4660 let winner = valid[0].kind;
4661 valid.append(&mut excluded);
4662 SelectedTopology {
4663 winner,
4664 ranking: valid,
4665 tie,
4666 }
4667}
4668
4669fn topology_selection_score(candidate: &TopologyCandidate, scale: TopologyScoreScale) -> f64 {
4670 match scale {
4671 TopologyScoreScale::PerObservation => {
4672 if candidate.n_obs == 0 {
4673 f64::NAN
4674 } else {
4675 candidate.negative_log_evidence / candidate.n_obs as f64
4676 }
4677 }
4678 TopologyScoreScale::PerEffectiveDim => {
4679 if !(candidate.effective_dim.is_finite() && candidate.effective_dim > 0.0) {
4680 f64::NAN
4681 } else {
4682 candidate.negative_log_evidence / candidate.effective_dim
4683 }
4684 }
4685 }
4686}
4687
4688pub fn cache_matches_system(cache: &ArrowFactorCache, sys: &ArrowSchurSystem) -> bool {
4696 cache.d == sys.d
4697 && cache.k == sys.k
4698 && cache.n_rows() == sys.rows.len()
4699 && cache.undamped_factor_count() == sys.rows.len()
4700 && cache.manifold_mode_fingerprint == sys.manifold_mode_fingerprint
4701 && cache.row_hessian_fingerprint == sys.current_row_hessian_fingerprint()
4702}
4703
4704#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
4757pub enum HybridAtomParam {
4758 Curved { latent_dim: usize },
4760 Linear,
4762}
4763
4764impl HybridAtomParam {
4765 pub const fn as_str(self) -> &'static str {
4767 match self {
4768 HybridAtomParam::Curved { .. } => "curved",
4769 HybridAtomParam::Linear => "linear",
4770 }
4771 }
4772
4773 pub const fn is_linear(self) -> bool {
4775 matches!(self, HybridAtomParam::Linear)
4776 }
4777}
4778
4779#[derive(Debug, Clone, Copy)]
4787pub struct HybridAtomCandidate {
4788 pub param: HybridAtomParam,
4789 pub negative_log_evidence: f64,
4791 pub num_parameters: usize,
4793 pub fitted_turning: Option<f64>,
4798}
4799
4800impl HybridAtomCandidate {
4801 pub fn linear(negative_log_evidence: f64, num_parameters: usize) -> Self {
4803 Self {
4804 param: HybridAtomParam::Linear,
4805 negative_log_evidence,
4806 num_parameters,
4807 fitted_turning: Some(0.0),
4808 }
4809 }
4810
4811 pub fn curved(
4813 latent_dim: usize,
4814 negative_log_evidence: f64,
4815 num_parameters: usize,
4816 fitted_turning: Option<f64>,
4817 ) -> Self {
4818 Self {
4819 param: HybridAtomParam::Curved { latent_dim },
4820 negative_log_evidence,
4821 num_parameters,
4822 fitted_turning,
4823 }
4824 }
4825}
4826
4827#[derive(Debug, Clone, Copy)]
4831pub struct HybridAtomChoice {
4832 pub param: HybridAtomParam,
4833 pub negative_log_evidence: f64,
4835 pub num_parameters: usize,
4837 pub curved_turning: Option<f64>,
4840 pub curved_evidence_margin: f64,
4845}
4846
4847pub const HYBRID_LINEAR_TURNING_FLOOR: f64 = 1e-9;
4855
4856pub fn select_hybrid_atom(candidates: &[HybridAtomCandidate]) -> Option<HybridAtomChoice> {
4884 if candidates.is_empty() {
4885 return None;
4886 }
4887 let linear = candidates.iter().find(|c| c.param.is_linear());
4888 let curved = candidates.iter().find(|c| !c.param.is_linear());
4889 let curved_turning = curved.and_then(|c| c.fitted_turning);
4890 let curved_evidence_margin = match (linear, curved) {
4891 (Some(l), Some(c)) => l.negative_log_evidence - c.negative_log_evidence,
4892 _ => 0.0,
4893 };
4894
4895 if let (Some(l), Some(turning)) = (linear, curved_turning)
4898 && turning <= HYBRID_LINEAR_TURNING_FLOOR
4899 {
4900 return Some(HybridAtomChoice {
4901 param: l.param,
4902 negative_log_evidence: l.negative_log_evidence,
4903 num_parameters: l.num_parameters,
4904 curved_turning,
4905 curved_evidence_margin,
4906 });
4907 }
4908
4909 let mut best = candidates[0];
4911 for cand in &candidates[1..] {
4912 let better_evidence = cand.negative_log_evidence < best.negative_log_evidence;
4913 let tied = cand.negative_log_evidence == best.negative_log_evidence;
4914 let cheaper_on_tie = tied && cand.num_parameters < best.num_parameters;
4915 if better_evidence || cheaper_on_tie {
4916 best = *cand;
4917 }
4918 }
4919 Some(HybridAtomChoice {
4920 param: best.param,
4921 negative_log_evidence: best.negative_log_evidence,
4922 num_parameters: best.num_parameters,
4923 curved_turning,
4924 curved_evidence_margin,
4925 })
4926}
4927
4928#[derive(Debug, Clone)]
4932pub struct HybridSplitSelection {
4933 pub atoms: Vec<HybridAtomChoice>,
4935 pub total_negative_log_evidence: f64,
4944 pub total_parameters: usize,
4947 pub curved_atom_count: usize,
4949}
4950
4951impl HybridSplitSelection {
4952 pub fn linear_atom_count(&self) -> usize {
4954 self.atoms.len() - self.curved_atom_count
4955 }
4956
4957 pub fn is_pure_linear(&self) -> bool {
4960 self.curved_atom_count == 0 && !self.atoms.is_empty()
4961 }
4962
4963 pub fn is_pure_curved(&self) -> bool {
4966 self.curved_atom_count == self.atoms.len() && !self.atoms.is_empty()
4967 }
4968}
4969
4970pub fn select_hybrid_split(
4984 slots: &[Vec<HybridAtomCandidate>],
4985) -> Result<HybridSplitSelection, String> {
4986 let mut atoms = Vec::with_capacity(slots.len());
4987 let mut total_nle = 0.0_f64;
4988 let mut total_parameters = 0usize;
4989 let mut curved_atom_count = 0usize;
4990 for (i, slot) in slots.iter().enumerate() {
4991 let choice = select_hybrid_atom(slot)
4992 .ok_or_else(|| format!("hybrid split slot {i} has no candidate parameterizations"))?;
4993 if !choice.negative_log_evidence.is_finite() {
4994 return Err(format!(
4995 "hybrid split slot {i} selected a non-finite evidence ({})",
4996 choice.negative_log_evidence
4997 ));
4998 }
4999 if !choice.param.is_linear() {
5000 curved_atom_count += 1;
5001 }
5002 total_nle += choice.negative_log_evidence;
5003 total_parameters += choice.num_parameters;
5004 atoms.push(choice);
5005 }
5006 Ok(HybridSplitSelection {
5007 atoms,
5008 total_negative_log_evidence: total_nle,
5009 total_parameters,
5010 curved_atom_count,
5011 })
5012}
5013
5014#[cfg(test)]
5024mod tests {
5025 use super::*;
5026 use crate::arrow_schur::ArrowFactorSlab;
5027 use ndarray::array;
5028
5029 fn dense_inverse(h: &Array2<f64>) -> Array2<f64> {
5031 let p = h.nrows();
5032 let mut aug = Array2::<f64>::zeros((p, 2 * p));
5033 for i in 0..p {
5034 for j in 0..p {
5035 aug[[i, j]] = h[[i, j]];
5036 }
5037 aug[[i, p + i]] = 1.0;
5038 }
5039 for col in 0..p {
5040 let mut pivot = col;
5041 for row in (col + 1)..p {
5042 if aug[[row, col]].abs() > aug[[pivot, col]].abs() {
5043 pivot = row;
5044 }
5045 }
5046 if pivot != col {
5047 for j in 0..(2 * p) {
5048 aug.swap([col, j], [pivot, j]);
5049 }
5050 }
5051 let d = aug[[col, col]];
5052 for j in 0..(2 * p) {
5053 aug[[col, j]] /= d;
5054 }
5055 for row in 0..p {
5056 if row == col {
5057 continue;
5058 }
5059 let f = aug[[row, col]];
5060 if f != 0.0 {
5061 for j in 0..(2 * p) {
5062 aug[[row, j]] -= f * aug[[col, j]];
5063 }
5064 }
5065 }
5066 }
5067 let mut inv = Array2::<f64>::zeros((p, p));
5068 for i in 0..p {
5069 for j in 0..p {
5070 inv[[i, j]] = aug[[i, p + j]];
5071 }
5072 }
5073 inv
5074 }
5075
5076 #[test]
5080 fn em_contraction_rate_recovers_a_planted_geometric_decay() {
5081 let planted = 0.98_f64;
5082 let mut window = std::collections::VecDeque::new();
5083 let mut residual = 1.0_f64;
5084 for _ in 0..EM_RATE_WINDOW {
5085 window.push_back(residual);
5086 residual *= planted;
5087 }
5088 assert_eq!(em_contraction_rate(&window), None);
5090 window.push_back(residual);
5091 let measured = em_contraction_rate(&window).expect("a full window yields a rate");
5092 assert!(
5093 (measured - planted).abs() < 1e-12,
5094 "measured {measured} should recover the planted {planted}"
5095 );
5096 }
5097
5098 #[test]
5101 fn em_contraction_rate_does_not_contract_on_a_flat_or_growing_residual() {
5102 let flat: std::collections::VecDeque<f64> =
5103 std::iter::repeat_n(1e-6, EM_RATE_WINDOW + 1).collect();
5104 let rate = em_contraction_rate(&flat).expect("a full window yields a rate");
5105 assert!(rate >= 1.0, "a flat residual must not look like contraction");
5106 let growing: std::collections::VecDeque<f64> = (0..=EM_RATE_WINDOW)
5107 .map(|i| 1e-6 * 1.01_f64.powi(i as i32))
5108 .collect();
5109 let rate = em_contraction_rate(&growing).expect("a full window yields a rate");
5110 assert!(rate > 1.0, "a growing residual must not look like contraction");
5111 }
5112
5113 #[test]
5116 fn em_projected_iterations_inverts_the_decay_and_declines_otherwise() {
5117 let steps = em_projected_iterations(1.0, 1e-8, 0.98).expect("a contracting rate projects");
5119 assert_eq!(steps, 912);
5120 assert!(0.98_f64.powi(steps as i32) <= 1e-8);
5122 assert_eq!(em_projected_iterations(1.0, 1e-8, 1.0), None);
5124 assert_eq!(em_projected_iterations(1.0, 1e-8, 1.05), None);
5125 assert_eq!(em_projected_iterations(1e-9, 1e-8, 0.98), None);
5126 }
5127
5128 #[test]
5129 fn coupling_components_block_diagonal_is_all_singletons_by_block() {
5130 let mut h = Array2::<f64>::eye(4);
5132 h[[0, 1]] = 0.3;
5133 h[[1, 0]] = 0.3;
5134 h[[2, 3]] = 0.7;
5135 h[[3, 2]] = 0.7;
5136 let labels = coupling_components(h.view());
5137 assert_eq!(labels[0], labels[1]);
5138 assert_eq!(labels[2], labels[3]);
5139 assert_ne!(labels[0], labels[2]);
5140 let mut uniq = labels.clone();
5142 uniq.sort_unstable();
5143 uniq.dedup();
5144 assert_eq!(uniq.len(), 2);
5145 }
5146
5147 #[test]
5148 fn coupling_components_fully_coupled_is_one_component() {
5149 let mut h = Array2::<f64>::eye(3);
5150 for i in 0..3 {
5151 for j in 0..3 {
5152 if i != j {
5153 h[[i, j]] = 0.1;
5154 }
5155 }
5156 }
5157 let labels = coupling_components(h.view());
5158 assert!(labels.iter().all(|&l| l == labels[0]));
5159 }
5160
5161 #[test]
5162 fn coupling_components_transitive_chain_merges() {
5163 let mut h = Array2::<f64>::eye(3);
5165 h[[0, 1]] = 0.5;
5166 h[[1, 0]] = 0.5;
5167 h[[1, 2]] = 0.5;
5168 h[[2, 1]] = 0.5;
5169 let labels = coupling_components(h.view());
5170 assert_eq!(labels[0], labels[1]);
5171 assert_eq!(labels[1], labels[2]);
5172 }
5173
5174 #[test]
5175 fn compare_reml_fits_delta_and_bayes_factor_never_contradict_winner_gh1465() {
5176 let cand = |name: &str, score: f64, edf: f64| RemlCandidate {
5188 index: 0,
5189 name: name.to_string(),
5190 score,
5191 edf: Some(edf),
5192 log_lik: Some(0.0),
5193 family: Some("gaussian".to_string()),
5194 n_obs: Some(100),
5195 };
5196 let candidates = vec![
5199 cand("m1", 53.748, 50.0),
5200 cand("m2", 41.605, 51.0),
5201 cand("m3", 120.011, 65.0),
5202 ];
5203 let cmp = compare_reml_fits(candidates).expect("comparison");
5204
5205 assert_eq!(cmp.winner, "m1", "AIC winner");
5206 for row in &cmp.ranking {
5208 assert!(
5209 row.delta >= 0.0,
5210 "ranking delta for {} must be >= 0, got {}",
5211 row.name,
5212 row.delta
5213 );
5214 assert!(
5215 row.bayes_factor >= 1.0 - 1e-12,
5216 "ranking bayes_factor for {} must be >= 1, got {}",
5217 row.name,
5218 row.bayes_factor
5219 );
5220 }
5221 let winner_row = cmp.ranking.iter().find(|r| r.name == "m1").unwrap();
5222 assert!(winner_row.delta.abs() < 1e-12, "winner delta == 0");
5223 assert!(
5224 (winner_row.bayes_factor - 1.0).abs() < 1e-9,
5225 "winner bayes_factor == 1"
5226 );
5227
5228 for row in &cmp.score_table {
5231 assert!(
5232 row.delta_reml >= 0.0,
5233 "score-table delta_reml for {} must be >= 0, got {}",
5234 row.name,
5235 row.delta_reml
5236 );
5237 assert!(
5238 row.bayes_factor_best_over_model >= 1.0 - 1e-12,
5239 "score-table bayes_factor for {} must be >= 1, got {}",
5240 row.name,
5241 row.bayes_factor_best_over_model
5242 );
5243 }
5244 let m2 = cmp.score_table.iter().find(|r| r.name == "m2").unwrap();
5246 assert!(
5247 m2.delta_reml.abs() < 1e-12,
5248 "the minimum-raw-REML row has delta_reml 0"
5249 );
5250 }
5251
5252 #[test]
5253 fn cone_of_influence_empty_support_is_empty() {
5254 let labels = vec![0usize, 0, 1, 1];
5255 assert!(cone_of_influence(&labels, &[]).is_empty());
5256 }
5257
5258 #[test]
5259 fn cone_of_influence_returns_full_component() {
5260 let labels = vec![0usize, 0, 1, 1];
5261 assert_eq!(cone_of_influence(&labels, &[0]), vec![0, 1]);
5263 assert_eq!(cone_of_influence(&labels, &[1, 2]), vec![0, 1, 2, 3]);
5265 }
5266
5267 #[test]
5268 fn coned_matches_full_solve_on_fully_coupled_hessian() {
5269 let h = Array2::from_shape_vec((3, 3), vec![4.0, 1.0, 0.5, 1.0, 3.0, 0.8, 0.5, 0.8, 2.5])
5272 .unwrap();
5273 let inv = dense_inverse(&h);
5274 let mut dg = Array2::<f64>::zeros((3, 2));
5276 dg[[0, 0]] = 1.3;
5277 dg[[2, 1]] = -0.7;
5278 let supports = vec![0..1usize, 2..3usize];
5279
5280 let eye: Array2<f64> = Array2::eye(3);
5281 let op = crate::sensitivity::FitSensitivity::from_projected(&eye, &inv);
5282 let full = op.mode_response(dg.view()).unwrap();
5283 let coned = op
5284 .mode_response_coned(h.view(), dg.view(), &supports)
5285 .unwrap();
5286 for i in 0..3 {
5287 for a in 0..2 {
5288 assert!(
5289 (full[[i, a]] - coned[[i, a]]).abs() < 1e-12,
5290 "fully-coupled mismatch at ({i},{a}): {} vs {}",
5291 full[[i, a]],
5292 coned[[i, a]]
5293 );
5294 }
5295 }
5296 }
5297
5298 #[test]
5299 fn coned_confines_to_component_on_decoupled_hessian() {
5300 let mut h = Array2::<f64>::zeros((4, 4));
5304 h[[0, 0]] = 4.0;
5306 h[[1, 1]] = 3.0;
5307 h[[0, 1]] = 1.0;
5308 h[[1, 0]] = 1.0;
5309 h[[2, 2]] = 2.0;
5311 h[[3, 3]] = 5.0;
5312 h[[2, 3]] = 0.6;
5313 h[[3, 2]] = 0.6;
5314 let inv = dense_inverse(&h);
5315
5316 let mut dg = Array2::<f64>::zeros((4, 1));
5317 dg[[0, 0]] = 0.9;
5318 dg[[1, 0]] = -0.4;
5319 let support_range = 0..2usize;
5320 let supports = std::slice::from_ref(&support_range);
5321
5322 let eye: Array2<f64> = Array2::eye(4);
5323 let coned = crate::sensitivity::FitSensitivity::from_projected(&eye, &inv)
5324 .mode_response_coned(h.view(), dg.view(), supports)
5325 .unwrap();
5326 let q = dg.column(0).to_owned();
5329 let exact = inv.dot(&q).mapv(|v| -v);
5330 for i in 0..4 {
5331 assert!(
5332 (coned[[i, 0]] - exact[[i]]).abs() < 1e-12,
5333 "decoupled mismatch at {i}: {} vs {}",
5334 coned[[i, 0]],
5335 exact[[i]]
5336 );
5337 }
5338 assert_eq!(coned[[2, 0]], 0.0);
5340 assert_eq!(coned[[3, 0]], 0.0);
5341 }
5342
5343 #[test]
5344 fn coned_skips_inactive_column_with_empty_support() {
5345 let h = Array2::<f64>::eye(2);
5346 let dg = Array2::<f64>::zeros((2, 1));
5347 let empty_support = 0..0usize;
5349 let supports = std::slice::from_ref(&empty_support);
5350 let eye: Array2<f64> = Array2::eye(2);
5355 let nan_inv = Array2::<f64>::from_elem((2, 2), f64::NAN);
5356 let coned = crate::sensitivity::FitSensitivity::from_projected(&eye, &nan_inv)
5357 .mode_response_coned(h.view(), dg.view(), supports)
5358 .unwrap();
5359 assert_eq!(coned[[0, 0]], 0.0);
5360 assert_eq!(coned[[1, 0]], 0.0);
5361 }
5362
5363 fn make_minimal_cache() -> ArrowFactorCache {
5364 let l_huu = Array2::from_shape_vec((1, 1), vec![std::f64::consts::SQRT_2]).unwrap();
5367 let l_schur = Array2::from_shape_vec((1, 1), vec![(1.875_f64).sqrt()]).unwrap();
5368 let htbeta = Array2::from_shape_vec((1, 1), vec![0.5]).unwrap();
5369 let mut cache = ArrowFactorCache {
5370 htt_factors: ArrowFactorSlab::from_blocks(vec![l_huu]),
5371 htt_factors_undamped: crate::arrow_schur::ArrowUndampedFactors::SameAsDamped,
5372 schur_factor: Some(l_schur),
5373 schur_factor_is_undamped: true,
5374 beta_schur_deflation: None,
5375 joint_hessian_log_det: None,
5376 solver_mode: crate::arrow_schur::ArrowSolverMode::Direct,
5377 ridge_t: 0.0,
5378 ridge_beta: 0.0,
5379 htbeta: crate::arrow_schur::ArrowHtbetaCache::Dense {
5380 blocks: std::sync::Arc::from(vec![htbeta]),
5381 estimated_bytes: std::mem::size_of::<f64>(),
5382 },
5383 d: 1,
5384 row_dims: std::sync::Arc::from(vec![1usize]),
5385 row_offsets: std::sync::Arc::from(vec![0usize, 1usize]),
5386 k: 1,
5387 manifold_mode_fingerprint: 0,
5388 row_hessian_fingerprint: 0,
5389 pcg_diagnostics: crate::arrow_schur::ArrowPcgDiagnostics::default(),
5390 gauge_deflated_directions: 0,
5391 deflated_row_directions: std::sync::Arc::from(Vec::new()),
5392 deflation_row_spectra: std::sync::Arc::from(Vec::new()),
5393 beta_gauge_quotient: None,
5394 };
5395 cache.joint_hessian_log_det = cache.compute_undamped_arrow_log_det();
5396 cache
5397 }
5398
5399 #[test]
5400 fn laplace_evidence_returns_finite_for_minimal_cache() {
5401 let cache = make_minimal_cache();
5402 let v = laplace_evidence(
5405 EvidenceLogDetSource::FactoredArrow {
5406 cache: &cache,
5407 fallback_hvp: None,
5408 },
5409 0.0,
5410 0.0,
5411 2.0,
5412 1.0,
5413 );
5414 assert!(v.is_finite());
5415 let expected =
5416 0.5 * (2.0_f64.ln() + 1.875_f64.ln()) - 0.5 * (2.0 * std::f64::consts::PI).ln();
5417 assert!((v - expected).abs() < 1e-12);
5418 }
5419
5420 fn k0_direct_cache_no_schur(latent_diag: f64) -> ArrowFactorCache {
5429 let l_huu = Array2::from_shape_vec((1, 1), vec![latent_diag.sqrt()]).unwrap();
5430 let mut cache = ArrowFactorCache {
5431 htt_factors: ArrowFactorSlab::from_blocks(vec![l_huu]),
5432 htt_factors_undamped: crate::arrow_schur::ArrowUndampedFactors::SameAsDamped,
5433 schur_factor: None,
5434 schur_factor_is_undamped: true,
5435 beta_schur_deflation: None,
5436 joint_hessian_log_det: None,
5437 solver_mode: crate::arrow_schur::ArrowSolverMode::Direct,
5438 ridge_t: 0.0,
5439 ridge_beta: 0.0,
5440 htbeta: crate::arrow_schur::ArrowHtbetaCache::Disabled { estimated_bytes: 0 },
5441 d: 1,
5442 row_dims: std::sync::Arc::from(vec![1usize]),
5443 row_offsets: std::sync::Arc::from(vec![0usize, 1usize]),
5444 k: 0,
5445 manifold_mode_fingerprint: 0,
5446 row_hessian_fingerprint: 0,
5447 pcg_diagnostics: crate::arrow_schur::ArrowPcgDiagnostics::default(),
5448 gauge_deflated_directions: 0,
5449 deflated_row_directions: std::sync::Arc::from(Vec::new()),
5450 deflation_row_spectra: std::sync::Arc::from(Vec::new()),
5451 beta_gauge_quotient: None,
5452 };
5453 cache.joint_hessian_log_det = cache.compute_undamped_arrow_log_det();
5454 cache
5455 }
5456
5457 #[test]
5458 fn arrow_log_det_some_for_k0_direct_cache_without_schur() {
5459 let cache = k0_direct_cache_no_schur(3.0);
5460 let log_det = arrow_log_det_from_cache(&cache)
5461 .expect("k==0 Direct cache must yield Some(per-row sum), not None (#1132)");
5462 assert!(
5464 (log_det - 3.0_f64.ln()).abs() < 1e-12,
5465 "log_det = {log_det}"
5466 );
5467 let cached = cache
5469 .compute_undamped_arrow_log_det()
5470 .expect("compute_undamped_arrow_log_det must be Some for k==0");
5471 assert!((cached - 3.0_f64.ln()).abs() < 1e-12, "cached = {cached}");
5472 }
5473
5474 #[test]
5475 fn arrow_log_det_none_for_kpos_cache_without_schur() {
5476 let mut cache = k0_direct_cache_no_schur(3.0);
5479 cache.k = 1;
5480 cache.solver_mode = crate::arrow_schur::ArrowSolverMode::InexactPCG;
5481 cache.joint_hessian_log_det = None;
5482 assert!(arrow_log_det_from_cache(&cache).is_none());
5483 assert!(cache.compute_undamped_arrow_log_det().is_none());
5484 }
5485
5486 #[test]
5487 fn laplace_evidence_nan_when_authoritative_logdet_missing() {
5488 let mut cache = make_minimal_cache();
5489 cache.ridge_t = 1e-3;
5490 cache.joint_hessian_log_det = None;
5491 assert!(
5492 laplace_evidence(
5493 EvidenceLogDetSource::FactoredArrow {
5494 cache: &cache,
5495 fallback_hvp: None,
5496 },
5497 0.0,
5498 0.0,
5499 2.0,
5500 1.0,
5501 )
5502 .is_nan()
5503 );
5504 }
5505
5506 #[test]
5507 fn laplace_evidence_uses_hvp_fallback_without_authoritative_logdet() {
5508 let mut cache = make_minimal_cache();
5509 cache.schur_factor = None;
5510 cache.joint_hessian_log_det = None;
5511 let hvp = |x: &[f64]| -> Vec<f64> { vec![2.0 * x[0], 1.875 * x[1]] };
5512 let v = laplace_evidence(
5513 EvidenceLogDetSource::FactoredArrow {
5514 cache: &cache,
5515 fallback_hvp: Some(EvidenceHvpLogDet {
5516 dim: 2,
5517 apply: &hvp,
5518 }),
5519 },
5520 0.0,
5521 0.0,
5522 2.0,
5523 1.0,
5524 );
5525 let expected =
5526 0.5 * (2.0_f64.ln() + 1.875_f64.ln()) - 0.5 * (2.0 * std::f64::consts::PI).ln();
5527 assert!((v - expected).abs() < 1e-12);
5528 }
5529
5530 #[test]
5531 fn ift_du_dbeta_has_expected_shape() {
5532 let cache = make_minimal_cache();
5533 let du_db = ift_du_dbeta(&cache);
5534 assert_eq!(du_db.shape(), &[1, 1]);
5535 assert!((du_db[[0, 0]] - (-0.25)).abs() < 1e-12);
5537 }
5538
5539 #[test]
5540 fn ift_dbeta_drho_returns_some_for_direct_cache() {
5541 let cache = make_minimal_cache();
5542 let q = Array2::from_shape_vec((1, 1), vec![1.0]).unwrap();
5543 let out = ift_dbeta_drho(&cache, q.view()).unwrap();
5544 assert_eq!(out.shape(), &[1, 1]);
5545 assert!((out[[0, 0]] + 1.0 / 1.875).abs() < 1e-12);
5547 }
5548
5549 #[test]
5550 fn topology_select_picks_lowest_negative_log_evidence() {
5551 let candidates = vec![
5552 TopologyCandidate {
5553 kind: TopologyKind::Flat,
5554 negative_log_evidence: 10.0,
5555 effective_dim: 4.0,
5556 n_obs: 100,
5557 converged: true,
5558 exclusion_reason: None,
5559 },
5560 TopologyCandidate {
5561 kind: TopologyKind::Sphere,
5562 negative_log_evidence: 8.0,
5563 effective_dim: 5.0,
5564 n_obs: 100,
5565 converged: true,
5566 exclusion_reason: None,
5567 },
5568 TopologyCandidate {
5569 kind: TopologyKind::Torus,
5570 negative_log_evidence: f64::NAN,
5571 effective_dim: 6.0,
5572 n_obs: 100,
5573 converged: false,
5574 exclusion_reason: Some("torus periods missing".to_string()),
5575 },
5576 ];
5577 let sel = select_topology(&candidates, TopologySelectOptions::default());
5578 assert_eq!(sel.winner, TopologyKind::Sphere);
5579 assert!(!sel.tie);
5580 }
5581
5582 #[test]
5583 fn topology_select_tie_breaks_to_simpler() {
5584 let candidates = vec![
5585 TopologyCandidate {
5586 kind: TopologyKind::Sphere,
5587 negative_log_evidence: 5.0,
5588 effective_dim: 5.0,
5589 n_obs: 100,
5590 converged: true,
5591 exclusion_reason: None,
5592 },
5593 TopologyCandidate {
5594 kind: TopologyKind::Flat,
5595 negative_log_evidence: 5.0 + 1e-6,
5596 effective_dim: 4.0,
5597 n_obs: 100,
5598 converged: true,
5599 exclusion_reason: None,
5600 },
5601 ];
5602 let sel = select_topology(&candidates, TopologySelectOptions::default());
5603 assert_eq!(sel.winner, TopologyKind::Flat);
5604 assert!(sel.tie);
5605 }
5606
5607 fn gaussian_logpdf(y: f64, mean: f64, sd: f64) -> f64 {
5608 let z = (y - mean) / sd;
5609 -0.5 * (2.0 * std::f64::consts::PI).ln() - sd.ln() - 0.5 * z * z
5610 }
5611
5612 #[test]
5613 fn stacking_single_candidate_gets_full_weight() {
5614 let log_density = Array2::from_shape_vec((3, 1), vec![-1.0, -2.0, -0.5]).unwrap();
5615 let out = solve_stacking_weights(log_density.view(), StackingConfig::default()).unwrap();
5616 assert!((out.weights[0] - 1.0).abs() < 1e-12);
5617 assert_eq!(out.weights.len(), 1);
5618 }
5619
5620 #[test]
5621 fn stacking_dominant_candidate_attracts_nearly_all_weight() {
5622 let mut log_density = Array2::<f64>::zeros((50, 2));
5623 for i in 0..50 {
5624 log_density[[i, 0]] = -0.1;
5625 log_density[[i, 1]] = -5.0;
5626 }
5627 let out = solve_stacking_weights(log_density.view(), StackingConfig::default()).unwrap();
5628 assert!(out.weights[0] > 0.99, "w0 = {}", out.weights[0]);
5629 assert!(out.weights[1] < 0.01, "w1 = {}", out.weights[1]);
5630 }
5631
5632 #[test]
5633 fn stacking_complementary_candidates_share_weight() {
5634 let n = 40;
5637 let mut log_density = Array2::<f64>::zeros((n, 2));
5638 for i in 0..n {
5639 if i < n / 2 {
5640 log_density[[i, 0]] = gaussian_logpdf(0.0, 0.0, 0.5);
5641 log_density[[i, 1]] = gaussian_logpdf(0.0, 1.5, 0.5);
5642 } else {
5643 log_density[[i, 0]] = gaussian_logpdf(0.0, 1.5, 0.5);
5644 log_density[[i, 1]] = gaussian_logpdf(0.0, 0.0, 0.5);
5645 }
5646 }
5647 let out = solve_stacking_weights(log_density.view(), StackingConfig::default()).unwrap();
5648 assert!(
5649 out.weights[0] > 0.2 && out.weights[0] < 0.8,
5650 "w0 = {}",
5651 out.weights[0]
5652 );
5653 assert!((out.weights.sum() - 1.0).abs() < 1e-9);
5654 }
5655
5656 #[test]
5657 fn stacking_weights_stay_on_the_simplex() {
5658 let log_density = Array2::from_shape_vec(
5659 (3, 3),
5660 vec![-1.0, -2.0, -3.0, -2.5, -1.0, -2.0, -3.0, -2.0, -1.0],
5661 )
5662 .unwrap();
5663 let out = solve_stacking_weights(log_density.view(), StackingConfig::default()).unwrap();
5664 assert!((out.weights.sum() - 1.0).abs() < 1e-9);
5665 assert!(out.weights.iter().all(|&w| w >= -1e-12));
5666 }
5667
5668 #[test]
5669 fn stacking_solution_satisfies_the_simplex_kkt_certificate() {
5670 let log_density = Array2::from_shape_vec(
5675 (5, 2),
5676 vec![-0.2, -3.0, -3.0, -0.2, -0.5, -1.5, -1.5, -0.5, -0.1, -2.0],
5677 )
5678 .unwrap();
5679 let config = StackingConfig::default();
5680 let out = solve_stacking_weights(log_density.view(), config).unwrap();
5681 assert!(out.certificate.residual() <= config.kkt_tol);
5682 let n = log_density.nrows();
5683 for k in 0..2 {
5684 let mut g = 0.0_f64;
5685 for i in 0..n {
5686 let mix: f64 = (0..2)
5687 .map(|c| out.weights[c] * log_density[[i, c]].exp())
5688 .sum();
5689 g += log_density[[i, k]].exp() / mix;
5690 }
5691 g /= n as f64;
5692 assert!(
5693 g <= 1.0 + config.kkt_tol,
5694 "stationarity violated for candidate {k}: g = {g}"
5695 );
5696 assert!(
5697 out.weights[k] * (g - 1.0).abs() <= config.kkt_tol * (1.0 + 1e-6),
5698 "complementary slackness violated for candidate {k}: w = {}, g = {g}",
5699 out.weights[k]
5700 );
5701 }
5702 }
5703
5704 #[test]
5705 fn stacking_exhaustion_without_certificate_is_an_error_not_weights() {
5706 let log_density = Array2::from_shape_vec(
5707 (6, 3),
5708 vec![
5709 0.0, -2.0, -4.0, -0.4, -0.1, -3.0, -2.0, 0.0, -0.3, -3.0, -1.0, 0.0, -0.2, -2.0,
5710 -0.5, -1.0, -0.3, -2.0,
5711 ],
5712 )
5713 .unwrap();
5714 let config = StackingConfig {
5715 max_iter: 1,
5716 ..StackingConfig::default()
5717 };
5718 let err = solve_stacking_weights(log_density.view(), config).unwrap_err();
5719 let checkpoint = match err {
5720 StackingError::DidNotConverge {
5721 certificate,
5722 checkpoint,
5723 ..
5724 } => {
5725 assert!(certificate.residual() > config.kkt_tol);
5726 assert_eq!(checkpoint.completed_iterations, 1);
5727 checkpoint
5728 }
5729 other => panic!("expected typed stacking exhaustion, got {other}"),
5730 };
5731 let encoded = serde_json::to_string(&checkpoint).unwrap();
5732 let checkpoint: StackingCheckpoint = serde_json::from_str(&encoded).unwrap();
5733 let mut other_density = log_density.clone();
5734 other_density[[0, 0]] += 0.25;
5735 assert!(matches!(
5736 resume_stacking_weights(other_density.view(), StackingConfig::default(), &checkpoint,),
5737 Err(StackingError::InvalidInput { .. })
5738 ));
5739 let resumed =
5740 resume_stacking_weights(log_density.view(), StackingConfig::default(), &checkpoint)
5741 .unwrap();
5742 let uninterrupted =
5743 solve_stacking_weights(log_density.view(), StackingConfig::default()).unwrap();
5744 for (resumed, uninterrupted) in resumed.weights.iter().zip(uninterrupted.weights.iter()) {
5745 assert!((resumed - uninterrupted).abs() <= 1.0e-10);
5746 }
5747 }
5748
5749 #[test]
5750 fn stacking_near_tied_boundary_uses_newton_not_millions_of_em_steps() {
5751 let log_density =
5752 Array2::from_shape_fn(
5753 (64, 2),
5754 |(_, candidate)| {
5755 if candidate == 0 { 0.0 } else { -1.0e-6 }
5756 },
5757 );
5758 let out = solve_stacking_weights(log_density.view(), StackingConfig::default()).unwrap();
5759 assert!(out.weights[0] >= 1.0 - StackingConfig::default().kkt_tol);
5760 assert!(out.iterations < 8, "iterations = {}", out.iterations);
5761 }
5762
5763 #[test]
5764 fn stacking_dead_candidate_column_gets_zero_weight() {
5765 let log_density = Array2::from_shape_vec(
5766 (3, 2),
5767 vec![
5768 -1.0,
5769 f64::NEG_INFINITY,
5770 -2.0,
5771 f64::NEG_INFINITY,
5772 -0.5,
5773 f64::NEG_INFINITY,
5774 ],
5775 )
5776 .unwrap();
5777 let out = solve_stacking_weights(log_density.view(), StackingConfig::default()).unwrap();
5778 assert_eq!(out.weights[1], 0.0);
5779 assert!((out.weights[0] - 1.0).abs() < 1e-12);
5780 }
5781
5782 #[test]
5783 fn stacking_rejects_invalid_and_unscorable_rows() {
5784 let log_density = Array2::from_shape_vec(
5785 (3, 2),
5786 vec![-1.0, -2.0, f64::NAN, f64::NEG_INFINITY, -2.0, -1.0],
5787 )
5788 .unwrap();
5789 assert!(matches!(
5790 solve_stacking_weights(log_density.view(), StackingConfig::default()),
5791 Err(StackingError::InvalidInput { .. })
5792 ));
5793 let unscorable = Array2::from_shape_vec(
5794 (2, 2),
5795 vec![-1.0, -2.0, f64::NEG_INFINITY, f64::NEG_INFINITY],
5796 )
5797 .unwrap();
5798 assert!(matches!(
5799 solve_stacking_weights(unscorable.view(), StackingConfig::default()),
5800 Err(StackingError::InvalidInput { .. })
5801 ));
5802 }
5803
5804 fn two_cluster_mixture_data() -> Array2<f64> {
5805 Array2::from_shape_vec(
5806 (12, 1),
5807 vec![
5808 -2.2, -2.0, -1.9, -2.1, -1.8, -2.05, 1.8, 2.0, 2.2, 1.9, 2.1, 2.05,
5809 ],
5810 )
5811 .unwrap()
5812 }
5813
5814 #[test]
5815 fn gaussian_mixture_monotonicity_resolves_composite_map_noise_2264() {
5816 let objective_scale = 1.0;
5817 let composite_resolution = f64::EPSILON.sqrt() * objective_scale;
5818 let uncertainty = gaussian_mixture_monotonicity_uncertainty(objective_scale, 0.0, 0.0);
5819 assert_eq!(uncertainty, composite_resolution);
5820
5821 let noise_scale_decrease = -0.5 * composite_resolution;
5822 assert!(noise_scale_decrease >= -uncertainty);
5823 let resolved_decrease = -2.0 * composite_resolution;
5824 assert!(resolved_decrease < -uncertainty);
5825
5826 let larger_reduction_bound = 2.0 * composite_resolution;
5827 assert_eq!(
5828 gaussian_mixture_monotonicity_uncertainty(objective_scale, larger_reduction_bound, 0.0),
5829 larger_reduction_bound,
5830 );
5831 }
5832
5833 #[test]
5834 fn gaussian_mixture_issue_scale_negative_step_is_within_computed_uncertainty_2264() {
5835 let objective_scale = 1.0;
5842 let recorded_step = -1.4e-13;
5843 let uncertainty = gaussian_mixture_monotonicity_uncertainty(objective_scale, 0.0, 0.0);
5844 let certificate = GaussianMixtureCertificate {
5845 mean_log_likelihood: -objective_scale,
5846 mean_log_likelihood_gain: recorded_step,
5847 monotonicity_uncertainty: uncertainty,
5848 objective_residual: recorded_step.abs() / objective_scale,
5849 objective_tolerance: f64::EPSILON.sqrt(),
5850 parameter_residual: 0.0,
5851 parameter_tolerance: f64::EPSILON.sqrt(),
5852 contraction_rate: None,
5853 projected_iterations_to_tolerance: None,
5854 };
5855
5856 assert_eq!(
5857 certificate.monotonicity_uncertainty,
5858 f64::EPSILON.sqrt() * objective_scale,
5859 "reported uncertainty must be the computed composite-map resolution"
5860 );
5861 assert!(
5862 certificate.mean_log_likelihood_gain >= -certificate.monotonicity_uncertainty,
5863 "the recorded noise-scale decrease must not be a monotonicity violation"
5864 );
5865 }
5866
5867 #[test]
5868 fn gaussian_mixture_below_roundoff_positive_gain_can_certify_2264() {
5869 let objective_scale = 1.0;
5874 let recorded_gain = 6.6e-15;
5875 let recorded_reduction_bound = 1.5e-14;
5876 let objective_tolerance = f64::EPSILON.sqrt();
5877 let parameter_tolerance = f64::EPSILON.sqrt();
5878 let uncertainty = gaussian_mixture_monotonicity_uncertainty(
5879 objective_scale,
5880 recorded_reduction_bound,
5881 0.0,
5882 );
5883 let certificate = GaussianMixtureCertificate {
5884 mean_log_likelihood: -objective_scale,
5885 mean_log_likelihood_gain: recorded_gain,
5886 monotonicity_uncertainty: uncertainty,
5887 objective_residual: recorded_gain / objective_scale,
5888 objective_tolerance,
5889 parameter_residual: 0.5 * parameter_tolerance,
5890 parameter_tolerance,
5891 contraction_rate: None,
5892 projected_iterations_to_tolerance: None,
5893 };
5894
5895 assert_eq!(
5896 certificate.monotonicity_uncertainty,
5897 (f64::EPSILON.sqrt() * objective_scale).max(recorded_reduction_bound),
5898 "reported uncertainty must come from the composite-map and reduction bounds"
5899 );
5900 assert!(certificate.mean_log_likelihood_gain >= -certificate.monotonicity_uncertainty);
5901 assert!(certificate.objective_residual <= certificate.objective_tolerance);
5902 assert!(certificate.parameter_residual <= certificate.parameter_tolerance);
5903 }
5904
5905 #[test]
5906 fn gaussian_mixture_certificate_quotients_duplicate_component_mass_exchange_2324() {
5907 let data = array![[-1.0], [0.0], [2.0]];
5913 let means = array![[0.0], [0.0]];
5914 let covariance = vec![array![[1.0]], array![[1.0]]];
5915 let weights = array![0.25, 0.75];
5916 let redistributed_weights = array![0.5, 0.5];
5917 let previous = mixture_e_step(data.view(), &weights, &means, &covariance).unwrap();
5918 let redistributed =
5919 mixture_e_step(data.view(), &redistributed_weights, &means, &covariance).unwrap();
5920 let residual = empirical_predictive_density_residual(
5921 &previous.row_log_likelihoods,
5922 &redistributed.row_log_likelihoods,
5923 )
5924 .unwrap();
5925 assert!(residual <= 4.0 * f64::EPSILON);
5926
5927 let shifted_means = array![[0.01], [0.0]];
5929 let shifted =
5930 mixture_e_step(data.view(), &weights, &shifted_means, &covariance).unwrap();
5931 let shifted_residual = empirical_predictive_density_residual(
5932 &previous.row_log_likelihoods,
5933 &shifted.row_log_likelihoods,
5934 )
5935 .unwrap();
5936 assert!(shifted_residual > f64::EPSILON.sqrt());
5937 }
5938
5939 #[test]
5940 fn gaussian_mixture_fit_certificate_describes_the_exact_returned_iterate() {
5941 let data = two_cluster_mixture_data();
5942 let config = GaussianMixtureConfig::default();
5943 let fit = fit_gaussian_mixture(data.view(), 2, config).unwrap();
5944 let certificate = fit.certificate();
5945 assert!(certificate.objective_residual <= certificate.objective_tolerance);
5946 assert!(certificate.parameter_residual <= certificate.parameter_tolerance);
5947
5948 let checkpoint = GaussianMixtureCheckpoint {
5949 weights: fit.weights.clone(),
5950 means: fit.means.clone(),
5951 covariances: fit.covariances.clone(),
5952 mean_log_likelihood: certificate.mean_log_likelihood,
5953 completed_iterations: fit.iterations,
5954 data_fingerprint: mixture_data_fingerprint(data.view()),
5955 covariance_floor: config.covariance_floor,
5956 };
5957 let current = mixture_e_step(
5958 data.view(),
5959 &checkpoint.weights,
5960 &checkpoint.means,
5961 &checkpoint.covariances,
5962 )
5963 .unwrap();
5964 let (weights, means, covariances) = mixture_m_step(
5965 data.view(),
5966 current.responsibilities.view(),
5967 config.covariance_floor,
5968 )
5969 .unwrap();
5970 let next = mixture_e_step(data.view(), &weights, &means, &covariances).unwrap();
5971 let residual = empirical_predictive_density_residual(
5972 ¤t.row_log_likelihoods,
5973 &next.row_log_likelihoods,
5974 )
5975 .unwrap();
5976 assert!(residual <= config.parameter_tol);
5977 assert_eq!(certificate.mean_log_likelihood, current.mean_log_likelihood);
5978 assert_eq!(
5979 certificate.mean_log_likelihood_gain,
5980 next.mean_log_likelihood - current.mean_log_likelihood
5981 );
5982 assert_eq!(
5983 certificate.monotonicity_uncertainty,
5984 gaussian_mixture_monotonicity_uncertainty(
5985 current
5986 .mean_log_likelihood
5987 .abs()
5988 .max(next.mean_log_likelihood.abs())
5989 .max(1.0),
5990 current.mean_log_likelihood_roundoff,
5991 next.mean_log_likelihood_roundoff,
5992 )
5993 );
5994 assert_eq!(certificate.parameter_residual, residual);
5995 assert!(
5996 (next.mean_log_likelihood - current.mean_log_likelihood).abs()
5997 / current
5998 .mean_log_likelihood
5999 .abs()
6000 .max(next.mean_log_likelihood.abs())
6001 .max(1.0)
6002 <= config.loglik_tol
6003 );
6004 }
6005
6006 #[test]
6007 fn gaussian_mixture_exhaustion_is_typed_and_resumable() {
6008 let data = two_cluster_mixture_data();
6009 let short = GaussianMixtureConfig {
6010 max_iter: 1,
6011 ..GaussianMixtureConfig::default()
6012 };
6013 let err = fit_gaussian_mixture(data.view(), 2, short).unwrap_err();
6014 let checkpoint = match err {
6015 GaussianMixtureError::DidNotConverge {
6016 certificate,
6017 checkpoint,
6018 ..
6019 } => {
6020 assert!(
6021 certificate.objective_residual > short.loglik_tol
6022 || certificate.parameter_residual > short.parameter_tol
6023 );
6024 assert_eq!(checkpoint.completed_iterations, 1);
6025 let at_checkpoint = mixture_e_step(
6026 data.view(),
6027 &checkpoint.weights,
6028 &checkpoint.means,
6029 &checkpoint.covariances,
6030 )
6031 .unwrap();
6032 assert_eq!(
6033 certificate.mean_log_likelihood, at_checkpoint.mean_log_likelihood,
6034 "exhaustion evidence and checkpoint must describe one iterate"
6035 );
6036 checkpoint
6037 }
6038 other => panic!("expected typed EM exhaustion, got {other}"),
6039 };
6040 let encoded = serde_json::to_string(&checkpoint).unwrap();
6041 let checkpoint: GaussianMixtureCheckpoint = serde_json::from_str(&encoded).unwrap();
6042 let mut other_data = data.clone();
6043 other_data[[0, 0]] += 0.01;
6044 assert!(matches!(
6045 resume_gaussian_mixture(
6046 other_data.view(),
6047 GaussianMixtureConfig::default(),
6048 checkpoint.clone(),
6049 ),
6050 Err(GaussianMixtureError::InvalidInput { .. })
6051 ));
6052 let resumed =
6053 resume_gaussian_mixture(data.view(), GaussianMixtureConfig::default(), checkpoint)
6054 .unwrap();
6055 let uninterrupted =
6056 fit_gaussian_mixture(data.view(), 2, GaussianMixtureConfig::default()).unwrap();
6057 for (resumed, uninterrupted) in resumed.weights.iter().zip(uninterrupted.weights.iter()) {
6058 assert!((resumed - uninterrupted).abs() <= 1.0e-10);
6059 }
6060
6061 assert!(resumed.bic().is_finite());
6062 }
6063
6064 #[test]
6065 fn gaussian_mixture_bic_is_finite_with_an_active_covariance_floor() {
6066 let per_cluster = 45usize;
6072 let mut data = Array2::<f64>::zeros((2 * per_cluster, 2));
6073 for sample in 0..per_cluster {
6074 let phase = std::f64::consts::TAU * sample as f64 / per_cluster as f64;
6075 data[[2 * sample, 0]] = -2.0;
6076 data[[2 * sample, 1]] = 0.08 * phase.sin();
6077 data[[2 * sample + 1, 0]] = 2.0 + 0.12 * phase.cos();
6078 data[[2 * sample + 1, 1]] = 0.08 * phase.sin();
6079 }
6080 let fit = fit_gaussian_mixture(data.view(), 2, GaussianMixtureConfig::default())
6081 .expect("the covariance floor defines a valid constrained mixture fit");
6082 let bic = fit.bic();
6083 assert!(bic.is_finite());
6084 assert_eq!(
6085 bic,
6086 -fit.loglik + 0.5 * fit.num_free_parameters() as f64 * (data.nrows() as f64).ln()
6087 );
6088 }
6089
6090 fn seven_clusters_on_a_circle_2262() -> Array2<f64> {
6091 let clusters = 7usize;
6092 let per_cluster = 32usize;
6093 let mut data = Array2::<f64>::zeros((clusters * per_cluster, 2));
6094 for cluster in 0..clusters {
6095 let angle = std::f64::consts::TAU * cluster as f64 / clusters as f64;
6096 let (sin_angle, cos_angle) = angle.sin_cos();
6097 for sample in 0..per_cluster {
6098 let phase = std::f64::consts::TAU * sample as f64 / per_cluster as f64;
6099 let local_radius = 0.035 * (1.0 + 0.3 * (3.0 * phase).cos());
6104 let radial_noise = local_radius * phase.cos();
6105 let tangent_noise = local_radius * phase.sin();
6106 let radius = 2.0 + radial_noise;
6107 let row = cluster * per_cluster + sample;
6108 data[[row, 0]] = 0.4 + radius * cos_angle - tangent_noise * sin_angle;
6109 data[[row, 1]] = -0.3 + radius * sin_angle + tangent_noise * cos_angle;
6110 }
6111 }
6112 data
6113 }
6114
6115 fn two_noisy_circles_for_union() -> Array2<f64> {
6116 let rows_per_circle = 96usize;
6117 let mut data = Array2::<f64>::zeros((2 * rows_per_circle, 2));
6118 for (circle, (center, radius)) in [([-4.0_f64, 0.3_f64], 1.2_f64), ([4.0, -0.2], 0.9)]
6119 .into_iter()
6120 .enumerate()
6121 {
6122 for sample in 0..rows_per_circle {
6123 let angle = std::f64::consts::TAU * sample as f64 / rows_per_circle as f64;
6124 let noisy_radius =
6125 radius + 0.045 * (3.0 * angle).cos() + 0.018 * (5.0 * angle).sin();
6126 let row = circle * rows_per_circle + sample;
6127 data[[row, 0]] = center[0] + noisy_radius * angle.cos();
6128 data[[row, 1]] = center[1] + noisy_radius * angle.sin();
6129 }
6130 }
6131 data
6132 }
6133
6134 #[test]
6135 fn circular_gaussian_density_avoids_extreme_scale_intermediate_overflow() {
6136 let noise_variance = f64::MAX / 2.0;
6137 let fit =
6138 CircularGaussianFit2d::from_parameters([0.0, 0.0], 1.1e154, noise_variance).unwrap();
6139 let center_log_density = fit.log_density(0.0, 0.0);
6142 let off_center_log_density = fit.log_density(1.7e154, 0.0);
6143 assert!(center_log_density.is_finite());
6144 assert!(off_center_log_density.is_finite());
6145 let expected_center = -std::f64::consts::TAU.ln()
6146 - noise_variance.ln()
6147 - 0.5 * (fit.radius() / noise_variance.sqrt()).powi(2);
6148 assert_eq!(center_log_density, expected_center);
6149 }
6150
6151 #[test]
6152 fn union_circles_use_the_shared_normalized_cartesian_density() {
6153 let data = two_noisy_circles_for_union();
6154 let config = GaussianMixtureConfig::default();
6155 let density_fit =
6156 fit_union_density(data.view(), UnionStructure::CircleCircle, config).unwrap();
6157 let union = fit_union_structure(data.view(), UnionStructure::CircleCircle, config).unwrap();
6158 assert_eq!(
6159 union.total_parameters,
6160 2 * CircularGaussianFit2d::NUM_FREE_PARAMETERS + 1
6161 );
6162 let component_weight_sum: f64 = union
6163 .components
6164 .iter()
6165 .map(|component| component.mixing_weight)
6166 .sum();
6167 assert!((component_weight_sum - 1.0).abs() <= 8.0 * f64::EPSILON);
6168
6169 let mut fitted_centers = Array2::<f64>::zeros((density_fit.components.len(), 2));
6170 for (index, component) in density_fit.components.iter().enumerate() {
6171 let UnionDensityModel::Circle(fit) = &component.model else {
6172 panic!("circle+circle union produced a non-circle density");
6173 };
6174 let center = fit.center();
6175 fitted_centers[[index, 0]] = center[0];
6176 fitted_centers[[index, 1]] = center[1];
6177 let at_center = fit.log_density(center[0], center[1]);
6178 let expected = -std::f64::consts::TAU.ln()
6179 - fit.noise_variance().ln()
6180 - 0.5 * (fit.radius() / fit.noise_variance().sqrt()).powi(2);
6181 assert!(at_center.is_finite());
6182 assert!((at_center - expected).abs() < 1.0e-12 * (1.0 + expected.abs()));
6183 }
6184
6185 let training_log_density = union_per_point_log_density(
6186 data.view(),
6187 data.view(),
6188 UnionStructure::CircleCircle,
6189 config,
6190 )
6191 .unwrap();
6192 let direct_log_likelihood = pairwise_sum(
6193 training_log_density
6194 .as_slice()
6195 .expect("owned score vector is contiguous"),
6196 );
6197 assert!(
6198 (union.log_likelihood - direct_log_likelihood).abs()
6199 <= 1.0e-12 * (1.0 + direct_log_likelihood.abs())
6200 );
6201 let expected_bic = -direct_log_likelihood
6202 + 0.5 * union.total_parameters as f64 * (data.nrows() as f64).ln();
6203 assert!((union.bic - expected_bic).abs() <= 1.0e-12 * (1.0 + expected_bic.abs()));
6204
6205 let held_out = union_per_point_log_density(
6206 data.view(),
6207 fitted_centers.view(),
6208 UnionStructure::CircleCircle,
6209 config,
6210 )
6211 .unwrap();
6212 assert!(held_out.iter().all(|value| value.is_finite()));
6213 }
6214
6215 fn circle_and_point_union_data() -> (Array2<f64>, Vec<Vec<usize>>) {
6216 let circle_rows = 32usize;
6217 let point_rows = 12usize;
6218 let mut data = Array2::<f64>::zeros((circle_rows + point_rows, 2));
6219 for row in 0..circle_rows {
6220 let angle = std::f64::consts::TAU * row as f64 / circle_rows as f64;
6221 let radius = 1.0 + 0.025 * (3.0 * angle).cos();
6222 data[[row, 0]] = -4.0 + radius * angle.cos();
6223 data[[row, 1]] = 0.2 + radius * angle.sin();
6224 }
6225 for offset in 0..point_rows {
6226 let phase = offset as f64;
6227 let row = circle_rows + offset;
6228 data[[row, 0]] = 4.0 + 0.055 * (1.7 * phase).cos() + 0.018 * (0.4 * phase).sin();
6229 data[[row, 1]] = -0.3 + 0.052 * (1.3 * phase).sin() - 0.015 * (0.9 * phase).cos();
6230 }
6231 (
6232 data,
6233 vec![
6234 (0..circle_rows).collect(),
6235 (circle_rows..circle_rows + point_rows).collect(),
6236 ],
6237 )
6238 }
6239
6240 #[test]
6241 fn heterogeneous_union_role_assignment_is_group_label_invariant() {
6242 let (data, groups) = circle_and_point_union_data();
6243 let config = GaussianMixtureConfig::default();
6244 let forward = fit_union_density_from_groups(
6245 data.view(),
6246 UnionStructure::CirclePointCluster,
6247 &groups,
6248 config,
6249 )
6250 .unwrap();
6251 let reversed_groups = vec![groups[1].clone(), groups[0].clone()];
6252 let reversed = fit_union_density_from_groups(
6253 data.view(),
6254 UnionStructure::CirclePointCluster,
6255 &reversed_groups,
6256 config,
6257 )
6258 .unwrap();
6259
6260 assert_eq!(forward.components[0].kind, UnionComponentKind::Circle);
6261 assert_eq!(forward.components[1].kind, UnionComponentKind::PointCluster);
6262 assert_eq!(
6263 reversed.components[0].kind,
6264 UnionComponentKind::PointCluster
6265 );
6266 assert_eq!(reversed.components[1].kind, UnionComponentKind::Circle);
6267 assert_eq!(forward.total_parameters, 4 + 3 + 1);
6268 assert_eq!(reversed.total_parameters, forward.total_parameters);
6269 assert!(
6270 (forward.log_likelihood - reversed.log_likelihood).abs()
6271 <= 1.0e-12 * (1.0 + forward.log_likelihood.abs())
6272 );
6273 assert!((forward.bic - reversed.bic).abs() <= 1.0e-12 * (1.0 + forward.bic.abs()));
6274 }
6275
6276 #[test]
6277 fn point_cluster_is_isotropic_and_line_remains_full_covariance() {
6278 let (mut data, mut groups) = circle_and_point_union_data();
6279 for row in 0..groups[0].len() {
6282 let coordinate = (row as f64 - 15.5) / 4.0;
6283 data[[row, 0]] = -4.0 + coordinate;
6284 data[[row, 1]] = 0.2 + 0.018 * coordinate + 0.006 * (1.9 * row as f64).sin();
6285 }
6286 let fit = fit_union_density_from_groups(
6287 data.view(),
6288 UnionStructure::LineCluster,
6289 &groups,
6290 GaussianMixtureConfig::default(),
6291 )
6292 .unwrap();
6293 assert_eq!(fit.components[0].kind, UnionComponentKind::Line);
6294 assert_eq!(fit.components[0].num_parameters, 5);
6295 assert_eq!(fit.components[1].kind, UnionComponentKind::PointCluster);
6296 assert_eq!(fit.components[1].num_parameters, 3);
6297 assert_eq!(fit.total_parameters, 5 + 3 + 1);
6298
6299 let UnionDensityModel::Gaussian(point) = &fit.components[1].model else {
6300 panic!("point cluster did not produce a Gaussian density");
6301 };
6302 assert_eq!(point.precision[[0, 1]], 0.0);
6303 assert_eq!(point.precision[[1, 0]], 0.0);
6304 assert_eq!(point.precision[[0, 0]], point.precision[[1, 1]]);
6305 let total_weight: f64 = fit
6306 .components
6307 .iter()
6308 .map(|component| component.mixing_weight)
6309 .sum();
6310 assert!((total_weight - 1.0).abs() <= 8.0 * f64::EPSILON);
6311
6312 groups.reverse();
6315 let reversed = fit_union_density_from_groups(
6316 data.view(),
6317 UnionStructure::LineCluster,
6318 &groups,
6319 GaussianMixtureConfig::default(),
6320 )
6321 .unwrap();
6322 assert_eq!(
6323 reversed.components[0].kind,
6324 UnionComponentKind::PointCluster
6325 );
6326 assert_eq!(reversed.components[1].kind, UnionComponentKind::Line);
6327 assert!((fit.bic - reversed.bic).abs() <= 1.0e-12 * (1.0 + fit.bic.abs()));
6328 }
6329
6330 #[test]
6331 fn isotropic_union_density_uses_the_same_fractional_mean_chart_as_its_mle() {
6332 let translated = ndarray::array![[1.0e16], [1.0e16 + 2.0], [1.0e16 + 2.0]];
6333 let fit = fit_isotropic_gaussian_component(translated.view(), 1.0e-12).unwrap();
6334 let variance = fit.precision[[0, 0]].recip();
6335 assert!((variance - 8.0 / 9.0).abs() <= 32.0 * f64::EPSILON);
6336
6337 let residuals = [-4.0 / 3.0, 2.0 / 3.0, 2.0 / 3.0];
6338 let expected_log_norm = -0.5 * ((2.0 * std::f64::consts::PI).ln() + variance.ln());
6339 for (row, residual) in residuals.into_iter().enumerate() {
6340 let expected = expected_log_norm - 0.5 * residual * residual / variance;
6341 let actual = fit.log_density(translated.row(row));
6342 assert!(
6343 (actual - expected).abs() <= 32.0 * f64::EPSILON * (1.0 + expected.abs()),
6344 "row {row}: density chart disagrees with fitted MLE residual: actual={actual}, expected={expected}"
6345 );
6346 }
6347
6348 let subnormal = f64::from_bits(1);
6349 let constant = ndarray::array![[subnormal], [subnormal], [subnormal]];
6350 let constant_fit = fit_isotropic_gaussian_component(constant.view(), 1.0).unwrap();
6351 assert_eq!(constant_fit.residual(constant.row(0)), vec![0.0]);
6352 assert_eq!(
6353 constant_fit.log_density(constant.row(0)),
6354 constant_fit.log_norm
6355 );
6356 }
6357
6358 #[test]
6359 fn union_ladder_fails_closed_when_one_declared_structure_fails() {
6360 let mut data = Array2::<f64>::zeros((8, 2));
6361 for row in 0..5 {
6362 let angle = std::f64::consts::TAU * row as f64 / 5.0;
6363 data[[row, 0]] = -5.0 + angle.cos();
6364 data[[row, 1]] = angle.sin();
6365 }
6366 data[[5, 0]] = 5.00;
6367 data[[5, 1]] = 0.00;
6368 data[[6, 0]] = 5.08;
6369 data[[6, 1]] = 0.02;
6370 data[[7, 0]] = 4.97;
6371 data[[7, 1]] = 0.07;
6372
6373 let error = fit_union_ladder(data.view(), GaussianMixtureConfig::default()).unwrap_err();
6374 assert!(error.contains("every declared structure must fit"));
6375 assert!(error.contains(UnionStructure::CircleCircle.as_str()));
6376 assert!(error.contains("needs at least 5 rows"));
6377 }
6378
6379 #[test]
6380 fn ring_of_clusters_fit_is_stationary_and_complexity_priced_2262() {
6381 let data = seven_clusters_on_a_circle_2262();
6382 let config = GaussianMixtureConfig::default();
6383 let fit = fit_ring_gaussian_mixture(data.view(), 7, config).unwrap();
6384 let certificate = fit.certificate();
6385 assert!(certificate.objective_residual <= certificate.objective_tolerance);
6386 assert!(certificate.parameter_residual <= certificate.parameter_tolerance);
6387 assert_eq!(fit.num_free_parameters(), 17);
6388 assert!((fit.center()[0] - 0.4).abs() < 0.05);
6389 assert!((fit.center()[1] + 0.3).abs() < 0.05);
6390 assert!((fit.radius() - 2.0).abs() < 0.05);
6391 assert!(fit.variance().is_finite() && fit.variance() > 0.0);
6392 assert!(
6393 fit.per_point_log_density(data.view())
6394 .unwrap()
6395 .iter()
6396 .all(|value| value.is_finite())
6397 );
6398 assert!(fit.bic().is_finite());
6399
6400 let free = fit_gaussian_mixture(data.view(), 7, config).unwrap();
6401 assert_eq!(free.num_free_parameters(), 41);
6402 assert!(fit.num_free_parameters() < free.num_free_parameters());
6403 }
6404
6405 #[test]
6406 fn ring_certificate_uses_identifiable_component_means() {
6407 let y = 0.91_f64.sqrt();
6413 let weights = Array1::from_vec(vec![0.2, 0.3, 0.5]);
6414 let previous = RingMixtureState {
6415 weights: weights.clone(),
6416 center: Array1::from_vec(vec![0.0, 0.0]),
6417 radius: 1.0,
6418 directions: Array2::from_shape_vec((3, 2), vec![0.3, y, 0.3, -y, 0.3, y]).unwrap(),
6419 variance: 0.25,
6420 mean_log_likelihood: -1.0,
6421 completed_iterations: 10,
6422 };
6423 let next = RingMixtureState {
6424 weights,
6425 center: Array1::from_vec(vec![0.6, 0.0]),
6426 radius: 1.0,
6427 directions: Array2::from_shape_vec((3, 2), vec![-0.3, y, -0.3, -y, -0.3, y]).unwrap(),
6428 variance: 0.25,
6429 mean_log_likelihood: -1.0,
6430 completed_iterations: 11,
6431 };
6432 assert!(relative_parameter_step(previous.center[0], next.center[0]) > 0.5);
6433 let data = array![[0.3, y], [0.3, -y], [1.0, 0.0]];
6434 let previous_e_step = ring_mixture_e_step(data.view(), &previous).unwrap();
6435 let next_e_step = ring_mixture_e_step(data.view(), &next).unwrap();
6436 let residual = empirical_predictive_density_residual(
6437 &previous_e_step.row_log_likelihoods,
6438 &next_e_step.row_log_likelihoods,
6439 )
6440 .unwrap();
6441 assert_eq!(residual, 0.0);
6442 }
6443
6444 #[test]
6445 fn ring_certificate_quotients_duplicate_component_mass_exchange_2324() {
6446 let previous = RingMixtureState {
6447 weights: array![0.2, 0.3, 0.5],
6448 center: array![0.0, 0.0],
6449 radius: 1.0,
6450 directions: array![[1.0, 0.0], [1.0, 0.0], [0.0, 1.0]],
6451 variance: 1.0,
6452 mean_log_likelihood: -1.0,
6453 completed_iterations: 10,
6454 };
6455 let next = RingMixtureState {
6456 weights: array![0.4, 0.1, 0.5],
6457 center: previous.center.clone(),
6458 radius: previous.radius,
6459 directions: previous.directions.clone(),
6460 variance: previous.variance,
6461 mean_log_likelihood: -1.0,
6462 completed_iterations: 11,
6463 };
6464 let data = array![[1.0, 0.0], [0.0, 1.0], [-1.0, 0.0]];
6465 let previous_e_step = ring_mixture_e_step(data.view(), &previous).unwrap();
6466 let next_e_step = ring_mixture_e_step(data.view(), &next).unwrap();
6467 let residual = empirical_predictive_density_residual(
6468 &previous_e_step.row_log_likelihoods,
6469 &next_e_step.row_log_likelihoods,
6470 )
6471 .unwrap();
6472 assert!(residual <= 4.0 * f64::EPSILON);
6473 }
6474
6475 #[test]
6476 fn stacked_mean_is_weighted_combination() {
6477 let weights = Array1::from_vec(vec![0.25, 0.75]);
6478 let means = vec![
6479 Array1::from_vec(vec![1.0, 2.0, 3.0]),
6480 Array1::from_vec(vec![5.0, 6.0, 7.0]),
6481 ];
6482 let out = stacked_predictive_mean(&weights, &means).unwrap();
6483 assert!((out[0] - (0.25 * 1.0 + 0.75 * 5.0)).abs() < 1e-12);
6484 assert!((out[2] - (0.25 * 3.0 + 0.75 * 7.0)).abs() < 1e-12);
6485 }
6486
6487 #[test]
6488 fn stacked_mean_rejects_shape_mismatch() {
6489 let weights = Array1::from_vec(vec![0.5, 0.5]);
6490 let means = vec![
6491 Array1::from_vec(vec![1.0, 2.0]),
6492 Array1::from_vec(vec![3.0]),
6493 ];
6494 assert!(stacked_predictive_mean(&weights, &means).is_err());
6495 }
6496
6497 fn hybrid_slot(
6512 linear_nle: f64,
6513 p_linear: usize,
6514 latent_dim: usize,
6515 p_curved: usize,
6516 theta: f64,
6517 curved_loglik_gain: f64,
6518 ) -> Vec<HybridAtomCandidate> {
6519 let param_price =
6520 0.5 * (p_curved as f64 - p_linear as f64) * (2.0 * std::f64::consts::PI).ln();
6521 let curved_nle = linear_nle - curved_loglik_gain + param_price;
6522 vec![
6523 HybridAtomCandidate::linear(linear_nle, p_linear),
6524 HybridAtomCandidate::curved(latent_dim, curved_nle, p_curved, Some(theta)),
6525 ]
6526 }
6527
6528 #[test]
6529 fn hybrid_dominance_floor_selects_linear_when_turning_is_zero() {
6530 let slot = hybrid_slot(100.0, 2, 1, 5, 0.0, 0.0);
6535 let choice = select_hybrid_atom(&slot).unwrap();
6536 assert!(choice.param.is_linear());
6537 assert_eq!(choice.param, HybridAtomParam::Linear);
6538 assert!(choice.curved_turning.unwrap() <= HYBRID_LINEAR_TURNING_FLOOR);
6540 }
6541
6542 #[test]
6543 fn hybrid_selects_curved_when_turning_pays_for_itself() {
6544 let slot = hybrid_slot(100.0, 2, 1, 5, 2.0 * std::f64::consts::PI, 30.0);
6548 let choice = select_hybrid_atom(&slot).unwrap();
6549 assert_eq!(choice.param, HybridAtomParam::Curved { latent_dim: 1 });
6550 assert!(choice.curved_evidence_margin > 0.0);
6552 }
6553
6554 #[test]
6555 fn hybrid_keeps_linear_when_curvature_doesnt_pay_its_price() {
6556 let slot = hybrid_slot(100.0, 2, 1, 5, 0.05, 0.1);
6560 let choice = select_hybrid_atom(&slot).unwrap();
6561 assert!(choice.param.is_linear());
6562 assert!(choice.curved_evidence_margin <= 0.0);
6563 }
6564
6565 #[test]
6566 fn hybrid_tie_breaks_to_the_cheaper_linear_atom() {
6567 let theta = 0.5; let nle = 42.0;
6572 let slot = vec![
6573 HybridAtomCandidate::linear(nle, 2),
6574 HybridAtomCandidate::curved(1, nle, 5, Some(theta)),
6575 ];
6576 let choice = select_hybrid_atom(&slot).unwrap();
6577 assert!(choice.param.is_linear());
6578 assert_eq!(choice.num_parameters, 2);
6579 }
6580
6581 #[test]
6582 fn hybrid_split_reduces_to_pure_linear_when_all_features_are_straight() {
6583 let slots: Vec<Vec<HybridAtomCandidate>> = (0..6)
6587 .map(|i| hybrid_slot(50.0 + i as f64, 2, 1, 5, 0.0, 0.0))
6588 .collect();
6589 let split = select_hybrid_split(&slots).unwrap();
6590 assert!(split.is_pure_linear());
6591 assert_eq!(split.curved_atom_count, 0);
6592 assert_eq!(split.linear_atom_count(), 6);
6593 let pure_linear: f64 = (0..6).map(|i| 50.0 + i as f64).sum();
6595 assert!((split.total_negative_log_evidence - pure_linear).abs() < 1e-12);
6596 }
6597
6598 #[test]
6599 fn hybrid_split_reduces_to_pure_curved_when_every_feature_curves() {
6600 let slots: Vec<Vec<HybridAtomCandidate>> = (0..5)
6603 .map(|i| hybrid_slot(80.0 + i as f64, 2, 1, 5, 2.0 * std::f64::consts::PI, 40.0))
6604 .collect();
6605 let split = select_hybrid_split(&slots).unwrap();
6606 assert!(split.is_pure_curved());
6607 assert_eq!(split.curved_atom_count, 5);
6608 assert_eq!(split.linear_atom_count(), 0);
6609 }
6610
6611 #[test]
6612 fn hybrid_split_on_mixed_dictionary_picks_curved_for_circles_linear_for_directions() {
6613 let mut slots: Vec<Vec<HybridAtomCandidate>> = Vec::new();
6623 let mut pure_linear_baseline = 0.0_f64;
6624 for i in 0..3 {
6627 let linear_nle = 120.0 + 3.0 * i as f64;
6628 pure_linear_baseline += linear_nle;
6629 slots.push(hybrid_slot(
6630 linear_nle,
6631 2,
6632 1,
6633 5,
6634 2.0 * std::f64::consts::PI,
6635 35.0,
6636 ));
6637 }
6638 for i in 0..4 {
6641 let linear_nle = 90.0 + 2.0 * i as f64;
6642 pure_linear_baseline += linear_nle;
6643 slots.push(hybrid_slot(linear_nle, 2, 1, 5, 0.0, 0.0));
6644 }
6645
6646 let split = select_hybrid_split(&slots).unwrap();
6647
6648 for (idx, choice) in split.atoms.iter().enumerate() {
6651 if idx < 3 {
6652 assert_eq!(
6653 choice.param,
6654 HybridAtomParam::Curved { latent_dim: 1 },
6655 "circle slot {idx} should select curved"
6656 );
6657 } else {
6658 assert!(
6659 choice.param.is_linear(),
6660 "direction slot {idx} should select linear"
6661 );
6662 }
6663 }
6664 assert_eq!(split.curved_atom_count, 3);
6665 assert_eq!(split.linear_atom_count(), 4);
6666
6667 assert!(
6673 split.total_negative_log_evidence <= pure_linear_baseline + 1e-9,
6674 "hybrid NLE {} must be <= summed linear-candidate NLE {}",
6675 split.total_negative_log_evidence,
6676 pure_linear_baseline
6677 );
6678 assert!(split.total_negative_log_evidence < pure_linear_baseline);
6680 }
6681
6682 #[test]
6683 fn hybrid_split_rejects_empty_slot() {
6684 let slots = vec![hybrid_slot(10.0, 2, 1, 5, 0.0, 0.0), Vec::new()];
6685 assert!(select_hybrid_split(&slots).is_err());
6686 }
6687
6688 fn cand(name: &str, score: f64, edf: f64, log_lik: f64) -> RemlCandidate {
6696 RemlCandidate {
6697 index: 0,
6698 name: name.to_string(),
6699 score,
6700 edf: Some(edf),
6701 log_lik: Some(log_lik),
6702 family: None,
6703 n_obs: None,
6704 }
6705 }
6706
6707 #[test]
6708 fn ranking_score_is_conditional_aic_when_loglik_and_edf_present() {
6709 let c = cand("m", 999.0, 6.748, -32.0866);
6711 let expected = -2.0 * -32.0866 + 2.0 * 6.748;
6712 assert!((c.ranking_score() - expected).abs() < 1e-9);
6713 }
6714
6715 #[test]
6716 fn ranking_score_falls_back_to_evidence_without_loglik() {
6717 let c = RemlCandidate {
6718 index: 0,
6719 name: "m".to_string(),
6720 score: 151.28,
6721 edf: Some(6.0),
6722 log_lik: None,
6723 family: None,
6724 n_obs: None,
6725 };
6726 assert_eq!(c.ranking_score(), 151.28);
6727 }
6728
6729 #[test]
6730 fn compare_models_rejects_pure_noise_smooth_despite_lower_evidence() {
6731 let small = cand("small", 180.526, 6.748, -32.0866);
6738 let big = cand("big", 177.404, 14.250, -32.1212);
6739
6740 assert!(big.score < small.score);
6742
6743 let cmp = compare_reml_fits(vec![small, big]).expect("compare");
6744 assert_eq!(
6745 cmp.winner, "small",
6746 "compare_models must Occam-penalise the pure-noise smooth and pick the smaller model"
6747 );
6748 let small_row = cmp
6751 .score_table
6752 .iter()
6753 .find(|r| r.name == "small")
6754 .expect("small row");
6755 let big_row = cmp
6756 .score_table
6757 .iter()
6758 .find(|r| r.name == "big")
6759 .expect("big row");
6760 assert!((small_row.reml_score - 180.526).abs() < 1e-9);
6761 assert!((big_row.reml_score - 177.404).abs() < 1e-9);
6762 }
6763
6764 #[test]
6765 fn ranking_bayes_factor_is_akaike_evidence_ratio_not_its_square() {
6766 let delta_aic = 27.68_f64;
6776 let winner = cand("winner", 100.0, 0.0, 0.0);
6777 let loser = cand("loser", 110.0, 0.0, -delta_aic / 2.0);
6778
6779 let cmp = compare_reml_fits(vec![winner, loser]).expect("compare");
6780 assert_eq!(cmp.winner, "winner");
6781
6782 let loser_row = cmp
6783 .ranking
6784 .iter()
6785 .find(|r| r.name == "loser")
6786 .expect("loser ranking row");
6787
6788 assert!((loser_row.delta - delta_aic).abs() < 1e-9);
6790
6791 let expected = (0.5 * delta_aic).exp();
6794 assert!(
6795 (loser_row.bayes_factor / expected - 1.0).abs() < 1e-9,
6796 "ranking bayes_factor {} should be exp(½ΔAIC)={}, not exp(ΔAIC)={}",
6797 loser_row.bayes_factor,
6798 expected,
6799 delta_aic.exp()
6800 );
6801 assert!(loser_row.bayes_factor < delta_aic.exp() * 0.5);
6803
6804 let loser_score_row = cmp
6809 .score_table
6810 .iter()
6811 .find(|r| r.name == "loser")
6812 .expect("loser score row");
6813 let expected_reml_bf = 10.0_f64.exp();
6814 assert!(
6815 (loser_score_row.bayes_factor_best_over_model / expected_reml_bf - 1.0).abs() < 1e-9,
6816 "raw-REML bayes_factor_best_over_model must stay exp(Δreml)=exp(10), got {}",
6817 loser_score_row.bayes_factor_best_over_model
6818 );
6819 }
6820
6821 #[test]
6822 fn compare_models_keeps_power_for_a_relevant_smooth() {
6823 let small = cand("small", 1025.067, 6.75, -368.985);
6829 let big = cand("big", 199.509, 14.25, -33.165);
6830 let cmp = compare_reml_fits(vec![small, big]).expect("compare");
6831 assert_eq!(
6832 cmp.winner, "big",
6833 "compare_models must retain power: the relevant smooth's model must win"
6834 );
6835 }
6836
6837 #[test]
6838 fn compare_models_rejects_mismatched_observation_counts() {
6839 let with_n = |name: &str, n: usize| RemlCandidate {
6843 index: 0,
6844 name: name.to_string(),
6845 score: 100.0,
6846 edf: Some(5.0),
6847 log_lik: Some(-40.0),
6848 family: Some("gaussian".to_string()),
6849 n_obs: Some(n),
6850 };
6851 let err = compare_reml_fits(vec![with_n("big", 500), with_n("small", 100)])
6852 .expect_err("cross-n comparison must be rejected");
6853 assert!(
6854 err.contains("number of observations") && err.contains("500") && err.contains("100"),
6855 "n-guard error should name the incomparable counts, got: {err}"
6856 );
6857
6858 compare_reml_fits(vec![with_n("a", 250), with_n("b", 250)])
6860 .expect("same-n comparison must succeed");
6861
6862 let without_n = RemlCandidate {
6865 index: 0,
6866 name: "legacy".to_string(),
6867 score: 90.0,
6868 edf: Some(4.0),
6869 log_lik: Some(-35.0),
6870 family: Some("gaussian".to_string()),
6871 n_obs: None,
6872 };
6873 compare_reml_fits(vec![with_n("counted", 500), without_n])
6874 .expect("an unconstrained (None) count must not trip the guard");
6875 }
6876}