Skip to content

Commit

Permalink
[Refactor]: CDAT Migration - Update dependencies and remove Dataset._…
Browse files Browse the repository at this point in the history
…add_cf_attrs_to_z_axes() (#891)
  • Loading branch information
tomvothecoder authored Nov 8, 2024
1 parent 3406bf2 commit cb1cfb2
Show file tree
Hide file tree
Showing 4 changed files with 30 additions and 42 deletions.
Original file line number Diff line number Diff line change
@@ -0,0 +1,22 @@
import os
from e3sm_diags.parameter.core_parameter import CoreParameter
from e3sm_diags.run import runner

param = CoreParameter()

# param.reference_data_path = '/global/cfs/cdirs/e3sm/diagnostics/observations/Atm/climatology'
# param.test_data_path = '/global/cfs/cdirs/e3sm/zhang40/e3sm_diags_for_EAMxx/data/Cess'
param.reference_data_path = (
"/global/cfs/cdirs/e3sm/diagnostics/observations/Atm/climatology"
)
param.test_data_path = "/global/cfs/cdirs/e3sm/chengzhu/eamxx/post/data/rgr"
param.test_name = "eamxx_decadal"
param.seasons = ["ANN"]
# param.save_netcdf = True

prefix = "/global/cfs/cdirs/e3sm/www/cdat-migration-fy24"
param.results_dir = os.path.join(prefix, "889-z-axis-bnds")

runner.sets_to_run = ["zonal_mean_2d"]

runner.run_diags([param])
8 changes: 4 additions & 4 deletions conda-env/ci.yml
Original file line number Diff line number Diff line change
Expand Up @@ -21,11 +21,11 @@ dependencies:
- matplotlib-base
- netcdf4
- output_viewer >=1.3.0
- numpy >=1.23.0
- numpy >=2.0.0,<3.0.0
- shapely >=2.0.0,<3.0.0
- xarray >=2023.02.0
- xcdat >=0.6.0
- xesmf >=0.7.0
- xarray >=2024.03.0
- xcdat >=0.7.3,<1.0.0
- xesmf >=0.8.7
- xskillscore >=0.0.20
# Testing
# ==================
Expand Down
8 changes: 4 additions & 4 deletions conda-env/dev.yml
Original file line number Diff line number Diff line change
Expand Up @@ -19,11 +19,11 @@ dependencies:
- matplotlib-base
- netcdf4
- output_viewer >=1.3.0
- numpy >=1.23.0
- numpy >=2.0.0,<3.0.0
- shapely >=2.0.0,<3.0.0
- xarray >=2023.02.0
- xcdat >=0.6.0
- xesmf >=0.7.0
- xarray >=2024.03.0
- xcdat >=0.7.3,<1.0.0
- xesmf >=0.8.7
- xskillscore >=0.0.20
# Testing
# =======================
Expand Down
34 changes: 0 additions & 34 deletions e3sm_diags/driver/utils/dataset_xr.py
Original file line number Diff line number Diff line change
Expand Up @@ -419,11 +419,6 @@ 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():
Expand All @@ -439,35 +434,6 @@ 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.
Expand Down

0 comments on commit cb1cfb2

Please sign in to comment.