Coverage for dragonfly_sim/core/Euler_model.py: 86%
636 statements
« prev ^ index » next coverage.py v7.16.2, created at 2026-10-17 03:37 +0000
1import ctypes
2import gc
3import contextlib
5import numpy as np
6from mpi4py import MPI
7from petsc4py import PETSc
8import ufl
9import ufl.classes
10import dolfinx
11import dolfinx.fem.petsc
13from dragonfly_sim.utils.Nonlinear_utils import NonlinearProblem_mod, SNESNewtonSolver
14from dragonfly_sim.utils.petsc_utils import set_petsc_vec_array
15from dragonfly_sim.utils.solver_utils import (locate_positivity_limiting_nodes, ksp_converged_reason_name,
16 ksp_norm_type_name, ksp_pc_side_name)
17from dragonfly_sim.core.form_manager import FormManager
18from dragonfly_sim.utils.Euler_utils import primitives, pressure, energy_density, flux, slip_wall_state, hll_flux, boundary_flux
19from dragonfly_sim.utils.NS_utils import (mean_flow, concat_state, extended_flux, extended_normal_flux,
20 extended_hll_flux, extended_boundary_flux)
21from dragonfly_sim.utils.farfield_utils import riemann_farfield_state
22from dragonfly_sim.utils.filewriter import FileWriter
25class CompressibleEulerModel():
26 """
27 The p = 0 DG (cell-centred finite-volume) compressible Euler model, and
28 the host every other flow model in the package is built on.
30 State: the conservative mean flow (rho, rho u, rho E), n_mean = 2 + dim
31 components per cell, optionally followed by `extra_variables`
32 transported scalars (n_state components in all). The scalars are
33 convected with the mean flow (NS_utils.extended_flux) and are not
34 positivity-limited. With no extra variables every form below is the
35 plain Euler one.
37 Plug-in points. CompressibleRANSModel fills these; for the plain Euler
38 model they are all empty or None, and the residual, Jacobian and
39 derivative forms are the same expressions as without them.
41 face_traces () -> (U+, U-) interior-facet traces of the
42 inviscid flux (MUSCL reconstruction)
43 residual_terms callables () -> UFL form added to the residual
44 boundary_terms callables (U_b, tag, kind) -> UFL form or None,
45 added to every boundary condition
46 boundary_scalars (kind, U, inflow) -> list of the extra variables'
47 boundary values (kind: slip, farfield,
48 wall); required when n_extra > 0
49 freestream_scalars () -> list of the extra variables' free-stream values
50 wall_bc (tag) -> None, replaces the slip wall in define_wall_bc
51 wall_traction (U, n) -> traction added to p n in wall_force_density
52 time_scale_terms callables (U) -> added to the PTC inverse time scale
53 derived_fields DerivedField objects (utils/derived_fields.py):
54 fields the residual reads that are not unknowns
55 (a reconstructed gradient, cell centres, the wall
56 distance). Every derivative assembler adds their
57 chain-rule terms.
58 problem_factory (F, J) -> nonlinear problem, replacing
59 NonlinearProblem_mod in set_up_solver
60 extra_scales () -> list of the extra variables' typical
61 magnitudes (numbers or UFL expressions of
62 Constants); required by state_scales when
63 n_extra > 0
65 Every boundary condition records (U_b, tag, kind) in boundary_states.
67 Time integration (core/time_integration.py): use_ptc makes solve_system
68 a pseudo-transient continuation solve; enable_time_stepping() adds a
69 BDF physical time term for unsteady runs. Both live in newton_forms()
70 only, so self.F stays the steady residual that the adjoint
71 differentiates.
72 """
73 def __init__(self, mesh, poly_order=0, gamma=1.4, flux_function=hll_flux, extra_variables=()): # llf_flux # hll_flux
74 if poly_order > 0:
75 raise ValueError("This release only support 0th-order DG bases for now")
77 self.mesh_obj = mesh
78 self.poly_order = poly_order
79 self.gamma = gamma
80 self.dimensions = self.mesh_obj.mesh._ufl_domain._geometric_dimension
82 self.flux_function = flux_function
84 # State layout: mean flow, then the transported scalars
85 self.extra_variables = tuple(extra_variables)
86 self.n_mean = 2 + self.dimensions
87 self.n_extra = len(self.extra_variables)
88 self.n_state = self.n_mean + self.n_extra
89 if self.n_extra and flux_function is not hll_flux:
90 raise NotImplementedError(
91 "Extra transported variables need the HLL flux (NS_utils.extended_hll_flux); "
92 "the other Riemann fluxes have no extended-state version.")
94 # Plug-in points (see the class docstring); all inert by default
95 self.boundary_states = []
96 self.face_traces = None
97 self.residual_terms = []
98 self.boundary_terms = []
99 self.boundary_scalars = None
100 self.freestream_scalars = None
101 self.wall_bc = None
102 self.wall_traction = None
103 self.time_scale_terms = []
104 self.derived_fields = []
105 self.problem_factory = None
106 self.extra_scales = None
108 # Time integration (core/time_integration.py). time_integrator is
109 # created by set_up_solver when use_ptc is set or by
110 # enable_time_stepping(); ptc_settings are CFLController keywords.
111 self.use_ptc = False
112 self.ptc_settings = {}
113 # a steady PTC solve that starts from the end state of the previous
114 # converged steady solve continues at that solve's final CFL instead
115 # of ramping up from cfl0 (TimeIntegrator.solve_steady)
116 self.ptc_warm_start = True
117 self.time_integrator = None
118 self._time_stepping = None
119 # Smallest positivity step length applied since the caller last
120 # reset it; the PTC CFL law reads it.
121 self.last_step_theta = 1.0
123 # Krylov solver of the Newton steps (see apply_krylov_solver_settings);
124 # direct_linear_solver switches every linear solve to MUMPS.
125 self.krylov_type = "gmres"
126 self.krylov_diagonal_scale = True
127 self.direct_linear_solver = False
128 # relative tolerance of that Krylov solve, set before set_up_solver
129 self.krylov_rtol = 1e-3
131 # Instantiate the FormManager
132 self.forms = FormManager(mesh, model=self,
133 measure_source=mesh.form_manager)
135 self.functionspaces = None
136 self.bcs = []
137 self.farfield_conditions = None
138 self.initial_conditions = None
139 self.F = None
140 self.J = None
141 self.solver_is_set_up = False
142 self.solver_converged = False
146 # Newton steps taken per pass of solve_system's loop
147 self.fom_newton_inner_max_it = 5
149 # Report which nodes bound the Newton step whenever the positivity
150 # limiter cuts theta below the step-length cap -- see
151 # _report_positivity_limiting_node. Costs three allreduces and one
152 # small allgather, and ONLY on iterations where theta was already cut,
153 # so it is on by default: a collapsed theta is the failure this most
154 # often needs explaining, and the limiter itself reports only the
155 # surviving step length. Set False once a case is solving cleanly and
156 # the per-iteration lines are just noise.
157 self.report_positivity_limiter = True
159 # ASM overlap and ILU fill for the preconditioner
160 self.asm_overlap = 2
161 self.ilu_levels = 2
163 self.fom_solution_writer = None
164 self._output_fields = []
165 # False until define_sol_export_files actually opens the writers, which
166 # keeps write_solution_output inert for a model that never opened them
167 self.export_solutions = False
170 def output_field_expressions(self):
171 """The fields of the solution file, as (name, value shape, UFL
172 expression of u_vec): the conservative state as separate fields --
173 density rho, momentum rho*u (a vector), energy rho*E, and each
174 transported scalar under its own name (e.g. rho_nu_tilde) -- then the
175 velocity u (a vector) and the pressure p."""
176 dim = self.dimensions
177 U = self.mean_flow(self.u_vec)
178 momentum = ufl.as_vector([U[1 + i] for i in range(dim)])
179 fields = [("density", (), U[0]),
180 ("momentum", (dim,), momentum),
181 ("energy", (), U[1 + dim])]
182 fields += [(name, (), self.u_vec[self.n_mean + k])
183 for k, name in enumerate(self.extra_variables)]
184 fields += [("velocity", (dim,), momentum / U[0]),
185 ("pressure", (), pressure(U, self.gamma))]
186 return fields
188 def define_sol_export_files(self, func_space=None, file_name_addendum=None, mesh_deformation=False):
189 """Open FOM_solution_<suffix>.bp, one file holding every output field
190 (output_field_expressions) as a separate, named field; the mesh is
191 stored once per frame for all of them. Collective.
193 It is the ALIGNED series: exactly one frame per design evaluation,
194 written by write_solution_output() at eval_idx, so step k here is step
195 k in mesh_deformation_<suffix>.bp. It carries the accepted solution --
196 see DG_windtunnel_model.solve_residual_equations. mesh_deformation=True
197 writes each frame on the deformed mesh it was solved on. func_space is
198 unused (the fields define their own spaces); it is kept for callers.
199 """
200 spaces = {}
201 self._output_fields = []
202 for name, shape, expr in self.output_field_expressions():
203 if shape not in spaces:
204 spaces[shape] = (self.functionspaces["V_scalar"] if shape == () else
205 dolfinx.fem.functionspace(self.mesh_obj.mesh,
206 ("DG", self.poly_order, shape)))
207 V = spaces[shape]
208 # compiled once here (collective), interpolated at every write
209 self._output_fields.append(
210 (name, dolfinx.fem.Function(V, name=name),
211 dolfinx.fem.Expression(expr, V.element.interpolation_points)))
212 self.fom_solution_writer = FileWriter(
213 "FOM_solution_{}".format(file_name_addendum), self.mesh_obj.mesh.comm,
214 {name: f.function_space for name, f, _ in self._output_fields},
215 mesh_deformation=mesh_deformation)
216 self.export_solutions = True
218 def _write_output_fields(self, write_counter):
219 """Interpolate the output fields off the current u_vec and write one
220 frame. u_vec must already hold the state being exported. Collective."""
221 for _, f, expr in self._output_fields:
222 f.interpolate(expr)
223 self.fom_solution_writer.interpolate_and_write(
224 {name: f for name, f, _ in self._output_fields}, write_counter=write_counter)
226 def write_solution_output(self, write_counter, solution_array=None):
227 """Write one aligned output frame at index `write_counter`.
229 Called once per design evaluation from
230 DG_windtunnel_model.solve_residual_equations with write_counter =
231 eval_idx, so FOM_solution_<suffix>.bp steps in lockstep with
232 mesh_deformation_<suffix>.bp.
234 `solution_array` is the accepted solution of the evaluation, passed
235 explicitly rather than read off self.u_vec. Passing None writes
236 whatever u_vec currently holds.
237 """
238 if not self.export_solutions:
239 return
241 if solution_array is not None:
242 self.interpolate_solution_vector(solution_array)
244 self._write_output_fields(write_counter)
246 def define_elements_functionspaces(self):
247 # Throughout this class we need access to various function spaces; this is where we define them
248 self.functionspaces = {
249 "V": dolfinx.fem.functionspace(self.mesh_obj.mesh, ("DG", self.poly_order, (self.n_state,))),
250 "V_scalar": dolfinx.fem.functionspace(self.mesh_obj.mesh, ("DG", self.poly_order)),
251 }
253 def define_trial_testfunctions(self):
254 # Define the trial and test functions for the weak form
255 self.u_vec = dolfinx.fem.Function(self.functionspaces["V"], name="u")
256 self.v_vec = ufl.TestFunction(self.functionspaces["V"])
258 def define_farfield_conditions(self, farfield):
259 """
260 Set the free-stream state the far-field boundary condition relaxes to.
262 `farfield` is {'rho', 'p', 'M', 'alpha'}: density, pressure, Mach
263 number and angle of attack (rad). This is all the user states; the
264 boundary treatment (inflow, outflow) follows from it per facet in
265 define_farfield_bc, through the characteristic state.
267 NOTE: the far field is the subsonic characteristic one, so M < 1.
268 A supersonic free stream needs a different boundary state in
269 define_farfield_bc (all variables from the free stream at inflow, all
270 from the interior at outflow); this is the one place that would change.
271 """
272 if not isinstance(farfield, dict) or set(farfield) != {'rho', 'p', 'M', 'alpha'}:
273 raise ValueError(
274 "The far-field conditions are a flat dict {{'rho': ..., 'p': ..., 'M': ..., "
275 "'alpha': ...}} (density, pressure, Mach number, angle of attack in rad); "
276 "got {}. The former {{'inlet': {{...}}, 'outlet': {{...}}}} form is gone: "
277 "pass the contents of 'inlet' and drop 'outlet'.".format(
278 sorted(farfield) if isinstance(farfield, dict) else type(farfield).__name__))
279 rho, p, M, alpha = (float(farfield[k]) for k in ('rho', 'p', 'M', 'alpha'))
280 if rho <= 0.0 or p <= 0.0:
281 raise ValueError("far-field density and pressure must be positive")
282 if not 0.0 <= M < 1.0:
283 raise NotImplementedError(
284 "far-field Mach number {} outside [0, 1): only the subsonic "
285 "characteristic far field is implemented".format(M))
287 c = (self.gamma * p / rho) ** 0.5 # speed of sound
289 # Since angle of attack is a design parameter we define it as a UFL variable;
290 # this allows us to use Dolfin's automatic differentiation for it
291 self.alpha_const = dolfinx.fem.Constant(self.mesh_obj.mesh, alpha)
292 self.alpha_var = ufl.variable(self.alpha_const)
294 inflow_dir = self._inflow_direction_ufl()
296 # unit vector along the free stream (the drag direction)
297 flow_dir = dolfinx.fem.Constant(
298 self.mesh_obj.mesh, self._inflow_direction_numpy(alpha))
300 # defining the reference velocity so we can update it later
301 # if the Mach number gets modified
302 self.u_ref_const = dolfinx.fem.Constant(self.mesh_obj.mesh, M*c)
303 u_inf = self.u_ref_const * inflow_dir
305 self.farfield_conditions = {'rho': rho,
306 'p': p,
307 'M': M,
308 'alpha': alpha,
309 'inflow_angle': self._inflow_direction_numpy(alpha),
310 'c': c,
311 'flow_dir': flow_dir,
312 'u': u_inf,
313 'u_ref': self.u_ref_const}
315 def _inflow_direction_ufl(self):
316 # unit freestream direction as a UFL expression
317 if self.dimensions == 2:
318 # x = chordwise, y = vertical
319 return ufl.as_vector([ufl.cos(self.alpha_var), ufl.sin(self.alpha_var)])
320 elif self.dimensions == 3:
321 # x = chordwise, y = spanwise, z = vertical
322 return ufl.as_vector([ufl.cos(self.alpha_var), 0., ufl.sin(self.alpha_var)])
323 else:
324 raise ValueError("Unsupported dimension {}".format(self.dimensions))
326 def _inflow_direction_numpy(self, alpha):
327 # unit freestream direction as a numpy array
328 if self.dimensions == 2:
329 return np.array([np.cos(alpha), np.sin(alpha)], dtype=np.double)
330 elif self.dimensions == 3:
331 return np.array([np.cos(alpha), 0., np.sin(alpha)], dtype=np.double)
332 else:
333 raise ValueError("Unsupported dimension {}".format(self.dimensions))
335 def set_angle_of_attack(self, alpha):
336 # propagate the new angle of attack `alpha` to all model dependencies
337 alpha = float(alpha)
339 self.alpha_const.value = alpha
341 inflow_angle = self._inflow_direction_numpy(alpha)
342 self.farfield_conditions['alpha'] = alpha
343 self.farfield_conditions['inflow_angle'] = inflow_angle
344 self.farfield_conditions['flow_dir'].value = inflow_angle
346 def set_mach(self, M):
347 # propagate the new Mach number `M` to all model dependencies
348 M = float(M)
349 c = self.farfield_conditions['c']
350 self.u_ref_const.value = M * c
351 self.farfield_conditions['M'] = M
353 def compute_initial_conditions_from_farfield(self):
354 # The initial guess used in the Newton solver: the free stream.
355 ff = self.farfield_conditions
356 rho_in = dolfinx.fem.Constant(self.mesh_obj.mesh, ff['rho'])
357 rhoE_in_guess = energy_density(ff['p'], rho_in, ff['u'], self.gamma)
358 self.initial_conditions = ufl.as_vector((rho_in, *[rho_in*ff['u'][i] for i in range(self.dimensions)], rhoE_in_guess))
359 if self.n_extra:
360 # plain numbers become Constants on this mesh, so the initial
361 # condition stays interpolable
362 extras = [v if isinstance(v, ufl.core.expr.Expr)
363 else dolfinx.fem.Constant(self.mesh_obj.mesh, float(v))
364 for v in self._freestream_scalars()]
365 self.initial_conditions = concat_state(self.initial_conditions, extras)
367 def interpolate_solution_vector(self, interpolant=None):
368 if interpolant is None:
369 interpolant = self.initial_conditions
371 if isinstance(interpolant, (np.ndarray, PETSc.Vec)):
372 set_petsc_vec_array(self.u_vec.x.petsc_vec, interpolant)
373 else:
374 for j in range(interpolant.ufl_shape[0]):
375 self.u_vec.sub(j).interpolate(
376 dolfinx.fem.Expression(interpolant[j], self.functionspaces['V'].sub(j).element.interpolation_points,
377 comm=self.mesh_obj.mesh.comm))
379 self.u_vec.x.scatter_forward()
381 # ------------------------------------------------------------------
382 # State layout
383 # ------------------------------------------------------------------
384 def mean_flow(self, U):
385 """The mean-flow slice (rho, rho u, rho E) of a state expression."""
386 return mean_flow(U, self.n_mean)
388 def _freestream_scalars(self):
389 if self.freestream_scalars is None:
390 raise RuntimeError("{} extra variables but no freestream_scalars plug-in".format(self.n_extra))
391 return list(self.freestream_scalars())
393 def _boundary_state(self, U_b_mean, kind, inflow=None):
394 """The full boundary state: the mean-flow state and, with extra
395 variables, their values from the boundary_scalars plug-in."""
396 if not self.n_extra:
397 return U_b_mean
398 if self.boundary_scalars is None:
399 raise RuntimeError("{} extra variables but no boundary_scalars plug-in".format(self.n_extra))
400 return concat_state(U_b_mean, list(self.boundary_scalars(kind, self.u_vec, inflow)))
402 # ------------------------------------------------------------------
403 # Boundary conditions
404 # ------------------------------------------------------------------
405 def add_boundary_condition(self, meshtag, U_b, kind, inviscid_flux="llf"):
406 """
407 Weakly impose the boundary state U_b on ds(meshtag) and register it.
409 inviscid_flux="llf" is the one-sided Lax-Friedrichs boundary flux
410 against U_b (Euler_utils.boundary_flux); "exact" is the physical flux
411 of U_b itself, for states that already carry the upwinding (the
412 characteristic far field) or that must have no mass flux (no-slip
413 wall). The boundary_terms plug-ins add their terms (e.g. viscous
414 fluxes) to the same integral.
415 """
416 U, n = self.u_vec, self.mesh_obj.n
417 if inviscid_flux == "llf":
418 if self.n_extra:
419 f = extended_boundary_flux(U, U_b, n, self.gamma, self.n_mean)
420 else:
421 f = boundary_flux(U, U_b, n, self.gamma)
422 elif inviscid_flux == "exact":
423 f = extended_normal_flux(U_b, n, self.gamma, self.n_mean)
424 else:
425 raise ValueError("inviscid_flux must be 'llf' or 'exact'")
426 F = ufl.inner(f, self.v_vec) * self.forms.ds(meshtag)
427 for term in self.boundary_terms:
428 extra = term(U_b, meshtag, kind)
429 if extra is not None:
430 F = F + extra
431 self.bcs += [F]
432 self.boundary_states.append((U_b, meshtag, kind))
434 def define_slipwall_bc(self, meshtag):
435 # define the weak form term for natural enforcement of the slip wall boundary condition
436 u_wall = slip_wall_state(self.mean_flow(self.u_vec), self.mesh_obj.n)
437 self.add_boundary_condition(meshtag, self._boundary_state(u_wall, "slip"), "slip")
439 def define_wall_bc(self, meshtag):
440 """The wall of the body: a slip wall, unless a wall_bc plug-in
441 (e.g. the RANS model's no-slip wall) replaces it."""
442 if self.wall_bc is not None:
443 self.wall_bc(meshtag)
444 else:
445 self.define_slipwall_bc(meshtag)
447 def define_farfield_bc(self, meshtag, transverse=None):
448 """
449 Characteristic (Riemann-invariant) far field on ds(meshtag), for the
450 whole outer boundary: no inflow/outflow division is made, each facet
451 takes its character from the sign of the characteristic speeds, so the
452 angle of attack can change freely. Outgoing acoustic waves leave the domain (exactly at normal incidence)
453 instead of reflecting off a fixed state. The inviscid flux is the
454 physical flux of the characteristic state, which already carries the
455 upwinding; extra variables take their free-stream values where the
456 flow enters and the interior values where it leaves.
458 `transverse` (core/farfield.py:TransverseFarfield) adds the
459 transverse correction to the incoming invariant: the 'riemann2' far
460 field of unsteady runs.
462 The state depends on alpha through the free-stream velocity, so
463 dR/dalpha follows automatically.
464 """
465 ff = self.farfield_conditions
466 R_in = None if transverse is None else transverse.R_in
467 U_b, inflow = riemann_farfield_state(self.mean_flow(self.u_vec), self.mesh_obj.n,
468 ff['rho'], ff['u'], ff['p'], self.gamma,
469 R_in=R_in)
470 self.add_boundary_condition(meshtag, self._boundary_state(U_b, "farfield", inflow),
471 "farfield", inviscid_flux="exact")
472 if transverse is not None:
473 self.transverse_farfield = transverse
475 def compute_interior_integral_terms(self):
476 if not self.n_extra and self.face_traces is None and not self.residual_terms:
477 # the plain Euler residual, written exactly as it always was
478 F_vol = ufl.inner(ufl.grad(self.v_vec), flux(self.u_vec, self.gamma)) * self.forms.dx
480 F_int = ufl.inner(
481 self.flux_function(self.u_vec('+'), self.u_vec('-'), self.mesh_obj.n('+'), self.gamma),
482 self.v_vec('+') - self.v_vec('-')) * self.forms.dS
483 return -F_vol + F_int
485 U, v, n = self.u_vec, self.v_vec, self.mesh_obj.n
486 F = -ufl.inner(ufl.grad(v), extended_flux(U, self.gamma, self.n_mean)) * self.forms.dx
487 U_p, U_m = self.face_traces() if self.face_traces is not None else (U('+'), U('-'))
488 if self.n_extra:
489 F_num = extended_hll_flux(U_p, U_m, n('+'), self.gamma, self.n_mean)
490 else:
491 F_num = self.flux_function(U_p, U_m, n('+'), self.gamma)
492 F += ufl.inner(F_num, v('+') - v('-')) * self.forms.dS
493 for term in self.residual_terms:
494 F += term()
495 return F
497 def _lambda_max_expr(self, U):
498 """The largest wave speed |u| + c of the state U (mean flow), with rho
499 and p floored as in Euler_utils.signed_wave_speeds."""
500 U = self.mean_flow(U)
501 rho, u_vel, _ = primitives(U)
502 p = pressure(U, self.gamma)
503 rho_safe = ufl.max_value(rho, 1e-8)
504 p_safe = ufl.max_value(p, 1e-8)
505 c = ufl.sqrt(self.gamma * p_safe / rho_safe)
506 return ufl.sqrt(ufl.dot(u_vel, u_vel)) + c
508 def wall_force_density(self, U, n):
509 """Force per unit wall area on the body: pressure, plus the
510 wall_traction plug-in's viscous traction when there is one."""
511 f = pressure(self.mean_flow(U), self.gamma) * n
512 if self.wall_traction is not None:
513 f = f + self.wall_traction(U, n)
514 return f
516 def assemble_dRdalpha_vec(self, u_vec_entries=None):
517 """
518 The derivative dR/d(alpha) of the converged residual, assembled as a Vec.
520 A VECTOR, not a matrix, because alpha is a single global scalar: the
521 derivative of the residual with respect to it is one column, and
522 ufl.diff against the stored alpha_var yields a rank-1 form that
523 assembles straight into it. Nothing has to be contracted afterwards.
525 Note this is ufl.diff, not the ufl.derivative used by every other
526 assembler here. ufl.derivative only accepts Coefficients and
527 SpatialCoordinate and rejects a Constant outright; ufl.diff against a
528 ufl.variable-wrapped Constant is the route that works. See
529 define_farfield_conditions for why alpha_var must be the single
530 stored instance.
532 Sparse in practice: alpha only enters the residual through the
533 free-stream velocity in the far-field boundary state, so the result is
534 supported on far-field-adjacent dofs alone. The interior terms and the
535 slip wall (which uses the facet normal) carry no alpha dependence.
537 Caches its own form and Vec.
538 """
539 if u_vec_entries is not None:
540 self.set_u_vec(u_vec_entries)
542 if self.derived_fields:
543 self.update_derived_fields()
544 if not hasattr(self, 'dRdalpha_vec'):
545 self.dRdalpha_vec = dolfinx.fem.petsc.create_vector(
546 dolfinx.fem.extract_function_spaces(self.F_physical_form))
547 return self._chain.residual_alpha(self.F, "F", self.dRdalpha_vec)
549 if not hasattr(self, 'dRdalpha_form'):
550 self.dRdalpha_form = dolfinx.fem.form(
551 ufl.diff(self.F, self.alpha_var))
552 self.dRdalpha_vec = dolfinx.fem.petsc.create_vector(
553 dolfinx.fem.extract_function_spaces(self.dRdalpha_form))
555 with self.dRdalpha_vec.localForm() as loc:
556 loc.set(0.0)
557 dolfinx.fem.petsc.assemble_vector(self.dRdalpha_vec, self.dRdalpha_form)
558 self.dRdalpha_vec.ghostUpdate(addv=PETSc.InsertMode.ADD,
559 mode=PETSc.ScatterMode.REVERSE)
561 return self.dRdalpha_vec
563 def compute_weakform(self):
564 if self.derived_fields and not getattr(self, "_geometry_listener_added", False):
565 # derived fields follow the mesh through every node motion
566 self.mesh_obj.add_geometry_listener(self.update_geometry_fields)
567 self._geometry_listener_added = True
568 if self.F is not None and self.J is not None:
569 return
571 self.F_physical = self.forms.build_residual()
573 self.F_physical_form = self.forms.compiled(
574 "F_physical", lambda: self.F_physical)
576 self.F = self.F_physical
578 self.J = self.compute_dRdu_mat()
580 # ------------------------------------------------------------------
581 # Time integration
582 # ------------------------------------------------------------------
583 def enable_time_stepping(self, dual_time=False):
584 """
585 Add the BDF physical time term to the Newton forms (and, with
586 dual_time, the pseudo-time term as well). Call before set_up_solver.
587 Returns the model's TimeIntegrator.
588 """
589 if self.solver_is_set_up:
590 raise RuntimeError("enable_time_stepping must be called before set_up_solver")
591 self._time_stepping = {"dual_time": bool(dual_time)}
592 return self._make_time_integrator()
594 def _make_time_integrator(self):
595 from dragonfly_sim.core.time_integration import TimeTerm, TimeIntegrator
596 if self.time_integrator is None:
597 physical = self._time_stepping is not None
598 pseudo = self.use_ptc or (physical and self._time_stepping["dual_time"])
599 if physical or pseudo:
600 self.time_integrator = TimeIntegrator(self, TimeTerm(self, physical=physical, pseudo=pseudo))
601 return self.time_integrator
603 def newton_forms(self):
604 """
605 The residual/Jacobian pair Newton drives: model.F with the time
606 terms the model uses (core/time_integration.py). Without PTC or time
607 stepping this is (self.F, self.J), the plain steady solve.
608 """
609 integrator = self._make_time_integrator()
610 if integrator is None:
611 return self.F, self.J
612 F = integrator.term.augment(self.F)
613 return F, ufl.derivative(F, self.u_vec)
615 def set_up_solver(self):
616 if not self.solver_is_set_up:
617 mat_type = "baij"
618 F_newton, J_newton = self.newton_forms()
619 if self.problem_factory is not None:
620 self.problem = self.problem_factory(F_newton, J_newton)
621 else:
622 self.problem = NonlinearProblem_mod(F_newton, J_newton, self.u_vec,
623 jit_options={"cffi_extra_compile_args": ["-O3", "-march=native", "-ffast-math"]},
624 mat_type=mat_type)
625 self.solver = SNESNewtonSolver(self.mesh_obj.mesh.comm, self.problem)
626 self.solver.convergence_criterion = "incremental" # convergence is determined through the norm of the iterative solution update
627 self.solver.rtol = 1e-6
628 # Decoupled from rtol: with ||x|| ~ O(1e4), the old atol=1e-8
629 # floor demanded ~1e-12 RELATIVE precision on ||dx|| -- far
630 # tighter than rtol=1e-6 -- so Newton was burning its full
631 # max_it budget chasing an unreachable absolute target on every
632 # solve instead of stopping once rtol was satisfied. Setting
633 # atol to a floor well below anything reachable makes rtol the
634 # sole criterion carrying the Eisenstat-Walker forcing, as
635 # intended.
636 self.solver.atol = 1e-11
637 self.solver.max_it = int(self.fom_newton_inner_max_it)
639 self.solver.error_on_nonconvergence = False
641 def limiter(solver, x, dx):
642 ksp = solver.krylov_solver
644 # Print solver summary (number of iterations, current residual norm, reason)
645 PETSc.Sys.Print("nls_solve_ solve: {} iterations, residual norm {} [{}]".format(
646 ksp.getIterationNumber(), ksp.getResidualNorm(),
647 ksp_converged_reason_name(ksp.getConvergedReason())))
648 dx_norm = dx.norm()
649 PETSc.Sys.Print("x pre norm: {}".format(x.norm()))
651 theta, initial_theta = self.positivity_step_length(x, dx)
652 self.last_step_theta = min(self.last_step_theta, theta)
654 # Print (density and pressure) positivity-preserving under-relaxation summary
655 PETSc.Sys.Print("dx norm: {}; Initial theta: {}, theta applied: {}".format(dx_norm, initial_theta, theta))
656 self._report_positivity_limiting_node(x, dx, theta, initial_theta)
658 # The solver lands on x - dx, so the limiter's job is to
659 # rescale dx in place; x itself must not be touched.
660 dx.scale(theta)
662 self.solver.set_step_limiter(limiter)
664 ksp = self.solver.krylov_solver
665 self.apply_krylov_solver_settings(ksp, monitor_convergence=False, max_it=300,
666 rtol=self.krylov_rtol)
667 self.solver_is_set_up = True
669 # Print Jacobian matrix block size
670 PETSc.Sys.Print("Euler Jacobian matrix type: {}, block size: {}".format(
671 self.solver.A.getType(), self.solver.A.getBlockSize()))
673 # Print summary of various solver parameters
674 rtol, _, _, ksp_max_it = ksp.getTolerances()
675 PETSc.Sys.Print(
676 "Euler forward KSP: type={}, pc={}, norm={}, pc_side={}, "
677 "rtol={:.1e}, max_it={}".format(
678 ksp.getType(), ksp.getPC().getType(),
679 ksp_norm_type_name(ksp.getNormType()),
680 ksp_pc_side_name(ksp.getPCSide()), rtol, ksp_max_it))
682 def positivity_step_length(self, x, dx, tau=1.0):
683 """
684 Step length for the update x - theta dx: (theta, cap), with the
685 step-length cap = min((tau / (||dx|| / max(||x||, 1)))^0.5, 1) and theta
686 <= cap the largest step keeping density and pressure positive in every
687 cell (transported extra variables are not limited). Used by the Newton
688 step limiter and by the reduced-order solve
689 (core/reduced_order_model.py). Collective.
690 """
691 dx_norm = dx.norm()
692 relative_dx = dx_norm / max(x.norm(), 1.)
693 initial_theta = min((tau / relative_dx)**0.5, 1.0) if dx_norm > 1e-16 else 0.
695 from dragonfly_sim.utils.solver_utils import compute_positivity_preserving_theta
696 theta = compute_positivity_preserving_theta(
697 self.mesh_obj.mesh.comm, x.array, dx.array, gamma=self.gamma, initial_theta=initial_theta,
698 block_size=self.n_state, n_mean=self.n_mean if self.n_extra else None)
699 return theta, initial_theta
701 @contextlib.contextmanager
702 def steady_residual(self):
703 """
704 Within this block the Newton problem (self.problem) assembles the
705 steady residual and its Jacobian: the pseudo-time and physical time
706 terms of core/time_integration.py are switched off (their Constants
707 set to zero) and restored on exit. A no-op without time terms.
708 """
709 term = None if self.time_integrator is None else self.time_integrator.term
710 if term is None:
711 yield
712 return
713 saved = (float(term.inv_cfl.value), float(term.inv_dt.value))
714 term.inv_cfl.value, term.inv_dt.value = 0.0, 0.0
715 try:
716 yield
717 finally:
718 term.inv_cfl.value, term.inv_dt.value = saved
720 def state_scale_expressions(self):
721 """
722 Typical magnitudes of the state components: the free-stream
723 [rho, rho U (per direction), p/(gamma - 1) + rho U^2/2] of the mean
724 flow, then the extra_scales plug-in's values for the extra variables
725 (numbers or UFL expressions of Constants).
726 """
727 ff = self.farfield_conditions
728 rho, p = ff['rho'], ff['p']
729 speed = ff['M'] * ff['c']
730 scales = ([rho] + [rho * speed] * self.dimensions
731 + [p / (self.gamma - 1.0) + 0.5 * rho * speed**2])
732 if self.n_extra:
733 if self.extra_scales is None:
734 raise RuntimeError("{} extra variables but no extra_scales plug-in".format(self.n_extra))
735 scales += list(self.extra_scales())
736 return scales
738 def state_scales(self):
739 """state_scale_expressions evaluated to floats, one per state
740 component. Collective (UFL expressions are evaluated by assembly)."""
741 msh = self.mesh_obj.mesh
742 dx = ufl.dx(domain=msh)
743 volume = None
744 values = []
745 for s in self.state_scale_expressions():
746 if isinstance(s, ufl.core.expr.Expr):
747 if volume is None:
748 volume = msh.comm.allreduce(
749 dolfinx.fem.assemble_scalar(dolfinx.fem.form(1.0 * dx)), op=MPI.SUM)
750 s = msh.comm.allreduce(dolfinx.fem.assemble_scalar(dolfinx.fem.form(s * dx)),
751 op=MPI.SUM) / volume
752 values.append(float(s))
753 return values
755 def apply_krylov_solver_settings(self, ksp, max_it=100, gmres_restart=None,
756 monitor_convergence=True,
757 rtol=1e-3, atol=None):
758 """
759 Configure `ksp` through the PETSc options DB.
761 rtol/atol are parameters rather than literals because this method is
762 shared by three very different solves. The DEFAULTS are the forward
763 Euler Newton inner KSP's settings -- a loose rtol=1e-3 carrying the
764 Eisenstat-Walker forcing (see set_up_solver's atol comment), and no
765 atol at all, i.e. PETSc's 1e-50 floor. The adjoint and tangent solves
766 pass the model's standardized linear-solve pair in instead.
768 atol=None means "leave the key unset" and NOT "set zero": writing 0
769 would be a real change from PETSc's 1e-50 default.
771 Both of those call sites previously called ksp.setTolerances() AFTER
772 this method, silently overriding the rtol written here -- so this
773 method's rtol only ever governed the forward solve, which was
774 readable from neither end. Passing tolerances in keeps one source of
775 truth per KSP.
776 """
777 opts = PETSc.Options()
778 option_prefix = ksp.getOptionsPrefix()
780 if monitor_convergence:
781 opts[f"{option_prefix}ksp_converged_reason"] = ""
782 opts[f"{option_prefix}ksp_monitor"] = ""
784 if self.direct_linear_solver:
785 opts[f"{option_prefix}ksp_type"] = "preonly"
786 opts[f"{option_prefix}pc_type"] = "lu"
787 opts[f"{option_prefix}pc_factor_mat_solver_type"] = "mumps"
788 ksp.setFromOptions()
789 return
791 # Plain gmres, not fgmres: fgmres only earns its double vector
792 # storage when the preconditioner itself varies iteration to
793 # iteration (e.g. an inner Krylov solve run to a loose, changing
794 # tolerance). The PC configured below (ASM+ILU) is fixed, so
795 # fgmres was paying for flexibility nothing here uses.
796 opts[f"{option_prefix}ksp_type"] = self.krylov_type
797 opts[f"{option_prefix}ksp_max_it"] = "{}".format(max_it)
798 opts[f"{option_prefix}ksp_rtol"] = rtol
799 if atol is not None:
800 opts[f"{option_prefix}ksp_atol"] = atol
801 # Default restart length to max_it: with a fixed (non-flexible)
802 # preconditioner there is no reason to discard the Krylov subspace
803 # before the iteration budget itself is exhausted -- a restart
804 # shorter than max_it just forces a mid-solve subspace discard for
805 # no memory savings that matter at this scale. Callers that pass
806 # gmres_restart explicitly (e.g. the adjoint solve, which pairs a
807 # much larger max_it with a deliberately smaller restart to bound
808 # memory) are unaffected.
809 if gmres_restart is None:
810 gmres_restart = max_it
811 opts[f"{option_prefix}ksp_gmres_restart"] = "{}".format(gmres_restart)
812 opts[f"{option_prefix}pc_type"] = "asm"
813 opts[f"{option_prefix}pc_asm_overlap"] = "{}".format(self.asm_overlap)
815 opts[f"{option_prefix}sub_pc_type"] = "ilu"
816 # opts[f"{option_prefix}sub_pc_factor_levels"] = "{}".format(int((self.poly_order + 1) ** self.dimensions))
817 opts[f"{option_prefix}sub_pc_factor_levels"] = "{}".format(self.ilu_levels)
818 # ILU(k)'s fill pattern is computed purely from the local
819 # block's sparsity GRAPH (level-of-fill is symbolic, not
820 # value-dependent), so reordering the unknowns before
821 # factoring can reduce the fill actually generated at a given
822 # k without changing accuracy -- a free lever, unlike k
823 # itself which is now pinned between "stalls" (k=3, see
824 # terminal_output_3Dtest_ILU_fillin=3.txt) and "OOM" (k=
825 # dimensions+2) in 3D. Was previously scaffolded but unused
826 # ("natural", i.e. no reordering, per that log's ksp_view).
827 opts[f"{option_prefix}sub_pc_factor_mat_ordering_type"] = "rcm"
829 # the options database is global and the prefix is shared by every KSP of that name:
830 # clear a setting an earlier KSP made (diagonal scaling cannot work on a matrix-free operator)
831 for name in ("ksp_diagonal_scale", "ksp_diagonal_scale_fix"):
832 if opts.hasName(option_prefix + name):
833 opts.delValue(option_prefix + name)
834 if self.krylov_diagonal_scale:
835 opts[f"{option_prefix}ksp_diagonal_scale"] = ""
836 opts[f"{option_prefix}ksp_diagonal_scale_fix"] = ""
838 ksp.setFromOptions()
840 def _assemble_physical_residual(self):
841 """
842 Re-assemble F_physical into self.residual_vec_physical and return its
843 norm. Derived fields (e.g. a reconstructed gradient) are brought up
844 to date with the current state first.
845 """
846 self.update_derived_fields()
847 with self.residual_vec_physical.localForm() as loc:
848 loc.set(0.0)
849 dolfinx.fem.petsc.assemble_vector(self.residual_vec_physical, self.F_physical_form)
850 self.residual_vec_physical.ghostUpdate(
851 addv=PETSc.InsertMode.ADD, mode=PETSc.ScatterMode.REVERSE)
852 return self.residual_vec_physical.norm()
854 def relative_residual_norm(self, full_norm):
855 """
856 eta = ||r|| / ||u||, the dimensionless residual measure.
858 Absolute residual norms are not comparable across problems: they
859 scale with sqrt(n_dofs) and with the magnitude of the conserved
860 variables, so a threshold calibrated on one mesh is meaningless on
861 the next. Dividing by the current ||u|| removes both, which is what
862 makes a single tolerance transferable between the 2D and 3D cases.
864 Collective (Vec.norm). Returns nan rather than raising when u_vec
865 has blown up, so a diverged solve still prints a number.
866 """
867 u_norm = self.u_vec.x.petsc_vec.norm()
868 if not np.isfinite(u_norm) or u_norm == 0.0:
869 return np.nan
870 return full_norm / u_norm
872 def linearized_solver(self, transpose=False, prefix="lin_", rtol=1e-10,
873 atol=1e-11, max_it=1000, gmres_restart=200, release_operands=False,
874 matrix_free=False, ksp_options=None):
875 """
876 The KSP of dR/dU x = rhs (or its transpose, for the adjoint) at the
877 current state, for models whose dR/dU is a condensed product
878 (has_condensed_jacobian). FGMRES on the exact operator, right
879 preconditioned with ASM/ILU of the compact A_UU part -- the forward
880 solve's compact_pc pairing -- or MUMPS with direct_linear_solver.
881 The transposes are formed explicitly. With release_operands and
882 transpose, the condensed dR/dU (and its blocks) is freed once the
883 transposes exist, before the preconditioner is set up: only the
884 transposes are used from then on (release_solver). The KSP holds its
885 operators and, after its first solve, the preconditioner's factors, so
886 it solves any number of right-hand sides at this state; the caller
887 destroys it. Collective.
889 matrix_free (with transpose): the operator is the python Mat of
890 TransposedCondensedOperator, A_UU^T x + B^T (A_UG^T x), which never
891 forms A or the product A_UG B; its context owns the matrices and is
892 destroyed with destroy_adjoint_operator(ksp).
893 """
894 if transpose and matrix_free:
895 return self._matrix_free_adjoint_solver(prefix, rtol, atol, max_it, gmres_restart, ksp_options)
896 A = self.assemble_dRdu_mat()
897 P = self._dRdu_condensed.A_UU
898 if transpose:
899 # out of place: Mat.transpose() without an argument transposes the
900 # cached matrices themselves
901 A, P = A.transpose(PETSc.Mat()), P.transpose(PETSc.Mat())
902 if release_operands:
903 self._release_condensed_jacobian()
904 self._return_freed_memory()
905 ksp = PETSc.KSP().create(self.mesh_obj.mesh.comm)
906 ksp.setOptionsPrefix(prefix)
907 ksp.setOperators(A, A if self.direct_linear_solver else P)
908 saved = (self.krylov_type, self.krylov_diagonal_scale)
909 self.krylov_type, self.krylov_diagonal_scale = "fgmres", False
910 try:
911 self.apply_krylov_solver_settings(ksp, max_it=max_it, gmres_restart=gmres_restart,
912 monitor_convergence=False, rtol=rtol, atol=atol)
913 finally:
914 self.krylov_type, self.krylov_diagonal_scale = saved
915 self._apply_ksp_options(ksp, ksp_options)
916 return ksp
918 # PETSc option keys written by _apply_ksp_options: the options database is global, so they are
919 # removed again when the next KSP is configured instead of leaking into it
920 _applied_ksp_option_keys = set()
922 @classmethod
923 def _apply_ksp_options(cls, ksp, ksp_options):
924 """Override KSP settings through the PETSc options DB, after the model's own:
925 a dict of option names without the prefix, e.g. {"ksp_type": "dgmres",
926 "ksp_gmres_restart": 400}. The overrides last until the next KSP is configured here
927 (the preconditioner reads its options at the first solve), then are deleted. An option
928 the PETSc build does not know (a KSP type it was not built with) raises a PETSc error."""
929 opts = PETSc.Options()
930 for key in cls._applied_ksp_option_keys:
931 opts.delValue(key)
932 cls._applied_ksp_option_keys = set()
933 if not ksp_options:
934 return
935 prefix = ksp.getOptionsPrefix()
936 for key, value in ksp_options.items():
937 opts[prefix + key] = value
938 cls._applied_ksp_option_keys.add(prefix + key)
939 ksp.setFromOptions()
941 def _matrix_free_adjoint_solver(self, prefix, rtol, atol, max_it, gmres_restart, ksp_options=None):
942 from dragonfly_sim.utils.gradient_condensation import TransposedCondensedOperator
943 self.update_derived_fields()
944 rec = self.reconstruction
945 forms = rec.compact_forms(self.F)
946 A_UU = dolfinx.fem.petsc.create_matrix(forms[0])
947 A_UG = dolfinx.fem.petsc.create_matrix(forms[1])
948 for M, a in ((A_UU, forms[0]), (A_UG, forms[1])):
949 M.zeroEntries()
950 dolfinx.fem.petsc.assemble_matrix(M, a)
951 M.assemble()
952 P = A_UU.transpose(PETSc.Mat())
953 A_UU.destroy()
954 gradient = rec.gradient
955 B = gradient.state_jacobian()
956 ctx = TransposedCondensedOperator(P, A_UG, B, gradient.discard)
957 A = ctx.create_mat(self.mesh_obj.mesh.comm)
958 self._return_freed_memory()
959 ksp = PETSc.KSP().create(self.mesh_obj.mesh.comm)
960 ksp.setOptionsPrefix(prefix)
961 ksp.setOperators(A, A if self.direct_linear_solver else P)
962 saved = (self.krylov_type, self.krylov_diagonal_scale)
963 self.krylov_type, self.krylov_diagonal_scale = "fgmres", False
964 try:
965 self.apply_krylov_solver_settings(ksp, max_it=max_it, gmres_restart=gmres_restart,
966 monitor_convergence=False, rtol=rtol, atol=atol)
967 finally:
968 self.krylov_type, self.krylov_diagonal_scale = saved
969 self._apply_ksp_options(ksp, ksp_options)
970 return ksp
972 @staticmethod
973 def destroy_adjoint_operator(ksp):
974 """Destroy a KSP of linearized_solver and, for a matrix-free one, what
975 its operator's context owns."""
976 A = ksp.getOperators()[0]
977 ksp.destroy()
978 if A.getType() == PETSc.Mat.Type.PYTHON:
979 A.destroy() # PETSc calls the context's destroy(mat)
981 def release_derived_caches(self):
982 """Free what the derived fields cache between calls (Jacobians kept
983 while the state, alpha and mesh are unchanged). Collective."""
984 for f in self.derived_fields:
985 f.release_cache()
987 def solve_linearized(self, rhs, out, **kwargs):
988 """Solve dR/dU x = rhs once (see linearized_solver for the arguments):
989 returns the KSP (iterations, residual norm and reason are read off it)."""
990 ksp = self.linearized_solver(**kwargs)
991 ksp.solve(rhs, out)
992 return ksp
994 def _report_max_residual_location(self):
995 """
996 Print the magnitude, solution component and physical coordinate of
997 the largest entry in the current residual_vec_physical.
999 Collective: every rank participates in the allgather, and the
1000 lookup that decides which rank owns the max dof is derived from a
1001 collective Vec.max(), so all ranks agree on the answer by
1002 construction.
1003 """
1004 PETSc.Sys.Print("max physical residual entry: {}".format(
1005 self.residual_vec_physical.max()[1]))
1006 PETSc.Sys.Print("max physical residual index: {}".format(
1007 self.residual_vec_physical.max()[0]))
1009 max_residual_global_idx, max_residual_val = self.residual_vec_physical.max()
1011 bs = self.functionspaces["V"].dofmap.index_map_bs
1012 block_global_idx = max_residual_global_idx // bs
1013 component = max_residual_global_idx % bs
1015 block_local_idx = self.functionspaces["V"].dofmap.index_map.global_to_local(
1016 np.array([block_global_idx], dtype=np.int32))[0]
1018 comm = self.mesh_obj.mesh.comm
1019 if block_local_idx >= 0:
1020 dof_coords = self.functionspaces["V"].tabulate_dof_coordinates()
1021 payload = (comm.rank, dof_coords[block_local_idx, :].tolist(),
1022 int(component), max_residual_val)
1023 else:
1024 payload = None
1026 owner = next(p for p in comm.allgather(payload) if p is not None)
1027 if comm.rank == 0:
1028 rank, coord, component, val = owner
1029 PETSc.Sys.Print(
1030 f"max residual {val:.6e} on rank {rank}, component {component}, coord {coord}")
1032 def _report_positivity_limiting_node(self, x, dx, theta, initial_theta,
1033 min_theta=1e-6):
1034 """
1035 Print which node bounded the Newton step length.
1036 """
1037 if not self.report_positivity_limiter or theta >= initial_theta:
1038 return
1040 comm = self.mesh_obj.mesh.comm
1042 if theta > 0.0:
1043 theta_probe = 2.0*theta
1044 else:
1045 # theta == 0 means the bisection fell BELOW min_theta and gave up.
1046 # The smallest step it actually tested is the last halving still
1047 # >= min_theta, which is strictly larger than min_theta itself
1048 # (1.907e-6 for the defaults, not 1e-6). Probing at min_theta
1049 # would evaluate a step the limiter never rejected, and then
1050 # truthfully report finding no violator -- which reads as a bug in
1051 # the diagnostic rather than as the deadlock it actually is.
1052 theta_probe = initial_theta
1053 while theta_probe * 0.5 >= min_theta:
1054 theta_probe *= 0.5
1056 local = locate_positivity_limiting_nodes(
1057 x.array, dx.array, theta_probe, gamma=self.gamma,
1058 block_size=self.n_state, n_mean=self.n_mean if self.n_extra else None)
1060 n_density = comm.allreduce(local["n_density"], op=MPI.SUM)
1061 n_pressure = comm.allreduce(local["n_pressure"], op=MPI.SUM)
1062 n_nonfinite = comm.allreduce(local["n_nonfinite"], op=MPI.SUM)
1064 # Same allgather-and-pick idiom as _report_max_residual_location. The
1065 # payload is a handful of floats, mpi4py returns it in rank order, and
1066 # `min` is stable, so every rank selects the same node.
1067 payload = None
1068 if local["block"] is not None:
1069 dof_coords = self.functionspaces["V"].tabulate_dof_coordinates()
1070 payload = (local["severity"], comm.rank,
1071 dof_coords[local["block"], :].tolist(), local)
1072 gathered = [p for p in comm.allgather(payload) if p is not None]
1074 PETSc.Sys.Print(
1075 "[theta-limit] theta {:.6e} (step-length cap {:.6e}); {} node(s) block "
1076 "the step at theta = {:.6e} -- {} density, {} pressure, {} non-finite".format(
1077 theta, initial_theta, n_density + n_pressure + n_nonfinite,
1078 theta_probe, n_density, n_pressure, n_nonfinite))
1080 if not gathered:
1081 # theta was cut, yet nothing violates at twice the accepted value.
1082 # Not reachable from the bisection alone, so it means x or dx moved
1083 # between the limiter call and this one.
1084 PETSc.Sys.Print(
1085 "[theta-limit] no violating node found at the probe step -- "
1086 "state or update changed since theta was computed")
1087 return
1089 _, rank, coord, info = min(gathered, key=lambda p: p[0])
1090 PETSc.Sys.Print(
1091 "[theta-limit] worst: rank {}, local block {}, coord {}, binding "
1092 "constraint: {}".format(rank, info["block"], coord, info["constraint"]))
1093 PETSc.Sys.Print(
1094 "[theta-limit] state rho = {:.6e}, p = {:.6e} -> trial rho = {:.6e}, "
1095 "p = {:.6e}".format(
1096 info["rho"], info["p"], info["rho_trial"], info["p_trial"]))
1097 PETSc.Sys.Print(
1098 "[theta-limit] update d_rho = {:.6e}, |d_rhou| = {:.6e}, "
1099 "d_rhoE = {:.6e}".format(
1100 info["d_rho"], info["d_rhou_norm"], info["d_rhoE"]))
1102 def incoming_state_converged(self):
1103 """
1104 True, with solver_converged set, if the current u_vec already
1105 satisfies residual_conv_limit. Needs only the residual form, not the
1106 solver: the caller can skip set_up_solver and solve_system. The
1107 incoming state routinely is converged, since every evaluation that
1108 re-solves at an unchanged design arrives here converged, and a full
1109 Newton step would walk it backwards. Leaves the residual in
1110 residual_vec_physical either way. Collective.
1111 """
1112 self._ensure_residual_vector()
1113 residual_norm = self._assemble_physical_residual()
1114 self.solver_converged = bool(residual_norm < self.residual_conv_limit)
1115 if self.solver_converged:
1116 PETSc.Sys.Print(
1117 "Incoming state already converged ({:.6e} < {:.1e}); "
1118 "skipping the solve.".format(residual_norm, self.residual_conv_limit))
1119 self.u_vec.x.scatter_forward()
1120 return self.solver_converged
1122 def solve_system(self):
1123 """
1124 Newton solve of the Euler system, starting from the current u_vec.
1126 SNESNewtonSolver runs in passes of fom_newton_inner_max_it steps, and
1127 the physical residual is tested after each pass, until it drops below
1128 residual_conv_limit or the max_newton_iterations budget is spent.
1129 With use_ptc it is instead a pseudo-transient continuation solve
1130 (TimeIntegrator.solve_steady), max_newton_iterations pseudo steps of
1131 one Newton step each.
1132 Returns the physical residual vector (residual_vec_physical).
1133 """
1134 if self.use_ptc:
1135 return self.time_integrator.solve_steady()
1137 if self.incoming_state_converged():
1138 return self.residual_vec_physical
1140 solver = self.solver
1142 conv_limit = self.residual_conv_limit
1143 # FOM: max_newton_iterations is a TOTAL Newton-step budget, and each
1144 # outer pass now runs fom_newton_inner_max_it steps, so divide to hold
1145 # the product fixed -- otherwise raising the inner count silently
1146 # multiplies the budget and any gain is confounded with 5x the work.
1147 max_outer = (max(1, int(self.max_newton_iterations)
1148 // max(1, int(self.fom_newton_inner_max_it))))
1150 residual_norm = self.residual_vec_physical.norm()
1151 self.solver_converged = False
1153 for _ in range(max_outer):
1154 solver.solve(self.u_vec)
1155 self.u_vec.x.scatter_forward()
1156 residual_norm = self._assemble_physical_residual()
1157 self._report_max_residual_location()
1158 if not np.isfinite(residual_norm):
1159 PETSc.Sys.Print("Non-finite residual; stopping.")
1160 break
1161 if residual_norm < conv_limit:
1162 self.solver_converged = True
1163 break
1164 else:
1165 PETSc.Sys.Print(
1166 "WARNING: {} Newton iterations exhausted ({} passes x {}) "
1167 "at {:.6e}".format(
1168 max_outer * (int(self.fom_newton_inner_max_it)),
1169 max_outer,
1170 int(self.fom_newton_inner_max_it),
1171 residual_norm))
1173 self.u_vec.x.scatter_forward()
1174 return self.residual_vec_physical
1176 def set_u_vec(self, u_vec_entries):
1177 # set `self.u_vec` to the input linearization point
1179 set_petsc_vec_array(self.u_vec.x.petsc_vec, u_vec_entries)
1180 self.u_vec.x.petsc_vec.ghostUpdate()
1182 def _ensure_residual_vector(self):
1183 if not hasattr(self, 'residual_vec_physical'):
1184 self.residual_vec_physical = dolfinx.fem.petsc.create_vector(
1185 dolfinx.fem.extract_function_spaces(self.F_physical_form))
1187 # ------------------------------------------------------------------
1188 # Derived fields (utils/derived_fields.py)
1189 # ------------------------------------------------------------------
1190 def update_derived_fields(self):
1191 """Bring every derived field up to date with the current state.
1192 Collective; a no-op without derived fields."""
1193 for f in self.derived_fields:
1194 f.update()
1196 def update_geometry_fields(self):
1197 """Refresh what depends on the mesh coordinates alone (cell centres,
1198 wall distance, the reconstruction's cell volumes). Runs automatically
1199 after every Mesh.apply_node_motions / reset_nodes once the model has
1200 derived fields (see compute_weakform). A no-op for the plain Euler
1201 model."""
1202 for f in self.derived_fields:
1203 f.update_geometry()
1205 @property
1206 def _chain(self):
1207 if getattr(self, "_chain_rule", None) is None:
1208 from dragonfly_sim.utils.derived_fields import DerivedFieldChainRule
1209 self._chain_rule = DerivedFieldChainRule(
1210 self.mesh_obj, self.u_vec, self.derived_fields,
1211 alpha_var=getattr(self, "alpha_var", None),
1212 jit_options={"cffi_extra_compile_args": ["-O3", "-march=native", "-ffast-math"]})
1213 return self._chain_rule
1215 def release_solver(self):
1216 """
1217 Free the forward solver (SNES, KSP and its preconditioner), the
1218 Newton problem's matrices, and the cached condensed dR/dU of
1219 assemble_dRdu_mat. set_up_solver and assemble_dRdu_mat rebuild them
1220 when next needed. For memory-bound runs, e.g. a 3D RANS shape
1221 optimization: the forward solve's matrices and the adjoint's then
1222 never coexist. Collective.
1223 """
1224 if self.solver_is_set_up:
1225 self.solver.destroy()
1226 condensed = getattr(self.problem, "_condensed", None)
1227 if condensed is not None:
1228 self._destroy_condensed(condensed)
1229 self.problem._A.destroy()
1230 self.problem._b.destroy()
1231 self.problem = None
1232 self.solver = None
1233 self.solver_is_set_up = False
1234 self._release_condensed_jacobian()
1235 self._return_freed_memory()
1237 def _return_freed_memory(self):
1238 """Collect garbage, flush petsc4py's queue of dropped parallel objects
1239 and hand freed heap back to the OS (glibc malloc_trim). Rebuilding the
1240 multi-GB 3D RANS matrices on every solve otherwise fragments the heap,
1241 and the resident size creeps up ~2 GB per optimization iteration."""
1242 gc.collect()
1243 PETSc.garbage_cleanup(self.mesh_obj.mesh.comm)
1244 try:
1245 ctypes.CDLL("libc.so.6").malloc_trim(0)
1246 except (OSError, AttributeError):
1247 pass # not glibc
1249 @staticmethod
1250 def _destroy_condensed(condensed):
1251 # the matrices a CondensedJacobian owns; _inv_mass is the reconstruction's
1252 for name in ("A_UU", "A_UG", "B_GU", "_C", "A"):
1253 getattr(condensed, name).destroy()
1255 def _release_condensed_jacobian(self):
1256 """Free assemble_dRdu_mat's cached condensed dR/dU; rebuilt when next needed."""
1257 if hasattr(self, "_dRdu_condensed"):
1258 self._destroy_condensed(self._dRdu_condensed)
1259 del self._dRdu_condensed
1261 @property
1262 def has_condensed_jacobian(self):
1263 """True when dR/dU is a condensed product (a derived field depends on
1264 U), which the default ASM/ILU-on-dR/dU solves cannot precondition;
1265 the linearized solves then go through solve_linearized."""
1266 return any(f.depends_on_state for f in self.derived_fields)
1268 @staticmethod
1269 def _form_key(form):
1270 """Cache key of a form: its signature AND the identity of every
1271 coefficient and constant in it. The signature alone is the same for
1272 two forms that differ only in which Constant they hold (e.g. the lift
1273 and drag directions), and a compiled form is bound to its objects."""
1274 return (form.signature(), tuple(id(c) for c in form.coefficients()),
1275 tuple(id(c) for c in form.constants()))
1277 def state_derivative(self, form):
1278 """dJ/dU of a scalar functional J (rank-0 UFL form), through the
1279 derived fields: a Vec in the state layout."""
1280 self.update_derived_fields()
1281 return self._chain.functional_state(form, self._form_key(form))
1283 def alpha_derivative(self, form):
1284 """dJ/dalpha of a scalar functional J at fixed U: its explicit alpha
1285 dependence plus that of the derived fields (e.g. a reconstructed
1286 gradient reading the alpha-dependent far-field states). A number.
1287 Collective."""
1288 self.update_derived_fields()
1289 return self._chain.functional_alpha(form, self._form_key(form))
1291 def mesh_derivative(self, form):
1292 """dJ/dx of a scalar functional J through the derived fields: a Vec on
1293 the mesh's coordinate_space."""
1294 self.update_derived_fields()
1295 return self._chain.functional_mesh(form, self._form_key(form))
1297 def compute_dRdu_mat(self, u_vec_entries=None):
1298 """
1299 The forward Newton Jacobian, dR/du, as a UFL expression.
1300 """
1301 if u_vec_entries is not None:
1302 self.set_u_vec(u_vec_entries)
1303 return ufl.derivative(self.F, self.u_vec)
1305 def assemble_dRdu_mat(self, u_vec_entries=None):
1306 """
1307 The forward Newton Jacobian, assembled.
1309 Caches its own form/matrix, mirroring compute_dRdxmesh_mat's
1310 convention. With a reconstructed gradient this is the exact
1311 condensed dR/dU of the steady residual (an aij Mat).
1312 """
1313 if u_vec_entries is not None:
1314 self.set_u_vec(u_vec_entries)
1316 if self.has_condensed_jacobian:
1317 self.update_derived_fields()
1318 if not hasattr(self, "_dRdu_condensed"):
1319 self._dRdu_condensed = self.reconstruction.condensed_jacobian(self.F)
1320 return self._dRdu_condensed.assemble()
1322 if not hasattr(self, "dRdu_form"):
1323 self.dRdu_form = dolfinx.fem.form(self.compute_dRdu_mat())
1324 self.dRdu_mat = dolfinx.fem.petsc.create_matrix(self.dRdu_form)
1326 self.dRdu_mat.zeroEntries()
1327 dolfinx.fem.petsc.assemble_matrix(self.dRdu_mat, self.dRdu_form)
1328 self.dRdu_mat.assemble()
1330 return self.dRdu_mat
1332 def compute_dRdxmesh_mat(self, u_vec_entries, V):
1333 self.set_u_vec(u_vec_entries)
1335 if self.derived_fields:
1336 # dR/dx through the derived fields, as a matrix-free operator
1337 # (mult / multTranspose); V must be the mesh's coordinate_space
1338 self.update_derived_fields()
1339 return self._chain.residual_mesh_operator(self.F, "F")
1341 if not hasattr(self, 'dRdxmesh_form'):
1342 # test out derivative calculation
1343 spatial_coordinate = ufl.SpatialCoordinate(self.mesh_obj.mesh)
1345 args = self.F.arguments()
1346 # # UFL arguments need unique indices within a form
1347 n = max(a.number() for a in args) if args else -1
1348 dX = ufl.Argument(V, n + 1)
1350 a_deriv = ufl.derivative(self.F, spatial_coordinate, dX)
1351 self.dRdxmesh_form = dolfinx.fem.form(a_deriv)
1352 self.dRdxmesh_mat = dolfinx.fem.petsc.create_matrix(self.dRdxmesh_form)
1354 self.dRdxmesh_mat.zeroEntries()
1355 dolfinx.fem.petsc.assemble_matrix(self.dRdxmesh_mat, self.dRdxmesh_form)
1356 self.dRdxmesh_mat.assemble()
1358 return self.dRdxmesh_mat