make all time path otuput over T periods

jdebacker · Sep 6, 2023 · e16ac76 · e16ac76
1 parent 0b48a79
commit e16ac76
Showing 1 changed file with 51 additions and 51 deletions.
diff --git a/ogcore/TPI.py b/ogcore/TPI.py
@@ -1243,60 +1243,60 @@ def run_TPI(p, client=None):
 
     output = {
         "Y": Y[: p.T],
-        "B": B,
-        "K": K,
-        "K_f": K_f,
-        "K_d": K_d,
-        "L": L,
-        "C": C,
-        "I": I,
-        "K_g": K_g,
-        "I_g": I_g,
-        "Y_vec": Y_vec,
-        "K_vec": K_vec,
-        "L_vec": L_vec,
-        "C_vec": C_vec,
-        "I_total": I_total,
-        "I_d": I_d,
-        "BQ": BQ,
-        "total_tax_revenue": total_tax_revenue,
-        "business_tax_revenue": business_tax_revenue,
-        "iit_payroll_tax_revenue": iit_payroll_tax_revenue,
-        "iit_revenue": iit_revenue,
-        "payroll_tax_revenue": payroll_tax_revenue,
-        "TR": TR,
-        "agg_pension_outlays": agg_pension_outlays,
-        "bequest_tax_revenue": bequest_tax_revenue,
-        "wealth_tax_revenue": wealth_tax_revenue,
-        "cons_tax_revenue": cons_tax_revenue,
-        "G": G,
-        "D": D,
-        "D_f": D_f,
-        "D_d": D_d,
-        "r": r,
-        "r_gov": r_gov,
-        "r_p": r_p,
-        "w": w,
-        "bmat_splus1": bmat_splus1,
-        "p_m": p_m,
+        "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,
-        "bq_path": bqmat,
-        "tr_path": trmat,
-        "y_before_tax_mat": y_before_tax_mat,
-        "tax_path": tax_mat,
-        "eul_savings": eul_savings,
-        "eul_laborleisure": eul_laborleisure,
-        "resource_constraint_error": RC_error,
-        "new_borrowing_f": new_borrowing_f,
-        "debt_service_f": debt_service_f,
-        "etr_path": etr_path,
-        "mtrx_path": mtrx_path,
-        "mtry_path": mtry_path,
-        "ubi_path": ubi,
-        "UBI_path": UBI,
+        "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")