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 649b04b7..a02f1f00 100755 --- a/gcpy/examples/plotting/plot_comparisons.py +++ b/gcpy/examples/plotting/plot_comparisons.py @@ -52,7 +52,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 1980ef12..39d8adb8 100755 --- a/gcpy/examples/plotting/plot_single_panel.py +++ b/gcpy/examples/plotting/plot_single_panel.py @@ -37,8 +37,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/file_regrid.py b/gcpy/file_regrid.py index cba259e3..41cfea01 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) @@ -98,17 +98,43 @@ 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( - 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 +155,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 +223,7 @@ def file_regrid( # Write dataset to file dset.to_netcdf( - fout, + fileout, mode="w", format="NETCDF4", engine="netcdf4", @@ -261,7 +287,7 @@ def prepare_cssg_input_grid( def regrid_cssg_to_cssg( - fout, + fileout, dset, dim_format_in, sg_params_in, @@ -277,7 +303,7 @@ def regrid_cssg_to_cssg( Args: ----- - fout : str + fileout : str File name template dset : xarray.Dataset Data on a cubed-sphere/stretched grid @@ -304,10 +330,18 @@ 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): + # 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( @@ -346,7 +380,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): @@ -404,6 +438,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( @@ -412,14 +458,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, @@ -439,7 +477,6 @@ def regrid_cssg_to_ll( sg_params_in, ll_res_out, verbose=False, - weightsdir="." ): """ @@ -477,6 +514,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( @@ -502,18 +553,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( @@ -572,8 +619,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] @@ -590,13 +652,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", @@ -616,22 +671,14 @@ 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 # 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 @@ -702,6 +749,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, @@ -724,13 +778,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, @@ -785,38 +832,35 @@ 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) 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 # ================================================================== # 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,29 +869,33 @@ 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"], + gchp_indices=True + ) + dset = dset.assign_coords({"lev": coord}) dset.lev.attrs["positive"] = "down" 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 # ================================================================== @@ -860,24 +908,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 @@ -921,14 +951,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)" @@ -1012,16 +1040,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. @@ -1042,18 +1087,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( @@ -1123,7 +1214,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, @@ -1166,7 +1257,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. @@ -1189,12 +1283,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", @@ -1209,32 +1297,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): @@ -1398,14 +1465,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 +1533,6 @@ def main(): "--verbose", metavar="VERB", type=bool, - nargs=1, default=False, help="Toggles verbose output on (True) or off (False)" ) @@ -1474,7 +1540,6 @@ def main(): "-w", "--weightsdir", metavar="WGT", type=str, - nargs=1, default=False, help="Directory where regridding weights are found (or will be created)" ) 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 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" diff --git a/gcpy/util.py b/gcpy/util.py index a8fced06..beb08474 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,49 @@ 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 + old_to_new = {} 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