Add pad_year function for xarray time slicing with years <1000 (#972)

Author: tomvothecoder

.github/workflows/build_workflow.yml

@@ -80,9 +80,9 @@ jobs:
         uses: conda-incubator/setup-miniconda@v3
         with:
           activate-environment: "e3sm_diags_ci"
+          environment-file: conda-env/ci.yml
           miniforge-variant: Miniforge3
           miniforge-version: latest
-          environment-file: conda-env/ci.yml
           channel-priority: strict
           auto-update-conda: true
           python-version: ${{ matrix.python-version }}

auxiliary_tools/debug/971-iso-8601/debug.py

@@ -0,0 +1,67 @@
+"""
+This script is designed to debug a specific issue in the E3SM Diagnostics tool
+(e3sm_diags) related to the tropical_subseasonal set. The bug occurs when
+processing year values less than 1000, resulting in the error:
+"ValueError: no ISO-8601 or cftime-string-like match for string: 1-01-01".
+This issue was reported in PR #971.
+
+The script replicates the behavior of the following command-line invocation
+of e3sm_diags:
+
+e3sm_diags tropical_subseasonal --no_viewer --reference_data_path '/lcrc/soft/climate/e3sm_diags_data/obs_for_e3sm_diags/time-series' --test_data_path '/lcrc/group/e3sm2/ac.xzheng/E3SMv3_dev/20250404.wcycl1850.ne120pg2_r025_RRSwISC6to18E3r5.test4.chrysalis//post/atm/180x360_traave/ts/daily/1yr' --results_dir '/lcrc/group/e3sm/public_html/diagnostic_output/ac.zhang40/tests/tropical_subseasonal_time_fix' --case_id 'wavenumber-frequency' --ref_timeseries_input --test_timeseries_input --run_type 'model_vs_obs' --sets 'tropical_subseasonal' --variables 'PRECT' --seasons 'ANN' 'DJF' 'MAM' 'JJA' 'SON' --regions '15S15N' --regrid_tool 'xesmf' --regrid_method 'conservative_normed' --multiprocessing --num_workers '32' --backend 'cartopy' --output_format 'png' --output_format_subplot 'pdf' --canvas_size_w '1212' --canvas_size_h '1628' --figsize '8.5' '11.0' --dpi '150' --arrows --test_name 'v3.HR_test4' --short_test_name 'v3.HR_test4' --test_colormap 'cet_rainbow.rgb' --ref_name 'IMERG_Daily' --reference_name 'IMERG Daily' --reference_colormap 'cet_rainbow.rgb' --diff_title 'percent difference' --diff_colormap 'diverging_bwr.rgb' --granulate 'variables' 'plevs' 'regions' --selectors 'sets' 'seasons' --test_start_yr 2 --test_end_yr 18 --ref_start_yr 2001 --ref_end_yr 2010
+
+The script uses the e3sm_diags Python API to configure and run the diagnostics
+with the same parameters as the command-line invocation. It is intended to
+help identify and resolve the issue with year values less than 1000.
+"""
+
+from e3sm_diags.parameter.core_parameter import CoreParameter
+from e3sm_diags.run import runner
+from e3sm_diags.parameter.tropical_subseasonal_parameter import TropicalSubseasonalParameter
+
+# Set up parameters
+param = CoreParameter()
+param.no_viewer = True
+param.reference_data_path = '/lcrc/soft/climate/e3sm_diags_data/obs_for_e3sm_diags/time-series'
+param.test_data_path = '/lcrc/group/e3sm2/ac.xzheng/E3SMv3_dev/20250404.wcycl1850.ne120pg2_r025_RRSwISC6to18E3r5.test4.chrysalis//post/atm/180x360_traave/ts/daily/1yr'
+param.results_dir = '/lcrc/group/e3sm/public_html/diagnostic_output/ac.tvo/tests/tropical_subseasonal_time_fix'
+param.case_id = 'wavenumber-frequency'
+param.ref_timeseries_input = True
+param.test_timeseries_input = True
+param.run_type = 'model_vs_obs'
+param.sets = ['tropical_subseasonal']
+param.variables = ['PRECT']
+param.seasons = ['ANN', 'DJF', 'MAM', 'JJA', 'SON']
+param.regions = ['15S15N']
+param.regrid_tool = 'xesmf'
+param.regrid_method = 'conservative_normed'
+param.multiprocessing = True
+param.num_workers = 32
+param.output_format_subplot = ['pdf']
+param.canvas_size_w = 1212
+param.canvas_size_h = 1628
+param.figsize = [8.5, 11.0]
+param.dpi = 150
+param.arrows = True
+param.short_test_name = 'v3.HR_test4'
+param.test_colormap = 'cet_rainbow.rgb'
+param.ref_name = 'IMERG_Daily'
+param.reference_name = 'IMERG Daily'
+param.reference_colormap = 'cet_rainbow.rgb'
+param.diff_title = 'percent difference'
+param.diff_colormap = 'diverging_bwr.rgb'
+param.granulate = ['variables', 'plevs', 'regions']
+param.selectors = ['sets', 'seasons']
+
+
+trop_param = TropicalSubseasonalParameter()
+trop_param.test_start_yr = 2
+trop_param.test_name = 'v3.HR_test4'
+
+trop_param.test_end_yr = 18
+trop_param.ref_start_yr = 2001
+trop_param.ref_end_yr = 2010
+
+# Run the diagnostics
+runner.sets_to_run = ['tropical_subseasonal']
+runner.run_diags([param, trop_param])

e3sm_diags/driver/tropical_subseasonal_driver.py

@@ -18,6 +18,7 @@
 from e3sm_diags.driver.utils import zwf_functions as wf
 from e3sm_diags.driver.utils.climo_xr import ClimoFreq
 from e3sm_diags.driver.utils.dataset_xr import Dataset
+from e3sm_diags.driver.utils.general import pad_year
 from e3sm_diags.logger import _setup_child_logger
 from e3sm_diags.plot.tropical_subseasonal_plot import plot
 
@@ -130,10 +131,15 @@ def calculate_spectrum(path, variable, start_year, end_year):
         "dosymmetries": True,
         "rmvLowFrq": True,
     }
-    # TODO the time subsetting and variable derivation should be replaced during cdat revamp
+    # TODO the time subsetting and variable derivation should be replaced during
+    # cdat revamp.
+
+    start_year_str = pad_year(start_year)
+    end_year_str = pad_year(end_year)
     try:
         var = xr.open_mfdataset(glob.glob(f"{path}/{variable}_*.nc")).sel(
-            lat=slice(-15, 15), time=slice(f"{start_year}-01-01", f"{end_year}-12-31")
+            lat=slice(-15, 15),
+            time=slice(f"{start_year_str}-01-01", f"{end_year_str}-12-31"),
         )[variable]
         actual_start = var.time.dt.year.values[0]
         actual_end = var.time.dt.year.values[-1]

e3sm_diags/driver/utils/dataset_xr.py

@@ -31,6 +31,7 @@
 )
 from e3sm_diags.driver import FRAC_REGION_KEYS, LAND_OCEAN_MASK_PATH
 from e3sm_diags.driver.utils.climo_xr import CLIMO_FREQS, ClimoFreq, climo
+from e3sm_diags.driver.utils.general import pad_year
 from e3sm_diags.driver.utils.regrid import HYBRID_SIGMA_KEYS
 from e3sm_diags.logger import _setup_child_logger
 
@@ -1197,13 +1198,14 @@ def _get_time_slice(self, ds: xr.Dataset) -> slice:
                 f"end_year ({end_yr_int}) > var_end_yr ({var_end_year})."
             )
 
-        start_yr_str = str(start_yr_int).zfill(4)
-        end_yr_str = str(end_yr_int).zfill(4)
+        start_yr_str = pad_year(start_yr_int)
+        end_yr_str = pad_year(end_yr_int)
 
         if self.is_sub_monthly:
             start_time = f"{start_yr_str}-01-01"
 
-            end_yr_str = str(int(end_yr_str) + 1).zfill(4)
+            new_end_year_int = int(end_yr_str) + 1
+            end_yr_str = pad_year(new_end_year_int)
             end_time = f"{end_yr_str}-01-01"
         else:
             start_time = self._get_slice_with_bounds(ds, start_yr_str, "start")

e3sm_diags/driver/utils/general.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 from e3sm_diags.logger import _setup_child_logger
 
 logger = _setup_child_logger(__name__)
@@ -15,3 +17,68 @@ def _monotonically_decreasing(L):
 def _monotonically_increasing(L):
     # FIXME: B905: zip() without an explicit strict= parameter
     return all(x <= y for x, y in zip(L, L[1:]))
+
+
+def pad_year(year: int | str) -> str:
+    """Pad the year with leading zeros to ensure it is 4 digits.
+
+    This function ensures that the input year is properly formatted as a
+    4-digit string, which is required for ISO-8601 date formats (YYYY-MM-DD).
+    If the year is less than 1000, it is padded with leading zeros.
+
+    Parameters
+    ----------
+    year : int or str
+        The year to pad. Must be a non-negative integer or a string representing
+        a non-negative integer. Floats are not allowed.
+
+    Returns
+    -------
+    str
+        The padded year as a 4-digit string (e.g., "0042" for year 42).
+
+    Raises
+    ------
+    ValueError
+        If the input year is not a non-negative integer or a string representing
+        a non-negative integer, or if it is a float, or if it is outside the
+        range 0 to 9999 inclusive.
+
+    Examples
+    --------
+    >>> pad_year(42)
+    '0042'
+    >>> pad_year("42")
+    '0042'
+    >>> pad_year(2023)
+    '2023'
+    >>> pad_year("2023")
+    '2023'
+    >>> pad_year(-1)
+    Traceback (most recent call last):
+        ...
+    ValueError: Year must be between 0 and 9999 inclusive.
+    >>> pad_year(10000)
+    Traceback (most recent call last):
+        ...
+    ValueError: Year must be between 0 and 9999 inclusive.
+    >>> pad_year(42.0)
+    Traceback (most recent call last):
+        ...
+    ValueError: Year must not be a float.
+    """
+    try:
+        if isinstance(year, float):
+            raise ValueError("Year must not be a float.")
+
+        year = int(year)
+    except (ValueError, TypeError) as e:
+        raise ValueError(
+            "Year must be a non-negative integer or a string representing a "
+            "non-negative integer."
+        ) from e
+
+    if year < 0 or year > 9999:
+        raise ValueError("Year must be between 0 and 9999 inclusive.")
+
+    return f"{year:04d}"

e3sm_diags/parameter/core_parameter.py

@@ -9,6 +9,7 @@
 
 from e3sm_diags.derivations.derivations import DerivedVariablesMap
 from e3sm_diags.driver.utils.climo_xr import ClimoFreq
+from e3sm_diags.driver.utils.general import pad_year
 from e3sm_diags.driver.utils.regrid import REGRID_TOOLS
 from e3sm_diags.logger import _setup_child_logger
 
@@ -42,6 +43,15 @@
     #    "tropical_subseasonal",
 ]
 
+YEAR_ATTRIBUTES = [
+    "start_yr",
+    "end_yr",
+    "test_start_yr",
+    "test_end_yr",
+    "ref_start_yr",
+    "ref_end_yr",
+]
+
 if TYPE_CHECKING:
     from e3sm_diags.driver.utils.dataset_xr import Dataset
 
@@ -346,3 +356,12 @@ def _run_diag(self) -> List[Any]:
                     sys.exit()
 
         return results
+
+    def __setattr__(self, name: str, value: Any) -> None:
+        """Override setattr to ensure year attributes are padded when set."""
+        if name in YEAR_ATTRIBUTES and value not in [None, ""]:
+            # Validate and pad the year before setting the attribute
+            value = pad_year(value)
+
+        # Set the attribute using the superclass method
+        super().__setattr__(name, value)

tests/e3sm_diags/driver/utils/test_dataset_xr.py

@@ -75,17 +75,17 @@ def _create_parameter_object(
             parameter.ref_timeseries_input = False
 
         parameter.reference_data_path = data_path
-        parameter.ref_start_yr = start_yr  # type: ignore
-        parameter.ref_end_yr = end_yr  # type: ignore
+        parameter.ref_start_yr = start_yr
+        parameter.ref_end_yr = end_yr
     elif dataset_type == "test":
         if data_type == "time_series":
             parameter.test_timeseries_input = True
         else:
             parameter.test_timeseries_input = False
 
         parameter.test_data_path = data_path
-        parameter.test_start_yr = start_yr  # type: ignore
-        parameter.test_end_yr = end_yr  # type: ignore
+        parameter.test_start_yr = start_yr
+        parameter.test_end_yr = end_yr
 
     return parameter
 

tests/e3sm_diags/driver/utils/test_general.py

@@ -0,0 +1,37 @@
+import pytest
+
+from e3sm_diags.driver.utils.general import pad_year
+
+
+class TestPadYear:
+    @pytest.mark.parametrize(
+        "input_year, expected_output",
+        [
+            (42, "0042"),
+            ("42", "0042"),
+            (2023, "2023"),
+            ("2023", "2023"),
+            (0, "0000"),
+            ("0", "0000"),
+            (9999, "9999"),
+            ("9999", "9999"),
+        ],
+    )
+    def test_valid_years(self, input_year, expected_output):
+        assert pad_year(input_year) == expected_output
+
+    @pytest.mark.parametrize(
+        "invalid_year",
+        [
+            -1,
+            "abcd",
+            None,
+            10000,
+            "10000",
+            1.5,
+            "1.5",
+        ],
+    )
+    def test_invalid_years(self, invalid_year):
+        with pytest.raises(ValueError):
+            pad_year(invalid_year)

tests/e3sm_diags/test_parameters.py

@@ -31,13 +31,13 @@ def test__add__copies_attributes_from_other_object(self):
         param2 = CoreParameter()
 
         # Add custom attributes to the second object
-        param2.test_start_yr = 2000  # type: ignore
-        param2.test_end_yr = 2001  # type: ignore
+        param2.test_start_yr = 2000
+        param2.test_end_yr = 2001
 
         new_param = param2 + param1
 
-        assert new_param.test_start_yr == 2000
-        assert new_param.test_end_yr == 2001
+        assert new_param.test_start_yr == "2000"
+        assert new_param.test_end_yr == "2001"
 
     def test_check_values_does_not_raise_error_if_required_args_are_set(self):
         param = CoreParameter()
@@ -105,6 +105,26 @@ def test_logs_error_if_driver_module_for_set_not_found(self, caplog):
             in caplog.text
         )
 
+    def test_year_properties_automatic_padding(self):
+        """Test that year properties are automatically zero-padded to 4 digits."""
+        param = CoreParameter()
+
+        # Test assigning various year formats
+        param.start_yr = "95"
+        param.end_yr = 2000
+        param.test_start_yr = "0"
+        param.test_end_yr = 95
+        param.ref_start_yr = "100"
+        param.ref_end_yr = "9999"
+
+        # Verify that all values have been properly padded
+        assert param.start_yr == "0095"  # type: ignore
+        assert param.end_yr == "2000"  # type: ignore
+        assert param.test_start_yr == "0000"  # type: ignore
+        assert param.test_end_yr == "0095"  # type: ignore
+        assert param.ref_start_yr == "0100"  # type: ignore
+        assert param.ref_end_yr == "9999"  # type: ignore
+
     @pytest.mark.xfail
     def test_logs_exception_if_driver_run_diag_function_fails(self, caplog):
         # TODO: Need to implement this test by raising an exception through