Refactor implementation of I/O for retrieving time-sliced datasets

- Refactor drivers for polar, meriodional_mean_2d, zonal_mean_2d, zonal_mean_xy to use `driver.utils.io._get_xarray_datasets()`
- Add `driver.io._get_xarray_datasets()` utility to simplify fetching of xarray datasets based on time selection
- Update references to `season` to `time_selection` with `TimeSelection` annotation
- Move `_set_time_slice_name_yrs_attrs()` back to `LatLonNativeParameter` because it is only used there -- consider refactoring later
- Remove `time_slice.py` as this functions were converted to `CoreParameter` methods

Author: tomvothecoder

e3sm_diags/driver/lat_lon_driver.py

@@ -5,7 +5,6 @@
 import xarray as xr
 
 from e3sm_diags.driver import METRICS_DEFAULT_VALUE
-from e3sm_diags.driver.utils.climo_xr import ClimoFreq
 from e3sm_diags.driver.utils.dataset_xr import Dataset
 from e3sm_diags.driver.utils.io import _save_data_metrics_and_plots
 from e3sm_diags.driver.utils.regrid import (
@@ -16,15 +15,14 @@
     regrid_z_axis_to_plevs,
     subset_and_align_datasets,
 )
-from e3sm_diags.driver.utils.time_slice import check_time_selection
-from e3sm_diags.driver.utils.type_annotations import MetricsDict
 from e3sm_diags.logger import _setup_child_logger
 from e3sm_diags.metrics.metrics import correlation, rmse, spatial_avg, std
 from e3sm_diags.plot.lat_lon_plot import plot as plot_func
 
 logger = _setup_child_logger(__name__)
 
 if TYPE_CHECKING:
+    from e3sm_diags.driver.utils.type_annotations import MetricsDict, TimeSelection
     from e3sm_diags.parameter.core_parameter import CoreParameter
 
 
@@ -51,57 +49,44 @@ def run_diag(parameter: CoreParameter) -> CoreParameter:
         (e.g., one is 2-D and the other is 3-D).
     """
     variables = parameter.variables
-    seasons = parameter.seasons
-    time_slices = getattr(parameter, "time_slices", [])
     ref_name = getattr(parameter, "ref_name", "")
     regions = parameter.regions
 
-    # Check that either seasons or time_slices is specified, but not both
-    has_seasons, has_time_slices = check_time_selection(
-        seasons, time_slices, require_one=True
-    )
-
-    # Determine which time selection to use
-    time_selections = time_slices if has_time_slices else seasons
-    is_time_slice_mode = has_time_slices
-
     # Variables storing xarray `Dataset` objects start with `ds_` and
     # variables storing e3sm_diags `Dataset` objects end with `_ds`. This
     # is to help distinguish both objects from each other.
     test_ds = Dataset(parameter, data_type="test")
     ref_ds = Dataset(parameter, data_type="ref")
 
+    time_selection_type, time_selections = parameter._get_time_selection_to_use()
+
     for var_key in variables:
         logger.info("Variable: {}".format(var_key))
         parameter.var_id = var_key
 
         for time_selection in time_selections:
-            # Set name/yrs attributes based on time selection mode
-            if is_time_slice_mode:
-                # For time slices, we set attributes after loading data
-                # since we need the actual time coordinates
-                pass
-            else:
-                parameter._set_name_yrs_attrs(test_ds, ref_ds, time_selection)
+            is_time_slice = time_selection_type == "time_slices"
 
-            # Get datasets - pass is_time_slice flag if it's a time slice
-            if is_time_slice_mode:
-                ds_test = test_ds.get_climo_dataset(
-                    var_key, time_selection, is_time_slice=True
-                )
-                ds_ref = _get_ref_dataset(
-                    ref_ds, var_key, time_selection, is_time_slice=True
-                )
-                # For time slices, set name_yrs after data is loaded
-                parameter._set_name_yrs_attrs(test_ds, ref_ds, time_selection)
-                # Use the climatology season (ANN) for land sea mask
+            # Get test and reference datasets.
+            # NOTE: lat_lon diagnostics get reference datasets differently than
+            # other sets using its own helper function `_get_ref_dataset`.
+            if is_time_slice:
+                ds_test = test_ds.get_time_sliced_dataset(var_key, time_selection)
+
+                # For time slices, always use the annual land-sea mask.
                 ds_land_sea_mask: xr.Dataset = test_ds._get_land_sea_mask("ANN")
             else:
-                ds_test = test_ds.get_climo_dataset(var_key, time_selection)
-                ds_ref = _get_ref_dataset(
-                    ref_ds, var_key, time_selection, is_time_slice=False
+                # time_selection will be ClimoFreq, so ignore type checking here.
+                ds_test = test_ds.get_climo_dataset(var_key, time_selection)  # type: ignore[arg-type]
+                ds_land_sea_mask: xr.Dataset = test_ds._get_land_sea_mask(  # type: ignore[no-redef]
+                    time_selection  # type: ignore[arg-type]
                 )
-                ds_land_sea_mask = test_ds._get_land_sea_mask(time_selection)
+
+            ds_ref = _get_ref_dataset(ref_ds, var_key, time_selection, is_time_slice)
+
+            # Set name_yrs after loading data because time sliced datasets
+            # have the required attributes only after loading the data.
+            parameter._set_name_yrs_attrs(test_ds, ref_ds, time_selection)
 
             if ds_ref is None:
                 is_vars_3d = has_z_axis(ds_test[var_key])
@@ -317,7 +302,7 @@ def _run_diags_2d(
     ds_test: xr.Dataset,
     ds_ref: xr.Dataset,
     ds_land_sea_mask: xr.Dataset,
-    season: str,
+    time_selection: TimeSelection,
     regions: list[str],
     var_key: str,
     ref_name: str,
@@ -338,8 +323,8 @@ def _run_diags_2d(
     ds_land_sea_mask : xr.Dataset
         The land sea mask dataset, which is only used for masking if the region
         is "land" or "ocean".
-    season : str
-        The season.
+    time_selection : TimeSelection
+        The time slice or season.
     regions : list[str]
         The list of regions.
     var_key : str
@@ -372,7 +357,9 @@ def _run_diags_2d(
             ds_diff_region,
         )
 
-        parameter._set_param_output_attrs(var_key, season, region, ref_name, ilev=None)
+        parameter._set_param_output_attrs(
+            var_key, time_selection, region, ref_name, ilev=None
+        )
         _save_data_metrics_and_plots(
             parameter,
             plot_func,
@@ -391,7 +378,7 @@ def _run_diags_3d(
     ds_test: xr.Dataset,
     ds_ref: xr.Dataset,
     ds_land_sea_mask: xr.Dataset,
-    season: str,
+    time_selection: str,
     regions: list[str],
     var_key: str,
     ref_name: str,
@@ -412,8 +399,8 @@ def _run_diags_3d(
     ds_land_sea_mask : xr.Dataset
         The land sea mask dataset, which is only used for masking if the region
         is "land" or "ocean".
-    season : str
-        The season.
+    time_selection : TimeSelection
+        The time slice or season.
     regions : list[str]
         The list of regions.
     var_key : str
@@ -457,7 +444,9 @@ def _run_diags_3d(
                 ds_diff_region,
             )
 
-            parameter._set_param_output_attrs(var_key, season, region, ref_name, ilev)
+            parameter._set_param_output_attrs(
+                var_key, time_selection, region, ref_name, ilev
+            )
             _save_data_metrics_and_plots(
                 parameter,
                 plot_func,
@@ -474,7 +463,7 @@ def _run_diags_3d(
 def _get_ref_dataset(
     dataset: Dataset,
     var_key: str,
-    time_selection: ClimoFreq | str,
+    time_selection: TimeSelection,
     is_time_slice: bool = False,
 ) -> xr.Dataset | None:
     """Get the reference dataset for the variable and time selection.
@@ -506,9 +495,11 @@ def _get_ref_dataset(
     """
     if dataset.data_type == "ref":
         try:
-            ds_ref = dataset.get_climo_dataset(
-                var_key, time_selection, is_time_slice=is_time_slice
-            )
+            if is_time_slice:
+                ds_ref = dataset.get_time_sliced_dataset(var_key, time_selection)
+            else:
+                # time_selection will be ClimoFreq, so ignore type checking here.
+                ds_ref = dataset.get_climo_dataset(var_key, time_selection)  # type: ignore[arg-type]
         except (RuntimeError, IOError):
             ds_ref = None
 

e3sm_diags/driver/meridional_mean_2d_driver.py

@@ -7,14 +7,15 @@
 import xcdat as xc
 
 from e3sm_diags.driver.utils.dataset_xr import Dataset
-from e3sm_diags.driver.utils.io import _save_data_metrics_and_plots
+from e3sm_diags.driver.utils.io import (
+    _get_xarray_datasets,
+    _save_data_metrics_and_plots,
+)
 from e3sm_diags.driver.utils.regrid import (
     align_grids_to_lower_res,
     has_z_axis,
     regrid_z_axis_to_plevs,
 )
-from e3sm_diags.driver.utils.time_slice import check_time_selection
-from e3sm_diags.driver.utils.type_annotations import MetricsDict
 from e3sm_diags.logger import _setup_child_logger
 from e3sm_diags.metrics.metrics import correlation, rmse, spatial_avg
 from e3sm_diags.parameter.zonal_mean_2d_parameter import DEFAULT_PLEVS
@@ -23,6 +24,7 @@
 logger = _setup_child_logger(__name__)
 
 if TYPE_CHECKING:
+    from e3sm_diags.driver.utils.type_annotations import MetricsDict, TimeSelection
     from e3sm_diags.parameter.meridional_mean_2d_parameter import (
         MeridionalMean2dParameter,
     )
@@ -51,47 +53,25 @@ def run_diag(parameter: MeridionalMean2dParameter) -> MeridionalMean2dParameter:
         If the test or ref variables do are not 3-D (no Z-axis).
     """
     variables = parameter.variables
-    seasons = parameter.seasons
-    time_slices = getattr(parameter, "time_slices", [])
     ref_name = getattr(parameter, "ref_name", "")
 
-    # Check that either seasons or time_slices is specified, but not both
-    has_seasons, has_time_slices = check_time_selection(
-        seasons, time_slices, require_one=True
-    )
-
-    # Determine which time selection to use
-    time_selections = time_slices if has_time_slices else seasons
-    is_time_slice_mode = has_time_slices
-
     test_ds = Dataset(parameter, data_type="test")
     ref_ds = Dataset(parameter, data_type="ref")
 
+    time_selection_type, time_selections = parameter._get_time_selection_to_use()
+
     for var_key in variables:
         logger.info("Variable: {}".format(var_key))
         parameter.var_id = var_key
 
         for time_selection in time_selections:
-            # Set name/yrs attributes based on time selection mode
-            if is_time_slice_mode:
-                # For time slices, we set attributes after loading data
-                pass
-            else:
-                parameter._set_name_yrs_attrs(test_ds, ref_ds, time_selection)
-
-            # Get datasets - pass is_time_slice flag if it's a time slice
-            if is_time_slice_mode:
-                ds_test = test_ds.get_climo_dataset(
-                    var_key, time_selection, is_time_slice=True
-                )
-                ds_ref = ref_ds.get_climo_dataset(
-                    var_key, time_selection, is_time_slice=True
-                )
-                # For time slices, set name_yrs after data is loaded
-                parameter._set_name_yrs_attrs(test_ds, ref_ds, time_selection)
-            else:
-                ds_test = test_ds.get_climo_dataset(var_key, time_selection)
-                ds_ref = ref_ds.get_climo_dataset(var_key, time_selection)
+            ds_test, ds_ref, _ = _get_xarray_datasets(
+                test_ds, ref_ds, var_key, time_selection_type, time_selection
+            )
+
+            # Set name_yrs after loading data because time sliced datasets
+            # have the required attributes only after loading the data.
+            parameter._set_name_yrs_attrs(test_ds, ref_ds, time_selection)
 
             dv_test = ds_test[var_key]
             dv_ref = ds_ref[var_key]
@@ -122,7 +102,7 @@ def _run_diags_3d(
     parameter: MeridionalMean2dParameter,
     ds_test: xr.Dataset,
     ds_ref: xr.Dataset,
-    season: str,
+    time_selection: TimeSelection,
     var_key: str,
     ref_name: str,
 ):
@@ -171,7 +151,7 @@ def _run_diags_3d(
     )
 
     parameter._set_param_output_attrs(
-        var_key, season, parameter.regions[0], ref_name, ilev=None
+        var_key, time_selection, parameter.regions[0], ref_name, ilev=None
     )
     _save_data_metrics_and_plots(
         parameter,