Skip to content
4 changes: 2 additions & 2 deletions docs/development/posting-issues.md
Original file line number Diff line number Diff line change
Expand Up @@ -46,8 +46,8 @@ As a user with access to your dataset, you would do:

# Generate an example dataset to zip. The user would use their own.
import xarray as xr
from parcels._datasets.structured.generic import datasets_comodo
datasets_comodo['ds_2d_left'].to_netcdf("my_dataset.nc")
from parcels._datasets.structured.generic import datasets
datasets['ds_2d_left'].to_netcdf("my_dataset.nc")
```

```{code-cell}
Expand Down
2 changes: 1 addition & 1 deletion src/parcels/_core/xgrid.py
Original file line number Diff line number Diff line change
Expand Up @@ -40,7 +40,7 @@ def _drop_field_data(ds: xr.Dataset) -> xr.Dataset:
when passed to the XGCM grid, the object only functions as an in memory representation
of the grid.
"""
return ds.drop_vars(ds.data_vars)
return ds.drop_vars(set(ds.data_vars) - {"grid"}) # don't drop sgrid metadata
Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This function will be removed in a future PR anyway (once from_dataset is removed)



def assert_all_field_dims_have_axis(da: xr.DataArray, xgcm_grid: xgcm.Grid) -> None:
Expand Down
277 changes: 185 additions & 92 deletions src/parcels/_datasets/structured/generic.py
Original file line number Diff line number Diff line change
Expand Up @@ -5,7 +5,7 @@

from . import T, X, Y, Z

__all__ = ["T", "X", "Y", "Z", "datasets_comodo", "datasets_sgrid"]
__all__ = ["T", "X", "Y", "Z", "datasets", "datasets_sgrid"]

TIME = xr.date_range("2000", "2001", T)

Expand Down Expand Up @@ -139,8 +139,21 @@ def _unrolled_cone_curvilinear_grid():
)


datasets_comodo = {
"2d_left_rotated": _rotated_curvilinear_grid(),
datasets = {
"2d_left_rotated": _rotated_curvilinear_grid().pipe(
sgrid._attach_sgrid_metadata,
sgrid.SGrid2DMetadata(
cf_role="grid_topology",
topology_dimension=2,
node_dimensions=("XG", "YG"),
face_dimensions=(
sgrid.FaceNodePadding("XC", "XG", sgrid.Padding.HIGH),
sgrid.FaceNodePadding("YC", "YG", sgrid.Padding.HIGH),
),
node_coordinates=("lon", "lat"),
vertical_dimensions=(sgrid.FaceNodePadding("ZC", "ZG", sgrid.Padding.HIGH),),
),
),
"ds_2d_left": xr.Dataset( # MITgcm indexing style
{
"data_g": (["time", "ZG", "YG", "XG"], np.random.rand(T, Z, Y, X)),
Expand Down Expand Up @@ -182,6 +195,19 @@ def _unrolled_cone_curvilinear_grid():
"depth": (["ZG"], np.arange(Z)),
"time": (["time"], TIME, {"axis": "T"}),
},
).pipe(
sgrid._attach_sgrid_metadata,
sgrid.SGrid2DMetadata(
cf_role="grid_topology",
topology_dimension=2,
node_dimensions=("XG", "YG"),
face_dimensions=(
sgrid.FaceNodePadding("XC", "XG", sgrid.Padding.HIGH),
sgrid.FaceNodePadding("YC", "YG", sgrid.Padding.HIGH),
),
node_coordinates=("lon", "lat"),
vertical_dimensions=(sgrid.FaceNodePadding("ZC", "ZG", sgrid.Padding.HIGH),),
),
),
"ds_2d_right": xr.Dataset( # NEMO indexing style
{
Expand Down Expand Up @@ -224,105 +250,172 @@ def _unrolled_cone_curvilinear_grid():
"depth": (["ZG"], np.arange(Z)),
"time": (["time"], TIME, {"axis": "T"}),
},
),
"2d_left_unrolled_cone": _unrolled_cone_curvilinear_grid(),
}

_COMODO_TO_2D_SGRID = { # Note "2D SGRID" here is meant in the context of SGRID convention (i.e., 1D depth)
"XG": "node_dimension1",
"YG": "node_dimension2",
"XC": "face_dimension1",
"YC": "face_dimension2",
"ZG": "vertical_dimensions_dim1",
"ZC": "vertical_dimensions_dim2",
}
datasets_sgrid = {
"ds_2d_padded_high": (
datasets_comodo["ds_2d_left"]
.pipe(
sgrid._attach_sgrid_metadata,
sgrid.SGrid2DMetadata(
cf_role="grid_topology",
topology_dimension=2,
node_dimensions=("XG", "YG"),
face_dimensions=(
sgrid.FaceNodePadding("XC", "XG", sgrid.Padding.HIGH),
sgrid.FaceNodePadding("YC", "YG", sgrid.Padding.HIGH),
),
node_coordinates=("lon", "lat"),
vertical_dimensions=(sgrid.FaceNodePadding("ZC", "ZG", sgrid.Padding.HIGH),),
).pipe(
sgrid._attach_sgrid_metadata,
sgrid.SGrid2DMetadata(
cf_role="grid_topology",
topology_dimension=2,
node_dimensions=("XG", "YG"),
face_dimensions=(
sgrid.FaceNodePadding("XC", "XG", sgrid.Padding.LOW),
sgrid.FaceNodePadding("YC", "YG", sgrid.Padding.LOW),
),
)
.sgrid.rename(
_COMODO_TO_2D_SGRID,
)
node_coordinates=("lon", "lat"),
vertical_dimensions=(sgrid.FaceNodePadding("ZC", "ZG", sgrid.Padding.LOW),),
),
),
"ds_2d_padded_low": (
datasets_comodo["ds_2d_right"]
.pipe(
sgrid._attach_sgrid_metadata,
sgrid.SGrid2DMetadata(
cf_role="grid_topology",
topology_dimension=2,
node_dimensions=("XG", "YG"),
face_dimensions=(
sgrid.FaceNodePadding("XC", "XG", sgrid.Padding.LOW),
sgrid.FaceNodePadding("YC", "YG", sgrid.Padding.LOW),
),
node_coordinates=("lon", "lat"),
vertical_dimensions=(sgrid.FaceNodePadding("ZC", "ZG", sgrid.Padding.LOW),),
),
)
.sgrid.rename(
_COMODO_TO_2D_SGRID,
)
),
"ds_2d_padded_none": xr.Dataset(
"ds_2d_inner": xr.Dataset(
{
"data_g": (["node_dimension1", "node_dimension2"], np.random.rand(10, 10)),
"data_c": (["face_dimension1", "face_dimension2"], np.random.rand(9, 9)),
"grid": (
[],
np.array(0),
sgrid.SGrid2DMetadata(
cf_role="grid_topology",
topology_dimension=2,
node_dimensions=("node_dimension1", "node_dimension2"),
face_dimensions=(
sgrid.FaceNodePadding("face_dimension1", "node_dimension1", sgrid.Padding.NONE),
sgrid.FaceNodePadding("face_dimension2", "node_dimension2", sgrid.Padding.NONE),
),
node_coordinates=("lon", "lat"),
).to_attrs(),
),
"data_g": (["time", "ZG", "YG", "XG"], np.random.rand(T, Z, Y, X)),
"data_c": (["time", "ZC", "YC", "XC"], np.random.rand(T, Z - 1, Y - 1, X - 1)),
"U_A_grid": (["time", "ZG", "YG", "XG"], np.random.rand(T, Z, Y, X)),
"V_A_grid": (["time", "ZG", "YG", "XG"], np.random.rand(T, Z, Y, X)),
"U_C_grid": (["time", "ZG", "YC", "XG"], np.random.rand(T, Z, Y - 1, X)),
"V_C_grid": (["time", "ZG", "YG", "XC"], np.random.rand(T, Z, Y, X - 1)),
},
coords={
"lon": (["node_dimension1"], np.linspace(0, 1, 10)),
"lat": (["node_dimension2"], np.linspace(0, 1, 10)),
"XG": (
["XG"],
2 * np.pi / X * np.arange(0, X),
{"axis": "X", "c_grid_axis_shift": 0.5},
),
"XC": (["XC"], 2 * np.pi / X * (np.arange(0, X - 1) - 0.5), {"axis": "X"}),
"YG": (
["YG"],
2 * np.pi / (Y) * np.arange(0, Y),
{"axis": "Y", "c_grid_axis_shift": 0.5},
),
"YC": (
["YC"],
2 * np.pi / (Y) * (np.arange(0, Y - 1) - 0.5),
{"axis": "Y"},
),
"ZG": (
["ZG"],
np.arange(Z),
{"axis": "Z", "c_grid_axis_shift": 0.5},
),
"ZC": (
["ZC"],
np.arange(Z - 1) - 0.5,
{"axis": "Z"},
),
"lon": (["XG"], 2 * np.pi / X * np.arange(0, X)),
"lat": (["YG"], 2 * np.pi / (Y) * np.arange(0, Y)),
"depth": (["ZG"], np.arange(Z)),
"time": (["time"], TIME, {"axis": "T"}),
},
).pipe(
sgrid._attach_sgrid_metadata,
sgrid.SGrid2DMetadata(
cf_role="grid_topology",
topology_dimension=2,
node_dimensions=("XG", "YG"),
face_dimensions=(
sgrid.FaceNodePadding("XC", "XG", sgrid.Padding.BOTH),
sgrid.FaceNodePadding("YC", "YG", sgrid.Padding.BOTH),
),
node_coordinates=("lon", "lat"),
vertical_dimensions=(sgrid.FaceNodePadding("ZC", "ZG", sgrid.Padding.BOTH),),
),
),
"ds_2d_padded_both": xr.Dataset(
"ds_2d_outer": xr.Dataset(
{
"data_g": (["node_dimension1", "node_dimension2"], np.random.rand(10, 10)),
"data_c": (["face_dimension1", "face_dimension2"], np.random.rand(11, 11)),
"grid": (
[],
np.array(0),
sgrid.SGrid2DMetadata(
cf_role="grid_topology",
topology_dimension=2,
node_dimensions=("node_dimension1", "node_dimension2"),
face_dimensions=(
sgrid.FaceNodePadding("face_dimension1", "node_dimension1", sgrid.Padding.BOTH),
sgrid.FaceNodePadding("face_dimension2", "node_dimension2", sgrid.Padding.BOTH),
),
node_coordinates=("lon", "lat"),
).to_attrs(),
),
"data_g": (["time", "ZG", "YG", "XG"], np.random.rand(T, Z, Y, X)),
"data_c": (["time", "ZC", "YC", "XC"], np.random.rand(T, Z + 1, Y + 1, X + 1)),
"U_A_grid": (["time", "ZG", "YG", "XG"], np.random.rand(T, Z, Y, X)),
"V_A_grid": (["time", "ZG", "YG", "XG"], np.random.rand(T, Z, Y, X)),
"U_C_grid": (["time", "ZG", "YC", "XG"], np.random.rand(T, Z, Y + 1, X)),
"V_C_grid": (["time", "ZG", "YG", "XC"], np.random.rand(T, Z, Y, X + 1)),
},
coords={
"lon": (["node_dimension1"], np.linspace(0, 1, 10)),
"lat": (["node_dimension2"], np.linspace(0, 1, 10)),
"XG": (
["XG"],
2 * np.pi / X * np.arange(0, X),
{"axis": "X", "c_grid_axis_shift": -0.5},
),
"XC": (["XC"], 2 * np.pi / X * (np.arange(0, X + 1) + 0.5), {"axis": "X"}),
"YG": (
["YG"],
2 * np.pi / (Y) * np.arange(0, Y),
{"axis": "Y", "c_grid_axis_shift": -0.5},
),
"YC": (
["YC"],
2 * np.pi / (Y) * (np.arange(0, Y + 1) + 0.5),
{"axis": "Y"},
),
"ZG": (
["ZG"],
np.arange(Z),
{"axis": "Z", "c_grid_axis_shift": -0.5},
),
"ZC": (
["ZC"],
np.arange(Z + 1) + 0.5,
{"axis": "Z"},
),
"lon": (["XG"], 2 * np.pi / X * np.arange(0, X)),
"lat": (["YG"], 2 * np.pi / (Y) * np.arange(0, Y)),
"depth": (["ZG"], np.arange(Z)),
"time": (["time"], TIME, {"axis": "T"}),
},
).pipe(
sgrid._attach_sgrid_metadata,
sgrid.SGrid2DMetadata(
cf_role="grid_topology",
topology_dimension=2,
node_dimensions=("XG", "YG"),
face_dimensions=(
sgrid.FaceNodePadding("XC", "XG", sgrid.Padding.NONE),
sgrid.FaceNodePadding("YC", "YG", sgrid.Padding.NONE),
),
node_coordinates=("lon", "lat"),
vertical_dimensions=(sgrid.FaceNodePadding("ZC", "ZG", sgrid.Padding.NONE),),
),
),
"2d_left_unrolled_cone": _unrolled_cone_curvilinear_grid().pipe(
sgrid._attach_sgrid_metadata,
sgrid.SGrid2DMetadata(
cf_role="grid_topology",
topology_dimension=2,
node_dimensions=("XG", "YG"),
face_dimensions=(
sgrid.FaceNodePadding("XC", "XG", sgrid.Padding.HIGH),
sgrid.FaceNodePadding("YC", "YG", sgrid.Padding.HIGH),
),
node_coordinates=("lon", "lat"),
vertical_dimensions=(sgrid.FaceNodePadding("ZC", "ZG", sgrid.Padding.HIGH),),
),
),
}

_COMODO_TO_2D_SGRID = { # Note "2D SGRID" here is meant in the context of SGRID convention (i.e., 1D depth)
"XG": "node_dimension1",
"YG": "node_dimension2",
"XC": "face_dimension1",
"YC": "face_dimension2",
"ZG": "vertical_dimensions_dim1",
"ZC": "vertical_dimensions_dim2",
}

_DATASET_NAME_TO_SGRID_NAME = {
"ds_2d_left": "ds_2d_padded_high",
"ds_2d_right": "ds_2d_padded_low",
"ds_2d_outer": "ds_2d_padded_none",
"ds_2d_inner": "ds_2d_padded_both",
}


def create_datasets_sgrid(datasets_: dict[str, xr.Dataset]) -> dict[str, xr.Dataset]:
ret = {}
for name, ds in datasets_.items():
if name not in _DATASET_NAME_TO_SGRID_NAME:
continue
ret[_DATASET_NAME_TO_SGRID_NAME[name]] = ds.sgrid.rename(
_COMODO_TO_2D_SGRID,
Comment thread
VeckoTheGecko marked this conversation as resolved.
)
return ret


datasets_sgrid = create_datasets_sgrid(datasets)
2 changes: 1 addition & 1 deletion src/parcels/_sgrid/core.py
Original file line number Diff line number Diff line change
Expand Up @@ -695,7 +695,7 @@ def _attach_sgrid_metadata(ds: xr.Dataset, grid: SGrid2DMetadata | SGrid3DMetada
0,
grid.to_attrs(),
)
ds.attrs["Conventions"] = "SGRID"
# ds.attrs["Conventions"] = "SGRID" # TODO: re-enable once XGrid.from_dataset is gone
Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This was needed to get around some xgcm internals - will be re-enabled in a future PR.

return ds


Expand Down
6 changes: 3 additions & 3 deletions tests/datasets/test_structured.py
Original file line number Diff line number Diff line change
@@ -1,12 +1,12 @@
import xgcm

from parcels._core.xgrid import _DEFAULT_XGCM_KWARGS
from parcels._datasets.structured.generic import datasets_comodo
from parcels._datasets.structured.generic import datasets


def test_left_indexed_dataset():
"""Checks that 'ds_2d_left' is right indexed on all variables."""
ds = datasets_comodo["ds_2d_left"]
ds = datasets["ds_2d_left"]
grid = xgcm.Grid(ds, **_DEFAULT_XGCM_KWARGS)

for _axis_name, axis in grid.axes.items():
Expand All @@ -16,7 +16,7 @@ def test_left_indexed_dataset():

def test_right_indexed_dataset():
"""Checks that 'ds_2d_right' is right indexed on all variables."""
ds = datasets_comodo["ds_2d_right"]
ds = datasets["ds_2d_right"]
grid = xgcm.Grid(ds, **_DEFAULT_XGCM_KWARGS)
for _axis_name, axis in grid.axes.items():
for pos, _dim_name in axis.coords.items():
Expand Down
4 changes: 2 additions & 2 deletions tests/datasets/test_utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -3,7 +3,7 @@
import xarray as xr

from parcels._datasets import utils
from parcels._datasets.structured.generic import datasets_comodo
from parcels._datasets.structured.generic import datasets


@pytest.fixture
Expand All @@ -28,7 +28,7 @@ def nonzero_ds():
)


@pytest.mark.parametrize("ds", [pytest.param(v, id=k) for k, v in datasets_comodo.items()])
@pytest.mark.parametrize("ds", [pytest.param(v, id=k) for k, v in datasets.items()])
@pytest.mark.parametrize("except_for", [None, "coords"])
def test_replace_arrays_with_zeros(ds, except_for):
# make sure doesn't error with range of datasets
Expand Down
2 changes: 1 addition & 1 deletion tests/test_field.py
Original file line number Diff line number Diff line change
Expand Up @@ -7,7 +7,7 @@

from parcels import Field, UxGrid, VectorField, XGrid
from parcels._datasets.structured.generic import T as T_structured
from parcels._datasets.structured.generic import datasets_comodo as datasets_structured
from parcels._datasets.structured.generic import datasets as datasets_structured
from parcels._datasets.unstructured.generic import datasets as datasets_unstructured
from parcels.interpolators import (
UxConstantFaceConstantZC,
Expand Down
Loading
Loading