Create example models folder

src/lcm/example_models/example_models_long.py

@@ -0,0 +1,88 @@
+"""Define example model specifications."""
+import jax.numpy as jnp
+
+RETIREMENT_AGE = 65
+N_CHOICE_GRID_POINTS = 5_000
+N_STATE_GRID_POINTS = 5_000
+
+
+def phelps_deaton_utility(consumption, working, delta):
+    return jnp.log(consumption) - delta * working
+
+
+def phelps_deaton_utility_with_filter(
+    consumption,
+    working,
+    delta,
+    lagged_retirement,  # noqa: ARG001
+):
+    return jnp.log(consumption) - delta * working
+
+
+def working(retirement):
+    return 1 - retirement
+
+
+def next_wealth_with_shock(
+    wealth,
+    consumption,
+    working,
+    wage,
+    wage_shock,
+    interest_rate,
+):
+    return interest_rate * (wealth - consumption) + wage * wage_shock * working
+
+
+def next_wealth(wealth, consumption, working, wage, interest_rate):
+    return (1 + interest_rate) * (wealth - consumption) + wage * working
+
+
+def consumption_constraint(consumption, wealth):
+    return consumption <= wealth
+
+
+def wage(age):
+    return 1 + 0.1 * age
+
+
+def age(_period):
+    return _period + 18
+
+
+PHELPS_DEATON = {
+    "functions": {
+        "utility": phelps_deaton_utility,
+        "next_wealth": next_wealth,
+        "consumption_constraint": consumption_constraint,
+        "working": working,
+        "wage": wage,
+        "age": age,
+    },
+    "choices": {
+        "retirement": {"options": [0, 1]},
+        "consumption": {
+            "grid_type": "linspace",
+            "start": 0,
+            "stop": 100,
+            "n_points": N_CHOICE_GRID_POINTS,
+        },
+    },
+    "states": {
+        "wealth": {
+            "grid_type": "linspace",
+            "start": 0,
+            "stop": 100,
+            "n_points": N_STATE_GRID_POINTS,
+        },
+        "health": {
+            "grid_type": "linspace",
+            "start": 0,
+            "stop": 100,
+            "n_points": N_STATE_GRID_POINTS,
+        },
+    },
+    "n_periods": RETIREMENT_AGE - 18,
+}
+
+PARAMS = {}

src/lcm/get_model.py

@@ -3,7 +3,7 @@
 
 from pybaum import tree_update
 
-from lcm.example_models import PHELPS_DEATON, PHELPS_DEATON_WITH_FILTERS
+from lcm.example_models.example_models import PHELPS_DEATON, PHELPS_DEATON_WITH_FILTERS
 
 
 class ModelAndParams(NamedTuple):

tests/test_entry_point.py

@@ -5,7 +5,7 @@
     create_compute_conditional_continuation_value,
     get_lcm_function,
 )
-from lcm.example_models import (
+from lcm.example_models.example_models import (
     PHELPS_DEATON,
     PHELPS_DEATON_FULLY_DISCRETE,
     PHELPS_DEATON_WITH_FILTERS,

tests/test_model_functions.py

@@ -1,6 +1,6 @@
 import jax.numpy as jnp
 import pandas as pd
-from lcm.example_models import PHELPS_DEATON, phelps_deaton_utility
+from lcm.example_models.example_models import PHELPS_DEATON, phelps_deaton_utility
 from lcm.interfaces import Model
 from lcm.model_functions import (
     get_combined_constraint,

tests/test_next_state.py

@@ -1,6 +1,6 @@
 import jax.numpy as jnp
 import pandas as pd
-from lcm.example_models import PHELPS_DEATON
+from lcm.example_models.example_models import PHELPS_DEATON
 from lcm.interfaces import Model
 from lcm.next_state import _get_stochastic_next_func, get_next_state_function
 from lcm.process_model import process_model

tests/test_process_model.py

@@ -4,7 +4,7 @@
 import numpy as np
 import pandas as pd
 import pytest
-from lcm.example_models import (
+from lcm.example_models.example_models import (
     N_CHOICE_GRID_POINTS,
     N_STATE_GRID_POINTS,
     PHELPS_DEATON,

tests/test_simulate.py

@@ -6,7 +6,7 @@
     create_compute_conditional_continuation_policy,
     get_lcm_function,
 )
-from lcm.example_models import (
+from lcm.example_models.example_models import (
     N_CHOICE_GRID_POINTS,
     PHELPS_DEATON,
     PHELPS_DEATON_WITH_FILTERS,

tests/test_state_space.py

@@ -2,7 +2,7 @@
 import numpy as np
 import pandas as pd
 import pytest
-from lcm.example_models import PHELPS_DEATON_WITH_FILTERS
+from lcm.example_models.example_models import PHELPS_DEATON_WITH_FILTERS
 from lcm.interfaces import Model
 from lcm.process_model import process_model
 from lcm.state_space import (

tests/test_stochastic.py

@@ -5,7 +5,7 @@
 from lcm.entry_point import (
     get_lcm_function,
 )
-from lcm.example_models_stochastic import MODEL, PARAMS
+from lcm.example_models.example_models_stochastic import MODEL, PARAMS
 
 # ======================================================================================
 # Simulate