Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
45 changes: 44 additions & 1 deletion src/parcels/_core/model.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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


Expand Down Expand Up @@ -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 = {
Expand Down
2 changes: 1 addition & 1 deletion src/parcels/_datasets/structured/generated.py
Original file line number Diff line number Diff line change
Expand Up @@ -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"}),
Expand Down
83 changes: 82 additions & 1 deletion tests/test_fieldset.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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"]

Expand Down Expand Up @@ -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)
Loading