diff --git a/src/parcels/convert.py b/src/parcels/convert.py index 0fa8017ae..3ef433f58 100644 --- a/src/parcels/convert.py +++ b/src/parcels/convert.py @@ -193,6 +193,12 @@ def _rename_nemo_depth_dims(ds: xr.Dataset) -> xr.Dataset: for old_depth, target in _NEMO_DEPTH_DIM_RENAMES.items(): if old_depth in ds[var].dims: ds[var] = ds[var].rename({old_depth: target}) + if "depth_center" in ds.dims and "depth" not in ds.dims: + raise ValueError( + "The fields have a vertical dimension (depthu, depthv or deptht), but 'depthw' is missing. " + "The depth levels of the W points are needed for the vertical grid. Either pass the W field " + "(e.g., fields={..., 'W': ds_w['wo']}) or add 'depthw' from the W files to coords." + ) return ds @@ -223,7 +229,7 @@ def _drop_unused_dimensions_and_coords(ds: xr.Dataset, dimension_and_coord_names for dim in ds.dims: if dim not in dimension_and_coord_names: dim = cast(str, dim) - ds = ds.drop_dims(dim, errors="ignore") + ds = ds.isel({dim: 0}) for coord in ds.coords: coord = cast(str, coord) if coord not in dimension_and_coord_names: @@ -238,13 +244,6 @@ def _set_coords(ds: xr.Dataset, dimension_names): return ds -def _maybe_remove_depth_from_lonlat(ds): - for coord in ["glamf", "gphif"]: - if coord in ds.coords and "depth" in ds[coord].dims: - ds[coord] = ds[coord].squeeze("depth", drop=True) - return ds - - def _set_axis_attrs(ds: xr.Dataset, dim_axis: dict[str, XgcmAxisDirection]): for dim, axis in dim_axis.items(): if dim in ds: @@ -367,34 +366,17 @@ def nemo_to_sgrid(*, fields: dict[str, xr.Dataset | xr.DataArray], coords: xr.Da raise ValueError("Time dimension in coords must be length 1 (i.e., no time-varying grid).") coords = coords.isel(time=0).drop_vars("time") - 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: #! 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(field_das.values()) + [coords]) + ds = xr.merge(list(field_das.values()) + [coords], join="exact", compat="override") ds.attrs.clear() # Clear global attributes from the merging ds = _maybe_rename_variables(ds, _NEMO.varnames_mapping) ds = _rename_nemo_depth_dims(ds) - if "depth_center" in ds.dims and "depth" not in ds.dims: - raise ValueError( - "The fields have a vertical dimension (depthu, depthv or deptht), but 'depthw' is missing. " - "The depth levels of the W points are needed for the vertical grid. Either pass the W field " - "(e.g., fields={..., 'W': ds_w['wo']}) or add 'depthw' from the W files to coords." - ) - ds = _maybe_expand_depth_dimension(ds) ds = _drop_unused_dimensions_and_coords(ds, _NEMO_DIMENSION_COORD_NAMES) ds = _assign_dims_as_coords(ds, _NEMO_DIMENSION_COORD_NAMES) ds = _set_coords(ds, _NEMO_DIMENSION_COORD_NAMES) - ds = _maybe_remove_depth_from_lonlat(ds) ds = _set_axis_attrs(ds, _NEMO.axis_varnames) - expected_axes = set("XYZT") # TODO: Update after we have support for 2D spatial fields + expected_axes = set("XY") if missing_axes := (expected_axes - set(ds.cf.axes)): raise ValueError( f"Dataset missing CF compliant metadata for axes " diff --git a/tests/test_convert.py b/tests/test_convert.py index f178597d3..ac7f9ceae 100644 --- a/tests/test_convert.py +++ b/tests/test_convert.py @@ -92,6 +92,17 @@ def test_nemo_to_sgrid_with_depth(U, V, depth, coords): # noqa: N803 parcels.FieldSet.from_sgrid_conventions(ds) +def test_convert_nemo_to_sgrid_without_time(): + ds = parcels.tutorial.open_dataset("NemoCurvilinear_data_zonal/mesh_mask") + ds_fset = parcels.convert.nemo_to_sgrid(fields={"fmask": ds.fmask}, coords=ds) + fieldset = parcels.FieldSet.from_sgrid_conventions(ds_fset) + + assert "mockT" in fieldset.fmask.data.dims + assert "mockZ" in fieldset.fmask.data.dims + # sampling time-constant field should yield the same value + assert fieldset.fmask[10, 0, 0, 0] == fieldset.fmask[1, 0, 0, 0] + + def test_convert_nemo_offsets(): U = parcels.tutorial.open_dataset("NemoCurvilinear_data_zonal/U") V = parcels.tutorial.open_dataset("NemoCurvilinear_data_zonal/V") @@ -262,13 +273,22 @@ def test_convert_sets_only_conventions_global_attr(): assert ds_fset.attrs == {"Conventions": "SGRID"} -def test_convert_nemo_surface_adds_depth_after_time(): +def test_convert_nemo_surface_adds_mockZ(): U = parcels.tutorial.open_dataset("NemoCurvilinear_data_zonal/U") V = parcels.tutorial.open_dataset("NemoCurvilinear_data_zonal/V") coords = parcels.tutorial.open_dataset("NemoCurvilinear_data_zonal/mesh_mask") ds = convert.nemo_to_sgrid(fields=dict(U=U, V=V), coords=coords) - assert ds["U"].dims[:2] == ("time", "depth") - assert ds["V"].dims[:2] == ("time", "depth") + + fieldset = FieldSet.from_sgrid_conventions(ds) + assert fieldset.U.data.dims[:2] == ("time", "mockZ") + assert fieldset.V.data.dims[:2] == ("time", "mockZ") + + +def test_convert_nemo_edges(): + ds = parcels.tutorial.open_dataset("NemoNorthSeaORCA025-N006_data/mesh_mask") + ds = convert.nemo_to_sgrid(fields={"e1u": ds["e1u"]}, coords=ds) + fieldset = FieldSet.from_sgrid_conventions(ds) + assert fieldset.e1u.data.dims[:2] == ("mockT", "mockZ") def test_convert_dataset_field_without_matching_variable():