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

chengzhuzhang · 2025-10-17T21:51:27Z

Description

Overview

The snapshot analysis was first implement for lat_lon_native set for native model output. This PR implemented time slice support for 6 additional core diagnostic sets in e3sm_diags,extending the existing time slice analysis feature from lat_lon_native to other core sets.

Diagnostic Sets Updated

✅ lat_lon - Latitude-longitude plots
✅ polar - Polar region plots
✅ zonal_mean_xy - Zonal mean XY plots
✅ zonal_mean_2d - 2D zonal mean plots
✅ zonal_mean_2d_stratosphere - Stratosphere zonal mean (inherits from zonal_mean_2d)
✅ meridional_mean_2d - 2D meridional mean plots

Key Architecture Changes

Shared Utilities Module (driver/utils/time_slice.py)

Created new module with reusable functions:

validate_time_slice_format() - Validates slice string format
set_time_slice_name_yrs_attrs() - Sets parameter attributes for time slice mode
check_time_selection() - Ensures mutual exclusivity between seasons and time_slices

Core Data Loading Enhancement (driver/utils/dataset_xr.py)

Extended Dataset.get_climo_dataset() method:

Added is_time_slice boolean parameter
New _get_full_dataset() - Opens raw data without time averaging
New _apply_time_slice_to_dataset() - Applies index-based selection using xr.isel()
Maintains backward compatibility with existing climatology workflow

Parameter System Updates

CoreParameter (parameter/core_parameter.py): Added time_slices: list[str] attribute
CoreParser (parser/core_parser.py): Added --time_slices command-line argument
Granulation already handles mutual exclusivity

Driver Pattern Implementation

Each driver now follows consistent pattern:

1. Get time selection parameters

time_slices = getattr(parameter, "time_slices", [])
has_seasons, has_time_slices = check_time_selection(seasons, time_slices, require_one=True)

2. Determine mode and data source

time_selections = time_slices if has_time_slices else seasons
is_time_slice_mode = has_time_slices

3. Load data conditionally

for time_selection in time_selections:
if is_time_slice_mode:
ds = dataset.get_climo_dataset(var, time_selection, is_time_slice=True)
parameter._set_name_yrs_attrs(dataset, time_selection) # After loading
else:
parameter._set_name_yrs_attrs(dataset, time_selection) # Before loading
ds = dataset.get_climo_dataset(var, time_selection)

Type System Improvements

Updated type annotations to support ClimoFreq | str union types
Added proper type guards and type ignores where needed
Fixed mypy errors across codebase

Test Updates

Updated test_lat_lon_driver.py: Renamed _get_ref_climo_dataset → _get_ref_dataset with new
signature
Updated test_dataset_xr.py: Fixed type ignore comments

Benefits

Snapshot Analysis: Users can now analyze specific time indices instead of climatological averages
Stride Support: Time slice notation supports strides (e.g., "0:100:5" for every 5th timestep)
Consistency: Same time slice feature available across all major diagnostic sets
Backward Compatible: Existing season-based workflows unchanged
Code Reuse: Shared utilities eliminate duplication across 6 drivers

Documentation with usage examples are added for v3.1.0

Checklist

My code follows the style guidelines of this project
I have performed a self-review of my own code
My changes generate no new warnings
Any dependent changes have been merged and published in downstream modules

If applicable:

New and existing unit tests pass with my changes (locally and CI/CD build)
I have added tests that prove my fix is effective or that my feature works
I have commented my code, particularly in hard-to-understand areas
I have made corresponding changes to the documentation
I have noted that this is a breaking change for a major release (fix or feature that would cause existing functionality to not work as expected)

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]>

chengzhuzhang · 2025-10-20T20:23:49Z

@tomvothecoder This PR extended the snapshot support implemented in lat_lon_native to other core sets. As well as the documentation update for v3.1.0. Could you please review. Things to review include if the code updates are reasonable, also for documentation update, could you test the added examples ex8 (native-grid-visualization) and ex9 (snap-shot-analysis) following the READme in each folder? Each example has command-line option and script options.

Let's try to get this in a new release candidate for Unified. Thanks!

tomvothecoder · 2025-10-21T17:52:54Z

Sounds good @chengzhuzhang. I will review today.

Copilot

Pull Request Overview

This PR extends the snapshot analysis feature from lat_lon_native to six additional core diagnostic sets, enabling time slice-based analysis across multiple diagnostic types. The implementation introduces shared utilities for time slice handling and updates data loading infrastructure to support both climatological and snapshot modes.

Key Changes

Added time slice support to 6 core diagnostic sets (lat_lon, polar, zonal_mean_xy, zonal_mean_2d, meridional_mean_2d, zonal_mean_2d_stratosphere)
Created shared utility module driver/utils/time_slice.py for common time slice functionality
Enhanced Dataset.get_climo_dataset() with time slice mode via is_time_slice parameter

Reviewed Changes

Copilot reviewed 31 out of 31 changed files in this pull request and generated 5 comments.

Show a summary per file

File	Description
e3sm_diags/parameter/core_parameter.py	Added `time_slices` attribute and granulation support
e3sm_diags/parser/core_parser.py	Added `--time_slices` CLI argument
e3sm_diags/driver/utils/time_slice.py	New shared utilities for validation and attribute setting
e3sm_diags/driver/utils/dataset_xr.py	Enhanced data loading with time slice support
e3sm_diags/driver/*_driver.py	Implemented time slice pattern in 6 diagnostic drivers
e3sm_diags/viewer/*.py	Updated sorting to handle time slices numerically
examples/ex9-snapshot-analysis/*	New example demonstrating snapshot analysis
docs/source/*.rst	Documentation updates for v3.1.0 features

_{Tip: Customize your code reviews with copilot-instructions.md. Create the file or learn how to get started.}

Copilot · 2025-10-21T20:44:26Z

e3sm_diags/driver/utils/time_slice.py

+    >>> validate_time_slice_format("5")
+    >>> validate_time_slice_format("42")
+    """
+    pattern = r"^\d+$"


The docstring states that time slices 'must be non-negative integer indices', but the current pattern allows negative integers via the regex. Consider updating the pattern to explicitly reject negative values or clarifying the docstring.

fyi @chengzhuzhang

examples/ex9-snapshot-analysis/ex9.py

e3sm_diags/driver/utils/dataset_xr.py

e3sm_diags/viewer/default_viewer.py

docs/source/available-parameters.rst

tomvothecoder

My initial review. I am addressing my comments in an upcoming commit.

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

e3sm_diags/driver/utils/dataset_xr.py

tomvothecoder · 2025-10-21T20:55:03Z

e3sm_diags/driver/utils/time_slice.py

+    >>> validate_time_slice_format("5")
+    >>> validate_time_slice_format("42")
+    """
+    pattern = r"^\d+$"


fyi @chengzhuzhang

examples/ex9-snapshot-analysis/ex9.py

e3sm_diags/driver/utils/dataset_xr.py

docs/source/available-parameters.rst

examples/ex9-snapshot-analysis/README.md

tomvothecoder · 2025-10-21T21:02:37Z

README.md

+| 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            |


Thanks for including me!

e3sm_diags/parameter/lat_lon_native_parameter.py

e3sm_diags/driver/utils/time_slice.py

- 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

tomvothecoder

Hi @chengzhuzhang, here's my initial refactor. Please review my changes in this commit f0e5426 (#1013). You can review my summary of changes below while referencing the related file(s) for each of the main topics.

Results directory (successful): https://web.lcrc.anl.gov/public/e3sm/diagnostic_output/ac.zhang40/tests/1013-snapshot-analysis-core-sets/viewer/
My next actions are to add unit tests and address any lingering items.

Summary of Changes

This refactor standardizes how time-sliced datasets are retrieved across multiple drivers, removes redundant I/O code, and improves consistency around time_selection handling.

New Utility

File: driver/utils/io.py
- Added _get_xarray_datasets() to centralize logic for fetching xarray datasets based on time_selection.
- Simplifies driver implementations by providing a single entry point for time-sliced I/O.

Updated Dataset class

File: driver/utils/dataset_xy.py
- Added get_time_sliced_dataset() which implementation logic focused on time sliced datasets, extracted from get_climo_dataset().
- Removed time slice logic from get_climo_dataset(), including is_time_slice argument.

Driver Refactors

Files: driver/polar_driver.py, driver/meridional_mean_2d_driver.py, driver/zonal_mean_2d_driver.py, driver/zonal_mean_xy_driver.py
- Updated these drivers to use _get_xarray_datasets() instead of duplicating dataset-loading logic.
- Ensures consistent behavior across drivers and reduces maintenance burden.
Note: driver/lat_lon_driver.py still uses _get_ref_dataset() since it has a unique reference dataset workflow. This was left unchanged intentionally.

Parameter and API Updates

Files: parameter/core_parameter.py, parameter/lat_lon_native_parameter.py
- Added _get_time_selection_to_use() to simplify process of determining the time selection type ( time_slices or seasons) and corresponding values.
- Replaced season with time_selection and added the TimeSelection annotation for clarity.
- Moved _set_time_slice_name_yrs_attrs() back into LatLonNativeParameter, as it’s only used there -- added a note suggesting future refactoring or consolidation.

Code Cleanup

File Removed: driver/utils/time_slice.py
- All functionality migrated into CoreParameter methods, eliminating redundancy and streamlining time-slice handling logic.

Overall Impact

Unified I/O logic for time-sliced dataset retrieval.
Reduced duplication across major drivers.
Improved naming and parameter clarity (time_selection).
Preserved lat_lon_driver’s specialized reference dataset logic.

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

tomvothecoder

Added unit tests in 38f05f9 (#1013). The test script ran successfully with all of my latest changes: https://web.lcrc.anl.gov/public/e3sm/diagnostic_output/ac.zhang40/tests/1013-snapshot-analysis-core-sets/viewer/. The GH Actions build is passing too.

I think this PR is ready for a final self-review and merge.

Copilot

Pull Request Overview

Copilot reviewed 34 out of 34 changed files in this pull request and generated 6 comments.

_{Tip: Customize your code reviews with copilot-instructions.md. Create the file or learn how to get started.}

e3sm_diags/parameter/lat_lon_native_parameter.py

e3sm_diags/driver/utils/dataset_xr.py

e3sm_diags/driver/lat_lon_driver.py

e3sm_diags/driver/utils/io.py

docs/source/available-parameters.rst

e3sm_diags/parameter/core_parameter.py

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

chengzhuzhang · 2025-10-22T21:39:05Z

@tomvothecoder thank you for the improvement for the PR. I'm doing a full run test, and will merge if things look okay.

chengzhuzhang · 2025-10-22T22:28:23Z

Test results looks good. Merge.

tomvothecoder · 2025-10-23T17:12:01Z

@chengzhuzhang Thanks! Great work on this PR.

chengzhuzhang and others added 9 commits October 17, 2025 16:43

Extend snapshot support for other core sets

32161ad

add new util module

c72549d

fix tests

47ff468

fix mean_2d_viewer

ca5c585

add testing scripts

5a81319

clean up; clarify on input format for time_slices

d0bf718

update testing scripts

294648b

update README.md for 3.1.0 new feature

5acb1eb

chengzhuzhang marked this pull request as ready for review October 19, 2025 20:31

fix examples; update default set for lat_lon_native

1d9c55e

chengzhuzhang requested a review from tomvothecoder October 20, 2025 20:19

chengzhuzhang changed the title ~~Extend snapshot support for other core sets~~ Extend snapshot support for other core sets; Documentation update for v.3.1.0 Oct 20, 2025

tomvothecoder requested a review from Copilot October 21, 2025 20:43

Copilot AI reviewed Oct 21, 2025

View reviewed changes

tomvothecoder changed the title ~~Extend snapshot support for other core sets; Documentation update for v.3.1.0~~ Extend snapshot support for other core sets and doc updates Oct 21, 2025

tomvothecoder reviewed Oct 21, 2025

View reviewed changes

tomvothecoder added 2 commits October 22, 2025 14:15

Add unit tests

38f05f9

- 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

9b598a1

tomvothecoder requested a review from Copilot October 22, 2025 19:25

tomvothecoder approved these changes Oct 22, 2025

View reviewed changes

Copilot AI reviewed Oct 22, 2025

View reviewed changes

tomvothecoder and others added 3 commits October 22, 2025 12:27

Update e3sm_diags/driver/lat_lon_driver.py

cd31762

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

Apply Copilot review suggestions

dc37922

minor doc change

2312b65

chengzhuzhang merged commit 813bc41 into main Oct 22, 2025
5 checks passed

chengzhuzhang deleted the feature/time-slice-support-core-sets branch October 22, 2025 22:28

chengzhuzhang mentioned this pull request Oct 22, 2025

Support short simulations #652

Closed

		\| 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 \|

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

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

Uh oh!

Conversation

chengzhuzhang commented Oct 17, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Description

1. Get time selection parameters

2. Determine mode and data source

3. Load data conditionally

Checklist

Uh oh!

chengzhuzhang commented Oct 20, 2025

Uh oh!

tomvothecoder commented Oct 21, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull Request Overview

Key Changes

Reviewed Changes

Uh oh!

Copilot AI Oct 21, 2025

Choose a reason for hiding this comment

Uh oh!

tomvothecoder Oct 21, 2025

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

tomvothecoder left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

tomvothecoder Oct 21, 2025

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

tomvothecoder Oct 21, 2025

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

tomvothecoder left a comment • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Summary of Changes

Uh oh!

tomvothecoder left a comment

Choose a reason for hiding this comment

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull Request Overview

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

chengzhuzhang commented Oct 22, 2025

Uh oh!

Uh oh!

chengzhuzhang commented Oct 22, 2025

Uh oh!

tomvothecoder commented Oct 23, 2025

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

chengzhuzhang commented Oct 17, 2025 •

edited

Loading

tomvothecoder commented Oct 21, 2025 •

edited

Loading

tomvothecoder left a comment •

edited

Loading