-
Notifications
You must be signed in to change notification settings - Fork 240
/
csr_matrix.py
49 lines (40 loc) · 1.34 KB
/
csr_matrix.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
40
41
42
43
44
45
46
47
48
49
"""Define the CSRmatrix class."""
from scipy.sparse import coo_matrix
from openmdao.matrices.coo_matrix import COOMatrix
class CSRMatrix(COOMatrix):
"""
Sparse matrix in Compressed Row Storage format.
Parameters
----------
comm : MPI.Comm or <FakeComm>
Communicator of the top-level system that owns the <Jacobian>.
is_internal : bool
If True, this is the int_mtx of an AssembledJacobian.
"""
def _build(self, num_rows, num_cols, in_ranges, out_ranges):
"""
Allocate the matrix.
Parameters
----------
num_rows : int
number of rows in the matrix.
num_cols : int
number of cols in the matrix.
in_ranges : dict
Maps input var name to column range.
out_ranges : dict
Maps output var name to row range.
"""
super()._build(num_rows, num_cols, in_ranges, out_ranges)
self._coo = self._matrix
def _pre_update(self):
"""
Do anything that needs to be done at the end of AssembledJacobian._update.
"""
self._matrix = self._coo
def _post_update(self):
"""
Do anything that needs to be done at the end of AssembledJacobian._update.
"""
# this will add any repeated entries together
self._matrix = self._coo.tocsc()