1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68


use ndarray::{Ix2, Array, RcArray, NdFloat, ArrayBase, DataMut};

use matrix::{Matrix, MFloat};
use square::SquareMatrix;
use error::LinalgError;
use solve::ImplSolve;

pub trait TriangularMatrix: Matrix + SquareMatrix {
    /// solve a triangular system with upper triangular matrix
    fn solve_upper(&self, Self::Vector) -> Result<Self::Vector, LinalgError>;
    /// solve a triangular system with lower triangular matrix
    fn solve_lower(&self, Self::Vector) -> Result<Self::Vector, LinalgError>;
}

impl<A: MFloat> TriangularMatrix for Array<A, Ix2> {
    fn solve_upper(&self, b: Self::Vector) -> Result<Self::Vector, LinalgError> {
        self.check_square()?;
        let (n, _) = self.size();
        let layout = self.layout()?;
        let a = self.as_slice_memory_order().unwrap();
        let x = ImplSolve::solve_triangle(layout, 'U' as u8, n, a, b.into_raw_vec())?;
        Ok(Array::from_vec(x))
    }
    fn solve_lower(&self, b: Self::Vector) -> Result<Self::Vector, LinalgError> {
        self.check_square()?;
        let (n, _) = self.size();
        let layout = self.layout()?;
        let a = self.as_slice_memory_order().unwrap();
        let x = ImplSolve::solve_triangle(layout, 'L' as u8, n, a, b.into_raw_vec())?;
        Ok(Array::from_vec(x))
    }
}

impl<A: MFloat> TriangularMatrix for RcArray<A, Ix2> {
    fn solve_upper(&self, b: Self::Vector) -> Result<Self::Vector, LinalgError> {
        // XXX unnecessary clone
        let x = self.to_owned().solve_upper(b.to_owned())?;
        Ok(x.into_shared())
    }
    fn solve_lower(&self, b: Self::Vector) -> Result<Self::Vector, LinalgError> {
        // XXX unnecessary clone
        let x = self.to_owned().solve_lower(b.to_owned())?;
        Ok(x.into_shared())
    }
}

pub fn drop_upper<A: NdFloat, S>(mut a: ArrayBase<S, Ix2>) -> ArrayBase<S, Ix2>
    where S: DataMut<Elem = A>
{
    for ((i, j), val) in a.indexed_iter_mut() {
        if i < j {
            *val = A::zero();
        }
    }
    a
}

pub fn drop_lower<A: NdFloat, S>(mut a: ArrayBase<S, Ix2>) -> ArrayBase<S, Ix2>
    where S: DataMut<Elem = A>
{
    for ((i, j), val) in a.indexed_iter_mut() {
        if i > j {
            *val = A::zero();
        }
    }
    a
}