diff --git a/src/underworld3/meshing/__init__.py b/src/underworld3/meshing/__init__.py index 9a7a7fa8..d64cdc44 100644 --- a/src/underworld3/meshing/__init__.py +++ b/src/underworld3/meshing/__init__.py @@ -50,6 +50,10 @@ FaultCollection, ) +from .smoothing import ( + smooth_mesh_interior, +) + # Make all functions available at module level for backward compatibility __all__ = [ # Cartesian meshes @@ -82,4 +86,6 @@ # Backward compatibility aliases "FaultSurface", "FaultCollection", + # Mesh smoothing + "smooth_mesh_interior", ] diff --git a/src/underworld3/meshing/smoothing.py b/src/underworld3/meshing/smoothing.py new file mode 100644 index 00000000..c5038132 --- /dev/null +++ b/src/underworld3/meshing/smoothing.py @@ -0,0 +1,404 @@ +"""Mesh smoothing utilities. + +Currently provides a Winslow-style Jacobi smoother for interior +vertex positions: each interior vertex is moved toward the average +position of its edge neighbours, with boundary vertices held fixed. + +Use after a mesh deformation has left some cells highly distorted +(e.g. free-surface evolution that has crushed cells near the +surface). Topology is unchanged — vertex indices, DOFs, and the +parallel partition are all preserved; only coordinates move. + +Parallel: a PETSc parallel AIJ matrix represents the vertex-vertex +adjacency. Each rank inserts entries for every edge it sees locally +using GLOBAL vertex indices; ``mat.assemble()`` combines cross-rank +contributions so that owned-vertex rows are complete after assembly. +Without this, UW3's default cell-overlap-0 distribution under-counts +neighbours for vertices on the rank partition boundary, producing +visibly wrong updates along the rank cut. + +Future extensions (separate PRs): + - PR B: nicer pinning API (per-boundary explicit lists, callable + masks) + - PR C: non-uniform metric (swarm-anchored target spacing, + mirroring ``mesh.adapt`` semantics) +""" + +from typing import Optional, Sequence + +import numpy as np + +import underworld3 as uw + + +# Cached adjacency keyed by (mesh-id, pinned-label-tuple, topology). +# Rebuilt automatically when the mesh topology changes. +_ADJ_CACHE: dict = {} + + +def _auto_pinned_labels(mesh) -> tuple: + """All non-sentinel geometric boundary labels on the mesh. + + Skips ``All_Boundaries`` / ``Null_Boundary`` (sentinels) and + known non-geometric pressure-pin markers such as ``Centre`` on + the Annulus (a single-point marker whose underlying ``DMLabel`` + has an invalid communicator and hard-crashes any + ``getNumValues`` / ``getValueIS`` / ``view`` call). + """ + skip = {"All_Boundaries", "Null_Boundary", "Centre"} + names = [] + for member in mesh.boundaries: + name = getattr(member, "name", None) + if name and name not in skip: + names.append(name) + return tuple(names) + + +def _owned_vertex_mask(dm): + """Local-chart boolean mask: True for owned vertices, False for + ghosts (leaves of the point StarForest). Used by the parallel + tests; the smoother itself derives ownership from the global + section attached to its scalar DM clone. + """ + pStart, pEnd = dm.getDepthStratum(0) + n_verts = pEnd - pStart + is_owned = np.ones(n_verts, dtype=bool) + sf = dm.getPointSF() + if sf is None: + return is_owned + try: + _n_roots, leaves, _remote = sf.getGraph() + except Exception: + return is_owned + if leaves is None or len(leaves) == 0: + return is_owned + for leaf in leaves: + if pStart <= leaf < pEnd: + is_owned[leaf - pStart] = False + return is_owned + + +def _pinned_mask(dm, pinned_labels): + """Local-chart boolean mask: True where the vertex belongs to (or + is the endpoint of an edge in) any of ``pinned_labels``. + + UW3 mesh generators tag boundaries by EDGE rather than by + vertex; the vertex stratum sometimes misses 1-2 endpoint + vertices at the gmsh seam (e.g. θ=0°/180° on the Annulus outer + rim). Pinning by vertex-stratum-only would leave those + "seam" vertices free, and the smoother would pull them + inward. Taking the closure of the tagged edges recovers them. + + Tolerates labels that are present but empty (e.g. the + ``Centre`` pressure-pin marker on an Annulus, whose underlying + ``DMLabel`` has no strata and hard-crashes any query).""" + pStart, pEnd = dm.getDepthStratum(0) + eStart, eEnd = dm.getDepthStratum(1) + n_verts = pEnd - pStart + is_pinned = np.zeros(n_verts, dtype=bool) + for lname in pinned_labels: + label = dm.getLabel(lname) + if label is None: + continue + try: + if label.getNumValues() == 0: + continue + vIS = label.getValueIS() + except Exception: + continue + if vIS is None: + continue + for val in vIS.getIndices(): + try: + iset = label.getStratumIS(int(val)) + except Exception: + continue + if iset is None: + continue + for idx in iset.getIndices(): + if pStart <= idx < pEnd: + # Tagged vertex — pin directly. + is_pinned[idx - pStart] = True + elif eStart <= idx < eEnd: + # Tagged edge — pin both endpoint vertices. + cone = dm.getCone(idx) + for c in cone: + if pStart <= c < pEnd: + is_pinned[c - pStart] = True + return is_pinned + + +def _build_scalar_dm(dm): + """A clone of the topological DM with a 1-dof-per-vertex local + section. Used to size the adjacency Mat and to produce the global + vertex numbering.""" + from petsc4py import PETSc + chart_start, chart_end = dm.getChart() + pStart, pEnd = dm.getDepthStratum(0) + section = PETSc.Section().create(comm=dm.getComm()) + section.setChart(chart_start, chart_end) + for p in range(chart_start, chart_end): + section.setDof(p, 1 if pStart <= p < pEnd else 0) + section.setUp() + dm_scalar = dm.clone() + dm_scalar.setLocalSection(section) + return dm_scalar + + +def _build_adjacency_matrix(mesh): + """Build the parallel vertex-vertex adjacency as a PETSc AIJ Mat. + + Each rank inserts entries for every locally-visible edge using + GLOBAL vertex indices; ``mat.assemble()`` combines cross-rank + contributions, so that after assembly an owned-vertex row has + every neighbour it would in a serial run — even when the + incident edge lives in a cell owned by another rank that is not + in this rank's overlap. + + Returns + ------- + A : PETSc.Mat + Unweighted vertex-vertex adjacency, entries are 1.0 where an + edge exists. Divide the result of ``A @ x`` by the degree + vector to get the neighbour average. + dm_scalar : PETSc.DMPlex + Clone of ``mesh.dm`` with a 1-dof-per-vertex section. Owns + the parallel layout for the Mat and any vectors of the same + shape. + local_to_global_owned : numpy.ndarray, shape (n_owned,) + ``local_to_global_owned[i]`` is the offset (in the *local* + owned portion of the global Vec) at which the ``i``-th + OWNED local vertex appears. Use this to pack/unpack between + ``coords[is_owned, d]`` and ``vec.array``. + """ + from petsc4py import PETSc + dm = mesh.dm + pStart, pEnd = dm.getDepthStratum(0) + eStart, eEnd = dm.getDepthStratum(1) + + dm_scalar = _build_scalar_dm(dm) + gsection = dm_scalar.getGlobalSection() + + def gidx(p): + off = gsection.getOffset(p) + return off if off >= 0 else -(off + 1) + + A = dm_scalar.createMatrix() + A.setOption(A.Option.NEW_NONZERO_LOCATION_ERR, False) + A.setOption(A.Option.IGNORE_OFF_PROC_ENTRIES, False) + + for e in range(eStart, eEnd): + cone = dm.getCone(e) + if len(cone) != 2: + continue + v0, v1 = cone[0], cone[1] + if not (pStart <= v0 < pEnd and pStart <= v1 < pEnd): + continue + g0, g1 = gidx(v0), gidx(v1) + A.setValues([g0], [g1], [1.0], PETSc.InsertMode.INSERT) + A.setValues([g1], [g0], [1.0], PETSc.InsertMode.INSERT) + A.assemble() + return A, dm_scalar, gsection + + +def _build_local_to_owned_map(dm, gsection, vec): + """Compute, for each local owned vertex, its position in the + rank's slice of the global Vec. + + Returns (owned_local_indices, owned_vec_positions, is_owned_local) + where: + * owned_local_indices : local-chart indices of owned vertices + (shape n_owned, dtype int64) + * owned_vec_positions : positions in vec.array (same shape) + * is_owned_local : bool mask over the local chart + """ + pStart, pEnd = dm.getDepthStratum(0) + n_local = pEnd - pStart + rstart, rend = vec.getOwnershipRange() + is_owned = np.zeros(n_local, dtype=bool) + owned_local = [] + owned_vec_pos = [] + for v in range(pStart, pEnd): + off = gsection.getOffset(v) + if off < 0: + continue # ghost + is_owned[v - pStart] = True + owned_local.append(v - pStart) + owned_vec_pos.append(off - rstart) + return (np.asarray(owned_local, dtype=np.int64), + np.asarray(owned_vec_pos, dtype=np.int64), + is_owned) + + +def smooth_mesh_interior( + mesh, + pinned_labels: Optional[Sequence[str]] = None, + n_iters: int = 5, + alpha: float = 0.5, + verbose: bool = False, +): + r"""Apply Winslow Jacobi smoothing to a mesh's interior vertices. + + Each interior vertex is replaced by a blend of its current + position and the unweighted mean of its edge-neighbour positions: + + .. math:: + + x_i^{n+1} = (1 - \alpha)\, x_i^n + + \alpha \cdot \frac{1}{|N(i)|} + \sum_{j \in N(i)} x_j^n + + Vertices in any of ``pinned_labels`` are held fixed (preserves + boundary geometry). The mesh's coordinate vector is updated in + place via ``mesh._deform_mesh`` once after all sweeps — so the + DM rebuild / cache invalidation cost is paid once rather than + per sweep. + + Parameters + ---------- + mesh : underworld3.discretisation.Mesh + The mesh to smooth. Modified in place. + pinned_labels : sequence of str, optional + Names of boundary labels whose vertices stay fixed. If + ``None`` (default), all non-sentinel labels on + ``mesh.boundaries`` are pinned — i.e. every named boundary + stays put. Pass an explicit list to release some boundaries. + n_iters : int, default 5 + Number of Jacobi sweeps. 5-10 is typical for surface- + deformation cleanup. + alpha : float, default 0.5 + Under-relaxation in ``(0, 1]``. 1.0 is pure Jacobi; smaller + is more damped (slower but safer on irregular meshes). + verbose : bool, default False + Print per-sweep RMS interior displacement. + + Notes + ----- + **Parallel implementation**: the vertex-vertex adjacency is + assembled as a parallel PETSc AIJ matrix; each rank inserts + entries for every locally-visible edge using GLOBAL vertex + indices and ``mat.assemble()`` routes cross-rank contributions + so that owned-vertex rows are complete after assembly. The + per-sweep update is then a per-component ``A.mult`` followed by + a pointwise divide by the precomputed degree vector. Results + are bit-identical (to a single ULP) between serial and parallel + runs at any rank count. + + **Topology preservation**: vertex IDs, DOF mappings, and the + rank partition are unchanged. Only coordinates move. Anything + cached against the topology version stays valid; anything + cached against coords is invalidated by the final + ``mesh._deform_mesh`` call. + + Examples + -------- + Pin all named boundaries (the usual case):: + + import underworld3 as uw + from underworld3.meshing import smooth_mesh_interior + + mesh = uw.meshing.Annulus(...) + # ... some deformation that leaves bad cells ... + smooth_mesh_interior(mesh, n_iters=5, alpha=0.5) + + Pin only the outer boundary, allowing the inner to drift:: + + smooth_mesh_interior(mesh, pinned_labels=["Upper"]) + + Pin nothing (free-floating; rare — boundary will collapse):: + + smooth_mesh_interior(mesh, pinned_labels=[]) + """ + if pinned_labels is None: + pinned_labels = _auto_pinned_labels(mesh) + pinned_labels = tuple(pinned_labels) + + dm = mesh.dm + pStart, pEnd = dm.getDepthStratum(0) + cStart, cEnd = dm.getHeightStratum(0) + cone_size = dm.getConeSize(cStart) if cEnd > cStart else 0 + cache_key = (id(mesh), pinned_labels, + pEnd - pStart, cEnd - cStart, cone_size) + + cache = _ADJ_CACHE.get(cache_key) + if cache is None: + A, dm_scalar, gsection = _build_adjacency_matrix(mesh) + # A scratch global Vec of the right shape — also used to read + # the ownership range when packing/unpacking coord components. + x_vec = A.createVecRight() + y_vec = A.createVecLeft() + ones = A.createVecLeft() + ones.set(1.0) + degrees = A.createVecLeft() + A.mult(ones, degrees) + owned_local, owned_vec_pos, is_owned = ( + _build_local_to_owned_map(dm, gsection, x_vec)) + is_pinned = _pinned_mask(dm, pinned_labels) + _ADJ_CACHE[cache_key] = ( + A, dm_scalar, gsection, x_vec, y_vec, degrees, + owned_local, owned_vec_pos, is_owned, is_pinned) + else: + (A, dm_scalar, gsection, x_vec, y_vec, degrees, + owned_local, owned_vec_pos, is_owned, is_pinned) = cache + + # is_int_owned over the LOCAL chart — selects interior owned + # vertices for displacement reporting. + is_int_owned = is_owned & ~is_pinned + # Subset of owned_local that's also interior (i.e. not pinned) + # — used to write the per-sweep updates into the numpy buffer. + int_mask_on_owned = ~is_pinned[owned_local] + int_owned_local = owned_local[int_mask_on_owned] + int_owned_vec_pos = owned_vec_pos[int_mask_on_owned] + + coord_dm = dm.getCoordinateDM() + local_vec = dm.getCoordinatesLocal() + global_vec = dm.getCoordinates() + cdim = mesh.cdim + parallel = uw.mpi.size > 1 + + coords = np.asarray( + local_vec.array, dtype=np.double).reshape(-1, cdim).copy() + + for sweep in range(n_iters): + new_int = np.empty((int_owned_local.shape[0], cdim), + dtype=np.double) + # For each coordinate component, do A @ coord_comp (PETSc + # handles cross-rank communication), then divide by degree + # to get the per-vertex neighbour average. + for d in range(cdim): + x_vec.array[owned_vec_pos] = coords[owned_local, d] + A.mult(x_vec, y_vec) + y_vec.pointwiseDivide(y_vec, degrees) + avg_owned = np.asarray(y_vec.array) + new_int[:, d] = ( + (1.0 - alpha) * coords[int_owned_local, d] + + alpha * avg_owned[int_owned_vec_pos]) + + if verbose: + disp = float(np.linalg.norm( + new_int - coords[int_owned_local])) + if parallel: + disp = uw.mpi.comm.allreduce( + disp ** 2) ** 0.5 + uw.pprint( + f" smooth_mesh_interior sweep " + f"{sweep+1}/{n_iters}: " + f"||Δx||_interior = {disp:.3e}") + + coords[int_owned_local] = new_int + + if parallel: + # Halo exchange so the next sweep sees updated owned + # values on every rank's ghost copies. (PETSc's mat.mult + # handles cross-rank READS internally via the matrix's + # column communication, so this halo exchange is only + # needed to keep the LOCAL coord array consistent for + # the final ``mesh._deform_mesh`` call.) + local_vec.array[:] = coords.ravel() + coord_dm.localToGlobal( + local_vec, global_vec, addv=False) + coord_dm.globalToLocal(global_vec, local_vec) + coords[:] = np.asarray( + local_vec.array).reshape(-1, cdim) + + mesh._deform_mesh(coords) diff --git a/src/underworld3/systems/solvers.py b/src/underworld3/systems/solvers.py index b3e36046..625cc4e9 100644 --- a/src/underworld3/systems/solvers.py +++ b/src/underworld3/systems/solvers.py @@ -2518,6 +2518,20 @@ class SNES_AdvectionDiffusion(SNES_Scalar): applied to the internally-constructed ``DFDt`` only — the user's ``DuDt`` already encodes whatever ``monotone_mode`` it was constructed with. + theta : float, default=0.5 + Adams-Moulton theta for the diffusive flux at order 1. + Forwarded to the internally-constructed + ``SemiLagrangian_DDt`` instances (same forwarding rule as + ``monotone_mode``). + + - ``0.5`` (default): Crank-Nicolson, A-stable but not + L-stable. Second-order accurate. Rings on stiff modes + with sharp gradients (negative amplification factor). + - ``1.0``: Backward Euler, L-stable, monotone for the + diffusive flux. First-order accurate. Recommended when + SLCN+CN ringing dominates the discretisation error. + - ``0.0``: Forward Euler — unstable for stiff diffusion; + included for completeness. Notes ----- @@ -2559,6 +2573,7 @@ def __init__( DuDt: Union[SemiLagrangian_DDt, Lagrangian_DDt] = None, DFDt: Union[SemiLagrangian_DDt, Lagrangian_DDt] = None, monotone_mode: Optional[str] = None, + theta: float = 0.5, ): ## Parent class will set up default values etc super().__init__( @@ -2605,6 +2620,7 @@ def __init__( order=1, smoothing=0.0, monotone_mode=monotone_mode, + theta=theta, ) else: @@ -2637,6 +2653,7 @@ def __init__( order=order, smoothing=0.0, monotone_mode=monotone_mode, + theta=theta, ) return diff --git a/tests/parallel/test_0855_mesh_smoothing_parallel.py b/tests/parallel/test_0855_mesh_smoothing_parallel.py new file mode 100644 index 00000000..6f98f22c --- /dev/null +++ b/tests/parallel/test_0855_mesh_smoothing_parallel.py @@ -0,0 +1,331 @@ +"""Parallel regression tests for ``smooth_mesh_interior``. + +The vertex-vertex adjacency is built as a parallel PETSc AIJ matrix; +each rank inserts its locally-visible edges using GLOBAL vertex +indices, and ``mat.assemble()`` routes cross-rank contributions so +that owned-vertex rows are complete after assembly. These tests +verify the parallel-safety properties: + + * Boundary vertices remain bit-identical on every rank + * After ``n_iters`` sweeps, every rank's ghost-vertex copies + agree exactly with the owner's value (halo exchange is doing + its job) + * Per-sweep interior displacement decreases monotonically using + a global reduction (matches the serial guarantee) + * Final coords from a parallel run match a serial reference to + a single ULP — i.e. the smoother is partition-independent + +Run with: + mpirun -n 2 python -m pytest --with-mpi \\ + tests/parallel/test_0855_mesh_smoothing_parallel.py + mpirun -n 4 python -m pytest --with-mpi \\ + tests/parallel/test_0855_mesh_smoothing_parallel.py +""" + +import os +import subprocess +import sys +import tempfile +import textwrap + +import numpy as np +import pytest +from mpi4py import MPI +from scipy.spatial import cKDTree + +import underworld3 as uw +from underworld3.meshing import smooth_mesh_interior + + +pytestmark = [pytest.mark.mpi(min_size=2), pytest.mark.timeout(120)] + + +def _box_mesh(resolution: int = 12): + return uw.meshing.UnstructuredSimplexBox( + minCoords=(0.0, 0.0), + maxCoords=(1.0, 1.0), + cellSize=1.0 / resolution, + ) + + +def _boundary_vertex_mask(mesh): + dm = mesh.dm + pStart, pEnd = dm.getDepthStratum(0) + n_verts = pEnd - pStart + skip = {"All_Boundaries", "Null_Boundary"} + mask = np.zeros(n_verts, dtype=bool) + for member in mesh.boundaries: + name = getattr(member, "name", None) + if not name or name in skip: + continue + label = dm.getLabel(name) + if label is None: + continue + vIS = label.getValueIS() + if vIS is None: + continue + for val in vIS.getIndices(): + iset = label.getStratumIS(int(val)) + if iset is None: + continue + for idx in iset.getIndices(): + if pStart <= idx < pEnd: + mask[idx - pStart] = True + return mask + + +@pytest.mark.mpi(min_size=2) +def test_parallel_boundary_pinned(): + """Boundary vertices stay bit-identical on every rank.""" + mesh = _box_mesh(resolution=12) + is_bnd = _boundary_vertex_mask(mesh) + before = np.asarray(mesh.X.coords).copy() + smooth_mesh_interior(mesh, n_iters=5, alpha=0.5) + after = np.asarray(mesh.X.coords) + if int(is_bnd.sum()) > 0: + assert np.allclose( + before[is_bnd], after[is_bnd], rtol=0, atol=0), ( + f"Rank {uw.mpi.rank}: boundary vertices moved.") + + +@pytest.mark.mpi(min_size=2) +def test_parallel_ghosts_agree_with_owners(): + """After smoothing, every ghost vertex's local coord must + match the owner's. Verifies the halo exchange did its job.""" + mesh = _box_mesh(resolution=12) + rng = np.random.default_rng(7 + uw.mpi.rank) + coords = np.asarray(mesh.X.coords).copy() + is_bnd = _boundary_vertex_mask(mesh) + if (~is_bnd).any(): + coords[~is_bnd] += 0.01 * rng.standard_normal( + (int((~is_bnd).sum()), coords.shape[1])) + mesh._deform_mesh(coords) + smooth_mesh_interior(mesh, n_iters=5, alpha=0.5) + # After the call, the DM's coordinate vector has been updated. + # We verify ghost-owner agreement by doing a fresh halo exchange: + # a properly-consistent state is invariant under another + # globalToLocal — i.e. ghost values are unchanged after refresh. + dm = mesh.dm + coord_dm = dm.getCoordinateDM() + local_vec = dm.getCoordinatesLocal() + global_vec = dm.getCoordinates() + before = np.asarray(local_vec.array).copy() + coord_dm.globalToLocal(global_vec, local_vec) + after = np.asarray(local_vec.array) + assert np.allclose(before, after, rtol=0, atol=1.0e-15), ( + f"Rank {uw.mpi.rank}: ghost-owner disagreement of " + f"max |Δ| = {np.abs(before - after).max():.3e} after " + f"smoothing — halo exchange did not propagate correctly") + + +@pytest.mark.mpi(min_size=2) +def test_parallel_sweep_displacement_decreases(): + """Global per-sweep displacement decreases, same as the + serial guarantee.""" + mesh = _box_mesh(resolution=12) + rng = np.random.default_rng(13 + uw.mpi.rank) + coords = np.asarray(mesh.X.coords).copy() + is_bnd = _boundary_vertex_mask(mesh) + if (~is_bnd).any(): + coords[~is_bnd] += 0.02 * rng.standard_normal( + (int((~is_bnd).sum()), coords.shape[1])) + mesh._deform_mesh(coords) + comm = MPI.COMM_WORLD + prev = np.asarray(mesh.X.coords).copy() + disps = [] + for _ in range(4): + smooth_mesh_interior(mesh, n_iters=1, alpha=0.5) + now = np.asarray(mesh.X.coords) + is_bnd_now = _boundary_vertex_mask(mesh) + is_int = ~is_bnd_now + local_sq = ( + float(np.linalg.norm((now - prev)[is_int]) ** 2) + if is_int.any() else 0.0) + global_disp = comm.allreduce(local_sq) ** 0.5 + disps.append(global_disp) + prev = now.copy() + for k in range(1, len(disps)): + assert disps[k] <= disps[k - 1] * 1.001, ( + f"Sweep {k+1} global displacement {disps[k]:.3e} > " + f"{disps[k-1]:.3e} (sweep {k}); series: {disps}") + + +# --------------------------------------------------------------------- +# Bit-identical-vs-serial regression test +# --------------------------------------------------------------------- + +# Deterministic, position-only perturbation. The smoother takes a +# perturbed mesh through 8 sweeps; we then compare the final owned +# coords against a serial reference. The perturbation is a pure +# function of (x, y) so that serial and parallel runs start from the +# same field, regardless of where the partition cut lands. +_SERIAL_REFERENCE_SCRIPT = textwrap.dedent(""" + import sys + import numpy as np + + import underworld3 as uw + from underworld3.meshing import ( + UnstructuredSimplexBox, smooth_mesh_interior) + + + def _boundary_vertex_mask(mesh): + dm = mesh.dm + pStart, pEnd = dm.getDepthStratum(0) + n = pEnd - pStart + skip = {"All_Boundaries", "Null_Boundary"} + mask = np.zeros(n, dtype=bool) + for b in mesh.boundaries: + nm = getattr(b, "name", None) + if not nm or nm in skip: + continue + lab = dm.getLabel(nm) + if lab is None: + continue + vIS = lab.getValueIS() + if vIS is None: + continue + for v in vIS.getIndices(): + iset = lab.getStratumIS(int(v)) + if iset is None: + continue + for idx in iset.getIndices(): + if pStart <= idx < pEnd: + mask[idx - pStart] = True + return mask + + + out_path = sys.argv[1] + mesh = UnstructuredSimplexBox( + minCoords=(0.0, 0.0), maxCoords=(1.0, 1.0), + cellSize=1.0 / 12) + is_bnd = _boundary_vertex_mask(mesh) + coords = np.asarray(mesh.X.coords).copy() + initial = coords.copy() + dx = 0.018 * np.sin(7.0 * np.pi * coords[:, 0]) \\ + * np.cos(5.0 * np.pi * coords[:, 1]) + dy = 0.018 * np.cos(3.0 * np.pi * coords[:, 0]) \\ + * np.sin(11.0 * np.pi * coords[:, 1]) + coords[~is_bnd, 0] += dx[~is_bnd] + coords[~is_bnd, 1] += dy[~is_bnd] + mesh._deform_mesh(coords) + smooth_mesh_interior(mesh, n_iters=8, alpha=0.5) + np.savez(out_path, + initial=initial, + final=np.asarray(mesh.X.coords)) +""") + + +def _owned_vertex_mask_local(dm): + pStart, pEnd = dm.getDepthStratum(0) + n = pEnd - pStart + owned = np.ones(n, dtype=bool) + sf = dm.getPointSF() + if sf is None: + return owned + try: + _, leaves, _ = sf.getGraph() + except Exception: + return owned + if leaves is None: + return owned + for L in leaves: + if pStart <= L < pEnd: + owned[L - pStart] = False + return owned + + +@pytest.mark.mpi(min_size=2) +def test_parallel_matches_serial_bit_identical(): + """Final coords from the parallel smoother match a serial + reference to a single ULP, partitioning-independent. + + Rank 0 spawns a serial Python subprocess that runs the same + mesh + perturbation + 8-sweep smoothing pipeline; rank 0 then + compares the parallel result (gathered from every rank's owned + vertices) against the serial reference, matching vertices by + their initial (pre-perturbation) coordinate. + """ + comm = MPI.COMM_WORLD + rank = comm.rank + + # 1. Compute the serial reference (rank 0 only) in a clean + # subprocess. We strip MPI/PMIX/PRTE env vars first so the + # subprocess's PETSc doesn't try to attach to this mpirun's + # MPI world (which would deadlock both processes). + ref_path = None + if rank == 0: + tmpdir = tempfile.mkdtemp(prefix="winslow_ref_") + ref_path = os.path.join(tmpdir, "ref.npz") + clean_env = { + k: v for k, v in os.environ.items() + if not (k.startswith("OMPI_") + or k.startswith("PMIX_") + or k.startswith("PRTE_") + or k.startswith("PRTERUN_") + or k == "OPAL_PREFIX") + } + proc = subprocess.run( + [sys.executable, "-c", + _SERIAL_REFERENCE_SCRIPT, ref_path], + capture_output=True, text=True, timeout=60, + env=clean_env) + assert proc.returncode == 0, ( + f"serial reference subprocess failed:\n" + f"stdout:\n{proc.stdout}\nstderr:\n{proc.stderr}") + assert os.path.exists(ref_path), ( + f"serial reference not written at {ref_path}") + + # 2. Run the parallel pipeline on this comm + mesh = uw.meshing.UnstructuredSimplexBox( + minCoords=(0.0, 0.0), maxCoords=(1.0, 1.0), + cellSize=1.0 / 12) + is_bnd = _boundary_vertex_mask(mesh) + coords = np.asarray(mesh.X.coords).copy() + initial_local = coords.copy() + dx = 0.018 * np.sin(7.0 * np.pi * coords[:, 0]) \ + * np.cos(5.0 * np.pi * coords[:, 1]) + dy = 0.018 * np.cos(3.0 * np.pi * coords[:, 0]) \ + * np.sin(11.0 * np.pi * coords[:, 1]) + coords[~is_bnd, 0] += dx[~is_bnd] + coords[~is_bnd, 1] += dy[~is_bnd] + mesh._deform_mesh(coords) + smooth_mesh_interior(mesh, n_iters=8, alpha=0.5) + + # 3. Gather owned final coords + their pre-perturbation initial + # coords (the latter is the matching key against the serial ref). + is_owned = _owned_vertex_mask_local(mesh.dm) + final_local = np.asarray(mesh.X.coords).copy() + own_initial = initial_local[is_owned] + own_final = final_local[is_owned] + g_initial = comm.gather(own_initial, root=0) + g_final = comm.gather(own_final, root=0) + + # 4. Compare on rank 0 + if rank == 0: + all_initial = np.vstack(g_initial) + all_final = np.vstack(g_final) + ref = np.load(ref_path) + ref_initial = ref["initial"] + ref_final = ref["final"] + # Match each parallel vertex to its serial counterpart by + # initial coordinate (bit-identical because the mesh + # generator runs serially before distribution). + tree = cKDTree(ref_initial) + match_dist, idx = tree.query(all_initial, k=1) + assert match_dist.max() < 1.0e-12, ( + "initial-coord matching failed: max mismatch " + f"{match_dist.max():.3e} — meshes must be identical " + "between serial and parallel runs") + drift = np.linalg.norm(all_final - ref_final[idx], axis=1) + size = MPI.COMM_WORLD.size + assert drift.max() < 1.0e-12, ( + f"np={size}: parallel smoother diverged from serial; " + f"max drift = {drift.max():.3e} " + f"mean = {drift.mean():.3e} " + f"({(drift > 1.0e-12).sum()}/{len(drift)} verts off)") + # Cleanup + try: + os.remove(ref_path) + os.rmdir(os.path.dirname(ref_path)) + except OSError: + pass diff --git a/tests/test_0850_mesh_smoothing.py b/tests/test_0850_mesh_smoothing.py new file mode 100644 index 00000000..de133119 --- /dev/null +++ b/tests/test_0850_mesh_smoothing.py @@ -0,0 +1,186 @@ +"""Regression tests for ``underworld3.meshing.smooth_mesh_interior``. + +The smoother relaxes interior vertex positions toward the average +of their edge neighbours, with boundary vertices held fixed. Tests: + +- Boundary vertices stay bit-identical after smoothing +- A uniform (already-equilibrated) mesh is a fixed point: smoothing + does not move interior vertices beyond floating-point noise +- A perturbed-interior mesh converges back toward the uniform state +- Pinned-label specification works (explicit list, ``None`` = + auto-detect, ``[]`` = nothing pinned) +- DM coordinate vector and ``mesh.X.coords`` are updated by the + single trailing ``_deform_mesh`` call +""" + +import numpy as np +import pytest + +import underworld3 as uw +from underworld3.meshing import smooth_mesh_interior + + +pytestmark = [pytest.mark.level_1, pytest.mark.tier_a] + + +def _box_mesh(resolution: int = 8): + return uw.meshing.UnstructuredSimplexBox( + minCoords=(0.0, 0.0), + maxCoords=(1.0, 1.0), + cellSize=1.0 / resolution, + ) + + +def _min_aspect_ratio(mesh): + """Compute (min altitude) / (longest edge) across all triangles + on the local rank. Closer to ~√3/2 ≈ 0.866 is more equilateral. + Returns a single float (min over all cells).""" + dm = mesh.dm + pStart, pEnd = dm.getDepthStratum(0) + cStart, cEnd = dm.getHeightStratum(0) + coords = np.asarray(mesh.X.coords) + worst = np.inf + for c in range(cStart, cEnd): + closure, _ = dm.getTransitiveClosure(c, useCone=True) + verts = [p - pStart for p in closure + if pStart <= p < pEnd] + if len(verts) != 3: + continue + pts = coords[verts] + a = np.linalg.norm(pts[1] - pts[0]) + b = np.linalg.norm(pts[2] - pts[1]) + c_ = np.linalg.norm(pts[0] - pts[2]) + Lmax = max(a, b, c_) + # 2D triangle signed area + area = 0.5 * abs((pts[1, 0] - pts[0, 0]) + * (pts[2, 1] - pts[0, 1]) + - (pts[2, 0] - pts[0, 0]) + * (pts[1, 1] - pts[0, 1])) + if Lmax > 0: + h_min = 2.0 * area / Lmax + ar = h_min / Lmax + if ar < worst: + worst = ar + return worst + + +def _boundary_vertex_mask(mesh): + """Local-chart boolean mask: True where the vertex lies on any + non-sentinel boundary label.""" + dm = mesh.dm + pStart, pEnd = dm.getDepthStratum(0) + n_verts = pEnd - pStart + skip = {"All_Boundaries", "Null_Boundary"} + mask = np.zeros(n_verts, dtype=bool) + for member in mesh.boundaries: + name = getattr(member, "name", None) + if not name or name in skip: + continue + label = dm.getLabel(name) + if label is None: + continue + vIS = label.getValueIS() + if vIS is None: + continue + for val in vIS.getIndices(): + iset = label.getStratumIS(int(val)) + if iset is None: + continue + for idx in iset.getIndices(): + if pStart <= idx < pEnd: + mask[idx - pStart] = True + return mask + + +class TestBasic: + + def test_boundary_vertices_held_fixed(self): + mesh = _box_mesh(resolution=6) + before = np.asarray(mesh.X.coords).copy() + is_bnd = _boundary_vertex_mask(mesh) + smooth_mesh_interior(mesh, n_iters=5, alpha=0.5) + after = np.asarray(mesh.X.coords) + # Boundary coords must be bit-identical + assert np.allclose( + before[is_bnd], after[is_bnd], rtol=0, atol=0), ( + "Boundary vertices moved during smoothing" + ) + + def test_sweep_displacement_decreases(self): + """Iterating Jacobi smoothing converges: the displacement + per sweep should decrease monotonically (in L2 norm) on the + approach to the topology-dependent equilibrium.""" + mesh = _box_mesh(resolution=6) + is_bnd = _boundary_vertex_mask(mesh) + is_interior = ~is_bnd + # Take 5 sweeps one at a time, recording the interior + # displacement at each step. + disps = [] + prev = np.asarray(mesh.X.coords).copy() + for _ in range(5): + smooth_mesh_interior(mesh, n_iters=1, alpha=0.5) + now = np.asarray(mesh.X.coords) + d = float(np.linalg.norm((now - prev)[is_interior])) + disps.append(d) + prev = now.copy() + # Strictly decreasing — graph Laplacian Jacobi is a + # contraction on the (pinned-boundary) problem. + for k in range(1, len(disps)): + assert disps[k] < disps[k - 1], ( + f"Sweep {k+1} displacement {disps[k]:.3e} not " + f"less than sweep {k} displacement " + f"{disps[k-1]:.3e}; full series: {disps}") + + def test_aspect_ratio_improves_on_perturbed_mesh(self): + """Perturb the interior, smooth, and check that the minimum + triangle altitude / longest edge ratio (a shape-quality + metric, closer to 1 = better) does not get worse — and + usually improves.""" + mesh = _box_mesh(resolution=6) + is_bnd = _boundary_vertex_mask(mesh) + rng = np.random.default_rng(42) + coords = np.asarray(mesh.X.coords).copy() + is_interior = ~is_bnd + coords[is_interior] += 0.04 * rng.standard_normal( + (int(is_interior.sum()), coords.shape[1])) + mesh._deform_mesh(coords) + ar_before = _min_aspect_ratio(mesh) + smooth_mesh_interior(mesh, n_iters=10, alpha=0.5) + ar_after = _min_aspect_ratio(mesh) + assert ar_after >= ar_before * 0.95, ( + f"Smoothing made the mesh worse: " + f"min aspect ratio {ar_before:.3f} -> {ar_after:.3f}") + + +class TestPinningAPI: + + def test_explicit_pinned_labels(self): + mesh = _box_mesh(resolution=6) + # Pin only Bottom + Top — Left and Right should move + smooth_mesh_interior( + mesh, pinned_labels=["Bottom", "Top"], + n_iters=1, alpha=0.5) + # Just check it didn't crash; explicit Left/Right motion is + # too small to verify reliably without a structured layout. + + def test_empty_pinned_list_is_legal(self): + """Empty pinned_labels: every vertex is interior. Smoothing + runs without error (boundary vertices drift).""" + mesh = _box_mesh(resolution=4) + before = np.asarray(mesh.X.coords).copy() + smooth_mesh_interior( + mesh, pinned_labels=[], n_iters=1, alpha=0.5) + after = np.asarray(mesh.X.coords) + # At least one vertex moved (no pinning → mesh contracts) + assert not np.allclose(before, after) + + def test_none_pinned_labels_auto_detects(self): + """None means 'pin every named boundary' — boundary stays + fixed even though we didn't enumerate the names.""" + mesh = _box_mesh(resolution=5) + is_bnd = _boundary_vertex_mask(mesh) + before = np.asarray(mesh.X.coords).copy() + smooth_mesh_interior(mesh, pinned_labels=None, + n_iters=1, alpha=0.5) + after = np.asarray(mesh.X.coords) + assert np.allclose(before[is_bnd], after[is_bnd])