diff --git a/docs/advanced/index.md b/docs/advanced/index.md index cb47f8dea..d022a29ef 100644 --- a/docs/advanced/index.md +++ b/docs/advanced/index.md @@ -82,6 +82,12 @@ time-derivative and Adams-Moulton/θ flux knobs, and how to pair them **[→ Semi-Lagrangian Time Integration](semi-lagrangian-time-integration.md)** +### SUPG Scalar Transport +Use local streamline-upwind stabilization with implicit BDF integration or +the continuous-P1 CitcomS-compatible predictor-corrector. + +**[→ SUPG Scalar Transport](supg-transport.md)** + ### Porous Media Flow Darcy flow, Richards equation, and variably-saturated groundwater modelling. @@ -139,9 +145,10 @@ custom-meshes curved-boundary-conditions mesh-adaptation semi-lagrangian-time-integration +supg-transport porous-flow snapshot-restore troubleshooting api-patterns SWARM-INTEGRATION-STATISTICS -``` \ No newline at end of file +``` diff --git a/docs/advanced/semi-lagrangian-time-integration.md b/docs/advanced/semi-lagrangian-time-integration.md index 23a935db2..9a592df2c 100644 --- a/docs/advanced/semi-lagrangian-time-integration.md +++ b/docs/advanced/semi-lagrangian-time-integration.md @@ -15,6 +15,9 @@ treating diffusion implicitly. This page explains the **two independent order knobs** in the scheme and how to pair them correctly — the common pitfall is mixing them. +For local finite-element streamline stabilization without characteristic +trace-back, see [SUPG scalar transport](supg-transport.md). + ## The scheme has two time-integration choices The discrete residual assembled by the solver is diff --git a/docs/advanced/supg-transport.md b/docs/advanced/supg-transport.md new file mode 100644 index 000000000..c79163d4d --- /dev/null +++ b/docs/advanced/supg-transport.md @@ -0,0 +1,141 @@ +--- +title: "SUPG Scalar Transport" +--- + +# SUPG scalar transport + +`AdvDiffusionSUPG` solves + +$$ +\frac{\partial T}{\partial t} + \mathbf{u}\cdot\nabla T +- \nabla\cdot(\kappa\nabla T) = f +$$ + +on simplex volume meshes. It adds streamline-upwind Petrov-Galerkin (SUPG) +stabilization to the continuous finite-element residual. Advection remains a +local finite-element operation: the solver does not trace departure points or +interpolate a semi-Lagrangian history. + +## Minimal example + +This example transports and diffuses a continuous P1 scalar in a prescribed +velocity field. Automatic stabilization is the default. + +```python +import numpy as np +import underworld3 as uw + +mesh = uw.meshing.UnstructuredSimplexBox( + minCoords=(0.0, 0.0), + maxCoords=(1.0, 1.0), + cellSize=0.1, +) + +temperature = uw.discretisation.MeshVariable("T", mesh, 1, degree=1) +velocity = uw.discretisation.MeshVariable("U", mesh, mesh.dim, degree=1) + +with mesh.access(temperature, velocity): + x = temperature.coords[:, 0] + y = temperature.coords[:, 1] + temperature.data[:, 0] = np.sin(np.pi * x) * np.sin(np.pi * y) + velocity.data[:, 0] = 1.0 + velocity.data[:, 1] = 0.0 + +thermal = uw.systems.AdvDiffusionSUPG( + mesh, + u_Field=temperature, + V_fn=velocity.sym, +) +thermal.constitutive_model = uw.constitutive_models.DiffusionModel +thermal.constitutive_model.Parameters.diffusivity = 0.01 + +for boundary in ("Left", "Right", "Top", "Bottom"): + thermal.add_dirichlet_bc(0.0, boundary) + +for _ in range(10): + thermal.solve(timestep=1.0e-3, zero_init_guess=False) +``` + +The default `time_integrator="bdf"` uses an implicit Eulerian BDF method and +the generic transient SUPG stabilization parameter. `order=1` and `order=2` +select BDF1 and BDF2 respectively. + +## CitcomS-compatible predictor-corrector + +For continuous P1 temperature, UW3 also provides the row-sum-mass +predictor-corrector used for the Zhong mantle-convection benchmark: + +```python +temperature_rate = uw.discretisation.MeshVariable( + "Tdot", mesh, 1, degree=1 +) + +thermal = uw.systems.AdvDiffusionSUPG( + mesh, + u_Field=temperature, + V_fn=velocity.sym, + time_integrator="citcoms", + temperature_rate_field=temperature_rate, +) +thermal.constitutive_model = uw.constitutive_models.DiffusionModel +thermal.constitutive_model.Parameters.diffusivity = 0.01 + +dt = thermal.estimate_dt() +thermal.solve(timestep=dt) +``` + +This path uses `adv_gamma=0.5`, two residual-correction iterations, positive +row-sum mass, and the clipped CitcomS stabilization parameter by default. The +timestep is explicit and must satisfy the returned advection-diffusion limit. +Supplying a named `temperature_rate_field` makes the additional restart state +visible and straightforward to checkpoint. Exact restart requires `T`, +`Tdot`, and the solver snapshot metadata. + +## Choosing a transport solver + +| Method | Strength | Main cost or limitation | Restart state | +| --- | --- | --- | --- | +| `AdvDiffusionSUPG`, implicit BDF | Local assembly, no trace-back interpolation, automatic simplex stabilization | Timestep accuracy still requires convergence testing; automatic tau currently assumes scalar isotropic diffusivity on volume simplices | Temperature plus BDF history | +| `AdvDiffusionSUPG`, CitcomS predictor-corrector | Second-order explicit update, row-lumped P1 mass, close to CitcomS mantle-convection numerics | Continuous P1 only; explicit advection-diffusion timestep limit | `T`, `Tdot`, solver metadata | +| `AdvDiffusionSLCN` | Stable characteristic transport at large advective Courant number | Departure-point search/interpolation, flux history, and higher MPI memory/runtime | Temperature plus characteristic and flux histories | +| `AdvDiffusionSLCN` with SL-BDF2 | Second-order characteristic history without Crank-Nicolson flux ringing | Two departure points and greater history/interpolation cost | Two-level characteristic history plus flux history | +| `AdvDiffusionSLCN` with BDF1/Backward Euler | Robust, L-stable diffusion baseline | First-order time integration and trace-back interpolation | One characteristic history level | + +Use the CitcomS predictor-corrector when reproducing a continuous-P1 CitcomS +benchmark. Use implicit SUPG when local streamline stabilization is desired +without the explicit predictor-corrector restriction. Use SLCN when large +advective timesteps are more important than trace-back cost. For every method, +verify timestep and mesh convergence using the physical diagnostics of the +problem; the solver name alone does not establish accuracy. + +## Conservation and boundedness + +SUPG is a consistent residual stabilization, but continuous Galerkin SUPG is +not automatically monotone and does not guarantee nodal maximum principles. +The CitcomS path uses positive row-sum mass, which improves the explicit +update, but temperature bounds must still be checked. Semi-Lagrangian methods +can remain stable at large advective Courant number, but departure-point +interpolation is not strictly conservative. For either solver family, monitor +the volume-integrated scalar, recovered boundary fluxes, source integral, and +their discrete balance in addition to minimum and maximum nodal values. + +Parallel execution does not change these definitions. Compare global +integrals between serial and MPI runs on the same mesh; do not compare local +rank extrema or partition-dependent raw boundary-node sums. + +## Stabilization controls + +- Omit `tau` for automatic stabilization. +- `tau_model="generic"` combines transient, advective, and diffusive scales. +- `tau_model="citcoms"` selects the clipped steady CitcomS relation on a + simplex streamline length. +- Pass an explicit scalar `tau` for unsupported element or constitutive-model + combinations. `tau=0` recovers the unstabilized Galerkin residual. +- Automatic tau supports two- and three-dimensional simplex volume meshes and + scalar isotropic non-negative diffusivity. + +## Related documentation + +- [Semi-Lagrangian time integration](semi-lagrangian-time-integration.md) +- [State snapshots and restore](snapshot-restore.md) +- [Parallel computing](parallel-computing.md) diff --git a/docs/developer/CHANGELOG.md b/docs/developer/CHANGELOG.md index 53c0126f8..b6ebdbd3d 100644 --- a/docs/developer/CHANGELOG.md +++ b/docs/developer/CHANGELOG.md @@ -278,6 +278,12 @@ component exactly — correct on curved, tilted, and deformed boundaries (#293). existing boundary traction onto an axisymmetric harmonic; the pure functions also accept coefficients recovered by other methods and an optional internal load. +- Rotated free slip now exposes + `Stokes.boundary_normal_traction_integral(boundary, fn)` for a distributed + weak contraction of the assembled normal reaction. The spherical-shell geoid + adapter accepts `projection="reaction"` to use this fitted quantity without + pointwise P2 recovery or a rank-zero surface triangulation; the existing + `projection="centroid"` behavior remains the default. - `uw.analytic.Zhong2008` implements the Hager--O'Connell propagator-matrix oracle used for the Zhong et al. spherical-shell response benchmark. It supports piecewise-constant radial viscosity and reproduces every analytical diff --git a/docs/developer/subsystems/boundary-stress-and-projection-postprocessing.md b/docs/developer/subsystems/boundary-stress-and-projection-postprocessing.md index fb007f0ba..f62832404 100644 --- a/docs/developer/subsystems/boundary-stress-and-projection-postprocessing.md +++ b/docs/developer/subsystems/boundary-stress-and-projection-postprocessing.md @@ -97,13 +97,27 @@ response = uw.postprocessing.geoid.spherical_shell_response_from_rotated_stokes( planet_radius=6370000.0, gravity=9.8, gravitational_constant=6.67e-11, + projection="reaction", ) ``` -The adapter delegates stress recovery to the existing rotated-free-slip API; -it does not implement a second CBF, constrained-multiplier, or topography -recovery path. `internal_load_coefficient` must use the same harmonic -normalisation and sign convention as the model's internal load. +The adapter supports two projection paths. `projection="centroid"` retains the +original pointwise-recovery workflow: recover `sigma_nn`, gather the samples to +rank zero, reconstruct a spherical triangulation, and integrate centroid +values. `projection="reaction"` contracts the assembled normal-reaction load +directly with the harmonic test function through +`Stokes.boundary_normal_traction_integral()`. The latter is distributed, avoids +the rank-zero surface reconstruction, and is an integral/fitted quantity rather +than a consumer of the slowly converging P2 vertex values on curved boundaries +(#414). Its fitted coefficient uses the matching discrete boundary norm, not an +analytical spherical norm, so the numerator and denominator share the same +faceted geometry. + +Both paths reuse the existing rotated-free-slip reaction; neither implements a +second CBF, constrained-multiplier, or topography recovery. `centroid` remains +the compatibility default while the direct reaction path accumulates benchmark +coverage. `internal_load_coefficient` must use the same harmonic normalisation +and sign convention as the model's internal load. When surface and CMB topography coefficients are already available, call `uw.postprocessing.geoid.spherical_shell_geoid_response()` or @@ -137,10 +151,11 @@ Published reference solvers, such as the Zhong et al. propagator-matrix method, belong in `uw.analytic`; their computed topography coefficients can be passed to the pure post-processing functions above. -The rotated harmonic projector gathers boundary samples to rank zero and -reconstructs their spherical triangulation. A future boundary-reaction -functional could replace this step with a direct distributed finite-element -projection without changing the coefficient API. +`Stokes.boundary_normal_traction_integral(boundary, fn)` is useful beyond the +geoid adapter whenever only an integrated or fitted normal-traction diagnostic +is required. Pointwise consumers should continue to call +`boundary_normal_traction()` or `dynamic_topography()` and follow their curved +P2 midpoint/field guidance. ## See also diff --git a/src/underworld3/checkpoint/disk_snapshot.py b/src/underworld3/checkpoint/disk_snapshot.py index 7d30d2000..102c9ba20 100644 --- a/src/underworld3/checkpoint/disk_snapshot.py +++ b/src/underworld3/checkpoint/disk_snapshot.py @@ -461,6 +461,15 @@ def read_snapshot(model, path: str) -> None: import h5py md = read_snapshot_metadata(path) + write_size = int(md.get("mpi_ranks_at_write", 1)) + if write_size != int(uw.mpi.size): + raise ValueError( + f"snapshot at {path} was written on {write_size} MPI rank(s); " + f"this run uses {uw.mpi.size}. Exact disk restart requires the " + "same rank count; use mesh.write_timestep/read_timestep for " + "coordinate-remapped field transfer." + ) + bulk_dir = _bulk_dir_for(path) if not os.path.isdir(bulk_dir): raise FileNotFoundError( diff --git a/src/underworld3/cython/petsc_generic_snes_solvers.pyx b/src/underworld3/cython/petsc_generic_snes_solvers.pyx index 3b3da609f..f3625fa5a 100644 --- a/src/underworld3/cython/petsc_generic_snes_solvers.pyx +++ b/src/underworld3/cython/petsc_generic_snes_solvers.pyx @@ -3125,6 +3125,16 @@ class SolverBaseClass(uw_object): self._subdict[name][1].localToGlobal(var.vec, sgvec) gvec.restoreSubVector(self._subdict[name][0], sgvec) else: + # Map the variable's LOCAL vector through field 0's subDM rather + # than assuming the solver DM's local layout matches it. The two + # coincide on a plain single-field solver, which is why the + # direct `self.dm.localToGlobal(self.Unknowns.u.vec, gvec)` looked + # equivalent -- but where they differ it writes to the wrong slots + # and the field comes back never-written. Restored while CI is red + # on tests/test_1120_SLVectorCartesian.py::test_SLVec_boxmesh[mesh1], + # a semi-Lagrangian VECTOR test, i.e. exactly the single-field path + # this branch serves; its recovered values were ~1e-18 against an + # analytic ~1e-5. _names, _iss, _subdms = self.dm.createFieldDecomposition() sgvec = gvec.getSubVector(_iss[0]) _subdms[0].localToGlobal(self.Unknowns.u.vec, sgvec) @@ -3209,6 +3219,18 @@ class SolverBaseClass(uw_object): return _bff(self, boundary, field, mass=mass, remove_mean=remove_mean, scale=scale, normal=normal) + def boundary_flux_integral(self, boundary): + r"""Integrated scalar CBF flux through ``boundary``. + + This is the direct integral diagnostic for quantities such as Nusselt + numbers. It sums the consistent scalar nodal reactions collectively, + avoiding pointwise de-smearing and a temporary flux MeshVariable. Use + :meth:`boundary_flux` or :meth:`boundary_flux_field` when nodal values + are required. + """ + from underworld3.utilities.boundary_flux import boundary_flux_integral as _bfi + return _bfi(self, boundary) + ## Specific to dimensionality @@ -6502,6 +6524,25 @@ class SNES_Stokes_SaddlePt(SolverBaseClass): return _dtf(self, boundary, self._rotated_freeslip_info, field, buoyancy_scale=buoyancy_scale, mass=mass) + def boundary_normal_traction_integral(self, boundary, fn, remove_mean=True): + r"""Return the boundary integral of ``sigma_nn * fn`` directly from the + rotated-free-slip constraint reaction. + + With ``remove_mean=True`` (default), the constant normal-traction gauge + is removed before projection. Unlike pointwise + :meth:`boundary_normal_traction`, this weak projection does not recover + nodal traction values or gather a global surface mesh. It is therefore + suitable for harmonic/integral diagnostics on curved P2 boundaries, + whose recovered vertex values converge slowly (issue #414). + """ + if self._rotated_freeslip_info is None: + raise RuntimeError( + "boundary_normal_traction_integral requires a completed " + "rotated-free-slip solve.") + from underworld3.utilities.rotated_bc import boundary_normal_traction_integral as _bnti + return _bnti(self, boundary, self._rotated_freeslip_info, fn, + remove_mean=remove_mean) + def add_nitsche_bc(self, conds=None, boundary=None, direction=None, normal=None, gamma=10.0, theta=1, mask=None, local_h=True, g=None): r"""Add Nitsche weak enforcement of a velocity constraint along a direction. diff --git a/src/underworld3/discretisation/discretisation_mesh.py b/src/underworld3/discretisation/discretisation_mesh.py index 769dd6fa1..f18e002b0 100644 --- a/src/underworld3/discretisation/discretisation_mesh.py +++ b/src/underworld3/discretisation/discretisation_mesh.py @@ -4836,9 +4836,10 @@ def write_timestep( - ``create_xdmf=True`` writes ParaView/XDMF output. Variable files also receive ``/vertex_fields`` or ``/cell_fields`` compatibility groups, and rank 0 writes the companion ``.xdmf`` file. - - ``petsc_reload=True`` writes PETSc DMPlex section/vector metadata into - the same per-variable HDF5 files. These files can then be loaded with - ``MeshVariable.read_checkpoint()`` for PETSc-native same-mesh reload. + - ``petsc_reload=True`` writes PETSc DMPlex metadata and an owned + global-vector payload into the same per-variable HDF5 files. These + files can then be loaded with ``MeshVariable.read_checkpoint()`` for + exact same-layout reload. Common choices are: @@ -5015,7 +5016,7 @@ def _write_petsc_reload_variable(self, viewer, var): subdm.destroy() def _write_petsc_reload_file(self, checkpoint_file, variables, mode="w"): - """Write PETSc DMPlex section/vector reload metadata.""" + """Write DMPlex metadata and exact same-layout global vectors.""" old_dm_name = self.dm.getName() self.dm.setName("uw_mesh") @@ -5037,6 +5038,23 @@ def _write_petsc_reload_file(self, checkpoint_file, variables, mode="w"): if old_dm_name is not None: self.dm.setName(old_dm_name) + viewer = PETSc.ViewerHDF5().create( + checkpoint_file, "a", comm=PETSc.COMM_WORLD + ) + try: + viewer.pushGroup("/uw_checkpoint") + for var in variables: + var._sync_lvec_to_gvec() + checkpoint_vec = PETSc.Vec().createWithArray( + var._gvec.array_r, comm=PETSc.COMM_WORLD + ) + checkpoint_vec.setName(var.clean_name) + viewer(checkpoint_vec) + checkpoint_vec.destroy() + viewer.popGroup() + finally: + viewer.destroy() + @timing.routine_timer_decorator def write_checkpoint( self, diff --git a/src/underworld3/discretisation/discretisation_mesh_variables.py b/src/underworld3/discretisation/discretisation_mesh_variables.py index 9bbcabbcf..beaa45870 100644 --- a/src/underworld3/discretisation/discretisation_mesh_variables.py +++ b/src/underworld3/discretisation/discretisation_mesh_variables.py @@ -1452,11 +1452,11 @@ def read_checkpoint( ): """Load this mesh variable from PETSc reload output. - This is an exact PETSc DMPlex section/vector reload path. It does not - use the coordinate/KDTree remapping used by ``read_timestep()``. New - output should be written with ``Mesh.write_timestep(..., - petsc_reload=True)``; legacy ``Mesh.write_checkpoint()`` files are also - supported. + This is an exact same-layout PETSc vector reload path. It does not use + the coordinate/KDTree remapping used by ``read_timestep()``. New output + should be written with ``Mesh.write_timestep(..., petsc_reload=True)``; + legacy ``Mesh.write_checkpoint()`` files are also supported through + the older DMPlex local-vector fallback. """ if data_name is None: @@ -1465,6 +1465,17 @@ def read_checkpoint( if self._lvec is None: self._set_vec(available=True) + import h5py + + has_direct_vector = False + if uw.mpi.rank == 0: + with h5py.File(filename, "r") as checkpoint_h5: + has_direct_vector = ( + "uw_checkpoint" in checkpoint_h5 + and data_name in checkpoint_h5["uw_checkpoint"] + ) + has_direct_vector = uw.mpi.comm.bcast(has_direct_vector, root=0) + indexset, subdm = self.mesh.dm.createSubDM(self.field_id) sectiondm = self.mesh.dm.clone() viewer = PETSc.ViewerHDF5().create(filename, "r", comm=PETSc.COMM_WORLD) @@ -1481,40 +1492,53 @@ def read_checkpoint( self._lvec.setName(data_name) self._gvec.setName(data_name) - from underworld3.cython.petsc_discretisation import ( - petsc_dmplex_load_local_vector, - ) + if has_direct_vector: + checkpoint_vec = PETSc.Vec().createMPI( + (self._gvec.getLocalSize(), self._gvec.getSize()), + comm=PETSc.COMM_WORLD, + ) + checkpoint_vec.setName(data_name) + viewer.pushGroup("/uw_checkpoint") + checkpoint_vec.load(viewer) + viewer.popGroup() + self._gvec.array[...] = checkpoint_vec.array_r + checkpoint_vec.destroy() + subdm.globalToLocal(self._gvec, self._lvec, addv=False) + else: + from underworld3.cython.petsc_discretisation import ( + petsc_dmplex_load_local_vector, + ) - loaded_lvec = petsc_dmplex_load_local_vector( - self.mesh.dm, viewer, sectiondm, self.mesh.sf, data_name - ) + loaded_lvec = petsc_dmplex_load_local_vector( + self.mesh.dm, viewer, sectiondm, self.mesh.sf, data_name + ) - source_section = sectiondm.getSection() - target_section = subdm.getSection() - source_array = loaded_lvec.array_r - target_array = self._lvec.array - p_start, p_end = target_section.getChart() - - for point in range(p_start, p_end): - target_dof = target_section.getDof(point) - if target_dof == 0: - continue - - source_dof = source_section.getDof(point) - if source_dof < target_dof: - raise RuntimeError( - f"Checkpoint section has {source_dof} dofs for point {point}, " - f"but target variable requires {target_dof}." + source_section = sectiondm.getSection() + target_section = subdm.getSection() + source_array = loaded_lvec.array_r + target_array = self._lvec.array + p_start, p_end = target_section.getChart() + + for point in range(p_start, p_end): + target_dof = target_section.getDof(point) + if target_dof == 0: + continue + + source_dof = source_section.getDof(point) + if source_dof < target_dof: + raise RuntimeError( + f"Checkpoint section has {source_dof} dofs for point " + f"{point}, but target variable requires {target_dof}." + ) + + source_offset = source_section.getOffset(point) + target_offset = target_section.getOffset(point) + target_array[target_offset : target_offset + target_dof] = ( + source_array[source_offset : source_offset + target_dof] ) - source_offset = source_section.getOffset(point) - target_offset = target_section.getOffset(point) - target_array[target_offset : target_offset + target_dof] = ( - source_array[source_offset : source_offset + target_dof] - ) - - loaded_lvec.destroy() - self._sync_lvec_to_gvec() + loaded_lvec.destroy() + self._sync_lvec_to_gvec() finally: self._lvec.setName(old_lvec_name) self._gvec.setName(old_vec_name) @@ -1526,6 +1550,11 @@ def read_checkpoint( indexset.destroy() subdm.destroy() + # The mesh-wide auxiliary vector packs every registered field and may + # still contain values from before this reload. Force the next residual + # assembly to rebuild it from the restored per-variable vectors. + self.mesh._stale_lvec = True + return @property diff --git a/src/underworld3/function/_function.pyx b/src/underworld3/function/_function.pyx index 6c785af02..4822863c4 100644 --- a/src/underworld3/function/_function.pyx +++ b/src/underworld3/function/_function.pyx @@ -48,6 +48,33 @@ cdef extern from "petsc.h" nogil: PetscErrorCode DMSwarmSetMigrateType(PetscDM dm, DMSwarmMigrateType mtype) PetscErrorCode DMSwarmGetMigrateType(PetscDM dm, DMSwarmMigrateType *mtype) + +_fallback_stats_enabled = False +_fallback_stats = {} + + +def _reset_global_evaluate_fallback_stats(enabled=True): + """Reset and optionally enable internal MPI fallback diagnostics.""" + + global _fallback_stats_enabled, _fallback_stats + _fallback_stats_enabled = bool(enabled) + _fallback_stats = { + "calls": 0, + "calls_with_points": 0, + "local_extrapolated_points": 0, + "replicated_points_per_rank": 0, + "max_replicated_points_per_call": 0, + "temporary_bytes_per_rank": 0, + "max_temporary_bytes_per_call": 0, + } + + +def _get_global_evaluate_fallback_stats(): + """Return one rank's internal MPI fallback diagnostics.""" + + return dict(_fallback_stats) + + class UnderworldAppliedFunction(sympy.core.function.AppliedUndef): """ Applied Underworld function representing a mesh variable evaluated at coordinates. @@ -347,6 +374,16 @@ def _lambdify_and_evaluate(expr, coords, interpolated_results, coord_sys=None, m return results.reshape(-1, *shape) +def _global_fallback_indices(return_value, return_mask): + """Indices requiring the parallel best-claim fallback. + + A point needs recovery when migration/location marks it extrapolated or + when interpolation returned a non-finite value despite a located flag. + """ + nonfinite = ~np.isfinite(return_value).all(axis=(1, 2)) + return np.where(return_mask[:, 0, 0] | nonfinite)[0] + + def global_evaluate_nd( expr, coords=None, coord_sys=None, @@ -589,13 +626,32 @@ def global_evaluate_nd( expr, from mpi4py import MPI comm = uw.mpi.comm - ext_idx = np.where(return_mask[:, 0, 0])[0] + # A failed interpolation can occasionally return NaN while reporting + # the point as located. Treat that exactly like an extrapolated/lost + # point so a finite value from the globally nearest rank replaces it. + # This is required by SLCN midpoint tracing: one silent NaN here makes + # the departure point and then the transported history non-finite. + ext_idx = _global_fallback_indices(return_value, return_mask) ext_coords = np.ascontiguousarray(coords_array[ext_idx], dtype=np.float64) counts = np.array(comm.allgather(ext_coords.shape[0]), dtype=int) n_ext_total = int(counts.sum()) + if _fallback_stats_enabled: + _fallback_stats["calls"] += 1 + _fallback_stats["local_extrapolated_points"] += int( + ext_coords.shape[0] + ) + _fallback_stats["replicated_points_per_rank"] += n_ext_total + _fallback_stats["max_replicated_points_per_call"] = max( + _fallback_stats["max_replicated_points_per_call"], + n_ext_total, + ) + if n_ext_total > 0: + if _fallback_stats_enabled: + _fallback_stats["calls_with_points"] += 1 + parts = comm.allgather(ext_coords) all_ext = np.concatenate( [p for p in parts if p.size], axis=0).reshape(n_ext_total, -1) @@ -632,6 +688,30 @@ def global_evaluate_nd( expr, best_flag = np.empty(n_ext_total, dtype=np.int32) comm.Allreduce([contrib_flag, MPI.INT], [best_flag, MPI.INT], op=MPI.SUM) + if _fallback_stats_enabled: + temporary_bytes = sum( + int(array.nbytes) + for array in ( + all_ext, + ext_vals, + ext_flag, + dist2, + min_dist2, + my_claim, + win_rank, + contrib_val, + best_val, + contrib_flag, + best_flag, + ) + ) + temporary_bytes += sum(int(part.nbytes) for part in parts) + _fallback_stats["temporary_bytes_per_rank"] += temporary_bytes + _fallback_stats["max_temporary_bytes_per_call"] = max( + _fallback_stats["max_temporary_bytes_per_call"], + temporary_bytes, + ) + # Scatter this rank's segment of the global set back to its points. offset = int(counts[:comm.rank].sum()) seg = slice(offset, offset + ext_coords.shape[0]) diff --git a/src/underworld3/postprocessing/geoid.py b/src/underworld3/postprocessing/geoid.py index 537a79fee..b35d201f0 100644 --- a/src/underworld3/postprocessing/geoid.py +++ b/src/underworld3/postprocessing/geoid.py @@ -275,11 +275,36 @@ def _rotated_topography_coefficient( harmonic_degree: int, buoyancy_scale: float, response_sign: float, + projection: str, ) -> float: buoyancy_scale = float(buoyancy_scale) if not np.isfinite(buoyancy_scale) or buoyancy_scale == 0.0: raise ValueError("Boundary buoyancy scales must be finite and nonzero.") + if projection == "reaction": + import sympy + from underworld3.maths import BdIntegral + + theta = stokes.mesh.CoordinateSystem.xR[1] + harmonic = sympy.assoc_legendre(harmonic_degree, 0, sympy.cos(theta)) + traction_integral = stokes.boundary_normal_traction_integral( + boundary, + harmonic, + remove_mean=True, + ) + # The reaction is the load functional assembled on the faceted FE + # boundary. Use the matching discrete inner product for the fitted + # coefficient; an analytical spherical norm would mix geometries and + # introduce a chord-area bias, especially on the smaller CMB. + harmonic_norm = float( + BdIntegral(stokes.mesh, fn=harmonic**2, boundary=boundary).evaluate() + ) + return float( + -response_sign * traction_integral / (buoyancy_scale * harmonic_norm) + ) + if projection != "centroid": + raise ValueError("projection must be 'centroid' or 'reaction'.") + coords, sigma_nn = stokes.boundary_normal_traction(boundary, mass="auto") local_rows = np.column_stack( ( @@ -339,16 +364,18 @@ def spherical_shell_response_from_rotated_stokes( planet_radius: float | None = None, gravity: float | None = None, gravitational_constant: float = 6.67430e-11, + projection: str = "centroid", ) -> SphericalShellResponse: r"""Compute spherical-shell response from a rotated-free-slip Stokes solve. - Normal traction recovery is delegated to the existing - :meth:`Stokes.boundary_normal_traction` implementation. This adapter only - projects the two boundary responses onto the unnormalised - axisymmetric :math:`P_l^0` harmonic. Use the pure coefficient functions - directly for other harmonic orders or topography-recovery methods. Density - contrasts, planet radius, and gravity are required when - ``include_self_gravity`` is true. + ``projection="centroid"`` (default) recovers pointwise traction and fits it + over a triangulation of the boundary samples. ``projection="reaction"`` + contracts the assembled nodal reaction directly with the harmonic test + function. The latter is a distributed weak/integral quantity that avoids + consuming slowly converging P2 vertex values on curved boundaries (issue + #414). Use the pure coefficient functions directly for other harmonic + orders or topography-recovery methods. Density contrasts, planet radius, + and gravity are required when ``include_self_gravity`` is true. """ ri, ro, degree = _validate_geometry( @@ -358,6 +385,8 @@ def spherical_shell_response_from_rotated_stokes( ) if not isinstance(include_self_gravity, bool): raise TypeError("include_self_gravity must be True or False.") + if projection not in ("centroid", "reaction"): + raise ValueError("projection must be 'centroid' or 'reaction'.") if degree == 0: raise ValueError( "The rotated-Stokes adapter requires harmonic_degree >= 1 because " @@ -386,6 +415,7 @@ def spherical_shell_response_from_rotated_stokes( degree, surface_buoyancy_scale, 1.0, + projection, ) cmb_topography = _rotated_topography_coefficient( stokes, @@ -394,6 +424,7 @@ def spherical_shell_response_from_rotated_stokes( degree, cmb_buoyancy_scale, -1.0, + projection, ) geoid = spherical_shell_geoid_response( radius_inner=ri, diff --git a/src/underworld3/systems/__init__.py b/src/underworld3/systems/__init__.py index ea6823ab5..c7351e610 100644 --- a/src/underworld3/systems/__init__.py +++ b/src/underworld3/systems/__init__.py @@ -23,6 +23,8 @@ Navier-Stokes equations with inertia. Diffusion : class Pure diffusion (no advection). +AdvDiffusionSUPG : class + Eulerian advection-diffusion with SUPG spatial stabilization. TransientDarcy : class Transient groundwater flow with constant storage. Richards : class @@ -67,6 +69,7 @@ # import diffusion-only solver from .solvers import SNES_Diffusion as Diffusion +from .advdiff_supg import SNES_AdvectionDiffusionSUPG as AdvDiffusionSUPG # Transient Darcy and Richards solvers from .solvers import SNES_TransientDarcy as TransientDarcy diff --git a/src/underworld3/systems/advdiff_supg.py b/src/underworld3/systems/advdiff_supg.py new file mode 100644 index 000000000..0e3cd5f68 --- /dev/null +++ b/src/underworld3/systems/advdiff_supg.py @@ -0,0 +1,717 @@ +r"""Streamline-upwind Petrov-Galerkin scalar transport.""" + +from typing import Optional +import math +from dataclasses import dataclass + +import numpy as np +import sympy +from petsc4py import PETSc + +import underworld3 as uw +import underworld3.timing as timing +from underworld3.function import expression +from underworld3.systems.ddt import Eulerian as Eulerian_DDt +from underworld3.systems.solvers import SNES_Diffusion, _centroid_velocities_nd +from underworld3.checkpoint.state import SnapshottableState + + +@dataclass +class AdvDiffusionSUPGState(SnapshottableState): + """Snapshot metadata not carried by SUPG mesh variables.""" + + time_integrator: str = "bdf" + rate_initialised: bool = False + + +class SNES_AdvectionDiffusionSUPG(SNES_Diffusion): + r"""Implicit scalar advection-diffusion with SUPG stabilization. + + The scalar residual is + + .. math:: + + R = \frac{\mathrm{BDF}(T)}{\Delta t} + + \mathbf{u}\cdot\nabla T - f, + + with pointwise residual terms + + .. math:: + + F_0 = R, \qquad + \mathbf{F}_1 = \boldsymbol{\kappa}\nabla T + + \tau\mathbf{u}R. + + The second contribution to :math:`\mathbf{F}_1` produces + :math:`\tau(\mathbf{u}\cdot\nabla w)R` in the weak form. Diffusion + remains standard Galerkin. + + Parameters + ---------- + mesh : Mesh + Computational mesh. + u_Field : MeshVariable + Scalar field being transported. + V_fn : MeshVariable or sympy Matrix + Prescribed advection velocity. It is frozen during each solve. + order : int, default=1 + Eulerian BDF history order. + theta : float, default=1.0 + Flux integration parameter. Only fully implicit fluxes are currently + supported. + tau : scalar expression, optional + User-provided stabilization parameter. When omitted, a transient + isotropic parameter is computed from a cell-constant streamline + length, local velocity, diffusivity, and timestep. + tau_model : {"generic", "citcoms"}, optional + Automatic stabilization model. ``generic`` uses the optimal 1-D + coth(Pe)-1/Pe relation with a transient scale. ``citcoms`` uses the + clipped steady relation on simplex streamline lengths. This option + does not change the implicit BDF time integrator. + time_integrator : {"bdf", "citcoms"}, default="bdf" + ``bdf`` uses the implicit Eulerian BDF solver. ``citcoms`` uses a + positive P1 row-sum mass, gamma=0.5 predictor, and two fixed residual + corrections. The latter is restricted to continuous P1 fields. + temperature_rate_field : MeshVariable, optional + Stored temperature derivative for the CitcomS predictor-corrector. + Supplying this field gives production checkpoint files a stable name. + DuDt, DFDt : optional + Existing history operators. A supplied ``DuDt`` must be Eulerian and + must not contain a velocity, because advection is represented in R. + + Notes + ----- + Automatic tau currently supports volume simplex meshes and scalar + isotropic diffusivity. Supply ``tau`` explicitly for other meshes or + constitutive models. The default ``bdf`` integrator is implicit. The + ``citcoms`` integrator provides the continuous-P1 row-lumped + predictor-corrector used by CitcomS-style mantle-convection benchmarks. + """ + + @timing.routine_timer_decorator + def __init__( + self, + mesh: uw.discretisation.Mesh, + u_Field: uw.discretisation.MeshVariable, + V_fn, + order: int = 1, + theta: float = 1.0, + tau=None, + tau_model: Optional[str] = None, + time_integrator: str = "bdf", + adv_gamma: float = 0.5, + corrector_steps: int = 2, + temperature_rate_field: Optional[uw.discretisation.MeshVariable] = None, + evalf: Optional[bool] = False, + verbose: bool = False, + DuDt: Optional[Eulerian_DDt] = None, + DFDt=None, + ): + if float(theta) != 1.0: + raise ValueError("AdvDiffusionSUPG currently requires theta=1.0.") + if time_integrator not in ("bdf", "citcoms"): + raise ValueError("time_integrator must be 'bdf' or 'citcoms'.") + if tau_model is None: + tau_model = "citcoms" if time_integrator == "citcoms" else "generic" + if tau_model not in ("generic", "citcoms"): + raise ValueError("tau_model must be 'generic' or 'citcoms'.") + if DuDt is not None and not isinstance(DuDt, Eulerian_DDt): + raise TypeError("AdvDiffusionSUPG requires an Eulerian DuDt operator.") + if DuDt is not None and DuDt.V_fn is not None: + raise ValueError( + "DuDt.V_fn must be None; AdvDiffusionSUPG includes advection " + "in its strong residual." + ) + if time_integrator == "citcoms": + if u_Field.degree != 1 or not u_Field.continuous: + raise ValueError( + "The CitcomS predictor-corrector requires continuous P1 " + "temperature." + ) + if DuDt is not None or DFDt is not None: + raise ValueError( + "The CitcomS predictor-corrector manages its own derivative " + "state; do not supply DuDt or DFDt." + ) + if not 0.0 < float(adv_gamma) <= 1.0: + raise ValueError("adv_gamma must be in (0, 1].") + if int(corrector_steps) < 1: + raise ValueError("corrector_steps must be positive.") + + super().__init__( + mesh, + u_Field, + order=order, + theta=theta, + evalf=evalf, + verbose=verbose, + DuDt=DuDt, + DFDt=DFDt, + ) + + self.V_fn = V_fn + self.tau_model = tau_model + self.time_integrator = time_integrator + self.adv_gamma = float(adv_gamma) + self.corrector_steps = int(corrector_steps) + self._automatic_tau = tau is None + self._supg_h = None + self._supg_tau = None + self._temperature_rate = None + 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._rate_initialised = False + + if self.time_integrator == "citcoms": + if temperature_rate_field is not None: + if ( + 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 continuous scalar P1 " + "variable on the solver mesh." + ) + self._temperature_rate = temperature_rate_field + else: + self._temperature_rate = uw.discretisation.MeshVariable( + f"_supg_dTdt_{self.instance_number}", + mesh, + 1, + degree=1, + continuous=True, + ) + + uw.get_default_model()._register_state_bearer(self) + + if self._automatic_tau: + suffix = self.instance_number + self._supg_h = uw.discretisation.MeshVariable( + f"_supg_h_{suffix}", mesh, 1, degree=0, continuous=False + ) + self._supg_tau = uw.discretisation.MeshVariable( + f"_supg_tau_{suffix}", mesh, 1, degree=0, continuous=False + ) + self._tau = self._supg_tau.sym[0] + else: + self._tau = sympy.sympify(tau) + + @property + def V_fn(self): + """Advection velocity expression.""" + return self._V_fn + + @V_fn.setter + def V_fn(self, value): + self.is_setup = False + self._V_fn = ( + value.sym if isinstance(value, uw.discretisation.MeshVariable) else value + ) + + @property + def tau(self): + """SUPG stabilization parameter used in the residual.""" + return self._tau + + @property + def temperature_rate(self): + """Stored derivative used by the CitcomS predictor-corrector.""" + return self._temperature_rate + + @property + def state(self): + """Return predictor-corrector initialization metadata.""" + return AdvDiffusionSUPGState( + time_integrator=self.time_integrator, + rate_initialised=self._rate_initialised, + ) + + @state.setter + def state(self, state): + if not isinstance(state, AdvDiffusionSUPGState): + raise TypeError("AdvDiffusionSUPG state has the wrong type.") + if state.time_integrator != self.time_integrator: + raise ValueError("AdvDiffusionSUPG time integrator changed since snapshot.") + self._rate_initialised = bool(state.rate_initialised) + + def _strong_transport_residual(self): + gradient = self.mesh.vector.gradient(self.u.sym) + advection = sympy.Matrix((self.V_fn.dot(gradient),)) + if self.time_integrator == "citcoms": + time_derivative = self._temperature_rate.sym + else: + time_derivative = self.DuDt.bdf() / self.delta_t + return time_derivative + advection - self.f + + def _simplex_data(self): + """Return local simplex connectivity, basis gradients, and volumes.""" + 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 + ) + if cells is None or self.mesh.dim != self.mesh.cdim: + raise NotImplementedError( + "Automatic SUPG operations require a 2-D or 3-D volume " "simplex mesh." + ) + + 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.constitutive_model.K) + 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 + + @property + def F0(self): + """Transient-advection-source strong residual.""" + value = expression( + r"f_0^{SUPG}", + self._strong_transport_residual(), + "SUPG transient-advection-source residual", + _unique_name_generation=True, + ) + self._f0 = value + return value + + @property + def F1(self): + """Galerkin diffusion flux plus streamline stabilization flux.""" + residual = self._strong_transport_residual()[0] + diffusion_flux = ( + self.constitutive_model.flux.T + if self.time_integrator == "citcoms" + else self.DFDt.adams_moulton_flux() + ) + value = expression( + r"\mathbf{F}_1^{SUPG}", + diffusion_flux + self.tau * self.V_fn * residual, + "Diffusive and SUPG streamline flux", + _unique_name_generation=True, + ) + self._f1 = value + return value + + def _update_automatic_tau(self): + """Update local simplex streamline lengths and automatic tau values.""" + if not self._automatic_tau: + return + if self.constitutive_model is None: + raise RuntimeError( + "Set constitutive_model before solving AdvDiffusionSUPG." + ) + + _, gradients, _ = self._simplex_data() + + 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]) + xi = np.empty_like(pe) + small = np.abs(pe) < 1.0e-3 + pe_small = pe[small] + xi[small] = pe_small / 3.0 - pe_small**3 / 45.0 + 2.0 * pe_small**5 / 945.0 + xi[~small] = 1.0 / np.tanh(pe[~small]) - 1.0 / pe[~small] + if self.tau_model == "generic": + tau_steady[diffusive] = ( + h_stream[diffusive] * xi / (2.0 * speed[diffusive]) + ) + else: + 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]) + + if self.tau_model == "generic": + dt = float(self.delta_t.data) + if dt <= 0.0: + raise ValueError("AdvDiffusionSUPG requires a positive timestep.") + tau_values = np.divide( + 1.0, + np.sqrt( + (2.0 / dt) ** 2 + + np.divide( + 1.0, + tau_steady**2, + out=np.full_like(tau_steady, np.inf), + where=tau_steady > 0.0, + ) + ), + out=np.zeros_like(tau_steady), + where=tau_steady > 0.0, + ) + else: + tau_values = tau_steady + + if self._supg_h.data.shape[0] != h_stream.size: + raise RuntimeError("SUPG P0 field and local simplex counts do not match.") + self._supg_h.data[:, 0] = h_stream + self._supg_tau.data[:, 0] = tau_values + + def _setup_citcoms_residual(self, verbose=False): + """Build the reusable residual assembler for predictor-corrector steps.""" + if not self.constitutive_model._solver_is_setup: + self._needs_function_rewire = True + self._build(verbose, False, None) + self.is_setup = True + self.constitutive_model._solver_is_setup = True + + 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.dm.createLocalVector() + global_mass = self.dm.createGlobalVector() + local_mass.set(0.0) + global_mass.set(0.0) + section = self.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.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.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_dt(self): + """Estimate a simplex advection-diffusion timestep. + + The CitcomS-compatible predictor-corrector uses + ``0.9 * min(1/max(lambda_adv), 2/max(rowsum(abs(M_L^-1 K))))``. + The same conservative value is also available for the implicit BDF + path as a resolution-accuracy estimate. + """ + 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 + elif self.time_integrator != "citcoms": + local_diff_rate = ( + float( + np.max( + 2.0 + * self.mesh.dim + * diffusivity + / np.maximum(self.mesh._radii**2, np.finfo(float).tiny) + ) + ) + if diffusivity.size + else 0.0 + ) + diff_rate = uw.mpi.comm.allreduce(local_diff_rate, op=MPI.MAX) + dt_diff = 2.0 / diff_rate + 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.dm.createMatrix() + stiffness.setOption(PETSc.Mat.Option.NEW_NONZERO_LOCATION_ERR, False) + section = self.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 _compute_citcoms_residual(self, solution=None, residual=None): + """Assemble the residual at the current temperature and rate.""" + if solution is None: + solution = self.dm.createGlobalVector() + if residual is None: + residual = solution.duplicate() + 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) + return solution, residual + + def _solve_citcoms(self, timestep, verbose=False): + """Advance one CitcomS-compatible predictor-corrector timestep.""" + if timestep is None: + timestep = float(self.delta_t.data) + self.delta_t = timestep + dt = float(self.delta_t.data) + if dt <= 0.0: + raise ValueError("AdvDiffusionSUPG requires a positive timestep.") + + 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.data[:, 0] = 0.0 + 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.dm.globalToLocal(delta_rate, self._temperature_rate.vec) + self.mesh._stale_lvec = True + self._rate_initialised = True + + self.u.data[:, 0] += ( + (1.0 - self.adv_gamma) * dt * self._temperature_rate.data[:, 0] + ) + self._temperature_rate.data[:, 0] = 0.0 + self.mesh._stale_lvec = True + + from underworld3.cython.petsc_discretisation import ( + petsc_dm_insert_boundary_values, + ) + + for _ in range(self.corrector_steps): + self._compute_citcoms_residual(temperature_global, residual) + delta_rate.pointwiseDivide(residual, mass) + delta_rate.scale(-1.0) + + rate_global.set(0.0) + self.dm.localToGlobal(self._temperature_rate.vec, rate_global, addv=False) + rate_global.axpy(1.0, delta_rate) + temperature_global.axpy(self.adv_gamma * dt, delta_rate) + + self._temperature_rate.vec.set(0.0) + self.u.vec.set(0.0) + self.dm.globalToLocal(rate_global, self._temperature_rate.vec) + self.dm.globalToLocal(temperature_global, self.u.vec) + petsc_dm_insert_boundary_values(self.dm, self.u.vec) + self.mesh._stale_lvec = True + + self.is_setup = True + self.constitutive_model._solver_is_setup = True + return + + @timing.routine_timer_decorator + def solve( + self, + zero_init_guess: bool = None, + timestep: float = None, + evalf: bool = False, + _force_setup: bool = False, + verbose: bool = False, + divergence_retries: int = 0, + ): + """Update automatic stabilization and solve one implicit timestep.""" + if self.time_integrator == "citcoms": + return self._solve_citcoms(timestep, verbose=verbose) + if timestep is not None: + self.delta_t = timestep + self._update_automatic_tau() + return super().solve( + zero_init_guess=zero_init_guess, + timestep=timestep, + evalf=evalf, + _force_setup=_force_setup, + verbose=verbose, + divergence_retries=divergence_retries, + ) diff --git a/src/underworld3/systems/ddt.py b/src/underworld3/systems/ddt.py index e198aad57..e5b93ed13 100644 --- a/src/underworld3/systems/ddt.py +++ b/src/underworld3/systems/ddt.py @@ -116,6 +116,29 @@ class DDtSymbolicState(_DDtCoreState): psi_star: list = field(default_factory=list) + def __deepcopy__(self, memo): + """Copy history containers without reconstructing symbolic atoms. + + SymPy matrices are value containers but their expression atoms include + UWexpression objects whose identity binds them to live parameter and + coefficient registries. Generic ``copy.deepcopy`` reconstructs those + Symbol subclasses without their wrapped value, producing invalid atoms + after snapshot restore. Matrix ``copy()`` keeps the immutable symbolic + atoms while separating the mutable history list and matrices. + """ + import copy + + duplicate = type(self)( + _schema_version=self._schema_version, + dt_history=copy.deepcopy(self.dt_history, memo), + history_initialised=self.history_initialised, + n_solves_completed=self.n_solves_completed, + dt=copy.deepcopy(self.dt, memo), + psi_star=[value.copy() for value in self.psi_star], + ) + memo[id(self)] = duplicate + return duplicate + @dataclass class DDtEulerianState(_DDtCoreState): @@ -3472,4 +3495,3 @@ def update_post_solve( self._n_solves_completed += 1 return - diff --git a/src/underworld3/systems/solvers.py b/src/underworld3/systems/solvers.py index 8043975bf..98d8c7e96 100644 --- a/src/underworld3/systems/solvers.py +++ b/src/underworld3/systems/solvers.py @@ -4151,6 +4151,32 @@ def adv_diff_slcn_problem_description(self): return + def _prepare_flux_history(self): + """Return whether the Adams-Moulton flux needs stored history. + + For first-order Backward Euler (theta=1), the flux expression is + exactly F(u[n+1]); every stored-history coefficient is zero. Refresh + the symbolic coefficients in case theta changed after construction, + then let solve() skip the otherwise unused projection and + characteristic trace-back. + """ + + flux_is_current_only = ( + getattr(self.DFDt, "order", None) == 1 + and float(getattr(self.DFDt, "theta", float("nan"))) == 1.0 + ) + if flux_is_current_only: + from underworld3.systems.ddt import _update_am_values + + _update_am_values( + self.DFDt._am_coeffs, + effective_order=1, + theta=1.0, + ) + return False + + return True + @property def f(self): r"""Source term for the advection-diffusion equation. @@ -4432,6 +4458,8 @@ def solve( self._needs_function_rewire = True self.DFDt.psi_fn = self.constitutive_model.flux.T + flux_history_active = self._prepare_flux_history() + if not self.is_setup: self._setup_pointwise_functions(verbose) self._setup_discretisation(verbose) @@ -4441,7 +4469,10 @@ def solve( # SemiLagrange and Lagrange may have different sequencing. self.DuDt.update_pre_solve(timestep, verbose=verbose, evalf=_evalf) - self.DFDt.update_pre_solve(timestep, verbose=verbose, evalf=_evalf) + if flux_history_active: + self.DFDt.update_pre_solve( + timestep, verbose=verbose, evalf=_evalf + ) super().solve(zero_init_guess, _force_setup, divergence_retries=divergence_retries) @@ -4449,7 +4480,10 @@ def solve( _invalidate_solution_cache(self.u) self.DuDt.update_post_solve(timestep, verbose=verbose, evalf=_evalf) - self.DFDt.update_post_solve(timestep, verbose=verbose, evalf=_evalf) + if flux_history_active: + self.DFDt.update_post_solve( + timestep, verbose=verbose, evalf=_evalf + ) self.is_setup = True self.constitutive_model._solver_is_setup = True diff --git a/src/underworld3/utilities/boundary_flux.py b/src/underworld3/utilities/boundary_flux.py index c0b8e5b2e..abb211a63 100644 --- a/src/underworld3/utilities/boundary_flux.py +++ b/src/underworld3/utilities/boundary_flux.py @@ -672,6 +672,37 @@ def boundary_flux(solver, boundary, mass="auto", remove_mean=False, normal=None) return xs, (np.column_stack(cols) if nodes else np.zeros((0, ncomp))) +def boundary_flux_integral(solver, boundary): + r"""Return the integrated scalar flux through ``boundary``. + + For a scalar solver, summing the consistent nodal reactions on the queried + boundary gives :math:`\int_\Gamma F\cdot\hat n\,d\Gamma` directly because + the boundary basis is a partition of unity. The raw reactions are partial + on partition-cut nodes, so the final sum is collective across ranks. + + This path is intended for integral diagnostics such as Nusselt numbers. It + avoids pointwise boundary-mass recovery and a temporary MeshVariable. Use + :meth:`boundary_flux` or :meth:`boundary_flux_field` when nodal values are + required. + """ + dm = solver.dm + ra = np.asarray(solver._assemble_volume_reaction()).ravel() + nodes, lsec, _csec, _cvec, _v0, _v1, _edge_nodes = _boundary_field_nodes( + solver, boundary, field_id=0 + ) + ncomp = lsec.getFieldComponents(0) + if ncomp != 1: + raise ValueError( + "boundary_flux_integral requires a scalar solver field; use " + "boundary_flux(..., normal=...) for vector traction." + ) + local_integral = sum( + float(ra[lsec.getFieldOffset(point, 0) + slot]) + for point, slot, _coordinate in nodes + ) + return float(dm.comm.tompi4py().allreduce(local_integral, op=MPI.SUM)) + + def write_boundary_scalar_field(solver, field, value_by_key, dim): """Write ``value_by_key`` (coordinate-key → scalar) onto a scalar MeshVariable ``field`` at the matching nodes; interior nodes untouched. Returns ``field``. diff --git a/src/underworld3/utilities/rotated_bc.py b/src/underworld3/utilities/rotated_bc.py index 47a06683a..646162681 100644 --- a/src/underworld3/utilities/rotated_bc.py +++ b/src/underworld3/utilities/rotated_bc.py @@ -2135,6 +2135,85 @@ def boundary_normal_traction(solver, boundary, solve_result, mass="auto"): remove_mean=True, partial_reaction=False) +def boundary_normal_traction_integral(solver, boundary, solve_result, fn, + remove_mean=True): + r"""Return ``integral((sigma_nn - mean) * fn, boundary)`` directly from the + assembled rotated-constraint reaction. + + This is the weak/integral counterpart of :func:`boundary_normal_traction`. + It contracts the nodal reaction with ``fn`` at the velocity interpolation + nodes before any pointwise boundary-mass recovery. On curved P2 boundaries + this is a fitted quantity and therefore does not consume the slowly + converging recovered vertex values described in issue #414. + + Each assembled reaction DOF is counted on its owning rank only, followed by + an MPI sum. The operation is distributed and does not gather boundary + topology or recovered values onto rank zero. ``remove_mean=True`` removes + the constant-traction gauge using boundary integrals of ``fn`` and one. + """ + if not isinstance(remove_mean, (bool, np.bool_)): + raise TypeError("remove_mean must be True or False.") + + import underworld3 as uw + + dm = solver.dm + comm = dm.comm.tompi4py() + dim = solver.mesh.dim + rc = solve_result["reaction"] + rstart, rend = rc.getOwnershipRange() + rcl = dm.getLocalVec() + dm.globalToLocal(rc, rcl) + + try: + rca = np.asarray(rcl.getArray()) + lsec = dm.getLocalSection() + l2g = dm.getLGMap() + csec = dm.getCoordinateSection() + cvec = np.asarray(dm.getCoordinatesLocal().array).reshape(-1, dim) + v0, v1 = dm.getDepthStratum(0) + normal = dict(_boundary_spec(s) for s in solve_result["boundaries"]).get(boundary) + nodes = _boundary_velocity_nodes(solver, boundary, normal=normal) + + owned_coords = [] + owned_reactions = [] + for q, nrm in nodes: + lo = lsec.getFieldOffset(q, _VELOCITY_FIELD) + global_row = int(l2g.apply([lo])[0]) + if not rstart <= global_row < rend: + continue + owned_coords.append(_point_coord(dm, dim, cvec, csec, v0, v1, q)) + # sigma_nn load = -n.r_c, matching boundary_normal_traction(). + owned_reactions.append(-float(np.dot(nrm, rca[lo:lo + dim]))) + finally: + dm.restoreLocalVec(rcl) + + if owned_coords: + coords = np.ascontiguousarray(owned_coords, dtype=float) + weights = np.asarray(uw.function.evaluate(fn, coords), dtype=float).reshape(-1) + if weights.size == 1 and len(owned_reactions) != 1: + weights = np.full(len(owned_reactions), float(weights[0])) + if weights.size != len(owned_reactions): + raise ValueError("fn must evaluate to one scalar per boundary node.") + local_weighted = float(np.dot(owned_reactions, weights)) + local_total = float(np.sum(owned_reactions)) + else: + local_weighted = 0.0 + local_total = 0.0 + + weighted = float(comm.allreduce(local_weighted)) + if not remove_mean: + return weighted + + total = float(comm.allreduce(local_total)) + area = float(uw.maths.BdIntegral(solver.mesh, fn=1.0, boundary=boundary).evaluate()) + if not np.isfinite(area) or area <= 0.0: + raise RuntimeError(f"Boundary {boundary!r} has non-positive area {area}.") + fn_integral = float( + uw.maths.BdIntegral(solver.mesh, fn=fn, boundary=boundary).evaluate() + ) + return weighted - (total / area) * fn_integral + + def dynamic_topography_field(solver, boundary, solve_result, field, buoyancy_scale=1.0, mass="auto"): """Populate a scalar MeshVariable ``field`` with the dynamic topography diff --git a/tests/parallel/ptest_0010_snapshot_disk.py b/tests/parallel/ptest_0010_snapshot_disk.py index b96823135..b4f53dfe6 100644 --- a/tests/parallel/ptest_0010_snapshot_disk.py +++ b/tests/parallel/ptest_0010_snapshot_disk.py @@ -1,10 +1,7 @@ """Parallel (MPI) test of the on-disk snapshot path (v1.1). -Phase 6 of the snapshot toolkit: per-rank swarm sidecars. The mesh -+ mesh-variable disk path is already parallel-correct via #146's -PETSc-collective HDF5 viewer; the swarm sidecar layer needs its -own per-rank file per swarm. This ptest exercises both together at -multi-rank. +Phase 6 of the snapshot toolkit: exact same-rank mesh-variable vectors and +per-rank swarm sidecars. This ptest exercises both layers together under MPI. Run (4 ranks exercises cross-rank distribution of swarm particles): @@ -19,7 +16,8 @@ state (verified by per-rank attrs on the sidecar). 3. Round-trip is exact: scribble all variables + swarm coords + swarm-var data, model.load_state(file=...), gathered (gid, x, y, - material) tables sorted by gid are np.array_equal. + material) tables sorted by gid are np.array_equal, and every restored + mesh-variable dof matches its coordinate-defined analytic value. """ import os @@ -66,8 +64,7 @@ def build(): def global_sorted_state(T, swarm, gid, material): - """Gather (gid, x, y, material, T-value-by-coord-bin) across ranks - + sort by gid → order/rank-independent canonical view.""" + """Return rank-independent swarm state and mesh-field error.""" g = gid.data[:, 0].copy() coords = swarm._particle_coordinates.data.copy() m = material.data[:, 0].copy() @@ -79,19 +76,12 @@ def global_sorted_state(T, swarm, gid, material): order = np.argsort(full[:, 0], kind="stable") swarm_state = full[order] - # T round-trip check: gather partition-invariant scalars - # (max, sum) rather than the full (coord, value) table — DOFs at - # partition boundaries are visible to multiple ranks and would - # appear duplicated/reordered in a gathered table, even though - # the underlying data is bit-exact. t_arr = np.asarray(T.array[...]).reshape(-1) - t_max = comm.allreduce(float(t_arr.max()) if t_arr.size else -np.inf, - op=MPI.MAX) - t_min = comm.allreduce(float(t_arr.min()) if t_arr.size else np.inf, - op=MPI.MIN) - # bit-exact float sum across ranks is non-deterministic in general - # (non-associative); use min/max as bit-exact invariants instead. - return swarm_state, (t_max, t_min) + t_coords = np.asarray(T.coords) + expected = t_coords[:, 0] - t_coords[:, 1] + local_error = float(np.max(np.abs(t_arr - expected))) if t_arr.size else 0.0 + global_error = comm.allreduce(local_error, op=MPI.MAX) + return swarm_state, global_error def main(): @@ -146,10 +136,7 @@ def main(): post_swarm, post_T = global_sorted_state(T, swarm, gid, material) swarm_ok = np.array_equal(pre_swarm, post_swarm) - # T is checked via partition-invariant min/max scalars (see note - # in global_sorted_state — gathered DOFs include partition- - # boundary duplicates that resist a global-table comparison). - T_ok = (pre_T == post_T) + T_ok = pre_T == 0.0 and post_T == 0.0 count_ok = pre_count == post_count tracker_ok = (model.tracker.time == 1.5 and model.tracker.step == 7) @@ -161,7 +148,8 @@ def main(): flush=True) print(f" P3 swarm (coords + gid + material) exact: {swarm_ok}", flush=True) - print(f" P4 T (mesh-variable DOFs) exact: {T_ok}", + print(f" P4 T (mesh-variable DOFs) exact: {T_ok} " + f"(max error={post_T:.3e})", flush=True) print(f" P5 tracker state restored: {tracker_ok}", flush=True) diff --git a/tests/parallel/test_0760_swarm_cache_migration.py b/tests/parallel/test_0760_swarm_cache_migration.py index d27475afa..10458a684 100644 --- a/tests/parallel/test_0760_swarm_cache_migration.py +++ b/tests/parallel/test_0760_swarm_cache_migration.py @@ -17,6 +17,7 @@ import pytest import numpy as np import underworld3 as uw +from underworld3.function._function import _global_fallback_indices from mpi4py import MPI pytestmark = [pytest.mark.mpi(min_size=2), pytest.mark.timeout(60)] @@ -101,3 +102,50 @@ def test_global_evaluate_displaced_nodes(): f"Rank {uw.mpi.rank}: expected {node_coords.shape[0]} results, " f"got {result.shape[0]}" ) + + +@pytest.mark.level_1 +@pytest.mark.tier_a +def test_global_fallback_includes_nonfinite_located_values(): + """Located NaN values must enter the same recovery path as lost points.""" + values = np.ones((4, 1, 2)) + values[1, 0, 0] = np.nan + mask = np.zeros((4, 1, 1), dtype=bool) + mask[2, 0, 0] = True + + assert np.array_equal( + _global_fallback_indices(values, mask), + np.array([1, 2]), + ) + + +@pytest.mark.mpi(min_size=2) +@pytest.mark.level_2 +@pytest.mark.tier_a +def test_global_evaluate_fallback_diagnostics_count_replicated_storage(): + """Opt-in diagnostics report actual collective fallback allocation.""" + from underworld3.function._function import ( + _get_global_evaluate_fallback_stats, + _reset_global_evaluate_fallback_stats, + ) + + mesh = uw.meshing.UnstructuredSimplexBox( + minCoords=(0.0, 0.0), maxCoords=(1.0, 1.0) + ) + field = uw.discretisation.MeshVariable("fallback_field", mesh, 1, degree=1) + field.data[:, 0] = field.coords[:, 0] + field.coords[:, 1] + coords = np.array([[-0.05, 0.5], [1.05, 0.5]], dtype=np.float64) + + _reset_global_evaluate_fallback_stats(enabled=True) + values = uw.function.global_evaluate(field.sym, coords) + stats = _get_global_evaluate_fallback_stats() + _reset_global_evaluate_fallback_stats(enabled=False) + + assert np.all(np.isfinite(values)) + assert stats["calls"] == 1 + assert stats["calls_with_points"] == 1 + assert stats["local_extrapolated_points"] == coords.shape[0] + assert stats["replicated_points_per_rank"] == coords.shape[0] * uw.mpi.size + assert stats["max_replicated_points_per_call"] == coords.shape[0] * uw.mpi.size + assert stats["temporary_bytes_per_rank"] > 0 + assert stats["max_temporary_bytes_per_call"] > 0 diff --git a/tests/parallel/test_1065_boundary_flux_parallel.py b/tests/parallel/test_1065_boundary_flux_parallel.py index 7f20a9636..37038f1e4 100644 --- a/tests/parallel/test_1065_boundary_flux_parallel.py +++ b/tests/parallel/test_1065_boundary_flux_parallel.py @@ -21,6 +21,8 @@ # SERIAL reference: BdIntegral of the flux field over Bottom. `python `. GOLDEN_BDFLUX = -1.731543e-01 +ANALYTIC_DIRECT_INTEGRAL = -2.0 / np.sinh(np.pi) +GOLDEN_DIRECT_INTEGRAL = -1.731790673330021e-01 def _flux_diagnostics(res=48): @@ -44,6 +46,7 @@ def _flux_diagnostics(res=48): xs, flux = poisson.boundary_flux("Bottom") poisson.boundary_flux_field("Bottom", q) bd_q = float(uw.maths.BdIntegral(mesh=mesh, fn=q.sym[0], boundary="Bottom").evaluate()) + direct_integral = poisson.boundary_flux_integral("Bottom") # gather + dedup for a whole-boundary relL2 vs analytic (on rank 0, then bcast) comm = uw.mpi.comm @@ -61,17 +64,26 @@ def _flux_diagnostics(res=48): c = np.dot(F, q_an) / (np.linalg.norm(F) * np.linalg.norm(q_an)) F = F if c >= 0 else -F relL2 = float(np.linalg.norm(F - q_an) / np.linalg.norm(q_an)) - return bd_q, comm.bcast(relL2, root=0) + return bd_q, direct_integral, comm.bcast(relL2, root=0) def test_boundary_flux_partition_independent(): """boundary_flux reproduces the serial reference at np=2 and np=4 (flux boundary cut at np=4): both the collective BdIntegral of the flux field and the whole-boundary accuracy vs analytic.""" - bd_q, relL2 = _flux_diagnostics(res=48) + bd_q, direct_integral, relL2 = _flux_diagnostics(res=48) assert np.isclose(bd_q, GOLDEN_BDFLUX, rtol=1e-5, atol=0), ( f"BdIntegral flux differs serial vs np={uw.mpi.size}: {GOLDEN_BDFLUX} vs {bd_q}") assert relL2 < 0.01, f"heat flux relL2 vs analytic {relL2:.4f} too large at np={uw.mpi.size}" + assert np.isclose( + direct_integral, GOLDEN_DIRECT_INTEGRAL, rtol=1.0e-10, atol=0.0 + ), ( + "Direct reaction integral differs from the serial reference at " + f"np={uw.mpi.size}: {GOLDEN_DIRECT_INTEGRAL} vs {direct_integral}" + ) + assert np.isclose( + direct_integral, ANALYTIC_DIRECT_INTEGRAL, rtol=1.0e-7, atol=0.0 + ) def _uniform_flux_3d_error(degree, mass): @@ -133,6 +145,6 @@ def test_boundary_flux_degree3_partition_independent(): if __name__ == "__main__": - _b, _r = _flux_diagnostics() + _b, _i, _r = _flux_diagnostics() if uw.mpi.rank == 0: print(f"DIAG_FLUX bd_q={_b:.9e} relL2={_r:.4f}") diff --git a/tests/parallel/test_1071_spherical_shell_geoid_parallel.py b/tests/parallel/test_1071_spherical_shell_geoid_parallel.py index 189d61217..9ad0b6e15 100644 --- a/tests/parallel/test_1071_spherical_shell_geoid_parallel.py +++ b/tests/parallel/test_1071_spherical_shell_geoid_parallel.py @@ -56,6 +56,7 @@ def test_rotated_spherical_shell_geoid_matches_serial_reference(): planet_radius=6370000.0, gravity=9.8, gravitational_constant=6.67e-11, + projection="reaction", ) values = np.array( [ @@ -71,7 +72,7 @@ def test_rotated_spherical_shell_geoid_matches_serial_reference(): ) for rank_values in uw.mpi.comm.allgather(values): - assert np.array_equal(rank_values, values) + np.testing.assert_allclose(rank_values, values, rtol=1.0e-13, atol=1.0e-14) zhong_table_2 = np.array( [0.41920, 0.77060, 0.02579, 0.03206, 0.49980, 0.93130, 0.04486, 0.05461] diff --git a/tests/test_0007_snapshot_inmemory.py b/tests/test_0007_snapshot_inmemory.py index d52fee33a..09dcc8674 100644 --- a/tests/test_0007_snapshot_inmemory.py +++ b/tests/test_0007_snapshot_inmemory.py @@ -402,6 +402,37 @@ def test_eulerian_ddt_roundtrip(): assert ddt.state.psi_star_var_names == state_pre.psi_star_var_names +def test_symbolic_flux_history_remains_valid_after_solver_restore(): + """Deep-copying symbolic history must preserve live UWexpression atoms.""" + import numpy as np + import underworld3 as uw + + uw.reset_default_model() + model = uw.get_default_model() + mesh = uw.meshing.UnstructuredSimplexBox( + minCoords=(0.0, 0.0), maxCoords=(1.0, 1.0), cellSize=0.3 + ) + temperature = uw.discretisation.MeshVariable( + "T_symbolic_restore", mesh, 1, degree=1 + ) + with mesh.access(temperature): + temperature.data[:, 0] = temperature.coords[:, 0] + + diffusion = uw.systems.Diffusion( + mesh, u_Field=temperature, order=2, theta=1.0 + ) + diffusion.constitutive_model = uw.constitutive_models.DiffusionModel + diffusion.constitutive_model.Parameters.diffusivity = 0.05 + for _ in range(3): + diffusion.solve(timestep=0.01, zero_init_guess=False) + + snapshot = model.save_state() + model.load_state(snapshot) + diffusion.solve(timestep=0.01, zero_init_guess=False) + + assert np.all(np.isfinite(temperature.data)) + + def test_semilagrangian_ddt_roundtrip(): import underworld3 as uw from underworld3.systems.ddt import DDtSemiLagrangianState @@ -766,4 +797,3 @@ def test_continuation_bit_identical_across_stash_and_recover(): _assert_bit_identical(ctrl, stash, "stash-and-recover") - diff --git a/tests/test_0010_snapshot_disk_format.py b/tests/test_0010_snapshot_disk_format.py index 56609dfee..46fdf89b8 100644 --- a/tests/test_0010_snapshot_disk_format.py +++ b/tests/test_0010_snapshot_disk_format.py @@ -193,6 +193,7 @@ def _fresh_model_mesh_and_vars(): def test_write_snapshot_produces_wrapper_and_bulk_dir(tmp_path): """The two artifacts the convention promises: wrapper file + sibling .bulk/ directory containing PETSc HDF5 files.""" + import h5py import os import underworld3 as uw @@ -213,6 +214,15 @@ def test_write_snapshot_produces_wrapper_and_bulk_dir(tmp_path): assert any("T.00000.h5" in f for f in files) assert any("V.00000.h5" in f for f in files) + for variable_name in ("T", "V"): + variable_file = next( + filename + for filename in files + if filename.endswith(f".{variable_name}.00000.h5") + ) + with h5py.File(os.path.join(bulk, variable_file), "r") as h5: + assert variable_name in h5["uw_checkpoint"] + def test_write_snapshot_populates_wrapper_layout(tmp_path): """The wrapper carries the per-mesh + per-variable metadata that @@ -250,8 +260,7 @@ def test_write_snapshot_populates_wrapper_layout(tmp_path): def test_write_read_snapshot_bit_exact_roundtrip(tmp_path): """The core phase-2 guarantee: write a snapshot, scribble all variables, read snapshot back, all variables match write-time - values bit-for-bit (#146's PETSc DMPlex same-rank reload, just - delivered via the wrapper).""" + values bit-for-bit through the exact same-layout PETSc vector payload.""" import underworld3 as uw uw, model, mesh, T, V = _fresh_model_mesh_and_vars() @@ -280,6 +289,49 @@ def test_write_read_snapshot_bit_exact_roundtrip(tmp_path): ) +def test_disk_restore_refreshes_packed_auxiliary_fields_before_solve(tmp_path): + """A solve after disk restore must use restored coefficient fields.""" + import underworld3 as uw + + uw.reset_default_model() + model = uw.get_default_model() + mesh = uw.meshing.UnstructuredSimplexBox( + minCoords=(0.0, 0.0), + maxCoords=(1.0, 1.0), + cellSize=0.5, + ) + solution = uw.discretisation.MeshVariable("U", mesh, 1, degree=1) + source = uw.discretisation.MeshVariable("source", mesh, 1, degree=1) + + poisson = uw.systems.Poisson(mesh, u_Field=solution) + poisson.constitutive_model = uw.constitutive_models.DiffusionModel + poisson.constitutive_model.Parameters.diffusivity = 1.0 + poisson.f = source.sym[0] + poisson.add_dirichlet_bc(0.0, "All_Boundaries") + + source.array[...] = 1.0 + poisson.solve(zero_init_guess=True) + reference = np.asarray(solution.array[...]).copy() + + path = str(tmp_path / "auxiliary.snap.h5") + model.save_state(file=path) + + source.array[...] = 4.0 + poisson.solve(zero_init_guess=True) + assert not np.allclose(np.asarray(solution.array[...]), reference) + + model.load_state(path) + assert mesh._stale_lvec is True + poisson.solve(zero_init_guess=True) + + assert np.allclose( + np.asarray(solution.array[...]), + reference, + rtol=0.0, + atol=1.0e-5, + ) + + def test_read_snapshot_rejects_missing_bulk_dir(tmp_path): """If the user moves the wrapper without the bulk dir, read fails with a clear pointer rather than an obscure h5py error.""" @@ -299,6 +351,22 @@ def test_read_snapshot_rejects_missing_bulk_dir(tmp_path): model.load_state(path) +def test_read_snapshot_rejects_different_mpi_rank_count(tmp_path): + """Exact disk restart must not silently remap a different MPI layout.""" + import h5py + import underworld3 as uw + + uw, model, mesh, T, V = _fresh_model_mesh_and_vars() + path = str(tmp_path / "rank_count.snap.h5") + model.save_state(file=path) + + with h5py.File(path, "r+") as h5: + h5["metadata"].attrs["mpi_ranks_at_write"] = uw.mpi.size + 1 + + with pytest.raises(ValueError, match="same rank count"): + model.load_state(path) + + def test_read_snapshot_rejects_mismatched_mesh(tmp_path): """If the target model's meshes don't match the snapshot's, raise clearly — mesh-rebuild on read is v1.2 scope.""" diff --git a/tests/test_1019_boundary_flux.py b/tests/test_1019_boundary_flux.py index 86e67faa1..944f712d5 100644 --- a/tests/test_1019_boundary_flux.py +++ b/tests/test_1019_boundary_flux.py @@ -38,17 +38,18 @@ def _heatflux_diagnostics(res=48): poisson.boundary_flux_field("Bottom", q) # field is symbolically usable bd_q = float(uw.maths.BdIntegral(mesh=mesh, fn=q.sym[0], boundary="Bottom").evaluate()) + direct_integral = poisson.boundary_flux_integral("Bottom") xc = np.asarray(xs)[:, 0] if len(xs) else np.zeros(0) q_an = np.pi * np.sin(np.pi * xc) / np.sinh(np.pi) # analytic outward flux - return np.asarray(flux), q_an, bd_q + return np.asarray(flux), q_an, bd_q, direct_integral @pytest.mark.skipif(uw.mpi.size > 1, reason="serial diagnostic: rank-local flux norms are 0/0 on non-owning ranks") def test_boundary_flux_scalar_heatflux_serial(): """Surface heat flux reproduces the analytic flux to high accuracy, and its mean is the (analytic) Nusselt number — NOT removed.""" - flux, q_an, bd_q = _heatflux_diagnostics(res=48) + flux, q_an, bd_q, direct_integral = _heatflux_diagnostics(res=48) corr = np.dot(flux, q_an) / (np.linalg.norm(flux) * np.linalg.norm(q_an)) fa = flux if corr >= 0 else -flux relL2 = np.linalg.norm(fa - q_an) / np.linalg.norm(q_an) @@ -58,6 +59,11 @@ def test_boundary_flux_scalar_heatflux_serial(): assert np.isclose(abs(fa.mean()), 2.0 / np.sinh(np.pi), rtol=0.02), ( f"mean flux {fa.mean():.4f} != Nusselt {2.0/np.sinh(np.pi):.4f}") assert abs(bd_q) > 0.0 # field populated + usable + # The reaction sum is the integral itself, whereas integrating the recovered + # pointwise field includes its finite-resolution projection error. + assert np.isclose( + abs(direct_integral), 2.0 / np.sinh(np.pi), rtol=1.0e-7, atol=0.0 + ) def _uniform_flux_3d(degree, mass): @@ -102,7 +108,7 @@ def test_boundary_flux_3d_p2_lumped_rejected(): if __name__ == "__main__": - _f, _a, _b = _heatflux_diagnostics() + _f, _a, _b, _i = _heatflux_diagnostics() c = np.dot(_f, _a) / (np.linalg.norm(_f) * np.linalg.norm(_a)) print(f"corr={abs(c):.4f} relL2={np.linalg.norm((_f if c>=0 else -_f)-_a)/np.linalg.norm(_a):.4f}") @@ -276,6 +282,16 @@ def collapsed(solver, degree): poisson.boundary_flux("Top") +def test_boundary_flux_scalar_reaction_needs_no_field_decomposition(): + """A scalar reaction scatters its sole field without allocating a sub-DM.""" + poisson = _unit_flux_2d(1, res=4) + + poisson.boundary_flux("Top") + poisson.boundary_flux("Bottom") + + assert not poisson._subdict + + def test_volume_residual_fields_insert_essential_values(): """#411: compute_volume_residual_fields (Stokes-only diagnostic) missed the #407 insert — its residual must now match _assemble_volume_reaction (the diff --git a/tests/test_1057_advdiff_flux_history_skip.py b/tests/test_1057_advdiff_flux_history_skip.py new file mode 100644 index 000000000..8509a2afa --- /dev/null +++ b/tests/test_1057_advdiff_flux_history_skip.py @@ -0,0 +1,86 @@ +"""Regression tests for inactive Backward-Euler flux history.""" + +import numpy as np +import pytest +import sympy + +import underworld3 as uw + + +pytestmark = [pytest.mark.level_1, pytest.mark.tier_a] + + +def _make_advdiff(theta=1.0, order=1): + mesh = uw.meshing.StructuredQuadBox( + elementRes=(4, 4), + minCoords=(0.0, 0.0), + maxCoords=(1.0, 1.0), + ) + temperature = uw.discretisation.MeshVariable( + "T_flux_history", mesh, 1, degree=1 + ) + x, y = mesh.X + velocity = sympy.Matrix([[-(y - 0.5), x - 0.5]]) + + temperature.data[:, 0] = np.asarray( + temperature.coords[:, 0] + 0.25 * temperature.coords[:, 1] + ) + + thermal = uw.systems.AdvDiffusionSLCN( + mesh, + u_Field=temperature, + V_fn=velocity, + order=order, + theta=theta, + ) + thermal.constitutive_model = uw.constitutive_models.DiffusionModel + thermal.constitutive_model.Parameters.diffusivity = 0.01 + return thermal, temperature + + +@pytest.mark.parametrize( + ("theta", "order", "expected_updates"), + ( + (1.0, 1, 0), + (0.5, 1, 1), + (1.0, 2, 1), + ), +) +def test_only_current_flux_skips_history_lifecycle( + theta, order, expected_updates +): + thermal, _ = _make_advdiff(theta=theta, order=order) + calls = {"pre": 0, "post": 0} + original_pre = thermal.DFDt.update_pre_solve + original_post = thermal.DFDt.update_post_solve + + def counted_pre(*args, **kwargs): + calls["pre"] += 1 + return original_pre(*args, **kwargs) + + def counted_post(*args, **kwargs): + calls["post"] += 1 + return original_post(*args, **kwargs) + + thermal.DFDt.update_pre_solve = counted_pre + thermal.DFDt.update_post_solve = counted_post + thermal.solve(timestep=0.01, zero_init_guess=False) + + assert calls == {"pre": expected_updates, "post": expected_updates} + + +def test_skipped_backward_euler_flux_matches_history_path(): + optimized, optimized_temperature = _make_advdiff(theta=1.0, order=1) + reference, reference_temperature = _make_advdiff(theta=1.0, order=1) + + reference._prepare_flux_history = lambda: True + + optimized.solve(timestep=0.01, zero_init_guess=False) + reference.solve(timestep=0.01, zero_init_guess=False) + + np.testing.assert_allclose( + optimized_temperature.data, + reference_temperature.data, + rtol=1.0e-12, + atol=1.0e-12, + ) diff --git a/tests/test_1070_postprocessing_geoid.py b/tests/test_1070_postprocessing_geoid.py index dce646a86..784b9440c 100644 --- a/tests/test_1070_postprocessing_geoid.py +++ b/tests/test_1070_postprocessing_geoid.py @@ -182,6 +182,17 @@ def test_rotated_adapter_requires_explicit_self_gravity_parameters(): ) +def test_rotated_adapter_rejects_unknown_projection(): + with pytest.raises(ValueError, match="projection must be"): + uw.postprocessing.geoid.spherical_shell_response_from_rotated_stokes( + stokes=object(), + radius_inner=0.55, + radius_outer=1.0, + harmonic_degree=2, + projection="nodal", + ) + + def test_rotated_stokes_adapter_matches_zhong_table_2(): radius_inner = 0.55 radius_outer = 1.0 @@ -229,6 +240,21 @@ def test_rotated_stokes_adapter_matches_zhong_table_2(): gravity=9.8, gravitational_constant=6.67e-11, ) + reaction_response = uw.postprocessing.geoid.spherical_shell_response_from_rotated_stokes( + stokes=stokes, + radius_inner=radius_inner, + radius_outer=radius_outer, + harmonic_degree=2, + internal_load_radius=rint, + internal_load_coefficient=1.0, + include_self_gravity=True, + surface_density_contrast=3300.0, + cmb_density_contrast=5400.0, + planet_radius=6370000.0, + gravity=9.8, + gravitational_constant=6.67e-11, + projection="reaction", + ) assert np.isclose(response.surface_topography, 0.41920, rtol=0.10) assert np.isclose(response.cmb_topography, 0.77060, rtol=0.10) @@ -238,3 +264,7 @@ def test_rotated_stokes_adapter_matches_zhong_table_2(): assert np.isclose(response.self_gravity.cmb_topography, 0.93130, rtol=0.10) assert np.isclose(response.self_gravity.surface_geoid, 0.04486, rtol=0.10) assert np.isclose(response.self_gravity.cmb_geoid, 0.05461, rtol=0.10) + assert np.isclose(reaction_response.surface_topography, 0.41920, rtol=0.03) + assert np.isclose(reaction_response.cmb_topography, 0.77060, rtol=0.03) + assert np.isclose(reaction_response.surface_geoid, 0.02579, rtol=0.03) + assert np.isclose(reaction_response.cmb_geoid, 0.03206, rtol=0.03) diff --git a/tests/test_1112_slcn_spherical_lifecycle.py b/tests/test_1112_slcn_spherical_lifecycle.py new file mode 100644 index 000000000..d6aa0c85e --- /dev/null +++ b/tests/test_1112_slcn_spherical_lifecycle.py @@ -0,0 +1,67 @@ +"""Spherical SLCN lifecycle regression for serial and MPI execution.""" + +import gc + +import numpy as np +import pytest + +import underworld3 as uw + + +pytestmark = pytest.mark.level_2 + + +def test_spherical_slcn_transient_state_is_bounded(): + mesh = uw.meshing.SphericalShell( + radiusInner=0.55, + radiusOuter=1.0, + cellSize=0.4, + qdegree=2, + ) + velocity = uw.discretisation.MeshVariable( + "U_slcn_lifecycle", mesh, mesh.dim, degree=1 + ) + temperature = uw.discretisation.MeshVariable( + "T_slcn_lifecycle", mesh, 1, degree=1 + ) + + with mesh.access(velocity, temperature): + coords = np.asarray(temperature.coords) + radii = np.linalg.norm(coords, axis=1) + temperature.data[:, 0] = (1.0 - radii) / 0.45 + velocity.data[:, 0] = -0.05 * coords[:, 1] + velocity.data[:, 1] = 0.05 * coords[:, 0] + velocity.data[:, 2] = 0.0 + + thermal = uw.systems.AdvDiffusionSLCN( + mesh, + u_Field=temperature, + V_fn=velocity.sym, + order=1, + theta=0.5, + ) + 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") + + live_swarms_before = len(mesh._registered_swarms) + history_lengths = [] + + for step in range(38): + thermal.solve(timestep=1.0e-3, zero_init_guess=False) + assert len(mesh._registered_swarms) == live_swarms_before + assert len(mesh._dminterpolation_cache._cache) <= ( + mesh._dminterpolation_cache.max_entries + ) + if step in (31, 37): + gc.collect() + history_lengths.append( + ( + len(thermal.solve_history), + len(mesh._eval_work_1x3_projector.solve_history), + ) + ) + + assert history_lengths == [(32, 32), (32, 32)] + assert np.all(np.isfinite(temperature.data)) diff --git a/tests/test_1113_advdiff_supg_residual.py b/tests/test_1113_advdiff_supg_residual.py new file mode 100644 index 000000000..939c117e8 --- /dev/null +++ b/tests/test_1113_advdiff_supg_residual.py @@ -0,0 +1,243 @@ +"""Focused tests for the implicit SUPG scalar transport residual.""" + +import numpy as np +import pytest + +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) + with mesh.access(temperature, flow): + temperature.data[:, 0] = temperature.coords[:, 0] + flow.data[:, 0] = velocity[0] + flow.data[:, 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.AdvDiffusionSUPG( + mesh, u_Field=temperature, V_fn=velocity.sym, tau=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) + assert thermal.tau == 0.0 + + +def test_rejects_double_counted_eulerian_advection(): + mesh, temperature, velocity = _mesh_temperature_velocity("double") + history = uw.systems.Eulerian_DDt( + mesh, + temperature, + vtype=uw.VarType.SCALAR, + degree=temperature.degree, + continuous=temperature.continuous, + V_fn=velocity.sym, + ) + + with pytest.raises(ValueError, match="DuDt.V_fn must be None"): + uw.systems.AdvDiffusionSUPG( + mesh, + u_Field=temperature, + V_fn=velocity.sym, + DuDt=history, + ) + + +@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.AdvDiffusionSUPG( + mesh, + u_Field=temperature, + V_fn=velocity.sym, + theta=theta, + ) + + +def test_automatic_tau_is_finite_and_bounded_by_transient_scale(): + mesh, temperature, velocity = _mesh_temperature_velocity("tau") + thermal = uw.systems.AdvDiffusionSUPG(mesh, u_Field=temperature, V_fn=velocity.sym) + _configure_diffusion(thermal, diffusivity=0.1) + thermal.delta_t = 0.02 + thermal._update_automatic_tau() + + assert np.all(thermal._supg_h.data > 0.0) + assert np.all(thermal._supg_tau.data > 0.0) + assert np.all(thermal._supg_tau.data <= 0.01) + + +def test_negative_diffusivity_is_rejected(): + mesh, temperature, velocity = _mesh_temperature_velocity("negative_k") + thermal = uw.systems.AdvDiffusionSUPG(mesh, u_Field=temperature, V_fn=velocity.sym) + _configure_diffusion(thermal, diffusivity=-0.1) + thermal.delta_t = 0.01 + + with pytest.raises(ValueError, match="non-negative"): + thermal._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) + ) + with mesh_a.access(temperature_a), mesh_b.access(temperature_b): + temperature_a.data[:, 0] = np.sin(np.pi * temperature_a.coords[:, 0]) + temperature_b.data[:, 0] = np.sin(np.pi * temperature_b.coords[:, 0]) + + supg = uw.systems.AdvDiffusionSUPG(mesh_a, u_Field=temperature_a, V_fn=velocity.sym) + 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) + + supg.solve(timestep=0.01, zero_init_guess=False) + diffusion.solve(timestep=0.01, zero_init_guess=False) + + assert np.all(supg._supg_tau.data == 0.0) + np.testing.assert_allclose( + temperature_a.data, + temperature_b.data, + 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 P1"): + uw.systems.AdvDiffusionSUPG( + mesh, + u_Field=temperature_p2, + V_fn=velocity.sym, + time_integrator="citcoms", + ) + + +def test_citcoms_lumped_mass_matches_constant_residual(): + mesh, temperature, velocity = _mesh_temperature_velocity( + "citcoms_mass", velocity=(0.0, 0.0) + ) + thermal = uw.systems.AdvDiffusionSUPG( + mesh, + u_Field=temperature, + V_fn=velocity.sym, + time_integrator="citcoms", + tau=0.0, + ) + _configure_diffusion(thermal, diffusivity=0.0) + thermal.delta_t = 0.01 + thermal._setup_citcoms_residual() + mass = thermal._assemble_lumped_mass() + thermal._temperature_rate.data[:, 0] = 1.0 + solution, residual = thermal._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.data[:, 0] = 0.0 + thermal = uw.systems.AdvDiffusionSUPG( + mesh, + u_Field=temperature, + V_fn=velocity.sym, + time_integrator="citcoms", + tau=0.0, + ) + _configure_diffusion(thermal, diffusivity=0.0) + thermal.f = 1.0 + + thermal.solve(timestep=0.1) + + np.testing.assert_allclose(temperature.data, 0.1, atol=1.0e-14) + np.testing.assert_allclose(thermal._temperature_rate.data, 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) + ) + thermal = uw.systems.AdvDiffusionSUPG( + mesh, + u_Field=temperature, + V_fn=velocity.sym, + time_integrator="citcoms", + tau=0.0, + ) + _configure_diffusion(thermal, diffusivity=0.0) + + thermal.solve(timestep=0.01) + vector_handles = tuple(vector.handle for vector in thermal._citcoms_work_vectors) + thermal.solve(timestep=0.01) + + assert ( + tuple(vector.handle for vector in thermal._citcoms_work_vectors) + == vector_handles + ) + + +def test_citcoms_timestep_uses_advection_and_lumped_diffusion_limits(): + mesh, temperature, velocity = _mesh_temperature_velocity("citcoms_dt") + thermal = uw.systems.AdvDiffusionSUPG( + mesh, + u_Field=temperature, + V_fn=velocity.sym, + time_integrator="citcoms", + ) + _configure_diffusion(thermal, diffusivity=0.1) + + timestep = thermal.estimate_dt() + + assert np.isfinite(timestep) + assert timestep == pytest.approx(0.9 * min(thermal.dt_adv, thermal.dt_diff)) + assert thermal.dt_adv > 0.0 + assert thermal.dt_diff > 0.0 + + +def test_timestep_diffusivity_branch_is_collective(): + mesh, temperature, velocity = _mesh_temperature_velocity("collective_diffusivity") + thermal = uw.systems.AdvDiffusionSUPG( + mesh, + u_Field=temperature, + V_fn=velocity.sym, + ) + _configure_diffusion(thermal, diffusivity=0.1) + thermal.delta_t = 0.01 + thermal._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.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 000000000..b5a71a6bb --- /dev/null +++ b/tests/test_1114_advdiff_supg.py @@ -0,0 +1,388 @@ +"""Numerical validation for implicit SUPG scalar transport.""" + +import numpy as np +import pytest +import sympy + +import underworld3 as uw + + +pytestmark = pytest.mark.level_2 + + +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 + ) + thermal = uw.systems.AdvDiffusionSUPG( + mesh, + u_Field=temperature, + V_fn=velocity.sym, + time_integrator="citcoms", + ) + thermal.constitutive_model = uw.constitutive_models.DiffusionModel + thermal.constitutive_model.Parameters.diffusivity = 1.0 + + first = thermal._simplex_data() + second = thermal._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._streamline_directional_rate( + first[1], sample_velocity + ) + second_rate = thermal._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._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._streamline_directional_rate( + third[1], sample_velocity + ) is not first_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 + ) + with mesh.access(temperature, velocity): + temperature.data[:, 0] = temperature.coords[:, 0] + velocity.data[:, 0] = 1.0 + velocity.data[:, 1] = 0.0 + + thermal = uw.systems.AdvDiffusionSUPG( + mesh, + u_Field=temperature, + V_fn=velocity.sym, + tau=tau, + ) + 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.data[:, 0] - exact) ** 2))) + return temperature.data.copy(), 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 + with mesh.access(temperature, velocity): + temperature.data[:, 0] = uw.function.evaluate( + exact, temperature.coords + ).reshape(-1) + velocity.data[:, 0] = 1.0 + velocity.data[:, 1] = 0.0 + + thermal = uw.systems.AdvDiffusionSUPG( + mesh, u_Field=temperature, V_fn=velocity.sym + ) + 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 test_spherical_shell_supg_is_parallel_safe(): + 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 + ) + with mesh.access(temperature, velocity): + coords = temperature.coords + radii = np.linalg.norm(coords, axis=1) + temperature.data[:, 0] = (1.0 - radii) / 0.45 + 0.01 * coords[:, 0] + velocity.data[:, 0] = -0.02 * coords[:, 1] + velocity.data[:, 1] = 0.02 * coords[:, 0] + velocity.data[:, 2] = 0.0 + + thermal = uw.systems.AdvDiffusionSUPG( + mesh, + u_Field=temperature, + V_fn=velocity.sym, + ) + 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.data)) + assert np.all(np.isfinite(thermal._supg_tau.data)) + assert temperature_l2_squared == pytest.approx(0.833491030982, rel=1.0e-8) + + +def test_bdf2_snapshot_restore_leaves_no_discarded_step_trace(): + uw.reset_default_model() + 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 + ) + with mesh.access(temperature, velocity): + temperature.data[:, 0] = np.sin(np.pi * temperature.coords[:, 0]) + velocity.data[:, 0] = 0.1 + velocity.data[:, 1] = 0.0 + + thermal = uw.systems.AdvDiffusionSUPG( + mesh, + u_Field=temperature, + V_fn=velocity.sym, + order=2, + ) + 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 = model.save_state() + + model.load_state(snapshot) + for _ in range(3): + thermal.solve(timestep=0.01, zero_init_guess=False) + reference = temperature.data.copy() + + model.load_state(snapshot) + thermal.solve(timestep=0.2, zero_init_guess=False) + model.load_state(snapshot) + for _ in range(3): + thermal.solve(timestep=0.01, zero_init_guess=False) + resumed = temperature.data.copy() + + np.testing.assert_array_equal(resumed, reference) + 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 + ) + with mesh.access(temperature, velocity): + temperature.data[:, 0] = temperature.coords[:, 0] + velocity.data[:, 0] = 0.1 + velocity.data[:, 1] = 0.0 + + thermal = uw.systems.AdvDiffusionSUPG( + mesh, u_Field=temperature, V_fn=velocity.sym + ) + 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.data)) + + +def test_citcoms_spherical_shell_is_parallel_safe(): + 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 + ) + with mesh.access(temperature, velocity): + coords = temperature.coords + radii = np.linalg.norm(coords, axis=1) + temperature.data[:, 0] = (1.0 - radii) / 0.45 + 0.01 * coords[:, 0] + velocity.data[:, 0] = -0.02 * coords[:, 1] + velocity.data[:, 1] = 0.02 * coords[:, 0] + velocity.data[:, 2] = 0.0 + + thermal = uw.systems.AdvDiffusionSUPG( + mesh, + u_Field=temperature, + V_fn=velocity.sym, + time_integrator="citcoms", + ) + 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._lumped_mass.getSize() > 0 + assert np.all(np.isfinite(temperature.data)) + assert temperature_l2_squared == pytest.approx(0.814491155536, rel=1.0e-8) + + +def test_citcoms_snapshot_restores_startup_state_exactly(): + uw.reset_default_model() + 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 + ) + with mesh.access(temperature, velocity): + temperature.data[:, 0] = 1.0 + + thermal = uw.systems.AdvDiffusionSUPG( + mesh, + u_Field=temperature, + V_fn=velocity.sym, + time_integrator="citcoms", + tau=0.0, + ) + thermal.constitutive_model = uw.constitutive_models.DiffusionModel + thermal.constitutive_model.Parameters.diffusivity = 0.0 + thermal.f = -temperature.sym[0] + + initial = model.save_state() + thermal.solve(timestep=0.05) + reference_temperature = temperature.data.copy() + reference_rate = thermal._temperature_rate.data.copy() + + model.load_state(initial) + assert not thermal._rate_initialised + thermal.solve(timestep=0.05) + + np.testing.assert_array_equal(temperature.data, reference_temperature) + np.testing.assert_array_equal( + thermal._temperature_rate.data, reference_rate + ) + uw.reset_default_model() diff --git a/tests/test_1115_advdiff_supg_transient.py b/tests/test_1115_advdiff_supg_transient.py new file mode 100644 index 000000000..aadc04685 --- /dev/null +++ b/tests/test_1115_advdiff_supg_transient.py @@ -0,0 +1,188 @@ +"""Temporal convergence validation for implicit SUPG transport.""" + +import numpy as np +import pytest +import sympy + +import underworld3 as uw + + +pytestmark = pytest.mark.level_3 + + +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 + with mesh.access(temperature, velocity): + temperature.data[:, 0] = uw.function.evaluate( + shape, temperature.coords + ).reshape(-1) + velocity.data[:, 0] = advection_speed + velocity.data[:, 1] = 0.0 + + thermal = uw.systems.AdvDiffusionSUPG( + mesh, + u_Field=temperature, + V_fn=velocity.sym, + order=order, + tau=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.data[:, 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.data[:, 0] = 1.0 + + thermal = uw.systems.AdvDiffusionSUPG( + mesh, + u_Field=temperature, + V_fn=velocity.sym, + time_integrator="citcoms", + tau=0.0, + ) + 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.data[:, 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)) + + with mesh.access(temperature, velocity): + temperature.data[:, 0] = uw.function.evaluate( + initial, temperature.coords + ).reshape(-1) + velocity.data[:, 0] = -2.0 * np.pi * velocity.coords[:, 1] + velocity.data[:, 1] = 2.0 * np.pi * velocity.coords[:, 0] + + thermal = uw.systems.AdvDiffusionSUPG( + mesh, + u_Field=temperature, + V_fn=velocity.sym, + time_integrator="citcoms", + ) + 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