Refactor plotting utilities for maintainability

Author: tomvothecoder

e3sm_diags/plot/aerosol_aeronet_plot.py

@@ -1,15 +1,12 @@
-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
+from e3sm_diags.plot.utils import _save_main_plot, _save_single_subplot
 
 matplotlib.use("Agg")
 import matplotlib.pyplot as plt  # isort:skip  # noqa: E402
@@ -37,14 +34,7 @@ def _save_plot_aerosol_aeronet(fig, parameter):
     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]
-        fnm = os.path.join(
-            _get_output_dir(parameter),
-            parameter.output_file + "." + f,
-        )
-        plt.savefig(fnm)
-        logger.info(f"Plot saved in: {fnm}")
+    _save_main_plot(parameter)
 
     # Save subplots with different border padding by calling general function
     # for each panel individually

e3sm_diags/plot/enso_diags_plot.py

@@ -1,6 +1,5 @@
 from __future__ import annotations
 
-import os
 from typing import TYPE_CHECKING, List, Tuple
 
 import cartopy.crs as ccrs
@@ -11,7 +10,6 @@
 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 (
@@ -25,6 +23,7 @@
     _get_x_ticks,
     _get_y_ticks,
     _make_lon_cyclic,
+    _save_main_plot,
     _save_plot,
     _save_single_subplot,
 )
@@ -51,15 +50,7 @@
 
 def _save_plot_scatter(fig: plt.Figure, parameter: EnsoDiagsParameter):
     """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(
-            _get_output_dir(parameter),
-            parameter.output_file + "." + f,
-        )
-        plt.savefig(fnm)
-        logger.info(f"Plot saved in: {fnm}")
+    _save_main_plot(parameter)
 
     # Save the single subplot using shared helper (panel_config=None for full figure)
     if parameter.output_format_subplot:

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/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):