-
Notifications
You must be signed in to change notification settings - Fork 240
/
dense_matrix.py
122 lines (103 loc) · 3.78 KB
/
dense_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
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
"""Define the DenseMatrix class."""
import numpy as np
from openmdao.matrices.coo_matrix import COOMatrix
# NOTE: DenseMatrix is inherited from COOMatrix so that we can easily handle use cases
# where partials overlap the same matrix entries, as in the case of repeated
# src_indices entries. This does require additional memory above storing just
# the dense matrix, but it's worth it because the code is simpler and more robust.
class DenseMatrix(COOMatrix):
"""
Dense global matrix.
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, system=None):
"""
Allocate the matrix.
Parameters
----------
num_rows : int
number of rows in the matrix.
num_cols : int
number of cols in the matrix.
system : <System>
owning system.
"""
super()._build(num_rows, num_cols)
self._coo = self._matrix
def _prod(self, in_vec, mode, mask=None):
"""
Perform a matrix vector product.
Parameters
----------
in_vec : ndarray[:]
incoming vector to multiply.
mode : str
'fwd' or 'rev'.
mask : ndarray of type bool, or None
Array used to mask out part of the input vector.
Returns
-------
ndarray[:]
vector resulting from the product.
"""
# when we have a derivative based solver at a level below the
# group that owns the AssembledJacobian, we need to use only
# the part of the matrix that is relevant to the lower level
# system.
mat = self._matrix
if mode == 'fwd':
if mask is None:
return mat.dot(in_vec)
else:
# Use the special dot product function from masking module so that we
# ignore masked parts.
return np.ma.dot(mat, np.ma.array(in_vec, mask=mask))
else: # rev
if mask is None:
return mat.T.dot(in_vec)
else:
# Mask need to be applied to ext_mtx so that we can ignore multiplication
# by certain columns.
mat_T = mat.T
arrmask = np.zeros(mat_T.shape, dtype=bool)
arrmask[mask, :] = True
masked_mtx = np.ma.array(mat_T, mask=arrmask, fill_value=0.0)
masked_product = np.ma.dot(masked_mtx, in_vec).flatten()
return np.ma.filled(masked_product, fill_value=0.0)
def _create_mask_cache(self, d_inputs):
"""
Create masking array for this matrix.
Note : this only applies when this Matrix is an 'ext_mtx' inside of a
Jacobian object.
Parameters
----------
d_inputs : Vector
The inputs linear vector.
Returns
-------
ndarray or None
The mask array or None.
"""
if d_inputs._in_matvec_context():
sub = d_inputs._names
mask = np.ones(len(d_inputs), dtype=bool)
for key, val in self._metadata.items():
if key[1] in sub:
mask[val[1]] = False
return mask
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.toarray()