Coverage for python/src/dolfinx_mpc/problem.py: 76%

277 statements  

« prev     ^ index     » next       coverage.py v7.14.1, created at 2026-06-03 19:21 +0000

1# -*- coding: utf-8 -*- 

2# Copyright (C) 2021-2025 Jørgen S. Dokken 

3# 

4# This file is part of DOLFINx MPC 

5# 

6# SPDX-License-Identifier: MIT 

7from __future__ import annotations 

8 

9from collections.abc import Iterable, Sequence 

10from functools import partial 

11 

12from petsc4py import PETSc 

13 

14import dolfinx.fem.petsc 

15import ufl 

16from dolfinx import fem as _fem 

17from dolfinx.la.petsc import _ghost_update, _zero_vector, create_vector 

18 

19from dolfinx_mpc.cpp import mpc as _cpp_mpc 

20 

21from .assemble_matrix import assemble_matrix, assemble_matrix_nest, create_matrix_nest 

22from .assemble_vector import apply_lifting, assemble_vector, assemble_vector_nest, create_vector_nest 

23from .multipointconstraint import MultiPointConstraint 

24 

25 

26def assemble_jacobian_mpc( 

27 u: Sequence[_fem.Function] | _fem.Function, 

28 jacobian: _fem.Form | Sequence[Sequence[_fem.Form]], 

29 preconditioner: _fem.Form | Sequence[Sequence[_fem.Form]] | None, 

30 bcs: Iterable[_fem.DirichletBC], 

31 mpc: MultiPointConstraint | Sequence[MultiPointConstraint], 

32 _snes: PETSc.SNES, # type: ignore 

33 x: PETSc.Vec, # type: ignore 

34 J: PETSc.Mat, # type: ignore 

35 P: PETSc.Mat, # type: ignore 

36): 

37 """Assemble the Jacobian matrix and preconditioner. 

38 

39 A function conforming to the interface expected by SNES.setJacobian can 

40 be created by fixing the first four arguments: 

41 

42 functools.partial(assemble_jacobian, u, jacobian, preconditioner, 

43 bcs) 

44 

45 Args: 

46 u: Function tied to the solution vector within the residual and 

47 jacobian 

48 jacobian: Form of the Jacobian 

49 preconditioner: Form of the preconditioner 

50 bcs: List of Dirichlet boundary conditions 

51 mpc: The multi point constraint or a sequence of multi point 

52 _snes: The solver instance 

53 x: The vector containing the point to evaluate at 

54 J: Matrix to assemble the Jacobian into 

55 P: Matrix to assemble the preconditioner into 

56 """ 

57 # Copy existing soultion into the function used in the residual and 

58 # Jacobian 

59 _ghost_update(x, PETSc.InsertMode.INSERT, PETSc.ScatterMode.FORWARD) # type: ignore 

60 # Assign the input vector to the unknowns 

61 _fem.petsc.assign(x, u) # type: ignore 

62 if isinstance(u, Sequence): 

63 assert isinstance(mpc, Sequence) 

64 for i in range(len(u)): 

65 mpc[i].homogenize(u[i]) 

66 mpc[i].backsubstitution(u[i]) 

67 else: 

68 assert isinstance(u, _fem.Function) 

69 assert isinstance(mpc, MultiPointConstraint) 

70 mpc.homogenize(u) 

71 mpc.backsubstitution(u) 

72 

73 # Assemble Jacobian 

74 J.zeroEntries() 

75 if J.getType() == "nest": 

76 assemble_matrix_nest(J, jacobian, mpc, bcs, diagval=1.0) # type: ignore 

77 else: 

78 assemble_matrix(jacobian, mpc, bcs, diagval=1.0, A=J) # type: ignore 

79 J.assemble() 

80 if preconditioner is not None: 

81 P.zeroEntries() 

82 if P.getType() == "nest": 

83 assemble_matrix_nest(P, preconditioner, mpc, bcs, diagval=1.0) # type: ignore 

84 else: 

85 assemble_matrix(mpc, preconditioner, bcs, diagval=1.0, A=P) # type: ignore 

86 

87 P.assemble() 

88 

89 

90def assemble_residual_mpc( 

91 u: _fem.Function | Sequence[_fem.Function], 

92 residual: _fem.Form | Sequence[_fem.Form], 

93 jacobian: _fem.Form | Sequence[Sequence[_fem.Form]], 

94 bcs: Sequence[_fem.DirichletBC], 

95 mpc: MultiPointConstraint | Sequence[MultiPointConstraint], 

96 _snes: PETSc.SNES, # type: ignore 

97 x: PETSc.Vec, # type: ignore 

98 F: PETSc.Vec, # type: ignore 

99): 

100 """Assemble the residual into the vector `F`. 

101 

102 A function conforming to the interface expected by SNES.setResidual can 

103 be created by fixing the first four arguments: 

104 

105 functools.partial(assemble_residual, u, jacobian, preconditioner, 

106 bcs) 

107 

108 Args: 

109 u: Function(s) tied to the solution vector within the residual and 

110 Jacobian. 

111 residual: Form of the residual. It can be a sequence of forms. 

112 jacobian: Form of the Jacobian. It can be a nested sequence of 

113 forms. 

114 bcs: List of Dirichlet boundary conditions. 

115 _snes: The solver instance. 

116 x: The vector containing the point to evaluate the residual at. 

117 F: Vector to assemble the residual into. 

118 """ 

119 # Update input vector before assigning 

120 _ghost_update(x, PETSc.InsertMode.INSERT, PETSc.ScatterMode.FORWARD) # type: ignore 

121 # Assign the input vector to the unknowns 

122 _fem.petsc.assign(x, u) # type: ignore 

123 if isinstance(u, Sequence): 

124 assert isinstance(mpc, Sequence) 

125 for i in range(len(u)): 

126 mpc[i].homogenize(u[i]) 

127 mpc[i].backsubstitution(u[i]) 

128 else: 

129 assert isinstance(u, _fem.Function) 

130 assert isinstance(mpc, MultiPointConstraint) 

131 mpc.homogenize(u) 

132 mpc.backsubstitution(u) 

133 # Assemble the residual 

134 _zero_vector(F) 

135 if x.getType() == "nest": 

136 assemble_vector_nest(F, residual, mpc) # type: ignore 

137 else: 

138 assert isinstance(residual, _fem.Form) 

139 assert isinstance(mpc, MultiPointConstraint) 

140 assemble_vector(residual, mpc, F) 

141 

142 # Lift vector 

143 try: 

144 # Nest and blocked lifting 

145 bcs1 = _fem.bcs.bcs_by_block(_fem.forms.extract_function_spaces(jacobian, 1), bcs) # type: ignore 

146 _fem.petsc._assign_block_data(residual, x) # type: ignore 

147 apply_lifting(F, jacobian, bcs=bcs1, constraint=mpc, x0=x, scale=-1.0) # type: ignore 

148 _ghost_update(F, PETSc.InsertMode.ADD, PETSc.ScatterMode.REVERSE) # type: ignore 

149 bcs0 = _fem.bcs.bcs_by_block(_fem.forms.extract_function_spaces(residual), bcs) # type: ignore 

150 _fem.petsc.set_bc(F, bcs0, x0=x, alpha=-1.0) 

151 except TypeError: 

152 # Single form lifting 

153 apply_lifting(F, [jacobian], bcs=[bcs], constraint=mpc, x0=[x], scale=-1.0) # type: ignore 

154 _ghost_update(F, PETSc.InsertMode.ADD, PETSc.ScatterMode.REVERSE) # type: ignore 

155 _fem.petsc.set_bc(F, bcs, x0=x, alpha=-1.0) 

156 _ghost_update(F, PETSc.InsertMode.INSERT, PETSc.ScatterMode.FORWARD) # type: ignore 

157 

158 

159class NonlinearProblem(dolfinx.fem.petsc.NonlinearProblem): 

160 def __init__( 

161 self, 

162 F: ufl.form.Form | Sequence[ufl.form.Form], 

163 u: _fem.Function | Sequence[_fem.Function], 

164 mpc: MultiPointConstraint | Sequence[MultiPointConstraint], 

165 bcs: Sequence[_fem.DirichletBC] | None = None, 

166 J: ufl.form.Form | Sequence[Sequence[ufl.form.Form]] | None = None, 

167 P: ufl.form.Form | Sequence[Sequence[ufl.form.Form]] | None = None, 

168 kind: str | Sequence[Sequence[str]] | None = None, 

169 form_compiler_options: dict | None = None, 

170 jit_options: dict | None = None, 

171 petsc_options_prefix: str = "dolfinx_mpc_nonlinear_problem_", 

172 petsc_options: dict | None = None, 

173 entity_maps: Sequence[dolfinx.mesh.EntityMap] | None = None, 

174 ): 

175 """Class for solving nonlinear problems with SNES. 

176 

177 Solves problems of the form 

178 :math:`F_i(u, v) = 0, i=0,...N\\ \\forall v \\in V` where 

179 :math:`u=(u_0,...,u_N), v=(v_0,...,v_N)` using PETSc SNES as the 

180 non-linear solver. 

181 

182 Note: The deprecated version of this class for use with 

183 NewtonSolver has been renamed NewtonSolverNonlinearProblem. 

184 

185 Args: 

186 F: UFL form(s) of residual :math:`F_i`. 

187 u: Function used to define the residual and Jacobian. 

188 bcs: Dirichlet boundary conditions. 

189 J: UFL form(s) representing the Jacobian 

190 :math:`J_ij = dF_i/du_j`. 

191 P: UFL form(s) representing the preconditioner. 

192 kind: The PETSc matrix type(s) for the Jacobian and 

193 preconditioner (``MatType``). 

194 See :func:`dolfinx.fem.petsc.create_matrix` for more 

195 information. 

196 form_compiler_options: Options used in FFCx compilation of all 

197 forms. Run ``ffcx --help`` at the command line to see all 

198 available options. 

199 jit_options: Options used in CFFI JIT compilation of C code 

200 generated by FFCx. See ``python/dolfinx/jit.py`` for all 

201 available options. Takes priority over all other option 

202 values. 

203 petsc_options_prefix: Options prefix used as the root prefix on 

204 all internally created PETSc objects (SNES, A, b, x and the 

205 preconditioner matrix). Typically ends with ``_``. Must be the 

206 same on all ranks and is usually unique within the 

207 programme. 

208 petsc_options: Options to pass to the PETSc SNES object. 

209 entity_maps: If any trial functions, test functions, or 

210 coefficients in the form are not defined over the same mesh 

211 as the integration domain, ``entity_maps`` must be 

212 supplied. For each key (a mesh, different to the 

213 integration domain mesh) a map should be provided relating 

214 the entities in the integration domain mesh to the entities 

215 in the key mesh e.g. for a key-value pair ``(msh, emap)`` 

216 in ``entity_maps``, ``emap[i]`` is the entity in ``msh`` 

217 corresponding to entity ``i`` in the integration domain 

218 mesh. 

219 """ 

220 # Compile residual and Jacobian forms 

221 self._F = _fem.form( 

222 F, 

223 form_compiler_options=form_compiler_options, 

224 jit_options=jit_options, 

225 entity_maps=entity_maps, 

226 ) 

227 

228 if J is None: 

229 if isinstance(F, ufl.form.Form): 

230 du = ufl.TrialFunction(self._F.arguments()[0].ufl_function_space()) 

231 J = ufl.derivative(F, du) 

232 else: 

233 dus = [ufl.TrialFunction(Fi.arguments()[0].ufl_function_space()) for Fi in F] 

234 J = _fem.forms.derivative_block(F, u, dus) 

235 

236 self._J = _fem.form( 

237 J, 

238 form_compiler_options=form_compiler_options, 

239 jit_options=jit_options, 

240 entity_maps=entity_maps, 

241 ) 

242 

243 if P is not None: 

244 self._preconditioner = _fem.form( 

245 P, 

246 form_compiler_options=form_compiler_options, 

247 jit_options=jit_options, 

248 entity_maps=entity_maps, 

249 ) 

250 else: 

251 self._preconditioner = None 

252 

253 self._u = u 

254 # Set default values if not supplied 

255 bcs = [] if bcs is None else bcs 

256 self.mpc = mpc 

257 # Create PETSc structures for the residual, Jacobian and solution 

258 # vector 

259 if kind == "nest" or isinstance(kind, Sequence): 

260 assert isinstance(mpc, Sequence) 

261 assert isinstance(self._J, Sequence) 

262 self._A = create_matrix_nest(self._J, mpc) 

263 elif kind is None: 

264 assert isinstance(mpc, MultiPointConstraint) 

265 self._A = _cpp_mpc.create_matrix(self._J._cpp_object, mpc._cpp_object) 

266 else: 

267 raise ValueError("Unsupported kind for matrix: {}".format(kind)) 

268 

269 kind = "nest" if self._A.getType() == "nest" else kind 

270 if kind == "nest": 

271 assert isinstance(mpc, Sequence) 

272 assert isinstance(self._F, Sequence) 

273 self._b = create_vector_nest(self._F, mpc) 

274 self._x = create_vector_nest(self._F, mpc) 

275 else: 

276 assert isinstance(mpc, MultiPointConstraint) 

277 self._b = create_vector([(mpc.function_space.dofmap.index_map, mpc.function_space.dofmap.index_map_bs)]) 

278 self._x = create_vector([(mpc.function_space.dofmap.index_map, mpc.function_space.dofmap.index_map_bs)]) 

279 

280 # Create PETSc structure for preconditioner if provided 

281 if self.preconditioner is not None: # type: ignore 

282 if kind == "nest": 

283 assert isinstance(self.preconditioner, Sequence) 

284 assert isinstance(self.mpc, Sequence) 

285 self._P_mat = create_matrix_nest(self.preconditioner, self.mpc) 

286 else: 

287 assert isinstance(self._preconditioner, _fem.Form) 

288 self._P_mat = _cpp_mpc.create_matrix(self.preconditioner, kind=kind) 

289 else: 

290 self._P_mat = None # type: ignore 

291 

292 # Create the SNES solver and attach the corresponding Jacobian and 

293 # residual computation functions 

294 self._snes = PETSc.SNES().create(comm=self.A.comm) # type: ignore 

295 self.solver.setJacobian( 

296 partial(assemble_jacobian_mpc, u, self.J, self.preconditioner, bcs, mpc), self._A, self.P_mat 

297 ) 

298 self.solver.setFunction(partial(assemble_residual_mpc, u, self.F, self.J, bcs, mpc), self.b) 

299 

300 # Set PETSc options 

301 self._petsc_options_prefix = petsc_options_prefix 

302 self.solver.setOptionsPrefix(petsc_options_prefix) 

303 self.A.setOptionsPrefix(f"{petsc_options_prefix}A_") 

304 self.b.setOptionsPrefix(f"{petsc_options_prefix}b_") 

305 self.x.setOptionsPrefix(f"{petsc_options_prefix}x_") 

306 if self.P_mat is not None: 

307 self.P_mat.setOptionsPrefix(f"{petsc_options_prefix}P_mat_") 

308 

309 # Set options on SNES only 

310 if petsc_options is not None: 

311 opts = PETSc.Options() # type: ignore 

312 opts.prefixPush(self.solver.getOptionsPrefix()) 

313 

314 for k, v in petsc_options.items(): 

315 opts.setValue(k, v) 

316 

317 self.solver.setFromOptions() 

318 

319 # Tidy up global options 

320 for k in petsc_options.keys(): 

321 opts.delValue(k) 

322 opts.prefixPop() 

323 

324 def solve(self) -> tuple[_fem.Function | Sequence[_fem.Function], PETSc.ConvergedReason, int]: # type: ignore 

325 """Solve the problem and update the solution in the problem 

326 instance. 

327 

328 Returns: 

329 The solution, convergence reason and number of iterations. 

330 """ 

331 

332 # Move current iterate into the work array. 

333 _fem.petsc.assign(self._u, self.x) 

334 

335 # Solve problem 

336 self.solver.solve(None, self.x) 

337 

338 # Move solution back to function 

339 dolfinx.fem.petsc.assign(self.x, self._u) # type: ignore 

340 if isinstance(self.mpc, Sequence): 

341 for i in range(len(self._u)): 

342 self.mpc[i].backsubstitution(self._u[i]) 

343 self.mpc[i].backsubstitution(self._u[i]) 

344 else: 

345 assert isinstance(self._u, _fem.Function) 

346 self.mpc.homogenize(self._u) 

347 self.mpc.backsubstitution(self._u) 

348 

349 converged_reason = self.solver.getConvergedReason() 

350 return self._u, converged_reason, self.solver.getIterationNumber() # type: ignore 

351 

352 

353class LinearProblem(dolfinx.fem.petsc.LinearProblem): 

354 """ 

355 Class for solving a linear variational problem with multi point constraints of the form 

356 a(u, v) = L(v) for all v using PETSc as a linear algebra backend. 

357 

358 Args: 

359 a: A bilinear UFL form, the left hand side of the variational problem. 

360 L: A linear UFL form, the right hand side of the variational problem. 

361 mpc: The multi point constraint. 

362 bcs: A list of Dirichlet boundary conditions. 

363 u: The solution function. It will be created if not provided. The function has 

364 to be based on the functionspace in the mpc, i.e. 

365 

366 .. highlight:: python 

367 .. code-block:: python 

368 

369 u = dolfinx.fem.Function(mpc.function_space) 

370 petsc_options: Parameters that is passed to the linear algebra backend PETSc. #type: ignore 

371 For available choices for the 'petsc_options' kwarg, see the PETSc-documentation 

372 https://www.mcs.anl.gov/petsc/documentation/index.html. 

373 form_compiler_options: Parameters used in FFCx compilation of this form. Run `ffcx --help` at 

374 the commandline to see all available options. Takes priority over all 

375 other parameter values, except for `scalar_type` which is determined by DOLFINx. 

376 jit_options: Parameters used in CFFI JIT compilation of C code generated by FFCx. 

377 See https://github.com/FEniCS/dolfinx/blob/main/python/dolfinx/jit.py#L22-L37 

378 for all available parameters. Takes priority over all other parameter values. 

379 P: A preconditioner UFL form. 

380 Examples: 

381 Example usage: 

382 

383 .. highlight:: python 

384 .. code-block:: python 

385 

386 problem = LinearProblem(a, L, mpc, [bc0, bc1], 

387 petsc_options={"ksp_type": "preonly", "pc_type": "lu"}) 

388 

389 """ 

390 

391 u: _fem.Function | list[_fem.Function] 

392 _a: _fem.Form | Sequence[Sequence[_fem.Form]] 

393 _L: _fem.Form | Sequence[_fem.Form] 

394 _preconditioner: _fem.Form | Sequence[Sequence[_fem.Form]] | None 

395 _mpc: MultiPointConstraint | Sequence[MultiPointConstraint] 

396 _A: PETSc.Mat 

397 _P: PETSc.Mat | None 

398 _b: PETSc.Vec 

399 _solver: PETSc.KSP 

400 _x: PETSc.Vec 

401 bcs: list[_fem.DirichletBC] 

402 __slots__ = tuple(__annotations__) 

403 

404 def __init__( 

405 self, 

406 a: ufl.Form | Sequence[Sequence[ufl.Form]], 

407 L: ufl.Form | Sequence[ufl.Form], 

408 mpc: MultiPointConstraint | Sequence[MultiPointConstraint], 

409 bcs: list[_fem.DirichletBC] | None = None, 

410 u: _fem.Function | Sequence[_fem.Function] | None = None, 

411 petsc_options_prefix: str = "dolfinx_mpc_linear_problem_", 

412 petsc_options: dict | None = None, 

413 form_compiler_options: dict | None = None, 

414 jit_options: dict | None = None, 

415 P: ufl.Form | Sequence[Sequence[ufl.Form]] | None = None, 

416 ): 

417 # Compile forms 

418 form_compiler_options = {} if form_compiler_options is None else form_compiler_options 

419 jit_options = {} if jit_options is None else jit_options 

420 self._a = _fem.form(a, jit_options=jit_options, form_compiler_options=form_compiler_options) 

421 self._L = _fem.form(L, jit_options=jit_options, form_compiler_options=form_compiler_options) 

422 

423 self._mpc = mpc 

424 # Nest assembly 

425 if isinstance(mpc, Sequence): 

426 is_nest = True 

427 # Sanity check 

428 for mpc_i in mpc: 

429 if not mpc_i.finalized: 

430 raise RuntimeError("The multi point constraint has to be finalized before calling initializer") 

431 # Create function containing solution vector 

432 else: 

433 is_nest = False 

434 if not mpc.finalized: 

435 raise RuntimeError("The multi point constraint has to be finalized before calling initializer") 

436 

437 # Create function(s) containing solution vector(s) 

438 if is_nest: 

439 if u is None: 

440 assert isinstance(self._mpc, Sequence) 

441 self.u = [_fem.Function(self._mpc[i].function_space) for i in range(len(self._mpc))] 

442 else: 

443 assert isinstance(self._mpc, Sequence) 

444 assert isinstance(u, Sequence) 

445 for i, (mpc_i, u_i) in enumerate(zip(self._mpc, u)): 

446 assert isinstance(u_i, _fem.Function) 

447 assert isinstance(mpc_i, MultiPointConstraint) 

448 if u_i.function_space is not mpc_i.function_space: 

449 raise ValueError( 

450 "The input function has to be in the function space in the multi-point constraint", 

451 "i.e. u = dolfinx.fem.Function(mpc.function_space)", 

452 ) 

453 self.u = list(u) 

454 else: 

455 if u is None: 

456 assert isinstance(self._mpc, MultiPointConstraint) 

457 self.u = _fem.Function(self._mpc.function_space) 

458 else: 

459 assert isinstance(u, _fem.Function) 

460 assert isinstance(self._mpc, MultiPointConstraint) 

461 if u.function_space is self._mpc.function_space: 

462 self.u = u 

463 else: 

464 raise ValueError( 

465 "The input function has to be in the function space in the multi-point constraint", 

466 "i.e. u = dolfinx.fem.Function(mpc.function_space)", 

467 ) 

468 

469 # Create MPC matrix and vector 

470 self._preconditioner = _fem.form(P, jit_options=jit_options, form_compiler_options=form_compiler_options) # type: ignore 

471 

472 if is_nest: 

473 assert isinstance(mpc, Sequence) 

474 assert isinstance(self._L, Sequence) 

475 assert isinstance(self._a, Sequence) 

476 self._A = create_matrix_nest(self._a, mpc) 

477 self._b = create_vector_nest(self._L, mpc) 

478 self._x = create_vector_nest(self._L, mpc) 

479 if self._preconditioner is None: 

480 self._P_mat = None 

481 else: 

482 assert isinstance(self._preconditioner, Sequence) 

483 self._P_mat = create_matrix_nest(self._preconditioner, mpc) 

484 else: 

485 assert isinstance(mpc, MultiPointConstraint) 

486 assert isinstance(self._L, _fem.Form) 

487 assert isinstance(self._a, _fem.Form) 

488 self._A = _cpp_mpc.create_matrix(self._a._cpp_object, mpc._cpp_object) 

489 self._b = create_vector([(mpc.function_space.dofmap.index_map, mpc.function_space.dofmap.index_map_bs)]) 

490 self._x = create_vector([(mpc.function_space.dofmap.index_map, mpc.function_space.dofmap.index_map_bs)]) 

491 if self._preconditioner is None: 

492 self._P_mat = None 

493 else: 

494 assert isinstance(self._preconditioner, _fem.Form) 

495 self._P_mat = _cpp_mpc.create_matrix(self._preconditioner._cpp_object, mpc._cpp_object) 

496 

497 self.bcs = [] if bcs is None else bcs 

498 

499 if is_nest: 

500 assert isinstance(self.u, Sequence) 

501 comm = self.u[0].function_space.mesh.comm 

502 else: 

503 assert isinstance(self.u, _fem.Function) 

504 comm = self.u.function_space.mesh.comm 

505 

506 self._solver = PETSc.KSP().create(comm) 

507 self._solver.setOperators(self._A, self._P_mat) 

508 

509 self._petsc_options_prefix = petsc_options_prefix 

510 self.solver.setOptionsPrefix(petsc_options_prefix) 

511 self.A.setOptionsPrefix(f"{petsc_options_prefix}A_") 

512 self.b.setOptionsPrefix(f"{petsc_options_prefix}b_") 

513 self.x.setOptionsPrefix(f"{petsc_options_prefix}x_") 

514 if self.P_mat is not None: 

515 self.P_mat.setOptionsPrefix(f"{petsc_options_prefix}P_mat_") 

516 

517 # Set options on KSP only 

518 if petsc_options is not None: 

519 opts = PETSc.Options() 

520 opts.prefixPush(self.solver.getOptionsPrefix()) 

521 

522 for k, v in petsc_options.items(): 

523 opts.setValue(k, v) 

524 

525 self.solver.setFromOptions() 

526 

527 # Tidy up global options 

528 for k in petsc_options.keys(): 

529 opts.delValue(k) 

530 opts.prefixPop() 

531 

532 def solve(self) -> _fem.Function | list[_fem.Function]: 

533 """Solve the problem. 

534 

535 Returns: 

536 Function containing the solution""" 

537 

538 # Assemble lhs 

539 self._A.zeroEntries() 

540 if self._A.getType() == "nest": 

541 assemble_matrix_nest(self._A, self._a, self._mpc, self.bcs, diagval=1.0) # type: ignore 

542 else: 

543 assert isinstance(self._a, _fem.Form) 

544 assemble_matrix(self._a, self._mpc, bcs=self.bcs, A=self._A) 

545 

546 self._A.assemble() 

547 assert self._A.assembled 

548 

549 # Assemble the preconditioner if provided 

550 if self._P_mat is not None: 

551 self._P_mat.zeroEntries() 

552 if self._P_mat.getType() == "nest": 

553 assert isinstance(self._preconditioner, Sequence) 

554 assemble_matrix_nest(self._P_mat, self._preconditioner, self._mpc, self.bcs) # type: ignore 

555 else: 

556 assert isinstance(self._preconditioner, _fem.Form) 

557 assemble_matrix(self._preconditioner, self._mpc, bcs=self.bcs, A=self._P_mat) 

558 self._P_mat.assemble() 

559 

560 # Assemble the residual 

561 _zero_vector(self._b) 

562 if self._x.getType() == "nest": 

563 assemble_vector_nest(self._b, self._L, self._mpc) # type: ignore 

564 else: 

565 assert isinstance(self._L, _fem.Form) 

566 assert isinstance(self._mpc, MultiPointConstraint) 

567 assemble_vector(self._L, self._mpc, self._b) 

568 

569 # Lift vector 

570 try: 

571 # Nest and blocked lifting 

572 bcs1 = _fem.bcs.bcs_by_block(_fem.forms.extract_function_spaces(self._a, 1), self.bcs) # type: ignore 

573 apply_lifting(self._b, self._a, bcs=bcs1, constraint=self._mpc) # type: ignore 

574 _ghost_update(self._b, PETSc.InsertMode.ADD, PETSc.ScatterMode.REVERSE) # type: ignore 

575 bcs0 = _fem.bcs.bcs_by_block(_fem.forms.extract_function_spaces(self._L), self.bcs) # type: ignore 

576 _fem.petsc.set_bc(self._b, bcs0) 

577 except TypeError: 

578 # Single form lifting 

579 apply_lifting(self._b, [self._a], bcs=[self.bcs], constraint=self._mpc) # type: ignore 

580 _ghost_update(self._b, PETSc.InsertMode.ADD, PETSc.ScatterMode.REVERSE) # type: ignore 

581 _fem.petsc.set_bc(self._b, self.bcs) 

582 _ghost_update(self._b, PETSc.InsertMode.INSERT, PETSc.ScatterMode.FORWARD) # type: ignore 

583 

584 # Solve linear system and update ghost values in the solution 

585 self._solver.solve(self._b, self._x) 

586 _ghost_update(self._x, PETSc.InsertMode.INSERT, PETSc.ScatterMode.FORWARD) # type: ignore 

587 _fem.petsc.assign(self._x, self.u) # type: ignore 

588 

589 if isinstance(self.u, Sequence): 

590 assert isinstance(self._mpc, Sequence) 

591 for i in range(len(self.u)): 

592 self._mpc[i].homogenize(self.u[i]) 

593 self._mpc[i].backsubstitution(self.u[i]) 

594 else: 

595 assert isinstance(self.u, _fem.Function) 

596 assert isinstance(self._mpc, MultiPointConstraint) 

597 self._mpc.homogenize(self.u) 

598 self._mpc.backsubstitution(self.u) 

599 

600 return self.u