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,
819 pub covariance_floor: f64,
823 pub kmeans_max_iter: usize,
825}
826
827impl Default for GaussianMixtureConfig {
828 fn default() -> Self {
829 Self {
830 max_iter: 1000,
831 loglik_tol: f64::EPSILON.sqrt(),
832 parameter_tol: f64::EPSILON.sqrt(),
833 covariance_floor: 1e-6,
834 kmeans_max_iter: 25,
835 }
836 }
837}
838
839#[derive(Debug, Clone, Copy, Serialize, Deserialize)]
842pub struct GaussianMixtureCertificate {
843 pub mean_log_likelihood: f64,
845 pub mean_log_likelihood_gain: f64,
847 pub monotonicity_uncertainty: f64,
851 pub objective_residual: f64,
852 pub objective_tolerance: f64,
853 pub parameter_residual: f64,
854 pub parameter_tolerance: f64,
855}
856
857#[derive(Debug, Clone, Serialize, Deserialize)]
859pub struct GaussianMixtureCheckpoint {
860 pub weights: Array1<f64>,
861 pub means: Array2<f64>,
862 pub covariances: Vec<Array2<f64>>,
863 pub mean_log_likelihood: f64,
864 pub completed_iterations: usize,
865 data_fingerprint: Fingerprint,
866 covariance_floor: f64,
867}
868
869#[derive(Debug, Clone)]
872pub enum GaussianMixtureError {
873 InvalidInput {
874 message: String,
875 },
876 NumericalFailure {
877 message: String,
878 checkpoint: Option<GaussianMixtureCheckpoint>,
879 },
880 MonotonicityViolation {
881 previous_mean_log_likelihood: f64,
882 next_mean_log_likelihood: f64,
883 numerical_uncertainty: f64,
884 checkpoint: GaussianMixtureCheckpoint,
885 },
886 DidNotConverge {
887 max_iterations: usize,
888 certificate: GaussianMixtureCertificate,
889 checkpoint: GaussianMixtureCheckpoint,
890 },
891}
892
893impl std::fmt::Display for GaussianMixtureError {
894 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
895 match self {
896 Self::InvalidInput { message } => write!(f, "invalid Gaussian mixture: {message}"),
897 Self::NumericalFailure {
898 message,
899 checkpoint,
900 } => write!(
901 f,
902 "Gaussian-mixture numerical failure: {message} (checkpoint iterations {})",
903 checkpoint
904 .as_ref()
905 .map_or(0, |value| value.completed_iterations)
906 ),
907 Self::MonotonicityViolation {
908 previous_mean_log_likelihood,
909 next_mean_log_likelihood,
910 numerical_uncertainty,
911 checkpoint,
912 } => write!(
913 f,
914 "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",
915 checkpoint.completed_iterations
916 ),
917 Self::DidNotConverge {
918 max_iterations,
919 certificate,
920 checkpoint,
921 } => write!(
922 f,
923 "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}; resume from the carried checkpoint, which is not comparable evidence",
924 checkpoint.completed_iterations,
925 certificate.mean_log_likelihood_gain,
926 certificate.monotonicity_uncertainty,
927 certificate.objective_residual,
928 certificate.objective_tolerance,
929 certificate.parameter_residual,
930 certificate.parameter_tolerance
931 ),
932 }
933 }
934}
935
936impl std::error::Error for GaussianMixtureError {}
937
938#[derive(Debug, Clone)]
940pub struct GaussianMixtureFit {
941 weights: Array1<f64>,
943 means: Array2<f64>,
945 covariances: Vec<Array2<f64>>,
947 k: usize,
949 d: usize,
951 n_obs: usize,
953 loglik: f64,
955 iterations: usize,
957 certificate: GaussianMixtureCertificate,
958}
959
960impl GaussianMixtureFit {
961 pub fn weights(&self) -> ArrayView1<'_, f64> {
962 self.weights.view()
963 }
964
965 pub fn means(&self) -> ArrayView2<'_, f64> {
966 self.means.view()
967 }
968
969 pub fn covariances(&self) -> &[Array2<f64>] {
970 &self.covariances
971 }
972
973 pub fn iterations(&self) -> usize {
974 self.iterations
975 }
976
977 pub fn certificate(&self) -> GaussianMixtureCertificate {
978 self.certificate
979 }
980
981 pub fn num_free_parameters(&self) -> usize {
986 let cov_per = self.d * (self.d + 1) / 2;
987 (self.k - 1) + self.k * self.d + self.k * cov_per
988 }
989
990 pub fn per_point_log_density(&self, data: ArrayView2<'_, f64>) -> Result<Array1<f64>, String> {
994 if data.ncols() != self.d {
995 return Err(format!(
996 "mixture log-density expects {} columns, got {}",
997 self.d,
998 data.ncols()
999 ));
1000 }
1001 let n = data.nrows();
1002 let mut comp = Vec::with_capacity(self.k);
1003 for j in 0..self.k {
1004 comp.push(GaussianComponentEval::factor(
1005 self.means.row(j),
1006 &self.covariances[j],
1007 )?);
1008 }
1009 let mut out = Array1::<f64>::zeros(n);
1010 let log_w: Vec<f64> = self.weights.iter().map(|w| w.ln()).collect();
1011 for i in 0..n {
1012 let row = data.row(i);
1013 let mut log_terms = vec![f64::NEG_INFINITY; self.k];
1014 let mut max_term = f64::NEG_INFINITY;
1015 for j in 0..self.k {
1016 let lt = log_w[j] + comp[j].log_density(row);
1017 log_terms[j] = lt;
1018 if lt > max_term {
1019 max_term = lt;
1020 }
1021 }
1022 out[i] = log_sum_exp(&log_terms, max_term);
1023 }
1024 Ok(out)
1025 }
1026
1027 pub fn bic(&self) -> f64 {
1031 -self.loglik + 0.5 * self.num_free_parameters() as f64 * (self.n_obs as f64).ln()
1032 }
1033}
1034
1035#[derive(Debug, Clone)]
1038struct GaussianComponentEval {
1039 residual_origin: Array1<f64>,
1040 residual_scale: Array1<f64>,
1041 residual_normalized_offset: Array1<f64>,
1042 precision: Array2<f64>,
1043 log_norm: f64,
1044 d: usize,
1045}
1046
1047impl GaussianComponentEval {
1048 fn factor(mean: ArrayView1<'_, f64>, cov: &Array2<f64>) -> Result<Self, String> {
1049 let d = mean.len();
1050 if mean.iter().any(|value| !value.is_finite()) {
1051 return Err("mixture component mean must be finite".to_string());
1052 }
1053 if cov.nrows() != d || cov.ncols() != d {
1054 return Err(format!(
1055 "mixture component covariance must be {d}x{d}, got {}x{}",
1056 cov.nrows(),
1057 cov.ncols()
1058 ));
1059 }
1060 let (evals, evecs) = cov
1061 .eigh(Side::Lower)
1062 .map_err(|e| format!("mixture component covariance eigendecomposition failed: {e}"))?;
1063 let mut log_det = 0.0_f64;
1064 let mut inv_evals = Array1::<f64>::zeros(d);
1065 for (idx, &ev) in evals.iter().enumerate() {
1066 if !ev.is_finite() || ev <= 0.0 {
1067 return Err(format!(
1068 "mixture component covariance is not SPD: eigenvalue {idx} is {ev:.3e}"
1069 ));
1070 }
1071 log_det += ev.ln();
1072 let inverse = ev.recip();
1073 if !inverse.is_finite() {
1074 return Err(format!(
1075 "mixture component precision is not representable: eigenvalue {idx} is {ev:.3e}"
1076 ));
1077 }
1078 inv_evals[idx] = inverse;
1079 }
1080 let mut precision = Array2::<f64>::zeros((d, d));
1082 for a in 0..d {
1083 for b in 0..d {
1084 let mut acc = 0.0_f64;
1085 for m in 0..d {
1086 acc += evecs[[a, m]] * inv_evals[m] * evecs[[b, m]];
1087 }
1088 precision[[a, b]] = acc;
1089 }
1090 }
1091 let log_norm = -0.5 * (d as f64 * (2.0 * std::f64::consts::PI).ln() + log_det);
1092 if precision.iter().any(|value| !value.is_finite()) || !log_norm.is_finite() {
1093 return Err(
1094 "mixture component factorization produced non-finite precision or log normalizer"
1095 .to_string(),
1096 );
1097 }
1098 Ok(Self {
1099 residual_origin: mean.to_owned(),
1100 residual_scale: Array1::zeros(d),
1101 residual_normalized_offset: Array1::zeros(d),
1102 precision,
1103 log_norm,
1104 d,
1105 })
1106 }
1107
1108 fn isotropic(charts: &[StableScalarMeanChart], variance: f64) -> Result<Self, String> {
1109 let d = charts.len();
1110 if d == 0 {
1111 return Err("isotropic Gaussian density requires positive dimension".to_string());
1112 }
1113 if !(variance.is_finite() && variance > 0.0) {
1114 return Err(format!(
1115 "isotropic Gaussian variance must be finite and positive, got {variance}"
1116 ));
1117 }
1118 let inverse_variance = variance.recip();
1119 if !inverse_variance.is_finite() {
1120 return Err(format!(
1121 "isotropic Gaussian precision is non-finite for variance {variance}"
1122 ));
1123 }
1124 let mut precision = Array2::<f64>::zeros((d, d));
1125 for axis in 0..d {
1126 precision[[axis, axis]] = inverse_variance;
1127 }
1128 let log_norm = -0.5 * d as f64 * ((2.0 * std::f64::consts::PI).ln() + variance.ln());
1129 if !log_norm.is_finite() {
1130 return Err("isotropic Gaussian log normalizer is non-finite".to_string());
1131 }
1132 Ok(Self {
1133 residual_origin: Array1::from_iter(charts.iter().map(|chart| chart.origin)),
1134 residual_scale: Array1::from_iter(charts.iter().map(|chart| chart.scale)),
1135 residual_normalized_offset: Array1::from_iter(
1136 charts.iter().map(|chart| chart.normalized_offset),
1137 ),
1138 precision,
1139 log_norm,
1140 d,
1141 })
1142 }
1143
1144 #[inline]
1145 fn log_density(&self, y: ArrayView1<'_, f64>) -> f64 {
1146 let residual = self.residual(y);
1147 let pv = self.precision_times_residual(&residual);
1148 let mut quad = 0.0_f64;
1149 for c in 0..self.d {
1150 quad += residual[c] * pv[c];
1151 }
1152 self.log_norm - 0.5 * quad
1153 }
1154
1155 #[inline]
1156 fn residual(&self, y: ArrayView1<'_, f64>) -> Vec<f64> {
1157 let mut residual = vec![0.0_f64; self.d];
1158 for axis in 0..self.d {
1159 residual[axis] = (-self.residual_normalized_offset[axis]).mul_add(
1160 self.residual_scale[axis],
1161 y[axis] - self.residual_origin[axis],
1162 );
1163 }
1164 residual
1165 }
1166
1167 #[inline]
1169 fn precision_times_residual(&self, residual: &[f64]) -> Vec<f64> {
1170 let mut out = vec![0.0_f64; self.d];
1171 for a in 0..self.d {
1172 let mut acc = 0.0_f64;
1173 for b in 0..self.d {
1174 acc += self.precision[[a, b]] * residual[b];
1175 }
1176 out[a] = acc;
1177 }
1178 out
1179 }
1180}
1181
1182#[inline]
1183fn log_sum_exp(terms: &[f64], max_term: f64) -> f64 {
1184 if !max_term.is_finite() {
1185 return f64::NEG_INFINITY;
1186 }
1187 let mut acc = 0.0_f64;
1188 for &t in terms {
1189 acc += (t - max_term).exp();
1190 }
1191 max_term + acc.ln()
1192}
1193
1194fn evidence_matrix_fingerprint(namespace: &str, values: ArrayView2<'_, f64>) -> Fingerprint {
1195 let mut hasher = Fingerprinter::new();
1196 hasher.write_str(namespace);
1197 hasher.write_usize(values.nrows());
1198 hasher.write_usize(values.ncols());
1199 for &value in values {
1202 hasher.write_f64(value);
1203 }
1204 hasher.finalize()
1205}
1206
1207fn mixture_data_fingerprint(data: ArrayView2<'_, f64>) -> Fingerprint {
1208 evidence_matrix_fingerprint("gaussian-mixture-em-v1", data)
1209}
1210
1211pub fn fit_gaussian_mixture(
1220 data: ArrayView2<'_, f64>,
1221 k: usize,
1222 config: GaussianMixtureConfig,
1223) -> Result<GaussianMixtureFit, GaussianMixtureError> {
1224 validate_gaussian_mixture_problem(data, k, config)?;
1225 let means = gam_terms::basis::select_centers_by_strategy(
1228 data,
1229 &gam_terms::basis::CenterStrategy::KMeans {
1230 num_centers: k,
1231 max_iter: config.kmeans_max_iter,
1232 },
1233 )
1234 .map_err(|error| GaussianMixtureError::NumericalFailure {
1235 message: format!("deterministic k-means seeding failed: {error}"),
1236 checkpoint: None,
1237 })?;
1238 if means.nrows() != k || means.ncols() != data.ncols() {
1239 return Err(GaussianMixtureError::NumericalFailure {
1240 message: format!(
1241 "seeding returned {}x{} centers, expected {k}x{}",
1242 means.nrows(),
1243 means.ncols(),
1244 data.ncols()
1245 ),
1246 checkpoint: None,
1247 });
1248 }
1249 let global_covariance =
1250 constrained_data_covariance(data, config.covariance_floor).map_err(|message| {
1251 GaussianMixtureError::NumericalFailure {
1252 message,
1253 checkpoint: None,
1254 }
1255 })?;
1256 let weights = Array1::<f64>::from_elem(k, 1.0 / k as f64);
1257 let covariances = vec![global_covariance; k];
1258 let initial_e_step =
1259 mixture_e_step(data, &weights, &means, &covariances).map_err(|message| {
1260 GaussianMixtureError::NumericalFailure {
1261 message,
1262 checkpoint: None,
1263 }
1264 })?;
1265 let data_fingerprint = mixture_data_fingerprint(data);
1266 let checkpoint = GaussianMixtureCheckpoint {
1267 weights,
1268 means,
1269 covariances,
1270 mean_log_likelihood: initial_e_step.mean_log_likelihood,
1271 completed_iterations: 0,
1272 data_fingerprint,
1273 covariance_floor: config.covariance_floor,
1274 };
1275 run_gaussian_mixture_em(data, config, checkpoint)
1276}
1277
1278pub fn resume_gaussian_mixture(
1280 data: ArrayView2<'_, f64>,
1281 config: GaussianMixtureConfig,
1282 checkpoint: GaussianMixtureCheckpoint,
1283) -> Result<GaussianMixtureFit, GaussianMixtureError> {
1284 let k = checkpoint.weights.len();
1285 validate_gaussian_mixture_problem(data, k, config)?;
1286 validate_gaussian_mixture_checkpoint(data, config.covariance_floor, &checkpoint)?;
1287 run_gaussian_mixture_em(data, config, checkpoint)
1288}
1289
1290fn validate_gaussian_mixture_problem(
1291 data: ArrayView2<'_, f64>,
1292 k: usize,
1293 config: GaussianMixtureConfig,
1294) -> Result<(), GaussianMixtureError> {
1295 let n = data.nrows();
1296 let d = data.ncols();
1297 if k == 0 {
1298 return Err(GaussianMixtureError::InvalidInput {
1299 message: "k must be positive".to_string(),
1300 });
1301 }
1302 if d == 0 {
1303 return Err(GaussianMixtureError::InvalidInput {
1304 message: "at least one data column is required".to_string(),
1305 });
1306 }
1307 if k > n {
1308 return Err(GaussianMixtureError::InvalidInput {
1309 message: format!("requested {k} components but data has {n} rows"),
1310 });
1311 }
1312 if data.iter().any(|value| !value.is_finite()) {
1313 return Err(GaussianMixtureError::InvalidInput {
1314 message: "data must be finite".to_string(),
1315 });
1316 }
1317 if config.max_iter == 0 || config.kmeans_max_iter == 0 {
1318 return Err(GaussianMixtureError::InvalidInput {
1319 message: "max_iter and kmeans_max_iter must be positive".to_string(),
1320 });
1321 }
1322 let numerical_floor = f64::EPSILON.sqrt();
1323 if !config.loglik_tol.is_finite()
1324 || config.loglik_tol < numerical_floor
1325 || !config.parameter_tol.is_finite()
1326 || config.parameter_tol < numerical_floor
1327 || !config.covariance_floor.is_finite()
1328 || config.covariance_floor <= 0.0
1329 {
1330 return Err(GaussianMixtureError::InvalidInput {
1331 message: format!(
1332 "loglik_tol and parameter_tol must be finite and >= {numerical_floor:.3e}, and covariance_floor must be finite and positive"
1333 ),
1334 });
1335 }
1336 Ok(())
1337}
1338
1339fn validate_gaussian_mixture_checkpoint(
1340 data: ArrayView2<'_, f64>,
1341 covariance_floor: f64,
1342 checkpoint: &GaussianMixtureCheckpoint,
1343) -> Result<(), GaussianMixtureError> {
1344 let d = data.ncols();
1345 let k = checkpoint.weights.len();
1346 let mass = checkpoint.weights.sum();
1347 if k == 0
1348 || checkpoint.data_fingerprint != mixture_data_fingerprint(data)
1349 || checkpoint.covariance_floor.to_bits() != covariance_floor.to_bits()
1350 || checkpoint.means.dim() != (k, d)
1351 || checkpoint.covariances.len() != k
1352 || checkpoint
1353 .covariances
1354 .iter()
1355 .any(|covariance| covariance.dim() != (d, d))
1356 || checkpoint
1357 .weights
1358 .iter()
1359 .chain(checkpoint.means.iter())
1360 .chain(checkpoint.covariances.iter().flat_map(|value| value.iter()))
1361 .any(|value| !value.is_finite())
1362 || checkpoint.weights.iter().any(|value| *value <= 0.0)
1363 || !mass.is_finite()
1364 || (mass - 1.0).abs() > f64::EPSILON.sqrt()
1365 || !checkpoint.mean_log_likelihood.is_finite()
1366 {
1367 return Err(GaussianMixtureError::InvalidInput {
1368 message: "checkpoint problem identity, dimensions, interior parameters, likelihood, or simplex mass are invalid".to_string(),
1369 });
1370 }
1371 Ok(())
1372}
1373
1374fn run_gaussian_mixture_em(
1375 data: ArrayView2<'_, f64>,
1376 config: GaussianMixtureConfig,
1377 mut checkpoint: GaussianMixtureCheckpoint,
1378) -> Result<GaussianMixtureFit, GaussianMixtureError> {
1379 validate_gaussian_mixture_checkpoint(data, config.covariance_floor, &checkpoint)?;
1380 let k = checkpoint.weights.len();
1381 let d = data.ncols();
1382 let data_fingerprint = mixture_data_fingerprint(data);
1383
1384 for additional_updates in 0..=config.max_iter {
1391 let current = mixture_e_step(
1392 data,
1393 &checkpoint.weights,
1394 &checkpoint.means,
1395 &checkpoint.covariances,
1396 )
1397 .map_err(|message| GaussianMixtureError::NumericalFailure {
1398 message,
1399 checkpoint: Some(checkpoint.clone()),
1400 })?;
1401 if (checkpoint.mean_log_likelihood - current.mean_log_likelihood).abs()
1402 > current.mean_log_likelihood_roundoff
1403 {
1404 return Err(GaussianMixtureError::InvalidInput {
1405 message: format!(
1406 "checkpoint mean log likelihood {:.12e} disagrees with its parameters ({:.12e} +/- {:.3e})",
1407 checkpoint.mean_log_likelihood,
1408 current.mean_log_likelihood,
1409 current.mean_log_likelihood_roundoff
1410 ),
1411 });
1412 }
1413 checkpoint.mean_log_likelihood = current.mean_log_likelihood;
1414
1415 let (next_weights, next_means, next_covariances) = mixture_m_step(
1416 data,
1417 current.responsibilities.view(),
1418 config.covariance_floor,
1419 )
1420 .map_err(|message| GaussianMixtureError::NumericalFailure {
1421 message,
1422 checkpoint: Some(checkpoint.clone()),
1423 })?;
1424 let next = mixture_e_step(data, &next_weights, &next_means, &next_covariances).map_err(
1425 |message| GaussianMixtureError::NumericalFailure {
1426 message,
1427 checkpoint: Some(checkpoint.clone()),
1428 },
1429 )?;
1430 let objective_scale = current
1431 .mean_log_likelihood
1432 .abs()
1433 .max(next.mean_log_likelihood.abs())
1434 .max(1.0);
1435 let objective_step = next.mean_log_likelihood - current.mean_log_likelihood;
1436 let objective_residual = objective_step.abs() / objective_scale;
1437 let parameter_residual = mixture_parameter_residual(
1438 &checkpoint.weights,
1439 &checkpoint.means,
1440 &checkpoint.covariances,
1441 &next_weights,
1442 &next_means,
1443 &next_covariances,
1444 );
1445 let monotonicity_uncertainty = gaussian_mixture_monotonicity_uncertainty(
1446 objective_scale,
1447 current.mean_log_likelihood_roundoff,
1448 next.mean_log_likelihood_roundoff,
1449 );
1450 let certificate = GaussianMixtureCertificate {
1451 mean_log_likelihood: current.mean_log_likelihood,
1452 mean_log_likelihood_gain: objective_step,
1453 monotonicity_uncertainty,
1454 objective_residual,
1455 objective_tolerance: config.loglik_tol,
1456 parameter_residual,
1457 parameter_tolerance: config.parameter_tol,
1458 };
1459 if objective_step < -monotonicity_uncertainty {
1460 return Err(GaussianMixtureError::MonotonicityViolation {
1461 previous_mean_log_likelihood: current.mean_log_likelihood,
1462 next_mean_log_likelihood: next.mean_log_likelihood,
1463 numerical_uncertainty: monotonicity_uncertainty,
1464 checkpoint,
1465 });
1466 }
1467 if objective_residual <= config.loglik_tol && parameter_residual <= config.parameter_tol {
1468 let loglik = current.mean_log_likelihood * data.nrows() as f64;
1469 if !loglik.is_finite() {
1470 return Err(GaussianMixtureError::NumericalFailure {
1471 message: "certified mean log likelihood overflows as a total likelihood"
1472 .to_string(),
1473 checkpoint: Some(checkpoint),
1474 });
1475 }
1476 return Ok(GaussianMixtureFit {
1477 weights: checkpoint.weights,
1478 means: checkpoint.means,
1479 covariances: checkpoint.covariances,
1480 k,
1481 d,
1482 n_obs: data.nrows(),
1483 loglik,
1484 iterations: checkpoint.completed_iterations,
1485 certificate,
1486 });
1487 }
1488 if additional_updates == config.max_iter {
1489 return Err(GaussianMixtureError::DidNotConverge {
1490 max_iterations: config.max_iter,
1491 certificate,
1492 checkpoint,
1493 });
1494 }
1495 checkpoint = GaussianMixtureCheckpoint {
1496 weights: next_weights,
1497 means: next_means,
1498 covariances: next_covariances,
1499 mean_log_likelihood: next.mean_log_likelihood,
1500 completed_iterations: checkpoint.completed_iterations + 1,
1501 data_fingerprint,
1502 covariance_floor: config.covariance_floor,
1503 };
1504 }
1505 Err(GaussianMixtureError::NumericalFailure {
1506 message: format!(
1507 "EM refinement exhausted its inclusive update budget ({}) without producing a \
1508 terminal verdict",
1509 config.max_iter
1510 ),
1511 checkpoint: Some(checkpoint),
1512 })
1513}
1514
1515struct GaussianMixtureEStep {
1516 responsibilities: Array2<f64>,
1517 mean_log_likelihood: f64,
1518 mean_log_likelihood_roundoff: f64,
1519}
1520
1521fn gaussian_mixture_monotonicity_uncertainty(
1533 objective_scale: f64,
1534 current_reduction_roundoff: f64,
1535 next_reduction_roundoff: f64,
1536) -> f64 {
1537 let reduction_roundoff = current_reduction_roundoff + next_reduction_roundoff;
1538 let composite_map_resolution = f64::EPSILON.sqrt() * objective_scale;
1539 reduction_roundoff.max(composite_map_resolution)
1540}
1541
1542fn pairwise_sum_max_depth(term_count: usize) -> usize {
1543 if term_count <= 1 {
1544 return 0;
1545 }
1546 let within_block = term_count.min(BASE_CHUNK) - 1;
1547 let blocks = term_count.div_ceil(BASE_CHUNK);
1548 let tree_levels = if blocks <= 1 {
1549 0
1550 } else {
1551 (usize::BITS - (blocks - 1).leading_zeros()) as usize
1552 };
1553 within_block.saturating_add(tree_levels)
1554}
1555
1556fn pairwise_mean_with_roundoff(mut values: Vec<f64>) -> Result<(f64, f64), String> {
1557 if values.is_empty() || values.iter().any(|value| !value.is_finite()) {
1558 return Err("mean log-likelihood terms must be nonempty and finite".to_string());
1559 }
1560 let sum = pairwise_sum(&values);
1561 for value in &mut values {
1562 *value = value.abs();
1563 }
1564 let magnitude_sum = pairwise_sum(&values);
1565 let unit_roundoff = 0.5 * f64::EPSILON;
1566 let accumulated = pairwise_sum_max_depth(values.len()) as f64 * unit_roundoff;
1567 let addition_bound = if accumulated < 1.0 {
1568 accumulated / (1.0 - accumulated) * magnitude_sum
1569 } else {
1570 f64::INFINITY
1571 };
1572 let count = values.len() as f64;
1573 let mean = sum / count;
1574 let roundoff = addition_bound / count + unit_roundoff * mean.abs();
1578 if !(mean.is_finite() && roundoff.is_finite()) {
1579 return Err("mean mixture log likelihood or its rounding bound is non-finite".to_string());
1580 }
1581 Ok((mean, roundoff))
1582}
1583
1584fn mixture_e_step(
1585 data: ArrayView2<'_, f64>,
1586 weights: &Array1<f64>,
1587 means: &Array2<f64>,
1588 covariances: &[Array2<f64>],
1589) -> Result<GaussianMixtureEStep, String> {
1590 let n = data.nrows();
1591 let k = weights.len();
1592 if weights
1593 .iter()
1594 .any(|weight| !weight.is_finite() || *weight <= 0.0)
1595 {
1596 return Err("mixture E-step requires strictly positive finite weights".to_string());
1597 }
1598 let mut components = Vec::with_capacity(k);
1599 for component in 0..k {
1600 components.push(GaussianComponentEval::factor(
1601 means.row(component),
1602 &covariances[component],
1603 )?);
1604 }
1605 let log_weights: Vec<f64> = weights.iter().map(|weight| weight.ln()).collect();
1606 let mut responsibilities = Array2::<f64>::zeros((n, k));
1607 let mut row_log_likelihoods = Vec::with_capacity(n);
1608 for row in 0..n {
1609 let observation = data.row(row);
1610 let mut log_terms = vec![f64::NEG_INFINITY; k];
1611 let mut max_term = f64::NEG_INFINITY;
1612 for component in 0..k {
1613 let term = log_weights[component] + components[component].log_density(observation);
1614 log_terms[component] = term;
1615 max_term = max_term.max(term);
1616 }
1617 let log_mixture = log_sum_exp(&log_terms, max_term);
1618 if !log_mixture.is_finite() {
1619 return Err(format!(
1620 "mixture density is non-finite at training row {row}"
1621 ));
1622 }
1623 row_log_likelihoods.push(log_mixture);
1624 for component in 0..k {
1625 responsibilities[[row, component]] = (log_terms[component] - log_mixture).exp();
1626 }
1627 }
1628 let (mean_log_likelihood, mean_log_likelihood_roundoff) =
1629 pairwise_mean_with_roundoff(row_log_likelihoods)?;
1630 Ok(GaussianMixtureEStep {
1631 responsibilities,
1632 mean_log_likelihood,
1633 mean_log_likelihood_roundoff,
1634 })
1635}
1636
1637fn mixture_m_step(
1638 data: ArrayView2<'_, f64>,
1639 responsibilities: ArrayView2<'_, f64>,
1640 covariance_floor: f64,
1641) -> Result<(Array1<f64>, Array2<f64>, Vec<Array2<f64>>), String> {
1642 let n = data.nrows();
1643 let d = data.ncols();
1644 let k = responsibilities.ncols();
1645 let mut component_mass = Array1::<f64>::zeros(k);
1646 for component in 0..k {
1647 component_mass[component] = responsibilities.column(component).sum();
1648 }
1649 if component_mass
1650 .iter()
1651 .any(|mass| !mass.is_finite() || *mass <= 0.0)
1652 {
1653 return Err(
1654 "M-step reached a zero-mass component; the requested mixture order has no interior fitted density"
1655 .to_string(),
1656 );
1657 }
1658 let mut weights = component_mass.mapv(|mass| mass / n as f64);
1659 let total_weight = weights.sum();
1660 if !(total_weight.is_finite() && total_weight > 0.0) {
1661 return Err("M-step produced invalid mixture-weight mass".to_string());
1662 }
1663 weights.mapv_inplace(|weight| weight / total_weight);
1664 let mut means = Array2::<f64>::zeros((k, d));
1665 let mut covariances = Vec::with_capacity(k);
1666 for component in 0..k {
1667 let mass = component_mass[component];
1668 let mut mean = Array1::<f64>::zeros(d);
1669 for row in 0..n {
1670 let responsibility = responsibilities[[row, component]];
1671 for col in 0..d {
1672 mean[col] += responsibility * data[[row, col]];
1673 }
1674 }
1675 mean.mapv_inplace(|value| value / mass);
1676 means.row_mut(component).assign(&mean);
1677 let mut covariance = Array2::<f64>::zeros((d, d));
1678 for row in 0..n {
1679 let responsibility = responsibilities[[row, component]];
1680 for left in 0..d {
1681 let left_residual = data[[row, left]] - mean[left];
1682 for right in 0..d {
1683 covariance[[left, right]] +=
1684 responsibility * left_residual * (data[[row, right]] - mean[right]);
1685 }
1686 }
1687 }
1688 covariance.mapv_inplace(|value| value / mass);
1689 covariances.push(constrain_covariance(covariance, covariance_floor)?);
1690 }
1691 Ok((weights, means, covariances))
1692}
1693
1694fn relative_parameter_step(previous: f64, next: f64) -> f64 {
1695 (next - previous).abs() / previous.abs().max(next.abs()).max(1.0)
1696}
1697
1698fn labeled_gaussian_component_measure_residual(
1713 previous_weights: &Array1<f64>,
1714 previous_means: &Array2<f64>,
1715 previous_covariance: impl Fn(usize, usize, usize) -> f64,
1716 next_weights: &Array1<f64>,
1717 next_means: &Array2<f64>,
1718 next_covariance: impl Fn(usize, usize, usize) -> f64,
1719) -> f64 {
1720 let k = previous_weights.len();
1721 let d = previous_means.ncols();
1722 let mut residual = 0.0_f64;
1723 for component in 0..k {
1724 let previous_weight = previous_weights[component];
1725 let next_weight = next_weights[component];
1726 residual = residual.max(relative_parameter_step(previous_weight, next_weight));
1727 for left in 0..d {
1728 let previous_first = previous_weight * previous_means[[component, left]];
1729 let next_first = next_weight * next_means[[component, left]];
1730 residual = residual.max(relative_parameter_step(previous_first, next_first));
1731 for right in 0..d {
1732 let previous_second = previous_weight
1733 * (previous_covariance(component, left, right)
1734 + previous_means[[component, left]] * previous_means[[component, right]]);
1735 let next_second = next_weight
1736 * (next_covariance(component, left, right)
1737 + next_means[[component, left]] * next_means[[component, right]]);
1738 residual = residual.max(relative_parameter_step(previous_second, next_second));
1739 }
1740 }
1741 }
1742 residual
1743}
1744
1745fn mixture_parameter_residual(
1746 previous_weights: &Array1<f64>,
1747 previous_means: &Array2<f64>,
1748 previous_covariances: &[Array2<f64>],
1749 next_weights: &Array1<f64>,
1750 next_means: &Array2<f64>,
1751 next_covariances: &[Array2<f64>],
1752) -> f64 {
1753 labeled_gaussian_component_measure_residual(
1754 previous_weights,
1755 previous_means,
1756 |component, left, right| previous_covariances[component][[left, right]],
1757 next_weights,
1758 next_means,
1759 |component, left, right| next_covariances[component][[left, right]],
1760 )
1761}
1762
1763fn constrain_covariance(covariance: Array2<f64>, floor: f64) -> Result<Array2<f64>, String> {
1764 let (eigenvalues, eigenvectors) = covariance
1765 .eigh(Side::Lower)
1766 .map_err(|error| format!("covariance eigendecomposition failed: {error}"))?;
1767 let d = covariance.nrows();
1768 let mut constrained = Array2::<f64>::zeros((d, d));
1769 for row in 0..d {
1770 for col in 0..d {
1771 let mut value = 0.0_f64;
1772 for index in 0..d {
1773 value += eigenvectors[[row, index]]
1774 * eigenvalues[index].max(floor)
1775 * eigenvectors[[col, index]];
1776 }
1777 constrained[[row, col]] = value;
1778 }
1779 }
1780 if constrained.iter().any(|value| !value.is_finite()) {
1781 return Err("constrained covariance became non-finite".to_string());
1782 }
1783 Ok(constrained)
1784}
1785
1786fn constrained_data_covariance(
1788 data: ArrayView2<'_, f64>,
1789 floor: f64,
1790) -> Result<Array2<f64>, String> {
1791 let n = data.nrows();
1792 let d = data.ncols();
1793 let mut mean = Array1::<f64>::zeros(d);
1794 for i in 0..n {
1795 for c in 0..d {
1796 mean[c] += data[[i, c]];
1797 }
1798 }
1799 mean.mapv_inplace(|v| v / n.max(1) as f64);
1800 let mut cov = Array2::<f64>::zeros((d, d));
1801 for i in 0..n {
1802 for a in 0..d {
1803 let da = data[[i, a]] - mean[a];
1804 for b in 0..d {
1805 cov[[a, b]] += da * (data[[i, b]] - mean[b]);
1806 }
1807 }
1808 }
1809 let inv = 1.0 / n as f64;
1810 cov.mapv_inplace(|v| v * inv);
1811 constrain_covariance(cov, floor)
1812}
1813
1814#[derive(Debug, Clone)]
1832pub struct RingGaussianMixtureFit {
1833 weights: Array1<f64>,
1834 center: Array1<f64>,
1835 radius: f64,
1836 directions: Array2<f64>,
1837 variance: f64,
1838 k: usize,
1839 n_obs: usize,
1840 loglik: f64,
1841 iterations: usize,
1842 certificate: GaussianMixtureCertificate,
1843}
1844
1845impl RingGaussianMixtureFit {
1846 pub fn weights(&self) -> ArrayView1<'_, f64> {
1847 self.weights.view()
1848 }
1849
1850 pub fn center(&self) -> ArrayView1<'_, f64> {
1851 self.center.view()
1852 }
1853
1854 pub fn radius(&self) -> f64 {
1855 self.radius
1856 }
1857
1858 pub fn directions(&self) -> ArrayView2<'_, f64> {
1859 self.directions.view()
1860 }
1861
1862 pub fn variance(&self) -> f64 {
1863 self.variance
1864 }
1865
1866 pub fn iterations(&self) -> usize {
1867 self.iterations
1868 }
1869
1870 pub fn certificate(&self) -> GaussianMixtureCertificate {
1871 self.certificate
1872 }
1873
1874 pub fn num_free_parameters(&self) -> usize {
1877 2 * self.k + 3
1878 }
1879
1880 pub fn per_point_log_density(&self, data: ArrayView2<'_, f64>) -> Result<Array1<f64>, String> {
1881 if data.ncols() != 2 {
1882 return Err(format!(
1883 "ring-of-clusters density expects two columns, got {}",
1884 data.ncols()
1885 ));
1886 }
1887 ring_mixture_log_density(
1888 data,
1889 &self.weights,
1890 &self.center,
1891 self.radius,
1892 &self.directions,
1893 self.variance,
1894 )
1895 }
1896
1897 pub fn bic(&self) -> f64 {
1900 -self.loglik + 0.5 * self.num_free_parameters() as f64 * (self.n_obs as f64).ln()
1901 }
1902}
1903
1904#[derive(Debug, Clone)]
1905struct RingMixtureState {
1906 weights: Array1<f64>,
1907 center: Array1<f64>,
1908 radius: f64,
1909 directions: Array2<f64>,
1910 variance: f64,
1911 mean_log_likelihood: f64,
1912 completed_iterations: usize,
1913}
1914
1915fn ring_component_means(
1916 center: &Array1<f64>,
1917 radius: f64,
1918 directions: &Array2<f64>,
1919) -> Array2<f64> {
1920 let mut means = Array2::<f64>::zeros((directions.nrows(), 2));
1921 for component in 0..directions.nrows() {
1922 means[[component, 0]] = center[0] + radius * directions[[component, 0]];
1923 means[[component, 1]] = center[1] + radius * directions[[component, 1]];
1924 }
1925 means
1926}
1927
1928fn ring_mixture_log_terms(
1929 data: ArrayView2<'_, f64>,
1930 weights: &Array1<f64>,
1931 center: &Array1<f64>,
1932 radius: f64,
1933 directions: &Array2<f64>,
1934 variance: f64,
1935) -> Result<(Array2<f64>, Vec<f64>), String> {
1936 if data.ncols() != 2
1937 || center.len() != 2
1938 || directions.ncols() != 2
1939 || directions.nrows() != weights.len()
1940 || weights
1941 .iter()
1942 .any(|weight| !weight.is_finite() || *weight <= 0.0)
1943 || !(radius.is_finite() && radius > 0.0)
1944 || !(variance.is_finite() && variance > 0.0)
1945 {
1946 return Err("invalid ring-of-clusters parameter state".to_string());
1947 }
1948 let means = ring_component_means(center, radius, directions);
1949 let log_normalizer = -(std::f64::consts::TAU).ln() - variance.ln();
1950 let mut terms = Array2::<f64>::zeros((data.nrows(), weights.len()));
1951 let mut row_log_likelihoods = Vec::with_capacity(data.nrows());
1952 for row in 0..data.nrows() {
1953 let mut max_term = f64::NEG_INFINITY;
1954 for component in 0..weights.len() {
1955 let dx = data[[row, 0]] - means[[component, 0]];
1956 let dy = data[[row, 1]] - means[[component, 1]];
1957 let term =
1958 weights[component].ln() + log_normalizer - 0.5 * (dx * dx + dy * dy) / variance;
1959 terms[[row, component]] = term;
1960 max_term = max_term.max(term);
1961 }
1962 let values = terms.row(row).to_vec();
1963 let log_likelihood = log_sum_exp(&values, max_term);
1964 if !log_likelihood.is_finite() {
1965 return Err(format!(
1966 "ring-of-clusters density is non-finite at training row {row}"
1967 ));
1968 }
1969 row_log_likelihoods.push(log_likelihood);
1970 }
1971 Ok((terms, row_log_likelihoods))
1972}
1973
1974fn ring_mixture_e_step(
1975 data: ArrayView2<'_, f64>,
1976 state: &RingMixtureState,
1977) -> Result<(Array2<f64>, f64, f64), String> {
1978 let (terms, row_log_likelihoods) = ring_mixture_log_terms(
1979 data,
1980 &state.weights,
1981 &state.center,
1982 state.radius,
1983 &state.directions,
1984 state.variance,
1985 )?;
1986 let mut responsibilities = Array2::<f64>::zeros(terms.raw_dim());
1987 for row in 0..terms.nrows() {
1988 for component in 0..terms.ncols() {
1989 responsibilities[[row, component]] =
1990 (terms[[row, component]] - row_log_likelihoods[row]).exp();
1991 }
1992 }
1993 let (mean, roundoff) = pairwise_mean_with_roundoff(row_log_likelihoods)?;
1994 Ok((responsibilities, mean, roundoff))
1995}
1996
1997fn ring_mixture_log_density(
1998 data: ArrayView2<'_, f64>,
1999 weights: &Array1<f64>,
2000 center: &Array1<f64>,
2001 radius: f64,
2002 directions: &Array2<f64>,
2003 variance: f64,
2004) -> Result<Array1<f64>, String> {
2005 let (_, row_log_likelihoods) =
2006 ring_mixture_log_terms(data, weights, center, radius, directions, variance)?;
2007 Ok(Array1::from_vec(row_log_likelihoods))
2008}
2009
2010fn ring_identifiable_parameter_residual(
2021 previous: &RingMixtureState,
2022 next: &RingMixtureState,
2023) -> f64 {
2024 let previous_means =
2025 ring_component_means(&previous.center, previous.radius, &previous.directions);
2026 let next_means = ring_component_means(&next.center, next.radius, &next.directions);
2027 let noise_scale = previous
2028 .variance
2029 .sqrt()
2030 .max(next.variance.sqrt())
2031 .max(f64::MIN_POSITIVE);
2032 let weight_residual = previous
2033 .weights
2034 .iter()
2035 .zip(next.weights.iter())
2036 .map(|(&left, &right)| (right - left).abs())
2037 .fold(0.0, f64::max);
2038 let mean_residual = previous_means
2039 .rows()
2040 .into_iter()
2041 .zip(next_means.rows())
2042 .zip(previous.weights.iter().zip(next.weights.iter()))
2043 .map(|((left, right), (&previous_weight, &next_weight))| {
2044 previous_weight.max(next_weight) * (right[0] - left[0]).hypot(right[1] - left[1])
2045 / noise_scale
2046 })
2047 .fold(0.0, f64::max);
2048 let variance_residual = (next.variance / previous.variance).ln().abs();
2049 weight_residual.max(mean_residual).max(variance_residual)
2050}
2051
2052fn fit_weighted_component_circle(
2053 component_means: &Array2<f64>,
2054 component_mass: &Array1<f64>,
2055 initial_center: &Array1<f64>,
2056 initial_radius: f64,
2057 parameter_tol: f64,
2058 max_iter: usize,
2059) -> Result<(Array1<f64>, f64, Array2<f64>), String> {
2060 let k = component_means.nrows();
2061 let total_mass = component_mass.sum();
2062 if component_means.ncols() != 2
2063 || component_mass.len() != k
2064 || component_mass
2065 .iter()
2066 .any(|mass| !mass.is_finite() || *mass <= 0.0)
2067 || !(total_mass.is_finite() && total_mass > 0.0)
2068 {
2069 return Err("ring M-step requires positive component masses and 2-D means".to_string());
2070 }
2071 let mut center = initial_center.clone();
2072 let mut radius = initial_radius;
2073 let mut directions = Array2::<f64>::zeros((k, 2));
2074 for _ in 0..max_iter {
2075 for component in 0..k {
2076 let dx = component_means[[component, 0]] - center[0];
2077 let dy = component_means[[component, 1]] - center[1];
2078 let norm = dx.hypot(dy);
2079 if !(norm.is_finite() && norm > 0.0) {
2080 return Err(
2081 "ring M-step reached a component centroid at the circle center; its angle is unidentified"
2082 .to_string(),
2083 );
2084 }
2085 directions[[component, 0]] = dx / norm;
2086 directions[[component, 1]] = dy / norm;
2087 }
2088
2089 let mut mean_point = Array1::<f64>::zeros(2);
2090 let mut mean_direction = Array1::<f64>::zeros(2);
2091 for component in 0..k {
2092 let weight = component_mass[component] / total_mass;
2093 for axis in 0..2 {
2094 mean_point[axis] += weight * component_means[[component, axis]];
2095 mean_direction[axis] += weight * directions[[component, axis]];
2096 }
2097 }
2098 let mut numerator = 0.0;
2099 let mut denominator = 0.0;
2100 for component in 0..k {
2101 let mass = component_mass[component];
2102 let dux = directions[[component, 0]] - mean_direction[0];
2103 let duy = directions[[component, 1]] - mean_direction[1];
2104 numerator += mass
2105 * (dux * (component_means[[component, 0]] - mean_point[0])
2106 + duy * (component_means[[component, 1]] - mean_point[1]));
2107 denominator += mass * (dux * dux + duy * duy);
2108 }
2109 if !(denominator.is_finite() && denominator > 0.0) {
2110 return Err(
2111 "ring M-step component directions are identical; radius and center are unidentified"
2112 .to_string(),
2113 );
2114 }
2115 let mut next_radius = numerator / denominator;
2116 if !next_radius.is_finite() || next_radius == 0.0 {
2117 return Err("ring M-step produced an unidentified zero radius".to_string());
2118 }
2119 if next_radius < 0.0 {
2120 next_radius = -next_radius;
2121 directions.mapv_inplace(|value| -value);
2122 }
2123 let next_center = Array1::from_vec(vec![
2124 mean_point[0] - next_radius * mean_direction[0],
2125 mean_point[1] - next_radius * mean_direction[1],
2126 ]);
2127 let residual = center
2128 .iter()
2129 .zip(next_center.iter())
2130 .map(|(&left, &right)| relative_parameter_step(left, right))
2131 .chain(std::iter::once(relative_parameter_step(
2132 radius,
2133 next_radius,
2134 )))
2135 .fold(0.0, f64::max);
2136 center = next_center;
2137 radius = next_radius;
2138 if residual <= parameter_tol {
2139 for component in 0..k {
2142 let dx = component_means[[component, 0]] - center[0];
2143 let dy = component_means[[component, 1]] - center[1];
2144 let norm = dx.hypot(dy);
2145 if !(norm.is_finite() && norm > 0.0) {
2146 return Err("ring M-step terminal component angle is unidentified".to_string());
2147 }
2148 directions[[component, 0]] = dx / norm;
2149 directions[[component, 1]] = dy / norm;
2150 }
2151 return Ok((center, radius, directions));
2152 }
2153 }
2154 Err(format!(
2155 "ring M-step did not certify its constrained center/radius fixed point after {max_iter} iterations"
2156 ))
2157}
2158
2159fn ring_mixture_m_step(
2160 data: ArrayView2<'_, f64>,
2161 responsibilities: ArrayView2<'_, f64>,
2162 previous: &RingMixtureState,
2163 config: GaussianMixtureConfig,
2164) -> Result<RingMixtureState, String> {
2165 let n = data.nrows();
2166 let k = responsibilities.ncols();
2167 let mut component_mass = Array1::<f64>::zeros(k);
2168 let mut component_means = Array2::<f64>::zeros((k, 2));
2169 for component in 0..k {
2170 let mass = responsibilities.column(component).sum();
2171 if !(mass.is_finite() && mass > 0.0) {
2172 return Err(
2173 "ring M-step reached a zero-mass component; the requested order is singular"
2174 .to_string(),
2175 );
2176 }
2177 component_mass[component] = mass;
2178 for row in 0..n {
2179 for axis in 0..2 {
2180 component_means[[component, axis]] +=
2181 responsibilities[[row, component]] * data[[row, axis]];
2182 }
2183 }
2184 for axis in 0..2 {
2185 component_means[[component, axis]] /= mass;
2186 }
2187 }
2188 let mut weights = component_mass.mapv(|mass| mass / n as f64);
2189 let weight_sum = weights.sum();
2190 weights.mapv_inplace(|weight| weight / weight_sum);
2191 let (center, radius, directions) = fit_weighted_component_circle(
2192 &component_means,
2193 &component_mass,
2194 &previous.center,
2195 previous.radius,
2196 config.parameter_tol,
2197 config.max_iter,
2198 )?;
2199 let means = ring_component_means(¢er, radius, &directions);
2200 let mut expected_squared_error = 0.0;
2201 for row in 0..n {
2202 for component in 0..k {
2203 let dx = data[[row, 0]] - means[[component, 0]];
2204 let dy = data[[row, 1]] - means[[component, 1]];
2205 expected_squared_error += responsibilities[[row, component]] * (dx * dx + dy * dy);
2206 }
2207 }
2208 let variance = (expected_squared_error / (2 * n) as f64).max(config.covariance_floor);
2209 if !variance.is_finite() {
2210 return Err("ring M-step produced non-finite shared variance".to_string());
2211 }
2212 Ok(RingMixtureState {
2213 weights,
2214 center,
2215 radius,
2216 directions,
2217 variance,
2218 mean_log_likelihood: f64::NAN,
2219 completed_iterations: previous.completed_iterations + 1,
2220 })
2221}
2222
2223pub fn fit_ring_gaussian_mixture(
2226 data: ArrayView2<'_, f64>,
2227 k: usize,
2228 config: GaussianMixtureConfig,
2229) -> Result<RingGaussianMixtureFit, String> {
2230 validate_gaussian_mixture_problem(data, k, config).map_err(|error| error.to_string())?;
2231 if data.ncols() != 2 {
2232 return Err(format!(
2233 "ring-of-clusters fitting requires exactly two columns, got {}",
2234 data.ncols()
2235 ));
2236 }
2237 if k < 3 {
2238 return Err(format!(
2239 "ring-of-clusters fitting requires at least three component centers, got {k}"
2240 ));
2241 }
2242 let seeded_means = gam_terms::basis::select_centers_by_strategy(
2243 data,
2244 &gam_terms::basis::CenterStrategy::KMeans {
2245 num_centers: k,
2246 max_iter: config.kmeans_max_iter,
2247 },
2248 )
2249 .map_err(|error| format!("ring-of-clusters deterministic seeding failed: {error}"))?;
2250 let component_mass = Array1::<f64>::ones(k);
2251 let mut initial_center = Array1::<f64>::zeros(2);
2252 for component in 0..k {
2253 initial_center[0] += seeded_means[[component, 0]] / k as f64;
2254 initial_center[1] += seeded_means[[component, 1]] / k as f64;
2255 }
2256 let mut initial_radius = 0.0;
2257 for component in 0..k {
2258 initial_radius += (seeded_means[[component, 0]] - initial_center[0])
2259 .hypot(seeded_means[[component, 1]] - initial_center[1])
2260 / k as f64;
2261 }
2262 if !(initial_radius.is_finite() && initial_radius > 0.0) {
2263 return Err("ring-of-clusters seed has an unidentified zero radius".to_string());
2264 }
2265 let (center, radius, directions) = fit_weighted_component_circle(
2266 &seeded_means,
2267 &component_mass,
2268 &initial_center,
2269 initial_radius,
2270 config.parameter_tol,
2271 config.max_iter,
2272 )?;
2273 let means = ring_component_means(¢er, radius, &directions);
2274 let mut squared_error = 0.0;
2275 for row in 0..data.nrows() {
2276 let mut nearest = f64::INFINITY;
2277 for component in 0..k {
2278 let dx = data[[row, 0]] - means[[component, 0]];
2279 let dy = data[[row, 1]] - means[[component, 1]];
2280 nearest = nearest.min(dx * dx + dy * dy);
2281 }
2282 squared_error += nearest;
2283 }
2284 let variance = (squared_error / (2 * data.nrows()) as f64).max(config.covariance_floor);
2285 let mut state = RingMixtureState {
2286 weights: Array1::from_elem(k, 1.0 / k as f64),
2287 center,
2288 radius,
2289 directions,
2290 variance,
2291 mean_log_likelihood: f64::NAN,
2292 completed_iterations: 0,
2293 };
2294 for additional_updates in 0..=config.max_iter {
2295 let (responsibilities, current_mean, current_roundoff) = ring_mixture_e_step(data, &state)?;
2296 state.mean_log_likelihood = current_mean;
2297 let mut next = ring_mixture_m_step(data, responsibilities.view(), &state, config)?;
2298 let (_, next_mean, next_roundoff) = ring_mixture_e_step(data, &next)?;
2299 next.mean_log_likelihood = next_mean;
2300 let objective_scale = current_mean.abs().max(next_mean.abs()).max(1.0);
2301 let objective_step = next_mean - current_mean;
2302 let objective_residual = objective_step.abs() / objective_scale;
2303 let parameter_residual = ring_identifiable_parameter_residual(&state, &next);
2304 let monotonicity_uncertainty = gaussian_mixture_monotonicity_uncertainty(
2305 objective_scale,
2306 current_roundoff,
2307 next_roundoff,
2308 );
2309 let certificate = GaussianMixtureCertificate {
2310 mean_log_likelihood: current_mean,
2311 mean_log_likelihood_gain: objective_step,
2312 monotonicity_uncertainty,
2313 objective_residual,
2314 objective_tolerance: config.loglik_tol,
2315 parameter_residual,
2316 parameter_tolerance: config.parameter_tol,
2317 };
2318 if objective_step < -monotonicity_uncertainty {
2319 return Err(format!(
2320 "ring-of-clusters generalized EM violated monotone ascent at iteration {}: {current_mean:.12e} -> {next_mean:.12e} (comparison uncertainty {monotonicity_uncertainty:.3e})",
2321 state.completed_iterations
2322 ));
2323 }
2324 if objective_residual <= config.loglik_tol && parameter_residual <= config.parameter_tol {
2325 let loglik = current_mean * data.nrows() as f64;
2326 if !loglik.is_finite() {
2327 return Err("ring-of-clusters total log likelihood overflowed".to_string());
2328 }
2329 return Ok(RingGaussianMixtureFit {
2330 weights: state.weights,
2331 center: state.center,
2332 radius: state.radius,
2333 directions: state.directions,
2334 variance: state.variance,
2335 k,
2336 n_obs: data.nrows(),
2337 loglik,
2338 iterations: state.completed_iterations,
2339 certificate,
2340 });
2341 }
2342 if additional_updates == config.max_iter {
2343 return Err(format!(
2344 "ring-of-clusters generalized EM did not certify after {} iterations: objective residual {:.6e}/{:.3e}, parameter-map residual {:.6e}/{:.3e}",
2345 config.max_iter,
2346 objective_residual,
2347 config.loglik_tol,
2348 parameter_residual,
2349 config.parameter_tol,
2350 ));
2351 }
2352 state = next;
2353 }
2354 Err("ring-of-clusters generalized EM exhausted without a terminal certificate".to_string())
2355}
2356
2357#[derive(Debug, Clone, Copy)]
2378pub struct CircularGaussianFit2d {
2379 center: [f64; 2],
2380 radius: f64,
2381 noise_variance: f64,
2382}
2383
2384impl CircularGaussianFit2d {
2385 pub const NUM_FREE_PARAMETERS: usize = 4;
2387
2388 pub fn from_parameters(
2390 center: [f64; 2],
2391 radius: f64,
2392 noise_variance: f64,
2393 ) -> Result<Self, String> {
2394 if !center.iter().all(|value| value.is_finite()) {
2395 return Err("circular Gaussian center must be finite".to_string());
2396 }
2397 if !(radius.is_finite() && radius >= 0.0) {
2398 return Err("circular Gaussian radius must be finite and nonnegative".to_string());
2399 }
2400 if !(noise_variance.is_finite() && noise_variance > 0.0) {
2401 return Err("circular Gaussian noise variance must be finite and positive".to_string());
2402 }
2403 Ok(Self {
2404 center,
2405 radius,
2406 noise_variance,
2407 })
2408 }
2409
2410 pub fn fit(coords: ArrayView2<'_, f64>, rows: &[usize]) -> Result<Self, String> {
2412 if coords.ncols() != 2 {
2413 return Err(format!(
2414 "circular Gaussian requires 2-D data, got {} columns",
2415 coords.ncols()
2416 ));
2417 }
2418 if rows.is_empty() {
2419 return Err("circular Gaussian requires a nonempty training set".to_string());
2420 }
2421 if rows.iter().any(|&row| row >= coords.nrows()) {
2422 return Err("circular Gaussian row index is out of bounds".to_string());
2423 }
2424 if rows
2425 .iter()
2426 .any(|&row| !coords[[row, 0]].is_finite() || !coords[[row, 1]].is_finite())
2427 {
2428 return Err("circular Gaussian requires finite training coordinates".to_string());
2429 }
2430
2431 let anchor_row = rows[0];
2435 let anchor = [coords[[anchor_row, 0]], coords[[anchor_row, 1]]];
2436 let mut scale = 0.0_f64;
2437 for &row in rows {
2438 let dx = coords[[row, 0]] - anchor[0];
2439 let dy = coords[[row, 1]] - anchor[1];
2440 if !(dx.is_finite() && dy.is_finite()) {
2441 return Err("circular Gaussian coordinate range exceeds f64".to_string());
2442 }
2443 scale = scale.max(dx.hypot(dy));
2444 }
2445 if !(scale.is_finite() && scale > 0.0) {
2446 return Err("circular Gaussian requires nonzero spatial extent".to_string());
2447 }
2448
2449 let mut points = Vec::with_capacity(rows.len());
2450 let mut mean = [0.0_f64; 2];
2451 for &row in rows {
2452 let point = [
2453 (coords[[row, 0]] - anchor[0]) / scale,
2454 (coords[[row, 1]] - anchor[1]) / scale,
2455 ];
2456 points.push(point);
2457 mean[0] += point[0];
2458 mean[1] += point[1];
2459 }
2460 let count = rows.len() as f64;
2461 mean[0] /= count;
2462 mean[1] /= count;
2463
2464 let mut squared_radii = Vec::with_capacity(rows.len());
2469 let mut mean_squared_radius = 0.0_f64;
2470 for point in &points {
2471 let dx = point[0] - mean[0];
2472 let dy = point[1] - mean[1];
2473 let squared_radius = dx * dx + dy * dy;
2474 squared_radii.push(squared_radius);
2475 mean_squared_radius += squared_radius;
2476 }
2477 mean_squared_radius /= count;
2478 let mut squared_radius_variance = 0.0_f64;
2479 for squared_radius in squared_radii {
2480 squared_radius_variance += (squared_radius - mean_squared_radius).powi(2);
2481 }
2482 squared_radius_variance /= count;
2483
2484 let variance_floor = (64.0 * f64::EPSILON * mean_squared_radius).max(f64::MIN_POSITIVE);
2488 let radius_squared = (mean_squared_radius * mean_squared_radius - squared_radius_variance)
2489 .max(0.0)
2490 .sqrt();
2491 let mut radius = radius_squared.sqrt();
2492 let mut noise_variance = (0.5 * (mean_squared_radius - radius_squared)).max(variance_floor);
2493 let mut center = mean;
2494
2495 const MAX_EM_ITERATIONS: usize = 4096;
2500 const EM_TOLERANCE: f64 = 2.0e-12;
2501 let mut posterior_means = vec![[0.0_f64; 2]; points.len()];
2502 let mut converged = false;
2503 for _ in 0..MAX_EM_ITERATIONS {
2504 let mut posterior_mean = [0.0_f64; 2];
2505 for (point, latent_mean) in points.iter().zip(&mut posterior_means) {
2506 let dx = point[0] - center[0];
2507 let dy = point[1] - center[1];
2508 let observed_radius = dx.hypot(dy);
2509 if observed_radius == 0.0 || radius == 0.0 {
2510 *latent_mean = [0.0, 0.0];
2511 } else {
2512 let (_, bessel_ratio) =
2513 circular_gaussian_bessel_terms(radius, observed_radius, noise_variance);
2514 if !(bessel_ratio.is_finite() && (0.0..=1.0).contains(&bessel_ratio)) {
2515 return Err("circular Gaussian Bessel ratio left [0, 1]".to_string());
2516 }
2517 let multiplier = bessel_ratio / observed_radius;
2518 *latent_mean = [multiplier * dx, multiplier * dy];
2519 }
2520 posterior_mean[0] += latent_mean[0];
2521 posterior_mean[1] += latent_mean[1];
2522 }
2523 posterior_mean[0] /= count;
2524 posterior_mean[1] /= count;
2525
2526 let denominator =
2527 1.0 - posterior_mean[0] * posterior_mean[0] - posterior_mean[1] * posterior_mean[1];
2528 if !(denominator.is_finite() && denominator > 0.0) {
2529 return Err("circular Gaussian EM radius update is singular".to_string());
2530 }
2531 let mut radius_numerator = 0.0_f64;
2532 for (point, latent_mean) in points.iter().zip(&posterior_means) {
2533 radius_numerator +=
2534 latent_mean[0] * (point[0] - mean[0]) + latent_mean[1] * (point[1] - mean[1]);
2535 }
2536 let next_radius = (radius_numerator / (count * denominator)).max(0.0);
2537 let next_center = [
2538 mean[0] - next_radius * posterior_mean[0],
2539 mean[1] - next_radius * posterior_mean[1],
2540 ];
2541
2542 let mut residual_sum = 0.0_f64;
2545 for (point, latent_mean) in points.iter().zip(&posterior_means) {
2546 let dx = point[0] - next_center[0];
2547 let dy = point[1] - next_center[1];
2548 let ex = dx - next_radius * latent_mean[0];
2549 let ey = dy - next_radius * latent_mean[1];
2550 let latent_norm_squared =
2551 latent_mean[0] * latent_mean[0] + latent_mean[1] * latent_mean[1];
2552 residual_sum += ex * ex
2553 + ey * ey
2554 + next_radius * next_radius * (1.0 - latent_norm_squared).max(0.0);
2555 }
2556 let next_noise_variance = (residual_sum / (2.0 * count)).max(variance_floor);
2557
2558 let parameter_change = (next_center[0] - center[0])
2559 .hypot(next_center[1] - center[1])
2560 .max((next_radius - radius).abs())
2561 .max(
2562 (next_noise_variance - noise_variance).abs()
2563 / (next_noise_variance + noise_variance),
2564 );
2565 center = next_center;
2566 radius = next_radius;
2567 noise_variance = next_noise_variance;
2568 if parameter_change <= EM_TOLERANCE {
2569 converged = true;
2570 break;
2571 }
2572 }
2573 if !converged {
2574 return Err("circular Gaussian maximum-likelihood fit did not converge".to_string());
2575 }
2576
2577 let fitted_noise_sd = scale * noise_variance.sqrt();
2578 Self::from_parameters(
2579 [anchor[0] + scale * center[0], anchor[1] + scale * center[1]],
2580 scale * radius,
2581 fitted_noise_sd * fitted_noise_sd,
2582 )
2583 .map_err(|error| format!("circular Gaussian fit produced invalid parameters: {error}"))
2584 }
2585
2586 pub const fn center(self) -> [f64; 2] {
2588 self.center
2589 }
2590
2591 pub const fn radius(self) -> f64 {
2593 self.radius
2594 }
2595
2596 pub const fn noise_variance(self) -> f64 {
2598 self.noise_variance
2599 }
2600
2601 pub fn log_density(self, x: f64, y: f64) -> f64 {
2603 let observed_radius = (x - self.center[0]).hypot(y - self.center[1]);
2604 let (log_i0_minus_kappa, _) =
2605 circular_gaussian_bessel_terms(self.radius, observed_radius, self.noise_variance);
2606 let standardized_radial_residual =
2607 (observed_radius - self.radius) / self.noise_variance.sqrt();
2608 -std::f64::consts::TAU.ln()
2611 - self.noise_variance.ln()
2612 - 0.5 * standardized_radial_residual.powi(2)
2613 + log_i0_minus_kappa
2614 }
2615
2616 pub fn log_likelihood(
2618 self,
2619 coords: ArrayView2<'_, f64>,
2620 rows: &[usize],
2621 ) -> Result<f64, String> {
2622 if coords.ncols() != 2 || rows.iter().any(|&row| row >= coords.nrows()) {
2623 return Err(
2624 "circular Gaussian likelihood received invalid coordinates or rows".to_string(),
2625 );
2626 }
2627 let mut log_densities = Vec::with_capacity(rows.len());
2628 for &row in rows {
2629 let value = self.log_density(coords[[row, 0]], coords[[row, 1]]);
2630 if !value.is_finite() {
2631 return Err("circular Gaussian likelihood is not finite".to_string());
2632 }
2633 log_densities.push(value);
2634 }
2635 let log_likelihood = pairwise_sum(&log_densities);
2636 if !log_likelihood.is_finite() {
2637 return Err("circular Gaussian likelihood sum is not finite".to_string());
2638 }
2639 Ok(log_likelihood)
2640 }
2641
2642 pub fn fit_with_bic(
2646 coords: ArrayView2<'_, f64>,
2647 rows: &[usize],
2648 ) -> Result<(Self, f64), String> {
2649 let fit = Self::fit(coords, rows)?;
2650 let log_likelihood = fit.log_likelihood(coords, rows)?;
2651 let bic =
2652 -log_likelihood + 0.5 * Self::NUM_FREE_PARAMETERS as f64 * (rows.len() as f64).ln();
2653 if !bic.is_finite() {
2654 return Err("circular Gaussian BIC is not finite".to_string());
2655 }
2656 Ok((fit, bic))
2657 }
2658}
2659
2660fn circular_gaussian_bessel_terms(
2666 radius: f64,
2667 observed_radius: f64,
2668 noise_variance: f64,
2669) -> (f64, f64) {
2670 if radius == 0.0 || observed_radius == 0.0 {
2671 return (0.0, 0.0);
2672 }
2673 let kappa = radius * observed_radius / noise_variance;
2674 if kappa.is_finite() {
2675 return bessel_i0_log_minus_abs_and_ratio(kappa);
2676 }
2677 let log_kappa = radius.ln() + observed_radius.ln() - noise_variance.ln();
2678 if log_kappa <= f64::MAX.ln() {
2679 return bessel_i0_log_minus_abs_and_ratio(log_kappa.exp());
2682 }
2683 (-0.5 * (std::f64::consts::TAU.ln() + log_kappa), 1.0)
2684}
2685#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
2706pub enum UnionStructure {
2707 CircleCircle,
2709 CirclePointCluster,
2711 LineCluster,
2713}
2714
2715pub const UNION_STRUCTURE_LADDER: &[UnionStructure] = &[
2717 UnionStructure::CircleCircle,
2718 UnionStructure::CirclePointCluster,
2719 UnionStructure::LineCluster,
2720];
2721
2722#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
2730pub enum UnionComponentKind {
2731 Circle,
2732 Line,
2733 PointCluster,
2734}
2735
2736impl UnionStructure {
2737 pub const fn as_str(self) -> &'static str {
2739 match self {
2740 UnionStructure::CircleCircle => "union_circle+circle",
2741 UnionStructure::CirclePointCluster => "union_circle+cluster",
2742 UnionStructure::LineCluster => "union_line+cluster",
2743 }
2744 }
2745
2746 pub const fn components(self) -> &'static [UnionComponentKind] {
2748 match self {
2749 UnionStructure::CircleCircle => {
2750 &[UnionComponentKind::Circle, UnionComponentKind::Circle]
2751 }
2752 UnionStructure::CirclePointCluster => {
2753 &[UnionComponentKind::Circle, UnionComponentKind::PointCluster]
2754 }
2755 UnionStructure::LineCluster => {
2756 &[UnionComponentKind::Line, UnionComponentKind::PointCluster]
2757 }
2758 }
2759 }
2760
2761 pub const fn num_components(self) -> usize {
2763 self.components().len()
2764 }
2765}
2766
2767#[derive(Debug, Clone)]
2773pub struct UnionComponentFit {
2774 pub kind: UnionComponentKind,
2775 pub row_count: usize,
2776 pub num_parameters: usize,
2777 pub mixing_weight: f64,
2778}
2779
2780#[derive(Debug, Clone)]
2784pub struct UnionStructureFit {
2785 pub structure: UnionStructure,
2786 pub components: Vec<UnionComponentFit>,
2787 pub log_likelihood: f64,
2789 pub bic: f64,
2791 pub total_parameters: usize,
2793}
2794
2795pub fn union_responsibility_split(
2800 data: ArrayView2<'_, f64>,
2801 m: usize,
2802 config: GaussianMixtureConfig,
2803) -> Result<Vec<Vec<usize>>, String> {
2804 let n = data.nrows();
2805 if m == 0 {
2806 return Err("union split requires at least one component".to_string());
2807 }
2808 if m > n {
2809 return Err(format!(
2810 "union split requested {m} groups but data has {n} rows"
2811 ));
2812 }
2813 if m == 1 {
2814 return Ok(vec![(0..n).collect()]);
2815 }
2816 let fit = fit_gaussian_mixture(data, m, config).map_err(|error| error.to_string())?;
2817 let mut groups: Vec<Vec<usize>> = vec![Vec::new(); m];
2818 let mut comp = Vec::with_capacity(m);
2820 for j in 0..m {
2821 comp.push(GaussianComponentEval::factor(
2822 fit.means.row(j),
2823 &fit.covariances[j],
2824 )?);
2825 }
2826 let log_w = fit
2827 .weights
2828 .iter()
2829 .enumerate()
2830 .map(|(component, &weight)| {
2831 if weight.is_finite() && weight > 0.0 {
2832 Ok(weight.ln())
2833 } else {
2834 Err(format!(
2835 "union split received invalid fitted weight {weight} for component {component}"
2836 ))
2837 }
2838 })
2839 .collect::<Result<Vec<_>, _>>()?;
2840 for i in 0..n {
2841 let row = data.row(i);
2842 let mut best_j = 0usize;
2843 let mut best_lt = f64::NEG_INFINITY;
2844 for j in 0..m {
2845 let lt = log_w[j] + comp[j].log_density(row);
2846 if lt > best_lt {
2847 best_lt = lt;
2848 best_j = j;
2849 }
2850 }
2851 if !best_lt.is_finite() {
2852 return Err(format!(
2853 "union split produced no finite component score at row {i}"
2854 ));
2855 }
2856 groups[best_j].push(i);
2857 }
2858 Ok(groups)
2859}
2860
2861pub fn fit_union_structure(
2870 data: ArrayView2<'_, f64>,
2871 structure: UnionStructure,
2872 config: GaussianMixtureConfig,
2873) -> Result<UnionStructureFit, String> {
2874 let fitted = fit_union_density(data, structure, config)?;
2875 Ok(UnionStructureFit {
2876 structure,
2877 components: fitted
2878 .components
2879 .iter()
2880 .map(UnionComponentDensity::summary)
2881 .collect(),
2882 log_likelihood: fitted.log_likelihood,
2883 bic: fitted.bic,
2884 total_parameters: fitted.total_parameters,
2885 })
2886}
2887
2888pub fn fit_union_ladder(
2893 data: ArrayView2<'_, f64>,
2894 config: GaussianMixtureConfig,
2895) -> Result<Vec<UnionStructureFit>, String> {
2896 let mut fits = Vec::new();
2897 let mut errors = Vec::new();
2898 for &structure in UNION_STRUCTURE_LADDER {
2899 match fit_union_structure(data, structure, config) {
2900 Ok(fit) => fits.push(fit),
2901 Err(e) => errors.push(format!("{}: {e}", structure.as_str())),
2902 }
2903 }
2904 if !errors.is_empty() {
2905 return Err(format!(
2906 "union ladder comparison failed; every declared structure must fit ({})",
2907 errors.join("; ")
2908 ));
2909 }
2910 if fits.is_empty() {
2911 return Err("union ladder is empty".to_string());
2912 }
2913 let ranked = rank_priority_candidates(
2914 fits.into_iter()
2915 .enumerate()
2916 .map(|(idx, row)| {
2917 let score = row.bic;
2918 let tie = row.total_parameters; PriorityCandidate::new(row, idx, score, tie)
2920 })
2921 .collect(),
2922 )
2923 .into_iter()
2924 .map(|row| row.item)
2925 .collect::<Vec<_>>();
2926 Ok(ranked)
2927}
2928
2929fn gather_union_rows(data: ArrayView2<'_, f64>, idx: &[usize]) -> Array2<f64> {
2930 let d = data.ncols();
2931 let mut out = Array2::<f64>::zeros((idx.len(), d));
2932 for (r, &i) in idx.iter().enumerate() {
2933 for c in 0..d {
2934 out[[r, c]] = data[[i, c]];
2935 }
2936 }
2937 out
2938}
2939
2940fn union_circle_rows(group: ArrayView2<'_, f64>) -> Result<Vec<usize>, String> {
2945 let minimum_rows = CircularGaussianFit2d::NUM_FREE_PARAMETERS + 1;
2946 if group.nrows() < minimum_rows {
2947 return Err(format!(
2948 "union circle component needs at least {minimum_rows} rows, got {}",
2949 group.nrows()
2950 ));
2951 }
2952 Ok((0..group.nrows()).collect())
2953}
2954
2955#[derive(Debug, Clone)]
2960enum UnionDensityModel {
2961 Gaussian(GaussianComponentEval),
2962 Circle(CircularGaussianFit2d),
2963}
2964
2965#[derive(Debug, Clone)]
2966struct UnionComponentDensity {
2967 kind: UnionComponentKind,
2968 row_count: usize,
2969 num_parameters: usize,
2970 mixing_weight: f64,
2971 log_weight: f64,
2972 model: UnionDensityModel,
2973}
2974
2975impl UnionComponentDensity {
2976 fn summary(&self) -> UnionComponentFit {
2977 UnionComponentFit {
2978 kind: self.kind,
2979 row_count: self.row_count,
2980 num_parameters: self.num_parameters,
2981 mixing_weight: self.mixing_weight,
2982 }
2983 }
2984
2985 fn dimension(&self) -> usize {
2986 match &self.model {
2987 UnionDensityModel::Gaussian(eval) => eval.d,
2988 UnionDensityModel::Circle(_) => 2,
2989 }
2990 }
2991
2992 fn weighted_log_density(&self, y: ArrayView1<'_, f64>) -> f64 {
2994 let component_log_density = match &self.model {
2995 UnionDensityModel::Gaussian(eval) => eval.log_density(y),
2996 UnionDensityModel::Circle(fit) => fit.log_density(y[0], y[1]),
2997 };
2998 self.log_weight + component_log_density
2999 }
3000}
3001
3002#[derive(Debug, Clone)]
3003struct FittedUnionDensity {
3004 components: Vec<UnionComponentDensity>,
3005 log_likelihood: f64,
3006 bic: f64,
3007 total_parameters: usize,
3008}
3009
3010fn fit_union_density(
3011 train: ArrayView2<'_, f64>,
3012 structure: UnionStructure,
3013 config: GaussianMixtureConfig,
3014) -> Result<FittedUnionDensity, String> {
3015 let groups = union_responsibility_split(train, structure.num_components(), config)?;
3016 fit_union_density_from_groups(train, structure, &groups, config)
3017}
3018
3019fn fit_union_density_from_groups(
3023 train: ArrayView2<'_, f64>,
3024 structure: UnionStructure,
3025 groups: &[Vec<usize>],
3026 config: GaussianMixtureConfig,
3027) -> Result<FittedUnionDensity, String> {
3028 validate_union_partition(train.nrows(), structure.num_components(), groups)?;
3029 let assignments = unique_union_role_assignments(structure.components());
3030 let mut best: Option<FittedUnionDensity> = None;
3031 let mut errors = Vec::new();
3032
3033 for roles in assignments {
3034 let candidate = (|| {
3035 let mut components = Vec::with_capacity(groups.len());
3036 let n_train = train.nrows() as f64;
3037 for (&kind, rows) in roles.iter().zip(groups) {
3038 let group = gather_union_rows(train, rows);
3039 let mixing_weight = rows.len() as f64 / n_train;
3040 components.push(fit_union_component_density(
3041 group.view(),
3042 kind,
3043 mixing_weight,
3044 config,
3045 )?);
3046 }
3047
3048 let component_parameters = components.iter().try_fold(0usize, |sum, component| {
3049 sum.checked_add(component.num_parameters)
3050 .ok_or_else(|| "union component parameter count overflowed usize".to_string())
3051 })?;
3052 let mixing_parameters = components.len() - 1;
3053 let total_parameters = component_parameters
3054 .checked_add(mixing_parameters)
3055 .ok_or_else(|| "union total parameter count overflowed usize".to_string())?;
3056 let per_point = score_union_components(&components, train)?;
3057 let log_likelihood = pairwise_sum(
3058 per_point
3059 .as_slice()
3060 .expect("owned union score vector must be contiguous"),
3061 );
3062 if !log_likelihood.is_finite() {
3063 return Err("union training log likelihood is non-finite".to_string());
3064 }
3065 let bic = -log_likelihood + 0.5 * total_parameters as f64 * (train.nrows() as f64).ln();
3066 if !bic.is_finite() {
3067 return Err("union normalized soft-mixture BIC is non-finite".to_string());
3068 }
3069 Ok(FittedUnionDensity {
3070 components,
3071 log_likelihood,
3072 bic,
3073 total_parameters,
3074 })
3075 })();
3076
3077 match candidate {
3078 Ok(candidate) => {
3079 let replace = match &best {
3080 Some(current) => candidate.bic.total_cmp(¤t.bic).is_lt(),
3081 None => true,
3082 };
3083 if replace {
3086 best = Some(candidate);
3087 }
3088 }
3089 Err(error) => errors.push(format!("{roles:?}: {error}")),
3090 }
3091 }
3092
3093 best.ok_or_else(|| {
3094 format!(
3095 "union {} has no finite role assignment ({})",
3096 structure.as_str(),
3097 errors.join("; ")
3098 )
3099 })
3100}
3101
3102fn validate_union_partition(
3103 n_rows: usize,
3104 expected_groups: usize,
3105 groups: &[Vec<usize>],
3106) -> Result<(), String> {
3107 if n_rows == 0 {
3108 return Err("union fitting requires at least one training row".to_string());
3109 }
3110 if groups.len() != expected_groups {
3111 return Err(format!(
3112 "union partition has {} groups, expected {expected_groups}",
3113 groups.len()
3114 ));
3115 }
3116 let mut seen = vec![false; n_rows];
3117 for (group_index, rows) in groups.iter().enumerate() {
3118 if rows.is_empty() {
3119 return Err(format!("union partition group {group_index} is empty"));
3120 }
3121 for &row in rows {
3122 if row >= n_rows {
3123 return Err(format!(
3124 "union partition group {group_index} contains out-of-range row {row} for {n_rows} rows"
3125 ));
3126 }
3127 if std::mem::replace(&mut seen[row], true) {
3128 return Err(format!("union partition contains duplicate row {row}"));
3129 }
3130 }
3131 }
3132 if let Some(missing) = seen.iter().position(|included| !included) {
3133 return Err(format!("union partition omits row {missing}"));
3134 }
3135 Ok(())
3136}
3137
3138fn unique_union_role_assignments(roles: &[UnionComponentKind]) -> Vec<Vec<UnionComponentKind>> {
3139 fn visit(
3140 roles: &[UnionComponentKind],
3141 used: &mut [bool],
3142 assignment: &mut Vec<UnionComponentKind>,
3143 out: &mut Vec<Vec<UnionComponentKind>>,
3144 ) {
3145 if assignment.len() == roles.len() {
3146 out.push(assignment.clone());
3147 return;
3148 }
3149 let mut used_at_depth = Vec::new();
3150 for (index, &role) in roles.iter().enumerate() {
3151 if used[index] || used_at_depth.contains(&role) {
3152 continue;
3153 }
3154 used_at_depth.push(role);
3155 used[index] = true;
3156 assignment.push(role);
3157 visit(roles, used, assignment, out);
3158 assignment.pop();
3159 used[index] = false;
3160 }
3161 }
3162
3163 let mut out = Vec::new();
3164 visit(
3165 roles,
3166 &mut vec![false; roles.len()],
3167 &mut Vec::with_capacity(roles.len()),
3168 &mut out,
3169 );
3170 out
3171}
3172
3173fn fit_union_component_density(
3174 group: ArrayView2<'_, f64>,
3175 kind: UnionComponentKind,
3176 mixing_weight: f64,
3177 config: GaussianMixtureConfig,
3178) -> Result<UnionComponentDensity, String> {
3179 if !(mixing_weight.is_finite() && mixing_weight > 0.0 && mixing_weight <= 1.0) {
3180 return Err(format!(
3181 "union component mixing weight must be finite and in (0, 1], got {mixing_weight}"
3182 ));
3183 }
3184 let row_count = group.nrows();
3185 let (model, num_parameters) = match kind {
3186 UnionComponentKind::Line => {
3187 if group.nrows() < group.ncols() + 1 {
3188 return Err(format!(
3189 "union line component needs >= {} rows, got {}",
3190 group.ncols() + 1,
3191 group.nrows()
3192 ));
3193 }
3194 let fit = fit_gaussian_mixture(group, 1, config).map_err(|error| error.to_string())?;
3195 let num_parameters = fit.num_free_parameters();
3196 let eval = GaussianComponentEval::factor(fit.means.row(0), &fit.covariances[0])?;
3197 (UnionDensityModel::Gaussian(eval), num_parameters)
3198 }
3199 UnionComponentKind::PointCluster => {
3200 if group.nrows() < group.ncols() + 1 {
3201 return Err(format!(
3202 "union isotropic point component needs >= {} rows, got {}",
3203 group.ncols() + 1,
3204 group.nrows()
3205 ));
3206 }
3207 let eval = fit_isotropic_gaussian_component(group, config.covariance_floor)?;
3208 (
3209 UnionDensityModel::Gaussian(eval),
3210 group
3211 .ncols()
3212 .checked_add(1)
3213 .ok_or_else(|| "union point parameter count overflowed usize".to_string())?,
3214 )
3215 }
3216 UnionComponentKind::Circle => {
3217 let rows = union_circle_rows(group)?;
3218 let fit = CircularGaussianFit2d::fit(group, &rows)?;
3219 (
3220 UnionDensityModel::Circle(fit),
3221 CircularGaussianFit2d::NUM_FREE_PARAMETERS,
3222 )
3223 }
3224 };
3225 Ok(UnionComponentDensity {
3226 kind,
3227 row_count,
3228 num_parameters,
3229 mixing_weight,
3230 log_weight: mixing_weight.ln(),
3231 model,
3232 })
3233}
3234
3235#[derive(Debug, Clone, Copy)]
3236struct StableScalarMeanChart {
3237 origin: f64,
3238 scale: f64,
3239 normalized_offset: f64,
3240}
3241
3242impl StableScalarMeanChart {
3243 #[inline]
3244 fn centered(self, value: f64) -> Result<f64, String> {
3245 let relative = value - self.origin;
3246 let centered = (-self.normalized_offset).mul_add(self.scale, relative);
3247 if centered.is_finite() {
3248 Ok(centered)
3249 } else {
3250 Err("union isotropic point residual is not representable".to_string())
3251 }
3252 }
3253}
3254
3255fn stable_scalar_mean_chart(values: ArrayView1<'_, f64>) -> Result<StableScalarMeanChart, String> {
3261 if values.is_empty() || values.iter().any(|value| !value.is_finite()) {
3262 return Err("stable scalar mean requires finite nonempty values".to_string());
3263 }
3264 let anchor = values[0];
3265 let anchor_chart_is_representable = values.iter().all(|&value| (value - anchor).is_finite());
3266 let origin = if anchor_chart_is_representable {
3267 anchor
3268 } else {
3269 0.0
3270 };
3271 let scale = values
3272 .iter()
3273 .map(|&value| (value - origin).abs())
3274 .fold(0.0_f64, f64::max);
3275 if scale == 0.0 {
3276 return Ok(StableScalarMeanChart {
3277 origin,
3278 scale: 0.0,
3279 normalized_offset: 0.0,
3280 });
3281 }
3282 let normalized = values
3283 .iter()
3284 .map(|&value| (value - origin) / scale)
3285 .collect::<Vec<_>>();
3286 let normalized_offset = pairwise_sum(&normalized) / values.len() as f64;
3287 let mean = normalized_offset.mul_add(scale, origin);
3288 if !(normalized_offset.is_finite() && mean.is_finite()) {
3289 return Err("union isotropic point mean is not representable".to_string());
3290 }
3291 Ok(StableScalarMeanChart {
3292 origin,
3293 scale,
3294 normalized_offset,
3295 })
3296}
3297
3298fn fit_isotropic_gaussian_component(
3300 group: ArrayView2<'_, f64>,
3301 covariance_floor: f64,
3302) -> Result<GaussianComponentEval, String> {
3303 let n = group.nrows();
3304 let d = group.ncols();
3305 if n == 0 || d == 0 {
3306 return Err("union isotropic point component requires a non-empty matrix".to_string());
3307 }
3308 if !(covariance_floor.is_finite() && covariance_floor > 0.0) {
3309 return Err(format!(
3310 "union isotropic covariance floor must be finite and positive, got {covariance_floor}"
3311 ));
3312 }
3313
3314 for row in group.rows() {
3315 for axis in 0..d {
3316 let value = row[axis];
3317 if !value.is_finite() {
3318 return Err(format!(
3319 "union isotropic point data contains non-finite coordinate {value}"
3320 ));
3321 }
3322 }
3323 }
3324
3325 let mut charts = Vec::with_capacity(d);
3326 for axis in 0..d {
3327 let chart = stable_scalar_mean_chart(group.column(axis))?;
3328 charts.push(chart);
3329 }
3330
3331 let scalar_count = n
3332 .checked_mul(d)
3333 .ok_or_else(|| "union isotropic residual count overflowed usize".to_string())?;
3334 let mut residuals = Vec::with_capacity(scalar_count);
3335 let mut residual_scale = 0.0_f64;
3336 for row in group.rows() {
3337 for axis in 0..d {
3338 let residual = charts[axis].centered(row[axis])?;
3339 residual_scale = residual_scale.max(residual.abs());
3340 residuals.push(residual);
3341 }
3342 }
3343 let variance = if residual_scale == 0.0 {
3344 covariance_floor
3345 } else {
3346 for residual in &mut residuals {
3347 *residual = (*residual / residual_scale).powi(2);
3348 }
3349 let normalized_mean_square = pairwise_sum(&residuals) / scalar_count as f64;
3350 let rms = residual_scale * normalized_mean_square.sqrt();
3351 let unconstrained = rms * rms;
3352 if !unconstrained.is_finite() {
3353 return Err("union isotropic point variance is non-finite".to_string());
3354 }
3355 unconstrained.max(covariance_floor)
3356 };
3357 GaussianComponentEval::isotropic(&charts, variance)
3358}
3359
3360fn score_union_components(
3361 components: &[UnionComponentDensity],
3362 eval: ArrayView2<'_, f64>,
3363) -> Result<Array1<f64>, String> {
3364 if components.is_empty() {
3365 return Err("union density requires at least one component".to_string());
3366 }
3367 if eval.iter().any(|coordinate| !coordinate.is_finite()) {
3368 return Err("union eval coordinates must be finite".to_string());
3369 }
3370 for component in components {
3371 if component.dimension() != eval.ncols() {
3372 return Err(format!(
3373 "union component {:?} has dimension {}, eval has {} columns",
3374 component.kind,
3375 component.dimension(),
3376 eval.ncols()
3377 ));
3378 }
3379 }
3380 let mut out = Array1::<f64>::zeros(eval.nrows());
3381 let mut terms = vec![f64::NEG_INFINITY; components.len()];
3382 for i in 0..eval.nrows() {
3383 let row = eval.row(i);
3384 let mut max_term = f64::NEG_INFINITY;
3385 for (component_index, component) in components.iter().enumerate() {
3386 let term = component.weighted_log_density(row);
3387 terms[component_index] = term;
3388 if term > max_term {
3389 max_term = term;
3390 }
3391 }
3392 let value = log_sum_exp(&terms, max_term);
3393 if !value.is_finite() {
3394 return Err(format!(
3395 "union density produced non-finite log density at eval row {i}"
3396 ));
3397 }
3398 out[i] = value;
3399 }
3400 Ok(out)
3401}
3402
3403pub fn union_per_point_log_density(
3408 train: ArrayView2<'_, f64>,
3409 eval: ArrayView2<'_, f64>,
3410 structure: UnionStructure,
3411 config: GaussianMixtureConfig,
3412) -> Result<Array1<f64>, String> {
3413 if train.ncols() != eval.ncols() {
3414 return Err(format!(
3415 "union held-out density: train has {} columns, eval has {}",
3416 train.ncols(),
3417 eval.ncols()
3418 ));
3419 }
3420 let fitted = fit_union_density(train, structure, config)?;
3421 score_union_components(&fitted.components, eval)
3422}
3423
3424#[derive(Clone, Debug)]
3426pub struct RemlCandidate {
3427 pub index: usize,
3428 pub name: String,
3429 pub score: f64,
3432 pub edf: Option<f64>,
3433 pub log_lik: Option<f64>,
3437 pub family: Option<String>,
3443 pub n_obs: Option<usize>,
3451}
3452
3453impl RemlCandidate {
3454 pub fn ranking_score(&self) -> f64 {
3473 match (self.log_lik, self.edf) {
3474 (Some(log_lik), Some(edf)) if log_lik.is_finite() && edf.is_finite() => {
3475 -2.0 * log_lik + 2.0 * edf
3476 }
3477 _ => self.score,
3478 }
3479 }
3480}
3481
3482#[derive(Clone, Debug)]
3483pub struct RemlComparison {
3484 pub ranking: Vec<RankedRow>,
3485 pub winner: String,
3486 pub evidence_summary: String,
3487 pub score_table: Vec<ScoreRow>,
3488}
3489
3490#[derive(Clone, Debug)]
3491pub struct RankedRow {
3492 pub name: String,
3493 pub score: f64,
3494 pub delta: f64,
3501 pub bayes_factor: f64,
3506 pub edf: Option<f64>,
3507}
3508
3509#[derive(Clone, Debug)]
3510pub struct ScoreRow {
3511 pub name: String,
3512 pub reml_score: f64,
3513 pub delta_reml: f64,
3514 pub bayes_factor_best_over_model: f64,
3515 pub effective_dof: Option<f64>,
3516}
3517
3518#[inline]
3520pub fn log_bayes_factor(reml_score_a: f64, reml_score_b: f64) -> f64 {
3521 reml_score_b - reml_score_a
3522}
3523
3524pub fn compare_reml_fits(mut candidates: Vec<RemlCandidate>) -> Result<RemlComparison, String> {
3528 if candidates.is_empty() {
3529 return Err("compare_models requires at least one fit".to_string());
3530 }
3531 {
3539 let mut seen_family: Option<&str> = None;
3540 for cand in &candidates {
3541 if let Some(fam) = cand.family.as_deref() {
3542 match seen_family {
3543 None => seen_family = Some(fam),
3544 Some(prev) if prev != fam => {
3545 return Err(format!(
3546 "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."
3547 ));
3548 }
3549 Some(_) => {}
3550 }
3551 }
3552 }
3553 }
3554 {
3565 let mut seen_n: Option<usize> = None;
3566 for cand in &candidates {
3567 if let Some(n) = cand.n_obs {
3568 match seen_n {
3569 None => seen_n = Some(n),
3570 Some(prev) if prev != n => {
3571 return Err(format!(
3572 "compare_models: cannot compare fits made on a different number of \
3573 observations (n={prev} vs n={n}); AIC / REML-LAML evidence scales \
3574 with the sample size, so their score difference is not a Bayes \
3575 factor. Compare models fit to the same response on the same data."
3576 ));
3577 }
3578 Some(_) => {}
3579 }
3580 }
3581 }
3582 }
3583 candidates = rank_priority_candidates(
3584 candidates
3585 .into_iter()
3586 .enumerate()
3587 .map(|(idx, row)| {
3588 let ranking = row.ranking_score();
3591 PriorityCandidate::new(row, idx, ranking, 0)
3592 })
3593 .collect(),
3594 )
3595 .into_iter()
3596 .map(|row| row.item)
3597 .collect();
3598
3599 let winner = candidates[0].name.clone();
3600 let best_ranking_score = candidates[0].ranking_score();
3610 let best_raw_score = candidates
3615 .iter()
3616 .map(|c| c.score)
3617 .fold(f64::INFINITY, f64::min);
3618 let mut ranking = Vec::with_capacity(candidates.len());
3619 let mut score_table = Vec::with_capacity(candidates.len());
3620 for row in &candidates {
3621 let delta = log_bayes_factor(best_ranking_score, row.ranking_score());
3622 let bayes_factor = (0.5 * delta).exp();
3629 let delta_reml = log_bayes_factor(best_raw_score, row.score);
3630 ranking.push(RankedRow {
3631 name: row.name.clone(),
3632 score: row.score,
3633 delta,
3634 bayes_factor,
3635 edf: row.edf,
3636 });
3637 score_table.push(ScoreRow {
3638 name: row.name.clone(),
3639 reml_score: row.score,
3640 delta_reml,
3641 bayes_factor_best_over_model: delta_reml.exp(),
3642 effective_dof: row.edf,
3643 });
3644 }
3645 let evidence_summary = if let Some(runner_up) = candidates.get(1) {
3650 let margin = runner_up.ranking_score() - candidates[0].ranking_score();
3651 format!(
3656 "{} wins by Bayes factor {} over {}",
3657 winner,
3658 format_bayes_factor(0.5 * margin),
3659 runner_up.name
3660 )
3661 } else {
3662 format!("{winner} (single fit; no comparison)")
3663 };
3664 Ok(RemlComparison {
3665 ranking,
3666 winner,
3667 evidence_summary,
3668 score_table,
3669 })
3670}
3671
3672pub fn format_bayes_factor(log_bf: f64) -> String {
3673 if !log_bf.is_finite() {
3674 return "inf".to_string();
3675 }
3676 if log_bf.abs() >= std::f64::consts::LN_10 * 3.0 {
3677 return format!("1e{:+.1}", log_bf / std::f64::consts::LN_10);
3678 }
3679 format_three_significant(log_bf.exp())
3680}
3681
3682pub fn format_three_significant(value: f64) -> String {
3683 if value == 0.0 {
3684 return "0".to_string();
3685 }
3686 if !value.is_finite() {
3687 return format!("{value}");
3688 }
3689 let exponent = value.abs().log10().floor() as i32;
3690 if exponent >= 3 {
3691 return format!("{value:.2e}");
3692 }
3693 let decimals = (2 - exponent).max(0) as usize;
3694 let scale = 10f64.powi(decimals as i32);
3695 let rounded = (value * scale).abs().round() / scale * value.signum();
3696 format!("{rounded:.decimals$}")
3697}
3698
3699impl Default for TopologySelectOptions {
3700 fn default() -> Self {
3701 Self {
3702 tie_tolerance: 1e-3,
3703 score_scale: TopologyScoreScale::PerObservation,
3704 }
3705 }
3706}
3707
3708pub fn laplace_evidence(
3753 logdet_source: EvidenceLogDetSource<'_>,
3754 penalty_log_det: f64,
3755 residual_objective: f64,
3756 effective_dim: f64,
3757 penalty_rank: f64,
3758) -> f64 {
3759 if !(effective_dim.is_finite() && penalty_rank.is_finite()) {
3760 return f64::NAN;
3761 }
3762 let log_det_h = match evidence_hessian_log_det(logdet_source) {
3763 Ok(v) => v,
3764 Err(_) => return f64::NAN,
3765 };
3766 let null_dim = effective_dim - penalty_rank;
3767 if !null_dim.is_finite() || null_dim < -1e-9 {
3768 return f64::NAN;
3769 }
3770 residual_objective + 0.5 * log_det_h
3771 - 0.5 * penalty_log_det
3772 - 0.5 * null_dim.max(0.0) * (2.0 * std::f64::consts::PI).ln()
3773}
3774
3775pub fn evidence_hessian_log_det(source: EvidenceLogDetSource<'_>) -> Result<f64, String> {
3777 match source {
3778 EvidenceLogDetSource::FactoredArrow {
3779 cache,
3780 fallback_hvp,
3781 } => match arrow_log_det_from_cache(cache) {
3782 Some(v) => Ok(v),
3783 None => match fallback_hvp {
3784 Some(hvp) => hessian_log_det_from_hvp(hvp),
3785 None => {
3786 Err("evidence Hessian logdet requires exact factors or HVP fallback".into())
3787 }
3788 },
3789 },
3790 EvidenceLogDetSource::Hvp(hvp) => hessian_log_det_from_hvp(hvp),
3791 }
3792}
3793
3794pub fn hessian_log_det_from_hvp(hvp: EvidenceHvpLogDet<'_>) -> Result<f64, String> {
3801 if hvp.dim == 0 {
3802 return Ok(0.0);
3803 }
3804 if hvp.dim <= ANALYTIC_LOGDET_DENSE_DIM_THRESHOLD {
3805 let mut dense = Array2::<f64>::zeros((hvp.dim, hvp.dim));
3806 let mut basis = vec![0.0_f64; hvp.dim];
3807 for j in 0..hvp.dim {
3808 basis[j] = 1.0;
3809 let col = (hvp.apply)(&basis);
3810 basis[j] = 0.0;
3811 if col.len() != hvp.dim || col.iter().any(|v| !v.is_finite()) {
3812 return Err(format!(
3813 "evidence HVP logdet expected finite column of length {}, got {}",
3814 hvp.dim,
3815 col.len()
3816 ));
3817 }
3818 for i in 0..hvp.dim {
3819 dense[[i, j]] = col[i];
3820 }
3821 }
3822 validate_dense_hvp_symmetry(&dense)?;
3823 for i in 0..hvp.dim {
3824 for j in (i + 1)..hvp.dim {
3825 let avg = 0.5 * (dense[[i, j]] + dense[[j, i]]);
3826 dense[[i, j]] = avg;
3827 dense[[j, i]] = avg;
3828 }
3829 }
3830 dense_spd_log_det(&dense)
3831 } else {
3832 stochastic_hvp_log_det(hvp)
3833 }
3834}
3835
3836fn dense_spd_log_det(matrix: &Array2<f64>) -> Result<f64, String> {
3837 if matrix.nrows() != matrix.ncols() {
3838 return Err(format!(
3839 "evidence dense logdet requires square matrix, got {}x{}",
3840 matrix.nrows(),
3841 matrix.ncols()
3842 ));
3843 }
3844 if gam_gpu::cuda_selected().map_err(|error| error.to_string())? {
3845 return crate::gpu::reml_gpu::evidence_derivatives_gpu(
3846 crate::gpu::reml_gpu::RemlGpuInput {
3847 penalized_hessian: matrix.view(),
3848 derivative_hessians: Vec::new(),
3849 },
3850 )
3851 .map(|evidence| evidence.logdet_hessian);
3852 }
3853 let (evals, _) = matrix
3854 .eigh(Side::Lower)
3855 .map_err(|e| format!("evidence dense logdet eigendecomposition failed: {e}"))?;
3856 let mut logdet = 0.0_f64;
3857 for (idx, &ev) in evals.iter().enumerate() {
3858 if !ev.is_finite() || ev <= 0.0 {
3859 return Err(format!(
3860 "evidence dense logdet expected SPD Hessian, eigenvalue {idx} is {ev:.3e}"
3861 ));
3862 }
3863 logdet += ev.ln();
3864 }
3865 Ok(logdet)
3866}
3867
3868fn validate_dense_hvp_symmetry(matrix: &Array2<f64>) -> Result<(), String> {
3869 let n = matrix.nrows();
3870 let mut norm_sq = 0.0_f64;
3871 for &value in matrix.iter() {
3872 norm_sq += value * value;
3873 }
3874
3875 let mut skew_sq = 0.0_f64;
3876 for i in 0..n {
3877 for j in (i + 1)..n {
3878 let skew = matrix[[i, j]] - matrix[[j, i]];
3879 skew_sq += 2.0 * skew * skew;
3880 }
3881 }
3882
3883 let rel_skew = skew_sq.sqrt() / norm_sq.sqrt().max(1.0);
3884 if !rel_skew.is_finite() || rel_skew > EVIDENCE_HVP_SYMMETRY_REL_TOL {
3885 return Err(format!(
3886 "evidence HVP logdet requires symmetric operator, relative skew norm is {rel_skew:.3e}"
3887 ));
3888 }
3889 Ok(())
3890}
3891
3892fn validate_hvp_randomized_symmetry(hvp: EvidenceHvpLogDet<'_>) -> Result<(), String> {
3893 let inv_norm = 1.0 / (hvp.dim as f64).sqrt();
3894 for probe in 0..EVIDENCE_HVP_SYMMETRY_PROBES.max(1) {
3895 let mut x = vec![0.0_f64; hvp.dim];
3896 let mut y = vec![0.0_f64; hvp.dim];
3897 rademacher_unit_probe_into_slice(&mut x, (2 * probe) as u64, inv_norm);
3898 rademacher_unit_probe_into_slice(&mut y, (2 * probe + 1) as u64, inv_norm);
3899
3900 let hx = (hvp.apply)(&x);
3901 let hy = (hvp.apply)(&y);
3902 if hx.len() != hvp.dim || hx.iter().any(|v| !v.is_finite()) {
3903 return Err(format!(
3904 "evidence HVP symmetry check expected finite vector of length {}, got {}",
3905 hvp.dim,
3906 hx.len()
3907 ));
3908 }
3909 if hy.len() != hvp.dim || hy.iter().any(|v| !v.is_finite()) {
3910 return Err(format!(
3911 "evidence HVP symmetry check expected finite vector of length {}, got {}",
3912 hvp.dim,
3913 hy.len()
3914 ));
3915 }
3916
3917 let lhs = dot_slice(&x, &hy);
3918 let rhs = dot_slice(&hx, &y);
3919 let scale = (norm2_slice(&hx) * norm2_slice(&y))
3920 .max(norm2_slice(&hy) * norm2_slice(&x))
3921 .max(lhs.abs())
3922 .max(rhs.abs())
3923 .max(1.0);
3924 let rel = (lhs - rhs).abs() / scale;
3925 if !rel.is_finite() || rel > EVIDENCE_HVP_SYMMETRY_REL_TOL {
3926 return Err(format!(
3927 "evidence HVP logdet requires symmetric operator, randomized symmetry probe {probe} has relative bilinear mismatch {rel:.3e}"
3928 ));
3929 }
3930 }
3931 Ok(())
3932}
3933
3934fn stochastic_hvp_log_det(hvp: EvidenceHvpLogDet<'_>) -> Result<f64, String> {
3935 validate_hvp_randomized_symmetry(hvp)?;
3936 let probes = EVIDENCE_LOGDET_SLQ_PROBES.max(1);
3937 let steps = EVIDENCE_LOGDET_LANCZOS_STEPS.min(hvp.dim).max(1);
3938 let inv_norm = 1.0 / (hvp.dim as f64).sqrt();
3939 let mut estimate = 0.0_f64;
3940 for probe in 0..probes {
3941 let mut q0 = vec![0.0_f64; hvp.dim];
3942 rademacher_unit_probe_into_slice(&mut q0, probe as u64, inv_norm);
3943 let quad = lanczos_log_quadrature_hvp(hvp, q0, steps)?;
3944 estimate += hvp.dim as f64 * quad;
3945 }
3946 Ok(estimate / probes as f64)
3947}
3948
3949fn lanczos_log_quadrature_hvp(
3950 hvp: EvidenceHvpLogDet<'_>,
3951 q: Vec<f64>,
3952 max_steps: usize,
3953) -> Result<f64, String> {
3954 let n = hvp.dim;
3955 let eigen = symmetric_lanczos_eigenpairs(
3956 n,
3957 &q,
3958 SymmetricLanczosOptions {
3959 max_steps,
3960 residual_tol: 1e-12,
3961 local_reorthogonalize: false,
3962 full_reorthogonalize: false,
3963 },
3964 |q, out| {
3965 let applied = (hvp.apply)(q);
3966 if applied.len() != n || applied.iter().any(|v| !v.is_finite()) {
3967 return Err(format!(
3968 "evidence HVP SLQ expected finite vector of length {n}, got {}",
3969 applied.len()
3970 ));
3971 }
3972 out.copy_from_slice(&applied);
3973 Ok(())
3974 },
3975 )
3976 .map_err(|e| format!("evidence HVP SLQ Lanczos failed: {e}"))?;
3977 symmetric_lanczos_log_quadrature(&eigen, "evidence HVP SLQ expected SPD Hessian")
3978}
3979
3980#[inline]
3981fn dot_slice(a: &[f64], b: &[f64]) -> f64 {
3982 assert_eq!(a.len(), b.len());
3983 let mut s = 0.0_f64;
3984 for i in 0..a.len() {
3985 s += a[i] * b[i];
3986 }
3987 s
3988}
3989
3990#[inline]
3991fn norm2_slice(a: &[f64]) -> f64 {
3992 dot_slice(a, a).sqrt()
3993}
3994
3995fn rademacher_unit_probe_into_slice(z: &mut [f64], probe: u64, scale: f64) {
3996 let mut state = 0x6A09E667F3BCC909_u64 ^ probe.wrapping_mul(0xD1B54A32D192ED03);
3997 let mut bits = 0_u64;
3998 let mut remaining_bits = 0_u32;
3999 for value in z.iter_mut() {
4000 if remaining_bits == 0 {
4001 bits = splitmix64(&mut state);
4002 remaining_bits = 64;
4003 }
4004 *value = if bits & 1 == 0 { scale } else { -scale };
4005 bits >>= 1;
4006 remaining_bits -= 1;
4007 }
4008}
4009
4010#[inline]
4011const fn splitmix64(state: &mut u64) -> u64 {
4012 gam_linalg::utils::splitmix64(state)
4013}
4014
4015pub fn arrow_log_det_from_cache(cache: &ArrowFactorCache) -> Option<f64> {
4023 if let Some(log_det) = cache.joint_hessian_log_det {
4024 return log_det.is_finite().then_some(log_det);
4025 }
4026 if cache.ridge_t != 0.0 || cache.ridge_beta != 0.0 {
4027 return None;
4028 }
4029 if cache.k > 0 && !cache.schur_factor_is_undamped {
4030 return None;
4031 }
4032 cache.compute_undamped_arrow_log_det()
4033}
4034
4035pub fn ift_du_dbeta(cache: &ArrowFactorCache) -> Array2<f64> {
4045 let n = cache.undamped_factor_count();
4046 let total_len = cache.delta_t_len();
4047 let k = cache.k;
4048 if !cache.htbeta_available() {
4049 return Array2::<f64>::from_elem((total_len, k), f64::NAN);
4050 }
4051 let mut out = Array2::<f64>::zeros((total_len, k));
4052 let mut beta_basis = Array1::<f64>::zeros(k);
4053 let mut rhs = Array1::<f64>::zeros(cache.d);
4055 for i in 0..n {
4056 let di = cache.row_dims[i];
4057 let row_base = cache.row_offsets[i];
4058 let factor = cache.undamped_factor(i);
4059 for col in 0..k {
4061 beta_basis.fill(0.0);
4062 beta_basis[col] = 1.0;
4063 let mut rhs_i = rhs.slice_mut(ndarray::s![..di]).to_owned();
4064 if !cache.apply_htbeta_row(i, beta_basis.view(), &mut rhs_i) {
4067 return Array2::<f64>::from_elem((total_len, k), f64::NAN);
4070 }
4071 let y = cholesky_solve_vector(factor, &rhs_i);
4072 for c in 0..di {
4073 out[[row_base + c, col]] = -y[c];
4074 }
4075 }
4076 }
4077 out
4078}
4079
4080pub fn coupling_components(hessian: ArrayView2<'_, f64>) -> Vec<usize> {
4101 let p = hessian.nrows();
4102 if p == 0 || hessian.ncols() != p {
4103 return Vec::new();
4104 }
4105 let mut parent: Vec<usize> = (0..p).collect();
4107 let mut size: Vec<usize> = vec![1; p];
4108
4109 fn find(parent: &mut [usize], mut x: usize) -> usize {
4110 while parent[x] != x {
4111 parent[x] = parent[parent[x]];
4112 x = parent[x];
4113 }
4114 x
4115 }
4116
4117 for i in 0..p {
4118 for j in (i + 1)..p {
4119 if hessian[[i, j]] != 0.0 || hessian[[j, i]] != 0.0 {
4122 let (ri, rj) = (find(&mut parent, i), find(&mut parent, j));
4123 if ri != rj {
4124 let (small, large) = if size[ri] < size[rj] {
4125 (ri, rj)
4126 } else {
4127 (rj, ri)
4128 };
4129 parent[small] = large;
4130 size[large] += size[small];
4131 }
4132 }
4133 }
4134 }
4135
4136 let mut label_of_root: Vec<Option<usize>> = vec![None; p];
4139 let mut next_label = 0usize;
4140 let mut labels = vec![0usize; p];
4141 for idx in 0..p {
4142 let root = find(&mut parent, idx);
4143 let label = match label_of_root[root] {
4144 Some(l) => l,
4145 None => {
4146 let l = next_label;
4147 label_of_root[root] = Some(l);
4148 next_label += 1;
4149 l
4150 }
4151 };
4152 labels[idx] = label;
4153 }
4154 labels
4155}
4156
4157pub fn cone_of_influence(labels: &[usize], support: &[usize]) -> Vec<usize> {
4168 if support.is_empty() {
4169 return Vec::new();
4170 }
4171 let mut in_cone_labels: Vec<usize> = support
4172 .iter()
4173 .filter_map(|&idx| labels.get(idx).copied())
4174 .collect();
4175 in_cone_labels.sort_unstable();
4176 in_cone_labels.dedup();
4177 if in_cone_labels.is_empty() {
4178 return Vec::new();
4179 }
4180 (0..labels.len())
4181 .filter(|idx| in_cone_labels.binary_search(&labels[*idx]).is_ok())
4182 .collect()
4183}
4184
4185pub fn ift_dbeta_drho(
4197 cache: &ArrowFactorCache,
4198 dg_red_drho: ArrayView2<'_, f64>,
4199) -> Option<Array2<f64>> {
4200 if !cache.schur_factor_is_undamped {
4201 return None;
4202 }
4203 let schur = cache.schur_factor.as_ref()?;
4204 if dg_red_drho.nrows() != cache.k || schur.nrows() != cache.k {
4205 return None;
4206 }
4207 crate::sensitivity::FitSensitivity::from_lower_triangular(schur).mode_response(dg_red_drho)
4208}
4209
4210#[derive(Clone)]
4228pub struct EvidenceIftGradientTerms<'a> {
4229 pub dbeta_drho: ArrayView2<'a, f64>,
4230 pub du_drho: ArrayView2<'a, f64>,
4231 pub value_beta: ArrayView1<'a, f64>,
4232 pub value_u: ArrayView1<'a, f64>,
4233 pub logdet_h_beta: ArrayView1<'a, f64>,
4234 pub logdet_h_u: ArrayView1<'a, f64>,
4235}
4236
4237pub fn evidence_ift_gradient_correction(terms: EvidenceIftGradientTerms<'_>) -> Array1<f64> {
4240 let k = terms.dbeta_drho.nrows();
4241 let nd = terms.du_drho.nrows();
4242 let r = terms.dbeta_drho.ncols();
4243 if terms.du_drho.ncols() != r
4244 || terms.value_beta.len() != k
4245 || terms.logdet_h_beta.len() != k
4246 || terms.value_u.len() != nd
4247 || terms.logdet_h_u.len() != nd
4248 {
4249 return Array1::<f64>::from_elem(r, f64::NAN);
4250 }
4251
4252 let mut out = Array1::<f64>::zeros(r);
4253 for a in 0..r {
4254 let mut acc = 0.0_f64;
4255 for j in 0..k {
4256 let mode = terms.dbeta_drho[[j, a]];
4257 acc += terms.value_beta[j] * mode;
4258 acc += 0.5 * terms.logdet_h_beta[j] * mode;
4259 }
4260 for j in 0..nd {
4261 let mode = terms.du_drho[[j, a]];
4262 acc += terms.value_u[j] * mode;
4263 acc += 0.5 * terms.logdet_h_u[j] * mode;
4264 }
4265 out[a] = acc;
4266 }
4267 out
4268}
4269
4270pub fn evidence_grad_rho(
4300 cache: &ArrowFactorCache,
4301 value_rho: ArrayView1<'_, f64>,
4302 huu_drho: &[Vec<Array2<f64>>],
4303 htbeta_drho: &[Vec<Array2<f64>>],
4304 hbb_drho: &[Array2<f64>],
4305 pen_logdet_drho: ArrayView1<'_, f64>,
4306 ift_terms: EvidenceIftGradientTerms<'_>,
4307) -> Array1<f64> {
4308 let r = value_rho.len();
4309 let n = cache.undamped_factor_count();
4310 let k = cache.k;
4311 let mut out = Array1::<f64>::zeros(r);
4312 if !cache.htbeta_available()
4313 || pen_logdet_drho.len() != r
4314 || huu_drho.len() != n
4315 || htbeta_drho.len() != n
4316 || hbb_drho.len() != r
4317 || huu_drho.iter().any(|row| row.len() != r)
4318 || htbeta_drho.iter().any(|row| row.len() != r)
4319 || hbb_drho.iter().any(|m| m.nrows() != k || m.ncols() != k)
4320 || huu_drho.iter().enumerate().any(|(i, row)| {
4321 let di = cache.row_dims[i];
4322 row.iter().any(|m| m.nrows() != di || m.ncols() != di)
4323 })
4324 || htbeta_drho.iter().enumerate().any(|(i, row)| {
4325 let di = cache.row_dims[i];
4326 row.iter().any(|m| m.nrows() != di || m.ncols() != k)
4327 })
4328 {
4329 out.fill(f64::NAN);
4330 return out;
4331 }
4332 let ift_correction = evidence_ift_gradient_correction(ift_terms);
4333 if ift_correction.len() != r || ift_correction.iter().any(|v| v.is_nan()) {
4334 out.fill(f64::NAN);
4335 return out;
4336 }
4337
4338 let schur = match cache.schur_factor.as_ref() {
4339 Some(s) => s,
4340 None => {
4341 for a in 0..r {
4342 out[a] = f64::NAN;
4343 }
4344 return out;
4345 }
4346 };
4347 if !cache.schur_factor_is_undamped {
4348 for a in 0..r {
4349 out[a] = f64::NAN;
4350 }
4351 return out;
4352 }
4353
4354 let mut y_blocks: Vec<Array2<f64>> = Vec::with_capacity(n);
4357 let mut beta_basis = Array1::<f64>::zeros(k);
4358 let mut rhs = Array1::<f64>::zeros(cache.d);
4360 for i in 0..n {
4361 let di = cache.row_dims[i];
4362 let factor = cache.undamped_factor(i);
4363 let mut yi = Array2::<f64>::zeros((di, k));
4364 for col in 0..k {
4365 beta_basis.fill(0.0);
4366 beta_basis[col] = 1.0;
4367 let mut rhs_i = rhs.slice_mut(ndarray::s![..di]).to_owned();
4368 if !cache.apply_htbeta_row(i, beta_basis.view(), &mut rhs_i) {
4370 out.fill(f64::NAN);
4373 return out;
4374 }
4375 let v = cholesky_solve_vector(factor, &rhs_i);
4376 for c in 0..di {
4377 yi[[c, col]] = v[c];
4378 }
4379 }
4380 y_blocks.push(yi);
4381 }
4382
4383 let mut trace_rhs = Array1::<f64>::zeros(cache.d);
4386 let mut da_tmp = Array2::<f64>::zeros((cache.d, k));
4387 let mut col_scratch = Array1::<f64>::zeros(k);
4388 for a in 0..r {
4389 let mut grad = value_rho[a];
4391
4392 let mut row_trace_acc = 0.0_f64;
4399 for i in 0..n {
4400 let di = cache.row_dims[i];
4401 let m_i = &huu_drho[i][a];
4402 assert_eq!(m_i.shape(), &[di, di]);
4403 for col in 0..di {
4404 let mut tr_rhs_i = trace_rhs.slice_mut(ndarray::s![..di]).to_owned();
4405 for r0 in 0..di {
4406 tr_rhs_i[r0] = m_i[[r0, col]];
4407 }
4408 let v = cholesky_solve_vector(cache.undamped_factor(i), &tr_rhs_i);
4409 row_trace_acc += v[col];
4410 }
4411 }
4412
4413 let mut da = hbb_drho[a].clone();
4422 assert_eq!(da.shape(), &[k, k]);
4423 for i in 0..n {
4424 let di = cache.row_dims[i];
4425 let dhtb = &htbeta_drho[i][a]; let yi = &y_blocks[i]; for r0 in 0..k {
4429 for c0 in 0..k {
4430 let mut acc = 0.0;
4431 for cc in 0..di {
4432 acc += dhtb[[cc, r0]] * yi[[cc, c0]];
4433 }
4434 da[[r0, c0]] -= acc;
4435 }
4436 }
4437 for r0 in 0..k {
4439 for c0 in 0..k {
4440 let mut acc = 0.0;
4441 for cc in 0..di {
4442 acc += yi[[cc, r0]] * dhtb[[cc, c0]];
4443 }
4444 da[[r0, c0]] -= acc;
4445 }
4446 }
4447 let dhuu = &huu_drho[i][a];
4449 let mut da_tmp_i = da_tmp.slice_mut(ndarray::s![..di, ..]).to_owned();
4451 for r0 in 0..di {
4452 for c0 in 0..k {
4453 let mut acc = 0.0;
4454 for cc in 0..di {
4455 acc += dhuu[[r0, cc]] * yi[[cc, c0]];
4456 }
4457 da_tmp_i[[r0, c0]] = acc;
4458 }
4459 }
4460 for r0 in 0..k {
4462 for c0 in 0..k {
4463 let mut acc = 0.0;
4464 for cc in 0..di {
4465 acc += yi[[cc, r0]] * da_tmp_i[[cc, c0]];
4466 }
4467 da[[r0, c0]] += acc;
4468 }
4469 }
4470 }
4471
4472 let mut schur_trace_acc = 0.0_f64;
4474 for j in 0..k {
4475 for r0 in 0..k {
4476 col_scratch[r0] = da[[r0, j]];
4477 }
4478 let v = cholesky_solve_vector(schur, &col_scratch);
4479 schur_trace_acc += v[j];
4480 }
4481
4482 grad += 0.5 * (row_trace_acc + schur_trace_acc);
4483 grad += ift_correction[a];
4484
4485 grad -= 0.5 * pen_logdet_drho[a];
4487
4488 out[a] = grad;
4489 }
4490 out
4491}
4492
4493pub fn select_topology(
4519 candidates: &[TopologyCandidate],
4520 options: TopologySelectOptions,
4521) -> SelectedTopology {
4522 let mut valid: Vec<TopologyCandidate> = candidates
4524 .iter()
4525 .filter(|c| {
4526 c.converged
4527 && c.exclusion_reason.is_none()
4528 && c.negative_log_evidence.is_finite()
4529 && topology_selection_score(c, options.score_scale).is_finite()
4530 })
4531 .cloned()
4532 .collect();
4533 let mut excluded: Vec<TopologyCandidate> = candidates
4534 .iter()
4535 .filter(|c| {
4536 !(c.converged && c.exclusion_reason.is_none() && c.negative_log_evidence.is_finite())
4537 || !topology_selection_score(c, options.score_scale).is_finite()
4538 })
4539 .cloned()
4540 .collect();
4541
4542 assert!(
4543 !valid.is_empty(),
4544 "select_topology: no finite valid candidates; proposal §6.11 forbids silent fallback"
4545 );
4546
4547 valid = rank_priority_candidates(
4552 valid
4553 .into_iter()
4554 .enumerate()
4555 .map(|(idx, row)| {
4556 let score = topology_selection_score(&row, options.score_scale);
4557 let tie_break = usize::from(row.kind.complexity_rank());
4558 PriorityCandidate::new(row, idx, score, tie_break)
4559 })
4560 .collect(),
4561 )
4562 .into_iter()
4563 .map(|row| row.item)
4564 .collect();
4565
4566 let tie = if valid.len() >= 2 {
4568 let top = topology_selection_score(&valid[0], options.score_scale);
4569 let next = topology_selection_score(&valid[1], options.score_scale);
4570 (next - top).abs() <= options.tie_tolerance
4571 } else {
4572 false
4573 };
4574
4575 if tie {
4577 let top_score = topology_selection_score(&valid[0], options.score_scale);
4578 let tied_end = valid
4580 .iter()
4581 .position(|c| {
4582 (topology_selection_score(c, options.score_scale) - top_score).abs()
4583 > options.tie_tolerance
4584 })
4585 .unwrap_or(valid.len());
4586 valid[..tied_end].sort_by_key(|c| c.kind.complexity_rank());
4588 }
4589
4590 let winner = valid[0].kind;
4591 valid.append(&mut excluded);
4592 SelectedTopology {
4593 winner,
4594 ranking: valid,
4595 tie,
4596 }
4597}
4598
4599fn topology_selection_score(candidate: &TopologyCandidate, scale: TopologyScoreScale) -> f64 {
4600 match scale {
4601 TopologyScoreScale::PerObservation => {
4602 if candidate.n_obs == 0 {
4603 f64::NAN
4604 } else {
4605 candidate.negative_log_evidence / candidate.n_obs as f64
4606 }
4607 }
4608 TopologyScoreScale::PerEffectiveDim => {
4609 if !(candidate.effective_dim.is_finite() && candidate.effective_dim > 0.0) {
4610 f64::NAN
4611 } else {
4612 candidate.negative_log_evidence / candidate.effective_dim
4613 }
4614 }
4615 }
4616}
4617
4618pub fn cache_matches_system(cache: &ArrowFactorCache, sys: &ArrowSchurSystem) -> bool {
4626 cache.d == sys.d
4627 && cache.k == sys.k
4628 && cache.n_rows() == sys.rows.len()
4629 && cache.undamped_factor_count() == sys.rows.len()
4630 && cache.manifold_mode_fingerprint == sys.manifold_mode_fingerprint
4631 && cache.row_hessian_fingerprint == sys.current_row_hessian_fingerprint()
4632}
4633
4634#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
4687pub enum HybridAtomParam {
4688 Curved { latent_dim: usize },
4690 Linear,
4692}
4693
4694impl HybridAtomParam {
4695 pub const fn as_str(self) -> &'static str {
4697 match self {
4698 HybridAtomParam::Curved { .. } => "curved",
4699 HybridAtomParam::Linear => "linear",
4700 }
4701 }
4702
4703 pub const fn is_linear(self) -> bool {
4705 matches!(self, HybridAtomParam::Linear)
4706 }
4707}
4708
4709#[derive(Debug, Clone, Copy)]
4717pub struct HybridAtomCandidate {
4718 pub param: HybridAtomParam,
4719 pub negative_log_evidence: f64,
4721 pub num_parameters: usize,
4723 pub fitted_turning: Option<f64>,
4728}
4729
4730impl HybridAtomCandidate {
4731 pub fn linear(negative_log_evidence: f64, num_parameters: usize) -> Self {
4733 Self {
4734 param: HybridAtomParam::Linear,
4735 negative_log_evidence,
4736 num_parameters,
4737 fitted_turning: Some(0.0),
4738 }
4739 }
4740
4741 pub fn curved(
4743 latent_dim: usize,
4744 negative_log_evidence: f64,
4745 num_parameters: usize,
4746 fitted_turning: Option<f64>,
4747 ) -> Self {
4748 Self {
4749 param: HybridAtomParam::Curved { latent_dim },
4750 negative_log_evidence,
4751 num_parameters,
4752 fitted_turning,
4753 }
4754 }
4755}
4756
4757#[derive(Debug, Clone, Copy)]
4761pub struct HybridAtomChoice {
4762 pub param: HybridAtomParam,
4763 pub negative_log_evidence: f64,
4765 pub num_parameters: usize,
4767 pub curved_turning: Option<f64>,
4770 pub curved_evidence_margin: f64,
4775}
4776
4777pub const HYBRID_LINEAR_TURNING_FLOOR: f64 = 1e-9;
4785
4786pub fn select_hybrid_atom(candidates: &[HybridAtomCandidate]) -> Option<HybridAtomChoice> {
4814 if candidates.is_empty() {
4815 return None;
4816 }
4817 let linear = candidates.iter().find(|c| c.param.is_linear());
4818 let curved = candidates.iter().find(|c| !c.param.is_linear());
4819 let curved_turning = curved.and_then(|c| c.fitted_turning);
4820 let curved_evidence_margin = match (linear, curved) {
4821 (Some(l), Some(c)) => l.negative_log_evidence - c.negative_log_evidence,
4822 _ => 0.0,
4823 };
4824
4825 if let (Some(l), Some(turning)) = (linear, curved_turning)
4828 && turning <= HYBRID_LINEAR_TURNING_FLOOR
4829 {
4830 return Some(HybridAtomChoice {
4831 param: l.param,
4832 negative_log_evidence: l.negative_log_evidence,
4833 num_parameters: l.num_parameters,
4834 curved_turning,
4835 curved_evidence_margin,
4836 });
4837 }
4838
4839 let mut best = candidates[0];
4841 for cand in &candidates[1..] {
4842 let better_evidence = cand.negative_log_evidence < best.negative_log_evidence;
4843 let tied = cand.negative_log_evidence == best.negative_log_evidence;
4844 let cheaper_on_tie = tied && cand.num_parameters < best.num_parameters;
4845 if better_evidence || cheaper_on_tie {
4846 best = *cand;
4847 }
4848 }
4849 Some(HybridAtomChoice {
4850 param: best.param,
4851 negative_log_evidence: best.negative_log_evidence,
4852 num_parameters: best.num_parameters,
4853 curved_turning,
4854 curved_evidence_margin,
4855 })
4856}
4857
4858#[derive(Debug, Clone)]
4862pub struct HybridSplitSelection {
4863 pub atoms: Vec<HybridAtomChoice>,
4865 pub total_negative_log_evidence: f64,
4874 pub total_parameters: usize,
4877 pub curved_atom_count: usize,
4879}
4880
4881impl HybridSplitSelection {
4882 pub fn linear_atom_count(&self) -> usize {
4884 self.atoms.len() - self.curved_atom_count
4885 }
4886
4887 pub fn is_pure_linear(&self) -> bool {
4890 self.curved_atom_count == 0 && !self.atoms.is_empty()
4891 }
4892
4893 pub fn is_pure_curved(&self) -> bool {
4896 self.curved_atom_count == self.atoms.len() && !self.atoms.is_empty()
4897 }
4898}
4899
4900pub fn select_hybrid_split(
4914 slots: &[Vec<HybridAtomCandidate>],
4915) -> Result<HybridSplitSelection, String> {
4916 let mut atoms = Vec::with_capacity(slots.len());
4917 let mut total_nle = 0.0_f64;
4918 let mut total_parameters = 0usize;
4919 let mut curved_atom_count = 0usize;
4920 for (i, slot) in slots.iter().enumerate() {
4921 let choice = select_hybrid_atom(slot)
4922 .ok_or_else(|| format!("hybrid split slot {i} has no candidate parameterizations"))?;
4923 if !choice.negative_log_evidence.is_finite() {
4924 return Err(format!(
4925 "hybrid split slot {i} selected a non-finite evidence ({})",
4926 choice.negative_log_evidence
4927 ));
4928 }
4929 if !choice.param.is_linear() {
4930 curved_atom_count += 1;
4931 }
4932 total_nle += choice.negative_log_evidence;
4933 total_parameters += choice.num_parameters;
4934 atoms.push(choice);
4935 }
4936 Ok(HybridSplitSelection {
4937 atoms,
4938 total_negative_log_evidence: total_nle,
4939 total_parameters,
4940 curved_atom_count,
4941 })
4942}
4943
4944#[cfg(test)]
4954mod tests {
4955 use super::*;
4956 use crate::arrow_schur::ArrowFactorSlab;
4957 use ndarray::array;
4958
4959 fn dense_inverse(h: &Array2<f64>) -> Array2<f64> {
4961 let p = h.nrows();
4962 let mut aug = Array2::<f64>::zeros((p, 2 * p));
4963 for i in 0..p {
4964 for j in 0..p {
4965 aug[[i, j]] = h[[i, j]];
4966 }
4967 aug[[i, p + i]] = 1.0;
4968 }
4969 for col in 0..p {
4970 let mut pivot = col;
4971 for row in (col + 1)..p {
4972 if aug[[row, col]].abs() > aug[[pivot, col]].abs() {
4973 pivot = row;
4974 }
4975 }
4976 if pivot != col {
4977 for j in 0..(2 * p) {
4978 aug.swap([col, j], [pivot, j]);
4979 }
4980 }
4981 let d = aug[[col, col]];
4982 for j in 0..(2 * p) {
4983 aug[[col, j]] /= d;
4984 }
4985 for row in 0..p {
4986 if row == col {
4987 continue;
4988 }
4989 let f = aug[[row, col]];
4990 if f != 0.0 {
4991 for j in 0..(2 * p) {
4992 aug[[row, j]] -= f * aug[[col, j]];
4993 }
4994 }
4995 }
4996 }
4997 let mut inv = Array2::<f64>::zeros((p, p));
4998 for i in 0..p {
4999 for j in 0..p {
5000 inv[[i, j]] = aug[[i, p + j]];
5001 }
5002 }
5003 inv
5004 }
5005
5006 #[test]
5007 fn coupling_components_block_diagonal_is_all_singletons_by_block() {
5008 let mut h = Array2::<f64>::eye(4);
5010 h[[0, 1]] = 0.3;
5011 h[[1, 0]] = 0.3;
5012 h[[2, 3]] = 0.7;
5013 h[[3, 2]] = 0.7;
5014 let labels = coupling_components(h.view());
5015 assert_eq!(labels[0], labels[1]);
5016 assert_eq!(labels[2], labels[3]);
5017 assert_ne!(labels[0], labels[2]);
5018 let mut uniq = labels.clone();
5020 uniq.sort_unstable();
5021 uniq.dedup();
5022 assert_eq!(uniq.len(), 2);
5023 }
5024
5025 #[test]
5026 fn coupling_components_fully_coupled_is_one_component() {
5027 let mut h = Array2::<f64>::eye(3);
5028 for i in 0..3 {
5029 for j in 0..3 {
5030 if i != j {
5031 h[[i, j]] = 0.1;
5032 }
5033 }
5034 }
5035 let labels = coupling_components(h.view());
5036 assert!(labels.iter().all(|&l| l == labels[0]));
5037 }
5038
5039 #[test]
5040 fn coupling_components_transitive_chain_merges() {
5041 let mut h = Array2::<f64>::eye(3);
5043 h[[0, 1]] = 0.5;
5044 h[[1, 0]] = 0.5;
5045 h[[1, 2]] = 0.5;
5046 h[[2, 1]] = 0.5;
5047 let labels = coupling_components(h.view());
5048 assert_eq!(labels[0], labels[1]);
5049 assert_eq!(labels[1], labels[2]);
5050 }
5051
5052 #[test]
5053 fn compare_reml_fits_delta_and_bayes_factor_never_contradict_winner_gh1465() {
5054 let cand = |name: &str, score: f64, edf: f64| RemlCandidate {
5066 index: 0,
5067 name: name.to_string(),
5068 score,
5069 edf: Some(edf),
5070 log_lik: Some(0.0),
5071 family: Some("gaussian".to_string()),
5072 n_obs: Some(100),
5073 };
5074 let candidates = vec![
5077 cand("m1", 53.748, 50.0),
5078 cand("m2", 41.605, 51.0),
5079 cand("m3", 120.011, 65.0),
5080 ];
5081 let cmp = compare_reml_fits(candidates).expect("comparison");
5082
5083 assert_eq!(cmp.winner, "m1", "AIC winner");
5084 for row in &cmp.ranking {
5086 assert!(
5087 row.delta >= 0.0,
5088 "ranking delta for {} must be >= 0, got {}",
5089 row.name,
5090 row.delta
5091 );
5092 assert!(
5093 row.bayes_factor >= 1.0 - 1e-12,
5094 "ranking bayes_factor for {} must be >= 1, got {}",
5095 row.name,
5096 row.bayes_factor
5097 );
5098 }
5099 let winner_row = cmp.ranking.iter().find(|r| r.name == "m1").unwrap();
5100 assert!(winner_row.delta.abs() < 1e-12, "winner delta == 0");
5101 assert!(
5102 (winner_row.bayes_factor - 1.0).abs() < 1e-9,
5103 "winner bayes_factor == 1"
5104 );
5105
5106 for row in &cmp.score_table {
5109 assert!(
5110 row.delta_reml >= 0.0,
5111 "score-table delta_reml for {} must be >= 0, got {}",
5112 row.name,
5113 row.delta_reml
5114 );
5115 assert!(
5116 row.bayes_factor_best_over_model >= 1.0 - 1e-12,
5117 "score-table bayes_factor for {} must be >= 1, got {}",
5118 row.name,
5119 row.bayes_factor_best_over_model
5120 );
5121 }
5122 let m2 = cmp.score_table.iter().find(|r| r.name == "m2").unwrap();
5124 assert!(
5125 m2.delta_reml.abs() < 1e-12,
5126 "the minimum-raw-REML row has delta_reml 0"
5127 );
5128 }
5129
5130 #[test]
5131 fn cone_of_influence_empty_support_is_empty() {
5132 let labels = vec![0usize, 0, 1, 1];
5133 assert!(cone_of_influence(&labels, &[]).is_empty());
5134 }
5135
5136 #[test]
5137 fn cone_of_influence_returns_full_component() {
5138 let labels = vec![0usize, 0, 1, 1];
5139 assert_eq!(cone_of_influence(&labels, &[0]), vec![0, 1]);
5141 assert_eq!(cone_of_influence(&labels, &[1, 2]), vec![0, 1, 2, 3]);
5143 }
5144
5145 #[test]
5146 fn coned_matches_full_solve_on_fully_coupled_hessian() {
5147 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])
5150 .unwrap();
5151 let inv = dense_inverse(&h);
5152 let mut dg = Array2::<f64>::zeros((3, 2));
5154 dg[[0, 0]] = 1.3;
5155 dg[[2, 1]] = -0.7;
5156 let supports = vec![0..1usize, 2..3usize];
5157
5158 let eye: Array2<f64> = Array2::eye(3);
5159 let op = crate::sensitivity::FitSensitivity::from_projected(&eye, &inv);
5160 let full = op.mode_response(dg.view()).unwrap();
5161 let coned = op
5162 .mode_response_coned(h.view(), dg.view(), &supports)
5163 .unwrap();
5164 for i in 0..3 {
5165 for a in 0..2 {
5166 assert!(
5167 (full[[i, a]] - coned[[i, a]]).abs() < 1e-12,
5168 "fully-coupled mismatch at ({i},{a}): {} vs {}",
5169 full[[i, a]],
5170 coned[[i, a]]
5171 );
5172 }
5173 }
5174 }
5175
5176 #[test]
5177 fn coned_confines_to_component_on_decoupled_hessian() {
5178 let mut h = Array2::<f64>::zeros((4, 4));
5182 h[[0, 0]] = 4.0;
5184 h[[1, 1]] = 3.0;
5185 h[[0, 1]] = 1.0;
5186 h[[1, 0]] = 1.0;
5187 h[[2, 2]] = 2.0;
5189 h[[3, 3]] = 5.0;
5190 h[[2, 3]] = 0.6;
5191 h[[3, 2]] = 0.6;
5192 let inv = dense_inverse(&h);
5193
5194 let mut dg = Array2::<f64>::zeros((4, 1));
5195 dg[[0, 0]] = 0.9;
5196 dg[[1, 0]] = -0.4;
5197 let support_range = 0..2usize;
5198 let supports = std::slice::from_ref(&support_range);
5199
5200 let eye: Array2<f64> = Array2::eye(4);
5201 let coned = crate::sensitivity::FitSensitivity::from_projected(&eye, &inv)
5202 .mode_response_coned(h.view(), dg.view(), supports)
5203 .unwrap();
5204 let q = dg.column(0).to_owned();
5207 let exact = inv.dot(&q).mapv(|v| -v);
5208 for i in 0..4 {
5209 assert!(
5210 (coned[[i, 0]] - exact[[i]]).abs() < 1e-12,
5211 "decoupled mismatch at {i}: {} vs {}",
5212 coned[[i, 0]],
5213 exact[[i]]
5214 );
5215 }
5216 assert_eq!(coned[[2, 0]], 0.0);
5218 assert_eq!(coned[[3, 0]], 0.0);
5219 }
5220
5221 #[test]
5222 fn coned_skips_inactive_column_with_empty_support() {
5223 let h = Array2::<f64>::eye(2);
5224 let dg = Array2::<f64>::zeros((2, 1));
5225 let empty_support = 0..0usize;
5227 let supports = std::slice::from_ref(&empty_support);
5228 let eye: Array2<f64> = Array2::eye(2);
5233 let nan_inv = Array2::<f64>::from_elem((2, 2), f64::NAN);
5234 let coned = crate::sensitivity::FitSensitivity::from_projected(&eye, &nan_inv)
5235 .mode_response_coned(h.view(), dg.view(), supports)
5236 .unwrap();
5237 assert_eq!(coned[[0, 0]], 0.0);
5238 assert_eq!(coned[[1, 0]], 0.0);
5239 }
5240
5241 fn make_minimal_cache() -> ArrowFactorCache {
5242 let l_huu = Array2::from_shape_vec((1, 1), vec![std::f64::consts::SQRT_2]).unwrap();
5245 let l_schur = Array2::from_shape_vec((1, 1), vec![(1.875_f64).sqrt()]).unwrap();
5246 let htbeta = Array2::from_shape_vec((1, 1), vec![0.5]).unwrap();
5247 let mut cache = ArrowFactorCache {
5248 htt_factors: ArrowFactorSlab::from_blocks(vec![l_huu]),
5249 htt_factors_undamped: crate::arrow_schur::ArrowUndampedFactors::SameAsDamped,
5250 schur_factor: Some(l_schur),
5251 schur_factor_is_undamped: true,
5252 beta_schur_deflation: None,
5253 joint_hessian_log_det: None,
5254 solver_mode: crate::arrow_schur::ArrowSolverMode::Direct,
5255 ridge_t: 0.0,
5256 ridge_beta: 0.0,
5257 htbeta: crate::arrow_schur::ArrowHtbetaCache::Dense {
5258 blocks: std::sync::Arc::from(vec![htbeta]),
5259 estimated_bytes: std::mem::size_of::<f64>(),
5260 },
5261 d: 1,
5262 row_dims: std::sync::Arc::from(vec![1usize]),
5263 row_offsets: std::sync::Arc::from(vec![0usize, 1usize]),
5264 k: 1,
5265 manifold_mode_fingerprint: 0,
5266 row_hessian_fingerprint: 0,
5267 pcg_diagnostics: crate::arrow_schur::ArrowPcgDiagnostics::default(),
5268 gauge_deflated_directions: 0,
5269 deflated_row_directions: std::sync::Arc::from(Vec::new()),
5270 deflation_row_spectra: std::sync::Arc::from(Vec::new()),
5271 beta_gauge_quotient: None,
5272 };
5273 cache.joint_hessian_log_det = cache.compute_undamped_arrow_log_det();
5274 cache
5275 }
5276
5277 #[test]
5278 fn laplace_evidence_returns_finite_for_minimal_cache() {
5279 let cache = make_minimal_cache();
5280 let v = laplace_evidence(
5283 EvidenceLogDetSource::FactoredArrow {
5284 cache: &cache,
5285 fallback_hvp: None,
5286 },
5287 0.0,
5288 0.0,
5289 2.0,
5290 1.0,
5291 );
5292 assert!(v.is_finite());
5293 let expected =
5294 0.5 * (2.0_f64.ln() + 1.875_f64.ln()) - 0.5 * (2.0 * std::f64::consts::PI).ln();
5295 assert!((v - expected).abs() < 1e-12);
5296 }
5297
5298 fn k0_direct_cache_no_schur(latent_diag: f64) -> ArrowFactorCache {
5307 let l_huu = Array2::from_shape_vec((1, 1), vec![latent_diag.sqrt()]).unwrap();
5308 let mut cache = ArrowFactorCache {
5309 htt_factors: ArrowFactorSlab::from_blocks(vec![l_huu]),
5310 htt_factors_undamped: crate::arrow_schur::ArrowUndampedFactors::SameAsDamped,
5311 schur_factor: None,
5312 schur_factor_is_undamped: true,
5313 beta_schur_deflation: None,
5314 joint_hessian_log_det: None,
5315 solver_mode: crate::arrow_schur::ArrowSolverMode::Direct,
5316 ridge_t: 0.0,
5317 ridge_beta: 0.0,
5318 htbeta: crate::arrow_schur::ArrowHtbetaCache::Disabled { estimated_bytes: 0 },
5319 d: 1,
5320 row_dims: std::sync::Arc::from(vec![1usize]),
5321 row_offsets: std::sync::Arc::from(vec![0usize, 1usize]),
5322 k: 0,
5323 manifold_mode_fingerprint: 0,
5324 row_hessian_fingerprint: 0,
5325 pcg_diagnostics: crate::arrow_schur::ArrowPcgDiagnostics::default(),
5326 gauge_deflated_directions: 0,
5327 deflated_row_directions: std::sync::Arc::from(Vec::new()),
5328 deflation_row_spectra: std::sync::Arc::from(Vec::new()),
5329 beta_gauge_quotient: None,
5330 };
5331 cache.joint_hessian_log_det = cache.compute_undamped_arrow_log_det();
5332 cache
5333 }
5334
5335 #[test]
5336 fn arrow_log_det_some_for_k0_direct_cache_without_schur() {
5337 let cache = k0_direct_cache_no_schur(3.0);
5338 let log_det = arrow_log_det_from_cache(&cache)
5339 .expect("k==0 Direct cache must yield Some(per-row sum), not None (#1132)");
5340 assert!(
5342 (log_det - 3.0_f64.ln()).abs() < 1e-12,
5343 "log_det = {log_det}"
5344 );
5345 let cached = cache
5347 .compute_undamped_arrow_log_det()
5348 .expect("compute_undamped_arrow_log_det must be Some for k==0");
5349 assert!((cached - 3.0_f64.ln()).abs() < 1e-12, "cached = {cached}");
5350 }
5351
5352 #[test]
5353 fn arrow_log_det_none_for_kpos_cache_without_schur() {
5354 let mut cache = k0_direct_cache_no_schur(3.0);
5357 cache.k = 1;
5358 cache.solver_mode = crate::arrow_schur::ArrowSolverMode::InexactPCG;
5359 cache.joint_hessian_log_det = None;
5360 assert!(arrow_log_det_from_cache(&cache).is_none());
5361 assert!(cache.compute_undamped_arrow_log_det().is_none());
5362 }
5363
5364 #[test]
5365 fn laplace_evidence_nan_when_authoritative_logdet_missing() {
5366 let mut cache = make_minimal_cache();
5367 cache.ridge_t = 1e-3;
5368 cache.joint_hessian_log_det = None;
5369 assert!(
5370 laplace_evidence(
5371 EvidenceLogDetSource::FactoredArrow {
5372 cache: &cache,
5373 fallback_hvp: None,
5374 },
5375 0.0,
5376 0.0,
5377 2.0,
5378 1.0,
5379 )
5380 .is_nan()
5381 );
5382 }
5383
5384 #[test]
5385 fn laplace_evidence_uses_hvp_fallback_without_authoritative_logdet() {
5386 let mut cache = make_minimal_cache();
5387 cache.schur_factor = None;
5388 cache.joint_hessian_log_det = None;
5389 let hvp = |x: &[f64]| -> Vec<f64> { vec![2.0 * x[0], 1.875 * x[1]] };
5390 let v = laplace_evidence(
5391 EvidenceLogDetSource::FactoredArrow {
5392 cache: &cache,
5393 fallback_hvp: Some(EvidenceHvpLogDet {
5394 dim: 2,
5395 apply: &hvp,
5396 }),
5397 },
5398 0.0,
5399 0.0,
5400 2.0,
5401 1.0,
5402 );
5403 let expected =
5404 0.5 * (2.0_f64.ln() + 1.875_f64.ln()) - 0.5 * (2.0 * std::f64::consts::PI).ln();
5405 assert!((v - expected).abs() < 1e-12);
5406 }
5407
5408 #[test]
5409 fn ift_du_dbeta_has_expected_shape() {
5410 let cache = make_minimal_cache();
5411 let du_db = ift_du_dbeta(&cache);
5412 assert_eq!(du_db.shape(), &[1, 1]);
5413 assert!((du_db[[0, 0]] - (-0.25)).abs() < 1e-12);
5415 }
5416
5417 #[test]
5418 fn ift_dbeta_drho_returns_some_for_direct_cache() {
5419 let cache = make_minimal_cache();
5420 let q = Array2::from_shape_vec((1, 1), vec![1.0]).unwrap();
5421 let out = ift_dbeta_drho(&cache, q.view()).unwrap();
5422 assert_eq!(out.shape(), &[1, 1]);
5423 assert!((out[[0, 0]] + 1.0 / 1.875).abs() < 1e-12);
5425 }
5426
5427 #[test]
5428 fn topology_select_picks_lowest_negative_log_evidence() {
5429 let candidates = vec![
5430 TopologyCandidate {
5431 kind: TopologyKind::Flat,
5432 negative_log_evidence: 10.0,
5433 effective_dim: 4.0,
5434 n_obs: 100,
5435 converged: true,
5436 exclusion_reason: None,
5437 },
5438 TopologyCandidate {
5439 kind: TopologyKind::Sphere,
5440 negative_log_evidence: 8.0,
5441 effective_dim: 5.0,
5442 n_obs: 100,
5443 converged: true,
5444 exclusion_reason: None,
5445 },
5446 TopologyCandidate {
5447 kind: TopologyKind::Torus,
5448 negative_log_evidence: f64::NAN,
5449 effective_dim: 6.0,
5450 n_obs: 100,
5451 converged: false,
5452 exclusion_reason: Some("torus periods missing".to_string()),
5453 },
5454 ];
5455 let sel = select_topology(&candidates, TopologySelectOptions::default());
5456 assert_eq!(sel.winner, TopologyKind::Sphere);
5457 assert!(!sel.tie);
5458 }
5459
5460 #[test]
5461 fn topology_select_tie_breaks_to_simpler() {
5462 let candidates = vec![
5463 TopologyCandidate {
5464 kind: TopologyKind::Sphere,
5465 negative_log_evidence: 5.0,
5466 effective_dim: 5.0,
5467 n_obs: 100,
5468 converged: true,
5469 exclusion_reason: None,
5470 },
5471 TopologyCandidate {
5472 kind: TopologyKind::Flat,
5473 negative_log_evidence: 5.0 + 1e-6,
5474 effective_dim: 4.0,
5475 n_obs: 100,
5476 converged: true,
5477 exclusion_reason: None,
5478 },
5479 ];
5480 let sel = select_topology(&candidates, TopologySelectOptions::default());
5481 assert_eq!(sel.winner, TopologyKind::Flat);
5482 assert!(sel.tie);
5483 }
5484
5485 fn gaussian_logpdf(y: f64, mean: f64, sd: f64) -> f64 {
5486 let z = (y - mean) / sd;
5487 -0.5 * (2.0 * std::f64::consts::PI).ln() - sd.ln() - 0.5 * z * z
5488 }
5489
5490 #[test]
5491 fn stacking_single_candidate_gets_full_weight() {
5492 let log_density = Array2::from_shape_vec((3, 1), vec![-1.0, -2.0, -0.5]).unwrap();
5493 let out = solve_stacking_weights(log_density.view(), StackingConfig::default()).unwrap();
5494 assert!((out.weights[0] - 1.0).abs() < 1e-12);
5495 assert_eq!(out.weights.len(), 1);
5496 }
5497
5498 #[test]
5499 fn stacking_dominant_candidate_attracts_nearly_all_weight() {
5500 let mut log_density = Array2::<f64>::zeros((50, 2));
5501 for i in 0..50 {
5502 log_density[[i, 0]] = -0.1;
5503 log_density[[i, 1]] = -5.0;
5504 }
5505 let out = solve_stacking_weights(log_density.view(), StackingConfig::default()).unwrap();
5506 assert!(out.weights[0] > 0.99, "w0 = {}", out.weights[0]);
5507 assert!(out.weights[1] < 0.01, "w1 = {}", out.weights[1]);
5508 }
5509
5510 #[test]
5511 fn stacking_complementary_candidates_share_weight() {
5512 let n = 40;
5515 let mut log_density = Array2::<f64>::zeros((n, 2));
5516 for i in 0..n {
5517 if i < n / 2 {
5518 log_density[[i, 0]] = gaussian_logpdf(0.0, 0.0, 0.5);
5519 log_density[[i, 1]] = gaussian_logpdf(0.0, 1.5, 0.5);
5520 } else {
5521 log_density[[i, 0]] = gaussian_logpdf(0.0, 1.5, 0.5);
5522 log_density[[i, 1]] = gaussian_logpdf(0.0, 0.0, 0.5);
5523 }
5524 }
5525 let out = solve_stacking_weights(log_density.view(), StackingConfig::default()).unwrap();
5526 assert!(
5527 out.weights[0] > 0.2 && out.weights[0] < 0.8,
5528 "w0 = {}",
5529 out.weights[0]
5530 );
5531 assert!((out.weights.sum() - 1.0).abs() < 1e-9);
5532 }
5533
5534 #[test]
5535 fn stacking_weights_stay_on_the_simplex() {
5536 let log_density = Array2::from_shape_vec(
5537 (3, 3),
5538 vec![-1.0, -2.0, -3.0, -2.5, -1.0, -2.0, -3.0, -2.0, -1.0],
5539 )
5540 .unwrap();
5541 let out = solve_stacking_weights(log_density.view(), StackingConfig::default()).unwrap();
5542 assert!((out.weights.sum() - 1.0).abs() < 1e-9);
5543 assert!(out.weights.iter().all(|&w| w >= -1e-12));
5544 }
5545
5546 #[test]
5547 fn stacking_solution_satisfies_the_simplex_kkt_certificate() {
5548 let log_density = Array2::from_shape_vec(
5553 (5, 2),
5554 vec![-0.2, -3.0, -3.0, -0.2, -0.5, -1.5, -1.5, -0.5, -0.1, -2.0],
5555 )
5556 .unwrap();
5557 let config = StackingConfig::default();
5558 let out = solve_stacking_weights(log_density.view(), config).unwrap();
5559 assert!(out.certificate.residual() <= config.kkt_tol);
5560 let n = log_density.nrows();
5561 for k in 0..2 {
5562 let mut g = 0.0_f64;
5563 for i in 0..n {
5564 let mix: f64 = (0..2)
5565 .map(|c| out.weights[c] * log_density[[i, c]].exp())
5566 .sum();
5567 g += log_density[[i, k]].exp() / mix;
5568 }
5569 g /= n as f64;
5570 assert!(
5571 g <= 1.0 + config.kkt_tol,
5572 "stationarity violated for candidate {k}: g = {g}"
5573 );
5574 assert!(
5575 out.weights[k] * (g - 1.0).abs() <= config.kkt_tol * (1.0 + 1e-6),
5576 "complementary slackness violated for candidate {k}: w = {}, g = {g}",
5577 out.weights[k]
5578 );
5579 }
5580 }
5581
5582 #[test]
5583 fn stacking_exhaustion_without_certificate_is_an_error_not_weights() {
5584 let log_density = Array2::from_shape_vec(
5585 (6, 3),
5586 vec![
5587 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,
5588 -0.5, -1.0, -0.3, -2.0,
5589 ],
5590 )
5591 .unwrap();
5592 let config = StackingConfig {
5593 max_iter: 1,
5594 ..StackingConfig::default()
5595 };
5596 let err = solve_stacking_weights(log_density.view(), config).unwrap_err();
5597 let checkpoint = match err {
5598 StackingError::DidNotConverge {
5599 certificate,
5600 checkpoint,
5601 ..
5602 } => {
5603 assert!(certificate.residual() > config.kkt_tol);
5604 assert_eq!(checkpoint.completed_iterations, 1);
5605 checkpoint
5606 }
5607 other => panic!("expected typed stacking exhaustion, got {other}"),
5608 };
5609 let encoded = serde_json::to_string(&checkpoint).unwrap();
5610 let checkpoint: StackingCheckpoint = serde_json::from_str(&encoded).unwrap();
5611 let mut other_density = log_density.clone();
5612 other_density[[0, 0]] += 0.25;
5613 assert!(matches!(
5614 resume_stacking_weights(other_density.view(), StackingConfig::default(), &checkpoint,),
5615 Err(StackingError::InvalidInput { .. })
5616 ));
5617 let resumed =
5618 resume_stacking_weights(log_density.view(), StackingConfig::default(), &checkpoint)
5619 .unwrap();
5620 let uninterrupted =
5621 solve_stacking_weights(log_density.view(), StackingConfig::default()).unwrap();
5622 for (resumed, uninterrupted) in resumed.weights.iter().zip(uninterrupted.weights.iter()) {
5623 assert!((resumed - uninterrupted).abs() <= 1.0e-10);
5624 }
5625 }
5626
5627 #[test]
5628 fn stacking_near_tied_boundary_uses_newton_not_millions_of_em_steps() {
5629 let log_density =
5630 Array2::from_shape_fn(
5631 (64, 2),
5632 |(_, candidate)| {
5633 if candidate == 0 { 0.0 } else { -1.0e-6 }
5634 },
5635 );
5636 let out = solve_stacking_weights(log_density.view(), StackingConfig::default()).unwrap();
5637 assert!(out.weights[0] >= 1.0 - StackingConfig::default().kkt_tol);
5638 assert!(out.iterations < 8, "iterations = {}", out.iterations);
5639 }
5640
5641 #[test]
5642 fn stacking_dead_candidate_column_gets_zero_weight() {
5643 let log_density = Array2::from_shape_vec(
5644 (3, 2),
5645 vec![
5646 -1.0,
5647 f64::NEG_INFINITY,
5648 -2.0,
5649 f64::NEG_INFINITY,
5650 -0.5,
5651 f64::NEG_INFINITY,
5652 ],
5653 )
5654 .unwrap();
5655 let out = solve_stacking_weights(log_density.view(), StackingConfig::default()).unwrap();
5656 assert_eq!(out.weights[1], 0.0);
5657 assert!((out.weights[0] - 1.0).abs() < 1e-12);
5658 }
5659
5660 #[test]
5661 fn stacking_rejects_invalid_and_unscorable_rows() {
5662 let log_density = Array2::from_shape_vec(
5663 (3, 2),
5664 vec![-1.0, -2.0, f64::NAN, f64::NEG_INFINITY, -2.0, -1.0],
5665 )
5666 .unwrap();
5667 assert!(matches!(
5668 solve_stacking_weights(log_density.view(), StackingConfig::default()),
5669 Err(StackingError::InvalidInput { .. })
5670 ));
5671 let unscorable = Array2::from_shape_vec(
5672 (2, 2),
5673 vec![-1.0, -2.0, f64::NEG_INFINITY, f64::NEG_INFINITY],
5674 )
5675 .unwrap();
5676 assert!(matches!(
5677 solve_stacking_weights(unscorable.view(), StackingConfig::default()),
5678 Err(StackingError::InvalidInput { .. })
5679 ));
5680 }
5681
5682 fn two_cluster_mixture_data() -> Array2<f64> {
5683 Array2::from_shape_vec(
5684 (12, 1),
5685 vec![
5686 -2.2, -2.0, -1.9, -2.1, -1.8, -2.05, 1.8, 2.0, 2.2, 1.9, 2.1, 2.05,
5687 ],
5688 )
5689 .unwrap()
5690 }
5691
5692 #[test]
5693 fn gaussian_mixture_monotonicity_resolves_composite_map_noise_2264() {
5694 let objective_scale = 1.0;
5695 let composite_resolution = f64::EPSILON.sqrt() * objective_scale;
5696 let uncertainty = gaussian_mixture_monotonicity_uncertainty(objective_scale, 0.0, 0.0);
5697 assert_eq!(uncertainty, composite_resolution);
5698
5699 let noise_scale_decrease = -0.5 * composite_resolution;
5700 assert!(noise_scale_decrease >= -uncertainty);
5701 let resolved_decrease = -2.0 * composite_resolution;
5702 assert!(resolved_decrease < -uncertainty);
5703
5704 let larger_reduction_bound = 2.0 * composite_resolution;
5705 assert_eq!(
5706 gaussian_mixture_monotonicity_uncertainty(objective_scale, larger_reduction_bound, 0.0),
5707 larger_reduction_bound,
5708 );
5709 }
5710
5711 #[test]
5712 fn gaussian_mixture_issue_scale_negative_step_is_within_computed_uncertainty_2264() {
5713 let objective_scale = 1.0;
5720 let recorded_step = -1.4e-13;
5721 let uncertainty = gaussian_mixture_monotonicity_uncertainty(objective_scale, 0.0, 0.0);
5722 let certificate = GaussianMixtureCertificate {
5723 mean_log_likelihood: -objective_scale,
5724 mean_log_likelihood_gain: recorded_step,
5725 monotonicity_uncertainty: uncertainty,
5726 objective_residual: recorded_step.abs() / objective_scale,
5727 objective_tolerance: f64::EPSILON.sqrt(),
5728 parameter_residual: 0.0,
5729 parameter_tolerance: f64::EPSILON.sqrt(),
5730 };
5731
5732 assert_eq!(
5733 certificate.monotonicity_uncertainty,
5734 f64::EPSILON.sqrt() * objective_scale,
5735 "reported uncertainty must be the computed composite-map resolution"
5736 );
5737 assert!(
5738 certificate.mean_log_likelihood_gain >= -certificate.monotonicity_uncertainty,
5739 "the recorded noise-scale decrease must not be a monotonicity violation"
5740 );
5741 }
5742
5743 #[test]
5744 fn gaussian_mixture_below_roundoff_positive_gain_can_certify_2264() {
5745 let objective_scale = 1.0;
5750 let recorded_gain = 6.6e-15;
5751 let recorded_reduction_bound = 1.5e-14;
5752 let objective_tolerance = f64::EPSILON.sqrt();
5753 let parameter_tolerance = f64::EPSILON.sqrt();
5754 let uncertainty = gaussian_mixture_monotonicity_uncertainty(
5755 objective_scale,
5756 recorded_reduction_bound,
5757 0.0,
5758 );
5759 let certificate = GaussianMixtureCertificate {
5760 mean_log_likelihood: -objective_scale,
5761 mean_log_likelihood_gain: recorded_gain,
5762 monotonicity_uncertainty: uncertainty,
5763 objective_residual: recorded_gain / objective_scale,
5764 objective_tolerance,
5765 parameter_residual: 0.5 * parameter_tolerance,
5766 parameter_tolerance,
5767 };
5768
5769 assert_eq!(
5770 certificate.monotonicity_uncertainty,
5771 (f64::EPSILON.sqrt() * objective_scale).max(recorded_reduction_bound),
5772 "reported uncertainty must come from the composite-map and reduction bounds"
5773 );
5774 assert!(certificate.mean_log_likelihood_gain >= -certificate.monotonicity_uncertainty);
5775 assert!(certificate.objective_residual <= certificate.objective_tolerance);
5776 assert!(certificate.parameter_residual <= certificate.parameter_tolerance);
5777 }
5778
5779 #[test]
5780 fn gaussian_mixture_parameter_map_uses_component_measure_geometry_2324() {
5781 let tolerance = f64::EPSILON.sqrt();
5787 let raw_coordinate_step: f64 = 4.6e-8;
5788 assert!(raw_coordinate_step > tolerance);
5789 let weights = array![0.25, 0.75];
5790 let previous_means = array![[0.0], [2.0]];
5791 let next_means = array![[raw_coordinate_step], [2.0]];
5792 let covariance = vec![array![[1.0]], array![[1.0]]];
5793 let residual = mixture_parameter_residual(
5794 &weights,
5795 &previous_means,
5796 &covariance,
5797 &weights,
5798 &next_means,
5799 &covariance,
5800 );
5801 assert_eq!(residual, weights[0] * raw_coordinate_step);
5802 assert!(residual <= tolerance);
5803
5804 let next_weights = array![
5807 weights[0] + raw_coordinate_step,
5808 weights[1] - raw_coordinate_step
5809 ];
5810 let mass_residual = mixture_parameter_residual(
5811 &weights,
5812 &previous_means,
5813 &covariance,
5814 &next_weights,
5815 &previous_means,
5816 &covariance,
5817 );
5818 assert!(mass_residual > tolerance);
5819 }
5820
5821 #[test]
5822 fn gaussian_mixture_fit_certificate_describes_the_exact_returned_iterate() {
5823 let data = two_cluster_mixture_data();
5824 let config = GaussianMixtureConfig::default();
5825 let fit = fit_gaussian_mixture(data.view(), 2, config).unwrap();
5826 let certificate = fit.certificate();
5827 assert!(certificate.objective_residual <= certificate.objective_tolerance);
5828 assert!(certificate.parameter_residual <= certificate.parameter_tolerance);
5829
5830 let checkpoint = GaussianMixtureCheckpoint {
5831 weights: fit.weights.clone(),
5832 means: fit.means.clone(),
5833 covariances: fit.covariances.clone(),
5834 mean_log_likelihood: certificate.mean_log_likelihood,
5835 completed_iterations: fit.iterations,
5836 data_fingerprint: mixture_data_fingerprint(data.view()),
5837 covariance_floor: config.covariance_floor,
5838 };
5839 let current = mixture_e_step(
5840 data.view(),
5841 &checkpoint.weights,
5842 &checkpoint.means,
5843 &checkpoint.covariances,
5844 )
5845 .unwrap();
5846 let (weights, means, covariances) = mixture_m_step(
5847 data.view(),
5848 current.responsibilities.view(),
5849 config.covariance_floor,
5850 )
5851 .unwrap();
5852 let residual = mixture_parameter_residual(
5853 &checkpoint.weights,
5854 &checkpoint.means,
5855 &checkpoint.covariances,
5856 &weights,
5857 &means,
5858 &covariances,
5859 );
5860 let next = mixture_e_step(data.view(), &weights, &means, &covariances).unwrap();
5861 assert!(residual <= config.parameter_tol);
5862 assert_eq!(certificate.mean_log_likelihood, current.mean_log_likelihood);
5863 assert_eq!(
5864 certificate.mean_log_likelihood_gain,
5865 next.mean_log_likelihood - current.mean_log_likelihood
5866 );
5867 assert_eq!(
5868 certificate.monotonicity_uncertainty,
5869 gaussian_mixture_monotonicity_uncertainty(
5870 current
5871 .mean_log_likelihood
5872 .abs()
5873 .max(next.mean_log_likelihood.abs())
5874 .max(1.0),
5875 current.mean_log_likelihood_roundoff,
5876 next.mean_log_likelihood_roundoff,
5877 )
5878 );
5879 assert_eq!(certificate.parameter_residual, residual);
5880 assert!(
5881 (next.mean_log_likelihood - current.mean_log_likelihood).abs()
5882 / current
5883 .mean_log_likelihood
5884 .abs()
5885 .max(next.mean_log_likelihood.abs())
5886 .max(1.0)
5887 <= config.loglik_tol
5888 );
5889 }
5890
5891 #[test]
5892 fn gaussian_mixture_exhaustion_is_typed_and_resumable() {
5893 let data = two_cluster_mixture_data();
5894 let short = GaussianMixtureConfig {
5895 max_iter: 1,
5896 ..GaussianMixtureConfig::default()
5897 };
5898 let err = fit_gaussian_mixture(data.view(), 2, short).unwrap_err();
5899 let checkpoint = match err {
5900 GaussianMixtureError::DidNotConverge {
5901 certificate,
5902 checkpoint,
5903 ..
5904 } => {
5905 assert!(
5906 certificate.objective_residual > short.loglik_tol
5907 || certificate.parameter_residual > short.parameter_tol
5908 );
5909 assert_eq!(checkpoint.completed_iterations, 1);
5910 let at_checkpoint = mixture_e_step(
5911 data.view(),
5912 &checkpoint.weights,
5913 &checkpoint.means,
5914 &checkpoint.covariances,
5915 )
5916 .unwrap();
5917 assert_eq!(
5918 certificate.mean_log_likelihood, at_checkpoint.mean_log_likelihood,
5919 "exhaustion evidence and checkpoint must describe one iterate"
5920 );
5921 checkpoint
5922 }
5923 other => panic!("expected typed EM exhaustion, got {other}"),
5924 };
5925 let encoded = serde_json::to_string(&checkpoint).unwrap();
5926 let checkpoint: GaussianMixtureCheckpoint = serde_json::from_str(&encoded).unwrap();
5927 let mut other_data = data.clone();
5928 other_data[[0, 0]] += 0.01;
5929 assert!(matches!(
5930 resume_gaussian_mixture(
5931 other_data.view(),
5932 GaussianMixtureConfig::default(),
5933 checkpoint.clone(),
5934 ),
5935 Err(GaussianMixtureError::InvalidInput { .. })
5936 ));
5937 let resumed =
5938 resume_gaussian_mixture(data.view(), GaussianMixtureConfig::default(), checkpoint)
5939 .unwrap();
5940 let uninterrupted =
5941 fit_gaussian_mixture(data.view(), 2, GaussianMixtureConfig::default()).unwrap();
5942 for (resumed, uninterrupted) in resumed.weights.iter().zip(uninterrupted.weights.iter()) {
5943 assert!((resumed - uninterrupted).abs() <= 1.0e-10);
5944 }
5945
5946 assert!(resumed.bic().is_finite());
5947 }
5948
5949 #[test]
5950 fn gaussian_mixture_bic_is_finite_with_an_active_covariance_floor() {
5951 let per_cluster = 45usize;
5957 let mut data = Array2::<f64>::zeros((2 * per_cluster, 2));
5958 for sample in 0..per_cluster {
5959 let phase = std::f64::consts::TAU * sample as f64 / per_cluster as f64;
5960 data[[2 * sample, 0]] = -2.0;
5961 data[[2 * sample, 1]] = 0.08 * phase.sin();
5962 data[[2 * sample + 1, 0]] = 2.0 + 0.12 * phase.cos();
5963 data[[2 * sample + 1, 1]] = 0.08 * phase.sin();
5964 }
5965 let fit = fit_gaussian_mixture(data.view(), 2, GaussianMixtureConfig::default())
5966 .expect("the covariance floor defines a valid constrained mixture fit");
5967 let bic = fit.bic();
5968 assert!(bic.is_finite());
5969 assert_eq!(
5970 bic,
5971 -fit.loglik + 0.5 * fit.num_free_parameters() as f64 * (data.nrows() as f64).ln()
5972 );
5973 }
5974
5975 fn seven_clusters_on_a_circle_2262() -> Array2<f64> {
5976 let clusters = 7usize;
5977 let per_cluster = 32usize;
5978 let mut data = Array2::<f64>::zeros((clusters * per_cluster, 2));
5979 for cluster in 0..clusters {
5980 let angle = std::f64::consts::TAU * cluster as f64 / clusters as f64;
5981 let (sin_angle, cos_angle) = angle.sin_cos();
5982 for sample in 0..per_cluster {
5983 let phase = std::f64::consts::TAU * sample as f64 / per_cluster as f64;
5984 let local_radius = 0.035 * (1.0 + 0.3 * (3.0 * phase).cos());
5989 let radial_noise = local_radius * phase.cos();
5990 let tangent_noise = local_radius * phase.sin();
5991 let radius = 2.0 + radial_noise;
5992 let row = cluster * per_cluster + sample;
5993 data[[row, 0]] = 0.4 + radius * cos_angle - tangent_noise * sin_angle;
5994 data[[row, 1]] = -0.3 + radius * sin_angle + tangent_noise * cos_angle;
5995 }
5996 }
5997 data
5998 }
5999
6000 fn two_noisy_circles_for_union() -> Array2<f64> {
6001 let rows_per_circle = 96usize;
6002 let mut data = Array2::<f64>::zeros((2 * rows_per_circle, 2));
6003 for (circle, (center, radius)) in [([-4.0_f64, 0.3_f64], 1.2_f64), ([4.0, -0.2], 0.9)]
6004 .into_iter()
6005 .enumerate()
6006 {
6007 for sample in 0..rows_per_circle {
6008 let angle = std::f64::consts::TAU * sample as f64 / rows_per_circle as f64;
6009 let noisy_radius =
6010 radius + 0.045 * (3.0 * angle).cos() + 0.018 * (5.0 * angle).sin();
6011 let row = circle * rows_per_circle + sample;
6012 data[[row, 0]] = center[0] + noisy_radius * angle.cos();
6013 data[[row, 1]] = center[1] + noisy_radius * angle.sin();
6014 }
6015 }
6016 data
6017 }
6018
6019 #[test]
6020 fn circular_gaussian_density_avoids_extreme_scale_intermediate_overflow() {
6021 let noise_variance = f64::MAX / 2.0;
6022 let fit =
6023 CircularGaussianFit2d::from_parameters([0.0, 0.0], 1.1e154, noise_variance).unwrap();
6024 let center_log_density = fit.log_density(0.0, 0.0);
6027 let off_center_log_density = fit.log_density(1.7e154, 0.0);
6028 assert!(center_log_density.is_finite());
6029 assert!(off_center_log_density.is_finite());
6030 let expected_center = -std::f64::consts::TAU.ln()
6031 - noise_variance.ln()
6032 - 0.5 * (fit.radius() / noise_variance.sqrt()).powi(2);
6033 assert_eq!(center_log_density, expected_center);
6034 }
6035
6036 #[test]
6037 fn union_circles_use_the_shared_normalized_cartesian_density() {
6038 let data = two_noisy_circles_for_union();
6039 let config = GaussianMixtureConfig::default();
6040 let density_fit =
6041 fit_union_density(data.view(), UnionStructure::CircleCircle, config).unwrap();
6042 let union = fit_union_structure(data.view(), UnionStructure::CircleCircle, config).unwrap();
6043 assert_eq!(
6044 union.total_parameters,
6045 2 * CircularGaussianFit2d::NUM_FREE_PARAMETERS + 1
6046 );
6047 let component_weight_sum: f64 = union
6048 .components
6049 .iter()
6050 .map(|component| component.mixing_weight)
6051 .sum();
6052 assert!((component_weight_sum - 1.0).abs() <= 8.0 * f64::EPSILON);
6053
6054 let mut fitted_centers = Array2::<f64>::zeros((density_fit.components.len(), 2));
6055 for (index, component) in density_fit.components.iter().enumerate() {
6056 let UnionDensityModel::Circle(fit) = &component.model else {
6057 panic!("circle+circle union produced a non-circle density");
6058 };
6059 let center = fit.center();
6060 fitted_centers[[index, 0]] = center[0];
6061 fitted_centers[[index, 1]] = center[1];
6062 let at_center = fit.log_density(center[0], center[1]);
6063 let expected = -std::f64::consts::TAU.ln()
6064 - fit.noise_variance().ln()
6065 - 0.5 * (fit.radius() / fit.noise_variance().sqrt()).powi(2);
6066 assert!(at_center.is_finite());
6067 assert!((at_center - expected).abs() < 1.0e-12 * (1.0 + expected.abs()));
6068 }
6069
6070 let training_log_density = union_per_point_log_density(
6071 data.view(),
6072 data.view(),
6073 UnionStructure::CircleCircle,
6074 config,
6075 )
6076 .unwrap();
6077 let direct_log_likelihood = pairwise_sum(
6078 training_log_density
6079 .as_slice()
6080 .expect("owned score vector is contiguous"),
6081 );
6082 assert!(
6083 (union.log_likelihood - direct_log_likelihood).abs()
6084 <= 1.0e-12 * (1.0 + direct_log_likelihood.abs())
6085 );
6086 let expected_bic = -direct_log_likelihood
6087 + 0.5 * union.total_parameters as f64 * (data.nrows() as f64).ln();
6088 assert!((union.bic - expected_bic).abs() <= 1.0e-12 * (1.0 + expected_bic.abs()));
6089
6090 let held_out = union_per_point_log_density(
6091 data.view(),
6092 fitted_centers.view(),
6093 UnionStructure::CircleCircle,
6094 config,
6095 )
6096 .unwrap();
6097 assert!(held_out.iter().all(|value| value.is_finite()));
6098 }
6099
6100 fn circle_and_point_union_data() -> (Array2<f64>, Vec<Vec<usize>>) {
6101 let circle_rows = 32usize;
6102 let point_rows = 12usize;
6103 let mut data = Array2::<f64>::zeros((circle_rows + point_rows, 2));
6104 for row in 0..circle_rows {
6105 let angle = std::f64::consts::TAU * row as f64 / circle_rows as f64;
6106 let radius = 1.0 + 0.025 * (3.0 * angle).cos();
6107 data[[row, 0]] = -4.0 + radius * angle.cos();
6108 data[[row, 1]] = 0.2 + radius * angle.sin();
6109 }
6110 for offset in 0..point_rows {
6111 let phase = offset as f64;
6112 let row = circle_rows + offset;
6113 data[[row, 0]] = 4.0 + 0.055 * (1.7 * phase).cos() + 0.018 * (0.4 * phase).sin();
6114 data[[row, 1]] = -0.3 + 0.052 * (1.3 * phase).sin() - 0.015 * (0.9 * phase).cos();
6115 }
6116 (
6117 data,
6118 vec![
6119 (0..circle_rows).collect(),
6120 (circle_rows..circle_rows + point_rows).collect(),
6121 ],
6122 )
6123 }
6124
6125 #[test]
6126 fn heterogeneous_union_role_assignment_is_group_label_invariant() {
6127 let (data, groups) = circle_and_point_union_data();
6128 let config = GaussianMixtureConfig::default();
6129 let forward = fit_union_density_from_groups(
6130 data.view(),
6131 UnionStructure::CirclePointCluster,
6132 &groups,
6133 config,
6134 )
6135 .unwrap();
6136 let reversed_groups = vec![groups[1].clone(), groups[0].clone()];
6137 let reversed = fit_union_density_from_groups(
6138 data.view(),
6139 UnionStructure::CirclePointCluster,
6140 &reversed_groups,
6141 config,
6142 )
6143 .unwrap();
6144
6145 assert_eq!(forward.components[0].kind, UnionComponentKind::Circle);
6146 assert_eq!(forward.components[1].kind, UnionComponentKind::PointCluster);
6147 assert_eq!(
6148 reversed.components[0].kind,
6149 UnionComponentKind::PointCluster
6150 );
6151 assert_eq!(reversed.components[1].kind, UnionComponentKind::Circle);
6152 assert_eq!(forward.total_parameters, 4 + 3 + 1);
6153 assert_eq!(reversed.total_parameters, forward.total_parameters);
6154 assert!(
6155 (forward.log_likelihood - reversed.log_likelihood).abs()
6156 <= 1.0e-12 * (1.0 + forward.log_likelihood.abs())
6157 );
6158 assert!((forward.bic - reversed.bic).abs() <= 1.0e-12 * (1.0 + forward.bic.abs()));
6159 }
6160
6161 #[test]
6162 fn point_cluster_is_isotropic_and_line_remains_full_covariance() {
6163 let (mut data, mut groups) = circle_and_point_union_data();
6164 for row in 0..groups[0].len() {
6167 let coordinate = (row as f64 - 15.5) / 4.0;
6168 data[[row, 0]] = -4.0 + coordinate;
6169 data[[row, 1]] = 0.2 + 0.018 * coordinate + 0.006 * (1.9 * row as f64).sin();
6170 }
6171 let fit = fit_union_density_from_groups(
6172 data.view(),
6173 UnionStructure::LineCluster,
6174 &groups,
6175 GaussianMixtureConfig::default(),
6176 )
6177 .unwrap();
6178 assert_eq!(fit.components[0].kind, UnionComponentKind::Line);
6179 assert_eq!(fit.components[0].num_parameters, 5);
6180 assert_eq!(fit.components[1].kind, UnionComponentKind::PointCluster);
6181 assert_eq!(fit.components[1].num_parameters, 3);
6182 assert_eq!(fit.total_parameters, 5 + 3 + 1);
6183
6184 let UnionDensityModel::Gaussian(point) = &fit.components[1].model else {
6185 panic!("point cluster did not produce a Gaussian density");
6186 };
6187 assert_eq!(point.precision[[0, 1]], 0.0);
6188 assert_eq!(point.precision[[1, 0]], 0.0);
6189 assert_eq!(point.precision[[0, 0]], point.precision[[1, 1]]);
6190 let total_weight: f64 = fit
6191 .components
6192 .iter()
6193 .map(|component| component.mixing_weight)
6194 .sum();
6195 assert!((total_weight - 1.0).abs() <= 8.0 * f64::EPSILON);
6196
6197 groups.reverse();
6200 let reversed = fit_union_density_from_groups(
6201 data.view(),
6202 UnionStructure::LineCluster,
6203 &groups,
6204 GaussianMixtureConfig::default(),
6205 )
6206 .unwrap();
6207 assert_eq!(
6208 reversed.components[0].kind,
6209 UnionComponentKind::PointCluster
6210 );
6211 assert_eq!(reversed.components[1].kind, UnionComponentKind::Line);
6212 assert!((fit.bic - reversed.bic).abs() <= 1.0e-12 * (1.0 + fit.bic.abs()));
6213 }
6214
6215 #[test]
6216 fn isotropic_union_density_uses_the_same_fractional_mean_chart_as_its_mle() {
6217 let translated = ndarray::array![[1.0e16], [1.0e16 + 2.0], [1.0e16 + 2.0]];
6218 let fit = fit_isotropic_gaussian_component(translated.view(), 1.0e-12).unwrap();
6219 let variance = fit.precision[[0, 0]].recip();
6220 assert!((variance - 8.0 / 9.0).abs() <= 32.0 * f64::EPSILON);
6221
6222 let residuals = [-4.0 / 3.0, 2.0 / 3.0, 2.0 / 3.0];
6223 let expected_log_norm = -0.5 * ((2.0 * std::f64::consts::PI).ln() + variance.ln());
6224 for (row, residual) in residuals.into_iter().enumerate() {
6225 let expected = expected_log_norm - 0.5 * residual * residual / variance;
6226 let actual = fit.log_density(translated.row(row));
6227 assert!(
6228 (actual - expected).abs() <= 32.0 * f64::EPSILON * (1.0 + expected.abs()),
6229 "row {row}: density chart disagrees with fitted MLE residual: actual={actual}, expected={expected}"
6230 );
6231 }
6232
6233 let subnormal = f64::from_bits(1);
6234 let constant = ndarray::array![[subnormal], [subnormal], [subnormal]];
6235 let constant_fit = fit_isotropic_gaussian_component(constant.view(), 1.0).unwrap();
6236 assert_eq!(constant_fit.residual(constant.row(0)), vec![0.0]);
6237 assert_eq!(
6238 constant_fit.log_density(constant.row(0)),
6239 constant_fit.log_norm
6240 );
6241 }
6242
6243 #[test]
6244 fn union_ladder_fails_closed_when_one_declared_structure_fails() {
6245 let mut data = Array2::<f64>::zeros((8, 2));
6246 for row in 0..5 {
6247 let angle = std::f64::consts::TAU * row as f64 / 5.0;
6248 data[[row, 0]] = -5.0 + angle.cos();
6249 data[[row, 1]] = angle.sin();
6250 }
6251 data[[5, 0]] = 5.00;
6252 data[[5, 1]] = 0.00;
6253 data[[6, 0]] = 5.08;
6254 data[[6, 1]] = 0.02;
6255 data[[7, 0]] = 4.97;
6256 data[[7, 1]] = 0.07;
6257
6258 let error = fit_union_ladder(data.view(), GaussianMixtureConfig::default()).unwrap_err();
6259 assert!(error.contains("every declared structure must fit"));
6260 assert!(error.contains(UnionStructure::CircleCircle.as_str()));
6261 assert!(error.contains("needs at least 5 rows"));
6262 }
6263
6264 #[test]
6265 fn ring_of_clusters_fit_is_stationary_and_complexity_priced_2262() {
6266 let data = seven_clusters_on_a_circle_2262();
6267 let config = GaussianMixtureConfig::default();
6268 let fit = fit_ring_gaussian_mixture(data.view(), 7, config).unwrap();
6269 let certificate = fit.certificate();
6270 assert!(certificate.objective_residual <= certificate.objective_tolerance);
6271 assert!(certificate.parameter_residual <= certificate.parameter_tolerance);
6272 assert_eq!(fit.num_free_parameters(), 17);
6273 assert!((fit.center()[0] - 0.4).abs() < 0.05);
6274 assert!((fit.center()[1] + 0.3).abs() < 0.05);
6275 assert!((fit.radius() - 2.0).abs() < 0.05);
6276 assert!(fit.variance().is_finite() && fit.variance() > 0.0);
6277 assert!(
6278 fit.per_point_log_density(data.view())
6279 .unwrap()
6280 .iter()
6281 .all(|value| value.is_finite())
6282 );
6283 assert!(fit.bic().is_finite());
6284
6285 let free = fit_gaussian_mixture(data.view(), 7, config).unwrap();
6286 assert_eq!(free.num_free_parameters(), 41);
6287 assert!(fit.num_free_parameters() < free.num_free_parameters());
6288 }
6289
6290 #[test]
6291 fn ring_certificate_uses_identifiable_component_means() {
6292 let y = 0.91_f64.sqrt();
6298 let weights = Array1::from_vec(vec![0.2, 0.3, 0.5]);
6299 let previous = RingMixtureState {
6300 weights: weights.clone(),
6301 center: Array1::from_vec(vec![0.0, 0.0]),
6302 radius: 1.0,
6303 directions: Array2::from_shape_vec((3, 2), vec![0.3, y, 0.3, -y, 0.3, y]).unwrap(),
6304 variance: 0.25,
6305 mean_log_likelihood: -1.0,
6306 completed_iterations: 10,
6307 };
6308 let next = RingMixtureState {
6309 weights,
6310 center: Array1::from_vec(vec![0.6, 0.0]),
6311 radius: 1.0,
6312 directions: Array2::from_shape_vec((3, 2), vec![-0.3, y, -0.3, -y, -0.3, y]).unwrap(),
6313 variance: 0.25,
6314 mean_log_likelihood: -1.0,
6315 completed_iterations: 11,
6316 };
6317 assert!(relative_parameter_step(previous.center[0], next.center[0]) > 0.5);
6318 assert_eq!(ring_identifiable_parameter_residual(&previous, &next), 0.0);
6319 }
6320
6321 #[test]
6322 fn ring_parameter_map_weights_component_motion_by_predictive_mass_2324() {
6323 let raw_coordinate_step: f64 = 4.6e-8;
6324 let (next_y, next_x) = raw_coordinate_step.sin_cos();
6325 let previous = RingMixtureState {
6326 weights: array![0.25, 0.75],
6327 center: array![0.0, 0.0],
6328 radius: 1.0,
6329 directions: array![[1.0, 0.0], [0.0, 1.0]],
6330 variance: 1.0,
6331 mean_log_likelihood: -1.0,
6332 completed_iterations: 10,
6333 };
6334 let next = RingMixtureState {
6335 weights: previous.weights.clone(),
6336 center: previous.center.clone(),
6337 radius: previous.radius,
6338 directions: array![[next_x, next_y], [0.0, 1.0]],
6339 variance: previous.variance,
6340 mean_log_likelihood: -1.0,
6341 completed_iterations: 11,
6342 };
6343 let raw_mean_step = (next_x - 1.0).hypot(next_y);
6344 assert!(raw_mean_step > f64::EPSILON.sqrt());
6345 assert!(ring_identifiable_parameter_residual(&previous, &next) <= f64::EPSILON.sqrt());
6346 }
6347
6348 #[test]
6349 fn stacked_mean_is_weighted_combination() {
6350 let weights = Array1::from_vec(vec![0.25, 0.75]);
6351 let means = vec![
6352 Array1::from_vec(vec![1.0, 2.0, 3.0]),
6353 Array1::from_vec(vec![5.0, 6.0, 7.0]),
6354 ];
6355 let out = stacked_predictive_mean(&weights, &means).unwrap();
6356 assert!((out[0] - (0.25 * 1.0 + 0.75 * 5.0)).abs() < 1e-12);
6357 assert!((out[2] - (0.25 * 3.0 + 0.75 * 7.0)).abs() < 1e-12);
6358 }
6359
6360 #[test]
6361 fn stacked_mean_rejects_shape_mismatch() {
6362 let weights = Array1::from_vec(vec![0.5, 0.5]);
6363 let means = vec![
6364 Array1::from_vec(vec![1.0, 2.0]),
6365 Array1::from_vec(vec![3.0]),
6366 ];
6367 assert!(stacked_predictive_mean(&weights, &means).is_err());
6368 }
6369
6370 fn hybrid_slot(
6385 linear_nle: f64,
6386 p_linear: usize,
6387 latent_dim: usize,
6388 p_curved: usize,
6389 theta: f64,
6390 curved_loglik_gain: f64,
6391 ) -> Vec<HybridAtomCandidate> {
6392 let param_price =
6393 0.5 * (p_curved as f64 - p_linear as f64) * (2.0 * std::f64::consts::PI).ln();
6394 let curved_nle = linear_nle - curved_loglik_gain + param_price;
6395 vec![
6396 HybridAtomCandidate::linear(linear_nle, p_linear),
6397 HybridAtomCandidate::curved(latent_dim, curved_nle, p_curved, Some(theta)),
6398 ]
6399 }
6400
6401 #[test]
6402 fn hybrid_dominance_floor_selects_linear_when_turning_is_zero() {
6403 let slot = hybrid_slot(100.0, 2, 1, 5, 0.0, 0.0);
6408 let choice = select_hybrid_atom(&slot).unwrap();
6409 assert!(choice.param.is_linear());
6410 assert_eq!(choice.param, HybridAtomParam::Linear);
6411 assert!(choice.curved_turning.unwrap() <= HYBRID_LINEAR_TURNING_FLOOR);
6413 }
6414
6415 #[test]
6416 fn hybrid_selects_curved_when_turning_pays_for_itself() {
6417 let slot = hybrid_slot(100.0, 2, 1, 5, 2.0 * std::f64::consts::PI, 30.0);
6421 let choice = select_hybrid_atom(&slot).unwrap();
6422 assert_eq!(choice.param, HybridAtomParam::Curved { latent_dim: 1 });
6423 assert!(choice.curved_evidence_margin > 0.0);
6425 }
6426
6427 #[test]
6428 fn hybrid_keeps_linear_when_curvature_doesnt_pay_its_price() {
6429 let slot = hybrid_slot(100.0, 2, 1, 5, 0.05, 0.1);
6433 let choice = select_hybrid_atom(&slot).unwrap();
6434 assert!(choice.param.is_linear());
6435 assert!(choice.curved_evidence_margin <= 0.0);
6436 }
6437
6438 #[test]
6439 fn hybrid_tie_breaks_to_the_cheaper_linear_atom() {
6440 let theta = 0.5; let nle = 42.0;
6445 let slot = vec![
6446 HybridAtomCandidate::linear(nle, 2),
6447 HybridAtomCandidate::curved(1, nle, 5, Some(theta)),
6448 ];
6449 let choice = select_hybrid_atom(&slot).unwrap();
6450 assert!(choice.param.is_linear());
6451 assert_eq!(choice.num_parameters, 2);
6452 }
6453
6454 #[test]
6455 fn hybrid_split_reduces_to_pure_linear_when_all_features_are_straight() {
6456 let slots: Vec<Vec<HybridAtomCandidate>> = (0..6)
6460 .map(|i| hybrid_slot(50.0 + i as f64, 2, 1, 5, 0.0, 0.0))
6461 .collect();
6462 let split = select_hybrid_split(&slots).unwrap();
6463 assert!(split.is_pure_linear());
6464 assert_eq!(split.curved_atom_count, 0);
6465 assert_eq!(split.linear_atom_count(), 6);
6466 let pure_linear: f64 = (0..6).map(|i| 50.0 + i as f64).sum();
6468 assert!((split.total_negative_log_evidence - pure_linear).abs() < 1e-12);
6469 }
6470
6471 #[test]
6472 fn hybrid_split_reduces_to_pure_curved_when_every_feature_curves() {
6473 let slots: Vec<Vec<HybridAtomCandidate>> = (0..5)
6476 .map(|i| hybrid_slot(80.0 + i as f64, 2, 1, 5, 2.0 * std::f64::consts::PI, 40.0))
6477 .collect();
6478 let split = select_hybrid_split(&slots).unwrap();
6479 assert!(split.is_pure_curved());
6480 assert_eq!(split.curved_atom_count, 5);
6481 assert_eq!(split.linear_atom_count(), 0);
6482 }
6483
6484 #[test]
6485 fn hybrid_split_on_mixed_dictionary_picks_curved_for_circles_linear_for_directions() {
6486 let mut slots: Vec<Vec<HybridAtomCandidate>> = Vec::new();
6496 let mut pure_linear_baseline = 0.0_f64;
6497 for i in 0..3 {
6500 let linear_nle = 120.0 + 3.0 * i as f64;
6501 pure_linear_baseline += linear_nle;
6502 slots.push(hybrid_slot(
6503 linear_nle,
6504 2,
6505 1,
6506 5,
6507 2.0 * std::f64::consts::PI,
6508 35.0,
6509 ));
6510 }
6511 for i in 0..4 {
6514 let linear_nle = 90.0 + 2.0 * i as f64;
6515 pure_linear_baseline += linear_nle;
6516 slots.push(hybrid_slot(linear_nle, 2, 1, 5, 0.0, 0.0));
6517 }
6518
6519 let split = select_hybrid_split(&slots).unwrap();
6520
6521 for (idx, choice) in split.atoms.iter().enumerate() {
6524 if idx < 3 {
6525 assert_eq!(
6526 choice.param,
6527 HybridAtomParam::Curved { latent_dim: 1 },
6528 "circle slot {idx} should select curved"
6529 );
6530 } else {
6531 assert!(
6532 choice.param.is_linear(),
6533 "direction slot {idx} should select linear"
6534 );
6535 }
6536 }
6537 assert_eq!(split.curved_atom_count, 3);
6538 assert_eq!(split.linear_atom_count(), 4);
6539
6540 assert!(
6546 split.total_negative_log_evidence <= pure_linear_baseline + 1e-9,
6547 "hybrid NLE {} must be <= summed linear-candidate NLE {}",
6548 split.total_negative_log_evidence,
6549 pure_linear_baseline
6550 );
6551 assert!(split.total_negative_log_evidence < pure_linear_baseline);
6553 }
6554
6555 #[test]
6556 fn hybrid_split_rejects_empty_slot() {
6557 let slots = vec![hybrid_slot(10.0, 2, 1, 5, 0.0, 0.0), Vec::new()];
6558 assert!(select_hybrid_split(&slots).is_err());
6559 }
6560
6561 fn cand(name: &str, score: f64, edf: f64, log_lik: f64) -> RemlCandidate {
6569 RemlCandidate {
6570 index: 0,
6571 name: name.to_string(),
6572 score,
6573 edf: Some(edf),
6574 log_lik: Some(log_lik),
6575 family: None,
6576 n_obs: None,
6577 }
6578 }
6579
6580 #[test]
6581 fn ranking_score_is_conditional_aic_when_loglik_and_edf_present() {
6582 let c = cand("m", 999.0, 6.748, -32.0866);
6584 let expected = -2.0 * -32.0866 + 2.0 * 6.748;
6585 assert!((c.ranking_score() - expected).abs() < 1e-9);
6586 }
6587
6588 #[test]
6589 fn ranking_score_falls_back_to_evidence_without_loglik() {
6590 let c = RemlCandidate {
6591 index: 0,
6592 name: "m".to_string(),
6593 score: 151.28,
6594 edf: Some(6.0),
6595 log_lik: None,
6596 family: None,
6597 n_obs: None,
6598 };
6599 assert_eq!(c.ranking_score(), 151.28);
6600 }
6601
6602 #[test]
6603 fn compare_models_rejects_pure_noise_smooth_despite_lower_evidence() {
6604 let small = cand("small", 180.526, 6.748, -32.0866);
6611 let big = cand("big", 177.404, 14.250, -32.1212);
6612
6613 assert!(big.score < small.score);
6615
6616 let cmp = compare_reml_fits(vec![small, big]).expect("compare");
6617 assert_eq!(
6618 cmp.winner, "small",
6619 "compare_models must Occam-penalise the pure-noise smooth and pick the smaller model"
6620 );
6621 let small_row = cmp
6624 .score_table
6625 .iter()
6626 .find(|r| r.name == "small")
6627 .expect("small row");
6628 let big_row = cmp
6629 .score_table
6630 .iter()
6631 .find(|r| r.name == "big")
6632 .expect("big row");
6633 assert!((small_row.reml_score - 180.526).abs() < 1e-9);
6634 assert!((big_row.reml_score - 177.404).abs() < 1e-9);
6635 }
6636
6637 #[test]
6638 fn ranking_bayes_factor_is_akaike_evidence_ratio_not_its_square() {
6639 let delta_aic = 27.68_f64;
6649 let winner = cand("winner", 100.0, 0.0, 0.0);
6650 let loser = cand("loser", 110.0, 0.0, -delta_aic / 2.0);
6651
6652 let cmp = compare_reml_fits(vec![winner, loser]).expect("compare");
6653 assert_eq!(cmp.winner, "winner");
6654
6655 let loser_row = cmp
6656 .ranking
6657 .iter()
6658 .find(|r| r.name == "loser")
6659 .expect("loser ranking row");
6660
6661 assert!((loser_row.delta - delta_aic).abs() < 1e-9);
6663
6664 let expected = (0.5 * delta_aic).exp();
6667 assert!(
6668 (loser_row.bayes_factor / expected - 1.0).abs() < 1e-9,
6669 "ranking bayes_factor {} should be exp(½ΔAIC)={}, not exp(ΔAIC)={}",
6670 loser_row.bayes_factor,
6671 expected,
6672 delta_aic.exp()
6673 );
6674 assert!(loser_row.bayes_factor < delta_aic.exp() * 0.5);
6676
6677 let loser_score_row = cmp
6682 .score_table
6683 .iter()
6684 .find(|r| r.name == "loser")
6685 .expect("loser score row");
6686 let expected_reml_bf = 10.0_f64.exp();
6687 assert!(
6688 (loser_score_row.bayes_factor_best_over_model / expected_reml_bf - 1.0).abs() < 1e-9,
6689 "raw-REML bayes_factor_best_over_model must stay exp(Δreml)=exp(10), got {}",
6690 loser_score_row.bayes_factor_best_over_model
6691 );
6692 }
6693
6694 #[test]
6695 fn compare_models_keeps_power_for_a_relevant_smooth() {
6696 let small = cand("small", 1025.067, 6.75, -368.985);
6702 let big = cand("big", 199.509, 14.25, -33.165);
6703 let cmp = compare_reml_fits(vec![small, big]).expect("compare");
6704 assert_eq!(
6705 cmp.winner, "big",
6706 "compare_models must retain power: the relevant smooth's model must win"
6707 );
6708 }
6709
6710 #[test]
6711 fn compare_models_rejects_mismatched_observation_counts() {
6712 let with_n = |name: &str, n: usize| RemlCandidate {
6716 index: 0,
6717 name: name.to_string(),
6718 score: 100.0,
6719 edf: Some(5.0),
6720 log_lik: Some(-40.0),
6721 family: Some("gaussian".to_string()),
6722 n_obs: Some(n),
6723 };
6724 let err = compare_reml_fits(vec![with_n("big", 500), with_n("small", 100)])
6725 .expect_err("cross-n comparison must be rejected");
6726 assert!(
6727 err.contains("number of observations") && err.contains("500") && err.contains("100"),
6728 "n-guard error should name the incomparable counts, got: {err}"
6729 );
6730
6731 compare_reml_fits(vec![with_n("a", 250), with_n("b", 250)])
6733 .expect("same-n comparison must succeed");
6734
6735 let without_n = RemlCandidate {
6738 index: 0,
6739 name: "legacy".to_string(),
6740 score: 90.0,
6741 edf: Some(4.0),
6742 log_lik: Some(-35.0),
6743 family: Some("gaussian".to_string()),
6744 n_obs: None,
6745 };
6746 compare_reml_fits(vec![with_n("counted", 500), without_n])
6747 .expect("an unconstrained (None) count must not trip the guard");
6748 }
6749}