Refactor example models

.pre-commit-config.yaml

@@ -33,6 +33,7 @@ repos:
       - id: name-tests-test
         args:
           - --pytest-test-first
+        exclude: ^tests/test_models/
       - id: no-commit-to-branch
         args:
           - --branch

MANIFEST.in

@@ -21,10 +21,11 @@ exclude *.yml
 exclude *.pickle
 exclude pytask.ini
 
-prune src/lcm/sandbox
+prune .envs
+prune examples
 prune docs
+prune src/lcm/sandbox
 prune tests
-prune .envs
 
 global-exclude __pycache__
 global-exclude *.py[co]

examples/README.md

@@ -0,0 +1,22 @@
+# Example model specifications
+
+## Choosing an example
+
+| Example name   | Description                                        | Runtime       |
+| -------------- | -------------------------------------------------- | ------------- |
+| `long_running` | Consumptions-savings model with health and leisure | a few minutes |
+
+## Running an example
+
+Say you want to run the [`long_running`](./long_running.py) example locally. In a Python
+shell, execute:
+
+```python
+from lcm.entry_point import get_lcm_function
+
+from long_running import MODEL_CONFIG, PARAMS
+
+
+solve_model, _ = get_lcm_function(model=MODEL_CONFIG)
+vf_arr = solve_model(PARAMS)
+```

examples/long_running.py

@@ -0,0 +1,121 @@
+"""Example specification for a consumption-savings model with health and leisure."""
+
+import jax.numpy as jnp
+
+# ======================================================================================
+# Numerical parameters and constants
+# ======================================================================================
+N_GRID_POINTS = {
+    "states": 100,
+    "choices": 200,
+}
+
+RETIREMENT_AGE = 65
+
+# ======================================================================================
+# Model functions
+# ======================================================================================
+
+
+# --------------------------------------------------------------------------------------
+# Utility function
+# --------------------------------------------------------------------------------------
+def utility(consumption, working, health, sport, delta):
+    return jnp.log(consumption) - (delta - health) * working - sport
+
+
+# --------------------------------------------------------------------------------------
+# Auxiliary variables
+# --------------------------------------------------------------------------------------
+def working(leisure):
+    return 1 - leisure
+
+
+def wage(age):
+    return 1 + 0.1 * age
+
+
+def age(_period):
+    return _period + 18
+
+
+# --------------------------------------------------------------------------------------
+# State transitions
+# --------------------------------------------------------------------------------------
+def next_wealth(wealth, consumption, working, wage, interest_rate):
+    return (1 + interest_rate) * (wealth - consumption) + wage * working
+
+
+def next_health(health, sport, working):
+    return health * (1 + sport - working / 2)
+
+
+def next_wealth_with_shock(
+    wealth,
+    consumption,
+    working,
+    wage,
+    wage_shock,
+    interest_rate,
+):
+    return interest_rate * (wealth - consumption) + wage * wage_shock * working
+
+
+# --------------------------------------------------------------------------------------
+# Constraints
+# --------------------------------------------------------------------------------------
+def consumption_constraint(consumption, wealth):
+    return consumption <= wealth
+
+
+# ======================================================================================
+# Model specification and parameters
+# ======================================================================================
+
+MODEL_CONFIG = {
+    "functions": {
+        "utility": utility,
+        "next_wealth": next_wealth,
+        "consumption_constraint": consumption_constraint,
+        "working": working,
+        "wage": wage,
+        "age": age,
+        "next_health": next_health,
+    },
+    "choices": {
+        "leisure": {"options": [0, 1]},
+        "consumption": {
+            "grid_type": "linspace",
+            "start": 1,
+            "stop": 100,
+            "n_points": N_GRID_POINTS["choices"],
+        },
+        "sport": {
+            "grid_type": "linspace",
+            "start": 0,
+            "stop": 1,
+            "n_points": N_GRID_POINTS["choices"],
+        },
+    },
+    "states": {
+        "wealth": {
+            "grid_type": "linspace",
+            "start": 1,
+            "stop": 100,
+            "n_points": N_GRID_POINTS["states"],
+        },
+        "health": {
+            "grid_type": "linspace",
+            "start": 0,
+            "stop": 1,
+            "n_points": N_GRID_POINTS["states"],
+        },
+    },
+    "n_periods": RETIREMENT_AGE - 18,
+}
+
+PARAMS = {
+    "beta": 0.95,
+    "utility": {"delta": 0.05},
+    "next_wealth": {"interest_rate": 0.05},
+}

pyproject.toml

@@ -9,6 +9,8 @@ write_to = "src/lcm/_version.py"
 [tool.ruff]
 target-version = "py311"
 fix = true
+
+[tool.ruff.lint]
 select = ["ALL"]
 extend-ignore = [
     # missing type annotation
@@ -57,6 +59,7 @@ extend-ignore = [
 [tool.ruff.lint.per-file-ignores]
 "docs/source/conf.py" = ["E501", "ERA001", "DTZ005"]
 "tests/test_*.py" = ["PLR2004"]
+"examples/*" = ["INP001"]
 
 [tool.ruff.lint.pydocstyle]
 convention = "google"
@@ -75,7 +78,7 @@ markers = [
     "slow: Tests that take a long time to run and are skipped in continuous integration.",
     "illustrative: Tests are designed for illustrative purposes",
 ]
-norecursedirs = ["docs", ".envs"]
+norecursedirs = ["docs", ".envs", "tests/test_models"]
 
 
 [tool.yamlfix]

src/lcm/get_model.py

@@ -4,7 +4,7 @@
 
 from pybaum import tree_update
 
-from lcm.example_models.basic_example_models import (
+from tests.test_models.phelps_deaton import (
     PHELPS_DEATON,
     PHELPS_DEATON_WITH_FILTERS,
 )

tests/test_entry_point.py

@@ -5,17 +5,18 @@
     create_compute_conditional_continuation_value,
     get_lcm_function,
 )
-from lcm.example_models.basic_example_models import (
-    PHELPS_DEATON,
-    PHELPS_DEATON_FULLY_DISCRETE,
-    PHELPS_DEATON_WITH_FILTERS,
-    phelps_deaton_utility,
-)
 from lcm.model_functions import get_utility_and_feasibility_function
 from lcm.process_model import process_model
 from lcm.state_space import create_state_choice_space
 from pybaum import tree_equal, tree_map
 
+from tests.test_models.phelps_deaton import (
+    PHELPS_DEATON,
+    PHELPS_DEATON_FULLY_DISCRETE,
+    PHELPS_DEATON_WITH_FILTERS,
+    utility,
+)
+
 MODELS = {
     "simple": PHELPS_DEATON,
     "with_filters": PHELPS_DEATON_WITH_FILTERS,
@@ -175,7 +176,7 @@ def test_create_compute_conditional_continuation_value():
         params=params,
         vf_arr=None,
     )
-    assert val == phelps_deaton_utility(consumption=30.0, working=0, delta=1.0)
+    assert val == utility(consumption=30.0, working=0, delta=1.0)
 
 
 def test_create_compute_conditional_continuation_value_with_discrete_model():
@@ -218,7 +219,7 @@ def test_create_compute_conditional_continuation_value_with_discrete_model():
         params=params,
         vf_arr=None,
     )
-    assert val == phelps_deaton_utility(consumption=2, working=0, delta=1.0)
+    assert val == utility(consumption=2, working=0, delta=1.0)
 
 
 # ======================================================================================
@@ -267,7 +268,7 @@ def test_create_compute_conditional_continuation_policy():
         vf_arr=None,
     )
     assert policy == 2
-    assert val == phelps_deaton_utility(consumption=30.0, working=0, delta=1.0)
+    assert val == utility(consumption=30.0, working=0, delta=1.0)
 
 
 def test_create_compute_conditional_continuation_policy_with_discrete_model():
@@ -311,4 +312,4 @@ def test_create_compute_conditional_continuation_policy_with_discrete_model():
         vf_arr=None,
     )
     assert policy == 1
-    assert val == phelps_deaton_utility(consumption=2, working=0, delta=1.0)
+    assert val == utility(consumption=2, working=0, delta=1.0)