{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# default_exp models.esrnn.mqesrnn"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#hide\n",
    "%load_ext autoreload\n",
    "%autoreload 2"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# MQESRNN model\n",
    "\n",
    "> API details."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# export\n",
    "from nixtlats.models.esrnn.esrnn import ESRNN"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# export\n",
    "class MQESRNN(ESRNN):\n",
    "    def __init__(self,\n",
    "                 n_series,\n",
    "                 n_x,\n",
    "                 n_s,\n",
    "                 sample_freq,\n",
    "                 input_size,\n",
    "                 output_size,\n",
    "                 es_component,\n",
    "                 cell_type,\n",
    "                 state_hsize,\n",
    "                 dilations,\n",
    "                 add_nl_layer,\n",
    "                 learning_rate,\n",
    "                 lr_scheduler_step_size,\n",
    "                 lr_decay,\n",
    "                 gradient_eps,\n",
    "                 gradient_clipping_threshold,\n",
    "                 rnn_weight_decay,\n",
    "                 noise_std,\n",
    "                 testing_percentiles,\n",
    "                 training_percentiles,\n",
    "                 loss,\n",
    "                 val_loss):\n",
    "        \n",
    "        allowed_losses = ['MQ', 'wMQ']\n",
    "        if loss not in allowed_losses:\n",
    "            raise Exception(f'Loss {loss} not allowed')\n",
    "            \n",
    "        if val_loss is not None and val_loss not in allowed_losses:\n",
    "            raise Exception(f'Val loss {val_loss} not allowed')\n",
    "        \n",
    "        allowed_es_component = ['median_residual', 'identity']\n",
    "        \n",
    "        if es_component not in allowed_es_component:\n",
    "            raise Exception(f'es component {es_component} not allowed')\n",
    "        \n",
    "        level_variability_penalty = 0\n",
    "        seasonality = []\n",
    "        per_series_lr_multip = 1\n",
    "        super(MQESRNN, self).__init__(n_series=n_series,\n",
    "                                      n_s=n_s,\n",
    "                                      n_x=n_x,\n",
    "                                      sample_freq=sample_freq,\n",
    "                                      input_size=input_size,\n",
    "                                      output_size=output_size,\n",
    "                                      es_component=es_component,\n",
    "                                      cell_type=cell_type,\n",
    "                                      state_hsize=state_hsize,\n",
    "                                      dilations=dilations,\n",
    "                                      add_nl_layer=add_nl_layer,\n",
    "                                      seasonality=seasonality,\n",
    "                                      learning_rate=learning_rate,\n",
    "                                      lr_scheduler_step_size=lr_scheduler_step_size,\n",
    "                                      lr_decay=lr_decay,\n",
    "                                      per_series_lr_multip=per_series_lr_multip,\n",
    "                                      gradient_eps=gradient_eps,\n",
    "                                      gradient_clipping_threshold=gradient_clipping_threshold,\n",
    "                                      rnn_weight_decay=rnn_weight_decay,\n",
    "                                      noise_std=noise_std,\n",
    "                                      level_variability_penalty=level_variability_penalty,\n",
    "                                      testing_percentile=testing_percentiles,\n",
    "                                      training_percentile=training_percentiles,\n",
    "                                      loss=loss,\n",
    "                                      val_loss=val_loss)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Tests MQESRNN "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import pandas as pd\n",
    "import pytorch_lightning as pl\n",
    "import matplotlib.pyplot as plt\n",
    "import torch as t\n",
    "\n",
    "from nixtlats.data.tsdataset import TimeSeriesDataset\n",
    "from nixtlats.data.tsloader import TimeSeriesLoader"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "n_ds = 100 \n",
    "n_ts = 1_000\n",
    "\n",
    "output_size = 2\n",
    "\n",
    "uids = [f'uid_{i + 1}' for i in range(n_ts)]\n",
    "dss = pd.date_range(end='2020-01-01', periods=n_ds)\n",
    "\n",
    "Y_df = pd.DataFrame({'unique_id': np.repeat(uids, n_ds), 'ds': np.tile(dss, n_ts)})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "np.random.seed(10)\n",
    "Y_df['y'] = Y_df.groupby('unique_id').transform(lambda x: np.random.uniform(1, 100, size=len(x)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "Y_df = Y_df.sort_values(['unique_id', 'ds'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_dataset = TimeSeriesDataset(Y_df=Y_df,\n",
    "                                  ds_in_test=3*output_size,\n",
    "                                  is_test=False,\n",
    "                                  input_size=7*output_size,\n",
    "                                  output_size=output_size,\n",
    "                                  verbose=True)\n",
    "\n",
    "valid_dataset = TimeSeriesDataset(Y_df=Y_df,\n",
    "                                  input_size=7*output_size,\n",
    "                                  output_size=output_size,\n",
    "                                  verbose=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_loader = TimeSeriesLoader(dataset=train_dataset,\n",
    "                                batch_size=32,\n",
    "                                eq_batch_size=True,\n",
    "                                shuffle=True)\n",
    "\n",
    "valid_loader = TimeSeriesLoader(dataset=valid_dataset,\n",
    "                                batch_size=32,\n",
    "                                shuffle=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = MQESRNN(n_series=train_dataset.n_series,\n",
    "                n_s=train_dataset.n_s,\n",
    "                n_x=train_dataset.n_x,\n",
    "                #sample_freq=dataset.sample_freq,\n",
    "                sample_freq=1,\n",
    "                input_size=7*2,\n",
    "                output_size=output_size,\n",
    "                learning_rate=1e-2,\n",
    "                lr_scheduler_step_size=30,\n",
    "                lr_decay=0.1,\n",
    "                gradient_eps=1e-8,\n",
    "                gradient_clipping_threshold=10,\n",
    "                rnn_weight_decay=0,\n",
    "                noise_std=0.001,\n",
    "                testing_percentiles=[30, 50, 70, 90],\n",
    "                training_percentiles=[30, 50, 70, 90],\n",
    "                es_component='median_residual',\n",
    "                cell_type='LSTM',\n",
    "                state_hsize=100,\n",
    "                dilations=[[1, 2], [4, 8]],\n",
    "                add_nl_layer=False,\n",
    "                loss='MQ',\n",
    "                val_loss='MQ')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "GPU available: True, used: False\n",
      "TPU available: False, using: 0 TPU cores\n",
      "\n",
      "  | Name  | Type   | Params\n",
      "---------------------------------\n",
      "0 | esrnn | _ESRNN | 289 K \n",
      "---------------------------------\n",
      "289 K     Trainable params\n",
      "1         Non-trainable params\n",
      "289 K     Total params\n",
      "1.158     Total estimated model params size (MB)\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "ec0d2d5f46a64ec98ef01a1e62c6168a",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation sanity check: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/ubuntu/anaconda3/envs/nixtla/lib/python3.7/site-packages/pytorch_lightning/utilities/distributed.py:69: UserWarning: The dataloader, val dataloader 0, does not have many workers which may be a bottleneck. Consider increasing the value of the `num_workers` argument` (try 8 which is the number of cpus on this machine) in the `DataLoader` init to improve performance.\n",
      "  warnings.warn(*args, **kwargs)\n",
      "/home/ubuntu/anaconda3/envs/nixtla/lib/python3.7/site-packages/pytorch_lightning/utilities/distributed.py:69: UserWarning: The dataloader, train dataloader, does not have many workers which may be a bottleneck. Consider increasing the value of the `num_workers` argument` (try 8 which is the number of cpus on this machine) in the `DataLoader` init to improve performance.\n",
      "  warnings.warn(*args, **kwargs)\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "0e32c139217444b29aaa44df98c61a98",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Training: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "8968f9a002904e75bc0759c238db1b30",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validating: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "trainer = pl.Trainer(max_epochs=1, progress_bar_refresh_rate=5, deterministic=True)\n",
    "trainer.fit(model, train_loader, valid_loader)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/ubuntu/anaconda3/envs/nixtla/lib/python3.7/site-packages/pytorch_lightning/utilities/distributed.py:69: UserWarning: The dataloader, predict dataloader 0, does not have many workers which may be a bottleneck. Consider increasing the value of the `num_workers` argument` (try 8 which is the number of cpus on this machine) in the `DataLoader` init to improve performance.\n",
      "  warnings.warn(*args, **kwargs)\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "b968fc50da034737ba326fe5b670a8eb",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Predicting: 32it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "outputs = trainer.predict(model, valid_loader)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "y_true, y_hat, mask = zip(*outputs)\n",
    "y_true = t.cat(y_true).numpy()\n",
    "y_hat = t.cat(y_hat).numpy()\n",
    "mask = t.cat(mask).numpy()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "y_true = y_true[-1000:, -1]\n",
    "y_hat = y_hat[-1000:, -1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "calibration p30:  0.303\n",
      "calibration p50:  0.49\n",
      "calibration p70:  0.693\n",
      "calibration p90:  0.952\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "<matplotlib.legend.Legend at 0x7fd3ac18acd0>"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 57688 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "plt.plot(y_true.flatten(), alpha=0.5, label='y')\n",
    "for idx, p in enumerate([30, 50, 70, 90]):\n",
    "    y_p = y_hat[:, :, idx]\n",
    "    plt.plot(y_p.flatten(), alpha=0.5, label=f'p{p}')\n",
    "    print(f'calibration p{p}: ', (y_true.flatten() <= y_p.flatten()).mean())\n",
    "plt.legend()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "nixtla",
   "language": "python",
   "name": "nixtla"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
