impl Index/IndexMut & impl opnorm

Author: doraneko94

src/lapack/tridiagonal.rs

@@ -10,14 +10,17 @@ use super::{into_result, Pivot, Transpose};
 
 use crate::error::*;
 use crate::layout::MatrixLayout;
+use crate::opnorm::*;
 use crate::tridiagonal::{LUFactorizedTriDiagonal, TriDiagonal};
 use crate::types::*;
 
 /// Wraps `*gttrf`, `*gtcon` and `*gttrs`
 pub trait TriDiagonal_: Scalar + Sized {
     /// Computes the LU factorization of a tridiagonal `m x n` matrix `a` using
     /// partial pivoting with row interchanges.
-    unsafe fn lu_tridiagonal(a: &mut TriDiagonal<Self>) -> Result<(Array1<Self>, Pivot)>;
+    unsafe fn lu_tridiagonal(
+        a: &mut TriDiagonal<Self>,
+    ) -> Result<(Array1<Self>, Self::Real, Pivot)>;
     /// Estimates the the reciprocal of the condition number of the tridiagonal matrix in 1-norm.
     unsafe fn rcond_tridiagonal(lu: &LUFactorizedTriDiagonal<Self>) -> Result<Self::Real>;
     unsafe fn solve_tridiagonal(
@@ -31,15 +34,18 @@ pub trait TriDiagonal_: Scalar + Sized {
 macro_rules! impl_tridiagonal {
     ($scalar:ty, $gttrf:path, $gtcon:path, $gttrs:path) => {
         impl TriDiagonal_ for $scalar {
-            unsafe fn lu_tridiagonal(a: &mut TriDiagonal<Self>) -> Result<(Array1<Self>, Pivot)> {
+            unsafe fn lu_tridiagonal(
+                a: &mut TriDiagonal<Self>,
+            ) -> Result<(Array1<Self>, Self::Real, Pivot)> {
                 let (n, _) = a.l.size();
+                let anom = a.opnorm_one()?;
                 let dl = a.dl.as_slice_mut().unwrap();
                 let d = a.d.as_slice_mut().unwrap();
                 let du = a.du.as_slice_mut().unwrap();
                 let mut du2 = vec![Zero::zero(); (n - 2) as usize];
                 let mut ipiv = vec![0; n as usize];
                 let info = $gttrf(n, dl, d, du, &mut du2, &mut ipiv);
-                into_result(info, (arr1(&du2), ipiv))
+                into_result(info, (arr1(&du2), anom, ipiv))
             }
 
             unsafe fn rcond_tridiagonal(lu: &LUFactorizedTriDiagonal<Self>) -> Result<Self::Real> {
@@ -49,7 +55,7 @@ macro_rules! impl_tridiagonal {
                 let du = lu.a.du.as_slice().unwrap();
                 let du2 = lu.du2.as_slice().unwrap();
                 let ipiv = &lu.ipiv;
-                let anorm = lu.a.n1;
+                let anorm = lu.anom;
                 let mut rcond = Self::Real::zero();
                 let info = $gtcon(
                     NormType::One as u8,

src/opnorm.rs

@@ -2,8 +2,10 @@
 
 use ndarray::*;
 
+use crate::convert::*;
 use crate::error::*;
 use crate::layout::*;
+use crate::tridiagonal::TriDiagonal;
 use crate::types::*;
 
 pub use crate::lapack::NormType;
@@ -46,3 +48,54 @@ where
         Ok(unsafe { A::opnorm(t, l, a) })
     }
 }
+
+impl<A> OperationNorm for TriDiagonal<A>
+where
+    A: Scalar + Lapack,
+{
+    type Output = A::Real;
+
+    fn opnorm(&self, t: NormType) -> Result<Self::Output> {
+        let arr = match t {
+            NormType::One => {
+                let zl: Array1<A> = Array::zeros(1);
+                let zu: Array1<A> = Array::zeros(1);
+                let dl = stack![Axis(0), self.dl.to_owned(), zl];
+                let du = stack![Axis(0), zu, self.du.to_owned()];
+                let arr = stack![
+                    Axis(0),
+                    into_row(du),
+                    into_row(self.d.to_owned()),
+                    into_row(dl)
+                ];
+                arr
+            }
+            NormType::Infinity => {
+                let zl: Array1<A> = Array::zeros(1);
+                let zu: Array1<A> = Array::zeros(1);
+                let dl = stack![Axis(0), zl, self.dl.to_owned()];
+                let du = stack![Axis(0), self.du.to_owned(), zu];
+                let arr = stack![
+                    Axis(1),
+                    into_col(dl),
+                    into_col(self.d.to_owned()),
+                    into_col(du)
+                ];
+                arr
+            }
+            NormType::Frobenius => {
+                let arr = stack![
+                    Axis(1),
+                    into_row(self.dl.to_owned()),
+                    into_row(self.d.to_owned()),
+                    into_row(self.du.to_owned())
+                ];
+                arr
+            }
+        };
+
+        let l = arr.layout()?;
+        let a = arr.as_allocated()?;
+        Ok(unsafe { A::opnorm(t, l, a) })
+    }
+}

src/tridiagonal.rs

@@ -2,26 +2,23 @@
 //! &
 //! Methods for tridiagonal matrices
 
+use std::ops::{Index, IndexMut};
+
 use cauchy::Scalar;
 use ndarray::*;
 use num_traits::One;
 
-use crate::opnorm::OperationNorm;
-
 use super::convert::*;
 use super::error::*;
 use super::lapack::*;
 use super::layout::*;
 
 /// Represents a tridiagonal matrix as 3 one-dimensional vectors.
-/// This struct also holds the layout and 1-norm of the raw matrix
-/// for some methods (eg. rcond_tridiagonal()).
-#[derive(Clone)]
+/// This struct also holds the layout of the raw matrix.
+#[derive(Clone, PartialEq)]
 pub struct TriDiagonal<A: Scalar> {
     /// layout of raw matrix
     pub l: MatrixLayout,
-    /// the one norm of raw matrix
-    pub n1: <A as Scalar>::Real,
     /// (n-1) sub-diagonal elements of matrix.
     pub dl: Array1<A>,
     /// (n) diagonal elements of matrix.
@@ -30,10 +27,73 @@ pub struct TriDiagonal<A: Scalar> {
     pub du: Array1<A>,
 }
 
+pub trait TridiagIndex {
+    fn to_tuple(&self) -> (i32, i32);
+}
+impl TridiagIndex for [Ix; 2] {
+    fn to_tuple(&self) -> (i32, i32) {
+        (self[0] as i32, self[1] as i32)
+    }
+}
+
+fn debug_bounds_check_tridiag(n: i32, row: i32, col: i32) {
+    if std::cmp::max(row, col) >= n {
+        panic!(
+            "ndarray: index {:?} is out of bounds for array of shape {}",
+            [row, col],
+            n
+        );
+    }
+}
+
+impl<A, I> Index<I> for TriDiagonal<A>
+where
+    A: Scalar,
+    I: TridiagIndex,
+{
+    type Output = A;
+    #[inline]
+    fn index(&self, index: I) -> &A {
+        let (n, _) = self.l.size();
+        let (row, col) = index.to_tuple();
+        debug_bounds_check_tridiag(n, row, col);
+        match row - col {
+            0 => &self.d[row as usize],
+            1 => &self.dl[col as usize],
+            -1 => &self.du[row as usize],
+            _ => panic!(
+                "ndarray-linalg::tridiagonal: index {:?} is not tridiagonal element",
+                [row, col]
+            ),
+        }
+    }
+}
+
+impl<A, I> IndexMut<I> for TriDiagonal<A>
+where
+    A: Scalar,
+    I: TridiagIndex,
+{
+    #[inline]
+    fn index_mut(&mut self, index: I) -> &mut A {
+        let (n, _) = self.l.size();
+        let (row, col) = index.to_tuple();
+        debug_bounds_check_tridiag(n, row, col);
+        match row - col {
+            0 => &mut self.d[row as usize],
+            1 => &mut self.dl[col as usize],
+            -1 => &mut self.du[row as usize],
+            _ => panic!(
+                "ndarray-linalg::tridiagonal: index {:?} is not tridiagonal element",
+                [row, col]
+            ),
+        }
+    }
+}
+
 /// An interface for making a TriDiagonal struct.
 pub trait ToTriDiagonal<A: Scalar> {
     /// Extract tridiagonal elements and layout of the raw matrix.
-    /// And also calculate 1-norm.
     ///
     /// If the raw matrix has some non-tridiagonal elements,
     /// they will be ignored.
@@ -53,12 +113,11 @@ where
         if n < 2 {
             panic!("Cannot make a tridiagonal matrix of shape=(1, 1)!");
         }
-        let n1 = self.opnorm_one()?;
 
         let dl = self.slice(s![1..n, 0..n - 1]).diag().to_owned();
         let d = self.diag().to_owned();
         let du = self.slice(s![0..n - 1, 1..n]).diag().to_owned();
-        Ok(TriDiagonal { l, n1, dl, d, du })
+        Ok(TriDiagonal { l, dl, d, du })
     }
 }
 
@@ -130,13 +189,14 @@ pub trait SolveTriDiagonalInplace<A: Scalar, D: Dimension> {
 pub struct LUFactorizedTriDiagonal<A: Scalar> {
     /// A tridiagonal matrix which consists of
     /// - l : layout of raw matrix
-    /// - n1: the one norm of raw matrix
     /// - dl: (n-1) multipliers that define the matrix L.
     /// - d : (n) diagonal elements of the upper triangular matrix U.
     /// - du: (n-1) elements of the first super-diagonal of U.
     pub a: TriDiagonal<A>,
     /// (n-2) elements of the second super-diagonal of U.
     pub du2: Array1<A>,
+    /// 1-norm of raw matrix (used in .rcond_tridiagonal()).
+    pub anom: A::Real,
     /// The pivot indices that define the permutation matrix `P`.
     pub ipiv: Pivot,
 }
@@ -598,10 +658,11 @@ where
     A: Scalar + Lapack,
 {
     fn factorize_tridiagonal_into(mut self) -> Result<LUFactorizedTriDiagonal<A>> {
-        let (du2, ipiv) = unsafe { A::lu_tridiagonal(&mut self)? };
+        let (du2, anom, ipiv) = unsafe { A::lu_tridiagonal(&mut self)? };
         Ok(LUFactorizedTriDiagonal {
             a: self,
             du2: du2,
+            anom: anom,
             ipiv: ipiv,
         })
     }
@@ -613,8 +674,8 @@ where
 {
     fn factorize_tridiagonal(&self) -> Result<LUFactorizedTriDiagonal<A>> {
         let mut a = self.clone();
-        let (du2, ipiv) = unsafe { A::lu_tridiagonal(&mut a)? };
-        Ok(LUFactorizedTriDiagonal { a, du2, ipiv })
+        let (du2, anom, ipiv) = unsafe { A::lu_tridiagonal(&mut a)? };
+        Ok(LUFactorizedTriDiagonal { a, du2, anom, ipiv })
     }
 }
 
@@ -625,8 +686,8 @@ where
 {
     fn factorize_tridiagonal(&self) -> Result<LUFactorizedTriDiagonal<A>> {
         let mut a = self.to_tridiagonal()?;
-        let (du2, ipiv) = unsafe { A::lu_tridiagonal(&mut a)? };
-        Ok(LUFactorizedTriDiagonal { a, du2, ipiv })
+        let (du2, anom, ipiv) = unsafe { A::lu_tridiagonal(&mut a)? };
+        Ok(LUFactorizedTriDiagonal { a, du2, anom, ipiv })
     }
 }
 

tests/tridiagonal.rs

@@ -10,6 +10,37 @@ fn to_tridiagonal() {
     assert_close_l2!(&t.du, &arr1(&[2.0, 6.0]), 1e-7);
 }
 
+#[test]
+fn tridiagonal_index() {
+    let a: Array2<f64> = arr2(&[[1.0, 2.0, 3.0], [4.0, 5.0, 6.0], [7.0, 8.0, 9.0]]);
+    let t1 = a.to_tridiagonal().unwrap();
+    let mut t2 = Array2::<f64>::eye(3).to_tridiagonal().unwrap();
+    t2[[0, 1]] = 2.0;
+    t2[[1, 0]] = 4.0;
+    t2[[1, 1]] += 4.0;
+    t2[[1, 2]] = 6.0;
+    t2[[2, 1]] = 8.0;
+    t2[[2, 2]] += 8.0;
+    assert_eq!(t1.dl, t2.dl);
+    assert_eq!(t1.d, t2.d);
+    assert_eq!(t1.du, t2.du);
+}
+
+#[test]
+fn opnorm_tridiagonal() {
+    let mut a: Array2<f64> = random((4, 4));
+    a[[0, 2]] = 0.0;
+    a[[0, 3]] = 0.0;
+    a[[1, 3]] = 0.0;
+    a[[2, 0]] = 0.0;
+    a[[3, 0]] = 0.0;
+    a[[3, 1]] = 0.0;
+    let t = a.to_tridiagonal().unwrap();
+    assert_aclose!(a.opnorm_one().unwrap(), t.opnorm_one().unwrap(), 1e-7);
+    assert_aclose!(a.opnorm_inf().unwrap(), t.opnorm_inf().unwrap(), 1e-7);
+    assert_aclose!(a.opnorm_fro().unwrap(), t.opnorm_fro().unwrap(), 1e-7);
+}
+
 #[test]
 fn solve_tridiagonal_f64() {
     // https://www.nag-j.co.jp/lapack/dgttrs.htm