use ndarray::{Array1, Array2, ArrayView1, ArrayView2};
use crate::linalg::faer_ndarray::FaerCholeskyFactor;
pub enum FittedInverse<'a> {
FaerCholesky(&'a FaerCholeskyFactor),
LowerTriangular(&'a Array2<f64>),
Projected {
basis: &'a Array2<f64>,
reduced_inverse: &'a Array2<f64>,
},
}
pub struct FitSensitivity<'a> {
inverse: FittedInverse<'a>,
dim: usize,
}
impl<'a> FitSensitivity<'a> {
pub fn from_faer_cholesky(factor: &'a FaerCholeskyFactor, dim: usize) -> Self {
Self {
inverse: FittedInverse::FaerCholesky(factor),
dim,
}
}
pub fn from_lower_triangular(factor: &'a Array2<f64>) -> Self {
let dim = factor.nrows();
Self {
inverse: FittedInverse::LowerTriangular(factor),
dim,
}
}
pub fn from_projected(basis: &'a Array2<f64>, reduced_inverse: &'a Array2<f64>) -> Self {
let dim = basis.nrows();
Self {
inverse: FittedInverse::Projected {
basis,
reduced_inverse,
},
dim,
}
}
pub fn dim(&self) -> usize {
self.dim
}
pub fn apply(&self, rhs: &Array1<f64>) -> Array1<f64> {
assert_eq!(rhs.len(), self.dim, "FitSensitivity rhs dimension");
match &self.inverse {
FittedInverse::FaerCholesky(factor) => factor.solvevec(rhs),
FittedInverse::LowerTriangular(factor) => {
crate::linalg::triangular::cholesky_solve_vector(factor.view(), rhs.view())
}
FittedInverse::Projected {
basis,
reduced_inverse,
} => {
let proj = crate::linalg::faer_ndarray::fast_atv(basis, rhs);
let reduced = reduced_inverse.dot(&proj);
crate::linalg::faer_ndarray::fast_av(basis, &reduced)
}
}
}
pub fn apply_multi(&self, rhs: ArrayView2<'_, f64>) -> Array2<f64> {
assert_eq!(rhs.nrows(), self.dim, "FitSensitivity RHS dimension");
match &self.inverse {
FittedInverse::FaerCholesky(factor) => {
let mut out = rhs.to_owned();
factor.solve_mat_in_place(&mut out);
out
}
FittedInverse::LowerTriangular(factor) => {
crate::linalg::triangular::cholesky_solve_matrix(*factor, rhs)
}
FittedInverse::Projected {
basis,
reduced_inverse,
} => {
let reduced = crate::linalg::faer_ndarray::fast_atb(basis, &rhs.to_owned());
crate::linalg::faer_ndarray::fast_ab(basis, &reduced_inverse.dot(&reduced))
}
}
}
pub fn mode_response(&self, dg_dt: ArrayView2<'_, f64>) -> Option<Array2<f64>> {
if dg_dt.nrows() != self.dim {
return None;
}
let mut out = self.apply_multi(dg_dt);
if out.iter().any(|value| !value.is_finite()) {
return None;
}
out.mapv_inplace(|value| -value);
Some(out)
}
pub fn mode_response_coned(
&self,
hessian: ArrayView2<'_, f64>,
dg_dt: ArrayView2<'_, f64>,
col_supports: &[std::ops::Range<usize>],
) -> Option<Array2<f64>> {
let p = self.dim;
let r = dg_dt.ncols();
if dg_dt.nrows() != p
|| hessian.nrows() != p
|| hessian.ncols() != p
|| col_supports.len() != r
{
return None;
}
let labels = crate::solver::evidence::coupling_components(hessian);
if labels.len() != p {
return None;
}
let mut active: Vec<(usize, Vec<usize>)> = Vec::new();
for a in 0..r {
let sr = &col_supports[a];
let support: Vec<usize> = (sr.start..sr.end)
.filter(|idx| *idx < p)
.filter(|idx| dg_dt[[*idx, a]] != 0.0)
.collect();
let cone = crate::solver::evidence::cone_of_influence(&labels, &support);
if !cone.is_empty() {
active.push((a, cone));
}
}
let mut out = Array2::<f64>::zeros((p, r));
if active.is_empty() {
return Some(out);
}
let mut rhs = Array2::<f64>::zeros((p, active.len()));
for (j, (a, _)) in active.iter().enumerate() {
rhs.column_mut(j).assign(&dg_dt.column(*a));
}
let solved = self.apply_multi(rhs.view());
if solved.iter().any(|value| !value.is_finite()) {
return None;
}
for (j, (a, cone)) in active.iter().enumerate() {
for &row in cone {
out[[row, *a]] = -solved[[row, j]];
}
}
Some(out)
}
pub fn leverage_block(&self, design: &Array2<f64>) -> Array2<f64> {
assert_eq!(design.ncols(), self.dim, "FitSensitivity design width");
self.apply_multi(design.t())
}
pub fn response_jacobian(
&self,
design: &Array2<f64>,
working_weights: ArrayView1<'_, f64>,
) -> Option<Array2<f64>> {
let n = design.nrows();
if design.ncols() != self.dim || working_weights.len() != n {
return None;
}
let mut dbeta_dy = self.leverage_block(design);
for i in 0..n {
let w_i = working_weights[i];
dbeta_dy.column_mut(i).mapv_inplace(|v| w_i * v);
}
Some(dbeta_dy)
}
pub fn case_deletion(
&self,
design: &Array2<f64>,
working_weights: ArrayView1<'_, f64>,
working_residual: ArrayView1<'_, f64>,
phi: f64,
) -> Option<CaseDeletionInfluence> {
let n = design.nrows();
let p = design.ncols();
if p != self.dim
|| working_weights.len() != n
|| working_residual.len() != n
|| !(phi.is_finite() && phi > 0.0)
|| p == 0
{
return None;
}
let h_inv_xt = self.leverage_block(design);
let mut dfbeta = Array2::<f64>::zeros((n, p));
let mut leverage = Array1::<f64>::zeros(n);
let mut cooks = Array1::<f64>::zeros(n);
let p_phi = p as f64 * phi;
for i in 0..n {
let hinv_xi = h_inv_xt.column(i);
let xhx = design.row(i).dot(&hinv_xi);
let h_ii = working_weights[i] * xhx;
let denom = 1.0 - h_ii;
if !denom.is_finite() || denom.abs() < f64::EPSILON {
return None;
}
let scale = working_weights[i] * working_residual[i] / denom;
dfbeta.row_mut(i).assign(&(&hinv_xi * scale));
leverage[i] = h_ii;
cooks[i] = scale * scale * xhx / p_phi;
}
Some(CaseDeletionInfluence {
dfbeta,
leverage,
cooks_distance: cooks,
})
}
}
pub struct CaseDeletionInfluence {
pub dfbeta: Array2<f64>,
pub leverage: Array1<f64>,
pub cooks_distance: Array1<f64>,
}
#[cfg(test)]
mod tests {
use super::*;
use crate::linalg::faer_ndarray::FaerCholesky;
use faer::Side;
use ndarray::array;
fn lower_cholesky_3x3(h: &Array2<f64>) -> Array2<f64> {
let mut l = Array2::<f64>::zeros((3, 3));
for i in 0..3 {
for j in 0..=i {
let mut acc = h[[i, j]];
for k in 0..j {
acc -= l[[i, k]] * l[[j, k]];
}
if i == j {
l[[i, j]] = acc.sqrt();
} else {
l[[i, j]] = acc / l[[j, j]];
}
}
}
l
}
#[test]
fn all_variants_agree_on_the_same_curvature() {
let h = array![[4.0, 1.0, 0.5], [1.0, 3.0, 0.2], [0.5, 0.2, 2.0]];
let faer = h.cholesky(Side::Lower).expect("SPD factor");
let lower = lower_cholesky_3x3(&h);
let eye = Array2::eye(3);
let h_inv = {
let mut out: Array2<f64> = Array2::eye(3);
faer.solve_mat_in_place(&mut out);
out
};
let s_faer = FitSensitivity::from_faer_cholesky(&faer, 3);
let s_tri = FitSensitivity::from_lower_triangular(&lower);
let s_proj = FitSensitivity::from_projected(&eye, &h_inv);
let rhs = array![0.7, -1.2, 0.4];
let block = array![[0.7, 1.0], [-1.2, 0.0], [0.4, -2.0]];
let a = s_faer.apply(&rhs);
let b = s_tri.apply(&rhs);
let c = s_proj.apply(&rhs);
for i in 0..3 {
assert!((a[i] - b[i]).abs() <= 1e-12, "faer vs triangular [{i}]");
assert!((a[i] - c[i]).abs() <= 1e-12, "faer vs projected [{i}]");
}
let ma = s_faer.apply_multi(block.view());
let mb = s_tri.apply_multi(block.view());
let mc = s_proj.apply_multi(block.view());
let resp = s_faer.mode_response(block.view()).expect("finite");
for i in 0..3 {
for j in 0..2 {
assert!((ma[[i, j]] - mb[[i, j]]).abs() <= 1e-12);
assert!((ma[[i, j]] - mc[[i, j]]).abs() <= 1e-12);
assert!(
(resp[[i, j]] + ma[[i, j]]).abs() <= 1e-15,
"mode response sign"
);
}
}
let back = h.dot(&a);
for i in 0..3 {
assert!((back[i] - rhs[i]).abs() <= 1e-12, "inverse residual [{i}]");
}
}
#[test]
fn case_deletion_matches_exact_loo_refit() {
let x = array![
[1.0, 0.2, -0.5],
[1.0, -1.1, 0.3],
[1.0, 0.7, 1.4],
[1.0, 2.0, -0.8],
[1.0, -0.4, 0.9],
[1.0, 1.3, 0.1],
];
let y = array![0.5, -1.2, 2.1, 0.3, -0.7, 1.0];
let n = x.nrows();
let p = x.ncols();
let mut s = Array2::<f64>::zeros((p, p));
s[[1, 1]] = 0.4;
s[[2, 2]] = 0.4;
let xtx = x.t().dot(&x);
let h = &xtx + &s;
let h_inv = {
let f = h.cholesky(Side::Lower).expect("SPD");
let mut out: Array2<f64> = Array2::eye(p);
f.solve_mat_in_place(&mut out);
out
};
let xty = x.t().dot(&y);
let beta = h_inv.dot(&xty);
let w = Array1::<f64>::ones(n);
let resid = &y - &x.dot(&beta);
let faer = h.cholesky(Side::Lower).expect("SPD");
let op = FitSensitivity::from_faer_cholesky(&faer, p);
let infl = op
.case_deletion(&x, w.view(), resid.view(), 1.0)
.expect("case deletion");
for i in 0..n {
let x_i = x.row(i).to_owned();
let mut h_del = h.clone();
for a in 0..p {
for b in 0..p {
h_del[[a, b]] -= x_i[a] * x_i[b];
}
}
let rhs_del = &xty - &(&x_i * y[i]);
let h_del_inv = {
let f = h_del.cholesky(Side::Lower).expect("SPD deleted");
let mut out: Array2<f64> = Array2::eye(p);
f.solve_mat_in_place(&mut out);
out
};
let beta_del = h_del_inv.dot(&rhs_del);
let exact_dfbeta = &beta - &beta_del;
for j in 0..p {
assert!(
(infl.dfbeta[[i, j]] - exact_dfbeta[j]).abs() < 1e-9,
"dfbeta[{i},{j}]: operator {} vs exact refit {}",
infl.dfbeta[[i, j]],
exact_dfbeta[j]
);
}
let cook_exact = exact_dfbeta.dot(&h.dot(&exact_dfbeta)) / (p as f64 * 1.0);
assert!(
(infl.cooks_distance[i] - cook_exact).abs() < 1e-9,
"cooks[{i}]: {} vs {}",
infl.cooks_distance[i],
cook_exact
);
let h_ii = x_i.dot(&h_inv.dot(&x_i));
assert!((infl.leverage[i] - h_ii).abs() < 1e-12, "leverage[{i}]");
}
}
#[test]
fn response_jacobian_matches_refit_perturbation() {
let x = array![
[1.0, 0.2, -0.5],
[1.0, -1.1, 0.3],
[1.0, 0.7, 1.4],
[1.0, 2.0, -0.8],
[1.0, -0.4, 0.9],
];
let y = array![0.5, -1.2, 2.1, 0.3, -0.7];
let n = x.nrows();
let p = x.ncols();
let mut s = Array2::<f64>::zeros((p, p));
s[[1, 1]] = 0.4;
s[[2, 2]] = 0.4;
let h = &x.t().dot(&x) + &s;
let faer = h.cholesky(Side::Lower).expect("SPD");
let solve = |rhs: &Array1<f64>| faer.solvevec(rhs);
let beta = solve(&x.t().dot(&y));
let op = FitSensitivity::from_faer_cholesky(&faer, p);
let w = Array1::<f64>::ones(n);
let dbeta_dy = op.response_jacobian(&x, w.view()).expect("jacobian");
let eps = 1e-6;
for j in 0..n {
let mut yp = y.clone();
yp[j] += eps;
let beta_p = solve(&x.t().dot(&yp));
for c in 0..p {
let fd = (beta_p[c] - beta[c]) / eps;
assert!(
(dbeta_dy[[c, j]] - fd).abs() < 1e-7,
"∂β̂_{c}/∂y_{j}: analytic {} vs refit {}",
dbeta_dy[[c, j]],
fd
);
}
}
}
#[test]
fn mode_response_refuses_non_finite_channels() {
let h = array![[2.0, 0.0], [0.0, 1.0]];
let faer = h.cholesky(Side::Lower).expect("SPD factor");
let s = FitSensitivity::from_faer_cholesky(&faer, 2);
let bad = array![[1.0], [f64::NAN]];
assert!(s.mode_response(bad.view()).is_none());
let wrong_dim = array![[1.0], [0.0], [0.0]];
assert!(s.mode_response(wrong_dim.view()).is_none());
}
}