use scirs2_core::ndarray::Array2;
use crate::error::{StatsError, StatsResult};
#[derive(Debug, Clone, PartialEq)]
#[non_exhaustive]
pub enum LatentFieldType {
IID,
RW1,
RW2,
AR1 {
phi: f64,
},
ICAR {
adjacency: Vec<(usize, usize)>,
},
Matern {
range: f64,
smoothness: f64,
},
}
pub fn validate_field_params(field_type: &LatentFieldType) -> StatsResult<()> {
match field_type {
LatentFieldType::IID | LatentFieldType::RW1 | LatentFieldType::RW2 => Ok(()),
LatentFieldType::AR1 { phi } => {
if phi.abs() >= 1.0 {
return Err(StatsError::InvalidArgument(format!(
"AR1 parameter phi must satisfy |phi| < 1, got phi = {phi}"
)));
}
if phi.is_nan() {
return Err(StatsError::InvalidArgument(
"AR1 parameter phi must not be NaN".to_string(),
));
}
Ok(())
}
LatentFieldType::ICAR { adjacency } => {
for &(i, j) in adjacency {
if i == j {
return Err(StatsError::InvalidArgument(format!(
"ICAR adjacency contains self-loop at node {i}"
)));
}
}
Ok(())
}
LatentFieldType::Matern { range, smoothness } => {
if *range <= 0.0 || range.is_nan() {
return Err(StatsError::InvalidArgument(format!(
"Matern range must be positive, got {range}"
)));
}
if *smoothness <= 0.0 || smoothness.is_nan() {
return Err(StatsError::InvalidArgument(format!(
"Matern smoothness must be positive, got {smoothness}"
)));
}
Ok(())
}
_ => Ok(()),
}
}
pub fn build_precision_matrix(
field_type: &LatentFieldType,
n: usize,
scale: f64,
) -> StatsResult<Array2<f64>> {
if n == 0 {
return Err(StatsError::InvalidArgument(
"Field dimension n must be at least 1".to_string(),
));
}
if scale <= 0.0 || scale.is_nan() {
return Err(StatsError::InvalidArgument(format!(
"Scale parameter must be positive, got {scale}"
)));
}
validate_field_params(field_type)?;
match field_type {
LatentFieldType::IID => build_iid_precision(n, scale),
LatentFieldType::RW1 => build_rw1_precision(n, scale),
LatentFieldType::RW2 => build_rw2_precision(n, scale),
LatentFieldType::AR1 { phi } => build_ar1_precision(n, scale, *phi),
LatentFieldType::ICAR { adjacency } => build_icar_precision(n, scale, adjacency),
LatentFieldType::Matern { range, smoothness } => {
build_matern_precision(n, scale, *range, *smoothness)
}
_ => Err(StatsError::NotImplementedError(
"Unknown latent field type".to_string(),
)),
}
}
fn build_iid_precision(n: usize, scale: f64) -> StatsResult<Array2<f64>> {
let mut q = Array2::zeros((n, n));
for i in 0..n {
q[[i, i]] = scale;
}
Ok(q)
}
fn build_rw1_precision(n: usize, scale: f64) -> StatsResult<Array2<f64>> {
let mut q = Array2::zeros((n, n));
if n == 1 {
q[[0, 0]] = scale;
return Ok(q);
}
for i in 0..n {
if i == 0 || i == n - 1 {
q[[i, i]] = scale;
} else {
q[[i, i]] = 2.0 * scale;
}
if i + 1 < n {
q[[i, i + 1]] = -scale;
q[[i + 1, i]] = -scale;
}
}
Ok(q)
}
fn build_rw2_precision(n: usize, scale: f64) -> StatsResult<Array2<f64>> {
if n < 3 {
return Err(StatsError::InvalidArgument(format!(
"RW2 requires n >= 3, got n = {n}"
)));
}
let mut q = Array2::zeros((n, n));
for row in 0..(n - 2) {
let indices = [row, row + 1, row + 2];
let coeffs = [1.0, -2.0, 1.0];
for (a, &ia) in indices.iter().enumerate() {
for (b, &ib) in indices.iter().enumerate() {
q[[ia, ib]] += scale * coeffs[a] * coeffs[b];
}
}
}
Ok(q)
}
fn build_ar1_precision(n: usize, scale: f64, phi: f64) -> StatsResult<Array2<f64>> {
let mut q = Array2::zeros((n, n));
if n == 1 {
q[[0, 0]] = scale;
return Ok(q);
}
let marginal_scale = scale / (1.0 - phi * phi);
for i in 0..n {
if i == 0 || i == n - 1 {
q[[i, i]] = marginal_scale;
} else {
q[[i, i]] = marginal_scale * (1.0 + phi * phi);
}
if i + 1 < n {
q[[i, i + 1]] = -marginal_scale * phi;
q[[i + 1, i]] = -marginal_scale * phi;
}
}
Ok(q)
}
fn build_icar_precision(
n: usize,
scale: f64,
adjacency: &[(usize, usize)],
) -> StatsResult<Array2<f64>> {
for &(i, j) in adjacency {
if i >= n || j >= n {
return Err(StatsError::InvalidArgument(format!(
"ICAR adjacency index out of bounds: ({i}, {j}) for n = {n}"
)));
}
}
let mut q = Array2::zeros((n, n));
for &(i, j) in adjacency {
q[[i, j]] -= scale;
q[[j, i]] -= scale;
q[[i, i]] += scale;
q[[j, j]] += scale;
}
Ok(q)
}
fn build_matern_precision(
n: usize,
scale: f64,
range: f64,
smoothness: f64,
) -> StatsResult<Array2<f64>> {
let kappa = (8.0 * smoothness).sqrt() / range;
let kappa2 = kappa * kappa;
let alpha = smoothness + 0.5;
if alpha <= 1.0 + 1e-10 {
let h = 1.0;
let h2_inv = 1.0 / (h * h);
let mut q = Array2::zeros((n, n));
for i in 0..n {
if i == 0 || i == n - 1 {
q[[i, i]] = scale * (kappa2 + h2_inv);
} else {
q[[i, i]] = scale * (kappa2 + 2.0 * h2_inv);
}
if i + 1 < n {
q[[i, i + 1]] = -scale * h2_inv;
q[[i + 1, i]] = -scale * h2_inv;
}
}
Ok(q)
} else {
let h = 1.0;
let h2_inv = 1.0 / (h * h);
let mut c = Array2::zeros((n, n));
for i in 0..n {
if i == 0 || i == n - 1 {
c[[i, i]] = kappa2 + h2_inv;
} else {
c[[i, i]] = kappa2 + 2.0 * h2_inv;
}
if i + 1 < n {
c[[i, i + 1]] = -h2_inv;
c[[i + 1, i]] = -h2_inv;
}
}
let ct = c.t();
let q = ct.dot(&c) * scale;
Ok(q)
}
}
pub fn kronecker_precision(q1: &Array2<f64>, q2: &Array2<f64>) -> Array2<f64> {
let (m1, n1) = (q1.nrows(), q1.ncols());
let (m2, n2) = (q2.nrows(), q2.ncols());
let mut result = Array2::zeros((m1 * m2, n1 * n2));
for i1 in 0..m1 {
for j1 in 0..n1 {
let val = q1[[i1, j1]];
if val.abs() < 1e-30 {
continue;
}
for i2 in 0..m2 {
for j2 in 0..n2 {
result[[i1 * m2 + i2, j1 * n2 + j2]] = val * q2[[i2, j2]];
}
}
}
}
result
}
#[cfg(test)]
mod tests {
use super::*;
use scirs2_core::ndarray::array;
#[test]
fn test_iid_precision() {
let q = build_precision_matrix(&LatentFieldType::IID, 4, 2.0)
.expect("IID precision should succeed");
assert_eq!(q.nrows(), 4);
assert_eq!(q.ncols(), 4);
for i in 0..4 {
assert!((q[[i, i]] - 2.0).abs() < 1e-12);
for j in 0..4 {
if i != j {
assert!(q[[i, j]].abs() < 1e-12);
}
}
}
}
#[test]
fn test_rw1_precision_n5() {
let q = build_precision_matrix(&LatentFieldType::RW1, 5, 1.0)
.expect("RW1 precision should succeed");
assert_eq!(q.nrows(), 5);
assert!((q[[0, 0]] - 1.0).abs() < 1e-12, "q[0,0] = {}", q[[0, 0]]);
assert!((q[[4, 4]] - 1.0).abs() < 1e-12, "q[4,4] = {}", q[[4, 4]]);
assert!((q[[1, 1]] - 2.0).abs() < 1e-12);
assert!((q[[2, 2]] - 2.0).abs() < 1e-12);
assert!((q[[3, 3]] - 2.0).abs() < 1e-12);
assert!((q[[0, 1]] - (-1.0)).abs() < 1e-12);
assert!((q[[1, 0]] - (-1.0)).abs() < 1e-12);
assert!((q[[3, 4]] - (-1.0)).abs() < 1e-12);
assert!(q[[0, 2]].abs() < 1e-12);
}
#[test]
fn test_rw2_precision_n6() {
let q = build_precision_matrix(&LatentFieldType::RW2, 6, 1.0)
.expect("RW2 precision should succeed");
assert_eq!(q.nrows(), 6);
for i in 0..6 {
for j in 0..6 {
assert!(
(q[[i, j]] - q[[j, i]]).abs() < 1e-12,
"Not symmetric at ({i},{j})"
);
}
}
for i in 0..6 {
for j in 0..6 {
if (i as isize - j as isize).unsigned_abs() > 2 {
assert!(
q[[i, j]].abs() < 1e-12,
"Non-zero at ({i},{j}): {}",
q[[i, j]]
);
}
}
}
assert!((q[[0, 0]] - 1.0).abs() < 1e-12);
assert!((q[[1, 1]] - 5.0).abs() < 1e-12);
assert!((q[[2, 2]] - 6.0).abs() < 1e-12);
assert!((q[[3, 3]] - 6.0).abs() < 1e-12);
assert!((q[[4, 4]] - 5.0).abs() < 1e-12);
assert!((q[[5, 5]] - 1.0).abs() < 1e-12);
}
#[test]
fn test_ar1_phi_zero_equals_iid() {
let n = 5;
let scale = 3.0;
let q_ar1 = build_precision_matrix(&LatentFieldType::AR1 { phi: 0.0 }, n, scale)
.expect("AR1(0) should succeed");
let q_iid =
build_precision_matrix(&LatentFieldType::IID, n, scale).expect("IID should succeed");
for i in 0..n {
for j in 0..n {
assert!(
(q_ar1[[i, j]] - q_iid[[i, j]]).abs() < 1e-10,
"Mismatch at ({i},{j}): AR1={}, IID={}",
q_ar1[[i, j]],
q_iid[[i, j]]
);
}
}
}
#[test]
fn test_ar1_banding_phi09() {
let n = 5;
let q = build_precision_matrix(&LatentFieldType::AR1 { phi: 0.9 }, n, 1.0)
.expect("AR1(0.9) should succeed");
let phi = 0.9;
let s = 1.0 / (1.0 - phi * phi);
assert!((q[[2, 2]] - s * (1.0 + phi * phi)).abs() < 1e-10);
assert!((q[[0, 0]] - s).abs() < 1e-10);
assert!((q[[0, 1]] - (-s * phi)).abs() < 1e-10);
assert!(q[[0, 2]].abs() < 1e-12);
}
#[test]
fn test_icar_line_graph_4() {
let adj = vec![(0, 1), (1, 2), (2, 3)];
let q = build_precision_matrix(&LatentFieldType::ICAR { adjacency: adj }, 4, 1.0)
.expect("ICAR line graph should succeed");
assert!((q[[0, 0]] - 1.0).abs() < 1e-12);
assert!((q[[1, 1]] - 2.0).abs() < 1e-12);
assert!((q[[2, 2]] - 2.0).abs() < 1e-12);
assert!((q[[3, 3]] - 1.0).abs() < 1e-12);
assert!((q[[0, 1]] - (-1.0)).abs() < 1e-12);
assert!((q[[1, 2]] - (-1.0)).abs() < 1e-12);
assert!(q[[0, 2]].abs() < 1e-12);
}
#[test]
fn test_icar_triangle_graph() {
let adj = vec![(0, 1), (1, 2), (0, 2)];
let q = build_precision_matrix(&LatentFieldType::ICAR { adjacency: adj }, 3, 1.0)
.expect("ICAR triangle should succeed");
for i in 0..3 {
assert!((q[[i, i]] - 2.0).abs() < 1e-12);
}
assert!((q[[0, 1]] - (-1.0)).abs() < 1e-12);
assert!((q[[0, 2]] - (-1.0)).abs() < 1e-12);
assert!((q[[1, 2]] - (-1.0)).abs() < 1e-12);
}
#[test]
fn test_matern_positive_semidefinite() {
let q = build_precision_matrix(
&LatentFieldType::Matern {
range: 2.0,
smoothness: 1.0,
},
8,
1.0,
)
.expect("Matern should succeed");
for i in 0..8 {
for j in 0..8 {
assert!(
(q[[i, j]] - q[[j, i]]).abs() < 1e-12,
"Matern not symmetric at ({i},{j})"
);
}
}
for i in 0..8 {
let off_sum: f64 = (0..8).filter(|&j| j != i).map(|j| q[[i, j]].abs()).sum();
assert!(
q[[i, i]] >= off_sum - 1e-10,
"Gershgorin violated at row {i}: diag={}, off_sum={}",
q[[i, i]],
off_sum
);
}
}
#[test]
fn test_invalid_ar1_phi() {
let result = build_precision_matrix(&LatentFieldType::AR1 { phi: 1.0 }, 5, 1.0);
assert!(result.is_err());
let result2 = build_precision_matrix(&LatentFieldType::AR1 { phi: -1.5 }, 5, 1.0);
assert!(result2.is_err());
}
#[test]
fn test_empty_adjacency_icar() {
let q = build_precision_matrix(&LatentFieldType::ICAR { adjacency: vec![] }, 3, 1.0)
.expect("Empty ICAR should produce zero matrix");
for i in 0..3 {
for j in 0..3 {
assert!(q[[i, j]].abs() < 1e-12);
}
}
}
#[test]
fn test_kronecker_product() {
let a = array![[1.0, 2.0], [3.0, 4.0]];
let b = array![[5.0, 6.0], [7.0, 8.0]];
let k = kronecker_precision(&a, &b);
assert_eq!(k.nrows(), 4);
assert_eq!(k.ncols(), 4);
assert!((k[[0, 0]] - 5.0).abs() < 1e-12);
assert!((k[[0, 1]] - 6.0).abs() < 1e-12);
assert!((k[[0, 2]] - 10.0).abs() < 1e-12);
assert!((k[[3, 3]] - 32.0).abs() < 1e-12);
}
#[test]
fn test_precision_matrix_symmetric() {
let types: Vec<LatentFieldType> = vec![
LatentFieldType::IID,
LatentFieldType::RW1,
LatentFieldType::AR1 { phi: 0.5 },
LatentFieldType::ICAR {
adjacency: vec![(0, 1), (1, 2), (2, 3)],
},
LatentFieldType::Matern {
range: 1.0,
smoothness: 0.5,
},
];
for ft in &types {
let q = build_precision_matrix(ft, 5, 1.0)
.unwrap_or_else(|_| panic!("Failed for {:?}", ft));
for i in 0..5 {
for j in 0..5 {
assert!(
(q[[i, j]] - q[[j, i]]).abs() < 1e-12,
"Not symmetric for {:?} at ({i},{j})",
ft
);
}
}
}
}
#[test]
fn test_scale_applied_correctly() {
let scale = 3.5;
let q1 = build_precision_matrix(&LatentFieldType::RW1, 4, 1.0).expect("RW1 scale=1");
let q_s = build_precision_matrix(&LatentFieldType::RW1, 4, scale).expect("RW1 scaled");
for i in 0..4 {
for j in 0..4 {
assert!(
(q_s[[i, j]] - scale * q1[[i, j]]).abs() < 1e-12,
"Scale mismatch at ({i},{j})"
);
}
}
}
#[test]
fn test_zero_dimension_rejected() {
let result = build_precision_matrix(&LatentFieldType::IID, 0, 1.0);
assert!(result.is_err());
}
#[test]
fn test_negative_scale_rejected() {
let result = build_precision_matrix(&LatentFieldType::IID, 3, -1.0);
assert!(result.is_err());
}
#[test]
fn test_icar_out_of_bounds_rejected() {
let result = build_precision_matrix(
&LatentFieldType::ICAR {
adjacency: vec![(0, 5)],
},
3,
1.0,
);
assert!(result.is_err());
}
#[test]
fn test_rw2_too_small_rejected() {
let result = build_precision_matrix(&LatentFieldType::RW2, 2, 1.0);
assert!(result.is_err());
}
#[test]
fn test_matern_invalid_range() {
let result = build_precision_matrix(
&LatentFieldType::Matern {
range: -1.0,
smoothness: 1.0,
},
5,
1.0,
);
assert!(result.is_err());
}
#[test]
fn test_validate_icar_self_loop() {
let result = validate_field_params(&LatentFieldType::ICAR {
adjacency: vec![(1, 1)],
});
assert!(result.is_err());
}
}