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` | Consumption-savings model with health and leisure | a few minutes |
+
+## Running an example
+
+Say you want to solve 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,117 @@
+"""Example specification for a consumption-savings model with health and leisure."""
+
+import jax.numpy as jnp
+
+# ======================================================================================
+# Numerical parameters and constants
+# ======================================================================================
+N_GRID_POINTS = {
+    "wealth": 100,
+    "health": 100,
+    "consumption": 100,
+    "exercise": 200,
+}
+
+RETIREMENT_AGE = 65
+
+# ======================================================================================
+# Model functions
+# ======================================================================================
+
+
+# --------------------------------------------------------------------------------------
+# Utility function
+# --------------------------------------------------------------------------------------
+def utility(consumption, working, health, exercise, disutility_of_work):
+    return jnp.log(consumption) - (disutility_of_work - health) * working - exercise
+
+
+# --------------------------------------------------------------------------------------
+# Auxiliary variables
+# --------------------------------------------------------------------------------------
+def labor_income(wage, working):
+    return wage * working
+
+
+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, labor_income, interest_rate):
+    return (1 + interest_rate) * (wealth + labor_income - consumption)
+
+
+def next_health(health, exercise, working):
+    return health * (1 + exercise - working / 2)
+
+
+# --------------------------------------------------------------------------------------
+# Constraints
+# --------------------------------------------------------------------------------------
+def consumption_constraint(consumption, wealth, labor_income):
+    return consumption <= wealth + labor_income
+
+
+# ======================================================================================
+# Model specification and parameters
+# ======================================================================================
+
+MODEL_CONFIG = {
+    "functions": {
+        "utility": utility,
+        "next_wealth": next_wealth,
+        "next_health": next_health,
+        "consumption_constraint": consumption_constraint,
+        "labor_income": labor_income,
+        "working": working,
+        "wage": wage,
+        "age": age,
+    },
+    "choices": {
+        "leisure": {"options": [0, 1]},
+        "consumption": {
+            "grid_type": "linspace",
+            "start": 1,
+            "stop": 100,
+            "n_points": N_GRID_POINTS["consumption"],
+        },
+        "exercise": {
+            "grid_type": "linspace",
+            "start": 0,
+            "stop": 1,
+            "n_points": N_GRID_POINTS["exercise"],
+        },
+    },
+    "states": {
+        "wealth": {
+            "grid_type": "linspace",
+            "start": 1,
+            "stop": 100,
+            "n_points": N_GRID_POINTS["wealth"],
+        },
+        "health": {
+            "grid_type": "linspace",
+            "start": 0,
+            "stop": 1,
+            "n_points": N_GRID_POINTS["health"],
+        },
+    },
+    "n_periods": RETIREMENT_AGE - 18,
+}
+
+PARAMS = {
+    "beta": 0.95,
+    "utility": {"disutility_of_work": 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
@@ -47,6 +49,9 @@ extend-ignore = [
     # use of `assert` detected
     "S101",
 
+    # `pickle` module is unsafe
+    "S301",
+
     # Private member accessed: `_stochastic_info`
     "SLF001",
 
@@ -57,6 +62,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"

src/lcm/_config.py

@@ -1,3 +1,3 @@
 from pathlib import Path
 
-TEST_DATA_PATH = Path(__file__).parent.parent.parent.resolve().joinpath("tests", "data")
+TEST_DATA = Path(__file__).parent.parent.parent.resolve().joinpath("tests", "data")