diff --git a/src/parcels/convert.py b/src/parcels/convert.py index e8f2b84b9..e324dcd11 100644 --- a/src/parcels/convert.py +++ b/src/parcels/convert.py @@ -13,6 +13,7 @@ from __future__ import annotations import typing +import warnings import numpy as np import xarray as xr @@ -23,9 +24,21 @@ if typing.TYPE_CHECKING: import uxarray as ux -_NEMO_EXPECTED_COORDS = ["glamf", "gphif"] - -_NEMO_DIMENSION_COORD_NAMES = ["x", "y", "time", "x", "x_center", "y", "y_center", "depth", "glamf", "gphif"] +_NEMO_EXPECTED_COORDS = ["glamf", "gphif", "depthw"] + +_NEMO_DIMENSION_COORD_NAMES = [ + "x", + "y", + "time", + "x", + "x_center", + "y", + "y_center", + "depth", + "depth_center", + "glamf", + "gphif", +] _NEMO_AXIS_VARNAMES = { "x": "X", @@ -33,6 +46,7 @@ "y": "Y", "y_center": "Y", "depth": "Z", + "depth_center": "Z", "time": "T", } @@ -94,16 +108,18 @@ def _pick_expected_coords(coords: xr.Dataset, expected_coord_names: list[str]) - def _maybe_bring_other_depths_to_depth(ds): - if "depth" in ds.coords: - for var in ds.data_vars: - for old_depth in ["depthu", "depthv", "deptht", "depthw"]: - if old_depth in ds[var].dims: - ds[var] = ds[var].assign_coords(**{old_depth: ds["depth"].values}).rename({old_depth: "depth"}) - return ds - + for var in ds.data_vars: + for old_depth, target in [ + ("depthu", "depth_center"), + ("depthv", "depth_center"), + ("deptht", "depth_center"), + ("depthw", "depth"), + ]: + if old_depth in ds[var].dims: + ds[var] = ds[var].rename({old_depth: target}) -def _maybe_create_depth_dim(ds): if "depth" not in ds.dims: + warnings.warn("No depth dimension found in your dataset. Assuming no depth (i.e., surface data).", stacklevel=1) ds = ds.expand_dims({"depth": [0]}) ds["depth"] = xr.DataArray([0], dims=["depth"]) return ds @@ -285,17 +301,18 @@ def nemo_to_sgrid(*, fields: dict[str, xr.Dataset | xr.DataArray], coords: xr.Da if coords.sizes["time"] != 1: raise ValueError("Time dimension in coords must be length 1 (i.e., no time-varying grid).") coords = coords.isel(time=0).drop("time") - if len(coords.dims) == 3: + + if ( + len(coords.dims) == 3 + ): #! This should really be looking at the dimensionality of the lons and lats arrays. Currently having 2D lon lat and 1D depth triggers this `if` clause for dim, len_ in coords.sizes.items(): if len_ == 1: # TODO: log statement about selecting along z dim of 1 coords = coords.isel({dim: 0}) - if len(coords.dims) != 2: - raise ValueError("Expected coordsinates to be 2 dimensional") - + # if len(coords.dims) != 2: #! This should really be looking at the dimensionality of the lons and lats arrays. Currently having 2D lon lat and 1D depth triggers this `if` clause + # raise ValueError("Expected coordinates to be 2 dimensional") ds = xr.merge(list(fields.values()) + [coords]) ds = _maybe_rename_variables(ds, _NEMO_VARNAMES_MAPPING) - ds = _maybe_create_depth_dim(ds) ds = _maybe_bring_other_depths_to_depth(ds) ds = _drop_unused_dimensions_and_coords(ds, _NEMO_DIMENSION_COORD_NAMES) ds = _assign_dims_as_coords(ds, _NEMO_DIMENSION_COORD_NAMES) @@ -331,7 +348,7 @@ def nemo_to_sgrid(*, fields: dict[str, xr.Dataset | xr.DataArray], coords: xr.Da sgrid.DimDimPadding("x_center", "x", sgrid.Padding.LOW), sgrid.DimDimPadding("y_center", "y", sgrid.Padding.LOW), ), - vertical_dimensions=(sgrid.DimDimPadding("z_center", "depth", sgrid.Padding.HIGH),), + vertical_dimensions=(sgrid.DimDimPadding("depth_center", "depth", sgrid.Padding.HIGH),), ).to_attrs(), )