Extend snapshot support for other core sets and doc updates (#1013)

* Extend snapshot support for other core sets

* add new util module

* fix tests

* fix mean_2d_viewer

* add testing scripts

* clean up; clarify on input format for time_slices

* update testing scripts

* update README.md for 3.1.0 new feature

* Update documentation for v3.1.0 features and parameters

This commit updates the documentation to reflect new features and
parameters added in v3.1.0:

1. Updated available-parameters.rst:
   - Added time_slices parameter for snapshot analysis
   - Added lat_lon_native set-specific parameters (test_grid_file,
     ref_grid_file, antialiased)
   - Added test_file and ref_file parameters
   - Updated regrid_tool default from 'esmf' to 'xesmf'
   - Updated granulate default to include 'time_slices'
   - Updated sets default list to include all current diagnostic sets
   - Added notes about mutual exclusivity of time_slices and seasons

2. Updated examples.rst:
   - Added Example 8: Native Grid Visualization
   - Added Example 9: Snapshot Analysis for Core Sets
   - Updated running instructions to include new examples
   - Updated batch script example with ex8 and ex9

3. Created example files:
   - examples/ex8-native-grid-visualization/ (ex8.py, diags.cfg, README.md)
   - examples/ex9-snapshot-analysis/ (ex9.py, diags.cfg, README.md)

These changes document the two major features introduced in v3.1.0:
- Native grid visualization using UXarray
- Snapshot analysis for core diagnostic sets using time_slices

🤖 Generated with [Claude Code](https://claude.com/claude-code)

Co-Authored-By: Claude <[email protected]>

* fix examples; update default set for lat_lon_native

* 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

* Add unit tests
- Test coverage for get_time_sliced_dataset(), _add_time_series_filepath_attr(), check_values(), and get_xarray_datasets()`

* Fix method reference in lat_lon_native_driver.py

* Update e3sm_diags/driver/lat_lon_driver.py

Co-authored-by: Copilot <[email protected]>

* Apply Copilot review suggestions

* minor doc change

---------

Co-authored-by: Claude <[email protected]>
Co-authored-by: tomvothecoder <[email protected]>
Co-authored-by: Copilot <[email protected]>

Author: 4 people

README.md

@@ -57,6 +57,15 @@ The modernization improves performance, usability, and maintainability, paving t
 for future enhancements to E3SM development. The refactored codebase is now more robust
 and extensively covered by unit tests, setting a solid foundation for ongoing testing and development.
 
+### New Features in v3.1.0
+
+v3.1.0 introduces significant enhancements to support advanced grid analysis and temporal snapshot capabilities:
+
+| Feature name <br />(set name)                    | Brief Introduction                                                                                                                                                                     | Developers Contributors\*                | Released version |
+| ------------------------------------------------ | -------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | ---------------------------------------- | ---------------- |
+| Native grid visualization (lat_lon_native)       | Support for plotting data on native grids (e.g., cubed-sphere, unstructured grids) using UXarray, enabling visualization without regridding to preserve native grid features         | Jill Zhang, Tom Vo                       | 3.1.0            |
+| Snapshot analysis for core sets                  | Index-based time selection for snapshot analysis on core diagnostic sets (lat_lon, lat_lon_native, polar, zonal_mean_2d, meridional_mean_2d, zonal_mean_2d_stratosphere), allowing analysis of individual time steps instead of climatological means | Jill Zhang, Tom Vo                       | 3.1.0            |
+
 ### New Feature added during v2 development
 
 | Feature name <br />(set name)                                       | Brief Introduction                                                                                                                                                                                     | Developers Contributors\*                                                                                                                                             | Released version |

auxiliary_tools/debug/1013-snapshot-analysis-core-sets/T.cfg

@@ -0,0 +1,50 @@
+[#]
+sets = ["lat_lon"]
+case_id = "model_vs_model"
+variables = ["T"]
+seasons = ["ANN", "01", "02", "03", "04", "05", "06", "07", "08", "09", "10", "11", "12", "DJF", "MAM", "JJA", "SON"]
+plevs = [850.0]
+contour_levels = [240, 245, 250, 255, 260, 265, 270, 275, 280, 285, 290, 295]
+diff_levels = [-5, -4, -3, -2, -1, -0.5, -0.25, 0.25, 0.5, 1, 2, 3, 4, 5]
+regrid_method = "bilinear"
+
+[#]
+sets = ["polar"]
+case_id = "model_vs_model"
+variables = ["T"]
+seasons = ["DJF", "MAM", "JJA", "SON"]
+regions = ["polar_S", "polar_N"]
+plevs = [850.0]
+contour_levels = [230, 240, 250, 260, 270, 280, 290, 300, 310]
+diff_levels = [-15, -10, -7.5, -5, -2.5, -1, 1, 2.5, 5, 7.5, 10, 15]
+
+[#]
+sets = ["zonal_mean_2d"]
+case_id = "model_vs_model"
+variables = ["T"]
+seasons = ["ANN", "01", "02", "03", "04", "05", "06", "07", "08", "09", "10", "11", "12", "DJF", "MAM", "JJA", "SON"]
+contour_levels = [180,185,190,200,210,220,230,240,250,260,270,280,290,295,300]
+diff_levels = [-3.0, -2.5, -2, -1.5, -1, -0.5, -0.25, 0.25, 0.5, 1, 1.5, 2, 2.5, 3.0]
+
+[#]
+sets = ["zonal_mean_2d_stratosphere"]
+case_id = "model_vs_model"
+variables = ["T"]
+seasons = ["ANN", "01", "02", "03", "04", "05", "06", "07", "08", "09", "10", "11", "12", "DJF", "MAM", "JJA", "SON"]
+contour_levels = [180,185,190,200,210,220,230,240,250,260,270,280,290,295,300]
+diff_levels = [-8, -6,-4,-2, -1, -0.5, 0.5, 1, 2, 4,6, 8]
+
+[#]
+sets = ["zonal_mean_xy"]
+case_id = "model_vs_model"
+variables = ["T"]
+seasons = ["ANN", "01", "02", "03", "04", "05", "06", "07", "08", "09", "10", "11", "12", "DJF", "MAM", "JJA", "SON"]
+plevs = [850.0]
+
+[#]
+sets = ["meridional_mean_2d"]
+case_id = "model_vs_model"
+variables = ["T"]
+seasons = ["ANN", "01", "02", "03", "04", "05", "06", "07", "08", "09", "10", "11", "12", "DJF", "MAM", "JJA", "SON"]
+contour_levels = [180,185,190,200,210,220,230,240,250,260,270,280,290,295,300]
+diff_levels = [-7,-6,-5,-4,-3,-2,-1,1,2,3,4,5,6,7]

auxiliary_tools/debug/1013-snapshot-analysis-core-sets/run_snapshot_core.py

@@ -0,0 +1,93 @@
+#!/usr/bin/env python3
+"""
+This script runs e3sm_diags with the core sets to visualize snap-shot data.
+"""
+
+import os
+import sys
+
+from e3sm_diags.parameter.core_parameter import CoreParameter
+from e3sm_diags.run import runner
+
+# Auto-detect username
+username = os.environ.get('USER', 'unknown_user')
+
+# Create parameter objects for 3 different runs
+params = []
+
+## (1) First test configuration
+#param1 = LatLonNativeParameter()
+#param1.results_dir = f"/lcrc/group/e3sm/public_html/diagnostic_output/{username}/tests/lat_lon_native_test_1"
+#param1.test_data_path = "/lcrc/group/e3sm/public_html/e3sm_diags_test_data/native_grid"
+#param1.test_name = "v3.LR.amip_0101"
+#param1.short_test_name = "v3.LR.amip_0101"
+#param1.reference_data_path = "/lcrc/group/e3sm/public_html/e3sm_diags_test_data/native_grid"
+#param1.ref_name = "v3.HR.test4"
+#param1.short_ref_name = "v3.HR.test4"
+#param1.seasons = ["DJF"]
+#param1.test_grid_file = "/lcrc/group/e3sm/diagnostics/grids/ne30pg2.nc"
+#param1.ref_grid_file = "/lcrc/group/e3sm/diagnostics/grids/ne120pg2.nc"
+#param1.case_id = "model_vs_model"
+#param1.run_type = "model_vs_model"
+#params.append(param1)
+#
+## (2) Second test configuration
+#param2 = LatLonNativeParameter()
+#param2.results_dir = f"/lcrc/group/e3sm/public_html/diagnostic_output/{username}/tests/lat_lon_native_test_2"
+#param2.test_data_path = "/lcrc/group/e3sm/public_html/e3sm_diags_test_data/native_grid"
+#param2.test_file = "v3.LR.amip_0101_DJF_climo.nc"
+#param2.short_test_name = "v3.LR.amip_0101"
+#param2.reference_data_path = "/lcrc/group/e3sm/public_html/e3sm_diags_test_data/native_grid"
+#param2.ref_file = "v3.HR.test4_DJF_climo.nc"
+#param2.short_ref_name = "v3.HR.test4"
+#param2.seasons = ["DJF"]
+#param2.test_grid_file = "/lcrc/group/e3sm/diagnostics/grids/ne30pg2.nc"
+#param2.ref_grid_file = "/lcrc/group/e3sm/diagnostics/grids/ne120pg2.nc"
+#param2.case_id = "model_vs_model"
+#param2.run_type = "model_vs_model"
+#params.append(param2)
+
+# (3) Third test configuration
+param3 = CoreParameter()
+param3.results_dir = f"/lcrc/group/e3sm/public_html/diagnostic_output/{username}/tests/1013-snapshot-analysis-core-sets"
+param3.test_data_path = "/lcrc/group/e3sm/public_html/e3sm_diags_test_data/postprocessed_e3sm_v2_data_for_e3sm_diags/20210528.v2rc3e.piControl.ne30pg2_EC30to60E2r2.chrysalis/time-series/rgr"
+param3.test_file = "T_005101_006012.nc"
+param3.reference_data_path = "/lcrc/group/e3sm/public_html/e3sm_diags_test_data/postprocessed_e3sm_v2_data_for_e3sm_diags/20210528.v2rc3e.piControl.ne30pg2_EC30to60E2r2.chrysalis/time-series/rgr"
+param3.ref_file = "T_005101_006012.nc"
+param3.short_test_name = "v2 test"
+param3.short_ref_name = "v2 test"
+param3.time_slices = ["0","1"]
+param3.case_id = "model_vs_model"
+param3.run_type = "model_vs_model"
+params.append(param3)
+
+# Run the single diagnostic, comment out for complete diagnostics.
+cfg_path = "auxiliary_tools/debug/1013-snapshot-analysis-core-sets/T.cfg"
+sys.argv.extend(["--diags", cfg_path])
+
+runner.sets_to_run = ["lat_lon",
+        "zonal_mean_xy",
+        "zonal_mean_2d",
+        "zonal_mean_2d_stratosphere",
+        "polar",
+        "meridional_mean_2d",]
+
+# Run each test sequentially
+for i, param in enumerate(params, 1):
+    print(f"\n{'='*60}")
+    print(f"Running Test {i}: {param.results_dir}")
+    print(f"{'='*60}")
+
+    # Create results directory
+    if not os.path.exists(param.results_dir):
+        os.makedirs(param.results_dir)
+
+    # Run the diagnostic
+    runner.run_diags([param])
+    print(f"Test {i} completed!")
+
+print(f"\n{'='*60}")
+print("All tests completed!")
+print(f"{'='*60}")
+
+

docs/source/available-parameters.rst

@@ -70,11 +70,18 @@ functionality of the diagnostics.
 -  **regrid_method**: What regrid method of the regrid tool to use.
    Possible values are ``'linear'`` or ``'conservative'``. Default is ``'conservative'``.
    Read the xCDAT documentation on `regridding`_ for more information.
--  **regrid_tool**: The regrid tool to use. Default is ``'esmf'``.
+-  **regrid_tool**: The regrid tool to use. Default is ``'xesmf'``.
    Read the xCDAT documentation on `regridding`_ for more information.
+   **Note:** The ``lat_lon_native`` set does not use regridding, so this parameter is ignored for that set.
 -  **seasons**: A list of season to use. Default is annual and all seasons: ``['ANN', 'DJF', 'MAM', 'JJA', 'SON']``.
+   **Note:** This parameter is mutually exclusive with ``time_slices``. When using ``time_slices``, do not set ``seasons``.
+-  **time_slices**: *(v3.1.0+)* A list of time indices for snapshot analysis. Examples: ``['0']``, ``['5']``, ``['0', '1', '2']``.
+   Time slices are zero-based indices into the time dimension of the input dataset.
+   This enables analysis of individual time steps instead of climatological means.
+   **Note:** This parameter is mutually exclusive with ``seasons``. When using ``time_slices``, do not set ``seasons``.
+   Supported sets: ``lat_lon``, ``lat_lon_native``, ``polar``, ``zonal_mean_xy``, ``zonal_mean_2d``, ``meridional_mean_2d``, ``zonal_mean_2d_stratosphere``.
 -  **sets**: A list of the sets to be run. Default is all sets:
-   ``['zonal_mean_xy', 'zonal_mean_2d', 'meridional_mean_2d', 'lat_lon', 'polar', 'area_mean_time_series', 'cosp_histogram', 'enso_diags', 'qbo', 'streamflow','diurnal_cycle']``.
+   ``['zonal_mean_xy', 'zonal_mean_2d', 'zonal_mean_2d_stratosphere', 'meridional_mean_2d', 'lat_lon', 'lat_lon_native', 'polar', 'area_mean_time_series', 'cosp_histogram', 'enso_diags', 'qbo', 'streamflow', 'diurnal_cycle', 'arm_diags', 'tc_analysis', 'annual_cycle_zonal_mean', 'lat_lon_land', 'lat_lon_river', 'aerosol_aeronet', 'aerosol_budget']``.
 -  **variables**: What variable(s) to use for this run. Ex: ``variables=["T", "PRECT"]``.
 
 .. _regridding: https://xcdat.readthedocs.io/en/latest/getting-started-guide/faqs.html#regridding
@@ -262,12 +269,30 @@ You can specify both ``test_end_yr`` and ``ref_end_yr`` or just ``end_yr``.
 -  **plot_log_plevs**: Log-scale the y-axis. Default ``False``.
 -  **plot_plevs**: Plot the pressure levels. Default ``False``.
 
+``'lat_lon_native'`` *(v3.1.0+)*:
+
+-  **test_grid_file**: *(Required)* Path to the grid file for test data in UGRID format (e.g., ``'/path/to/ne30pg2.nc'``).
+   The grid file defines the native grid structure used by UXarray for visualization.
+-  **ref_grid_file**: Path to the grid file for reference data in UGRID format (e.g., ``'/path/to/ne30pg2.nc'``).
+   Required for model vs model comparisons. Can be omitted for model-only runs.
+-  **antialiased**: Apply antialiasing to the plot. Default ``False``. Setting to ``True`` may improve visual quality but can impact performance.
+-  **time_slices**: Time indices for snapshot analysis (same as in core parameters). See ``time_slices`` description above.
+
+**Notes for lat_lon_native:**
+
+-  Native grid visualization requires UXarray (included in E3SM Unified environment)
+-  Grid files must be in UGRID format
+-  Regridding parameters (``regrid_tool``, ``regrid_method``) are ignored for this set
+-  Can use either ``seasons`` for climatology or ``time_slices`` for snapshot analysis (mutually exclusive)
 
 Other parameters
 ~~~~~~~~~~~~~~~~
 
 -  **dataset**: Default is ``''``.
--  **granulate**: Default is ``['variables', 'seasons', 'plevs', 'regions']``.
+-  **granulate**: Default is ``['variables', 'seasons', 'plevs', 'regions', 'time_slices']``.
+   This parameter controls how diagnostics are split into separate runs.
 -  **selectors**: Default is ``['sets', 'seasons']``. See :ref:`Using the selectors parameter <selector-ex>`.
 -  **viewer_descr**: Used to specify values in the viewer. Default ``{}``.
 -  **fail_on_incomplete**: Exit status will reflect failure if any parameter fails to complete. Default is ``False`` (e.g., a failing parameter will not create a failing exit code).
+-  **test_file**: *(v3.1.0+)* Specify the exact file name for test data. Useful for snapshot analysis with ``time_slices`` or when using specific data files.
+-  **ref_file**: *(v3.1.0+)* Specify the exact file name for reference data. Useful for snapshot analysis with ``time_slices`` or when using specific data files.

docs/source/examples.rst

@@ -93,6 +93,38 @@ on two different sets: ``zonal_mean_2d`` an ``lat_lon``.
 so you can compare different version of the data, or the same variable from different datasets.
 We are comparing CERES EBAF TOA version 2.8 and 4.0.
 
+8. Native Grid Visualization (v3.1.0)
+--------------------------------------
+`This example <https://github.com/E3SM-Project/e3sm_diags/tree/master/examples/ex8-native-grid-visualization>`__ demonstrates how to visualize model data on its native grid
+(e.g., cubed-sphere, unstructured grids) without regridding to a regular lat-lon grid.
+This feature uses UXarray to preserve native grid features and is particularly useful for high-resolution models with complex grid structures.
+The example shows model vs model comparison using snapshot analysis on native grids.
+
+**Key features:**
+
+- Visualize data on native grids without regridding
+- Preserve native grid features and characteristics
+- Support for cubed-sphere and unstructured grids
+- Uses UXarray for grid-aware operations
+
+9. Snapshot Analysis for Core Sets (v3.1.0)
+--------------------------------------------
+`This example <https://github.com/E3SM-Project/e3sm_diags/tree/master/examples/ex9-snapshot-analysis>`__ demonstrates time slice analysis on core diagnostic sets.
+Instead of computing climatological seasonal means, this analyzes individual time steps from model output using index-based time selection.
+
+This is useful for analyzing specific events, comparing model states at particular time points,
+or understanding temporal evolution without time averaging. The example shows how to use the ``time_slices`` parameter
+on multiple diagnostic sets (lat_lon, zonal_mean_2d, polar, meridional_mean_2d, zonal_mean_2d_stratosphere).
+
+**Key features:**
+
+- Index-based time selection (e.g., time_slices = ["0", "1", "2"])
+- Analyze individual time steps without temporal averaging
+- Event-based or process-oriented diagnostics
+- Works across multiple core diagnostic sets
+
+**Note:** ``time_slices`` and ``seasons`` parameters are mutually exclusive.
+
 Running the Examples
 ====================
 
@@ -117,7 +149,7 @@ The parameters file contains information related to the location
 of the data, what years to run the diagnostics on, what plots to create, and more.
 
 The configuration file provides information about the diagnostics you are running.
-This is used in Ex.4,5,7.
+This is used in Ex.4, 5, 7, 8, 9.
 
 Parameters for each example can be found in
 `this directory <https://github.com/E3SM-Project/e3sm_diags/tree/master/examples>`__.
@@ -134,7 +166,7 @@ Use the code below to run the diagnostics.
         salloc --nodes 1 --qos interactive --time 01:00:00 --constraint cpu --account=e3sm
         # Enter the E3SM Unified environment. For Perlmutter CPU, the command to do this is:
         source /global/common/software/e3sm/anaconda_envs/load_latest_e3sm_unified_pm-cpu.sh
-        # Running Ex.1. For examples 4,5,7 append ``-d diags.cfg``.
+        # Running Ex.1. For examples 4, 5, 7, 8, 9 append ``-d diags.cfg``.
         python ex1.py --multiprocessing --num_workers=32
         # You may need to change permissions on your web directory to see the example output.
         chmod -R 755 <your web directory>
@@ -174,6 +206,8 @@ These were generated with the following script:
        # emacs ex5-model-vs-obs/ex5.py
        # emacs ex6-model-vs-obs-custom/ex6.py
        # emacs ex7-obs-vs-obs/ex7.py
+       # emacs ex8-native-grid-visualization/ex8.py
+       # emacs ex9-snapshot-analysis/ex9.py
 
        source /global/common/software/e3sm/anaconda_envs/load_latest_e3sm_unified_pm-cpu.sh
        cd ex1-model_ts-vs-model_ts
@@ -187,9 +221,13 @@ These were generated with the following script:
        cd ../ex5-model-vs-obs
        python ex5.py --multiprocessing --num_workers=32 -d diags.cfg
        cd ../ex6-model-vs-obs-custom
-       python ex6.py --multiprocessing --num_workers=32
+       python ex6.py 
        cd ../ex7-obs-vs-obs
-       python ex7.py --multiprocessing --num_workers=32 -d diags.cfg
+       python ex7.py -d diags.cfg
+       cd ../ex8-native-grid-visualization
+       python ex8.py -d diags.cfg
+       cd ../ex9-snapshot-analysis
+       python ex9.py -d diags.cfg
        cd ../
 
        chmod -R 755 /global/cfs/cdirs/e3sm/www/forsyth/examples