Coverage for python/src/dolfinx_mpc/assemble_vector.py: 100%

1# Copyright (C) 2020 Jørgen S. Dokken 

2# 

3# This file is part of DOLFINX_MPC 

4# 

5# SPDX-License-Identifier: MIT 

6from __future__ import annotations 

7 

8import contextlib 

9from typing import Iterable, Optional, Sequence, Union 

10 

11from petsc4py import PETSc as _PETSc 

12 

13import dolfinx.cpp as _cpp 

14import dolfinx.fem as _fem 

15import numpy 

16from dolfinx import default_scalar_type 

17from dolfinx.common import Timer 

18from dolfinx.la.petsc import create_vector 

19 

20import dolfinx_mpc.cpp 

21 

22from .multipointconstraint import MultiPointConstraint, _float_classes 

23 

24 

25def apply_lifting( 

26 b: _PETSc.Vec, 

27 form: Union[Iterable[Sequence[_fem.Form]], Iterable[_fem.Form]], # type: ignore 

28 bcs: Union[Sequence[_fem.DirichletBC], Sequence[Sequence[_fem.DirichletBC]]], 

29 constraint: Union[MultiPointConstraint, Sequence[MultiPointConstraint]], 

30 x0: Optional[Sequence[_PETSc.Vec]] = None, 

31 scale: _float_classes = default_scalar_type(1.0), # type: ignore 

32 num_threads: Optional[int] = 1, 

33): # type: ignore 

34 """ 

35 Apply lifting to vector b, i.e. 

36 :math:`b = b - scale \\cdot K^T (A_j (g_j - x0_j))` 

37 

38 Args: 

39 b: PETSc vector to assemble into 

40 form: The linear form 

41 bcs: List of Dirichlet boundary conditions 

42 constraint: The multi point constraint 

43 x0: List of vectors 

44 scale: Scaling for lifting 

45 num_threads: The number of threads to use for certain operations 

46 """ 

47 t = Timer("~MPC: Apply lifting (C++)") 

48 if isinstance(scale, numpy.generic): # nanobind conversion of numpy dtypes to general Python types 

49 scale = scale.item() # type: ignore 

50 

51 if b.getType() == "nest": 

52 try: 

53 bcs = _fem.bcs_by_block(_fem.extract_function_spaces(form, 1), bcs) # type: ignore 

54 except AttributeError: 

55 pass 

56 x0 = [] if x0 is None else x0.getNestSubVecs() # type: ignore 

57 assert isinstance(form, Sequence) and isinstance(constraint, Sequence) 

58 for b_sub, a_sub, mpc_i in zip(b.getNestSubVecs(), form, constraint): 

59 _a = [None if form is None else form._cpp_object for form in a_sub] # type:ignore 

60 _bcs = [[bc._cpp_object for bc in bcs0] for bcs0 in bcs] # type: ignore 

61 dolfinx_mpc.cpp.mpc.apply_lifting(b_sub.array_w, _a, _bcs, x0, scale, mpc_i._cpp_object, num_threads) 

62 else: 

63 with contextlib.ExitStack() as stack: 

64 if x0 is None: 

65 x0 = [] 

66 else: 

67 x0 = [stack.enter_context(x.localForm()) for x in x0] 

68 x0_r = [x.array_r for x in x0] 

69 b_local = stack.enter_context(b.localForm()) 

70 _forms = [f._cpp_object for f in form] # type: ignore 

71 _bcs = [[bc._cpp_object for bc in bcs0] for bcs0 in bcs] # type: ignore 

72 assert isinstance(constraint, MultiPointConstraint) 

73 dolfinx_mpc.cpp.mpc.apply_lifting( 

74 b_local.array_w, _forms, _bcs, x0_r, scale, constraint._cpp_object, num_threads 

75 ) 

76 t.stop() 

77 

78 

79def apply_mpc_lifting( 

80 b: _PETSc.Vec, # type: ignore 

81 form: Sequence[_fem.Form], 

82 constraint: Union[MultiPointConstraint, Sequence[MultiPointConstraint]], 

83 constraint1: Optional[Sequence[MultiPointConstraint]] = None, 

84 scale: _float_classes = default_scalar_type(1.0), # type: ignore 

85 num_threads: Optional[int] = 1, 

86): 

87 """ 

88 Lift the inhomogeneity of a multi point constraint into the vector `b`, i.e. 

89 

90 :math:`b = b - scale \\cdot K^T (A_j g_j)` 

91 

92 where :math:`g` is the constraint offset of the constraint on the trial space. This is 

93 the term arising in :math:`K^T A K x_{red} = K^T (b - A g)` for the affine constraint 

94 :math:`x = K x_{red} + g`, and is a no-op for a homogeneous constraint. 

95 

96 Note: 

97 Only required when solving directly for :math:`x_{red}`. A residual assembled at an 

98 iterate that already satisfies the constraint contains :math:`K^T A g` already, so 

99 the Newton/SNES path must not call this. 

100 

101 Args: 

102 b: PETSc vector to assemble into 

103 form: The bilinear forms, one per block column 

104 constraint: The multi point constraint for the rows of `b` 

105 constraint1: The multi point constraints for the columns, one per block. Defaults 

106 to `constraint`, which is correct for a square problem. 

107 scale: Scaling for lifting 

108 num_threads: The number of threads to use for certain operations 

109 """ 

110 t = Timer("~MPC: Apply MPC lifting (C++)") 

111 if isinstance(scale, numpy.generic): # nanobind conversion of numpy dtypes to general Python types 

112 scale = scale.item() # type: ignore 

113 

114 if b.getType() == "nest": 

115 assert isinstance(form, Sequence) and isinstance(constraint, Sequence) 

116 cols = constraint if constraint1 is None else constraint1 

117 for b_sub, a_sub, mpc_i in zip(b.getNestSubVecs(), form, constraint): 

118 _a = [None if f is None else f._cpp_object for f in a_sub] # type: ignore 

119 _mpc1 = [c._cpp_object for c in cols] # type: ignore 

120 dolfinx_mpc.cpp.mpc.apply_mpc_lifting(b_sub.array_w, _a, scale, mpc_i._cpp_object, _mpc1, num_threads) 

121 else: 

122 assert isinstance(constraint, MultiPointConstraint) 

123 cols = [constraint] if constraint1 is None else constraint1 

124 with b.localForm() as b_local: 

125 _forms = [f._cpp_object for f in form] # type: ignore 

126 _mpc1 = [c._cpp_object for c in cols] # type: ignore 

127 dolfinx_mpc.cpp.mpc.apply_mpc_lifting( 

128 b_local.array_w, _forms, scale, constraint._cpp_object, _mpc1, num_threads 

129 ) 

130 t.stop() 

131 

132 

133def assemble_vector( 

134 form: _fem.Form, 

135 constraint: MultiPointConstraint, 

136 b: Optional[_PETSc.Vec] = None, # type: ignore 

137 num_threads: Optional[int] = 1, 

138) -> _PETSc.Vec: # type: ignore 

139 """ 

140 Assemble a linear form into vector `b` with corresponding multi point constraint 

141 

142 Args: 

143 form: The linear form 

144 constraint: The multi point constraint 

145 b: PETSc vector to assemble 

146 

147 Returns: 

148 The vector with the assembled linear form (`b` if supplied) 

149 """ 

150 

151 if b is None: 

152 b = create_vector([(constraint.function_space.dofmap.index_map, constraint.function_space.dofmap.index_map_bs)]) 

153 t = Timer("~MPC: Assemble vector (C++)") 

154 with b.localForm() as b_local: 

155 b_local.set(0.0) 

156 dolfinx_mpc.cpp.mpc.assemble_vector(b_local.array_w, form._cpp_object, constraint._cpp_object, num_threads) 

157 t.stop() 

158 return b 

159 

160 

161def create_vector_nest(L: Sequence[_fem.Form], constraints: Sequence[MultiPointConstraint]) -> _PETSc.Vec: # type: ignore 

162 """ 

163 Create a PETSc vector of type "nest" appropriate for the provided multi 

164 point constraints 

165 

166 Args: 

167 L: A sequence of linear forms 

168 constraints: An ordered list of multi point constraints 

169 

170 Returns: 

171 PETSc.Vec: A PETSc vector of type "nest" #type: ignore 

172 """ 

173 assert len(constraints) == len(L) 

174 

175 maps = [ 

176 (constraint.function_space.dofmap.index_map._cpp_object, constraint.function_space.dofmap.index_map_bs) 

177 for constraint in constraints 

178 ] 

179 return _cpp.fem.petsc.create_vector_nest(maps) 

180 

181 

182def assemble_vector_nest( 

183 b: _PETSc.Vec, # type: ignore 

184 L: Sequence[_fem.Form], 

185 constraints: Sequence[MultiPointConstraint], 

186 num_threads: Optional[int] = 1, 

187): 

188 """ 

189 Assemble a linear form into a PETSc vector of type "nest" 

190 

191 Args: 

192 b: A PETSc vector of type "nest" 

193 L: A sequence of linear forms 

194 constraints: An ordered list of multi point constraints 

195 """ 

196 assert len(constraints) == len(L) 

197 assert b.getType() == "nest" 

198 

199 b_sub_vecs = b.getNestSubVecs() 

200 for i, L_row in enumerate(L): 

201 assemble_vector(L_row, constraints[i], b=b_sub_vecs[i], num_threads=num_threads)