Update colab notebook to use juliaup

examples/pysr_demo.ipynb

@@ -19,7 +19,7 @@
     "## Instructions\n",
     "1. Work on a copy of this notebook: _File_ > _Save a copy in Drive_ (you will need a Google account).\n",
     "2. (Optional) If you would like to do the deep learning component of this tutorial, turn on the GPU with Edit->Notebook settings->Hardware accelerator->GPU\n",
-    "3. Execute the following cell (click on it and press Ctrl+Enter) to install Julia, IJulia and other packages (if needed, update `JULIA_VERSION` and the other parameters). This takes a couple of minutes.\n",
+    "3. Execute the following cell (click on it and press Ctrl+Enter) to install Julia. This may take a minute or so.\n",
     "4. Continue to the next section.\n",
     "\n",
     "_Notes_:\n",
@@ -40,33 +40,34 @@
    "cell_type": "code",
    "execution_count": null,
    "metadata": {
-    "id": "GIeFXS0F0zww"
+    "colab": {
+     "base_uri": "https://localhost:8080/"
+    },
+    "id": "GIeFXS0F0zww",
+    "outputId": "5399ed75-f77f-47c5-e53b-4b2f231f2839"
    },
    "outputs": [],
    "source": [
-    "%%shell\n",
-    "set -e\n",
-    "\n",
-    "#---------------------------------------------------#\n",
-    "JULIA_VERSION=\"1.8.5\"\n",
-    "export JULIA_PKG_PRECOMPILE_AUTO=0\n",
-    "#---------------------------------------------------#\n",
-    "\n",
-    "if [ -z `which julia` ]; then\n",
-    "  # Install Julia\n",
-    "  JULIA_VER=`cut -d '.' -f -2 <<< \"$JULIA_VERSION\"`\n",
-    "  echo \"Installing Julia $JULIA_VERSION on the current Colab Runtime...\"\n",
-    "  BASE_URL=\"https://julialang-s3.julialang.org/bin/linux/x64\"\n",
-    "  URL=\"$BASE_URL/$JULIA_VER/julia-$JULIA_VERSION-linux-x86_64.tar.gz\"\n",
-    "  wget -nv $URL -O /tmp/julia.tar.gz # -nv means \"not verbose\"\n",
-    "  tar -x -f /tmp/julia.tar.gz -C /usr/local --strip-components 1\n",
-    "  rm /tmp/julia.tar.gz\n",
-    "\n",
-    "  echo \"Installing PyCall.jl...\"\n",
-    "  julia -e 'using Pkg; Pkg.add(\"PyCall\"); Pkg.build(\"PyCall\")'\n",
-    "  julia -e 'println(\"Success\")'\n",
+    "!curl -fsSL https://install.julialang.org | sh -s -- -y --default-channel 1.10"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "colab": {
+     "base_uri": "https://localhost:8080/"
+    },
+    "id": "Iu9X-Y-YNmwM",
+    "outputId": "ee14af65-043a-4ad6-efa0-3cdcc48a4eb8"
+   },
+   "outputs": [],
+   "source": [
+    "# Make julia available on PATH:\n",
+    "!ln -s $HOME/.juliaup/bin/julia /usr/local/bin/julia\n",
     "\n",
-    "fi"
+    "# Test it works:\n",
+    "!julia --version"
    ]
   },
   {
@@ -75,18 +76,22 @@
     "id": "ORv1c6xvbDgV"
    },
    "source": [
-    "Install PySR and PyTorch-Lightning:"
+    "Install PySR"
    ]
   },
   {
    "cell_type": "code",
    "execution_count": null,
    "metadata": {
-    "id": "EhMRSZEYFPLz"
+    "colab": {
+     "base_uri": "https://localhost:8080/"
+    },
+    "id": "EhMRSZEYFPLz",
+    "outputId": "e3aad3cb-d921-473e-b77b-8fa6a3a9e2e8"
    },
    "outputs": [],
    "source": [
-    "%pip install -Uq pysr pytorch_lightning"
+    "!pip install pysr && python -m pysr install"
    ]
   },
   {
@@ -95,7 +100,7 @@
     "id": "etTMEV0wDqld"
    },
    "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"
+    "Colab's printing is non-standard, so we need to manually initialize Julia and redirect its printing. Normally, however, this is not required, and PySR will automatically start Julia during the first call to `.fit`:"
    ]
   },
   {
@@ -106,38 +111,25 @@
    },
    "outputs": [],
    "source": [
-    "from julia import Julia\n",
+    "def init_colab_printing():\n",
+    "    from pysr.julia_helpers import init_julia\n",
+    "    from julia.tools import redirect_output_streams\n",
     "\n",
-    "julia = Julia(compiled_modules=False, threads=\"auto\")\n",
-    "from julia import Main\n",
-    "from julia.tools import redirect_output_streams\n",
+    "    julia_kwargs = dict(optimize=3, threads=\"auto\", compiled_modules=False)\n",
+    "    init_julia(julia_kwargs=julia_kwargs)\n",
+    "    redirect_output_streams()\n",
     "\n",
-    "redirect_output_streams()"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {
-    "id": "6u2WhbVhht-G"
-   },
-   "source": [
-    "Let's install the backend of PySR, and all required libraries.\n",
     "\n",
-    "**(This may take some time)**"
+    "init_colab_printing()"
    ]
   },
   {
-   "cell_type": "code",
-   "execution_count": null,
+   "cell_type": "markdown",
    "metadata": {
-    "id": "J-0QbxyK1_51"
+    "id": "qeCPKd9wldEK"
    },
-   "outputs": [],
    "source": [
-    "import pysr\n",
-    "\n",
-    "# We don't precompile in colab because compiled modules are incompatible static Python libraries:\n",
-    "pysr.install(precompile=False)"
+    "Now, let's import all of our libraries:"
    ]
   },
   {
@@ -227,14 +219,19 @@
    "cell_type": "code",
    "execution_count": null,
    "metadata": {
-    "id": "p4PSrO-NK1Wa"
+    "colab": {
+     "base_uri": "https://localhost:8080/",
+     "height": 1000
+    },
+    "id": "p4PSrO-NK1Wa",
+    "outputId": "55910ab3-895d-400b-e9ce-c75aef639c68"
    },
    "outputs": [],
    "source": [
     "# Learn equations\n",
     "model = PySRRegressor(\n",
     "    niterations=30,\n",
-    "    binary_operators=[\"plus\", \"mult\"],\n",
+    "    binary_operators=[\"+\", \"-\"],\n",
     "    unary_operators=[\"cos\", \"exp\", \"sin\"],\n",
     "    **default_pysr_params\n",
     ")\n",
@@ -255,7 +252,12 @@
    "cell_type": "code",
    "execution_count": null,
    "metadata": {
-    "id": "4HR8gknlZz4W"
+    "colab": {
+     "base_uri": "https://localhost:8080/",
+     "height": 252
+    },
+    "id": "4HR8gknlZz4W",
+    "outputId": "496283bd-a743-4cc6-a2f9-9619ba91d870"
    },
    "outputs": [],
    "source": [
@@ -275,7 +277,12 @@
    "cell_type": "code",
    "execution_count": null,
    "metadata": {
-    "id": "IQKOohdpztS7"
+    "colab": {
+     "base_uri": "https://localhost:8080/",
+     "height": 38
+    },
+    "id": "IQKOohdpztS7",
+    "outputId": "0e7d058a-cce1-45ae-db94-6625f7e53a06"
    },
    "outputs": [],
    "source": [
@@ -295,7 +302,12 @@
    "cell_type": "code",
    "execution_count": null,
    "metadata": {
-    "id": "GRcxq-TTlpRX"
+    "colab": {
+     "base_uri": "https://localhost:8080/",
+     "height": 39
+    },
+    "id": "GRcxq-TTlpRX",
+    "outputId": "50bda367-1ed1-4860-8fcf-c940f2e4d935"
    },
    "outputs": [],
    "source": [
@@ -324,7 +336,12 @@
    "cell_type": "code",
    "execution_count": null,
    "metadata": {
-    "id": "HFGaNL6tbDgi"
+    "colab": {
+     "base_uri": "https://localhost:8080/",
+     "height": 35
+    },
+    "id": "HFGaNL6tbDgi",
+    "outputId": "0f364da5-e18d-4e31-cadf-087d641a3aed"
    },
    "outputs": [],
    "source": [
@@ -346,7 +363,11 @@
    "cell_type": "code",
    "execution_count": null,
    "metadata": {
-    "id": "Vbz4IMsk2NYH"
+    "colab": {
+     "base_uri": "https://localhost:8080/"
+    },
+    "id": "Vbz4IMsk2NYH",
+    "outputId": "361d4b6e-ac23-479d-b511-5001af05ca43"
    },
    "outputs": [],
    "source": [
@@ -406,14 +427,19 @@
    "cell_type": "code",
    "execution_count": null,
    "metadata": {
-    "id": "PoEkpvYuGUdy"
+    "colab": {
+     "base_uri": "https://localhost:8080/",
+     "height": 339
+    },
+    "id": "PoEkpvYuGUdy",
+    "outputId": "02834373-a054-400b-8247-2bf33a5c5beb"
    },
    "outputs": [],
    "source": [
     "model = PySRRegressor(\n",
     "    niterations=5,\n",
     "    populations=40,\n",
-    "    binary_operators=[\"plus\", \"mult\"],\n",
+    "    binary_operators=[\"+\", \"-\"],\n",
     "    unary_operators=[\"cos\", \"exp\", \"sin\", \"quart(x) = x^4\"],\n",
     "    extra_sympy_mappings={\"quart\": lambda x: x**4},\n",
     ")\n",
@@ -424,7 +450,12 @@
    "cell_type": "code",
    "execution_count": null,
    "metadata": {
-    "id": "emn2IajKbDgy"
+    "colab": {
+     "base_uri": "https://localhost:8080/",
+     "height": 38
+    },
+    "id": "emn2IajKbDgy",
+    "outputId": "11d5d3cf-de43-4f2b-f653-30016e09bdd0"
    },
    "outputs": [],
    "source": [
@@ -546,7 +577,12 @@
    "cell_type": "code",
    "execution_count": null,
    "metadata": {
-    "id": "sqMqb4nJ5ZR5"
+    "colab": {
+     "base_uri": "https://localhost:8080/",
+     "height": 467
+    },
+    "id": "sqMqb4nJ5ZR5",
+    "outputId": "aa24922b-2395-4e00-dce3-268fc8e603dc"
    },
    "outputs": [],
    "source": [
@@ -579,7 +615,11 @@
    "cell_type": "code",
    "execution_count": null,
    "metadata": {
-    "id": "v8WBYtcZbDhC"
+    "colab": {
+     "base_uri": "https://localhost:8080/"
+    },
+    "id": "v8WBYtcZbDhC",
+    "outputId": "37d4002f-e9d6-40c0-9a24-c671d9c384e6"
    },
    "outputs": [],
    "source": [
@@ -599,7 +639,11 @@
    "cell_type": "code",
    "execution_count": null,
    "metadata": {
-    "id": "a07K3KUjOxcp"
+    "colab": {
+     "base_uri": "https://localhost:8080/"
+    },
+    "id": "a07K3KUjOxcp",
+    "outputId": "41d11915-78b7-4446-c153-b92a5e2abd4c"
    },
    "outputs": [],
    "source": [
@@ -607,7 +651,7 @@
     "    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",
+    "    binary_operators=[\"+\", \"-\"],\n",
     "    unary_operators=[\"cos\"],\n",
     ")\n",
     "model.fit(X, y, weights=weights)"
@@ -688,17 +732,19 @@
    ]
   },
   {
-   "attachments": {},
    "cell_type": "markdown",
-   "metadata": {},
+   "metadata": {
+    "id": "2x-8M8W4G-KM"
+   },
    "source": [
     "# Multiple outputs"
    ]
   },
   {
-   "attachments": {},
    "cell_type": "markdown",
-   "metadata": {},
+   "metadata": {
+    "id": "LIJcWqBQG-KM"
+   },
    "source": [
     "For multiple outputs, multiple equations are returned:"
    ]
@@ -706,7 +752,9 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "metadata": {},
+   "metadata": {
+    "id": "_Aar1ZJwG-KM"
+   },
    "outputs": [],
    "source": [
     "X = 2 * np.random.randn(100, 5)\n",
@@ -716,7 +764,9 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "metadata": {},
+   "metadata": {
+    "id": "9Znwq40PG-KM"
+   },
    "outputs": [],
    "source": [
     "model = PySRRegressor(\n",
@@ -730,24 +780,28 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "metadata": {},
+   "metadata": {
+    "id": "0Y_vy0sqG-KM"
+   },
    "outputs": [],
    "source": [
     "model"
    ]
   },
   {
-   "attachments": {},
    "cell_type": "markdown",
-   "metadata": {},
+   "metadata": {
+    "id": "-UP49CsGG-KN"
+   },
    "source": [
     "# Julia packages and types"
    ]
   },
   {
-   "attachments": {},
    "cell_type": "markdown",
-   "metadata": {},
+   "metadata": {
+    "id": "tOdNHheUG-KN"
+   },
    "source": [
     "PySR uses [SymbolicRegression.jl](https://github.com/MilesCranmer/SymbolicRegression.jl)\n",
     "as its search backend. This is a pure Julia package, and so can interface easily with any other\n",
@@ -771,20 +825,23 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "metadata": {},
+   "metadata": {
+    "id": "yUC4BMuHG-KN"
+   },
    "outputs": [],
    "source": [
     "import pysr\n",
     "\n",
     "jl = pysr.julia_helpers.init_julia(\n",
-    "    julia_kwargs={\"threads\": \"auto\", \"optimize\": 2, \"compiled_modules\": False}\n",
+    "    julia_kwargs=dict(optimize=3, threads=\"auto\", compiled_modules=\"False\")\n",
     ")"
    ]
   },
   {
-   "attachments": {},
    "cell_type": "markdown",
-   "metadata": {},
+   "metadata": {
+    "id": "af07m4uBG-KN"
+   },
    "source": [
     "\n",
     "\n",
@@ -797,7 +854,9 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "metadata": {},
+   "metadata": {
+    "id": "xBlMY-s4G-KN"
+   },
    "outputs": [],
    "source": [
     "jl.eval(\n",
@@ -809,9 +868,10 @@
    ]
   },
   {
-   "attachments": {},
    "cell_type": "markdown",
-   "metadata": {},
+   "metadata": {
+    "id": "1rJFukD6G-KN"
+   },
    "source": [
     "This imports the Julia package manager, and uses it to install\n",
     "`Primes.jl`. Now let's import `Primes.jl`:"
@@ -820,16 +880,19 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "metadata": {},
+   "metadata": {
+    "id": "1PQl1rIaG-KN"
+   },
    "outputs": [],
    "source": [
     "jl.eval(\"import Primes\")"
    ]
   },
   {
-   "attachments": {},
    "cell_type": "markdown",
-   "metadata": {},
+   "metadata": {
+    "id": "edGdMxKnG-KN"
+   },
    "source": [
     "\n",
     "Now, we define a custom operator:\n"
@@ -838,7 +901,9 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "metadata": {},
+   "metadata": {
+    "id": "9Ut3HcW3G-KN"
+   },
    "outputs": [],
    "source": [
     "jl.eval(\n",
@@ -855,9 +920,10 @@
    ]
   },
   {
-   "attachments": {},
    "cell_type": "markdown",
-   "metadata": {},
+   "metadata": {
+    "id": "_wcV8889G-KN"
+   },
    "source": [
     "\n",
     "We have created a function `p`, which takes a number `i` of type `T` (e.g., `T=Float64`).\n",
@@ -887,16 +953,19 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "metadata": {},
+   "metadata": {
+    "id": "giqwisEPG-KN"
+   },
    "outputs": [],
    "source": [
     "primes = {i: jl.p(i * 1.0) for i in range(1, 999)}"
    ]
   },
   {
-   "attachments": {},
    "cell_type": "markdown",
-   "metadata": {},
+   "metadata": {
+    "id": "MPAqARj6G-KO"
+   },
    "source": [
     "Next, let's use this list of primes to create a dataset of $x, y$ pairs:"
    ]
@@ -904,7 +973,9 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "metadata": {},
+   "metadata": {
+    "id": "jab4tRRRG-KO"
+   },
    "outputs": [],
    "source": [
     "import numpy as np\n",
@@ -914,9 +985,10 @@
    ]
   },
   {
-   "attachments": {},
    "cell_type": "markdown",
-   "metadata": {},
+   "metadata": {
+    "id": "3eFgWrjcG-KO"
+   },
    "source": [
     "Note that we have also added a tiny bit of noise to the dataset.\n",
     "\n",
@@ -926,7 +998,9 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "metadata": {},
+   "metadata": {
+    "id": "pEYskM2_G-KO"
+   },
    "outputs": [],
    "source": [
     "from pysr import PySRRegressor\n",
@@ -947,8 +1021,9 @@
   },
   {
    "cell_type": "markdown",
-   "id": "ee30bd41",
-   "metadata": {},
+   "metadata": {
+    "id": "ee30bd41"
+   },
    "source": [
     "We are all set to go! Let's see if we can find the true relation:"
    ]
@@ -956,16 +1031,19 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "metadata": {},
+   "metadata": {
+    "id": "li-TB19iG-KO"
+   },
    "outputs": [],
    "source": [
     "model.fit(X, y)"
    ]
   },
   {
-   "attachments": {},
    "cell_type": "markdown",
-   "metadata": {},
+   "metadata": {
+    "id": "jwhTWZryG-KO"
+   },
    "source": [
     "if all works out, you should be able to see the true relation (note that the constant offset might not be exactly 1, since it is allowed to round to the nearest integer).\n",
     "\n",
@@ -975,7 +1053,9 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "metadata": {},
+   "metadata": {
+    "id": "bSlpX9xAG-KO"
+   },
    "outputs": [],
    "source": [
     "model.sympy()"
@@ -991,7 +1071,6 @@
    ]
   },
   {
-   "attachments": {},
    "cell_type": "markdown",
    "metadata": {
     "id": "3hS2kTAbbDhL"
@@ -1068,6 +1147,17 @@
     "> We import torch *after* already starting PyJulia. This is required due to interference between their C bindings. If you use torch, and then run PyJulia, you will likely hit a segfault. So keep this in mind for mixed deep learning + PyJulia/PySR workflows."
    ]
   },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "id": "k-Od8b9DlkHK"
+   },
+   "outputs": [],
+   "source": [
+    "!pip install pytorch_lightning"
+   ]
+  },
   {
    "cell_type": "code",
    "execution_count": null,
@@ -1083,7 +1173,7 @@
     "import pytorch_lightning as pl\n",
     "\n",
     "hidden = 128\n",
-    "total_steps = 30_000\n",
+    "total_steps = 100_000\n",
     "\n",
     "\n",
     "def mlp(size_in, size_out, act=nn.ReLU):\n",
@@ -1298,7 +1388,9 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "metadata": {},
+   "metadata": {
+    "id": "UX7Am6mZG-KT"
+   },
    "outputs": [],
    "source": [
     "nnet_recordings = {\n",
@@ -1316,9 +1408,10 @@
    ]
   },
   {
-   "attachments": {},
    "cell_type": "markdown",
-   "metadata": {},
+   "metadata": {
+    "id": "krhaNlwFG-KT"
+   },
    "source": [
     "We can now load the data, including after a crash (be sure to re-run the import cells at the top of this notebook, including the one that starts PyJulia)."
    ]
@@ -1326,7 +1419,9 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "metadata": {},
+   "metadata": {
+    "id": "NF9aSFXHG-KT"
+   },
    "outputs": [],
    "source": [
     "import pickle as pkl\n",
@@ -1339,9 +1434,10 @@
    ]
   },
   {
-   "attachments": {},
    "cell_type": "markdown",
-   "metadata": {},
+   "metadata": {
+    "id": "_hTYHhDGG-KT"
+   },
    "source": [
     "And now fit using a subsample of the data (symbolic regression only needs a small sample to find the best equation):"
    ]
@@ -1358,9 +1454,9 @@
     "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",
+    "    niterations=50,\n",
+    "    binary_operators=[\"+\", \"-\", \"*\"],\n",
+    "    unary_operators=[\"cos\", \"square\"],\n",
     ")\n",
     "model.fit(g_input[f_sample_idx], g_output[f_sample_idx])"
    ]
@@ -1384,14 +1480,13 @@
    ]
   },
   {
-   "attachments": {},
    "cell_type": "markdown",
    "metadata": {
     "id": "6WuaeqyqbDhe"
    },
    "source": [
     "Recall we are searching for $f$ and $g$ such that:\n",
-    "$$z=f(\\sum g(x_i))$$ \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",
@@ -1459,7 +1554,6 @@
  "metadata": {
   "accelerator": "GPU",
   "colab": {
-   "name": "pysr_demo.ipynb",
    "provenance": []
   },
   "gpuClass": "standard",