diff --git a/data/ocean_preprocessing/simulation_preprocessing/gfdl_om4.py b/data/ocean_preprocessing/simulation_preprocessing/gfdl_om4.py index 07193fad..f1d6c1c6 100644 --- a/data/ocean_preprocessing/simulation_preprocessing/gfdl_om4.py +++ b/data/ocean_preprocessing/simulation_preprocessing/gfdl_om4.py @@ -34,6 +34,24 @@ def convert_super_grid(ds_super_grid: xr.Dataset): return angle_h, lon_h, lat_h, lon_b, lat_b +def normalize_vertical_coords(ds: xr.Dataset) -> xr.Dataset: + """Rename raw OM4 vertical coordinates to the pipeline's expected names. + + In OM4 output, ``z_l`` is the depth of the cell centers (equivalent to + ``lev``) and ``z_i`` is the depth of the cell interfaces (equivalent to + ``ilev``, which holds one extra entry). The 5-daily snapshot sources expose + ``z_l`` without a matching ``z_i``, so each coordinate is renamed + independently rather than gating the center rename (``z_l`` -> ``lev``) on + the interface coordinate being present. + """ + vertical_rename = { + raw: new + for raw, new in (("z_l", "lev"), ("z_i", "ilev")) + if raw in ds.coords or raw in ds.dims + } + return ds.rename(vertical_rename) if vertical_rename else ds + + def om4_preprocessing( zarr_data_path, nc_grid_path, nc_mosaic_path, fs=fsspec, backend_kwargs=None ): @@ -42,8 +60,9 @@ def om4_preprocessing( zarr_data_path, engine="zarr", chunks={}, backend_kwargs=backend_kwargs ) - if "z_i" in ds.coords: - ds = ds.rename({"z_i": "ilev", "z_l": "lev"}) + ds = normalize_vertical_coords(ds) + + if "ilev" in ds.coords: dz = xr.DataArray( ds.ilev.diff("ilev").values, dims=["lev"], diff --git a/data/tests/test_simulation_preprocessing.py b/data/tests/test_simulation_preprocessing.py new file mode 100644 index 00000000..a7e0fe00 --- /dev/null +++ b/data/tests/test_simulation_preprocessing.py @@ -0,0 +1,39 @@ +# SPDX-FileCopyrightText: 2026 Ocean Emulator Authors +# +# SPDX-License-Identifier: Apache-2.0 + +import numpy as np +import xarray as xr +from ocean_preprocessing.simulation_preprocessing.gfdl_om4 import ( + normalize_vertical_coords, +) + + +def _ds_with_vertical(names): + """Build a tiny dataset whose vertical dimension coordinate(s) are ``names``.""" + coords = {name: (name, np.arange(3)) for name in names} + return xr.Dataset(coords=coords) + + +def test_normalize_renames_z_l_without_z_i(): + # Snapshot sources expose z_l (cell centers) but no z_i; z_l must still + # become lev so downstream rechunking on "lev" succeeds. + ds = _ds_with_vertical(["z_l"]) + out = normalize_vertical_coords(ds) + assert "lev" in out.coords + assert "z_l" not in out.coords + + +def test_normalize_renames_both_when_present(): + # Averaged sources carry both interfaces and centers. + ds = _ds_with_vertical(["z_l", "z_i"]) + out = normalize_vertical_coords(ds) + assert {"lev", "ilev"} <= set(out.coords) + assert "z_l" not in out.coords and "z_i" not in out.coords + + +def test_normalize_is_noop_without_raw_vertical_coords(): + # Already-normalized data passes through untouched. + ds = _ds_with_vertical(["lev"]) + out = normalize_vertical_coords(ds) + assert "lev" in out.coords