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
« 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
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 apply_lifting, assemble_vector, assemble_vector_nest, create_vector_nest
23from .multipointconstraint import MultiPointConstraint
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.
39 A function conforming to the interface expected by SNES.setJacobian can
40 be created by fixing the first four arguments:
42 functools.partial(assemble_jacobian, u, jacobian, preconditioner,
43 bcs)
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)
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
87 P.assemble()
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`.
102 A function conforming to the interface expected by SNES.setResidual can
103 be created by fixing the first four arguments:
105 functools.partial(assemble_residual, u, jacobian, preconditioner,
106 bcs)
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)
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
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.
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.
182 Note: The deprecated version of this class for use with
183 NewtonSolver has been renamed NewtonSolverNonlinearProblem.
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 )
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)
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 )
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
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))
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)])
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
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)
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_")
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())
314 for k, v in petsc_options.items():
315 opts.setValue(k, v)
317 self.solver.setFromOptions()
319 # Tidy up global options
320 for k in petsc_options.keys():
321 opts.delValue(k)
322 opts.prefixPop()
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.
328 Returns:
329 The solution, convergence reason and number of iterations.
330 """
332 # Move current iterate into the work array.
333 _fem.petsc.assign(self._u, self.x)
335 # Solve problem
336 self.solver.solve(None, self.x)
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)
349 converged_reason = self.solver.getConvergedReason()
350 return self._u, converged_reason, self.solver.getIterationNumber() # type: ignore
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.
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.
366 .. highlight:: python
367 .. code-block:: python
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:
383 .. highlight:: python
384 .. code-block:: python
386 problem = LinearProblem(a, L, mpc, [bc0, bc1],
387 petsc_options={"ksp_type": "preonly", "pc_type": "lu"})
389 """
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__)
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)
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")
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 )
469 # Create MPC matrix and vector
470 self._preconditioner = _fem.form(P, jit_options=jit_options, form_compiler_options=form_compiler_options) # type: ignore
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)
497 self.bcs = [] if bcs is None else bcs
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
506 self._solver = PETSc.KSP().create(comm)
507 self._solver.setOperators(self._A, self._P_mat)
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_")
517 # Set options on KSP only
518 if petsc_options is not None:
519 opts = PETSc.Options()
520 opts.prefixPush(self.solver.getOptionsPrefix())
522 for k, v in petsc_options.items():
523 opts.setValue(k, v)
525 self.solver.setFromOptions()
527 # Tidy up global options
528 for k in petsc_options.keys():
529 opts.delValue(k)
530 opts.prefixPop()
532 def solve(self) -> _fem.Function | list[_fem.Function]:
533 """Solve the problem.
535 Returns:
536 Function containing the solution"""
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)
546 self._A.assemble()
547 assert self._A.assembled
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()
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)
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
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
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)
600 return self.u