From d0f3fec2cf5b20c8ac6ebff831bb6d01c3fd0c2d Mon Sep 17 00:00:00 2001 From: Tyagi Date: Tue, 25 Aug 2026 20:28:54 +0530 Subject: [PATCH 01/23] Add distributed harmonic projection of rotated reactions Expose boundary_normal_traction_integral() to contract the assembled rotated free-slip normal reaction directly with a scalar boundary test function. Count only owned reaction DOFs, reduce the weak functional across MPI ranks, and remove the constant traction gauge without recovering pointwise P2 values. Add an opt-in reaction projection to spherical-shell geoid postprocessing. Normalize with the matching discrete boundary inner product so the reaction numerator and harmonic norm use the same faceted geometry. Retain the centroid projection as the compatibility default. Document the relationship to curved-boundary P2 midpoint/fitted guidance from issue #414 and add focused serial and MPI Zhong regressions for the new path. --- docs/developer/CHANGELOG.md | 6 ++ ...ry-stress-and-projection-postprocessing.md | 31 ++++++-- .../cython/petsc_generic_snes_solvers.pyx | 19 +++++ src/underworld3/postprocessing/geoid.py | 45 +++++++++-- src/underworld3/utilities/rotated_bc.py | 79 +++++++++++++++++++ ...est_1071_spherical_shell_geoid_parallel.py | 1 + tests/test_1070_postprocessing_geoid.py | 30 +++++++ 7 files changed, 196 insertions(+), 15 deletions(-) diff --git a/docs/developer/CHANGELOG.md b/docs/developer/CHANGELOG.md index 53c0126f..b6ebdbd3 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 fb007f0b..f6283240 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/cython/petsc_generic_snes_solvers.pyx b/src/underworld3/cython/petsc_generic_snes_solvers.pyx index c42bbb47..67715d29 100644 --- a/src/underworld3/cython/petsc_generic_snes_solvers.pyx +++ b/src/underworld3/cython/petsc_generic_snes_solvers.pyx @@ -6500,6 +6500,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/postprocessing/geoid.py b/src/underworld3/postprocessing/geoid.py index 537a79fe..b35d201f 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/utilities/rotated_bc.py b/src/underworld3/utilities/rotated_bc.py index 8400a887..337483bb 100644 --- a/src/underworld3/utilities/rotated_bc.py +++ b/src/underworld3/utilities/rotated_bc.py @@ -2096,6 +2096,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/test_1071_spherical_shell_geoid_parallel.py b/tests/parallel/test_1071_spherical_shell_geoid_parallel.py index 189d6121..ac3c525c 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( [ diff --git a/tests/test_1070_postprocessing_geoid.py b/tests/test_1070_postprocessing_geoid.py index dce646a8..784b9440 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) From 0c0e65f2c35c78073f7194b69965fdfd0712a8a9 Mon Sep 17 00:00:00 2001 From: Tyagi Date: Tue, 25 Aug 2026 20:46:52 +0530 Subject: [PATCH 02/23] Use tolerant MPI geoid rank comparison Collective reductions may differ by the final floating-point bits across MPI implementations. Compare projected geoid and topography coefficients at near-machine precision instead of requiring bitwise-identical arrays. --- tests/parallel/test_1071_spherical_shell_geoid_parallel.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/parallel/test_1071_spherical_shell_geoid_parallel.py b/tests/parallel/test_1071_spherical_shell_geoid_parallel.py index ac3c525c..9ad0b6e1 100644 --- a/tests/parallel/test_1071_spherical_shell_geoid_parallel.py +++ b/tests/parallel/test_1071_spherical_shell_geoid_parallel.py @@ -72,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] From a17586b31532c2d6c93af84f7fecdcb0a0b3eeac Mon Sep 17 00:00:00 2001 From: Tyagi Date: Wed, 26 Aug 2026 14:13:09 +0530 Subject: [PATCH 03/23] Skip inactive Backward-Euler flux history updates For order-one advection-diffusion with theta=1, the Adams-Moulton flux has coefficients [1, 0]. Avoid projecting and tracing the zero-weight stored flux on every solve while still refreshing the symbolic coefficients when theta changes after construction. Preserve the existing DFDt lifecycle for Crank-Nicolson, forward Euler, and higher-order Adams-Moulton configurations. Add focused regression coverage for lifecycle call counts and numerical equivalence with the previous forced-history path. --- src/underworld3/systems/solvers.py | 38 ++++++++- tests/test_1057_advdiff_flux_history_skip.py | 86 ++++++++++++++++++++ 2 files changed, 122 insertions(+), 2 deletions(-) create mode 100644 tests/test_1057_advdiff_flux_history_skip.py diff --git a/src/underworld3/systems/solvers.py b/src/underworld3/systems/solvers.py index 8043975b..98d8c7e9 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/tests/test_1057_advdiff_flux_history_skip.py b/tests/test_1057_advdiff_flux_history_skip.py new file mode 100644 index 00000000..8509a2af --- /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, + ) From 3438b5afadecf4bff19c764d3394aebf5fc748b3 Mon Sep 17 00:00:00 2001 From: Tyagi Date: Wed, 26 Aug 2026 14:28:32 +0530 Subject: [PATCH 04/23] Add spherical SLCN lifecycle regression Exercise 38 canonical Crank-Nicolson transport solves on a three-dimensional spherical shell. Verify in serial and MPI that transient global-evaluation swarms do not survive a solve, interpolation cache entries remain bounded, solver report histories cap at 32, and temperature remains finite. --- tests/test_1112_slcn_spherical_lifecycle.py | 67 +++++++++++++++++++++ 1 file changed, 67 insertions(+) create mode 100644 tests/test_1112_slcn_spherical_lifecycle.py diff --git a/tests/test_1112_slcn_spherical_lifecycle.py b/tests/test_1112_slcn_spherical_lifecycle.py new file mode 100644 index 00000000..d6aa0c85 --- /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)) From e238a255844a9e7422ca585130a027a885cbe082 Mon Sep 17 00:00:00 2001 From: Tyagi Date: Wed, 26 Aug 2026 14:47:06 +0530 Subject: [PATCH 05/23] Add generic implicit SUPG transport solver Introduce uw.systems.AdvDiffusionSUPG as a Python systems-layer solver without changing the existing semi-Lagrangian aliases. Assemble the streamline test contribution through the scalar F1 flux, keep advection out of Eulerian history to avoid double counting, and compute coordinate-invariant simplex streamline lengths in internal P0 fields. Support generic transient and CitcomS-compatible steady tau models for isotropic diffusion plus explicit user tau expressions. Add Level 1 residual and limiting-case tests and Level 2 manufactured, high-Peclet, and spherical MPI validation. --- src/underworld3/systems/__init__.py | 3 + src/underworld3/systems/advdiff_supg.py | 318 +++++++++++++++++++++++ tests/test_1113_advdiff_supg_residual.py | 135 ++++++++++ tests/test_1114_advdiff_supg.py | 160 ++++++++++++ 4 files changed, 616 insertions(+) create mode 100644 src/underworld3/systems/advdiff_supg.py create mode 100644 tests/test_1113_advdiff_supg_residual.py create mode 100644 tests/test_1114_advdiff_supg.py diff --git a/src/underworld3/systems/__init__.py b/src/underworld3/systems/__init__.py index ea6823ab..c7351e61 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 00000000..08ca9e5e --- /dev/null +++ b/src/underworld3/systems/advdiff_supg.py @@ -0,0 +1,318 @@ +r"""Streamline-upwind Petrov-Galerkin scalar transport.""" + +from typing import Optional + +import numpy as np +import sympy + +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 + + +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"}, default="generic" + 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. + 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. This class provides the generic implicit SUPG path; + it is not CitcomS's row-lumped predictor-corrector time integrator. + """ + + @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: str = "generic", + 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 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." + ) + + 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._automatic_tau = tau is None + self._supg_h = None + self._supg_tau = None + + 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 + + def _strong_transport_residual(self): + gradient = self.mesh.vector.gradient(self.u.sym) + advection = sympy.Matrix((self.V_fn.dot(gradient),)) + return self.DuDt.bdf() / self.delta_t + advection - self.f + + @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] + value = expression( + r"\mathbf{F}_1^{SUPG}", + self.DFDt.adams_moulton_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.") + + from underworld3.meshing.smoothing import _tet_cells, _tri_cells + + if self.mesh.dim == 2: + cells = _tri_cells(self.mesh.dm) + elif self.mesh.dim == 3: + cells = _tet_cells(self.mesh.dm) + else: + cells = None + if cells is None or self.mesh.dim != self.mesh.cdim: + raise NotImplementedError( + "Automatic SUPG tau currently requires 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 compute SUPG length 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) + + velocity = _centroid_velocities_nd(self.V_fn, self.mesh) + speed = np.linalg.norm(velocity, axis=1) + directional_rate = np.abs( + np.einsum("cad,cd->ca", gradients, velocity) + ).sum(axis=1) + h_stream = np.divide( + 2.0 * speed, + directional_rate, + out=np.zeros_like(speed), + where=directional_rate > 0.0, + ) + + diffusivity_expr = sympy.sympify(self.constitutive_model.K) + if isinstance(diffusivity_expr, sympy.MatrixBase): + raise NotImplementedError( + "Automatic SUPG tau requires 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_like(speed, diffusivity.item()) + if diffusivity.shape != speed.shape: + raise ValueError("Diffusivity must evaluate to one scalar per simplex cell.") + if np.any(diffusivity < 0.0): + raise ValueError("SUPG diffusivity must be non-negative.") + + 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 + + @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 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/tests/test_1113_advdiff_supg_residual.py b/tests/test_1113_advdiff_supg_residual.py new file mode 100644 index 00000000..e2d3f3a8 --- /dev/null +++ b/tests/test_1113_advdiff_supg_residual.py @@ -0,0 +1,135 @@ +"""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, + ) diff --git a/tests/test_1114_advdiff_supg.py b/tests/test_1114_advdiff_supg.py new file mode 100644 index 00000000..64631e72 --- /dev/null +++ b/tests/test_1114_advdiff_supg.py @@ -0,0 +1,160 @@ +"""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 _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) From 11348b3fe881154595ac5ff721d68dfe915e85ec Mon Sep 17 00:00:00 2001 From: Tyagi Date: Wed, 26 Aug 2026 15:00:30 +0530 Subject: [PATCH 06/23] Make symbolic flux history snapshots restart safe Preserve live UWexpression atoms when deep-copying Symbolic_DDt history state while copying the mutable history containers and timestep metadata. Generic deepcopy reconstructed parameter symbols without their wrapped values, causing Jacobian rebuilds after model restore to fail in existing Diffusion and the new SUPG solver. Add a base diffusion regression, SUPG BDF2 discarded-step equivalence, bounded repeated-solve state checks, and measured BDF1/BDF2 temporal convergence. --- src/underworld3/systems/ddt.py | 24 ++++++- tests/test_0007_snapshot_inmemory.py | 32 ++++++++- tests/test_1114_advdiff_supg.py | 84 ++++++++++++++++++++++ tests/test_1115_advdiff_supg_transient.py | 85 +++++++++++++++++++++++ 4 files changed, 223 insertions(+), 2 deletions(-) create mode 100644 tests/test_1115_advdiff_supg_transient.py diff --git a/src/underworld3/systems/ddt.py b/src/underworld3/systems/ddt.py index e198aad5..e5b93ed1 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/tests/test_0007_snapshot_inmemory.py b/tests/test_0007_snapshot_inmemory.py index d52fee33..09dcc867 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_1114_advdiff_supg.py b/tests/test_1114_advdiff_supg.py index 64631e72..a773b75e 100644 --- a/tests/test_1114_advdiff_supg.py +++ b/tests/test_1114_advdiff_supg.py @@ -158,3 +158,87 @@ def test_spherical_shell_supg_is_parallel_safe(): 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)) diff --git a/tests/test_1115_advdiff_supg_transient.py b/tests/test_1115_advdiff_supg_transient.py new file mode 100644 index 00000000..bca63aff --- /dev/null +++ b/tests/test_1115_advdiff_supg_transient.py @@ -0,0 +1,85 @@ +"""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 From b2c9ca1f1c963330e89cdf76a13c0a1fcdd35f27 Mon Sep 17 00:00:00 2001 From: Tyagi Date: Wed, 26 Aug 2026 15:23:51 +0530 Subject: [PATCH 07/23] Add CitcomS-compatible SUPG predictor corrector Extend AdvDiffusionSUPG with a continuous-P1 CitcomS-compatible time integrator using geometric row-sum lumped mass, initialized temperature-rate state, gamma=0.5 prediction, and two fixed residual corrections. Assemble the SUPG residual through the existing scalar SNES callback, retain constrained boundary handling, snapshot integrator startup metadata, and compute conservative advection and discrete M_L^-1 K diffusion timestep limits. Add constant-residual mass verification, exact source startup, second-order scalar decay, snapshot restart, and serial/MPI spherical tests. --- src/underworld3/systems/advdiff_supg.py | 392 +++++++++++++++++++--- tests/test_1113_advdiff_supg_residual.py | 82 +++++ tests/test_1114_advdiff_supg.py | 86 +++++ tests/test_1115_advdiff_supg_transient.py | 44 +++ 4 files changed, 558 insertions(+), 46 deletions(-) diff --git a/src/underworld3/systems/advdiff_supg.py b/src/underworld3/systems/advdiff_supg.py index 08ca9e5e..6223ce55 100644 --- a/src/underworld3/systems/advdiff_supg.py +++ b/src/underworld3/systems/advdiff_supg.py @@ -1,15 +1,27 @@ 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): @@ -51,11 +63,15 @@ class SNES_AdvectionDiffusionSUPG(SNES_Diffusion): 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"}, default="generic" + 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. DuDt, DFDt : optional Existing history operators. A supplied ``DuDt`` must be Eulerian and must not contain a velocity, because advection is represented in R. @@ -77,7 +93,10 @@ def __init__( order: int = 1, theta: float = 1.0, tau=None, - tau_model: str = "generic", + tau_model: Optional[str] = None, + time_integrator: str = "bdf", + adv_gamma: float = 0.5, + corrector_steps: int = 2, evalf: Optional[bool] = False, verbose: bool = False, DuDt: Optional[Eulerian_DDt] = None, @@ -85,6 +104,10 @@ def __init__( ): 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): @@ -94,6 +117,21 @@ def __init__( "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, @@ -108,9 +146,26 @@ def __init__( 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._rate_initialised = False + + if self.time_integrator == "citcoms": + 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 @@ -143,10 +198,83 @@ def tau(self): """SUPG stabilization parameter used in the residual.""" return self._tau + @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),)) - return self.DuDt.bdf() / self.delta_t + advection - self.f + 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 + + 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) + return cells, gradients, volumes + + 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): @@ -164,9 +292,14 @@ def F0(self): 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}", - self.DFDt.adams_moulton_flux() + self.tau * self.V_fn * residual, + diffusion_flux + self.tau * self.V_fn * residual, "Diffusive and SUPG streamline flux", _unique_name_generation=True, ) @@ -180,30 +313,7 @@ def _update_automatic_tau(self): if self.constitutive_model is None: raise RuntimeError("Set constitutive_model before solving AdvDiffusionSUPG.") - from underworld3.meshing.smoothing import _tet_cells, _tri_cells - - if self.mesh.dim == 2: - cells = _tri_cells(self.mesh.dm) - elif self.mesh.dim == 3: - cells = _tet_cells(self.mesh.dm) - else: - cells = None - if cells is None or self.mesh.dim != self.mesh.cdim: - raise NotImplementedError( - "Automatic SUPG tau currently requires 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 compute SUPG length 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) + _, gradients, _ = self._simplex_data() velocity = _centroid_velocities_nd(self.V_fn, self.mesh) speed = np.linalg.norm(velocity, axis=1) @@ -217,24 +327,7 @@ def _update_automatic_tau(self): where=directional_rate > 0.0, ) - diffusivity_expr = sympy.sympify(self.constitutive_model.K) - if isinstance(diffusivity_expr, sympy.MatrixBase): - raise NotImplementedError( - "Automatic SUPG tau requires 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_like(speed, diffusivity.item()) - if diffusivity.shape != speed.shape: - raise ValueError("Diffusivity must evaluate to one scalar per simplex cell.") - if np.any(diffusivity < 0.0): - raise ValueError("SUPG diffusivity must be non-negative.") + diffusivity = self._cell_diffusivity(speed.size) tau_steady = np.zeros_like(speed) moving = speed > np.finfo(float).eps @@ -294,6 +387,211 @@ def _update_automatic_tau(self): 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.""" + if self._lumped_mass is not None: + return self._lumped_mass + + 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 + return self._lumped_mass + + @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 = np.abs( + np.einsum("cad,cd->ca", gradients, velocity) + ).sum(axis=1) + 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)) + if not np.any(diffusivity > 0.0): + dt_diff = np.inf + elif self.time_integrator != "citcoms": + local_diff_rate = np.max( + 2.0 * self.mesh.dim * diffusivity / np.maximum( + self.mesh._radii**2, np.finfo(float).tiny + ) + ) + diff_rate = uw.mpi.comm.allreduce( + float(local_diff_rate), op=MPI.MAX + ) + dt_diff = 2.0 / diff_rate + else: + self._setup_citcoms_residual() + mass = self._assemble_lumped_mass() + 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.dt_adv = dt_adv + self.dt_diff = dt_diff + return 0.9 * min(dt_adv, dt_diff) + + def _compute_citcoms_residual(self): + """Return the globally assembled residual at the current T and dT/dt.""" + solution = self.dm.createGlobalVector() + solution.set(0.0) + self.dm.localToGlobal(self.u.vec, solution, addv=False) + residual = solution.duplicate() + 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() + + if not self._rate_initialised: + self._temperature_rate.data[:, 0] = 0.0 + temperature_global, residual = self._compute_citcoms_residual() + initial_rate = residual.duplicate() + initial_rate.pointwiseDivide(residual, mass) + initial_rate.scale(-1.0) + self._temperature_rate.vec.set(0.0) + self.dm.globalToLocal(initial_rate, self._temperature_rate.vec) + self.mesh._stale_lvec = True + temperature_global.destroy() + residual.destroy() + initial_rate.destroy() + 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): + temperature_global, residual = self._compute_citcoms_residual() + delta_rate = residual.duplicate() + delta_rate.pointwiseDivide(residual, mass) + delta_rate.scale(-1.0) + + rate_global = temperature_global.duplicate() + 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 + + residual.destroy() + delta_rate.destroy() + rate_global.destroy() + temperature_global.destroy() + + self.is_setup = True + self.constitutive_model._solver_is_setup = True + return + @timing.routine_timer_decorator def solve( self, @@ -305,6 +603,8 @@ def solve( 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() diff --git a/tests/test_1113_advdiff_supg_residual.py b/tests/test_1113_advdiff_supg_residual.py index e2d3f3a8..4cc76b8c 100644 --- a/tests/test_1113_advdiff_supg_residual.py +++ b/tests/test_1113_advdiff_supg_residual.py @@ -133,3 +133,85 @@ def test_zero_velocity_matches_diffusion_solver(): 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_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 diff --git a/tests/test_1114_advdiff_supg.py b/tests/test_1114_advdiff_supg.py index a773b75e..ac31b6fa 100644 --- a/tests/test_1114_advdiff_supg.py +++ b/tests/test_1114_advdiff_supg.py @@ -242,3 +242,89 @@ def test_repeated_solves_keep_histories_and_transient_state_bounded(): 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 index bca63aff..3353dd19 100644 --- a/tests/test_1115_advdiff_supg_transient.py +++ b/tests/test_1115_advdiff_supg_transient.py @@ -83,3 +83,47 @@ def test_bdf_temporal_convergence(order, minimum_rate): 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 From a7e197eaa8402e98d0cc0da46a23fe04ed79ad18 Mon Sep 17 00:00:00 2001 From: Tyagi Date: Wed, 26 Aug 2026 15:29:10 +0530 Subject: [PATCH 08/23] Support stable predictor state checkpoints Allow callers to supply a named continuous-P1 temperature-rate field for production checkpoint workflows and expose it through the solver. Rebuild cached lumped mass after mesh changes and cache the discrete diffusion timestep limit by mesh version and diffusivity so coupled convection steps only recompute the velocity-dependent advective limit. --- src/underworld3/systems/advdiff_supg.py | 62 +++++++++++++++++++++---- 1 file changed, 54 insertions(+), 8 deletions(-) diff --git a/src/underworld3/systems/advdiff_supg.py b/src/underworld3/systems/advdiff_supg.py index 6223ce55..804bdcfb 100644 --- a/src/underworld3/systems/advdiff_supg.py +++ b/src/underworld3/systems/advdiff_supg.py @@ -72,6 +72,9 @@ class SNES_AdvectionDiffusionSUPG(SNES_Diffusion): ``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. @@ -97,6 +100,7 @@ def __init__( 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, @@ -154,16 +158,31 @@ def __init__( self._supg_tau = None self._temperature_rate = None self._lumped_mass = None + self._lumped_mass_mesh_version = None + self._diffusion_dt_cache = None self._rate_initialised = False if self.time_integrator == "citcoms": - self._temperature_rate = uw.discretisation.MeshVariable( - f"_supg_dTdt_{self.instance_number}", - mesh, - 1, - degree=1, - continuous=True, - ) + 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) @@ -198,6 +217,11 @@ 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.""" @@ -397,8 +421,15 @@ def _setup_citcoms_residual(self, verbose=False): def _assemble_lumped_mass(self): """Assemble positive P1 simplex row-sum masses on free global DOFs.""" - if self._lumped_mass is not None: + 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 @@ -430,6 +461,7 @@ def _assemble_lumped_mass(self): 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 @timing.routine_timer_decorator @@ -471,6 +503,19 @@ def estimate_dt(self): else: self._setup_citcoms_residual() mass = self._assemble_lumped_mass() + diffusion_signature = ( + getattr(self.mesh, "_mesh_version", 0), + hash(diffusivity.tobytes()), + ) + if ( + self._diffusion_dt_cache is not None + and self._diffusion_dt_cache[0] == diffusion_signature + ): + 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() @@ -507,6 +552,7 @@ def estimate_dt(self): 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 From a46e37fc0dea3dfe22226261475db12d1be5c25d Mon Sep 17 00:00:00 2001 From: Tyagi Date: Wed, 26 Aug 2026 15:43:02 +0530 Subject: [PATCH 09/23] Reuse CitcomS SUPG predictor-corrector work vectors Allocate the global temperature, residual, correction-rate, and stored-rate PETSc vectors once per mesh version instead of recreating and destroying them for every corrector sweep. Preserve the one-off residual helper ownership contract for diagnostics and tests. Add a focused regression that advances two timesteps and verifies the PETSc workspace handles are reused. Serial SUPG tests pass (14 tests), the focused spherical and snapshot tests pass on eight MPI ranks, and the isolated 120-step MPI memory diagnostic has no persistent object or RSS growth. --- src/underworld3/systems/advdiff_supg.py | 57 +++++++++++++++--------- tests/test_1113_advdiff_supg_residual.py | 22 +++++++++ 2 files changed, 59 insertions(+), 20 deletions(-) diff --git a/src/underworld3/systems/advdiff_supg.py b/src/underworld3/systems/advdiff_supg.py index 804bdcfb..8202edcc 100644 --- a/src/underworld3/systems/advdiff_supg.py +++ b/src/underworld3/systems/advdiff_supg.py @@ -159,6 +159,8 @@ def __init__( 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._diffusion_dt_cache = None self._rate_initialised = False @@ -464,6 +466,27 @@ def _assemble_lumped_mass(self): 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. @@ -558,12 +581,14 @@ def estimate_dt(self): self.dt_diff = dt_diff return 0.9 * min(dt_adv, dt_diff) - def _compute_citcoms_residual(self): - """Return the globally assembled residual at the current T and dT/dt.""" - solution = self.dm.createGlobalVector() + 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 = solution.duplicate() residual.set(0.0) self.mesh.update_lvec() self.dm.setAuxiliaryVec(self.mesh.lvec, None) @@ -583,19 +608,18 @@ def _solve_citcoms(self, timestep, verbose=False): 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 - temperature_global, residual = self._compute_citcoms_residual() - initial_rate = residual.duplicate() - initial_rate.pointwiseDivide(residual, mass) - initial_rate.scale(-1.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(initial_rate, self._temperature_rate.vec) + self.dm.globalToLocal(delta_rate, self._temperature_rate.vec) self.mesh._stale_lvec = True - temperature_global.destroy() - residual.destroy() - initial_rate.destroy() self._rate_initialised = True self.u.data[:, 0] += ( @@ -609,12 +633,10 @@ def _solve_citcoms(self, timestep, verbose=False): ) for _ in range(self.corrector_steps): - temperature_global, residual = self._compute_citcoms_residual() - delta_rate = residual.duplicate() + self._compute_citcoms_residual(temperature_global, residual) delta_rate.pointwiseDivide(residual, mass) delta_rate.scale(-1.0) - rate_global = temperature_global.duplicate() rate_global.set(0.0) self.dm.localToGlobal( self._temperature_rate.vec, rate_global, addv=False @@ -629,11 +651,6 @@ def _solve_citcoms(self, timestep, verbose=False): petsc_dm_insert_boundary_values(self.dm, self.u.vec) self.mesh._stale_lvec = True - residual.destroy() - delta_rate.destroy() - rate_global.destroy() - temperature_global.destroy() - self.is_setup = True self.constitutive_model._solver_is_setup = True return diff --git a/tests/test_1113_advdiff_supg_residual.py b/tests/test_1113_advdiff_supg_residual.py index 4cc76b8c..59703e56 100644 --- a/tests/test_1113_advdiff_supg_residual.py +++ b/tests/test_1113_advdiff_supg_residual.py @@ -197,6 +197,28 @@ def test_citcoms_constant_source_is_exact_from_first_step(): ) +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( From 81d359bf4c64e537ae676001e90f98ee006b738b Mon Sep 17 00:00:00 2001 From: Tyagi Date: Wed, 26 Aug 2026 15:57:09 +0530 Subject: [PATCH 10/23] Make SUPG timestep limits MPI collective Base the zero-diffusivity path and diffusion-cache return on communicator-wide decisions so ranks with empty or locally zero-diffusivity partitions cannot skip collectives reached by their peers. Handle empty local diffusivity arrays without a local maximum. Add a rank-varying diffusivity regression that would deadlock before this fix. The collective guard scan passes, the regression passes in serial and on eight MPI ranks, and focused SUPG tests remain green. --- src/underworld3/systems/advdiff_supg.py | 84 ++++++++++++------------ tests/test_1113_advdiff_supg_residual.py | 58 ++++++++-------- 2 files changed, 72 insertions(+), 70 deletions(-) diff --git a/src/underworld3/systems/advdiff_supg.py b/src/underworld3/systems/advdiff_supg.py index 8202edcc..359aea7a 100644 --- a/src/underworld3/systems/advdiff_supg.py +++ b/src/underworld3/systems/advdiff_supg.py @@ -209,9 +209,7 @@ def V_fn(self): def V_fn(self, value): self.is_setup = False self._V_fn = ( - value.sym - if isinstance(value, uw.discretisation.MeshVariable) - else value + value.sym if isinstance(value, uw.discretisation.MeshVariable) else value ) @property @@ -256,14 +254,11 @@ def _simplex_data(self): cells = ( _tri_cells(self.mesh.dm) if self.mesh.dim == 2 - else _tet_cells(self.mesh.dm) - if self.mesh.dim == 3 - else None + 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." + "Automatic SUPG operations require a 2-D or 3-D volume " "simplex mesh." ) coords = np.asarray(self.mesh.X.coords) @@ -337,15 +332,17 @@ def _update_automatic_tau(self): if not self._automatic_tau: return if self.constitutive_model is None: - raise RuntimeError("Set constitutive_model before solving AdvDiffusionSUPG.") + 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 = np.abs( - np.einsum("cad,cd->ca", gradients, velocity) - ).sum(axis=1) + directional_rate = np.abs(np.einsum("cad,cd->ca", gradients, velocity)).sum( + axis=1 + ) h_stream = np.divide( 2.0 * speed, directional_rate, @@ -361,19 +358,11 @@ def _update_automatic_tau(self): nondiffusive = moving & ~diffusive if np.any(diffusive): - pe = ( - speed[diffusive] - * h_stream[diffusive] - / (2.0 * diffusivity[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] = 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] = ( @@ -501,9 +490,9 @@ def estimate_dt(self): cells, gradients, volumes = self._simplex_data() velocity = _centroid_velocities_nd(self.V_fn, self.mesh) - directional_rate = np.abs( - np.einsum("cad,cd->ca", gradients, velocity) - ).sum(axis=1) + directional_rate = np.abs(np.einsum("cad,cd->ca", gradients, velocity)).sum( + axis=1 + ) local_adv_rate = ( float(np.max(directional_rate)) if directional_rate.size else 0.0 ) @@ -511,17 +500,28 @@ def estimate_dt(self): dt_adv = 1.0 / adv_rate if adv_rate > 0.0 else np.inf diffusivity = self._cell_diffusivity(len(cells)) - if not np.any(diffusivity > 0.0): + 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 = np.max( - 2.0 * self.mesh.dim * diffusivity / np.maximum( - self.mesh._radii**2, np.finfo(float).tiny + 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( - float(local_diff_rate), op=MPI.MAX - ) + diff_rate = uw.mpi.comm.allreduce(local_diff_rate, op=MPI.MAX) dt_diff = 2.0 / diff_rate else: self._setup_citcoms_residual() @@ -530,10 +530,14 @@ def estimate_dt(self): getattr(self.mesh, "_mesh_version", 0), hash(diffusivity.tobytes()), ) - if ( + 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 @@ -546,9 +550,7 @@ def estimate_dt(self): 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] - ] + 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] @@ -608,9 +610,7 @@ def _solve_citcoms(self, timestep, verbose=False): self._update_automatic_tau() self._setup_citcoms_residual(verbose) mass = self._assemble_lumped_mass() - temperature_global, residual, delta_rate, rate_global = ( - self._citcoms_vectors() - ) + temperature_global, residual, delta_rate, rate_global = self._citcoms_vectors() if not self._rate_initialised: self._temperature_rate.data[:, 0] = 0.0 @@ -638,9 +638,7 @@ def _solve_citcoms(self, timestep, verbose=False): delta_rate.scale(-1.0) rate_global.set(0.0) - self.dm.localToGlobal( - self._temperature_rate.vec, rate_global, addv=False - ) + 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) diff --git a/tests/test_1113_advdiff_supg_residual.py b/tests/test_1113_advdiff_supg_residual.py index 59703e56..939c117e 100644 --- a/tests/test_1113_advdiff_supg_residual.py +++ b/tests/test_1113_advdiff_supg_residual.py @@ -16,12 +16,8 @@ def _mesh_temperature_velocity(prefix, velocity=(1.0, 0.0)): cellSize=0.3, regular=False, ) - temperature = uw.discretisation.MeshVariable( - f"T_{prefix}", mesh, 1, degree=1 - ) - flow = uw.discretisation.MeshVariable( - f"U_{prefix}", mesh, mesh.dim, degree=1 - ) + temperature = 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] @@ -81,9 +77,7 @@ def test_rejects_nonimplicit_flux_history(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 - ) + 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() @@ -95,9 +89,7 @@ def test_automatic_tau_is_finite_and_bounded_by_transient_scale(): 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 - ) + thermal = uw.systems.AdvDiffusionSUPG(mesh, u_Field=temperature, V_fn=velocity.sym) _configure_diffusion(thermal, diffusivity=-0.1) thermal.delta_t = 0.01 @@ -116,9 +108,7 @@ def test_zero_velocity_matches_diffusion_solver(): 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 - ) + 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) @@ -137,9 +127,7 @@ def test_zero_velocity_matches_diffusion_solver(): 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 - ) + temperature_p2 = uw.discretisation.MeshVariable("T_citcoms_p2", mesh, 1, degree=2) with pytest.raises(ValueError, match="continuous P1"): uw.systems.AdvDiffusionSUPG( @@ -192,9 +180,7 @@ def test_citcoms_constant_source_is_exact_from_first_step(): 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 - ) + np.testing.assert_allclose(thermal._temperature_rate.data, 1.0, atol=1.0e-14) def test_citcoms_reuses_predictor_corrector_work_vectors(): @@ -214,9 +200,10 @@ def test_citcoms_reuses_predictor_corrector_work_vectors(): 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 + assert ( + tuple(vector.handle for vector in thermal._citcoms_work_vectors) + == vector_handles + ) def test_citcoms_timestep_uses_advection_and_lumped_diffusion_limits(): @@ -232,8 +219,25 @@ def test_citcoms_timestep_uses_advection_and_lumped_diffusion_limits(): timestep = thermal.estimate_dt() assert np.isfinite(timestep) - assert timestep == pytest.approx( - 0.9 * min(thermal.dt_adv, thermal.dt_diff) - ) + 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 From a77d595c9730b36d0b89c5837d23c33a078b4f4e Mon Sep 17 00:00:00 2001 From: Tyagi Date: Wed, 26 Aug 2026 17:09:23 +0530 Subject: [PATCH 11/23] Document SUPG scalar transport workflows Add a runnable SUPG guide covering the implicit BDF and CitcomS-compatible predictor-corrector paths, stabilization controls, restart state, and method-selection tradeoffs. Cross-link the guide from the advanced documentation and SLCN time-integration page, and update the solver docstring to describe both supported integrators. --- docs/advanced/index.md | 9 +- .../semi-lagrangian-time-integration.md | 3 + docs/advanced/supg-transport.md | 127 ++++++++++++++++++ src/underworld3/systems/advdiff_supg.py | 5 +- 4 files changed, 141 insertions(+), 3 deletions(-) create mode 100644 docs/advanced/supg-transport.md diff --git a/docs/advanced/index.md b/docs/advanced/index.md index cb47f8de..d022a29e 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 23a935db..9a592df2 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 00000000..f48fec9b --- /dev/null +++ b/docs/advanced/supg-transport.md @@ -0,0 +1,127 @@ +--- +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. + +## 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/src/underworld3/systems/advdiff_supg.py b/src/underworld3/systems/advdiff_supg.py index 359aea7a..6590a423 100644 --- a/src/underworld3/systems/advdiff_supg.py +++ b/src/underworld3/systems/advdiff_supg.py @@ -83,8 +83,9 @@ class SNES_AdvectionDiffusionSUPG(SNES_Diffusion): ----- Automatic tau currently supports volume simplex meshes and scalar isotropic diffusivity. Supply ``tau`` explicitly for other meshes or - constitutive models. This class provides the generic implicit SUPG path; - it is not CitcomS's row-lumped predictor-corrector time integrator. + 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 From cff1dcc81505598e47312f34699ff3a70c208db1 Mon Sep 17 00:00:00 2001 From: Tyagi Date: Wed, 26 Aug 2026 17:10:18 +0530 Subject: [PATCH 12/23] Clarify transport conservation diagnostics Distinguish residual consistency, nodal boundedness, and integral conservation for SUPG and semi-Lagrangian transport, and state the partition-independent diagnostics required for serial/MPI comparison. --- docs/advanced/supg-transport.md | 16 +++++++++++++++- 1 file changed, 15 insertions(+), 1 deletion(-) diff --git a/docs/advanced/supg-transport.md b/docs/advanced/supg-transport.md index f48fec9b..c79163d4 100644 --- a/docs/advanced/supg-transport.md +++ b/docs/advanced/supg-transport.md @@ -108,6 +108,21 @@ 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. @@ -124,4 +139,3 @@ problem; the solver name alone does not establish accuracy. - [Semi-Lagrangian time integration](semi-lagrangian-time-integration.md) - [State snapshots and restore](snapshot-restore.md) - [Parallel computing](parallel-computing.md) - From 1b6990f19b1600a6caa35857d6709299e0b131ea Mon Sep 17 00:00:00 2001 From: Tyagi Date: Wed, 26 Aug 2026 17:15:09 +0530 Subject: [PATCH 13/23] Add SUPG curved-streamline return regression Advect a Gaussian through one rigid-body revolution in an annulus with the public CitcomS-compatible predictor-corrector and estimate_dt(). Verify that the finite-element L2 return error decreases under mesh refinement, covering curved streamlines and the zero-diffusivity timestep path. --- tests/test_1115_advdiff_supg_transient.py | 59 +++++++++++++++++++++++ 1 file changed, 59 insertions(+) diff --git a/tests/test_1115_advdiff_supg_transient.py b/tests/test_1115_advdiff_supg_transient.py index 3353dd19..aadc0468 100644 --- a/tests/test_1115_advdiff_supg_transient.py +++ b/tests/test_1115_advdiff_supg_transient.py @@ -127,3 +127,62 @@ def test_citcoms_predictor_corrector_is_second_order_for_scalar_decay(): 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 From 7119cd144242a146a3730d17b12a46e42f3cb74b Mon Sep 17 00:00:00 2001 From: Tyagi Date: Wed, 26 Aug 2026 18:20:08 +0530 Subject: [PATCH 14/23] Fix exact MPI disk snapshot field reload Write an owned global-vector payload alongside existing DMPlex checkpoint metadata so same-layout restarts do not depend on invalid local point-number assumptions after parallel redistribution. Reload that payload through the live global layout, scatter it to local and ghost dofs, and invalidate the mesh-wide packed auxiliary vector before the next residual assembly. Reject snapshot restarts on a different MPI rank count, retain the legacy DMPlex local-vector fallback for older files, and document the distinction from coordinate-remapped timestep reads. Strengthen snapshot regressions to verify coordinate-defined pointwise field values under four MPI ranks, scalar and P2 vector payload presence, post-reload coefficient use in a solve, and rank-count validation. Focused serial tests, the four-rank snapshot test, all three eight-rank Zhong transport replay gates, and the full Level 1 suite pass. --- src/underworld3/checkpoint/disk_snapshot.py | 9 ++ .../discretisation/discretisation_mesh.py | 26 ++++- .../discretisation_mesh_variables.py | 99 ++++++++++++------- tests/parallel/ptest_0010_snapshot_disk.py | 38 +++---- tests/test_0010_snapshot_disk_format.py | 72 +++++++++++++- 5 files changed, 178 insertions(+), 66 deletions(-) diff --git a/src/underworld3/checkpoint/disk_snapshot.py b/src/underworld3/checkpoint/disk_snapshot.py index 7d30d200..102c9ba2 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/discretisation/discretisation_mesh.py b/src/underworld3/discretisation/discretisation_mesh.py index 769dd6fa..f18e002b 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 9bbcabbc..beaa4587 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/tests/parallel/ptest_0010_snapshot_disk.py b/tests/parallel/ptest_0010_snapshot_disk.py index b9682313..b4f53dfe 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/test_0010_snapshot_disk_format.py b/tests/test_0010_snapshot_disk_format.py index 56609dfe..46fdf89b 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.""" From 5763497da2a4ff1ce25ff0cc958f06242a9e19d9 Mon Sep 17 00:00:00 2001 From: Tyagi Date: Wed, 26 Aug 2026 18:44:48 +0530 Subject: [PATCH 15/23] Instrument MPI evaluation fallback allocation Add disabled-by-default internal counters for global_evaluate's best-claim fallback. When explicitly enabled by a diagnostic, report call counts, local and globally replicated extrapolated points, cumulative temporary-array bytes per rank, and per-call peaks without changing the public evaluator API or production behavior. Add a two-rank regression with exterior query points that verifies finite results and exact point/replica accounting. The focused MPI test and seven evaluator tests pass; the Level 1 suite passes with the known divergent-rank fixture excluded (1659 passed, 35 skipped, 2 xfailed). --- src/underworld3/function/_function.pyx | 65 +++++++++++++++++++ .../test_0760_swarm_cache_migration.py | 32 +++++++++ 2 files changed, 97 insertions(+) diff --git a/src/underworld3/function/_function.pyx b/src/underworld3/function/_function.pyx index 6c785af0..d65425a5 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. @@ -595,7 +622,21 @@ def global_evaluate_nd( expr, 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 +673,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/tests/parallel/test_0760_swarm_cache_migration.py b/tests/parallel/test_0760_swarm_cache_migration.py index d27475af..fbe75612 100644 --- a/tests/parallel/test_0760_swarm_cache_migration.py +++ b/tests/parallel/test_0760_swarm_cache_migration.py @@ -101,3 +101,35 @@ def test_global_evaluate_displaced_nodes(): f"Rank {uw.mpi.rank}: expected {node_coords.shape[0]} results, " f"got {result.shape[0]}" ) + + +@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 From 0841085decb32e93fd72abdf6e8c7eb862804d7a Mon Sep 17 00:00:00 2001 From: Tyagi Date: Tue, 28 Jul 2026 00:05:13 +0530 Subject: [PATCH 16/23] Recover non-finite parallel point evaluations Include located-but-nonfinite interpolation results in global_evaluate's parallel best-claim fallback. This prevents finite SLCN midpoint coordinates from receiving NaN velocities when a rank-local interpolation reports a false located status.\n\nAdd a focused regression for fallback index selection. Validate with the two-rank migration suite and an eight-rank Zhong A1 cellsize=1/16 step that previously diverged with DIVERGED_FNORM_NAN. --- src/underworld3/function/_function.pyx | 17 ++++++++++++++++- .../parallel/test_0760_swarm_cache_migration.py | 16 ++++++++++++++++ 2 files changed, 32 insertions(+), 1 deletion(-) diff --git a/src/underworld3/function/_function.pyx b/src/underworld3/function/_function.pyx index d65425a5..4822863c 100644 --- a/src/underworld3/function/_function.pyx +++ b/src/underworld3/function/_function.pyx @@ -374,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, @@ -616,7 +626,12 @@ 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) diff --git a/tests/parallel/test_0760_swarm_cache_migration.py b/tests/parallel/test_0760_swarm_cache_migration.py index fbe75612..10458a68 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)] @@ -103,6 +104,21 @@ def test_global_evaluate_displaced_nodes(): ) +@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 From c77b23013e0d424b8c85fca63f96daec17a54afa Mon Sep 17 00:00:00 2001 From: Tyagi Date: Wed, 26 Aug 2026 22:13:02 +0530 Subject: [PATCH 17/23] Reuse scalar reaction field decomposition Cache the scalar solver field decomposition on first volume-reaction recovery instead of creating a new PETSc IS and sub-DM for every boundary-flux call. The existing solver reset lifecycle now owns and destroys these objects. Add a focused regression proving repeated upper/lower boundary recovery retains the same decomposition objects. Validate the complete serial boundary-flux suite and the two-rank parallel recovery suite. --- .../cython/petsc_generic_snes_solvers.pyx | 18 ++++++++++++++---- tests/test_1019_boundary_flux.py | 14 ++++++++++++++ 2 files changed, 28 insertions(+), 4 deletions(-) diff --git a/src/underworld3/cython/petsc_generic_snes_solvers.pyx b/src/underworld3/cython/petsc_generic_snes_solvers.pyx index dadb3f84..fa0453db 100644 --- a/src/underworld3/cython/petsc_generic_snes_solvers.pyx +++ b/src/underworld3/cython/petsc_generic_snes_solvers.pyx @@ -3125,10 +3125,20 @@ class SolverBaseClass(uw_object): self._subdict[name][1].localToGlobal(var.vec, sgvec) gvec.restoreSubVector(self._subdict[name][0], sgvec) else: - _names, _iss, _subdms = self.dm.createFieldDecomposition() - sgvec = gvec.getSubVector(_iss[0]) - _subdms[0].localToGlobal(self.Unknowns.u.vec, sgvec) - gvec.restoreSubVector(_iss[0], sgvec) + if not self._subdict: + _names, _iss, _subdms = self.dm.createFieldDecomposition() + self._subdict = { + name: (_iss[index], _subdms[index]) + for index, name in enumerate(_names) + } + if len(self._subdict) != 1: + raise RuntimeError( + "Scalar volume reaction requires one cached field decomposition." + ) + _field_is, _field_dm = next(iter(self._subdict.values())) + sgvec = gvec.getSubVector(_field_is) + _field_dm.localToGlobal(self.Unknowns.u.vec, sgvec) + gvec.restoreSubVector(_field_is, sgvec) self.dm.globalToLocal(gvec, xlocal) diff --git a/tests/test_1019_boundary_flux.py b/tests/test_1019_boundary_flux.py index 86e67faa..82727b5f 100644 --- a/tests/test_1019_boundary_flux.py +++ b/tests/test_1019_boundary_flux.py @@ -276,6 +276,20 @@ def collapsed(solver, degree): poisson.boundary_flux("Top") +def test_boundary_flux_reuses_scalar_field_decomposition(): + """Repeated reaction recovery must retain one solver-owned decomposition.""" + poisson = _unit_flux_2d(1, res=4) + + poisson.boundary_flux("Top") + first = tuple(poisson._subdict.values()) + poisson.boundary_flux("Bottom") + second = tuple(poisson._subdict.values()) + + assert len(first) == 1 + assert first[0][0] is second[0][0] + assert first[0][1] is second[0][1] + + 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 From 419677592491a9d14561d3aa7cf98cae59a024b7 Mon Sep 17 00:00:00 2001 From: Tyagi Date: Wed, 26 Aug 2026 22:13:10 +0530 Subject: [PATCH 18/23] Reuse static simplex geometry in SUPG operations Cache local simplex connectivity, basis gradients, and volumes by mesh version so automatic stabilization and timestep estimation do not rebuild large geometry arrays every step. Mesh deformation or adaptation invalidates the cache through the existing mesh-version lifecycle. Add a focused identity regression for repeated automatic operations. The change removes repeated allocator high-water growth from the coupled A1 timestep-estimation stage without changing the CFL or diffusion limits. --- src/underworld3/systems/advdiff_supg.py | 13 +++++++++++- tests/test_1114_advdiff_supg.py | 28 +++++++++++++++++++++++++ 2 files changed, 40 insertions(+), 1 deletion(-) diff --git a/src/underworld3/systems/advdiff_supg.py b/src/underworld3/systems/advdiff_supg.py index 6590a423..6a69a023 100644 --- a/src/underworld3/systems/advdiff_supg.py +++ b/src/underworld3/systems/advdiff_supg.py @@ -162,6 +162,8 @@ def __init__( 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._diffusion_dt_cache = None self._rate_initialised = False @@ -252,6 +254,13 @@ 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 @@ -274,7 +283,9 @@ def _simplex_data(self): 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) - return cells, gradients, volumes + self._simplex_data_cache = (cells, gradients, volumes) + self._simplex_data_mesh_version = mesh_version + return self._simplex_data_cache def _cell_diffusivity(self, cell_count): """Evaluate non-negative scalar diffusivity at cell centroids.""" diff --git a/tests/test_1114_advdiff_supg.py b/tests/test_1114_advdiff_supg.py index ac31b6fa..8280858c 100644 --- a/tests/test_1114_advdiff_supg.py +++ b/tests/test_1114_advdiff_supg.py @@ -10,6 +10,34 @@ 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)) + + def _high_peclet_solution(tau, name): mesh = uw.meshing.UnstructuredSimplexBox( minCoords=(0.0, 0.0), From 6f3f8c62b223081ce8bbad168d80199ad12c0a5c Mon Sep 17 00:00:00 2001 From: Tyagi Date: Wed, 26 Aug 2026 22:13:58 +0530 Subject: [PATCH 19/23] Test SUPG geometry cache invalidation Extend the simplex-geometry reuse regression through a public mesh deformation. Confirm the cache retains array identities on an unchanged mesh and rebuilds connectivity-derived arrays and volumes after the mesh version advances. --- tests/test_1114_advdiff_supg.py | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/tests/test_1114_advdiff_supg.py b/tests/test_1114_advdiff_supg.py index 8280858c..44cc6114 100644 --- a/tests/test_1114_advdiff_supg.py +++ b/tests/test_1114_advdiff_supg.py @@ -37,6 +37,14 @@ def test_simplex_geometry_is_reused_between_automatic_operations(): assert all(a is b for a, b in zip(first, second)) + 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()) + def _high_peclet_solution(tau, name): mesh = uw.meshing.UnstructuredSimplexBox( From af58a2e40b2cea5415b2e80d617f2f1b18d2a501 Mon Sep 17 00:00:00 2001 From: Tyagi Date: Wed, 26 Aug 2026 22:48:07 +0530 Subject: [PATCH 20/23] Reuse SUPG directional-rate work arrays Replace the per-call cells-by-basis temporary used by automatic SUPG stabilization and timestep estimation with two mesh-versioned one-dimensional work arrays. Compute each basis-direction contribution in place and accumulate it without changing the streamline-rate formula. Validate numerical equivalence against the vectorized expression, workspace reuse, deformation invalidation, and a five-step eight-rank A1 run with unchanged diagnostics. This targets the remaining native allocator growth observed in the Gadi timestep-estimation stage. --- src/underworld3/systems/advdiff_supg.py | 38 +++++++++++++++++++++---- tests/test_1114_advdiff_supg.py | 22 ++++++++++++++ 2 files changed, 54 insertions(+), 6 deletions(-) diff --git a/src/underworld3/systems/advdiff_supg.py b/src/underworld3/systems/advdiff_supg.py index 6a69a023..0e3cd5f6 100644 --- a/src/underworld3/systems/advdiff_supg.py +++ b/src/underworld3/systems/advdiff_supg.py @@ -164,6 +164,8 @@ def __init__( 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 @@ -287,6 +289,34 @@ def _simplex_data(self): 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) @@ -352,9 +382,7 @@ def _update_automatic_tau(self): velocity = _centroid_velocities_nd(self.V_fn, self.mesh) speed = np.linalg.norm(velocity, axis=1) - directional_rate = np.abs(np.einsum("cad,cd->ca", gradients, velocity)).sum( - axis=1 - ) + directional_rate = self._streamline_directional_rate(gradients, velocity) h_stream = np.divide( 2.0 * speed, directional_rate, @@ -502,9 +530,7 @@ def estimate_dt(self): cells, gradients, volumes = self._simplex_data() velocity = _centroid_velocities_nd(self.V_fn, self.mesh) - directional_rate = np.abs(np.einsum("cad,cd->ca", gradients, velocity)).sum( - axis=1 - ) + directional_rate = self._streamline_directional_rate(gradients, velocity) local_adv_rate = ( float(np.max(directional_rate)) if directional_rate.size else 0.0 ) diff --git a/tests/test_1114_advdiff_supg.py b/tests/test_1114_advdiff_supg.py index 44cc6114..b5a71a6b 100644 --- a/tests/test_1114_advdiff_supg.py +++ b/tests/test_1114_advdiff_supg.py @@ -37,6 +37,25 @@ def test_simplex_geometry_is_reused_between_automatic_operations(): 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) @@ -44,6 +63,9 @@ def test_simplex_geometry_is_reused_between_automatic_operations(): 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): From 380ee7bd1ca4895190ee3a225a10b2a887de1c5b Mon Sep 17 00:00:00 2001 From: Tyagi Date: Wed, 26 Aug 2026 22:51:35 +0530 Subject: [PATCH 21/23] Avoid scalar reaction field decomposition Scatter a scalar solver's sole local unknown directly through its solver DM when assembling volume reactions. This removes the need to construct or retain PETSc field-decomposition IS and sub-DM objects for boundary heat-flux recovery. Update the lifecycle regression to require an empty scalar decomposition cache. Validate all 15 serial boundary-flux tests, the two-rank parallel suite, and a five-step eight-rank A1 run with unchanged thermal diagnostics and a 4.28 MiB aggregate repeat-diagnostic delta. --- .../cython/petsc_generic_snes_solvers.pyx | 15 +-------------- tests/test_1019_boundary_flux.py | 10 +++------- 2 files changed, 4 insertions(+), 21 deletions(-) diff --git a/src/underworld3/cython/petsc_generic_snes_solvers.pyx b/src/underworld3/cython/petsc_generic_snes_solvers.pyx index fa0453db..58fe0b3d 100644 --- a/src/underworld3/cython/petsc_generic_snes_solvers.pyx +++ b/src/underworld3/cython/petsc_generic_snes_solvers.pyx @@ -3125,20 +3125,7 @@ class SolverBaseClass(uw_object): self._subdict[name][1].localToGlobal(var.vec, sgvec) gvec.restoreSubVector(self._subdict[name][0], sgvec) else: - if not self._subdict: - _names, _iss, _subdms = self.dm.createFieldDecomposition() - self._subdict = { - name: (_iss[index], _subdms[index]) - for index, name in enumerate(_names) - } - if len(self._subdict) != 1: - raise RuntimeError( - "Scalar volume reaction requires one cached field decomposition." - ) - _field_is, _field_dm = next(iter(self._subdict.values())) - sgvec = gvec.getSubVector(_field_is) - _field_dm.localToGlobal(self.Unknowns.u.vec, sgvec) - gvec.restoreSubVector(_field_is, sgvec) + self.dm.localToGlobal(self.Unknowns.u.vec, gvec) self.dm.globalToLocal(gvec, xlocal) diff --git a/tests/test_1019_boundary_flux.py b/tests/test_1019_boundary_flux.py index 82727b5f..3103da57 100644 --- a/tests/test_1019_boundary_flux.py +++ b/tests/test_1019_boundary_flux.py @@ -276,18 +276,14 @@ def collapsed(solver, degree): poisson.boundary_flux("Top") -def test_boundary_flux_reuses_scalar_field_decomposition(): - """Repeated reaction recovery must retain one solver-owned decomposition.""" +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") - first = tuple(poisson._subdict.values()) poisson.boundary_flux("Bottom") - second = tuple(poisson._subdict.values()) - assert len(first) == 1 - assert first[0][0] is second[0][0] - assert first[0][1] is second[0][1] + assert not poisson._subdict def test_volume_residual_fields_insert_essential_values(): From 6177be86345f26054ff48a6374aa2bef876dff68 Mon Sep 17 00:00:00 2001 From: Tyagi Date: Thu, 27 Aug 2026 07:53:46 +0530 Subject: [PATCH 22/23] Add direct scalar boundary flux integral Expose SolverBaseClass.boundary_flux_integral() for integral diagnostics such as Nusselt numbers. The implementation sums consistent scalar nodal reactions collectively, using the boundary basis partition of unity, and avoids pointwise mass recovery, a temporary MeshVariable, and a second boundary quadrature. Add serial and MPI regressions against the analytic manufactured heat flux, the established recovered-field result, and the serial direct-integral reference. Validate partitions that cut the measured boundary on two and four ranks. --- .../cython/petsc_generic_snes_solvers.pyx | 12 +++++++ src/underworld3/utilities/boundary_flux.py | 31 +++++++++++++++++++ .../test_1065_boundary_flux_parallel.py | 18 +++++++++-- tests/test_1019_boundary_flux.py | 12 +++++-- 4 files changed, 67 insertions(+), 6 deletions(-) diff --git a/src/underworld3/cython/petsc_generic_snes_solvers.pyx b/src/underworld3/cython/petsc_generic_snes_solvers.pyx index 58fe0b3d..5f7cde8c 100644 --- a/src/underworld3/cython/petsc_generic_snes_solvers.pyx +++ b/src/underworld3/cython/petsc_generic_snes_solvers.pyx @@ -3206,6 +3206,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 diff --git a/src/underworld3/utilities/boundary_flux.py b/src/underworld3/utilities/boundary_flux.py index c0b8e5b2..abb211a6 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/tests/parallel/test_1065_boundary_flux_parallel.py b/tests/parallel/test_1065_boundary_flux_parallel.py index 7f20a963..37038f1e 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/test_1019_boundary_flux.py b/tests/test_1019_boundary_flux.py index 3103da57..944f712d 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}") From fb268de3cef68e08db7cb558d42247fdf5a35968 Mon Sep 17 00:00:00 2001 From: lmoresi Date: Thu, 27 Aug 2026 16:49:59 +1000 Subject: [PATCH 23/23] Restore the field-decomposition path for single-field solvers MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit CI was 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 simplified: - _names, _iss, _subdms = self.dm.createFieldDecomposition() - sgvec = gvec.getSubVector(_iss[0]) - _subdms[0].localToGlobal(self.Unknowns.u.vec, sgvec) - gvec.restoreSubVector(_iss[0], sgvec) + self.dm.localToGlobal(self.Unknowns.u.vec, gvec) `self.Unknowns.u.vec` is the VARIABLE's local vector. The old code mapped it through field 0's subDM; the direct call assumes the solver DM's local layout matches the variable's. On a plain single-field solver the two coincide, which is why the simplification looked equivalent — where they differ it writes to the wrong slots and the field comes back never-written. The failure is consistent with that: recovered values ~1e-18 against an analytic ~1e-5, and both arrays sit far inside the assertion's atol=0.01, so `allclose` returned False on a non-finite entry rather than on a magnitude error. This restores the original mapping and leaves everything else on the branch untouched. The harmonic projection this PR is actually for is unaffected. NOT independently confirmed as the cause: the failure does not reproduce on macOS/arm64 (test_1120 passes 3/3 there both with and without the change), so this is the cheapest decisive experiment rather than a verified fix. If CI is still red after it, the next candidate is a NaN out of `uw.function.evaluate` feeding `vec_prof_uw` (see #604, #641). Verified before pushing: test_1120 3 passed; test_1070 (this PR's own geoid suite) 13 passed, so the revert does not undo what the branch is for; and `level_1 and tier_a` 1093 passed, 0 failed. Underworld development team with AI support from Claude Code --- .../cython/petsc_generic_snes_solvers.pyx | 15 ++++++++++++++- 1 file changed, 14 insertions(+), 1 deletion(-) diff --git a/src/underworld3/cython/petsc_generic_snes_solvers.pyx b/src/underworld3/cython/petsc_generic_snes_solvers.pyx index 5f7cde8c..f3625fa5 100644 --- a/src/underworld3/cython/petsc_generic_snes_solvers.pyx +++ b/src/underworld3/cython/petsc_generic_snes_solvers.pyx @@ -3125,7 +3125,20 @@ class SolverBaseClass(uw_object): self._subdict[name][1].localToGlobal(var.vec, sgvec) gvec.restoreSubVector(self._subdict[name][0], sgvec) else: - self.dm.localToGlobal(self.Unknowns.u.vec, gvec) + # 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) + gvec.restoreSubVector(_iss[0], sgvec) self.dm.globalToLocal(gvec, xlocal)