CDAT Migration: Refactor annual_cycle_zonal_mean set (#798)

* Refactor `annual_cycle_zonal_mean` set

* Address PR review comments

* Add lat lon regression testing

* Add debugging scripts

* Update `_open_climo_dataset()` to decode times as workaround to misaligned time coords
- Update `annual_cycle_zonal_mean_plot.py` to convert time coordinates to month integers

* Fix unit tests

* Remove old plotter

* Add script to debug decode_times=True and ncclimo file

* Update plotter time values to month integers

* Fix slow `.load()` and multiprocessing issue
- Due to incorrectly updating `keep_bnds` logic
- Add `_encode_time_coords()` to workaround cftime issue `ValueError: "months since" units only allowed for "360_day" calendar`

* Update `_encode_time_coords()` docstring

* Add AODVIS debug script

* update AODVIS obs datasets; regression test results

---------

Co-authored-by: Tom Vo <[email protected]>

e3sm_diags/driver/utils/dataset_xr.py

@@ -447,6 +447,18 @@ def _get_climo_dataset(self, season: str) -> xr.Dataset:
         if "slat" in ds.dims:
             ds = ds.drop_dims(["slat", "slon"])
 
+        all_vars = list(ds.data_vars.keys())
+        keep_bnds = [var for var in all_vars if "bnd" in var or "bounds" in var]
+        ds = ds[[self.var] + keep_bnds]
+
+        # NOTE: There seems to be an issue with `open_mfdataset()` and
+        # using the multiprocessing scheduler defined in e3sm_diags,
+        # resulting in timeouts and resource locking.
+        # To avoid this, we load the multi-file dataset into memory before
+        # performing downstream operations.
+        # Related GH issue: https://github.com/pydata/xarray/issues/3781
+        ds.load(scheduler="sync")
+
         return ds
 
     def _open_climo_dataset(self, filepath: str) -> xr.Dataset: