added Z axis to cmor axis

cordex/cmor/cmor.py

@@ -144,7 +144,7 @@ def _get_time_axis_name(time_cell_method):
     return time_axis_names.get(time_cell_method, "time")
 
 
-def _define_axes(ds, table_id):
+def _define_grid(ds, table_id):
     if "domain_id" in ds.attrs:
         grid = cordex_domain(ds.attrs["domain_id"], bounds=True)
         lon_vertices = grid.lon_vertices.to_numpy()
@@ -164,6 +164,7 @@ def _define_axes(ds, table_id):
         coord_vals=ds.rlon.to_numpy(),
         units=ds.rlon.units,
     )
+
     cmorGrid = cmor.grid(
         [cmorLat, cmorLon],
         latitude=ds.lat.to_numpy(),
@@ -220,29 +221,51 @@ def _define_time(ds, table_id, time_cell_method=None):
     )
 
 
-def _define_grid(ds, table_ids, time_cell_method="point"):
-    cmorGrid = _define_axes(ds, table_ids["grid"])
+def _define_axis(ds, table_ids, time_cell_method="point"):
+    cmorGrid = _define_grid(ds, table_ids["grid"])
 
     if "time" in ds:
         cmorTime = _define_time(ds, table_ids["mip"], time_cell_method)
     else:
         cmorTime = None
 
-    return cmorTime, cmorGrid
+    # add z axis if required
+    if "Z" in ds.cf.dims:
+        z = ds.cf["Z"]
+        bounds = ds.cf.add_bounds("Z").cf.get_bounds("Z").to_numpy()
+        cmorZ = cmor.axis(
+            table_entry=z.name,
+            coord_vals=z.to_numpy(),
+            units=z.attrs.get("units"),
+            cell_bounds=bounds,
+        )
+    else:
+        cmorZ = None
+
+    return cmorTime, cmorZ, cmorGrid
+
 
+def _cmor_write(da, table_id, cmorTime, cmorZ, cmorGrid, file_name=True):
+    """write to netcdf via cmor python API"""
 
-def _cmor_write(da, table_id, cmorTime, cmorGrid, file_name=True):
     cmor.set_table(table_id)
-    if cmorTime is None:
-        coords = [cmorGrid]
-    else:
-        coords = [cmorTime, cmorGrid]
+
+    # create coordinate ids
+    coords = []
+    if cmorTime:
+        coords.append(cmorTime)
+    if cmorZ:
+        coords.append(cmorZ)
+    coords.append(cmorGrid)
+
     cmor_var = cmor.variable(da.name, da.units, coords)
+
     if "time" in da.coords:
         ntimes_passed = da.time.size
     else:
         ntimes_passed = None
     cmor.write(cmor_var, da.to_numpy(), ntimes_passed=ntimes_passed)
+
     return cmor.close(cmor_var, file_name=file_name)
 
 
@@ -444,9 +467,11 @@ def cmorize_cmor(
         dataset_table_json, cmor_table_json, grids_table=grids_table, inpath=inpath
     )
 
-    cmorTime, cmorGrid = _define_grid(ds, table_ids, time_cell_method=time_cell_method)
+    cmorTime, cmorZ, cmorGrid = _define_axis(
+        ds, table_ids, time_cell_method=time_cell_method
+    )
 
-    return _cmor_write(ds[out_name], table_ids["mip"], cmorTime, cmorGrid)
+    return _cmor_write(ds[out_name], table_ids["mip"], cmorTime, cmorZ, cmorGrid)
 
 
 def prepare_variable(
@@ -470,6 +495,8 @@ def prepare_variable(
         mapping_table = {}
 
     ds = ds.copy(deep=False)
+    # use cf_xarray to guess coordinate meta data
+    ds = ds.cf.guess_coord_axis(verbose=True)
 
     if isinstance(cmor_table, str):
         cmor_table = _read_table(cmor_table)
@@ -485,13 +512,12 @@ def prepare_variable(
     # ds = xr.decode_cf(ds, decode_coords="all")
 
     # no mapping table provided, we assume datasets has already correct out_names and units.
-    if out_name not in mapping_table:
-        try:
-            var_ds = ds[[out_name]]
-        except Exception:
-            raise Exception(
-                f"Could not find {out_name} in dataset. Please make sure, variable names and units have CF standard or pass a mapping table."
-            )
+    if out_name in ds.data_vars:
+        var_ds = ds[[out_name]]
+    elif mapping_table and out_name not in mapping_table:
+        raise Exception(
+            f"Could not find {out_name} in dataset. Please make sure, variable names and units have CF standard or pass a mapping table."
+        )
     else:
         varname = mapping_table[out_name]["varname"]
         # cf_name = varinfo["cf_name"]
@@ -565,6 +591,13 @@ def cmorize_variable(
 ):
     """Cmorizes a variable.
 
+    This functions call the python cmor API and creates a cmorized NetCDF file from the input
+    dataset. Coordinates of the input dataset should be understandable by ``cf_xarray`` if they
+    follow basic CF conventions. If a vertical coordinate is available, it should have an ``axis="Z"``
+    attribute so it can be understood by ``cf_xarray`` and it should be named after the unique key of the
+    coordinate in the cmor grids table (not the out_name), e.g., name it ``sdepth`` if you
+    need a soil layer coordinate instead of ``depth``.
+
     Parameters
     ----------
     ds : xr.Dataset
@@ -620,6 +653,27 @@ def cmorize_variable(
     filename
         Filepath to cmorized file.
 
+    Example
+    -------
+
+    To cmorize a dataset, you can use ,e.g.,::
+
+        import cordex as cx
+        from cordex.cmor import cmorize_variable
+        from cordex.tables import cordex_cmor_table
+
+        ds = cx.cordex_domain("EUR-11", dummy="topo").rename(topo="orog")
+        dataset_table = cordex_cmor_table(f"CORDEX-CMIP6_remo_example")
+
+        filename = cmorize_variable(
+            ds,
+            "orog",
+            cmor_table=cordex_cmor_table("CORDEX-CMIP6_fx"),
+            dataset_table=dataset_table,
+            replace_coords=True,
+            allow_units_convert=True,
+        )
+
     """
     ds = ds.copy()
 

cordex/domain.py

@@ -67,7 +67,6 @@ def domain_names(table_name=None):
 def domain(
     domain_id,
     dummy=False,
-    add_vertices=False,
     tables=None,
     attrs=None,
     mapping_name=None,

docs/whats_new.rst

@@ -6,14 +6,20 @@ What's new
 v0.8.0 (UNRELEASED)
 -------------------
 
-Unreleased.
+This version introduces :py:meth:`domain` which should tighten the API call in the future.
+The cmor module was also updated to make more usage of ``cf_xarray``.
 
 .. note::
    This version introduces :py:meth:`domain` which is equivalent to the former
-   `cordex_domain` function. Both API calls are available, however, future
+   ``cordex_domain`` function. Both API calls are available, however, future
    improvements might only be done on the new ``domain`` API and it is
    recommended to use :py:meth:`domain` in the future.
 
+New Features
+~~~~~~~~~~~~
+
+- Cmor function :py:meth:`cmor.cmorize_variable` can now handle 3D fields if a Z-axis is available (:pull:`230`).
+
 Internal Changes
 ~~~~~~~~~~~~~~~~
 

pyproject.toml

@@ -57,6 +57,7 @@ exclude = [
 # E402: module level import not at top of file
 # E501: line too long - let black worry about that
 # E731: do not assign a lambda expression, use a def
+[lint]
 ignore = [
     "E402",
     "E501",

tests/test_cmor.py

@@ -4,6 +4,7 @@
 import cftime
 import cftime as cfdt
 import numpy as np
+import pandas as pd
 import pytest
 import xarray as xr
 
@@ -21,6 +22,20 @@
 cmor.set_options(table_prefix=table_prefix)
 
 
+def create_sdepth_ds():
+    ds = cx.domain("EUR-11", dummy="topo")
+    ds = ds.drop("topo").assign(
+        tsl=ds.topo.expand_dims(
+            time=pd.date_range("2000-01-01T12:00:00", periods=3, freq="D"),
+            sdepth=[0.05, 0.10, 0.2],
+        )
+    )
+    ds.tsl.attrs["units"] = "K"
+    ds.sdepth.attrs["axis"] = "Z"
+    ds.sdepth.attrs["units"] = "m"
+    return ds
+
+
 def test_cfdt():
     assert cmor.to_cftime(dt.datetime(2000, 1, 1, 1)) == cfdt.datetime(2000, 1, 1, 1)
     assert cmor.to_cftime(dt.date(2000, 1, 1)) == cfdt.datetime(2000, 1, 1)
@@ -203,6 +218,12 @@ def test_cmorizer_mon():
     assert "tas" in output
 
 
+def test_cmorizer_mon_sdepth():
+    ds = create_sdepth_ds()
+    filename = run_cmorizer(ds, "tsl", "EUR-11", "day")
+    return filename
+
+
 @pytest.mark.parametrize("table_id, tdim", [("day", 3), ("1hr", 49)])
 def test_cmorizer_subdaily(table_id, tdim):
     ds = cx.tutorial.open_dataset("remo_EUR-11_TEMP2_1hr")