diff --git a/src/underworld3/coordinates.py b/src/underworld3/coordinates.py index 5040d3003..a2b968c61 100644 --- a/src/underworld3/coordinates.py +++ b/src/underworld3/coordinates.py @@ -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 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 @@ -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): @@ -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): @@ -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 @@ -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)") @@ -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 diff --git a/src/underworld3/discretisation/discretisation_mesh.py b/src/underworld3/discretisation/discretisation_mesh.py index c3a37b00f..e79cc3976 100644 --- a/src/underworld3/discretisation/discretisation_mesh.py +++ b/src/underworld3/discretisation/discretisation_mesh.py @@ -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 + + # 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 @@ -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 @@ -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 + g.attrs["ellipsoid"] = json.dumps(ellipsoid_ser) + # Add coordinate units metadata if hasattr(self, "coordinate_units"): coord_units_dict = { diff --git a/src/underworld3/meshing/geographic.py b/src/underworld3/meshing/geographic.py index 70e157ecf..d33f2800e 100644 --- a/src/underworld3/meshing/geographic.py +++ b/src/underworld3/meshing/geographic.py @@ -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 diff --git a/src/underworld3/meshing/surfaces.py b/src/underworld3/meshing/surfaces.py index 0029d29b9..455f2153c 100644 --- a/src/underworld3/meshing/surfaces.py +++ b/src/underworld3/meshing/surfaces.py @@ -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 = [] diff --git a/src/underworld3/systems/solvers.py b/src/underworld3/systems/solvers.py index 7406e6d22..3ac2c7603 100644 --- a/src/underworld3/systems/solvers.py +++ b/src/underworld3/systems/solvers.py @@ -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) return