clm5/python/ctsm/crop_calendars/generate_gdds.py
2024-05-09 15:14:01 +08:00

496 lines
17 KiB
Python

"""
Generate maturity requirements (GDD) from outputs of a GDD-generating run
"""
import os
import sys
import pickle
import datetime as dt
import argparse
import logging
import numpy as np
import xarray as xr
# Import the CTSM Python utilities.
# sys.path.insert() is necessary for RXCROPMATURITY to work. The fact that it's calling this script
# in the RUN phase seems to require the python/ directory to be manually added to path.
_CTSM_PYTHON = os.path.join(
os.path.dirname(os.path.realpath(__file__)), os.pardir, os.pardir, os.pardir, "python"
)
sys.path.insert(1, _CTSM_PYTHON)
import ctsm.crop_calendars.cropcal_module as cc # pylint: disable=wrong-import-position
import ctsm.crop_calendars.generate_gdds_functions as gddfn # pylint: disable=wrong-import-position
# Global constants
PARAMFILE_DIR = "/glade/campaign/cesm/cesmdata/cseg/inputdata/lnd/clm2/paramdata"
MY_CLM_VER = 51
MY_CLM_SUBVER = "c211112"
def main(
input_dir=None,
first_season=None,
last_season=None,
sdates_file=None,
hdates_file=None,
output_dir=None,
save_figs=True,
only_make_figs=False,
run1_name=None,
run2_name=None,
land_use_file=None,
first_land_use_year=None,
last_land_use_year=None,
unlimited_season_length=False,
skip_crops=None,
logger=None,
):
# pylint: disable=missing-function-docstring,too-many-statements
# Directories to save output files and figures
if not output_dir:
if only_make_figs:
output_dir = input_dir
else:
output_dir = os.path.join(input_dir, "generate_gdds")
if not unlimited_season_length:
output_dir += ".mxmat"
output_dir += "." + dt.datetime.now().strftime("%Y-%m-%d-%H%M%S")
if not os.path.exists(output_dir):
os.makedirs(output_dir)
outdir_figs = os.path.join(output_dir, "figs")
# Set up log file and function, if needed
if logger is None:
logging.basicConfig(
level=logging.INFO,
format="",
filename=os.path.join(output_dir, "generate_gdds.log"),
filemode="a",
)
logger = logging.getLogger("")
# Disable plotting if any plotting module is unavailable
if save_figs:
try:
# pylint: disable=import-outside-toplevel,unused-import,import-error
import cartopy
import matplotlib
except ModuleNotFoundError as exc:
if only_make_figs:
raise RuntimeError(
"only_make_figs True but not all plotting modules are available"
) from exc
gddfn.log(logger, "Not all plotting modules are available; disabling save_figs")
save_figs = False
# Print some info
gddfn.log(logger, f"Saving to {output_dir}")
# Parse list of crops to skip
if "," in skip_crops:
skip_crops = skip_crops.split(",")
else:
skip_crops = skip_crops.split(" ")
##########################
### Import and process ###
##########################
if not only_make_figs:
# Keep 1 extra year to avoid incomplete final growing season for crops
# harvested after Dec. 31.
yr_1_import_str = f"{first_season+1}-01-01"
yr_n_import_str = f"{last_season+2}-01-01"
gddfn.log(
logger,
f"Importing netCDF time steps {yr_1_import_str} through {yr_n_import_str} "
+ "(years are +1 because of CTSM output naming)",
)
pickle_file = os.path.join(output_dir, f"{first_season}-{last_season}.pickle")
h2_ds_file = os.path.join(output_dir, f"{first_season}-{last_season}.h2_ds.nc")
if os.path.exists(pickle_file):
with open(pickle_file, "rb") as file:
(
first_season,
last_season,
pickle_year,
gddaccum_yp_list,
gddharv_yp_list,
skip_patches_for_isel_nan_lastyear,
lastyear_active_patch_indices_list,
incorrectly_daily,
save_figs,
incl_vegtypes_str,
incl_patches1d_itype_veg,
mxsowings,
skip_crops,
) = pickle.load(file)
print(f"Will resume import at {pickle_year+1}")
h2_ds = None
else:
incorrectly_daily = False
skip_patches_for_isel_nan_lastyear = np.ndarray([])
pickle_year = -np.inf
gddaccum_yp_list = []
gddharv_yp_list = []
incl_vegtypes_str = None
lastyear_active_patch_indices_list = None
sdates_rx = sdates_file
hdates_rx = hdates_file
if not unlimited_season_length:
mxmats = cc.import_max_gs_length(PARAMFILE_DIR, MY_CLM_VER, MY_CLM_SUBVER)
else:
mxmats = None
for yr_index, this_yr in enumerate(np.arange(first_season + 1, last_season + 3)):
if this_yr <= pickle_year:
continue
(
h2_ds,
sdates_rx,
hdates_rx,
gddaccum_yp_list,
gddharv_yp_list,
skip_patches_for_isel_nan_lastyear,
lastyear_active_patch_indices_list,
incorrectly_daily,
incl_vegtypes_str,
incl_patches1d_itype_veg,
mxsowings,
) = gddfn.import_and_process_1yr(
first_season,
last_season,
yr_index,
this_yr,
sdates_rx,
hdates_rx,
gddaccum_yp_list,
gddharv_yp_list,
skip_patches_for_isel_nan_lastyear,
lastyear_active_patch_indices_list,
incorrectly_daily,
input_dir,
incl_vegtypes_str,
h2_ds_file,
mxmats,
cc.get_gs_len_da,
skip_crops,
logger,
)
gddfn.log(logger, f" Saving pickle file ({pickle_file})...")
with open(pickle_file, "wb") as file:
pickle.dump(
[
first_season,
last_season,
this_yr,
gddaccum_yp_list,
gddharv_yp_list,
skip_patches_for_isel_nan_lastyear,
lastyear_active_patch_indices_list,
incorrectly_daily,
save_figs,
incl_vegtypes_str,
incl_patches1d_itype_veg,
mxsowings,
skip_crops,
],
file,
protocol=-1,
)
if isinstance(incl_vegtypes_str, list):
incl_vegtypes_str = np.array(incl_vegtypes_str)
plot_vegtypes_str = incl_vegtypes_str[
[i for i, c in enumerate(gddaccum_yp_list) if not isinstance(c, type(None))]
]
gddfn.log(logger, "Done")
if not h2_ds:
h2_ds = xr.open_dataset(h2_ds_file)
######################################################
### Get and grid mean GDDs in GGCMI growing season ###
######################################################
if not only_make_figs:
longname_prefix = "GDD harvest target for "
# Could skip this by saving sdates_rx['time_bounds']
sdates_rx = gddfn.import_rx_dates(
"s", sdates_rx, incl_patches1d_itype_veg, mxsowings, logger
)
gddfn.log(logger, "Getting and gridding mean GDDs...")
gdd_maps_ds = gddfn.yp_list_to_ds(
gddaccum_yp_list, h2_ds, incl_vegtypes_str, sdates_rx, longname_prefix, logger
)
gddharv_maps_ds = gddfn.yp_list_to_ds(
gddharv_yp_list, h2_ds, incl_vegtypes_str, sdates_rx, longname_prefix, logger
)
# Fill NAs with dummy values
dummy_fill = -1
gdd_maps_ds = gdd_maps_ds.fillna(dummy_fill)
gddfn.log(logger, "Done getting and gridding means.")
# Add dummy variables for crops not actually simulated
gddfn.log(logger, "Adding dummy variables...")
# Unnecessary?
template_ds = xr.open_dataset(sdates_file, decode_times=True)
all_vars = [v.replace("sdate", "gdd") for v in template_ds if "sdate" in v]
all_longnames = [
template_ds[v].attrs["long_name"].replace("Planting day ", longname_prefix) + " (dummy)"
for v in template_ds
if "sdate" in v
]
dummy_vars = []
dummy_longnames = []
for var_index, this_var in enumerate(all_vars):
if this_var not in gdd_maps_ds:
dummy_vars.append(this_var)
dummy_longnames.append(all_longnames[var_index])
def make_dummy(this_crop_gridded, addend):
dummy_gridded = this_crop_gridded
dummy_gridded.values = dummy_gridded.values * 0 + addend
return dummy_gridded
for var in gdd_maps_ds:
this_crop_gridded = gdd_maps_ds[var].copy()
break
dummy_gridded = make_dummy(this_crop_gridded, -1)
for var_index, this_var in enumerate(dummy_vars):
if this_var in gdd_maps_ds:
gddfn.error(
logger, f"{this_var} is already in gdd_maps_ds. Why overwrite it with dummy?"
)
dummy_gridded.name = this_var
dummy_gridded.attrs["long_name"] = dummy_longnames[var_index]
gdd_maps_ds[this_var] = dummy_gridded
# Add lon/lat attributes
def add_lonlat_attrs(this_ds):
this_ds.lon.attrs = {"long_name": "coordinate_longitude", "units": "degrees_east"}
this_ds.lat.attrs = {"long_name": "coordinate_latitude", "units": "degrees_north"}
return this_ds
gdd_maps_ds = add_lonlat_attrs(gdd_maps_ds)
gddharv_maps_ds = add_lonlat_attrs(gddharv_maps_ds)
gddfn.log(logger, "Done.")
######################
### Save to netCDF ###
######################
if not only_make_figs:
gddfn.log(logger, "Saving...")
# Get output file path
datestr = dt.datetime.now().strftime("%Y%m%d_%H%M%S")
outfile = os.path.join(output_dir, "gdds_" + datestr + ".nc")
def save_gdds(sdates_file, hdates_file, outfile, gdd_maps_ds, sdates_rx):
# Set up output file from template (i.e., prescribed sowing dates).
template_ds = xr.open_dataset(sdates_file, decode_times=True)
for var in template_ds:
if "sdate" in var:
template_ds = template_ds.drop(var)
template_ds.to_netcdf(path=outfile, format="NETCDF3_CLASSIC")
template_ds.close()
# Add global attributes
comment = (
"Derived from CLM run plus crop calendar input files "
+ f"{os.path.basename(sdates_file) and {os.path.basename(hdates_file)}}."
)
gdd_maps_ds.attrs = {
"author": "Sam Rabin (sam.rabin@gmail.com)",
"comment": comment,
"created": dt.datetime.now().astimezone().isoformat(),
}
# Add time_bounds
if "time_bounds" in sdates_rx:
gdd_maps_ds["time_bounds"] = sdates_rx.time_bounds
# Save cultivar GDDs
gdd_maps_ds.to_netcdf(outfile, mode="w", format="NETCDF3_CLASSIC")
save_gdds(sdates_file, hdates_file, outfile, gdd_maps_ds, sdates_rx)
gddfn.log(logger, "Done saving.")
########################################
### Save things needed for mapmaking ###
########################################
def add_attrs_to_map_ds(
map_ds, incl_vegtypes_str, dummy_fill, outdir_figs, first_season, last_season
):
return map_ds.assign_attrs(
{
"incl_vegtypes_str": incl_vegtypes_str,
"dummy_fill": dummy_fill,
"outdir_figs": outdir_figs,
"y1": first_season,
"yN": last_season,
}
)
if not only_make_figs:
if not os.path.exists(outdir_figs):
os.makedirs(outdir_figs)
gdd_maps_ds = add_attrs_to_map_ds(
gdd_maps_ds, plot_vegtypes_str, dummy_fill, outdir_figs, first_season, last_season
)
gddharv_maps_ds = add_attrs_to_map_ds(
gddharv_maps_ds, plot_vegtypes_str, dummy_fill, outdir_figs, first_season, last_season
)
gdd_maps_ds.to_netcdf(os.path.join(outdir_figs, "gdd_maps.nc"))
gddharv_maps_ds.to_netcdf(os.path.join(outdir_figs, "gddharv_maps.nc"))
#################################################
### Save before/after map and boxplot figures ###
#################################################
if save_figs:
if only_make_figs:
gdd_maps_ds = xr.open_dataset(os.path.join(input_dir, "figs", "gdd_maps.nc"))
gddharv_maps_ds = xr.open_dataset(os.path.join(input_dir, "figs", "gddharv_maps.nc"))
gddfn.make_figures(
first_land_use_year,
last_land_use_year,
land_use_file,
run1_name,
run2_name,
logger,
gdd_maps_ds=gdd_maps_ds,
gddharv_maps_ds=gddharv_maps_ds,
outdir_figs=outdir_figs,
)
if __name__ == "__main__":
###############################
### Process input arguments ###
###############################
parser = argparse.ArgumentParser(description="ADD DESCRIPTION HERE")
# Required
parser.add_argument(
"-i",
"--input-dir",
help=(
"Directory where run outputs can be found (and where outputs will go). If "
+ "--only-make-figs, this is the directory with the preprocessed files (e.g., *.pickle "
+ "file)."
),
required=True,
)
parser.add_argument(
"-1",
"--first-season",
help="First growing season to include in calculation of mean",
required=True,
type=int,
)
parser.add_argument(
"-n",
"-N",
"--last-season",
help="Last growing season to include in calculation of mean",
required=True,
type=int,
)
parser.add_argument(
"-o",
"--output-dir",
help="Output directory. Default is auto-generated subdir of -i/--input-dir.",
)
parser.add_argument(
"-sd", "--sdates-file", help="File of prescribed sowing dates", required=True
)
parser.add_argument(
"-hd", "--hdates-file", help="File of prescribed harvest dates", required=True
)
# Optional
figsgroup = parser.add_mutually_exclusive_group()
figsgroup.add_argument(
"--dont-save-figs", help="Do not save figures", action="store_true", default=False
)
figsgroup.add_argument(
"--only-make-figs",
help="Use preprocessed files to make figures only",
action="store_true",
default=False,
)
parser.add_argument(
"--run1-name", help="Name of original values to show in figures", default="Old"
)
parser.add_argument("--run2-name", help="Name of new values to show in figures", default="New")
parser.add_argument(
"-lu",
"--land-use-file",
help="Path to CLM land use timeseries file, for masking figures",
default=None,
)
parser.add_argument(
"--first-land-use-year",
help="First year in land use file to use for masking. Default --first-season.",
default=None,
type=int,
)
parser.add_argument(
"--last-land-use-year",
help="Last year in land use file to use for masking. Default --last-season.",
default=None,
type=int,
)
parser.add_argument(
"--unlimited-season-length",
help="Limit mean growing season length based on CLM CFT parameter mxmat.",
action="store_true",
default=False,
)
parser.add_argument(
"--skip-crops",
help="Skip processing of these crops. Comma- or space-separated list.",
type=str,
default="",
)
# Get arguments
args = parser.parse_args(sys.argv[1:])
for k, v in sorted(vars(args).items()):
print(f"{k}: {v}")
# Call main()
main(
input_dir=args.input_dir,
first_season=args.first_season,
last_season=args.last_season,
sdates_file=args.sdates_file,
hdates_file=args.hdates_file,
output_dir=args.output_dir,
save_figs=not args.dont_save_figs,
only_make_figs=args.only_make_figs,
run1_name=args.run1_name,
run2_name=args.run2_name,
land_use_file=args.land_use_file,
first_land_use_year=args.first_land_use_year,
last_land_use_year=args.last_land_use_year,
unlimited_season_length=args.unlimited_season_length,
skip_crops=args.skip_crops,
)