diff --git a/docs/developer/design/submesh-solver-architecture.md b/docs/developer/design/submesh-solver-architecture.md index 5de1793d..70ea7e51 100644 --- a/docs/developer/design/submesh-solver-architecture.md +++ b/docs/developer/design/submesh-solver-architecture.md @@ -21,27 +21,40 @@ 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 +## Three Submesh Flavours: Subdomain, Resolution Level, Surface 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 +and supports explicit field transfer back and forth. There are three 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. +| | Subdomain | Resolution level | Surface | +|---|---|---|---| +| Constructor | `mesh.extract_region("Inner")` | `coarsened_companion(fine, levels=1)` | `extract_surface(mesh, "Upper")` | +| PETSc mechanism | `DMPlexFilter` (cell subset) | `dm.refine()` nested hierarchy | `DMPlexCreateSubmesh` then `DMPlexFilter(depth, 2)` (strip phantom DAG) | +| dim, cdim | `parent.dim`, `parent.cdim` | `parent.dim`, `parent.cdim` | `parent.dim − 1`, `parent.cdim` | +| Parent ↔ child map | `getSubpointIS()` (point-level) | nested `createInjection` / `createInterpolation` (DOF-level) | `getSubpointIS()` (point-level) | +| Transfer fidelity | exact (shared nodes) | exact (nested FE) | exact (shared nodes) | +| Example | `docs/examples/submesh_investigation/test_region_ds_submesh.py` | `docs/examples/submesh_investigation/example_refined_companion.py` | `docs/examples/submesh_investigation/example_surface_extraction.py` | + +The subdomain and resolution-level examples solve the **same** annulus + +radial-buoyancy Stokes problem so those two 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. + +The surface flavour produces a 2-manifold embedded in 3-space +(`dim=2, cdim=3`) and is the natural carrier for *lateral surface +processes*: surface diffusion, surface advection, surface evolution. +Phase 1 (the lineage layer: extract, restrict/prolongate round-trip, +registration with the parent, navigation, `evaluate(expr, surface_pts)`) +lands as the investigation under +`docs/examples/submesh_investigation/surface_submesh_prototype.py`. +Phase 2 — making a solver actually run on the manifold (Laplace– +Beltrami) — is tracked in the *Surface submesh: solver path* section +below. ### Design contract for the resolution-level flavour: refine-DM mode only @@ -108,6 +121,152 @@ Status: prototype + both parallel examples + gating tests promotion of `coarsened_companion` into the `Mesh` API proper is a follow-up. +### Design contract for the surface flavour: two-stage strip, no fallback + +A surface submesh is extractable **only when the parent carries a +non-empty boundary-face stratum** for the named label. Mirrors the +loud-failure stance of the other two flavours: `extract_surface` raises +on an unknown label, on a label whose value isn't in the parent's live +value set, or on an empty stratum. There is deliberately no geometric +fallback that might silently produce a degenerate mesh. + +The PETSc construction is **two-stage** — both primitives are already +wrapped in UW3's cython layer: + +``` +sub1 = petsc_dm_create_submesh_from_label(parent.dm, label, value, marked_faces=True) +sub2 = petsc_dm_filter_by_label(sub1, "depth", 2) # for parent.dim == 3 +subpoint_is = compose(sub2.getSubpointIS(), sub1.getSubpointIS()) +``` + +Stage 1 (`DMPlexCreateSubmesh`) gives a cd-1 DM with the right cells, +edges and vertices, but PETSc additionally retains an **upward-DAG +phantom stratum**: depth-3 (= parent.dim) points, one per parent +volume cell, celltype 12, included inside each surface cell's downward +closure tuple. This linkage is there so the resulting DM can still be +navigated as "a slice of the parent" (assemble surface integrals from +parent volume cells). For a standalone surface mesh we don't use that +linkage, and the phantom points are actively harmful — every consumer +that does `closure[-n:]` slicing to grab vertices (the kd-tree +builder, centroid pickers, control-point face markers — three copies +of the same idiom in `discretisation_mesh.py` alone) gets sometimes +the right answer and sometimes a phantom point inside the slice. + +Stage 2 (`DMPlexFilter(depth, 2)`) keeps only the cells of `sub1` and +their downward closure (edges, vertices). The result is a **genuine +standalone 2-manifold mesh** with a clean 3-stratum chart +(`depth = 0, 1, 2`; celltypes `[0, 1, 3]`). Cell closures are length 7 +(1 cell + 3 edges + 3 vertices) and `closure[-3:]` reliably returns +vertices. The standard `Mesh._build_kd_tree_index` and downstream +navigation code run on this DM without any cd-1 special-casing. + +The composed subpoint IS gives a direct surface → parent point map +(point IDs in the parent's chart). The properties carry through: + +- `surf.dim = parent.dim − 1`, `surf.cdim = parent.cdim` (2 and 3 for + the sphere case); +- vertex coordinates land on the parent surface to machine precision + (we measured 2.2e-16 on `SphericalShell.Upper`); +- parent boundary labels propagate onto the submesh: the *selecting* + label (e.g. `"Upper"`) survives carrying the surface itself; + sibling boundary labels (`"Lower"`) survive as empty strata which + the submesh constructor filters out. + +**One PETSc footgun, worth knowing about.** Calling +`label.getStratumIS(value)` on a `DMPlexCreateSubmesh` DM for a value +that isn't in `getValueIS()` hard-aborts PETSc (no Python-catchable +exception). The prototype enumerates `surface_dm.getNumLabels()` by +index and checks each label's live value set first. + +**Navigation on the manifold.** For an on-surface query point (the +contract assumption — query points come *from* the manifold), the +centroid kd-tree ranks faces by chord distance, which agrees with +the owning-face ranking to first order on any convex manifold (sphere, +ellipsoid). `get_closest_cells(surface_dof_coords)` returns valid +surface cell IDs and `evaluate(expr, surface_pts)` works to machine +precision. The cell-containment-by-half-space-rule helpers +(`_test_if_points_in_cells_internal`, `_mark_faces_inside_and_out`) +use a 2-D perpendicular construction that doesn't make sense for a +curved cell, so the cd-1 mesh path skips that rejection step and +trusts the centroid kd-tree result. A proper manifold in-cell test +(project the query into the cell's tangent plane, then barycentric) +is Phase 2 work — only needed if we ever support off-surface queries. + +```python +shell = uw.meshing.SphericalShell(radiusOuter=1.0, radiusInner=0.5, + cellSize=0.2) +surface = extract_surface(shell, "Upper") # parent = shell + +# Round-trip a parent scalar via shared-vertex KDTree (1e-10 match) +T_p = uw.discretisation.MeshVariable("T_p", shell, 1, degree=1) +T_s = uw.discretisation.MeshVariable("T_s", surface, 1, degree=1) +surface.restrict(T_p, T_s) # parent -> surface, bit-exact +# … work on the surface … +surface.prolongate(T_s, T_p) # surface -> parent, bit-exact at surface DOFs + +# Symbolic expressions can be evaluated at surface points directly +vals = uw.function.evaluate(T_s.sym[0], surface.X.coords) # to machine precision +``` + +Status (Phase 1): prototype + documented example + +contract test +(`docs/examples/submesh_investigation/{surface_submesh_prototype.py, +example_surface_extraction.py, test_surface_submesh_contract.py}`) all +passing. + +### Surface submesh: solver path (Phase 2) + +Phase 1 produces the *mesh*. Phase 2 makes a solver run on it. UW3's +solver stack was built when `dim == cdim` was implicit; a surface +submesh — `dim = parent.dim − 1`, `cdim = parent.cdim` — exercises +that path from a new angle. + +Concrete deliverable: lateral surface diffusion (Laplace–Beltrami, +`∫_M ∇_M T · ∇_M w dA`) on the upper surface of a `SphericalShell`, +verified against the analytic spherical-harmonic decay +`exp(-l(l+1) κ t / R²)`. Either a `SurfaceDiffusion`/`LaplaceBeltrami` +sibling of `Poisson`, or a flag on `Poisson` — decide after the JIT +audit. The Phase 1 example +(`example_surface_extraction.py`) is the natural site to add this once +the solver path is operational. + +The investigation needs to clear (at least) these knowns: + +1. **JIT pointwise functions on `dim != cdim`.** `petsc_x[i]` indexing + should run to `cdim`; gradient-of-basis indexing to `dim` + (tangent-space). Audit `utilities/_jitextension.py` for any + `range(self.dim)` that should be `range(self.cdim)` (or vice + versa). Phase 1 status: navigation and `evaluate()` work via the + standard volume-mesh code path on the stripped chart — no JIT + change was needed for that. Assembly is the open question. +2. **FE assembly with non-square Jacobian.** PETSc DMPlex computes + element Jacobians from the embedded coordinates and produces the + correct surface metric automatically *when the FE machinery is + exercised with* `dim != cdim`. Confirm this on a trivial bilinear + form before assuming it. +3. **Coordinate-system symbols.** `mesh.X` on the surface submesh is a + 3-vector; `mesh.dim` is 2. Code that iterates `range(mesh.dim)` + over `mesh.X` components silently misses the third — exactly the + kind of cdim/dim conflation the audit needs to catch. +4. **Manifold in-cell test.** `_test_if_points_in_cells_internal` / + `_mark_faces_inside_and_out` use a 2-D perpendicular construction + to place face-relative control points and run the half-space test. + The cd-1 path currently skips that step (centroid kd-tree only). + A proper version projects the query into the cell's tangent plane + then runs barycentric. Only needed if we ever support off-surface + queries; on-surface queries already work via the centroid path. +5. **Surface outward normal.** For intrinsic surface diffusion the + bilinear form is metric-only, no explicit normal needed. For + coupled problems (surface advection driven by a 3D flow), we need + the surface's outward unit normal in 3-space — distinct from + `mesh.Gamma_N` which is the normal to the *boundary* of the mesh + (a closed sphere has none). Possibly a new symbol + (`mesh.surface_normal`); deferred until needed. +6. **Boundary conditions on a closed surface.** `SphericalShell.Upper` + has no boundary curve, so no Dirichlet/Neumann is needed for the + first pass. A partial-surface submesh would add another dim/cdim + layer (BCs on a 1D curve in 3D); out of scope for the first pass. + ## Design Principles ### 1. Separate meshes, separate variables, explicit copies @@ -182,7 +341,7 @@ The label names are preserved from the parent (they survive `DMPlexFilter`). The | `DMPlex.getSubpointIS()` | IS mapping submesh → parent points | Available in petsc4py | | `DMSetRegionDS(dm, label, fields, ds, dsIn)` | Per-region discrete system | **Segfaults**, no examples | | `DMGetCellDS(dm, point, &ds, &dsIn)` | Per-cell DS dispatch in assembly | Works but requires Region DS | -| `DMPlexCreateSubmesh(dm, label, value, ...)` | Co-dimension 1 submesh (boundaries) | Works but wrong dimension | +| `DMPlexCreateSubmesh(dm, label, value, ...)` | Co-dimension 1 submesh (boundaries) | **Works**, used by `extract_surface` (dim=parent.dim−1, cdim=parent.cdim) | | `VecScatter` / `PetscSF` | Parallel data transfer | Standard PETSc | ### PETSc Alternatives Investigated (2026-04-05) diff --git a/docs/examples/manifold_pdes/example_manifold_diffusion.py b/docs/examples/manifold_pdes/example_manifold_diffusion.py new file mode 100644 index 00000000..f29d99ec --- /dev/null +++ b/docs/examples/manifold_pdes/example_manifold_diffusion.py @@ -0,0 +1,208 @@ +"""Time-dependent scalar diffusion on a SphericalManifold. + +The canonical Laplace–Beltrami demonstrator: a pole-localised +Gaussian on a unit sphere relaxes toward its spatial mean. Each +spherical-harmonic component decays at its own analytic rate + + T_lm(t) = T_lm(0) · exp( -l(l+1) κ t / R^2 ) + +so high-l (small-scale) features die first and the low-l components +linger longest. Numerically the field flattens monotonically and +the L2 deviation from the mean decays exponentially. + +Three pieces: + (A) Time-step the Eulerian ``Diffusion`` solver forward, snapshot + the field at several times, write a pyvista frame for each. + (B) Track ``L2(T - )`` against the dominant Y_10 decay rate + ``exp(-2 κ t)`` (for the centered Gaussian, l=1 dominates). + (C) Assemble a 2×3 grid of snapshots into a single figure. + +Run: + pixi run -e amr-dev python -u \\ + docs/examples/manifold_pdes/example_manifold_diffusion.py + +Notes +----- +* ``Diffusion(...)`` is a linear time-step solve. Configuring it with + ``snes_type = "ksponly"`` skips Newton's line-search machinery so + the per-step iteration is a single linear KSP — clean output, no + spurious ``DIVERGED_LINE_SEARCH`` warnings from a residual that's + already at quadrature noise level. +* ``petsc_use_constant_nullspace = True`` handles the constant + Laplace–Beltrami kernel on the closed manifold (conservative + diffusion preserves the spatial mean exactly in the continuum; + PETSc's nullspace projection enforces it discretely). +""" + +from __future__ import annotations + +import os + +import numpy as np + +import underworld3 as uw +from underworld3.systems import Diffusion + + +# --------------------------------------------------------------------------- +# Set up +# --------------------------------------------------------------------------- + +CELL_SIZE = 0.15 +DIFFUSIVITY = 0.05 +DT = 0.1 +N_STEPS = 40 +SNAPSHOT_STEPS = (0, 4, 10, 20, 30, 40) +SIGMA = 0.4 + + +def build_initial_gaussian(T, sphere): + """Plant a Gaussian centred at (1, 0, 0) with std-dev ``SIGMA`` + (geodesic units on the unit sphere). Subtracts the discrete mean + so the field is mean-zero — the conservative diffusion preserves + that exactly.""" + coords = np.asarray(T.coords) + arc = np.arccos(np.clip(coords[:, 0], -1.0, 1.0)) + blob = np.exp(-arc ** 2 / (2.0 * SIGMA ** 2)) + new_T = np.zeros_like(T.data) + new_T[:, 0] = blob - blob.mean() + T.pack_raw_data_to_petsc(new_T, sync=True) + + +def build_solver(sphere, T): + diff = Diffusion(sphere, u_Field=T, order=1, theta=1.0) + diff.constitutive_model = uw.constitutive_models.DiffusionModel + diff.constitutive_model.Parameters.diffusivity = DIFFUSIVITY + diff.petsc_use_constant_nullspace = True + # Linear time-step solve — skip Newton's line search entirely. + diff.petsc_options["snes_type"] = "ksponly" + return diff + + +# --------------------------------------------------------------------------- +# pyvista — one snapshot per step + composite grid +# --------------------------------------------------------------------------- + +def write_snapshot(sphere, T, step_idx, t, out_dir, lims): + try: + import pyvista as pv + import underworld3.visualisation as vis + except ImportError: + return None + + pvmesh = vis.mesh_to_pv_mesh(sphere) + pvmesh.point_data["T"] = vis.scalar_fn_to_pv_points(pvmesh, T.sym) + + plotter = pv.Plotter(window_size=(700, 600), off_screen=True) + plotter.set_background("white") + plotter.add_mesh( + pvmesh, + scalars="T", + cmap="RdBu_r", + clim=lims, + show_edges=False, + smooth_shading=False, + scalar_bar_args={"title": "T", "color": "black"}, + ) + plotter.add_text(f"t = {t:.2f}", position="upper_left", color="black", font_size=12) + plotter.camera_position = "yz" + plotter.camera.azimuth = 30 + plotter.camera.elevation = 20 + + out_path = os.path.join(out_dir, f"_diffusion_step{step_idx:03d}.png") + plotter.screenshot(out_path) + plotter.close() + return out_path + + +def make_composite(snapshot_paths, out_path): + """Stitch the snapshots into a single 2 × 3 grid via PIL.""" + try: + from PIL import Image + except ImportError: + print(" PIL not available — skipping composite.") + return + imgs = [Image.open(p) for p in snapshot_paths] + w, h = imgs[0].size + rows, cols = 2, 3 + grid = Image.new("RGB", (cols * w, rows * h), "white") + for i, img in enumerate(imgs): + r, c = divmod(i, cols) + grid.paste(img, (c * w, r * h)) + grid.save(out_path) + + +# --------------------------------------------------------------------------- +# main +# --------------------------------------------------------------------------- + +def main(): + print("=" * 72) + print("Pole-localised Gaussian → relaxation toward zero mean on a") + print(f"unit SphericalManifold (cellSize={CELL_SIZE}, κ={DIFFUSIVITY}, dt={DT})") + print("=" * 72) + + out_dir = os.path.dirname(__file__) + sphere = uw.meshing.SphericalManifold(radius=1.0, cellSize=CELL_SIZE) + + T = uw.discretisation.MeshVariable("T", sphere, 1, degree=2) + build_initial_gaussian(T, sphere) + diff = build_solver(sphere, T) + + # Color limits fixed by the initial field so the colormap is + # comparable across snapshots. + lims = (float(T.data.min()), float(T.data.max())) + print(f"\nIC: T range = [{lims[0]:.4f}, {lims[1]:.4f}], " + f"mean = {T.data[:, 0].mean():.6f}") + + # Decay diagnostic: ||T - ||_2 over time. + snapshot_paths = [] + times = [] + rms_history = [] + + if 0 in SNAPSHOT_STEPS: + p = write_snapshot(sphere, T, 0, 0.0, out_dir, lims) + if p is not None: + snapshot_paths.append(p) + times.append(0.0) + rms_history.append(float(np.sqrt((T.data[:, 0] ** 2).mean()))) + + print(f"\n{'step':>5} {'t':>8} {'||T||_2':>14} {'max':>14} " + f"{'Y_10 decay':>14}") + print("-" * 65) + print(f"{0:>5} {0.0:>8.2f} {rms_history[0]:>14.4e} " + f"{T.data.max():>14.4e} {1.0:>14.4e}") + + for step in range(1, N_STEPS + 1): + diff.solve(timestep=DT) + t = step * DT + rms = float(np.sqrt((T.data[:, 0] ** 2).mean())) + # Y_10 has eigenvalue 2; ||Y_10(t)|| ~ exp(-2 κ t) + y10_factor = float(np.exp(-2.0 * DIFFUSIVITY * t)) + rel = rms / rms_history[0] + times.append(t) + rms_history.append(rms) + + if step % 4 == 0 or step in SNAPSHOT_STEPS: + print(f"{step:>5} {t:>8.2f} {rms:>14.4e} " + f"{T.data.max():>14.4e} {y10_factor:>14.4e}") + + if step in SNAPSHOT_STEPS: + p = write_snapshot(sphere, T, step, t, out_dir, lims) + if p is not None: + snapshot_paths.append(p) + + print("-" * 65) + print("Y_10 column is the analytic decay envelope of the l=1 mode;") + print("the numerical ||T||_2 / ||T||_0 should approach it as high-l") + print("modes die out and the l=1 component dominates the residual.") + + # Composite figure + composite_path = os.path.join(out_dir, "_diffusion_composite.png") + if snapshot_paths: + make_composite(snapshot_paths, composite_path) + print(f"\nComposite snapshot grid: {composite_path}") + + +if __name__ == "__main__": + main() diff --git a/docs/examples/manifold_pdes/example_manifold_helmholtz.py b/docs/examples/manifold_pdes/example_manifold_helmholtz.py new file mode 100644 index 00000000..3c749375 --- /dev/null +++ b/docs/examples/manifold_pdes/example_manifold_helmholtz.py @@ -0,0 +1,269 @@ +"""Manifold-mesh scalar PDE demonstrations on a SphericalManifold. + +Three demonstrations of scalar PDE solves on a 2-manifold sphere +mesh (``dim=2, cdim=3``): + + (A) **Spherical-harmonic spectrum recovery** (steady Helmholtz). + Solve :math:`(-\\Delta_S + I) T = Y_{lm}` for several ``(l, m)`` + and verify ``T = Y_{lm} / (l(l+1) + 1)`` to FE accuracy. + + (B) **h-convergence sweep**. Same problem (Y_10), refine cellSize + from 0.4 → 0.05, fit the convergence rate. + + (C) **Closed-manifold Poisson with nullspace**. Solve + :math:`-\\Delta_S T = z` on the closed unit sphere. The + Laplace–Beltrami operator has a constant nullspace; PETSc + handles it via ``petsc_use_constant_nullspace``. Recovers + ``T = z/2 + const`` to FE accuracy. + +A pyvista figure is rendered at the end showing T on the sphere +surface for the Y_10 Poisson case, with a few iso-contour lines. + +Run: + pixi run -e amr-dev python -u \\ + docs/examples/manifold_pdes/example_manifold_helmholtz.py + +Optional headless rendering: + PYVISTA_OFF_SCREEN=true pixi run -e amr-dev python -u \\ + docs/examples/manifold_pdes/example_manifold_helmholtz.py +""" + +from __future__ import annotations + +import os + +import numpy as np +import sympy + +import underworld3 as uw +from underworld3.systems import Poisson + + +# --------------------------------------------------------------------------- +# Spherical-harmonic helpers (real, orthonormal on the unit sphere) +# --------------------------------------------------------------------------- + +def real_spherical_harmonic(l: int, m: int, x, y, z): + r"""Real, unnormalised spherical-harmonic basis as a sympy expression. + + The 2-manifold sphere mesh stores Cartesian ``(x, y, z)`` coordinates, + so we express :math:`Y_{lm}(\theta, \phi)` directly as polynomials + in those — the natural form for this prototype (no charts, no + pole-singularity worries). Each function is an eigenfunction of the + Laplace–Beltrami operator on the unit sphere with eigenvalue + :math:`-l(l+1)`: + + .. math:: + \Delta_S Y_{lm} = -l(l+1) Y_{lm} + + Conventions used here: each function is a spherical harmonic up + to a normalisation constant. The L2 error against the analytic + answer is computed relative to the dynamic range of the actual + field on the mesh, so the absolute normalisation doesn't matter + for the convergence assessment. + + Supported (l, m): + (1, 0): z + (1, 1): x + (2, 0): 3 z^2 - 1 + (2, 1): x z + (2, 2): x^2 - y^2 + (3, 0): z (5 z^2 - 3) + (3, 2): z (x^2 - y^2) + """ + if (l, m) == (1, 0): + return z + if (l, m) == (1, 1): + return x + if (l, m) == (2, 0): + return 3 * z**2 - 1 + if (l, m) == (2, 1): + return x * z + if (l, m) == (2, 2): + return x**2 - y**2 + if (l, m) == (3, 0): + return z * (5 * z**2 - 3) + if (l, m) == (3, 2): + return z * (x**2 - y**2) + raise ValueError(f"unsupported (l, m) = ({l}, {m})") + + +def _eval_at_dofs(expr_fn, mesh, dof_coords): + """Evaluate a sympy expression at given physical coords via mesh.X. + + Substitutes ``mesh.X`` symbols with the coordinate columns so the + analytic spherical harmonic can be compared element-wise against + a MeshVariable's data array. + """ + x_sym, y_sym, z_sym = mesh.X + expr = expr_fn(x_sym, y_sym, z_sym) + f = sympy.lambdify((x_sym, y_sym, z_sym), expr, modules="numpy") + return f(dof_coords[:, 0], dof_coords[:, 1], dof_coords[:, 2]) + + +# --------------------------------------------------------------------------- +# (A) Spherical-harmonic spectrum recovery +# --------------------------------------------------------------------------- + +def helmholtz_recovery(sphere, mode, alpha=1.0, degree=2): + r"""Solve :math:`(-\Delta_S + \alpha) T = Y_{lm}` and return the L2 + error against the analytic answer ``T = Y_{lm} / (l(l+1) + alpha)``. + + Non-singular for ``alpha > 0`` (the operator's constant kernel is + broken by the reaction term), so no nullspace handling is needed. + """ + l, m = mode + T = uw.discretisation.MeshVariable( + f"T_l{l}m{m}", sphere, num_components=1, degree=degree, + ) + + poisson = Poisson(sphere, u_Field=T) + poisson.constitutive_model = uw.constitutive_models.DiffusionModel + poisson.constitutive_model.Parameters.diffusivity = 1.0 + + x_sym, y_sym, z_sym = sphere.X + Y_lm = real_spherical_harmonic(l, m, x_sym, y_sym, z_sym) + poisson.f = Y_lm - alpha * T.sym[0] # residual: -ΔT - Y_lm + αT = 0 + poisson.solve() + + eig = l * (l + 1) + factor = eig + alpha + coords = np.asarray(T.coords) + expected = _eval_at_dofs( + lambda x, y, z: real_spherical_harmonic(l, m, x, y, z) / factor, + sphere, coords, + ) + err = T.data[:, 0] - expected + l2 = float(np.sqrt((err ** 2).mean())) + scale = float(np.abs(expected).max()) + return l2, scale, T + + +def demo_A(sphere): + print("\n" + "=" * 72) + print("(A) Spherical-harmonic spectrum recovery on the unit sphere") + print(" -Δ_S T + T = Y_lm → T = Y_lm / (l(l+1) + 1)") + print("=" * 72) + modes = [(1, 0), (1, 1), (2, 0), (2, 1), (2, 2), (3, 0), (3, 2)] + print(f"\n{'mode':>10} {'L2 err':>12} {'rel L2':>12} {'eigenvalue':>14}") + print("-" * 54) + for mode in modes: + l2, scale, _ = helmholtz_recovery(sphere, mode) + rel = l2 / scale if scale > 0 else l2 + eig = mode[0] * (mode[0] + 1) + print(f" ({mode[0]},{mode[1]}){'':<3} {l2:>12.3e} {rel:>12.3e} {eig:>14.1f}") + + +# --------------------------------------------------------------------------- +# (B) h-convergence sweep +# --------------------------------------------------------------------------- + +def demo_B(): + print("\n" + "=" * 72) + print("(B) h-convergence sweep — (-Δ_S + I) T = Y_10 → T = z/2") + print("=" * 72) + print(f"\n{'cellSize':>10} {'cells':>8} {'L2 err':>12} {'rel L2':>12} {'order':>8}") + print("-" * 56) + prev_err = None + prev_h = None + for cs in (0.4, 0.2, 0.1, 0.05): + sphere = uw.meshing.SphericalManifold(radius=1.0, cellSize=cs) + l2, scale, _ = helmholtz_recovery(sphere, (1, 0)) + rel = l2 / scale if scale > 0 else l2 + cS, cE = sphere.dm.getHeightStratum(0) + order = ( + np.log(prev_err / l2) / np.log(prev_h / cs) if prev_err is not None else float("nan") + ) + order_str = f"{order:>8.2f}" if not np.isnan(order) else " -- " + print(f" {cs:>8} {cE-cS:>8} {l2:>12.3e} {rel:>12.3e} {order_str}") + prev_err = l2 + prev_h = cs + + +# --------------------------------------------------------------------------- +# (C) Closed Poisson with nullspace +# --------------------------------------------------------------------------- + +def demo_C(sphere): + print("\n" + "=" * 72) + print("(C) Closed-manifold Poisson with constant-nullspace handling") + print(" -Δ_S T = z → T = z/2 + const") + print("=" * 72) + T = uw.discretisation.MeshVariable("T_C", sphere, 1, degree=2) + poisson = Poisson(sphere, u_Field=T) + poisson.constitutive_model = uw.constitutive_models.DiffusionModel + poisson.constitutive_model.Parameters.diffusivity = 1.0 + poisson.petsc_use_constant_nullspace = True + x_sym, y_sym, z_sym = sphere.X + poisson.f = z_sym + poisson.solve() + + coords = np.asarray(T.coords) + expected = coords[:, 2] / 2.0 + diff = T.data[:, 0] - expected + centred = diff - diff.mean() + print(f"\n T range: [{T.data.min():.4f}, {T.data.max():.4f}]") + print(f" expected range: [{expected.min():.4f}, {expected.max():.4f}]") + print(f" L_inf(T - z/2 - ): {np.abs(centred).max():.3e}") + print(f" L_2 / |z/2|_max: {np.sqrt((centred**2).mean()) / 0.5:.3e}") + return T + + +# --------------------------------------------------------------------------- +# (D) pyvista visualisation +# --------------------------------------------------------------------------- + +def render_pyvista(sphere, T_C, out_path=None): + print("\n" + "=" * 72) + print("(D) pyvista figure") + print("=" * 72) + try: + import pyvista as pv + import underworld3.visualisation as vis + except ImportError: + print(" pyvista not available — skipping render.") + return + + pvmesh = vis.mesh_to_pv_mesh(sphere) + pvmesh.point_data["T"] = vis.scalar_fn_to_pv_points(pvmesh, T_C.sym) + + plotter = pv.Plotter(window_size=(900, 700), off_screen=True) + plotter.set_background("white") + plotter.add_mesh( + pvmesh, + scalars="T", + cmap="coolwarm", + show_edges=True, + edge_color=(0.3, 0.3, 0.3), + line_width=0.4, + smooth_shading=False, + scalar_bar_args={"title": "T = z/2", "color": "black"}, + ) + plotter.add_axes() + plotter.camera_position = "yz" + plotter.camera.azimuth = 30 + plotter.camera.elevation = 20 + + if out_path is None: + out_path = os.path.join(os.path.dirname(__file__), "_manifold_poisson.png") + plotter.screenshot(out_path) + print(f" wrote {out_path}") + + +# --------------------------------------------------------------------------- +# main +# --------------------------------------------------------------------------- + +def main(): + sphere_med = uw.meshing.SphericalManifold(radius=1.0, cellSize=0.15) + demo_A(sphere_med) + demo_B() + T_C = demo_C(sphere_med) + render_pyvista(sphere_med, T_C) + print("\n" + "=" * 72) + print("Manifold PDE demonstration complete.") + print("=" * 72) + + +if __name__ == "__main__": + main() diff --git a/docs/examples/submesh_investigation/example_surface_extraction.py b/docs/examples/submesh_investigation/example_surface_extraction.py new file mode 100644 index 00000000..10c71c64 --- /dev/null +++ b/docs/examples/submesh_investigation/example_surface_extraction.py @@ -0,0 +1,157 @@ +"""Surface submesh extraction: documented example. + +Demonstrates the third submesh flavour: extracting the upper surface of +a ``SphericalShell`` as a 2-manifold embedded in 3-space, with the same +submesh lineage state as ``Mesh.extract_region`` and +``coarsened_companion``. + +Parallel to: + test_region_ds_submesh.py (subdomain via DMPlexFilter) + example_refined_companion.py (resolution level via dm.refine()) + +This example deliberately stops at the *mesh* layer -- solving on the +surface (lateral diffusion / Laplace-Beltrami) is the Session 2 +problem. Once that lands, this example will be the natural place to +add an end-to-end "diffuse a Gaussian on the sphere" demo. + +Run: + pixi run -e amr-dev python -u \\ + docs/examples/submesh_investigation/example_surface_extraction.py +""" + +import numpy as np +import underworld3 as uw + +import surface_submesh_prototype as ssp + + +def banner(msg): + uw.pprint(0, f"\n{'='*70}\n{msg}\n{'='*70}") + + +# --------------------------------------------------------------------------- +# 1. Parent SphericalShell (3D) +# --------------------------------------------------------------------------- + +banner("Parent mesh: SphericalShell") + +r_outer, r_inner = 1.0, 0.5 +shell = uw.meshing.SphericalShell( + radiusOuter=r_outer, + radiusInner=r_inner, + cellSize=0.2, +) +cS, cE = shell.dm.getHeightStratum(0) +uw.pprint( + 0, + f"shell: dim={shell.dim}, cdim={shell.cdim}, " + f"cells={cE - cS}, boundaries={[b.name for b in shell.boundaries]}", +) + + +# --------------------------------------------------------------------------- +# 2. Extract the upper surface as a uw.Mesh +# --------------------------------------------------------------------------- + +banner("Extract upper surface (DMPlexCreateSubmesh on Upper)") + +surface = ssp.extract_surface(shell, "Upper") +sS, sE = surface.dm.getHeightStratum(0) +uw.pprint( + 0, + f"surface: dim={surface.dim}, cdim={surface.cdim}, " + f"cells={sE - sS}, vertices={surface.X.coords.shape[0]}", +) +uw.pprint(0, f" parent linked: {surface.parent is shell}") +uw.pprint(0, f" registered with parent: {surface in shell._registered_submeshes}") +uw.pprint( + 0, + f" surviving boundaries: " + f"{[b.name for b in surface.boundaries] if surface.boundaries else 'none'}", +) + + +# --------------------------------------------------------------------------- +# 3. Geometric sanity: every surface vertex sits on the parent sphere +# --------------------------------------------------------------------------- + +banner("Geometric sanity") + +radii = np.linalg.norm(surface.X.coords, axis=1) +uw.pprint( + 0, + f"vertex radii: min={radii.min():.12e}, max={radii.max():.12e}", +) +max_dev = np.abs(radii - r_outer).max() +uw.pprint(0, f"|r - r_outer| max = {max_dev:.3e} (expect < 1e-10)") +assert max_dev < 1e-10, "surface vertices not on the parent sphere" + + +# --------------------------------------------------------------------------- +# 4. Round-trip a scalar via the standard restrict/prolongate path +# --------------------------------------------------------------------------- + +banner("Scalar round-trip: parent -> surface -> parent") + +# Scalar fields on parent and surface +T_parent = uw.discretisation.MeshVariable( + "T_parent", shell, num_components=1, degree=1, +) +T_surface = uw.discretisation.MeshVariable( + "T_surface", surface, num_components=1, degree=1, +) +T_back = uw.discretisation.MeshVariable( + "T_back", shell, num_components=1, degree=1, +) + +# Plant a recognisable field on the parent: latitude as a function of z. +# (Any 3-coord function does; pick something that varies smoothly.) +parent_coords = np.asarray(T_parent.coords) # (N, 3) +new_T = np.zeros_like(T_parent.data) +new_T[:, 0] = np.arctan2( + np.sqrt(parent_coords[:, 0]**2 + parent_coords[:, 1]**2), + parent_coords[:, 2], +) +T_parent.pack_raw_data_to_petsc(new_T, sync=True) + +# parent -> surface (KDTree at 1e-10 picks the 252 surface DOFs) +surface.restrict(T_parent, T_surface) + +# surface -> a fresh parent variable +surface.prolongate(T_surface, T_back) + +# Compare on surface DOFs only (rest of parent should be zero in T_back) +nonzero = np.where(np.abs(T_back.data[:, 0]) > 0)[0] +diff = T_parent.data[nonzero, 0] - T_back.data[nonzero, 0] +uw.pprint( + 0, + f"parent surface DOFs touched: {nonzero.shape[0]} " + f"(surface has {surface.X.coords.shape[0]} vertices)", +) +uw.pprint(0, f"max |T_parent - T_back| on touched DOFs: {np.abs(diff).max():.3e}") +assert np.abs(diff).max() < 1e-12, "round-trip not bit-exact at surface DOFs" + + +# --------------------------------------------------------------------------- +# 5. Subpoint IS sanity (point-level parent <-> surface map) +# --------------------------------------------------------------------------- + +banner("Subpoint IS sanity") + +subpts = surface.subpoint_is +uw.pprint(0, f"subpoint_is size: {subpts.getSize()}") +parent_chart = shell.dm.getChart() +indices = subpts.getIndices() +uw.pprint(0, f"parent chart: {parent_chart}") +uw.pprint(0, f"subpoint index range: [{indices.min()}, {indices.max()}]") +assert indices.min() >= parent_chart[0] +assert indices.max() < parent_chart[1] + + +banner("EXAMPLE PASSED (Session 1 lineage + transfer)") +uw.pprint( + 0, + "Solving on the surface (lateral diffusion / Laplace-Beltrami) is\n" + "Session 2 -- it needs manifold-aware FE assembly through the JIT\n" + "and solver layers and is tracked in the surface-submesh plan.", +) diff --git a/docs/examples/submesh_investigation/probe_surface_extraction.py b/docs/examples/submesh_investigation/probe_surface_extraction.py new file mode 100644 index 00000000..ffa57e59 --- /dev/null +++ b/docs/examples/submesh_investigation/probe_surface_extraction.py @@ -0,0 +1,245 @@ +""" +Phase 0 probe: surface submesh extraction via DMPlexCreateSubmesh. + +Answers four investigation questions before any prototype is written: + +1. Does ``petsc_dm_create_submesh_from_label(marked_faces=True)`` produce + a DM with ``getDimension()=2`` and ``getCoordinateDim()=3`` on a + spherical shell's upper boundary? +2. Does that surface DM expose ``getSubpointIS()`` the same way the + ``DMPlexFilter``-derived submeshes do? +3. What labels survive on the surface DM? A closed sphere surface has + no boundary of its own, so we expect parent boundary labels to land + on the surface either as empty strata or as the surface's *own* + cells (depending on how PETSc maps them). +4. Do surface vertex coordinates sit at ``r = radiusOuter`` to + machine precision? + +This is investigation code under docs/examples/, run via: + + pixi run -e amr-dev python -u \ + docs/examples/submesh_investigation/probe_surface_extraction.py +""" + +import numpy as np +from petsc4py import PETSc + +import underworld3 as uw +from underworld3.cython.petsc_discretisation import ( + petsc_dm_create_submesh_from_label, +) + + +def _hdr(title): + uw.pprint(0, "") + uw.pprint(0, f"=== {title} ===") + + +# --------------------------------------------------------------------------- +# 1. Build a parent SphericalShell and inspect its boundary labels +# --------------------------------------------------------------------------- + +_hdr("Parent SphericalShell") + +r_outer = 1.0 +r_inner = 0.5 +cellsize = 0.25 # coarse — we only need a few cells for the probe + +shell = uw.meshing.SphericalShell( + radiusOuter=r_outer, + radiusInner=r_inner, + cellSize=cellsize, +) + +uw.pprint(0, f"parent dim = {shell.dm.getDimension()}") +uw.pprint(0, f"parent cdim = {shell.dm.getCoordinateDim()}") +uw.pprint(0, f"parent cells = " + f"{shell.dm.getHeightStratum(0)[1] - shell.dm.getHeightStratum(0)[0]}") +uw.pprint(0, f"parent boundaries: {[b.name for b in shell.boundaries]}") + +upper_value = shell.boundaries.Upper.value +uw.pprint(0, f" Upper.value = {upper_value}") + +# Confirm the Upper label has a non-empty face stratum on the parent +upper_label = shell.dm.getLabel("Upper") +upper_is = upper_label.getStratumIS(upper_value) +upper_count = upper_is.getSize() if upper_is is not None else 0 +uw.pprint(0, f" parent has {upper_count} marked points (faces) for Upper") + + +# --------------------------------------------------------------------------- +# 2. Extract the surface via DMPlexCreateSubmesh +# --------------------------------------------------------------------------- + +_hdr("DMPlexCreateSubmesh(Upper, marked_faces=True)") + +surf_dm = petsc_dm_create_submesh_from_label( + shell.dm, "Upper", upper_value, marked_faces=True, +) + +uw.pprint(0, f"surface DM type = {type(surf_dm).__name__}") +uw.pprint(0, f"surface dim = {surf_dm.getDimension()}") +uw.pprint(0, f"surface cdim = {surf_dm.getCoordinateDim()}") + +sS, sE = surf_dm.getHeightStratum(0) +uw.pprint(0, f"surface cells = {sE - sS}") + +vS, vE = surf_dm.getDepthStratum(0) +uw.pprint(0, f"surface vertices = {vE - vS}") + + +# --------------------------------------------------------------------------- +# 3. Subpoint IS — point-level parent ↔ surface map +# --------------------------------------------------------------------------- + +_hdr("Subpoint IS") + +try: + subpoint_is = surf_dm.getSubpointIS() +except Exception as e: + uw.pprint(0, f"getSubpointIS FAILED: {e!r}") + subpoint_is = None + +if subpoint_is is not None: + indices = subpoint_is.getIndices() + uw.pprint(0, f"subpoint_is size = {subpoint_is.getSize()}") + uw.pprint(0, f"first 10 indices = {indices[:10].tolist()}") + uw.pprint(0, f"index range = [{indices.min()}, {indices.max()}]") + + p_chart = shell.dm.getChart() + uw.pprint(0, f"parent chart = {p_chart}") + uw.pprint(0, + " (indices should lie inside parent chart — they are " + "point IDs in the parent's numbering)") + + +# --------------------------------------------------------------------------- +# 4. Surface vertex coordinates — embedded in 3-space at r=r_outer +# --------------------------------------------------------------------------- + +_hdr("Surface vertex coordinates") + +coords_local = surf_dm.getCoordinatesLocal().array +# Layout: flat array of cdim-tuples per vertex (cdim is on the SURFACE DM) +cdim = surf_dm.getCoordinateDim() +coords = coords_local.reshape(-1, cdim) +uw.pprint(0, f"coords shape = {coords.shape}") +uw.pprint(0, f"cdim used = {cdim}") + +radii = np.linalg.norm(coords, axis=1) +uw.pprint(0, f"radius range = [{radii.min():.10e}, {radii.max():.10e}]") +uw.pprint(0, f"|r - r_outer| max = {np.abs(radii - r_outer).max():.3e}") +uw.pprint(0, f"|r - r_outer| < 1e-10: " + f"{bool(np.all(np.abs(radii - r_outer) < 1e-10))}") + + +# --------------------------------------------------------------------------- +# 5. Label survival +# --------------------------------------------------------------------------- + +_hdr("Label survival on surface DM") + +# Enumerate all labels on the surface DM directly first — don't risk +# probing parent label *names* against the surface DM (we hit a silent +# PETSc abort doing that). This loop reaches the real survivor list +# without touching any label by name. +import sys + +uw.pprint(0, " Enumerating submesh labels by index:") +sys.stdout.flush() +n_labels = surf_dm.getNumLabels() +uw.pprint(0, f" getNumLabels = {n_labels}") +sys.stdout.flush() +all_label_names = [] +for i in range(n_labels): + nm = surf_dm.getLabelName(i) + all_label_names.append(nm) + uw.pprint(0, f" [{i}] {nm}") + sys.stdout.flush() + +uw.pprint(0, " Stratum sizes (only labels listed above):") +sys.stdout.flush() +# Skip "Centre" — known hard-abort pseudo-label +# (project_annulus_centre_pseudo_label memory) +for nm in all_label_names: + if nm == "Centre": + uw.pprint(0, f" {nm:<24} skipped (pseudo-label, would hard-abort)") + sys.stdout.flush() + continue + lab = surf_dm.getLabel(nm) + if lab is None: + uw.pprint(0, f" {nm:<24} getLabel returned None") + sys.stdout.flush() + continue + # Some labels have multiple strata; iterate values + try: + n_vals = lab.getNumValues() + except Exception as e: + uw.pprint(0, f" {nm:<24} getNumValues failed: {e!r}") + sys.stdout.flush() + continue + try: + vals_is = lab.getValueIS() + vals = vals_is.getIndices().tolist() if vals_is is not None else [] + except Exception as e: + uw.pprint(0, f" {nm:<24} getValueIS failed: {e!r}") + sys.stdout.flush() + continue + sizes = {} + for v in vals: + sis = lab.getStratumIS(int(v)) + sizes[int(v)] = sis.getSize() if sis is not None else 0 + uw.pprint(0, f" {nm:<24} values={vals} sizes={sizes}") + sys.stdout.flush() + + +# --------------------------------------------------------------------------- +# 6. Probe chart layout of the surface DM +# --------------------------------------------------------------------------- + +_hdr("Surface DM chart layout") + +cS_, cE_ = surf_dm.getHeightStratum(0) +eS_, eE_ = surf_dm.getHeightStratum(1) # edges (in 2D mesh) +vS_, vE_ = surf_dm.getDepthStratum(0) + +uw.pprint(0, f" cells : [{cS_}, {cE_}) count={cE_-cS_}") +uw.pprint(0, f" edges : [{eS_}, {eE_}) count={eE_-eS_}") +uw.pprint(0, f" vertices: [{vS_}, {vE_}) count={vE_-vS_}") + +# Try one cell's closure to see point ordering +cell_id = cS_ +closure_pts, _ = surf_dm.getTransitiveClosure(cell_id) +uw.pprint(0, f" closure(cell {cell_id}) = {closure_pts.tolist()}") +uw.pprint(0, f" cell_num_points (entities[dim]=3) → last 3 points:") +uw.pprint(0, f" {closure_pts[-3:].tolist()}") +uw.pprint(0, f" Are last-3 points in [{vS_}, {vE_})? " + f"{bool(np.all((closure_pts[-3:] >= vS_) & (closure_pts[-3:] < vE_)))}") + +# --------------------------------------------------------------------------- +# 7. Wrap the surface DM in uw.Mesh (loud failures are diagnostic) +# --------------------------------------------------------------------------- + +_hdr("Probe: can we wrap surf_dm in uw.discretisation.Mesh?") + +try: + surf_mesh = uw.discretisation.Mesh( + surf_dm, + degree=shell.degree, + qdegree=shell.qdegree, + coordinate_system_type=shell.CoordinateSystemType, + verbose=False, + ) + uw.pprint(0, f" uw.Mesh wrap OK: dim={surf_mesh.dim}, cdim={surf_mesh.cdim}") + uw.pprint(0, f" surf_mesh.X.coords shape = {surf_mesh.X.coords.shape}") + radii_uw = np.linalg.norm(surf_mesh.X.coords, axis=1) + uw.pprint(0, + f" X.coords radii: [{radii_uw.min():.10e}, " + f"{radii_uw.max():.10e}]") +except Exception as e: + uw.pprint(0, f" uw.Mesh wrap FAILED: {type(e).__name__}: {e}") + import traceback + traceback.print_exc() + +uw.pprint(0, "") +uw.pprint(0, "Probe complete.") diff --git a/docs/examples/submesh_investigation/probe_surface_strip.py b/docs/examples/submesh_investigation/probe_surface_strip.py new file mode 100644 index 00000000..cd932ab8 --- /dev/null +++ b/docs/examples/submesh_investigation/probe_surface_strip.py @@ -0,0 +1,130 @@ +"""Phase 0 probe: stripping the phantom depth-3 stratum. + +Compose DMPlexCreateSubmesh + DMPlexFilter(depth, 2) and verify that: + + 1. The resulting DM has a clean 3-stratum chart (depth = 0, 1, 2). + 2. ``getSubpointIS()`` on the filtered DM maps to points of the + un-filtered submesh — which we can then compose with the original + submesh's subpoint IS to get the parent map. + 3. Coordinates are preserved: vertices still at r = r_outer. + 4. Cell closures are now exactly (cell, edges, vertices) — no phantom + point in the middle. + 5. Standard ``Mesh._build_kd_tree_index`` (without the dim!=cdim + early-exit branch) survives the clean chart. + +Run: + pixi run -e amr-dev python -u \\ + docs/examples/submesh_investigation/probe_surface_strip.py +""" + +import numpy as np + +import underworld3 as uw +from underworld3.cython.petsc_discretisation import ( + petsc_dm_create_submesh_from_label, + petsc_dm_filter_by_label, +) + + +def hdr(t): + uw.pprint(0, ""); uw.pprint(0, f"=== {t} ===") + + +hdr("Parent SphericalShell") +shell = uw.meshing.SphericalShell( + radiusOuter=1.0, radiusInner=0.5, cellSize=0.25, +) +uw.pprint(0, f"shell dim={shell.dim} cdim={shell.cdim}") + + +hdr("Stage 1: DMPlexCreateSubmesh(Upper, marked_faces=True)") +sub1 = petsc_dm_create_submesh_from_label( + shell.dm, "Upper", shell.boundaries.Upper.value, marked_faces=True, +) +uw.pprint(0, f"sub1 dim={sub1.getDimension()} cdim={sub1.getCoordinateDim()}") +uw.pprint(0, f"sub1 chart={sub1.getChart()}") +for d in range(4): + try: + s, e = sub1.getDepthStratum(d) + uw.pprint(0, f" depth={d}: [{s}, {e}) count={e-s}") + except Exception: + pass +uw.pprint(0, f"sub1 num labels = {sub1.getNumLabels()}") + + +hdr("Stage 2: DMPlexFilter on (depth, 2) to drop the phantom stratum") +sub2 = petsc_dm_filter_by_label(sub1, "depth", 2) +uw.pprint(0, f"sub2 dim={sub2.getDimension()} cdim={sub2.getCoordinateDim()}") +uw.pprint(0, f"sub2 chart={sub2.getChart()}") +for d in range(4): + try: + s, e = sub2.getDepthStratum(d) + uw.pprint(0, f" depth={d}: [{s}, {e}) count={e-s}") + except Exception: + pass + + +hdr("Subpoint IS composition") +sub1_sp = sub1.getSubpointIS() # sub1 point -> parent point +sub2_sp = sub2.getSubpointIS() # sub2 point -> sub1 point +uw.pprint(0, f"sub1->parent IS size = {sub1_sp.getSize()}") +uw.pprint(0, f"sub2->sub1 IS size = {sub2_sp.getSize()}") +sub1_idx = sub1_sp.getIndices() +sub2_idx = sub2_sp.getIndices() +# Compose: sub2 point i -> sub1 point sub2_idx[i] -> parent point sub1_idx[sub2_idx[i]] +parent_idx = sub1_idx[sub2_idx] +uw.pprint(0, f"composed sub2->parent map size = {parent_idx.shape[0]}") +uw.pprint(0, + f"composed index range = [{parent_idx.min()}, {parent_idx.max()}] " + f"(parent chart = {shell.dm.getChart()})") + + +hdr("Surface vertex coordinates after strip") +coords = sub2.getCoordinatesLocal().array.reshape(-1, sub2.getCoordinateDim()) +radii = np.linalg.norm(coords, axis=1) +uw.pprint(0, f"sub2 vertex count = {coords.shape[0]}") +uw.pprint(0, f"radii: [{radii.min():.10e}, {radii.max():.10e}]") +uw.pprint(0, f"|r - 1| max = {np.abs(radii - 1.0).max():.3e}") + + +hdr("Cell closure layout on the stripped DM") +cS, cE = sub2.getHeightStratum(0) +for cid in [cS, cS+1, cS+10, cE-1]: + pts, _ = sub2.getTransitiveClosure(cid) + uw.pprint(0, f" closure(cell {cid}) = {pts.tolist()} (len={len(pts)})") +uw.pprint(0, " Expect length 7 = 1 cell + 3 edges + 3 vertices (no phantom).") + + +hdr("Label survival on the stripped DM") +names = [sub2.getLabelName(i) for i in range(sub2.getNumLabels())] +uw.pprint(0, f"labels: {names}") +for nm in names: + if nm == "Centre": + continue + lab = sub2.getLabel(nm) + if lab is None: continue + try: + vis = lab.getValueIS() + vals = vis.getIndices().tolist() if vis is not None else [] + sizes = {int(v): lab.getStratumIS(int(v)).getSize() for v in vals} + uw.pprint(0, f" {nm:<24} values={vals} sizes={sizes}") + except Exception as e: + uw.pprint(0, f" {nm:<24} value probe failed: {e!r}") + + +hdr("Wrap the stripped DM in uw.Mesh") +try: + surf = uw.discretisation.Mesh( + sub2, degree=shell.degree, qdegree=shell.qdegree, + coordinate_system_type=shell.CoordinateSystemType, verbose=False, + ) + uw.pprint(0, f"OK: dim={surf.dim}, cdim={surf.cdim}, " + f"vertices={surf.X.coords.shape[0]}") + uw.pprint(0, + f"X.coords radii: [{np.linalg.norm(np.asarray(surf.X.coords), axis=1).min():.6e}, " + f"{np.linalg.norm(np.asarray(surf.X.coords), axis=1).max():.6e}]") +except Exception as e: + uw.pprint(0, f"FAILED: {type(e).__name__}: {e}") + import traceback; traceback.print_exc() + +uw.pprint(0, ""); uw.pprint(0, "Probe complete.") diff --git a/docs/examples/submesh_investigation/surface_submesh_prototype.py b/docs/examples/submesh_investigation/surface_submesh_prototype.py new file mode 100644 index 00000000..053707d7 --- /dev/null +++ b/docs/examples/submesh_investigation/surface_submesh_prototype.py @@ -0,0 +1,327 @@ +"""Prototype: surface submesh extraction (third submesh flavour). + +A *surface submesh* is a real ``uw.Mesh`` extracted from a parent mesh's +codimension-1 boundary, sharing exact vertex positions with the parent. +On a 3D ``SphericalShell``, ``extract_surface(shell, "Upper")`` returns +a 2-manifold embedded in 3-space: ``dim=2, cdim=3``. + +This is the third flavour, alongside: + +| Flavour | Constructor | PETSc mechanism | dim, cdim | +|------------------|-----------------------------------|------------------------------------------|-------------------| +| Subdomain | ``mesh.extract_region(label)`` | ``DMPlexFilter`` | parent.dim, cdim | +| Resolution level | ``coarsened_companion(...)`` | ``dm.refine()`` hierarchy | parent.dim, cdim | +| **Surface** | ``extract_surface(mesh, label)`` | ``DMPlexCreateSubmesh`` + ``DMPlexFilter`` | parent.dim-1, cdim | + +All three share the same usage pattern: + + get a submesh -> build a solver on it -> map fields back and forth. + +This is investigation code under docs/examples/, not a merged API. + +Design notes +------------ + +* **Two PETSc primitives, composed.** ``DMPlexCreateSubmesh`` on a face + label produces a cd-1 DM, but PETSc retains an upward-DAG phantom + depth-3 stratum on it (one point per parent volume cell, celltype 12) + so the result can still be navigated as "a slice of the parent". For + a standalone surface mesh we don't use that parent-DAG linkage, and + the phantom points break naive ``closure[-n:]`` slicing in the + volume-mesh kd-tree builder and centroid pickers. The second stage, + ``DMPlexFilter(depth, 2)``, drops the phantom: the result is a clean + 3-stratum chart (depth = 0/1/2, celltypes ``[0, 1, 3]``). The two + subpoint IS's compose into a single surface->parent point map. +* **Design contract (loud-failure).** Surface extraction *requires* a + non-empty boundary-face stratum on the parent. If the named label + is absent or its face stratum is empty on this rank, + ``extract_surface`` raises rather than returning a degenerate mesh. +* **Manifold navigation.** With the chart clean, the standard + ``Mesh._build_kd_tree_index`` runs on the cd-1 surface: a centroid + kd-tree on a convex manifold (sphere) ranks faces by chord distance, + which agrees with the owning-face ranking to first order. Surface + ``evaluate(expr, pts)`` works for query points on the surface (the + only meaningful regime — querying R^3 points off the surface is a + separate, ill-posed problem). +""" + +from enum import Enum + +import numpy as np +from petsc4py import PETSc + +import underworld3 as uw +from underworld3.cython.petsc_discretisation import ( + petsc_dm_create_submesh_from_label, + petsc_dm_filter_by_label, +) + + +def _attach_vertex_map(surf_mesh, parent_mesh): + """Build sub_rows / parent_rows for coincident-vertex pairs. + + Workaround for UW3 issue #197 (``Mesh._build_vertex_map`` calls + ``self.X._get_kdtree()``, but ``mesh.X`` is a ``CoordinateSystem`` + with no such method — ``extract_region`` is affected too). Once the + core fix lands, this and the call site can collapse back to + ``surf_mesh._build_vertex_map()``. + + Surface vertex coords ARE an exact subset of the parent's vertex + coords (DMPlexCreateSubmesh preserves vertex positions), so a 1e-10 + coordinate match is bit-exact. + """ + sub_coords = np.asarray(surf_mesh._coords) + parent_coords = np.asarray(parent_mesh._coords) + + tree = uw.kdtree.KDTree(sub_coords) + dists, indices = tree.query(parent_coords, sqr_dists=False) + dists = np.asarray(dists).reshape(-1) + indices = np.asarray(indices).reshape(-1) + matched = dists < 1.0e-10 + + parent_rows = np.where(matched)[0] + sub_rows = indices[matched] + surf_mesh._vertex_map = (sub_rows, parent_rows) + + +def _resolve_label_value(parent_mesh, label_name, label_value): + """Resolve a boundary label name to its integer stratum value. + + Mirrors the resolve step in ``Mesh.extract_region`` but pulls from + ``boundaries`` (face labels) rather than ``regions`` (cell labels). + """ + if label_value is not None: + return label_value + + if parent_mesh.boundaries is None: + raise ValueError( + "No boundaries defined on parent mesh. " + "Provide label_value explicitly." + ) + try: + return parent_mesh.boundaries[label_name].value + except KeyError: + raise ValueError( + f"Boundary '{label_name}' not found on parent mesh. " + f"Available: {[b.name for b in parent_mesh.boundaries]}" + ) + + +def _require_non_empty_face_stratum(parent_dm, label_name, label_value): + """Loud-fail if the parent has no faces marked with this label. + + Cd-1 submesh extraction is meaningless on an empty stratum and + silently producing a zero-cell mesh is exactly the kind of soft + failure the established submesh pattern rejects. + + Important: we must *not* call ``getStratumIS(value)`` for a value + that isn't in the label's live value set — that hard-aborts PETSc + on at least some labels (the same class of abort that protects the + annulus/shell ``"Centre"`` pseudo-label). Check via ``getValueIS()`` + first. + """ + label = parent_dm.getLabel(label_name) + if label is None: + raise ValueError( + f"Parent DM has no label '{label_name}'. " + f"Cannot extract a surface submesh from it." + ) + vals_is = label.getValueIS() + live_values = ( + set(int(v) for v in vals_is.getIndices()) + if vals_is is not None + else set() + ) + if int(label_value) not in live_values: + raise ValueError( + f"Label '{label_name}' has no stratum with value " + f"{label_value} on the parent (live values: " + f"{sorted(live_values)}). There is no surface to extract." + ) + sis = label.getStratumIS(label_value) + size = sis.getSize() if sis is not None else 0 + if size == 0: + raise ValueError( + f"Label '{label_name}' (value {label_value}) has an empty " + f"face stratum on the parent. There is no surface to extract." + ) + return size + + +def _surviving_boundaries(surface_dm, parent_boundaries): + """Build a boundaries Enum from labels that landed on the surface. + + DMPlexCreateSubmesh propagates parent labels onto the submesh: the + label that selected the surface (e.g. ``"Upper"``) carries the + surface cells themselves; sibling boundary labels (e.g. ``"Lower"`` + on a SphericalShell) survive as empty strata. We keep only the + labels whose stratum is non-empty on the submesh. + + Two contract details: + + 1. We enumerate the *submesh's* labels (by index, ``getLabelName``) + rather than probing parent labels by name. ``surface_dm.getLabel( + name)`` for some parent labels (notably the annulus/shell + ``"Centre"`` pseudo-label) hits PETSc machinery that hard-aborts + rather than raises a Python error — the probe revealed this. + 2. Within a label, we obtain the live value set via ``getValueIS()`` + *before* asking for any stratum — calling ``getStratumIS(v)`` + with a value that has an empty stratum on this submesh likewise + hard-aborts (the probe path avoided this by getting values + first). + """ + if parent_boundaries is None: + return None + + parent_by_name = {b.name: b.value for b in parent_boundaries} + + # Enumerate submesh labels by index (the only safe path on a + # submesh DM — see docstring). + sm_label_names = [ + surface_dm.getLabelName(i) + for i in range(surface_dm.getNumLabels()) + ] + + surviving = {} + for name in sm_label_names: + if name not in parent_by_name: + continue # internal PETSc label (celltype, depth) — ignore + if name in ("Null_Boundary", "All_Boundaries", "Centre"): + continue + lab = surface_dm.getLabel(name) + if lab is None: + continue + # Get the live value set on this submesh BEFORE asking for any + # stratum — empty-stratum getStratumIS hard-aborts. + try: + vals_is = lab.getValueIS() + vals = ( + set(int(v) for v in vals_is.getIndices()) + if vals_is is not None + else set() + ) + except Exception: + continue + parent_value = parent_by_name[name] + if parent_value not in vals: + continue + sis = lab.getStratumIS(parent_value) + if sis is not None and sis.getSize() > 0: + surviving[name] = parent_value + + return Enum("Boundaries", surviving) if surviving else None + + +def extract_surface(parent_mesh, label_name, label_value=None, verbose=False): + """Extract the codimension-1 surface marked by ``label_name`` as a uw.Mesh. + + Wraps ``DMPlexCreateSubmesh`` to produce a real ``uw.Mesh`` for the + surface stratum. The submesh carries: + + - ``dim = parent.dim - 1`` (one less than the parent: 2 for a sphere + surface extracted from a 3D shell); + - ``cdim = parent.cdim`` (the embedding space is preserved; surface + vertices keep their 3-component coordinates); + - ``parent``, ``subpoint_is`` and registration with + ``parent._registered_submeshes`` (standard submesh lineage). + + Parent ↔ submesh DOF transfer reuses ``Mesh.restrict`` / + ``Mesh.prolongate`` (the same KDTree-coord-match-at-1e-10 path that + ``extract_region`` uses); surface vertex coordinates are an *exact* + subset of the parent's, so this is bit-exact. + + Parameters + ---------- + parent_mesh : uw.discretisation.Mesh + The parent (typically 3D) mesh with a boundary-face label. + label_name : str + Name of the parent boundary label whose marked faces become the + cells of the surface submesh. E.g. ``"Upper"`` on a + ``SphericalShell``. + label_value : int, optional + Stratum value within the label. If ``None``, resolved from + ``parent_mesh.boundaries[label_name].value``. + + Returns + ------- + uw.discretisation.Mesh + A solver-ready surface mesh with ``parent`` set to + ``parent_mesh`` and the standard submesh state attached. + + Raises + ------ + ValueError + If ``label_name`` is missing from the parent or its face stratum + is empty (loud-fail contract — there is deliberately no + geometric fallback). + """ + label_value = _resolve_label_value(parent_mesh, label_name, label_value) + _require_non_empty_face_stratum(parent_mesh.dm, label_name, label_value) + + # Stage 1: DMPlexCreateSubmesh produces a cd-1 DM with the right + # cells, edges, and vertices, but PETSc retains an upward-DAG + # phantom depth-3 stratum (one point per parent volume cell, + # celltype 12). The phantom points show up inside surface cell + # closures and break any `closure[-n:]` slicing in downstream + # code — kd-tree builder, cell-centroid pickers, anything that + # walks the closure naively. We don't use the parent-volume-cell + # linkage for anything in the surface-only mesh use case, so we + # strip it cleanly. + sub_with_phantoms = petsc_dm_create_submesh_from_label( + parent_mesh.dm, label_name, label_value, marked_faces=True, + ) + + # Stage 2: DMPlexFilter on (depth, 2) keeps only the surface cells + # and their downward closure (edges, vertices). The result is a + # genuine standalone 2-manifold mesh with a clean 3-stratum chart + # (depth = 0, 1, 2) — celltypes [0, 1, 3] only, no phantom. + surf_dm = petsc_dm_filter_by_label(sub_with_phantoms, "depth", 2) + + # Get subpoint IS's *before* the Mesh constructor wraps the DM, + # since constructor side-effects (createDS, projectCoordinates) can + # invalidate cached IS handles otherwise. Compose to get a single + # surface -> parent point map. + stage1_sp = sub_with_phantoms.getSubpointIS() # sub1 point -> parent point + stage2_sp = surf_dm.getSubpointIS() # surf point -> sub1 point + composed_indices = stage1_sp.getIndices()[stage2_sp.getIndices()] + subpoint_is = PETSc.IS().createGeneral( + composed_indices, comm=surf_dm.getComm(), + ) + + sub_boundaries = _surviving_boundaries(surf_dm, parent_mesh.boundaries) + + # Construct the Mesh. The constructor's kd-tree builder uses a + # vertex-only placeholder when dim != cdim (the chart layout of a + # DMPlexCreateSubmesh result breaks the volume-mesh closure + # assumptions); point-location-based evaluate() on a surface + # submesh is not supported by this prototype — that path needs the + # manifold-aware solver work scheduled for Session 2. + surf_mesh = uw.discretisation.Mesh( + surf_dm, + degree=parent_mesh.degree, + qdegree=parent_mesh.qdegree, + boundaries=sub_boundaries, + coordinate_system_type=parent_mesh.CoordinateSystemType, + verbose=verbose, + ) + + # Submesh lineage — same shape as extract_region + surf_mesh.parent = parent_mesh + surf_mesh.subpoint_is = subpoint_is + surf_mesh._parent_mesh_version = parent_mesh._mesh_version + surf_mesh._extract_label_name = label_name + surf_mesh._extract_label_value = label_value + surf_mesh._is_surface_submesh = True # disambiguates from extract_region + surf_mesh._dof_maps = {} + + # Build the vertex map (sub_rows -> parent_rows for shared verts). + # Inlined here rather than calling Mesh._build_vertex_map() because + # that method is broken on development (see UW3 issue #197 — + # extract_region is also affected). Once #197 lands, this can + # collapse back to surf_mesh._build_vertex_map(). + _attach_vertex_map(surf_mesh, parent_mesh) + + # Register with parent for coordinate-sync notifications + parent_mesh._registered_submeshes.add(surf_mesh) + + return surf_mesh diff --git a/docs/examples/submesh_investigation/test_surface_submesh_contract.py b/docs/examples/submesh_investigation/test_surface_submesh_contract.py new file mode 100644 index 00000000..5830d4bb --- /dev/null +++ b/docs/examples/submesh_investigation/test_surface_submesh_contract.py @@ -0,0 +1,183 @@ +"""Contract test: surface submesh extraction. + +Mirrors the loud-failure stance of ``test_refined_pair_contract.py``. +Surface extraction requires: + + 1. A boundary label that exists on the parent. + 2. That label's face stratum is non-empty. + +Both contracts must raise rather than degrade — there is deliberately no +geometric / KDTree fallback that would silently produce a degenerate +submesh. + +Also asserts that on a valid extraction, the submesh has the expected +geometry: ``dim = parent.dim - 1``, ``cdim = parent.cdim``, vertex +coordinates lie on the parent surface to machine precision. + +Run: + pixi run -e amr-dev python -u \\ + docs/examples/submesh_investigation/test_surface_submesh_contract.py +""" + +import numpy as np + +import underworld3 as uw + +import surface_submesh_prototype as ssp + + +def banner(msg): + uw.pprint(0, f"\n{'='*70}\n{msg}\n{'='*70}") + + +def main(): + banner("Parent mesh") + shell = uw.meshing.SphericalShell( + radiusOuter=1.0, + radiusInner=0.5, + cellSize=0.25, + ) + uw.pprint( + 0, + f"shell: dim={shell.dim}, cdim={shell.cdim}, " + f"boundaries={[b.name for b in shell.boundaries]}", + ) + + banner("CONTRACT 1: unknown label -> ValueError, no degraded mesh") + raised = False + try: + ssp.extract_surface(shell, "Bogus") + except ValueError as e: + raised = True + uw.pprint(0, f"raised ValueError as required:\n {e}") + assert raised, "extract_surface must raise on an unknown label" + + banner("CONTRACT 2: empty face stratum -> ValueError") + # Construct a parent boundary label name that exists but has an + # empty face stratum. We test this by passing an out-of-range + # label_value to bypass the name-lookup and hit the empty-stratum + # check directly. + raised = False + try: + ssp.extract_surface(shell, "Upper", label_value=99999) + except ValueError as e: + raised = True + uw.pprint(0, f"raised ValueError as required:\n {e}") + assert raised, "extract_surface must raise on an empty face stratum" + + banner("CONTRACT 3: valid extraction -> dim=2, cdim=3, on the sphere") + surface = ssp.extract_surface(shell, "Upper") + assert surface.parent is shell, "surface must reference its parent" + assert surface in shell._registered_submeshes, "surface must be registered" + assert surface.dim == shell.dim - 1, ( + f"surface dim should be parent.dim - 1, got {surface.dim} " + f"(parent {shell.dim})" + ) + assert surface.cdim == shell.cdim, ( + f"surface cdim should equal parent.cdim, got {surface.cdim} " + f"(parent {shell.cdim})" + ) + + radii = np.linalg.norm(np.asarray(surface.X.coords), axis=1) + max_dev = np.abs(radii - 1.0).max() + assert max_dev < 1.0e-10, ( + f"surface vertices not on parent sphere: max deviation {max_dev:.3e}" + ) + uw.pprint( + 0, + f"valid extraction: dim={surface.dim}, cdim={surface.cdim}, " + f"max|r-1| = {max_dev:.3e}", + ) + + banner("CONTRACT 4: round-trip is bit-exact at surface DOFs") + T_p = uw.discretisation.MeshVariable("T_p", shell, 1, degree=1) + T_s = uw.discretisation.MeshVariable("T_s", surface, 1, degree=1) + T_back = uw.discretisation.MeshVariable("T_back", shell, 1, degree=1) + + parent_coords = np.asarray(T_p.coords) + new_T = np.zeros_like(T_p.data) + # plant a recognisable 3-coord function + new_T[:, 0] = parent_coords[:, 0] + 2 * parent_coords[:, 1] - parent_coords[:, 2] + T_p.pack_raw_data_to_petsc(new_T, sync=True) + + surface.restrict(T_p, T_s) + surface.prolongate(T_s, T_back) + + nonzero = np.where(np.abs(T_back.data[:, 0]) > 0)[0] + diff = T_p.data[nonzero, 0] - T_back.data[nonzero, 0] + max_err = np.abs(diff).max() if diff.size else 0.0 + n_surface_dofs = T_s.data.shape[0] + uw.pprint( + 0, + f"DOFs touched in round-trip: {nonzero.shape[0]} " + f"(surface has {n_surface_dofs} DOFs)", + ) + uw.pprint(0, f"max round-trip error: {max_err:.3e}") + assert nonzero.shape[0] >= n_surface_dofs, ( + "round-trip should touch at least every surface DOF on the parent" + ) + assert max_err < 1.0e-12, ( + f"round-trip not bit-exact: max error {max_err:.3e}" + ) + + banner("CONTRACT 5: cell navigation returns valid surface cells") + # On a convex manifold (sphere) with on-surface query points, + # centroid kd-tree ranks faces by chord distance which agrees with + # the owning-face ranking to first order. With the phantom DAG + # stripped (two-stage filter), the chart is clean and the standard + # closest-cells path works. + cS, cE = surface.dm.getHeightStratum(0) + n_surface_cells = cE - cS + # Use surface vertex coordinates as query points — they are + # on-surface by construction. + query = np.asarray(surface.X.coords)[:5] + cells = surface.get_closest_cells(query) + uw.pprint(0, + f"get_closest_cells on 5 surface vertices -> {cells.tolist()}") + assert cells.shape == (5,), ( + f"closest_cells should return a 1-D array, got shape {cells.shape}" + ) + assert (cells >= 0).all() and (cells < n_surface_cells).all(), ( + f"returned cell ids {cells.tolist()} fall outside surface cell " + f"range [0, {n_surface_cells})" + ) + + banner("CONTRACT 6: evaluate(expr, surface_coords) works") + # Plant a 3-coord expression on a surface variable, then evaluate + # the symbolic expression at the surface DOF coords. Result should + # match the direct numpy eval at those coords. + T_evaluate = uw.discretisation.MeshVariable( + "T_evaluate", surface, 1, degree=1, + ) + surface_coords = np.asarray(T_evaluate.coords) + expected = ( + surface_coords[:, 0] + + 2.0 * surface_coords[:, 1] + - surface_coords[:, 2] + ) + new_T = np.zeros_like(T_evaluate.data) + new_T[:, 0] = expected + T_evaluate.pack_raw_data_to_petsc(new_T, sync=True) + + # Evaluate at a subset of the DOF coordinates — query points lie + # on the manifold by construction. + query_pts = surface_coords[::13] # sample + vals = np.asarray( + uw.function.evaluate(T_evaluate.sym[0], query_pts) + ).reshape(-1) + expected_at_query = ( + query_pts[:, 0] + 2.0 * query_pts[:, 1] - query_pts[:, 2] + ) + max_err = np.abs(vals - expected_at_query).max() + uw.pprint(0, + f"evaluate at {query_pts.shape[0]} surface points: " + f"max |val - expected| = {max_err:.3e}") + assert max_err < 1.0e-10, ( + f"evaluate disagreement on the surface: max error {max_err:.3e}" + ) + + banner("CONTRACT TESTS PASSED") + + +if __name__ == "__main__": + main() diff --git a/src/underworld3/ckdtree.pyx b/src/underworld3/ckdtree.pyx index 19acaedb..2a7164fc 100644 --- a/src/underworld3/ckdtree.pyx +++ b/src/underworld3/ckdtree.pyx @@ -350,6 +350,7 @@ cdef class KDTree: # For consistency with pykdtree if k==1: i = i.reshape(-1) + d = d.reshape(-1) if sqr_dists: return d, i diff --git a/src/underworld3/constitutive_models.py b/src/underworld3/constitutive_models.py index fbec591a..cc75c20b 100644 --- a/src/underworld3/constitutive_models.py +++ b/src/underworld3/constitutive_models.py @@ -311,7 +311,15 @@ def __init__(self, unknowns, material_name: str = None): self._DFDt = self.Unknowns.DFDt self._DuDt = self.Unknowns.DuDt - self.dim = u.mesh.dim + # Constitutive tensors relate gradients to fluxes; both live in + # the embedded coordinate space (cdim-dimensional), not the + # topological one. ``u.sym.jacobian(mesh.N)`` produces a + # (u_dim x cdim) matrix, so the conductivity / viscosity + # tensor must also be cdim-sized to multiply against it. For + # volume meshes ``dim == cdim`` so this is a no-op; for + # manifold meshes (e.g. SphericalManifold: dim=2, cdim=3) the + # constitutive tensor correctly acts on the 3-component flux. + self.dim = u.mesh.cdim self.u_dim = u.num_components self.Parameters = self._Parameters(self) diff --git a/src/underworld3/coordinates.py b/src/underworld3/coordinates.py index a2b968c6..3f50b90b 100644 --- a/src/underworld3/coordinates.py +++ b/src/underworld3/coordinates.py @@ -964,7 +964,12 @@ def __init__(self, coordinate_system): """ self.cs = coordinate_system self.mesh = coordinate_system.mesh - self._dim = self.mesh.dim # 2 for polar, 3 for spherical + # 2 for polar (cdim=2 mesh), 3 for spherical (cdim=3 mesh — + # either a 3-D volume mesh or a 2-manifold in 3-space). + # Always branch on cdim, not dim: spherical coords are a + # function of the embedded Cartesian (x, y, z) which has + # cdim components. + self._dim = self.mesh.cdim # Cache for coordinate arrays self._r_cache = None @@ -1499,7 +1504,13 @@ def __init__( self._rRotN = self._rRotN_sym.subs(th, t) self._xRotN = sympy.eye(self.mesh.dim) - elif system == CoordinateSystemType.SPHERICAL and self.mesh.dim == 3: + elif system == CoordinateSystemType.SPHERICAL and self.mesh.cdim == 3: + # Spherical coords (r, θ, φ) are a function of the + # embedded Cartesian coordinates (x, y, z). The right test + # for "can we expose spherical?" is therefore cdim == 3, + # not dim == 3 — this branch fires on both the 3-D volume + # mesh case (SphericalShell, dim=3, cdim=3) and the + # 2-manifold case (SphericalManifold, dim=2, cdim=3). self.type = "Spherical" # _X already contains UWCoordinates, _x is a copy for the "lowercase" alias @@ -1578,7 +1589,11 @@ def __init__( self._rRotN_sym = rRotN_sym self._rRotN = rRotN - self._xRotN = sympy.eye(self.mesh.dim) + # The Cartesian basis is cdim-dimensional. For a 3-D volume + # mesh dim == cdim so this is unchanged; for a 2-manifold + # in 3-space (cdim=3, dim=2) we still want a 3x3 identity + # for the embedded-Cartesian basis rotation. + self._xRotN = sympy.eye(self.mesh.cdim) elif system == CoordinateSystemType.GEOGRAPHIC and self.mesh.dim == 3: """ diff --git a/src/underworld3/cython/petsc_generic_snes_solvers.pyx b/src/underworld3/cython/petsc_generic_snes_solvers.pyx index 6bb6b71b..8c8ea8b5 100644 --- a/src/underworld3/cython/petsc_generic_snes_solvers.pyx +++ b/src/underworld3/cython/petsc_generic_snes_solvers.pyx @@ -1646,6 +1646,15 @@ class SNES_Scalar(SolverBaseClass): self.boundary_conditions = False # self._constitutive_model = None + # Constant-nullspace handling. When True, attach a MatNullSpace + # with the constant mode to the Jacobian operator before each + # KSP solve. Needed for scalar problems on a closed manifold or + # a fully-Neumann domain where the operator (e.g. Laplacian) + # has a 1-D constant kernel and the linear system is otherwise + # singular. Default False — every solver with any Dirichlet BC + # or reaction term is non-singular and shouldn't pay the cost. + self._petsc_use_constant_nullspace = False + self.verbose = verbose self._rebuild_after_mesh_update = self._build # Maybe just reboot the dm @@ -1656,6 +1665,25 @@ class SNES_Scalar(SolverBaseClass): self.is_setup = False + @property + def petsc_use_constant_nullspace(self): + """Whether to attach a constant MatNullSpace to the Jacobian. + + Set to ``True`` for scalar problems on closed manifolds (e.g. + Poisson on a ``SphericalManifold``) or fully-Neumann domains + where the linear operator has a constant kernel. PETSc projects + the right-hand side onto the orthogonal complement of the + nullspace and selects the minimum-norm solution from the + affine null-affine family, so the system becomes uniquely + solvable up to that nullspace. + """ + return self._petsc_use_constant_nullspace + + @petsc_use_constant_nullspace.setter + def petsc_use_constant_nullspace(self, value): + self._petsc_use_constant_nullspace = bool(value) + self.is_setup = False + @property def tolerance(self): """ @@ -1866,7 +1894,11 @@ class SNES_Scalar(SolverBaseClass): # Don't unwrap here — let getext()'s two-phase unwrap handle it. # This preserves constant UWexpressions as symbols for the constants[] mechanism. f0 = sympy.Array(self.F0.sym).reshape(1).as_immutable() - F1 = sympy.Array(self.F1.sym).reshape(dim).as_immutable() + # F1 is the flux vector, which lives in the embedded coordinate + # space (cdim components). For volume meshes dim==cdim so this + # is unchanged; for manifold meshes (dim=2, cdim=3) the flux + # is genuinely 3-component. + F1 = sympy.Array(self.F1.sym).reshape(cdim).as_immutable() self._u_f0 = f0 self._u_F1 = F1 @@ -2081,9 +2113,47 @@ class SNES_Scalar(SolverBaseClass): UW_DMPlexSetSNESLocalFEM(cdm.dm, PETSC_FALSE, NULL) + if self._petsc_use_constant_nullspace: + self._attach_constant_nullspace() + self.is_setup = True self.constitutive_model._solver_is_setup = True + def _attach_constant_nullspace(self): + """Attach a constant MatNullSpace to the Jacobian. + + Calls ``snes.setUp()`` first to ensure the Jacobian template + exists, then sets a constant-mode nullspace on both the + operator and preconditioner matrices (and the transpose + nullspace, since the projector is symmetric). PETSc projects + each KSP right-hand side onto the orthogonal complement of + the nullspace before solving, and returns the minimum-norm + solution within the affine null space. + + Used for scalar Poisson on closed manifolds and fully-Neumann + domains. See ``petsc_use_constant_nullspace``. + """ + self.snes.setUp() + jacobian = self.snes.getJacobian() + operator_matrix = jacobian[0] + preconditioner_matrix = jacobian[1] if len(jacobian) > 1 else None + + nullspace = PETSc.NullSpace().create( + constant=True, vectors=(), comm=self.dm.comm, + ) + + operator_matrix.setNullSpace(nullspace) + operator_matrix.setTransposeNullSpace(nullspace) + if preconditioner_matrix is not None and preconditioner_matrix.handle != operator_matrix.handle: + preconditioner_matrix.setNullSpace(nullspace) + preconditioner_matrix.setTransposeNullSpace(nullspace) + + if self.verbose and uw.mpi.rank == 0: + print( + f"SNES_Scalar ({self.name}): attached constant nullspace", + flush=True, + ) + @timing.routine_timer_decorator def solve(self, zero_init_guess: bool =True, diff --git a/src/underworld3/discretisation/discretisation_mesh.py b/src/underworld3/discretisation/discretisation_mesh.py index e8de2b57..f9e9d8c4 100644 --- a/src/underworld3/discretisation/discretisation_mesh.py +++ b/src/underworld3/discretisation/discretisation_mesh.py @@ -3187,6 +3187,23 @@ def _mark_local_boundary_faces_inside_and_out(self): ): return + # On a cd-1 surface mesh (dim != cdim — e.g. extract_surface on + # a SphericalShell: dim=2, cdim=3), the "inside vs outside" + # control-point trick doesn't apply: the 2-D perpendicular + # rule produces a 2-vector that can't dot with the 3-D + # face-to-cell vector, and "outside the domain" for a closed + # manifold has no natural geometric meaning. Surface query + # points come from the manifold itself (no R^3 → surface + # projection in scope), so points_in_domain reduces to "is + # this point close enough to the surface to be inside a cell?" + # — which the closest-local-cell test already answers. Leave + # the kd-tree as None and let points_in_domain short-circuit. + if self.dim != self.cdim: + self.boundary_face_control_points_kdtree = None + self.boundary_face_control_points_sign = None + self._domain_radius_squared = float("inf") + return + cStart, cEnd = self.dm.getHeightStratum(0) fStart, fEnd = self.dm.getHeightStratum(1) pStart, pEnd = self.dm.getDepthStratum(0) @@ -3296,6 +3313,14 @@ def points_in_domain(self, points, strict_validation=True): if model_points.shape[0] == 0: return numpy.array([], dtype=bool) + # Cd-1 surface mesh: no boundary-face control points exist + # (see _mark_local_boundary_faces_inside_and_out). Per the + # surface-mesh contract, query points are assumed to lie on + # the manifold; the closest-local-cell test is the right + # filter, not an inside/outside split. + if self.boundary_face_control_points_kdtree is None: + return self._get_closest_local_cells_internal(model_points) != -1 + dist2, closest_control_points_ext = self.boundary_face_control_points_kdtree.query( model_points, k=1, sqr_dists=True ) @@ -3407,12 +3432,23 @@ def _get_closest_local_cells_internal(self, coords: numpy.ndarray) -> numpy.ndar else: return np.zeros((0,)) - # We need to filter points that lie outside the mesh but - # still are allocated a nearby element by this distance-only check. - cells = self._indexMap[closest_points] cStart, cEnd = self.dm.getHeightStratum(0) + # Cd-1 surface mesh: the in-cell test (face-normal half-space + # rule) needs the cell-plane projection of the query for a + # curved 3-D triangle, not the 2-D perpendicular rule the + # existing helpers use. For a convex manifold with on-surface + # queries (the surface-mesh contract) the centroid kd-tree + # already returns the owning cell to first order, so skip the + # rejection step. A proper manifold in-cell test is Session 2 + # work for the surface-submesh investigation. + if self.dim != self.cdim: + return cells + + # We need to filter points that lie outside the mesh but + # still are allocated a nearby element by this distance-only check. + inside = self._test_if_points_in_cells_internal(coords, cells) cells[~inside] = -1 lost_points = np.where(inside == False)[0] diff --git a/src/underworld3/discretisation/discretisation_mesh_variables.py b/src/underworld3/discretisation/discretisation_mesh_variables.py index 78ad8249..07f46ec4 100644 --- a/src/underworld3/discretisation/discretisation_mesh_variables.py +++ b/src/underworld3/discretisation/discretisation_mesh_variables.py @@ -230,14 +230,21 @@ def __init__( self.clean_name = re.sub(r"[^a-zA-Z0-9_]", "", name) - # Variable type inference + # Variable type inference. On a cd-1 mesh (dim < cdim, e.g. + # SphericalManifold), vector fields are stored with ``cdim`` + # components — the natural embedded-Cartesian representation + # of a surface-tangent vector. We accept either dim or cdim + # as a vector match; on volume meshes dim == cdim and the + # two branches coincide. if vtype == None: if isinstance(num_components, int) and num_components == 1: vtype = uw.VarType.SCALAR - elif isinstance(num_components, int) and num_components == mesh.dim: + elif (isinstance(num_components, int) + and (num_components == mesh.dim or num_components == mesh.cdim)): vtype = uw.VarType.VECTOR elif isinstance(num_components, tuple): - if num_components[0] == mesh.dim and num_components[1] == mesh.dim: + if ((num_components[0] == mesh.dim and num_components[1] == mesh.dim) + or (num_components[0] == mesh.cdim and num_components[1] == mesh.cdim)): vtype = uw.VarType.TENSOR else: vtype = uw.VarType.MATRIX @@ -280,19 +287,26 @@ def __init__( else: self._units = None - # Component and shape handling + # Component and shape handling. + # Vector / tensor fields are sized by ``cdim`` (the embedded + # coordinate space) rather than ``dim`` (topological). For + # volume meshes dim == cdim so this is unchanged; for manifold + # meshes (e.g. SphericalManifold: dim=2, cdim=3) vectors are + # 3-component (tangent-constrained) and rank-2 tensors are + # 3x3 — matching the gradient / flux dimensions the JIT and + # solver layers expect. if vtype == uw.VarType.SCALAR: self.shape = (1, 1) self.num_components = 1 elif vtype == uw.VarType.VECTOR: - self.shape = (1, mesh.dim) - self.num_components = mesh.dim + self.shape = (1, mesh.cdim) + self.num_components = mesh.cdim elif vtype == uw.VarType.TENSOR: - self.num_components = mesh.dim * mesh.dim - self.shape = (mesh.dim, mesh.dim) + self.num_components = mesh.cdim * mesh.cdim + self.shape = (mesh.cdim, mesh.cdim) elif vtype == uw.VarType.SYM_TENSOR: - self.num_components = math.comb(mesh.dim + 1, 2) - self.shape = (mesh.dim, mesh.dim) + self.num_components = math.comb(mesh.cdim + 1, 2) + self.shape = (mesh.cdim, mesh.cdim) elif vtype == uw.VarType.MATRIX: self.num_components = self.shape[0] * self.shape[1] @@ -313,8 +327,11 @@ def __init__( self._ijk = self._sym[0] elif vtype == uw.VarType.VECTOR: - self._sym = sympy.Matrix.zeros(1, mesh.dim) - for comp in range(mesh.dim): + # cdim components — embedded-coord vector. Volume meshes + # have dim==cdim so unchanged; manifold meshes get the + # extra component(s) the gradient/flux need. + self._sym = sympy.Matrix.zeros(1, mesh.cdim) + for comp in range(mesh.cdim): self._sym[0, comp] = UnderworldFunction( self.symbol, self, @@ -326,11 +343,12 @@ def __init__( self._ijk = sympy.vector.matrix_to_vector(self._sym, self.mesh.N) elif vtype == uw.VarType.TENSOR: - self._sym = sympy.Matrix.zeros(mesh.dim, mesh.dim) + # cdim x cdim — embedded-coord rank-2 tensor. + self._sym = sympy.Matrix.zeros(mesh.cdim, mesh.cdim) # Matrix form (any number of components) - for i in range(mesh.dim): - for j in range(mesh.dim): + for i in range(mesh.cdim): + for j in range(mesh.cdim): self._sym[i, j] = UnderworldFunction( self.symbol, self, @@ -340,11 +358,12 @@ def __init__( )(*self.mesh.r) elif vtype == uw.VarType.SYM_TENSOR: - self._sym = sympy.Matrix.zeros(mesh.dim, mesh.dim) + # cdim x cdim symmetric. + self._sym = sympy.Matrix.zeros(mesh.cdim, mesh.cdim) # Matrix form (any number of components) - for i in range(mesh.dim): - for j in range(0, mesh.dim): + for i in range(mesh.cdim): + for j in range(0, mesh.cdim): if j >= i: self._sym[i, j] = UnderworldFunction( self.symbol, diff --git a/src/underworld3/meshing/__init__.py b/src/underworld3/meshing/__init__.py index 9a7a7fa8..499d0b89 100644 --- a/src/underworld3/meshing/__init__.py +++ b/src/underworld3/meshing/__init__.py @@ -17,6 +17,7 @@ SphericalShellInternalBoundary, SegmentofSphere, CubedSphere, + SphericalManifold, ) from .annulus import ( @@ -61,6 +62,7 @@ "SphericalShellInternalBoundary", "SegmentofSphere", "CubedSphere", + "SphericalManifold", # Annulus/cylindrical meshes "Annulus", "QuarterAnnulus", diff --git a/src/underworld3/meshing/spherical.py b/src/underworld3/meshing/spherical.py index 9bcb5c91..eb17cea2 100644 --- a/src/underworld3/meshing/spherical.py +++ b/src/underworld3/meshing/spherical.py @@ -280,6 +280,215 @@ class boundary_normals(Enum): return new_mesh +@timing.routine_timer_decorator +def SphericalManifold( + radius: float = 1.0, + cellSize: float = 0.1, + degree: int = 1, + qdegree: int = 2, + filename=None, + refinement=None, + gmsh_verbosity=0, + verbose=False, +): + r""" + Create a 2-manifold mesh of a spherical surface, embedded in 3-D. + + Produces a triangular surface mesh with topological dimension + ``dim=2`` and coordinate dimension ``cdim=3``: a 2-sphere lying in + Euclidean 3-space. This is the manifold sibling of + :func:`SphericalShell` (which builds a 3-D volume between two + radii); both share the spherical coordinate system so + :math:`(r, \theta, \phi)` accessors are available on either. + + The mesh is generated by building a 3-D ball in gmsh, extracting + its boundary surface, and meshing only the surface — no volume + cells are produced. The resulting DM has a clean three-stratum + chart (cells, edges, vertices) and supports PDEs solved on the + manifold (Laplace–Beltrami, surface advection-diffusion, surface + flow). + + Parameters + ---------- + radius : float, default=1.0 + Radius of the sphere. + cellSize : float, default=0.1 + Target mesh element size (chord length). + degree : int, default=1 + Polynomial degree of finite element basis functions. + qdegree : int, default=2 + Quadrature degree for numerical integration. + filename : str, optional + Path to save the gmsh ``.msh`` file. + refinement : int, optional + Number of uniform refinement levels to apply. + gmsh_verbosity : int, default=0 + Gmsh output verbosity level. + verbose : bool, default=False + Print diagnostic information. + + Returns + ------- + Mesh + A 2-manifold mesh with ``dim=2``, ``cdim=3``. The closed + sphere has no boundary curve; ``boundaries`` carries a single + ``Surface`` label covering the whole manifold for + convenience. The mesh registers spherical coordinate + accessors (``mesh.X.spherical.r``, ``.theta``, ``.phi``; + symbolic ``mesh.r`` etc.) the same way :func:`SphericalShell` + does. + + See Also + -------- + SphericalShell : 3-D volume sibling. + extract_surface : extract a manifold from an existing volume mesh's + boundary (alternative construction path). + + Notes + ----- + The closed sphere has 3 rigid-body rotation modes (about each + Cartesian axis) which are recorded on ``mesh._nullspace_rotations`` + just as for ``SphericalShell``. These tangent-to-surface vector + fields lie in the kernel of the surface Stokes / shallow-water + operators and must be projected out in solvers. + + Examples + -------- + Construct a unit-sphere manifold mesh and inspect its dimensions:: + + >>> import underworld3 as uw + >>> sphere = uw.meshing.SphericalManifold(radius=1.0, cellSize=0.1) + >>> sphere.dim, sphere.cdim + (2, 3) + + Access spherical coordinates:: + + >>> r = sphere.X.spherical.r # constant == radius + >>> theta = sphere.X.spherical.theta # co-latitude + >>> phi = sphere.X.spherical.phi # longitude + """ + + class boundaries(Enum): + Surface = 12 # the manifold itself (closed sphere has no boundary curve) + + import gmsh + + if filename is None: + if uw.mpi.rank == 0: + os.makedirs(".meshes", exist_ok=True) + uw_filename = ( + f".meshes/uw_spherical_manifold_r{radius}_csize{cellSize}.msh" + ) + else: + uw_filename = filename + + if uw.mpi.rank == 0: + gmsh.initialize() + gmsh.option.setNumber("General.Verbosity", gmsh_verbosity) + gmsh.model.add("SphericalManifold") + + ball_tag = gmsh.model.occ.addSphere(0, 0, 0, radius) + gmsh.model.occ.synchronize() + + # Extract the 2-D boundary surfaces of the ball, then discard + # the volume — we want a surface-only mesh. + surfaces = gmsh.model.getEntities(2) + gmsh.model.occ.remove([(3, ball_tag)], recursive=False) + gmsh.model.occ.synchronize() + + # Tag the surface so it has a name we can refer to. + surface_dim_tags = [s[1] for s in surfaces] + gmsh.model.addPhysicalGroup( + 2, + surface_dim_tags, + boundaries.Surface.value, + name=boundaries.Surface.name, + ) + + gmsh.option.setNumber("Mesh.CharacteristicLengthMax", cellSize) + gmsh.model.mesh.generate(2) # 2-D mesh in 3-D space + gmsh.write(uw_filename) + gmsh.finalize() + + # Force the PETSc gmsh reader to preserve the 3-D embedding when + # loading the 2-D surface mesh. Without this, DMPlexCreateGmsh + # defaults coordDim = dim (= 2) and drops the z-coordinate, which + # collapses the manifold to its xy-projection. We load the DM + # directly here (rather than going through ``_from_gmsh``'s + # h5 round-trip, which currently loses the explicit coordinate + # dimension), then hand the DM to Mesh() which accepts either a + # filename or a DMPlex object. + PETSc.Options().setValue("dm_plex_gmsh_spacedim", 3) + surface_dm = PETSc.DMPlex().createFromFile( + uw_filename, interpolate=True, comm=PETSc.COMM_WORLD, + ) + surface_dm.setName("uw_mesh") + surface_dm.markBoundaryFaces("All_Boundaries", 1001) + assert surface_dm.getCoordinateDim() == 3, ( + f"SphericalManifold expected cdim=3, got " + f"cdim={surface_dm.getCoordinateDim()} after gmsh load" + ) + + def spherical_manifold_refinement_callback(dm): + """Snap refined vertices back onto the true sphere.""" + c = dm.getCoordinatesLocal() + coords = c.array.reshape(-1, 3) + R = np.sqrt((coords ** 2).sum(axis=1)) + # Every vertex should lie on the sphere — push it back exactly. + coords *= (radius / np.maximum(R, 1.0e-30)).reshape(-1, 1) + c.array[...] = coords.reshape(-1) + dm.setCoordinatesLocal(c) + + def sphere_manifold_return_coords_to_bounds(coords): + """Project arbitrary R^3 points onto the sphere surface. + + Plugged into ``Mesh.return_coords_to_bounds`` so the SLCN + semi-Lagrangian trace-back machinery automatically maps each + Euler-step launch point back to the manifold after sampling + a tangent velocity field. The Euler step in R^3 leaves the + surface by O((|v| dt)^2 / R) per substep; this projection + wipes that drift between substeps. The closed-form sphere + projection is exact — no kd-tree or per-cell barycentric + needed for the spherical case. + + Coords are modified in-place AND returned, matching the + contract used by box/annulus/geographic factories. + """ + R = np.sqrt((coords ** 2).sum(axis=1)) + scale = radius / np.maximum(R, 1.0e-30) + coords *= scale.reshape(-1, 1) + return coords + + new_mesh = Mesh( + surface_dm, + degree=degree, + qdegree=qdegree, + coordinate_system_type=CoordinateSystemType.SPHERICAL, + useMultipleTags=True, + useRegions=True, + markVertices=True, + boundaries=boundaries, + boundary_normals=None, + refinement=refinement, + refinement_callback=spherical_manifold_refinement_callback, + return_coords_to_bounds=sphere_manifold_return_coords_to_bounds, + verbose=verbose, + ) + + # The closed sphere has 3 rigid-body rotation modes — same as + # SphericalShell. These tangent-to-surface vector fields live in + # the kernel of surface Stokes / shallow-water operators and need + # projecting out at solve time. + x, y, z = new_mesh.X + new_mesh._nullspace_rotations = [ + sympy.Matrix([0, -z, y]), # rotation about x + sympy.Matrix([z, 0, -x]), # rotation about y + sympy.Matrix([-y, x, 0]), # rotation about z + ] + + return new_mesh + + @timing.routine_timer_decorator def SphericalShellInternalBoundary( radiusOuter: float = 1.0, diff --git a/src/underworld3/swarm.py b/src/underworld3/swarm.py index 2a8350cd..9705f2ab 100644 --- a/src/underworld3/swarm.py +++ b/src/underworld3/swarm.py @@ -175,12 +175,22 @@ def __init__( mesh = swarm.mesh if vtype == None: + # Note: on a cd-1 mesh (dim < cdim, e.g. SphericalManifold), + # vector fields are stored with ``cdim`` components in the + # embedded coordinate space (tangent-constrained 3-vectors + # on a 2-manifold) and the SwarmVariable coord cache is + # built with size == mesh.cdim. We accept either dim or + # cdim as a vector match; on volume meshes dim == cdim so + # the two branches coincide. if isinstance(size, int) and size == 1: vtype = uw.VarType.SCALAR - elif isinstance(size, int) and size == mesh.dim: + elif isinstance(size, int) and (size == mesh.dim or size == mesh.cdim): vtype = uw.VarType.VECTOR elif isinstance(size, tuple): - if size[0] == mesh.dim and size[1] == mesh.dim: + if ( + (size[0] == mesh.dim and size[1] == mesh.dim) + or (size[0] == mesh.cdim and size[1] == mesh.cdim) + ): vtype = uw.VarType.TENSOR else: vtype = uw.VarType.MATRIX diff --git a/src/underworld3/systems/ddt.py b/src/underworld3/systems/ddt.py index 9b2fe244..494b3f47 100644 --- a/src/underworld3/systems/ddt.py +++ b/src/underworld3/systems/ddt.py @@ -1354,9 +1354,14 @@ def __init__( units=psi_units, # Inherit units from psi_fn ) - # We just need one swarm since this is inherently a sequential operation - nswarm = uw.swarm.NodalPointSwarm(self._workVar, verbose) - self._nswarm_psi = nswarm + # Historically this allocated a NodalPointSwarm cache here, but + # the actual trace-back path uses ``uw.function.global_evaluate`` + # on the upstream coords directly — the swarm was vestigial and + # nothing reads ``_nswarm_psi`` anywhere in the codebase. Skip + # the allocation; on manifold meshes it would fail anyway because + # DMSwarm's built-in coord field is dim-sized while manifold + # coords are cdim-sized. + self._nswarm_psi = None # The projection operator for mapping swarm values to the mesh - needs to be different for # each variable type, unfortunately ... diff --git a/src/underworld3/systems/solvers.py b/src/underworld3/systems/solvers.py index b3e36046..95a82a8a 100644 --- a/src/underworld3/systems/solvers.py +++ b/src/underworld3/systems/solvers.py @@ -2620,9 +2620,13 @@ def __init__( self.Unknowns.DuDt = DuDt + # Flux placeholder: cdim-sized to match the embedded-coord + # flux vector (volume meshes have dim==cdim so this is + # unchanged; manifold meshes have cdim > dim and need the + # extra component). self.Unknowns.DFDt = SemiLagrangian_DDt( self.mesh, - sympy.Matrix([[0] * self.mesh.dim]), # Actual function is not defined at this point + sympy.Matrix([[0] * self.mesh.cdim]), # Actual function is not defined at this point self._V_fn, vtype=uw.VarType.VECTOR, degree=u_Field.degree, @@ -3128,8 +3132,11 @@ def __init__( self.Unknowns.DuDt = DuDt if DFDt is None: + # Flux placeholder must match the flux dimension (cdim — the + # embedded coordinate space), not the topological dim. On + # volume meshes dim==cdim so this is unchanged. self.Unknowns.DFDt = Symbolic_DDt( - sympy.Matrix([[0] * self.mesh.dim]), + sympy.Matrix([[0] * self.mesh.cdim]), varsymbol=rf"{{F[ {self.u.symbol} ] }}", theta=theta, bcs=None, diff --git a/src/underworld3/utilities/_jitextension.py b/src/underworld3/utilities/_jitextension.py index 298c02ed..7c26686e 100644 --- a/src/underworld3/utilities/_jitextension.py +++ b/src/underworld3/utilities/_jitextension.py @@ -830,8 +830,14 @@ def ccode_patch_fns(varlist, prefix_str): type(var.fn)._ccodestr = f"{prefix_str}[{u_i}]" type(var.fn)._ccode = lambdafunc u_i += 1 - # now patch gradient guy into varfn guy - for ind in range(mesh.dim): + # Now patch the gradient components. The gradient of a + # field on the mesh lives in the embedded coordinate + # space (cdim-dim), so iterate to cdim — not dim. For + # volume meshes ``dim == cdim`` so this is unchanged; + # for manifold meshes (e.g. SphericalManifold dim=2, + # cdim=3) the third partial ``f_{,2}`` exists and + # needs to be wired to ``u_x[2]``. + for ind in range(mesh.cdim): # Note that var.fn._diff[ind] returns the class, so we don't need type(var.fn._diff[ind]) var.fn._diff[ind]._ccodestr = f"{prefix_str}_x[{u_x_i}]" var.fn._diff[ind]._ccode = lambdafunc @@ -848,8 +854,9 @@ def ccode_patch_fns(varlist, prefix_str): type(comp)._ccodestr = f"{prefix_str}[{u_i}]" type(comp)._ccode = lambdafunc u_i += 1 - # and also patch gradient guy into varfn guy's comp guy # Argh ... too much Mansourness - for ind in range(mesh.dim): + # Iterate to cdim (embedded coord dim) — see the + # scalar branch above for the dim != cdim reason. + for ind in range(mesh.cdim): # Note that var.fn._diff[ind] returns the class, so we don't need type(var.fn._diff[ind]) comp._diff[ind]._ccodestr = f"{prefix_str}_x[{u_x_i}]" comp._diff[ind]._ccode = lambdafunc