diff --git a/src/parcels/_core/model.py b/src/parcels/_core/model.py index e1c44de30..f72b48146 100644 --- a/src/parcels/_core/model.py +++ b/src/parcels/_core/model.py @@ -25,6 +25,7 @@ from parcels._logger import logger from parcels._python import NOTSET, NotSetType from parcels.interpolators import ( + CGrid_Tracer, CGrid_Velocity, Ux_Velocity, UxConstantEdgeConstantZC, @@ -216,7 +217,7 @@ def construct_fields(self) -> list[Field | VectorField]: for vfield_name, components in self.vector_field_components.items(): interp_method = ( - XLinear_Velocity() if _is_agrid(self.data, u=components[0], v=components[1]) else CGrid_Velocity() + CGrid_Velocity() if _is_cgrid(self.data, u=components[0], v=components[1]) else XLinear_Velocity() ) component_fields = [single_fields[name] for name in components] @@ -225,6 +226,17 @@ def construct_fields(self) -> list[Field | VectorField]: fields: dict[str, Field | VectorField] = {**single_fields, **vector_fields} return list(fields.values()) + @property + def is_cgrid(self) -> bool: + """Whether the model grid is a staggered C-grid.""" + if self.vector_field_components: + return any( + _is_cgrid(self.data, u=comps[0], v=comps[1]) + for comps in self.vector_field_components.values() + if len(comps) >= 2 and comps[0] in self.data and comps[1] in self.data + ) + return _is_cgrid(self.data, sgrid_metadata=self.grid.sgrid_metadata) + @classmethod def from_sgrid_conventions( cls, @@ -273,9 +285,14 @@ def from_sgrid_conventions( skip_field_data_validation=skip_field_data_validation, ) model._fields = model.construct_fields() + is_cgrid = model.is_cgrid for f in model._fields: if isinstance(f, Field): - f.interp_method = XLinear() + f.interp_method = _select_xinterpolator( + model.data[f.name], + model.grid.sgrid_metadata, + is_cgrid=is_cgrid, + )() return model @@ -512,6 +529,26 @@ def _discover_ux_U_and_V(ds: ux.UxDataset) -> ux.UxDataset: return ds +def _select_xinterpolator( + da: xr.DataArray, + sgrid_metadata: sgrid.SGrid2DMetadata, + *, + is_cgrid: bool = True, +) -> type[ScalarInterpolator]: + """Selects the appropriate interpolator for a given xarray DataArray on an SGRID grid. + + Following Delandmeter and Van Sebille (2019), tracer fields on an Arakawa C-grid + are located at cell centers (face dimensions) and default to CGrid_Tracer. + All other fields default to XLinear. + """ + if is_cgrid: + fnp_x, fnp_y = sgrid_metadata.face_dimensions + if {fnp_x.face, fnp_y.face}.issubset(da.dims): + return CGrid_Tracer + + return XLinear + + def _select_uxinterpolator(da: ux.UxDataArray): """Selects the appropriate uxarray interpolator for a given uxarray UxDataArray""" supported_uxinterp_mapping = { @@ -553,10 +590,56 @@ def _select_uxinterpolator(da: ux.UxDataArray): return None +def _is_cgrid( + ds: xr.Dataset, + u: str | None = None, + v: str | None = None, + sgrid_metadata: sgrid.SGrid2DMetadata | None = None, +) -> bool: + """Check if a dataset or specified (u, v) components represent a staggered C-grid. + + Following Arakawa C-grid conventions, horizontal velocity components U and V + are defined on different dimensions (edges). If u and v are defined on identical + dimensions, the grid is an A-grid or B-grid. + + When u and v are not specified (e.g. for scalar-only FieldSets), the dataset + is checked for default velocity components or edge-staggered variables. + """ + if u is not None and v is not None: + # check if U and V are defined on different dimensions + # if yes, interpret as C grid (if no, it could be an A or B grid) + return set(ds[u].dims) != set(ds[v].dims) + + # If dataset has default U and V variables, check their dimensions + default_vfields = _default_vector_field_components(list(ds.data_vars)) + for comps in default_vfields.values(): + if len(comps) >= 2 and comps[0] in ds and comps[1] in ds: + return _is_cgrid(ds, u=comps[0], v=comps[1]) + + # If SGRID metadata is available, check for edge-staggered variables + if sgrid_metadata is None and hasattr(ds, "sgrid"): + try: + sgrid_metadata = ds.sgrid.metadata + except Exception: + pass + + if sgrid_metadata is not None: + fnp_x, fnp_y = sgrid_metadata.face_dimensions + if fnp_x.face != fnp_x.node and fnp_y.face != fnp_y.node: + edge1 = {fnp_x.face, fnp_y.node} + edge2 = {fnp_x.node, fnp_y.face} + has_edge1 = any(edge1.issubset(da.dims) for da in ds.data_vars.values()) + has_edge2 = any(edge2.issubset(da.dims) for da in ds.data_vars.values()) + if has_edge1 or has_edge2: + return True + + return False + + def _is_agrid(ds: xr.Dataset, u: str, v: str) -> bool: # check if U and V are defined on the same dimensions # if yes, interpret as A grid - return set(ds[u].dims) == set(ds[v].dims) + return not _is_cgrid(ds, u=u, v=v) def _get_time_interval(data: xr.DataArray | ux.UxDataArray) -> TimeInterval | None: diff --git a/src/parcels/_datasets/structured/generated.py b/src/parcels/_datasets/structured/generated.py index 04c6fe473..8c696bb43 100644 --- a/src/parcels/_datasets/structured/generated.py +++ b/src/parcels/_datasets/structured/generated.py @@ -310,7 +310,7 @@ def peninsula_dataset(xdim=100, ydim=50, mesh_str="flat", grid_type="A"): { "U": (Udims, U), "V": (Vdims, V), - "P": (["YC", "XC"], P), + "P": (["YG", "XG"], P), }, coords={ "YC": (["YC"], np.arange(ydim) - 0.5, {"axis": "Y"}), diff --git a/tests/test_fieldset.py b/tests/test_fieldset.py index 7fd63f5c9..598e47484 100644 --- a/tests/test_fieldset.py +++ b/tests/test_fieldset.py @@ -12,10 +12,14 @@ import parcels.tutorial import tests -from parcels import Particle, ParticleFile, ParticleSet, Variable, convert +from parcels import Particle, ParticleFile, ParticleSet, Variable, VectorField, convert from parcels._core.fieldset import CalendarError, FieldSet, IncompatibleMeshesException, _datetime_to_msg -from parcels._core.mesh import SphericalMesh -from parcels._core.model import _default_vector_field_components, _infer_mesh_from_sgrid_data +from parcels._core.mesh import FlatMesh, SphericalMesh +from parcels._core.model import ( + _default_vector_field_components, + _infer_mesh_from_sgrid_data, + _select_xinterpolator, +) from parcels._core.utils.time import TimeInterval from parcels._datasets.structured.generated import simple_UV_dataset, simple_UV_sigma_dataset from parcels._datasets.structured.generic import datasets as datasets_structured @@ -24,6 +28,7 @@ from parcels._datasets.unstructured.generated import simple_small_delaunay from parcels._datasets.unstructured.generic import datasets as datasets_unstructured from parcels._sgrid.accessor import SGridDatasetInconsistency +from parcels.interpolators import CGrid_Tracer, CGrid_Velocity, XLinear ds = datasets_structured["ds_2d_left"] @@ -722,3 +727,119 @@ def test_fieldset_from_simple_UV_sigma_dataset(n_snapshots): assert depth.dims == ("time", "ZF", "YC", "XC") assert depth.sizes["time"] == n_snapshots assert fieldset.U.grid.get_axis_dim("Z") == nz + + +def test_cgrid_tracer_default_interpolator(): + """Verify that tracer fields on C-grid default to CGrid_Tracer while non-tracers default to XLinear.""" + ds = datasets_structured["ds_2d_left"] + fieldset = FieldSet.from_sgrid_conventions( + ds, + vector_fields={"UV": ("U_C_grid", "V_C_grid")}, + ) + + # data_c is defined on cell centers (face dimensions) -> CGrid_Tracer + assert isinstance(fieldset.data_c.interp_method, CGrid_Tracer) + + # data_g is defined on cell vertices (node dimensions) -> XLinear + assert isinstance(fieldset.data_g.interp_method, XLinear) + + # Scalar velocity components default to XLinear + assert isinstance(fieldset.U_C_grid.interp_method, XLinear) + assert isinstance(fieldset.V_C_grid.interp_method, XLinear) + assert isinstance(fieldset.U_A_grid.interp_method, XLinear) + assert isinstance(fieldset.V_A_grid.interp_method, XLinear) + + # Vector velocity field defaults to CGrid_Velocity + assert isinstance(fieldset.UV.interp_method, CGrid_Velocity) + + +def test_select_xinterpolator_edge_cases(): + """Directly test _select_xinterpolator behavior across various dimensional configurations.""" + metadata = datasets_structured["ds_2d_left"].sgrid.metadata + da_3d_tracer = datasets_structured["ds_2d_left"]["data_c"] + da_node = datasets_structured["ds_2d_left"]["data_g"] + da_edge_u = datasets_structured["ds_2d_left"]["U_C_grid"] + + # 3D cell-center tracer -> CGrid_Tracer + assert _select_xinterpolator(da_3d_tracer, metadata) is CGrid_Tracer + + # 2D cell-center tracer (no vertical dim) -> CGrid_Tracer + da_2d_tracer = da_3d_tracer.isel(ZC=0, drop=True) + assert _select_xinterpolator(da_2d_tracer, metadata) is CGrid_Tracer + + # Node-defined variable -> XLinear + assert _select_xinterpolator(da_node, metadata) is XLinear + + # Edge-defined velocity variable -> XLinear + assert _select_xinterpolator(da_edge_u, metadata) is XLinear + + # A-grid configuration -> XLinear + assert _select_xinterpolator(da_3d_tracer, metadata, is_cgrid=False) is XLinear + + +def test_cgrid_tracer_without_vector_fields(): + """Verify that cell-centred tracers on C-grid default to CGrid_Tracer even when no VectorFields are present.""" + ds = datasets_structured["ds_2d_left"] + fieldset = FieldSet.from_sgrid_conventions( + ds, + vector_fields={}, + ) + + # data_c is defined on cell centers (face dimensions) -> CGrid_Tracer + assert isinstance(fieldset.data_c.interp_method, CGrid_Tracer) + + # data_g is defined on cell vertices (node dimensions) -> XLinear + assert isinstance(fieldset.data_g.interp_method, XLinear) + + # Scalar velocity components retain XLinear + assert isinstance(fieldset.U_C_grid.interp_method, XLinear) + assert isinstance(fieldset.V_C_grid.interp_method, XLinear) + + # FieldSet contains no VectorFields + assert not any(isinstance(f, VectorField) for f in fieldset.fields.values()) + + +def test_agrid_and_bgrid_without_vector_fields(): + """Verify that A-grid and B-grid fields default to XLinear when no VectorFields are present.""" + from parcels._datasets.structured.generated import peninsula_dataset + from parcels._sgrid.core import FaceNodePadding, Padding, SGrid2DMetadata, _attach_sgrid_metadata + + # A-grid dataset with no VectorFields + ds_a = peninsula_dataset(grid_type="A") + fieldset_a = FieldSet.from_sgrid_conventions(ds_a, vector_fields={}) + assert isinstance(fieldset_a.U.interp_method, XLinear) + assert isinstance(fieldset_a.V.interp_method, XLinear) + assert isinstance(fieldset_a.P.interp_method, XLinear) + + # B-grid dataset where U and V are both collocated on nodes (YG, XG) + ds_b = xr.Dataset( + { + "tracer": (["YC", "XC"], np.ones((5, 5))), + "U": (["YG", "XG"], np.ones((6, 6))), + "V": (["YG", "XG"], np.ones((6, 6))), + }, + coords={ + "YC": (["YC"], np.arange(5) + 0.5, {"axis": "Y"}), + "YG": (["YG"], np.arange(6), {"axis": "Y"}), + "XC": (["XC"], np.arange(5) + 0.5, {"axis": "X"}), + "XG": (["XG"], np.arange(6), {"axis": "X"}), + "lon": (["XG"], np.linspace(0, 10, 6), {"axis": "X", "units": "m"}), + "lat": (["YG"], np.linspace(0, 10, 6), {"axis": "Y", "units": "m"}), + }, + ).pipe( + _attach_sgrid_metadata, + SGrid2DMetadata( + cf_role="grid_topology", + topology_dimension=2, + node_dimensions=("XG", "YG"), + node_coordinates=("lon", "lat"), + face_dimensions=( + FaceNodePadding("XC", "XG", Padding.NONE), + FaceNodePadding("YC", "YG", Padding.NONE), + ), + ), + ) + fieldset_b = FieldSet.from_sgrid_conventions(ds_b, mesh=FlatMesh(), vector_fields={}) + assert isinstance(fieldset_b.tracer.interp_method, XLinear) + assert isinstance(fieldset_b.U.interp_method, XLinear) + assert isinstance(fieldset_b.V.interp_method, XLinear)