DemoMSSP.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "source": [
        "# MSSP: Multi-Set Symbolic Skeleton Prediction for Symbolic Regression\n",
        "\n",
        "## Installation\n",
        "\n",
        "Execute `!pip install git+https://github.com/NISL-MSU/MultiSetSR`\n",
        "\n",
        "**IMPORTANT:** This code is implemented using Pytorch and CUDA. If you're running this on Google Colab, change the runtime type to GPU."
      ],
      "metadata": {
        "collapsed": false,
        "id": "FSxK4oaVIzKd"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "!pip install -q git+https://github.com/NISL-MSU/MultiSetSR\n",
        "import warnings\n",
        "warnings.filterwarnings(\"ignore\")"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "Y__1Bwv3KSq3",
        "outputId": "9ab2cb89-6148-49ff-93cb-bd45fd25ee3d"
      },
      "execution_count": 1,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Found existing installation: MultiSetSR 0.0.1\n",
            "Uninstalling MultiSetSR-0.0.1:\n",
            "  Successfully uninstalled MultiSetSR-0.0.1\n",
            "  Installing build dependencies ... \u001b[?25l\u001b[?25hdone\n",
            "  Getting requirements to build wheel ... \u001b[?25l\u001b[?25hdone\n",
            "  Preparing metadata (pyproject.toml) ... \u001b[?25l\u001b[?25hdone\n",
            "  Building wheel for MultiSetSR (pyproject.toml) ... \u001b[?25l\u001b[?25hdone\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Example using pre-determined datasets\n",
        "\n",
        "In this example, we will predict the symbolic skeletons corresponding to each variable of a system whose underlying equation is one of the following:\n",
        "\n",
        "<br>\n",
        "\n",
        "| Eq. | Underlying equation________________________________________________________|\n",
        "|-----|------------------------|\n",
        "| E1  | $ (3.0375 x_1 x_2 + 5.5 \\sin (9/4 (x_1 - 2/3)(x_2 - 2/3)))/5 $|\n",
        "| E2  | $ 5.5 + (1- x_1/4) ^ 2 + \\sqrt{x_2 + 10} \\sin( x_3/5)$|\n",
        "| E3  | $(1.5 e^{1.5  x_1} + 5 \\cos(3 x_2)) / 10$|\n",
        "| E4  | $((1- x_1)^2 + (1- x_3) ^ 2 + 100 (x_2 - x_1 ^ 2) ^ 2 + 100 (x_4 - x_3 ^ 2) ^ 2)/10000$|\n",
        "| E5  | $\\sin(x_1 + x_2 x_3) + \\exp{(1.2  x_4)}$|\n",
        "| E6  | $\\tanh(x_1 / 2) + \\text{abs}(x_2) \\cos(x_3^2/5)$|\n",
        "| E7  | $(1 - x_2^2) / (\\sin(2 \\pi \\, x_1) + 1.5)$|\n",
        "| E8  | $x_1^4 / (x_1^4 + 1) + x_2^4 / (x_2^4 + 1)$|\n",
        "| E9  | $\\log(2 x_2 + 1) - \\log(4 x_1 ^ 2 + 1)$|\n",
        "| E10 | $\\sin(x_1 \\, e^{x_2})$|\n",
        "| E11 | $x_1 \\, \\log(x_2 ^ 4)$|\n",
        "| E12 | $1 + x_1 \\, \\sin(1 / x_2)$|\n",
        "| E13 | $\\sqrt{x_1}\\, \\log(x_2 ^ 2)$|"
      ],
      "metadata": {
        "id": "pOyWruXjKWAK"
      }
    },
    {
      "cell_type": "code",
      "execution_count": 2,
      "outputs": [],
      "source": [
        "from EquationLearning.SymbolicRegressor.MSSP import *\n",
        "\n",
        "datasetName = 'E6'\n",
        "data_loader = DataLoader(name=datasetName)\n",
        "data = data_loader.dataset"
      ],
      "metadata": {
        "id": "gPjx1EfiIzKf"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "**Define NN and load weights**\n",
        "\n",
        "For this example, we have already trained a feedforward neural network on the generated dataset so we only load their corresponding weights."
      ],
      "metadata": {
        "collapsed": false,
        "id": "UJddKsl-IzKg"
      }
    },
    {
      "cell_type": "code",
      "execution_count": 3,
      "outputs": [],
      "source": [
        "device = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")\n",
        "root = get_project_root()\n",
        "folder = os.path.join(root, \"EquationLearning//saved_models//saved_NNs//\" + datasetName)\n",
        "filepath = folder + \"//weights-NN-\" + datasetName\n",
        "nn_model = NNModel(device=device, n_features=data.n_features, NNtype=data_loader.modelType)\n",
        "nn_model.loadModel(filepath)"
      ],
      "metadata": {
        "id": "p_HcTwFlIzKg"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "**Get Skeletons**\n",
        "\n",
        "The following method will generate some candidate symbolic skeletons and select the most appropriate for each variable"
      ],
      "metadata": {
        "collapsed": false,
        "id": "BZDdLquCIzKh"
      }
    },
    {
      "cell_type": "code",
      "execution_count": 4,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "********************************\n",
            "Analyzing variable x0\n",
            "********************************\n",
            "Predicted skeleton 1 for variable x0: c*tanh(c*x0) + c\n",
            "Predicted skeleton 2 for variable x0: c*tanh(c*x0 + c) + c\n",
            "Predicted skeleton 3 for variable x0: c + tanh(c*x0)\n",
            "Predicted skeleton 4 for variable x0: c + tanh(c*x0 + c)\n",
            "Predicted skeleton 5 for variable x0: c*sqrt(c*tanh(c*x0) + c) + c\n",
            "\n",
            "Choosing the best skeleton... (skeletons ordered based on number of nodes)\n",
            "\tSkeleton: c + tanh(c*x0). Correlation: 0.9997416536770057. Expr: tanh(0.536603*x0)\n",
            "-----------------------------------------------------------\n",
            "Selected skeleton: c*tanh(c*x0) + c\n",
            "\n",
            "********************************\n",
            "Analyzing variable x1\n",
            "********************************\n",
            "Predicted skeleton 1 for variable x1: c*Abs(x1) + c\n",
            "Predicted skeleton 2 for variable x1: c*x1*tanh(c*x1) + c\n",
            "Predicted skeleton 3 for variable x1: c + x1*tanh(c*x1 + c)\n",
            "Predicted skeleton 4 for variable x1: c + x1*tanh(c*x1)\n",
            "Predicted skeleton 5 for variable x1: c*x1*tanh(c*x1 + c) + c\n",
            "\n",
            "Choosing the best skeleton... (skeletons ordered based on number of nodes)\n",
            "\tSkeleton: c*Abs(x1) + c. Correlation: 0.9987919523613259. Expr: Abs(x1)\n",
            "-----------------------------------------------------------\n",
            "Selected skeleton: c*Abs(x1) + c\n",
            "\n",
            "********************************\n",
            "Analyzing variable x2\n",
            "********************************\n",
            "Predicted skeleton 1 for variable x2: c*cos(c*(c + x2)**2 + c) + c\n",
            "Predicted skeleton 2 for variable x2: c*cos(c*x2 + c) + c\n",
            "Predicted skeleton 3 for variable x2: c*(c*cos(c*x2 + c) + c)**3 + c\n",
            "Predicted skeleton 4 for variable x2: c*(c*cos(x2) + c)**3 + c\n",
            "Predicted skeleton 5 for variable x2: c*cos(c*(c + x2)**3 + c) + c\n",
            "\n",
            "Choosing the best skeleton... (skeletons ordered based on number of nodes)\n",
            "\tSkeleton: c*cos(c*x2 + c) + c. Correlation: 0.45346070794917664. Expr: cos(2.419372*x2 - 6.283185)\n",
            "\tSkeleton: c*cos(x2)**3 + c*cos(x2)**2 + c*cos(x2) + c. Correlation: 0.38832039013792896. Expr: -15.209976*cos(x2)**3 + cos(x2)**2 + 20.0*cos(x2)\n",
            "\tSkeleton: c*cos(c*x2**2 + c*x2 + c) + c. Correlation: 0.9990666723733511. Expr: cos(0.199791*x2**2 - 0.001698*x2 + 0.016807)\n",
            "-----------------------------------------------------------\n",
            "Selected skeleton: c*cos(c*x2**2 + c*x2 + c) + c\n",
            "\n"
          ]
        },
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "[c*tanh(c*x0) + c, c*Abs(x1) + c, c*cos(c*x2**2 + c*x2 + c) + c]"
            ]
          },
          "metadata": {},
          "execution_count": 4
        }
      ],
      "source": [
        "regressor = MSSP(dataset=data, bb_model=nn_model)\n",
        "regressor.get_skeletons()"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "c2zsE5vFIzKh",
        "outputId": "74eddb50-885d-4afc-bb22-74f9e34332fa"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Example using custom equations\n",
        "\n",
        "Here we will show how to use data generated from your own equations. Alternatively, you can bring your dataset (e.g., a CSV file) and load the matrix $X$ (explainable variables) and $Y$ (response variable).\n",
        "\n",
        "In this example, consider the simple equation $y = \\frac{\\sin(x_1 + 1.2 \\, x_2) \\, x_3^2}{2}$. Suppose that $x_1$ and $x_2$ are continuous variables and $x_3$ is discrete and can take 100 possible values ($x_1 \\in [-10, 10]$, $x_2 \\in [-5, 5]$, and $x_3 \\in [-8, ... , 8]$)"
      ],
      "metadata": {
        "collapsed": false,
        "id": "TawNtDurIzKi"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "**Generate and format data**"
      ],
      "metadata": {
        "collapsed": false,
        "id": "Uwpany30IzKi"
      }
    },
    {
      "cell_type": "code",
      "execution_count": 5,
      "outputs": [],
      "source": [
        "np.random.seed(7)\n",
        "n = 10000\n",
        "# Generate data from the equation\n",
        "x1 = np.random.uniform(-10, 10, size=n)\n",
        "x2 = np.random.uniform(-5, 5, size=n)\n",
        "x3 = np.array([np.random.choice(np.linspace(-8, 8, 100)) for _ in range(n)])  # Example of discrete variable\n",
        "X = np.array([x1, x2, x3]).T\n",
        "Y = np.sin(x1 + 1.2 * x2) * (x3**2 / 2)  # Or load matrices X and Y from a CSV file\n",
        "\n",
        "# Format the dataset\n",
        "names = ['x0', 'x1', 'x2']  # Specify the names of the variables\n",
        "types = ['continuous', 'continuous', 'discrete']  # Specify if the variables are continuous or discrete\n",
        "dataset = InputData(X=X, Y=Y, names=names, types=types)"
      ],
      "metadata": {
        "id": "cwfuSiHWIzKi"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "**Train a NN**\n",
        "\n",
        "Unlike the previous example, we haven't trained a NN for this problem so let's train it now. If you're not satisfied with the validation MSE, you can try increasing the number of epochs or try a different architecture. By default, we use the `modelType='NN'`; if you need less complexity, try `modelType='NN2'`; or if you need more complexity, try `modelType='NN3'`."
      ],
      "metadata": {
        "collapsed": false,
        "id": "kF5AwoDEIzKj"
      }
    },
    {
      "cell_type": "code",
      "execution_count": 6,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "*****************************************\n",
            "Start MLP training\n",
            "*****************************************\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "100%|██████████| 3000/3000 [09:20<00:00,  5.35it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Val MSE: 0.07133777567923286\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\n"
          ]
        }
      ],
      "source": [
        "from EquationLearning.Trainer.TrainNNmodel import Trainer\n",
        "\n",
        "predictor = Trainer(dataset=dataset, modelType='NN')\n",
        "predictor.train(batch_size=128, epochs=3000, printProcess=False)\n",
        "# Save the model\n",
        "# predictor.model.saveModel(path)  # Specify your own path"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "OK3syrkvIzKj",
        "outputId": "9d6e3b30-1fd8-477e-f75e-d03efc6d4144"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "**Get Skeletons**\n",
        "\n",
        "The following method will generate some candidate symbolic skeletons and select the most appropriate for each variable"
      ],
      "metadata": {
        "collapsed": false,
        "id": "rRwLwX_-IzKj"
      }
    },
    {
      "cell_type": "code",
      "execution_count": 7,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "********************************\n",
            "Analyzing variable x0\n",
            "********************************\n",
            "Predicted skeleton 1 for variable x0: c*cos(c + x0) + c\n",
            "Predicted skeleton 2 for variable x0: c*sin(c + x0) + c\n",
            "Predicted skeleton 3 for variable x0: c*cos(c*x0 + c) + c\n",
            "Predicted skeleton 4 for variable x0: c*cos(c*x0) + c\n",
            "Predicted skeleton 5 for variable x0: c*cos(x0) + c\n",
            "\n",
            "Choosing the best skeleton... (skeletons ordered based on number of nodes)\n",
            "\tSkeleton: c*cos(x0) + c. Correlation: 0.9549434622348927. Expr: cos(x0)\n",
            "\tSkeleton: c*cos(c + x0) + c. Correlation: 0.9998902707120207. Expr: cos(x0 - 5.968739)\n",
            "-----------------------------------------------------------\n",
            "Selected skeleton: c*cos(c + x0) + c\n",
            "\n",
            "********************************\n",
            "Analyzing variable x1\n",
            "********************************\n",
            "Predicted skeleton 1 for variable x1: c*sin(c*x1 + c) + c\n",
            "Predicted skeleton 2 for variable x1: c*cos(c*x1 + c) + c\n",
            "Predicted skeleton 3 for variable x1: c*sin(c + x1) + c\n",
            "Predicted skeleton 4 for variable x1: c*cos(c*x1) + c\n",
            "Predicted skeleton 5 for variable x1: c*cos(c + x1) + c\n",
            "\n",
            "Choosing the best skeleton... (skeletons ordered based on number of nodes)\n",
            "\tSkeleton: c*sin(c + x1) + c. Correlation: 0.8109687552029223. Expr: sin(x1 - 5.021683)\n",
            "\tSkeleton: c*cos(c*x1) + c. Correlation: 0.9437250798178028. Expr: cos(1.195566*x1)\n",
            "\tSkeleton: c*cos(c + x1) + c. Correlation: 0.8109687552029744. Expr: cos(x1 + 5.973891)\n",
            "\tSkeleton: c*sin(c*x1 + c) + c. Correlation: 0.9969592793660724. Expr: sin(1.19879*x1 + 1.25925)\n",
            "\tSkeleton: c*cos(c*x1 + c) + c. Correlation: 0.9969591651624546. Expr: cos(1.19903*x1 - 0.311585)\n",
            "-----------------------------------------------------------\n",
            "Selected skeleton: c*sin(c*x1 + c) + c\n",
            "\n",
            "********************************\n",
            "Analyzing variable x2\n",
            "********************************\n",
            "Predicted skeleton 1 for variable x2: c*x2**2 + c\n",
            "Predicted skeleton 2 for variable x2: c*x2**2 + c*x2 + c\n",
            "Predicted skeleton 3 for variable x2: c*(c + x2)**2 + c\n",
            "Predicted skeleton 4 for variable x2: c*x2**2 + c*tanh(c*x2 + c) + c\n",
            "Predicted skeleton 5 for variable x2: c*x2**3 + c\n",
            "\n",
            "Choosing the best skeleton... (skeletons ordered based on number of nodes)\n",
            "\tSkeleton: c*x2**2 + c. Correlation: 0.9997445507920887. Expr: x2**2\n",
            "-----------------------------------------------------------\n",
            "Selected skeleton: c*x2**2 + c\n",
            "\n"
          ]
        },
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "[c*cos(c + x0) + c, c*sin(c*x1 + c) + c, c*x2**2 + c]"
            ]
          },
          "metadata": {},
          "execution_count": 7
        }
      ],
      "source": [
        "regressor = MSSP(dataset=dataset, bb_model=predictor.model)\n",
        "regressor.get_skeletons()"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "_Hsjb4o5IzKj",
        "outputId": "20a34b87-1706-4a12-d8fd-327eb16cb73b"
      }
    }
  ],
  "metadata": {
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "codemirror_mode": {
        "name": "ipython",
        "version": 2
      },
      "file_extension": ".py",
      "mimetype": "text/x-python",
      "name": "python",
      "nbconvert_exporter": "python",
      "pygments_lexer": "ipython2",
      "version": "2.7.6"
    },
    "colab": {
      "provenance": [],
      "gpuType": "T4"
    },
    "accelerator": "GPU"
  },
  "nbformat": 4,
  "nbformat_minor": 0
}