Coverage for /usr/share/miniconda3/envs/dolfin/lib/python3.8/site-packages/block/algebraic/petsc/matrix.py: 52%
256 statements
« prev ^ index » next coverage.py v7.2.1, created at 2023-03-20 13:03 +0000
« prev ^ index » next coverage.py v7.2.1, created at 2023-03-20 13:03 +0000
1from __future__ import division
3from builtins import map
4from builtins import str
5from builtins import range
6"""Module implementing algebraic operations on PETSc matrices: Diag, InvDiag
7etc, as well as the collapse() method which performs matrix addition and
8multiplication.
9"""
11from block.block_base import block_base
12from petsc4py import PETSc
14class diag_op(block_base):
15 """Base class for diagonal PETSc operators (represented by an PETSc vector)."""
16 from block.object_pool import shared_vec_pool
18 def __init__(self, v):
19 import dolfin
20 if isinstance(v, dolfin.GenericVector): 20 ↛ 21line 20 didn't jump to line 21, because the condition on line 20 was never true
21 v = v.down_cast().vec()
22 assert isinstance(v, PETSc.Vec)
23 self.v = v
25 def _transpose(self):
26 return self
28 def matvec(self, b):
29 try:
30 b_vec = b.down_cast().vec()
31 except AttributeError:
32 return NotImplemented
34 x = self.create_vec(dim=1)
35 if len(x) != len(b): 35 ↛ 36line 35 didn't jump to line 36, because the condition on line 35 was never true
36 raise RuntimeError(
37 'incompatible dimensions for %s matvec, %d != %d'%(self.__class__.__name__,len(x),len(b)))
39 x.down_cast().vec().pointwiseMult(self.v, b_vec)
40 return x
42 transpmult = matvec
44 def matmat(self, other):
45 try:
46 from block.block_util import isscalar
47 if isscalar(other):
48 x = self.v.copy()
49 x.scale(other)
50 return diag_op(x)
51 other = other.down_cast()
52 if hasattr(other, 'mat'):
53 C = other.mat().copy()
54 C.diagonalScale(L=self.v)
55 return matrix_op(C)
56 else:
57 x = self.v.copy()
58 x.pointwiseMult(self.v, other.vec())
59 return diag_op(x)
60 except AttributeError:
61 raise TypeError("can't extract matrix data from type '%s'"%str(type(other)))
63 def add(self, other, lscale=1.0, rscale=1.0):
64 from block.block_util import isscalar
65 try:
66 if isscalar(other):
67 x = self.v.copy()
68 x.shift(rscale*other)
69 return 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 = self.v.copy()
75 x.axpby(lscale, rscale, other.vec())
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 PETScVector
83 if dim > 1: 83 ↛ 84line 83 didn't jump to line 84, because the condition on line 83 was never true
84 raise ValueError('dim must be <= 1')
85 return PETScVector(self.v.copy())
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__,len(self.v),len(self.v))
95class matrix_op(block_base):
96 """Base class for PETSc operators (represented by an PETSc matrix)."""
97 from block.object_pool import vec_pool
99 def __init__(self, M, _transposed=False):
100 assert isinstance(M, PETSc.Mat)
101 self.M = M
102 self.transposed = _transposed
104 def _transpose(self):
105 # Note: An object with transposed set should only be used internally
106 # (in matmat() / matvec() / ...), and not returned to the user.
107 return matrix_op(self.M, not self.transposed)
109 def matvec(self, b):
110 from dolfin import GenericVector
111 if not isinstance(b, GenericVector):
112 return NotImplemented
113 if self.transposed:
114 domainlen = self.M.size[0]
115 else:
116 domainlen = self.M.size[1]
117 x = self.create_vec(dim=0)
118 if len(b) != domainlen:
119 raise RuntimeError(
120 'incompatible dimensions for %s matvec, %d != %d'%(self.__class__.__name__,domainlen,len(b)))
121 if self.transposed:
122 self.M.multTranspose(b.down_cast().vec(), x.down_cast().vec())
123 else:
124 self.M.mult(b.down_cast().vec(), x.down_cast().vec())
125 return x
127 def transpmult(self, b):
128 return self._transpose().matvec(b)
130 def matmat(self, other):
131 from block.block_util import isscalar
132 try:
133 if isscalar(other): 133 ↛ 134line 133 didn't jump to line 134, because the condition on line 133 was never true
134 C = self.M.copy()
135 C.scale(other)
136 return matrix_op(C, self.transposed)
137 other = other.down_cast()
138 if hasattr(other, 'mat'):
139 if self.transposed and other.transposed: 139 ↛ 140line 139 didn't jump to line 140, because the condition on line 139 was never true
140 return other.transpose().matmat(self.transpose()).transpose()
142 if self.transposed: 142 ↛ 143line 142 didn't jump to line 143, because the condition on line 142 was never true
143 C = self.M.transposeMatMult(other.mat())
144 elif other.transposed: 144 ↛ 145line 144 didn't jump to line 145, because the condition on line 144 was never true
145 C = self.M.matTransposeMult(other.mat())
146 else:
147 C = self.M.matMult(other.mat())
149 return matrix_op(C)
150 else:
151 C = self.M.copy()
152 if self.transposed: 152 ↛ 153line 152 didn't jump to line 153, because the condition on line 152 was never true
153 C.diagonalScale(L=other.vec())
154 else:
155 C.diagonalScale(R=other.vec())
156 return matrix_op(C, self.transposed)
157 except AttributeError:
158 raise TypeError("can't extract matrix data from type '%s'"%str(type(other)))
160 def add(self, other, lscale=1.0, rscale=1.0):
161 try:
162 other = other.down_cast()
163 except AttributeError:
164 #raise TypeError("can't extract matrix data from type '%s'"%str(type(other)))
165 pass
167 if hasattr(other, 'mat'): 167 ↛ 186line 167 didn't jump to line 186, because the condition on line 167 was never false
168 if self.transposed and other.transposed: 168 ↛ 169line 168 didn't jump to line 169, because the condition on line 168 was never true
169 return add(self, other, lscale=lscale, rscale=rscale).transpose()
171 A = self.M
172 if self.transposed: 172 ↛ 173line 172 didn't jump to line 173, because the condition on line 172 was never true
173 A = self.M.transpose()
174 if lscale != 1.0: 174 ↛ 175line 174 didn't jump to line 175, because the condition on line 174 was never true
175 A = A.copy()
176 A.scale(lscale)
177 B = other.mat()
178 if other.transposed: 178 ↛ 179line 178 didn't jump to line 179, because the condition on line 178 was never true
179 B = B.transpose()
180 if rscale != 1.0:
181 B = B.copy()
182 B.scale(rscale)
184 return matrix_op(A+B)
185 else:
186 lhs = self.matmat(lscale)
187 D = Diag(lhs).add(other, rscale=rscale)
188 lhs.M.setDiagonal(D.vec())
189 return lhs
191 @vec_pool
192 def create_vec(self, dim=1):
193 from dolfin import PETScVector
194 if self.transposed:
195 dim = 1-dim
196 if dim == 0:
197 m = self.M.createVecRight()
198 elif dim == 1:
199 m = self.M.createVecLeft()
200 else:
201 raise ValueError('dim must be <= 1')
202 return PETScVector(m)
204 def down_cast(self):
205 return self
206 def mat(self):
207 return self.M
209 def __str__(self):
210 format = '<%s transpose(%dx%d)>' if self.transposed else '<%s %dx%d>'
211 return format%(self.__class__.__name__, self.M.size[0], self.M.size[1])
213class Diag(diag_op):
214 """Extract the diagonal entries of a matrix"""
215 def __init__(self, A):
216 diag_op.__init__(self, A.down_cast().mat().getDiagonal())
218class InvDiag(Diag):
219 """Extract the inverse of the diagonal entries of a matrix"""
220 def __init__(self, A):
221 Diag.__init__(self, A)
222 self.v.reciprocal()
224class LumpedInvDiag(diag_op):
225 """Extract the inverse of the lumped diagonal of a matrix (i.e., sum of the
226 absolute values in the row)."""
227 def __init__(self, A):
228 v = A.create_vec().down_cast().vec()
229 a = A.down_cast().mat()
230 for row in range(*a.getOwnershipRange()):
231 data = a.getRow(row)[1]
232 v.setValue(row, sum(abs(d) for d in data))
233 v.assemblyBegin()
234 v.assemblyEnd()
235 v.reciprocal()
236 diag_op.__init__(self, v)
238def _collapse(x):
239 # Works by calling the matmat(), transpose() and add() methods of
240 # diag_op/matrix_op, depending on the input types. The input is a tree
241 # structure of block.block_mul objects, which is collapsed recursively.
243 # This method knows too much about the internal variables of the
244 # block_mul objects... should convert to accessor functions.
245 from dolfin import Matrix
247 from block.block_compose import block_mul, block_add, block_sub, block_transpose
248 from block.block_mat import block_mat
249 from block.block_util import isscalar
251 if isinstance(x, (matrix_op, diag_op)):
252 return x
253 elif isinstance(x, Matrix):
254 return matrix_op(x.down_cast().mat())
255 elif isinstance(x, block_mat): 255 ↛ 256line 255 didn't jump to line 256, because the condition on line 255 was never true
256 if x.blocks.shape != (1,1):
257 raise NotImplementedError("collapse() for block_mat with shape != (1,1)")
258 return _collapse(x[0,0])
259 elif isinstance(x, block_mul):
260 factors = list(map(_collapse, x))
261 while len(factors) > 1:
262 A = factors.pop(0)
263 B = factors[0]
264 if isscalar(A): 264 ↛ 265line 264 didn't jump to line 265, because the condition on line 264 was never true
265 C = A*B if isscalar(B) else B.matmat(A)
266 else:
267 C = A.matmat(B)
268 factors[0] = C
270 return factors[0]
271 elif isinstance(x, block_add):
272 A,B = list(map(_collapse, x))
273 if isscalar(A) and isscalar(B): 273 ↛ 274line 273 didn't jump to line 274, because the condition on line 273 was never true
274 return A+B
275 else:
276 return B.add(A) if isscalar(A) else A.add(B)
277 elif isinstance(x, block_sub): 277 ↛ 283line 277 didn't jump to line 283, because the condition on line 277 was never false
278 A,B = list(map(_collapse, x))
279 if isscalar(A) and isscalar(B): 279 ↛ 280line 279 didn't jump to line 280, because the condition on line 279 was never true
280 return A-B
281 else:
282 return B.add(A, lscale=-1.0) if isscalar(A) else A.add(B, rscale=-1.0)
283 elif isinstance(x, block_transpose):
284 A = _collapse(x.A)
285 return A if isscalar(A) else A._transpose()
286 elif isscalar(x):
287 return x
288 else:
289 raise NotImplementedError("collapse() for type '%s'"%str(type(x)))
291def collapse(x):
292 """Compute an explicit matrix representation of an operator. For example,
293 given a block_ mul object M=A*B, collapse(M) performs the actual matrix
294 multiplication.
295 """
296 # Since _collapse works recursively, this method is a user-visible wrapper
297 # to print timing, and to check input/output arguments.
298 from time import time
299 from dolfin import PETScMatrix, info, warning
300 T = time()
301 res = _collapse(x)
302 if getattr(res, 'transposed', False): 302 ↛ 306line 302 didn't jump to line 306, because the condition on line 302 was never true
303 # transposed matrices will normally be converted to a non-transposed
304 # one by matrix multiplication or addition, but if the transpose is the
305 # outermost operation then this doesn't work.
306 res.M.transpose()
307 info('computed explicit matrix representation %s in %.2f s'%(str(res),time()-T))
309 result = PETScMatrix(res.M)
311 # Sanity check.
312 from block import block_vec
313 v = x.create_vec()
314 block_vec([v]).randomize()
315 xv = x*v
316 err = (xv-result*v).norm('l2')/(xv).norm('l2')
317 if (err > 1e-3): 317 ↛ 318line 317 didn't jump to line 318, because the condition on line 317 was never true
318 raise RuntimeError('collapse computed wrong result; ||(a-a\')x||/||ax|| = %g'%err)
320 return result
323def create_identity(vec, val=1):
324 raise NotImplementedError()
326 """Create an identity matrix with the layout given by the supplied
327 GenericVector. The values of the vector are NOT used."""
328 import numpy
329 from PyTrilinos import Epetra
330 from dolfin import EpetraMatrix
331 rowmap = vec.down_cast().vec().Map()
332 graph = Epetra.CrsGraph(Epetra.Copy, rowmap, 1, True)
333 indices = numpy.array([0], dtype=numpy.intc)
334 for row in rowmap.MyGlobalElements():
335 indices[0] = row
336 graph.InsertGlobalIndices(row, indices)
337 graph.FillComplete()
339 matrix = EpetraMatrix(graph)
340 indices = numpy.array(rowmap.MyGlobalElements())
341 if val == 0:
342 matrix.zero(indices)
343 else:
344 matrix.ident(indices)
345 if val != 1:
346 matrix *= val
348 return matrix