diff --git a/src/parcels/_core/model.py b/src/parcels/_core/model.py index e1c44de30..a2d91512a 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, @@ -275,7 +276,12 @@ def from_sgrid_conventions( model._fields = model.construct_fields() 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, + model.vector_field_components, + dataset=model.data, + )() return model @@ -512,6 +518,43 @@ def _discover_ux_U_and_V(ds: ux.UxDataset) -> ux.UxDataset: return ds +def _select_xinterpolator( + da: xr.DataArray, + sgrid_metadata: sgrid.SGrid2DMetadata, + vector_fields: ptyping.VectorFields | None = None, + dataset: xr.Dataset | None = None, +) -> 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 should be interpolated constant + over the grid cell using CGrid_Tracer. Non-tracer fields, node-located variables, + or fields on unstaggered/A-grids default to XLinear. + """ + # If the model grid is an A-grid (velocity components defined on identical dimensions), + # all fields are collocated; default to XLinear. + if dataset is not None and vector_fields: + for comps in vector_fields.values(): + if len(comps) >= 2 and comps[0] in dataset and comps[1] in dataset: + if _is_agrid(dataset, u=comps[0], v=comps[1]): + return XLinear + + # Tracers on a C-grid must be defined at the face (cell center) dimensions + fnp_x, fnp_y = sgrid_metadata.face_dimensions + face_dims = {fnp_x.face, fnp_y.face} + if not face_dims.issubset(da.dims): + return XLinear + + # If vertical dimensions are present and staggered, interface-defined variables + # (e.g. vertical velocity W defined on node) are not cell-centered tracers + if sgrid_metadata.vertical_dimensions: + fnp_z = sgrid_metadata.vertical_dimensions[0] + if fnp_z.face != fnp_z.node and fnp_z.node in da.dims: + return XLinear + + return CGrid_Tracer + + def _select_uxinterpolator(da: ux.UxDataArray): """Selects the appropriate uxarray interpolator for a given uxarray UxDataArray""" supported_uxinterp_mapping = { 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..221559f7e 100644 --- a/tests/test_fieldset.py +++ b/tests/test_fieldset.py @@ -15,7 +15,11 @@ from parcels import Particle, ParticleFile, ParticleSet, Variable, 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.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, XLinear_Velocity ds = datasets_structured["ds_2d_left"] @@ -722,3 +727,79 @@ 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 + + # 3D vertical interface variable (node along Z) -> XLinear + da_z_interface = da_3d_tracer.rename({"ZC": "ZG"}) + assert _select_xinterpolator(da_z_interface, metadata) is XLinear + + # 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 dataset with collocated velocity components -> XLinear for all scalar fields + ds_agrid = datasets_structured["ds_2d_left"] + assert ( + _select_xinterpolator( + da_3d_tracer, + metadata, + vector_fields={"UV": ("U_A_grid", "V_A_grid")}, + dataset=ds_agrid, + ) + is XLinear + ) + + +def test_agrid_all_fields_default_to_xlinear(): + """Verify that when a dataset is configured as an A-grid, all scalar fields default to XLinear.""" + ds = datasets_structured["ds_2d_left"] + fieldset = FieldSet.from_sgrid_conventions( + ds, + vector_fields={"UV": ("U_A_grid", "V_A_grid")}, + ) + # Even though data_c has face dimensions, on an A-grid it must default to XLinear + assert isinstance(fieldset.data_c.interp_method, XLinear) + assert isinstance(fieldset.data_g.interp_method, XLinear) + assert isinstance(fieldset.U_A_grid.interp_method, XLinear) + assert isinstance(fieldset.V_A_grid.interp_method, XLinear) + assert isinstance(fieldset.UV.interp_method, XLinear_Velocity)