4_How_to_contribute.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Example Run of an Augmented Simulator"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This tutorial notebook provides a guidance for installing the required packages and testing implemented augmented simulators using LIPS platform. \n",
    "\n",
    "**A quick walkthrough:**\n",
    "\n",
    "- Install the required packages using the requirements.txt file in the github repository for the required used case.\n",
    "\n",
    "- Some baseline are already implemented in the LIPS platform that could be seen to have some inspiration.\n",
    "\n",
    "- The augmented simulators related hyperparameters could be modified via dedicated configuration files.\n",
    "\n",
    "- The LIPS platform will be used to evaluate the trained augmented simulators from different evaluation criteria categories and attribute a score to each run."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## How to implement your own augmented simulator\n",
    "\n",
    "In the following, we show 3 ways to implement an augmented simulator (based on ML or a hybrid physics-AI model):\n",
    "\n",
    "1- Using an existing augmented simulator (baseline) in LIPS platform, train it and then evaluate the results;\n",
    "\n",
    "2- Implement an augmented simulator using LIPS framework template to take the advantage of existing training loop and other offered features;\n",
    "\n",
    "3- Implement an augmented simulator independently from LIPS platform and plug the trained model into LIPS to evaluate its results."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As so, in order to adress the augmented simulator handling when running this notebook, it is divided into the following four sections:\n",
    "1. [Generic step (Load the required data)](#generic_step)\n",
    "2. [Evaluate an existing augmented simulator](#existing_sim) (Beginner users)\n",
    "3. [Train and evaluate a custom augmented simulators developed using LIPS framework](#train_using_lips) (Intermediate level users)\n",
    "4. [Train a custom augmented simulator independently from LIPS and use the framwork to evaluate the final results](#train_custom) (Advanced users)\n",
    "\n",
    "Depending on the user level, it conveniently point to the dedicated section."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Prerequisites"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Install the LIPS framework if it is not already done. For more information look at the LIPS framework [Github repository](https://github.com/IRT-SystemX/LIPS) "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# !pip install -r requirements.txt\n",
    "# or \n",
    "# !pip install -U ."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Install the AirfRANS package"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# !pip install airfrans"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Generic Step (Load the required data) <a id='generic_step'></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "from lips import get_root_path"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "# indicate required paths\n",
    "LIPS_PATH = get_root_path()\n",
    "DIRECTORY_NAME = 'Dataset'\n",
    "BENCHMARK_NAME = \"Case1\"\n",
    "LOG_PATH = LIPS_PATH + \"lips_logs.log\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Define the configuration files path, that aim to describe specific caracteristics of the use case or the augmented simulator."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "BENCH_CONFIG_PATH = os.path.join(\"airfoilConfigurations\",\"benchmarks\",\"confAirfoil.ini\") #Configuration file related to the benchmark\n",
    "SIM_CONFIG_PATH = os.path.join(\"airfoilConfigurations\",\"simulators\",\"torch_fc.ini\") #Configuration file re"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Download the data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "from lips.dataset.airfransDataSet import download_data\n",
    "if not os.path.isdir(DIRECTORY_NAME):\n",
    "    download_data(root_path=\".\", directory_name=DIRECTORY_NAME)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Loading the dataset using the dedicated class used by LIPS platform offers a list of advantages:\n",
    "\n",
    "1. Ease the importing of datasets\n",
    "1. A set of functions to organize the `inputs` and `outputs` required by augmented simulators\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load the required benchmark datasets\n",
    "from lips.benchmark.airfransBenchmark import AirfRANSBenchmark\n",
    "\n",
    "benchmark=AirfRANSBenchmark(benchmark_path = DIRECTORY_NAME,\n",
    "                            config_path = BENCH_CONFIG_PATH,\n",
    "                            benchmark_name = BENCHMARK_NAME,\n",
    "                            log_path = LOG_PATH)\n",
    "benchmark.load(path=DIRECTORY_NAME)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Section-I (Evaluate an existing augmented simulator) <a id='existing_sim'></a>\n",
    "For beginners."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We start by importing an architecture from exisiting set of architectures and instantiate the `TorchSimulator` class which offers a set of utilities to train and analyze the selected augmented simulator. User could play with the configuration file of an existing augmented simulator to modify the model hyperparameters.\n",
    "\n",
    "The configuration file could be found at `./configurations/airfoil/simulators/torch_fc.ini`:\n",
    "\n",
    "```output\n",
    "[DEFAULT]\n",
    "name = \"torch_fc\"\n",
    "layers = (64,64,8,64,64,64,8,64,64)\n",
    "activation = \"relu\"\n",
    "layer = \"linear\"\n",
    "input_dropout = 0.0\n",
    "dropout = 0.0\n",
    "metrics = (\"MAELoss\",)\n",
    "loss = {\"name\": \"MSELoss\",\n",
    "        \"params\": {\"size_average\": None,\n",
    "                   \"reduce\": None,\n",
    "                   \"reduction\": 'mean'}}\n",
    "device = \"cpu\"\n",
    "optimizer = {\"name\": \"adam\",\n",
    "             \"params\": {\"lr\": 2e-4}}\n",
    "train_batch_size = 128000\n",
    "eval_batch_size = 256000\n",
    "epochs = 200\n",
    "shuffle = False\n",
    "save_freq = False\n",
    "ckpt_freq = 50\n",
    "```\n",
    "\n",
    "In the example below we select the configuration provided in `[DEFAULT]` section and new configuration could be created using a new section name and modifying the existing parameters."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "N.B.: In this context, `train_batch_size` and `eval_batch_size` refer to the number of nodes, not the number of simulations."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The next step is to instantiate a simulator with the `[DEFAULT]` configuration"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [],
   "source": [
    "from lips.augmented_simulators.torch_models.fully_connected import TorchFullyConnected\n",
    "from lips.augmented_simulators.torch_simulator import TorchSimulator\n",
    "from lips.dataset.scaler.standard_scaler_iterative import StandardScalerIterative\n",
    "from lips.dataset.scaler.standard_scaler import StandardScaler\n",
    "\n",
    "# chunk_sizes=benchmark.train_dataset.get_simulations_sizes()\n",
    "# no_norm_x=benchmark.train_dataset.get_no_normalization_axis_indices()\n",
    "# scalerParams={\"chunk_sizes\":chunk_sizes,\"no_norm_x\":no_norm_x}\n",
    "\n",
    "torch_sim = TorchSimulator(name=\"torch_fc\",\n",
    "                           model=TorchFullyConnected,\n",
    "                           scaler=StandardScaler,\n",
    "                           #scalerParams=scalerParams,\n",
    "                           log_path=\"log_benchmark\",\n",
    "                           device=\"cuda:1\",\n",
    "                           bench_config_path=BENCH_CONFIG_PATH,\n",
    "                           bench_config_name=BENCHMARK_NAME,\n",
    "                           sim_config_path=SIM_CONFIG_PATH,\n",
    "                           sim_config_name=\"DEFAULT\",\n",
    "                          )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "torch_sim.params"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Train the augmented simulator using the benchmark datasets."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "torch_sim.train(benchmark.train_dataset, \n",
    "                save_path=None, \n",
    "                epochs=200, \n",
    "                pin_memory=True, \n",
    "                non_blocking=True, \n",
    "                num_workers=16 \n",
    "                )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "\n",
    "plt.plot(torch_sim.train_losses)\n",
    "plt.xlabel(\"Epochs\")\n",
    "plt.ylabel(\"MSE Loss\")\n",
    "plt.grid()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You can also save and load the model fitted parameters alongside its meta data using the following functions."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [],
   "source": [
    "TRAINED_MODELS = \"AirfRANSModel\"\n",
    "SAVE_PATH = TRAINED_MODELS+os.sep+ \"fully_connected\"\n",
    "torch_sim.save(SAVE_PATH)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Load the trained model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from lips.augmented_simulators.torch_models.fully_connected import TorchFullyConnected\n",
    "from lips.augmented_simulators.torch_simulator import TorchSimulator\n",
    "from lips.dataset.scaler.standard_scaler import StandardScaler\n",
    "\n",
    "torch_sim = TorchSimulator(name=\"torch_fc\",\n",
    "                           model=TorchFullyConnected,\n",
    "                           scaler=StandardScaler,\n",
    "                           log_path=\"log_benchmark\",\n",
    "                           device=\"cuda:1\",\n",
    "                           bench_config_path=BENCH_CONFIG_PATH,\n",
    "                           bench_config_name=BENCHMARK_NAME,\n",
    "                           sim_config_path=SIM_CONFIG_PATH,\n",
    "                           sim_config_name=\"DEFAULT\",\n",
    "                          )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "LOAD_PATH = TRAINED_MODELS +os.sep+ \"fully_connected\"\n",
    "torch_sim.restore(path=LOAD_PATH)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Finally, the trained augmented simulator could be evaluated using the `evaluate_simulator` function of the `Benchmark` class. You can set on which dataset you want to evaluate your trained augmented simulator. The possibilites are `all`, `val`, `test`, `test_ood_topo`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [],
   "source": [
    "torch_sim_metrics = benchmark.evaluate_simulator(augmented_simulator=torch_sim,\n",
    "                                                 eval_batch_size=256000,\n",
    "                                                 dataset=\"test\", #all, test_ood\n",
    "                                                 pin_memory=True,\n",
    "                                                 non_blocking=True,\n",
    "                                                 num_workers=16,\n",
    "                                                 save_path=None,\n",
    "                                                 save_predictions=False\n",
    "                                                )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You can see how your model performs directly by looking at the evaluation metrics resulted by from the last step."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "torch_sim_metrics[\"test\"]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Section-II (Train and Evaluate a new augmented simulator using LIPS platform) <a id='train_using_lips'></a>\n",
    "For intermediate level users."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You can implement an augmented simulator respecting the following template. Some of the functions are mandatory (`build_model`, `forward`, `process_dataset`, `post_process`) and others are optional (function to get metadata, save, load the model parameters).\n",
    "\n",
    "A best way to take advantage of all the offered functionalities by LIPS platform, is to keep the constructor `__init__` as it is presented and to customize the mandatory functions to construct your own architecture."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {},
   "outputs": [],
   "source": [
    "\"\"\"\n",
    "Torch fully connected model\n",
    "\"\"\"\n",
    "import os\n",
    "import pathlib\n",
    "from typing import Union\n",
    "import json\n",
    "\n",
    "import numpy as np\n",
    "import numpy.typing as npt\n",
    "\n",
    "import torch\n",
    "from torch import nn\n",
    "import torch.nn.functional as F\n",
    "from torch.utils.data import TensorDataset, DataLoader\n",
    "\n",
    "from lips.dataset import DataSet\n",
    "from lips.dataset.scaler import Scaler\n",
    "from lips.logger import CustomLogger\n",
    "from lips.config import ConfigManager\n",
    "from lips.utils import NpEncoder\n",
    "from lips.augmented_simulators.torch_models.utils import LOSSES\n",
    "\n",
    "class MyCustomFullyConnected(nn.Module):\n",
    "    \"\"\"_summary_\n",
    "\n",
    "    Parameters\n",
    "    ----------\n",
    "    sim_config_path : Union[``pathlib.Path``, ``str``]\n",
    "        The path to the configuration file for simulator.\n",
    "        It should contain all the required hyperparameters for this model.\n",
    "    sim_config_name : Union[str, None], optional\n",
    "        the simulator config section name, by default None\n",
    "    name : Union[str, None], optional\n",
    "        the simulator name used for save and load, by default None\n",
    "    scaler : Union[Scaler, None], optional\n",
    "        A scaler used to normalize the data, by default None\n",
    "    bench_config_path : Union[str, pathlib.Path, None], optional\n",
    "        a path to the benchmark configuration file, by default None\n",
    "    bench_config_name : Union[str, None], optional\n",
    "        the section name of the benchmark configuration, by default None\n",
    "    log_path : Union[None, str], optional\n",
    "        a path where the logs should be saved, by default None\n",
    "\n",
    "    Raises\n",
    "    ------\n",
    "    RuntimeError\n",
    "        You should provide a path to the configuration file for this augmented simulator\n",
    "    \"\"\"\n",
    "    def __init__(self,\n",
    "                 sim_config_path: Union[pathlib.Path, str],\n",
    "                 bench_config_path: Union[str, pathlib.Path],\n",
    "                 sim_config_name: Union[str, None]=None,\n",
    "                 bench_config_name: Union[str, None]=None,\n",
    "                 name: Union[str, None]=None,\n",
    "                 scaler: Union[Scaler, None]=None,\n",
    "                 log_path: Union[None, pathlib.Path, str]=None,\n",
    "                 **kwargs):\n",
    "        super().__init__()\n",
    "        if not os.path.exists(sim_config_path):\n",
    "            raise RuntimeError(\"Configuration path for the simulator not found!\")\n",
    "        if not str(sim_config_path).endswith(\".ini\"):\n",
    "            raise RuntimeError(\"The configuration file should have `.ini` extension!\")\n",
    "        sim_config_name = sim_config_name if sim_config_name is not None else \"DEFAULT\"\n",
    "        self.sim_config = ConfigManager(section_name=sim_config_name, path=sim_config_path)\n",
    "        self.bench_config = ConfigManager(section_name=bench_config_name, path=bench_config_path)\n",
    "        self.name = name if name is not None else self.sim_config.get_option(\"name\")\n",
    "        # scaler\n",
    "        self.scaler = scaler\n",
    "        # Logger\n",
    "        self.log_path = log_path\n",
    "        self.logger = CustomLogger(__class__.__name__, log_path).logger\n",
    "        # model parameters\n",
    "        self.params = self.sim_config.get_options_dict()\n",
    "        self.params.update(kwargs)\n",
    "\n",
    "        self.activation = {\n",
    "            \"relu\": F.relu,\n",
    "            \"sigmoid\": F.sigmoid,\n",
    "            \"tanh\": F.tanh\n",
    "        }\n",
    "\n",
    "        self.input_size = None if kwargs.get(\"input_size\") is None else kwargs[\"input_size\"]\n",
    "        self.output_size = None if kwargs.get(\"output_size\") is None else kwargs[\"output_size\"]\n",
    "\n",
    "        self.input_layer = None\n",
    "        self.input_dropout = None\n",
    "        self.fc_layers = None\n",
    "        self.dropout_layers = None\n",
    "        self.output_layer = None\n",
    "\n",
    "        # batch information\n",
    "        self._data = None\n",
    "        self._target = None\n",
    "\n",
    "    def build_model(self):\n",
    "        \"\"\"Build the model flow\n",
    "        \"\"\"\n",
    "        linear_sizes = list(self.params[\"layers\"])\n",
    "\n",
    "        self.input_layer = nn.Linear(self.input_size, linear_sizes[0])\n",
    "        self.input_dropout = nn.Dropout(p=self.params[\"input_dropout\"])\n",
    "\n",
    "        self.fc_layers = nn.ModuleList([nn.Linear(in_f, out_f) \\\n",
    "            for in_f, out_f in zip(linear_sizes[:-1], linear_sizes[1:])])\n",
    "\n",
    "        self.dropout_layers = nn.ModuleList([nn.Dropout(p=self.params[\"dropout\"]) \\\n",
    "            for _ in range(len(self.fc_layers))])\n",
    "\n",
    "        self.output_layer = nn.Linear(linear_sizes[-1], self.output_size)\n",
    "\n",
    "    def forward(self, data):\n",
    "        out = self.input_layer(data)\n",
    "        out = self.input_dropout(out)\n",
    "        for _, (fc_, dropout) in enumerate(zip(self.fc_layers, self.dropout_layers)):\n",
    "            out = fc_(out)\n",
    "            out = self.activation[self.params[\"activation\"]](out)\n",
    "            out = dropout(out)\n",
    "        out = self.output_layer(out)\n",
    "        return out\n",
    "\n",
    "    def process_dataset(self, dataset: DataSet, training: bool, **kwargs) -> DataLoader:\n",
    "        \"\"\"process the datasets for training and evaluation\n",
    "\n",
    "        This function transforms all the dataset into something that can be used by the neural network (for example)\n",
    "\n",
    "        Parameters\n",
    "        ----------\n",
    "        dataset : DataSet\n",
    "            an object of the DataSet class including the required data\n",
    "        scaler : Scaler, optional\n",
    "            A scaler instance to be used for normalization, by default True\n",
    "        training : bool, optional\n",
    "            A boolean indicating whether we are in training or evaluation phases, by default False\n",
    "            If `True`, the scaler will be fit to the data to estimate the parameters\n",
    "            If `False`, the estimated parameters of the scaler during training will be used to normalize the \n",
    "            validation/test/test_ood data\n",
    "\n",
    "        kwargs : dict\n",
    "            The supplementary arguments to be used for acceleration of DataLoader which are:\n",
    "                pin_memory : `bool`, optional\n",
    "                    refere to pytorch documentation for more information\n",
    "                num_workers : Union[None, int], optional\n",
    "                    the number of CPU workers to be used to transfer the batches to device\n",
    "                dtype : torch.types\n",
    "                    the data type that will be used to transform the processed dataset\n",
    "        Returns\n",
    "        -------\n",
    "        DataLoader\n",
    "            A pytorch data loader from which the batches of data could be loaded for training\n",
    "        \"\"\"\n",
    "        pin_memory = kwargs.get(\"pin_memory\", True)\n",
    "        num_workers = kwargs.get(\"num_workers\", None)\n",
    "        dtype = kwargs.get(\"dtype\", torch.float32)\n",
    "\n",
    "        if training:\n",
    "            self._infer_size(dataset)\n",
    "            batch_size = self.params[\"train_batch_size\"]\n",
    "            extract_x, extract_y = dataset.extract_data()\n",
    "            if self.scaler is not None:\n",
    "                extract_x, extract_y = self.scaler.fit_transform(extract_x, extract_y)\n",
    "        else:\n",
    "            batch_size = self.params[\"eval_batch_size\"]\n",
    "            extract_x, extract_y = dataset.extract_data()\n",
    "            if self.scaler is not None:\n",
    "                extract_x, extract_y = self.scaler.transform(extract_x, extract_y)\n",
    "\n",
    "        torch_dataset = TensorDataset(torch.tensor(extract_x, dtype=dtype), torch.tensor(extract_y, dtype=dtype))\n",
    "        if num_workers is None:\n",
    "            data_loader = DataLoader(torch_dataset, batch_size=batch_size, shuffle=self.params[\"shuffle\"], pin_memory=pin_memory)\n",
    "        else:\n",
    "            data_loader = DataLoader(torch_dataset, batch_size=batch_size, shuffle=self.params[\"shuffle\"], pin_memory=pin_memory, num_workers=num_workers)\n",
    "        #data_loader = DataLoader(torch_dataset, batch_size=batch_size, shuffle=self.params[\"shuffle\"])\n",
    "        return data_loader\n",
    "\n",
    "    def _post_process(self, data):\n",
    "        if self.scaler is not None:\n",
    "            try:\n",
    "                processed = self.scaler.inverse_transform(data)\n",
    "            except TypeError:\n",
    "                processed = self.scaler.inverse_transform(data.cpu())\n",
    "        else:\n",
    "            processed = data\n",
    "        return processed\n",
    "    \n",
    "    def _reconstruct_output(self, dataset: DataSet, data: npt.NDArray[np.float64]) -> dict:\n",
    "        \"\"\"Reconstruct the outputs to obtain the desired shape for evaluation\n",
    "\n",
    "        In the simplest form, this function is implemented in DataSet class. It supposes that the predictions \n",
    "        obtained by the augmented simulator are exactly the same as the one indicated in the configuration file\n",
    "\n",
    "        However, if some transformations required by each specific model, the extra operations to obtained the\n",
    "        desired output shape should be done in this function.\n",
    "\n",
    "        Parameters\n",
    "        ----------\n",
    "        dataset : DataSet\n",
    "            An object of the `DataSet` class \n",
    "        data : npt.NDArray[np.float64]\n",
    "            the data which should be reconstructed to the desired form\n",
    "        \"\"\"\n",
    "        data_rec = dataset.reconstruct_output(data)\n",
    "        return data_rec\n",
    "\n",
    "    def _infer_size(self, dataset: DataSet):\n",
    "        \"\"\"Infer the size of the model\n",
    "\n",
    "        Parameters\n",
    "        ----------\n",
    "        dataset : DataSet\n",
    "            An object of the dataset class providing some functionalities to get sizes of inputs/outputs\n",
    "\n",
    "        \"\"\"\n",
    "        *dim_inputs, self.output_size = dataset.get_sizes()\n",
    "        self.input_size = np.sum(dim_inputs)\n",
    "\n",
    "    def get_metadata(self):\n",
    "        \"\"\"getting the augmented simulator meta data\n",
    "\n",
    "        Returns\n",
    "        -------\n",
    "        dict\n",
    "            a dictionary containing the meta data for the augmented simulator\n",
    "        \"\"\"\n",
    "        res_json = {}\n",
    "        res_json[\"input_size\"] = self.input_size\n",
    "        res_json[\"output_size\"] = self.output_size\n",
    "        return res_json\n",
    "\n",
    "    def _save_metadata(self, path: str):\n",
    "        \"\"\"Save the augmented simulator specific meta data\n",
    "\n",
    "        These information are required to restore a saved model\n",
    "\n",
    "        Parameters\n",
    "        ----------\n",
    "        path : str\n",
    "            A path where the meta data should be saved\n",
    "        \"\"\"\n",
    "        #super()._save_metadata(path)\n",
    "        #if self.scaler is not None:\n",
    "        #    self.scaler.save(path)\n",
    "        res_json = {}\n",
    "        res_json[\"input_size\"] = self.input_size\n",
    "        res_json[\"output_size\"] = self.output_size\n",
    "        with open((path / \"metadata.json\"), \"w\", encoding=\"utf-8\") as f:\n",
    "            json.dump(obj=res_json, fp=f, indent=4, sort_keys=True, cls=NpEncoder)\n",
    "\n",
    "    def _load_metadata(self, path: str):\n",
    "        \"\"\"Load the metadata for the augmentd simulator\n",
    "\n",
    "        Parameters\n",
    "        ----------\n",
    "        path : str\n",
    "            a path where the meta data are saved\n",
    "        \"\"\"\n",
    "        if not isinstance(path, pathlib.Path):\n",
    "            path = pathlib.Path(path)\n",
    "        #super()._load_metadata(path)\n",
    "        #if self.scaler is not None:\n",
    "        #    self.scaler.load(path)\n",
    "        with open((path / \"metadata.json\"), \"r\", encoding=\"utf-8\") as f:\n",
    "            res_json = json.load(fp=f)\n",
    "        self.input_size = res_json[\"input_size\"]\n",
    "        self.output_size = res_json[\"output_size\"]\n",
    "\n",
    "    def _do_forward(self, batch, **kwargs):\n",
    "        \"\"\"Do the forward step through a batch of data\n",
    "\n",
    "        This step could be very specific to each augmented simulator as each architecture\n",
    "        takes various inputs during the learning procedure. \n",
    "\n",
    "        Parameters\n",
    "        ----------\n",
    "        batch : _type_\n",
    "            A batch of data including various information required by an architecture\n",
    "        device : _type_\n",
    "            the device on which the data should be processed\n",
    "\n",
    "        Returns\n",
    "        -------\n",
    "        ``tuple``\n",
    "            returns the predictions made by the augmented simulator and also the real targets\n",
    "            on which the loss function should be computed\n",
    "        \"\"\"\n",
    "        non_blocking = kwargs.get(\"non_blocking\", True)\n",
    "        device = self.params.get(\"device\", \"cpu\")\n",
    "        self._data, self._target = batch\n",
    "        self._data = self._data.to(device, non_blocking=non_blocking)\n",
    "        self._target = self._target.to(device, non_blocking=non_blocking)\n",
    "\n",
    "        predictions = self.forward(self._data)\n",
    "        \n",
    "        return self._data, predictions, self._target\n",
    "\n",
    "    def get_loss_func(self, loss_name: str, **kwargs) -> torch.Tensor:\n",
    "        \"\"\"\n",
    "        Helper to get loss. It is specific to each architecture\n",
    "        \"\"\"\n",
    "        # if len(args) > 0:\n",
    "        #     # for Masked RNN loss. args[0] is the list of sequence lengths\n",
    "        #     loss_func = LOSSES[self.params[\"loss\"][\"name\"]](args[0], self.params[\"device\"])\n",
    "        # else:\n",
    "        loss_func = LOSSES[loss_name](**kwargs)\n",
    "        \n",
    "        return loss_func"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Once, the augmented simulator is implemented, you should also create a configuration which indicate all the hyper parameters required by this augmented simulator. An example of configuration file is shown in `configs/simulators/fully_connected.ini` and its content is shown below. \n",
    "\n",
    "The path and the section name of this configuration file should be given to your architecture as an argument (`sim_config_path`, `sim_config_name`) in order that it could be able to import all its required hyper-parameters."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "```\n",
    "[DEFAULT]\n",
    "name = \"torch_fc\"\n",
    "layers = (300, 300, 300, 300)\n",
    "activation = \"relu\"\n",
    "layer = \"linear\"\n",
    "input_dropout = 0.0\n",
    "dropout = 0.0\n",
    "metrics = (\"MAELoss\",)\n",
    "loss = {\"name\": \"MSELoss\",\n",
    "        \"params\": {\"size_average\": None,\n",
    "                   \"reduce\": None,\n",
    "                   \"reduction\": 'mean'}}\n",
    "device = \"cpu\"\n",
    "optimizer = {\"name\": \"adam\",\n",
    "             \"params\": {\"lr\": 3e-4}}\n",
    "train_batch_size = 128\n",
    "eval_batch_size = 128\n",
    "epochs = 10\n",
    "shuffle = False\n",
    "save_freq = False\n",
    "ckpt_freq = 50\n",
    "\n",
    "[CONFIG1]\n",
    "layers = (100, 100)\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Use the torch simulator which give as input your implemented augmented simulator (`MyCustomFullyConnected`) and offers a set of functionalities to train it and analyze its results. Optinally, you can also give a scaler (from the existing list of scalers or implement it yourself if you require a more advanced scaler), which is used by the `TorchSimulator` class to normalize your data before training your model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "metadata": {},
   "outputs": [],
   "source": [
    "from lips.augmented_simulators.torch_simulator import TorchSimulator\n",
    "from lips.dataset.scaler.standard_scaler_iterative import StandardScalerIterative\n",
    "\n",
    "chunk_sizes=benchmark.train_dataset.get_simulations_sizes()\n",
    "no_norm_x=benchmark.train_dataset.get_no_normalization_axis_indices()\n",
    "scalerParams={\"chunk_sizes\":chunk_sizes,\"no_norm_x\":no_norm_x}\n",
    "\n",
    "torch_sim = TorchSimulator(name=\"torch_fc\",\n",
    "                           model=MyCustomFullyConnected,\n",
    "                           scaler=StandardScalerIterative,\n",
    "                           scalerParams=scalerParams,\n",
    "                           log_path=None,\n",
    "                           device=\"cuda:1\", # use \"cpu\" if you don't have a GPU available on your machine\n",
    "                           seed=42,\n",
    "                           bench_config_path=BENCH_CONFIG_PATH,\n",
    "                           bench_config_name=\"Benchmark1\",\n",
    "                           sim_config_path=SIM_CONFIG_PATH,\n",
    "                           sim_config_name=\"DEFAULT\", # use the default set of hyper parameters\n",
    "                          )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "torch_sim.train(benchmark.train_dataset, \n",
    "                save_path=None, \n",
    "                epochs=10, \n",
    "                train_batch_size=128000,\n",
    "                pin_memory=True, \n",
    "                non_blocking=True, \n",
    "                num_workers=16\n",
    "                )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "torch_sim_metrics = benchmark.evaluate_simulator(augmented_simulator=torch_sim,\n",
    "                                                  eval_batch_size=256000,\n",
    "                                                  dataset=\"all\",\n",
    "                                                  shuffle=False,\n",
    "                                                  save_path=None,\n",
    "                                                  save_predictions=False\n",
    "                                                 )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "torch_sim_metrics[\"test\"]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Section-III (Train an augmented simulator independently and evaluate it through LIPS) <a id='train_custom'></a>\n",
    "For advanced users."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If you require more functionalities that are not offered by LIPS platform (e.g., adding advanced regularizations into the training loop, or adding physics constraints in your model) you can implement your architecture independently from LIPS platform and use only the evaluation part of the framework to assess your model performance. \n",
    "\n",
    "In the following, we show a simple architecture with a training loop and how it can be evaluated by the LIPS platform."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##### STEP 1: Implement your architecture based on Pytorch library in this Example\n",
    "\n",
    "**NB.** For Tensorflow users, there are also some examples provided in LIPS platform (see LIPS documentation)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "\n",
    "import torch\n",
    "from torch import nn\n",
    "from torch import optim\n",
    "import torch.nn.functional as F\n",
    "from torch.utils.data import TensorDataset, DataLoader\n",
    "\n",
    "class MyCustomFullyConnected(nn.Module):\n",
    "    def __init__(self,\n",
    "                 name: str=\"MyCustomFC\",\n",
    "                 input_size: int=None,\n",
    "                 output_size: int=None,\n",
    "                 hidden_sizes: tuple=(100,100,),\n",
    "                 activation=F.relu\n",
    "                ):\n",
    "        super().__init__()\n",
    "        self.name = name\n",
    "        \n",
    "        self.activation = activation\n",
    "\n",
    "        self.input_size = input_size\n",
    "        self.output_size = output_size\n",
    "        self.hidden_sizes = hidden_sizes\n",
    "\n",
    "        # model architecture\n",
    "        self.input_layer = nn.Linear(self.input_size, self.hidden_sizes[0])\n",
    "        self.fc_layers = nn.ModuleList([nn.Linear(in_f, out_f) \\\n",
    "                                        for in_f, out_f in zip(hidden_sizes[:-1], self.hidden_sizes[1:])])\n",
    "        self.output_layer = nn.Linear(hidden_sizes[-1], self.output_size)\n",
    "\n",
    "    def forward(self, data):\n",
    "        \"\"\"The forward pass of the model\n",
    "        \"\"\"\n",
    "        out = self.input_layer(data)\n",
    "        for _, fc_ in enumerate(self.fc_layers):\n",
    "            out = fc_(out)\n",
    "            out = self.activation(out)\n",
    "        out = self.output_layer(out)\n",
    "        return out"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##### STEP 2: Process the data to acquire the right Inputs and Outputs for your model alongside their dimensions\n",
    "This function uses a functionality offered by the Dataset class to extract the required inputs and outputs for the problem in hand, which facilitate the task. \n",
    "\n",
    "It also allows to create DataLoader from existing datasets.\n",
    "\n",
    "**NB.** However, the users could use their own extraction if they require to add further inputs (feature engineering or other operations)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "def process_dataset(dataset, batch_size: int=128000, training: bool=False, shuffle: bool=False, pin_memory=True, num_workers=0, dtype=torch.float32) -> DataLoader:\n",
    "    if training:\n",
    "        batch_size = batch_size\n",
    "        extract_x, extract_y = dataset.extract_data()\n",
    "    else:\n",
    "        batch_size = batch_size\n",
    "        extract_x, extract_y = dataset.extract_data()\n",
    "\n",
    "    torch_dataset = TensorDataset(torch.tensor(extract_x, dtype=dtype), torch.tensor(extract_y, dtype=dtype))\n",
    "    data_loader = DataLoader(torch_dataset, batch_size=batch_size, shuffle=shuffle, pin_memory=pin_memory, num_workers=num_workers)\n",
    "    return data_loader\n",
    "\n",
    "def infer_input_output_size(dataset):\n",
    "    *dim_inputs, output_size = dataset.get_sizes()\n",
    "    input_size = np.sum(dim_inputs)\n",
    "    return input_size, output_size"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##### STEP 3: Implement your own Training, Validation and Prediction functions\n",
    "\n",
    "**train.** This function allows to train (adjust the parameters of) your defined model using the provided datasets.\n",
    "\n",
    "**validate.** This function allows to validate your model on a validation dataset. The validation step is not mendatory and is used only to trace the model behavior (overfitting or not). \n",
    "\n",
    "**predict.** This function allows to predict using the trained model. The `DataSet` class provides a function `reconstruct_output` which allows to reshape the predictions in the correct form which will be comparable with ground truth. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "from tqdm import tqdm\n",
    "\n",
    "def train(model, train_loader, val_loader=None, epochs=100, lr=3e-4, device=\"cpu\"):\n",
    "    train_losses = []\n",
    "    val_losses = []\n",
    "    # select your optimizer\n",
    "    optimizer = optim.Adam(model.parameters(), lr=lr)\n",
    "    # select your loss function\n",
    "    loss_function = nn.MSELoss()\n",
    "    pbar = tqdm(range(epochs))\n",
    "    for epoch in pbar:\n",
    "        # set your model for training\n",
    "        model.train()\n",
    "        total_loss = 0\n",
    "        # iterate over the batches of data\n",
    "        pbar_batch=tqdm(train_loader)\n",
    "        for batch in pbar_batch:\n",
    "            data, target = batch\n",
    "            # transfer your data on proper device. The model and your data should be on the same device\n",
    "            data = data.to(device)\n",
    "            target = target.to(device)\n",
    "            # reset the gradient\n",
    "            optimizer.zero_grad()\n",
    "            # predict using your model on the current batch of data\n",
    "            prediction = model(data)\n",
    "            # compute the loss between prediction and real target\n",
    "            loss = loss_function(prediction, target)\n",
    "            # compute the gradient (backward pass of back propagation algorithm)\n",
    "            loss.backward()\n",
    "            # update the parameters of your model\n",
    "            optimizer.step()\n",
    "            total_loss += loss * len(data)\n",
    "        # the validation step is optional\n",
    "        if val_loader is not None:\n",
    "            val_loss = validate(model, val_loader, device)\n",
    "            val_losses.append(val_loss)\n",
    "        mean_loss = total_loss.item() / len(train_loader.dataset)\n",
    "        print(f\"Train Epoch: {epoch}   Avg_Loss: {mean_loss:.5f}\")\n",
    "        train_losses.append(mean_loss)\n",
    "    return model, train_losses, val_losses\n",
    "\n",
    "def validate(model, val_loader, device):\n",
    "    # set the model for evaluation (no update of the parameters)\n",
    "    model.eval()\n",
    "    total_loss = 0\n",
    "    loss_function = nn.MSELoss()\n",
    "    with torch.no_grad():\n",
    "        pbar_batch=tqdm(val_loader)\n",
    "        for batch in val_loader:\n",
    "            data, target = batch\n",
    "            data.to(device)\n",
    "            target.to(device)\n",
    "            prediction = model(data)\n",
    "            loss = loss_function(prediction, target)\n",
    "            total_loss += loss.item()*len(data)\n",
    "        mean_loss = total_loss / len(val_loader.dataset)\n",
    "        print(f\"Eval:   Avg_Loss: {mean_loss:.5f}\")\n",
    "    return mean_loss\n",
    "\n",
    "def predict(model, dataset, device, dtype=torch.float32, pin_memory=True, num_workers=0):\n",
    "    # set the model for the evaluation\n",
    "    model.eval()\n",
    "    predictions = []\n",
    "    observations = []\n",
    "    test_loader = process_dataset(dataset, training=False, dtype=dtype, pin_memory=pin_memory, num_workers=num_workers)\n",
    "    # we dont require the computation of the gradient\n",
    "    with torch.no_grad():\n",
    "        for batch in test_loader:\n",
    "            data, target = batch\n",
    "            data = data.to(device)\n",
    "            target = target.to(device)\n",
    "            prediction = model(data)\n",
    "            \n",
    "            if device == torch.device(\"cpu\"):\n",
    "                predictions.append(prediction.numpy())\n",
    "                observations.append(target.numpy())\n",
    "            else:\n",
    "                predictions.append(prediction.cpu().data.numpy())\n",
    "                observations.append(target.cpu().data.numpy())\n",
    "    # reconstruct the prediction in the proper required shape of target variables\n",
    "    predictions = np.concatenate(predictions)\n",
    "    predictions = dataset.reconstruct_output(predictions)\n",
    "    # Do the same for the real observations\n",
    "    observations = np.concatenate(observations)\n",
    "    observations = dataset.reconstruct_output(observations)\n",
    "\n",
    "    return predictions, observations"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "dtype = torch.float32\n",
    "train_loader = process_dataset(benchmark.train_dataset, training=True, shuffle=True, dtype=dtype, pin_memory=True, num_workers=16)\n",
    "input_size, output_size = infer_input_output_size(benchmark.train_dataset)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "device = torch.device(\"cuda:0\")\n",
    "\n",
    "model = MyCustomFullyConnected(input_size=input_size,\n",
    "                               output_size=output_size,\n",
    "                               hidden_sizes=(50,100,50),\n",
    "                               activation=F.relu\n",
    "                               )\n",
    "model.to(device)\n",
    "model.to(dtype)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model, train_losses, _ = train(model, train_loader, epochs=10, device=device, lr=3e-4)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##### prediction on `test_dataset`\n",
    "This dataset has the same distribution as the training set"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [],
   "source": [
    "predictions, observations = predict(model, benchmark._test_dataset, device=device, num_workers=16)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from lips.evaluation.airfrans_evaluation import AirfRANSEvaluation\n",
    "\n",
    "evaluator = AirfRANSEvaluation(config_path = BENCH_CONFIG_PATH,\n",
    "                               scenario = BENCHMARK_NAME,\n",
    "                               data_path = DIRECTORY_NAME,\n",
    "                               log_path = LOG_PATH)\n",
    "\n",
    "observation_metadata = benchmark._test_dataset.extra_data\n",
    "metrics = evaluator.evaluate(observations=observations,\n",
    "                             predictions=predictions,\n",
    "                             observation_metadata=observation_metadata)\n",
    "print(metrics)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##### Prediction on `test_ood_dataset`\n",
    "This dataset has a different distribution in comparison to the training set. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "predictions, observations = predict(model, benchmark._test_ood_dataset, device=device, num_workers=16)\n",
    "evaluator = AirfRANSEvaluation(config_path = BENCH_CONFIG_PATH,\n",
    "                               scenario = BENCHMARK_NAME,\n",
    "                               data_path = DIRECTORY_NAME,\n",
    "                               log_path = LOG_PATH)\n",
    "\n",
    "observation_metadata = benchmark._test_ood_dataset.extra_data\n",
    "metrics = evaluator.evaluate(observations=observations,\n",
    "                             predictions=predictions,\n",
    "                             observation_metadata=observation_metadata)\n",
    "print(metrics)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "ml4phy",
   "language": "python",
   "name": "python3"
  },
  "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.8.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}