use more general slicing

PSLmodels · Sep 7, 2023 · 25492a0 · 25492a0
1 parent e16ac76
commit 25492a0
Showing 1 changed file with 54 additions and 54 deletions.
diff --git 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")