gam-problem 0.3.157

//! Joint (cross-block) penalty specifications.
//!
//! After the `T^T S_j T` pullback used by the V+M / SMGS-exact compile path,
//! a single penalty `S_j` no longer has its nonzero region confined to one
//! `ParameterBlockSpec`: the pullback by the inter-block coupling matrix `T`
//! distributes weight across the *entire* compiled parameter vector. The
//! existing `ParameterBlockSpec.penalties: Vec<PenaltyMatrix>` model encodes
//! a per-block-local penalty (its dim equals the owning block's column count),
//! so it cannot represent these full-width operators.
//!
//! [`JointPenaltySpec`] is the carrier for that case: one dense
//! `total_compiled × total_compiled` matrix with its own initial smoothing
//! parameter and structural nullspace dimension. It lives *alongside*, not
//! *inside*, the per-block specs.
//!
//! ## Inner-solve integration
//!
//! `inner_blockwise_fit` and the joint-Newton kernels in `custom_family`
//! consume ordinary block-local penalties as a `&[Array2<f64>]` paired with
//! per-block `(start, end)` ranges:
//!
//! * `apply_joint_block_penalty_into(ranges, s_lambdas, …)` (≈ line 19960)
//! * `joint_penalty_preconditioner_diag(…)` (≈ line 20067)
//! * `add_joint_penalty_to_matrix(matrix, ranges, s_lambdas, …)` (≈ line 20132)
//!
//! A cross-block dense `S` has no single owning block range, so the solver also
//! threads a `JointPenaltyBundle` through those helpers as a full-width path
//! that:
//!
//! 1. computes `S · v` as a full `total × total` mat-vec (cf. `fast_av`),
//! 2. accumulates `diag(S)` into the Jacobi preconditioner over the full
//!    parameter vector, and
//! 3. adds `λ · S` to the dense joint Hessian without slicing.
//!
//! The remaining construction-site work is to produce the correct
//! `JointPenaltySpec` instances for each coupled-family compile path; once a
//! bundle is supplied through `BlockwiseFitOptions::joint_penalties`, the inner
//! solve consumes its objective, mat-vec, preconditioner, and dense-Hessian
//! contributions.

use ndarray::{Array2, ArrayView1};

/// A penalty whose support spans the entire compiled parameter vector.
///
/// Unlike `crate::families::custom_family::PenaltyMatrix`, this carries a
/// single dense `total_compiled × total_compiled` quadratic form — the
/// shape produced by `T^T S_j T` pullback after the V+M / SMGS-exact
/// compile. The `nullspace_dim` is the structural dimension of `ker(S)`
/// as reported by the construction site (rank-revealing on the *pulled-back*
/// operator, not the pre-pullback `S_j`), so the REML pseudo-logdet can
/// avoid numerical rank thresholds.
#[derive(Debug, Clone)]
pub struct JointPenaltySpec {
    /// Optional user-visible precision label. Joint penalties that share a
    /// label share one smoothing parameter (same convention as
    /// `crate::families::custom_family::PenaltyMatrix::Labeled`).
    pub label: Option<String>,
    /// Dense symmetric PSD matrix of shape `(total_compiled, total_compiled)`.
    pub matrix: Array2<f64>,
    /// Initial value of `log λ` for this penalty.
    pub initial_log_lambda: f64,
    /// Structural nullspace dimension of `matrix` (i.e. `total_compiled - rank`).
    pub nullspace_dim: usize,
    /// Optional term grouping, declared by the producing family.
    ///
    /// Specs sharing a group are the SAME smooth term seen through different
    /// class contrasts, so a relabeling permutes them among themselves. Any
    /// consumer that needs a reference-INVARIANT quantity per term must
    /// aggregate over the group rather than read one spec: an individual spec's
    /// matrix is expressed in the stacked ALR basis, whose meaning depends on
    /// which class is the reference (#2579). This is deliberately a declared
    /// integer and not something recovered by parsing [`Self::label`] — a
    /// substring classifier over a formatted name is exactly the failure #2593
    /// was closed for.
    ///
    /// `None` means "stands alone", which is every family that does not group.
    pub group: Option<usize>,
}

/// Reason a [`JointPenaltySpec`] failed validation.
#[derive(Debug, Clone, PartialEq)]
pub enum JointPenaltyError {
    NotSquare {
        nrows: usize,
        ncols: usize,
    },
    NonFiniteEntry {
        row: usize,
        col: usize,
        value: f64,
    },
    InitialLogStrengthOutOfDomain {
        value: f64,
    },
    NotSymmetric {
        row: usize,
        col: usize,
        asymmetry: f64,
    },
    NullspaceTooLarge {
        total: usize,
        nullspace_dim: usize,
    },
    NotPositiveSemidefinite {
        min_eigenvalue: f64,
        max_abs_eigenvalue: f64,
    },
    NullspaceMismatch {
        declared: usize,
        numerical: usize,
    },
    EigendecompositionFailed {
        reason: String,
    },
}

impl std::fmt::Display for JointPenaltyError {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            Self::NotSquare { nrows, ncols } => {
                write!(f, "joint penalty matrix is not square: {nrows}x{ncols}")
            }
            Self::NonFiniteEntry { row, col, value } => write!(
                f,
                "joint penalty matrix has non-finite entry at ({row},{col}): {value}"
            ),
            Self::InitialLogStrengthOutOfDomain { value } => {
                write!(
                    f,
                    "joint penalty initial_log_lambda is outside the exact strength domain: {value}"
                )
            }
            Self::NotSymmetric {
                row,
                col,
                asymmetry,
            } => write!(
                f,
                "joint penalty matrix is not symmetric at ({row},{col}): |S - Sᵀ|={asymmetry:.3e}"
            ),
            Self::NullspaceTooLarge {
                total,
                nullspace_dim,
            } => write!(
                f,
                "joint penalty nullspace_dim={nullspace_dim} exceeds dim={total}"
            ),
            Self::NotPositiveSemidefinite {
                min_eigenvalue,
                max_abs_eigenvalue,
            } => write!(
                f,
                "joint penalty matrix is not positive semidefinite: min eigenvalue \
                 {min_eigenvalue:.6e} (max |eigenvalue| {max_abs_eigenvalue:.6e}); the \
                 penalized objective is unbounded below along the negative mode"
            ),
            Self::NullspaceMismatch {
                declared,
                numerical,
            } => write!(
                f,
                "joint penalty declares nullspace_dim={declared} but the eigenspectrum has \
                 {numerical} numerical-zero direction(s); the REML pseudo-logdet rank would \
                 be wrong"
            ),
            Self::EigendecompositionFailed { reason } => write!(
                f,
                "joint penalty eigendecomposition failed during validation: {reason}"
            ),
        }
    }
}

impl std::error::Error for JointPenaltyError {}

impl JointPenaltySpec {
    /// Symmetry tolerance for [`validate`]. Cross-block pullbacks via `T`
    /// accumulate roundoff, so an exact symmetric requirement is too tight;
    /// this matches the floor used by the surrounding penalty code paths.
    const SYMMETRY_TOL: f64 = 1e-10;

    /// Total compiled parameter count this penalty acts on.
    #[inline]
    pub fn dim(&self) -> usize {
        self.matrix.nrows()
    }

    /// Trace of the penalty matrix (`Σ_i S[i,i]`).
    pub fn trace(&self) -> f64 {
        self.matrix.diag().iter().copied().sum()
    }

    /// Structural pseudo-rank, derived from the declared `nullspace_dim`.
    /// This is the rank used by the REML pseudo-logdet under the
    /// no-numerical-thresholds policy in the surrounding code.
    #[inline]
    pub fn pseudo_rank(&self) -> usize {
        self.dim().saturating_sub(self.nullspace_dim)
    }

    /// Quadratic form `βᵀ S β`. Mirrors
    /// `crate::families::custom_family::PenaltyMatrix::quadratic_form` for
    /// the full-width case.
    pub fn quadratic_form(&self, beta: ArrayView1<'_, f64>) -> f64 {
        assert_eq!(
            beta.len(),
            self.dim(),
            "joint penalty quadratic form: beta length {} != dim {}",
            beta.len(),
            self.dim()
        );
        beta.dot(&self.matrix.dot(&beta))
    }

    /// Validate shape, finiteness, symmetry, PSD, and nullspace bookkeeping,
    /// returning the invariant thin root `R` such that `matrix = RᵀR`.
    ///
    /// Joint-penalty strengths vary throughout an outer optimization, but the
    /// component matrices do not. Returning the root from the same
    /// eigendecomposition that validates the component lets callers retain it
    /// as penalty geometry instead of repeating one O(p³) decomposition per
    /// component on every objective evaluation.
    pub fn validated_root(&self) -> Result<Array2<f64>, JointPenaltyError> {
        let (nrows, ncols) = self.matrix.dim();
        if nrows != ncols {
            return Err(JointPenaltyError::NotSquare { nrows, ncols });
        }
        if crate::validate_log_strength(self.initial_log_lambda).is_err() {
            return Err(JointPenaltyError::InitialLogStrengthOutOfDomain {
                value: self.initial_log_lambda,
            });
        }
        if self.nullspace_dim > nrows {
            return Err(JointPenaltyError::NullspaceTooLarge {
                total: nrows,
                nullspace_dim: self.nullspace_dim,
            });
        }
        for ((row, col), &value) in self.matrix.indexed_iter() {
            if !value.is_finite() {
                return Err(JointPenaltyError::NonFiniteEntry { row, col, value });
            }
        }
        for row in 0..nrows {
            for col in (row + 1)..ncols {
                let asymmetry = (self.matrix[[row, col]] - self.matrix[[col, row]]).abs();
                if asymmetry > Self::SYMMETRY_TOL {
                    return Err(JointPenaltyError::NotSymmetric {
                        row,
                        col,
                        asymmetry,
                    });
                }
            }
        }
        // PSD + declared-nullity honesty. An indefinite joint penalty makes
        // the penalized objective unbounded below along its negative mode
        // while the pseudo-logdet's positive-eigenspace filter would silently
        // drop that mode; a wrong declared nullity mis-ranks the REML
        // pseudo-logdet (the whole point of declaring it is to avoid runtime
        // thresholds, so it must agree with the spectrum at construction).
        if nrows == 0 {
            return Ok(Array2::zeros((0, 0)));
        }
        use gam_linalg::faer_ndarray::FaerEigh;
        let (eigenvalues, eigenvectors) =
            FaerEigh::eigh(&self.matrix, faer::Side::Lower).map_err(|e| {
                JointPenaltyError::EigendecompositionFailed {
                    reason: e.to_string(),
                }
            })?;
        let max_abs_eigenvalue = eigenvalues
            .iter()
            .fold(0.0_f64, |acc, &ev| acc.max(ev.abs()));
        // Same relative classification as the REML pseudo-logdet kernel:
        // the eigensolver noise floor is O(p·ε·‖S‖), never an absolute cut.
        let tol = 100.0 * (nrows as f64) * f64::EPSILON * max_abs_eigenvalue;
        if let Some(&min_eigenvalue) = eigenvalues
            .iter()
            .filter(|&&ev| ev < -tol)
            .min_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal))
        {
            return Err(JointPenaltyError::NotPositiveSemidefinite {
                min_eigenvalue,
                max_abs_eigenvalue,
            });
        }
        let active: Vec<usize> = eigenvalues
            .iter()
            .enumerate()
            .filter_map(|(index, &value)| (value > tol).then_some(index))
            .collect();
        let numerical = nrows - active.len();
        if numerical != self.nullspace_dim {
            return Err(JointPenaltyError::NullspaceMismatch {
                declared: self.nullspace_dim,
                numerical,
            });
        }
        let mut root = Array2::<f64>::zeros((active.len(), nrows));
        for (root_row, &eigen_index) in active.iter().enumerate() {
            let scale = eigenvalues[eigen_index].sqrt();
            for column in 0..nrows {
                root[[root_row, column]] = scale * eigenvectors[[column, eigen_index]];
            }
        }
        Ok(root)
    }

    /// Validate this joint penalty without retaining its spectral root.
    pub fn validate(&self) -> Result<(), JointPenaltyError> {
        self.validated_root().map(|_| ())
    }
}

/// Per-evaluation bundle of cross-block penalties paired with their current
/// log-smoothing parameters.
///
/// The outer optimizer concatenates joint penalty `log λ` values onto the
/// per-block ρ vector; the inner solver receives this bundle via
/// `crate::families::custom_family::BlockwiseFitOptions::joint_penalties`
/// and adds the full-width quadratic / matvec / preconditioner / Hessian
/// contributions to the joint-Newton primitives.
#[derive(Clone, Debug)]
pub struct JointPenaltyBundle {
    specs: std::sync::Arc<Vec<JointPenaltySpec>>,
    roots: std::sync::Arc<Vec<Array2<f64>>>,
    log_lambdas: Vec<f64>,
    lambdas: Vec<f64>,
}

impl JointPenaltyBundle {
    /// Build a bundle, validating the per-penalty `log λ` count and dimension
    /// agreement against `total_compiled`.
    pub fn new(
        specs: std::sync::Arc<Vec<JointPenaltySpec>>,
        log_lambdas: Vec<f64>,
        total_compiled: usize,
    ) -> Result<Self, String> {
        let roots = specs
            .iter()
            .enumerate()
            .map(|(index, spec)| {
                spec.validated_root()
                    .map_err(|error| format!("joint penalty {index}: {error}"))
            })
            .collect::<Result<Vec<_>, _>>()?;
        Self::from_validated_geometry(
            specs,
            std::sync::Arc::new(roots),
            log_lambdas,
            total_compiled,
        )
    }

    /// Build a rho-specific bundle from already-validated invariant geometry.
    ///
    /// `specs` and `roots` are constructed together by the label-layout
    /// compiler and retained across every outer evaluation. Only
    /// `log_lambdas` changes here. Shape and finiteness are still checked at
    /// this boundary; the expensive spectral identity was certified when the
    /// roots were created.
    pub fn from_validated_geometry(
        specs: std::sync::Arc<Vec<JointPenaltySpec>>,
        roots: std::sync::Arc<Vec<Array2<f64>>>,
        log_lambdas: Vec<f64>,
        total_compiled: usize,
    ) -> Result<Self, String> {
        if specs.len() != log_lambdas.len() {
            return Err(format!(
                "joint penalty bundle: {} specs vs {} log_lambdas",
                specs.len(),
                log_lambdas.len(),
            ));
        }
        if roots.len() != specs.len() {
            return Err(format!(
                "joint penalty bundle: {} specs vs {} cached roots",
                specs.len(),
                roots.len(),
            ));
        }
        let mut lambdas = Vec::with_capacity(log_lambdas.len());
        for (i, ((spec, root), &log_lambda)) in specs
            .iter()
            .zip(roots.iter())
            .zip(log_lambdas.iter())
            .enumerate()
        {
            if spec.dim() != total_compiled {
                return Err(format!(
                    "joint penalty {i}: dim {} != total_compiled {}",
                    spec.dim(),
                    total_compiled,
                ));
            }
            if root.dim() != (spec.pseudo_rank(), total_compiled) {
                return Err(format!(
                    "joint penalty {i}: cached root shape {}x{} != rank-by-dimension {}x{}",
                    root.nrows(),
                    root.ncols(),
                    spec.pseudo_rank(),
                    total_compiled,
                ));
            }
            if let Some(((row, column), &value)) =
                root.indexed_iter().find(|(_, value)| !value.is_finite())
            {
                return Err(format!(
                    "joint penalty {i}: cached root has non-finite entry at ({row},{column}): {value}"
                ));
            }
            lambdas.push(
                crate::checked_exp_log_strength(log_lambda)
                    .map_err(|error| format!("joint penalty {i} current log-precision: {error}"))?,
            );
        }
        Ok(Self {
            specs,
            roots,
            log_lambdas,
            lambdas,
        })
    }

    #[inline]
    pub fn len(&self) -> usize {
        self.specs.len()
    }

    #[inline]
    pub fn is_empty(&self) -> bool {
        self.specs.is_empty()
    }

    #[inline]
    pub fn specs(&self) -> &[JointPenaltySpec] {
        self.specs.as_slice()
    }

    #[inline]
    pub fn roots(&self) -> &[Array2<f64>] {
        self.roots.as_slice()
    }

    #[inline]
    pub fn log_lambdas(&self) -> &[f64] {
        self.log_lambdas.as_slice()
    }

    #[inline]
    pub fn lambdas(&self) -> &[f64] {
        self.lambdas.as_slice()
    }

    /// Total joint-penalty contribution to the objective:
    ///   `½ Σ_j exp(ρ_j) · βᵀ S_j β`.
    pub fn quadratic(&self, beta: ArrayView1<'_, f64>) -> f64 {
        let mut total = 0.0;
        for (spec, &lam) in self.specs.iter().zip(self.lambdas.iter()) {
            total += 0.5 * lam * spec.quadratic_form(beta);
        }
        total
    }

    /// Accumulate `Σ_j exp(ρ_j) · S_j · v` into `out` (additive).
    pub fn add_apply_into(&self, vector: ArrayView1<'_, f64>, out: &mut ndarray::Array1<f64>) {
        assert_eq!(out.len(), vector.len());
        for (spec, &lam) in self.specs.iter().zip(self.lambdas.iter()) {
            let sv = spec.matrix.dot(&vector);
            out.scaled_add(lam, &sv);
        }
    }

    /// Accumulate `Σ_j exp(ρ_j) · diag(S_j)` into `diag` (additive).
    pub fn add_diag(&self, diag: &mut ndarray::Array1<f64>) {
        for (spec, &lam) in self.specs.iter().zip(self.lambdas.iter()) {
            for (i, value) in spec.matrix.diag().iter().enumerate() {
                diag[i] += lam * *value;
            }
        }
    }

    /// Accumulate `Σ_j exp(ρ_j) · S_j` into the full `matrix` (additive).
    pub fn add_to_matrix(&self, matrix: &mut Array2<f64>) {
        assert_eq!(matrix.nrows(), matrix.ncols());
        for (spec, &lam) in self.specs.iter().zip(self.lambdas.iter()) {
            matrix.scaled_add(lam, &spec.matrix);
        }
    }

}

#[cfg(test)]
mod tests {
    use super::*;
    use ndarray::{Array1, Array2, array};

    /// 4-dim cross-block dense penalty: a rank-2 operator that couples
    /// indices {0,1} to {2,3} (i.e. nonzero off the 2×2 block diagonal),
    /// which is exactly the shape that defeats a per-block `PenaltyMatrix`.
    fn cross_block_spec() -> JointPenaltySpec {
        // Build S = vᵀv + wᵀw where v and w span across both 2-blocks.
        let v: Array1<f64> = array![1.0, 0.0, -1.0, 0.0];
        let w: Array1<f64> = array![0.0, 1.0, 0.0, -1.0];
        let mut matrix: Array2<f64> = Array2::zeros((4, 4));
        for i in 0..4 {
            for j in 0..4 {
                matrix[[i, j]] = v[i] * v[j] + w[i] * w[j];
            }
        }
        JointPenaltySpec {
            label: Some("cross_block_pullback".to_string()),
            matrix,
            initial_log_lambda: -1.5,
            nullspace_dim: 2,
            group: None,
        }
    }

    #[test]
    fn cross_block_dense_validates() {
        let result = cross_block_spec().validate();
        assert!(
            result.is_ok(),
            "valid cross-block spec rejected: {result:?}"
        );
    }

    #[test]
    fn trace_matches_diagonal_sum() {
        let spec = cross_block_spec();
        // diag(S) = [v0^2+w0^2, v1^2+w1^2, v2^2+w2^2, v3^2+w3^2] = [1,1,1,1]
        assert!((spec.trace() - 4.0).abs() < 1e-12);
    }

    #[test]
    fn pseudo_rank_uses_declared_nullspace() {
        let spec = cross_block_spec();
        assert_eq!(spec.dim(), 4);
        assert_eq!(spec.pseudo_rank(), 2);
    }

    #[test]
    fn quadratic_form_matches_explicit_mat_vec() {
        let spec = cross_block_spec();
        // Pick a beta that has support in both 2-blocks.
        let beta: Array1<f64> = array![0.5, -0.25, 1.0, 0.75];
        // v·β = 0.5 - 1.0 = -0.5; w·β = -0.25 - 0.75 = -1.0
        // βᵀSβ = (v·β)^2 + (w·β)^2 = 0.25 + 1.0 = 1.25
        let q = spec.quadratic_form(beta.view());
        assert!((q - 1.25).abs() < 1e-12, "got {q}");
    }

    #[test]
    fn determinant_zero_for_rank_deficient_matches_nullspace() {
        use gam_linalg::faer_ndarray::FaerEigh;
        let spec = cross_block_spec();
        // Symmetric eigendecomposition; expect exactly nullspace_dim
        // zeros (up to floating-point), matching the declared rank.
        let (eigvals, _) =
            FaerEigh::eigh(&spec.matrix, faer::Side::Lower).expect("symmetric eigh succeeds");
        let mut sorted: Vec<f64> = eigvals.iter().copied().collect();
        sorted.sort_by(|a, b| a.partial_cmp(b).unwrap());
        let zeros = sorted.iter().take_while(|&&v| v.abs() < 1e-10).count();
        assert_eq!(
            zeros, spec.nullspace_dim,
            "spectrum {sorted:?} should have {} near-zeros",
            spec.nullspace_dim
        );
        // Determinant = product of eigenvalues; with a real nullspace
        // it is exactly zero modulo roundoff.
        let det: f64 = sorted.iter().product();
        assert!(det.abs() < 1e-10, "expected ~0 determinant, got {det}");
    }

    #[test]
    fn validate_rejects_non_square() {
        let spec = JointPenaltySpec {
            label: None,
            matrix: Array2::zeros((3, 4)),
            initial_log_lambda: 0.0,
            nullspace_dim: 0,
            group: None,
        };
        assert!(matches!(
            spec.validate(),
            Err(JointPenaltyError::NotSquare { nrows: 3, ncols: 4 })
        ));
    }

    #[test]
    fn validate_rejects_non_symmetric() {
        let mut matrix = Array2::<f64>::zeros((3, 3));
        matrix[[0, 1]] = 1.0;
        matrix[[1, 0]] = -1.0;
        let spec = JointPenaltySpec {
            label: None,
            matrix,
            initial_log_lambda: 0.0,
            nullspace_dim: 0,
            group: None,
        };
        assert!(matches!(
            spec.validate(),
            Err(JointPenaltyError::NotSymmetric { .. })
        ));
    }

    #[test]
    fn validate_rejects_oversized_nullspace() {
        let spec = JointPenaltySpec {
            label: None,
            matrix: Array2::zeros((3, 3)),
            initial_log_lambda: 0.0,
            nullspace_dim: 4,
            group: None,
        };
        assert!(matches!(
            spec.validate(),
            Err(JointPenaltyError::NullspaceTooLarge {
                total: 3,
                nullspace_dim: 4
            })
        ));
    }

    #[test]
    fn validate_rejects_initial_log_strength_outside_exact_domain() {
        let spec = JointPenaltySpec {
            label: None,
            matrix: Array2::zeros((2, 2)),
            initial_log_lambda: f64::NAN,
            nullspace_dim: 0,
            group: None,
        };
        assert!(matches!(
            spec.validate(),
            Err(JointPenaltyError::InitialLogStrengthOutOfDomain { .. })
        ));

        let mut finite_but_too_large = cross_block_spec();
        finite_but_too_large.initial_log_lambda = crate::LOG_STRENGTH_MAX + 1.0;
        assert!(matches!(
            finite_but_too_large.validate(),
            Err(JointPenaltyError::InitialLogStrengthOutOfDomain { .. })
        ));
    }

    #[test]
    fn bundle_construction_is_atomic_at_exact_log_strength_faces() {
        let specs = std::sync::Arc::new(vec![cross_block_spec(), cross_block_spec()]);
        let bundle = JointPenaltyBundle::new(
            specs.clone(),
            vec![crate::LOG_STRENGTH_MIN, crate::LOG_STRENGTH_MAX],
            4,
        )
        .expect("closed endpoints");
        for ((&actual, &log_strength), expected) in bundle
            .lambdas()
            .iter()
            .zip(bundle.log_lambdas())
            .zip([crate::LOG_STRENGTH_MIN.exp(), crate::LOG_STRENGTH_MAX.exp()])
        {
            assert_eq!(actual.to_bits(), expected.to_bits());
            assert_eq!(actual.to_bits(), log_strength.exp().to_bits());
        }

        let error = JointPenaltyBundle::new(specs, vec![0.0, crate::LOG_STRENGTH_MAX + 1.0], 4)
            .expect_err("one invalid coordinate refuses the whole bundle");
        assert!(error.contains("joint penalty 1 current log-precision"));
    }

    #[test]
    fn bundle_rejects_dim_mismatch() {
        let spec = JointPenaltySpec {
            label: None,
            matrix: Array2::<f64>::eye(3),
            initial_log_lambda: 0.0,
            nullspace_dim: 0,
            group: None,
        };
        let err = JointPenaltyBundle::new(std::sync::Arc::new(vec![spec]), vec![0.0], 4)
            .expect_err("dim mismatch must reject");
        assert!(err.contains("total_compiled"));
    }

    #[test]
    fn bundle_rejects_lambda_count_mismatch() {
        let spec = JointPenaltySpec {
            label: None,
            matrix: Array2::<f64>::eye(2),
            initial_log_lambda: 0.0,
            nullspace_dim: 0,
            group: None,
        };
        let err = JointPenaltyBundle::new(std::sync::Arc::new(vec![spec]), vec![], 2)
            .expect_err("count mismatch must reject");
        assert!(err.contains("specs vs"));
    }
}