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    /// Max-norm tolerance for the EM map in empirical predictive-density
815    /// coordinates: the largest absolute change in any training row's log
816    /// density. This quotients component permutations, duplicate-component
817    /// mass exchange, and non-identifiable ring factorizations while retaining
818    /// sensitivity to likelihood changes that cancel in the mean objective.
819    pub parameter_tol: f64,
820    /// Lower eigenvalue constraint for every component covariance. The M-step
821    /// solves this constrained likelihood problem exactly by spectral clipping;
822    /// it is not an additive ridge or an unmodelled prior.
823    pub covariance_floor: f64,
824    /// Maximum iterations for the deterministic k-means seeding pass.
825    pub kmeans_max_iter: usize,
826}
827
828impl Default for GaussianMixtureConfig {
829    fn default() -> Self {
830        Self {
831            max_iter: 1000,
832            loglik_tol: f64::EPSILON.sqrt(),
833            parameter_tol: f64::EPSILON.sqrt(),
834            covariance_floor: 1e-6,
835            kmeans_max_iter: 25,
836        }
837    }
838}
839
840/// Residual evidence proving that the returned mixture is a fixed point of its
841/// likelihood EM map rather than merely the iterate present at a work cap.
842#[derive(Debug, Clone, Copy, Serialize, Deserialize)]
843pub struct GaussianMixtureCertificate {
844    /// Mean log likelihood at the exact parameter state being certified.
845    pub mean_log_likelihood: f64,
846    /// Signed gain produced by one further EM map application from that state.
847    pub mean_log_likelihood_gain: f64,
848    /// Absolute numerical uncertainty on the comparison of the two likelihoods.
849    /// This includes both the final likelihood-reduction error and the
850    /// scale-derived resolution floor of the composite EM map.
851    pub monotonicity_uncertainty: f64,
852    pub objective_residual: f64,
853    pub objective_tolerance: f64,
854    pub parameter_residual: f64,
855    pub parameter_tolerance: f64,
856    /// Measured per-iteration contraction rate `ρ` of the parameter residual
857    /// over the trailing [`EM_RATE_WINDOW`], or `None` before a full window has
858    /// accumulated. `ρ < 1` is the evidence that the iterate is still
859    /// descending; `ρ ≥ 1` is the evidence that it has stalled.
860    pub contraction_rate: Option<f64>,
861    /// Iterations still required to reach `parameter_tolerance` at the measured
862    /// `contraction_rate`, or `None` when no rate is available, the rate does
863    /// not contract, or the tolerance is already met. This is what makes a
864    /// refusal PRICEABLE: a caller can see whether it was interrupted mid-
865    /// descent and by how much.
866    pub projected_iterations_to_tolerance: Option<usize>,
867}
868
869/// Trailing window, in EM updates, over which the parameter residual's
870/// contraction rate is measured.
871///
872/// DERIVATION. A single step ratio `r_t / r_{t-1}` carries the full relative
873/// noise of both residuals, and near a fixed point that noise is comparable to
874/// the step itself — one ratio cannot distinguish descent from a stall. The
875/// geometric mean over `W` steps averages `W` independent log-ratios, so its
876/// log-jitter falls as `1/√W`: `W = 64` suppresses per-step jitter eightfold
877/// while costing 6.4% of the base update budget to establish. It is also the
878/// re-validation cadence during an extension, so a stall is caught within one
879/// window of appearing rather than at the end of the projection.
880const EM_RATE_WINDOW: usize = 64;
881
882/// Geometric per-iteration contraction rate of the parameter residual across
883/// the window: `ρ = (r_last / r_first)^{1/(W)}`.
884///
885/// `None` until the window is full, or when either endpoint is not strictly
886/// positive and finite — a zero residual is convergence, not a rate, and the
887/// caller's tolerance test has already handled it.
888fn em_contraction_rate(window: &std::collections::VecDeque<f64>) -> Option<f64> {
889    if window.len() < EM_RATE_WINDOW + 1 {
890        return None;
891    }
892    let first = *window.front()?;
893    let last = *window.back()?;
894    if !(first.is_finite() && last.is_finite() && first > 0.0 && last > 0.0) {
895        return None;
896    }
897    let steps = (window.len() - 1) as f64;
898    let rate = (last / first).powf(1.0 / steps);
899    rate.is_finite().then_some(rate)
900}
901
902/// Iterations still required to bring `residual` to `tolerance` at contraction
903/// rate `rate`, i.e. the `N*` solving `residual·ρ^{N*} = tolerance`.
904///
905/// `None` when the rate does not contract (`ρ ∉ (0, 1)`), when the tolerance is
906/// already met, or when the inputs are not finite — in every one of those cases
907/// there is no projection to make, and inventing one would be the fabrication
908/// this certificate exists to prevent.
909fn em_projected_iterations(residual: f64, tolerance: f64, rate: f64) -> Option<usize> {
910    if !(residual.is_finite() && tolerance.is_finite() && rate.is_finite()) {
911        return None;
912    }
913    if !(rate > 0.0 && rate < 1.0) || !(residual > tolerance) || tolerance <= 0.0 {
914        return None;
915    }
916    let steps = (tolerance / residual).ln() / rate.ln();
917    (steps.is_finite() && steps >= 0.0).then(|| steps.ceil() as usize)
918}
919
920/// Exact parameter state carried across an EM exhaustion boundary.
921#[derive(Debug, Clone, Serialize, Deserialize)]
922pub struct GaussianMixtureCheckpoint {
923    pub weights: Array1<f64>,
924    pub means: Array2<f64>,
925    pub covariances: Vec<Array2<f64>>,
926    pub mean_log_likelihood: f64,
927    pub completed_iterations: usize,
928    data_fingerprint: Fingerprint,
929    covariance_floor: f64,
930}
931
932/// Typed Gaussian-mixture optimization failure. Exhaustion and a broken EM
933/// monotonicity invariant both carry the last internally consistent state.
934#[derive(Debug, Clone)]
935pub enum GaussianMixtureError {
936    InvalidInput {
937        message: String,
938    },
939    NumericalFailure {
940        message: String,
941        checkpoint: Option<GaussianMixtureCheckpoint>,
942    },
943    MonotonicityViolation {
944        previous_mean_log_likelihood: f64,
945        next_mean_log_likelihood: f64,
946        numerical_uncertainty: f64,
947        checkpoint: GaussianMixtureCheckpoint,
948    },
949    DidNotConverge {
950        max_iterations: usize,
951        certificate: GaussianMixtureCertificate,
952        checkpoint: GaussianMixtureCheckpoint,
953    },
954}
955
956impl std::fmt::Display for GaussianMixtureError {
957    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
958        match self {
959            Self::InvalidInput { message } => write!(f, "invalid Gaussian mixture: {message}"),
960            Self::NumericalFailure {
961                message,
962                checkpoint,
963            } => write!(
964                f,
965                "Gaussian-mixture numerical failure: {message} (checkpoint iterations {})",
966                checkpoint
967                    .as_ref()
968                    .map_or(0, |value| value.completed_iterations)
969            ),
970            Self::MonotonicityViolation {
971                previous_mean_log_likelihood,
972                next_mean_log_likelihood,
973                numerical_uncertainty,
974                checkpoint,
975            } => write!(
976                f,
977                "Gaussian-mixture EM violated monotone ascent at iteration {}: mean log likelihood {previous_mean_log_likelihood:.12e} -> {next_mean_log_likelihood:.12e} (comparison uncertainty {numerical_uncertainty:.3e}); resume from the carried checkpoint only after diagnosing the numerical failure",
978                checkpoint.completed_iterations
979            ),
980            Self::DidNotConverge {
981                max_iterations,
982                certificate,
983                checkpoint,
984            } => write!(
985                f,
986                "Gaussian-mixture EM did not certify after {max_iterations} additional iterations (total {}): signed mean-log-likelihood gain {:.6e} (numerical uncertainty {:.3e}), objective residual {:.6e}/{:.3e}, parameter-map residual {:.6e}/{:.3e}, contraction rate {} per iteration, projected iterations to tolerance {}; resume from the carried checkpoint, which is not comparable evidence",
987                checkpoint.completed_iterations,
988                certificate.mean_log_likelihood_gain,
989                certificate.monotonicity_uncertainty,
990                certificate.objective_residual,
991                certificate.objective_tolerance,
992                certificate.parameter_residual,
993                certificate.parameter_tolerance,
994                match certificate.contraction_rate {
995                    Some(rate) => format!("{rate:.6}"),
996                    None => "unmeasured".to_string(),
997                },
998                match certificate.projected_iterations_to_tolerance {
999                    Some(steps) => steps.to_string(),
1000                    None => "none (not contracting)".to_string(),
1001                }
1002            ),
1003        }
1004    }
1005}
1006
1007impl std::error::Error for GaussianMixtureError {}
1008
1009/// A fitted `k`-component full-covariance Gaussian mixture.
1010#[derive(Debug, Clone)]
1011pub struct GaussianMixtureFit {
1012    /// Mixing weights, length `k`, on the simplex.
1013    weights: Array1<f64>,
1014    /// Component means, `k × d`.
1015    means: Array2<f64>,
1016    /// Component covariances, `k` matrices of shape `d × d` (SPD).
1017    covariances: Vec<Array2<f64>>,
1018    /// Number of mixture components.
1019    k: usize,
1020    /// Data dimension.
1021    d: usize,
1022    /// Number of rows used to fit.
1023    n_obs: usize,
1024    /// Maximised total log-likelihood `Σ_i log Σ_j w_j N(y_i; μ_j, Σ_j)`.
1025    loglik: f64,
1026    /// EM iterations taken.
1027    iterations: usize,
1028    certificate: GaussianMixtureCertificate,
1029}
1030
1031impl GaussianMixtureFit {
1032    pub fn weights(&self) -> ArrayView1<'_, f64> {
1033        self.weights.view()
1034    }
1035
1036    pub fn means(&self) -> ArrayView2<'_, f64> {
1037        self.means.view()
1038    }
1039
1040    pub fn covariances(&self) -> &[Array2<f64>] {
1041        &self.covariances
1042    }
1043
1044    pub fn iterations(&self) -> usize {
1045        self.iterations
1046    }
1047
1048    pub fn certificate(&self) -> GaussianMixtureCertificate {
1049        self.certificate
1050    }
1051
1052    /// Free-parameter count `P` of a `k`-component full-covariance mixture in
1053    /// `d` dimensions: `(k − 1)` mixing weights on the simplex, `k·d` mean
1054    /// coordinates, and `k · d(d+1)/2` covariance entries. This is the exact
1055    /// quantity that enters the rank-aware normalizer as `dim(H) − rank(S)`.
1056    pub fn num_free_parameters(&self) -> usize {
1057        let cov_per = self.d * (self.d + 1) / 2;
1058        (self.k - 1) + self.k * self.d + self.k * cov_per
1059    }
1060
1061    /// Per-observation log predictive density `log p(y_i)` under the fitted
1062    /// mixture, length `n`. This is the held-out-density column source for
1063    /// cross-class stacking when the mixture is evaluated on a held-out fold.
1064    pub fn per_point_log_density(&self, data: ArrayView2<'_, f64>) -> Result<Array1<f64>, String> {
1065        if data.ncols() != self.d {
1066            return Err(format!(
1067                "mixture log-density expects {} columns, got {}",
1068                self.d,
1069                data.ncols()
1070            ));
1071        }
1072        let n = data.nrows();
1073        let mut comp = Vec::with_capacity(self.k);
1074        for j in 0..self.k {
1075            comp.push(GaussianComponentEval::factor(
1076                self.means.row(j),
1077                &self.covariances[j],
1078            )?);
1079        }
1080        let mut out = Array1::<f64>::zeros(n);
1081        let log_w: Vec<f64> = self.weights.iter().map(|w| w.ln()).collect();
1082        for i in 0..n {
1083            let row = data.row(i);
1084            let mut log_terms = vec![f64::NEG_INFINITY; self.k];
1085            let mut max_term = f64::NEG_INFINITY;
1086            for j in 0..self.k {
1087                let lt = log_w[j] + comp[j].log_density(row);
1088                log_terms[j] = lt;
1089                if lt > max_term {
1090                    max_term = lt;
1091                }
1092            }
1093            out[i] = log_sum_exp(&log_terms, max_term);
1094        }
1095        Ok(out)
1096    }
1097
1098    /// Schwarz BIC approximation to negative log evidence, divided by two so
1099    /// it shares the ordinary negative-log-likelihood scale used by the shape
1100    /// race. Lower is better.
1101    pub fn bic(&self) -> f64 {
1102        -self.loglik + 0.5 * self.num_free_parameters() as f64 * (self.n_obs as f64).ln()
1103    }
1104}
1105
1106/// Cached per-component Gaussian evaluator: mean, precision `Σ⁻¹`, and the
1107/// log-normalizing constant `−½(d log 2π + log|Σ|)`.
1108#[derive(Debug, Clone)]
1109struct GaussianComponentEval {
1110    residual_origin: Array1<f64>,
1111    residual_scale: Array1<f64>,
1112    residual_normalized_offset: Array1<f64>,
1113    precision: Array2<f64>,
1114    log_norm: f64,
1115    d: usize,
1116}
1117
1118impl GaussianComponentEval {
1119    fn factor(mean: ArrayView1<'_, f64>, cov: &Array2<f64>) -> Result<Self, String> {
1120        let d = mean.len();
1121        if mean.iter().any(|value| !value.is_finite()) {
1122            return Err("mixture component mean must be finite".to_string());
1123        }
1124        if cov.nrows() != d || cov.ncols() != d {
1125            return Err(format!(
1126                "mixture component covariance must be {d}x{d}, got {}x{}",
1127                cov.nrows(),
1128                cov.ncols()
1129            ));
1130        }
1131        let (evals, evecs) = cov
1132            .eigh(Side::Lower)
1133            .map_err(|e| format!("mixture component covariance eigendecomposition failed: {e}"))?;
1134        let mut log_det = 0.0_f64;
1135        let mut inv_evals = Array1::<f64>::zeros(d);
1136        for (idx, &ev) in evals.iter().enumerate() {
1137            if !ev.is_finite() || ev <= 0.0 {
1138                return Err(format!(
1139                    "mixture component covariance is not SPD: eigenvalue {idx} is {ev:.3e}"
1140                ));
1141            }
1142            log_det += ev.ln();
1143            let inverse = ev.recip();
1144            if !inverse.is_finite() {
1145                return Err(format!(
1146                    "mixture component precision is not representable: eigenvalue {idx} is {ev:.3e}"
1147                ));
1148            }
1149            inv_evals[idx] = inverse;
1150        }
1151        // Σ⁻¹ = V diag(1/λ) Vᵀ.
1152        let mut precision = Array2::<f64>::zeros((d, d));
1153        for a in 0..d {
1154            for b in 0..d {
1155                let mut acc = 0.0_f64;
1156                for m in 0..d {
1157                    acc += evecs[[a, m]] * inv_evals[m] * evecs[[b, m]];
1158                }
1159                precision[[a, b]] = acc;
1160            }
1161        }
1162        let log_norm = -0.5 * (d as f64 * (2.0 * std::f64::consts::PI).ln() + log_det);
1163        if precision.iter().any(|value| !value.is_finite()) || !log_norm.is_finite() {
1164            return Err(
1165                "mixture component factorization produced non-finite precision or log normalizer"
1166                    .to_string(),
1167            );
1168        }
1169        Ok(Self {
1170            residual_origin: mean.to_owned(),
1171            residual_scale: Array1::zeros(d),
1172            residual_normalized_offset: Array1::zeros(d),
1173            precision,
1174            log_norm,
1175            d,
1176        })
1177    }
1178
1179    fn isotropic(charts: &[StableScalarMeanChart], variance: f64) -> Result<Self, String> {
1180        let d = charts.len();
1181        if d == 0 {
1182            return Err("isotropic Gaussian density requires positive dimension".to_string());
1183        }
1184        if !(variance.is_finite() && variance > 0.0) {
1185            return Err(format!(
1186                "isotropic Gaussian variance must be finite and positive, got {variance}"
1187            ));
1188        }
1189        let inverse_variance = variance.recip();
1190        if !inverse_variance.is_finite() {
1191            return Err(format!(
1192                "isotropic Gaussian precision is non-finite for variance {variance}"
1193            ));
1194        }
1195        let mut precision = Array2::<f64>::zeros((d, d));
1196        for axis in 0..d {
1197            precision[[axis, axis]] = inverse_variance;
1198        }
1199        let log_norm = -0.5 * d as f64 * ((2.0 * std::f64::consts::PI).ln() + variance.ln());
1200        if !log_norm.is_finite() {
1201            return Err("isotropic Gaussian log normalizer is non-finite".to_string());
1202        }
1203        Ok(Self {
1204            residual_origin: Array1::from_iter(charts.iter().map(|chart| chart.origin)),
1205            residual_scale: Array1::from_iter(charts.iter().map(|chart| chart.scale)),
1206            residual_normalized_offset: Array1::from_iter(
1207                charts.iter().map(|chart| chart.normalized_offset),
1208            ),
1209            precision,
1210            log_norm,
1211            d,
1212        })
1213    }
1214
1215    #[inline]
1216    fn log_density(&self, y: ArrayView1<'_, f64>) -> f64 {
1217        let residual = self.residual(y);
1218        let pv = self.precision_times_residual(&residual);
1219        let mut quad = 0.0_f64;
1220        for c in 0..self.d {
1221            quad += residual[c] * pv[c];
1222        }
1223        self.log_norm - 0.5 * quad
1224    }
1225
1226    #[inline]
1227    fn residual(&self, y: ArrayView1<'_, f64>) -> Vec<f64> {
1228        let mut residual = vec![0.0_f64; self.d];
1229        for axis in 0..self.d {
1230            residual[axis] = (-self.residual_normalized_offset[axis]).mul_add(
1231                self.residual_scale[axis],
1232                y[axis] - self.residual_origin[axis],
1233            );
1234        }
1235        residual
1236    }
1237
1238    /// `Σ⁻¹ (y − μ)`.
1239    #[inline]
1240    fn precision_times_residual(&self, residual: &[f64]) -> Vec<f64> {
1241        let mut out = vec![0.0_f64; self.d];
1242        for a in 0..self.d {
1243            let mut acc = 0.0_f64;
1244            for b in 0..self.d {
1245                acc += self.precision[[a, b]] * residual[b];
1246            }
1247            out[a] = acc;
1248        }
1249        out
1250    }
1251}
1252
1253#[inline]
1254fn log_sum_exp(terms: &[f64], max_term: f64) -> f64 {
1255    if !max_term.is_finite() {
1256        return f64::NEG_INFINITY;
1257    }
1258    let mut acc = 0.0_f64;
1259    for &t in terms {
1260        acc += (t - max_term).exp();
1261    }
1262    max_term + acc.ln()
1263}
1264
1265fn evidence_matrix_fingerprint(namespace: &str, values: ArrayView2<'_, f64>) -> Fingerprint {
1266    let mut hasher = Fingerprinter::new();
1267    hasher.write_str(namespace);
1268    hasher.write_usize(values.nrows());
1269    hasher.write_usize(values.ncols());
1270    // Hash logical row-major iteration order rather than the backing storage,
1271    // so an equivalent strided view resumes the same mathematical problem.
1272    for &value in values {
1273        hasher.write_f64(value);
1274    }
1275    hasher.finalize()
1276}
1277
1278fn mixture_data_fingerprint(data: ArrayView2<'_, f64>) -> Fingerprint {
1279    evidence_matrix_fingerprint("gaussian-mixture-em-v1", data)
1280}
1281
1282/// Fit a `k`-component full-covariance Gaussian mixture by deterministic
1283/// k-means++-style seeding (reusing the `terms::basis` farthest-point k-means,
1284/// a pure function of the data — no clock randomness) followed by EM to the
1285/// configured tolerance.
1286///
1287/// The fit is deterministic given `(data, k, config)`: the seed is the
1288/// farthest-point/k-means center selection, EM is a deterministic map, so
1289/// re-running yields the identical mixture.
1290pub fn fit_gaussian_mixture(
1291    data: ArrayView2<'_, f64>,
1292    k: usize,
1293    config: GaussianMixtureConfig,
1294) -> Result<GaussianMixtureFit, GaussianMixtureError> {
1295    validate_gaussian_mixture_problem(data, k, config)?;
1296    // Deterministic k-means++-style seeding via the shared basis k-means
1297    // (farthest-point init + Lloyd iterations).
1298    let means = gam_terms::basis::select_centers_by_strategy(
1299        data,
1300        &gam_terms::basis::CenterStrategy::KMeans {
1301            num_centers: k,
1302            max_iter: config.kmeans_max_iter,
1303        },
1304    )
1305    .map_err(|error| GaussianMixtureError::NumericalFailure {
1306        message: format!("deterministic k-means seeding failed: {error}"),
1307        checkpoint: None,
1308    })?;
1309    if means.nrows() != k || means.ncols() != data.ncols() {
1310        return Err(GaussianMixtureError::NumericalFailure {
1311            message: format!(
1312                "seeding returned {}x{} centers, expected {k}x{}",
1313                means.nrows(),
1314                means.ncols(),
1315                data.ncols()
1316            ),
1317            checkpoint: None,
1318        });
1319    }
1320    let global_covariance =
1321        constrained_data_covariance(data, config.covariance_floor).map_err(|message| {
1322            GaussianMixtureError::NumericalFailure {
1323                message,
1324                checkpoint: None,
1325            }
1326        })?;
1327    let weights = Array1::<f64>::from_elem(k, 1.0 / k as f64);
1328    let covariances = vec![global_covariance; k];
1329    let initial_e_step =
1330        mixture_e_step(data, &weights, &means, &covariances).map_err(|message| {
1331            GaussianMixtureError::NumericalFailure {
1332                message,
1333                checkpoint: None,
1334            }
1335        })?;
1336    let data_fingerprint = mixture_data_fingerprint(data);
1337    let checkpoint = GaussianMixtureCheckpoint {
1338        weights,
1339        means,
1340        covariances,
1341        mean_log_likelihood: initial_e_step.mean_log_likelihood,
1342        completed_iterations: 0,
1343        data_fingerprint,
1344        covariance_floor: config.covariance_floor,
1345    };
1346    run_gaussian_mixture_em(data, config, checkpoint)
1347}
1348
1349/// Resume EM from the exact state carried by [`GaussianMixtureError`].
1350pub fn resume_gaussian_mixture(
1351    data: ArrayView2<'_, f64>,
1352    config: GaussianMixtureConfig,
1353    checkpoint: GaussianMixtureCheckpoint,
1354) -> Result<GaussianMixtureFit, GaussianMixtureError> {
1355    let k = checkpoint.weights.len();
1356    validate_gaussian_mixture_problem(data, k, config)?;
1357    validate_gaussian_mixture_checkpoint(data, config.covariance_floor, &checkpoint)?;
1358    run_gaussian_mixture_em(data, config, checkpoint)
1359}
1360
1361fn validate_gaussian_mixture_problem(
1362    data: ArrayView2<'_, f64>,
1363    k: usize,
1364    config: GaussianMixtureConfig,
1365) -> Result<(), GaussianMixtureError> {
1366    let n = data.nrows();
1367    let d = data.ncols();
1368    if k == 0 {
1369        return Err(GaussianMixtureError::InvalidInput {
1370            message: "k must be positive".to_string(),
1371        });
1372    }
1373    if d == 0 {
1374        return Err(GaussianMixtureError::InvalidInput {
1375            message: "at least one data column is required".to_string(),
1376        });
1377    }
1378    if k > n {
1379        return Err(GaussianMixtureError::InvalidInput {
1380            message: format!("requested {k} components but data has {n} rows"),
1381        });
1382    }
1383    if data.iter().any(|value| !value.is_finite()) {
1384        return Err(GaussianMixtureError::InvalidInput {
1385            message: "data must be finite".to_string(),
1386        });
1387    }
1388    if config.max_iter == 0 || config.kmeans_max_iter == 0 {
1389        return Err(GaussianMixtureError::InvalidInput {
1390            message: "max_iter and kmeans_max_iter must be positive".to_string(),
1391        });
1392    }
1393    let numerical_floor = f64::EPSILON.sqrt();
1394    if !config.loglik_tol.is_finite()
1395        || config.loglik_tol < numerical_floor
1396        || !config.parameter_tol.is_finite()
1397        || config.parameter_tol < numerical_floor
1398        || !config.covariance_floor.is_finite()
1399        || config.covariance_floor <= 0.0
1400    {
1401        return Err(GaussianMixtureError::InvalidInput {
1402            message: format!(
1403                "loglik_tol and parameter_tol must be finite and >= {numerical_floor:.3e}, and covariance_floor must be finite and positive"
1404            ),
1405        });
1406    }
1407    Ok(())
1408}
1409
1410fn validate_gaussian_mixture_checkpoint(
1411    data: ArrayView2<'_, f64>,
1412    covariance_floor: f64,
1413    checkpoint: &GaussianMixtureCheckpoint,
1414) -> Result<(), GaussianMixtureError> {
1415    let d = data.ncols();
1416    let k = checkpoint.weights.len();
1417    let mass = checkpoint.weights.sum();
1418    if k == 0
1419        || checkpoint.data_fingerprint != mixture_data_fingerprint(data)
1420        || checkpoint.covariance_floor.to_bits() != covariance_floor.to_bits()
1421        || checkpoint.means.dim() != (k, d)
1422        || checkpoint.covariances.len() != k
1423        || checkpoint
1424            .covariances
1425            .iter()
1426            .any(|covariance| covariance.dim() != (d, d))
1427        || checkpoint
1428            .weights
1429            .iter()
1430            .chain(checkpoint.means.iter())
1431            .chain(checkpoint.covariances.iter().flat_map(|value| value.iter()))
1432            .any(|value| !value.is_finite())
1433        || checkpoint.weights.iter().any(|value| *value <= 0.0)
1434        || !mass.is_finite()
1435        || (mass - 1.0).abs() > f64::EPSILON.sqrt()
1436        || !checkpoint.mean_log_likelihood.is_finite()
1437    {
1438        return Err(GaussianMixtureError::InvalidInput {
1439            message: "checkpoint problem identity, dimensions, interior parameters, likelihood, or simplex mass are invalid".to_string(),
1440        });
1441    }
1442    Ok(())
1443}
1444
1445fn run_gaussian_mixture_em(
1446    data: ArrayView2<'_, f64>,
1447    config: GaussianMixtureConfig,
1448    mut checkpoint: GaussianMixtureCheckpoint,
1449) -> Result<GaussianMixtureFit, GaussianMixtureError> {
1450    validate_gaussian_mixture_checkpoint(data, config.covariance_floor, &checkpoint)?;
1451    let k = checkpoint.weights.len();
1452    let d = data.ncols();
1453    let data_fingerprint = mixture_data_fingerprint(data);
1454
1455    // Certify the CURRENT checkpoint before accepting another EM update. The
1456    // inclusive bound permits exactly `max_iter` accepted updates and then one
1457    // final map evaluation at the resulting checkpoint. Consequently every
1458    // success and every exhaustion pairs its certificate with the exact same
1459    // parameter state; a certificate for theta_t can never be attached to
1460    // theta_{t+1} merely because the work boundary was reached.
1461    // The base cap means "give up when not progressing", not "interrupt provable
1462    // progress". `budget` therefore starts at `max_iter` and may be extended
1463    // ONCE, by the iterate's OWN projection, and only while the residual is
1464    // measurably contracting. See the gate at the bottom of the loop.
1465    let mut budget = config.max_iter;
1466    let mut extension: Option<usize> = None;
1467    let mut residual_window: std::collections::VecDeque<f64> =
1468        std::collections::VecDeque::with_capacity(EM_RATE_WINDOW + 1);
1469    let mut additional_updates = 0usize;
1470    loop {
1471        let current = mixture_e_step(
1472            data,
1473            &checkpoint.weights,
1474            &checkpoint.means,
1475            &checkpoint.covariances,
1476        )
1477        .map_err(|message| GaussianMixtureError::NumericalFailure {
1478            message,
1479            checkpoint: Some(checkpoint.clone()),
1480        })?;
1481        if (checkpoint.mean_log_likelihood - current.mean_log_likelihood).abs()
1482            > current.mean_log_likelihood_roundoff
1483        {
1484            return Err(GaussianMixtureError::InvalidInput {
1485                message: format!(
1486                    "checkpoint mean log likelihood {:.12e} disagrees with its parameters ({:.12e} +/- {:.3e})",
1487                    checkpoint.mean_log_likelihood,
1488                    current.mean_log_likelihood,
1489                    current.mean_log_likelihood_roundoff
1490                ),
1491            });
1492        }
1493        checkpoint.mean_log_likelihood = current.mean_log_likelihood;
1494
1495        let (next_weights, next_means, next_covariances) = mixture_m_step(
1496            data,
1497            current.responsibilities.view(),
1498            config.covariance_floor,
1499        )
1500        .map_err(|message| GaussianMixtureError::NumericalFailure {
1501            message,
1502            checkpoint: Some(checkpoint.clone()),
1503        })?;
1504        let next = mixture_e_step(data, &next_weights, &next_means, &next_covariances).map_err(
1505            |message| GaussianMixtureError::NumericalFailure {
1506                message,
1507                checkpoint: Some(checkpoint.clone()),
1508            },
1509        )?;
1510        let objective_scale = current
1511            .mean_log_likelihood
1512            .abs()
1513            .max(next.mean_log_likelihood.abs())
1514            .max(1.0);
1515        let objective_step = next.mean_log_likelihood - current.mean_log_likelihood;
1516        let objective_residual = objective_step.abs() / objective_scale;
1517        let parameter_residual = empirical_predictive_density_residual(
1518            &current.row_log_likelihoods,
1519            &next.row_log_likelihoods,
1520        )
1521        .map_err(|message| GaussianMixtureError::NumericalFailure {
1522            message,
1523            checkpoint: Some(checkpoint.clone()),
1524        })?;
1525        let monotonicity_uncertainty = gaussian_mixture_monotonicity_uncertainty(
1526            objective_scale,
1527            current.mean_log_likelihood_roundoff,
1528            next.mean_log_likelihood_roundoff,
1529        );
1530        residual_window.push_back(parameter_residual);
1531        if residual_window.len() > EM_RATE_WINDOW + 1 {
1532            residual_window.pop_front();
1533        }
1534        let contraction_rate = em_contraction_rate(&residual_window);
1535        let projected_iterations_to_tolerance = contraction_rate
1536            .and_then(|rate| em_projected_iterations(parameter_residual, config.parameter_tol, rate));
1537        let certificate = GaussianMixtureCertificate {
1538            mean_log_likelihood: current.mean_log_likelihood,
1539            mean_log_likelihood_gain: objective_step,
1540            monotonicity_uncertainty,
1541            objective_residual,
1542            objective_tolerance: config.loglik_tol,
1543            parameter_residual,
1544            parameter_tolerance: config.parameter_tol,
1545            contraction_rate,
1546            projected_iterations_to_tolerance,
1547        };
1548        if objective_step < -monotonicity_uncertainty {
1549            return Err(GaussianMixtureError::MonotonicityViolation {
1550                previous_mean_log_likelihood: current.mean_log_likelihood,
1551                next_mean_log_likelihood: next.mean_log_likelihood,
1552                numerical_uncertainty: monotonicity_uncertainty,
1553                checkpoint,
1554            });
1555        }
1556        if objective_residual <= config.loglik_tol && parameter_residual <= config.parameter_tol {
1557            let loglik = current.mean_log_likelihood * data.nrows() as f64;
1558            if !loglik.is_finite() {
1559                return Err(GaussianMixtureError::NumericalFailure {
1560                    message: "certified mean log likelihood overflows as a total likelihood"
1561                        .to_string(),
1562                    checkpoint: Some(checkpoint),
1563                });
1564            }
1565            return Ok(GaussianMixtureFit {
1566                weights: checkpoint.weights,
1567                means: checkpoint.means,
1568                covariances: checkpoint.covariances,
1569                k,
1570                d,
1571                n_obs: data.nrows(),
1572                loglik,
1573                iterations: checkpoint.completed_iterations,
1574                certificate,
1575            });
1576        }
1577        if additional_updates >= budget {
1578            // At the budget the question is NOT "have we run long enough?" but
1579            // "is this iterate stuck, or was it interrupted mid-descent?" — and
1580            // the residual window answers it. A rate at or above 1 is a genuine
1581            // stall and refuses exactly as before, now with the evidence
1582            // attached. A contracting rate earns ONE extension, bounded by the
1583            // iterate's own projection `N*`: if it cannot meet the deadline it
1584            // set for itself, that failure is the honest verdict, and the
1585            // certificate reports the rate and projection that priced it.
1586            let extend = match (contraction_rate, projected_iterations_to_tolerance) {
1587                (Some(rate), Some(steps)) if rate < 1.0 && extension.is_none() => Some(steps),
1588                _ => None,
1589            };
1590            match extend {
1591                Some(steps) => {
1592                    // Hard secondary ceiling, derived from the caller's own
1593                    // budget rather than picked: the extension may not exceed
1594                    // the work already authorized. `max_iter` IS the caller's
1595                    // stated work tolerance, so spending at most that much
1596                    // again to finish a provably-converging descent is
1597                    // proportionate, while an iterate whose own projection
1598                    // exceeds it is not "nearly there" and its refusal is
1599                    // honest. Without this a rate of 0.9999 would project six
1600                    // figures of iterations and silently convert a refusal into
1601                    // a hang.
1602                    let steps = steps.min(config.max_iter);
1603                    budget = budget.saturating_add(steps);
1604                    extension = Some(steps);
1605                }
1606                None => {
1607                    return Err(GaussianMixtureError::DidNotConverge {
1608                        max_iterations: budget,
1609                        certificate,
1610                        checkpoint,
1611                    });
1612                }
1613            }
1614        } else if extension.is_some() && additional_updates.is_multiple_of(EM_RATE_WINDOW) {
1615            // Re-validate on the window cadence so a stall inside the extension
1616            // is caught within one window of appearing, not at the projection's
1617            // end. Progress that stops being progress ends the extension.
1618            if !matches!(contraction_rate, Some(rate) if rate < 1.0) {
1619                return Err(GaussianMixtureError::DidNotConverge {
1620                    max_iterations: budget,
1621                    certificate,
1622                    checkpoint,
1623                });
1624            }
1625        }
1626        checkpoint = GaussianMixtureCheckpoint {
1627            weights: next_weights,
1628            means: next_means,
1629            covariances: next_covariances,
1630            mean_log_likelihood: next.mean_log_likelihood,
1631            completed_iterations: checkpoint.completed_iterations + 1,
1632            data_fingerprint,
1633            covariance_floor: config.covariance_floor,
1634        };
1635        additional_updates += 1;
1636    }
1637}
1638
1639struct GaussianMixtureEStep {
1640    responsibilities: Array2<f64>,
1641    row_log_likelihoods: Vec<f64>,
1642    mean_log_likelihood: f64,
1643    mean_log_likelihood_roundoff: f64,
1644}
1645
1646/// Resolution of one observed EM likelihood comparison.
1647///
1648/// `pairwise_mean_with_roundoff` bounds only the final reduction of already
1649/// rounded row log likelihoods. An EM comparison also traverses covariance
1650/// eigendecompositions, precision quadratics, log-sum-exp, the M-step, and a
1651/// second E-step. Treating the reduction bound as a bound for that whole map
1652/// is false precision and turns cancellation at a stationary point into a
1653/// spurious monotonicity violation. The square root of machine epsilon is the
1654/// numerical resolution already required of every configured EM tolerance;
1655/// scaling it by the observed objective magnitude makes the invariant
1656/// independent of data units and of user-selected stopping knobs.
1657fn gaussian_mixture_monotonicity_uncertainty(
1658    objective_scale: f64,
1659    current_reduction_roundoff: f64,
1660    next_reduction_roundoff: f64,
1661) -> f64 {
1662    let reduction_roundoff = current_reduction_roundoff + next_reduction_roundoff;
1663    let composite_map_resolution = f64::EPSILON.sqrt() * objective_scale;
1664    reduction_roundoff.max(composite_map_resolution)
1665}
1666
1667fn pairwise_sum_max_depth(term_count: usize) -> usize {
1668    if term_count <= 1 {
1669        return 0;
1670    }
1671    let within_block = term_count.min(BASE_CHUNK) - 1;
1672    let blocks = term_count.div_ceil(BASE_CHUNK);
1673    let tree_levels = if blocks <= 1 {
1674        0
1675    } else {
1676        (usize::BITS - (blocks - 1).leading_zeros()) as usize
1677    };
1678    within_block.saturating_add(tree_levels)
1679}
1680
1681fn pairwise_mean_with_roundoff(values: &[f64]) -> Result<(f64, f64), String> {
1682    if values.is_empty() || values.iter().any(|value| !value.is_finite()) {
1683        return Err("mean log-likelihood terms must be nonempty and finite".to_string());
1684    }
1685    let sum = pairwise_sum(values);
1686    let magnitudes: Vec<f64> = values.iter().map(|value| value.abs()).collect();
1687    let magnitude_sum = pairwise_sum(&magnitudes);
1688    let unit_roundoff = 0.5 * f64::EPSILON;
1689    let accumulated = pairwise_sum_max_depth(values.len()) as f64 * unit_roundoff;
1690    let addition_bound = if accumulated < 1.0 {
1691        accumulated / (1.0 - accumulated) * magnitude_sum
1692    } else {
1693        f64::INFINITY
1694    };
1695    let count = values.len() as f64;
1696    let mean = sum / count;
1697    // The first term bounds the deterministic pairwise additions; the second
1698    // bounds the final division. This tolerance is derived from the actual
1699    // reduction depth and magnitudes, independently of the EM stopping knob.
1700    let roundoff = addition_bound / count + unit_roundoff * mean.abs();
1701    if !(mean.is_finite() && roundoff.is_finite()) {
1702        return Err("mean mixture log likelihood or its rounding bound is non-finite".to_string());
1703    }
1704    Ok((mean, roundoff))
1705}
1706
1707fn mixture_e_step(
1708    data: ArrayView2<'_, f64>,
1709    weights: &Array1<f64>,
1710    means: &Array2<f64>,
1711    covariances: &[Array2<f64>],
1712) -> Result<GaussianMixtureEStep, String> {
1713    let n = data.nrows();
1714    let k = weights.len();
1715    if weights
1716        .iter()
1717        .any(|weight| !weight.is_finite() || *weight <= 0.0)
1718    {
1719        return Err("mixture E-step requires strictly positive finite weights".to_string());
1720    }
1721    let mut components = Vec::with_capacity(k);
1722    for component in 0..k {
1723        components.push(GaussianComponentEval::factor(
1724            means.row(component),
1725            &covariances[component],
1726        )?);
1727    }
1728    let log_weights: Vec<f64> = weights.iter().map(|weight| weight.ln()).collect();
1729    let mut responsibilities = Array2::<f64>::zeros((n, k));
1730    let mut row_log_likelihoods = Vec::with_capacity(n);
1731    for row in 0..n {
1732        let observation = data.row(row);
1733        let mut log_terms = vec![f64::NEG_INFINITY; k];
1734        let mut max_term = f64::NEG_INFINITY;
1735        for component in 0..k {
1736            let term = log_weights[component] + components[component].log_density(observation);
1737            log_terms[component] = term;
1738            max_term = max_term.max(term);
1739        }
1740        let log_mixture = log_sum_exp(&log_terms, max_term);
1741        if !log_mixture.is_finite() {
1742            return Err(format!(
1743                "mixture density is non-finite at training row {row}"
1744            ));
1745        }
1746        row_log_likelihoods.push(log_mixture);
1747        for component in 0..k {
1748            responsibilities[[row, component]] = (log_terms[component] - log_mixture).exp();
1749        }
1750    }
1751    let (mean_log_likelihood, mean_log_likelihood_roundoff) =
1752        pairwise_mean_with_roundoff(&row_log_likelihoods)?;
1753    Ok(GaussianMixtureEStep {
1754        responsibilities,
1755        row_log_likelihoods,
1756        mean_log_likelihood,
1757        mean_log_likelihood_roundoff,
1758    })
1759}
1760
1761fn mixture_m_step(
1762    data: ArrayView2<'_, f64>,
1763    responsibilities: ArrayView2<'_, f64>,
1764    covariance_floor: f64,
1765) -> Result<(Array1<f64>, Array2<f64>, Vec<Array2<f64>>), String> {
1766    let n = data.nrows();
1767    let d = data.ncols();
1768    let k = responsibilities.ncols();
1769    let mut component_mass = Array1::<f64>::zeros(k);
1770    for component in 0..k {
1771        component_mass[component] = responsibilities.column(component).sum();
1772    }
1773    if component_mass
1774        .iter()
1775        .any(|mass| !mass.is_finite() || *mass <= 0.0)
1776    {
1777        return Err(
1778            "M-step reached a zero-mass component; the requested mixture order has no interior fitted density"
1779                .to_string(),
1780        );
1781    }
1782    let mut weights = component_mass.mapv(|mass| mass / n as f64);
1783    let total_weight = weights.sum();
1784    if !(total_weight.is_finite() && total_weight > 0.0) {
1785        return Err("M-step produced invalid mixture-weight mass".to_string());
1786    }
1787    weights.mapv_inplace(|weight| weight / total_weight);
1788    let mut means = Array2::<f64>::zeros((k, d));
1789    let mut covariances = Vec::with_capacity(k);
1790    for component in 0..k {
1791        let mass = component_mass[component];
1792        let mut mean = Array1::<f64>::zeros(d);
1793        for row in 0..n {
1794            let responsibility = responsibilities[[row, component]];
1795            for col in 0..d {
1796                mean[col] += responsibility * data[[row, col]];
1797            }
1798        }
1799        mean.mapv_inplace(|value| value / mass);
1800        means.row_mut(component).assign(&mean);
1801        let mut covariance = Array2::<f64>::zeros((d, d));
1802        for row in 0..n {
1803            let responsibility = responsibilities[[row, component]];
1804            for left in 0..d {
1805                let left_residual = data[[row, left]] - mean[left];
1806                for right in 0..d {
1807                    covariance[[left, right]] +=
1808                        responsibility * left_residual * (data[[row, right]] - mean[right]);
1809                }
1810            }
1811        }
1812        covariance.mapv_inplace(|value| value / mass);
1813        covariances.push(constrain_covariance(covariance, covariance_floor)?);
1814    }
1815    Ok((weights, means, covariances))
1816}
1817
1818fn relative_parameter_step(previous: f64, next: f64) -> f64 {
1819    (next - previous).abs() / previous.abs().max(next.abs()).max(1.0)
1820}
1821
1822/// Distance between two EM states in the quotient space the empirical
1823/// likelihood can identify.
1824///
1825/// A finite mixture density is invariant to component relabeling and to
1826/// exchanging mass among duplicate components. Ring center/radius/direction
1827/// tuples have additional factorizations of the same component means. No
1828/// component-coordinate norm can therefore be a necessary convergence
1829/// condition. The likelihood sees the vector `(log p(y_i))`; its max-norm
1830/// change is the exact empirical predictive-density residual. Taking the
1831/// maximum (rather than only the mean objective gain) detects row-wise changes
1832/// that cancel, while identical fitted densities have residual zero regardless
1833/// of their internal representation.
1834fn empirical_predictive_density_residual(
1835    previous_row_log_density: &[f64],
1836    next_row_log_density: &[f64],
1837) -> Result<f64, String> {
1838    if previous_row_log_density.is_empty()
1839        || previous_row_log_density.len() != next_row_log_density.len()
1840        || previous_row_log_density
1841            .iter()
1842            .chain(next_row_log_density)
1843            .any(|value| !value.is_finite())
1844    {
1845        return Err(
1846            "predictive-density residual requires equal, nonempty, finite log-density vectors"
1847                .to_string(),
1848        );
1849    }
1850    Ok(previous_row_log_density
1851        .iter()
1852        .zip(next_row_log_density)
1853        .map(|(&previous, &next)| (next - previous).abs())
1854        .fold(0.0_f64, f64::max))
1855}
1856
1857fn constrain_covariance(covariance: Array2<f64>, floor: f64) -> Result<Array2<f64>, String> {
1858    let (eigenvalues, eigenvectors) = covariance
1859        .eigh(Side::Lower)
1860        .map_err(|error| format!("covariance eigendecomposition failed: {error}"))?;
1861    let d = covariance.nrows();
1862    let mut constrained = Array2::<f64>::zeros((d, d));
1863    for row in 0..d {
1864        for col in 0..d {
1865            let mut value = 0.0_f64;
1866            for index in 0..d {
1867                value += eigenvectors[[row, index]]
1868                    * eigenvalues[index].max(floor)
1869                    * eigenvectors[[col, index]];
1870            }
1871            constrained[[row, col]] = value;
1872        }
1873    }
1874    if constrained.iter().any(|value| !value.is_finite()) {
1875        return Err("constrained covariance became non-finite".to_string());
1876    }
1877    Ok(constrained)
1878}
1879
1880/// Global constrained covariance used to seed EM.
1881fn constrained_data_covariance(
1882    data: ArrayView2<'_, f64>,
1883    floor: f64,
1884) -> Result<Array2<f64>, String> {
1885    let n = data.nrows();
1886    let d = data.ncols();
1887    let mut mean = Array1::<f64>::zeros(d);
1888    for i in 0..n {
1889        for c in 0..d {
1890            mean[c] += data[[i, c]];
1891        }
1892    }
1893    mean.mapv_inplace(|v| v / n.max(1) as f64);
1894    let mut cov = Array2::<f64>::zeros((d, d));
1895    for i in 0..n {
1896        for a in 0..d {
1897            let da = data[[i, a]] - mean[a];
1898            for b in 0..d {
1899                cov[[a, b]] += da * (data[[i, b]] - mean[b]);
1900            }
1901        }
1902    }
1903    let inv = 1.0 / n as f64;
1904    cov.mapv_inplace(|v| v * inv);
1905    constrain_covariance(cov, floor)
1906}
1907
1908// ---------------------------------------------------------------------------
1909// Ring-of-clusters candidate (#2262)
1910// ---------------------------------------------------------------------------
1911//
1912// A free Gaussian mixture treats the component means as unrelated points. That
1913// is the wrong null for a discrete cyclic concept: weekdays and months form
1914// tight clusters, but their component means share a low-dimensional circular
1915// constraint. `RingGaussianMixtureFit` models exactly that density,
1916//
1917//     x | z=j ~ N(c + r u_j, sigma^2 I_2),  ||u_j|| = 1,
1918//
1919// with free mixture weights, a shared center/radius, one angle per component,
1920// and a shared isotropic variance. Its `2k + 3` continuous parameters are
1921// priced by the same BIC-form criterion as the unconstrained mixture's `6k - 1`
1922// parameters in two dimensions.
1923
1924/// Certified Gaussian mixture whose component centers lie on one fitted circle.
1925#[derive(Debug, Clone)]
1926pub struct RingGaussianMixtureFit {
1927    weights: Array1<f64>,
1928    center: Array1<f64>,
1929    radius: f64,
1930    directions: Array2<f64>,
1931    variance: f64,
1932    k: usize,
1933    n_obs: usize,
1934    loglik: f64,
1935    iterations: usize,
1936    certificate: GaussianMixtureCertificate,
1937}
1938
1939impl RingGaussianMixtureFit {
1940    pub fn weights(&self) -> ArrayView1<'_, f64> {
1941        self.weights.view()
1942    }
1943
1944    pub fn center(&self) -> ArrayView1<'_, f64> {
1945        self.center.view()
1946    }
1947
1948    pub fn radius(&self) -> f64 {
1949        self.radius
1950    }
1951
1952    pub fn directions(&self) -> ArrayView2<'_, f64> {
1953        self.directions.view()
1954    }
1955
1956    pub fn variance(&self) -> f64 {
1957        self.variance
1958    }
1959
1960    pub fn iterations(&self) -> usize {
1961        self.iterations
1962    }
1963
1964    pub fn certificate(&self) -> GaussianMixtureCertificate {
1965        self.certificate
1966    }
1967
1968    /// Free parameters: `k-1` weight logits, center(2), radius(1), `k`
1969    /// component angles, and shared log standard deviation(1).
1970    pub fn num_free_parameters(&self) -> usize {
1971        2 * self.k + 3
1972    }
1973
1974    pub fn per_point_log_density(&self, data: ArrayView2<'_, f64>) -> Result<Array1<f64>, String> {
1975        if data.ncols() != 2 {
1976            return Err(format!(
1977                "ring-of-clusters density expects two columns, got {}",
1978                data.ncols()
1979            ));
1980        }
1981        ring_mixture_log_density(
1982            data,
1983            &self.weights,
1984            &self.center,
1985            self.radius,
1986            &self.directions,
1987            self.variance,
1988        )
1989    }
1990
1991    /// Schwarz BIC approximation to negative log evidence, divided by two so
1992    /// it is on the ordinary negative-log-likelihood scale. Lower is better.
1993    pub fn bic(&self) -> f64 {
1994        -self.loglik + 0.5 * self.num_free_parameters() as f64 * (self.n_obs as f64).ln()
1995    }
1996}
1997
1998#[derive(Debug, Clone)]
1999struct RingMixtureState {
2000    weights: Array1<f64>,
2001    center: Array1<f64>,
2002    radius: f64,
2003    directions: Array2<f64>,
2004    variance: f64,
2005    mean_log_likelihood: f64,
2006    completed_iterations: usize,
2007}
2008
2009fn ring_component_means(
2010    center: &Array1<f64>,
2011    radius: f64,
2012    directions: &Array2<f64>,
2013) -> Array2<f64> {
2014    let mut means = Array2::<f64>::zeros((directions.nrows(), 2));
2015    for component in 0..directions.nrows() {
2016        means[[component, 0]] = center[0] + radius * directions[[component, 0]];
2017        means[[component, 1]] = center[1] + radius * directions[[component, 1]];
2018    }
2019    means
2020}
2021
2022fn ring_mixture_log_terms(
2023    data: ArrayView2<'_, f64>,
2024    weights: &Array1<f64>,
2025    center: &Array1<f64>,
2026    radius: f64,
2027    directions: &Array2<f64>,
2028    variance: f64,
2029) -> Result<(Array2<f64>, Vec<f64>), String> {
2030    if data.ncols() != 2
2031        || center.len() != 2
2032        || directions.ncols() != 2
2033        || directions.nrows() != weights.len()
2034        || weights
2035            .iter()
2036            .any(|weight| !weight.is_finite() || *weight <= 0.0)
2037        || !(radius.is_finite() && radius > 0.0)
2038        || !(variance.is_finite() && variance > 0.0)
2039    {
2040        return Err("invalid ring-of-clusters parameter state".to_string());
2041    }
2042    let means = ring_component_means(center, radius, directions);
2043    let log_normalizer = -(std::f64::consts::TAU).ln() - variance.ln();
2044    let mut terms = Array2::<f64>::zeros((data.nrows(), weights.len()));
2045    let mut row_log_likelihoods = Vec::with_capacity(data.nrows());
2046    for row in 0..data.nrows() {
2047        let mut max_term = f64::NEG_INFINITY;
2048        for component in 0..weights.len() {
2049            let dx = data[[row, 0]] - means[[component, 0]];
2050            let dy = data[[row, 1]] - means[[component, 1]];
2051            let term =
2052                weights[component].ln() + log_normalizer - 0.5 * (dx * dx + dy * dy) / variance;
2053            terms[[row, component]] = term;
2054            max_term = max_term.max(term);
2055        }
2056        let values = terms.row(row).to_vec();
2057        let log_likelihood = log_sum_exp(&values, max_term);
2058        if !log_likelihood.is_finite() {
2059            return Err(format!(
2060                "ring-of-clusters density is non-finite at training row {row}"
2061            ));
2062        }
2063        row_log_likelihoods.push(log_likelihood);
2064    }
2065    Ok((terms, row_log_likelihoods))
2066}
2067
2068fn ring_mixture_e_step(
2069    data: ArrayView2<'_, f64>,
2070    state: &RingMixtureState,
2071) -> Result<GaussianMixtureEStep, String> {
2072    let (terms, row_log_likelihoods) = ring_mixture_log_terms(
2073        data,
2074        &state.weights,
2075        &state.center,
2076        state.radius,
2077        &state.directions,
2078        state.variance,
2079    )?;
2080    let mut responsibilities = Array2::<f64>::zeros(terms.raw_dim());
2081    for row in 0..terms.nrows() {
2082        for component in 0..terms.ncols() {
2083            responsibilities[[row, component]] =
2084                (terms[[row, component]] - row_log_likelihoods[row]).exp();
2085        }
2086    }
2087    let (mean_log_likelihood, mean_log_likelihood_roundoff) =
2088        pairwise_mean_with_roundoff(&row_log_likelihoods)?;
2089    Ok(GaussianMixtureEStep {
2090        responsibilities,
2091        row_log_likelihoods,
2092        mean_log_likelihood,
2093        mean_log_likelihood_roundoff,
2094    })
2095}
2096
2097fn ring_mixture_log_density(
2098    data: ArrayView2<'_, f64>,
2099    weights: &Array1<f64>,
2100    center: &Array1<f64>,
2101    radius: f64,
2102    directions: &Array2<f64>,
2103    variance: f64,
2104) -> Result<Array1<f64>, String> {
2105    let (_, row_log_likelihoods) =
2106        ring_mixture_log_terms(data, weights, center, radius, directions, variance)?;
2107    Ok(Array1::from_vec(row_log_likelihoods))
2108}
2109
2110fn fit_weighted_component_circle(
2111    component_means: &Array2<f64>,
2112    component_mass: &Array1<f64>,
2113    initial_center: &Array1<f64>,
2114    initial_radius: f64,
2115    parameter_tol: f64,
2116    max_iter: usize,
2117) -> Result<(Array1<f64>, f64, Array2<f64>), String> {
2118    let k = component_means.nrows();
2119    let total_mass = component_mass.sum();
2120    if component_means.ncols() != 2
2121        || component_mass.len() != k
2122        || component_mass
2123            .iter()
2124            .any(|mass| !mass.is_finite() || *mass <= 0.0)
2125        || !(total_mass.is_finite() && total_mass > 0.0)
2126    {
2127        return Err("ring M-step requires positive component masses and 2-D means".to_string());
2128    }
2129    let mut center = initial_center.clone();
2130    let mut radius = initial_radius;
2131    let mut directions = Array2::<f64>::zeros((k, 2));
2132    for _ in 0..max_iter {
2133        for component in 0..k {
2134            let dx = component_means[[component, 0]] - center[0];
2135            let dy = component_means[[component, 1]] - center[1];
2136            let norm = dx.hypot(dy);
2137            if !(norm.is_finite() && norm > 0.0) {
2138                return Err(
2139                    "ring M-step reached a component centroid at the circle center; its angle is unidentified"
2140                        .to_string(),
2141                );
2142            }
2143            directions[[component, 0]] = dx / norm;
2144            directions[[component, 1]] = dy / norm;
2145        }
2146
2147        let mut mean_point = Array1::<f64>::zeros(2);
2148        let mut mean_direction = Array1::<f64>::zeros(2);
2149        for component in 0..k {
2150            let weight = component_mass[component] / total_mass;
2151            for axis in 0..2 {
2152                mean_point[axis] += weight * component_means[[component, axis]];
2153                mean_direction[axis] += weight * directions[[component, axis]];
2154            }
2155        }
2156        let mut numerator = 0.0;
2157        let mut denominator = 0.0;
2158        for component in 0..k {
2159            let mass = component_mass[component];
2160            let dux = directions[[component, 0]] - mean_direction[0];
2161            let duy = directions[[component, 1]] - mean_direction[1];
2162            numerator += mass
2163                * (dux * (component_means[[component, 0]] - mean_point[0])
2164                    + duy * (component_means[[component, 1]] - mean_point[1]));
2165            denominator += mass * (dux * dux + duy * duy);
2166        }
2167        if !(denominator.is_finite() && denominator > 0.0) {
2168            return Err(
2169                "ring M-step component directions are identical; radius and center are unidentified"
2170                    .to_string(),
2171            );
2172        }
2173        let mut next_radius = numerator / denominator;
2174        if !next_radius.is_finite() || next_radius == 0.0 {
2175            return Err("ring M-step produced an unidentified zero radius".to_string());
2176        }
2177        if next_radius < 0.0 {
2178            next_radius = -next_radius;
2179            directions.mapv_inplace(|value| -value);
2180        }
2181        let next_center = Array1::from_vec(vec![
2182            mean_point[0] - next_radius * mean_direction[0],
2183            mean_point[1] - next_radius * mean_direction[1],
2184        ]);
2185        let residual = center
2186            .iter()
2187            .zip(next_center.iter())
2188            .map(|(&left, &right)| relative_parameter_step(left, right))
2189            .chain(std::iter::once(relative_parameter_step(
2190                radius,
2191                next_radius,
2192            )))
2193            .fold(0.0, f64::max);
2194        center = next_center;
2195        radius = next_radius;
2196        if residual <= parameter_tol {
2197            // Recompute directions at the returned center so the stored angles
2198            // are the exact angular block update belonging to that center.
2199            for component in 0..k {
2200                let dx = component_means[[component, 0]] - center[0];
2201                let dy = component_means[[component, 1]] - center[1];
2202                let norm = dx.hypot(dy);
2203                if !(norm.is_finite() && norm > 0.0) {
2204                    return Err("ring M-step terminal component angle is unidentified".to_string());
2205                }
2206                directions[[component, 0]] = dx / norm;
2207                directions[[component, 1]] = dy / norm;
2208            }
2209            return Ok((center, radius, directions));
2210        }
2211    }
2212    Err(format!(
2213        "ring M-step did not certify its constrained center/radius fixed point after {max_iter} iterations"
2214    ))
2215}
2216
2217fn ring_mixture_m_step(
2218    data: ArrayView2<'_, f64>,
2219    responsibilities: ArrayView2<'_, f64>,
2220    previous: &RingMixtureState,
2221    config: GaussianMixtureConfig,
2222) -> Result<RingMixtureState, String> {
2223    let n = data.nrows();
2224    let k = responsibilities.ncols();
2225    let mut component_mass = Array1::<f64>::zeros(k);
2226    let mut component_means = Array2::<f64>::zeros((k, 2));
2227    for component in 0..k {
2228        let mass = responsibilities.column(component).sum();
2229        if !(mass.is_finite() && mass > 0.0) {
2230            return Err(
2231                "ring M-step reached a zero-mass component; the requested order is singular"
2232                    .to_string(),
2233            );
2234        }
2235        component_mass[component] = mass;
2236        for row in 0..n {
2237            for axis in 0..2 {
2238                component_means[[component, axis]] +=
2239                    responsibilities[[row, component]] * data[[row, axis]];
2240            }
2241        }
2242        for axis in 0..2 {
2243            component_means[[component, axis]] /= mass;
2244        }
2245    }
2246    let mut weights = component_mass.mapv(|mass| mass / n as f64);
2247    let weight_sum = weights.sum();
2248    weights.mapv_inplace(|weight| weight / weight_sum);
2249    let (center, radius, directions) = fit_weighted_component_circle(
2250        &component_means,
2251        &component_mass,
2252        &previous.center,
2253        previous.radius,
2254        config.parameter_tol,
2255        config.max_iter,
2256    )?;
2257    let means = ring_component_means(&center, radius, &directions);
2258    let mut expected_squared_error = 0.0;
2259    for row in 0..n {
2260        for component in 0..k {
2261            let dx = data[[row, 0]] - means[[component, 0]];
2262            let dy = data[[row, 1]] - means[[component, 1]];
2263            expected_squared_error += responsibilities[[row, component]] * (dx * dx + dy * dy);
2264        }
2265    }
2266    let variance = (expected_squared_error / (2 * n) as f64).max(config.covariance_floor);
2267    if !variance.is_finite() {
2268        return Err("ring M-step produced non-finite shared variance".to_string());
2269    }
2270    Ok(RingMixtureState {
2271        weights,
2272        center,
2273        radius,
2274        directions,
2275        variance,
2276        mean_log_likelihood: f64::NAN,
2277        completed_iterations: previous.completed_iterations + 1,
2278    })
2279}
2280
2281/// Fit a deterministic, certified `k`-component isotropic Gaussian mixture
2282/// whose component centers are constrained to a common circle.
2283pub fn fit_ring_gaussian_mixture(
2284    data: ArrayView2<'_, f64>,
2285    k: usize,
2286    config: GaussianMixtureConfig,
2287) -> Result<RingGaussianMixtureFit, String> {
2288    validate_gaussian_mixture_problem(data, k, config).map_err(|error| error.to_string())?;
2289    if data.ncols() != 2 {
2290        return Err(format!(
2291            "ring-of-clusters fitting requires exactly two columns, got {}",
2292            data.ncols()
2293        ));
2294    }
2295    if k < 3 {
2296        return Err(format!(
2297            "ring-of-clusters fitting requires at least three component centers, got {k}"
2298        ));
2299    }
2300    let seeded_means = gam_terms::basis::select_centers_by_strategy(
2301        data,
2302        &gam_terms::basis::CenterStrategy::KMeans {
2303            num_centers: k,
2304            max_iter: config.kmeans_max_iter,
2305        },
2306    )
2307    .map_err(|error| format!("ring-of-clusters deterministic seeding failed: {error}"))?;
2308    let component_mass = Array1::<f64>::ones(k);
2309    let mut initial_center = Array1::<f64>::zeros(2);
2310    for component in 0..k {
2311        initial_center[0] += seeded_means[[component, 0]] / k as f64;
2312        initial_center[1] += seeded_means[[component, 1]] / k as f64;
2313    }
2314    let mut initial_radius = 0.0;
2315    for component in 0..k {
2316        initial_radius += (seeded_means[[component, 0]] - initial_center[0])
2317            .hypot(seeded_means[[component, 1]] - initial_center[1])
2318            / k as f64;
2319    }
2320    if !(initial_radius.is_finite() && initial_radius > 0.0) {
2321        return Err("ring-of-clusters seed has an unidentified zero radius".to_string());
2322    }
2323    let (center, radius, directions) = fit_weighted_component_circle(
2324        &seeded_means,
2325        &component_mass,
2326        &initial_center,
2327        initial_radius,
2328        config.parameter_tol,
2329        config.max_iter,
2330    )?;
2331    let means = ring_component_means(&center, radius, &directions);
2332    let mut squared_error = 0.0;
2333    for row in 0..data.nrows() {
2334        let mut nearest = f64::INFINITY;
2335        for component in 0..k {
2336            let dx = data[[row, 0]] - means[[component, 0]];
2337            let dy = data[[row, 1]] - means[[component, 1]];
2338            nearest = nearest.min(dx * dx + dy * dy);
2339        }
2340        squared_error += nearest;
2341    }
2342    let variance = (squared_error / (2 * data.nrows()) as f64).max(config.covariance_floor);
2343    let mut state = RingMixtureState {
2344        weights: Array1::from_elem(k, 1.0 / k as f64),
2345        center,
2346        radius,
2347        directions,
2348        variance,
2349        mean_log_likelihood: f64::NAN,
2350        completed_iterations: 0,
2351    };
2352    for additional_updates in 0..=config.max_iter {
2353        let current = ring_mixture_e_step(data, &state)?;
2354        state.mean_log_likelihood = current.mean_log_likelihood;
2355        let mut next =
2356            ring_mixture_m_step(data, current.responsibilities.view(), &state, config)?;
2357        let next_e_step = ring_mixture_e_step(data, &next)?;
2358        next.mean_log_likelihood = next_e_step.mean_log_likelihood;
2359        let current_mean = current.mean_log_likelihood;
2360        let next_mean = next_e_step.mean_log_likelihood;
2361        let objective_scale = current_mean.abs().max(next_mean.abs()).max(1.0);
2362        let objective_step = next_mean - current_mean;
2363        let objective_residual = objective_step.abs() / objective_scale;
2364        let parameter_residual = empirical_predictive_density_residual(
2365            &current.row_log_likelihoods,
2366            &next_e_step.row_log_likelihoods,
2367        )?;
2368        let monotonicity_uncertainty = gaussian_mixture_monotonicity_uncertainty(
2369            objective_scale,
2370            current.mean_log_likelihood_roundoff,
2371            next_e_step.mean_log_likelihood_roundoff,
2372        );
2373        let certificate = GaussianMixtureCertificate {
2374            mean_log_likelihood: current_mean,
2375            mean_log_likelihood_gain: objective_step,
2376            monotonicity_uncertainty,
2377            objective_residual,
2378            objective_tolerance: config.loglik_tol,
2379            parameter_residual,
2380            parameter_tolerance: config.parameter_tol,
2381            // The ring-of-clusters rung does not yet measure a contraction rate,
2382            // so its exhaustion stays un-priced. Reporting `None` says exactly
2383            // that; fabricating a rate here would be the invention the rest of
2384            // this certificate exists to prevent.
2385            contraction_rate: None,
2386            projected_iterations_to_tolerance: None,
2387        };
2388        if objective_step < -monotonicity_uncertainty {
2389            return Err(format!(
2390                "ring-of-clusters generalized EM violated monotone ascent at iteration {}: {current_mean:.12e} -> {next_mean:.12e} (comparison uncertainty {monotonicity_uncertainty:.3e})",
2391                state.completed_iterations
2392            ));
2393        }
2394        if objective_residual <= config.loglik_tol && parameter_residual <= config.parameter_tol {
2395            let loglik = current_mean * data.nrows() as f64;
2396            if !loglik.is_finite() {
2397                return Err("ring-of-clusters total log likelihood overflowed".to_string());
2398            }
2399            return Ok(RingGaussianMixtureFit {
2400                weights: state.weights,
2401                center: state.center,
2402                radius: state.radius,
2403                directions: state.directions,
2404                variance: state.variance,
2405                k,
2406                n_obs: data.nrows(),
2407                loglik,
2408                iterations: state.completed_iterations,
2409                certificate,
2410            });
2411        }
2412        if additional_updates == config.max_iter {
2413            return Err(format!(
2414                "ring-of-clusters generalized EM did not certify after {} iterations: objective residual {:.6e}/{:.3e}, parameter-map residual {:.6e}/{:.3e}",
2415                config.max_iter,
2416                objective_residual,
2417                config.loglik_tol,
2418                parameter_residual,
2419                config.parameter_tol,
2420            ));
2421        }
2422        state = next;
2423    }
2424    Err("ring-of-clusters generalized EM exhausted without a terminal certificate".to_string())
2425}
2426
2427// ---------------------------------------------------------------------------
2428// Circular Gaussian density and structured-union candidates (#907)
2429// ---------------------------------------------------------------------------
2430
2431/// Maximum-likelihood fit of a Gaussian-blurred circle in two dimensions.
2432///
2433/// The generative model is
2434///
2435/// `X = center + radius * U + epsilon`,
2436///
2437/// where `U` is uniform on the unit circle and
2438/// `epsilon ~ N(0, noise_variance * I_2)`. Integrating out `U` gives the proper
2439/// Cartesian density
2440///
2441/// `p(x) = exp(-(r^2 + R^2)/(2s)) I0(Rr/s) / (2 pi s)`.
2442///
2443/// Unlike a Gaussian density assigned directly to the nonnegative radius, this
2444/// density is normalized on the plane, remains finite at the center, and has no
2445/// artificial `1/r` singularity. The center is fitted jointly with `(R, s)` by
2446/// latent-angle EM instead of being frozen at the coordinate mean.
2447#[derive(Debug, Clone, Copy)]
2448pub struct CircularGaussianFit2d {
2449    center: [f64; 2],
2450    radius: f64,
2451    noise_variance: f64,
2452}
2453
2454impl CircularGaussianFit2d {
2455    /// Two center coordinates, one radius, and one isotropic noise variance.
2456    pub const NUM_FREE_PARAMETERS: usize = 4;
2457
2458    /// Construct a circular Gaussian from validated model parameters.
2459    pub fn from_parameters(
2460        center: [f64; 2],
2461        radius: f64,
2462        noise_variance: f64,
2463    ) -> Result<Self, String> {
2464        if !center.iter().all(|value| value.is_finite()) {
2465            return Err("circular Gaussian center must be finite".to_string());
2466        }
2467        if !(radius.is_finite() && radius >= 0.0) {
2468            return Err("circular Gaussian radius must be finite and nonnegative".to_string());
2469        }
2470        if !(noise_variance.is_finite() && noise_variance > 0.0) {
2471            return Err("circular Gaussian noise variance must be finite and positive".to_string());
2472        }
2473        Ok(Self {
2474            center,
2475            radius,
2476            noise_variance,
2477        })
2478    }
2479
2480    /// Fit selected rows of a finite two-column coordinate matrix.
2481    pub fn fit(coords: ArrayView2<'_, f64>, rows: &[usize]) -> Result<Self, String> {
2482        if coords.ncols() != 2 {
2483            return Err(format!(
2484                "circular Gaussian requires 2-D data, got {} columns",
2485                coords.ncols()
2486            ));
2487        }
2488        if rows.is_empty() {
2489            return Err("circular Gaussian requires a nonempty training set".to_string());
2490        }
2491        if rows.iter().any(|&row| row >= coords.nrows()) {
2492            return Err("circular Gaussian row index is out of bounds".to_string());
2493        }
2494        if rows
2495            .iter()
2496            .any(|&row| !coords[[row, 0]].is_finite() || !coords[[row, 1]].is_finite())
2497        {
2498            return Err("circular Gaussian requires finite training coordinates".to_string());
2499        }
2500
2501        // Work in a dimensionless chart relative to one observed point. This
2502        // preserves the low-order bits of a small translated circle and makes
2503        // the stopping rule and variance floor scale equivariant.
2504        let anchor_row = rows[0];
2505        let anchor = [coords[[anchor_row, 0]], coords[[anchor_row, 1]]];
2506        let mut scale = 0.0_f64;
2507        for &row in rows {
2508            let dx = coords[[row, 0]] - anchor[0];
2509            let dy = coords[[row, 1]] - anchor[1];
2510            if !(dx.is_finite() && dy.is_finite()) {
2511                return Err("circular Gaussian coordinate range exceeds f64".to_string());
2512            }
2513            scale = scale.max(dx.hypot(dy));
2514        }
2515        if !(scale.is_finite() && scale > 0.0) {
2516            return Err("circular Gaussian requires nonzero spatial extent".to_string());
2517        }
2518
2519        let mut points = Vec::with_capacity(rows.len());
2520        let mut mean = [0.0_f64; 2];
2521        for &row in rows {
2522            let point = [
2523                (coords[[row, 0]] - anchor[0]) / scale,
2524                (coords[[row, 1]] - anchor[1]) / scale,
2525            ];
2526            points.push(point);
2527            mean[0] += point[0];
2528            mean[1] += point[1];
2529        }
2530        let count = rows.len() as f64;
2531        mean[0] /= count;
2532        mean[1] /= count;
2533
2534        // Moment initialization is exact at the population level. For
2535        // q = ||X-E X||^2,
2536        //   E[q] = R^2 + 2s,  Var(q) = 4s(R^2+s),
2537        // hence R^4 = E[q]^2-Var(q) and s=(E[q]-R^2)/2.
2538        let mut squared_radii = Vec::with_capacity(rows.len());
2539        let mut mean_squared_radius = 0.0_f64;
2540        for point in &points {
2541            let dx = point[0] - mean[0];
2542            let dy = point[1] - mean[1];
2543            let squared_radius = dx * dx + dy * dy;
2544            squared_radii.push(squared_radius);
2545            mean_squared_radius += squared_radius;
2546        }
2547        mean_squared_radius /= count;
2548        let mut squared_radius_variance = 0.0_f64;
2549        for squared_radius in squared_radii {
2550            squared_radius_variance += (squared_radius - mean_squared_radius).powi(2);
2551        }
2552        squared_radius_variance /= count;
2553
2554        // A noiseless observed circle is an unbounded-likelihood boundary.
2555        // Keep the numerical optimizer in a scale-relative interior whose
2556        // width is roundoff, rather than imposing a floor in data units.
2557        let variance_floor = (64.0 * f64::EPSILON * mean_squared_radius).max(f64::MIN_POSITIVE);
2558        let radius_squared = (mean_squared_radius * mean_squared_radius - squared_radius_variance)
2559            .max(0.0)
2560            .sqrt();
2561        let mut radius = radius_squared.sqrt();
2562        let mut noise_variance = (0.5 * (mean_squared_radius - radius_squared)).max(variance_floor);
2563        let mut center = mean;
2564
2565        // Exact EM for the latent circle angle. Given current parameters, the
2566        // conditional mean of U is A(kappa) * (x-c)/||x-c|| with
2567        // A=I1/I0 and kappa=R||x-c||/s. Solving the joint quadratic M-step for
2568        // center and radius avoids the biased `center = sample mean` plug-in.
2569        const MAX_EM_ITERATIONS: usize = 4096;
2570        const EM_TOLERANCE: f64 = 2.0e-12;
2571        let mut posterior_means = vec![[0.0_f64; 2]; points.len()];
2572        let mut converged = false;
2573        for _ in 0..MAX_EM_ITERATIONS {
2574            let mut posterior_mean = [0.0_f64; 2];
2575            for (point, latent_mean) in points.iter().zip(&mut posterior_means) {
2576                let dx = point[0] - center[0];
2577                let dy = point[1] - center[1];
2578                let observed_radius = dx.hypot(dy);
2579                if observed_radius == 0.0 || radius == 0.0 {
2580                    *latent_mean = [0.0, 0.0];
2581                } else {
2582                    let (_, bessel_ratio) =
2583                        circular_gaussian_bessel_terms(radius, observed_radius, noise_variance);
2584                    if !(bessel_ratio.is_finite() && (0.0..=1.0).contains(&bessel_ratio)) {
2585                        return Err("circular Gaussian Bessel ratio left [0, 1]".to_string());
2586                    }
2587                    let multiplier = bessel_ratio / observed_radius;
2588                    *latent_mean = [multiplier * dx, multiplier * dy];
2589                }
2590                posterior_mean[0] += latent_mean[0];
2591                posterior_mean[1] += latent_mean[1];
2592            }
2593            posterior_mean[0] /= count;
2594            posterior_mean[1] /= count;
2595
2596            let denominator =
2597                1.0 - posterior_mean[0] * posterior_mean[0] - posterior_mean[1] * posterior_mean[1];
2598            if !(denominator.is_finite() && denominator > 0.0) {
2599                return Err("circular Gaussian EM radius update is singular".to_string());
2600            }
2601            let mut radius_numerator = 0.0_f64;
2602            for (point, latent_mean) in points.iter().zip(&posterior_means) {
2603                radius_numerator +=
2604                    latent_mean[0] * (point[0] - mean[0]) + latent_mean[1] * (point[1] - mean[1]);
2605            }
2606            let next_radius = (radius_numerator / (count * denominator)).max(0.0);
2607            let next_center = [
2608                mean[0] - next_radius * posterior_mean[0],
2609                mean[1] - next_radius * posterior_mean[1],
2610            ];
2611
2612            // Evaluate E||X-c-RU||^2 in an explicitly nonnegative form to
2613            // avoid catastrophic cancellation on a very thin ring.
2614            let mut residual_sum = 0.0_f64;
2615            for (point, latent_mean) in points.iter().zip(&posterior_means) {
2616                let dx = point[0] - next_center[0];
2617                let dy = point[1] - next_center[1];
2618                let ex = dx - next_radius * latent_mean[0];
2619                let ey = dy - next_radius * latent_mean[1];
2620                let latent_norm_squared =
2621                    latent_mean[0] * latent_mean[0] + latent_mean[1] * latent_mean[1];
2622                residual_sum += ex * ex
2623                    + ey * ey
2624                    + next_radius * next_radius * (1.0 - latent_norm_squared).max(0.0);
2625            }
2626            let next_noise_variance = (residual_sum / (2.0 * count)).max(variance_floor);
2627
2628            let parameter_change = (next_center[0] - center[0])
2629                .hypot(next_center[1] - center[1])
2630                .max((next_radius - radius).abs())
2631                .max(
2632                    (next_noise_variance - noise_variance).abs()
2633                        / (next_noise_variance + noise_variance),
2634                );
2635            center = next_center;
2636            radius = next_radius;
2637            noise_variance = next_noise_variance;
2638            if parameter_change <= EM_TOLERANCE {
2639                converged = true;
2640                break;
2641            }
2642        }
2643        if !converged {
2644            return Err("circular Gaussian maximum-likelihood fit did not converge".to_string());
2645        }
2646
2647        let fitted_noise_sd = scale * noise_variance.sqrt();
2648        Self::from_parameters(
2649            [anchor[0] + scale * center[0], anchor[1] + scale * center[1]],
2650            scale * radius,
2651            fitted_noise_sd * fitted_noise_sd,
2652        )
2653        .map_err(|error| format!("circular Gaussian fit produced invalid parameters: {error}"))
2654    }
2655
2656    /// Fitted circle center.
2657    pub const fn center(self) -> [f64; 2] {
2658        self.center
2659    }
2660
2661    /// Fitted latent-circle radius.
2662    pub const fn radius(self) -> f64 {
2663        self.radius
2664    }
2665
2666    /// Fitted isotropic Cartesian noise variance per coordinate.
2667    pub const fn noise_variance(self) -> f64 {
2668        self.noise_variance
2669    }
2670
2671    /// Proper Cartesian log density at `(x, y)`.
2672    pub fn log_density(self, x: f64, y: f64) -> f64 {
2673        let observed_radius = (x - self.center[0]).hypot(y - self.center[1]);
2674        let (log_i0_minus_kappa, _) =
2675            circular_gaussian_bessel_terms(self.radius, observed_radius, self.noise_variance);
2676        let standardized_radial_residual =
2677            (observed_radius - self.radius) / self.noise_variance.sqrt();
2678        // Algebraically this is -(r^2+R^2)/(2s)+log I0(kappa), but this
2679        // rearrangement preserves log I0(kappa) ~= kappa cancellation.
2680        -std::f64::consts::TAU.ln()
2681            - self.noise_variance.ln()
2682            - 0.5 * standardized_radial_residual.powi(2)
2683            + log_i0_minus_kappa
2684    }
2685
2686    /// Sum the fitted log density over selected rows.
2687    pub fn log_likelihood(
2688        self,
2689        coords: ArrayView2<'_, f64>,
2690        rows: &[usize],
2691    ) -> Result<f64, String> {
2692        if coords.ncols() != 2 || rows.iter().any(|&row| row >= coords.nrows()) {
2693            return Err(
2694                "circular Gaussian likelihood received invalid coordinates or rows".to_string(),
2695            );
2696        }
2697        let mut log_densities = Vec::with_capacity(rows.len());
2698        for &row in rows {
2699            let value = self.log_density(coords[[row, 0]], coords[[row, 1]]);
2700            if !value.is_finite() {
2701                return Err("circular Gaussian likelihood is not finite".to_string());
2702            }
2703            log_densities.push(value);
2704        }
2705        let log_likelihood = pairwise_sum(&log_densities);
2706        if !log_likelihood.is_finite() {
2707            return Err("circular Gaussian likelihood sum is not finite".to_string());
2708        }
2709        Ok(log_likelihood)
2710    }
2711
2712    /// Fit selected rows and return both the fit and its BIC/2 (lower is
2713    /// better). Keeping fitting and evidence evaluation in one operation makes
2714    /// it impossible to label a likelihood on unrelated data as fitted BIC.
2715    pub fn fit_with_bic(
2716        coords: ArrayView2<'_, f64>,
2717        rows: &[usize],
2718    ) -> Result<(Self, f64), String> {
2719        let fit = Self::fit(coords, rows)?;
2720        let log_likelihood = fit.log_likelihood(coords, rows)?;
2721        let bic =
2722            -log_likelihood + 0.5 * Self::NUM_FREE_PARAMETERS as f64 * (rows.len() as f64).ln();
2723        if !bic.is_finite() {
2724            return Err("circular Gaussian BIC is not finite".to_string());
2725        }
2726        Ok((fit, bic))
2727    }
2728}
2729
2730/// Stable Bessel terms for `kappa = radius * observed_radius / variance`.
2731/// The ordinary finite branch retains the shared approximation exactly. If the
2732/// product itself overflows, only the leading asymptotic
2733/// `log I0(kappa)-kappa = -log(2 pi kappa)/2 + O(1/kappa)` is representable;
2734/// `I1/I0` rounds to one at that scale.
2735fn circular_gaussian_bessel_terms(
2736    radius: f64,
2737    observed_radius: f64,
2738    noise_variance: f64,
2739) -> (f64, f64) {
2740    if radius == 0.0 || observed_radius == 0.0 {
2741        return (0.0, 0.0);
2742    }
2743    let kappa = radius * observed_radius / noise_variance;
2744    if kappa.is_finite() {
2745        return bessel_i0_log_minus_abs_and_ratio(kappa);
2746    }
2747    let log_kappa = radius.ln() + observed_radius.ln() - noise_variance.ln();
2748    if log_kappa <= f64::MAX.ln() {
2749        // The left-to-right product overflowed before a large variance divided
2750        // it back into range. Reconstruct the representable ratio from its log.
2751        return bessel_i0_log_minus_abs_and_ratio(log_kappa.exp());
2752    }
2753    (-0.5 * (std::f64::consts::TAU.ln() + log_kappa), 1.0)
2754}
2755//
2756// A *union* candidate is a small FIXED composite of named component structures
2757// joined by a hard row-responsibility split. Unlike the discrete-mixture rung
2758// (which is one free k-component Gaussian density), a union pins each component
2759// to a specific generative STRUCTURE (a circle, a line, a point cluster) and
2760// asks whether the data is better explained as the disjoint sum of those
2761// structures than by any single pure rung.
2762//
2763// The hard responsibility groups only determine which rows fit each component.
2764// The resulting candidate is one normalized, unlabeled soft-mixture density
2765// `p(y) = Σ_c π_c p_c(y)`, with `π_c = n_c / n`. It is scored on every
2766// training row by that same mixture density used for held-out evaluation. Its
2767// BIC/2 complexity price is `½(Σ_c P_c + m - 1) log(n)`: component
2768// parameters plus the `m - 1` free mixing weights, all on the common sample
2769// scale. This makes the union directly comparable to the other normalized
2770// parametric candidates in the topology race.
2771
2772/// The fixed ladder of structured-union composites. Deterministic and closed:
2773/// open-ended structure search stays owned by #976's move set; these three are
2774/// the only composites the topology race may select.
2775#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
2776pub enum UnionStructure {
2777    /// Two circles (two well-separated periodic loops).
2778    CircleCircle,
2779    /// One circle plus one isolated point cluster (a loop with an outlier blob).
2780    CirclePointCluster,
2781    /// One line (anisotropic cluster) plus one isolated point cluster.
2782    LineCluster,
2783}
2784
2785/// The fixed structured-union ladder, in stable order.
2786pub const UNION_STRUCTURE_LADDER: &[UnionStructure] = &[
2787    UnionStructure::CircleCircle,
2788    UnionStructure::CirclePointCluster,
2789    UnionStructure::LineCluster,
2790];
2791
2792/// The per-component generative structure a union pins each responsibility group
2793/// to. `Line` is a full-covariance Gaussian, while `PointCluster` is the nested
2794/// isotropic Gaussian with `d + 1` parameters. The covariance constraint makes
2795/// line+cluster a genuine structured alternative to a generic two-component
2796/// full-covariance mixture instead of a duplicate candidate. `Circle` is the
2797/// proper Cartesian density of a uniform latent circle convolved with isotropic
2798/// Gaussian noise.
2799#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
2800pub enum UnionComponentKind {
2801    Circle,
2802    Line,
2803    PointCluster,
2804}
2805
2806impl UnionStructure {
2807    /// Stable display name, e.g. `"union_circle+circle"`.
2808    pub const fn as_str(self) -> &'static str {
2809        match self {
2810            UnionStructure::CircleCircle => "union_circle+circle",
2811            UnionStructure::CirclePointCluster => "union_circle+cluster",
2812            UnionStructure::LineCluster => "union_line+cluster",
2813        }
2814    }
2815
2816    /// The fixed ordered component structures of this union.
2817    pub const fn components(self) -> &'static [UnionComponentKind] {
2818        match self {
2819            UnionStructure::CircleCircle => {
2820                &[UnionComponentKind::Circle, UnionComponentKind::Circle]
2821            }
2822            UnionStructure::CirclePointCluster => {
2823                &[UnionComponentKind::Circle, UnionComponentKind::PointCluster]
2824            }
2825            UnionStructure::LineCluster => {
2826                &[UnionComponentKind::Line, UnionComponentKind::PointCluster]
2827            }
2828        }
2829    }
2830
2831    /// Number of components (= the responsibility-split order `m`).
2832    pub const fn num_components(self) -> usize {
2833        self.components().len()
2834    }
2835}
2836
2837/// One fitted component of a union: its pinned structure, the rows used to fit
2838/// it after the hard responsibility split, its free-parameter count, and its
2839/// normalized soft-mixture weight. A component has no standalone BIC inside a
2840/// union: the likelihood is the indivisible `log Σ_c π_c p_c(y)` scored on
2841/// every row.
2842#[derive(Debug, Clone)]
2843pub struct UnionComponentFit {
2844    pub kind: UnionComponentKind,
2845    pub row_count: usize,
2846    pub num_parameters: usize,
2847    pub mixing_weight: f64,
2848}
2849
2850/// A fitted structured-union candidate: the composite kind, the per-component
2851/// fits, its normalized soft-mixture training likelihood, the corresponding
2852/// BIC-form negative-log-evidence, and the complete free-parameter count.
2853#[derive(Debug, Clone)]
2854pub struct UnionStructureFit {
2855    pub structure: UnionStructure,
2856    pub components: Vec<UnionComponentFit>,
2857    /// `Σ_i log(Σ_c π_c p_c(y_i))` over all training rows.
2858    pub log_likelihood: f64,
2859    /// `-log_likelihood + ½ total_parameters log(n)` (lower wins).
2860    pub bic: f64,
2861    /// `Σ_c P_c + (m - 1)`, including the free mixing weights.
2862    pub total_parameters: usize,
2863}
2864
2865/// Hard responsibility split of `0..n` into `m` groups by argmax of the
2866/// deterministic `m`-component Gaussian-mixture responsibilities. Reuses the
2867/// mixture rung's seeding + EM so the split is a pure function of the data and
2868/// `m` (no clock). Returns one row-index vector per component.
2869pub fn union_responsibility_split(
2870    data: ArrayView2<'_, f64>,
2871    m: usize,
2872    config: GaussianMixtureConfig,
2873) -> Result<Vec<Vec<usize>>, String> {
2874    let n = data.nrows();
2875    if m == 0 {
2876        return Err("union split requires at least one component".to_string());
2877    }
2878    if m > n {
2879        return Err(format!(
2880            "union split requested {m} groups but data has {n} rows"
2881        ));
2882    }
2883    if m == 1 {
2884        return Ok(vec![(0..n).collect()]);
2885    }
2886    let fit = fit_gaussian_mixture(data, m, config).map_err(|error| error.to_string())?;
2887    let mut groups: Vec<Vec<usize>> = vec![Vec::new(); m];
2888    // Hard assignment by argmax per-component log responsibility.
2889    let mut comp = Vec::with_capacity(m);
2890    for j in 0..m {
2891        comp.push(GaussianComponentEval::factor(
2892            fit.means.row(j),
2893            &fit.covariances[j],
2894        )?);
2895    }
2896    let log_w = fit
2897        .weights
2898        .iter()
2899        .enumerate()
2900        .map(|(component, &weight)| {
2901            if weight.is_finite() && weight > 0.0 {
2902                Ok(weight.ln())
2903            } else {
2904                Err(format!(
2905                    "union split received invalid fitted weight {weight} for component {component}"
2906                ))
2907            }
2908        })
2909        .collect::<Result<Vec<_>, _>>()?;
2910    for i in 0..n {
2911        let row = data.row(i);
2912        let mut best_j = 0usize;
2913        let mut best_lt = f64::NEG_INFINITY;
2914        for j in 0..m {
2915            let lt = log_w[j] + comp[j].log_density(row);
2916            if lt > best_lt {
2917                best_lt = lt;
2918                best_j = j;
2919            }
2920        }
2921        if !best_lt.is_finite() {
2922            return Err(format!(
2923                "union split produced no finite component score at row {i}"
2924            ));
2925        }
2926        groups[best_j].push(i);
2927    }
2928    Ok(groups)
2929}
2930
2931/// Fit one structured-union candidate. The hard split identifies component
2932/// training groups; heterogeneous component roles are assigned by evaluating
2933/// every unique role permutation. Each assignment is then scored as one
2934/// normalized soft mixture on every training row.
2935///
2936/// Returns an error if no role assignment identifies every component. The
2937/// fixed-ladder caller propagates that failure for the whole declared family;
2938/// it never ranks a survivor subset.
2939pub fn fit_union_structure(
2940    data: ArrayView2<'_, f64>,
2941    structure: UnionStructure,
2942    config: GaussianMixtureConfig,
2943) -> Result<UnionStructureFit, String> {
2944    let fitted = fit_union_density(data, structure, config)?;
2945    Ok(UnionStructureFit {
2946        structure,
2947        components: fitted
2948            .components
2949            .iter()
2950            .map(UnionComponentDensity::summary)
2951            .collect(),
2952        log_likelihood: fitted.log_likelihood,
2953        bic: fitted.bic,
2954        total_parameters: fitted.total_parameters,
2955    })
2956}
2957
2958/// Fit the whole fixed union ladder and rank in-class by normalized soft-mixture
2959/// BIC (lower wins). The declared ladder is one selection family: if any
2960/// eligible structure fails, the whole comparison fails with every per-structure
2961/// error rather than silently selecting from an easier survivor set.
2962pub fn fit_union_ladder(
2963    data: ArrayView2<'_, f64>,
2964    config: GaussianMixtureConfig,
2965) -> Result<Vec<UnionStructureFit>, String> {
2966    let mut fits = Vec::new();
2967    let mut errors = Vec::new();
2968    for &structure in UNION_STRUCTURE_LADDER {
2969        match fit_union_structure(data, structure, config) {
2970            Ok(fit) => fits.push(fit),
2971            Err(e) => errors.push(format!("{}: {e}", structure.as_str())),
2972        }
2973    }
2974    if !errors.is_empty() {
2975        return Err(format!(
2976            "union ladder comparison failed; every declared structure must fit ({})",
2977            errors.join("; ")
2978        ));
2979    }
2980    if fits.is_empty() {
2981        return Err("union ladder is empty".to_string());
2982    }
2983    let ranked = rank_priority_candidates(
2984        fits.into_iter()
2985            .enumerate()
2986            .map(|(idx, row)| {
2987                let score = row.bic;
2988                let tie = row.total_parameters; // cheaper composite wins ties
2989                PriorityCandidate::new(row, idx, score, tie)
2990            })
2991            .collect(),
2992    )
2993    .into_iter()
2994    .map(|row| row.item)
2995    .collect::<Vec<_>>();
2996    Ok(ranked)
2997}
2998
2999fn gather_union_rows(data: ArrayView2<'_, f64>, idx: &[usize]) -> Array2<f64> {
3000    let d = data.ncols();
3001    let mut out = Array2::<f64>::zeros((idx.len(), d));
3002    for (r, &i) in idx.iter().enumerate() {
3003        for c in 0..d {
3004            out[[r, c]] = data[[i, c]];
3005        }
3006    }
3007    out
3008}
3009
3010/// Apply the structured-union admission policy to one circle group and return
3011/// its complete deterministic row set.
3012/// The `n > P` rule belongs to this composite ladder; it is not an intrinsic
3013/// restriction of circular-Gaussian likelihood fitting or BIC itself.
3014fn union_circle_rows(group: ArrayView2<'_, f64>) -> Result<Vec<usize>, String> {
3015    let minimum_rows = CircularGaussianFit2d::NUM_FREE_PARAMETERS + 1;
3016    if group.nrows() < minimum_rows {
3017        return Err(format!(
3018            "union circle component needs at least {minimum_rows} rows, got {}",
3019            group.nrows()
3020        ));
3021    }
3022    Ok((0..group.nrows()).collect())
3023}
3024
3025/// The shared fitted-density representation used by both in-sample BIC and
3026/// held-out predictive scoring. `Line` and `PointCluster` deliberately share
3027/// the Gaussian evaluator after fitting, but differ in covariance constraints
3028/// and parameter count.
3029#[derive(Debug, Clone)]
3030enum UnionDensityModel {
3031    Gaussian(GaussianComponentEval),
3032    Circle(CircularGaussianFit2d),
3033}
3034
3035#[derive(Debug, Clone)]
3036struct UnionComponentDensity {
3037    kind: UnionComponentKind,
3038    row_count: usize,
3039    num_parameters: usize,
3040    mixing_weight: f64,
3041    log_weight: f64,
3042    model: UnionDensityModel,
3043}
3044
3045impl UnionComponentDensity {
3046    fn summary(&self) -> UnionComponentFit {
3047        UnionComponentFit {
3048            kind: self.kind,
3049            row_count: self.row_count,
3050            num_parameters: self.num_parameters,
3051            mixing_weight: self.mixing_weight,
3052        }
3053    }
3054
3055    fn dimension(&self) -> usize {
3056        match &self.model {
3057            UnionDensityModel::Gaussian(eval) => eval.d,
3058            UnionDensityModel::Circle(_) => 2,
3059        }
3060    }
3061
3062    /// `log π_c + log p_c(y)` for one eval row.
3063    fn weighted_log_density(&self, y: ArrayView1<'_, f64>) -> f64 {
3064        let component_log_density = match &self.model {
3065            UnionDensityModel::Gaussian(eval) => eval.log_density(y),
3066            UnionDensityModel::Circle(fit) => fit.log_density(y[0], y[1]),
3067        };
3068        self.log_weight + component_log_density
3069    }
3070}
3071
3072#[derive(Debug, Clone)]
3073struct FittedUnionDensity {
3074    components: Vec<UnionComponentDensity>,
3075    log_likelihood: f64,
3076    bic: f64,
3077    total_parameters: usize,
3078}
3079
3080fn fit_union_density(
3081    train: ArrayView2<'_, f64>,
3082    structure: UnionStructure,
3083    config: GaussianMixtureConfig,
3084) -> Result<FittedUnionDensity, String> {
3085    let groups = union_responsibility_split(train, structure.num_components(), config)?;
3086    fit_union_density_from_groups(train, structure, &groups, config)
3087}
3088
3089/// Fit and score a union for an already-established hard partition. This seam
3090/// makes role assignment explicitly independent of the arbitrary component
3091/// labels emitted by the responsibility split.
3092fn fit_union_density_from_groups(
3093    train: ArrayView2<'_, f64>,
3094    structure: UnionStructure,
3095    groups: &[Vec<usize>],
3096    config: GaussianMixtureConfig,
3097) -> Result<FittedUnionDensity, String> {
3098    validate_union_partition(train.nrows(), structure.num_components(), groups)?;
3099    let assignments = unique_union_role_assignments(structure.components());
3100    let mut best: Option<FittedUnionDensity> = None;
3101    let mut errors = Vec::new();
3102
3103    for roles in assignments {
3104        let candidate = (|| {
3105            let mut components = Vec::with_capacity(groups.len());
3106            let n_train = train.nrows() as f64;
3107            for (&kind, rows) in roles.iter().zip(groups) {
3108                let group = gather_union_rows(train, rows);
3109                let mixing_weight = rows.len() as f64 / n_train;
3110                components.push(fit_union_component_density(
3111                    group.view(),
3112                    kind,
3113                    mixing_weight,
3114                    config,
3115                )?);
3116            }
3117
3118            let component_parameters = components.iter().try_fold(0usize, |sum, component| {
3119                sum.checked_add(component.num_parameters)
3120                    .ok_or_else(|| "union component parameter count overflowed usize".to_string())
3121            })?;
3122            let mixing_parameters = components.len() - 1;
3123            let total_parameters = component_parameters
3124                .checked_add(mixing_parameters)
3125                .ok_or_else(|| "union total parameter count overflowed usize".to_string())?;
3126            let per_point = score_union_components(&components, train)?;
3127            let log_likelihood = pairwise_sum(
3128                per_point
3129                    .as_slice()
3130                    .expect("owned union score vector must be contiguous"),
3131            );
3132            if !log_likelihood.is_finite() {
3133                return Err("union training log likelihood is non-finite".to_string());
3134            }
3135            let bic = -log_likelihood + 0.5 * total_parameters as f64 * (train.nrows() as f64).ln();
3136            if !bic.is_finite() {
3137                return Err("union normalized soft-mixture BIC is non-finite".to_string());
3138            }
3139            Ok(FittedUnionDensity {
3140                components,
3141                log_likelihood,
3142                bic,
3143                total_parameters,
3144            })
3145        })();
3146
3147        match candidate {
3148            Ok(candidate) => {
3149                let replace = match &best {
3150                    Some(current) => candidate.bic.total_cmp(&current.bic).is_lt(),
3151                    None => true,
3152                };
3153                // Unique assignments are generated in canonical order. Keeping
3154                // the earlier assignment on an exact score tie is deterministic.
3155                if replace {
3156                    best = Some(candidate);
3157                }
3158            }
3159            Err(error) => errors.push(format!("{roles:?}: {error}")),
3160        }
3161    }
3162
3163    best.ok_or_else(|| {
3164        format!(
3165            "union {} has no finite role assignment ({})",
3166            structure.as_str(),
3167            errors.join("; ")
3168        )
3169    })
3170}
3171
3172fn validate_union_partition(
3173    n_rows: usize,
3174    expected_groups: usize,
3175    groups: &[Vec<usize>],
3176) -> Result<(), String> {
3177    if n_rows == 0 {
3178        return Err("union fitting requires at least one training row".to_string());
3179    }
3180    if groups.len() != expected_groups {
3181        return Err(format!(
3182            "union partition has {} groups, expected {expected_groups}",
3183            groups.len()
3184        ));
3185    }
3186    let mut seen = vec![false; n_rows];
3187    for (group_index, rows) in groups.iter().enumerate() {
3188        if rows.is_empty() {
3189            return Err(format!("union partition group {group_index} is empty"));
3190        }
3191        for &row in rows {
3192            if row >= n_rows {
3193                return Err(format!(
3194                    "union partition group {group_index} contains out-of-range row {row} for {n_rows} rows"
3195                ));
3196            }
3197            if std::mem::replace(&mut seen[row], true) {
3198                return Err(format!("union partition contains duplicate row {row}"));
3199            }
3200        }
3201    }
3202    if let Some(missing) = seen.iter().position(|included| !included) {
3203        return Err(format!("union partition omits row {missing}"));
3204    }
3205    Ok(())
3206}
3207
3208fn unique_union_role_assignments(roles: &[UnionComponentKind]) -> Vec<Vec<UnionComponentKind>> {
3209    fn visit(
3210        roles: &[UnionComponentKind],
3211        used: &mut [bool],
3212        assignment: &mut Vec<UnionComponentKind>,
3213        out: &mut Vec<Vec<UnionComponentKind>>,
3214    ) {
3215        if assignment.len() == roles.len() {
3216            out.push(assignment.clone());
3217            return;
3218        }
3219        let mut used_at_depth = Vec::new();
3220        for (index, &role) in roles.iter().enumerate() {
3221            if used[index] || used_at_depth.contains(&role) {
3222                continue;
3223            }
3224            used_at_depth.push(role);
3225            used[index] = true;
3226            assignment.push(role);
3227            visit(roles, used, assignment, out);
3228            assignment.pop();
3229            used[index] = false;
3230        }
3231    }
3232
3233    let mut out = Vec::new();
3234    visit(
3235        roles,
3236        &mut vec![false; roles.len()],
3237        &mut Vec::with_capacity(roles.len()),
3238        &mut out,
3239    );
3240    out
3241}
3242
3243fn fit_union_component_density(
3244    group: ArrayView2<'_, f64>,
3245    kind: UnionComponentKind,
3246    mixing_weight: f64,
3247    config: GaussianMixtureConfig,
3248) -> Result<UnionComponentDensity, String> {
3249    if !(mixing_weight.is_finite() && mixing_weight > 0.0 && mixing_weight <= 1.0) {
3250        return Err(format!(
3251            "union component mixing weight must be finite and in (0, 1], got {mixing_weight}"
3252        ));
3253    }
3254    let row_count = group.nrows();
3255    let (model, num_parameters) = match kind {
3256        UnionComponentKind::Line => {
3257            if group.nrows() < group.ncols() + 1 {
3258                return Err(format!(
3259                    "union line component needs >= {} rows, got {}",
3260                    group.ncols() + 1,
3261                    group.nrows()
3262                ));
3263            }
3264            let fit = fit_gaussian_mixture(group, 1, config).map_err(|error| error.to_string())?;
3265            let num_parameters = fit.num_free_parameters();
3266            let eval = GaussianComponentEval::factor(fit.means.row(0), &fit.covariances[0])?;
3267            (UnionDensityModel::Gaussian(eval), num_parameters)
3268        }
3269        UnionComponentKind::PointCluster => {
3270            if group.nrows() < group.ncols() + 1 {
3271                return Err(format!(
3272                    "union isotropic point component needs >= {} rows, got {}",
3273                    group.ncols() + 1,
3274                    group.nrows()
3275                ));
3276            }
3277            let eval = fit_isotropic_gaussian_component(group, config.covariance_floor)?;
3278            (
3279                UnionDensityModel::Gaussian(eval),
3280                group
3281                    .ncols()
3282                    .checked_add(1)
3283                    .ok_or_else(|| "union point parameter count overflowed usize".to_string())?,
3284            )
3285        }
3286        UnionComponentKind::Circle => {
3287            let rows = union_circle_rows(group)?;
3288            let fit = CircularGaussianFit2d::fit(group, &rows)?;
3289            (
3290                UnionDensityModel::Circle(fit),
3291                CircularGaussianFit2d::NUM_FREE_PARAMETERS,
3292            )
3293        }
3294    };
3295    Ok(UnionComponentDensity {
3296        kind,
3297        row_count,
3298        num_parameters,
3299        mixing_weight,
3300        log_weight: mixing_weight.ln(),
3301        model,
3302    })
3303}
3304
3305#[derive(Debug, Clone, Copy)]
3306struct StableScalarMeanChart {
3307    origin: f64,
3308    scale: f64,
3309    normalized_offset: f64,
3310}
3311
3312impl StableScalarMeanChart {
3313    #[inline]
3314    fn centered(self, value: f64) -> Result<f64, String> {
3315        let relative = value - self.origin;
3316        let centered = (-self.normalized_offset).mul_add(self.scale, relative);
3317        if centered.is_finite() {
3318            Ok(centered)
3319        } else {
3320            Err("union isotropic point residual is not representable".to_string())
3321        }
3322    }
3323}
3324
3325/// Range-safe and translation-accurate scalar mean chart. The normalized mean
3326/// offset is retained separately from the rounded absolute mean so residuals
3327/// use `(x-origin)-offset` rather than losing a fractional offset at a large
3328/// common translation. FMA also lets a subnormal offset affect the correctly
3329/// rounded absolute mean without first rounding that offset to zero.
3330fn stable_scalar_mean_chart(values: ArrayView1<'_, f64>) -> Result<StableScalarMeanChart, String> {
3331    if values.is_empty() || values.iter().any(|value| !value.is_finite()) {
3332        return Err("stable scalar mean requires finite nonempty values".to_string());
3333    }
3334    let anchor = values[0];
3335    let anchor_chart_is_representable = values.iter().all(|&value| (value - anchor).is_finite());
3336    let origin = if anchor_chart_is_representable {
3337        anchor
3338    } else {
3339        0.0
3340    };
3341    let scale = values
3342        .iter()
3343        .map(|&value| (value - origin).abs())
3344        .fold(0.0_f64, f64::max);
3345    if scale == 0.0 {
3346        return Ok(StableScalarMeanChart {
3347            origin,
3348            scale: 0.0,
3349            normalized_offset: 0.0,
3350        });
3351    }
3352    let normalized = values
3353        .iter()
3354        .map(|&value| (value - origin) / scale)
3355        .collect::<Vec<_>>();
3356    let normalized_offset = pairwise_sum(&normalized) / values.len() as f64;
3357    let mean = normalized_offset.mul_add(scale, origin);
3358    if !(normalized_offset.is_finite() && mean.is_finite()) {
3359        return Err("union isotropic point mean is not representable".to_string());
3360    }
3361    Ok(StableScalarMeanChart {
3362        origin,
3363        scale,
3364        normalized_offset,
3365    })
3366}
3367
3368/// Maximum-likelihood isotropic Gaussian fit in stable per-column mean charts.
3369fn fit_isotropic_gaussian_component(
3370    group: ArrayView2<'_, f64>,
3371    covariance_floor: f64,
3372) -> Result<GaussianComponentEval, String> {
3373    let n = group.nrows();
3374    let d = group.ncols();
3375    if n == 0 || d == 0 {
3376        return Err("union isotropic point component requires a non-empty matrix".to_string());
3377    }
3378    if !(covariance_floor.is_finite() && covariance_floor > 0.0) {
3379        return Err(format!(
3380            "union isotropic covariance floor must be finite and positive, got {covariance_floor}"
3381        ));
3382    }
3383
3384    for row in group.rows() {
3385        for axis in 0..d {
3386            let value = row[axis];
3387            if !value.is_finite() {
3388                return Err(format!(
3389                    "union isotropic point data contains non-finite coordinate {value}"
3390                ));
3391            }
3392        }
3393    }
3394
3395    let mut charts = Vec::with_capacity(d);
3396    for axis in 0..d {
3397        let chart = stable_scalar_mean_chart(group.column(axis))?;
3398        charts.push(chart);
3399    }
3400
3401    let scalar_count = n
3402        .checked_mul(d)
3403        .ok_or_else(|| "union isotropic residual count overflowed usize".to_string())?;
3404    let mut residuals = Vec::with_capacity(scalar_count);
3405    let mut residual_scale = 0.0_f64;
3406    for row in group.rows() {
3407        for axis in 0..d {
3408            let residual = charts[axis].centered(row[axis])?;
3409            residual_scale = residual_scale.max(residual.abs());
3410            residuals.push(residual);
3411        }
3412    }
3413    let variance = if residual_scale == 0.0 {
3414        covariance_floor
3415    } else {
3416        for residual in &mut residuals {
3417            *residual = (*residual / residual_scale).powi(2);
3418        }
3419        let normalized_mean_square = pairwise_sum(&residuals) / scalar_count as f64;
3420        let rms = residual_scale * normalized_mean_square.sqrt();
3421        let unconstrained = rms * rms;
3422        if !unconstrained.is_finite() {
3423            return Err("union isotropic point variance is non-finite".to_string());
3424        }
3425        unconstrained.max(covariance_floor)
3426    };
3427    GaussianComponentEval::isotropic(&charts, variance)
3428}
3429
3430fn score_union_components(
3431    components: &[UnionComponentDensity],
3432    eval: ArrayView2<'_, f64>,
3433) -> Result<Array1<f64>, String> {
3434    if components.is_empty() {
3435        return Err("union density requires at least one component".to_string());
3436    }
3437    if eval.iter().any(|coordinate| !coordinate.is_finite()) {
3438        return Err("union eval coordinates must be finite".to_string());
3439    }
3440    for component in components {
3441        if component.dimension() != eval.ncols() {
3442            return Err(format!(
3443                "union component {:?} has dimension {}, eval has {} columns",
3444                component.kind,
3445                component.dimension(),
3446                eval.ncols()
3447            ));
3448        }
3449    }
3450    let mut out = Array1::<f64>::zeros(eval.nrows());
3451    let mut terms = vec![f64::NEG_INFINITY; components.len()];
3452    for i in 0..eval.nrows() {
3453        let row = eval.row(i);
3454        let mut max_term = f64::NEG_INFINITY;
3455        for (component_index, component) in components.iter().enumerate() {
3456            let term = component.weighted_log_density(row);
3457            terms[component_index] = term;
3458            if term > max_term {
3459                max_term = term;
3460            }
3461        }
3462        let value = log_sum_exp(&terms, max_term);
3463        if !value.is_finite() {
3464            return Err(format!(
3465                "union density produced non-finite log density at eval row {i}"
3466            ));
3467        }
3468        out[i] = value;
3469    }
3470    Ok(out)
3471}
3472
3473/// Per-point held-out log predictive density of a structured-union candidate:
3474/// fit the component densities on `train` and score each row of `eval` as the
3475/// soft mixture `log Σ_c π_c p_c(y)`. This is the cross-class stacking column
3476/// source for a union (the analogue of [`GaussianMixtureFit::per_point_log_density`]).
3477pub fn union_per_point_log_density(
3478    train: ArrayView2<'_, f64>,
3479    eval: ArrayView2<'_, f64>,
3480    structure: UnionStructure,
3481    config: GaussianMixtureConfig,
3482) -> Result<Array1<f64>, String> {
3483    if train.ncols() != eval.ncols() {
3484        return Err(format!(
3485            "union held-out density: train has {} columns, eval has {}",
3486            train.ncols(),
3487            eval.ncols()
3488        ));
3489    }
3490    let fitted = fit_union_density(train, structure, config)?;
3491    score_union_components(&fitted.components, eval)
3492}
3493
3494/// One fitted model in a REML/LAML evidence comparison.
3495#[derive(Clone, Debug)]
3496pub struct RemlCandidate {
3497    pub index: usize,
3498    pub name: String,
3499    /// Minimised REML/LAML cost. Lower is better. This is the model's reported
3500    /// evidence headline (`Model.evidence`), kept verbatim in the score table.
3501    pub score: f64,
3502    pub edf: Option<f64>,
3503    /// Log-likelihood at the converged mode, on the engine's
3504    /// constants-omitted scale (same as `gam_inference::model_comparison`).
3505    /// Present when the fit carries it; `None` for legacy payloads.
3506    pub log_lik: Option<f64>,
3507    /// Response-family tag (e.g. "gaussian", "gamma", "binomial"). Carried so
3508    /// `compare_reml_fits` can REFUSE to rank fits whose REML/LAML scores are on
3509    /// incomparable base measures (a cross-family comparison is meaningless;
3510    /// #1384). `None` for legacy payloads that did not record it — those are not
3511    /// guarded (back-compatible), but every current FFI candidate carries it.
3512    pub family: Option<String>,
3513    /// Number of observations the fit was trained on. Carried so
3514    /// `compare_reml_fits` can REFUSE to rank fits made on a different number of
3515    /// observations (hence different data): `−2·loglik` and the REML/LAML
3516    /// evidence grow with `n`, so a score difference between two fits with
3517    /// different `n` is not a Bayes factor — the same incomparability the family
3518    /// guard already rejects. `None` for payloads that do not record it (legacy /
3519    /// O(n) scan smoothers), which the guard treats as unconstrained.
3520    pub n_obs: Option<usize>,
3521}
3522
3523impl RemlCandidate {
3524    /// Cost used to RANK candidates and pick the winner.
3525    ///
3526    /// The REML/LAML marginal-likelihood evidence headline (`score`) does NOT
3527    /// reliably Occam-penalise an added pure-noise smooth: on `y ~ s(x)` vs
3528    /// `y ~ s(x) + s(z)` with `z ⟂ y`, the augmented model's evidence is
3529    /// *lower* (apparently better) by a few nats on essentially every dataset,
3530    /// because the Gaussian REML Occam pair `½(log|H| − log|S|₊)` collapses
3531    /// toward zero for a finite-`λ̂` null term while that term still spends a
3532    /// few effective degrees of freedom fitting noise (issue #1362).
3533    ///
3534    /// The conditional AIC `−2ℓ + 2·edf` prices exactly those spent degrees of
3535    /// freedom and discriminates correctly: it penalises the noise smooth
3536    /// (Δ ≈ +15 nats) yet rewards a genuinely relevant smooth (Δ ≈ −650),
3537    /// preserving power. We therefore rank on the conditional AIC whenever both
3538    /// the log-likelihood and the effective degrees of freedom are available,
3539    /// and fall back to the raw evidence headline otherwise. The reported
3540    /// `score_table` still carries the unaltered evidence (`reml_score`), so
3541    /// `Model.evidence` / `bayes_factor_vs` stay consistent with the table.
3542    pub fn ranking_score(&self) -> f64 {
3543        match (self.log_lik, self.edf) {
3544            (Some(log_lik), Some(edf)) if log_lik.is_finite() && edf.is_finite() => {
3545                -2.0 * log_lik + 2.0 * edf
3546            }
3547            _ => self.score,
3548        }
3549    }
3550}
3551
3552#[derive(Clone, Debug)]
3553pub struct RemlComparison {
3554    pub ranking: Vec<RankedRow>,
3555    pub winner: String,
3556    pub evidence_summary: String,
3557    pub score_table: Vec<ScoreRow>,
3558}
3559
3560#[derive(Clone, Debug)]
3561pub struct RankedRow {
3562    pub name: String,
3563    pub score: f64,
3564    /// Cost gap from the winning model on the SAME scale used to order the
3565    /// ranking (`ranking_score`, the Occam-penalised conditional AIC where
3566    /// available, issue #1362). The winner is `argmin ranking_score`, so this
3567    /// is `>= 0` for every row by construction — it never contradicts the
3568    /// declared winner (issue #1465). `score` still carries the raw REML/LAML
3569    /// evidence so it stays consistent with `Model.evidence`.
3570    pub delta: f64,
3571    /// Akaike evidence ratio of the winner over this row on the ranking scale.
3572    /// `delta` is a conditional-AIC gap (a −2·log / deviance-scale quantity), so
3573    /// the evidence ratio is `exp(½·delta) >= 1` (Burnham & Anderson), NOT
3574    /// `exp(delta)` — the latter squares the intended ratio (issues #1465, #2124).
3575    pub bayes_factor: f64,
3576    pub edf: Option<f64>,
3577}
3578
3579#[derive(Clone, Debug)]
3580pub struct ScoreRow {
3581    pub name: String,
3582    pub reml_score: f64,
3583    pub delta_reml: f64,
3584    pub bayes_factor_best_over_model: f64,
3585    pub effective_dof: Option<f64>,
3586}
3587
3588/// Log Bayes factor of model `a` over model `b` from minimised REML/LAML costs.
3589#[inline]
3590pub fn log_bayes_factor(reml_score_a: f64, reml_score_b: f64) -> f64 {
3591    reml_score_b - reml_score_a
3592}
3593
3594/// Compare fitted models by the single evidence ordering contract used by
3595/// topology ranking and seed screening: lower finite cost wins, with stable
3596/// original-order tie handling.
3597pub fn compare_reml_fits(mut candidates: Vec<RemlCandidate>) -> Result<RemlComparison, String> {
3598    if candidates.is_empty() {
3599        return Err("compare_models requires at least one fit".to_string());
3600    }
3601    // Fail-loud comparability guard (#1384): REML/LAML evidence scores are only
3602    // comparable across fits of the SAME response family — a Gaussian score and
3603    // a Gamma score live on different log-density base measures, so their
3604    // difference is not a Bayes factor. Ranking them anyway returns a confident
3605    // but meaningless winner. Refuse when two candidates carry DIFFERENT family
3606    // tags. Candidates with no family tag (`None`, legacy payloads) are not
3607    // constrained, so this never spuriously rejects an older saved model.
3608    {
3609        let mut seen_family: Option<&str> = None;
3610        for cand in &candidates {
3611            if let Some(fam) = cand.family.as_deref() {
3612                match seen_family {
3613                    None => seen_family = Some(fam),
3614                    Some(prev) if prev != fam => {
3615                        return Err(format!(
3616                            "compare_models: cannot compare fits of different response families                              ('{prev}' vs '{fam}'); their REML/LAML evidence scores are on                              incomparable base measures. Compare models fit to the same response                              under the same family."
3617                        ));
3618                    }
3619                    Some(_) => {}
3620                }
3621            }
3622        }
3623    }
3624    // Fail-loud comparability guard (#1384 sibling): AIC / REML-LAML evidence are
3625    // only comparable across fits of the SAME response on the SAME observations.
3626    // `−2·loglik` (and the marginal-likelihood headline) grow with the number of
3627    // observations `n`, so two fits with different `n` live on incomparable
3628    // scales and their score gap is not a Bayes factor — comparing an n=500 and
3629    // an n=100 fit of the same DGP otherwise declares the n=100 model the winner
3630    // purely because fewer points give a less-negative total log-likelihood.
3631    // Refuse when two candidates carry DIFFERENT observation counts. Candidates
3632    // with no count (`None`, legacy / O(n) scan payloads) are unconstrained, so
3633    // this never spuriously rejects a fit that simply did not record `n`.
3634    {
3635        let mut seen_n: Option<usize> = None;
3636        for cand in &candidates {
3637            if let Some(n) = cand.n_obs {
3638                match seen_n {
3639                    None => seen_n = Some(n),
3640                    Some(prev) if prev != n => {
3641                        return Err(format!(
3642                            "compare_models: cannot compare fits made on a different number of \
3643                             observations (n={prev} vs n={n}); AIC / REML-LAML evidence scales \
3644                             with the sample size, so their score difference is not a Bayes \
3645                             factor. Compare models fit to the same response on the same data."
3646                        ));
3647                    }
3648                    Some(_) => {}
3649                }
3650            }
3651        }
3652    }
3653    candidates = rank_priority_candidates(
3654        candidates
3655            .into_iter()
3656            .enumerate()
3657            .map(|(idx, row)| {
3658                // Rank/winner on the Occam-penalised conditional AIC where it is
3659                // available (issue #1362); falls back to the raw evidence score.
3660                let ranking = row.ranking_score();
3661                PriorityCandidate::new(row, idx, ranking, 0)
3662            })
3663            .collect(),
3664    )
3665    .into_iter()
3666    .map(|row| row.item)
3667    .collect();
3668
3669    let winner = candidates[0].name.clone();
3670    // The ranking `delta` / `bayes_factor` must be measured on the SAME scale
3671    // that orders the table — the `ranking_score` (Occam-penalised conditional
3672    // AIC where available, issue #1362). `candidates[0]` is the winner =
3673    // `argmin ranking_score`, so its ranking score IS the minimum; every row's
3674    // ranking-scale gap is then `>= 0` and its Bayes factor `>= 1`, never
3675    // contradicting the declared winner (issue #1465). Computing these against
3676    // the AIC winner's *raw REML* — which is not the minimum raw REML once AIC
3677    // and REML disagree — produced negative deltas and Bayes factors < 1 for
3678    // non-winner rows.
3679    let best_ranking_score = candidates[0].ranking_score();
3680    // The raw-REML `score_table` stays on the raw evidence scale (consistent
3681    // with `Model.evidence` / `bayes_factor_vs`), but is referenced to the
3682    // genuine minimum raw REML so its best-over-model Bayes factors are also
3683    // coherent (`>= 1`), rather than to whichever row happens to sit at index 0.
3684    let best_raw_score = candidates
3685        .iter()
3686        .map(|c| c.score)
3687        .fold(f64::INFINITY, f64::min);
3688    let mut ranking = Vec::with_capacity(candidates.len());
3689    let mut score_table = Vec::with_capacity(candidates.len());
3690    for row in &candidates {
3691        let delta = log_bayes_factor(best_ranking_score, row.ranking_score());
3692        // `ranking_score` is the conditional AIC (`−2·loglik + 2·edf`), a −2·log /
3693        // deviance-scale cost, so `delta` is a full ΔAIC gap. The Akaike evidence
3694        // ratio for an AIC gap Δ is `exp(−½Δ)` (Burnham & Anderson evidence ratio),
3695        // hence the winner-over-row Bayes factor is `exp(½·delta)`. Reporting
3696        // `delta.exp()` squared the intended ratio (issue #2124). `delta` itself is
3697        // left on the AIC scale on purpose — only its exp() conversion is halved.
3698        let bayes_factor = (0.5 * delta).exp();
3699        let delta_reml = log_bayes_factor(best_raw_score, row.score);
3700        ranking.push(RankedRow {
3701            name: row.name.clone(),
3702            score: row.score,
3703            delta,
3704            bayes_factor,
3705            edf: row.edf,
3706        });
3707        score_table.push(ScoreRow {
3708            name: row.name.clone(),
3709            reml_score: row.score,
3710            delta_reml,
3711            bayes_factor_best_over_model: delta_reml.exp(),
3712            effective_dof: row.edf,
3713        });
3714    }
3715    // The winner is decided by `ranking_score` (the Occam-penalised conditional
3716    // AIC where available, issue #1362), which can disagree in sign with the raw
3717    // evidence Bayes factor for a noise-augmented model. Summarise the actual
3718    // decision margin so the headline never contradicts the chosen winner.
3719    let evidence_summary = if let Some(runner_up) = candidates.get(1) {
3720        let margin = runner_up.ranking_score() - candidates[0].ranking_score();
3721        // `margin` is a conditional-AIC gap (−2·log scale), so the Akaike evidence
3722        // ratio is `exp(−½·margin)`; `format_bayes_factor` formats `exp()` of its
3723        // argument, so pass the halved margin to headline `exp(½·margin)` rather
3724        // than the squared `exp(margin)` (issue #2124).
3725        format!(
3726            "{} wins by Bayes factor {} over {}",
3727            winner,
3728            format_bayes_factor(0.5 * margin),
3729            runner_up.name
3730        )
3731    } else {
3732        format!("{winner} (single fit; no comparison)")
3733    };
3734    Ok(RemlComparison {
3735        ranking,
3736        winner,
3737        evidence_summary,
3738        score_table,
3739    })
3740}
3741
3742pub fn format_bayes_factor(log_bf: f64) -> String {
3743    if !log_bf.is_finite() {
3744        return "inf".to_string();
3745    }
3746    if log_bf.abs() >= std::f64::consts::LN_10 * 3.0 {
3747        return format!("1e{:+.1}", log_bf / std::f64::consts::LN_10);
3748    }
3749    format_three_significant(log_bf.exp())
3750}
3751
3752pub fn format_three_significant(value: f64) -> String {
3753    if value == 0.0 {
3754        return "0".to_string();
3755    }
3756    if !value.is_finite() {
3757        return format!("{value}");
3758    }
3759    let exponent = value.abs().log10().floor() as i32;
3760    if exponent >= 3 {
3761        return format!("{value:.2e}");
3762    }
3763    let decimals = (2 - exponent).max(0) as usize;
3764    let scale = 10f64.powi(decimals as i32);
3765    let rounded = (value * scale).abs().round() / scale * value.signum();
3766    format!("{rounded:.decimals$}")
3767}
3768
3769impl Default for TopologySelectOptions {
3770    fn default() -> Self {
3771        Self {
3772            tie_tolerance: 1e-3,
3773            score_scale: TopologyScoreScale::PerObservation,
3774        }
3775    }
3776}
3777
3778// ---------------------------------------------------------------------------
3779// Laplace evidence
3780// ---------------------------------------------------------------------------
3781
3782/// Single canonical Laplace evidence at the inner-loop fixed point.
3783///
3784/// Returns negative log evidence:
3785///
3786/// ```text
3787/// V(ρ, T) = F(β*, u*; ρ, T)
3788///         + 0.5 log|H|
3789///         - 0.5 log|S_pen(ρ)|+
3790///         - 0.5 (dim(H) - rank(S_pen)) log(2π).
3791/// ```
3792///
3793/// The last term is the rank-aware Tierney-Kadane normalizer:
3794/// `log p(y|T) ≈ -V`, with `0.5 log|2πH⁻¹| - 0.5 log|2πS⁻¹|`.
3795///
3796/// The `H` log-determinant is computed from the arrow factorization
3797///
3798/// ```text
3799/// log|H| = Σ_i log|H_uu_i| + log|A|
3800/// ```
3801///
3802/// (proposal §3.4 / §7) using the **undamped** per-row Cholesky factors
3803/// `cache.htt_factors_undamped` and the **undamped** Schur factor.
3804///
3805/// `penalty_log_det` is `log|S_pen(ρ)|+` — the prior penalty
3806/// pseudo-logdet from `crate::reml::penalty_logdet` (proposal
3807/// §3.6). It must NOT be confused with the arrow Schur log-det, which
3808/// this function recomputes internally from `logdet_source`.
3809///
3810/// `residual_objective` is `F(β*, u*; ρ, T)` at the inner optimum. The
3811/// envelope theorem (proposal §3.2) makes this the only `F`-related
3812/// contribution.
3813///
3814/// `effective_dim` is `dim(H)` after constraints/projections and
3815/// `penalty_rank` is `rank(S_pen)`. Their difference is the unpenalized
3816/// nullspace dimension that remains in the Laplace integral.
3817///
3818/// # Errors
3819///
3820/// Returns `f64::NAN` if the exact factor path is incoherent and no HVP
3821/// fallback is supplied, or if the supplied dimensions are non-finite.
3822pub fn laplace_evidence(
3823    logdet_source: EvidenceLogDetSource<'_>,
3824    penalty_log_det: f64,
3825    residual_objective: f64,
3826    effective_dim: f64,
3827    penalty_rank: f64,
3828) -> f64 {
3829    if !(effective_dim.is_finite() && penalty_rank.is_finite()) {
3830        return f64::NAN;
3831    }
3832    let log_det_h = match evidence_hessian_log_det(logdet_source) {
3833        Ok(v) => v,
3834        Err(_) => return f64::NAN,
3835    };
3836    let null_dim = effective_dim - penalty_rank;
3837    if !null_dim.is_finite() || null_dim < -1e-9 {
3838        return f64::NAN;
3839    }
3840    residual_objective + 0.5 * log_det_h
3841        - 0.5 * penalty_log_det
3842        - 0.5 * null_dim.max(0.0) * (2.0 * std::f64::consts::PI).ln()
3843}
3844
3845/// Compute the Hessian logdet from exact arrow factors or an HVP fallback.
3846pub fn evidence_hessian_log_det(source: EvidenceLogDetSource<'_>) -> Result<f64, String> {
3847    match source {
3848        EvidenceLogDetSource::FactoredArrow {
3849            cache,
3850            fallback_hvp,
3851        } => match arrow_log_det_from_cache(cache) {
3852            Some(v) => Ok(v),
3853            None => match fallback_hvp {
3854                Some(hvp) => hessian_log_det_from_hvp(hvp),
3855                None => {
3856                    Err("evidence Hessian logdet requires exact factors or HVP fallback".into())
3857                }
3858            },
3859        },
3860        EvidenceLogDetSource::Hvp(hvp) => hessian_log_det_from_hvp(hvp),
3861    }
3862}
3863
3864/// Log determinant of an SPD operator supplied by HVP callback.
3865///
3866/// The dispatch boundary intentionally matches
3867/// `ANALYTIC_LOGDET_DENSE_DIM_THRESHOLD` in `terms::analytic_penalties`:
3868/// small operators are materialized and diagonalized exactly; larger ones use
3869/// Rademacher stochastic Lanczos quadrature.
3870pub fn hessian_log_det_from_hvp(hvp: EvidenceHvpLogDet<'_>) -> Result<f64, String> {
3871    if hvp.dim == 0 {
3872        return Ok(0.0);
3873    }
3874    if hvp.dim <= ANALYTIC_LOGDET_DENSE_DIM_THRESHOLD {
3875        let mut dense = Array2::<f64>::zeros((hvp.dim, hvp.dim));
3876        let mut basis = vec![0.0_f64; hvp.dim];
3877        for j in 0..hvp.dim {
3878            basis[j] = 1.0;
3879            let col = (hvp.apply)(&basis);
3880            basis[j] = 0.0;
3881            if col.len() != hvp.dim || col.iter().any(|v| !v.is_finite()) {
3882                return Err(format!(
3883                    "evidence HVP logdet expected finite column of length {}, got {}",
3884                    hvp.dim,
3885                    col.len()
3886                ));
3887            }
3888            for i in 0..hvp.dim {
3889                dense[[i, j]] = col[i];
3890            }
3891        }
3892        validate_dense_hvp_symmetry(&dense)?;
3893        for i in 0..hvp.dim {
3894            for j in (i + 1)..hvp.dim {
3895                let avg = 0.5 * (dense[[i, j]] + dense[[j, i]]);
3896                dense[[i, j]] = avg;
3897                dense[[j, i]] = avg;
3898            }
3899        }
3900        dense_spd_log_det(&dense)
3901    } else {
3902        stochastic_hvp_log_det(hvp)
3903    }
3904}
3905
3906fn dense_spd_log_det(matrix: &Array2<f64>) -> Result<f64, String> {
3907    if matrix.nrows() != matrix.ncols() {
3908        return Err(format!(
3909            "evidence dense logdet requires square matrix, got {}x{}",
3910            matrix.nrows(),
3911            matrix.ncols()
3912        ));
3913    }
3914    if gam_gpu::cuda_selected().map_err(|error| error.to_string())? {
3915        return crate::gpu::reml_gpu::evidence_derivatives_gpu(
3916            crate::gpu::reml_gpu::RemlGpuInput {
3917                penalized_hessian: matrix.view(),
3918                derivative_hessians: Vec::new(),
3919            },
3920        )
3921        .map(|evidence| evidence.logdet_hessian);
3922    }
3923    let (evals, _) = matrix
3924        .eigh(Side::Lower)
3925        .map_err(|e| format!("evidence dense logdet eigendecomposition failed: {e}"))?;
3926    let mut logdet = 0.0_f64;
3927    for (idx, &ev) in evals.iter().enumerate() {
3928        if !ev.is_finite() || ev <= 0.0 {
3929            return Err(format!(
3930                "evidence dense logdet expected SPD Hessian, eigenvalue {idx} is {ev:.3e}"
3931            ));
3932        }
3933        logdet += ev.ln();
3934    }
3935    Ok(logdet)
3936}
3937
3938fn validate_dense_hvp_symmetry(matrix: &Array2<f64>) -> Result<(), String> {
3939    let n = matrix.nrows();
3940    let mut norm_sq = 0.0_f64;
3941    for &value in matrix.iter() {
3942        norm_sq += value * value;
3943    }
3944
3945    let mut skew_sq = 0.0_f64;
3946    for i in 0..n {
3947        for j in (i + 1)..n {
3948            let skew = matrix[[i, j]] - matrix[[j, i]];
3949            skew_sq += 2.0 * skew * skew;
3950        }
3951    }
3952
3953    let rel_skew = skew_sq.sqrt() / norm_sq.sqrt().max(1.0);
3954    if !rel_skew.is_finite() || rel_skew > EVIDENCE_HVP_SYMMETRY_REL_TOL {
3955        return Err(format!(
3956            "evidence HVP logdet requires symmetric operator, relative skew norm is {rel_skew:.3e}"
3957        ));
3958    }
3959    Ok(())
3960}
3961
3962fn validate_hvp_randomized_symmetry(hvp: EvidenceHvpLogDet<'_>) -> Result<(), String> {
3963    let inv_norm = 1.0 / (hvp.dim as f64).sqrt();
3964    for probe in 0..EVIDENCE_HVP_SYMMETRY_PROBES.max(1) {
3965        let mut x = vec![0.0_f64; hvp.dim];
3966        let mut y = vec![0.0_f64; hvp.dim];
3967        rademacher_unit_probe_into_slice(&mut x, (2 * probe) as u64, inv_norm);
3968        rademacher_unit_probe_into_slice(&mut y, (2 * probe + 1) as u64, inv_norm);
3969
3970        let hx = (hvp.apply)(&x);
3971        let hy = (hvp.apply)(&y);
3972        if hx.len() != hvp.dim || hx.iter().any(|v| !v.is_finite()) {
3973            return Err(format!(
3974                "evidence HVP symmetry check expected finite vector of length {}, got {}",
3975                hvp.dim,
3976                hx.len()
3977            ));
3978        }
3979        if hy.len() != hvp.dim || hy.iter().any(|v| !v.is_finite()) {
3980            return Err(format!(
3981                "evidence HVP symmetry check expected finite vector of length {}, got {}",
3982                hvp.dim,
3983                hy.len()
3984            ));
3985        }
3986
3987        let lhs = dot_slice(&x, &hy);
3988        let rhs = dot_slice(&hx, &y);
3989        let scale = (norm2_slice(&hx) * norm2_slice(&y))
3990            .max(norm2_slice(&hy) * norm2_slice(&x))
3991            .max(lhs.abs())
3992            .max(rhs.abs())
3993            .max(1.0);
3994        let rel = (lhs - rhs).abs() / scale;
3995        if !rel.is_finite() || rel > EVIDENCE_HVP_SYMMETRY_REL_TOL {
3996            return Err(format!(
3997                "evidence HVP logdet requires symmetric operator, randomized symmetry probe {probe} has relative bilinear mismatch {rel:.3e}"
3998            ));
3999        }
4000    }
4001    Ok(())
4002}
4003
4004fn stochastic_hvp_log_det(hvp: EvidenceHvpLogDet<'_>) -> Result<f64, String> {
4005    validate_hvp_randomized_symmetry(hvp)?;
4006    let probes = EVIDENCE_LOGDET_SLQ_PROBES.max(1);
4007    let steps = EVIDENCE_LOGDET_LANCZOS_STEPS.min(hvp.dim).max(1);
4008    let inv_norm = 1.0 / (hvp.dim as f64).sqrt();
4009    let mut estimate = 0.0_f64;
4010    for probe in 0..probes {
4011        let mut q0 = vec![0.0_f64; hvp.dim];
4012        rademacher_unit_probe_into_slice(&mut q0, probe as u64, inv_norm);
4013        let quad = lanczos_log_quadrature_hvp(hvp, q0, steps)?;
4014        estimate += hvp.dim as f64 * quad;
4015    }
4016    Ok(estimate / probes as f64)
4017}
4018
4019fn lanczos_log_quadrature_hvp(
4020    hvp: EvidenceHvpLogDet<'_>,
4021    q: Vec<f64>,
4022    max_steps: usize,
4023) -> Result<f64, String> {
4024    let n = hvp.dim;
4025    let eigen = symmetric_lanczos_eigenpairs(
4026        n,
4027        &q,
4028        SymmetricLanczosOptions {
4029            max_steps,
4030            residual_tol: 1e-12,
4031            local_reorthogonalize: false,
4032            full_reorthogonalize: false,
4033        },
4034        |q, out| {
4035            let applied = (hvp.apply)(q);
4036            if applied.len() != n || applied.iter().any(|v| !v.is_finite()) {
4037                return Err(format!(
4038                    "evidence HVP SLQ expected finite vector of length {n}, got {}",
4039                    applied.len()
4040                ));
4041            }
4042            out.copy_from_slice(&applied);
4043            Ok(())
4044        },
4045    )
4046    .map_err(|e| format!("evidence HVP SLQ Lanczos failed: {e}"))?;
4047    symmetric_lanczos_log_quadrature(&eigen, "evidence HVP SLQ expected SPD Hessian")
4048}
4049
4050#[inline]
4051fn dot_slice(a: &[f64], b: &[f64]) -> f64 {
4052    assert_eq!(a.len(), b.len());
4053    let mut s = 0.0_f64;
4054    for i in 0..a.len() {
4055        s += a[i] * b[i];
4056    }
4057    s
4058}
4059
4060#[inline]
4061fn norm2_slice(a: &[f64]) -> f64 {
4062    dot_slice(a, a).sqrt()
4063}
4064
4065fn rademacher_unit_probe_into_slice(z: &mut [f64], probe: u64, scale: f64) {
4066    let mut state = 0x6A09E667F3BCC909_u64 ^ probe.wrapping_mul(0xD1B54A32D192ED03);
4067    let mut bits = 0_u64;
4068    let mut remaining_bits = 0_u32;
4069    for value in z.iter_mut() {
4070        if remaining_bits == 0 {
4071            bits = splitmix64(&mut state);
4072            remaining_bits = 64;
4073        }
4074        *value = if bits & 1 == 0 { scale } else { -scale };
4075        bits >>= 1;
4076        remaining_bits -= 1;
4077    }
4078}
4079
4080#[inline]
4081const fn splitmix64(state: &mut u64) -> u64 {
4082    gam_linalg::utils::splitmix64(state)
4083}
4084
4085/// Authoritative factored-arrow evidence log-determinant.
4086///
4087/// This reads the stored joint value that cache construction records for the
4088/// operator used by selected-inverse/adjointers. It intentionally does not
4089/// reconstruct a determinant from row and Schur pieces for damped caches; if
4090/// construction did not record an exact joint log-det, evidence must refuse the
4091/// cache and route to an explicit matrix-free fallback.
4092pub fn arrow_log_det_from_cache(cache: &ArrowFactorCache) -> Option<f64> {
4093    if let Some(log_det) = cache.joint_hessian_log_det {
4094        return log_det.is_finite().then_some(log_det);
4095    }
4096    if cache.ridge_t != 0.0 || cache.ridge_beta != 0.0 {
4097        return None;
4098    }
4099    if cache.k > 0 && !cache.schur_factor_is_undamped {
4100        return None;
4101    }
4102    cache.compute_undamped_arrow_log_det()
4103}
4104
4105// ---------------------------------------------------------------------------
4106// IFT cascade: ∂u*/∂β → ∂β*/∂ρ → ∂u*/∂ρ
4107// ---------------------------------------------------------------------------
4108
4109/// Tier-1 IFT sensitivity `∂u_i*/∂β = -H_uu_i⁻¹ H_uβ_i`.
4110///
4111/// Concatenated row-major to a single `(N·d) × K` dense matrix. Each
4112/// row block is solved with the **undamped** Cholesky factor. Proposal
4113/// §2.2 / §7.
4114pub fn ift_du_dbeta(cache: &ArrowFactorCache) -> Array2<f64> {
4115    let n = cache.undamped_factor_count();
4116    let total_len = cache.delta_t_len();
4117    let k = cache.k;
4118    if !cache.htbeta_available() {
4119        return Array2::<f64>::from_elem((total_len, k), f64::NAN);
4120    }
4121    let mut out = Array2::<f64>::zeros((total_len, k));
4122    let mut beta_basis = Array1::<f64>::zeros(k);
4123    // Allocate scratch at max_d; per-row slice is ..di.
4124    let mut rhs = Array1::<f64>::zeros(cache.d);
4125    for i in 0..n {
4126        let di = cache.row_dims[i];
4127        let row_base = cache.row_offsets[i];
4128        let factor = cache.undamped_factor(i);
4129        // Solve H_uu_i Y = H_uβ_i column by column.
4130        for col in 0..k {
4131            beta_basis.fill(0.0);
4132            beta_basis[col] = 1.0;
4133            let mut rhs_i = rhs.slice_mut(ndarray::s![..di]).to_owned();
4134            // The Tier-2 IFT assembler is built only when the family's
4135            // capability surface promises cached `H_tβ` row products.
4136            if !cache.apply_htbeta_row(i, beta_basis.view(), &mut rhs_i) {
4137                // SAFETY: reaching `false` means a family declared the cache
4138                // available but failed to populate it — contract violation.
4139                return Array2::<f64>::from_elem((total_len, k), f64::NAN);
4140            }
4141            let y = cholesky_solve_vector(factor, &rhs_i);
4142            for c in 0..di {
4143                out[[row_base + c, col]] = -y[c];
4144            }
4145        }
4146    }
4147    out
4148}
4149
4150/// Coupling components of a symmetric coefficient Hessian: the connected
4151/// components of the graph whose vertices are coefficient indices `0..p` and
4152/// whose edges are the structurally nonzero off-diagonal entries of `H` (#779).
4153///
4154/// Returns a length-`p` vector of component labels in `0..num_components`,
4155/// where two indices share a label iff they are connected through a chain of
4156/// nonzero `H[i,j]` couplings. This is the exact structural partition the
4157/// cone-of-influence sensitivity reuse is keyed on: a smoothing-parameter move
4158/// whose stationarity-gradient derivative `∂g/∂ρ` is supported only inside one
4159/// component can change `β = -H⁻¹ ∂g/∂ρ` only inside that same component, so
4160/// the sensitivity of every *other* component is provably unchanged and may be
4161/// reused unrecomputed (lazy/local propagation).
4162///
4163/// The nonzero test is exact (`!= 0.0`), matching the structural-coupling gate
4164/// used elsewhere for the joint inner Hessian: a tolerance would risk dropping a
4165/// genuine (small) coupling edge and silently biasing the propagated sensitivity
4166/// — the failure mode #779/#740 explicitly guard against. A block-diagonal `H`
4167/// yields the all-singletons partition (one component per block-decoupled
4168/// coordinate); a fully coupled `H` yields a single component (no shortcut, the
4169/// full joint solve is required — and is what the non-coned path performs).
4170pub fn coupling_components(hessian: ArrayView2<'_, f64>) -> Vec<usize> {
4171    let p = hessian.nrows();
4172    if p == 0 || hessian.ncols() != p {
4173        return Vec::new();
4174    }
4175    // Union-find with path compression and union by size.
4176    let mut parent: Vec<usize> = (0..p).collect();
4177    let mut size: Vec<usize> = vec![1; p];
4178
4179    fn find(parent: &mut [usize], mut x: usize) -> usize {
4180        while parent[x] != x {
4181            parent[x] = parent[parent[x]];
4182            x = parent[x];
4183        }
4184        x
4185    }
4186
4187    for i in 0..p {
4188        for j in (i + 1)..p {
4189            // Symmetric structure: an edge exists if either triangle is nonzero,
4190            // so a numerically one-sided fill still couples the two indices.
4191            if hessian[[i, j]] != 0.0 || hessian[[j, i]] != 0.0 {
4192                let (ri, rj) = (find(&mut parent, i), find(&mut parent, j));
4193                if ri != rj {
4194                    let (small, large) = if size[ri] < size[rj] {
4195                        (ri, rj)
4196                    } else {
4197                        (rj, ri)
4198                    };
4199                    parent[small] = large;
4200                    size[large] += size[small];
4201                }
4202            }
4203        }
4204    }
4205
4206    // Relabel roots to a dense `0..num_components` range, preserving
4207    // first-seen order so labels are deterministic.
4208    let mut label_of_root: Vec<Option<usize>> = vec![None; p];
4209    let mut next_label = 0usize;
4210    let mut labels = vec![0usize; p];
4211    for idx in 0..p {
4212        let root = find(&mut parent, idx);
4213        let label = match label_of_root[root] {
4214            Some(l) => l,
4215            None => {
4216                let l = next_label;
4217                label_of_root[root] = Some(l);
4218                next_label += 1;
4219                l
4220            }
4221        };
4222        labels[idx] = label;
4223    }
4224    labels
4225}
4226
4227/// The cone of influence of a single stationarity-gradient derivative column
4228/// whose support (the coefficient indices where `∂g/∂ρ_k` is nonzero) lies in
4229/// `support`: the set of coefficient indices in the same coupling component(s)
4230/// as that support, given precomputed `labels` from [`coupling_components`].
4231///
4232/// `β_k = -H⁻¹ ∂g/∂ρ_k` is exactly zero outside this cone, so a confined solve
4233/// (or reuse of a cached zero) is exact, not an approximation. An empty support
4234/// (a structurally inactive `ρ_k`, e.g. a rank-0 or out-of-range penalty block)
4235/// yields an empty cone: the sensitivity is identically zero and no solve is
4236/// needed at all.
4237pub fn cone_of_influence(labels: &[usize], support: &[usize]) -> Vec<usize> {
4238    if support.is_empty() {
4239        return Vec::new();
4240    }
4241    let mut in_cone_labels: Vec<usize> = support
4242        .iter()
4243        .filter_map(|&idx| labels.get(idx).copied())
4244        .collect();
4245    in_cone_labels.sort_unstable();
4246    in_cone_labels.dedup();
4247    if in_cone_labels.is_empty() {
4248        return Vec::new();
4249    }
4250    (0..labels.len())
4251        .filter(|idx| in_cone_labels.binary_search(&labels[*idx]).is_ok())
4252        .collect()
4253}
4254
4255/// Tier-2 IFT sensitivity `∂β*/∂ρ = -A⁻¹ ∂g_red/∂ρ` (proposal §2.4 /
4256/// §7).
4257///
4258/// `dg_red_drho` is the `K × R` matrix whose `a`-th column is `q_a =
4259/// ∂g_red/∂ρ_a`. Returns the `K × R` matrix `β_ρ`.
4260///
4261/// Returns `None` if the Schur factor is unavailable (PCG mode) or was
4262/// built from a damped operator, or if any solved entry is non-finite;
4263/// callers must not silently substitute an approximation. The solve is
4264/// the one sensitivity operator (#935) — this site holds no private H⁻¹
4265/// convention of its own.
4266pub fn ift_dbeta_drho(
4267    cache: &ArrowFactorCache,
4268    dg_red_drho: ArrayView2<'_, f64>,
4269) -> Option<Array2<f64>> {
4270    if !cache.schur_factor_is_undamped {
4271        return None;
4272    }
4273    let schur = cache.schur_factor.as_ref()?;
4274    if dg_red_drho.nrows() != cache.k || schur.nrows() != cache.k {
4275        return None;
4276    }
4277    crate::sensitivity::FitSensitivity::from_lower_triangular(schur).mode_response(dg_red_drho)
4278}
4279
4280// ---------------------------------------------------------------------------
4281// ∂V/∂ρ — analytic optimized-evidence gradient via IFT mode response
4282// ---------------------------------------------------------------------------
4283
4284/// IFT terms needed to differentiate the optimized Laplace evidence through
4285/// the fitted mode `(β*(ρ), u*(ρ))`.
4286///
4287/// For each hyperparameter `ρ_a`, the correction added to the direct trace is
4288///
4289/// ```text
4290/// F_β · β_a + F_u · u_a
4291/// + 0.5 (∂_β log|H| · β_a + ∂_u log|H| · u_a).
4292/// ```
4293///
4294/// At an exact KKT point the value-gradient pieces are zero, but they are
4295/// explicit here so the exported gradient matches the optimized objective
4296/// whenever callers carry a certified nonzero residual correction.
4297#[derive(Clone)]
4298pub struct EvidenceIftGradientTerms<'a> {
4299    pub dbeta_drho: ArrayView2<'a, f64>,
4300    pub du_drho: ArrayView2<'a, f64>,
4301    pub value_beta: ArrayView1<'a, f64>,
4302    pub value_u: ArrayView1<'a, f64>,
4303    pub logdet_h_beta: ArrayView1<'a, f64>,
4304    pub logdet_h_u: ArrayView1<'a, f64>,
4305}
4306
4307/// Contract the IFT mode-response columns into the optimized-evidence
4308/// gradient correction.
4309pub fn evidence_ift_gradient_correction(terms: EvidenceIftGradientTerms<'_>) -> Array1<f64> {
4310    let k = terms.dbeta_drho.nrows();
4311    let nd = terms.du_drho.nrows();
4312    let r = terms.dbeta_drho.ncols();
4313    if terms.du_drho.ncols() != r
4314        || terms.value_beta.len() != k
4315        || terms.logdet_h_beta.len() != k
4316        || terms.value_u.len() != nd
4317        || terms.logdet_h_u.len() != nd
4318    {
4319        return Array1::<f64>::from_elem(r, f64::NAN);
4320    }
4321
4322    let mut out = Array1::<f64>::zeros(r);
4323    for a in 0..r {
4324        let mut acc = 0.0_f64;
4325        for j in 0..k {
4326            let mode = terms.dbeta_drho[[j, a]];
4327            acc += terms.value_beta[j] * mode;
4328            acc += 0.5 * terms.logdet_h_beta[j] * mode;
4329        }
4330        for j in 0..nd {
4331            let mode = terms.du_drho[[j, a]];
4332            acc += terms.value_u[j] * mode;
4333            acc += 0.5 * terms.logdet_h_u[j] * mode;
4334        }
4335        out[a] = acc;
4336    }
4337    out
4338}
4339
4340/// Per-`ρ` optimized-evidence gradient (proposal §3.7 / §3.8 split):
4341///
4342/// ```text
4343/// ∂V/∂ρ_a =
4344///       F_{ρ_a}                                  (value part)
4345///   + 0.5 tr(H⁻¹ H_{ρ_a})                        (direct Hessian)
4346///   + F_x · x_{ρ_a}
4347///   + 0.5 (∂_x log|H|) · x_{ρ_a}                 (IFT mode response)
4348///   - 0.5 tr(S_pen⁺ S_{pen,ρ_a})                 (penalty pseudo-logdet)
4349/// ```
4350/// where `x = (β, u)`.
4351///
4352/// The `tr(H⁻¹ H_{ρ_a})` trace is computed via the arrow structure
4353/// (proposal §3.5 / §3.10):
4354///
4355/// ```text
4356/// tr(H⁻¹ H_{ρ_a}) = Σ_i tr(H_uu_i⁻¹ ∂_{ρ_a} H_uu_i) + tr(A⁻¹ ∂_{ρ_a} A).
4357/// ```
4358///
4359/// `value_rho[a] = F_{ρ_a}` (envelope theorem, proposal §3.2).
4360/// `huu_drho[i][a]` is `∂H_uu_i/∂ρ_a` as a `d × d` matrix.
4361/// `hbb_drho[a]` is `∂H_ββ/∂ρ_a` as a `K × K` matrix.
4362/// `htbeta_drho[i][a]` is `∂H_uβ_i/∂ρ_a` as a `d × K` matrix.
4363/// `pen_logdet_drho[a]` is `∂_{ρ_a} log|S_pen|+`.
4364/// `ift_terms` carries `∂β*/∂ρ`, `∂u*/∂ρ`, and the already-contracted
4365/// mode derivatives of `F` and `log|H|`.
4366///
4367/// Returns the per-`ρ` gradient. Returns a NaN-filled vector when the
4368/// cache has no undamped Schur factor (PCG mode).
4369pub fn evidence_grad_rho(
4370    cache: &ArrowFactorCache,
4371    value_rho: ArrayView1<'_, f64>,
4372    huu_drho: &[Vec<Array2<f64>>],
4373    htbeta_drho: &[Vec<Array2<f64>>],
4374    hbb_drho: &[Array2<f64>],
4375    pen_logdet_drho: ArrayView1<'_, f64>,
4376    ift_terms: EvidenceIftGradientTerms<'_>,
4377) -> Array1<f64> {
4378    let r = value_rho.len();
4379    let n = cache.undamped_factor_count();
4380    let k = cache.k;
4381    let mut out = Array1::<f64>::zeros(r);
4382    if !cache.htbeta_available()
4383        || pen_logdet_drho.len() != r
4384        || huu_drho.len() != n
4385        || htbeta_drho.len() != n
4386        || hbb_drho.len() != r
4387        || huu_drho.iter().any(|row| row.len() != r)
4388        || htbeta_drho.iter().any(|row| row.len() != r)
4389        || hbb_drho.iter().any(|m| m.nrows() != k || m.ncols() != k)
4390        || huu_drho.iter().enumerate().any(|(i, row)| {
4391            let di = cache.row_dims[i];
4392            row.iter().any(|m| m.nrows() != di || m.ncols() != di)
4393        })
4394        || htbeta_drho.iter().enumerate().any(|(i, row)| {
4395            let di = cache.row_dims[i];
4396            row.iter().any(|m| m.nrows() != di || m.ncols() != k)
4397        })
4398    {
4399        out.fill(f64::NAN);
4400        return out;
4401    }
4402    let ift_correction = evidence_ift_gradient_correction(ift_terms);
4403    if ift_correction.len() != r || ift_correction.iter().any(|v| v.is_nan()) {
4404        out.fill(f64::NAN);
4405        return out;
4406    }
4407
4408    let schur = match cache.schur_factor.as_ref() {
4409        Some(s) => s,
4410        None => {
4411            for a in 0..r {
4412                out[a] = f64::NAN;
4413            }
4414            return out;
4415        }
4416    };
4417    if !cache.schur_factor_is_undamped {
4418        for a in 0..r {
4419            out[a] = f64::NAN;
4420        }
4421        return out;
4422    }
4423
4424    // Precompute Y_i = H_uu_i⁻¹ H_uβ_i (di × K). Used by both the Schur
4425    // derivative formula (§3.5) and the row trace `tr(H_uu_i⁻¹ ∂H_uu_i)`.
4426    let mut y_blocks: Vec<Array2<f64>> = Vec::with_capacity(n);
4427    let mut beta_basis = Array1::<f64>::zeros(k);
4428    // Scratch sized to max_d; per-row slice is ..di.
4429    let mut rhs = Array1::<f64>::zeros(cache.d);
4430    for i in 0..n {
4431        let di = cache.row_dims[i];
4432        let factor = cache.undamped_factor(i);
4433        let mut yi = Array2::<f64>::zeros((di, k));
4434        for col in 0..k {
4435            beta_basis.fill(0.0);
4436            beta_basis[col] = 1.0;
4437            let mut rhs_i = rhs.slice_mut(ndarray::s![..di]).to_owned();
4438            // Same H_tβ cache contract as the IFT du/dβ and du/dρ paths.
4439            if !cache.apply_htbeta_row(i, beta_basis.view(), &mut rhs_i) {
4440                // SAFETY: `false` means the family declared the cache
4441                // available but did not populate it — contract violation.
4442                out.fill(f64::NAN);
4443                return out;
4444            }
4445            let v = cholesky_solve_vector(factor, &rhs_i);
4446            for c in 0..di {
4447                yi[[c, col]] = v[c];
4448            }
4449        }
4450        y_blocks.push(yi);
4451    }
4452
4453    // Outer-hoisted scratch reused across all (a, i) iterations.
4454    // Sized to max_d for trace_rhs and da_tmp; per-row slices used below.
4455    let mut trace_rhs = Array1::<f64>::zeros(cache.d);
4456    let mut da_tmp = Array2::<f64>::zeros((cache.d, k));
4457    let mut col_scratch = Array1::<f64>::zeros(k);
4458    for a in 0..r {
4459        // Part 1: F_{ρ_a} envelope contribution.
4460        let mut grad = value_rho[a];
4461
4462        // Part 2a: Σ_i tr(H_uu_i⁻¹ ∂H_uu_i).
4463        // tr(H_uu_i⁻¹ M_i) = tr(L_iᵀ⁻¹ L_i⁻¹ M_i). Compute as the sum
4464        // over columns: solve L_i Lᵀ x = e_c for the c-th column of
4465        // M_i, then take its c-th component. Equivalently and more
4466        // cheaply, build (H_uu_i⁻¹ M_i) by solving column-by-column
4467        // and take its diagonal sum.
4468        let mut row_trace_acc = 0.0_f64;
4469        for i in 0..n {
4470            let di = cache.row_dims[i];
4471            let m_i = &huu_drho[i][a];
4472            assert_eq!(m_i.shape(), &[di, di]);
4473            for col in 0..di {
4474                let mut tr_rhs_i = trace_rhs.slice_mut(ndarray::s![..di]).to_owned();
4475                for r0 in 0..di {
4476                    tr_rhs_i[r0] = m_i[[r0, col]];
4477                }
4478                let v = cholesky_solve_vector(cache.undamped_factor(i), &tr_rhs_i);
4479                row_trace_acc += v[col];
4480            }
4481        }
4482
4483        // Part 2b: tr(A⁻¹ ∂A) where (proposal §3.5)
4484        //     ∂A = ∂H_ββ
4485        //          - Σ_i (∂H_uβ_i)ᵀ Y_i
4486        //          - Σ_i Y_iᵀ (∂H_uβ_i)
4487        //          + Σ_i Y_iᵀ (∂H_uu_i) Y_i.
4488        // We accumulate ∂A as a dense `K × K` matrix, then evaluate
4489        // tr(A⁻¹ ∂A) by `Σ_j (A⁻¹ ∂A)[j, j]` via column solves of the
4490        // Schur Cholesky.
4491        let mut da = hbb_drho[a].clone();
4492        assert_eq!(da.shape(), &[k, k]);
4493        for i in 0..n {
4494            let di = cache.row_dims[i];
4495            let dhtb = &htbeta_drho[i][a]; // di × K
4496            let yi = &y_blocks[i]; // di × K
4497            // - (∂H_uβ_i)ᵀ Y_i
4498            for r0 in 0..k {
4499                for c0 in 0..k {
4500                    let mut acc = 0.0;
4501                    for cc in 0..di {
4502                        acc += dhtb[[cc, r0]] * yi[[cc, c0]];
4503                    }
4504                    da[[r0, c0]] -= acc;
4505                }
4506            }
4507            // - Y_iᵀ (∂H_uβ_i)
4508            for r0 in 0..k {
4509                for c0 in 0..k {
4510                    let mut acc = 0.0;
4511                    for cc in 0..di {
4512                        acc += yi[[cc, r0]] * dhtb[[cc, c0]];
4513                    }
4514                    da[[r0, c0]] -= acc;
4515                }
4516            }
4517            // + Y_iᵀ (∂H_uu_i) Y_i
4518            let dhuu = &huu_drho[i][a];
4519            // tmp = (∂H_uu_i) Y_i  (di × K) — use a slice of the hoisted buffer.
4520            let mut da_tmp_i = da_tmp.slice_mut(ndarray::s![..di, ..]).to_owned();
4521            for r0 in 0..di {
4522                for c0 in 0..k {
4523                    let mut acc = 0.0;
4524                    for cc in 0..di {
4525                        acc += dhuu[[r0, cc]] * yi[[cc, c0]];
4526                    }
4527                    da_tmp_i[[r0, c0]] = acc;
4528                }
4529            }
4530            // da += Y_iᵀ tmp
4531            for r0 in 0..k {
4532                for c0 in 0..k {
4533                    let mut acc = 0.0;
4534                    for cc in 0..di {
4535                        acc += yi[[cc, r0]] * da_tmp_i[[cc, c0]];
4536                    }
4537                    da[[r0, c0]] += acc;
4538                }
4539            }
4540        }
4541
4542        // tr(A⁻¹ ∂A) via column solves.
4543        let mut schur_trace_acc = 0.0_f64;
4544        for j in 0..k {
4545            for r0 in 0..k {
4546                col_scratch[r0] = da[[r0, j]];
4547            }
4548            let v = cholesky_solve_vector(schur, &col_scratch);
4549            schur_trace_acc += v[j];
4550        }
4551
4552        grad += 0.5 * (row_trace_acc + schur_trace_acc);
4553        grad += ift_correction[a];
4554
4555        // Part 3: -0.5 ∂_{ρ_a} log|S_pen|+.
4556        grad -= 0.5 * pen_logdet_drho[a];
4557
4558        out[a] = grad;
4559    }
4560    out
4561}
4562
4563// ---------------------------------------------------------------------------
4564// Topology selection
4565// ---------------------------------------------------------------------------
4566
4567/// Enumerate the candidate topologies, rank by normalized negative log
4568/// evidence, and return the winner. Failed/excluded candidates (proposal
4569/// §6.11) are appended at the end of `ranking` and are never the winner.
4570///
4571/// The caller fits each topology separately (proposal §4.2) and supplies
4572/// the resulting `TopologyCandidate` records. This function is purely
4573/// the discrete comparator + tie breaker.
4574///
4575/// # Tie-breaking
4576///
4577/// Per proposal §4.6: if normalized `|score_a - score_b| <= tie_tolerance`,
4578/// prefer the simpler topology by `TopologyKind::complexity_rank` (flat <
4579/// periodic < sphere < torus). The `tie` flag in the result records whether
4580/// such a tie occurred at the top of the ranking.
4581///
4582/// # Panics
4583///
4584/// Panics if `candidates` is empty after filtering out non-finite
4585/// scores. Proposal §6.11 explicitly forbids silent fallback to a
4586/// default topology; callers must handle the empty-candidate case
4587/// before invocation.
4588pub fn select_topology(
4589    candidates: &[TopologyCandidate],
4590    options: TopologySelectOptions,
4591) -> SelectedTopology {
4592    // Split valid and excluded.
4593    let mut valid: Vec<TopologyCandidate> = candidates
4594        .iter()
4595        .filter(|c| {
4596            c.converged
4597                && c.exclusion_reason.is_none()
4598                && c.negative_log_evidence.is_finite()
4599                && topology_selection_score(c, options.score_scale).is_finite()
4600        })
4601        .cloned()
4602        .collect();
4603    let mut excluded: Vec<TopologyCandidate> = candidates
4604        .iter()
4605        .filter(|c| {
4606            !(c.converged && c.exclusion_reason.is_none() && c.negative_log_evidence.is_finite())
4607                || !topology_selection_score(c, options.score_scale).is_finite()
4608        })
4609        .cloned()
4610        .collect();
4611
4612    assert!(
4613        !valid.is_empty(),
4614        "select_topology: no finite valid candidates; proposal §6.11 forbids silent fallback"
4615    );
4616
4617    // Sort by normalized negative log evidence (ascending = best first),
4618    // breaking ties by complexity_rank (smaller wins). The shared selector is
4619    // the single lower-is-better ordering contract used by topology ranking,
4620    // seed screening, and REML model comparison (#782).
4621    valid = rank_priority_candidates(
4622        valid
4623            .into_iter()
4624            .enumerate()
4625            .map(|(idx, row)| {
4626                let score = topology_selection_score(&row, options.score_scale);
4627                let tie_break = usize::from(row.kind.complexity_rank());
4628                PriorityCandidate::new(row, idx, score, tie_break)
4629            })
4630            .collect(),
4631    )
4632    .into_iter()
4633    .map(|row| row.item)
4634    .collect();
4635
4636    // Detect numerical tie at the top.
4637    let tie = if valid.len() >= 2 {
4638        let top = topology_selection_score(&valid[0], options.score_scale);
4639        let next = topology_selection_score(&valid[1], options.score_scale);
4640        (next - top).abs() <= options.tie_tolerance
4641    } else {
4642        false
4643    };
4644
4645    // If tied, prefer simpler topology among the tied prefix.
4646    if tie {
4647        let top_score = topology_selection_score(&valid[0], options.score_scale);
4648        // Find the tied prefix range.
4649        let tied_end = valid
4650            .iter()
4651            .position(|c| {
4652                (topology_selection_score(c, options.score_scale) - top_score).abs()
4653                    > options.tie_tolerance
4654            })
4655            .unwrap_or(valid.len());
4656        // Sort the tied prefix by complexity_rank ascending.
4657        valid[..tied_end].sort_by_key(|c| c.kind.complexity_rank());
4658    }
4659
4660    let winner = valid[0].kind;
4661    valid.append(&mut excluded);
4662    SelectedTopology {
4663        winner,
4664        ranking: valid,
4665        tie,
4666    }
4667}
4668
4669fn topology_selection_score(candidate: &TopologyCandidate, scale: TopologyScoreScale) -> f64 {
4670    match scale {
4671        TopologyScoreScale::PerObservation => {
4672            if candidate.n_obs == 0 {
4673                f64::NAN
4674            } else {
4675                candidate.negative_log_evidence / candidate.n_obs as f64
4676            }
4677        }
4678        TopologyScoreScale::PerEffectiveDim => {
4679            if !(candidate.effective_dim.is_finite() && candidate.effective_dim > 0.0) {
4680                f64::NAN
4681            } else {
4682                candidate.negative_log_evidence / candidate.effective_dim
4683            }
4684        }
4685    }
4686}
4687
4688// ---------------------------------------------------------------------------
4689// Cache verification helpers
4690// ---------------------------------------------------------------------------
4691
4692/// Verifies the `ArrowSchurSystem` dimensions match the cache. Used as
4693/// a debug-time precondition; never silently masks shape errors
4694/// (proposal §6.9 — sign and shape errors must be loud).
4695pub fn cache_matches_system(cache: &ArrowFactorCache, sys: &ArrowSchurSystem) -> bool {
4696    cache.d == sys.d
4697        && cache.k == sys.k
4698        && cache.n_rows() == sys.rows.len()
4699        && cache.undamped_factor_count() == sys.rows.len()
4700        && cache.manifold_mode_fingerprint == sys.manifold_mode_fingerprint
4701        && cache.row_hessian_fingerprint == sys.current_row_hessian_fingerprint()
4702}
4703
4704// ---------------------------------------------------------------------------
4705// #1026 hybrid curved + linear-tail dictionary split-selection
4706// ---------------------------------------------------------------------------
4707//
4708// COMMON-EVIDENCE NOTE (#1202): the candidates BOTH fit the same data — the
4709// atom's leave-this-atom-out response residual `y_resp` (the response with every
4710// other atom's contribution removed). The curved candidate predicts the atom's
4711// actual mass-scaled contribution `a_k·γ_k`, the linear candidate the best
4712// mass-weighted straight line fit to `y_resp`. Because the curved family's
4713// `Θ = 0` member reproduces the linear prediction exactly, linear IS the nested
4714// `Θ = 0` sub-model on common data, so the "match-or-beat" statements below are a
4715// genuine data-level comparison: the curved candidate wins only when fitting the
4716// response residual better than its own straight projection pays for its extra
4717// parameters. See `crate::terms::sae::hybrid_split` for the residual assembly.
4718//
4719// The per-slot adjudication uses the SAME rank-aware Laplace evidence criterion
4720// the union/mixture rungs use (`−V = NLE`, lower wins), comparing the data-fit +
4721// complexity cost of the curved contribution against that of the straight line.
4722//
4723// ## The turning floor (Θ → 0) and the curved ceiling (Θ large)
4724//
4725// Per slot, the curved candidate fits the response residual with its actual
4726// mass-scaled contribution `a_k·γ_k` (data-fit `½·curved_rss`) and pays a larger
4727// free-parameter price `P_curved > P_linear`; the linear candidate fits the same
4728// residual with its best straight line (data-fit `½·linear_rss ≥ ½·curved_rss`
4729// whenever the curve beats its own straight projection) at a smaller price,
4730// charged with its genuine weighted Gram logdet `p·(log w_sum + log s_tt)`
4731// (#1203). Hence:
4732//
4733//   * Θ → 0 (the residual is straight): the curve and the line fit it equally, so
4734//     the cheaper LINEAR candidate wins — the turning floor / nested dominance. A
4735//     curved parameterization "buys nothing" on an already-straight residual.
4736//   * Θ large (a genuinely turning residual): the line's data-fit residual
4737//     exceeds the curved atom's extra parameter price, so CURVED wins. (Whether
4738//     curved wins also depends on the coordinate spread `s_tt` and amplitude, via
4739//     the honest logdet — a tightly-spread, mildly-curved residual can still
4740//     prefer the cheaper line.)
4741//
4742// The crossover is governed by the documented shatter law: a linear SAE shatters
4743// a feature of total turning Θ into `N(ε) ≈ Θ/(2√(2ε))` rank-1 directions at
4744// relative reconstruction error ε, so the curved advantage scales as `Θ/√ε`. We
4745// use the fitted turning Θ (`sae::chart_canonicalization::d1_atom_fitted_turning`)
4746// as the decision FEATURE: it both (a) sharpens the evidence comparison into a
4747// falsifiable per-atom prediction and (b) provides the exact-zero dominance
4748// guard — when an atom's fitted turning is identically zero, the curved fit has
4749// no curvature to price and the linear special case is selected by construction,
4750// independent of finite-sample evidence noise.
4751
4752/// Which atom parameterization a hybrid-dictionary slot selects: a CURVED atom
4753/// (a `latent_dim ≥ 1` curved basis whose decoded image may turn) or its LINEAR
4754/// special case (the euclidean-d=1-linear atom — one straight decoder direction,
4755/// `γ(t) = t·b`, fitted turning `Θ = 0`).
4756#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
4757pub enum HybridAtomParam {
4758    /// The curved atom (`latent_dim ≥ 1`), priced at its full coefficient count.
4759    Curved { latent_dim: usize },
4760    /// The linear special case: one decoder direction, zero turning.
4761    Linear,
4762}
4763
4764impl HybridAtomParam {
4765    /// Stable display name for logs and tests.
4766    pub const fn as_str(self) -> &'static str {
4767        match self {
4768            HybridAtomParam::Curved { .. } => "curved",
4769            HybridAtomParam::Linear => "linear",
4770        }
4771    }
4772
4773    /// `true` iff this is the linear special case (the linear tail).
4774    pub const fn is_linear(self) -> bool {
4775        matches!(self, HybridAtomParam::Linear)
4776    }
4777}
4778
4779/// One fitted candidate parameterization for a single hybrid-dictionary atom
4780/// slot, scored on the COMMON rank-aware Laplace scale (`−V = NLE`, lower wins,
4781/// identical to the union/mixture rungs). The curved and linear candidates for
4782/// the SAME slot are fit on the same rows AND the same data (the atom's response
4783/// residual, #1202), so their NLEs are directly comparable; the structural
4784/// difference is the curved candidate's larger free-parameter price and whatever
4785/// data-fit it buys with its curvature.
4786#[derive(Debug, Clone, Copy)]
4787pub struct HybridAtomCandidate {
4788    pub param: HybridAtomParam,
4789    /// Rank-aware Laplace negative-log-evidence on the common scale (lower wins).
4790    pub negative_log_evidence: f64,
4791    /// Free-parameter count this candidate is charged for (the complexity price).
4792    pub num_parameters: usize,
4793    /// The candidate's fitted total turning `Θ = ∫κ ds` of its decoded curve, if
4794    /// the basis admits an analytic second jet. `Some(0.0)` for a linear atom (a
4795    /// straight image has no turning); `None` when the turning is honestly
4796    /// unavailable (no second jet / degenerate curve) — never fabricated.
4797    pub fitted_turning: Option<f64>,
4798}
4799
4800impl HybridAtomCandidate {
4801    /// A linear special-case candidate: exact zero turning by construction.
4802    pub fn linear(negative_log_evidence: f64, num_parameters: usize) -> Self {
4803        Self {
4804            param: HybridAtomParam::Linear,
4805            negative_log_evidence,
4806            num_parameters,
4807            fitted_turning: Some(0.0),
4808        }
4809    }
4810
4811    /// A curved candidate of the given latent dimension, with its fitted turning.
4812    pub fn curved(
4813        latent_dim: usize,
4814        negative_log_evidence: f64,
4815        num_parameters: usize,
4816        fitted_turning: Option<f64>,
4817    ) -> Self {
4818        Self {
4819            param: HybridAtomParam::Curved { latent_dim },
4820            negative_log_evidence,
4821            num_parameters,
4822            fitted_turning,
4823        }
4824    }
4825}
4826
4827/// The evidence-selected parameterization for one hybrid-dictionary atom slot:
4828/// the winning candidate, plus the curved/linear NLEs that decided it (for the
4829/// EV-vs-Θ diagnostic and the tie-break audit trail).
4830#[derive(Debug, Clone, Copy)]
4831pub struct HybridAtomChoice {
4832    pub param: HybridAtomParam,
4833    /// The winning candidate's NLE.
4834    pub negative_log_evidence: f64,
4835    /// The winning candidate's free-parameter price.
4836    pub num_parameters: usize,
4837    /// The curved candidate's fitted turning `Θ` (the decision feature). `None`
4838    /// when no curved candidate offered an analytic turning.
4839    pub curved_turning: Option<f64>,
4840    /// `NLE_linear − NLE_curved`: the evidence margin the curved fit won (or lost,
4841    /// if negative) over the linear special case at this slot. Positive ⇒ curved
4842    /// bought more evidence than its parameter price; ≤ 0 ⇒ the dominance floor
4843    /// keeps the linear tail.
4844    pub curved_evidence_margin: f64,
4845}
4846
4847/// Below this fitted turning the curved candidate is treated as straight: its
4848/// curvature is numerically indistinguishable from zero, so the dominance floor
4849/// (the linear special case is cheaper at equal likelihood) is enforced by
4850/// construction rather than left to finite-sample evidence noise. This is the
4851/// exact-zero guard from the `Θ → 0 ⇒ N(ε) → 0` limit of the shatter law, not a
4852/// tunable knob: it is the curvature scale below which `‖γ' ∧ γ''‖` is at the
4853/// floor of the Simpson quadrature for a genuinely straight image.
4854pub const HYBRID_LINEAR_TURNING_FLOOR: f64 = 1e-9;
4855
4856/// Adjudicate the curved-vs-linear parameterization for ONE hybrid-dictionary
4857/// atom slot by the common rank-aware Laplace evidence criterion.
4858///
4859/// Selection rule (all on the single `NLE = −V` scale, lower wins):
4860///
4861///  1. **Dominance floor (Θ → 0).** If the curved candidate's fitted turning is
4862///     `Some(Θ)` with `Θ ≤ HYBRID_LINEAR_TURNING_FLOOR` and a linear candidate
4863///     exists, select LINEAR. A straight curved fit recovers no likelihood the
4864///     linear special case does not, and the linear atom is strictly cheaper, so
4865///     it cannot lose — we enforce that exactly instead of trusting evidence
4866///     noise at the floor.
4867///  2. **Evidence comparison.** Otherwise select the candidate with the smaller
4868///     `NLE`. The curved candidate wins only when its extra curvature lowers the
4869///     NLE by MORE than its extra parameter price — the `Θ/√ε` crossover, decided
4870///     here by the evidence numbers themselves, not by fiat. This is a
4871///     common-data comparison (both candidates fit the atom's response residual,
4872///     see `crate::terms::sae::hybrid_split`) in which linear is the curved
4873///     family's nested `Θ = 0` sub-model (#1202): the curved candidate cannot be
4874///     charged its extra parameters to fit the residual no better than its own
4875///     straight projection, and a tightly-spread, mildly-curved residual can
4876///     still prefer the cheaper line.
4877///  3. **Tie-break.** Exact NLE ties go to the cheaper (fewer-parameter)
4878///     candidate — i.e. linear — preserving the strict-generalization guarantee
4879///     that the hybrid never pays for curvature it does not need.
4880///
4881/// `candidates` must contain at most one linear and at most one curved candidate
4882/// for the slot; returns `None` only if `candidates` is empty.
4883pub fn select_hybrid_atom(candidates: &[HybridAtomCandidate]) -> Option<HybridAtomChoice> {
4884    if candidates.is_empty() {
4885        return None;
4886    }
4887    let linear = candidates.iter().find(|c| c.param.is_linear());
4888    let curved = candidates.iter().find(|c| !c.param.is_linear());
4889    let curved_turning = curved.and_then(|c| c.fitted_turning);
4890    let curved_evidence_margin = match (linear, curved) {
4891        (Some(l), Some(c)) => l.negative_log_evidence - c.negative_log_evidence,
4892        _ => 0.0,
4893    };
4894
4895    // (1) Exact-zero dominance floor: a straight curved fit yields to the linear
4896    // special case by construction.
4897    if let (Some(l), Some(turning)) = (linear, curved_turning)
4898        && turning <= HYBRID_LINEAR_TURNING_FLOOR
4899    {
4900        return Some(HybridAtomChoice {
4901            param: l.param,
4902            negative_log_evidence: l.negative_log_evidence,
4903            num_parameters: l.num_parameters,
4904            curved_turning,
4905            curved_evidence_margin,
4906        });
4907    }
4908
4909    // (2)+(3) Evidence argmin with the cheaper candidate winning exact ties.
4910    let mut best = candidates[0];
4911    for cand in &candidates[1..] {
4912        let better_evidence = cand.negative_log_evidence < best.negative_log_evidence;
4913        let tied = cand.negative_log_evidence == best.negative_log_evidence;
4914        let cheaper_on_tie = tied && cand.num_parameters < best.num_parameters;
4915        if better_evidence || cheaper_on_tie {
4916            best = *cand;
4917        }
4918    }
4919    Some(HybridAtomChoice {
4920        param: best.param,
4921        negative_log_evidence: best.negative_log_evidence,
4922        num_parameters: best.num_parameters,
4923        curved_turning,
4924        curved_evidence_margin,
4925    })
4926}
4927
4928/// The evidence-selected split for a whole hybrid dictionary: the per-atom
4929/// curved-vs-linear choices and the dictionary-level aggregates the EV-vs-Θ
4930/// frontier reports against.
4931#[derive(Debug, Clone)]
4932pub struct HybridSplitSelection {
4933    /// One adjudicated choice per atom slot, in slot order.
4934    pub atoms: Vec<HybridAtomChoice>,
4935    /// `Σ NLE` across the selected per-atom parameterizations — the dictionary's
4936    /// summed rank-aware Laplace negative-log-evidence (lower wins). Because each
4937    /// slot picks the argmin over {curved contribution, best straight line to the
4938    /// response residual}, this is ≤ the sum of the per-slot LINEAR-candidate
4939    /// NLEs. The linear baseline is the best straight line fit to each atom's
4940    /// leave-this-atom-out RESPONSE residual (#1202), the curved family's nested
4941    /// `Θ = 0` member on common data — so this is a genuine data-level
4942    /// match-or-beat dominance, not a post-hoc curve-simplification one.
4943    pub total_negative_log_evidence: f64,
4944    /// `Σ P` across the selected parameterizations — the dictionary's total
4945    /// free-parameter price (the matched-active-budget accounting).
4946    pub total_parameters: usize,
4947    /// Count of slots that selected the curved parameterization.
4948    pub curved_atom_count: usize,
4949}
4950
4951impl HybridSplitSelection {
4952    /// Count of slots that selected the linear special case (the linear tail).
4953    pub fn linear_atom_count(&self) -> usize {
4954        self.atoms.len() - self.curved_atom_count
4955    }
4956
4957    /// `true` iff every slot selected linear — the pure-linear limit, reached
4958    /// when every feature is straight (all `Θ → 0`).
4959    pub fn is_pure_linear(&self) -> bool {
4960        self.curved_atom_count == 0 && !self.atoms.is_empty()
4961    }
4962
4963    /// `true` iff every slot selected curved — the pure-curved limit, reached
4964    /// when every feature turns enough to pay for curvature.
4965    pub fn is_pure_curved(&self) -> bool {
4966        self.curved_atom_count == self.atoms.len() && !self.atoms.is_empty()
4967    }
4968}
4969
4970/// Adjudicate the curved-vs-linear split across a whole hybrid dictionary by the
4971/// common evidence criterion. `slots[i]` holds the curved/linear candidates for
4972/// atom slot `i` (each scored on the same rows, on the common Laplace scale).
4973///
4974/// The result reduces EXACTLY to pure-linear when every slot's curved candidate
4975/// has `Θ → 0` (the turning floor fires everywhere) and to pure-curved when
4976/// every slot's curved candidate wins the evidence comparison. (Common-data
4977/// criterion, #1202 — both candidates fit the atom's response residual, with
4978/// linear nested as the curved family's `Θ = 0` sub-model; see the module header
4979/// above and `crate::terms::sae::hybrid_split`.)
4980///
4981/// Returns an error only if some slot has no candidates to adjudicate (an empty
4982/// dictionary slot is a caller bug, not a silent skip).
4983pub fn select_hybrid_split(
4984    slots: &[Vec<HybridAtomCandidate>],
4985) -> Result<HybridSplitSelection, String> {
4986    let mut atoms = Vec::with_capacity(slots.len());
4987    let mut total_nle = 0.0_f64;
4988    let mut total_parameters = 0usize;
4989    let mut curved_atom_count = 0usize;
4990    for (i, slot) in slots.iter().enumerate() {
4991        let choice = select_hybrid_atom(slot)
4992            .ok_or_else(|| format!("hybrid split slot {i} has no candidate parameterizations"))?;
4993        if !choice.negative_log_evidence.is_finite() {
4994            return Err(format!(
4995                "hybrid split slot {i} selected a non-finite evidence ({})",
4996                choice.negative_log_evidence
4997            ));
4998        }
4999        if !choice.param.is_linear() {
5000            curved_atom_count += 1;
5001        }
5002        total_nle += choice.negative_log_evidence;
5003        total_parameters += choice.num_parameters;
5004        atoms.push(choice);
5005    }
5006    Ok(HybridSplitSelection {
5007        atoms,
5008        total_negative_log_evidence: total_nle,
5009        total_parameters,
5010        curved_atom_count,
5011    })
5012}
5013
5014// ---------------------------------------------------------------------------
5015// Tests
5016//
5017// These are type-level / structural tests: per the task contract we do
5018// not compile or run them in this session. They document the expected
5019// shapes and degenerate-case behavior so a future maintainer running
5020// `cargo test` sees the contract written down.
5021// ---------------------------------------------------------------------------
5022
5023#[cfg(test)]
5024mod tests {
5025    use super::*;
5026    use crate::arrow_schur::ArrowFactorSlab;
5027    use ndarray::array;
5028
5029    // Dense `H⁻¹` apply via explicit inverse (test-only reference solver).
5030    fn dense_inverse(h: &Array2<f64>) -> Array2<f64> {
5031        let p = h.nrows();
5032        let mut aug = Array2::<f64>::zeros((p, 2 * p));
5033        for i in 0..p {
5034            for j in 0..p {
5035                aug[[i, j]] = h[[i, j]];
5036            }
5037            aug[[i, p + i]] = 1.0;
5038        }
5039        for col in 0..p {
5040            let mut pivot = col;
5041            for row in (col + 1)..p {
5042                if aug[[row, col]].abs() > aug[[pivot, col]].abs() {
5043                    pivot = row;
5044                }
5045            }
5046            if pivot != col {
5047                for j in 0..(2 * p) {
5048                    aug.swap([col, j], [pivot, j]);
5049                }
5050            }
5051            let d = aug[[col, col]];
5052            for j in 0..(2 * p) {
5053                aug[[col, j]] /= d;
5054            }
5055            for row in 0..p {
5056                if row == col {
5057                    continue;
5058                }
5059                let f = aug[[row, col]];
5060                if f != 0.0 {
5061                    for j in 0..(2 * p) {
5062                        aug[[row, j]] -= f * aug[[col, j]];
5063                    }
5064                }
5065            }
5066        }
5067        let mut inv = Array2::<f64>::zeros((p, p));
5068        for i in 0..p {
5069            for j in 0..p {
5070                inv[[i, j]] = aug[[i, p + j]];
5071            }
5072        }
5073        inv
5074    }
5075
5076    /// The rate is what separates "interrupted mid-descent" from "stuck", so it
5077    /// must read a clean geometric decay exactly and refuse to speak before it
5078    /// has a full window.
5079    #[test]
5080    fn em_contraction_rate_recovers_a_planted_geometric_decay() {
5081        let planted = 0.98_f64;
5082        let mut window = std::collections::VecDeque::new();
5083        let mut residual = 1.0_f64;
5084        for _ in 0..EM_RATE_WINDOW {
5085            window.push_back(residual);
5086            residual *= planted;
5087        }
5088        // One short of a full window: no rate may be claimed yet.
5089        assert_eq!(em_contraction_rate(&window), None);
5090        window.push_back(residual);
5091        let measured = em_contraction_rate(&window).expect("a full window yields a rate");
5092        assert!(
5093            (measured - planted).abs() < 1e-12,
5094            "measured {measured} should recover the planted {planted}"
5095        );
5096    }
5097
5098    /// A residual that is flat or growing must NOT produce a rate below 1, or a
5099    /// stalled iterate would earn an extension it cannot use.
5100    #[test]
5101    fn em_contraction_rate_does_not_contract_on_a_flat_or_growing_residual() {
5102        let flat: std::collections::VecDeque<f64> =
5103            std::iter::repeat_n(1e-6, EM_RATE_WINDOW + 1).collect();
5104        let rate = em_contraction_rate(&flat).expect("a full window yields a rate");
5105        assert!(rate >= 1.0, "a flat residual must not look like contraction");
5106        let growing: std::collections::VecDeque<f64> = (0..=EM_RATE_WINDOW)
5107            .map(|i| 1e-6 * 1.01_f64.powi(i as i32))
5108            .collect();
5109        let rate = em_contraction_rate(&growing).expect("a full window yields a rate");
5110        assert!(rate > 1.0, "a growing residual must not look like contraction");
5111    }
5112
5113    /// The projection is the deadline an extension is held to, so it must invert
5114    /// the decay exactly and decline to exist when there is nothing to project.
5115    #[test]
5116    fn em_projected_iterations_inverts_the_decay_and_declines_otherwise() {
5117        // 1.0 -> 1e-8 at rate 0.98 needs ln(1e-8)/ln(0.98) = 911.6 -> 912.
5118        let steps = em_projected_iterations(1.0, 1e-8, 0.98).expect("a contracting rate projects");
5119        assert_eq!(steps, 912);
5120        // Applying the rate for that many steps must actually reach tolerance.
5121        assert!(0.98_f64.powi(steps as i32) <= 1e-8);
5122        // No projection without contraction, or when already inside tolerance.
5123        assert_eq!(em_projected_iterations(1.0, 1e-8, 1.0), None);
5124        assert_eq!(em_projected_iterations(1.0, 1e-8, 1.05), None);
5125        assert_eq!(em_projected_iterations(1e-9, 1e-8, 0.98), None);
5126    }
5127
5128    #[test]
5129    fn coupling_components_block_diagonal_is_all_singletons_by_block() {
5130        // Two decoupled 2x2 blocks: {0,1} and {2,3}.
5131        let mut h = Array2::<f64>::eye(4);
5132        h[[0, 1]] = 0.3;
5133        h[[1, 0]] = 0.3;
5134        h[[2, 3]] = 0.7;
5135        h[[3, 2]] = 0.7;
5136        let labels = coupling_components(h.view());
5137        assert_eq!(labels[0], labels[1]);
5138        assert_eq!(labels[2], labels[3]);
5139        assert_ne!(labels[0], labels[2]);
5140        // Exactly two components.
5141        let mut uniq = labels.clone();
5142        uniq.sort_unstable();
5143        uniq.dedup();
5144        assert_eq!(uniq.len(), 2);
5145    }
5146
5147    #[test]
5148    fn coupling_components_fully_coupled_is_one_component() {
5149        let mut h = Array2::<f64>::eye(3);
5150        for i in 0..3 {
5151            for j in 0..3 {
5152                if i != j {
5153                    h[[i, j]] = 0.1;
5154                }
5155            }
5156        }
5157        let labels = coupling_components(h.view());
5158        assert!(labels.iter().all(|&l| l == labels[0]));
5159    }
5160
5161    #[test]
5162    fn coupling_components_transitive_chain_merges() {
5163        // 0-1 and 1-2 coupled (but no direct 0-2 edge) must form one component.
5164        let mut h = Array2::<f64>::eye(3);
5165        h[[0, 1]] = 0.5;
5166        h[[1, 0]] = 0.5;
5167        h[[1, 2]] = 0.5;
5168        h[[2, 1]] = 0.5;
5169        let labels = coupling_components(h.view());
5170        assert_eq!(labels[0], labels[1]);
5171        assert_eq!(labels[1], labels[2]);
5172    }
5173
5174    #[test]
5175    fn compare_reml_fits_delta_and_bayes_factor_never_contradict_winner_gh1465() {
5176        // Regression for #1465: the ranking `delta` / `bayes_factor` must be
5177        // measured on the SAME scale that orders the table (the Occam-penalised
5178        // conditional AIC `ranking_score`), so every row's delta is >= 0 and its
5179        // Bayes factor >= 1 — the table must never claim a non-winner beats the
5180        // declared winner. The scenario is exactly the case the comparison
5181        // exists to handle: AIC and raw REML DISAGREE. `m1` is the AIC winner
5182        // but does NOT carry the minimum raw REML (`m2` does) — the noise
5183        // extra-term case from the issue.
5184        //
5185        // `ranking_score` = -2*log_lik + 2*edf; with log_lik = 0 it is `2*edf`,
5186        // so the AIC order is m1 < m2 < m3 while the raw-REML order has m2 lowest.
5187        let cand = |name: &str, score: f64, edf: f64| RemlCandidate {
5188            index: 0,
5189            name: name.to_string(),
5190            score,
5191            edf: Some(edf),
5192            log_lik: Some(0.0),
5193            family: Some("gaussian".to_string()),
5194            n_obs: Some(100),
5195        };
5196        // raw REML : m2 (41.605) < m1 (53.748) < m3 (120.011)
5197        // AIC=2*edf: m1 (100)    < m2 (102)    < m3 (130)
5198        let candidates = vec![
5199            cand("m1", 53.748, 50.0),
5200            cand("m2", 41.605, 51.0),
5201            cand("m3", 120.011, 65.0),
5202        ];
5203        let cmp = compare_reml_fits(candidates).expect("comparison");
5204
5205        assert_eq!(cmp.winner, "m1", "AIC winner");
5206        // No ranking row may contradict the declared winner.
5207        for row in &cmp.ranking {
5208            assert!(
5209                row.delta >= 0.0,
5210                "ranking delta for {} must be >= 0, got {}",
5211                row.name,
5212                row.delta
5213            );
5214            assert!(
5215                row.bayes_factor >= 1.0 - 1e-12,
5216                "ranking bayes_factor for {} must be >= 1, got {}",
5217                row.name,
5218                row.bayes_factor
5219            );
5220        }
5221        let winner_row = cmp.ranking.iter().find(|r| r.name == "m1").unwrap();
5222        assert!(winner_row.delta.abs() < 1e-12, "winner delta == 0");
5223        assert!(
5224            (winner_row.bayes_factor - 1.0).abs() < 1e-9,
5225            "winner bayes_factor == 1"
5226        );
5227
5228        // The raw-REML score table is referenced to the genuine minimum raw REML
5229        // (m2), so its best-over-model Bayes factors are also coherent (>= 1).
5230        for row in &cmp.score_table {
5231            assert!(
5232                row.delta_reml >= 0.0,
5233                "score-table delta_reml for {} must be >= 0, got {}",
5234                row.name,
5235                row.delta_reml
5236            );
5237            assert!(
5238                row.bayes_factor_best_over_model >= 1.0 - 1e-12,
5239                "score-table bayes_factor for {} must be >= 1, got {}",
5240                row.name,
5241                row.bayes_factor_best_over_model
5242            );
5243        }
5244        // m2 carries the minimum raw REML, so its raw delta is exactly 0.
5245        let m2 = cmp.score_table.iter().find(|r| r.name == "m2").unwrap();
5246        assert!(
5247            m2.delta_reml.abs() < 1e-12,
5248            "the minimum-raw-REML row has delta_reml 0"
5249        );
5250    }
5251
5252    #[test]
5253    fn cone_of_influence_empty_support_is_empty() {
5254        let labels = vec![0usize, 0, 1, 1];
5255        assert!(cone_of_influence(&labels, &[]).is_empty());
5256    }
5257
5258    #[test]
5259    fn cone_of_influence_returns_full_component() {
5260        let labels = vec![0usize, 0, 1, 1];
5261        // Support in component 0 -> cone is {0,1}.
5262        assert_eq!(cone_of_influence(&labels, &[0]), vec![0, 1]);
5263        // Support spanning both -> cone is everything.
5264        assert_eq!(cone_of_influence(&labels, &[1, 2]), vec![0, 1, 2, 3]);
5265    }
5266
5267    #[test]
5268    fn coned_matches_full_solve_on_fully_coupled_hessian() {
5269        // Fully coupled SPD H: cone is the whole space, result must equal the
5270        // unconfined sensitivity-operator mode response bit-for-bit.
5271        let h = Array2::from_shape_vec((3, 3), vec![4.0, 1.0, 0.5, 1.0, 3.0, 0.8, 0.5, 0.8, 2.5])
5272            .unwrap();
5273        let inv = dense_inverse(&h);
5274        // Two ρ-columns, each supported on a single coefficient.
5275        let mut dg = Array2::<f64>::zeros((3, 2));
5276        dg[[0, 0]] = 1.3;
5277        dg[[2, 1]] = -0.7;
5278        let supports = vec![0..1usize, 2..3usize];
5279
5280        let eye: Array2<f64> = Array2::eye(3);
5281        let op = crate::sensitivity::FitSensitivity::from_projected(&eye, &inv);
5282        let full = op.mode_response(dg.view()).unwrap();
5283        let coned = op
5284            .mode_response_coned(h.view(), dg.view(), &supports)
5285            .unwrap();
5286        for i in 0..3 {
5287            for a in 0..2 {
5288                assert!(
5289                    (full[[i, a]] - coned[[i, a]]).abs() < 1e-12,
5290                    "fully-coupled mismatch at ({i},{a}): {} vs {}",
5291                    full[[i, a]],
5292                    coned[[i, a]]
5293                );
5294            }
5295        }
5296    }
5297
5298    #[test]
5299    fn coned_confines_to_component_on_decoupled_hessian() {
5300        // Block-decoupled H: blocks {0,1} and {2,3}. A column supported only in
5301        // block {0,1} must produce sensitivity zero in block {2,3}, and match
5302        // the exact solution within its own block.
5303        let mut h = Array2::<f64>::zeros((4, 4));
5304        // Block A.
5305        h[[0, 0]] = 4.0;
5306        h[[1, 1]] = 3.0;
5307        h[[0, 1]] = 1.0;
5308        h[[1, 0]] = 1.0;
5309        // Block B.
5310        h[[2, 2]] = 2.0;
5311        h[[3, 3]] = 5.0;
5312        h[[2, 3]] = 0.6;
5313        h[[3, 2]] = 0.6;
5314        let inv = dense_inverse(&h);
5315
5316        let mut dg = Array2::<f64>::zeros((4, 1));
5317        dg[[0, 0]] = 0.9;
5318        dg[[1, 0]] = -0.4;
5319        let support_range = 0..2usize;
5320        let supports = std::slice::from_ref(&support_range);
5321
5322        let eye: Array2<f64> = Array2::eye(4);
5323        let coned = crate::sensitivity::FitSensitivity::from_projected(&eye, &inv)
5324            .mode_response_coned(h.view(), dg.view(), supports)
5325            .unwrap();
5326        // Exact reference: -H⁻¹ q. Off-block entries are exactly zero already
5327        // (decoupled inverse), and the cone must preserve the in-block ones.
5328        let q = dg.column(0).to_owned();
5329        let exact = inv.dot(&q).mapv(|v| -v);
5330        for i in 0..4 {
5331            assert!(
5332                (coned[[i, 0]] - exact[[i]]).abs() < 1e-12,
5333                "decoupled mismatch at {i}: {} vs {}",
5334                coned[[i, 0]],
5335                exact[[i]]
5336            );
5337        }
5338        // Block B is outside the cone -> exactly zero.
5339        assert_eq!(coned[[2, 0]], 0.0);
5340        assert_eq!(coned[[3, 0]], 0.0);
5341    }
5342
5343    #[test]
5344    fn coned_skips_inactive_column_with_empty_support() {
5345        let h = Array2::<f64>::eye(2);
5346        let dg = Array2::<f64>::zeros((2, 1));
5347        // Inactive ρ: empty support, must be skipped without solving.
5348        let empty_support = 0..0usize;
5349        let supports = std::slice::from_ref(&empty_support);
5350        // A NaN inverse: an empty-support column must be skipped WITHOUT
5351        // solving, so the operator's finite-check never sees the NaN and the
5352        // result is `Some(zeros)`. Were the inactive column ever solved, the
5353        // NaN would propagate and `mode_response_coned` would return `None`.
5354        let eye: Array2<f64> = Array2::eye(2);
5355        let nan_inv = Array2::<f64>::from_elem((2, 2), f64::NAN);
5356        let coned = crate::sensitivity::FitSensitivity::from_projected(&eye, &nan_inv)
5357            .mode_response_coned(h.view(), dg.view(), supports)
5358            .unwrap();
5359        assert_eq!(coned[[0, 0]], 0.0);
5360        assert_eq!(coned[[1, 0]], 0.0);
5361    }
5362
5363    fn make_minimal_cache() -> ArrowFactorCache {
5364        // d = 1, k = 1, n = 1, H_uu_1 = [[2.0]] => L = [[sqrt(2)]],
5365        // H_uβ_1 = [[0.5]], A = 2 - 0.5 * 0.5 / 2 = 1.875.
5366        let l_huu = Array2::from_shape_vec((1, 1), vec![std::f64::consts::SQRT_2]).unwrap();
5367        let l_schur = Array2::from_shape_vec((1, 1), vec![(1.875_f64).sqrt()]).unwrap();
5368        let htbeta = Array2::from_shape_vec((1, 1), vec![0.5]).unwrap();
5369        let mut cache = ArrowFactorCache {
5370            htt_factors: ArrowFactorSlab::from_blocks(vec![l_huu]),
5371            htt_factors_undamped: crate::arrow_schur::ArrowUndampedFactors::SameAsDamped,
5372            schur_factor: Some(l_schur),
5373            schur_factor_is_undamped: true,
5374            beta_schur_deflation: None,
5375            joint_hessian_log_det: None,
5376            solver_mode: crate::arrow_schur::ArrowSolverMode::Direct,
5377            ridge_t: 0.0,
5378            ridge_beta: 0.0,
5379            htbeta: crate::arrow_schur::ArrowHtbetaCache::Dense {
5380                blocks: std::sync::Arc::from(vec![htbeta]),
5381                estimated_bytes: std::mem::size_of::<f64>(),
5382            },
5383            d: 1,
5384            row_dims: std::sync::Arc::from(vec![1usize]),
5385            row_offsets: std::sync::Arc::from(vec![0usize, 1usize]),
5386            k: 1,
5387            manifold_mode_fingerprint: 0,
5388            row_hessian_fingerprint: 0,
5389            pcg_diagnostics: crate::arrow_schur::ArrowPcgDiagnostics::default(),
5390            gauge_deflated_directions: 0,
5391            deflated_row_directions: std::sync::Arc::from(Vec::new()),
5392            deflation_row_spectra: std::sync::Arc::from(Vec::new()),
5393            beta_gauge_quotient: None,
5394        };
5395        cache.joint_hessian_log_det = cache.compute_undamped_arrow_log_det();
5396        cache
5397    }
5398
5399    #[test]
5400    fn laplace_evidence_returns_finite_for_minimal_cache() {
5401        let cache = make_minimal_cache();
5402        // log|H| = log(2) + log(1.875). With dim(H)=2 and rank(S)=1,
5403        // V includes the rank-aware TK nullspace normalizer.
5404        let v = laplace_evidence(
5405            EvidenceLogDetSource::FactoredArrow {
5406                cache: &cache,
5407                fallback_hvp: None,
5408            },
5409            0.0,
5410            0.0,
5411            2.0,
5412            1.0,
5413        );
5414        assert!(v.is_finite());
5415        let expected =
5416            0.5 * (2.0_f64.ln() + 1.875_f64.ln()) - 0.5 * (2.0 * std::f64::consts::PI).ln();
5417        assert!((v - expected).abs() < 1e-12);
5418    }
5419
5420    /// #1132 bug 2: a β-profiled atom (no shared `β` block, `k == 0`) reaches
5421    /// `arrow_log_det_from_cache` in the dense Direct path with
5422    /// `schur_factor = None` — there is no reduced Schur complement to form. The
5423    /// joint Hessian is then block-diagonal in the latent rows, so its log-det
5424    /// is exactly the per-row sum with NO Schur term. Before the fix this
5425    /// returned `None` (the `schur_factor.as_ref()?` bail), starving the REML
5426    /// Laplace normaliser and erroring "arrow_log_det_from_cache returned None
5427    /// at ridge=0 Direct mode". Now it returns `Some(Σ_i log|H_tt^(i)|)`.
5428    fn k0_direct_cache_no_schur(latent_diag: f64) -> ArrowFactorCache {
5429        let l_huu = Array2::from_shape_vec((1, 1), vec![latent_diag.sqrt()]).unwrap();
5430        let mut cache = ArrowFactorCache {
5431            htt_factors: ArrowFactorSlab::from_blocks(vec![l_huu]),
5432            htt_factors_undamped: crate::arrow_schur::ArrowUndampedFactors::SameAsDamped,
5433            schur_factor: None,
5434            schur_factor_is_undamped: true,
5435            beta_schur_deflation: None,
5436            joint_hessian_log_det: None,
5437            solver_mode: crate::arrow_schur::ArrowSolverMode::Direct,
5438            ridge_t: 0.0,
5439            ridge_beta: 0.0,
5440            htbeta: crate::arrow_schur::ArrowHtbetaCache::Disabled { estimated_bytes: 0 },
5441            d: 1,
5442            row_dims: std::sync::Arc::from(vec![1usize]),
5443            row_offsets: std::sync::Arc::from(vec![0usize, 1usize]),
5444            k: 0,
5445            manifold_mode_fingerprint: 0,
5446            row_hessian_fingerprint: 0,
5447            pcg_diagnostics: crate::arrow_schur::ArrowPcgDiagnostics::default(),
5448            gauge_deflated_directions: 0,
5449            deflated_row_directions: std::sync::Arc::from(Vec::new()),
5450            deflation_row_spectra: std::sync::Arc::from(Vec::new()),
5451            beta_gauge_quotient: None,
5452        };
5453        cache.joint_hessian_log_det = cache.compute_undamped_arrow_log_det();
5454        cache
5455    }
5456
5457    #[test]
5458    fn arrow_log_det_some_for_k0_direct_cache_without_schur() {
5459        let cache = k0_direct_cache_no_schur(3.0);
5460        let log_det = arrow_log_det_from_cache(&cache)
5461            .expect("k==0 Direct cache must yield Some(per-row sum), not None (#1132)");
5462        // Single latent block H_tt = [[3.0]]; no Schur term for k == 0.
5463        assert!(
5464            (log_det - 3.0_f64.ln()).abs() < 1e-12,
5465            "log_det = {log_det}"
5466        );
5467        // The cache's own computation must agree bit-for-bit.
5468        let cached = cache
5469            .compute_undamped_arrow_log_det()
5470            .expect("compute_undamped_arrow_log_det must be Some for k==0");
5471        assert!((cached - 3.0_f64.ln()).abs() < 1e-12, "cached = {cached}");
5472    }
5473
5474    #[test]
5475    fn arrow_log_det_none_for_kpos_cache_without_schur() {
5476        // k > 0 but no dense Schur factor is the genuine InexactPCG case and
5477        // must still reject (the guard must not over-broaden to all `None`).
5478        let mut cache = k0_direct_cache_no_schur(3.0);
5479        cache.k = 1;
5480        cache.solver_mode = crate::arrow_schur::ArrowSolverMode::InexactPCG;
5481        cache.joint_hessian_log_det = None;
5482        assert!(arrow_log_det_from_cache(&cache).is_none());
5483        assert!(cache.compute_undamped_arrow_log_det().is_none());
5484    }
5485
5486    #[test]
5487    fn laplace_evidence_nan_when_authoritative_logdet_missing() {
5488        let mut cache = make_minimal_cache();
5489        cache.ridge_t = 1e-3;
5490        cache.joint_hessian_log_det = None;
5491        assert!(
5492            laplace_evidence(
5493                EvidenceLogDetSource::FactoredArrow {
5494                    cache: &cache,
5495                    fallback_hvp: None,
5496                },
5497                0.0,
5498                0.0,
5499                2.0,
5500                1.0,
5501            )
5502            .is_nan()
5503        );
5504    }
5505
5506    #[test]
5507    fn laplace_evidence_uses_hvp_fallback_without_authoritative_logdet() {
5508        let mut cache = make_minimal_cache();
5509        cache.schur_factor = None;
5510        cache.joint_hessian_log_det = None;
5511        let hvp = |x: &[f64]| -> Vec<f64> { vec![2.0 * x[0], 1.875 * x[1]] };
5512        let v = laplace_evidence(
5513            EvidenceLogDetSource::FactoredArrow {
5514                cache: &cache,
5515                fallback_hvp: Some(EvidenceHvpLogDet {
5516                    dim: 2,
5517                    apply: &hvp,
5518                }),
5519            },
5520            0.0,
5521            0.0,
5522            2.0,
5523            1.0,
5524        );
5525        let expected =
5526            0.5 * (2.0_f64.ln() + 1.875_f64.ln()) - 0.5 * (2.0 * std::f64::consts::PI).ln();
5527        assert!((v - expected).abs() < 1e-12);
5528    }
5529
5530    #[test]
5531    fn ift_du_dbeta_has_expected_shape() {
5532        let cache = make_minimal_cache();
5533        let du_db = ift_du_dbeta(&cache);
5534        assert_eq!(du_db.shape(), &[1, 1]);
5535        // ∂u/∂β = -H_uu⁻¹ H_uβ = -0.5 / 2 = -0.25.
5536        assert!((du_db[[0, 0]] - (-0.25)).abs() < 1e-12);
5537    }
5538
5539    #[test]
5540    fn ift_dbeta_drho_returns_some_for_direct_cache() {
5541        let cache = make_minimal_cache();
5542        let q = Array2::from_shape_vec((1, 1), vec![1.0]).unwrap();
5543        let out = ift_dbeta_drho(&cache, q.view()).unwrap();
5544        assert_eq!(out.shape(), &[1, 1]);
5545        // ∂β/∂ρ = -A⁻¹ · 1 = -1/1.875.
5546        assert!((out[[0, 0]] + 1.0 / 1.875).abs() < 1e-12);
5547    }
5548
5549    #[test]
5550    fn topology_select_picks_lowest_negative_log_evidence() {
5551        let candidates = vec![
5552            TopologyCandidate {
5553                kind: TopologyKind::Flat,
5554                negative_log_evidence: 10.0,
5555                effective_dim: 4.0,
5556                n_obs: 100,
5557                converged: true,
5558                exclusion_reason: None,
5559            },
5560            TopologyCandidate {
5561                kind: TopologyKind::Sphere,
5562                negative_log_evidence: 8.0,
5563                effective_dim: 5.0,
5564                n_obs: 100,
5565                converged: true,
5566                exclusion_reason: None,
5567            },
5568            TopologyCandidate {
5569                kind: TopologyKind::Torus,
5570                negative_log_evidence: f64::NAN,
5571                effective_dim: 6.0,
5572                n_obs: 100,
5573                converged: false,
5574                exclusion_reason: Some("torus periods missing".to_string()),
5575            },
5576        ];
5577        let sel = select_topology(&candidates, TopologySelectOptions::default());
5578        assert_eq!(sel.winner, TopologyKind::Sphere);
5579        assert!(!sel.tie);
5580    }
5581
5582    #[test]
5583    fn topology_select_tie_breaks_to_simpler() {
5584        let candidates = vec![
5585            TopologyCandidate {
5586                kind: TopologyKind::Sphere,
5587                negative_log_evidence: 5.0,
5588                effective_dim: 5.0,
5589                n_obs: 100,
5590                converged: true,
5591                exclusion_reason: None,
5592            },
5593            TopologyCandidate {
5594                kind: TopologyKind::Flat,
5595                negative_log_evidence: 5.0 + 1e-6,
5596                effective_dim: 4.0,
5597                n_obs: 100,
5598                converged: true,
5599                exclusion_reason: None,
5600            },
5601        ];
5602        let sel = select_topology(&candidates, TopologySelectOptions::default());
5603        assert_eq!(sel.winner, TopologyKind::Flat);
5604        assert!(sel.tie);
5605    }
5606
5607    fn gaussian_logpdf(y: f64, mean: f64, sd: f64) -> f64 {
5608        let z = (y - mean) / sd;
5609        -0.5 * (2.0 * std::f64::consts::PI).ln() - sd.ln() - 0.5 * z * z
5610    }
5611
5612    #[test]
5613    fn stacking_single_candidate_gets_full_weight() {
5614        let log_density = Array2::from_shape_vec((3, 1), vec![-1.0, -2.0, -0.5]).unwrap();
5615        let out = solve_stacking_weights(log_density.view(), StackingConfig::default()).unwrap();
5616        assert!((out.weights[0] - 1.0).abs() < 1e-12);
5617        assert_eq!(out.weights.len(), 1);
5618    }
5619
5620    #[test]
5621    fn stacking_dominant_candidate_attracts_nearly_all_weight() {
5622        let mut log_density = Array2::<f64>::zeros((50, 2));
5623        for i in 0..50 {
5624            log_density[[i, 0]] = -0.1;
5625            log_density[[i, 1]] = -5.0;
5626        }
5627        let out = solve_stacking_weights(log_density.view(), StackingConfig::default()).unwrap();
5628        assert!(out.weights[0] > 0.99, "w0 = {}", out.weights[0]);
5629        assert!(out.weights[1] < 0.01, "w1 = {}", out.weights[1]);
5630    }
5631
5632    #[test]
5633    fn stacking_complementary_candidates_share_weight() {
5634        // Each candidate is the better predictor on its own half of the data;
5635        // stacking keeps both, unlike winner-take-all.
5636        let n = 40;
5637        let mut log_density = Array2::<f64>::zeros((n, 2));
5638        for i in 0..n {
5639            if i < n / 2 {
5640                log_density[[i, 0]] = gaussian_logpdf(0.0, 0.0, 0.5);
5641                log_density[[i, 1]] = gaussian_logpdf(0.0, 1.5, 0.5);
5642            } else {
5643                log_density[[i, 0]] = gaussian_logpdf(0.0, 1.5, 0.5);
5644                log_density[[i, 1]] = gaussian_logpdf(0.0, 0.0, 0.5);
5645            }
5646        }
5647        let out = solve_stacking_weights(log_density.view(), StackingConfig::default()).unwrap();
5648        assert!(
5649            out.weights[0] > 0.2 && out.weights[0] < 0.8,
5650            "w0 = {}",
5651            out.weights[0]
5652        );
5653        assert!((out.weights.sum() - 1.0).abs() < 1e-9);
5654    }
5655
5656    #[test]
5657    fn stacking_weights_stay_on_the_simplex() {
5658        let log_density = Array2::from_shape_vec(
5659            (3, 3),
5660            vec![-1.0, -2.0, -3.0, -2.5, -1.0, -2.0, -3.0, -2.0, -1.0],
5661        )
5662        .unwrap();
5663        let out = solve_stacking_weights(log_density.view(), StackingConfig::default()).unwrap();
5664        assert!((out.weights.sum() - 1.0).abs() < 1e-9);
5665        assert!(out.weights.iter().all(|&w| w >= -1e-12));
5666    }
5667
5668    #[test]
5669    fn stacking_solution_satisfies_the_simplex_kkt_certificate() {
5670        // Recompute the KKT residual from scratch at the returned weights:
5671        // g_k = mean_i p_ik / mix_i must satisfy g_k <= 1 (+tol) everywhere
5672        // and w_k * |g_k - 1| <= tol. The objective is concave, so this
5673        // certifies the GLOBAL optimum, not merely a stationary iterate.
5674        let log_density = Array2::from_shape_vec(
5675            (5, 2),
5676            vec![-0.2, -3.0, -3.0, -0.2, -0.5, -1.5, -1.5, -0.5, -0.1, -2.0],
5677        )
5678        .unwrap();
5679        let config = StackingConfig::default();
5680        let out = solve_stacking_weights(log_density.view(), config).unwrap();
5681        assert!(out.certificate.residual() <= config.kkt_tol);
5682        let n = log_density.nrows();
5683        for k in 0..2 {
5684            let mut g = 0.0_f64;
5685            for i in 0..n {
5686                let mix: f64 = (0..2)
5687                    .map(|c| out.weights[c] * log_density[[i, c]].exp())
5688                    .sum();
5689                g += log_density[[i, k]].exp() / mix;
5690            }
5691            g /= n as f64;
5692            assert!(
5693                g <= 1.0 + config.kkt_tol,
5694                "stationarity violated for candidate {k}: g = {g}"
5695            );
5696            assert!(
5697                out.weights[k] * (g - 1.0).abs() <= config.kkt_tol * (1.0 + 1e-6),
5698                "complementary slackness violated for candidate {k}: w = {}, g = {g}",
5699                out.weights[k]
5700            );
5701        }
5702    }
5703
5704    #[test]
5705    fn stacking_exhaustion_without_certificate_is_an_error_not_weights() {
5706        let log_density = Array2::from_shape_vec(
5707            (6, 3),
5708            vec![
5709                0.0, -2.0, -4.0, -0.4, -0.1, -3.0, -2.0, 0.0, -0.3, -3.0, -1.0, 0.0, -0.2, -2.0,
5710                -0.5, -1.0, -0.3, -2.0,
5711            ],
5712        )
5713        .unwrap();
5714        let config = StackingConfig {
5715            max_iter: 1,
5716            ..StackingConfig::default()
5717        };
5718        let err = solve_stacking_weights(log_density.view(), config).unwrap_err();
5719        let checkpoint = match err {
5720            StackingError::DidNotConverge {
5721                certificate,
5722                checkpoint,
5723                ..
5724            } => {
5725                assert!(certificate.residual() > config.kkt_tol);
5726                assert_eq!(checkpoint.completed_iterations, 1);
5727                checkpoint
5728            }
5729            other => panic!("expected typed stacking exhaustion, got {other}"),
5730        };
5731        let encoded = serde_json::to_string(&checkpoint).unwrap();
5732        let checkpoint: StackingCheckpoint = serde_json::from_str(&encoded).unwrap();
5733        let mut other_density = log_density.clone();
5734        other_density[[0, 0]] += 0.25;
5735        assert!(matches!(
5736            resume_stacking_weights(other_density.view(), StackingConfig::default(), &checkpoint,),
5737            Err(StackingError::InvalidInput { .. })
5738        ));
5739        let resumed =
5740            resume_stacking_weights(log_density.view(), StackingConfig::default(), &checkpoint)
5741                .unwrap();
5742        let uninterrupted =
5743            solve_stacking_weights(log_density.view(), StackingConfig::default()).unwrap();
5744        for (resumed, uninterrupted) in resumed.weights.iter().zip(uninterrupted.weights.iter()) {
5745            assert!((resumed - uninterrupted).abs() <= 1.0e-10);
5746        }
5747    }
5748
5749    #[test]
5750    fn stacking_near_tied_boundary_uses_newton_not_millions_of_em_steps() {
5751        let log_density =
5752            Array2::from_shape_fn(
5753                (64, 2),
5754                |(_, candidate)| {
5755                    if candidate == 0 { 0.0 } else { -1.0e-6 }
5756                },
5757            );
5758        let out = solve_stacking_weights(log_density.view(), StackingConfig::default()).unwrap();
5759        assert!(out.weights[0] >= 1.0 - StackingConfig::default().kkt_tol);
5760        assert!(out.iterations < 8, "iterations = {}", out.iterations);
5761    }
5762
5763    #[test]
5764    fn stacking_dead_candidate_column_gets_zero_weight() {
5765        let log_density = Array2::from_shape_vec(
5766            (3, 2),
5767            vec![
5768                -1.0,
5769                f64::NEG_INFINITY,
5770                -2.0,
5771                f64::NEG_INFINITY,
5772                -0.5,
5773                f64::NEG_INFINITY,
5774            ],
5775        )
5776        .unwrap();
5777        let out = solve_stacking_weights(log_density.view(), StackingConfig::default()).unwrap();
5778        assert_eq!(out.weights[1], 0.0);
5779        assert!((out.weights[0] - 1.0).abs() < 1e-12);
5780    }
5781
5782    #[test]
5783    fn stacking_rejects_invalid_and_unscorable_rows() {
5784        let log_density = Array2::from_shape_vec(
5785            (3, 2),
5786            vec![-1.0, -2.0, f64::NAN, f64::NEG_INFINITY, -2.0, -1.0],
5787        )
5788        .unwrap();
5789        assert!(matches!(
5790            solve_stacking_weights(log_density.view(), StackingConfig::default()),
5791            Err(StackingError::InvalidInput { .. })
5792        ));
5793        let unscorable = Array2::from_shape_vec(
5794            (2, 2),
5795            vec![-1.0, -2.0, f64::NEG_INFINITY, f64::NEG_INFINITY],
5796        )
5797        .unwrap();
5798        assert!(matches!(
5799            solve_stacking_weights(unscorable.view(), StackingConfig::default()),
5800            Err(StackingError::InvalidInput { .. })
5801        ));
5802    }
5803
5804    fn two_cluster_mixture_data() -> Array2<f64> {
5805        Array2::from_shape_vec(
5806            (12, 1),
5807            vec![
5808                -2.2, -2.0, -1.9, -2.1, -1.8, -2.05, 1.8, 2.0, 2.2, 1.9, 2.1, 2.05,
5809            ],
5810        )
5811        .unwrap()
5812    }
5813
5814    #[test]
5815    fn gaussian_mixture_monotonicity_resolves_composite_map_noise_2264() {
5816        let objective_scale = 1.0;
5817        let composite_resolution = f64::EPSILON.sqrt() * objective_scale;
5818        let uncertainty = gaussian_mixture_monotonicity_uncertainty(objective_scale, 0.0, 0.0);
5819        assert_eq!(uncertainty, composite_resolution);
5820
5821        let noise_scale_decrease = -0.5 * composite_resolution;
5822        assert!(noise_scale_decrease >= -uncertainty);
5823        let resolved_decrease = -2.0 * composite_resolution;
5824        assert!(resolved_decrease < -uncertainty);
5825
5826        let larger_reduction_bound = 2.0 * composite_resolution;
5827        assert_eq!(
5828            gaussian_mixture_monotonicity_uncertainty(objective_scale, larger_reduction_bound, 0.0),
5829            larger_reduction_bound,
5830        );
5831    }
5832
5833    #[test]
5834    fn gaussian_mixture_issue_scale_negative_step_is_within_computed_uncertainty_2264() {
5835        // Recorded issue mechanism: a signed mean-log-likelihood step of
5836        // -1.4e-13 was rejected even though it is unresolved at the scale of
5837        // the composite EM map. The admissible decrease below is derived only
5838        // from that map's observed objective scale and arithmetic reduction
5839        // bounds; the recorded step is an input to the decision, not a new
5840        // tolerance.
5841        let objective_scale = 1.0;
5842        let recorded_step = -1.4e-13;
5843        let uncertainty = gaussian_mixture_monotonicity_uncertainty(objective_scale, 0.0, 0.0);
5844        let certificate = GaussianMixtureCertificate {
5845            mean_log_likelihood: -objective_scale,
5846            mean_log_likelihood_gain: recorded_step,
5847            monotonicity_uncertainty: uncertainty,
5848            objective_residual: recorded_step.abs() / objective_scale,
5849            objective_tolerance: f64::EPSILON.sqrt(),
5850            parameter_residual: 0.0,
5851            parameter_tolerance: f64::EPSILON.sqrt(),
5852            contraction_rate: None,
5853            projected_iterations_to_tolerance: None,
5854        };
5855
5856        assert_eq!(
5857            certificate.monotonicity_uncertainty,
5858            f64::EPSILON.sqrt() * objective_scale,
5859            "reported uncertainty must be the computed composite-map resolution"
5860        );
5861        assert!(
5862            certificate.mean_log_likelihood_gain >= -certificate.monotonicity_uncertainty,
5863            "the recorded noise-scale decrease must not be a monotonicity violation"
5864        );
5865    }
5866
5867    #[test]
5868    fn gaussian_mixture_below_roundoff_positive_gain_can_certify_2264() {
5869        // Recorded issue mechanism: a +6.6e-15 gain with a 1.5e-14 reduction
5870        // bound exhausted instead of certifying. A below-resolution gain is a
5871        // valid objective fixed point only when the independent parameter-map
5872        // residual also clears its configured tolerance.
5873        let objective_scale = 1.0;
5874        let recorded_gain = 6.6e-15;
5875        let recorded_reduction_bound = 1.5e-14;
5876        let objective_tolerance = f64::EPSILON.sqrt();
5877        let parameter_tolerance = f64::EPSILON.sqrt();
5878        let uncertainty = gaussian_mixture_monotonicity_uncertainty(
5879            objective_scale,
5880            recorded_reduction_bound,
5881            0.0,
5882        );
5883        let certificate = GaussianMixtureCertificate {
5884            mean_log_likelihood: -objective_scale,
5885            mean_log_likelihood_gain: recorded_gain,
5886            monotonicity_uncertainty: uncertainty,
5887            objective_residual: recorded_gain / objective_scale,
5888            objective_tolerance,
5889            parameter_residual: 0.5 * parameter_tolerance,
5890            parameter_tolerance,
5891            contraction_rate: None,
5892            projected_iterations_to_tolerance: None,
5893        };
5894
5895        assert_eq!(
5896            certificate.monotonicity_uncertainty,
5897            (f64::EPSILON.sqrt() * objective_scale).max(recorded_reduction_bound),
5898            "reported uncertainty must come from the composite-map and reduction bounds"
5899        );
5900        assert!(certificate.mean_log_likelihood_gain >= -certificate.monotonicity_uncertainty);
5901        assert!(certificate.objective_residual <= certificate.objective_tolerance);
5902        assert!(certificate.parameter_residual <= certificate.parameter_tolerance);
5903    }
5904
5905    #[test]
5906    fn gaussian_mixture_certificate_quotients_duplicate_component_mass_exchange_2324() {
5907        // Two identical components can exchange arbitrary mass without
5908        // changing the mixture density. Labeled component measures are
5909        // therefore not identifiable even when every component has positive
5910        // mass; the empirical predictive-density certificate must quotient
5911        // this singular direction exactly.
5912        let data = array![[-1.0], [0.0], [2.0]];
5913        let means = array![[0.0], [0.0]];
5914        let covariance = vec![array![[1.0]], array![[1.0]]];
5915        let weights = array![0.25, 0.75];
5916        let redistributed_weights = array![0.5, 0.5];
5917        let previous = mixture_e_step(data.view(), &weights, &means, &covariance).unwrap();
5918        let redistributed =
5919            mixture_e_step(data.view(), &redistributed_weights, &means, &covariance).unwrap();
5920        let residual = empirical_predictive_density_residual(
5921            &previous.row_log_likelihoods,
5922            &redistributed.row_log_likelihoods,
5923        )
5924        .unwrap();
5925        assert!(residual <= 4.0 * f64::EPSILON);
5926
5927        // A resolved change in the represented density remains visible.
5928        let shifted_means = array![[0.01], [0.0]];
5929        let shifted =
5930            mixture_e_step(data.view(), &weights, &shifted_means, &covariance).unwrap();
5931        let shifted_residual = empirical_predictive_density_residual(
5932            &previous.row_log_likelihoods,
5933            &shifted.row_log_likelihoods,
5934        )
5935        .unwrap();
5936        assert!(shifted_residual > f64::EPSILON.sqrt());
5937    }
5938
5939    #[test]
5940    fn gaussian_mixture_fit_certificate_describes_the_exact_returned_iterate() {
5941        let data = two_cluster_mixture_data();
5942        let config = GaussianMixtureConfig::default();
5943        let fit = fit_gaussian_mixture(data.view(), 2, config).unwrap();
5944        let certificate = fit.certificate();
5945        assert!(certificate.objective_residual <= certificate.objective_tolerance);
5946        assert!(certificate.parameter_residual <= certificate.parameter_tolerance);
5947
5948        let checkpoint = GaussianMixtureCheckpoint {
5949            weights: fit.weights.clone(),
5950            means: fit.means.clone(),
5951            covariances: fit.covariances.clone(),
5952            mean_log_likelihood: certificate.mean_log_likelihood,
5953            completed_iterations: fit.iterations,
5954            data_fingerprint: mixture_data_fingerprint(data.view()),
5955            covariance_floor: config.covariance_floor,
5956        };
5957        let current = mixture_e_step(
5958            data.view(),
5959            &checkpoint.weights,
5960            &checkpoint.means,
5961            &checkpoint.covariances,
5962        )
5963        .unwrap();
5964        let (weights, means, covariances) = mixture_m_step(
5965            data.view(),
5966            current.responsibilities.view(),
5967            config.covariance_floor,
5968        )
5969        .unwrap();
5970        let next = mixture_e_step(data.view(), &weights, &means, &covariances).unwrap();
5971        let residual = empirical_predictive_density_residual(
5972            &current.row_log_likelihoods,
5973            &next.row_log_likelihoods,
5974        )
5975        .unwrap();
5976        assert!(residual <= config.parameter_tol);
5977        assert_eq!(certificate.mean_log_likelihood, current.mean_log_likelihood);
5978        assert_eq!(
5979            certificate.mean_log_likelihood_gain,
5980            next.mean_log_likelihood - current.mean_log_likelihood
5981        );
5982        assert_eq!(
5983            certificate.monotonicity_uncertainty,
5984            gaussian_mixture_monotonicity_uncertainty(
5985                current
5986                    .mean_log_likelihood
5987                    .abs()
5988                    .max(next.mean_log_likelihood.abs())
5989                    .max(1.0),
5990                current.mean_log_likelihood_roundoff,
5991                next.mean_log_likelihood_roundoff,
5992            )
5993        );
5994        assert_eq!(certificate.parameter_residual, residual);
5995        assert!(
5996            (next.mean_log_likelihood - current.mean_log_likelihood).abs()
5997                / current
5998                    .mean_log_likelihood
5999                    .abs()
6000                    .max(next.mean_log_likelihood.abs())
6001                    .max(1.0)
6002                <= config.loglik_tol
6003        );
6004    }
6005
6006    #[test]
6007    fn gaussian_mixture_exhaustion_is_typed_and_resumable() {
6008        let data = two_cluster_mixture_data();
6009        let short = GaussianMixtureConfig {
6010            max_iter: 1,
6011            ..GaussianMixtureConfig::default()
6012        };
6013        let err = fit_gaussian_mixture(data.view(), 2, short).unwrap_err();
6014        let checkpoint = match err {
6015            GaussianMixtureError::DidNotConverge {
6016                certificate,
6017                checkpoint,
6018                ..
6019            } => {
6020                assert!(
6021                    certificate.objective_residual > short.loglik_tol
6022                        || certificate.parameter_residual > short.parameter_tol
6023                );
6024                assert_eq!(checkpoint.completed_iterations, 1);
6025                let at_checkpoint = mixture_e_step(
6026                    data.view(),
6027                    &checkpoint.weights,
6028                    &checkpoint.means,
6029                    &checkpoint.covariances,
6030                )
6031                .unwrap();
6032                assert_eq!(
6033                    certificate.mean_log_likelihood, at_checkpoint.mean_log_likelihood,
6034                    "exhaustion evidence and checkpoint must describe one iterate"
6035                );
6036                checkpoint
6037            }
6038            other => panic!("expected typed EM exhaustion, got {other}"),
6039        };
6040        let encoded = serde_json::to_string(&checkpoint).unwrap();
6041        let checkpoint: GaussianMixtureCheckpoint = serde_json::from_str(&encoded).unwrap();
6042        let mut other_data = data.clone();
6043        other_data[[0, 0]] += 0.01;
6044        assert!(matches!(
6045            resume_gaussian_mixture(
6046                other_data.view(),
6047                GaussianMixtureConfig::default(),
6048                checkpoint.clone(),
6049            ),
6050            Err(GaussianMixtureError::InvalidInput { .. })
6051        ));
6052        let resumed =
6053            resume_gaussian_mixture(data.view(), GaussianMixtureConfig::default(), checkpoint)
6054                .unwrap();
6055        let uninterrupted =
6056            fit_gaussian_mixture(data.view(), 2, GaussianMixtureConfig::default()).unwrap();
6057        for (resumed, uninterrupted) in resumed.weights.iter().zip(uninterrupted.weights.iter()) {
6058            assert!((resumed - uninterrupted).abs() <= 1.0e-10);
6059        }
6060
6061        assert!(resumed.bic().is_finite());
6062    }
6063
6064    #[test]
6065    fn gaussian_mixture_bic_is_finite_with_an_active_covariance_floor() {
6066        // Component zero is exactly one-dimensional: its x coordinate never
6067        // changes, so the constrained MLE has one covariance eigenvalue at the
6068        // configured floor. The old BHHH determinant had identically-zero
6069        // mean-x and covariance-xy score columns and therefore rejected this
6070        // perfectly valid constrained predictive density as non-SPD.
6071        let per_cluster = 45usize;
6072        let mut data = Array2::<f64>::zeros((2 * per_cluster, 2));
6073        for sample in 0..per_cluster {
6074            let phase = std::f64::consts::TAU * sample as f64 / per_cluster as f64;
6075            data[[2 * sample, 0]] = -2.0;
6076            data[[2 * sample, 1]] = 0.08 * phase.sin();
6077            data[[2 * sample + 1, 0]] = 2.0 + 0.12 * phase.cos();
6078            data[[2 * sample + 1, 1]] = 0.08 * phase.sin();
6079        }
6080        let fit = fit_gaussian_mixture(data.view(), 2, GaussianMixtureConfig::default())
6081            .expect("the covariance floor defines a valid constrained mixture fit");
6082        let bic = fit.bic();
6083        assert!(bic.is_finite());
6084        assert_eq!(
6085            bic,
6086            -fit.loglik + 0.5 * fit.num_free_parameters() as f64 * (data.nrows() as f64).ln()
6087        );
6088    }
6089
6090    fn seven_clusters_on_a_circle_2262() -> Array2<f64> {
6091        let clusters = 7usize;
6092        let per_cluster = 32usize;
6093        let mut data = Array2::<f64>::zeros((clusters * per_cluster, 2));
6094        for cluster in 0..clusters {
6095            let angle = std::f64::consts::TAU * cluster as f64 / clusters as f64;
6096            let (sin_angle, cos_angle) = angle.sin_cos();
6097            for sample in 0..per_cluster {
6098                let phase = std::f64::consts::TAU * sample as f64 / per_cluster as f64;
6099                // Vary the within-cluster radius while preserving its angular
6100                // symmetry. A literal constant-radius micro-circle makes the
6101                // Gaussian scale score identically zero and its empirical
6102                // Fisher singular, which is not a Gaussian-cluster fixture.
6103                let local_radius = 0.035 * (1.0 + 0.3 * (3.0 * phase).cos());
6104                let radial_noise = local_radius * phase.cos();
6105                let tangent_noise = local_radius * phase.sin();
6106                let radius = 2.0 + radial_noise;
6107                let row = cluster * per_cluster + sample;
6108                data[[row, 0]] = 0.4 + radius * cos_angle - tangent_noise * sin_angle;
6109                data[[row, 1]] = -0.3 + radius * sin_angle + tangent_noise * cos_angle;
6110            }
6111        }
6112        data
6113    }
6114
6115    fn two_noisy_circles_for_union() -> Array2<f64> {
6116        let rows_per_circle = 96usize;
6117        let mut data = Array2::<f64>::zeros((2 * rows_per_circle, 2));
6118        for (circle, (center, radius)) in [([-4.0_f64, 0.3_f64], 1.2_f64), ([4.0, -0.2], 0.9)]
6119            .into_iter()
6120            .enumerate()
6121        {
6122            for sample in 0..rows_per_circle {
6123                let angle = std::f64::consts::TAU * sample as f64 / rows_per_circle as f64;
6124                let noisy_radius =
6125                    radius + 0.045 * (3.0 * angle).cos() + 0.018 * (5.0 * angle).sin();
6126                let row = circle * rows_per_circle + sample;
6127                data[[row, 0]] = center[0] + noisy_radius * angle.cos();
6128                data[[row, 1]] = center[1] + noisy_radius * angle.sin();
6129            }
6130        }
6131        data
6132    }
6133
6134    #[test]
6135    fn circular_gaussian_density_avoids_extreme_scale_intermediate_overflow() {
6136        let noise_variance = f64::MAX / 2.0;
6137        let fit =
6138            CircularGaussianFit2d::from_parameters([0.0, 0.0], 1.1e154, noise_variance).unwrap();
6139        // Both `2πs` and `Rr` overflow if formed directly, although their log
6140        // and the ratio `Rr/s` are representable.
6141        let center_log_density = fit.log_density(0.0, 0.0);
6142        let off_center_log_density = fit.log_density(1.7e154, 0.0);
6143        assert!(center_log_density.is_finite());
6144        assert!(off_center_log_density.is_finite());
6145        let expected_center = -std::f64::consts::TAU.ln()
6146            - noise_variance.ln()
6147            - 0.5 * (fit.radius() / noise_variance.sqrt()).powi(2);
6148        assert_eq!(center_log_density, expected_center);
6149    }
6150
6151    #[test]
6152    fn union_circles_use_the_shared_normalized_cartesian_density() {
6153        let data = two_noisy_circles_for_union();
6154        let config = GaussianMixtureConfig::default();
6155        let density_fit =
6156            fit_union_density(data.view(), UnionStructure::CircleCircle, config).unwrap();
6157        let union = fit_union_structure(data.view(), UnionStructure::CircleCircle, config).unwrap();
6158        assert_eq!(
6159            union.total_parameters,
6160            2 * CircularGaussianFit2d::NUM_FREE_PARAMETERS + 1
6161        );
6162        let component_weight_sum: f64 = union
6163            .components
6164            .iter()
6165            .map(|component| component.mixing_weight)
6166            .sum();
6167        assert!((component_weight_sum - 1.0).abs() <= 8.0 * f64::EPSILON);
6168
6169        let mut fitted_centers = Array2::<f64>::zeros((density_fit.components.len(), 2));
6170        for (index, component) in density_fit.components.iter().enumerate() {
6171            let UnionDensityModel::Circle(fit) = &component.model else {
6172                panic!("circle+circle union produced a non-circle density");
6173            };
6174            let center = fit.center();
6175            fitted_centers[[index, 0]] = center[0];
6176            fitted_centers[[index, 1]] = center[1];
6177            let at_center = fit.log_density(center[0], center[1]);
6178            let expected = -std::f64::consts::TAU.ln()
6179                - fit.noise_variance().ln()
6180                - 0.5 * (fit.radius() / fit.noise_variance().sqrt()).powi(2);
6181            assert!(at_center.is_finite());
6182            assert!((at_center - expected).abs() < 1.0e-12 * (1.0 + expected.abs()));
6183        }
6184
6185        let training_log_density = union_per_point_log_density(
6186            data.view(),
6187            data.view(),
6188            UnionStructure::CircleCircle,
6189            config,
6190        )
6191        .unwrap();
6192        let direct_log_likelihood = pairwise_sum(
6193            training_log_density
6194                .as_slice()
6195                .expect("owned score vector is contiguous"),
6196        );
6197        assert!(
6198            (union.log_likelihood - direct_log_likelihood).abs()
6199                <= 1.0e-12 * (1.0 + direct_log_likelihood.abs())
6200        );
6201        let expected_bic = -direct_log_likelihood
6202            + 0.5 * union.total_parameters as f64 * (data.nrows() as f64).ln();
6203        assert!((union.bic - expected_bic).abs() <= 1.0e-12 * (1.0 + expected_bic.abs()));
6204
6205        let held_out = union_per_point_log_density(
6206            data.view(),
6207            fitted_centers.view(),
6208            UnionStructure::CircleCircle,
6209            config,
6210        )
6211        .unwrap();
6212        assert!(held_out.iter().all(|value| value.is_finite()));
6213    }
6214
6215    fn circle_and_point_union_data() -> (Array2<f64>, Vec<Vec<usize>>) {
6216        let circle_rows = 32usize;
6217        let point_rows = 12usize;
6218        let mut data = Array2::<f64>::zeros((circle_rows + point_rows, 2));
6219        for row in 0..circle_rows {
6220            let angle = std::f64::consts::TAU * row as f64 / circle_rows as f64;
6221            let radius = 1.0 + 0.025 * (3.0 * angle).cos();
6222            data[[row, 0]] = -4.0 + radius * angle.cos();
6223            data[[row, 1]] = 0.2 + radius * angle.sin();
6224        }
6225        for offset in 0..point_rows {
6226            let phase = offset as f64;
6227            let row = circle_rows + offset;
6228            data[[row, 0]] = 4.0 + 0.055 * (1.7 * phase).cos() + 0.018 * (0.4 * phase).sin();
6229            data[[row, 1]] = -0.3 + 0.052 * (1.3 * phase).sin() - 0.015 * (0.9 * phase).cos();
6230        }
6231        (
6232            data,
6233            vec![
6234                (0..circle_rows).collect(),
6235                (circle_rows..circle_rows + point_rows).collect(),
6236            ],
6237        )
6238    }
6239
6240    #[test]
6241    fn heterogeneous_union_role_assignment_is_group_label_invariant() {
6242        let (data, groups) = circle_and_point_union_data();
6243        let config = GaussianMixtureConfig::default();
6244        let forward = fit_union_density_from_groups(
6245            data.view(),
6246            UnionStructure::CirclePointCluster,
6247            &groups,
6248            config,
6249        )
6250        .unwrap();
6251        let reversed_groups = vec![groups[1].clone(), groups[0].clone()];
6252        let reversed = fit_union_density_from_groups(
6253            data.view(),
6254            UnionStructure::CirclePointCluster,
6255            &reversed_groups,
6256            config,
6257        )
6258        .unwrap();
6259
6260        assert_eq!(forward.components[0].kind, UnionComponentKind::Circle);
6261        assert_eq!(forward.components[1].kind, UnionComponentKind::PointCluster);
6262        assert_eq!(
6263            reversed.components[0].kind,
6264            UnionComponentKind::PointCluster
6265        );
6266        assert_eq!(reversed.components[1].kind, UnionComponentKind::Circle);
6267        assert_eq!(forward.total_parameters, 4 + 3 + 1);
6268        assert_eq!(reversed.total_parameters, forward.total_parameters);
6269        assert!(
6270            (forward.log_likelihood - reversed.log_likelihood).abs()
6271                <= 1.0e-12 * (1.0 + forward.log_likelihood.abs())
6272        );
6273        assert!((forward.bic - reversed.bic).abs() <= 1.0e-12 * (1.0 + forward.bic.abs()));
6274    }
6275
6276    #[test]
6277    fn point_cluster_is_isotropic_and_line_remains_full_covariance() {
6278        let (mut data, mut groups) = circle_and_point_union_data();
6279        // Replace the first group by a narrow, genuinely anisotropic line so
6280        // the line role is identifiable without changing the point group.
6281        for row in 0..groups[0].len() {
6282            let coordinate = (row as f64 - 15.5) / 4.0;
6283            data[[row, 0]] = -4.0 + coordinate;
6284            data[[row, 1]] = 0.2 + 0.018 * coordinate + 0.006 * (1.9 * row as f64).sin();
6285        }
6286        let fit = fit_union_density_from_groups(
6287            data.view(),
6288            UnionStructure::LineCluster,
6289            &groups,
6290            GaussianMixtureConfig::default(),
6291        )
6292        .unwrap();
6293        assert_eq!(fit.components[0].kind, UnionComponentKind::Line);
6294        assert_eq!(fit.components[0].num_parameters, 5);
6295        assert_eq!(fit.components[1].kind, UnionComponentKind::PointCluster);
6296        assert_eq!(fit.components[1].num_parameters, 3);
6297        assert_eq!(fit.total_parameters, 5 + 3 + 1);
6298
6299        let UnionDensityModel::Gaussian(point) = &fit.components[1].model else {
6300            panic!("point cluster did not produce a Gaussian density");
6301        };
6302        assert_eq!(point.precision[[0, 1]], 0.0);
6303        assert_eq!(point.precision[[1, 0]], 0.0);
6304        assert_eq!(point.precision[[0, 0]], point.precision[[1, 1]]);
6305        let total_weight: f64 = fit
6306            .components
6307            .iter()
6308            .map(|component| component.mixing_weight)
6309            .sum();
6310        assert!((total_weight - 1.0).abs() <= 8.0 * f64::EPSILON);
6311
6312        // Reversing group labels must reverse the selected roles, not the
6313        // fitted unlabeled mixture density or its common-scale score.
6314        groups.reverse();
6315        let reversed = fit_union_density_from_groups(
6316            data.view(),
6317            UnionStructure::LineCluster,
6318            &groups,
6319            GaussianMixtureConfig::default(),
6320        )
6321        .unwrap();
6322        assert_eq!(
6323            reversed.components[0].kind,
6324            UnionComponentKind::PointCluster
6325        );
6326        assert_eq!(reversed.components[1].kind, UnionComponentKind::Line);
6327        assert!((fit.bic - reversed.bic).abs() <= 1.0e-12 * (1.0 + fit.bic.abs()));
6328    }
6329
6330    #[test]
6331    fn isotropic_union_density_uses_the_same_fractional_mean_chart_as_its_mle() {
6332        let translated = ndarray::array![[1.0e16], [1.0e16 + 2.0], [1.0e16 + 2.0]];
6333        let fit = fit_isotropic_gaussian_component(translated.view(), 1.0e-12).unwrap();
6334        let variance = fit.precision[[0, 0]].recip();
6335        assert!((variance - 8.0 / 9.0).abs() <= 32.0 * f64::EPSILON);
6336
6337        let residuals = [-4.0 / 3.0, 2.0 / 3.0, 2.0 / 3.0];
6338        let expected_log_norm = -0.5 * ((2.0 * std::f64::consts::PI).ln() + variance.ln());
6339        for (row, residual) in residuals.into_iter().enumerate() {
6340            let expected = expected_log_norm - 0.5 * residual * residual / variance;
6341            let actual = fit.log_density(translated.row(row));
6342            assert!(
6343                (actual - expected).abs() <= 32.0 * f64::EPSILON * (1.0 + expected.abs()),
6344                "row {row}: density chart disagrees with fitted MLE residual: actual={actual}, expected={expected}"
6345            );
6346        }
6347
6348        let subnormal = f64::from_bits(1);
6349        let constant = ndarray::array![[subnormal], [subnormal], [subnormal]];
6350        let constant_fit = fit_isotropic_gaussian_component(constant.view(), 1.0).unwrap();
6351        assert_eq!(constant_fit.residual(constant.row(0)), vec![0.0]);
6352        assert_eq!(
6353            constant_fit.log_density(constant.row(0)),
6354            constant_fit.log_norm
6355        );
6356    }
6357
6358    #[test]
6359    fn union_ladder_fails_closed_when_one_declared_structure_fails() {
6360        let mut data = Array2::<f64>::zeros((8, 2));
6361        for row in 0..5 {
6362            let angle = std::f64::consts::TAU * row as f64 / 5.0;
6363            data[[row, 0]] = -5.0 + angle.cos();
6364            data[[row, 1]] = angle.sin();
6365        }
6366        data[[5, 0]] = 5.00;
6367        data[[5, 1]] = 0.00;
6368        data[[6, 0]] = 5.08;
6369        data[[6, 1]] = 0.02;
6370        data[[7, 0]] = 4.97;
6371        data[[7, 1]] = 0.07;
6372
6373        let error = fit_union_ladder(data.view(), GaussianMixtureConfig::default()).unwrap_err();
6374        assert!(error.contains("every declared structure must fit"));
6375        assert!(error.contains(UnionStructure::CircleCircle.as_str()));
6376        assert!(error.contains("needs at least 5 rows"));
6377    }
6378
6379    #[test]
6380    fn ring_of_clusters_fit_is_stationary_and_complexity_priced_2262() {
6381        let data = seven_clusters_on_a_circle_2262();
6382        let config = GaussianMixtureConfig::default();
6383        let fit = fit_ring_gaussian_mixture(data.view(), 7, config).unwrap();
6384        let certificate = fit.certificate();
6385        assert!(certificate.objective_residual <= certificate.objective_tolerance);
6386        assert!(certificate.parameter_residual <= certificate.parameter_tolerance);
6387        assert_eq!(fit.num_free_parameters(), 17);
6388        assert!((fit.center()[0] - 0.4).abs() < 0.05);
6389        assert!((fit.center()[1] + 0.3).abs() < 0.05);
6390        assert!((fit.radius() - 2.0).abs() < 0.05);
6391        assert!(fit.variance().is_finite() && fit.variance() > 0.0);
6392        assert!(
6393            fit.per_point_log_density(data.view())
6394                .unwrap()
6395                .iter()
6396                .all(|value| value.is_finite())
6397        );
6398        assert!(fit.bic().is_finite());
6399
6400        let free = fit_gaussian_mixture(data.view(), 7, config).unwrap();
6401        assert_eq!(free.num_free_parameters(), 41);
6402        assert!(fit.num_free_parameters() < free.num_free_parameters());
6403    }
6404
6405    #[test]
6406    fn ring_certificate_uses_identifiable_component_means() {
6407        // The points (.3, ±sqrt(.91)) lie on both unit circles centered at
6408        // (0, 0) and (.6, 0). Repeating one point gives three labelled
6409        // components. Thus center and directions move by O(1) while every
6410        // component mean—and therefore the represented mixture density—is
6411        // bit-identical.
6412        let y = 0.91_f64.sqrt();
6413        let weights = Array1::from_vec(vec![0.2, 0.3, 0.5]);
6414        let previous = RingMixtureState {
6415            weights: weights.clone(),
6416            center: Array1::from_vec(vec![0.0, 0.0]),
6417            radius: 1.0,
6418            directions: Array2::from_shape_vec((3, 2), vec![0.3, y, 0.3, -y, 0.3, y]).unwrap(),
6419            variance: 0.25,
6420            mean_log_likelihood: -1.0,
6421            completed_iterations: 10,
6422        };
6423        let next = RingMixtureState {
6424            weights,
6425            center: Array1::from_vec(vec![0.6, 0.0]),
6426            radius: 1.0,
6427            directions: Array2::from_shape_vec((3, 2), vec![-0.3, y, -0.3, -y, -0.3, y]).unwrap(),
6428            variance: 0.25,
6429            mean_log_likelihood: -1.0,
6430            completed_iterations: 11,
6431        };
6432        assert!(relative_parameter_step(previous.center[0], next.center[0]) > 0.5);
6433        let data = array![[0.3, y], [0.3, -y], [1.0, 0.0]];
6434        let previous_e_step = ring_mixture_e_step(data.view(), &previous).unwrap();
6435        let next_e_step = ring_mixture_e_step(data.view(), &next).unwrap();
6436        let residual = empirical_predictive_density_residual(
6437            &previous_e_step.row_log_likelihoods,
6438            &next_e_step.row_log_likelihoods,
6439        )
6440        .unwrap();
6441        assert_eq!(residual, 0.0);
6442    }
6443
6444    #[test]
6445    fn ring_certificate_quotients_duplicate_component_mass_exchange_2324() {
6446        let previous = RingMixtureState {
6447            weights: array![0.2, 0.3, 0.5],
6448            center: array![0.0, 0.0],
6449            radius: 1.0,
6450            directions: array![[1.0, 0.0], [1.0, 0.0], [0.0, 1.0]],
6451            variance: 1.0,
6452            mean_log_likelihood: -1.0,
6453            completed_iterations: 10,
6454        };
6455        let next = RingMixtureState {
6456            weights: array![0.4, 0.1, 0.5],
6457            center: previous.center.clone(),
6458            radius: previous.radius,
6459            directions: previous.directions.clone(),
6460            variance: previous.variance,
6461            mean_log_likelihood: -1.0,
6462            completed_iterations: 11,
6463        };
6464        let data = array![[1.0, 0.0], [0.0, 1.0], [-1.0, 0.0]];
6465        let previous_e_step = ring_mixture_e_step(data.view(), &previous).unwrap();
6466        let next_e_step = ring_mixture_e_step(data.view(), &next).unwrap();
6467        let residual = empirical_predictive_density_residual(
6468            &previous_e_step.row_log_likelihoods,
6469            &next_e_step.row_log_likelihoods,
6470        )
6471        .unwrap();
6472        assert!(residual <= 4.0 * f64::EPSILON);
6473    }
6474
6475    #[test]
6476    fn stacked_mean_is_weighted_combination() {
6477        let weights = Array1::from_vec(vec![0.25, 0.75]);
6478        let means = vec![
6479            Array1::from_vec(vec![1.0, 2.0, 3.0]),
6480            Array1::from_vec(vec![5.0, 6.0, 7.0]),
6481        ];
6482        let out = stacked_predictive_mean(&weights, &means).unwrap();
6483        assert!((out[0] - (0.25 * 1.0 + 0.75 * 5.0)).abs() < 1e-12);
6484        assert!((out[2] - (0.25 * 3.0 + 0.75 * 7.0)).abs() < 1e-12);
6485    }
6486
6487    #[test]
6488    fn stacked_mean_rejects_shape_mismatch() {
6489        let weights = Array1::from_vec(vec![0.5, 0.5]);
6490        let means = vec![
6491            Array1::from_vec(vec![1.0, 2.0]),
6492            Array1::from_vec(vec![3.0]),
6493        ];
6494        assert!(stacked_predictive_mean(&weights, &means).is_err());
6495    }
6496
6497    // -----------------------------------------------------------------------
6498    // #1026 hybrid curved + linear-tail split-selection
6499    // -----------------------------------------------------------------------
6500
6501    /// Build the two candidate parameterizations for one atom slot the way the
6502    /// fit would: the linear special case (one decoder direction, `Θ = 0`,
6503    /// `P_linear` params) and the curved candidate (`latent_dim` ≥ 1, more
6504    /// params, fitted turning `theta`). The curved candidate's likelihood is the
6505    /// linear likelihood MINUS `curved_loglik_gain` of NLE (curvature it captures
6506    /// the secant cannot), so the nesting invariant `curved_loglik ≥ linear` is
6507    /// honored: a straight feature has zero gain, a turning feature a positive
6508    /// gain that grows with Θ. The rank-aware Laplace normalizer charges the
6509    /// extra `½(P_curved − P_linear)·log(2π)` for the curved parameters, so the
6510    /// evidence comparison is the real `Θ/√ε` crossover.
6511    fn hybrid_slot(
6512        linear_nle: f64,
6513        p_linear: usize,
6514        latent_dim: usize,
6515        p_curved: usize,
6516        theta: f64,
6517        curved_loglik_gain: f64,
6518    ) -> Vec<HybridAtomCandidate> {
6519        let param_price =
6520            0.5 * (p_curved as f64 - p_linear as f64) * (2.0 * std::f64::consts::PI).ln();
6521        let curved_nle = linear_nle - curved_loglik_gain + param_price;
6522        vec![
6523            HybridAtomCandidate::linear(linear_nle, p_linear),
6524            HybridAtomCandidate::curved(latent_dim, curved_nle, p_curved, Some(theta)),
6525        ]
6526    }
6527
6528    #[test]
6529    fn hybrid_dominance_floor_selects_linear_when_turning_is_zero() {
6530        // A perfectly straight curved fit (Θ = 0) gains no likelihood over its
6531        // linear sub-model but pays more parameters → linear must win, by
6532        // construction, even if finite-sample evidence noise nudged the curved
6533        // NLE slightly below linear.
6534        let slot = hybrid_slot(100.0, 2, 1, 5, 0.0, 0.0);
6535        let choice = select_hybrid_atom(&slot).unwrap();
6536        assert!(choice.param.is_linear());
6537        assert_eq!(choice.param, HybridAtomParam::Linear);
6538        // The exact-zero guard fires regardless of the evidence margin sign.
6539        assert!(choice.curved_turning.unwrap() <= HYBRID_LINEAR_TURNING_FLOOR);
6540    }
6541
6542    #[test]
6543    fn hybrid_selects_curved_when_turning_pays_for_itself() {
6544        // A genuinely turning feature (Θ = 2π, a full loop): the curved fit
6545        // captures enough curvature that, even charged the extra-parameter price,
6546        // its NLE drops below the linear secant's → curved wins.
6547        let slot = hybrid_slot(100.0, 2, 1, 5, 2.0 * std::f64::consts::PI, 30.0);
6548        let choice = select_hybrid_atom(&slot).unwrap();
6549        assert_eq!(choice.param, HybridAtomParam::Curved { latent_dim: 1 });
6550        // The curved fit won a strictly positive evidence margin.
6551        assert!(choice.curved_evidence_margin > 0.0);
6552    }
6553
6554    #[test]
6555    fn hybrid_keeps_linear_when_curvature_doesnt_pay_its_price() {
6556        // A barely-curved feature (small Θ): the curved fit recovers only a sliver
6557        // of likelihood, not enough to cover the extra-parameter price → the
6558        // dominance floor keeps the linear tail.
6559        let slot = hybrid_slot(100.0, 2, 1, 5, 0.05, 0.1);
6560        let choice = select_hybrid_atom(&slot).unwrap();
6561        assert!(choice.param.is_linear());
6562        assert!(choice.curved_evidence_margin <= 0.0);
6563    }
6564
6565    #[test]
6566    fn hybrid_tie_breaks_to_the_cheaper_linear_atom() {
6567        // Exact NLE tie (above the turning floor so the evidence path decides):
6568        // the cheaper linear atom wins, preserving strict generalization — the
6569        // hybrid never pays for curvature it does not need.
6570        let theta = 0.5; // above the floor → evidence path, not the exact guard
6571        let nle = 42.0;
6572        let slot = vec![
6573            HybridAtomCandidate::linear(nle, 2),
6574            HybridAtomCandidate::curved(1, nle, 5, Some(theta)),
6575        ];
6576        let choice = select_hybrid_atom(&slot).unwrap();
6577        assert!(choice.param.is_linear());
6578        assert_eq!(choice.num_parameters, 2);
6579    }
6580
6581    #[test]
6582    fn hybrid_split_reduces_to_pure_linear_when_all_features_are_straight() {
6583        // Every slot's curved candidate has Θ → 0 (flat features everywhere): the
6584        // dominance floor fires at every slot → the hybrid recovers the pure-
6585        // linear dictionary exactly. This is the `all Θ → 0` limit (3).
6586        let slots: Vec<Vec<HybridAtomCandidate>> = (0..6)
6587            .map(|i| hybrid_slot(50.0 + i as f64, 2, 1, 5, 0.0, 0.0))
6588            .collect();
6589        let split = select_hybrid_split(&slots).unwrap();
6590        assert!(split.is_pure_linear());
6591        assert_eq!(split.curved_atom_count, 0);
6592        assert_eq!(split.linear_atom_count(), 6);
6593        // Summed NLE equals the pure-linear baseline (every slot chose linear).
6594        let pure_linear: f64 = (0..6).map(|i| 50.0 + i as f64).sum();
6595        assert!((split.total_negative_log_evidence - pure_linear).abs() < 1e-12);
6596    }
6597
6598    #[test]
6599    fn hybrid_split_reduces_to_pure_curved_when_every_feature_curves() {
6600        // Every slot's feature turns enough (Θ = 2π, large likelihood gain) that
6601        // curved beats linear everywhere → the pure-curved limit (3).
6602        let slots: Vec<Vec<HybridAtomCandidate>> = (0..5)
6603            .map(|i| hybrid_slot(80.0 + i as f64, 2, 1, 5, 2.0 * std::f64::consts::PI, 40.0))
6604            .collect();
6605        let split = select_hybrid_split(&slots).unwrap();
6606        assert!(split.is_pure_curved());
6607        assert_eq!(split.curved_atom_count, 5);
6608        assert_eq!(split.linear_atom_count(), 0);
6609    }
6610
6611    #[test]
6612    fn hybrid_split_on_mixed_dictionary_picks_curved_for_circles_linear_for_directions() {
6613        // Mixed synthetic: slots 0..3 are CIRCLE features (high turning Θ = 2π,
6614        // the curved fit captures the loop), slots 3..7 are LINEAR DIRECTIONS
6615        // (straight, Θ = 0). The evidence split must select curved for the
6616        // circles and linear for the directions — and the hybrid's summed
6617        // evidence must be ≤ the summed per-slot LINEAR-candidate NLE (each
6618        // slot's best straight line fit to its response residual). This is a
6619        // data-level match-or-beat dominance (#1202: linear is the curved
6620        // family's nested Θ = 0 sub-model on common data), and holds because each
6621        // slot picks the argmin of its two common-data candidates.
6622        let mut slots: Vec<Vec<HybridAtomCandidate>> = Vec::new();
6623        let mut pure_linear_baseline = 0.0_f64;
6624        // Three circle features: a curved atom replaces ~10-30 linear secants, so
6625        // the curved fit buys a large likelihood gain that dwarfs its param price.
6626        for i in 0..3 {
6627            let linear_nle = 120.0 + 3.0 * i as f64;
6628            pure_linear_baseline += linear_nle;
6629            slots.push(hybrid_slot(
6630                linear_nle,
6631                2,
6632                1,
6633                5,
6634                2.0 * std::f64::consts::PI,
6635                35.0,
6636            ));
6637        }
6638        // Four straight linear directions: zero turning, the linear special case
6639        // is optimal — a curved atom buys nothing and only costs parameters.
6640        for i in 0..4 {
6641            let linear_nle = 90.0 + 2.0 * i as f64;
6642            pure_linear_baseline += linear_nle;
6643            slots.push(hybrid_slot(linear_nle, 2, 1, 5, 0.0, 0.0));
6644        }
6645
6646        let split = select_hybrid_split(&slots).unwrap();
6647
6648        // The first three (circles) chose curved; the last four (directions) chose
6649        // linear.
6650        for (idx, choice) in split.atoms.iter().enumerate() {
6651            if idx < 3 {
6652                assert_eq!(
6653                    choice.param,
6654                    HybridAtomParam::Curved { latent_dim: 1 },
6655                    "circle slot {idx} should select curved"
6656                );
6657            } else {
6658                assert!(
6659                    choice.param.is_linear(),
6660                    "direction slot {idx} should select linear"
6661                );
6662            }
6663        }
6664        assert_eq!(split.curved_atom_count, 3);
6665        assert_eq!(split.linear_atom_count(), 4);
6666
6667        // The hybrid's summed negative-log-evidence is ≤ the summed per-slot
6668        // LINEAR-candidate NLE (each slot's best straight line fit to its response
6669        // residual): the per-slot argmin can only lower the sum. This is a
6670        // data-level match-or-beat dominance (#1202): linear is the curved
6671        // family's nested Θ = 0 sub-model on common data.
6672        assert!(
6673            split.total_negative_log_evidence <= pure_linear_baseline + 1e-9,
6674            "hybrid NLE {} must be <= summed linear-candidate NLE {}",
6675            split.total_negative_log_evidence,
6676            pure_linear_baseline
6677        );
6678        // And strictly better, because the curved circle slots paid off.
6679        assert!(split.total_negative_log_evidence < pure_linear_baseline);
6680    }
6681
6682    #[test]
6683    fn hybrid_split_rejects_empty_slot() {
6684        let slots = vec![hybrid_slot(10.0, 2, 1, 5, 0.0, 0.0), Vec::new()];
6685        assert!(select_hybrid_split(&slots).is_err());
6686    }
6687
6688    // ── #1362: compare_models must Occam-penalise a pure-noise smooth ────────
6689    //
6690    // These tests pin the ranking contract directly on `compare_reml_fits` with
6691    // controlled (score, edf, log_lik) inputs taken from the actual #1362
6692    // reproduction (Rust `reml_score` of `y ~ s(x)` vs `y ~ s(x) + s(z)` at
6693    // n=700). They do not need a fitted GAM or a Python wheel.
6694
6695    fn cand(name: &str, score: f64, edf: f64, log_lik: f64) -> RemlCandidate {
6696        RemlCandidate {
6697            index: 0,
6698            name: name.to_string(),
6699            score,
6700            edf: Some(edf),
6701            log_lik: Some(log_lik),
6702            family: None,
6703            n_obs: None,
6704        }
6705    }
6706
6707    #[test]
6708    fn ranking_score_is_conditional_aic_when_loglik_and_edf_present() {
6709        // AIC = -2ℓ + 2·edf.
6710        let c = cand("m", /*score (ignored)*/ 999.0, 6.748, -32.0866);
6711        let expected = -2.0 * -32.0866 + 2.0 * 6.748;
6712        assert!((c.ranking_score() - expected).abs() < 1e-9);
6713    }
6714
6715    #[test]
6716    fn ranking_score_falls_back_to_evidence_without_loglik() {
6717        let c = RemlCandidate {
6718            index: 0,
6719            name: "m".to_string(),
6720            score: 151.28,
6721            edf: Some(6.0),
6722            log_lik: None,
6723            family: None,
6724            n_obs: None,
6725        };
6726        assert_eq!(c.ranking_score(), 151.28);
6727    }
6728
6729    #[test]
6730    fn compare_models_rejects_pure_noise_smooth_despite_lower_evidence() {
6731        // Seed-3000 numbers from the #1362 Rust reproduction:
6732        //   small (y ~ s(x)):      reml=180.526, edf=6.748,  loglik=-32.0866
6733        //   big   (y ~ s(x)+s(z)): reml=177.404, edf=14.250, loglik=-32.1212
6734        // The big (noise-augmented) model has the LOWER (apparently better) raw
6735        // REML evidence, yet it spends ~7.5 extra EDF fitting noise without
6736        // improving the likelihood. The winner must be the SMALL model.
6737        let small = cand("small", 180.526, 6.748, -32.0866);
6738        let big = cand("big", 177.404, 14.250, -32.1212);
6739
6740        // Sanity: raw evidence (the broken headline) prefers big.
6741        assert!(big.score < small.score);
6742
6743        let cmp = compare_reml_fits(vec![small, big]).expect("compare");
6744        assert_eq!(
6745            cmp.winner, "small",
6746            "compare_models must Occam-penalise the pure-noise smooth and pick the smaller model"
6747        );
6748        // The score table still reports the raw evidence headline unchanged, so
6749        // Model.evidence / bayes_factor_vs stay consistent with the table.
6750        let small_row = cmp
6751            .score_table
6752            .iter()
6753            .find(|r| r.name == "small")
6754            .expect("small row");
6755        let big_row = cmp
6756            .score_table
6757            .iter()
6758            .find(|r| r.name == "big")
6759            .expect("big row");
6760        assert!((small_row.reml_score - 180.526).abs() < 1e-9);
6761        assert!((big_row.reml_score - 177.404).abs() < 1e-9);
6762    }
6763
6764    #[test]
6765    fn ranking_bayes_factor_is_akaike_evidence_ratio_not_its_square() {
6766        // Issue #2124: `ranking_score` is the conditional AIC (`−2ℓ + 2·edf`), a
6767        // −2·log / deviance-scale cost. For an AIC gap Δ the Akaike evidence ratio
6768        // (Burnham & Anderson) is `exp(−½Δ)`, so the winner-over-loser
6769        // `bayes_factor` must be `exp(½Δ)` — NOT `exp(Δ)`, which squares it.
6770        //
6771        // Winner: AIC 0 (loglik 0, edf 0). Loser: AIC = 27.68 (loglik −13.84,
6772        // edf 0), matching the ΔAIC in the issue repro. Raw REML scores are set
6773        // distinct (100 vs 110) to lock the scoping: the raw score_table path
6774        // must stay `exp(Δreml)` with NO halving.
6775        let delta_aic = 27.68_f64;
6776        let winner = cand("winner", 100.0, 0.0, 0.0);
6777        let loser = cand("loser", 110.0, 0.0, -delta_aic / 2.0);
6778
6779        let cmp = compare_reml_fits(vec![winner, loser]).expect("compare");
6780        assert_eq!(cmp.winner, "winner");
6781
6782        let loser_row = cmp
6783            .ranking
6784            .iter()
6785            .find(|r| r.name == "loser")
6786            .expect("loser ranking row");
6787
6788        // The AIC gap FIELD stays on the AIC scale, unchanged (issue #2124).
6789        assert!((loser_row.delta - delta_aic).abs() < 1e-9);
6790
6791        // The Bayes factor is the Akaike evidence ratio exp(½·ΔAIC) = exp(13.84)
6792        // ≈ 1.03e6 — NOT the squared exp(27.68) ≈ 1.05e12 the bug reported.
6793        let expected = (0.5 * delta_aic).exp();
6794        assert!(
6795            (loser_row.bayes_factor / expected - 1.0).abs() < 1e-9,
6796            "ranking bayes_factor {} should be exp(½ΔAIC)={}, not exp(ΔAIC)={}",
6797            loser_row.bayes_factor,
6798            expected,
6799            delta_aic.exp()
6800        );
6801        // Explicit anti-regression: it must not be the squared ratio.
6802        assert!(loser_row.bayes_factor < delta_aic.exp() * 0.5);
6803
6804        // Scoping lock (issue #2124): the RAW-REML score_table path is untouched —
6805        // its best-over-model Bayes factor is `exp(Δreml)` with NO halving. Raw
6806        // scores 100 (winner) vs 110 (loser) give Δreml = 10, so the loser's raw
6807        // Bayes factor is exp(10), not exp(5).
6808        let loser_score_row = cmp
6809            .score_table
6810            .iter()
6811            .find(|r| r.name == "loser")
6812            .expect("loser score row");
6813        let expected_reml_bf = 10.0_f64.exp();
6814        assert!(
6815            (loser_score_row.bayes_factor_best_over_model / expected_reml_bf - 1.0).abs() < 1e-9,
6816            "raw-REML bayes_factor_best_over_model must stay exp(Δreml)=exp(10), got {}",
6817            loser_score_row.bayes_factor_best_over_model
6818        );
6819    }
6820
6821    #[test]
6822    fn compare_models_keeps_power_for_a_relevant_smooth() {
6823        // Seed-3000 relevant-z numbers from the same reproduction:
6824        //   small: reml=1025.067, edf≈6.75,  loglik≈-368.99 (aic≈751.5)
6825        //   big:   reml=199.509,  edf≈14.25, loglik≈-33.16  (aic≈94.8)
6826        // A genuinely relevant smooth lowers BOTH the evidence and the AIC, so
6827        // the bigger model must still win — a fix cannot just always pick small.
6828        let small = cand("small", 1025.067, 6.75, -368.985);
6829        let big = cand("big", 199.509, 14.25, -33.165);
6830        let cmp = compare_reml_fits(vec![small, big]).expect("compare");
6831        assert_eq!(
6832            cmp.winner, "big",
6833            "compare_models must retain power: the relevant smooth's model must win"
6834        );
6835    }
6836
6837    #[test]
6838    fn compare_models_rejects_mismatched_observation_counts() {
6839        // Two same-family fits on different-sized data are not comparable by
6840        // AIC / evidence; the comparison must fail loud, mirroring the family
6841        // guard, rather than declare a sample-size-driven winner.
6842        let with_n = |name: &str, n: usize| RemlCandidate {
6843            index: 0,
6844            name: name.to_string(),
6845            score: 100.0,
6846            edf: Some(5.0),
6847            log_lik: Some(-40.0),
6848            family: Some("gaussian".to_string()),
6849            n_obs: Some(n),
6850        };
6851        let err = compare_reml_fits(vec![with_n("big", 500), with_n("small", 100)])
6852            .expect_err("cross-n comparison must be rejected");
6853        assert!(
6854            err.contains("number of observations") && err.contains("500") && err.contains("100"),
6855            "n-guard error should name the incomparable counts, got: {err}"
6856        );
6857
6858        // Same n is comparable.
6859        compare_reml_fits(vec![with_n("a", 250), with_n("b", 250)])
6860            .expect("same-n comparison must succeed");
6861
6862        // A missing count (`None`) is unconstrained: it must not block a
6863        // comparison against a fit that does carry one (legacy / scan payloads).
6864        let without_n = RemlCandidate {
6865            index: 0,
6866            name: "legacy".to_string(),
6867            score: 90.0,
6868            edf: Some(4.0),
6869            log_lik: Some(-35.0),
6870            family: Some("gaussian".to_string()),
6871            n_obs: None,
6872        };
6873        compare_reml_fits(vec![with_n("counted", 500), without_n])
6874            .expect("an unconstrained (None) count must not trip the guard");
6875    }
6876}