rustebra 0.3.2

A hybrid no_std/alloc linear algebra crate for Rust, scaling from embedded targets to dynamic Krylov subspace solvers.
Documentation 
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
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122

use alloc::vec;
use alloc::vec::Vec;

use crate::scalar::Scalar;

use super::add::DimensionMismatch;
use super::{CscMatrix, CsrMatrix};

/// An abstract sparse linear operator that can be applied to a dense vector.
///
/// Implementing this trait for a storage format lets iterative Krylov solvers (CG, GMRES,
/// Lanczos, etc.) be written once, generically, rather than once per format.
///
/// # Examples
///
/// ```
/// use rustebra::sparse::{CsrMatrix, SparseLinearOp};
///
/// let eye = CsrMatrix::new(2, 2, vec![0, 1, 2], vec![0, 1], vec![1.0_f64, 1.0]).unwrap();
/// let y = eye.apply(&[3.0, 5.0]).unwrap();
/// assert_eq!(y, vec![3.0, 5.0]);
/// ```
pub trait SparseLinearOp<T: Scalar> {
    /// Number of rows (length of the output vector from [`apply`](SparseLinearOp::apply)).
    ///
    /// # Examples
    ///
    /// ```
    /// use rustebra::sparse::{CsrMatrix, SparseLinearOp};
    ///
    /// let m = CsrMatrix::<f64>::new(3, 5, vec![0; 4], vec![], vec![]).unwrap();
    /// assert_eq!(m.rows(), 3);
    /// ```
    fn rows(&self) -> usize;

    /// Number of columns (required length of the input vector to [`apply`](SparseLinearOp::apply)).
    ///
    /// # Examples
    ///
    /// ```
    /// use rustebra::sparse::{CsrMatrix, SparseLinearOp};
    ///
    /// let m = CsrMatrix::<f64>::new(3, 5, vec![0; 4], vec![], vec![]).unwrap();
    /// assert_eq!(m.cols(), 5);
    /// ```
    fn cols(&self) -> usize;

    /// Multiplies the operator by the dense column vector `x`, returning a new `Vec<T>`
    /// of length `self.rows()`.
    ///
    /// # Errors
    ///
    /// Returns `Err(DimensionMismatch)` when `x.len() != self.cols()`.
    ///
    /// # Examples
    ///
    /// ```
    /// use rustebra::sparse::{CscMatrix, SparseLinearOp};
    ///
    /// // [ 1  0 ]   [ 2 ]   [ 2 ]
    /// // [ 0  3 ] × [ 4 ] = [ 12 ]
    /// let m = CscMatrix::new(2, 2, vec![0, 1, 2], vec![0, 1], vec![1.0_f64, 3.0]).unwrap();
    /// let y = m.apply(&[2.0, 4.0]).unwrap();
    /// assert_eq!(y, vec![2.0, 12.0]);
    /// ```
    fn apply(&self, x: &[T]) -> Result<Vec<T>, DimensionMismatch>;
}

impl<T: Scalar> SparseLinearOp<T> for CsrMatrix<T> {
    fn rows(&self) -> usize {
        CsrMatrix::rows(self)
    }

    fn cols(&self) -> usize {
        CsrMatrix::cols(self)
    }

    fn apply(&self, x: &[T]) -> Result<Vec<T>, DimensionMismatch> {
        if x.len() != CsrMatrix::cols(self) {
            return Err(DimensionMismatch);
        }
        let mut out = vec![T::zero(); CsrMatrix::rows(self)];
        let row_ptr = self.row_ptr();
        let col_idx = self.col_indices();
        let vals = self.values();
        for r in 0..CsrMatrix::rows(self) {
            for k in row_ptr[r]..row_ptr[r + 1] {
                let prev = out[r];
                out[r] = prev.add(vals[k].mul(x[col_idx[k]]));
            }
        }
        Ok(out)
    }
}

impl<T: Scalar> SparseLinearOp<T> for CscMatrix<T> {
    fn rows(&self) -> usize {
        CscMatrix::rows(self)
    }

    fn cols(&self) -> usize {
        CscMatrix::cols(self)
    }

    fn apply(&self, x: &[T]) -> Result<Vec<T>, DimensionMismatch> {
        if x.len() != CscMatrix::cols(self) {
            return Err(DimensionMismatch);
        }
        let mut out = vec![T::zero(); CscMatrix::rows(self)];
        let col_ptr = self.col_ptr();
        let row_idx = self.row_indices();
        let vals = self.values();
        for c in 0..CscMatrix::cols(self) {
            for k in col_ptr[c]..col_ptr[c + 1] {
                let r = row_idx[k];
                let prev = out[r];
                out[r] = prev.add(vals[k].mul(x[c]));
            }
        }
        Ok(out)
    }
}