use ndarray::{Array1, Array2, Axis};
use rayon::iter::{IndexedParallelIterator, IntoParallelIterator, ParallelIterator};
#[derive(Clone, Debug)]
pub enum PenaltyMatrix {
Dense(Array2<f64>),
Diagonal(Array1<f64>),
KroneckerFactored {
left: Array2<f64>,
right: Array2<f64>,
},
Blockwise {
local: Array2<f64>,
col_range: std::ops::Range<usize>,
total_dim: usize,
},
Labeled {
label: String,
inner: Box<PenaltyMatrix>,
},
Fixed {
log_lambda: f64,
inner: Box<PenaltyMatrix>,
},
}
impl PenaltyMatrix {
pub fn dim(&self) -> usize {
match self {
Self::Dense(m) => m.nrows(),
Self::Diagonal(diagonal) => diagonal.len(),
Self::KroneckerFactored { left, right } => left.nrows() * right.nrows(),
Self::Blockwise { total_dim, .. } => *total_dim,
Self::Labeled { inner, .. } | Self::Fixed { inner, .. } => inner.dim(),
}
}
pub fn shape(&self) -> (usize, usize) {
match self {
Self::Dense(m) => m.dim(),
Self::Diagonal(diagonal) => (diagonal.len(), diagonal.len()),
Self::KroneckerFactored { left, right } => {
(left.nrows() * right.nrows(), left.ncols() * right.ncols())
}
Self::Blockwise { total_dim, .. } => (*total_dim, *total_dim),
Self::Labeled { inner, .. } | Self::Fixed { inner, .. } => inner.shape(),
}
}
pub fn validate(&self, expected_dim: usize) -> Result<(), String> {
let (nrows, ncols) = self.shape();
if nrows != ncols || nrows != expected_dim {
return Err(format!(
"penalty must be {expected_dim}x{expected_dim}, got {nrows}x{ncols}"
));
}
match self {
Self::Dense(m) => validate_symmetric_psd_core(m, "dense penalty"),
Self::Diagonal(diagonal) => {
let max_abs = diagonal
.iter()
.try_fold(0.0_f64, |scale, &value| {
value
.is_finite()
.then_some(scale.max(value.abs()))
.ok_or_else(|| {
format!("diagonal penalty has non-finite entry: {value}")
})
})?;
let tolerance =
100.0 * diagonal.len() as f64 * f64::EPSILON * max_abs;
if let Some((index, value)) = diagonal
.iter()
.copied()
.enumerate()
.find(|(_, value)| *value < -tolerance)
{
return Err(format!(
"diagonal penalty is not positive semidefinite: entry {index} is {value:.6e}"
));
}
Ok(())
}
Self::KroneckerFactored { left, right } => {
validate_symmetric_psd_core(left, "Kronecker left factor")?;
validate_symmetric_psd_core(right, "Kronecker right factor")
}
Self::Blockwise {
local,
col_range,
total_dim,
} => {
if col_range.end > *total_dim || col_range.len() != local.nrows() {
return Err(format!(
"blockwise penalty embedding is inconsistent: local {}x{} at columns \
{}..{} of total_dim {}",
local.nrows(),
local.ncols(),
col_range.start,
col_range.end,
total_dim
));
}
validate_symmetric_psd_core(local, "blockwise local penalty")
}
Self::Labeled { inner, .. } => inner.validate(expected_dim),
Self::Fixed { log_lambda, inner } => {
crate::validate_log_strength(*log_lambda)
.map_err(|error| format!("fixed penalty log-precision: {error}"))?;
inner.validate(expected_dim)
}
}
}
pub fn to_dense(&self) -> Array2<f64> {
match self {
Self::Dense(m) => m.clone(),
Self::Diagonal(diagonal) => Array2::from_diag(diagonal),
Self::KroneckerFactored { left, right } => kronecker_product(left, right),
Self::Blockwise {
local,
col_range,
total_dim,
} => {
let mut g = Array2::zeros((*total_dim, *total_dim));
g.slice_mut(ndarray::s![
col_range.start..col_range.end,
col_range.start..col_range.end
])
.assign(local);
g
}
Self::Labeled { inner, .. } | Self::Fixed { inner, .. } => inner.to_dense(),
}
}
pub fn as_dense_cow(&self) -> std::borrow::Cow<'_, Array2<f64>> {
match self {
Self::Dense(m) => std::borrow::Cow::Borrowed(m),
Self::Diagonal(_)
| Self::KroneckerFactored { .. }
| Self::Blockwise { .. }
| Self::Labeled { .. }
| Self::Fixed { .. } => std::borrow::Cow::Owned(self.to_dense()),
}
}
pub fn as_dense_ref(&self) -> Option<&Array2<f64>> {
match self {
Self::Dense(m) => Some(m),
Self::Fixed { inner, .. } => inner.as_dense_ref(),
Self::Diagonal(_)
| Self::KroneckerFactored { .. }
| Self::Blockwise { .. }
| Self::Labeled { .. } => None,
}
}
pub fn with_precision_label(self, label: impl Into<String>) -> Self {
Self::Labeled {
label: label.into(),
inner: Box::new(self),
}
}
pub fn precision_label(&self) -> Option<&str> {
match self {
Self::Labeled { label, .. } => Some(label.as_str()),
Self::Fixed { .. } => None,
_ => None,
}
}
pub fn with_fixed_log_lambda(self, log_lambda: f64) -> Self {
Self::Fixed {
log_lambda,
inner: Box::new(self),
}
}
pub fn fixed_log_lambda(&self) -> Option<f64> {
match self {
Self::Fixed { log_lambda, .. } => Some(*log_lambda),
Self::Labeled { inner, .. } => inner.fixed_log_lambda(),
_ => None,
}
}
pub fn dot(&self, v: &Array1<f64>) -> Array1<f64> {
match self {
Self::Dense(m) => m.dot(v),
Self::Diagonal(diagonal) => diagonal * v,
Self::KroneckerFactored { left, right } => {
let p_left = left.nrows();
let p_right = right.nrows();
let v_mat = ndarray::ArrayView2::from_shape(
(p_left, p_right),
v.as_slice()
.expect("penalty operand is a contiguous Array1"),
)
.expect("operand length equals p_left * p_right for this Kronecker factorization");
let avbt = left.dot(&v_mat).dot(&right.t());
let standard = avbt.as_standard_layout();
Array1::from_iter(standard.iter().copied())
}
Self::Blockwise {
local,
col_range,
total_dim,
} => {
let mut out = Array1::zeros(*total_dim);
let v_block = v.slice(ndarray::s![col_range.clone()]);
let result_block = local.dot(&v_block);
out.slice_mut(ndarray::s![col_range.clone()])
.assign(&result_block);
out
}
Self::Labeled { inner, .. } | Self::Fixed { inner, .. } => inner.dot(v),
}
}
pub fn add_scaled_to(&self, lambda: f64, target: &mut Array2<f64>) {
match self {
Self::Dense(m) => {
target.scaled_add(lambda, m);
}
Self::Diagonal(diagonal) => {
assert_eq!(target.dim(), (diagonal.len(), diagonal.len()));
for (index, &value) in diagonal.iter().enumerate() {
target[[index, index]] += lambda * value;
}
}
Self::KroneckerFactored { left, right } => {
let p_left = left.nrows();
let p_right = right.nrows();
for i1 in 0..p_left {
for j1 in 0..p_left {
let a_ij = left[[i1, j1]];
if a_ij == 0.0 {
continue;
}
let scaled_a = lambda * a_ij;
for i2 in 0..p_right {
let row = i1 * p_right + i2;
for j2 in 0..p_right {
let col = j1 * p_right + j2;
target[[row, col]] += scaled_a * right[[i2, j2]];
}
}
}
}
}
Self::Blockwise {
local, col_range, ..
} => {
target
.slice_mut(ndarray::s![col_range.clone(), col_range.clone()])
.scaled_add(lambda, local);
}
Self::Labeled { inner, .. } | Self::Fixed { inner, .. } => {
inner.add_scaled_to(lambda, target)
}
}
}
pub fn nrows(&self) -> usize {
self.dim()
}
pub fn ncols(&self) -> usize {
self.dim()
}
}
impl From<Array2<f64>> for PenaltyMatrix {
fn from(m: Array2<f64>) -> Self {
Self::Dense(m)
}
}
impl From<Array1<f64>> for PenaltyMatrix {
fn from(diagonal: Array1<f64>) -> Self {
Self::Diagonal(diagonal)
}
}
fn validate_symmetric_psd_core(matrix: &Array2<f64>, what: &str) -> Result<(), String> {
use gam_linalg::faer_ndarray::FaerEigh;
let (nrows, ncols) = matrix.dim();
if nrows != ncols {
return Err(format!("{what} is not square: {nrows}x{ncols}"));
}
let mut max_abs = 0.0_f64;
for ((row, col), &value) in matrix.indexed_iter() {
if !value.is_finite() {
return Err(format!(
"{what} has non-finite entry at ({row},{col}): {value}"
));
}
max_abs = max_abs.max(value.abs());
}
let sym_tol = 1e-10 * max_abs.max(1.0);
for row in 0..nrows {
for col in (row + 1)..ncols {
let asymmetry = (matrix[[row, col]] - matrix[[col, row]]).abs();
if asymmetry > sym_tol {
return Err(format!(
"{what} is not symmetric at ({row},{col}): |S - Sᵀ| = {asymmetry:.3e}; \
the gradient of βᵀSβ/2 is sym(S)β, so a skew component would make the \
implemented objective and gradient describe different functions"
));
}
}
}
if nrows == 0 || max_abs == 0.0 {
return Ok(()); }
let (eigenvalues, _) = matrix
.eigh(faer::Side::Lower)
.map_err(|e| format!("{what} eigendecomposition failed during validation: {e}"))?;
let max_abs_eval = eigenvalues
.iter()
.fold(0.0_f64, |acc, &ev| acc.max(ev.abs()));
let psd_tol = 100.0 * (nrows as f64) * f64::EPSILON * max_abs_eval;
if let Some(&min_eval) = eigenvalues
.iter()
.filter(|&&ev| ev < -psd_tol)
.min_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal))
{
return Err(format!(
"{what} is not positive semidefinite: min eigenvalue {min_eval:.6e} \
(max |eigenvalue| {max_abs_eval:.6e}); the penalized objective is unbounded \
below along the negative mode while rank/logdet filtering would silently \
drop it"
));
}
Ok(())
}
pub fn kronecker_product(a: &Array2<f64>, b: &Array2<f64>) -> Array2<f64> {
let (arows, a_cols) = a.dim();
let (brows, b_cols) = b.dim();
if arows == 0 || a_cols == 0 || brows == 0 || b_cols == 0 {
return Array2::zeros((arows * brows, a_cols * b_cols));
}
let mut result = Array2::zeros((arows * brows, a_cols * b_cols));
result
.axis_chunks_iter_mut(Axis(0), brows)
.into_par_iter()
.enumerate()
.for_each(|(i, mut row_block)| {
let arow = a.row(i);
let col_chunks = row_block.axis_chunks_iter_mut(Axis(1), b_cols);
for (j, mut block) in col_chunks.into_iter().enumerate() {
let aval = arow[j];
if aval == 0.0 {
continue;
}
for (dest, &src) in block.iter_mut().zip(b.iter()) {
*dest = aval * src;
}
}
});
result
}
#[cfg(test)]
mod tests {
use super::*;
use ndarray::array;
#[test]
fn dense_dim_and_shape() {
let m = array![[1.0, 0.0], [0.0, 2.0]];
let p = PenaltyMatrix::Dense(m);
assert_eq!(p.dim(), 2);
assert_eq!(p.shape(), (2, 2));
assert_eq!(p.nrows(), 2);
assert_eq!(p.ncols(), 2);
}
#[test]
fn dense_to_dense_is_clone() {
let m = array![[3.0, 1.0], [1.0, 4.0]];
let p = PenaltyMatrix::Dense(m.clone());
assert_eq!(p.to_dense(), m);
}
#[test]
fn dense_dot_product() {
let m = array![[1.0, 0.0], [0.0, 2.0]];
let p = PenaltyMatrix::Dense(m);
let v = ndarray::array![3.0, 5.0];
let result = p.dot(&v);
assert_eq!(result.as_slice().unwrap(), &[3.0, 10.0]);
}
#[test]
fn dense_add_scaled_to() {
let s = array![[1.0, 0.0], [0.0, 1.0]];
let p = PenaltyMatrix::Dense(s);
let mut acc = ndarray::Array2::<f64>::zeros((2, 2));
p.add_scaled_to(3.0, &mut acc);
assert_eq!(acc, array![[3.0, 0.0], [0.0, 3.0]]);
}
#[test]
fn diagonal_carrier_rejects_nonfinite_and_negative_precision() {
assert!(
PenaltyMatrix::Diagonal(array![1.0, f64::NAN])
.validate(2)
.unwrap_err()
.contains("non-finite")
);
assert!(
PenaltyMatrix::Diagonal(array![1.0, -0.25])
.validate(2)
.unwrap_err()
.contains("not positive semidefinite")
);
}
#[test]
fn kronecker_dim_is_product() {
let left = array![[1.0, 0.0], [0.0, 1.0]]; let right = array![[1.0, 0.0, 0.0], [0.0, 1.0, 0.0], [0.0, 0.0, 1.0]]; let p = PenaltyMatrix::KroneckerFactored { left, right };
assert_eq!(p.dim(), 6);
}
#[test]
fn kronecker_to_dense_identity_x_identity() {
let eye2 = ndarray::Array2::<f64>::eye(2);
let p = PenaltyMatrix::KroneckerFactored {
left: eye2.clone(),
right: eye2,
};
let dense = p.to_dense();
assert_eq!(dense, ndarray::Array2::<f64>::eye(4));
}
#[test]
fn kronecker_dot_matches_dense_dot() {
let left = array![[2.0, 0.0], [0.0, 3.0]];
let right = array![[1.0, 1.0], [0.0, 1.0]];
let p = PenaltyMatrix::KroneckerFactored {
left: left.clone(),
right: right.clone(),
};
let dense = p.to_dense();
let v = ndarray::array![1.0, 2.0, 3.0, 4.0];
let got = p.dot(&v);
let expected = dense.dot(&v);
for (a, b) in got.iter().zip(expected.iter()) {
assert!((a - b).abs() < 1e-14, "got={a} expected={b}");
}
}
#[test]
fn blockwise_dim_is_total() {
let local = array![[1.0, 0.0], [0.0, 1.0]];
let p = PenaltyMatrix::Blockwise {
local,
col_range: 1..3,
total_dim: 5,
};
assert_eq!(p.dim(), 5);
}
#[test]
fn blockwise_to_dense_embeds_local_block() {
let local = array![[2.0, 1.0], [1.0, 3.0]];
let p = PenaltyMatrix::Blockwise {
local,
col_range: 1..3,
total_dim: 3,
};
let dense = p.to_dense();
assert_eq!(dense[[0, 0]], 0.0);
assert_eq!(dense[[1, 1]], 2.0);
assert_eq!(dense[[1, 2]], 1.0);
assert_eq!(dense[[2, 1]], 1.0);
assert_eq!(dense[[2, 2]], 3.0);
}
#[test]
fn blockwise_dot_only_touches_block() {
let local = array![[2.0, 0.0], [0.0, 3.0]];
let p = PenaltyMatrix::Blockwise {
local,
col_range: 1..3,
total_dim: 4,
};
let v = ndarray::array![7.0, 1.0, 2.0, 9.0];
let out = p.dot(&v);
assert_eq!(out.as_slice().unwrap(), &[0.0, 2.0, 6.0, 0.0]);
}
#[test]
fn labeled_inherits_dim_and_delegates_dot() {
let m = array![[1.0, 0.0], [0.0, 2.0]];
let p = PenaltyMatrix::Dense(m).with_precision_label("smooth");
assert_eq!(p.dim(), 2);
assert_eq!(p.precision_label(), Some("smooth"));
let v = ndarray::array![3.0, 4.0];
let out = p.dot(&v);
assert_eq!(out.as_slice().unwrap(), &[3.0, 8.0]);
}
#[test]
fn fixed_inherits_dim_and_exposes_log_lambda() {
let m = array![[5.0, 0.0], [0.0, 5.0]];
let p = PenaltyMatrix::Dense(m).with_fixed_log_lambda(2.5);
assert_eq!(p.dim(), 2);
assert_eq!(p.fixed_log_lambda(), Some(2.5));
}
#[test]
fn shape_reports_actual_storage_not_fabricated_square() {
let p = PenaltyMatrix::Dense(Array2::<f64>::zeros((2, 3)));
assert_eq!(p.shape(), (2, 3));
assert!(p.validate(2).is_err());
assert!(p.validate(3).is_err());
}
#[test]
fn validate_accepts_canonical_carriers() {
let dense = PenaltyMatrix::Dense(array![[2.0, -1.0], [-1.0, 2.0]]);
assert_eq!(dense.validate(2), Ok(()));
let kron = PenaltyMatrix::KroneckerFactored {
left: array![[1.0, -1.0], [-1.0, 1.0]],
right: ndarray::Array2::<f64>::eye(3),
};
assert_eq!(kron.validate(6), Ok(()));
let blockwise = PenaltyMatrix::Blockwise {
local: array![[1.0, 0.0], [0.0, 1.0]],
col_range: 1..3,
total_dim: 4,
};
assert_eq!(blockwise.validate(4), Ok(()));
}
#[test]
fn validate_rejects_asymmetric_indefinite_and_nonfinite() {
let skew = PenaltyMatrix::Dense(array![[1.0, 1.0], [0.0, 1.0]]);
assert!(skew.validate(2).unwrap_err().contains("not symmetric"));
let indefinite = PenaltyMatrix::Dense(array![[1.0, 0.0], [0.0, -1.0]]);
assert!(
indefinite
.validate(2)
.unwrap_err()
.contains("not positive semidefinite")
);
let nan = PenaltyMatrix::Dense(array![[f64::NAN, 0.0], [0.0, 1.0]]);
assert!(nan.validate(2).unwrap_err().contains("non-finite"));
let bad_fixed = PenaltyMatrix::Dense(ndarray::Array2::<f64>::eye(2))
.with_fixed_log_lambda(f64::INFINITY);
assert!(
bad_fixed
.validate(2)
.unwrap_err()
.contains("must be finite")
);
let finite_but_out_of_domain = PenaltyMatrix::Dense(ndarray::Array2::<f64>::eye(2))
.with_fixed_log_lambda(crate::LOG_STRENGTH_MAX + 1.0);
assert!(
finite_but_out_of_domain
.validate(2)
.unwrap_err()
.contains("must be finite and in")
);
}
#[test]
fn validate_rejects_inconsistent_blockwise_embedding() {
let p = PenaltyMatrix::Blockwise {
local: ndarray::Array2::<f64>::eye(3),
col_range: 1..3,
total_dim: 4,
};
assert!(p.validate(4).is_err());
let q = PenaltyMatrix::Blockwise {
local: ndarray::Array2::<f64>::eye(2),
col_range: 3..5,
total_dim: 4,
};
assert!(q.validate(4).is_err());
}
}