generate_gdds bugfix: Use correct h2 years.

Author: samsrabin

python/ctsm/crop_calendars/generate_gdds.py

@@ -30,6 +30,9 @@
 # fixed. For now, we'll just disable the warning.
 # pylint: disable=too-many-positional-arguments
 
+# Mapping of history tape number to output frequency
+H_FREQ_DICT = {1: "annual", 2: "daily"}
+
 
 def _get_max_growing_season_lengths(max_season_length_from_hdates_file, paramfile, cushion):
     """
@@ -46,19 +49,31 @@ def _get_max_growing_season_lengths(max_season_length_from_hdates_file, paramfil
     return mxmats
 
 
-def _get_history_yr_range(first_season, last_season):
+def _get_history_yr_range(first_season, last_season, freq):
     """
-    Get a range object that can be used for looping over all years we need to process timestamps
-    from.
+    Get range objects that can be used for looping over all years we need to process timestamps
+    from. Different handling for annual vs. daily history files. Assumption is that all history
+    files are instantanous.
     """
-    # Saving at the end of a year receive the timestamp of the END of the year's final timestep,
-    # which means it will actually be 00:00 of Jan. 1 of the next year.
-    first_history_yr = first_season + 1
 
-    # Same deal for the last history timestep, but we have to read an extra year in that case,
-    # because in some places the last growing season won't complete until the year after it was
-    # planted.
-    last_history_yr = last_season + 2
+    if freq == "annual":
+        # Saving at the end of a year gives the timestamp of the END of the year's final timestep,
+        # which means it will actually be 00:00 of Jan. 1 of the next year.
+        first_history_yr = first_season + 1
+        last_history_yr = last_season + 1
+    elif freq == "daily":
+        # Saving at the end of a day/beginning of the next day (i.e., 00:00:00) means that the year
+        # will be correct for all but the Dec. 31 save, which will receive a timestamp of Jan. 1
+        # 00:00:00. That will be handled in _get_time_slice_lists(), so we don't need to account for
+        # it here.
+        first_history_yr = first_season
+        last_history_yr = last_season
+    else:
+        raise NotImplementedError(f"Not sure how to handle freq '{freq}'")
+
+    # For the last season, we have to read an extra year, because in some places the last growing
+    # season won't complete until the year after it was planted.
+    last_history_yr += 1
 
     # last_history_yr + 1 because range() will iterate up to but not including the second value.
     history_yr_range = range(first_history_yr, last_history_yr + 1)
@@ -70,7 +85,7 @@ def _get_time_slice_lists(first_season, last_season):
     """
     Given the requested first and last seasons, get the list of time slices that the script should
     look for. The assumption here, as in _get_file_lists() and as instructed in the docs, is
-    that the user is saving instantaneous files.
+    that the user is saving instantaneous tapes.
     """
 
     # Input checks
@@ -79,10 +94,15 @@ def _get_time_slice_lists(first_season, last_season):
     if first_season > last_season:
         raise ValueError(f"first_season ({first_season}) > last_season ({last_season})")
 
-    slice_lists_list = [None, None]
-    for i, h in enumerate([1, 2]):
+    # Initialize list with None for each history tape. Could avoid by starting with empty list and
+    # doing .append(), but pylint gets confused by that for some reason.
+    slice_lists_list = [None for x in range(len(H_FREQ_DICT))]
+
+    # Get time slice for each required history year in each history tape.
+    for i, h in enumerate(list(H_FREQ_DICT.keys())):
         slice_list = []
-        for history_yr in _get_history_yr_range(first_season, last_season):
+        freq = H_FREQ_DICT[h]
+        for history_yr in _get_history_yr_range(first_season, last_season, freq):
             if h == 1:
                 # Annual timesteps
                 slice_start = f"{history_yr}-01-01"
@@ -98,7 +118,10 @@ def _get_time_slice_lists(first_season, last_season):
 
         # We should be reading one more than the total number of years in
         # [first_season, last_season].
-        assert len(slice_list) == last_season - first_season + 2
+        ns_exp = last_season - first_season + 2
+        ns_actual = len(slice_list)
+        msg = f"Expected {ns_exp} time slices in list for h{h}; got {ns_actual}"
+        assert ns_exp == ns_actual, msg
 
         # Save
         slice_lists_list[i] = slice_list
@@ -244,26 +267,29 @@ def main(
         h1_file_lists, h2_file_lists = _get_file_lists(
             input_dir, (h1_time_slices, h2_time_slices), logger
         )
-        history_yr_range = _get_history_yr_range(first_season, last_season)
+        history_yr_range_h1 = _get_history_yr_range(first_season, last_season, "annual")
+        history_yr_range_h2 = _get_history_yr_range(first_season, last_season, "daily")
         # Check
+        assert len(history_yr_range_h1) == len(history_yr_range_h2)
         log(logger, "Checking h1 files")
-        gddfn.check_file_lists(history_yr_range, h1_file_lists, h1_time_slices, "annual", logger)
+        gddfn.check_file_lists(history_yr_range_h1, h1_file_lists, h1_time_slices, "annual", logger)
         log(logger, "Checking h2 files")
-        gddfn.check_file_lists(history_yr_range, h2_file_lists, h2_time_slices, "daily", logger)
+        gddfn.check_file_lists(history_yr_range_h2, h2_file_lists, h2_time_slices, "daily", logger)
         log(logger, "Done")
 
-        for y, history_yr in enumerate(history_yr_range):
+        for y, history_yr_h1 in enumerate(history_yr_range_h1):
             # If resuming from a pickled file, we continue until we reach a year that hasn't yet
             # been processed.
-            if history_yr <= pickle_season:
+            if history_yr_h1 <= pickle_season:
                 continue
-            log(logger, f"History year {history_yr}...")
+            log(logger, f"History year {history_yr_h1}...")
 
             # Get time slice and files to read for this year
             h1_time_slice = h1_time_slices[y]  # pylint: disable=unsubscriptable-object
             h2_time_slice = h2_time_slices[y]  # pylint: disable=unsubscriptable-object
             h1_file_list = h1_file_lists[y]  # pylint: disable=unsubscriptable-object
             h2_file_list = h2_file_lists[y]  # pylint: disable=unsubscriptable-object
+            history_yr_h2 = list(history_yr_range_h2)[y]
 
             (
                 h2_ds,
@@ -293,7 +319,8 @@ def main(
                 skip_crops,
                 outdir_figs,
                 logger,
-                history_yr,
+                history_yr_h1,
+                history_yr_h2,
                 h1_file_list,
                 h2_file_list,
                 h1_time_slice,
@@ -306,7 +333,7 @@ def main(
                     [
                         first_season,
                         last_season,
-                        history_yr,
+                        history_yr_h1,
                         gddaccum_yp_list,
                         gddharv_yp_list,
                         skip_patches_for_isel_nan_lastyear,

python/ctsm/crop_calendars/generate_gdds_functions.py

@@ -371,7 +371,8 @@ def import_and_process_1yr(
     skip_crops,
     outdir_figs,
     logger,
-    history_yr,
+    history_yr_h1,
+    history_yr_h2,
     h1_filelist,
     h2_filelist,
     h1_time_slice,
@@ -405,7 +406,7 @@ def import_and_process_1yr(
     )
 
     # Check included timesteps
-    _check_time_da("annual", history_yr, dates_ds["time"], logger)
+    _check_time_da("annual", history_yr_h1, dates_ds["time"], logger)
 
     # Should now just be one timestep, so select it to remove dimension.
     dates_ds = dates_ds.isel(time=0)
@@ -651,7 +652,7 @@ def import_and_process_1yr(
     )
 
     # Check included timesteps
-    _check_time_da("daily", history_yr, h2_ds["time"], logger)
+    _check_time_da("daily", history_yr_h2, h2_ds["time"], logger)
 
     # Restrict to patches we're including
     if skipping_patches_for_isel_nan:

python/ctsm/test/test_unit_generate_gdds.py

@@ -253,10 +253,11 @@ def test_generate_gdds_get_time_slice_lists(self):
         self.assertEqual(h1_slices, expected_h1)
 
         # Check h2 slices (daily timesteps - full year)
+        # For daily, starts at first_season (1986), not first_season + 1
         expected_h2 = [
+            slice("1986-01-02", "1987-01-01"),
             slice("1987-01-02", "1988-01-01"),
             slice("1988-01-02", "1989-01-01"),
-            slice("1989-01-02", "1990-01-01"),
         ]
         self.assertEqual(h2_slices, expected_h2)
 
@@ -272,9 +273,10 @@ def test_generate_gdds_get_time_slice_lists_1yr(self):
         self.assertEqual(h1_slices, expected_h1)
 
         # Check h2 slices
+        # For daily, starts at first_season (1987), not first_season + 1
         expected_h2 = [
+            slice("1987-01-02", "1988-01-01"),
             slice("1988-01-02", "1989-01-01"),
-            slice("1989-01-02", "1990-01-01"),
         ]
         self.assertEqual(h2_slices, expected_h2)
 
@@ -324,6 +326,51 @@ def test_generate_gdds_get_time_slice_lists_h2_year_long(self):
             self.assertEqual(stop_year, start_year + 1)
 
 
+class TestGetHistoryYrRange(unittest.TestCase):
+    """Tests for _get_history_yr_range()"""
+
+    def test_get_history_yr_range_annual(self):
+        """Test _get_history_yr_range with annual frequency"""
+        result = gg._get_history_yr_range(1986, 1987, "annual")
+        # For annual: first_season + 1 through last_season + 2
+        expected = range(1987, 1990)
+        self.assertEqual(result, expected)
+
+    def test_get_history_yr_range_daily(self):
+        """Test _get_history_yr_range with daily frequency"""
+        result = gg._get_history_yr_range(1986, 1987, "daily")
+        # For daily: first_season through last_season + 1
+        expected = range(1986, 1989)
+        self.assertEqual(result, expected)
+
+    def test_get_history_yr_range_annual_single_year(self):
+        """Test _get_history_yr_range with annual frequency and single year"""
+        result = gg._get_history_yr_range(2000, 2000, "annual")
+        # Should give 2001, 2002
+        expected = range(2001, 2003)
+        self.assertEqual(result, expected)
+
+    def test_get_history_yr_range_daily_single_year(self):
+        """Test _get_history_yr_range with daily frequency and single year"""
+        result = gg._get_history_yr_range(2000, 2000, "daily")
+        # Should give 2000, 2001
+        expected = range(2000, 2002)
+        self.assertEqual(result, expected)
+
+    def test_get_history_yr_range_unknown_freq(self):
+        """Test _get_history_yr_range with unknown frequency"""
+        with self.assertRaises(NotImplementedError):
+            gg._get_history_yr_range(2000, 2001, "monthly")
+
+    def test_get_history_yr_range_lengths_match(self):
+        """Test that annual and daily ranges have the same length"""
+        annual_range = gg._get_history_yr_range(2000, 2005, "annual")
+        daily_range = gg._get_history_yr_range(2000, 2005, "daily")
+        self.assertEqual(len(annual_range), len(daily_range))
+        # Should be last_season - first_season + 2
+        self.assertEqual(len(annual_range), 2005 - 2000 + 2)
+
+
 class TestCheckGridMatch(unittest.TestCase):
     """Tests check_grid_match()"""
 
@@ -617,14 +664,14 @@ def _create_test_file(self, filename):
     def test_get_file_lists_single_year(self):
         """Test _get_file_lists with a single year of data"""
         # Create h1 and h2 files for 2000 and 2001
-        # (first_season=1999, last_season=1999 will request slices for both years)
+        # Also need h2 file for 1999 since daily starts at first_season
+        h2_file_1999 = self._create_test_file("test.clm2.h2i.1999-01-02-00000.nc")
         h1_file_2000 = self._create_test_file("test.clm2.h1i.2000-01-01-00000.nc")
         h2_file_2000 = self._create_test_file("test.clm2.h2i.2000-01-02-00000.nc")
         h1_file_2001 = self._create_test_file("test.clm2.h1i.2001-01-01-00000.nc")
-        h2_file_2001 = self._create_test_file("test.clm2.h2i.2001-01-02-00000.nc")
 
         # Get time slice lists for first_season=1999, last_season=1999
-        # This will give us slices for 2000 and 2001
+        # This will give us slices for 2000 and 2001 (h1), and 1999, 2000, 2001 (h2)
         time_slice_lists_list = gg._get_time_slice_lists(1999, 1999)
 
         h1_file_lists, h2_file_lists = gg._get_file_lists(
@@ -640,23 +687,24 @@ def test_get_file_lists_single_year(self):
         self.assertEqual(len(h1_file_lists[0]), 1)
         self.assertEqual(h1_file_lists[0], [h1_file_2000])
         self.assertEqual(len(h2_file_lists[0]), 1)
-        self.assertEqual(h2_file_lists[0], [h2_file_2000])
+        self.assertEqual(h2_file_lists[0], [h2_file_1999])
         self.assertEqual(len(h1_file_lists[1]), 1)
         self.assertEqual(h1_file_lists[1], [h1_file_2001])
         self.assertEqual(len(h2_file_lists[1]), 1)
-        self.assertEqual(h2_file_lists[1], [h2_file_2001])
+        self.assertEqual(h2_file_lists[1], [h2_file_2000])
 
     def test_get_file_lists_multiple_years(self):
         """Test _get_file_lists with multiple years of data"""
         # Create h1 and h2 files for 2000-2002
+        # Also need h2 file for 1999 since daily starts at first_season
         h1_files = []
-        h2_files = []
+        h2_files = [self._create_test_file("test.clm2.h2i.1999-01-02-00000.nc")]
         for year in [2000, 2001, 2002]:
             h1_files.append(self._create_test_file(f"test.clm2.h1i.{year}-01-01-00000.nc"))
             h2_files.append(self._create_test_file(f"test.clm2.h2i.{year}-01-02-00000.nc"))
 
         # Get time slice lists for first_season=1999, last_season=2000
-        # This will give us slices for 2000, 2001, 2002
+        # This will give us slices for 2000, 2001, 2002 (h1) and 1999, 2000, 2001 (h2)
         time_slice_lists_list = gg._get_time_slice_lists(1999, 2000)
 
         h1_file_lists, h2_file_lists = gg._get_file_lists(
@@ -681,14 +729,15 @@ def test_get_file_lists_multiple_files_per_slice(self):
         h1_file_2000 = self._create_test_file("test.clm2.h1i.2000-01-01-00000.nc")
         h1_file_2001 = self._create_test_file("test.clm2.h1i.2001-01-01-00000.nc")
 
-        # Create multiple h2 files for 2000 (daily throughout the year)
+        # Create multiple h2 files for 1999 and 2000 (daily throughout the year)
+        h2_files_1999 = []
+        for month in ["01", "06", "12"]:
+            h2_files_1999.append(self._create_test_file(f"test.clm2.h2i.1999-{month}-15-00000.nc"))
+
         h2_files_2000 = []
         for month in ["01", "06", "12"]:
             h2_files_2000.append(self._create_test_file(f"test.clm2.h2i.2000-{month}-15-00000.nc"))
 
-        # Create h2 file for 2001
-        h2_file_2001 = self._create_test_file("test.clm2.h2i.2001-01-02-00000.nc")
-
         # Get time slice lists for first_season=1999, last_season=1999
         time_slice_lists_list = gg._get_time_slice_lists(1999, 1999)
 
@@ -700,18 +749,18 @@ def test_get_file_lists_multiple_files_per_slice(self):
         self.assertEqual(len(h1_file_lists), 2)
         self.assertEqual(len(h2_file_lists), 2)
 
-        # Check contents of file lists for first year (2000)
+        # Check contents of file lists
         # pylint: disable=unsubscriptable-object
         self.assertEqual(len(h1_file_lists[0]), 1)
         self.assertEqual(h1_file_lists[0], [h1_file_2000])
         self.assertEqual(len(h2_file_lists[0]), 3)
-        self.assertEqual(h2_file_lists[0], sorted(h2_files_2000))
+        self.assertEqual(h2_file_lists[0], sorted(h2_files_1999))
 
         # Check second year (2001)
         self.assertEqual(len(h1_file_lists[1]), 1)
         self.assertEqual(h1_file_lists[1], [h1_file_2001])
-        self.assertEqual(len(h2_file_lists[1]), 1)
-        self.assertEqual(h2_file_lists[1], [h2_file_2001])
+        self.assertEqual(len(h2_file_lists[1]), 3)
+        self.assertEqual(h2_file_lists[1], sorted(h2_files_2000))
 
     def test_get_file_lists_no_h1_files(self):
         """Test _get_file_lists when h1 files are missing"""