Skip to main content

gam_solve/
evidence.rs

1//! Canonical Laplace evidence, IFT cascade, and topology selection.
2//!
3//! This module is the single canonical entry point for:
4//!
5//!   1. Laplace evidence `V(ρ, T) = F + (1/2) log|H| - (1/2) log|S(ρ)|+
6//!      - ((dim(H)-rank(S))/2) log(2π)`
7//!      evaluated at the arrow-Schur inner-loop fixed point.
8//!   2. The full IFT cascade `∂u*/∂β → ∂β*/∂ρ → ∂u*/∂ρ` through the three
9//!      continuous tiers `(u, β, ρ)`, per §2.2 / §2.4 / §2.6.
10//!   3. The per-`ρ` evidence gradient `∂V/∂ρ` via the arrow trace formula,
11//!      per §3.5 / §3.7 / §3.8.
12//!   4. Discrete topology selection across `{periodic, flat, sphere, torus}`,
13//!      per §4 (4.1 / 4.5 / 4.6).
14//!
15//! ## Crucial numerical invariants (proposal §1.7, §6.4, §6.5)
16//!
17//!   * Evidence log-determinants use **undamped** factors. The cached
18//!     `ArrowFactorCache::htt_factors_undamped` Cholesky factors of
19//!     `H_uu_i` (no `ridge_u`) are the ones that must enter
20//!     `Σ_i log|H_uu_i|`. Likewise a factored Schur log-det must be of
21//!     `A(0, 0) = H_ββ - Σ_i H_uβ_iᵀ H_uu_i⁻¹ H_uβ_i`, not the LM-damped
22//!     surrogate. Matrix-free evidence callers must provide the matching
23//!     undamped HVP so the same log-det is estimated by SLQ.
24//!   * IFT solves invert `H_uu`, not `H_uu + ridge_u I` (proposal §1.7,
25//!     §6.6). The evidence-side IFT predictor loop here uses the undamped
26//!     `htt_factors_undamped` factors for exactly this reason.
27//!   * Penalty pseudo-logdet `log|S(ρ)|+` is the prior penalty, distinct
28//!     from the arrow Schur complement (proposal §3.1, §3.6). The variable
29//!     names below preserve that distinction:
30//!       `arrow_schur_log_det`   = `log|A|` where `A` is the arrow Schur.
31//!       `penalty_log_det`       = `log|S_pen(ρ)|+` where `S_pen` is the
32//!                                 prior penalty matrix pseudo-logdet.
33//!
34//! ## Sign discipline (proposal §3.1, §4.3)
35//!
36//! `V` as written is the *negative log evidence* when `F` is the
37//! penalized negative log posterior. The maximizer of evidence is the
38//! minimizer of `V`. For the public API we expose **negative log
39//! evidence** under `laplace_evidence` and rank topologies by the
40//! **minimum** of the configured per-row or per-effective-dimension
41//! normalization (see `select_topology` below); equivalently the caller can
42//! negate and `argmax`.
43
44use 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/// Matrix-free SPD Hessian logdet source used when the arrow Schur factor is
66/// not materialized. The callback must apply the same undamped Hessian whose
67/// determinant enters the Laplace evidence.
68#[derive(Clone, Copy)]
69pub struct EvidenceHvpLogDet<'a> {
70    pub dim: usize,
71    pub apply: &'a dyn Fn(&[f64]) -> Vec<f64>,
72}
73
74/// Source for the Hessian log determinant in [`laplace_evidence`].
75#[derive(Clone, Copy)]
76pub enum EvidenceLogDetSource<'a> {
77    /// Use the exact arrow Cholesky factors, falling back to `fallback_hvp`
78    /// when the Schur factor is absent on a matrix-free solve.
79    FactoredArrow {
80        cache: &'a ArrowFactorCache,
81        fallback_hvp: Option<EvidenceHvpLogDet<'a>>,
82    },
83    /// Use an HVP callback directly. Dimensions at or below
84    /// [`ANALYTIC_LOGDET_DENSE_DIM_THRESHOLD`] are materialized exactly;
85    /// larger operators use the same Rademacher-Lanczos SLQ constants as
86    /// `FrozenAnalyticPenaltyOp`.
87    Hvp(EvidenceHvpLogDet<'a>),
88}
89
90// ---------------------------------------------------------------------------
91// Topology candidate enum and selection result
92// ---------------------------------------------------------------------------
93
94/// Discrete topology choice for the latent coordinate domain.
95///
96/// Maps directly to the set `{periodic, flat, sphere, torus}`. No additional
97/// variants — unused candidate variants are deliberately not carried
98/// alongside the four-way selector.
99#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
100pub enum TopologyKind {
101    /// `S¹` or periodic interval (cyclic B-spline / periodic Duchon).
102    Periodic,
103    /// `Rᵈ` Euclidean Duchon / Matérn / thin-plate patch.
104    Flat,
105    /// `S²` embedded in `R³`, spherical Wahba/Sobolev basis.
106    Sphere,
107    /// `S¹ × S¹` mixed-periodicity Duchon.
108    Torus,
109}
110
111impl TopologyKind {
112    /// Tie-break priority — smaller wins. Per §4.6: `flat < periodic <
113    /// sphere < torus`.
114    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/// One topology candidate together with the evidence ingredients it
125/// produced at its own fitted optimum.
126#[derive(Debug, Clone)]
127pub struct TopologyCandidate {
128    pub kind: TopologyKind,
129    /// Negative-log-evidence `V(ρ_T*, T)` evaluated at the candidate's own
130    /// fitted `(ρ_T*, β_T*, u_T*)`.
131    pub negative_log_evidence: f64,
132    /// Effective integrated dimension after rank/nullspace accounting. This
133    /// is the dimension used for per-complexity topology normalization.
134    pub effective_dim: f64,
135    /// Number of response rows used to fit this topology candidate. This is
136    /// the dimension used for per-observation topology normalization.
137    pub n_obs: usize,
138    /// `True` iff the candidate's continuous inner+outer fit converged
139    /// cleanly. Failed candidates are excluded from ranking (proposal
140    /// §4.4 item 7 and §6.11).
141    pub converged: bool,
142    /// Optional rationale string for excluded candidates (proposal
143    /// §6.11): `"sphere input not on S²"`, `"torus periods missing"`, etc.
144    pub exclusion_reason: Option<String>,
145}
146
147/// Outcome of [`select_topology`].
148#[derive(Debug, Clone)]
149pub struct SelectedTopology {
150    pub winner: TopologyKind,
151    /// All candidates sorted from best (lowest negative log evidence)
152    /// to worst, with excluded candidates appended last.
153    pub ranking: Vec<TopologyCandidate>,
154    /// `True` iff the top two finite scores fall within `tie_tolerance`.
155    /// Per §4.6 we still pick one — the simpler topology — but expose
156    /// the tie so callers can warn.
157    pub tie: bool,
158}
159
160/// Tolerance options for the topology comparator.
161#[derive(Debug, Clone, Copy)]
162pub struct TopologySelectOptions {
163    /// Maximum `|V_a - V_b|` for which two candidates are treated as
164    /// numerically tied after [`TopologyScoreScale`] normalization. Default
165    /// `1e-3` per proposal §4.6 examples.
166    pub tie_tolerance: f64,
167    /// Score scale used for discrete topology comparison. Raw evidence is
168    /// intentionally not a selector because candidates may have different
169    /// row counts and basis/nullspace dimensions.
170    pub score_scale: TopologyScoreScale,
171}
172
173/// Normalization applied before ranking topology candidates.
174#[derive(Debug, Clone, Copy, PartialEq, Eq)]
175pub enum TopologyScoreScale {
176    /// Compare negative log evidence per observation row.
177    PerObservation,
178    /// Compare negative log evidence per effective integrated dimension.
179    PerEffectiveDim,
180}
181
182/// Convergence controls for stacking retained topology predictive densities.
183#[derive(Debug, Clone, Copy, Serialize, Deserialize)]
184pub struct StackingConfig {
185    /// Exhaustion-escalation bound on solver iterations. It never selects the
186    /// weights: exhausting it without the KKT certificate is an error, so an
187    /// uncertified iterate can never feed model selection.
188    pub max_iter: usize,
189    /// Simplex-KKT residual the solution must certify before it is returned.
190    /// The residual is scale-free: the KKT multiplier of the stacking problem
191    /// is exactly 1, so `g_k − 1` is already a relative stationarity measure.
192    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/// Auditable global-optimality certificate for a stacking solution.
205#[derive(Debug, Clone, Copy, Serialize, Deserialize)]
206pub struct StackingCertificate {
207    /// Achieved held-out mean log predictive density.
208    pub mean_log_score: f64,
209    /// Frank-Wolfe/KKT duality gap `max_k g_k - w·g`. Concavity makes this
210    /// an upper bound on objective suboptimality.
211    pub duality_gap: f64,
212    /// Absolute simplex mass residual `|Σw - 1|`.
213    pub simplex_residual: f64,
214    /// Error in the analytic multiplier identity `w·g = 1`.
215    pub multiplier_residual: f64,
216    /// Largest complementary-slackness residual `w_k |g_k - w·g|`.
217    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/// Serializable work state carried by a stacking exhaustion error and accepted
230/// by [`resume_stacking_weights`]. Weights stay aligned to the original input
231/// columns; no candidate is silently dropped.
232#[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/// Typed stacking failure. In particular, exhaustion carries both its
240/// certificate evidence and an exact checkpoint rather than returning weights.
241#[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/// Simplex weights for retained topology candidates plus their verified global
294/// optimum certificate.
295#[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
434/// Solve the stacking-of-predictive-distributions weight problem from a
435/// per-observation held-out log-density table `log_density[i, k] = log p_k(y_i)`.
436///
437/// This belongs on the evidence surface rather than in a separate solver: it is
438/// the topology/evidence consumer that replaces winner-take-all only when the
439/// caller has retained candidate fits and per-point held-out densities.
440///
441/// ## Optimality certificate
442///
443/// The objective `f(w) = mean_i log Σ_k w_k p_ik` is concave on the simplex,
444/// so first-order KKT conditions are necessary AND sufficient for the global
445/// optimum. With `g_k = ∂f/∂w_k = mean_i p_ik / mix_i`, the simplex multiplier
446/// is exactly `Σ_k w_k g_k = 1`, so the KKT system is `g_k ≤ 1` for every
447/// candidate with `w_k · (1 − g_k) = 0` (complementary slackness). Iterates
448/// use an analytic reduced-space Newton step on the current simplex face; an
449/// exact concave line solve toward the most violated vertex activates a
450/// candidate or globalizes a singular Newton system. The solve returns only
451/// after the Frank-Wolfe duality gap and all primal/KKT residuals are verified
452/// below `config.kkt_tol`. Exhaustion carries the full certificate and a
453/// resumable weights checkpoint; uncertified weights never reach selection.
454pub 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
461/// Resume a previously exhausted stacking solve without redoing its accepted
462/// Newton/active-set work.
463pub 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
751/// Combine retained candidate response-scale means with stacking weights.
752pub 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// ---------------------------------------------------------------------------
782// Discrete mixture rung (Object 3a / WP-C)
783// ---------------------------------------------------------------------------
784//
785// A `k`-component full-covariance Gaussian mixture fitted by deterministic
786// k-means++-style seeding (reusing `terms::basis` farthest-point k-means) plus
787// EM to a tolerance. It is priced by its free-parameter count with the
788// invariant BIC approximation to negative log evidence,
789//
790//     BIC/2 = -loglik + (P/2) log(n).
791//
792// BIC is intentional here. An outer product of per-observation scores is not
793// an observed Hessian and need not be full-rank even when the likelihood has
794// curvature (an exactly centered Gaussian mean is the simplest counterexample).
795// Moreover, the covariance-floor constraint can put a component on a boundary,
796// where an interior SPD Laplace expansion is mathematically invalid. Without a
797// declared parameter prior and its Jacobian, a raw Hessian determinant would
798// also change under reparameterization. The smooth parametric shape candidates
799// use this same BIC-form score, so every shape-race corroborating score now has
800// one finite, parameterization-invariant meaning.
801
802/// Convergence + ladder controls for the discrete-mixture rung. All fields are
803/// fixed (no clock randomness, no env): deterministic seeding makes the fitted
804/// mixture a pure function of the data and `k`.
805#[derive(Debug, Clone, Copy, Serialize, Deserialize)]
806pub struct GaussianMixtureConfig {
807    /// Exhaustion-escalation bound on EM iterations. It never selects the
808    /// estimator: exhausting it without the convergence certificate
809    /// (monotone ascent + relative objective step below `loglik_tol`) is an
810    /// error, so an uncertified mixture never enters evidence comparison.
811    pub max_iter: usize,
812    /// Relative mean-log-likelihood improvement tolerance for EM stopping.
813    pub loglik_tol: f64,
814    /// Relative max-norm tolerance for the EM map in identifiable
815    /// predictive-density coordinates. Free mixtures use labeled component
816    /// measures `(w, w μ, w(Σ + μμᵀ))`; ring mixtures use mass-weighted,
817    /// noise-standardized component means plus their shared variance.
818    pub parameter_tol: f64,
819    /// Lower eigenvalue constraint for every component covariance. The M-step
820    /// solves this constrained likelihood problem exactly by spectral clipping;
821    /// it is not an additive ridge or an unmodelled prior.
822    pub covariance_floor: f64,
823    /// Maximum iterations for the deterministic k-means seeding pass.
824    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/// Residual evidence proving that the returned mixture is a fixed point of its
840/// likelihood EM map rather than merely the iterate present at a work cap.
841#[derive(Debug, Clone, Copy, Serialize, Deserialize)]
842pub struct GaussianMixtureCertificate {
843    /// Mean log likelihood at the exact parameter state being certified.
844    pub mean_log_likelihood: f64,
845    /// Signed gain produced by one further EM map application from that state.
846    pub mean_log_likelihood_gain: f64,
847    /// Absolute numerical uncertainty on the comparison of the two likelihoods.
848    /// This includes both the final likelihood-reduction error and the
849    /// scale-derived resolution floor of the composite EM map.
850    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/// Exact parameter state carried across an EM exhaustion boundary.
858#[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/// Typed Gaussian-mixture optimization failure. Exhaustion and a broken EM
870/// monotonicity invariant both carry the last internally consistent state.
871#[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/// A fitted `k`-component full-covariance Gaussian mixture.
939#[derive(Debug, Clone)]
940pub struct GaussianMixtureFit {
941    /// Mixing weights, length `k`, on the simplex.
942    weights: Array1<f64>,
943    /// Component means, `k × d`.
944    means: Array2<f64>,
945    /// Component covariances, `k` matrices of shape `d × d` (SPD).
946    covariances: Vec<Array2<f64>>,
947    /// Number of mixture components.
948    k: usize,
949    /// Data dimension.
950    d: usize,
951    /// Number of rows used to fit.
952    n_obs: usize,
953    /// Maximised total log-likelihood `Σ_i log Σ_j w_j N(y_i; μ_j, Σ_j)`.
954    loglik: f64,
955    /// EM iterations taken.
956    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    /// Free-parameter count `P` of a `k`-component full-covariance mixture in
982    /// `d` dimensions: `(k − 1)` mixing weights on the simplex, `k·d` mean
983    /// coordinates, and `k · d(d+1)/2` covariance entries. This is the exact
984    /// quantity that enters the rank-aware normalizer as `dim(H) − rank(S)`.
985    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    /// Per-observation log predictive density `log p(y_i)` under the fitted
991    /// mixture, length `n`. This is the held-out-density column source for
992    /// cross-class stacking when the mixture is evaluated on a held-out fold.
993    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    /// Schwarz BIC approximation to negative log evidence, divided by two so
1028    /// it shares the ordinary negative-log-likelihood scale used by the shape
1029    /// race. Lower is better.
1030    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/// Cached per-component Gaussian evaluator: mean, precision `Σ⁻¹`, and the
1036/// log-normalizing constant `−½(d log 2π + log|Σ|)`.
1037#[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        // Σ⁻¹ = V diag(1/λ) Vᵀ.
1081        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    /// `Σ⁻¹ (y − μ)`.
1168    #[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    // Hash logical row-major iteration order rather than the backing storage,
1200    // so an equivalent strided view resumes the same mathematical problem.
1201    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
1211/// Fit a `k`-component full-covariance Gaussian mixture by deterministic
1212/// k-means++-style seeding (reusing the `terms::basis` farthest-point k-means,
1213/// a pure function of the data — no clock randomness) followed by EM to the
1214/// configured tolerance.
1215///
1216/// The fit is deterministic given `(data, k, config)`: the seed is the
1217/// farthest-point/k-means center selection, EM is a deterministic map, so
1218/// re-running yields the identical mixture.
1219pub 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    // Deterministic k-means++-style seeding via the shared basis k-means
1226    // (farthest-point init + Lloyd iterations).
1227    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
1278/// Resume EM from the exact state carried by [`GaussianMixtureError`].
1279pub 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    // Certify the CURRENT checkpoint before accepting another EM update. The
1385    // inclusive bound permits exactly `max_iter` accepted updates and then one
1386    // final map evaluation at the resulting checkpoint. Consequently every
1387    // success and every exhaustion pairs its certificate with the exact same
1388    // parameter state; a certificate for theta_t can never be attached to
1389    // theta_{t+1} merely because the work boundary was reached.
1390    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
1521/// Resolution of one observed EM likelihood comparison.
1522///
1523/// `pairwise_mean_with_roundoff` bounds only the final reduction of already
1524/// rounded row log likelihoods. An EM comparison also traverses covariance
1525/// eigendecompositions, precision quadratics, log-sum-exp, the M-step, and a
1526/// second E-step. Treating the reduction bound as a bound for that whole map
1527/// is false precision and turns cancellation at a stationary point into a
1528/// spurious monotonicity violation. The square root of machine epsilon is the
1529/// numerical resolution already required of every configured EM tolerance;
1530/// scaling it by the observed objective magnitude makes the invariant
1531/// independent of data units and of user-selected stopping knobs.
1532fn 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    // The first term bounds the deterministic pairwise additions; the second
1575    // bounds the final division. This tolerance is derived from the actual
1576    // reduction depth and magnitudes, independently of the EM stopping knob.
1577    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
1698/// Distance between two label-aligned Gaussian-mixture states in identifiable
1699/// component-measure coordinates.
1700///
1701/// Raw `(μ, Σ)` coordinates cease to identify the density as a component mass
1702/// tends to zero: an arbitrarily large raw update can then have arbitrarily
1703/// small predictive effect. The finite measures
1704///
1705/// `w`, `w μ`, and `w (Σ + μμᵀ)`
1706///
1707/// are respectively each labeled component's mass and first two raw moments.
1708/// For positive mass they are a bijection with `(w, μ, Σ)`; at zero mass they
1709/// correctly collapse all inert component parameters to the same zero measure.
1710/// EM preserves responsibility-column labels, so no relabeling alignment is
1711/// needed between consecutive states.
1712fn 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
1786/// Global constrained covariance used to seed EM.
1787fn 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// ---------------------------------------------------------------------------
1815// Ring-of-clusters candidate (#2262)
1816// ---------------------------------------------------------------------------
1817//
1818// A free Gaussian mixture treats the component means as unrelated points. That
1819// is the wrong null for a discrete cyclic concept: weekdays and months form
1820// tight clusters, but their component means share a low-dimensional circular
1821// constraint. `RingGaussianMixtureFit` models exactly that density,
1822//
1823//     x | z=j ~ N(c + r u_j, sigma^2 I_2),  ||u_j|| = 1,
1824//
1825// with free mixture weights, a shared center/radius, one angle per component,
1826// and a shared isotropic variance. Its `2k + 3` continuous parameters are
1827// priced by the same BIC-form criterion as the unconstrained mixture's `6k - 1`
1828// parameters in two dimensions.
1829
1830/// Certified Gaussian mixture whose component centers lie on one fitted circle.
1831#[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    /// Free parameters: `k-1` weight logits, center(2), radius(1), `k`
1875    /// component angles, and shared log standard deviation(1).
1876    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    /// Schwarz BIC approximation to negative log evidence, divided by two so
1898    /// it is on the ordinary negative-log-likelihood scale. Lower is better.
1899    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
2010/// Distance between two ring-mixture states in identifiable density space.
2011///
2012/// `center`, `radius`, and `directions` are only a factorization of the actual
2013/// component means `m_j = center + radius * direction_j`. We map those means
2014/// and the shared isotropic variance into density coordinates. Component-mean
2015/// movement is standardized by the common noise scale and weighted by the
2016/// component's predictive mass; the shared variance remains unweighted because
2017/// it changes every component. Thus factorization drift and vanishing-component
2018/// motion cannot block certification, while a resolved change in the fitted
2019/// density still does.
2020fn 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            // Recompute directions at the returned center so the stored angles
2140            // are the exact angular block update belonging to that center.
2141            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(&center, 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
2223/// Fit a deterministic, certified `k`-component isotropic Gaussian mixture
2224/// whose component centers are constrained to a common circle.
2225pub 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(&center, 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// ---------------------------------------------------------------------------
2358// Circular Gaussian density and structured-union candidates (#907)
2359// ---------------------------------------------------------------------------
2360
2361/// Maximum-likelihood fit of a Gaussian-blurred circle in two dimensions.
2362///
2363/// The generative model is
2364///
2365/// `X = center + radius * U + epsilon`,
2366///
2367/// where `U` is uniform on the unit circle and
2368/// `epsilon ~ N(0, noise_variance * I_2)`. Integrating out `U` gives the proper
2369/// Cartesian density
2370///
2371/// `p(x) = exp(-(r^2 + R^2)/(2s)) I0(Rr/s) / (2 pi s)`.
2372///
2373/// Unlike a Gaussian density assigned directly to the nonnegative radius, this
2374/// density is normalized on the plane, remains finite at the center, and has no
2375/// artificial `1/r` singularity. The center is fitted jointly with `(R, s)` by
2376/// latent-angle EM instead of being frozen at the coordinate mean.
2377#[derive(Debug, Clone, Copy)]
2378pub struct CircularGaussianFit2d {
2379    center: [f64; 2],
2380    radius: f64,
2381    noise_variance: f64,
2382}
2383
2384impl CircularGaussianFit2d {
2385    /// Two center coordinates, one radius, and one isotropic noise variance.
2386    pub const NUM_FREE_PARAMETERS: usize = 4;
2387
2388    /// Construct a circular Gaussian from validated model parameters.
2389    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    /// Fit selected rows of a finite two-column coordinate matrix.
2411    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        // Work in a dimensionless chart relative to one observed point. This
2432        // preserves the low-order bits of a small translated circle and makes
2433        // the stopping rule and variance floor scale equivariant.
2434        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        // Moment initialization is exact at the population level. For
2465        // q = ||X-E X||^2,
2466        //   E[q] = R^2 + 2s,  Var(q) = 4s(R^2+s),
2467        // hence R^4 = E[q]^2-Var(q) and s=(E[q]-R^2)/2.
2468        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        // A noiseless observed circle is an unbounded-likelihood boundary.
2485        // Keep the numerical optimizer in a scale-relative interior whose
2486        // width is roundoff, rather than imposing a floor in data units.
2487        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        // Exact EM for the latent circle angle. Given current parameters, the
2496        // conditional mean of U is A(kappa) * (x-c)/||x-c|| with
2497        // A=I1/I0 and kappa=R||x-c||/s. Solving the joint quadratic M-step for
2498        // center and radius avoids the biased `center = sample mean` plug-in.
2499        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            // Evaluate E||X-c-RU||^2 in an explicitly nonnegative form to
2543            // avoid catastrophic cancellation on a very thin ring.
2544            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    /// Fitted circle center.
2587    pub const fn center(self) -> [f64; 2] {
2588        self.center
2589    }
2590
2591    /// Fitted latent-circle radius.
2592    pub const fn radius(self) -> f64 {
2593        self.radius
2594    }
2595
2596    /// Fitted isotropic Cartesian noise variance per coordinate.
2597    pub const fn noise_variance(self) -> f64 {
2598        self.noise_variance
2599    }
2600
2601    /// Proper Cartesian log density at `(x, y)`.
2602    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        // Algebraically this is -(r^2+R^2)/(2s)+log I0(kappa), but this
2609        // rearrangement preserves log I0(kappa) ~= kappa cancellation.
2610        -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    /// Sum the fitted log density over selected rows.
2617    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    /// Fit selected rows and return both the fit and its BIC/2 (lower is
2643    /// better). Keeping fitting and evidence evaluation in one operation makes
2644    /// it impossible to label a likelihood on unrelated data as fitted BIC.
2645    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
2660/// Stable Bessel terms for `kappa = radius * observed_radius / variance`.
2661/// The ordinary finite branch retains the shared approximation exactly. If the
2662/// product itself overflows, only the leading asymptotic
2663/// `log I0(kappa)-kappa = -log(2 pi kappa)/2 + O(1/kappa)` is representable;
2664/// `I1/I0` rounds to one at that scale.
2665fn 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        // The left-to-right product overflowed before a large variance divided
2680        // it back into range. Reconstruct the representable ratio from its log.
2681        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//
2686// A *union* candidate is a small FIXED composite of named component structures
2687// joined by a hard row-responsibility split. Unlike the discrete-mixture rung
2688// (which is one free k-component Gaussian density), a union pins each component
2689// to a specific generative STRUCTURE (a circle, a line, a point cluster) and
2690// asks whether the data is better explained as the disjoint sum of those
2691// structures than by any single pure rung.
2692//
2693// The hard responsibility groups only determine which rows fit each component.
2694// The resulting candidate is one normalized, unlabeled soft-mixture density
2695// `p(y) = Σ_c π_c p_c(y)`, with `π_c = n_c / n`. It is scored on every
2696// training row by that same mixture density used for held-out evaluation. Its
2697// BIC/2 complexity price is `½(Σ_c P_c + m - 1) log(n)`: component
2698// parameters plus the `m - 1` free mixing weights, all on the common sample
2699// scale. This makes the union directly comparable to the other normalized
2700// parametric candidates in the topology race.
2701
2702/// The fixed ladder of structured-union composites. Deterministic and closed:
2703/// open-ended structure search stays owned by #976's move set; these three are
2704/// the only composites the topology race may select.
2705#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
2706pub enum UnionStructure {
2707    /// Two circles (two well-separated periodic loops).
2708    CircleCircle,
2709    /// One circle plus one isolated point cluster (a loop with an outlier blob).
2710    CirclePointCluster,
2711    /// One line (anisotropic cluster) plus one isolated point cluster.
2712    LineCluster,
2713}
2714
2715/// The fixed structured-union ladder, in stable order.
2716pub const UNION_STRUCTURE_LADDER: &[UnionStructure] = &[
2717    UnionStructure::CircleCircle,
2718    UnionStructure::CirclePointCluster,
2719    UnionStructure::LineCluster,
2720];
2721
2722/// The per-component generative structure a union pins each responsibility group
2723/// to. `Line` is a full-covariance Gaussian, while `PointCluster` is the nested
2724/// isotropic Gaussian with `d + 1` parameters. The covariance constraint makes
2725/// line+cluster a genuine structured alternative to a generic two-component
2726/// full-covariance mixture instead of a duplicate candidate. `Circle` is the
2727/// proper Cartesian density of a uniform latent circle convolved with isotropic
2728/// Gaussian noise.
2729#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
2730pub enum UnionComponentKind {
2731    Circle,
2732    Line,
2733    PointCluster,
2734}
2735
2736impl UnionStructure {
2737    /// Stable display name, e.g. `"union_circle+circle"`.
2738    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    /// The fixed ordered component structures of this union.
2747    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    /// Number of components (= the responsibility-split order `m`).
2762    pub const fn num_components(self) -> usize {
2763        self.components().len()
2764    }
2765}
2766
2767/// One fitted component of a union: its pinned structure, the rows used to fit
2768/// it after the hard responsibility split, its free-parameter count, and its
2769/// normalized soft-mixture weight. A component has no standalone BIC inside a
2770/// union: the likelihood is the indivisible `log Σ_c π_c p_c(y)` scored on
2771/// every row.
2772#[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/// A fitted structured-union candidate: the composite kind, the per-component
2781/// fits, its normalized soft-mixture training likelihood, the corresponding
2782/// BIC-form negative-log-evidence, and the complete free-parameter count.
2783#[derive(Debug, Clone)]
2784pub struct UnionStructureFit {
2785    pub structure: UnionStructure,
2786    pub components: Vec<UnionComponentFit>,
2787    /// `Σ_i log(Σ_c π_c p_c(y_i))` over all training rows.
2788    pub log_likelihood: f64,
2789    /// `-log_likelihood + ½ total_parameters log(n)` (lower wins).
2790    pub bic: f64,
2791    /// `Σ_c P_c + (m - 1)`, including the free mixing weights.
2792    pub total_parameters: usize,
2793}
2794
2795/// Hard responsibility split of `0..n` into `m` groups by argmax of the
2796/// deterministic `m`-component Gaussian-mixture responsibilities. Reuses the
2797/// mixture rung's seeding + EM so the split is a pure function of the data and
2798/// `m` (no clock). Returns one row-index vector per component.
2799pub 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    // Hard assignment by argmax per-component log responsibility.
2819    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
2861/// Fit one structured-union candidate. The hard split identifies component
2862/// training groups; heterogeneous component roles are assigned by evaluating
2863/// every unique role permutation. Each assignment is then scored as one
2864/// normalized soft mixture on every training row.
2865///
2866/// Returns an error if no role assignment identifies every component. The
2867/// fixed-ladder caller propagates that failure for the whole declared family;
2868/// it never ranks a survivor subset.
2869pub 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
2888/// Fit the whole fixed union ladder and rank in-class by normalized soft-mixture
2889/// BIC (lower wins). The declared ladder is one selection family: if any
2890/// eligible structure fails, the whole comparison fails with every per-structure
2891/// error rather than silently selecting from an easier survivor set.
2892pub 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; // cheaper composite wins ties
2919                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
2940/// Apply the structured-union admission policy to one circle group and return
2941/// its complete deterministic row set.
2942/// The `n > P` rule belongs to this composite ladder; it is not an intrinsic
2943/// restriction of circular-Gaussian likelihood fitting or BIC itself.
2944fn 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/// The shared fitted-density representation used by both in-sample BIC and
2956/// held-out predictive scoring. `Line` and `PointCluster` deliberately share
2957/// the Gaussian evaluator after fitting, but differ in covariance constraints
2958/// and parameter count.
2959#[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    /// `log π_c + log p_c(y)` for one eval row.
2993    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
3019/// Fit and score a union for an already-established hard partition. This seam
3020/// makes role assignment explicitly independent of the arbitrary component
3021/// labels emitted by the responsibility split.
3022fn 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(&current.bic).is_lt(),
3081                    None => true,
3082                };
3083                // Unique assignments are generated in canonical order. Keeping
3084                // the earlier assignment on an exact score tie is deterministic.
3085                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
3255/// Range-safe and translation-accurate scalar mean chart. The normalized mean
3256/// offset is retained separately from the rounded absolute mean so residuals
3257/// use `(x-origin)-offset` rather than losing a fractional offset at a large
3258/// common translation. FMA also lets a subnormal offset affect the correctly
3259/// rounded absolute mean without first rounding that offset to zero.
3260fn 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
3298/// Maximum-likelihood isotropic Gaussian fit in stable per-column mean charts.
3299fn 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
3403/// Per-point held-out log predictive density of a structured-union candidate:
3404/// fit the component densities on `train` and score each row of `eval` as the
3405/// soft mixture `log Σ_c π_c p_c(y)`. This is the cross-class stacking column
3406/// source for a union (the analogue of [`GaussianMixtureFit::per_point_log_density`]).
3407pub 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/// One fitted model in a REML/LAML evidence comparison.
3425#[derive(Clone, Debug)]
3426pub struct RemlCandidate {
3427    pub index: usize,
3428    pub name: String,
3429    /// Minimised REML/LAML cost. Lower is better. This is the model's reported
3430    /// evidence headline (`Model.evidence`), kept verbatim in the score table.
3431    pub score: f64,
3432    pub edf: Option<f64>,
3433    /// Log-likelihood at the converged mode, on the engine's
3434    /// constants-omitted scale (same as `gam_inference::model_comparison`).
3435    /// Present when the fit carries it; `None` for legacy payloads.
3436    pub log_lik: Option<f64>,
3437    /// Response-family tag (e.g. "gaussian", "gamma", "binomial"). Carried so
3438    /// `compare_reml_fits` can REFUSE to rank fits whose REML/LAML scores are on
3439    /// incomparable base measures (a cross-family comparison is meaningless;
3440    /// #1384). `None` for legacy payloads that did not record it — those are not
3441    /// guarded (back-compatible), but every current FFI candidate carries it.
3442    pub family: Option<String>,
3443    /// Number of observations the fit was trained on. Carried so
3444    /// `compare_reml_fits` can REFUSE to rank fits made on a different number of
3445    /// observations (hence different data): `−2·loglik` and the REML/LAML
3446    /// evidence grow with `n`, so a score difference between two fits with
3447    /// different `n` is not a Bayes factor — the same incomparability the family
3448    /// guard already rejects. `None` for payloads that do not record it (legacy /
3449    /// O(n) scan smoothers), which the guard treats as unconstrained.
3450    pub n_obs: Option<usize>,
3451}
3452
3453impl RemlCandidate {
3454    /// Cost used to RANK candidates and pick the winner.
3455    ///
3456    /// The REML/LAML marginal-likelihood evidence headline (`score`) does NOT
3457    /// reliably Occam-penalise an added pure-noise smooth: on `y ~ s(x)` vs
3458    /// `y ~ s(x) + s(z)` with `z ⟂ y`, the augmented model's evidence is
3459    /// *lower* (apparently better) by a few nats on essentially every dataset,
3460    /// because the Gaussian REML Occam pair `½(log|H| − log|S|₊)` collapses
3461    /// toward zero for a finite-`λ̂` null term while that term still spends a
3462    /// few effective degrees of freedom fitting noise (issue #1362).
3463    ///
3464    /// The conditional AIC `−2ℓ + 2·edf` prices exactly those spent degrees of
3465    /// freedom and discriminates correctly: it penalises the noise smooth
3466    /// (Δ ≈ +15 nats) yet rewards a genuinely relevant smooth (Δ ≈ −650),
3467    /// preserving power. We therefore rank on the conditional AIC whenever both
3468    /// the log-likelihood and the effective degrees of freedom are available,
3469    /// and fall back to the raw evidence headline otherwise. The reported
3470    /// `score_table` still carries the unaltered evidence (`reml_score`), so
3471    /// `Model.evidence` / `bayes_factor_vs` stay consistent with the table.
3472    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    /// Cost gap from the winning model on the SAME scale used to order the
3495    /// ranking (`ranking_score`, the Occam-penalised conditional AIC where
3496    /// available, issue #1362). The winner is `argmin ranking_score`, so this
3497    /// is `>= 0` for every row by construction — it never contradicts the
3498    /// declared winner (issue #1465). `score` still carries the raw REML/LAML
3499    /// evidence so it stays consistent with `Model.evidence`.
3500    pub delta: f64,
3501    /// Akaike evidence ratio of the winner over this row on the ranking scale.
3502    /// `delta` is a conditional-AIC gap (a −2·log / deviance-scale quantity), so
3503    /// the evidence ratio is `exp(½·delta) >= 1` (Burnham & Anderson), NOT
3504    /// `exp(delta)` — the latter squares the intended ratio (issues #1465, #2124).
3505    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/// Log Bayes factor of model `a` over model `b` from minimised REML/LAML costs.
3519#[inline]
3520pub fn log_bayes_factor(reml_score_a: f64, reml_score_b: f64) -> f64 {
3521    reml_score_b - reml_score_a
3522}
3523
3524/// Compare fitted models by the single evidence ordering contract used by
3525/// topology ranking and seed screening: lower finite cost wins, with stable
3526/// original-order tie handling.
3527pub 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    // Fail-loud comparability guard (#1384): REML/LAML evidence scores are only
3532    // comparable across fits of the SAME response family — a Gaussian score and
3533    // a Gamma score live on different log-density base measures, so their
3534    // difference is not a Bayes factor. Ranking them anyway returns a confident
3535    // but meaningless winner. Refuse when two candidates carry DIFFERENT family
3536    // tags. Candidates with no family tag (`None`, legacy payloads) are not
3537    // constrained, so this never spuriously rejects an older saved model.
3538    {
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    // Fail-loud comparability guard (#1384 sibling): AIC / REML-LAML evidence are
3555    // only comparable across fits of the SAME response on the SAME observations.
3556    // `−2·loglik` (and the marginal-likelihood headline) grow with the number of
3557    // observations `n`, so two fits with different `n` live on incomparable
3558    // scales and their score gap is not a Bayes factor — comparing an n=500 and
3559    // an n=100 fit of the same DGP otherwise declares the n=100 model the winner
3560    // purely because fewer points give a less-negative total log-likelihood.
3561    // Refuse when two candidates carry DIFFERENT observation counts. Candidates
3562    // with no count (`None`, legacy / O(n) scan payloads) are unconstrained, so
3563    // this never spuriously rejects a fit that simply did not record `n`.
3564    {
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                // Rank/winner on the Occam-penalised conditional AIC where it is
3589                // available (issue #1362); falls back to the raw evidence score.
3590                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    // The ranking `delta` / `bayes_factor` must be measured on the SAME scale
3601    // that orders the table — the `ranking_score` (Occam-penalised conditional
3602    // AIC where available, issue #1362). `candidates[0]` is the winner =
3603    // `argmin ranking_score`, so its ranking score IS the minimum; every row's
3604    // ranking-scale gap is then `>= 0` and its Bayes factor `>= 1`, never
3605    // contradicting the declared winner (issue #1465). Computing these against
3606    // the AIC winner's *raw REML* — which is not the minimum raw REML once AIC
3607    // and REML disagree — produced negative deltas and Bayes factors < 1 for
3608    // non-winner rows.
3609    let best_ranking_score = candidates[0].ranking_score();
3610    // The raw-REML `score_table` stays on the raw evidence scale (consistent
3611    // with `Model.evidence` / `bayes_factor_vs`), but is referenced to the
3612    // genuine minimum raw REML so its best-over-model Bayes factors are also
3613    // coherent (`>= 1`), rather than to whichever row happens to sit at index 0.
3614    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        // `ranking_score` is the conditional AIC (`−2·loglik + 2·edf`), a −2·log /
3623        // deviance-scale cost, so `delta` is a full ΔAIC gap. The Akaike evidence
3624        // ratio for an AIC gap Δ is `exp(−½Δ)` (Burnham & Anderson evidence ratio),
3625        // hence the winner-over-row Bayes factor is `exp(½·delta)`. Reporting
3626        // `delta.exp()` squared the intended ratio (issue #2124). `delta` itself is
3627        // left on the AIC scale on purpose — only its exp() conversion is halved.
3628        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    // The winner is decided by `ranking_score` (the Occam-penalised conditional
3646    // AIC where available, issue #1362), which can disagree in sign with the raw
3647    // evidence Bayes factor for a noise-augmented model. Summarise the actual
3648    // decision margin so the headline never contradicts the chosen winner.
3649    let evidence_summary = if let Some(runner_up) = candidates.get(1) {
3650        let margin = runner_up.ranking_score() - candidates[0].ranking_score();
3651        // `margin` is a conditional-AIC gap (−2·log scale), so the Akaike evidence
3652        // ratio is `exp(−½·margin)`; `format_bayes_factor` formats `exp()` of its
3653        // argument, so pass the halved margin to headline `exp(½·margin)` rather
3654        // than the squared `exp(margin)` (issue #2124).
3655        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
3708// ---------------------------------------------------------------------------
3709// Laplace evidence
3710// ---------------------------------------------------------------------------
3711
3712/// Single canonical Laplace evidence at the inner-loop fixed point.
3713///
3714/// Returns negative log evidence:
3715///
3716/// ```text
3717/// V(ρ, T) = F(β*, u*; ρ, T)
3718///         + 0.5 log|H|
3719///         - 0.5 log|S_pen(ρ)|+
3720///         - 0.5 (dim(H) - rank(S_pen)) log(2π).
3721/// ```
3722///
3723/// The last term is the rank-aware Tierney-Kadane normalizer:
3724/// `log p(y|T) ≈ -V`, with `0.5 log|2πH⁻¹| - 0.5 log|2πS⁻¹|`.
3725///
3726/// The `H` log-determinant is computed from the arrow factorization
3727///
3728/// ```text
3729/// log|H| = Σ_i log|H_uu_i| + log|A|
3730/// ```
3731///
3732/// (proposal §3.4 / §7) using the **undamped** per-row Cholesky factors
3733/// `cache.htt_factors_undamped` and the **undamped** Schur factor.
3734///
3735/// `penalty_log_det` is `log|S_pen(ρ)|+` — the prior penalty
3736/// pseudo-logdet from `crate::reml::penalty_logdet` (proposal
3737/// §3.6). It must NOT be confused with the arrow Schur log-det, which
3738/// this function recomputes internally from `logdet_source`.
3739///
3740/// `residual_objective` is `F(β*, u*; ρ, T)` at the inner optimum. The
3741/// envelope theorem (proposal §3.2) makes this the only `F`-related
3742/// contribution.
3743///
3744/// `effective_dim` is `dim(H)` after constraints/projections and
3745/// `penalty_rank` is `rank(S_pen)`. Their difference is the unpenalized
3746/// nullspace dimension that remains in the Laplace integral.
3747///
3748/// # Errors
3749///
3750/// Returns `f64::NAN` if the exact factor path is incoherent and no HVP
3751/// fallback is supplied, or if the supplied dimensions are non-finite.
3752pub 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
3775/// Compute the Hessian logdet from exact arrow factors or an HVP fallback.
3776pub 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
3794/// Log determinant of an SPD operator supplied by HVP callback.
3795///
3796/// The dispatch boundary intentionally matches
3797/// `ANALYTIC_LOGDET_DENSE_DIM_THRESHOLD` in `terms::analytic_penalties`:
3798/// small operators are materialized and diagonalized exactly; larger ones use
3799/// Rademacher stochastic Lanczos quadrature.
3800pub 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
4015/// Authoritative factored-arrow evidence log-determinant.
4016///
4017/// This reads the stored joint value that cache construction records for the
4018/// operator used by selected-inverse/adjointers. It intentionally does not
4019/// reconstruct a determinant from row and Schur pieces for damped caches; if
4020/// construction did not record an exact joint log-det, evidence must refuse the
4021/// cache and route to an explicit matrix-free fallback.
4022pub 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
4035// ---------------------------------------------------------------------------
4036// IFT cascade: ∂u*/∂β → ∂β*/∂ρ → ∂u*/∂ρ
4037// ---------------------------------------------------------------------------
4038
4039/// Tier-1 IFT sensitivity `∂u_i*/∂β = -H_uu_i⁻¹ H_uβ_i`.
4040///
4041/// Concatenated row-major to a single `(N·d) × K` dense matrix. Each
4042/// row block is solved with the **undamped** Cholesky factor. Proposal
4043/// §2.2 / §7.
4044pub 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    // Allocate scratch at max_d; per-row slice is ..di.
4054    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        // Solve H_uu_i Y = H_uβ_i column by column.
4060        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            // The Tier-2 IFT assembler is built only when the family's
4065            // capability surface promises cached `H_tβ` row products.
4066            if !cache.apply_htbeta_row(i, beta_basis.view(), &mut rhs_i) {
4067                // SAFETY: reaching `false` means a family declared the cache
4068                // available but failed to populate it — contract violation.
4069                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
4080/// Coupling components of a symmetric coefficient Hessian: the connected
4081/// components of the graph whose vertices are coefficient indices `0..p` and
4082/// whose edges are the structurally nonzero off-diagonal entries of `H` (#779).
4083///
4084/// Returns a length-`p` vector of component labels in `0..num_components`,
4085/// where two indices share a label iff they are connected through a chain of
4086/// nonzero `H[i,j]` couplings. This is the exact structural partition the
4087/// cone-of-influence sensitivity reuse is keyed on: a smoothing-parameter move
4088/// whose stationarity-gradient derivative `∂g/∂ρ` is supported only inside one
4089/// component can change `β = -H⁻¹ ∂g/∂ρ` only inside that same component, so
4090/// the sensitivity of every *other* component is provably unchanged and may be
4091/// reused unrecomputed (lazy/local propagation).
4092///
4093/// The nonzero test is exact (`!= 0.0`), matching the structural-coupling gate
4094/// used elsewhere for the joint inner Hessian: a tolerance would risk dropping a
4095/// genuine (small) coupling edge and silently biasing the propagated sensitivity
4096/// — the failure mode #779/#740 explicitly guard against. A block-diagonal `H`
4097/// yields the all-singletons partition (one component per block-decoupled
4098/// coordinate); a fully coupled `H` yields a single component (no shortcut, the
4099/// full joint solve is required — and is what the non-coned path performs).
4100pub 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    // Union-find with path compression and union by size.
4106    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            // Symmetric structure: an edge exists if either triangle is nonzero,
4120            // so a numerically one-sided fill still couples the two indices.
4121            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    // Relabel roots to a dense `0..num_components` range, preserving
4137    // first-seen order so labels are deterministic.
4138    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
4157/// The cone of influence of a single stationarity-gradient derivative column
4158/// whose support (the coefficient indices where `∂g/∂ρ_k` is nonzero) lies in
4159/// `support`: the set of coefficient indices in the same coupling component(s)
4160/// as that support, given precomputed `labels` from [`coupling_components`].
4161///
4162/// `β_k = -H⁻¹ ∂g/∂ρ_k` is exactly zero outside this cone, so a confined solve
4163/// (or reuse of a cached zero) is exact, not an approximation. An empty support
4164/// (a structurally inactive `ρ_k`, e.g. a rank-0 or out-of-range penalty block)
4165/// yields an empty cone: the sensitivity is identically zero and no solve is
4166/// needed at all.
4167pub 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
4185/// Tier-2 IFT sensitivity `∂β*/∂ρ = -A⁻¹ ∂g_red/∂ρ` (proposal §2.4 /
4186/// §7).
4187///
4188/// `dg_red_drho` is the `K × R` matrix whose `a`-th column is `q_a =
4189/// ∂g_red/∂ρ_a`. Returns the `K × R` matrix `β_ρ`.
4190///
4191/// Returns `None` if the Schur factor is unavailable (PCG mode) or was
4192/// built from a damped operator, or if any solved entry is non-finite;
4193/// callers must not silently substitute an approximation. The solve is
4194/// the one sensitivity operator (#935) — this site holds no private H⁻¹
4195/// convention of its own.
4196pub 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// ---------------------------------------------------------------------------
4211// ∂V/∂ρ — analytic optimized-evidence gradient via IFT mode response
4212// ---------------------------------------------------------------------------
4213
4214/// IFT terms needed to differentiate the optimized Laplace evidence through
4215/// the fitted mode `(β*(ρ), u*(ρ))`.
4216///
4217/// For each hyperparameter `ρ_a`, the correction added to the direct trace is
4218///
4219/// ```text
4220/// F_β · β_a + F_u · u_a
4221/// + 0.5 (∂_β log|H| · β_a + ∂_u log|H| · u_a).
4222/// ```
4223///
4224/// At an exact KKT point the value-gradient pieces are zero, but they are
4225/// explicit here so the exported gradient matches the optimized objective
4226/// whenever callers carry a certified nonzero residual correction.
4227#[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
4237/// Contract the IFT mode-response columns into the optimized-evidence
4238/// gradient correction.
4239pub 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
4270/// Per-`ρ` optimized-evidence gradient (proposal §3.7 / §3.8 split):
4271///
4272/// ```text
4273/// ∂V/∂ρ_a =
4274///       F_{ρ_a}                                  (value part)
4275///   + 0.5 tr(H⁻¹ H_{ρ_a})                        (direct Hessian)
4276///   + F_x · x_{ρ_a}
4277///   + 0.5 (∂_x log|H|) · x_{ρ_a}                 (IFT mode response)
4278///   - 0.5 tr(S_pen⁺ S_{pen,ρ_a})                 (penalty pseudo-logdet)
4279/// ```
4280/// where `x = (β, u)`.
4281///
4282/// The `tr(H⁻¹ H_{ρ_a})` trace is computed via the arrow structure
4283/// (proposal §3.5 / §3.10):
4284///
4285/// ```text
4286/// tr(H⁻¹ H_{ρ_a}) = Σ_i tr(H_uu_i⁻¹ ∂_{ρ_a} H_uu_i) + tr(A⁻¹ ∂_{ρ_a} A).
4287/// ```
4288///
4289/// `value_rho[a] = F_{ρ_a}` (envelope theorem, proposal §3.2).
4290/// `huu_drho[i][a]` is `∂H_uu_i/∂ρ_a` as a `d × d` matrix.
4291/// `hbb_drho[a]` is `∂H_ββ/∂ρ_a` as a `K × K` matrix.
4292/// `htbeta_drho[i][a]` is `∂H_uβ_i/∂ρ_a` as a `d × K` matrix.
4293/// `pen_logdet_drho[a]` is `∂_{ρ_a} log|S_pen|+`.
4294/// `ift_terms` carries `∂β*/∂ρ`, `∂u*/∂ρ`, and the already-contracted
4295/// mode derivatives of `F` and `log|H|`.
4296///
4297/// Returns the per-`ρ` gradient. Returns a NaN-filled vector when the
4298/// cache has no undamped Schur factor (PCG mode).
4299pub 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    // Precompute Y_i = H_uu_i⁻¹ H_uβ_i (di × K). Used by both the Schur
4355    // derivative formula (§3.5) and the row trace `tr(H_uu_i⁻¹ ∂H_uu_i)`.
4356    let mut y_blocks: Vec<Array2<f64>> = Vec::with_capacity(n);
4357    let mut beta_basis = Array1::<f64>::zeros(k);
4358    // Scratch sized to max_d; per-row slice is ..di.
4359    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            // Same H_tβ cache contract as the IFT du/dβ and du/dρ paths.
4369            if !cache.apply_htbeta_row(i, beta_basis.view(), &mut rhs_i) {
4370                // SAFETY: `false` means the family declared the cache
4371                // available but did not populate it — contract violation.
4372                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    // Outer-hoisted scratch reused across all (a, i) iterations.
4384    // Sized to max_d for trace_rhs and da_tmp; per-row slices used below.
4385    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        // Part 1: F_{ρ_a} envelope contribution.
4390        let mut grad = value_rho[a];
4391
4392        // Part 2a: Σ_i tr(H_uu_i⁻¹ ∂H_uu_i).
4393        // tr(H_uu_i⁻¹ M_i) = tr(L_iᵀ⁻¹ L_i⁻¹ M_i). Compute as the sum
4394        // over columns: solve L_i Lᵀ x = e_c for the c-th column of
4395        // M_i, then take its c-th component. Equivalently and more
4396        // cheaply, build (H_uu_i⁻¹ M_i) by solving column-by-column
4397        // and take its diagonal sum.
4398        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        // Part 2b: tr(A⁻¹ ∂A) where (proposal §3.5)
4414        //     ∂A = ∂H_ββ
4415        //          - Σ_i (∂H_uβ_i)ᵀ Y_i
4416        //          - Σ_i Y_iᵀ (∂H_uβ_i)
4417        //          + Σ_i Y_iᵀ (∂H_uu_i) Y_i.
4418        // We accumulate ∂A as a dense `K × K` matrix, then evaluate
4419        // tr(A⁻¹ ∂A) by `Σ_j (A⁻¹ ∂A)[j, j]` via column solves of the
4420        // Schur Cholesky.
4421        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]; // di × K
4426            let yi = &y_blocks[i]; // di × K
4427            // - (∂H_uβ_i)ᵀ Y_i
4428            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            // - Y_iᵀ (∂H_uβ_i)
4438            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            // + Y_iᵀ (∂H_uu_i) Y_i
4448            let dhuu = &huu_drho[i][a];
4449            // tmp = (∂H_uu_i) Y_i  (di × K) — use a slice of the hoisted buffer.
4450            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            // da += Y_iᵀ tmp
4461            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        // tr(A⁻¹ ∂A) via column solves.
4473        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        // Part 3: -0.5 ∂_{ρ_a} log|S_pen|+.
4486        grad -= 0.5 * pen_logdet_drho[a];
4487
4488        out[a] = grad;
4489    }
4490    out
4491}
4492
4493// ---------------------------------------------------------------------------
4494// Topology selection
4495// ---------------------------------------------------------------------------
4496
4497/// Enumerate the candidate topologies, rank by normalized negative log
4498/// evidence, and return the winner. Failed/excluded candidates (proposal
4499/// §6.11) are appended at the end of `ranking` and are never the winner.
4500///
4501/// The caller fits each topology separately (proposal §4.2) and supplies
4502/// the resulting `TopologyCandidate` records. This function is purely
4503/// the discrete comparator + tie breaker.
4504///
4505/// # Tie-breaking
4506///
4507/// Per proposal §4.6: if normalized `|score_a - score_b| <= tie_tolerance`,
4508/// prefer the simpler topology by `TopologyKind::complexity_rank` (flat <
4509/// periodic < sphere < torus). The `tie` flag in the result records whether
4510/// such a tie occurred at the top of the ranking.
4511///
4512/// # Panics
4513///
4514/// Panics if `candidates` is empty after filtering out non-finite
4515/// scores. Proposal §6.11 explicitly forbids silent fallback to a
4516/// default topology; callers must handle the empty-candidate case
4517/// before invocation.
4518pub fn select_topology(
4519    candidates: &[TopologyCandidate],
4520    options: TopologySelectOptions,
4521) -> SelectedTopology {
4522    // Split valid and excluded.
4523    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    // Sort by normalized negative log evidence (ascending = best first),
4548    // breaking ties by complexity_rank (smaller wins). The shared selector is
4549    // the single lower-is-better ordering contract used by topology ranking,
4550    // seed screening, and REML model comparison (#782).
4551    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    // Detect numerical tie at the top.
4567    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 tied, prefer simpler topology among the tied prefix.
4576    if tie {
4577        let top_score = topology_selection_score(&valid[0], options.score_scale);
4578        // Find the tied prefix range.
4579        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        // Sort the tied prefix by complexity_rank ascending.
4587        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
4618// ---------------------------------------------------------------------------
4619// Cache verification helpers
4620// ---------------------------------------------------------------------------
4621
4622/// Verifies the `ArrowSchurSystem` dimensions match the cache. Used as
4623/// a debug-time precondition; never silently masks shape errors
4624/// (proposal §6.9 — sign and shape errors must be loud).
4625pub 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// ---------------------------------------------------------------------------
4635// #1026 hybrid curved + linear-tail dictionary split-selection
4636// ---------------------------------------------------------------------------
4637//
4638// COMMON-EVIDENCE NOTE (#1202): the candidates BOTH fit the same data — the
4639// atom's leave-this-atom-out response residual `y_resp` (the response with every
4640// other atom's contribution removed). The curved candidate predicts the atom's
4641// actual mass-scaled contribution `a_k·γ_k`, the linear candidate the best
4642// mass-weighted straight line fit to `y_resp`. Because the curved family's
4643// `Θ = 0` member reproduces the linear prediction exactly, linear IS the nested
4644// `Θ = 0` sub-model on common data, so the "match-or-beat" statements below are a
4645// genuine data-level comparison: the curved candidate wins only when fitting the
4646// response residual better than its own straight projection pays for its extra
4647// parameters. See `crate::terms::sae::hybrid_split` for the residual assembly.
4648//
4649// The per-slot adjudication uses the SAME rank-aware Laplace evidence criterion
4650// the union/mixture rungs use (`−V = NLE`, lower wins), comparing the data-fit +
4651// complexity cost of the curved contribution against that of the straight line.
4652//
4653// ## The turning floor (Θ → 0) and the curved ceiling (Θ large)
4654//
4655// Per slot, the curved candidate fits the response residual with its actual
4656// mass-scaled contribution `a_k·γ_k` (data-fit `½·curved_rss`) and pays a larger
4657// free-parameter price `P_curved > P_linear`; the linear candidate fits the same
4658// residual with its best straight line (data-fit `½·linear_rss ≥ ½·curved_rss`
4659// whenever the curve beats its own straight projection) at a smaller price,
4660// charged with its genuine weighted Gram logdet `p·(log w_sum + log s_tt)`
4661// (#1203). Hence:
4662//
4663//   * Θ → 0 (the residual is straight): the curve and the line fit it equally, so
4664//     the cheaper LINEAR candidate wins — the turning floor / nested dominance. A
4665//     curved parameterization "buys nothing" on an already-straight residual.
4666//   * Θ large (a genuinely turning residual): the line's data-fit residual
4667//     exceeds the curved atom's extra parameter price, so CURVED wins. (Whether
4668//     curved wins also depends on the coordinate spread `s_tt` and amplitude, via
4669//     the honest logdet — a tightly-spread, mildly-curved residual can still
4670//     prefer the cheaper line.)
4671//
4672// The crossover is governed by the documented shatter law: a linear SAE shatters
4673// a feature of total turning Θ into `N(ε) ≈ Θ/(2√(2ε))` rank-1 directions at
4674// relative reconstruction error ε, so the curved advantage scales as `Θ/√ε`. We
4675// use the fitted turning Θ (`sae::chart_canonicalization::d1_atom_fitted_turning`)
4676// as the decision FEATURE: it both (a) sharpens the evidence comparison into a
4677// falsifiable per-atom prediction and (b) provides the exact-zero dominance
4678// guard — when an atom's fitted turning is identically zero, the curved fit has
4679// no curvature to price and the linear special case is selected by construction,
4680// independent of finite-sample evidence noise.
4681
4682/// Which atom parameterization a hybrid-dictionary slot selects: a CURVED atom
4683/// (a `latent_dim ≥ 1` curved basis whose decoded image may turn) or its LINEAR
4684/// special case (the euclidean-d=1-linear atom — one straight decoder direction,
4685/// `γ(t) = t·b`, fitted turning `Θ = 0`).
4686#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
4687pub enum HybridAtomParam {
4688    /// The curved atom (`latent_dim ≥ 1`), priced at its full coefficient count.
4689    Curved { latent_dim: usize },
4690    /// The linear special case: one decoder direction, zero turning.
4691    Linear,
4692}
4693
4694impl HybridAtomParam {
4695    /// Stable display name for logs and tests.
4696    pub const fn as_str(self) -> &'static str {
4697        match self {
4698            HybridAtomParam::Curved { .. } => "curved",
4699            HybridAtomParam::Linear => "linear",
4700        }
4701    }
4702
4703    /// `true` iff this is the linear special case (the linear tail).
4704    pub const fn is_linear(self) -> bool {
4705        matches!(self, HybridAtomParam::Linear)
4706    }
4707}
4708
4709/// One fitted candidate parameterization for a single hybrid-dictionary atom
4710/// slot, scored on the COMMON rank-aware Laplace scale (`−V = NLE`, lower wins,
4711/// identical to the union/mixture rungs). The curved and linear candidates for
4712/// the SAME slot are fit on the same rows AND the same data (the atom's response
4713/// residual, #1202), so their NLEs are directly comparable; the structural
4714/// difference is the curved candidate's larger free-parameter price and whatever
4715/// data-fit it buys with its curvature.
4716#[derive(Debug, Clone, Copy)]
4717pub struct HybridAtomCandidate {
4718    pub param: HybridAtomParam,
4719    /// Rank-aware Laplace negative-log-evidence on the common scale (lower wins).
4720    pub negative_log_evidence: f64,
4721    /// Free-parameter count this candidate is charged for (the complexity price).
4722    pub num_parameters: usize,
4723    /// The candidate's fitted total turning `Θ = ∫κ ds` of its decoded curve, if
4724    /// the basis admits an analytic second jet. `Some(0.0)` for a linear atom (a
4725    /// straight image has no turning); `None` when the turning is honestly
4726    /// unavailable (no second jet / degenerate curve) — never fabricated.
4727    pub fitted_turning: Option<f64>,
4728}
4729
4730impl HybridAtomCandidate {
4731    /// A linear special-case candidate: exact zero turning by construction.
4732    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    /// A curved candidate of the given latent dimension, with its fitted turning.
4742    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/// The evidence-selected parameterization for one hybrid-dictionary atom slot:
4758/// the winning candidate, plus the curved/linear NLEs that decided it (for the
4759/// EV-vs-Θ diagnostic and the tie-break audit trail).
4760#[derive(Debug, Clone, Copy)]
4761pub struct HybridAtomChoice {
4762    pub param: HybridAtomParam,
4763    /// The winning candidate's NLE.
4764    pub negative_log_evidence: f64,
4765    /// The winning candidate's free-parameter price.
4766    pub num_parameters: usize,
4767    /// The curved candidate's fitted turning `Θ` (the decision feature). `None`
4768    /// when no curved candidate offered an analytic turning.
4769    pub curved_turning: Option<f64>,
4770    /// `NLE_linear − NLE_curved`: the evidence margin the curved fit won (or lost,
4771    /// if negative) over the linear special case at this slot. Positive ⇒ curved
4772    /// bought more evidence than its parameter price; ≤ 0 ⇒ the dominance floor
4773    /// keeps the linear tail.
4774    pub curved_evidence_margin: f64,
4775}
4776
4777/// Below this fitted turning the curved candidate is treated as straight: its
4778/// curvature is numerically indistinguishable from zero, so the dominance floor
4779/// (the linear special case is cheaper at equal likelihood) is enforced by
4780/// construction rather than left to finite-sample evidence noise. This is the
4781/// exact-zero guard from the `Θ → 0 ⇒ N(ε) → 0` limit of the shatter law, not a
4782/// tunable knob: it is the curvature scale below which `‖γ' ∧ γ''‖` is at the
4783/// floor of the Simpson quadrature for a genuinely straight image.
4784pub const HYBRID_LINEAR_TURNING_FLOOR: f64 = 1e-9;
4785
4786/// Adjudicate the curved-vs-linear parameterization for ONE hybrid-dictionary
4787/// atom slot by the common rank-aware Laplace evidence criterion.
4788///
4789/// Selection rule (all on the single `NLE = −V` scale, lower wins):
4790///
4791///  1. **Dominance floor (Θ → 0).** If the curved candidate's fitted turning is
4792///     `Some(Θ)` with `Θ ≤ HYBRID_LINEAR_TURNING_FLOOR` and a linear candidate
4793///     exists, select LINEAR. A straight curved fit recovers no likelihood the
4794///     linear special case does not, and the linear atom is strictly cheaper, so
4795///     it cannot lose — we enforce that exactly instead of trusting evidence
4796///     noise at the floor.
4797///  2. **Evidence comparison.** Otherwise select the candidate with the smaller
4798///     `NLE`. The curved candidate wins only when its extra curvature lowers the
4799///     NLE by MORE than its extra parameter price — the `Θ/√ε` crossover, decided
4800///     here by the evidence numbers themselves, not by fiat. This is a
4801///     common-data comparison (both candidates fit the atom's response residual,
4802///     see `crate::terms::sae::hybrid_split`) in which linear is the curved
4803///     family's nested `Θ = 0` sub-model (#1202): the curved candidate cannot be
4804///     charged its extra parameters to fit the residual no better than its own
4805///     straight projection, and a tightly-spread, mildly-curved residual can
4806///     still prefer the cheaper line.
4807///  3. **Tie-break.** Exact NLE ties go to the cheaper (fewer-parameter)
4808///     candidate — i.e. linear — preserving the strict-generalization guarantee
4809///     that the hybrid never pays for curvature it does not need.
4810///
4811/// `candidates` must contain at most one linear and at most one curved candidate
4812/// for the slot; returns `None` only if `candidates` is empty.
4813pub 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    // (1) Exact-zero dominance floor: a straight curved fit yields to the linear
4826    // special case by construction.
4827    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    // (2)+(3) Evidence argmin with the cheaper candidate winning exact ties.
4840    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/// The evidence-selected split for a whole hybrid dictionary: the per-atom
4859/// curved-vs-linear choices and the dictionary-level aggregates the EV-vs-Θ
4860/// frontier reports against.
4861#[derive(Debug, Clone)]
4862pub struct HybridSplitSelection {
4863    /// One adjudicated choice per atom slot, in slot order.
4864    pub atoms: Vec<HybridAtomChoice>,
4865    /// `Σ NLE` across the selected per-atom parameterizations — the dictionary's
4866    /// summed rank-aware Laplace negative-log-evidence (lower wins). Because each
4867    /// slot picks the argmin over {curved contribution, best straight line to the
4868    /// response residual}, this is ≤ the sum of the per-slot LINEAR-candidate
4869    /// NLEs. The linear baseline is the best straight line fit to each atom's
4870    /// leave-this-atom-out RESPONSE residual (#1202), the curved family's nested
4871    /// `Θ = 0` member on common data — so this is a genuine data-level
4872    /// match-or-beat dominance, not a post-hoc curve-simplification one.
4873    pub total_negative_log_evidence: f64,
4874    /// `Σ P` across the selected parameterizations — the dictionary's total
4875    /// free-parameter price (the matched-active-budget accounting).
4876    pub total_parameters: usize,
4877    /// Count of slots that selected the curved parameterization.
4878    pub curved_atom_count: usize,
4879}
4880
4881impl HybridSplitSelection {
4882    /// Count of slots that selected the linear special case (the linear tail).
4883    pub fn linear_atom_count(&self) -> usize {
4884        self.atoms.len() - self.curved_atom_count
4885    }
4886
4887    /// `true` iff every slot selected linear — the pure-linear limit, reached
4888    /// when every feature is straight (all `Θ → 0`).
4889    pub fn is_pure_linear(&self) -> bool {
4890        self.curved_atom_count == 0 && !self.atoms.is_empty()
4891    }
4892
4893    /// `true` iff every slot selected curved — the pure-curved limit, reached
4894    /// when every feature turns enough to pay for curvature.
4895    pub fn is_pure_curved(&self) -> bool {
4896        self.curved_atom_count == self.atoms.len() && !self.atoms.is_empty()
4897    }
4898}
4899
4900/// Adjudicate the curved-vs-linear split across a whole hybrid dictionary by the
4901/// common evidence criterion. `slots[i]` holds the curved/linear candidates for
4902/// atom slot `i` (each scored on the same rows, on the common Laplace scale).
4903///
4904/// The result reduces EXACTLY to pure-linear when every slot's curved candidate
4905/// has `Θ → 0` (the turning floor fires everywhere) and to pure-curved when
4906/// every slot's curved candidate wins the evidence comparison. (Common-data
4907/// criterion, #1202 — both candidates fit the atom's response residual, with
4908/// linear nested as the curved family's `Θ = 0` sub-model; see the module header
4909/// above and `crate::terms::sae::hybrid_split`.)
4910///
4911/// Returns an error only if some slot has no candidates to adjudicate (an empty
4912/// dictionary slot is a caller bug, not a silent skip).
4913pub 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// ---------------------------------------------------------------------------
4945// Tests
4946//
4947// These are type-level / structural tests: per the task contract we do
4948// not compile or run them in this session. They document the expected
4949// shapes and degenerate-case behavior so a future maintainer running
4950// `cargo test` sees the contract written down.
4951// ---------------------------------------------------------------------------
4952
4953#[cfg(test)]
4954mod tests {
4955    use super::*;
4956    use crate::arrow_schur::ArrowFactorSlab;
4957    use ndarray::array;
4958
4959    // Dense `H⁻¹` apply via explicit inverse (test-only reference solver).
4960    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        // Two decoupled 2x2 blocks: {0,1} and {2,3}.
5009        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        // Exactly two components.
5019        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        // 0-1 and 1-2 coupled (but no direct 0-2 edge) must form one component.
5042        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        // Regression for #1465: the ranking `delta` / `bayes_factor` must be
5055        // measured on the SAME scale that orders the table (the Occam-penalised
5056        // conditional AIC `ranking_score`), so every row's delta is >= 0 and its
5057        // Bayes factor >= 1 — the table must never claim a non-winner beats the
5058        // declared winner. The scenario is exactly the case the comparison
5059        // exists to handle: AIC and raw REML DISAGREE. `m1` is the AIC winner
5060        // but does NOT carry the minimum raw REML (`m2` does) — the noise
5061        // extra-term case from the issue.
5062        //
5063        // `ranking_score` = -2*log_lik + 2*edf; with log_lik = 0 it is `2*edf`,
5064        // so the AIC order is m1 < m2 < m3 while the raw-REML order has m2 lowest.
5065        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        // raw REML : m2 (41.605) < m1 (53.748) < m3 (120.011)
5075        // AIC=2*edf: m1 (100)    < m2 (102)    < m3 (130)
5076        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        // No ranking row may contradict the declared winner.
5085        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        // The raw-REML score table is referenced to the genuine minimum raw REML
5107        // (m2), so its best-over-model Bayes factors are also coherent (>= 1).
5108        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        // m2 carries the minimum raw REML, so its raw delta is exactly 0.
5123        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        // Support in component 0 -> cone is {0,1}.
5140        assert_eq!(cone_of_influence(&labels, &[0]), vec![0, 1]);
5141        // Support spanning both -> cone is everything.
5142        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        // Fully coupled SPD H: cone is the whole space, result must equal the
5148        // unconfined sensitivity-operator mode response bit-for-bit.
5149        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        // Two ρ-columns, each supported on a single coefficient.
5153        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        // Block-decoupled H: blocks {0,1} and {2,3}. A column supported only in
5179        // block {0,1} must produce sensitivity zero in block {2,3}, and match
5180        // the exact solution within its own block.
5181        let mut h = Array2::<f64>::zeros((4, 4));
5182        // Block A.
5183        h[[0, 0]] = 4.0;
5184        h[[1, 1]] = 3.0;
5185        h[[0, 1]] = 1.0;
5186        h[[1, 0]] = 1.0;
5187        // Block B.
5188        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        // Exact reference: -H⁻¹ q. Off-block entries are exactly zero already
5205        // (decoupled inverse), and the cone must preserve the in-block ones.
5206        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        // Block B is outside the cone -> exactly zero.
5217        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        // Inactive ρ: empty support, must be skipped without solving.
5226        let empty_support = 0..0usize;
5227        let supports = std::slice::from_ref(&empty_support);
5228        // A NaN inverse: an empty-support column must be skipped WITHOUT
5229        // solving, so the operator's finite-check never sees the NaN and the
5230        // result is `Some(zeros)`. Were the inactive column ever solved, the
5231        // NaN would propagate and `mode_response_coned` would return `None`.
5232        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        // d = 1, k = 1, n = 1, H_uu_1 = [[2.0]] => L = [[sqrt(2)]],
5243        // H_uβ_1 = [[0.5]], A = 2 - 0.5 * 0.5 / 2 = 1.875.
5244        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        // log|H| = log(2) + log(1.875). With dim(H)=2 and rank(S)=1,
5281        // V includes the rank-aware TK nullspace normalizer.
5282        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    /// #1132 bug 2: a β-profiled atom (no shared `β` block, `k == 0`) reaches
5299    /// `arrow_log_det_from_cache` in the dense Direct path with
5300    /// `schur_factor = None` — there is no reduced Schur complement to form. The
5301    /// joint Hessian is then block-diagonal in the latent rows, so its log-det
5302    /// is exactly the per-row sum with NO Schur term. Before the fix this
5303    /// returned `None` (the `schur_factor.as_ref()?` bail), starving the REML
5304    /// Laplace normaliser and erroring "arrow_log_det_from_cache returned None
5305    /// at ridge=0 Direct mode". Now it returns `Some(Σ_i log|H_tt^(i)|)`.
5306    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        // Single latent block H_tt = [[3.0]]; no Schur term for k == 0.
5341        assert!(
5342            (log_det - 3.0_f64.ln()).abs() < 1e-12,
5343            "log_det = {log_det}"
5344        );
5345        // The cache's own computation must agree bit-for-bit.
5346        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        // k > 0 but no dense Schur factor is the genuine InexactPCG case and
5355        // must still reject (the guard must not over-broaden to all `None`).
5356        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        // ∂u/∂β = -H_uu⁻¹ H_uβ = -0.5 / 2 = -0.25.
5414        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        // ∂β/∂ρ = -A⁻¹ · 1 = -1/1.875.
5424        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        // Each candidate is the better predictor on its own half of the data;
5513        // stacking keeps both, unlike winner-take-all.
5514        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        // Recompute the KKT residual from scratch at the returned weights:
5549        // g_k = mean_i p_ik / mix_i must satisfy g_k <= 1 (+tol) everywhere
5550        // and w_k * |g_k - 1| <= tol. The objective is concave, so this
5551        // certifies the GLOBAL optimum, not merely a stationary iterate.
5552        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        // Recorded issue mechanism: a signed mean-log-likelihood step of
5714        // -1.4e-13 was rejected even though it is unresolved at the scale of
5715        // the composite EM map. The admissible decrease below is derived only
5716        // from that map's observed objective scale and arithmetic reduction
5717        // bounds; the recorded step is an input to the decision, not a new
5718        // tolerance.
5719        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        // Recorded issue mechanism: a +6.6e-15 gain with a 1.5e-14 reduction
5746        // bound exhausted instead of certifying. A below-resolution gain is a
5747        // valid objective fixed point only when the independent parameter-map
5748        // residual also clears its configured tolerance.
5749        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        // Recorded #2324 scale: a raw mean-coordinate update of 4.6e-8 sat
5782        // just above sqrt(epsilon) and refused every adjudication even though
5783        // the component carries only one quarter of the predictive measure.
5784        // The stopping rule must measure the density parameter, not an inert
5785        // conditional coordinate whose meaning disappears with component mass.
5786        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        // Component mass itself is an identifiable density coordinate and is
5805        // therefore never discounted by the component's old mass.
5806        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        // Component zero is exactly one-dimensional: its x coordinate never
5952        // changes, so the constrained MLE has one covariance eigenvalue at the
5953        // configured floor. The old BHHH determinant had identically-zero
5954        // mean-x and covariance-xy score columns and therefore rejected this
5955        // perfectly valid constrained predictive density as non-SPD.
5956        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                // Vary the within-cluster radius while preserving its angular
5985                // symmetry. A literal constant-radius micro-circle makes the
5986                // Gaussian scale score identically zero and its empirical
5987                // Fisher singular, which is not a Gaussian-cluster fixture.
5988                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        // Both `2πs` and `Rr` overflow if formed directly, although their log
6025        // and the ratio `Rr/s` are representable.
6026        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        // Replace the first group by a narrow, genuinely anisotropic line so
6165        // the line role is identifiable without changing the point group.
6166        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        // Reversing group labels must reverse the selected roles, not the
6198        // fitted unlabeled mixture density or its common-scale score.
6199        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        // The points (.3, ±sqrt(.91)) lie on both unit circles centered at
6293        // (0, 0) and (.6, 0). Repeating one point gives three labelled
6294        // components. Thus center and directions move by O(1) while every
6295        // component mean—and therefore the represented mixture density—is
6296        // bit-identical.
6297        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    // -----------------------------------------------------------------------
6371    // #1026 hybrid curved + linear-tail split-selection
6372    // -----------------------------------------------------------------------
6373
6374    /// Build the two candidate parameterizations for one atom slot the way the
6375    /// fit would: the linear special case (one decoder direction, `Θ = 0`,
6376    /// `P_linear` params) and the curved candidate (`latent_dim` ≥ 1, more
6377    /// params, fitted turning `theta`). The curved candidate's likelihood is the
6378    /// linear likelihood MINUS `curved_loglik_gain` of NLE (curvature it captures
6379    /// the secant cannot), so the nesting invariant `curved_loglik ≥ linear` is
6380    /// honored: a straight feature has zero gain, a turning feature a positive
6381    /// gain that grows with Θ. The rank-aware Laplace normalizer charges the
6382    /// extra `½(P_curved − P_linear)·log(2π)` for the curved parameters, so the
6383    /// evidence comparison is the real `Θ/√ε` crossover.
6384    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        // A perfectly straight curved fit (Θ = 0) gains no likelihood over its
6404        // linear sub-model but pays more parameters → linear must win, by
6405        // construction, even if finite-sample evidence noise nudged the curved
6406        // NLE slightly below linear.
6407        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        // The exact-zero guard fires regardless of the evidence margin sign.
6412        assert!(choice.curved_turning.unwrap() <= HYBRID_LINEAR_TURNING_FLOOR);
6413    }
6414
6415    #[test]
6416    fn hybrid_selects_curved_when_turning_pays_for_itself() {
6417        // A genuinely turning feature (Θ = 2π, a full loop): the curved fit
6418        // captures enough curvature that, even charged the extra-parameter price,
6419        // its NLE drops below the linear secant's → curved wins.
6420        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        // The curved fit won a strictly positive evidence margin.
6424        assert!(choice.curved_evidence_margin > 0.0);
6425    }
6426
6427    #[test]
6428    fn hybrid_keeps_linear_when_curvature_doesnt_pay_its_price() {
6429        // A barely-curved feature (small Θ): the curved fit recovers only a sliver
6430        // of likelihood, not enough to cover the extra-parameter price → the
6431        // dominance floor keeps the linear tail.
6432        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        // Exact NLE tie (above the turning floor so the evidence path decides):
6441        // the cheaper linear atom wins, preserving strict generalization — the
6442        // hybrid never pays for curvature it does not need.
6443        let theta = 0.5; // above the floor → evidence path, not the exact guard
6444        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        // Every slot's curved candidate has Θ → 0 (flat features everywhere): the
6457        // dominance floor fires at every slot → the hybrid recovers the pure-
6458        // linear dictionary exactly. This is the `all Θ → 0` limit (3).
6459        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        // Summed NLE equals the pure-linear baseline (every slot chose linear).
6467        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        // Every slot's feature turns enough (Θ = 2π, large likelihood gain) that
6474        // curved beats linear everywhere → the pure-curved limit (3).
6475        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        // Mixed synthetic: slots 0..3 are CIRCLE features (high turning Θ = 2π,
6487        // the curved fit captures the loop), slots 3..7 are LINEAR DIRECTIONS
6488        // (straight, Θ = 0). The evidence split must select curved for the
6489        // circles and linear for the directions — and the hybrid's summed
6490        // evidence must be ≤ the summed per-slot LINEAR-candidate NLE (each
6491        // slot's best straight line fit to its response residual). This is a
6492        // data-level match-or-beat dominance (#1202: linear is the curved
6493        // family's nested Θ = 0 sub-model on common data), and holds because each
6494        // slot picks the argmin of its two common-data candidates.
6495        let mut slots: Vec<Vec<HybridAtomCandidate>> = Vec::new();
6496        let mut pure_linear_baseline = 0.0_f64;
6497        // Three circle features: a curved atom replaces ~10-30 linear secants, so
6498        // the curved fit buys a large likelihood gain that dwarfs its param price.
6499        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        // Four straight linear directions: zero turning, the linear special case
6512        // is optimal — a curved atom buys nothing and only costs parameters.
6513        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        // The first three (circles) chose curved; the last four (directions) chose
6522        // linear.
6523        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        // The hybrid's summed negative-log-evidence is ≤ the summed per-slot
6541        // LINEAR-candidate NLE (each slot's best straight line fit to its response
6542        // residual): the per-slot argmin can only lower the sum. This is a
6543        // data-level match-or-beat dominance (#1202): linear is the curved
6544        // family's nested Θ = 0 sub-model on common data.
6545        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        // And strictly better, because the curved circle slots paid off.
6552        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    // ── #1362: compare_models must Occam-penalise a pure-noise smooth ────────
6562    //
6563    // These tests pin the ranking contract directly on `compare_reml_fits` with
6564    // controlled (score, edf, log_lik) inputs taken from the actual #1362
6565    // reproduction (Rust `reml_score` of `y ~ s(x)` vs `y ~ s(x) + s(z)` at
6566    // n=700). They do not need a fitted GAM or a Python wheel.
6567
6568    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        // AIC = -2ℓ + 2·edf.
6583        let c = cand("m", /*score (ignored)*/ 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        // Seed-3000 numbers from the #1362 Rust reproduction:
6605        //   small (y ~ s(x)):      reml=180.526, edf=6.748,  loglik=-32.0866
6606        //   big   (y ~ s(x)+s(z)): reml=177.404, edf=14.250, loglik=-32.1212
6607        // The big (noise-augmented) model has the LOWER (apparently better) raw
6608        // REML evidence, yet it spends ~7.5 extra EDF fitting noise without
6609        // improving the likelihood. The winner must be the SMALL model.
6610        let small = cand("small", 180.526, 6.748, -32.0866);
6611        let big = cand("big", 177.404, 14.250, -32.1212);
6612
6613        // Sanity: raw evidence (the broken headline) prefers big.
6614        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        // The score table still reports the raw evidence headline unchanged, so
6622        // Model.evidence / bayes_factor_vs stay consistent with the table.
6623        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        // Issue #2124: `ranking_score` is the conditional AIC (`−2ℓ + 2·edf`), a
6640        // −2·log / deviance-scale cost. For an AIC gap Δ the Akaike evidence ratio
6641        // (Burnham & Anderson) is `exp(−½Δ)`, so the winner-over-loser
6642        // `bayes_factor` must be `exp(½Δ)` — NOT `exp(Δ)`, which squares it.
6643        //
6644        // Winner: AIC 0 (loglik 0, edf 0). Loser: AIC = 27.68 (loglik −13.84,
6645        // edf 0), matching the ΔAIC in the issue repro. Raw REML scores are set
6646        // distinct (100 vs 110) to lock the scoping: the raw score_table path
6647        // must stay `exp(Δreml)` with NO halving.
6648        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        // The AIC gap FIELD stays on the AIC scale, unchanged (issue #2124).
6662        assert!((loser_row.delta - delta_aic).abs() < 1e-9);
6663
6664        // The Bayes factor is the Akaike evidence ratio exp(½·ΔAIC) = exp(13.84)
6665        // ≈ 1.03e6 — NOT the squared exp(27.68) ≈ 1.05e12 the bug reported.
6666        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        // Explicit anti-regression: it must not be the squared ratio.
6675        assert!(loser_row.bayes_factor < delta_aic.exp() * 0.5);
6676
6677        // Scoping lock (issue #2124): the RAW-REML score_table path is untouched —
6678        // its best-over-model Bayes factor is `exp(Δreml)` with NO halving. Raw
6679        // scores 100 (winner) vs 110 (loser) give Δreml = 10, so the loser's raw
6680        // Bayes factor is exp(10), not exp(5).
6681        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        // Seed-3000 relevant-z numbers from the same reproduction:
6697        //   small: reml=1025.067, edf≈6.75,  loglik≈-368.99 (aic≈751.5)
6698        //   big:   reml=199.509,  edf≈14.25, loglik≈-33.16  (aic≈94.8)
6699        // A genuinely relevant smooth lowers BOTH the evidence and the AIC, so
6700        // the bigger model must still win — a fix cannot just always pick small.
6701        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        // Two same-family fits on different-sized data are not comparable by
6713        // AIC / evidence; the comparison must fail loud, mirroring the family
6714        // guard, rather than declare a sample-size-driven winner.
6715        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        // Same n is comparable.
6732        compare_reml_fits(vec![with_n("a", 250), with_n("b", 250)])
6733            .expect("same-n comparison must succeed");
6734
6735        // A missing count (`None`) is unconstrained: it must not block a
6736        // comparison against a fit that does carry one (legacy / scan payloads).
6737        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}