Coverage for python/src/dolfinx_mpc/problem.py: 77%
286 statements
« prev ^ index » next coverage.py v7.16.1, created at 2026-09-16 10:47 +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
12from petsc4py import PETSc
14import dolfinx.fem.petsc
15import ufl
16from dolfinx import fem as _fem
17from dolfinx.la.petsc import _ghost_update, _zero_vector, create_vector
19from dolfinx_mpc.cpp import mpc as _cpp_mpc
21from .assemble_matrix import assemble_matrix, assemble_matrix_nest, create_matrix_nest
22from .assemble_vector import (
23 apply_lifting,
24 apply_mpc_lifting,
25 assemble_vector,
26 assemble_vector_nest,
27 create_vector_nest,
28)
29from .multipointconstraint import MultiPointConstraint
32def assemble_jacobian_mpc(
33 u: Sequence[_fem.Function] | _fem.Function,
34 jacobian: _fem.Form | Sequence[Sequence[_fem.Form]],
35 preconditioner: _fem.Form | Sequence[Sequence[_fem.Form]] | None,
36 bcs: Iterable[_fem.DirichletBC],
37 mpc: MultiPointConstraint | Sequence[MultiPointConstraint],
38 _snes: PETSc.SNES, # type: ignore
39 x: PETSc.Vec, # type: ignore
40 J: PETSc.Mat, # type: ignore
41 P: PETSc.Mat, # type: ignore
42):
43 """Assemble the Jacobian matrix and preconditioner.
45 A function conforming to the interface expected by SNES.setJacobian can
46 be created by fixing the first four arguments:
48 functools.partial(assemble_jacobian, u, jacobian, preconditioner,
49 bcs)
51 Args:
52 u: Function tied to the solution vector within the residual and
53 jacobian
54 jacobian: Form of the Jacobian
55 preconditioner: Form of the preconditioner
56 bcs: List of Dirichlet boundary conditions
57 mpc: The multi point constraint or a sequence of multi point
58 _snes: The solver instance
59 x: The vector containing the point to evaluate at
60 J: Matrix to assemble the Jacobian into
61 P: Matrix to assemble the preconditioner into
62 """
63 # Copy existing soultion into the function used in the residual and
64 # Jacobian
65 _ghost_update(x, PETSc.InsertMode.INSERT, PETSc.ScatterMode.FORWARD) # type: ignore
66 # Assign the input vector to the unknowns
67 _fem.petsc.assign(x, u) # type: ignore
68 if isinstance(u, Sequence):
69 assert isinstance(mpc, Sequence)
70 for i in range(len(u)):
71 mpc[i].homogenize(u[i])
72 mpc[i].backsubstitution(u[i])
73 else:
74 assert isinstance(u, _fem.Function)
75 assert isinstance(mpc, MultiPointConstraint)
76 mpc.homogenize(u)
77 mpc.backsubstitution(u)
79 # Assemble Jacobian
80 J.zeroEntries()
81 if J.getType() == "nest":
82 assemble_matrix_nest(J, jacobian, mpc, bcs, diagval=1.0) # type: ignore
83 else:
84 assemble_matrix(jacobian, mpc, bcs, diagval=1.0, A=J) # type: ignore
85 J.assemble()
86 if preconditioner is not None:
87 P.zeroEntries()
88 if P.getType() == "nest":
89 assemble_matrix_nest(P, preconditioner, mpc, bcs, diagval=1.0) # type: ignore
90 else:
91 assemble_matrix(mpc, preconditioner, bcs, diagval=1.0, A=P) # type: ignore
93 P.assemble()
96def assemble_residual_mpc(
97 u: _fem.Function | Sequence[_fem.Function],
98 residual: _fem.Form | Sequence[_fem.Form],
99 jacobian: _fem.Form | Sequence[Sequence[_fem.Form]],
100 bcs: Sequence[_fem.DirichletBC],
101 mpc: MultiPointConstraint | Sequence[MultiPointConstraint],
102 _snes: PETSc.SNES, # type: ignore
103 x: PETSc.Vec, # type: ignore
104 F: PETSc.Vec, # type: ignore
105):
106 """Assemble the residual into the vector `F`.
108 A function conforming to the interface expected by SNES.setResidual can
109 be created by fixing the first four arguments:
111 functools.partial(assemble_residual, u, jacobian, preconditioner,
112 bcs)
114 Args:
115 u: Function(s) tied to the solution vector within the residual and
116 Jacobian.
117 residual: Form of the residual. It can be a sequence of forms.
118 jacobian: Form of the Jacobian. It can be a nested sequence of
119 forms.
120 bcs: List of Dirichlet boundary conditions.
121 _snes: The solver instance.
122 x: The vector containing the point to evaluate the residual at.
123 F: Vector to assemble the residual into.
124 """
125 # Update input vector before assigning
126 _ghost_update(x, PETSc.InsertMode.INSERT, PETSc.ScatterMode.FORWARD) # type: ignore
127 # Assign the input vector to the unknowns
128 _fem.petsc.assign(x, u) # type: ignore
129 if isinstance(u, Sequence):
130 assert isinstance(mpc, Sequence)
131 for i in range(len(u)):
132 mpc[i].homogenize(u[i])
133 mpc[i].backsubstitution(u[i])
134 else:
135 assert isinstance(u, _fem.Function)
136 assert isinstance(mpc, MultiPointConstraint)
137 mpc.homogenize(u)
138 mpc.backsubstitution(u)
139 # Assemble the residual
140 _zero_vector(F)
141 if x.getType() == "nest":
142 assemble_vector_nest(F, residual, mpc) # type: ignore
143 else:
144 assert isinstance(residual, _fem.Form)
145 assert isinstance(mpc, MultiPointConstraint)
146 assemble_vector(residual, mpc, F)
148 # Lift vector
149 try:
150 # Nest and blocked lifting
151 bcs1 = _fem.bcs.bcs_by_block(_fem.forms.extract_function_spaces(jacobian, 1), bcs) # type: ignore
152 _fem.petsc._assign_block_data(residual, x) # type: ignore
153 apply_lifting(F, jacobian, bcs=bcs1, constraint=mpc, x0=x, scale=-1.0) # type: ignore
154 _ghost_update(F, PETSc.InsertMode.ADD, PETSc.ScatterMode.REVERSE) # type: ignore
155 bcs0 = _fem.bcs.bcs_by_block(_fem.forms.extract_function_spaces(residual), bcs) # type: ignore
156 _fem.petsc.set_bc(F, bcs0, x0=x, alpha=-1.0)
157 except (TypeError, ValueError):
158 # Single form lifting
159 apply_lifting(F, [jacobian], bcs=[bcs], constraint=mpc, x0=[x], scale=-1.0) # type: ignore
160 _ghost_update(F, PETSc.InsertMode.ADD, PETSc.ScatterMode.REVERSE) # type: ignore
161 _fem.petsc.set_bc(F, bcs, x0=x, alpha=-1.0)
162 _ghost_update(F, PETSc.InsertMode.INSERT, PETSc.ScatterMode.FORWARD) # type: ignore
165class NonlinearProblem(dolfinx.fem.petsc.NonlinearProblem):
166 def __init__(
167 self,
168 F: ufl.form.Form | Sequence[ufl.form.Form],
169 u: _fem.Function | Sequence[_fem.Function],
170 mpc: MultiPointConstraint | Sequence[MultiPointConstraint],
171 bcs: Sequence[_fem.DirichletBC] | None = None,
172 J: ufl.form.Form | Sequence[Sequence[ufl.form.Form]] | None = None,
173 P: ufl.form.Form | Sequence[Sequence[ufl.form.Form]] | None = None,
174 kind: str | Sequence[Sequence[str]] | None = None,
175 form_compiler_options: dict | None = None,
176 jit_options: dict | None = None,
177 petsc_options_prefix: str = "dolfinx_mpc_nonlinear_problem_",
178 petsc_options: dict | None = None,
179 entity_maps: Sequence[dolfinx.mesh.EntityMap] | None = None,
180 ):
181 """Class for solving nonlinear problems with SNES.
183 Solves problems of the form
184 :math:`F_i(u, v) = 0, i=0,...N\\ \\forall v \\in V` where
185 :math:`u=(u_0,...,u_N), v=(v_0,...,v_N)` using PETSc SNES as the
186 non-linear solver.
188 Note: The deprecated version of this class for use with
189 NewtonSolver has been renamed NewtonSolverNonlinearProblem.
191 Args:
192 F: UFL form(s) of residual :math:`F_i`.
193 u: Function used to define the residual and Jacobian.
194 bcs: Dirichlet boundary conditions.
195 J: UFL form(s) representing the Jacobian
196 :math:`J_ij = dF_i/du_j`.
197 P: UFL form(s) representing the preconditioner.
198 kind: The PETSc matrix type(s) for the Jacobian and
199 preconditioner (``MatType``).
200 See :func:`dolfinx.fem.petsc.create_matrix` for more
201 information.
202 form_compiler_options: Options used in FFCx compilation of all
203 forms. Run ``ffcx --help`` at the command line to see all
204 available options.
205 jit_options: Options used in CFFI JIT compilation of C code
206 generated by FFCx. See ``python/dolfinx/jit.py`` for all
207 available options. Takes priority over all other option
208 values.
209 petsc_options_prefix: Options prefix used as the root prefix on
210 all internally created PETSc objects (SNES, A, b, x and the
211 preconditioner matrix). Typically ends with ``_``. Must be the
212 same on all ranks and is usually unique within the
213 programme.
214 petsc_options: Options to pass to the PETSc SNES object.
215 entity_maps: If any trial functions, test functions, or
216 coefficients in the form are not defined over the same mesh
217 as the integration domain, ``entity_maps`` must be
218 supplied. For each key (a mesh, different to the
219 integration domain mesh) a map should be provided relating
220 the entities in the integration domain mesh to the entities
221 in the key mesh e.g. for a key-value pair ``(msh, emap)``
222 in ``entity_maps``, ``emap[i]`` is the entity in ``msh``
223 corresponding to entity ``i`` in the integration domain
224 mesh.
225 """
226 # Compile residual and Jacobian forms
227 self._F = _fem.form(
228 F,
229 form_compiler_options=form_compiler_options,
230 jit_options=jit_options,
231 entity_maps=entity_maps,
232 )
234 if J is None:
235 if isinstance(F, ufl.form.Form):
236 du = ufl.TrialFunction(self._F.arguments()[0].ufl_function_space())
237 J = ufl.derivative(F, du)
238 else:
239 dus = [ufl.TrialFunction(Fi.arguments()[0].ufl_function_space()) for Fi in F]
240 J = _fem.forms.derivative_block(F, u, dus)
242 self._J = _fem.form(
243 J,
244 form_compiler_options=form_compiler_options,
245 jit_options=jit_options,
246 entity_maps=entity_maps,
247 )
249 if P is not None:
250 self._preconditioner = _fem.form(
251 P,
252 form_compiler_options=form_compiler_options,
253 jit_options=jit_options,
254 entity_maps=entity_maps,
255 )
256 else:
257 self._preconditioner = None
259 self._u = u
260 # Set default values if not supplied
261 bcs = [] if bcs is None else bcs
262 self.mpc = mpc
263 # Create PETSc structures for the residual, Jacobian and solution
264 # vector
265 if kind == "nest" or isinstance(kind, Sequence):
266 assert isinstance(mpc, Sequence)
267 assert isinstance(self._J, Sequence)
268 self._A = create_matrix_nest(self._J, mpc)
269 elif kind is None:
270 assert isinstance(mpc, MultiPointConstraint)
271 self._A = _cpp_mpc.create_matrix(self._J._cpp_object, mpc._cpp_object)
272 else:
273 raise ValueError("Unsupported kind for matrix: {}".format(kind))
275 kind = "nest" if self._A.getType() == "nest" else kind
276 if kind == "nest":
277 assert isinstance(mpc, Sequence)
278 assert isinstance(self._F, Sequence)
279 self._b = create_vector_nest(self._F, mpc)
280 self._x = create_vector_nest(self._F, mpc)
281 else:
282 assert isinstance(mpc, MultiPointConstraint)
283 self._b = create_vector([(mpc.function_space.dofmap.index_map, mpc.function_space.dofmap.index_map_bs)])
284 self._x = create_vector([(mpc.function_space.dofmap.index_map, mpc.function_space.dofmap.index_map_bs)])
286 # Create PETSc structure for preconditioner if provided
287 if self.preconditioner is not None: # type: ignore
288 if kind == "nest":
289 assert isinstance(self.preconditioner, Sequence)
290 assert isinstance(self.mpc, Sequence)
291 self._P_mat = create_matrix_nest(self.preconditioner, self.mpc)
292 else:
293 assert isinstance(self._preconditioner, _fem.Form)
294 self._P_mat = _cpp_mpc.create_matrix(self.preconditioner, kind=kind)
295 else:
296 self._P_mat = None # type: ignore
298 # Create the SNES solver and attach the corresponding Jacobian and
299 # residual computation functions
300 self._snes = PETSc.SNES().create(comm=self.A.comm) # type: ignore
301 self.solver.setJacobian(
302 partial(assemble_jacobian_mpc, u, self.J, self.preconditioner, bcs, mpc), self._A, self.P_mat
303 )
304 self.solver.setFunction(partial(assemble_residual_mpc, u, self.F, self.J, bcs, mpc), self.b)
306 # Set PETSc options
307 self._petsc_options_prefix = petsc_options_prefix
308 self.solver.setOptionsPrefix(petsc_options_prefix)
309 self.A.setOptionsPrefix(f"{petsc_options_prefix}A_")
310 self.b.setOptionsPrefix(f"{petsc_options_prefix}b_")
311 self.x.setOptionsPrefix(f"{petsc_options_prefix}x_")
312 if self.P_mat is not None:
313 self.P_mat.setOptionsPrefix(f"{petsc_options_prefix}P_mat_")
315 # Set options on SNES only
316 if petsc_options is not None:
317 opts = PETSc.Options() # type: ignore
318 opts.prefixPush(self.solver.getOptionsPrefix())
320 for k, v in petsc_options.items():
321 opts.setValue(k, v)
323 self.solver.setFromOptions()
325 # Tidy up global options
326 for k in petsc_options.keys():
327 opts.delValue(k)
328 opts.prefixPop()
330 def solve(self) -> tuple[_fem.Function | Sequence[_fem.Function], PETSc.ConvergedReason, int]: # type: ignore
331 """Solve the problem and update the solution in the problem
332 instance.
334 Returns:
335 The solution, convergence reason and number of iterations.
336 """
338 # Move current iterate into the work array.
339 _fem.petsc.assign(self._u, self.x)
341 # Solve problem
342 self.solver.solve(None, self.x)
344 # Move solution back to function
345 dolfinx.fem.petsc.assign(self.x, self._u) # type: ignore
346 if isinstance(self.mpc, Sequence):
347 for i in range(len(self._u)):
348 self.mpc[i].backsubstitution(self._u[i])
349 self.mpc[i].backsubstitution(self._u[i])
350 else:
351 assert isinstance(self._u, _fem.Function)
352 self.mpc.homogenize(self._u)
353 self.mpc.backsubstitution(self._u)
355 converged_reason = self.solver.getConvergedReason()
356 return self._u, converged_reason, self.solver.getIterationNumber() # type: ignore
359class LinearProblem(dolfinx.fem.petsc.LinearProblem):
360 """
361 Class for solving a linear variational problem with multi point constraints of the form
362 a(u, v) = L(v) for all v using PETSc as a linear algebra backend.
364 Args:
365 a: A bilinear UFL form, the left hand side of the variational problem.
366 L: A linear UFL form, the right hand side of the variational problem.
367 mpc: The multi point constraint.
368 bcs: A list of Dirichlet boundary conditions.
369 u: The solution function. It will be created if not provided. The function has
370 to be based on the functionspace in the mpc, i.e.
372 .. highlight:: python
373 .. code-block:: python
375 u = dolfinx.fem.Function(mpc.function_space)
376 petsc_options: Parameters that is passed to the linear algebra backend PETSc. #type: ignore
377 For available choices for the 'petsc_options' kwarg, see the PETSc-documentation
378 https://www.mcs.anl.gov/petsc/documentation/index.html.
379 form_compiler_options: Parameters used in FFCx compilation of this form. Run `ffcx --help` at
380 the commandline to see all available options. Takes priority over all
381 other parameter values, except for `scalar_type` which is determined by DOLFINx.
382 jit_options: Parameters used in CFFI JIT compilation of C code generated by FFCx.
383 See https://github.com/FEniCS/dolfinx/blob/main/python/dolfinx/jit.py#L22-L37
384 for all available parameters. Takes priority over all other parameter values.
385 P: A preconditioner UFL form.
386 Examples:
387 Example usage:
389 .. highlight:: python
390 .. code-block:: python
392 problem = LinearProblem(a, L, mpc, [bc0, bc1],
393 petsc_options={"ksp_type": "preonly", "pc_type": "lu"})
395 """
397 u: _fem.Function | list[_fem.Function]
398 _a: _fem.Form | Sequence[Sequence[_fem.Form]]
399 _L: _fem.Form | Sequence[_fem.Form]
400 _preconditioner: _fem.Form | Sequence[Sequence[_fem.Form]] | None
401 _mpc: MultiPointConstraint | Sequence[MultiPointConstraint]
402 _A: PETSc.Mat
403 _P: PETSc.Mat | None
404 _b: PETSc.Vec
405 _solver: PETSc.KSP
406 _x: PETSc.Vec
407 bcs: list[_fem.DirichletBC]
408 __slots__ = tuple(__annotations__)
410 def __init__(
411 self,
412 a: ufl.Form | Sequence[Sequence[ufl.Form]],
413 L: ufl.Form | Sequence[ufl.Form],
414 mpc: MultiPointConstraint | Sequence[MultiPointConstraint],
415 bcs: list[_fem.DirichletBC] | None = None,
416 u: _fem.Function | Sequence[_fem.Function] | None = None,
417 petsc_options_prefix: str = "dolfinx_mpc_linear_problem_",
418 petsc_options: dict | None = None,
419 form_compiler_options: dict | None = None,
420 jit_options: dict | None = None,
421 P: ufl.Form | Sequence[Sequence[ufl.Form]] | None = None,
422 ):
423 # Compile forms
424 form_compiler_options = {} if form_compiler_options is None else form_compiler_options
425 jit_options = {} if jit_options is None else jit_options
426 self._a = _fem.form(a, jit_options=jit_options, form_compiler_options=form_compiler_options)
427 self._L = _fem.form(L, jit_options=jit_options, form_compiler_options=form_compiler_options)
429 self._mpc = mpc
430 # Nest assembly
431 if isinstance(mpc, Sequence):
432 is_nest = True
433 # Sanity check
434 for mpc_i in mpc:
435 if not mpc_i.finalized:
436 raise RuntimeError("The multi point constraint has to be finalized before calling initializer")
437 # Create function containing solution vector
438 else:
439 is_nest = False
440 if not mpc.finalized:
441 raise RuntimeError("The multi point constraint has to be finalized before calling initializer")
443 # Create function(s) containing solution vector(s)
444 if is_nest:
445 if u is None:
446 assert isinstance(self._mpc, Sequence)
447 self.u = [_fem.Function(self._mpc[i].function_space) for i in range(len(self._mpc))]
448 else:
449 assert isinstance(self._mpc, Sequence)
450 assert isinstance(u, Sequence)
451 for i, (mpc_i, u_i) in enumerate(zip(self._mpc, u)):
452 assert isinstance(u_i, _fem.Function)
453 assert isinstance(mpc_i, MultiPointConstraint)
454 if u_i.function_space is not mpc_i.function_space:
455 raise ValueError(
456 "The input function has to be in the function space in the multi-point constraint",
457 "i.e. u = dolfinx.fem.Function(mpc.function_space)",
458 )
459 self.u = list(u)
460 else:
461 if u is None:
462 assert isinstance(self._mpc, MultiPointConstraint)
463 self.u = _fem.Function(self._mpc.function_space)
464 else:
465 assert isinstance(u, _fem.Function)
466 assert isinstance(self._mpc, MultiPointConstraint)
467 if u.function_space is self._mpc.function_space:
468 self.u = u
469 else:
470 raise ValueError(
471 "The input function has to be in the function space in the multi-point constraint",
472 "i.e. u = dolfinx.fem.Function(mpc.function_space)",
473 )
475 # Create MPC matrix and vector
476 self._preconditioner = _fem.form(P, jit_options=jit_options, form_compiler_options=form_compiler_options) # type: ignore
478 if is_nest:
479 assert isinstance(mpc, Sequence)
480 assert isinstance(self._L, Sequence)
481 assert isinstance(self._a, Sequence)
482 self._A = create_matrix_nest(self._a, mpc)
483 self._b = create_vector_nest(self._L, mpc)
484 self._x = create_vector_nest(self._L, mpc)
485 if self._preconditioner is None:
486 self._P_mat = None
487 else:
488 assert isinstance(self._preconditioner, Sequence)
489 self._P_mat = create_matrix_nest(self._preconditioner, mpc)
490 else:
491 assert isinstance(mpc, MultiPointConstraint)
492 assert isinstance(self._L, _fem.Form)
493 assert isinstance(self._a, _fem.Form)
494 self._A = _cpp_mpc.create_matrix(self._a._cpp_object, mpc._cpp_object)
495 self._b = create_vector([(mpc.function_space.dofmap.index_map, mpc.function_space.dofmap.index_map_bs)])
496 self._x = create_vector([(mpc.function_space.dofmap.index_map, mpc.function_space.dofmap.index_map_bs)])
497 if self._preconditioner is None:
498 self._P_mat = None
499 else:
500 assert isinstance(self._preconditioner, _fem.Form)
501 self._P_mat = _cpp_mpc.create_matrix(self._preconditioner._cpp_object, mpc._cpp_object)
503 self.bcs = [] if bcs is None else bcs
505 if is_nest:
506 assert isinstance(self.u, Sequence)
507 comm = self.u[0].function_space.mesh.comm
508 else:
509 assert isinstance(self.u, _fem.Function)
510 comm = self.u.function_space.mesh.comm
512 self._solver = PETSc.KSP().create(comm)
513 self._solver.setOperators(self._A, self._P_mat)
515 self._petsc_options_prefix = petsc_options_prefix
516 self.solver.setOptionsPrefix(petsc_options_prefix)
517 self.A.setOptionsPrefix(f"{petsc_options_prefix}A_")
518 self.b.setOptionsPrefix(f"{petsc_options_prefix}b_")
519 self.x.setOptionsPrefix(f"{petsc_options_prefix}x_")
520 if self.P_mat is not None:
521 self.P_mat.setOptionsPrefix(f"{petsc_options_prefix}P_mat_")
523 # Set options on KSP only
524 if petsc_options is not None:
525 opts = PETSc.Options()
526 opts.prefixPush(self.solver.getOptionsPrefix())
528 for k, v in petsc_options.items():
529 opts.setValue(k, v)
531 self.solver.setFromOptions()
533 # Tidy up global options
534 for k in petsc_options.keys():
535 opts.delValue(k)
536 opts.prefixPop()
538 def solve(self) -> _fem.Function | list[_fem.Function]:
539 """Solve the problem.
541 Returns:
542 Function containing the solution"""
544 # Refresh the constraint offsets, so that a change in the values of the
545 # Dirichlet conditions held by the constraint is picked up
546 if isinstance(self._mpc, Sequence):
547 for mpc_i in self._mpc:
548 mpc_i.update_constants()
549 else:
550 self._mpc.update_constants()
552 # Assemble lhs
553 self._A.zeroEntries()
554 if self._A.getType() == "nest":
555 assemble_matrix_nest(self._A, self._a, self._mpc, self.bcs, diagval=1.0) # type: ignore
556 else:
557 assert isinstance(self._a, _fem.Form)
558 assemble_matrix(self._a, self._mpc, bcs=self.bcs, A=self._A)
560 self._A.assemble()
561 assert self._A.assembled
563 # Assemble the preconditioner if provided
564 if self._P_mat is not None:
565 self._P_mat.zeroEntries()
566 if self._P_mat.getType() == "nest":
567 assert isinstance(self._preconditioner, Sequence)
568 assemble_matrix_nest(self._P_mat, self._preconditioner, self._mpc, self.bcs) # type: ignore
569 else:
570 assert isinstance(self._preconditioner, _fem.Form)
571 assemble_matrix(self._preconditioner, self._mpc, bcs=self.bcs, A=self._P_mat)
572 self._P_mat.assemble()
574 # Assemble the residual
575 _zero_vector(self._b)
576 if self._x.getType() == "nest":
577 assemble_vector_nest(self._b, self._L, self._mpc) # type: ignore
578 else:
579 assert isinstance(self._L, _fem.Form)
580 assert isinstance(self._mpc, MultiPointConstraint)
581 assemble_vector(self._L, self._mpc, self._b)
583 # Lift vector
584 # Decide between nest/blocked and single form lifting up front, so that a
585 # failure in one of the lifting calls cannot fall through to the other
586 # branch and apply the lifting twice
587 try:
588 bcs1 = _fem.bcs.bcs_by_block(_fem.forms.extract_function_spaces(self._a, 1), self.bcs) # type: ignore
589 bcs0 = _fem.bcs.bcs_by_block(_fem.forms.extract_function_spaces(self._L), self.bcs) # type: ignore
590 blocked = True
591 except ValueError:
592 blocked = False
594 if blocked:
595 # Nest and blocked lifting
596 apply_lifting(self._b, self._a, bcs=bcs1, constraint=self._mpc) # type: ignore
597 apply_mpc_lifting(self._b, self._a, constraint=self._mpc) # type: ignore
598 _ghost_update(self._b, PETSc.InsertMode.ADD, PETSc.ScatterMode.REVERSE) # type: ignore
599 _fem.petsc.set_bc(self._b, bcs0)
600 else:
601 # Single form lifting
602 apply_lifting(self._b, [self._a], bcs=[self.bcs], constraint=self._mpc) # type: ignore
603 apply_mpc_lifting(self._b, [self._a], constraint=self._mpc) # type: ignore
604 _ghost_update(self._b, PETSc.InsertMode.ADD, PETSc.ScatterMode.REVERSE) # type: ignore
605 _fem.petsc.set_bc(self._b, self.bcs)
606 _ghost_update(self._b, PETSc.InsertMode.INSERT, PETSc.ScatterMode.FORWARD) # type: ignore
608 # Solve linear system and update ghost values in the solution
609 self._solver.solve(self._b, self._x)
610 _ghost_update(self._x, PETSc.InsertMode.INSERT, PETSc.ScatterMode.FORWARD) # type: ignore
611 _fem.petsc.assign(self._x, self.u) # type: ignore
613 if isinstance(self.u, Sequence):
614 assert isinstance(self._mpc, Sequence)
615 for i in range(len(self.u)):
616 self._mpc[i].homogenize(self.u[i])
617 self._mpc[i].backsubstitution(self.u[i])
618 else:
619 assert isinstance(self.u, _fem.Function)
620 assert isinstance(self._mpc, MultiPointConstraint)
621 self._mpc.homogenize(self.u)
622 self._mpc.backsubstitution(self.u)
624 return self.u