From 18d75a8aaea152f434a93619af8233fb8fe4780a Mon Sep 17 00:00:00 2001 From: Tom Vo Date: Wed, 9 Oct 2024 09:49:13 -0700 Subject: [PATCH] Add function to add CF axis attr to Z axis if missing for downstream xCDAT operations (#865) --- e3sm_diags/driver/utils/dataset_xr.py | 36 ++++++++++++++++++++++++++- 1 file changed, 35 insertions(+), 1 deletion(-) diff --git a/e3sm_diags/driver/utils/dataset_xr.py b/e3sm_diags/driver/utils/dataset_xr.py index d94c8df96..32a4b6f63 100644 --- a/e3sm_diags/driver/utils/dataset_xr.py +++ b/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", }