diff --git a/docs/book/content/api/parameter_plots.rst b/docs/book/content/api/parameter_plots.rst index 977b628f8..16c10f463 100644 --- a/docs/book/content/api/parameter_plots.rst +++ b/docs/book/content/api/parameter_plots.rst @@ -11,6 +11,7 @@ ogcore.parameter_plots .. automodule:: ogcore.parameter_plots :members: plot_imm_rates, plot_mort_rates, plot_pop_growth, plot_ability_profiles, plot_elliptical_u, plot_chi_n, - plot_fert_rates, plot_mort_rates_data, plot_omega_fixed, + plot_fert_rates, plot_mort_rates_data, plot_g_n, plot_omega_fixed, plot_imm_fixed, plot_population_path, gen_3Dscatters_hist, txfunc_graph, txfunc_sse_plot, plot_income_data, plot_2D_taxfunc + diff --git a/ogcore/output_plots.py b/ogcore/output_plots.py index bf7af7fc5..661044c3f 100644 --- a/ogcore/output_plots.py +++ b/ogcore/output_plots.py @@ -20,6 +20,7 @@ def plot_aggregates( reform_params=None, var_list=["Y", "C", "K", "L"], plot_type="pct_diff", + stationarized=True, num_years_to_plot=50, start_year=DEFAULT_START_YEAR, forecast_data=None, @@ -40,13 +41,15 @@ def plot_aggregates( object var_list (list): names of variable to plot plot_type (string): type of plot, can be: - 'pct_diff': plots percentage difference between baselien + 'pct_diff': plots percentage difference between baseline and reform ((reform-base)/base) 'diff': plots difference between baseline and reform (reform-base) 'levels': plot variables in model units 'forecast': plots variables in levels relative to baseline economic forecast + stationarized (bool): whether used stationarized variables (False + only affects pct_diff right now) num_years_to_plot (integer): number of years to include in plot start_year (integer): year to start plot forecast_data (array_like): baseline economic forecast series, @@ -78,11 +81,21 @@ def plot_aggregates( # Compute just percentage point changes for rates plot_var = reform_tpi[v] - base_tpi[v] else: - plot_var = (reform_tpi[v] - base_tpi[v]) / base_tpi[v] + if stationarized: + plot_var = (reform_tpi[v] - base_tpi[v]) / base_tpi[v] + else: + pct_changes = utils.pct_change_unstationarized( + base_tpi, + base_params, + reform_tpi, + reform_params, + output_vars=[v], + ) + plot_var = pct_changes[v] ylabel = r"Pct. change" plt.plot( year_vec, - plot_var[start_index : start_index + num_years_to_plot], + plot_var[start_index : start_index + num_years_to_plot] * 100, label=VAR_LABELS[v], ) elif plot_type == "diff": @@ -185,7 +198,7 @@ def plot_industry_aggregates( object var_list (list): names of variable to plot plot_type (string): type of plot, can be: - 'pct_diff': plots percentage difference between baselien + 'pct_diff': plots percentage difference between baseline and reform ((reform-base)/base) 'diff': plots difference between baseline and reform (reform-base) @@ -343,7 +356,7 @@ def ss_3Dplot( reform_ss (dictionary): SS output from reform run var (string): name of variable to plot plot_type (string): type of plot, can be: - 'pct_diff': plots percentage difference between baselien + 'pct_diff': plots percentage difference between baseline and reform ((reform-base)/base) 'diff': plots difference between baseline and reform (reform-base) 'levels': plot variables in model units diff --git a/ogcore/output_tables.py b/ogcore/output_tables.py index d0514831b..5411af995 100644 --- a/ogcore/output_tables.py +++ b/ogcore/output_tables.py @@ -4,6 +4,7 @@ from ogcore.constants import VAR_LABELS, DEFAULT_START_YEAR from ogcore import tax from ogcore.utils import save_return_table, Inequality +from ogcore.utils import pct_change_unstationarized cur_path = os.path.split(os.path.abspath(__file__))[0] @@ -15,6 +16,7 @@ def macro_table( reform_params=None, var_list=["Y", "C", "K", "L", "r", "w"], output_type="pct_diff", + stationarized=True, num_years=10, include_SS=True, include_overall=True, @@ -38,6 +40,8 @@ def macro_table( and reform ((reform-base)/base) 'diff': plots difference between baseline and reform (reform-base) 'levels': variables in model units + stationarized (bool): whether used stationarized variables (False + only affects pct_diff right now) num_years (integer): number of years to include in table include_SS (bool): whether to include the steady-state results in the table @@ -72,7 +76,17 @@ def macro_table( for i, v in enumerate(var_list): if output_type == "pct_diff": # multiple by 100 so in percentage points - results = ((reform_tpi[v] - base_tpi[v]) / base_tpi[v]) * 100 + if stationarized: + results = ((reform_tpi[v] - base_tpi[v]) / base_tpi[v]) * 100 + else: + pct_changes = pct_change_unstationarized( + base_tpi, + base_params, + reform_tpi, + reform_params, + output_vars=[v], + ) + results = pct_changes[v] * 100 results_years = results[start_index : start_index + num_years] results_overall = ( ( diff --git a/ogcore/parameter_plots.py b/ogcore/parameter_plots.py index df188f652..963eafbfb 100644 --- a/ogcore/parameter_plots.py +++ b/ogcore/parameter_plots.py @@ -64,36 +64,62 @@ def plot_imm_rates( def plot_mort_rates( - p, years=[DEFAULT_START_YEAR], include_title=False, path=None + p_list, + labels=[""], + years=[DEFAULT_START_YEAR], + survival_rates=False, + include_title=False, + path=None, ): """ Create a plot of mortality rates from OG-Core parameterization. Args: - p (OG-Core Specifications class): parameters object + p_list (list): list of parameters objects + labels (list): list of labels for the legend + survival_rates (bool): whether to plot survival rates instead + of mortality rates include_title (bool): whether to include a title in the plot path (string): path to save figure to Returns: - fig (Matplotlib plot object): plot of immigration rates + fig (Matplotlib plot object): plot of mortality rates """ - age_per = np.linspace(p.E, p.E + p.S, p.S) - years = np.array(years) - p.start_year + p0 = p_list[0] + age_per = np.linspace(p0.E, p0.E + p0.S, p0.S) fig, ax = plt.subplots() for y in years: - plt.plot(age_per, p.rho[y, :], label=str(y + p.start_year)) + t = y - p0.start_year + for i, p in enumerate(p_list): + if survival_rates: + plt.plot( + age_per, + np.cumprod(1 - p.rho[t, :]), + label=labels[i] + " " + str(y), + ) + else: + plt.plot(age_per, p.rho[t, :], label=labels[i] + " " + str(y)) plt.xlabel(r"Age $s$ (model periods)") - plt.ylabel(r"Mortality Rates $\rho_{s}$") - plt.legend(loc="upper right") + if survival_rates: + plt.ylabel(r"Cumulative Survival Rates") + plt.legend(loc="lower left") + title = "Survival Rates" + else: + plt.ylabel(r"Mortality Rates $\rho_{s}$") + plt.legend(loc="upper right") + title = "Mortality Rates" vals = ax.get_yticks() ax.set_yticklabels(["{:,.0%}".format(x) for x in vals]) if include_title: - plt.title("Mortality Rates") + plt.title(title) if path is None: return fig else: - fig_path = os.path.join(path, "mortality_rates") + if survival_rates: + fig_path = os.path.join(path, "survival_rates") + else: + fig_path = os.path.join(path, "mortality_rates") plt.savefig(fig_path, dpi=300) @@ -176,13 +202,17 @@ 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, 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. Args: p (OG-Core Specifications class): parameters object t (int): model period for year, if None, then plot ability matrix for SS + log_scale (bool): whether to plot in log points + include_title (bool): whether to include a title in the plot path (string): path to save figure to Returns: @@ -196,10 +226,32 @@ def plot_ability_profiles(p, t=None, include_title=False, path=None): cm = plt.get_cmap("coolwarm") ax.set_prop_cycle(color=[cm(1.0 * i / p.J) for i in range(p.J)]) for j in range(p.J): - plt.plot(age_vec, p.e[t, :, j], label=GROUP_LABELS[p.J][j]) + if log_scale: + 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") - plt.ylabel(r"Earnings ability") - plt.legend(loc=9, bbox_to_anchor=(0.5, -0.15), ncol=2) + if log_scale: + plt.ylabel(r"ln(Earnings ability)") + else: + plt.ylabel(r"Earnings ability") + plt.legend(loc=9, bbox_to_anchor=(0.5, -0.15), ncols=5) if include_title: plt.title("Lifecycle Profiles of Effective Labor Units") if path is None: @@ -267,21 +319,37 @@ def plot_elliptical_u(p, plot_MU=True, include_title=False, path=None): plt.savefig(fig_path, dpi=300) -def plot_chi_n(p, include_title=False, path=None): +def plot_chi_n( + p_list, + labels=[""], + years_to_plot=[DEFAULT_START_YEAR], + include_title=False, + path=None, +): """ Create a plot of showing the values of the chi_n parameters. Args: - p (OG-Core Specifications class): parameters object + p_list (list): parameters objects + labels (list): labels for legend + years_to_plot (list): list of years to plot + include_title (boolean): whether to include a title in the plot path (string): path to save figure to Returns: fig (Matplotlib plot object): plot of chi_n parameters """ - age = np.linspace(p.starting_age, p.ending_age, p.S) + p0 = p_list[0] + age = np.linspace(p0.starting_age, p0.ending_age, p0.S) fig, ax = plt.subplots() - plt.plot(age, p.chi_n) + for y in years_to_plot: + for i, p in enumerate(p_list): + plt.plot( + age, + p.chi_n[y - p.start_year, :], + label=labels[i] + " " + str(y), + ) if include_title: plt.title("Utility Weight on the Disutility of Labor Supply") plt.xlabel("Age, $s$") @@ -294,9 +362,11 @@ def plot_chi_n(p, include_title=False, path=None): def plot_fert_rates( - fert_rates, + fert_rates_list, + labels=[""], start_year=DEFAULT_START_YEAR, years_to_plot=[DEFAULT_START_YEAR], + include_title=False, source="United Nations, World Population Prospects", path=None, ): @@ -304,10 +374,12 @@ def plot_fert_rates( Plot fertility rates from the data Args: - fert_rates (NumPy array): fertility rates for each of - totpers + fert_rates_list (list): list of Numpy arrays of fertility rates + for each model period and age + labels (list): list of labels for the legend start_year (int): first year of data years_to_plot (list): list of years to plot + include_title (bool): whether to include a title in the plot source (str): data source for fertility rates path (str): path to save figure to, if None then figure is returned @@ -321,9 +393,10 @@ def plot_fert_rates( fig, ax = plt.subplots() for y in years_to_plot: i = start_year - y - plt.plot(fert_rates[i, :], c="blue", label="Year " + str(y)) - # plt.title('Fertility rates by age ($f_{s}$)', - # fontsize=20) + for i, fert_rates in enumerate(fert_rates_list): + plt.plot(fert_rates[i, :], label=labels[i] + " " + str(y)) + if include_title: + plt.title("Fertility rates by age ($f_{s}$)", fontsize=20) plt.xlabel(r"Age $s$") plt.ylabel(r"Fertility rate $f_{s}$") plt.legend(loc="upper right") @@ -395,6 +468,39 @@ def plot_mort_rates_data( return fig +def plot_g_n(p_list, label_list=[""], include_title=False, path=None): + """ + Create a plot of population growth rates from OG-Core parameterization. + + Args: + p_list (list): list of OG-Core Specifications objects + label_list (list): list of labels for the legend + include_title (bool): whether to include a title in the plot + path (string): path to save figure to + + Returns: + fig (Matplotlib plot object): plot of immigration rates + + """ + p0 = p_list[0] + years = np.arange(p0.start_year, p0.start_year + p0.T) + fig, ax = plt.subplots() + for i, p in enumerate(p_list): + plt.plot(years, p.g_n[: p.T], label=label_list[i]) + plt.xlabel(r"Year $s$ (model periods)") + plt.ylabel(r"Population Growth Rate $g_{n,t}$") + plt.legend(loc="upper right") + vals = ax.get_yticks() + ax.set_yticklabels(["{:,.0%}".format(x) for x in vals]) + if include_title: + plt.title("Population Growth Rates") + if path is None: + return fig + else: + fig_path = os.path.join(path, "pop_growth_rates") + plt.savefig(fig_path, dpi=300) + + def plot_omega_fixed(age_per_EpS, omega_SS_orig, omega_SSfx, E, S, path=None): """ Plot the steady-state population distribution implied by the data diff --git a/ogcore/utils.py b/ogcore/utils.py index 14a6bf194..623ec67f1 100644 --- a/ogcore/utils.py +++ b/ogcore/utils.py @@ -1218,7 +1218,40 @@ def pct_change_unstationarized( pct_changes = {} T = param_base.T for var in output_vars: - if var in ["K", "C", "Y", "BQ", "TR", "UBI", "D", "total_tax_revenue"]: + if var in [ + "Y", + "B", + "K", + "K_f", + "K_d", + "C", + "I", + "K_g", + "I_g", + "Y_vec", + "K_vec", + "C_vec", + "I_total", + "I_d", + "BQ", + "TR", + "total_tax_revenue", + "business_tax_revenue", + "iit_payroll_tax_revenue", + "iit_revenue", + "payroll_tax_revenue", + "agg_pension_outlays", + "bequest_tax_revenue", + "wealth_tax_revenue", + "cons_tax_revenue", + "G", + "D", + "D_f", + "D_d", + "UBI_path", + "new_borrowing_f", + "debt_service_f", + ]: non_stationary_output["base"][var] = ( tpi_base[var][:T] * np.cumprod(1 + param_base.g_n[:T]) @@ -1229,14 +1262,27 @@ def pct_change_unstationarized( * np.cumprod(1 + param_reform.g_n[:T]) * np.exp(param_reform.g_y * np.arange(param_reform.T)) ) - elif var in ["L"]: + elif var in [ + "L", + "L_vec", + ]: non_stationary_output["base"][var] = tpi_base[var][ :T ] * np.cumprod(1 + param_base.g_n[:T]) non_stationary_output["reform"][var] = tpi_reform[var][ :T ] * np.cumprod(1 + param_reform.g_n[:T]) - elif var in ["w", "ubi", "tr", "bq"]: + elif var in [ + "w", + "ubi_path", + "tr_path", + "bq_path", + "bmat_splus1", + "bmat_s", + "c_path", + "y_before_tax_path", + "tax_path", + ]: non_stationary_output["base"][var] = tpi_base[var][:T] * np.exp( param_base.g_y * np.arange(param_base.T) ) diff --git a/tests/test_output_plots.py b/tests/test_output_plots.py index bee057acc..fa9ff88b9 100644 --- a/tests/test_output_plots.py +++ b/tests/test_output_plots.py @@ -48,19 +48,74 @@ test_data = [ - (base_tpi, base_params, reform_tpi, reform_params, "pct_diff", None, None), - (base_tpi, base_params, reform_tpi, reform_params, "diff", None, None), - (base_tpi, base_params, reform_tpi, reform_params, "forecast", None, None), - (base_tpi, base_params, reform_tpi, reform_params, "levels", None, None), - (base_tpi, base_params, None, None, "levels", None, None), - (base_tpi, base_params, None, None, "levels", [2040, 2060], None), - (base_tpi, base_params, None, None, "levels", None, "Test plot title"), + ( + base_tpi, + base_params, + reform_tpi, + reform_params, + "pct_diff", + False, + None, + None, + ), + ( + base_tpi, + base_params, + reform_tpi, + reform_params, + "diff", + False, + None, + None, + ), + ( + base_tpi, + base_params, + reform_tpi, + reform_params, + "forecast", + False, + None, + None, + ), + ( + base_tpi, + base_params, + reform_tpi, + reform_params, + "levels", + False, + None, + None, + ), + (base_tpi, base_params, None, None, "levels", False, None, None), + (base_tpi, base_params, None, None, "levels", False, [2040, 2060], None), + ( + base_tpi, + base_params, + None, + None, + "levels", + False, + None, + "Test plot title", + ), + ( + base_tpi, + base_params, + reform_tpi, + reform_params, + "pct_diff", + True, + None, + None, + ), ] @pytest.mark.parametrize( "base_tpi,base_params,reform_tpi,reform_parms,plot_type," - + "vertical_line_years,plot_title", + + "stationarized,vertical_line_years,plot_title", test_data, ids=[ "Pct Diff", @@ -70,6 +125,7 @@ "Levels w/o reform", "Vertical line included", "Plot title included", + "Stationarized pct diff", ], ) def test_plot_aggregates( @@ -78,6 +134,7 @@ def test_plot_aggregates( reform_tpi, reform_parms, plot_type, + stationarized, vertical_line_years, plot_title, ): @@ -88,6 +145,7 @@ def test_plot_aggregates( reform_params=reform_params, var_list=["Y", "r"], plot_type=plot_type, + stationarized=stationarized, num_years_to_plot=20, start_year=2023, forecast_data=np.ones(20), @@ -98,6 +156,17 @@ def test_plot_aggregates( assert fig +test_data = [ + (base_tpi, base_params, reform_tpi, reform_params, "pct_diff", None, None), + (base_tpi, base_params, reform_tpi, reform_params, "diff", None, None), + (base_tpi, base_params, reform_tpi, reform_params, "forecast", None, None), + (base_tpi, base_params, reform_tpi, reform_params, "levels", None, None), + (base_tpi, base_params, None, None, "levels", None, None), + (base_tpi, base_params, None, None, "levels", [2040, 2060], None), + (base_tpi, base_params, None, None, "levels", None, "Test plot title"), +] + + @pytest.mark.parametrize( "base_tpi,base_params,reform_tpi,reform_parms,plot_type," + "vertical_line_years,plot_title", diff --git a/tests/test_output_tables.py b/tests/test_output_tables.py index 69dca8272..aa8521938 100644 --- a/tests/test_output_tables.py +++ b/tests/test_output_tables.py @@ -51,19 +51,25 @@ ) test_data = [ - (base_tpi, base_params, reform_tpi, reform_params, "pct_diff"), - (base_tpi, base_params, reform_tpi, reform_params, "diff"), - (base_tpi, base_params, reform_tpi, reform_params, "levels"), + (base_tpi, base_params, reform_tpi, reform_params, "pct_diff", False), + (base_tpi, base_params, reform_tpi, reform_params, "diff", False), + (base_tpi, base_params, reform_tpi, reform_params, "levels", False), + (base_tpi, base_params, reform_tpi, reform_params, "pct_diff", True), ] @pytest.mark.parametrize( - "base_tpi,base_params,reform_tpi,reform_params,output_type", + "base_tpi,base_params,reform_tpi,reform_params,output_type,stationarized", test_data, - ids=["Pct Diff", "Diff", "Levels"], + ids=["Pct Diff", "Diff", "Levels", "Unstationary pct diff"], ) def test_macro_table( - base_tpi, base_params, reform_tpi, reform_params, output_type + base_tpi, + base_params, + reform_tpi, + reform_params, + output_type, + stationarized, ): df = output_tables.macro_table( base_tpi, @@ -72,6 +78,7 @@ def test_macro_table( reform_params=reform_params, start_year=2023, output_type=output_type, + stationarized=stationarized, include_SS=True, include_overall=True, ) diff --git a/tests/test_parameter_plots.py b/tests/test_parameter_plots.py index 993a06d0e..178370bd8 100644 --- a/tests/test_parameter_plots.py +++ b/tests/test_parameter_plots.py @@ -66,17 +66,33 @@ def test_plot_imm_rates_save_fig(tmpdir): def test_plot_mort_rates(): - fig = parameter_plots.plot_mort_rates(base_params, include_title=True) + fig = parameter_plots.plot_mort_rates([base_params], include_title=True) + assert fig + + +def test_plot_surv_rates(): + fig = parameter_plots.plot_mort_rates( + [base_params], survival_rates=True, include_title=True + ) assert fig def test_plot_mort_rates_save_fig(tmpdir): - parameter_plots.plot_mort_rates(base_params, path=tmpdir) + parameter_plots.plot_mort_rates([base_params], path=tmpdir) img = mpimg.imread(os.path.join(tmpdir, "mortality_rates.png")) assert isinstance(img, np.ndarray) +def test_plot_surv_rates_save_fig(tmpdir): + parameter_plots.plot_mort_rates( + [base_params], survival_rates=True, path=tmpdir + ) + img = mpimg.imread(os.path.join(tmpdir, "survival_rates.png")) + + assert isinstance(img, np.ndarray) + + def test_plot_pop_growth(): fig = parameter_plots.plot_pop_growth( base_params, start_year=2023, include_title=True @@ -92,19 +108,22 @@ def test_plot_pop_growth_rates_save_fig(tmpdir): def test_plot_ability_profiles(): - # make save e matrix 3D - base_params.e = np.tile( - base_params.e.reshape(1, base_params.S, base_params.J), - (base_params.T, 1, 1), - ) + p = Specifications() + fig = parameter_plots.plot_ability_profiles(p, p2=p, include_title=True) + assert fig + + +def test_plot_log_ability_profiles(): + p = Specifications() fig = parameter_plots.plot_ability_profiles( - base_params, include_title=True + p, p2=p, log_scale=True, include_title=True ) assert fig def test_plot_ability_profiles_save_fig(tmpdir): - parameter_plots.plot_ability_profiles(base_params, path=tmpdir) + p = Specifications() + parameter_plots.plot_ability_profiles(p, path=tmpdir) img = mpimg.imread(os.path.join(tmpdir, "ability_profiles.png")) assert isinstance(img, np.ndarray) @@ -127,12 +146,14 @@ def test_plot_elliptical_u_save_fig(tmpdir): def test_plot_chi_n(): - fig = parameter_plots.plot_chi_n(base_params, include_title=True) + p = Specifications() + fig = parameter_plots.plot_chi_n([p], include_title=True) assert fig def test_plot_chi_n_save_fig(tmpdir): - parameter_plots.plot_chi_n(base_params, path=tmpdir) + p = Specifications() + parameter_plots.plot_chi_n([p], path=tmpdir) img = mpimg.imread(os.path.join(tmpdir, "chi_n_values.png")) assert isinstance(img, np.ndarray) @@ -184,7 +205,7 @@ def test_plot_fert_rates(): age_midp = np.array([9, 10, 12, 16, 18.5, 22, 27, 32, 37, 42, 47, 55, 56]) fert_func = si.interp1d(age_midp, fert_data, kind="cubic") fert_rates = np.random.uniform(size=totpers).reshape((1, totpers)) - fig = parameter_plots.plot_fert_rates(fert_rates) + fig = parameter_plots.plot_fert_rates([fert_rates], include_title=True) assert fig @@ -216,7 +237,8 @@ def test_plot_fert_rates_save_fig(tmpdir): fert_func = si.interp1d(age_midp, fert_data, kind="cubic") fert_rates = np.random.uniform(size=totpers).reshape((1, totpers)) parameter_plots.plot_fert_rates( - fert_rates, + [fert_rates], + include_title=True, path=tmpdir, ) img = mpimg.imread(os.path.join(tmpdir, "fert_rates.png")) @@ -224,6 +246,20 @@ def test_plot_fert_rates_save_fig(tmpdir): assert isinstance(img, np.ndarray) +def test_plot_g_n(): + p = Specifications() + fig = parameter_plots.plot_g_n([p], include_title=True) + assert fig + + +def test_plot_g_n_savefig(tmpdir): + p = Specifications() + parameter_plots.plot_g_n([p], include_title=True, path=tmpdir) + img = mpimg.imread(os.path.join(tmpdir, "pop_growth_rates.png")) + + assert isinstance(img, np.ndarray) + + def test_plot_mort_rates_data(): totpers = base_params.S - 1 mort_rates = base_params.rho[-1, 1:].reshape((1, totpers))