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

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

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

3# 

4# This file is part of DOLFINx MPC 

5# 

6# SPDX-License-Identifier: MIT 

7from __future__ import annotations 

8 

9from collections.abc import Iterable, Sequence 

10from functools import partial 

11 

12from petsc4py import PETSc 

13 

14import dolfinx.fem.petsc 

15import ufl 

16from dolfinx import fem as _fem 

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

18 

19from dolfinx_mpc.cpp import mpc as _cpp_mpc 

20 

21from .assemble_matrix import assemble_matrix, assemble_matrix_nest, create_matrix_nest 

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

23from .multipointconstraint import MultiPointConstraint 

24 

25 

26def assemble_jacobian_mpc( 

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

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

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

30 bcs: Iterable[_fem.DirichletBC], 

31 mpc: MultiPointConstraint | Sequence[MultiPointConstraint], 

32 _snes: PETSc.SNES, # type: ignore 

33 x: PETSc.Vec, # type: ignore 

34 J: PETSc.Mat, # type: ignore 

35 P: PETSc.Mat, # type: ignore 

36): 

37 """Assemble the Jacobian matrix and preconditioner. 

38 

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

40 be created by fixing the first four arguments: 

41 

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

43 bcs) 

44 

45 Args: 

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

47 jacobian 

48 jacobian: Form of the Jacobian 

49 preconditioner: Form of the preconditioner 

50 bcs: List of Dirichlet boundary conditions 

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

52 _snes: The solver instance 

53 x: The vector containing the point to evaluate at 

54 J: Matrix to assemble the Jacobian into 

55 P: Matrix to assemble the preconditioner into 

56 """ 

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

58 # Jacobian 

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

60 _fem.petsc.assign(x, u) 

61 

62 # Assemble Jacobian 

63 J.zeroEntries() 

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

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

66 else: 

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

68 J.assemble() 

69 if preconditioner is not None: 

70 P.zeroEntries() 

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

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

73 else: 

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

75 

76 P.assemble() 

77 

78 

79def assemble_residual_mpc( 

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

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

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

83 bcs: Sequence[_fem.DirichletBC], 

84 mpc: MultiPointConstraint | Sequence[MultiPointConstraint], 

85 _snes: PETSc.SNES, # type: ignore 

86 x: PETSc.Vec, # type: ignore 

87 F: PETSc.Vec, # type: ignore 

88): 

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

90 

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

92 be created by fixing the first four arguments: 

93 

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

95 bcs) 

96 

97 Args: 

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

99 Jacobian. 

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

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

102 forms. 

103 bcs: List of Dirichlet boundary conditions. 

104 _snes: The solver instance. 

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

106 F: Vector to assemble the residual into. 

107 """ 

108 # Update input vector before assigning 

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

110 # Assign the input vector to the unknowns 

111 _fem.petsc.assign(x, u) 

112 if isinstance(u, Sequence): 

113 assert isinstance(mpc, Sequence) 

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

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

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

117 else: 

118 assert isinstance(u, _fem.Function) 

119 assert isinstance(mpc, MultiPointConstraint) 

120 mpc.homogenize(u) 

121 mpc.backsubstitution(u) 

122 # Assemble the residual 

123 _zero_vector(F) 

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

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

126 else: 

127 assert isinstance(residual, _fem.Form) 

128 assert isinstance(mpc, MultiPointConstraint) 

129 assemble_vector(residual, mpc, F) 

130 

131 # Lift vector 

132 try: 

133 # Nest and blocked lifting 

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

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

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

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

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

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

140 except TypeError: 

141 # Single form lifting 

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

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

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

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

146 

147 

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

149 def __init__( 

150 self, 

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

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

153 mpc: MultiPointConstraint | Sequence[MultiPointConstraint], 

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

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

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

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

158 form_compiler_options: dict | None = None, 

159 jit_options: dict | None = None, 

160 petsc_options: dict | None = None, 

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

162 ): 

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

164 

165 Solves problems of the form 

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

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

168 non-linear solver. 

169 

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

171 NewtonSolver has been renamed NewtonSolverNonlinearProblem. 

172 

173 Args: 

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

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

176 bcs: Dirichlet boundary conditions. 

177 J: UFL form(s) representing the Jacobian 

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

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

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

181 preconditioner (``MatType``). 

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

183 information. 

184 form_compiler_options: Options used in FFCx compilation of all 

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

186 available options. 

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

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

189 available options. Takes priority over all other option 

190 values. 

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

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

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

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

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

196 integration domain mesh) a map should be provided relating 

197 the entities in the integration domain mesh to the entities 

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

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

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

201 mesh. 

202 """ 

203 # Compile residual and Jacobian forms 

204 self._F = _fem.form( 

205 F, 

206 form_compiler_options=form_compiler_options, 

207 jit_options=jit_options, 

208 entity_maps=entity_maps, 

209 ) 

210 

211 if J is None: 

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

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

214 

215 self._J = _fem.form( 

216 J, 

217 form_compiler_options=form_compiler_options, 

218 jit_options=jit_options, 

219 entity_maps=entity_maps, 

220 ) 

221 

222 if P is not None: 

223 self._preconditioner = _fem.form( 

224 P, 

225 form_compiler_options=form_compiler_options, 

226 jit_options=jit_options, 

227 entity_maps=entity_maps, 

228 ) 

229 else: 

230 self._preconditioner = None 

231 

232 self._u = u 

233 # Set default values if not supplied 

234 bcs = [] if bcs is None else bcs 

235 self.mpc = mpc 

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

237 # vector 

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

239 assert isinstance(mpc, Sequence) 

240 assert isinstance(self._J, Sequence) 

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

242 elif kind is None: 

243 assert isinstance(mpc, MultiPointConstraint) 

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

245 else: 

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

247 

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

249 if kind == "nest": 

250 assert isinstance(mpc, Sequence) 

251 assert isinstance(self._F, Sequence) 

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

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

254 else: 

255 assert isinstance(mpc, MultiPointConstraint) 

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

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

258 

259 # Create PETSc structure for preconditioner if provided 

260 if self.preconditioner is not None: 

261 if kind == "nest": 

262 assert isinstance(self.preconditioner, Sequence) 

263 assert isinstance(self.mpc, Sequence) 

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

265 else: 

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

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

268 else: 

269 self._P_mat = None 

270 

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

272 # residual computation functions 

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

274 self.solver.setJacobian( 

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

276 ) 

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

278 

279 # Set PETSc options 

280 problem_prefix = f"dolfinx_nonlinearproblem_{id(self)}" 

281 self.solver.setOptionsPrefix(problem_prefix) 

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

283 opts.prefixPush(problem_prefix) 

284 

285 if petsc_options is not None: 

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

287 opts[k] = v 

288 self.solver.setFromOptions() 

289 

290 for k in petsc_options.keys(): 

291 del opts[k] 

292 

293 opts.prefixPop() 

294 

295 def solve(self) -> tuple[PETSc.Vec, int, int]: # type: ignore 

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

297 instance. 

298 

299 Returns: 

300 The solution, convergence reason and number of iterations. 

301 """ 

302 

303 # Move current iterate into the work array. 

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

305 

306 # Solve problem 

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

308 

309 # Move solution back to function 

310 dolfinx.fem.petsc.assign(self.x, self._u) 

311 if isinstance(self.mpc, Sequence): 

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

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

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

315 else: 

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

317 self.mpc.homogenize(self._u) 

318 self.mpc.backsubstitution(self._u) 

319 

320 converged_reason = self.solver.getConvergedReason() 

321 return self._u, converged_reason, self.solver.getIterationNumber() 

322 

323 

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

325 """ 

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

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

328 

329 Args: 

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

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

332 mpc: The multi point constraint. 

333 bcs: A list of Dirichlet boundary conditions. 

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

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

336 

337 .. highlight:: python 

338 .. code-block:: python 

339 

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

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

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

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

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

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

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

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

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

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

350 P: A preconditioner UFL form. 

351 Examples: 

352 Example usage: 

353 

354 .. highlight:: python 

355 .. code-block:: python 

356 

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

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

359 

360 """ 

361 

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

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

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

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

366 _mpc: MultiPointConstraint | Sequence[MultiPointConstraint] 

367 _A: PETSc.Mat # type: ignore 

368 _P: PETSc.Mat | None # type: ignore 

369 _b: PETSc.Vec # type: ignore 

370 _solver: PETSc.KSP # type: ignore 

371 _x: PETSc.Vec # type: ignore 

372 bcs: list[_fem.DirichletBC] 

373 __slots__ = tuple(__annotations__) 

374 

375 def __init__( 

376 self, 

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

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

379 mpc: MultiPointConstraint | Sequence[MultiPointConstraint], 

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

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

382 petsc_options_prefix: str = "dolfinx_mpc_linear_problem_", 

383 petsc_options: dict | None = None, 

384 form_compiler_options: dict | None = None, 

385 jit_options: dict | None = None, 

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

387 ): 

388 # Compile forms 

389 form_compiler_options = {} if form_compiler_options is None else form_compiler_options 

390 jit_options = {} if jit_options is None else jit_options 

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

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

393 

394 self._mpc = mpc 

395 # Nest assembly 

396 if isinstance(mpc, Sequence): 

397 is_nest = True 

398 # Sanity check 

399 for mpc_i in mpc: 

400 if not mpc_i.finalized: 

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

402 # Create function containing solution vector 

403 else: 

404 is_nest = False 

405 if not mpc.finalized: 

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

407 

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

409 if is_nest: 

410 if u is None: 

411 assert isinstance(self._mpc, Sequence) 

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

413 else: 

414 assert isinstance(self._mpc, Sequence) 

415 assert isinstance(u, Sequence) 

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

417 assert isinstance(u_i, _fem.Function) 

418 assert isinstance(mpc_i, MultiPointConstraint) 

419 if u_i.function_space is not mpc_i.function_space: 

420 raise ValueError( 

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

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

423 ) 

424 self.u = list(u) 

425 else: 

426 if u is None: 

427 assert isinstance(self._mpc, MultiPointConstraint) 

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

429 else: 

430 assert isinstance(u, _fem.Function) 

431 assert isinstance(self._mpc, MultiPointConstraint) 

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

433 self.u = u 

434 else: 

435 raise ValueError( 

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

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

438 ) 

439 

440 # Create MPC matrix and vector 

441 self._preconditioner = _fem.form(P, jit_options=jit_options, form_compiler_options=form_compiler_options) 

442 

443 if is_nest: 

444 assert isinstance(mpc, Sequence) 

445 assert isinstance(self._L, Sequence) 

446 assert isinstance(self._a, Sequence) 

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

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

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

450 if self._preconditioner is None: 

451 self._P_mat = None 

452 else: 

453 assert isinstance(self._preconditioner, Sequence) 

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

455 else: 

456 assert isinstance(mpc, MultiPointConstraint) 

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

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

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

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

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

462 if self._preconditioner is None: 

463 self._P_mat = None 

464 else: 

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

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

467 

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

469 

470 if is_nest: 

471 assert isinstance(self.u, Sequence) 

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

473 else: 

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

475 comm = self.u.function_space.mesh.comm 

476 

477 self._solver = PETSc.KSP().create(comm) # type: ignore 

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

479 

480 self._petsc_options_prefix = petsc_options_prefix 

481 self.solver.setOptionsPrefix(petsc_options_prefix) 

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

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

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

485 if self.P_mat is not None: 

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

487 

488 # Set options on KSP only 

489 if petsc_options is not None: 

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

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

492 

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

494 opts[k] = v 

495 

496 self.solver.setFromOptions() 

497 

498 # Tidy up global options 

499 for k in petsc_options.keys(): 

500 del opts[k] 

501 opts.prefixPop() 

502 

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

504 """Solve the problem. 

505 

506 Returns: 

507 Function containing the solution""" 

508 

509 # Assemble lhs 

510 self._A.zeroEntries() 

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

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

513 else: 

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

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

516 

517 self._A.assemble() 

518 assert self._A.assembled 

519 

520 # Assemble the preconditioner if provided 

521 if self._P_mat is not None: 

522 self._P_mat.zeroEntries() 

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

524 assert isinstance(self._preconditioner, Sequence) 

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

526 else: 

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

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

529 self._P_mat.assemble() 

530 

531 # Assemble the residual 

532 _zero_vector(self._b) 

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

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

535 else: 

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

537 assert isinstance(self._mpc, MultiPointConstraint) 

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

539 

540 # Lift vector 

541 try: 

542 # Nest and blocked lifting 

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

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

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

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

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

548 except TypeError: 

549 # Single form lifting 

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

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

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

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

554 

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

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

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

558 _fem.petsc.assign(self._x, self.u) 

559 

560 if isinstance(self.u, Sequence): 

561 assert isinstance(self._mpc, Sequence) 

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

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

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

565 else: 

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

567 assert isinstance(self._mpc, MultiPointConstraint) 

568 self._mpc.homogenize(self.u) 

569 self._mpc.backsubstitution(self.u) 

570 

571 return self.u