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, ArrayViewMut2, Axis, 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 pub fn block_dimension(&self) -> usize {
49 match self {
50 PenaltyRepresentation::Dense(matrix) => matrix.nrows(),
51 PenaltyRepresentation::Banded { bands, offsets } => {
52 let mut dim = 0usize;
53 for (band, &offset) in bands.iter().zip(offsets.iter()) {
54 let len = band.len();
55 let extent = if offset >= 0 {
56 len + offset as usize
57 } else {
58 len + (-offset) as usize
59 };
60 dim = dim.max(extent);
61 }
62 dim
63 }
64 PenaltyRepresentation::Kronecker { left, right } => left.nrows() * right.nrows(),
65 }
66 }
67
68}
69
70#[derive(Clone)]
71pub struct PenaltyMatrix {
72 pub col_range: Range<usize>,
73 pub representation: PenaltyRepresentation,
74}
75
76impl PenaltyMatrix {
77 fn accumulate_into(&self, mut dest: ArrayViewMut2<'_, f64>, weight: f64) {
78 if weight == 0.0 {
79 return;
80 }
81 match &self.representation {
82 PenaltyRepresentation::Dense(block) => {
83 dest.scaled_add(weight, block);
84 }
85 PenaltyRepresentation::Banded { bands, offsets } => {
86 let positive_offsets: HashSet<usize> = offsets
87 .iter()
88 .filter_map(|&off| (off >= 0).then_some(off as usize))
89 .collect();
90 for (band, &offset) in bands.iter().zip(offsets.iter()) {
91 let off = offset.unsigned_abs() as usize;
92 if offset < 0 && positive_offsets.contains(&off) {
93 continue;
94 }
95 for (idx, &value) in band.iter().enumerate() {
96 let (i, j) = if offset >= 0 {
97 (idx, idx + off)
98 } else {
99 (idx + off, idx)
100 };
101 let Some(entry_ij) = dest.get_mut((i, j)) else {
102 continue;
103 };
104 *entry_ij += weight * value;
105 if i != j
106 && let Some(entry_ji) = dest.get_mut((j, i))
107 {
108 *entry_ji += weight * value;
109 }
110 }
111 }
112 }
113 PenaltyRepresentation::Kronecker { left, right } => {
114 let (lrows, l_cols) = left.dim();
115 let (rrows, r_cols) = right.dim();
116 for i in 0..lrows {
117 for j in 0..l_cols {
118 let scale = left[(i, j)] * weight;
119 if scale == 0.0 {
120 continue;
121 }
122 let mut block = dest.slice_mut(s![
123 i * rrows..(i + 1) * rrows,
124 j * r_cols..(j + 1) * r_cols
125 ]);
126 block.scaled_add(scale, right);
127 }
128 }
129 }
130 }
131 }
132
133 pub fn to_dense(&self, total_dim: usize) -> Array2<f64> {
134 let mut dense = Array2::<f64>::zeros((total_dim, total_dim));
135 self.accumulate_into(
136 dense.slice_mut(s![self.col_range.clone(), self.col_range.clone()]),
137 1.0,
138 );
139 dense
140 }
141}
142
143pub(crate) fn array_to_faer(array: &Array2<f64>) -> Mat<f64> {
144 let (rows, cols) = array.dim();
145 Mat::from_fn(rows, cols, |i, j| array[[i, j]])
146}
147
148pub(crate) fn mat_to_array(mat: &Mat<f64>) -> Array2<f64> {
149 let mut out = Array2::<f64>::zeros((mat.nrows(), mat.ncols()));
150 for i in 0..mat.nrows() {
151 for j in 0..mat.ncols() {
152 out[[i, j]] = mat[(i, j)];
153 }
154 }
155 out
156}
157
158fn mat_max_abs_element(matrix: MatRef<'_, f64>) -> f64 {
159 let (rows, cols) = matrix.shape();
160 let mut maxval = 0.0_f64;
161 for i in 0..rows {
162 for j in 0..cols {
163 let val = matrix[(i, j)];
164 if val.is_finite() {
165 maxval = maxval.max(val.abs());
166 }
167 }
168 }
169 maxval
170}
171
172fn sanitize_symmetric_faer(matrix: &Mat<f64>) -> Mat<f64> {
173 let (rows, cols) = matrix.as_ref().shape();
174 assert_eq!(rows, cols, "Matrix must be square for sanitization");
175
176 let mut sanitized = matrix.clone();
177
178 for i in 0..rows {
179 let diag = sanitized[(i, i)];
180 if !diag.is_finite() {
181 sanitized[(i, i)] = 0.0;
182 }
183 for j in (i + 1)..cols {
184 let mut upper = sanitized[(i, j)];
185 let mut lower = sanitized[(j, i)];
186 if !upper.is_finite() {
187 upper = 0.0;
188 }
189 if !lower.is_finite() {
190 lower = 0.0;
191 }
192 let avg = 0.5 * (upper + lower);
193 sanitized[(i, j)] = avg;
194 sanitized[(j, i)] = avg;
195 }
196 }
197
198 let scale = mat_max_abs_element(sanitized.as_ref());
199 let tiny = (scale * 1e-14).max(1e-30);
200 for i in 0..rows {
201 for j in 0..cols {
202 let val = sanitized[(i, j)];
203 if !val.is_finite() {
204 sanitized[(i, j)] = 0.0;
205 } else if val.abs() < tiny {
206 sanitized[(i, j)] = 0.0;
207 }
208 }
209 }
210
211 sanitized
212}
213
214fn penalty_from_root_faer(root: &Mat<f64>) -> Mat<f64> {
215 let cols = root.ncols();
216 let mut full = Mat::<f64>::zeros(cols, cols);
217 let root_ref = root.as_ref();
218 let root_t = root_ref.transpose();
219 matmul(
220 full.as_mut(),
221 Accum::Replace,
222 root_t,
223 root_ref,
224 1.0,
225 Par::Seq,
226 );
227 sanitize_symmetric_faer(&full)
228}
229
230fn symmetrize_faer_matrix_in_place(matrix: &mut Mat<f64>) {
231 let n = matrix.nrows().min(matrix.ncols());
232 for i in 0..n {
233 for j in 0..i {
234 let avg = 0.5 * (matrix[(i, j)] + matrix[(j, i)]);
235 matrix[(i, j)] = avg;
236 matrix[(j, i)] = avg;
237 }
238 }
239}
240
241fn orthogonal_similarity_transform_faer(
242 matrix: &Mat<f64>,
243 block_dim: usize,
244 orthogonal: &Mat<f64>,
245) -> Mat<f64> {
246 let matrix_block = matrix.as_ref().submatrix(0, 0, block_dim, block_dim);
247 let cols = orthogonal.ncols();
248 let mut temp = Mat::<f64>::zeros(block_dim, cols);
249 matmul(
250 temp.as_mut(),
251 Accum::Replace,
252 matrix_block,
253 orthogonal.as_ref(),
254 1.0,
255 Par::Seq,
256 );
257 let mut rotated = Mat::<f64>::zeros(cols, cols);
258 matmul(
259 rotated.as_mut(),
260 Accum::Replace,
261 orthogonal.transpose(),
262 temp.as_ref(),
263 1.0,
264 Par::Seq,
265 );
266 symmetrize_faer_matrix_in_place(&mut rotated);
267 rotated
268}
269
270fn trace_penalty_in_orthogonal_basis(
271 matrix: &Mat<f64>,
272 block_dim: usize,
273 orthogonal: &Mat<f64>,
274 rotated_eigenvalues: &[f64],
275 delta: f64,
276) -> f64 {
277 let matrix_block = matrix.as_ref().submatrix(0, 0, block_dim, block_dim);
278 let cols = orthogonal.ncols();
279 assert!(rotated_eigenvalues.len() >= cols);
280 let mut projected = Mat::<f64>::zeros(block_dim, cols);
281 matmul(
282 projected.as_mut(),
283 Accum::Replace,
284 matrix_block,
285 orthogonal.as_ref(),
286 1.0,
287 Par::Seq,
288 );
289 let mut trace = KahanSum::default();
290 for l in 0..cols {
291 let mut diag_ll = KahanSum::default();
292 for i in 0..block_dim {
293 diag_ll.add(orthogonal[(i, l)] * projected[(i, l)]);
294 }
295 trace.add(diag_ll.sum() / (rotated_eigenvalues[l] + delta));
296 }
297 trace.sum()
298}
299
300pub fn trace_reduced_penalty_covariance(
301 reduced_penalty: &Array2<f64>,
302 covariance_basis: &Array2<f64>,
303) -> f64 {
304 assert_eq!(
305 reduced_penalty.dim(),
306 covariance_basis.dim(),
307 "trace_reduced_penalty_covariance dimension mismatch"
308 );
309 let r = covariance_basis.nrows();
310 let mut trace = KahanSum::default();
311 for i in 0..r {
312 for j in 0..r {
313 trace.add(covariance_basis[[i, j]] * reduced_penalty[[j, i]]);
314 }
315 }
316 trace.sum()
317}
318
319pub fn trace_penalty_covariance_in_orthogonal_basis(
320 matrix: &Array2<f64>,
321 orthogonal: &Array2<f64>,
322 covariance_basis: &Array2<f64>,
323) -> f64 {
324 let reduced = gam_linalg::faer_ndarray::fast_ab(
325 &gam_linalg::faer_ndarray::fast_atb(orthogonal, matrix),
326 orthogonal,
327 );
328 trace_reduced_penalty_covariance(&reduced, covariance_basis)
329}
330
331fn classify_eigenvalues_strict(
351 eigenvalues: &mut [f64],
352 context: &str,
353) -> Result<(), EstimationError> {
354 const C_EPS_P_FACTOR: f64 = 64.0;
355 let p = eigenvalues.len();
360
361 let mut scale = 0.0_f64;
362 for (idx, &val) in eigenvalues.iter().enumerate() {
363 if !val.is_finite() {
364 return Err(EstimationError::PenaltySpectrumNonFinite {
365 context: context.to_string(),
366 index: idx,
367 value: val,
368 });
369 }
370 scale = scale.max(val.abs());
371 }
372
373 let machine_floor = C_EPS_P_FACTOR * f64::EPSILON * (p.max(1) as f64) * scale;
380 let tolerance = machine_floor
381 .max(REL_PSD_FLOOR * scale)
382 .max(f64::MIN_POSITIVE);
383
384 for (idx, val) in eigenvalues.iter_mut().enumerate() {
385 if val.abs() <= tolerance {
386 *val = 0.0;
387 } else if *val < 0.0 {
388 return Err(EstimationError::PenaltySpectrumIndefinite {
389 context: context.to_string(),
390 index: idx,
391 value: *val,
392 tolerance,
393 scale,
394 });
395 }
396 }
397 Ok(())
398}
399
400fn robust_eighwith_policy<M, V, E, Validate, Sanitize, EigCall, MapErr>(
401 matrix: &M,
402 context: &str,
403 validate_input: Validate,
404 sanitize: Sanitize,
405 mut eig_call: EigCall,
406 map_error: MapErr,
407) -> Result<(Vec<f64>, V), EstimationError>
408where
409 Validate: Fn(&M, &str) -> Result<(), EstimationError>,
410 Sanitize: Fn(&M) -> M,
411 EigCall: FnMut(&M) -> Result<(Vec<f64>, V), E>,
412 MapErr: Fn(E, &str) -> EstimationError,
413{
414 validate_input(matrix, context)?;
415
416 let candidate = sanitize(matrix);
422 match eig_call(&candidate) {
423 Ok((mut eigenvalues, eigenvectors)) => {
424 classify_eigenvalues_strict(&mut eigenvalues, context)?;
425 Ok((eigenvalues, eigenvectors))
426 }
427 Err(err) => Err(map_error(err, context)),
428 }
429}
430
431pub(crate) fn robust_eigh_faer(
432 matrix: &Mat<f64>,
433 side: Side,
434 context: &str,
435) -> Result<(Vec<f64>, Mat<f64>), EstimationError> {
436 robust_eighwith_policy(
437 matrix,
438 context,
439 |mat, ctx| {
440 let (rows, cols) = mat.as_ref().shape();
441 for i in 0..rows {
442 for j in 0..cols {
443 let val = mat[(i, j)];
444 if !val.is_finite() {
445 let max_abs = mat_max_abs_element(mat.as_ref());
446 crate::bail_invalid_estim!(
447 "{} contains non-finite entries (max finite magnitude {:.3e})",
448 ctx,
449 max_abs
450 );
451 }
452 }
453 }
454 Ok(())
455 },
456 sanitize_symmetric_faer,
457 |candidate| {
458 let eig = candidate.as_ref().self_adjoint_eigen(side)?;
459 let diag = eig.S();
460 let mut eigenvalues = Vec::with_capacity(diag.dim());
461 for idx in 0..diag.dim() {
462 eigenvalues.push(diag[idx]);
463 }
464
465 let vectors_ref = eig.U();
466 let mut eigenvectors = Mat::<f64>::zeros(vectors_ref.nrows(), vectors_ref.ncols());
467 for i in 0..vectors_ref.nrows() {
468 for j in 0..vectors_ref.ncols() {
469 eigenvectors[(i, j)] = vectors_ref[(i, j)];
470 }
471 }
472 Ok((eigenvalues, eigenvectors))
473 },
474 |err, _| {
475 EstimationError::EigendecompositionFailed(FaerLinalgError::SelfAdjointEigen(err))
476 },
477 )
478}
479
480fn robust_eigh(
481 matrix: &Array2<f64>,
482 side: Side,
483 context: &str,
484) -> Result<(Array1<f64>, Array2<f64>), EstimationError> {
485 let matrix_faer = array_to_faer(matrix);
486 let (eigenvalues, eigenvectors) = robust_eigh_faer(&matrix_faer, side, context)?;
487 Ok((Array1::from_vec(eigenvalues), mat_to_array(&eigenvectors)))
488}
489
490pub(crate) fn kronecker_marginal_eigensystems(
491 marginal_penalties: &[Array2<f64>],
492 context: &str,
493) -> Result<Vec<(Array1<f64>, Array2<f64>)>, EstimationError> {
494 let mut eigensystems = Vec::with_capacity(marginal_penalties.len());
495 for (k, penalty) in marginal_penalties.iter().enumerate() {
496 eigensystems.push(robust_eigh(
497 penalty,
498 Side::Lower,
499 &format!("{context} marginal {k}"),
500 )?);
501 }
502 Ok(eigensystems)
503}
504
505#[derive(Debug, Clone, Copy)]
506struct SubspaceLeakageMetrics {
507 max_abs_sq: f64,
508 max_rel_sq: f64,
509 worst_penalty: usize,
510 max_cross_gram_abs: f64,
511}
512
513fn assess_subspace_leakage(
514 qs: &Mat<f64>,
515 rs_transformed: &[Mat<f64>],
516 structural_rank: usize,
517 p: usize,
518) -> SubspaceLeakageMetrics {
519 let mut max_abs_sq = 0.0_f64;
520 let mut max_rel_sq = 0.0_f64;
521 let mut worst_penalty = 0usize;
522
523 for (k, rs) in rs_transformed.iter().enumerate() {
524 let rows = rs.nrows();
525 let cols = rs.ncols().min(p);
526 let null_start = structural_rank.min(cols);
527 let mut abs_sq = 0.0_f64;
528 let mut total_sq = 0.0_f64;
529 for i in 0..rows {
530 for j in 0..cols {
531 let v = rs[(i, j)];
532 let vv = v * v;
533 total_sq += vv;
534 if j >= null_start {
535 abs_sq += vv;
536 }
537 }
538 }
539 let rel_sq = if total_sq > 0.0 {
540 abs_sq / total_sq
541 } else {
542 0.0
543 };
544 if rel_sq > max_rel_sq {
545 max_rel_sq = rel_sq;
546 worst_penalty = k;
547 }
548 max_abs_sq = max_abs_sq.max(abs_sq);
549 }
550
551 let mut max_cross_gram_abs = 0.0_f64;
552 let null_count = p.saturating_sub(structural_rank);
553 if structural_rank > 0 && null_count > 0 {
554 for i in 0..structural_rank {
555 for j in 0..null_count {
556 let qn_col = structural_rank + j;
557 let mut dot = 0.0_f64;
558 for r in 0..p {
559 dot += qs[(r, i)] * qs[(r, qn_col)];
560 }
561 max_cross_gram_abs = max_cross_gram_abs.max(dot.abs());
562 }
563 }
564 }
565
566 SubspaceLeakageMetrics {
567 max_abs_sq,
568 max_rel_sq,
569 worst_penalty,
570 max_cross_gram_abs,
571 }
572}
573
574fn subspace_split_is_consistent(leakage: &SubspaceLeakageMetrics, p: usize) -> bool {
601 let leakage_rel_tol = (p.max(1) as f64) * REL_PSD_FLOOR;
602 let leakage_abs_tol = 1e-12;
603 let orth_tol = 1e-10;
604 let root_leaks = leakage.max_rel_sq > leakage_rel_tol && leakage.max_abs_sq > leakage_abs_tol;
605 let split_nonorthogonal = leakage.max_cross_gram_abs > orth_tol;
606 !(root_leaks || split_nonorthogonal)
607}
608
609fn compose_qs_from_split(q_pen: &Mat<f64>, q_null: &Mat<f64>, p: usize) -> Mat<f64> {
610 let rank = q_pen.ncols();
611 let null_count = q_null.ncols();
612 let mut qs = Mat::<f64>::zeros(p, p);
613 for i in 0..p {
614 for j in 0..rank {
615 qs[(i, j)] = q_pen[(i, j)];
616 }
617 for j in 0..null_count {
618 qs[(i, rank + j)] = q_null[(i, j)];
619 }
620 }
621 qs
622}
623
624pub fn kronecker_product(a: &Array2<f64>, b: &Array2<f64>) -> Array2<f64> {
628 let (arows, a_cols) = a.dim();
629 let (brows, b_cols) = b.dim();
630 if arows == 0 || a_cols == 0 || brows == 0 || b_cols == 0 {
631 return Array2::zeros((arows * brows, a_cols * b_cols));
632 }
633 let mut result = Array2::zeros((arows * brows, a_cols * b_cols));
634
635 result
636 .axis_chunks_iter_mut(Axis(0), brows)
637 .into_par_iter()
638 .enumerate()
639 .for_each(|(i, mut row_block)| {
640 let arow = a.row(i);
641 let col_chunks = row_block.axis_chunks_iter_mut(Axis(1), b_cols);
642 for (j, mut block) in col_chunks.into_iter().enumerate() {
643 let aval = arow[j];
644 if aval == 0.0 {
645 continue;
646 }
647 for (dest, &src) in block.iter_mut().zip(b.iter()) {
648 *dest = aval * src;
649 }
650 }
651 });
652
653 result
654}
655
656#[derive(Clone)]
658pub struct ReparamResult {
659 pub s_transformed: Array2<f64>,
663 pub log_det: f64,
665 pub det1: Array1<f64>,
667 pub qs: Array2<f64>,
669 pub canonical_transformed: Vec<CanonicalPenalty>,
674 pub e_transformed: Array2<f64>,
677 pub u_truncated: Array2<f64>,
687}
688
689#[derive(Clone, Copy, Debug, PartialEq, Eq)]
696pub enum PenaltyFrame {
697 Original,
698 Transformed,
699}
700
701#[derive(Clone, Debug)]
736pub struct PenaltyNullSplit {
737 transformed: Option<Array2<f64>>,
738 original: Option<Array2<f64>>,
739}
740
741impl PenaltyNullSplit {
742 pub fn identity() -> Self {
746 Self {
747 transformed: None,
748 original: None,
749 }
750 }
751
752 pub fn declared_null_dim(&self) -> usize {
754 self.transformed.as_ref().map_or(0, |n| n.ncols())
755 }
756
757 pub fn basis(&self, frame: PenaltyFrame) -> Option<&Array2<f64>> {
759 match frame {
760 PenaltyFrame::Original => self.original.as_ref(),
761 PenaltyFrame::Transformed => self.transformed.as_ref(),
762 }
763 }
764
765 pub fn project_canonical(
771 &self,
772 penalty: &CanonicalPenalty,
773 frame: PenaltyFrame,
774 ) -> Result<CanonicalPenalty, EstimationError> {
775 match self.basis(frame) {
776 Some(n) if n.nrows() == penalty.total_dim => {
777 penalty.project_out_null_directions(n.view())
778 }
779 _ => Ok(penalty.clone()),
780 }
781 }
782
783 pub fn project_coordinate(
785 &self,
786 coord: &gam_problem::PenaltyCoordinate,
787 frame: PenaltyFrame,
788 ) -> gam_problem::PenaltyCoordinate {
789 match self.basis(frame) {
790 Some(n) if n.nrows() == coord.dim() => coord.project_out_null_directions(n.view()),
791 _ => coord.clone(),
792 }
793 }
794}
795
796impl ReparamResult {
797 pub fn null_split(&self) -> PenaltyNullSplit {
804 let u = &self.u_truncated;
805 if u.ncols() == 0 || u.nrows() == 0 {
806 return PenaltyNullSplit::identity();
807 }
808 let qs = &self.qs;
809 let original = if qs.nrows() == u.nrows() && qs.ncols() == u.nrows() {
810 qs.dot(u)
811 } else {
812 u.clone()
816 };
817 PenaltyNullSplit {
818 transformed: Some(u.clone()),
819 original: Some(original),
820 }
821 }
822}
823
824struct KroneckerFactorDecomp {
830 root: Array2<f64>, positive_eigenvalues: Vec<f64>, rank: usize,
833 dim: usize,
834}
835
836fn decompose_kronecker_factors(
839 factors: &[Array2<f64>],
840 context: &str,
841) -> Result<Option<Vec<KroneckerFactorDecomp>>, EstimationError> {
842 let mut decomps = Vec::with_capacity(factors.len());
843 for (j, factor) in factors.iter().enumerate() {
844 let q_j = factor.nrows();
845 if q_j != factor.ncols() {
846 crate::bail_invalid_estim!(
847 "{context}: Kronecker factor {j} must be square, got {}x{}",
848 factor.nrows(),
849 factor.ncols()
850 );
851 }
852 let is_identity = {
853 let mut is_id = true;
854 'outer: for r in 0..q_j {
855 for c in 0..q_j {
856 let expected = if r == c { 1.0 } else { 0.0 };
857 if (factor[[r, c]] - expected).abs() > 1e-12 {
858 is_id = false;
859 break 'outer;
860 }
861 }
862 }
863 is_id
864 };
865 if is_identity {
866 decomps.push(KroneckerFactorDecomp {
867 root: Array2::eye(q_j),
868 positive_eigenvalues: vec![1.0; q_j],
869 rank: q_j,
870 dim: q_j,
871 });
872 continue;
873 }
874 let analysis = analyze_penalty_block(factor).map_err(|err| {
875 EstimationError::InvalidInput(format!(
876 "{context}: Kronecker factor {j} eigendecomp failed: {err}"
877 ))
878 })?;
879 if analysis.rank == 0 {
880 return Ok(None);
881 }
882 let factor_classes =
886 crate::basis::SpectralClassification::new(
887 &analysis.eigenvalues,
888 analysis.rank_tol,
889 analysis.noise_tol,
890 );
891 let mut root_j = Array2::zeros((analysis.rank, q_j));
892 let mut pos_eigs = Vec::with_capacity(analysis.rank);
893 for (row_idx, &i) in factor_classes.range_idx.iter().enumerate() {
894 let eigenval = analysis.eigenvalues[i];
895 let sqrt_ev = eigenval.sqrt();
896 let evec = analysis.eigenvectors.column(i);
897 for (col, &v) in evec.iter().enumerate() {
898 root_j[[row_idx, col]] = sqrt_ev * v;
899 }
900 pos_eigs.push(eigenval);
901 }
902 decomps.push(KroneckerFactorDecomp {
903 root: root_j,
904 positive_eigenvalues: pos_eigs,
905 rank: analysis.rank,
906 dim: q_j,
907 });
908 }
909 Ok(Some(decomps))
910}
911
912fn assemble_kronecker_root_local(decomps: &[KroneckerFactorDecomp]) -> Array2<f64> {
914 let mut kron_root = decomps[0].root.clone();
915 for fr in &decomps[1..] {
916 let (r1, c1) = kron_root.dim();
917 let (r2, c2) = (fr.rank, fr.dim);
918 let mut new_root = Array2::zeros((r1 * r2, c1 * c2));
919 for i1 in 0..r1 {
920 for i2 in 0..r2 {
921 for j1 in 0..c1 {
922 for j2 in 0..c2 {
923 new_root[[i1 * r2 + i2, j1 * c2 + j2]] =
924 kron_root[[i1, j1]] * fr.root[[i2, j2]];
925 }
926 }
927 }
928 }
929 kron_root = new_root;
930 }
931 kron_root
932}
933
934fn kronecker_eigenvalues(decomps: &[KroneckerFactorDecomp], block_dim: usize) -> (Vec<f64>, usize) {
936 let mut kron_eigs = decomps[0].positive_eigenvalues.clone();
937 for fd in &decomps[1..] {
938 let mut new_eigs = Vec::with_capacity(kron_eigs.len() * fd.positive_eigenvalues.len());
939 for &a in &kron_eigs {
940 for &b in &fd.positive_eigenvalues {
941 new_eigs.push(a * b);
942 }
943 }
944 kron_eigs = new_eigs;
945 }
946 let max_ev = kron_eigs.iter().copied().fold(0.0_f64, f64::max);
947 let tol = max_ev * 1e-10 * (block_dim as f64);
948 let positive: Vec<f64> = kron_eigs.into_iter().filter(|&ev| ev > tol).collect();
949 let nullity = block_dim - positive.len();
950 (positive, nullity)
951}
952
953#[derive(Clone)]
963pub struct CanonicalPenalty {
964 pub root: Array2<f64>,
967 pub col_range: std::ops::Range<usize>,
970 pub total_dim: usize,
972 pub nullity: usize,
974 pub local: Array2<f64>,
978 pub prior_mean: Array1<f64>,
980 pub positive_eigenvalues: Vec<f64>,
983 pub op: Option<std::sync::Arc<dyn crate::analytic_penalties::PenaltyOp>>,
987}
988
989impl std::fmt::Debug for CanonicalPenalty {
990 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
991 f.debug_struct("CanonicalPenalty")
992 .field(
993 "root",
994 &format_args!("{}×{}", self.root.nrows(), self.root.ncols()),
995 )
996 .field("col_range", &self.col_range)
997 .field("total_dim", &self.total_dim)
998 .field("nullity", &self.nullity)
999 .field(
1000 "local",
1001 &format_args!("{}×{}", self.local.nrows(), self.local.ncols()),
1002 )
1003 .field("prior_mean_len", &self.prior_mean.len())
1004 .field("positive_eigenvalues", &self.positive_eigenvalues)
1005 .field("op", &self.op.as_ref().map(|o| o.dim()))
1006 .finish()
1007 }
1008}
1009
1010impl CanonicalPenalty {
1011 pub fn from_dense_root(root: Array2<f64>, p: usize) -> Self {
1015 Self::from_dense_root_with_mean(root, p, Array1::zeros(p))
1016 }
1017
1018 pub fn from_dense_root_with_mean(root: Array2<f64>, p: usize, prior_mean: Array1<f64>) -> Self {
1019 assert_eq!(prior_mean.len(), p);
1020 let local = root.t().dot(&root);
1021 let positive_eigenvalues = Vec::new(); Self {
1023 root,
1024 col_range: 0..p,
1025 total_dim: p,
1026 nullity: 0,
1027 local,
1028 prior_mean,
1029 positive_eigenvalues,
1030 op: None,
1031 }
1032 }
1033
1034 pub fn full_width_root(&self) -> Array2<f64> {
1037 if self.col_range.start == 0 && self.col_range.end == self.total_dim {
1038 return self.root.clone();
1039 }
1040 let rank = self.root.nrows();
1041 let mut full = Array2::<f64>::zeros((rank, self.total_dim));
1042 full.slice_mut(ndarray::s![.., self.col_range.clone()])
1043 .assign(&self.root);
1044 full
1045 }
1046
1047 pub fn rank(&self) -> usize {
1049 self.root.nrows()
1050 }
1051
1052 pub fn block_dim(&self) -> usize {
1054 self.col_range.len()
1055 }
1056
1057 pub const fn is_block_local(&self) -> bool {
1059 self.col_range.start != 0 || self.col_range.end != self.total_dim
1060 }
1061
1062 pub fn local_ref(&self) -> &Array2<f64> {
1065 &self.local
1066 }
1067
1068 pub fn local_penalty(&self) -> Array2<f64> {
1071 self.local.clone()
1072 }
1073
1074 pub fn accumulate_weighted(&self, target: &mut Array2<f64>, lambda: f64) {
1077 if lambda == 0.0 || self.rank() == 0 {
1078 return;
1079 }
1080 let r = &self.col_range;
1081 target
1082 .slice_mut(s![r.start..r.end, r.start..r.end])
1083 .scaled_add(lambda, &self.local);
1084 }
1085
1086 pub fn trace_product(&self, m: &Array2<f64>, scale: f64) -> f64 {
1089 if self.rank() == 0 || scale == 0.0 {
1090 return 0.0;
1091 }
1092 let r = &self.col_range;
1093 let m_block = m.slice(s![r.start..r.end, r.start..r.end]);
1094 let rm = self.root.dot(&m_block);
1095 scale
1096 * rm.iter()
1097 .zip(self.root.iter())
1098 .map(|(&a, &b)| a * b)
1099 .sum::<f64>()
1100 }
1101
1102 pub fn quadratic(&self, v: &Array1<f64>, scale: f64) -> f64 {
1105 if self.rank() == 0 || scale == 0.0 {
1106 return 0.0;
1107 }
1108 let v_block = v.slice(s![self.col_range.start..self.col_range.end]);
1109 let rv = self.root.dot(&v_block);
1110 scale * rv.dot(&rv)
1111 }
1112
1113 pub fn prior_linear_shift(&self, scale: f64) -> Array1<f64> {
1115 let mut out = Array1::<f64>::zeros(self.total_dim);
1116 if self.rank() == 0 || scale == 0.0 || self.prior_mean.iter().all(|&v| v == 0.0) {
1117 return out;
1118 }
1119 let block = self.local.dot(&self.prior_mean) * scale;
1120 out.slice_mut(s![self.col_range.start..self.col_range.end])
1121 .assign(&block);
1122 out
1123 }
1124
1125 pub fn prior_constant_shift(&self, scale: f64) -> f64 {
1127 if self.rank() == 0 || scale == 0.0 || self.prior_mean.iter().all(|&v| v == 0.0) {
1128 return 0.0;
1129 }
1130 scale * self.prior_mean.dot(&self.local.dot(&self.prior_mean))
1131 }
1132
1133 pub fn full_width_prior_mean(&self) -> Array1<f64> {
1135 if self.col_range.start == 0 && self.col_range.end == self.total_dim {
1136 return self.prior_mean.clone();
1137 }
1138 let mut out = Array1::<f64>::zeros(self.total_dim);
1139 out.slice_mut(s![self.col_range.start..self.col_range.end])
1140 .assign(&self.prior_mean);
1141 out
1142 }
1143
1144 pub fn project_out_null_directions(
1170 &self,
1171 null_basis: ndarray::ArrayView2<'_, f64>,
1172 ) -> Result<Self, EstimationError> {
1173 if null_basis.ncols() == 0 || self.rank() == 0 {
1174 return Ok(self.clone());
1175 }
1176 if null_basis.nrows() != self.total_dim {
1177 return Err(EstimationError::LayoutError(format!(
1178 "CanonicalPenalty::project_out_null_directions: null-basis row count {} does \
1179 not match the penalty's total dimension {}",
1180 null_basis.nrows(),
1181 self.total_dim
1182 )));
1183 }
1184 let root_scale = self.root.iter().fold(0.0_f64, |a, v| a.max(v.abs()));
1185 let (projected_full, stays_block_local) =
1186 gam_problem::project_block_root_out_of_null_directions(
1187 self.root.view(),
1188 self.col_range.start,
1189 self.col_range.end,
1190 self.total_dim,
1191 null_basis,
1192 );
1193 let mut moved = 0.0_f64;
1196 {
1197 let original_full = self.full_width_root();
1198 for (a, b) in projected_full.iter().zip(original_full.iter()) {
1199 moved = moved.max((a - b).abs());
1200 }
1201 }
1202 if moved <= root_scale * f64::EPSILON * (self.total_dim as f64) {
1203 return Ok(self.clone());
1204 }
1205
1206 let (root, col_range) = if stays_block_local {
1207 (
1208 projected_full
1209 .slice(s![.., self.col_range.start..self.col_range.end])
1210 .to_owned(),
1211 self.col_range.clone(),
1212 )
1213 } else {
1214 (projected_full, 0..self.total_dim)
1215 };
1216 let prior_mean = if col_range == self.col_range {
1217 self.prior_mean.clone()
1218 } else {
1219 let mut full = Array1::<f64>::zeros(self.total_dim);
1220 full.slice_mut(s![self.col_range.start..self.col_range.end])
1221 .assign(&self.prior_mean);
1222 full
1223 };
1224 let positive_eigenvalues = match root.svd(false, false) {
1229 Ok((_, singular_values, _)) => {
1230 let mut eigs = vec![0.0_f64; root.nrows()];
1231 for (slot, &sigma) in eigs.iter_mut().zip(singular_values.iter()) {
1232 *slot = sigma * sigma;
1233 }
1234 eigs
1235 }
1236 Err(error) => {
1237 return Err(EstimationError::LayoutError(format!(
1238 "CanonicalPenalty::project_out_null_directions: SVD of the projected root \
1239 failed: {error:?}"
1240 )));
1241 }
1242 };
1243 let local = root.t().dot(&root);
1244 Ok(Self {
1245 root,
1246 col_range,
1247 total_dim: self.total_dim,
1248 nullity: self.nullity,
1251 local,
1252 prior_mean,
1253 positive_eigenvalues,
1254 op: None,
1255 })
1256 }
1257
1258 pub fn to_penalty_coordinate(&self) -> gam_problem::PenaltyCoordinate {
1260 use gam_problem::PenaltyCoordinate;
1261 if self.is_block_local() {
1262 PenaltyCoordinate::from_block_root_with_mean(
1263 self.root.clone(),
1264 self.col_range.start,
1265 self.col_range.end,
1266 self.total_dim,
1267 self.prior_mean.clone(),
1268 )
1269 } else {
1270 PenaltyCoordinate::from_dense_root_with_mean(self.root.clone(), self.prior_mean.clone())
1271 }
1272 }
1273}
1274
1275pub fn report_penalty_pair_redundancy(canonical: &[CanonicalPenalty]) -> Vec<(usize, usize, f64)> {
1362 const SIMILARITY_REPORTING_DEFECT: f64 = 1e-1;
1366 const LARGE_SCALE_K_THRESHOLD: usize = 64;
1367 const TOP_SIMILARITY_PAIRS: usize = 3;
1368
1369 let k = canonical.len();
1370 let mut redundant: Vec<(usize, usize, f64)> = Vec::new();
1371 let mut similar: Vec<(usize, usize, f64, f64)> = Vec::new();
1372
1373 let trace_sq: Vec<f64> = canonical
1376 .iter()
1377 .map(|p| p.local.iter().map(|&v| v * v).sum::<f64>())
1378 .collect();
1379
1380 for i in 0..k {
1381 if trace_sq[i] == 0.0 {
1382 continue;
1383 }
1384 for j in (i + 1)..k {
1385 if trace_sq[j] == 0.0 {
1386 continue;
1387 }
1388 if canonical[i].col_range != canonical[j].col_range {
1392 continue;
1393 }
1394 assert_eq!(canonical[i].local.dim(), canonical[j].local.dim());
1397
1398 let inner: f64 = canonical[i]
1399 .local
1400 .iter()
1401 .zip(canonical[j].local.iter())
1402 .map(|(&a, &b)| a * b)
1403 .sum();
1404 let scale = inner / trace_sq[i];
1410 let residual_sq: f64 = canonical[i]
1411 .local
1412 .iter()
1413 .zip(canonical[j].local.iter())
1414 .map(|(&a, &b)| {
1415 let residual = b - scale * a;
1416 residual * residual
1417 })
1418 .sum();
1419 let defect = (residual_sq / trace_sq[i]).sqrt();
1420 let entries = canonical[i].local.len() as f64;
1421 let arithmetic_floor = entries.sqrt() * f64::EPSILON;
1422
1423 if defect <= arithmetic_floor {
1424 redundant.push((i, j, defect));
1425 } else if defect <= SIMILARITY_REPORTING_DEFECT {
1426 similar.push((i, j, defect, scale));
1427 }
1428 }
1429 }
1430
1431 for &(i, j, defect) in &redundant {
1433 log::warn!(
1434 "[PENALTY-REDUNDANCY] penalties i={i} j={j} are proportional to the arithmetic \
1435 that formed them (relative defect min_c ||S_{j} - c S_{i}||_F / ||S_{i}||_F = \
1436 {defect:.6e}) — only their COMBINATION is identified, so the criterion is exactly \
1437 constant along their antisymmetric direction and lambda_{i} / lambda_{j} are not \
1438 separately estimable. The fit itself is unaffected (every point of that manifold \
1439 assembles the same penalty), and the curvature certificate deflates the direction \
1440 rather than judging a chain-rule term there (#2676). Consider re-specifying (e.g. \
1441 anisotropic→isotropic for spatial smoothers with weak axis signal) if you need the \
1442 individual smoothing parameters to mean something."
1443 );
1444 }
1445
1446 if k > LARGE_SCALE_K_THRESHOLD && similar.len() > TOP_SIMILARITY_PAIRS {
1448 similar.sort_by(|a, b| a.2.partial_cmp(&b.2).unwrap_or(std::cmp::Ordering::Equal));
1449 similar.truncate(TOP_SIMILARITY_PAIRS);
1450 }
1451 for (i, j, defect, scale) in similar {
1452 log::info!(
1453 "[PENALTY-SIMILARITY] penalties i={i} j={j} are close but MEASURABLY distinct \
1454 (relative defect {defect:.6e} at the best scale c={scale:.6e}) — the outer Hessian \
1455 may be ill-conditioned along their antisymmetric direction, and the criterion \
1456 carries genuine curvature of order defect^2 there. This is NOT an invariance: \
1457 nothing is exactly constant along it and no direction is deflated for it (#2676)."
1458 );
1459 }
1460
1461 redundant
1462}
1463
1464pub fn canonicalize_penalty_spec(
1470 spec: &crate::PenaltySpec,
1471 p: usize,
1472 idx: usize,
1473 context: &str,
1474) -> Result<Option<CanonicalPenalty>, EstimationError> {
1475 use crate::PenaltySpec;
1476
1477 crate::validate_penalty_spec_shape(idx, spec, p, context)?;
1478
1479 let (local_matrix, col_range, prior_mean_spec, hint, op) = match spec {
1480 PenaltySpec::Block {
1481 local,
1482 col_range,
1483 prior_mean,
1484 structure_hint,
1485 op,
1486 } => (
1487 local.view(),
1488 col_range.clone(),
1489 prior_mean,
1490 structure_hint.as_ref(),
1491 op.clone(),
1492 ),
1493 PenaltySpec::Dense(m) => (
1494 m.view(),
1495 0..p,
1496 &gam_problem::CoefficientPriorMean::Zero,
1497 None,
1498 None,
1499 ),
1500 PenaltySpec::DenseWithMean { matrix, prior_mean } => {
1501 (matrix.view(), 0..p, prior_mean, None, None)
1502 }
1503 };
1504
1505 let block_dim = col_range.len();
1506 let prior_mean = prior_mean_spec
1507 .evaluate(block_dim, &format!("{context}: penalty {idx}"))
1508 .map_err(|e| EstimationError::InvalidInput(e.0))?;
1509
1510 if let Some(PenaltyStructureHint::Ridge(scale)) = hint {
1512 if *scale <= 0.0 {
1513 return Ok(None);
1514 }
1515 let sqrt_scale = scale.sqrt();
1516 let mut root = Array2::zeros((block_dim, block_dim));
1517 for i in 0..block_dim {
1518 root[[i, i]] = sqrt_scale;
1519 }
1520 let mut local_sym = local_matrix.to_owned();
1524 symmetrize_in_place(&mut local_sym);
1525 return Ok(Some(CanonicalPenalty {
1526 root,
1527 col_range,
1528 total_dim: p,
1529 nullity: 0,
1530 local: local_sym,
1531 prior_mean,
1532 positive_eigenvalues: vec![*scale; block_dim],
1533 op,
1534 }));
1535 }
1536
1537 if let Some(PenaltyStructureHint::Kronecker(factors)) = hint {
1539 let decomps =
1540 match decompose_kronecker_factors(factors, &format!("{context} penalty {idx}"))? {
1541 None => return Ok(None),
1542 Some(d) => d,
1543 };
1544 let (positive_eigenvalues, nullity) = kronecker_eigenvalues(&decomps, block_dim);
1545 if positive_eigenvalues.is_empty() {
1546 return Ok(None);
1547 }
1548 let root = assemble_kronecker_root_local(&decomps);
1549 let mut local_sym = local_matrix.to_owned();
1550 symmetrize_in_place(&mut local_sym);
1551 return Ok(Some(CanonicalPenalty {
1552 root,
1553 col_range,
1554 total_dim: p,
1555 nullity,
1556 local: local_sym,
1557 prior_mean,
1558 positive_eigenvalues,
1559 op,
1560 }));
1561 }
1562
1563 let local_owned = local_matrix.to_owned();
1565 let analysis = analyze_penalty_block(&local_owned).map_err(|err| {
1566 EstimationError::InvalidInput(format!(
1567 "{context}: penalty canonicalization failed at index {idx}: {err}"
1568 ))
1569 })?;
1570
1571 if analysis.rank == 0 {
1572 log::debug!(
1573 "Dropped inactive penalty block idx={idx} reason={}",
1574 if analysis.iszero {
1575 "ZeroMatrix"
1576 } else {
1577 "NumericalRankZero"
1578 }
1579 );
1580 return Ok(None);
1581 }
1582
1583 let tolerance = analysis.rank_tol;
1589 let classes = crate::basis::SpectralClassification::new(
1590 &analysis.eigenvalues,
1591 tolerance,
1592 analysis.noise_tol,
1593 );
1594 let rank_k = classes.rank();
1595 assert_eq!(
1596 rank_k, analysis.rank,
1597 "penalty-root rank disagreement: SpectralClassification rank={rank_k} vs analyze_penalty_block rank={} (#1425 canonical-classifier invariant)",
1598 analysis.rank
1599 );
1600
1601 let mut root = Array2::zeros((rank_k, block_dim));
1609 let mut positive_eigenvalues = Vec::with_capacity(rank_k);
1610 for (row_idx, &i) in classes.range_idx.iter().enumerate() {
1611 let eigenval = analysis.eigenvalues[i];
1612 let eigenvec = analysis.eigenvectors.column(i);
1613 root.row_mut(row_idx).assign(&(&eigenvec * eigenval.sqrt()));
1614 positive_eigenvalues.push(eigenval);
1615 }
1616
1617 if classes.is_indefinite() {
1623 log::debug!(
1624 "{context}: penalty block idx={idx} carries {} negative-curvature \
1625 eigendirection(s) below -tol={tolerance:e}; dropped from the canonical \
1626 root and NOT counted as null space (rank={rank_k}, nullity={})",
1627 classes.negative_dim(),
1628 classes.nullity()
1629 );
1630 }
1631
1632 let local = root.t().dot(&root);
1636 Ok(Some(CanonicalPenalty {
1637 root,
1638 col_range,
1639 total_dim: p,
1640 nullity: classes.nullity(),
1641 local,
1642 prior_mean,
1643 positive_eigenvalues,
1644 op,
1645 }))
1646}
1647
1648pub fn canonicalize_penalty_specs(
1651 specs: &[crate::PenaltySpec],
1652 nullspace_dims: &[usize],
1653 p: usize,
1654 context: &str,
1655) -> Result<(Vec<CanonicalPenalty>, Vec<usize>), EstimationError> {
1656 if specs.len() != nullspace_dims.len() {
1657 crate::bail_invalid_estim!(
1658 "{context}: nullspace_dims length mismatch: penalties={}, nullspace_dims={}",
1659 specs.len(),
1660 nullspace_dims.len()
1661 );
1662 }
1663
1664 let mut active = Vec::with_capacity(specs.len());
1665 let mut active_nullspace = Vec::with_capacity(specs.len());
1666 for (idx, spec) in specs.iter().enumerate() {
1667 if let Some(canonical) = canonicalize_penalty_spec(spec, p, idx, context)? {
1668 active_nullspace.push(nullspace_dims[idx]);
1669 active.push(canonical);
1670 }
1671 }
1672 Ok((active, active_nullspace))
1673}
1674
1675pub(crate) const OVERLAPPING_PENALTY_DENSE_FALLBACK_MAX_P: usize = 4096;
1685
1686pub fn create_balanced_penalty_root_from_canonical(
1693 penalties: &[CanonicalPenalty],
1694 p: usize,
1695) -> Result<Array2<f64>, EstimationError> {
1696 if penalties.is_empty() {
1697 return Ok(Array2::zeros((0, p)));
1698 }
1699
1700 let mut block_groups: BTreeMap<(usize, usize), Vec<&CanonicalPenalty>> = BTreeMap::new();
1702 for cp in penalties {
1703 if cp.rank() == 0 {
1704 continue;
1705 }
1706 let key = (cp.col_range.start, cp.col_range.end);
1707 block_groups.entry(key).or_default().push(cp);
1708 }
1709
1710 if block_groups.is_empty() {
1711 return Ok(Array2::zeros((0, p)));
1712 }
1713
1714 let ranges: Vec<(usize, usize)> = block_groups.keys().copied().collect();
1716 let mut overlapping = false;
1717 for i in 1..ranges.len() {
1718 if ranges[i].0 < ranges[i - 1].1 {
1719 overlapping = true;
1720 break;
1721 }
1722 }
1723
1724 if overlapping {
1725 if p > OVERLAPPING_PENALTY_DENSE_FALLBACK_MAX_P {
1726 return Err(EstimationError::LayoutError(format!(
1727 "overlapping penalty root would require dense {}x{} eigendecomposition; \
1728 large-model dense fallback is disabled. Keep penalties structured or \
1729 extend the overlapping-penalty solver path",
1730 p, p
1731 )));
1732 }
1733 let mut s_balanced = Array2::zeros((p, p));
1735 for cp in penalties {
1736 if cp.rank() == 0 {
1737 continue;
1738 }
1739 let local = cp.local_ref();
1740 let frob_norm = local.iter().map(|&x| x * x).sum::<f64>().sqrt();
1741 if frob_norm > 1e-12 {
1742 let r = &cp.col_range;
1743 s_balanced
1744 .slice_mut(s![r.start..r.end, r.start..r.end])
1745 .scaled_add(1.0 / frob_norm, local);
1746 }
1747 }
1748 let (eigenvalues, eigenvectors) =
1749 robust_eigh(&s_balanced, Side::Lower, "balanced penalty matrix")?;
1750 let max_eig = eigenvalues.iter().fold(0.0f64, |max, &val| max.max(val));
1751 let tolerance = if max_eig > 0.0 {
1752 max_eig * 1e-12
1753 } else {
1754 1e-12
1755 };
1756 let penalty_rank = eigenvalues.iter().filter(|&&ev| ev > tolerance).count();
1757 if penalty_rank == 0 {
1758 return Ok(Array2::zeros((0, p)));
1759 }
1760 let mut eb = Array2::zeros((p, penalty_rank));
1761 let mut col_idx = 0;
1762 for (i, &eigenval) in eigenvalues.iter().enumerate() {
1763 if eigenval > tolerance {
1764 let sqrt_ev = eigenval.sqrt();
1765 let evec = eigenvectors.column(i);
1766 eb.column_mut(col_idx).assign(&(&evec * sqrt_ev));
1767 col_idx += 1;
1768 }
1769 }
1770 return Ok(eb.t().to_owned());
1771 }
1772
1773 struct BlockRoot {
1775 col_range: Range<usize>,
1776 root: Array2<f64>, }
1778 let ordered_blocks: Vec<((usize, usize), Vec<&CanonicalPenalty>)> =
1783 block_groups.into_iter().collect();
1784 let block_roots: Vec<BlockRoot> = ordered_blocks
1785 .into_par_iter()
1786 .map(
1787 |((start, end), cps)| -> Result<Option<BlockRoot>, EstimationError> {
1788 let block_dim = end - start;
1789 let mut s_balanced_local = Array2::zeros((block_dim, block_dim));
1790
1791 for cp in cps {
1792 let local = cp.local_ref();
1793 let frob_norm = local.iter().map(|&x| x * x).sum::<f64>().sqrt();
1794 if frob_norm > 1e-12 {
1795 s_balanced_local.scaled_add(1.0 / frob_norm, local);
1796 }
1797 }
1798
1799 let (eigenvalues, eigenvectors) =
1800 robust_eigh(&s_balanced_local, Side::Lower, "balanced penalty block")?;
1801 let max_eig = eigenvalues.iter().fold(0.0f64, |max, &val| max.max(val));
1802 let tolerance = if max_eig > 0.0 {
1803 max_eig * 1e-12
1804 } else {
1805 1e-12
1806 };
1807 let block_rank = eigenvalues.iter().filter(|&&ev| ev > tolerance).count();
1808
1809 if block_rank == 0 {
1810 return Ok(None);
1811 }
1812
1813 let mut root = Array2::zeros((block_rank, block_dim));
1814 let mut row_idx = 0;
1815 for (i, &eigenval) in eigenvalues.iter().enumerate() {
1816 if eigenval > tolerance {
1817 let sqrt_ev = eigenval.sqrt();
1818 let evec = eigenvectors.column(i);
1819 root.row_mut(row_idx).assign(&(&evec * sqrt_ev));
1820 row_idx += 1;
1821 }
1822 }
1823
1824 Ok(Some(BlockRoot {
1825 col_range: start..end,
1826 root,
1827 }))
1828 },
1829 )
1830 .collect::<Result<Vec<_>, _>>()?
1831 .into_iter()
1832 .flatten()
1833 .collect();
1834 let total_rank: usize = block_roots.iter().map(|br| br.root.nrows()).sum();
1835
1836 if total_rank == 0 {
1837 return Ok(Array2::zeros((0, p)));
1838 }
1839
1840 let mut eb = Array2::zeros((total_rank, p));
1842 let mut row_offset = 0;
1843 for br in &block_roots {
1844 let rank_b = br.root.nrows();
1845 eb.slice_mut(s![
1846 row_offset..(row_offset + rank_b),
1847 br.col_range.start..br.col_range.end
1848 ])
1849 .assign(&br.root);
1850 row_offset += rank_b;
1851 }
1852
1853 Ok(eb)
1854}
1855
1856#[derive(Clone)]
1858struct SubspaceSplit {
1859 q_pen: Array2<f64>,
1860 q_null: Array2<f64>,
1861}
1862
1863impl SubspaceSplit {
1864 fn identity(p: usize) -> Self {
1865 Self {
1866 q_pen: Array2::zeros((p, 0)),
1867 q_null: Array2::eye(p),
1868 }
1869 }
1870
1871 fn from_ordered_qs(
1872 qs: &Mat<f64>,
1873 penalized_rank: usize,
1874 p: usize,
1875 ) -> Result<Self, EstimationError> {
1876 if qs.nrows() != p || qs.ncols() != p {
1877 return Err(EstimationError::LayoutError(format!(
1878 "Invalid Q basis dimensions: expected {p}x{p}, got {}x{}",
1879 qs.nrows(),
1880 qs.ncols()
1881 )));
1882 }
1883 if penalized_rank > p {
1884 return Err(EstimationError::LayoutError(format!(
1885 "Invalid penalized rank {penalized_rank} for p={p}"
1886 )));
1887 }
1888
1889 let null_count = p - penalized_rank;
1890 let mut q_pen = Array2::<f64>::zeros((p, penalized_rank));
1891 let mut q_null = Array2::<f64>::zeros((p, null_count));
1892 for i in 0..p {
1893 for j in 0..penalized_rank {
1894 q_pen[(i, j)] = qs[(i, j)];
1895 }
1896 for j in 0..null_count {
1897 q_null[(i, j)] = qs[(i, penalized_rank + j)];
1898 }
1899 }
1900
1901 Ok(Self { q_pen, q_null })
1902 }
1903
1904 fn rank(&self) -> usize {
1905 self.q_pen.ncols()
1906 }
1907
1908 fn p(&self) -> usize {
1909 self.q_pen.nrows()
1910 }
1911
1912 fn compose_qs(&self) -> Array2<f64> {
1913 let p = self.p();
1914 let rank = self.rank();
1915 let null_count = self.q_null.ncols();
1916 let mut qs = Array2::<f64>::zeros((p, p));
1917 for i in 0..p {
1918 for j in 0..rank {
1919 qs[(i, j)] = self.q_pen[(i, j)];
1920 }
1921 for j in 0..null_count {
1922 qs[(i, rank + j)] = self.q_null[(i, j)];
1923 }
1924 }
1925 qs
1926 }
1927}
1928
1929#[derive(Clone)]
1931pub struct ReparamInvariant {
1932 split: SubspaceSplit,
1933 qs_base: Array2<f64>,
1937 has_nonzero: bool,
1938 max_balanced_eigenvalue: f64,
1941}
1942
1943impl ReparamInvariant {
1944 pub const fn max_balanced_eigenvalue(&self) -> f64 {
1947 self.max_balanced_eigenvalue
1948 }
1949}
1950
1951pub fn precompute_reparam_invariant_from_canonical(
1958 penalties: &[CanonicalPenalty],
1959 p_total: usize,
1960) -> Result<ReparamInvariant, EstimationError> {
1961 use std::cmp::Ordering;
1962
1963 let m = penalties.len();
1964
1965 if m == 0 {
1966 return Ok(ReparamInvariant {
1967 split: SubspaceSplit::identity(p_total),
1968 qs_base: Array2::eye(p_total),
1969 has_nonzero: false,
1970 max_balanced_eigenvalue: 0.0,
1971 });
1972 }
1973
1974 struct PenRef {
1976 penalty_index: usize,
1977 }
1978 let mut block_groups: BTreeMap<(usize, usize), Vec<PenRef>> = BTreeMap::new();
1979 let mut has_nonzero = false;
1980 for (i, cp) in penalties.iter().enumerate() {
1981 if cp.rank() == 0 {
1982 continue;
1983 }
1984 let local = cp.local_ref();
1985 let frob_norm = local.iter().map(|&x| x * x).sum::<f64>().sqrt();
1986 if frob_norm > 1e-12 {
1987 has_nonzero = true;
1988 }
1989 let key = (cp.col_range.start, cp.col_range.end);
1990 block_groups
1991 .entry(key)
1992 .or_default()
1993 .push(PenRef { penalty_index: i });
1994 }
1995
1996 if !has_nonzero {
1997 return Ok(ReparamInvariant {
1998 split: SubspaceSplit::identity(p_total),
1999 qs_base: Array2::eye(p_total),
2000 has_nonzero: false,
2001 max_balanced_eigenvalue: 0.0,
2002 });
2003 }
2004
2005 let ranges: Vec<(usize, usize)> = block_groups.keys().copied().collect();
2007 let mut overlapping = false;
2008 for i in 1..ranges.len() {
2009 if ranges[i].0 < ranges[i - 1].1 {
2010 overlapping = true;
2011 break;
2012 }
2013 }
2014
2015 if overlapping {
2016 if p_total > OVERLAPPING_PENALTY_DENSE_FALLBACK_MAX_P {
2022 return Err(EstimationError::LayoutError(format!(
2023 "overlapping penalty reparameterization would require dense {}x{} eigendecomposition; \
2024 large-model dense fallback is disabled. Keep penalties structured or \
2025 extend the overlapping-penalty solver path",
2026 p_total, p_total
2027 )));
2028 }
2029 let mut s_balanced = Mat::<f64>::zeros(p_total, p_total);
2031 for cp in penalties {
2032 if cp.rank() == 0 {
2033 continue;
2034 }
2035 let local = cp.local_ref();
2036 let frob_norm = local.iter().map(|&x| x * x).sum::<f64>().sqrt();
2037 if frob_norm > 1e-12 {
2038 let scale = 1.0 / frob_norm;
2039 let r = &cp.col_range;
2040 for i in 0..local.nrows() {
2041 for j in 0..local.ncols() {
2042 s_balanced[(r.start + i, r.start + j)] += scale * local[[i, j]];
2043 }
2044 }
2045 }
2046 }
2047
2048 let (bal_eigenvalues, bal_eigenvectors) =
2049 robust_eigh_faer(&s_balanced, Side::Lower, "balanced penalty matrix")?;
2050
2051 let mut order: Vec<usize> = (0..p_total).collect();
2052 order.sort_by(|&i, &j| {
2053 bal_eigenvalues[j]
2054 .partial_cmp(&bal_eigenvalues[i])
2055 .unwrap_or(Ordering::Equal)
2056 .then(i.cmp(&j))
2057 });
2058
2059 let mut qs = Mat::<f64>::zeros(p_total, p_total);
2060 for (col_idx, &idx) in order.iter().enumerate() {
2061 for row in 0..p_total {
2062 qs[(row, col_idx)] = bal_eigenvectors[(row, idx)];
2063 }
2064 }
2065
2066 let max_bal = order
2067 .iter()
2068 .map(|&idx| bal_eigenvalues[idx].abs())
2069 .fold(0.0_f64, f64::max);
2070 let rank_tol = if max_bal > 0.0 {
2071 max_bal * 1e-12
2072 } else {
2073 1e-12
2074 };
2075 let penalized_rank = order
2076 .iter()
2077 .take_while(|&&idx| bal_eigenvalues[idx] > rank_tol)
2078 .count();
2079 let split = SubspaceSplit::from_ordered_qs(&qs, penalized_rank, p_total)?;
2080
2081 return Ok(ReparamInvariant {
2082 split,
2083 qs_base: mat_to_array(&qs),
2084 has_nonzero,
2085 max_balanced_eigenvalue: max_bal,
2086 });
2087 }
2088
2089 let mut covered = vec![false; p_total];
2097 for cp in penalties {
2098 for j in cp.col_range.clone() {
2099 covered[j] = true;
2100 }
2101 }
2102 let uncovered_cols: Vec<usize> = (0..p_total).filter(|j| !covered[*j]).collect();
2103
2104 struct BlockResult {
2105 col_range: Range<usize>,
2106 q_pen_local: Array2<f64>, q_null_local: Array2<f64>, max_balanced_eigenvalue: f64,
2110 pen_col_offset: usize,
2112 null_col_offset: usize,
2114 }
2115
2116 let block_specs: Vec<_> = block_groups.iter().collect();
2120 let mut block_results: Vec<BlockResult> = block_specs
2121 .into_par_iter()
2122 .map(
2123 |(&(start, end), refs)| -> Result<BlockResult, EstimationError> {
2124 let block_dim = end - start;
2125
2126 let mut s_balanced_local = Array2::zeros((block_dim, block_dim));
2128 let mut block_has_nonzero = false;
2129 for pref in refs {
2130 let cp = &penalties[pref.penalty_index];
2131 let local = cp.local_ref();
2132 let frob_norm = local.iter().map(|&x| x * x).sum::<f64>().sqrt();
2133 if frob_norm > 1e-12 {
2134 s_balanced_local.scaled_add(1.0 / frob_norm, local);
2135 block_has_nonzero = true;
2136 }
2137 }
2138
2139 if !block_has_nonzero {
2140 return Ok(BlockResult {
2141 col_range: start..end,
2142 q_pen_local: Array2::zeros((block_dim, 0)),
2143 q_null_local: Array2::eye(block_dim),
2144 max_balanced_eigenvalue: 0.0,
2145 pen_col_offset: 0, null_col_offset: 0, });
2148 }
2149
2150 let (bal_eigenvalues, bal_eigenvectors) =
2152 robust_eigh(&s_balanced_local, Side::Lower, "balanced penalty block")?;
2153
2154 let mut order: Vec<usize> = (0..block_dim).collect();
2155 order.sort_by(|&i, &j| {
2156 bal_eigenvalues[j]
2157 .partial_cmp(&bal_eigenvalues[i])
2158 .unwrap_or(Ordering::Equal)
2159 .then(i.cmp(&j))
2160 });
2161
2162 let max_bal = order
2163 .iter()
2164 .map(|&idx| bal_eigenvalues[idx].abs())
2165 .fold(0.0_f64, f64::max);
2166 let rank_tol = if max_bal > 0.0 {
2167 max_bal * 1e-12
2168 } else {
2169 1e-12
2170 };
2171 let penalized_rank = order
2172 .iter()
2173 .take_while(|&&idx| bal_eigenvalues[idx] > rank_tol)
2174 .count();
2175 let null_count = block_dim - penalized_rank;
2176
2177 let mut q_pen_local = Array2::zeros((block_dim, penalized_rank));
2178 let mut q_null_local = Array2::zeros((block_dim, null_count));
2179 for (col_idx, &idx) in order.iter().enumerate() {
2180 if col_idx < penalized_rank {
2181 for row in 0..block_dim {
2182 q_pen_local[[row, col_idx]] = bal_eigenvectors[[row, idx]];
2183 }
2184 } else {
2185 let null_col = col_idx - penalized_rank;
2186 for row in 0..block_dim {
2187 q_null_local[[row, null_col]] = bal_eigenvectors[[row, idx]];
2188 }
2189 }
2190 }
2191
2192 Ok(BlockResult {
2193 col_range: start..end,
2194 q_pen_local,
2195 q_null_local,
2196 max_balanced_eigenvalue: max_bal,
2197 pen_col_offset: 0, null_col_offset: 0, })
2200 },
2201 )
2202 .collect::<Result<_, _>>()?;
2203 let global_max_bal = block_results
2204 .iter()
2205 .map(|br| br.max_balanced_eigenvalue)
2206 .fold(0.0_f64, f64::max);
2207
2208 let total_pen_rank: usize = block_results.iter().map(|br| br.q_pen_local.ncols()).sum();
2210 let total_null: usize = block_results
2211 .iter()
2212 .map(|br| br.q_null_local.ncols())
2213 .sum::<usize>()
2214 + uncovered_cols.len();
2215 {
2216 let mut pen_off = 0usize;
2217 let mut null_off = 0usize;
2218 for br in &mut block_results {
2219 br.pen_col_offset = pen_off;
2220 br.null_col_offset = null_off;
2221 pen_off += br.q_pen_local.ncols();
2222 null_off += br.q_null_local.ncols();
2223 }
2224 }
2225
2226 let mut q_pen = Array2::zeros((p_total, total_pen_rank));
2227 let mut q_null = Array2::zeros((p_total, total_null));
2228
2229 for br in &block_results {
2230 let start = br.col_range.start;
2231 let bd = br.q_pen_local.nrows();
2232 let pen_r = br.q_pen_local.ncols();
2233 let null_r = br.q_null_local.ncols();
2234 if pen_r > 0 {
2235 q_pen
2236 .slice_mut(s![
2237 start..(start + bd),
2238 br.pen_col_offset..(br.pen_col_offset + pen_r)
2239 ])
2240 .assign(&br.q_pen_local);
2241 }
2242 if null_r > 0 {
2243 q_null
2244 .slice_mut(s![
2245 start..(start + bd),
2246 br.null_col_offset..(br.null_col_offset + null_r)
2247 ])
2248 .assign(&br.q_null_local);
2249 }
2250 }
2251 let mut null_col = block_results
2252 .iter()
2253 .map(|br| br.q_null_local.ncols())
2254 .sum::<usize>();
2255 for &j in &uncovered_cols {
2256 q_null[[j, null_col]] = 1.0;
2257 null_col += 1;
2258 }
2259
2260 let split = SubspaceSplit { q_pen, q_null };
2261
2262 let qs_global = split.compose_qs();
2266
2267 Ok(ReparamInvariant {
2268 split,
2269 qs_base: qs_global,
2270 has_nonzero,
2271 max_balanced_eigenvalue: global_max_bal,
2272 })
2273}
2274
2275fn absorb_right_singular_factorization(
2286 singular_values: &Array1<f64>,
2287 vt: &Array2<f64>,
2288 penalized_rank: usize,
2289 range_eigenvalues_sorted: &mut Vec<f64>,
2290 range_rotation: &mut Mat<f64>,
2291) {
2292 let n_sv = singular_values.len().min(penalized_rank).min(vt.nrows());
2295 *range_eigenvalues_sorted = (0..penalized_rank)
2296 .map(|i| {
2297 if i < n_sv {
2298 let s = singular_values[i];
2299 s * s
2300 } else {
2301 0.0
2302 }
2303 })
2304 .collect();
2305 for col_idx in 0..n_sv {
2306 for row in 0..penalized_rank {
2307 range_rotation[(row, col_idx)] = vt[[col_idx, row]];
2308 }
2309 }
2310}
2311
2312fn describe_stacked_roots(e_stacked: &Array2<f64>, lambdas: &[f64]) -> String {
2318 let mut max_abs = 0.0_f64;
2319 let mut nonfinite = 0usize;
2320 for &value in e_stacked.iter() {
2321 if value.is_finite() {
2322 max_abs = max_abs.max(value.abs());
2323 } else {
2324 nonfinite += 1;
2325 }
2326 }
2327 let finite_lambdas: Vec<f64> = lambdas.iter().copied().filter(|l| l.is_finite()).collect();
2328 let lambda_min = finite_lambdas.iter().copied().fold(f64::INFINITY, f64::min);
2329 let lambda_max = finite_lambdas
2330 .iter()
2331 .copied()
2332 .fold(f64::NEG_INFINITY, f64::max);
2333 format!(
2334 "rows={} cols={} nonfinite_entries={} max_abs={:.6e} n_lambda={} lambda_min={:.6e} \
2335 lambda_max={:.6e}",
2336 e_stacked.nrows(),
2337 e_stacked.ncols(),
2338 nonfinite,
2339 max_abs,
2340 lambdas.len(),
2341 lambda_min,
2342 lambda_max,
2343 )
2344}
2345
2346pub fn stable_reparameterizationwith_invariant(
2347 penalties: &[CanonicalPenalty],
2348 lambdas: &[f64],
2349 p: usize,
2350 invariant: &ReparamInvariant,
2351) -> Result<ReparamResult, EstimationError> {
2352 let m = penalties.len();
2353
2354 if lambdas.len() != m {
2355 return Err(EstimationError::ParameterConstraintViolation(format!(
2356 "Lambda count mismatch: expected {} lambdas for {} penalties, got {}",
2357 m,
2358 m,
2359 lambdas.len()
2360 )));
2361 }
2362
2363 if m == 0 {
2375 return Ok(ReparamResult {
2376 s_transformed: Array2::zeros((p, p)),
2377 log_det: 0.0,
2378 det1: Array1::zeros(0),
2379 qs: Array2::eye(p),
2380 canonical_transformed: vec![],
2381 e_transformed: Array2::zeros((0, p)),
2382 u_truncated: Array2::eye(p),
2384 });
2385 }
2386
2387 if !invariant.has_nonzero {
2388 let qs = invariant.split.compose_qs();
2389 let u_truncated = qs.t().dot(&invariant.split.q_null);
2390 let canonical_transformed: Vec<CanonicalPenalty> = penalties.to_vec();
2392 return Ok(ReparamResult {
2393 s_transformed: Array2::zeros((p, p)),
2394 log_det: 0.0,
2395 det1: Array1::zeros(m),
2396 qs,
2397 canonical_transformed,
2398 e_transformed: Array2::zeros((0, p)),
2399 u_truncated,
2400 });
2401 }
2402
2403 let q_pen = array_to_faer(&invariant.split.q_pen);
2404 let q_null = array_to_faer(&invariant.split.q_null);
2405 let qs_base = array_to_faer(&invariant.qs_base);
2406 let penalty_transforms: Vec<(Mat<f64>, Mat<f64>)> = penalties
2411 .par_iter()
2412 .map(|cp| {
2413 let r = &cp.col_range;
2414 let root_faer = array_to_faer(&cp.root);
2415 let q_block = qs_base.submatrix(r.start, 0, cp.block_dim(), p);
2416 let mut product = Mat::<f64>::zeros(cp.rank(), p);
2417 matmul(
2418 product.as_mut(),
2419 Accum::Replace,
2420 root_faer.as_ref(),
2421 q_block,
2422 1.0,
2423 Par::Seq,
2424 );
2425 let s_k = penalty_from_root_faer(&product);
2426 (product, s_k)
2427 })
2428 .collect();
2429 let (rs_transformed, s_k_penalized_cache): (Vec<Mat<f64>>, Vec<Mat<f64>>) =
2430 penalty_transforms.into_iter().unzip();
2431
2432 let penalized_rank = invariant.split.rank();
2433
2434 let mut range_eigenvalues_sorted: Vec<f64> = Vec::new();
2435 let mut range_rotation = Mat::<f64>::zeros(penalized_rank, penalized_rank);
2436 if penalized_rank > 0 {
2437 let total_root_rows: usize = rs_transformed.iter().map(Mat::nrows).sum();
2470 let mut have_rotation = false;
2496 let mut rescued_by_r_svd = false;
2497 let mut svd_refusal: Option<String> = None;
2498 if total_root_rows >= penalized_rank {
2499 let mut e_stacked = Array2::<f64>::zeros((total_root_rows, penalized_rank));
2500 let mut row_off = 0usize;
2501 for (lambda, root) in lambdas.iter().zip(rs_transformed.iter()) {
2502 let sqrt_lambda = lambda.max(0.0).sqrt();
2503 let rk = root.nrows();
2504 for r in 0..rk {
2505 for c in 0..penalized_rank {
2506 e_stacked[[row_off + r, c]] = sqrt_lambda * root[(r, c)];
2507 }
2508 }
2509 row_off += rk;
2510 }
2511 match e_stacked.svd(false, true) {
2512 Ok((_, singular_values, Some(vt))) => {
2513 absorb_right_singular_factorization(
2514 &singular_values,
2515 &vt,
2516 penalized_rank,
2517 &mut range_eigenvalues_sorted,
2518 &mut range_rotation,
2519 );
2520 have_rotation = true;
2521 }
2522 direct => {
2523 let facts = describe_stacked_roots(&e_stacked, lambdas);
2524 svd_refusal = Some(match direct {
2525 Err(err) => format!("failed: {err:?}; {facts}"),
2526 _ => format!("returned no right singular vectors; {facts}"),
2527 });
2528 if let Ok((_, r_factor)) = e_stacked.qr()
2544 && let Ok((_, singular_values, Some(vt))) = r_factor.svd(false, true)
2545 {
2546 absorb_right_singular_factorization(
2547 &singular_values,
2548 &vt,
2549 penalized_rank,
2550 &mut range_eigenvalues_sorted,
2551 &mut range_rotation,
2552 );
2553 have_rotation = true;
2554 rescued_by_r_svd = true;
2555 }
2556 }
2557 }
2558 }
2559 if let Some(reason) = svd_refusal.as_deref() {
2560 if rescued_by_r_svd {
2561 log::warn!(
2562 "penalized-block rotation: stacked-root SVD {reason}. Recovered the SAME \
2563 right-singular basis from the Householder QR of `E` followed by the SVD \
2564 of its triangular factor `R`: `EᵀE = RᵀR`, so no accuracy is given up."
2565 );
2566 } else {
2567 log::warn!(
2571 "penalized-block rotation: stacked-root SVD {reason}, and so did the R-SVD \
2572 of its Householder QR factor. Recomputing it from the Gram `Σₖ λₖ Sₖ`, \
2573 which squares the condition number: recessive eigenvalues are resolved to \
2574 O(ε·d_max) rather than O(ε²·d_max)."
2575 );
2576 }
2577 }
2578 if !have_rotation {
2579 let mut range_block = Mat::<f64>::zeros(penalized_rank, penalized_rank);
2584 for (lambda, s_k) in lambdas.iter().zip(s_k_penalized_cache.iter()) {
2585 for i in 0..penalized_rank {
2586 for j in 0..penalized_rank {
2587 range_block[(i, j)] += *lambda * s_k[(i, j)];
2588 }
2589 }
2590 }
2591 let (range_eigenvalues, range_eigenvectors) =
2592 robust_eigh_faer(&range_block, Side::Lower, "range penalty block")?;
2593 let mut range_order: Vec<usize> = (0..penalized_rank).collect();
2594 range_order.sort_by(|&i, &j| {
2595 range_eigenvalues[j]
2596 .partial_cmp(&range_eigenvalues[i])
2597 .unwrap_or(std::cmp::Ordering::Equal)
2598 .then(i.cmp(&j))
2599 });
2600 range_eigenvalues_sorted = range_order
2601 .iter()
2602 .map(|&idx| range_eigenvalues[idx])
2603 .collect();
2604 for (col_idx, &idx) in range_order.iter().enumerate() {
2605 for row in 0..penalized_rank {
2606 range_rotation[(row, col_idx)] = range_eigenvectors[(row, idx)];
2607 }
2608 }
2609 }
2610 }
2611
2612 let structural_rank = penalized_rank;
2617 let range_eigs_sorted: Vec<f64> = range_eigenvalues_sorted;
2618
2619 let eigenvalue_floor = invariant.max_balanced_eigenvalue.max(1.0) * 1e-12;
2620 let qs = compose_qs_from_split(&q_pen, &q_null, p);
2621
2622 let leakage = assess_subspace_leakage(&qs, &rs_transformed, structural_rank, p);
2625 if !subspace_split_is_consistent(&leakage, p) {
2626 return Err(EstimationError::LayoutError(format!(
2627 "Reparameterization subspace split is inconsistent: max null leakage {:.3e} (rel {:.3e}, worst penalty {}), max |Qp'Qn| {:.3e}",
2628 leakage.max_abs_sq.sqrt(),
2629 leakage.max_rel_sq.sqrt(),
2630 leakage.worst_penalty,
2631 leakage.max_cross_gram_abs,
2632 )));
2633 }
2634
2635 let mut u_truncated_mat = Mat::<f64>::zeros(p, q_null.ncols());
2638 matmul(
2639 u_truncated_mat.as_mut(),
2640 Accum::Replace,
2641 qs.transpose(),
2642 q_null.as_ref(),
2643 1.0,
2644 Par::Seq,
2645 );
2646
2647 let mut e_transformed_mat = Mat::<f64>::zeros(structural_rank, p);
2653 for row_idx in 0..structural_rank {
2654 let safe_eigenval = range_eigs_sorted[row_idx].max(eigenvalue_floor);
2655 let sqrt_eigenval = safe_eigenval.sqrt();
2656 for j in 0..penalized_rank {
2658 e_transformed_mat[(row_idx, j)] = sqrt_eigenval * range_rotation[(j, row_idx)];
2659 }
2660 }
2661
2662 let mut floored_eigs: Vec<f64> = Vec::with_capacity(range_eigs_sorted.len());
2678 let mut log_det_sum = KahanSum::default();
2679 for (idx, &ev) in range_eigs_sorted.iter().enumerate() {
2680 if !ev.is_finite() || ev < -eigenvalue_floor {
2681 return Err(EstimationError::LayoutError(format!(
2682 "Penalty pseudo-logdet has a non-finite or large-negative structural eigenvalue at index {idx}: {ev:.3e}"
2683 )));
2684 }
2685 let safe_ev = ev.max(eigenvalue_floor);
2686 floored_eigs.push(safe_ev);
2687 if idx < penalized_rank {
2688 log_det_sum.add(safe_ev.ln());
2689 }
2690 }
2691 let log_det = log_det_sum.sum();
2692 let delta = 0.0;
2693
2694 let det1vec: Vec<f64> = (0..lambdas.len())
2697 .into_par_iter()
2698 .map(|k| {
2699 let s_k = &s_k_penalized_cache[k];
2700 let trace = trace_penalty_in_orthogonal_basis(
2704 s_k,
2705 penalized_rank,
2706 &range_rotation,
2707 &floored_eigs,
2708 delta,
2709 );
2710 lambdas[k] * trace
2711 })
2712 .collect();
2713
2714 {
2715 let mut maxdet1_mismatch = 0.0_f64;
2719 let mut det1_scale = 0.0_f64;
2720 for (k, lambda) in lambdas.iter().enumerate() {
2721 let s_k_penalized = &s_k_penalized_cache[k];
2722 let s_k_eigenbasis = orthogonal_similarity_transform_faer(
2723 s_k_penalized,
2724 penalized_rank,
2725 &range_rotation,
2726 );
2727 let mut trace = KahanSum::default();
2728 for l in 0..penalized_rank {
2729 trace.add(s_k_eigenbasis[(l, l)] / (floored_eigs[l] + delta));
2730 }
2731 let reference = *lambda * trace.sum();
2732 maxdet1_mismatch = maxdet1_mismatch.max((reference - det1vec[k]).abs());
2733 det1_scale = det1_scale.max(reference.abs()).max(det1vec[k].abs());
2734 }
2735 let det1_tolerance = 1e-7 * det1_scale.max(1.0);
2736 assert!(
2737 maxdet1_mismatch <= det1_tolerance,
2738 "det1 mismatch between optimized and reference formulas: max_abs={maxdet1_mismatch:.3e}, tol={det1_tolerance:.3e}"
2739 );
2740 }
2741
2742 let mut s_truncated = Mat::<f64>::zeros(p, p);
2753 matmul(
2754 s_truncated.as_mut(),
2755 Accum::Replace,
2756 e_transformed_mat.transpose(),
2757 e_transformed_mat.as_ref(),
2758 1.0,
2759 Par::Seq,
2760 );
2761
2762 {
2763 let mut max_null_diag = 0.0_f64;
2765 let mut max_null_offdiag = 0.0_f64;
2766 for i in structural_rank..p {
2767 max_null_diag = max_null_diag.max(s_truncated[(i, i)].abs());
2768 for j in 0..p {
2769 if i != j {
2770 max_null_offdiag = max_null_offdiag.max(s_truncated[(i, j)].abs());
2771 }
2772 }
2773 }
2774 assert!(
2775 max_null_diag <= 1e-10 && max_null_offdiag <= 1e-10,
2776 "null-space leakage in transformed penalty: max_null_diag={max_null_diag:.3e}, max_null_offdiag={max_null_offdiag:.3e}"
2777 );
2778 }
2779
2780 let qs_array = mat_to_array(&qs);
2781 let canonical_transformed: Vec<CanonicalPenalty> = rs_transformed
2782 .par_iter()
2783 .zip(penalties.par_iter())
2784 .map(|(r, cp)| {
2785 let mean_transformed = qs_array.t().dot(&cp.full_width_prior_mean());
2786 CanonicalPenalty::from_dense_root_with_mean(mat_to_array(r), p, mean_transformed)
2787 })
2788 .collect();
2789 Ok(ReparamResult {
2790 s_transformed: mat_to_array(&s_truncated),
2791 log_det,
2792 det1: Array1::from(det1vec),
2793 qs: qs_array,
2794 canonical_transformed,
2795 e_transformed: mat_to_array(&e_transformed_mat),
2796 u_truncated: mat_to_array(&u_truncated_mat),
2797 })
2798}
2799
2800#[derive(Debug, Clone, Copy, PartialEq, Eq)]
2802pub struct EngineDims {
2803 pub p: usize,
2804 pub k: usize,
2805}
2806
2807impl EngineDims {
2808 pub fn new(p: usize, k: usize) -> Self {
2809 Self { p, k }
2810 }
2811}
2812
2813pub fn stable_reparameterization_engine_canonical(
2820 penalties: &[CanonicalPenalty],
2821 lambdas: &[f64],
2822 dims: EngineDims,
2823 cached_invariant: Option<&ReparamInvariant>,
2824) -> Result<ReparamResult, EstimationError> {
2825 let owned;
2826 let invariant = match cached_invariant {
2827 Some(inv) => inv,
2828 None => {
2829 owned = precompute_reparam_invariant_from_canonical(penalties, dims.p)?;
2830 &owned
2831 }
2832 };
2833 stable_reparameterizationwith_invariant(penalties, lambdas, dims.p, invariant)
2834}
2835
2836#[derive(Clone)]
2846pub struct KroneckerReparamResult {
2847 pub reparameterized_marginals: Arc<Vec<Array2<f64>>>,
2853 pub marginal_eigenvalues: Arc<Vec<Array1<f64>>>,
2855 pub marginal_qs: Arc<Vec<Array2<f64>>>,
2857 pub log_det: f64,
2859 pub det1: Array1<f64>,
2861 pub det2: Array2<f64>,
2863 pub has_double_penalty: bool,
2865 pub marginal_dims: Vec<usize>,
2867}
2868
2869impl KroneckerReparamResult {
2870 pub fn materialize_qs(&self) -> Array2<f64> {
2873 let mut qs = Array2::<f64>::eye(1);
2874 for u_k in self.marginal_qs.iter() {
2875 qs = kronecker_product(&qs, u_k);
2876 }
2877 qs
2878 }
2879
2880 pub fn materialize_s_transformed(&self, lambdas: &[f64]) -> Array2<f64> {
2883 let d = self.marginal_dims.len();
2884 let p: usize = self.marginal_dims.iter().copied().product();
2885 let mut s = Array2::<f64>::zeros((p, p));
2886
2887 let eigenvalue_views: Vec<ArrayView1<'_, f64>> =
2891 self.marginal_eigenvalues.iter().map(|m| m.view()).collect();
2892 let has_double = self.has_double_penalty && lambdas.len() > d;
2893 let mut multi_idx = vec![0usize; d];
2894 let mut flat = 0usize;
2895 loop {
2896 let (sigma, _structural_sigma, _joint_null) = kronecker_cell_sigma(
2897 &eigenvalue_views,
2898 &multi_idx,
2899 lambdas,
2900 d,
2901 has_double,
2902 0.0,
2903 );
2904 s[[flat, flat]] = sigma;
2905 flat += 1;
2906
2907 if kronecker_multi_index_advance(&mut multi_idx, &self.marginal_dims) {
2908 break;
2909 }
2910 }
2911 s
2912 }
2913
2914 pub fn materialize_dense_artifact_result(
2917 &self,
2918 rs_list: &[Array2<f64>],
2919 lambdas: &[f64],
2920 p: usize,
2921 ) -> Result<ReparamResult, EstimationError> {
2922 const KRONECKER_DENSE_COMPAT_FALLBACK_MAX_P: usize = 4096;
2923 if p > KRONECKER_DENSE_COMPAT_FALLBACK_MAX_P {
2924 return Err(EstimationError::LayoutError(format!(
2925 "Kronecker reparameterization would materialize dense {}x{} compatibility tensors; \
2926 large-model dense fallback is disabled. Wire the downstream solver to consume \
2927 the factored Kronecker result directly",
2928 p, p
2929 )));
2930 }
2931 let qs = self.materialize_qs();
2932 let s_transformed = self.materialize_s_transformed(lambdas);
2933
2934 let rs_transformed: Vec<Array2<f64>> = if rs_list.len() >= 2 {
2936 use rayon::prelude::*;
2937 rs_list
2938 .par_iter()
2939 .map(|r| gam_linalg::faer_ndarray::fast_ab(r, &qs))
2940 .collect()
2941 } else {
2942 rs_list
2943 .iter()
2944 .map(|r| gam_linalg::faer_ndarray::fast_ab(r, &qs))
2945 .collect()
2946 };
2947 let d = self.marginal_dims.len();
2953 let eigenvalue_views: Vec<ArrayView1<'_, f64>> =
2960 self.marginal_eigenvalues.iter().map(|m| m.view()).collect();
2961 let has_double = self.has_double_penalty && lambdas.len() > d;
2962 let diag_vals: Vec<f64> = {
2963 let mut vals = Vec::with_capacity(p);
2964 let mut multi_idx = vec![0usize; d];
2965 loop {
2966 let (sigma, _structural_sigma, _joint_null) = kronecker_cell_sigma(
2967 &eigenvalue_views,
2968 &multi_idx,
2969 lambdas,
2970 d,
2971 has_double,
2972 0.0,
2973 );
2974 vals.push(if sigma > 0.0 { sigma.sqrt() } else { 0.0 });
2975
2976 if kronecker_multi_index_advance(&mut multi_idx, &self.marginal_dims) {
2977 break;
2978 }
2979 }
2980 vals
2981 };
2982 let rank = diag_vals.iter().filter(|&&v| v > 1e-12).count();
2983 let mut e_transformed = Array2::<f64>::zeros((rank, p));
2984 let mut row = 0;
2985 for (j, &v) in diag_vals.iter().enumerate() {
2986 if v > 1e-12 {
2987 e_transformed[[row, j]] = v;
2988 row += 1;
2989 }
2990 }
2991
2992 let null_count = p - rank;
2994 let mut u_truncated = Array2::<f64>::zeros((p, null_count));
2995 let mut col = 0;
2996 for (j, &v) in diag_vals.iter().enumerate() {
2997 if v <= 1e-12 {
2998 u_truncated[[j, col]] = 1.0; col += 1;
3000 }
3001 }
3002
3003 let canonical_transformed: Vec<CanonicalPenalty> = rs_transformed
3004 .iter()
3005 .map(|r| CanonicalPenalty::from_dense_root(r.clone(), p))
3006 .collect();
3007 Ok(ReparamResult {
3008 s_transformed,
3009 log_det: self.log_det,
3010 det1: self.det1.clone(),
3011 qs,
3012 canonical_transformed,
3013 e_transformed,
3014 u_truncated,
3015 })
3016 }
3017}
3018
3019const KRONECKER_STRUCTURAL_ZERO_TOL: f64 = 1e-12;
3026
3027#[inline]
3041fn kronecker_cell_sigma(
3042 marginal_eigenvalues: &[ArrayView1<'_, f64>],
3043 multi_idx: &[usize],
3044 lambdas: &[f64],
3045 d: usize,
3046 has_double_penalty: bool,
3047 objective_ridge: f64,
3048) -> (f64, f64, bool) {
3049 let mut sigma = 0.0;
3050 let mut structural_sigma = 0.0;
3051 for k in 0..d {
3052 let marginal_eigenvalue = marginal_eigenvalues[k][multi_idx[k]];
3053 structural_sigma += marginal_eigenvalue;
3054 sigma += lambdas[k] * marginal_eigenvalue;
3055 }
3056 let joint_null = structural_sigma <= KRONECKER_STRUCTURAL_ZERO_TOL;
3057 if has_double_penalty && joint_null {
3058 sigma += lambdas[d];
3059 }
3060 if structural_sigma > KRONECKER_STRUCTURAL_ZERO_TOL {
3061 sigma += objective_ridge;
3062 }
3063 (sigma, structural_sigma, joint_null)
3064}
3065
3066#[inline]
3069fn kronecker_multi_index_advance(multi_idx: &mut [usize], dims: &[usize]) -> bool {
3070 let mut carry = true;
3071 for dim in (0..dims.len()).rev() {
3072 if carry {
3073 multi_idx[dim] += 1;
3074 if multi_idx[dim] < dims[dim] {
3075 carry = false;
3076 } else {
3077 multi_idx[dim] = 0;
3078 }
3079 }
3080 }
3081 carry
3082}
3083
3084pub fn kronecker_logdet_and_derivatives(
3085 marginal_eigenvalues: &[ArrayView1<'_, f64>],
3086 marginal_dims: &[usize],
3087 lambdas: &[f64],
3088 has_double_penalty: bool,
3089 objective_ridge: f64,
3090) -> (f64, Array1<f64>, Array2<f64>) {
3091 let d = marginal_dims.len();
3092 let n_pen = d + if has_double_penalty { 1 } else { 0 };
3093
3094 let mut logdet = 0.0;
3095 let mut grad = Array1::<f64>::zeros(n_pen);
3096 let mut hess = Array2::<f64>::zeros((n_pen, n_pen));
3097 let tol = 1e-12;
3098
3099 let mut multi_idx = vec![0usize; d];
3100 loop {
3101 let (sigma, _structural_sigma, joint_null) = kronecker_cell_sigma(
3102 marginal_eigenvalues,
3103 &multi_idx,
3104 lambdas,
3105 d,
3106 has_double_penalty,
3107 objective_ridge,
3108 );
3109
3110 if sigma > tol {
3111 logdet += sigma.ln();
3112 let inv_sigma = 1.0 / sigma;
3113 let inv_sigma2 = inv_sigma * inv_sigma;
3114
3115 for k in 0..d {
3116 let ck = lambdas[k] * marginal_eigenvalues[k][multi_idx[k]];
3117 grad[k] += ck * inv_sigma;
3118 }
3119 if has_double_penalty && joint_null {
3120 grad[d] += lambdas[d] * inv_sigma;
3121 }
3122
3123 for k in 0..n_pen {
3124 let ck = if k < d {
3125 lambdas[k] * marginal_eigenvalues[k][multi_idx[k]]
3126 } else if joint_null {
3127 lambdas[d]
3128 } else {
3129 0.0
3130 };
3131 if ck == 0.0 {
3138 continue;
3139 }
3140 hess[[k, k]] += ck * inv_sigma - ck * ck * inv_sigma2;
3141 for l in (k + 1)..n_pen {
3142 let cl = if l < d {
3143 lambdas[l] * marginal_eigenvalues[l][multi_idx[l]]
3144 } else if joint_null {
3145 lambdas[d]
3146 } else {
3147 0.0
3148 };
3149 let off = -ck * cl * inv_sigma2;
3150 hess[[k, l]] += off;
3151 hess[[l, k]] += off;
3152 }
3153 }
3154 }
3155
3156 if kronecker_multi_index_advance(&mut multi_idx, marginal_dims) {
3157 break;
3158 }
3159 }
3160
3161 (logdet, grad, hess)
3162}
3163
3164use crate::kronecker::KroneckerInvariantStructure;
3168
3169pub fn kronecker_reparameterization_engine(
3175 marginal_designs: &[Array2<f64>],
3176 marginal_penalties: &[Array2<f64>],
3177 marginal_dims: &[usize],
3178 lambdas: &[f64],
3179 has_double_penalty: bool,
3180) -> Result<KroneckerReparamResult, EstimationError> {
3181 let d = marginal_dims.len();
3182 if marginal_designs.len() != d || marginal_penalties.len() != d {
3183 return Err(EstimationError::LayoutError(format!(
3184 "kronecker_reparameterization_engine: dimension mismatch: designs={}, penalties={}, dims={}",
3185 marginal_designs.len(),
3186 marginal_penalties.len(),
3187 d
3188 )));
3189 }
3190
3191 let invariant =
3192 KroneckerInvariantStructure::compute(marginal_designs, marginal_penalties, marginal_dims)?;
3193 kronecker_reparameterization_engine_with_invariant(
3194 &invariant,
3195 marginal_dims,
3196 lambdas,
3197 has_double_penalty,
3198 )
3199}
3200
3201pub fn kronecker_reparameterization_engine_with_invariant(
3209 invariant: &KroneckerInvariantStructure,
3210 marginal_dims: &[usize],
3211 lambdas: &[f64],
3212 has_double_penalty: bool,
3213) -> Result<KroneckerReparamResult, EstimationError> {
3214 let marginal_eigenvalues = Arc::clone(&invariant.marginal_eigenvalues);
3217 let marginal_qs = Arc::clone(&invariant.marginal_qs);
3218 let reparameterized_marginals = Arc::clone(&invariant.reparameterized_marginals);
3219
3220 let marginal_eigenvalue_views: Vec<_> = marginal_eigenvalues
3221 .iter()
3222 .map(|evals| evals.view())
3223 .collect();
3224 let (log_det, det1, det2) = kronecker_logdet_and_derivatives(
3225 &marginal_eigenvalue_views,
3226 marginal_dims,
3227 lambdas,
3228 has_double_penalty,
3229 0.0,
3230 );
3231
3232 Ok(KroneckerReparamResult {
3233 reparameterized_marginals,
3234 marginal_eigenvalues,
3235 marginal_qs,
3236 log_det,
3237 det1,
3238 det2,
3239 has_double_penalty,
3240 marginal_dims: marginal_dims.to_vec(),
3241 })
3242}
3243
3244#[cfg(test)]
3245mod tests {
3246 use super::{
3247 CanonicalPenalty, REL_PSD_FLOOR, SubspaceLeakageMetrics, assess_subspace_leakage,
3248 classify_eigenvalues_strict, precompute_reparam_invariant_from_canonical,
3249 report_penalty_pair_redundancy, stable_reparameterizationwith_invariant,
3250 subspace_split_is_consistent,
3251 };
3252 use crate::EstimationError;
3253 use crate::construction::kronecker_product;
3254 use faer::Mat;
3255 use gam_linalg::faer_ndarray::FaerEigh;
3256 use gam_linalg::utils::inf_norm;
3257 use ndarray::{Array1, Array2, array};
3258
3259 fn canonical_from_roots(rs_list: &[Array2<f64>], p: usize) -> Vec<CanonicalPenalty> {
3261 rs_list
3262 .iter()
3263 .map(|r| {
3264 let local = r.t().dot(r);
3265 CanonicalPenalty {
3266 root: r.clone(),
3267 col_range: 0..p,
3268 total_dim: p,
3269 nullity: 0,
3270 local,
3271 prior_mean: Array1::zeros(p),
3272 positive_eigenvalues: Vec::new(),
3273 op: None,
3274 }
3275 })
3276 .collect()
3277 }
3278
3279 fn metrics_for(
3280 qs: &Mat<f64>,
3281 rs: &[Mat<f64>],
3282 structural_rank: usize,
3283 p: usize,
3284 ) -> SubspaceLeakageMetrics {
3285 assess_subspace_leakage(qs, rs, structural_rank, p)
3286 }
3287
3288 #[test]
3289 fn subspace_leakage_iszero_for_clean_split() {
3290 let p = 4usize;
3291 let structural_rank = 2usize;
3292 let qs = Mat::<f64>::identity(p, p);
3293 let mut r0 = Mat::<f64>::zeros(2, p);
3294 r0[(0, 0)] = 1.0;
3295 r0[(1, 1)] = 2.0;
3296
3297 let m = metrics_for(&qs, &[r0], structural_rank, p);
3298 assert!(m.max_abs_sq <= 1e-16);
3299 assert!(m.max_rel_sq <= 1e-16);
3300 assert!(m.max_cross_gram_abs <= 1e-16);
3301 }
3302
3303 #[test]
3304 fn subspace_leakage_detects_null_column_energy() {
3305 let p = 4usize;
3306 let structural_rank = 2usize;
3307 let qs = Mat::<f64>::identity(p, p);
3308 let mut r0 = Mat::<f64>::zeros(1, p);
3309 r0[(0, 2)] = 3.0;
3310
3311 let m = metrics_for(&qs, &[r0], structural_rank, p);
3312 assert!(m.max_abs_sq > 0.0);
3313 assert!(m.max_rel_sq > 0.99);
3314 }
3315
3316 #[test]
3317 fn subspace_leakage_detects_qp_qn_nonorthogonality() {
3318 let p = 3usize;
3319 let structural_rank = 1usize;
3320 let mut qs = Mat::<f64>::identity(p, p);
3321 qs[(0, 1)] = 0.2;
3322 let r0 = Mat::<f64>::zeros(1, p);
3323
3324 let m = metrics_for(&qs, &[r0], structural_rank, p);
3325 assert!(m.max_cross_gram_abs > 1e-3);
3326 }
3327
3328 #[test]
3329 fn subspace_split_admits_near_threshold_manifold_leakage_1802() {
3330 let p = 40usize;
3344 let structural_rank = p - 1;
3345 let null_amp = (1.06e-8_f64).sqrt();
3348 let mut rs = Mat::<f64>::zeros(1, p);
3349 rs[(0, 0)] = 1.0;
3350 rs[(0, p - 1)] = null_amp;
3351 let qs = Mat::<f64>::identity(p, p);
3352 let leakage = metrics_for(&qs, &[rs], structural_rank, p);
3353 assert!(
3356 leakage.max_rel_sq > 1e-10 && leakage.max_rel_sq < 1e-6,
3357 "reproduced leakage should sit in the near-REL_PSD_FLOOR band, got {:.3e}",
3358 leakage.max_rel_sq
3359 );
3360 assert!(leakage.max_cross_gram_abs <= 1e-12);
3361 assert!(
3362 subspace_split_is_consistent(&leakage, p),
3363 "near-REL_PSD_FLOOR null leakage on a manifold basis must be admitted \
3364 (rel_sq={:.3e}, tol={:.3e})",
3365 leakage.max_rel_sq,
3366 (p as f64) * REL_PSD_FLOOR,
3367 );
3368 }
3369
3370 #[test]
3371 fn subspace_split_still_rejects_genuine_inconsistency_1802() {
3372 let p = 40usize;
3377 let structural_rank = p - 1;
3378 let mut rs = Mat::<f64>::zeros(1, p);
3379 rs[(0, p - 1)] = 1.0; let qs = Mat::<f64>::identity(p, p);
3381 let leakage = metrics_for(&qs, &[rs], structural_rank, p);
3382 assert!(leakage.max_rel_sq > 0.99);
3383 assert!(
3384 !subspace_split_is_consistent(&leakage, p),
3385 "an O(1) null-block leakage is a real inconsistency and must be rejected"
3386 );
3387
3388 let mut qs_bad = Mat::<f64>::identity(3, 3);
3390 qs_bad[(0, 1)] = 0.2;
3391 let clean = Mat::<f64>::zeros(1, 3);
3392 let leakage2 = metrics_for(&qs_bad, &[clean], 1, 3);
3393 assert!(leakage2.max_cross_gram_abs > 1e-3);
3394 assert!(
3395 !subspace_split_is_consistent(&leakage2, 3),
3396 "a non-orthogonal Qp/Qn split must be rejected"
3397 );
3398 }
3399
3400 #[test]
3401 fn u_truncated_is_transformed_frame_in_nonzero_case() {
3402 let p = 3usize;
3403 let rs_list = vec![array![[1.0, 0.0, 0.0]]];
3404 let canonical = canonical_from_roots(&rs_list, p);
3405 let lambdas = vec![2.0];
3406 let inv = precompute_reparam_invariant_from_canonical(&canonical, p)
3407 .expect("precompute invariant");
3408 let rep = stable_reparameterizationwith_invariant(&canonical, &lambdas, p, &inv)
3409 .expect("stable reparam");
3410
3411 let expected = rep.qs.t().dot(&inv.split.q_null);
3412 let diff = &rep.u_truncated - &expected;
3413 let max_abs = inf_norm(diff.iter().copied());
3414 assert!(
3415 max_abs <= 1e-10,
3416 "u_truncated frame mismatch: max_abs={max_abs}"
3417 );
3418 }
3419
3420 #[test]
3421 fn infinite_lambda_keeps_range_penalty_block_finite_1379() {
3422 let p = 3usize;
3439 let rs_list = vec![array![[1.0, 0.0, 0.0]], array![[0.0, 1.0, 0.0]]];
3440 let canonical = canonical_from_roots(&rs_list, p);
3441 let inv = precompute_reparam_invariant_from_canonical(&canonical, p)
3442 .expect("precompute invariant");
3443
3444 let lambdas_inf = vec![f64::INFINITY, 3.0];
3445 let inf_result =
3446 stable_reparameterizationwith_invariant(&canonical, &lambdas_inf, p, &inv);
3447 assert!(
3448 inf_result.is_err(),
3449 "an infinite lambda must surface as an error, not be silently clamped (#1074)"
3450 );
3451
3452 let lambdas_big = vec![1e300_f64, 3.0];
3456 let rep = stable_reparameterizationwith_invariant(&canonical, &lambdas_big, p, &inv)
3457 .expect("stable reparam at large-but-finite lambda");
3458 assert!(
3459 rep.s_transformed.iter().all(|v| v.is_finite()),
3460 "transformed penalty must be finite at large-but-finite lambda"
3461 );
3462 assert!(
3463 rep.qs.iter().all(|v| v.is_finite()),
3464 "reparam rotation must be finite at large-but-finite lambda"
3465 );
3466 assert!(
3467 rep.log_det.is_finite(),
3468 "penalty log-det must be finite at large-but-finite lambda"
3469 );
3470 assert!(
3471 rep.det1.iter().all(|v| v.is_finite()),
3472 "penalty log-det derivatives must be finite at large-but-finite lambda"
3473 );
3474 }
3475
3476 #[test]
3477 fn u_truncated_is_identitywhen_no_penalties() {
3478 let p = 4usize;
3479 let canonical: Vec<CanonicalPenalty> = Vec::new();
3480 let lambdas: Vec<f64> = Vec::new();
3481 let inv = precompute_reparam_invariant_from_canonical(&canonical, p)
3482 .expect("precompute invariant");
3483 let rep = stable_reparameterizationwith_invariant(&canonical, &lambdas, p, &inv)
3484 .expect("stable reparam");
3485 assert_eq!(rep.u_truncated, Array2::<f64>::eye(p));
3486 }
3487
3488 #[test]
3489 fn kronecker_memoized_invariant_is_bit_identical_to_unmemoized_engine() {
3490 let marginal_designs = vec![
3497 array![[1.0, 0.3, -0.2], [0.4, 1.0, 0.1], [-0.1, 0.2, 1.0]],
3498 array![[1.0, -0.5], [0.2, 1.0], [0.7, 0.3]],
3499 ];
3500 let marginal_penalties = vec![
3501 array![[2.0, -1.0, 0.0], [-1.0, 2.0, -1.0], [0.0, -1.0, 1.0]],
3502 array![[3.0, -1.5], [-1.5, 3.0]],
3503 ];
3504 let marginal_dims = vec![3usize, 2usize];
3505
3506 let invariant = super::KroneckerInvariantStructure::compute(
3507 &marginal_designs,
3508 &marginal_penalties,
3509 &marginal_dims,
3510 )
3511 .expect("invariant structure");
3512
3513 for lambdas in [
3514 vec![5.0, 7.0],
3515 vec![0.0, 7.0],
3516 vec![5.0, 0.0],
3517 vec![1e-3, 1e3],
3518 ] {
3519 {
3520 let unmemoized = super::kronecker_reparameterization_engine(
3521 &marginal_designs,
3522 &marginal_penalties,
3523 &marginal_dims,
3524 &lambdas,
3525 true,
3526 )
3527 .expect("unmemoized engine");
3528 let memoized = super::kronecker_reparameterization_engine_with_invariant(
3529 &invariant,
3530 &marginal_dims,
3531 &lambdas,
3532 true,
3533 )
3534 .expect("memoized engine");
3535
3536 assert_eq!(memoized.log_det.to_bits(), unmemoized.log_det.to_bits());
3537 for (a, b) in memoized.det1.iter().zip(unmemoized.det1.iter()) {
3538 assert_eq!(a.to_bits(), b.to_bits());
3539 }
3540 for (a, b) in memoized.det2.iter().zip(unmemoized.det2.iter()) {
3541 assert_eq!(a.to_bits(), b.to_bits());
3542 }
3543 for (ma, ua) in memoized
3544 .reparameterized_marginals
3545 .iter()
3546 .zip(unmemoized.reparameterized_marginals.iter())
3547 {
3548 for (a, b) in ma.iter().zip(ua.iter()) {
3549 assert_eq!(a.to_bits(), b.to_bits());
3550 }
3551 }
3552 for (mq, uq) in memoized
3553 .marginal_qs
3554 .iter()
3555 .zip(unmemoized.marginal_qs.iter())
3556 {
3557 for (a, b) in mq.iter().zip(uq.iter()) {
3558 assert_eq!(a.to_bits(), b.to_bits());
3559 }
3560 }
3561 }
3562 }
3563 }
3564
3565 #[test]
3566 fn kronecker_double_penalty_shrinks_only_joint_null_space() {
3567 let marginal_designs = vec![Array2::<f64>::eye(2), Array2::<f64>::eye(2)];
3568 let marginal_penalties = vec![
3569 array![[0.0, 0.0], [0.0, 2.0]],
3570 array![[0.0, 0.0], [0.0, 3.0]],
3571 ];
3572 let marginal_dims = vec![2usize, 2usize];
3573 let lambdas = vec![5.0, 7.0, 11.0];
3574
3575 let rep = super::kronecker_reparameterization_engine(
3576 &marginal_designs,
3577 &marginal_penalties,
3578 &marginal_dims,
3579 &lambdas,
3580 true,
3581 )
3582 .expect("kronecker reparameterization");
3583
3584 let s = rep.materialize_s_transformed(&lambdas);
3585 let expected = [11.0, 21.0, 10.0, 31.0];
3586 for (idx, expected_diag) in expected.iter().copied().enumerate() {
3587 assert!(
3588 (s[[idx, idx]] - expected_diag).abs() <= 1e-12,
3589 "diagonal {idx} got {}, expected {expected_diag}",
3590 s[[idx, idx]]
3591 );
3592 }
3593
3594 let expected_logdet: f64 = expected.iter().map(|v| f64::ln(*v)).sum();
3595 assert!((rep.log_det - expected_logdet).abs() <= 1e-12);
3596 assert!(
3597 (rep.det1[2] - 1.0).abs() <= 1e-12,
3598 "double-penalty derivative must come only from the joint null mode, got {}",
3599 rep.det1[2]
3600 );
3601 assert!(rep.det2[[2, 2]].abs() <= 1e-12);
3602
3603 let tensor_roots = vec![
3604 array![
3605 [0.0, 0.0, 2.0_f64.sqrt(), 0.0],
3606 [0.0, 0.0, 0.0, 2.0_f64.sqrt()]
3607 ],
3608 array![
3609 [0.0, 3.0_f64.sqrt(), 0.0, 0.0],
3610 [0.0, 0.0, 0.0, 3.0_f64.sqrt()]
3611 ],
3612 ];
3613 let dense = rep
3614 .materialize_dense_artifact_result(&tensor_roots, &lambdas, 4)
3615 .expect("dense artifact materialization");
3616 for (idx, expected_diag) in expected.iter().copied().enumerate() {
3617 assert!(
3618 (dense.s_transformed[[idx, idx]] - expected_diag).abs() <= 1e-12,
3619 "dense artifact diagonal {idx} got {}, expected {expected_diag}",
3620 dense.s_transformed[[idx, idx]]
3621 );
3622 }
3623 }
3624
3625 #[test]
3626 fn transformed_penalty_is_diagonal_in_transformed_frame() {
3627 let p = 3usize;
3628 let inv_sqrt2 = 2.0_f64.sqrt().recip();
3629 let rs_list = vec![array![[inv_sqrt2, inv_sqrt2, 0.0]]];
3631 let canonical = canonical_from_roots(&rs_list, p);
3632 let lambdas = vec![4.0];
3633 let inv = precompute_reparam_invariant_from_canonical(&canonical, p)
3634 .expect("precompute invariant");
3635 let rep = stable_reparameterizationwith_invariant(&canonical, &lambdas, p, &inv)
3636 .expect("stable reparam");
3637
3638 assert_eq!(rep.e_transformed.nrows(), 1);
3639 assert!(rep.e_transformed[[0, 0]].abs() > 0.0);
3640 assert!(rep.e_transformed[[0, 1]].abs() <= 1e-12);
3641 assert!(rep.e_transformed[[0, 2]].abs() <= 1e-12);
3642 let expected_det1 = 1.0_f64;
3645 assert!((rep.det1[0] - expected_det1).abs() <= 1e-12);
3646
3647 let s = rep.s_transformed;
3648 let mut max_offdiag = 0.0_f64;
3649 for i in 0..p {
3650 for j in 0..p {
3651 if i != j {
3652 max_offdiag = max_offdiag.max(s[[i, j]].abs());
3653 }
3654 }
3655 }
3656 assert!(
3657 max_offdiag <= 1e-10,
3658 "transformed penalty should be diagonal, max offdiag={max_offdiag}"
3659 );
3660 assert!(s[[1, 1]].abs() <= 1e-10);
3661 assert!(s[[2, 2]].abs() <= 1e-10);
3662 }
3663
3664 #[test]
3665 fn det1_matches_rank_for_single_full_rank_penalty() {
3666 let p = 2usize;
3667 let inv_sqrt2 = 2.0_f64.sqrt().recip();
3668 let q_t = [[inv_sqrt2, inv_sqrt2], [-inv_sqrt2, inv_sqrt2]];
3670 let rs = array![
3672 [3.0 * q_t[0][0], 3.0 * q_t[0][1]],
3673 [1.0 * q_t[1][0], 1.0 * q_t[1][1]]
3674 ];
3675 let rs_list = vec![rs];
3676 let canonical = canonical_from_roots(&rs_list, p);
3677 let lambdas = vec![5.0];
3678
3679 let inv = precompute_reparam_invariant_from_canonical(&canonical, p)
3680 .expect("precompute invariant");
3681 let rep = stable_reparameterizationwith_invariant(&canonical, &lambdas, p, &inv)
3682 .expect("stable reparam");
3683
3684 assert_eq!(rep.e_transformed.nrows(), p);
3685 let det1 = rep.det1[0];
3686 let s_k_eigs = [9.0_f64, 1.0_f64];
3690 let lambda = 5.0_f64;
3691 let expected_det1: f64 = s_k_eigs.iter().map(|&d| lambda * d / (lambda * d)).sum();
3692 assert!(
3693 (det1 - expected_det1).abs() <= 1e-12,
3694 "expected det1={expected_det1}, got {det1}",
3695 );
3696
3697 let s = rep.s_transformed;
3698 assert!(s[[0, 1]].abs() <= 1e-10);
3699 assert!(s[[1, 0]].abs() <= 1e-10);
3700 assert!(s[[0, 0]] > 0.0);
3701 assert!(s[[1, 1]] > 0.0);
3702 }
3703
3704 #[test]
3705 fn kronecker_reparam_logdet_matches_dense() {
3706 let q1 = 3;
3709 let q2 = 4;
3710 let s1 = {
3711 let mut s = Array2::<f64>::zeros((q1, q1));
3712 s[[0, 0]] = 1.0;
3714 s[[0, 1]] = -1.0;
3715 s[[1, 0]] = -1.0;
3716 s[[1, 1]] = 2.0;
3717 s[[1, 2]] = -1.0;
3718 s[[2, 1]] = -1.0;
3719 s[[2, 2]] = 1.0;
3720 s
3721 };
3722 let s2 = {
3723 let mut s = Array2::<f64>::zeros((q2, q2));
3724 s[[0, 0]] = 1.0;
3725 s[[0, 1]] = -1.0;
3726 s[[1, 0]] = -1.0;
3727 s[[1, 1]] = 2.0;
3728 s[[1, 2]] = -1.0;
3729 s[[2, 1]] = -1.0;
3730 s[[2, 2]] = 2.0;
3731 s[[2, 3]] = -1.0;
3732 s[[3, 2]] = -1.0;
3733 s[[3, 3]] = 1.0;
3734 s
3735 };
3736
3737 let lambdas = [2.5, 1.3];
3738 let p = q1 * q2;
3740 let i1 = Array2::<f64>::eye(q1);
3741 let i2 = Array2::<f64>::eye(q2);
3742 let pen0 = kronecker_product(&s1, &i2);
3743 let pen1 = kronecker_product(&i1, &s2);
3744 let mut s_dense = Array2::<f64>::zeros((p, p));
3745 s_dense.scaled_add(lambdas[0], &pen0);
3746 s_dense.scaled_add(lambdas[1], &pen1);
3747
3748 let (evals_dense, _): (ndarray::Array1<f64>, ndarray::Array2<f64>) =
3750 s_dense.eigh(faer::Side::Lower).unwrap();
3751 let tol = 1e-12;
3752 let ref_logdet: f64 = evals_dense
3753 .iter()
3754 .filter(|&&v: &&f64| v > tol)
3755 .map(|&v: &f64| v.ln())
3756 .sum();
3757
3758 let marginal_designs = vec![
3760 Array2::<f64>::eye(q1), Array2::<f64>::eye(q2),
3762 ];
3763 let marginal_penalties = vec![s1, s2];
3764 let kron_result = super::kronecker_reparameterization_engine(
3765 &marginal_designs,
3766 &marginal_penalties,
3767 &[q1, q2],
3768 &lambdas,
3769 false,
3770 )
3771 .unwrap();
3772
3773 let diff = (kron_result.log_det - ref_logdet).abs();
3774 assert!(
3775 diff < 1e-8,
3776 "Kronecker logdet {:.10} vs dense {:.10}, diff={:.3e}",
3777 kron_result.log_det,
3778 ref_logdet,
3779 diff,
3780 );
3781
3782 let rhos: Vec<f64> = lambdas.iter().map(|&l| l.ln()).collect();
3784 let eps = 1e-5;
3785 for k in 0..2 {
3786 let mut rho_plus = rhos.clone();
3787 rho_plus[k] += eps;
3788 let mut rho_minus = rhos.clone();
3789 rho_minus[k] -= eps;
3790 let lam_plus: Vec<f64> = rho_plus.iter().map(|&r| r.exp()).collect();
3791 let lam_minus: Vec<f64> = rho_minus.iter().map(|&r| r.exp()).collect();
3792 let result_plus = super::kronecker_reparameterization_engine(
3793 &marginal_designs,
3794 &marginal_penalties,
3795 &[q1, q2],
3796 &lam_plus,
3797 false,
3798 )
3799 .unwrap();
3800 let result_minus = super::kronecker_reparameterization_engine(
3801 &marginal_designs,
3802 &marginal_penalties,
3803 &[q1, q2],
3804 &lam_minus,
3805 false,
3806 )
3807 .unwrap();
3808 let fd_deriv = (result_plus.log_det - result_minus.log_det) / (2.0 * eps);
3809 let analytic_deriv = kron_result.det1[k];
3810 let rel_err = if analytic_deriv.abs() > 1e-10 {
3811 (fd_deriv - analytic_deriv).abs() / analytic_deriv.abs()
3812 } else {
3813 (fd_deriv - analytic_deriv).abs()
3814 };
3815 assert!(
3816 rel_err < 1e-4,
3817 "det1[{k}] mismatch: analytic={:.8}, fd={:.8}, rel_err={:.3e}",
3818 analytic_deriv,
3819 fd_deriv,
3820 rel_err,
3821 );
3822 }
3823 }
3824
3825 #[test]
3826 fn classify_strict_rejects_nan_eigenvalue() {
3827 let mut eigs = [1.0, f64::NAN, 0.5];
3828 match classify_eigenvalues_strict(&mut eigs, "test_nan") {
3829 Err(EstimationError::PenaltySpectrumNonFinite {
3830 context,
3831 index,
3832 value,
3833 }) => {
3834 assert_eq!(context, "test_nan");
3835 assert_eq!(index, 1);
3836 assert!(value.is_nan());
3837 }
3838 other => panic!("expected PenaltySpectrumNonFinite, got {:?}", other),
3839 }
3840 }
3841
3842 #[test]
3843 fn classify_strict_rejects_inf_eigenvalue() {
3844 let mut eigs = [1.0, 0.5, f64::INFINITY];
3845 match classify_eigenvalues_strict(&mut eigs, "test_inf") {
3846 Err(EstimationError::PenaltySpectrumNonFinite { index, value, .. }) => {
3847 assert_eq!(index, 2);
3848 assert!(value.is_infinite());
3849 }
3850 other => panic!("expected PenaltySpectrumNonFinite, got {:?}", other),
3851 }
3852 }
3853
3854 #[test]
3855 fn classify_strict_rejects_materially_indefinite() {
3856 let mut eigs = [1.0, -1e-2, 0.5];
3858 match classify_eigenvalues_strict(&mut eigs, "test_indef") {
3859 Err(EstimationError::PenaltySpectrumIndefinite {
3860 context,
3861 index,
3862 value,
3863 ..
3864 }) => {
3865 assert_eq!(context, "test_indef");
3866 assert_eq!(index, 1);
3867 assert!((value + 1e-2).abs() <= 1e-15);
3868 }
3869 other => panic!("expected PenaltySpectrumIndefinite, got {:?}", other),
3870 }
3871 }
3872
3873 #[test]
3874 fn classify_strict_accepts_roundoff_negative() {
3875 let scale = 1.0_f64;
3877 let roundoff = -1e-16 * scale;
3878 let mut eigs = [scale, 0.5 * scale, roundoff, 0.25 * scale];
3879 classify_eigenvalues_strict(&mut eigs, "test_roundoff").expect("roundoff must classify");
3880 assert_eq!(eigs[2], 0.0);
3882 assert!(eigs[0] > 0.0 && eigs[1] > 0.0 && eigs[3] > 0.0);
3884 }
3885
3886 #[test]
3887 fn classify_strict_accepts_extreme_lambda_assembly_noise_1619() {
3888 let scale = 8.509e12_f64;
3895 let noise = -6.546e2_f64;
3897 assert!(
3898 (noise.abs() / scale) < 1.0e-10,
3899 "fixture must reproduce the ~1e-11-relative noise from #1619"
3900 );
3901 let mut eigs = vec![scale, 0.5 * scale, noise, 0.1 * scale];
3902 classify_eigenvalues_strict(&mut eigs, "range penalty block")
3903 .expect("a ~1e-11-relative roundoff-negative eigenvalue must be accepted (#1619)");
3904 assert_eq!(eigs[2], 0.0);
3906 assert!(eigs[0] > 0.0 && eigs[1] > 0.0 && eigs[3] > 0.0);
3908 }
3909
3910 #[test]
3911 fn classify_strict_snaps_subtol_positive_to_zero() {
3912 let scale = 10.0_f64;
3915 let subtol = 1e-15 * scale;
3916 let mut eigs = [scale, subtol];
3917 classify_eigenvalues_strict(&mut eigs, "test_sub_pos").expect("sub-tol positive ok");
3918 assert_eq!(eigs[1], 0.0);
3919 }
3920
3921 fn canonical_from_local(
3925 local: Array2<f64>,
3926 col_range: std::ops::Range<usize>,
3927 total_dim: usize,
3928 ) -> CanonicalPenalty {
3929 let block_dim = local.nrows();
3930 let root = Array2::<f64>::zeros((0, block_dim));
3932 CanonicalPenalty {
3933 root,
3934 col_range,
3935 total_dim,
3936 nullity: 0,
3937 local,
3938 prior_mean: Array1::zeros(block_dim),
3939 positive_eigenvalues: Vec::new(),
3940 op: None,
3941 }
3942 }
3943
3944 #[test]
3945 fn report_penalty_pair_redundancy_detects_identical_pair() {
3946 let s0 = ndarray::array![[2.0, 0.5, 0.0], [0.5, 1.0, 0.25], [0.0, 0.25, 1.5],];
3948 let s_shared = ndarray::array![[1.0, -0.5, 0.0], [-0.5, 2.0, -0.5], [0.0, -0.5, 1.0],];
3951
3952 let bundle = vec![
3953 canonical_from_local(s0, 0..3, 3),
3954 canonical_from_local(s_shared.clone(), 0..3, 3),
3955 canonical_from_local(s_shared, 0..3, 3),
3956 ];
3957
3958 let redundant = report_penalty_pair_redundancy(&bundle);
3959
3960 assert_eq!(
3963 redundant.len(),
3964 1,
3965 "expected exactly one redundant pair, got {:?}",
3966 redundant
3967 );
3968 let (i, j, defect) = redundant[0];
3969 assert_eq!((i, j), (1, 2));
3970 assert_eq!(
3971 defect, 0.0,
3972 "bit-identical penalties leave an EXACTLY zero residual, not a small one"
3973 );
3974 }
3975
3976 #[test]
3985 fn a_pair_1p9e_minus_5_apart_is_not_proportional_2676() {
3986 const DEFECT: f64 = 1.874_020e-5;
3987 let base = ndarray::array![[1.0, -0.5, 0.0], [-0.5, 2.0, -0.5], [0.0, -0.5, 1.0]];
3988 let base_norm = base.iter().map(|v| v * v).sum::<f64>().sqrt();
3989 let mut direction = ndarray::array![[0.0, 0.0, 1.0], [0.0, 0.0, 0.0], [1.0, 0.0, 0.0]];
3993 let overlap = direction
3994 .iter()
3995 .zip(base.iter())
3996 .map(|(a, b)| a * b)
3997 .sum::<f64>()
3998 / (base_norm * base_norm);
3999 direction = &direction - &(overlap * &base);
4000 let direction_norm = direction.iter().map(|v| v * v).sum::<f64>().sqrt();
4001 let perturbed = &base + &(direction * (DEFECT * base_norm / direction_norm));
4002
4003 let bundle = vec![
4004 canonical_from_local(base.clone(), 0..3, 3),
4005 canonical_from_local(perturbed.clone(), 0..3, 3),
4006 ];
4007 assert!(
4008 report_penalty_pair_redundancy(&bundle).is_empty(),
4009 "a pair {DEFECT:.3e} apart is measurably distinct and must not be called proportional"
4010 );
4011
4012 let inner = base
4016 .iter()
4017 .zip(perturbed.iter())
4018 .map(|(a, b)| a * b)
4019 .sum::<f64>();
4020 let cosine =
4021 inner / (base_norm * perturbed.iter().map(|v| v * v).sum::<f64>().sqrt());
4022 assert!(
4023 cosine > 1.0 - 1e-8,
4024 "the fixture must be one the OLD cos > 1 - 1e-8 screen accepted; got cos = {cosine}"
4025 );
4026 assert_eq!(
4027 format!("{cosine:.6}"),
4028 "1.000000",
4029 "and one whose six-decimal print is indistinguishable from an exact identity"
4030 );
4031 }
4032
4033 #[test]
4034 fn report_penalty_pair_redundancy_skips_different_col_ranges() {
4035 let s = ndarray::array![[1.0, 0.0], [0.0, 1.0]];
4039 let bundle = vec![
4040 canonical_from_local(s.clone(), 0..2, 4),
4041 canonical_from_local(s, 2..4, 4),
4042 ];
4043 let redundant = report_penalty_pair_redundancy(&bundle);
4044 assert!(
4045 redundant.is_empty(),
4046 "different col_ranges must not be flagged"
4047 );
4048 }
4049}