Coverage for /usr/share/miniconda3/envs/dolfin/lib/python3.8/site-packages/block/algebraic/trilinos/Epetra.py: 0%
264 statements
« prev ^ index » next coverage.py v7.2.1, created at 2023-03-20 13:03 +0000
1from __future__ import division
2from __future__ import print_function
4from builtins import map
5from builtins import str
6"""Module implementing algebraic operations on Epetra matrices: Diag, InvDiag
7etc, as well as the collapse() method which performs matrix addition and
8multiplication.
9"""
11from block.block_base import block_base
13class diag_op(block_base):
14 """Base class for diagonal Epetra operators (represented by an Epetra vector)."""
15 from block.object_pool import shared_vec_pool
17 def __init__(self, v):
18 from PyTrilinos import Epetra
19 import dolfin
20 assert isinstance(v, (Epetra.MultiVector, Epetra.Epetra_MultiVector))
21 self.v = v
23 def _transpose(self):
24 return self
26 def matvec(self, b):
27 try:
28 b_vec = b.down_cast().vec()
29 except AttributeError:
30 return NotImplemented
32 x = self.create_vec(dim=1)
33 if len(x) != len(b):
34 raise RuntimeError(
35 'incompatible dimensions for %s matvec, %d != %d'%(self.__class__.__name__,len(x),len(b)))
37 x.down_cast().vec().Multiply(1.0, self.v, b_vec, 0.0)
38 return x
40 transpmult = matvec
42 def matmat(self, other):
43 from PyTrilinos import Epetra
44 try:
45 from block.block_util import isscalar
46 if isscalar(other):
47 x = Epetra.Vector(self.v)
48 x.Scale(other)
49 return diag_op(x)
50 other = other.down_cast()
51 if hasattr(other, 'mat'):
52 C = Epetra.FECrsMatrix(other.mat())
53 C.LeftScale(self.v)
54 return matrix_op(C)
55 else:
56 x = Epetra.Vector(self.v.Map())
57 x.Multiply(1.0, self.v, other.vec(), 0.0)
58 return diag_op(x)
59 except AttributeError:
60 raise TypeError("can't extract matrix data from type '%s'"%str(type(other)))
62 def add(self, other, lscale=1.0, rscale=1.0):
63 from block.block_util import isscalar
64 from PyTrilinos import Epetra
65 try:
66 if isscalar(other):
67 x = Epetra.Vector(self.v.Map())
68 x.PutScalar(other)
69 other = diag_op(x)
70 other = other.down_cast()
71 if isinstance(other, matrix_op):
72 return other.add(self, lscale=rscale, rscale=lscale)
73 else:
74 x = Epetra.Vector(self.v)
75 x.Update(rscale, other.vec(), lscale)
76 return diag_op(x)
77 except AttributeError:
78 raise TypeError("can't extract matrix data from type '%s'"%str(type(other)))
80 @shared_vec_pool
81 def create_vec(self, dim=1):
82 from dolfin import EpetraVector
83 if dim > 1:
84 raise ValueError('dim must be <= 1')
85 return EpetraVector(self.v.Map())
87 def down_cast(self):
88 return self
89 def vec(self):
90 return self.v
92 def __str__(self):
93 return '<%s %dx%d>'%(self.__class__.__name__,self.v.GlobalLength(),self.v.GlobalLength())
95class matrix_op(block_base):
96 """Base class for Epetra operators (represented by an Epetra matrix)."""
97 from block.object_pool import vec_pool
99 def __init__(self, M, _transposed=False):
100 from PyTrilinos import Epetra
101 #M = Epetra.FECrsMatrix(M)
102 assert isinstance(M, (Epetra.CrsMatrix, Epetra.FECrsMatrix))
103 self.M = M
104 self.transposed = _transposed
106 def _transpose(self):
107 # Note: An object with transposed set should only be used internally
108 # (in matmat() / matvec() / ...), and not returned to the user.
109 return matrix_op(self.M, not self.transposed)
111 def rowmap(self):
112 # Note: Tried ColMap for the transposed matrix, but that
113 # crashes when the result is used by ML in parallel
114 return self.M.DomainMap() if self.transposed else self.M.RowMap()
116 def matvec(self, b):
117 from dolfin import GenericVector
118 if not isinstance(b, GenericVector):
119 return NotImplemented
120 if self.transposed:
121 domainlen = self.M.NumGlobalRows()
122 else:
123 domainlen = self.M.NumGlobalCols()
124 x = self.create_vec(dim=0)
125 if len(b) != domainlen:
126 raise RuntimeError(
127 'incompatible dimensions for %s matvec, %d != %d'%(self.__class__.__name__,domainlen,len(b)))
128 self.M.SetUseTranspose(self.transposed)
129 self.M.Apply(b.down_cast().vec(), x.down_cast().vec())
130 self.M.SetUseTranspose(False)
131 return x
133 def transpmult(self, b):
134 return self._transpose().matvec(b)
136 def matmat(self, other):
138 from PyTrilinos import Epetra
139 from block.block_util import isscalar
140 try:
141 if isscalar(other):
142 C = type(self.M)(self.M)
143 if other != 1:
144 C.Scale(other)
145 return matrix_op(C, self.transposed)
146 other = other.down_cast()
147 if hasattr(other, 'mat'):
148 from PyTrilinos import EpetraExt
150 # Create result matrix C. This is done in a contorted way, to
151 # ensure all diagonals are present.
152 A = type(self.M)(self.M); A.PutScalar(1.0)
153 B = type(other.mat())(other.mat()); B.PutScalar(1.0)
154 C = Epetra.FECrsMatrix(Epetra.Copy, self.rowmap(), 100)
155 EpetraExt.Multiply(A, self.transposed, B, other.transposed, C)
156 C.OptimizeStorage()
157 # C is now finalised, we can use it to store the real mat-mat product.
159 assert (0 == EpetraExt.Multiply(self.M, self.transposed, other.mat(), other.transposed, C))
161 result = matrix_op(C)
163 # Sanity check. Cannot trust EpetraExt.Multiply always, it seems.
164 from block import block_vec
165 v = result.create_vec()
166 block_vec([v]).randomize()
167 xv = self*other*v
168 err = (xv-result*v).norm('l2')/(xv).norm('l2')
169 if (err > 1e-3):
170 print(('++ a :',self*other))
171 print(('++ a\':',result))
172 print ('++ EpetraExt.Multiply computed wrong result; ||(a-a\')x||/||ax|| = %g'%err )
174 return result
175 else:
176 C = type(self.M)(self.M)
177 if self.transposed:
178 C.LeftScale(other.vec())
179 else:
180 C.RightScale(other.vec())
181 return matrix_op(C, self.transposed)
182 except AttributeError:
183 raise TypeError("can't extract matrix data from type '%s'"%str(type(other)))
185 def add(self, other, lscale=1.0, rscale=1.0):
186 try:
187 other = other.down_cast()
188 except AttributeError:
189 #raise TypeError("can't extract matrix data from type '%s'"%str(type(other)))
190 pass
192 if hasattr(other, 'mat'):
193 from PyTrilinos import Epetra, EpetraExt
194 C = Epetra.FECrsMatrix(Epetra.Copy, self.rowmap(), 100)
195 assert (0 == EpetraExt.Add(self.M, self.transposed, lscale, C, 0.0))
196 assert (0 == EpetraExt.Add(other.mat(), other.transposed, rscale, C, 1.0))
197 C.FillComplete()
198 C.OptimizeStorage()
199 return matrix_op(C)
200 else:
201 lhs = self.matmat(lscale)
202 D = Diag(lhs).add(other, rscale=rscale)
203 lhs.M.ReplaceDiagonalValues(D.vec())
204 return lhs
206 @vec_pool
207 def create_vec(self, dim=1):
208 from dolfin import EpetraVector
209 if self.transposed:
210 dim = 1-dim
211 if dim == 0:
212 m = self.M.RangeMap()
213 elif dim == 1:
214 m = self.M.DomainMap()
215 else:
216 raise ValueError('dim must be <= 1')
217 return EpetraVector(m)
219 def down_cast(self):
220 return self
221 def mat(self):
222 return self.M
224 def __str__(self):
225 format = '<%s transpose(%dx%d)>' if self.transposed else '<%s %dx%d>'
226 return format%(self.__class__.__name__, self.M.NumGlobalRows(), self.M.NumGlobalCols())
228class Diag(diag_op):
229 """Extract the diagonal entries of a matrix"""
230 def __init__(self, A):
231 from PyTrilinos import Epetra
232 A = A.down_cast().mat()
233 v = Epetra.Vector(A.RowMap())
234 A.ExtractDiagonalCopy(v)
235 diag_op.__init__(self, v)
237class InvDiag(Diag):
238 """Extract the inverse of the diagonal entries of a matrix"""
239 def __init__(self, A):
240 Diag.__init__(self, A)
241 self.v.Reciprocal(self.v)
243class LumpedInvDiag(diag_op):
244 """Extract the inverse of the lumped diagonal of a matrix (i.e., sum of the
245 def __init__(self, A):
246 absolute values in the row)."""
247 def __init__(self, A):
248 from PyTrilinos import Epetra
249 A = A.down_cast().mat()
250 v = Epetra.Vector(A.RowMap())
251 A.InvRowSums(v)
252 diag_op.__init__(self, v)
254def _collapse(x):
255 # Works by calling the matmat(), transpose() and add() methods of
256 # diag_op/matrix_op, depending on the input types. The input is a tree
257 # structure of block.block_mul objects, which is collapsed recursively.
259 # This method knows too much about the internal variables of the
260 # block_mul objects... should convert to accessor functions.
262 from block.block_compose import block_mul, block_add, block_sub, block_transpose
263 from block.block_mat import block_mat
264 from block.block_util import isscalar
265 from dolfin import Matrix
266 if isinstance(x, (matrix_op, diag_op)):
267 return x
268 elif isinstance(x, Matrix):
269 return matrix_op(x.down_cast().mat())
270 elif isinstance(x, block_mat):
271 if x.blocks.shape != (1,1):
272 raise NotImplementedError("collapse() for block_mat with shape != (1,1)")
273 return _collapse(x[0,0])
274 elif isinstance(x, block_mul):
275 factors = list(map(_collapse, x))
276 while len(factors) > 1:
277 A = factors.pop(0)
278 B = factors[0]
279 if isscalar(A):
280 C = A*B if isscalar(B) else B.matmat(A)
281 else:
282 C = A.matmat(B)
283 factors[0] = C
285 return factors[0]
286 elif isinstance(x, block_add):
287 A,B = list(map(_collapse, x))
288 if isscalar(A) and isscalar(B):
289 return A+B
290 else:
291 return B.add(A) if isscalar(A) else A.add(B)
292 elif isinstance(x, block_sub):
293 A,B = list(map(_collapse, x))
294 if isscalar(A) and isscalar(B):
295 return A-B
296 else:
297 return B.add(A, lscale=-1.0) if isscalar(A) else A.add(B, rscale=-1.0)
298 elif isinstance(x, block_transpose):
299 A = _collapse(x.A)
300 return A if isscalar(A) else A._transpose()
301 elif isscalar(x):
302 return x
303 else:
304 raise NotImplementedError("collapse() for type '%s'"%str(type(x)))
306def collapse(x):
307 """Compute an explicit matrix representation of an operator. For example,
308 given a block_ mul object M=A*B, collapse(M) performs the actual matrix
309 multiplication.
310 """
311 # Since _collapse works recursively, this method is a user-visible wrapper
312 # to print timing, and to check input/output arguments.
313 from time import time
314 from dolfin import EpetraMatrix, info, warning
315 T = time()
316 res = _collapse(x)
317 if getattr(res, 'transposed', False):
318 # transposed matrices will normally be converted to a non-transposed
319 # one by matrix multiplication or addition, but if the transpose is the
320 # outermost operation then this doesn't work.
321 warning('Transposed matrix returned from collapse() -- this matrix can be used for multiplications, '
322 + 'but not (for example) as input to ML. Try to convert from (A*B)^T to B^T*A^T in your call.')
323 from block.block_compose import block_transpose
324 return block_transpose(EpetraMatrix(res.M))
325 info('computed explicit matrix representation %s in %.2f s'%(str(res),time()-T))
327 result = EpetraMatrix(res.M)
329 # Sanity check. Cannot trust EpetraExt.Multiply always, it seems.
330 from block import block_vec
331 v = x.create_vec()
332 block_vec([v]).randomize()
333 xv = x*v
334 err = (xv-result*v).norm('l2')/(xv).norm('l2')
335 if (err > 1e-3):
336 raise RuntimeError('collapse computed wrong result; ||(a-a\')x||/||ax|| = %g'%err)
338 return result
341def create_identity(vec, val=1):
342 """Create an identity matrix with the layout given by the supplied
343 GenericVector. The values of the vector are NOT used."""
344 import numpy
345 from PyTrilinos import Epetra
346 from dolfin import EpetraMatrix
347 rowmap = vec.down_cast().vec().Map()
348 graph = Epetra.CrsGraph(Epetra.Copy, rowmap, 1, True)
349 indices = numpy.array([0], dtype=numpy.intc)
350 for row in rowmap.MyGlobalElements():
351 indices[0] = row
352 graph.InsertGlobalIndices(row, indices)
353 graph.FillComplete()
355 matrix = EpetraMatrix(graph)
356 indices = numpy.array(rowmap.MyGlobalElements())
357 if val == 0:
358 matrix.zero(indices)
359 else:
360 matrix.ident(indices)
361 if val != 1:
362 matrix *= val
364 return matrix