Coverage for dragonfly_sim/utils/solver_utils.py: 94%
113 statements
« prev ^ index » next coverage.py v7.16.2, created at 2026-10-17 03:37 +0000
1from petsc4py import PETSc
2import numpy as np
5_KSP_REASON_NAMES = None
8_KSP_NORM_TYPE_NAMES = None
11def _enum_name_map(enum_cls):
12 """
13 Reverse map {code: name} for a petsc4py enum-like class.
15 petsc4py exposes most codes under several aliases -- PC.Side has both "L"
16 and "LEFT" for 0, KSP.ConvergedReason both "DIVERGED_PC_FAILED" and
17 "DIVERGED_PCSETUP_FAILED" for -11 -- so the map has to choose. Longest name
18 first, ties broken alphabetically: longest picks the descriptive alias over
19 the terse one (which is the whole point of printing a name instead of the
20 integer), and the tie-break makes the result identical on every run rather
21 than dependent on `dir()` ordering.
22 """
23 names = {}
24 for name in sorted((n for n in dir(enum_cls) if not n.startswith("_")),
25 key=lambda n: (-len(n), n)):
26 value = getattr(enum_cls, name)
27 if isinstance(value, int):
28 names.setdefault(value, name)
29 return names
32def ksp_converged_reason_name(reason):
33 """
34 Human-readable name for a PETSc KSP convergence reason code.
36 Worth printing next to the iteration count because the count alone cannot
37 distinguish the two failures that call for opposite responses:
38 DIVERGED_MAX_IT means a hard but well-posed operator that simply exhausted
39 its Krylov budget (tune the solver), while DIVERGED_BREAKDOWN or
40 DIVERGED_PCSETUP_FAILED means the Arnoldi process or the preconditioner
41 itself failed (the operator or the coarse space is broken). Both show up as
42 "ran to max_it" in a log that reports only iterations.
44 Unknown codes come back as "REASON(<n>)" rather than raising -- this is a
45 diagnostic and must never be the thing that breaks a solve.
46 """
47 global _KSP_REASON_NAMES
48 if _KSP_REASON_NAMES is None:
49 _KSP_REASON_NAMES = _enum_name_map(PETSc.KSP.ConvergedReason)
50 return _KSP_REASON_NAMES.get(int(reason), "REASON({})".format(reason))
53def ksp_norm_type_name(norm_type):
54 """
55 Human-readable name for a PETSc KSP norm type.
57 This is what says whether a reported residual norm is the true ||b - Ax||
58 or the preconditioned one, which decides whether a norm LARGER than ||b||
59 is impossible (UNPRECONDITIONED, as fgmres defaults to) or unremarkable
60 (PRECONDITIONED, as gmres defaults to). Without it the number cannot be
61 read at all.
62 """
63 global _KSP_NORM_TYPE_NAMES
64 if _KSP_NORM_TYPE_NAMES is None:
65 _KSP_NORM_TYPE_NAMES = _enum_name_map(PETSc.KSP.NormType)
66 return _KSP_NORM_TYPE_NAMES.get(int(norm_type), "NORM({})".format(norm_type))
69def ksp_pc_side_name(pc_side):
70 """Human-readable name for a PETSc preconditioner side (LEFT/RIGHT/SYMMETRIC)."""
71 return _enum_name_map(PETSc.PC.Side).get(
72 int(pc_side), "SIDE({})".format(pc_side))
75# Floor on density and pressure enforced by the positivity limiter below.
76# Shared with locate_positivity_limiting_nodes so the diagnostic reports on
77# exactly the constraint the limiter applied -- a diagnostic running against a
78# floor of its own would name the wrong nodes as soon as the two drifted.
79# Lower bound density and pressure are held above. A state is REJECTED at
80# p <= POSITIVITY_FLOOR, so this is not merely a guard against negative
81# pressure -- it reserves a margin, and a solve whose transient legitimately
82# passes through a low-pressure state is blocked by it.
83POSITIVITY_FLOOR = 1e-4
86# Smallest step length the halving search will accept. Below this the search
87# gives up and returns theta = 0.0, freezing the state: see the tail of
88# compute_positivity_preserving_theta. Module-level rather than a keyword
89# default because no call site passes min_theta -- every path takes the
90# default -- so this is the only place it can be varied.
91MIN_POSITIVITY_THETA = 1e-6
94def _mean_flow_blocks(x_arr, dx_arr, block_size, n_mean):
95 """
96 Drop the trailing transported scalars (turbulence variables) from every
97 block, so the density/pressure checks below see (rho, rho u, rho E) only.
98 Returns (x, dx, block size); a no-op when n_mean is None or the blocks
99 are already mean flow.
100 """
101 if n_mean is None or n_mean >= block_size:
102 return x_arr, dx_arr, block_size
103 n_blocks = len(x_arr) // block_size
104 x = x_arr[:n_blocks * block_size].reshape(-1, block_size)[:, :n_mean].ravel()
105 dx = dx_arr[:n_blocks * block_size].reshape(-1, block_size)[:, :n_mean].ravel()
106 return x, dx, n_mean
109def compute_positivity_preserving_theta(comm, x_arr, dx_arr, gamma=1.4, initial_theta=1.0, min_theta=None, block_size=4,
110 n_mean=None):
111 """
112 Computes the maximum theta <= initial_theta such that x_new = x_arr - theta * dx_arr
113 has positive mass density and positive pressure everywhere.
115 Returns the step length only. To find out WHICH nodes bounded it -- the
116 question a collapsed theta raises -- pass twice the returned value to
117 locate_positivity_limiting_nodes.
119 With n_mean given, each block of block_size entries is (rho, rho u, rho E)
120 followed by transported scalars, which are not limited.
121 """
122 import numpy as np
123 if min_theta is None:
124 min_theta = MIN_POSITIVITY_THETA
125 x_arr, dx_arr, block_size = _mean_flow_blocks(x_arr, dx_arr, block_size, n_mean)
126 theta = initial_theta
128 # Determine the number of valid full blocks
129 n_blocks = len(x_arr) // block_size
130 if n_blocks == 0:
131 # No local degrees of freedom
132 local_theta = theta
133 else:
134 x_nodes = x_arr[:n_blocks*block_size].reshape(-1, block_size)
135 dx_nodes = dx_arr[:n_blocks*block_size].reshape(-1, block_size)
137 rho = x_nodes[:, 0]
138 rhou = x_nodes[:, 1:block_size-1]
139 rhoE = x_nodes[:, block_size-1]
141 d_rho = dx_nodes[:, 0]
142 d_rhou = dx_nodes[:, 1:block_size-1]
143 d_rhoE = dx_nodes[:, block_size-1]
145 local_theta = theta
146 while local_theta >= min_theta:
147 rho_new = rho - local_theta * d_rho
148 rhou_new = rhou - local_theta * d_rhou
149 rhoE_new = rhoE - local_theta * d_rhoE
151 # Check density positivity
152 if np.any(rho_new <= POSITIVITY_FLOOR):
153 local_theta *= 0.5
154 continue
156 # Check pressure positivity
157 kin_energy_new = 0.5 * np.sum(rhou_new**2, axis=1) / rho_new
158 p_new = (gamma - 1.0) * (rhoE_new - kin_energy_new)
160 if np.any(p_new <= POSITIVITY_FLOOR):
161 local_theta *= 0.5
162 continue
164 break
166 if local_theta < min_theta:
167 local_theta = 0.0
169 if comm.Get_size() > 1:
170 from mpi4py import MPI
171 global_theta = comm.allreduce(local_theta, op=MPI.MIN)
172 else:
173 global_theta = local_theta
175 return global_theta
178def locate_positivity_limiting_nodes(x_arr, dx_arr, theta_probe, gamma=1.4,
179 block_size=4, floor=POSITIVITY_FLOOR, n_mean=None):
180 """
181 Find which local nodes make x_arr - theta_probe*dx_arr unphysical.
183 Diagnostic companion to compute_positivity_preserving_theta, which returns
184 the surviving step length and says nothing about what bounded it. Pass
185 theta_probe = twice the accepted theta -- the smallest step that bisection
186 actually rejected -- and the nodes reported here are exactly the ones that
187 forced the final halving.
189 Rank-local and pure numpy: the caller owns the MPI reduction and the
190 dof-coordinate lookup. `block` is an index into the BLOCKS of x_arr, which
191 for a blocked function space is also a row index into that space's local
192 dof coordinates.
194 Returns a dict:
195 n_density, n_pressure, n_nonfinite
196 local counts of nodes violating each way, for the caller to sum
197 across ranks. One violating node out of millions and a whole
198 region going unphysical at once are different problems, and the
199 count is what tells them apart.
200 block
201 block index of this rank's WORST violator, or None if this rank
202 has none -- in which case the keys below are absent.
203 severity
204 trial value over current value at that node, so that a density
205 and a pressure violation are comparable and one argmin ranks the
206 whole set; more negative is a worse overshoot, -inf for a
207 non-finite trial state.
208 constraint
209 "density", "pressure" or "non-finite".
210 rho, p, rho_trial, p_trial, d_rho, d_rhou_norm, d_rhoE
211 the state, the trial state and the update at that node.
212 """
213 x_arr, dx_arr, block_size = _mean_flow_blocks(x_arr, dx_arr, block_size, n_mean)
214 n_blocks = len(x_arr) // block_size
215 empty = {"n_density": 0, "n_pressure": 0, "n_nonfinite": 0, "block": None}
216 if n_blocks == 0:
217 # No local degrees of freedom.
218 return empty
220 x_nodes = x_arr[:n_blocks*block_size].reshape(-1, block_size)
221 dx_nodes = dx_arr[:n_blocks*block_size].reshape(-1, block_size)
222 trial = x_nodes - theta_probe * dx_nodes
224 rho = x_nodes[:, 0]
225 rhou = x_nodes[:, 1:block_size-1]
226 rhoE = x_nodes[:, block_size-1]
228 rho_trial = trial[:, 0]
229 rhou_trial = trial[:, 1:block_size-1]
230 rhoE_trial = trial[:, block_size-1]
232 # Checked first and excluded from the two comparisons below, because a NaN
233 # fails `<= floor` and would otherwise be silently counted as healthy --
234 # the one state that most needs reporting.
235 nonfinite = ~(np.isfinite(rho_trial) & np.isfinite(rhoE_trial)
236 & np.all(np.isfinite(rhou_trial), axis=1))
238 bad_rho = (~nonfinite) & (rho_trial <= floor)
240 # Pressure is only meaningful where the trial density survived; where it
241 # did not, the density violation is the one worth reporting anyway.
242 rho_trial_safe = np.where(rho_trial > floor, rho_trial, 1.0)
243 p_trial = (gamma - 1.0) * (rhoE_trial
244 - 0.5*np.sum(rhou_trial**2, axis=1) / rho_trial_safe)
245 bad_p = (~nonfinite) & (~bad_rho) & (p_trial <= floor)
247 result = dict(empty)
248 result["n_density"] = int(np.count_nonzero(bad_rho))
249 result["n_pressure"] = int(np.count_nonzero(bad_p))
250 result["n_nonfinite"] = int(np.count_nonzero(nonfinite))
252 violating = np.flatnonzero(nonfinite | bad_rho | bad_p)
253 if violating.size == 0:
254 return result
256 # The incoming state has already passed the limiter, so rho > floor there;
257 # the guard is only so a caller handing in a blown-up state still gets a
258 # report rather than a divide-by-zero.
259 rho_safe = np.where(rho > floor, rho, 1.0)
260 p = (gamma - 1.0) * (rhoE - 0.5*np.sum(rhou**2, axis=1) / rho_safe)
262 reference = np.where(bad_rho, np.abs(rho), np.abs(p))
263 reference = np.where(reference > 0.0, reference, 1.0)
264 severity = np.where(bad_rho, rho_trial, p_trial) / reference
265 severity = np.where(nonfinite, -np.inf, severity)
267 worst = int(violating[np.argmin(severity[violating])])
269 result["block"] = worst
270 result["severity"] = float(severity[worst])
271 result["constraint"] = ("non-finite" if nonfinite[worst]
272 else "density" if bad_rho[worst] else "pressure")
273 result["rho"] = float(rho[worst])
274 result["p"] = float(p[worst])
275 result["rho_trial"] = float(rho_trial[worst])
276 result["p_trial"] = float(p_trial[worst])
277 result["d_rho"] = float(dx_nodes[worst, 0])
278 result["d_rhou_norm"] = float(np.linalg.norm(dx_nodes[worst, 1:block_size-1]))
279 result["d_rhoE"] = float(dx_nodes[worst, block_size-1])
280 return result