From 6e4ffc8970411f5b0c16d376955bc146c01fe06f Mon Sep 17 00:00:00 2001 From: jdebacker Date: Fri, 19 Jan 2024 14:25:07 -0500 Subject: [PATCH] black format --- ogcore/SS.py | 10 ++++++---- ogcore/TPI.py | 2 +- ogcore/aggregates.py | 8 ++++++-- ogcore/household.py | 32 ++++++++++++++++++++++++++++---- ogcore/parameter_plots.py | 7 ++++++- tests/test_SS.py | 4 ++-- tests/test_aggregates.py | 4 +++- tests/test_firm.py | 12 ++++++------ tests/test_household.py | 8 ++++++-- tests/test_tax.py | 4 +++- 10 files changed, 67 insertions(+), 24 deletions(-) diff --git a/ogcore/SS.py b/ogcore/SS.py index e7e1fa231..e4183a09f 100644 --- a/ogcore/SS.py +++ b/ogcore/SS.py @@ -481,9 +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) + 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, @@ -821,7 +821,9 @@ def SS_solver( np.squeeze(p.e[-1, :, :]), p, ) - yss_before_tax_mat = household.get_y(r_p_ss, wss, bssmat_s, nssmat, p, "SS") + 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 diff --git a/ogcore/TPI.py b/ogcore/TPI.py index 0b5c7b3fe..6797fb8d9 100644 --- a/ogcore/TPI.py +++ b/ogcore/TPI.py @@ -818,7 +818,7 @@ def run_TPI(p, client=None): bmat_s[: p.T, :, :], n_mat[: p.T, :, :], p, - "TPI" + "TPI", ) L[: p.T] = aggr.get_L(n_mat[: p.T], p, "TPI") diff --git a/ogcore/aggregates.py b/ogcore/aggregates.py index 6b022d780..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 = np.squeeze(p.e[-1, :, :]) * 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( @@ -341,7 +345,7 @@ def revenue( """ inc_pay_tax_liab = tax.income_tax_liab( - r, w, b, n, factor, 0, None, method, 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, e, p diff --git a/ogcore/household.py b/ogcore/household.py index 7580ee0af..11f5ae7cf 100644 --- a/ogcore/household.py +++ b/ogcore/household.py @@ -387,7 +387,13 @@ def FOC_savings( 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_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 @@ -403,7 +409,13 @@ def FOC_savings( 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_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": @@ -567,7 +579,13 @@ def FOC_labor( 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_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": @@ -580,7 +598,13 @@ def FOC_labor( 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_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] diff --git a/ogcore/parameter_plots.py b/ogcore/parameter_plots.py index 7fc531d1f..a5c79277e 100644 --- a/ogcore/parameter_plots.py +++ b/ogcore/parameter_plots.py @@ -914,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[t, :, :], 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$") diff --git a/tests/test_SS.py b/tests/test_SS.py index bb577bdfc..652b92601 100644 --- a/tests/test_SS.py +++ b/tests/test_SS.py @@ -585,8 +585,8 @@ def test_inner_loop(baseline, r_p, param_updates, filename, dask_client): 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) + 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_aggregates.py b/tests/test_aggregates.py index f040d4974..1142c712c 100644 --- a/tests/test_aggregates.py +++ b/tests/test_aggregates.py @@ -461,7 +461,9 @@ 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.S * p.J).reshape(p.S, p.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) diff --git a/tests/test_firm.py b/tests/test_firm.py index cd77d8d6a..d97977657 100644 --- a/tests/test_firm.py +++ b/tests/test_firm.py @@ -732,7 +732,7 @@ def test_get_K(L, r, p, method, expected): "epsilon": [1.0], "Z": [[2.0]], "T": 3, - "e": p2.e[0, :, :] + "e": p2.e[0, :, :], } # update parameters instance with new values for test p2.update_specifications(new_param_values2) @@ -778,7 +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, :, :] + "e": p1.e[0, :, :], } # update parameters instance with new values for test p1.update_specifications(new_param_values1) @@ -793,7 +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, :, :] + "e": p2.e[0, :, :], } # update parameters instance with new values for test p2.update_specifications(new_param_values2) @@ -810,7 +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, :, :] + "e": p3.e[0, :, :], } # update parameters instance with new values for test p3.update_specifications(new_param_values3) @@ -920,7 +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, :, :] + "e": p5.e[0, :, :], } # update parameters instance with new values for test p5.update_specifications(new_param_values5) @@ -938,7 +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, :, :] + "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 8a113d60b..e7a26eb48 100644 --- a/tests/test_household.py +++ b/tests/test_household.py @@ -652,7 +652,9 @@ def test_FOC_savings(model_vars, in_params, expected): ) = 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)) + 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: @@ -935,7 +937,9 @@ def test_FOC_labor(model_vars, params, expected): method, ) = model_vars # reshape e matrix for 3D - params.e = np.tile(params.e.reshape(1, params.S, params.J), (params.T, 1, 1)) + 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: diff --git a/tests/test_tax.py b/tests/test_tax.py index 073d3b09d..0ba7d9984 100644 --- a/tests/test_tax.py +++ b/tests/test_tax.py @@ -64,7 +64,9 @@ def test_replacement_rate_vals(n, w, factor, j, p_in, expected): # 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)) + 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)