From ae942ea7879faa81bec24c795057b234d8de1c6e Mon Sep 17 00:00:00 2001 From: Po-han Li <34739111+d31003@users.noreply.github.com> Date: Thu, 25 Apr 2024 18:05:43 -0500 Subject: [PATCH] Refactor save_fig function to add support for tight layout in general_utils.py (#27) --- swarm_visualizer/utility/general_utils.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/swarm_visualizer/utility/general_utils.py b/swarm_visualizer/utility/general_utils.py index b101d68..30c5913 100644 --- a/swarm_visualizer/utility/general_utils.py +++ b/swarm_visualizer/utility/general_utils.py @@ -54,17 +54,18 @@ def set_plot_properties( sns.set_style(style="darkgrid") -def save_fig(fig, save_loc: str = None, dpi: int = 600) -> None: +def save_fig(fig, save_loc: str = None, dpi: int = 600, tight_layout: bool=True) -> None: """Save figure. :param fig: figure :param save_loc: location to save the figure :param dpi: dpi + :param tight_layout: tight layout or not :return: None. """ + plt.tight_layout() if tight_layout else None # If save location doesn't exist, create it - if not os.path.exists(os.path.dirname(save_loc)): - os.makedirs(os.path.dirname(save_loc)) + os.makedirs(os.path.dirname(save_loc), exist_ok=True) # Save the figure fig.savefig(save_loc, dpi=dpi)