Incorporate comments from review

OpenSourceEconomics · Feb 13, 2024 · cf34ca0 · cf34ca0
1 parent 0545080
commit cf34ca0
Show file tree

Hide file tree

Showing 2 changed files with 58 additions and 9 deletions.
diff --git a/pyproject.toml b/pyproject.toml
@@ -54,11 +54,11 @@ extend-ignore = [
     "TRY003",
 ]
 
-[tool.ruff.per-file-ignores]
+[tool.ruff.lint.per-file-ignores]
 "docs/source/conf.py" = ["E501", "ERA001", "DTZ005"]
 "tests/test_*.py" = ["PLR2004"]
 
-[tool.ruff.pydocstyle]
+[tool.ruff.lint.pydocstyle]
 convention = "google"
 
 

diff --git a/src/lcm/example_models_stochastic.py b/src/lcm/example_models_stochastic.py
@@ -9,7 +9,7 @@
 
 
 def utility(consumption, working, health, partner, delta, gamma):  # noqa: ARG001
-    return jnp.log(consumption + 1) + (gamma * health - delta) * working
+    return jnp.log(consumption) + (gamma * health - delta) * working
 
 
 def next_wealth(wealth, consumption, working, wage, interest_rate):
@@ -42,7 +42,7 @@ def consumption_constraint(consumption, wealth):
         "working": {"options": [0, 1]},
         "consumption": {
             "grid_type": "linspace",
-            "start": 0,
+            "start": 1,
             "stop": 100,
             "n_points": N_CHOICE_GRID_POINTS,
         },
@@ -52,7 +52,7 @@ def consumption_constraint(consumption, wealth):
         "partner": {"options": [0, 1]},
         "wealth": {
             "grid_type": "linspace",
-            "start": 0,
+            "start": 1,
             "stop": 100,
             "n_points": N_STATE_GRID_POINTS,
         },
@@ -62,17 +62,66 @@ def consumption_constraint(consumption, wealth):
 
 
 PARAMS = {
-    "beta": 1.0,
+    "beta": 0.95,
     "utility": {"delta": 0.5, "gamma": 0.25},
     "next_wealth": {"interest_rate": 0.05, "wage": 10.0},
     "next_health": {},
     "consumption_constraint": {},
     "shocks": {
-        "health": jnp.array([[[0.5, 0.5], [0.5, 0.5]], [[0.5, 0.5], [0.5, 0.5]]]),
+        # Health shock:
+        # ------------------------------------------------------------------------------
+        # 1st dimension: Current health state
+        # 2nd dimension: Current Partner state
+        # 3rd dimension: Probability distribution over next period's health state
+        "health": jnp.array(
+            [
+                # Current health state 0
+                [
+                    # Current Partner state 0
+                    [0.9, 0.1],
+                    # Current Partner state 1
+                    [0.5, 0.5],
+                ],
+                # Current health state 1
+                [
+                    # Current Partner state 0
+                    [0.5, 0.5],
+                    # Current Partner state 1
+                    [0.1, 0.9],
+                ],
+            ],
+        ),
+        # Partner shock:
+        # ------------------------------------------------------------------------------
+        # 1st dimension: The period
+        # 2nd dimension: Current working decision
+        # 3rd dimension: Current partner state
+        # 4th dimension: Probability distribution over next period's partner state
         "partner": jnp.array(
             [
-                [[[0, 1.0], [1.0, 0]], [[0, 1.0], [0.0, 1.0]]],
-                [[[0, 1.0], [1.0, 0]], [[0, 1.0], [0.0, 1.0]]],
+                # Transition from period 0 to period 1
+                [
+                    # Current working decision 0
+                    [
+                        # Current partner state 0
+                        [0, 1.0],
+                        # Current partner state 1
+                        [1.0, 0],
+                    ],
+                    # Current working decision 1
+                    [
+                        # Current partner state 0
+                        [0, 1.0],
+                        # Current partner state 1
+                        [0.0, 1.0],
+                    ],
+                ],
+                # Transition from period 1 to period 2
+                [
+                    # Description is the same as above
+                    [[0, 1.0], [1.0, 0]],
+                    [[0, 1.0], [0.0, 1.0]],
+                ],
             ],
         ),
     },