Expand colab notebook

examples/pysr_demo.ipynb

@@ -1262,6 +1262,7 @@
       ]
     },
     {
+      "attachments": {},
       "cell_type": "markdown",
       "metadata": {
         "id": "nCCIvvAGuyFi"
@@ -1269,7 +1270,60 @@
       "source": [
         "## Learning over the network:\n",
         "\n",
-        "Now, let's fit `g` using PySR:"
+        "Now, let's fit `g` using PySR.\n",
+        "\n",
+        "> **Warning**\n",
+        ">\n",
+        "> First, let's save the data, because sometimes PyTorch and PyJulia's C bindings interfere and cause the colab kernel to crash. If we need to restart, we can just load the data without having to retrain the network:"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {},
+      "outputs": [],
+      "source": [
+        "nnet_recordings = {\n",
+        "    \"g_input\": X_for_pysr.detach().cpu().numpy().reshape(-1, 5),\n",
+        "    \"g_output\": y_i_for_pysr.detach().cpu().numpy().reshape(-1),\n",
+        "    \"f_input\": y_for_pysr.detach().cpu().numpy().reshape(-1, 1),\n",
+        "    \"f_output\": z_for_pysr.detach().cpu().numpy().reshape(-1),\n",
+        "}\n",
+        "\n",
+        "# Save the data for later use:\n",
+        "import pickle as pkl\n",
+        "\n",
+        "with open(\"nnet_recordings.pkl\", \"wb\") as f:\n",
+        "    pkl.dump(nnet_recordings, f)"
+      ]
+    },
+    {
+      "attachments": {},
+      "cell_type": "markdown",
+      "metadata": {},
+      "source": [
+        "We can now load the data:"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {},
+      "outputs": [],
+      "source": [
+        "nnet_recordings = pkl.load(open(\"nnet_recordings.pkl\", \"rb\"))\n",
+        "f_input = nnet_recordings[\"f_input\"]\n",
+        "f_output = nnet_recordings[\"f_output\"]\n",
+        "g_input = nnet_recordings[\"g_input\"]\n",
+        "g_output = nnet_recordings[\"g_output\"]"
+      ]
+    },
+    {
+      "attachments": {},
+      "cell_type": "markdown",
+      "metadata": {},
+      "source": [
+        "And now fit using a subsample of the data (symbolic regression only needs a small sample to find the best equation):"
       ]
     },
     {
@@ -1281,17 +1335,15 @@
       },
       "outputs": [],
       "source": [
-        "np.random.seed(1)\n",
-        "tmpX = X_for_pysr.detach().numpy().reshape(-1, 5)\n",
-        "tmpy = y_i_for_pysr.detach().numpy().reshape(-1)\n",
-        "idx2 = np.random.randint(0, tmpy.shape[0], size=500)\n",
+        "rstate = np.random.RandomState(0)\n",
+        "f_sample_idx = rstate.choice(f_input.shape[0], size=500, replace=False)\n",
         "\n",
         "model = PySRRegressor(\n",
         "    niterations=20,\n",
         "    binary_operators=[\"plus\", \"sub\", \"mult\"],\n",
         "    unary_operators=[\"cos\", \"square\", \"neg\"],\n",
         ")\n",
-        "model.fit(X=tmpX[idx2], y=tmpy[idx2])"
+        "model.fit(g_input[f_sample_idx], g_output[f_sample_idx])"
       ]
     },
     {
@@ -1310,9 +1362,12 @@
         "id": "6WuaeqyqbDhe"
       },
       "source": [
-        "Recall we are searching for $y_i$ above:\n",
+        "Recall we are searching for $f$ and $g$ such that:\n",
+        "$$z=f(\\sum g(x_i))$$ \n",
+        "which approximates the true relation:\n",
+        "$$ z = y^2,\\quad y = \\frac{1}{10} \\sum(y_i),\\quad y_i = x_{i0}^2 + 6 \\cos(2 x_{i2})$$\n",
         "\n",
-        "$$ z = y^2,\\quad y = \\frac{1}{10} \\sum(y_i),\\quad y_i = x_{i0}^2 + 6 \\cos(2 x_{i2})$$"
+        "Let's see how well we did in recovering $g$:"
       ]
     },
     {
@@ -1384,7 +1439,15 @@
       "name": "main_ipynb"
     },
     "language_info": {
+      "codemirror_mode": {
+        "name": "ipython",
+        "version": 3
+      },
+      "file_extension": ".py",
+      "mimetype": "text/x-python",
       "name": "python",
+      "nbconvert_exporter": "python",
+      "pygments_lexer": "ipython3",
       "version": "3.10.9"
     }
   },