diff --git a/docs/developer/design/submesh-solver-architecture.md b/docs/developer/design/submesh-solver-architecture.md new file mode 100644 index 000000000..b258dce73 --- /dev/null +++ b/docs/developer/design/submesh-solver-architecture.md @@ -0,0 +1,248 @@ +# Submesh Solver Architecture: Multi-Domain Equation Systems + +## Context + +Underworld3 needs to support solving different equations on different subsets of a mesh while maintaining a unified field representation. Use cases include: + +- **Air/rock**: Stokes on rock only, full mesh for temperature/gravity +- **Surface evolution**: deforming air mesh coupled to rock Stokes +- **Gravity**: Poisson on full domain, density source from rock only +- **Multi-physics**: different equations on different subdomains (Stokes, Darcy, etc.) + +### What we've established (2026-04-05) + +1. **`DMPlexFilter`** extracts a submesh with exact shared nodes. The submesh carries a subpoint IS mapping back to the parent via `getSubpointIS()`. + +2. **PETSc Region DS** (`DMSetRegionDS`) segfaults during assembly — no examples exist in PETSc, likely incomplete infrastructure. Dead end for now. + +3. **Solver `part` parameter** in PETSc boundary assembly (`support[key.part]`) — controls which cell's closure is used for internal boundary integrals. Useful for one-sided boundary assembly but doesn't address the core problem of restricting volume assembly to a subdomain. + +4. **Low-viscosity air layer** with discontinuous pressure works reasonably but the air's incompressibility constraint acts as an unintended physical boundary condition. Not equivalent to solving on rock alone. + +5. **Normalised `Gamma_N`** (merged) — `mesh.Gamma_N` now returns a unit normal. Penalty and Nitsche BCs are mesh-independent. + +## Design Principles + +### 1. Separate meshes, separate variables, explicit copies + +Each mesh has its own MeshVariables. The user decides when data moves between meshes. There are no hidden globals or auto-managed shared fields. + +```python +# Each mesh owns its own variables +v_rock = MeshVariable("v", rock_mesh, ...) +v_full = MeshVariable("v", full_mesh, ...) + +# Solver works on submesh variables directly +stokes = Stokes(rock_mesh, velocityField=v_rock, ...) +stokes.solve() + +# Explicit copy to full mesh when needed (e.g., for visualisation or coupling) +rock_mesh.prolongate(v_rock, v_full) +``` + +### 2. Meshes know their lineage + +Every mesh has a `parent` attribute and a `subpoint_is` mapping. Top-level meshes have `parent=None` and `subpoint_is=None`. Submeshes reference their parent and carry the IS. + +```python +full_mesh.parent # None +full_mesh.subpoint_is # None + +rock_mesh = full_mesh.extract_region("Inner") +rock_mesh.parent # full_mesh +rock_mesh.subpoint_is # IS mapping submesh points -> parent points +``` + +### 3. Restrict/prolongate as mesh operations + +```python +mesh.restrict(var) # parent -> submesh DOFs (no-op if parent is None) +mesh.prolongate(var) # submesh DOFs -> parent (no-op if parent is None) +``` + +Solvers call these uniformly. On a top-level mesh they're no-ops. On a submesh they gather/scatter via the subpoint IS. The solver code doesn't branch. + +### 4. One mesh per expression + +An expression passed to a solver must only contain MeshVariable symbols from that solver's mesh. The JIT compiler evaluates all symbols against one DM's auxiliary vector and one coordinate system — mixing meshes is undefined. + +The user must restrict cross-mesh data before building expressions: + +```python +# T lives on full_mesh, but Stokes is on rock_mesh +rock_mesh.restrict(T_full, T_rock) + +# Expression uses only rock_mesh variables — safe +stokes.bodyforce = rho_rock.sym * alpha * T_rock.sym * gravity +``` + +If meshes are mixed in an expression, detect it (check `var.mesh` for all MeshVariable atoms) and raise an error at solver setup. + +### 5. Boundary mapping is automatic + +When `extract_region("Inner")` creates a submesh, boundaries are remapped: +- Full mesh "Lower" (r=r_inner) → submesh "Lower" +- Full mesh "Internal" (r=r_internal) → submesh outer boundary +- Full mesh "Upper" (r=r_outer) → not present on submesh + +The label names are preserved from the parent (they survive `DMPlexFilter`). The user refers to boundaries by the same names. + +## PETSc Infrastructure Available + +| API | What it does | Status | +|-----|-------------|--------| +| `DMPlexFilter(dm, label, value, ...)` | Extract cells by label → new DMPlex | **Works**, tested | +| `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 | +| `VecScatter` / `PetscSF` | Parallel data transfer | Standard PETSc | + +### PETSc Alternatives Investigated (2026-04-05) + +**DMComposite** — packs multiple DMs into one composite. Tested 2026-04-05. + +- Accepts DMPlex sub-DMs from DMPlexFilter. Scatter/gather works correctly. +- Interface nodes appear in both sub-DMs (102 shared vertices + 102 shared edges confirmed). +- Composite Vec concatenates sub-DM DOFs — interface DOFs are **duplicated**, not shared. Synchronisation after each solve is still required. +- **Verdict**: Designed for **combining** separate problems (fluid + structure), not **subdividing** one mesh. Doesn't simplify our use case — the core challenge (interface DOF ownership, restrict/prolongate) remains the same either way. The direct subpoint IS approach is simpler and more natural. + +**PCFIELDSPLIT with spatial IS** — split by region, not field. + +- `PCFieldSplitSetIS()` accepts arbitrary IS — confirmed no restriction to field-based splits. +- Supports Schur complement strategies between spatial blocks. +- **Problem**: This is a preconditioner, not an assembly strategy. Both blocks still assemble from the same DS. Doesn't let you have different equations per region. +- **Verdict**: Useful for preconditioning variable-viscosity systems, but doesn't solve the core problem. + +**DMCreateDomainDecomposition** — PETSc's native spatial decomposition. + +- `DMCreateDomainDecomposition_Plex()` returns inner/outer IS with configurable overlap. +- `DMCreateDomainDecompositionScatters_Plex()` creates VecScatter for restrict/prolongate. +- **Problem**: Designed for PCASM/PCGASM where the *same* equations are solved on each subdomain. Not for different physics per region. +- **Verdict**: Scatter infrastructure is useful but intent doesn't match multi-physics. + +### Assessment + +None of the PETSc mechanisms directly solve "different equations on different subsets of the same mesh with shared fields." They each address adjacent problems: + +| Mechanism | Different equations? | Shared fields? | Fits? | +|-----------|---------------------|----------------|-------| +| DMComposite | Yes | No (different vector layout) | Partial | +| PCFIELDSPLIT | No (same assembly) | Yes | No | +| DomainDecomp | No (same equations) | Yes | No | +| Region DS | Yes (in theory) | Yes | Segfaults | + +The **DMPlexFilter + subpoint IS + UW3-level restrict/prolongate** approach remains the best fit. PETSc provides the building blocks (mesh filtering, IS mapping, parallel SF), UW3 handles the multi-physics orchestration. + +## Open Questions + +1. **DM lifecycle**: The solver currently clones DMs freely (`clone_dm_hierarchy`). If the submesh also clones, DMs proliferate with no clear ownership. Need a cleanup strategy. + +2. **Mesh adaptation**: If the full mesh adapts (refinement, coarsening, surface deformation), the submesh must be re-extracted and the IS rebuilt. All in-flight MeshVariables need re-projection. How does this interact with the existing `refinement_callback` infrastructure? + +3. **Parallel decomposition**: `DMPlexFilter` builds a new SF for the submesh. If the partition differs from the parent, restrict/prolongate need MPI communication. How expensive is this? Does it matter for the target use cases? + +4. **Coupled solves**: If two solvers on different submeshes need to iterate (e.g., rock Stokes + air transport), the restrict/prolongate happens every outer iteration. Is the data copy overhead acceptable, or do we need shared vectors? + +5. **Pressure space**: Discontinuous pressure (dP1) is required for viscosity contrasts at internal boundaries. Should this be the default for submesh solvers, or should the user choose? + +## Implementation Plan + +### Immediate: `Mesh.extract_region()` + +The minimum viable feature. Everything else follows from existing UW3 patterns. + +```python +rock_mesh = full_mesh.extract_region("Inner") +``` + +Wraps `DMPlexFilter`, returns a new `Mesh` with: +- `parent` reference to the full mesh +- `subpoint_is` from `getSubpointIS()` (stored for future optimisation) +- Boundaries inherited from parent labels (they survive DMPlexFilter) +- Coordinate system inherited from parent + +The extracted mesh is fully independent — users create their own MeshVariables on it, set up solvers normally, and transfer data between parent and submesh via restrict/prolongate: + +```python +# Separate variables on separate meshes +v_rock = MeshVariable("v", rock_mesh, ...) +rho_rock = MeshVariable("rho", rock_mesh, ...) +rho_full = MeshVariable("rho", full_mesh, ...) + +# Transfer density from full mesh to rock submesh +rock_mesh.restrict(rho_full, rho_rock) + +# Stokes on rock submesh — standard solver, nothing special +stokes = Stokes(rock_mesh, velocityField=v_rock, ...) +stokes.add_natural_bc(penalty * Gamma_N.dot(v_rock.sym) * Gamma_N, "Internal") +stokes.solve() + +# Transfer rock velocity back to full mesh +rock_mesh.prolongate(v_rock, v_full) + +# Gravity on full mesh using transferred data +gravity = Poisson(full_mesh, ...) +gravity.solve() +``` + +The restrict/prolongate use the subpoint IS from `DMPlexFilter` — a direct index mapping with exact point correspondence. No kd-tree search, no interpolation, no error. This is the preferred transfer mechanism between parent and submesh. + +For transfer between unrelated meshes (no parent relationship), the existing `uw.function.evaluate(expr, coords)` path still works. + +### Restrict / Prolongate + +```python +rock_mesh.restrict(parent_var, sub_var) # gather parent DOFs at subpoint IS +rock_mesh.prolongate(sub_var, parent_var) # scatter submesh DOFs back to parent +``` + +- No-op when `parent is None` (top-level mesh) +- The subpoint IS maps submesh points → parent points +- Translation from point IS to DOF IS uses the PETSc section (offset lookup per point) +- Exact — same nodes, no interpolation + +### Why not auto-managed globals? + +We considered having MeshVariables live on the parent mesh with solvers auto-restricting/prolongating. This hides data flow, makes the solver more complex, and the user loses track of where data lives. The explicit approach is clearer: each mesh owns its variables, copies are visible. + +### Mesh deformation and adaptation + +Changes to the parent mesh must propagate to submeshes. Two cases: + +**Coordinate deformation** (ALE, surface evolution): Parent node positions change but topology is unchanged. The subpoint IS remains valid — restrict the parent's coordinate Vec to update submesh node positions. The submesh DM's internal geometry (Jacobians, normals, quadrature) must then be rebuilt. + +```python +# After deforming parent mesh coordinates +rock_mesh.sync_coordinates() # restrict parent coords via subpoint IS, rebuild geometry +``` + +This should be automatic: if the submesh detects that its parent's coordinates have changed (version counter on the parent mesh, which we already have via `_mesh_version`), it updates on next access. + +**Topology change** (adaptation, remeshing): The parent mesh gains/loses cells and vertices. The subpoint IS is invalidated — the submesh must be re-extracted from scratch. All submesh MeshVariables need re-projection onto the new submesh (interpolation from old to new via the usual adaptation path). + +```python +# After parent mesh adapts +rock_mesh = full_mesh.extract_region("Inner") # fresh extraction +# Old submesh variables are orphaned — user must re-create and re-project +``` + +This is the expensive case. The parent mesh already has `refinement_callback` infrastructure for post-adaptation fixups. The submesh re-extraction could hook into this: the parent notifies registered submeshes that topology has changed, and they invalidate themselves. + +The parent `Mesh` should track its submeshes (weak references, like the existing `_registered_swarms` pattern) so it can notify them of coordinate or topology changes. + +### Other items + +- **Boundary remapping**: Document which parent labels map to submesh boundaries. DMPlexFilter preserves labels; "Internal" on the parent becomes an exterior boundary on the submesh. +- **DM lifecycle**: Audit clone/destroy patterns, ensure submesh DMs are cleaned up. +- **Parallel**: `DMPlexFilter` builds a new SF. Test in MPI before relying on it. + +## Additional Findings + +### Discontinuous pressure required for viscosity contrasts + +Continuous P1 pressure cannot represent the pressure jump at a viscosity discontinuity (scales with viscosity ratio). With eta_rock/eta_air = 1000, the pressure smears across interface elements and corrupts velocity direction up to 177 degrees. Discontinuous P1 handles each side independently — velocity direction error drops to <5 degrees. + +### Normalised boundary normal (Gamma_N) + +`mesh.Gamma_N` now returns `Gamma / |Gamma|` — a unit normal regardless of element size. The raw `mesh.Gamma` magnitude scales with edge length (2D) / face area (3D). This affects penalty scaling: `penalty * Gamma.dot(v) * Gamma` has effective penalty ~ penalty * h², while `penalty * Gamma_N.dot(v) * Gamma_N` is mesh-independent. Nitsche's `gamma * mu / h` term now has correct 1/h scaling with normalised normals. diff --git a/docs/examples/submesh_investigation/test_bootstrap_viscosity.py b/docs/examples/submesh_investigation/test_bootstrap_viscosity.py new file mode 100644 index 000000000..25681f7e5 --- /dev/null +++ b/docs/examples/submesh_investigation/test_bootstrap_viscosity.py @@ -0,0 +1,110 @@ +""" +Bootstrap through decreasing air viscosity contrasts. + +Start from eta_air=1e-3 checkpoint, solve at 1e-4, use that to +initialise 1e-5, and so on down to 1e-6. Each step uses the +previous solution as initial guess (zero_init_guess=False). + +All use dP1 pressure, normalised Gamma_N, penalty=1e4. +""" + +import underworld3 as uw +from underworld3.systems import Stokes +import numpy as np +import sympy +import os + +r_internal = 1.0; r_inner = 0.5; r_outer_full = 1.5; cellsize = 1/16 +n = 2; k = 1; vel_penalty = 1e4; stokes_tol = 1e-4 + +# --- Create mesh --- +print("Creating mesh...", flush=True) +mesh = uw.meshing.AnnulusInternalBoundary( + radiusOuter=r_outer_full, radiusInternal=r_internal, + radiusInner=r_inner, cellSize=cellsize) + +v = uw.discretisation.MeshVariable("V", mesh, mesh.dim, degree=2) +p = uw.discretisation.MeshVariable("P", mesh, 1, degree=1, continuous=False) +eta_var = uw.discretisation.MeshVariable("eta", mesh, 1, degree=1, continuous=False) +bf_mask = uw.discretisation.MeshVariable("mask", mesh, 1, degree=1, continuous=False) + +r_at = np.sqrt(eta_var.coords[:, 0]**2 + eta_var.coords[:, 1]**2) +is_rock = r_at < r_internal +bf_mask.data[is_rock, 0] = 1.0 +bf_mask.data[~is_rock, 0] = 0.0 + +r_f, th_f = mesh.CoordinateSystem.xR +unit_r_f = mesh.CoordinateSystem.unit_e_0 +G_N = mesh.Gamma_N +v_theta = r_f * mesh.CoordinateSystem.rRotN.T * sympy.Matrix((0, 1)) + +# --- Load eta=1e-3 checkpoint as starting point --- +print("Loading eta=1e-3 checkpoint...", flush=True) +v.read_timestep("nitsche", "V", 0, outputPath="output/normalised_nitsche/") +p.read_timestep("nitsche", "P", 0, outputPath="output/normalised_nitsche/") + +# Set initial viscosity +eta_var.data[is_rock, 0] = 1.0 +eta_var.data[~is_rock, 0] = 1e-3 + +# --- Bootstrap through decreasing viscosity --- +eta_steps = [1e-4, 1e-5, 1e-6] + +for eta_air in eta_steps: + print(f"\n{'='*60}", flush=True) + print(f"Solving with eta_air = {eta_air:.0e}", flush=True) + print(f"{'='*60}", flush=True) + + # Update viscosity + eta_var.data[~is_rock, 0] = eta_air + + # Create fresh solver (needed because constitutive model refs change) + stokes = Stokes(mesh, velocityField=v, pressureField=p) + stokes.constitutive_model = uw.constitutive_models.ViscousFlowModel + stokes.constitutive_model.Parameters.shear_viscosity_0 = eta_var.sym[0, 0] + stokes.saddle_preconditioner = 1.0 / eta_var.sym[0, 0] + + rho_f = ((r_f / r_internal) ** k) * sympy.cos(n * th_f) + stokes.bodyforce = bf_mask.sym[0, 0] * rho_f * (-1.0 * unit_r_f) + + stokes.add_natural_bc(vel_penalty * G_N.dot(v.sym) * G_N, "Upper") + stokes.add_natural_bc(vel_penalty * G_N.dot(v.sym) * G_N, "Lower") + stokes.add_natural_bc(vel_penalty * v.sym.dot(unit_r_f) * unit_r_f, "Internal") + + stokes.tolerance = stokes_tol + stokes.petsc_options["snes_type"] = "newtonls" + stokes.petsc_options["ksp_type"] = "fgmres" + stokes.petsc_options["ksp_monitor"] = None + stokes.petsc_options.setValue("fieldsplit_velocity_pc_mg_type", "kaskade") + stokes.petsc_options.setValue("fieldsplit_velocity_pc_mg_cycle_type", "w") + stokes.petsc_options["fieldsplit_velocity_mg_coarse_pc_type"] = "svd" + stokes.petsc_options["fieldsplit_velocity_ksp_type"] = "fcg" + stokes.petsc_options["fieldsplit_velocity_mg_levels_ksp_type"] = "chebyshev" + stokes.petsc_options["fieldsplit_velocity_mg_levels_ksp_max_it"] = 5 + stokes.petsc_options["fieldsplit_velocity_mg_levels_ksp_converged_maxits"] = None + + # Use previous solution as initial guess + stokes.solve(zero_init_guess=False, verbose=True) + + # Null space removal + I0 = uw.maths.Integral(mesh, v_theta.dot(v.sym)) + ns = I0.evaluate() + I0.fn = v_theta.dot(v_theta) + nn = I0.evaluate() + dv = uw.function.evaluate(ns * v_theta, v.coords).reshape(-1, 2) / nn + v.data[...] -= dv + + # Norms + rock_mask = sympy.Piecewise((1.0, r_f < r_internal), (0.0, True)) + v_l2 = np.sqrt(uw.maths.Integral(mesh, rock_mask * v.sym.dot(v.sym)).evaluate()) + print(f" Inner velocity L2: {v_l2:.6e}", flush=True) + + # Checkpoint + eta_str = f"{eta_air:.0e}".replace("-", "m") + out_dir = f"./output/bootstrap_eta{eta_str}/" + if uw.mpi.rank == 0: + os.makedirs(out_dir, exist_ok=True) + mesh.write_timestep(f"eta{eta_str}", meshVars=[v, p, eta_var], outputPath=out_dir, index=0) + print(f" Checkpoint: {out_dir}", flush=True) + +print("\nDone.", flush=True) diff --git a/docs/examples/submesh_investigation/test_coupled_submesh_gravity.py b/docs/examples/submesh_investigation/test_coupled_submesh_gravity.py new file mode 100644 index 000000000..a50d4e958 --- /dev/null +++ b/docs/examples/submesh_investigation/test_coupled_submesh_gravity.py @@ -0,0 +1,184 @@ +""" +Coupled submesh demonstration: thermal-Stokes with gravity. + +Workflow: +1. Create full mesh with internal boundary +2. Extract rock submesh +3. Set temperature on rock submesh (analytical) +4. Prolongate temperature to full mesh (zero in air) +5. Solve Poisson gravity on full mesh using T-derived density +6. Restrict gravity to rock submesh +7. Solve Stokes on rock submesh with gravity as buoyancy + +This exercises the full data flow: extract_region, prolongate, +restrict, copy_into, and solving on both meshes. + +Usage: + pixi run -e default python tests/test_coupled_submesh_gravity.py +""" + +import underworld3 as uw +from underworld3.systems import Stokes +import numpy as np +import sympy +import os + +# --- Parameters --- + +r_outer = 1.5 +r_internal = 1.0 +r_inner = 0.5 +cellsize = 1/12 +vel_penalty = 1e4 + +# --- Step 1: Create meshes --- + +uw.pprint(0, "Step 1: Creating meshes...") + +full_mesh = uw.meshing.AnnulusInternalBoundary( + radiusOuter=r_outer, + radiusInternal=r_internal, + radiusInner=r_inner, + cellSize=cellsize, +) + +rock_mesh = full_mesh.extract_region("Inner") + +uw.pprint(0, f" Full mesh: {full_mesh.dm.getChart()}") +uw.pprint(0, f" Rock submesh: {rock_mesh.dm.getChart()}") + +# --- Step 2: Variables --- + +# Temperature on rock submesh +T_rock = uw.discretisation.MeshVariable("T_rock", rock_mesh, 1, degree=2) + +# Temperature on full mesh (for gravity source) +T_full = uw.discretisation.MeshVariable("T_full", full_mesh, 1, degree=2) + +# Gravity potential on full mesh +phi_full = uw.discretisation.MeshVariable("phi", full_mesh, 1, degree=2) + +# Gravity potential restricted to rock submesh +phi_rock = uw.discretisation.MeshVariable("phi_rock", rock_mesh, 1, degree=2) + +# Stokes variables on rock submesh +v_rock = uw.discretisation.MeshVariable("V_rock", rock_mesh, rock_mesh.dim, degree=2) +p_rock = uw.discretisation.MeshVariable("P_rock", rock_mesh, 1, degree=1, continuous=True) + +# --- Step 3: Set temperature on rock submesh --- + +uw.pprint(0, "Step 3: Setting temperature on rock submesh...") + +r_rock_coords = np.sqrt(T_rock.coords[:, 0]**2 + T_rock.coords[:, 1]**2) +th_rock_coords = np.arctan2(T_rock.coords[:, 1], T_rock.coords[:, 0]) + +# Temperature: hot blob near the inner boundary +T_rock.data[:, 0] = np.cos(2 * th_rock_coords) * (1.0 - (r_rock_coords - r_inner) / (r_internal - r_inner)) + +uw.pprint(0, f" T_rock range: [{T_rock.data.min():.4f}, {T_rock.data.max():.4f}]") + +# --- Step 4: Prolongate temperature to full mesh --- + +uw.pprint(0, "Step 4: Prolongating T to full mesh...") + +T_full.data[:] = 0.0 # zero in air +rock_mesh.prolongate(T_rock, T_full) + +r_full_coords = np.sqrt(T_full.coords[:, 0]**2 + T_full.coords[:, 1]**2) +rock_mask = r_full_coords < r_internal + 1e-6 +uw.pprint(0, f" T_full rock region: [{T_full.data[rock_mask].min():.4f}, {T_full.data[rock_mask].max():.4f}]") +uw.pprint(0, f" T_full air region max: {np.abs(T_full.data[~rock_mask]).max():.2e}") + +# --- Step 5: Solve Poisson gravity on full mesh --- + +uw.pprint(0, "Step 5: Solving Poisson gravity on full mesh...") + +gravity = uw.systems.Poisson(full_mesh, u_Field=phi_full) +gravity.constitutive_model = uw.constitutive_models.DiffusionModel +gravity.constitutive_model.Parameters.diffusivity = 1.0 +gravity.f = T_full.sym[0, 0] # density source = temperature + +# Zero potential on outer boundary +gravity.add_dirichlet_bc(0.0, "Upper") + +gravity.tolerance = 1e-6 +gravity.petsc_options["snes_type"] = "newtonls" +gravity.petsc_options["ksp_type"] = "fgmres" + +gravity.solve(verbose=False) + +uw.pprint(0, f" phi range: [{phi_full.data.min():.4e}, {phi_full.data.max():.4e}]") + +# --- Step 6: Restrict gravity to rock submesh --- + +uw.pprint(0, "Step 6: Restricting gravity to rock submesh...") + +rock_mesh.restrict(phi_full, phi_rock) + +err = np.abs(phi_rock.data[:, 0] - phi_full.data[rock_mask, 0][:phi_rock.data.shape[0]]).max() +uw.pprint(0, f" phi_rock range: [{phi_rock.data.min():.4e}, {phi_rock.data.max():.4e}]") + +# --- Step 7: Solve Stokes on rock submesh --- + +uw.pprint(0, "Step 7: Solving Stokes on rock submesh...") + +r_s, th_s = rock_mesh.CoordinateSystem.xR +G_N = rock_mesh.Gamma_N +unit_r = rock_mesh.CoordinateSystem.unit_e_0 + +stokes = Stokes(rock_mesh, velocityField=v_rock, pressureField=p_rock) +stokes.constitutive_model = uw.constitutive_models.ViscousFlowModel +stokes.constitutive_model.Parameters.shear_viscosity_0 = 1.0 +stokes.saddle_preconditioner = 1.0 + +# Buoyancy from gravity gradient (simplified: use T directly as density) +# In a real model: bodyforce = -rho * grad(phi) +# Here we use T as a proxy for density-driven flow +stokes.bodyforce = T_rock.sym[0, 0] * (-unit_r) + +stokes.add_natural_bc(vel_penalty * G_N.dot(v_rock.sym) * G_N, "Internal") +stokes.add_natural_bc(vel_penalty * G_N.dot(v_rock.sym) * G_N, "Lower") + +stokes.tolerance = 1e-4 +stokes.petsc_options["snes_type"] = "newtonls" +stokes.petsc_options["ksp_type"] = "fgmres" +stokes.petsc_options.setValue("fieldsplit_velocity_pc_mg_type", "kaskade") +stokes.petsc_options["fieldsplit_velocity_mg_coarse_pc_type"] = "svd" + +stokes.solve(verbose=False) + +# Null space removal +v_theta = r_s * rock_mesh.CoordinateSystem.rRotN.T * sympy.Matrix((0, 1)) +I0 = uw.maths.Integral(rock_mesh, v_theta.dot(v_rock.sym)) +ns = I0.evaluate() +I0.fn = v_theta.dot(v_theta) +nn = I0.evaluate() +dv = uw.function.evaluate(ns * v_theta, v_rock.coords).reshape(-1, 2) / nn +v_rock.data[...] -= dv + +vmag = np.sqrt(v_rock.data[:, 0]**2 + v_rock.data[:, 1]**2) +uw.pprint(0, f" Stokes max |v|: {vmag.max():.6e}") + +# --- Step 8: Prolongate velocity back to full mesh for visualisation --- + +uw.pprint(0, "Step 8: Prolongating velocity to full mesh...") + +v_full = uw.discretisation.MeshVariable("V_full", full_mesh, full_mesh.dim, degree=2) +v_full.data[:] = 0.0 +rock_mesh.prolongate(v_rock, v_full) + +vmag_full = np.sqrt(v_full.data[:, 0]**2 + v_full.data[:, 1]**2) +uw.pprint(0, f" Full mesh: rock |v| max={vmag_full[rock_mask[:v_full.data.shape[0]]].max():.6e}") +uw.pprint(0, f" Full mesh: air |v| max={vmag_full[~rock_mask[:v_full.data.shape[0]]].max():.2e}") + +# --- Checkpoint --- + +out = "./output/coupled_submesh/" +if uw.mpi.rank == 0: + os.makedirs(out, exist_ok=True) + +full_mesh.write_timestep("coupled", meshVars=[T_full, phi_full, v_full], outputPath=out, index=0) +rock_mesh.write_timestep("coupled_rock", meshVars=[T_rock, phi_rock, v_rock, p_rock], outputPath=out, index=0) + +uw.pprint(0, f"\nCheckpoints saved to {out}") +uw.pprint(0, "Done — coupled submesh workflow complete.") diff --git a/docs/examples/submesh_investigation/test_dmcomposite_probe.py b/docs/examples/submesh_investigation/test_dmcomposite_probe.py new file mode 100644 index 000000000..f47d988e1 --- /dev/null +++ b/docs/examples/submesh_investigation/test_dmcomposite_probe.py @@ -0,0 +1,152 @@ +""" +Probe: Can DMComposite manage rock/air sub-DMs from DMPlexFilter? + +Tests: +1. Create full mesh, filter into rock + air sub-DMs +2. Wrap in DMComposite +3. Check: global Vec size, scatter to sub-Vecs, IS mappings +4. Check: do interface nodes appear in both sub-DMs? +5. Can we set up fields on the rock sub-DM and solve within the composite? + +This is an investigation — not a production pattern. +""" + +from petsc4py import PETSc +import underworld3 as uw +from underworld3.cython.petsc_discretisation import petsc_dm_filter_by_label +import numpy as np + +r_internal = 1.0; r_inner = 0.5; r_outer_full = 1.5; cellsize = 1/16 + +# --- Create full mesh and filter --- + +print("Creating full mesh...", flush=True) +full_mesh = uw.meshing.AnnulusInternalBoundary( + radiusOuter=r_outer_full, radiusInternal=r_internal, + radiusInner=r_inner, cellSize=cellsize) + +full_dm = full_mesh.dm + +print("Filtering rock and air sub-DMs...", flush=True) +rock_dm = petsc_dm_filter_by_label(full_dm, "Inner", 101) +air_dm = petsc_dm_filter_by_label(full_dm, "Outer", 102) + +# Basic info +print(f"\nFull mesh chart: {full_dm.getChart()}", flush=True) +print(f"Rock submesh chart: {rock_dm.getChart()}", flush=True) +print(f"Air submesh chart: {air_dm.getChart()}", flush=True) + +# Count vertices +for name, dm in [("Full", full_dm), ("Rock", rock_dm), ("Air", air_dm)]: + depth = dm.getLabel("depth") + v_is = depth.getStratumIS(0) + n_verts = v_is.getSize() if v_is else 0 + c_is = depth.getStratumIS(2) + n_cells = c_is.getSize() if c_is else 0 + print(f" {name}: {n_verts} vertices, {n_cells} cells", flush=True) + +# --- Check subpoint IS (interface overlap) --- + +print("\n--- Subpoint IS (submesh -> parent mapping) ---", flush=True) +rock_subpoint = rock_dm.getSubpointIS() +air_subpoint = air_dm.getSubpointIS() + +if rock_subpoint: + rock_pts = set(rock_subpoint.getIndices()) + print(f"Rock subpoint IS size: {len(rock_pts)}", flush=True) +else: + rock_pts = set() + print("Rock subpoint IS: None", flush=True) + +if air_subpoint: + air_pts = set(air_subpoint.getIndices()) + print(f"Air subpoint IS size: {len(air_pts)}", flush=True) +else: + air_pts = set() + print("Air subpoint IS: None", flush=True) + +if rock_pts and air_pts: + overlap = rock_pts & air_pts + print(f"Overlap (shared points): {len(overlap)}", flush=True) + + # What depth are the shared points? + full_depth = full_dm.getLabel("depth") + overlap_by_depth = {} + for pt in overlap: + for d in range(3): + d_is = full_depth.getStratumIS(d) + if d_is and pt in set(d_is.getIndices()): + overlap_by_depth[d] = overlap_by_depth.get(d, 0) + 1 + print(f" By depth: {overlap_by_depth} (0=vertices, 1=edges, 2=cells)", flush=True) + +# --- Try DMComposite --- + +print("\n--- DMComposite test ---", flush=True) + +# DMComposite needs sub-DMs with sections (fields defined) +# Let's add a simple scalar field to each +rock_dm.setNumFields(1) +fe_rock = PETSc.FE().createDefault(2, 1, True, 1, comm=PETSc.COMM_WORLD) +rock_dm.setField(0, fe_rock) +rock_dm.createDS() + +air_dm.setNumFields(1) +fe_air = PETSc.FE().createDefault(2, 1, True, 1, comm=PETSc.COMM_WORLD) +air_dm.setField(0, fe_air) +air_dm.createDS() + +# Create composite +comp = PETSc.DMComposite().create(comm=PETSc.COMM_WORLD) +comp.addDM(rock_dm) +comp.addDM(air_dm) +comp.setUp() + +# Global vector +gvec = comp.createGlobalVec() +print(f"Composite global Vec size: {gvec.getSize()}", flush=True) + +# Check individual sub-DM vector sizes +rock_gvec = rock_dm.createGlobalVector() +air_gvec = air_dm.createGlobalVector() +print(f"Rock global Vec size: {rock_gvec.getSize()}", flush=True) +print(f"Air global Vec size: {air_gvec.getSize()}", flush=True) +print(f"Sum: {rock_gvec.getSize() + air_gvec.getSize()}", flush=True) +print(f"Full mesh would have: {full_dm.getChart()[1]} points (but DOFs depend on section)", flush=True) + +# Get IS mappings +gISs = comp.getGlobalISs() +print(f"\nGlobal IS count: {len(gISs)}", flush=True) +for i, gis in enumerate(gISs): + print(f" IS[{i}]: size={gis.getSize()}, range=[{gis.getIndices().min()}, {gis.getIndices().max()}]", flush=True) + +# Scatter test: set rock values to 1, air to 2, scatter back +rock_gvec.set(1.0) +air_gvec.set(2.0) + +# Gather into composite (petsc4py uses scatterArray/gatherArray) +comp.scatter(gvec, [rock_gvec, air_gvec]) +print(f"\nAfter scatter: rock sum={rock_gvec.sum():.0f}, air sum={air_gvec.sum():.0f}", flush=True) + +# Set sub-vecs and gather back +rock_gvec.set(1.0) +air_gvec.set(2.0) +comp.gather(gvec, PETSc.InsertMode.INSERT_VALUES, [rock_gvec, air_gvec]) +arr = gvec.getArray() +print(f"Composite Vec after gather: min={arr.min()}, max={arr.max()}", flush=True) +print(f" Values==1 (rock): {(arr == 1.0).sum()}", flush=True) +print(f" Values==2 (air): {(arr == 2.0).sum()}", flush=True) + +# The key question: can we map composite DOFs back to full mesh DOFs? +# rock subpoint IS maps rock_dm point -> full_dm point +# air subpoint IS maps air_dm point -> full_dm point +# But the composite IS maps composite index -> concatenated index +# We need: composite index -> full mesh DOF +print("\n--- Mapping composite -> full mesh ---", flush=True) +print(f"Rock subpoint IS gives rock_dm points -> full_dm points", flush=True) +print(f" e.g. rock point 0 -> full point {rock_subpoint.getIndices()[0]}", flush=True) +print(f" e.g. rock point 100 -> full point {rock_subpoint.getIndices()[100]}", flush=True) +if air_subpoint: + print(f" e.g. air point 0 -> full point {air_subpoint.getIndices()[0]}", flush=True) + print(f" e.g. air point 100 -> full point {air_subpoint.getIndices()[100]}", flush=True) + +print("\n--- Done ---", flush=True) diff --git a/docs/examples/submesh_investigation/test_investigation.py b/docs/examples/submesh_investigation/test_investigation.py new file mode 100644 index 000000000..914431ad4 --- /dev/null +++ b/docs/examples/submesh_investigation/test_investigation.py @@ -0,0 +1,162 @@ +"""Investigation: why penalty solution differs between submesh and full mesh.""" + +import underworld3 as uw +from underworld3.systems import Stokes +from underworld3.cython.petsc_discretisation import petsc_dm_filter_by_label +from underworld3.discretisation import Mesh +from underworld3.coordinates import CoordinateSystemType +import numpy as np +import sympy +from enum import Enum +from scipy.spatial import cKDTree + +r_internal = 1.0; r_inner = 0.5; r_outer_full = 1.5; cellsize = 1/16 +n = 2; k = 1; vel_penalty = 1e4; stokes_tol = 1e-4 + +print("=" * 70, flush=True) +print("INVESTIGATION: Penalty comparison submesh vs full mesh", flush=True) +print("=" * 70, flush=True) + +# Create both meshes +full_mesh = uw.meshing.AnnulusInternalBoundary( + radiusOuter=r_outer_full, radiusInternal=r_internal, + radiusInner=r_inner, cellSize=cellsize) + +subdm = petsc_dm_filter_by_label(full_mesh.dm, "Inner", 101) +subdm.markBoundaryFaces("All_Boundaries", 1001) + +class sub_bd(Enum): + Lower = 1; Upper = 2 + +rock_mesh = Mesh(subdm, degree=1, qdegree=2, boundaries=sub_bd, + coordinate_system_type=CoordinateSystemType.CYLINDRICAL2D) + +r_s, th_s = rock_mesh.CoordinateSystem.xR +r_f, th_f = full_mesh.CoordinateSystem.xR +unit_r_s = rock_mesh.CoordinateSystem.unit_e_0 +unit_r_f = full_mesh.CoordinateSystem.unit_e_0 +Gamma_s = rock_mesh.Gamma +Gamma_f = full_mesh.Gamma +v_theta_s = r_s * rock_mesh.CoordinateSystem.rRotN.T * sympy.Matrix((0, 1)) +v_theta_f = r_f * full_mesh.CoordinateSystem.rRotN.T * sympy.Matrix((0, 1)) + +# ===================================================================== +# Solve submesh with penalty=1e4 +# ===================================================================== +v_ref = uw.discretisation.MeshVariable("V_ref", rock_mesh, rock_mesh.dim, degree=2) +p_ref = uw.discretisation.MeshVariable("P_ref", rock_mesh, 1, degree=1, continuous=True) + +stokes_s = Stokes(rock_mesh, velocityField=v_ref, pressureField=p_ref) +stokes_s.constitutive_model = uw.constitutive_models.ViscousFlowModel +stokes_s.constitutive_model.Parameters.shear_viscosity_0 = 1.0 +stokes_s.saddle_preconditioner = 1.0 +rho_s = ((r_s / r_internal) ** k) * sympy.cos(n * th_s) +stokes_s.bodyforce = rho_s * (-1.0 * unit_r_s) +stokes_s.add_natural_bc(vel_penalty * Gamma_s.dot(v_ref.sym) * Gamma_s, "Upper") +stokes_s.add_natural_bc(vel_penalty * Gamma_s.dot(v_ref.sym) * Gamma_s, "Lower") +stokes_s.tolerance = stokes_tol +stokes_s.petsc_options["snes_type"] = "newtonls" +stokes_s.petsc_options["ksp_type"] = "fgmres" +stokes_s.petsc_options.setValue("fieldsplit_velocity_pc_mg_type", "kaskade") +stokes_s.petsc_options["fieldsplit_velocity_mg_coarse_pc_type"] = "svd" + +print("\nSolving submesh (penalty=1e4)...", flush=True) +stokes_s.solve(verbose=False) + +# ===================================================================== +# POINT 3: Null space +# ===================================================================== +print("\n--- POINT 3: Null space ---", flush=True) + +# Submesh null space BEFORE removal +I0 = uw.maths.Integral(rock_mesh, v_theta_s.dot(v_ref.sym)) +ns_sub = I0.evaluate() +I0.fn = v_theta_s.dot(v_theta_s) +ns_norm = I0.evaluate() +print(f"Submesh NS component (before removal): {ns_sub:.6e}", flush=True) +print(f" NS norm: {ns_norm:.6e}, relative: {abs(ns_sub/ns_norm):.6e}", flush=True) + +# Remove +dv = uw.function.evaluate(ns_sub * v_theta_s, v_ref.coords).reshape(-1, 2) / ns_norm +v_ref.data[...] -= dv + +v_mag_ref = np.sqrt(v_ref.data[:, 0]**2 + v_ref.data[:, 1]**2) +print(f"Submesh |v| after NS removal: mean={v_mag_ref.mean():.4e}, max={v_mag_ref.max():.4e}", flush=True) + +# Full mesh null space (from checkpoint) +v_pen = uw.discretisation.MeshVariable("V_pen", full_mesh, full_mesh.dim, degree=2) +v_pen.read_timestep("air_layer", "V", 0, outputPath="output/region_ds_air_layer/") + +I0_f = uw.maths.Integral(full_mesh, v_theta_f.dot(v_pen.sym)) +ns_full = I0_f.evaluate() +I0_f.fn = v_theta_f.dot(v_theta_f) +ns_norm_f = I0_f.evaluate() +print(f"\nFull mesh NS component (checkpoint, already removed): {ns_full:.6e}", flush=True) +print(f" NS norm: {ns_norm_f:.6e}, relative: {abs(ns_full/ns_norm_f):.6e}", flush=True) + +# ===================================================================== +# POINT 1: Air incompressibility — radial velocity at interface +# ===================================================================== +print("\n--- POINT 1: Radial velocity at r=1.0 ---", flush=True) + +# Full mesh +r_at_v = np.sqrt(v_pen.coords[:, 0]**2 + v_pen.coords[:, 1]**2) +int_mask = np.abs(r_at_v - r_internal) < cellsize * 0.3 +v_int = v_pen.data[int_mask] +c_int = v_pen.coords[int_mask] +r_hat = c_int / np.linalg.norm(c_int, axis=1, keepdims=True) +vr_full = np.sum(v_int * r_hat, axis=1) +print(f"Full mesh v_r at r=1.0: mean={vr_full.mean():.4e}, rms={np.sqrt((vr_full**2).mean()):.4e}, max|vr|={np.abs(vr_full).max():.4e}", flush=True) + +# Submesh +r_at_vs = np.sqrt(v_ref.coords[:, 0]**2 + v_ref.coords[:, 1]**2) +int_mask_s = np.abs(r_at_vs - r_internal) < cellsize * 0.3 +v_int_s = v_ref.data[int_mask_s] +c_int_s = v_ref.coords[int_mask_s] +r_hat_s = c_int_s / np.linalg.norm(c_int_s, axis=1, keepdims=True) +vr_sub = np.sum(v_int_s * r_hat_s, axis=1) +print(f"Submesh v_r at r=1.0: mean={vr_sub.mean():.4e}, rms={np.sqrt((vr_sub**2).mean()):.4e}, max|vr|={np.abs(vr_sub).max():.4e}", flush=True) + +print(f"\nRatio rms(vr) submesh/full: {np.sqrt((vr_sub**2).mean()) / np.sqrt((vr_full**2).mean()):.2f}", flush=True) +print(" >1 means submesh leaks MORE radially (no air resistance)", flush=True) + +# ===================================================================== +# POINT 2: Effective penalty — compare Gamma vs unit_rvec +# ===================================================================== +print("\n--- POINT 2: Penalty form ---", flush=True) +print(f"Submesh BC: vel_penalty * Gamma.dot(v) * Gamma (PETSc face normal)", flush=True) +print(f"Full mesh BC on Internal: vel_penalty * v.dot(unit_rvec) * unit_rvec (analytical radial)", flush=True) +print(f"These are DIFFERENT penalty forms. Gamma may not align with radial on the submesh.", flush=True) + +# ===================================================================== +# Match and compare +# ===================================================================== +print("\n--- MATCHED NODE COMPARISON ---", flush=True) + +tree = cKDTree(v_ref.coords) +dists, idx = tree.query(v_pen.coords) +matched = dists < 1e-10 + +v_ref_m = v_ref.data[idx[matched]] +v_pen_m = v_pen.data[matched] +coords_m = v_ref.coords[idx[matched]] + +def l2(a, b): + return np.sqrt(np.sum((a - b)**2)) / np.sqrt(np.sum(b**2)) + +vmag_r = np.sqrt(v_ref_m[:, 0]**2 + v_ref_m[:, 1]**2) +vmag_p = np.sqrt(v_pen_m[:, 0]**2 + v_pen_m[:, 1]**2) + +print(f"L2 rel error: {l2(v_pen_m, v_ref_m):.4e}", flush=True) +print(f"|v_ref| mean: {vmag_r.mean():.4e}", flush=True) +print(f"|v_pen| mean: {vmag_p.mean():.4e}", flush=True) +print(f"Ratio pen/ref: {vmag_p.mean()/vmag_r.mean():.4f}", flush=True) + +print("\n" + "=" * 70, flush=True) +print("SUMMARY", flush=True) +print("=" * 70, flush=True) +print(f"1. Radial velocity at interface: submesh leaks {np.sqrt((vr_sub**2).mean()) / np.sqrt((vr_full**2).mean()):.1f}x more than full mesh", flush=True) +print(f" -> Air incompressibility constrains radial flow even with low penalty", flush=True) +print(f"2. Different penalty forms: Gamma.dot(v)*Gamma vs v.dot(r_hat)*r_hat", flush=True) +print(f"3. Null space: submesh component = {abs(ns_sub/ns_norm):.2e}", flush=True) +print("=" * 70, flush=True) diff --git a/docs/examples/submesh_investigation/test_normalised_comparison.py b/docs/examples/submesh_investigation/test_normalised_comparison.py new file mode 100644 index 000000000..5bba88858 --- /dev/null +++ b/docs/examples/submesh_investigation/test_normalised_comparison.py @@ -0,0 +1,148 @@ +""" +Re-run rock-only submesh and Nitsche air-layer with normalised Gamma_N. +Both use identical penalty=1e4, tol=1e-4. +Checkpoints saved for notebook visualisation. +""" + +import underworld3 as uw +from underworld3.systems import Stokes +from underworld3.cython.petsc_discretisation import petsc_dm_filter_by_label +from underworld3.discretisation import Mesh +from underworld3.coordinates import CoordinateSystemType +import numpy as np +import sympy +import os +from enum import Enum + +r_outer_full = 1.5; r_internal = 1.0; r_inner = 0.5 +cellsize = 1/16; n = 2; k = 1 +vel_penalty = 1e4; stokes_tol = 1e-4; eta_air = 1e-3 + +# ===================================================================== +# Full mesh (shared by both solves) +# ===================================================================== +print("Creating full mesh...", flush=True) +full_mesh = uw.meshing.AnnulusInternalBoundary( + radiusOuter=r_outer_full, radiusInternal=r_internal, + radiusInner=r_inner, cellSize=cellsize) + +# ===================================================================== +# 1. Rock-only submesh solve +# ===================================================================== +print("\n--- Rock-only submesh ---", flush=True) +subdm = petsc_dm_filter_by_label(full_mesh.dm, "Inner", 101) +subdm.markBoundaryFaces("All_Boundaries", 1001) + +class sub_bd(Enum): + Lower = 1 # r = r_inner (from full mesh "Lower") + Internal = 2 # r = r_internal (from full mesh "Internal" — submesh outer boundary) + +rock_mesh = Mesh(subdm, degree=1, qdegree=2, boundaries=sub_bd, + coordinate_system_type=CoordinateSystemType.CYLINDRICAL2D) + +v_rock = uw.discretisation.MeshVariable("V", rock_mesh, rock_mesh.dim, degree=2) +p_rock = uw.discretisation.MeshVariable("P", rock_mesh, 1, degree=1, continuous=True) + +r_s, th_s = rock_mesh.CoordinateSystem.xR +unit_r_s = rock_mesh.CoordinateSystem.unit_e_0 +G_N_s = rock_mesh.Gamma_N +v_theta_s = r_s * rock_mesh.CoordinateSystem.rRotN.T * sympy.Matrix((0, 1)) + +stokes_rock = Stokes(rock_mesh, velocityField=v_rock, pressureField=p_rock) +stokes_rock.constitutive_model = uw.constitutive_models.ViscousFlowModel +stokes_rock.constitutive_model.Parameters.shear_viscosity_0 = 1.0 +stokes_rock.saddle_preconditioner = 1.0 +rho_s = ((r_s / r_internal) ** k) * sympy.cos(n * th_s) +stokes_rock.bodyforce = rho_s * (-1.0 * unit_r_s) +stokes_rock.add_natural_bc(vel_penalty * G_N_s.dot(v_rock.sym) * G_N_s, "Internal") +stokes_rock.add_natural_bc(vel_penalty * G_N_s.dot(v_rock.sym) * G_N_s, "Lower") +stokes_rock.tolerance = stokes_tol +stokes_rock.petsc_options["snes_type"] = "newtonls" +stokes_rock.petsc_options["ksp_type"] = "fgmres" +stokes_rock.petsc_options.setValue("fieldsplit_velocity_pc_mg_type", "kaskade") +stokes_rock.petsc_options["fieldsplit_velocity_mg_coarse_pc_type"] = "svd" + +print("Solving submesh...", flush=True) +stokes_rock.solve(verbose=False) + +# Null space removal +I0 = uw.maths.Integral(rock_mesh, v_theta_s.dot(v_rock.sym)) +ns = I0.evaluate() +I0.fn = v_theta_s.dot(v_theta_s) +ns_norm = I0.evaluate() +dv = uw.function.evaluate(ns * v_theta_s, v_rock.coords).reshape(-1, 2) / ns_norm +v_rock.data[...] -= dv + +v_l2_rock = np.sqrt(uw.maths.Integral(rock_mesh, v_rock.sym.dot(v_rock.sym)).evaluate()) +print(f"Rock submesh velocity L2: {v_l2_rock:.6e}", flush=True) + +out_rock = "./output/normalised_rock/" +if uw.mpi.rank == 0: + os.makedirs(out_rock, exist_ok=True) +rock_mesh.write_timestep("rock", meshVars=[v_rock, p_rock], outputPath=out_rock, index=0) + +# ===================================================================== +# 2. Nitsche air-layer solve on full mesh +# ===================================================================== +print("\n--- Nitsche air-layer (full mesh) ---", flush=True) +v_nit = uw.discretisation.MeshVariable("V", full_mesh, full_mesh.dim, degree=2) +p_nit = uw.discretisation.MeshVariable("P", full_mesh, 1, degree=1, continuous=False) +eta_var = uw.discretisation.MeshVariable("eta", full_mesh, 1, degree=1, continuous=False) +bf_mask = uw.discretisation.MeshVariable("mask", full_mesh, 1, degree=1, continuous=False) + +r_at = np.sqrt(eta_var.coords[:, 0]**2 + eta_var.coords[:, 1]**2) +is_rock = r_at < r_internal +eta_var.data[is_rock, 0] = 1.0 +eta_var.data[~is_rock, 0] = eta_air +bf_mask.data[is_rock, 0] = 1.0 +bf_mask.data[~is_rock, 0] = 0.0 + +r_f, th_f = full_mesh.CoordinateSystem.xR +unit_r_f = full_mesh.CoordinateSystem.unit_e_0 +G_N_f = full_mesh.Gamma_N +v_theta_f = r_f * full_mesh.CoordinateSystem.rRotN.T * sympy.Matrix((0, 1)) + +stokes_nit = Stokes(full_mesh, velocityField=v_nit, pressureField=p_nit) +stokes_nit.constitutive_model = uw.constitutive_models.ViscousFlowModel +stokes_nit.constitutive_model.Parameters.shear_viscosity_0 = eta_var.sym[0, 0] +stokes_nit.saddle_preconditioner = 1.0 / eta_var.sym[0, 0] +rho_f = ((r_f / r_internal) ** k) * sympy.cos(n * th_f) +stokes_nit.bodyforce = bf_mask.sym[0, 0] * rho_f * (-1.0 * unit_r_f) +stokes_nit.add_natural_bc(vel_penalty * G_N_f.dot(v_nit.sym) * G_N_f, "Upper") +stokes_nit.add_natural_bc(vel_penalty * G_N_f.dot(v_nit.sym) * G_N_f, "Lower") +stokes_nit.add_natural_bc(vel_penalty * v_nit.sym.dot(unit_r_f) * unit_r_f, "Internal") +stokes_nit.tolerance = stokes_tol +stokes_nit.petsc_options["snes_type"] = "newtonls" +stokes_nit.petsc_options["ksp_type"] = "fgmres" +stokes_nit.petsc_options.setValue("fieldsplit_velocity_pc_mg_type", "kaskade") +stokes_nit.petsc_options.setValue("fieldsplit_velocity_pc_mg_cycle_type", "w") +stokes_nit.petsc_options["fieldsplit_velocity_mg_coarse_pc_type"] = "svd" +stokes_nit.petsc_options["fieldsplit_velocity_ksp_type"] = "fcg" +stokes_nit.petsc_options["fieldsplit_velocity_mg_levels_ksp_type"] = "chebyshev" +stokes_nit.petsc_options["fieldsplit_velocity_mg_levels_ksp_max_it"] = 5 +stokes_nit.petsc_options["fieldsplit_velocity_mg_levels_ksp_converged_maxits"] = None + +print("Solving Nitsche...", flush=True) +stokes_nit.solve(verbose=False) + +# Null space removal +I0_f = uw.maths.Integral(full_mesh, v_theta_f.dot(v_nit.sym)) +ns_f = I0_f.evaluate() +I0_f.fn = v_theta_f.dot(v_theta_f) +ns_norm_f = I0_f.evaluate() +dv_f = uw.function.evaluate(ns_f * v_theta_f, v_nit.coords).reshape(-1, 2) / ns_norm_f +v_nit.data[...] -= dv_f + +v_l2_nit_inner = np.sqrt(uw.maths.Integral( + full_mesh, + sympy.Piecewise((1.0, r_f < r_internal), (0.0, True)) * v_nit.sym.dot(v_nit.sym) +).evaluate()) +print(f"Nitsche inner velocity L2: {v_l2_nit_inner:.6e}", flush=True) +print(f"Relative error: {abs(v_l2_nit_inner - v_l2_rock) / v_l2_rock:.4e}", flush=True) + +out_nit = "./output/normalised_nitsche/" +if uw.mpi.rank == 0: + os.makedirs(out_nit, exist_ok=True) +full_mesh.write_timestep("nitsche", meshVars=[v_nit, p_nit, eta_var], outputPath=out_nit, index=0) + +print("\nCheckpoints saved.", flush=True) diff --git a/docs/examples/submesh_investigation/test_region_ds_air_layer.py b/docs/examples/submesh_investigation/test_region_ds_air_layer.py new file mode 100644 index 000000000..fa6c0d49c --- /dev/null +++ b/docs/examples/submesh_investigation/test_region_ds_air_layer.py @@ -0,0 +1,193 @@ +""" +Low-viscosity air layer comparison for Region DS verification. + +Solves Stokes on the full AnnulusInternalBoundary mesh with: + - Rock (inner): viscosity=1.0, body force active + - Air (outer): viscosity=eta_air (very low), zero body force + +Viscosity is set element-by-element using a P0 (discontinuous degree-1) +MeshVariable, assigned from the cell region labels. + +Compares inner-region velocity/pressure norms against the rock-only +reference solution. As eta_air -> 0, the inner-region solution should +converge to the rock-only reference. + +Usage: + pixi run -e default python tests/test_region_ds_air_layer.py +""" + +import underworld3 as uw +from underworld3.systems import Stokes +import numpy as np +import sympy +import os + +# --- Parameters --- + +r_outer_full = 1.5 # Full mesh outer radius +r_internal = 1.0 # Internal boundary (rock/air interface) +r_inner = 0.5 # Inner radius +cellsize = 1/16 +n = 2 +k = 1 +vel_penalty = 1.0e4 +stokes_tol = 1.0e-4 +eta_air = 1.0e-3 # Low viscosity for air layer + +output_dir = "./output/region_ds_air_layer/" +if uw.mpi.rank == 0: + os.makedirs(output_dir, exist_ok=True) + +# --- Mesh --- + +uw.pprint(0, f"Creating full mesh: r_inner={r_inner}, r_internal={r_internal}, r_outer={r_outer_full}") + +mesh = uw.meshing.AnnulusInternalBoundary( + radiusOuter=r_outer_full, + radiusInternal=r_internal, + radiusInner=r_inner, + cellSize=cellsize, +) + +uw.pprint(0, f"Mesh chart: {mesh.dm.getChart()}") +uw.pprint(0, f"Regions: {[r.name for r in mesh.regions]}") + +# --- Variables --- + +v = uw.discretisation.MeshVariable("V", mesh, mesh.dim, degree=2) +p = uw.discretisation.MeshVariable("P", mesh, 1, degree=1, continuous=True) + +# P0-like viscosity field (discontinuous, degree 1 — lowest available) +eta_var = uw.discretisation.MeshVariable("eta", mesh, 1, degree=1, continuous=False) + +# P0-like body force mask +bf_mask_var = uw.discretisation.MeshVariable("mask", mesh, 1, degree=1, continuous=False) + +# --- Assign viscosity and mask by radius --- + +r_at_eta = np.sqrt(eta_var.coords[:, 0]**2 + eta_var.coords[:, 1]**2) +is_rock = r_at_eta < r_internal + +eta_var.data[is_rock, 0] = 1.0 +eta_var.data[~is_rock, 0] = eta_air + +bf_mask_var.data[is_rock, 0] = 1.0 +bf_mask_var.data[~is_rock, 0] = 0.0 + +n_rock = is_rock.sum() +n_air = (~is_rock).sum() +uw.pprint(0, f"Viscosity assigned: {n_rock} rock DOFs (eta=1), {n_air} air DOFs (eta={eta_air})") + +# --- Coordinate system --- + +unit_rvec = mesh.CoordinateSystem.unit_e_0 +r, th = mesh.CoordinateSystem.xR +Gamma = mesh.Gamma + +# Null space: constant v_theta in x,y coordinates +v_theta_fn_xy = r * mesh.CoordinateSystem.rRotN.T * sympy.Matrix((0, 1)) + +# --- Stokes solver --- + +stokes = Stokes(mesh, velocityField=v, pressureField=p) +stokes.constitutive_model = uw.constitutive_models.ViscousFlowModel +stokes.constitutive_model.Parameters.shear_viscosity_0 = eta_var.sym[0, 0] +stokes.saddle_preconditioner = 1.0 / eta_var.sym[0, 0] + +# Body force only in rock region +rho = ((r / r_internal) ** k) * sympy.cos(n * th) +gravity_fn = -1.0 * unit_rvec +stokes.bodyforce = bf_mask_var.sym[0, 0] * rho * gravity_fn + +# Free-slip on outer and inner boundaries (Gamma-based) +stokes.add_natural_bc(vel_penalty * Gamma.dot(v.sym) * Gamma, "Upper") +stokes.add_natural_bc(vel_penalty * Gamma.dot(v.sym) * Gamma, "Lower") + +# Penalty on radial velocity at internal boundary (analytical radial direction +# avoids Gamma support-ordering ambiguity on internal faces) +stokes.add_natural_bc(vel_penalty * v.sym.dot(unit_rvec) * unit_rvec, "Internal") + +# --- Solver options --- + +stokes.tolerance = stokes_tol +stokes.petsc_options["ksp_monitor"] = None +stokes.petsc_options["snes_type"] = "newtonls" +stokes.petsc_options["ksp_type"] = "fgmres" + +stokes.petsc_options.setValue("fieldsplit_velocity_pc_mg_type", "kaskade") +stokes.petsc_options.setValue("fieldsplit_velocity_pc_mg_cycle_type", "w") +stokes.petsc_options["fieldsplit_velocity_mg_coarse_pc_type"] = "svd" +stokes.petsc_options["fieldsplit_velocity_ksp_type"] = "fcg" +stokes.petsc_options["fieldsplit_velocity_mg_levels_ksp_type"] = "chebyshev" +stokes.petsc_options["fieldsplit_velocity_mg_levels_ksp_max_it"] = 5 +stokes.petsc_options["fieldsplit_velocity_mg_levels_ksp_converged_maxits"] = None + +# --- Solve --- + +uw.pprint(0, f"Solving Stokes with eta_air={eta_air}...") +stokes.solve(verbose=True) + +# --- Null space removal --- + +I0 = uw.maths.Integral(mesh, v_theta_fn_xy.dot(v.sym)) +norm = I0.evaluate() +I0.fn = v_theta_fn_xy.dot(v_theta_fn_xy) +vnorm = I0.evaluate() + +dv = uw.function.evaluate(norm * v_theta_fn_xy, v.coords).reshape(-1, 2) / vnorm +v.data[...] -= dv + +# --- Compute norms on INNER region only --- + +# Use Piecewise mask for integration over inner region +sharp_mask = sympy.Piecewise((1.0, r < r_internal), (0.0, True)) + +# Inner-region velocity L2 norm +v_l2_inner_integral = uw.maths.Integral(mesh, sharp_mask * v.sym.dot(v.sym)) +v_l2_inner = np.sqrt(v_l2_inner_integral.evaluate()) + +# Inner-region pressure L2 norm +p_l2_inner_integral = uw.maths.Integral(mesh, sharp_mask * p.sym.dot(p.sym)) +p_l2_inner = np.sqrt(p_l2_inner_integral.evaluate()) + +# Full-domain norms +v_l2_full = np.sqrt(uw.maths.Integral(mesh, v.sym.dot(v.sym)).evaluate()) +p_l2_full = np.sqrt(uw.maths.Integral(mesh, p.sym.dot(p.sym)).evaluate()) + +# Velocity magnitude stats in inner region +r_vals = uw.function.evaluate(r, v.coords) +inner_mask = r_vals.flatten() < r_internal +v_mag = np.sqrt(v.data[:, 0] ** 2 + v.data[:, 1] ** 2) +v_max_inner = v_mag[inner_mask].max() if inner_mask.any() else 0.0 +v_max_air = v_mag[~inner_mask].max() if (~inner_mask).any() else 0.0 + +# --- Report --- + +# Reference values from rock-only solve (cellsize=1/16, n=2, k=1) +ref_v_l2 = 1.8061681957e-03 +ref_p_l2 = 1.1796447277e-01 +ref_v_max = 2.1782171120e-03 + +uw.pprint(0, "=" * 60) +uw.pprint(0, f"Air layer comparison (eta_air={eta_air})") +uw.pprint(0, f" r_inner={r_inner}, r_internal={r_internal}, r_outer={r_outer_full}") +uw.pprint(0, "") +uw.pprint(0, " Inner-region norms:") +uw.pprint(0, f" Velocity L2: {v_l2_inner:.10e} (ref: {ref_v_l2:.10e})") +uw.pprint(0, f" Pressure L2: {p_l2_inner:.10e} (ref: {ref_p_l2:.10e})") +uw.pprint(0, f" Max |v|: {v_max_inner:.10e} (ref: {ref_v_max:.10e})") +uw.pprint(0, "") +uw.pprint(0, " Relative errors:") +uw.pprint(0, f" Velocity L2: {abs(v_l2_inner - ref_v_l2) / ref_v_l2:.4e}") +uw.pprint(0, f" Pressure L2: {abs(p_l2_inner - ref_p_l2) / ref_p_l2:.4e}") +uw.pprint(0, f" Max |v|: {abs(v_max_inner - ref_v_max) / ref_v_max:.4e}") +uw.pprint(0, "") +uw.pprint(0, " Full-domain norms:") +uw.pprint(0, f" Velocity L2: {v_l2_full:.10e}") +uw.pprint(0, f" Pressure L2: {p_l2_full:.10e}") +uw.pprint(0, f" Max |v| air: {v_max_air:.10e}") +uw.pprint(0, "=" * 60) + +# --- Save checkpoint --- +mesh.write_timestep("air_layer", meshVars=[v, p, eta_var], outputPath=output_dir, index=0) +uw.pprint(0, f"Checkpoint saved to {output_dir}") diff --git a/docs/examples/submesh_investigation/test_region_ds_nitsche.py b/docs/examples/submesh_investigation/test_region_ds_nitsche.py new file mode 100644 index 000000000..8fd84b023 --- /dev/null +++ b/docs/examples/submesh_investigation/test_region_ds_nitsche.py @@ -0,0 +1,149 @@ +""" +Nitsche BC on internal boundary with viscosity contrast. + +Same setup as test_region_ds_air_layer.py but uses add_nitsche_bc() +on the internal boundary instead of the simple velocity penalty. + +With a viscosity contrast, the Nitsche consistency term (sigma.n.d) +may correctly weight the stress from each side, potentially giving +better results than the simple penalty. + +Usage: + pixi run -e default python tests/test_region_ds_nitsche.py +""" + +import underworld3 as uw +from underworld3.systems import Stokes +import numpy as np +import sympy +import os + +# --- Parameters --- + +r_outer_full = 1.5 +r_internal = 1.0 +r_inner = 0.5 +cellsize = 1/16 +n = 2 +k = 1 +stokes_tol = 1.0e-4 +eta_air = 1.0e-3 + +output_dir = "./output/region_ds_nitsche/" +if uw.mpi.rank == 0: + os.makedirs(output_dir, exist_ok=True) + +# --- Mesh --- + +uw.pprint(0, f"Creating full mesh: r_inner={r_inner}, r_internal={r_internal}, r_outer={r_outer_full}") + +mesh = uw.meshing.AnnulusInternalBoundary( + radiusOuter=r_outer_full, + radiusInternal=r_internal, + radiusInner=r_inner, + cellSize=cellsize, +) + +# --- Variables --- + +v = uw.discretisation.MeshVariable("V", mesh, mesh.dim, degree=2) +p = uw.discretisation.MeshVariable("P", mesh, 1, degree=1, continuous=True) +eta_var = uw.discretisation.MeshVariable("eta", mesh, 1, degree=1, continuous=False) +bf_mask_var = uw.discretisation.MeshVariable("mask", mesh, 1, degree=1, continuous=False) + +# --- Assign viscosity and mask by radius --- + +r_at_eta = np.sqrt(eta_var.coords[:, 0]**2 + eta_var.coords[:, 1]**2) +is_rock = r_at_eta < r_internal +eta_var.data[is_rock, 0] = 1.0 +eta_var.data[~is_rock, 0] = eta_air +bf_mask_var.data[is_rock, 0] = 1.0 +bf_mask_var.data[~is_rock, 0] = 0.0 + +uw.pprint(0, f"Viscosity: {is_rock.sum()} rock DOFs (eta=1), {(~is_rock).sum()} air DOFs (eta={eta_air})") + +# --- Coordinate system --- + +unit_rvec = mesh.CoordinateSystem.unit_e_0 +r, th = mesh.CoordinateSystem.xR +Gamma = mesh.Gamma + +v_theta_fn_xy = r * mesh.CoordinateSystem.rRotN.T * sympy.Matrix((0, 1)) + +# --- Stokes solver --- + +stokes = Stokes(mesh, velocityField=v, pressureField=p) +stokes.constitutive_model = uw.constitutive_models.ViscousFlowModel +stokes.constitutive_model.Parameters.shear_viscosity_0 = eta_var.sym[0, 0] +stokes.saddle_preconditioner = 1.0 / eta_var.sym[0, 0] + +rho = ((r / r_internal) ** k) * sympy.cos(n * th) +stokes.bodyforce = bf_mask_var.sym[0, 0] * rho * (-1.0 * unit_rvec) + +# Free-slip on outer and inner (penalty — these are exterior boundaries) +stokes.add_natural_bc(1e4 * Gamma.dot(v.sym) * Gamma, "Upper") +stokes.add_natural_bc(1e4 * Gamma.dot(v.sym) * Gamma, "Lower") + +# Nitsche free-slip on internal boundary +# Uses constitutive model viscosity, so it sees the contrast +stokes.add_nitsche_bc("Internal", direction=unit_rvec, gamma=10.0, theta=1) + +# --- Solver options --- + +stokes.tolerance = stokes_tol +stokes.petsc_options["ksp_monitor"] = None +stokes.petsc_options["snes_type"] = "newtonls" +stokes.petsc_options["ksp_type"] = "fgmres" + +stokes.petsc_options.setValue("fieldsplit_velocity_pc_mg_type", "kaskade") +stokes.petsc_options.setValue("fieldsplit_velocity_pc_mg_cycle_type", "w") +stokes.petsc_options["fieldsplit_velocity_mg_coarse_pc_type"] = "svd" +stokes.petsc_options["fieldsplit_velocity_ksp_type"] = "fcg" +stokes.petsc_options["fieldsplit_velocity_mg_levels_ksp_type"] = "chebyshev" +stokes.petsc_options["fieldsplit_velocity_mg_levels_ksp_max_it"] = 5 +stokes.petsc_options["fieldsplit_velocity_mg_levels_ksp_converged_maxits"] = None + +# --- Solve --- + +uw.pprint(0, "Solving with Nitsche BC on internal boundary...") +stokes.solve(verbose=True) + +# --- Null space removal --- + +I0 = uw.maths.Integral(mesh, v_theta_fn_xy.dot(v.sym)) +norm = I0.evaluate() +I0.fn = v_theta_fn_xy.dot(v_theta_fn_xy) +vnorm = I0.evaluate() +dv = uw.function.evaluate(norm * v_theta_fn_xy, v.coords).reshape(-1, 2) / vnorm +v.data[...] -= dv + +# --- Norms --- + +sharp_mask = sympy.Piecewise((1.0, r < r_internal), (0.0, True)) +v_l2_inner = np.sqrt(uw.maths.Integral(mesh, sharp_mask * v.sym.dot(v.sym)).evaluate()) +p_l2_inner = np.sqrt(uw.maths.Integral(mesh, sharp_mask * p.sym.dot(p.sym)).evaluate()) + +r_vals = uw.function.evaluate(r, v.coords) +inner_mask = r_vals.flatten() < r_internal +v_mag = np.sqrt(v.data[:, 0]**2 + v.data[:, 1]**2) +v_max_inner = v_mag[inner_mask].max() + +ref_v_l2 = 1.8061681957e-03 +ref_p_l2 = 1.1796447277e-01 +ref_v_max = 2.1782171120e-03 + +uw.pprint(0, "=" * 60) +uw.pprint(0, f"Nitsche BC on Internal (eta_air={eta_air})") +uw.pprint(0, f" Rock-region norms:") +uw.pprint(0, f" Velocity L2: {v_l2_inner:.10e} (ref: {ref_v_l2:.10e})") +uw.pprint(0, f" Pressure L2: {p_l2_inner:.10e} (ref: {ref_p_l2:.10e})") +uw.pprint(0, f" Max |v|: {v_max_inner:.10e} (ref: {ref_v_max:.10e})") +uw.pprint(0, f" Relative errors:") +uw.pprint(0, f" Velocity L2: {abs(v_l2_inner - ref_v_l2) / ref_v_l2:.4e}") +uw.pprint(0, f" Pressure L2: {abs(p_l2_inner - ref_p_l2) / ref_p_l2:.4e}") +uw.pprint(0, f" Max |v|: {abs(v_max_inner - ref_v_max) / ref_v_max:.4e}") +uw.pprint(0, "=" * 60) + +# --- Checkpoint --- +mesh.write_timestep("nitsche", meshVars=[v, p, eta_var], outputPath=output_dir, index=0) +uw.pprint(0, f"Checkpoint saved to {output_dir}") diff --git a/docs/examples/submesh_investigation/test_region_ds_phase3.py b/docs/examples/submesh_investigation/test_region_ds_phase3.py new file mode 100644 index 000000000..3841f95cb --- /dev/null +++ b/docs/examples/submesh_investigation/test_region_ds_phase3.py @@ -0,0 +1,187 @@ +""" +Phase 3: Region DS — restrict Stokes assembly to rock cells only. + +Uses DMSetRegionDS to register a trivial (empty) DS on air cells. +Air DOFs are pinned to zero via Dirichlet on the "Outer" label. +The internal boundary penalty acts one-sided because air cells +contribute nothing to the residual/Jacobian. + +Usage: + pixi run -e default python tests/test_region_ds_phase3.py +""" + +import underworld3 as uw +from underworld3.systems import Stokes +import numpy as np +import sympy +import os +from enum import Enum + +# --- Parameters --- + +r_outer_full = 1.5 +r_internal = 1.0 +r_inner = 0.5 +cellsize = 1/16 +n = 2 +k = 1 +stokes_tol = 1.0e-4 +vel_penalty = 1.0e4 + +output_dir = "./output/region_ds_phase3/" +if uw.mpi.rank == 0: + os.makedirs(output_dir, exist_ok=True) + +# --- Mesh --- + +uw.pprint(0, "Creating full mesh...") + +mesh = uw.meshing.AnnulusInternalBoundary( + radiusOuter=r_outer_full, + radiusInternal=r_internal, + radiusInner=r_inner, + cellSize=cellsize, +) + +# Build a complete "AirDOFs" label that includes ALL points (vertices, edges, cells) +# in the outer region, so P2 Dirichlet BCs cover every DOF. +from underworld3.discretisation.discretisation_mesh import extend_enum +from petsc4py import PETSc as _PETSc + +dm = mesh.dm +outer_label = dm.getLabel("Outer") +outer_is = outer_label.getStratumIS(mesh.regions.Outer.value) +outer_cells = set(outer_is.getIndices()) if outer_is else set() + +# Get all cells (depth == mesh.dim) +depth_label = dm.getLabel("depth") +cell_is = depth_label.getStratumIS(mesh.dim) +all_cells = set(cell_is.getIndices()) +outer_cells_only = outer_cells & all_cells + +# For each outer cell, get its closure (vertices + edges) and label them +AIR_DOFS_VAL = 200 +dm.createLabel("AirDOFs") +air_label = dm.getLabel("AirDOFs") + +air_points = set() +for cell in outer_cells_only: + closure = dm.getTransitiveClosure(cell)[0] + air_points.update(closure) + +for pt in sorted(air_points): + air_label.setValue(pt, AIR_DOFS_VAL) + +uw.pprint(0, f"AirDOFs label: {len(air_points)} points (cells+edges+vertices in outer region)") + +# Add to boundaries enum +@extend_enum([mesh.boundaries]) +class extended_boundaries(Enum): + Outer = mesh.regions.Outer.value + AirDOFs = AIR_DOFS_VAL + +mesh.boundaries = extended_boundaries + +# Stack into UW_Boundaries +uw_bc_label = dm.getLabel("UW_Boundaries") +air_is = air_label.getStratumIS(AIR_DOFS_VAL) +if air_is: + uw_bc_label.setStratumIS(AIR_DOFS_VAL, air_is) + +# --- Variables --- + +v = uw.discretisation.MeshVariable("V", mesh, mesh.dim, degree=2) +p = uw.discretisation.MeshVariable("P", mesh, 1, degree=1, continuous=True) + +# --- Coordinate system --- + +unit_rvec = mesh.CoordinateSystem.unit_e_0 +r, th = mesh.CoordinateSystem.xR +Gamma = mesh.Gamma +v_theta_fn_xy = r * mesh.CoordinateSystem.rRotN.T * sympy.Matrix((0, 1)) + +# --- Stokes solver --- + +stokes = Stokes(mesh, velocityField=v, pressureField=p) +stokes.constitutive_model = uw.constitutive_models.ViscousFlowModel +stokes.constitutive_model.Parameters.shear_viscosity_0 = 1.0 +stokes.saddle_preconditioner = 1.0 + +# Body force everywhere (air cells won't assemble it due to Region DS) +rho = ((r / r_internal) ** k) * sympy.cos(n * th) +stokes.bodyforce = rho * (-1.0 * unit_rvec) + +# Free-slip on outer and inner boundaries +stokes.add_natural_bc(vel_penalty * Gamma.dot(v.sym) * Gamma, "Upper") +stokes.add_natural_bc(vel_penalty * Gamma.dot(v.sym) * Gamma, "Lower") + +# Penalty on internal boundary +stokes.add_natural_bc(vel_penalty * v.sym.dot(unit_rvec) * unit_rvec, "Internal") + +# Pin air DOFs to zero (using complete label with vertices+edges+cells) +# Velocity (field 0) pinned to zero in air region +stokes.add_dirichlet_bc([0.0, 0.0], "AirDOFs") + +# Configure Region DS: "Outer" cells get trivial DS (no assembly) +stokes.set_active_region("Outer", mesh.regions.Outer.value) + +# --- Solver options --- + +stokes.tolerance = stokes_tol +stokes.petsc_options["ksp_monitor"] = None +stokes.petsc_options["snes_monitor"] = None +stokes.petsc_options["snes_converged_reason"] = None +stokes.petsc_options["snes_type"] = "newtonls" +stokes.petsc_options["ksp_type"] = "fgmres" +stokes.petsc_options.setValue("fieldsplit_velocity_pc_mg_type", "kaskade") +stokes.petsc_options.setValue("fieldsplit_velocity_pc_mg_cycle_type", "w") +stokes.petsc_options["fieldsplit_velocity_mg_coarse_pc_type"] = "svd" +stokes.petsc_options["fieldsplit_velocity_ksp_type"] = "fcg" +stokes.petsc_options["fieldsplit_velocity_mg_levels_ksp_type"] = "chebyshev" +stokes.petsc_options["fieldsplit_velocity_mg_levels_ksp_max_it"] = 5 +stokes.petsc_options["fieldsplit_velocity_mg_levels_ksp_converged_maxits"] = None + +# --- Solve --- + +uw.pprint(0, "Solving with Region DS (air cells: trivial DS)...") +stokes.solve(verbose=True) + +# --- Null space removal --- + +I0 = uw.maths.Integral(mesh, v_theta_fn_xy.dot(v.sym)) +norm = I0.evaluate() +I0.fn = v_theta_fn_xy.dot(v_theta_fn_xy) +vnorm = I0.evaluate() +dv = uw.function.evaluate(norm * v_theta_fn_xy, v.coords).reshape(-1, 2) / vnorm +v.data[...] -= dv + +# --- Norms --- + +rock_mask = sympy.Piecewise((1.0, r < r_internal), (0.0, True)) +v_l2_rock = np.sqrt(uw.maths.Integral(mesh, rock_mask * v.sym.dot(v.sym)).evaluate()) +p_l2_rock = np.sqrt(uw.maths.Integral(mesh, rock_mask * p.sym.dot(p.sym)).evaluate()) + +r_vals = uw.function.evaluate(r, v.coords) +inner_mask = r_vals.flatten() < r_internal +v_mag = np.sqrt(v.data[:, 0]**2 + v.data[:, 1]**2) +v_max_rock = v_mag[inner_mask].max() +v_max_air = v_mag[~inner_mask].max() + +ref_v_l2 = 1.8061681957e-03 +ref_p_l2 = 1.1796447277e-01 + +uw.pprint(0, "=" * 60) +uw.pprint(0, "Region DS approach (trivial DS on air cells)") +uw.pprint(0, f" Rock-region norms:") +uw.pprint(0, f" Velocity L2: {v_l2_rock:.10e} (ref: {ref_v_l2:.10e})") +uw.pprint(0, f" Pressure L2: {p_l2_rock:.10e} (ref: {ref_p_l2:.10e})") +uw.pprint(0, f" Relative errors:") +uw.pprint(0, f" Velocity L2: {abs(v_l2_rock - ref_v_l2) / ref_v_l2:.4e}") +uw.pprint(0, f" Pressure L2: {abs(p_l2_rock - ref_p_l2) / ref_p_l2:.4e}") +uw.pprint(0, f" Max |v| rock: {v_max_rock:.10e}") +uw.pprint(0, f" Max |v| air: {v_max_air:.10e} (should be ~0)") +uw.pprint(0, "=" * 60) + +# --- Checkpoint --- +mesh.write_timestep("phase3", meshVars=[v, p], outputPath=output_dir, index=0) +uw.pprint(0, f"Checkpoint saved to {output_dir}") diff --git a/docs/examples/submesh_investigation/test_region_ds_pinned_air.py b/docs/examples/submesh_investigation/test_region_ds_pinned_air.py new file mode 100644 index 000000000..b858824be --- /dev/null +++ b/docs/examples/submesh_investigation/test_region_ds_pinned_air.py @@ -0,0 +1,156 @@ +""" +Pinned-air approach: Dirichlet-constrain all air DOFs to zero. + +Uses the existing solver infrastructure but adds essential BCs on the +"Outer" region label. This requires temporarily adding "Outer" to the +mesh boundaries enum so the solver's BC registration can find it. + +With air velocity pinned to zero, the penalty on the internal boundary +becomes effectively one-sided — air-side closure data contributes zero. + +Usage: + pixi run -e default python tests/test_region_ds_pinned_air.py +""" + +import underworld3 as uw +from underworld3.systems import Stokes +import numpy as np +import sympy +import os +from enum import Enum + +# --- Parameters --- + +r_outer_full = 1.5 +r_internal = 1.0 +r_inner = 0.5 +cellsize = 1/16 +n = 2 +k = 1 +stokes_tol = 1.0e-4 +vel_penalty = 1.0e4 + +output_dir = "./output/region_ds_pinned/" +if uw.mpi.rank == 0: + os.makedirs(output_dir, exist_ok=True) + +# --- Mesh --- + +uw.pprint(0, "Creating full mesh...") + +mesh = uw.meshing.AnnulusInternalBoundary( + radiusOuter=r_outer_full, + radiusInternal=r_internal, + radiusInner=r_inner, + cellSize=cellsize, +) + +# Add region labels to the boundaries enum so the solver can find them +# for essential BC registration. This is a workaround until proper +# Region DS support is added to the solver. +from underworld3.discretisation.discretisation_mesh import extend_enum + +@extend_enum([mesh.boundaries]) +class extended_boundaries(Enum): + Outer = mesh.regions.Outer.value # 102 + +mesh.boundaries = extended_boundaries + +uw.pprint(0, f"Boundaries: {[b.name for b in mesh.boundaries]}") +uw.pprint(0, f"Outer label value: {mesh.boundaries.Outer.value}") + +# Verify the "Outer" DM label exists +outer_label = mesh.dm.getLabel("Outer") +uw.pprint(0, f"Outer DM label exists: {outer_label is not None}") + +# --- Variables --- + +v = uw.discretisation.MeshVariable("V", mesh, mesh.dim, degree=2) +p = uw.discretisation.MeshVariable("P", mesh, 1, degree=1, continuous=True) + +# --- Coordinate system --- + +unit_rvec = mesh.CoordinateSystem.unit_e_0 +r, th = mesh.CoordinateSystem.xR +Gamma = mesh.Gamma +v_theta_fn_xy = r * mesh.CoordinateSystem.rRotN.T * sympy.Matrix((0, 1)) + +# --- Stokes solver --- + +stokes = Stokes(mesh, velocityField=v, pressureField=p) +stokes.constitutive_model = uw.constitutive_models.ViscousFlowModel +stokes.constitutive_model.Parameters.shear_viscosity_0 = 1.0 +stokes.saddle_preconditioner = 1.0 + +# Body force: same smooth density, applied everywhere (air DOFs are pinned anyway) +rho = ((r / r_internal) ** k) * sympy.cos(n * th) +stokes.bodyforce = rho * (-1.0 * unit_rvec) + +# Free-slip on outer and inner boundaries +stokes.add_natural_bc(vel_penalty * Gamma.dot(v.sym) * Gamma, "Upper") +stokes.add_natural_bc(vel_penalty * Gamma.dot(v.sym) * Gamma, "Lower") + +# Penalty on internal boundary (radial direction) +stokes.add_natural_bc(vel_penalty * v.sym.dot(unit_rvec) * unit_rvec, "Internal") + +# Pin all air DOFs to zero (v=0, p=0 in outer region) +stokes.add_dirichlet_bc([0.0, 0.0], "Outer") + +# --- Solver options --- + +stokes.tolerance = stokes_tol +stokes.petsc_options["ksp_monitor"] = None +stokes.petsc_options["snes_type"] = "newtonls" +stokes.petsc_options["ksp_type"] = "fgmres" +stokes.petsc_options.setValue("fieldsplit_velocity_pc_mg_type", "kaskade") +stokes.petsc_options.setValue("fieldsplit_velocity_pc_mg_cycle_type", "w") +stokes.petsc_options["fieldsplit_velocity_mg_coarse_pc_type"] = "svd" +stokes.petsc_options["fieldsplit_velocity_ksp_type"] = "fcg" +stokes.petsc_options["fieldsplit_velocity_mg_levels_ksp_type"] = "chebyshev" +stokes.petsc_options["fieldsplit_velocity_mg_levels_ksp_max_it"] = 5 +stokes.petsc_options["fieldsplit_velocity_mg_levels_ksp_converged_maxits"] = None + +# --- Solve --- + +uw.pprint(0, "Solving with pinned air DOFs...") +stokes.solve(verbose=True) + +# --- Null space removal --- + +I0 = uw.maths.Integral(mesh, v_theta_fn_xy.dot(v.sym)) +norm = I0.evaluate() +I0.fn = v_theta_fn_xy.dot(v_theta_fn_xy) +vnorm = I0.evaluate() +dv = uw.function.evaluate(norm * v_theta_fn_xy, v.coords).reshape(-1, 2) / vnorm +v.data[...] -= dv + +# --- Norms --- + +rock_mask = sympy.Piecewise((1.0, r < r_internal), (0.0, True)) +v_l2_rock = np.sqrt(uw.maths.Integral(mesh, rock_mask * v.sym.dot(v.sym)).evaluate()) +p_l2_rock = np.sqrt(uw.maths.Integral(mesh, rock_mask * p.sym.dot(p.sym)).evaluate()) + +r_vals = uw.function.evaluate(r, v.coords) +inner_mask = r_vals.flatten() < r_internal +v_mag = np.sqrt(v.data[:, 0]**2 + v.data[:, 1]**2) +v_max_rock = v_mag[inner_mask].max() +v_max_air = v_mag[~inner_mask].max() + +ref_v_l2 = 1.8061681957e-03 +ref_p_l2 = 1.1796447277e-01 + +uw.pprint(0, "=" * 60) +uw.pprint(0, "Pinned-air approach (Dirichlet v=0 on Outer region)") +uw.pprint(0, f" Rock-region norms:") +uw.pprint(0, f" Velocity L2: {v_l2_rock:.10e} (ref: {ref_v_l2:.10e})") +uw.pprint(0, f" Pressure L2: {p_l2_rock:.10e} (ref: {ref_p_l2:.10e})") +uw.pprint(0, f" Relative errors:") +uw.pprint(0, f" Velocity L2: {abs(v_l2_rock - ref_v_l2) / ref_v_l2:.4e}") +uw.pprint(0, f" Pressure L2: {abs(p_l2_rock - ref_p_l2) / ref_p_l2:.4e}") +uw.pprint(0, f" Max |v| rock: {v_max_rock:.10e}") +uw.pprint(0, f" Max |v| air: {v_max_air:.10e} (should be ~0)") +uw.pprint(0, "=" * 60) + +# --- Checkpoint --- +mesh.write_timestep("pinned", meshVars=[v, p], outputPath=output_dir, index=0) +uw.pprint(0, f"Checkpoint saved to {output_dir}") diff --git a/docs/examples/submesh_investigation/test_region_ds_pinned_interior.py b/docs/examples/submesh_investigation/test_region_ds_pinned_interior.py new file mode 100644 index 000000000..2dc64a40e --- /dev/null +++ b/docs/examples/submesh_investigation/test_region_ds_pinned_interior.py @@ -0,0 +1,185 @@ +""" +Pinned air-interior approach: Dirichlet only on air DOFs that are NOT +on the internal boundary. + +The previous pinned-air test applied Dirichlet v=0 on ALL points in the +"Outer" label, including vertices shared with the internal boundary. +Those interface vertices should be free to participate in the rock solve. + +This test creates an "AirInterior" label excluding interface points, +and applies Dirichlet only there. + +Usage: + pixi run -e default python tests/test_region_ds_pinned_interior.py +""" + +import underworld3 as uw +from underworld3.systems import Stokes +import numpy as np +import sympy +import os +from enum import Enum +from petsc4py import PETSc + +# --- Parameters --- + +r_outer_full = 1.5 +r_internal = 1.0 +r_inner = 0.5 +cellsize = 1/16 +n = 2 +k = 1 +stokes_tol = 1.0e-4 +vel_penalty = 1.0e4 + +output_dir = "./output/region_ds_pinned_interior/" +if uw.mpi.rank == 0: + os.makedirs(output_dir, exist_ok=True) + +# --- Mesh --- + +uw.pprint(0, "Creating full mesh...") + +mesh = uw.meshing.AnnulusInternalBoundary( + radiusOuter=r_outer_full, + radiusInternal=r_internal, + radiusInner=r_inner, + cellSize=cellsize, +) + +# --- Create AirInterior label: Outer points minus Internal points --- + +dm = mesh.dm +outer_label = dm.getLabel("Outer") +internal_label = dm.getLabel("Internal") + +# Get point sets +outer_is = outer_label.getStratumIS(mesh.regions.Outer.value) +internal_is = internal_label.getStratumIS(mesh.boundaries.Internal.value) + +outer_points = set(outer_is.getIndices()) if outer_is else set() +internal_points = set(internal_is.getIndices()) if internal_is else set() + +# Air interior = outer minus internal boundary +air_interior_points = outer_points - internal_points + +uw.pprint(0, f"Outer points: {len(outer_points)}") +uw.pprint(0, f"Internal points: {len(internal_points)}") +uw.pprint(0, f"Air interior points: {len(air_interior_points)}") +uw.pprint(0, f"Interface points removed: {len(outer_points) - len(air_interior_points)}") + +# Create DM label +AIR_INTERIOR_VAL = 200 +dm.createLabel("AirInterior") +air_label = dm.getLabel("AirInterior") +for pt in sorted(air_interior_points): + air_label.setValue(pt, AIR_INTERIOR_VAL) + +# Add to mesh boundaries so solver can find it +from underworld3.discretisation.discretisation_mesh import extend_enum + +@extend_enum([mesh.boundaries]) +class extended_boundaries(Enum): + AirInterior = AIR_INTERIOR_VAL + +mesh.boundaries = extended_boundaries + +# Also stack into UW_Boundaries +uw_bc_label = dm.getLabel("UW_Boundaries") +air_is = air_label.getStratumIS(AIR_INTERIOR_VAL) +if air_is: + uw_bc_label.setStratumIS(AIR_INTERIOR_VAL, air_is) + +uw.pprint(0, f"AirInterior label created with {len(air_interior_points)} points, value={AIR_INTERIOR_VAL}") + +# --- Variables --- + +v = uw.discretisation.MeshVariable("V", mesh, mesh.dim, degree=2) +p = uw.discretisation.MeshVariable("P", mesh, 1, degree=1, continuous=True) + +# --- Coordinate system --- + +unit_rvec = mesh.CoordinateSystem.unit_e_0 +r, th = mesh.CoordinateSystem.xR +Gamma = mesh.Gamma +v_theta_fn_xy = r * mesh.CoordinateSystem.rRotN.T * sympy.Matrix((0, 1)) + +# --- Stokes solver --- + +stokes = Stokes(mesh, velocityField=v, pressureField=p) +stokes.constitutive_model = uw.constitutive_models.ViscousFlowModel +stokes.constitutive_model.Parameters.shear_viscosity_0 = 1.0 +stokes.saddle_preconditioner = 1.0 + +# Body force everywhere (air DOFs are pinned anyway) +rho = ((r / r_internal) ** k) * sympy.cos(n * th) +stokes.bodyforce = rho * (-1.0 * unit_rvec) + +# Free-slip on outer and inner boundaries +stokes.add_natural_bc(vel_penalty * Gamma.dot(v.sym) * Gamma, "Upper") +stokes.add_natural_bc(vel_penalty * Gamma.dot(v.sym) * Gamma, "Lower") + +# Penalty on internal boundary +stokes.add_natural_bc(vel_penalty * v.sym.dot(unit_rvec) * unit_rvec, "Internal") + +# Pin air-interior DOFs to zero (NOT interface DOFs) +stokes.add_dirichlet_bc([0.0, 0.0], "AirInterior") + +# --- Solver options --- + +stokes.tolerance = stokes_tol +stokes.petsc_options["ksp_monitor"] = None +stokes.petsc_options["snes_type"] = "newtonls" +stokes.petsc_options["ksp_type"] = "fgmres" +stokes.petsc_options.setValue("fieldsplit_velocity_pc_mg_type", "kaskade") +stokes.petsc_options.setValue("fieldsplit_velocity_pc_mg_cycle_type", "w") +stokes.petsc_options["fieldsplit_velocity_mg_coarse_pc_type"] = "svd" +stokes.petsc_options["fieldsplit_velocity_ksp_type"] = "fcg" +stokes.petsc_options["fieldsplit_velocity_mg_levels_ksp_type"] = "chebyshev" +stokes.petsc_options["fieldsplit_velocity_mg_levels_ksp_max_it"] = 5 +stokes.petsc_options["fieldsplit_velocity_mg_levels_ksp_converged_maxits"] = None + +# --- Solve --- + +uw.pprint(0, "Solving with pinned air-interior DOFs...") +stokes.solve(verbose=True) + +# --- Null space removal --- + +I0 = uw.maths.Integral(mesh, v_theta_fn_xy.dot(v.sym)) +norm = I0.evaluate() +I0.fn = v_theta_fn_xy.dot(v_theta_fn_xy) +vnorm = I0.evaluate() +dv = uw.function.evaluate(norm * v_theta_fn_xy, v.coords).reshape(-1, 2) / vnorm +v.data[...] -= dv + +# --- Norms --- + +rock_mask = sympy.Piecewise((1.0, r < r_internal), (0.0, True)) +v_l2_rock = np.sqrt(uw.maths.Integral(mesh, rock_mask * v.sym.dot(v.sym)).evaluate()) +p_l2_rock = np.sqrt(uw.maths.Integral(mesh, rock_mask * p.sym.dot(p.sym)).evaluate()) + +r_vals = uw.function.evaluate(r, v.coords) +inner_mask = r_vals.flatten() < r_internal +v_mag = np.sqrt(v.data[:, 0]**2 + v.data[:, 1]**2) +v_max_rock = v_mag[inner_mask].max() +v_max_air = v_mag[~inner_mask].max() + +ref_v_l2 = 1.8061681957e-03 +ref_p_l2 = 1.1796447277e-01 + +uw.pprint(0, "=" * 60) +uw.pprint(0, "Pinned air-interior (Dirichlet on Outer minus Internal)") +uw.pprint(0, f" Rock-region norms:") +uw.pprint(0, f" Velocity L2: {v_l2_rock:.10e} (ref: {ref_v_l2:.10e})") +uw.pprint(0, f" Pressure L2: {p_l2_rock:.10e} (ref: {ref_p_l2:.10e})") +uw.pprint(0, f" Relative errors:") +uw.pprint(0, f" Velocity L2: {abs(v_l2_rock - ref_v_l2) / ref_v_l2:.4e}") +uw.pprint(0, f" Pressure L2: {abs(p_l2_rock - ref_p_l2) / ref_p_l2:.4e}") +uw.pprint(0, f" Max |v| rock: {v_max_rock:.10e}") +uw.pprint(0, f" Max |v| air: {v_max_air:.10e}") +uw.pprint(0, "=" * 60) + +# --- Checkpoint --- +mesh.write_timestep("pinned_interior", meshVars=[v, p], outputPath=output_dir, index=0) +uw.pprint(0, f"Checkpoint saved to {output_dir}") diff --git a/docs/examples/submesh_investigation/test_region_ds_reference.py b/docs/examples/submesh_investigation/test_region_ds_reference.py new file mode 100644 index 000000000..b373266d6 --- /dev/null +++ b/docs/examples/submesh_investigation/test_region_ds_reference.py @@ -0,0 +1,135 @@ +""" +Reference Stokes solution on a rock-only annulus mesh. + +This establishes ground-truth velocity and pressure norms for verifying +the Region DS subdomain solving approach. The rock-only annulus here +corresponds to the Inner region of an AnnulusInternalBoundary mesh. + +Test problem: Isoviscous Stokes with smooth density, free-slip BCs. + + Rock region: r_inner=0.5, r_outer=1.0 (= r_internal of full mesh) + Density: rho = cos(n*theta) * (r/r_outer)^k + Body force: -rho * unit_r (radial gravity) + BCs: Free-slip (penalty) on both boundaries + Viscosity: 1.0 + +Usage: + pixi run -e default python tests/test_region_ds_reference.py +""" + +import underworld3 as uw +from underworld3.systems import Stokes +import numpy as np +import sympy +import os + +# --- Parameters --- + +r_outer = 1.0 # Outer radius (= r_internal of full mesh) +r_inner = 0.5 # Inner radius +cellsize = 1/16 # Mesh resolution +n = 2 # Wave number +k = 1 # Power exponent for density +vel_penalty = 1.0e6 +stokes_tol = 1.0e-6 + +output_dir = "./output/region_ds_reference/" +if uw.mpi.rank == 0: + os.makedirs(output_dir, exist_ok=True) + +# --- Mesh --- + +uw.pprint(0, f"Creating rock-only annulus: r_inner={r_inner}, r_outer={r_outer}") + +mesh = uw.meshing.Annulus( + radiusOuter=r_outer, + radiusInner=r_inner, + cellSize=cellsize, +) + +uw.pprint(0, f"Mesh chart: {mesh.dm.getChart()}") + +# --- Variables --- + +v = uw.discretisation.MeshVariable("V", mesh, mesh.dim, degree=2) +p = uw.discretisation.MeshVariable("P", mesh, 1, degree=1, continuous=True) + +# --- Coordinate system --- + +unit_rvec = mesh.CoordinateSystem.unit_e_0 +r, th = mesh.CoordinateSystem.xR +Gamma = mesh.Gamma + +# Null space: constant v_theta in x,y coordinates +v_theta_fn_xy = r * mesh.CoordinateSystem.rRotN.T * sympy.Matrix((0, 1)) + +# --- Stokes solver --- + +stokes = Stokes(mesh, velocityField=v, pressureField=p) +stokes.constitutive_model = uw.constitutive_models.ViscousFlowModel +stokes.constitutive_model.Parameters.shear_viscosity_0 = 1.0 +stokes.saddle_preconditioner = 1.0 + +# Smooth density anomaly +rho = ((r / r_outer) ** k) * sympy.cos(n * th) +gravity_fn = -1.0 * unit_rvec +stokes.bodyforce = rho * gravity_fn + +# Free-slip on both boundaries (penalty on normal velocity) +stokes.add_natural_bc(vel_penalty * Gamma.dot(v.sym) * Gamma, "Upper") +stokes.add_natural_bc(vel_penalty * Gamma.dot(v.sym) * Gamma, "Lower") + +# --- Solver options --- + +stokes.tolerance = stokes_tol +stokes.petsc_options["ksp_monitor"] = None +stokes.petsc_options["snes_type"] = "newtonls" +stokes.petsc_options["ksp_type"] = "fgmres" + +stokes.petsc_options.setValue("fieldsplit_velocity_pc_mg_type", "kaskade") +stokes.petsc_options.setValue("fieldsplit_velocity_pc_mg_cycle_type", "w") +stokes.petsc_options["fieldsplit_velocity_mg_coarse_pc_type"] = "svd" +stokes.petsc_options["fieldsplit_velocity_ksp_type"] = "fcg" +stokes.petsc_options["fieldsplit_velocity_mg_levels_ksp_type"] = "chebyshev" +stokes.petsc_options["fieldsplit_velocity_mg_levels_ksp_max_it"] = 5 +stokes.petsc_options["fieldsplit_velocity_mg_levels_ksp_converged_maxits"] = None + +# --- Solve --- + +uw.pprint(0, "Solving Stokes...") +stokes.solve(verbose=True) + +# --- Null space removal (rigid body rotation) --- + +I0 = uw.maths.Integral(mesh, v_theta_fn_xy.dot(v.sym)) +norm = I0.evaluate() +I0.fn = v_theta_fn_xy.dot(v_theta_fn_xy) +vnorm = I0.evaluate() + +dv = uw.function.evaluate(norm * v_theta_fn_xy, v.coords).reshape(-1, 2) / vnorm +v.data[...] -= dv + +# --- Compute norms --- + +v_l2_integral = uw.maths.Integral(mesh, v.sym.dot(v.sym)) +v_l2 = np.sqrt(v_l2_integral.evaluate()) + +p_l2_integral = uw.maths.Integral(mesh, p.sym.dot(p.sym)) +p_l2 = np.sqrt(p_l2_integral.evaluate()) + +# Velocity magnitude stats +v_mag = np.sqrt(v.data[:, 0] ** 2 + v.data[:, 1] ** 2) +v_max = v_mag.max() + +uw.pprint(0, "=" * 60) +uw.pprint(0, "Reference solution norms (rock-only annulus)") +uw.pprint(0, f" r_inner={r_inner}, r_outer={r_outer}") +uw.pprint(0, f" n={n}, k={k}, cellsize={cellsize}") +uw.pprint(0, f" Velocity L2 norm: {v_l2:.10e}") +uw.pprint(0, f" Pressure L2 norm: {p_l2:.10e}") +uw.pprint(0, f" Max |v|: {v_max:.10e}") +uw.pprint(0, "=" * 60) + +# --- Save checkpoint --- +mesh.write_timestep("reference", meshVars=[v, p], outputPath=output_dir, index=0) +uw.pprint(0, f"Checkpoint saved to {output_dir}") diff --git a/docs/examples/submesh_investigation/test_region_ds_submesh.py b/docs/examples/submesh_investigation/test_region_ds_submesh.py new file mode 100644 index 000000000..4e3aca1bf --- /dev/null +++ b/docs/examples/submesh_investigation/test_region_ds_submesh.py @@ -0,0 +1,217 @@ +""" +Submesh approach: extract the inner region from AnnulusInternalBoundary +and solve Stokes on it directly. + +The submesh shares exact node positions with the full mesh, so solutions +can be mapped back by coordinate matching without interpolation. + +Usage: + pixi run -e default python tests/test_region_ds_submesh.py +""" + +import underworld3 as uw +from underworld3.systems import Stokes +from underworld3.cython.petsc_discretisation import petsc_dm_filter_by_label +from underworld3.discretisation import Mesh +import numpy as np +import sympy +import os +from enum import Enum + +# --- Parameters --- + +r_outer_full = 1.5 +r_internal = 1.0 +r_inner = 0.5 +cellsize = 1/16 +n = 2 +k = 1 +stokes_tol = 1.0e-6 +vel_penalty = 1.0e6 + +output_dir = "./output/region_ds_submesh/" +if uw.mpi.rank == 0: + os.makedirs(output_dir, exist_ok=True) + +# --- Full mesh --- + +uw.pprint(0, "Creating full mesh...") +full_mesh = uw.meshing.AnnulusInternalBoundary( + radiusOuter=r_outer_full, + radiusInternal=r_internal, + radiusInner=r_inner, + cellSize=cellsize, +) +uw.pprint(0, f"Full mesh: {full_mesh.dm.getChart()}") + +# --- Extract inner region submesh --- + +uw.pprint(0, "Extracting inner region submesh via DMPlexFilter...") +subdm = petsc_dm_filter_by_label(full_mesh.dm, "Inner", 101) + +# Mark boundary faces on the submesh +subdm.markBoundaryFaces("All_Boundaries", 1001) + +# The submesh needs boundary labels. The internal boundary (r=r_internal) +# becomes the outer boundary of the submesh. We'll set up boundaries +# by radius. + +# Wrap in a UW3 Mesh +class submesh_boundaries(Enum): + Lower = 1 # r = r_inner + Upper = 2 # r = r_internal (was Internal on full mesh) + +from underworld3.coordinates import CoordinateSystemType + +rock_mesh = Mesh( + subdm, + degree=1, + qdegree=2, + boundaries=submesh_boundaries, + coordinate_system_type=CoordinateSystemType.CYLINDRICAL2D, + verbose=False, +) + +uw.pprint(0, f"Rock submesh: {rock_mesh.dm.getChart()}") + +# Check coordinates +coords = rock_mesh.X.coords +r_coords = np.sqrt(coords[:, 0]**2 + coords[:, 1]**2) +uw.pprint(0, f"Rock mesh r range: [{r_coords.min():.6f}, {r_coords.max():.6f}]") + +# --- Label boundaries by radius --- +# The submesh lost the original boundary labels. Re-label by radius. + +dm = rock_mesh.dm +dm.createLabel("UW_Boundaries") +uw_label = dm.getLabel("UW_Boundaries") +all_bd_label = dm.getLabel("All_Boundaries") + +if all_bd_label: + bd_is = all_bd_label.getStratumIS(1001) + if bd_is: + bd_points = bd_is.getIndices() + uw.pprint(0, f"Boundary points: {len(bd_points)}") + + # Get vertex coordinates for boundary points + # Only process vertices (depth 0) + depth_label = dm.getLabel("depth") + vert_is = depth_label.getStratumIS(0) + verts = set(vert_is.getIndices()) if vert_is else set() + + coord_sec = dm.getCoordinateSection() + coord_vec = dm.getCoordinatesLocal() + + n_lower = 0 + n_upper = 0 + for pt in bd_points: + if pt in verts: + off = coord_sec.getOffset(pt) + x = coord_vec.getArray()[off] + y = coord_vec.getArray()[off + 1] + radius = np.sqrt(x**2 + y**2) + + if abs(radius - r_inner) < cellsize * 0.5: + uw_label.setValue(pt, submesh_boundaries.Lower.value) + n_lower += 1 + elif abs(radius - r_internal) < cellsize * 0.5: + uw_label.setValue(pt, submesh_boundaries.Upper.value) + n_upper += 1 + else: + # Edges/faces: classify by checking if they're on inner or outer boundary + # Use closure to find vertices and determine which boundary + closure = dm.getTransitiveClosure(pt)[0] + radii = [] + for cpt in closure: + if cpt in verts: + off = coord_sec.getOffset(cpt) + x = coord_vec.getArray()[off] + y = coord_vec.getArray()[off + 1] + radii.append(np.sqrt(x**2 + y**2)) + if radii: + mean_r = np.mean(radii) + if abs(mean_r - r_inner) < cellsize * 0.5: + uw_label.setValue(pt, submesh_boundaries.Lower.value) + elif abs(mean_r - r_internal) < cellsize * 0.5: + uw_label.setValue(pt, submesh_boundaries.Upper.value) + + uw.pprint(0, f"Labeled: {n_lower} lower vertices, {n_upper} upper vertices") + +# --- Variables --- + +v = uw.discretisation.MeshVariable("V", rock_mesh, rock_mesh.dim, degree=2) +p = uw.discretisation.MeshVariable("P", rock_mesh, 1, degree=1, continuous=True) + +# --- Coordinate system --- + +unit_rvec = rock_mesh.CoordinateSystem.unit_e_0 +r, th = rock_mesh.CoordinateSystem.xR +Gamma = rock_mesh.Gamma +v_theta_fn_xy = r * rock_mesh.CoordinateSystem.rRotN.T * sympy.Matrix((0, 1)) + +# --- Stokes solver --- + +stokes = Stokes(rock_mesh, velocityField=v, pressureField=p) +stokes.constitutive_model = uw.constitutive_models.ViscousFlowModel +stokes.constitutive_model.Parameters.shear_viscosity_0 = 1.0 +stokes.saddle_preconditioner = 1.0 + +rho = ((r / r_internal) ** k) * sympy.cos(n * th) +stokes.bodyforce = rho * (-1.0 * unit_rvec) + +# Free-slip on both boundaries +stokes.add_natural_bc(vel_penalty * Gamma.dot(v.sym) * Gamma, "Upper") +stokes.add_natural_bc(vel_penalty * Gamma.dot(v.sym) * Gamma, "Lower") + +# --- Solver options --- + +stokes.tolerance = stokes_tol +stokes.petsc_options["ksp_monitor"] = None +stokes.petsc_options["snes_type"] = "newtonls" +stokes.petsc_options["ksp_type"] = "fgmres" +stokes.petsc_options.setValue("fieldsplit_velocity_pc_mg_type", "kaskade") +stokes.petsc_options.setValue("fieldsplit_velocity_pc_mg_cycle_type", "w") +stokes.petsc_options["fieldsplit_velocity_mg_coarse_pc_type"] = "svd" +stokes.petsc_options["fieldsplit_velocity_ksp_type"] = "fcg" +stokes.petsc_options["fieldsplit_velocity_mg_levels_ksp_type"] = "chebyshev" +stokes.petsc_options["fieldsplit_velocity_mg_levels_ksp_max_it"] = 5 +stokes.petsc_options["fieldsplit_velocity_mg_levels_ksp_converged_maxits"] = None + +# --- Solve --- + +uw.pprint(0, "Solving Stokes on rock submesh...") +stokes.solve(verbose=True) + +# --- Null space removal --- + +I0 = uw.maths.Integral(rock_mesh, v_theta_fn_xy.dot(v.sym)) +norm = I0.evaluate() +I0.fn = v_theta_fn_xy.dot(v_theta_fn_xy) +vnorm = I0.evaluate() +dv = uw.function.evaluate(norm * v_theta_fn_xy, v.coords).reshape(-1, 2) / vnorm +v.data[...] -= dv + +# --- Norms --- + +v_l2 = np.sqrt(uw.maths.Integral(rock_mesh, v.sym.dot(v.sym)).evaluate()) +p_l2 = np.sqrt(uw.maths.Integral(rock_mesh, p.sym.dot(p.sym)).evaluate()) +v_mag = np.sqrt(v.data[:, 0]**2 + v.data[:, 1]**2) + +ref_v_l2 = 1.8061681957e-03 +ref_p_l2 = 1.1796447277e-01 +ref_v_max = 2.1782171120e-03 + +uw.pprint(0, "=" * 60) +uw.pprint(0, "Submesh approach (DMPlexFilter inner region)") +uw.pprint(0, f" Velocity L2: {v_l2:.10e} (ref: {ref_v_l2:.10e})") +uw.pprint(0, f" Pressure L2: {p_l2:.10e} (ref: {ref_p_l2:.10e})") +uw.pprint(0, f" Max |v|: {v_mag.max():.10e} (ref: {ref_v_max:.10e})") +uw.pprint(0, f" Relative errors:") +uw.pprint(0, f" Velocity L2: {abs(v_l2 - ref_v_l2) / ref_v_l2:.4e}") +uw.pprint(0, f" Pressure L2: {abs(p_l2 - ref_p_l2) / ref_p_l2:.4e}") +uw.pprint(0, f" Max |v|: {abs(v_mag.max() - ref_v_max) / ref_v_max:.4e}") +uw.pprint(0, "=" * 60) + +# --- Checkpoint --- +rock_mesh.write_timestep("submesh", meshVars=[v, p], outputPath=output_dir, index=0) +uw.pprint(0, f"Checkpoint saved to {output_dir}") diff --git a/docs/examples/submesh_investigation/viz_region_ds_comparison.py b/docs/examples/submesh_investigation/viz_region_ds_comparison.py new file mode 100644 index 000000000..f328bc3b0 --- /dev/null +++ b/docs/examples/submesh_investigation/viz_region_ds_comparison.py @@ -0,0 +1,237 @@ +# --- +# jupyter: +# jupytext: +# formats: py:percent +# text_representation: +# extension: .py +# format_name: percent +# format_version: '1.3' +# jupytext_version: 1.18.1 +# kernelspec: +# display_name: Python 3 (ipykernel) +# language: python +# name: python3 +# --- + +# %% [markdown] +""" +# Rock-Only vs Air-Layer (dP1) Comparison + +All solves use normalised Gamma_N, penalty=1e4, tol=1e-4, discontinuous pressure. + +Three cases: +- Rock-only submesh (extracted via DMPlexFilter) +- Air-layer with eta_air=1e-3 +- Air-layer with eta_air=1e-6 +""" + +# %% +import underworld3 as uw +import underworld3.visualisation as vis +from underworld3.cython.petsc_discretisation import petsc_dm_filter_by_label +from underworld3.discretisation import Mesh +from underworld3.coordinates import CoordinateSystemType +import numpy as np +import sympy +from enum import Enum +from scipy.spatial import cKDTree + +if uw.mpi.size == 1: + import pyvista as pv + +# %% +r_inner = 0.5 +r_internal = 1.0 +r_outer_full = 1.5 +cellsize = 1/16 + +# %% [markdown] +""" +## Create meshes and load checkpoints +""" + +# %% +full_mesh = uw.meshing.AnnulusInternalBoundary( + radiusOuter=r_outer_full, radiusInternal=r_internal, + radiusInner=r_inner, cellSize=cellsize, +) + +subdm = petsc_dm_filter_by_label(full_mesh.dm, "Inner", 101) +subdm.markBoundaryFaces("All_Boundaries", 1001) + +class sub_bd(Enum): + Lower = 1; Internal = 2 + +rock_mesh = Mesh(subdm, degree=1, qdegree=2, boundaries=sub_bd, + coordinate_system_type=CoordinateSystemType.CYLINDRICAL2D) + +# Rock-only +v_rock = uw.discretisation.MeshVariable("V_rock", rock_mesh, rock_mesh.dim, degree=2) +p_rock = uw.discretisation.MeshVariable("P_rock", rock_mesh, 1, degree=1, continuous=True) +v_rock.read_timestep("rock", "V", 0, outputPath="../output/normalised_rock/") +p_rock.read_timestep("rock", "P", 0, outputPath="../output/normalised_rock/") +print(f"Rock submesh: {v_rock.data.shape[0]} v-nodes") + +# Air-layer eta=1e-3 (dP1) +v_dg3 = uw.discretisation.MeshVariable("V_dg3", full_mesh, full_mesh.dim, degree=2) +p_dg3 = uw.discretisation.MeshVariable("P_dg3", full_mesh, 1, degree=1, continuous=False) +v_dg3.read_timestep("nitsche", "V", 0, outputPath="../output/normalised_nitsche/") +p_dg3.read_timestep("nitsche", "P", 0, outputPath="../output/normalised_nitsche/") +print(f"Air-layer eta=1e-3 (dP1): {v_dg3.data.shape[0]} v-nodes") + +# Air-layer eta=1e-5 (dP1, from bootstrap) +v_dg5 = uw.discretisation.MeshVariable("V_dg5", full_mesh, full_mesh.dim, degree=2) +p_dg5 = uw.discretisation.MeshVariable("P_dg5", full_mesh, 1, degree=1, continuous=False) +v_dg5.read_timestep("eta1em05", "V", 0, outputPath="../output/bootstrap_eta1em05/") +p_dg5.read_timestep("eta1em05", "P", 0, outputPath="../output/bootstrap_eta1em05/") +print(f"Air-layer eta=1e-5 (dP1): {v_dg5.data.shape[0]} v-nodes") + +# %% [markdown] +""" +## Rock-only: velocity and pressure +""" + +# %% +if uw.mpi.size == 1: + vmag = np.sqrt(v_rock.data[:, 0]**2 + v_rock.data[:, 1]**2) + vis.plot_vector(rock_mesh, v_rock, vector_name="V_rock", vfreq=1, vmag=2e1, + clip_angle=0., cpos="xy", show_arrows=True, + clim=[0., float(vmag.max())], cmap="coolwarm") + +# %% +if uw.mpi.size == 1: + pvals = uw.function.evaluate(p_rock.sym[0, 0], p_rock.coords).flatten() + plim = float(max(abs(pvals.min()), abs(pvals.max()))) + vis.plot_scalar(rock_mesh, p_rock.sym, "P_rock", + clip_angle=0., cpos="xy", cmap="RdBu", + clim=[-plim, plim]) + +# %% [markdown] +""" +## Air-layer eta=1e-3 (dP1): velocity and pressure +""" + +# %% +if uw.mpi.size == 1: + vmag3 = np.sqrt(v_dg3.data[:, 0]**2 + v_dg3.data[:, 1]**2) + vis.plot_vector(full_mesh, v_dg3, vector_name="V_dg3", vfreq=1, vmag=2e1, + clip_angle=0., cpos="xy", show_arrows=True, + clim=[0., float(vmag3.max())], cmap="coolwarm") + +# %% +if uw.mpi.size == 1: + pvals3 = uw.function.evaluate(p_dg3.sym[0, 0], p_dg3.coords).flatten() + plim3 = float(max(abs(pvals3.min()), abs(pvals3.max()))) + vis.plot_scalar(full_mesh, p_dg3.sym, "P_dg3", + clip_angle=0., cpos="xy", cmap="RdBu", + clim=[-plim3, plim3]) + +# %% [markdown] +""" +## Air-layer eta=1e-5 (dP1): velocity and pressure +""" + +# %% +if uw.mpi.size == 1: + vmag5 = np.sqrt(v_dg5.data[:, 0]**2 + v_dg5.data[:, 1]**2) + vis.plot_vector(full_mesh, v_dg5, vector_name="V_dg5", vfreq=1, vmag=2e1, + clip_angle=0., cpos="xy", show_arrows=True, + clim=[0., float(vmag5.max())], cmap="coolwarm") + +# %% +if uw.mpi.size == 1: + pvals5 = uw.function.evaluate(p_dg5.sym[0, 0], p_dg5.coords).flatten() + plim5 = float(max(abs(pvals5.min()), abs(pvals5.max()))) + vis.plot_scalar(full_mesh, p_dg5.sym, "P_dg5", + clip_angle=0., cpos="xy", cmap="RdBu", + clim=[-plim5, plim5]) + +# %% [markdown] +""" +## Overlay: rock-only (blue), eta=1e-3 (red), eta=1e-5 (green) +""" + +# %% +if uw.mpi.size == 1: + tree = cKDTree(v_rock.coords) + + dists3, idx3 = tree.query(v_dg3.coords) + matched3 = dists3 < 1e-10 + + dists5, idx5 = tree.query(v_dg5.coords) + matched5 = dists5 < 1e-10 + + # Rock submesh + rock_pts = pv.PolyData(np.column_stack([v_rock.coords, np.zeros(len(v_rock.coords))])) + rock_pts["vectors"] = np.column_stack([v_rock.data, np.zeros(len(v_rock.data))]) + + # eta=1e-3 at matched nodes + c3 = v_dg3.coords[matched3] + d3 = v_dg3.data[matched3] + pts3 = pv.PolyData(np.column_stack([c3, np.zeros(len(c3))])) + pts3["vectors"] = np.column_stack([d3, np.zeros(len(d3))]) + + # eta=1e-5 at matched nodes + c5 = v_dg5.coords[matched5] + d5 = v_dg5.data[matched5] + pts5 = pv.PolyData(np.column_stack([c5, np.zeros(len(c5))])) + pts5["vectors"] = np.column_stack([d5, np.zeros(len(d5))]) + + vmax = max(np.sqrt(v_rock.data[:, 0]**2 + v_rock.data[:, 1]**2).max(), + np.sqrt(d3[:, 0]**2 + d3[:, 1]**2).max(), + np.sqrt(d5[:, 0]**2 + d5[:, 1]**2).max()) + factor = 0.1 / vmax if vmax > 0 else 1.0 + + rock_arrows = rock_pts.glyph(orient="vectors", scale="vectors", factor=factor) + arrows3 = pts3.glyph(orient="vectors", scale="vectors", factor=factor) + arrows5 = pts5.glyph(orient="vectors", scale="vectors", factor=factor) + + pl = pv.Plotter() + pl.add_mesh(rock_arrows, color="blue", opacity=0.7, label="Rock-only submesh") + #pl.add_mesh(arrows3, color="red", opacity=0.7, label="Air-layer eta=1e-3") + pl.add_mesh(arrows5, color="green", opacity=0.7, label="Air-layer eta=1e-5") + + theta = np.linspace(0, 2*np.pi, 200) + circle = pv.lines_from_points(np.column_stack([ + 1.0 * np.cos(theta), 1.0 * np.sin(theta), np.zeros(200) + ])) + pl.add_mesh(circle, color="black", line_width=2) + + pl.add_legend() + pl.camera_position = "xy" + pl.show() + +# %% [markdown] +""" +## Norm comparison at matched nodes +""" + +# %% +tree = cKDTree(v_rock.coords) + +dists3, idx3 = tree.query(v_dg3.coords) +matched3 = dists3 < 1e-10 + +dists5, idx5 = tree.query(v_dg5.coords) +matched5 = dists5 < 1e-10 + +v_ref = v_rock.data +v3_m = v_dg3.data[matched3] +v5_m = v_dg5.data[matched5] +v_ref3 = v_ref[idx3[matched3]] +v_ref5 = v_ref[idx5[matched5]] + +def l2(a, b): + return np.sqrt(np.sum((a - b)**2)) / np.sqrt(np.sum(b**2)) + +print(f"Matched: eta=1e-3: {matched3.sum()} nodes, eta=1e-5: {matched5.sum()} nodes") +print() +print(f"{'Metric':<22} {'eta=1e-3':>12} {'eta=1e-5':>12}") +print("-" * 48) +print(f"{'Velocity L2 rel':<22} {l2(v3_m, v_ref3):>12.4e} {l2(v5_m, v_ref5):>12.4e}") + +vmag_r3 = np.sqrt(v_ref3[:, 0]**2 + v_ref3[:, 1]**2) +vmag_3 = np.sqrt(v3_m[:, 0]**2 + v3_m[:, 1]**2) +vmag_r5 = np.sqrt(v_ref5[:, 0]**2 + v_ref5[:, 1]**2) +vmag_5 = np.sqrt(v5_m[:, 0]**2 + v5_m[:, 1]**2) +print(f"{'|v| ratio (air/rock)':<22} {vmag_3.mean()/vmag_r3.mean():>12.4f} {vmag_5.mean()/vmag_r5.mean():>12.4f}") diff --git a/src/underworld3/cython/petsc_compat.h b/src/underworld3/cython/petsc_compat.h index ac4efa517..2fe0c649e 100644 --- a/src/underworld3/cython/petsc_compat.h +++ b/src/underworld3/cython/petsc_compat.h @@ -1,5 +1,19 @@ #include "petsc.h" +// Version-compatible wrapper for DMPlexFilter. +// PETSc 3.25 added an MPI_Comm argument before the SF pointer. +static inline PetscErrorCode UW_DMPlexFilter(DM dm, DMLabel label, PetscInt value, + PetscBool useClosure, PetscBool ignoreClosure, + DM *subdm) +{ +#if PETSC_VERSION_GE(3, 25, 0) + return DMPlexFilter(dm, label, value, useClosure, ignoreClosure, + PetscObjectComm((PetscObject)dm), NULL, subdm); +#else + return DMPlexFilter(dm, label, value, useClosure, ignoreClosure, NULL, subdm); +#endif +} + // Add 1 boundary condition at a time (1 boundary, 1 component etc etc) PetscErrorCode PetscDSAddBoundary_UW(DM dm, diff --git a/src/underworld3/cython/petsc_discretisation.pyx b/src/underworld3/cython/petsc_discretisation.pyx index a28212e8a..1a581381d 100644 --- a/src/underworld3/cython/petsc_discretisation.pyx +++ b/src/underworld3/cython/petsc_discretisation.pyx @@ -60,24 +60,79 @@ def petsc_fvm_get_local_cell_sizes(mesh) -> np.array: return cell_radii, cell_centroids -def petsc_dm_create_submesh_from_label(incoming_dm, boundary_label_name, boundary_label_value, marked_faces=True) -> float: +def petsc_dm_create_submesh_from_label(incoming_dm, label_name, label_value, marked_faces=False): """ - Wraps DMPlexCreateSubmesh + Extract a submesh from a DMPlex using a label. + + Wraps DMPlexCreateSubmesh: returns a new DMPlex containing only + cells (and their closures) that have the given value in the + specified label. + + Parameters + ---------- + incoming_dm : PETSc.DM + The source DMPlex. + label_name : str + Name of the DM label to filter on. + label_value : int + Stratum value to select. + marked_faces : bool + If True, the label marks faces; if False, marks cells. + + Returns + ------- + PETSc.DM + The submesh DMPlex. """ + cdef DM c_dm = incoming_dm + cdef DM subdm = PETSc.DMPlex() + cdef PetscDMLabel dmlabel + cdef PetscInt value = label_value + cdef PetscBool mf = marked_faces + + DMGetLabel(c_dm.dm, label_name.encode('utf8'), &dmlabel) + if dmlabel == NULL: + raise ValueError(f"Label '{label_name}' not found on DM") + + CHKERRQ( DMPlexCreateSubmesh(c_dm.dm, dmlabel, value, mf, &subdm.dm) ) + + return subdm + + +def petsc_dm_filter_by_label(incoming_dm, label_name, label_value): + """ + Extract a full-dimension submesh containing only cells with the + given label value. Uses DMPlexFilter. + + Parameters + ---------- + incoming_dm : PETSc.DM + The source DMPlex. + label_name : str + Name of the DM label to filter on. + label_value : int + Stratum value to select. + + Returns + ------- + PETSc.DM + The filtered submesh (same dimension as input). + """ cdef DM c_dm = incoming_dm - cdef DM subdm + cdef DM subdm = PETSc.DMPlex() cdef PetscDMLabel dmlabel - cdef PetscInt value = boundary_label_value - cdef PetscBool markedFaces = marked_faces + cdef PetscInt value = label_value - subdm = PETSc.DM() + DMGetLabel(c_dm.dm, label_name.encode('utf8'), &dmlabel) + if dmlabel == NULL: + raise ValueError(f"Label '{label_name}' not found on DM") - DMGetLabel(c_dm.dm, "Boundary", &dmlabel) - # DMPlexCreateSubmesh(dm.dm, dmlabel, value, markedFaces, &subdm.dm) + # UW_DMPlexFilter handles the PETSc version difference (3.25 added MPI_Comm arg) + CHKERRQ( UW_DMPlexFilter(c_dm.dm, dmlabel, value, PETSC_TRUE, PETSC_FALSE, &subdm.dm) ) - return + return subdm diff --git a/src/underworld3/cython/petsc_extras.pxi b/src/underworld3/cython/petsc_extras.pxi index 91bcb1763..f36ae1d30 100644 --- a/src/underworld3/cython/petsc_extras.pxi +++ b/src/underworld3/cython/petsc_extras.pxi @@ -70,8 +70,18 @@ cdef extern from "petsc.h" nogil: PetscErrorCode PetscDSAddBdResidual( PetscDS, PetscInt, PetscDSBdResidualFn, PetscDSBdResidualFn ) PetscErrorCode DMPlexCreateSubmesh(PetscDM, PetscDMLabel label, PetscInt value, PetscBool markedFaces, PetscDM *subdm) + PetscErrorCode UW_DMPlexFilter(PetscDM, PetscDMLabel, PetscInt, PetscBool, PetscBool, PetscDM *) PetscErrorCode DMGetLabel(PetscDM dm, const char name[], PetscDMLabel *label) + # Region DS — per-cell discrete system dispatch + PetscErrorCode DMSetRegionDS(PetscDM dm, PetscDMLabel label, PetscIS fields, PetscDS ds, PetscDS dsIn) + PetscErrorCode DMGetRegionDS(PetscDM dm, PetscDMLabel label, PetscIS *fields, PetscDS *ds, PetscDS *dsIn) + PetscErrorCode DMGetRegionNumDS(PetscDM dm, PetscInt num, PetscDMLabel *label, PetscIS *fields, PetscDS *ds, PetscDS *dsIn) + PetscErrorCode DMSetRegionNumDS(PetscDM dm, PetscInt num, PetscDMLabel label, PetscIS fields, PetscDS ds, PetscDS dsIn) + PetscErrorCode DMGetNumDS(PetscDM dm, PetscInt *num) + PetscErrorCode DMGetCellDS(PetscDM dm, PetscInt point, PetscDS *ds, PetscDS *dsIn) + PetscErrorCode PetscDSSetCoordinateDimension(PetscDS ds, PetscInt dim) + # These do not appear to be in the 3.17.2 release PetscErrorCode DMProjectCoordinates(PetscDM dm, PetscFE disc) PetscErrorCode DMCreateSubDM(PetscDM, PetscInt, const PetscInt *, PetscIS *, PetscDM *) diff --git a/src/underworld3/cython/petsc_generic_snes_solvers.pyx b/src/underworld3/cython/petsc_generic_snes_solvers.pyx index 58a8b0752..079e31717 100644 --- a/src/underworld3/cython/petsc_generic_snes_solvers.pyx +++ b/src/underworld3/cython/petsc_generic_snes_solvers.pyx @@ -56,6 +56,62 @@ class SolverBaseClass(uw_object): self.petsc_options_prefix = self.name self.petsc_options = PETSc.Options(self.petsc_options_prefix) + def _check_expression_meshes(self): + """Check that all MeshVariable symbols in solver expressions + belong to this solver's mesh. + + Raises a clear error if variables from a different mesh are + found, rather than letting the JIT fail with a cryptic message. + """ + from underworld3.function.expressions import extract_meshes + + solver_mesh = self.mesh + + # Collect all sympy expressions from the solver + exprs = [] + + if hasattr(self, 'bodyforce') and self.bodyforce is not None: + if hasattr(self.bodyforce, 'atoms'): + exprs.append(self.bodyforce) + + for bc in getattr(self, 'natural_bcs', []): + for attr in ('fn_f', 'fn_F', 'fn_p'): + fn = getattr(bc, attr, None) + if fn is not None and hasattr(fn, 'atoms'): + exprs.append(fn) + + if hasattr(self, '_constitutive_model') and self._constitutive_model is not None: + cm = self._constitutive_model + # Check parameter expressions rather than cm.flux — the flux + # property triggers tensor contraction which can fail for + # some model/solver combinations before setup is complete. + if hasattr(cm, 'Parameters'): + for attr_name in dir(cm.Parameters): + if attr_name.startswith('_'): + continue + try: + val = getattr(cm.Parameters, attr_name) + if hasattr(val, 'atoms'): + exprs.append(val) + except (AttributeError, TypeError): + pass + + # Extract all meshes from all expressions + foreign_meshes = set() + for expr in exprs: + meshes = extract_meshes(expr) + for m in meshes: + if m is not solver_mesh: + foreign_meshes.add(m) + + if foreign_meshes: + raise ValueError( + f"Solver expressions contain MeshVariable symbols from " + f"{len(foreign_meshes)} foreign mesh(es). All variables in " + f"a solver expression must belong to the solver's mesh. " + f"Use var.copy_into() to transfer data before building expressions." + ) + return @@ -485,6 +541,8 @@ class SolverBaseClass(uw_object): debug_name: str = None, ): + self._check_expression_meshes() + if self.is_setup: return @@ -2329,9 +2387,10 @@ class SNES_Vector(SolverBaseClass): mesh = self.mesh dim = mesh.dim - # Surface normal components (normalised) - Gamma_N = mesh.Gamma_N - n = [Gamma_N[i] for i in range(dim)] + # Surface normal components — use projected P1 normals by default. + # These are smooth, consistently oriented, and converge in 3D. + Gamma_P1 = mesh.Gamma_P1 + n = [Gamma_P1[i] for i in range(dim)] # Constraint direction: defaults to surface normal if direction is not None: @@ -4006,7 +4065,7 @@ class SNES_Stokes_SaddlePt(SolverBaseClass): # BC = namedtuple('EssentialBC', ['components', 'fn', 'boundary', 'boundary_label_val', 'type', 'PETScID']) # self.essential_p_bcs.append(BC(components, sympy_fn, boundary, -1, 'essential', -1)) - def add_nitsche_bc(self, boundary, g=None, direction=None, normal=None, gamma=10.0, theta=1): + def add_nitsche_bc(self, boundary, g=None, direction=None, normal=None, gamma=10.0, theta=1, mask=None): r"""Add Nitsche weak enforcement of a velocity constraint along a direction. Nitsche's method provides a variationally consistent alternative to @@ -4050,6 +4109,11 @@ class SNES_Stokes_SaddlePt(SolverBaseClass): 1: symmetric (default — optimal convergence and solver efficiency) 0: incomplete (no symmetry term) -1: skew-symmetric (unconditionally stable but slower convergence) + mask : sympy expression, optional + Element-wise mask for one-sided application on internal + boundaries. Use a DG MeshVariable that is 1 on the active + side and 0 on the inactive side. The mask multiplies all + Nitsche terms so that only the active-side cell contributes. Examples -------- @@ -4084,15 +4148,16 @@ class SNES_Stokes_SaddlePt(SolverBaseClass): mesh = self.mesh dim = mesh.dim - # Surface normal components. By default use normalised PETSc facet normal. + # Surface normal components. By default use projected P1 normals + # (smooth, consistently oriented, converges in 3D). if normal is not None: if isinstance(normal, sympy.MatrixBase): n = [normal[i] for i in range(dim)] else: n = list(normal) else: - Gamma_N = mesh.Gamma_N - n = [Gamma_N[i] for i in range(dim)] + Gamma_P1 = mesh.Gamma_P1 + n = [Gamma_P1[i] for i in range(dim)] # Constraint direction: defaults to surface normal if direction is not None: @@ -4115,7 +4180,11 @@ class SNES_Stokes_SaddlePt(SolverBaseClass): # n.d — how much of the constraint direction is normal to the surface # Controls pressure coupling (vanishes when d is purely tangential) + # Use Abs to ensure sign-consistent pressure coupling regardless + # of whether PETSc face normal points inward or outward n_dot_d = sum(n[i] * d[i] for i in range(dim)) + # n_dot_d is always positive when n = d (it's |n|²), + # so sign of PETSc face normal doesn't affect this term # Mesh size (global estimate via UWexpression constant) h = uw.function.expression( @@ -4164,6 +4233,17 @@ class SNES_Stokes_SaddlePt(SolverBaseClass): # Vanishes when constraint direction is purely tangential fn_p = sympy.Matrix([n_dot_d * constraint]).as_immutable() + # Apply mask for one-sided internal boundary application + if mask is not None: + if hasattr(mask, 'sym'): + mask_expr = mask.sym[0, 0] + else: + mask_expr = mask + fn_f = (fn_f * mask_expr).as_immutable() + if fn_F is not None: + fn_F = (fn_F * mask_expr).as_immutable() + fn_p = (fn_p * mask_expr).as_immutable() + # Create the NaturalBC with all terms populated BC = namedtuple('NaturalBC', [ 'f_id', 'components', 'fn_f', 'fn_F', 'fn_p', @@ -5422,9 +5502,72 @@ class SNES_Stokes_SaddlePt(SolverBaseClass): self._attach_stokes_nullspace() + # Apply region DS if active_region was set + if hasattr(self, '_inactive_region_label') and self._inactive_region_label is not None: + self._setup_region_ds() + self.is_setup = True self.constitutive_model._solver_is_setup = True + def set_active_region(self, region_label_name, region_label_value): + """Configure the solver to assemble only on cells in the given region. + + Cells NOT in the specified region get a trivial DS (no volume + contributions) and their DOFs should be pinned via Dirichlet BCs. + + Parameters + ---------- + region_label_name : str + DM label name for the INACTIVE region (e.g., "Outer"). + region_label_value : int + Label stratum value for the inactive region (e.g., 102). + """ + self._inactive_region_label = region_label_name + self._inactive_region_value = region_label_value + self.is_setup = False + + def _setup_region_ds(self): + """Register a trivial DS for the inactive region. + + After _setup_solver populates the default DS with Stokes weak forms, + this creates an empty DS for cells in the inactive region. PETSc's + DMGetCellDS dispatches per-cell: inactive cells get the empty DS + (zero volume contributions), active cells get the default DS. + """ + cdef DM c_dm = self.dm + cdef DS ds_default = self.dm.getDS() + cdef DMLabel c_label + + # Get the inactive region label + label_name = self._inactive_region_label + bc_label = self.dm.getLabel(label_name) + if bc_label is None: + raise ValueError(f"DM label '{label_name}' not found") + c_label = bc_label + + # Create a new DS with the same fields but no weak forms + cdef DS air_ds = PETSc.DS().create(comm=self.dm.comm) + + # Copy field discretisations from the DM's fields + cdef PetscInt nfields + nfields = self.dm.getNumFields() + + for f in range(nfields): + fe, _ = self.dm.getField(f) + air_ds.setDiscretisation(f, fe) + + # Set coordinate dimension to match the default DS + CHKERRQ( PetscDSSetCoordinateDimension(air_ds.ds, self.mesh.dim) ) + + # Register the empty DS for the inactive region + CHKERRQ( DMSetRegionDS(c_dm.dm, c_label.dmlabel, NULL, air_ds.ds, NULL) ) + + # Copy to coarse levels too + for coarse_dm in self.dm_hierarchy: + self.dm.copyDS(coarse_dm) + + if uw.mpi.rank == 0 and self.verbose: + print(f"Region DS: inactive region '{label_name}' gets trivial DS", flush=True) @timing.routine_timer_decorator def solve(self, diff --git a/src/underworld3/discretisation/discretisation_mesh.py b/src/underworld3/discretisation/discretisation_mesh.py index 5545a2062..ea7a7124c 100644 --- a/src/underworld3/discretisation/discretisation_mesh.py +++ b/src/underworld3/discretisation/discretisation_mesh.py @@ -284,9 +284,11 @@ def __init__( self._mesh_version = 0 self._registered_swarms = weakref.WeakSet() self._registered_surfaces = weakref.WeakSet() # Surfaces using this mesh + self._registered_submeshes = weakref.WeakSet() # Submeshes from extract_region self._mesh_update_lock = threading.RLock() comm = PETSc.COMM_WORLD + regions = None # May be set from h5 metadata or mesh generator if isinstance(plex_or_meshfile, PETSc.DMPlex): isDistributed = plex_or_meshfile.isDistributed() @@ -411,6 +413,9 @@ class replacement_boundaries(Enum): self.filename = filename self.boundaries = boundaries self.boundary_normals = boundary_normals + self.regions = regions + self.parent = None # Set by extract_region() for submeshes + self.subpoint_is = None # IS mapping submesh points -> parent points # Wrapped imported DMPlex meshes may only expose generic Gmsh labels # such as "Face Sets". Rebuild named boundary labels from those sets so @@ -1118,6 +1123,419 @@ def clone_dm_hierarchy(self): return new_dm_hierarchy + def extract_region(self, label_name, label_value=None): + """Extract a submesh containing only cells with the given region label. + + Uses ``DMPlexFilter`` to create a new mesh sharing exact node + positions with the parent. The submesh carries a ``subpoint_is`` + mapping back to the parent for restrict/prolongate operations, + and a ``parent`` reference. + + Boundary labels from the parent survive the filter. For example, + an "Internal" boundary on the parent becomes an exterior boundary + on the submesh and can be referenced by the same name. + + Parameters + ---------- + label_name : str + DM label name identifying the region (e.g., ``"Inner"``). + label_value : int, optional + Stratum value within the label. If ``None``, uses + ``mesh.regions..value`` when available. + + Returns + ------- + Mesh + A new mesh covering only the specified region. + + Examples + -------- + >>> full_mesh = uw.meshing.AnnulusInternalBoundary(...) + >>> rock_mesh = full_mesh.extract_region("Inner") + >>> rock_mesh.parent is full_mesh + True + """ + from underworld3.cython.petsc_discretisation import petsc_dm_filter_by_label + + # Resolve label value + if label_value is None: + if self.regions is not None: + try: + label_value = self.regions[label_name].value + except KeyError: + raise ValueError( + f"Region '{label_name}' not found. " + f"Available: {[r.name for r in self.regions]}" + ) + else: + raise ValueError( + "No regions defined on this mesh. Provide label_value explicitly." + ) + + # Filter the DM + subdm = petsc_dm_filter_by_label(self.dm, label_name, label_value) + subdm.markBoundaryFaces("All_Boundaries", 1001) + + # Build boundaries enum from labels that survived the filter + # (DMPlexFilter preserves parent labels on the submesh) + surviving = {} + if self.boundaries is not None: + for b in self.boundaries: + if b.name in ("Null_Boundary", "All_Boundaries"): + continue + label = subdm.getLabel(b.name) + if label: + sis = label.getStratumIS(b.value) + if sis and sis.getSize() > 0: + surviving[b.name] = b.value + + if self.regions is not None: + for r in self.regions: + label = subdm.getLabel(r.name) + if label: + sis = label.getStratumIS(r.value) + if sis and sis.getSize() > 0: + surviving[r.name] = r.value + + sub_boundaries = Enum("Boundaries", surviving) if surviving else None + + # Get the subpoint IS before wrapping (the Mesh constructor may modify the DM) + subpoint_is = subdm.getSubpointIS() + + # Construct the submesh + sub_mesh = Mesh( + subdm, + degree=self.degree, + qdegree=self.qdegree, + boundaries=sub_boundaries, + coordinate_system_type=self.CoordinateSystemType, + verbose=False, + ) + + # Store lineage + sub_mesh.parent = self + sub_mesh.subpoint_is = subpoint_is + sub_mesh._parent_mesh_version = self._mesh_version + sub_mesh._extract_label_name = label_name + sub_mesh._extract_label_value = label_value + + # Inherit regions from parent (for nested extraction) + sub_mesh.regions = self.regions + + # Cache for DOF mappings (built lazily on first restrict/prolongate) + sub_mesh._dof_maps = {} + + # Build and cache the vertex map now (before any deformation) + sub_mesh._build_vertex_map() + + # Register with parent for coordinate sync notifications + self._registered_submeshes.add(sub_mesh) + + return sub_mesh + + def _build_vertex_map(self): + """Build vertex index mapping between submesh and parent. + + Uses coordinate matching at extraction time (before any + deformation). Cached permanently since topology doesn't change. + """ + if hasattr(self, '_vertex_map') and self._vertex_map is not None: + return self._vertex_map + + from scipy.spatial import cKDTree + + tree = cKDTree(self.X.coords) + dists, indices = tree.query(self.parent.X.coords) + matched = dists < 1.0e-10 + + # parent_rows[i] -> sub_rows[i]: matched vertex pairs + parent_rows = numpy.where(matched)[0] + sub_rows = indices[matched] + + self._vertex_map = (sub_rows, parent_rows) + return self._vertex_map + + def sync_coordinates_from_parent(self): + """Update submesh coordinates from the parent mesh. + + Called automatically when the parent mesh deforms. Uses the + cached vertex map to copy parent vertex positions to the + submesh, then calls ``_deform_mesh`` to rebuild geometry. + + Raises + ------ + ValueError + If this mesh has no parent. + """ + if self.parent is None: + raise ValueError("sync_coordinates_from_parent requires a submesh") + + sub_rows, parent_rows = self._build_vertex_map() + + new_sub_coords = numpy.array(self.X.coords) + new_sub_coords[sub_rows] = self.parent.X.coords[parent_rows] + + self._deform_mesh(new_sub_coords) + self._parent_mesh_version = self.parent._mesh_version + + def _re_extract_from_parent(self, verbose=False): + """Re-extract this submesh from the adapted parent mesh. + + Called automatically when the parent mesh adapts. Replaces the + DM, rebuilds coordinates and vertex map, and reinitialises all + MeshVariables on the new submesh (reset to zero). + + The Python object is updated in-place — external references + to this submesh remain valid. + """ + import underworld3 as uw + from underworld3.cython.petsc_discretisation import petsc_dm_filter_by_label + + if self.parent is None: + raise ValueError("_re_extract_from_parent requires a submesh") + + # Find which region label this submesh was extracted from + # (stored at extraction time) + if not hasattr(self, '_extract_label_name') or not hasattr(self, '_extract_label_value'): + raise RuntimeError( + "Cannot re-extract: submesh doesn't know its extraction label. " + "Was it created with extract_region()?" + ) + + label_name = self._extract_label_name + label_value = self._extract_label_value + + if verbose: + uw.pprint(0, f"Re-extracting submesh '{label_name}' from adapted parent...") + + # Extract new DM + new_subdm = petsc_dm_filter_by_label(self.parent.dm, label_name, label_value) + new_subdm.markBoundaryFaces("All_Boundaries", 1001) + + # Back up old variable data and coordinates for interpolation + old_vars = {} + old_var_backups = {} + for var_name, var in self._vars.items(): + if var is not None: + old_vars[var_name] = var + try: + if var._lvec is not None and var.data.size > 0: + old_var_backups[var_name] = ( + numpy.array(var.coords), # old DOF coords + numpy.array(var.data), # old DOF values + ) + except Exception: + pass + + # Update DM in-place + with self._mesh_update_lock: + self.dm = new_subdm + self.subpoint_is = new_subdm.getSubpointIS() + + # Rebuild coordinates + self._coords = uw.utilities.NDArray_With_Callback( + numpy.ndarray.view(self.dm.getCoordinatesLocal().array.reshape(-1, self.cdim)), + owner=self, + ) + + def mesh_update_callback(array, change_context): + coords = array.reshape(-1, array.owner.cdim) + self._deform_mesh(coords, verbose=False) + with self._mesh_update_lock: + self._mesh_version += 1 + return + + self._coords.add_callback(mesh_update_callback) + + self._mesh_version += 1 + self._topology_version += 1 + self.nuke_coords_and_rebuild(verbose=False) + + # Rebuild vertex map (for restrict/prolongate) + self._vertex_map = None + self._build_vertex_map() + + # Invalidate DOF maps + self._dof_maps = {} + + # Reinitialise variables on the new DM + for var_name, old_var in old_vars.items(): + try: + if old_var._lvec is not None: + old_var._lvec.destroy() + old_var._lvec = None + if old_var._gvec is not None: + old_var._gvec.destroy() + old_var._gvec = None + if hasattr(old_var, '_canonical_data'): + old_var._canonical_data = None + if hasattr(old_var, '_cached_data_array'): + old_var._cached_data_array = None + + old_var._setup_ds() + old_var._set_vec(available=True) + + # Interpolate from backed-up data via kd-tree IDW + if var_name in old_var_backups: + try: + from scipy.spatial import cKDTree + old_coords, old_data = old_var_backups[var_name] + new_coords = old_var.coords + + tree = cKDTree(old_coords) + nnn = 3 if self.dim == 2 else 4 + dists, indices = tree.query(new_coords, k=nnn) + + # Inverse distance weighting + weights = 1.0 / (dists + 1e-30) + weights /= weights.sum(axis=1, keepdims=True) + new_data = numpy.zeros_like(old_var.data) + for i in range(nnn): + new_data += weights[:, i:i+1] * old_data[indices[:, i]] + + old_var.pack_raw_data_to_petsc(new_data, sync=True) + if verbose: + uw.pprint(0, f" Submesh variable '{var_name}' transferred") + except Exception as e2: + if verbose: + uw.pprint(0, f" Submesh variable '{var_name}' reset (transfer failed: {e2})") + else: + if verbose: + uw.pprint(0, f" Submesh variable '{var_name}' reset") + except Exception as e: + if verbose: + uw.pprint(0, f" Warning: failed to reinitialise '{var_name}': {e}") + + # Mark solvers for rebuild + for solver in self._equation_systems_register: + if solver is not None and hasattr(solver, 'is_setup'): + solver.is_setup = False + + # Clear caches + self._evaluation_hash = None + self._evaluation_interpolated_results = None + if hasattr(self, '_dminterpolation_cache'): + self._dminterpolation_cache.invalidate_all(reason="submesh_re_extraction") + + self._parent_mesh_version = self.parent._mesh_version + + if verbose: + uw.pprint(0, f" Submesh re-extracted: {self.dm.getChart()}") + + def _build_dof_map(self, parent_var, sub_var): + """Build a DOF-level index mapping between parent and submesh variables. + + Uses coordinate matching on DOF coordinates (exact match from + DMPlexFilter shared nodes). Cached per variable pair. + + Returns (sub_rows, parent_rows) — numpy arrays of matching DOF indices. + """ + import numpy as np + from scipy.spatial import cKDTree + + key = (id(parent_var), id(sub_var)) + if key in self._dof_maps: + return self._dof_maps[key] + + tree = cKDTree(sub_var.coords) + dists, indices = tree.query(parent_var.coords) + matched = dists < 1.0e-10 + + # indices[matched] maps parent row → sub row + parent_rows = np.where(matched)[0] + sub_rows = indices[matched] + + if len(sub_rows) != sub_var.data.shape[0]: + import warnings + warnings.warn( + f"DOF mapping: matched {len(sub_rows)} of " + f"{sub_var.data.shape[0]} submesh DOFs" + ) + + result = (sub_rows, parent_rows) + self._dof_maps[key] = result + return result + + def restrict(self, parent_var, sub_var, mode="replace"): + """Copy data from a parent-mesh variable to a submesh variable. + + Parameters + ---------- + parent_var : MeshVariable + Source variable on the parent mesh. + sub_var : MeshVariable + Destination variable on this (sub)mesh. + mode : str + ``"replace"`` overwrites submesh values (INSERT_VALUES). + ``"add"`` adds parent values into submesh (ADD_VALUES). + + Raises + ------ + ValueError + If this mesh has no parent, or the variable meshes don't match. + """ + if self.parent is None: + raise ValueError("restrict requires a submesh (parent is None)") + if parent_var.mesh is not self.parent: + raise ValueError("parent_var must be on this mesh's parent") + if sub_var.mesh is not self: + raise ValueError("sub_var must be on this mesh") + + sub_rows, parent_rows = self._build_dof_map(parent_var, sub_var) + + # Copy, modify, then write through pack_raw_data_to_petsc + # to properly sync the PETSc Vec without callback issues + new_data = numpy.array(sub_var.data) + + if mode == "replace": + new_data[sub_rows] = parent_var.data[parent_rows] + elif mode == "add": + new_data[sub_rows] += parent_var.data[parent_rows] + else: + raise ValueError(f"mode must be 'replace' or 'add', got '{mode}'") + + sub_var.pack_raw_data_to_petsc(new_data, sync=True) + + def prolongate(self, sub_var, parent_var, mode="replace"): + """Copy data from a submesh variable to a parent-mesh variable. + + Parameters + ---------- + sub_var : MeshVariable + Source variable on this (sub)mesh. + parent_var : MeshVariable + Destination variable on the parent mesh. + mode : str + ``"replace"`` overwrites parent values at submesh DOFs. + ``"add"`` adds submesh values into parent. + + Raises + ------ + ValueError + If this mesh has no parent, or the variable meshes don't match. + """ + if self.parent is None: + raise ValueError("prolongate requires a submesh (parent is None)") + if parent_var.mesh is not self.parent: + raise ValueError("parent_var must be on this mesh's parent") + if sub_var.mesh is not self: + raise ValueError("sub_var must be on this mesh") + + sub_rows, parent_rows = self._build_dof_map(parent_var, sub_var) + + new_data = numpy.array(parent_var.data) + + if mode == "replace": + new_data[parent_rows] = sub_var.data[sub_rows] + elif mode == "add": + new_data[parent_rows] += sub_var.data[sub_rows] + else: + raise ValueError(f"mode must be 'replace' or 'add', got '{mode}'") + + parent_var.pack_raw_data_to_petsc(new_data, sync=True) + + parent_var._data_is_dirty = True + def nuke_coords_and_rebuild( self, verbose=False, @@ -1232,6 +1650,9 @@ def nuke_coords_and_rebuild( if self.dm is not self.dm_hierarchy[-1]: self.dm.copyDS(self.dm_hierarchy[-1]) + # Invalidate projected boundary normals (rebuilt lazily on access) + self._projected_normals = None + if verbose and uw.mpi.rank == 0: print( f"Mesh Spatial Discretisation Complete", @@ -1240,6 +1661,46 @@ def nuke_coords_and_rebuild( return + def _update_projected_normals(self): + """Project PETSc face normals (Gamma) onto a P1 field and normalise. + + Creates ``_projected_normals`` on first call, updates in-place + thereafter. The result is a smooth, consistently-oriented unit + normal field that works well for penalty and Nitsche BCs on + curved boundaries. + """ + import underworld3 as uw + + Gamma = self.Gamma + + if not hasattr(self, '_projected_normals') or self._projected_normals is None: + self._projected_normals = uw.discretisation.MeshVariable( + "_n_proj", self, self.cdim, degree=1, + ) + + n = self._projected_normals + for i in range(self.cdim): + n.data[:, i] = uw.function.evaluate(Gamma[i], n.coords).flatten() + + mag = numpy.sqrt(numpy.sum(n.data ** 2, axis=1)) + nonzero = mag > 1.0e-30 + n.data[nonzero] /= mag[nonzero, numpy.newaxis] + + @property + def Gamma_P1(self): + """Projected P1 boundary normals as a sympy Matrix. + + Returns the normalised, vertex-averaged PETSc face normals + as a smooth P1 field. Preferred over :attr:`Gamma_N` for + penalty and Nitsche BCs on curved boundaries — gives + consistent orientation and better convergence in 3D. + + Automatically updated when the mesh deforms. + """ + if not hasattr(self, '_projected_normals') or self._projected_normals is None: + self._update_projected_normals() + return self._projected_normals.sym + @timing.routine_timer_decorator def update_lvec(self): """ @@ -1322,6 +1783,10 @@ def _deform_mesh(self, new_coords: numpy.ndarray, verbose=False): for cb in old_callbacks: self._coords.add_callback(cb) + # Propagate coordinate changes to registered submeshes + for submesh in self._registered_submeshes: + submesh.sync_coordinates_from_parent() + return def _legacy_access(self, *writeable_vars: "MeshVariable"): @@ -2116,6 +2581,10 @@ def write(self, filename: str, index: Optional[int] = None): boundaries_dict = {i.name: i.value for i in self.boundaries} g.attrs["boundaries"] = json.dumps(boundaries_dict) + if self.regions is not None: + regions_dict = {i.name: i.value for i in self.regions} + g.attrs["regions"] = json.dumps(regions_dict) + coordinates_type_dict = { "name": self.CoordinateSystemType.name, "value": self.CoordinateSystemType.value, @@ -3162,6 +3631,26 @@ def adapt(self, metric_field, verbose=False): # Stack boundary labels for adaptation adaptivity._dm_stack_bcs(self.dm, self.boundaries, "CombinedBoundaries") + # Create cell region label if regions exist — this tells MMG to + # preserve the interface between regions during adaptation + rgLabel_name = None + if self.regions is not None: + depth_label = self.dm.getLabel("depth") + cell_is = depth_label.getStratumIS(self.dim) + if cell_is: + cells = cell_is.getIndices() + self.dm.createLabel("_CellRegions_") + rg = self.dm.getLabel("_CellRegions_") + for region in self.regions: + lab = self.dm.getLabel(region.name) + if lab: + region_is = lab.getStratumIS(region.value) + if region_is: + region_cells = set(region_is.getIndices()) & set(cells) + for c in region_cells: + rg.setValue(c, region.value) + rgLabel_name = "_CellRegions_" + # Create the metric from the field hvec = metric_field._lvec metric_vec = self.dm.metricCreateIsotropic(hvec, metric_field.field_id) @@ -3171,11 +3660,26 @@ def adapt(self, metric_field, verbose=False): print(f"[{uw.mpi.rank}] Mesh adaptation starting (nodes: ~{n_nodes_old})...", flush=True) # Perform the actual mesh adaptation - new_dm = self.dm.adaptMetric(metric_vec, bdLabel="CombinedBoundaries") + new_dm = self.dm.adaptMetric( + metric_vec, + bdLabel="CombinedBoundaries", + rgLabel=rgLabel_name, + ) # Unstack boundary labels on the new dm adaptivity._dm_unstack_bcs(new_dm, self.boundaries, "CombinedBoundaries") + # Reconstruct region labels from cell tags on the adapted mesh + if rgLabel_name and self.regions is not None: + rg_new = new_dm.getLabel(rgLabel_name) + if rg_new: + for region in self.regions: + new_dm.createLabel(region.name) + region_label = new_dm.getLabel(region.name) + region_is = rg_new.getStratumIS(region.value) + if region_is: + region_label.setStratumIS(region.value, region_is) + if verbose: n_nodes_new = new_dm.getChart()[1] - new_dm.getChart()[0] print(f"[{uw.mpi.rank}] Mesh adapted (nodes: ~{n_nodes_new})", flush=True) @@ -3192,24 +3696,24 @@ def adapt(self, metric_field, verbose=False): boundaries=self.boundaries, ) - # Note: Variable transfer is complex and may hang with large meshes. - # For now, we skip automatic transfer. Users can reinitialize variables - # after adaptation using old_var.rbf_interpolate() if needed. - if verbose and old_vars_data: - print(f"[{uw.mpi.rank}] Found {len(old_vars_data)} variables. " - "Variables will be reset; reinitialize manually if needed.", flush=True) + # Transfer variable data from old mesh to new mesh via evaluate. + # The old variables are still on `self` (old DM). Evaluate them at + # the new mesh coordinates (from temp_mesh) to get interpolated values. + new_coords = temp_mesh.X.coords + transferred_data = {} - # Store old data for potential manual recovery - old_var_data_backup = {} for var_name, old_var in old_vars_data.items(): try: - # Back up old data before adaptation - if old_var._lvec is not None: - old_var_data_backup[var_name] = old_var._lvec.array.copy() - except Exception: - pass + if old_var._lvec is not None and old_var.data.size > 0: + if verbose: + print(f"[{uw.mpi.rank}] Transferring '{var_name}'...", flush=True) + transferred_data[var_name] = uw.function.evaluate( + old_var.sym, new_coords + ) + except Exception as e: + if verbose: + print(f"[{uw.mpi.rank}] Warning: transfer of '{var_name}' failed: {e}", flush=True) - # Clean up temp mesh (we created it but won't use it for transfer) del temp_mesh # Now update this mesh's internal state @@ -3266,8 +3770,21 @@ def mesh_update_callback(array, change_context): old_var._setup_ds() old_var._set_vec(available=True) - if verbose: - print(f"[{uw.mpi.rank}] Variable '{var_name}' reset on adapted mesh", flush=True) + # Restore transferred data if available + if var_name in transferred_data: + try: + data = transferred_data[var_name] + # evaluate returns (N, a, b) shaped array; pack to (N, ncomp) + data_flat = data.reshape(old_var.data.shape) + old_var.pack_raw_data_to_petsc(data_flat, sync=True) + if verbose: + print(f"[{uw.mpi.rank}] Variable '{var_name}' transferred to adapted mesh", flush=True) + except Exception as e2: + if verbose: + print(f"[{uw.mpi.rank}] Variable '{var_name}' reset (transfer failed: {e2})", flush=True) + else: + if verbose: + print(f"[{uw.mpi.rank}] Variable '{var_name}' reset on adapted mesh", flush=True) except Exception as e: if verbose: print(f"[{uw.mpi.rank}] Warning: Failed to reinitialize '{var_name}': {e}", flush=True) @@ -3288,6 +3805,14 @@ def mesh_update_callback(array, change_context): if hasattr(self, '_dminterpolation_cache'): self._dminterpolation_cache.invalidate_all(reason="mesh_adaptation") + # Re-extract registered submeshes from the adapted parent + for submesh in list(self._registered_submeshes): + try: + submesh._re_extract_from_parent(verbose=verbose) + except Exception as e: + if verbose: + print(f"[{uw.mpi.rank}] Warning: submesh re-extraction failed: {e}", flush=True) + if verbose: print(f"[{uw.mpi.rank}] Mesh adaptation complete", flush=True) diff --git a/src/underworld3/discretisation/discretisation_mesh_variables.py b/src/underworld3/discretisation/discretisation_mesh_variables.py index 952bbd41d..d5b69aabe 100644 --- a/src/underworld3/discretisation/discretisation_mesh_variables.py +++ b/src/underworld3/discretisation/discretisation_mesh_variables.py @@ -1039,6 +1039,11 @@ def write( Write variable data to the specified mesh hdf5 data file. The file will be over-written. + Note: This is a low-level method intended to be called by wrapper + functions such as ``mesh.write_timestep()`` which handle output paths, + XDMF generation, and multi-variable coordination. Prefer using + ``mesh.write_timestep()`` for normal checkpoint and visualisation output. + Note: This is a COLLECTIVE operation - all MPI ranks must call it. Parameters diff --git a/src/underworld3/discretisation/enhanced_variables.py b/src/underworld3/discretisation/enhanced_variables.py index da4fdc75c..6b2aea316 100644 --- a/src/underworld3/discretisation/enhanced_variables.py +++ b/src/underworld3/discretisation/enhanced_variables.py @@ -479,6 +479,69 @@ def read_timestep(self, *args, **kwargs): """Read timestep data.""" return self._base_var.read_timestep(*args, **kwargs) + def copy_into(self, target): + """Copy this variable's data into a variable on a related mesh. + + Detects the parent/submesh relationship and calls restrict or + prolongate as appropriate. Both meshes must be related via + ``extract_region``. + + Parameters + ---------- + target : MeshVariable + Destination variable. Must be on the parent or a submesh + of this variable's mesh. + + Examples + -------- + >>> v_full.copy_into(v_rock) # restrict: parent → submesh + >>> v_rock.copy_into(v_full) # prolongate: submesh → parent + """ + src_mesh = self._base_var.mesh + tgt_mesh = target._base_var.mesh if hasattr(target, '_base_var') else target.mesh + + if hasattr(tgt_mesh, 'parent') and tgt_mesh.parent is src_mesh: + # target is submesh of source → restrict + tgt_mesh.restrict(self, target, mode="replace") + elif hasattr(src_mesh, 'parent') and src_mesh.parent is tgt_mesh: + # source is submesh of target → prolongate + src_mesh.prolongate(self, target, mode="replace") + else: + raise ValueError( + "copy_into requires a parent/submesh relationship between " + "the two variables' meshes. Use uw.function.evaluate() " + "for unrelated meshes." + ) + + def add_into(self, target): + """Add this variable's data into a variable on a related mesh. + + Like ``copy_into`` but uses ADD_VALUES — adds to existing + values in the target rather than replacing them. + + Parameters + ---------- + target : MeshVariable + Destination variable. Must be on the parent or a submesh + of this variable's mesh. + + Examples + -------- + >>> v_rock.add_into(v_full) # prolongate with ADD + """ + src_mesh = self._base_var.mesh + tgt_mesh = target._base_var.mesh if hasattr(target, '_base_var') else target.mesh + + if hasattr(tgt_mesh, 'parent') and tgt_mesh.parent is src_mesh: + tgt_mesh.restrict(self, target, mode="add") + elif hasattr(src_mesh, 'parent') and src_mesh.parent is tgt_mesh: + src_mesh.prolongate(self, target, mode="add") + else: + raise ValueError( + "add_into requires a parent/submesh relationship between " + "the two variables' meshes." + ) + def stats(self, *args, **kwargs): """Get statistics for the variable.""" return self._base_var.stats(*args, **kwargs) diff --git a/src/underworld3/function/expressions.py b/src/underworld3/function/expressions.py index c0e87eaef..d6e6f8cb8 100644 --- a/src/underworld3/function/expressions.py +++ b/src/underworld3/function/expressions.py @@ -284,6 +284,50 @@ def extract_expressions(fn): return atoms +def extract_meshes(fn): + """Extract all meshes referenced by MeshVariable symbols in an expression. + + Searches for UnderworldFunction (applied function) atoms and + coordinate BaseScalar atoms, collecting the meshes they belong to. + + Parameters + ---------- + fn : sympy.Expr, sympy.Matrix, or UWexpression + Expression to search. + + Returns + ------- + set + Set of Mesh objects referenced by the expression. + """ + import underworld3 + + if isinstance(fn, underworld3.function.expression): + fn = fn.sym + + if not hasattr(fn, 'atoms'): + return set() + + meshes = set() + + # Check applied functions (e.g., {Tf}(N.x, N.y)) — the function CLASS + # carries a weakref to the MeshVariable via 'meshvar' + for atom in fn.atoms(sympy.Function): + func_class = type(atom) + if hasattr(func_class, 'meshvar'): + ref = func_class.meshvar + var = ref() if callable(ref) else ref # dereference weakref + if var is not None and hasattr(var, 'mesh') and var.mesh is not None: + meshes.add(var.mesh) + + # Check coordinate base scalars (N.x, N.y, Gamma.x, etc.) + for atom in fn.atoms(sympy.vector.scalar.BaseScalar): + if hasattr(atom, 'mesh'): + meshes.add(atom.mesh) + + return meshes + + def extract_expressions_and_functions(fn): """Extract all UWexpression, Function, and coordinate atoms. diff --git a/src/underworld3/function/functions_unit_system.py b/src/underworld3/function/functions_unit_system.py index c4d275f0f..1385eb9cf 100644 --- a/src/underworld3/function/functions_unit_system.py +++ b/src/underworld3/function/functions_unit_system.py @@ -143,6 +143,17 @@ def evaluate( rbf_flag = rbf if rbf is not None else False force_l2_flag = force_l2 if force_l2 is not None else False + # Step 0: CHECK for mixed-mesh expressions + # All MeshVariable symbols must belong to the same mesh. + from .expressions import extract_meshes + expr_meshes = extract_meshes(expr) + if len(expr_meshes) > 1: + raise ValueError( + f"Expression contains MeshVariable symbols from {len(expr_meshes)} " + f"different meshes. All variables in an expression must belong to " + f"the same mesh. Use var.copy_into() to transfer data first." + ) + # Step 1: UNWRAP to canonical form (preprocessing/compiler IR) # This converts ALL expressions to a standardized form: # - UWexpressions substituted with base SI numeric values diff --git a/src/underworld3/meshing/annulus.py b/src/underworld3/meshing/annulus.py index cf8d73e91..b495b72fe 100644 --- a/src/underworld3/meshing/annulus.py +++ b/src/underworld3/meshing/annulus.py @@ -264,7 +264,7 @@ class boundary_normals(Enum): Right = new_mesh.CoordinateSystem.unit_e_1 Centre = None - new_mesh.boundary_normals = boundary_normals + # boundary_normals deprecated — use mesh.Gamma_P1 for boundary normals return new_mesh @@ -536,7 +536,7 @@ class boundary_normals(Enum): Upper = new_mesh.CoordinateSystem.unit_e_0 Centre = None - new_mesh.boundary_normals = boundary_normals + # boundary_normals deprecated — use mesh.Gamma_P1 for boundary normals # Full annulus: rigid rotation about z-axis x, y = new_mesh.X @@ -782,7 +782,7 @@ class boundary_normals(Enum): Upper = new_mesh.CoordinateSystem.unit_e_0 Centre = None - new_mesh.boundary_normals = boundary_normals + # boundary_normals deprecated — use mesh.Gamma_P1 for boundary normals return new_mesh @@ -1112,7 +1112,7 @@ class boundary_normals(Enum): ) Centre = None - new_mesh.boundary_normals = boundary_normals + # boundary_normals deprecated — use mesh.Gamma_P1 for boundary normals # Full annulus with spokes: rigid rotation about z-axis x, y = new_mesh.X @@ -1226,6 +1226,10 @@ class boundaries(Enum): Upper = 3 Centre = 10 + class regions(Enum): + Inner = 101 + Outer = 102 + if cellSize_Inner is None: cellSize_Inner = cellSize @@ -1273,10 +1277,6 @@ class boundaries(Enum): cl2 = gmsh.model.geo.add_curve_loop([c3, c4], tag=boundaries.Internal.value) - ### adding this curve loop results in the mesh not being generated correctly - ### although the internal boundary is still defined in the mesh dm - # loops = [cl2] + loops - # Outermost mesh p6 = gmsh.model.geo.add_point(radiusOuter, 0.0, 0.0, meshSize=cellSize_Outer) @@ -1287,20 +1287,21 @@ class boundaries(Enum): cl3 = gmsh.model.geo.add_curve_loop([c5, c6], tag=boundaries.Upper.value) - loops = [cl3] + loops + # Create two surfaces sharing the internal boundary (no embed needed) + if radiusInner > 0.0: + s_inner = gmsh.model.geo.add_plane_surface([cl2, cl1]) + else: + s_inner = gmsh.model.geo.add_plane_surface([cl2]) - s = gmsh.model.geo.add_plane_surface(loops) + s_outer = gmsh.model.geo.add_plane_surface([cl3, cl2]) gmsh.model.geo.synchronize() if radiusInner == 0.0: - gmsh.model.mesh.embed(0, [p1], 2, s) - - gmsh.model.geo.synchronize() - gmsh.model.mesh.embed(1, [c3, c4], 2, s) - - gmsh.model.geo.synchronize() + gmsh.model.mesh.embed(0, [p1], 2, s_inner) + gmsh.model.geo.synchronize() + # Boundary physical groups (1D) if radiusInner > 0.0: gmsh.model.addPhysicalGroup( 1, [c1, c2], boundaries.Lower.value, name=boundaries.Lower.name @@ -1324,7 +1325,15 @@ class boundaries(Enum): name=boundaries.Upper.name, ) - gmsh.model.addPhysicalGroup(2, [s], 666666, "Elements") + # Region physical groups (2D) — labels cells by region + gmsh.model.addPhysicalGroup( + 2, [s_inner], tag=regions.Inner.value, name=regions.Inner.name + ) + gmsh.model.addPhysicalGroup( + 2, [s_outer], tag=regions.Outer.value, name=regions.Outer.name + ) + gmsh.model.addPhysicalGroup(2, [s_inner, s_outer], 666666, "Elements") + gmsh.model.geo.synchronize() gmsh.model.mesh.generate(2) @@ -1398,7 +1407,8 @@ class boundary_normals(Enum): Internal = new_mesh.CoordinateSystem.unit_e_0 Centre = None - new_mesh.boundary_normals = boundary_normals + # boundary_normals deprecated — use mesh.Gamma_P1 for boundary normals + new_mesh.regions = regions # Full annulus with internal boundary: rigid rotation about z-axis x, y = new_mesh.X @@ -1681,7 +1691,7 @@ class boundary_normals(Enum): Internal = new_mesh.CoordinateSystem.unit_e_0 Centre = None - new_mesh.boundary_normals = boundary_normals + # boundary_normals deprecated — use mesh.Gamma_P1 for boundary normals # Full disc with internal boundaries: rigid rotation about z-axis x, y = new_mesh.X diff --git a/src/underworld3/meshing/cartesian.py b/src/underworld3/meshing/cartesian.py index 5c9440301..9e63dd92f 100644 --- a/src/underworld3/meshing/cartesian.py +++ b/src/underworld3/meshing/cartesian.py @@ -482,6 +482,10 @@ class boundaries_2D(Enum): Left = 14 Internal = 15 + class regions_2D(Enum): + Inner = 101 # Below internal boundary + Outer = 102 # Above internal boundary + class boundary_normals_2D(Enum): Bottom = sympy.Matrix([0, 1]) Top = sympy.Matrix([0, -1]) @@ -498,6 +502,10 @@ class boundaries_3D(Enum): Back = 16 Internal = 17 + class regions_3D(Enum): + Inner = 101 # Below internal boundary + Outer = 102 # Above internal boundary + class boundary_normals_3D(Enum): Bottom = sympy.Matrix([0, 0, 1]) Top = sympy.Matrix([0, 0, -1]) @@ -580,6 +588,9 @@ class boundary_normals_3D(Enum): gmsh.model.set_physical_name(1, l56, boundaries.Right.name) gmsh.model.add_physical_group(1, [l7], boundaries.Internal.value) gmsh.model.set_physical_name(1, l7, boundaries.Internal.name) + # Region physical groups — surface1 is below, surface2 is above + gmsh.model.addPhysicalGroup(2, [surface1], regions_2D.Inner.value, name=regions_2D.Inner.name) + gmsh.model.addPhysicalGroup(2, [surface2], regions_2D.Outer.value, name=regions_2D.Outer.name) gmsh.model.addPhysicalGroup(2, [surface1, surface2], 99999) gmsh.model.setPhysicalName(2, 99999, "Elements") @@ -731,6 +742,9 @@ class boundary_normals_3D(Enum): gmsh.model.add_physical_group(2, [back_t, back_b], boundaries.Back.value) gmsh.model.set_physical_name(2, back, boundaries.Back.name) + # Region physical groups — volume_b is below, volume_t is above + gmsh.model.addPhysicalGroup(3, [volume_b], regions_3D.Inner.value, name=regions_3D.Inner.name) + gmsh.model.addPhysicalGroup(3, [volume_t], regions_3D.Outer.value, name=regions_3D.Outer.name) gmsh.model.addPhysicalGroup(3, [volume_t, volume_b], 99999) gmsh.model.setPhysicalName(3, 99999, "Elements") @@ -882,7 +896,7 @@ def box_return_coords_to_bounds(coords): boundary_normals=boundary_normals, coordinate_system_type=CoordinateSystemType.CARTESIAN, useMultipleTags=True, - useRegions=False, + useRegions=False, # BoxInternalBoundary uses _dm_unstack_bcs instead markVertices=True, refinement=0.0, refinement_callback=None, @@ -891,6 +905,50 @@ def box_return_coords_to_bounds(coords): verbose=verbose, ) uw.adaptivity._dm_unstack_bcs(new_mesh.dm, new_mesh.boundaries, "Face Sets") + + # Create region labels by classifying cells based on centroid position + # relative to the internal boundary coordinate + if dim == 2: + new_mesh.regions = regions_2D + else: + new_mesh.regions = regions_3D + + dm = new_mesh.dm + depth_label = dm.getLabel("depth") + cell_is = depth_label.getStratumIS(dim) + + if cell_is: + cells = cell_is.getIndices() + coord_sec = dm.getCoordinateSection() + coord_vec = dm.getCoordinatesLocal() + coord_arr = coord_vec.array + + for region in new_mesh.regions: + dm.createLabel(region.name) + + inner_label = dm.getLabel(new_mesh.regions.Inner.name) + outer_label = dm.getLabel(new_mesh.regions.Outer.name) + + # z-coordinate index: 1 for 2D (y), 2 for 3D (z) + z_idx = dim - 1 + + for cell in cells: + # Compute centroid from cell vertex coordinates + closure = dm.getTransitiveClosure(cell)[0] + vert_coords = [] + for pt in closure: + ndof = coord_sec.getDof(pt) + if ndof > 0: + off = coord_sec.getOffset(pt) + vert_coords.append(coord_arr[off + z_idx]) + + if vert_coords: + centroid_z = sum(vert_coords) / len(vert_coords) + if centroid_z < zintCoord: + inner_label.setValue(cell, new_mesh.regions.Inner.value) + else: + outer_label.setValue(cell, new_mesh.regions.Outer.value) + return new_mesh diff --git a/src/underworld3/meshing/geographic.py b/src/underworld3/meshing/geographic.py index d33f2800e..760ef7ee9 100644 --- a/src/underworld3/meshing/geographic.py +++ b/src/underworld3/meshing/geographic.py @@ -385,7 +385,7 @@ class boundary_normals(Enum): sympy.Piecewise((1.0, new_mesh.CoordinateSystem.R[0] > 0.99 * radiusOuter), (0.0, True)) ) - new_mesh.boundary_normals = boundary_normals + # boundary_normals deprecated — use mesh.Gamma_P1 for boundary normals return new_mesh @@ -823,6 +823,6 @@ class boundary_normals(Enum): East = new_mesh.CoordinateSystem.geo.unit_east # Eastward at east boundary West = new_mesh.CoordinateSystem.geo.unit_west # Westward at west boundary - new_mesh.boundary_normals = boundary_normals + # boundary_normals deprecated — use mesh.Gamma_P1 for boundary normals return new_mesh diff --git a/src/underworld3/meshing/segmented.py b/src/underworld3/meshing/segmented.py index b5d264415..662c5a15a 100644 --- a/src/underworld3/meshing/segmented.py +++ b/src/underworld3/meshing/segmented.py @@ -674,7 +674,7 @@ class boundary_normals(Enum): ) Centre = None - new_mesh.boundary_normals = boundary_normals + # boundary_normals deprecated — use mesh.Gamma_P1 for boundary normals # Full segmented spherical shell: 3 rigid rotation modes x, y, z = new_mesh.X @@ -1094,7 +1094,7 @@ class boundary_normals(Enum): ) Centre = None - new_mesh.boundary_normals = boundary_normals + # boundary_normals deprecated — use mesh.Gamma_P1 for boundary normals # Solid sphere: 3 rigid rotation modes x, y, z = new_mesh.X diff --git a/src/underworld3/meshing/spherical.py b/src/underworld3/meshing/spherical.py index 7d74116db..9bcb5c91e 100644 --- a/src/underworld3/meshing/spherical.py +++ b/src/underworld3/meshing/spherical.py @@ -341,6 +341,10 @@ class boundaries(Enum): Internal = 12 Upper = 13 + class regions(Enum): + Inner = 101 + Outer = 102 + import gmsh if filename is None: @@ -360,67 +364,96 @@ class boundaries(Enum): gmsh.option.setNumber("General.Verbosity", gmsh_verbosity) gmsh.model.add("SphereShell_with_Internal_Surface") - p1 = gmsh.model.geo.add_point(0.0, 0.0, 0.0, meshSize=cellSize) + # Create three concentric spheres and use OCC fragment to split + # into two non-overlapping shell volumes sharing the internal surface + ball_outer = gmsh.model.occ.addSphere(0, 0, 0, radiusOuter) + ball_internal = gmsh.model.occ.addSphere(0, 0, 0, radiusInternal) + ball_inner = gmsh.model.occ.addSphere(0, 0, 0, radiusInner) - ball1_tag = gmsh.model.occ.addSphere(0, 0, 0, radiusOuter) - ball2_tag = gmsh.model.occ.addSphere(0, 0, 0, radiusInner) - # Cut the inner sphere from the outer sphere to create a shell - gmsh.model.occ.cut([(3, ball1_tag)], [(3, ball2_tag)], removeObject=True, removeTool=True) - - ball3_tag = gmsh.model.occ.addSphere(0.0, 0.0, 0.0, radiusInternal) - ball4_tag = gmsh.model.occ.addSphere(0, 0, 0, radiusInner) - # Create another inner sphere with radius r_i (for the internal sphere) - gmsh.model.occ.cut([(3, ball3_tag)], [(3, ball4_tag)], removeObject=True, removeTool=True) + # Fragment creates non-overlapping pieces from the boolean intersection + out_dimtags, out_map = gmsh.model.occ.fragment( + [(3, ball_outer)], + [(3, ball_internal), (3, ball_inner)], + ) - # Set the maximum characteristic length (mesh size) for the mesh elements - gmsh.option.setNumber("Mesh.CharacteristicLengthMax", cellSize) gmsh.model.occ.synchronize() + gmsh.option.setNumber("Mesh.CharacteristicLengthMax", cellSize) - # Embed a 2D surface into a 3D volume - # Here, 2D entities with tag 6 are embedded into a 3D entity with tag 1 - gmsh.model.mesh.embed(2, [6], 3, 1) - # Remove specific entities from the model (these repetitions) - gmsh.model.remove_entities([(3, 2)], [(2, 5)]) - gmsh.model.occ.remove([(3, 2)], [(2, 5)]) + # Identify volumes and surfaces by bounding box + # For a sphere, bbox diagonal = sqrt(3) * radius + volumes = gmsh.model.getEntities(3) + surfaces = gmsh.model.getEntities(2) - # Get all surface entities (2D) and the first volume entity (3D) + def bbox_radius(dimtag): + """Estimate the sphere radius from a bounding box diagonal.""" + bb = gmsh.model.get_bounding_box(dimtag[0], dimtag[1]) + return np.sqrt(bb[3]**2 + bb[4]**2 + bb[5]**2) / np.sqrt(3.0) + + inner_vols = [] + outer_vols = [] + solid_ball_vols = [] # r < radiusInner — to be removed + + for vol in volumes: + r_est = bbox_radius(vol) + if np.isclose(r_est, radiusInner, atol=cellSize): + solid_ball_vols.append(vol) + elif np.isclose(r_est, radiusInternal, atol=cellSize): + inner_vols.append(vol) + elif np.isclose(r_est, radiusOuter, atol=cellSize): + outer_vols.append(vol) + + # Remove the solid inner ball (r < radiusInner) + if solid_ball_vols: + gmsh.model.occ.remove(solid_ball_vols, recursive=True) + gmsh.model.occ.synchronize() + + # Re-query after removal + volumes = gmsh.model.getEntities(3) surfaces = gmsh.model.getEntities(2) - volume = gmsh.model.getEntities(3)[0] - # Loop through all surface entities to categorize them based on their bounding box + # Classify surfaces by bounding box radius for surface in surfaces: - if np.isclose(gmsh.model.get_bounding_box(surface[0], surface[1])[-1], radiusInner): + r_est = bbox_radius(surface) + if np.isclose(r_est, radiusInner, atol=cellSize * 0.5): gmsh.model.addPhysicalGroup( - surface[0], - [surface[1]], - boundaries.Lower.value, - name=boundaries.Lower.name, + surface[0], [surface[1]], + boundaries.Lower.value, name=boundaries.Lower.name, ) - print("Created inner boundary surface") - elif np.isclose(gmsh.model.get_bounding_box(surface[0], surface[1])[-1], radiusOuter): + elif np.isclose(r_est, radiusOuter, atol=cellSize * 0.5): gmsh.model.addPhysicalGroup( - surface[0], - [surface[1]], - boundaries.Upper.value, - name=boundaries.Upper.name, + surface[0], [surface[1]], + boundaries.Upper.value, name=boundaries.Upper.name, ) - print("Created outer boundary surface") - elif np.isclose( - gmsh.model.get_bounding_box(surface[0], surface[1])[-1], radiusInternal - ): + elif np.isclose(r_est, radiusInternal, atol=cellSize * 0.5): gmsh.model.addPhysicalGroup( - surface[0], - [surface[1]], - boundaries.Internal.value, - name=boundaries.Internal.name, + surface[0], [surface[1]], + boundaries.Internal.value, name=boundaries.Internal.name, ) - print("Created internal boundary surface") - # Add the volume entity to a physical group with a high tag number (99999) and name it "Elements" - gmsh.model.addPhysicalGroup(volume[0], [volume[1]], 99999) - gmsh.model.setPhysicalName(volume[1], 99999, "Elements") - - gmsh.model.occ.synchronize() + # Classify remaining volumes into Inner and Outer + inner_vol_tags = [v[1] for v in inner_vols if v not in solid_ball_vols] + outer_vol_tags = [v[1] for v in outer_vols] + # Re-classify from current volumes in case tags changed after removal + inner_vol_tags = [] + outer_vol_tags = [] + for vol in volumes: + r_est = bbox_radius(vol) + if r_est < radiusInternal + cellSize * 0.5: + inner_vol_tags.append(vol[1]) + else: + outer_vol_tags.append(vol[1]) + + # Region physical groups + if inner_vol_tags: + gmsh.model.addPhysicalGroup(3, inner_vol_tags, + regions.Inner.value, name=regions.Inner.name) + if outer_vol_tags: + gmsh.model.addPhysicalGroup(3, outer_vol_tags, + regions.Outer.value, name=regions.Outer.name) + + # Combined elements group + all_vol_tags = inner_vol_tags + outer_vol_tags + gmsh.model.addPhysicalGroup(3, all_vol_tags, 99999, "Elements") gmsh.model.mesh.generate(3) gmsh.write(uw_filename) @@ -477,11 +510,9 @@ def spherical_mesh_refinement_callback(dm): verbose=verbose, ) - class boundary_normals(Enum): - Lower = 11 - Internal = 12 - Upper = 13 - Centre = 1 + # boundary_normals deprecated — use mesh.Gamma_P1 for boundary normals + + new_mesh.regions = regions # Full spherical shell with internal boundary: 3 rigid rotation modes x, y, z = new_mesh.X @@ -1025,7 +1056,7 @@ class boundary_normals(Enum): Lower = new_mesh.CoordinateSystem.unit_e_0 Upper = new_mesh.CoordinateSystem.unit_e_0 - new_mesh.boundary_normals = boundary_normals + # boundary_normals deprecated — use mesh.Gamma_P1 for boundary normals # Full cubed sphere: 3 rigid rotation modes x, y, z = new_mesh.X diff --git a/tests/parallel/test_0770_submesh_extract_mpi.py b/tests/parallel/test_0770_submesh_extract_mpi.py new file mode 100644 index 000000000..93c4219e9 --- /dev/null +++ b/tests/parallel/test_0770_submesh_extract_mpi.py @@ -0,0 +1,219 @@ +""" +MPI tests for submesh extraction and data transfer. + +Tests that extract_region, restrict, prolongate, copy_into, +and coordinate sync all work correctly in parallel. +""" + +import numpy as np +import pytest +import underworld3 as uw + +pytestmark = [ + pytest.mark.level_2, + pytest.mark.tier_b, + pytest.mark.mpi(min_size=2), + pytest.mark.timeout(120), +] + + +def _make_meshes(): + full = uw.meshing.AnnulusInternalBoundary( + radiusOuter=1.5, + radiusInternal=1.0, + radiusInner=0.5, + cellSize=1.0 / 8.0, + ) + rock = full.extract_region("Inner") + return full, rock + + +# ------------------------------------------------------------------ +# 1. extract_region produces a valid submesh in parallel +# ------------------------------------------------------------------ + +def test_extract_region_parallel(): + full, rock = _make_meshes() + + # Submesh should exist and have correct dimension + assert rock.dm.getDimension() == 2 + assert rock.parent is full + assert rock.subpoint_is is not None + + # All submesh vertex radii should be in [r_inner, r_internal] + coords = rock.X.coords + r = np.sqrt(coords[:, 0] ** 2 + coords[:, 1] ** 2) + assert r.min() >= 0.5 - 1e-10, f"r_min={r.min()} < 0.5" + assert r.max() <= 1.0 + 1e-10, f"r_max={r.max()} > 1.0" + + +# ------------------------------------------------------------------ +# 2. restrict: parent → submesh data transfer +# ------------------------------------------------------------------ + +def test_restrict_parallel(): + full, rock = _make_meshes() + + v_full = uw.discretisation.MeshVariable("Vf", full, full.dim, degree=2) + v_rock = uw.discretisation.MeshVariable("Vr", rock, rock.dim, degree=2) + + # Set parent to a known function of coordinates + r_f = np.sqrt(v_full.coords[:, 0] ** 2 + v_full.coords[:, 1] ** 2) + v_full.data[:, 0] = r_f + v_full.data[:, 1] = -r_f + + rock.restrict(v_full, v_rock) + + # Check submesh values match the function at submesh coordinates + r_r = np.sqrt(v_rock.coords[:, 0] ** 2 + v_rock.coords[:, 1] ** 2) + err = np.abs(v_rock.data[:, 0] - r_r).max() + + # Gather max error across ranks + from mpi4py import MPI + + global_err = MPI.COMM_WORLD.allreduce(err, op=MPI.MAX) + assert global_err < 1e-10, f"restrict error: {global_err}" + + +# ------------------------------------------------------------------ +# 3. prolongate: submesh → parent data transfer +# ------------------------------------------------------------------ + +def test_prolongate_parallel(): + full, rock = _make_meshes() + + v_full = uw.discretisation.MeshVariable("Vf", full, full.dim, degree=2) + v_rock = uw.discretisation.MeshVariable("Vr", rock, rock.dim, degree=2) + + # Set submesh to known function + r_r = np.sqrt(v_rock.coords[:, 0] ** 2 + v_rock.coords[:, 1] ** 2) + v_rock.data[:, 0] = r_r + v_rock.data[:, 1] = -r_r + + # Clear parent and prolongate + v_full.data[:] = 0.0 + rock.prolongate(v_rock, v_full) + + # Check: rock-region DOFs should be set, air DOFs should be zero + r_f = np.sqrt(v_full.coords[:, 0] ** 2 + v_full.coords[:, 1] ** 2) + rock_mask = r_f < 1.0 + 1e-6 + air_mask = ~rock_mask + + if rock_mask.any(): + rock_err = np.abs(v_full.data[rock_mask, 0] - r_f[rock_mask]).max() + else: + rock_err = 0.0 + + if air_mask.any(): + air_max = np.abs(v_full.data[air_mask]).max() + else: + air_max = 0.0 + + from mpi4py import MPI + + global_rock_err = MPI.COMM_WORLD.allreduce(rock_err, op=MPI.MAX) + global_air_max = MPI.COMM_WORLD.allreduce(air_max, op=MPI.MAX) + + assert global_rock_err < 1e-10, f"prolongate rock error: {global_rock_err}" + assert global_air_max < 1e-10, f"prolongate air leakage: {global_air_max}" + + +# ------------------------------------------------------------------ +# 4. copy_into works in both directions +# ------------------------------------------------------------------ + +def test_copy_into_parallel(): + full, rock = _make_meshes() + + v_full = uw.discretisation.MeshVariable("Vf", full, full.dim, degree=2) + v_rock = uw.discretisation.MeshVariable("Vr", rock, rock.dim, degree=2) + + # Parent → submesh + r_f = np.sqrt(v_full.coords[:, 0] ** 2 + v_full.coords[:, 1] ** 2) + v_full.data[:, 0] = r_f + + v_full.copy_into(v_rock) + + r_r = np.sqrt(v_rock.coords[:, 0] ** 2 + v_rock.coords[:, 1] ** 2) + err1 = np.abs(v_rock.data[:, 0] - r_r).max() + + # Submesh → parent + v_full.data[:] = 0.0 + v_rock.copy_into(v_full) + + rock_mask = r_f < 1.0 + 1e-6 + if rock_mask.any(): + err2 = np.abs(v_full.data[rock_mask, 0] - r_f[rock_mask]).max() + else: + err2 = 0.0 + + from mpi4py import MPI + + global_err1 = MPI.COMM_WORLD.allreduce(err1, op=MPI.MAX) + global_err2 = MPI.COMM_WORLD.allreduce(err2, op=MPI.MAX) + + assert global_err1 < 1e-10, f"copy_into restrict error: {global_err1}" + assert global_err2 < 1e-10, f"copy_into prolongate error: {global_err2}" + + +# ------------------------------------------------------------------ +# 5. Stokes solve on extracted submesh +# ------------------------------------------------------------------ + +def test_stokes_on_submesh_parallel(): + import sympy + + full, rock = _make_meshes() + + v = uw.discretisation.MeshVariable("V", rock, rock.dim, degree=2) + p = uw.discretisation.MeshVariable("P", rock, 1, degree=1, continuous=True) + + r, th = rock.CoordinateSystem.xR + G_N = rock.Gamma_N + + stokes = uw.systems.Stokes(rock, velocityField=v, pressureField=p) + stokes.constitutive_model = uw.constitutive_models.ViscousFlowModel + stokes.constitutive_model.Parameters.shear_viscosity_0 = 1.0 + stokes.bodyforce = sympy.cos(2 * th) * (-rock.CoordinateSystem.unit_e_0) + stokes.add_natural_bc(1e4 * G_N.dot(v.sym) * G_N, "Internal") + stokes.add_natural_bc(1e4 * G_N.dot(v.sym) * G_N, "Lower") + stokes.tolerance = 1e-4 + stokes.petsc_options["snes_type"] = "newtonls" + stokes.petsc_options["ksp_type"] = "fgmres" + stokes.petsc_options.setValue("fieldsplit_velocity_pc_mg_type", "kaskade") + stokes.petsc_options["fieldsplit_velocity_mg_coarse_pc_type"] = "svd" + + stokes.solve(verbose=False) + + vmag = np.sqrt(v.data[:, 0] ** 2 + v.data[:, 1] ** 2) + + from mpi4py import MPI + + global_max = MPI.COMM_WORLD.allreduce(vmag.max(), op=MPI.MAX) + + # Solution should be non-trivial + assert global_max > 1e-6, f"Stokes solution is zero: max|v|={global_max}" + # And bounded + assert global_max < 1.0, f"Stokes solution unbounded: max|v|={global_max}" + + +# ------------------------------------------------------------------ +# 6. Expression safety check works in parallel +# ------------------------------------------------------------------ + +def test_mixed_mesh_error_parallel(): + full, rock = _make_meshes() + + v_rock = uw.discretisation.MeshVariable("Vr", rock, rock.dim, degree=2) + T_full = uw.discretisation.MeshVariable("Tf", full, 1, degree=1) + + # This should raise ValueError, not a JIT error + with pytest.raises(ValueError, match="foreign mesh"): + stokes = uw.systems.Stokes( + rock, + velocityField=v_rock, + pressureField=uw.discretisation.MeshVariable("Pr", rock, 1, degree=1), + ) + stokes.constitutive_model = uw.constitutive_models.ViscousFlowModel + stokes.bodyforce = T_full.sym * rock.CoordinateSystem.unit_e_0 + stokes.solve(verbose=False)