rustebra 0.4.0

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

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

use crate::scalar::Scalar;

use super::CsrMatrix;
use super::add::DimensionMismatch;
use super::sorted_csr::SortedCsrMatrix;

/// Multiplies two CSR sparse matrices (`a × b`), returning the product as a
/// [`SortedCsrMatrix`].
///
/// `a` is `m × k` and `b` is `k × n`; the result is `m × n`.  The algorithm uses a
/// dense row accumulator of size `n` so the working memory is `O(n)` per output row.
/// Output column indices within each row are guaranteed to be in ascending order, which
/// is why the result is a [`SortedCsrMatrix`] rather than a plain `CsrMatrix`.
///
/// # Errors
///
/// Returns `Err(DimensionMismatch::Shape)` when `a.cols() != b.rows()`.
///
/// # Examples
///
/// ```
/// use rustebra::sparse::{CsrMatrix, spmm_csr};
///
/// // [1  0]   [1  2]   [1  2]
/// // [0  2] × [3  4] = [6  8]
/// let a = CsrMatrix::new(2, 2, vec![0, 1, 2], vec![0, 1], vec![1.0_f64, 2.0]).unwrap();
/// let b = CsrMatrix::new(2, 2, vec![0, 2, 4], vec![0, 1, 0, 1],
///                        vec![1.0_f64, 2.0, 3.0, 4.0]).unwrap();
/// let c = spmm_csr(&a, &b).unwrap();
/// assert_eq!(c.row_ptr(),     &[0, 2, 4]);
/// assert_eq!(c.col_indices(), &[0, 1, 0, 1]);
/// assert_eq!(c.values(),      &[1.0, 2.0, 6.0, 8.0]);
/// ```
pub fn spmm_csr<T: Scalar>(
    a: &CsrMatrix<T>,
    b: &CsrMatrix<T>,
) -> Result<SortedCsrMatrix<T>, DimensionMismatch> {
    if a.cols() != b.rows() {
        return Err(DimensionMismatch::Shape);
    }
    let m = a.rows();
    let n = b.cols();

    let mut out_row_ptr = vec![0u32; m + 1];
    let mut out_col: Vec<u32> = Vec::new();
    let mut out_val: Vec<T> = Vec::new();

    // Dense accumulator: `dense[j]` holds the accumulated value for column j in the
    // current output row.  `touched` tracks which columns are non-zero so we can reset
    // only those entries after each row, keeping the overall work proportional to nnz.
    let mut dense = vec![T::zero(); n];
    let mut touched: Vec<usize> = Vec::new();
    // `in_touched[j]` is true when column j has been placed in `touched`.
    let mut in_touched = vec![false; n];

    let a_row_ptr = a.row_ptr();
    let a_col = a.col_indices();
    let a_val = a.values();
    let b_row_ptr = b.row_ptr();
    let b_col = b.col_indices();
    let b_val = b.values();

    for r in 0..m {
        for ka in a_row_ptr[r] as usize..a_row_ptr[r + 1] as usize {
            let k = a_col[ka] as usize;
            let av = a_val[ka];
            for kb in b_row_ptr[k] as usize..b_row_ptr[k + 1] as usize {
                let j = b_col[kb] as usize;
                let bv = b_val[kb];
                let prev = dense[j];
                dense[j] = prev.add(av.mul(bv));
                if !in_touched[j] {
                    touched.push(j);
                    in_touched[j] = true;
                }
            }
        }
        // Emit output entries sorted by column index.
        touched.sort_unstable();
        for &col in &touched {
            if dense[col] != T::zero() {
                out_col.push(col as u32);
                out_val.push(dense[col]);
            }
            dense[col] = T::zero();
            in_touched[col] = false;
        }
        touched.clear();
        out_row_ptr[r + 1] =
            u32::try_from(out_col.len()).map_err(|_| DimensionMismatch::DimensionOverflow)?;
    }

    Ok(SortedCsrMatrix::from_sorted_unchecked(CsrMatrix::new_raw(
        m,
        n,
        out_row_ptr,
        out_col,
        out_val,
    )))
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn s2_spmm_exact_cancellation_zeros() {
        let a = CsrMatrix::new(1, 2, vec![0, 2], vec![0, 1], vec![1.0_f64, -1.0]).unwrap();
        let b = CsrMatrix::new(2, 1, vec![0, 1, 2], vec![0, 0], vec![1.0_f64, 1.0]).unwrap();
        let c = spmm_csr(&a, &b).unwrap();
        assert_eq!(c.nnz(), 0);
    }
}