refactor for review comments; more refining for enso

Author: chengzhuzhang

e3sm_diags/plot/aerosol_aeronet_plot.py

@@ -1,11 +1,15 @@
+import os
+
 import matplotlib
 import numpy as np
 import xarray as xr
 
+from e3sm_diags.driver.utils.io import _get_output_dir
 from e3sm_diags.driver.utils.type_annotations import MetricsDict
 from e3sm_diags.logger import _setup_child_logger
 from e3sm_diags.parameter.core_parameter import CoreParameter
 from e3sm_diags.plot.lat_lon_plot import _add_colormap
+from e3sm_diags.plot.utils import _save_single_subplot
 
 matplotlib.use("Agg")
 import matplotlib.pyplot as plt  # isort:skip  # noqa: E402
@@ -26,13 +30,12 @@
 
 
 def _save_plot_aerosol_aeronet(fig, parameter):
-    """Save aerosol_aeronet plots with different border padding for each panel."""
-    import os
-
-    from matplotlib.transforms import Bbox
-
-    from e3sm_diags.driver.utils.io import _get_output_dir
+    """Save aerosol_aeronet plots using the _save_single_subplot helper function.
 
+    This function handles the special case where different border padding is needed
+    for each panel by calling _save_single_subplot twice with panel-specific
+    configurations (BORDER_PADDING_COLORMAP for panel 0, BORDER_PADDING_SCATTER for panel 1).
+    """
     # Save the main plot
     for f in parameter.output_format:
         f = f.lower().split(".")[-1]
@@ -43,28 +46,14 @@ def _save_plot_aerosol_aeronet(fig, parameter):
         plt.savefig(fnm)
         logger.info(f"Plot saved in: {fnm}")
 
-    # Save individual subplots with different border padding
-    border_paddings = [BORDER_PADDING_COLORMAP, BORDER_PADDING_SCATTER]
+    # Save subplots with different border padding by calling general function
+    # for each panel individually
+    if parameter.output_format_subplot:
+        # Save colormap panel (panel 0) with its specific border padding
+        _save_single_subplot(fig, parameter, 0, PANEL_CFG[0], BORDER_PADDING_COLORMAP)
 
-    for f in parameter.output_format_subplot:
-        fnm = os.path.join(
-            _get_output_dir(parameter),
-            parameter.output_file,
-        )
-        page = fig.get_size_inches()
-
-        for idx, (panel, border_padding) in enumerate(zip(PANEL_CFG, border_paddings)):
-            # Extent of subplot
-            subpage = np.array(panel).reshape(2, 2)
-            subpage[1, :] = subpage[0, :] + subpage[1, :]
-            subpage = subpage + np.array(border_padding).reshape(2, 2)
-            subpage_list = list(((subpage) * page).flatten())
-            extent = Bbox.from_extents(*subpage_list)
-
-            # Save subplot
-            fname = fnm + ".%i." % idx + f
-            plt.savefig(fname, bbox_inches=extent)
-            logger.info(f"Sub-plot saved in: {fname}")
+        # Save scatter panel (panel 1) with its specific border padding
+        _save_single_subplot(fig, parameter, 1, PANEL_CFG[1], BORDER_PADDING_SCATTER)
 
 
 def plot(

e3sm_diags/plot/enso_diags_plot.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import os
 from typing import TYPE_CHECKING, List, Tuple
 
 import cartopy.crs as ccrs
@@ -10,6 +11,7 @@
 from numpy.polynomial.polynomial import polyfit
 
 from e3sm_diags.derivations.default_regions_xr import REGION_SPECS
+from e3sm_diags.driver.utils.io import _get_output_dir
 from e3sm_diags.logger import _setup_child_logger
 from e3sm_diags.parameter.enso_diags_parameter import EnsoDiagsParameter
 from e3sm_diags.plot.utils import (
@@ -24,6 +26,7 @@
     _get_y_ticks,
     _make_lon_cyclic,
     _save_plot,
+    _save_single_subplot,
 )
 
 matplotlib.use("Agg")
@@ -43,15 +46,12 @@
 
 # Border padding relative to subplot axes for saving individual panels
 # (left, bottom, right, top) in page coordinates
-ENSO_BORDER_PADDING_MAP = (-0.07, -0.025, 0.2, 0.035)
+ENSO_BORDER_PADDING_MAP = (-0.07, -0.025, 0.17, 0.022)
 
 
 def _save_plot_scatter(fig: plt.Figure, parameter: EnsoDiagsParameter):
-    """Save the scatter plot using a simplified approach for single panel plots."""
-    import os
-
-    from e3sm_diags.driver.utils.io import _get_output_dir
-
+    """Save the scatter plot using the shared _save_single_subplot function."""
+    # Save the main plot
     for f in parameter.output_format:
         f = f.lower().split(".")[-1]
         fnm = os.path.join(
@@ -61,14 +61,9 @@ def _save_plot_scatter(fig: plt.Figure, parameter: EnsoDiagsParameter):
         plt.savefig(fnm)
         logger.info(f"Plot saved in: {fnm}")
 
-    # Save individual subplots (single panel for scatter)
-    for f in parameter.output_format_subplot:
-        fnm = os.path.join(
-            _get_output_dir(parameter),
-            parameter.output_file + ".0." + f,
-        )
-        plt.savefig(fnm)
-        logger.info(f"Sub-plot saved in: {fnm}")
+    # Save the single subplot using shared helper (panel_config=None for full figure)
+    if parameter.output_format_subplot:
+        _save_single_subplot(fig, parameter, 0, None, None)
 
 
 def plot_scatter(
@@ -330,17 +325,17 @@ def _add_colormap(
     top_text = "Max\nMin\nMean\nSTD"
     fig.text(
         DEFAULT_PANEL_CFG[subplot_num][0] + 0.6635,
-        DEFAULT_PANEL_CFG[subplot_num][1] + 0.2107,
+        DEFAULT_PANEL_CFG[subplot_num][1] + 0.2,
         top_text,
         ha="left",
-        fontdict={"fontsize": SECONDARY_TITLE_FONTSIZE},
+        fontdict={"fontsize": 9},
     )
     fig.text(
         DEFAULT_PANEL_CFG[subplot_num][0] + 0.7635,
-        DEFAULT_PANEL_CFG[subplot_num][1] + 0.2107,
+        DEFAULT_PANEL_CFG[subplot_num][1] + 0.2,
         "%.2f\n%.2f\n%.2f\n%.2f" % metrics_values,  # type: ignore
         ha="right",
-        fontdict={"fontsize": SECONDARY_TITLE_FONTSIZE},
+        fontdict={"fontsize": 9},
     )
 
     # Hatch text

e3sm_diags/plot/tropical_subseasonal_plot.py

@@ -11,6 +11,7 @@
 from e3sm_diags.driver.utils.io import _get_output_dir
 from e3sm_diags.logger import _setup_child_logger
 from e3sm_diags.parameter.core_parameter import CoreParameter
+from e3sm_diags.plot.utils import _save_plot
 
 matplotlib.use("Agg")
 import matplotlib.pyplot as plt  # isort:skip  # noqa: E402
@@ -227,61 +228,6 @@ def create_colormap_clevs(cmapSpec, clevs):
     return cmapSpecUse, normSpecUse
 
 
-def _save_plot(fig: plt.figure, parameter: CoreParameter):
-    """Save the plot using the figure object and parameter configs.
-
-    This function creates the output filename to save the plot. It also
-    saves each individual subplot if the reference name is an empty string ("").
-
-    Parameters
-    ----------
-    fig : plt.figure
-        The plot figure.
-    parameter : CoreParameter
-        The CoreParameter with file configurations.
-    """
-    for f in parameter.output_format:
-        f = f.lower().split(".")[-1]
-        fnm = os.path.join(
-            _get_output_dir(parameter),
-            parameter.output_file + "." + f,
-        )
-        plt.savefig(fnm)
-        logger.info(f"Plot saved in: {fnm}")
-
-    # Save individual subplots
-    if parameter.ref_name == "":
-        panels = [PANEL[0]]
-    else:
-        panels = PANEL
-
-    for f in parameter.output_format_subplot:
-        fnm = os.path.join(
-            _get_output_dir(parameter),
-            parameter.output_file,
-        )
-        page = fig.get_size_inches()
-
-        for idx, panel in enumerate(panels):
-            # Extent of subplot
-            subpage = np.array(panel).reshape(2, 2)
-            subpage[1, :] = subpage[0, :] + subpage[1, :]
-            subpage = subpage + np.array(BORDER_PADDING).reshape(2, 2)
-            subpage = list(((subpage) * page).flatten())  # type: ignore
-            extent = matplotlib.transforms.Bbox.from_extents(*subpage)
-
-            # Save subplot
-            fname = fnm + ".%i." % idx + f
-            plt.savefig(fname, bbox_inches=extent)
-
-            orig_fnm = os.path.join(
-                _get_output_dir(parameter),
-                parameter.output_file,
-            )
-            fname = orig_fnm + ".%i." % idx + f
-            logger.info(f"Sub-plot saved in: {fname}")
-
-
 def _wave_frequency_plot(  # noqa: C901
     subplot_num: int,
     var: xr.DataArray,
@@ -770,6 +716,6 @@ def plot(
         do_zoom=do_zoom,
     )
 
-    _save_plot(fig, parameter)
+    _save_plot(fig, parameter, PANEL, BORDER_PADDING)
 
     plt.close()

e3sm_diags/plot/utils.py

@@ -57,6 +57,50 @@
 ]
 
 
+def _save_single_subplot(fig, parameter, panel_idx, panel_config, border_padding):
+    """Save a single subplot with specific border padding.
+
+    Parameters
+    ----------
+    fig : plt.Figure
+        The plot figure.
+    parameter : CoreParameter
+        The CoreParameter with file configurations.
+    panel_idx : int
+        The panel index for filename generation.
+    panel_config : tuple or None
+        Panel configuration (left, bottom, width, height). If None, saves entire figure.
+    border_padding : tuple
+        Border padding (left, bottom, right, top) in page coordinates.
+    """
+    for f in parameter.output_format_subplot:
+        fnm = os.path.join(
+            _get_output_dir(parameter),
+            parameter.output_file,
+        )
+
+        if panel_config is None:
+            # Save entire figure (for full-figure plots like scatter)
+            fname = fnm + f".{panel_idx}." + f
+            plt.savefig(fname)
+            logger.info(f"Sub-plot saved in: {fname}")
+        else:
+            # Save cropped subplot with border padding
+            page = fig.get_size_inches()
+
+            # Extent of subplot
+            subpage = np.array(panel_config).reshape(2, 2)
+            subpage[1, :] = subpage[0, :] + subpage[1, :]
+            subpage = subpage + np.array(border_padding).reshape(2, 2)
+            subpage_list = list(((subpage) * page).flatten())
+            extent = Bbox.from_extents(*subpage_list)
+
+            # Save subplot
+            fname = fnm + f".{panel_idx}." + f
+            plt.savefig(fname, bbox_inches=extent)
+            logger.info(f"Sub-plot saved in: {fname}")
+
+
 def _save_plot(
     fig: plt.Figure,
     parameter: CoreParameter,