diff --git a/environment.yml b/environment.yml index d3c8724d2..e90cd8f2b 100644 --- a/environment.yml +++ b/environment.yml @@ -2,7 +2,7 @@ name: ogcore-dev channels: - conda-forge dependencies: -- python>=3.7.7 +- python>=3.7.7, <3.12 - ipython - setuptools - scipy>=1.7.1 diff --git a/ogcore/SS.py b/ogcore/SS.py index b8b66afb0..e4183a09f 100644 --- a/ogcore/SS.py +++ b/ogcore/SS.py @@ -78,7 +78,6 @@ def euler_equation_solver(guesses, *args): tr, ubi, theta, - p.e[:, j], p.rho[-1, :], p.etr_params[-1], p.mtry_params[-1], @@ -100,7 +99,6 @@ def euler_equation_solver(guesses, *args): ubi, theta, p.chi_n, - p.e[:, j], p.etr_params[-1], p.mtrx_params[-1], None, @@ -138,7 +136,7 @@ def euler_equation_solver(guesses, *args): j, False, "SS", - p.e[:, j], + p.e[-1, :, j], p.etr_params[-1], p, ) @@ -151,7 +149,7 @@ def euler_equation_solver(guesses, *args): n_guess, bq, taxes, - p.e[:, j], + p.e[-1, :, j], p, ) mask6 = cons < 0 @@ -293,7 +291,7 @@ def inner_loop(outer_loop_vars, p, client): None, False, "SS", - p.e, + np.squeeze(p.e[-1, :, :]), etr_params_3D, p, ) @@ -306,7 +304,7 @@ def inner_loop(outer_loop_vars, p, client): nssmat, bq, net_tax, - p.e, + np.squeeze(p.e[-1, :, :]), p, ) c_i = household.get_ci(c_s, p_i, p_tilde, p.tau_c[-1, :], p.alpha_c) @@ -383,7 +381,7 @@ def inner_loop(outer_loop_vars, p, client): new_r, new_r_gov, p_m, K_vec, K_g, D, MPKg_vec, p, "SS" ) average_income_model = ( - (new_r_p * b_s + new_w * p.e * nssmat) + (new_r_p * b_s + new_w * np.squeeze(p.e[-1, :, :]) * nssmat) * p.omega_SS.reshape(p.S, 1) * p.lambdas.reshape(1, p.J) ).sum() @@ -425,7 +423,7 @@ def inner_loop(outer_loop_vars, p, client): None, False, "SS", - p.e, + np.squeeze(p.e[-1, :, :]), etr_params_3D, p, ) @@ -438,7 +436,7 @@ def inner_loop(outer_loop_vars, p, client): nssmat, new_bq, taxss, - p.e, + np.squeeze(p.e[-1, :, :]), p, ) ( @@ -467,6 +465,7 @@ def inner_loop(outer_loop_vars, p, client): ubi, theta, etr_params_3D, + np.squeeze(p.e[-1, :, :]), p, None, "SS", @@ -482,6 +481,9 @@ def inner_loop(outer_loop_vars, p, client): debt_service, p, ) + print("Agg tax = ", total_tax_revenue) + print("Agg pension outlays = ", agg_pension_outlays) + print("Agg UBI outlays = ", UBI_outlays) new_TR = fiscal.get_TR( Y, TR, @@ -757,7 +759,7 @@ def SS_solver( nssmat, factor, True, - p.e, + np.squeeze(p.e[-1, :, :]), etr_params_3D, mtry_params_3D, capital_noncompliance_rate_2D, @@ -770,7 +772,7 @@ def SS_solver( nssmat, factor, False, - p.e, + np.squeeze(p.e[-1, :, :]), etr_params_3D, mtrx_params_3D, labor_noncompliance_rate_2D, @@ -782,7 +784,7 @@ def SS_solver( bssmat_s, nssmat, factor, - p.e, + np.squeeze(p.e[-1, :, :]), etr_params_3D, labor_noncompliance_rate_2D, capital_noncompliance_rate_2D, @@ -803,7 +805,7 @@ def SS_solver( None, False, "SS", - p.e, + np.squeeze(p.e[-1, :, :]), etr_params_3D, p, ) @@ -816,10 +818,12 @@ def SS_solver( nssmat, bqssmat, taxss, - p.e, + np.squeeze(p.e[-1, :, :]), p, ) - yss_before_tax_mat = household.get_y(r_p_ss, wss, bssmat_s, nssmat, p) + yss_before_tax_mat = household.get_y( + r_p_ss, wss, bssmat_s, nssmat, p, "SS" + ) Css = aggr.get_C(cssmat, p, "SS") c_i_ss_mat = household.get_ci( cssmat, p_i_ss, p_tilde_ss, p.tau_c[-1, :], p.alpha_c @@ -860,6 +864,7 @@ def SS_solver( ubissmat, theta, etr_params_3D, + np.squeeze(p.e[-1, :, :]), p, None, "SS", diff --git a/ogcore/TPI.py b/ogcore/TPI.py index 4c3bdc5ff..6797fb8d9 100644 --- a/ogcore/TPI.py +++ b/ogcore/TPI.py @@ -154,7 +154,6 @@ def firstdoughnutring( np.array([tr]), np.array([ubi]), theta[j], - p.e[-1, j], p.rho[0, -1], p.etr_params[0][-1], p.mtry_params[0][-1], @@ -177,7 +176,6 @@ def firstdoughnutring( np.array([ubi]), theta[j], p.chi_n[-1], - p.e[-1, j], p.etr_params[0][-1], p.mtrx_params[0][-1], None, @@ -261,7 +259,6 @@ def twist_doughnut( p_tilde_s = p_tilde[t : t + length] n_s = n_guess chi_n_s = p.chi_n[-length:] - e_s = p.e[-length:, j] rho_s = np.diag(p.rho[t : t + p.S, :], max(p.S - length, 0)) error1 = household.FOC_savings( @@ -276,7 +273,6 @@ def twist_doughnut( tr, ubi, theta, - e_s, rho_s, etr_params, mtry_params, @@ -299,7 +295,6 @@ def twist_doughnut( ubi, theta, chi_n_s, - e_s, etr_params, mtrx_params, t, @@ -823,6 +818,7 @@ def run_TPI(p, client=None): bmat_s[: p.T, :, :], n_mat[: p.T, :, :], p, + "TPI", ) L[: p.T] = aggr.get_L(n_mat[: p.T], p, "TPI") @@ -908,6 +904,7 @@ def run_TPI(p, client=None): ubi[: p.T, :, :], theta, etr_params_4D, + p.e, p, None, "TPI", @@ -1014,6 +1011,7 @@ def run_TPI(p, client=None): ubi[: p.T, :, :], theta, etr_params_4D, + p.e, p, None, "TPI", @@ -1143,7 +1141,7 @@ def run_TPI(p, client=None): ), (1, p.S, 1), ) - e_3D = np.tile(p.e.reshape(1, p.S, p.J), (p.T, 1, 1)) + e_3D = p.e mtry_path = tax.MTR_income( r_p_path[: p.T], wpath[: p.T], diff --git a/ogcore/aggregates.py b/ogcore/aggregates.py index 9aaa0664b..5a7c9a9a4 100644 --- a/ogcore/aggregates.py +++ b/ogcore/aggregates.py @@ -33,7 +33,11 @@ def get_L(n, p, method): """ if method == "SS": - L_presum = p.e * np.transpose(p.omega_SS * p.lambdas) * n + L_presum = ( + np.squeeze(p.e[-1, :, :]) + * np.transpose(p.omega_SS * p.lambdas) + * n + ) L = L_presum.sum() elif method == "TPI": L_presum = (n * (p.e * np.squeeze(p.lambdas))) * np.tile( @@ -284,6 +288,7 @@ def revenue( ubi, theta, etr_params, + e, p, m, method, @@ -317,6 +322,7 @@ def revenue( lifetime income group etr_params (list): list of parameters of the effective tax rate functions + e (Numpy array): effective labor units p (OG-Core Specifications object): model parameters method (str): adjusts calculation dimensions based on 'SS' or 'TPI' @@ -339,10 +345,10 @@ def revenue( """ inc_pay_tax_liab = tax.income_tax_liab( - r, w, b, n, factor, 0, None, method, p.e, etr_params, p + r, w, b, n, factor, 0, None, method, e, etr_params, p ) pension_benefits = tax.pension_amount( - w, n, theta, 0, None, False, method, p.e, p + w, n, theta, 0, None, False, method, e, p ) bq_tax_liab = tax.bequest_tax_liab(r, b, bq, 0, None, method, p) w_tax_liab = tax.wealth_tax_liab(r, b, 0, None, method, p) diff --git a/ogcore/household.py b/ogcore/household.py index c323cbcc3..11f5ae7cf 100644 --- a/ogcore/household.py +++ b/ogcore/household.py @@ -324,7 +324,6 @@ def FOC_savings( tr, ubi, theta, - e, rho, etr_params, mtry_params, @@ -359,7 +358,6 @@ def FOC_savings( ubi (Numpy array): universal basic income payment theta (Numpy array): social security replacement rate for each lifetime income group - e (Numpy array): effective labor units rho (Numpy array): mortality rates etr_params (list): parameters of the effective tax rate functions @@ -380,25 +378,46 @@ def FOC_savings( beta = p.beta[j] if method == "SS": tax_noncompliance = p.capital_income_tax_noncompliance_rate[-1, j] + e = np.squeeze(p.e[-1, :, j]) elif method == "TPI_scalar": tax_noncompliance = p.capital_income_tax_noncompliance_rate[0, j] + e = np.squeeze(p.e[0, :, j]) else: length = r.shape[0] tax_noncompliance = p.capital_income_tax_noncompliance_rate[ t : t + length, j ] + e_long = np.concatenate( + ( + p.e, + np.tile(p.e[-1, :, :].reshape(1, p.S, p.J), (p.S, 1, 1)), + ), + axis=0, + ) + e = np.diag(e_long[t : t + p.S, :, j], max(p.S - length, 0)) else: chi_b = p.chi_b beta = p.beta if method == "SS": tax_noncompliance = p.capital_income_tax_noncompliance_rate[-1, :] + e = np.squeeze(p.e[-1, :, :]) elif method == "TPI_scalar": tax_noncompliance = p.capital_income_tax_noncompliance_rate[0, :] + e = np.squeeze(p.e[0, :, :]) else: length = r.shape[0] tax_noncompliance = p.capital_income_tax_noncompliance_rate[ t : t + length, : ] + e_long = np.concatenate( + ( + p.e, + np.tile(p.e[-1, :, :].reshape(1, p.S, p.J), (p.S, 1, 1)), + ), + axis=0, + ) + e = np.diag(e_long[t : t + p.S, :, :], max(p.S - length, 0)) + e = np.squeeze(e) if method == "SS": h_wealth = p.h_wealth[-1] m_wealth = p.m_wealth[-1] @@ -412,7 +431,6 @@ def FOC_savings( h_wealth = p.h_wealth[t] m_wealth = p.m_wealth[t] p_wealth = p.p_wealth[t] - taxes = tax.net_taxes( r, w, @@ -493,7 +511,6 @@ def FOC_labor( ubi, theta, chi_n, - e, etr_params, mtrx_params, t, @@ -554,21 +571,41 @@ def FOC_labor( if j is not None: if method == "SS": tax_noncompliance = p.labor_income_tax_noncompliance_rate[-1, j] + e = np.squeeze(p.e[-1, :, j]) elif method == "TPI_scalar": tax_noncompliance = p.labor_income_tax_noncompliance_rate[0, j] + e = np.squeeze(p.e[0, -1, j]) else: tax_noncompliance = p.labor_income_tax_noncompliance_rate[ t : t + length, j ] + e_long = np.concatenate( + ( + p.e, + np.tile(p.e[-1, :, :].reshape(1, p.S, p.J), (p.S, 1, 1)), + ), + axis=0, + ) + e = np.diag(e_long[t : t + p.S, :, j], max(p.S - length, 0)) else: if method == "SS": tax_noncompliance = p.labor_income_tax_noncompliance_rate[-1, :] + e = np.squeeze(p.e[-1, :, :]) elif method == "TPI_scalar": tax_noncompliance = p.labor_income_tax_noncompliance_rate[0, :] + e = np.squeeze(p.e[0, -1, :]) else: tax_noncompliance = p.labor_income_tax_noncompliance_rate[ t : t + length, : ] + e_long = np.concatenate( + ( + p.e, + np.tile(p.e[-1, :, :].reshape(1, p.S, p.J), (p.S, 1, 1)), + ), + axis=0, + ) + e = np.diag(e_long[t : t + p.S, :, j], max(p.S - length, 0)) if method == "SS": tau_payroll = p.tau_payroll[-1] elif method == "TPI_scalar": # for 1st donut ring only @@ -625,7 +662,7 @@ def FOC_labor( return FOC_error -def get_y(r_p, w, b_s, n, p): +def get_y(r_p, w, b_s, n, p, method): r""" Compute household income before taxes. @@ -638,9 +675,13 @@ def get_y(r_p, w, b_s, n, p): b_s (Numpy array): household savings coming into the period n (Numpy array): household labor supply p (OG-Core Specifications object): model parameters + method (str): adjusts calculation dimensions based on 'SS' or 'TPI' """ - - y = r_p * b_s + w * p.e * n + if method == "SS": + e = np.squeeze(p.e[-1, :, :]) + elif method == "TPI": + e = p.e + y = r_p * b_s + w * e * n return y diff --git a/ogcore/parameter_plots.py b/ogcore/parameter_plots.py index f9ef3b3c5..a5c79277e 100644 --- a/ogcore/parameter_plots.py +++ b/ogcore/parameter_plots.py @@ -153,24 +153,27 @@ def plot_population(p, years_to_plot=["SS"], include_title=False, path=None): plt.savefig(fig_path, dpi=300) -def plot_ability_profiles(p, include_title=False, path=None): +def plot_ability_profiles(p, t=None, 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 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): - plt.plot(age_vec, p.e[:, j], label=GROUP_LABELS[p.J][j]) + plt.plot(age_vec, p.e[t, :, j], 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) @@ -863,7 +866,7 @@ def txfunc_sse_plot(age_vec, sse_mat, start_year, varstr, output_dir, round): def plot_income_data( - ages, abil_midp, abil_pcts, emat, output_dir=None, filesuffix="" + ages, abil_midp, abil_pcts, emat, t=None, output_dir=None, filesuffix="" ): """ This function graphs ability matrix in 3D, 2D, log, and nolog @@ -875,13 +878,16 @@ def plot_income_data( 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 SxJ + 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.cm.get_cmap("summer") @@ -891,7 +897,7 @@ def plot_income_data( if J == 1: # Plot of 2D, J=1 in levels plt.figure() - plt.plot(ages, emat) + plt.plot(ages, emat[t, :, :]) filename = "ability_2D_lev" + filesuffix fullpath = os.path.join(output_dir, filename) plt.savefig(fullpath, dpi=300) @@ -899,7 +905,7 @@ def plot_income_data( # Plot of 2D, J=1 in logs plt.figure() - plt.plot(ages, np.log(emat)) + plt.plot(ages, np.log(emat[t, :, :])) filename = "ability_2D_log" + filesuffix fullpath = os.path.join(output_dir, filename) plt.savefig(fullpath, dpi=300) @@ -908,7 +914,12 @@ def plot_income_data( # Plot of 3D, J>1 in levels fig10, ax10 = plt.subplots(subplot_kw={"projection": "3d"}) ax10.plot_surface( - age_mesh, abil_mesh, emat, rstride=8, cstride=1, cmap=cmap1 + 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$") @@ -923,7 +934,7 @@ def plot_income_data( ax11.plot_surface( age_mesh, abil_mesh, - np.log(emat), + np.log(emat[t, :, :]), rstride=8, cstride=1, cmap=cmap1, @@ -960,7 +971,7 @@ def plot_income_data( if j <= 3: ax.plot( ages, - np.log(emat[:, j]), + np.log(emat[t, :, j]), label=this_label, linestyle=linestyles[j], color="black", @@ -968,7 +979,7 @@ def plot_income_data( elif j > 3: ax.plot( ages, - np.log(emat[:, j]), + np.log(emat[t, :, j]), label=this_label, marker=markers[j - 4], color="black", @@ -1008,7 +1019,7 @@ def plot_income_data( if j <= 3: ax.plot( ages, - np.log(emat[:, j]), + np.log(emat[t, :, j]), label=this_label, linestyle=linestyles[j], color="black", @@ -1016,7 +1027,7 @@ def plot_income_data( elif j > 3: ax.plot( ages, - np.log(emat[:, j]), + np.log(emat[t, :, j]), label=this_label, marker=markers[j - 4], color="black", diff --git a/ogcore/parameters.py b/ogcore/parameters.py index 360e53aee..51999981f 100644 --- a/ogcore/parameters.py +++ b/ogcore/parameters.py @@ -257,6 +257,11 @@ def compute_default_params(self): param_in, dims=(self.T + self.S, self.S, self.J), item="eta" ) setattr(self, "eta", param_out) + param_in = getattr(self, "e") + param_out = extrapolate_arrays( + param_in, dims=(self.T, self.S, self.J), item="e" + ) + setattr(self, "e", param_out) # make sure zeta matrix sums to one (e.g., default off due to rounding) self.zeta = self.zeta / self.zeta.sum() diff --git a/ogcore/tax.py b/ogcore/tax.py index 1f0c0ed08..eb988ed60 100644 --- a/ogcore/tax.py +++ b/ogcore/tax.py @@ -34,9 +34,9 @@ def replacement_rate_vals(nssmat, wss, factor_ss, j, p): """ if j is not None: - e = p.e[:, j] + e = np.squeeze(p.e[-1, :, j]) # Only computes using SS earnings else: - e = p.e + e = np.squeeze(p.e[-1, :, :]) # Only computes using SS earnings # adjust number of calendar years AIME computed from int model periods equiv_periods = int(round((p.S / 80.0) * p.AIME_num_years)) - 1 if e.ndim == 2: @@ -405,7 +405,6 @@ def income_tax_liab(r, w, b, n, factor, t, j, method, e, etr_params, p): capital_income_tax_compliance_rate = ( p.capital_income_tax_noncompliance_rate[-1, :] ) - income = r * b + w * e * n labor_income = w * e * n T_I = ( diff --git a/ogcore/utils.py b/ogcore/utils.py index 1a12187ab..2fe130ec1 100644 --- a/ogcore/utils.py +++ b/ogcore/utils.py @@ -995,19 +995,9 @@ def extrapolate_arrays(param_in, dims=None, item="Parameter Name"): elif param_in.ndim == 2: # case if S by J input if param_in.shape[0] == dims[1]: - param_in = np.tile( + param_out = np.tile( param_in.reshape(1, dims[1], dims[2]), - (dims[0] - dims[1], 1, 1), - ) - param_out = np.concatenate( - ( - param_in, - np.tile( - param_in[-1, :, :].reshape(1, dims[1], dims[2]), - (dims[1], 1, 1), - ), - ), - axis=0, + (dims[0], 1, 1), ) # case if T by S input elif param_in.shape[0] == dims[0] - dims[1]: @@ -1034,15 +1024,18 @@ def extrapolate_arrays(param_in, dims=None, item="Parameter Name"): assert False elif param_in.ndim == 3: # this is the case where input varies by T, S, J - param_out = np.concatenate( - ( - param_in, - np.tile( - param_in[-1, :, :].reshape(1, dims[1], dims[2]), - (dims[1], 1, 1), + if param_in.shape[0] > dims[0]: + param_out = param_in[: dims[0], :, :] + else: + param_out = np.concatenate( + ( + param_in, + np.tile( + param_in[-1, :, :].reshape(1, dims[1], dims[2]), + (dims[0] - param_in.shape[0], 1, 1), + ), ), - ), - axis=0, - ) + axis=0, + ) return param_out diff --git a/tests/test_SS.py b/tests/test_SS.py index 215f57392..652b92601 100644 --- a/tests/test_SS.py +++ b/tests/test_SS.py @@ -584,6 +584,9 @@ def test_inner_loop(baseline, r_p, param_updates, filename, dask_client): for i, v in enumerate(expected_tuple): print("Max diff = ", np.absolute(test_tuple[i] - v).max()) print("Checking item = ", i) + if np.absolute(test_tuple[i] - v).max() > 1.0: + print("test_value = ", test_tuple[i]) + print("expected_value = ", v) for i, v in enumerate(expected_tuple): print("Max diff = ", np.absolute(test_tuple[i] - v).max()) diff --git a/tests/test_TPI.py b/tests/test_TPI.py index 180c5bce9..2e9039b3f 100644 --- a/tests/test_TPI.py +++ b/tests/test_TPI.py @@ -184,7 +184,6 @@ def test_twist_doughnut(file_inputs, file_outputs): initial_b, p, ) - print("t = ", t, s, " p.S = ", p.S, p.rho.shape) test_list = TPI.twist_doughnut(*input_tuple) expected_list = utils.safe_read_pickle(file_outputs) assert np.allclose(np.array(test_list), np.array(expected_list), atol=1e-5) diff --git a/tests/test_aggregates.py b/tests/test_aggregates.py index b80f0549c..1142c712c 100644 --- a/tests/test_aggregates.py +++ b/tests/test_aggregates.py @@ -29,7 +29,7 @@ for i in range(p.S): for k in range(p.J): L_loop[t, i, k] *= ( - p.omega[t, i] * p.lambdas[k] * n[t, i, k] * p.e[i, k] + p.omega[t, i] * p.lambdas[k] * n[t, i, k] * p.e[t, i, k] ) expected1 = L_loop[-1, :, :].sum() expected2 = L_loop.sum(1).sum(1) @@ -53,6 +53,7 @@ def test_get_L(n, p, method, expected): "S": 40, "rho": rho_vec.tolist(), "J": 2, + "e": np.ones((40, 2)), "labor_income_tax_noncompliance_rate": [[0.0]], "capital_income_tax_noncompliance_rate": [[0.0]], "eta": (np.ones((40, 2)) / (40 * 2)), @@ -133,6 +134,7 @@ def test_get_I(b_splus1, K_p1, K, p, method, expected): "S": 40, "rho": rho_vec.tolist(), "J": 2, + "e": np.ones((40, 2)), "labor_income_tax_noncompliance_rate": [[0.0]], "capital_income_tax_noncompliance_rate": [[0.0]], "eta": (np.ones((40, 2)) / (40 * 2)), @@ -189,6 +191,7 @@ def test_get_B(b, p, method, PreTP, expected): "S": 40, "rho": rho_vec.tolist(), "J": 2, + "e": np.ones((40, 2)), "labor_income_tax_noncompliance_rate": [[0.0]], "capital_income_tax_noncompliance_rate": [[0.0]], "eta": (np.ones((40, 2)) / (40 * 2)), @@ -270,6 +273,7 @@ def test_get_BQ(r, b_splus1, j, p, method, PreTP, expected): "S": 40, "rho": rho_vec.tolist(), "J": 2, + "e": np.ones((40, 2)), "M": 3, "labor_income_tax_noncompliance_rate": [[0.0]], "capital_income_tax_noncompliance_rate": [[0.0]], @@ -319,6 +323,7 @@ def test_get_C(c, p, method, expected): "S": 20, "rho": rho_vec.tolist(), "J": 2, + "e": np.ones((20, 2)), "labor_income_tax_noncompliance_rate": [[0.0]], "capital_income_tax_noncompliance_rate": [[0.0]], "eta": (np.ones((20, 2)) / (20 * 2)), @@ -358,6 +363,7 @@ def test_get_C(c, p, method, expected): factor = 140000.0 # update parameters instance with new values for test p.e = 0.263 + (2.024 - 0.263) * random_state.rand(p.S * p.J).reshape(p.S, p.J) +p.e = np.tile(p.e.reshape(1, p.S, p.J), (p.T, 1, 1)) p.omega = 0.039 * random_state.rand(p.T * p.S * 1).reshape(p.T, p.S) p.omega = p.omega / p.omega.sum(axis=1).reshape(p.T, 1) p.omega_SS = p.omega[-1, :] @@ -374,6 +380,7 @@ def test_get_C(c, p, method, expected): "S": 20, "rho": rho_vec.tolist(), "J": 2, + "e": np.ones((20, 2)), "labor_income_tax_noncompliance_rate": [[0.0]], "capital_income_tax_noncompliance_rate": [[0.0]], "eta": (np.ones((20, 2)) / (20 * 2)), @@ -408,6 +415,7 @@ def test_get_C(c, p, method, expected): "S": 20, "rho": rho_vec.tolist(), "J": 2, + "e": np.ones((20, 2)), "labor_income_tax_noncompliance_rate": [[0.0]], "capital_income_tax_noncompliance_rate": [[0.0]], "eta": (np.ones((20, 2)) / (20 * 2)), @@ -453,9 +461,10 @@ def test_get_C(c, p, method, expected): factor_u = 140000.0 ubi_u = p_u.ubi_nom_array / factor_u # update parameters instance with new values for test -p_u.e = 0.263 + (2.024 - 0.263) * random_state.rand(p_u.S * p_u.J).reshape( - p_u.S, p_u.J +p_u.e = 0.263 + (2.024 - 0.263) * random_state.rand(p.S * p.J).reshape( + p.S, p.J ) +p_u.e = np.tile(p_u.e.reshape(1, p_u.S, p_u.J), (p_u.T, 1, 1)) p_u.omega = 0.039 * random_state.rand(p_u.T * p_u.S * 1).reshape(p_u.T, p_u.S) p_u.omega = p_u.omega / p_u.omega.sum(axis=1).reshape(p_u.T, 1) p_u.omega_SS = p_u.omega[-1, :] @@ -589,6 +598,7 @@ def test_get_C(c, p, method, expected): ubi[0, :, :], theta, etr_params[-1, :, :, :], + p.e[-1, :, :], p, None, "SS", @@ -609,6 +619,7 @@ def test_get_C(c, p, method, expected): ubi, theta, etr_params, + p.e, p, None, "TPI", @@ -629,6 +640,7 @@ def test_get_C(c, p, method, expected): ubi, theta, etr_params, + p.e, p3, None, "TPI", @@ -649,6 +661,7 @@ def test_get_C(c, p, method, expected): ubi_u[0, :, :], theta_u, etr_params_u[-1, :, :, :], + p.e[0, :, :], p_u, None, "SS", @@ -669,6 +682,7 @@ def test_get_C(c, p, method, expected): ubi_u, theta_u, etr_params_u, + p.e, p_u, None, "TPI", @@ -678,7 +692,7 @@ def test_get_C(c, p, method, expected): @pytest.mark.parametrize( - "r,w,b,n,bq,c,Y,L,K,p_m,factor,ubi,theta,etr_params,p,m,method,expected", + "r,w,b,n,bq,c,Y,L,K,p_m,factor,ubi,theta,etr_params,e,p,m,method,expected", test_data, ids=["SS", "TPI", "TPI, replace rate adjust", "SS UBI>0", "TPI UBI>0"], ) @@ -697,6 +711,7 @@ def test_revenue( ubi, theta, etr_params, + e, p, m, method, @@ -720,6 +735,7 @@ def test_revenue( ubi, theta, etr_params, + e, p, m, method, diff --git a/tests/test_firm.py b/tests/test_firm.py index 3bcbc8553..d97977657 100644 --- a/tests/test_firm.py +++ b/tests/test_firm.py @@ -35,6 +35,7 @@ "epsilon": [1.0], "T": 3, "S": 3, + "e": np.ones((3, p4.J)), "rho": rho_vec.tolist(), "eta": (np.ones((3, p4.J)) / (3 * p4.J)), } @@ -53,6 +54,7 @@ "S": 3, "rho": rho_vec.tolist(), "eta": (np.ones((3, p5.J)) / (3 * p5.J)), + "e": np.ones((3, p5.J)), } # update parameters instance with new values for test p5.update_specifications(new_param_values5) @@ -64,6 +66,7 @@ "epsilon": [1.0], "T": 3, "S": 3, + "e": np.ones((3, p5.J)), "rho": rho_vec.tolist(), "eta": (np.ones((3, p5.J)) / (3 * p5.J)), "gamma_g": [0.2], @@ -84,6 +87,7 @@ "T": 3, "S": 3, "M": 2, + "e": np.ones((3, p4.J)), "rho": rho_vec.tolist(), "eta": (np.ones((3, p4.J)) / (3 * p4.J)), "initial_Kg_ratio": 0.01, @@ -105,6 +109,7 @@ "T": 3, "S": 3, "M": 2, + "e": np.ones((3, p4.J)), "rho": rho_vec.tolist(), "eta": (np.ones((3, p4.J)) / (3 * p4.J)), "initial_Kg_ratio": 0.01, @@ -122,6 +127,7 @@ "T": 3, "S": 3, "M": 3, + "e": np.ones((3, p5.J)), "rho": rho_vec.tolist(), "eta": (np.ones((3, p5.J)) / (3 * p5.J)), "gamma_g": [0.2, 0.1, 0.25], @@ -147,6 +153,7 @@ "T": 3, "S": 3, "M": 3, + "e": np.ones((3, p5.J)), "rho": rho_vec.tolist(), "eta": (np.ones((3, p5.J)) / (3 * p5.J)), "gamma_g": [0.2, 0.1, 0.25], @@ -267,6 +274,7 @@ def test_get_Y(K, K_g, L, p, method, m, expected): "delta_annual": 0.5, "T": 3, "S": 3, + "e": np.ones((3, p4.J)), "rho": rho_vec.tolist(), "eta": (np.ones((3, p4.J)) / (3 * p4.J)), } @@ -287,6 +295,7 @@ def test_get_Y(K, K_g, L, p, method, m, expected): "inv_tax_credit": [[0.07]], "T": 3, "S": 3, + "e": np.ones((3, p4.J)), "rho": rho_vec.tolist(), "eta": (np.ones((3, p5.J)) / (3 * p5.J)), } @@ -361,6 +370,7 @@ def test_get_r(Y, K, p_m, p, method, expected): "epsilon": [1.2], "T": 3, "S": 3, + "e": np.ones((3, p4.J)), "rho": rho_vec.tolist(), "eta": (np.ones((3, p4.J)) / (3 * p4.J)), } @@ -450,6 +460,7 @@ def test_get_w(Y, L, p_m, p, method, expected): "cit_rate": [[(0.0357 / 0.55) * (0.055 / 0.017)]], "T": 3, "S": 3, + "e": np.ones((3, p4.J)), "rho": rho_vec.tolist(), "eta": (np.ones((3, p4.J)) / (3 * p4.J)), } @@ -558,6 +569,7 @@ def test_get_KLratio(r, p, method, expected): "cit_rate": [[(0.0357 / 0.55) * (0.055 / 0.017)]], "T": 3, "S": 3, + "e": np.ones((3, p4.J)), "rho": rho_vec.tolist(), "eta": (np.ones((3, p4.J)) / (3 * p4.J)), } @@ -639,6 +651,7 @@ def test_get_w_from_r(r, p, method, expected): "cit_rate": [[0.5]], "T": 3, "S": 3, + "e": np.ones((3, p4.J)), "rho": rho_vec.tolist(), "eta": (np.ones((3, p4.J)) / (3 * p4.J)), } @@ -719,6 +732,7 @@ def test_get_K(L, r, p, method, expected): "epsilon": [1.0], "Z": [[2.0]], "T": 3, + "e": p2.e[0, :, :], } # update parameters instance with new values for test p2.update_specifications(new_param_values2) @@ -764,6 +778,7 @@ def test_get_MPx(Y, x, share, p, method, expected): "adjustment_factor_for_cit_receipts": [1.0], "c_corp_share_of_assets": 1.0, "T": 3, + "e": p1.e[0, :, :], } # update parameters instance with new values for test p1.update_specifications(new_param_values1) @@ -778,6 +793,7 @@ def test_get_MPx(Y, x, share, p, method, expected): "adjustment_factor_for_cit_receipts": [1.0], "c_corp_share_of_assets": 1.0, "T": 3, + "e": p2.e[0, :, :], } # update parameters instance with new values for test p2.update_specifications(new_param_values2) @@ -794,6 +810,7 @@ def test_get_MPx(Y, x, share, p, method, expected): "c_corp_share_of_assets": 1.0, "T": 3, "M": 2, + "e": p3.e[0, :, :], } # update parameters instance with new values for test p3.update_specifications(new_param_values3) @@ -903,6 +920,7 @@ def test_get_pm(w, Y, L, p, method, expected): "c_corp_share_of_assets": 1.0, "initial_Kg_ratio": 0.01, "T": 3, + "e": p5.e[0, :, :], } # update parameters instance with new values for test p5.update_specifications(new_param_values5) @@ -920,6 +938,7 @@ def test_get_pm(w, Y, L, p, method, expected): "c_corp_share_of_assets": 1.0, "initial_Kg_ratio": 0.01, "T": 3, + "e": p6.e[0, :, :], } # update parameters instance with new values for test p6.update_specifications(new_param_values6) diff --git a/tests/test_household.py b/tests/test_household.py index 0664f5025..e7a26eb48 100644 --- a/tests/test_household.py +++ b/tests/test_household.py @@ -618,7 +618,7 @@ def test_get_cons(model_args, expected): @pytest.mark.parametrize( - "model_vars,params,expected", + "model_vars,in_params,expected", test_data, ids=[ "SS", @@ -631,7 +631,7 @@ def test_get_cons(model_args, expected): "TPI, j=0, noncomply", ], ) -def test_FOC_savings(model_vars, params, expected): +def test_FOC_savings(model_vars, in_params, expected): # Test FOC condition for household's choice of savings ( r, @@ -650,6 +650,11 @@ def test_FOC_savings(model_vars, params, expected): j, method, ) = model_vars + params = copy.deepcopy(in_params) + # reshape e matrix to be 3D + params.e = np.tile( + params.e.reshape(1, params.S, params.J), (params.T, 1, 1) + ) if method == "TPI": p_tilde = np.ones_like(w) else: @@ -667,7 +672,6 @@ def test_FOC_savings(model_vars, params, expected): tr, ubi, theta, - params.e[:, j], params.rho[-1, :], etr_params, mtry_params, @@ -689,7 +693,6 @@ def test_FOC_savings(model_vars, params, expected): tr, ubi, theta, - np.squeeze(params.e), params.rho[-1, :], etr_params, mtry_params, @@ -933,6 +936,10 @@ def test_FOC_labor(model_vars, params, expected): j, method, ) = model_vars + # reshape e matrix for 3D + params.e = np.tile( + params.e.reshape(1, params.S, params.J), (params.T, 1, 1) + ) if method == "TPI": p_tilde = np.ones_like(w) else: @@ -950,7 +957,6 @@ def test_FOC_labor(model_vars, params, expected): ubi, theta, params.chi_n, - params.e[:, j], etr_params, mtrx_params, t, @@ -976,8 +982,10 @@ def test_get_y(): p.S = 3 p.rho = np.array([[0.0, 0.0, 1.0]]) p.e = np.array([0.99, 1.5, 0.2]) + p.e = np.tile(p.e.reshape(1, p.S, 1), (p.T, 1, 1)) - test_y = household.get_y(r_p, w, b_s, n, p) + test_y = household.get_y(r_p, w, b_s, n, p, "SS") + # TODO: test with "TPI" assert np.allclose(test_y, expected_y) diff --git a/tests/test_parameter_plots.py b/tests/test_parameter_plots.py index 90bd17773..3f9defbb8 100644 --- a/tests/test_parameter_plots.py +++ b/tests/test_parameter_plots.py @@ -82,6 +82,11 @@ 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), + ) fig = parameter_plots.plot_ability_profiles( base_params, include_title=True ) diff --git a/tests/test_parameters.py b/tests/test_parameters.py index 835968e07..6da427db6 100644 --- a/tests/test_parameters.py +++ b/tests/test_parameters.py @@ -41,19 +41,19 @@ def test_compute_default_params(): "T": 4, "S": 3, "rho": [[0.0, 0.0, 1.0]], - "J": 1, + "e": np.ones((3, 7)), "ubi_nom_017": 1000, "eta": np.ones((4, 3, 1)) / 12, "ubi_nom_1864": 1200, "ubi_nom_65p": 400, "ubi_growthadj": True, } -expected1 = np.ones((7, 3, 1)) * 2180 +expected1 = np.ones((7, 3, 7)) * 2180 param_updates2 = { "T": 4, "S": 3, "rho": [[0.0, 0.0, 1.0]], - "J": 1, + "e": np.ones((3, 7)), "ubi_nom_017": 1000, "eta": np.ones((4, 3, 1)) / 12, "ubi_nom_1864": 1200, @@ -61,12 +61,12 @@ def test_compute_default_params(): "ubi_nom_65p": 400, "ubi_growthadj": True, } -expected2 = np.ones((7, 3, 1)) * 2000 +expected2 = np.ones((7, 3, 7)) * 2000 param_updates3 = { "T": 4, "S": 3, "rho": [[0.0, 0.0, 1.0]], - "J": 1, + "e": np.ones((3, 7)), "ubi_nom_017": 1000, "eta": np.ones((4, 3, 1)) / 12, "ubi_nom_1864": 1200, diff --git a/tests/test_tax.py b/tests/test_tax.py index ddc622ded..0ba7d9984 100644 --- a/tests/test_tax.py +++ b/tests/test_tax.py @@ -17,6 +17,7 @@ "J": 1, "T": 4, "eta": (np.ones((4, 1)) / (4 * 1)), + "e": np.ones((4, 1)), } p.update_specifications(new_param_values) p.retire = [3, 3, 3, 3, 3, 3, 3, 3] @@ -53,14 +54,19 @@ @pytest.mark.parametrize( - "n,w,factor,j,p,expected", + "n,w,factor,j,p_in,expected", test_data, ids=["1D e", "2D e", "AIME case 2", "AIME case 3", "Min PIA case"], ) -def test_replacement_rate_vals(n, w, factor, j, p, expected): +def test_replacement_rate_vals(n, w, factor, j, p_in, expected): # Test replacement rate function, making sure to trigger all three # cases of AIME - + # make e 3D + p = copy.deepcopy(p_in) + # p.e = np.tile(np.reshape(p.e, (1, p.S, p.J)), (p.T, 1, 1)) + p.e = np.tile( + np.reshape(p.e, (1, p.e.shape[0], p.e.shape[1])), (p.T, 1, 1) + ) theta = tax.replacement_rate_vals(n, w, factor, j, p) assert np.allclose(theta, expected) @@ -75,6 +81,7 @@ def test_replacement_rate_vals(n, w, factor, j, p, expected): "lambdas": [1.0], "J": 1, "T": 3, + "e": np.ones((3, 1)), "eta": (np.ones((3, 1)) / (3 * 1)), "h_wealth": [2], "p_wealth": [3], @@ -91,6 +98,7 @@ def test_replacement_rate_vals(n, w, factor, j, p, expected): "lambdas": [1.0], "J": 1, "T": 3, + "e": np.ones((3, 1)), "eta": (np.ones((3, 1)) / (3 * 1)), "h_wealth": [1.2, 1.1, 2.3], "p_wealth": [2.2, 2.3, 1.8], @@ -124,6 +132,7 @@ def test_ETR_wealth(b, p, expected): "lambdas": [1.0], "J": 1, "T": 3, + "e": np.ones((3, 1)), "eta": (np.ones((3, 1)) / (3 * 1)), "h_wealth": [3], "p_wealth": [4], @@ -141,6 +150,7 @@ def test_ETR_wealth(b, p, expected): "lambdas": [1.0], "J": 1, "T": 3, + "e": np.ones((3, 1)), "eta": (np.ones((3, 1)) / (3 * 1)), "h_wealth": [1.2, 1.1, 2.3], "p_wealth": [2.2, 2.3, 1.8], @@ -875,6 +885,7 @@ def test_MTR_income(etr_params, mtr_params, params, mtr_capital, expected): "inv_tax_credit": [[0.02]], "T": 3, "S": 3, + "e": np.ones((3, p1.J)), "rho": rho_vec.tolist(), "eta": (np.ones((3, p1.J)) / (3 * p1.J)), "labor_income_tax_noncompliance_rate": [[0.0]], @@ -1203,6 +1214,7 @@ def test_get_biz_tax(w, Y, L, K, p_m, p, m, method, expected): "T": 3, "S": 3, "J": 2, + "e": np.ones((3, 2)), "rho": rho_vec.tolist(), "lambdas": [0.65, 0.35], "eta": (np.ones((3, 2)) / (3 * 2)),