From 25492a0c846a59aa433c61f7e0a48fd544ca6d57 Mon Sep 17 00:00:00 2001 From: jdebacker Date: Thu, 7 Sep 2023 12:15:08 -0400 Subject: [PATCH] use more general slicing --- ogcore/TPI.py | 108 +++++++++++++++++++++++++------------------------- 1 file changed, 54 insertions(+), 54 deletions(-) diff --git a/ogcore/TPI.py b/ogcore/TPI.py index 47583f7e9..661e4e454 100644 --- a/ogcore/TPI.py +++ b/ogcore/TPI.py @@ -1242,61 +1242,61 @@ def run_TPI(p, client=None): """ output = { - "Y": Y[: p.T], - "B": B[: p.T], - "K": K[: p.T], + "Y": Y[: p.T, ...], + "B": B[: p.T, ...], + "K": K[: p.T, ...], "K_f": K_f[: p.T], - "K_d": K_d[: p.T], - "L": L[: p.T], - "C": C[: p.T], - "I": I[: p.T], - "K_g": K_g[: p.T], - "I_g": I_g[: p.T], - "Y_vec": Y_vec[: p.T, :], - "K_vec": K_vec[: p.T, :], - "L_vec": L_vec[: p.T, :], - "C_vec": C_vec[: p.T, :], - "I_total": I_total[: p.T], - "I_d": I_d[: p.T], - "BQ": BQ[: p.T], - "total_tax_revenue": total_tax_revenue[: p.T], - "business_tax_revenue": business_tax_revenue[: p.T], - "iit_payroll_tax_revenue": iit_payroll_tax_revenue[: p.T], - "iit_revenue": iit_revenue[: p.T], - "payroll_tax_revenue": payroll_tax_revenue[: p.T], - "TR": TR[: p.T], - "agg_pension_outlays": agg_pension_outlays[: p.T], - "bequest_tax_revenue": bequest_tax_revenue[: p.T], - "wealth_tax_revenue": wealth_tax_revenue[: p.T], - "cons_tax_revenue": cons_tax_revenue[: p.T], - "G": G[: p.T], - "D": D[: p.T], - "D_f": D_f[: p.T], - "D_d": D_d[: p.T], - "r": r[: p.T], - "r_gov": r_gov[: p.T], - "r_p": r_p[: p.T], - "w": w[: p.T], - "bmat_splus1": bmat_splus1[: p.T, :, :], - "p_m": p_m[: p.T, :], - "p_tilde": p_tilde, - "bmat_s": bmat_s[: p.T, :, :], - "n_mat": n_mat[: p.T, :, :], - "c_path": c_mat[: p.T, :, :, :], - "bq_path": bqmat[: p.T, :, :], - "tr_path": trmat[: p.T, :, :], - "y_before_tax_mat": y_before_tax_mat[: p.T, :, :], - "tax_path": tax_mat[: p.T, :, :], - "eul_savings": eul_savings[: p.T, :, :], - "eul_laborleisure": eul_laborleisure[: p.T, :, :], - "resource_constraint_error": RC_error[: p.T], - "new_borrowing_f": new_borrowing_f[: p.T], - "debt_service_f": debt_service_f[: p.T], - "etr_path": etr_path[: p.T, :, :], - "mtrx_path": mtrx_path[: p.T, :, :], - "mtry_path": mtry_path[: p.T, :, :], - "ubi_path": ubi[: p.T, :, :], - "UBI_path": UBI[: p.T], + "K_d": K_d[: p.T, ...], + "L": L[: p.T, ...], + "C": C[: p.T, ...], + "I": I[: p.T, ...], + "K_g": K_g[: p.T, ...], + "I_g": I_g[: p.T, ...], + "Y_vec": Y_vec[: p.T, ...], + "K_vec": K_vec[: p.T, ...], + "L_vec": L_vec[: p.T, ...], + "C_vec": C_vec[: p.T, ...], + "I_total": I_total[: p.T, ...], + "I_d": I_d[: p.T, ...], + "BQ": BQ[: p.T, ...], + "total_tax_revenue": total_tax_revenue[: p.T, ...], + "business_tax_revenue": business_tax_revenue[: p.T, ...], + "iit_payroll_tax_revenue": iit_payroll_tax_revenue[: p.T, ...], + "iit_revenue": iit_revenue[: p.T, ...], + "payroll_tax_revenue": payroll_tax_revenue[: p.T, ...], + "TR": TR[: p.T, ...], + "agg_pension_outlays": agg_pension_outlays[: p.T, ...], + "bequest_tax_revenue": bequest_tax_revenue[: p.T, ...], + "wealth_tax_revenue": wealth_tax_revenue[: p.T, ...], + "cons_tax_revenue": cons_tax_revenue[: p.T, ...], + "G": G[: p.T, ...], + "D": D[: p.T, ...], + "D_f": D_f[: p.T, ...], + "D_d": D_d[: p.T, ...], + "r": r[: p.T, ...], + "r_gov": r_gov[: p.T, ...], + "r_p": r_p[: p.T, ...], + "w": w[: p.T, ...], + "bmat_splus1": bmat_splus1[: p.T, ...], + "p_m": p_m[: p.T, ...], + "p_tilde": p_tilde[: p.T, ...], + "bmat_s": bmat_s[: p.T, ...], + "n_mat": n_mat[: p.T, ...], + "c_path": c_mat[: p.T, ...], + "bq_path": bqmat[: p.T, ...], + "tr_path": trmat[: p.T, ...], + "y_before_tax_mat": y_before_tax_mat[: p.T, ...], + "tax_path": tax_mat[: p.T, ...], + "eul_savings": eul_savings[: p.T, ...], + "eul_laborleisure": eul_laborleisure[: p.T, ...], + "resource_constraint_error": RC_error[: p.T, ...], + "new_borrowing_f": new_borrowing_f[: p.T, ...], + "debt_service_f": debt_service_f[: p.T, ...], + "etr_path": etr_path[: p.T, ...], + "mtrx_path": mtrx_path[: p.T, ...], + "mtry_path": mtry_path[: p.T, ...], + "ubi_path": ubi[: p.T, ...], + "UBI_path": UBI[: p.T, ...], } tpi_dir = os.path.join(p.output_base, "TPI")