update SS and TPI algos to apply haircut to baseline spending

ogcore/SS.py

@@ -325,7 +325,7 @@ def inner_loop(outer_loop_vars, p, client):
     D, D_d, D_f, new_borrowing, _, new_borrowing_f = fiscal.get_D_ss(
         r_gov, Y, p
     )
-    I_g = fiscal.get_I_g(Y, p.alpha_I[-1])
+    I_g = fiscal.get_I_g(Y, p, "SS")
     K_g = fiscal.get_K_g(0, I_g, p, "SS")
 
     # Find wage rate consistent with open economy interest rate
@@ -364,7 +364,7 @@ def inner_loop(outer_loop_vars, p, client):
     Y_vec[-1] = firm.get_Y(K_vec[-1], K_g, L_vec[-1], p, "SS", -1)
     # Find GDP
     Y = (p_m * Y_vec).sum()
-    I_g = fiscal.get_I_g(Y, p.alpha_I[-1])
+    I_g = fiscal.get_I_g(Y, p, "SS")
     K_g = fiscal.get_K_g(0, I_g, p, "SS")
     if p.zeta_K[-1] == 1.0:
         new_r = p.world_int_rate[-1]
@@ -574,14 +574,14 @@ def SS_solver(
     if fsolve_flag:  # case where already solved via SS_fsolve
         maxiter_ss = 1
     if p.baseline_spending:
-        TR_ss = TR
+        TR_baseline = p.alpha_bs_T[-1] * TR
     if not p.budget_balance and not p.baseline_spending:
         Y = TR / p.alpha_T[-1]
     while (dist > p.mindist_SS) and (iteration < maxiter_ss):
         # Solve for the steady state levels of b and n, given w, r,
         # Y, BQ, TR, and factor
         if p.baseline_spending:
-            TR = TR_ss
+            TR = p.alpha_bs_T[-1] * TR_baseline
         if not p.budget_balance and not p.baseline_spending:
             Y = TR / p.alpha_T[-1]
 
@@ -641,7 +641,7 @@ def SS_solver(
                 ).max()
         else:
             if p.baseline_spending:
-                TR = TR_ss
+                TR = p.alpha_bs_T[-1] * TR_baseline
             else:
                 TR = utils.convex_combo(new_TR, TR, nu_ss)
             dist = np.array(
@@ -681,7 +681,7 @@ def SS_solver(
     p_tilde_ss = aggr.get_ptilde(p_i_ss, p.tau_c[-1, :], p.alpha_c)
     TR_ss = new_TR
     Yss = new_Y
-    I_g_ss = fiscal.get_I_g(Yss, p.alpha_I[-1])
+    I_g_ss = fiscal.get_I_g(Yss, p, "SS")
     K_g_ss = fiscal.get_K_g(0, I_g_ss, p, "SS")
     Lss = aggr.get_L(nssmat, p, "SS")
     Bss = aggr.get_B(bssmat_splus1, p, "SS", False)
@@ -704,7 +704,7 @@ def SS_solver(
         Bss, K_demand_open_ss.sum(), D_d_ss, p.zeta_K[-1]
     )
     # Yss = firm.get_Y(Kss, K_g_ss, Lss, p, 'SS')
-    I_g_ss = fiscal.get_I_g(Yss, p.alpha_I[-1])
+    I_g_ss = fiscal.get_I_g(Yss, p, "SS")
     K_g_ss = fiscal.get_K_g(0, I_g_ss, p, "SS")
     MPKg_vec = np.zeros(p.M)
     for m in range(p.M):

ogcore/TPI.py

@@ -68,9 +68,9 @@ def get_initial_SS_values(p):
     if p.baseline_spending:
         baseline_tpi = os.path.join(p.baseline_dir, "TPI", "TPI_vars.pkl")
         tpi_baseline_vars = utils.safe_read_pickle(baseline_tpi)
-        TRbaseline = tpi_baseline_vars["TR"]
-        Gbaseline = tpi_baseline_vars["G"]
-        Ig_baseline = tpi_baseline_vars["I_g"]
+        TRbaseline = p.alpha_bs_T * tpi_baseline_vars["TR"]
+        Gbaseline = p.alpha_bs_G * tpi_baseline_vars["G"]
+        Ig_baseline = p.alpha_I * tpi_baseline_vars["I_g"]
 
     if p.baseline:
         ss_vars = ss_baseline_vars
@@ -593,7 +593,7 @@ def run_TPI(p, client=None):
     if p.baseline_spending:
         I_g[: p.T] = Ig_baseline[: p.T]
     else:
-        I_g = fiscal.get_I_g(Y[: p.T], p.alpha_I[: p.T])
+        I_g = fiscal.get_I_g(Y[: p.T], p, "TPI")
     if p.baseline:
         K_g0 = p.initial_Kg_ratio * Y[0]
     else:
@@ -945,7 +945,7 @@ def run_TPI(p, client=None):
             B[: p.T], K_demand_open_vec.sum(-1), D_d[: p.T], p.zeta_K[: p.T]
         )
         if not p.baseline_spending:
-            I_g = fiscal.get_I_g(Y[: p.T], p.alpha_I[: p.T])
+            I_g = fiscal.get_I_g(Y[: p.T], p, "TPI")
         if p.baseline:
             K_g0 = p.initial_Kg_ratio * Y[0]
         K_g = fiscal.get_K_g(K_g0, I_g, p, "TPI")

ogcore/fiscal.py

@@ -78,7 +78,7 @@ def D_G_path(r_gov, dg_fixed_values, p):
         D[0] = D0_baseline
 
     if p.baseline_spending:
-        G = p.alpha_bs_G * Gbaseline[: p.T]
+        G = Gbaseline[: p.T]
     else:
         G = p.alpha_G[: p.T] * Y[: p.T]
 
@@ -375,7 +375,7 @@ def get_r_gov(r, p, method):
     return r_gov
 
 
-def get_I_g(Y, alpha_I):
+def get_I_g(Y, p, method="SS"):
     r"""
     Find investment in public capital
 
@@ -384,12 +384,22 @@ def get_I_g(Y, alpha_I):
 
     Args:
         Y (array_like): aggregate output
-        alpha_I (array_like): percentage of output invested in public capital
+        p (OG-Core Specifications object): model parameters
+        method (str): either 'SS' for steady-state or 'TPI' for transition path
 
     Returns
         I_g (array_like): investment in public capital
     """
-    I_g = alpha_I * Y
+    if p.baseline_spending:
+        if method == "SS":
+            I_g = p.alpha_bs_I * Y
+        else:
+            I_g = p.alpha_bs_I[:p.T] * Y
+    else:
+        if method == "SS":
+            I_g = p.alpha_I[-1] * Y
+        else:
+            I_g = p.alpha_I[:p.T] * Y
 
     return I_g