Clean up and format example

examples/pysr_demo.ipynb

@@ -33,9 +33,7 @@
         "id": "COndi88gbDgO"
       },
       "source": [
-        "**Run the following code if you need to install Julia**\n",
-        "\n",
-        "**(select all lines -> Option-/)**"
+        "**Run the following code to install Julia**"
       ]
     },
     {
@@ -95,7 +93,14 @@
       },
       "outputs": [],
       "source": [
-        "!pip install -Uq pysr pytorch_lightning"
+        "%pip install -Uq pysr pytorch_lightning"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {},
+      "source": [
+        "The following step is not normally required, but colab's printing is non-standard and we need to manually set it up PyJulia:\n"
       ]
     },
     {
@@ -106,13 +111,13 @@
       },
       "outputs": [],
       "source": [
-        "# Required to get printing from Julia working in colab \n",
-        "# (you don't need to normally do this)\n",
         "from julia import Julia\n",
+        "\n",
         "julia = Julia(compiled_modules=False)\n",
         "from julia import Main\n",
         "from julia.tools import redirect_output_streams\n",
-        "redirect_output_streams()"
+        "\n",
+        "redirect_output_streams()\n"
       ]
     },
     {
@@ -123,7 +128,7 @@
       "source": [
         "Let's install the backend of PySR, and all required libraries. We will also precompile them so they are faster at startup.\n",
         "\n",
-        "**This will take ~5 minutes.**"
+        "**(This may take some time)**"
       ]
     },
     {
@@ -139,7 +144,8 @@
       "outputs": [],
       "source": [
         "import pysr\n",
-        "pysr.install()"
+        "\n",
+        "pysr.install()\n"
       ]
     },
     {
@@ -159,7 +165,7 @@
         "from torch.nn import functional as F\n",
         "from torch.utils.data import DataLoader, TensorDataset\n",
         "import pytorch_lightning as pl\n",
-        "from sklearn.model_selection import train_test_split"
+        "from sklearn.model_selection import train_test_split\n"
       ]
     },
     {
@@ -192,8 +198,8 @@
       "source": [
         "# Dataset\n",
         "np.random.seed(0)\n",
-        "X = 2*np.random.randn(100, 5)\n",
-        "y = 2.5382*np.cos(X[:, 3]) + X[:, 0]**2 - 2"
+        "X = 2 * np.random.randn(100, 5)\n",
+        "y = 2.5382 * np.cos(X[:, 3]) + X[:, 0] ** 2 - 2\n"
       ]
     },
     {
@@ -217,7 +223,7 @@
         "    populations=30,\n",
         "    procs=4,\n",
         "    model_selection=\"best\",\n",
-        ")"
+        ")\n"
       ]
     },
     {
@@ -246,12 +252,13 @@
       "source": [
         "# Learn equations\n",
         "model = PySRRegressor(\n",
-        "        niterations=30,\n",
-        "        binary_operators=[\"plus\", \"mult\"],\n",
-        "        unary_operators=[\"cos\", \"exp\", \"sin\"],\n",
-        "        **default_pysr_params)\n",
+        "    niterations=30,\n",
+        "    binary_operators=[\"plus\", \"mult\"],\n",
+        "    unary_operators=[\"cos\", \"exp\", \"sin\"],\n",
+        "    **default_pysr_params\n",
+        ")\n",
         "\n",
-        "model.fit(X, y)"
+        "model.fit(X, y)\n"
       ]
     },
     {
@@ -275,7 +282,7 @@
       },
       "outputs": [],
       "source": [
-        "model"
+        "model\n"
       ]
     },
     {
@@ -300,7 +307,7 @@
       },
       "outputs": [],
       "source": [
-        "model.sympy()"
+        "model.sympy()\n"
       ]
     },
     {
@@ -309,7 +316,7 @@
         "id": "EHIIPlmClltn"
       },
       "source": [
-        "We can also view the SymPy of any other expression in the list, using the index of it in `model.equation`."
+        "We can also view the SymPy of any other expression in the list, using the index of it in `model.equations_`."
       ]
     },
     {
@@ -325,7 +332,7 @@
       },
       "outputs": [],
       "source": [
-        "model.sympy(2)"
+        "model.sympy(2)\n"
       ]
     },
     {
@@ -359,7 +366,7 @@
       },
       "outputs": [],
       "source": [
-        "model.latex()"
+        "model.latex()\n"
       ]
     },
     {
@@ -389,7 +396,7 @@
         "ypredict_simpler = model.predict(X, 2)\n",
         "\n",
         "print(\"Default selection MSE:\", np.power(ypredict - y, 2).mean())\n",
-        "print(\"Manual selection MSE for index 2:\", np.power(ypredict_simpler - y, 2).mean())"
+        "print(\"Manual selection MSE for index 2:\", np.power(ypredict_simpler - y, 2).mean())\n"
       ]
     },
     {
@@ -423,7 +430,7 @@
       },
       "outputs": [],
       "source": [
-        "y = X[:, 0]**4 - 2"
+        "y = X[:, 0] ** 4 - 2\n"
       ]
     },
     {
@@ -451,12 +458,13 @@
       "outputs": [],
       "source": [
         "model = PySRRegressor(\n",
-        "        niterations=5, populations=40,\n",
-        "        binary_operators=[\"plus\", \"mult\"],\n",
-        "        unary_operators=[\"cos\", \"exp\", \"sin\",\n",
-        "                         \"quart(x) = x^4\"],\n",
-        "        extra_sympy_mappings={\"quart\": lambda x: x**4})\n",
-        "model.fit(X, y)                  "
+        "    niterations=5,\n",
+        "    populations=40,\n",
+        "    binary_operators=[\"plus\", \"mult\"],\n",
+        "    unary_operators=[\"cos\", \"exp\", \"sin\", \"quart(x) = x^4\"],\n",
+        "    extra_sympy_mappings={\"quart\": lambda x: x**4},\n",
+        ")\n",
+        "model.fit(X, y)\n"
       ]
     },
     {
@@ -472,7 +480,7 @@
       },
       "outputs": [],
       "source": [
-        "model.sympy()"
+        "model.sympy()\n"
       ]
     },
     {
@@ -571,10 +579,10 @@
         "np.random.seed(0)\n",
         "N = 3000\n",
         "upper_sigma = 5\n",
-        "X = 2*np.random.rand(N, 5)\n",
-        "sigma = np.random.rand(N)*(5-0.1) + 0.1\n",
-        "eps = sigma*np.random.randn(N)\n",
-        "y = 5*np.cos(3.5*X[:, 0]) - 1.3 + eps"
+        "X = 2 * np.random.rand(N, 5)\n",
+        "sigma = np.random.rand(N) * (5 - 0.1) + 0.1\n",
+        "eps = sigma * np.random.randn(N)\n",
+        "y = 5 * np.cos(3.5 * X[:, 0]) - 1.3 + eps\n"
       ]
     },
     {
@@ -601,7 +609,7 @@
       "source": [
         "plt.scatter(X[:, 0], y, alpha=0.2)\n",
         "plt.xlabel(\"$x_0$\")\n",
-        "plt.ylabel(\"$y$\")"
+        "plt.ylabel(\"$y$\")\n"
       ]
     },
     {
@@ -621,7 +629,7 @@
       },
       "outputs": [],
       "source": [
-        "weights = 1/sigma**2"
+        "weights = 1 / sigma[:, None] ** 2\n"
       ]
     },
     {
@@ -636,7 +644,7 @@
       },
       "outputs": [],
       "source": [
-        "weights[:5]"
+        "weights[:5, 0]\n"
       ]
     },
     {
@@ -662,13 +670,13 @@
       "outputs": [],
       "source": [
         "model = PySRRegressor(\n",
-        "    loss='myloss(x, y, w) = w * abs(x - y)',  # Custom loss function with weights.\n",
+        "    loss=\"myloss(x, y, w) = w * abs(x - y)\",  # Custom loss function with weights.\n",
         "    niterations=20,\n",
         "    populations=20,  # Use more populations\n",
         "    binary_operators=[\"plus\", \"mult\"],\n",
-        "    unary_operators=[\"cos\"]\n",
+        "    unary_operators=[\"cos\"],\n",
         ")\n",
-        "model.fit(X, y, weights=weights)"
+        "model.fit(X, y, weights=weights)\n"
       ]
     },
     {
@@ -692,7 +700,7 @@
       },
       "outputs": [],
       "source": [
-        "model"
+        "model\n"
       ]
     },
     {
@@ -717,8 +725,10 @@
       },
       "outputs": [],
       "source": [
-        "best_idx = model.equations_.query(f\"loss < {2 * model.equations_.loss.min()}\").score.idxmax()\n",
-        "model.sympy(best_idx)"
+        "best_idx = model.equations_.query(\n",
+        "    f\"loss < {2 * model.equations_.loss.min()}\"\n",
+        ").score.idxmax()\n",
+        "model.sympy(best_idx)\n"
       ]
     },
     {
@@ -753,7 +763,8 @@
       "outputs": [],
       "source": [
         "plt.scatter(X[:, 0], y, alpha=0.1)\n",
-        "plt.scatter(X[:, 0], model.predict(X, index=best_idx));"
+        "y_prediction = model.predict(X, index=best_idx)\n",
+        "plt.scatter(X[:, 0], y_prediction)\n"
       ]
     },
     {
@@ -800,10 +811,10 @@
         "N = 100000\n",
         "Nt = 100\n",
         "X = 6 * np.random.rand(N, Nt, 5) - 3\n",
-        "y_i = X[..., 0]**2 + 6*np.cos(2*X[..., 2])\n",
-        "y = np.sum(y_i, axis=1)/y_i.shape[1]\n",
+        "y_i = X[..., 0] ** 2 + 6 * np.cos(2 * X[..., 2])\n",
+        "y = np.sum(y_i, axis=1) / y_i.shape[1]\n",
         "z = y**2\n",
-        "X.shape, y.shape"
+        "X.shape, y.shape\n"
       ]
     },
     {
@@ -843,6 +854,7 @@
         "hidden = 128\n",
         "total_steps = 50000\n",
         "\n",
+        "\n",
         "def mlp(size_in, size_out, act=nn.ReLU):\n",
         "    return nn.Sequential(\n",
         "        nn.Linear(size_in, hidden),\n",
@@ -851,13 +863,14 @@
         "        act(),\n",
         "        nn.Linear(hidden, hidden),\n",
         "        act(),\n",
-        "        nn.Linear(hidden, size_out))\n",
+        "        nn.Linear(hidden, size_out),\n",
+        "    )\n",
         "\n",
         "\n",
         "class SumNet(pl.LightningModule):\n",
         "    def __init__(self):\n",
         "        super().__init__()\n",
-        "        \n",
+        "\n",
         "    ########################################################\n",
         "    # The same inductive bias as above!\n",
         "        self.g = mlp(5, 1)\n",
@@ -865,11 +878,12 @@
         "\n",
         "    def forward(self, x):\n",
         "        y_i = self.g(x)[:, :, 0]\n",
-        "        y = torch.sum(y_i, dim=1, keepdim=True)/y_i.shape[1]\n",
+        "        y = torch.sum(y_i, dim=1, keepdim=True) / y_i.shape[1]\n",
         "        z = self.f(y)\n",
         "        return z[:, 0]\n",
+        "\n",
         "    ########################################################\n",
-        "    \n",
+        "\n",
         "    # PyTorch Lightning bookkeeping:\n",
         "    def training_step(self, batch, batch_idx):\n",
         "        x, z = batch\n",
@@ -882,14 +896,18 @@
         "\n",
         "    def configure_optimizers(self):\n",
         "        self.trainer.reset_train_dataloader()\n",
-        "        # self.train_dataloader.loaders  # access it here.\n",
         "\n",
         "        optimizer = torch.optim.Adam(self.parameters(), lr=self.max_lr)\n",
-        "        scheduler = {'scheduler': torch.optim.lr_scheduler.OneCycleLR(optimizer, max_lr=self.max_lr,\n",
-        "                                             total_steps=self.total_steps,\n",
-        "                                             final_div_factor=1e4),\n",
-        "                     'interval': 'step'}\n",
-        "        return [optimizer], [scheduler]"
+        "        scheduler = {\n",
+        "            \"scheduler\": torch.optim.lr_scheduler.OneCycleLR(\n",
+        "                optimizer,\n",
+        "                max_lr=self.max_lr,\n",
+        "                total_steps=self.total_steps,\n",
+        "                final_div_factor=1e4,\n",
+        "            ),\n",
+        "            \"interval\": \"step\",\n",
+        "        }\n",
+        "        return [optimizer], [scheduler]\n"
       ]
     },
     {
@@ -924,7 +942,7 @@
         "train_set = TensorDataset(X_train, z_train)\n",
         "train = DataLoader(train_set, batch_size=128, num_workers=2)\n",
         "test_set = TensorDataset(X_test, z_test)\n",
-        "test = DataLoader(test_set, batch_size=256, num_workers=2)"
+        "test = DataLoader(test_set, batch_size=256, num_workers=2)\n"
       ]
     },
     {
@@ -960,7 +978,7 @@
         "pl.seed_everything(0)\n",
         "model = SumNet()\n",
         "model.total_steps = total_steps\n",
-        "model.max_lr = 1e-2"
+        "model.max_lr = 1e-2\n"
       ]
     },
     {
@@ -984,7 +1002,7 @@
       },
       "outputs": [],
       "source": [
-        "trainer = pl.Trainer(max_steps=total_steps, gpus=1, benchmark=True)"
+        "trainer = pl.Trainer(max_steps=total_steps, gpus=1, benchmark=True)\n"
       ]
     },
     {
@@ -1033,7 +1051,7 @@
       },
       "outputs": [],
       "source": [
-        "trainer.fit(model, train_dataloaders=train, val_dataloaders=test)"
+        "trainer.fit(model, train_dataloaders=train, val_dataloaders=test)\n"
       ]
     },
     {
@@ -1064,10 +1082,10 @@
         "\n",
         "X_for_pysr = Xt[idx]\n",
         "y_i_for_pysr = model.g(X_for_pysr)[:, :, 0]\n",
-        "y_for_pysr = torch.sum(y_i_for_pysr, dim=1)/y_i_for_pysr.shape[1]\n",
-        "z_for_pysr = zt[idx] #Use true values.\n",
+        "y_for_pysr = torch.sum(y_i_for_pysr, dim=1) / y_i_for_pysr.shape[1]\n",
+        "z_for_pysr = zt[idx]  # Use true values.\n",
         "\n",
-        "X_for_pysr.shape, y_i_for_pysr.shape"
+        "X_for_pysr.shape, y_i_for_pysr.shape\n"
       ]
     },
     {
@@ -1102,9 +1120,9 @@
         "model = PySRRegressor(\n",
         "    niterations=20,\n",
         "    binary_operators=[\"plus\", \"sub\", \"mult\"],\n",
-        "    unary_operators=[\"cos\", \"square\", \"neg\"]\n",
+        "    unary_operators=[\"cos\", \"square\", \"neg\"],\n",
         ")\n",
-        "model.fit(X=tmpX[idx2], y=tmpy[idx2])"
+        "model.fit(X=tmpX[idx2], y=tmpy[idx2])\n"
       ]
     },
     {
@@ -1135,7 +1153,7 @@
       },
       "outputs": [],
       "source": [
-        "model"
+        "model\n"
       ]
     },
     {