Fix z_l->lev rename for OM4 snapshot sources

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"],

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