nalgebra-sparse 0.11.0

Sparse matrix computation based on nalgebra.
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
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149

use crate::csr::CsrMatrix;
use crate::ops::Op;
use crate::ops::serial::OperationError;
use crate::ops::serial::cs::{
    spadd_cs_prealloc, spmm_cs_dense, spmm_cs_prealloc, spmm_cs_prealloc_unchecked,
};
use nalgebra::{ClosedAddAssign, ClosedMulAssign, DMatrixView, DMatrixViewMut, Scalar};
use num_traits::{One, Zero};
use std::borrow::Cow;

/// Sparse-dense matrix-matrix multiplication `C <- beta * C + alpha * op(A) * op(B)`.
pub fn spmm_csr_dense<'a, T>(
    beta: T,
    c: impl Into<DMatrixViewMut<'a, T>>,
    alpha: T,
    a: Op<&CsrMatrix<T>>,
    b: Op<impl Into<DMatrixView<'a, T>>>,
) where
    T: Scalar + ClosedAddAssign + ClosedMulAssign + Zero + One,
{
    let b = b.convert();
    spmm_csr_dense_(beta, c.into(), alpha, a, b)
}

fn spmm_csr_dense_<T>(
    beta: T,
    c: DMatrixViewMut<'_, T>,
    alpha: T,
    a: Op<&CsrMatrix<T>>,
    b: Op<DMatrixView<'_, T>>,
) where
    T: Scalar + ClosedAddAssign + ClosedMulAssign + Zero + One,
{
    assert_compatible_spmm_dims!(c, a, b);
    spmm_cs_dense(beta, c, alpha, a.map_same_op(|a| &a.cs), b)
}

/// Sparse matrix addition `C <- beta * C + alpha * op(A)`.
///
/// # Errors
///
/// If the pattern of `c` does not accommodate all the non-zero entries in `a`, an error is
/// returned.
///
/// # Panics
///
/// Panics if the dimensions of the matrices involved are not compatible with the expression.
pub fn spadd_csr_prealloc<T>(
    beta: T,
    c: &mut CsrMatrix<T>,
    alpha: T,
    a: Op<&CsrMatrix<T>>,
) -> Result<(), OperationError>
where
    T: Scalar + ClosedAddAssign + ClosedMulAssign + Zero + One,
{
    assert_compatible_spadd_dims!(c, a);
    spadd_cs_prealloc(beta, &mut c.cs, alpha, a.map_same_op(|a| &a.cs))
}

/// Sparse-sparse matrix multiplication, `C <- beta * C + alpha * op(A) * op(B)`.
///
/// # Errors
///
/// If the pattern of `C` is not able to hold the result of the operation, an error is returned.
///
/// # Panics
///
/// Panics if the dimensions of the matrices involved are not compatible with the expression.
pub fn spmm_csr_prealloc<T>(
    beta: T,
    c: &mut CsrMatrix<T>,
    alpha: T,
    a: Op<&CsrMatrix<T>>,
    b: Op<&CsrMatrix<T>>,
) -> Result<(), OperationError>
where
    T: Scalar + ClosedAddAssign + ClosedMulAssign + Zero + One,
{
    assert_compatible_spmm_dims!(c, a, b);

    use Op::NoOp;

    match (&a, &b) {
        (NoOp(a), NoOp(b)) => spmm_cs_prealloc(beta, &mut c.cs, alpha, &a.cs, &b.cs),
        _ => spmm_csr_transposed(beta, c, alpha, a, b, spmm_csr_prealloc),
    }
}

/// Faster sparse-sparse matrix multiplication, `C <- beta * C + alpha * op(A) * op(B)`.
/// This will not return an error even if the patterns don't match.
/// Should be used for situations where pattern creation immediately precedes multiplication.
///
/// Panics if the dimensions of the matrices involved are not compatible with the expression.
pub fn spmm_csr_prealloc_unchecked<T>(
    beta: T,
    c: &mut CsrMatrix<T>,
    alpha: T,
    a: Op<&CsrMatrix<T>>,
    b: Op<&CsrMatrix<T>>,
) -> Result<(), OperationError>
where
    T: Scalar + ClosedAddAssign + ClosedMulAssign + Zero + One,
{
    assert_compatible_spmm_dims!(c, a, b);

    use Op::NoOp;

    match (&a, &b) {
        (NoOp(a), NoOp(b)) => spmm_cs_prealloc_unchecked(beta, &mut c.cs, alpha, &a.cs, &b.cs),
        _ => spmm_csr_transposed(beta, c, alpha, a, b, spmm_csr_prealloc_unchecked),
    }
}

fn spmm_csr_transposed<T, F>(
    beta: T,
    c: &mut CsrMatrix<T>,
    alpha: T,
    a: Op<&CsrMatrix<T>>,
    b: Op<&CsrMatrix<T>>,
    spmm_kernel: F,
) -> Result<(), OperationError>
where
    T: Scalar + ClosedAddAssign + ClosedMulAssign + Zero + One,
    F: Fn(
        T,
        &mut CsrMatrix<T>,
        T,
        Op<&CsrMatrix<T>>,
        Op<&CsrMatrix<T>>,
    ) -> Result<(), OperationError>,
{
    use Op::{NoOp, Transpose};

    // Currently we handle transposition by explicitly precomputing transposed matrices
    // and calling the operation again without transposition
    let a_ref: &CsrMatrix<T> = a.inner_ref();
    let b_ref: &CsrMatrix<T> = b.inner_ref();
    let (a, b) = {
        use Cow::*;
        match (&a, &b) {
            (NoOp(_), NoOp(_)) => unreachable!(),
            (Transpose(a), NoOp(_)) => (Owned(a.transpose()), Borrowed(b_ref)),
            (NoOp(_), Transpose(b)) => (Borrowed(a_ref), Owned(b.transpose())),
            (Transpose(a), Transpose(b)) => (Owned(a.transpose()), Owned(b.transpose())),
        }
    };
    spmm_kernel(beta, c, alpha, NoOp(a.as_ref()), NoOp(b.as_ref()))
}