{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "670316b8-460c-4009-a5da-94278f4ac9a9",
   "metadata": {
    "papermill": {
     "duration": null,
     "end_time": null,
     "exception": null,
     "start_time": null,
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "# Optimize model hyperparameters using Optuna"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "52af59bb-083c-46c6-989a-bd4c65137a1a",
   "metadata": {
    "papermill": {
     "duration": null,
     "end_time": null,
     "exception": null,
     "start_time": null,
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Import to be able to import python package from src\n",
    "import sys\n",
    "sys.path.insert(0, '../src')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "d6fc731f-3f50-4e9a-a24c-b2ab01d4fa31",
   "metadata": {
    "papermill": {
     "duration": null,
     "end_time": null,
     "exception": null,
     "start_time": null,
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "import ontime as on\n",
    "import pandas as pd"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2ffac0d9-4247-4ac2-af8e-aad77eb1f960",
   "metadata": {},
   "source": [
    "## Prerequisite"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b211f78b-2de8-4a0c-8ce5-0806e504207b",
   "metadata": {},
   "source": [
    "Install Optuna within your project"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ef375aa9-89b6-4ba8-8f06-d874526dda32",
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install optuna optuna-integration optuna-dashboard"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7873abeb",
   "metadata": {},
   "source": [
    "## Create the dataset\n",
    "\n",
    "We will directly use the BenchmarkDataset class, so that we can quickly create splits and samples for evaluating the model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "218c5e0f",
   "metadata": {},
   "outputs": [],
   "source": [
    "from ontime.module.benchmarking import BenchmarkDataset\n",
    "from sklearn.preprocessing import StandardScaler\n",
    "from darts.dataprocessing.transformers import Scaler"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "831f1944-599b-4761-a071-2a682346610a",
   "metadata": {
    "papermill": {
     "duration": null,
     "end_time": null,
     "exception": null,
     "start_time": null,
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "# generate a random time series\n",
    "ts = on.generators.random_walk().generate(start=pd.Timestamp('2019-01-01'), end=pd.Timestamp('2023-12-31'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "ef3e03e1-c247-4b5a-a27a-a13361e673b0",
   "metadata": {
    "papermill": {
     "duration": null,
     "end_time": null,
     "exception": null,
     "start_time": null,
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "# create the dataset\n",
    "dataset = BenchmarkDataset(\n",
    "    ts,\n",
    "    \"Random series\",\n",
    "    input_length=120,\n",
    "    target_length=48,\n",
    "    gap=0,\n",
    "    stride=48,\n",
    "    scaler_type=StandardScaler # add normalization cause why not\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "3084ca5c",
   "metadata": {},
   "outputs": [],
   "source": [
    "# split dataset (we only need train and val)\n",
    "train, val = dataset.get_train_val_split()\n",
    "scaler = Scaler(dataset.scaler_type())\n",
    "train = scaler.fit_transform(train)\n",
    "val = scaler.transform(val)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "402d43be",
   "metadata": {},
   "source": [
    "## Setup the model with its hyperparameters\n",
    "\n",
    "The setup method defines the hyperparameters and their values that must be tested in during the optimization. This method is dependent of the model, and should therefore be changed for your use case."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "2ecf519f",
   "metadata": {},
   "outputs": [],
   "source": [
    "from darts.models import TCNModel\n",
    "from pytorch_lightning.callbacks.early_stopping import EarlyStopping\n",
    "from optuna_integration import PyTorchLightningPruningCallback"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "4ff66b88",
   "metadata": {},
   "outputs": [],
   "source": [
    "# avoid issue with optuna https://github.com/optuna/optuna/issues/4689\n",
    "import pytorch_lightning as pl\n",
    "\n",
    "class OptunaPruning(PyTorchLightningPruningCallback, pl.Callback):\n",
    "    def __init__(self, *args, **kwargs):\n",
    "        super().__init__(*args, **kwargs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "5f0c1dd6",
   "metadata": {},
   "outputs": [],
   "source": [
    "def setup_model(trial):\n",
    "    # define the model\n",
    "\n",
    "    pl_trainer_kwargs = {\n",
    "        \"accelerator\": \"auto\",\n",
    "        \"callbacks\": [\n",
    "            EarlyStopping(\n",
    "                monitor=\"val_loss\",\n",
    "                patience=10,\n",
    "                mode=\"min\",\n",
    "            ),\n",
    "            OptunaPruning(\n",
    "                trial,\n",
    "                monitor=\"val_loss\",\n",
    "            ),\n",
    "        ]\n",
    "    }\n",
    "\n",
    "    tcn_model = TCNModel(\n",
    "        model_name=\"optuna_tcn_model\",\n",
    "        input_chunk_length=dataset.input_length,\n",
    "        output_chunk_length=trial.suggest_int(\"output_chunk_length\", 1, 48),\n",
    "        kernel_size=trial.suggest_int(\"kernel_size\", 2, 5),\n",
    "        num_filters=trial.suggest_int(\"num_filters\", 8, 64),\n",
    "        num_layers=trial.suggest_int(\"num_layers\", 1, 4),\n",
    "        dropout=trial.suggest_float(\"dropout\", 0.0, 0.5),\n",
    "        weight_norm=trial.suggest_categorical(\"weight_norm\", [True, False]),\n",
    "        n_epochs=5,\n",
    "        random_state=42,\n",
    "        optimizer_kwargs={\n",
    "            \"lr\": trial.suggest_float(\"lr\", 5e-5, 1e-3, log=True)\n",
    "        },\n",
    "        pl_trainer_kwargs=pl_trainer_kwargs,\n",
    "        batch_size=trial.suggest_categorical(\"batch_size\", [16, 32, 64]),\n",
    "        save_checkpoints=True,\n",
    "        force_reset=True,\n",
    "    )\n",
    "\n",
    "    return on.Model(tcn_model)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9342ecf5",
   "metadata": {},
   "source": [
    "## Create objective method\n",
    "\n",
    "The objective method defines the objective function to optimized. Therefore, it must return a metric.\n",
    "We use the BenchmarkEvaluator to compute the metric to optimized, as it is able to directly create the samples according to the defined dataset configuration."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "94698100",
   "metadata": {},
   "outputs": [],
   "source": [
    "from darts.metrics import mse\n",
    "from ontime.module.benchmarking import BenchmarkEvaluator, BenchmarkMetric"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "5d88b2f2-c34f-45ca-b737-ac4307dcdf5e",
   "metadata": {},
   "outputs": [],
   "source": [
    "def objective(trial, metric=mse):\n",
    "    # setup the model\n",
    "    tcn_model = setup_model(trial)\n",
    "\n",
    "    # fit the model\n",
    "    tcn_model.fit(train, val_series=val)\n",
    "\n",
    "    # make predictions\n",
    "    evaluator = BenchmarkEvaluator(dataset, [BenchmarkMetric(\"val_metric\", metric)], on_val_ts=True)\n",
    "\n",
    "    results = evaluator.evaluate(tcn_model, scaler=scaler, scaled_evaluation=True)\n",
    "\n",
    "    return results[\"val_metric\"]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5f7b3df6-f66e-47f8-a49d-517ac1aadf2a",
   "metadata": {},
   "source": [
    "## Run the optimization\n",
    "\n",
    "Once the optimization is run, you can access to the dashboard on your web browser using the command `optuna-dashboard sqlite:///relative/path/to/file.db`.  \n",
    "If you are using Visual Studio Code, you can install the [Optuna Dashboard extension](https://marketplace.visualstudio.com/items/?itemName=Optuna.optuna-dashboard), right click on your db file and select \"Open in Optuna Dashboard\". However, dashboard in VS Code is not as complete as the one opened in the browser."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "c100bcf2",
   "metadata": {},
   "outputs": [],
   "source": [
    "import optuna"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a83dfc4e",
   "metadata": {},
   "outputs": [],
   "source": [
    "study = optuna.create_study(\n",
    "    storage=\"sqlite:///db.sqlite3\", # for dashboard\n",
    "    study_name=\"tcn_optuna_tuto\",\n",
    "    load_if_exists=True,\n",
    "    direction=\"minimize\")\n",
    "study.optimize(objective, n_trials=5)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "3b2199d4",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Best trial: 0.16961007783457446\n",
      "Best hyperparameters: {'output_chunk_length': 6, 'kernel_size': 4, 'num_filters': 61, 'num_layers': 2, 'dropout': 0.2106525223389889, 'weight_norm': True, 'lr': 0.0009837089623902606, 'batch_size': 32}\n"
     ]
    }
   ],
   "source": [
    "print(f\"Best trial: {study.best_trial.value}\")\n",
    "print(f\"Best hyperparameters: {study.best_trial.params}\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "ontime-LQHiZBFd-py3.11",
   "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.11.12"
  },
  "papermill": {
   "default_parameters": {},
   "duration": 60.248854,
   "end_time": "2024-01-31T17:51:31.161244",
   "environment_variables": {},
   "exception": null,
   "input_path": "docs/user_guide/0_core/0.1-time-series-custom-class.ipynb",
   "output_path": "docs/user_guide/0_core/0.1-time-series-custom-class.ipynb",
   "parameters": {},
   "start_time": "2024-01-31T17:50:30.912390",
   "version": "2.5.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}