Add function to add CF axis attr to Z axis

- This is a temporary workaround for variables that have Z axes without CF axis attr, which is required for xCDAT to map to for bounds

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,32 @@ 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 axes of the dataset.
+
+        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.
+        """
+        dim = xc.get_dim_keys(ds, axis="Z")
+
+        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 +480,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",
         }