Skip to main content

gam_terms/
construction.rs

1use crate::EstimationError;
2use crate::basis::analyze_penalty_block;
3use crate::smooth::PenaltyStructureHint;
4use faer::linalg::matmul::matmul;
5use faer::{Accum, Mat, MatRef, Par, Side};
6use gam_linalg::faer_ndarray::{FaerLinalgError, FaerQr, FaerSvd};
7use gam_linalg::matrix::symmetrize_in_place;
8use gam_linalg::utils::KahanSum;
9use ndarray::{Array1, Array2, ArrayView1, ArrayView2, ArrayViewMut2, s};
10use rayon::iter::{
11    IndexedParallelIterator, IntoParallelIterator, IntoParallelRefIterator, ParallelIterator,
12};
13use std::collections::{BTreeMap, HashSet};
14use std::ops::Range;
15use std::sync::Arc;
16
17/// Relative "numerically-PSD" floor of the symmetric eigensolver: an eigenvalue
18/// smaller than this fraction of the spectrum scale is roundoff, not genuine
19/// curvature (~`sqrt(machine ε)`; #1619). This single floor governs BOTH which
20/// eigenvalues are snapped to zero (declared null) in `classify_eigenvalues_strict`
21/// AND how much per-mode relative energy the transformed penalty root may retain
22/// in the resulting null block (`subspace_split_is_consistent`). Keeping the two
23/// tied to one constant makes the null-space definition and the leakage-consistency
24/// guard provably mutually consistent, rather than the guard demanding a precision
25/// the classifier never promised.
26const REL_PSD_FLOOR: f64 = 1.0e-8;
27
28#[derive(Clone)]
29pub enum PenaltyRepresentation {
30    Dense(Array2<f64>),
31    Banded {
32        bands: Vec<Array1<f64>>,
33        offsets: Vec<i32>,
34    },
35    Kronecker {
36        /// Full penalty-block Kronecker product `left ⊗ right`: each entry of
37        /// `left` scales an entire copy of `right` in the dense expansion.
38        ///
39        /// This is distinct from chunked kernel design assembly, where center
40        /// rows are kernel-evaluation arguments rather than matrix factors.
41        left: Array2<f64>,
42        right: Array2<f64>,
43    },
44}
45
46impl PenaltyRepresentation {
47
48}
49
50#[derive(Clone)]
51pub struct PenaltyMatrix {
52    pub col_range: Range<usize>,
53    pub representation: PenaltyRepresentation,
54}
55
56impl PenaltyMatrix {
57    fn accumulate_into(&self, mut dest: ArrayViewMut2<'_, f64>, weight: f64) {
58        if weight == 0.0 {
59            return;
60        }
61        match &self.representation {
62            PenaltyRepresentation::Dense(block) => {
63                dest.scaled_add(weight, block);
64            }
65            PenaltyRepresentation::Banded { bands, offsets } => {
66                let positive_offsets: HashSet<usize> = offsets
67                    .iter()
68                    .filter_map(|&off| (off >= 0).then_some(off as usize))
69                    .collect();
70                for (band, &offset) in bands.iter().zip(offsets.iter()) {
71                    let off = offset.unsigned_abs() as usize;
72                    if offset < 0 && positive_offsets.contains(&off) {
73                        continue;
74                    }
75                    for (idx, &value) in band.iter().enumerate() {
76                        let (i, j) = if offset >= 0 {
77                            (idx, idx + off)
78                        } else {
79                            (idx + off, idx)
80                        };
81                        let Some(entry_ij) = dest.get_mut((i, j)) else {
82                            continue;
83                        };
84                        *entry_ij += weight * value;
85                        if i != j
86                            && let Some(entry_ji) = dest.get_mut((j, i))
87                        {
88                            *entry_ji += weight * value;
89                        }
90                    }
91                }
92            }
93            PenaltyRepresentation::Kronecker { left, right } => {
94                let (lrows, l_cols) = left.dim();
95                let (rrows, r_cols) = right.dim();
96                for i in 0..lrows {
97                    for j in 0..l_cols {
98                        let scale = left[(i, j)] * weight;
99                        if scale == 0.0 {
100                            continue;
101                        }
102                        let mut block = dest.slice_mut(s![
103                            i * rrows..(i + 1) * rrows,
104                            j * r_cols..(j + 1) * r_cols
105                        ]);
106                        block.scaled_add(scale, right);
107                    }
108                }
109            }
110        }
111    }
112
113    pub fn to_dense(&self, total_dim: usize) -> Array2<f64> {
114        let mut dense = Array2::<f64>::zeros((total_dim, total_dim));
115        self.accumulate_into(
116            dense.slice_mut(s![self.col_range.clone(), self.col_range.clone()]),
117            1.0,
118        );
119        dense
120    }
121}
122
123pub(crate) fn array_to_faer(array: &Array2<f64>) -> Mat<f64> {
124    let (rows, cols) = array.dim();
125    Mat::from_fn(rows, cols, |i, j| array[[i, j]])
126}
127
128pub(crate) fn mat_to_array(mat: &Mat<f64>) -> Array2<f64> {
129    let mut out = Array2::<f64>::zeros((mat.nrows(), mat.ncols()));
130    for i in 0..mat.nrows() {
131        for j in 0..mat.ncols() {
132            out[[i, j]] = mat[(i, j)];
133        }
134    }
135    out
136}
137
138fn mat_max_abs_element(matrix: MatRef<'_, f64>) -> f64 {
139    let (rows, cols) = matrix.shape();
140    let mut maxval = 0.0_f64;
141    for i in 0..rows {
142        for j in 0..cols {
143            let val = matrix[(i, j)];
144            if val.is_finite() {
145                maxval = maxval.max(val.abs());
146            }
147        }
148    }
149    maxval
150}
151
152fn sanitize_symmetric_faer(matrix: &Mat<f64>) -> Mat<f64> {
153    let (rows, cols) = matrix.as_ref().shape();
154    assert_eq!(rows, cols, "Matrix must be square for sanitization");
155
156    let mut sanitized = matrix.clone();
157
158    for i in 0..rows {
159        let diag = sanitized[(i, i)];
160        if !diag.is_finite() {
161            sanitized[(i, i)] = 0.0;
162        }
163        for j in (i + 1)..cols {
164            let mut upper = sanitized[(i, j)];
165            let mut lower = sanitized[(j, i)];
166            if !upper.is_finite() {
167                upper = 0.0;
168            }
169            if !lower.is_finite() {
170                lower = 0.0;
171            }
172            let avg = 0.5 * (upper + lower);
173            sanitized[(i, j)] = avg;
174            sanitized[(j, i)] = avg;
175        }
176    }
177
178    let scale = mat_max_abs_element(sanitized.as_ref());
179    let tiny = (scale * 1e-14).max(1e-30);
180    for i in 0..rows {
181        for j in 0..cols {
182            let val = sanitized[(i, j)];
183            if !val.is_finite() {
184                sanitized[(i, j)] = 0.0;
185            } else if val.abs() < tiny {
186                sanitized[(i, j)] = 0.0;
187            }
188        }
189    }
190
191    sanitized
192}
193
194fn penalty_from_root_faer(root: &Mat<f64>) -> Mat<f64> {
195    let cols = root.ncols();
196    let mut full = Mat::<f64>::zeros(cols, cols);
197    let root_ref = root.as_ref();
198    let root_t = root_ref.transpose();
199    matmul(
200        full.as_mut(),
201        Accum::Replace,
202        root_t,
203        root_ref,
204        1.0,
205        Par::Seq,
206    );
207    sanitize_symmetric_faer(&full)
208}
209
210fn symmetrize_faer_matrix_in_place(matrix: &mut Mat<f64>) {
211    let n = matrix.nrows().min(matrix.ncols());
212    for i in 0..n {
213        for j in 0..i {
214            let avg = 0.5 * (matrix[(i, j)] + matrix[(j, i)]);
215            matrix[(i, j)] = avg;
216            matrix[(j, i)] = avg;
217        }
218    }
219}
220
221fn orthogonal_similarity_transform_faer(
222    matrix: &Mat<f64>,
223    block_dim: usize,
224    orthogonal: &Mat<f64>,
225) -> Mat<f64> {
226    let matrix_block = matrix.as_ref().submatrix(0, 0, block_dim, block_dim);
227    let cols = orthogonal.ncols();
228    let mut temp = Mat::<f64>::zeros(block_dim, cols);
229    matmul(
230        temp.as_mut(),
231        Accum::Replace,
232        matrix_block,
233        orthogonal.as_ref(),
234        1.0,
235        Par::Seq,
236    );
237    let mut rotated = Mat::<f64>::zeros(cols, cols);
238    matmul(
239        rotated.as_mut(),
240        Accum::Replace,
241        orthogonal.transpose(),
242        temp.as_ref(),
243        1.0,
244        Par::Seq,
245    );
246    symmetrize_faer_matrix_in_place(&mut rotated);
247    rotated
248}
249
250fn trace_penalty_in_orthogonal_basis(
251    matrix: &Mat<f64>,
252    block_dim: usize,
253    orthogonal: &Mat<f64>,
254    rotated_eigenvalues: &[f64],
255    delta: f64,
256) -> f64 {
257    let matrix_block = matrix.as_ref().submatrix(0, 0, block_dim, block_dim);
258    let cols = orthogonal.ncols();
259    assert!(rotated_eigenvalues.len() >= cols);
260    let mut projected = Mat::<f64>::zeros(block_dim, cols);
261    matmul(
262        projected.as_mut(),
263        Accum::Replace,
264        matrix_block,
265        orthogonal.as_ref(),
266        1.0,
267        Par::Seq,
268    );
269    let mut trace = KahanSum::default();
270    for l in 0..cols {
271        let mut diag_ll = KahanSum::default();
272        for i in 0..block_dim {
273            diag_ll.add(orthogonal[(i, l)] * projected[(i, l)]);
274        }
275        trace.add(diag_ll.sum() / (rotated_eigenvalues[l] + delta));
276    }
277    trace.sum()
278}
279
280pub fn trace_reduced_penalty_covariance(
281    reduced_penalty: &Array2<f64>,
282    covariance_basis: &Array2<f64>,
283) -> f64 {
284    assert_eq!(
285        reduced_penalty.dim(),
286        covariance_basis.dim(),
287        "trace_reduced_penalty_covariance dimension mismatch"
288    );
289    let r = covariance_basis.nrows();
290    let mut trace = KahanSum::default();
291    for i in 0..r {
292        for j in 0..r {
293            trace.add(covariance_basis[[i, j]] * reduced_penalty[[j, i]]);
294        }
295    }
296    trace.sum()
297}
298
299pub fn trace_penalty_covariance_in_orthogonal_basis(
300    matrix: &Array2<f64>,
301    orthogonal: &Array2<f64>,
302    covariance_basis: &Array2<f64>,
303) -> f64 {
304    let reduced = gam_linalg::faer_ndarray::fast_ab(
305        &gam_linalg::faer_ndarray::fast_atb(orthogonal, matrix),
306        orthogonal,
307    );
308    trace_reduced_penalty_covariance(&reduced, covariance_basis)
309}
310
311/// Strict spectral classifier used as a final guard on penalty eigendecompositions.
312///
313/// Penalty matrices fed to the GAM solver are required to be PSD by construction.
314/// This routine snaps roundoff-zero eigenvalues to exact zero, accepts strictly
315/// positive eigenvalues, and rejects materially-indefinite or non-finite spectra
316/// with a hard error rather than silently rewriting them. The previous behaviour
317/// (mass-zeroing negative or non-finite eigenvalues) hid construction bugs and
318/// changed the optimisation objective downstream.
319///
320/// The acceptance tolerance is the larger of a machine-ε floor
321/// (`C_EPS_P_FACTOR * eps_machine * p * scale`, with `C_EPS_P_FACTOR = 64`
322/// absorbing the rounding accumulated in a symmetric eigendecomposition of a
323/// moderate-dimension matrix) and a relative "numerically PSD" floor
324/// `REL_PSD_FLOOR * scale`. The latter dominates for large high-rank penalties
325/// assembled / reparameterized at extreme λ, where roundoff produces
326/// ~1e-11-relative negative eigenvalues that are PSD to any reasonable precision
327/// yet exceeded the bare ~12×ε machine floor and spuriously failed the inner
328/// P-IRLS solve (#1619). Genuine indefiniteness is O(1) relative and is still
329/// rejected far above either floor.
330fn classify_eigenvalues_strict(
331    eigenvalues: &mut [f64],
332    context: &str,
333) -> Result<(), EstimationError> {
334    const C_EPS_P_FACTOR: f64 = 64.0;
335    // `REL_PSD_FLOOR` (module-level): the relative threshold below which a
336    // (possibly slightly negative) eigenvalue is roundoff and is snapped to zero
337    // rather than rejected. Shared with the subspace-leakage guard so the null
338    // definition and the leakage tolerance stay mutually consistent.
339    let p = eigenvalues.len();
340
341    let mut scale = 0.0_f64;
342    for (idx, &val) in eigenvalues.iter().enumerate() {
343        if !val.is_finite() {
344            return Err(EstimationError::PenaltySpectrumNonFinite {
345                context: context.to_string(),
346                index: idx,
347                value: val,
348            });
349        }
350        scale = scale.max(val.abs());
351    }
352
353    // p * eps captures the rounding floor of a symmetric eigendecomposition of a
354    // p-dimensional matrix; multiplying by `scale` lifts the floor to the actual
355    // magnitude of the spectrum. For large high-rank penalties assembled at
356    // extreme λ this machine floor (~12×ε relative) is tighter than the roundoff
357    // actually produced, so we take the larger of it and a relative numerically-PSD
358    // floor `REL_PSD_FLOOR * scale` (#1619).
359    let machine_floor = C_EPS_P_FACTOR * f64::EPSILON * (p.max(1) as f64) * scale;
360    let tolerance = machine_floor
361        .max(REL_PSD_FLOOR * scale)
362        .max(f64::MIN_POSITIVE);
363
364    for (idx, val) in eigenvalues.iter_mut().enumerate() {
365        if val.abs() <= tolerance {
366            *val = 0.0;
367        } else if *val < 0.0 {
368            return Err(EstimationError::PenaltySpectrumIndefinite {
369                context: context.to_string(),
370                index: idx,
371                value: *val,
372                tolerance,
373                scale,
374            });
375        }
376    }
377    Ok(())
378}
379
380fn robust_eighwith_policy<M, V, E, Validate, Sanitize, EigCall, MapErr>(
381    matrix: &M,
382    context: &str,
383    validate_input: Validate,
384    sanitize: Sanitize,
385    mut eig_call: EigCall,
386    map_error: MapErr,
387) -> Result<(Vec<f64>, V), EstimationError>
388where
389    Validate: Fn(&M, &str) -> Result<(), EstimationError>,
390    Sanitize: Fn(&M) -> M,
391    EigCall: FnMut(&M) -> Result<(Vec<f64>, V), E>,
392    MapErr: Fn(E, &str) -> EstimationError,
393{
394    validate_input(matrix, context)?;
395
396    // The sanitize step only enforces exact symmetry by averaging M and M^T and
397    // zeros sub-eps noise; it never adds a diagonal ridge. Adding ridge changes
398    // the matrix being decomposed, which silently changes the optimisation
399    // objective downstream. If eigh genuinely fails on a finite symmetric input,
400    // surface the error instead of mutating the spectrum.
401    let candidate = sanitize(matrix);
402    match eig_call(&candidate) {
403        Ok((mut eigenvalues, eigenvectors)) => {
404            classify_eigenvalues_strict(&mut eigenvalues, context)?;
405            Ok((eigenvalues, eigenvectors))
406        }
407        Err(err) => Err(map_error(err, context)),
408    }
409}
410
411pub(crate) fn robust_eigh_faer(
412    matrix: &Mat<f64>,
413    side: Side,
414    context: &str,
415) -> Result<(Vec<f64>, Mat<f64>), EstimationError> {
416    robust_eighwith_policy(
417        matrix,
418        context,
419        |mat, ctx| {
420            let (rows, cols) = mat.as_ref().shape();
421            for i in 0..rows {
422                for j in 0..cols {
423                    let val = mat[(i, j)];
424                    if !val.is_finite() {
425                        let max_abs = mat_max_abs_element(mat.as_ref());
426                        crate::bail_invalid_estim!(
427                            "{} contains non-finite entries (max finite magnitude {:.3e})",
428                            ctx,
429                            max_abs
430                        );
431                    }
432                }
433            }
434            Ok(())
435        },
436        sanitize_symmetric_faer,
437        |candidate| {
438            let eig = candidate.as_ref().self_adjoint_eigen(side)?;
439            let diag = eig.S();
440            let mut eigenvalues = Vec::with_capacity(diag.dim());
441            for idx in 0..diag.dim() {
442                eigenvalues.push(diag[idx]);
443            }
444
445            let vectors_ref = eig.U();
446            let mut eigenvectors = Mat::<f64>::zeros(vectors_ref.nrows(), vectors_ref.ncols());
447            for i in 0..vectors_ref.nrows() {
448                for j in 0..vectors_ref.ncols() {
449                    eigenvectors[(i, j)] = vectors_ref[(i, j)];
450                }
451            }
452            Ok((eigenvalues, eigenvectors))
453        },
454        |err, _| {
455            EstimationError::EigendecompositionFailed(FaerLinalgError::SelfAdjointEigen(err))
456        },
457    )
458}
459
460fn robust_eigh(
461    matrix: &Array2<f64>,
462    side: Side,
463    context: &str,
464) -> Result<(Array1<f64>, Array2<f64>), EstimationError> {
465    let matrix_faer = array_to_faer(matrix);
466    let (eigenvalues, eigenvectors) = robust_eigh_faer(&matrix_faer, side, context)?;
467    Ok((Array1::from_vec(eigenvalues), mat_to_array(&eigenvectors)))
468}
469
470pub(crate) fn kronecker_marginal_eigensystems(
471    marginal_penalties: &[Array2<f64>],
472    context: &str,
473) -> Result<Vec<(Array1<f64>, Array2<f64>)>, EstimationError> {
474    let mut eigensystems = Vec::with_capacity(marginal_penalties.len());
475    for (k, penalty) in marginal_penalties.iter().enumerate() {
476        eigensystems.push(robust_eigh(
477            penalty,
478            Side::Lower,
479            &format!("{context} marginal {k}"),
480        )?);
481    }
482    Ok(eigensystems)
483}
484
485#[derive(Debug, Clone, Copy)]
486struct SubspaceLeakageMetrics {
487    max_abs_sq: f64,
488    max_rel_sq: f64,
489    worst_penalty: usize,
490    max_cross_gram_abs: f64,
491}
492
493fn assess_subspace_leakage(
494    qs: &Mat<f64>,
495    rs_transformed: &[Mat<f64>],
496    structural_rank: usize,
497    p: usize,
498) -> SubspaceLeakageMetrics {
499    let mut max_abs_sq = 0.0_f64;
500    let mut max_rel_sq = 0.0_f64;
501    let mut worst_penalty = 0usize;
502
503    for (k, rs) in rs_transformed.iter().enumerate() {
504        let rows = rs.nrows();
505        let cols = rs.ncols().min(p);
506        let null_start = structural_rank.min(cols);
507        let mut abs_sq = 0.0_f64;
508        let mut total_sq = 0.0_f64;
509        for i in 0..rows {
510            for j in 0..cols {
511                let v = rs[(i, j)];
512                let vv = v * v;
513                total_sq += vv;
514                if j >= null_start {
515                    abs_sq += vv;
516                }
517            }
518        }
519        let rel_sq = if total_sq > 0.0 {
520            abs_sq / total_sq
521        } else {
522            0.0
523        };
524        if rel_sq > max_rel_sq {
525            max_rel_sq = rel_sq;
526            worst_penalty = k;
527        }
528        max_abs_sq = max_abs_sq.max(abs_sq);
529    }
530
531    let mut max_cross_gram_abs = 0.0_f64;
532    let null_count = p.saturating_sub(structural_rank);
533    if structural_rank > 0 && null_count > 0 {
534        for i in 0..structural_rank {
535            for j in 0..null_count {
536                let qn_col = structural_rank + j;
537                let mut dot = 0.0_f64;
538                for r in 0..p {
539                    dot += qs[(r, i)] * qs[(r, qn_col)];
540                }
541                max_cross_gram_abs = max_cross_gram_abs.max(dot.abs());
542            }
543        }
544    }
545
546    SubspaceLeakageMetrics {
547        max_abs_sq,
548        max_rel_sq,
549        worst_penalty,
550        max_cross_gram_abs,
551    }
552}
553
554/// True when the penalized/null subspace split is numerically self-consistent.
555///
556/// The split has two independent invariants:
557///
558/// 1. **Orthogonality** — `Qs = [Q_p | Q_n]` must be orthonormal, so the range
559///    and null blocks share no direction (`max |Qp'Qn| ≤ orth_tol`). This is the
560///    structural correctness guarantee and is checked at machine precision.
561///
562/// 2. **Bounded root leakage** — the transformed penalty root must have
563///    negligible energy on the null columns. The admissible relative-energy
564///    leakage is DERIVED from `REL_PSD_FLOOR`, the same numerically-PSD floor
565///    that `classify_eigenvalues_strict` uses to decide which modes are null:
566///    a mode the classifier is entitled to call null can, by the symmetric
567///    eigensolver's own eigenvector accuracy on a small-gap spectrum, retain up
568///    to `REL_PSD_FLOOR` of the penalty root's relative energy in that direction.
569///    Summed over the at-most-`p` near-threshold null modes the relative leakage
570///    cannot exceed `p · REL_PSD_FLOOR` without signalling a genuine
571///    (non-numerical) inconsistency, so that is the derived tolerance — matching
572///    the `p`-scaling `classify_eigenvalues_strict` already applies to its
573///    machine floor. Demanding a leakage tighter than the very floor that
574///    defined the null space is self-contradictory: it rejects well-posed smooth
575///    manifold / Duchon / sphere penalties whose Laplace-Beltrami spectrum decays
576///    through the classification threshold with no clean rank gap (#1802), even
577///    though the downstream penalty (`E'E`) is rebuilt with an EXACTLY clean null
578///    block regardless. The absolute floor keeps a vanishing-scale penalty from
579///    tripping the relative test on pure roundoff.
580fn subspace_split_is_consistent(leakage: &SubspaceLeakageMetrics, p: usize) -> bool {
581    let leakage_rel_tol = (p.max(1) as f64) * REL_PSD_FLOOR;
582    let leakage_abs_tol = 1e-12;
583    let orth_tol = 1e-10;
584    let root_leaks = leakage.max_rel_sq > leakage_rel_tol && leakage.max_abs_sq > leakage_abs_tol;
585    let split_nonorthogonal = leakage.max_cross_gram_abs > orth_tol;
586    !(root_leaks || split_nonorthogonal)
587}
588
589fn compose_qs_from_split(q_pen: &Mat<f64>, q_null: &Mat<f64>, p: usize) -> Mat<f64> {
590    let rank = q_pen.ncols();
591    let null_count = q_null.ncols();
592    let mut qs = Mat::<f64>::zeros(p, p);
593    for i in 0..p {
594        for j in 0..rank {
595            qs[(i, j)] = q_pen[(i, j)];
596        }
597        for j in 0..null_count {
598            qs[(i, rank + j)] = q_null[(i, j)];
599        }
600    }
601    qs
602}
603
604pub use crate::kronecker::kronecker_product;
605
606/// Result of the stable reparameterization algorithm from Wood (2011) Appendix B
607#[derive(Clone)]
608pub struct ReparamResult {
609    /// Penalty matrix in TRANSFORMED coefficient coordinates.
610    ///
611    /// This must be compatible with `beta_transformed` and `X_transformed = X * Qs`.
612    pub s_transformed: Array2<f64>,
613    /// Log-determinant of the penalty matrix (stable computation)
614    pub log_det: f64,
615    /// First derivatives of log-determinant w.r.t. log-smoothing parameters
616    pub det1: Array1<f64>,
617    /// Orthogonal transformation matrix Qs
618    pub qs: Array2<f64>,
619    /// Canonical penalties in the TRANSFORMED coordinate frame.
620    /// The single source of truth for penalty roots in the transformed frame.
621    /// Downstream consumers use these for block-local `PenaltyCoordinate`
622    /// construction, TK correction, and ext-coord paths.
623    pub canonical_transformed: Vec<CanonicalPenalty>,
624    /// Lambda-dependent penalty square root in TRANSFORMED coordinates (rank x p matrix).
625    /// This is used for applying the actual penalty in the least squares solve.
626    pub e_transformed: Array2<f64>,
627    /// Truncated eigenvectors (p × m where m = p - structural_rank).
628    ///
629    /// Coordinate frame note:
630    /// - This matrix is stored in the TRANSFORMED coefficient frame (post-`Qs`),
631    ///   i.e. it is compatible with `canonical_transformed`, `beta_transformed`,
632    ///   and transformed Hessians without additional coordinate mapping.
633    ///
634    /// These vectors span the structural null space used by positive-part
635    /// log-determinant conventions.
636    pub u_truncated: Array2<f64>,
637}
638
639/// The coefficient frame a set of penalty roots is expressed in.
640///
641/// The reparameterization's declared-null basis is stored in the TRANSFORMED
642/// (post-`Qs`) frame, so projecting an ORIGINAL-frame penalty against it
643/// requires rotating the basis by `Qs` first. Naming the frame at the call site
644/// is what keeps that rotation from being a coin flip.
645#[derive(Clone, Copy, Debug, PartialEq, Eq)]
646pub enum PenaltyFrame {
647    Original,
648    Transformed,
649}
650
651/// The λ-invariant declared-null subspace of a penalty reparameterization,
652/// materialized in both coefficient frames (#2454).
653///
654/// # The invariant this type exists to hold
655///
656/// `gam.reparam` keeps only balanced-penalty eigendirections above a relative
657/// rank tolerance and rebuilds the penalty the model applies as
658/// `S̃(λ) = E(λ)ᵀE(λ)` on that subspace alone. **Everything the criterion is
659/// made of is a function of `S̃`**: the inner solve minimizes `−ℓ + ½βᵀS̃β`,
660/// `H = −∇²ℓ + S̃`, and the reported penalty energy is `‖E β̂‖²`.
661///
662/// A per-block `S_k` whose own root rank exceeds the split's penalized rank
663/// therefore describes a DIFFERENT penalty. Two things go wrong if one is used
664/// anywhere the criterion is differentiated or normalized:
665///
666///  1. `½ λ_k β̂ᵀS_kβ̂` charges β̂'s energy in directions `S̃` never penalized —
667///     and `β̂` is free there, so that energy is `O(1)` and `∂/∂ρ_k` multiplies
668///     it by `λ_k`. That is an additive `c·λ` in the outer gradient, invisible
669///     at `‖ρ‖ ≤ 1` and sign-flipping it a dozen e-folds up.
670///  2. `−½log|Σ_k λ_k S_k|₊` charges MORE directions than `½log|H|` can ever
671///     inflate, because `H`'s penalty part is `S̃`. The two halves of the LAML
672///     ratio then saturate at different rates and the criterion acquires an
673///     asymptotic slope of `½(rank(S̃) − rank(S))` per unit ρ — unbounded
674///     below, with no interior optimum and no certifiable λ=∞ rail.
675///
676/// Projecting restores the identity every outer derivative is built on:
677/// `Σ_k λ_k (Π S_k Π) = Π (Σ_k λ_k S_k) Π = S̃(λ)`, and because `Π` is
678/// λ-invariant by construction, `∂S̃/∂ρ_k = λ_k · Π S_k Π` exactly. One penalty
679/// object then serves the value, the quadratic, the per-block scores, the
680/// `tr(H⁻¹Ḣ_k)` drift, the `log|S|₊` and its rank, and the outer Hessian.
681///
682/// A `None` basis means "nothing declared null" and every `project` is the
683/// identity, which is the case for any penalty whose numerical rank agrees
684/// with the split's.
685#[derive(Clone, Debug)]
686pub struct PenaltyNullSplit {
687    transformed: Option<Array2<f64>>,
688    original: Option<Array2<f64>>,
689}
690
691impl PenaltyNullSplit {
692    /// The identity split: nothing is declared null, every projection is a
693    /// no-op. This is what a penalty system with no dense reparameterization
694    /// (Kronecker marginal grids, sparse-native identity frames) carries.
695    pub fn identity() -> Self {
696        Self {
697            transformed: None,
698            original: None,
699        }
700    }
701
702    /// The number of declared-null directions, or 0 for the identity split.
703    pub fn declared_null_dim(&self) -> usize {
704        self.transformed.as_ref().map_or(0, |n| n.ncols())
705    }
706
707    /// The declared-null basis in `frame`, if this split declares anything.
708    pub fn basis(&self, frame: PenaltyFrame) -> Option<&Array2<f64>> {
709        match frame {
710            PenaltyFrame::Original => self.original.as_ref(),
711            PenaltyFrame::Transformed => self.transformed.as_ref(),
712        }
713    }
714
715    /// `Π S_k Π` for a canonical penalty expressed in `frame`.
716    ///
717    /// The penalty is returned unchanged when this split declares nothing null,
718    /// or when the basis does not match the penalty's dimension — a shape
719    /// disagreement the projection must not paper over by guessing.
720    pub fn project_canonical(
721        &self,
722        penalty: &CanonicalPenalty,
723        frame: PenaltyFrame,
724    ) -> Result<CanonicalPenalty, EstimationError> {
725        match self.basis(frame) {
726            Some(n) if n.nrows() == penalty.total_dim => {
727                penalty.project_out_null_directions(n.view())
728            }
729            _ => Ok(penalty.clone()),
730        }
731    }
732
733    /// `Π S_k Π` for a solver-side penalty coordinate expressed in `frame`.
734    pub fn project_coordinate(
735        &self,
736        coord: &gam_problem::PenaltyCoordinate,
737        frame: PenaltyFrame,
738    ) -> gam_problem::PenaltyCoordinate {
739        match self.basis(frame) {
740            Some(n) if n.nrows() == coord.dim() => coord.project_out_null_directions(n.view()),
741            _ => coord.clone(),
742        }
743    }
744}
745
746impl ReparamResult {
747    /// The λ-invariant subspace split this reparameterization declared.
748    ///
749    /// `u_truncated` is `p × m` in the transformed frame; the original-frame
750    /// basis is `Qs · u_truncated` (`Qs` orthogonal, and `u_truncated` is
751    /// itself defined as `Qsᵀ Q_null`). A degenerate or absent basis yields the
752    /// identity split rather than a guessed rotation.
753    pub fn null_split(&self) -> PenaltyNullSplit {
754        let u = &self.u_truncated;
755        if u.ncols() == 0 || u.nrows() == 0 {
756            return PenaltyNullSplit::identity();
757        }
758        let qs = &self.qs;
759        let original = if qs.nrows() == u.nrows() && qs.ncols() == u.nrows() {
760            qs.dot(u)
761        } else {
762            // No usable `Qs`: the frames coincide (sparse-native / identity
763            // reparameterization) or the shapes are inconsistent, in which case
764            // `project` declines on the dimension check anyway.
765            u.clone()
766        };
767        PenaltyNullSplit {
768            transformed: Some(u.clone()),
769            original: Some(original),
770        }
771    }
772}
773
774// ---------------------------------------------------------------------------
775// Kronecker factor decomposition primitives
776// ---------------------------------------------------------------------------
777
778/// Per-factor decomposition result for Kronecker penalties.
779struct KroneckerFactorDecomp {
780    root: Array2<f64>,              // rank_j × q_j
781    positive_eigenvalues: Vec<f64>, // length = rank_j
782    rank: usize,
783    dim: usize,
784}
785
786/// Eigendecompose each Kronecker factor separately at O(Σ q_j³).
787/// Returns per-factor decompositions, or `None` if any factor is zero.
788fn decompose_kronecker_factors(
789    factors: &[Array2<f64>],
790    context: &str,
791) -> Result<Option<Vec<KroneckerFactorDecomp>>, EstimationError> {
792    let mut decomps = Vec::with_capacity(factors.len());
793    for (j, factor) in factors.iter().enumerate() {
794        let q_j = factor.nrows();
795        if q_j != factor.ncols() {
796            crate::bail_invalid_estim!(
797                "{context}: Kronecker factor {j} must be square, got {}x{}",
798                factor.nrows(),
799                factor.ncols()
800            );
801        }
802        let is_identity = {
803            let mut is_id = true;
804            'outer: for r in 0..q_j {
805                for c in 0..q_j {
806                    let expected = if r == c { 1.0 } else { 0.0 };
807                    if (factor[[r, c]] - expected).abs() > 1e-12 {
808                        is_id = false;
809                        break 'outer;
810                    }
811                }
812            }
813            is_id
814        };
815        if is_identity {
816            decomps.push(KroneckerFactorDecomp {
817                root: Array2::eye(q_j),
818                positive_eigenvalues: vec![1.0; q_j],
819                rank: q_j,
820                dim: q_j,
821            });
822            continue;
823        }
824        let analysis = analyze_penalty_block(factor).map_err(|err| {
825            EstimationError::InvalidInput(format!(
826                "{context}: Kronecker factor {j} eigendecomp failed: {err}"
827            ))
828        })?;
829        if analysis.rank == 0 {
830            return Ok(None);
831        }
832        // Build the factor root from ONLY the range (positive-curvature)
833        // directions via the canonical classifier — never the null or
834        // negative-curvature directions (#1425).
835        let factor_classes =
836            crate::basis::SpectralClassification::new(
837                &analysis.eigenvalues,
838                analysis.rank_tol,
839                analysis.noise_tol,
840            );
841        let mut root_j = Array2::zeros((analysis.rank, q_j));
842        let mut pos_eigs = Vec::with_capacity(analysis.rank);
843        for (row_idx, &i) in factor_classes.range_idx.iter().enumerate() {
844            let eigenval = analysis.eigenvalues[i];
845            let sqrt_ev = eigenval.sqrt();
846            let evec = analysis.eigenvectors.column(i);
847            for (col, &v) in evec.iter().enumerate() {
848                root_j[[row_idx, col]] = sqrt_ev * v;
849            }
850            pos_eigs.push(eigenval);
851        }
852        decomps.push(KroneckerFactorDecomp {
853            root: root_j,
854            positive_eigenvalues: pos_eigs,
855            rank: analysis.rank,
856            dim: q_j,
857        });
858    }
859    Ok(Some(decomps))
860}
861
862/// Build the block-local Kronecker root from pre-computed factor decompositions.
863fn assemble_kronecker_root_local(decomps: &[KroneckerFactorDecomp]) -> Array2<f64> {
864    let mut kron_root = decomps[0].root.clone();
865    for fr in &decomps[1..] {
866        let (r1, c1) = kron_root.dim();
867        let (r2, c2) = (fr.rank, fr.dim);
868        let mut new_root = Array2::zeros((r1 * r2, c1 * c2));
869        for i1 in 0..r1 {
870            for i2 in 0..r2 {
871                for j1 in 0..c1 {
872                    for j2 in 0..c2 {
873                        new_root[[i1 * r2 + i2, j1 * c2 + j2]] =
874                            kron_root[[i1, j1]] * fr.root[[i2, j2]];
875                    }
876                }
877            }
878        }
879        kron_root = new_root;
880    }
881    kron_root
882}
883
884/// Compute eigenvalues of the Kronecker product from per-factor eigenvalues.
885fn kronecker_eigenvalues(decomps: &[KroneckerFactorDecomp], block_dim: usize) -> (Vec<f64>, usize) {
886    let mut kron_eigs = decomps[0].positive_eigenvalues.clone();
887    for fd in &decomps[1..] {
888        let mut new_eigs = Vec::with_capacity(kron_eigs.len() * fd.positive_eigenvalues.len());
889        for &a in &kron_eigs {
890            for &b in &fd.positive_eigenvalues {
891                new_eigs.push(a * b);
892            }
893        }
894        kron_eigs = new_eigs;
895    }
896    let max_ev = kron_eigs.iter().copied().fold(0.0_f64, f64::max);
897    let tol = max_ev * 1e-10 * (block_dim as f64);
898    let positive: Vec<f64> = kron_eigs.into_iter().filter(|&ev| ev > tol).collect();
899    let nullity = block_dim - positive.len();
900    (positive, nullity)
901}
902
903// ---------------------------------------------------------------------------
904// CanonicalPenalty — block-local processed penalty for the solver
905// ---------------------------------------------------------------------------
906
907/// A canonicalized penalty with block-local root, ready for the solver.
908///
909/// Instead of storing a full `p x p` penalty matrix, this stores only the
910/// `rank x block_dim` root and the column range, enabling O(p_k^2) operations
911/// instead of O(p^2).
912#[derive(Clone)]
913pub struct CanonicalPenalty {
914    /// Square root matrix: S_k = root^T * root.
915    /// Shape: `rank x block_dim` for block-local, `rank x p` for dense.
916    pub root: Array2<f64>,
917    /// Column range in the global coefficient vector [start..end).
918    /// For dense penalties this is `0..p`.
919    pub col_range: std::ops::Range<usize>,
920    /// Full parameter dimension p.
921    pub total_dim: usize,
922    /// Structural nullity of the local penalty.
923    pub nullity: usize,
924    /// The symmetrized block-local penalty matrix (block_dim × block_dim).
925    /// Cached at construction time to avoid recomputing root^T * root
926    /// in hot paths (penalty assembly, trace products).
927    pub local: Array2<f64>,
928    /// Block-local prior mean used to center this penalty.
929    pub prior_mean: Array1<f64>,
930    /// Positive eigenvalues of the local penalty matrix (length = rank).
931    /// Cached at construction time for REML logdet block-factored paths.
932    pub positive_eigenvalues: Vec<f64>,
933    /// Optional operator-form handle bit-equivalent to `local`. Propagated
934    /// from `PenaltySpec::Block.op`. Downstream PIRLS and REML exact operator
935    /// algebra route through this for dense-Gram-free matvec when present.
936    pub op: Option<std::sync::Arc<dyn crate::analytic_penalties::PenaltyOp>>,
937}
938
939impl std::fmt::Debug for CanonicalPenalty {
940    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
941        f.debug_struct("CanonicalPenalty")
942            .field(
943                "root",
944                &format_args!("{}×{}", self.root.nrows(), self.root.ncols()),
945            )
946            .field("col_range", &self.col_range)
947            .field("total_dim", &self.total_dim)
948            .field("nullity", &self.nullity)
949            .field(
950                "local",
951                &format_args!("{}×{}", self.local.nrows(), self.local.ncols()),
952            )
953            .field("prior_mean_len", &self.prior_mean.len())
954            .field("positive_eigenvalues", &self.positive_eigenvalues)
955            .field("op", &self.op.as_ref().map(|o| o.dim()))
956            .finish()
957    }
958}
959
960impl CanonicalPenalty {
961    /// Construct a dense (full-width) canonical penalty from a `rank x p` root.
962    /// Used to wrap reparam-transformed roots for consumers that expect
963    /// `&[CanonicalPenalty]`.
964    pub fn from_dense_root(root: Array2<f64>, p: usize) -> Self {
965        Self::from_dense_root_with_mean(root, p, Array1::zeros(p))
966    }
967
968    pub fn from_dense_root_with_mean(root: Array2<f64>, p: usize, prior_mean: Array1<f64>) -> Self {
969        assert_eq!(prior_mean.len(), p);
970        let local = root.t().dot(&root);
971        let positive_eigenvalues = Vec::new(); // not needed for TK paths
972        Self {
973            root,
974            col_range: 0..p,
975            total_dim: p,
976            nullity: 0,
977            local,
978            prior_mean,
979            positive_eigenvalues,
980            op: None,
981        }
982    }
983
984    /// Embed the block-local root into a full-width `rank × total_dim` matrix.
985    /// For dense penalties (col_range = 0..p), returns the root unchanged.
986    pub fn full_width_root(&self) -> Array2<f64> {
987        if self.col_range.start == 0 && self.col_range.end == self.total_dim {
988            return self.root.clone();
989        }
990        let rank = self.root.nrows();
991        let mut full = Array2::<f64>::zeros((rank, self.total_dim));
992        full.slice_mut(ndarray::s![.., self.col_range.clone()])
993            .assign(&self.root);
994        full
995    }
996
997    /// Numerical rank of this penalty.
998    pub fn rank(&self) -> usize {
999        self.root.nrows()
1000    }
1001
1002    /// Block dimension (number of columns this penalty covers).
1003    pub fn block_dim(&self) -> usize {
1004        self.col_range.len()
1005    }
1006
1007    /// Whether this penalty is block-local (col_range != 0..total_dim).
1008    pub const fn is_block_local(&self) -> bool {
1009        self.col_range.start != 0 || self.col_range.end != self.total_dim
1010    }
1011
1012    /// Return a reference to the cached local penalty matrix.
1013    /// Shape: `block_dim x block_dim`.
1014    pub fn local_ref(&self) -> &Array2<f64> {
1015        &self.local
1016    }
1017
1018    /// Return an owned copy of the local penalty matrix.
1019    /// Prefer `local_ref()` when a reference suffices.
1020    pub fn local_penalty(&self) -> Array2<f64> {
1021        self.local.clone()
1022    }
1023
1024    /// Accumulate lambda * S_k into a pre-allocated `p x p` target matrix.
1025    /// Only touches the block [col_range × col_range].
1026    pub fn accumulate_weighted(&self, target: &mut Array2<f64>, lambda: f64) {
1027        if lambda == 0.0 || self.rank() == 0 {
1028            return;
1029        }
1030        let r = &self.col_range;
1031        target
1032            .slice_mut(s![r.start..r.end, r.start..r.end])
1033            .scaled_add(lambda, &self.local);
1034    }
1035
1036    /// Compute `scale * v^T S_k v` (quadratic form).
1037    /// Only reads `v[start..end]` — O(rank × block_dim) not O(rank × p).
1038    pub fn quadratic(&self, v: &Array1<f64>, scale: f64) -> f64 {
1039        if self.rank() == 0 || scale == 0.0 {
1040            return 0.0;
1041        }
1042        let v_block = v.slice(s![self.col_range.start..self.col_range.end]);
1043        let rv = self.root.dot(&v_block);
1044        scale * rv.dot(&rv)
1045    }
1046
1047    /// Compute `scale * S_k * prior_mean` embedded into the global basis.
1048    pub fn prior_linear_shift(&self, scale: f64) -> Array1<f64> {
1049        let mut out = Array1::<f64>::zeros(self.total_dim);
1050        if self.rank() == 0 || scale == 0.0 || self.prior_mean.iter().all(|&v| v == 0.0) {
1051            return out;
1052        }
1053        let block = self.local.dot(&self.prior_mean) * scale;
1054        out.slice_mut(s![self.col_range.start..self.col_range.end])
1055            .assign(&block);
1056        out
1057    }
1058
1059    /// Compute `scale * prior_mean' S_k prior_mean`.
1060    pub fn prior_constant_shift(&self, scale: f64) -> f64 {
1061        if self.rank() == 0 || scale == 0.0 || self.prior_mean.iter().all(|&v| v == 0.0) {
1062            return 0.0;
1063        }
1064        scale * self.prior_mean.dot(&self.local.dot(&self.prior_mean))
1065    }
1066
1067    /// Embed this block's prior mean into the global coefficient basis.
1068    pub fn full_width_prior_mean(&self) -> Array1<f64> {
1069        if self.col_range.start == 0 && self.col_range.end == self.total_dim {
1070            return self.prior_mean.clone();
1071        }
1072        let mut out = Array1::<f64>::zeros(self.total_dim);
1073        out.slice_mut(s![self.col_range.start..self.col_range.end])
1074            .assign(&self.prior_mean);
1075        out
1076    }
1077
1078    /// `Π S_k Π` for the declared-null basis `N` (orthonormal, `total_dim × m`)
1079    /// in THIS penalty's coefficient frame, with `Π = I − N Nᵀ` (#2454).
1080    ///
1081    /// See [`PenaltyNullSplit`] for why every penalty the criterion touches has
1082    /// to be projected onto the reparameterization's penalized subspace. This
1083    /// is the `CanonicalPenalty` face of [`gam_problem::PenaltyCoordinate::
1084    /// project_out_null_directions`]; both go through the same root primitive.
1085    ///
1086    /// Returns `self` unchanged — `op` handle, cached spectrum and all — when
1087    /// the projection is below the root's own representation noise
1088    /// (`‖R N‖_max ≤ ‖R‖_max · ε · p`). That is not a shortcut but a statement
1089    /// of exactness: a null direction the root does not resolve carries no
1090    /// energy the criterion's `log|S|₊` can see either, so the projected and
1091    /// unprojected penalties agree to every digit either object represents.
1092    /// It is also the ordinary case — for any penalty whose structural null
1093    /// space is exact, `N` spans `∩_k null(S_k)` and `Π S_k Π = S_k` in exact
1094    /// arithmetic.
1095    ///
1096    /// When the projection does bite, block locality is preserved if the null
1097    /// basis is itself block-local, `local` is rebuilt as `RᵀR` from the
1098    /// projected root, and `positive_eigenvalues` becomes the projected root's
1099    /// squared singular values — so `Σ positive_eigenvalues = tr(Π S_k Π)`
1100    /// stays exact for the consumers that read it as a trace. The `op` handle
1101    /// is dropped: it is declared bit-equivalent to `local`, and the projected
1102    /// `local` is a different operator.
1103    pub fn project_out_null_directions(
1104        &self,
1105        null_basis: ndarray::ArrayView2<'_, f64>,
1106    ) -> Result<Self, EstimationError> {
1107        if null_basis.ncols() == 0 || self.rank() == 0 {
1108            return Ok(self.clone());
1109        }
1110        if null_basis.nrows() != self.total_dim {
1111            return Err(EstimationError::LayoutError(format!(
1112                "CanonicalPenalty::project_out_null_directions: null-basis row count {} does \
1113                 not match the penalty's total dimension {}",
1114                null_basis.nrows(),
1115                self.total_dim
1116            )));
1117        }
1118        let root_scale = self.root.iter().fold(0.0_f64, |a, v| a.max(v.abs()));
1119        let (projected_full, stays_block_local) =
1120            gam_problem::project_block_root_out_of_null_directions(
1121                self.root.view(),
1122                self.col_range.start,
1123                self.col_range.end,
1124                self.total_dim,
1125                null_basis,
1126            );
1127        // How much the projection actually moved, against the root's own
1128        // representation noise.
1129        let mut moved = 0.0_f64;
1130        {
1131            let original_full = self.full_width_root();
1132            for (a, b) in projected_full.iter().zip(original_full.iter()) {
1133                moved = moved.max((a - b).abs());
1134            }
1135        }
1136        if moved <= root_scale * f64::EPSILON * (self.total_dim as f64) {
1137            return Ok(self.clone());
1138        }
1139
1140        let (root, col_range) = if stays_block_local {
1141            (
1142                projected_full
1143                    .slice(s![.., self.col_range.start..self.col_range.end])
1144                    .to_owned(),
1145                self.col_range.clone(),
1146            )
1147        } else {
1148            (projected_full, 0..self.total_dim)
1149        };
1150        let prior_mean = if col_range == self.col_range {
1151            self.prior_mean.clone()
1152        } else {
1153            let mut full = Array1::<f64>::zeros(self.total_dim);
1154            full.slice_mut(s![self.col_range.start..self.col_range.end])
1155                .assign(&self.prior_mean);
1156            full
1157        };
1158        // `RᵀR` eigenvalues are the squared singular values of `R`; taking them
1159        // from the root rather than from an eigendecomposition of the assembled
1160        // `local` keeps the small ones (the whole point of a root-scale
1161        // representation) and costs O(rank² · block_dim).
1162        let positive_eigenvalues = match root.svd(false, false) {
1163            Ok((_, singular_values, _)) => {
1164                let mut eigs = vec![0.0_f64; root.nrows()];
1165                for (slot, &sigma) in eigs.iter_mut().zip(singular_values.iter()) {
1166                    *slot = sigma * sigma;
1167                }
1168                eigs
1169            }
1170            Err(error) => {
1171                return Err(EstimationError::LayoutError(format!(
1172                    "CanonicalPenalty::project_out_null_directions: SVD of the projected root \
1173                     failed: {error:?}"
1174                )));
1175            }
1176        };
1177        let local = root.t().dot(&root);
1178        Ok(Self {
1179            root,
1180            col_range,
1181            total_dim: self.total_dim,
1182            // The projection can only remove curvature, so the block's nullity
1183            // can only grow; the honest floor is what it already declared.
1184            nullity: self.nullity,
1185            local,
1186            prior_mean,
1187            positive_eigenvalues,
1188            op: None,
1189        })
1190    }
1191
1192    /// Convert to a PenaltyCoordinate for the unified REML evaluator.
1193    pub fn to_penalty_coordinate(&self) -> gam_problem::PenaltyCoordinate {
1194        use gam_problem::PenaltyCoordinate;
1195        if self.is_block_local() {
1196            PenaltyCoordinate::from_block_root_with_mean(
1197                self.root.clone(),
1198                self.col_range.start,
1199                self.col_range.end,
1200                self.total_dim,
1201                self.prior_mean.clone(),
1202            )
1203        } else {
1204            PenaltyCoordinate::from_dense_root_with_mean(self.root.clone(), self.prior_mean.clone())
1205        }
1206    }
1207}
1208
1209/// Measure and report how close each pair of penalties in the canonical bundle
1210/// comes to being proportional.
1211///
1212/// Two penalties `S_i`, `S_j` with the same `col_range` are compared by the
1213/// **relative defect of the proportionality claim itself**,
1214///
1215/// ```text
1216///     delta(S_i, S_j) = min_c ||S_j - c S_i||_F / ||S_i||_F,
1217///     attained at      c = <S_i, S_j>_F / ||S_i||_F^2,
1218/// ```
1219///
1220/// which is what `sum_i w_i S_i = 0` actually asks about. The cosine is the
1221/// same information (`delta^2 = 1 - cos^2` for unit-normalized operators), but
1222/// it is the WRONG COORDINATE to report it in and this issue is the cost of
1223/// having done so: `delta` enters the cosine SQUARED, so a pair that is
1224/// `1.9e-5` apart — five orders above any arithmetic — prints as
1225/// `cos = 1.000000` and reads as an exact identity. #2676 ran on that reading
1226/// for its whole life. Measured on `geo_disease_matern` (centers=10, n=1500,
1227/// n_pcs=16, the #2676 penalty-map probe):
1228///
1229/// ```text
1230///     length_scale   delta       certified nullity
1231///       2.05e-2      2.079e-15          1     <- exact, to the residual's own round-off
1232///       1.64e-1      1.874e-5           0     <- prints cos = 1.000000
1233///       1.27e0       3.396e-1           0     <- what the fit realizes
1234/// ```
1235///
1236/// so the "structural identity" that issue is named for is the small-length-
1237/// scale LIMIT of two genuinely different operators, and it is not present at
1238/// the geometry any of those fits settle on.
1239/// Because `local` is symmetric, `<A, B>_F = sum_{r,c} A[r,c] * B[r,c] =
1240/// tr(A·B)`. Pairs with different `col_range` cannot be proportional by
1241/// construction and are skipped.
1242///
1243/// # ⚠ This detector is a SCREEN, not the authority (#2676)
1244///
1245/// It sees proportional PAIRS. The criterion's actual invariance is the null
1246/// space of the penalty map's Gram, `{w : sum_i w_i S_i = 0}`, which contains
1247/// every linear redundancy — including three-term ones no pairwise measure can
1248/// find. `gam_solve::penalty_invariance::PenaltyMapInvariance` computes that
1249/// subspace and is what the curvature certificate and the smoothing correction
1250/// consult. Use this function for the human-facing report; do not use it to
1251/// decide anything numeric.
1252///
1253/// # ⚠ And an EXACT redundancy is not a saddle
1254///
1255/// This warning used to say the redundancy makes the LAML cost carry "a
1256/// Z₂-symmetric saddle". It does not. Where the redundancy is EXACT the
1257/// criterion is exactly constant along it in `lambda` — a flat direction, not
1258/// a descending one — and every apparent negative curvature there is the
1259/// chain-rule term of
1260/// `H_rho = diag(lambda) H_lambda diag(lambda) + diag(g_rho)`, whose sign is
1261/// the sign of a rounding residual. An optimizer converging there has not
1262/// converged to a saddle; it has converged to a point on a manifold of points
1263/// that all give the SAME fit, because they all assemble the same penalty.
1264/// That is a non-identifiability of `lambda`, worth reporting, and not a defect
1265/// in the answer.
1266///
1267/// Where the redundancy is only NEAR-exact none of that holds: the criterion is
1268/// not constant along the direction, it carries genuine curvature of order
1269/// `delta^2`, and calling it "structurally identical" asserts a flatness the
1270/// evidence does not support. Which of the two a pair is is decided below, at
1271/// the residual's own arithmetic floor, and it is said in the message.
1272///
1273/// # The exactness bar, and why it is derived rather than chosen
1274///
1275/// `delta` is computed as a residual norm over `m = block_dim^2` products of
1276/// `O(1)` entries, so its own round-off is `sqrt(m) * EPSILON` — the error of
1277/// the norm, not of the operators. A pair at or under that is proportional to
1278/// everything this arithmetic can see; a pair above it is measurably not.
1279/// Measured, the two populations sit ten orders apart (`2.079e-15` against
1280/// `1.874e-5` on the sweep above, at `sqrt(m) * EPSILON = 2.2e-15` for
1281/// `block_dim = 10`), so nothing rides on where in that gap the bar falls.
1282///
1283/// Logging policy, `delta`-denominated throughout:
1284/// - `delta <= sqrt(m) * EPSILON` → `log::warn!` with `[PENALTY-REDUNDANCY]`.
1285///   Only their combination `lambda_i + c*lambda_j` is identified.
1286/// - `sqrt(m) * EPSILON < delta <= 1e-1` → `log::info!` with
1287///   `[PENALTY-SIMILARITY]`, carrying `delta`. At large scale (`k > 64`) only
1288///   the three smallest-`delta` such pairs are logged to bound log volume.
1289///
1290/// Returns `Vec<(i, j, delta)>` for the **exactly proportional** pairs,
1291/// primarily to make this function unit-testable without a log capture.
1292///
1293/// Performance: this is O(k² · block_dim²); intended to be called exactly
1294/// once per fit (e.g. from `RemlState::newwith_offset_shared`).
1295pub fn report_penalty_pair_redundancy(canonical: &[CanonicalPenalty]) -> Vec<(usize, usize, f64)> {
1296    // A pair `delta` above this is measurably NOT proportional, and calling it
1297    // "similar" past a tenth of the operator's own size says nothing. Purely a
1298    // log-volume bound: no verdict is taken on it.
1299    const SIMILARITY_REPORTING_DEFECT: f64 = 1e-1;
1300    const LARGE_SCALE_K_THRESHOLD: usize = 64;
1301    const TOP_SIMILARITY_PAIRS: usize = 3;
1302
1303    let k = canonical.len();
1304    let mut redundant: Vec<(usize, usize, f64)> = Vec::new();
1305    let mut similar: Vec<(usize, usize, f64, f64)> = Vec::new();
1306
1307    // Pre-compute tr(S_i^2) = sum of squares of S_i entries (Frobenius norm
1308    // squared). `local` is symmetric, so this equals tr(S_i^T S_i) = tr(S_i^2).
1309    let trace_sq: Vec<f64> = canonical
1310        .iter()
1311        .map(|p| p.local.iter().map(|&v| v * v).sum::<f64>())
1312        .collect();
1313
1314    for i in 0..k {
1315        if trace_sq[i] == 0.0 {
1316            continue;
1317        }
1318        for j in (i + 1)..k {
1319            if trace_sq[j] == 0.0 {
1320                continue;
1321            }
1322            // Different col_range → cannot be proportional by construction (the
1323            // block-local matrices live in disjoint or mismatched parameter
1324            // subspaces).
1325            if canonical[i].col_range != canonical[j].col_range {
1326                continue;
1327            }
1328            // Shapes must match — they do when col_range matches because
1329            // `local` is `block_dim × block_dim` and `block_dim = col_range.len()`.
1330            assert_eq!(canonical[i].local.dim(), canonical[j].local.dim());
1331
1332            let inner: f64 = canonical[i]
1333                .local
1334                .iter()
1335                .zip(canonical[j].local.iter())
1336                .map(|(&a, &b)| a * b)
1337                .sum();
1338            // The least-squares proportionality constant and the residual it
1339            // leaves, formed DIRECTLY rather than through `1 - cos`: the cosine
1340            // route squares `delta` and then subtracts from one, which loses
1341            // every digit of a defect under `sqrt(EPSILON)` — the loss this
1342            // whole issue is made of.
1343            let scale = inner / trace_sq[i];
1344            let residual_sq: f64 = canonical[i]
1345                .local
1346                .iter()
1347                .zip(canonical[j].local.iter())
1348                .map(|(&a, &b)| {
1349                    let residual = b - scale * a;
1350                    residual * residual
1351                })
1352                .sum();
1353            let defect = (residual_sq / trace_sq[i]).sqrt();
1354            let entries = canonical[i].local.len() as f64;
1355            let arithmetic_floor = entries.sqrt() * f64::EPSILON;
1356
1357            if defect <= arithmetic_floor {
1358                redundant.push((i, j, defect));
1359            } else if defect <= SIMILARITY_REPORTING_DEFECT {
1360                similar.push((i, j, defect, scale));
1361            }
1362        }
1363    }
1364
1365    // Always emit every exact redundancy — these are structural model errors.
1366    for &(i, j, defect) in &redundant {
1367        log::warn!(
1368            "[PENALTY-REDUNDANCY] penalties i={i} j={j} are proportional to the arithmetic \
1369             that formed them (relative defect min_c ||S_{j} - c S_{i}||_F / ||S_{i}||_F = \
1370             {defect:.6e}) — only their COMBINATION is identified, so the criterion is exactly \
1371             constant along their antisymmetric direction and lambda_{i} / lambda_{j} are not \
1372             separately estimable. The fit itself is unaffected (every point of that manifold \
1373             assembles the same penalty), and the curvature certificate deflates the direction \
1374             rather than judging a chain-rule term there (#2676). Consider re-specifying (e.g. \
1375             anisotropic→isotropic for spatial smoothers with weak axis signal) if you need the \
1376             individual smoothing parameters to mean something."
1377        );
1378    }
1379
1380    // Cap similarity log volume at large scale.
1381    if k > LARGE_SCALE_K_THRESHOLD && similar.len() > TOP_SIMILARITY_PAIRS {
1382        similar.sort_by(|a, b| a.2.partial_cmp(&b.2).unwrap_or(std::cmp::Ordering::Equal));
1383        similar.truncate(TOP_SIMILARITY_PAIRS);
1384    }
1385    for (i, j, defect, scale) in similar {
1386        log::info!(
1387            "[PENALTY-SIMILARITY] penalties i={i} j={j} are close but MEASURABLY distinct \
1388             (relative defect {defect:.6e} at the best scale c={scale:.6e}) — the outer Hessian \
1389             may be ill-conditioned along their antisymmetric direction, and the criterion \
1390             carries genuine curvature of order defect^2 there. This is NOT an invariance: \
1391             nothing is exactly constant along it and no direction is deflated for it (#2676)."
1392        );
1393    }
1394
1395    redundant
1396}
1397
1398/// Canonicalize a single `PenaltySpec` into a `CanonicalPenalty` by computing
1399/// the block-local eigendecomposition and extracting the root.
1400///
1401/// This is O(block_dim^3) instead of O(p^3) for block-local penalties.
1402/// Returns `None` if the penalty has rank zero (should be dropped).
1403pub fn canonicalize_penalty_spec(
1404    spec: &crate::PenaltySpec,
1405    p: usize,
1406    idx: usize,
1407    context: &str,
1408) -> Result<Option<CanonicalPenalty>, EstimationError> {
1409    use crate::PenaltySpec;
1410
1411    crate::validate_penalty_spec_shape(idx, spec, p, context)?;
1412
1413    let (local_matrix, col_range, prior_mean_spec, hint, op) = match spec {
1414        PenaltySpec::Block {
1415            local,
1416            col_range,
1417            prior_mean,
1418            structure_hint,
1419            op,
1420        } => (
1421            local.view(),
1422            col_range.clone(),
1423            prior_mean,
1424            structure_hint.as_ref(),
1425            op.clone(),
1426        ),
1427        PenaltySpec::Dense(m) => (
1428            m.view(),
1429            0..p,
1430            &gam_problem::CoefficientPriorMean::Zero,
1431            None,
1432            None,
1433        ),
1434        PenaltySpec::DenseWithMean { matrix, prior_mean } => {
1435            (matrix.view(), 0..p, prior_mean, None, None)
1436        }
1437    };
1438
1439    let block_dim = col_range.len();
1440    let prior_mean = prior_mean_spec
1441        .evaluate(block_dim, &format!("{context}: penalty {idx}"))
1442        .map_err(|e| EstimationError::InvalidInput(e.0))?;
1443
1444    // ── Ridge fast path: closed-form, no eigendecomposition ──
1445    if let Some(PenaltyStructureHint::Ridge(scale)) = hint {
1446        if *scale <= 0.0 {
1447            return Ok(None);
1448        }
1449        let sqrt_scale = scale.sqrt();
1450        let mut root = Array2::zeros((block_dim, block_dim));
1451        for i in 0..block_dim {
1452            root[[i, i]] = sqrt_scale;
1453        }
1454        // Ridge penalties are diagonal by construction, but still route through
1455        // the crate-wide ndarray symmetrizer so every construction variant uses
1456        // the same "average the transpose" cleanup instead of a local copy.
1457        let mut local_sym = local_matrix.to_owned();
1458        symmetrize_in_place(&mut local_sym);
1459        return Ok(Some(CanonicalPenalty {
1460            root,
1461            col_range,
1462            total_dim: p,
1463            nullity: 0,
1464            local: local_sym,
1465            prior_mean,
1466            positive_eigenvalues: vec![*scale; block_dim],
1467            op,
1468        }));
1469    }
1470
1471    // ── Kronecker fast path: single per-factor eigendecomposition ──
1472    if let Some(PenaltyStructureHint::Kronecker(factors)) = hint {
1473        let decomps =
1474            match decompose_kronecker_factors(factors, &format!("{context} penalty {idx}"))? {
1475                None => return Ok(None),
1476                Some(d) => d,
1477            };
1478        let (positive_eigenvalues, nullity) = kronecker_eigenvalues(&decomps, block_dim);
1479        if positive_eigenvalues.is_empty() {
1480            return Ok(None);
1481        }
1482        let root = assemble_kronecker_root_local(&decomps);
1483        let mut local_sym = local_matrix.to_owned();
1484        symmetrize_in_place(&mut local_sym);
1485        return Ok(Some(CanonicalPenalty {
1486            root,
1487            col_range,
1488            total_dim: p,
1489            nullity,
1490            local: local_sym,
1491            prior_mean,
1492            positive_eigenvalues,
1493            op,
1494        }));
1495    }
1496
1497    // ── Generic block-local path: eigendecompose at O(block_dim³) ──
1498    let local_owned = local_matrix.to_owned();
1499    let analysis = analyze_penalty_block(&local_owned).map_err(|err| {
1500        EstimationError::InvalidInput(format!(
1501            "{context}: penalty canonicalization failed at index {idx}: {err}"
1502        ))
1503    })?;
1504
1505    if analysis.rank == 0 {
1506        log::debug!(
1507            "Dropped inactive penalty block idx={idx} reason={}",
1508            if analysis.iszero {
1509                "ZeroMatrix"
1510            } else {
1511                "NumericalRankZero"
1512            }
1513        );
1514        return Ok(None);
1515    }
1516
1517    // Reuse the eigendecomposition from analyze_penalty_block and route the
1518    // range / null / negative-curvature split through the one canonical
1519    // classifier, so this root construction cannot disagree with the block's
1520    // own `rank` / `nullity` / `negative_dim` about which directions are
1521    // penalized, unpenalized, or non-PSD (#1425).
1522    let tolerance = analysis.rank_tol;
1523    let classes = crate::basis::SpectralClassification::new(
1524        &analysis.eigenvalues,
1525        tolerance,
1526        analysis.noise_tol,
1527    );
1528    let rank_k = classes.rank();
1529    assert_eq!(
1530        rank_k, analysis.rank,
1531        "penalty-root rank disagreement: SpectralClassification rank={rank_k} vs analyze_penalty_block rank={} (#1425 canonical-classifier invariant)",
1532        analysis.rank
1533    );
1534
1535    // Build the penalty root R from ONLY the range directions (positive
1536    // curvature): R has one row per range eigenpair, scaled by sqrt(ev), so
1537    // RᵀR reconstructs S on range(S). Null directions contribute nothing
1538    // (their eigenvalue is zero); negative-curvature directions are NEVER
1539    // square-rooted into R (their sqrt is imaginary) and are NOT null — they
1540    // are simply dropped from R, exactly as the closed-form Duchon kernels at
1541    // high d require to preserve the q_pen / q_null invariant downstream.
1542    let mut root = Array2::zeros((rank_k, block_dim));
1543    let mut positive_eigenvalues = Vec::with_capacity(rank_k);
1544    for (row_idx, &i) in classes.range_idx.iter().enumerate() {
1545        let eigenval = analysis.eigenvalues[i];
1546        let eigenvec = analysis.eigenvectors.column(i);
1547        root.row_mut(row_idx).assign(&(&eigenvec * eigenval.sqrt()));
1548        positive_eigenvalues.push(eigenval);
1549    }
1550
1551    // Surface any genuine negative curvature honestly: it is neither range
1552    // (dropped from R) nor null (excluded from `nullity`), so it would
1553    // otherwise vanish without a trace. A non-PSD penalty reaching this path
1554    // is a real geometric fact (e.g. high-d Duchon kernels) the operator
1555    // should be able to see.
1556    if classes.is_indefinite() {
1557        log::debug!(
1558            "{context}: penalty block idx={idx} carries {} negative-curvature \
1559             eigendirection(s) below -tol={tolerance:e}; dropped from the canonical \
1560             root and NOT counted as null space (rank={rank_k}, nullity={})",
1561            classes.negative_dim(),
1562            classes.nullity()
1563        );
1564    }
1565
1566    // Store the PSD reconstruction RᵀR rather than the raw symmetrised input so
1567    // the cached `local` matches the rank truncation embedded in `root`
1568    // (negative-curvature directions are excluded from both, as above).
1569    let local = root.t().dot(&root);
1570    Ok(Some(CanonicalPenalty {
1571        root,
1572        col_range,
1573        total_dim: p,
1574        nullity: classes.nullity(),
1575        local,
1576        prior_mean,
1577        positive_eigenvalues,
1578        op,
1579    }))
1580}
1581
1582/// Canonicalize a batch of penalty specs, dropping zero-rank penalties.
1583/// Returns (active_penalties, active_nullspace_dims).
1584pub fn canonicalize_penalty_specs(
1585    specs: &[crate::PenaltySpec],
1586    nullspace_dims: &[usize],
1587    p: usize,
1588    context: &str,
1589) -> Result<(Vec<CanonicalPenalty>, Vec<usize>), EstimationError> {
1590    if specs.len() != nullspace_dims.len() {
1591        crate::bail_invalid_estim!(
1592            "{context}: nullspace_dims length mismatch: penalties={}, nullspace_dims={}",
1593            specs.len(),
1594            nullspace_dims.len()
1595        );
1596    }
1597
1598    let mut active = Vec::with_capacity(specs.len());
1599    let mut active_nullspace = Vec::with_capacity(specs.len());
1600    for (idx, spec) in specs.iter().enumerate() {
1601        if let Some(canonical) = canonicalize_penalty_spec(spec, p, idx, context)? {
1602            active_nullspace.push(nullspace_dims[idx]);
1603            active.push(canonical);
1604        }
1605    }
1606    Ok((active, active_nullspace))
1607}
1608
1609/// Hard cap on the dimension `p` allowed to fall back to a dense p × p
1610/// eigendecomposition of the overlapping balanced penalty.
1611///
1612/// Beyond this cap the overlapping-penalty path errors out instead of
1613/// allocating an O(p²) workspace whose eigendecomposition would dominate the
1614/// solve. ResourcePolicy threading is the long-term home for this cap (the
1615/// resource_serialize agent is widening ResourcePolicy coverage); until that
1616/// lands, both overlapping branches share this single constant so they can't
1617/// drift apart.
1618pub(crate) const OVERLAPPING_PENALTY_DENSE_FALLBACK_MAX_P: usize = 4096;
1619
1620/// Creates a balanced penalty root from canonical penalties.
1621///
1622/// When all penalties have non-overlapping col_ranges, the balanced sum is
1623/// block-diagonal and eigendecomposition is done per-block at O(Σ p_k³)
1624/// instead of the global O(p³). Falls back to the global path when penalties
1625/// overlap.
1626pub fn create_balanced_penalty_root_from_canonical(
1627    penalties: &[CanonicalPenalty],
1628    p: usize,
1629) -> Result<Array2<f64>, EstimationError> {
1630    if penalties.is_empty() {
1631        return Ok(Array2::zeros((0, p)));
1632    }
1633
1634    // Group penalties by col_range.
1635    let mut block_groups: BTreeMap<(usize, usize), Vec<&CanonicalPenalty>> = BTreeMap::new();
1636    for cp in penalties {
1637        if cp.rank() == 0 {
1638            continue;
1639        }
1640        let key = (cp.col_range.start, cp.col_range.end);
1641        block_groups.entry(key).or_default().push(cp);
1642    }
1643
1644    if block_groups.is_empty() {
1645        return Ok(Array2::zeros((0, p)));
1646    }
1647
1648    // Check for overlapping ranges.
1649    let ranges: Vec<(usize, usize)> = block_groups.keys().copied().collect();
1650    let mut overlapping = false;
1651    for i in 1..ranges.len() {
1652        if ranges[i].0 < ranges[i - 1].1 {
1653            overlapping = true;
1654            break;
1655        }
1656    }
1657
1658    if overlapping {
1659        if p > OVERLAPPING_PENALTY_DENSE_FALLBACK_MAX_P {
1660            return Err(EstimationError::LayoutError(format!(
1661                "overlapping penalty root would require dense {}x{} eigendecomposition; \
1662                 large-model dense fallback is disabled. Keep penalties structured or \
1663                 extend the overlapping-penalty solver path",
1664                p, p
1665            )));
1666        }
1667        // Fallback: accumulate into p × p and eigendecompose globally, under
1668        // the ONE balanced rule the reparameterization split and the
1669        // criterion's rank hint share (gam#2454).
1670        let s_balanced = balanced_penalty_sum(
1671            penalties
1672                .iter()
1673                .filter(|cp| cp.rank() != 0)
1674                .map(|cp| (cp.local_ref().view(), cp.col_range.clone())),
1675            p,
1676        );
1677        let (eigenvalues, eigenvectors) =
1678            robust_eigh(&s_balanced, Side::Lower, "balanced penalty matrix")?;
1679        let max_eig = eigenvalues.iter().fold(0.0f64, |max, &val| max.max(val));
1680        let tolerance = balanced_penalty_rank_tolerance(max_eig);
1681        let penalty_rank = eigenvalues.iter().filter(|&&ev| ev > tolerance).count();
1682        if penalty_rank == 0 {
1683            return Ok(Array2::zeros((0, p)));
1684        }
1685        let mut eb = Array2::zeros((p, penalty_rank));
1686        let mut col_idx = 0;
1687        for (i, &eigenval) in eigenvalues.iter().enumerate() {
1688            if eigenval > tolerance {
1689                let sqrt_ev = eigenval.sqrt();
1690                let evec = eigenvectors.column(i);
1691                eb.column_mut(col_idx).assign(&(&evec * sqrt_ev));
1692                col_idx += 1;
1693            }
1694        }
1695        return Ok(eb.t().to_owned());
1696    }
1697
1698    // Non-overlapping: eigendecompose per block at O(Σ p_k³).
1699    struct BlockRoot {
1700        col_range: Range<usize>,
1701        root: Array2<f64>, // rank_b × block_dim
1702    }
1703    // Materialize the BTreeMap order first. Rayon preserves Vec collection
1704    // order for indexed parallel iterators, so assembly below remains stable by
1705    // ascending column range while independent block eigendecompositions run in
1706    // parallel.
1707    let ordered_blocks: Vec<((usize, usize), Vec<&CanonicalPenalty>)> =
1708        block_groups.into_iter().collect();
1709    let block_roots: Vec<BlockRoot> = ordered_blocks
1710        .into_par_iter()
1711        .map(
1712            |((start, end), cps)| -> Result<Option<BlockRoot>, EstimationError> {
1713                let block_dim = end - start;
1714                // The block's balanced sum under the same shared rule as the
1715                // dense fallback above and the reparameterization split.
1716                let s_balanced_local = balanced_penalty_sum(
1717                    cps.iter()
1718                        .map(|cp| (cp.local_ref().view(), 0..block_dim)),
1719                    block_dim,
1720                );
1721
1722                let (eigenvalues, eigenvectors) =
1723                    robust_eigh(&s_balanced_local, Side::Lower, "balanced penalty block")?;
1724                let max_eig = eigenvalues.iter().fold(0.0f64, |max, &val| max.max(val));
1725                let tolerance = balanced_penalty_rank_tolerance(max_eig);
1726                let block_rank = eigenvalues.iter().filter(|&&ev| ev > tolerance).count();
1727
1728                if block_rank == 0 {
1729                    return Ok(None);
1730                }
1731
1732                let mut root = Array2::zeros((block_rank, block_dim));
1733                let mut row_idx = 0;
1734                for (i, &eigenval) in eigenvalues.iter().enumerate() {
1735                    if eigenval > tolerance {
1736                        let sqrt_ev = eigenval.sqrt();
1737                        let evec = eigenvectors.column(i);
1738                        root.row_mut(row_idx).assign(&(&evec * sqrt_ev));
1739                        row_idx += 1;
1740                    }
1741                }
1742
1743                Ok(Some(BlockRoot {
1744                    col_range: start..end,
1745                    root,
1746                }))
1747            },
1748        )
1749        .collect::<Result<Vec<_>, _>>()?
1750        .into_iter()
1751        .flatten()
1752        .collect();
1753    let total_rank: usize = block_roots.iter().map(|br| br.root.nrows()).sum();
1754
1755    if total_rank == 0 {
1756        return Ok(Array2::zeros((0, p)));
1757    }
1758
1759    // Assemble global balanced root: total_rank × p
1760    let mut eb = Array2::zeros((total_rank, p));
1761    let mut row_offset = 0;
1762    for br in &block_roots {
1763        let rank_b = br.root.nrows();
1764        eb.slice_mut(s![
1765            row_offset..(row_offset + rank_b),
1766            br.col_range.start..br.col_range.end
1767        ])
1768        .assign(&br.root);
1769        row_offset += rank_b;
1770    }
1771
1772    Ok(eb)
1773}
1774
1775/// Lambda-independent reparameterization invariants derived from penalty structure.
1776/// Relative cut of the balanced penalty spectrum below which a direction is
1777/// structurally unpenalized: `eigenvalue ≤ BALANCED_PENALTY_RANK_RELATIVE_TOL ·
1778/// max|eigenvalue|` of [`balanced_penalty_sum`].
1779pub const BALANCED_PENALTY_RANK_RELATIVE_TOL: f64 = 1.0e-12;
1780
1781/// The λ-invariant penalty operator the structural rank is read from:
1782/// `Σ_k S_k / ‖S_k‖_F`, each component embedded on its own column range.
1783///
1784/// # One rule, two readers
1785///
1786/// The reparameterization splits the coefficient space into the penalized
1787/// subspace `H + S̃` carries and the λ-invariant null space from THIS operator's
1788/// spectrum. The outer criterion's `−½ log|S(λ)|₊` must range over exactly that
1789/// penalized subspace — its asymptotic slope in `ρ_k` is `½(rank(S̃) −
1790/// rank|S|₊)` per unit `ρ_k`, so any disagreement in rank is a criterion with no
1791/// interior optimum along that coordinate. The criterion used to take its
1792/// structural rank from the UNWEIGHTED sum `Σ_k S_k` at a `100·p·ε·max` cut,
1793/// which is a different function of the same penalties: a component whose
1794/// Frobenius norm is small against its neighbours (a double-penalty null-space
1795/// term beside a Matérn range penalty) can sit above one cut and below the
1796/// other. Measured on the iso-κ Matérn double-penalty ladder (gam#2454):
1797/// `logdet_rank = 10` against `penalized_rank = 9`, and the gate
1798/// `penalty_logdet_ranks_the_same_subspace_the_hessian_carries_2454` red on
1799/// main. Normalizing every component first makes the rank a property of the
1800/// penalties' supports, not of their scales, which is what "structural" means.
1801pub fn balanced_penalty_sum<'a, I>(components: I, p_total: usize) -> Array2<f64>
1802where
1803    I: IntoIterator<Item = (ArrayView2<'a, f64>, std::ops::Range<usize>)>,
1804{
1805    let mut balanced = Array2::<f64>::zeros((p_total, p_total));
1806    for (local, range) in components {
1807        let frob_norm = local.iter().map(|&x| x * x).sum::<f64>().sqrt();
1808        if !(frob_norm > 1e-12) {
1809            continue;
1810        }
1811        let scale = 1.0 / frob_norm;
1812        for i in 0..local.nrows() {
1813            for j in 0..local.ncols() {
1814                balanced[[range.start + i, range.start + j]] += scale * local[[i, j]];
1815            }
1816        }
1817    }
1818    balanced
1819}
1820
1821/// The rank cut for a balanced penalty spectrum whose largest magnitude is
1822/// `max_balanced_eigenvalue`.
1823pub fn balanced_penalty_rank_tolerance(max_balanced_eigenvalue: f64) -> f64 {
1824    if max_balanced_eigenvalue > 0.0 {
1825        max_balanced_eigenvalue * BALANCED_PENALTY_RANK_RELATIVE_TOL
1826    } else {
1827        BALANCED_PENALTY_RANK_RELATIVE_TOL
1828    }
1829}
1830
1831/// Structural rank of a set of penalty components: the number of eigenvalues of
1832/// [`balanced_penalty_sum`] above [`balanced_penalty_rank_tolerance`]. This is
1833/// the rank the reparameterization's penalized subspace has, and the rank the
1834/// criterion's `log|S(λ)|₊` must range over.
1835pub fn balanced_penalty_structural_rank<'a, I>(
1836    components: I,
1837    p_total: usize,
1838) -> Result<usize, EstimationError>
1839where
1840    I: IntoIterator<Item = (ArrayView2<'a, f64>, std::ops::Range<usize>)>,
1841{
1842    if p_total == 0 {
1843        return Ok(0);
1844    }
1845    let balanced = array_to_faer(&balanced_penalty_sum(components, p_total));
1846    let (eigenvalues, _) = robust_eigh_faer(&balanced, Side::Lower, "balanced penalty matrix")?;
1847    let max_bal = eigenvalues.iter().fold(0.0_f64, |acc, &value| acc.max(value.abs()));
1848    let tol = balanced_penalty_rank_tolerance(max_bal);
1849    Ok(eigenvalues.iter().filter(|&&value| value > tol).count())
1850}
1851
1852#[derive(Clone)]
1853struct SubspaceSplit {
1854    q_pen: Array2<f64>,
1855    q_null: Array2<f64>,
1856}
1857
1858impl SubspaceSplit {
1859    fn identity(p: usize) -> Self {
1860        Self {
1861            q_pen: Array2::zeros((p, 0)),
1862            q_null: Array2::eye(p),
1863        }
1864    }
1865
1866    fn from_ordered_qs(
1867        qs: &Mat<f64>,
1868        penalized_rank: usize,
1869        p: usize,
1870    ) -> Result<Self, EstimationError> {
1871        if qs.nrows() != p || qs.ncols() != p {
1872            return Err(EstimationError::LayoutError(format!(
1873                "Invalid Q basis dimensions: expected {p}x{p}, got {}x{}",
1874                qs.nrows(),
1875                qs.ncols()
1876            )));
1877        }
1878        if penalized_rank > p {
1879            return Err(EstimationError::LayoutError(format!(
1880                "Invalid penalized rank {penalized_rank} for p={p}"
1881            )));
1882        }
1883
1884        let null_count = p - penalized_rank;
1885        let mut q_pen = Array2::<f64>::zeros((p, penalized_rank));
1886        let mut q_null = Array2::<f64>::zeros((p, null_count));
1887        for i in 0..p {
1888            for j in 0..penalized_rank {
1889                q_pen[(i, j)] = qs[(i, j)];
1890            }
1891            for j in 0..null_count {
1892                q_null[(i, j)] = qs[(i, penalized_rank + j)];
1893            }
1894        }
1895
1896        Ok(Self { q_pen, q_null })
1897    }
1898
1899    fn rank(&self) -> usize {
1900        self.q_pen.ncols()
1901    }
1902
1903    fn p(&self) -> usize {
1904        self.q_pen.nrows()
1905    }
1906
1907    fn compose_qs(&self) -> Array2<f64> {
1908        let p = self.p();
1909        let rank = self.rank();
1910        let null_count = self.q_null.ncols();
1911        let mut qs = Array2::<f64>::zeros((p, p));
1912        for i in 0..p {
1913            for j in 0..rank {
1914                qs[(i, j)] = self.q_pen[(i, j)];
1915            }
1916            for j in 0..null_count {
1917                qs[(i, rank + j)] = self.q_null[(i, j)];
1918            }
1919        }
1920        qs
1921    }
1922}
1923
1924/// Lambda-independent reparameterization invariants derived from penalty structure.
1925#[derive(Clone)]
1926pub struct ReparamInvariant {
1927    split: SubspaceSplit,
1928    /// The balanced eigenvector matrix Q (p x p). Block-local roots are
1929    /// transformed on-the-fly as `R_block @ Q[start..end, :]` instead of
1930    /// storing pre-multiplied full-width roots.
1931    qs_base: Array2<f64>,
1932    has_nonzero: bool,
1933    /// Largest eigenvalue of the balanced (unit-Frobenius) penalty matrix.
1934    /// Used as the scale reference for the shrinkage floor.
1935    max_balanced_eigenvalue: f64,
1936}
1937
1938/// Precompute the lambda-invariant reparameterization structure from canonical penalties.
1939///
1940/// Uses block-local roots directly instead of requiring rank x p global roots.
1941/// Each `CanonicalPenalty` carries its own block-local root and column range,
1942/// so the balanced sum can be assembled without ever materializing full-size
1943/// penalty matrices.
1944pub fn precompute_reparam_invariant_from_canonical(
1945    penalties: &[CanonicalPenalty],
1946    p_total: usize,
1947) -> Result<ReparamInvariant, EstimationError> {
1948    use std::cmp::Ordering;
1949
1950    let m = penalties.len();
1951
1952    if m == 0 {
1953        return Ok(ReparamInvariant {
1954            split: SubspaceSplit::identity(p_total),
1955            qs_base: Array2::eye(p_total),
1956            has_nonzero: false,
1957            max_balanced_eigenvalue: 0.0,
1958        });
1959    }
1960
1961    // Group penalties by col_range to detect block-diagonal structure.
1962    struct PenRef {
1963        penalty_index: usize,
1964    }
1965    let mut block_groups: BTreeMap<(usize, usize), Vec<PenRef>> = BTreeMap::new();
1966    let mut has_nonzero = false;
1967    for (i, cp) in penalties.iter().enumerate() {
1968        if cp.rank() == 0 {
1969            continue;
1970        }
1971        let local = cp.local_ref();
1972        let frob_norm = local.iter().map(|&x| x * x).sum::<f64>().sqrt();
1973        if frob_norm > 1e-12 {
1974            has_nonzero = true;
1975        }
1976        let key = (cp.col_range.start, cp.col_range.end);
1977        block_groups
1978            .entry(key)
1979            .or_default()
1980            .push(PenRef { penalty_index: i });
1981    }
1982
1983    if !has_nonzero {
1984        return Ok(ReparamInvariant {
1985            split: SubspaceSplit::identity(p_total),
1986            qs_base: Array2::eye(p_total),
1987            has_nonzero: false,
1988            max_balanced_eigenvalue: 0.0,
1989        });
1990    }
1991
1992    // Check for overlapping ranges.
1993    let ranges: Vec<(usize, usize)> = block_groups.keys().copied().collect();
1994    let mut overlapping = false;
1995    for i in 1..ranges.len() {
1996        if ranges[i].0 < ranges[i - 1].1 {
1997            overlapping = true;
1998            break;
1999        }
2000    }
2001
2002    if overlapping {
2003        // Mirror the dense-fallback guard from
2004        // `create_balanced_penalty_root_from_canonical`. Without this, large-scale-
2005        // scale models with overlapping penalties allocated a full
2006        // p_total × p_total workspace and ran an O(p³) eigendecomposition
2007        // before any solver code saw the problem size.
2008        if p_total > OVERLAPPING_PENALTY_DENSE_FALLBACK_MAX_P {
2009            return Err(EstimationError::LayoutError(format!(
2010                "overlapping penalty reparameterization would require dense {}x{} eigendecomposition; \
2011                 large-model dense fallback is disabled. Keep penalties structured or \
2012                 extend the overlapping-penalty solver path",
2013                p_total, p_total
2014            )));
2015        }
2016        // Fallback: global p×p eigendecomposition.
2017        let balanced = balanced_penalty_sum(
2018            penalties
2019                .iter()
2020                .filter(|cp| cp.rank() > 0)
2021                .map(|cp| (cp.local_ref().view(), cp.col_range.clone())),
2022            p_total,
2023        );
2024        let s_balanced = array_to_faer(&balanced);
2025        let (bal_eigenvalues, bal_eigenvectors) =
2026            robust_eigh_faer(&s_balanced, Side::Lower, "balanced penalty matrix")?;
2027
2028        let mut order: Vec<usize> = (0..p_total).collect();
2029        order.sort_by(|&i, &j| {
2030            bal_eigenvalues[j]
2031                .partial_cmp(&bal_eigenvalues[i])
2032                .unwrap_or(Ordering::Equal)
2033                .then(i.cmp(&j))
2034        });
2035
2036        let mut qs = Mat::<f64>::zeros(p_total, p_total);
2037        for (col_idx, &idx) in order.iter().enumerate() {
2038            for row in 0..p_total {
2039                qs[(row, col_idx)] = bal_eigenvectors[(row, idx)];
2040            }
2041        }
2042
2043        let max_bal = order
2044            .iter()
2045            .map(|&idx| bal_eigenvalues[idx].abs())
2046            .fold(0.0_f64, f64::max);
2047        let rank_tol = balanced_penalty_rank_tolerance(max_bal);
2048        let penalized_rank = order
2049            .iter()
2050            .take_while(|&&idx| bal_eigenvalues[idx] > rank_tol)
2051            .count();
2052        let split = SubspaceSplit::from_ordered_qs(&qs, penalized_rank, p_total)?;
2053
2054        return Ok(ReparamInvariant {
2055            split,
2056            qs_base: mat_to_array(&qs),
2057            has_nonzero,
2058            max_balanced_eigenvalue: max_bal,
2059        });
2060    }
2061
2062    // -----------------------------------------------------------------------
2063    // Non-overlapping: block-diagonal eigendecomposition at O(Σ p_k³).
2064    // -----------------------------------------------------------------------
2065    // The balanced sum is block-diagonal ⟹ its eigenvectors are block-local.
2066    // Q_pen and Q_null are assembled by embedding block-local eigenvectors.
2067
2068    // Track which columns are covered by any penalty.
2069    let mut covered = vec![false; p_total];
2070    for cp in penalties {
2071        for j in cp.col_range.clone() {
2072            covered[j] = true;
2073        }
2074    }
2075    let uncovered_cols: Vec<usize> = (0..p_total).filter(|j| !covered[*j]).collect();
2076
2077    struct BlockResult {
2078        col_range: Range<usize>,
2079        q_pen_local: Array2<f64>,  // block_dim × pen_rank
2080        q_null_local: Array2<f64>, // block_dim × null_rank
2081        /// Largest balanced eigenvalue contributed by this block.
2082        max_balanced_eigenvalue: f64,
2083        /// Column offset of this block's penalized directions within global Q_pen.
2084        pen_col_offset: usize,
2085        /// Column offset of this block's null directions within global Q_null.
2086        null_col_offset: usize,
2087    }
2088
2089    // BTreeMap iteration defines the deterministic block order; collecting the
2090    // indexed parallel iterator preserves that order while eigendecomposing each
2091    // independent canonical penalty block concurrently.
2092    let block_specs: Vec<_> = block_groups.iter().collect();
2093    let mut block_results: Vec<BlockResult> = block_specs
2094        .into_par_iter()
2095        .map(
2096            |(&(start, end), refs)| -> Result<BlockResult, EstimationError> {
2097                let block_dim = end - start;
2098
2099                // Build the local balanced sum under the ONE shared rule
2100                // (gam#2454); a component the rule keeps has unit Frobenius
2101                // norm, so the block is non-zero iff any entry is.
2102                let s_balanced_local = balanced_penalty_sum(
2103                    refs.iter().map(|pref| {
2104                        (
2105                            penalties[pref.penalty_index].local_ref().view(),
2106                            0..block_dim,
2107                        )
2108                    }),
2109                    block_dim,
2110                );
2111                let block_has_nonzero = s_balanced_local.iter().any(|&value| value != 0.0);
2112
2113                if !block_has_nonzero {
2114                    return Ok(BlockResult {
2115                        col_range: start..end,
2116                        q_pen_local: Array2::zeros((block_dim, 0)),
2117                        q_null_local: Array2::eye(block_dim),
2118                        max_balanced_eigenvalue: 0.0,
2119                        pen_col_offset: 0,  // set later
2120                        null_col_offset: 0, // set later
2121                    });
2122                }
2123
2124                // Eigendecompose the local balanced penalty.
2125                let (bal_eigenvalues, bal_eigenvectors) =
2126                    robust_eigh(&s_balanced_local, Side::Lower, "balanced penalty block")?;
2127
2128                let mut order: Vec<usize> = (0..block_dim).collect();
2129                order.sort_by(|&i, &j| {
2130                    bal_eigenvalues[j]
2131                        .partial_cmp(&bal_eigenvalues[i])
2132                        .unwrap_or(Ordering::Equal)
2133                        .then(i.cmp(&j))
2134                });
2135
2136                let max_bal = order
2137                    .iter()
2138                    .map(|&idx| bal_eigenvalues[idx].abs())
2139                    .fold(0.0_f64, f64::max);
2140                let rank_tol = balanced_penalty_rank_tolerance(max_bal);
2141                let penalized_rank = order
2142                    .iter()
2143                    .take_while(|&&idx| bal_eigenvalues[idx] > rank_tol)
2144                    .count();
2145                let null_count = block_dim - penalized_rank;
2146
2147                let mut q_pen_local = Array2::zeros((block_dim, penalized_rank));
2148                let mut q_null_local = Array2::zeros((block_dim, null_count));
2149                for (col_idx, &idx) in order.iter().enumerate() {
2150                    if col_idx < penalized_rank {
2151                        for row in 0..block_dim {
2152                            q_pen_local[[row, col_idx]] = bal_eigenvectors[[row, idx]];
2153                        }
2154                    } else {
2155                        let null_col = col_idx - penalized_rank;
2156                        for row in 0..block_dim {
2157                            q_null_local[[row, null_col]] = bal_eigenvectors[[row, idx]];
2158                        }
2159                    }
2160                }
2161
2162                Ok(BlockResult {
2163                    col_range: start..end,
2164                    q_pen_local,
2165                    q_null_local,
2166                    max_balanced_eigenvalue: max_bal,
2167                    pen_col_offset: 0,  // set later
2168                    null_col_offset: 0, // set later
2169                })
2170            },
2171        )
2172        .collect::<Result<_, _>>()?;
2173    let global_max_bal = block_results
2174        .iter()
2175        .map(|br| br.max_balanced_eigenvalue)
2176        .fold(0.0_f64, f64::max);
2177
2178    // Compute column offsets for each block in the global Q_pen / Q_null layout.
2179    let total_pen_rank: usize = block_results.iter().map(|br| br.q_pen_local.ncols()).sum();
2180    let total_null: usize = block_results
2181        .iter()
2182        .map(|br| br.q_null_local.ncols())
2183        .sum::<usize>()
2184        + uncovered_cols.len();
2185    {
2186        let mut pen_off = 0usize;
2187        let mut null_off = 0usize;
2188        for br in &mut block_results {
2189            br.pen_col_offset = pen_off;
2190            br.null_col_offset = null_off;
2191            pen_off += br.q_pen_local.ncols();
2192            null_off += br.q_null_local.ncols();
2193        }
2194    }
2195
2196    let mut q_pen = Array2::zeros((p_total, total_pen_rank));
2197    let mut q_null = Array2::zeros((p_total, total_null));
2198
2199    for br in &block_results {
2200        let start = br.col_range.start;
2201        let bd = br.q_pen_local.nrows();
2202        let pen_r = br.q_pen_local.ncols();
2203        let null_r = br.q_null_local.ncols();
2204        if pen_r > 0 {
2205            q_pen
2206                .slice_mut(s![
2207                    start..(start + bd),
2208                    br.pen_col_offset..(br.pen_col_offset + pen_r)
2209                ])
2210                .assign(&br.q_pen_local);
2211        }
2212        if null_r > 0 {
2213            q_null
2214                .slice_mut(s![
2215                    start..(start + bd),
2216                    br.null_col_offset..(br.null_col_offset + null_r)
2217                ])
2218                .assign(&br.q_null_local);
2219        }
2220    }
2221    let mut null_col = block_results
2222        .iter()
2223        .map(|br| br.q_null_local.ncols())
2224        .sum::<usize>();
2225    for &j in &uncovered_cols {
2226        q_null[[j, null_col]] = 1.0;
2227        null_col += 1;
2228    }
2229
2230    let split = SubspaceSplit { q_pen, q_null };
2231
2232    // Store the global Q_s = [Q_pen | Q_null] from the split.
2233    // Block-local roots are transformed on-the-fly as R_block @ Q[start..end, :]
2234    // inside the reparam engine, avoiding O(k * rank * p) storage.
2235    let qs_global = split.compose_qs();
2236
2237    Ok(ReparamInvariant {
2238        split,
2239        qs_base: qs_global,
2240        has_nonzero,
2241        max_balanced_eigenvalue: global_max_bal,
2242    })
2243}
2244
2245/// Apply stable reparameterization using precomputed lambda-invariant structures.
2246///
2247/// Write the penalized-block spectrum and rotation from a right-singular
2248/// factorization of the stacked scaled roots `E = [√λₖ Rₖ]ₖ`.
2249///
2250/// Both admissible routes to that factorization — the SVD of `E` and the SVD of
2251/// its Householder QR factor `R` — produce the identical pair, because
2252/// `EᵀE = RᵀR`. Absorbing them through one function is what makes "the same two
2253/// outputs by a different computation" a fact of the code rather than a claim
2254/// about it.
2255fn absorb_right_singular_factorization(
2256    singular_values: &Array1<f64>,
2257    vt: &Array2<f64>,
2258    penalized_rank: usize,
2259    range_eigenvalues_sorted: &mut Vec<f64>,
2260    range_rotation: &mut Mat<f64>,
2261) {
2262    // `singular_values` descending → eigenvalues `d_i = σ_i²` descending;
2263    // `vt` is `penalized_rank × penalized_rank` with row i = vᵢᵀ.
2264    let n_sv = singular_values.len().min(penalized_rank).min(vt.nrows());
2265    *range_eigenvalues_sorted = (0..penalized_rank)
2266        .map(|i| {
2267            if i < n_sv {
2268                let s = singular_values[i];
2269                s * s
2270            } else {
2271                0.0
2272            }
2273        })
2274        .collect();
2275    for col_idx in 0..n_sv {
2276        for row in 0..penalized_rank {
2277            range_rotation[(row, col_idx)] = vt[[col_idx, row]];
2278        }
2279    }
2280}
2281
2282/// The facts a refusing stacked-root SVD was decided against (#2465): the shape
2283/// it was handed, whether that input was even finite, its magnitude, and the λ
2284/// dynamic range that set it. A refusal naming only "no convergence" cannot
2285/// distinguish a genuinely unusable pencil from an iteration that gave up on a
2286/// well-formed one, and those two want opposite responses.
2287fn describe_stacked_roots(e_stacked: &Array2<f64>, lambdas: &[f64]) -> String {
2288    let mut max_abs = 0.0_f64;
2289    let mut nonfinite = 0usize;
2290    for &value in e_stacked.iter() {
2291        if value.is_finite() {
2292            max_abs = max_abs.max(value.abs());
2293        } else {
2294            nonfinite += 1;
2295        }
2296    }
2297    let finite_lambdas: Vec<f64> = lambdas.iter().copied().filter(|l| l.is_finite()).collect();
2298    let lambda_min = finite_lambdas.iter().copied().fold(f64::INFINITY, f64::min);
2299    let lambda_max = finite_lambdas
2300        .iter()
2301        .copied()
2302        .fold(f64::NEG_INFINITY, f64::max);
2303    format!(
2304        "rows={} cols={} nonfinite_entries={} max_abs={:.6e} n_lambda={} lambda_min={:.6e} \
2305         lambda_max={:.6e}",
2306        e_stacked.nrows(),
2307        e_stacked.ncols(),
2308        nonfinite,
2309        max_abs,
2310        lambdas.len(),
2311        lambda_min,
2312        lambda_max,
2313    )
2314}
2315
2316pub fn stable_reparameterizationwith_invariant(
2317    penalties: &[CanonicalPenalty],
2318    lambdas: &[f64],
2319    p: usize,
2320    invariant: &ReparamInvariant,
2321) -> Result<ReparamResult, EstimationError> {
2322    let m = penalties.len();
2323
2324    if lambdas.len() != m {
2325        return Err(EstimationError::ParameterConstraintViolation(format!(
2326            "Lambda count mismatch: expected {} lambdas for {} penalties, got {}",
2327            m,
2328            m,
2329            lambdas.len()
2330        )));
2331    }
2332
2333    // No separate length check needed — penalties are matched against lambdas above,
2334    // and the invariant's qs_base is p x p (dimension-checked by the split).
2335
2336    // #1074: the gam#1379 finite-ceiling on λ_k = exp(ρ_k) (clamp to 1e300 to
2337    // avoid `∞·0 = NaN` when the outer optimizer drives a redundant penalty
2338    // direction's log-λ past ~709) was DELETED. It masked the real defect: the
2339    // optimizer drives a redundant/unidentified penalty direction off to ∞
2340    // instead of that direction being detected and dropped from the model.
2341    // The root fix (detect+drop the redundant penalty direction at construction)
2342    // is tracked separately; λ now passes through raw.
2343
2344    if m == 0 {
2345        return Ok(ReparamResult {
2346            s_transformed: Array2::zeros((p, p)),
2347            log_det: 0.0,
2348            det1: Array1::zeros(0),
2349            qs: Array2::eye(p),
2350            canonical_transformed: vec![],
2351            e_transformed: Array2::zeros((0, p)),
2352            // All modes truncated when no penalties; already in transformed frame.
2353            u_truncated: Array2::eye(p),
2354        });
2355    }
2356
2357    if !invariant.has_nonzero {
2358        let qs = invariant.split.compose_qs();
2359        let u_truncated = qs.t().dot(&invariant.split.q_null);
2360        // All penalties are zero — canonical_transformed = originals (no rotation needed).
2361        let canonical_transformed: Vec<CanonicalPenalty> = penalties.to_vec();
2362        return Ok(ReparamResult {
2363            s_transformed: Array2::zeros((p, p)),
2364            log_det: 0.0,
2365            det1: Array1::zeros(m),
2366            qs,
2367            canonical_transformed,
2368            e_transformed: Array2::zeros((0, p)),
2369            u_truncated,
2370        });
2371    }
2372
2373    let q_pen = array_to_faer(&invariant.split.q_pen);
2374    let q_null = array_to_faer(&invariant.split.q_null);
2375    let qs_base = array_to_faer(&invariant.qs_base);
2376    // Each penalty root transform is independent: R_k_block @ Q[start..end, :].
2377    // Run those per-penalty products (and their S_k = R_k'R_k caches) in
2378    // parallel, then collect in slice order so all downstream accumulation stays
2379    // deterministic and bit-for-bit stable with respect to penalty ordering.
2380    let penalty_transforms: Vec<(Mat<f64>, Mat<f64>)> = penalties
2381        .par_iter()
2382        .map(|cp| {
2383            let r = &cp.col_range;
2384            let root_faer = array_to_faer(&cp.root);
2385            let q_block = qs_base.submatrix(r.start, 0, cp.block_dim(), p);
2386            let mut product = Mat::<f64>::zeros(cp.rank(), p);
2387            matmul(
2388                product.as_mut(),
2389                Accum::Replace,
2390                root_faer.as_ref(),
2391                q_block,
2392                1.0,
2393                Par::Seq,
2394            );
2395            let s_k = penalty_from_root_faer(&product);
2396            (product, s_k)
2397        })
2398        .collect();
2399    let (rs_transformed, s_k_penalized_cache): (Vec<Mat<f64>>, Vec<Mat<f64>>) =
2400        penalty_transforms.into_iter().unzip();
2401
2402    let penalized_rank = invariant.split.rank();
2403
2404    let mut range_eigenvalues_sorted: Vec<f64> = Vec::new();
2405    let mut range_rotation = Mat::<f64>::zeros(penalized_rank, penalized_rank);
2406    if penalized_rank > 0 {
2407        // Penalized-block spectrum `Σ_k λ_k S_k = U diag(d) Uᵀ` (restricted to the
2408        // λ-invariant penalized subspace).  Compute it from the SVD of the STACKED
2409        // SCALED ROOTS `E = [√λ_k R_k]_k` (root rows stacked), NOT from an
2410        // eigendecomposition of the assembled Gram `range_block = EᵀE`.
2411        //
2412        // Why (#2123): `S_k = R_kᵀ R_k`, so `Σ_k λ_k S_k = EᵀE` with
2413        // `E = vstack_k(√λ_k R_k)`.  Assembling the Gram and eigendecomposing it
2414        // SQUARES the condition number (`κ(EᵀE) = κ(E)²`).  When the outer optimizer
2415        // drives one margin toward its null space the λ dynamic range is enormous
2416        // (here `te(x,z)` with the near-linear z axis reaches `λ_ratio ≳ 1e8`), so a
2417        // recessive-penalty eigenvalue `d_min ≈ λ_min·σ²` is swamped by the
2418        // eigensolver's `O(ε·d_max) = O(ε·λ_max)` absolute floor — its eigenVECTOR
2419        // rotates into numerical noise, the genuinely-penalized direction is lost
2420        // from `e_transformed`, and the inner P-IRLS solve then fits that direction
2421        // to the data (a WIGGLY β̂ with `βᵀSβ̂ ≈ 0` despite `λ → ∞`).  That silent
2422        // loss-of-penalty is a discontinuous function of ρ (it flips as the noise
2423        // floor crosses `d_min`), so it injects spurious cliffs into the REML/LAML
2424        // objective and a false low-cost basin in the high-λ corner — which the
2425        // outer optimizer then lands in for some training-row orders but not others
2426        // (the row-order-dependent EDF/SE of #2123).
2427        //
2428        // The SVD operates on `E` directly, so `σ_min(E) = √d_min` is resolved
2429        // whenever `√λ_min·σ ≳ ε·√λ_max·σ` — a λ dynamic range up to ~1e32 instead
2430        // of ~1e16 — keeping the reparameterized penalty faithful across the whole ρ
2431        // box and the outer objective smooth and permutation-invariant.  `EᵀE`
2432        // equals the previous `range_block` bit-for-bit, so well-conditioned fits
2433        // are unchanged; only the ill-conditioned corner is corrected.
2434        //
2435        // As before, the right singular vectors `V` (= the penalized-block rotation)
2436        // are used ONLY to build `E`/`S⁺`/traces below; they are NOT applied to
2437        // `q_pen` or `rs_transformed`, so `Q_s` stays λ-independent and the
2438        // quasi-Newton coordinate system does not drift at eigenvalue crossings.
2439        let total_root_rows: usize = rs_transformed.iter().map(Mat::nrows).sum();
2440        // Thin SVD yields a COMPLETE orthonormal `V` (all `penalized_rank`
2441        // directions, including exactly-zero σ) only when `E` is tall or square
2442        // (`total_root_rows ≥ penalized_rank`).  That always holds structurally —
2443        // the union of the penalty root ranges spans the penalized subspace — but a
2444        // pathological degenerate layout is handled by falling back to the Gram
2445        // eigendecomposition so `range_rotation` is never left rank-deficient.
2446        //
2447        // ROUTE LADDER (#2581).  The shape test gates only the ATTEMPT, not the
2448        // choice.  Both routes below compute the SAME two outputs
2449        // (`range_eigenvalues_sorted`, `range_rotation`), so a stacked-root SVD
2450        // that fails to CONVERGE is exactly the pathological case the Gram route
2451        // was written for: non-convergence is a property of the bidiagonal
2452        // iteration, not evidence that the pencil is unusable.  Selecting on the
2453        // shape alone turned that into a fatal `LayoutError` raised one branch
2454        // away from a trusted routine for the same quantity, aborting a whole
2455        // converging fit (measured: nottem `cc(month, k=12)`, one 75% partition,
2456        // λ ≈ 5, an 11×11 `E` — the outer BFGS had already certified a nearby ρ
2457        // before the refinement pass hit it).
2458        //
2459        // The ladder has three rungs, in accuracy order: the direct SVD of `E`;
2460        // the R-SVD of its Householder QR factor, which is EXACTLY as accurate
2461        // and merely a different computation; and only then the Gram route,
2462        // whose real cost is the squared condition number.  Reaching either
2463        // lower rung is REPORTED, so the route that produced the answer is never
2464        // silently substituted.
2465        let mut have_rotation = false;
2466        let mut rescued_by_r_svd = false;
2467        let mut svd_refusal: Option<String> = None;
2468        if total_root_rows >= penalized_rank {
2469            let mut e_stacked = Array2::<f64>::zeros((total_root_rows, penalized_rank));
2470            let mut row_off = 0usize;
2471            for (lambda, root) in lambdas.iter().zip(rs_transformed.iter()) {
2472                let sqrt_lambda = lambda.max(0.0).sqrt();
2473                let rk = root.nrows();
2474                for r in 0..rk {
2475                    for c in 0..penalized_rank {
2476                        e_stacked[[row_off + r, c]] = sqrt_lambda * root[(r, c)];
2477                    }
2478                }
2479                row_off += rk;
2480            }
2481            match e_stacked.svd(false, true) {
2482                Ok((_, singular_values, Some(vt))) => {
2483                    absorb_right_singular_factorization(
2484                        &singular_values,
2485                        &vt,
2486                        penalized_rank,
2487                        &mut range_eigenvalues_sorted,
2488                        &mut range_rotation,
2489                    );
2490                    have_rotation = true;
2491                }
2492                direct => {
2493                    let facts = describe_stacked_roots(&e_stacked, lambdas);
2494                    svd_refusal = Some(match direct {
2495                        Err(err) => format!("failed: {err:?}; {facts}"),
2496                        _ => format!("returned no right singular vectors; {facts}"),
2497                    });
2498                    // RUNG 2, and the reason the Gram route is a LAST resort
2499                    // rather than the only alternative.  `E = QR` by Householder
2500                    // reflections is direct — it has no convergence criterion to
2501                    // miss — and `EᵀE = RᵀR`, so `R`'s right singular vectors
2502                    // and singular values ARE `E`'s.  It therefore delivers the
2503                    // same two outputs at the same `O(ε²·d_max)` resolution the
2504                    // direct SVD promises, on a `penalized_rank`-square problem
2505                    // instead of a `total_root_rows`-tall one.
2506                    //
2507                    // Measured on the refusing input (nottem `cc(month, k=12)`,
2508                    // split 1, an 11×11 `E`, one λ = 6.068, no non-finite
2509                    // entries): a power-of-two rescale of `E` refuses
2510                    // identically — so the failure is not a scaling artefact —
2511                    // while this route and `svd(Eᵀ)` both converge and agree on
2512                    // `σ₀ = 1.653863e0`.
2513                    if let Ok((_, r_factor)) = e_stacked.qr()
2514                        && let Ok((_, singular_values, Some(vt))) = r_factor.svd(false, true)
2515                    {
2516                        absorb_right_singular_factorization(
2517                            &singular_values,
2518                            &vt,
2519                            penalized_rank,
2520                            &mut range_eigenvalues_sorted,
2521                            &mut range_rotation,
2522                        );
2523                        have_rotation = true;
2524                        rescued_by_r_svd = true;
2525                    }
2526                }
2527            }
2528        }
2529        if let Some(reason) = svd_refusal.as_deref() {
2530            if rescued_by_r_svd {
2531                log::warn!(
2532                    "penalized-block rotation: stacked-root SVD {reason}. Recovered the SAME \
2533                     right-singular basis from the Householder QR of `E` followed by the SVD \
2534                     of its triangular factor `R`: `EᵀE = RᵀR`, so no accuracy is given up."
2535                );
2536            } else {
2537                // The accuracy downgrade is observable rather than silent: the
2538                // Gram route resolves a recessive penalized eigenvalue only down
2539                // to `O(ε·d_max)`, where the SVD of `E` reaches `O(ε²·d_max)`.
2540                log::warn!(
2541                    "penalized-block rotation: stacked-root SVD {reason}, and so did the R-SVD \
2542                     of its Householder QR factor. Recomputing it from the Gram `Σₖ λₖ Sₖ`, \
2543                     which squares the condition number: recessive eigenvalues are resolved to \
2544                     O(ε·d_max) rather than O(ε²·d_max)."
2545                );
2546            }
2547        }
2548        if !have_rotation {
2549            // Assemble the Gram and eigendecompose it. This route serves a layout
2550            // whose penalty roots cannot span the penalized subspace
2551            // (`total_root_rows < penalized_rank`) AND a stacked-root SVD that
2552            // did not converge.
2553            let mut range_block = Mat::<f64>::zeros(penalized_rank, penalized_rank);
2554            for (lambda, s_k) in lambdas.iter().zip(s_k_penalized_cache.iter()) {
2555                for i in 0..penalized_rank {
2556                    for j in 0..penalized_rank {
2557                        range_block[(i, j)] += *lambda * s_k[(i, j)];
2558                    }
2559                }
2560            }
2561            let (range_eigenvalues, range_eigenvectors) =
2562                robust_eigh_faer(&range_block, Side::Lower, "range penalty block")?;
2563            let mut range_order: Vec<usize> = (0..penalized_rank).collect();
2564            range_order.sort_by(|&i, &j| {
2565                range_eigenvalues[j]
2566                    .partial_cmp(&range_eigenvalues[i])
2567                    .unwrap_or(std::cmp::Ordering::Equal)
2568                    .then(i.cmp(&j))
2569            });
2570            range_eigenvalues_sorted = range_order
2571                .iter()
2572                .map(|&idx| range_eigenvalues[idx])
2573                .collect();
2574            for (col_idx, &idx) in range_order.iter().enumerate() {
2575                for row in 0..penalized_rank {
2576                    range_rotation[(row, col_idx)] = range_eigenvectors[(row, idx)];
2577                }
2578            }
2579        }
2580    }
2581
2582    // Subspace-invariant penalty spectral calculus:
2583    // - Penalized and null spaces are fixed by the lambda-invariant basis `qs_base`.
2584    // - Runtime lambda dependence only appears in the penalized block eigenvalues.
2585    // This avoids basis mixing inside the degenerate zero-eigenspace.
2586    let structural_rank = penalized_rank;
2587    let range_eigs_sorted: Vec<f64> = range_eigenvalues_sorted;
2588
2589    let eigenvalue_floor = invariant.max_balanced_eigenvalue.max(1.0) * 1e-12;
2590    let qs = compose_qs_from_split(&q_pen, &q_null, p);
2591
2592    // Guard against any accidental penalized/null mixing. The transformed penalty
2593    // roots must have negligible support on null columns by construction.
2594    let leakage = assess_subspace_leakage(&qs, &rs_transformed, structural_rank, p);
2595    if !subspace_split_is_consistent(&leakage, p) {
2596        return Err(EstimationError::LayoutError(format!(
2597            "Reparameterization subspace split is inconsistent: max null leakage {:.3e} (rel {:.3e}, worst penalty {}), max |Qp'Qn| {:.3e}",
2598            leakage.max_abs_sq.sqrt(),
2599            leakage.max_rel_sq.sqrt(),
2600            leakage.worst_penalty,
2601            leakage.max_cross_gram_abs,
2602        )));
2603    }
2604
2605    // Truncated basis in transformed coordinates:
2606    //   U_⊥^(t) = Qs^T U_⊥^(orig) = Qs^T Q_n.
2607    let mut u_truncated_mat = Mat::<f64>::zeros(p, q_null.ncols());
2608    matmul(
2609        u_truncated_mat.as_mut(),
2610        Accum::Replace,
2611        qs.transpose(),
2612        q_null.as_ref(),
2613        1.0,
2614        Par::Seq,
2615    );
2616
2617    // E is represented in TRANSFORMED coordinates (beta_t).  Because the
2618    // penalized subspace is NOT rotated by the lambda-dependent eigenvectors
2619    // (to keep Q_s stable across BFGS iterations), E is no longer diagonal.
2620    // Instead E = diag(√d) · U' embedded in structural_rank × p, so that
2621    // E'E = U diag(d) U' = Σ λ_k S_k in the invariant penalized basis.
2622    let mut e_transformed_mat = Mat::<f64>::zeros(structural_rank, p);
2623    for row_idx in 0..structural_rank {
2624        let safe_eigenval = range_eigs_sorted[row_idx].max(eigenvalue_floor);
2625        let sqrt_eigenval = safe_eigenval.sqrt();
2626        // E[row, j] = sqrt(d_row) * U'[row, j] = sqrt(d_row) * U[j, row]
2627        for j in 0..penalized_rank {
2628            e_transformed_mat[(row_idx, j)] = sqrt_eigenval * range_rotation[(j, row_idx)];
2629        }
2630    }
2631
2632    // Pseudo-logdet on the structural penalized block.  The null block is split
2633    // out above, so there is no nullspace normalization here.  Spectrally-noisy
2634    // directions (eigenvalues snapped to 0 by `classify_eigenvalues_strict`
2635    // because they fall below the `c * eps_machine * p * scale` tolerance in
2636    // the lambda-weighted sum) are floored to `eigenvalue_floor` to keep the
2637    // log-det finite and consistent with the floored values used to construct
2638    // `e_transformed_mat` above.  This avoids spurious P-IRLS failures when the
2639    // lambda dynamic range is wide (e.g. during BFGS line search probing extreme
2640    // rho candidates).  Materially negative or non-finite spectra are already
2641    // rejected by the strict classifier upstream; this loop re-checks the
2642    // *post-shrinkage* range eigenvalues against the same floor.
2643    //
2644    // The same floored spectrum is used in the trace formula tr(S⁺ S_k) below,
2645    // matching the rank structure embedded in `e_transformed_mat` and avoiding
2646    // a 1/0 in the trace contraction when an eigenvalue was floored to 0.
2647    let mut floored_eigs: Vec<f64> = Vec::with_capacity(range_eigs_sorted.len());
2648    let mut log_det_sum = KahanSum::default();
2649    for (idx, &ev) in range_eigs_sorted.iter().enumerate() {
2650        if !ev.is_finite() || ev < -eigenvalue_floor {
2651            return Err(EstimationError::LayoutError(format!(
2652                "Penalty pseudo-logdet has a non-finite or large-negative structural eigenvalue at index {idx}: {ev:.3e}"
2653            )));
2654        }
2655        let safe_ev = ev.max(eigenvalue_floor);
2656        floored_eigs.push(safe_ev);
2657        if idx < penalized_rank {
2658            log_det_sum.add(safe_ev.ln());
2659        }
2660    }
2661    let log_det = log_det_sum.sum();
2662    let delta = 0.0;
2663
2664    // The det1 contractions are independent once the eigensystem is fixed.  Use
2665    // indexed parallel collection so the output vector preserves lambda order.
2666    let det1vec: Vec<f64> = (0..lambdas.len())
2667        .into_par_iter()
2668        .map(|k| {
2669            let s_k = &s_k_penalized_cache[k];
2670            // Compute tr((S+δI)⁻¹ S_k) in the range eigenbasis without ever
2671            // materializing (S+δI)⁻¹. Using faer's matmul keeps this contraction
2672            // aligned with the orthogonal-similarity debug reference path.
2673            let trace = trace_penalty_in_orthogonal_basis(
2674                s_k,
2675                penalized_rank,
2676                &range_rotation,
2677                &floored_eigs,
2678                delta,
2679            );
2680            lambdas[k] * trace
2681        })
2682        .collect();
2683
2684    {
2685        // Guardrail: cross-check the primary Rayleigh-quotient contraction
2686        // against a full orthogonal similarity transform, while staying in
2687        // the same numerically stable eigenbasis coordinates.
2688        let mut maxdet1_mismatch = 0.0_f64;
2689        let mut det1_scale = 0.0_f64;
2690        for (k, lambda) in lambdas.iter().enumerate() {
2691            let s_k_penalized = &s_k_penalized_cache[k];
2692            let s_k_eigenbasis = orthogonal_similarity_transform_faer(
2693                s_k_penalized,
2694                penalized_rank,
2695                &range_rotation,
2696            );
2697            let mut trace = KahanSum::default();
2698            for l in 0..penalized_rank {
2699                trace.add(s_k_eigenbasis[(l, l)] / (floored_eigs[l] + delta));
2700            }
2701            let reference = *lambda * trace.sum();
2702            maxdet1_mismatch = maxdet1_mismatch.max((reference - det1vec[k]).abs());
2703            det1_scale = det1_scale.max(reference.abs()).max(det1vec[k].abs());
2704        }
2705        let det1_tolerance = 1e-7 * det1_scale.max(1.0);
2706        assert!(
2707            maxdet1_mismatch <= det1_tolerance,
2708            "det1 mismatch between optimized and reference formulas: max_abs={maxdet1_mismatch:.3e}, tol={det1_tolerance:.3e}"
2709        );
2710    }
2711
2712    // Rebuild s_transformed from e_transformed to ensure rank consistency.
2713    //
2714    // The sum of λ*S_k may contain numerical noise modes (eigenvalues ~1e-15) that
2715    // become significant when λ is large (e.g., 10^12). These modes would appear in H
2716    // but are truncated from log|S|_+, creating a "phantom penalty" in the objective.
2717    //
2718    // By reconstructing s_transformed = E^T * E, we force the penalty matrix used
2719    // in H to have the EXACT same rank structure as the one used for log|S|_+.
2720    // Any mode truncated from the prior is now strictly zero in the Hessian
2721    // calculation, ensuring mathematical consistency of the gradients.
2722    let mut s_truncated = Mat::<f64>::zeros(p, p);
2723    matmul(
2724        s_truncated.as_mut(),
2725        Accum::Replace,
2726        e_transformed_mat.transpose(),
2727        e_transformed_mat.as_ref(),
2728        1.0,
2729        Par::Seq,
2730    );
2731
2732    {
2733        // Structural check: transformed S must not leak into declared null coordinates.
2734        let mut max_null_diag = 0.0_f64;
2735        let mut max_null_offdiag = 0.0_f64;
2736        for i in structural_rank..p {
2737            max_null_diag = max_null_diag.max(s_truncated[(i, i)].abs());
2738            for j in 0..p {
2739                if i != j {
2740                    max_null_offdiag = max_null_offdiag.max(s_truncated[(i, j)].abs());
2741                }
2742            }
2743        }
2744        assert!(
2745            max_null_diag <= 1e-10 && max_null_offdiag <= 1e-10,
2746            "null-space leakage in transformed penalty: max_null_diag={max_null_diag:.3e}, max_null_offdiag={max_null_offdiag:.3e}"
2747        );
2748    }
2749
2750    let qs_array = mat_to_array(&qs);
2751    let canonical_transformed: Vec<CanonicalPenalty> = rs_transformed
2752        .par_iter()
2753        .zip(penalties.par_iter())
2754        .map(|(r, cp)| {
2755            let mean_transformed = qs_array.t().dot(&cp.full_width_prior_mean());
2756            CanonicalPenalty::from_dense_root_with_mean(mat_to_array(r), p, mean_transformed)
2757        })
2758        .collect();
2759    Ok(ReparamResult {
2760        s_transformed: mat_to_array(&s_truncated),
2761        log_det,
2762        det1: Array1::from(det1vec),
2763        qs: qs_array,
2764        canonical_transformed,
2765        e_transformed: mat_to_array(&e_transformed_mat),
2766        u_truncated: mat_to_array(&u_truncated_mat),
2767    })
2768}
2769
2770/// Minimal engine layout descriptor that avoids domain-specific layout coupling.
2771#[derive(Debug, Clone, Copy, PartialEq, Eq)]
2772pub struct EngineDims {
2773    pub p: usize,
2774    pub k: usize,
2775}
2776
2777impl EngineDims {
2778    pub fn new(p: usize, k: usize) -> Self {
2779        Self { p, k }
2780    }
2781}
2782
2783/// Engine-facing stable reparameterization API using only `(p, k)`.
2784///
2785/// When `cached_invariant` is `Some`, reuses the precomputed eigendecomposition
2786/// (the hot path inside the REML loop). When `None`, computes the invariant on
2787/// the fly (the post-REML refit path).
2788/// Stable reparameterization from block-local canonical penalties.
2789pub fn stable_reparameterization_engine_canonical(
2790    penalties: &[CanonicalPenalty],
2791    lambdas: &[f64],
2792    dims: EngineDims,
2793    cached_invariant: Option<&ReparamInvariant>,
2794) -> Result<ReparamResult, EstimationError> {
2795    let owned;
2796    let invariant = match cached_invariant {
2797        Some(inv) => inv,
2798        None => {
2799            owned = precompute_reparam_invariant_from_canonical(penalties, dims.p)?;
2800            &owned
2801        }
2802    };
2803    stable_reparameterizationwith_invariant(penalties, lambdas, dims.p, invariant)
2804}
2805
2806// ---------------------------------------------------------------------------
2807// Kronecker-factored reparameterization for tensor-product smooths
2808// ---------------------------------------------------------------------------
2809
2810/// Result of Kronecker-factored reparameterization.
2811///
2812/// Exploits the fact that for Kronecker-structured penalties, the joint
2813/// eigenvector matrix is `U_1 ⊗ ... ⊗ U_d` and the reparameterized design
2814/// is a rowwise Kronecker of `(B_k U_k)` — all remaining factored.
2815#[derive(Clone)]
2816pub struct KroneckerReparamResult {
2817    /// Reparameterized marginal designs: `B_k · U_k` for each marginal k.
2818    ///
2819    /// `Arc`-shared with the λ-invariant cache so the per-outer-iterate
2820    /// memoized engine bumps a refcount instead of deep-copying the
2821    /// (n × q) reparameterized marginals every call.
2822    pub reparameterized_marginals: Arc<Vec<Array2<f64>>>,
2823    /// Marginal eigenvalues from each marginal penalty eigendecomposition.
2824    pub marginal_eigenvalues: Arc<Vec<Array1<f64>>>,
2825    /// Marginal eigenvector matrices U_k.
2826    pub marginal_qs: Arc<Vec<Array2<f64>>>,
2827    /// log|S|₊ computed from marginal eigenvalue grid.
2828    pub log_det: f64,
2829    /// First derivatives of log|S|₊ w.r.t. ρ_k = log(λ_k).
2830    pub det1: Array1<f64>,
2831    /// Second derivatives of log|S|₊ w.r.t. ρ.
2832    pub det2: Array2<f64>,
2833    /// Whether a double penalty (global ridge) is present.
2834    pub has_double_penalty: bool,
2835    /// Marginal basis dimensions.
2836    pub marginal_dims: Vec<usize>,
2837}
2838
2839impl KroneckerReparamResult {
2840    /// Materialize the joint Qs matrix (U_1 ⊗ ... ⊗ U_d) as dense p×p.
2841    /// Only for fallback paths — avoid in hot loops.
2842    pub fn materialize_qs(&self) -> Array2<f64> {
2843        let mut qs = Array2::<f64>::eye(1);
2844        for u_k in self.marginal_qs.iter() {
2845            qs = kronecker_product(&qs, u_k);
2846        }
2847        qs
2848    }
2849
2850    /// Materialize s_transformed (the penalty in the reparameterized basis).
2851    /// In the eigenbasis, this is diagonal with entries Σ_k λ_k μ_{k,j_k}.
2852    pub fn materialize_s_transformed(&self, lambdas: &[f64]) -> Array2<f64> {
2853        let d = self.marginal_dims.len();
2854        let p: usize = self.marginal_dims.iter().copied().product();
2855        let mut s = Array2::<f64>::zeros((p, p));
2856
2857        // Delegate the per-cell tensor-penalty accumulation to the shared
2858        // `kronecker_cell_sigma` (#1172/#1185 single source of truth). Fold the
2859        // `lambdas.len() > d` guard into `has_double` to preserve exact gating.
2860        let eigenvalue_views: Vec<ArrayView1<'_, f64>> =
2861            self.marginal_eigenvalues.iter().map(|m| m.view()).collect();
2862        let has_double = self.has_double_penalty && lambdas.len() > d;
2863        let mut multi_idx = vec![0usize; d];
2864        let mut flat = 0usize;
2865        loop {
2866            let (sigma, _structural_sigma, _joint_null) = kronecker_cell_sigma(
2867                &eigenvalue_views,
2868                &multi_idx,
2869                lambdas,
2870                d,
2871                has_double,
2872                0.0,
2873            );
2874            s[[flat, flat]] = sigma;
2875            flat += 1;
2876
2877            if kronecker_multi_index_advance(&mut multi_idx, &self.marginal_dims) {
2878                break;
2879            }
2880        }
2881        s
2882    }
2883
2884    /// Explicitly materialize the dense artifact bundle expected by legacy
2885    /// downstream consumers. This is not part of the native Kronecker solve path.
2886    pub fn materialize_dense_artifact_result(
2887        &self,
2888        rs_list: &[Array2<f64>],
2889        lambdas: &[f64],
2890        p: usize,
2891    ) -> Result<ReparamResult, EstimationError> {
2892        const KRONECKER_DENSE_COMPAT_FALLBACK_MAX_P: usize = 4096;
2893        if p > KRONECKER_DENSE_COMPAT_FALLBACK_MAX_P {
2894            return Err(EstimationError::LayoutError(format!(
2895                "Kronecker reparameterization would materialize dense {}x{} compatibility tensors; \
2896                 large-model dense fallback is disabled. Wire the downstream solver to consume \
2897                 the factored Kronecker result directly",
2898                p, p
2899            )));
2900        }
2901        let qs = self.materialize_qs();
2902        let s_transformed = self.materialize_s_transformed(lambdas);
2903
2904        // Transform penalty roots: R_k_transformed = R_k · Qs
2905        let rs_transformed: Vec<Array2<f64>> = if rs_list.len() >= 2 {
2906            use rayon::prelude::*;
2907            rs_list
2908                .par_iter()
2909                .map(|r| gam_linalg::faer_ndarray::fast_ab(r, &qs))
2910                .collect()
2911        } else {
2912            rs_list
2913                .iter()
2914                .map(|r| gam_linalg::faer_ndarray::fast_ab(r, &qs))
2915                .collect()
2916        };
2917        // rs_transposed removed — canonical_transformed is the single source of truth.
2918
2919        // Build e_transformed: combined penalty square root in transformed coords.
2920        // For Kronecker structure, the penalty is diagonal in the eigenbasis.
2921        // e_transformed rows are the nonzero rows of sqrt(Σ_k λ_k S_k)^{1/2}.
2922        let d = self.marginal_dims.len();
2923        // Delegate the per-cell tensor-penalty accumulation to the shared
2924        // `kronecker_cell_sigma` (the #1172/#1185 single source of truth). The
2925        // double-penalty term is only valid when `lambdas` actually carries the
2926        // λ_d entry, so fold the original `lambdas.len() > d` guard into the
2927        // `has_double_penalty` flag passed to the helper — preserving the exact
2928        // gating behavior.
2929        let eigenvalue_views: Vec<ArrayView1<'_, f64>> =
2930            self.marginal_eigenvalues.iter().map(|m| m.view()).collect();
2931        let has_double = self.has_double_penalty && lambdas.len() > d;
2932        let diag_vals: Vec<f64> = {
2933            let mut vals = Vec::with_capacity(p);
2934            let mut multi_idx = vec![0usize; d];
2935            loop {
2936                let (sigma, _structural_sigma, _joint_null) = kronecker_cell_sigma(
2937                    &eigenvalue_views,
2938                    &multi_idx,
2939                    lambdas,
2940                    d,
2941                    has_double,
2942                    0.0,
2943                );
2944                vals.push(if sigma > 0.0 { sigma.sqrt() } else { 0.0 });
2945
2946                if kronecker_multi_index_advance(&mut multi_idx, &self.marginal_dims) {
2947                    break;
2948                }
2949            }
2950            vals
2951        };
2952        let rank = diag_vals.iter().filter(|&&v| v > 1e-12).count();
2953        let mut e_transformed = Array2::<f64>::zeros((rank, p));
2954        let mut row = 0;
2955        for (j, &v) in diag_vals.iter().enumerate() {
2956            if v > 1e-12 {
2957                e_transformed[[row, j]] = v;
2958                row += 1;
2959            }
2960        }
2961
2962        // u_truncated: null-space eigenvectors (columns with zero eigenvalue).
2963        let null_count = p - rank;
2964        let mut u_truncated = Array2::<f64>::zeros((p, null_count));
2965        let mut col = 0;
2966        for (j, &v) in diag_vals.iter().enumerate() {
2967            if v <= 1e-12 {
2968                u_truncated[[j, col]] = 1.0; // standard basis vector in eigenbasis
2969                col += 1;
2970            }
2971        }
2972
2973        let canonical_transformed: Vec<CanonicalPenalty> = rs_transformed
2974            .iter()
2975            .map(|r| CanonicalPenalty::from_dense_root(r.clone(), p))
2976            .collect();
2977        Ok(ReparamResult {
2978            s_transformed,
2979            log_det: self.log_det,
2980            det1: self.det1.clone(),
2981            qs,
2982            canonical_transformed,
2983            e_transformed,
2984            u_truncated,
2985        })
2986    }
2987}
2988
2989/// Compute `log|S|₊` and its first/second derivatives w.r.t. `ρ_k = log(λ_k)`
2990/// from factored marginal eigenvalues.
2991///
2992/// Shared implementation for `KroneckerPenaltySystem::logdet_and_derivatives`
2993/// and `kronecker_reparameterization_engine`.  Iterates over the ∏q_j
2994/// multi-index grid in O(d · ∏q_j) time with no O(p²) storage.
2995const KRONECKER_STRUCTURAL_ZERO_TOL: f64 = 1e-12;
2996
2997/// Per-cell Kronecker eigenvalue accumulation — the single source of truth for
2998/// the #1172/#1185 tensor-penalty math.
2999///
3000/// For the multi-index cell `multi_idx`, accumulates:
3001///   - `sigma`            = Σ_k λ_k · μ_k  (+ joint-null double-penalty term + declared objective ridge)
3002///   - `structural_sigma` = Σ_k μ_k        (unweighted; classifies joint-null cells)
3003///   - `joint_null`       = whether the cell lies in the joint null space
3004///
3005/// `marginal_eigenvalues[k][multi_idx[k]]` is the k-th marginal eigenvalue μ_k.
3006/// The double-penalty (global ridge) term `λ_d` is added only on joint-null
3007/// cells; an explicit objective ridge is added only on structurally penalized
3008/// cells. This mirrors the gated logic fixed in #1172/#1185 and MUST be kept
3009/// identical across every caller.
3010#[inline]
3011fn kronecker_cell_sigma(
3012    marginal_eigenvalues: &[ArrayView1<'_, f64>],
3013    multi_idx: &[usize],
3014    lambdas: &[f64],
3015    d: usize,
3016    has_double_penalty: bool,
3017    objective_ridge: f64,
3018) -> (f64, f64, bool) {
3019    let mut sigma = 0.0;
3020    let mut structural_sigma = 0.0;
3021    for k in 0..d {
3022        let marginal_eigenvalue = marginal_eigenvalues[k][multi_idx[k]];
3023        structural_sigma += marginal_eigenvalue;
3024        sigma += lambdas[k] * marginal_eigenvalue;
3025    }
3026    let joint_null = structural_sigma <= KRONECKER_STRUCTURAL_ZERO_TOL;
3027    if has_double_penalty && joint_null {
3028        sigma += lambdas[d];
3029    }
3030    if structural_sigma > KRONECKER_STRUCTURAL_ZERO_TOL {
3031        sigma += objective_ridge;
3032    }
3033    (sigma, structural_sigma, joint_null)
3034}
3035
3036/// Advance a row-major multi-index over the `dims` grid in place.
3037/// Returns `true` when the grid is exhausted (the index wrapped back to all-zero).
3038#[inline]
3039pub(crate) fn kronecker_multi_index_advance(multi_idx: &mut [usize], dims: &[usize]) -> bool {
3040    let mut carry = true;
3041    for dim in (0..dims.len()).rev() {
3042        if carry {
3043            multi_idx[dim] += 1;
3044            if multi_idx[dim] < dims[dim] {
3045                carry = false;
3046            } else {
3047                multi_idx[dim] = 0;
3048            }
3049        }
3050    }
3051    carry
3052}
3053
3054pub fn kronecker_logdet_and_derivatives(
3055    marginal_eigenvalues: &[ArrayView1<'_, f64>],
3056    marginal_dims: &[usize],
3057    lambdas: &[f64],
3058    has_double_penalty: bool,
3059    objective_ridge: f64,
3060) -> (f64, Array1<f64>, Array2<f64>) {
3061    let d = marginal_dims.len();
3062    let n_pen = d + if has_double_penalty { 1 } else { 0 };
3063
3064    let mut logdet = 0.0;
3065    let mut grad = Array1::<f64>::zeros(n_pen);
3066    let mut hess = Array2::<f64>::zeros((n_pen, n_pen));
3067    let tol = 1e-12;
3068
3069    let mut multi_idx = vec![0usize; d];
3070    loop {
3071        let (sigma, _structural_sigma, joint_null) = kronecker_cell_sigma(
3072            marginal_eigenvalues,
3073            &multi_idx,
3074            lambdas,
3075            d,
3076            has_double_penalty,
3077            objective_ridge,
3078        );
3079
3080        if sigma > tol {
3081            logdet += sigma.ln();
3082            let inv_sigma = 1.0 / sigma;
3083            let inv_sigma2 = inv_sigma * inv_sigma;
3084
3085            for k in 0..d {
3086                let ck = lambdas[k] * marginal_eigenvalues[k][multi_idx[k]];
3087                grad[k] += ck * inv_sigma;
3088            }
3089            if has_double_penalty && joint_null {
3090                grad[d] += lambdas[d] * inv_sigma;
3091            }
3092
3093            for k in 0..n_pen {
3094                let ck = if k < d {
3095                    lambdas[k] * marginal_eigenvalues[k][multi_idx[k]]
3096                } else if joint_null {
3097                    lambdas[d]
3098                } else {
3099                    0.0
3100                };
3101                // When ck == 0 (a zero λ, a zero marginal eigenvalue, or a cell
3102                // outside the joint null for the ridge penalty) every term this
3103                // index k contributes — `ck·inv_sigma − ck²·inv_sigma2` on the
3104                // diagonal and `−ck·cl·inv_sigma2` on every off-diagonal — is
3105                // exactly 0.0, so adding them to the finite running accumulators
3106                // is a bit-identical no-op. Skip the inner sweep entirely.
3107                if ck == 0.0 {
3108                    continue;
3109                }
3110                hess[[k, k]] += ck * inv_sigma - ck * ck * inv_sigma2;
3111                for l in (k + 1)..n_pen {
3112                    let cl = if l < d {
3113                        lambdas[l] * marginal_eigenvalues[l][multi_idx[l]]
3114                    } else if joint_null {
3115                        lambdas[d]
3116                    } else {
3117                        0.0
3118                    };
3119                    let off = -ck * cl * inv_sigma2;
3120                    hess[[k, l]] += off;
3121                    hess[[l, k]] += off;
3122                }
3123            }
3124        }
3125
3126        if kronecker_multi_index_advance(&mut multi_idx, marginal_dims) {
3127            break;
3128        }
3129    }
3130
3131    (logdet, grad, hess)
3132}
3133
3134// #1521: `KroneckerInvariantStructure` is defined once in `crate::kronecker`
3135// (the leaf data+compute module). The byte-identical copy that the carve left
3136// here is replaced by an import so the cache and this engine share one type.
3137use crate::kronecker::KroneckerInvariantStructure;
3138
3139/// Kronecker-factored reparameterization reusing a precomputed λ-invariant
3140/// structure (eigensystems, reparameterized marginals, shrinkage scale).
3141///
3142/// Bit-identical to `kronecker_reparameterization_engine` for the same marginal
3143/// data — the only difference is that the `eigh()` / `B_k U_k` work was hoisted
3144/// out of the per-iterate path into the cached `invariant`. Only the λ-dependent
3145/// `kronecker_logdet_and_derivatives` sweep runs here.
3146pub fn kronecker_reparameterization_engine_with_invariant(
3147    invariant: &KroneckerInvariantStructure,
3148    marginal_dims: &[usize],
3149    lambdas: &[f64],
3150    has_double_penalty: bool,
3151) -> Result<KroneckerReparamResult, EstimationError> {
3152    // Arc refcount bumps — the underlying eigensystems / reparameterized
3153    // marginals are λ-invariant and shared with the cache, not deep-copied.
3154    let marginal_eigenvalues = Arc::clone(&invariant.marginal_eigenvalues);
3155    let marginal_qs = Arc::clone(&invariant.marginal_qs);
3156    let reparameterized_marginals = Arc::clone(&invariant.reparameterized_marginals);
3157
3158    let marginal_eigenvalue_views: Vec<_> = marginal_eigenvalues
3159        .iter()
3160        .map(|evals| evals.view())
3161        .collect();
3162    let (log_det, det1, det2) = kronecker_logdet_and_derivatives(
3163        &marginal_eigenvalue_views,
3164        marginal_dims,
3165        lambdas,
3166        has_double_penalty,
3167        0.0,
3168    );
3169
3170    Ok(KroneckerReparamResult {
3171        reparameterized_marginals,
3172        marginal_eigenvalues,
3173        marginal_qs,
3174        log_det,
3175        det1,
3176        det2,
3177        has_double_penalty,
3178        marginal_dims: marginal_dims.to_vec(),
3179    })
3180}
3181
3182#[cfg(test)]
3183mod tests {
3184    /// gam#2454: the reparameterization's penalized rank and the shared
3185    /// balanced structural rank are one number. Three unit directions penalized
3186    /// at Frobenius norms `1e6`, `1e-9` and `1` (overlapping full-width
3187    /// components, the dense fallback path): the unweighted sum would drop the
3188    /// `1e-9` direction at a `100·p·ε·max` cut; the balanced rule keeps all three
3189    /// and the split's `q_pen` spans all three.
3190    #[test]
3191    fn reparam_split_rank_is_the_balanced_structural_rank_2454() {
3192        use ndarray::array;
3193        let p = 3usize;
3194        let roots = [
3195            array![[1.0e3, 0.0, 0.0]],
3196            array![[0.0, 1.0e-9_f64.sqrt(), 0.0]],
3197            array![[0.0, 0.0, 1.0]],
3198        ];
3199        let penalties: Vec<super::CanonicalPenalty> = roots
3200            .iter()
3201            .map(|root| super::CanonicalPenalty::from_dense_root(root.clone(), p))
3202            .collect();
3203        let balanced = super::balanced_penalty_structural_rank(
3204            penalties
3205                .iter()
3206                .map(|cp| (cp.local_ref().view(), cp.col_range.clone())),
3207            p,
3208        )
3209        .expect("balanced rank");
3210        assert_eq!(balanced, 3);
3211        let invariant = super::precompute_reparam_invariant_from_canonical(&penalties, p)
3212            .expect("reparam invariant");
3213        assert_eq!(
3214            invariant.split.q_pen.ncols(),
3215            balanced,
3216            "the split's penalized subspace has the balanced structural rank"
3217        );
3218        assert_eq!(invariant.split.q_null.ncols(), 0);
3219    }
3220
3221    use super::{
3222        CanonicalPenalty, REL_PSD_FLOOR, SubspaceLeakageMetrics, assess_subspace_leakage,
3223        classify_eigenvalues_strict, precompute_reparam_invariant_from_canonical,
3224        report_penalty_pair_redundancy, stable_reparameterizationwith_invariant,
3225        subspace_split_is_consistent,
3226    };
3227    use crate::EstimationError;
3228    use faer::Mat;
3229    use gam_linalg::utils::inf_norm;
3230    use ndarray::{Array1, Array2, array};
3231
3232    /// Build CanonicalPenalty values from full-width roots for tests.
3233    fn canonical_from_roots(rs_list: &[Array2<f64>], p: usize) -> Vec<CanonicalPenalty> {
3234        rs_list
3235            .iter()
3236            .map(|r| {
3237                let local = r.t().dot(r);
3238                CanonicalPenalty {
3239                    root: r.clone(),
3240                    col_range: 0..p,
3241                    total_dim: p,
3242                    nullity: 0,
3243                    local,
3244                    prior_mean: Array1::zeros(p),
3245                    positive_eigenvalues: Vec::new(),
3246                    op: None,
3247                }
3248            })
3249            .collect()
3250    }
3251
3252    fn metrics_for(
3253        qs: &Mat<f64>,
3254        rs: &[Mat<f64>],
3255        structural_rank: usize,
3256        p: usize,
3257    ) -> SubspaceLeakageMetrics {
3258        assess_subspace_leakage(qs, rs, structural_rank, p)
3259    }
3260
3261    #[test]
3262    fn subspace_leakage_iszero_for_clean_split() {
3263        let p = 4usize;
3264        let structural_rank = 2usize;
3265        let qs = Mat::<f64>::identity(p, p);
3266        let mut r0 = Mat::<f64>::zeros(2, p);
3267        r0[(0, 0)] = 1.0;
3268        r0[(1, 1)] = 2.0;
3269
3270        let m = metrics_for(&qs, &[r0], structural_rank, p);
3271        assert!(m.max_abs_sq <= 1e-16);
3272        assert!(m.max_rel_sq <= 1e-16);
3273        assert!(m.max_cross_gram_abs <= 1e-16);
3274    }
3275
3276    #[test]
3277    fn subspace_leakage_detects_null_column_energy() {
3278        let p = 4usize;
3279        let structural_rank = 2usize;
3280        let qs = Mat::<f64>::identity(p, p);
3281        let mut r0 = Mat::<f64>::zeros(1, p);
3282        r0[(0, 2)] = 3.0;
3283
3284        let m = metrics_for(&qs, &[r0], structural_rank, p);
3285        assert!(m.max_abs_sq > 0.0);
3286        assert!(m.max_rel_sq > 0.99);
3287    }
3288
3289    #[test]
3290    fn subspace_leakage_detects_qp_qn_nonorthogonality() {
3291        let p = 3usize;
3292        let structural_rank = 1usize;
3293        let mut qs = Mat::<f64>::identity(p, p);
3294        qs[(0, 1)] = 0.2;
3295        let r0 = Mat::<f64>::zeros(1, p);
3296
3297        let m = metrics_for(&qs, &[r0], structural_rank, p);
3298        assert!(m.max_cross_gram_abs > 1e-3);
3299    }
3300
3301    #[test]
3302    fn subspace_split_admits_near_threshold_manifold_leakage_1802() {
3303        // #1802: on sphere / Duchon / spline-on-sphere bases the REML outer
3304        // startup rejected EVERY candidate seed with
3305        //   "Reparameterization subspace split is inconsistent:
3306        //    max null leakage 1.174e-4 (rel 1.031e-4, worst penalty 0),
3307        //    max |Qp'Qn| 4.316e-16"
3308        // The split is perfectly ORTHOGONAL (|Qp'Qn| ≈ 4e-16); the only tripped
3309        // quantity is the transformed-root null-block leakage, whose RELATIVE
3310        // energy (1.031e-4)² ≈ 1.06e-8 sits right at `REL_PSD_FLOOR`. That is the
3311        // eigensolver's own numerically-PSD floor — the same floor
3312        // `classify_eigenvalues_strict` uses to declare a mode null — so a
3313        // manifold penalty whose Laplace-Beltrami spectrum decays through the
3314        // rank threshold with no clean gap MUST be admitted. Reproduce that exact
3315        // signature and assert the guard accepts it.
3316        let p = 40usize;
3317        let structural_rank = p - 1;
3318        // Unit-amplitude range column + a null column at √(REL_PSD_FLOOR)
3319        // amplitude, so the null-block relative energy lands at ~REL_PSD_FLOOR.
3320        let null_amp = (1.06e-8_f64).sqrt();
3321        let mut rs = Mat::<f64>::zeros(1, p);
3322        rs[(0, 0)] = 1.0;
3323        rs[(0, p - 1)] = null_amp;
3324        let qs = Mat::<f64>::identity(p, p);
3325        let leakage = metrics_for(&qs, &[rs], structural_rank, p);
3326        // The leakage reproduces the observed band: above the OLD fixed 1e-10
3327        // tolerance (which rejected the sphere fit) yet a benign ~REL_PSD_FLOOR.
3328        assert!(
3329            leakage.max_rel_sq > 1e-10 && leakage.max_rel_sq < 1e-6,
3330            "reproduced leakage should sit in the near-REL_PSD_FLOOR band, got {:.3e}",
3331            leakage.max_rel_sq
3332        );
3333        assert!(leakage.max_cross_gram_abs <= 1e-12);
3334        assert!(
3335            subspace_split_is_consistent(&leakage, p),
3336            "near-REL_PSD_FLOOR null leakage on a manifold basis must be admitted \
3337             (rel_sq={:.3e}, tol={:.3e})",
3338            leakage.max_rel_sq,
3339            (p as f64) * REL_PSD_FLOOR,
3340        );
3341    }
3342
3343    #[test]
3344    fn subspace_split_still_rejects_genuine_inconsistency_1802() {
3345        // The #1802 relaxation only widens the leakage tolerance to the
3346        // classifier's own `p · REL_PSD_FLOOR` floor; it must NOT admit a
3347        // genuinely broken split. A whole penalized mode dumped into the null
3348        // block (O(1) relative leakage) is still rejected.
3349        let p = 40usize;
3350        let structural_rank = p - 1;
3351        let mut rs = Mat::<f64>::zeros(1, p);
3352        rs[(0, p - 1)] = 1.0; // all penalty-root energy in the null column
3353        let qs = Mat::<f64>::identity(p, p);
3354        let leakage = metrics_for(&qs, &[rs], structural_rank, p);
3355        assert!(leakage.max_rel_sq > 0.99);
3356        assert!(
3357            !subspace_split_is_consistent(&leakage, p),
3358            "an O(1) null-block leakage is a real inconsistency and must be rejected"
3359        );
3360
3361        // A non-orthogonal split is rejected regardless of root leakage.
3362        let mut qs_bad = Mat::<f64>::identity(3, 3);
3363        qs_bad[(0, 1)] = 0.2;
3364        let clean = Mat::<f64>::zeros(1, 3);
3365        let leakage2 = metrics_for(&qs_bad, &[clean], 1, 3);
3366        assert!(leakage2.max_cross_gram_abs > 1e-3);
3367        assert!(
3368            !subspace_split_is_consistent(&leakage2, 3),
3369            "a non-orthogonal Qp/Qn split must be rejected"
3370        );
3371    }
3372
3373    #[test]
3374    fn u_truncated_is_transformed_frame_in_nonzero_case() {
3375        let p = 3usize;
3376        let rs_list = vec![array![[1.0, 0.0, 0.0]]];
3377        let canonical = canonical_from_roots(&rs_list, p);
3378        let lambdas = vec![2.0];
3379        let inv = precompute_reparam_invariant_from_canonical(&canonical, p)
3380            .expect("precompute invariant");
3381        let rep = stable_reparameterizationwith_invariant(&canonical, &lambdas, p, &inv)
3382            .expect("stable reparam");
3383
3384        let expected = rep.qs.t().dot(&inv.split.q_null);
3385        let diff = &rep.u_truncated - &expected;
3386        let max_abs = inf_norm(diff.iter().copied());
3387        assert!(
3388            max_abs <= 1e-10,
3389            "u_truncated frame mismatch: max_abs={max_abs}"
3390        );
3391    }
3392
3393    #[test]
3394    fn infinite_lambda_keeps_range_penalty_block_finite_1379() {
3395        // gam#1379 / gam#1074: a genuinely infinite λ = exp(ρ) is NOT silently
3396        // clamped to a finite ceiling. The original #1379 fix added a 1e300
3397        // ceiling so `∞ · 0` could not poison the range block Σ_k λ_k S_k, but
3398        // #1074 DELETED that clamp on purpose (see the comment at the top of
3399        // `stable_reparameterizationwith_invariant`): masking ∞ hid the real
3400        // defect — the outer optimizer driving a redundant/unidentified penalty
3401        // direction off to ∞ instead of that direction being detected and
3402        // dropped. With the clamp gone, a literal `f64::INFINITY` λ surfaces as
3403        // a clean, detectable error (the eigensolver rejects the NaN-poisoned
3404        // block) rather than a silent finite success. Pin that contract: ∞ must
3405        // ERROR, not be quietly clamped.
3406        //
3407        // Fixture: two penalties on a 3-wide block. The first penalizes only
3408        // coordinate 0 (so its block S_k has structural zeros everywhere except
3409        // [0,0]); give it λ = +∞. The second penalizes coordinate 1 at a normal
3410        // λ.
3411        let p = 3usize;
3412        let rs_list = vec![array![[1.0, 0.0, 0.0]], array![[0.0, 1.0, 0.0]]];
3413        let canonical = canonical_from_roots(&rs_list, p);
3414        let inv = precompute_reparam_invariant_from_canonical(&canonical, p)
3415            .expect("precompute invariant");
3416
3417        let lambdas_inf = vec![f64::INFINITY, 3.0];
3418        let inf_result =
3419            stable_reparameterizationwith_invariant(&canonical, &lambdas_inf, p, &inv);
3420        assert!(
3421            inf_result.is_err(),
3422            "an infinite lambda must surface as an error, not be silently clamped (#1074)"
3423        );
3424
3425        // A finite (even very large) λ must still produce an all-finite reparam:
3426        // the function is robust to large-but-finite penalties; only the
3427        // non-finite input is rejected.
3428        let lambdas_big = vec![1e300_f64, 3.0];
3429        let rep = stable_reparameterizationwith_invariant(&canonical, &lambdas_big, p, &inv)
3430            .expect("stable reparam at large-but-finite lambda");
3431        assert!(
3432            rep.s_transformed.iter().all(|v| v.is_finite()),
3433            "transformed penalty must be finite at large-but-finite lambda"
3434        );
3435        assert!(
3436            rep.qs.iter().all(|v| v.is_finite()),
3437            "reparam rotation must be finite at large-but-finite lambda"
3438        );
3439        assert!(
3440            rep.log_det.is_finite(),
3441            "penalty log-det must be finite at large-but-finite lambda"
3442        );
3443        assert!(
3444            rep.det1.iter().all(|v| v.is_finite()),
3445            "penalty log-det derivatives must be finite at large-but-finite lambda"
3446        );
3447    }
3448
3449    #[test]
3450    fn u_truncated_is_identitywhen_no_penalties() {
3451        let p = 4usize;
3452        let canonical: Vec<CanonicalPenalty> = Vec::new();
3453        let lambdas: Vec<f64> = Vec::new();
3454        let inv = precompute_reparam_invariant_from_canonical(&canonical, p)
3455            .expect("precompute invariant");
3456        let rep = stable_reparameterizationwith_invariant(&canonical, &lambdas, p, &inv)
3457            .expect("stable reparam");
3458        assert_eq!(rep.u_truncated, Array2::<f64>::eye(p));
3459    }
3460
3461    #[test]
3462    fn transformed_penalty_is_diagonal_in_transformed_frame() {
3463        let p = 3usize;
3464        let inv_sqrt2 = 2.0_f64.sqrt().recip();
3465        // Penalize a rotated direction in original space so Qs is non-trivial.
3466        let rs_list = vec![array![[inv_sqrt2, inv_sqrt2, 0.0]]];
3467        let canonical = canonical_from_roots(&rs_list, p);
3468        let lambdas = vec![4.0];
3469        let inv = precompute_reparam_invariant_from_canonical(&canonical, p)
3470            .expect("precompute invariant");
3471        let rep = stable_reparameterizationwith_invariant(&canonical, &lambdas, p, &inv)
3472            .expect("stable reparam");
3473
3474        assert_eq!(rep.e_transformed.nrows(), 1);
3475        assert!(rep.e_transformed[[0, 0]].abs() > 0.0);
3476        assert!(rep.e_transformed[[0, 1]].abs() <= 1e-12);
3477        assert!(rep.e_transformed[[0, 2]].abs() <= 1e-12);
3478        // Exact pseudo-logdet on the structural penalized block has no
3479        // delta-dependent nullspace normalization.
3480        let expected_det1 = 1.0_f64;
3481        assert!((rep.det1[0] - expected_det1).abs() <= 1e-12);
3482
3483        let s = rep.s_transformed;
3484        let mut max_offdiag = 0.0_f64;
3485        for i in 0..p {
3486            for j in 0..p {
3487                if i != j {
3488                    max_offdiag = max_offdiag.max(s[[i, j]].abs());
3489                }
3490            }
3491        }
3492        assert!(
3493            max_offdiag <= 1e-10,
3494            "transformed penalty should be diagonal, max offdiag={max_offdiag}"
3495        );
3496        assert!(s[[1, 1]].abs() <= 1e-10);
3497        assert!(s[[2, 2]].abs() <= 1e-10);
3498    }
3499
3500    #[test]
3501    fn det1_matches_rank_for_single_full_rank_penalty() {
3502        let p = 2usize;
3503        let inv_sqrt2 = 2.0_f64.sqrt().recip();
3504        // Q^T for a 45-degree rotation.
3505        let q_t = [[inv_sqrt2, inv_sqrt2], [-inv_sqrt2, inv_sqrt2]];
3506        // R = diag(3, 1) * Q^T gives S = Q * diag(9, 1) * Q^T.
3507        let rs = array![
3508            [3.0 * q_t[0][0], 3.0 * q_t[0][1]],
3509            [1.0 * q_t[1][0], 1.0 * q_t[1][1]]
3510        ];
3511        let rs_list = vec![rs];
3512        let canonical = canonical_from_roots(&rs_list, p);
3513        let lambdas = vec![5.0];
3514
3515        let inv = precompute_reparam_invariant_from_canonical(&canonical, p)
3516            .expect("precompute invariant");
3517        let rep = stable_reparameterizationwith_invariant(&canonical, &lambdas, p, &inv)
3518            .expect("stable reparam");
3519
3520        assert_eq!(rep.e_transformed.nrows(), p);
3521        let det1 = rep.det1[0];
3522        // Exact pseudo-logdet on the structural penalized block:
3523        //   det1 = lambda * sum_l d_l / (lambda*d_l)
3524        // where d_l are eigenvalues of S_k.
3525        let s_k_eigs = [9.0_f64, 1.0_f64];
3526        let lambda = 5.0_f64;
3527        let expected_det1: f64 = s_k_eigs.iter().map(|&d| lambda * d / (lambda * d)).sum();
3528        assert!(
3529            (det1 - expected_det1).abs() <= 1e-12,
3530            "expected det1={expected_det1}, got {det1}",
3531        );
3532
3533        let s = rep.s_transformed;
3534        assert!(s[[0, 1]].abs() <= 1e-10);
3535        assert!(s[[1, 0]].abs() <= 1e-10);
3536        assert!(s[[0, 0]] > 0.0);
3537        assert!(s[[1, 1]] > 0.0);
3538    }
3539
3540    #[test]
3541    fn classify_strict_rejects_nan_eigenvalue() {
3542        let mut eigs = [1.0, f64::NAN, 0.5];
3543        match classify_eigenvalues_strict(&mut eigs, "test_nan") {
3544            Err(EstimationError::PenaltySpectrumNonFinite {
3545                context,
3546                index,
3547                value,
3548            }) => {
3549                assert_eq!(context, "test_nan");
3550                assert_eq!(index, 1);
3551                assert!(value.is_nan());
3552            }
3553            other => panic!("expected PenaltySpectrumNonFinite, got {:?}", other),
3554        }
3555    }
3556
3557    #[test]
3558    fn classify_strict_rejects_inf_eigenvalue() {
3559        let mut eigs = [1.0, 0.5, f64::INFINITY];
3560        match classify_eigenvalues_strict(&mut eigs, "test_inf") {
3561            Err(EstimationError::PenaltySpectrumNonFinite { index, value, .. }) => {
3562                assert_eq!(index, 2);
3563                assert!(value.is_infinite());
3564            }
3565            other => panic!("expected PenaltySpectrumNonFinite, got {:?}", other),
3566        }
3567    }
3568
3569    #[test]
3570    fn classify_strict_rejects_materially_indefinite() {
3571        // -1e-2 with scale ~1.0 is well above any reasonable roundoff tolerance.
3572        let mut eigs = [1.0, -1e-2, 0.5];
3573        match classify_eigenvalues_strict(&mut eigs, "test_indef") {
3574            Err(EstimationError::PenaltySpectrumIndefinite {
3575                context,
3576                index,
3577                value,
3578                ..
3579            }) => {
3580                assert_eq!(context, "test_indef");
3581                assert_eq!(index, 1);
3582                assert!((value + 1e-2).abs() <= 1e-15);
3583            }
3584            other => panic!("expected PenaltySpectrumIndefinite, got {:?}", other),
3585        }
3586    }
3587
3588    #[test]
3589    fn classify_strict_accepts_roundoff_negative() {
3590        // -1e-16 * scale is well within tol = 64 * eps * p * scale.
3591        let scale = 1.0_f64;
3592        let roundoff = -1e-16 * scale;
3593        let mut eigs = [scale, 0.5 * scale, roundoff, 0.25 * scale];
3594        classify_eigenvalues_strict(&mut eigs, "test_roundoff").expect("roundoff must classify");
3595        // The roundoff eigenvalue is snapped to exact zero.
3596        assert_eq!(eigs[2], 0.0);
3597        // Strictly positive entries must be preserved.
3598        assert!(eigs[0] > 0.0 && eigs[1] > 0.0 && eigs[3] > 0.0);
3599    }
3600
3601    #[test]
3602    fn classify_strict_accepts_extreme_lambda_assembly_noise_1619() {
3603        // #1619: high-rank thin-plate / Duchon penalties (p≈200) assembled and
3604        // reparameterized at extreme λ produce float-noise negative eigenvalues at
3605        // ~1e-11 relative to the spectrum scale (~1e13 there). The bare machine-ε
3606        // floor (~12×ε relative ≈ 1e-12) rejected these as "indefinite" and
3607        // spuriously failed the inner P-IRLS solve. They are PSD to numerical
3608        // precision and must be snapped to zero, not rejected.
3609        let scale = 8.509e12_f64;
3610        // Worst (eig, scale) pair observed in the issue: eig/scale ≈ -7.7e-11.
3611        let noise = -6.546e2_f64;
3612        assert!(
3613            (noise.abs() / scale) < 1.0e-10,
3614            "fixture must reproduce the ~1e-11-relative noise from #1619"
3615        );
3616        let mut eigs = vec![scale, 0.5 * scale, noise, 0.1 * scale];
3617        classify_eigenvalues_strict(&mut eigs, "range penalty block")
3618            .expect("a ~1e-11-relative roundoff-negative eigenvalue must be accepted (#1619)");
3619        // The roundoff-negative eigenvalue is snapped to exact zero.
3620        assert_eq!(eigs[2], 0.0);
3621        // Strictly positive entries are preserved.
3622        assert!(eigs[0] > 0.0 && eigs[1] > 0.0 && eigs[3] > 0.0);
3623    }
3624
3625    #[test]
3626    fn classify_strict_snaps_subtol_positive_to_zero() {
3627        // Positive eigenvalues below the tolerance are also snapped to exact 0
3628        // so downstream rank counts and pseudo-logdets are deterministic.
3629        let scale = 10.0_f64;
3630        let subtol = 1e-15 * scale;
3631        let mut eigs = [scale, subtol];
3632        classify_eigenvalues_strict(&mut eigs, "test_sub_pos").expect("sub-tol positive ok");
3633        assert_eq!(eigs[1], 0.0);
3634    }
3635
3636    /// Build a `CanonicalPenalty` directly from a symmetric `local` matrix.
3637    /// Bypasses root extraction — the redundancy diagnostic only reads `local`
3638    /// and `col_range`, so the rest is filler.
3639    fn canonical_from_local(
3640        local: Array2<f64>,
3641        col_range: std::ops::Range<usize>,
3642        total_dim: usize,
3643    ) -> CanonicalPenalty {
3644        let block_dim = local.nrows();
3645        // A trivially valid root: zero rank. The diagnostic doesn't read root.
3646        let root = Array2::<f64>::zeros((0, block_dim));
3647        CanonicalPenalty {
3648            root,
3649            col_range,
3650            total_dim,
3651            nullity: 0,
3652            local,
3653            prior_mean: Array1::zeros(block_dim),
3654            positive_eigenvalues: Vec::new(),
3655            op: None,
3656        }
3657    }
3658
3659    #[test]
3660    fn report_penalty_pair_redundancy_detects_identical_pair() {
3661        // Penalty 0: a "generic" SPD matrix on cols 0..3.
3662        let s0 = ndarray::array![[2.0, 0.5, 0.0], [0.5, 1.0, 0.25], [0.0, 0.25, 1.5],];
3663        // Penalties 1 and 2: identical block-local penalty on the SAME col_range.
3664        // This is the Z₂-symmetric saddle scenario.
3665        let s_shared = ndarray::array![[1.0, -0.5, 0.0], [-0.5, 2.0, -0.5], [0.0, -0.5, 1.0],];
3666
3667        let bundle = vec![
3668            canonical_from_local(s0, 0..3, 3),
3669            canonical_from_local(s_shared.clone(), 0..3, 3),
3670            canonical_from_local(s_shared, 0..3, 3),
3671        ];
3672
3673        let redundant = report_penalty_pair_redundancy(&bundle);
3674
3675        // Exactly one redundant pair: (1, 2). Pairs (0, 1) and (0, 2) involve
3676        // distinct matrices and must NOT be flagged.
3677        assert_eq!(
3678            redundant.len(),
3679            1,
3680            "expected exactly one redundant pair, got {:?}",
3681            redundant
3682        );
3683        let (i, j, defect) = redundant[0];
3684        assert_eq!((i, j), (1, 2));
3685        assert_eq!(
3686            defect, 0.0,
3687            "bit-identical penalties leave an EXACTLY zero residual, not a small one"
3688        );
3689    }
3690
3691    /// The regression this whole screen was rewritten for (#2676): a pair whose
3692    /// relative defect is `1.9e-5` — the measured `geo_disease_matern` figure at
3693    /// its cold geometry — must NOT be reported as proportional, even though its
3694    /// cosine rounds to `1.000000` at six decimals.
3695    ///
3696    /// The two assertions are a pair. The first is the fix; the second is the
3697    /// evidence that the OLD screen would have failed it, so the test cannot go
3698    /// green for the boring reason that the fixture is not marginal.
3699    #[test]
3700    fn a_pair_1p9e_minus_5_apart_is_not_proportional_2676() {
3701        const DEFECT: f64 = 1.874_020e-5;
3702        let base = ndarray::array![[1.0, -0.5, 0.0], [-0.5, 2.0, -0.5], [0.0, -0.5, 1.0]];
3703        let base_norm = base.iter().map(|v| v * v).sum::<f64>().sqrt();
3704        // A perturbation ORTHOGONAL to `base` in the Frobenius inner product,
3705        // so the defect is exactly the perturbation's relative size and no part
3706        // of it is absorbed into the best scale `c`.
3707        let mut direction = ndarray::array![[0.0, 0.0, 1.0], [0.0, 0.0, 0.0], [1.0, 0.0, 0.0]];
3708        let overlap = direction
3709            .iter()
3710            .zip(base.iter())
3711            .map(|(a, b)| a * b)
3712            .sum::<f64>()
3713            / (base_norm * base_norm);
3714        direction = &direction - &(overlap * &base);
3715        let direction_norm = direction.iter().map(|v| v * v).sum::<f64>().sqrt();
3716        let perturbed = &base + &(direction * (DEFECT * base_norm / direction_norm));
3717
3718        let bundle = vec![
3719            canonical_from_local(base.clone(), 0..3, 3),
3720            canonical_from_local(perturbed.clone(), 0..3, 3),
3721        ];
3722        assert!(
3723            report_penalty_pair_redundancy(&bundle).is_empty(),
3724            "a pair {DEFECT:.3e} apart is measurably distinct and must not be called proportional"
3725        );
3726
3727        // Negative control: the cosine the old screen thresholded on cannot see
3728        // it. `1 - cos = defect^2/2 = 1.8e-10`, six orders under the `1e-8` bar
3729        // that used to declare the pair "structurally identical".
3730        let inner = base
3731            .iter()
3732            .zip(perturbed.iter())
3733            .map(|(a, b)| a * b)
3734            .sum::<f64>();
3735        let cosine =
3736            inner / (base_norm * perturbed.iter().map(|v| v * v).sum::<f64>().sqrt());
3737        assert!(
3738            cosine > 1.0 - 1e-8,
3739            "the fixture must be one the OLD cos > 1 - 1e-8 screen accepted; got cos = {cosine}"
3740        );
3741        assert_eq!(
3742            format!("{cosine:.6}"),
3743            "1.000000",
3744            "and one whose six-decimal print is indistinguishable from an exact identity"
3745        );
3746    }
3747
3748    #[test]
3749    fn report_penalty_pair_redundancy_skips_different_col_ranges() {
3750        // Two identical local matrices but on disjoint col_ranges. The
3751        // function must NOT flag them — they live in different parameter
3752        // subspaces by construction.
3753        let s = ndarray::array![[1.0, 0.0], [0.0, 1.0]];
3754        let bundle = vec![
3755            canonical_from_local(s.clone(), 0..2, 4),
3756            canonical_from_local(s, 2..4, 4),
3757        ];
3758        let redundant = report_penalty_pair_redundancy(&bundle);
3759        assert!(
3760            redundant.is_empty(),
3761            "different col_ranges must not be flagged"
3762        );
3763    }
3764}