allow for 2 sets of earnings profiles

jdebacker · Mar 19, 2024 · e9f0d7c · e9f0d7c
1 parent 3c29174
commit e9f0d7c
Show file tree

Hide file tree

Showing 2 changed files with 19 additions and 3 deletions.
diff --git a/ogcore/parameter_plots.py b/ogcore/parameter_plots.py
@@ -193,7 +193,7 @@ def plot_population(p, years_to_plot=["SS"], include_title=False, path=None):
         plt.savefig(fig_path, dpi=300)
 
 
-def plot_ability_profiles(p, t=None, log_scale=False, include_title=False, path=None):
+def plot_ability_profiles(p, p2=None, t=None, log_scale=False, include_title=False, path=None):
     """
     Create a plot of earnings ability profiles.
 
@@ -219,6 +219,22 @@ def plot_ability_profiles(p, t=None, log_scale=False, include_title=False, path=
             plt.plot(age_vec, np.log(p.e[t, :, j]), label=GROUP_LABELS[p.J][j])
         else:
             plt.plot(age_vec, p.e[t, :, j], label=GROUP_LABELS[p.J][j])
+    if p2 is not None:
+        for j in range(p.J):
+            if log_scale:
+                plt.plot(
+                    age_vec,
+                    np.log(p2.e[t, :, j]),
+                    linestyle="--",
+                    label=GROUP_LABELS[p.J][j],
+                )
+            else:
+                plt.plot(
+                    age_vec,
+                    p2.e[t, :, j],
+                    linestyle="--",
+                    label=GROUP_LABELS[p.J][j],
+                )
     plt.xlabel(r"Age")
     if log_scale:
         plt.ylabel(r"ln(Earnings ability)")

diff --git a/tests/test_parameter_plots.py b/tests/test_parameter_plots.py
@@ -106,15 +106,15 @@ def test_plot_pop_growth_rates_save_fig(tmpdir):
 def test_plot_ability_profiles():
     p = Specifications()
     fig = parameter_plots.plot_ability_profiles(
-        p, include_title=True
+        p, p2=p, include_title=True
     )
     assert fig
 
 
 def test_plot_log_ability_profiles():
     p = Specifications()
     fig = parameter_plots.plot_ability_profiles(
-        p, log_scale=True, include_title=True
+        p, p2=p, log_scale=True, include_title=True
     )
     assert fig