Time path varying parameters

ogusa/SS.py

@@ -122,12 +122,14 @@ def euler_equation_solver(guesses, *args):
 
     error1 = household.FOC_savings(r, w, b_s, b_splus1, n_guess, BQ,
                                    factor, T_H, theta, p.e[:, j], p.rho,
-                                   p.retire, p.etr_params[-1, :, :],
-                                   p.mtry_params[-1, :, :], j, p, 'SS')
+                                   p.etr_params[-1, :, :],
+                                   p.mtry_params[-1, :, :], None, j, p,
+                                   'SS')
     error2 = household.FOC_labor(r, w, b_s, b_splus1, n_guess, BQ,
                                  factor, T_H, theta, p.chi_n, p.e[:, j],
-                                 p.retire, p.etr_params[-1, :, :],
-                                 p.mtrx_params[-1, :, :], j, p, 'SS')
+                                 p.etr_params[-1, :, :],
+                                 p.mtrx_params[-1, :, :], None, j, p,
+                                 'SS')
 
     # Put in constraints for consumption and savings.
     # According to the euler equations, they can be negative.  When
@@ -145,7 +147,7 @@ def euler_equation_solver(guesses, *args):
     error1[mask5] = 1e14
     error2[mask4] = 1e14
     taxes = tax.total_taxes(r, w, b_s, n_guess, BQ, factor, T_H, theta,
-                            j, False, 'SS', p.e[:, j], p.retire,
+                            None, j, False, 'SS', p.e[:, j],
                             p.etr_params[-1, :, :], p)
     cons = household.get_cons(r, w, b_s, b_splus1, n_guess, BQ, taxes,
                               p.e[:, j], j, p)
@@ -199,7 +201,7 @@ def inner_loop(outer_loop_vars, p, client):
 
     euler_errors = np.zeros((2 * p.S, p.J))
 
-    w = firm.get_w_from_r(r, p)
+    w = firm.get_w_from_r(r, p, 'SS')
 
     lazy_values = []
     for j in range(p.J):
@@ -228,16 +230,15 @@ def inner_loop(outer_loop_vars, p, client):
         else:
             K = B - p.debt_ratio_ss * Y
     else:
-        K = firm.get_K(L, r, p)
-    new_Y = firm.get_Y(K, L, p)
+        K = firm.get_K(L, r, p, 'SS')
+    new_Y = firm.get_Y(K, L, p, 'SS')
     if p.budget_balance:
         Y = new_Y
     if not p.small_open:
-        new_r = firm.get_r(Y, K, p)
-        # new_w = firm.get_w_from_r(new_r, p)
+        new_r = firm.get_r(Y, K, p, 'SS')
     else:
-        new_r = p.ss_hh_r
-    new_w = firm.get_w_from_r(new_r, p)
+        new_r = p.hh_r[-1]
+    new_w = firm.get_w_from_r(new_r, p, 'SS')
     print('inner factor prices: ', new_r, new_w)
 
     b_s = np.array(list(np.zeros(p.J).reshape(1, p.J)) +
@@ -262,7 +263,7 @@ def inner_loop(outer_loop_vars, p, client):
     elif p.baseline_spending:
         new_T_H = T_H
     else:
-        new_T_H = p.alpha_T * new_Y
+        new_T_H = p.alpha_T[-1] * new_Y
 
     return euler_errors, bssmat, nssmat, new_r, new_w, \
         new_T_H, new_Y, new_factor, new_BQ, average_income_model
@@ -336,9 +337,9 @@ def SS_solver(bmat, nmat, r, T_H, factor, Y, p, client,
     # Rename the inputs
     if not p.budget_balance:
         if not p.baseline_spending:
-            Y = T_H / p.alpha_T
+            Y = T_H / p.alpha_T[-1]
     if p.small_open:
-        r = p.ss_hh_r
+        r = p.ss_hh_r[-1]
 
     dist = 10
     iteration = 0
@@ -392,7 +393,7 @@ def SS_solver(bmat, nmat, r, T_H, factor, Y, p, client,
                 nu_ss /= 2.0
                 print('New value of nu:', nu_ss)
         iteration += 1
-        print("Iteration: %02d" % iteration, " Distance: ", dist)
+        print('Iteration: %02d' % iteration, ' Distance: ', dist)
 
     '''
     ------------------------------------------------------------------------
@@ -419,7 +420,7 @@ def SS_solver(bmat, nmat, r, T_H, factor, Y, p, client,
         Iss = aggr.get_I(bssmat_splus1, Kss, Kss, p, 'SS')
     else:
         # Compute capital (K) and wealth (B) separately
-        Kss = firm.get_K(Lss, p.ss_firm_r, p)
+        Kss = firm.get_K(Lss, p.ss_firm_r[-1], p, 'SS')
         InvestmentPlaceholder = np.zeros(bssmat_splus1.shape)
         Iss = aggr.get_I(InvestmentPlaceholder, Kss, Kss, p, 'SS')
         Bss = aggr.get_K(bssmat_splus1, p, 'SS', False)
@@ -430,15 +431,15 @@ def SS_solver(bmat, nmat, r, T_H, factor, Y, p, client,
         else:
             debt_ss = p.debt_ratio_ss * Y
 
-    Yss = firm.get_Y(Kss, Lss, p)
+    Yss = firm.get_Y(Kss, Lss, p, 'SS')
     BQss = new_BQ
     theta = tax.replacement_rate_vals(nssmat, wss, factor_ss, None, p)
 
     etr_params_3D = np.tile(np.reshape(
         p.etr_params[-1, :, :], (p.S, 1, p.etr_params.shape[2])), (1, p.J, 1))
-    total_revenue_ss = aggr.revenue(rss, wss, bssmat_s, nssmat, BQss, Yss,
-                              Lss, Kss, factor, theta,
-                              etr_params_3D, p, 'SS')
+    total_revenue_ss = aggr.revenue(
+        rss, wss, bssmat_s, nssmat, BQss, Yss, Lss, Kss, factor, theta,
+        etr_params_3D, p, 'SS')
     r_gov_ss = rss
     debt_service_ss = r_gov_ss * p.debt_ratio_ss * Yss
     new_borrowing = p.debt_ratio_ss * Yss * ((1 + p.g_n_ss) *
@@ -469,13 +470,13 @@ def SS_solver(bmat, nmat, r, T_H, factor, Y, p, client,
 
     # solve resource constraint
     taxss = tax.total_taxes(rss, wss, bssmat_s, nssmat, BQss, factor_ss,
-                            T_Hss, theta, None, False, 'SS',
-                            p.e, p.retire, etr_params_3D, p)
+                            T_Hss, theta, None, None, False, 'SS',
+                            p.e, etr_params_3D, p)
     cssmat = household.get_cons(rss, wss, bssmat_s, bssmat_splus1,
                                 nssmat, BQss.reshape(1, p.J), taxss,
                                 p.e, None, p)
 
-    business_revenue = tax.get_biz_tax(wss, Yss, Lss, Kss, p)
+    business_revenue = tax.get_biz_tax(wss, Yss, Lss, Kss, p, 'SS')
     IITpayroll_revenue = total_revenue_ss - business_revenue
 
     Css = aggr.get_C(cssmat, p, 'SS')
@@ -484,7 +485,8 @@ def SS_solver(bmat, nmat, r, T_H, factor, Y, p, client,
     if p.small_open:
         # include term for current account
         resource_constraint = (Yss + new_borrowing - (Css + BIss + Gss)
-                               + (p.ss_hh_r * Bss - (p.delta + p.ss_firm_r) *
+                               + (p.ss_hh_r[-1] * Bss -
+                                  (p.delta + p.ss_firm_r[-1]) *
                                   Kss - debt_service_ss))
         print('Yss= ', Yss, '\n', 'Css= ', Css, '\n', 'Bss = ', Bss,
               '\n', 'BIss = ', BIss, '\n', 'Kss = ', Kss, '\n', 'Iss = ',
@@ -515,7 +517,7 @@ def SS_solver(bmat, nmat, r, T_H, factor, Y, p, client,
     if ENFORCE_SOLUTION_CHECKS and (np.absolute(resource_constraint) >
                                     p.mindist_SS):
         print('Resource Constraint Difference:', resource_constraint)
-        err = "Steady state aggregate resource constraint not satisfied"
+        err = 'Steady state aggregate resource constraint not satisfied'
         raise RuntimeError(err)
 
     # check constraints
@@ -586,7 +588,7 @@ def SS_fsolve(guesses, *args):
     if p.budget_balance:
         outer_loop_vars = (bssmat, nssmat, r, T_H, factor)
     else:
-        Y = T_H / p.alpha_T
+        Y = T_H / p.alpha_T[-1]
         outer_loop_vars = (bssmat, nssmat, r, Y, T_H, factor)
     (euler_errors, bssmat, nssmat, new_r, new_w, new_T_H, new_Y,
      new_factor, new_BQ, average_income_model) =\
@@ -649,7 +651,7 @@ def SS_fsolve_reform(guesses, *args):
     if p.budget_balance:
         outer_loop_vars = (bssmat, nssmat, r, T_H, factor)
     else:
-        Y = T_H / p.alpha_T
+        Y = T_H / p.alpha_T[-1]
         outer_loop_vars = (bssmat, nssmat, r, Y, T_H, factor)
 
     (euler_errors, bssmat, nssmat, new_r, new_w, new_T_H, new_Y,
@@ -789,18 +791,18 @@ def run_SS(p, client=None):
             opt.fsolve(SS_fsolve, guesses, args=ss_params_baseline,
                        xtol=p.mindist_SS, full_output=True)
         if ENFORCE_SOLUTION_CHECKS and not ier == 1:
-            raise RuntimeError("Steady state equilibrium not found")
+            raise RuntimeError('Steady state equilibrium not found')
         [rss, T_Hss, factor_ss] = solutions_fsolve
-        Yss = T_Hss/p.alpha_T  # may not be right - if budget_balance
+        Yss = T_Hss/p.alpha_T[-1]  # may not be right - if budget_balance
         # = True, but that's ok - will be fixed in SS_solver
         fsolve_flag = True
         # Return SS values of variables
         output = SS_solver(b_guess, n_guess, rss, T_Hss, factor_ss, Yss, p,
                            client, fsolve_flag)
     else:
-        baseline_ss_dir = os.path.join(p.baseline_dir, "SS/SS_vars.pkl")
-        ss_solutions = pickle.load(open(baseline_ss_dir, "rb"),
-                                   encoding="latin1")
+        baseline_ss_dir = os.path.join(p.baseline_dir, 'SS/SS_vars.pkl')
+        ss_solutions = pickle.load(open(baseline_ss_dir, 'rb'),
+                                   encoding='latin1')
         (rguess, T_Hguess, Yguess, factor) =\
             (ss_solutions['rss'], ss_solutions['T_Hss'],
              ss_solutions['Yss'], ss_solutions['factor_ss'])
@@ -821,11 +823,11 @@ def run_SS(p, client=None):
                            args=ss_params_reform, xtol=p.mindist_SS,
                            full_output=True)
             [rss, T_Hss] = solutions_fsolve
-            Yss = T_Hss/p.alpha_T  # may not be right - if
+            Yss = T_Hss/p.alpha_T[-1]  # may not be right - if
             # budget_balance = True, but that's ok - will be fixed in
             # SS_solver
         if ENFORCE_SOLUTION_CHECKS and not ier == 1:
-            raise RuntimeError("Steady state equilibrium not found")
+            raise RuntimeError('Steady state equilibrium not found')
         # Return SS values of variables
         fsolve_flag = True
         # Return SS values of variables