Add function to add CF axis attr to Z axis if missing for downstream xCDAT operations (#865)

Author: tomvothecoder

e3sm_diags/driver/utils/dataset_xr.py

@@ -393,6 +393,11 @@ def _get_climo_dataset(self, season: str) -> xr.Dataset:
         filepath = self._get_climo_filepath(season)
         ds = self._open_climo_dataset(filepath)
 
+        # Add CF attributes to Z axes if they are missing.
+        # NOTE: This is a temporary workaround for xCDAT.
+        # Refer to https://github.com/xCDAT/xcdat/pull/708
+        ds = self._add_cf_attrs_to_z_axes(ds)
+
         if self.var in self.derived_vars_map:
             ds = self._get_dataset_with_derived_climo_var(ds)
         elif self.var in ds.data_vars.keys():
@@ -408,6 +413,35 @@ def _get_climo_dataset(self, season: str) -> xr.Dataset:
 
         return ds
 
+    def _add_cf_attrs_to_z_axes(self, ds: xr.Dataset) -> xr.Dataset:
+        """Add CF attributes to the Z axis of the dataset if the Z axis exists.
+
+        This method is a temporary solution to enable xCDAT to properly
+        retrieve bounds for Z axes that do not have CF attributes, which
+        is required for downstream regridding operations.
+
+        Parameters
+        ----------
+        ds : xr.Dataset
+            The dataset.
+
+        Returns
+        -------
+        xr.Dataset
+            The dataset with CF attributes added to the Z axes.
+        """
+        try:
+            dim = xc.get_dim_keys(ds, axis="Z")
+        except KeyError:
+            pass
+        else:
+            axis_attr = ds[dim].attrs.get("axis")
+
+            if axis_attr is None:
+                ds[dim].attrs["axis"] = "Z"
+
+        return ds
+
     def _open_climo_dataset(self, filepath: str) -> xr.Dataset:
         """Open a climatology dataset.
 
@@ -449,7 +483,7 @@ def _open_climo_dataset(self, filepath: str) -> xr.Dataset:
         args = {
             "paths": filepath,
             "decode_times": True,
-            "add_bounds": ["X", "Y"],
+            "add_bounds": ["X", "Y", "Z"],
             "coords": "minimal",
             "compat": "override",
         }