Coverage for python/src/dolfinx_mpc/problem.py: 97%
252 statements
« prev ^ index » next coverage.py v7.17.0, created at 2026-11-03 12:11 +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
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 _row_spaces(L, mpc) -> list:
32 """The space of each block row of a linear form, or of its constraint where the block has none.
34 A block without a linear form still has rows, on which its Dirichlet conditions are set.
35 """
36 spaces = _fem.forms.extract_function_spaces(L)
37 return [c.input_space if V is None else V for V, c in zip(spaces, mpc)]
40def _backsubstitute(
41 mpc: MultiPointConstraint | Sequence[MultiPointConstraint],
42 u: _fem.Function | Sequence[_fem.Function],
43):
44 """Set the slave values of `u` from its masters, for one block or for every block.
46 A constraint with masters in another block reads them from the functions of all blocks, so
47 every block is homogenized before any is back-substituted.
48 """
49 if isinstance(mpc, MultiPointConstraint):
50 assert isinstance(u, _fem.Function)
51 mpc.homogenize(u)
52 mpc.backsubstitution(u)
53 return
54 assert isinstance(u, Sequence)
55 for mpc_i, u_i in zip(mpc, u):
56 mpc_i.homogenize(u_i)
57 for mpc_i, u_i in zip(mpc, u):
58 mpc_i.backsubstitution(u if mpc_i.has_cross_block_masters else u_i)
61def assemble_jacobian_mpc(
62 u: Sequence[_fem.Function] | _fem.Function,
63 jacobian: _fem.Form | Sequence[Sequence[_fem.Form | None]],
64 preconditioner: _fem.Form | Sequence[Sequence[_fem.Form | None]] | None,
65 bcs: Iterable[_fem.DirichletBC],
66 mpc: MultiPointConstraint | Sequence[MultiPointConstraint],
67 bc_data: BCData,
68 _snes: PETSc.SNES, # type: ignore
69 x: PETSc.Vec, # type: ignore
70 J: PETSc.Mat, # type: ignore
71 P: PETSc.Mat, # type: ignore
72 _blocks: tuple[tuple[int, ...], tuple[int, ...]] | None = None,
73):
74 """Assemble the Jacobian matrix and preconditioner.
76 A function conforming to the interface expected by SNES.setJacobian can
77 be created by fixing the first four arguments:
79 functools.partial(assemble_jacobian, u, jacobian, preconditioner,
80 bcs)
82 Args:
83 u: Function tied to the solution vector within the residual and
84 jacobian
85 jacobian: Form of the Jacobian
86 preconditioner: Form of the preconditioner
87 bcs: List of Dirichlet boundary conditions
88 mpc: The multi point constraint or a sequence of multi point
89 bc_data: Dof marker and diagonal row cache, built once by the caller so
90 that every Newton iteration reuses it. The Jacobian and the
91 preconditioner share it: entries are keyed by function space, and
92 the two forms are over the same spaces.
93 _snes: The solver instance
94 x: The vector containing the point to evaluate at
95 J: Matrix to assemble the Jacobian into
96 P: Matrix to assemble the preconditioner into
97 _blocks: For a monolithic system of several blocks, the offsets of its blocks, see
98 :func:`dolfinx_mpc.create_vector`. SNES may pass a vector other than the one the
99 problem created, so they are set on `x` here.
100 """
101 if _blocks is not None:
102 x.setAttr("_blocks", _blocks)
103 # Copy existing soultion into the function used in the residual and
104 # Jacobian
105 _ghost_update(x, PETSc.InsertMode.INSERT, PETSc.ScatterMode.FORWARD) # type: ignore
106 # Assign the input vector to the unknowns
107 _fem.petsc.assign(x, u) # type: ignore
108 _backsubstitute(mpc, u) # type: ignore
110 # Assemble Jacobian
111 J.zeroEntries()
112 assemble_matrix(jacobian, mpc, bcs, diagval=1.0, A=J, bc_data=bc_data) # type: ignore
113 J.assemble()
114 if preconditioner is not None:
115 P.zeroEntries()
116 assemble_matrix(preconditioner, mpc, bcs, diagval=1.0, A=P, bc_data=bc_data) # type: ignore
118 P.assemble()
121def assemble_residual_mpc(
122 u: _fem.Function | Sequence[_fem.Function],
123 residual: _fem.Form | Sequence[_fem.Form],
124 jacobian: _fem.Form | Sequence[Sequence[_fem.Form]],
125 bcs: Sequence[_fem.DirichletBC],
126 mpc: MultiPointConstraint | Sequence[MultiPointConstraint],
127 bc_data: BCData,
128 _snes: PETSc.SNES, # type: ignore
129 x: PETSc.Vec, # type: ignore
130 F: PETSc.Vec, # type: ignore
131 _blocks: tuple[tuple[int, ...], tuple[int, ...]] | None = None,
132):
133 """Assemble the residual into the vector `F`.
135 A function conforming to the interface expected by SNES.setResidual can
136 be created by fixing the first four arguments:
138 functools.partial(assemble_residual, u, jacobian, preconditioner,
139 bcs)
141 Args:
142 u: Function(s) tied to the solution vector within the residual and
143 Jacobian.
144 residual: Form of the residual. It can be a sequence of forms.
145 jacobian: Form of the Jacobian. It can be a nested sequence of
146 forms.
147 bcs: List of Dirichlet boundary conditions.
148 mpc: The multi point constraint or a sequence of multi point
149 constraints.
150 bc_data: Dof marker cache shared with the Jacobian, built once by the
151 caller so that every Newton iteration reuses it.
152 _snes: The solver instance.
153 x: The vector containing the point to evaluate the residual at.
154 F: Vector to assemble the residual into.
155 _blocks: For a monolithic system of several blocks, the offsets of its blocks, see
156 :func:`dolfinx_mpc.create_vector`. A line search evaluates the residual in vectors
157 duplicated from the one given to SNES, so they are set on `x` and `F` here.
158 """
159 if _blocks is not None:
160 x.setAttr("_blocks", _blocks)
161 F.setAttr("_blocks", _blocks)
162 # Update input vector before assigning
163 _ghost_update(x, PETSc.InsertMode.INSERT, PETSc.ScatterMode.FORWARD) # type: ignore
164 # Assign the input vector to the unknowns
165 _fem.petsc.assign(x, u) # type: ignore
166 _backsubstitute(mpc, u) # type: ignore
167 # Assemble the residual
168 _zero_vector(F)
169 assemble_vector(residual, mpc, F) # type: ignore
171 # Lift vector. Decide between blocked (nest or monolithic) and single form lifting up front, so
172 # that a failure in one of the lifting calls cannot fall through to the other branch
173 if isinstance(jacobian, Sequence):
174 bcs1 = _fem.bcs.bcs_by_block(_fem.forms.extract_function_spaces(jacobian, 1), bcs) # type: ignore
175 apply_lifting(F, jacobian, bcs=bcs1, constraint=mpc, x0=x, scale=-1.0, bc_data=bc_data) # type: ignore
176 _ghost_update(F, PETSc.InsertMode.ADD, PETSc.ScatterMode.REVERSE) # type: ignore
177 bcs0 = _fem.bcs.bcs_by_block(_row_spaces(residual, mpc), bcs) # type: ignore
178 _fem.petsc.set_bc(F, bcs0, x0=x, alpha=-1.0)
179 else:
180 apply_lifting(F, [jacobian], bcs=[bcs], constraint=mpc, x0=[x], scale=-1.0, bc_data=bc_data) # type: ignore
181 _ghost_update(F, PETSc.InsertMode.ADD, PETSc.ScatterMode.REVERSE) # type: ignore
182 _fem.petsc.set_bc(F, bcs, x0=x, alpha=-1.0)
183 _ghost_update(F, PETSc.InsertMode.INSERT, PETSc.ScatterMode.FORWARD) # type: ignore
186class NonlinearProblem(dolfinx.fem.petsc.NonlinearProblem):
187 def __init__(
188 self,
189 F: ufl.form.Form | Sequence[ufl.form.Form],
190 u: _fem.Function | Sequence[_fem.Function],
191 mpc: MultiPointConstraint | Sequence[MultiPointConstraint],
192 bcs: Sequence[_fem.DirichletBC] | None = None,
193 J: ufl.form.Form | Sequence[Sequence[ufl.form.Form]] | None = None,
194 P: ufl.form.Form | Sequence[Sequence[ufl.form.Form]] | None = None,
195 kind: str | Sequence[Sequence[str]] | None = None,
196 form_compiler_options: dict | None = None,
197 jit_options: dict | None = None,
198 petsc_options_prefix: str = "dolfinx_mpc_nonlinear_problem_",
199 petsc_options: dict | None = None,
200 entity_maps: Sequence[dolfinx.mesh.EntityMap] | None = None,
201 ):
202 """Class for solving nonlinear problems with SNES.
204 Solves problems of the form
205 :math:`F_i(u, v) = 0, i=0,...N\\ \\forall v \\in V` where
206 :math:`u=(u_0,...,u_N), v=(v_0,...,v_N)` using PETSc SNES as the
207 non-linear solver.
209 Note: The deprecated version of this class for use with
210 NewtonSolver has been renamed NewtonSolverNonlinearProblem.
212 Args:
213 F: UFL form(s) of residual :math:`F_i`.
214 u: Function used to define the residual and Jacobian.
215 bcs: Dirichlet boundary conditions.
216 J: UFL form(s) representing the Jacobian
217 :math:`J_ij = dF_i/du_j`.
218 P: UFL form(s) representing the preconditioner.
219 kind: The kind of Jacobian and preconditioner matrix, as in
220 :func:`dolfinx_mpc.create_matrix`. For a single constraint, a PETSc matrix type
221 (``MatType``). For a sequence of constraints, one per block, ``"nest"`` or a nested
222 sequence of matrix types gives a ``nest`` system, and ``None``, ``"mpi"`` or a
223 matrix type a single monolithic matrix and vector with the blocks one after another.
224 form_compiler_options: Options used in FFCx compilation of all
225 forms. Run ``ffcx --help`` at the command line to see all
226 available options.
227 jit_options: Options used in CFFI JIT compilation of C code
228 generated by FFCx. See ``python/dolfinx/jit.py`` for all
229 available options. Takes priority over all other option
230 values.
231 petsc_options_prefix: Options prefix used as the root prefix on
232 all internally created PETSc objects (SNES, A, b, x and the
233 preconditioner matrix). Typically ends with ``_``. Must be the
234 same on all ranks and is usually unique within the
235 programme.
236 petsc_options: Options to pass to the PETSc SNES object.
237 entity_maps: If any trial functions, test functions, or
238 coefficients in the form are not defined over the same mesh
239 as the integration domain, ``entity_maps`` must be
240 supplied. For each key (a mesh, different to the
241 integration domain mesh) a map should be provided relating
242 the entities in the integration domain mesh to the entities
243 in the key mesh e.g. for a key-value pair ``(msh, emap)``
244 in ``entity_maps``, ``emap[i]`` is the entity in ``msh``
245 corresponding to entity ``i`` in the integration domain
246 mesh.
247 """
248 # Compile residual and Jacobian forms
249 self._F = _fem.form(
250 F,
251 form_compiler_options=form_compiler_options,
252 jit_options=jit_options,
253 entity_maps=entity_maps,
254 )
256 if J is None:
257 if isinstance(F, ufl.form.Form):
258 du = ufl.TrialFunction(F.arguments()[0].ufl_function_space())
259 J = ufl.derivative(F, u, du)
260 else:
261 dus = [ufl.TrialFunction(Fi.arguments()[0].ufl_function_space()) for Fi in F]
262 J = _fem.forms.derivative_block(F, u, dus)
264 self._J = _fem.form(
265 J,
266 form_compiler_options=form_compiler_options,
267 jit_options=jit_options,
268 entity_maps=entity_maps,
269 )
271 if P is not None:
272 self._preconditioner = _fem.form(
273 P,
274 form_compiler_options=form_compiler_options,
275 jit_options=jit_options,
276 entity_maps=entity_maps,
277 )
278 else:
279 self._preconditioner = None
281 self._u = u
282 # Set default values if not supplied
283 bcs = [] if bcs is None else bcs
284 self.mpc = mpc
285 # Create PETSc structures for the residual, Jacobian and solution vector
286 if not (kind is None or isinstance(kind, (str, Sequence))):
287 raise ValueError(f"Unsupported kind for matrix: {kind!r}")
288 if not isinstance(mpc, Sequence) and (kind == "nest" or (kind is not None and not isinstance(kind, str))):
289 raise ValueError(f"kind={kind!r} needs a sequence of constraints, one per block")
290 self._A = create_matrix(self._J, mpc, kind)
292 # The vectors are nested if the matrix is, and monolithic or single otherwise
293 vector_kind = "nest" if self._A.getType() == "nest" else None
294 self._b = create_vector(self._F, mpc, vector_kind)
295 self._x = create_vector(self._F, mpc, vector_kind)
297 # Create PETSc structure for preconditioner if provided
298 prec = self.preconditioner
299 if prec is not None: # type: ignore
300 self._P_mat = create_matrix(prec, mpc, kind)
301 else:
302 self._P_mat = None # type: ignore
304 # Create the SNES solver and attach the corresponding Jacobian and
305 # residual computation functions
306 self._snes = PETSc.SNES().create(comm=self.A.comm) # type: ignore
308 # Markers and diagonal rows depend only on the function spaces and the
309 # conditions, so one cache covers the Jacobian and the preconditioner
310 # and every Newton iteration reuses it.
311 bc_data = BCData(bcs)
313 # The block layout of a monolithic system is an attribute of the vectors, which SNES may
314 # replace by duplicates in the callbacks, so it is handed to them
315 blocks = cast("tuple[tuple[int, ...], tuple[int, ...]] | None", self._b.getAttr("_blocks"))
316 self.solver.setJacobian(
317 partial(assemble_jacobian_mpc, u, self.J, prec, bcs, mpc, bc_data, _blocks=blocks),
318 self._A,
319 self.P_mat,
320 )
321 self.solver.setFunction(
322 partial(assemble_residual_mpc, u, self.F, self.J, bcs, mpc, bc_data, _blocks=blocks), self.b
323 )
325 # Set PETSc options
326 self._petsc_options_prefix = petsc_options_prefix
327 self.solver.setOptionsPrefix(petsc_options_prefix)
328 self.A.setOptionsPrefix(f"{petsc_options_prefix}A_")
329 self.b.setOptionsPrefix(f"{petsc_options_prefix}b_")
330 self.x.setOptionsPrefix(f"{petsc_options_prefix}x_")
331 if self.P_mat is not None:
332 self.P_mat.setOptionsPrefix(f"{petsc_options_prefix}P_mat_")
334 # Set options on SNES only
335 if petsc_options is not None:
336 opts = PETSc.Options() # type: ignore
337 opts.prefixPush(self.solver.getOptionsPrefix())
339 for k, v in petsc_options.items():
340 opts.setValue(k, v)
342 self.solver.setFromOptions()
344 # Tidy up global options
345 for k in petsc_options.keys():
346 opts.delValue(k)
347 opts.prefixPop()
349 def solve(self) -> tuple[_fem.Function | Sequence[_fem.Function], PETSc.ConvergedReason, int]: # type: ignore
350 """Solve the problem and update the solution in the problem
351 instance.
353 Returns:
354 The solution, convergence reason and number of iterations.
355 """
357 # Move current iterate into the work array.
358 _fem.petsc.assign(self._u, self.x)
360 # Solve problem
361 self.solver.solve(None, self.x)
363 # Move solution back to function
364 dolfinx.fem.petsc.assign(self.x, self._u) # type: ignore
365 _backsubstitute(self.mpc, self._u) # type: ignore
367 converged_reason = self.solver.getConvergedReason()
368 return self._u, converged_reason, self.solver.getIterationNumber() # type: ignore
371class LinearProblem(dolfinx.fem.petsc.LinearProblem):
372 """
373 Class for solving a linear variational problem with multi point constraints of the form
374 a(u, v) = L(v) for all v using PETSc as a linear algebra backend.
376 Args:
377 a: A bilinear UFL form, the left hand side of the variational problem.
378 L: A linear UFL form, the right hand side of the variational problem.
379 mpc: The multi point constraint.
380 bcs: A list of Dirichlet boundary conditions.
381 u: The solution function. It will be created if not provided. The function has
382 to be based on the functionspace in the mpc, i.e.
384 .. highlight:: python
385 .. code-block:: python
387 u = dolfinx.fem.Function(mpc.function_space)
388 petsc_options: Parameters that is passed to the linear algebra backend PETSc. #type: ignore
389 For available choices for the 'petsc_options' kwarg, see the PETSc-documentation
390 https://www.mcs.anl.gov/petsc/documentation/index.html.
391 form_compiler_options: Parameters used in FFCx compilation of this form. Run `ffcx --help` at
392 the commandline to see all available options. Takes priority over all
393 other parameter values, except for `scalar_type` which is determined by DOLFINx.
394 jit_options: Parameters used in CFFI JIT compilation of C code generated by FFCx.
395 See https://github.com/FEniCS/dolfinx/blob/main/python/dolfinx/jit.py#L22-L37
396 for all available parameters. Takes priority over all other parameter values.
397 P: A preconditioner UFL form.
398 entity_maps: If any trial functions, test functions, or
399 coefficients in the form are not defined over the same mesh
400 as the integration domain, ``entity_maps`` must be
401 supplied. For each key (a mesh, different to the
402 integration domain mesh) a map should be provided relating
403 the entities in the integration domain mesh to the entities
404 in the key mesh e.g. for a key-value pair ``(msh, emap)``
405 in ``entity_maps``, ``emap[i]`` is the entity in ``msh``
406 corresponding to entity ``i`` in the integration domain
407 mesh.
408 kind: The kind of matrix, as in :func:`dolfinx.fem.petsc.LinearProblem`. For a single
409 constraint, a PETSc matrix type, with ``None`` the default. When ``mpc`` is a sequence,
410 one constraint per block, ``"nest"`` or a nested sequence of matrix types assembles a PETSc
411 ``nest`` matrix and vector, whose blocks have those types, and any other kind, such as
412 ``"mpi"``, a single monolithic matrix and vector with the blocks one after another.
413 ``None``, the default, is ``"nest"``, which is what a problem with one constraint per
414 block has always used.
415 Examples:
416 Example usage:
418 .. highlight:: python
419 .. code-block:: python
421 problem = LinearProblem(a, L, mpc, [bc0, bc1],
422 petsc_options={"ksp_type": "preonly", "pc_type": "lu"})
424 """
426 _u: _fem.Function | list[_fem.Function]
427 _a: _fem.Form | Sequence[Sequence[_fem.Form]]
428 _L: _fem.Form | Sequence[_fem.Form]
429 _jacobian: _fem.Form | Sequence[Sequence[_fem.Form | None]]
430 _preconditioner: _fem.Form | Sequence[Sequence[_fem.Form | None]] | None # type: ignore
431 _mpc: MultiPointConstraint | Sequence[MultiPointConstraint]
432 _bc_data: BCData
433 _A: PETSc.Mat
434 _P: PETSc.Mat | None
435 _b: PETSc.Vec
436 _solver: PETSc.KSP
437 _x: PETSc.Vec
438 bcs: list[_fem.DirichletBC]
440 def __init__(
441 self,
442 a: ufl.Form | Sequence[Sequence[ufl.Form]],
443 L: ufl.Form | Sequence[ufl.Form],
444 mpc: MultiPointConstraint | Sequence[MultiPointConstraint],
445 bcs: list[_fem.DirichletBC] | None = None,
446 u: _fem.Function | Sequence[_fem.Function] | None = None,
447 petsc_options_prefix: str = "dolfinx_mpc_linear_problem_",
448 petsc_options: dict | None = None,
449 form_compiler_options: dict | None = None,
450 jit_options: dict | None = None,
451 P: ufl.Form | Sequence[Sequence[ufl.Form]] | None = None,
452 entity_maps: Sequence[dolfinx.mesh.EntityMap] | None = None,
453 kind: str | Sequence[Sequence[str | None]] | None = None,
454 ):
455 # One constraint per block gives a nest or a monolithic system, one constraint a single matrix
456 if not (kind is None or isinstance(kind, (str, Sequence))):
457 raise ValueError(
458 f"Unsupported kind {kind!r}, expected None, a PETSc matrix type or a nested sequence of them"
459 )
460 if isinstance(mpc, Sequence):
461 # Without a kind a problem with one constraint per block keeps the nest layout it always had
462 kind = "nest" if kind is None else kind
463 elif kind == "nest" or (kind is not None and not isinstance(kind, str)):
464 raise ValueError(f"kind={kind!r} needs a sequence of constraints, one per block")
465 # Compile forms
466 form_compiler_options = {} if form_compiler_options is None else form_compiler_options
467 jit_options = {} if jit_options is None else jit_options
468 self._a = _fem.form(
469 a,
470 jit_options=jit_options,
471 form_compiler_options=form_compiler_options,
472 entity_maps=entity_maps,
473 )
474 self._L = _fem.form(
475 L,
476 jit_options=jit_options,
477 form_compiler_options=form_compiler_options,
478 entity_maps=entity_maps,
479 )
481 self._mpc = mpc
482 # Blocked problems
483 if isinstance(mpc, Sequence):
484 is_blocked = True
485 # Sanity check
486 for mpc_i in mpc:
487 if not mpc_i.finalized:
488 raise RuntimeError("The multi point constraint has to be finalized before calling initializer")
489 # Create function containing solution vector
490 else:
491 is_blocked = False
492 if not mpc.finalized:
493 raise RuntimeError("The multi point constraint has to be finalized before calling initializer")
495 # Create function(s) containing solution vector(s)
496 if is_blocked:
497 if u is None:
498 assert isinstance(self._mpc, Sequence)
499 self._u = [_fem.Function(self._mpc[i].function_space) for i in range(len(self._mpc))]
500 else:
501 assert isinstance(self._mpc, Sequence)
502 assert isinstance(u, Sequence)
503 for i, (mpc_i, u_i) in enumerate(zip(self._mpc, u)):
504 assert isinstance(u_i, _fem.Function)
505 assert isinstance(mpc_i, MultiPointConstraint)
506 if u_i.function_space is not mpc_i.function_space:
507 raise ValueError(
508 "The input function has to be in the function space in the multi-point constraint",
509 "i.e. u = dolfinx.fem.Function(mpc.function_space)",
510 )
511 self._u = list(u)
512 else:
513 if u is None:
514 assert isinstance(self._mpc, MultiPointConstraint)
515 self._u = _fem.Function(self._mpc.function_space)
516 else:
517 assert isinstance(u, _fem.Function)
518 assert isinstance(self._mpc, MultiPointConstraint)
519 if u.function_space is self._mpc.function_space:
520 self._u = u
521 else:
522 raise ValueError(
523 "The input function has to be in the function space in the multi-point constraint",
524 "i.e. u = dolfinx.fem.Function(mpc.function_space)",
525 )
527 # Markers and diagonal rows depend only on the function spaces and the
528 # conditions, both fixed for this object, so one cache covers the
529 # operator and the preconditioner and every solve reuses it.
530 self._bc_data = BCData(bcs)
532 # Create MPC matrix and vector
533 self._preconditioner = _fem.form( # type: ignore
534 P,
535 jit_options=jit_options,
536 form_compiler_options=form_compiler_options,
537 entity_maps=entity_maps,
538 )
540 if is_blocked:
541 assert isinstance(mpc, Sequence)
542 assert isinstance(self._L, Sequence)
543 assert isinstance(self._a, Sequence)
544 self._A = create_matrix(self._a, mpc, kind)
545 self._b = create_vector(self._L, mpc, kind)
546 self._x = create_vector(self._L, mpc, kind)
547 if self._preconditioner is None:
548 self._P_mat = None
549 else:
550 assert isinstance(self._preconditioner, Sequence)
551 self._P_mat = create_matrix(self._preconditioner, mpc, kind)
552 else:
553 assert isinstance(mpc, MultiPointConstraint)
554 assert isinstance(self._L, _fem.Form)
555 assert isinstance(self._a, _fem.Form)
556 self._A = create_matrix(self._a, mpc, kind)
557 self._b = create_vector(self._L, mpc)
558 self._x = create_vector(self._L, mpc)
559 if self._preconditioner is None:
560 self._P_mat = None
561 else:
562 assert isinstance(self._preconditioner, _fem.Form)
563 self._P_mat = create_matrix(self._preconditioner, mpc, kind)
565 self.bcs = [] if bcs is None else bcs
567 if is_blocked:
568 assert isinstance(self.u, Sequence)
569 comm = self.u[0].function_space.mesh.comm
570 else:
571 assert isinstance(self.u, _fem.Function)
572 comm = self.u.function_space.mesh.comm
574 self._solver = PETSc.KSP().create(comm)
575 self._solver.setOperators(self._A, self._P_mat)
577 self._petsc_options_prefix = petsc_options_prefix
578 self.solver.setOptionsPrefix(petsc_options_prefix)
579 self.A.setOptionsPrefix(f"{petsc_options_prefix}A_")
580 self.b.setOptionsPrefix(f"{petsc_options_prefix}b_")
581 self.x.setOptionsPrefix(f"{petsc_options_prefix}x_")
582 if self.P_mat is not None:
583 self.P_mat.setOptionsPrefix(f"{petsc_options_prefix}P_mat_")
585 # Set options on KSP only
586 if petsc_options is not None:
587 opts = PETSc.Options()
588 opts.prefixPush(self.solver.getOptionsPrefix())
590 for k, v in petsc_options.items():
591 opts.setValue(k, v)
593 self.solver.setFromOptions()
595 # Tidy up global options
596 for k in petsc_options.keys():
597 opts.delValue(k)
598 opts.prefixPop()
600 def solve(self) -> _fem.Function | Sequence[_fem.Function]:
601 """Solve the problem.
603 Returns:
604 Function containing the solution"""
606 # Refresh the constraint offsets, so that a change in the values of the
607 # Dirichlet conditions held by the constraint is picked up
608 if isinstance(self._mpc, Sequence):
609 for mpc_i in self._mpc:
610 mpc_i.update_constants()
611 else:
612 self._mpc.update_constants()
614 # Assemble lhs. The layout, single, nest or monolithic, is that of the matrix
615 self._A.zeroEntries()
616 assemble_matrix(self._a, self._mpc, bcs=self.bcs, diagval=1.0, A=self._A, bc_data=self._bc_data) # type: ignore
618 self._A.assemble()
619 assert self._A.assembled
621 # Assemble the preconditioner if provided
622 if self._P_mat is not None:
623 self._P_mat.zeroEntries()
624 assemble_matrix(self._preconditioner, self._mpc, bcs=self.bcs, A=self._P_mat, bc_data=self._bc_data) # type: ignore
625 self._P_mat.assemble()
627 # Assemble the residual
628 _zero_vector(self._b)
629 assemble_vector(self._L, self._mpc, self._b) # type: ignore
631 # Lift vector
632 # Decide between nest/blocked and single form lifting up front, so that a
633 # failure in one of the lifting calls cannot fall through to the other
634 # branch and apply the lifting twice
635 try:
636 bcs1 = _fem.bcs.bcs_by_block(_fem.forms.extract_function_spaces(self._a, 1), self.bcs) # type: ignore
637 bcs0 = _fem.bcs.bcs_by_block(_row_spaces(self._L, self._mpc), self.bcs) # type: ignore
638 blocked = True
639 except ValueError:
640 blocked = False
642 if blocked:
643 # Nest and blocked lifting
644 apply_lifting(self._b, self._a, bcs=bcs1, constraint=self._mpc, bc_data=self._bc_data) # type: ignore
645 apply_mpc_lifting(self._b, self._a, constraint=self._mpc) # type: ignore
646 _ghost_update(self._b, PETSc.InsertMode.ADD, PETSc.ScatterMode.REVERSE) # type: ignore
647 _fem.petsc.set_bc(self._b, bcs0)
648 else:
649 # Single form lifting
650 apply_lifting(self._b, [self._a], bcs=[self.bcs], constraint=self._mpc, bc_data=self._bc_data) # type: ignore
651 apply_mpc_lifting(self._b, [self._a], constraint=self._mpc) # type: ignore
652 _ghost_update(self._b, PETSc.InsertMode.ADD, PETSc.ScatterMode.REVERSE) # type: ignore
653 _fem.petsc.set_bc(self._b, self.bcs)
654 _ghost_update(self._b, PETSc.InsertMode.INSERT, PETSc.ScatterMode.FORWARD) # type: ignore
656 # Solve linear system and update ghost values in the solution
657 self._solver.solve(self._b, self._x)
658 _ghost_update(self._x, PETSc.InsertMode.INSERT, PETSc.ScatterMode.FORWARD) # type: ignore
659 _fem.petsc.assign(self._x, self.u) # type: ignore
661 _backsubstitute(self._mpc, self.u) # type: ignore
663 return self._u