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

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

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

3# 

4# This file is part of DOLFINx MPC 

5# 

6# SPDX-License-Identifier: MIT 

7from __future__ import annotations 

8 

9import typing 

10 

11from petsc4py import PETSc 

12 

13import dolfinx.fem.petsc 

14import ufl 

15from dolfinx import cpp as _cpp 

16from dolfinx import fem as _fem 

17from dolfinx.la.petsc import create_vector 

18 

19from .assemble_matrix import assemble_matrix, create_sparsity_pattern 

20from .assemble_vector import apply_lifting, assemble_vector 

21from .multipointconstraint import MultiPointConstraint 

22 

23 

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

25 """ 

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

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

28 

29 Args: 

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

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

32 mpc: The multi point constraint. 

33 bcs: A list of Dirichlet boundary conditions. 

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

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

36 

37 .. highlight:: python 

38 .. code-block:: python 

39 

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

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

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

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

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

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

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

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

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

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

50 Examples: 

51 Example usage: 

52 

53 .. highlight:: python 

54 .. code-block:: python 

55 

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

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

58 

59 """ 

60 

61 u: _fem.Function 

62 _a: _fem.Form 

63 _L: _fem.Form 

64 _mpc: MultiPointConstraint 

65 _A: PETSc.Mat # type: ignore 

66 _b: PETSc.Vec # type: ignore 

67 _solver: PETSc.KSP # type: ignore 

68 _x: PETSc.Vec # type: ignore 

69 bcs: typing.List[_fem.DirichletBC] 

70 __slots__ = tuple(__annotations__) 

71 

72 def __init__( 

73 self, 

74 a: ufl.Form, 

75 L: ufl.Form, 

76 mpc: MultiPointConstraint, 

77 bcs: typing.Optional[typing.List[_fem.DirichletBC]] = None, 

78 u: typing.Optional[_fem.Function] = None, 

79 petsc_options: typing.Optional[dict] = None, 

80 form_compiler_options: typing.Optional[dict] = None, 

81 jit_options: typing.Optional[dict] = None, 

82 ): 

83 # Compile forms 

84 form_compiler_options = {} if form_compiler_options is None else form_compiler_options 

85 jit_options = {} if jit_options is None else jit_options 

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

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

88 

89 if not mpc.finalized: 

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

91 self._mpc = mpc 

92 # Create function containing solution vector 

93 if u is None: 

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

95 else: 

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

97 self.u = u 

98 else: 

99 raise ValueError( 

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

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

102 ) 

103 self._x = self.u.x.petsc_vec 

104 

105 # Create MPC matrix 

106 pattern = create_sparsity_pattern(self._a, self._mpc) 

107 pattern.finalize() 

108 self._A = _cpp.la.petsc.create_matrix(self._mpc.function_space.mesh.comm, pattern) 

109 

110 self._b = create_vector(self._mpc.function_space.dofmap.index_map, self._mpc.function_space.dofmap.index_map_bs) 

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

112 

113 self._solver = PETSc.KSP().create(self.u.function_space.mesh.comm) # type: ignore 

114 self._solver.setOperators(self._A) 

115 

116 # Give PETSc solver options a unique prefix 

117 solver_prefix = "dolfinx_mpc_solve_{}".format(id(self)) 

118 self._solver.setOptionsPrefix(solver_prefix) 

119 

120 # Set PETSc options 

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

122 opts.prefixPush(solver_prefix) 

123 if petsc_options is not None: 

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

125 opts[k] = v 

126 opts.prefixPop() 

127 self._solver.setFromOptions() 

128 

129 def solve(self) -> _fem.Function: 

130 """Solve the problem. 

131 

132 Returns: 

133 Function containing the solution""" 

134 

135 # Assemble lhs 

136 self._A.zeroEntries() 

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

138 self._A.assemble() 

139 assert self._A.assembled 

140 

141 # Assemble rhs 

142 with self._b.localForm() as b_loc: 

143 b_loc.set(0) 

144 assemble_vector(self._L, self._mpc, b=self._b) 

145 

146 # Apply boundary conditions to the rhs 

147 apply_lifting(self._b, [self._a], [self.bcs], self._mpc) 

148 self._b.ghostUpdate(addv=PETSc.InsertMode.ADD, mode=PETSc.ScatterMode.REVERSE) # type: ignore 

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

150 

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

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

153 self.u.x.scatter_forward() 

154 self._mpc.backsubstitution(self.u) 

155 

156 return self.u