Allow chi_n parameter to vary over the time path

ogcore/SS.py

@@ -98,7 +98,7 @@ def euler_equation_solver(guesses, *args):
         tr,
         ubi,
         theta,
-        p.chi_n,
+        p.chi_n[-1, :],
         p.etr_params[-1],
         p.mtrx_params[-1],
         None,

ogcore/TPI.py

@@ -175,7 +175,7 @@ def firstdoughnutring(
         np.array([tr]),
         np.array([ubi]),
         theta[j],
-        p.chi_n[-1],
+        p.chi_n[0, -1],
         p.etr_params[0][-1],
         p.mtrx_params[0][-1],
         None,
@@ -258,7 +258,7 @@ def twist_doughnut(
     r_s = r[t : t + length]
     p_tilde_s = p_tilde[t : t + length]
     n_s = n_guess
-    chi_n_s = p.chi_n[-length:]
+    chi_n_s = np.diag(p.chi_n[t : t + p.S, :], max(p.S - length, 0))
     rho_s = np.diag(p.rho[t : t + p.S, :], max(p.S - length, 0))
 
     error1 = household.FOC_savings(

ogcore/parameters.py

@@ -128,6 +128,30 @@ def compute_default_params(self):
                 np.ones(self.T + self.S - self.BW) * self.frac_tax_payroll[-1],
             )
 
+        # Interpolate chi_n and create omega_SS_80 if necessary
+        if self.S < 80 and self.chi_n.shape[-1] == 80:
+            self.age_midp_80 = np.linspace(20.5, 99.5, 80)
+            reshape_chi_n = np.zeros((self.T + self.S, self.S))
+            for t in range(self.T + self.S):
+                if self.chi_n.ndim == 1:
+                    self.chi_n_interp = si.interp1d(
+                        self.age_midp_80,
+                        np.squeeze(self.chi_n),
+                        kind="cubic",
+                    )
+                else:
+                    self.chi_n_interp = si.interp1d(
+                        self.age_midp_80,
+                        np.squeeze(self.chi_n[t, :]),
+                        kind="cubic",
+                    )
+                self.newstep = 80.0 / self.S
+                self.age_midp_S = np.linspace(
+                    20 + 0.5 * self.newstep, 100 - 0.5 * self.newstep, self.S
+                )
+                reshape_chi_n[t, :] = self.chi_n_interp(self.age_midp_S)
+            self.chi_n = reshape_chi_n
+
         # Extend parameters that may vary over the time path
         tp_param_list = [
             "alpha_G",
@@ -262,6 +286,12 @@ def compute_default_params(self):
             param_in, dims=(self.T, self.S, self.J), item="e"
         )
         setattr(self, "e", param_out)
+        # Extrapolate chi_n over T + S
+        param_in = getattr(self, "chi_n")
+        param_out = extrapolate_arrays(
+            param_in, dims=(self.T + self.S, self.S), item="chi_n"
+        )
+        setattr(self, "chi_n", param_out)
 
         # make sure zeta matrix sums to one (e.g., default off due to rounding)
         self.zeta = self.zeta / self.zeta.sum()
@@ -319,18 +349,6 @@ def compute_default_params(self):
             )
             self.omega_S_preTP = self.omega_SS
 
-        # Interpolate chi_n and create omega_SS_80 if necessary
-        if self.S < 80 and len(self.chi_n) == 80:
-            self.age_midp_80 = np.linspace(20.5, 99.5, 80)
-            self.chi_n_interp = si.interp1d(
-                self.age_midp_80, np.squeeze(self.chi_n), kind="cubic"
-            )
-            self.newstep = 80.0 / self.S
-            self.age_midp_S = np.linspace(
-                20 + 0.5 * self.newstep, 100 - 0.5 * self.newstep, self.S
-            )
-            self.chi_n = self.chi_n_interp(self.age_midp_S)
-
         # Create time series of stationarized UBI transfers
         self.ubi_nom_array = self.get_ubi_nom_objs()
 

ogcore/utils.py

@@ -932,20 +932,21 @@ def extrapolate_arrays(param_in, dims=None, item="Parameter Name"):
                 param_out = np.ones((dims)) * param_in[0]
             # case where user enters just one year for all types in 2nd dim
             if param_in.shape[0] == dims[1]:
-                param_in = np.tile(param_in.reshape(1, dims[1]), (dims[0], 1))
+                param_out = np.tile(param_in.reshape(1, dims[1]), (dims[0], 1))
             else:
                 # case where user enters multiple years, but not full time
                 # path
                 # will assume they implied values the same across 2nd dim
                 # and will fill in all periods
-                param_in = np.concatenate(
+                param_out = np.concatenate(
                     (
                         param_in,
                         np.ones((dims[0] - param_in.size)) * param_in[-1],
                     )
                 )
-                param_in = np.tile(param_in.reshape(dims[0], 1), (1, dims[1]))
-            param_out = np.squeeze(param_in, axis=2)
+                param_out = np.tile(
+                    param_out.reshape(dims[0], 1), (1, dims[1])
+                )
         elif param_in.ndim == 2:
             # case where enter values along 2 dimensions, but those aren't
             # complete

tests/test_aggregates.py

@@ -13,6 +13,7 @@
     "S": 40,
     "rho": rho_vec.tolist(),
     "J": 2,
+    "chi_n": np.ones(2),
     "labor_income_tax_noncompliance_rate": [[0.0]],
     "capital_income_tax_noncompliance_rate": [[0.0]],
     "eta": (np.ones((40, 2)) / (40 * 2)),
@@ -53,6 +54,7 @@ def test_get_L(n, p, method, expected):
     "S": 40,
     "rho": rho_vec.tolist(),
     "J": 2,
+    "chi_n": np.ones(2),
     "e": np.ones((40, 2)),
     "labor_income_tax_noncompliance_rate": [[0.0]],
     "capital_income_tax_noncompliance_rate": [[0.0]],
@@ -134,6 +136,7 @@ def test_get_I(b_splus1, K_p1, K, p, method, expected):
     "S": 40,
     "rho": rho_vec.tolist(),
     "J": 2,
+    "chi_n": np.ones(2),
     "e": np.ones((40, 2)),
     "labor_income_tax_noncompliance_rate": [[0.0]],
     "capital_income_tax_noncompliance_rate": [[0.0]],
@@ -191,6 +194,7 @@ def test_get_B(b, p, method, PreTP, expected):
     "S": 40,
     "rho": rho_vec.tolist(),
     "J": 2,
+    "chi_n": np.ones(2),
     "e": np.ones((40, 2)),
     "labor_income_tax_noncompliance_rate": [[0.0]],
     "capital_income_tax_noncompliance_rate": [[0.0]],
@@ -273,6 +277,7 @@ def test_get_BQ(r, b_splus1, j, p, method, PreTP, expected):
     "S": 40,
     "rho": rho_vec.tolist(),
     "J": 2,
+    "chi_n": np.ones(2),
     "e": np.ones((40, 2)),
     "M": 3,
     "labor_income_tax_noncompliance_rate": [[0.0]],
@@ -323,6 +328,7 @@ def test_get_C(c, p, method, expected):
     "S": 20,
     "rho": rho_vec.tolist(),
     "J": 2,
+    "chi_n": np.ones(2),
     "e": np.ones((20, 2)),
     "labor_income_tax_noncompliance_rate": [[0.0]],
     "capital_income_tax_noncompliance_rate": [[0.0]],
@@ -380,6 +386,7 @@ def test_get_C(c, p, method, expected):
     "S": 20,
     "rho": rho_vec.tolist(),
     "J": 2,
+    "chi_n": np.ones(2),
     "e": np.ones((20, 2)),
     "labor_income_tax_noncompliance_rate": [[0.0]],
     "capital_income_tax_noncompliance_rate": [[0.0]],
@@ -415,6 +422,7 @@ def test_get_C(c, p, method, expected):
     "S": 20,
     "rho": rho_vec.tolist(),
     "J": 2,
+    "chi_n": np.ones(2),
     "e": np.ones((20, 2)),
     "labor_income_tax_noncompliance_rate": [[0.0]],
     "capital_income_tax_noncompliance_rate": [[0.0]],

tests/test_firm.py

@@ -35,6 +35,7 @@
     "epsilon": [1.0],
     "T": 3,
     "S": 3,
+    "chi_n": np.ones(3),
     "e": np.ones((3, p4.J)),
     "rho": rho_vec.tolist(),
     "eta": (np.ones((3, p4.J)) / (3 * p4.J)),
@@ -52,6 +53,7 @@
     "epsilon": [1.0],
     "T": 3,
     "S": 3,
+    "chi_n": np.ones(3),
     "rho": rho_vec.tolist(),
     "eta": (np.ones((3, p5.J)) / (3 * p5.J)),
     "e": np.ones((3, p5.J)),
@@ -66,6 +68,7 @@
     "epsilon": [1.0],
     "T": 3,
     "S": 3,
+    "chi_n": np.ones(3),
     "e": np.ones((3, p5.J)),
     "rho": rho_vec.tolist(),
     "eta": (np.ones((3, p5.J)) / (3 * p5.J)),
@@ -87,6 +90,7 @@
     "T": 3,
     "S": 3,
     "M": 2,
+    "chi_n": np.ones(3),
     "e": np.ones((3, p4.J)),
     "rho": rho_vec.tolist(),
     "eta": (np.ones((3, p4.J)) / (3 * p4.J)),
@@ -109,6 +113,7 @@
     "T": 3,
     "S": 3,
     "M": 2,
+    "chi_n": np.ones(3),
     "e": np.ones((3, p4.J)),
     "rho": rho_vec.tolist(),
     "eta": (np.ones((3, p4.J)) / (3 * p4.J)),
@@ -127,6 +132,7 @@
     "T": 3,
     "S": 3,
     "M": 3,
+    "chi_n": np.ones(3),
     "e": np.ones((3, p5.J)),
     "rho": rho_vec.tolist(),
     "eta": (np.ones((3, p5.J)) / (3 * p5.J)),
@@ -153,6 +159,7 @@
     "T": 3,
     "S": 3,
     "M": 3,
+    "chi_n": np.ones(3),
     "e": np.ones((3, p5.J)),
     "rho": rho_vec.tolist(),
     "eta": (np.ones((3, p5.J)) / (3 * p5.J)),
@@ -274,6 +281,7 @@ def test_get_Y(K, K_g, L, p, method, m, expected):
     "delta_annual": 0.5,
     "T": 3,
     "S": 3,
+    "chi_n": np.ones(3),
     "e": np.ones((3, p4.J)),
     "rho": rho_vec.tolist(),
     "eta": (np.ones((3, p4.J)) / (3 * p4.J)),
@@ -295,6 +303,7 @@ def test_get_Y(K, K_g, L, p, method, m, expected):
     "inv_tax_credit": [[0.07]],
     "T": 3,
     "S": 3,
+    "chi_n": np.ones(3),
     "e": np.ones((3, p4.J)),
     "rho": rho_vec.tolist(),
     "eta": (np.ones((3, p5.J)) / (3 * p5.J)),
@@ -370,6 +379,7 @@ def test_get_r(Y, K, p_m, p, method, expected):
     "epsilon": [1.2],
     "T": 3,
     "S": 3,
+    "chi_n": np.ones(3),
     "e": np.ones((3, p4.J)),
     "rho": rho_vec.tolist(),
     "eta": (np.ones((3, p4.J)) / (3 * p4.J)),
@@ -460,6 +470,7 @@ def test_get_w(Y, L, p_m, p, method, expected):
     "cit_rate": [[(0.0357 / 0.55) * (0.055 / 0.017)]],
     "T": 3,
     "S": 3,
+    "chi_n": np.ones(3),
     "e": np.ones((3, p4.J)),
     "rho": rho_vec.tolist(),
     "eta": (np.ones((3, p4.J)) / (3 * p4.J)),
@@ -569,6 +580,7 @@ def test_get_KLratio(r, p, method, expected):
     "cit_rate": [[(0.0357 / 0.55) * (0.055 / 0.017)]],
     "T": 3,
     "S": 3,
+    "chi_n": np.ones(3),
     "e": np.ones((3, p4.J)),
     "rho": rho_vec.tolist(),
     "eta": (np.ones((3, p4.J)) / (3 * p4.J)),
@@ -651,6 +663,7 @@ def test_get_w_from_r(r, p, method, expected):
     "cit_rate": [[0.5]],
     "T": 3,
     "S": 3,
+    "chi_n": np.ones(3),
     "e": np.ones((3, p4.J)),
     "rho": rho_vec.tolist(),
     "eta": (np.ones((3, p4.J)) / (3 * p4.J)),

tests/test_tax.py

@@ -16,6 +16,7 @@
     "capital_income_tax_noncompliance_rate": [[0.0]],
     "J": 1,
     "T": 4,
+    "chi_n": np.ones(4),
     "eta": (np.ones((4, 1)) / (4 * 1)),
     "e": np.ones((4, 1)),
 }
@@ -81,6 +82,7 @@ def test_replacement_rate_vals(n, w, factor, j, p_in, expected):
     "lambdas": [1.0],
     "J": 1,
     "T": 3,
+    "chi_n": np.ones(3),
     "e": np.ones((3, 1)),
     "eta": (np.ones((3, 1)) / (3 * 1)),
     "h_wealth": [2],
@@ -98,6 +100,7 @@ def test_replacement_rate_vals(n, w, factor, j, p_in, expected):
     "lambdas": [1.0],
     "J": 1,
     "T": 3,
+    "chi_n": np.ones(3),
     "e": np.ones((3, 1)),
     "eta": (np.ones((3, 1)) / (3 * 1)),
     "h_wealth": [1.2, 1.1, 2.3],
@@ -132,6 +135,7 @@ def test_ETR_wealth(b, p, expected):
     "lambdas": [1.0],
     "J": 1,
     "T": 3,
+    "chi_n": np.ones(3),
     "e": np.ones((3, 1)),
     "eta": (np.ones((3, 1)) / (3 * 1)),
     "h_wealth": [3],
@@ -150,6 +154,7 @@ def test_ETR_wealth(b, p, expected):
     "lambdas": [1.0],
     "J": 1,
     "T": 3,
+    "chi_n": np.ones(3),
     "e": np.ones((3, 1)),
     "eta": (np.ones((3, 1)) / (3 * 1)),
     "h_wealth": [1.2, 1.1, 2.3],
@@ -885,6 +890,7 @@ def test_MTR_income(etr_params, mtr_params, params, mtr_capital, expected):
     "inv_tax_credit": [[0.02]],
     "T": 3,
     "S": 3,
+    "chi_n": np.ones(3),
     "e": np.ones((3, p1.J)),
     "rho": rho_vec.tolist(),
     "eta": (np.ones((3, p1.J)) / (3 * p1.J)),
@@ -1214,6 +1220,7 @@ def test_get_biz_tax(w, Y, L, K, p_m, p, m, method, expected):
     "T": 3,
     "S": 3,
     "J": 2,
+    "chi_n": np.ones(3),
     "e": np.ones((3, 2)),
     "rho": rho_vec.tolist(),
     "lambdas": [0.65, 0.35],