[Bug]: CDAT Migration Phase 2: enso_diags plot fixes (#841)

E3SM-Project · Oct 25, 2024 · 9779112 · 9779112
1 parent a00eca8
commit 9779112
Show file tree

Hide file tree

Showing 3 changed files with 31 additions and 6 deletions.
diff --git a/auxiliary_tools/cdat_regression_testing/663-enso-diags/regression_test_png.ipynb b/auxiliary_tools/cdat_regression_testing/663-enso-diags/regression_test_png.ipynb
@@ -161,7 +161,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 6,
    "metadata": {},
    "outputs": [
     {
@@ -192,7 +192,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 6,
+   "execution_count": 7,
    "metadata": {},
    "outputs": [
     {
@@ -224,7 +224,7 @@
        " '/global/cfs/cdirs/e3sm/www/cdat-migration-fy24/663-enso-diags/enso_diags/TAUY-response_diff/regression-coefficient-tauy-over-nino34.png']"
       ]
      },
-     "execution_count": 6,
+     "execution_count": 7,
      "metadata": {},
      "output_type": "execute_result"
     }
@@ -235,7 +235,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 7,
+   "execution_count": 8,
    "metadata": {},
    "outputs": [
     {

diff --git a/e3sm_diags/driver/enso_diags_driver.py b/e3sm_diags/driver/enso_diags_driver.py
@@ -17,7 +17,7 @@
 )
 from e3sm_diags.driver.utils.regrid import _subset_on_region, align_grids_to_lower_res
 from e3sm_diags.logger import custom_logger
-from e3sm_diags.metrics.metrics import spatial_avg, std
+from e3sm_diags.metrics.metrics import correlation, rmse, spatial_avg, std
 from e3sm_diags.plot.enso_diags_plot import plot_map, plot_scatter
 
 if TYPE_CHECKING:
@@ -44,6 +44,8 @@ class MetricsSubDict(TypedDict):
     max: float
     mean: List[float]
     std: List[float]
+    rmse: float | None
+    corr: float | None
 
 
 # A type annotation representing the metrics dictionary.
@@ -531,6 +533,8 @@ def _create_metrics_dict(
     metrics_dict["test"] = get_metrics_subdict(ds_test, var_key)
     metrics_dict["test_regrid"] = get_metrics_subdict(ds_test_regrid, var_key)
     metrics_dict["diff"] = get_metrics_subdict(ds_diff, var_key)
+    metrics_dict["diff"]["rmse"] = rmse(ds_test_regrid, ds_ref_regrid, var_key)  # type: ignore
+    metrics_dict["diff"]["corr"] = correlation(ds_test_regrid, ds_ref_regrid, var_key)  # type: ignore
 
     metrics_dict["unit"] = ds_test[var_key].units
 

diff --git a/e3sm_diags/plot/enso_diags_plot.py b/e3sm_diags/plot/enso_diags_plot.py
@@ -190,6 +190,7 @@ def plot_map(
         (None, parameter.diff_title, da_test.units),
         metrics_dict["diff"],  # type: ignore
     )
+    _plot_diff_rmse_and_corr(fig, metrics_dict["diff"])  # type: ignore
 
     _save_plot(fig, parameter)
 
@@ -296,7 +297,7 @@ def _add_colormap(
 
     # Add metrics text to the figure.
     # --------------------------------------------------------------------------
-    metrics_values = tuple(metrics.values())
+    metrics_values = (metrics["max"], metrics["min"], metrics["mean"], metrics["std"])
     top_text = "Max\nMin\nMean\nSTD"
     fig.text(
         DEFAULT_PANEL_CFG[subplot_num][0] + 0.6635,
@@ -325,6 +326,26 @@ def _add_colormap(
         )
 
 
+def _plot_diff_rmse_and_corr(fig: plt.Figure, metrics_dict: MetricsSubDict):
+    bottom_stats = (metrics_dict["rmse"], metrics_dict["corr"])
+    bottom_text = "RMSE\nCORR"
+
+    fig.text(
+        DEFAULT_PANEL_CFG[2][0] + 0.6635,
+        DEFAULT_PANEL_CFG[2][1] - 0.0205,
+        bottom_text,
+        ha="left",
+        fontdict={"fontsize": SECONDARY_TITLE_FONTSIZE},
+    )
+    fig.text(
+        DEFAULT_PANEL_CFG[2][0] + 0.7635,
+        DEFAULT_PANEL_CFG[2][1] - 0.0205,
+        "%.2f\n%.2f" % bottom_stats,  # type: ignore
+        ha="right",
+        fontdict={"fontsize": SECONDARY_TITLE_FONTSIZE},
+    )
+
+
 def _determine_tick_step(degrees_covered: float) -> int:
     """Determine the number of tick steps based on the degrees covered by the axis.