New Feature: Support plotting native grid data in lat_lon plots #968

chengzhuzhang · 2025-04-16T22:59:21Z

Description

This feature is to address one of the new feature code review requirement: certain patterns (e.g., checkerboard) can only be detected on native grids. Developers have individual approaches/scripts for this capability. The code review required standard tests to provide a lat-lon view of a subset of these fields for the modeler to view along with standard seasonal average lat-lon images. Details see (internal E3SM doc):
https://acme-climate.atlassian.net/wiki/spaces/EIDMG/pages/3330179146/Native+Grid+Visualization

This PR will start with adding template scripts; and make the capability available through e3sm_diags command-line, as well as run script.

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)

whannah1 · 2025-04-23T22:53:58Z

To see that you're working on this makes me so happy!

chengzhuzhang · 2025-04-23T23:16:13Z

To see that you're working on this makes me so happy!

Great to hear!! I know you've been a strong advocate for this feature in standard analysis and I was going to reach out to you about the variables and regions we should normally look at, especially with checkerboard issue in mind.. Or maybe we should just support the same set of regular lat-lon variables (i.e. PRECT, PRECC, TREFT and surface fluxes etc.), but add additional regions in addition to global maps?

whannah1 · 2025-04-24T14:13:42Z

I was going to reach out to you about the variables and regions we should normally look at, especially with checkerboard issue in mind.. Or maybe we should just support the same set of regular lat-lon variables (i.e. PRECT, PRECC, TREFT and surface fluxes etc.), but add additional regions in addition to global maps?

I've found that precip and liquid water path (TGCLDLWP) are good variables for visualizing the checkerboard issue. The surface variables like reference temperature and latent/sensible flux usually don't exhibit small scale grid noise, but it wouldn't hurt to include those. Global maps seem fine. Maybe in HR cases we might want regional maps to focus on specific issues like the precip biases around central America, but we can add that as the need arises.

tomvothecoder

Hi @chengzhuzhang, I made some initial comments for review. Can you provide update the paths in the scripts so that I can run it? They are currently local paths to your machine.

Once I can run the code, I can do a deeper dive in my review and look at refactoring opportunities.

tomvothecoder · 2025-08-28T18:23:23Z

e3sm_diags/derivations/derivations.py

+    "TGCLDLWP": OrderedDict(
+        [
+            (
+                ("TGCLDLWP",),
+                lambda x: convert_units(x, target_units="g/m^2"),
+            ),
+            (
+                ("LiqWaterPath",),
+                lambda x: convert_units(x, target_units="g/m^2"),
+            ),  # EAMxx
+        ]
+    ),


Replace with regular dictionary

tomvothecoder · 2025-08-28T18:23:57Z

conda-env/dev.yml

  - numpy >=2.0.0,<3.0.0
  - pywavelets
  - scipy
+  - uxarray >=2023.3.0,<2025.6.0


Any reason why we're constraining <2025.6.0 here? Also ci.yml doesn't have this constraint.

tomvothecoder · 2025-08-28T18:24:09Z

e3sm_diags/driver/default_diags/lat_lon_native_model_vs_model.cfg

+
+

Suggested change

tomvothecoder · 2025-08-28T18:24:19Z

e3sm_diags/driver/default_diags/lat_lon_native_model_vs_model.cfg

+contour_levels = [-170, -150, -135, -120, -105, -90, -75, -60, -45, -30, -15, 0, 15, 30, 45]
+diff_levels = [-30, -25, -20, -15, -10, -5, -2, 2, 5, 10, 15, 20, 25, 30]
+
+


Suggested change

tomvothecoder · 2025-08-28T18:24:26Z

e3sm_diags/driver/default_diags/lat_lon_native_model_vs_model.cfg

+contour_levels = [-100, -75, -50, -25, -10, 0, 10, 25, 50, 75, 100, 125, 150]
+diff_levels = [-75, -50, -25, -10, -5, -2, 2, 5, 10, 25, 50, 75]
+
+


Suggested change

e3sm_diags/viewer/main.py

e3sm_diags/driver/utils/dataset_xr.py

e3sm_diags/driver/lat_lon_native_driver.py

e3sm_diags/parameter/lat_lon_native_parameter.py

chengzhuzhang · 2025-08-29T19:25:54Z

@tomvothecoder thank you for starting the initial review. I have a new commit to push the run script and cfg file for testing.
I had to pin uxarray version (the stable verson), because I notice that one plotting feature (e.g. periodic_elements in pc = uxds["PRECT"].to_polycollection(periodic_elements="split")) is no longer available in latest version. Pinging @philipc2 if this change is permanent or perhaps can guide us if we need a different approach for plotting..Thanks!

tomvothecoder · 2025-09-11T17:07:26Z

Hey @chengzhuzhang, are you aiming to get this feature in a new version before the upcoming E3SM Unified release? Or does it require more work over a longer timeline? This helps me understand the task priority. I'll review again once Chrysalis is back online.

chengzhuzhang · 2025-09-11T17:10:12Z

@tomvothecoder I think this PR is ready, please help review. in the mean time, I will do more testing as well.

philipc2 · 2025-09-11T17:15:08Z

@tomvothecoder thank you for starting the initial review. I have a new commit to push the run script and cfg file for testing. I had to pin uxarray version (the stable verson), because I notice that one plotting feature (e.g. periodic_elements in pc = uxds["PRECT"].to_polycollection(periodic_elements="split")) is no longer available in latest version. Pinging @philipc2 if this change is permanent or perhaps can guide us if we need a different approach for plotting..Thanks!

Hi @chengzhuzhang

Apologies for not getting back to you sooner about this. We have deprecated the periodic_elements parameter and added a fix to handle the deprecation more smoothly (see UXARRAY/uxarray#1351)

This feature will be released on Tuesday in our September release, in addition to some other Matplotlib/Cartopy plotting improvements UXARRAY/uxarray#1354

chengzhuzhang · 2025-09-11T19:18:42Z

@philipc2 thank you for confirming and for the additional information on new plotting improvement.
I will remove the dependency on periodic_elements and remove uxarray version specification.

Copilot

Pull Request Overview

This PR adds support for plotting native grid data in lat_lon plots to address visualization requirements for detecting patterns (e.g., checkerboard) that can only be detected on native grids. The feature provides standard tests for modelers to view lat-lon representations of native grid fields alongside seasonal average lat-lon images.

Key changes:

New lat_lon_native diagnostic set with plotting capabilities for unstructured grid data
Command-line interface and configuration support for native grid visualization
Model-only and model-vs-model comparison support

Reviewed Changes

Copilot reviewed 24 out of 25 changed files in this pull request and generated 10 comments.

Show a summary per file

File	Description
e3sm_diags/plot/lat_lon_native_plot.py	Core plotting functionality for native grid visualization using uxarray
e3sm_diags/driver/lat_lon_native_driver.py	Driver logic for native grid diagnostics and data processing
e3sm_diags/parameter/lat_lon_native_parameter.py	Parameter class with native grid-specific configuration options
pyproject.toml	Added uxarray dependency and data file configuration
e3sm_diags/viewer/	Updated viewer components to support lat_lon_native diagnostic set

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

e3sm_diags/viewer/default_viewer.py

e3sm_diags/parser/core_parser.py

auxiliary_tools/LCF_mixed-phase_partition_diagram.ipynb

tests/test_native_grid.cfg

run_native_grid_test.py

tomvothecoder · 2025-09-15T17:21:16Z

@tomvothecoder I think this PR is ready, please help review. in the mean time, I will do more testing as well.

Hi @chengzhuzhang, I am nearly done refactoring and reviewing this PR. I just have the driver and plotter left. I will test my code changes with the full run, push the commit when ready, and tag you for another review + test run on your end.

add testing script and cfg file Move METRICS_DEFAULT_VALUE to /driver/__init__.py Refactor `default_viewer.py` and `main.py` Move `run_native_grid_test.py` to `debug/968-native-grid-vis` Refactor `lat_lon_native_parameter.py` and `core_parameter.py` - Update type annotations, re-order methods in `LatLonNativeParameter` class Add a comment for granulation in `core_parser.py` Extract code from LatLonNativeDriver into NativeDataset class - Move `TimeSlice` and `TimeSelection` into `type_annotations.py` and update all references Add type annotations and comments Add fixme and todo comments in driver and plotter Update run script to use single run

tomvothecoder

@chengzhuzhang, here is my PR review. I currently have a full test run going with these changes. I will report back with the status soon.

EDIT: Full run complete here: https://web.lcrc.anl.gov/public/e3sm/diagnostic_output/ac.tvo/tests/lat_lon_native_file/viewer/index.html

I think it is important to refactor lat_lon_native_plot.py to re-use existing utilities, as there is a lot of repeated code from the plotter utilities. I left FIXME comments in that file to section off parts that should be refactored. I currently don't have bandwidth to dive deeper to refactor this file.

Changes:

Extracted NativeDataset class from lat_lon_native_driver.py to store native-grid related functionalities and removed unused functions (e.g., metrics functions, processing test dataset)
Added FIXME comments through lat_lon_native_driver.py and lat_lon_native_plot.py to refactor code and re-use utilities
Refactored methods in class Dataset which now reference new methods in CoreParameter

tomvothecoder · 2025-09-15T21:56:15Z

e3sm_diags/driver/lat_lon_native_driver.py

+    test_ds = NativeDataset(parameter, data_type="test")
+    ref_ds = NativeDataset(parameter, data_type="ref")


Use NativeDataset class instead of Dataset class.

tomvothecoder · 2025-09-15T22:03:36Z

e3sm_diags/plot/lat_lon_native_plot.py

+        # ------------------------------------------------------------
+        # FIXME: Metrics extraction for test, ref, and diff datasets are duplicated,
+        # extract to a helper function.
+        # ------------------------------------------------------------


tomvothecoder · 2025-09-15T22:03:42Z

e3sm_diags/plot/lat_lon_native_plot.py

+        # ------------------------------------------------------------
+        # FIXME: Metrics extraction for test, ref, and diff datasets are duplicated,
+        # extract to a helper function.
+        # ------------------------------------------------------------


tomvothecoder · 2025-09-15T22:03:48Z

e3sm_diags/plot/lat_lon_native_plot.py

+            # ------------------------------------------------------------
+            # FIXME: Metrics extraction for test, ref, and diff datasets are duplicated,
+            # extract to a helper function.
+            # ------------------------------------------------------------


tomvothecoder · 2025-09-15T22:03:54Z

e3sm_diags/plot/lat_lon_native_plot.py

+                    # ------------------------------------------------------------
+                    # FIXME: Re-use utils.add_rmse_corr_text here
+                    # ------------------------------------------------------------


tomvothecoder · 2025-09-15T22:04:36Z

e3sm_diags/driver/lat_lon_native_driver.py

+                # TODO: What is this supposed to do? Just check if the variables
+                # can be derived? It returns a bool but the return value is never
+                # used downstream.
+                test_ds._process_variable_derivations(var_key)


@chengzhuzhang What is this supposed to do?

This explicit call in the driver actually derive variables which is not covered by dataset_native class. The comment is changed to Apply variable derivations if needed.

tomvothecoder · 2025-09-15T22:04:45Z

e3sm_diags/driver/lat_lon_native_driver.py

+                    # TODO: What is this supposed to do? Just check if the variables
+                    # can be derived? It returns a bool but the return value is never
+                    # used downstream.
+                    ref_ds._process_variable_derivations(var_key)


What is this supposed to do?

same as above but for reference data.

tomvothecoder · 2025-09-15T22:05:06Z

e3sm_diags/driver/utils/dataset_native.py

+class NativeDataset:
+    """
+    A class for handling native grid datasets using xarray for raw data
+    and uxarray for grid-aware operations.
+
+    NOTE: NativeDataset uses composition instead of inheritance to wrap Dataset
+    with additional native-grid specific functionalities. It does not inherit
+    from Dataset to avoid confusion with existing Dataset methods and prevent
+    tight coupling. If needed, we can refactor to inherit from Dataset or
+    create a parent abstract class in the future.
+    """


Note about this new class.

tomvothecoder · 2025-09-15T22:05:36Z

e3sm_diags/driver/utils/dataset_native.py

+    A class for handling native grid datasets using xarray for raw data
+    and uxarray for grid-aware operations.


TODO: Add unit tests for this class.

tomvothecoder · 2025-09-15T22:06:09Z

e3sm_diags/parameter/core_parameter.py

+    def _add_time_series_file_path_attr(
+        self,
+        data_type: Literal["test", "ref"],
+        ds: xr.Dataset,
+    ):
+        """Add file path attributes to the parameter object.
+
+        Parameters
+        ----------
+        data_type : Literal["test", "ref"]
+            The type of data, either "test" or "ref".
+        ds : xr.Dataset
+            The dataset object containing the file path attribute.
+
+        Raises
+        ------
+        ValueError
+            If `data_type` is not "test" or "ref".
+        """
+        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"
+
+        setattr(self, file_path_attr, getattr(ds, "file_path", "Unknown"))
+
+    def _add_climatology_file_path_attr(
+        self,
+        data_type: Literal["test", "ref"],
+        filepath: str | None = None,
+    ):
+        """Add file path attributes to the parameter object.
+
+        Parameters
+        ----------
+        data_type : Literal["test", "ref"]
+            The type of data, either "test" or "ref".
+        filepath : str | None, optional
+            The file path for climatology data.
+
+        Raises
+        ------
+        ValueError
+            If `data_type` is not "test" or "ref".
+        ValueError
+            If `filepath` is not provided for climatology data.
+        """
+        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"
+
+        if not filepath:
+            raise ValueError("Filepath must be provided for climatology data.")
+
+        setattr(self, file_path_attr, os.path.abspath(filepath))


New methods in CoreParameter used in class Dataset

chengzhuzhang · 2025-09-17T22:59:14Z

@tomvothecoder thanks for helping reviewing and refactoring the code! I tested with a couple of configuration and had a minor fix for the case of using season DJF as input. Next, I will relax the uxarray pinned version by removing the dependency on a deprecated function. And will try to refactor the plotting functions to following what we have for lat-lon. I may tag you for another review once the changes are done.

… into native_grids_viz

tomvothecoder · 2025-09-18T18:49:01Z

e3sm_diags/driver/utils/dataset_native.py

+
+        Parameters
+        ----------
+


Suggested change

… into native_grids_viz

chengzhuzhang · 2025-10-05T18:05:26Z

Merging for a release candidate test of e3sm_unified.
Additional documentation updates and final refinements will be submitted in a follow-up PR before the official release.

chengzhuzhang force-pushed the native_grids_viz branch from 0d033db to 18771b7 Compare April 18, 2025 02:43

chengzhuzhang added 18 commits July 28, 2025 18:15

add notebook to visualize native grid data

7bc760c

add skeleton for a the new lat_lon_native set

cc21bd2

add dependencies; make simple plotting/viewer work

e902660

fix variable description in viewer

5e2741b

make variable derivation work

c33df95

make 2d model only run plotting work

cb49bf5

model vs model works

4740a8b

fix for region e.g. 60S60N work

108a660

add uxarray's nearest neighbor regridding and make diff plot work

05cb53a

clean up debug lines

1b2c5ff

reorg plot funcs

f3a9a89

create subset region metrics

5291d78

update standard variable list for native grid plotting

99fdf26

fix sets in cfg

fefb25f

more code clean up

bab6b04

temporal fix for input data with n time snapshots

2898e8d

add TGCLDLWP variable

bb7dd22

update default var in cfg

4605fac

chengzhuzhang force-pushed the native_grids_viz branch from 458b14a to 4605fac Compare July 28, 2025 23:23

chengzhuzhang added 6 commits July 28, 2025 18:50

fix pre-commit errors

da3b9a3

pin uxarray to avoid latest version(2025.6.0)

4340524

simplify derived var logic;code cleanup

399a95e

create new get_native_dataset function

8f9521e

add time snapshot selection function for visualize native history file

8fa1800

update granulation list to avoid looping over seasons; fix viewer

60f1084

tomvothecoder reviewed Aug 28, 2025

View reviewed changes

chengzhuzhang marked this pull request as ready for review September 11, 2025 17:09

tomvothecoder requested a review from Copilot September 12, 2025 19:36

Copilot AI reviewed Sep 12, 2025

View reviewed changes

tomvothecoder force-pushed the native_grids_viz branch from 5c3a5c8 to 0233432 Compare September 15, 2025 21:52

tomvothecoder reviewed Sep 15, 2025

View reviewed changes

fix importing ClimoFreq and using it

54ab380

chengzhuzhang added 4 commits September 17, 2025 18:29

relax uxarray version by remove dependency on periodic_elements

4f70de2

Merge branch 'main' into native_grids_viz

8e5b847

refactor to use utils in plotting

d887088

Merge branch 'native_grids_viz' of github.com:E3SM-Project/e3sm_diags…

8944b82

… into native_grids_viz

tomvothecoder reviewed Sep 18, 2025

View reviewed changes

e3sm_diags/driver/utils/dataset_native.py

Parameters

----------

Copy link

Collaborator

tomvothecoder Sep 18, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Suggested change

chengzhuzhang added 8 commits September 18, 2025 15:04

refactor metrics generation and using weighted mean,rsme,corr

8b72480

reformat testing run script

5a71782

make metrics work on regional subset

3af5471

Merge branch 'main' into native_grids_viz

bf31962

clean up

ed2b5f0

update example script

33ba31e

Merge branch 'native_grids_viz' of github.com:E3SM-Project/e3sm_diags…

11f8cd9

… into native_grids_viz

fix provenance command-line; fix metrics for model_only case

2cca3ec

chengzhuzhang merged commit e8a7412 into main Oct 5, 2025
5 checks passed

chengzhuzhang deleted the native_grids_viz branch October 5, 2025 18:23

		contour_levels = [-170, -150, -135, -120, -105, -90, -75, -60, -45, -30, -15, 0, 15, 30, 45]
		diff_levels = [-30, -25, -20, -15, -10, -5, -2, 2, 5, 10, 15, 20, 25, 30]

		contour_levels = [-100, -75, -50, -25, -10, 0, 10, 25, 50, 75, 100, 125, 150]
		diff_levels = [-75, -50, -25, -10, -5, -2, 2, 5, 10, 25, 50, 75]

		test_ds = NativeDataset(parameter, data_type="test")
		ref_ds = NativeDataset(parameter, data_type="ref")

		A class for handling native grid datasets using xarray for raw data
		and uxarray for grid-aware operations.

New Feature: Support plotting native grid data in lat_lon plots #968

New Feature: Support plotting native grid data in lat_lon plots #968

Uh oh!

Conversation

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

Uh oh!

Description

Checklist

Uh oh!

whannah1 commented Apr 23, 2025

Uh oh!

chengzhuzhang commented Apr 23, 2025

Uh oh!

whannah1 commented Apr 24, 2025

Uh oh!

tomvothecoder left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

chengzhuzhang commented Aug 29, 2025

Uh oh!

tomvothecoder commented Sep 11, 2025

Uh oh!

chengzhuzhang commented Sep 11, 2025

Uh oh!

philipc2 commented Sep 11, 2025

Uh oh!

chengzhuzhang commented Sep 11, 2025

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull Request Overview

Reviewed Changes

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

tomvothecoder commented Sep 15, 2025

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

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

chengzhuzhang commented Apr 16, 2025 •

edited

Loading

tomvothecoder left a comment •

edited

Loading