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

1from __future__ import division 

2 

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""" 

10 

11from block.block_base import block_base 

12from petsc4py import PETSc 

13 

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 

17 

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 

24 

25 def _transpose(self): 

26 return self 

27 

28 def matvec(self, b): 

29 try: 

30 b_vec = b.down_cast().vec() 

31 except AttributeError: 

32 return NotImplemented 

33 

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))) 

38 

39 x.down_cast().vec().pointwiseMult(self.v, b_vec) 

40 return x 

41 

42 transpmult = matvec 

43 

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))) 

62 

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))) 

79 

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()) 

86 

87 def down_cast(self): 

88 return self 

89 def vec(self): 

90 return self.v 

91 

92 def __str__(self): 

93 return '<%s %dx%d>'%(self.__class__.__name__,len(self.v),len(self.v)) 

94 

95class matrix_op(block_base): 

96 """Base class for PETSc operators (represented by an PETSc matrix).""" 

97 from block.object_pool import vec_pool 

98 

99 def __init__(self, M, _transposed=False): 

100 assert isinstance(M, PETSc.Mat) 

101 self.M = M 

102 self.transposed = _transposed 

103 

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) 

108 

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 

126 

127 def transpmult(self, b): 

128 return self._transpose().matvec(b) 

129 

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() 

141 

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()) 

148 

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))) 

159 

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 

166 

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() 

170 

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) 

183 

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 

190 

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) 

203 

204 def down_cast(self): 

205 return self 

206 def mat(self): 

207 return self.M 

208 

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]) 

212 

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()) 

217 

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() 

223 

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) 

237 

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. 

242 

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 

246 

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 

250 

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 

269 

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))) 

290 

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)) 

308 

309 result = PETScMatrix(res.M) 

310 

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) 

319 

320 return result 

321 

322 

323def create_identity(vec, val=1): 

324 raise NotImplementedError() 

325 

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() 

338 

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 

347 

348 return matrix