diff --git a/docs/developer/CHANGELOG.md b/docs/developer/CHANGELOG.md index 0b6f1607..7e658218 100644 --- a/docs/developer/CHANGELOG.md +++ b/docs/developer/CHANGELOG.md @@ -273,6 +273,12 @@ component exactly — correct on curved, tilted, and deformed boundaries (#293). seam sensitivity is unchanged (#404). - Recorded as the preferred free-slip BC in the project guidance (#300); conda PETSc floor raised to ≥ 3.25 for FMG/rotation API consistency (#304). +- `Stokes.geoid(...)` is the compact spherical-benchmark entry point. It + validates the solved spherical shell and delegates to + `uw.postprocessing.spherical_shell_dynamic_response(...)`, which auto-selects + rotated reaction, constrained multiplier, or CBF topography and shares one + Appendix A linear operator between geoid and self-gravity calculations. The + direct postprocessing function remains available for advanced controls. ### Generalized Geometric Multigrid via Custom Prolongation (July 2026) diff --git a/docs/developer/subsystems/boundary-stress-and-projection-postprocessing.md b/docs/developer/subsystems/boundary-stress-and-projection-postprocessing.md index e23e2e2f..c355dab5 100644 --- a/docs/developer/subsystems/boundary-stress-and-projection-postprocessing.md +++ b/docs/developer/subsystems/boundary-stress-and-projection-postprocessing.md @@ -74,6 +74,54 @@ Combine the two rules — project the τ components with `linear_solver()` for t cheap linear solve, then compose `σ_rr` analytically — for accurate boundary stress recovery that also scales. +## 3. Use `Stokes.geoid()` for spherical-shell response + +For Zhong-style spherical-shell response functions, solve the Stokes system and +call the solver facade: + +```python +stokes.solve() + +response = stokes.geoid( + radius_inner=0.55, + radius_outer=1.0, + radius_internal=0.775, + harmonic_degree=2, + self_gravity=True, +) +``` + +`Stokes.geoid()` validates the solved spherical-shell model and internally +selects rotated-reaction, constrained-multiplier, or CBF residual topography. +Users do not choose the boundary-reaction mass recovery or duplicate the +velocity field in benchmark scripts. + +The implementation remains in +`uw.postprocessing.spherical_shell_dynamic_response(...)`. Call that function +directly only when an advanced workflow must override boundary names, +topography source, constrained reference, or physical self-gravity constants. +The facade and direct API return the same `SphericalShellDynamicResponse`. + +The Zhong Appendix A calculation is expressed as one linear operator: + +```text +N = G h + n_load +(I - Q G) h_self_gravity = h + Q n_load +``` + +where `h` contains surface and CMB topography, `N` contains their geoid +responses, and `Q` contains the two self-gravity density factors. Sharing `G` +and `n_load` between the no-self-gravity and self-gravity paths avoids separate +scalar implementations of the same coefficients. + +The current rotated/CBF harmonic projector gathers boundary samples to rank +zero and reconstructs a spherical triangulation. This is validated for the +benchmark resolutions but is not the final large-model scaling strategy. A +direct finite-element mass-weighted projection of the boundary reaction onto +the harmonic basis would avoid the gather and reconstructed triangulation; it +should replace the internal projector once the boundary-reaction API exposes +that functional directly. + ## See also - Issues [#156] (projection solver settings), [#157] (projection memory), diff --git a/src/underworld3/__init__.py b/src/underworld3/__init__.py index 818ff233..450e2368 100644 --- a/src/underworld3/__init__.py +++ b/src/underworld3/__init__.py @@ -201,6 +201,7 @@ def view(): import underworld3.maths import underworld3.swarm import underworld3.systems +import underworld3.postprocessing import underworld3.maths import underworld3.utilities import underworld3.model diff --git a/src/underworld3/postprocessing/__init__.py b/src/underworld3/postprocessing/__init__.py new file mode 100644 index 00000000..caa59eda --- /dev/null +++ b/src/underworld3/postprocessing/__init__.py @@ -0,0 +1,12 @@ +r""" +Post-processing helpers for Underworld3 models. + +The post-processing package contains derived diagnostics that are useful across +multiple model scripts but do not belong in a solver implementation. +""" + +from . import geoid +from . import topography +from .geoid import spherical_shell_dynamic_response + +__all__ = ["geoid", "topography", "spherical_shell_dynamic_response"] diff --git a/src/underworld3/postprocessing/geoid.py b/src/underworld3/postprocessing/geoid.py new file mode 100644 index 00000000..bac44640 --- /dev/null +++ b/src/underworld3/postprocessing/geoid.py @@ -0,0 +1,232 @@ +r"""Spherical-shell geoid and dynamic-response post-processing helpers.""" + +from __future__ import annotations + +from dataclasses import dataclass + +import numpy as np + + +@dataclass(frozen=True) +class SphericalShellGeoidResponse: + """No-self-gravity geoid response at the surface and CMB.""" + + surface_geoid: float + cmb_geoid: float + + def __iter__(self): + yield self.surface_geoid + yield self.cmb_geoid + + +@dataclass(frozen=True) +class SphericalShellSelfGravityResponse: + """Self-gravity corrected topography and geoid response.""" + + surface_topography: float + cmb_topography: float + surface_geoid: float + cmb_geoid: float + q_surface: float + q_cmb: float + + +@dataclass(frozen=True) +class SphericalShellDynamicResponse: + """Dynamic topography plus geoid response for a spherical shell.""" + + surface_topography: float + cmb_topography: float + surface_geoid: float + cmb_geoid: float + topography_source: str + self_gravity: SphericalShellSelfGravityResponse | None = None + + +def _spherical_shell_geoid_operator( + radius_inner: float, + radius_outer: float, + radius_internal: float, + harmonic_degree: int, +) -> tuple[np.ndarray, np.ndarray]: + """Return the Zhong Appendix A topography operator and load vector.""" + + ri = float(radius_inner) + ro = float(radius_outer) + rint = float(radius_internal) + degree = int(harmonic_degree) + if not 0.0 < ri < rint < ro: + raise ValueError("Expected 0 < radius_inner < radius_internal < radius_outer.") + if degree < 0: + raise ValueError("harmonic_degree must be non-negative.") + + denominator = float(2 * degree + 1) + operator = np.array( + [ + [ + ro / denominator, + ri * (ri / ro) ** (degree + 1) / denominator, + ], + [ + ro * (ri / ro) ** degree / denominator, + ri / denominator, + ], + ], + dtype=float, + ) + load = np.array( + [ + -rint * (rint / ro) ** (degree + 1) / denominator, + -rint * (ri / rint) ** degree / denominator, + ], + dtype=float, + ) + return operator, load + + +def spherical_shell_geoid_response( + *, + radius_inner: float, + radius_outer: float, + radius_internal: float, + harmonic_degree: int, + surface_topography: float, + cmb_topography: float, +) -> SphericalShellGeoidResponse: + """Return Zhong-style no-self-gravity surface and CMB geoid responses.""" + + operator, load = _spherical_shell_geoid_operator( + radius_inner, + radius_outer, + radius_internal, + harmonic_degree, + ) + geoid = ( + operator @ np.array([float(surface_topography), float(cmb_topography)]) + load + ) + return SphericalShellGeoidResponse(float(geoid[0]), float(geoid[1])) + + +def spherical_shell_self_gravity_response( + *, + radius_inner: float, + radius_outer: float, + radius_internal: float, + harmonic_degree: int, + surface_topography: float, + cmb_topography: float, + density_mantle: float = 3300.0, + density_cmb: float = 5400.0, + planet_radius: float = 6370000.0, + gravity: float = 9.8, + gravitational_constant: float = 6.67e-11, +) -> SphericalShellSelfGravityResponse: + """Return Zhong-style self-gravity corrected topography/geoid response.""" + + operator, load = _spherical_shell_geoid_operator( + radius_inner, + radius_outer, + radius_internal, + harmonic_degree, + ) + gravity_scale = float( + 4.0 + * np.pi + * float(gravitational_constant) + * float(planet_radius) + / float(gravity) + ) + q = gravity_scale * np.array( + [float(density_mantle), float(density_cmb)], dtype=float + ) + topography = np.array( + [float(surface_topography), float(cmb_topography)], dtype=float + ) + q_matrix = np.diag(q) + corrected_topography = np.linalg.solve( + np.eye(2) - q_matrix @ operator, + topography + q_matrix @ load, + ) + corrected_geoid = operator @ corrected_topography + load + + return SphericalShellSelfGravityResponse( + surface_topography=float(corrected_topography[0]), + cmb_topography=float(corrected_topography[1]), + surface_geoid=float(corrected_geoid[0]), + cmb_geoid=float(corrected_geoid[1]), + q_surface=float(q[0]), + q_cmb=float(q[1]), + ) + + +def spherical_shell_dynamic_response( + *, + stokes, + velocity=None, + radius_inner: float, + radius_outer: float, + radius_internal: float, + harmonic_degree: int, + include_self_gravity: bool = False, + boundary_tolerance: float | None = None, + surface_boundary: str = "Upper", + cmb_boundary: str = "Lower", + topography_source: str = "auto", + constrained_reference=None, + **self_gravity_kwargs, +) -> SphericalShellDynamicResponse: + """Return topography and geoid response for a solved spherical-shell Stokes model. + + ``topography_source="auto"`` uses constrained multiplier topography when + the solved Stokes system has boundary multipliers, rotated constraint + reactions for rotated strong free slip, and CBF-lumped residual topography + otherwise. Rotated and constrained recovery need only ``stokes``; + ``velocity`` is optional and used only by the CBF fallback. + """ + + from .topography import spherical_shell_topography_coefficients + + topography = spherical_shell_topography_coefficients( + stokes=stokes, + velocity=velocity, + radius_inner=radius_inner, + radius_outer=radius_outer, + harmonic_degree=harmonic_degree, + source=topography_source, + boundary_tolerance=boundary_tolerance, + surface_boundary=surface_boundary, + cmb_boundary=cmb_boundary, + constrained_reference=constrained_reference, + ) + surface_topography = topography.surface + cmb_topography = topography.cmb + + geoid = spherical_shell_geoid_response( + radius_inner=radius_inner, + radius_outer=radius_outer, + radius_internal=radius_internal, + harmonic_degree=harmonic_degree, + surface_topography=surface_topography, + cmb_topography=cmb_topography, + ) + + self_gravity = None + if include_self_gravity: + self_gravity = spherical_shell_self_gravity_response( + radius_inner=radius_inner, + radius_outer=radius_outer, + radius_internal=radius_internal, + harmonic_degree=harmonic_degree, + surface_topography=surface_topography, + cmb_topography=cmb_topography, + **self_gravity_kwargs, + ) + + return SphericalShellDynamicResponse( + surface_topography=surface_topography, + cmb_topography=cmb_topography, + surface_geoid=geoid.surface_geoid, + cmb_geoid=geoid.cmb_geoid, + topography_source=topography.source, + self_gravity=self_gravity, + ) diff --git a/src/underworld3/postprocessing/topography.py b/src/underworld3/postprocessing/topography.py new file mode 100644 index 00000000..87e38cff --- /dev/null +++ b/src/underworld3/postprocessing/topography.py @@ -0,0 +1,511 @@ +r"""Dynamic-topography recovery and spherical-harmonic projection utilities.""" + +from __future__ import annotations + +from dataclasses import dataclass + +from mpi4py import MPI +import numpy as np + + +@dataclass(frozen=True) +class CbfLumpedBoundaryTopography: + """CBF-lumped topography samples on a spherical boundary. + + The arrays are populated on rank 0 and set to ``None`` on other ranks. Use + :func:`cbf_lumped_spherical_harmonic_coefficient` when only a projected + scalar coefficient is required on all ranks. + """ + + coords: np.ndarray | None + topography: np.ndarray | None + harmonic_values: np.ndarray | None = None + + +@dataclass(frozen=True) +class SphericalShellTopographyCoefficients: + """Surface and CMB topography coefficients plus the source used.""" + + surface: float + cmb: float + source: str + + +def boundary_adjacent_cells(mesh, boundary: str) -> list[int]: + """Return local cells adjacent to a named boundary label.""" + + dm = mesh.dm + boundary_obj = getattr(mesh.boundaries, boundary) + label = dm.getLabel(boundary_obj.name) + if label is None: + raise ValueError(f"Boundary label '{boundary_obj.name}' was not found.") + + cell_start, cell_end = dm.getHeightStratum(0) + cells = set() + value_is = label.getValueIS() + values = [] if value_is is None else list(value_is.getIndices()) + + for value in values: + point_is = label.getStratumIS(int(value)) + if point_is is None: + continue + for point in point_is.getIndices(): + found_cell = False + for support_point in dm.getSupport(int(point)): + if cell_start <= int(support_point) < cell_end: + cells.add(int(support_point)) + found_cell = True + + if not found_cell: + star_points, _ = dm.getTransitiveClosure(int(point), useCone=False) + for star_point in star_points: + if cell_start <= int(star_point) < cell_end: + cells.add(int(star_point)) + + return sorted(cells) + + +def _boundary_node_mask(coords: np.ndarray, radius: float, tolerance: float | None): + coord_radius = np.linalg.norm(coords, axis=1) + if tolerance is None: + tolerance = max(1.0e-8, 1.0e-6 * max(1.0, abs(float(radius)))) + return coord_radius, np.abs(coord_radius - radius) < tolerance + + +def _spherical_triangle_area(a, b, c, radius: float) -> float: + det_abc = abs(float(np.dot(a, np.cross(b, c)))) + denom_abc = float(1.0 + np.dot(a, b) + np.dot(b, c) + np.dot(c, a)) + return 2.0 * np.arctan2(det_abc, denom_abc) * radius**2 + + +def _project_spherical_harmonic_samples( + coords: np.ndarray, + values: np.ndarray, + radius: float, + harmonic_degree: int, + response_sign: float, +) -> float: + """Project nodal samples over their spherical convex-hull triangulation.""" + + from scipy.spatial import ConvexHull + + unit_coords = coords / np.linalg.norm(coords, axis=1)[:, None] + hull = ConvexHull(unit_coords, qhull_options="QJ") + projected_integral = 0.0 + + for simplex in hull.simplices: + area = _spherical_triangle_area( + unit_coords[simplex[0]], + unit_coords[simplex[1]], + unit_coords[simplex[2]], + float(radius), + ) + centroid = unit_coords[simplex].mean(axis=0) + centroid /= np.linalg.norm(centroid) + centroid_harmonic = np.polynomial.legendre.Legendre.basis(int(harmonic_degree))( + np.clip(centroid[2], -1.0, 1.0) + ) + projected_integral += area * float(values[simplex].mean()) * centroid_harmonic + + harmonic_angular_norm = 4.0 * np.pi / (2 * int(harmonic_degree) + 1) + return float( + float(response_sign) + * projected_integral + / (float(radius) ** 2 * harmonic_angular_norm) + ) + + +def _gather_boundary_samples( + coords: np.ndarray, + values: np.ndarray, +) -> tuple[np.ndarray | None, np.ndarray | None]: + """Gather and average coordinate-keyed boundary samples on rank zero.""" + + local_rows = np.column_stack( + (np.asarray(coords, dtype=float), np.asarray(values, dtype=float)) + ) + gathered_rows = MPI.COMM_WORLD.gather(local_rows, root=0) + if MPI.COMM_WORLD.rank != 0: + return None, None + + rows = np.vstack([item for item in gathered_rows if item.size > 0]) + merged = {} + for row in rows: + key = tuple(np.round(row[:3], 12)) + if key not in merged: + merged[key] = [row[:3].copy(), 0.0, 0] + merged[key][1] += float(row[3]) + merged[key][2] += 1 + + global_coords = np.array([item[0] for item in merged.values()]) + global_values = np.array([item[1] / item[2] for item in merged.values()]) + return global_coords, global_values + + +def cbf_lumped_boundary_topography( + stokes, + velocity, + boundary: str, + radius: float, + normal_sign: float, + *, + harmonic_degree: int | None = None, + boundary_tolerance: float | None = None, + residual_field_id: int = 0, + velocity_residual_key: str = "velocity", +) -> CbfLumpedBoundaryTopography: + """Recover CBF-lumped topography samples on a spherical boundary. + + Parameters + ---------- + stokes + Solved Stokes system. It must provide ``compute_volume_residual_fields``. + velocity + Velocity ``MeshVariable`` used by ``stokes``. + boundary + Named mesh boundary, for example ``"Upper"`` or ``"Lower"``. + radius + Spherical boundary radius. + normal_sign + Sign of the outward unit normal relative to the radial unit vector. + Use ``+1`` on an outer surface and ``-1`` on an inner spherical surface. + harmonic_degree + Optional Legendre degree used to attach ``P_l^0(cos(theta))`` values to + the returned samples. + boundary_tolerance + Coordinate-radius tolerance used to identify boundary velocity nodes. + If omitted, a small radius-relative tolerance is used. + residual_field_id + PETSc residual field id for the velocity block. + velocity_residual_key + Key used by ``compute_volume_residual_fields`` for the velocity residual. + """ + + from scipy.spatial import ConvexHull + + mesh = stokes.mesh + adjacent_cells = boundary_adjacent_cells(mesh, boundary) + residual_fields = stokes.compute_volume_residual_fields( + cell_indices=adjacent_cells, + residual_field_id=residual_field_id, + ) + if velocity_residual_key not in residual_fields: + available = ", ".join(sorted(residual_fields)) + raise KeyError( + f"Residual field '{velocity_residual_key}' was not returned. " + f"Available fields: {available}" + ) + + velocity_residual = np.asarray(residual_fields[velocity_residual_key]).reshape( + velocity.data.shape + ) + velocity_coords = np.asarray(velocity.coords) + velocity_r, boundary_nodes = _boundary_node_mask( + velocity_coords, float(radius), boundary_tolerance + ) + if not np.any(boundary_nodes): + raise ValueError( + f"No velocity nodes found on boundary '{boundary}' at radius {radius}." + ) + + velocity_unit_r = velocity_coords / velocity_r[:, None] + normal_residual = np.einsum( + "ij,ij->i", + velocity_residual[boundary_nodes], + float(normal_sign) * velocity_unit_r[boundary_nodes], + ) + + if harmonic_degree is None: + harmonic_values = np.zeros(normal_residual.shape, dtype=float) + else: + cos_theta = np.clip( + velocity_coords[boundary_nodes, 2] / velocity_r[boundary_nodes], + -1.0, + 1.0, + ) + harmonic_values = np.polynomial.legendre.Legendre.basis(int(harmonic_degree))( + cos_theta + ) + + local_rows = np.column_stack( + ( + velocity_coords[boundary_nodes], + normal_residual, + harmonic_values, + ) + ) + gathered_rows = MPI.COMM_WORLD.gather(local_rows, root=0) + + if MPI.COMM_WORLD.rank != 0: + return CbfLumpedBoundaryTopography(None, None, None) + + rows = np.vstack([item for item in gathered_rows if item.size > 0]) + merged = {} + for row in rows: + key = tuple(np.round(row[:3], 12)) + if key not in merged: + merged[key] = [row[:3].copy(), 0.0, 0.0, 0] + merged[key][1] += float(row[3]) + merged[key][2] += float(row[4]) + merged[key][3] += 1 + + coords = np.array([item[0] for item in merged.values()]) + nodal_residual = np.array([item[1] for item in merged.values()]) + nodal_harmonic = np.array([item[2] / item[3] for item in merged.values()]) + unit_coords = coords / np.linalg.norm(coords, axis=1)[:, None] + + hull = ConvexHull(unit_coords, qhull_options="QJ") + nodal_mass = np.zeros(coords.shape[0]) + + for simplex in hull.simplices: + a, b, c = unit_coords[simplex] + area = _spherical_triangle_area(a, b, c, float(radius)) + nodal_mass[simplex] += area / 3.0 + + nodal_topography = np.zeros_like(nodal_residual) + valid = nodal_mass > 0.0 + nodal_topography[valid] = nodal_residual[valid] / nodal_mass[valid] + + return CbfLumpedBoundaryTopography(coords, nodal_topography, nodal_harmonic) + + +def cbf_lumped_spherical_harmonic_coefficient( + stokes, + velocity, + boundary: str, + radius: float, + harmonic_degree: int, + normal_sign: float, + *, + response_sign: float = 1.0, + boundary_tolerance: float | None = None, + residual_field_id: int = 0, + velocity_residual_key: str = "velocity", +) -> float: + """Return a CBF-lumped ``P_l^0`` topography coefficient on a sphere. + + The coefficient is projected with the Zhong et al. convention + ``4*pi / (2*l + 1)`` for the angular norm of ``P_l^0``. + """ + + samples = cbf_lumped_boundary_topography( + stokes=stokes, + velocity=velocity, + boundary=boundary, + radius=radius, + normal_sign=normal_sign, + harmonic_degree=harmonic_degree, + boundary_tolerance=boundary_tolerance, + residual_field_id=residual_field_id, + velocity_residual_key=velocity_residual_key, + ) + + result = None + if MPI.COMM_WORLD.rank == 0: + result = _project_spherical_harmonic_samples( + samples.coords, + samples.topography, + radius, + harmonic_degree, + response_sign, + ) + + return MPI.COMM_WORLD.bcast(result, root=0) + + +def rotated_reaction_spherical_harmonic_coefficient( + stokes, + boundary: str, + radius: float, + harmonic_degree: int, + *, + response_sign: float = 1.0, + buoyancy_scale: float = 1.0, + mass: str = "auto", +) -> float: + """Return a spherical-harmonic topography coefficient from rotated free slip. + + The boundary normal traction is recovered directly from the strong + free-slip constraint reaction. The returned topography follows + ``h = -sigma_nn / buoyancy_scale`` after the traction mean has been removed + by :meth:`stokes.boundary_normal_traction`. + + ``mass="auto"`` selects consistent surface-mass recovery for a 3D P2 + triangular trace and lumped recovery where it is valid. On faceted curved + boundaries, raw P2 vertex values are more sensitive to geometry error than + edge-midpoint values; use fitted harmonic coefficients such as this result, + or midpoint samples, for convergence studies. + """ + + xs, sigma_nn = stokes.boundary_normal_traction(boundary, mass=mass) + coords, topography = _gather_boundary_samples( + xs, + -np.asarray(sigma_nn, dtype=float) / float(buoyancy_scale), + ) + + result = None + if MPI.COMM_WORLD.rank == 0: + result = _project_spherical_harmonic_samples( + coords, + topography, + radius, + harmonic_degree, + response_sign, + ) + + return MPI.COMM_WORLD.bcast(result, root=0) + + +def spherical_harmonic_boundary_coefficient( + *, + mesh, + fn, + boundary: str, + radius: float, + harmonic_degree: int, + response_sign: float = 1.0, +) -> float: + """Project a boundary scalar function onto ``P_l^0`` on a sphere.""" + + import sympy + import underworld3 as uw + + theta = mesh.CoordinateSystem.xR[1] + cos_theta = sympy.cos(theta) + harmonic = sympy.assoc_legendre(int(harmonic_degree), 0, cos_theta) + harmonic_angular_norm = 4.0 * np.pi / (2 * int(harmonic_degree) + 1) + integral = float( + uw.maths.BdIntegral(mesh, fn=fn * harmonic, boundary=boundary).evaluate() + ) + return float( + float(response_sign) * integral / (float(radius) ** 2 * harmonic_angular_norm) + ) + + +def _has_constrained_multiplier(stokes, boundary: str) -> bool: + if not hasattr(stokes, "multiplier"): + return False + try: + return stokes.multiplier(boundary) is not None + except (AttributeError, ValueError): + return False + + +def _has_rotated_reaction(stokes, boundary: str) -> bool: + solve_info = getattr(stokes, "_rotated_freeslip_info", None) + if not solve_info: + return False + boundaries = [ + item[0] if isinstance(item, tuple) else item + for item in solve_info.get("boundaries", ()) + ] + return boundary in boundaries + + +def spherical_shell_topography_coefficients( + *, + stokes, + velocity=None, + radius_inner: float, + radius_outer: float, + harmonic_degree: int, + source: str = "auto", + boundary_tolerance: float | None = None, + surface_boundary: str = "Upper", + cmb_boundary: str = "Lower", + constrained_reference=None, +) -> SphericalShellTopographyCoefficients: + """Return surface and CMB topography coefficients from the best source. + + ``source="auto"`` uses constrained multiplier topography when the solved + Stokes system has multipliers on both requested boundaries, rotated + constraint reactions when both boundaries use rotated strong free slip, + and CBF-lumped residual recovery otherwise. Explicit sources are + ``"constrained_multiplier"``, ``"rotated_reaction"``, and + ``"cbf_lumped_residual"``. ``velocity`` is needed only by the CBF fallback; + when omitted, the Stokes velocity field is used. + """ + + if source == "auto": + if _has_constrained_multiplier( + stokes, surface_boundary + ) and _has_constrained_multiplier(stokes, cmb_boundary): + source = "constrained_multiplier" + elif _has_rotated_reaction(stokes, surface_boundary) and _has_rotated_reaction( + stokes, cmb_boundary + ): + source = "rotated_reaction" + else: + source = "cbf_lumped_residual" + + if source == "constrained_multiplier": + surface = spherical_harmonic_boundary_coefficient( + mesh=stokes.mesh, + fn=stokes.topography(surface_boundary, reference=constrained_reference), + boundary=surface_boundary, + radius=radius_outer, + harmonic_degree=harmonic_degree, + response_sign=1.0, + ) + cmb = spherical_harmonic_boundary_coefficient( + mesh=stokes.mesh, + fn=stokes.topography(cmb_boundary, reference=constrained_reference), + boundary=cmb_boundary, + radius=radius_inner, + harmonic_degree=harmonic_degree, + response_sign=1.0, + ) + elif source == "rotated_reaction": + surface = rotated_reaction_spherical_harmonic_coefficient( + stokes=stokes, + boundary=surface_boundary, + radius=radius_outer, + harmonic_degree=harmonic_degree, + response_sign=1.0, + ) + cmb = rotated_reaction_spherical_harmonic_coefficient( + stokes=stokes, + boundary=cmb_boundary, + radius=radius_inner, + harmonic_degree=harmonic_degree, + response_sign=-1.0, + ) + elif source == "cbf_lumped_residual": + if velocity is None: + velocity = getattr(getattr(stokes, "Unknowns", None), "u", None) + if velocity is None: + raise ValueError( + "CBF residual topography requires a velocity field, either on " + "stokes.Unknowns.u or through the velocity argument." + ) + surface = cbf_lumped_spherical_harmonic_coefficient( + stokes=stokes, + velocity=velocity, + boundary=surface_boundary, + radius=radius_outer, + harmonic_degree=harmonic_degree, + normal_sign=1.0, + response_sign=1.0, + boundary_tolerance=boundary_tolerance, + ) + cmb = cbf_lumped_spherical_harmonic_coefficient( + stokes=stokes, + velocity=velocity, + boundary=cmb_boundary, + radius=radius_inner, + harmonic_degree=harmonic_degree, + normal_sign=-1.0, + response_sign=-1.0, + boundary_tolerance=boundary_tolerance, + ) + else: + raise ValueError( + "source must be 'auto', 'constrained_multiplier', " + f"'rotated_reaction', or 'cbf_lumped_residual', got {source!r}" + ) + + return SphericalShellTopographyCoefficients( + surface=float(surface), + cmb=float(cmb), + source=source, + ) diff --git a/src/underworld3/systems/solvers.py b/src/underworld3/systems/solvers.py index 77e1aecd..7f5ad7a1 100644 --- a/src/underworld3/systems/solvers.py +++ b/src/underworld3/systems/solvers.py @@ -1659,6 +1659,146 @@ def solve( divergence_retries=divergence_retries, ) + def geoid( + self, + *, + radius_inner: float, + radius_outer: float, + radius_internal: float, + harmonic_degree: int, + self_gravity: bool = False, + surface_boundary: str = "Upper", + cmb_boundary: str = "Lower", + density_mantle: float = 3300.0, + density_cmb: float = 5400.0, + planet_radius: float = 6370000.0, + gravity: float = 9.8, + gravitational_constant: float = 6.67e-11, + ): + r"""Return spherical-shell topography and geoid response after a solve. + + This is a convenience facade over + :func:`uw.postprocessing.spherical_shell_dynamic_response`; the geoid + and self-gravity mathematics remain in the post-processing package. + Boundary-condition-specific topography recovery is selected internally: + constrained multiplier, rotated constraint reaction, or CBF residual. + + Parameters + ---------- + radius_inner, radius_outer, radius_internal : float + CMB, surface, and internal-load radii, ordered as + ``0 < radius_inner < radius_internal < radius_outer``. + harmonic_degree : int + Non-negative spherical-harmonic degree. + self_gravity : bool, default False + Include the Zhong-style self-gravity correction. + surface_boundary, cmb_boundary : str + Surface and CMB mesh boundary labels. + density_mantle, density_cmb : float + Density factors used when ``self_gravity=True``. + planet_radius, gravity, gravitational_constant : float + Physical constants used when ``self_gravity=True``. + + Returns + ------- + uw.postprocessing.geoid.SphericalShellDynamicResponse + Surface/CMB topography and geoid coefficients, plus the optional + self-gravity response. + + Notes + ----- + The method currently supports three-dimensional spherical-coordinate + meshes and requires a completed, converged solve. Advanced users may + call :func:`uw.postprocessing.spherical_shell_dynamic_response` + directly; both entry points use the same implementation. + """ + + from numbers import Integral + from underworld3.coordinates import CoordinateSystemType + + if self.mesh.dim != 3 or ( + self.mesh.CoordinateSystemType != CoordinateSystemType.SPHERICAL + ): + raise ValueError( + "Stokes.geoid() requires a three-dimensional spherical-coordinate mesh." + ) + + available_boundaries = {boundary.name for boundary in self.mesh.boundaries} + for role, boundary in ( + ("surface", surface_boundary), + ("CMB", cmb_boundary), + ): + if boundary not in available_boundaries: + raise ValueError( + f"Unknown {role} boundary {boundary!r}; " + f"available boundaries are {sorted(available_boundaries)}." + ) + + try: + ri = float(radius_inner) + ro = float(radius_outer) + rint = float(radius_internal) + except (TypeError, ValueError) as error: + raise TypeError( + "radius_inner, radius_internal, and radius_outer must be real numbers." + ) from error + if not np.all(np.isfinite((ri, rint, ro))) or not 0.0 < ri < rint < ro: + raise ValueError( + "Expected finite radii ordered as " + "0 < radius_inner < radius_internal < radius_outer." + ) + + if isinstance(harmonic_degree, bool) or not isinstance( + harmonic_degree, Integral + ): + raise TypeError("harmonic_degree must be an integer.") + degree = int(harmonic_degree) + if degree < 0: + raise ValueError("harmonic_degree must be non-negative.") + if not isinstance(self_gravity, bool): + raise TypeError("self_gravity must be True or False.") + + rotated_info = getattr(self, "_rotated_freeslip_info", None) + if rotated_info is not None: + converged = int(rotated_info.get("ksp_reason", 0)) > 0 and bool( + rotated_info.get("converged", True) + ) + else: + try: + converged = int(self.snes.getConvergedReason()) > 0 + except (AttributeError, TypeError): + converged = False + if not converged: + raise RuntimeError( + "Stokes.geoid() requires a completed, converged Stokes solve." + ) + + from underworld3.postprocessing import spherical_shell_dynamic_response + + self_gravity_kwargs = {} + if self_gravity: + self_gravity_kwargs = { + "density_mantle": density_mantle, + "density_cmb": density_cmb, + "planet_radius": planet_radius, + "gravity": gravity, + "gravitational_constant": gravitational_constant, + } + + return spherical_shell_dynamic_response( + stokes=self, + radius_inner=ri, + radius_outer=ro, + radius_internal=rint, + harmonic_degree=degree, + include_self_gravity=self_gravity, + surface_boundary=surface_boundary, + cmb_boundary=cmb_boundary, + topography_source="auto", + constrained_reference="mean", + **self_gravity_kwargs, + ) + @property def tau(self): r"""Deviatoric stress from the most recent solve. diff --git a/src/underworld3/utilities/boundary_flux.py b/src/underworld3/utilities/boundary_flux.py index 13824a9d..c0b8e5b2 100644 --- a/src/underworld3/utilities/boundary_flux.py +++ b/src/underworld3/utilities/boundary_flux.py @@ -12,8 +12,10 @@ The equation-specific bit — extracting the nodal reaction — is the solver method ``_assemble_volume_reaction``, which returns each rank's RAW (per-rank) volume FEM residual; the complete reaction at a boundary node shared across a partition cut is then -assembled here in ``_desmear`` by SUMMING each rank's partial contribution by coordinate -(the same rock-solid gather used for the boundary mass — no hand-rolled global assembly). +assembled here in ``_desmear`` by SUMMING each rank's partial contribution by coordinate. +In 3D, the complete boundary mass is assembled and solved once on rank zero; only the +recovered values needed by each rank are scattered back. This avoids replicating the +global P2 surface mesh and sparse solve on every rank. ``mass="auto"`` (default) uses a diagonal lumped mass where lumping is pointwise sound (the 2D P1/P2 line traces and the 3D P1 triangle trace), and the consistent @@ -280,8 +282,10 @@ def _node_reactions(xs, R, dim, boundary): def _desmear(solver, boundary, xs, R, mass, remove_mean, partial_reaction=True, edge_node_coords=None): """De-smear per-node reaction loads R (aligned with xs) into a pointwise flux via the - boundary mass, assembled globally by a coordinate-keyed allgather so every rank forms - the identical system. Returns the flux at this rank's local nodes (xs order). + boundary mass. In 3D, coordinate-keyed reactions and trace elements are gathered to + rank zero, which forms and solves the global system once; the flux values requested + by each rank are then scattered in local ``xs`` order. Returns the flux at this + rank's local nodes. ``partial_reaction`` controls how a boundary node shared across a partition cut is reconciled across ranks: ``True`` (default) SUMS each rank's contribution — correct @@ -369,126 +373,147 @@ def coord(q): ) ) - R_by = {} - for rank_values in comm.allgather(nodeR): - for key, value in rank_values.items(): - R_by[key] = ( - R_by.get(key, 0.0) + value if partial_reaction else value - ) - - elements = {} - for rank_elements in comm.allgather(local_elements): - for order, nodes, area in rank_elements: - elements[(order, tuple(sorted(nodes)))] = (order, nodes, area) - - orders = {order for order, _nodes, _area in elements.values()} - if len(orders) != 1: - raise RuntimeError( - f"Expected one trace order on boundary {boundary!r}, found {sorted(orders)}." - ) - order = orders.pop() - if mass == "auto": - mass = "consistent" if order == 2 else "lumped" - if order == 2 and mass == "lumped": - raise ValueError( - "A 3D P2 triangular trace has zero row-sum mass at its vertices; " - "use mass='consistent' (pointwise, carries the vertex-integral " - "checkerboard risk) or mass='p1' (P1-projected, monotone — the " - "choice for driving a P1 surface field) for boundary-flux recovery." - ) - mid_owners = {} - if mass == "p1": - if order != 2: - mass = "lumped" # P1 trace: p1 IS lumped - else: - # P1-PROJECTED recovery on a P2 trace: the consistent P2 path has - # the ∫φ_vertex = 0 vertex checkerboard (#404 hold), while the P1 - # trace is sound — and a P1 surface field only consumes vertex - # values anyway. Fold each edge-midpoint load onto its two edge - # vertices (φ^{P1}(edge-mid) = 1/2 exactly, P1 ⊂ P2 — the load - # transfer is the interpolation transpose, so the total load is - # conserved), then de-smear with the P1 lumped triangle mass. - # Midpoint outputs are read back as the P1 interpolant (vertex - # average). - new_elements = {} - for _order, nodes, area in elements.values(): - vk = nodes[:3] - m01, m12, m20 = nodes[3:] - mid_owners[m01] = (vk[0], vk[1]) - mid_owners[m12] = (vk[1], vk[2]) - mid_owners[m20] = (vk[2], vk[0]) - new_elements[(1, tuple(sorted(vk)))] = (1, vk, area) - folded = {} - for key, value in R_by.items(): - if key in mid_owners: - va, vb = mid_owners[key] - folded[va] = folded.get(va, 0.0) + 0.5 * value - folded[vb] = folded.get(vb, 0.0) + 0.5 * value + local_keys = [_key(x, dim) for x in xs] + gathered = comm.gather((nodeR, local_elements, local_keys), root=0) + flux_by_rank = None + root_error = None + + if comm.rank == 0: + try: + requested_keys = [rank_keys for _rank_r, _rank_e, rank_keys in gathered] + + R_by = {} + elements = {} + for rank_values, rank_elements, _rank_keys in gathered: + for key, value in rank_values.items(): + R_by[key] = R_by.get(key, 0.0) + value if partial_reaction else value + for order, nodes, area in rank_elements: + elements[(order, tuple(sorted(nodes)))] = (order, nodes, area) + del gathered + + orders = {order for order, _nodes, _area in elements.values()} + if len(orders) != 1: + raise RuntimeError( + f"Expected one trace order on boundary {boundary!r}, " + f"found {sorted(orders)}." + ) + order = orders.pop() + if mass == "auto": + mass = "consistent" if order == 2 else "lumped" + if order == 2 and mass == "lumped": + raise ValueError( + "A 3D P2 triangular trace has zero row-sum mass at its " + "vertices; use mass='consistent' (pointwise, carries the " + "vertex-integral checkerboard risk) or mass='p1' " + "(P1-projected, monotone — the choice for driving a P1 " + "surface field) for boundary-flux recovery." + ) + mid_owners = {} + if mass == "p1": + if order != 2: + mass = "lumped" # P1 trace: p1 IS lumped else: - folded[key] = folded.get(key, 0.0) + value - R_by = folded - elements = new_elements - order = 1 - mass = "lumped" - - keys = sorted(R_by) - global_index = {key: i for i, key in enumerate(keys)} - reaction = np.array([R_by[key] for key in keys], dtype=float) - if mass == "lumped": - boundary_mass = np.zeros(len(keys), dtype=float) - for _order, nodes, area in elements.values(): - for key in nodes: - boundary_mass[global_index[key]] += area / 3.0 - missing = np.flatnonzero(boundary_mass <= 0.0) - if missing.size: - raise RuntimeError( - f"Boundary mass is zero at {missing.size} nodes on {boundary!r}." - ) - flux = reaction / boundary_mass - elif mass == "consistent": - from scipy.sparse import coo_matrix - from scipy.sparse.linalg import spsolve - - rows = [] - cols = [] - values = [] - reference_mass = ( - _P1_TRIANGLE_MASS if order == 1 else _P2_TRIANGLE_MASS - ) - mass_scale = 12.0 if order == 1 else 180.0 - for _order, nodes, area in elements.values(): - indices = [global_index[key] for key in nodes] - element_mass = (area / mass_scale) * reference_mass - for i, row in enumerate(indices): - for j, col in enumerate(indices): - rows.append(row) - cols.append(col) - values.append(element_mass[i, j]) - surface_mass = coo_matrix( - (values, (rows, cols)), shape=(len(keys), len(keys)) - ).tocsr() - surface_mass.sum_duplicates() - flux = np.asarray(spsolve(surface_mass, reaction), dtype=float) - boundary_mass = np.asarray( - surface_mass @ np.ones(len(keys), dtype=float) - ) - if not np.all(np.isfinite(flux)): - raise RuntimeError( - f"Consistent boundary-mass solve failed on boundary {boundary!r}." - ) - if remove_mean: - mean = float(np.dot(flux, boundary_mass) / np.sum(boundary_mass)) - flux -= mean - - def value_at(x): - key = _key(x, dim) - if key in global_index: - return flux[global_index[key]] - # P1-projected mode: a P2 edge midpoint reads the P1 interpolant - va, vb = mid_owners[key] - return 0.5 * (flux[global_index[va]] + flux[global_index[vb]]) - - return np.array([value_at(x) for x in xs]) + # Fold each P2 midpoint load onto its two P1 vertices, then + # de-smear with the monotone P1 lumped triangle mass. + new_elements = {} + for _order, nodes, area in elements.values(): + vk = nodes[:3] + m01, m12, m20 = nodes[3:] + mid_owners[m01] = (vk[0], vk[1]) + mid_owners[m12] = (vk[1], vk[2]) + mid_owners[m20] = (vk[2], vk[0]) + new_elements[(1, tuple(sorted(vk)))] = (1, vk, area) + folded = {} + for key, value in R_by.items(): + if key in mid_owners: + va, vb = mid_owners[key] + folded[va] = folded.get(va, 0.0) + 0.5 * value + folded[vb] = folded.get(vb, 0.0) + 0.5 * value + else: + folded[key] = folded.get(key, 0.0) + value + R_by = folded + elements = new_elements + order = 1 + mass = "lumped" + + keys = sorted(R_by) + global_index = {key: i for i, key in enumerate(keys)} + reaction = np.array([R_by[key] for key in keys], dtype=float) + if mass == "lumped": + boundary_mass = np.zeros(len(keys), dtype=float) + for _order, nodes, area in elements.values(): + for key in nodes: + boundary_mass[global_index[key]] += area / 3.0 + missing = np.flatnonzero(boundary_mass <= 0.0) + if missing.size: + raise RuntimeError( + f"Boundary mass is zero at {missing.size} nodes on " f"{boundary!r}." + ) + flux = reaction / boundary_mass + elif mass == "consistent": + from scipy.sparse import coo_matrix + from scipy.sparse.linalg import spsolve + + reference_mass = _P1_TRIANGLE_MASS if order == 1 else _P2_TRIANGLE_MASS + mass_scale = 12.0 if order == 1 else 180.0 + nodes_per_element = reference_mass.shape[0] + entries_per_element = nodes_per_element**2 + entry_count = len(elements) * entries_per_element + rows = np.empty(entry_count, dtype=np.int64) + cols = np.empty(entry_count, dtype=np.int64) + values = np.empty(entry_count, dtype=float) + cursor = 0 + for _order, nodes, area in elements.values(): + indices = np.fromiter( + (global_index[key] for key in nodes), + dtype=np.int64, + count=nodes_per_element, + ) + next_cursor = cursor + entries_per_element + rows[cursor:next_cursor] = np.repeat(indices, nodes_per_element) + cols[cursor:next_cursor] = np.tile(indices, nodes_per_element) + values[cursor:next_cursor] = ((area / mass_scale) * reference_mass).ravel() + cursor = next_cursor + surface_mass = coo_matrix( + (values, (rows, cols)), shape=(len(keys), len(keys)) + ).tocsr() + del rows, cols, values + surface_mass.sum_duplicates() + flux = np.asarray(spsolve(surface_mass, reaction), dtype=float) + boundary_mass = np.asarray(surface_mass @ np.ones(len(keys), dtype=float)) + if not np.all(np.isfinite(flux)): + raise RuntimeError( + "Consistent boundary-mass solve failed on boundary " f"{boundary!r}." + ) + if remove_mean: + mean = float(np.dot(flux, boundary_mass) / np.sum(boundary_mass)) + flux -= mean + + def value_at(key): + if key in global_index: + return flux[global_index[key]] + # P1-projected mode: a P2 midpoint reads the P1 interpolant. + va, vb = mid_owners[key] + return 0.5 * (flux[global_index[va]] + flux[global_index[vb]]) + + flux_by_rank = [ + np.array([value_at(key) for key in rank_keys], dtype=float) + for rank_keys in requested_keys + ] + except Exception as error: + root_error = (type(error).__name__, str(error)) + + root_error = comm.bcast(root_error, root=0) + if root_error is not None: + error_name, error_message = root_error + error_type = { + "ValueError": ValueError, + "NotImplementedError": NotImplementedError, + "RuntimeError": RuntimeError, + }.get(error_name, RuntimeError) + raise error_type(error_message) + + return np.asarray(comm.scatter(flux_by_rank, root=0), dtype=float) if dim != 2: raise NotImplementedError( diff --git a/src/underworld3/utilities/rotated_bc.py b/src/underworld3/utilities/rotated_bc.py index 640e0864..f582a485 100644 --- a/src/underworld3/utilities/rotated_bc.py +++ b/src/underworld3/utilities/rotated_bc.py @@ -2022,9 +2022,10 @@ def boundary_normal_traction(solver, boundary, solve_result, mass="auto"): ``"lumped"``. Parallel-safe: r_c is scattered to a local vector (ghosts included) and read by LOCAL - section offset; the boundary mass is assembled globally by a coordinate-keyed - allgather of the boundary elements, so every rank produces the same de-smear and the - mean-removal gauge is global. In 3D, only triangular P1/P2 traces are supported. + section offset. In 3D, coordinate-keyed reactions and boundary elements are gathered + to rank zero, which assembles and solves the global boundary-mass system once before + scattering each rank's recovered values. The mean-removal gauge remains global. Only + triangular P1/P2 traces are supported in 3D. """ dm = solver.dm dim = solver.mesh.dim diff --git a/tests/test_1070_postprocessing_geoid.py b/tests/test_1070_postprocessing_geoid.py new file mode 100644 index 00000000..a91ac027 --- /dev/null +++ b/tests/test_1070_postprocessing_geoid.py @@ -0,0 +1,213 @@ +import math + +import numpy as np +import pytest +import sympy +import underworld3 as uw + +pytestmark = pytest.mark.level_2 + + +def test_postprocessing_module_is_public(): + assert hasattr(uw, "postprocessing") + assert hasattr(uw.postprocessing, "topography") + assert hasattr(uw.postprocessing, "geoid") + assert hasattr(uw.postprocessing, "spherical_shell_dynamic_response") + assert hasattr(uw.systems.Stokes, "geoid") + + +def test_spherical_shell_geoid_response_matches_direct_formula(): + response = uw.postprocessing.geoid.spherical_shell_geoid_response( + radius_inner=0.55, + radius_outer=1.0, + radius_internal=0.775, + harmonic_degree=2, + surface_topography=0.3918264884592169, + cmb_topography=0.2653834107104151, + ) + + surface_expected = ( + 0.55 * (0.55 / 1.0) ** 3 * 0.2653834107104151 + + 1.0 * 0.3918264884592169 + - 0.775 * (0.775 / 1.0) ** 3 + ) / 5.0 + cmb_expected = ( + 0.55 * 0.2653834107104151 + + 1.0 * (0.55 / 1.0) ** 2 * 0.3918264884592169 + - 0.775 * (0.55 / 0.775) ** 2 + ) / 5.0 + + assert math.isclose(response.surface_geoid, surface_expected) + assert math.isclose(response.cmb_geoid, cmb_expected) + + +def test_spherical_shell_self_gravity_response_matches_direct_solve(): + response = uw.postprocessing.geoid.spherical_shell_self_gravity_response( + radius_inner=0.55, + radius_outer=1.0, + radius_internal=0.775, + harmonic_degree=2, + surface_topography=0.3918264884592169, + cmb_topography=0.2653834107104151, + ) + + q_surface = 4.0 * np.pi * 6.67e-11 * 6370000.0 * 3300.0 / 9.8 + q_cmb = 4.0 * np.pi * 6.67e-11 * 6370000.0 * 5400.0 / 9.8 + denominator = 5.0 + surface_b = 0.55 * (0.55 / 1.0) ** 3 / denominator + surface_s = 1.0 / denominator + surface_load = -0.775 * (0.775 / 1.0) ** 3 / denominator + cmb_b = 0.55 / denominator + cmb_s = 1.0 * (0.55 / 1.0) ** 2 / denominator + cmb_load = -0.775 * (0.55 / 0.775) ** 2 / denominator + + matrix = np.array( + [ + [1.0 - q_surface * surface_s, -q_surface * surface_b], + [-q_cmb * cmb_s, 1.0 - q_cmb * cmb_b], + ] + ) + rhs = np.array( + [ + 0.3918264884592169 + q_surface * surface_load, + 0.2653834107104151 + q_cmb * cmb_load, + ] + ) + surface_topography, cmb_topography = np.linalg.solve(matrix, rhs) + surface_geoid = ( + surface_b * cmb_topography + surface_s * surface_topography + surface_load + ) + cmb_geoid = cmb_b * cmb_topography + cmb_s * surface_topography + cmb_load + + assert math.isclose(response.q_surface, q_surface) + assert math.isclose(response.q_cmb, q_cmb) + assert math.isclose(response.surface_topography, surface_topography) + assert math.isclose(response.cmb_topography, cmb_topography) + assert math.isclose(response.surface_geoid, surface_geoid) + assert math.isclose(response.cmb_geoid, cmb_geoid) + + +def test_stokes_geoid_rejects_non_spherical_geometry(): + mesh = uw.meshing.UnstructuredSimplexBox( + minCoords=(0.0, 0.0, 0.0), + maxCoords=(1.0, 1.0, 1.0), + cellSize=0.5, + ) + stokes = uw.systems.Stokes(mesh, degree=2) + + with pytest.raises(ValueError, match="spherical-coordinate"): + stokes.geoid( + radius_inner=0.55, + radius_outer=1.0, + radius_internal=0.775, + harmonic_degree=2, + ) + + +def test_stokes_geoid_matches_dynamic_response_for_rotated_reaction(): + """The Stokes facade exactly matches the rotated-reaction implementation.""" + + radius_inner = 0.55 + radius_outer = 1.0 + radius_internal = 0.775 + mesh = uw.meshing.SphericalShellInternalBoundary( + radiusOuter=radius_outer, + radiusInternal=radius_internal, + radiusInner=radius_inner, + cellSize=0.25, + qdegree=2, + degree=1, + ) + velocity = uw.discretisation.MeshVariable( + "U_rotated_geoid", mesh, mesh.dim, degree=2 + ) + pressure = uw.discretisation.MeshVariable( + "P_rotated_geoid", mesh, 1, degree=1, continuous=True + ) + stokes = uw.systems.Stokes(mesh, velocityField=velocity, pressureField=pressure) + stokes.constitutive_model = uw.constitutive_models.ViscousFlowModel + stokes.constitutive_model.Parameters.shear_viscosity_0 = 1.0 + + theta = mesh.CoordinateSystem.xR[1] + unit_r = mesh.CoordinateSystem.unit_e_0 + load = sympy.assoc_legendre(2, 0, sympy.cos(theta)) * unit_r + stokes.add_natural_bc(load, "Internal") + stokes.add_rotated_freeslip_bc(0, "Upper", normal=unit_r) + stokes.add_rotated_freeslip_bc(0, "Lower", normal=-unit_r) + stokes.petsc_use_nullspace = True + stokes.petsc_options["snes_type"] = "ksponly" + stokes.tolerance = 1.0e-5 + + with pytest.raises(ValueError, match="ordered"): + stokes.geoid( + radius_inner=radius_outer, + radius_outer=radius_inner, + radius_internal=radius_internal, + harmonic_degree=2, + ) + with pytest.raises(TypeError, match="harmonic_degree"): + stokes.geoid( + radius_inner=radius_inner, + radius_outer=radius_outer, + radius_internal=radius_internal, + harmonic_degree=2.0, + ) + with pytest.raises(TypeError, match="self_gravity"): + stokes.geoid( + radius_inner=radius_inner, + radius_outer=radius_outer, + radius_internal=radius_internal, + harmonic_degree=2, + self_gravity="on", + ) + with pytest.raises(RuntimeError, match="completed, converged"): + stokes.geoid( + radius_inner=radius_inner, + radius_outer=radius_outer, + radius_internal=radius_internal, + harmonic_degree=2, + ) + + stokes.solve() + + direct_response = uw.postprocessing.spherical_shell_dynamic_response( + stokes=stokes, + radius_inner=radius_inner, + radius_outer=radius_outer, + radius_internal=radius_internal, + harmonic_degree=2, + include_self_gravity=True, + topography_source="auto", + ) + response = stokes.geoid( + radius_inner=radius_inner, + radius_outer=radius_outer, + radius_internal=radius_internal, + harmonic_degree=2, + self_gravity=True, + ) + response_without_self_gravity = stokes.geoid( + radius_inner=radius_inner, + radius_outer=radius_outer, + radius_internal=radius_internal, + harmonic_degree=2, + self_gravity=False, + ) + + assert response.topography_source == "rotated_reaction" + assert response == direct_response + assert response_without_self_gravity.self_gravity is None + assert ( + response_without_self_gravity.surface_topography == response.surface_topography + ) + assert response_without_self_gravity.cmb_topography == response.cmb_topography + assert response_without_self_gravity.surface_geoid == response.surface_geoid + assert response_without_self_gravity.cmb_geoid == response.cmb_geoid + assert np.isclose(response.surface_topography, 0.41920, rtol=0.10) + assert np.isclose(response.cmb_topography, 0.77060, rtol=0.10) + assert np.isclose(response.surface_geoid, 0.02579, rtol=0.10) + assert np.isclose(response.cmb_geoid, 0.03206, rtol=0.10) + assert np.isclose(response.self_gravity.surface_topography, 0.49980, rtol=0.10) + assert np.isclose(response.self_gravity.cmb_topography, 0.93130, rtol=0.10) + assert np.isclose(response.self_gravity.surface_geoid, 0.04486, rtol=0.10) + assert np.isclose(response.self_gravity.cmb_geoid, 0.05461, rtol=0.10)