Add support for defining planetary radius (#2739)

Co-authored-by: Erik van Sebille <e.vansebille@uu.nl>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Vecko <36369090+VeckoTheGecko@users.noreply.github.com>

Author: 4 people

src/parcels/__init__.py

@@ -22,6 +22,7 @@
 from parcels._core.basegrid import BaseGrid
 from parcels._core.uxgrid import UxGrid
 from parcels._core.xgrid import XGrid
+from parcels._core.mesh import SphericalMesh
 
 from parcels._core.statuscodes import (
     AllParcelsErrorCodes,
@@ -54,6 +55,7 @@
     "BaseGrid",
     "UxGrid",
     "XGrid",
+    "SphericalMesh",
     # Status codes and errors
     "AllParcelsErrorCodes",
     "FieldInterpolationError",

src/parcels/_core/basegrid.py

@@ -8,7 +8,7 @@
 
 import numpy as np
 
-import parcels._typing as ptyping
+from parcels._core.mesh import FlatMesh, SphericalMesh
 from parcels._core.spatialhash import SpatialHash
 
 if TYPE_CHECKING:
@@ -26,7 +26,7 @@ class BaseGrid(ABC):
     """Base class for parcels.XGrid and parcels.UxGrid defining common methods and properties"""
 
     _spatialhash: SpatialHash | None
-    _mesh: ptyping.Mesh
+    _mesh: FlatMesh | SphericalMesh
 
     @abstractmethod
     def search(self, z: float, y: float, x: float, ei=None) -> dict[str, tuple[int, float | np.ndarray]]:

src/parcels/_core/fieldset.py

@@ -189,7 +189,7 @@ def to_windowed_arrays(self, *, max_levels: int | None = None):
             model.to_windowed_arrays(max_levels=max_levels)
         return self
 
-    def add_constant_field(self, name: str, value, mesh: ptyping.Mesh = "spherical"):
+    def add_constant_field(self, name: str, value, mesh: ptyping.TMesh = "spherical"):
         """Wrapper function to add a Field that is constant in space,
            useful e.g. when using constant horizontal diffusivity
 
@@ -287,7 +287,7 @@ def from_ugrid_conventions(
     def from_sgrid_conventions(
         cls,
         ds: xr.Dataset,
-        mesh: ptyping.Mesh | None = None,
+        mesh: ptyping.TMesh | None = None,
         vector_fields: ptyping.VectorFields | NotSetType = NOTSET,
     ):  # TODO: Update mesh to be discovered from the dataset metadata
         """Create a FieldSet from a dataset using SGRID convention metadata.

src/parcels/_core/index_search.py

@@ -109,7 +109,7 @@ def curvilinear_point_in_cell(grid, y: np.ndarray, x: np.ndarray, yi: np.ndarray
         dtype=float,
     )
 
-    if grid._mesh == "spherical":
+    if grid._mesh.is_spherical():
         xsi, eta = _bilinear_inverse_tangent_plane(clon, clat, x, y)
         is_in_cell = np.where((xsi >= 0) & (xsi <= 1) & (eta >= 0) & (eta <= 1), 1, 0)
     else:
@@ -318,7 +318,7 @@ def uxgrid_point_in_cell(grid, y: np.ndarray, x: np.ndarray, yi: np.ndarray, xi:
     coords : np.ndarray
         Barycentric coordinates of the points within their respective cells.
     """
-    if grid._mesh == "spherical":
+    if grid._mesh.is_spherical():
         lon_rad = np.deg2rad(x)
         lat_rad = np.deg2rad(y)
         x_cart, y_cart, z_cart = _latlon_rad_to_xyz(lat_rad, lon_rad)

src/parcels/_core/kernel.py

@@ -141,9 +141,9 @@ def check_fieldsets_in_kernels(self, kernel):  # TODO v4: this can go into anoth
                         stacklevel=2,
                     )
                     self.fieldset.add_context("RK45_tol", 10)
-                if self.fieldset.U.grid._mesh == "spherical":
+                if self.fieldset.U.grid._mesh.is_spherical():
                     self.fieldset.RK45_tol /= (
-                        1852 * 60
+                        self.fieldset.U.grid.deg2m
                     )  # TODO does not account for zonal variation in meter -> degree conversion
                 if not hasattr(self.fieldset, "RK45_min_dt"):
                     warnings.warn(

src/parcels/_core/mesh.py

@@ -0,0 +1,73 @@
+from abc import ABC, abstractmethod
+from typing import Literal
+
+import numpy as np
+
+EARTH_RADIUS = 6366707.019493707
+
+
+class BaseMesh(ABC):
+    radius: float | None
+
+    @abstractmethod
+    def is_spherical(self) -> bool: ...
+
+
+class SphericalMesh(BaseMesh):
+    """Spherical mesh object with configurable planetary radius.
+
+    Pass to FieldSet object as ``mesh=SphericalMesh(radius=...)``.
+    radius is in meters; defaults to Earth radius.
+    """
+
+    def __init__(self, radius: float = EARTH_RADIUS):
+        if not isinstance(radius, (int, float, np.number)):
+            raise TypeError(f"radius must be a number, got {type(radius).__name__}")
+        if radius <= 0:
+            raise ValueError(f"radius must be positive, got {radius}")
+        self.radius = radius
+
+    @property
+    def deg2m(self) -> float:
+        """Meters per degree of arc."""
+        assert self.radius is not None
+        return self.radius * np.pi / 180.0
+
+    def is_spherical(self):
+        return True
+
+    def __repr__(self) -> str:
+        return f"SphericalMesh(radius={self.radius})"
+
+
+class FlatMesh(BaseMesh):
+    """Flat mesh object."""
+
+    def __init__(self):
+        self.radius = None
+        return
+
+    def __repr__(self) -> str:
+        return "FlatMesh()"
+
+    def is_spherical(self):
+        return False
+
+
+TMesh = SphericalMesh | Literal["spherical", "flat"]  # corresponds with `mesh`
+
+
+def get_mesh(mesh: TMesh):
+    if isinstance(mesh, SphericalMesh):
+        return mesh
+    if mesh == "flat":
+        return FlatMesh()
+    if mesh == "spherical":
+        return SphericalMesh(EARTH_RADIUS)
+    raise ValueError(f"mesh must be 'flat', 'spherical', or a SphericalMesh object. Got {mesh=!r}")
+
+
+def is_spherical(mesh: FlatMesh | SphericalMesh):
+    if isinstance(mesh, SphericalMesh):
+        return True
+    return False

src/parcels/_core/model.py

@@ -23,7 +23,6 @@
 )
 from parcels._logger import logger
 from parcels._python import NOTSET, NotSetType
-from parcels._typing import Mesh
 from parcels.convert import _ds_rename_using_standard_names
 from parcels.interpolators import (
     CGrid_Velocity,
@@ -133,7 +132,7 @@ def preprocess_sgrid_model_data(ds: xr.Dataset) -> xr.Dataset:
 
 
 class StructuredModelData(ModelData):
-    def __init__(self, data: xr.Dataset, mesh: Mesh, vector_field_components: ptyping.VectorFields):
+    def __init__(self, data: xr.Dataset, mesh: ptyping.TMesh, vector_field_components: ptyping.VectorFields):
         if not isinstance(data, xr.Dataset):
             raise ValueError(f"Expected `data` to be an xarray.Dataset . Got {type(data)}")
 
@@ -182,7 +181,7 @@ def construct_fields(self) -> list[Field | VectorField]:
 
     @classmethod
     def from_sgrid_conventions(
-        cls, ds: xr.Dataset, mesh: Mesh | None, vector_fields: ptyping.VectorFields | NotSetType
+        cls, ds: xr.Dataset, mesh: ptyping.TMesh | None, vector_fields: ptyping.VectorFields | NotSetType
     ) -> Self:
         ds = ds.copy()
         if mesh is None:
@@ -329,7 +328,9 @@ def scalar_field_names(self) -> list[str]:
         return list(self.data.data_vars)
 
     @classmethod
-    def from_ugrid_conventions(cls, ds: ux.UxDataset, mesh: Mesh, vector_fields: ptyping.VectorFields | NotSetType):
+    def from_ugrid_conventions(
+        cls, ds: ux.UxDataset, mesh: ptyping.TMesh, vector_fields: ptyping.VectorFields | NotSetType
+    ):
         ds_dims = list(ds.dims)
         if not all(dim in ds_dims for dim in ["time", "zf", "zc"]):
             raise ValueError(
@@ -353,7 +354,7 @@ def from_ugrid_conventions(cls, ds: ux.UxDataset, mesh: Mesh, vector_fields: pty
 
 
 # TODO: Refactor later into something like `parcels._metadata.discover(dataset)` helper that can be used to discover important metadata like this. I think this whole metadata handling should be refactored into its own module.
-def _get_mesh_type_from_sgrid_dataset(ds_sgrid: xr.Dataset) -> Mesh:
+def _get_mesh_type_from_sgrid_dataset(ds_sgrid: xr.Dataset) -> ptyping.TMesh:
     """Small helper to inspect SGRID metadata and dataset metadata to determine mesh type."""
     sgrid_metadata = ds_sgrid.sgrid.metadata
 

src/parcels/_core/particlefile.py

@@ -14,6 +14,7 @@
 import xarray as xr
 
 import parcels
+from parcels._core.mesh import BaseMesh
 from parcels._core.particle import ParticleClass
 from parcels._core.particlesetview import ParticleSetView
 from parcels._core.utils.time import timedelta_to_float
@@ -119,14 +120,14 @@ def __init__(
     def __repr__(self) -> str:
         return particlefile_repr(self)
 
-    def set_metadata(self, parcels_grid_mesh: Literal["spherical", "flat"]):
+    def set_metadata(self, parcels_grid_mesh: BaseMesh):
         self.metadata.update(
             {
                 "feature_type": "trajectory",
                 "Conventions": "CF-1.6/CF-1.7",
                 "ncei_template_version": "NCEI_NetCDF_Trajectory_Template_v2.0",
                 "parcels_version": parcels.__version__,
-                "parcels_grid_mesh": parcels_grid_mesh,
+                "parcels_grid_mesh": repr(parcels_grid_mesh),
             }
         )
 

src/parcels/_core/spatialhash.py

@@ -55,7 +55,7 @@ def __init__(
 
         if isinstance_noimport(grid, "XGrid"):
             self._coord_dim = 2  # Number of computational coordinates is 2 (bilinear interpolation)
-            if self._source_grid._mesh == "spherical":
+            if self._source_grid._mesh.is_spherical():
                 lon = np.deg2rad(self._source_grid.lon)
                 lat = np.deg2rad(self._source_grid.lat)
                 x, y, z = _latlon_rad_to_xyz(lat, lon)
@@ -160,7 +160,7 @@ def __init__(
 
         elif isinstance_noimport(grid, "UxGrid"):
             self._coord_dim = grid.uxgrid.n_max_face_nodes  # Number of barycentric coordinates
-            if self._source_grid._mesh == "spherical":
+            if self._source_grid._mesh.is_spherical():
                 # Reshape node coordinates to (nfaces, nnodes_per_face)
                 nids = self._source_grid.uxgrid.face_node_connectivity.values
                 lon = self._source_grid.uxgrid.node_lon.values[nids]
@@ -404,7 +404,7 @@ def query(self, y, x):
 
         y = np.asarray(y)
         x = np.asarray(x)
-        if self._source_grid._mesh == "spherical":
+        if self._source_grid._mesh.is_spherical():
             # Convert coords to Cartesian coordinates (x, y, z)
             lat = np.deg2rad(y)
             lon = np.deg2rad(x)

src/parcels/_core/utils/interpolation.py

@@ -3,7 +3,7 @@
 
 import numpy as np
 
-from parcels._typing import Mesh
+from parcels._core.mesh import BaseMesh
 
 __all__ = []
 
@@ -61,12 +61,12 @@ def dphidxsi3D_lin(zeta: float, eta: float, xsi: float) -> tuple[list[float], li
 
 
 def dxdxsi3D_lin(
-    hexa_z: list[float], hexa_y: list[float], hexa_x: list[float], zeta: float, eta: float, xsi: float, mesh: Mesh
+    hexa_z: list[float], hexa_y: list[float], hexa_x: list[float], zeta: float, eta: float, xsi: float, mesh: BaseMesh,
+    deg2m: float = 1852 * 60.0
 ) -> tuple[float, float, float, float, float, float, float, float, float]:
     dphidxsi, dphideta, dphidzet = dphidxsi3D_lin(zeta, eta, xsi)
 
-    if mesh == 'spherical':
-        deg2m = 1852 * 60.
+    if mesh.is_spherical():
         rad = np.pi / 180.
         lat = (1-xsi) * (1-eta) * hexa_y[0] + \
                  xsi  * (1-eta) * hexa_y[1] + \
@@ -92,9 +92,10 @@ def dxdxsi3D_lin(
 
 
 def jacobian3D_lin(
-    hexa_z: list[float], hexa_y: list[float], hexa_x: list[float], zeta: float, eta: float, xsi: float, mesh: Mesh
+    hexa_z: list[float], hexa_y: list[float], hexa_x: list[float], zeta: float, eta: float, xsi: float, mesh: BaseMesh,
+    deg2m: float = 1852 * 60.0
 ) -> float:
-    dxdxsi, dxdeta, dxdzet, dydxsi, dydeta, dydzet, dzdxsi, dzdeta, dzdzet = dxdxsi3D_lin(hexa_z, hexa_y, hexa_x, zeta, eta, xsi, mesh)
+    dxdxsi, dxdeta, dxdzet, dydxsi, dydeta, dydzet, dzdxsi, dzdeta, dzdzet = dxdxsi3D_lin(hexa_z, hexa_y, hexa_x, zeta, eta, xsi, mesh, deg2m)
 
     jac = (
         dxdxsi * (dydeta * dzdzet - dzdeta * dydzet)
@@ -112,7 +113,7 @@ def jacobian3D_lin_face(
     eta: float,
     xsi: float,
     orientation: Literal["zonal", "meridional", "vertical"],
-    mesh: Mesh,
+    mesh: BaseMesh,
 ) -> float:
     dxdxsi, dxdeta, dxdzet, dydxsi, dydeta, dydzet, dzdxsi, dzdeta, dzdzet = dxdxsi3D_lin(hexa_z, hexa_y, hexa_x, zeta, eta, xsi, mesh)
 
@@ -174,10 +175,11 @@ def interpolate(phi: Callable[[float], list[float]], f: list[float], xsi: float)
     return np.dot(phi(xsi), f)
 
 
-def _geodetic_distance(lat1: float, lat2: float, lon1: float, lon2: float, mesh: Mesh, lat: float) -> float:
-    if mesh == "spherical":
+def _geodetic_distance(
+    lat1: float, lat2: float, lon1: float, lon2: float, mesh: BaseMesh, lat: float, deg2m: float = 1852 * 60.0
+) -> float:
+    if mesh.is_spherical():
         rad = np.pi / 180.0
-        deg2m = 1852 * 60.0
         return np.sqrt(((lon2 - lon1) * deg2m * np.cos(rad * lat)) ** 2 + ((lat2 - lat1) * deg2m) ** 2)
     else:
         return np.sqrt((lon2 - lon1) ** 2 + (lat2 - lat1) ** 2)