#![allow(non_snake_case)]
use ndarray::prelude::*;
use ndarray::Zip;
fn main() {
println!("Example 1:");
let A: Array2<f64> = arr2(&[
[1.0, 3.0, 5.0],
[2.0, 4.0, 7.0],
[1.0, 1.0, 0.0],
]);
println!("A \n {}", A);
let (L, U, P) = lu_decomp(&A);
println!("L \n {}", L);
println!("U \n {}", U);
println!("P \n {}", P);
println!("\nExample 2:");
let A: Array2<f64> = arr2(&[
[11.0, 9.0, 24.0, 2.0],
[1.0, 5.0, 2.0, 6.0],
[3.0, 17.0, 18.0, 1.0],
[2.0, 5.0, 7.0, 1.0],
]);
let A:Array2<f64> = array![
[4.3552 , 6.25851, 4.12662, 1.93708, 0.21272, 3.25683, 6.53326],
[4.24746, 1.84137, 6.71904, 0.59754, 3.5806 , 3.63597, 5.347 ],
[2.30479, 1.70591, 3.05354, 1.82188, 5.27839, 7.9166 , 2.04607],
[2.40158, 6.38524, 7.90296, 4.69683, 6.63801, 7.32958, 1.45936],
[0.42456, 6.47456, 1.55398, 8.28979, 4.20987, 0.90401, 4.94587],
[5.78903, 1.92032, 6.20261, 5.78543, 1.94331, 8.25178, 7.47273],
[1.44797, 7.41157, 7.69495, 8.90113, 3.05983, 0.41582, 6.42932]];
println!("A \n {}", A);
let (mut L, mut U, mut P) = lu_decomp(&A);
println!("L \n {}", L);
println!("U \n {}", U);
println!("P \n {}", P);
let mut _q = lu_decomp(&A);
for i in 0 .. 1000000 {
_q = lu_decomp(&A);
}
}
fn lu_decomp<T: NdFloat>(A: &Array2<T>) -> (Array2<T>, Array2<T>, Array2<T>) {
fn pivot<T: NdFloat>(A: &Array2<T>) -> Array2<T> {
fn swap<T: NdFloat>(A: &mut Array2<T>, ir1: usize, ir2: usize) {
let (.., mut rest) = A.view_mut().split_at(Axis(0), ir1);
let (r0, mut rest) = rest.view_mut().split_at(Axis(0), 1);
let (.., mut rest) = rest.view_mut().split_at(Axis(0), ir2 - ir1 - 1);
let (r1, ..) = rest.view_mut().split_at(Axis(0), 1);
Zip::from(r0).and(r1).for_each(std::mem::swap);
}
let n = A.raw_dim()[0];
let mut P: Array2<T> = Array::eye(n);
for (idx, col) in A.axis_iter(Axis(1)).enumerate() {
// find idx of maximum value in column i
let mut mp = idx;
for i in idx .. n {
if col[mp].abs() < col[i].abs() {
mp = i;
}
}
// swap rows if necessary
if mp != idx {
swap(&mut P, idx, mp);
}
}
P
}
let d = A.raw_dim();
let n = d[0];
assert_eq!(n, d[1], "LU decomposition must take a square matrix.");
let P = pivot(&A);
let pA = P.dot(A);
let mut L: Array2<T> = Array::eye(n);
let mut U: Array2<T> = Array::zeros((n, n));
for c in 0 .. n {
for r in 0 .. n {
let pAs = pA[[r, c]] - U.slice(s![0..r, c]).dot(&L.slice(s![r, 0..r]));
if r < c + 1 { // U
U[[r, c]] = pAs;
} else { // L
L[[r, c]] = (pAs) / U[[c, c]];
}
}
}
(L, U, P)
}