Coverage for dragonfly_sim/utils/Nonlinear_utils.py: 94%

1from __future__ import annotations 

2 

3import typing 

4 

5import numpy as np 

6 

7import ufl 

8from mpi4py import MPI 

9 

10from dolfinx.fem.bcs import DirichletBC 

11from dolfinx.fem.forms import Form 

12from dolfinx.fem.forms import form as _create_form 

13from dolfinx.fem.forms import extract_function_spaces 

14from dolfinx.fem.petsc import (assemble_vector, assemble_matrix, create_vector, create_matrix, set_bc, apply_lifting) 

15from dolfinx.fem.function import Function as _Function 

16 

17from petsc4py import PETSc 

18 

19 

20class NonlinearProblem_mod: 

21 """Nonlinear problem class for solving the non-linear problem 

22 :math:`F(u, v) = 0 \\ \\forall v \\in V` using PETSc as the linear algebra 

23 backend. 

24 

25 Copied from DOLFINx's built-in `NonlinearProblem` and extended with a 

26 caller-selectable Jacobian matrix type (`mat_type`, e.g. "baij" for 

27 point-block ILU). 

28 """ 

29 

30 def __init__(self, F: ufl.form.Form, J: ufl.form.Form, u: _Function, bcs: typing.List[DirichletBC] = [], 

31 form_compiler_options={}, jit_options={}, mat_type: str = None): 

32 """Initialize solver for solving a non-linear problem using Newton's method, :math:`dF/du(u) du = -F(u)`. 

33 

34 Args: 

35 F: The PDE residual F(u, v) 

36 u: The unknown 

37 bcs: List of Dirichlet boundary conditions 

38 J: UFL representation of the Jacobian (Optional) 

39 form_compiler_options: Options used in FFCx 

40 compilation of this form. Run ``ffcx --help`` at the 

41 commandline to see all available options. 

42 jit_options: Options used in CFFI JIT compilation of C 

43 code generated by FFCx. See ``python/dolfinx/jit.py`` 

44 for all available options. Takes priority over all 

45 other option values. 

46 

47 Example:: 

48 

49 problem = LinearProblem(F, u, [bc0, bc1]) 

50 

51 """ 

52 # TODO: Figure out why form_compiler_options and jit_options below are unexpected keywords.  

53 # According to the DOLFINx 0.6.0 documentation this should be correct 

54 

55 self._L = _create_form(F, form_compiler_options=form_compiler_options, 

56 jit_options=jit_options) 

57 

58 self.u = u 

59 

60 # Create the Jacobian matrix, dF/du 

61 # if J is None: 

62 # V = u.function_space 

63 # du = ufl.TrialFunction(V) 

64 # J = ufl.derivative(F, u, du) 

65 

66 self._a = _create_form(J, form_compiler_options=form_compiler_options, 

67 jit_options=jit_options) 

68 self.bcs = bcs 

69 

70 # PETSc matrix type for the Jacobian (e.g. "baij" for point-block 

71 # ILU). None -> PETSc default (aij). self._A below is the one 

72 # full-space Jacobian: SNESNewtonSolver uses it directly as its 

73 # operator. 

74 self.mat_type = mat_type 

75 

76 # save some vectors and matrices with appropriate sparsity structures for use later 

77 self._A = create_matrix(self.a, kind=self.mat_type) 

78 self._b = create_vector(extract_function_spaces(self.L)) 

79 

80 @property 

81 def L(self) -> Form: 

82 """Compiled linear form (the residual form)""" 

83 return self._L 

84 

85 @property 

86 def a(self) -> Form: 

87 """Compiled bilinear form (the Jacobian form)""" 

88 return self._a 

89 

90 def form(self, x: PETSc.Vec): 

91 """ 

92 This function is called before the residual or Jacobian is 

93 computed. This is usually used to update ghost values. 

94 

95 Args: 

96 x: The vector containing the latest solution 

97 

98 """ 

99 # print("Call form in NonlinearProblem_mod") 

100 # PETSc.Sys.Print("u pre norm: {}".format(self.u.x.petsc_vec.norm())) 

101 x.ghostUpdate(addv=PETSc.InsertMode.INSERT, mode=PETSc.ScatterMode.FORWARD) 

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

103 # PETSc.Sys.Print("u post norm: {}".format(self.u.x.petsc_vec.norm())) 

104 

105 def F(self, x: PETSc.Vec, b: PETSc.Vec): 

106 """ 

107 Assemble the residual F into the vector b. 

108 

109 Args: 

110 x: The vector containing the latest solution 

111 b: Vector to assemble the residual into 

112 """ 

113 # PETSc.Sys.Print("Start construct b") 

114 # # Reset the residual vector 

115 # print("Call F in NonlinearProblem_mod") 

116 with self._b.localForm() as b_local: 

117 b_local.set(0.0) 

118 

119 assemble_vector(self._b, self._L) 

120 

121 apply_lifting(self._b, [self._a], bcs=[self.bcs], x0=[self.u.x.petsc_vec], alpha=-1.0) 

122 self._b.ghostUpdate(addv=PETSc.InsertMode.ADD, mode=PETSc.ScatterMode.REVERSE) 

123 set_bc(self._b, self.bcs, self.u.x.petsc_vec, -1.0) 

124 self._b.ghostUpdate(addv=PETSc.InsertMode.INSERT_VALUES, mode=PETSc.ScatterMode.FORWARD) 

125 

126 self._b.copy(b) # copy content of self._b into vector b 

127 

128 def J(self, x: PETSc.Vec, A: PETSc.Mat): 

129 """ 

130 Assemble the Jacobian matrix. 

131 

132 Args: 

133 x: The vector containing the latest solution 

134 A: The matrix that will hold the Jacobian matrix 

135 """ 

136 # print("Call J in NonlinearProblem_mod") 

137 # print("A shape in J: {}".format(A.getSize())) 

138 

139 self._A.zeroEntries() 

140 assemble_matrix(self._A, self._a, self.bcs) 

141 self._A.assemble() 

142 

143 # SNESNewtonSolver hands in self._A itself; any other matrix gets a 

144 # copy. 

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

146 A.assemble() 

147 self._A.copy(A) 

148 

149 

150class SNESNewtonSolver: 

151 """ 

152 Full-space Newton on PETSc SNES, driving a NonlinearProblem_mod. 

153 

154 Replaces NewtonSolver_mod, which subclassed dolfinx's C++ NewtonSolver -- 

155 deprecated since dolfinx 0.10 and slated for removal. The attribute surface 

156 is the old one (rtol, atol, max_it, convergence_criterion, 

157 relaxation_parameter, error_on_nonconvergence, krylov_solver, A, b, 

158 solve(u) -> (n, converged)), and the defaults below reproduce the old 

159 solver step for step rather than adopting SNES's own behaviour: 

160 

161 * newtonls with the BASIC line search, so every step is taken in full, 

162 as NewtonSolver did. SNES's default (bt) would evaluate the residual at 

163 extra trial points, and could shorten a step the old solver took in 

164 full. 

165 * dolfinx's convergence test, not SNES's: "incremental" judges ||dx_k|| 

166 against ||dx_1|| and never stops on the first step; "residual" judges 

167 ||F|| against ||F_0||. ||dx|| is the RAW Newton step, recorded before 

168 the step limiter rescales it. 

169 * a KSP that fails to converge does not stop the solve. NewtonSolver took 

170 the step regardless; SNES by default aborts with DIVERGED_LINEAR_SOLVE 

171 before taking any step (checked on this build). 

172 * the step limiter (set_step_limiter) replaces set_update. It runs as the 

173 line search's pre-check and rescales dx IN PLACE; the basic line search 

174 then lands on x - dx, the same step the old x.axpy(-theta, dx) took. 

175 

176 Checked against NewtonSolver_mod on a nonlinear Poisson problem with 

177 Dirichlet BCs (1 and 3 ranks; both criteria, relaxation, a step limiter, 

178 max_it=1, a KSP capped short of convergence): iteration counts and 

179 converged flags identical, iterates equal to round-off (<= 4.4e-16). Not 

180 bit-identical: the Krylov solve itself differs in the last bit, from 

181 identical residual, settings and iteration count. 

182 

183 The Jacobian is assembled straight into problem._A, which is also the 

184 KSP's operator. NewtonSolver_mod kept a second full-size matrix and copied 

185 problem._A into it every iteration. 

186 

187 The SNES iterate is its own vector, not u's: the forms read u, and SNES 

188 evaluates the residual at a line-search work vector, so F and J copy their 

189 argument into u (owned entries, then a forward ghost update) before 

190 assembling. Handing SNES u's vector instead would let that copy overwrite 

191 the iterate mid-line-search under any line search other than basic. 

192 

193 SNES options use the prefix `nls_solve_`, which the KSP inherits -- the 

194 prefix the C++ NewtonSolver gave its KSP, so every -nls_solve_* option 

195 set by apply_krylov_solver_settings (or on the command line) still lands. 

196 """ 

197 def __init__(self, comm: MPI.Intracomm, problem: NonlinearProblem_mod, 

198 options_prefix: str = "nls_solve_"): 

199 self.comm = comm 

200 self._problem = problem 

201 

202 # dolfinx NewtonSolver's defaults, so a caller that sets only some of 

203 # these gets what it got before. 

204 self.rtol = 1e-9 

205 self.atol = 1e-10 

206 self.max_it = 50 

207 self.convergence_criterion = "residual" 

208 self.relaxation_parameter = 1.0 

209 self.error_on_nonconvergence = True 

210 

211 self._step_limiter = None 

212 self._dx_norm = None 

213 self._residual0 = None 

214 

215 self._x = problem.u.x.petsc_vec.duplicate() 

216 self._f = problem._b.duplicate() 

217 

218 snes = PETSc.SNES().create(comm) 

219 snes.setOptionsPrefix(options_prefix) 

220 snes.setType(PETSc.SNES.Type.NEWTONLS) 

221 snes.getLineSearch().setType(PETSc.SNESLineSearch.Type.BASIC) 

222 snes.setFunction(self._residual, self._f) 

223 snes.setJacobian(self._jacobian, problem._A, problem._A) 

224 snes.setLineSearchPreCheck(self._precheck) 

225 snes.setConvergenceTest(self._converged) 

226 snes.setMaxLinearSolveFailures(int(np.iinfo(PETSc.IntType).max)) 

227 # Before the caller's KSP configuration, so that configuration is the 

228 # last word; -nls_solve_snes_* on the command line can still override 

229 # the choices above. 

230 snes.setFromOptions() 

231 self._snes = snes 

232 

233 # --- callbacks ------------------------------------------------------ 

234 def _sync(self, x): 

235 u = self._problem.u.x.petsc_vec 

236 x.copy(u) 

237 u.ghostUpdate(addv=PETSc.InsertMode.INSERT, mode=PETSc.ScatterMode.FORWARD) 

238 return u 

239 

240 def _residual(self, snes, x, b): 

241 self._problem.F(self._sync(x), b) 

242 

243 def _jacobian(self, snes, x, A, P): 

244 self._problem.J(self._sync(x), A) 

245 

246 def _precheck(self, x, dx): 

247 # Before the limiter touches dx: the incremental test is on the raw 

248 # Newton step, as it was in NewtonSolver. 

249 self._dx_norm = dx.norm() 

250 if self._step_limiter is not None: 

251 self._step_limiter(self, x, dx) 

252 elif self.relaxation_parameter != 1.0: 

253 dx.scale(self.relaxation_parameter) 

254 

255 def _converged(self, snes, it, norms): 

256 _, _, fnorm = norms 

257 if self.convergence_criterion == "residual": 

258 residual = fnorm 

259 if it == 0: 

260 self._residual0 = residual 

261 elif self.convergence_criterion == "incremental": 

262 # NewtonSolver took at least one step under this criterion, and 

263 # measured the later steps against the first. 

264 if it == 0: 

265 return PETSc.SNES.ConvergedReason.ITERATING 

266 residual = self._dx_norm 

267 if it == 1: 

268 self._residual0 = residual 

269 return PETSc.SNES.ConvergedReason.ITERATING 

270 else: 

271 raise ValueError("convergence_criterion must be 'residual' or " 

272 "'incremental', got {!r}".format( 

273 self.convergence_criterion)) 

274 if residual < self.atol: 

275 return PETSc.SNES.ConvergedReason.CONVERGED_FNORM_ABS 

276 if self._residual0 > 0 and residual / self._residual0 < self.rtol: 

277 return PETSc.SNES.ConvergedReason.CONVERGED_FNORM_RELATIVE 

278 # max_it is enforced by SNES itself, as DIVERGED_MAX_IT. 

279 return PETSc.SNES.ConvergedReason.ITERATING 

280 

281 # --- public API ----------------------------------------------------- 

282 def set_step_limiter(self, limiter): 

283 """ 

284 Install `limiter(solver, x, dx)`, called once per Newton step with the 

285 current iterate x and the Newton step dx (the solve of J dx = F). 

286 It must rescale dx IN PLACE to the step it wants taken: the new 

287 iterate is x - dx. It must not modify x. Replaces relaxation_parameter 

288 when set. 

289 """ 

290 self._step_limiter = limiter 

291 

292 def solve(self, u: _Function): 

293 """Solve the non-linear problem into function u. Returns the number 

294 of Newton iterations and whether the solver converged.""" 

295 self._snes.setTolerances(max_it=int(self.max_it)) 

296 u.x.petsc_vec.copy(self._x) 

297 self._snes.solve(None, self._x) 

298 self._x.copy(u.x.petsc_vec) 

299 u.x.scatter_forward() 

300 n = self._snes.getIterationNumber() 

301 converged = self._snes.getConvergedReason() > 0 

302 if not converged and self.error_on_nonconvergence: 

303 raise RuntimeError( 

304 "Newton solver did not converge ({} after {} iterations)".format( 

305 self._snes.getConvergedReason(), n)) 

306 return n, converged 

307 

308 @property 

309 def krylov_solver(self) -> PETSc.KSP: 

310 """The Newton step's linear solver.""" 

311 return self._snes.getKSP() 

312 

313 @property 

314 def snes(self) -> PETSc.SNES: 

315 return self._snes 

316 

317 @property 

318 def A(self) -> PETSc.Mat: 

319 """Jacobian matrix (the problem's own, not a copy).""" 

320 return self._problem._A 

321 

322 @property 

323 def b(self) -> PETSc.Vec: 

324 """Residual vector.""" 

325 return self._f 

326 

327 def destroy(self): 

328 """Free the SNES and this solver's two vectors. The Jacobian belongs to 

329 the problem and is left alone.""" 

330 self._snes.destroy() 

331 self._x.destroy() 

332 self._f.destroy()