Fix layout for sub-plot saved in pdf format (#984)

* fix padding for zonal_mean_2d and meridional_mean_2d

* fix qbo and tropical subseasonal plots

* fix parser/plot for enso_diags

* update for aerosol_aeronet

* fix pre-commit errors

* refactor for review comments; more refining for enso

* Add a temporary workaround to rename nv bounds dim

* Fix pre-commit issue

* Refactor plotting utilities for maintainability

---------

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

Author: chengzhuzhang

e3sm_diags/driver/arm_diags_driver.py

@@ -278,6 +278,11 @@ def _run_diag_annual_cycle(parameter: ARMDiagsParameter) -> ARMDiagsParameter:
                     vars_funcs = _get_vars_funcs_for_derived_var(ds_ref, var)
                     target_var = list(vars_funcs.keys())[0][0]
 
+                    # NOTE: The bounds dimension can be "nv", which is not
+                    # currently recognized as a valid bounds dimension
+                    # by xcdat. We rename it to "bnds" to make it compatible.
+                    ds_ref = _rename_bounds_dim(ds_ref)
+
                     ds_ref_climo = ds_ref.temporal.climatology(target_var, "month")
                     da_ref = vars_funcs[(target_var,)](ds_ref_climo[target_var]).rename(
                         var
@@ -530,3 +535,33 @@ def _save_metrics_to_json(parameter: ARMDiagsParameter, metrics_dict: Dict[str,
         json.dump(metrics_dict, outfile)
 
     logger.info(f"Metrics saved in: {abs_path}")
+
+
+def _rename_bounds_dim(ds: xr.Dataset) -> xr.Dataset:
+    """
+    Renames the bounds dimension "nv" to "bnds" in the given xarray.Dataset for
+    xCDAT compatibility.
+
+    This is a temporary workaround to ensure compatibility with xCDAT's bounds
+    handling. The bounds dimension "nv" is commonly used in datasets to
+    represent the number of vertices in a polygon, but xCDAT expects the
+    bounds dimension to be in `xcdat.bounds.VALID_BOUNDS_DIMS`. This function
+    renames "nv" to "bnds" to align with xCDAT's expectations.
+
+    Parameters
+    ----------
+    ds : xr.Dataset
+        The input xarray.Dataset which may contain a bounds dimension named "nv".
+
+    Returns
+    -------
+    xr.Dataset
+        A new xarray.Dataset with the "nv" dimension renamed to "bnds" if it
+        existed; otherwise, the original dataset copy.
+    """
+    ds_new = ds.copy()
+
+    if "nv" in ds_new.dims:
+        ds_new = ds_new.rename({"nv": "bnds"})
+
+    return ds_new

e3sm_diags/parser/enso_diags_parser.py

@@ -50,3 +50,10 @@ def add_arguments(self):
             help="End year for the timeseries files.",
             required=False,
         )
+
+        self.parser.add_argument(
+            "--plot_type",
+            dest="plot_type",
+            help="Type of plot to generate: 'map' or 'scatter'.",
+            required=False,
+        )

e3sm_diags/plot/aerosol_aeronet_plot.py

@@ -6,7 +6,7 @@
 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_plot
+from e3sm_diags.plot.utils import _save_main_plot, _save_single_subplot
 
 matplotlib.use("Agg")
 import matplotlib.pyplot as plt  # isort:skip  # noqa: E402
@@ -22,7 +22,28 @@
 ]
 # Border padding relative to subplot axes for saving individual panels
 # (left, bottom, right, top) in page coordinates.
-BORDER_PADDING = (-0.06, -0.03, 0.13, 0.03)
+BORDER_PADDING_COLORMAP = (-0.06, 0.25, 0.13, 0.25)
+BORDER_PADDING_SCATTER = (-0.08, -0.04, 0.15, 0.04)
+
+
+def _save_plot_aerosol_aeronet(fig, parameter):
+    """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
+    _save_main_plot(parameter)
+
+    # 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)
+
+        # Save scatter panel (panel 1) with its specific border padding
+        _save_single_subplot(fig, parameter, 1, PANEL_CFG[1], BORDER_PADDING_SCATTER)
 
 
 def plot(
@@ -108,4 +129,4 @@ def plot(
 
     plt.loglog(ref_site_arr, test_site_arr, "kx", markersize=3.0, mfc="none")
 
-    _save_plot(fig, parameter, PANEL_CFG, BORDER_PADDING)
+    _save_plot_aerosol_aeronet(fig, parameter)

e3sm_diags/plot/enso_diags_plot.py

@@ -23,7 +23,9 @@
     _get_x_ticks,
     _get_y_ticks,
     _make_lon_cyclic,
+    _save_main_plot,
     _save_plot,
+    _save_single_subplot,
 )
 
 matplotlib.use("Agg")
@@ -41,6 +43,19 @@
 # Use 179.99 as central longitude due to https://github.com/SciTools/cartopy/issues/946
 PROJECTION = ccrs.PlateCarree(central_longitude=179.99)
 
+# 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.17, 0.022)
+
+
+def _save_plot_scatter(fig: plt.Figure, parameter: EnsoDiagsParameter):
+    """Save the scatter plot using the shared _save_single_subplot function."""
+    _save_main_plot(parameter)
+
+    # 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(
     parameter: EnsoDiagsParameter, x: MetricsDictScatter, y: MetricsDictScatter
@@ -138,7 +153,7 @@ def plot_scatter(
     plt.ylabel("{} anomaly ({})".format(y["var"], y["units"]))
     plt.legend()
 
-    _save_plot(fig, parameter)
+    _save_plot_scatter(fig, parameter)
 
     plt.close()
 
@@ -192,7 +207,7 @@ def plot_map(
     )
     _plot_diff_rmse_and_corr(fig, metrics_dict["diff"])  # type: ignore
 
-    _save_plot(fig, parameter)
+    _save_plot(fig, parameter, DEFAULT_PANEL_CFG, ENSO_BORDER_PADDING_MAP)
 
     plt.close()
 
@@ -301,17 +316,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/mp_partition_plot.py

@@ -1,9 +1,7 @@
-import os
-
 import matplotlib
 
-from e3sm_diags.driver.utils.io import _get_output_dir
 from e3sm_diags.logger import _setup_child_logger
+from e3sm_diags.plot.utils import _save_main_plot
 
 matplotlib.use("Agg")
 import matplotlib.pyplot as plt  # isort:skip  # noqa: E402
@@ -89,11 +87,4 @@ def plot(metrics_dict, parameter):
     ax.legend(loc="upper left")
     ax.set_title("Mixed-phase Partition LCF [30S - 70S]")
 
-    for f in parameter.output_format:
-        f = f.lower().split(".")[-1]
-        fnm = os.path.join(
-            _get_output_dir(parameter),
-            f"{parameter.output_file}" + "." + f,
-        )
-        plt.savefig(fnm)
-        logger.info(f"Plot saved in: {fnm}")
+    _save_main_plot(parameter)

e3sm_diags/plot/qbo_plot.py

@@ -1,6 +1,6 @@
 from __future__ import annotations
 
-from typing import List, Literal, TypedDict
+from typing import List, Literal, Tuple, TypedDict
 
 import matplotlib
 import numpy as np
@@ -18,11 +18,15 @@
 PANEL_CFG = [
     (0.075, 0.75, 0.6, 0.175),
     (0.075, 0.525, 0.6, 0.175),
-    (0.725, 0.525, 0.2, 0.4),
+    (0.735, 0.525, 0.2, 0.4),
     (0.075, 0.285, 0.85, 0.175),
     (0.075, 0.04, 0.85, 0.175),
 ]
 
+# Border padding relative to subplot axes for saving individual panels
+# (left, bottom, right, top) in page coordinates
+QBO_BORDER_PADDING: Tuple[float, float, float, float] = (-0.07, -0.03, 0.009, 0.03)
+
 LABEL_SIZE = 14
 CMAP = plt.cm.RdBu_r
 
@@ -203,7 +207,7 @@ def plot(parameter: QboParameter, test_dict, ref_dict):
     fig.suptitle(parameter.main_title, x=0.5, y=0.97, fontsize=15)
 
     # Save figure
-    _save_plot(fig, parameter, PANEL_CFG)
+    _save_plot(fig, parameter, PANEL_CFG, QBO_BORDER_PADDING)
 
     plt.close()
 

e3sm_diags/plot/tc_analysis_plot.py

@@ -1,15 +1,12 @@
-import os
-
 import cartopy.crs as ccrs
 import cartopy.feature as cfeature
 import matplotlib
 import numpy as np
 import xcdat as xc
 from cartopy.mpl.ticker import LatitudeFormatter, LongitudeFormatter
 
-from e3sm_diags.driver.utils.io import _get_output_dir
 from e3sm_diags.logger import _setup_child_logger
-from e3sm_diags.plot.utils import MAIN_TITLE_FONTSIZE
+from e3sm_diags.plot.utils import MAIN_TITLE_FONTSIZE, _save_main_plot
 
 matplotlib.use("agg")
 import matplotlib.pyplot as plt  # isort:skip  # noqa: E402
@@ -103,16 +100,8 @@ def plot(test, ref, parameter, basin_dict):
         y=0.99,
     )
 
-    output_file_name = "tc-intensity"
-    for f in parameter.output_format:
-        f = f.lower().split(".")[-1]
-        fnm = os.path.join(
-            _get_output_dir(parameter),
-            output_file_name + "." + f,
-        )
-        plt.savefig(fnm)
-        logger.info(f"Plot saved in: {fnm}")
-        plt.close()
+    parameter.output_file = "tc-intensity"
+    _save_main_plot(parameter)
 
     # TC frequency of each basins
     fig = plt.figure(figsize=(12, 7))
@@ -150,16 +139,8 @@ def plot(test, ref, parameter, basin_dict):
     ax.set_ylabel("Fraction")
     ax.set_title("Relative frequency of TCs for each ocean basins")
 
-    output_file_name = "tc-frequency"
-    for f in parameter.output_format:
-        f = f.lower().split(".")[-1]
-        fnm = os.path.join(
-            _get_output_dir(parameter),
-            output_file_name + "." + f,
-        )
-        plt.savefig(fnm)
-        logger.info(f"Plot saved in: {fnm}")
-        plt.close()
+    parameter.output_file = "tc-frequency"
+    _save_main_plot(parameter)
 
     fig1 = plt.figure(figsize=(12, 6))
     ax = fig1.add_subplot(111)
@@ -190,16 +171,9 @@ def plot(test, ref, parameter, basin_dict):
     ax.set_title(
         "Distribution of accumulated cyclone energy (ACE) among various ocean basins"
     )
-    output_file_name = "ace-distribution"
-    for f in parameter.output_format:
-        f = f.lower().split(".")[-1]
-        fnm = os.path.join(
-            _get_output_dir(parameter),
-            output_file_name + "." + f,
-        )
-        plt.savefig(fnm)
-        logger.info(f"Plot saved in: {fnm}")
-        plt.close()
+
+    parameter.output_file = "ace-distribution"
+    _save_main_plot(parameter)
 
     fig, axes = plt.subplots(2, 3, figsize=(12, 6), sharex=True, sharey=True)
     fig.subplots_adjust(hspace=0.4, wspace=0.15)
@@ -231,16 +205,8 @@ def plot(test, ref, parameter, basin_dict):
         y=0.99,
     )
 
-    output_file_name = "tc-frequency-annual-cycle"
-    for f in parameter.output_format:
-        f = f.lower().split(".")[-1]
-        fnm = os.path.join(
-            _get_output_dir(parameter),
-            output_file_name + "." + f,
-        )
-        plt.savefig(fnm)
-        logger.info(f"Plot saved in: {fnm}")
-        plt.close()
+    parameter.output_file = "tc-frequency-annual-cycle"
+    _save_main_plot(parameter)
 
     ##########################################################
     # Plot TC tracks density
@@ -286,17 +252,9 @@ def plot_map(test_data, ref_data, region, parameter):
 
     # Figure title
     fig.suptitle(PLOT_INFO[region]["title"], x=0.5, y=0.9, fontsize=14)
-    output_file_name = "{}-density-map".format(region)
+    parameter.output_file = "{}-density-map".format(region)
 
-    for f in parameter.output_format:
-        f = f.lower().split(".")[-1]
-        fnm = os.path.join(
-            _get_output_dir(parameter),
-            output_file_name + "." + f,
-        )
-        plt.savefig(fnm)
-        logger.info(f"Plot saved in: {fnm}")
-        plt.close()
+    _save_main_plot(parameter)
 
 
 def plot_panel(n, fig, proj, var, var_num_years, region, title):