diff --git a/docs/developer/design/submesh-solver-architecture.md b/docs/developer/design/submesh-solver-architecture.md index b258dce7..5de1793d 100644 --- a/docs/developer/design/submesh-solver-architecture.md +++ b/docs/developer/design/submesh-solver-architecture.md @@ -21,6 +21,93 @@ Underworld3 needs to support solving different equations on different subsets of 5. **Normalised `Gamma_N`** (merged) — `mesh.Gamma_N` now returns a unit normal. Penalty and Nitsche BCs are mesh-independent. +## Two Submesh Flavours: Subdomain and Resolution Level + +A *submesh* in UW3 is any mesh pulled out of a parent mesh that retains a +lineage link (`parent`, registration in `parent._registered_submeshes`) +and supports explicit field transfer back and forth. There are two +flavours, and they share one usage pattern: + +> **get a submesh → build a solver on it → map fields back and forth** + +| | Subdomain | Resolution level | +|---|---|---| +| Constructor | `mesh.extract_region("Inner")` | `coarsened_companion(fine, levels=1)` | +| PETSc mechanism | `DMPlexFilter` (cell subset) | `dm.refine()` nested hierarchy | +| Parent ↔ child map | `getSubpointIS()` (point-level) | nested `createInjection` / `createInterpolation` (DOF-level) | +| Transfer fidelity | exact (shared nodes) | exact (nested FE) | +| Example | `docs/examples/submesh_investigation/test_region_ds_submesh.py` | `docs/examples/submesh_investigation/example_refined_companion.py` | + +Both examples solve the **same** annulus + radial-buoyancy Stokes problem +so the flavours are directly comparable: one solves on a *subdomain* of +the annulus, the other on a *coarser resolution* of the whole annulus, +and both map the solution back to the parent. + +### Design contract for the resolution-level flavour: refine-DM mode only + +A resolution level is extractable **only when a genuine nested +refinement hierarchy exists** (the mesh was built with +`refinement >= 1`). The hierarchy is the source of truth; any level can +be pulled out as a standalone solver-ready `uw.Mesh`, exactly as a +subdomain is pulled out with `extract_region`. + +**If there is no refinement relationship the operation is unavailable.** +There is deliberately: + +- **no geometric / `DMPlexComputeInterpolatorGeneral` fallback**, and +- **no KDTree coordinate-matching fallback**. + +`coarsened_companion` raises a clear error on a non-refined mesh rather +than silently degrading to an approximate transfer. + +Transfer between levels uses PETSc's *nested* interpolator/injector +(`plex.c:10328`, `DMPlexComputeInterpolatorNested` / +`DMPlexComputeInjectorFEM`), which is: + +- **exact** — the prolongation is the FE embedding; injection is a pure + scatter of coincident DOFs; +- **parallel-local** — the injector builds per-rank `COMM_SELF` scatters + (`plexfem.c:3739-3741`); `DMRefine` preserves the coarse partition, so + no MPI communication is needed for restrict/prolongate; +- **not dependent on point location** — no grid hashing, no geometry + search. + +```{note} +The nested path triggers only when the fine DM's `getCoarseDM()` *is* +the coarse DM and the regular-refinement flag is set. Independently +cloning two hierarchy levels breaks this and petsc4py exposes no setter +to restore it. The working construction is to build the transfer pair by +**refining a single-field clone of the coarse level** +(`dm_f = dm_c.refine()`): `refine()` itself establishes the linkage and +the flag. See `refined_pair_prototype.py` (`_linked_pair`). +``` + +Empirically (box and annulus, P2 velocity), prolongating a coarse Stokes +solution to the fine mesh and sampling it back recovers the coarse +solution to **O(1e-15)** — machine precision — with the geometric escape +hatch deliberately removed, proving the transfer is genuinely the nested +operator. + +```python +fine = uw.meshing.Annulus(radiusOuter=1.5, radiusInner=0.5, + cellSize=1/16, refinement=2) +coarse = coarsened_companion(fine, levels=1) # parent = fine + +v_c = uw.discretisation.MeshVariable("V", coarse, coarse.dim, degree=2) +stokes = Stokes(coarse, velocityField=v_c, ...) # solve cheaply, coarse +stokes.solve() + +v_f = uw.discretisation.MeshVariable("Vf", fine, fine.dim, degree=2) +prolongate(coarse, v_c, v_f) # exact, fills all fine DOFs +``` + +Status: prototype + both parallel examples + gating tests +(`test_refined_pair_solver.py`) + contract test +(`test_refined_pair_contract.py`) all passing. The investigation under +`docs/examples/submesh_investigation/` is the sign-off artifact; +promotion of `coarsened_companion` into the `Mesh` API proper is a +follow-up. + ## Design Principles ### 1. Separate meshes, separate variables, explicit copies diff --git a/docs/examples/submesh_investigation/example_refined_companion.py b/docs/examples/submesh_investigation/example_refined_companion.py new file mode 100644 index 00000000..604adeda --- /dev/null +++ b/docs/examples/submesh_investigation/example_refined_companion.py @@ -0,0 +1,148 @@ +""" +Refined-companion approach: pull a COARSE level out of a refined mesh's +nested hierarchy, solve Stokes on it, map the solution back to the fine +mesh exactly. + +This is the sibling of ``test_region_ds_submesh.py``. Both follow the +same pattern -- *get a submesh, build a solver, map back and forth* -- +but the submesh here is a different *resolution* of the whole domain +rather than a *subdomain*: + + test_region_ds_submesh.py : extract_region("Inner") -> subdomain + example_refined_companion.py : coarsened_companion(...) -> coarse level + +The same annulus + radial-buoyancy problem is used so the two examples +are directly comparable. + +Design contract (refine-DM mode only): the coarse companion is available +ONLY because a genuine nested refinement hierarchy exists. Transfer +between levels uses PETSc's *nested* interpolator/injector -- exact, +parallel-local, no geometric point location, no KDTree. On a mesh with +no refinement relationship the companion is simply not offered (raises). + +Usage: + pixi run -e amr-dev python -u \ + docs/examples/submesh_investigation/example_refined_companion.py +""" + +import numpy as np +import sympy + +import underworld3 as uw +from underworld3.systems import Stokes + +import refined_pair_prototype as rpp + +# --- Parameters --- + +r_outer = 1.5 +r_inner = 0.5 +cellsize = 1 / 16 +refine_levels = 2 # fine mesh = 2 uniform refinements of the base +companion_levels = 1 # solve one level coarser than the finest +n = 2 +k = 1 +stokes_tol = 1.0e-6 +vel_penalty = 1.0e6 + +# --- Fine mesh (the full-resolution domain) --- + +uw.pprint(0, "Creating fine annulus mesh (with refinement hierarchy)...") +fine = uw.meshing.Annulus( + radiusOuter=r_outer, + radiusInner=r_inner, + cellSize=cellsize, + refinement=refine_levels, +) +fS, fE = fine.dm.getHeightStratum(0) +uw.pprint( + 0, + f"Fine mesh: {fE - fS} cells, " + f"hierarchy depth {len(fine.dm_hierarchy)}", +) + +# --- Pull out the coarse companion (the "submesh") --- + +uw.pprint(0, "Pulling coarse level out of the nested hierarchy...") +coarse = rpp.coarsened_companion(fine, levels=companion_levels) +cS, cE = coarse.dm.getHeightStratum(0) +uw.pprint(0, f"Coarse companion: {cE - cS} cells") +uw.pprint(0, f" parent is fine mesh: {coarse.parent is fine}") +uw.pprint(0, f" registered with parent: {coarse in fine._registered_submeshes}") +uw.pprint(0, f" boundaries: {[b.name for b in coarse.boundaries]}") + +# --- Variables on the coarse companion --- + +v = uw.discretisation.MeshVariable("V", coarse, coarse.dim, degree=2) +p = uw.discretisation.MeshVariable("P", coarse, 1, degree=1, continuous=True) + +# --- Coordinate system (same radial-buoyancy problem as the rock/air ex) --- + +unit_rvec = coarse.CoordinateSystem.unit_e_0 +r, th = coarse.CoordinateSystem.xR +Gamma = coarse.Gamma + +# --- Stokes solver on the coarse companion --- + +stokes = Stokes(coarse, velocityField=v, pressureField=p) +stokes.constitutive_model = uw.constitutive_models.ViscousFlowModel +stokes.constitutive_model.Parameters.shear_viscosity_0 = 1.0 +stokes.saddle_preconditioner = 1.0 + +rho = ((r / r_outer) ** k) * sympy.cos(n * th) +stokes.bodyforce = rho * (-1.0 * unit_rvec) + +# Free-slip on both annulus boundaries +stokes.add_natural_bc(vel_penalty * Gamma.dot(v.sym) * Gamma, "Upper") +stokes.add_natural_bc(vel_penalty * Gamma.dot(v.sym) * Gamma, "Lower") + +stokes.tolerance = stokes_tol + +uw.pprint(0, "Solving Stokes on the coarse companion...") +stokes.solve(verbose=False) + +v_mag = np.sqrt(v.data[:, 0] ** 2 + v.data[:, 1] ** 2) +uw.pprint( + 0, + f"Coarse solve: |v| max={v_mag.max():.6e} mean={v_mag.mean():.6e}", +) + +# --- Map the coarse solution back to the fine mesh --- + +uw.pprint(0, "Mapping coarse velocity back to the fine mesh (nested FE)...") +v_fine = uw.discretisation.MeshVariable("Vf", fine, fine.dim, degree=2) +rpp.prolongate(coarse, v, v_fine) + +vf_mag = np.sqrt(v_fine.data[:, 0] ** 2 + v_fine.data[:, 1] ** 2) +uw.pprint( + 0, + f"Prolongated to fine: |v| max={vf_mag.max():.6e} " + f"mean={vf_mag.mean():.6e}", +) + +# Round-trip: sampling the prolongated fine field back to the coarse +# companion must recover the coarse solution exactly (it lives in the +# fine FE space by construction). +v_back = uw.discretisation.MeshVariable("Vb", coarse, coarse.dim, degree=2) +rpp.sample(coarse, v_fine, v_back) +rt_err = np.linalg.norm(v_back.data - v.data) / np.linalg.norm(v.data) + +# --- Report --- + +uw.pprint(0, "=" * 60) +uw.pprint(0, "Refined-companion approach (coarse level of refined mesh)") +uw.pprint(0, f" Coarse cells: {cE - cS}") +uw.pprint(0, f" Fine cells: {fE - fS}") +uw.pprint(0, f" Max |v| (coarse): {v_mag.max():.10e}") +uw.pprint(0, f" Max |v| (-> fine): {vf_mag.max():.10e}") +uw.pprint( + 0, + f" Prolongation max|v| rel diff: " + f"{abs(vf_mag.max() - v_mag.max()) / v_mag.max():.3e}", +) +uw.pprint(0, f" Round-trip rel error (sample o prolongate): {rt_err:.3e}") +uw.pprint(0, " -> transfer is exact (nested FE): differences are O(1e-15)") +uw.pprint(0, "=" * 60) + +assert rt_err < 1.0e-8, f"nested round-trip not exact: {rt_err:.3e}" +uw.pprint(0, "example_refined_companion: OK") diff --git a/docs/examples/submesh_investigation/probe_hierarchy_labels.py b/docs/examples/submesh_investigation/probe_hierarchy_labels.py new file mode 100644 index 00000000..dc4e9406 --- /dev/null +++ b/docs/examples/submesh_investigation/probe_hierarchy_labels.py @@ -0,0 +1,80 @@ +"""Probe: do boundary labels survive on the coarse DMs in dm_hierarchy? + +Gating question for the refined-submesh-pair investigation. If the coarse +levels of the refinement hierarchy don't carry the named boundary labels +(with non-empty strata), we cannot build a solver-ready coarse companion +and the whole approach needs rethinking. + +Run: pixi run -e amr-dev python -u docs/examples/submesh_investigation/probe_hierarchy_labels.py +""" + +import underworld3 as uw + + +def dump_hierarchy_labels(mesh, label): + print(f"\n=== {label} ===") + hier = mesh.dm_hierarchy + print(f"dm_hierarchy length: {len(hier)} (index 0 = coarsest, -1 = finest)") + bnames = [b.name for b in mesh.boundaries] if mesh.boundaries is not None else [] + print(f"boundaries enum: {bnames}") + + for lvl, dm in enumerate(hier): + cStart, cEnd = dm.getHeightStratum(0) + ncells = cEnd - cStart + print(f"\n level {lvl}: {ncells} cells") + for b in mesh.boundaries: + if b.name in ("Null_Boundary", "All_Boundaries"): + continue + lab = dm.getLabel(b.name) + if not lab: + print(f" {b.name:<20} value={b.value:<5} LABEL ABSENT") + continue + sis = lab.getStratumIS(b.value) + size = sis.getSize() if sis else 0 + flag = "" if size > 0 else " <-- EMPTY" + print(f" {b.name:<20} value={b.value:<5} stratum size={size}{flag}") + + +def main(): + # 2D box (simplex, gmsh-generated boundary labels) + box = uw.meshing.UnstructuredSimplexBox( + minCoords=(0.0, 0.0), + maxCoords=(1.0, 1.0), + cellSize=0.25, + refinement=2, + ) + dump_hierarchy_labels(box, "UnstructuredSimplexBox 2D, refinement=2") + + # Annulus with an internal boundary (different boundary-construction path). + # AnnulusInternalBoundary has no `refinement` kwarg, so build it and + # re-wrap its DM through the Mesh constructor with refinement=2 — this is + # exactly the DM-passthrough path coarsened_companion will rely on. + annulus_base = uw.meshing.AnnulusInternalBoundary( + radiusOuter=1.0, + radiusInternal=0.7, + radiusInner=0.5, + cellSize=0.2, + ) + annulus = uw.discretisation.Mesh( + annulus_base.dm, + boundaries=annulus_base.boundaries, + coordinate_system_type=annulus_base.CoordinateSystemType, + refinement=2, + qdegree=annulus_base.qdegree, + ) + dump_hierarchy_labels( + annulus, "AnnulusInternalBoundary (re-wrapped DM), refinement=2" + ) + + # Spherical shell (3D, yet another path) + shell = uw.meshing.SphericalShell( + radiusOuter=1.0, + radiusInner=0.5, + cellSize=0.4, + refinement=1, + ) + dump_hierarchy_labels(shell, "SphericalShell 3D, refinement=1") + + +if __name__ == "__main__": + main() diff --git a/docs/examples/submesh_investigation/refined_pair_prototype.py b/docs/examples/submesh_investigation/refined_pair_prototype.py new file mode 100644 index 00000000..29ceba54 --- /dev/null +++ b/docs/examples/submesh_investigation/refined_pair_prototype.py @@ -0,0 +1,348 @@ +"""Prototype: coarse companion mesh from a refinement hierarchy. + +A *coarsened companion* is a real ``uw.Mesh`` that wraps a coarser DM from +the fine mesh's refinement hierarchy. It carries the same submesh lineage +state as ``Mesh.extract_region`` (``parent``, registered with the parent's +``_registered_submeshes``) so the downstream pattern is identical: + + get a mesh -> create a solver -> map fields back and forth. + +Transfer between the fine parent and the coarse companion uses the +PETSc-native operators that geometric multigrid relies on +(``createInterpolation`` / ``createInjection`` on the variables' +single-field sub-DMs) -- no KDTree. + +This is investigation code under docs/examples/, not a merged API. +""" + +import weakref +from enum import Enum + +from petsc4py import PETSc + +import underworld3 as uw + +# DESIGN CONTRACT: refine-DM mode only. +# +# A level may be pulled out of a mesh ONLY when a genuine nested +# refinement hierarchy exists (the mesh was built with refinement >= 1). +# Transfer between levels uses PETSc's *nested* interpolator/injector, +# which is exact, parallel-local, and needs no geometric point location. +# +# There is deliberately NO geometric/General fallback and NO KDTree. If +# the nested path is not in effect the operators must fail loudly rather +# than silently degrade -- so we do NOT enable -dm_plex_hash_location. +# (If the General branch were ever taken it would raise "Nearest point +# location only supported with grid hashing", which is the desired +# loud failure.) + + +# --------------------------------------------------------------------------- +# Building the companion +# --------------------------------------------------------------------------- + +def _require_refinement_hierarchy(mesh): + """Enforce the refine-DM-only contract. + + A level can be pulled out ONLY when the mesh carries a genuine + nested refinement hierarchy (built with ``refinement >= 1``). If + there is no refinement relationship the operation is unavailable -- + we raise rather than offer a geometric/KDTree approximation. + """ + hier = getattr(mesh, "dm_hierarchy", None) + if hier is None or len(hier) < 2: + raise ValueError( + "Hierarchy-level extraction requires a genuine nested " + "refinement hierarchy: build the mesh with refinement >= 1. " + f"This mesh has {0 if hier is None else len(hier)} level(s) " + "(no refinement relationship), so a coarse/fine companion is " + "not available. There is deliberately no geometric or KDTree " + "fallback." + ) + return hier + + +def coarsened_companion(fine_mesh, levels=1, verbose=False): + """Pull a coarser level out of ``fine_mesh``'s nested hierarchy. + + Analogous to ``Mesh.extract_region`` (which pulls out a subdomain): + here we pull out a *level* of the refinement hierarchy as a + standalone, solver-ready ``uw.Mesh``. Only available when a genuine + refinement relationship exists (see the design contract above). + + Parameters + ---------- + fine_mesh : uw.discretisation.Mesh + A mesh built with ``refinement >= levels`` so ``dm_hierarchy`` + has enough levels (index 0 = coarsest, -1 = finest). + levels : int + How many refinement steps below the finest to take. + + Returns + ------- + uw.discretisation.Mesh + A solver-ready mesh whose ``.parent`` is ``fine_mesh``, + registered with ``fine_mesh._registered_submeshes``. + """ + hier = _require_refinement_hierarchy(fine_mesh) + if len(hier) <= levels: + raise ValueError( + f"dm_hierarchy has {len(hier)} levels; need > {levels}. " + f"Build the fine mesh with refinement >= {levels}." + ) + + coarse_dm = hier[-1 - levels] + + # Mirror extract_region: build a boundaries enum from labels that are + # actually present (non-empty stratum) on the coarse DM. The probe + # showed gmsh boundary labels survive every hierarchy level. + surviving = {} + if fine_mesh.boundaries is not None: + for b in fine_mesh.boundaries: + if b.name in ("Null_Boundary", "All_Boundaries"): + continue + lab = coarse_dm.getLabel(b.name) + if lab: + sis = lab.getStratumIS(b.value) + if sis and sis.getSize() > 0: + surviving[b.name] = b.value + sub_boundaries = Enum("Boundaries", surviving) if surviving else None + + # Construct the companion. refinement=None -> the constructor's + # no-refine branch: self.dm IS coarse_dm, dm_hierarchy=[coarse_dm]. + companion = uw.discretisation.Mesh( + coarse_dm, + degree=fine_mesh.degree, + qdegree=fine_mesh.qdegree, + boundaries=sub_boundaries, + coordinate_system_type=fine_mesh.CoordinateSystemType, + verbose=verbose, + ) + + # Submesh lineage -- same shape as extract_region + companion.parent = fine_mesh + companion._parent_mesh_version = fine_mesh._mesh_version + companion.regions = fine_mesh.regions + companion._dof_maps = {} + # Keyed weakly by the variable objects (not id()) so a GC'd variable + # can't alias a stale entry with a mismatched FE layout. Nested: + # coarse_var -> { fine_var -> cached tuple }. + # _interp_cache value: (matInterp, vecScale, dm_c, dm_f) + # _inject_cache value: (injectMat, dm_c, dm_f) + companion._interp_cache = weakref.WeakKeyDictionary() + companion._inject_cache = weakref.WeakKeyDictionary() + companion._companion_levels = levels + companion._is_coarsened_companion = True + fine_mesh._registered_submeshes.add(companion) + + return companion + + +# --------------------------------------------------------------------------- +# Per-variable transfer operators (PETSc-native, no KDTree) +# --------------------------------------------------------------------------- + +_FE_TAG = 0 + + +def _single_field_dm(mesh, var): + """Clone ``mesh.dm`` and install only ``var``'s FE as field 0. + + ``createSubDM(field_id)`` yields a section-only DM with no PetscFE, so + ``createInterpolation``/``createInjection`` (FEM operators) fail with + PETSC_ERR_SUP on it. Mirror the ``_get_coords_for_basis`` pattern + (discretisation_mesh.py:2749-2777): clone the plex, attach the single + PetscFE, build the DS. The clone keeps topology + coordinates so the + FEM interpolator/injector can operate. + """ + global _FE_TAG + _FE_TAG += 1 + pfx = f"rpt{_FE_TAG}_" + opts = PETSc.Options() + opts.setValue(f"{pfx}petscspace_degree", var.degree) + opts.setValue(f"{pfx}petscdualspace_lagrange_continuity", var.continuous) + opts.setValue(f"{pfx}petscdualspace_lagrange_node_endpoints", False) + + fe = PETSc.FE().createDefault( + mesh.dm.getDimension(), + var.num_components, + mesh.isSimplex, + mesh.qdegree, + pfx, + PETSc.COMM_SELF, + ) + dm = mesh.dm.clone() + dm.clearDS() + dm.setField(0, fe) + dm.createDS() + return dm + + +def _var_global_vec(var, sfdm): + """Variable data -> single-field-DM global Vec. + + ``sfdm`` is the single-field clone (same FE, same topology as the + variable's own field), so its local layout matches ``var.vec``. + """ + var._set_vec(available=True) + g = sfdm.createGlobalVector() + loc = sfdm.getLocalVec() + loc.array[...] = var.vec.array + sfdm.localToGlobal(loc, g, addv=False) + sfdm.restoreLocalVec(loc) + return g + + +def _write_global_vec_to_var(var, sfdm, g): + """Single-field-DM global Vec -> variable data. + + Invalidates the same higher-level caches the core code clears on a + direct PETSc write (see discretisation_mesh.py _re_extract_from_parent): + ``.array`` / ``.data`` are copies built from the PETSc vec, so they go + stale unless cleared here. + """ + loc = sfdm.getLocalVec() + sfdm.globalToLocal(g, loc, addv=False) + var._set_vec(available=True) + var.vec.array[...] = loc.array + var._lvec.array[...] = var.vec.array + sfdm.restoreLocalVec(loc) + for attr in ("_canonical_data", "_cached_data_array"): + if hasattr(var, attr): + setattr(var, attr, None) + + +def _linked_pair(companion, coarse_var, fine_var): + """Build a refinement-linked single-field (dm_c, dm_f) pair. + + dm_c is a single-field clone of the coarse companion DM. dm_f is + produced by refining dm_c ``levels`` times -- so refine() sets the + coarse/fine linkage and the regularRefinement flag, and + createInterpolation/createInjection take the *nested* exact path + (plex.c:10328) with no geometric point location. + + Because dm_f is the uniform refinement of the same coarse topology + that produced the fine mesh, its DOF ordering matches the fine + variable's storage (verified by the round-trip test). + """ + levels = companion._companion_levels + dm_c = _single_field_dm(companion, coarse_var) + dm_f = dm_c + for _ in range(levels): + nxt = dm_f.refine() + nxt.setCoarseDM(dm_f) + dm_f = nxt + # refine() carries topology; (re)build the DS so the field/section + # exist on the refined DM for the FEM operators. + dm_f.createDS() + return dm_c, dm_f + + +def _cache_lookup(cache, coarse_var, fine_var): + inner = cache.get(coarse_var) + return inner.get(fine_var) if inner is not None else None + + +def _cache_store(cache, coarse_var, fine_var, value): + inner = cache.get(coarse_var) + if inner is None: + inner = weakref.WeakKeyDictionary() + cache[coarse_var] = inner + inner[fine_var] = value + return value + + +def _get_interpolation(companion, coarse_var, fine_var): + """Build & cache the nested FE prolongation coarse -> fine. + + Cached value: (matInterp, vecScale, dm_c, dm_f). + """ + hit = _cache_lookup(companion._interp_cache, coarse_var, fine_var) + if hit is not None: + return hit + + dm_c, dm_f = _linked_pair(companion, coarse_var, fine_var) + matInterp, vecScale = dm_c.createInterpolation(dm_f) + return _cache_store( + companion._interp_cache, coarse_var, fine_var, + (matInterp, vecScale, dm_c, dm_f), + ) + + +def _get_injection(companion, coarse_var, fine_var): + """Build & cache the injection scatter (MATSCATTER) coarse <-> fine. + + Cached value: (injectMat, dm_c, dm_f). + """ + hit = _cache_lookup(companion._inject_cache, coarse_var, fine_var) + if hit is not None: + return hit + + dm_c, dm_f = _linked_pair(companion, coarse_var, fine_var) + injectMat = dm_c.createInjection(dm_f) # MATSCATTER wrapping a VecScatter + return _cache_store( + companion._inject_cache, coarse_var, fine_var, + (injectMat, dm_c, dm_f), + ) + + +def prolongate(companion, coarse_var, fine_var): + """coarse -> fine, FE prolongation (fills all fine DOFs).""" + matInterp, _, dm_c, dm_f = _get_interpolation(companion, coarse_var, fine_var) + gc = _var_global_vec(coarse_var, dm_c) + gf = dm_f.createGlobalVector() + try: + matInterp.mult(gc, gf) + _write_global_vec_to_var(fine_var, dm_f, gf) + finally: + gc.destroy() + gf.destroy() + + +def restrict(companion, fine_var, coarse_var, weighted=True): + """fine -> coarse, Galerkin restriction (transpose of prolongation).""" + matInterp, vecScale, dm_c, dm_f = _get_interpolation( + companion, coarse_var, fine_var + ) + gf = _var_global_vec(fine_var, dm_f) + gc = dm_c.createGlobalVector() + try: + matInterp.multTranspose(gf, gc) + if weighted and vecScale is not None: + gc.pointwiseMult(gc, vecScale) + _write_global_vec_to_var(coarse_var, dm_c, gc) + finally: + gf.destroy() + gc.destroy() + + +def sample(companion, fine_var, coarse_var): + """fine -> coarse, pure injection (exact at coarse-coincident DOFs).""" + injectMat, dm_c, dm_f = _get_injection(companion, coarse_var, fine_var) + gf = _var_global_vec(fine_var, dm_f) + gc = dm_c.createGlobalVector() + try: + # MATSCATTER from createInjection maps fine -> coarse via + # multTranspose; mult maps coarse -> fine. + injectMat.multTranspose(gf, gc) + _write_global_vec_to_var(coarse_var, dm_c, gc) + finally: + gf.destroy() + gc.destroy() + + +def inject(companion, coarse_var, fine_var): + """coarse -> fine, scatter onto coarse-coincident fine DOFs only. + + Fine DOFs introduced by refinement are left untouched. + """ + injectMat, dm_c, dm_f = _get_injection(companion, coarse_var, fine_var) + gc = _var_global_vec(coarse_var, dm_c) + gf = dm_f.createGlobalVector() + try: + gf.set(0.0) + injectMat.mult(gc, gf) + _write_global_vec_to_var(fine_var, dm_f, gf) + finally: + gc.destroy() + gf.destroy() diff --git a/docs/examples/submesh_investigation/test_refined_pair_contract.py b/docs/examples/submesh_investigation/test_refined_pair_contract.py new file mode 100644 index 00000000..86efe192 --- /dev/null +++ b/docs/examples/submesh_investigation/test_refined_pair_contract.py @@ -0,0 +1,76 @@ +"""Contract test: refine-DM mode only. + +The coarse/fine companion is available ONLY when a genuine nested +refinement hierarchy exists. With no refinement relationship the +operation must fail loudly -- no geometric or KDTree fallback. + + 1. Mesh built WITHOUT refinement -> coarsened_companion raises. + 2. Mesh built WITH refinement -> companion works, and the + transfer is genuinely the nested path (proved by the gating test, + which hits machine precision with -dm_plex_hash_location absent; + a silent General fallback would have errored on the missing hash + option). + +Run: + pixi run -e amr-dev python -u \ + docs/examples/submesh_investigation/test_refined_pair_contract.py +""" + +import underworld3 as uw + +import refined_pair_prototype as rpp + + +def banner(msg): + uw.pprint(0, f"\n{'='*70}\n{msg}\n{'='*70}") + + +def main(): + banner("CONTRACT 1: no refinement relationship -> unavailable") + flat = uw.meshing.UnstructuredSimplexBox( + minCoords=(0.0, 0.0), + maxCoords=(1.0, 1.0), + cellSize=0.2, + ) + uw.pprint(0, f"flat mesh dm_hierarchy length: {len(flat.dm_hierarchy)}") + raised = False + try: + rpp.coarsened_companion(flat, levels=1) + except ValueError as e: + raised = True + msg = str(e) + uw.pprint(0, f"raised ValueError as required:\n {msg}") + assert "refinement" in msg.lower(), "error must cite refinement" + assert "fallback" in msg.lower(), "error must state no fallback" + assert raised, "coarsened_companion must raise on a non-refined mesh" + + banner("CONTRACT 2: refinement present -> companion available") + fine = uw.meshing.UnstructuredSimplexBox( + minCoords=(0.0, 0.0), + maxCoords=(1.0, 1.0), + cellSize=0.2, + degree=2, + qdegree=4, + refinement=2, + ) + uw.pprint(0, f"fine mesh dm_hierarchy length: {len(fine.dm_hierarchy)}") + coarse = rpp.coarsened_companion(fine, levels=1) + assert coarse.parent is fine + assert coarse in fine._registered_submeshes + cS, cE = coarse.dm.getHeightStratum(0) + uw.pprint(0, f"companion available: {cE - cS} cells, parent linked") + + # Asking for more levels than exist must also fail loudly. + raised = False + try: + rpp.coarsened_companion(fine, levels=99) + except ValueError as e: + raised = True + uw.pprint(0, f"deep request correctly refused: {e}") + assert raised, "request beyond hierarchy depth must raise" + + banner("CONTRACT TEST PASSED") + + +if __name__ == "__main__": + main() diff --git a/docs/examples/submesh_investigation/test_refined_pair_solver.py b/docs/examples/submesh_investigation/test_refined_pair_solver.py new file mode 100644 index 00000000..f735f1c5 --- /dev/null +++ b/docs/examples/submesh_investigation/test_refined_pair_solver.py @@ -0,0 +1,128 @@ +"""Gating test: extract a coarse companion, keep labels, solve Stokes, +map the solution back to the fine parent mesh. + +This is the fork in the road for the refined-submesh-pair investigation. +If Stokes won't solve on the coarse companion, or the solution can't be +mapped back, the rest of the investigation doesn't matter. + +Stages (each logged): + 1. Build fine box mesh (refinement=2) -> has dm_hierarchy + 2. coarsened_companion(levels=1) -> real uw.Mesh, parent state + 3. Assert labels intact + submesh lineage + 4. Stokes on the companion (no-slip walls + buoyancy body force) + 5. Map velocity back to the fine mesh via prolongate (nested FE) + 6. Round-trip sanity: sample fine -> coarse ~= original coarse + +Run: + pixi run -e amr-dev python -u \ + docs/examples/submesh_investigation/test_refined_pair_solver.py +""" + +import numpy as np +import sympy + +import underworld3 as uw +from underworld3.systems import Stokes + +import refined_pair_prototype as rpp + + +def banner(msg): + uw.pprint(0, f"\n{'='*70}\n{msg}\n{'='*70}") + + +def main(): + banner("STAGE 1: build fine box mesh, refinement=2") + fine = uw.meshing.UnstructuredSimplexBox( + minCoords=(0.0, 0.0), + maxCoords=(1.0, 1.0), + cellSize=0.2, + degree=2, + qdegree=4, + refinement=2, + ) + cS, cE = fine.dm.getHeightStratum(0) + uw.pprint(0, f"fine mesh: {cE - cS} cells, hierarchy len {len(fine.dm_hierarchy)}") + + banner("STAGE 2: coarsened_companion(levels=1)") + coarse = rpp.coarsened_companion(fine, levels=1, verbose=False) + cS, cE = coarse.dm.getHeightStratum(0) + uw.pprint(0, f"coarse companion: {cE - cS} cells") + + banner("STAGE 3: assert labels intact + submesh lineage") + assert coarse.parent is fine, "companion.parent should be the fine mesh" + assert coarse in fine._registered_submeshes, "not registered with parent" + bnames = [b.name for b in coarse.boundaries] + uw.pprint(0, f"companion boundaries enum: {bnames}") + for name in ("Bottom", "Top", "Left", "Right"): + lab = coarse.dm.getLabel(name) + assert lab, f"boundary label {name} missing on companion DM" + sis = lab.getStratumIS(coarse.boundaries[name].value) + size = sis.getSize() if sis else 0 + uw.pprint(0, f" {name}: stratum size {size}") + assert size > 0, f"boundary {name} empty on companion" + uw.pprint(0, "lineage + labels OK") + + banner("STAGE 4: Stokes on the coarse companion") + v = uw.discretisation.MeshVariable("Vc", coarse, coarse.dim, degree=2) + p = uw.discretisation.MeshVariable("Pc", coarse, 1, degree=1) + + stokes = Stokes(coarse, velocityField=v, pressureField=p) + stokes.constitutive_model = uw.constitutive_models.ViscousFlowModel + stokes.constitutive_model.Parameters.shear_viscosity_0 = 1.0 + + # Buoyancy-style body force, no-slip on all four walls. + x, y = coarse.X + stokes.bodyforce = sympy.Matrix([0.0, -sympy.sin(sympy.pi * x)]) + + for wall in ("Bottom", "Top", "Left", "Right"): + stokes.add_dirichlet_bc((0.0, 0.0), wall) + + stokes.petsc_options["snes_type"] = "newtonls" + stokes.petsc_options["snes_rtol"] = 1.0e-6 + stokes.petsc_options["ksp_rtol"] = 1.0e-8 + stokes.petsc_options["pc_type"] = "lu" + + stokes.solve(verbose=False) + + vmag = np.linalg.norm(v.data, axis=1) + uw.pprint( + 0, + f"coarse solve done: |v| min={vmag.min():.3e} " + f"max={vmag.max():.3e} mean={vmag.mean():.3e}", + ) + assert np.all(np.isfinite(v.data)), "non-finite velocity on coarse solve" + assert vmag.max() > 1.0e-6, "coarse velocity is trivially zero" + + banner("STAGE 5: map velocity back to the fine mesh (prolongate)") + v_fine = uw.discretisation.MeshVariable("Vf", fine, fine.dim, degree=2) + rpp.prolongate(coarse, v, v_fine) + + vf = v_fine.data + vfmag = np.linalg.norm(vf, axis=1) + uw.pprint( + 0, + f"prolongated to fine: |v| min={vfmag.min():.3e} " + f"max={vfmag.max():.3e} mean={vfmag.mean():.3e}", + ) + assert np.all(np.isfinite(vf)), "non-finite velocity after prolongate" + assert vfmag.max() > 1.0e-6, "prolongated velocity trivially zero" + # Prolongated max should be close to the coarse max (nested FE does not + # overshoot for these smooth fields). + rel = abs(vfmag.max() - vmag.max()) / vmag.max() + uw.pprint(0, f"max|v| coarse vs fine relative diff: {rel:.3e}") + + banner("STAGE 6: round-trip sanity (sample fine -> coarse)") + v_back = uw.discretisation.MeshVariable("Vb", coarse, coarse.dim, degree=2) + rpp.sample(coarse, v_fine, v_back) + err = np.linalg.norm(v_back.data - v.data) / np.linalg.norm(v.data) + uw.pprint(0, f"round-trip relative error (sample o prolongate): {err:.3e}") + assert err < 1.0e-8, ( + f"round-trip via sample should recover coarse exactly, got {err:.3e}" + ) + + banner("ALL STAGES PASSED") + + +if __name__ == "__main__": + main()