From b6997b64d1dc78af6f983307aada8981ef3d1bef Mon Sep 17 00:00:00 2001
From: jdebacker <jason.debacker@gmail.com>
Date: Sun, 1 Sep 2024 08:27:42 -0400
Subject: [PATCH] close figure plots after tests

---
 ogcore/parameter_plots.py     |  4 ++--
 tests/test_parameter_plots.py | 17 +++++++++++++++++
 2 files changed, 19 insertions(+), 2 deletions(-)

diff --git a/ogcore/parameter_plots.py b/ogcore/parameter_plots.py
index e6a6b82e3..96648f36b 100644
--- a/ogcore/parameter_plots.py
+++ b/ogcore/parameter_plots.py
@@ -34,7 +34,7 @@ def plot_imm_rates(
 
     """
     # create line styles to cycle through
-    fig, ax = plt.subplots(layout="constrained")
+    fig, ax = plt.subplots()
     for y in years_to_plot:
         i = start_year - y
         plt.plot(imm_rates[i, :], c="blue", label="Year " + str(y))
@@ -49,7 +49,7 @@ def plot_imm_rates(
         "Source: " + source,
         fontsize=9,
     )
-    # plt.tight_layout(rect=(0, 0.035, 1, 1))
+    plt.tight_layout(rect=(0, 0.035, 1, 1))
     if include_title:
         plt.title("Immigration Rates")
     # Save or return figure
diff --git a/tests/test_parameter_plots.py b/tests/test_parameter_plots.py
index c5b1b51cf..12c9eb694 100644
--- a/tests/test_parameter_plots.py
+++ b/tests/test_parameter_plots.py
@@ -6,6 +6,7 @@
 import os
 import sys
 import numpy as np
+import matplotlib.pyplot as plt
 import scipy.interpolate as si
 import matplotlib.image as mpimg
 from ogcore import utils, parameter_plots, Specifications
@@ -74,6 +75,7 @@ def test_plot_imm_rates_save_fig(tmpdir):
 def test_plot_mort_rates():
     fig = parameter_plots.plot_mort_rates([base_params], include_title=True)
     assert fig
+    plt.close()
 
 
 def test_plot_surv_rates():
@@ -81,6 +83,7 @@ def test_plot_surv_rates():
         [base_params], survival_rates=True, include_title=True
     )
     assert fig
+    plt.close()
 
 
 def test_plot_mort_rates_save_fig(tmpdir):
@@ -104,6 +107,7 @@ def test_plot_pop_growth():
         base_params, start_year=int(base_params.start_year), include_title=True
     )
     assert fig
+    plt.close()
 
 
 def test_plot_pop_growth_rates_save_fig(tmpdir):
@@ -119,6 +123,7 @@ def test_plot_ability_profiles():
     p = Specifications()
     fig = parameter_plots.plot_ability_profiles(p, p2=p, include_title=True)
     assert fig
+    plt.close()
 
 
 def test_plot_log_ability_profiles():
@@ -127,6 +132,7 @@ def test_plot_log_ability_profiles():
         p, p2=p, log_scale=True, include_title=True
     )
     assert fig
+    plt.close()
 
 
 def test_plot_ability_profiles_save_fig(tmpdir):
@@ -144,6 +150,7 @@ def test_plot_elliptical_u():
     )
     assert fig1
     assert fig2
+    plt.close()
 
 
 def test_plot_elliptical_u_save_fig(tmpdir):
@@ -157,6 +164,7 @@ def test_plot_chi_n():
     p = Specifications()
     fig = parameter_plots.plot_chi_n([p], include_title=True)
     assert fig
+    plt.close()
 
 
 def test_plot_chi_n_save_fig(tmpdir):
@@ -177,6 +185,7 @@ def test_plot_population(years_to_plot):
         base_params, years_to_plot=years_to_plot, include_title=True
     )
     assert fig
+    plt.close()
 
 
 def test_plot_population_save_fig(tmpdir):
@@ -215,6 +224,7 @@ def test_plot_fert_rates():
     fert_rates = np.random.uniform(size=totpers).reshape((1, totpers))
     fig = parameter_plots.plot_fert_rates([fert_rates], include_title=True)
     assert fig
+    plt.close()
 
 
 def test_plot_fert_rates_save_fig(tmpdir):
@@ -258,6 +268,7 @@ def test_plot_g_n():
     p = Specifications()
     fig = parameter_plots.plot_g_n([p], include_title=True)
     assert fig
+    plt.close()
 
 
 def test_plot_g_n_savefig(tmpdir):
@@ -276,6 +287,7 @@ def test_plot_mort_rates_data():
         path=None,
     )
     assert fig
+    plt.close()
 
 
 def test_plot_mort_rates_data_save_fig(tmpdir):
@@ -300,6 +312,7 @@ def test_plot_omega_fixed():
         age_per_EpS, omega_SS_orig, omega_SSfx, E, S
     )
     assert fig
+    plt.close()
 
 
 def test_plot_omega_fixed_save_fig(tmpdir):
@@ -326,6 +339,7 @@ def test_plot_imm_fixed():
         age_per_EpS, imm_rates_orig, imm_rates_adj, E, S
     )
     assert fig
+    plt.close()
 
 
 def test_plot_imm_fixed_save_fig(tmpdir):
@@ -360,6 +374,7 @@ def test_plot_population_path():
         S,
     )
     assert fig
+    plt.close()
 
 
 def test_plot_population_path_save_fig(tmpdir):
@@ -398,6 +413,7 @@ def test_plot_income_data():
     fig = parameter_plots.plot_income_data(ages, abil_midp, abil_pcts, emat)
 
     assert fig
+    plt.close()
 
 
 def test_plot_income_data_save_fig(tmpdir):
@@ -481,6 +497,7 @@ def test_plot_2D_taxfunc(
         )
 
         assert fig
+        plt.close()
     else:
         assert True