1use ndarray::{Array2, ArrayView1};
42
43#[derive(Debug, Clone)]
53pub struct JointPenaltySpec {
54 pub label: Option<String>,
58 pub matrix: Array2<f64>,
60 pub initial_log_lambda: f64,
62 pub nullspace_dim: usize,
64 pub group: Option<usize>,
78}
79
80#[derive(Debug, Clone, PartialEq)]
82pub enum JointPenaltyError {
83 NotSquare {
84 nrows: usize,
85 ncols: usize,
86 },
87 NonFiniteEntry {
88 row: usize,
89 col: usize,
90 value: f64,
91 },
92 InitialLogStrengthOutOfDomain {
93 value: f64,
94 },
95 NotSymmetric {
96 row: usize,
97 col: usize,
98 asymmetry: f64,
99 },
100 NullspaceTooLarge {
101 total: usize,
102 nullspace_dim: usize,
103 },
104 NotPositiveSemidefinite {
105 min_eigenvalue: f64,
106 max_abs_eigenvalue: f64,
107 },
108 NullspaceMismatch {
109 declared: usize,
110 numerical: usize,
111 },
112 EigendecompositionFailed {
113 reason: String,
114 },
115}
116
117impl std::fmt::Display for JointPenaltyError {
118 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
119 match self {
120 Self::NotSquare { nrows, ncols } => {
121 write!(f, "joint penalty matrix is not square: {nrows}x{ncols}")
122 }
123 Self::NonFiniteEntry { row, col, value } => write!(
124 f,
125 "joint penalty matrix has non-finite entry at ({row},{col}): {value}"
126 ),
127 Self::InitialLogStrengthOutOfDomain { value } => {
128 write!(
129 f,
130 "joint penalty initial_log_lambda is outside the exact strength domain: {value}"
131 )
132 }
133 Self::NotSymmetric {
134 row,
135 col,
136 asymmetry,
137 } => write!(
138 f,
139 "joint penalty matrix is not symmetric at ({row},{col}): |S - Sᵀ|={asymmetry:.3e}"
140 ),
141 Self::NullspaceTooLarge {
142 total,
143 nullspace_dim,
144 } => write!(
145 f,
146 "joint penalty nullspace_dim={nullspace_dim} exceeds dim={total}"
147 ),
148 Self::NotPositiveSemidefinite {
149 min_eigenvalue,
150 max_abs_eigenvalue,
151 } => write!(
152 f,
153 "joint penalty matrix is not positive semidefinite: min eigenvalue \
154 {min_eigenvalue:.6e} (max |eigenvalue| {max_abs_eigenvalue:.6e}); the \
155 penalized objective is unbounded below along the negative mode"
156 ),
157 Self::NullspaceMismatch {
158 declared,
159 numerical,
160 } => write!(
161 f,
162 "joint penalty declares nullspace_dim={declared} but the eigenspectrum has \
163 {numerical} numerical-zero direction(s); the REML pseudo-logdet rank would \
164 be wrong"
165 ),
166 Self::EigendecompositionFailed { reason } => write!(
167 f,
168 "joint penalty eigendecomposition failed during validation: {reason}"
169 ),
170 }
171 }
172}
173
174impl std::error::Error for JointPenaltyError {}
175
176impl JointPenaltySpec {
177 const SYMMETRY_TOL: f64 = 1e-10;
181
182 #[inline]
184 pub fn dim(&self) -> usize {
185 self.matrix.nrows()
186 }
187
188 pub fn trace(&self) -> f64 {
190 self.matrix.diag().iter().copied().sum()
191 }
192
193 #[inline]
197 pub fn pseudo_rank(&self) -> usize {
198 self.dim().saturating_sub(self.nullspace_dim)
199 }
200
201 pub fn quadratic_form(&self, beta: ArrayView1<'_, f64>) -> f64 {
205 assert_eq!(
206 beta.len(),
207 self.dim(),
208 "joint penalty quadratic form: beta length {} != dim {}",
209 beta.len(),
210 self.dim()
211 );
212 beta.dot(&self.matrix.dot(&beta))
213 }
214
215 pub fn validated_root(&self) -> Result<Array2<f64>, JointPenaltyError> {
224 let (nrows, ncols) = self.matrix.dim();
225 if nrows != ncols {
226 return Err(JointPenaltyError::NotSquare { nrows, ncols });
227 }
228 if crate::validate_log_strength(self.initial_log_lambda).is_err() {
229 return Err(JointPenaltyError::InitialLogStrengthOutOfDomain {
230 value: self.initial_log_lambda,
231 });
232 }
233 if self.nullspace_dim > nrows {
234 return Err(JointPenaltyError::NullspaceTooLarge {
235 total: nrows,
236 nullspace_dim: self.nullspace_dim,
237 });
238 }
239 for ((row, col), &value) in self.matrix.indexed_iter() {
240 if !value.is_finite() {
241 return Err(JointPenaltyError::NonFiniteEntry { row, col, value });
242 }
243 }
244 for row in 0..nrows {
245 for col in (row + 1)..ncols {
246 let asymmetry = (self.matrix[[row, col]] - self.matrix[[col, row]]).abs();
247 if asymmetry > Self::SYMMETRY_TOL {
248 return Err(JointPenaltyError::NotSymmetric {
249 row,
250 col,
251 asymmetry,
252 });
253 }
254 }
255 }
256 if nrows == 0 {
263 return Ok(Array2::zeros((0, 0)));
264 }
265 use gam_linalg::faer_ndarray::FaerEigh;
266 let (eigenvalues, eigenvectors) =
267 FaerEigh::eigh(&self.matrix, faer::Side::Lower).map_err(|e| {
268 JointPenaltyError::EigendecompositionFailed {
269 reason: e.to_string(),
270 }
271 })?;
272 let max_abs_eigenvalue = eigenvalues
273 .iter()
274 .fold(0.0_f64, |acc, &ev| acc.max(ev.abs()));
275 let tol = 100.0 * (nrows as f64) * f64::EPSILON * max_abs_eigenvalue;
278 if let Some(&min_eigenvalue) = eigenvalues
279 .iter()
280 .filter(|&&ev| ev < -tol)
281 .min_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal))
282 {
283 return Err(JointPenaltyError::NotPositiveSemidefinite {
284 min_eigenvalue,
285 max_abs_eigenvalue,
286 });
287 }
288 let active: Vec<usize> = eigenvalues
289 .iter()
290 .enumerate()
291 .filter_map(|(index, &value)| (value > tol).then_some(index))
292 .collect();
293 let numerical = nrows - active.len();
294 if numerical != self.nullspace_dim {
295 return Err(JointPenaltyError::NullspaceMismatch {
296 declared: self.nullspace_dim,
297 numerical,
298 });
299 }
300 let mut root = Array2::<f64>::zeros((active.len(), nrows));
301 for (root_row, &eigen_index) in active.iter().enumerate() {
302 let scale = eigenvalues[eigen_index].sqrt();
303 for column in 0..nrows {
304 root[[root_row, column]] = scale * eigenvectors[[column, eigen_index]];
305 }
306 }
307 Ok(root)
308 }
309
310 pub fn validate(&self) -> Result<(), JointPenaltyError> {
312 self.validated_root().map(|_| ())
313 }
314}
315
316#[derive(Clone, Debug)]
325pub struct JointPenaltyBundle {
326 specs: std::sync::Arc<Vec<JointPenaltySpec>>,
327 roots: std::sync::Arc<Vec<Array2<f64>>>,
328 log_lambdas: Vec<f64>,
329 lambdas: Vec<f64>,
330}
331
332impl JointPenaltyBundle {
333 pub fn new(
336 specs: std::sync::Arc<Vec<JointPenaltySpec>>,
337 log_lambdas: Vec<f64>,
338 total_compiled: usize,
339 ) -> Result<Self, String> {
340 let roots = specs
341 .iter()
342 .enumerate()
343 .map(|(index, spec)| {
344 spec.validated_root()
345 .map_err(|error| format!("joint penalty {index}: {error}"))
346 })
347 .collect::<Result<Vec<_>, _>>()?;
348 Self::from_validated_geometry(
349 specs,
350 std::sync::Arc::new(roots),
351 log_lambdas,
352 total_compiled,
353 )
354 }
355
356 pub fn from_validated_geometry(
364 specs: std::sync::Arc<Vec<JointPenaltySpec>>,
365 roots: std::sync::Arc<Vec<Array2<f64>>>,
366 log_lambdas: Vec<f64>,
367 total_compiled: usize,
368 ) -> Result<Self, String> {
369 if specs.len() != log_lambdas.len() {
370 return Err(format!(
371 "joint penalty bundle: {} specs vs {} log_lambdas",
372 specs.len(),
373 log_lambdas.len(),
374 ));
375 }
376 if roots.len() != specs.len() {
377 return Err(format!(
378 "joint penalty bundle: {} specs vs {} cached roots",
379 specs.len(),
380 roots.len(),
381 ));
382 }
383 let mut lambdas = Vec::with_capacity(log_lambdas.len());
384 for (i, ((spec, root), &log_lambda)) in specs
385 .iter()
386 .zip(roots.iter())
387 .zip(log_lambdas.iter())
388 .enumerate()
389 {
390 if spec.dim() != total_compiled {
391 return Err(format!(
392 "joint penalty {i}: dim {} != total_compiled {}",
393 spec.dim(),
394 total_compiled,
395 ));
396 }
397 if root.dim() != (spec.pseudo_rank(), total_compiled) {
398 return Err(format!(
399 "joint penalty {i}: cached root shape {}x{} != rank-by-dimension {}x{}",
400 root.nrows(),
401 root.ncols(),
402 spec.pseudo_rank(),
403 total_compiled,
404 ));
405 }
406 if let Some(((row, column), &value)) =
407 root.indexed_iter().find(|(_, value)| !value.is_finite())
408 {
409 return Err(format!(
410 "joint penalty {i}: cached root has non-finite entry at ({row},{column}): {value}"
411 ));
412 }
413 lambdas.push(
414 crate::checked_exp_log_strength(log_lambda)
415 .map_err(|error| format!("joint penalty {i} current log-precision: {error}"))?,
416 );
417 }
418 Ok(Self {
419 specs,
420 roots,
421 log_lambdas,
422 lambdas,
423 })
424 }
425
426 #[inline]
427 pub fn len(&self) -> usize {
428 self.specs.len()
429 }
430
431 #[inline]
432 pub fn is_empty(&self) -> bool {
433 self.specs.is_empty()
434 }
435
436 #[inline]
437 pub fn specs(&self) -> &[JointPenaltySpec] {
438 self.specs.as_slice()
439 }
440
441 #[inline]
442 pub fn roots(&self) -> &[Array2<f64>] {
443 self.roots.as_slice()
444 }
445
446 #[inline]
447 pub fn log_lambdas(&self) -> &[f64] {
448 self.log_lambdas.as_slice()
449 }
450
451 #[inline]
452 pub fn lambdas(&self) -> &[f64] {
453 self.lambdas.as_slice()
454 }
455
456 pub fn quadratic(&self, beta: ArrayView1<'_, f64>) -> f64 {
459 let mut total = 0.0;
460 for (spec, &lam) in self.specs.iter().zip(self.lambdas.iter()) {
461 total += 0.5 * lam * spec.quadratic_form(beta);
462 }
463 total
464 }
465
466 pub fn add_apply_into(&self, vector: ArrayView1<'_, f64>, out: &mut ndarray::Array1<f64>) {
468 assert_eq!(out.len(), vector.len());
469 for (spec, &lam) in self.specs.iter().zip(self.lambdas.iter()) {
470 let sv = spec.matrix.dot(&vector);
471 out.scaled_add(lam, &sv);
472 }
473 }
474
475 pub fn add_diag(&self, diag: &mut ndarray::Array1<f64>) {
477 for (spec, &lam) in self.specs.iter().zip(self.lambdas.iter()) {
478 for (i, value) in spec.matrix.diag().iter().enumerate() {
479 diag[i] += lam * *value;
480 }
481 }
482 }
483
484 pub fn add_to_matrix(&self, matrix: &mut Array2<f64>) {
486 assert_eq!(matrix.nrows(), matrix.ncols());
487 for (spec, &lam) in self.specs.iter().zip(self.lambdas.iter()) {
488 matrix.scaled_add(lam, &spec.matrix);
489 }
490 }
491
492}
493
494#[cfg(test)]
495mod tests {
496 use super::*;
497 use ndarray::{Array1, Array2, array};
498
499 fn cross_block_spec() -> JointPenaltySpec {
503 let v: Array1<f64> = array![1.0, 0.0, -1.0, 0.0];
505 let w: Array1<f64> = array![0.0, 1.0, 0.0, -1.0];
506 let mut matrix: Array2<f64> = Array2::zeros((4, 4));
507 for i in 0..4 {
508 for j in 0..4 {
509 matrix[[i, j]] = v[i] * v[j] + w[i] * w[j];
510 }
511 }
512 JointPenaltySpec {
513 label: Some("cross_block_pullback".to_string()),
514 matrix,
515 initial_log_lambda: -1.5,
516 nullspace_dim: 2,
517 group: None,
518 }
519 }
520
521 #[test]
522 fn cross_block_dense_validates() {
523 let result = cross_block_spec().validate();
524 assert!(
525 result.is_ok(),
526 "valid cross-block spec rejected: {result:?}"
527 );
528 }
529
530 #[test]
531 fn trace_matches_diagonal_sum() {
532 let spec = cross_block_spec();
533 assert!((spec.trace() - 4.0).abs() < 1e-12);
535 }
536
537 #[test]
538 fn pseudo_rank_uses_declared_nullspace() {
539 let spec = cross_block_spec();
540 assert_eq!(spec.dim(), 4);
541 assert_eq!(spec.pseudo_rank(), 2);
542 }
543
544 #[test]
545 fn quadratic_form_matches_explicit_mat_vec() {
546 let spec = cross_block_spec();
547 let beta: Array1<f64> = array![0.5, -0.25, 1.0, 0.75];
549 let q = spec.quadratic_form(beta.view());
552 assert!((q - 1.25).abs() < 1e-12, "got {q}");
553 }
554
555 #[test]
556 fn determinant_zero_for_rank_deficient_matches_nullspace() {
557 use gam_linalg::faer_ndarray::FaerEigh;
558 let spec = cross_block_spec();
559 let (eigvals, _) =
562 FaerEigh::eigh(&spec.matrix, faer::Side::Lower).expect("symmetric eigh succeeds");
563 let mut sorted: Vec<f64> = eigvals.iter().copied().collect();
564 sorted.sort_by(|a, b| a.partial_cmp(b).unwrap());
565 let zeros = sorted.iter().take_while(|&&v| v.abs() < 1e-10).count();
566 assert_eq!(
567 zeros, spec.nullspace_dim,
568 "spectrum {sorted:?} should have {} near-zeros",
569 spec.nullspace_dim
570 );
571 let det: f64 = sorted.iter().product();
574 assert!(det.abs() < 1e-10, "expected ~0 determinant, got {det}");
575 }
576
577 #[test]
578 fn validate_rejects_non_square() {
579 let spec = JointPenaltySpec {
580 label: None,
581 matrix: Array2::zeros((3, 4)),
582 initial_log_lambda: 0.0,
583 nullspace_dim: 0,
584 group: None,
585 };
586 assert!(matches!(
587 spec.validate(),
588 Err(JointPenaltyError::NotSquare { nrows: 3, ncols: 4 })
589 ));
590 }
591
592 #[test]
593 fn validate_rejects_non_symmetric() {
594 let mut matrix = Array2::<f64>::zeros((3, 3));
595 matrix[[0, 1]] = 1.0;
596 matrix[[1, 0]] = -1.0;
597 let spec = JointPenaltySpec {
598 label: None,
599 matrix,
600 initial_log_lambda: 0.0,
601 nullspace_dim: 0,
602 group: None,
603 };
604 assert!(matches!(
605 spec.validate(),
606 Err(JointPenaltyError::NotSymmetric { .. })
607 ));
608 }
609
610 #[test]
611 fn validate_rejects_oversized_nullspace() {
612 let spec = JointPenaltySpec {
613 label: None,
614 matrix: Array2::zeros((3, 3)),
615 initial_log_lambda: 0.0,
616 nullspace_dim: 4,
617 group: None,
618 };
619 assert!(matches!(
620 spec.validate(),
621 Err(JointPenaltyError::NullspaceTooLarge {
622 total: 3,
623 nullspace_dim: 4
624 })
625 ));
626 }
627
628 #[test]
629 fn validate_rejects_initial_log_strength_outside_exact_domain() {
630 let spec = JointPenaltySpec {
631 label: None,
632 matrix: Array2::zeros((2, 2)),
633 initial_log_lambda: f64::NAN,
634 nullspace_dim: 0,
635 group: None,
636 };
637 assert!(matches!(
638 spec.validate(),
639 Err(JointPenaltyError::InitialLogStrengthOutOfDomain { .. })
640 ));
641
642 let mut finite_but_too_large = cross_block_spec();
643 finite_but_too_large.initial_log_lambda = crate::LOG_STRENGTH_MAX + 1.0;
644 assert!(matches!(
645 finite_but_too_large.validate(),
646 Err(JointPenaltyError::InitialLogStrengthOutOfDomain { .. })
647 ));
648 }
649
650 #[test]
651 fn bundle_construction_is_atomic_at_exact_log_strength_faces() {
652 let specs = std::sync::Arc::new(vec![cross_block_spec(), cross_block_spec()]);
653 let bundle = JointPenaltyBundle::new(
654 specs.clone(),
655 vec![crate::LOG_STRENGTH_MIN, crate::LOG_STRENGTH_MAX],
656 4,
657 )
658 .expect("closed endpoints");
659 for ((&actual, &log_strength), expected) in bundle
660 .lambdas()
661 .iter()
662 .zip(bundle.log_lambdas())
663 .zip([crate::LOG_STRENGTH_MIN.exp(), crate::LOG_STRENGTH_MAX.exp()])
664 {
665 assert_eq!(actual.to_bits(), expected.to_bits());
666 assert_eq!(actual.to_bits(), log_strength.exp().to_bits());
667 }
668
669 let error = JointPenaltyBundle::new(specs, vec![0.0, crate::LOG_STRENGTH_MAX + 1.0], 4)
670 .expect_err("one invalid coordinate refuses the whole bundle");
671 assert!(error.contains("joint penalty 1 current log-precision"));
672 }
673
674 #[test]
675 fn bundle_rejects_dim_mismatch() {
676 let spec = JointPenaltySpec {
677 label: None,
678 matrix: Array2::<f64>::eye(3),
679 initial_log_lambda: 0.0,
680 nullspace_dim: 0,
681 group: None,
682 };
683 let err = JointPenaltyBundle::new(std::sync::Arc::new(vec![spec]), vec![0.0], 4)
684 .expect_err("dim mismatch must reject");
685 assert!(err.contains("total_compiled"));
686 }
687
688 #[test]
689 fn bundle_rejects_lambda_count_mismatch() {
690 let spec = JointPenaltySpec {
691 label: None,
692 matrix: Array2::<f64>::eye(2),
693 initial_log_lambda: 0.0,
694 nullspace_dim: 0,
695 group: None,
696 };
697 let err = JointPenaltyBundle::new(std::sync::Arc::new(vec![spec]), vec![], 2)
698 .expect_err("count mismatch must reject");
699 assert!(err.contains("specs vs"));
700 }
701}