Coverage for dragonfly_sim/core/time_integration.py: 92%
277 statements
« prev ^ index » next coverage.py v7.16.2, created at 2026-10-17 03:37 +0000
1"""
2Time integration: pseudo-transient continuation (PTC) for steady solves and
3BDF1/BDF2 for unsteady ones, on one residual.
5 R*(U) = R(U) + inv_dt (a0 U + a1 U^n + a2 U^{n-1}) physical BDF term
6 + (w_K / CFL) (U - U_old) pseudo-time term
8R is the model's steady residual (model.F, unchanged, which is what the
9adjoint differentiates). w_K = lambda_K / h_K (+ the model's
10time_scale_terms, e.g. viscous 2 nu / h^2) is the cell's inverse local time
11step at CFL 1, evaluated at U_old. Both terms are lumped DG0 mass terms and
12both live only in TimeTerm.augment(), so they never enter model.F.
14Modes (all one step loop, TimeIntegrator._newton_step):
16 steady PTC physical term off (inv_dt = 0); one Newton step per pseudo
17 step, U_old <- U before each, the CFL law (CFLController)
18 adapts the step from the steady |R|. At the start of a
19 pseudo step U = U_old, so the residual assembled there is the
20 steady one.
21 unsteady physical term on, pseudo term off (1/CFL = 0): Newton per
22 physical step to the residual criterion; Delta t ramp and
23 halving on failure (StepSizeController).
24 dual time physical term on, and the pseudo term on inside every
25 physical step, with the CFL law restarted each step: PTC on
26 R* instead of R.
28Switching a term off sets its Constant to exactly zero, so the compiled
29forms stay valid across modes; a model with neither term enabled never
30builds them at all (the plain Newton solve is unchanged).
31"""
32from dataclasses import dataclass, field
34import numpy as np
35from mpi4py import MPI
36from petsc4py import PETSc
37import ufl
38import dolfinx
41# (a0, a1, a2) in (a0 U^{n+1} + a1 U^n + a2 U^{n-1}) / dt at a constant dt.
42BDF_COEFFS = {1: (1.0, -1.0, 0.0), 2: (1.5, -2.0, 0.5)}
45def bdf_coefficients(order, ratio=None):
46 """
47 (a0, a1, a2) for (a0 U^{n+1} + a1 U^n + a2 U^{n-1}) / dt_n. Variable-step
48 BDF2 with ratio w = dt_n / dt_{n-1}:
50 a0 = (1 + 2w)/(1 + w), a1 = -(1 + w), a2 = w^2/(1 + w),
52 which is BDF_COEFFS[2] at w = 1. BDF1 ignores the ratio.
53 """
54 if order == 1:
55 return BDF_COEFFS[1]
56 w = 1.0 if ratio is None else float(ratio)
57 return ((1.0 + 2.0 * w) / (1.0 + w), -(1.0 + w), w * w / (1.0 + w))
60# ----------------------------------------------------------------------
61# CFL law
62# ----------------------------------------------------------------------
63@dataclass
64class CFLController:
65 """
66 Switched-evolution-relaxation CFL law with the safeguards the p = 0 RANS
67 runs needed. Pure Python: update() takes the step's outcome and returns
68 the decision, so the law is unit-tested without a solver.
70 After a step from residual r_prev to r, with a converged linear solve
71 (ksp_ok) and positivity step length theta:
73 reject r non-finite or r > reject_factor r_prev: undo the step,
74 CFL *= cut (stop when CFL < 1e-6 cfl0)
75 limited theta < theta_cut: CFL *= cut, down to cfl_min_limited.
76 Without it a limited step leaves |R| flat, the band below
77 holds CFL, and the limiter goes on to zero every step (p=0
78 MUSCL from freestream froze at M 0.9 and 0.95; one
79 leading-edge cell at p = 1e-4 blocked the step at CFL 1 but
80 not at CFL 0.1, hence a floor below cfl_min).
81 no Krylov an inexact direction never grows CFL
82 progress r < (1 - 1e-3) r_prev: grow by max(r_prev/r, min_growth),
83 capped at max_growth
84 rise r > 1.05 r_prev: CFL *= r_prev/r, clipped to
85 [rise_factor_min, rise_factor_max]; the defaults, 0.1 and
86 1/max_growth = 0.5, cut it by at least half
87 band otherwise (|R| flat within -0.1%..+5%): grow by
88 band_growth on a clean step (theta = 1), else hold. Holding
89 kept the STW L3 wing at CFL 3.15 for 118 of 152 steps.
91 CFL stays within [floor, cfl_max]; the floor is cfl_min_limited after a
92 limited step and min(cfl_min, CFL) otherwise, so a CFL below cfl_min
93 regrows by the usual rule instead of jumping back to cfl_min.
95 exhausted: CFL fell below 1e-6 cfl0 after rejections, or the last
96 max_frozen_steps steps all had theta = 0. A step with theta = 0 leaves the
97 state unchanged, so once a cell sits at the positivity floor every later
98 step repeats it; without this stop PTC spent its whole step budget frozen
99 (OAT15A, Cadence grid level 1: ~950 identical steps at CFL cfl_min_limited).
100 """
101 cfl0: float = 5.0
102 cfl_max: float = 1e12
103 max_growth: float = 2.0
104 min_growth: float = 1.5
105 cut: float = 0.1
106 reject_factor: float = 10.0
107 cfl_min: float = 1.0
108 theta_cut: float = 0.1
109 cfl_min_limited: float = 1e-3
110 band_growth: float = 1.2
111 rise_factor_min: float = 0.1
112 rise_factor_max: float = 0.5
113 max_frozen_steps: int = 5
114 cfl: float = field(init=False)
115 frozen_steps: int = field(init=False, default=0)
117 def __post_init__(self):
118 self.reset()
120 def reset(self):
121 self.cfl = float(self.cfl0)
122 self.frozen_steps = 0
124 def update(self, r_prev, r, ksp_ok=True, theta=1.0):
125 """Returns (accept, limited). Updates self.cfl."""
126 self.frozen_steps = self.frozen_steps + 1 if theta <= 0.0 else 0
127 limited = theta < self.theta_cut
128 if not np.isfinite(r) or r > self.reject_factor * r_prev:
129 self.cfl *= self.cut
130 return False, limited
131 growth = min(max(r_prev / r, self.cut), self.max_growth) if r > 0 else self.max_growth
132 if limited:
133 growth = self.cut
134 elif not ksp_ok:
135 growth = min(growth, 1.0)
136 elif r < (1.0 - 1e-3) * r_prev:
137 growth = max(growth, self.min_growth)
138 elif r > 1.05 * r_prev:
139 growth = min(max(r_prev / r, self.rise_factor_min), self.rise_factor_max)
140 elif theta >= 1.0:
141 growth = self.band_growth
142 else:
143 growth = 1.0
144 floor = self.cfl_min_limited if limited else min(self.cfl_min, self.cfl)
145 self.cfl = max(min(self.cfl * growth, self.cfl_max), floor)
146 return True, limited
148 @property
149 def exhausted(self):
150 return self.cfl < 1e-6 * self.cfl0 or self.frozen_steps >= self.max_frozen_steps
153@dataclass
154class StepSizeController:
155 """Physical time step: ramp from dt_start by `growth` per step up to dt,
156 and halve on a failed step, up to max_retries times."""
157 dt: float
158 dt_start: float = None
159 growth: float = 1.5
160 max_retries: int = 8
162 def first(self, dt_prev):
163 if dt_prev is None:
164 return self.dt if self.dt_start is None else self.dt_start
165 return min(dt_prev * self.growth, self.dt)
168# ----------------------------------------------------------------------
169# Forms
170# ----------------------------------------------------------------------
171class TimeTerm:
172 """
173 The physical and pseudo-time terms of R* as UFL forms over Constants
174 and Functions owned here (see the module docstring).
175 """
177 def __init__(self, model, physical=False, pseudo=False):
178 self.model = model
179 msh = model.mesh_obj.mesh
180 V = model.functionspaces["V"]
181 self.physical, self.pseudo = bool(physical), bool(pseudo)
182 self.inv_dt = dolfinx.fem.Constant(msh, 0.0)
183 self.a = tuple(dolfinx.fem.Constant(msh, c) for c in BDF_COEFFS[1])
184 self.inv_cfl = dolfinx.fem.Constant(msh, 0.0)
185 self.U_n = dolfinx.fem.Function(V, name="u_n")
186 self.U_nm1 = dolfinx.fem.Function(V, name="u_nm1")
187 self.U_old = dolfinx.fem.Function(V, name="u_old")
188 for f in (self.U_n, self.U_nm1, self.U_old):
189 f.x.array[:] = model.u_vec.x.array
191 def local_inverse_time_scale(self, U):
192 """w_K = lambda_max / h_min plus the model's time_scale_terms."""
193 m = self.model
194 w = m._lambda_max_expr(U) / m.mesh_obj.h_min
195 for term in m.time_scale_terms:
196 w = w + term(U)
197 return w
199 def augment(self, F):
200 m = self.model
201 u, v, dx = m.u_vec, m.v_vec, m.forms.dx
202 if self.pseudo:
203 w = self.local_inverse_time_scale(self.U_old)
204 F = F + w * self.inv_cfl * ufl.inner(u - self.U_old, v) * dx
205 if self.physical:
206 a0, a1, a2 = self.a
207 F = F + self.inv_dt * ufl.inner(a0 * u + a1 * self.U_n + a2 * self.U_nm1, v) * dx
208 return F
210 def set_bdf(self, order, dt, dt_prev=None):
211 ratio = None if (order == 1 or dt_prev is None) else dt / dt_prev
212 for c, val in zip(self.a, bdf_coefficients(order, ratio)):
213 c.value = val
214 self.inv_dt.value = 1.0 / dt
215 return ratio
217 def physical_off(self):
218 self.inv_dt.value = 0.0
220 def set_cfl(self, cfl):
221 self.inv_cfl.value = 0.0 if cfl is None or np.isinf(cfl) else 1.0 / cfl
224# ----------------------------------------------------------------------
225# Integrator
226# ----------------------------------------------------------------------
227@dataclass
228class StepRecord:
229 accepted: bool
230 newton_its: int = 0
231 ksp_its: int = 0
232 res_initial: float = np.nan
233 res_final: float = np.nan
234 dt: float = np.nan
235 order: int = 0
236 retries: int = 0
239class TimeIntegrator:
240 """
241 Drives a CompressibleEulerModel (or a model hosted on one) through
242 steady PTC, unsteady BDF or dual-time steps. Built by the model
243 (set_up_solver / enable_time_stepping); `term` is its TimeTerm.
245 Time-dependent boundary conditions (core/farfield.py:TransverseFarfield)
246 register through `boundary_updaters`: objects with prepare(order) before a
247 physical step's Newton solve and commit() after it is accepted.
248 """
250 def __init__(self, model, term):
251 self.model = model
252 self.term = term
253 self.boundary_updaters = []
254 self.time = 0.0
255 self.step_index = 0
256 self.dt_prev = None
257 self.verbose = True
258 # (state, CFL) at the end of the last converged steady solve
259 self._steady_end = None
261 # -- helpers ---------------------------------------------------------
262 def _print(self, *args):
263 if self.verbose:
264 PETSc.Sys.Print(*args)
266 def _copy(self, dst, src):
267 dst.x.array[:] = src.x.array
268 dst.x.scatter_forward()
270 # -- steady PTC ------------------------------------------------------
271 def _warm_start_cfl(self):
272 """The final CFL of the last converged steady solve if the state now in
273 u_vec is that solve's end state (the warm start of an optimization
274 step), else None. Collective."""
275 end = self._steady_end
276 if end is None:
277 return None
278 u = self.model.u_vec.x.array
279 same = end[0].shape == u.shape and bool(np.array_equal(end[0], u))
280 return end[1] if self.model.mesh_obj.mesh.comm.allreduce(same, op=MPI.MIN) else None
282 def solve_steady(self, controller=None, max_steps=None, tol=None):
283 """
284 Pseudo-transient continuation to the steady state from the current
285 u_vec. Returns the model's physical residual Vec.
287 With model.ptc_warm_start (default) and no explicit controller, a
288 start from the end state of the previous converged steady solve
289 continues at that solve's final CFL: the design (or alpha) has moved
290 a little, so the flow is near its new steady state, and ramping the
291 CFL up from cfl0 again would only repeat the early steps. A step
292 that is too bold is rejected and the CFL cut, as always.
293 """
294 m = self.model
295 if self.term.physical:
296 self.term.physical_off()
297 if not self.term.pseudo:
298 raise RuntimeError("steady PTC needs the pseudo-time term (model.use_ptc = True "
299 "before set_up_solver)")
300 if self.boundary_updaters:
301 raise RuntimeError("time-dependent boundary conditions (riemann2) have no steady state")
302 ctrl = controller if controller is not None else CFLController(**m.ptc_settings)
303 if controller is None and getattr(m, "ptc_warm_start", False):
304 warm_cfl = self._warm_start_cfl()
305 if warm_cfl is not None:
306 ctrl.cfl = float(min(max(warm_cfl, ctrl.cfl), ctrl.cfl_max))
307 self._print("[ptc] warm start: continuing at CFL {:.3e}".format(ctrl.cfl))
308 tol = m.residual_conv_limit if tol is None else tol
309 max_steps = int(m.max_newton_iterations if max_steps is None else max_steps)
310 m._ensure_residual_vector()
311 # one Newton step per pseudo step
312 max_it_saved, m.solver.max_it = m.solver.max_it, 1
313 r_prev = m._assemble_physical_residual()
314 self._print("[ptc] initial |R| = {:.6e}".format(r_prev))
315 m.solver_converged = r_prev < tol
316 it = 0
317 while not m.solver_converged and it < max_steps:
318 it += 1
319 cfl = ctrl.cfl
320 self.term.set_cfl(cfl)
321 self._copy(self.term.U_old, m.u_vec)
322 m.last_step_theta = 1.0
323 m.solver.solve(m.u_vec)
324 m.u_vec.x.scatter_forward()
325 r = m._assemble_physical_residual()
326 ksp_ok = m.solver.krylov_solver.getConvergedReason() > 0
327 theta = m.last_step_theta
328 accept, limited = ctrl.update(r_prev, r, ksp_ok, theta)
329 self._print("[ptc] step {:3d} CFL {:.3e} |R| = {:.6e}{}{}".format(
330 it, cfl, r, "" if ksp_ok else " (linear solve not converged)",
331 " (limited, theta {:.1e})".format(theta) if limited else ""))
332 if not accept:
333 self._copy(m.u_vec, self.term.U_old)
334 self._print("[ptc] rejected; CFL -> {:.3e}".format(ctrl.cfl))
335 if ctrl.exhausted:
336 break
337 continue
338 m.solver_converged = r < tol
339 r_prev = r
340 if ctrl.exhausted:
341 self._print("[ptc] no update for {} steps (positivity limiter theta = 0); "
342 "stopping".format(ctrl.frozen_steps))
343 break
344 m.solver.max_it = max_it_saved
345 if not m.solver_converged:
346 self._print("WARNING: PTC stopped after {} steps at |R| = {:.6e}".format(it, r_prev))
347 self.last_steady_steps = it
348 self._steady_end = ((m.u_vec.x.array.copy(), ctrl.cfl) if m.solver_converged else None)
349 m._assemble_physical_residual()
350 m.u_vec.x.scatter_forward()
351 return m.residual_vec_physical
353 # -- unsteady --------------------------------------------------------
354 def start(self, time=0.0, step_index=0, dt_prev=None):
355 """Begin (or restart) time stepping from the current u_vec, which is
356 taken as U^n (and U^{n-1} unless restored from a checkpoint)."""
357 self.time, self.step_index, self.dt_prev = time, step_index, dt_prev
358 self._copy(self.term.U_n, self.model.u_vec)
359 if dt_prev is None:
360 self._copy(self.term.U_nm1, self.model.u_vec)
362 def advance(self, step_control, order=2, newton_rtol=1e-8, newton_atol=1e-10,
363 newton_max_it=20, dual_time=None):
364 """
365 One physical time step from U^n. The step size comes from
366 step_control (StepSizeController); a failed Newton solve restores
367 U^n and halves dt, up to step_control.max_retries times.
368 `dual_time` (a CFLController) switches the inner solve to PTC on R*.
369 Returns a StepRecord; record.accepted is False if every retry failed.
370 """
371 m, term = self.model, self.term
372 if not term.physical:
373 raise RuntimeError("unsteady steps need the physical time term (enable_time_stepping)")
374 solver = m.solver
375 dt = step_control.first(self.dt_prev)
376 for attempt in range(step_control.max_retries + 1):
377 k = 1 if (self.dt_prev is None or order == 1) else 2
378 ratio = term.set_bdf(k, dt, self.dt_prev)
379 for b in self.boundary_updaters:
380 b.prepare(k, ratio, dt)
381 if dual_time is None:
382 term.set_cfl(None)
383 ok, rec = self._newton_solve(newton_rtol, newton_atol, newton_max_it)
384 else:
385 ok, rec = self._dual_time_solve(dual_time, newton_rtol, newton_atol, newton_max_it)
386 if ok:
387 rec.dt, rec.order, rec.retries = dt, k, attempt
388 break
389 self._print("step {:4d} dt={:.3e}: Newton failed (|R| {:.3e} -> {:.3e}); retrying at "
390 "dt={:.3e}".format(self.step_index + 1, dt, rec.res_initial, rec.res_final, 0.5 * dt))
391 self._copy(m.u_vec, term.U_n)
392 dt *= 0.5
393 else:
394 return StepRecord(accepted=False, dt=dt, retries=step_control.max_retries)
395 for b in self.boundary_updaters:
396 b.commit()
397 self.time += dt
398 self.step_index += 1
399 self.dt_prev = dt
400 self._copy(term.U_nm1, term.U_n)
401 self._copy(term.U_n, m.u_vec)
402 return rec
404 def _newton_solve(self, rtol, atol, max_it):
405 m = self.model
406 solver = m.solver
407 crit = (solver.convergence_criterion, solver.rtol, solver.atol, solver.max_it)
408 solver.convergence_criterion, solver.rtol, solver.atol, solver.max_it = "residual", rtol, atol, max_it
409 snes = solver.snes
410 snes.setConvergenceHistory(reset=True)
411 try:
412 n, converged = solver.solve(m.u_vec)
413 finally:
414 solver.convergence_criterion, solver.rtol, solver.atol, solver.max_it = crit
415 hist, _ = snes.getConvergenceHistory()
416 r0, r1 = (hist[0], hist[-1]) if len(hist) else (np.nan, np.nan)
417 rec = StepRecord(accepted=True, newton_its=n, ksp_its=snes.getLinearSolveIterations(),
418 res_initial=r0, res_final=r1)
419 return bool(converged and np.isfinite(r1)), rec
421 def _dual_time_solve(self, ctrl, rtol, atol, max_it):
422 """PTC on R* inside one physical step; converged when |R*| < max(atol,
423 rtol |R*_0|). Each pseudo step is one Newton step."""
424 m, term = self.model, self.term
425 ctrl.reset()
426 b = m.solver.b
427 r_prev = self._unsteady_residual_norm(b)
428 r0, its, ksp = r_prev, 0, 0
429 target = max(atol, rtol * r0)
430 while r_prev > target and its < max_it:
431 its += 1
432 term.set_cfl(ctrl.cfl)
433 self._copy(term.U_old, m.u_vec)
434 m.last_step_theta = 1.0
435 m.solver.solve(m.u_vec)
436 ksp += m.solver.snes.getLinearSolveIterations()
437 r = self._unsteady_residual_norm(b)
438 ok_ksp = m.solver.krylov_solver.getConvergedReason() > 0
439 accept, _ = ctrl.update(r_prev, r, ok_ksp, m.last_step_theta)
440 if not accept:
441 self._copy(m.u_vec, term.U_old)
442 if ctrl.exhausted:
443 break
444 continue
445 r_prev = r
446 term.set_cfl(None)
447 rec = StepRecord(accepted=True, newton_its=its, ksp_its=ksp, res_initial=r0, res_final=r_prev)
448 return bool(r_prev <= target and np.isfinite(r_prev)), rec
450 def _unsteady_residual_norm(self, b):
451 """|R*| with the pseudo term off (at U = U_old it vanishes anyway)."""
452 m = self.model
453 cfl_saved = self.term.inv_cfl.value
454 self.term.inv_cfl.value = 0.0
455 m.problem.F(m.u_vec.x.petsc_vec, b)
456 self.term.inv_cfl.value = cfl_saved
457 return b.norm()