Fix generate_gdd20_baseline

python/ctsm/crop_calendars/check_rx_obeyed.py

@@ -144,7 +144,7 @@ def check_rx_obeyed(
             continue
         ds_thisveg = dates_ds.isel(patch=thisveg_patches)
 
-        vegtype_int = utils.vegtype_str2int(vegtype_str)[0]
+        vegtype_int = utils.vegtype_str2int(vegtype_str)
         rx_da = rx_ds[f"gs1_{vegtype_int}"]
         rx_array = rx_da.values[
             ds_thisveg.patches1d_jxy.values.astype(int) - 1,

python/ctsm/crop_calendars/cropcal_utils.py

@@ -265,6 +265,8 @@ def vegtype_str2int(vegtype_str, vegtype_mainlist=None):
     convert_to_ndarray = not isinstance(vegtype_str, np.ndarray)
     if convert_to_ndarray:
         vegtype_str = np.array(vegtype_str)
+    was_0d = vegtype_str.ndim == 0
+    vegtype_str = np.atleast_1d(vegtype_str)
 
     if isinstance(vegtype_mainlist, xr.Dataset):
         vegtype_mainlist = vegtype_mainlist.vegtype_str.values
@@ -289,6 +291,10 @@ def vegtype_str2int(vegtype_str, vegtype_mainlist=None):
         indices[np.where(vegtype_str == vegtype_str_2)] = vegtype_mainlist.index(vegtype_str_2)
     if convert_to_ndarray:
         indices = [int(x) for x in indices]
+
+    if was_0d:
+        indices = indices[0]
+
     return indices
 
 

python/ctsm/crop_calendars/generate_gdd20_baseline.py

@@ -109,22 +109,7 @@ def _parse_args():
         if not os.path.exists(filename):
             raise FileNotFoundError(f"Input file not found: {filename}")
 
-    # Process time slice
-    # Assumes CESM behavior where data for e.g. 1987 is saved as 1988-01-01.
-    # It would be more robust, accounting for upcoming behavior (where timestamp for a year is the
-    # middle of that year), to do slice("YEAR1-01-03", "YEARN-01-02"), but that's not compatible
-    # with ctsm_pylib as of the version using python 3.7.9. See safer_timeslice() in cropcal_utils.
-    if args.first_year is not None:
-        date_1 = f"{args.first_year+1}-01-01"
-    else:
-        date_1 = "0000-01-01"
-    if args.last_year is not None:
-        date_n = f"{args.last_year+1}-01-01"
-    else:
-        date_n = "9999-12-31"
-    time_slice = slice(date_1, date_n)
-
-    return args, time_slice
+    return args
 
 
 def _get_cft_list(crop_list):
@@ -182,7 +167,7 @@ def _get_gddn_for_cft(cft_str, variable):
 
 
 def _get_output_varname(cft_str):
-    cft_int = utils.vegtype_str2int(cft_str)[0]
+    cft_int = utils.vegtype_str2int(cft_str)
     return f"gdd20bl_{cft_int}"
 
 
@@ -232,7 +217,23 @@ def setup_output_dataset(input_files, author, variable, year_args, ds_in):
     return ds_out
 
 
-def generate_gdd20_baseline(input_files, output_file, author, time_slice, variable, year_args):
+def _get_time_slice(year_args):
+    """
+    Based on years from input arguments, return a time slice for selecting from dataset
+    """
+    first_year = year_args[0]
+    last_year = year_args[1]
+    date_1 = f"{first_year}-01-01"
+    date_n = f"{last_year}-12-31"
+    if first_year is None:
+        date_1 = "0000-01-01"
+    if last_year is None:
+        date_n = "9999-12-31"
+    time_slice = slice(date_1, date_n)
+    return time_slice
+
+
+def generate_gdd20_baseline(input_files, output_file, author, variable, year_args):
     """
     Generate stream_fldFileName_gdd20_baseline file from CTSM outputs
     """
@@ -252,6 +253,9 @@ def generate_gdd20_baseline(input_files, output_file, author, time_slice, variab
     input_files = list(set(input_files))
     input_files.sort()
 
+    # Process time slice
+    time_slice = _get_time_slice(year_args)
+
     # Import history files and ensure they have lat/lon dims
     ds_in = import_ds(input_files, my_vars=var_list_in + GRIDDING_VAR_LIST, time_slice=time_slice)
     if not all(x in ds_in.dims for x in ["lat", "lon"]):
@@ -275,7 +279,7 @@ def generate_gdd20_baseline(input_files, output_file, author, time_slice, variab
     # Process all crops
     encoding_dict = {}
     for cft_str in MGDCROP_LIST:
-        cft_int = utils.vegtype_str2int(cft_str)[0]
+        cft_int = utils.vegtype_str2int(cft_str)
         print(f"{cft_str} ({cft_int})")
 
         # Which GDDN history variable does this crop use? E.g., GDD0, GDD10
@@ -323,12 +327,11 @@ def main():
     """
     main() function for calling generate_gdd20_baseline.py from command line.
     """
-    args, time_slice = _parse_args()
+    args = _parse_args()
     generate_gdd20_baseline(
         args.input_files,
         args.output_file,
         args.author,
-        time_slice,
         args.variable,
         [args.first_year, args.last_year],
     )

python/ctsm/crop_calendars/generate_gdds_functions.py

@@ -607,7 +607,7 @@ def import_and_process_1yr(
             log(logger, f"      SKIPPING {vegtype_str}")
             continue
 
-        vegtype_int = utils.vegtype_str2int(vegtype_str)[0]
+        vegtype_int = utils.vegtype_str2int(vegtype_str)
         this_crop_full_patchlist = list(xr_flexsel(h2_ds, vegtype=vegtype_str).patch.values)
 
         # Get time series for each patch of this type
@@ -1166,7 +1166,7 @@ def make_figures(
                 raise RuntimeError(f"If mapping {vegtype_str}, you must provide land use dataset")
             else:
                 vegtypes_str = [x for x in incl_vegtypes_str if vegtype_str.lower() in x]
-            vegtypes_int = [utils.vegtype_str2int(x)[0] for x in vegtypes_str]
+            vegtypes_int = [utils.vegtype_str2int(x) for x in vegtypes_str]
 
             # Crop fraction map (for masking and weighting)
             if lu_ds:

python/ctsm/crop_calendars/grid_one_variable.py

@@ -87,6 +87,10 @@ def create_filled_array(this_ds, fill_value, thisvar_da, new_dims):
     """
     Create a Numpy array to be filled with gridded data
     """
+
+    if fill_value is None:
+        fill_value = np.nan
+
     dim_size_list = []
     for dim in new_dims:
         if dim == "ivt_str":
@@ -97,10 +101,7 @@ def create_filled_array(this_ds, fill_value, thisvar_da, new_dims):
             dim_size = this_ds.sizes[dim]
         dim_size_list = dim_size_list + [dim_size]
     thisvar_gridded = np.empty(dim_size_list)
-    if fill_value:
-        thisvar_gridded[:] = fill_value
-    else:
-        thisvar_gridded[:] = np.NaN
+    thisvar_gridded[:] = fill_value
     return thisvar_gridded
 
 
@@ -160,7 +161,7 @@ def grid_one_variable(this_ds, var, fill_value=None, **kwargs):
     # Get DataArrays needed for gridding
     thisvar_da, vt_da, spatial_unit, ixy_da, jxy_da = get_ixy_jxy_das(this_ds, var)
 
-    if not fill_value and "_FillValue" in thisvar_da.attrs:
+    if fill_value is None and "_FillValue" in thisvar_da.attrs:
         fill_value = thisvar_da.attrs["_FillValue"]
 
     # Renumber vt_da to work as indices on new ivt dimension, if needed.

python/ctsm/test/test_unit_cropcal_utils.py

@@ -0,0 +1,38 @@
+#!/usr/bin/env python3
+
+"""Unit tests for cropcal_utils.py"""
+
+import unittest
+
+from ctsm import unit_testing
+from ctsm.crop_calendars import cropcal_utils as ccu
+
+# Allow names that pylint doesn't like, because otherwise I find it hard
+# to make readable unit test names
+# pylint: disable=invalid-name
+
+
+class TestCropCalUtils(unittest.TestCase):
+    """Tests of cropcal_utils.py"""
+
+    def setUp(self):
+        self.vegtype_mainlist = ["crop_1", "crop_2", "crop_3"]
+
+    def test_vegtype_str2int_1string(self):
+        """
+        Tests vegtype_str2int() for a single string. Result should be an int.
+        """
+        result = ccu.vegtype_str2int("crop_1", vegtype_mainlist=self.vegtype_mainlist)
+        self.assertEqual(result, 0)
+
+    def test_vegtype_str2int_2strings(self):
+        """
+        Tests vegtype_str2int() for two strings. result should be a list of ints.
+        """
+        result = ccu.vegtype_str2int(["crop_1", "crop_3"], vegtype_mainlist=self.vegtype_mainlist)
+        self.assertListEqual(result, [0, 2])
+
+
+if __name__ == "__main__":
+    unit_testing.setup_for_tests()
+    unittest.main()

python/ctsm/test/test_unit_grid_one_variable.py

@@ -0,0 +1,91 @@
+#!/usr/bin/env python3
+
+"""
+Unit tests for grid_one_variable
+"""
+
+import unittest
+
+import numpy as np
+import xarray as xr
+
+from ctsm import unit_testing
+from ctsm.crop_calendars import grid_one_variable as g1v
+
+# Allow test names that pylint doesn't like; otherwise hard to make them
+# readable
+# pylint: disable=invalid-name
+
+# pylint: disable=protected-access
+
+## Too many instant variables as part of the class (too many self.<varible> in the SetUp)
+# pylint: disable=too-many-instance-attributes
+
+
+class TestCreateFilledArray(unittest.TestCase):
+    """Unit tests for create_filled_array"""
+
+    def setUp(self):
+        # Set up this_ds, which will provide us with sizes of dimensions in most cases
+        lat_vals = [55.0, 56.0, 57.0]
+        lat_da = xr.DataArray(
+            data=lat_vals,
+            dims=["lat"],
+            coords={"lat": lat_vals},
+        )
+        lon_vals = [255.0, 256.0, 257.0]
+        lon_da = xr.DataArray(
+            data=lon_vals,
+            dims=["lon"],
+            coords={"lon": lon_vals},
+        )
+        self.this_ds = xr.Dataset(
+            data_vars={
+                "lat": lat_da,
+                "lon": lon_da,
+            }
+        )
+
+    def test_create_filled_array_fillNone(self):
+        """
+        Test create_filled_array() with fill_value None: Should be filled with NaN
+        """
+
+        fill_value = None
+        thisvar_da_dummy = xr.DataArray()
+        new_dims = ["lat", "lon"]
+
+        result = g1v.create_filled_array(self.this_ds, fill_value, thisvar_da_dummy, new_dims)
+
+        self.assertTrue(np.all(np.isnan(result)))
+
+    def test_create_filled_array_fill1(self):
+        """
+        Test create_filled_array() with fill_value 1: Should be filled with 1
+        """
+
+        fill_value = 1.0
+        thisvar_da_dummy = xr.DataArray()
+        new_dims = ["lat", "lon"]
+
+        result = g1v.create_filled_array(self.this_ds, fill_value, thisvar_da_dummy, new_dims)
+
+        self.assertTrue(np.all(result == fill_value))
+
+    def test_create_filled_array_fill0(self):
+        """
+        Test create_filled_array() with fill_value 0: Should be filled with 0
+        """
+
+        fill_value = 0.0
+        thisvar_da_dummy = xr.DataArray()
+        new_dims = ["lat", "lon"]
+
+        result = g1v.create_filled_array(self.this_ds, fill_value, thisvar_da_dummy, new_dims)
+
+        self.assertTrue(np.all(result == fill_value))
+
+
+if __name__ == "__main__":
+    unit_testing.setup_for_tests()
+    unittest.main()