diff --git a/docs/advanced/eulerian-advection-diffusion.md b/docs/advanced/eulerian-advection-diffusion.md index 2b7a1baf..dc114ad0 100644 --- a/docs/advanced/eulerian-advection-diffusion.md +++ b/docs/advanced/eulerian-advection-diffusion.md @@ -38,6 +38,10 @@ that sizes its step in Courant numbers can still ask for ## What carries over +The comparisons and implicit timestep policies below describe the default +`EulerianSUPG` manager. The optional `EulerianSUPGPC` manager has a different +update and timestep policy; see [Predictor-corrector transport](#predictor-corrector-transport). + | SLCN | SUPG | note | |---|---|---| | `order=1, theta=0.5` | same | Crank-Nicolson, the default for both | @@ -146,6 +150,143 @@ field here, the traced-back flux there). The manager works for a vector or tenso as well (`vtype`), applying the advection component by component, which is how the Navier-Stokes solver and a transported stress use it. +## Predictor-corrector transport + +For continuous-P1 scalar transport, select the predictor-corrector manager +explicitly and pass it to the same `uw.systems.AdvDiffusion` solver: + +```python +Tdot = uw.discretisation.MeshVariable("Tdot", mesh, 1, degree=1) +transport = uw.systems.ddt.EulerianSUPGPC( + mesh, T, v.sym, + method="citcoms", + temperature_rate_field=Tdot, +) +adv = uw.systems.AdvDiffusion(mesh, T, v.sym, DuDt=transport) +adv.constitutive_model = uw.constitutive_models.DiffusionModel +adv.constitutive_model.Parameters.diffusivity = 1.0 +adv.add_dirichlet_bc(0.0, "Upper") +adv.add_dirichlet_bc(1.0, "Lower") +adv.solve(timestep=adv.estimate_dt()) +``` + +Here `T` must also be continuous P1. The manager owns the rate, startup state, +correction controls and transport policy; the solver owns the constitutive +model, source and boundary conditions. Select `method="citcoms"` for the +fixed-correction benchmark method or `method="pc_converged"` for a +residual-converged reference. These are manager methods, not solver time +integrator arguments. CN and BDF remain on the default `EulerianSUPG` manager. + +### Predictor and corrections + +Writing the temperature rate as $q$, the predictor is +$T^{(0)} = T^n + (1-\gamma)\Delta t\,q^n$, followed by resetting $q$ to zero. +Each correction assembles the full finite-element residual $F(T,q)$ and applies + +$$ +\delta q = -D^{-1}F(T,q), \qquad +q \leftarrow q + \delta q, \qquad +T \leftarrow T + \gamma\Delta t\,\delta q, +$$ + +where $D$ is the positive row-lumped mass. Dirichlet values are reinserted at +each correction. `method="citcoms"` defaults to `adv_gamma=0.5` and +`corrector_steps=2`, with a single lumped correction to initialise the rate. + +Both PC methods use the steady directional simplex stabilisation + +$$ +\tau = \frac{h}{2|\mathbf{u}|}\max(0,1-1/Pe), \qquad +Pe = \frac{|\mathbf{u}|h}{2\kappa}, \qquad +h = \frac{2|\mathbf{u}|}{\sum_a |\mathbf{u}\cdot\nabla N_a|}. +$$ + +Zero velocity gives zero tau; zero diffusivity uses the advective limit. +This is not the default manager's transient norm tau or cell-Peclet weighting. +`tau=None` selects this automatic rule; a scalar symbolic expression or number +overrides it. `supg_weight=1.0` is the PC default. +Automatic geometry supports 2-D triangles and 3-D tetrahedra and currently +requires a non-empty volume partition on every rank. Unsupported layouts +are rejected collectively; use fewer ranks or a sufficiently resolved mesh. + +After attachment to the solver, `transport.estimate_dt()` returns +`0.9*min(dt_adv, dt_diff)`, using the directional advective rate and a row-sum +bound on the lumped diffusion operator; `adv.estimate_dt()` delegates to it. +The implicit field-change estimate described above is not a stability bound +for this update. Fixed comparison timesteps must respect the PC bound; +residual convergence does not make diagonal correction converge for arbitrary +steps. SUPG does not guarantee a nodal maximum principle: check temperature +bounds and heat balance. + +Diffusion remains in the Galerkin flux but is absent from the strong SUPG +residual. This omission is exact for affine P1 fields with elementwise +constant diffusivity, not for arbitrary curved mappings, variable +coefficients or P2 temperature. + +### Finite-correction accuracy + +The correction mass is lumped, but the time derivative in the residual uses +the consistent finite-element mass. Consequently `adv_gamma=0.5` and two +corrections do **not** guarantee second-order temporal convergence for a +nonuniform field at fixed mesh. For pure diffusion, let $M$ be the consistent +mass, $K$ the stiffness and $D=\operatorname{diag}(M\mathbf{1})$. Two corrections +approach the operator $(2I-D^{-1}M)D^{-1}K$ as $\Delta t$ vanishes, generally +different from both $M^{-1}K$ and $D^{-1}K$. The startup rate $-D^{-1}KT$ is +also only an approximation to the consistent rate $-M^{-1}KT$. + +The regression in `tests/test_1118_pc2_diffusion_time.py` isolates these +effects with independently integrated element +matrices and exact discrete eigenmode/matrix-exponential solutions on tiny +triangular and tetrahedral meshes. It records first-order timestep +differences in serial and MPI. Uniform scalar decay, where the two masses +agree, is not sufficient evidence of PDE time accuracy. The same reference +records temporal order 2.00 for an actual UW3 consistent-mass CN update in +both geometries, in serial and on eight ranks, with its nodal amplification +map agreeing within 1.6e-14. The DDt-manager migration reproduced these +results on 8 September 2026. These are isolated numerical checks, not +production-scale validation. + +### Residual-converged reference + +Construct `EulerianSUPGPC` with `method="pc_converged"` to keep the same SUPG +residual and gamma update while using the lumped mass only as an iterative +preconditioner. At startup it converges the consistent Petrov-Galerkin rate +equation with temperature held fixed; after prediction it converges the +coupled rate/temperature correction. The full residual must be no larger +than `max(corrector_atol, corrector_rtol*initial_residual)`. Defaults are +`corrector_rtol=1e-10`, `corrector_atol=1e-12` and +`max_corrector_steps=100`. Non-convergence raises `RuntimeError` instead of +accepting the step. Inspect `transport.temperature_rate`, +`transport.last_corrector_iterations`, `transport.last_corrector_residual` +and `transport.corrector_target` on the manager. + +With `adv_gamma=0.5`, this supplies a separate second-order reference rather +than changing the fixed-correction CitcomS method. The discrete diffusion +regression records order 2.00 in 2-D and 3-D, in serial and on +eight ranks, and agreement with the trapezoidal amplification map below +5.2e-14. At relative tolerance `1e-12`, those small meshes needed 48-63 +corrections per step in 2-D and 63-81 in 3-D. This is an accuracy reference, +not evidence that diagonal iteration is the most efficient production +consistent-mass solve. Changing residual mass or correction count changes +the fixed-correction method and must not be presented as unchanged paper +reproduction. + +### Checkpoint state + +```python +orchestration_model = uw.get_default_model() +orchestration_model.save_state(file="checkpoint.h5") +# Reconstruct the matching model, fields and manager before loading. +orchestration_model.load_state("checkpoint.h5") +``` + +An exact PC restart needs temperature plus the manager's rate, startup state +and correction controls. A temperature-only checkpoint is insufficient. +Implicit integration instead needs its DDt fields, timestep history, theta +and field-change estimator state. Disk snapshots require the same model +layout and MPI rank count; snapshots with a different solver/manager layout +require migration. The manager owns PC restart state, not a solver alias. + ## Further reading - Design note and measurements: `docs/developer/design/eulerian-supg-transport.md` diff --git a/docs/advanced/supg-transport.md b/docs/advanced/supg-transport.md new file mode 100644 index 00000000..6338899d --- /dev/null +++ b/docs/advanced/supg-transport.md @@ -0,0 +1,9 @@ +# SUPG Scalar Transport + +This page is obsolete. See +[Eulerian advection-diffusion](eulerian-advection-diffusion.md) for +`uw.systems.AdvDiffusion` and its transport managers: the default +`uw.systems.ddt.EulerianSUPG` for implicit CN/BDF transport and +`uw.systems.ddt.EulerianSUPGPC` for `method="citcoms"` or +`method="pc_converged"`, supplied through `DuDt=`. The guide covers the PC +algorithm, accuracy limitations, timestep policies and checkpoint state. diff --git a/docs/developer/CHANGELOG.md b/docs/developer/CHANGELOG.md index 5aaa5696..c13aad44 100644 --- a/docs/developer/CHANGELOG.md +++ b/docs/developer/CHANGELOG.md @@ -6,6 +6,17 @@ This log tracks significant development work at a conceptual level, suitable for ## 2026 Q3 (July – September) +### Predictor-Corrector Transport Manager (September 2026, #689) + +`uw.systems.ddt.EulerianSUPGPC`, supplied as `DuDt=` to +`uw.systems.AdvDiffusion`, owns the P1 predictor-corrector update, rate state +and restart controls. `method="citcoms"` retains fixed corrections and the +explicit timestep bound; `method="pc_converged"` provides a residual-converged +accuracy reference. The default `EulerianSUPG` CN/BDF transport is unchanged. +The [transport guide](../advanced/eulerian-advection-diffusion.md) explains +why two lumped corrections do not generally establish second-order time +accuracy and records the separate consistent-mass reference measurements. + ### The Multiplier Was Not the Whole Traction (August 2026) **`Stokes_Constrained.topography()` now returns the traction the boundary is diff --git a/docs/developer/design/eulerian-supg-transport.md b/docs/developer/design/eulerian-supg-transport.md index 39146431..9b248df5 100644 --- a/docs/developer/design/eulerian-supg-transport.md +++ b/docs/developer/design/eulerian-supg-transport.md @@ -772,3 +772,58 @@ garbage-collected variable leaves its PETSc field in the DM, and both the DM fields line up by position. Every later variable was then packed into, and read from, the wrong slots. Fixed in the same branch (pack by field name, offsets from the DM's field list) with `tests/test_1058_dropped_meshvariable_aux_layout.py`. + +## Predictor-corrector manager (#689) + +`uw.systems.ddt.EulerianSUPGPC(mesh, T, V_fn, method="citcoms", ...)` is a +scalar continuous-P1 transport manager supplied as `DuDt=` to +`uw.systems.AdvDiffusion`. `method="pc_converged"` selects a separate +residual-converged reference. The existing `EulerianSUPG` manager retains +CN/BDF, its transient tau and field-change timestep policy; the measurements +above describe that implicit path, not the PC update. + +The ownership follows the transport contract: the manager supplies the +rate time derivative, advection, stabilisation and current spatial state; +the solver assembles the constitutive diffusion flux, source and boundary +terms. A generic stepping hook lets the manager execute corrections using +the solver's assembled residual and boundary handling. PC-specific rate, +startup state, correction controls, geometry and reusable workspaces belong +to the manager, without solver-side method aliases or unused BDF history. + +For rate $q$, predict $T^{(0)}=T^n+(1-\gamma)\Delta t q^n$, reset $q=0$, +then apply $\delta q=-D^{-1}F(T,q)$, $q\leftarrow q+\delta q$ and +$T\leftarrow T+\gamma\Delta t\delta q$, reinserting Dirichlet values. +Here $D$ is positive row-lumped mass but $F$ retains the consistent +Petrov-Galerkin time derivative. `citcoms` defaults to `adv_gamma=0.5` and +`corrector_steps=2`, with one lumped startup correction. `pc_converged` +converges the rate equation at fixed temperature at startup and the coupled +correction after prediction, using $D$ only as a preconditioner. It stops at +`max(corrector_atol, corrector_rtol*initial_residual)` (defaults `1e-12`, +`1e-10`) or raises `RuntimeError` after `max_corrector_steps` (default 100). + +Both PC methods retain the steady directional simplex tau and the +`0.9*min(dt_adv, dt_diff)` timestep estimate. Automatic geometry is limited +to triangles/tetrahedra with non-empty volume partitions on every rank; +unsupported layouts must be rejected collectively. The missing strong +diffusion term is exact only for affine P1 with elementwise constant +diffusivity. Neither SUPG nor residual convergence guarantees a nodal +maximum principle or unrestricted diagonal-iteration timesteps. + +Finite corrections are not a consistent-mass solve: for pure diffusion, +two corrections approach $(2I-D^{-1}M)D^{-1}K$, not generally $M^{-1}K$, +as $\Delta t\to0$. A lumped startup rate adds another discrepancy. The +[user guide](../../advanced/eulerian-advection-diffusion.md#finite-correction-accuracy) +records the mathematical regression results: first-order timestep +differences for fixed corrections, and order 2.00 for consistent CN and +`pc_converged` at gamma 0.5 on tiny triangles/tetrahedra in serial and on +eight ranks. The DDt-manager migration reproduced these isolated results +on 8 September 2026; they are not coupled production acceptance. + +Restart registration must capture the manager's rate, startup flag and +correction controls as well as temperature. Derived PETSc workspaces can be +rebuilt; rate history cannot be replaced with zero on continuation. A +matching model/manager layout and MPI rank count are required for disk +replay. Migration checks should cover frozen numerical equivalence, exact +discrete time-order references, snapshot continuation in a fresh process, +and workspace reuse in serial and MPI, separately from coupled production +benchmarks. diff --git a/src/underworld3/systems/advection_diffusion_eulerian.py b/src/underworld3/systems/advection_diffusion_eulerian.py index 06d27494..24e9d5a2 100644 --- a/src/underworld3/systems/advection_diffusion_eulerian.py +++ b/src/underworld3/systems/advection_diffusion_eulerian.py @@ -33,6 +33,7 @@ """ import warnings +from dataclasses import dataclass import numpy as np import sympy @@ -40,6 +41,7 @@ import underworld3 as uw import underworld3.timing as timing +from underworld3.checkpoint.state import SnapshottableState from underworld3.systems import SNES_Scalar from underworld3.utilities._api_tools import Template from underworld3.function import expression as public_expression @@ -54,6 +56,16 @@ ) +@dataclass +class AdvDiffusionState(SnapshottableState): + """Solver timestep/estimator metadata; transport history belongs to DuDt.""" + + last_timestep: Optional[float] = None + last_change_rate: Optional[float] = None + order: int = 1 + theta: float = 0.5 + + def _check_supplied_manager(DuDt, order, theta): """A supplied history manager fixes the scheme: the arguments must agree with it.""" if DuDt.order != order: @@ -263,6 +275,7 @@ def __init__( "the semi-Lagrangian trace-back and the Eulerian scheme has none.", stacklevel=2, ) + requested_theta = theta order = int(order) if order not in (1, 2, 3): raise ValueError(f"order must be 1, 2 or 3, not {order}.") @@ -312,7 +325,7 @@ def __init__( raise TypeError(f"DuDt must be a DDt history manager, not {type(DuDt).__name__}.") if sympy.Matrix(DuDt.psi_fn).shape != u_Field.sym.shape: raise ValueError("DuDt tracks a different unknown from u_Field.") - _check_supplied_manager(DuDt, order, theta) + _check_supplied_manager(DuDt, order, requested_theta) self.Unknowns.DuDt = DuDt self._theta = float(getattr(self.DuDt, "theta", theta)) @@ -331,6 +344,8 @@ def __init__( self.petsc_options["snes_rtol"] = 1.0e-8 self.petsc_options["ksp_rtol"] = 1.0e-9 self.petsc_options["snes_max_it"] = 20 + self._bind_transport_manager(self.DuDt) + uw.get_default_model()._register_state_bearer(self) # ------------------------------------------------------------------ # Linear solver @@ -423,6 +438,27 @@ def _object_viewer(self): # Scheme description # ------------------------------------------------------------------ + def _bind_transport_manager(self, manager): + bind_transport = getattr(manager, "_bind_transport_solver", None) + if bind_transport is not None: + bind_transport(self) + + @property + def DuDt(self): + """Transport manager bound to this solver's unknown.""" + return self.Unknowns.DuDt + + @DuDt.setter + def DuDt(self, manager): + if not isinstance(manager, _DDtBase): + raise TypeError("DuDt must be a DDt transport manager.") + if sympy.Matrix(manager.psi_fn).shape != self.u.sym.shape: + raise ValueError("DuDt tracks a different unknown from u_Field.") + self._bind_transport_manager(manager) + self.Unknowns.DuDt = manager + self._theta = float(getattr(manager, "theta", self._theta)) + self._last_timestep = manager._dt + @property def integrator(self) -> str: """The multistep family in use: ``"am"`` (the theta rule) at order 1, ``"bdf"`` above.""" @@ -453,8 +489,8 @@ def theta(self, value): ) if not hasattr(self.DuDt, "theta"): raise AttributeError(f"{type(self.DuDt).__name__} has no theta to set.") - self._theta = value self.DuDt.theta = value + self._theta = value @property def delta_t(self): @@ -470,11 +506,11 @@ def delta_t(self): @delta_t.setter def delta_t(self, value): dt = float(_nondimensionalise_timestep(value)) - if dt <= 0.0: + if not np.isfinite(dt) or dt <= 0.0: raise ValueError(f"timestep must be positive, not {dt}.") if dt != self._last_timestep: - self.DuDt.delta_t.sym = dt self._last_timestep = dt + self.DuDt._dt = dt @property def V_fn(self): @@ -571,9 +607,12 @@ def _stabilisation_flux(self): # ------------------------------------------------------------------ @timing.routine_timer_decorator - def estimate_dt(self, fraction: float = 0.02, basis: str = "accuracy", + def estimate_dt(self, fraction: float = 0.02, basis: Optional[str] = None, direction_aware: bool = False, percentile: float = 0.0): - r"""A timestep for this scheme, chosen for accuracy. + r"""A timestep selected by the transport manager or implicit estimator. + + A rate-based manager supplies its stability estimate. The default + implicit manager uses the accuracy estimate described below. The implicit scheme has no stability limit, so the cell-crossing time the semi-Lagrangian solver reports says nothing about how large a step @@ -602,7 +641,9 @@ def estimate_dt(self, fraction: float = 0.02, basis: str = "accuracy", ---------- fraction : float, default 0.02 Allowed change of the field per step as a fraction of its range. - basis : {"accuracy", "resolution"} + basis : {None, "accuracy", "resolution", "stability"} + None selects the manager default; "stability" applies to + rate-based transport only. ``"resolution"`` returns the cell-crossing / diffusion time the semi-Lagrangian solver's ``estimate_dt`` returns, for scripts that size the step in Courant numbers. @@ -617,6 +658,13 @@ def estimate_dt(self, fraction: float = 0.02, basis: str = "accuracy", """ from mpi4py import MPI + self._bind_transport_manager(self.DuDt) + estimate_transport = getattr(self.DuDt, "_estimate_transport_dt", None) + if estimate_transport is not None: + return estimate_transport(fraction=fraction, basis=basis, + direction_aware=direction_aware, percentile=percentile) + if basis is None: + basis = "accuracy" if basis == "resolution": dt_estimate, dt_adv, dt_diff = _advective_diffusive_dt( self.constitutive_model.K, self.V_fn, self.mesh, @@ -665,6 +713,48 @@ def _advective_rate(self): local = 0.0 return uw.mpi.comm.allreduce(local, op=MPI.MAX) + @property + def state(self): + """Snapshot the timestep and measured change used by estimate_dt.""" + return AdvDiffusionState( + last_timestep=self._last_timestep, + last_change_rate=self._last_change_rate, + order=self.order, theta=self.theta, + ) + + @state.setter + def state(self, state): + if not isinstance(state, AdvDiffusionState): + raise TypeError("AdvDiffusion state has the wrong type.") + if state._schema_version != AdvDiffusionState._schema_version: + raise ValueError("AdvDiffusion state schema changed since snapshot.") + if state.order != self.order: + raise ValueError("AdvDiffusion order changed since snapshot.") + if hasattr(self.DuDt, "theta"): + self.theta = state.theta + self._last_timestep = None + if state.last_timestep is not None: + self.delta_t = state.last_timestep + self._last_change_rate = state.last_change_rate + + def _prepare_transport_residual(self, verbose=False): + """Build the solver-owned PDE residual and constrained scalar DM.""" + if not self.constitutive_model._solver_is_setup: + self._needs_function_rewire = True + self._build(verbose) + self.is_setup = True + self.constitutive_model._solver_is_setup = True + + def _compute_transport_residual(self, solution, residual): + """Assemble at the current unknown with refreshed auxiliary fields.""" + solution.set(0.0) + self.dm.localToGlobal(self.u.vec, solution, addv=False) + residual.set(0.0) + self.mesh.update_lvec() + self.dm.setAuxiliaryVec(self.mesh.lvec, None) + self._update_constants() + self.snes.computeFunction(solution, residual) + def solve( self, zero_init_guess: Optional[bool] = None, @@ -681,8 +771,11 @@ def solve( between calls updates a runtime constant of the compiled kernels; nothing is recompiled. """ + self._bind_transport_manager(self.DuDt) if timestep is not None: self.delta_t = timestep + elif self.DuDt._dt is not None: + self.delta_t = self.DuDt._dt elif self._last_timestep is None: raise ValueError( "solve() needs a timestep: pass timestep=
or set solver.delta_t first." @@ -691,6 +784,11 @@ def solve( if _force_setup: self._needs_function_rewire = True + solve_transport = getattr(self.DuDt, "_solve_transport", None) + if solve_transport is not None: + solve_transport(self, dt, zero_init_guess=zero_init_guess, + verbose=verbose, divergence_retries=divergence_retries) + return if not self.constitutive_model._solver_is_setup: self._needs_function_rewire = True # The base ``_build`` resolves the preconditioner choice against the diff --git a/src/underworld3/systems/ddt.py b/src/underworld3/systems/ddt.py index f82f58c4..2b432c8a 100644 --- a/src/underworld3/systems/ddt.py +++ b/src/underworld3/systems/ddt.py @@ -67,6 +67,7 @@ from underworld3.utilities._api_tools import uw_object from underworld3.utilities.unit_aware_array import UnitAwareArray from underworld3.checkpoint.state import SnapshottableState +from underworld3.systems.ddt_pc import _EulerianSUPGPCMethods from underworld3.discretisation.remesh import RemeshPolicy, remap_var_set from petsc4py import PETSc @@ -1673,6 +1674,8 @@ def __init__( ] self.supg_weight = supg_weight self.tau_weights = tau_weights + # Snapshots restore fields before the first residual is built. + self.mesh.cell_size() # ----- data ----- @@ -1796,6 +1799,45 @@ def _object_viewer(self): display(Latex(rf"$\quad$ integrator: {self.integrator}, tau shape: {self.tau_shape}")) +class EulerianSUPGPC(_EulerianSUPGPCMethods, _DDtBase): + r"""Scalar P1 Eulerian SUPG with persistent rate history. + + Parameters + ---------- + mesh : Mesh + A two- or three-dimensional volume simplex mesh. + psi_fn : MeshVariable + Continuous scalar P1 temperature, identical to the solver unknown. + V_fn : MeshVariable or sympy Matrix + Advecting velocity at the current step. + method : {"citcoms", "pc_converged"}, default "citcoms" + Fixed residual corrections, or corrections to a residual tolerance. + temperature_rate_field : MeshVariable, optional + Separate continuous scalar P1 derivative field on the same mesh. + adv_gamma : float, default 0.5 + Predictor/corrector weight; pc_converged requires 0.5. + corrector_steps : int, default 2 + Number of fixed CitcomS corrections, not a temporal accuracy order. + corrector_rtol, corrector_atol : float + Converged-mode residual tolerances, default 1e-10 and 1e-12. + max_corrector_steps : int, default 100 + Maximum corrections for pc_converged; failure raises RuntimeError. + tau : scalar expression, optional + Override the automatic steady CitcomS directional-simplex tau. + supg_weight : float, default 1.0 + Runtime stabilization multiplier; zero gives the Galerkin residual. + + Notes + ----- + Attach with ``AdvDiffusion(mesh, T, V_fn, DuDt=transport)``. This manager + owns the derivative, corrector diagnostics, stability estimate and restart + metadata. The solver supplies its PDE residual and constrained scalar DM. + No implicit history fields are allocated. Snapshot after completed steps; + rebuild the same object graph before a fresh-process ``load_state``. + Fixed two-correction CitcomS is not guaranteed second-order in time. + """ + + class SemiLagrangian(_DDtBase): r""" Semi-Lagrangian history manager using nodal swarm. @@ -3878,4 +3920,3 @@ def update_post_solve( self._n_solves_completed += 1 return - diff --git a/src/underworld3/systems/ddt_pc.py b/src/underworld3/systems/ddt_pc.py new file mode 100644 index 00000000..b993b505 --- /dev/null +++ b/src/underworld3/systems/ddt_pc.py @@ -0,0 +1,755 @@ +"""Rate-based Eulerian SUPG integration, internal implementation for :mod:`ddt`.""" + +import math +import weakref +from dataclasses import dataclass +from typing import Optional + +import numpy as np +import sympy +from petsc4py import PETSc + +import underworld3 as uw +import underworld3.timing as timing +from underworld3.checkpoint.state import SnapshottableState + + +@dataclass +class DDtEulerianSUPGPCState(SnapshottableState): + """Rate metadata; rate and stabilization DOFs travel with the mesh.""" + + _schema_version: int = 1 + method: str = "citcoms" + rate_var_name: str = "" + rate_initialised: bool = False + dt: Optional[float] = None + adv_gamma: float = 0.5 + corrector_steps: int = 2 + corrector_rtol: float = 1.0e-10 + corrector_atol: float = 1.0e-12 + max_corrector_steps: int = 100 + supg_weight: float = 1.0 + tau_override: Optional[str] = None + last_corrector_iterations: int = 0 + last_corrector_residual: float = np.inf + corrector_target: float = np.inf + + +class _EulerianSUPGPCMethods: + """PC storage and execution; the public class supplies the DDt base.""" + + def __init__( + self, mesh, psi_fn, V_fn, *, method="citcoms", + temperature_rate_field=None, adv_gamma=0.5, corrector_steps=2, + corrector_rtol=1.0e-10, corrector_atol=1.0e-12, + max_corrector_steps=100, tau=None, supg_weight=1.0, + ): + from underworld3.systems.ddt import _as_row_vector, _UWexpression + + if method not in ("citcoms", "pc_converged"): + raise ValueError("method must be 'citcoms' or 'pc_converged'.") + if not isinstance(psi_fn, uw.discretisation.MeshVariable): + raise TypeError("EulerianSUPGPC requires a temperature MeshVariable.") + if (psi_fn.mesh is not mesh or psi_fn.num_components != 1 + or psi_fn.degree != 1 or not psi_fn.continuous): + raise ValueError("Predictor-corrector transport requires a continuous scalar P1 field on mesh.") + if mesh.dim != mesh.cdim or mesh.dim not in (2, 3): + raise NotImplementedError("Predictor-corrector transport requires a 2-D or 3-D volume mesh.") + if not 0.0 < float(adv_gamma) <= 1.0: + raise ValueError("adv_gamma must be in (0, 1].") + if int(corrector_steps) != corrector_steps or corrector_steps < 1: + raise ValueError("corrector_steps must be a positive integer.") + if temperature_rate_field is not None and ( + temperature_rate_field is psi_fn + or temperature_rate_field.mesh is not mesh + or temperature_rate_field.degree != 1 + or not temperature_rate_field.continuous + or temperature_rate_field.num_components != 1 + ): + raise ValueError("temperature_rate_field must be a separate continuous scalar P1 variable on mesh.") + if method == "pc_converged": + if float(adv_gamma) != 0.5: + raise ValueError("pc_converged requires adv_gamma=0.5 for second-order time accuracy.") + if corrector_steps != 2: + raise ValueError("corrector_steps configures fixed CitcomS corrections only.") + if not np.isfinite(float(corrector_rtol)) or float(corrector_rtol) <= 0.0: + raise ValueError("corrector_rtol must be finite and positive.") + if not np.isfinite(float(corrector_atol)) or float(corrector_atol) < 0.0: + raise ValueError("corrector_atol must be finite and non-negative.") + if int(max_corrector_steps) != max_corrector_steps or max_corrector_steps < 1: + raise ValueError("max_corrector_steps must be a positive integer.") + elif (corrector_rtol != 1.0e-10 or corrector_atol != 1.0e-12 + or max_corrector_steps != 100): + raise ValueError("corrector tolerances and max_corrector_steps configure pc_converged only.") + + super().__init__() + self.mesh = mesh + self._psi_meshVar = psi_fn + self.psi_fn = psi_fn.sym + self.V_fn = _as_row_vector(V_fn, mesh.dim) + self.method = method + self.order = 1 + self.psi_star = [] + self._init_history_tracking(1) + self.adv_gamma = float(adv_gamma) + self.corrector_steps = int(corrector_steps) + self.corrector_rtol = float(corrector_rtol) + self.corrector_atol = float(corrector_atol) + self.max_corrector_steps = int(max_corrector_steps) + self.last_corrector_iterations = 0 + self.last_corrector_residual = np.inf + self.corrector_target = np.inf + self._rate_initialised = False + self._solver_ref = None + self._assembly_dm = None + tag = self.instance_number + self._supg_weight = _UWexpression( + rf"w^{{SUPG}}_{{{tag}}}", float(supg_weight), "SUPG term weight", + _unique_name_generation=True) + self._tau_override = None if tau is None else sympy.sympify(tau) + if isinstance(self._tau_override, sympy.MatrixBase): + raise ValueError("tau must be a scalar expression.") + self._automatic_tau = tau is None + self._temperature_rate = temperature_rate_field + if self._temperature_rate is None: + self._temperature_rate = uw.discretisation.MeshVariable( + f"_supg_dTdt_{tag}", mesh, 1, degree=1, continuous=True) + self._supg_h = None + self._supg_tau = None + if self._automatic_tau: + self._supg_h = uw.discretisation.MeshVariable( + f"_supg_h_{tag}", mesh, 1, degree=0, continuous=False) + self._supg_tau = uw.discretisation.MeshVariable( + f"_supg_tau_{tag}", mesh, 1, degree=0, continuous=False) + self._lumped_mass = None + self._lumped_mass_mesh_version = None + self._citcoms_work_vectors = None + self._citcoms_work_mesh_version = None + self._simplex_data_cache = None + self._simplex_data_mesh_version = None + self._directional_rate_work = None + self._directional_rate_mesh_version = None + self._diffusion_dt_cache = None + self._register_with_default_model() + + @property + def V_fn(self): + """Advecting velocity as a row matrix.""" + return self._V_fn + + @V_fn.setter + def V_fn(self, value): + from underworld3.systems.ddt import _as_row_vector + + self._V_fn = _as_row_vector(value, self.mesh.dim) + solver_ref = getattr(self, "_solver_ref", None) + solver = None if solver_ref is None else solver_ref() + if solver is not None: + solver._needs_function_rewire = True + solver.is_setup = False + + @property + def integrator(self): + """Rate integration method, not a BDF history depth.""" + return self.method + + @property + def theta(self): + """Spatial residuals are evaluated at the current iterate.""" + return 1.0 + + @theta.setter + def theta(self, value): + if float(value) != 1.0: + raise ValueError("Predictor-corrector transport uses adv_gamma, not theta.") + + @property + def temperature_rate(self): + """Persistent derivative field used by the next predictor.""" + return self._temperature_rate + + @property + def supg_weight(self): + """Runtime multiplier of the stabilization term.""" + return float(self._supg_weight.sym) + + @supg_weight.setter + def supg_weight(self, value): + self._supg_weight.sym = float(value) + + def states(self): + return [sympy.Matrix(self.psi_fn)] + + def spatial_weights(self): + return [sympy.Integer(1)] + + def time_derivative(self): + return sympy.Matrix(self._temperature_rate.sym) + + def advecting_velocity(self, level=0): + return self.V_fn + + def advection(self): + from underworld3.systems.ddt import EulerianSUPG + + return EulerianSUPG._convective(self, self.V_fn, self.psi_fn) + + def tau(self): + """Steady CitcomS tau, or the supplied scalar override.""" + value = self._tau_override if self._tau_override is not None else self._supg_tau.sym[0] + return self._supg_weight * value + + def stabilisation_flux(self, R): + from underworld3.systems.ddt import EulerianSUPG + + return EulerianSUPG.stabilisation_flux(self, R) + + def update_pre_solve(self, dt, evalf=False, verbose=False): + raise NotImplementedError("EulerianSUPGPC needs a solver supporting the transport execution hook.") + + def update_post_solve(self, dt, evalf=False, verbose=False): + raise NotImplementedError("EulerianSUPGPC commits its rate through the transport execution hook.") + + def _bind_transport_solver(self, solver): + if solver.mesh is not self.mesh or solver.u is not self._psi_meshVar: + raise ValueError("EulerianSUPGPC must track the solver's temperature field on its mesh.") + if self._solver_ref is not None and self._solver_ref() not in (None, solver): + raise ValueError("EulerianSUPGPC is already attached to another solver.") + self._solver_ref = weakref.ref(solver) + + def _solver(self): + solver = None if self._solver_ref is None else self._solver_ref() + if solver is None: + raise RuntimeError("Attach EulerianSUPGPC to AdvDiffusion with DuDt= before using solver services.") + return solver + + def _invalidate_assembly_cache(self): + if self._lumped_mass is not None: + self._lumped_mass.destroy() + if self._citcoms_work_vectors is not None: + for vector in self._citcoms_work_vectors: + vector.destroy() + self._lumped_mass = None + self._citcoms_work_vectors = None + self._diffusion_dt_cache = None + self._assembly_dm = None + + def _setup_citcoms_residual(self, verbose=False): + solver = self._solver() + solver._prepare_transport_residual(verbose) + if self._assembly_dm is not solver.dm: + self._invalidate_assembly_cache() + self._assembly_dm = solver.dm + + def _compute_citcoms_residual(self, solution=None, residual=None): + """Assemble at the current temperature/rate; caller owns returned vectors.""" + solver = self._solver() + if solution is None: + solution = solver.dm.createGlobalVector() + if residual is None: + residual = solution.duplicate() + solver._compute_transport_residual(solution, residual) + return solution, residual + + def _solve_transport(self, solver, dt, *, zero_init_guess=None, + verbose=False, divergence_retries=0): + self._bind_transport_solver(solver) + if zero_init_guess or divergence_retries: + raise ValueError("Rate-based transport does not support zero_init_guess or SNES divergence retries.") + self._solve_predictor_corrector(dt, verbose=verbose) + + def estimate_dt(self, fraction=0.02, basis=None, + direction_aware=False, percentile=0.0): + """Return the CitcomS stability timestep, dimensionalised when applicable.""" + from underworld3.systems.solvers import _dimensionalise_dt + + if basis not in (None, "stability"): + raise ValueError("Predictor-corrector transport requires basis='stability'.") + if fraction != 0.02 or direction_aware or percentile != 0.0: + raise ValueError("Predictor-corrector transport uses its fixed 0.9 stability factor and directional simplex length.") + return _dimensionalise_dt(self._estimate_citcoms_dt()) + + def _estimate_transport_dt(self, **kwargs): + return self.estimate_dt(**kwargs) + + @property + def state(self): + return DDtEulerianSUPGPCState( + method=self.method, rate_var_name=self.temperature_rate.clean_name, + rate_initialised=self._rate_initialised, dt=self._dt, + adv_gamma=self.adv_gamma, corrector_steps=self.corrector_steps, + corrector_rtol=self.corrector_rtol, corrector_atol=self.corrector_atol, + max_corrector_steps=self.max_corrector_steps, + supg_weight=self.supg_weight, + tau_override=None if self._tau_override is None else str(self._tau_override), + last_corrector_iterations=self.last_corrector_iterations, + last_corrector_residual=self.last_corrector_residual, + corrector_target=self.corrector_target, + ) + + @state.setter + def state(self, state): + if not isinstance(state, DDtEulerianSUPGPCState): + raise TypeError("EulerianSUPGPC state has the wrong type.") + current = self.state + for name in ("_schema_version", "method", "rate_var_name", "adv_gamma", + "corrector_steps", "corrector_rtol", "corrector_atol", + "max_corrector_steps", "tau_override"): + if getattr(state, name) != getattr(current, name): + raise ValueError(f"EulerianSUPGPC {name} changed since snapshot.") + self._dt = state.dt + self._rate_initialised = bool(state.rate_initialised) + self.supg_weight = state.supg_weight + self.last_corrector_iterations = state.last_corrector_iterations + self.last_corrector_residual = state.last_corrector_residual + self.corrector_target = state.corrector_target + self._invalidate_assembly_cache() + self._simplex_data_cache = None + self._directional_rate_work = None + self.mesh._stale_lvec = True + + def _simplex_data(self): + """Return local simplex data; validate the layout collectively on rebuild.""" + from underworld3.meshing.smoothing import _tet_cells, _tri_cells + + mesh_version = getattr(self.mesh, "_mesh_version", 0) + if ( + self._simplex_data_cache is not None + and self._simplex_data_mesh_version == mesh_version + ): + return self._simplex_data_cache + + cells = ( + _tri_cells(self.mesh.dm) + if self.mesh.dim == 2 + else _tet_cells(self.mesh.dm) if self.mesh.dim == 3 else None + ) + cell_start, cell_end = self.mesh.dm.getHeightStratum(0) + invalid = ( + (uw.mpi.rank, cell_end - cell_start, self.mesh.dim, self.mesh.cdim) + if cells is None or self.mesh.dim != self.mesh.cdim else None + ) + invalid_ranks = [item for item in uw.mpi.comm.allgather(invalid) if item is not None] + if invalid_ranks: + raise NotImplementedError( + "Automatic CitcomS operations require a non-empty 2-D or 3-D " + "volume simplex partition on every rank. Unsupported local " + f"layouts (rank, cells, dim, cdim): {invalid_ranks}." + ) + + coords = np.asarray(self.mesh.X.coords) + cell_coords = coords[cells] + edges = cell_coords[:, 1:, :] - cell_coords[:, :1, :] + try: + inverse_edges = np.linalg.inv(edges) + except np.linalg.LinAlgError as error: + raise RuntimeError("Cannot operate on a singular simplex.") from error + + gradients = np.empty_like(cell_coords) + gradients[:, 1:, :] = np.transpose(inverse_edges, (0, 2, 1)) + gradients[:, 0, :] = -gradients[:, 1:, :].sum(axis=1) + volumes = np.abs(np.linalg.det(edges)) / math.factorial(self.mesh.dim) + self._simplex_data_cache = (cells, gradients, volumes) + self._simplex_data_mesh_version = mesh_version + return self._simplex_data_cache + + def _streamline_directional_rate(self, gradients, velocity): + """Return ``sum_a |u.grad(N_a)|`` using reusable cell work arrays.""" + mesh_version = getattr(self.mesh, "_mesh_version", 0) + cell_count = velocity.shape[0] + if ( + self._directional_rate_work is None + or self._directional_rate_mesh_version != mesh_version + or self._directional_rate_work[0].shape != (cell_count,) + ): + self._directional_rate_work = ( + np.empty(cell_count, dtype=float), + np.empty(cell_count, dtype=float), + ) + self._directional_rate_mesh_version = mesh_version + + directional_rate, projection = self._directional_rate_work + directional_rate.fill(0.0) + for basis_index in range(gradients.shape[1]): + np.einsum( + "cd,cd->c", + gradients[:, basis_index, :], + velocity, + out=projection, + ) + np.abs(projection, out=projection) + np.add(directional_rate, projection, out=directional_rate) + return directional_rate + + def _cell_diffusivity(self, cell_count): + """Evaluate non-negative scalar diffusivity at cell centroids.""" + diffusivity_expr = sympy.sympify(self._solver()._scalar_diffusivity()) + if isinstance(diffusivity_expr, sympy.MatrixBase): + raise NotImplementedError( + "Automatic SUPG operations require scalar isotropic " + "diffusivity; supply tau explicitly for tensor diffusivity." + ) + diffusivity = uw.function.evaluate(diffusivity_expr, self.mesh._centroids) + if hasattr(diffusivity, "units") and diffusivity.units is not None: + diffusivity = uw.non_dimensionalise(diffusivity) + elif hasattr(diffusivity, "magnitude"): + diffusivity = diffusivity.magnitude + diffusivity = np.asarray(diffusivity, dtype=float).reshape(-1) + if diffusivity.size == 1: + diffusivity = np.full(cell_count, diffusivity.item()) + if diffusivity.shape != (cell_count,): + raise ValueError("Diffusivity must evaluate to one scalar per cell.") + if np.any(diffusivity < 0.0): + raise ValueError("SUPG diffusivity must be non-negative.") + return diffusivity + + def _update_automatic_tau(self): + """Update local simplex streamline lengths and automatic tau values.""" + if not self._automatic_tau: + return + + _, gradients, _ = self._simplex_data() + + from underworld3.systems.solvers import _centroid_velocities_nd + + velocity = _centroid_velocities_nd(self.V_fn, self.mesh) + speed = np.linalg.norm(velocity, axis=1) + directional_rate = self._streamline_directional_rate(gradients, velocity) + h_stream = np.divide( + 2.0 * speed, + directional_rate, + out=np.zeros_like(speed), + where=directional_rate > 0.0, + ) + + diffusivity = self._cell_diffusivity(speed.size) + + tau_steady = np.zeros_like(speed) + moving = speed > np.finfo(float).eps + diffusive = moving & (diffusivity > 0.0) + nondiffusive = moving & ~diffusive + + if np.any(diffusive): + pe = speed[diffusive] * h_stream[diffusive] / (2.0 * diffusivity[diffusive]) + tau_steady[diffusive] = ( + h_stream[diffusive] + * np.maximum(0.0, 1.0 - 1.0 / pe) + / (2.0 * speed[diffusive]) + ) + tau_steady[nondiffusive] = h_stream[nondiffusive] / (2.0 * speed[nondiffusive]) + + tau_values = tau_steady + + if self._supg_h.array.shape[0] != h_stream.size: + raise RuntimeError("SUPG P0 field and local simplex counts do not match.") + self._supg_h.array[:, 0, 0] = h_stream + self._supg_tau.array[:, 0, 0] = tau_values + + def _assemble_lumped_mass(self): + """Assemble positive P1 simplex row-sum masses on free global DOFs.""" + mesh_version = getattr(self.mesh, "_mesh_version", 0) + if ( + self._lumped_mass is not None + and self._lumped_mass_mesh_version == mesh_version + ): + return self._lumped_mass + if self._lumped_mass is not None: + self._lumped_mass.destroy() + self._lumped_mass = None + + from underworld3.meshing.smoothing import _owned_cell_mask + + cells, _, volumes = self._simplex_data() + owned = _owned_cell_mask(self.mesh.dm) + + local_mass = self._solver().dm.createLocalVector() + global_mass = self._solver().dm.createGlobalVector() + local_mass.set(0.0) + global_mass.set(0.0) + section = self._solver().dm.getLocalSection() + vertex_start, _ = self.mesh.dm.getDepthStratum(0) + + for cell_index in np.flatnonzero(owned): + contribution = volumes[cell_index] / (self.mesh.dim + 1) + for vertex_index in cells[cell_index]: + offset = section.getOffset(vertex_start + int(vertex_index)) + if offset >= 0: + local_mass.array[offset] += contribution + + self._solver().dm.localToGlobal( + local_mass, + global_mass, + addv=PETSc.InsertMode.ADD_VALUES, + ) + local_mass.destroy() + if global_mass.getLocalSize() and np.any(global_mass.array <= 0.0): + global_mass.destroy() + raise RuntimeError("CitcomS P1 lumped mass contains non-positive rows.") + + self._lumped_mass = global_mass + self._lumped_mass_mesh_version = mesh_version + return self._lumped_mass + + def _citcoms_vectors(self): + """Return reusable global vectors for predictor-corrector updates.""" + mesh_version = getattr(self.mesh, "_mesh_version", 0) + if ( + self._citcoms_work_vectors is not None + and self._citcoms_work_mesh_version == mesh_version + ): + return self._citcoms_work_vectors + + if self._citcoms_work_vectors is not None: + for vector in self._citcoms_work_vectors: + vector.destroy() + + solution = self._solver().dm.createGlobalVector() + residual = solution.duplicate() + delta_rate = solution.duplicate() + rate = solution.duplicate() + self._citcoms_work_vectors = (solution, residual, delta_rate, rate) + self._citcoms_work_mesh_version = mesh_version + return self._citcoms_work_vectors + + @timing.routine_timer_decorator + def _estimate_citcoms_dt(self): + """Estimate a simplex advection-diffusion timestep. + + The predictor-corrector modes use + ``0.9 * min(1/max(lambda_adv), 2/max(rowsum(abs(M_L^-1 K))))``. + Generic implicit transport retains its separate Eulerian estimator. + """ + from underworld3.systems.solvers import _centroid_velocities_nd + from mpi4py import MPI + from underworld3.meshing.smoothing import _owned_cell_mask + + cells, gradients, volumes = self._simplex_data() + velocity = _centroid_velocities_nd(self.V_fn, self.mesh) + directional_rate = self._streamline_directional_rate(gradients, velocity) + local_adv_rate = ( + float(np.max(directional_rate)) if directional_rate.size else 0.0 + ) + adv_rate = uw.mpi.comm.allreduce(local_adv_rate, op=MPI.MAX) + dt_adv = 1.0 / adv_rate if adv_rate > 0.0 else np.inf + + diffusivity = self._cell_diffusivity(len(cells)) + has_diffusivity = bool( + uw.mpi.comm.allreduce( + int(np.any(diffusivity > 0.0)), + op=MPI.MAX, + ) + ) + if not has_diffusivity: + dt_diff = np.inf + else: + self._setup_citcoms_residual() + mass = self._assemble_lumped_mass() + diffusion_signature = ( + getattr(self.mesh, "_mesh_version", 0), + hash(diffusivity.tobytes()), + ) + local_cache_valid = ( + self._diffusion_dt_cache is not None + and self._diffusion_dt_cache[0] == diffusion_signature + ) + cache_valid = bool( + uw.mpi.comm.allreduce(int(local_cache_valid), op=MPI.MIN) + ) + if cache_valid: + dt_diff = self._diffusion_dt_cache[1] + self.dt_adv = dt_adv + self.dt_diff = dt_diff + return 0.9 * min(dt_adv, dt_diff) + + stiffness = self._solver().dm.createMatrix() + stiffness.setOption(PETSc.Mat.Option.NEW_NONZERO_LOCATION_ERR, False) + section = self._solver().dm.getLocalSection() + vertex_start, _ = self.mesh.dm.getDepthStratum(0) + owned = _owned_cell_mask(self.mesh.dm) + + for cell_index in np.flatnonzero(owned): + points = [vertex_start + int(index) for index in cells[cell_index]] + local_dofs = [section.getOffset(point) for point in points] + element_stiffness = ( + diffusivity[cell_index] + * volumes[cell_index] + * gradients[cell_index].dot(gradients[cell_index].T) + ) + stiffness.setValuesLocal( + local_dofs, + local_dofs, + element_stiffness, + addv=PETSc.InsertMode.ADD_VALUES, + ) + stiffness.assemble() + + row_start, row_end = stiffness.getOwnershipRange() + local_diff_rate = 0.0 + for row in range(row_start, row_end): + _, values = stiffness.getRow(row) + row_sum = float(np.sum(np.abs(values))) + local_diff_rate = max( + local_diff_rate, + row_sum / mass.array[row - row_start], + ) + diff_rate = uw.mpi.comm.allreduce(local_diff_rate, op=MPI.MAX) + stiffness.destroy() + dt_diff = 2.0 / diff_rate if diff_rate > 0.0 else np.inf + self._diffusion_dt_cache = (diffusion_signature, dt_diff) + + self.dt_adv = dt_adv + self.dt_diff = dt_diff + return 0.9 * min(dt_adv, dt_diff) + + def _apply_pc_correction( + self, + temperature_global, + residual, + delta_rate, + rate_global, + mass, + dt, + *, + advance_temperature, + ): + """Apply one lumped-preconditioned correction to rate and temperature.""" + delta_rate.pointwiseDivide(residual, mass) + delta_rate.scale(-1.0) + rate_global.set(0.0) + self._solver().dm.localToGlobal(self._temperature_rate.vec, rate_global, addv=False) + rate_global.axpy(1.0, delta_rate) + if advance_temperature: + temperature_global.axpy(self.adv_gamma * dt, delta_rate) + + self._temperature_rate.vec.set(0.0) + self._solver().dm.globalToLocal(rate_global, self._temperature_rate.vec) + if advance_temperature: + from underworld3.cython.petsc_discretisation import ( + petsc_dm_insert_boundary_values, + ) + + self._psi_meshVar.vec.set(0.0) + self._solver().dm.globalToLocal(temperature_global, self._psi_meshVar.vec) + petsc_dm_insert_boundary_values(self._solver().dm, self._psi_meshVar.vec) + self.mesh._stale_lvec = True + + def _converge_pc_residual( + self, + temperature_global, + residual, + delta_rate, + rate_global, + mass, + dt, + *, + advance_temperature, + ): + """Iterate the predictor-corrector residual to its configured tolerance.""" + initial_norm = None + for corrections in range(self.max_corrector_steps + 1): + self._compute_citcoms_residual(temperature_global, residual) + residual_norm = float(residual.norm(PETSc.NormType.NORM_2)) + if not np.isfinite(residual_norm): + raise RuntimeError("pc_converged produced a non-finite residual norm.") + if initial_norm is None: + initial_norm = residual_norm + self.corrector_target = max( + self.corrector_atol, + self.corrector_rtol * initial_norm, + ) + self.last_corrector_iterations = corrections + self.last_corrector_residual = residual_norm + if residual_norm <= self.corrector_target: + return + if corrections == self.max_corrector_steps: + break + self._apply_pc_correction( + temperature_global, + residual, + delta_rate, + rate_global, + mass, + dt, + advance_temperature=advance_temperature, + ) + raise RuntimeError( + "pc_converged did not reach its predictor-corrector residual " + f"tolerance after {self.max_corrector_steps} corrections: " + f"residual={self.last_corrector_residual:.6e}, " + f"target={self.corrector_target:.6e}." + ) + + def _solve_predictor_corrector(self, timestep, verbose=False): + """Advance one fixed or residual-converged predictor-corrector step.""" + from underworld3.systems.solvers import _invalidate_solution_cache + + dt = self._dt if timestep is None else float(timestep) + if dt is None or not np.isfinite(dt) or dt <= 0.0: + raise ValueError("EulerianSUPGPC requires a finite positive timestep.") + self._dt = dt + + self._update_automatic_tau() + self._setup_citcoms_residual(verbose) + mass = self._assemble_lumped_mass() + temperature_global, residual, delta_rate, rate_global = self._citcoms_vectors() + + if not self._rate_initialised: + self._temperature_rate.array[:, 0, 0] = 0.0 + if self.method == "pc_converged": + self.mesh._stale_lvec = True + self._converge_pc_residual( + temperature_global, + residual, + delta_rate, + rate_global, + mass, + dt, + advance_temperature=False, + ) + else: + self._compute_citcoms_residual(temperature_global, residual) + delta_rate.pointwiseDivide(residual, mass) + delta_rate.scale(-1.0) + self._temperature_rate.vec.set(0.0) + self._solver().dm.globalToLocal(delta_rate, self._temperature_rate.vec) + self.mesh._stale_lvec = True + self._rate_initialised = True + + self._psi_meshVar.array[:, 0, 0] += ( + (1.0 - self.adv_gamma) * dt * self._temperature_rate.array[:, 0, 0] + ) + self._temperature_rate.array[:, 0, 0] = 0.0 + self.mesh._stale_lvec = True + + if self.method == "pc_converged": + from underworld3.cython.petsc_discretisation import ( + petsc_dm_insert_boundary_values, + ) + + petsc_dm_insert_boundary_values(self._solver().dm, self._psi_meshVar.vec) + self.mesh._stale_lvec = True + self._converge_pc_residual( + temperature_global, + residual, + delta_rate, + rate_global, + mass, + dt, + advance_temperature=True, + ) + else: + for _ in range(self.corrector_steps): + self._compute_citcoms_residual(temperature_global, residual) + self._apply_pc_correction( + temperature_global, + residual, + delta_rate, + rate_global, + mass, + dt, + advance_temperature=True, + ) + + _invalidate_solution_cache(self._psi_meshVar) + _invalidate_solution_cache(self._temperature_rate) + return diff --git a/tests/parallel/ptest_1119_supg_restart.py b/tests/parallel/ptest_1119_supg_restart.py new file mode 100644 index 00000000..bacc0542 --- /dev/null +++ b/tests/parallel/ptest_1119_supg_restart.py @@ -0,0 +1,115 @@ +"""Worker for a fresh-process SUPG restart; no Stokes or A1 setup.""" + +from dataclasses import asdict +import h5py +import numpy as np + +import underworld3 as uw + + +params = uw.Params( + uw_method=uw.Param("pc2", type=uw.ParamType.STRING), + uw_phase=uw.Param("full", type=uw.ParamType.STRING), +) +assert params.uw_method in ("pc2", "pc_converged", "cn", "bdf2") +assert params.uw_phase in ("full", "write", "resume") +uw.reset_default_model() +orchestration_model = uw.get_default_model() +mesh = uw.meshing.UnstructuredSimplexBox( + minCoords=(0.0, 0.0, 0.0), maxCoords=(1.0, 1.0, 1.0), + cellSize=0.25, qdegree=4, regular=False, filename="mesh.msh", +) +temperature = uw.discretisation.MeshVariable("T", mesh, 1, degree=1) +velocity = uw.discretisation.MeshVariable("U", mesh, 3, degree=1) +temperature.array[:, 0, 0] = np.prod(np.sin(np.pi * np.asarray(temperature.coords)), axis=1) +velocity.array[...] = 0.0 +velocity.array[:, 0, 0] = 0.2 +if params.uw_method in ("pc2", "pc_converged"): + manager = uw.systems.ddt.EulerianSUPGPC( + mesh, temperature, velocity.sym, + method="citcoms" if params.uw_method == "pc2" else params.uw_method, + ) + thermal = uw.systems.AdvDiffusion(mesh, temperature, velocity.sym, DuDt=manager) +else: + thermal = uw.systems.AdvDiffusion( + mesh, temperature, velocity.sym, + order=1 if params.uw_method == "cn" else 2, + theta=0.5 if params.uw_method == "cn" else 1.0, + peclet_weight=0.0, + ) + manager = thermal.DuDt +thermal.constitutive_model.Parameters.diffusivity = 0.01 +for boundary in mesh.boundaries: + if boundary.name not in ("All_Boundaries", "Null_Boundary"): + thermal.add_dirichlet_bc(0.0, boundary.name) +thermal.petsc_options["ksp_rtol"] = 1e-14 +thermal.petsc_options["ksp_atol"] = 0.0 +thermal.petsc_options["snes_rtol"] = 1e-13 +thermal.petsc_options["snes_atol"] = 1e-14 +orchestration_model.tracker.step = 0 +orchestration_model.tracker.time = 0.0 + + +def capture(): + """All evolving fields and numerical metadata, separately from PETSc files.""" + fields = [temperature, velocity] + if isinstance(manager, uw.systems.ddt.EulerianSUPGPC): + fields.append(manager.temperature_rate) + else: + fields.extend(thermal.DuDt.psi_star) + record = {field.clean_name: np.array(field.array) for field in fields} + record["coords"] = np.asarray(temperature.coords) + record["step"] = orchestration_model.tracker.step + record["time"] = orchestration_model.tracker.time + record["estimate_dt"] = float(thermal.estimate_dt()) + for name, value in asdict(thermal.state).items(): + record["solver_" + name] = "None" if value is None else value + for name, value in asdict(manager.state).items(): + if name != "psi_star_var_names": + record["history_" + name] = "None" if value is None else value + return record + + +if params.uw_phase == "resume": + orchestration_model.load_state("checkpoint.h5") + restored = capture() + failure = None + try: + with h5py.File(f"write_rank{uw.mpi.rank}.h5", "r") as saved: + assert set(saved) == set(restored) + for name, actual in restored.items(): + expected = saved[name][()] + if isinstance(expected, bytes): + expected = expected.decode() + np.testing.assert_array_equal(actual, expected, err_msg=name) + except Exception as error: + # Every rank must report a failed restore before peers enter a solve. + failure = str(error) + failures = uw.mpi.comm.allgather(failure) + assert not any(failures), failures + uw.pprint(f"SUPG_RESTORE_EXACT method={params.uw_method} ranks={uw.mpi.size}") + +end_step = 5 if params.uw_phase == "write" else 12 +for step in range(orchestration_model.tracker.step, end_step): + dt = (0.002, 0.003, 0.0015, 0.0025)[step % 4] + velocity.array[:, 0, 0] = 0.2 * (1.0 + 0.1 * np.sin(step)) + thermal.solve(timestep=dt) + orchestration_model.tracker.step = step + 1 + orchestration_model.tracker.time += dt + orchestration_model.tracker.dt = dt + +if params.uw_phase == "write": + orchestration_model.save_state(file="checkpoint.h5") + +record = capture() +failure = None +try: + with h5py.File(f"{params.uw_phase}_rank{uw.mpi.rank}.h5", "w") as output: + for name, value in record.items(): + output[name] = value +except Exception as error: + # Rank-local test output errors must not strand peers on shutdown. + failure = str(error) +failures = uw.mpi.comm.allgather(failure) +assert not any(failures), failures +uw.pprint(f"SUPG_RESTART_STAGE phase={params.uw_phase} method={params.uw_method} step={end_step}") diff --git a/tests/test_1113_advdiff_supg_residual.py b/tests/test_1113_advdiff_supg_residual.py new file mode 100644 index 00000000..4df32756 --- /dev/null +++ b/tests/test_1113_advdiff_supg_residual.py @@ -0,0 +1,296 @@ +"""Focused tests for the implicit SUPG scalar transport residual.""" + +import numpy as np +import pytest +import sympy + +import underworld3 as uw + + +pytestmark = [pytest.mark.level_1, pytest.mark.tier_a] + + +def _mesh_temperature_velocity(prefix, velocity=(1.0, 0.0)): + mesh = uw.meshing.UnstructuredSimplexBox( + minCoords=(0.0, 0.0), + maxCoords=(1.0, 1.0), + cellSize=0.3, + regular=False, + ) + temperature = uw.discretisation.MeshVariable(f"T_{prefix}", mesh, 1, degree=1) + flow = uw.discretisation.MeshVariable(f"U_{prefix}", mesh, mesh.dim, degree=1) + temperature.array[:, 0, 0] = temperature.coords[:, 0] + flow.array[:, 0, 0] = velocity[0] + flow.array[:, 0, 1] = velocity[1] + return mesh, temperature, flow + + +def _configure_diffusion(solver, diffusivity=0.1): + solver.constitutive_model = uw.constitutive_models.DiffusionModel + solver.constitutive_model.Parameters.diffusivity = diffusivity + + +def test_public_api_and_residual_shapes(): + mesh, temperature, velocity = _mesh_temperature_velocity("api") + thermal = uw.systems.AdvDiffusion( + mesh, temperature, velocity.sym, theta=0.5, peclet_weight=0.0 + ) + thermal.DuDt.supg_weight = 0.0 + _configure_diffusion(thermal) + thermal.delta_t = 0.01 + + assert thermal.F0.sym.shape == (1, 1) + assert thermal.F1.sym.shape == (1, mesh.cdim) + np.testing.assert_array_equal( + uw.function.evaluate(thermal.DuDt.tau(), mesh._centroids), 0.0 + ) + + +def test_manager_owns_eulerian_advection_once(): + mesh, temperature, velocity = _mesh_temperature_velocity("double") + manager = uw.systems.ddt.EulerianSUPG( + mesh, + temperature, + velocity.sym, + vtype=uw.VarType.SCALAR, + degree=temperature.degree, + continuous=temperature.continuous, + theta=1.0, + peclet_weight=0.0, + ) + thermal = uw.systems.AdvDiffusion( + mesh, temperature, velocity.sym, DuDt=manager, theta=1.0 + ) + _configure_diffusion(thermal) + thermal.delta_t = 0.01 + expected = manager.time_derivative() + manager.advection() - thermal.f + assert thermal.DuDt is manager + assert sympy.simplify(thermal.F0.sym - expected) == sympy.zeros(1, 1) + + +@pytest.mark.parametrize("theta", (0.0, 0.5)) +def test_rejects_nonimplicit_flux_history(theta): + mesh, temperature, velocity = _mesh_temperature_velocity(f"theta_{theta}") + with pytest.raises(ValueError, match="theta=1.0"): + uw.systems.AdvDiffusion( + mesh, temperature, velocity.sym, order=2, theta=theta, peclet_weight=0.0 + ) + + +def test_automatic_tau_is_finite_and_bounded_by_transient_scale(): + mesh, temperature, velocity = _mesh_temperature_velocity("tau") + thermal = uw.systems.AdvDiffusion( + mesh, temperature, velocity.sym, theta=0.5, peclet_weight=0.0 + ) + _configure_diffusion(thermal, diffusivity=0.1) + thermal.delta_t = 0.02 + thermal.DuDt.diffusivity = 0.1 + + tau = uw.function.evaluate(thermal.DuDt.tau(), mesh._centroids) + assert np.all(np.isfinite(tau)) + assert np.all(tau > 0.0) + assert np.all(tau <= 0.01) + + +def test_negative_diffusivity_is_rejected(): + mesh, temperature, velocity = _mesh_temperature_velocity("negative_k") + manager = uw.systems.ddt.EulerianSUPGPC( + mesh, temperature, velocity.sym, method="citcoms" + ) + thermal = uw.systems.AdvDiffusion(mesh, temperature, velocity.sym, DuDt=manager) + _configure_diffusion(thermal, diffusivity=-0.1) + thermal.delta_t = 0.01 + + with pytest.raises(ValueError, match="non-negative"): + thermal.DuDt._update_automatic_tau() + + +def test_zero_velocity_matches_diffusion_solver(): + mesh_a, temperature_a, velocity = _mesh_temperature_velocity( + "supg_zero", velocity=(0.0, 0.0) + ) + mesh_b, temperature_b, _ = _mesh_temperature_velocity( + "diffusion", velocity=(0.0, 0.0) + ) + temperature_a.array[:, 0, 0] = np.sin(np.pi * temperature_a.coords[:, 0]) + temperature_b.array[:, 0, 0] = np.sin(np.pi * temperature_b.coords[:, 0]) + + supg = uw.systems.AdvDiffusion( + mesh_a, temperature_a, velocity.sym, theta=1.0, peclet_weight=0.0 + ) + diffusion = uw.systems.Diffusion(mesh_b, u_Field=temperature_b, theta=1.0) + _configure_diffusion(supg, diffusivity=0.1) + _configure_diffusion(diffusion, diffusivity=0.1) + # Compare equations at the same solve accuracy, not two preconditioners' + # different default stopping criteria. + for solver in (supg, diffusion): + solver.petsc_options["ksp_rtol"] = 1.0e-13 + solver.petsc_options["snes_rtol"] = 1.0e-12 + solver.petsc_options["snes_atol"] = 1.0e-13 + + supg.solve(timestep=0.01, zero_init_guess=False) + diffusion.solve(timestep=0.01, zero_init_guess=False) + + np.testing.assert_allclose( + temperature_a.array, + temperature_b.array, + rtol=1.0e-11, + atol=1.0e-11, + ) + + +def test_citcoms_integrator_requires_continuous_p1_temperature(): + mesh, temperature, velocity = _mesh_temperature_velocity("citcoms_p1") + temperature_p2 = uw.discretisation.MeshVariable("T_citcoms_p2", mesh, 1, degree=2) + + with pytest.raises(ValueError, match="continuous scalar P1"): + uw.systems.ddt.EulerianSUPGPC( + mesh, temperature_p2, velocity.sym, + method="citcoms", + ) + + +def test_converged_pc_validates_correction_controls(): + mesh, temperature, velocity = _mesh_temperature_velocity("pc_converged_api") + manager = uw.systems.ddt.EulerianSUPGPC( + mesh, temperature, velocity.sym, + method="pc_converged", + ) + thermal = uw.systems.AdvDiffusion(mesh, temperature, velocity.sym, DuDt=manager) + assert manager.method == "pc_converged" + assert manager.integrator == "pc_converged" + assert thermal.DuDt.corrector_rtol == pytest.approx(1.0e-10) + assert thermal.DuDt.corrector_atol == pytest.approx(1.0e-12) + assert thermal.DuDt.max_corrector_steps == 100 + + invalid = ( + ({"corrector_rtol": 0.0}, "corrector_rtol"), + ({"corrector_atol": -1.0}, "corrector_atol"), + ({"max_corrector_steps": 0}, "max_corrector_steps"), + ({"adv_gamma": 0.6}, "adv_gamma=0.5"), + ) + for kwargs, message in invalid: + with pytest.raises(ValueError, match=message): + uw.systems.ddt.EulerianSUPGPC( + mesh, temperature, velocity.sym, + method="pc_converged", **kwargs, + ) + + +def test_converged_pc_fails_when_residual_tolerance_is_not_reached(): + mesh, temperature, velocity = _mesh_temperature_velocity( + "pc_converged_failure", velocity=(0.0, 0.0) + ) + temperature.array[:, 0, 0] = np.prod( + np.sin(np.pi * np.asarray(temperature.coords)), axis=1 + ) + manager = uw.systems.ddt.EulerianSUPGPC( + mesh, temperature, velocity.sym, + method="pc_converged", + corrector_rtol=1.0e-15, + corrector_atol=0.0, + max_corrector_steps=1, + ) + thermal = uw.systems.AdvDiffusion(mesh, temperature, velocity.sym, DuDt=manager) + _configure_diffusion(thermal, diffusivity=0.1) + + with pytest.raises(RuntimeError, match="did not reach"): + thermal.solve(timestep=0.01) + + +def test_citcoms_lumped_mass_matches_constant_residual(): + mesh, temperature, velocity = _mesh_temperature_velocity( + "citcoms_mass", velocity=(0.0, 0.0) + ) + manager = uw.systems.ddt.EulerianSUPGPC( + mesh, temperature, velocity.sym, + method="citcoms", tau=0.0, + ) + thermal = uw.systems.AdvDiffusion(mesh, temperature, velocity.sym, DuDt=manager) + _configure_diffusion(thermal, diffusivity=0.0) + thermal.delta_t = 0.01 + thermal.DuDt._setup_citcoms_residual() + mass = thermal.DuDt._assemble_lumped_mass() + thermal.DuDt.temperature_rate.array[:, 0, 0] = 1.0 + solution, residual = thermal.DuDt._compute_citcoms_residual() + + np.testing.assert_allclose(residual.array / mass.array, 1.0, atol=1.0e-14) + assert mass.min()[1] > 0.0 + solution.destroy() + residual.destroy() + + +def test_citcoms_constant_source_is_exact_from_first_step(): + mesh, temperature, velocity = _mesh_temperature_velocity( + "citcoms_source", velocity=(0.0, 0.0) + ) + temperature.array[:, 0, 0] = 0.0 + manager = uw.systems.ddt.EulerianSUPGPC( + mesh, temperature, velocity.sym, + method="citcoms", tau=0.0, + ) + thermal = uw.systems.AdvDiffusion(mesh, temperature, velocity.sym, DuDt=manager) + _configure_diffusion(thermal, diffusivity=0.0) + thermal.f = 1.0 + + thermal.solve(timestep=0.1) + + np.testing.assert_allclose(temperature.array, 0.1, atol=1.0e-14) + np.testing.assert_allclose(thermal.DuDt.temperature_rate.array, 1.0, atol=1.0e-14) + + +def test_citcoms_reuses_predictor_corrector_work_vectors(): + mesh, temperature, velocity = _mesh_temperature_velocity( + "citcoms_workspace", velocity=(0.0, 0.0) + ) + manager = uw.systems.ddt.EulerianSUPGPC( + mesh, temperature, velocity.sym, + method="citcoms", tau=0.0, + ) + thermal = uw.systems.AdvDiffusion(mesh, temperature, velocity.sym, DuDt=manager) + _configure_diffusion(thermal, diffusivity=0.0) + + thermal.solve(timestep=0.01) + vector_handles = tuple(vector.handle for vector in thermal.DuDt._citcoms_work_vectors) + thermal.solve(timestep=0.01) + + assert ( + tuple(vector.handle for vector in thermal.DuDt._citcoms_work_vectors) + == vector_handles + ) + + +def test_citcoms_timestep_uses_advection_and_lumped_diffusion_limits(): + mesh, temperature, velocity = _mesh_temperature_velocity("citcoms_dt") + manager = uw.systems.ddt.EulerianSUPGPC( + mesh, temperature, velocity.sym, + method="citcoms", + ) + thermal = uw.systems.AdvDiffusion(mesh, temperature, velocity.sym, DuDt=manager) + _configure_diffusion(thermal, diffusivity=0.1) + + timestep = thermal.estimate_dt() + + assert np.isfinite(timestep) + assert timestep == pytest.approx(0.9 * min(thermal.DuDt.dt_adv, thermal.DuDt.dt_diff)) + assert thermal.DuDt.dt_adv > 0.0 + assert thermal.DuDt.dt_diff > 0.0 + + +def test_timestep_diffusivity_branch_is_collective(): + mesh, temperature, velocity = _mesh_temperature_velocity("collective_diffusivity") + manager = uw.systems.ddt.EulerianSUPGPC( + mesh, temperature, velocity.sym, + method="citcoms", + ) + thermal = uw.systems.AdvDiffusion(mesh, temperature, velocity.sym, DuDt=manager) + _configure_diffusion(thermal, diffusivity=0.1) + thermal.delta_t = 0.01 + thermal.DuDt._cell_diffusivity = lambda count: ( + np.ones(count) if uw.mpi.rank == 0 else np.zeros(count) + ) + + timestep = thermal.estimate_dt() + + assert np.isfinite(timestep) + assert thermal.DuDt.dt_diff > 0.0 diff --git a/tests/test_1114_advdiff_supg.py b/tests/test_1114_advdiff_supg.py new file mode 100644 index 00000000..768aa878 --- /dev/null +++ b/tests/test_1114_advdiff_supg.py @@ -0,0 +1,458 @@ +"""Numerical validation for implicit and predictor-corrector SUPG transport.""" + +from pathlib import Path +import sys + +import numpy as np +import pytest +import sympy + +import underworld3 as uw + + +pytestmark = [pytest.mark.level_2, pytest.mark.tier_b] + + +def test_simplex_geometry_is_reused_between_automatic_operations(): + mesh = uw.meshing.UnstructuredSimplexBox( + minCoords=(0.0, 0.0), + maxCoords=(1.0, 1.0), + cellSize=0.25, + regular=True, + ) + temperature = uw.discretisation.MeshVariable( + "T_geometry_cache", mesh, 1, degree=1 + ) + velocity = uw.discretisation.MeshVariable( + "U_geometry_cache", mesh, mesh.dim, degree=1 + ) + manager = uw.systems.ddt.EulerianSUPGPC( + mesh, temperature, velocity.sym, + method="citcoms", + ) + thermal = uw.systems.AdvDiffusion(mesh, temperature, velocity.sym, DuDt=manager) + thermal.constitutive_model = uw.constitutive_models.DiffusionModel + thermal.constitutive_model.Parameters.diffusivity = 1.0 + + first = thermal.DuDt._simplex_data() + second = thermal.DuDt._simplex_data() + + assert all(a is b for a, b in zip(first, second)) + + sample_velocity = np.column_stack( + ( + np.linspace(0.1, 0.9, len(first[0])), + np.linspace(-0.3, 0.4, len(first[0])), + ) + ) + expected_rate = np.abs( + np.einsum("cad,cd->ca", first[1], sample_velocity) + ).sum(axis=1) + first_rate = thermal.DuDt._streamline_directional_rate( + first[1], sample_velocity + ) + second_rate = thermal.DuDt._streamline_directional_rate( + first[1], sample_velocity + ) + + np.testing.assert_allclose(first_rate, expected_rate) + assert first_rate is second_rate + + deformed = mesh.X.coords.copy() + deformed[:, 0] *= 1.1 + mesh.deform(deformed) + third = thermal.DuDt._simplex_data() + + assert all(a is not b for a, b in zip(first, third)) + assert not np.isclose(first[2].sum(), third[2].sum()) + assert thermal.DuDt._streamline_directional_rate( + third[1], sample_velocity + ) is not first_rate + + +def test_tetrahedron_streamline_length_is_geometry_invariant(): + mesh = uw.meshing.UnstructuredSimplexBox( + minCoords=(0.0, 0.0, 0.0), + maxCoords=(1.0, 1.0, 1.0), + cellSize=1.0, + regular=True, + ) + temperature = uw.discretisation.MeshVariable( + "T_tet_streamline", mesh, 1, degree=1 + ) + velocity_field = uw.discretisation.MeshVariable( + "U_tet_streamline", mesh, mesh.dim, degree=1 + ) + manager = uw.systems.ddt.EulerianSUPGPC( + mesh, temperature, velocity_field.sym, + method="citcoms", + ) + thermal = uw.systems.AdvDiffusion(mesh, temperature, velocity_field.sym, DuDt=manager) + + lengths = np.array((2.0, 1.0, 0.5)) + gradients = np.vstack((-1.0 / lengths, np.diag(1.0 / lengths)))[None, :, :] + velocity = np.array(((0.8, 0.3, 0.2),)) + speed = np.linalg.norm(velocity, axis=1) + expected_length = speed / np.sum(velocity / lengths, axis=1) + + directional_rate = thermal.DuDt._streamline_directional_rate( + gradients, velocity + ).copy() + streamline_length = 2.0 * speed / directional_rate + np.testing.assert_allclose(streamline_length, expected_length) + + permutation = (2, 0, 3, 1) + permuted_rate = thermal.DuDt._streamline_directional_rate( + gradients[:, permutation, :], velocity + ) + np.testing.assert_allclose(permuted_rate, directional_rate) + + rotation = np.array( + ( + (0.0, -1.0, 0.0), + (1.0, 0.0, 0.0), + (0.0, 0.0, 1.0), + ) + ) + rotated_gradients = gradients @ rotation.T + rotated_velocity = velocity @ rotation.T + rotated_rate = thermal.DuDt._streamline_directional_rate( + rotated_gradients, rotated_velocity + ) + np.testing.assert_allclose(rotated_rate, directional_rate) + + +def _high_peclet_solution(tau, name): + mesh = uw.meshing.UnstructuredSimplexBox( + minCoords=(0.0, 0.0), + maxCoords=(1.0, 1.0), + cellSize=0.22, + regular=True, + ) + temperature = uw.discretisation.MeshVariable( + f"T_layer_{name}", mesh, 1, degree=1 + ) + velocity = uw.discretisation.MeshVariable( + f"U_layer_{name}", mesh, mesh.dim, degree=1 + ) + temperature.array[:, 0, 0] = temperature.coords[:, 0] + velocity.array[:, 0, 0] = 1.0 + velocity.array[:, 0, 1] = 0.0 + + thermal = uw.systems.AdvDiffusion( + mesh, temperature, velocity.sym, theta=1.0, peclet_weight=0.0 + ) + thermal.DuDt.supg_weight = 0.0 if tau == 0.0 else 1.0 + thermal.constitutive_model = uw.constitutive_models.DiffusionModel + thermal.constitutive_model.Parameters.diffusivity = 0.01 + thermal.add_dirichlet_bc(0.0, "Left") + thermal.add_dirichlet_bc(1.0, "Right") + thermal.solve(timestep=1.0e6, zero_init_guess=False) + + x = temperature.coords[:, 0] + exact = np.expm1(100.0 * x) / np.expm1(100.0) + rms_error = float(np.sqrt(np.mean((temperature.array[:, 0, 0] - exact) ** 2))) + return np.array(temperature.array), rms_error + + +def test_supg_reduces_high_peclet_oscillation_and_error(): + galerkin, galerkin_error = _high_peclet_solution(0.0, "galerkin") + supg, supg_error = _high_peclet_solution(None, "supg") + + galerkin_overshoot = max(0.0, float(galerkin.max() - 1.0)) + galerkin_undershoot = max(0.0, float(-galerkin.min())) + supg_overshoot = max(0.0, float(supg.max() - 1.0)) + supg_undershoot = max(0.0, float(-supg.min())) + + assert supg_overshoot < galerkin_overshoot + assert supg_undershoot < galerkin_undershoot + assert supg_error < 0.2 * galerkin_error + + +def _manufactured_error(cell_size, degree): + mesh = uw.meshing.UnstructuredSimplexBox( + minCoords=(0.0, 0.0), + maxCoords=(1.0, 1.0), + cellSize=cell_size, + regular=True, + qdegree=4, + ) + temperature = uw.discretisation.MeshVariable( + f"T_mms_{degree}_{cell_size}", mesh, 1, degree=degree + ) + velocity = uw.discretisation.MeshVariable( + f"U_mms_{degree}_{cell_size}", mesh, mesh.dim, degree=1 + ) + x, y = mesh.X + exact = sympy.sin(sympy.pi * x) * sympy.sin(sympy.pi * y) + diffusivity = 0.1 + temperature.array[:, 0, 0] = uw.function.evaluate( + exact, temperature.coords + ).reshape(-1) + velocity.array[:, 0, 0] = 1.0 + velocity.array[:, 0, 1] = 0.0 + + thermal = uw.systems.AdvDiffusion( + mesh, temperature, velocity.sym, theta=1.0, peclet_weight=0.0 + ) + thermal.constitutive_model = uw.constitutive_models.DiffusionModel + thermal.constitutive_model.Parameters.diffusivity = diffusivity + thermal.f = ( + sympy.pi * sympy.cos(sympy.pi * x) * sympy.sin(sympy.pi * y) + + 2.0 * diffusivity * sympy.pi**2 * exact + ) + for boundary in ("Left", "Right", "Top", "Bottom"): + thermal.add_dirichlet_bc(0.0, boundary) + thermal.solve(timestep=1.0e8, zero_init_guess=False) + + return float( + np.sqrt( + uw.maths.Integral( + mesh, fn=(temperature.sym[0] - exact) ** 2 + ).evaluate() + ) + ) + + +@pytest.mark.parametrize("degree", (1, 2)) +def test_manufactured_solution_converges_under_refinement(degree): + cell_sizes = (0.3, 0.2, 0.13) + errors = [_manufactured_error(cell_size, degree) for cell_size in cell_sizes] + final_rate = np.log(errors[-2] / errors[-1]) / np.log( + cell_sizes[-2] / cell_sizes[-1] + ) + + assert errors[0] > errors[1] > errors[2] + assert final_rate > 1.5 + + +def _spherical_implicit_response(): + mesh = uw.meshing.SphericalShell( + radiusInner=0.55, + radiusOuter=1.0, + cellSize=0.4, + qdegree=2, + ) + temperature = uw.discretisation.MeshVariable( + "T_supg_spherical", mesh, 1, degree=1 + ) + velocity = uw.discretisation.MeshVariable( + "U_supg_spherical", mesh, mesh.dim, degree=1 + ) + coords = temperature.coords + radii = np.linalg.norm(coords, axis=1) + temperature.array[:, 0, 0] = (1.0 - radii) / 0.45 + 0.01 * coords[:, 0] + velocity.array[:, 0, 0] = -0.02 * coords[:, 1] + velocity.array[:, 0, 1] = 0.02 * coords[:, 0] + velocity.array[:, 0, 2] = 0.0 + + thermal = uw.systems.AdvDiffusion( + mesh, temperature, velocity.sym, theta=0.5, peclet_weight=0.0 + ) + thermal.constitutive_model = uw.constitutive_models.DiffusionModel + thermal.constitutive_model.Parameters.diffusivity = 0.01 + thermal.add_dirichlet_bc(0.0, "Upper") + thermal.add_dirichlet_bc(1.0, "Lower") + + for _ in range(3): + thermal.solve(timestep=1.0e-3, zero_init_guess=False) + + temperature_l2_squared = float( + uw.maths.Integral(mesh, fn=temperature.sym[0] ** 2).evaluate() + ) + assert np.all(np.isfinite(temperature.array)) + assert np.all(np.isfinite(uw.function.evaluate(thermal.DuDt.tau(), mesh._centroids))) + return temperature_l2_squared, mesh + + +def _compare_spherical_with_serial(run, kind): + sys.path.insert(0, str(Path(__file__).parent / "parallel")) + from serial_reference import compare, mesh_fingerprint, serial_reference + + value, mesh = run() + # Absolute accuracy is covered by the manufactured-solution tests above; + # this gate compares partitions of the same cached spherical triangulation. + assert np.isfinite(value) and value > 0.0 + compare([value], serial_reference(__file__, kind), [1e-8], ["integral T^2"], + mesh_fingerprint(mesh), f"SUPG {kind}") + + +def test_spherical_shell_supg_is_parallel_safe(): + _compare_spherical_with_serial(_spherical_implicit_response, "implicit") + + +def test_bdf2_snapshot_restore_leaves_no_discarded_step_trace(): + uw.reset_default_model() + orchestration_model = uw.get_default_model() + mesh = uw.meshing.UnstructuredSimplexBox( + minCoords=(0.0, 0.0), + maxCoords=(1.0, 1.0), + cellSize=0.3, + regular=True, + ) + temperature = uw.discretisation.MeshVariable( + "T_supg_restart", mesh, 1, degree=1 + ) + velocity = uw.discretisation.MeshVariable( + "U_supg_restart", mesh, mesh.dim, degree=1 + ) + temperature.array[:, 0, 0] = np.sin(np.pi * temperature.coords[:, 0]) + velocity.array[:, 0, 0] = 0.1 + velocity.array[:, 0, 1] = 0.0 + + thermal = uw.systems.AdvDiffusion( + mesh, temperature, velocity.sym, order=2, theta=1.0, peclet_weight=0.0 + ) + thermal.petsc_options["ksp_rtol"] = 1e-14 + thermal.petsc_options["ksp_atol"] = 0.0 + thermal.petsc_options["snes_rtol"] = 1e-13 + thermal.petsc_options["snes_atol"] = 1e-14 + thermal.constitutive_model = uw.constitutive_models.DiffusionModel + thermal.constitutive_model.Parameters.diffusivity = 0.05 + thermal.add_dirichlet_bc(0.0, "Left") + thermal.add_dirichlet_bc(0.0, "Right") + + for _ in range(3): + thermal.solve(timestep=0.01, zero_init_guess=False) + snapshot = orchestration_model.save_state() + + orchestration_model.load_state(snapshot) + for _ in range(3): + thermal.solve(timestep=0.01, zero_init_guess=False) + reference = np.array(temperature.array) + + orchestration_model.load_state(snapshot) + thermal.solve(timestep=0.2, zero_init_guess=False) + orchestration_model.load_state(snapshot) + for _ in range(3): + thermal.solve(timestep=0.01, zero_init_guess=False) + resumed = np.array(temperature.array) + + np.testing.assert_allclose(resumed, reference, rtol=2e-14, atol=2e-14) + uw.reset_default_model() + + +def test_repeated_solves_keep_histories_and_transient_state_bounded(): + mesh = uw.meshing.UnstructuredSimplexBox( + minCoords=(0.0, 0.0), + maxCoords=(1.0, 1.0), + cellSize=0.3, + regular=True, + ) + temperature = uw.discretisation.MeshVariable( + "T_supg_lifecycle", mesh, 1, degree=1 + ) + velocity = uw.discretisation.MeshVariable( + "U_supg_lifecycle", mesh, mesh.dim, degree=1 + ) + temperature.array[:, 0, 0] = temperature.coords[:, 0] + velocity.array[:, 0, 0] = 0.1 + velocity.array[:, 0, 1] = 0.0 + + thermal = uw.systems.AdvDiffusion( + mesh, temperature, velocity.sym, theta=0.5, peclet_weight=0.0 + ) + thermal.constitutive_model = uw.constitutive_models.DiffusionModel + thermal.constitutive_model.Parameters.diffusivity = 0.05 + live_swarms = len(mesh._registered_swarms) + + for _ in range(38): + thermal.solve(timestep=0.001, zero_init_guess=False) + assert len(mesh._registered_swarms) == live_swarms + + assert len(thermal.solve_history) == 32 + assert np.all(np.isfinite(temperature.array)) + + +def _spherical_citcoms_response(): + mesh = uw.meshing.SphericalShell( + radiusInner=0.55, + radiusOuter=1.0, + cellSize=0.25, + qdegree=2, + ) + temperature = uw.discretisation.MeshVariable( + "T_citcoms_spherical", mesh, 1, degree=1 + ) + velocity = uw.discretisation.MeshVariable( + "U_citcoms_spherical", mesh, mesh.dim, degree=1 + ) + coords = temperature.coords + radii = np.linalg.norm(coords, axis=1) + temperature.array[:, 0, 0] = (1.0 - radii) / 0.45 + 0.01 * coords[:, 0] + velocity.array[:, 0, 0] = -0.02 * coords[:, 1] + velocity.array[:, 0, 1] = 0.02 * coords[:, 0] + velocity.array[:, 0, 2] = 0.0 + + manager = uw.systems.ddt.EulerianSUPGPC( + mesh, temperature, velocity.sym, + method="citcoms", + ) + thermal = uw.systems.AdvDiffusion(mesh, temperature, velocity.sym, DuDt=manager) + thermal.constitutive_model = uw.constitutive_models.DiffusionModel + thermal.constitutive_model.Parameters.diffusivity = 0.01 + thermal.add_dirichlet_bc(0.0, "Upper") + thermal.add_dirichlet_bc(1.0, "Lower") + thermal.solve(timestep=1.0e-3) + + temperature_l2_squared = float( + uw.maths.Integral(mesh, fn=temperature.sym[0] ** 2).evaluate() + ) + assert thermal.DuDt._lumped_mass.getSize() > 0 + assert np.all(np.isfinite(temperature.array)) + return temperature_l2_squared, mesh + + +def test_citcoms_spherical_shell_is_parallel_safe(): + _compare_spherical_with_serial(_spherical_citcoms_response, "citcoms") + + +def test_citcoms_snapshot_restores_startup_state_exactly(): + uw.reset_default_model() + orchestration_model = uw.get_default_model() + mesh = uw.meshing.UnstructuredSimplexBox( + minCoords=(0.0, 0.0), + maxCoords=(1.0, 1.0), + cellSize=0.3, + regular=True, + ) + temperature = uw.discretisation.MeshVariable( + "T_citcoms_restart", mesh, 1, degree=1 + ) + velocity = uw.discretisation.MeshVariable( + "U_citcoms_restart", mesh, mesh.dim, degree=1 + ) + temperature.array[:, 0, 0] = 1.0 + + manager = uw.systems.ddt.EulerianSUPGPC( + mesh, temperature, velocity.sym, + method="citcoms", tau=0.0, + ) + thermal = uw.systems.AdvDiffusion(mesh, temperature, velocity.sym, DuDt=manager) + thermal.constitutive_model = uw.constitutive_models.DiffusionModel + thermal.constitutive_model.Parameters.diffusivity = 0.0 + thermal.f = -temperature.sym[0] + + initial = orchestration_model.save_state() + thermal.solve(timestep=0.05) + reference_temperature = np.array(temperature.array) + reference_rate = np.array(thermal.DuDt.temperature_rate.array) + + orchestration_model.load_state(initial) + assert not thermal.DuDt._rate_initialised + thermal.solve(timestep=0.05) + + np.testing.assert_array_equal(temperature.array, reference_temperature) + np.testing.assert_array_equal( + thermal.DuDt.temperature_rate.array, reference_rate + ) + uw.reset_default_model() + + +if __name__ == "__main__": + sys.path.insert(0, str(Path(__file__).parent / "parallel")) + from serial_reference import emit, mesh_fingerprint + + run = _spherical_citcoms_response if sys.argv[1] == "citcoms" else _spherical_implicit_response + value, mesh = run() + emit([value], mesh_fingerprint(mesh)) diff --git a/tests/test_1115_advdiff_supg_transient.py b/tests/test_1115_advdiff_supg_transient.py new file mode 100644 index 00000000..e325b8c4 --- /dev/null +++ b/tests/test_1115_advdiff_supg_transient.py @@ -0,0 +1,180 @@ +"""Temporal convergence validation for implicit SUPG transport.""" + +import numpy as np +import pytest +import sympy + +import underworld3 as uw + + +pytestmark = [pytest.mark.level_3, pytest.mark.tier_b] + + +def _transient_state(timestep, order): + mesh = uw.meshing.UnstructuredSimplexBox( + minCoords=(0.0, 0.0), + maxCoords=(1.0, 1.0), + cellSize=0.22, + regular=True, + qdegree=3, + ) + token = str(timestep).replace(".", "p") + temperature = uw.discretisation.MeshVariable( + f"T_supg_time_{order}_{token}", mesh, 1, degree=1 + ) + velocity = uw.discretisation.MeshVariable( + f"U_supg_time_{order}_{token}", mesh, mesh.dim, degree=1 + ) + x, y = mesh.X + shape = sympy.sin(sympy.pi * x) * sympy.sin(sympy.pi * y) + diffusivity = 0.05 + advection_speed = 0.4 + temperature.array[:, 0, 0] = uw.function.evaluate( + shape, temperature.coords + ).reshape(-1) + velocity.array[:, 0, 0] = advection_speed + velocity.array[:, 0, 1] = 0.0 + + thermal = uw.systems.AdvDiffusion( + mesh, temperature, velocity.sym, order=order, theta=1.0, peclet_weight=0.0 + ) + thermal.DuDt.supg_weight = 0.0 + thermal.constitutive_model = uw.constitutive_models.DiffusionModel + thermal.constitutive_model.Parameters.diffusivity = diffusivity + for boundary in ("Left", "Right", "Top", "Bottom"): + thermal.add_dirichlet_bc(0.0, boundary) + + final_time = 0.2 + for step in range(round(final_time / timestep)): + new_time = (step + 1) * timestep + amplitude = np.exp(-new_time) + thermal.f = amplitude * ( + (-1.0 + 2.0 * diffusivity * sympy.pi**2) * shape + + advection_speed + * sympy.pi + * sympy.cos(sympy.pi * x) + * sympy.sin(sympy.pi * y) + ) + thermal.solve(timestep=timestep, zero_init_guess=False) + + return temperature.array[:, 0, 0].copy() + + +@pytest.mark.parametrize( + ("order", "minimum_rate"), + ((1, 0.9), (2, 1.8)), +) +def test_bdf_temporal_convergence(order, minimum_rate): + reference = _transient_state(0.003125, 2) + timesteps = (0.05, 0.025, 0.0125) + errors = [ + np.linalg.norm(_transient_state(timestep, order) - reference) + / np.sqrt(reference.size) + for timestep in timesteps + ] + rates = [ + np.log(errors[index] / errors[index + 1]) / np.log(2.0) + for index in range(2) + ] + + assert errors[0] > errors[1] > errors[2] + assert min(rates) > minimum_rate + + +def _citcoms_decay_error(timestep): + mesh = uw.meshing.UnstructuredSimplexBox( + minCoords=(0.0, 0.0), + maxCoords=(1.0, 1.0), + cellSize=0.5, + regular=True, + ) + token = str(timestep).replace(".", "p") + temperature = uw.discretisation.MeshVariable( + f"T_citcoms_decay_{token}", mesh, 1, degree=1 + ) + velocity = uw.discretisation.MeshVariable( + f"U_citcoms_decay_{token}", mesh, mesh.dim, degree=1 + ) + temperature.array[:, 0, 0] = 1.0 + + manager = uw.systems.ddt.EulerianSUPGPC( + mesh, temperature, velocity.sym, + method="citcoms", tau=0.0, + ) + thermal = uw.systems.AdvDiffusion(mesh, temperature, velocity.sym, DuDt=manager) + thermal.constitutive_model = uw.constitutive_models.DiffusionModel + thermal.constitutive_model.Parameters.diffusivity = 0.0 + thermal.f = -temperature.sym[0] + + for _ in range(round(1.0 / timestep)): + thermal.solve(timestep=timestep) + + return abs(float(np.mean(temperature.array[:, 0, 0])) - np.exp(-1.0)) + + +def test_citcoms_predictor_corrector_is_second_order_for_scalar_decay(): + errors = [_citcoms_decay_error(dt) for dt in (0.1, 0.05, 0.025)] + rates = [ + np.log(errors[index] / errors[index + 1]) / np.log(2.0) + for index in range(2) + ] + + assert errors[0] > errors[1] > errors[2] + assert min(rates) > 1.9 + + +def _citcoms_rotation_return_error(cell_size): + mesh = uw.meshing.Annulus( + radiusOuter=1.0, + radiusInner=0.5, + cellSize=cell_size, + qdegree=4, + ) + token = str(cell_size).replace(".", "p") + temperature = uw.discretisation.MeshVariable( + f"T_citcoms_rotation_{token}", mesh, 1, degree=1 + ) + velocity = uw.discretisation.MeshVariable( + f"U_citcoms_rotation_{token}", mesh, mesh.dim, degree=1 + ) + x, y = mesh.X + initial = sympy.exp(-30.0 * (x**2 + (y - 0.75) ** 2)) + + temperature.array[:, 0, 0] = uw.function.evaluate( + initial, temperature.coords + ).reshape(-1) + velocity.array[:, 0, 0] = -2.0 * np.pi * velocity.coords[:, 1] + velocity.array[:, 0, 1] = 2.0 * np.pi * velocity.coords[:, 0] + + manager = uw.systems.ddt.EulerianSUPGPC( + mesh, temperature, velocity.sym, + method="citcoms", + ) + thermal = uw.systems.AdvDiffusion(mesh, temperature, velocity.sym, DuDt=manager) + thermal.constitutive_model = uw.constitutive_models.DiffusionModel + thermal.constitutive_model.Parameters.diffusivity = 0.0 + + step_count = int(np.ceil(1.0 / thermal.estimate_dt())) + timestep = 1.0 / step_count + for _ in range(step_count): + thermal.solve(timestep=timestep) + + error = float( + np.sqrt( + uw.maths.Integral( + mesh, fn=(temperature.sym[0] - initial) ** 2 + ).evaluate() + ) + ) + initial_norm = float( + np.sqrt(uw.maths.Integral(mesh, fn=initial**2).evaluate()) + ) + return error / initial_norm + + +def test_citcoms_rotation_return_error_decreases_with_refinement(): + coarse_error = _citcoms_rotation_return_error(0.2) + fine_error = _citcoms_rotation_return_error(0.1) + + assert fine_error < 0.9 * coarse_error + assert fine_error < 0.7 diff --git a/tests/test_1116_supg_unified.py b/tests/test_1116_supg_unified.py new file mode 100644 index 00000000..2600a4a9 --- /dev/null +++ b/tests/test_1116_supg_unified.py @@ -0,0 +1,176 @@ +"""Shared SUPG integration, restart, and pre-migration equivalence.""" + +import importlib.util +import os +import sys + +import numpy as np +import pytest +import sympy + +import underworld3 as uw + +pytestmark = [pytest.mark.level_2, pytest.mark.tier_b] + + +def _problem(dim, tag, cellsize=0.25): + mesh = uw.meshing.UnstructuredSimplexBox( + minCoords=(0.0,) * dim, maxCoords=(1.0,) * dim, + cellSize=cellsize, qdegree=4, regular=False, + ) + temperature = uw.discretisation.MeshVariable(f"T_{tag}", mesh, 1, degree=1) + velocity = uw.discretisation.MeshVariable(f"U_{tag}", mesh, dim, degree=1) + temperature.array[:, 0, 0] = temperature.coords[:, 0] + velocity.array[:, 0, :] = 0.2 + return mesh, temperature, velocity + + +@pytest.mark.parametrize("dim", [2, 3]) +@pytest.mark.parametrize("method", ["citcoms", "pc_converged"]) +def test_pc_matches_pre_migration_implementation(dim, method): + """Optional release gate against the last PR689 head, commit f41bcd2f. + + Both assemblers see the same mesh, partition, source, changing velocity + and timestep sequence, with separate temperature and rate fields. + The frozen source is an external test artifact, not another installed solver. + """ + baseline = os.environ.get("UW_SUPG_BASELINE_FILE") + if baseline is None: + pytest.skip("Set UW_SUPG_BASELINE_FILE to the frozen f41bcd2f module.") + spec = importlib.util.spec_from_file_location("_supg_baseline", baseline) + module = importlib.util.module_from_spec(spec) + sys.modules[spec.name] = module + spec.loader.exec_module(module) + + mesh, temperature, velocity = _problem(dim, f"migration_{method}_{dim}") + reference = uw.discretisation.MeshVariable("T_reference", mesh, 1, degree=1) + rate = uw.discretisation.MeshVariable("Tdot", mesh, 1, degree=1) + reference_rate = uw.discretisation.MeshVariable("Tdot_reference", mesh, 1, degree=1) + shape = sympy.prod(sympy.sin(sympy.pi * x) for x in mesh.X) + temperature.array[:, 0, 0] = uw.function.evaluate(shape, temperature.coords).reshape(-1) + reference.array[...] = temperature.array + manager = uw.systems.ddt.EulerianSUPGPC( + mesh, temperature, velocity.sym, + method=method, temperature_rate_field=rate, + ) + current = uw.systems.AdvDiffusion(mesh, temperature, velocity.sym, DuDt=manager) + previous = module.SNES_AdvectionDiffusion_SUPG( + mesh, reference, velocity.sym, time_integrator=method, + temperature_rate_field=reference_rate, + ) + for solver in (current, previous): + solver.constitutive_model = uw.constitutive_models.DiffusionModel + solver.constitutive_model.Parameters.diffusivity = 0.01 + solver.f = 0.1 * shape + for boundary in mesh.boundaries: + if boundary.name not in ("All_Boundaries", "Null_Boundary"): + solver.add_dirichlet_bc(0.0, boundary.name) + + for step in range(6): + velocity.array[:, 0, :] = 0.2 * (1.0 + step / 10.0) + dt = min((0.002, 0.003, 0.0015)[step % 3], float(current.estimate_dt())) + np.testing.assert_allclose( + current.estimate_dt(), previous.estimate_dt(), rtol=1e-12, atol=1e-14) + current.solve(timestep=dt) + previous.solve(timestep=dt) + np.testing.assert_allclose(temperature.array, reference.array, rtol=1e-11, atol=1e-12) + np.testing.assert_allclose(rate.array, reference_rate.array, rtol=1e-10, atol=1e-11) + + +@pytest.mark.parametrize("method", ["citcoms", "pc_converged", "be", "cn", "bdf2"]) +@pytest.mark.parametrize("disk", [False, True]) +def test_snapshot_restores_fields_and_timestep_estimator(method, disk, tmp_path): + uw.reset_default_model() + orchestration_model = uw.get_default_model() + mesh, temperature, velocity = _problem(2, "snapshot") + is_pc = method in ("citcoms", "pc_converged") + if is_pc: + manager = uw.systems.ddt.EulerianSUPGPC( + mesh, temperature, velocity.sym, method=method + ) + thermal = uw.systems.AdvDiffusion(mesh, temperature, velocity.sym, DuDt=manager) + else: + thermal = uw.systems.AdvDiffusion( + mesh, temperature, velocity.sym, + order=2 if method == "bdf2" else 1, + theta=0.5 if method == "cn" else 1.0, + peclet_weight=0.0, + ) + manager = thermal.DuDt + # Replay is compared near machine precision, independently of the default + # stopping tolerance and the preconditioner rebuilt after a discarded step. + thermal.petsc_options["ksp_rtol"] = 1e-14 + thermal.petsc_options["ksp_atol"] = 0.0 + thermal.petsc_options["snes_rtol"] = 1e-13 + thermal.petsc_options["snes_atol"] = 1e-14 + thermal.constitutive_model.Parameters.diffusivity = 0.05 + thermal.add_dirichlet_bc(0.0, "Left") + thermal.add_dirichlet_bc(1.0, "Right") + for _ in range(3): + thermal.solve(timestep=0.002) + if disk: + path = uw.mpi.comm.bcast(str(tmp_path / "thermal.h5"), root=0) + snapshot = orchestration_model.save_state(file=path) + else: + snapshot = orchestration_model.save_state() + saved_temperature = np.array(temperature.array) + estimate = thermal.estimate_dt() + thermal.solve(timestep=0.003) + expected = np.array(temperature.array) + expected_state = thermal.state + expected_manager_state = manager.state + expected_rate = np.array(manager.temperature_rate.array) if is_pc else None + orchestration_model.load_state(snapshot) + np.testing.assert_array_equal(temperature.array, saved_temperature) + assert thermal.estimate_dt() == pytest.approx(estimate, rel=1e-14) + thermal.solve(timestep=0.01) + orchestration_model.load_state(snapshot) + thermal.solve(timestep=0.003) + # Rebuilding an implicit Krylov solve can change final rounding, but the + # restored fields above must be exact and replay must agree near machine precision. + np.testing.assert_allclose(temperature.array, expected, rtol=2e-14, atol=2e-14) + actual_state = thermal.state + assert actual_state.last_timestep == expected_state.last_timestep + if expected_state.last_change_rate is not None: + # The estimator is max(|T_new - T_old|) / dt. Propagate the field + # assertion's absolute-plus-relative bound through that division. + field_bound = max(uw.mpi.comm.allgather( + 2e-14 * (1.0 + float(np.max(np.abs(expected), initial=0.0))))) + assert actual_state.last_change_rate == pytest.approx( + expected_state.last_change_rate, rel=0.0, + abs=field_bound / expected_state.last_timestep) + if expected_rate is not None: + assert manager.state == expected_manager_state + np.testing.assert_array_equal(manager.temperature_rate.array, expected_rate) + uw.reset_default_model() + + +def test_citcoms_does_not_allocate_unused_multistep_history(): + mesh, temperature, velocity = _problem(2, "history") + manager = uw.systems.ddt.EulerianSUPGPC( + mesh, temperature, velocity.sym, + method="citcoms", + ) + thermal = uw.systems.AdvDiffusion(mesh, temperature, velocity.sym, DuDt=manager) + assert thermal.DuDt is manager + assert manager.psi_star == [] + assert manager.temperature_rate is not None + with pytest.raises(ValueError, match="stability"): + thermal.estimate_dt(basis="accuracy") + with pytest.raises(ValueError, match="theta"): + thermal.theta = 0.5 + + +def test_empty_partition_is_rejected_on_every_rank(): + mesh, temperature, velocity = _problem(2, "empty_partition", cellsize=0.5) + counts = uw.mpi.comm.allgather( + mesh.dm.getHeightStratum(0)[1] - mesh.dm.getHeightStratum(0)[0]) + if min(counts) > 0: + pytest.skip(f"This partition has no empty ranks: {counts}") + manager = uw.systems.ddt.EulerianSUPGPC( + mesh, temperature, velocity.sym, + method="citcoms", + ) + thermal = uw.systems.AdvDiffusion(mesh, temperature, velocity.sym, DuDt=manager) + with pytest.raises(NotImplementedError, match="on every rank"): + thermal.estimate_dt() diff --git a/tests/test_1117_supg_pc2_analytical.py b/tests/test_1117_supg_pc2_analytical.py new file mode 100644 index 00000000..87204591 --- /dev/null +++ b/tests/test_1117_supg_pc2_analytical.py @@ -0,0 +1,253 @@ +"""Independent exact-solution gates for P1 CitcomS predictor-corrector transport. + +No Stokes solve or fine numerical reference is used. The channel pulse is the +translated, initially smoothed pulse of Calhoun & LeVeque (2000), section 6.1, +equations 45-48, with rescaled width, origin and time: +https://doi.org/10.1006/jcph.1999.6369 + +The rotation uses uw.analytic.RotatingGaussian. The spherical test follows +directly from (r*T)_t = kappa*(r*T)_rr. All spatial refinements use one fixed +timestep selected from the most restrictive mesh. Set UW_PC2_RESULTS to retain +small HDF5 metrics files. Separate jobs use separate UW_MESH_CACHE_DIR paths; +Gmsh SHA256 fingerprints must match before comparing their numerical results. +""" + +import hashlib +import math +import os +from pathlib import Path +import time + +import numpy as np +import pytest +import sympy +from mpi4py import MPI + +import underworld3 as uw +from underworld3.meshing._mesh_files import mesh_file_path + +pytestmark = [pytest.mark.level_2, pytest.mark.tier_b] + + +def _pulse(x, time_value, speed, diffusivity): + # Nonzero initial smoothing avoids the unresolved top-hat singularity. + width = sympy.sqrt(4.0 * (0.01 + diffusivity * time_value)) + distance = x - speed * time_value + return (sympy.erf((0.25 - distance) / width) + + sympy.erf((0.25 + distance) / width)) / 2 + + +def _problem(case, h, dim=2, speed=0.0, diffusivity=0.0): + filename = mesh_file_path(f"pc2_{case}_{dim}d_h{h:g}.msh") + if case == "shell": + mesh = uw.meshing.SphericalShell( + radiusInner=0.55, radiusOuter=1.0, cellSize=h, qdegree=4, filename=filename) + else: + lower = (-2.0, -2.0) if case == "rotation" else (-1.5,) + (-0.25,) * (dim - 1) + upper = tuple(-value for value in lower) + mesh = uw.meshing.UnstructuredSimplexBox( + minCoords=lower, maxCoords=upper, cellSize=h, + qdegree=4, regular=False, filename=filename) + fingerprint = None + if uw.mpi.rank == 0: + try: + fingerprint = (None, hashlib.sha256(Path(filename).read_bytes()).hexdigest()) + except OSError as exc: + fingerprint = (str(exc), None) + failure, mesh_sha256 = uw.mpi.comm.bcast(fingerprint, root=0) + assert failure is None, failure + temperature = uw.discretisation.MeshVariable("T", mesh, 1, degree=1) + velocity = uw.discretisation.MeshVariable("U", mesh, mesh.dim, degree=1) + velocity.array[...] = 0.0 + + if case == "rotation": + oracle = uw.analytic.RotatingGaussian( + mesh, sigma=0.2, centre_radius=0.5, omega=1.0, + diffusivity=diffusivity) + initial = oracle.at(0.0) + end_time = float(sympy.pi / 2) + exact = oracle.at(end_time) + velocity.array[:, 0, 0] = -velocity.coords[:, 1] + velocity.array[:, 0, 1] = velocity.coords[:, 0] + boundaries = ("Left", "Right", "Top", "Bottom") + # Max Gaussian tail on the square throughout this quarter turn. + variance = 0.2**2 + 2 * diffusivity * end_time + boundary_tail = 0.2**2 / variance * math.exp(-1.5**2 / (2 * variance)) + elif case == "shell": + radius = sympy.sqrt(sum(x**2 for x in mesh.X)) + initial = 0.55 / radius * sympy.sin(sympy.pi * (radius - 0.55) / 0.45) + end_time = 0.2 + exact = initial * sympy.exp(-diffusivity * (sympy.pi / 0.45)**2 * end_time) + boundaries = ("Lower", "Upper") + boundary_tail = 0.0 + else: + initial = _pulse(mesh.X[0], 0.0, speed, diffusivity) + end_time = 0.2 + exact = _pulse(mesh.X[0], end_time, speed, diffusivity) + velocity.array[:, 0, 0] = speed + # Zero transverse diffusive flux is exact. End-wall tails are bounded + # analytically, not silently treated as exactly zero. + boundaries = ("Left", "Right") + boundary_tail = math.erfc( + (1.5 - 0.25 - abs(speed) * end_time) + / math.sqrt(4 * (0.01 + diffusivity * end_time))) + assert boundary_tail < 1e-10, boundary_tail + temperature.array[:, 0, 0] = uw.function.evaluate( + initial, temperature.coords).reshape(-1) + manager = uw.systems.ddt.EulerianSUPGPC( + mesh, temperature, velocity.sym, + method="citcoms", adv_gamma=0.5, corrector_steps=2, + ) + thermal = uw.systems.AdvDiffusion(mesh, temperature, velocity.sym, DuDt=manager) + thermal.constitutive_model.Parameters.diffusivity = diffusivity + for boundary in boundaries: + thermal.add_dirichlet_bc(0.0, boundary) + return mesh, temperature, thermal, initial, exact, end_time, boundary_tail, mesh_sha256 + + +def _integral(mesh, expression): + return float(uw.maths.Integral(mesh, fn=expression).evaluate()) + + +def _save_result(name, metrics): + """Optional rank-zero output, with write failures propagated collectively.""" + error = None + if uw.mpi.rank == 0: + print("PC2_ANALYTICAL " + name + " " + " ".join( + f"{key}={value}" if isinstance(value, str) else f"{key}={value:.12g}" + for key, value in metrics.items()), flush=True) + directory = os.environ.get("UW_PC2_RESULTS") + if directory: + try: + import h5py + + target = Path(directory) / f"ncpus_{uw.mpi.size}" + target.mkdir(parents=True, exist_ok=True) + with h5py.File(target / f"{name}.h5", "w") as output: + output.attrs["method"] = "citcoms_pc2" + for key, value in metrics.items(): + output[key] = value + except Exception as exc: + error = f"Cannot write PC2 analytical metrics: {exc}" + error = uw.mpi.comm.bcast(error, root=0) + assert error is None, error + + +def _step_count(problems): + dt_limit = min(float(problem[2].estimate_dt()) for problem in problems) + end_time = problems[0][5] + # Round down before choosing an integer number of steps, avoiding a + # partition-dependent ceil when a stability estimate differs by roundoff. + dt_cap = 2.0**math.floor(math.log2(min(0.005, 0.5 * dt_limit))) + return math.ceil(end_time / dt_cap) + + +def _advance(problem, h, steps): + mesh, temperature, thermal, initial, exact, end_time, tail, mesh_sha256 = problem + timestep = end_time / steps + norm_squared = _integral(mesh, exact**2) + interpolation_error = math.sqrt(_integral( + mesh, (temperature.sym[0] - initial)**2) / _integral(mesh, initial**2)) + start = time.perf_counter() + for _ in range(steps): + thermal.solve(timestep=timestep) + solve_seconds = uw.mpi.comm.allreduce(time.perf_counter() - start, op=MPI.MAX) + error = math.sqrt(_integral(mesh, (temperature.sym[0] - exact)**2) / norm_squared) + heat = _integral(mesh, temperature.sym[0]) + exact_heat = _integral(mesh, exact) + nodal = temperature.array[:, 0, 0] + minimum = uw.mpi.comm.allreduce(float(np.min(nodal, initial=np.inf)), op=MPI.MIN) + maximum = uw.mpi.comm.allreduce(float(np.max(nodal, initial=-np.inf)), op=MPI.MAX) + return dict( + cellsize=h, dim=mesh.dim, ncpus=uw.mpi.size, timestep=timestep, + steps=steps, end_time=end_time, relative_l2=error, + initial_relative_l2=interpolation_error, + temperature_integral=heat, exact_temperature_integral=exact_heat, + relative_heat_error=(heat - exact_heat) / exact_heat, + minimum=minimum, maximum=maximum, solve_seconds=solve_seconds, + boundary_tail_bound=tail, volume=_integral(mesh, sympy.Integer(1)), + mesh_sha256=mesh_sha256) + + +def _spatial_refinement(case, sizes, dim=2, speed=0.0, diffusivity=0.0): + problems = [_problem(case, h, dim, speed, diffusivity) for h in sizes] + steps = _step_count(problems) + errors = [] + for h, problem in zip(sizes, problems): + mesh, temperature, thermal = problem[:3] + metrics = _advance(problem, h, steps) + if case == "rotation": + centre = [_integral(mesh, coordinate * temperature.sym[0]) / metrics["temperature_integral"] + for coordinate in mesh.X] + metrics["phase_error_radians"] = math.atan2(centre[1], centre[0]) - metrics["end_time"] + if speed != 0 or case == "rotation": + tau = uw.function.evaluate(thermal.DuDt.tau(), mesh._centroids) + assert uw.mpi.comm.allreduce(bool(np.any(tau > 0)), op=MPI.LOR) + name = f"{case}_{mesh.dim}d_u{speed:g}_k{diffusivity:g}_h{h:g}" + _save_result(name, metrics) + assert (np.isfinite(metrics["relative_l2"]) + and metrics["minimum"] > -0.05 and metrics["maximum"] < 1.05), metrics + errors.append(metrics["relative_l2"]) + return errors + + +@pytest.mark.parametrize("dim", [2, 3]) +@pytest.mark.parametrize("speed,diffusivity", [(1.0, 0.0), (0.0, 0.01), (1.0, 0.01)]) +def test_pc2_exact_channel_pulse(dim, speed, diffusivity): + errors = _spatial_refinement("pulse", (0.125, 0.0625), dim, speed, diffusivity) + assert errors[1] < 0.7 * errors[0], errors + assert errors[1] < 0.05, errors + + +def test_pc2_exact_rotating_gaussian(): + errors = _spatial_refinement("rotation", (0.125, 0.0625)) + assert errors[1] < 0.7 * errors[0], errors + assert errors[1] < 0.10, errors + + +def test_pc2_exact_spherical_diffusion(): + errors = _spatial_refinement("shell", (0.125, 0.0625), dim=3, diffusivity=0.02) + assert errors[1] < 0.7 * errors[0], errors + assert errors[1] < 0.08, errors + + +def test_pc2_spherical_diffusion_timestep_sensitivity(): + """Separate finite-step changes from the continuum spatial error at h=1/8.""" + problem = _problem("shell", 0.125, dim=3, diffusivity=0.02) + mesh, temperature, thermal = problem[:3] + difference = uw.discretisation.MeshVariable("T_difference", mesh, 1, degree=1) + initial_values = np.array(temperature.array) + initial_state = thermal.DuDt.state + base_steps = _step_count([problem]) + metrics, solutions = [], [] + for factor in (1, 2, 4): + temperature.array[...] = initial_values + thermal.DuDt.temperature_rate.array[...] = 0.0 + thermal.DuDt.state = initial_state + metrics.append(_advance(problem, 0.125, base_steps * factor)) + solutions.append(np.array(temperature.array)) + norm_squared = _integral(mesh, problem[4]**2) + changes = [] + for index in (0, 1): + difference.array[...] = solutions[index] - solutions[index + 1] + changes.append(math.sqrt(_integral(mesh, difference.sym[0]**2) / norm_squared)) + metrics[index]["relative_difference_to_next_dt"] = changes[-1] + if min(changes) > 0: + metrics[-1]["observed_time_order"] = math.log2(changes[0] / changes[1]) + for factor, result in zip((1, 2, 4), metrics): + _save_result(f"shell_time_h0.125_dtdiv{factor}", result) + assert all(np.isfinite(item["relative_l2"]) for item in metrics), metrics + assert all(item["minimum"] > -0.05 and item["maximum"] < 1.05 for item in metrics), metrics + assert changes[1] < changes[0], changes + assert changes[1] < 0.05 * metrics[-1]["relative_l2"], (changes, metrics) + + +def test_exact_spherical_diffusion_satisfies_radial_heat_equation(): + r, t, ri, thickness, kappa = sympy.symbols("r t ri d kappa", positive=True) + exact = ri / r * sympy.sin(sympy.pi * (r - ri) / thickness) * sympy.exp( + -kappa * (sympy.pi / thickness)**2 * t) + residual = sympy.diff(exact, t) - kappa * ( + sympy.diff(exact, r, 2) + 2 / r * sympy.diff(exact, r)) + assert sympy.simplify(residual) == 0 + assert sympy.simplify(exact.subs(r, ri)) == 0 + assert sympy.simplify(exact.subs(r, ri + thickness)) == 0 diff --git a/tests/test_1118_pc2_diffusion_time.py b/tests/test_1118_pc2_diffusion_time.py new file mode 100644 index 00000000..84235ac2 --- /dev/null +++ b/tests/test_1118_pc2_diffusion_time.py @@ -0,0 +1,263 @@ +"""Isolate finite-correction PC2 time error without Stokes or a fine mesh. + +Analytical P1 element integrals independently supply M, K and D=diag(M*1). +The exact semidiscrete solution is a generalized eigenmode exp(-lambda*t). +Small dense SciPy matrices are an independent test oracle, not solver code. + +With consistent M in the residual and two D-preconditioned corrections, +the dt->0 operator is (2I-D^-1 M)D^-1 K, not M^-1 K or D^-1 K. This test +documents that limitation; it does NOT certify second-order PDE accuracy. +""" + +import hashlib +import math + +import numpy as np +import pytest +from scipy.linalg import eigh, expm + +import underworld3 as uw +from underworld3.meshing.smoothing import _owned_cell_mask, _tet_cells, _tri_cells + +pytestmark = [pytest.mark.level_2, pytest.mark.tier_b] + + +def _p1_matrices(mesh): + """Integrate affine basis functions independently of the solver assembly.""" + cells = (_tri_cells if mesh.dim == 2 else _tet_cells)(mesh.dm) + local_cells = np.asarray(mesh.X.coords)[cells[_owned_cell_mask(mesh.dm)]] + vertices = np.concatenate(uw.mpi.comm.allgather(local_cells)) + coords = np.unique(vertices.reshape(-1, mesh.dim).round(12), axis=0) + indices = {tuple(point): index for index, point in enumerate(coords)} + connectivity = np.array([ + [indices[tuple(point.round(12))] for point in cell] for cell in vertices + ]) + canonical = np.sort(connectivity, axis=1) + canonical = canonical[np.lexsort(canonical.T[::-1])] + fingerprint = hashlib.sha256(coords.tobytes() + canonical.tobytes()).hexdigest() + mass = np.zeros((len(coords), len(coords))) + stiffness = np.zeros_like(mass) + for ids, cell in zip(connectivity, vertices): + affine = np.column_stack([np.ones(mesh.dim + 1), cell]) + gradients = np.linalg.inv(affine)[1:, :].T + volume = abs(np.linalg.det(affine)) / math.factorial(mesh.dim) + mass[np.ix_(ids, ids)] += volume * ( + np.ones((mesh.dim + 1, mesh.dim + 1)) + np.eye(mesh.dim + 1) + ) / ((mesh.dim + 1) * (mesh.dim + 2)) + stiffness[np.ix_(ids, ids)] += 0.1 * volume * gradients @ gradients.T + np.testing.assert_allclose(mass.sum(), 1.0, atol=1e-12) + np.testing.assert_allclose(stiffness.sum(axis=1), 0.0, atol=1e-12) + return coords, mass, stiffness, len(vertices), fingerprint + + +def _norm(values, mass): + return float(np.sqrt(values @ mass @ values)) + + +def _orders(values): + return np.log2(np.asarray(values[:-1]) / values[1:]) + + +@pytest.fixture(params=[2, 3]) +def diffusion(request): + dim = request.param + mesh = uw.meshing.UnstructuredSimplexBox( + minCoords=(0.0,) * dim, maxCoords=(1.0,) * dim, + cellSize=0.25, qdegree=4, regular=False, + ) + temperature = uw.discretisation.MeshVariable("T", mesh, 1, degree=1) + velocity = uw.discretisation.MeshVariable("U", mesh, dim, degree=1) + velocity.array[...] = 0.0 + manager = uw.systems.ddt.EulerianSUPGPC( + mesh, temperature, velocity.sym, + method="citcoms", + ) + thermal = uw.systems.AdvDiffusion(mesh, temperature, velocity.sym, DuDt=manager) + thermal.constitutive_model.Parameters.diffusivity = 0.1 + # Natural homogeneous Neumann boundaries remove constrained-DOF effects. + coords, mass, stiffness, cell_count, fingerprint = _p1_matrices(mesh) + indices = {tuple(point): index for index, point in enumerate(coords)} + local_ids = np.array([ + indices[tuple(point.round(12))] for point in np.asarray(temperature.coords) + ]) + eigenvalues, eigenvectors = eigh(stiffness, mass) + np.testing.assert_allclose(eigenvalues[0], 0.0, atol=1e-12) + initial = eigenvectors[:, 1].copy() + initial *= np.sign(initial[np.argmax(np.abs(initial))]) + initial /= np.max(np.abs(initial)) + uw.pprint( + f"PC2_ORACLE dim={dim} cells={cell_count} vertices={len(coords)} " + f"mesh_sha256={fingerprint} lambda={eigenvalues[1]:.12g}") + return thermal, temperature, local_ids, mass, stiffness, initial, eigenvalues, eigenvectors + + +def test_two_corrections_match_independent_diffusion_map(diffusion): + thermal, temperature, ids, mass, stiffness, initial, eigenvalues, _ = diffusion + lumped = mass.sum(axis=1) + H = mass / lumped[:, None] + J = stiffness / lumped[:, None] + identity = np.eye(len(initial)) + final_time = 0.1 + exact = initial * np.exp(-eigenvalues[1] * final_time) + effective = expm(-final_time * (2 * identity - H) @ J) @ initial + initial_state = thermal.DuDt.state + dt_limit = thermal.estimate_dt() + solutions, effective_errors = [], [] + for steps in (16, 32, 64, 128): + dt = final_time / steps + assert dt < dt_limit + temperature.array[:, 0, 0] = initial[ids] + thermal.DuDt.temperature_rate.array[...] = 0.0 + thermal.DuDt.state = initial_state + expected = initial.copy() + rate = -J @ initial + # Algebraically eliminate both corrections; do not call UW3 residuals. + B = (2 * identity - H - 0.5 * dt * J) @ J + discrepancy = np.zeros(2) + for _ in range(steps): + predictor = expected + 0.5 * dt * rate + rate = -B @ predictor + expected = predictor + 0.5 * dt * rate + thermal.solve(timestep=dt) + discrepancy = np.maximum(discrepancy, [ + np.max(np.abs(temperature.array[:, 0, 0] - expected[ids])), + np.max(np.abs(thermal.DuDt.temperature_rate.array[:, 0, 0] - rate[ids])), + ]) + discrepancy = np.max(uw.mpi.comm.allgather(discrepancy), axis=0) + assert discrepancy[0] < 1e-11 and discrepancy[1] < 1e-10, discrepancy + actual = np.zeros(len(initial)) + for local_ids, values in uw.mpi.comm.allgather( + (ids, np.array(temperature.array[:, 0, 0]))): + actual[local_ids] = values + solutions.append(actual) + effective_error = _norm(actual - effective, mass) / _norm(effective, mass) + effective_errors.append(effective_error) + uw.pprint( + f"PC2_DIFFUSION dim={thermal.mesh.dim} steps={steps} dt={dt:.12g} " + f"consistent_error={_norm(actual-exact, mass)/_norm(exact, mass):.12g} " + f"effective_error={effective_error:.12g} " + f"T_map_error={discrepancy[0]:.12g} Tdot_map_error={discrepancy[1]:.12g}") + changes = [_norm(a - b, mass) for a, b in zip(solutions, solutions[1:])] + rates = _orders(changes) + # Detect the known finite-correction limit, not a general accuracy guarantee. + assert np.all((0.9 < rates) & (rates < 1.15)), rates + assert np.all((0.9 < _orders(effective_errors)) & (_orders(effective_errors) < 1.15)) + uw.pprint(f"PC2_TIME_ORDER dim={thermal.mesh.dim} rates={rates.tolist()}") + + +def test_exact_matrix_controls_separate_mass_and_startup(diffusion): + thermal, _, _, mass, stiffness, initial, eigenvalues, eigenvectors = diffusion + D = mass.sum(axis=1) + J = stiffness / D[:, None] + final_time = 0.1 + consistent_exact = initial * np.exp(-eigenvalues[1] * final_time) + lumped_exact = expm(-final_time * J) @ initial + errors = {"consistent_cn": [], "cn_lumped_startup": [], "lumped_pc2": []} + for steps in (16, 32, 64, 128): + dt = final_time / steps + factors = (1 - 0.5 * dt * eigenvalues) / (1 + 0.5 * dt * eigenvalues) + consistent = initial * factors[1]**steps + errors["consistent_cn"].append(_norm(consistent - consistent_exact, mass)) + # Exactly converged corrections cannot repair an inconsistent first rate. + predictor = initial - 0.5 * dt * J @ initial + amplitudes = (eigenvectors.T @ mass @ predictor) / (1 + 0.5 * dt * eigenvalues) + bad_start = eigenvectors @ (factors**(steps - 1) * amplitudes) + errors["cn_lumped_startup"].append(_norm(bad_start - consistent_exact, mass)) + # A genuinely lumped residual has H=I; this is a diagnostic alternative, + # not a replacement of the CitcomS mode installed in UW3. + values, rate = initial.copy(), -J @ initial + B = (np.eye(len(initial)) - 0.5 * dt * J) @ J + for _ in range(steps): + predictor = values + 0.5 * dt * rate + rate = -B @ predictor + values = predictor + 0.5 * dt * rate + errors["lumped_pc2"].append(_norm(values - lumped_exact, mass)) + for name, values in errors.items(): + rates = _orders(values) + uw.pprint(f"PC2_CONTROL dim={thermal.mesh.dim} name={name} errors={values} rates={rates.tolist()}") + if name == "cn_lumped_startup": + assert np.all((0.9 < rates) & (rates < 1.15)), rates + else: + assert np.all((1.9 < rates) & (rates < 2.2)), rates + + +def test_uw3_cn_is_second_order_for_discrete_diffusion(diffusion): + pc2, _, ids, mass, _, initial, eigenvalues, _ = diffusion + exact = initial * np.exp(-0.1 * eigenvalues[1]) + errors = [] + for steps in (16, 32, 64, 128): + temperature = uw.discretisation.MeshVariable( + f"T_cn_{steps}", pc2.mesh, 1, degree=1) + temperature.array[:, 0, 0] = initial[ids] + thermal = uw.systems.AdvDiffusion( + pc2.mesh, temperature, pc2.V_fn, order=1, theta=0.5, peclet_weight=0.0 + ) + thermal.constitutive_model.Parameters.diffusivity = 0.1 + # Time error on the finest dt is about 1e-8; solver error must be smaller. + thermal.petsc_options["ksp_rtol"] = 1e-14 + thermal.petsc_options["ksp_atol"] = 0.0 + thermal.petsc_options["snes_rtol"] = 1e-13 + thermal.petsc_options["snes_atol"] = 1e-14 + dt = 0.1 / steps + for _ in range(steps): + thermal.solve(timestep=dt) + factor = (1 - 0.5 * dt * eigenvalues[1]) / (1 + 0.5 * dt * eigenvalues[1]) + expected = initial * factor**steps + discrepancy = max(uw.mpi.comm.allgather(float(np.max( + np.abs(temperature.array[:, 0, 0] - expected[ids]))))) + assert discrepancy < 1e-10, discrepancy + actual = np.zeros(len(initial)) + for local_ids, values in uw.mpi.comm.allgather( + (ids, np.array(temperature.array[:, 0, 0]))): + actual[local_ids] = values + errors.append(_norm(actual - exact, mass) / _norm(exact, mass)) + uw.pprint( + f"UW3_CN_DIFFUSION dim={pc2.mesh.dim} steps={steps} dt={dt:.12g} " + f"relative_error={errors[-1]:.12g} map_error={discrepancy:.12g}") + rates = _orders(errors) + assert np.all((1.9 < rates) & (rates < 2.2)), rates + uw.pprint(f"UW3_CN_TIME_ORDER dim={pc2.mesh.dim} rates={rates.tolist()}") + + +def test_converged_pc_is_second_order_for_discrete_diffusion(diffusion): + pc2, _, ids, mass, _, initial, eigenvalues, _ = diffusion + temperature = uw.discretisation.MeshVariable( + "T_pc_converged", pc2.mesh, 1, degree=1) + manager = uw.systems.ddt.EulerianSUPGPC( + pc2.mesh, temperature, pc2.V_fn, + method="pc_converged", + corrector_rtol=1.0e-12, + corrector_atol=1.0e-14, + max_corrector_steps=200, + ) + thermal = uw.systems.AdvDiffusion(pc2.mesh, temperature, pc2.V_fn, DuDt=manager) + thermal.constitutive_model.Parameters.diffusivity = 0.1 + initial_state = thermal.DuDt.state + exact = initial * np.exp(-0.1 * eigenvalues[1]) + errors = [] + for steps in (4, 8, 16): + temperature.array[:, 0, 0] = initial[ids] + thermal.DuDt.temperature_rate.array[...] = 0.0 + thermal.DuDt.state = initial_state + dt = 0.1 / steps + for _ in range(steps): + thermal.solve(timestep=dt) + actual = np.zeros(len(initial)) + for local_ids, values in uw.mpi.comm.allgather( + (ids, np.array(temperature.array[:, 0, 0]))): + actual[local_ids] = values + factor = (1 - 0.5 * dt * eigenvalues[1]) / (1 + 0.5 * dt * eigenvalues[1]) + expected = initial * factor**steps + map_error = _norm(actual - expected, mass) / _norm(expected, mass) + errors.append(_norm(actual - exact, mass) / _norm(exact, mass)) + assert map_error < 1.0e-9, map_error + assert thermal.DuDt.last_corrector_iterations <= thermal.DuDt.max_corrector_steps + assert thermal.DuDt.last_corrector_residual <= thermal.DuDt.corrector_target + uw.pprint( + f"PC_CONVERGED_DIFFUSION dim={pc2.mesh.dim} steps={steps} " + f"dt={dt:.12g} relative_error={errors[-1]:.12g} " + f"map_error={map_error:.12g} corrections={thermal.DuDt.last_corrector_iterations} " + f"residual={thermal.DuDt.last_corrector_residual:.12g}") + rates = _orders(errors) + assert np.all((1.9 < rates) & (rates < 2.2)), rates + uw.pprint(f"PC_CONVERGED_TIME_ORDER dim={pc2.mesh.dim} rates={rates.tolist()}") diff --git a/tests/test_1119_supg_process_restart.py b/tests/test_1119_supg_process_restart.py new file mode 100644 index 00000000..d0ca1c9b --- /dev/null +++ b/tests/test_1119_supg_process_restart.py @@ -0,0 +1,73 @@ +"""Fresh-process transport snapshots on tiny tetrahedra. + +Run this parent pytest in serial. UW_SUPG_TEST_RANKS=8 requests eight-rank +workers; the default uses singleton workers. Every phase starts a fresh +interpreter. No forked in-memory snapshot can satisfy the restore check. +""" + +import os +from pathlib import Path +import shutil +import subprocess +import sys + +import h5py +import numpy as np +import pytest + +import underworld3 as uw +from parallel.serial_reference import _MPI_ENV_PREFIXES + +pytestmark = [pytest.mark.level_2, pytest.mark.tier_b] + + +@pytest.mark.parametrize("method", ["pc2", "pc_converged", "cn", "bdf2"]) +def test_fresh_process_transport_restart(method, tmp_path): + if uw.mpi.size != 1: + pytest.skip("Run the parent in serial; UW_SUPG_TEST_RANKS selects worker ranks.") + ranks = int(os.environ.get("UW_SUPG_TEST_RANKS", "1")) + root = Path(__file__).resolve().parents[1] + worker = root / "tests/parallel/ptest_1119_supg_restart.py" + supervisor = root / "scripts/mpi_supervisor.py" + env = {key: value for key, value in os.environ.items() + if not key.startswith(_MPI_ENV_PREFIXES)} + launcher = [] + if ranks > 1: + executable = Path(sys.executable).with_name("mpirun") + if not executable.is_file(): + executable = shutil.which("mpirun") + assert executable, "Activate the matching MPI environment before this test." + launcher = [str(executable), "-np", str(ranks)] + if ranks > 1 and not supervisor.is_file(): + pytest.skip("MPI restart supervision requires development commit #678.") + for phase in ("full", "write", "resume"): + worker_command = [*launcher, sys.executable, "-m", "mpi4py", str(worker), + "-uw_method", method, "-uw_phase", phase] + command = ([sys.executable, str(supervisor), "--silence", "45", + "--hard-cap", "60", "--", *worker_command] + if supervisor.is_file() else worker_command) + with (tmp_path / f"{phase}.log").open("w") as log: + status = subprocess.run(command, cwd=tmp_path, env=env, + stdout=log, stderr=subprocess.STDOUT, + timeout=65).returncode + assert status == 0, (tmp_path / f"{phase}.log").read_text(errors="replace") + maxima = {"field": 0.0, "estimate": 0.0} + for rank in range(ranks): + with h5py.File(tmp_path / f"full_rank{rank}.h5") as full, h5py.File( + tmp_path / f"resume_rank{rank}.h5") as resumed: + assert set(full) == set(resumed) + for name in full: + expected, actual = full[name][()], resumed[name][()] + if name == "estimate_dt" or name == "solver_last_change_rate": + if isinstance(expected, bytes): + assert actual == expected + continue + np.testing.assert_allclose(actual, expected, rtol=1e-9, atol=1e-10, err_msg=name) + maxima["estimate"] = max(maxima["estimate"], float(np.max(np.abs(actual-expected)))) + elif name in ("coords", "step", "time") or name.startswith(("solver_", "history_")): + np.testing.assert_array_equal(actual, expected, err_msg=name) + else: + np.testing.assert_allclose(actual, expected, rtol=1e-11, atol=1e-12, err_msg=name) + maxima["field"] = max(maxima["field"], float(np.max(np.abs(actual-expected)))) + print(f"SUPG_FRESH_RESTART method={method} ranks={ranks} " + f"max_field_error={maxima['field']:.12g} max_estimator_error={maxima['estimate']:.12g}", flush=True) diff --git a/tests/test_1120_supg_memory.py b/tests/test_1120_supg_memory.py new file mode 100644 index 00000000..f658cbe6 --- /dev/null +++ b/tests/test_1120_supg_memory.py @@ -0,0 +1,132 @@ +"""Fast workspace reuse and opt-in transport memory soak tests. + +No Stokes, reaction diagnostics, checkpoints, or forced garbage collection +occur in either loop. The default test checks stable object identities over +eight updates. Set UW_RUN_SUPG_MEMORY_SOAK=1 to run the 200-update RSS test. +""" + +import os +import time + +import numpy as np +import pytest + +import underworld3 as uw +from underworld3.utilities import memprobe + +pytestmark = pytest.mark.tier_b + + +def _workspace(thermal): + """Record identities, not contents that should change during transport.""" + identity = [thermal.snes.handle, thermal.dm.handle, + tuple((name, field.vec.handle) for name, field in thermal.mesh.vars.items())] + if isinstance(thermal.DuDt, uw.systems.ddt.EulerianSUPGPC): + identity.extend([ + thermal.DuDt._lumped_mass.handle, + tuple(vector.handle for vector in thermal.DuDt._citcoms_work_vectors), + tuple(id(array) for array in thermal.DuDt._simplex_data_cache), + tuple(id(array) for array in thermal.DuDt._directional_rate_work), + ]) + return identity + + +def _transport_problem(dim, method): + mesh = uw.meshing.UnstructuredSimplexBox( + minCoords=(0.0,) * dim, maxCoords=(1.0,) * dim, + cellSize=0.25, qdegree=4, regular=False, + ) + temperature = uw.discretisation.MeshVariable("T", mesh, 1, degree=1) + velocity = uw.discretisation.MeshVariable("U", mesh, dim, degree=1) + temperature.array[:, 0, 0] = np.prod( + np.sin(np.pi * np.asarray(temperature.coords)), axis=1) + velocity.array[...] = 0.0 + velocity.array[:, 0, 0] = 0.2 + if method in ("pc2", "pc_converged"): + manager = uw.systems.ddt.EulerianSUPGPC( + mesh, temperature, velocity.sym, + method="citcoms" if method == "pc2" else method, + ) + thermal = uw.systems.AdvDiffusion(mesh, temperature, velocity.sym, DuDt=manager) + else: + thermal = uw.systems.AdvDiffusion( + mesh, temperature, velocity.sym, + order=1 if method == "cn" else 2, + theta=0.5 if method == "cn" else 1.0, + peclet_weight=0.0, + ) + thermal.constitutive_model.Parameters.diffusivity = 0.01 + for boundary in mesh.boundaries: + if boundary.name not in ("All_Boundaries", "Null_Boundary"): + thermal.add_dirichlet_bc(0.0, boundary.name) + return thermal, temperature, velocity + + +def _advance(thermal, velocity, step): + velocity.array[:, 0, 0] = 0.2 * (1.0 + 0.1 * np.sin(step)) + dt = (0.001, 0.0015, 0.002, 0.0025)[step % 4] + thermal.estimate_dt() + thermal.solve(timestep=dt) + + +@pytest.mark.level_2 +@pytest.mark.parametrize("dim", [2, 3]) +@pytest.mark.parametrize("method", ["pc2", "pc_converged", "cn", "bdf2"]) +def test_transport_workspace_reuse(dim, method): + thermal, temperature, velocity = _transport_problem(dim, method) + _advance(thermal, velocity, 1) + workspace = _workspace(thermal) + for step in range(2, 9): + _advance(thermal, velocity, step) + unchanged = _workspace(thermal) == workspace + assert all(uw.mpi.comm.allgather(unchanged)), "Solver workspace was reallocated" + nodal = np.asarray(temperature.array) + assert all(uw.mpi.comm.allgather(bool(np.isfinite(nodal).all()))) + minimum = min(uw.mpi.comm.allgather(float(np.min(nodal)))) + maximum = max(uw.mpi.comm.allgather(float(np.max(nodal)))) + assert -0.05 < minimum and maximum < 1.05 + + +@pytest.mark.level_3 +@pytest.mark.slow +@pytest.mark.skipif( + os.environ.get("UW_RUN_SUPG_MEMORY_SOAK") != "1", + reason="Set UW_RUN_SUPG_MEMORY_SOAK=1 to run the 200-update RSS regression.", +) +@pytest.mark.parametrize("dim", [2, 3]) +@pytest.mark.parametrize("method", ["pc2", "cn", "bdf2"]) +def test_repeated_transport_memory_and_workspace_reuse(dim, method): + pytest.importorskip("psutil", reason="This test requires current RSS, not peak RSS.") + thermal, temperature, velocity = _transport_problem(dim, method) + samples = [] + start = time.perf_counter() + for step in range(1, 201): + _advance(thermal, velocity, step) + if step == 40: + workspace = _workspace(thermal) + if step >= 40 and step % 10 == 0: + unchanged = _workspace(thermal) == workspace + assert all(uw.mpi.comm.allgather(unchanged)), "Solver workspace was reallocated" + rss = uw.mpi.comm.allgather(memprobe.snapshot()["rss_mb"]) + samples.append((step, *rss)) + elapsed = max(uw.mpi.comm.allgather(time.perf_counter() - start)) + samples = np.asarray(samples) + late = samples[samples[:, 0] >= 120] + slopes = np.polyfit(late[:, 0], late[:, 1:], 1)[0] + growth = samples[-1, 1:] - samples[0, 1:] + nodal = np.asarray(temperature.array) + assert all(uw.mpi.comm.allgather(bool(np.isfinite(nodal).all()))) + minimum = min(uw.mpi.comm.allgather(float(np.min(nodal)))) + maximum = max(uw.mpi.comm.allgather(float(np.max(nodal)))) + uw.pprint( + f"SUPG_MEMORY method={method} dim={dim} ranks={uw.mpi.size} seconds={elapsed:.6f} " + f"rss_start_mib={samples[0, 1:].sum():.6f} rss_end_mib={samples[-1, 1:].sum():.6f} " + f"growth_mib={growth.sum():.6f} late_slope_mib_per_step={slopes.sum():.9f} " + f"max_rank_growth_mib={growth.max():.6f} max_rank_slope={slopes.max():.9f} " + f"Tmin={minimum:.9g} Tmax={maximum:.9g}") + uw.pprint(f"SUPG_MEMORY_SAMPLES method={method} dim={dim} values={samples.tolist()}") + assert np.isfinite(samples).all() + assert -0.05 < minimum and maximum < 1.05 + # Fixed pre-run bounds allow allocator noise but reject sustained growth. + assert growth.max() < 16.0, growth + assert slopes.max() < 0.05, slopes diff --git a/tests/test_1121_supg_pc_manager.py b/tests/test_1121_supg_pc_manager.py new file mode 100644 index 00000000..9010f080 --- /dev/null +++ b/tests/test_1121_supg_pc_manager.py @@ -0,0 +1,98 @@ +"""Predictor-corrector composition with the scalar transport solver.""" + +import numpy as np +import pytest +import sympy + +import underworld3 as uw + +pytestmark = [pytest.mark.level_1, pytest.mark.tier_b] + + +@pytest.fixture +def fields(): + mesh = uw.meshing.UnstructuredSimplexBox( + minCoords=(0.0, 0.0), maxCoords=(1.0, 1.0), cellSize=0.25, + qdegree=3, + ) + temperature = uw.discretisation.MeshVariable("T", mesh, 1, degree=1) + return mesh, temperature, sympy.zeros(1, mesh.dim) + + +def test_manager_rejects_a_different_solver_unknown(fields): + mesh, temperature, velocity = fields + other = uw.discretisation.MeshVariable("Other", mesh, 1, degree=1) + transport = uw.systems.ddt.EulerianSUPGPC(mesh, temperature, velocity) + with pytest.raises(ValueError, match="unknown|u_Field|field"): + uw.systems.AdvDiffusion(mesh, other, velocity, DuDt=transport) + + +@pytest.mark.parametrize("method", ["citcoms", "pc_converged"]) +def test_manager_uses_live_boundary_conditions_and_timestep(fields, method): + mesh, temperature, velocity = fields + temperature.array[:, 0, 0] = 1.0 + transport = uw.systems.ddt.EulerianSUPGPC( + mesh, temperature, velocity, method=method, + ) + thermal = uw.systems.AdvDiffusion(mesh, temperature, velocity, DuDt=transport) + thermal.constitutive_model.Parameters.diffusivity = 0.1 + # Conditions are added after the manager is bound to the solver. + for boundary in ("Left", "Right", "Top", "Bottom"): + thermal.add_dirichlet_bc(1.0, boundary) + thermal.solve(timestep=0.001) + thermal.solve() + assert thermal.DuDt is transport + assert float(transport.delta_t.sym) == 0.001 + np.testing.assert_allclose(temperature.array, 1.0, rtol=0, atol=1e-12) + + +def test_rejected_theta_does_not_change_snapshot_metadata(fields): + mesh, temperature, velocity = fields + transport = uw.systems.ddt.EulerianSUPGPC(mesh, temperature, velocity) + thermal = uw.systems.AdvDiffusion(mesh, temperature, velocity, DuDt=transport) + previous = thermal.state + with pytest.raises(ValueError, match="theta"): + thermal.theta = 0.5 + assert thermal.state == previous + + +def test_manager_cannot_execute_against_another_solver(fields): + mesh, temperature, velocity = fields + transport = uw.systems.ddt.EulerianSUPGPC(mesh, temperature, velocity) + first = uw.systems.AdvDiffusion(mesh, temperature, velocity, DuDt=transport) + other = uw.discretisation.MeshVariable("Other", mesh, 1, degree=1) + second = uw.systems.AdvDiffusion(mesh, other, velocity) + with pytest.raises(ValueError, match="field|solver|unknown"): + second.DuDt = first.DuDt + second.solve(timestep=0.001) + + +def test_replacing_velocity_expression_matches_updating_velocity_field(fields): + mesh, temperature, velocity = fields + reference = uw.discretisation.MeshVariable("Reference", mesh, 1, degree=1) + vector = uw.discretisation.MeshVariable("Velocity", mesh, mesh.dim, degree=1) + temperature.array[:, 0, 0] = temperature.coords[:, 0] + reference.array[...] = temperature.array + vector.array[...] = 0.0 + actual_manager = uw.systems.ddt.EulerianSUPGPC(mesh, temperature, velocity) + reference_manager = uw.systems.ddt.EulerianSUPGPC(mesh, reference, vector.sym) + actual = uw.systems.AdvDiffusion(mesh, temperature, velocity, DuDt=actual_manager) + expected = uw.systems.AdvDiffusion(mesh, reference, vector.sym, DuDt=reference_manager) + actual.solve(timestep=0.001) + expected.solve(timestep=0.001) + actual_manager.V_fn = sympy.Matrix([[0.2, 0.0]]) + vector.array[:, 0, 0] = 0.2 + actual.solve(timestep=0.001) + expected.solve(timestep=0.001) + np.testing.assert_allclose(temperature.array, reference.array, rtol=0, atol=1e-12) + + +def test_explicit_tau_accepts_directional_diffusion(fields): + mesh, temperature, velocity = fields + temperature.array[:, 0, 0] = 1.0 + transport = uw.systems.ddt.EulerianSUPGPC(mesh, temperature, velocity, tau=0) + thermal = uw.systems.AdvDiffusion(mesh, temperature, velocity, DuDt=transport) + thermal.constitutive_model = uw.constitutive_models.AnisotropicDiffusionModel + thermal.constitutive_model.Parameters.diffusivity = sympy.Matrix([0.1, 0.2]) + thermal.solve(timestep=0.001) + np.testing.assert_allclose(temperature.array, 1.0, rtol=0, atol=1e-12)