Merge pull request #895 from jdebacker/emat

Merging
PSLmodels · Jan 22, 2024 · 4b894bd · 4b894bd
2 parents d8a0fa4 + 5fa18cb
commit 4b894bd
Show file tree

Hide file tree

Showing 17 changed files with 207 additions and 87 deletions.
diff --git a/environment.yml b/environment.yml
@@ -2,7 +2,7 @@ name: ogcore-dev
 channels:
 - conda-forge
 dependencies:
-- python>=3.7.7
+- python>=3.7.7, <3.12
 - ipython
 - setuptools
 - scipy>=1.7.1

diff --git a/ogcore/SS.py b/ogcore/SS.py
@@ -78,7 +78,6 @@ def euler_equation_solver(guesses, *args):
         tr,
         ubi,
         theta,
-        p.e[:, j],
         p.rho[-1, :],
         p.etr_params[-1],
         p.mtry_params[-1],
@@ -100,7 +99,6 @@ def euler_equation_solver(guesses, *args):
         ubi,
         theta,
         p.chi_n,
-        p.e[:, j],
         p.etr_params[-1],
         p.mtrx_params[-1],
         None,
@@ -138,7 +136,7 @@ def euler_equation_solver(guesses, *args):
         j,
         False,
         "SS",
-        p.e[:, j],
+        p.e[-1, :, j],
         p.etr_params[-1],
         p,
     )
@@ -151,7 +149,7 @@ def euler_equation_solver(guesses, *args):
         n_guess,
         bq,
         taxes,
-        p.e[:, j],
+        p.e[-1, :, j],
         p,
     )
     mask6 = cons < 0
@@ -293,7 +291,7 @@ def inner_loop(outer_loop_vars, p, client):
         None,
         False,
         "SS",
-        p.e,
+        np.squeeze(p.e[-1, :, :]),
         etr_params_3D,
         p,
     )
@@ -306,7 +304,7 @@ def inner_loop(outer_loop_vars, p, client):
         nssmat,
         bq,
         net_tax,
-        p.e,
+        np.squeeze(p.e[-1, :, :]),
         p,
     )
     c_i = household.get_ci(c_s, p_i, p_tilde, p.tau_c[-1, :], p.alpha_c)
@@ -383,7 +381,7 @@ def inner_loop(outer_loop_vars, p, client):
         new_r, new_r_gov, p_m, K_vec, K_g, D, MPKg_vec, p, "SS"
     )
     average_income_model = (
-        (new_r_p * b_s + new_w * p.e * nssmat)
+        (new_r_p * b_s + new_w * np.squeeze(p.e[-1, :, :]) * nssmat)
         * p.omega_SS.reshape(p.S, 1)
         * p.lambdas.reshape(1, p.J)
     ).sum()
@@ -425,7 +423,7 @@ def inner_loop(outer_loop_vars, p, client):
         None,
         False,
         "SS",
-        p.e,
+        np.squeeze(p.e[-1, :, :]),
         etr_params_3D,
         p,
     )
@@ -438,7 +436,7 @@ def inner_loop(outer_loop_vars, p, client):
         nssmat,
         new_bq,
         taxss,
-        p.e,
+        np.squeeze(p.e[-1, :, :]),
         p,
     )
     (
@@ -467,6 +465,7 @@ def inner_loop(outer_loop_vars, p, client):
         ubi,
         theta,
         etr_params_3D,
+        np.squeeze(p.e[-1, :, :]),
         p,
         None,
         "SS",
@@ -482,6 +481,9 @@ def inner_loop(outer_loop_vars, p, client):
         debt_service,
         p,
     )
+    print("Agg tax = ", total_tax_revenue)
+    print("Agg pension outlays = ", agg_pension_outlays)
+    print("Agg UBI outlays = ", UBI_outlays)
     new_TR = fiscal.get_TR(
         Y,
         TR,
@@ -757,7 +759,7 @@ def SS_solver(
         nssmat,
         factor,
         True,
-        p.e,
+        np.squeeze(p.e[-1, :, :]),
         etr_params_3D,
         mtry_params_3D,
         capital_noncompliance_rate_2D,
@@ -770,7 +772,7 @@ def SS_solver(
         nssmat,
         factor,
         False,
-        p.e,
+        np.squeeze(p.e[-1, :, :]),
         etr_params_3D,
         mtrx_params_3D,
         labor_noncompliance_rate_2D,
@@ -782,7 +784,7 @@ def SS_solver(
         bssmat_s,
         nssmat,
         factor,
-        p.e,
+        np.squeeze(p.e[-1, :, :]),
         etr_params_3D,
         labor_noncompliance_rate_2D,
         capital_noncompliance_rate_2D,
@@ -803,7 +805,7 @@ def SS_solver(
         None,
         False,
         "SS",
-        p.e,
+        np.squeeze(p.e[-1, :, :]),
         etr_params_3D,
         p,
     )
@@ -816,10 +818,12 @@ def SS_solver(
         nssmat,
         bqssmat,
         taxss,
-        p.e,
+        np.squeeze(p.e[-1, :, :]),
         p,
     )
-    yss_before_tax_mat = household.get_y(r_p_ss, wss, bssmat_s, nssmat, p)
+    yss_before_tax_mat = household.get_y(
+        r_p_ss, wss, bssmat_s, nssmat, p, "SS"
+    )
     Css = aggr.get_C(cssmat, p, "SS")
     c_i_ss_mat = household.get_ci(
         cssmat, p_i_ss, p_tilde_ss, p.tau_c[-1, :], p.alpha_c
@@ -860,6 +864,7 @@ def SS_solver(
         ubissmat,
         theta,
         etr_params_3D,
+        np.squeeze(p.e[-1, :, :]),
         p,
         None,
         "SS",

diff --git a/ogcore/TPI.py b/ogcore/TPI.py
@@ -154,7 +154,6 @@ def firstdoughnutring(
         np.array([tr]),
         np.array([ubi]),
         theta[j],
-        p.e[-1, j],
         p.rho[0, -1],
         p.etr_params[0][-1],
         p.mtry_params[0][-1],
@@ -177,7 +176,6 @@ def firstdoughnutring(
         np.array([ubi]),
         theta[j],
         p.chi_n[-1],
-        p.e[-1, j],
         p.etr_params[0][-1],
         p.mtrx_params[0][-1],
         None,
@@ -261,7 +259,6 @@ def twist_doughnut(
     p_tilde_s = p_tilde[t : t + length]
     n_s = n_guess
     chi_n_s = p.chi_n[-length:]
-    e_s = p.e[-length:, j]
     rho_s = np.diag(p.rho[t : t + p.S, :], max(p.S - length, 0))
 
     error1 = household.FOC_savings(
@@ -276,7 +273,6 @@ def twist_doughnut(
         tr,
         ubi,
         theta,
-        e_s,
         rho_s,
         etr_params,
         mtry_params,
@@ -299,7 +295,6 @@ def twist_doughnut(
         ubi,
         theta,
         chi_n_s,
-        e_s,
         etr_params,
         mtrx_params,
         t,
@@ -823,6 +818,7 @@ def run_TPI(p, client=None):
             bmat_s[: p.T, :, :],
             n_mat[: p.T, :, :],
             p,
+            "TPI",
         )
 
         L[: p.T] = aggr.get_L(n_mat[: p.T], p, "TPI")
@@ -908,6 +904,7 @@ def run_TPI(p, client=None):
             ubi[: p.T, :, :],
             theta,
             etr_params_4D,
+            p.e,
             p,
             None,
             "TPI",
@@ -1014,6 +1011,7 @@ def run_TPI(p, client=None):
             ubi[: p.T, :, :],
             theta,
             etr_params_4D,
+            p.e,
             p,
             None,
             "TPI",
@@ -1143,7 +1141,7 @@ def run_TPI(p, client=None):
         ),
         (1, p.S, 1),
     )
-    e_3D = np.tile(p.e.reshape(1, p.S, p.J), (p.T, 1, 1))
+    e_3D = p.e
     mtry_path = tax.MTR_income(
         r_p_path[: p.T],
         wpath[: p.T],

diff --git a/ogcore/aggregates.py b/ogcore/aggregates.py
@@ -33,7 +33,11 @@ def get_L(n, p, method):
 
     """
     if method == "SS":
-        L_presum = p.e * np.transpose(p.omega_SS * p.lambdas) * n
+        L_presum = (
+            np.squeeze(p.e[-1, :, :])
+            * np.transpose(p.omega_SS * p.lambdas)
+            * n
+        )
         L = L_presum.sum()
     elif method == "TPI":
         L_presum = (n * (p.e * np.squeeze(p.lambdas))) * np.tile(
@@ -284,6 +288,7 @@ def revenue(
     ubi,
     theta,
     etr_params,
+    e,
     p,
     m,
     method,
@@ -317,6 +322,7 @@ def revenue(
             lifetime income group
         etr_params (list): list of parameters of the effective tax rate
             functions
+        e (Numpy array): effective labor units
         p (OG-Core Specifications object): model parameters
         method (str): adjusts calculation dimensions based on 'SS' or
             'TPI'
@@ -339,10 +345,10 @@ def revenue(
 
     """
     inc_pay_tax_liab = tax.income_tax_liab(
-        r, w, b, n, factor, 0, None, method, p.e, etr_params, p
+        r, w, b, n, factor, 0, None, method, e, etr_params, p
     )
     pension_benefits = tax.pension_amount(
-        w, n, theta, 0, None, False, method, p.e, p
+        w, n, theta, 0, None, False, method, e, p
     )
     bq_tax_liab = tax.bequest_tax_liab(r, b, bq, 0, None, method, p)
     w_tax_liab = tax.wealth_tax_liab(r, b, 0, None, method, p)