Merge pull request #1 from jdebacker/pr-871

Reshaping input rho
PSLmodels · Sep 24, 2023 · f3e4879 · f3e4879
2 parents 4c299ba + 89013bd
commit f3e4879
Show file tree

Hide file tree

Showing 12 changed files with 142 additions and 19 deletions.
diff --git a/ogcore/TPI.py b/ogcore/TPI.py
@@ -262,7 +262,7 @@ def twist_doughnut(
     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 - (s+1), 0))
+    rho_s = np.diag(p.rho[t: t + p.S, :], max(p.S - length, 0))
 
     error1 = household.FOC_savings(
         r_s,

diff --git a/ogcore/default_parameters.json b/ogcore/default_parameters.json
@@ -3547,10 +3547,10 @@
         "section_1": "Demographic Parameters",
         "notes": "",
         "type": "float",
-        "number_dims": 1,
+        "number_dims": 2,
         "value": [
             {
-                "value": [0.0007328533962063233, 0.0008171197422418564, 0.0008788729669386708, 0.0009116525208591186, 0.0009229802325645409, 0.0009273406814831464, 0.0009357102393970917, 0.0009500953703120985, 0.0009759735565497696, 0.0010103554581823992, 0.0010497459745855142, 0.0010896505583020133, 0.0011290564585360041, 0.0011654456037331062, 0.0012043239258441307, 0.0012511773865916398, 0.0013110045346917953, 0.0013828255727224015, 0.0014696353786847194, 0.0015729202258729647, 0.0016901599123994027, 0.0018253151315732463, 0.00198785868083029, 0.002182248418529209, 0.0024054583493769988, 0.002649495698085591, 0.002912359902066397, 0.003196521891964843, 0.0035014404113965503, 0.0038265853053018883, 0.004175977199939118, 0.004546000079867052, 0.004929386691983617, 0.005323440679881664, 0.005732162610849301, 0.00618294920155571, 0.0066650630096521946, 0.007142892752724284, 0.007605334723393753, 0.008079293606839011, 0.008607183802693097, 0.009219540427197637, 0.009916977317307896, 0.01071379466147504, 0.011618680680630211, 0.012651175630041145, 0.013806200722411788, 0.015069231422104257, 0.01643990870680101, 0.017946285365270986, 0.019682839990789902, 0.02164309127204489, 0.02374663347580741, 0.025983719293213747, 0.028428273440626484, 0.03125859803249409, 0.03451112898875497, 0.038087062505788305, 0.04199247322144961, 0.04632683864241138, 0.05133101803558171, 0.05704185782117355, 0.06332015480431863, 0.07017696315508493, 0.0777688358372085, 0.0862903314228447, 0.09590618992081834, 0.1067242554403307, 0.11879857491346835, 0.13212525643676043, 0.14668282278042066, 0.16243247757898882, 0.17932696648992474, 0.19731341445765527, 0.21633944779689218, 0.23535763270336796, 0.2540598444887924, 0.2721191714344474, 0.28919517981339793, 1.0]
+                "value": [[0.0007328533962063233, 0.0008171197422418564, 0.0008788729669386708, 0.0009116525208591186, 0.0009229802325645409, 0.0009273406814831464, 0.0009357102393970917, 0.0009500953703120985, 0.0009759735565497696, 0.0010103554581823992, 0.0010497459745855142, 0.0010896505583020133, 0.0011290564585360041, 0.0011654456037331062, 0.0012043239258441307, 0.0012511773865916398, 0.0013110045346917953, 0.0013828255727224015, 0.0014696353786847194, 0.0015729202258729647, 0.0016901599123994027, 0.0018253151315732463, 0.00198785868083029, 0.002182248418529209, 0.0024054583493769988, 0.002649495698085591, 0.002912359902066397, 0.003196521891964843, 0.0035014404113965503, 0.0038265853053018883, 0.004175977199939118, 0.004546000079867052, 0.004929386691983617, 0.005323440679881664, 0.005732162610849301, 0.00618294920155571, 0.0066650630096521946, 0.007142892752724284, 0.007605334723393753, 0.008079293606839011, 0.008607183802693097, 0.009219540427197637, 0.009916977317307896, 0.01071379466147504, 0.011618680680630211, 0.012651175630041145, 0.013806200722411788, 0.015069231422104257, 0.01643990870680101, 0.017946285365270986, 0.019682839990789902, 0.02164309127204489, 0.02374663347580741, 0.025983719293213747, 0.028428273440626484, 0.03125859803249409, 0.03451112898875497, 0.038087062505788305, 0.04199247322144961, 0.04632683864241138, 0.05133101803558171, 0.05704185782117355, 0.06332015480431863, 0.07017696315508493, 0.0777688358372085, 0.0862903314228447, 0.09590618992081834, 0.1067242554403307, 0.11879857491346835, 0.13212525643676043, 0.14668282278042066, 0.16243247757898882, 0.17932696648992474, 0.19731341445765527, 0.21633944779689218, 0.23535763270336796, 0.2540598444887924, 0.2721191714344474, 0.28919517981339793, 1.0]]
             }
         ],
         "validators": {

diff --git a/ogcore/parameter_plots.py b/ogcore/parameter_plots.py
@@ -54,10 +54,13 @@ def plot_mort_rates(p, years=[DEFAULT_START_YEAR], include_title=False, path=Non
 
     """
     age_per = np.linspace(p.E, p.E + p.S, p.S)
+    years = np.array(years) - p.start_year
     fig, ax = plt.subplots()
-    plt.plot(age_per, p.rho)
+    for y in years:
+        plt.plot(age_per, p.rho[y, :], label=str(y + p.start_year))
     plt.xlabel(r"Age $s$ (model periods)")
     plt.ylabel(r"Mortality Rates $\rho_{s}$")
+    plt.legend(loc="upper right")
     vals = ax.get_yticks()
     ax.set_yticklabels(["{:,.0%}".format(x) for x in vals])
     if include_title:

diff --git a/ogcore/parameters.py b/ogcore/parameters.py
@@ -347,6 +347,66 @@ def compute_default_params(self):
                     )
                 )
             setattr(self, item, this_attr)
+        # Deal with parameters that vary across age and over time
+        tp_param_list4 = [
+            "rho",
+        ]
+        for item in tp_param_list4:
+            this_attr = getattr(self, item)
+            if this_attr.ndim == 1:
+                # case where enter single number, so assume constant
+                # across years and age
+                if this_attr.shape[0] == 1:
+                    this_attr = (
+                        np.ones((self.T + self.S, self.S)) * this_attr[0]
+                    )
+                # case where user enters just one year for all ages
+                if this_attr.shape[0] == self.S:
+                    this_attr = np.tile(
+                        this_attr.reshape(1, self.S), (self.T + self.S, 1)
+                    )
+                else:
+                    # case where user enters multiple years for one age
+                    # will assume they implied values the same across
+                    # time for all periods after the last year of values
+                    this_attr = np.concatenate(
+                        (
+                            this_attr,
+                            np.ones((self.T + self.S - this_attr.size))
+                            * this_attr[-1],
+                        )
+                    )
+                    this_attr = np.tile(
+                        this_attr.reshape(self.T + self.S, 1), (1, self.S)
+                    )
+            elif this_attr.ndim == 2:
+                if this_attr.shape[1] > 1 and this_attr.shape[1] != self.S:
+                    print(
+                        "please provide values of "
+                        + item
+                        + " for each age (or one if common across "
+                        + "S}"
+                    )
+                    assert False
+                if this_attr.shape[1] == 1:
+                    this_attr = np.tile(
+                        this_attr.reshape(this_attr.shape[0], 1), (1, self.S)
+                    )
+                if this_attr.shape[0] > self.T + self.S:
+                    this_attr = this_attr[: self.T + self.S, :]
+                this_attr = np.concatenate(
+                    (
+                        this_attr,
+                        np.ones(
+                            (
+                                self.T + self.S - this_attr.shape[0],
+                                this_attr.shape[1],
+                            )
+                        )
+                        * this_attr[-1, :],
+                    )
+                )
+            setattr(self, item, this_attr)
         # Deal with tax parameters that maybe age and time specific
         tax_params_to_TP = [
             "etr_params",
@@ -631,8 +691,8 @@ def update_specifications(self, revision, raise_errors=True):
         """
         if not (isinstance(revision, dict) or isinstance(revision, str)):
             raise ValueError("ERROR: revision is not a dictionary or string")
-        # Skip over the adjust method if the tax paraemeters passed in
-        # are fucntions (e.g., in the case of tax_func_type = mono)
+        # Skip over the adjust method if the tax parameters passed in
+        # are functions (e.g., in the case of tax_func_type = mono)
         tax_update_dict = {}
         tax_func_params_functions = False
         try:

diff --git a/tests/test_TPI.py b/tests/test_TPI.py
@@ -183,6 +183,7 @@ def test_twist_doughnut(file_inputs, file_outputs):
         initial_b,
         p,
     )
+    print("t = ", t, s,  " p.S = ", p.S, p.rho.shape)
     test_list = TPI.twist_doughnut(*input_tuple)
     expected_list = utils.safe_read_pickle(file_outputs)
     assert np.allclose(np.array(test_list), np.array(expected_list), atol=1e-5)

diff --git a/tests/test_aggregates.py b/tests/test_aggregates.py
@@ -6,9 +6,12 @@
 
 
 p = Specifications()
+rho_vec = np.zeros((1, 40))
+rho_vec[0, -1] = 1.0
 new_param_values = {
     "T": 160,
     "S": 40,
+    "rho": rho_vec.tolist(),
     "J": 2,
     "labor_income_tax_noncompliance_rate": [[0.0]],
     "capital_income_tax_noncompliance_rate": [[0.0]],
@@ -43,9 +46,12 @@ def test_get_L(n, p, method, expected):
 
 
 p = Specifications()
+rho_vec = np.zeros((1, 40))
+rho_vec[0, -1] = 1.0
 new_param_values = {
     "T": 160,
     "S": 40,
+    "rho": rho_vec.tolist(),
     "J": 2,
     "labor_income_tax_noncompliance_rate": [[0.0]],
     "capital_income_tax_noncompliance_rate": [[0.0]],
@@ -120,9 +126,12 @@ def test_get_I(b_splus1, K_p1, K, p, method, expected):
 
 
 p = Specifications()
+rho_vec = np.zeros((1, 40))
+rho_vec[0, -1] = 1.0
 new_param_values = {
     "T": 160,
     "S": 40,
+    "rho": rho_vec.tolist(),
     "J": 2,
     "labor_income_tax_noncompliance_rate": [[0.0]],
     "capital_income_tax_noncompliance_rate": [[0.0]],
@@ -173,9 +182,12 @@ def test_get_B(b, p, method, PreTP, expected):
 
 
 p = Specifications()
+rho_vec = np.zeros((1, 40))
+rho_vec[0, -1] = 1.0
 new_param_values = {
     "T": 160,
     "S": 40,
+    "rho": rho_vec.tolist(),
     "J": 2,
     "labor_income_tax_noncompliance_rate": [[0.0]],
     "capital_income_tax_noncompliance_rate": [[0.0]],
@@ -184,7 +196,7 @@ def test_get_B(b, p, method, PreTP, expected):
     "omega": np.ones((160, 40)) / 40,
     "omega_SS": np.ones(40) / 40,
     "imm_rates": np.zeros((160, 40)),
-    "rho": np.zeros(40),
+    "rho": rho_vec.tolist(),
 }
 # update parameters instance with new values for test
 p.update_specifications(new_param_values)
@@ -194,7 +206,7 @@ def test_get_B(b, p, method, PreTP, expected):
 b_splus1 = 0.06 + 7 * np.random.rand(p.T, p.S, p.J)
 pop = np.append(p.omega_S_preTP.reshape(1, p.S), p.omega[: p.T - 1, :], axis=0)
 BQ_presum = (b_splus1 * np.squeeze(p.lambdas)) * np.tile(
-    np.reshape(p.rho * pop, (p.T, p.S, 1)), (1, 1, p.J)
+    np.reshape(p.rho[0, :] * pop, (p.T, p.S, 1)), (1, 1, p.J)
 )
 growth_adj = (1.0 + r) / (1.0 + p.g_n[: p.T])
 
@@ -251,9 +263,12 @@ def test_get_BQ(r, b_splus1, j, p, method, PreTP, expected):
 
 
 p = Specifications()
+rho_vec = np.zeros((1, 40))
+rho_vec[0, -1] = 1.0
 new_param_values = {
     "T": 160,
     "S": 40,
+    "rho": rho_vec.tolist(),
     "J": 2,
     "M": 3,
     "labor_income_tax_noncompliance_rate": [[0.0]],
@@ -297,9 +312,12 @@ def test_get_C(c, p, method, expected):
 """
 p = Specifications()
 num_tax_params = 12
+rho_vec = np.zeros((1, 20))
+rho_vec[0, -1] = 1.0
 new_param_values = {
     "T": 30,
     "S": 20,
+    "rho": rho_vec.tolist(),
     "J": 2,
     "labor_income_tax_noncompliance_rate": [[0.0]],
     "capital_income_tax_noncompliance_rate": [[0.0]],
@@ -349,9 +367,12 @@ def test_get_C(c, p, method, expected):
 theta = 0.101 + (0.156 - 0.101) * random_state.rand(p.J)
 
 p3 = Specifications()
+rho_vec = np.zeros((1, 20))
+rho_vec[0, -1] = 1.0
 new_param_values3 = {
     "T": 30,
     "S": 20,
+    "rho": rho_vec.tolist(),
     "J": 2,
     "labor_income_tax_noncompliance_rate": [[0.0]],
     "capital_income_tax_noncompliance_rate": [[0.0]],
@@ -380,9 +401,12 @@ def test_get_C(c, p, method, expected):
 inv_tax_cred_rev3 = np.squeeze(p3.inv_tax_credit[: p3.T] * Inv3)
 
 p_u = Specifications()
+rho_vec = np.zeros((1, 20))
+rho_vec[0, -1] = 1.0
 new_param_values_ubi = {
     "T": 30,
     "S": 20,
+    "rho": rho_vec.tolist(),
     "J": 2,
     "labor_income_tax_noncompliance_rate": [[0.0]],
     "capital_income_tax_noncompliance_rate": [[0.0]],

diff --git a/tests/test_firm.py b/tests/test_firm.py
@@ -27,12 +27,15 @@
 expected3 = np.array([0.592030917])
 # update parameters instance with new values for test
 p4 = Specifications()
+rho_vec = np.zeros((1, 3))
+rho_vec[0, -1] = 1.0
 new_param_values4 = {
     "Z": [[2.0]],
     "gamma": [0.5],
     "epsilon": [1.0],
     "T": 3,
     "S": 3,
+    "rho": rho_vec.tolist(),
     "eta": (np.ones((3, p4.J)) / (3 * p4.J)),
 }
 # update parameters instance with new values for test
@@ -48,6 +51,7 @@
     "epsilon": [1.0],
     "T": 3,
     "S": 3,
+    "rho": rho_vec.tolist(),
     "eta": (np.ones((3, p5.J)) / (3 * p5.J)),
 }
 # update parameters instance with new values for test
@@ -60,6 +64,7 @@
     "epsilon": [1.0],
     "T": 3,
     "S": 3,
+    "rho": rho_vec.tolist(),
     "eta": (np.ones((3, p5.J)) / (3 * p5.J)),
     "gamma_g": [0.2],
     "initial_Kg_ratio": 0.01,
@@ -79,6 +84,7 @@
     "T": 3,
     "S": 3,
     "M": 2,
+    "rho": rho_vec.tolist(),
     "eta": (np.ones((3, p4.J)) / (3 * p4.J)),
     "initial_Kg_ratio": 0.01,
 }
@@ -99,6 +105,7 @@
     "T": 3,
     "S": 3,
     "M": 2,
+    "rho": rho_vec.tolist(),
     "eta": (np.ones((3, p4.J)) / (3 * p4.J)),
     "initial_Kg_ratio": 0.01,
 }
@@ -115,6 +122,7 @@
     "T": 3,
     "S": 3,
     "M": 3,
+    "rho": rho_vec.tolist(),
     "eta": (np.ones((3, p5.J)) / (3 * p5.J)),
     "gamma_g": [0.2, 0.1, 0.25],
     "initial_Kg_ratio": 0.01,
@@ -139,6 +147,7 @@
     "T": 3,
     "S": 3,
     "M": 3,
+    "rho": rho_vec.tolist(),
     "eta": (np.ones((3, p5.J)) / (3 * p5.J)),
     "gamma_g": [0.2, 0.1, 0.25],
     "initial_Kg_ratio": 0.01,
@@ -258,6 +267,7 @@ def test_get_Y(K, K_g, L, p, method, m, expected):
     "delta_annual": 0.5,
     "T": 3,
     "S": 3,
+    "rho": rho_vec.tolist(),
     "eta": (np.ones((3, p4.J)) / (3 * p4.J)),
 }
 # update parameters instance with new values for test
@@ -277,6 +287,7 @@ def test_get_Y(K, K_g, L, p, method, m, expected):
     "inv_tax_credit": [[0.07]],
     "T": 3,
     "S": 3,
+    "rho": rho_vec.tolist(),
     "eta": (np.ones((3, p5.J)) / (3 * p5.J)),
 }
 # update parameters instance with new values for test
@@ -350,6 +361,7 @@ def test_get_r(Y, K, p_m, p, method, expected):
     "epsilon": [1.2],
     "T": 3,
     "S": 3,
+    "rho": rho_vec.tolist(),
     "eta": (np.ones((3, p4.J)) / (3 * p4.J)),
 }
 # update parameters instance with new values for test
@@ -438,6 +450,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,
+    "rho": rho_vec.tolist(),
     "eta": (np.ones((3, p4.J)) / (3 * p4.J)),
 }
 # update parameters instance with new values for test
@@ -545,6 +558,7 @@ def test_get_KLratio(r, p, method, expected):
     "cit_rate": [[(0.0357 / 0.55) * (0.055 / 0.017)]],
     "T": 3,
     "S": 3,
+    "rho": rho_vec.tolist(),
     "eta": (np.ones((3, p4.J)) / (3 * p4.J)),
 }
 # update parameters instance with new values for test
@@ -625,6 +639,7 @@ def test_get_w_from_r(r, p, method, expected):
     "cit_rate": [[0.5]],
     "T": 3,
     "S": 3,
+    "rho": rho_vec.tolist(),
     "eta": (np.ones((3, p4.J)) / (3 * p4.J)),
 }
 # update parameters instance with new values for test