Compare the runtime

Author: xuanguang-li

lectures/markov_asset.ipynb

@@ -3,16 +3,16 @@
     {
       "cell_type": "markdown",
       "metadata": {
-        "id": "view-in-github",
-        "colab_type": "text"
+        "colab_type": "text",
+        "id": "view-in-github"
       },
       "source": [
         "<a href=\"https://colab.research.google.com/github/QuantEcon/lecture-python.myst/blob/update_markov_asset/lectures/markov_asset.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
       ]
     },
     {
       "cell_type": "code",
-      "execution_count": 2,
+      "execution_count": 12,
       "id": "ce7fce98",
       "metadata": {
         "id": "ce7fce98"
@@ -30,7 +30,25 @@
     },
     {
       "cell_type": "code",
-      "execution_count": 3,
+      "execution_count": 13,
+      "id": "50d12e94",
+      "metadata": {},
+      "outputs": [
+        {
+          "name": "stdout",
+          "output_type": "stream",
+          "text": [
+            "[CpuDevice(id=0)]\n"
+          ]
+        }
+      ],
+      "source": [
+        "print(jax.devices())"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": 14,
       "id": "b7bcf31f",
       "metadata": {
         "id": "b7bcf31f"
@@ -135,7 +153,7 @@
     },
     {
       "cell_type": "code",
-      "execution_count": 4,
+      "execution_count": 15,
       "id": "29a2f7a6",
       "metadata": {
         "id": "29a2f7a6"
@@ -180,7 +198,7 @@
     },
     {
       "cell_type": "code",
-      "execution_count": 5,
+      "execution_count": 16,
       "id": "8cfa0f41",
       "metadata": {
         "id": "8cfa0f41"
@@ -220,7 +238,6 @@
         "\n",
         "    # Compute option price\n",
         "    p = consol_price(ap, ζ)\n",
-        "    err.throw()\n",
         "    n = M.shape[0]\n",
         "    w = jnp.zeros(n)\n",
         "    error = ϵ + 1\n",
@@ -247,37 +264,78 @@
     },
     {
       "cell_type": "code",
-      "execution_count": 6,
+      "execution_count": 17,
       "id": "7179ea30",
       "metadata": {
+        "id": "7179ea30",
         "mystnb": {
           "figure": {
             "caption": "Consol price and call option value\n",
             "name": "fig_consol_call"
           }
-        },
-        "id": "7179ea30"
+        }
       },
       "outputs": [],
       "source": [
         "ap = create_ap_model(β=0.9)\n",
         "ζ = 1.0\n",
         "strike_price = 40\n",
-        "\n",
         "x = ap.mc.state_values\n",
-        "err, p = consol_price_jit(ap, ζ)\n",
-        "err.throw()\n",
-        "err, w = call_option_jit(ap, ζ, strike_price)\n",
-        "err.throw()\n"
+        "\n",
+        "def timer_function():\n",
+        "  err, p = consol_price_jit(ap, ζ)\n",
+        "  err, w = call_option_jit(ap, ζ, strike_price)\n",
+        "\n",
+        "result = qe.timeit(\n",
+        "  timer_function, runs=1000, verbose=False, \n",
+        "  results=True, unit=\"milliseconds\"\n",
+        "  )"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": 18,
+      "id": "1ec35c54",
+      "metadata": {},
+      "outputs": [
+        {
+          "name": "stdout",
+          "output_type": "stream",
+          "text": [
+            "Results saved to timer_results_jax_cpu.json\n"
+          ]
+        }
+      ],
+      "source": [
+        "import json\n",
+        "\n",
+        "# result is already a dictionary, no need to call _asdict()\n",
+        "result_dict = result\n",
+        "\n",
+        "# Define the filename\n",
+        "filename = 'timer_results_jax_cpu.json'\n",
+        "\n",
+        "# Save the dictionary to a JSON file\n",
+        "with open(filename, 'w') as f:\n",
+        "    json.dump(result_dict, f, indent=4)\n",
+        "\n",
+        "print(f\"Results saved to {filename}\")"
       ]
     }
   ],
   "metadata": {
+    "accelerator": "GPU",
+    "colab": {
+      "gpuType": "T4",
+      "include_colab_link": true,
+      "provenance": []
+    },
     "jupytext": {
       "default_lexer": "ipython"
     },
     "kernelspec": {
-      "display_name": "Python 3",
+      "display_name": "quantecon",
+      "language": "python",
       "name": "python3"
     },
     "language_info": {
@@ -291,14 +349,8 @@
       "nbconvert_exporter": "python",
       "pygments_lexer": "ipython3",
       "version": "3.13.5"
-    },
-    "colab": {
-      "provenance": [],
-      "gpuType": "T4",
-      "include_colab_link": true
-    },
-    "accelerator": "GPU"
+    }
   },
   "nbformat": 4,
   "nbformat_minor": 5
-}
+}