Add unit tests

- Test coverage for get_time_sliced_dataset(), _add_time_series_filepath_attr(), check_values(), and get_xarray_datasets()`

Author: tomvothecoder

e3sm_diags/driver/utils/dataset_xr.py

@@ -360,8 +360,6 @@ def get_time_sliced_dataset(self, var: str, time_slice: TimeSlice) -> xr.Dataset
 
         These variables can either be from the test data or reference data.
 
-        TODO: Add unit test for this method.
-
         Parameters
         ----------
         var : str
@@ -381,34 +379,31 @@ def get_time_sliced_dataset(self, var: str, time_slice: TimeSlice) -> xr.Dataset
         if not isinstance(self.var, str) or self.var == "":
             raise ValueError("The `var` argument is not a valid string.")
 
-        ds = self._get_full_dataset(var)
-        ds = self._apply_time_slice_to_dataset(ds, time_slice)
+        filepath = self._get_filepath_with_params()
 
-        try:
-            filepath = self._get_filepath_with_params()
+        if filepath is None:
+            raise RuntimeError(
+                f"Unable to get file path for {self.data_type} dataset. "
+                f"For time slicing, please ensure that "
+                f"{'ref_file' if self.data_type == 'ref' else 'test_file'} parameter is set."
+            )
 
-            if filepath:
-                self.parameter._add_filepath_attr(self.data_type, filepath)
+        if not os.path.exists(filepath):
+            raise RuntimeError(f"File not found: {filepath}")
 
-        except Exception as e:
-            logger.warning(f"Failed to store absolute file path: {e}")
+        self.parameter._add_filepath_attr(self.data_type, filepath)
+
+        ds = self._get_full_dataset()
+        ds = self._apply_time_slice_to_dataset(ds, time_slice)
 
         return ds
 
-    def _get_full_dataset(self, var: str) -> xr.Dataset:
+    def _get_full_dataset(self) -> xr.Dataset:
         """Get the full dataset without any time averaging for time slicing.
 
         This function uses the dataset's file path parameters to directly open
         the raw data file for time slicing operations.
 
-        TODO: Add unit test for this method.
-        FIXME: The `var` parameter is currently unused.
-
-        Parameters
-        ----------
-        var : str
-            The key of the variable.
-
         Returns
         -------
         xr.Dataset
@@ -419,31 +414,20 @@ def _get_full_dataset(self, var: str) -> xr.Dataset:
         RuntimeError
             If unable to get the full dataset or file not found.
         """
-        filepath = self._get_filepath_with_params()
-
-        if filepath is None:
-            raise RuntimeError(
-                f"Unable to get file path for {self.data_type} dataset. "
-                f"For time slicing, please ensure that "
-                f"{'ref_file' if self.data_type == 'ref' else 'test_file'} parameter is set."
-            )
-
-        if not os.path.exists(filepath):
-            raise RuntimeError(f"File not found: {filepath}")
+        filepath = getattr(self.parameter, f"{self.data_type}_data_file_path")
 
         logger.info(f"Opening full dataset from: {filepath}")
 
         try:
-            ds = xr.open_dataset(filepath, decode_times=True)
+            ds = xc.open_dataset(filepath, add_bounds=["X", "Y", "T"])
 
             logger.info(
                 f"Successfully opened dataset with time dimension size: {ds.sizes.get('time', 'N/A')}"
             )
-
-            return ds
-        # FIXME: This except block is too broad; it should catch specific exceptions.
-        except Exception as e:
+        except (FileNotFoundError, OSError, ValueError) as e:
             raise RuntimeError(f"Failed to open dataset {filepath}: {e}") from e
+        else:
+            return ds
 
     def _get_filepath_with_params(self) -> str | None:
         """Get the filepath using parameters.
@@ -483,14 +467,10 @@ def _apply_time_slice_to_dataset(
         xr.Dataset
             The dataset with time slice applied.
         """
-        # FIXME: Use xcdat.get_dim_keys to find time dimension.
-        time_dim = None
-
-        for dim in ds.dims:
-            if str(dim).lower() in ["time", "t"]:
-                time_dim = dim
-
-                break
+        try:
+            time_dim = xc.get_dim_keys(ds, axis="T")
+        except (ValueError, KeyError):
+            time_dim = None
 
         if time_dim is None:
             logger.warning(
@@ -499,9 +479,15 @@ def _apply_time_slice_to_dataset(
 
             return ds
 
-        # Single index selection
         index = int(time_slice)
-        ds_sliced = ds.isel({time_dim: index})
+
+        try:
+            ds_sliced = ds.isel({time_dim: index})
+        except IndexError as e:
+            raise IndexError(
+                f"Time slice index {index} is out of bounds for time dimension "
+                f"of size {ds.sizes[time_dim]}."
+            ) from e
 
         logger.info(
             f"Applied time slice '{time_slice}' to dataset. "
@@ -562,7 +548,7 @@ def get_climo_dataset(self, var: str, season: ClimoFreq) -> xr.Dataset:
             ds_climo = climo(ds, self.var, season).to_dataset()
             ds_climo = ds_climo.bounds.add_missing_bounds(axes=["X", "Y"])
 
-            self.parameter._add_time_series_file_path_attr(self.data_type, ds)
+            self.parameter._add_time_series_filepath_attr(self.data_type, ds)
 
             return ds_climo
 

e3sm_diags/parameter/core_parameter.py

@@ -412,7 +412,7 @@ def _run_diag(self) -> list[Any]:
 
         return results
 
-    def _add_time_series_file_path_attr(
+    def _add_time_series_filepath_attr(
         self,
         data_type: Literal["test", "ref"],
         ds: xr.Dataset,
@@ -434,9 +434,9 @@ def _add_time_series_file_path_attr(
         if data_type not in {"test", "ref"}:
             raise ValueError("data_type must be either 'test' or 'ref'.")
 
-        file_path_attr = f"{data_type}_data_file_path"
+        filepath_attr = f"{data_type}_data_file_path"
 
-        setattr(self, file_path_attr, getattr(ds, "file_path", "Unknown"))
+        setattr(self, filepath_attr, getattr(ds, "file_path", "Unknown"))
 
     def _add_filepath_attr(
         self,
@@ -462,14 +462,14 @@ def _add_filepath_attr(
         if data_type not in {"test", "ref"}:
             raise ValueError("data_type must be either 'test' or 'ref'.")
 
-        file_path_attr = f"{data_type}_data_file_path"
+        filepath_attr = f"{data_type}_data_file_path"
 
         if not filepath:
             raise ValueError(
                 "Filepath must be provided for the climatology or time-slice data."
             )
 
-        setattr(self, file_path_attr, os.path.abspath(filepath))
+        setattr(self, filepath_attr, os.path.abspath(filepath))
 
     def _get_time_selection_to_use(
         self, require_one: bool = True
@@ -479,8 +479,6 @@ def _get_time_selection_to_use(
 
         If ``time_slices`` are specified, they take precedence over ``seasons``.
 
-        TODO: Add unit tests for this method.
-
         Parameters
         ----------
         require_one : bool, optional

examples/ex8-native-grid-visualization/README.md

@@ -12,6 +12,7 @@ This example demonstrates the **native grid visualization** feature introduced i
 ## Key Features
 
 The native grid visualization capability:
+
 - Supports various native grid formats (cubed-sphere, unstructured, etc.)
 - Eliminates artifacts introduced by regridding
 - Enables comparison of models with different native grids
@@ -28,10 +29,12 @@ The native grid visualization capability:
 ## Data Requirements
 
 This example uses test data located at LCRC:
+
 - Data path: `/lcrc/group/e3sm/public_html/e3sm_diags_test_data/native_grid`
 - Grid files: `/lcrc/group/e3sm/diagnostics/grids/`
 
 For your own data, ensure you have:
+
 1. Model output files on native grid
 2. Corresponding grid files in UGRID format
 
@@ -79,6 +82,7 @@ e3sm_diags lat_lon_native \
 ## Configuration File
 
 The `diags.cfg` file allows you to customize:
+
 - Variables to plot (e.g., TGCLDLWP)
 - Regions of interest
 - Colormap settings
@@ -87,6 +91,7 @@ The `diags.cfg` file allows you to customize:
 ## Expected Output
 
 The diagnostic will generate:
+
 - Native grid visualizations for specified variables
 - Test model plot
 - Reference model plot
@@ -97,14 +102,15 @@ Results will be saved in: `/lcrc/group/e3sm/public_html/diagnostic_output/$USER/
 
 ## Notes
 
-- Native grid visualization requires UXarray library (included in E3SM Unified environment)
+- Native grid visualization requires the UXarray library, which is included as a dependency of E3SM diagnostics and the E3SM Unified environment.
 - Grid files must be in UGRID format
 - This example uses `time_slices` for snapshot analysis; you can also use `seasons` for climatology
 - For model-only runs (no reference data), set `model_only = True` and omit ref_grid_file
 
 ## Differences from Regular lat_lon Set
 
 Unlike the standard `lat_lon` set which regrids data to a regular lat-lon grid:
+
 - `lat_lon_native` preserves the original grid structure
 - No interpolation artifacts
 - Better representation of native grid features
@@ -113,5 +119,6 @@ Unlike the standard `lat_lon` set which regrids data to a regular lat-lon grid:
 ## More Information
 
 For more details, see:
+
 - [E3SM Diags Documentation](https://e3sm-project.github.io/e3sm_diags)
 - [UXarray Documentation](https://uxarray.readthedocs.io/)