From cf34ca087408a61e4bc080c8c6f6bfbb53f5fdeb Mon Sep 17 00:00:00 2001
From: Tim Mensinger <mensingertim@gmail.com>
Date: Tue, 13 Feb 2024 12:52:13 +0100
Subject: [PATCH] Incorporate comments from review

---
 pyproject.toml                       |  4 +-
 src/lcm/example_models_stochastic.py | 63 ++++++++++++++++++++++++----
 2 files changed, 58 insertions(+), 9 deletions(-)

diff --git a/pyproject.toml b/pyproject.toml
index 98e55818..bfba52a5 100644
--- 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
index 169e9c1b..fb139424 100644
--- 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]],
+                ],
             ],
         ),
     },