Refactorings after review

- Adds a FlatMesh object
- Remove assert_valid_mesh

Author: VeckoTheGecko

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

@@ -152,7 +152,7 @@ def add_field(self, field: Field, name: str | None = None):
 
         self.fields[name] = field
 
-    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
 
@@ -249,7 +249,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,7 +141,7 @@ 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 /= (
                         self.fieldset.U.grid.deg2m
                     )  # TODO does not account for zonal variation in meter -> degree conversion

src/parcels/_core/mesh.py

@@ -1,9 +1,19 @@
+from abc import ABC, abstractmethod
+from typing import Literal
+
 import numpy as np
 
 EARTH_RADIUS = 6366707.019493707
 
 
-class SphericalMesh:
+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=...)``.
@@ -22,5 +32,41 @@ def deg2m(self) -> float:
         """Meters per degree of arc."""
         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

@@ -21,7 +21,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,
@@ -83,7 +82,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)}")
 
@@ -132,7 +131,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:
@@ -279,7 +278,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(
@@ -303,7 +304,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

@@ -45,7 +45,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():
                 # Boundaries of the hash grid are the unit cube
                 self._xmin = -1.0
                 self._ymin = -1.0
@@ -148,7 +148,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():
                 # Boundaries of the hash grid are the unit cube
                 self._xmin = -1.0
                 self._ymin = -1.0
@@ -346,7 +346,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':
+    if mesh.is_spherical():
         rad = np.pi / 180.
         lat = (1-xsi) * (1-eta) * hexa_y[0] + \
                  xsi  * (1-eta) * hexa_y[1] + \
@@ -92,7 +92,7 @@ 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, deg2m)
@@ -113,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)
 
@@ -176,9 +176,9 @@ def interpolate(phi: Callable[[float], list[float]], f: list[float], xsi: float)
 
 
 def _geodetic_distance(
-    lat1: float, lat2: float, lon1: float, lon2: float, mesh: Mesh, lat: float, deg2m: float = 1852 * 60.0
+    lat1: float, lat2: float, lon1: float, lon2: float, mesh: BaseMesh, lat: float, deg2m: float = 1852 * 60.0
 ) -> float:
-    if mesh == "spherical":
+    if mesh.is_spherical():
         rad = np.pi / 180.0
         return np.sqrt(((lon2 - lon1) * deg2m * np.cos(rad * lat)) ** 2 + ((lat2 - lat1) * deg2m) ** 2)
     else:

src/parcels/_core/uxgrid.py

@@ -7,8 +7,7 @@
 
 from parcels._core.basegrid import BaseGrid
 from parcels._core.index_search import GRID_SEARCH_ERROR, _search_1d_array, uxgrid_point_in_cell
-from parcels._core.mesh import EARTH_RADIUS, SphericalMesh
-from parcels._typing import assert_valid_mesh
+from parcels._core.mesh import SphericalMesh, get_mesh
 
 _UXGRID_AXES = Literal["Z", "FACE"]
 
@@ -19,7 +18,9 @@ class UxGrid(BaseGrid):
     for interpolation on unstructured grids.
     """
 
-    def __init__(self, grid: ux.grid.Grid, z: ux.UxDataArray, mesh) -> None:
+    def __init__(
+        self, grid: ux.grid.Grid, z: ux.UxDataArray, mesh: Literal["flat", "spherical"] | SphericalMesh
+    ) -> None:
         """
         Initializes the UxGrid with a uxarray grid and vertical coordinate array.
 
@@ -42,16 +43,9 @@ def __init__(self, grid: ux.grid.Grid, z: ux.UxDataArray, mesh) -> None:
         if z.ndim != 1:
             raise ValueError("z must be a 1D array of vertical coordinates")
         self.z = z
-        if isinstance(mesh, SphericalMesh):
-            self._mesh = "spherical"
-            self._radius = mesh.radius
-        else:
-            self._mesh = mesh
-            self._radius = EARTH_RADIUS if mesh == "spherical" else None
+        self._mesh = get_mesh(mesh)
         self._spatialhash = None
 
-        assert_valid_mesh(mesh)
-
     @property
     def depth(self):
         """
@@ -82,9 +76,9 @@ def get_axis_dim(self, axis: _UXGRID_AXES) -> int:
     @property
     def deg2m(self) -> float:
         """Metres per arcdegree for this grid's mesh."""
-        if self._radius is None:  # flat mesh; None causes crash in advection
+        if not self._mesh.is_spherical():
             return 1.0
-        return self._radius * np.pi / 180.0
+        return self._mesh.deg2m
 
     def search(self, z, y, x, ei=None, tol=1e-6):
         """