Coverage for python/src/dolfinx

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

4# This file is part of DOLFINx MPC

6# SPDX-License-Identifier: MIT

7from __future__ import annotations

9from collections.abc import Iterable, Sequence

10from functools import partial

11from typing import cast

13from petsc4py import PETSc

15import dolfinx.fem.petsc

16import ufl

17from dolfinx import fem as _fem

18from dolfinx.la.petsc import _ghost_update, _zero_vector

20from .assemble_matrix import assemble_matrix, create_matrix

21from .assemble_vector import (

22 apply_lifting,

23 apply_mpc_lifting,

24 assemble_vector,

25 create_vector,

26)

27from .dirichletbc import BCData

28from .multipointconstraint import MultiPointConstraint

31def _backsubstitute(

32 mpc: MultiPointConstraint | Sequence[MultiPointConstraint],

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

34):

35 """Set the slave values of `u` from its masters, for one block or for every block.

37 A constraint with masters in another block reads them from the functions of all blocks, so

38 every block is homogenized before any is back-substituted.

39 """

40 if isinstance(mpc, MultiPointConstraint):

41 assert isinstance(u, _fem.Function)

42 mpc.homogenize(u)

43 mpc.backsubstitution(u)

44 return

45 assert isinstance(u, Sequence)

46 for mpc_i, u_i in zip(mpc, u):

47 mpc_i.homogenize(u_i)

48 for mpc_i, u_i in zip(mpc, u):

49 mpc_i.backsubstitution(u if mpc_i.has_cross_block_masters else u_i)

52def assemble_jacobian_mpc(

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

54 jacobian: _fem.Form | Sequence[Sequence[_fem.Form | None]],

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

56 bcs: Iterable[_fem.DirichletBC],

57 mpc: MultiPointConstraint | Sequence[MultiPointConstraint],

58 bc_data: BCData,

59 _snes: PETSc.SNES, # type: ignore

60 x: PETSc.Vec, # type: ignore

61 J: PETSc.Mat, # type: ignore

62 P: PETSc.Mat, # type: ignore

63 _blocks: tuple[tuple[int, ...], tuple[int, ...]] | None = None,

64):

65 """Assemble the Jacobian matrix and preconditioner.

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

68 be created by fixing the first four arguments:

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

71 bcs)

73 Args:

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

75 jacobian

76 jacobian: Form of the Jacobian

77 preconditioner: Form of the preconditioner

78 bcs: List of Dirichlet boundary conditions

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

80 bc_data: Dof marker and diagonal row cache, built once by the caller so

81 that every Newton iteration reuses it. The Jacobian and the

82 preconditioner share it: entries are keyed by function space, and

83 the two forms are over the same spaces.

84 _snes: The solver instance

85 x: The vector containing the point to evaluate at

86 J: Matrix to assemble the Jacobian into

87 P: Matrix to assemble the preconditioner into

88 _blocks: For a monolithic system of several blocks, the offsets of its blocks, see

89 :func:`dolfinx_mpc.create_vector`. SNES may pass a vector other than the one the

90 problem created, so they are set on `x` here.

91 """

92 if _blocks is not None:

93 x.setAttr("_blocks", _blocks)

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

95 # Jacobian

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

97 # Assign the input vector to the unknowns

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

99 _backsubstitute(mpc, u) # type: ignore

100

101 # Assemble Jacobian

102 J.zeroEntries()

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

104 J.assemble()

105 if preconditioner is not None:

106 P.zeroEntries()

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

108

109 P.assemble()

110

111

112def assemble_residual_mpc(

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

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

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

116 bcs: Sequence[_fem.DirichletBC],

117 mpc: MultiPointConstraint | Sequence[MultiPointConstraint],

118 bc_data: BCData,

119 _snes: PETSc.SNES, # type: ignore

120 x: PETSc.Vec, # type: ignore

121 F: PETSc.Vec, # type: ignore

122 _blocks: tuple[tuple[int, ...], tuple[int, ...]] | None = None,

123):

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

125

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

127 be created by fixing the first four arguments:

128

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

130 bcs)

131

132 Args:

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

134 Jacobian.

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

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

137 forms.

138 bcs: List of Dirichlet boundary conditions.

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

140 constraints.

141 bc_data: Dof marker cache shared with the Jacobian, built once by the

142 caller so that every Newton iteration reuses it.

143 _snes: The solver instance.

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

145 F: Vector to assemble the residual into.

146 _blocks: For a monolithic system of several blocks, the offsets of its blocks, see

147 :func:`dolfinx_mpc.create_vector`. A line search evaluates the residual in vectors

148 duplicated from the one given to SNES, so they are set on `x` and `F` here.

149 """

150 if _blocks is not None:

151 x.setAttr("_blocks", _blocks)

152 F.setAttr("_blocks", _blocks)

153 # Update input vector before assigning

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

155 # Assign the input vector to the unknowns

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

157 _backsubstitute(mpc, u) # type: ignore

158 # Assemble the residual

159 _zero_vector(F)

160 assemble_vector(residual, mpc, F) # type: ignore

161

162 # Lift vector. Decide between blocked (nest or monolithic) and single form lifting up front, so

163 # that a failure in one of the lifting calls cannot fall through to the other branch

164 if isinstance(jacobian, Sequence):

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

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

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

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

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

170 else:

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

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

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

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

175

176

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

178 def __init__(

179 self,

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

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

182 mpc: MultiPointConstraint | Sequence[MultiPointConstraint],

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

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

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

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

187 form_compiler_options: dict | None = None,

188 jit_options: dict | None = None,

189 petsc_options_prefix: str = "dolfinx_mpc_nonlinear_problem_",

190 petsc_options: dict | None = None,

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

192 ):

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

194

195 Solves problems of the form

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

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

198 non-linear solver.

199

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

201 NewtonSolver has been renamed NewtonSolverNonlinearProblem.

202

203 Args:

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

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

206 bcs: Dirichlet boundary conditions.

207 J: UFL form(s) representing the Jacobian

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

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

210 kind: The kind of Jacobian and preconditioner matrix, as in

211 :func:`dolfinx_mpc.create_matrix`. For a single constraint, a PETSc matrix type

212 (``MatType``). For a sequence of constraints, one per block, ``"nest"`` or a nested

213 sequence of matrix types gives a ``nest`` system, and ``None``, ``"mpi"`` or a

214 matrix type a single monolithic matrix and vector with the blocks one after another.

215 form_compiler_options: Options used in FFCx compilation of all

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

217 available options.

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

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

220 available options. Takes priority over all other option

221 values.

222 petsc_options_prefix: Options prefix used as the root prefix on

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

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

225 same on all ranks and is usually unique within the

226 programme.

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

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

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

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

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

232 integration domain mesh) a map should be provided relating

233 the entities in the integration domain mesh to the entities

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

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

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

237 mesh.

238 """

239 # Compile residual and Jacobian forms

240 self._F = _fem.form(

241 F,

242 form_compiler_options=form_compiler_options,

243 jit_options=jit_options,

244 entity_maps=entity_maps,

245 )

246

247 if J is None:

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

249 du = ufl.TrialFunction(F.arguments()[0].ufl_function_space())

250 J = ufl.derivative(F, u, du)

251 else:

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

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

254

255 self._J = _fem.form(

256 J,

257 form_compiler_options=form_compiler_options,

258 jit_options=jit_options,

259 entity_maps=entity_maps,

260 )

261

262 if P is not None:

263 self._preconditioner = _fem.form(

264 P,

265 form_compiler_options=form_compiler_options,

266 jit_options=jit_options,

267 entity_maps=entity_maps,

268 )

269 else:

270 self._preconditioner = None

271

272 self._u = u

273 # Set default values if not supplied

274 bcs = [] if bcs is None else bcs

275 self.mpc = mpc

276 # Create PETSc structures for the residual, Jacobian and solution vector

277 if not (kind is None or isinstance(kind, (str, Sequence))):

278 raise ValueError(f"Unsupported kind for matrix: {kind!r}")

279 if not isinstance(mpc, Sequence) and (kind == "nest" or (kind is not None and not isinstance(kind, str))):

280 raise ValueError(f"kind={kind!r} needs a sequence of constraints, one per block")

281 self._A = create_matrix(self._J, mpc, kind)

282

283 # The vectors are nested if the matrix is, and monolithic or single otherwise

284 vector_kind = "nest" if self._A.getType() == "nest" else None

285 self._b = create_vector(self._F, mpc, vector_kind)

286 self._x = create_vector(self._F, mpc, vector_kind)

287

288 # Create PETSc structure for preconditioner if provided

289 prec = self.preconditioner

290 if prec is not None: # type: ignore

291 self._P_mat = create_matrix(prec, mpc, kind)

292 else:

293 self._P_mat = None # type: ignore

294

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

296 # residual computation functions

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

298

299 # Markers and diagonal rows depend only on the function spaces and the

300 # conditions, so one cache covers the Jacobian and the preconditioner

301 # and every Newton iteration reuses it.

302 bc_data = BCData(bcs)

303

304 # The block layout of a monolithic system is an attribute of the vectors, which SNES may

305 # replace by duplicates in the callbacks, so it is handed to them

306 blocks = cast("tuple[tuple[int, ...], tuple[int, ...]] | None", self._b.getAttr("_blocks"))

307 self.solver.setJacobian(

308 partial(assemble_jacobian_mpc, u, self.J, prec, bcs, mpc, bc_data, _blocks=blocks),

309 self._A,

310 self.P_mat,

311 )

312 self.solver.setFunction(

313 partial(assemble_residual_mpc, u, self.F, self.J, bcs, mpc, bc_data, _blocks=blocks), self.b

314 )

315

316 # Set PETSc options

317 self._petsc_options_prefix = petsc_options_prefix

318 self.solver.setOptionsPrefix(petsc_options_prefix)

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

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

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

322 if self.P_mat is not None:

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

324

325 # Set options on SNES only

326 if petsc_options is not None:

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

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

329

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

331 opts.setValue(k, v)

332

333 self.solver.setFromOptions()

334

335 # Tidy up global options

336 for k in petsc_options.keys():

337 opts.delValue(k)

338 opts.prefixPop()

339

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

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

342 instance.

343

344 Returns:

345 The solution, convergence reason and number of iterations.

346 """

347

348 # Move current iterate into the work array.

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

350

351 # Solve problem

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

353

354 # Move solution back to function

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

356 _backsubstitute(self.mpc, self._u) # type: ignore

357

358 converged_reason = self.solver.getConvergedReason()

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

360

361

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

363 """

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

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

366

367 Args:

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

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

370 mpc: The multi point constraint.

371 bcs: A list of Dirichlet boundary conditions.

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

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

374

375 .. highlight:: python

376 .. code-block:: python

377

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

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

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

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

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

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

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

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

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

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

388 P: A preconditioner UFL form.

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

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

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

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

393 integration domain mesh) a map should be provided relating

394 the entities in the integration domain mesh to the entities

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

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

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

398 mesh.

399 kind: The kind of matrix, as in :func:`dolfinx.fem.petsc.LinearProblem`. For a single

400 constraint, a PETSc matrix type, with ``None`` the default. When ``mpc`` is a sequence,

401 one constraint per block, ``"nest"`` or a nested sequence of matrix types assembles a PETSc

402 ``nest`` matrix and vector, whose blocks have those types, and any other kind, such as

403 ``"mpi"``, a single monolithic matrix and vector with the blocks one after another.

404 ``None``, the default, is ``"nest"``, which is what a problem with one constraint per

405 block has always used.

406 Examples:

407 Example usage:

408

409 .. highlight:: python

410 .. code-block:: python

411

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

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

414

415 """

416

417 _u: _fem.Function | list[_fem.Function]

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

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

420 _jacobian: _fem.Form | Sequence[Sequence[_fem.Form | None]]

421 _preconditioner: _fem.Form | Sequence[Sequence[_fem.Form | None]] | None # type: ignore

422 _mpc: MultiPointConstraint | Sequence[MultiPointConstraint]

423 _bc_data: BCData

424 _A: PETSc.Mat

425 _P: PETSc.Mat | None

426 _b: PETSc.Vec

427 _solver: PETSc.KSP

428 _x: PETSc.Vec

429 bcs: list[_fem.DirichletBC]

430

431 def __init__(

432 self,

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

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

435 mpc: MultiPointConstraint | Sequence[MultiPointConstraint],

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

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

438 petsc_options_prefix: str = "dolfinx_mpc_linear_problem_",

439 petsc_options: dict | None = None,

440 form_compiler_options: dict | None = None,

441 jit_options: dict | None = None,

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

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

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

445 ):

446 # One constraint per block gives a nest or a monolithic system, one constraint a single matrix

447 if not (kind is None or isinstance(kind, (str, Sequence))):

448 raise ValueError(

449 f"Unsupported kind {kind!r}, expected None, a PETSc matrix type or a nested sequence of them"

450 )

451 if isinstance(mpc, Sequence):

452 # Without a kind a problem with one constraint per block keeps the nest layout it always had

453 kind = "nest" if kind is None else kind

454 elif kind == "nest" or (kind is not None and not isinstance(kind, str)):

455 raise ValueError(f"kind={kind!r} needs a sequence of constraints, one per block")

456 # Compile forms

457 form_compiler_options = {} if form_compiler_options is None else form_compiler_options

458 jit_options = {} if jit_options is None else jit_options

459 self._a = _fem.form(

460 a,

461 jit_options=jit_options,

462 form_compiler_options=form_compiler_options,

463 entity_maps=entity_maps,

464 )

465 self._L = _fem.form(

466 L,

467 jit_options=jit_options,

468 form_compiler_options=form_compiler_options,

469 entity_maps=entity_maps,

470 )

471

472 self._mpc = mpc

473 # Blocked problems

474 if isinstance(mpc, Sequence):

475 is_blocked = True

476 # Sanity check

477 for mpc_i in mpc:

478 if not mpc_i.finalized:

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

480 # Create function containing solution vector

481 else:

482 is_blocked = False

483 if not mpc.finalized:

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

485

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

487 if is_blocked:

488 if u is None:

489 assert isinstance(self._mpc, Sequence)

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

491 else:

492 assert isinstance(self._mpc, Sequence)

493 assert isinstance(u, Sequence)

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

495 assert isinstance(u_i, _fem.Function)

496 assert isinstance(mpc_i, MultiPointConstraint)

497 if u_i.function_space is not mpc_i.function_space:

498 raise ValueError(

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

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

501 )

502 self._u = list(u)

503 else:

504 if u is None:

505 assert isinstance(self._mpc, MultiPointConstraint)

506 self._u = _fem.Function(self._mpc.function_space)

507 else:

508 assert isinstance(u, _fem.Function)

509 assert isinstance(self._mpc, MultiPointConstraint)

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

511 self._u = u

512 else:

513 raise ValueError(

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

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

516 )

517

518 # Markers and diagonal rows depend only on the function spaces and the

519 # conditions, both fixed for this object, so one cache covers the

520 # operator and the preconditioner and every solve reuses it.

521 self._bc_data = BCData(bcs)

522

523 # Create MPC matrix and vector

524 self._preconditioner = _fem.form( # type: ignore

525 P,

526 jit_options=jit_options,

527 form_compiler_options=form_compiler_options,

528 entity_maps=entity_maps,

529 )

530

531 if is_blocked:

532 assert isinstance(mpc, Sequence)

533 assert isinstance(self._L, Sequence)

534 assert isinstance(self._a, Sequence)

535 self._A = create_matrix(self._a, mpc, kind)

536 self._b = create_vector(self._L, mpc, kind)

537 self._x = create_vector(self._L, mpc, kind)

538 if self._preconditioner is None:

539 self._P_mat = None

540 else:

541 assert isinstance(self._preconditioner, Sequence)

542 self._P_mat = create_matrix(self._preconditioner, mpc, kind)

543 else:

544 assert isinstance(mpc, MultiPointConstraint)

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

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

547 self._A = create_matrix(self._a, mpc, kind)

548 self._b = create_vector(self._L, mpc)

549 self._x = create_vector(self._L, mpc)

550 if self._preconditioner is None:

551 self._P_mat = None

552 else:

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

554 self._P_mat = create_matrix(self._preconditioner, mpc, kind)

555

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

557

558 if is_blocked:

559 assert isinstance(self.u, Sequence)

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

561 else:

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

563 comm = self.u.function_space.mesh.comm

564

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

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

567

568 self._petsc_options_prefix = petsc_options_prefix

569 self.solver.setOptionsPrefix(petsc_options_prefix)

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

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

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

573 if self.P_mat is not None:

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

575

576 # Set options on KSP only

577 if petsc_options is not None:

578 opts = PETSc.Options()

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

580

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

582 opts.setValue(k, v)

583

584 self.solver.setFromOptions()

585

586 # Tidy up global options

587 for k in petsc_options.keys():

588 opts.delValue(k)

589 opts.prefixPop()

590

591 def solve(self) -> _fem.Function | Sequence[_fem.Function]:

592 """Solve the problem.

593

594 Returns:

595 Function containing the solution"""

596

597 # Refresh the constraint offsets, so that a change in the values of the

598 # Dirichlet conditions held by the constraint is picked up

599 if isinstance(self._mpc, Sequence):

600 for mpc_i in self._mpc:

601 mpc_i.update_constants()

602 else:

603 self._mpc.update_constants()

604

605 # Assemble lhs. The layout, single, nest or monolithic, is that of the matrix

606 self._A.zeroEntries()

607 assemble_matrix(self._a, self._mpc, bcs=self.bcs, diagval=1.0, A=self._A, bc_data=self._bc_data) # type: ignore

608

609 self._A.assemble()

610 assert self._A.assembled

611

612 # Assemble the preconditioner if provided

613 if self._P_mat is not None:

614 self._P_mat.zeroEntries()

615 assemble_matrix(self._preconditioner, self._mpc, bcs=self.bcs, A=self._P_mat, bc_data=self._bc_data) # type: ignore

616 self._P_mat.assemble()

617

618 # Assemble the residual

619 _zero_vector(self._b)

620 assemble_vector(self._L, self._mpc, self._b) # type: ignore

621

622 # Lift vector

623 # Decide between nest/blocked and single form lifting up front, so that a

624 # failure in one of the lifting calls cannot fall through to the other

625 # branch and apply the lifting twice

626 try:

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

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

629 blocked = True

630 except ValueError:

631 blocked = False

632

633 if blocked:

634 # Nest and blocked lifting

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

636 apply_mpc_lifting(self._b, self._a, constraint=self._mpc) # type: ignore

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

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

639 else:

640 # Single form lifting

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

642 apply_mpc_lifting(self._b, [self._a], constraint=self._mpc) # type: ignore

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

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

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

646

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

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

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

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

651

652 _backsubstitute(self._mpc, self.u) # type: ignore

653

654 return self._u

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

249 statements