-
Notifications
You must be signed in to change notification settings - Fork 112
/
sparse.py
39 lines (27 loc) · 1016 Bytes
/
sparse.py
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
"""Sparse matrix interface to internal sparse matrix operations."""
import numpy as np
from scipy.sparse import csr_matrix
try:
# scipy >=1.8
from scipy.sparse._sputils import upcast_char
except ImportError:
# scipy <1.8
from scipy.sparse.sputils import upcast_char
from .. import amg_core
class csr(csr_matrix): # noqa: N801
"""CSR class to redefine operations.
The purpose of this class is to redefine the matvec in scipy.sparse
"""
def _mul_vector(self, other):
"""Matrix-vector multiplication.
Identical to scipy.sparse with an in internal call to
pyamg.amg_core.sparse.csr_matvec
"""
M, N = self.shape
# output array
result = np.zeros(M, dtype=upcast_char(self.dtype.char,
other.dtype.char))
amg_core.csr_matvec(M, N, self.indptr, self.indices, self.data,
other, result)
return result
csr.__doc__ += csr_matrix.__doc__