diff --git a/python/ctsm/crop_calendars/check_rx_obeyed.py b/python/ctsm/crop_calendars/check_rx_obeyed.py index c8c5410faf..68a37fb146 100644 --- a/python/ctsm/crop_calendars/check_rx_obeyed.py +++ b/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, diff --git a/python/ctsm/crop_calendars/cropcal_utils.py b/python/ctsm/crop_calendars/cropcal_utils.py index c7e8b6ac52..4f2929de37 100644 --- a/python/ctsm/crop_calendars/cropcal_utils.py +++ b/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 diff --git a/python/ctsm/crop_calendars/generate_gdd20_baseline.py b/python/ctsm/crop_calendars/generate_gdd20_baseline.py index ee942eba21..fa42eb58b4 100644 --- a/python/ctsm/crop_calendars/generate_gdd20_baseline.py +++ b/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], ) diff --git a/python/ctsm/crop_calendars/generate_gdds_functions.py b/python/ctsm/crop_calendars/generate_gdds_functions.py index be724659dd..ec60a85beb 100644 --- a/python/ctsm/crop_calendars/generate_gdds_functions.py +++ b/python/ctsm/crop_calendars/generate_gdds_functions.py @@ -591,7 +591,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 @@ -1152,7 +1152,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: diff --git a/python/ctsm/crop_calendars/grid_one_variable.py b/python/ctsm/crop_calendars/grid_one_variable.py index d7c7126e54..c06ed50b16 100644 --- a/python/ctsm/crop_calendars/grid_one_variable.py +++ b/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. diff --git a/python/ctsm/test/test_unit_cropcal_utils.py b/python/ctsm/test/test_unit_cropcal_utils.py new file mode 100755 index 0000000000..edbb81bd7b --- /dev/null +++ b/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() diff --git a/python/ctsm/test/test_unit_grid_one_variable.py b/python/ctsm/test/test_unit_grid_one_variable.py new file mode 100755 index 0000000000..400b4ede7e --- /dev/null +++ b/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. 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()