Update explanation notebooks

OpenSourceEconomics · Jul 11, 2024 · 2a1b955 · 2a1b955
1 parent d981304
commit 2a1b955
Show file tree

Hide file tree

Showing 2 changed files with 53 additions and 69 deletions.
diff --git a/explanations/dispatchers.ipynb b/explanations/dispatchers.ipynb
@@ -273,6 +273,9 @@
    "metadata": {},
    "outputs": [],
    "source": [
+    "from lcm import DiscreteGrid, LinspaceGrid, Model\n",
+    "\n",
+    "\n",
     "def utility(consumption, retirement, lagged_retirement, wealth):\n",
     "    working = 1 - retirement\n",
     "    retirement_habit = lagged_retirement * wealth\n",
@@ -283,33 +286,23 @@
     "    return jnp.logical_or(retirement == 1, lagged_retirement == 0)\n",
     "\n",
     "\n",
-    "MODEL_CONFIG = {\n",
-    "    \"functions\": {\n",
+    "model = Model(\n",
+    "    functions={\n",
     "        \"utility\": utility,\n",
     "        \"next_lagged_retirement\": lambda retirement: retirement,\n",
     "        \"next_wealth\": lambda wealth, consumption: wealth - consumption,\n",
     "        \"absorbing_retirement_filter\": absorbing_retirement_filter,\n",
     "    },\n",
-    "    \"choices\": {\n",
-    "        \"retirement\": {\"options\": [0, 1]},\n",
-    "        \"consumption\": {\n",
-    "            \"grid_type\": \"linspace\",\n",
-    "            \"start\": 1,\n",
-    "            \"stop\": 2,\n",
-    "            \"n_points\": 2,\n",
-    "        },\n",
+    "    n_periods=1,\n",
+    "    choices={\n",
+    "        \"retirement\": DiscreteGrid([0, 1]),\n",
+    "        \"consumption\": LinspaceGrid(start=1, stop=2, n_points=2),\n",
     "    },\n",
-    "    \"states\": {\n",
-    "        \"lagged_retirement\": {\"options\": [0, 1]},\n",
-    "        \"wealth\": {\n",
-    "            \"grid_type\": \"linspace\",\n",
-    "            \"start\": 1,\n",
-    "            \"stop\": 4,\n",
-    "            \"n_points\": 4,\n",
-    "        },\n",
+    "    states={\n",
+    "        \"lagged_retirement\": DiscreteGrid([0, 1]),\n",
+    "        \"wealth\": LinspaceGrid(start=1, stop=4, n_points=4),\n",
     "    },\n",
-    "    \"n_periods\": 1,\n",
-    "}"
+    ")"
    ]
   },
   {
@@ -321,10 +314,10 @@
     "from lcm.process_model import process_model\n",
     "from lcm.state_space import create_state_choice_space\n",
     "\n",
-    "model = process_model(user_model=MODEL_CONFIG)\n",
+    "processed_model = process_model(model)\n",
     "\n",
     "sc_space, space_info, state_indexer, segments = create_state_choice_space(\n",
-    "    model,\n",
+    "    processed_model,\n",
     "    period=2,\n",
     "    is_last_period=False,\n",
     "    jit_filter=False,\n",
@@ -807,7 +800,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.12.0"
+   "version": "3.12.3"
   }
  },
  "nbformat": 4,