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
17const 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 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
311fn classify_eigenvalues_strict(
331 eigenvalues: &mut [f64],
332 context: &str,
333) -> Result<(), EstimationError> {
334 const C_EPS_P_FACTOR: f64 = 64.0;
335 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 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 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
554fn 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#[derive(Clone)]
608pub struct ReparamResult {
609 pub s_transformed: Array2<f64>,
613 pub log_det: f64,
615 pub det1: Array1<f64>,
617 pub qs: Array2<f64>,
619 pub canonical_transformed: Vec<CanonicalPenalty>,
624 pub e_transformed: Array2<f64>,
627 pub u_truncated: Array2<f64>,
637}
638
639#[derive(Clone, Copy, Debug, PartialEq, Eq)]
646pub enum PenaltyFrame {
647 Original,
648 Transformed,
649}
650
651#[derive(Clone, Debug)]
686pub struct PenaltyNullSplit {
687 transformed: Option<Array2<f64>>,
688 original: Option<Array2<f64>>,
689}
690
691impl PenaltyNullSplit {
692 pub fn identity() -> Self {
696 Self {
697 transformed: None,
698 original: None,
699 }
700 }
701
702 pub fn declared_null_dim(&self) -> usize {
704 self.transformed.as_ref().map_or(0, |n| n.ncols())
705 }
706
707 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 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 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 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 u.clone()
766 };
767 PenaltyNullSplit {
768 transformed: Some(u.clone()),
769 original: Some(original),
770 }
771 }
772}
773
774struct KroneckerFactorDecomp {
780 root: Array2<f64>, positive_eigenvalues: Vec<f64>, rank: usize,
783 dim: usize,
784}
785
786fn 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 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
862fn 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
884fn 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#[derive(Clone)]
913pub struct CanonicalPenalty {
914 pub root: Array2<f64>,
917 pub col_range: std::ops::Range<usize>,
920 pub total_dim: usize,
922 pub nullity: usize,
924 pub local: Array2<f64>,
928 pub prior_mean: Array1<f64>,
930 pub positive_eigenvalues: Vec<f64>,
933 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 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(); 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 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 pub fn rank(&self) -> usize {
999 self.root.nrows()
1000 }
1001
1002 pub fn block_dim(&self) -> usize {
1004 self.col_range.len()
1005 }
1006
1007 pub const fn is_block_local(&self) -> bool {
1009 self.col_range.start != 0 || self.col_range.end != self.total_dim
1010 }
1011
1012 pub fn local_ref(&self) -> &Array2<f64> {
1015 &self.local
1016 }
1017
1018 pub fn local_penalty(&self) -> Array2<f64> {
1021 self.local.clone()
1022 }
1023
1024 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 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 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 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 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 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 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 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 nullity: self.nullity,
1185 local,
1186 prior_mean,
1187 positive_eigenvalues,
1188 op: None,
1189 })
1190 }
1191
1192 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
1209pub fn report_penalty_pair_redundancy(canonical: &[CanonicalPenalty]) -> Vec<(usize, usize, f64)> {
1296 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 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 if canonical[i].col_range != canonical[j].col_range {
1326 continue;
1327 }
1328 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 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 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 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
1398pub 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 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 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 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 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 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 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 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 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
1582pub 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
1609pub(crate) const OVERLAPPING_PENALTY_DENSE_FALLBACK_MAX_P: usize = 4096;
1619
1620pub 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 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 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 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 struct BlockRoot {
1700 col_range: Range<usize>,
1701 root: Array2<f64>, }
1703 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 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 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
1775pub const BALANCED_PENALTY_RANK_RELATIVE_TOL: f64 = 1.0e-12;
1780
1781pub 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
1821pub 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
1831pub 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#[derive(Clone)]
1926pub struct ReparamInvariant {
1927 split: SubspaceSplit,
1928 qs_base: Array2<f64>,
1932 has_nonzero: bool,
1933 max_balanced_eigenvalue: f64,
1936}
1937
1938pub 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 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 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 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 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 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>, q_null_local: Array2<f64>, max_balanced_eigenvalue: f64,
2083 pen_col_offset: usize,
2085 null_col_offset: usize,
2087 }
2088
2089 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 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, null_col_offset: 0, });
2122 }
2123
2124 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, null_col_offset: 0, })
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 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 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
2245fn 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 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
2282fn 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 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 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 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 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 let total_root_rows: usize = rs_transformed.iter().map(Mat::nrows).sum();
2440 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 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 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 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 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 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 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 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 for j in 0..penalized_rank {
2628 e_transformed_mat[(row_idx, j)] = sqrt_eigenval * range_rotation[(j, row_idx)];
2629 }
2630 }
2631
2632 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 let det1vec: Vec<f64> = (0..lambdas.len())
2667 .into_par_iter()
2668 .map(|k| {
2669 let s_k = &s_k_penalized_cache[k];
2670 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 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 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 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#[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
2783pub 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#[derive(Clone)]
2816pub struct KroneckerReparamResult {
2817 pub reparameterized_marginals: Arc<Vec<Array2<f64>>>,
2823 pub marginal_eigenvalues: Arc<Vec<Array1<f64>>>,
2825 pub marginal_qs: Arc<Vec<Array2<f64>>>,
2827 pub log_det: f64,
2829 pub det1: Array1<f64>,
2831 pub det2: Array2<f64>,
2833 pub has_double_penalty: bool,
2835 pub marginal_dims: Vec<usize>,
2837}
2838
2839impl KroneckerReparamResult {
2840 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 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 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 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 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 let d = self.marginal_dims.len();
2923 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 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; 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
2989const KRONECKER_STRUCTURAL_ZERO_TOL: f64 = 1e-12;
2996
2997#[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#[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 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
3134use crate::kronecker::KroneckerInvariantStructure;
3138
3139pub 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 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 #[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 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 let p = 40usize;
3317 let structural_rank = p - 1;
3318 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 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 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; 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 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 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 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 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 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 let q_t = [[inv_sqrt2, inv_sqrt2], [-inv_sqrt2, inv_sqrt2]];
3506 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 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 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 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 assert_eq!(eigs[2], 0.0);
3597 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 let scale = 8.509e12_f64;
3610 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 assert_eq!(eigs[2], 0.0);
3621 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 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 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 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 let s0 = ndarray::array![[2.0, 0.5, 0.0], [0.5, 1.0, 0.25], [0.0, 0.25, 1.5],];
3663 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 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 #[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 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 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 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}