diff --git a/docs/developer/CHANGELOG.md b/docs/developer/CHANGELOG.md index 0b6f1607..d417a24a 100644 --- a/docs/developer/CHANGELOG.md +++ b/docs/developer/CHANGELOG.md @@ -273,6 +273,11 @@ 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). +- `uw.postprocessing.geoid` provides generic spherical-shell geoid and + self-gravity coefficient functions. Its rotated-Stokes adapter projects the + existing boundary traction onto an axisymmetric harmonic; the pure functions + also accept coefficients recovered by other methods and an optional internal + load. ### 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..fb007f0b 100644 --- a/docs/developer/subsystems/boundary-stress-and-projection-postprocessing.md +++ b/docs/developer/subsystems/boundary-stress-and-projection-postprocessing.md @@ -74,6 +74,74 @@ 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. Spherical-shell geoid and self-gravity response + +Rotated free slip already exposes normal traction through +`Stokes.boundary_normal_traction()`. The convenience adapter projects that +recovered traction onto the unnormalised axisymmetric `P_l^0` harmonic and +applies the spherical-shell geoid operator: + +```python +stokes.solve() + +response = uw.postprocessing.geoid.spherical_shell_response_from_rotated_stokes( + stokes=stokes, + radius_inner=0.55, + radius_outer=1.0, + harmonic_degree=2, + internal_load_radius=0.775, + internal_load_coefficient=1.0, + include_self_gravity=True, + surface_density_contrast=3300.0, + cmb_density_contrast=5400.0, + planet_radius=6370000.0, + gravity=9.8, + gravitational_constant=6.67e-11, +) +``` + +The adapter delegates stress recovery to the existing rotated-free-slip API; +it does not implement a second CBF, constrained-multiplier, or topography +recovery path. `internal_load_coefficient` must use the same harmonic +normalisation and sign convention as the model's internal load. + +When surface and CMB topography coefficients are already available, call +`uw.postprocessing.geoid.spherical_shell_geoid_response()` or +`uw.postprocessing.geoid.spherical_shell_self_gravity_response()` directly. +These functions are pure post-processing, work for any spherical-harmonic +order with a consistent coefficient normalisation, and do not require a Stokes +object. They support non-negative harmonic degrees, and the internal load is +optional. The rotated-Stokes adapter requires degree one or greater because +normal-traction recovery removes the degree-zero mean. + +The density contrasts, dimensional outer-radius scale, and gravity are required +when self-gravity is enabled. They deliberately have no Earth- or +benchmark-specific defaults. The universal gravitational constant defaults to +the current CODATA value and can be overridden when reproducing a paper's +rounded constant. + +The 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. + +This module does not compute a benchmark's semi-analytical Stokes solution. +Published reference solvers, such as the Zhong et al. propagator-matrix method, +belong in `uw.analytic`; their computed topography coefficients can be passed to +the pure post-processing functions above. + +The rotated harmonic projector gathers boundary samples to rank zero and +reconstructs their spherical triangulation. A future boundary-reaction +functional could replace this step with a direct distributed finite-element +projection without changing the coefficient API. + ## 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..4e4325a1 --- /dev/null +++ b/src/underworld3/postprocessing/__init__.py @@ -0,0 +1,10 @@ +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 + +__all__ = ["geoid"] diff --git a/src/underworld3/postprocessing/geoid.py b/src/underworld3/postprocessing/geoid.py new file mode 100644 index 00000000..537a79fe --- /dev/null +++ b/src/underworld3/postprocessing/geoid.py @@ -0,0 +1,430 @@ +r"""Spherical-harmonic geoid and self-gravity response functions. + +The pure coefficient functions in this module are independent of a particular +Stokes discretisation. They combine surface, CMB, and optional internal-load +coefficients through the radial Green's function for one spherical-harmonic +degree. A separate convenience adapter obtains the two topography coefficients +from a completed rotated-free-slip Stokes solve. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from numbers import Integral + +from mpi4py import MPI +import numpy as np + + +__all__ = [ + "GeoidResponse", + "SelfGravityResponse", + "SphericalShellResponse", + "spherical_shell_geoid_response", + "spherical_shell_self_gravity_response", + "spherical_shell_response_from_rotated_stokes", +] + + +@dataclass(frozen=True) +class GeoidResponse: + """Surface and CMB geoid coefficients without self-gravity feedback.""" + + surface_geoid: float + cmb_geoid: float + + +@dataclass(frozen=True) +class SelfGravityResponse: + """Self-gravity-corrected topography and geoid coefficients.""" + + surface_topography: float + cmb_topography: float + surface_geoid: float + cmb_geoid: float + q_surface: float + q_cmb: float + + +@dataclass(frozen=True) +class SphericalShellResponse: + """Spherical-shell topography and corresponding geoid coefficients.""" + + surface_topography: float + cmb_topography: float + surface_geoid: float + cmb_geoid: float + self_gravity: SelfGravityResponse | None = None + + +def _validate_geometry( + radius_inner: float, + radius_outer: float, + harmonic_degree: int, +) -> tuple[float, float, int]: + try: + ri = float(radius_inner) + ro = float(radius_outer) + except (TypeError, ValueError) as error: + raise TypeError("The shell radii must be real numbers.") from error + if not np.all(np.isfinite((ri, ro))) or not 0.0 < ri < ro: + raise ValueError("Expected finite radii ordered as 0 < radius_inner < 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.") + return ri, ro, degree + + +def _spherical_shell_geoid_operator( + radius_inner: float, + radius_outer: float, + harmonic_degree: int, + internal_load_radius: float | None, + internal_load_coefficient: float, +) -> tuple[np.ndarray, np.ndarray]: + """Return the boundary-topography operator and fixed-load vector.""" + + ri, ro, degree = _validate_geometry( + radius_inner, + radius_outer, + harmonic_degree, + ) + internal_load_coefficient = float(internal_load_coefficient) + if not np.isfinite(internal_load_coefficient): + raise ValueError("internal_load_coefficient must be finite.") + + rint = None + if internal_load_radius is not None: + try: + rint = float(internal_load_radius) + except (TypeError, ValueError) as error: + raise TypeError("internal_load_radius must be a real number or None.") from error + if not np.isfinite(rint) or not ri < rint < ro: + raise ValueError("internal_load_radius must lie strictly between the shell radii.") + elif internal_load_coefficient != 0.0: + raise ValueError( + "internal_load_radius is required when internal_load_coefficient is nonzero." + ) + + 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.zeros(2, dtype=float) + if rint is not None and internal_load_coefficient != 0.0: + load = -internal_load_coefficient * 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, + harmonic_degree: int, + surface_topography_coefficient: float, + cmb_topography_coefficient: float, + internal_load_radius: float | None = None, + internal_load_coefficient: float = 0.0, +) -> GeoidResponse: + r"""Return spherical-shell geoid coefficients without self-gravity. + + All source coefficients must use one consistent spherical-harmonic + normalisation. The radial potential kernel depends on degree but not order, + so coefficients for any order :math:`m` may be used. A positive optional + internal-load coefficient follows the outward radial-load convention and + enters the potential with the opposite sign to boundary topography. + """ + + operator, load = _spherical_shell_geoid_operator( + radius_inner, + radius_outer, + harmonic_degree, + internal_load_radius, + internal_load_coefficient, + ) + topography = np.array( + [surface_topography_coefficient, cmb_topography_coefficient], + dtype=float, + ) + if not np.all(np.isfinite(topography)): + raise ValueError("The topography coefficients must be finite.") + geoid = operator @ topography + load + return GeoidResponse(float(geoid[0]), float(geoid[1])) + + +def spherical_shell_self_gravity_response( + *, + radius_inner: float, + radius_outer: float, + harmonic_degree: int, + surface_topography_coefficient: float, + cmb_topography_coefficient: float, + surface_density_contrast: float, + cmb_density_contrast: float, + planet_radius: float, + gravity: float, + internal_load_radius: float | None = None, + internal_load_coefficient: float = 0.0, + gravitational_constant: float = 6.67430e-11, +) -> SelfGravityResponse: + r"""Return self-gravity-corrected spherical-shell response coefficients. + + Density contrasts are signed and in kg/m³, ``planet_radius`` is in metres, + ``gravity`` is in m/s², and ``gravitational_constant`` is in SI units. + Shell radii and response coefficients remain nondimensional. + """ + + operator, load = _spherical_shell_geoid_operator( + radius_inner, + radius_outer, + harmonic_degree, + internal_load_radius, + internal_load_coefficient, + ) + topography = np.array( + [surface_topography_coefficient, cmb_topography_coefficient], + dtype=float, + ) + density_contrasts = np.array([surface_density_contrast, cmb_density_contrast], dtype=float) + physical_constants = np.array([planet_radius, gravity, gravitational_constant], dtype=float) + if not np.all(np.isfinite(topography)): + raise ValueError("The topography coefficients must be finite.") + if not np.all(np.isfinite(density_contrasts)): + raise ValueError("Density contrasts must be finite.") + if not np.all(np.isfinite(physical_constants)) or np.any(physical_constants <= 0.0): + raise ValueError("Physical constants must be finite and positive.") + + gravity_scale = 4.0 * np.pi * gravitational_constant * planet_radius / gravity + q = gravity_scale * density_contrasts + 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 SelfGravityResponse( + 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_triangle_area(a, b, c, radius: float) -> float: + determinant = abs(float(np.dot(a, np.cross(b, c)))) + denominator = float(1.0 + np.dot(a, b) + np.dot(b, c) + np.dot(c, a)) + return 2.0 * np.arctan2(determinant, denominator) * radius**2 + + +def _project_spherical_harmonic_samples( + coords: np.ndarray, + values: np.ndarray, + radius: float, + harmonic_degree: int, + response_sign: float, +) -> float: + 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 + harmonic = np.polynomial.legendre.Legendre.basis(harmonic_degree) + + for simplex in hull.simplices: + area = _spherical_triangle_area( + unit_coords[simplex[0]], + unit_coords[simplex[1]], + unit_coords[simplex[2]], + radius, + ) + centroid = unit_coords[simplex].mean(axis=0) + centroid /= np.linalg.norm(centroid) + projected_integral += ( + area * float(values[simplex].mean()) * harmonic(np.clip(centroid[2], -1, 1)) + ) + + harmonic_norm = 4.0 * np.pi / (2 * harmonic_degree + 1) + return float(response_sign * projected_integral / (radius**2 * harmonic_norm)) + + +def _rotated_topography_coefficient( + stokes, + boundary: str, + radius: float, + harmonic_degree: int, + buoyancy_scale: float, + response_sign: float, +) -> float: + buoyancy_scale = float(buoyancy_scale) + if not np.isfinite(buoyancy_scale) or buoyancy_scale == 0.0: + raise ValueError("Boundary buoyancy scales must be finite and nonzero.") + + coords, sigma_nn = stokes.boundary_normal_traction(boundary, mass="auto") + local_rows = np.column_stack( + ( + np.asarray(coords, dtype=float), + -np.asarray(sigma_nn, dtype=float) / buoyancy_scale, + ) + ) + gathered_rows = MPI.COMM_WORLD.gather(local_rows, root=0) + + coefficient = None + root_error = None + if MPI.COMM_WORLD.rank == 0: + try: + nonempty_rows = [rows for rows in gathered_rows if rows.size] + if not nonempty_rows: + raise RuntimeError(f"No samples found on boundary {boundary!r}.") + merged = {} + for row in np.vstack(nonempty_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()]) + coefficient = _project_spherical_harmonic_samples( + global_coords, + global_values, + radius, + harmonic_degree, + response_sign, + ) + except Exception as error: + root_error = f"{type(error).__name__}: {error}" + + root_error = MPI.COMM_WORLD.bcast(root_error, root=0) + if root_error is not None: + raise RuntimeError(f"Spherical-harmonic projection failed: {root_error}") + return float(MPI.COMM_WORLD.bcast(coefficient, root=0)) + + +def spherical_shell_response_from_rotated_stokes( + *, + stokes, + radius_inner: float, + radius_outer: float, + harmonic_degree: int, + internal_load_radius: float | None = None, + internal_load_coefficient: float = 0.0, + surface_boundary: str = "Upper", + cmb_boundary: str = "Lower", + surface_buoyancy_scale: float = 1.0, + cmb_buoyancy_scale: float = 1.0, + include_self_gravity: bool = False, + surface_density_contrast: float | None = None, + cmb_density_contrast: float | None = None, + planet_radius: float | None = None, + gravity: float | None = None, + gravitational_constant: float = 6.67430e-11, +) -> SphericalShellResponse: + r"""Compute spherical-shell response from a rotated-free-slip Stokes solve. + + Normal traction recovery is delegated to the existing + :meth:`Stokes.boundary_normal_traction` implementation. This adapter only + projects the two boundary responses onto the unnormalised + axisymmetric :math:`P_l^0` harmonic. Use the pure coefficient functions + directly for other harmonic orders or topography-recovery methods. Density + contrasts, planet radius, and gravity are required when + ``include_self_gravity`` is true. + """ + + ri, ro, degree = _validate_geometry( + radius_inner, + radius_outer, + harmonic_degree, + ) + if not isinstance(include_self_gravity, bool): + raise TypeError("include_self_gravity must be True or False.") + if degree == 0: + raise ValueError( + "The rotated-Stokes adapter requires harmonic_degree >= 1 because " + "boundary_normal_traction() removes the degree-zero mean." + ) + if include_self_gravity: + missing = [ + name + for name, value in ( + ("surface_density_contrast", surface_density_contrast), + ("cmb_density_contrast", cmb_density_contrast), + ("planet_radius", planet_radius), + ("gravity", gravity), + ) + if value is None + ] + if missing: + raise ValueError( + "Self-gravity requires explicit values for " + ", ".join(missing) + "." + ) + + surface_topography = _rotated_topography_coefficient( + stokes, + surface_boundary, + ro, + degree, + surface_buoyancy_scale, + 1.0, + ) + cmb_topography = _rotated_topography_coefficient( + stokes, + cmb_boundary, + ri, + degree, + cmb_buoyancy_scale, + -1.0, + ) + geoid = spherical_shell_geoid_response( + radius_inner=ri, + radius_outer=ro, + harmonic_degree=degree, + surface_topography_coefficient=surface_topography, + cmb_topography_coefficient=cmb_topography, + internal_load_radius=internal_load_radius, + internal_load_coefficient=internal_load_coefficient, + ) + self_gravity = None + if include_self_gravity: + self_gravity = spherical_shell_self_gravity_response( + radius_inner=ri, + radius_outer=ro, + harmonic_degree=degree, + surface_topography_coefficient=surface_topography, + cmb_topography_coefficient=cmb_topography, + internal_load_radius=internal_load_radius, + internal_load_coefficient=internal_load_coefficient, + surface_density_contrast=float(surface_density_contrast), + cmb_density_contrast=float(cmb_density_contrast), + planet_radius=float(planet_radius), + gravity=float(gravity), + gravitational_constant=gravitational_constant, + ) + + return SphericalShellResponse( + surface_topography=surface_topography, + cmb_topography=cmb_topography, + surface_geoid=geoid.surface_geoid, + cmb_geoid=geoid.cmb_geoid, + self_gravity=self_gravity, + ) diff --git a/tests/parallel/test_1071_spherical_shell_geoid_parallel.py b/tests/parallel/test_1071_spherical_shell_geoid_parallel.py new file mode 100644 index 00000000..189d6121 --- /dev/null +++ b/tests/parallel/test_1071_spherical_shell_geoid_parallel.py @@ -0,0 +1,79 @@ +import numpy as np +import pytest +import sympy +import underworld3 as uw + + +pytestmark = pytest.mark.level_2 + + +def test_rotated_spherical_shell_geoid_matches_serial_reference(): + if uw.mpi.size == 1: + pytest.skip("Run this regression with at least two MPI ranks.") + + radius_inner = 0.55 + radius_outer = 1.0 + rint = 0.775 + mesh = uw.meshing.SphericalShellInternalBoundary( + radiusOuter=radius_outer, + radiusInternal=rint, + radiusInner=radius_inner, + cellSize=0.25, + qdegree=2, + degree=1, + ) + velocity = uw.discretisation.MeshVariable("U_geoid_mpi", mesh, mesh.dim, degree=2) + pressure = uw.discretisation.MeshVariable("P_geoid_mpi", 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 + stokes.solve() + + response = uw.postprocessing.geoid.spherical_shell_response_from_rotated_stokes( + stokes=stokes, + radius_inner=radius_inner, + radius_outer=radius_outer, + harmonic_degree=2, + internal_load_radius=rint, + internal_load_coefficient=1.0, + include_self_gravity=True, + surface_density_contrast=3300.0, + cmb_density_contrast=5400.0, + planet_radius=6370000.0, + gravity=9.8, + gravitational_constant=6.67e-11, + ) + values = np.array( + [ + response.surface_topography, + response.cmb_topography, + response.surface_geoid, + response.cmb_geoid, + response.self_gravity.surface_topography, + response.self_gravity.cmb_topography, + response.self_gravity.surface_geoid, + response.self_gravity.cmb_geoid, + ] + ) + + for rank_values in uw.mpi.comm.allgather(values): + assert np.array_equal(rank_values, values) + + zhong_table_2 = np.array( + [0.41920, 0.77060, 0.02579, 0.03206, 0.49980, 0.93130, 0.04486, 0.05461] + ) + assert np.allclose(values, zhong_table_2, rtol=0.10, atol=1.0e-12) diff --git a/tests/test_1070_postprocessing_geoid.py b/tests/test_1070_postprocessing_geoid.py new file mode 100644 index 00000000..dce646a8 --- /dev/null +++ b/tests/test_1070_postprocessing_geoid.py @@ -0,0 +1,240 @@ +import math + +import numpy as np +import pytest +import sympy +import underworld3 as uw + + +pytestmark = pytest.mark.level_2 + + +def test_spherical_shell_postprocessing_is_public(): + assert hasattr(uw, "postprocessing") + assert uw.postprocessing.__all__ == ["geoid"] + assert hasattr(uw.postprocessing.geoid, "spherical_shell_geoid_response") + assert hasattr(uw.postprocessing.geoid, "spherical_shell_self_gravity_response") + assert hasattr(uw.postprocessing.geoid, "spherical_shell_response_from_rotated_stokes") + + +def test_spherical_shell_geoid_response_matches_direct_formula_and_load(): + radius_inner = 0.55 + radius_outer = 1.0 + rint = 0.775 + degree = 2 + internal_load_coefficient = 1.75 + surface_topography = 0.3918264884592169 + cmb_topography = 0.2653834107104151 + + response = uw.postprocessing.geoid.spherical_shell_geoid_response( + radius_inner=radius_inner, + radius_outer=radius_outer, + harmonic_degree=degree, + surface_topography_coefficient=surface_topography, + cmb_topography_coefficient=cmb_topography, + internal_load_radius=rint, + internal_load_coefficient=internal_load_coefficient, + ) + + denominator = 2 * degree + 1 + surface_expected = ( + radius_inner * (radius_inner / radius_outer) ** (degree + 1) * cmb_topography + + radius_outer * surface_topography + - internal_load_coefficient * rint * (rint / radius_outer) ** (degree + 1) + ) / denominator + cmb_expected = ( + radius_inner * cmb_topography + + radius_outer * (radius_inner / radius_outer) ** degree * surface_topography + - internal_load_coefficient * rint * (radius_inner / rint) ** degree + ) / denominator + + assert math.isclose(response.surface_geoid, surface_expected) + assert math.isclose(response.cmb_geoid, cmb_expected) + + +def test_spherical_shell_geoid_response_without_internal_load(): + response = uw.postprocessing.geoid.spherical_shell_geoid_response( + radius_inner=0.55, + radius_outer=1.0, + harmonic_degree=3, + surface_topography_coefficient=0.4, + cmb_topography_coefficient=-0.2, + ) + + denominator = 7.0 + expected = np.array( + [ + (0.4 + 0.55 * 0.55**4 * -0.2) / denominator, + (0.55**3 * 0.4 + 0.55 * -0.2) / denominator, + ] + ) + assert np.allclose([response.surface_geoid, response.cmb_geoid], expected) + + +def test_spherical_shell_self_gravity_response_matches_direct_solve(): + surface_topography = 0.3918264884592169 + cmb_topography = 0.2653834107104151 + response = uw.postprocessing.geoid.spherical_shell_self_gravity_response( + radius_inner=0.55, + radius_outer=1.0, + harmonic_degree=2, + surface_topography_coefficient=surface_topography, + cmb_topography_coefficient=cmb_topography, + internal_load_radius=0.775, + internal_load_coefficient=1.0, + surface_density_contrast=3300.0, + cmb_density_contrast=5400.0, + planet_radius=6370000.0, + gravity=9.8, + gravitational_constant=6.67e-11, + ) + + 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 + operator = np.array( + [ + [1.0 / denominator, 0.55 * 0.55**3 / denominator], + [0.55**2 / denominator, 0.55 / denominator], + ] + ) + load = np.array( + [ + -0.775 * 0.775**3 / denominator, + -0.775 * (0.55 / 0.775) ** 2 / denominator, + ] + ) + q_matrix = np.diag([q_surface, q_cmb]) + expected_topography = np.linalg.solve( + np.eye(2) - q_matrix @ operator, + np.array([surface_topography, cmb_topography]) + q_matrix @ load, + ) + expected_geoid = operator @ expected_topography + load + + assert math.isclose(response.q_surface, q_surface) + assert math.isclose(response.q_cmb, q_cmb) + assert np.allclose( + [response.surface_topography, response.cmb_topography], + expected_topography, + ) + assert np.allclose( + [response.surface_geoid, response.cmb_geoid], + expected_geoid, + ) + + +def test_spherical_shell_geoid_supports_degree_zero(): + response = uw.postprocessing.geoid.spherical_shell_geoid_response( + radius_inner=0.55, + radius_outer=1.0, + harmonic_degree=0, + surface_topography_coefficient=0.4, + cmb_topography_coefficient=-0.2, + ) + + assert math.isclose(response.surface_geoid, 0.4 - 0.55**2 * 0.2) + assert math.isclose(response.cmb_geoid, 0.4 - 0.55 * 0.2) + + +@pytest.mark.parametrize("harmonic_degree", [-1, 2.0, True]) +def test_spherical_shell_geoid_rejects_invalid_harmonic_degree(harmonic_degree): + error = TypeError if isinstance(harmonic_degree, (float, bool)) else ValueError + with pytest.raises(error): + uw.postprocessing.geoid.spherical_shell_geoid_response( + radius_inner=0.55, + radius_outer=1.0, + harmonic_degree=harmonic_degree, + surface_topography_coefficient=0.4, + cmb_topography_coefficient=0.7, + ) + + +def test_rotated_adapter_rejects_degree_zero(): + with pytest.raises(ValueError, match="requires harmonic_degree >= 1"): + uw.postprocessing.geoid.spherical_shell_response_from_rotated_stokes( + stokes=object(), + radius_inner=0.55, + radius_outer=1.0, + harmonic_degree=0, + ) + + +def test_internal_load_requires_a_radius(): + with pytest.raises(ValueError, match="internal_load_radius is required"): + uw.postprocessing.geoid.spherical_shell_geoid_response( + radius_inner=0.55, + radius_outer=1.0, + harmonic_degree=2, + surface_topography_coefficient=0.4, + cmb_topography_coefficient=0.7, + internal_load_coefficient=1.0, + ) + + +def test_rotated_adapter_requires_explicit_self_gravity_parameters(): + with pytest.raises(ValueError, match="Self-gravity requires explicit values"): + uw.postprocessing.geoid.spherical_shell_response_from_rotated_stokes( + stokes=object(), + radius_inner=0.55, + radius_outer=1.0, + harmonic_degree=2, + include_self_gravity=True, + ) + + +def test_rotated_stokes_adapter_matches_zhong_table_2(): + radius_inner = 0.55 + radius_outer = 1.0 + rint = 0.775 + mesh = uw.meshing.SphericalShellInternalBoundary( + radiusOuter=radius_outer, + radiusInternal=rint, + radiusInner=radius_inner, + cellSize=0.25, + qdegree=2, + degree=1, + ) + velocity = uw.discretisation.MeshVariable("U_geoid", mesh, mesh.dim, degree=2) + pressure = uw.discretisation.MeshVariable("P_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 + stokes.solve() + + response = uw.postprocessing.geoid.spherical_shell_response_from_rotated_stokes( + stokes=stokes, + radius_inner=radius_inner, + radius_outer=radius_outer, + harmonic_degree=2, + internal_load_radius=rint, + internal_load_coefficient=1.0, + include_self_gravity=True, + surface_density_contrast=3300.0, + cmb_density_contrast=5400.0, + planet_radius=6370000.0, + gravity=9.8, + gravitational_constant=6.67e-11, + ) + + 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)