"""The Earth's crust: a cubed sphere times a graded radial segment.
A spherical shell of thickness ``L`` under the surface of radius ``R``, built
the way ``tokamak_mesh.py`` builds a torus: a mesh that IS a product of meshes,
curved by a map on the product space.
* **The surface** is :func:`cubed_sphere`: six mapped squares, no pole, cells
within a factor of about two of each other (a latitude-longitude grid would
collapse its cells at the poles).
* **The depth** is a 1-D segment ``s in [0, 1]``, ``s = 0`` at the surface,
mapped to a depth ``d(s)`` in ``[0, L]``.
* **The product map** sends ``(x, d)`` -- a point of the unit sphere and a depth
-- to ``(R - d) x / |x|``: from the product's four coordinates to three.
**The radial grading.** The crust is to be resolved finely near the surface,
so the depth map is a geometric progression (exponential stretching),
d(s) = L (exp(beta s) - 1) / (exp(beta) - 1),
s(d) = log(1 + d (exp(beta) - 1) / L) / beta,
the usual choice for a boundary layer:
* two neighbouring layers differ in thickness by the same ratio
``exp(beta / n_r)`` everywhere -- the quality criterion one wants to keep
below about 1.3;
* ``beta`` is solved for from the thickness wanted for the TOP layer, which is
the number one actually has in mind;
* its derivative never vanishes, where a power law ``L (1 - (1 - s)^p)`` would
give the surface layer a zero thickness and the map a singular Jacobian;
* the inverse is in closed form, which point location needs.
**The controls**, as for the tokamak: the volume against
``4/3 pi (R^3 - (R - L)^3)``; the surface's area against ``4 pi R^2``; and the
divergence theorem, ``flux of x = 3 x volume``, which checks that the faces are
listed once each with outward normals. The boundary is named: ``"surface"``
outside, ``"moho"`` at the base of the crust.
**The picture.** The crust is 0.5 % of the radius: drawn to scale, it is a line.
The depths are therefore EXAGGERATED on the display only (``--exaggeration``,
stated in the title); the mesh itself is never touched. An octant is cut away
so the graded layers show on the cut, coloured by their true depth; the surface
is coloured by cube patch. Interactive (plotly): rotate, zoom, hover.
Usage::
python earth_crust_mesh.py [--n-face 8] [--n-radial 10] [--crust 35]
[--top-layer 1] [--exaggeration 25] [--html out.html] [--vtu out.vtu]
[--no-show]
``--vtu`` writes the mesh for ParaView (the cells' eight corners, with the
layer and the patch as cell data).
"""
import argparse
import jax
import jax.numpy as jnp
import numpy as np
import plotly.graph_objects as go
import tokamak_mesh as tokamak
from plotly.subplots import make_subplots
from scimba_jax.linear_approximation.meshes.manifold_mesh import cubed_sphere
from scimba_jax.linear_approximation.meshes.mesh import Mesh
from scimba_jax.linear_approximation.meshes.tensor_mesh import tensor_mesh
from scimba_jax.linear_approximation.quad.gauss_quad import UnitSquareTensorized
from scimba_jax.mapping.mapping import InvertibleFunction, Mapping
EARTH_RADIUS = 6371.0 # km
CRUST = 35.0 # km, a mean continental crust
N_FACE = 8 # cells per edge of a cube face: 6 N_FACE^2 surface cells
N_RADIAL = 10 # layers: with a 1 km top layer, a ratio of 1.26
TOP_LAYER = 1.0 # km, thickness wanted for the surface layer
EXAGGERATION = 25.0 # depths on the DISPLAY only
ORDER = 2
QUAD_ORDER = 4
# Camera presets, in plotly's scene units (the data box spans [-0.5, 0.5]):
# the whole globe, the rim of the cut where the layers show, the pole.
VIEWS = {
"globe": dict(eye=dict(x=1.35, y=1.35, z=0.95), center=dict(x=0, y=0, z=0)),
"zoom on the cut": dict(
eye=dict(x=0.28, y=0.45, z=0.45), center=dict(x=0.0, y=0.35, z=0.35)
),
"from the pole": dict(eye=dict(x=0.0, y=0.001, z=1.9), center=dict(x=0, y=0, z=0)),
}
PATCH_COLOURS = ["#8fb9a8", "#c9b37e", "#9fb0cf", "#d29a8a", "#a8c686", "#b9a3c9"]
# ── The radial grading ──────────────────────────────────────────────────────
def grading_rate(thickness, n_layers, top_layer):
"""``beta`` such that the top layer of the progression is ``top_layer`` thick.
The top layer, ``L (exp(beta / n) - 1) / (exp(beta) - 1)``, decreases from
``L / n`` (uniform, ``beta -> 0``) to zero as ``beta`` grows, so bisection
finds it. Asking for a top layer at least ``L / n`` thick means uniform.
Args:
thickness: The crust thickness ``L``.
n_layers: Number of layers.
top_layer: Thickness wanted for the surface layer.
Returns:
``beta``, ``0`` for uniform layers.
"""
if top_layer >= thickness / n_layers:
return 0.0
def top(beta):
return thickness * np.expm1(beta / n_layers) / np.expm1(beta)
low, high = 1e-12, 1.0
while top(high) > top_layer:
high *= 2.0
for _ in range(200):
middle = 0.5 * (low + high)
low, high = (middle, high) if top(middle) > top_layer else (low, middle)
return 0.5 * (low + high)
def depth_map(thickness, beta):
"""``s -> d(s)``, the geometric progression, and its closed-form inverse."""
if beta == 0.0:
return InvertibleFunction(lambda s: thickness * s, lambda d: d / thickness)
scale = float(np.expm1(beta))
def forward(s):
return thickness * jnp.expm1(beta * s) / scale
def inverse(d):
return jnp.log1p(d * scale / thickness) / beta
return InvertibleFunction(forward, inverse)
def interfaces(thickness, n_layers, beta):
"""Depths of the ``n_layers + 1`` layer interfaces, surface first."""
s = np.linspace(0.0, 1.0, n_layers + 1)
if beta == 0.0:
return thickness * s
return thickness * np.expm1(beta * s) / np.expm1(beta)
def radial_segment(n_layers, thickness, beta):
"""The depth, as a 1-D structured mesh: the map carries the grading."""
return Mesh(
dim=1,
n_cells=(n_layers,),
ref_quad=UnitSquareTensorized(dim=1, order=QUAD_ORDER),
mapping=Mapping(mappings=[depth_map(thickness, beta)]),
)
# ── The shell ───────────────────────────────────────────────────────────────
def shell_map(radius):
"""``(x, d)`` -- a point of the unit sphere and a depth -> a point of the shell.
Four coordinates in, three out. The inverse is in closed form: direction
and depth.
"""
def forward(z):
x, depth = z[..., :3], z[..., 3:]
return (radius - depth) * x / jnp.linalg.norm(x, axis=-1, keepdims=True)
def inverse(p):
r = jnp.linalg.norm(p, axis=-1, keepdims=True)
return jnp.concatenate([p / r, radius - r], axis=-1)
return InvertibleFunction(forward, inverse)
def crust_mesh(
n_face=N_FACE,
n_layers=N_RADIAL,
radius=EARTH_RADIUS,
thickness=CRUST,
top_layer=TOP_LAYER,
):
"""The crust: a cubed sphere times the graded depth, curved by `shell_map`.
Returns:
``(mesh, beta)``: the 3-D :class:`UnstructuredMesh`, its boundary named
``"surface"`` and ``"moho"``, and the grading rate.
"""
beta = grading_rate(thickness, n_layers, top_layer)
sphere = cubed_sphere(
n_face,
order=ORDER,
radius=1.0,
ref_quad=UnitSquareTensorized(dim=2, order=QUAD_ORDER),
)
mesh = tensor_mesh(
sphere,
radial_segment(n_layers, thickness, beta),
mapping=shell_map(radius),
)
# The controls read the boundary faces' quadrature: declared, as a scheme
# declares what it reads -- a bare mesh stores nothing.
mesh.request_storage("faces")
mesh.label_boundary("surface", outer_of_each_column)
mesh.label_boundary("moho", lambda x: ~outer_of_each_column(x))
return mesh, beta
def outer_of_each_column(midpoints):
"""True for the upper of the two boundary faces of each radial column.
⚠ Not a threshold on the midpoint's radius. A midpoint is the MEAN of the
face's quadrature points, and a face spanning ~11 degrees of a 6371 km
sphere has a sagitta of 20-30 km: its mean lies that far below the surface
-- as deep as the crust itself, so "radius above ``R - L/2``" selected no
face at all. The shell has no lateral boundary, so its boundary faces come
in pairs, top and bottom of each column, and the bottom one is the top one
scaled by ``(R - L) / R``: same direction exactly. Pairing by direction
and keeping the farther one is free of the sagitta whatever the sizes.
Args:
midpoints: ``(n, 3)`` face midpoints, as ``label_boundary`` hands them.
Returns:
``(n,)`` bool.
Raises:
ValueError: If the faces do not pair up by direction.
"""
radii = np.linalg.norm(midpoints, axis=-1)
directions = np.round(midpoints / radii[:, None], 9)
order = np.lexsort(directions.T[::-1])
if order.size % 2:
raise ValueError("an odd number of boundary faces cannot pair up.")
pairs = order.reshape(-1, 2)
if np.abs(directions[pairs[:, 0]] - directions[pairs[:, 1]]).max() > 1e-8:
raise ValueError("the boundary faces do not pair up by direction.")
upper = np.where(radii[pairs[:, 0]] > radii[pairs[:, 1]], pairs[:, 0], pairs[:, 1])
outer = np.zeros(radii.shape, dtype=bool)
outer[upper] = True
return outer
# ── What the mesh believes about itself ─────────────────────────────────────
def surface_area(mesh):
"""Quadrature of ``1`` over the faces labelled ``"surface"``."""
weights = np.asarray(mesh._store_ext_face_wts)
outer = np.isin(mesh.external_faces_idx, mesh.boundary_groups["surface"])
return float(weights[outer].sum())
def report(mesh, beta, n_face, n_layers, radius, thickness):
"""Size, the three controls, and the layers."""
volume = tokamak.volume(mesh)
exact_volume = 4.0 / 3.0 * np.pi * (radius**3 - (radius - thickness) ** 3)
area = surface_area(mesh)
exact_area = 4.0 * np.pi * radius**2
print(
f"crust: R = {radius:g} km, L = {thickness:g} km -- "
f"{mesh.n_cells_total} Q{ORDER} cells ({6 * n_face**2} surface cells x "
f"{n_layers} layers), {len(mesh.external_faces_idx)} boundary faces"
)
print(
f" volume {volume:.6e} km^3 exact {exact_volume:.6e} "
f"relative gap {abs(volume - exact_volume) / exact_volume:.2e}"
)
print(
f" surface {area:.6e} km^2 exact {exact_area:.6e} "
f"relative gap {abs(area - exact_area) / exact_area:.2e}"
)
print(
f" |flux of x - 3 volume| / (3 volume) = "
f"{tokamak.divergence_defect(mesh) / (3.0 * volume):.2e}"
)
depths = interfaces(thickness, n_layers, beta)
layers = np.diff(depths)
lateral = 0.5 * np.pi * radius / n_face
print(
f" grading beta = {beta:.4f}: ratio between neighbouring layers "
f"{np.exp(beta / n_layers):.3f}; lateral cell size ~{lateral:.0f} km, "
f"so aspect ratio {lateral / layers[0]:.0f} at the top, "
f"{lateral / layers[-1]:.0f} at the base"
)
if np.exp(beta / n_layers) > 1.3:
print(
" ⚠ neighbouring layers differ by more than 1.3: add layers or "
"thicken the top one."
)
for k, (top, width) in enumerate(zip(depths[:-1], layers)):
print(
f" layer {k}: depth {top:7.3f} -> {top + width:7.3f} km ({width:.3f} km)"
)
# ── Drawing ─────────────────────────────────────────────────────────────────
def exaggerated(points, radius, factor):
"""Depths multiplied by ``factor``, directions kept -- the display only."""
r = np.linalg.norm(points, axis=-1, keepdims=True)
return points / r * (radius - factor * (radius - r))
def face_samples(mesh, faces, k=5):
"""``k x k`` points on each listed face, through the cell map, ``(F, k^2, 3)``.
From the map, not from the corners, so curved faces are drawn curved.
Vectorised over the faces for the reason given in `tokamak.face_wireframe`.
"""
t = np.linspace(0.0, 1.0, k)
u, v = np.meshgrid(t, t, indexing="ij")
tangential = jnp.asarray(np.stack([u.ravel(), v.ravel()], axis=-1))
def per_face(face_id, cell):
return mesh._unit_hypercube_to_cell(
cell, mesh._face_reference_points(face_id, tangential)
)
return np.asarray(
jax.vmap(per_face)(
jnp.asarray(mesh.faces_left_edge)[faces],
jnp.asarray(mesh.faces_left)[faces],
)
)
def grid_triangles(k, n_faces):
"""The two triangles of every square of a ``k x k`` grid, for each face."""
i, j = np.meshgrid(np.arange(k - 1), np.arange(k - 1), indexing="ij")
a = (i * k + j).ravel()
b, c, d = a + k, a + k + 1, a + 1
one = np.concatenate([np.stack([a, b, c], -1), np.stack([a, c, d], -1)])
offsets = (np.arange(n_faces) * k * k)[:, None, None]
return (one[None] + offsets).reshape(-1, 3)
def surface_trace(points, triangles, name, **style):
"""One `Mesh3d` of sampled faces."""
return go.Mesh3d(
x=points[:, 0],
y=points[:, 1],
z=points[:, 2],
i=triangles[:, 0],
j=triangles[:, 1],
k=triangles[:, 2],
name=name,
flatshading=False,
lighting=dict(ambient=0.55, diffuse=0.7, specular=0.15, roughness=0.8),
**style,
)
def figure(mesh, beta, n_face, n_layers, radius, thickness, factor):
"""The cut-away crust (3-D, interactive) and its layers against depth."""
def keep(centre):
# The octant facing the camera is taken out.
return not (centre[0] > 0.0 and centre[1] > 0.0 and centre[2] > 0.0)
faces = tokamak.surface_of(mesh, keep)
samples = face_samples(mesh, faces)
boundary = np.asarray(mesh.faces_right)[faces] < 0
outer = np.isin(faces, mesh.boundary_groups["surface"])
base = boundary & ~outer
cut = ~boundary
k = int(round(samples.shape[1] ** 0.5))
shown = exaggerated(samples, radius, factor)
traces = []
# The surface, one colour per cube patch: which patch a product cell comes
# from is its surface cell, ``cell // n_layers``, and the cubed sphere lists
# its cells patch by patch, ``n_face^2`` each.
cells = np.asarray(mesh.faces_left)[faces]
patch = (cells // n_layers) // n_face**2
colours = np.array(PATCH_COLOURS)[patch[outer]]
traces.append(
surface_trace(
shown[outer].reshape(-1, 3),
grid_triangles(k, int(outer.sum())),
"surface (cube patches)",
facecolor=np.repeat(colours, 2 * (k - 1) ** 2),
)
)
traces.append(
surface_trace(
shown[base].reshape(-1, 3),
grid_triangles(k, int(base.sum())),
"moho (base of the crust)",
color="#6b6f76",
)
)
# The cut, coloured by TRUE depth, so the grading reads off the colours.
depth = radius - np.linalg.norm(samples[cut], axis=-1)
traces.append(
surface_trace(
shown[cut].reshape(-1, 3),
grid_triangles(k, int(cut.sum())),
"cut (true depth)",
intensity=depth.reshape(-1),
colorscale="Cividis",
reversescale=True,
colorbar=dict(title="depth (km)", x=0.66, len=0.7),
hovertemplate="depth %{intensity:.2f} km",
)
)
# The cells' edges, drawn from the map so they curve as the cells do.
curves = exaggerated(tokamak.face_wireframe(mesh, faces), radius, factor)
segments = curves.reshape(-1, curves.shape[-2], 3)
gaps = np.full((segments.shape[0], 1, 3), np.nan)
lines = np.concatenate([segments, gaps], axis=1).reshape(-1, 3)
traces.append(
go.Scatter3d(
x=lines[:, 0],
y=lines[:, 1],
z=lines[:, 2],
mode="lines",
line=dict(color="rgba(30,30,30,0.55)", width=1.5),
name="cell edges",
hoverinfo="skip",
)
)
depths = interfaces(thickness, n_layers, beta)
layers = np.diff(depths)
fig = make_subplots(
rows=1,
cols=2,
column_widths=[0.72, 0.28],
specs=[[{"type": "scene"}, {"type": "xy"}]],
subplot_titles=(
f"{mesh.n_cells_total} cells -- depths x{factor:g} on the display",
"layer thickness",
),
)
for trace in traces:
fig.add_trace(trace, row=1, col=1)
fig.add_trace(
go.Bar(
x=layers,
y=0.5 * (depths[:-1] + depths[1:]),
width=0.92 * layers,
orientation="h",
marker_color="#7d8fa3",
name="layers",
hovertemplate="%{x:.3f} km thick at %{y:.2f} km",
),
row=1,
col=2,
)
fig.update_yaxes(autorange="reversed", title="depth (km)", row=1, col=2)
fig.update_xaxes(title="thickness (km)", row=1, col=2)
fig.update_scenes(
aspectmode="data",
xaxis_visible=False,
yaxis_visible=False,
zaxis_visible=False,
camera=dict(eye=dict(x=1.35, y=1.35, z=0.95)),
)
fig.update_layout(
title=(
f"Earth's crust: cubed sphere ({n_face}x{n_face} per face) x "
f"{n_layers} graded layers -- R = {radius:g} km, L = {thickness:g} km, "
f"top layer {layers[0]:.2f} km, ratio {np.exp(beta / n_layers):.2f}"
),
legend=dict(x=0.0, y=0.02),
updatemenus=[
dict(
type="buttons",
direction="right",
x=0.0,
y=1.07,
xanchor="left",
buttons=[
dict(label=label, method="relayout", args=[{"scene.camera": view}])
for label, view in VIEWS.items()
],
)
],
margin=dict(l=10, r=10, t=70, b=10),
height=780,
)
return fig
def write_vtu(mesh, path, n_layers, n_face):
"""The mesh for ParaView: each cell's eight corners, layer and patch as data.
The cells are in tensor node order (first axis slowest), so the corner of
local index ``(i, j, k)`` in ``{0, p}^3`` is node ``(i (p+1) + j) (p+1) + k``;
listed here in VTK's hexahedron order.
"""
import meshio
side = ORDER + 1
corners = [(0, 0, 0), (1, 0, 0), (1, 1, 0), (0, 1, 0)]
corners += [(i, j, 1) for i, j, _ in corners]
local = [((i * ORDER) * side + j * ORDER) * side + k * ORDER for i, j, k in corners]
cells = np.asarray(mesh.cells)[:, local]
index = np.arange(mesh.n_cells_total)
meshio.write(
path,
meshio.Mesh(
np.asarray(mesh.nodes),
[("hexahedron", cells)],
cell_data={
"layer": [index % n_layers],
"patch": [(index // n_layers) // n_face**2],
},
),
)
def main():
parser = argparse.ArgumentParser(description=__doc__.split("\n")[0])
parser.add_argument("--n-face", type=int, default=N_FACE)
parser.add_argument("--n-radial", type=int, default=N_RADIAL)
parser.add_argument("--radius", type=float, default=EARTH_RADIUS)
parser.add_argument("--crust", type=float, default=CRUST)
parser.add_argument("--top-layer", type=float, default=TOP_LAYER)
parser.add_argument("--exaggeration", type=float, default=EXAGGERATION)
parser.add_argument("--html", default=None, help="write the figure there")
parser.add_argument("--vtu", default=None, help="write the mesh there")
parser.add_argument("--no-show", action="store_true")
args = parser.parse_args()
mesh, beta = crust_mesh(
args.n_face, args.n_radial, args.radius, args.crust, args.top_layer
)
report(mesh, beta, args.n_face, args.n_radial, args.radius, args.crust)
if args.vtu:
write_vtu(mesh, args.vtu, args.n_radial, args.n_face)
print(f"mesh written to {args.vtu}")
fig = figure(
mesh,
beta,
args.n_face,
args.n_radial,
args.radius,
args.crust,
args.exaggeration,
)
if args.html:
fig.write_html(args.html, include_plotlyjs="cdn")
print(f"figure written to {args.html}")
if not args.no_show:
fig.show()
if __name__ == "__main__":
main()