CDAT Migration Phase 2: Refactor core utilities and `lat_lon` set

[tomvothecoder]

Description

• Closes E3SM-Project/e3sm_diags#658

This PR refactors most core utilities for e3sm_diags to use Xarray/xCDAT since CDAT is deprecated at the end of the year. The lat_lon set is also refactored to use these core utilities.

How to Review

This is a large PR that involves many changes. I am interested in code review for areas of improvement, optimization, potential bugs, etc. The functionality itself is unchanged and working (refer to viewer outputs at the bottom of the description).

To help streamline the review process, there is section below for the main focus areas. If you are interested and have time, feel free to review other files and functions that are not listed.

I anticipate review will take the next few weeks (targeting end of Sep/2023). I'll try my best to respond to any comments, suggestions, or questions you might have. Thanks for the help!

Main Focus Areas

Each focus area includes the most important files and functions to review.

• Most refactored functions include a docstring with a "Notes" section at the bottom. The "Notes" section lists the old function it replaces (if you wanted to compare them). Here's an example (look at the bottom of the docstring):
• https://github.com/E3SM-Project/e3sm_diags/blob/3e0a0871736f89b9bb24968b25c58539f54ce975/e3sm_diags/driver/utils/regrid.py#L180-L200

(1) Dataset class

• driver/dataset_xr.py
  • Methods have been reordered to make the class more readable
  • Methods have been renamed, refactored, and updated with docstrings. I tried to group related time series and climatology functions together.
  • Most methods now return xr.Dataset objects, which are crucial because they contain bounds for downstream operations (e.g., generating weights). Variables in xr.Dataset objects should be operated on (regridding, regional subsetting, metrics calculation) until the last step, plotting (which will use the variables directly as xr.DataArray).

(2) Computation, Regridding, and I/O Utilities

• driver/utils/climo_xr.py
  • climo() function has the same logic as the older version but operates on Xarray objects
• driver/utils/regrid.py
  • align_grids_to_lower_res (formerly regrid_to_lower_res())
  • regrid_z_axis_to_plevs() (formerly convert_to_plevs())
  • _hybrid_to_plevs()
  • _pressure_to_plevs()
• driver/metrics/metrics.py
  • spatial_avg() (formerly mean())
  • std()
  • correlation()
  • rmse()
• driver/utils/io.py
  • Moved _write_vars_to_netcdf() and _get_output_dir() from general.py to this file

(3) Lat Lon Driver

• _driver/lat_lon_driver.py_
  • Moved run_diag() to the top level of the file
  • Most of the code is similar. I replaced function calls to use the new refactored functions, renamed things, and tried to refactor logic to make it more readable.
  • xr.Dataset objects are mainly used because they contain bounds. The test/reference variable gets updated in the xr.Dataset objects when applying land sea mask, regridding, etc. These objects are passed to _create_and_save_metrics().

(4) Lat Lon Plotter

• _plot/lat_lon_plot.py_
  • Refactored arguments to expect xr.DataArray type
  • Extract code into separate utility functions (added to plot/utils.py), which reduced the length of the file significantly (compare to plot/cartopy/deprecated_lat_lon_plot.py)

(5) Derived Variables

• derivations/acme.py
  • Split acme.py file into derivations.py, formulas.py, and utils.py
  • derivations.py -- stores derived variables dictionary.
  • derivations/formulas.py -- stores formulas for deriving variables, all functions have been updated to operate on xr.DataArray objects
  • utils.py -- a general file for storing utility functions (e.g,. unit conversion)
  • _apply_land_sea_mask() replaces the old mask_by() function (removed unnecessary logic and renamed for clarity)
  • _TODO: Some functions have not been refactored (e.g., determine_tau())_
• _derivations/default_regions_xr.py_
  • Updated structure of REGION_SPECS dict (formerly region_specs)
  • Updated structure of ARM_SITE_SPECS dict (formerly points_specs)

(6) Misc.

• e3sm_diags/driver/utils/general.py
  • Removed the mask_by() function because it was overly complex and select_region() only uses one part of the code.
• e3sm_diags/driver/utils/core_parameter.py
  • Updated self.backend="mpl" to self.backend=cartopy because "mpl" gets updated to "cartopy" in e3sm_diags/plot/__init__.py anyways

Checklist

• [x] My code follows the style guidelines of this project
• [x] I have performed a self-review of my own code
• [x] My changes generate no new warnings
• [ ] Any dependent changes have been merged and published in downstream modules

If applicable:

• [x] I have added tests that prove my fix is effective or that my feature works
• [x] New and existing unit tests pass with my changes (locally and CI/CD build)
• [x] I have commented my code, particularly in hard-to-understand areas
• [x] I have made corresponding changes to the documentation
• [ ] I have noted that this is a breaking change for a major release (fix or feature that would cause existing functionality to not work as expected)

Viewer Outputs

• This branch: https://portal.nersc.gov/cfs/e3sm/vo13/examples_658/ex1_modTS_vs_modTS_3years/viewer/lat_lon/index.html
• main branch: https://portal.nersc.gov/cfs/e3sm/vo13/examples/ex1_modTS_vs_modTS_3years/viewer/lat_lon/index.html

Future Items

• E3SM-Project/e3sm_diags#734

[tomvothecoder]

Hi @chengzhuzhang, I'm spending more time managing the effort and progress for this PR.

Here are a few notes:

1. This is a large and high priority task so we might need to pull another person in to help refactor other components if we want it merged within the next 1-2 months. xCDAT and other projects are currently taking up a lot of my bandwidth.
   • Workaround -- I refactor the core driver and utilities, and somebody else refactors the plotter and viewer
2. I'm sinking too much time on investigating the diffs between xCDAT's climatology vs. E3SM diags climatology
   • Workaround -- Refactor E3SM diags climatology to operate on xr.DataArray instead of TransientVariable. Open up a new issue to integrate xCDAT's climatology later, and attach my investigation Jupyter Notebook
3. I don't have a clear timeline on when Jason can work on incorporating a version of cdutil.reconstructPressureFromHybrid into cf-xarray.
   • Temporary Workaround -- Create a function in e3sm_diags that is a copy of this code and convert it to xarray directly

[chengzhuzhang]

Hi @chengzhuzhang, I'm spending more time managing the effort and progress for this PR.

Here are a few notes:

1. This is a large and high priority task so we might need to pull another person in to help refactor other components if we want it merged within the next 1-2 months. xCDAT and other projects are currently taking up a lot of my bandwidth.

   * **Workaround -- I refactor the core driver and utilities, and somebody else refactors the plotter and viewer**

2. I'm sinking too much time on investigating the diffs between xCDAT's climatology vs. E3SM diags climatology

   * **Workaround -- Refactor E3SM diags climatology to operate on xr.DataArray instead of TransientVariable. Open up a new issue to integrate xCDAT's climatology later, and attach my investigation Jupyter Notebook**

3. I don't have a clear timeline on when Jason can work on incorporating a version of [cdutil.reconstructPressureFromHybrid](https://github.com/CDAT/cdutil/blob/b823b69db46bb76536db7d435e72075fc3975c65/cdutil/vertical.py#L8) into `cf-xarray`.

   * **Temporary Workaround -- Create a function in e3sm_diags that is a copy of this code and convert it to xarray directly**

Hey @tomvothecoder thank you for the thoughtful suggestions. This is a high priority PR, since we have a crunching time line for the CDAT support. The workaround for climatology and and pressure level construction sounds good to me! Both @forsyth2 and I are familiar with the plotter and viewer code, so I think distribute the refactor work to either would make sense. We can talk about how to get start splitting the work.

[tomvothecoder]

Hey @chengzhuzhang and @forsyth2, this PR is ready for code review. The PR description includes all of the information needed to help make the review process smooth since it's a large PR.

I have a few items left including validating the metrics for closeness compared to main, fixing the unit tests, and making sure integration tests pass. After everything checks out, I will merge this PR into a cdat-refactor-fy23 development branch, which we will merge into main periodically until all sets are refactored.

[chengzhuzhang]

Hey @chengzhuzhang and @forsyth2, this PR is ready for code review. The PR description includes all of the information needed to help make the review process smooth since it's a large PR.

I have a few items left including validating the metrics for closeness compared to main, fixing the unit tests, and making sure integration tests pass. After everything checks out, I will merge this PR into a cdat-refactor-fy23 development branch, which we will merge into main periodically until all sets are refactored.

Super! Thank you for the notes. I will start review soon!