Coverage for dragonfly_sim/utils/gradient_condensation.py: 73%

1""" 

2Newton problems whose residual reads a gradient field that is eliminated 

3from the unknowns. 

4 

5At p = 0 the RANS model reconstructs a DG0 gradient G of the state by 

6Green-Gauss (core/reconstruction.py), 

7 

8 M_G G = b(U), b(U; tau) = sum over facets of int u_f (x) n : tau ds, 

9 

10with M_G the diagonal (cell-volume) DG0 mass matrix. G is linear in the 

11facet states, so it is eliminated exactly instead of being carried as an 

12unknown: 

13 

14 R(U) = R_U(U, G(U)), dR/dU = A_UU + A_UG M_G^-1 B_GU, 

15 

16with A_UU = dR_U/dU at fixed G, A_UG = dR_U/dG and B_GU = db/dU. The product 

17widens the Jacobian's stencil to the neighbours' neighbours, which is the 

18stencil of the reconstructed scheme. Everything downstream (SNES, the 

19positivity limiter, the time integrator) still sees a single-field problem 

20in U. 

21 

22jacobian_mode selects what Newton's linear solve uses (configure_snes wires 

23it into an existing SNES): 

24 

25 exact operator = preconditioning matrix = A_UU + A_UG M^-1 B_GU 

26 compact_pc operator exact, preconditioning matrix A_UU: the compact, 

27 nearest-neighbour part (gradients frozen), as finite-volume 

28 codes precondition with their first-order Jacobian. The 

29 default: on the coarse SA NACA0012 C-grid it takes the same 

30 PTC path as MUMPS, 5x faster, and scales to the 181k-cell STW 

31 wing where every exact factorization runs out of memory. 

32 compact operator = preconditioning matrix = A_UU: approximate Newton 

33 (defect correction); converges linearly 

34""" 

35import ufl 

36from petsc4py import PETSc 

37from dolfinx.fem.forms import form as _create_form 

38from dolfinx.fem.petsc import assemble_matrix, create_matrix 

39 

40from dragonfly_sim.utils.Nonlinear_utils import NonlinearProblem_mod 

41 

42 

43class CondensedJacobian: 

44 """ 

45 A = A_UU + A_UG M^-1 B_GU for a residual form F(U, G), with `b_G` the UFL 

46 Green-Gauss right-hand side (test function on G's space) and `inv_mass` 

47 M_G^-1 as a Vec in G's layout. The matrices are plain aij: MatMatMult 

48 does not take baij here. The product's sparsity is fixed by the 

49 operands', so one symbolic product sizes A for good. 

50 """ 

51 

52 def __init__(self, F, u, G, b_G, inv_mass, jit_options={}, A_UU=None, a_UU=None): 

53 self._inv_mass = inv_mass 

54 self._a = a_UU if a_UU is not None else _create_form(ufl.derivative(F, u), jit_options=jit_options) 

55 self._a_UG = _create_form(ufl.derivative(F, G, ufl.TrialFunction(G.function_space)), 

56 jit_options=jit_options) 

57 self._a_GU = _create_form(ufl.derivative(b_G, u, ufl.TrialFunction(u.function_space)), 

58 jit_options=jit_options) 

59 self.A_UU = A_UU if A_UU is not None else create_matrix(self._a) 

60 self.A_UG = create_matrix(self._a_UG) 

61 self.B_GU = create_matrix(self._a_GU) 

62 self._assemble_blocks() 

63 self._C = self.A_UG.matMult(self.B_GU) 

64 A = self.A_UU.duplicate(copy=True) 

65 A.axpy(1.0, self._C, structure=PETSc.Mat.Structure.DIFFERENT_NONZERO_PATTERN) 

66 A.zeroEntries() 

67 self.A = A 

68 

69 def _assemble_blocks(self): 

70 for M, a in ((self.A_UU, self._a), (self.A_UG, self._a_UG), (self.B_GU, self._a_GU)): 

71 M.zeroEntries() 

72 assemble_matrix(M, a) 

73 M.assemble() 

74 self.B_GU.diagonalScale(L=self._inv_mass) 

75 

76 def assemble(self): 

77 """Assemble A_UU, A_UG, M^-1 B_GU and A at the current state; G must 

78 be up to date.""" 

79 self._assemble_blocks() 

80 self.A_UG.matMult(self.B_GU, result=self._C) 

81 self.A.zeroEntries() 

82 self.A.axpy(1.0, self.A_UU, structure=PETSc.Mat.Structure.SUBSET_NONZERO_PATTERN) 

83 self.A.axpy(1.0, self._C, structure=PETSc.Mat.Structure.SUBSET_NONZERO_PATTERN) 

84 self.A.assemble() 

85 return self.A 

86 

87 

88class TransposedCondensedOperator: 

89 """ 

90 Python-Mat context of the transposed condensed Jacobian without forming it: 

91 

92 A^T x = (A_UU + A_UG B)^T x = P x + B^T (A_UG^T x), P = A_UU^T 

93 

94 with B = M_G^-1 dG/dU (any Mat with mult / multTranspose, e.g. the 

95 structured GreenGaussStateOperator). Neither A nor the product A_UG B 

96 is ever assembled. It owns P (also the preconditioning matrix), A_UG and 

97 B (released through `release_B`, a callable) and frees them in destroy(). 

98 """ 

99 

100 def __init__(self, P, A_UG, B, release_B): 

101 self.P, self.A_UG, self.B, self._release_B = P, A_UG, B, release_B 

102 self._w = A_UG.createVecRight() 

103 

104 def mult(self, mat, x, y): 

105 self.A_UG.multTranspose(x, self._w) 

106 self.B.multTranspose(self._w, y) 

107 self.P.multAdd(x, y, y) 

108 

109 def destroy(self, mat=None): 

110 # called by PETSc when the python Mat is destroyed 

111 if self.P is None: 

112 return 

113 self._release_B(self.B) 

114 for obj in (self._w, self.A_UG, self.P): 

115 obj.destroy() 

116 self.P = self.A_UG = self.B = None 

117 

118 def create_mat(self, comm): 

119 M = PETSc.Mat().createPython(self.P.getSizes(), comm=comm) 

120 M.setPythonContext(self) 

121 M.setUp() 

122 return M 

123 

124 

125class CondensedGradientProblem(NonlinearProblem_mod): 

126 """ 

127 NonlinearProblem_mod for R_U(U, G(U)). 

128 

129 `G` is the gradient Function the residual reads, `b_G` the UFL form of 

130 the Green-Gauss right-hand side, `update_G()` recomputes G from the 

131 current u in place and `inv_mass` is M_G^-1 as a Vec in G's layout. 

132 """ 

133 

134 MODES = ("exact", "compact_pc", "compact") 

135 

136 def __init__(self, F, J, u, G, b_G, update_G, inv_mass, jit_options={}): 

137 super().__init__(F, J, u, jit_options=jit_options, mat_type=None) 

138 self.G = G 

139 self.update_G = update_G 

140 self.update_G() 

141 self._condensed = CondensedJacobian(F, u, G, b_G, inv_mass, jit_options=jit_options, 

142 A_UU=self._A, a_UU=self._a) 

143 self._A_UU = self._condensed.A_UU 

144 self._A = self._condensed.A 

145 self.jacobian_mode = "exact" 

146 

147 def configure_snes(self, snes, mode): 

148 """Re-wire `snes`'s Jacobian for jacobian_mode `mode` (see the module 

149 docstring). Call after the solver is built.""" 

150 if mode not in self.MODES: 

151 raise ValueError("jacobian_mode must be one of {}".format(self.MODES)) 

152 self.jacobian_mode = mode 

153 if mode == "exact": 

154 return 

155 

156 def sync(x): 

157 x.copy(self.u.x.petsc_vec) 

158 self.u.x.petsc_vec.ghostUpdate(addv=PETSc.InsertMode.INSERT, 

159 mode=PETSc.ScatterMode.FORWARD) 

160 

161 if mode == "compact_pc": 

162 def jac(snes_, x, A, P): 

163 sync(x) 

164 self.J(x, self._A) # assembles A_UU on the way 

165 snes.setJacobian(jac, self._A, self._A_UU) 

166 return 

167 

168 def compact(snes_, x, A, P): 

169 sync(x) 

170 self.update_G() 

171 self._A_UU.zeroEntries() 

172 assemble_matrix(self._A_UU, self._a) 

173 self._A_UU.assemble() 

174 snes.setJacobian(compact, self._A_UU, self._A_UU) 

175 

176 def F(self, x, b): 

177 self.update_G() 

178 super().F(x, b) 

179 

180 def J(self, x, A): 

181 self.update_G() 

182 self._condensed.assemble() 

183 if A.handle != self._A.handle: 

184 A.assemble() 

185 self._A.copy(A)