generate_gdds: Work on arguments

cime_config/SystemTests/rxcropmaturity.py

@@ -448,6 +448,7 @@ def _run_generate_gdds(self, case_gddgen):
                 f"--hdates-file {hdates_file}",
                 f"--output-dir generate_gdds_out",
                 f"--skip-crops miscanthus,irrigated_miscanthus,switchgrass,irrigated_switchgrass",
+                "--max-season-length-from-hdates-file",
             ]
         )
         stu.run_python_script(

python/ctsm/crop_calendars/cropcal_module.py

@@ -2,8 +2,8 @@
 Helper functions for various crop calendar stuff
 """
 
-import os
-import glob
+import warnings
+
 import numpy as np
 import xarray as xr
 
@@ -203,16 +203,12 @@ def get_gs_len_da(this_da):
     return this_da
 
 
-def import_max_gs_length(paramfile_dir, my_clm_ver, my_clm_subver):
+def import_max_gs_length(paramfile):
     """
     Import maximum growing season length
     """
     # Get parameter file
-    pattern = os.path.join(paramfile_dir, f"*{my_clm_ver}_params.{my_clm_subver}.nc")
-    paramfile = glob.glob(pattern)
-    if len(paramfile) != 1:
-        raise RuntimeError(f"Expected to find 1 match of {pattern}; found {len(paramfile)}")
-    paramfile_ds = xr.open_dataset(paramfile[0])
+    paramfile_ds = xr.open_dataset(paramfile)
 
     # Import max growing season length (stored in netCDF as nanoseconds!)
     paramfile_mxmats = paramfile_ds["mxmat"].values / np.timedelta64(1, "D")
@@ -234,6 +230,56 @@ def import_max_gs_length(paramfile_dir, my_clm_ver, my_clm_subver):
     return mxmat_dict
 
 
+def cushion_gs_length(mxmat_dict_in, cushion, *, min_mxmat=1, max_mxmat=365):
+    """
+    Given a dictionary of maximum growing season lengths, apply a "cushion". This is useful for
+    generating crop maturity requirements: If observed growing season is longer than maximum,
+    and you're limiting growing seasons based on the maximum, then you'd expect about 50% of
+    growing seasons to fail to reach maturity (assuming a normal distribution of seasonal
+    growing degree-days).
+    """
+
+    mxmat_dict_out = mxmat_dict_in.copy()
+
+    for pftname, mxmat in mxmat_dict_in.items():
+        # Skip PFTs without max growing season length
+        if np.isinf(mxmat):
+            continue
+
+        assert mxmat >= min_mxmat, f"{pftname} input mxmat ({mxmat}) is < min_mxmat ({min_mxmat})"
+        assert mxmat <= max_mxmat, f"{pftname} input mxmat ({mxmat}) is > max_mxmat ({max_mxmat})"
+
+        new_mxmat = mxmat - cushion
+
+        # Apply limits
+        msg = None
+        if new_mxmat < min_mxmat:
+            msg = (
+                f"Applying cushion of {cushion} to {pftname}'s mxmat ({mxmat}) resulted in new"
+                f" mxmat of {new_mxmat}; increasing that to min_mxmat {min_mxmat}"
+            )
+            new_mxmat = min_mxmat
+        elif new_mxmat > max_mxmat:
+            msg = (
+                f"Applying cushion of {cushion} to {pftname}'s mxmat ({mxmat}) resulted in new"
+                f" mxmat of {new_mxmat}; decreasing that to max_mxmat {max_mxmat}"
+            )
+            new_mxmat = max_mxmat
+        if msg:
+            warnings.warn(msg, RuntimeWarning)
+
+        assert (
+            new_mxmat >= min_mxmat
+        ), f"{pftname} new_mxmat ({new_mxmat}) < min_mxmat ({min_mxmat})"
+        assert (
+            new_mxmat <= max_mxmat
+        ), f"{pftname} new_mxmat ({new_mxmat}) > max_mxmat ({max_mxmat})"
+
+        mxmat_dict_out[pftname] = new_mxmat
+
+    return mxmat_dict_out
+
+
 def unexpected_negative_rx_gdd(data_array):
     """
     Return True if there's a negative value not matching the designated missing value

python/ctsm/crop_calendars/generate_gdds.py

@@ -26,10 +26,20 @@
 # fixed. For now, we'll just disable the warning.
 # pylint: disable=too-many-positional-arguments
 
-# Global constants
-PARAMFILE_DIR = "/glade/campaign/cesm/cesmdata/cseg/inputdata/lnd/clm2/paramdata"
-MY_CLM_VER = 51
-MY_CLM_SUBVER = "c211112"
+
+def _get_max_growing_season_lengths(max_season_length_from_hdates_file, paramfile, cushion):
+    """
+    Import maximum growing season lengths from paramfile, if doing so.
+    """
+    if max_season_length_from_hdates_file:
+        return None
+
+    mxmats = cc.import_max_gs_length(paramfile)
+
+    if cushion:
+        mxmats = cc.cushion_gs_length(mxmats, cushion)
+
+    return mxmats
 
 
 def main(
@@ -47,18 +57,20 @@ def main(
     land_use_file=None,
     first_land_use_year=None,
     last_land_use_year=None,
-    unlimited_season_length=False,
+    max_season_length_from_hdates_file=False,
     skip_crops=None,
     logger=None,
     no_pickle=None,
+    paramfile=None,
+    max_season_length_cushion=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:
+            if not max_season_length_from_hdates_file:
                 output_dir += ".mxmat"
             output_dir += "." + dt.datetime.now().strftime("%Y-%m-%d-%H%M%S")
     if not os.path.exists(output_dir):
@@ -146,10 +158,9 @@ def main(
         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
+        mxmats = _get_max_growing_season_lengths(
+            max_season_length_from_hdates_file, paramfile, max_season_length_cushion
+        )
 
         h1_instantaneous = None
         for yr_index, this_yr in enumerate(np.arange(first_season + 1, last_season + 3)):
@@ -390,11 +401,34 @@ def add_attrs_to_map_ds(
         )
 
 
-if __name__ == "__main__":
-    ###############################
-    ### Process input arguments ###
-    ###############################
-    parser = argparse.ArgumentParser(description="ADD DESCRIPTION HERE")
+def _parse_args(argv):
+    parser = argparse.ArgumentParser(
+        description=(
+            "A script to generate maturity requirements for CLM crops in units of growing degree-"
+            "days (GDDs)."
+        )
+    )
+
+    # Required but mutually exclusive
+    max_growing_season_length_group = parser.add_mutually_exclusive_group(required=True)
+    max_growing_season_length_group.add_argument(
+        "--paramfile",
+        help=(
+            "Path to parameter file with maximum growing season lengths (mxmat)."
+            " Mutually exclusive with --max-season-length-from-hdates-file."
+        ),
+    )
+    max_growing_season_length_group.add_argument(
+        "--max-season-length-from-hdates-file",
+        help=(
+            "Rather than limiting growing season length based on mxmat values from a CLM parameter"
+            " file, use the season lengths from --hdates-file. Not recommended unless you use the"
+            "results of this script in a run with sufficiently long mxmat values!"
+            " Mutually exclusive with --paramfile."
+        ),
+        action="store_true",
+        default=False,
+    )
 
     # Required
     parser.add_argument(
@@ -467,12 +501,6 @@ def add_attrs_to_map_ds(
         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.",
@@ -485,12 +513,40 @@ def add_attrs_to_map_ds(
         action="store_true",
         default=False,
     )
+    parser.add_argument(
+        "--max-season-length-cushion",
+        help=(
+            "How much to reduce the maximum growing season length (mxmat) for each crop in the"
+            " parameter file. This might be useful for helping avoid high rates of immature"
+            " harvests for gridcells where the observed harvest date is longer than mxmat."
+            " Incompatible with --max-season-length-from-hdates-file."
+        ),
+        default=0,
+        type=int,
+    )
 
     # Get arguments
-    args = parser.parse_args(sys.argv[1:])
-    for k, v in sorted(vars(args).items()):
+    args_parsed = parser.parse_args(argv)
+    for k, v in sorted(vars(args_parsed).items()):
         print(f"{k}: {v}")
 
+    # Check arguments
+    if args_parsed.max_season_length_from_hdates_file and args_parsed.max_season_length_cushion:
+        raise argparse.ArgumentError(
+            None,
+            "--max-season-length-from-hdates-file is incompatible with --max-season-length-cushion"
+            " ≠ 0.",
+        )
+
+    return args_parsed
+
+
+if __name__ == "__main__":
+    ###############################
+    ### Process input arguments ###
+    ###############################
+    args = _parse_args(sys.argv[1:])
+
     # Call main()
     main(
         input_dir=args.input_dir,
@@ -506,7 +562,9 @@ def add_attrs_to_map_ds(
         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,
+        max_season_length_from_hdates_file=args.max_season_length_from_hdates_file,
         skip_crops=args.skip_crops,
         no_pickle=args.no_pickle,
+        paramfile=args.paramfile,
+        max_season_length_cushion=args.max_season_length_cushion,
     )

python/ctsm/crop_calendars/generate_gdds_functions.py

@@ -22,6 +22,9 @@
 # fixed. For now, we'll just disable the warning.
 # pylint: disable=too-many-positional-arguments
 
+# Tolerance (degrees) for checking lat/lon grid matches
+GRID_TOL_DEG = 1e-6
+
 CAN_PLOT = True
 try:
     # pylint: disable=wildcard-import,unused-wildcard-import
@@ -54,6 +57,31 @@
     CAN_PLOT = False
 
 
+def check_grid_match(grid0, grid1, tol=GRID_TOL_DEG):
+    """
+    Check whether latitude or longitude values match
+    """
+    if grid0.shape != grid1.shape:
+        return False, None
+
+    if hasattr(grid0, "values"):
+        grid0 = grid0.values
+    if hasattr(grid1, "values"):
+        grid1 = grid1.values
+
+    abs_diff = np.abs(grid1 - grid0)
+    if np.any(np.isnan(abs_diff)):
+        if np.any(np.isnan(grid0) != np.isnan(grid1)):
+            warnings.warn("NaN(s) in grid don't match", RuntimeWarning)
+            return False, None
+        warnings.warn("NaN(s) in grid", RuntimeWarning)
+
+    max_abs_diff = np.nanmax(abs_diff)
+    match = max_abs_diff < tol
+
+    return match, max_abs_diff
+
+
 def check_sdates(dates_ds, sdates_rx, outdir_figs, logger, verbose=False):
     """
     Checking that input and output sdates match
@@ -62,6 +90,29 @@ def check_sdates(dates_ds, sdates_rx, outdir_figs, logger, verbose=False):
 
     sdates_grid = grid_one_variable(dates_ds, "SDATES")
 
+    # In this script, we assume that you used prescribed sowing dates on the same grid as the
+    # CLM run
+    # Check that latitudes match, coercing a match if difference is acceptable
+    match, max_abs_diff = check_grid_match(sdates_rx["lat"], sdates_grid["lat"])
+    assert bool(
+        match
+    ), f"CLM lat grid doesn't match rx sdates's within {GRID_TOL_DEG}; max abs diff {max_abs_diff}"
+    if max_abs_diff > 0:
+        log(logger, f"Max lat abs diff: {max_abs_diff}. Coercing match.")
+        # sdates_grid comes from the CLM outputs, so we coerce the prescribed sowing date file
+        # coordinate to match, because we want the outputs of this script to have CLM's coordinates.
+        sdates_rx["lat"] = sdates_grid["lat"]
+    # Check that longitudes match, coercing a match if difference is acceptable
+    match, max_abs_diff = check_grid_match(sdates_rx["lon"], sdates_grid["lon"])
+    assert bool(
+        check_grid_match(sdates_rx["lon"], sdates_grid["lon"])
+    ), f"CLM lon grid doesn't match rx sdates's within {GRID_TOL_DEG}; max abs diff {max_abs_diff}"
+    if max_abs_diff > 0:
+        log(logger, f"Max lon abs diff: {max_abs_diff}. Coercing match.")
+        # sdates_grid comes from the CLM outputs, so we coerce the prescribed sowing date file
+        # coordinate to match, because we want the outputs of this script to have CLM's coordinates.
+        sdates_rx["lon"] = sdates_grid["lon"]
+
     all_ok = True
     any_found = False
     vegtypes_skipped = []
@@ -83,8 +134,16 @@ def check_sdates(dates_ds, sdates_rx, outdir_figs, logger, verbose=False):
         # Output
         out_map = sdates_grid.sel(ivt_str=vegtype_str).squeeze(drop=True)
 
-        # Check for differences
+        # Calculate differences
         diff_map = out_map - in_map
+        assert (
+            diff_map.shape == in_map.shape
+        ), f"Diff map shape {diff_map.shape} doesn't match in_map shape {in_map.shape}"
+        assert (
+            diff_map.shape == out_map.shape
+        ), f"Diff map shape {diff_map.shape} doesn't match out_map shape {out_map.shape}"
+
+        # Check for differences
         diff_map_notnan = diff_map.values[np.invert(np.isnan(diff_map.values))]
         if np.any(diff_map_notnan):
             log(logger, f"Difference(s) found in {vegtype_str}")
@@ -496,6 +555,29 @@ def import_and_process_1yr(
         "h", hdates_rx, incl_patches1d_itype_veg, mxsowings, logger
     )  # Yes, mxsowings even when importing harvests
 
+    # In this script, we assume that you have prescribed harvest dates on the same grid as the
+    # CLM run
+    # Check that latitudes match, coercing a match if difference is acceptable
+    match, max_abs_diff = check_grid_match(hdates_rx_orig["lat"], dates_incl_ds["lat"])
+    assert bool(
+        match
+    ), f"CLM lat grid doesn't match rx hdates's within {GRID_TOL_DEG}; max abs diff {max_abs_diff}"
+    if max_abs_diff > 0:
+        log(logger, f"Max lat abs diff: {max_abs_diff}. Coercing match.")
+        # dates_incl_ds comes from the CLM outputs, so we coerce the prescribed harvest date file
+        # coordinate to match, because we want the outputs of this script to have CLM's coordinates.
+        hdates_rx_orig["lat"] = dates_incl_ds["lat"]
+    # Check that longitudes match, coercing a match if difference is acceptable
+    match, max_abs_diff = check_grid_match(hdates_rx_orig["lon"], dates_incl_ds["lon"])
+    assert bool(
+        match
+    ), f"CLM lon grid doesn't match rx hdates's within {GRID_TOL_DEG}; max abs diff {max_abs_diff}"
+    if max_abs_diff > 0:
+        log(logger, f"Max lon abs diff: {max_abs_diff}. Coercing match.")
+        # dates_incl_ds comes from the CLM outputs, so we coerce the prescribed harvest date file
+        # coordinate to match, because we want the outputs of this script to have CLM's coordinates.
+        hdates_rx_orig["lon"] = dates_incl_ds["lon"]
+
     # Limit growing season to CLM max growing season length, if needed
     if mxmats and (imported_sdates or imported_hdates):
         print("   Limiting growing season length...")