Skip to content

Fix geographic mesh units: ellipsoid quantities, checkpoint, Darcy sign #131

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
lmoresi merged 1 commit into from
Apr 22, 2026
Merged

Fix geographic mesh units: ellipsoid quantities, checkpoint, Darcy sign #131

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
66 changes: 39 additions & 27 deletions src/underworld3/coordinates.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -533,18 +533,21 @@ def _compute_coordinates(self):
x, y, z = dm_coords[:, 0], dm_coords[:, 1], dm_coords[:, 2]

# Get ellipsoid parameters
# Use nondimensional if available (when mesh created with units)
# Otherwise use raw km values
# ellipsoid["a"] is uw.quantity when units active, float (km) when not.
# DM coordinates are nondimensional, so nondimensionalise to match.
ellipsoid = self.cs.ellipsoid
if "a_nd" in ellipsoid and "b_nd" in ellipsoid:
# Units were active - use nondimensional ellipsoid
a = ellipsoid["a_nd"]
b = ellipsoid["b_nd"]
a_raw = ellipsoid["a"]
b_raw = ellipsoid["b"]
if hasattr(a_raw, "to"):
# Units active — nondimensionalise
import underworld3 as uw
a = float(uw.non_dimensionalise(a_raw))
b = float(uw.non_dimensionalise(b_raw))
self._nondimensional = True

Copilot AI Apr 22, 2026

Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

GeographicCoordinateAccessor.depth still dimensionalises nondimensional depth using self.cs.ellipsoid.get('L_ref_km', 1000), but RegionalGeographicBox no longer stores L_ref_km (or a_nd/b_nd). This will silently default to 1000 km and return incorrect dimensional depths when units are active. Consider deriving the length scale from the mesh/model (e.g., mesh.length_scale converted to km, or model.get_scale_for_dimensionality('[length]')) instead of relying on an ellipsoid dict field, or reintroduce a reliably-populated reference length entry during mesh creation/checkpoint restore.

Suggested change
self._nondimensional = True
self._nondimensional = True
# Preserve the reference length used for nondimensionalisation so
# downstream dimensional depth conversion does not silently fall
# back to an incorrect default.
if "L_ref_km" not in ellipsoid:
ref_scales_km = []
if a != 0.0:
ref_scales_km.append(float(a_raw.to("kilometer").magnitude) / a)
if b != 0.0:
ref_scales_km.append(float(b_raw.to("kilometer").magnitude) / b)
if ref_scales_km:
ellipsoid["L_ref_km"] = sum(ref_scales_km) / len(ref_scales_km)

Copilot uses AI. Check for mistakes.
else:
# No units - use km values
a = ellipsoid["a"]
b = ellipsoid["b"]
# No units — km floats
a = float(a_raw)
b = float(b_raw)
self._nondimensional = False

# Convert to geographic using our utility function
Expand Down Expand Up @@ -728,14 +731,16 @@ def to_cartesian(self, lon, lat, depth):
>>> tomo_depth = np.array([10.0, 20.0, 30.0]) # km or nondimensional
>>> x, y, z = mesh.geo.to_cartesian(tomo_lon, tomo_lat, tomo_depth)
"""
# Use nondimensional ellipsoid if available
# Nondimensionalise ellipsoid for numeric coordinate conversion
ellipsoid = self.cs.ellipsoid
if "a_nd" in ellipsoid and "b_nd" in ellipsoid:
a = ellipsoid["a_nd"]
b = ellipsoid["b_nd"]
a_raw, b_raw = ellipsoid["a"], ellipsoid["b"]
if hasattr(a_raw, "to"):
import underworld3 as uw
a = float(uw.non_dimensionalise(a_raw))
b = float(uw.non_dimensionalise(b_raw))
else:
a = ellipsoid["a"]
b = ellipsoid["b"]
a = float(a_raw)
b = float(b_raw)
return geographic_to_cartesian(lon, lat, depth, a, b)

def from_cartesian(self, x, y, z):
Expand Down Expand Up @@ -768,14 +773,16 @@ def from_cartesian(self, x, y, z):
>>> x, y, z = mesh.data[:, 0], mesh.data[:, 1], mesh.data[:, 2]
>>> lon, lat, depth = mesh.geo.from_cartesian(x, y, z)
"""
# Use nondimensional ellipsoid if available
# Nondimensionalise ellipsoid for numeric coordinate conversion
ellipsoid = self.cs.ellipsoid
if "a_nd" in ellipsoid and "b_nd" in ellipsoid:
a = ellipsoid["a_nd"]
b = ellipsoid["b_nd"]
a_raw, b_raw = ellipsoid["a"], ellipsoid["b"]
if hasattr(a_raw, "to"):
import underworld3 as uw
a = float(uw.non_dimensionalise(a_raw))
b = float(uw.non_dimensionalise(b_raw))
else:
a = ellipsoid["a"]
b = ellipsoid["b"]
a = float(a_raw)
b = float(b_raw)
return cartesian_to_geographic(x, y, z, a, b)

def points_to_cartesian(self, points_geo):
Expand Down Expand Up @@ -1588,9 +1595,12 @@ def __init__(

self.type = "Geographic"

# Store ellipsoid parameters (will be set by mesh creation function)
# Default to WGS84 if not specified
if not hasattr(self, "ellipsoid"):
# Store ellipsoid parameters.
# Priority: checkpoint-restored ellipsoid > already set > WGS84 default
if hasattr(self.mesh, "_checkpoint_ellipsoid_pending") and self.mesh._checkpoint_ellipsoid_pending is not None:
self.ellipsoid = self.mesh._checkpoint_ellipsoid_pending
self.mesh._checkpoint_ellipsoid_pending = None
elif not hasattr(self, "ellipsoid"):
self.ellipsoid = ELLIPSOIDS["WGS84"].copy()

# _X already contains UWCoordinates, _x is a copy for the "lowercase" alias
Expand Down Expand Up @@ -1636,7 +1646,9 @@ def __init__(
lat_approx = sympy.atan2(z, rxy) * 180 / sympy.pi

# Depth (simplified - proper depth calculated in accessor)
depth_approx = self.ellipsoid["a"] - r
# self.ellipsoid["a"] is uw.quantity when units active, float when not
a_sym = expression(r"a_{ellipsoid}", self.ellipsoid["a"], "Semi-major axis")
depth_approx = a_sym - r

# Geographic coordinate expressions (approximations for symbolic work)
lambda_lon = expression(R"\lambda_{lon}", lon_deg, "Longitude (degrees East)")
Expand All @@ -1654,8 +1666,8 @@ def __init__(
# This matters at regional scales (10-100 km) where the difference
# between geodetic and geocentric latitude (~10 arcmin) is significant.

a = sympy.sympify(self.ellipsoid["a"])
b = sympy.sympify(self.ellipsoid["b"])
a = expression(r"a_{ellipsoid}", self.ellipsoid["a"], "Semi-major axis")
b = expression(r"b_{ellipsoid}", self.ellipsoid["b"], "Semi-minor axis")

# Geodetic normal components (gradient of ellipsoid equation)
# ∇F = (2x/a², 2y/a², 2z/b²) - we drop the factor of 2 since we normalize
Expand Down
38 changes: 38 additions & 0 deletions src/underworld3/discretisation/discretisation_mesh.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -355,6 +355,26 @@ def __init__(
except KeyError:
pass

regions = None
try:
json_str = f["metadata"].attrs["regions"]
rgn_dict = json.loads(json_str)
regions = Enum("Regions", rgn_dict)
except KeyError:
pass

Comment on lines +358 to +365

Copilot AI Apr 22, 2026

Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The regions metadata is loaded from the checkpoint and assigned to a local regions variable, but it is never used afterwards (unlike boundaries and coordinate_system_type). Either wire this through (e.g., store it on self and/or use it to rebuild labels) or remove the dead code to avoid confusion about whether region restoration is supported.

Suggested change
regions = None
try:
json_str = f["metadata"].attrs["regions"]
rgn_dict = json.loads(json_str)
regions = Enum("Regions", rgn_dict)
except KeyError:
pass

Copilot uses AI. Check for mistakes.
# Restore ellipsoid with quantities for geographic meshes
self._checkpoint_ellipsoid = None
try:
json_str = f["metadata"].attrs["ellipsoid"]
ellipsoid_raw = json.loads(json_str)
for k, v in ellipsoid_raw.items():
if isinstance(v, dict) and "value" in v and "unit" in v:
ellipsoid_raw[k] = uw.quantity(v["value"], v["unit"])
self._checkpoint_ellipsoid = ellipsoid_raw
except KeyError:
pass

f.close()

# This needs to be done when reading a dm from a checkpoint
Expand Down Expand Up @@ -610,6 +630,11 @@ def mesh_update_callback(array, change_context):

# Now add the appropriate coordinate system for the mesh's natural geometry
# This step will usually over-write the defaults we just defined
# For geographic meshes loaded from checkpoint, pre-set the ellipsoid
# so the CoordinateSystem __init__ picks it up.
if hasattr(self, "_checkpoint_ellipsoid") and self._checkpoint_ellipsoid is not None:
self._checkpoint_ellipsoid_pending = self._checkpoint_ellipsoid

self._CoordinateSystem = CoordinateSystem(self, coordinate_system_type)

# This was in the _jit extension but ... if
Expand Down Expand Up @@ -2083,6 +2108,19 @@ def write(self, filename: str, index: Optional[int] = None):
}
g.attrs["coordinate_system_type"] = json.dumps(coordinates_type_dict)

# Save ellipsoid metadata for geographic meshes
if hasattr(self.CoordinateSystem, "ellipsoid"):
ellipsoid_ser = {}
for k, v in self.CoordinateSystem.ellipsoid.items():
if hasattr(v, "to"): # uw.quantity
ellipsoid_ser[k] = {
"value": float(v.magnitude),
"unit": str(v.units),
}
else:
ellipsoid_ser[k] = v
Comment on lines +2111 to +2121

Copilot AI Apr 22, 2026

Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Ellipsoid serialization/deserialization is new checkpoint behavior for geographic meshes (including quantity handling). Please add a checkpoint round-trip test that: creates a GEOGRAPHIC mesh with units active, writes it, reloads it, and asserts the restored ellipsoid values/units and coordinate conversions (geo↔cartesian) are consistent. This will protect the HDF5 metadata schema and the quantity reconstruction logic from future breakage.

Suggested change
# Save ellipsoid metadata for geographic meshes
if hasattr(self.CoordinateSystem, "ellipsoid"):
ellipsoid_ser = {}
for k, v in self.CoordinateSystem.ellipsoid.items():
if hasattr(v, "to"): # uw.quantity
ellipsoid_ser[k] = {
"value": float(v.magnitude),
"unit": str(v.units),
}
else:
ellipsoid_ser[k] = v
def _serialise_checkpoint_metadata_value(value):
if (
hasattr(value, "magnitude")
and hasattr(value, "units")
and hasattr(value, "to")
):
magnitude = value.magnitude
if isinstance(magnitude, numpy.ndarray):
magnitude = magnitude.tolist()
elif isinstance(magnitude, numpy.generic):
magnitude = magnitude.item()
return {
"__uw_quantity__": True,
"value": magnitude,
"unit": str(value.units),
}
if isinstance(value, dict):
return {
key: _serialise_checkpoint_metadata_value(val)
for key, val in value.items()
}
if isinstance(value, (list, tuple)):
return [
_serialise_checkpoint_metadata_value(val) for val in value
]
if isinstance(value, numpy.ndarray):
return value.tolist()
if isinstance(value, numpy.generic):
return value.item()
return value
# Save ellipsoid metadata for geographic meshes
if hasattr(self.CoordinateSystem, "ellipsoid"):
ellipsoid_ser = {
k: _serialise_checkpoint_metadata_value(v)
for k, v in self.CoordinateSystem.ellipsoid.items()
}

Copilot uses AI. Check for mistakes.
g.attrs["ellipsoid"] = json.dumps(ellipsoid_ser)

# Add coordinate units metadata
if hasattr(self, "coordinate_units"):
coord_units_dict = {
Expand Down
51 changes: 14 additions & 37 deletions src/underworld3/meshing/geographic.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -555,50 +555,27 @@ def RegionalGeographicBox(
"Use str name, (a, b) tuple, True for WGS84, or False for sphere."
)

# Get ellipsoid parameters (in km)
# Get ellipsoid parameters (raw km floats from ELLIPSOIDS dict)
a_km = ellipsoid_dict["a"]
b_km = ellipsoid_dict["b"]

# Nondimensionalize ellipsoid if units active
if units_active:
# Get reference length in km
ref_length = model.get_fundamental_scales().get("length")
if ref_length is not None:
# Convert reference length to km
if hasattr(ref_length, "to"):
L_ref_km = float(ref_length.to("km").magnitude)
elif hasattr(ref_length, "magnitude"):
# Assume it's in base SI (meters)
L_ref_km = float(ref_length.magnitude) / 1000.0
else:
L_ref_km = float(ref_length) / 1000.0

# Nondimensional ellipsoid parameters
a_nd = a_km / L_ref_km
b_nd = b_km / L_ref_km

# Store both in ellipsoid dict
ellipsoid_dict["a_nd"] = a_nd
ellipsoid_dict["b_nd"] = b_nd
ellipsoid_dict["L_ref_km"] = L_ref_km

# Use nondimensional values for mesh generation
a = a_nd
b = b_nd
depth_min = depth_min_nd
depth_max = depth_max_nd
else:
# No reference length available - use km
a = a_km
b = b_km
depth_min = depth_min_nd
depth_max = depth_max_nd
# Store as quantities — the units system handles nondimensionalisation
# downstream (symbolic expressions via JIT, numeric via non_dimensionalise)
ellipsoid_dict["a"] = uw.quantity(a_km, "km")
ellipsoid_dict["b"] = uw.quantity(b_km, "km")

# For mesh generation (gmsh addPoint), use nondimensional floats
a = float(uw.non_dimensionalise(ellipsoid_dict["a"]))
b = float(uw.non_dimensionalise(ellipsoid_dict["b"]))
depth_min = float(depth_min_nd)
depth_max = float(depth_max_nd)
else:
# No units - use km directly
# No units — keep as km floats
a = a_km
b = b_km
depth_min = depth_min_nd
depth_max = depth_max_nd
depth_min = float(depth_min_nd)
depth_max = float(depth_max_nd)

# Unpack element counts (angles already processed above)
numLon, numLat, numDepth = numElements
Expand Down
7 changes: 5 additions & 2 deletions src/underworld3/meshing/surfaces.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1690,8 +1690,11 @@ def from_trace(

# --- Build 3D points for each depth layer ---
if is_geographic:
a_km = ellipsoid["a"]
b_km = ellipsoid["b"]
# Extract km floats — ellipsoid["a"] is uw.quantity when units active
a_raw = ellipsoid["a"]
b_raw = ellipsoid["b"]
a_km = float(a_raw.to("km").magnitude) if hasattr(a_raw, "to") else float(a_raw)
b_km = float(b_raw.to("km").magnitude) if hasattr(b_raw, "to") else float(b_raw)
from underworld3.coordinates import geographic_to_cartesian

all_points_km = []
Expand Down
4 changes: 2 additions & 2 deletions src/underworld3/systems/solvers.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -537,11 +537,11 @@ def solve(

super().solve(zero_init_guess, _force_setup)

# Now solve flow field
# Now solve flow field: v = -flux = -K(grad(h) - s)

# self._v_projector.petsc_options["snes_rtol"] = 1.0e-6
# self._v_projector.petsc_options.delValue("ksp_monitor")
self._v_projector.uw_function = self.darcy_flux
self._v_projector.uw_function = -self.darcy_flux
self._v_projector.solve(zero_init_guess)
Comment on lines +540 to 545

Copilot AI Apr 22, 2026

Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This sign change affects the meaning/direction of the projected Darcy velocity field (self.v) and could easily regress again because the existing Darcy solver tests primarily validate pressure, not velocity direction. Add/extend a test to assert that the projected velocity matches -darcy_flux (e.g., check sign/direction at a few sample points for a known pressure gradient).

Copilot uses AI. Check for mistakes.

return
Expand Down
Loading