use ndarray::Array1;
use ndarray::Array2;
use crate::traits::FloatExt;
#[derive(Debug, Clone)]
pub struct TransitionMatrix<T: FloatExt> {
matrix: Array2<T>,
}
impl<T: FloatExt> TransitionMatrix<T> {
pub fn new(matrix: Array2<T>) -> Self {
let (rows, cols) = matrix.dim();
assert_eq!(rows, cols, "transition matrix must be square");
Self { matrix }
}
pub fn states(&self) -> usize {
self.matrix.nrows()
}
pub fn default_state(&self) -> usize {
self.matrix.nrows() - 1
}
pub fn matrix(&self) -> &Array2<T> {
&self.matrix
}
pub fn power(&self, k: usize) -> TransitionMatrix<T> {
assert!(k >= 1, "power must be at least 1");
let mut result = self.matrix.clone();
for _ in 1..k {
result = result.dot(&self.matrix);
}
TransitionMatrix::new(result)
}
pub fn default_probabilities(&self, k: usize) -> Array1<T> {
let powered = self.power(k);
powered.matrix.column(self.default_state()).to_owned()
}
pub fn check_row_stochastic(&self, tol: T) -> Result<(), String> {
for (i, row) in self.matrix.rows().into_iter().enumerate() {
let mut s = T::zero();
for &v in row.iter() {
if v < T::zero() && v.abs() > tol {
return Err(format!("row {i} contains negative entry {v:?}"));
}
s += v;
}
if (s - T::one()).abs() > tol {
return Err(format!("row {i} sums to {s:?}, expected 1"));
}
}
Ok(())
}
}
#[derive(Debug, Clone)]
pub struct GeneratorMatrix<T: FloatExt> {
matrix: Array2<T>,
}
impl<T: FloatExt> GeneratorMatrix<T> {
pub fn new(matrix: Array2<T>) -> Self {
let (rows, cols) = matrix.dim();
assert_eq!(rows, cols, "generator matrix must be square");
Self { matrix }
}
pub fn states(&self) -> usize {
self.matrix.nrows()
}
pub fn default_state(&self) -> usize {
self.matrix.nrows() - 1
}
pub fn matrix(&self) -> &Array2<T> {
&self.matrix
}
pub fn transition_at(&self, t: T) -> TransitionMatrix<T> {
let scaled = &self.matrix * t;
let p = expm(&scaled);
TransitionMatrix::new(p)
}
pub fn from_yearly_transition(p: &TransitionMatrix<T>) -> Self {
let n = p.states();
let identity = identity_matrix(n);
let x = p.matrix() - &identity;
let tolerance = T::from_f64_fast(1e-14);
let mut result: Array2<T> = Array2::zeros((n, n));
let mut term = x.clone();
let mut sign = T::one();
for k in 1..=200 {
let k_t = T::from_usize_(k);
let contribution = &term * (sign / k_t);
result = &result + &contribution;
let norm = contribution.iter().fold(T::zero(), |acc, v| acc + v.abs());
if norm < tolerance && k > 3 {
break;
}
term = term.dot(&x);
sign = -sign;
}
Self::new(result)
}
pub fn project_to_generator(&self) -> Self {
let n = self.states();
let mut m: Array2<T> = self.matrix.clone();
for i in 0..n {
let mut off_sum = T::zero();
let mut neg_mass = T::zero();
for j in 0..n {
if i == j {
continue;
}
if m[[i, j]] < T::zero() {
neg_mass += m[[i, j]];
m[[i, j]] = T::zero();
} else {
off_sum += m[[i, j]];
}
}
let total_off = off_sum + neg_mass;
if off_sum > T::zero() && neg_mass < T::zero() {
let scale = (off_sum + neg_mass) / off_sum;
for j in 0..n {
if i != j && m[[i, j]] > T::zero() {
m[[i, j]] = m[[i, j]] * scale;
}
}
}
let new_off_sum: T = (0..n)
.filter(|&j| j != i)
.map(|j| m[[i, j]])
.fold(T::zero(), |acc, v| acc + v);
m[[i, i]] = -new_off_sum;
let _ = total_off;
}
Self::new(m)
}
pub fn default_probabilities(&self, t: T) -> Array1<T> {
let p = self.transition_at(t);
p.matrix.column(self.default_state()).to_owned()
}
pub fn check_generator(&self, tol: T) -> Result<(), String> {
let n = self.states();
for i in 0..n {
let mut s = T::zero();
for j in 0..n {
let v = self.matrix[[i, j]];
if i != j && v < -tol {
return Err(format!("off-diagonal ({i},{j}) negative: {v:?}"));
}
if i == j && v > tol {
return Err(format!("diagonal ({i},{i}) positive: {v:?}"));
}
s += v;
}
if s.abs() > tol {
return Err(format!("row {i} sum {s:?} not zero"));
}
}
Ok(())
}
}
fn identity_matrix<T: FloatExt>(n: usize) -> Array2<T> {
let mut m = Array2::zeros((n, n));
for i in 0..n {
m[[i, i]] = T::one();
}
m
}
fn expm<T: FloatExt>(a: &Array2<T>) -> Array2<T> {
let n = a.nrows();
if n == 0 {
return a.clone();
}
let theta13 = T::from_f64_fast(5.371920351148152);
let norm_a = one_norm(a);
if norm_a <= theta13 {
return pade13(a);
}
let s = (norm_a / theta13)
.ln()
.to_f64()
.unwrap_or(0.0)
.max(0.0)
.ceil() as i32;
let s = s.max(1);
let scale = T::from_f64_fast(2f64.powi(s));
let a_scaled = a / scale;
let mut result = pade13(&a_scaled);
for _ in 0..s {
result = result.dot(&result);
}
result
}
fn one_norm<T: FloatExt>(a: &Array2<T>) -> T {
let mut max_col = T::zero();
for j in 0..a.ncols() {
let mut s = T::zero();
for i in 0..a.nrows() {
s += a[[i, j]].abs();
}
if s > max_col {
max_col = s;
}
}
max_col
}
fn pade13<T: FloatExt>(a: &Array2<T>) -> Array2<T> {
let b: [T; 14] = [
T::from_f64_fast(64764752532480000.0),
T::from_f64_fast(32382376266240000.0),
T::from_f64_fast(7771770303897600.0),
T::from_f64_fast(1187353796428800.0),
T::from_f64_fast(129060195264000.0),
T::from_f64_fast(10559470521600.0),
T::from_f64_fast(670442572800.0),
T::from_f64_fast(33522128640.0),
T::from_f64_fast(1323241920.0),
T::from_f64_fast(40840800.0),
T::from_f64_fast(960960.0),
T::from_f64_fast(16380.0),
T::from_f64_fast(182.0),
T::from_f64_fast(1.0),
];
let n = a.nrows();
let id = identity_matrix::<T>(n);
let a2 = a.dot(a);
let a4 = a2.dot(&a2);
let a6 = a4.dot(&a2);
let u_outer = &a6 * b[13] + &a4 * b[11] + &a2 * b[9];
let u_inner = a6.dot(&u_outer);
let u_low = &a6 * b[7] + &a4 * b[5] + &a2 * b[3] + &id * b[1];
let u = a.dot(&(&u_inner + &u_low));
let v_outer = &a6 * b[12] + &a4 * b[10] + &a2 * b[8];
let v_inner = a6.dot(&v_outer);
let v_low = &a6 * b[6] + &a4 * b[4] + &a2 * b[2] + &id * b[0];
let v = &v_inner + &v_low;
let numer = &v + &u;
let denom = &v - &u;
let denom_inv = invert_matrix(&denom);
denom_inv.dot(&numer)
}
fn invert_matrix<T: FloatExt>(a: &Array2<T>) -> Array2<T> {
let n = a.nrows();
assert_eq!(n, a.ncols(), "matrix must be square for inversion");
let mut aug = Array2::<T>::zeros((n, 2 * n));
for i in 0..n {
for j in 0..n {
aug[[i, j]] = a[[i, j]];
}
aug[[i, n + i]] = T::one();
}
for col in 0..n {
let mut pivot_row = col;
let mut pivot_val = aug[[col, col]].abs();
for row in (col + 1)..n {
let v = aug[[row, col]].abs();
if v > pivot_val {
pivot_val = v;
pivot_row = row;
}
}
assert!(pivot_val > T::min_positive_val(), "singular matrix");
if pivot_row != col {
for j in 0..(2 * n) {
let tmp = aug[[col, j]];
aug[[col, j]] = aug[[pivot_row, j]];
aug[[pivot_row, j]] = tmp;
}
}
let inv_pivot = T::one() / aug[[col, col]];
for j in 0..(2 * n) {
aug[[col, j]] = aug[[col, j]] * inv_pivot;
}
for row in 0..n {
if row == col {
continue;
}
let factor = aug[[row, col]];
if factor == T::zero() {
continue;
}
for j in 0..(2 * n) {
let sub = factor * aug[[col, j]];
aug[[row, j]] -= sub;
}
}
}
let mut inv = Array2::<T>::zeros((n, n));
for i in 0..n {
for j in 0..n {
inv[[i, j]] = aug[[i, n + j]];
}
}
inv
}