Add more examples to colab notebook

examples/pysr_demo.ipynb

@@ -711,6 +711,276 @@
         "plt.scatter(X[:, 0], y_prediction)\n"
       ]
     },
+    {
+      "attachments": {},
+      "cell_type": "markdown",
+      "metadata": {},
+      "source": [
+        "# Multiple outputs"
+      ]
+    },
+    {
+      "attachments": {},
+      "cell_type": "markdown",
+      "metadata": {},
+      "source": [
+        "For multiple outputs, multiple equations are returned:"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {},
+      "outputs": [],
+      "source": [
+        "X = 2 * np.random.randn(100, 5)\n",
+        "y = 1 / X[:, [0, 1, 2]]\n"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {},
+      "outputs": [],
+      "source": [
+        "model = PySRRegressor(\n",
+        "    binary_operators=[\"+\", \"*\"],\n",
+        "    unary_operators=[\"inv(x) = 1/x\"],\n",
+        "    extra_sympy_mappings={\"inv\": lambda x: 1/x},\n",
+        ")\n",
+        "model.fit(X, y)"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {},
+      "outputs": [],
+      "source": [
+        "model"
+      ]
+    },
+    {
+      "attachments": {},
+      "cell_type": "markdown",
+      "metadata": {},
+      "source": [
+        "# Julia packages and types"
+      ]
+    },
+    {
+      "attachments": {},
+      "cell_type": "markdown",
+      "metadata": {},
+      "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",
+        "Julia package.\n",
+        "For some tasks, it may be necessary to load such a package.\n",
+        "\n",
+        "For example, let's say we wish to discovery the following relationship:\n",
+        "\n",
+        "$$ y = p_{3x + 1} - 5, $$\n",
+        "\n",
+        "where $p_i$ is the $i$th prime number, and $x$ is the input feature.\n",
+        "\n",
+        "Let's see if we can discover this using\n",
+        "the [Primes.jl](https://github.com/JuliaMath/Primes.jl) package.\n",
+        "\n",
+        "First, let's get the Julia backend\n",
+        "(here, we manually specify 4 threads and `-O3` - although this will only work if PySR has not yet started):"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {},
+      "outputs": [],
+      "source": [
+        "import pysr\n",
+        "jl = pysr.julia_helpers.init_julia(julia_kwargs={\"threads\": 8, \"optimize\": 3})"
+      ]
+    },
+    {
+      "attachments": {},
+      "cell_type": "markdown",
+      "metadata": {},
+      "source": [
+        "\n",
+        "\n",
+        "`jl` stores the Julia runtime.\n",
+        "\n",
+        "Now, let's run some Julia code to add the Primes.jl\n",
+        "package to the PySR environment:"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {},
+      "outputs": [],
+      "source": [
+        "jl.eval(\"\"\"\n",
+        "import Pkg\n",
+        "Pkg.add(\"Primes\")\n",
+        "\"\"\")"
+      ]
+    },
+    {
+      "attachments": {},
+      "cell_type": "markdown",
+      "metadata": {},
+      "source": [
+        "This imports the Julia package manager, and uses it to install\n",
+        "`Primes.jl`. Now let's import `Primes.jl`:"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {},
+      "outputs": [],
+      "source": [
+        "jl.eval(\"import Primes\")"
+      ]
+    },
+    {
+      "attachments": {},
+      "cell_type": "markdown",
+      "metadata": {},
+      "source": [
+        "\n",
+        "Now, we define a custom operator:\n"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {},
+      "outputs": [],
+      "source": [
+        "jl.eval(\"\"\"\n",
+        "function p(i::T) where T\n",
+        "    if (0.5 < i < 1000)\n",
+        "        return T(Primes.prime(round(Int, i)))\n",
+        "    else\n",
+        "        return T(NaN)\n",
+        "    end\n",
+        "end\n",
+        "\"\"\")"
+      ]
+    },
+    {
+      "attachments": {},
+      "cell_type": "markdown",
+      "metadata": {},
+      "source": [
+        "\n",
+        "We have created a a function `p`, which takes an arbitrary number as input.\n",
+        "`p` first checks whether the input is between 0.5 and 1000.\n",
+        "If out-of-bounds, it returns `NaN`.\n",
+        "If in-bounds, it rounds it to the nearest integer, compures the corresponding prime number, and then\n",
+        "converts it to the same type as input.\n",
+        "\n",
+        "Next, let's generate a list of primes for our test dataset.\n",
+        "Since we are using PyJulia, we can just call `p` directly to do this:\n"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {},
+      "outputs": [],
+      "source": [
+        "primes = {i: jl.p(i*1.0) for i in range(1, 999)}"
+      ]
+    },
+    {
+      "attachments": {},
+      "cell_type": "markdown",
+      "metadata": {},
+      "source": [
+        "Next, let's use this list of primes to create a dataset of $x, y$ pairs:"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {},
+      "outputs": [],
+      "source": [
+        "import numpy as np\n",
+        "\n",
+        "X = np.random.randint(0, 100, 100)[:, None]\n",
+        "y = [primes[3*X[i, 0] + 1] - 5 + np.random.randn()*0.001 for i in range(100)]"
+      ]
+    },
+    {
+      "attachments": {},
+      "cell_type": "markdown",
+      "metadata": {},
+      "source": [
+        "Note that we have also added a tiny bit of noise to the dataset.\n",
+        "\n",
+        "Finally, let's create a PySR model, and pass the custom operator. We also need to define the sympy equivalent, which we can leave as a placeholder for now:"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {},
+      "outputs": [],
+      "source": [
+        "from pysr import PySRRegressor\n",
+        "import sympy\n",
+        "\n",
+        "class sympy_p(sympy.Function):\n",
+        "    pass\n",
+        "\n",
+        "model = PySRRegressor(\n",
+        "    binary_operators=[\"+\", \"-\", \"*\", \"/\"],\n",
+        "    unary_operators=[\"p\"],\n",
+        "    niterations=100,\n",
+        "    extra_sympy_mappings={\"p\": sympy_p}\n",
+        ")"
+      ]
+    },
+    {
+      "attachments": {},
+      "cell_type": "markdown",
+      "metadata": {},
+      "source": [
+        "We are all set to go! Let's see if we can find the true relation:"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {},
+      "outputs": [],
+      "source": [
+        "model.fit(X, y)"
+      ]
+    },
+    {
+      "attachments": {},
+      "cell_type": "markdown",
+      "metadata": {},
+      "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",
+        "You can get the sympy version of the best equation with:"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {},
+      "outputs": [],
+      "source": [
+        "model.sympy()"
+      ]
+    },
     {
       "cell_type": "markdown",
       "metadata": {