CDAT Migration Phase 2: Refactor core utilities and lat_lon set (#677)

Refer to the PR for more information because the changelog is massive. Update build workflow to run on `cdat-migration-fy24` branch CDAT Migration Phase 2: Add CDAT regression test notebook template and fix GH Actions build (#743) - Add Makefile for quick access to multiple Python-based commands such as linting, testing, cleaning up cache and build files - Fix some lingering unit tests failure - Update `xcdat=0.6.0rc1` to `xcdat >=0.6.0` in `ci.yml`, `dev.yml` and `dev-nompi.yml` - Add `xskillscore` to `ci.yml` - Fix `pre-commit` issues CDAT Migration Phase 2: Regression testing for `lat_lon`, `lat_lon_land`, and `lat_lon_river` (#744) - Add Makefile that simplifies common development commands (building and installing, testing, etc.) - Write unit tests to cover all new code for utility functions - `dataset_xr.py`, `metrics.py`, `climo_xr.py`, `io.py`, `regrid.py` - Metrics comparison for `cdat-migration-fy24` `lat_lon` and `main` branch of `lat_lon` -- `NET_FLUX_SRF` and `RESTOM` have the highest spatial average diffs - Test run with 3D variables (`_run_3d_diags()`) - Fix Python 3.9 bug with using pipe command to represent Union -- doesn't work with `from __future__ import annotations` still - Fix subsetting syntax bug using ilev - Fix regridding bug where a single plev is passed and xCDAT does not allow generating bounds for coordinates of len <= 1 -- add conditional that just ignores adding new bounds for regridded output datasets, fix related tests - Fix accidentally calling save plots and metrics twice in `_get_metrics_by_region()` - Fix failing integration tests pass in CI/CD - Refactor `test_diags.py` -- replace unittest with pytest - Refactor `test_all_sets.py` -- replace unittest with pytest - Test climatology datasets -- tested with 3d variables using `test_all_sets.py` CDAT Migration Phase 2: Refactor utilities and CoreParameter methods for reusability across diagnostic sets (#746) - Move driver type annotations to `type_annotations.py` - Move `lat_lon_driver._save_data_metrics_and_plots()` to `io.py` - Update `_save_data_metrics_and_plots` args to accept `plot_func` callable - Update `metrics.spatial_avg` to return an optionally `xr.DataArray` with `as_list=False` - Move `parameter` arg to the top in `lat_lon_plot.plot` - Move `_set_param_output_attrs` and `_set_name_yr_attrs` from `lat_lon_driver` to `CoreParameter` class Regression testing for lat_lon variables `NET_FLUX_SRF` and `RESTOM` (#754) Update regression test notebook to show validation of all vars Add `subset_and_align_datasets()` to regrid.py (#776) Add template run scripts CDAT Migration Phase: Refactor `cosp_histogram` set (#748) - Refactor `cosp_histogram_driver.py` and `cosp_histogram_plot.py` - `formulas_cosp.py` (new file) - Includes refactored, Xarray-based `cosp_histogram_standard()` and `cosp_bin_sum()` functions - I wrote a lot of new code in `formulas_cosp.py` to clean up `derivations.py` and the old equivalent functions in `utils.py` - `derivations.py` - Cleaned up portions of `DERIVED_VARIABLES` dictionary - Removed unnecessary `OrderedDict` usage for `cosp_histogram` related variables (we should do this for the rest of the variables in in #716) - Remove unnecessary `convert_units()` function calls - Move cloud levels passed to derived variable formulas to `formulas_cosp.CLOUD_BIN_SUM_MAP` - `utils.py` - Delete deprecated, CDAT-based `cosp_histogram` functions - `dataset_xr.py` - Add `dataset_xr.Dataset._open_climo_dataset()` method with a catch for dataset quality issues where "time" is a scalar variable that does not match the "time" dimension array length, drops this variable and replaces it with the correct coordinate - Update `_get_dataset_with_derivation_func()` to handle derivation functions that require the `xr.Dataset` and `target_var_key` args (e.g., `cosp_histogram_standardize()` and `cosp_bin_sum()`) - `io.py` - Update `_write_vars_to_netcdf()` to write test, ref, and diff variables to individual netCDF (required for easy comparison to CDAT-based code that does the same thing) - Add `cdat_migration_regression_test_netcdf.ipynb` validation notebook template for comparing `.nc` files CDAT Migration Phase 2: Refactor `zonal_mean_2d()` and `zonal_mean_2d_stratosphere()` sets (#774) Refactor 654 zonal mean xy (#752) Co-authored-by: Tom Vo <tomvothecoder@gmail.com> CDAT Migration - Update run script output directory to NERSC public webserver (#793) [PR]: CDAT Migration: Refactor `aerosol_aeronet` set (#788) CDAT Migration: Test `lat_lon` set with run script and debug any issues (#794) CDAT Migration: Refactor `polar` set (#749) Co-authored-by: Tom Vo <tomvothecoder@gmail.com> Align order of calls to `_set_param_output_attrs` CDAT Migration: Refactor `meridional_mean_2d` set (#795) CDAT Migration: Refactor `aerosol_budget` (#800) Add `acme.py` changes from PR #712 (#814) * Add `acme.py` changes from PR #712 * Replace unnecessary lambda call Refactor area_mean_time_series and add ccb slice flag feature (#750) Co-authored-by: Tom Vo <tomvothecoder@gmail.com> [Refactor]: Validate fix in PR #750 for #759 (#815) CDAT Migration Phase 2: Refactor `diurnal_cycle` set (#819) CDAT Migration: Refactor annual_cycle_zonal_mean set (#798) * Refactor `annual_cycle_zonal_mean` set * Address PR review comments * Add lat lon regression testing * Add debugging scripts * Update `_open_climo_dataset()` to decode times as workaround to misaligned time coords - Update `annual_cycle_zonal_mean_plot.py` to convert time coordinates to month integers * Fix unit tests * Remove old plotter * Add script to debug decode_times=True and ncclimo file * Update plotter time values to month integers * Fix slow `.load()` and multiprocessing issue - Due to incorrectly updating `keep_bnds` logic - Add `_encode_time_coords()` to workaround cftime issue `ValueError: "months since" units only allowed for "360_day" calendar` * Update `_encode_time_coords()` docstring * Add AODVIS debug script * update AODVIS obs datasets; regression test results --------- Co-authored-by: Tom Vo <tomvothecoder@gmail.com> CDAT Migration Phase 2: Refactor `qbo` set (#826) CDAT Migration Phase 2: Refactor tc_analysis set (#829) * start tc_analysis_refactor * update driver * update plotting * Clean up plotter - Remove unused variables - Make `plot_info` a constant called `PLOT_INFO`, which is now a dict of dicts - Reorder functions for top-down readability * Remove unused notebook --------- Co-authored-by: tomvothecoder <tomvothecoder@gmail.com> CDAT Migration Phase 2: Refactor `enso_diags` set (#832) CDAT Migration Phase 2: Refactor `streamflow` set (#837) [Bug]: CDAT Migration Phase 2: enso_diags plot fixes (#841) [Refactor]: CDAT Migration Phase 3: testing and documentation update (#846) CDAT Migration Phase 3 - Port QBO Wavelet feature to Xarray/xCDAT codebase (#860) CDAT Migration Phase 2: Refactor arm_diags set (#842) Add performance benchmark material (#864) Add function to add CF axis attr to Z axis if missing for downstream xCDAT operations (#865) CDAT Migration Phase 3: Add Convective Precipitation Fraction in lat-lon (#875) CDAT Migration Phase 3: Fix LHFLX name and add catch for non-existent or empty TE stitch file (#876) Add support for time series datasets via glob and fix `enso_diags` set (#866) Add fix for checking `is_time_series()` property based on `data_type` attr (#881) CDAT migration: Fix African easterly wave density plots in TC analysis and convert H20LNZ units to ppm/volume (#882) CDAT Migration: Update `mp_partition_driver.py` to use Dataset from `dataset_xr.py` (#883) CDAT Migration - Port JJB tropical subseasonal diags to Xarray/xCDAT (#887) CDAT Migration: Prepare branch for merge to `main` (#885) [Refactor]: CDAT Migration - Update dependencies and remove Dataset._add_cf_attrs_to_z_axes() (#891) CDAT Migration Phase 2: Refactor core utilities and `lat_lon` set (#677) Refer to the PR for more information because the changelog is massive. Update build workflow to run on `cdat-migration-fy24` branch CDAT Migration Phase 2: Add CDAT regression test notebook template and fix GH Actions build (#743) - Add Makefile for quick access to multiple Python-based commands such as linting, testing, cleaning up cache and build files - Fix some lingering unit tests failure - Update `xcdat=0.6.0rc1` to `xcdat >=0.6.0` in `ci.yml`, `dev.yml` and `dev-nompi.yml` - Add `xskillscore` to `ci.yml` - Fix `pre-commit` issues CDAT Migration Phase 2: Regression testing for `lat_lon`, `lat_lon_land`, and `lat_lon_river` (#744) - Add Makefile that simplifies common development commands (building and installing, testing, etc.) - Write unit tests to cover all new code for utility functions - `dataset_xr.py`, `metrics.py`, `climo_xr.py`, `io.py`, `regrid.py` - Metrics comparison for `cdat-migration-fy24` `lat_lon` and `main` branch of `lat_lon` -- `NET_FLUX_SRF` and `RESTOM` have the highest spatial average diffs - Test run with 3D variables (`_run_3d_diags()`) - Fix Python 3.9 bug with using pipe command to represent Union -- doesn't work with `from __future__ import annotations` still - Fix subsetting syntax bug using ilev - Fix regridding bug where a single plev is passed and xCDAT does not allow generating bounds for coordinates of len <= 1 -- add conditional that just ignores adding new bounds for regridded output datasets, fix related tests - Fix accidentally calling save plots and metrics twice in `_get_metrics_by_region()` - Fix failing integration tests pass in CI/CD - Refactor `test_diags.py` -- replace unittest with pytest - Refactor `test_all_sets.py` -- replace unittest with pytest - Test climatology datasets -- tested with 3d variables using `test_all_sets.py` CDAT Migration Phase 2: Refactor utilities and CoreParameter methods for reusability across diagnostic sets (#746) - Move driver type annotations to `type_annotations.py` - Move `lat_lon_driver._save_data_metrics_and_plots()` to `io.py` - Update `_save_data_metrics_and_plots` args to accept `plot_func` callable - Update `metrics.spatial_avg` to return an optionally `xr.DataArray` with `as_list=False` - Move `parameter` arg to the top in `lat_lon_plot.plot` - Move `_set_param_output_attrs` and `_set_name_yr_attrs` from `lat_lon_driver` to `CoreParameter` class CDAT Migration Phase 2: Refactor `zonal_mean_2d()` and `zonal_mean_2d_stratosphere()` sets (#774) CDAT Migration Phase 2: Refactor `qbo` set (#826)
E3SM-Project · Dec 5, 2024 · a60ee3f · a60ee3f
1 parent 9f021d0
commit a60ee3f
Show file tree

Hide file tree

Showing 13 changed files with 2,094 additions and 4 deletions.
diff --git a/.github/workflows/build_workflow.yml b/.github/workflows/build_workflow.yml
@@ -5,7 +5,7 @@ on:
     branches: [main]
 
   pull_request:
-    branches: [main]
+    branches: [main, cdat-migration-fy24]
 
   workflow_dispatch:
 

diff --git a/auxiliary_tools/template_cdat_regression_test.ipynb b/auxiliary_tools/template_cdat_regression_test.ipynb
diff --git a/e3sm_diags/driver/aerosol_budget_driver.py b/e3sm_diags/driver/aerosol_budget_driver.py
@@ -3,6 +3,7 @@
 script is integrated in e3sm_diags by Jill Zhang, with input from Kai Zhang,
 Taufiq Hassan, Xue Zheng, Ziming Ke, Susannah Burrows, and Naser Mahfouz.
 """
+
 from __future__ import annotations
 
 import csv

diff --git a/e3sm_diags/driver/qbo_driver.py b/e3sm_diags/driver/qbo_driver.py
@@ -125,6 +125,11 @@ def run_diag(parameter: QboParameter) -> QboParameter:
         test_dict["name"] = test_ds._get_test_name()
         ref_dict["name"] = ref_ds._get_ref_name()
 
+        try:
+            ref_dict["name"] = ref_ds._get_ref_name()
+        except AttributeError:
+            ref_dict["name"] = parameter.ref_name
+
         _save_metrics_to_json(parameter, test_dict, "test")  # type: ignore
         _save_metrics_to_json(parameter, ref_dict, "ref")  # type: ignore
 

diff --git a/e3sm_diags/driver/utils/climo_xr.py b/e3sm_diags/driver/utils/climo_xr.py
@@ -1,8 +1,8 @@
 """This module stores climatology functions operating on Xarray objects.
 
-
 This file will eventually be refactored to use xCDAT's climatology API.
 """
+
 from typing import Dict, List, Literal, get_args
 
 import numpy as np

diff --git a/e3sm_diags/driver/utils/dataset_xr.py b/e3sm_diags/driver/utils/dataset_xr.py
@@ -1,13 +1,13 @@
 """This module stores the Dataset class, which is the primary class for I/O.
 
-
 This Dataset class operates on `xr.Dataset` objects, which are created using
 netCDF files. These `xr.Dataset` contain either the reference or test variable.
 This variable can either be from a climatology file or a time series file.
 If the variable is from a time series file, the climatology of the variable is
 calculated. Reference and test variables can also be derived using other
 variables from dataset files.
 """
+
 from __future__ import annotations
 
 import collections

diff --git a/e3sm_diags/metrics/metrics.py b/e3sm_diags/metrics/metrics.py
@@ -1,4 +1,5 @@
 """This module stores functions to calculate metrics using Xarray objects."""
+
 from __future__ import annotations
 
 from typing import List, Literal

diff --git a/e3sm_diags/plot/cartopy/aerosol_aeronet_plot.py b/e3sm_diags/plot/cartopy/aerosol_aeronet_plot.py
@@ -0,0 +1,132 @@
+import os
+
+import cartopy.crs as ccrs
+import matplotlib
+import numpy as np
+
+from e3sm_diags.driver.utils.general import get_output_dir
+from e3sm_diags.logger import custom_logger
+from e3sm_diags.metrics import mean
+from e3sm_diags.plot.cartopy.deprecated_lat_lon_plot import plot_panel
+
+matplotlib.use("Agg")
+import matplotlib.pyplot as plt  # isort:skip  # noqa: E402
+
+logger = custom_logger(__name__)
+
+plotTitle = {"fontsize": 11.5}
+plotSideTitle = {"fontsize": 9.5}
+
+
+def plot(test, test_site, ref_site, parameter):
+    # Plot scatter plot
+    # Position and sizes of subplot axes in page coordinates (0 to 1)
+    # (left, bottom, width, height) in page coordinates
+    panel = [
+        (0.09, 0.40, 0.72, 0.30),
+        (0.19, 0.2, 0.62, 0.30),
+    ]
+    # Border padding relative to subplot axes for saving individual panels
+    # (left, bottom, right, top) in page coordinates
+    border = (-0.06, -0.03, 0.13, 0.03)
+
+    fig = plt.figure(figsize=parameter.figsize, dpi=parameter.dpi)
+    fig.suptitle(parameter.var_id, x=0.5, y=0.97)
+    proj = ccrs.PlateCarree()
+    max1 = test.max()
+    min1 = test.min()
+    mean1 = mean(test)
+    # TODO: Replace this function call with `e3sm_diags.plot.utils._add_colormap()`.
+    plot_panel(
+        0,
+        fig,
+        proj,
+        test,
+        parameter.contour_levels,
+        parameter.test_colormap,
+        (parameter.test_name_yrs, None, None),
+        parameter,
+        stats=(max1, mean1, min1),
+    )
+
+    ax = fig.add_axes(panel[1])
+    ax.set_title(f"{parameter.var_id} from AERONET sites")
+
+    # define 1:1 line, and x y axis limits
+
+    if parameter.var_id == "AODVIS":
+        x1 = np.arange(0.01, 3.0, 0.1)
+        y1 = np.arange(0.01, 3.0, 0.1)
+        plt.xlim(0.03, 1)
+        plt.ylim(0.03, 1)
+    else:
+        x1 = np.arange(0.0001, 1.0, 0.01)
+        y1 = np.arange(0.0001, 1.0, 0.01)
+        plt.xlim(0.001, 0.3)
+        plt.ylim(0.001, 0.3)
+
+    plt.loglog(x1, y1, "-k", linewidth=0.5)
+    plt.loglog(x1, y1 * 0.5, "--k", linewidth=0.5)
+    plt.loglog(x1 * 0.5, y1, "--k", linewidth=0.5)
+
+    corr = np.corrcoef(ref_site, test_site)
+    xmean = np.mean(ref_site)
+    ymean = np.mean(test_site)
+    ax.text(
+        0.3,
+        0.9,
+        f"Mean (test): {ymean:.3f} \n Mean (ref): {xmean:.3f}\n Corr: {corr[0, 1]:.2f}",
+        horizontalalignment="right",
+        verticalalignment="top",
+        transform=ax.transAxes,
+    )
+
+    # axis ticks
+    plt.tick_params(axis="both", which="major")
+    plt.tick_params(axis="both", which="minor")
+
+    # axis labels
+    plt.xlabel(f"ref: {parameter.ref_name_yrs}")
+    plt.ylabel(f"test: {parameter.test_name_yrs}")
+
+    plt.loglog(ref_site, test_site, "kx", markersize=3.0, mfc="none")
+
+    # legend
+    plt.legend(frameon=False, prop={"size": 5})
+
+    # TODO: This section can be refactored to use `plot.utils._save_plot()`.
+    for f in parameter.output_format:
+        f = f.lower().split(".")[-1]
+        fnm = os.path.join(
+            get_output_dir(parameter.current_set, parameter),
+            f"{parameter.output_file}" + "." + f,
+        )
+        plt.savefig(fnm)
+        logger.info(f"Plot saved in: {fnm}")
+
+    for f in parameter.output_format_subplot:
+        fnm = os.path.join(
+            get_output_dir(parameter.current_set, parameter),
+            parameter.output_file,
+        )
+        page = fig.get_size_inches()
+        i = 0
+        for p in panel:
+            # Extent of subplot
+            subpage = np.array(p).reshape(2, 2)
+            subpage[1, :] = subpage[0, :] + subpage[1, :]
+            subpage = subpage + np.array(border).reshape(2, 2)
+            subpage = list(((subpage) * page).flatten())  # type: ignore
+            extent = matplotlib.transforms.Bbox.from_extents(*subpage)
+            # Save subplot
+            fname = fnm + ".%i." % (i) + f
+            plt.savefig(fname, bbox_inches=extent)
+
+            orig_fnm = os.path.join(
+                get_output_dir(parameter.current_set, parameter),
+                parameter.output_file,
+            )
+            fname = orig_fnm + ".%i." % (i) + f
+            logger.info(f"Sub-plot saved in: {fname}")
+
+            i += 1
diff --git a/e3sm_diags/plot/cartopy/zonal_mean_2d_plot.py b/e3sm_diags/plot/cartopy/zonal_mean_2d_plot.py
@@ -0,0 +1,187 @@
+from typing import List, Optional, Tuple
+
+import matplotlib
+import numpy as np
+import xarray as xr
+import xcdat as xc
+
+from e3sm_diags.driver.utils.type_annotations import MetricsDict
+from e3sm_diags.logger import custom_logger
+from e3sm_diags.parameter.core_parameter import CoreParameter
+from e3sm_diags.parameter.zonal_mean_2d_parameter import DEFAULT_PLEVS
+from e3sm_diags.plot.utils import (
+    DEFAULT_PANEL_CFG,
+    _add_colorbar,
+    _add_contour_plot,
+    _add_min_mean_max_text,
+    _add_rmse_corr_text,
+    _configure_titles,
+    _configure_x_and_y_axes,
+    _get_c_levels_and_norm,
+    _save_plot,
+)
+
+matplotlib.use("Agg")
+import matplotlib.pyplot as plt  # isort:skip  # noqa: E402
+
+logger = custom_logger(__name__)
+
+
+# Configs for x axis ticks and x axis limits.
+X_TICKS = np.array([-90, -60, -30, 0, 30, 60, 90])
+X_LIM = -90, 90
+
+
+def plot(
+    parameter: CoreParameter,
+    da_test: xr.DataArray,
+    da_ref: xr.DataArray,
+    da_diff: xr.DataArray,
+    metrics_dict: MetricsDict,
+):
+    """Plot the variable's metrics generated by the zonal_mean_2d set.
+
+    Parameters
+    ----------
+    parameter : CoreParameter
+        The CoreParameter object containing plot configurations.
+    da_test : xr.DataArray
+        The test data.
+    da_ref : xr.DataArray
+        The reference data.
+    da_diff : xr.DataArray
+        The difference between `da_test` and `da_ref` (both are regridded to
+        the lower resolution of the two beforehand).
+    metrics_dict : Metrics
+        The metrics.
+    """
+    fig = plt.figure(figsize=parameter.figsize, dpi=parameter.dpi)
+    fig.suptitle(parameter.main_title, x=0.5, y=0.96, fontsize=18)
+
+    # The variable units.
+    units = metrics_dict["units"]
+
+    # Add the first subplot for test data.
+    min1 = metrics_dict["test"]["min"]  # type: ignore
+    mean1 = metrics_dict["test"]["mean"]  # type: ignore
+    max1 = metrics_dict["test"]["max"]  # type: ignore
+
+    _add_colormap(
+        0,
+        da_test,
+        fig,
+        parameter,
+        parameter.test_colormap,
+        parameter.contour_levels,
+        title=(parameter.test_name_yrs, parameter.test_title, units),  # type: ignore
+        metrics=(max1, mean1, min1),  # type: ignore
+    )
+
+    # Add the second and third subplots for ref data and the differences,
+    # respectively.
+    min2 = metrics_dict["ref"]["min"]  # type: ignore
+    mean2 = metrics_dict["ref"]["mean"]  # type: ignore
+    max2 = metrics_dict["ref"]["max"]  # type: ignore
+
+    _add_colormap(
+        1,
+        da_ref,
+        fig,
+        parameter,
+        parameter.reference_colormap,
+        parameter.contour_levels,
+        title=(parameter.ref_name_yrs, parameter.reference_title, units),  # type: ignore
+        metrics=(max2, mean2, min2),  # type: ignore
+    )
+
+    min3 = metrics_dict["diff"]["min"]  # type: ignore
+    mean3 = metrics_dict["diff"]["mean"]  # type: ignore
+    max3 = metrics_dict["diff"]["max"]  # type: ignore
+    r = metrics_dict["misc"]["rmse"]  # type: ignore
+    c = metrics_dict["misc"]["corr"]  # type: ignore
+
+    _add_colormap(
+        2,
+        da_diff,
+        fig,
+        parameter,
+        parameter.diff_colormap,
+        parameter.diff_levels,
+        title=(None, parameter.diff_title, da_diff.attrs["units"]),  #
+        metrics=(max3, mean3, min3, r, c),  # type: ignore
+    )
+
+    _save_plot(fig, parameter)
+
+    plt.close()
+
+
+def _add_colormap(
+    subplot_num: int,
+    var: xr.DataArray,
+    fig: plt.Figure,
+    parameter: CoreParameter,
+    color_map: str,
+    contour_levels: List[float],
+    title: Tuple[Optional[str], str, str],
+    metrics: Tuple[float, ...],
+):
+    lat = xc.get_dim_coords(var, axis="Y")
+    plev = xc.get_dim_coords(var, axis="Z")
+    var = var.squeeze()
+
+    # Configure contour levels
+    # --------------------------------------------------------------------------
+    c_levels, norm = _get_c_levels_and_norm(contour_levels)
+
+    # Add the contour plot
+    # --------------------------------------------------------------------------
+    ax = fig.add_axes(DEFAULT_PANEL_CFG[subplot_num], projection=None)
+
+    contour_plot = _add_contour_plot(
+        ax, parameter, var, lat, plev, color_map, None, norm, c_levels
+    )
+
+    # Configure the aspect ratio and plot titles.
+    # --------------------------------------------------------------------------
+    ax.set_aspect("auto")
+    _configure_titles(ax, title)
+
+    # Configure x and y axis.
+    # --------------------------------------------------------------------------
+    _configure_x_and_y_axes(ax, X_TICKS, None, None, parameter.current_set)
+    ax.set_xlim(X_LIM)
+
+    if parameter.plot_log_plevs:
+        ax.set_yscale("log")
+
+    if parameter.plot_plevs:
+        plev_ticks = parameter.plevs
+        plt.yticks(plev_ticks, plev_ticks)
+
+    # For default plevs, specify the pressure axis and show the 50 mb tick
+    # at the top.
+    if (
+        not parameter.plot_log_plevs
+        and not parameter.plot_plevs
+        and parameter.plevs == DEFAULT_PLEVS
+    ):
+        plev_ticks = parameter.plevs
+        new_ticks = [plev_ticks[0]] + plev_ticks[1::2]
+        new_ticks = [int(x) for x in new_ticks]
+        plt.yticks(new_ticks, new_ticks)
+
+    plt.ylabel("pressure (mb)")
+    ax.invert_yaxis()
+
+    # Add and configure the color bar.
+    # --------------------------------------------------------------------------
+    _add_colorbar(fig, subplot_num, DEFAULT_PANEL_CFG, contour_plot, c_levels)
+
+    # Add metrics text.
+    # --------------------------------------------------------------------------
+    # Min, Mean, Max
+    _add_min_mean_max_text(fig, subplot_num, DEFAULT_PANEL_CFG, metrics)
+
+    if len(metrics) == 5:
+        _add_rmse_corr_text(fig, subplot_num, DEFAULT_PANEL_CFG, metrics)
diff --git a/e3sm_diags/plot/cartopy/zonal_mean_2d_stratosphere_plot.py b/e3sm_diags/plot/cartopy/zonal_mean_2d_stratosphere_plot.py
@@ -0,0 +1,15 @@
+import xarray as xr
+
+from e3sm_diags.driver.utils.type_annotations import MetricsDict
+from e3sm_diags.parameter.core_parameter import CoreParameter
+from e3sm_diags.plot.cartopy.zonal_mean_2d_plot import plot as base_plot
+
+
+def plot(
+    parameter: CoreParameter,
+    da_test: xr.DataArray,
+    da_ref: xr.DataArray,
+    da_diff: xr.DataArray,
+    metrics_dict: MetricsDict,
+):
+    return base_plot(parameter, da_test, da_ref, da_diff, metrics_dict)