Merge PR #266 (Fixed online regridding code)

This merge brings PR #266 (Fixed online regridding code (file_regrid.py) for restart-to-restart file regridding only, by @yantosca) into the GCPy 1.4.0 development stream. This PR fixes several (but not all) issues with the file_regrid.py routine, which uses xESMF and online regridding weights. Regridding between GCClassic and GCHP checkpoint files is OK; there may still be some issues with cubed-sphere to cubed-sphere regridding. We will address those in a separate PR. Signed-off-by: Bob Yantosca <[email protected]>
geoschem · Oct 11, 2023 · 7949550 · 7949550
2 parents b05ea76 + 922f233
commit 7949550
Show file tree

Hide file tree

Showing 10 changed files with 1,594 additions and 543 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -32,6 +32,8 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 - Added statement `from dask.array import Array as DaskArray` in `gcpy plot.py`
 - Added SLURM run script `gcpy/benchmark/benchmark_slurm.sh`
 - Added `gcpy/gcpy_plot_style` style sheet for title and label default settings
+- Added new cubed-sphere grid inquiry functions to `gcpy/cstools.py`
+- Added functions `get_ilev_coord` and `get_lev_coord` to `gcpy/grid.py`
 
 ### Changed
 - Simplified the Github issues templates into two options: `new-feature-or-discussion.md` and `question-issue.md`
@@ -61,6 +63,9 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 - Now allow `plot_val` to be of type `dask.array.Array` in `plot.py` routines `six_plot` and `single_panel`
 - Now add `if` statements to turn of `Parallel()` commands when `n_jobs==1`.
 - Do not hardwire fontsize in `gcpy/plot.py`; get defaults from `gcpy_plot_style`
+- Updated and cleaned up code in `gcpy/regrid.py`
+- Example scripts`plot_single_level` and `plot_comparisons` can now accept command-line arguments
+- Example scripts `plot_single_level.py`, `plot_comparisons.py`, `compare_diags.py` now handle GCHP restart files properly
 
 ### Fixed
 - Generalized test for GCHP or GCClassic restart file in `regrid_restart_file.py`
@@ -69,6 +74,8 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 - Prevent plot panels from overlapping in six-panel plots
 - Prevent colorbar tick labels from overlapping in dynamic-range ratio plots
 - Updated `seaborn` plot style names to conform to the latest matplotlib
+- Set `lev:positive` and/or `ilev:positive` properly in `regrid_restart_file.py` and `file_regrid.py`
+- Prevent overwriting of `lev` coord in `file_regrid.py` at netCDF write time
 
 ### Removed
 - Removed `gchp_is_pre_13_1` arguments & code from benchmarking routines
@@ -366,16 +373,3 @@ This is the first labeled version of GCPy. The primary functionality of GCPy is
 - Support for plotting benchmark output for both GEOS-Chem Classic (lat/lon data) and GCHP (cubed-sphere data).
 
 The first official release version of GCPy, v1.0.0, will correspond with the release of GEOS-Chem 13.0.0.
-
-
-## [Unreleased]
-
-### Added
-
-### Changed
-
-### Deprecated
-
-### Fixed
-
-### Removed
diff --git a/gcpy/cstools.py b/gcpy/cstools.py
@@ -37,15 +37,10 @@
     import pyproj
     import shapely.ops
     import shapely.geometry
-# HOTFIX: Don't raise an ImportError until we can add pyproj to
-# the AWS cloud container.  This is causing errors that prevent
-# benchmark plotting jobs from finishing. -- Bob Y. (01 Aug 2023)
-#except ImportError as exc:
-#    raise ImportError(
-#        "gcpy.cstools needs packages 'pyproj' and 'shapely'!"
-#    ) from exc
-except:
-    print("pyproj is not available")
+except ImportError as exc:
+    raise ImportError(
+        "gcpy.cstools needs packages 'pyproj' and 'shapely'!"
+    ) from exc
 
 # Constants
 RAD_TO_DEG = 180.0 / np.pi
@@ -74,8 +69,8 @@ def extract_grid(
     if not is_cubed_sphere(data):
         return None
 
-    n_cs = data["Xdim"].shape[-1]
-    return gen_grid(n_cs)
+    cs_res = get_cubed_sphere_res(data)
+    return gen_grid(cs_res)
 
 
 def read_gridspec(gs_obj):
@@ -111,21 +106,23 @@ def read_gridspec(gs_obj):
         area[i,...] = tile.area[...]
 
     data = xr.Dataset(
-        data_vars=dict(
-            area=(['nf','Ydim','Xdim'],area),
-            lon=(['nf','Ydim','Xdim'],lon),
-            lat=(['nf','Ydim','Xdim'],lat),
-            lon_b=(['nf','Ydim_b','Xdim_b'],lon_b),
-            lat_b=(['nf','Ydim_b','Xdim_b'],lat_b),
-        ),
-        coords=dict(
-            nf=(['nf'],list(range(6))),
-            Ydim=(['Ydim'],list(range(n_cs))),
-            Xdim=(['Xdim'],list(range(n_cs))),
-            Ydim_b=(['Ydim_b'],list(range(n_cs+1))),
-            Xdim_b=(['Xdim_b'],list(range(n_cs+1))),
-        ),
-        attrs=dict(description=f'c{n_cs:d} grid data'),
+        data_vars={
+            "area": (['nf','Ydim','Xdim'], area),
+            "lon": (['nf','Ydim','Xdim'], lon),
+            "lat": (['nf','Ydim','Xdim'], lat),
+            "lon_b": (['nf','Ydim_b','Xdim_b'], lon_b),
+            "lat_b": (['nf','Ydim_b','Xdim_b'], lat_b),
+        },
+        coords={
+            "nf": (['nf'], list(range(6))),
+            "Ydim": (['Ydim'], list(range(n_cs))),
+            "Xdim": (['Xdim'], list(range(n_cs))),
+            "Ydim_b": (['Ydim_b'], list(range(n_cs+1))),
+            "Xdim_b": (['Xdim_b'], list(range(n_cs+1))),
+        },
+        attrs={
+            "description": f"c{n_cs:d} grid data"
+        },
     )
     return data
 
@@ -348,7 +345,7 @@ def gen_grid(
         where each value has an extra face dimension of length 6.
     """
     if stretch_factor is not None:
-        cs_temp, ignore = gcpy.make_grid_SG(
+        cs_temp, _ = gcpy.make_grid_SG(
             n_cs,
             stretch_factor,
             target_lon,
@@ -680,7 +677,7 @@ def is_cubed_sphere(
 ):
     """
     Given an xarray Dataset or DataArray object, determines if the
-    data is placed on a cubed-sphere grid
+    data is placed on a cubed-sphere grid.
 
     Args:
     -----
@@ -692,10 +689,123 @@ def is_cubed_sphere(
     is_gchp : bool
         Returns True if data is placed on a cubed-sphere grid,
         and False otherwise.
+
+    Remarks:
+    --------
+    A cubed-sphere data file has one of the following attributes
+    (1) A dimension named "nf" (GCHP/GEOS diagnostic files)
+    (2) The lat/lon ratio is exactly 6 (GCHP/GEOS checkpoints)
+    """
+    gcpy.util.verify_variable_type(data, (xr.DataArray, xr.Dataset))
+
+    if is_cubed_sphere_diag_grid(data):
+        return True
+    if is_cubed_sphere_rst_grid(data):
+        return True
+    return False
+
+
+def is_cubed_sphere_diag_grid(data):
+    """
+    Determines if a cubed-sphere grid has History file dimensions.
+    (i.e. a dimension named "nf", aka number of grid faces).
+
+    Args:
+    -----
+    data : xarray.DataArray or xarray.Dataset
+        The input data.
+
+    Returns:
+    --------
+    True if the grid has History diagnostic dimensions,
+    False otherwise.
+    """
+    if "nf" in data.dims:
+        return True
+    return False
+
+
+def is_cubed_sphere_rst_grid(data):
+    """
+    Determines if a cubed-sphere grid has restart file dimensions.
+    (i.e. lat and lon, with lat = lon*6).
+
+    Args:
+    -----
+    data : xarray.DataArray or xarray.Dataset
+        The input data.
+
+    Returns:
+    --------
+    True if the grid has restart dimensions, False otherwise.
     """
     gcpy.util.verify_variable_type(data, (xr.DataArray, xr.Dataset))
 
-    if "nf" in data.dims:   # nf = number of cubed-sphere faces
+    # TODO: Rethink this if we ever end up changing the GC-Classic
+    # restart variables to start with SPC, or if we ever rename the
+    # internal state variables in GCHP. A more robust back-up check
+    # could be to see if all the lats and lons are integer, since
+    # that will be the case with the GCHP restart file format.
+    if "lat" in data.dims:
+        if data.dims["lat"] == data.dims["lon"] * 6:
+            return True
+    if "SPC_" in data.data_vars.keys():
         return True
+    return False
+
 
+def get_cubed_sphere_res(data):
+    """
+    Given a Dataset or DataArray object, returns the number of
+    grid cells along one side of the cubed-sphere grid face
+    (e.g. 24 for grid resolution C24, which has 24x25 grid cells
+    per face).
+
+    Args:
+    -----
+    data : xarray.DataArray or xarray.Dataset
+        The input data.
+
+    Returns:
+    --------
+    cs_res : int
+        The cubed-sphere resolution.  Will return 0 if the data
+        is not placed on a cubed-sphere grid.
+    """
+    gcpy.util.verify_variable_type(data, (xr.DataArray, xr.Dataset))
+
+    if not is_cubed_sphere(data):
+        return 0
+    if is_cubed_sphere_rst_grid(data):
+        return data.dims["lon"]
+    return data.dims["Xdim"]
+
+
+def is_gchp_lev_positive_down(data):
+    """
+    Determines if GCHP data is arranged vertically from the top of the
+    atmosphere downwards or from the surface upwards, according to:
+
+    (1) Checkpoint files:     lev:positive="down
+    (2) Emissions collection: lev:positive="down"
+    (3) Other collections     lev:positive="up"
+
+    Args:
+    -----
+    data : xarray.DataArray or xarray.Dataset
+       The input data
+
+    Returns:
+    --------
+    True if the data is arranged from top-of-atm downwards.
+    False if the data is arranged from the surface upwards.
+    """
+    gcpy.util.verify_variable_type(data, (xr.DataArray, xr.Dataset))
+
+    if is_cubed_sphere_rst_grid(data):
+        return True
+    if is_cubed_sphere_diag_grid(data):
+        emis_vars = [var for var in data.data_vars if var.startswith("Emis")]
+        if len(emis_vars) > 0:
+            return True
     return False
diff --git a/gcpy/examples/diagnostics/compare_diags.py b/gcpy/examples/diagnostics/compare_diags.py
@@ -4,8 +4,6 @@
 collections.  Similar to compute_diagnostics.ipynb, but can be used
 without having to open a Jupyter notebook.
 """
-
-# Imports
 import os
 import sys
 import warnings
@@ -99,6 +97,10 @@ def read_data(config):
         msg = "Error reading " + dev_file
         raise Exception(msg) from exc
 
+    # Special handling for GCHP restart files
+    refdata = util.rename_and_flip_gchp_rst_vars(refdata)
+    devdata = util.rename_and_flip_gchp_rst_vars(devdata)
+
     # Define dictionary for return
     data = {
         "ref": refdata,
@@ -303,8 +305,7 @@ def compare_data(config, data):
     # ==================================================================
     # Print totals for each quantity
     # ==================================================================
-    if config["options"]["totals_and_diffs"]["create_table"] or \
-       config["options"]["totals_and_diffs"]["print_to_screen"]:
+    if config["options"]["totals_and_diffs"]["create_table"]:
         print('... Printing totals and differences')
         print_totals_and_diffs(
             config,