Implement suggestions from review

OpenSourceEconomics · Feb 28, 2024 · dcaf4a6 · dcaf4a6
1 parent 7f6acef
commit dcaf4a6
Show file tree

Hide file tree

Showing 5 changed files with 24 additions and 22 deletions.
diff --git a/examples/long_running.py b/examples/long_running.py
@@ -6,8 +6,10 @@
 # Numerical parameters and constants
 # ======================================================================================
 N_GRID_POINTS = {
-    "states": 100,
-    "choices": 200,
+    "wealth": 100,
+    "health": 100,
+    "consumption": 100,
+    "exericse": 200,
 }
 
 RETIREMENT_AGE = 65
@@ -77,27 +79,27 @@ def consumption_constraint(consumption, wealth):
             "grid_type": "linspace",
             "start": 1,
             "stop": 100,
-            "n_points": N_GRID_POINTS["choices"],
+            "n_points": N_GRID_POINTS["consumption"],
         },
         "exercise": {
             "grid_type": "linspace",
             "start": 0,
             "stop": 1,
-            "n_points": N_GRID_POINTS["choices"],
+            "n_points": N_GRID_POINTS["exercise"],
         },
     },
     "states": {
         "wealth": {
             "grid_type": "linspace",
             "start": 1,
             "stop": 100,
-            "n_points": N_GRID_POINTS["states"],
+            "n_points": N_GRID_POINTS["wealth"],
         },
         "health": {
             "grid_type": "linspace",
             "start": 0,
             "stop": 1,
-            "n_points": N_GRID_POINTS["states"],
+            "n_points": N_GRID_POINTS["health"],
         },
     },
     "n_periods": RETIREMENT_AGE - 18,

diff --git a/tests/test_models/phelps_deaton.py b/tests/test_models/phelps_deaton.py
@@ -7,8 +7,8 @@
 # ======================================================================================
 
 N_GRID_POINTS = {
-    "states": 100,
-    "choices": 500,
+    "wealth": 100,
+    "consumption": 500,
 }
 
 RETIREMENT_AGE = 65
@@ -93,15 +93,15 @@ def absorbing_retirement_filter(retirement, lagged_retirement):
             "grid_type": "linspace",
             "start": 0,
             "stop": 100,
-            "n_points": N_GRID_POINTS["choices"],
+            "n_points": N_GRID_POINTS["consumption"],
         },
     },
     "states": {
         "wealth": {
             "grid_type": "linspace",
             "start": 0,
             "stop": 100,
-            "n_points": N_GRID_POINTS["states"],
+            "n_points": N_GRID_POINTS["wealth"],
         },
     },
     "n_periods": 3,
@@ -124,7 +124,7 @@ def absorbing_retirement_filter(retirement, lagged_retirement):
             "grid_type": "linspace",
             "start": 0,
             "stop": 100,
-            "n_points": N_GRID_POINTS["states"],
+            "n_points": N_GRID_POINTS["wealth"],
         },
     },
     "n_periods": 3,
@@ -146,15 +146,15 @@ def absorbing_retirement_filter(retirement, lagged_retirement):
             "grid_type": "linspace",
             "start": 1,
             "stop": 100,
-            "n_points": N_GRID_POINTS["choices"],
+            "n_points": N_GRID_POINTS["consumption"],
         },
     },
     "states": {
         "wealth": {
             "grid_type": "linspace",
             "start": 0,
             "stop": 100,
-            "n_points": N_GRID_POINTS["states"],
+            "n_points": N_GRID_POINTS["wealth"],
         },
         "lagged_retirement": {"options": [0, 1]},
     },

diff --git a/tests/test_models/stochastic.py b/tests/test_models/stochastic.py
@@ -8,8 +8,8 @@
 # ======================================================================================
 
 N_GRID_POINTS = {
-    "states": 100,
-    "choices": 200,
+    "wealth": 100,
+    "consumption": 200,
 }
 
 # ======================================================================================
@@ -69,7 +69,7 @@ def consumption_constraint(consumption, wealth):
             "grid_type": "linspace",
             "start": 1,
             "stop": 100,
-            "n_points": N_GRID_POINTS["choices"],
+            "n_points": N_GRID_POINTS["consumption"],
         },
     },
     "states": {
@@ -79,7 +79,7 @@ def consumption_constraint(consumption, wealth):
             "grid_type": "linspace",
             "start": 1,
             "stop": 100,
-            "n_points": N_GRID_POINTS["states"],
+            "n_points": N_GRID_POINTS["wealth"],
         },
     },
     "n_periods": 3,

diff --git a/tests/test_process_model.py b/tests/test_process_model.py
@@ -122,14 +122,14 @@ def test_process_phelps_deaton_with_filters():
     # Gridspecs
     wealth_specs = GridSpec(
         kind="linspace",
-        specs={"start": 0, "stop": 100, "n_points": N_GRID_POINTS["states"]},
+        specs={"start": 0, "stop": 100, "n_points": N_GRID_POINTS["wealth"]},
     )
 
     assert model.gridspecs["wealth"] == wealth_specs
 
     consumption_specs = GridSpec(
         kind="linspace",
-        specs={"start": 1, "stop": 100, "n_points": N_GRID_POINTS["choices"]},
+        specs={"start": 1, "stop": 100, "n_points": N_GRID_POINTS["consumption"]},
     )
     assert model.gridspecs["consumption"] == consumption_specs
 
@@ -173,14 +173,14 @@ def test_process_phelps_deaton():
     # Gridspecs
     wealth_specs = GridSpec(
         kind="linspace",
-        specs={"start": 0, "stop": 100, "n_points": N_GRID_POINTS["states"]},
+        specs={"start": 0, "stop": 100, "n_points": N_GRID_POINTS["wealth"]},
     )
 
     assert model.gridspecs["wealth"] == wealth_specs
 
     consumption_specs = GridSpec(
         kind="linspace",
-        specs={"start": 0, "stop": 100, "n_points": N_GRID_POINTS["choices"]},
+        specs={"start": 0, "stop": 100, "n_points": N_GRID_POINTS["consumption"]},
     )
     assert model.gridspecs["consumption"] == consumption_specs
 

diff --git a/tests/test_simulate.py b/tests/test_simulate.py
@@ -69,7 +69,7 @@ def simulate_inputs():
     return {
         "state_indexers": [{}],
         "continuous_choice_grids": [
-            {"consumption": jnp.linspace(1, 100, num=N_GRID_POINTS["choices"])},
+            {"consumption": jnp.linspace(1, 100, num=N_GRID_POINTS["consumption"])},
         ],
         "compute_ccv_policy_functions": compute_ccv_policy_functions,
         "model": model,