import numpy as np
import pandas as pd
import os
import matplotlib.pyplot as plt
import matplotlib
import matplotlib.ticker as mticker
from ogcore.constants import GROUP_LABELS
from ogcore import utils, txfunc
from ogcore.constants import DEFAULT_START_YEAR, VAR_LABELS
[docs]
def plot_imm_rates(
imm_rates,
start_year=DEFAULT_START_YEAR,
years_to_plot=[DEFAULT_START_YEAR],
include_title=False,
source="United Nations, World Population Prospects",
path=None,
):
"""
Plot fertility rates from the data
Args:
imm_rates (NumPy array): immigration rates for each of
totpers
start_year (int): first year of data
years_to_plot (list): list of years to plot
source (str): data source for fertility rates
path (str): path to save figure to, if None then figure
is returned
Returns:
fig (Matplotlib plot object): plot of fertility rates
"""
# create line styles to cycle through
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))
# plt.title('Fertility rates by age ($f_{s}$)',
# fontsize=20)
plt.xlabel(r"Age $s$")
plt.ylabel(r"Immigration rate $i_{s}$")
plt.legend(loc="upper left")
plt.text(
-5,
-0.05,
"Source: " + source,
fontsize=9,
)
plt.tight_layout(rect=(0, 0.035, 1, 1))
if include_title:
plt.title("Immigration Rates")
# Save or return figure
if path:
output_path = os.path.join(path, "imm_rates")
plt.savefig(output_path, dpi=300)
plt.close()
else:
fig.show()
return fig
[docs]
def plot_mort_rates(
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_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 mortality rates
"""
p0 = p_list[0]
age_per = np.linspace(p0.E, p0.E + p0.S, p0.S)
fig, ax = plt.subplots()
for y in years:
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)")
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 left")
title = "Mortality Rates"
ticks_loc = ax.get_yticks().tolist()
ax.yaxis.set_major_locator(mticker.FixedLocator(ticks_loc))
ax.set_yticklabels(["{:,.0%}".format(x) for x in ticks_loc])
if include_title:
plt.title(title)
if path is None:
return fig
else:
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)
[docs]
def plot_pop_growth(
p,
start_year=DEFAULT_START_YEAR,
num_years_to_plot=150,
include_title=False,
path=None,
):
"""
Create a plot of population growth rates by year.
Args:
p (OG-Core Specifications class): parameters object
start_year (integer): year to begin plotting
num_years_to_plot (integer): number of years to plot
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
"""
assert isinstance(start_year, int)
assert isinstance(num_years_to_plot, int)
year_vec = np.arange(start_year, start_year + num_years_to_plot)
start_index = start_year - p.start_year
fig, ax = plt.subplots()
plt.plot(year_vec, p.g_n[start_index : start_index + num_years_to_plot])
plt.xlabel(r"Year $t$")
plt.ylabel(r"Population Growth Rate $g_{n, t}$")
ticks_loc = ax.get_yticks().tolist()
ax.yaxis.set_major_locator(mticker.FixedLocator(ticks_loc))
ax.set_yticklabels(["{:,.2%}".format(x) for x in ticks_loc])
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_population(p, years_to_plot=["SS"], include_title=False, path=None):
"""
Plot the distribution of the population over age for various years.
Args:
p (OG-Core Specifications class): parameters object
years_to_plot (list): list of years to plot, 'SS' will denote
the steady-state period
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 population distribution
"""
for i, v in enumerate(years_to_plot):
assert isinstance(v, int) | (v == "SS")
if isinstance(v, int):
assert v >= p.start_year
age_vec = np.arange(p.E, p.S + p.E)
fig, ax = plt.subplots()
for i, v in enumerate(years_to_plot):
if v == "SS":
pop_dist = p.omega_SS
else:
pop_dist = p.omega[v - p.start_year, :]
plt.plot(age_vec, pop_dist, label=str(v) + " pop.")
plt.xlabel(r"Age $s$")
plt.ylabel(r"Pop. dist'n $\omega_{s}$")
plt.legend(loc="lower left")
if include_title:
plt.title("Population Distribution by Year")
if path is None:
return fig
else:
fig_path = os.path.join(path, "pop_distribution")
plt.savefig(fig_path, dpi=300)
[docs]
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:
fig (Matplotlib plot object): plot of earnings ability profiles
"""
if t is None:
t = -1
age_vec = np.arange(p.starting_age, p.starting_age + p.S)
fig, ax = plt.subplots()
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):
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")
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:
return fig
else:
fig_path = os.path.join(path, "ability_profiles")
plt.savefig(fig_path, bbox_inches="tight", dpi=300)
[docs]
def plot_elliptical_u(p, plot_MU=True, include_title=False, path=None):
"""
Create a plot of showing the fit of the elliptical utility function.
Args:
p (OG-Core Specifications class): parameters object
plot_MU (boolean): whether plot marginal utility or utility in
levels
path (string): path to save figure to
Returns:
fig (Matplotlib plot object): plot of elliptical vs CFE utility
"""
theta = 1 / p.frisch
N = 101
n_grid = np.linspace(0.01, 0.8, num=N)
if plot_MU:
CFE = (1.0 / p.ltilde) * ((n_grid / p.ltilde) ** theta)
ellipse = (
1.0
* p.b_ellipse
* (1.0 / p.ltilde)
* (
(1.0 - (n_grid / p.ltilde) ** p.upsilon)
** ((1.0 / p.upsilon) - 1.0)
)
* (n_grid / p.ltilde) ** (p.upsilon - 1.0)
)
else:
CFE = ((n_grid / p.ltilde) ** (1 + theta)) / (1 + theta)
k = 1.0 # we don't estimate k, so not in parameters
ellipse = (
p.b_ellipse
* ((1 - ((n_grid / p.ltilde) ** p.upsilon)) ** (1 / p.upsilon))
+ k
)
fig, ax = plt.subplots()
plt.plot(n_grid, CFE, label="Constant Frisch elasticity")
plt.plot(n_grid, ellipse, label="Elliptical disutility")
if include_title:
if plot_MU:
plt.title("Marginal Utility of CFE and Elliptical")
else:
plt.title("Constant Frisch Elasticity vs. Elliptical Utility")
plt.xlabel(r"Labor Supply $n_{j,s,t}$")
if plot_MU:
plt.ylabel(r"Marginal disutility")
else:
plt.ylabel(r"Disutility")
plt.legend(loc="upper left")
plt.grid(color="gray", linestyle=":", linewidth=1, alpha=0.5)
if path is None:
return fig
else:
fig_path = os.path.join(path, "ellipse_v_CFE")
plt.savefig(fig_path, dpi=300)
[docs]
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_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
"""
p0 = p_list[0]
age = np.linspace(p0.starting_age, p0.ending_age, p0.S)
fig, ax = plt.subplots()
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$")
plt.ylabel(r"$\chi^{n}_{s}$")
if path is None:
return fig
else:
fig_path = os.path.join(path, "chi_n_values")
plt.savefig(fig_path, dpi=300)
[docs]
def plot_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,
):
"""
Plot fertility rates from the data
Args:
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
Returns:
fig (Matplotlib plot object): plot of fertility rates
"""
# create line styles to cycle through
fig, ax = plt.subplots()
for y in years_to_plot:
i = start_year - y
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")
plt.text(
-5,
-0.023,
"Source: " + source,
fontsize=9,
)
plt.tight_layout(rect=(0, 0.035, 1, 1))
# Save or return figure
if path:
output_path = os.path.join(path, "fert_rates")
plt.savefig(output_path, dpi=300)
plt.close()
else:
fig.show()
return fig
[docs]
def plot_mort_rates_data(
mort_rates,
start_year=DEFAULT_START_YEAR,
years_to_plot=[DEFAULT_START_YEAR],
source="United Nations, World Population Prospects",
path=None,
):
"""
Plots mortality rates from the data.
Args:
mort_rates (array_like): mortality rates for each of
totpers
start_year (int): first year of data
years_to_plot (list): list of years to plot
source (str): data source for fertility rates
path (str): path to save figure to, if None then figure
is returned
Returns:
fig (Matplotlib plot object): plot of mortality rates
"""
# create line styles to cycle through
fig, ax = plt.subplots()
for y in years_to_plot:
i = start_year - y
plt.plot(mort_rates[i, :], c="blue", label="Year " + str(y))
# plt.title('Fertility rates by age ($f_{s}$)',
# fontsize=20)
plt.xlabel(r"Age $s$")
plt.ylabel(r"Mortality rate $rho_{s}$")
plt.legend(loc="upper left")
plt.text(
-5,
-0.223,
"Source: " + source,
fontsize=9,
)
plt.tight_layout(rect=(0, 0.035, 1, 1))
# Save or return figure
if path:
output_path = os.path.join(path, "mort_rates")
plt.savefig(output_path, dpi=300)
plt.close()
else:
fig.show()
return fig
[docs]
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}$")
if label_list[0] != "":
plt.legend(loc="upper right")
ticks_loc = ax.get_yticks().tolist()
ax.yaxis.set_major_locator(mticker.FixedLocator(ticks_loc))
ax.set_yticklabels(["{:,.0%}".format(x) for x in ticks_loc])
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)
[docs]
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
on fertility and mortality rates versus the the steady-state
population distribution after adjusting immigration rates so that
the stationary distribution is achieved a reasonable number of
model periods.
Args:
age_per_EpS (array_like): list of ages over which to plot
population distribution
omega_SS_orig (Numpy array): population distribution in SS
without adjustment to immigration rates
omega_SSfx (Numpy array): population distribution in SS
after adjustment to immigration rates
E (int): age at which household becomes economically active
S (int): number of years which household is economically active
path (str): path to save figure to, if None then figure
is returned
Returns:
fig (Matplotlib plot object): plot of SS population distribution
before and after adjustment to immigration rates
"""
fig, ax = plt.subplots()
plt.plot(age_per_EpS, omega_SS_orig, label="Original Dist'n")
plt.plot(age_per_EpS, omega_SSfx, label="Fixed Dist'n")
plt.title("Original steady-state population distribution vs. fixed")
plt.xlabel(r"Age $s$")
plt.ylabel(r"Pop. dist'n $\omega_{s}$")
plt.xlim((0, E + S + 1))
plt.legend(loc="upper right")
# Save or return figure
if path:
output_path = os.path.join(path, "OrigVsFixSSpop")
plt.savefig(output_path, dpi=300)
plt.close()
else:
return fig
[docs]
def plot_imm_fixed(
age_per_EpS, imm_rates_orig, imm_rates_adj, E, S, path=None
):
"""
Plot the immigration rates implied by the data on population,
mortality, and fertility versus the adjusted immigration rates
needed to achieve a stationary distribution of the population in a
reasonable number of model periods.
Args:
age_per_EpS (array_like): list of ages over which to plot
population distribution
imm_rates_orig (Numpy array): immigration rates by age
imm_rates_adj (Numpy array): adjusted immigration rates by age
E (int): age at which household becomes economically active
S (int): number of years which household is economically active
path (str): path to save figure to, if None then figure
is returned
Returns:
fig (Matplotlib plot object): plot of immigration rates found
from residuals and the adjusted rates to hit SS sooner
"""
fig, ax = plt.subplots()
plt.plot(age_per_EpS, imm_rates_orig, label="Original Imm. Rates")
plt.plot(age_per_EpS, imm_rates_adj, label="Adj. Imm. Rates")
plt.title("Original immigration rates vs. adjusted")
plt.xlabel(r"Age $s$")
plt.ylabel(r"Imm. rates $i_{s}$")
plt.xlim((0, E + S + 1))
plt.legend(loc="upper center")
# Save or return figure
if path:
output_path = os.path.join(path, "OrigVsAdjImm")
plt.savefig(output_path, dpi=300)
plt.close()
else:
return fig
[docs]
def plot_population_path(
age_per_EpS,
omega_path_lev,
omega_SSfx,
start_year,
year1,
year2,
S,
path=None,
):
"""
Plot the distribution of the population over age for various years.
Args:
age_per_EpS (array_like): list of ages over which to plot
population distribution
initial_pop_pct (array_like): initial year population distribution
omega_path_lev (Numpy array): number of households by age
over the transition path
omega_SSfx (Numpy array): number of households by age
in the SS
start_year (int): first year of data (so can get index of year1
and year2)
year1 (int): first year of data to plot
year2 (int): second year of data to plot
S (int): number of years which household is economically active
path (str): path to save figure to, if None then figure
is returned
Returns:
fig (Matplotlib plot object): plot of population distribution
at points along the time path
"""
fig, ax = plt.subplots()
plt.plot(
age_per_EpS,
(
omega_path_lev[start_year - year1, :]
/ omega_path_lev[start_year - year1, :].sum()
),
label=str(year1) + " pop.",
)
plt.plot(
age_per_EpS,
(
omega_path_lev[start_year - year2, :]
/ omega_path_lev[start_year - year2, :].sum()
),
label=str(year2) + " pop.",
)
plt.plot(
age_per_EpS,
(
omega_path_lev[int(0.5 * S), :]
/ omega_path_lev[int(0.5 * S), :].sum()
),
label="T=" + str(int(0.5 * S)) + " pop.",
)
plt.plot(
age_per_EpS,
(omega_path_lev[int(S), :] / omega_path_lev[int(S), :].sum()),
label="T=" + str(int(S)) + " pop.",
)
plt.plot(age_per_EpS, omega_SSfx, label="Adj. SS pop.")
plt.title("Population distribution at points in time path")
plt.xlabel(r"Age $s$")
plt.ylabel(r"Pop. dist'n $\omega_{s}$")
plt.legend(loc="lower left")
# Save or return figure
if path:
output_path = os.path.join(path, "PopDistPath")
plt.savefig(output_path, dpi=300)
plt.close()
else:
return fig
[docs]
def gen_3Dscatters_hist(df, s, t, output_dir):
"""
Create 3-D scatterplots and corresponding 3D histogram of ETR, MTRx,
and MTRy as functions of labor income and capital income with
truncated data in the income dimension
Args:
df (Pandas DataFrame): 11 variables with N observations of tax
rates
s (int): age of individual, >= 21
t (int): year of analysis, >= 2016
path (str): output directory for saving plot files
Returns:
None
"""
from ogcore.txfunc import MAX_INC_GRAPH, MIN_INC_GRAPH
# Truncate the data
df_trnc = df[
(df["total_labinc"] > MIN_INC_GRAPH)
& (df["total_labinc"] < MAX_INC_GRAPH)
& (df["total_capinc"] > MIN_INC_GRAPH)
& (df["total_capinc"] < MAX_INC_GRAPH)
]
inc_lab = df_trnc["total_labinc"]
inc_cap = df_trnc["total_capinc"]
etr_data = df_trnc["etr"]
mtrx_data = df_trnc["mtr_labinc"]
mtry_data = df_trnc["mtr_capinc"]
# Plot 3D scatterplot of ETR data
fig = plt.figure()
ax = fig.add_subplot(111, projection="3d")
ax.scatter(inc_lab, inc_cap, etr_data, c="r", marker="o")
ax.set_xlabel("Total Labor Income")
ax.set_ylabel("Total Capital Income")
ax.set_zlabel("ETR")
plt.title(
"ETR, Lab. Inc., and Cap. Inc., Age=" + str(s) + ", Year=" + str(t)
)
filename = "ETR_age_" + str(s) + "_Year_" + str(t) + "_data.png"
fullpath = os.path.join(output_dir, filename)
fig.savefig(fullpath, bbox_inches="tight", dpi=300)
plt.close()
# Plot 3D histogram for all data
fig = plt.figure()
ax = fig.add_subplot(111, projection="3d")
bin_num = int(30)
hist, xedges, yedges = np.histogram2d(inc_lab, inc_cap, bins=bin_num)
hist = hist / hist.sum()
x_midp = xedges[:-1] + 0.5 * (xedges[1] - xedges[0])
y_midp = yedges[:-1] + 0.5 * (yedges[1] - yedges[0])
elements = (len(xedges) - 1) * (len(yedges) - 1)
ypos, xpos = np.meshgrid(y_midp, x_midp)
xpos = xpos.flatten()
ypos = ypos.flatten()
zpos = np.zeros(elements)
dx = (xedges[1] - xedges[0]) * np.ones_like(bin_num)
dy = (yedges[1] - yedges[0]) * np.ones_like(bin_num)
dz = hist.flatten()
ax.bar3d(xpos, ypos, zpos, dx, dy, dz, color="b", zsort="average")
ax.set_xlabel("Total Labor Income")
ax.set_ylabel("Total Capital Income")
ax.set_zlabel("Percent of obs.")
plt.title(
"Histogram by lab. inc., and cap. inc., Age="
+ str(s)
+ ", Year="
+ str(t)
)
filename = "Hist_Age_" + str(s) + "_Year_" + str(t) + ".png"
fullpath = os.path.join(output_dir, filename)
fig.savefig(fullpath, bbox_inches="tight", dpi=300)
plt.close()
# Plot 3D scatterplot of MTRx data
fig = plt.figure()
ax = fig.add_subplot(111, projection="3d")
ax.scatter(inc_lab, inc_cap, mtrx_data, c="r", marker="o")
ax.set_xlabel("Total Labor Income")
ax.set_ylabel("Total Capital Income")
ax.set_zlabel("Marginal Tax Rate, Labor Inc.)")
plt.title(
"MTR Labor Income, Lab. Inc., and Cap. Inc., Age="
+ str(s)
+ ", Year="
+ str(t)
)
filename = "MTRx_Age_" + str(s) + "_Year_" + str(t) + "_data.png"
fullpath = os.path.join(output_dir, filename)
fig.savefig(fullpath, bbox_inches="tight", dpi=300)
plt.close()
# Plot 3D scatterplot of MTRy data
fig = plt.figure()
ax = fig.add_subplot(111, projection="3d")
ax.scatter(inc_lab, inc_cap, mtry_data, c="r", marker="o")
ax.set_xlabel("Total Labor Income")
ax.set_ylabel("Total Capital Income")
ax.set_zlabel("Marginal Tax Rate (Capital Inc.)")
plt.title(
"MTR Capital Income, Cap. Inc., and Cap. Inc., Age="
+ str(s)
+ ", Year="
+ str(t)
)
filename = "MTRy_Age_" + str(s) + "_Year_" + str(t) + "_data.png"
fullpath = os.path.join(output_dir, filename)
fig.savefig(fullpath, bbox_inches="tight", dpi=300)
plt.close()
# Garbage collection
del df, df_trnc, inc_lab, inc_cap, etr_data, mtrx_data, mtry_data
[docs]
def txfunc_graph(
s,
t,
df,
X,
Y,
txrates,
rate_type,
tax_func_type,
params_to_plot,
output_dir,
):
"""
This function creates a 3D plot of the fitted tax function against
the data.
Args:
s (int): age of individual, >= 21
t (int): year of analysis, >= 2016
df (Pandas DataFrame): 11 variables with N observations of tax
rates
X (Pandas DataSeries): labor income
Y (Pandas DataSeries): capital income
Y (Pandas DataSeries): tax rates from the data
rate_type (str): type of tax rate: mtrx, mtry, etr
tax_func_type (str): functional form of tax functions
params_to_plot (array_like or function): tax function parameters or
nonparametric function
path (str): output directory for saving plot files
Returns:
None
"""
cmap1 = matplotlib.cm.get_cmap("summer")
# Make comparison plot with full income domains
fig = plt.figure()
ax = fig.add_subplot(111, projection="3d")
ax.scatter(X, Y, txrates, c="r", marker="o")
ax.set_xlabel("Total Labor Income")
ax.set_ylabel("Total Capital Income")
if rate_type == "etr":
tx_label = "ETR"
elif rate_type == "mtrx":
tx_label = "MTRx"
elif rate_type == "mtry":
tx_label = "MTRy"
ax.set_zlabel(tx_label)
plt.title(
tx_label
+ " vs. Predicted "
+ tx_label
+ ": Age="
+ str(s)
+ ", Year="
+ str(t)
)
gridpts = 50
X_vec = np.exp(np.linspace(np.log(5), np.log(X.max()), gridpts))
Y_vec = np.exp(np.linspace(np.log(5), np.log(Y.max()), gridpts))
X_grid, Y_grid = np.meshgrid(X_vec, Y_vec)
txrate_grid = txfunc.get_tax_rates(
params_to_plot,
X_grid,
Y_grid,
None,
tax_func_type,
rate_type,
for_estimation=False,
)
ax.plot_surface(X_grid, Y_grid, txrate_grid, cmap=cmap1, linewidth=0)
filename = tx_label + "_age_" + str(s) + "_Year_" + str(t) + "_vsPred.png"
fullpath = os.path.join(output_dir, filename)
fig.savefig(fullpath, bbox_inches="tight", dpi=300)
plt.close()
# Make comparison plot with truncated income domains
df_trnc_gph = df[
(df["total_labinc"] > 5)
& (df["total_labinc"] < 800000)
& (df["total_capinc"] > 5)
& (df["total_capinc"] < 800000)
]
X_gph = df_trnc_gph["total_labinc"]
Y_gph = df_trnc_gph["total_capinc"]
if rate_type == "etr":
txrates_gph = df_trnc_gph["etr"]
elif rate_type == "mtrx":
txrates_gph = df_trnc_gph["mtr_labinc"]
elif rate_type == "mtry":
txrates_gph = df_trnc_gph["mtr_capinc"]
fig = plt.figure()
ax = fig.add_subplot(111, projection="3d")
ax.scatter(X_gph, Y_gph, txrates_gph, c="r", marker="o")
ax.set_xlabel("Total Labor Income")
ax.set_ylabel("Total Capital Income")
ax.set_zlabel(tx_label)
plt.title(
"Truncated "
+ tx_label
+ ", Lab. Inc., and Cap. "
+ "Inc., Age="
+ str(s)
+ ", Year="
+ str(t)
)
gridpts = 50
X_vec = np.exp(np.linspace(np.log(5), np.log(X_gph.max()), gridpts))
Y_vec = np.exp(np.linspace(np.log(5), np.log(Y_gph.max()), gridpts))
X_grid, Y_grid = np.meshgrid(X_vec, Y_vec)
txrate_grid = txfunc.get_tax_rates(
params_to_plot,
X_grid,
Y_grid,
None,
tax_func_type,
rate_type,
for_estimation=False,
)
ax.plot_surface(X_grid, Y_grid, txrate_grid, cmap=cmap1, linewidth=0)
filename = (
tx_label + "trunc_age_" + str(s) + "_Year_" + str(t) + "_vsPred.png"
)
fullpath = os.path.join(output_dir, filename)
fig.savefig(fullpath, bbox_inches="tight", dpi=300)
plt.close()
[docs]
def txfunc_sse_plot(age_vec, sse_mat, start_year, varstr, output_dir, round):
"""
Plot sum of squared errors of tax functions over age for each year
of budget window.
Args:
age_vec (numpy array): vector of ages, length S
sse_mat (Numpy array): SSE for each estimated tax function,
size is BW x S
start_year (int): first year of budget window
varstr (str): name of tax function being evaluated
path (str): path to save graph to
round (int): which round of sweeping for outliers (0, 1, or 2)
Returns:
None
"""
fig, ax = plt.subplots()
BW = sse_mat.shape[0]
for y in range(BW):
plt.plot(age_vec, sse_mat[y, :], label=str(start_year + y))
plt.legend(loc="upper left")
titletext = (
"Sum of Squared Errors by age and Tax Year"
+ " minus outliers (Round "
+ str(round)
+ "): "
+ varstr
)
plt.title(titletext)
plt.xlabel(r"age $s$")
plt.ylabel(r"SSE")
graphname = "SSE_" + varstr + "_Round" + str(round)
output_path = os.path.join(output_dir, graphname)
plt.savefig(output_path, bbox_inches="tight", dpi=300)
plt.close()
[docs]
def plot_income_data(
ages, abil_midp, abil_pcts, emat, t=None, path=None, filesuffix=""
):
"""
This function graphs ability matrix in 3D, 2D, log, and nolog
Args:
ages (Numpy array) ages represented in sample, length S
abil_midp (Numpy array): midpoints of income percentile bins in
each ability group
abil_pcts (Numpy array): percent of population in each lifetime
income group, length J
emat (Numpy array): effective labor units by age and lifetime
income group, size TxSxJ
t (int): model period for year, if None, then plot SS values
filesuffix (str): suffix to be added to plot files
Returns:
None
"""
if t is None:
t = -1
J = abil_midp.shape[0]
abil_mesh, age_mesh = np.meshgrid(abil_midp, ages)
cmap1 = matplotlib.colormaps["summer"]
if path:
# Make sure that directory is created
utils.mkdirs(path)
if J == 1:
# Plot of 2D, J=1 in levels
plt.figure()
plt.plot(ages, emat[t, :, :])
filename = "ability_2D_lev" + filesuffix
fullpath = os.path.join(path, filename)
plt.savefig(fullpath, dpi=300)
plt.close()
# Plot of 2D, J=1 in logs
plt.figure()
plt.plot(ages, np.log(emat[t, :, :]))
filename = "ability_2D_log" + filesuffix
fullpath = os.path.join(path, filename)
plt.savefig(fullpath, dpi=300)
plt.close()
else:
# Plot of 3D, J>1 in levels
fig10, ax10 = plt.subplots(subplot_kw={"projection": "3d"})
ax10.plot_surface(
age_mesh,
abil_mesh,
emat[t, :, :],
rstride=8,
cstride=1,
cmap=cmap1,
)
ax10.set_xlabel(r"age-$s$")
ax10.set_ylabel(r"ability type -$j$")
ax10.set_zlabel(r"ability $e_{j,s}$")
filename = "ability_3D_lev" + filesuffix
fullpath = os.path.join(path, filename)
plt.savefig(fullpath, dpi=300)
plt.close()
# Plot of 3D, J>1 in logs
fig11, ax11 = plt.subplots(subplot_kw={"projection": "3d"})
ax11.plot_surface(
age_mesh,
abil_mesh,
np.log(emat[t, :, :]),
rstride=8,
cstride=1,
cmap=cmap1,
)
ax11.set_xlabel(r"age-$s$")
ax11.set_ylabel(r"ability type -$j$")
ax11.set_zlabel(r"log ability $log(e_{j,s})$")
filename = "ability_3D_log" + filesuffix
fullpath = os.path.join(path, filename)
plt.savefig(fullpath, dpi=300)
plt.close()
if J <= 10: # Restricted because of line and marker types
# Plot of 2D lines from 3D version in logs
ax = plt.subplot(111)
linestyles = np.array(
[
"-",
"--",
"-.",
":",
]
)
markers = np.array(["x", "v", "o", "d", ">", "|"])
pct_lb = 0
for j in range(J):
this_label = (
str(int(np.rint(pct_lb)))
+ " - "
+ str(int(np.rint(pct_lb + 100 * abil_pcts[j])))
+ "%"
)
pct_lb += 100 * abil_pcts[j]
if j <= 3:
ax.plot(
ages,
np.log(emat[t, :, j]),
label=this_label,
linestyle=linestyles[j],
color="black",
)
elif j > 3:
ax.plot(
ages,
np.log(emat[t, :, j]),
label=this_label,
marker=markers[j - 4],
color="black",
)
ax.axvline(x=80, color="black", linestyle="--")
box = ax.get_position()
ax.set_position([box.x0, box.y0, box.width * 0.8, box.height])
ax.legend(loc="center left", bbox_to_anchor=(1, 0.5))
ax.set_xlabel(r"age-$s$")
ax.set_ylabel(r"log ability $log(e_{j,s})$")
filename = "ability_2D_log" + filesuffix
fullpath = os.path.join(path, filename)
plt.savefig(fullpath, dpi=300)
plt.close()
else:
if J <= 10: # Restricted because of line and marker types
# Plot of 2D lines from 3D version in logs
ax = plt.subplot(111)
linestyles = np.array(
[
"-",
"--",
"-.",
":",
]
)
markers = np.array(["x", "v", "o", "d", ">", "|"])
pct_lb = 0
for j in range(J):
this_label = (
str(int(np.rint(pct_lb)))
+ " - "
+ str(int(np.rint(pct_lb + 100 * abil_pcts[j])))
+ "%"
)
pct_lb += 100 * abil_pcts[j]
if j <= 3:
ax.plot(
ages,
np.log(emat[t, :, j]),
label=this_label,
linestyle=linestyles[j],
color="black",
)
elif j > 3:
ax.plot(
ages,
np.log(emat[t, :, j]),
label=this_label,
marker=markers[j - 4],
color="black",
)
ax.axvline(x=80, color="black", linestyle="--")
box = ax.get_position()
ax.set_position([box.x0, box.y0, box.width * 0.8, box.height])
ax.legend(loc="center left", bbox_to_anchor=(1, 0.5))
ax.set_xlabel(r"age-$s$")
ax.set_ylabel(r"log ability $log(e_{j,s})$")
return ax
[docs]
def plot_2D_taxfunc(
year,
start_year,
tax_param_list,
age=None,
E=21, # Age at which agents become economically active in the model
tax_func_type=["DEP"],
rate_type="etr",
over_labinc=True,
other_inc_val=1000,
max_inc_amt=1000000,
data_list=None,
labels=["1st Functions"],
title=None,
path=None,
):
"""
This function plots OG-Core tax functions in two dimensions.
The tax rates are plotted over capital or labor income, as
entered by the user.
Args:
year (int): year of policy tax functions represent
start_year (int): first year tax functions estimated for in
tax_param_list elements
tax_param_list (list): list of arrays containing tax function
parameters
age (int): age for tax functions to plot, use None if tax
function parameters were not age specific
tax_func_type (list): list of strings in ["DEP", "DEP_totalinc",
"GS", "linear"] and specifies functional form of tax functions
in tax_param_list
rate_type (str): string that is in ["etr", "mtrx", "mtry"] and
determines the type of tax rate that is plotted
over_labinc (bool): indicates that x-axis of the plot is over
labor income, if False then plot is over capital income
other_inc_val (scalar): dollar value at which to hold constant
the amount of income that is not represented on the x-axis
max_inc_amt (scalar): largest income amount to represent on the
x-axis of the plot
data_list (list): list of DataFrames with data to scatter plot
with tax functions, needs to be of format output from
ogcore.get_micro_data.get_data
labels (list): list of labels for tax function parameters
title (str): title for the plot
path (str): path to which to save plot, if None then figure
returned
Returns:
fig (Matplotlib plot object): plot of tax functions
"""
# Check that inputs are valid
assert isinstance(start_year, int)
assert isinstance(year, int)
assert year >= start_year
# if list of tax function types less than list of params, assume
# all the same functional form
if len(tax_func_type) < len(tax_param_list):
tax_func_type = [tax_func_type[0]] * len(tax_param_list)
for i, v in enumerate(tax_func_type):
assert v in ["DEP", "DEP_totalinc", "GS", "linear", "mono", "mono2D"]
assert rate_type in ["etr", "mtrx", "mtry"]
assert len(tax_param_list) == len(labels)
# Set age and year to look at
if age is not None:
assert isinstance(age, int)
assert age >= E
s = (
age - E
) # Note: assumed age is given in E + model periods (but age below is also assumed to be calendar years)
else:
s = 0 # if not age-specific, all ages have the same values
t = year - start_year
# create rate_key to correspond to keys in tax func dicts
rate_key = "tfunc_" + rate_type + "_params_S"
# Set income range to plot over (min income value hard coded to 5)
inc_sup = np.exp(np.linspace(np.log(5), np.log(max_inc_amt), 100))
# Set income value for other income
inc_fix = other_inc_val
if over_labinc:
key1 = "total_labinc"
X = inc_sup
Y = inc_fix
else:
key1 = "total_capinc"
X = inc_fix
Y = inc_sup
# get tax rates for each point in the income support and plot
fig, ax = plt.subplots()
for i, tax_params in enumerate(tax_param_list):
tax_params = tax_params[rate_key][t][s]
rates = txfunc.get_tax_rates(
tax_params,
X,
Y,
None,
tax_func_type[i],
rate_type,
for_estimation=False,
)
plt.plot(inc_sup, rates, label=labels[i])
# plot raw data (if passed)
if data_list is not None:
rate_type_dict = {
"etr": "etr",
"mtrx": "mtr_labinc",
"mtry": "mtr_capinc",
}
# censor data to range of the plot
for d, data in enumerate(data_list):
data_to_plot = data[str(year)].copy()
if age is not None:
data_to_plot.drop(
data_to_plot[data_to_plot["age"] != age].index,
inplace=True,
)
# other censoring
data_to_plot.drop(
data_to_plot[data_to_plot[key1] > max_inc_amt].index,
inplace=True,
)
# other censoring used in txfunc.py
data_to_plot = txfunc.tax_data_sample(data_to_plot)
# set number of bins to 100 or bins of $1000 dollars
n_bins = min(100, np.floor_divide(max_inc_amt, 1000))
# need to compute weighted averages by group...
def weighted_mean(x, cols, w="weight"):
try:
return pd.Series(
np.average(x[cols], weights=x[w], axis=0), cols
)
except ZeroDivisionError:
return 0
data_to_plot["inc_bin"] = pd.cut(data_to_plot[key1], n_bins)
groups = data_to_plot.groupby("inc_bin", observed=True).apply(
weighted_mean, [rate_type_dict[rate_type], key1]
)
plt.scatter(
groups[key1], groups[rate_type_dict[rate_type]], alpha=0.1
)
# add legend, labels, etc to plot
plt.legend(loc="center right")
if title:
plt.title(title)
if over_labinc:
plt.xlabel(r"Labor income")
else:
plt.xlabel(r"Capital income")
plt.ylabel(VAR_LABELS[rate_type])
if path is None:
return fig
else:
plt.savefig(path, dpi=300)
