From e9f0d7c8d786163026d3d9174043ec9b1e0d6157 Mon Sep 17 00:00:00 2001 From: jdebacker Date: Mon, 18 Mar 2024 21:43:30 -0400 Subject: [PATCH] allow for 2 sets of earnings profiles --- ogcore/parameter_plots.py | 18 +++++++++++++++++- tests/test_parameter_plots.py | 4 ++-- 2 files changed, 19 insertions(+), 3 deletions(-) diff --git a/ogcore/parameter_plots.py b/ogcore/parameter_plots.py index 45e3e0d39..b4f2258bf 100644 --- 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 index e5006032d..0e80fc392 100644 --- a/tests/test_parameter_plots.py +++ b/tests/test_parameter_plots.py @@ -106,7 +106,7 @@ 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 @@ -114,7 +114,7 @@ def test_plot_ability_profiles(): 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