Skip to main content

gam_problem/
joint_penalty.rs

1//! Joint (cross-block) penalty specifications.
2//!
3//! After the `T^T S_j T` pullback used by the V+M / SMGS-exact compile path,
4//! a single penalty `S_j` no longer has its nonzero region confined to one
5//! `ParameterBlockSpec`: the pullback by the inter-block coupling matrix `T`
6//! distributes weight across the *entire* compiled parameter vector. The
7//! existing `ParameterBlockSpec.penalties: Vec<PenaltyMatrix>` model encodes
8//! a per-block-local penalty (its dim equals the owning block's column count),
9//! so it cannot represent these full-width operators.
10//!
11//! [`JointPenaltySpec`] is the carrier for that case: one dense
12//! `total_compiled × total_compiled` matrix with its own initial smoothing
13//! parameter and structural nullspace dimension. It lives *alongside*, not
14//! *inside*, the per-block specs.
15//!
16//! ## Inner-solve integration
17//!
18//! `inner_blockwise_fit` and the joint-Newton kernels in `custom_family`
19//! consume ordinary block-local penalties as a `&[Array2<f64>]` paired with
20//! per-block `(start, end)` ranges:
21//!
22//! * `apply_joint_block_penalty_into(ranges, s_lambdas, …)` (≈ line 19960)
23//! * `joint_penalty_preconditioner_diag(…)` (≈ line 20067)
24//! * `add_joint_penalty_to_matrix(matrix, ranges, s_lambdas, …)` (≈ line 20132)
25//!
26//! A cross-block dense `S` has no single owning block range, so the solver also
27//! threads a `JointPenaltyBundle` through those helpers as a full-width path
28//! that:
29//!
30//! 1. computes `S · v` as a full `total × total` mat-vec (cf. `fast_av`),
31//! 2. accumulates `diag(S)` into the Jacobi preconditioner over the full
32//!    parameter vector, and
33//! 3. adds `λ · S` to the dense joint Hessian without slicing.
34//!
35//! The remaining construction-site work is to produce the correct
36//! `JointPenaltySpec` instances for each coupled-family compile path; once a
37//! bundle is supplied through `BlockwiseFitOptions::joint_penalties`, the inner
38//! solve consumes its objective, mat-vec, preconditioner, and dense-Hessian
39//! contributions.
40
41use ndarray::{Array2, ArrayView1};
42
43/// A penalty whose support spans the entire compiled parameter vector.
44///
45/// Unlike `crate::families::custom_family::PenaltyMatrix`, this carries a
46/// single dense `total_compiled × total_compiled` quadratic form — the
47/// shape produced by `T^T S_j T` pullback after the V+M / SMGS-exact
48/// compile. The `nullspace_dim` is the structural dimension of `ker(S)`
49/// as reported by the construction site (rank-revealing on the *pulled-back*
50/// operator, not the pre-pullback `S_j`), so the REML pseudo-logdet can
51/// avoid numerical rank thresholds.
52#[derive(Debug, Clone)]
53pub struct JointPenaltySpec {
54    /// Optional user-visible precision label. Joint penalties that share a
55    /// label share one smoothing parameter (same convention as
56    /// `crate::families::custom_family::PenaltyMatrix::Labeled`).
57    pub label: Option<String>,
58    /// Dense symmetric PSD matrix of shape `(total_compiled, total_compiled)`.
59    pub matrix: Array2<f64>,
60    /// Initial value of `log λ` for this penalty.
61    pub initial_log_lambda: f64,
62    /// Structural nullspace dimension of `matrix` (i.e. `total_compiled - rank`).
63    pub nullspace_dim: usize,
64    /// Optional term grouping, declared by the producing family.
65    ///
66    /// Specs sharing a group are the SAME smooth term seen through different
67    /// class contrasts, so a relabeling permutes them among themselves. Any
68    /// consumer that needs a reference-INVARIANT quantity per term must
69    /// aggregate over the group rather than read one spec: an individual spec's
70    /// matrix is expressed in the stacked ALR basis, whose meaning depends on
71    /// which class is the reference (#2579). This is deliberately a declared
72    /// integer and not something recovered by parsing [`Self::label`] — a
73    /// substring classifier over a formatted name is exactly the failure #2593
74    /// was closed for.
75    ///
76    /// `None` means "stands alone", which is every family that does not group.
77    pub group: Option<usize>,
78}
79
80/// Reason a [`JointPenaltySpec`] failed validation.
81#[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    /// Symmetry tolerance for [`validate`]. Cross-block pullbacks via `T`
178    /// accumulate roundoff, so an exact symmetric requirement is too tight;
179    /// this matches the floor used by the surrounding penalty code paths.
180    const SYMMETRY_TOL: f64 = 1e-10;
181
182    /// Total compiled parameter count this penalty acts on.
183    #[inline]
184    pub fn dim(&self) -> usize {
185        self.matrix.nrows()
186    }
187
188    /// Trace of the penalty matrix (`Σ_i S[i,i]`).
189    pub fn trace(&self) -> f64 {
190        self.matrix.diag().iter().copied().sum()
191    }
192
193    /// Structural pseudo-rank, derived from the declared `nullspace_dim`.
194    /// This is the rank used by the REML pseudo-logdet under the
195    /// no-numerical-thresholds policy in the surrounding code.
196    #[inline]
197    pub fn pseudo_rank(&self) -> usize {
198        self.dim().saturating_sub(self.nullspace_dim)
199    }
200
201    /// Quadratic form `βᵀ S β`. Mirrors
202    /// `crate::families::custom_family::PenaltyMatrix::quadratic_form` for
203    /// the full-width case.
204    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    /// Validate shape, finiteness, symmetry, PSD, and nullspace bookkeeping,
216    /// returning the invariant thin root `R` such that `matrix = RᵀR`.
217    ///
218    /// Joint-penalty strengths vary throughout an outer optimization, but the
219    /// component matrices do not. Returning the root from the same
220    /// eigendecomposition that validates the component lets callers retain it
221    /// as penalty geometry instead of repeating one O(p³) decomposition per
222    /// component on every objective evaluation.
223    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        // PSD + declared-nullity honesty. An indefinite joint penalty makes
257        // the penalized objective unbounded below along its negative mode
258        // while the pseudo-logdet's positive-eigenspace filter would silently
259        // drop that mode; a wrong declared nullity mis-ranks the REML
260        // pseudo-logdet (the whole point of declaring it is to avoid runtime
261        // thresholds, so it must agree with the spectrum at construction).
262        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        // Same relative classification as the REML pseudo-logdet kernel:
276        // the eigensolver noise floor is O(p·ε·‖S‖), never an absolute cut.
277        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    /// Validate this joint penalty without retaining its spectral root.
311    pub fn validate(&self) -> Result<(), JointPenaltyError> {
312        self.validated_root().map(|_| ())
313    }
314}
315
316/// Per-evaluation bundle of cross-block penalties paired with their current
317/// log-smoothing parameters.
318///
319/// The outer optimizer concatenates joint penalty `log λ` values onto the
320/// per-block ρ vector; the inner solver receives this bundle via
321/// `crate::families::custom_family::BlockwiseFitOptions::joint_penalties`
322/// and adds the full-width quadratic / matvec / preconditioner / Hessian
323/// contributions to the joint-Newton primitives.
324#[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    /// Build a bundle, validating the per-penalty `log λ` count and dimension
334    /// agreement against `total_compiled`.
335    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    /// Build a rho-specific bundle from already-validated invariant geometry.
357    ///
358    /// `specs` and `roots` are constructed together by the label-layout
359    /// compiler and retained across every outer evaluation. Only
360    /// `log_lambdas` changes here. Shape and finiteness are still checked at
361    /// this boundary; the expensive spectral identity was certified when the
362    /// roots were created.
363    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    /// Total joint-penalty contribution to the objective:
457    ///   `½ Σ_j exp(ρ_j) · βᵀ S_j β`.
458    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    /// Accumulate `Σ_j exp(ρ_j) · S_j · v` into `out` (additive).
467    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    /// Accumulate `Σ_j exp(ρ_j) · diag(S_j)` into `diag` (additive).
476    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    /// Accumulate `Σ_j exp(ρ_j) · S_j` into the full `matrix` (additive).
485    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    /// 4-dim cross-block dense penalty: a rank-2 operator that couples
500    /// indices {0,1} to {2,3} (i.e. nonzero off the 2×2 block diagonal),
501    /// which is exactly the shape that defeats a per-block `PenaltyMatrix`.
502    fn cross_block_spec() -> JointPenaltySpec {
503        // Build S = vᵀv + wᵀw where v and w span across both 2-blocks.
504        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        // diag(S) = [v0^2+w0^2, v1^2+w1^2, v2^2+w2^2, v3^2+w3^2] = [1,1,1,1]
534        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        // Pick a beta that has support in both 2-blocks.
548        let beta: Array1<f64> = array![0.5, -0.25, 1.0, 0.75];
549        // v·β = 0.5 - 1.0 = -0.5; w·β = -0.25 - 0.75 = -1.0
550        // βᵀSβ = (v·β)^2 + (w·β)^2 = 0.25 + 1.0 = 1.25
551        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        // Symmetric eigendecomposition; expect exactly nullspace_dim
560        // zeros (up to floating-point), matching the declared rank.
561        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        // Determinant = product of eigenvalues; with a real nullspace
572        // it is exactly zero modulo roundoff.
573        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}