Add example model with very large state-choice space (#54)

environment.yml

@@ -34,6 +34,7 @@ dependencies:
   - pre-commit
   - pydot
   - snakeviz
+  - memory_profiler
 
   # Documentation
   - sphinx

src/lcm/example_models/testing_example_models.py

@@ -0,0 +1,95 @@
+"""Define example model specifications."""
+
+import jax.numpy as jnp
+
+RETIREMENT_AGE = 65
+N_CHOICE_GRID_POINTS = 200
+N_STATE_GRID_POINTS = 100
+
+
+def phelps_deaton_utility(consumption, working, health, sport, delta):
+    return jnp.log(consumption) - (delta - health) * working - sport
+
+
+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 next_health(health, sport, working):
+    return health * (1 + sport - working / 2)
+
+
+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,
+        "next_health": next_health,
+    },
+    "choices": {
+        "retirement": {"options": [0, 1]},
+        "consumption": {
+            "grid_type": "linspace",
+            "start": 1,
+            "stop": 100,
+            "n_points": N_CHOICE_GRID_POINTS,
+        },
+        "sport": {
+            "grid_type": "linspace",
+            "start": 0,
+            "stop": 1,
+            "n_points": N_CHOICE_GRID_POINTS,
+        },
+    },
+    "states": {
+        "wealth": {
+            "grid_type": "linspace",
+            "start": 1,
+            "stop": 100,
+            "n_points": N_STATE_GRID_POINTS,
+        },
+        "health": {
+            "grid_type": "linspace",
+            "start": 0,
+            "stop": 1,
+            "n_points": N_STATE_GRID_POINTS,
+        },
+    },
+    "n_periods": RETIREMENT_AGE - 18,
+}
+
+PARAMS = {
+    "beta": 0.95,
+    "utility": {"delta": 0.05},
+    "next_wealth": {"interest_rate": 0.05},
+}

src/lcm/get_model.py

@@ -4,7 +4,10 @@
 
 from pybaum import tree_update
 
-from lcm.example_models import PHELPS_DEATON, PHELPS_DEATON_WITH_FILTERS
+from lcm.example_models.basic_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.basic_example_models import (
     PHELPS_DEATON,
     PHELPS_DEATON_FULLY_DISCRETE,
     PHELPS_DEATON_WITH_FILTERS,

tests/test_long.py

@@ -0,0 +1,13 @@
+import pytest
+from lcm.entry_point import get_lcm_function
+from lcm.example_models.testing_example_models import PARAMS, PHELPS_DEATON
+
+SKIP_REASON = """The test is designed to run approximately 1 minute on a standard
+laptop, such that we can differentiate the performance of running LCM on a GPU versus
+on the CPU."""
+
+
+@pytest.mark.skip(reason=SKIP_REASON)
+def test_long():
+    solve_model, template = get_lcm_function(PHELPS_DEATON, targets="solve")
+    solve_model(PARAMS)

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.basic_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.basic_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.basic_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.basic_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.basic_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.stochastic_example_models import MODEL, PARAMS
 
 # ======================================================================================
 # Simulate