From e243d64fa34128a456c39c7bb5ff552429c2c2ce Mon Sep 17 00:00:00 2001 From: Bob Yantosca Date: Mon, 16 Oct 2023 15:49:03 -0400 Subject: [PATCH 1/7] Fixed several issues in file_regrid.py gcpy/file_regrid.py - Renamed "fin" to "filein", to match the command-line argument - Renamed "fout" to "fileout", to match the command-line argument - Added verbose printout of arguments - Remove nargs=1 from weightsdir and verbose, which turns them into a list of 1 element Signed-off-by: Bob Yantosca --- gcpy/file_regrid.py | 72 +++++++++++++++++++++++++++------------------ 1 file changed, 43 insertions(+), 29 deletions(-) diff --git a/gcpy/file_regrid.py b/gcpy/file_regrid.py index cba259e3..74b3606b 100644 --- a/gcpy/file_regrid.py +++ b/gcpy/file_regrid.py @@ -19,8 +19,8 @@ def file_regrid( - fin, - fout, + filein, + fileout, dim_format_in, dim_format_out, cs_res_out=0, @@ -36,9 +36,9 @@ def file_regrid( Args: ----- - fin: str + filein: str The input filename - fout: str + fileout: str The output filename (file will be overwritten if it already exists) dim_format_in: str Format of the input file's dimensions (choose from: classic, @@ -75,8 +75,8 @@ def file_regrid( Path to the directory containing regridding weights (or where weights will be created). Default value: "." """ - verify_variable_type(fin, str) - verify_variable_type(fout, str) + verify_variable_type(filein, str) + verify_variable_type(fileout, str) verify_variable_type(dim_format_in, str) verify_variable_type(dim_format_out, str) @@ -100,15 +100,28 @@ def file_regrid( # Load dataset dset = xr.open_dataset( - fin, + filein, decode_cf=False, engine="netcdf4" ).load() cs_res_in = get_cubed_sphere_res(dset) + # Verbose printout of inputs if verbose: - print(f"file_regrid.py: cs_res_in: {cs_res_in}") - print(f"file_regrid.py: cs_res_out: {cs_res_out}") + print("Inputs to file_regrid.py") + print(f"filein : {filein}") + print(f"dim_format_in : {dim_format_in}") + if "classic" not in dim_format_in: + print(f"sg_params_in : {sg_params_in}") + print(f"fileout : {fileout}") + print(f"dim_format_out : {dim_format_out}"), + if "classic" in dim_format_out: + print(f"ll_res_out : {ll_res_out}") + else: + print(f"cs_res_out : {cs_res_out}") + print(f"sg_params_out : {sg_params_out}") + print(f"verbose : {verbose}") + print(f"weightsdir : {weightsdir}") # Make sure all xarray.Dataset global & variable attributes are kept with xr.set_options(keep_attrs=True): @@ -129,7 +142,7 @@ def file_regrid( # Input grid is CS/SG; Output grid is CS/SG # ---------------------------------------------------------- dset = regrid_cssg_to_cssg( - fout, + fileout, dset, dim_format_in, sg_params_in, @@ -197,7 +210,7 @@ def file_regrid( # Write dataset to file dset.to_netcdf( - fout, + fileout, mode="w", format="NETCDF4", engine="netcdf4", @@ -261,7 +274,7 @@ def prepare_cssg_input_grid( def regrid_cssg_to_cssg( - fout, + fileout, dset, dim_format_in, sg_params_in, @@ -277,7 +290,7 @@ def regrid_cssg_to_cssg( Args: ----- - fout : str + fileout : str File name template dset : xarray.Dataset Data on a cubed-sphere/stretched grid @@ -346,7 +359,7 @@ def regrid_cssg_to_cssg( ) # Save temporary output face files to minimize RAM usage - oface_files = [os.path.join(".",fout+str(x)) for x in range(6)] + oface_files = [os.path.join(".",fileout+str(x)) for x in range(6)] # For each output face, sum regridded input faces for oface in range(6): @@ -630,8 +643,8 @@ def regrid_ll_to_cssg( # if "checkpoint" output are requested. dset = adjust_cssg_grid_and_coords( dset, - dim_format_in, - dim_format_out + dim_format_in="diagnostic", + dim_format_out=dim_format_out ) # Save stretched-grid metadata as global attrs @@ -785,7 +798,6 @@ def flip_lev_coord_if_necessary( # Case 1: checkpoint/diagnostic to classic # lev, ilev need to be in ascending order # ================================================================== - print(f"in {dim_format_in} out {dim_format_out}") if dim_format_in != "classic" and dim_format_out == "classic": # Flip lev and set to eta values at midpoints (if necessary) @@ -816,7 +828,7 @@ def flip_lev_coord_if_necessary( # ================================================================== # Case 2: classic/diagnostic to checkpoint - # lev needs to be in descending order + # lev needs to be in descending order (with ascending indices) # # TODO: Check for Emissions diagnostic (not a common use case) # ================================================================== @@ -825,12 +837,11 @@ def flip_lev_coord_if_necessary( if "lev" in dset.coords: if not is_gchp_lev_positive_down(dset): dset = dset.reindex(lev=dset.lev[::-1]) - if any(var > 1.0 for var in dset.lev): - coord = get_lev_coord( - n_lev=dset.dims["lev"], - top_down=True - ) - dset = dset.assign_coords({"lev": coord}) + coord = get_lev_coord( + n_lev=dset.dims["lev"], + top_down=True + ) + dset = dset.assign_coords({"lev": coord}) dset.lev.attrs["positive"] = "down" return dset @@ -1123,7 +1134,7 @@ def adjust_cssg_grid_and_coords( # Add "fake" Xdim and Ydim coordinates as done by MAPL, # which is needed for the GMAO GrADS visualization software. # NOTE: Use .values to convert to numpy.ndarray type in - # order to avoid xarray from trying to redefine dim "nf". + # order to avoid xarray from trying to redefileine dim "nf". if "lons" in dset.coords and "lats" in dset.coords: dset = dset.assign_coords({ "Xdim": dset.lons.isel(nf=0, Ydim=0).values, @@ -1328,6 +1339,11 @@ def reshape_cssg_diag_to_chkpt( # We then have to unpack that into a linear list that # ranges from 1..nf*ydim. # ============================================================== + print(dset.dims) + print(dset.dims["nf"]) + print(dset.dims["Ydim"]) + print(dset.Ydim) + print(ydim) if "nf" in dset.dims and "Ydim" in dset.dims: dset = dset.stack(lat=("nf", "Ydim")) multi_index_list = dset.lat.values @@ -1398,14 +1414,14 @@ def main(): ) parser.add_argument( "-i", "--filein", - metavar="FIN", + metavar="FILEIN", type=str, required=True, help="input NetCDF file" ) parser.add_argument( "-o", "--fileout", - metavar="FOUT", + metavar="FILEOUT", type=str, required=True, help="name of output file" @@ -1466,7 +1482,6 @@ def main(): "--verbose", metavar="VERB", type=bool, - nargs=1, default=False, help="Toggles verbose output on (True) or off (False)" ) @@ -1474,7 +1489,6 @@ def main(): "-w", "--weightsdir", metavar="WGT", type=str, - nargs=1, default=False, help="Directory where regridding weights are found (or will be created)" ) From 21ca8164a379367f656ec27b9375c149344995be Mon Sep 17 00:00:00 2001 From: Bob Yantosca Date: Thu, 19 Oct 2023 17:23:56 -0400 Subject: [PATCH 2/7] Fix issues in file_regrid.py and regrid_restart_file.py gcpy/file_regrid.py - Removed a double-flipping of vertical levels in LL -> CS/SG regridding gcpy/regrid_restart_file.py - Removed function check_lev_attribute, it is not defined Signed-off-by: Bob Yantosca --- gcpy/file_regrid.py | 8 -------- gcpy/regrid_restart_file.py | 2 -- 2 files changed, 10 deletions(-) diff --git a/gcpy/file_regrid.py b/gcpy/file_regrid.py index 74b3606b..9126a9b3 100644 --- a/gcpy/file_regrid.py +++ b/gcpy/file_regrid.py @@ -629,14 +629,6 @@ def regrid_ll_to_cssg( towards_common=False ) - # Flip vertical levels (if necessary) and - # set the lev:positive attribute accordingly - dset = flip_lev_coord_if_necessary( - dset, - dim_format_in="classic", - dim_format_out=dim_format_out - ) - # Fix names and attributes of of coordinate variables depending # on the format of the ouptut grid (checkpoint or diagnostic). # Also convert the "diagnostic" grid to the "checkpoint" grid diff --git a/gcpy/regrid_restart_file.py b/gcpy/regrid_restart_file.py index c0c96878..e811f1ad 100644 --- a/gcpy/regrid_restart_file.py +++ b/gcpy/regrid_restart_file.py @@ -533,8 +533,6 @@ def regrid_restart_file( "Error when processing your stretched-grid parameters - are they correct?" ) from exception - dataset = check_lev_attribute(output_is_gchp) - dataset.to_netcdf("new_restart_file.nc") info_message = "Wrote 'new_restart_file.nc' with %d variables" From b9a2b30c604aee369b87d566d310aef8806555c0 Mon Sep 17 00:00:00 2001 From: Bob Yantosca Date: Fri, 20 Oct 2023 10:58:05 -0400 Subject: [PATCH 3/7] Remove debug print statements in file_regrid.py gcpy/file_regrid.py - Removed several debug printout statements at about line 1333 Signed-off-by: Bob Yantosca --- gcpy/file_regrid.py | 5 ----- 1 file changed, 5 deletions(-) diff --git a/gcpy/file_regrid.py b/gcpy/file_regrid.py index 9126a9b3..e361dc2b 100644 --- a/gcpy/file_regrid.py +++ b/gcpy/file_regrid.py @@ -1331,11 +1331,6 @@ def reshape_cssg_diag_to_chkpt( # We then have to unpack that into a linear list that # ranges from 1..nf*ydim. # ============================================================== - print(dset.dims) - print(dset.dims["nf"]) - print(dset.dims["Ydim"]) - print(dset.Ydim) - print(ydim) if "nf" in dset.dims and "Ydim" in dset.dims: dset = dset.stack(lat=("nf", "Ydim")) multi_index_list = dset.lat.values From a466c2ecf02d5064e735d562b3e4fa98c91650d3 Mon Sep 17 00:00:00 2001 From: Bob Yantosca Date: Fri, 20 Oct 2023 15:16:11 -0400 Subject: [PATCH 4/7] rename_and_flip_gchp_rst_vars returns if data is not from GCHP restart gcpy/util.py - Routine "rename_and_flip_gchp_rst_vars" now uses the inquiry function "is_cubed_sphere_rst_grid" from cstools.py to determine if the dataset is from a GCHP restart file. If not, the original dataset will be returned unchanged. - Updated Pydoc and comments - Updated algorithm in "rename_and_flip_gchp_rst_vars" so that we create a dictionary of variable name replacements so that the replacement can be done as a single operation. This is more efficient. - If we have flipped the dataset, also set lev.positive="up" gcpy/benchmark_funcs.py gcpy/examples/diagnostics/compare_diags.py gcpy/examples/plotting/plot_comparisons.py gcpy/examples/plotting/plot_single_panel.py - Remove external tests for the GCHP restart grid before calling the "rename_and_flip_gchp_rst_vars" function from util.py. The function can now detemrine internally whether or not the passed dataset is from a GCHP restart file. If not the original data will be returned unchnaged. - Updated comments for clarity. Signed-off-by: Bob Yantosca --- gcpy/benchmark_funcs.py | 30 +++++-------- gcpy/examples/diagnostics/compare_diags.py | 16 ++++++- gcpy/examples/plotting/plot_comparisons.py | 4 +- gcpy/examples/plotting/plot_single_panel.py | 7 ++- gcpy/util.py | 49 ++++++++++++++------- 5 files changed, 68 insertions(+), 38 deletions(-) diff --git a/gcpy/benchmark_funcs.py b/gcpy/benchmark_funcs.py index aa1cc252..cca9690f 100644 --- a/gcpy/benchmark_funcs.py +++ b/gcpy/benchmark_funcs.py @@ -3016,13 +3016,11 @@ def make_benchmark_mass_tables( # Update GCHP restart dataset (if any) # ================================================================== - # Ref - if any(v.startswith("SPC_") for v in refds.data_vars.keys()): - refds = util.rename_and_flip_gchp_rst_vars(refds) - - # Dev - if any(v.startswith("SPC_") for v in devds.data_vars.keys()): - devds = util.rename_and_flip_gchp_rst_vars(devds) + # If the data is from a GCHP restart file, rename variables and + # flip levels to match the GEOS-Chem Classic naming and level + # conventions. Otherwise no changes will be made. + refds = util.rename_and_flip_gchp_rst_vars(refds) + devds = util.rename_and_flip_gchp_rst_vars(devds) # ================================================================== # Make sure that all necessary meteorological variables are found @@ -3288,17 +3286,13 @@ def make_benchmark_mass_accumulation_tables( # Update GCHP restart dataset if needed # ================================================================== - # Ref - if any(v.startswith("SPC_") for v in refSds.data_vars.keys()): - refSds = util.rename_and_flip_gchp_rst_vars(refSds) - if any(v.startswith("SPC_") for v in refEds.data_vars.keys()): - refEds = util.rename_and_flip_gchp_rst_vars(refEds) - - # Dev - if any(v.startswith("SPC_") for v in devSds.data_vars.keys()): - devSds = util.rename_and_flip_gchp_rst_vars(devSds) - if any(v.startswith("SPC_") for v in devEds.data_vars.keys()): - devEds = util.rename_and_flip_gchp_rst_vars(devEds) + # If the data is from a GCHP restart file, rename variables and + # flip levels to match the GEOS-Chem Classic naming and level + # conventions. Otherwise no changes will be made. + refSds = util.rename_and_flip_gchp_rst_vars(refSds) + refEds = util.rename_and_flip_gchp_rst_vars(refEds) + devSds = util.rename_and_flip_gchp_rst_vars(devSds) + devEds = util.rename_and_flip_gchp_rst_vars(devEds) # Add area to start restart dataset if area in end but not start # Need to consider area variable names used in both GC-Classic and GCHP diff --git a/gcpy/examples/diagnostics/compare_diags.py b/gcpy/examples/diagnostics/compare_diags.py index 35d05e53..d764d33e 100755 --- a/gcpy/examples/diagnostics/compare_diags.py +++ b/gcpy/examples/diagnostics/compare_diags.py @@ -97,7 +97,9 @@ def read_data(config): msg = "Error reading " + dev_file raise Exception(msg) from exc - # Special handling for GCHP restart files + # If the data is from a GCHP restart file, rename variables and + # flip levels to match the GEOS-Chem Classic naming and level + # conventions. Otherwise no changes will be made. refdata = util.rename_and_flip_gchp_rst_vars(refdata) devdata = util.rename_and_flip_gchp_rst_vars(devdata) @@ -261,6 +263,14 @@ def compare_data(config, data): varlist_level = [v for v in varlist_level if v in restrict_vars] varlist_zonal = [v for v in varlist_zonal if v in restrict_vars] + # Determine if we need to flip levels in the vertical + flip_ref = False + flip_dev = False + if "flip_levels" in config["data"]["ref"]: + flip_ref = config["data"]["ref"]["flip_levels"] + if "flip_levels" in config["data"]["dev"]: + flip_dev = config["data"]["dev"]["flip_levels"] + # ================================================================== # Generate the single level comparison plot # ================================================================== @@ -275,6 +285,8 @@ def compare_data(config, data): config["data"]["ref"]["label"], devdata, config["data"]["dev"]["label"], + flip_ref=flip_ref, + flip_dev=flip_dev, ilev=config["options"]["level_plot"]["level_to_plot"], varlist=varlist_level, pdfname=pdfname, @@ -296,6 +308,8 @@ def compare_data(config, data): config["data"]["ref"]["label"], devdata, config["data"]["dev"]["label"], + flip_ref=flip_ref, + flip_dev=flip_dev, varlist=varlist_zonal, pdfname=pdfname, weightsdir=config["paths"]["weights_dir"], diff --git a/gcpy/examples/plotting/plot_comparisons.py b/gcpy/examples/plotting/plot_comparisons.py index be56a708..47d6f5bd 100755 --- a/gcpy/examples/plotting/plot_comparisons.py +++ b/gcpy/examples/plotting/plot_comparisons.py @@ -54,7 +54,9 @@ def plot_comparisons( drop_variables=skip_these_vars ) - # Special handling is needed for GCHP restart files + # If the data is from a GCHP restart file, rename variables and + # flip levels to match the GEOS-Chem Classic naming and level + # conventions. Otherwise no changes will be made. ref_ds = rename_and_flip_gchp_rst_vars(ref_ds) dev_ds = rename_and_flip_gchp_rst_vars(dev_ds) diff --git a/gcpy/examples/plotting/plot_single_panel.py b/gcpy/examples/plotting/plot_single_panel.py index 7ec54dff..777de7ac 100755 --- a/gcpy/examples/plotting/plot_single_panel.py +++ b/gcpy/examples/plotting/plot_single_panel.py @@ -39,8 +39,11 @@ def plot_single_panel(infile, varname, level): # xarray allows us to read in any NetCDF file dset = xr.open_dataset(infile) - # Special handling if the file is a GCHP restart file - dset = rename_and_flip_gchp_rst_vars(dset) + # If the data is from a GCHP restart file, rename variables and + # flip levels to match the GEOS-Chem Classic naming and level + # conventions. Otherwise no changes will be made. + ref_ds = rename_and_flip_gchp_rst_vars(ref_ds) + dev_ds = rename_and_flip_gchp_rst_vars(dev_ds) # You can easily view the variables available for plotting # using xarray. Each of these variables has its own xarray diff --git a/gcpy/util.py b/gcpy/util.py index a8fced06..dd66bede 100644 --- a/gcpy/util.py +++ b/gcpy/util.py @@ -11,6 +11,7 @@ import xarray as xr from pypdf import PdfWriter, PdfReader from gcpy.constants import ENCODING, TABLE_WIDTH +from gcpy.cstools import is_cubed_sphere_rst_grid # ====================================================================== # %%%%% METHODS %%%%% @@ -807,32 +808,48 @@ def rename_and_flip_gchp_rst_vars( dset ): ''' - Transforms a GCHP restart dataset to match GCC names and level convention + Transforms a GCHP restart dataset to match GCClassic names + and level conventions. Args: dset: xarray Dataset - Dataset containing GCHP restart file data, such as variables - SPC_{species}, BXHEIGHT, DELP_DRY, and TropLev, with level - convention down (level 0 is top-of-atmosphere). + The input dataset. Returns: dset: xarray Dataset - Dataset containing GCHP restart file data with names and level - convention matching GCC restart. Variables include - SpeciesRst_{species}, Met_BXHEIGHT, Met_DELPDRY, and Met_TropLev, - with level convention up (level 0 is surface). + If the input dataset is from a GCHP restart file, then + dset will contain the original data with variables renamed + to match the GEOS-Chem Classic naming conventions, and + with levels indexed as lev:positive="up". Otherwise, the + original data will be returned. ''' + verify_variable_type(dset, xr.Dataset) + + # Return if this dataset is not from a GCHP checkpoint/restart file + if not is_cubed_sphere_rst_grid(dset): + return dset + + # Create dictionary of variable name replacements for var in dset.data_vars.keys(): - if var.startswith('SPC_'): + # TODO: Think of better algorithm in case we ever change + # the internal state to start with something else than "SPC_". + if var.startswith("SPC_"): spc = var.replace('SPC_', '') - dset = dset.rename({var: 'SpeciesRst_' + spc}) - elif var == 'DELP_DRY': - dset = dset.rename({"DELP_DRY": "Met_DELPDRY"}) - elif var == 'BXHEIGHT': - dset = dset.rename({"BXHEIGHT": "Met_BXHEIGHT"}) - elif var == 'TropLev': - dset = dset.rename({"TropLev": "Met_TropLev"}) + old_to_new[var] = 'SpeciesRst_' + spc + if var == "DELP_DRY": + old_to_new["DELP_DRY"] = "Met_DELPDRY" + if var == "BXHEIGHT": + old_to_new["BXHEIGHT"] = "Met_BXHEIGHT" + if var == "TropLev": + old_to_new["TropLev"] = "Met_TropLev" + + # Replace variable names in one operation + dset = dset.rename(old_to_new) + + # Flip levels dset = dset.sortby('lev', ascending=False) + dset.lev.attrs["positive"] = "up" + return dset From 5d83b5ed1dc2c032f5aaaa6c120c9cf170b5bc24 Mon Sep 17 00:00:00 2001 From: Bob Yantosca Date: Fri, 20 Oct 2023 15:18:48 -0400 Subject: [PATCH 5/7] get_lev_coords, get_ilev_coords can return an array of indices gcpy/grid.py - Remove unused import statement - Added "gchp_indices" argument to get_lev_coords and get_ilev_coords functions. If gchp_indices=True, the lev/ilev coordinates will be returned as an array of indices (as is the custom for GCHP netCDF files). - Updated comments and Pydoc Signed-off-by: Bob Yantosca --- gcpy/grid.py | 34 +++++++++++++++++++++++++--------- 1 file changed, 25 insertions(+), 9 deletions(-) diff --git a/gcpy/grid.py b/gcpy/grid.py index 428ebc5b..d9a51805 100644 --- a/gcpy/grid.py +++ b/gcpy/grid.py @@ -2,10 +2,10 @@ Module containing variables and functions that define and manipulate GEOS-Chem horizontal and vertical grids """ +from itertools import product import xarray as xr import numpy as np import scipy.sparse -from itertools import product from gcpy.util import get_shape_of_data, verify_variable_type from .grid_stretching_transforms import scs_transform from gcpy.constants import R_EARTH_m @@ -296,7 +296,8 @@ def get_ilev_coord( n_lev=72, AP_edge=None, BP_edge=None, - top_down=False + top_down=False, + gchp_indices=False ): """ Returns the eta values (defined as (A/P0) + B) at vertical @@ -316,8 +317,11 @@ def get_ilev_coord( edges. If not specified, values from the _GEOS_72L_BP array in this module will be used. top_down : bool - Set this to true if the eta coordinate will be arranged from + Set this to True if the eta coordinate will be arranged from top-of-atm downward (True) or from the surface upward (False). + gchp_indices : bool + Set this to True to return an array of indices (as is used + in GCHP files). Returns: -------- @@ -326,6 +330,12 @@ def get_ilev_coord( """ if n_lev is None: n_lev = 72 + + # Return GCHP-style indices for the level dimension + if gchp_indices: + return np.linspace(1, n_lev+1, n_lev+1, dtype=np.float64) + + # Get eta values at vertical level edges if AP_edge is None and n_lev == 72: AP_edge = _GEOS_72L_AP if AP_edge is None and n_lev == 47: @@ -334,8 +344,6 @@ def get_ilev_coord( BP_edge = _GEOS_72L_BP if BP_edge is None and n_lev == 47: BP_edge = _GEOS_47L_BP - - # Get eta values at vertical level edges ilev = np.array((AP_edge/1000.0) + BP_edge, dtype=np.float64) if top_down: ilev = ilev[::-1] @@ -345,7 +353,8 @@ def get_lev_coord( n_lev=72, AP_edge=None, BP_edge=None, - top_down=False + top_down=False, + gchp_indices=False ): """ Returns the eta values (defined as (A/P0) + B) at vertical @@ -367,6 +376,9 @@ def get_lev_coord( top_down : bool Set this to true if the eta coordinate will be arranged from top-of-atm downward (True) or from the surface upward (False). + gchp_indices : bool + Set this to True to return an array of indices (as is used + in GCHP files). Returns: -------- @@ -375,6 +387,13 @@ def get_lev_coord( """ if n_lev is None: n_lev = 72 + + # Return GCHP-style indices for the level dimension + if gchp_indices: + return np.linspace(1, n_lev, n_lev, dtype=np.float64) + + # Compute AP, BP at midpoints. + # Convert inputs to numpy.ndarray for fast computation if AP_edge is None and n_lev == 72: AP_edge = _GEOS_72L_AP if AP_edge is None and n_lev == 47: @@ -383,9 +402,6 @@ def get_lev_coord( BP_edge = _GEOS_72L_BP if BP_edge is None and n_lev == 47: BP_edge = _GEOS_47L_BP - - # Compute AP, BP at midpoints. - # Convert inputs to numpy.ndarray for fast computation AP_edge = np.array(AP_edge) BP_edge = np.array(BP_edge) AP_mid = (AP_edge[0:n_lev:1] + AP_edge[1:n_lev+1:1]) * 0.5 From 5db734757cc865cac35f07961bd848eba7476c58 Mon Sep 17 00:00:00 2001 From: Bob Yantosca Date: Fri, 20 Oct 2023 18:19:30 -0400 Subject: [PATCH 6/7] Fix several regridding issues in file_regrid.py gcpy_file_regrid.py - Move flip_lev_coord_if_necessary so that it occurs before variable renaming/regridding. This prevents the flipping algorithm from being spoofed by renamed variables. - Renamed "drop_and_rename_classic_variables" to "drop_classic_vars", which only drops the non-regriddable variables. - Move variable renaming into "rename_restart_variables". Also improve logic so that variables are renamed properly - Updated logic in the flip_lev_coord_if_necessary routine. Also make sure to save coordinates in the proper order. GCHP level coordinates are now saved as indices. - Updated Pydoc and comments. - Trimmed trailing whitespace. gcpy/util.py - Initialize "old_to_new" to a blank dictionary object Signed-off-by: Bob Yantosca --- gcpy/file_regrid.py | 291 ++++++++++++++++++++++++++------------------ gcpy/util.py | 1 + 2 files changed, 172 insertions(+), 120 deletions(-) diff --git a/gcpy/file_regrid.py b/gcpy/file_regrid.py index e361dc2b..0dcba97c 100644 --- a/gcpy/file_regrid.py +++ b/gcpy/file_regrid.py @@ -114,7 +114,7 @@ def file_regrid( if "classic" not in dim_format_in: print(f"sg_params_in : {sg_params_in}") print(f"fileout : {fileout}") - print(f"dim_format_out : {dim_format_out}"), + print(f"dim_format_out : {dim_format_out}") if "classic" in dim_format_out: print(f"ll_res_out : {ll_res_out}") else: @@ -321,6 +321,14 @@ def regrid_cssg_to_cssg( # Keep all xarray attributes with xr.set_options(keep_attrs=True): + # Flip vertical levels (if necessary) and + # set the lev:positive attribute accordingly + dset = flip_lev_coord_if_necessary( + dset, + dim_format_in=dim_format_in, + dim_format_out=dim_format_out + ) + # Change CS/SG dimensions to universal format # and drop non-regriddable variables dset, cs_res_in = prepare_cssg_input_grid( @@ -417,6 +425,18 @@ def regrid_cssg_to_cssg( towards_common=False ) + # Rename variables if we are going between different grid types + if "checkpoint" in dim_format_in and "diagnostic" in dim_format_out: + dset = rename_restart_variables( + dset, + towards_gchp=False + ) + if "diagnostic" in dim_format_in and "checkpoint" in dim_format_out: + dset = rename_restart_variables( + dset, + towards_gchp=True + ) + # Fix names and attributes of of coordinate variables depending # on the format of the ouptut grid (checkpoint or diagnostic). dset = adjust_cssg_grid_and_coords( @@ -425,14 +445,6 @@ def regrid_cssg_to_cssg( dim_format_out ) - # Flip vertical levels (if necessary) and - # set the lev:positive attribute accordingly - dset = flip_lev_coord_if_necessary( - dset, - dim_format_in=dim_format_in, - dim_format_out=dim_format_out - ) - # Save stretched-grid metadata as global attrs dset = save_cssg_metadata( dset, @@ -452,7 +464,6 @@ def regrid_cssg_to_ll( sg_params_in, ll_res_out, verbose=False, - weightsdir="." ): """ @@ -490,6 +501,20 @@ def regrid_cssg_to_ll( with xr.set_options(keep_attrs=True): + # Flip vertical levels (if necessary) and + # set the lev:positive attribute accordingly + dset = flip_lev_coord_if_necessary( + dset, + dim_format_in=dim_format_in, + dim_format_out="classic" + ) + + # Drop non-regriddable variables (if any) + dset = drop_classic_vars( + dset, + towards_gchp=False + ) + # Change CS/SG dimensions to universal format # and drop non-regriddable variables dset, cs_res_in = prepare_cssg_input_grid( @@ -515,18 +540,14 @@ def regrid_cssg_to_ll( "T": "time", "Z": "lev" }) - dset = drop_and_rename_classic_vars( - dset, - towards_gchp=False - ) - # Flip vertical levels (if necessary) and - # set the lev:positive attribute accordingly - dset = flip_lev_coord_if_necessary( - dset, - dim_format_in=dim_format_in, - dim_format_out="classic" - ) + # If regridding from a GCHP checkpoint/restart file, then + # rename variables to adhere GCClassic name conventions. + if "checkpoint" in dim_format_in: + dset = rename_restart_variables( + dset, + towards_gchp=False + ) # Save lat/lon coordinate metadata dset = save_ll_metadata( @@ -585,8 +606,23 @@ def regrid_ll_to_cssg( with xr.set_options(keep_attrs=True): - # Drop non-regriddable variables when going from ll -> cs - dset = drop_and_rename_classic_vars(dset) + # Flip vertical levels (if necessary) and set lev:positive + dset = flip_lev_coord_if_necessary( + dset, + dim_format_in="classic", + dim_format_out=dim_format_out + ) + + # Drop non- regriddable variables when going from ll -> cs + dset = drop_classic_vars(dset) + + # If regridding to a GCHP checkpoint/restart file, then + # rename variables to adhere to GCHP naming conventions. + if "checkpoint" in dim_format_out: + dset = rename_restart_variables( + dset, + towards_gchp=True + ) # Input lat/lon grid resolution llres_in = get_input_res(dset)[0] @@ -603,13 +639,6 @@ def regrid_ll_to_cssg( dim="nf" ) - # Flip vertical levels (if necessary) and set lev:positive - dset = flip_lev_coord_if_necessary( - dset, - dim_format_in="classic", - dim_format_out=dim_format_out - ) - # Rename dimensions to the "common dimension format" dset = dset.rename({ "time": "T", @@ -707,6 +736,13 @@ def regrid_ll_to_ll( non_fields = dset[non_fields] dset = dset.drop(non_fields) + # Set the lev:positive attribute accordingly + dset = flip_lev_coord_if_necessary( + dset, + dim_format_in="classic", + dim_format_out="classic" + ) + # Create the regridder and regrid the data regridder = make_regridder_L2L( ll_res_in, @@ -729,13 +765,6 @@ def regrid_ll_to_ll( "time", "lev", "ilev", "lat", "lon", ... ) - # Set the lev:positive attribute accordingly - dset = flip_lev_coord_if_necessary( - dset, - dim_format_in="classic", - dim_format_out="classic" - ) - # Save lat/lon coordinate metadata dset = save_ll_metadata( dset, @@ -796,24 +825,22 @@ def flip_lev_coord_if_necessary( if "ilev" in dset.coords: if is_gchp_lev_positive_down(dset): dset = dset.reindex(ilev=dset.ilev[::-1]) - if any(var > 1.0 for var in dset.ilev): - coord = get_ilev_coord( - n_lev=dset.dims["ilev"], - top_down=False - ) - dset = dset.assign_coords({"ilev": coord}) + coord = get_ilev_coord( + n_lev=dset.dims["ilev"], + top_down=False + ) + dset = dset.assign_coords({"ilev": coord}) dset.ilev.attrs["positive"] = "up" # Flip lev and set to eta values at midpoints (if necessary) if "lev" in dset.coords: if is_gchp_lev_positive_down(dset): dset = dset.reindex(lev=dset.lev[::-1]) - if any(var > 1.0 for var in dset.lev): - coord = get_lev_coord( - n_lev=dset.dims["lev"], - top_down=False - ) - dset = dset.assign_coords({"lev": coord}) + coord = get_lev_coord( + n_lev=dset.dims["lev"], + top_down=False + ) + dset = dset.assign_coords({"lev": coord}) dset.lev.attrs["positive"] = "up" return dset @@ -831,7 +858,7 @@ def flip_lev_coord_if_necessary( dset = dset.reindex(lev=dset.lev[::-1]) coord = get_lev_coord( n_lev=dset.dims["lev"], - top_down=True + gchp_indices=True ) dset = dset.assign_coords({"lev": coord}) dset.lev.attrs["positive"] = "down" @@ -839,18 +866,23 @@ def flip_lev_coord_if_necessary( return dset # ================================================================== - # Case 3: classic/diagnostic to classic/diagnostic: - # No flipping, but add lev:positive="up" and ilev:positive="up" + # Case 3: classic/checkpoint to diagnostic: + # lev, ilev need to be in ascending order # # TODO: Check for Emissions diagnostic (not a common use case) # ================================================================== - if dim_format_in == "classic" and dim_format_out == "diagnostic" or \ - dim_format_in == "diagnostic" and dim_format_out == "classic": + if dim_format_in != "diagnostic" and dim_format_out == "diagnostic": - if "ilev" in dset.coords: - dset.ilev.attrs["positive"] = "up" if "lev" in dset.coords: + if is_gchp_lev_positive_down(dset): + dset = dset.reindex(lev=dset.lev[::-1]) + coord = get_lev_coord( + n_lev=dset.dims["lev"], + gchp_indices=True + ) + dset = dset.assign_coords({"lev": coord}) dset.lev.attrs["positive"] = "up" + return dset # ================================================================== @@ -863,24 +895,6 @@ def flip_lev_coord_if_necessary( dset.lev.attrs["positive"] = "down" return dset - # ================================================================== - # Case 5: checkpoint to diagnostic: - # lev needs to be in ascending order - # ================================================================== - if dim_format_in == "checkpoint" and dim_format_out == "diagnostic": - - if "lev" in dset.coords: - dset = dset.reindex(lev=dset.lev[::-1]) - if any(var > 1.0 for var in dset.lev): - coord = get_lev_coord( - n_lev=dset.dims["lev"], - top_down=False - ) - dset = dset.assign_coords({"lev": coord}) - dset.lev.attrs["positive"] = "up" - - return dset - return dset @@ -924,14 +938,12 @@ def save_ll_metadata( "axis": "X" } - # ilev:positive is set by flip_lev_coord_if_necessary if "ilev" in dset.coords: dset.ilev.attrs["long_name"] = \ "hybrid level at interfaces ((A/P0)+B)" dset.ilev.attrs["units"] = "level" dset.ilev.attrs["axis"] = "Z" - # lev:positive is set by flip_lev_coord_if_necessary if "lev" in dset.coords: dset.lev.attrs["long_name"] = \ "hybrid level at midpoints ((A/P0)+B)" @@ -1015,16 +1027,33 @@ def save_cssg_metadata( } if "lats" in dset.dims: dset.lats.attrs = { - "standard_name": "latitude", + "standard_name": "la7titude", "long_name": "Latitude", "units": "degrees_north", "axis": "Y" } + # ilev:positive is set by flip_lev_coord_if_necessary + if "ilev" in dset.coords: + dset.ilev.attrs["long_name"] = \ + "hybrid level at interfaces ((A/P0)+B)" + dset.ilev.attrs["units"] = "level" + dset.ilev.attrs["axis"] = "Z" + + # lev:positive is set by flip_lev_coord_if_necessary + if "lev" in dset.coords: + dset.lev.attrs["long_name"] = \ + "hybrid level at midpoints ((A/P0)+B)" + dset.lev.attrs["units"] = "level" + dset.lev.attrs["axis"] = "Z" + return dset -def rename_restart_variables(dset, towards_gchp=True): +def rename_restart_variables( + dset, + towards_gchp=True +): """ Renames restart variables according to GEOS-Chem Classic and GCHP conventions. @@ -1045,18 +1074,64 @@ def rename_restart_variables(dset, towards_gchp=True): dset : xarray.Dataset The modified dataset. """ + verify_variable_type(dset, xr.Dataset) - if towards_gchp: - old_str = "SpeciesRst" - new_str = "SPC" - else: - old_str = "SPC" - new_str = "SpeciesRst" - return dset.rename({ - name: name.replace(old_str, new_str, 1) - for name in list(dset.data_vars) - if name.startswith(old_str) - }) + # Keep all xarray attribute settings + with xr.set_options(keep_attrs=True): + + # Dictionary for name replacements + old_to_new = {} + + # ============================================================== + # classic/diagnostic -> checkpoint + # ============================================================== + if towards_gchp: + for var in dset.data_vars.keys(): + if "Met_DELPDRY" in var: + old_to_new[var] = "DELP_DRY" + if var.startswith("Met_"): + old_to_new[var] = var.replace("Met_", "") + if var.startswith("Chem_"): + old_to_new[var] = var.replace("Chem_", "") + if var.startswith("SpeciesRst_"): + old_to_new[var] = var.replace("SpeciesRst_", "SPC_") + if var.startswith("SpeciesConcVV_"): + old_to_new[var] = var.replace("SpeciesConcVV_", "SPC_") + + return dset.rename(old_to_new) + + # ============================================================== + # checkpoint -> classic/diagnostic + # ============================================================== + for var in dset.data_vars.keys(): + if var == "DELP_DRY": + old_to_new[var] = "Met_DELPDRY" + if var == "BXHEIGHT": + old_to_new[var] = "Met_BXHEIGHT" + if var == "StatePSC": + old_to_new[var] = "Chem_StatePSC" + if var == "KPPHvalue": + old_to_new[var] = "Chem_KPPHvalue" + if var == "DryDepNitrogen": + old_to_new[var] = "ChemDryDepNitrogen" + if var == "WetDepNitrogen": + old_to_new[var] = "Chem_WetDepNitrogen" + if var == "SO2AfterChem": + old_to_new[var] = "Chem_SO2AfterChem" + if var == "JNO2": + old_to_new[var] = "Chem_JNO2" + if var == "JOH": + old_to_new[var] = "Chem_JOH" + if var == "H2O2AfterChem": + old_to_new[var] = "Chem_H2O2AfterChem" + if var == "ORVCsesq": + old_to_new[var] = "Chem_ORVCsesq" + if var == "AeroH2OSNA": + old_to_new[var] = "Chem_AeroH2OSNA" + if var.startswith("SPC_"): + old_to_new[var] = var.replace("SPC_", "SpeciesRst_") + + return dset.rename(old_to_new) def adjust_cssg_grid_and_coords( @@ -1169,7 +1244,10 @@ def adjust_cssg_grid_and_coords( return dset -def drop_and_rename_classic_vars(dset, towards_gchp=True): +def drop_classic_vars( + dset, + towards_gchp=True +): """ Renames and drops certain restart variables according to GEOS-Chem Classic and GCHP conventions. @@ -1192,12 +1270,6 @@ def drop_and_rename_classic_vars(dset, towards_gchp=True): """ with xr.set_options(keep_attrs=True): if towards_gchp: - dset = dset.rename( - {name: name.replace("Met_", "", 1).replace("Chem_", "", 1) - for name in list(dset.data_vars) - if name.startswith("Met_") or name.startswith("Chem_")}) - if "DELPDRY" in list(dset.data_vars): - dset = dset.rename({"DELPDRY": "DELP_DRY"}) dset = dset.drop_vars( ["P0", "hyam", @@ -1212,32 +1284,11 @@ def drop_and_rename_classic_vars(dset, towards_gchp=True): "SPHU1", "StatePSC", "lon_bnds", - "lat_bnds" - ], + "lat_bnds"], errors="ignore" ) - else: - renames = { - "DELP_DRY": "Met_DELPDRY", - "BXHEIGHT": "Met_BXHEIGHT", - "TropLev": "Met_TropLev", - "DryDepNitrogen": "Chem_DryDepNitrogen", - "WetDepNitrogen": "Chem_WetDepNitrogen", - "H2O2AfterChem": "Chem_H2O2AfterChem", - "SO2AfterChem": "Chem_SO2AfterChem", - "KPPHvalue": "Chem_KPPHvalue" - } - data_vars = list(dset.data_vars) - new_renames = renames.copy() - for items in renames.items(): - if items[0] not in data_vars: - del new_renames[items[0]] - dset = dset.rename(new_renames) - - return rename_restart_variables( - dset, - towards_gchp=towards_gchp - ) + + return dset def order_dims_time_lev_lat_lon(dset): diff --git a/gcpy/util.py b/gcpy/util.py index dd66bede..beb08474 100644 --- a/gcpy/util.py +++ b/gcpy/util.py @@ -830,6 +830,7 @@ def rename_and_flip_gchp_rst_vars( return dset # Create dictionary of variable name replacements + old_to_new = {} for var in dset.data_vars.keys(): # TODO: Think of better algorithm in case we ever change # the internal state to start with something else than "SPC_". From 1cbc921a49486aa70e2cc7d2fdd69e5f73ac958d Mon Sep 17 00:00:00 2001 From: Bob Yantosca Date: Tue, 24 Oct 2023 11:35:47 -0400 Subject: [PATCH 7/7] Prevent regridding to/from stretched grids in file_regrid.py gcpy/file_regrid.py - Because there is an issue in regridding to stretched-grids, we now throw a RuntimeError if sg_params_in or sg_params_out does not match the default (i.e. non-stretched) values. Signed-off-by: Bob Yantosca --- gcpy/file_regrid.py | 15 ++++++++++++++- 1 file changed, 14 insertions(+), 1 deletion(-) diff --git a/gcpy/file_regrid.py b/gcpy/file_regrid.py index 0dcba97c..41cfea01 100644 --- a/gcpy/file_regrid.py +++ b/gcpy/file_regrid.py @@ -98,6 +98,19 @@ def file_regrid( if sg_params_out is None: sg_params_out = [1.0, 170.0, -90.0] + # ------------------------------------------------------------------ + # There still seem to be a few issues with regridding to cubed- + # sphere stretched grids. For time time being, stop with error + # if sg_params_in or sg_params_out do not equal the defaults. + # -- Bob Yantosca & Lizzie Lundgren (24 Oct 2023) + if not np.array_equal(sg_params_in, [1.0, 170.0, -90.0]) or \ + not np.array_equal(sg_params_out, [1.0, 170.0, -90.0]): + msg = "Regridding to or from cubed-sphere stretched grids is\n" + \ + "currently not supported. Please use the offline regridding\n" + \ + "method described in the Regridding section of gcpy.readthedocs.io." + raise RuntimeError(msg) + # ------------------------------------------------------------------ + # Load dataset dset = xr.open_dataset( filein, @@ -317,7 +330,7 @@ def regrid_cssg_to_cssg( """ if verbose: print("file_regrid.py: Regridding from CS/SG to CS/SG") - + # Keep all xarray attributes with xr.set_options(keep_attrs=True):