{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# default_exp models.esrnn.esrnn"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "The autoreload extension is already loaded. To reload it, use:\n",
      "  %reload_ext autoreload\n"
     ]
    }
   ],
   "source": [
    "#hide\n",
    "%load_ext autoreload\n",
    "%autoreload 2"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# ES-RNN: Exponential Smoothing Recurrent Neural Network\n",
    "\n",
    "> API details.\n",
    "\n",
    "The winning submission of the M4 competition was the <em>Exponential Smoothing Recurrent Neural Network</em> (ESRNN). This hybrid model combined and simultaneously optimized the standard exponential smoothing method as a preprocessing component for the inputs of recurrent neural networks. The [original model](https://github.com/Mcompetitions/M4-methods/tree/master/118%20-%20slaweks17) is implemented in dynet.\n",
    "\n",
    "[Slawek Smyl. A hybrid method of exponential smoothing and recurrent neural networks for time series forecasting. International Journal of Forecasting, 07 2019.](https://www.sciencedirect.com/science/article/pii/S0169207019301153)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "import numpy as np\n",
    "import torch as t\n",
    "import torch.nn as nn\n",
    "\n",
    "from nixtlats.models.components.drnn import DRNN"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "#TODO: rnn con canales\n",
    "#TODO: notacion de todo, windows_y_insample\n",
    "class _ES(nn.Module):\n",
    "    def __init__(self, \n",
    "                 n_series: int, \n",
    "                 input_size: int,\n",
    "                 output_size: int,\n",
    "                 output_size_m: int,\n",
    "                 n_t: int, n_s: int, seasonality: list, noise_std: float):\n",
    "        super(_ES, self).__init__()\n",
    "\n",
    "        self.n_series = n_series\n",
    "        self.input_size = input_size\n",
    "        self.output_size = output_size\n",
    "        self.output_size_m = output_size_m\n",
    "        self.n_t = n_t\n",
    "        self.n_s = n_s\n",
    "        self.seasonality = seasonality\n",
    "        assert len(self.seasonality) in [0, 1, 2]\n",
    "\n",
    "        self.noise_std = noise_std\n",
    "\n",
    "    def gaussian_noise(self, Y: t.Tensor, std: float=0.2):\n",
    "        size = Y.size()\n",
    "        noise = t.autograd.Variable(Y.data.new(size).normal_(0, std))\n",
    "        return Y + noise\n",
    "\n",
    "    def compute_levels_seasons(self, Y: t.Tensor, idxs: t.Tensor):\n",
    "        pass\n",
    "\n",
    "    def normalize(self, Y: t.Tensor, level: t.Tensor, seasonalities: t.Tensor, start: int, end: int):\n",
    "        pass\n",
    "\n",
    "    def predict(self, trend: t.Tensor, levels: t.Tensor, seasonalities: t.Tensor):\n",
    "        pass\n",
    "\n",
    "    def forward(self, S: t.Tensor, Y: t.Tensor, X: t.Tensor, \n",
    "                idxs: t.Tensor, sample_mask: t.Tensor):\n",
    "        # parse attributes\n",
    "        input_size = self.input_size\n",
    "        output_size = self.output_size\n",
    "        n_t = self.n_t\n",
    "        n_s = self.n_s\n",
    "        context_size = input_size + (input_size+output_size)*n_t + n_s\n",
    "        noise_std = self.noise_std\n",
    "        seasonality = self.seasonality\n",
    "        batch_size = len(idxs)\n",
    "\n",
    "        n_series, n_time = Y.shape\n",
    "\n",
    "        # Explicacion: windows_end es el idx delultimo inicio de windows. Como se necesitan windows completas con input\n",
    "        # + output, se pierden (input_size+output_size-1) windows del len total de la serie.\n",
    "        windows_end = n_time - input_size - output_size + 1\n",
    "        #windows_range = range(0, windows_end, step_size)\n",
    "        windows_range = range(0, windows_end)\n",
    "        n_windows = len(windows_range)\n",
    "        \n",
    "        assert n_windows>0\n",
    "\n",
    "        # Initialize windows, levels and seasonalities\n",
    "        levels, seasonalities = self.compute_levels_seasons(Y=Y, idxs=idxs)\n",
    "        windows_y_insample = t.zeros((n_windows, batch_size, context_size),\n",
    "                                     device=Y.device)\n",
    "        windows_y_outsample = t.zeros((n_windows, batch_size, output_size),\n",
    "                                      device=Y.device)\n",
    "\n",
    "        for i, window in enumerate(windows_range):\n",
    "            # Windows yhat\n",
    "            y_insample_start = window\n",
    "            y_insample_end   = input_size + window\n",
    "\n",
    "            # Y_hat deseasonalization and normalization\n",
    "            window_y_insample = self.normalize(Y=Y[:, y_insample_start:y_insample_end],\n",
    "                                               level=levels[:, [y_insample_end-1]],\n",
    "                                               seasonalities=seasonalities,\n",
    "                                               start=y_insample_start, end=y_insample_end) #TODO: improve this inputs\n",
    "\n",
    "            if self.training:\n",
    "                window_y_insample = self.gaussian_noise(window_y_insample, std=noise_std)\n",
    "\n",
    "            if n_t > 0:\n",
    "                window_x_t = X[:, :, y_insample_start:(y_insample_end+output_size)]\n",
    "                window_x_t = window_x_t.reshape(batch_size, -1)\n",
    "                window_y_insample = t.cat((window_y_insample, window_x_t), 1)\n",
    "\n",
    "            # Concatenate S static variables matrix\n",
    "            if n_s > 0:\n",
    "                window_y_insample = t.cat((window_y_insample, S), 1)\n",
    "\n",
    "            windows_y_insample[i, :, :] += window_y_insample\n",
    "\n",
    "            # Windows_y_outsample\n",
    "            y_outsample_start = y_insample_end\n",
    "            y_outsample_end = y_outsample_start + output_size\n",
    "            window_y_outsample = Y[:, y_outsample_start:y_outsample_end]\n",
    "            # If training, normalize outsample for loss on normalized data\n",
    "            if self.training:\n",
    "                # Y deseasonalization and normalization\n",
    "                window_y_outsample = self.normalize(Y=window_y_outsample,\n",
    "                                                    level=levels[:, [y_outsample_start]],\n",
    "                                                    seasonalities=seasonalities,\n",
    "                                                    start=y_outsample_start, end=y_outsample_end) #TODO: improve this inputs\n",
    "            windows_y_outsample[i, :, :] += window_y_outsample\n",
    "        \n",
    "        \n",
    "        # Wrangles the sample_mask with same indexes as windows_y_outsample\n",
    "        n_windows, n_batch, n_output = windows_y_outsample.shape\n",
    "        sample_mask = sample_mask.unfold(dimension=-1, size=self.output_size, step=1)\n",
    "        sample_mask = sample_mask[:, -n_windows:, :]\n",
    "        sample_mask = sample_mask.permute(1,0,2)\n",
    "        \n",
    "        return windows_y_insample, windows_y_outsample, levels, seasonalities, sample_mask"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "class _ESI(_ES):\n",
    "    def __init__(self, n_series: int, input_size: int, output_size: int,\n",
    "                 output_size_m: int,\n",
    "                 n_t: int, n_s: int, seasonality: list, noise_std: float):\n",
    "        super(_ESI, self).__init__(n_series, input_size, output_size, \n",
    "                                   output_size_m,\n",
    "                                   n_t, n_s, seasonality, noise_std)\n",
    "        self.W = t.nn.Parameter(t.randn(1))\n",
    "        self.W.requires_grad = False\n",
    "\n",
    "    def compute_levels_seasons(self, Y: t.Tensor, idxs: t.Tensor):\n",
    "        levels = t.ones(Y.shape, device=Y.device)\n",
    "        seasonalities = None\n",
    "        return levels, None\n",
    "    \n",
    "    def normalize(self, Y: t.Tensor, level: t.Tensor, seasonalities: t.Tensor, \n",
    "                  start: int, end: int):\n",
    "        return Y\n",
    "\n",
    "    def predict(self, \n",
    "                trends: t.Tensor, \n",
    "                levels: t.Tensor, \n",
    "                seasonalities: t.Tensor):\n",
    "        return trends"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "class _MedianResidual(_ES):\n",
    "    def __init__(self, n_series: int, input_size: int, output_size: int, \n",
    "                 output_size_m: int,\n",
    "                 n_t: int, n_s: int, seasonality: list, noise_std: float):\n",
    "        super(_MedianResidual, self).__init__(n_series, input_size, output_size, output_size_m,\n",
    "                                   n_t, n_s, seasonality, noise_std)\n",
    "        self.W = t.nn.Parameter(t.randn(1))\n",
    "        self.W.requires_grad = False\n",
    "\n",
    "    def compute_levels_seasons(self, Y: t.Tensor, idxs: t.Tensor):\n",
    "        \"\"\"\n",
    "        Computes levels and seasons\n",
    "        \"\"\"\n",
    "        y_transformed, _ = Y.median(1)\n",
    "        y_transformed = y_transformed.reshape(-1, 1)\n",
    "        levels = y_transformed.repeat(1, Y.shape[1])\n",
    "        seasonalities = None\n",
    "        \n",
    "        return levels, None\n",
    "    \n",
    "    def normalize(self, Y: t.Tensor, level: t.Tensor, \n",
    "                  seasonalities: t.Tensor, \n",
    "                  start: int, end: int):\n",
    "        \n",
    "        return Y - level\n",
    "\n",
    "    def predict(self, \n",
    "                trends: t.Tensor, \n",
    "                levels: t.Tensor, \n",
    "                seasonalities: t.Tensor):\n",
    "        levels = levels[:, (self.input_size-1):-self.output_size]\n",
    "        levels = levels.unsqueeze(2)\n",
    "\n",
    "        return trends + levels"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "class _ESM(_ES):\n",
    "    def __init__(self, n_series: int, input_size: int, output_size: int,\n",
    "                 output_size_m: int,\n",
    "                 n_t: int, n_s: int, seasonality: list, noise_std: float):\n",
    "        super(_ESM, self).__init__(n_series, input_size, output_size, \n",
    "                                   output_size_m, \n",
    "                                   n_t, n_s, seasonality, noise_std)\n",
    "        # Level and Seasonality Smoothing parameters\n",
    "        # 1 level, S seasonalities, S init_seas\n",
    "        embeds_size = 1 + len(self.seasonality) + sum(self.seasonality)\n",
    "        init_embeds = t.ones((self.n_series, embeds_size)) * 0.5\n",
    "        self.embeds = nn.Embedding(self.n_series, embeds_size)\n",
    "        self.embeds.weight.data.copy_(init_embeds)\n",
    "        self.seasonality = t.LongTensor(self.seasonality)\n",
    "\n",
    "    def compute_levels_seasons(self, Y: t.Tensor, idxs: t.Tensor):\n",
    "        \"\"\"\n",
    "        Computes levels and seasons\n",
    "        \"\"\"\n",
    "        # Lookup parameters per serie\n",
    "        #seasonality = self.seasonality\n",
    "        embeds = self.embeds(idxs)\n",
    "        lev_sms = t.sigmoid(embeds[:, 0])\n",
    "\n",
    "        # Initialize seasonalities\n",
    "        seas_prod = t.ones(len(Y[:,0]), device=Y.device)\n",
    "        seasonalities1 = []\n",
    "        seasonalities2 = []\n",
    "        seas_sms1 = t.ones(1, device=Y.device)\n",
    "        seas_sms2 = t.ones(1, device=Y.device)\n",
    "\n",
    "        if len(self.seasonality)>0:\n",
    "            seas_sms1 = t.sigmoid(embeds[:, 1])\n",
    "            init_seas1 = t.exp(embeds[:, 2:(2+self.seasonality[0])]).unbind(1)\n",
    "            assert len(init_seas1) == self.seasonality[0]\n",
    "\n",
    "            for i in range(len(init_seas1)):\n",
    "                seasonalities1 += [init_seas1[i]]\n",
    "            seasonalities1 += [init_seas1[0]]\n",
    "            seas_prod = seas_prod * init_seas1[0]\n",
    "\n",
    "        if len(self.seasonality)==2:\n",
    "            seas_sms2 = t.sigmoid(embeds[:, 2+self.seasonality[0]])\n",
    "            init_seas2 = t.exp(embeds[:, 3+self.seasonality[0]:]).unbind(1)\n",
    "            assert len(init_seas2) == self.seasonality[1]\n",
    "\n",
    "            for i in range(len(init_seas2)):\n",
    "                seasonalities2 += [init_seas2[i]]\n",
    "            seasonalities2 += [init_seas2[0]]\n",
    "            seas_prod = seas_prod * init_seas2[0]\n",
    "\n",
    "        # Initialize levels\n",
    "        levels = []\n",
    "        levels += [Y[:,0]/seas_prod]\n",
    "\n",
    "        # Recursive seasonalities and levels\n",
    "        ys = Y.unbind(1)\n",
    "        n_time = len(ys)\n",
    "        for t_idx in range(1, n_time):\n",
    "            seas_prod_t = t.ones(len(Y[:,t_idx]), device=Y.device)\n",
    "            if len(self.seasonality)>0:\n",
    "                seas_prod_t = seas_prod_t * seasonalities1[t_idx]\n",
    "            if len(self.seasonality)==2:\n",
    "                seas_prod_t = seas_prod_t * seasonalities2[t_idx]\n",
    "\n",
    "            newlev = lev_sms * (ys[t_idx] / seas_prod_t) + (1-lev_sms) * levels[t_idx-1]\n",
    "            levels += [newlev]\n",
    "\n",
    "            if len(self.seasonality)==1:\n",
    "                newseason1 = seas_sms1 * (ys[t_idx] / newlev) + (1-seas_sms1) * seasonalities1[t_idx]\n",
    "                seasonalities1 += [newseason1]\n",
    "\n",
    "            if len(self.seasonality)==2:\n",
    "                newseason1 = seas_sms1 * (ys[t_idx] / (newlev * seasonalities2[t_idx])) + \\\n",
    "                                         (1-seas_sms1) * seasonalities1[t_idx]\n",
    "                seasonalities1 += [newseason1]\n",
    "                newseason2 = seas_sms2 * (ys[t_idx] / (newlev * seasonalities1[t_idx])) + \\\n",
    "                                         (1-seas_sms2) * seasonalities2[t_idx]\n",
    "                seasonalities2 += [newseason2]\n",
    "\n",
    "        levels = t.stack(levels).transpose(1,0)\n",
    "\n",
    "        seasonalities = []\n",
    "\n",
    "        if len(self.seasonality)>0:\n",
    "            seasonalities += [t.stack(seasonalities1).transpose(1,0)]\n",
    "\n",
    "        if len(self.seasonality)==2:\n",
    "            seasonalities += [t.stack(seasonalities2).transpose(1,0)]\n",
    "\n",
    "        return levels, seasonalities\n",
    "\n",
    "    def normalize(self, Y: t.Tensor, level: t.Tensor, seasonalities: t.Tensor, \n",
    "                  start: int, end: int):\n",
    "        # Deseasonalization and normalization\n",
    "        y_n = Y / level\n",
    "        for s in range(len(self.seasonality)):\n",
    "            y_n /= seasonalities[s][:, start:end]\n",
    "        y_n = t.log(y_n)\n",
    "        return y_n\n",
    "\n",
    "    def predict(self, trends: t.Tensor, levels: t.Tensor, seasonalities: t.Tensor):\n",
    "\n",
    "        # First trend uses last value of first y_insample of length self.input_size.\n",
    "        # Last self.output_size levels are not used (leakeage!!!)\n",
    "        levels = levels[:, (self.input_size-1):-self.output_size]\n",
    "        levels = levels.unsqueeze(2)\n",
    "\n",
    "        # Seasonalities are unfolded, because each element of trends must be multiplied\n",
    "        # by the corresponding seasonality.\n",
    "        for i in range(len(seasonalities)):\n",
    "            seasonalities[i] = seasonalities[i][:, self.seasonality[i]:]\n",
    "            seasonalities[i] = seasonalities[i].unfold(dimension=-1, \n",
    "                                                       size=self.input_size + self.output_size,\n",
    "                                                       step=1)\n",
    "            seasonalities[i] = seasonalities[i][:,:,:self.input_size] #avoid leakage\n",
    "            \n",
    "            # Fill seasonalities with NaiveSeasonal, to avoid leakage.\n",
    "            if self.output_size > seasonalities[i].shape[2]:\n",
    "                repetitions = int(np.ceil(self.output_size / seasonalities[i].shape[2]))\n",
    "                seasonalities[i] = seasonalities[i].repeat((1, 1, repetitions))\n",
    "            seasonalities[i] = seasonalities[i][:, :, :self.output_size]\n",
    "\n",
    "        # Denormalize\n",
    "        trends = t.exp(trends)\n",
    "        # Deseasonalization and normalization (inverse)\n",
    "        y_hat = trends * levels\n",
    "        for s in range(len(self.seasonality)):\n",
    "            seas = seasonalities[s]\n",
    "            y_hat *= t.vstack([seas.T for _ in range(self.output_size_m)]).T\n",
    "\n",
    "        return y_hat"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "class _RNN(nn.Module):\n",
    "    def __init__(self, input_size: int, output_size: int,\n",
    "                 output_size_m: int,\n",
    "                 n_t: int, n_s: int, cell_type: str, dilations: list, state_hsize: int, add_nl_layer: bool):\n",
    "        super(_RNN, self).__init__()\n",
    "\n",
    "        self.input_size = input_size\n",
    "        self.output_size = output_size\n",
    "        self.output_size_m = output_size_m\n",
    "        self.n_t = n_t\n",
    "        self.n_s = n_s\n",
    "\n",
    "        self.cell_type = cell_type\n",
    "        self.dilations = dilations\n",
    "        self.state_hsize = state_hsize\n",
    "        self.add_nl_layer = add_nl_layer\n",
    "        self.layers = len(dilations)\n",
    "\n",
    "        layers = []\n",
    "        for grp_num in range(len(self.dilations)):\n",
    "            if grp_num == 0:\n",
    "                input_size = self.input_size + (self.input_size + self.output_size)*self.n_t + self.n_s\n",
    "            else:\n",
    "                input_size = self.state_hsize\n",
    "            layer = DRNN(input_size,\n",
    "                         self.state_hsize,\n",
    "                         n_layers=len(self.dilations[grp_num]),\n",
    "                         dilations=self.dilations[grp_num],\n",
    "                         cell_type=self.cell_type)\n",
    "            layers.append(layer)\n",
    "\n",
    "        self.rnn_stack = nn.Sequential(*layers)\n",
    "\n",
    "        if self.add_nl_layer:\n",
    "            self.MLPW  = nn.Linear(self.state_hsize, self.state_hsize)\n",
    "\n",
    "        self.adapterW  = nn.Linear(self.state_hsize, self.output_size * self.output_size_m)\n",
    "\n",
    "    def forward(self, input_data: t.Tensor):\n",
    "        for layer_num in range(len(self.rnn_stack)):\n",
    "            residual = input_data\n",
    "            output, _ = self.rnn_stack[layer_num](input_data)\n",
    "            if layer_num > 0:\n",
    "                output += residual\n",
    "            input_data = output\n",
    "\n",
    "        if self.add_nl_layer:\n",
    "            input_data = self.MLPW(input_data)\n",
    "            input_data = t.tanh(input_data)\n",
    "\n",
    "        input_data = self.adapterW(input_data)\n",
    "        \n",
    "        return input_data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "class _ESRNN(nn.Module):\n",
    "    def __init__(self, n_series, input_size, output_size, \n",
    "                 output_size_m, n_t, n_s,\n",
    "                 es_component, seasonality, noise_std, cell_type,\n",
    "                 dilations, state_hsize, add_nl_layer):\n",
    "        super(_ESRNN, self).__init__()\n",
    "        allowed_componets = ['multiplicative', 'identity', 'median_residual']\n",
    "        assert es_component in allowed_componets, f'es_component {es_component} not valid.'\n",
    "        self.es_component = es_component\n",
    "\n",
    "        if es_component == 'multiplicative':\n",
    "            self.es = _ESM(n_series=n_series, input_size=input_size, output_size=output_size,\n",
    "                           output_size_m=output_size_m, n_t=n_t, n_s=n_s, \n",
    "                           seasonality=seasonality, noise_std=noise_std)\n",
    "        elif es_component == 'identity':\n",
    "            self.es = _ESI(n_series=n_series, input_size=input_size, output_size=output_size,\n",
    "                           output_size_m=output_size_m, n_t=n_t, n_s=n_s, \n",
    "                           seasonality=seasonality, noise_std=noise_std)\n",
    "        elif es_component == 'median_residual':\n",
    "            self.es = _MedianResidual(n_series=n_series, input_size=input_size, output_size=output_size,\n",
    "                                      output_size_m=output_size_m, n_t=n_t, n_s=n_s, \n",
    "                                      seasonality=seasonality, noise_std=noise_std)\n",
    "            \n",
    "        self.rnn = _RNN(input_size=input_size, output_size=output_size,\n",
    "                        output_size_m=output_size_m,\n",
    "                        n_t=n_t, n_s=n_s,\n",
    "                        cell_type=cell_type, dilations=dilations, \n",
    "                        state_hsize=state_hsize,\n",
    "                        add_nl_layer=add_nl_layer)\n",
    "\n",
    "    def forward(self, S: t.Tensor, Y: t.Tensor, X: t.Tensor, \n",
    "                idxs: t.Tensor, sample_mask: t.Tensor):\n",
    "        # Multiplicative model protection\n",
    "        if self.es_component == 'multiplicative' and t.min(Y) == 0:\n",
    "            raise Exception(\n",
    "                'Check your Y data, multiplicative model only deals with Y>0. \\n'\n",
    "                f'Series: {idxs} \\n'\n",
    "                f'Min Y: {t.min(Y)}'\n",
    "            )\n",
    "        \n",
    "        # ES Forward\n",
    "        y_in, y_out, levels, seasonalities, sample_mask = self.es(S=S,Y=Y,X=X,idxs=idxs,\n",
    "                                                                  sample_mask=sample_mask)\n",
    "        \n",
    "        # RNN Forward\n",
    "        y_hat = self.rnn(y_in)\n",
    "        \n",
    "        if self.rnn.output_size_m > 1:\n",
    "            n_w, n_ts, _ = y_in.shape\n",
    "            y_hat = y_hat.view(n_w, n_ts, -1, self.rnn.output_size_m)\n",
    "\n",
    "        return y_out, y_hat, levels, sample_mask\n",
    "\n",
    "    def predict(self, S: t.Tensor, Y: t.Tensor, X: t.Tensor,\n",
    "                idxs: t.Tensor, sample_mask: t.Tensor):\n",
    "        # ES Forward\n",
    "        y_in, y_out, levels, seasonalities, sample_mask = self.es(S=S,Y=Y,X=X,idxs=idxs,\n",
    "                                                                  sample_mask=sample_mask)\n",
    "        # RNN Forward\n",
    "        trends = self.rnn(y_in)\n",
    "\n",
    "        # (n_windows, n_batch, n_input) -> (n_batch, n_windows, n_input)\n",
    "        y_out = y_out.permute(1,0,2)\n",
    "        trends = trends.permute(1,0,2)\n",
    "        sample_mask = sample_mask.permute(1,0,2)\n",
    "\n",
    "        y_hat = self.es.predict(trends, levels, seasonalities)\n",
    "        \n",
    "        if self.rnn.output_size_m > 1:\n",
    "            n_ts, n_w, _ = y_out.shape\n",
    "            y_hat = y_hat.view(n_ts, n_w, -1, self.rnn.output_size_m)\n",
    "\n",
    "        return y_out, y_hat, sample_mask"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Test ES-RNN model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def test_no_leakage_season(n_series, input_size, output_size, seasonality):\n",
    "    \"\"\"\n",
    "    This test checks no leakage for seasonality.\n",
    "    \"\"\"\n",
    "    t.manual_seed(1)\n",
    "    model = _ESRNN(n_series=n_series, input_size=input_size, output_size=output_size, \n",
    "                   output_size_m=1, n_t=0, n_s=0, \n",
    "                   es_component='multiplicative', seasonality=seasonality, noise_std=0.000000001,\n",
    "                   cell_type='GRU',\n",
    "                   add_nl_layer=False, dilations=[[1, 2]], state_hsize=30)\n",
    "    S = t.empty((n_series, 0))\n",
    "    X = t.empty((n_series, 0))\n",
    "    Y = t.normal(0, 1, (n_series, 2 * (input_size + output_size)))\n",
    "    Y += Y.min().abs() + 10\n",
    "    sample_mask = t.ones_like(Y)\n",
    "    sample_mask[:, -output_size] = 0\n",
    "    idxs = t.arange(n_series)\n",
    "    \n",
    "    # Testing different values for output_size\n",
    "    # forecasts should be the same in all cases\n",
    "    # except for numeric exceptions\n",
    "    Y_to_test = [Y] * 4\n",
    "    Y_to_test[1][:, -output_size:] = 10_000\n",
    "    Y_to_test[2][:, -output_size:] = 0\n",
    "    Y_to_test[3][:, -output_size:] = -10_000\n",
    "    \n",
    "    forecasts = []\n",
    "    for Y in Y_to_test:\n",
    "        # forward es and rnn\n",
    "        windows_y_insample, windows_y_outsample, levels, seasonalities, sample_mask_w = model.es(S, Y, X, idxs, sample_mask=sample_mask)\n",
    "        trends = model.rnn(windows_y_insample)\n",
    "        trends = trends.permute(1, 0, 2)\n",
    "\n",
    "        # predict es\n",
    "        y_hat = model.es.predict(trends=trends, levels=levels, seasonalities=seasonalities)\n",
    "\n",
    "        forecasts.append(y_hat)\n",
    "        \n",
    "    assert all(t.allclose(forecasts[0], forecast) for forecast in forecasts), (\n",
    "        'Season leakage detected ',\n",
    "        'please check.'\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_no_leakage_season(n_series=10, input_size=10, output_size=8, seasonality=[6])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_no_leakage_season(n_series=10, input_size=10, output_size=2, seasonality=[6])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_no_leakage_season(n_series=10, input_size=10, output_size=8, seasonality=[12, 24])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_no_leakage_season(n_series=10, input_size=10, output_size=8, seasonality=[6, 12])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_no_leakage_season(n_series=10, input_size=10, output_size=18, seasonality=[12])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_no_leakage_season(n_series=10, input_size=10, output_size=4, seasonality=[1])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_no_leakage_season(n_series=10, input_size=10, output_size=18, seasonality=[12, 24])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## ES-RNN model wrapper"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "from typing import Union, List\n",
    "\n",
    "import pytorch_lightning as pl\n",
    "import torch.nn as nn\n",
    "from torch.nn.utils import clip_grad_norm_\n",
    "from torch.optim import Adam\n",
    "from torch.optim.lr_scheduler import StepLR\n",
    "\n",
    "from nixtlats.losses.utils import LossFunction"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "class ESRNN(pl.LightningModule):\n",
    "    def __init__(self, n_series: int,\n",
    "                 n_x: int, n_s: int, \n",
    "                 input_size: int, output_size: int,\n",
    "                 sample_freq: int,\n",
    "                 es_component: str = 'multiplicative', \n",
    "                 cell_type: str = 'LSTM', state_hsize: int = 50, \n",
    "                 dilations: List[List[int]] = [[1, 2], [4, 8]], \n",
    "                 add_nl_layer: bool = False, seasonality: List[int] = [],\n",
    "                 learning_rate: float = 1e-3, lr_scheduler_step_size: int = 9,\n",
    "                 lr_decay: float = 0.9, per_series_lr_multip: float = 1.,\n",
    "                 gradient_eps: float = 1e-8, \n",
    "                 gradient_clipping_threshold: float = 20.,\n",
    "                 rnn_weight_decay: float = 0., noise_std: float = 1e-3,\n",
    "                 level_variability_penalty: float = 20.,\n",
    "                 testing_percentile: Union[int, List] = 50, \n",
    "                 training_percentile: Union[int, List] = 50,\n",
    "                 loss: str = 'SMYL', val_loss: str = 'MAE'):\n",
    "        super(ESRNN, self).__init__()\n",
    "        \"\"\" Exponential Smoothing Recurrent Neural Network\n",
    "\n",
    "        Pytorch Implementation of the M4 time series forecasting competition winner.\n",
    "        Proposed by Smyl. The model uses a hybrid approach of Machine Learning and\n",
    "        statistical methods by combining recurrent neural networks to model a common\n",
    "        trend with shared parameters across series, and multiplicative Holt-Winter\n",
    "        exponential smoothing.\n",
    "\n",
    "        Parameters\n",
    "        ----------\n",
    "        n_series: int\n",
    "            Number of time series.\n",
    "        n_x: int\n",
    "            Number of temporal exogenous variables.\n",
    "        n_s: int\n",
    "            Number of static variables.\n",
    "        input_size: int\n",
    "            input size of the recurrent neural network, usually a\n",
    "            multiple of seasonality\n",
    "        output_size: int\n",
    "            output_size or forecast horizon of the recurrent neural\n",
    "            network, usually multiple of seasonality\n",
    "        sample_freq: int\n",
    "            Step size between windows.\n",
    "        es_component: str\n",
    "            Exponential Smoothing component.\n",
    "            Default multiplicative.\n",
    "        cell_type: str\n",
    "            Type of RNN cell, available GRU, LSTM, RNN, ResidualLSTM.\n",
    "        state_hsize: int\n",
    "            dimension of hidden state of the recurrent neural network\n",
    "        dilations: int list\n",
    "            each list represents one chunk of Dilated LSTMS, connected in\n",
    "            standard ResNet fashion\n",
    "        add_nl_layer: bool\n",
    "            whether to insert a tanh() layer between the RNN stack and the\n",
    "            linear adaptor (output) layers\n",
    "        seasonality: int list\n",
    "            list of seasonalities of the time series\n",
    "            Hourly [24, 168], Daily [7], Weekly [52], Monthly [12],\n",
    "            Quarterly [4], Yearly [].\n",
    "        learning_rate: float\n",
    "            size of the stochastic gradient descent steps\n",
    "        lr_scheduler_step_size: int\n",
    "            this step_size is the period for each learning rate decay\n",
    "        lr_decay: float\n",
    "            Learning rate decay.\n",
    "        per_series_lr_multip: float\n",
    "            multiplier for per-series parameters smoothing and initial\n",
    "            seasonalities learning rate (default 1.0)\n",
    "        gradient_eps: float\n",
    "            term added to the Adam optimizer denominator to improve\n",
    "            numerical stability (default: 1e-8)\n",
    "        gradient_clipping_threshold: float\n",
    "            max norm of gradient vector, with all parameters treated\n",
    "            as a single vector\n",
    "        rnn_weight_decay: float\n",
    "            parameter to control classic L2/Tikhonov regularization\n",
    "            of the rnn parameters\n",
    "        noise_std: float\n",
    "            standard deviation of white noise added to input during\n",
    "            fit to avoid the model from memorizing the train data\n",
    "        level_variability_penalty: float\n",
    "            this parameter controls the strength of the penalization\n",
    "            to the wigglines of the level vector, induces smoothness\n",
    "            in the output\n",
    "        testing_percentile: int\n",
    "            This value is only for diagnostic evaluation.\n",
    "            In case of percentile predictions this parameter controls\n",
    "            for the value predicted, when forecasting point value,\n",
    "            the forecast is the median, so percentile=50.\n",
    "        training_percentile: float\n",
    "            To reduce the model's tendency to over estimate, the\n",
    "            training_percentile can be set to fit a smaller value\n",
    "            through the Pinball Loss.\n",
    "        loss: str\n",
    "            Loss used to train.\n",
    "        val_loss: str\n",
    "            Loss used to validate.\n",
    "        \n",
    "        Notes\n",
    "        -----\n",
    "        **References:**\n",
    "        `M4 Competition Conclusions\n",
    "        <https://rpubs.com/fotpetr/m4competition>`__\n",
    "        `Original Dynet Implementation of ESRNN\n",
    "        <https://github.com/M4Competition/M4-methods/tree/master/118%20-%20slaweks17>`__\n",
    "        \"\"\"\n",
    "\n",
    "        #------------------------ Model Attributes ------------------------#\n",
    "        # Architecture parameters\n",
    "        self.n_series = n_series\n",
    "        self.n_x = n_x\n",
    "        self.n_s = n_s \n",
    "        self.sample_freq = sample_freq\n",
    "        self.input_size = input_size\n",
    "        self.output_size = output_size\n",
    "        self.es_component = es_component\n",
    "        self.cell_type = cell_type\n",
    "        self.state_hsize = state_hsize\n",
    "        self.dilations = dilations\n",
    "        self.add_nl_layer = add_nl_layer\n",
    "        self.seasonality = seasonality\n",
    "\n",
    "        # Regularization and optimization parameters\n",
    "        self.learning_rate = learning_rate\n",
    "        self.lr_scheduler_step_size = lr_scheduler_step_size\n",
    "        self.lr_decay = lr_decay\n",
    "        self.per_series_lr_multip = per_series_lr_multip\n",
    "        self.gradient_eps = gradient_eps\n",
    "        self.gradient_clipping_threshold = gradient_clipping_threshold\n",
    "        self.rnn_weight_decay = rnn_weight_decay\n",
    "        self.noise_std = noise_std\n",
    "        self.level_variability_penalty = level_variability_penalty\n",
    "        self.testing_percentile = testing_percentile\n",
    "        self.training_percentile = training_percentile\n",
    "        self.loss = loss\n",
    "        self.val_loss = val_loss\n",
    "        self.loss_fn = LossFunction(loss, \n",
    "                                    percentile=self.training_percentile,\n",
    "                                    level_variability_penalty=self.level_variability_penalty)\n",
    "        self.val_loss_fn = LossFunction(val_loss,\n",
    "                                        percentile=self.testing_percentile,\n",
    "                                        level_variability_penalty=self.level_variability_penalty)\n",
    "\n",
    "        # MQESRNN\n",
    "        self.mq = isinstance(self.training_percentile, list)\n",
    "        self.output_size_m = len(self.training_percentile) if self.mq else 1\n",
    "\n",
    "        #Defining model\n",
    "        self.model = _ESRNN(n_series=self.n_series, \n",
    "                            input_size=self.input_size,\n",
    "                            output_size=self.output_size, \n",
    "                            output_size_m=self.output_size_m,\n",
    "                            n_t=self.n_x, n_s=self.n_s,\n",
    "                            es_component=self.es_component, \n",
    "                            seasonality=self.seasonality,\n",
    "                            noise_std=self.noise_std, \n",
    "                            cell_type=self.cell_type,\n",
    "                            dilations=self.dilations, \n",
    "                            state_hsize=self.state_hsize,\n",
    "                            add_nl_layer=self.add_nl_layer)\n",
    "        \n",
    "        self.automatic_optimization = False\n",
    "        \n",
    "    def parse_batch(self, batch):\n",
    "        S = batch['S']\n",
    "        Y = batch['Y']\n",
    "        X = batch['X']\n",
    "        idxs = batch['idxs']\n",
    "        sample_mask = batch['sample_mask']\n",
    "        available_mask = batch['available_mask']\n",
    "        \n",
    "        av_condition = t.nonzero(t.min(available_mask, axis=0).values)\n",
    "        min_time_stamp = int(av_condition.min())\n",
    "        sample_condition = t.nonzero(t.min(sample_mask, axis=0).values)\n",
    "        \n",
    "        if sample_condition.nelement() == 0:\n",
    "            max_time_stamp = int(av_condition.max())\n",
    "        else:\n",
    "            max_time_stamp = int(sample_condition.max())\n",
    "        available_ts = max_time_stamp - min_time_stamp + 1 # +1, inclusive counting\n",
    "        if available_ts < self.input_size + self.output_size:\n",
    "            raise Exception(\n",
    "                'Time series too short for given input and output size. \\n'\n",
    "                f'Available timestamps: {available_ts}'\n",
    "            )\n",
    "        \n",
    "        Y = Y[:, min_time_stamp:max_time_stamp + 1] #+1 because is not inclusive\n",
    "        X = X[:, :, min_time_stamp:max_time_stamp + 1]\n",
    "        sample_mask = sample_mask[:, min_time_stamp:max_time_stamp + 1]\n",
    "        available_mask = available_mask[:, min_time_stamp:max_time_stamp + 1]\n",
    "        \n",
    "        return S, Y, X, idxs, sample_mask, available_mask\n",
    "            \n",
    "    def training_step(self, batch, batch_idx):\n",
    "        es_opt, rnn_opt = self.optimizers()\n",
    "        \n",
    "        # Parsing batch\n",
    "        S, Y, X, idxs, sample_mask, available_mask = self.parse_batch(batch)\n",
    "        \n",
    "        target, forecast, levels, sample_mask = self.model(S=S, Y=Y, X=X, idxs=idxs,\n",
    "                                                           sample_mask=sample_mask)\n",
    "        loss = self.loss_fn(y=target,\n",
    "                            y_hat=forecast,\n",
    "                            y_insample=Y, \n",
    "                            levels=levels,\n",
    "                            mask=sample_mask) \n",
    "        \n",
    "        es_opt.zero_grad()\n",
    "        rnn_opt.zero_grad()\n",
    "        self.manual_backward(loss)\n",
    "        clip_grad_norm_(parameters=self.model.rnn.parameters(),\n",
    "                        max_norm=self.gradient_clipping_threshold)\n",
    "        clip_grad_norm_(parameters=self.model.es.parameters(),\n",
    "                        max_norm=self.gradient_clipping_threshold)\n",
    "        es_opt.step()\n",
    "        rnn_opt.step()\n",
    "\n",
    "        self.log('train_loss', loss, prog_bar=True, on_epoch=True)\n",
    "        \n",
    "        return loss\n",
    "    \n",
    "    def on_train_epoch_end(self):\n",
    "        for lr_scheduler in self.lr_schedulers():\n",
    "            lr_scheduler.step()\n",
    "    \n",
    "    def validation_step(self, batch, idx):\n",
    "        # Parsing batch\n",
    "        S, Y, X, idxs, sample_mask, available_mask = self.parse_batch(batch)\n",
    "        \n",
    "        y_true, y_hat, sample_mask = self.model.predict(S=S, Y=Y, X=X, idxs=idxs, \n",
    "                                                        sample_mask=sample_mask)\n",
    "        \n",
    "        y_true = y_true[:, ::self.sample_freq]\n",
    "        y_hat = y_hat[:, ::self.sample_freq]\n",
    "        sample_mask = sample_mask[:, ::self.sample_freq]\n",
    "        \n",
    "        loss = self.val_loss_fn(y=y_true,\n",
    "                                y_hat=y_hat,\n",
    "                                y_insample=Y)\n",
    "        self.log('val_loss', loss, prog_bar=True)\n",
    "        \n",
    "        return loss\n",
    "    \n",
    "    def forward(self, batch):\n",
    "        # Parsing batch\n",
    "        S, Y, X, idxs, sample_mask, available_mask = self.parse_batch(batch)\n",
    "        \n",
    "        y_true, y_hat, sample_mask = self.model.predict(S=S, Y=Y, X=X, idxs=idxs,\n",
    "                                                        sample_mask=sample_mask)\n",
    "        \n",
    "        y_true = y_true[:, ::self.sample_freq]\n",
    "        y_hat = y_hat[:, ::self.sample_freq]\n",
    "        sample_mask = sample_mask[:, ::self.sample_freq]\n",
    "        \n",
    "        return y_true, y_hat, sample_mask\n",
    "    \n",
    "    def configure_optimizers(self):\n",
    "        es_optimizer = Adam(params=self.model.es.parameters(),\n",
    "                            lr=self.learning_rate * self.per_series_lr_multip,\n",
    "                            eps=self.gradient_eps)\n",
    "        rnn_optimizer = Adam(params=self.model.rnn.parameters(),\n",
    "                             lr=self.learning_rate,\n",
    "                             eps=self.gradient_eps,\n",
    "                             weight_decay=self.rnn_weight_decay)\n",
    "        \n",
    "        lr_es = StepLR(optimizer=es_optimizer,\n",
    "                       step_size=self.lr_scheduler_step_size,\n",
    "                       gamma=self.lr_decay)\n",
    "        lr_rnn = StepLR(optimizer=rnn_optimizer,\n",
    "                        step_size=self.lr_scheduler_step_size,\n",
    "                        gamma=self.lr_decay)\n",
    "        \n",
    "        return [es_optimizer, rnn_optimizer], [lr_es, lr_rnn]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### ES-RNN Univariate Example"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "\n",
    "from nixtlats.data.datasets.epf import EPF, EPFInfo\n",
    "from nixtlats.data.tsdataset import TimeSeriesDataset\n",
    "from nixtlats.data.tsloader import TimeSeriesLoader\n",
    "\n",
    "from pytorch_lightning.callbacks import EarlyStopping"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "INFO:root:Train Validation splits\n",
      "\n",
      "INFO:root:                              ds                    \n",
      "                             min                 max\n",
      "unique_id sample_mask                               \n",
      "NP        0           2018-11-25 2018-12-24 23:00:00\n",
      "          1           2018-08-27 2018-11-24 23:00:00\n",
      "INFO:root:\n",
      "Total data \t\t\t2880 time stamps \n",
      "Available percentage=100.0, \t2880 time stamps \n",
      "Insample  percentage=75.0, \t2160 time stamps \n",
      "Outsample percentage=25.0, \t720 time stamps \n",
      "\n",
      "INFO:root:Train Validation splits\n",
      "\n",
      "INFO:root:                              ds                    \n",
      "                             min                 max\n",
      "unique_id sample_mask                               \n",
      "NP        0           2018-08-27 2018-11-24 23:00:00\n",
      "          1           2018-11-25 2018-12-24 23:00:00\n",
      "INFO:root:\n",
      "Total data \t\t\t2880 time stamps \n",
      "Available percentage=100.0, \t2880 time stamps \n",
      "Insample  percentage=25.0, \t720 time stamps \n",
      "Outsample percentage=75.0, \t2160 time stamps \n",
      "\n"
     ]
    }
   ],
   "source": [
    "Y_df, X_df, _ = EPF.load(directory='./data', group=EPFInfo.groups[0])\n",
    "\n",
    "X_df = X_df[['unique_id', 'ds', 'week_day']]\n",
    "\n",
    "# Trimming series to avoid slow backprop through time\n",
    "Y_df = Y_df.groupby('unique_id').tail(90*24+30*24)\n",
    "X_df = X_df.groupby('unique_id').tail(90*24+30*24)\n",
    "\n",
    "# Leveling Y_df (multiplicative model)\n",
    "Y_min = Y_df.y.min()\n",
    "Y_df.y = Y_df.y - Y_min + 20\n",
    "\n",
    "train_dataset = TimeSeriesDataset(Y_df=Y_df, X_df=X_df,\n",
    "                                  ds_in_test=30*24,\n",
    "                                  is_test=False,\n",
    "                                  input_size=7*24,\n",
    "                                  output_size=24,\n",
    "                                  verbose=True)\n",
    "\n",
    "valid_dataset = TimeSeriesDataset(Y_df=Y_df, X_df=X_df,\n",
    "                                  ds_in_test=30*24,\n",
    "                                  is_test=True,\n",
    "                                  input_size=7*24,\n",
    "                                  output_size=24,\n",
    "                                  verbose=True)\n",
    "\n",
    "train_loader = TimeSeriesLoader(dataset=train_dataset,\n",
    "                                num_workers=4,\n",
    "                                batch_size=32,\n",
    "                                shuffle=True)\n",
    "\n",
    "valid_loader = TimeSeriesLoader(dataset=valid_dataset,\n",
    "                                num_workers=4,\n",
    "                                batch_size=1,\n",
    "                                shuffle=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = ESRNN(# Architecture parameters\n",
    "              n_series=train_dataset.n_series,\n",
    "              n_s=train_dataset.n_s,\n",
    "              n_x=train_dataset.n_x,\n",
    "              input_size=train_dataset.input_size,\n",
    "              output_size=train_dataset.output_size,\n",
    "              sample_freq=train_dataset.output_size,\n",
    "              es_component='multiplicative',\n",
    "              cell_type='LSTM',\n",
    "              state_hsize=10,\n",
    "              dilations=[[1, 24], [48, 168]],\n",
    "              add_nl_layer=False,\n",
    "              # Regularization and optimization parameters\n",
    "              learning_rate=5e-3,\n",
    "              lr_scheduler_step_size=100,\n",
    "              lr_decay=0.9,\n",
    "              per_series_lr_multip=1.5,\n",
    "              gradient_eps=1e-8,\n",
    "              gradient_clipping_threshold=10,\n",
    "              rnn_weight_decay=0,\n",
    "              noise_std=0.001,\n",
    "              level_variability_penalty=10,\n",
    "              testing_percentile=50,\n",
    "              training_percentile=51,\n",
    "              loss='SMYL',\n",
    "              val_loss='MAE',\n",
    "              seasonality=[24])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "GPU available: False, used: False\n",
      "TPU available: False, using: 0 TPU cores\n",
      "\n",
      "  | Name  | Type   | Params\n",
      "---------------------------------\n",
      "0 | model | _ESRNN | 17.8 K\n",
      "---------------------------------\n",
      "17.8 K    Trainable params\n",
      "0         Non-trainable params\n",
      "17.8 K    Total params\n",
      "0.071     Total estimated model params size (MB)\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation sanity check: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "860e5950449a4bc589fd42b3bb659df8",
       "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": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validating: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Metric val_loss improved. New best score: 14.236\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validating: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Metric val_loss improved by 1.447 >= min_delta = 0.0001. New best score: 12.790\n"
     ]
    }
   ],
   "source": [
    "early_stopping = EarlyStopping(monitor=\"val_loss\", \n",
    "                               min_delta=1e-4, \n",
    "                               patience=1, verbose=True, \n",
    "                               mode=\"min\")\n",
    "\n",
    "trainer = pl.Trainer(max_epochs=2, progress_bar_refresh_rate=1, \n",
    "                     log_every_n_steps=1, check_val_every_n_epoch=1,\n",
    "                     callbacks=[early_stopping])\n",
    "trainer.fit(model, train_loader, valid_loader)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "3159e313a90d41b987b34d51639035dc",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Predicting: 1it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Original\n",
      "y_true.shape torch.Size([1, 113, 24])\n",
      "y_hat.shape torch.Size([1, 113, 24])\n",
      "\n",
      "Flatten\n",
      "y_true.shape torch.Size([1, 2712])\n",
      "y_hat.shape torch.Size([1, 2712])\n",
      "sample_mask.shape torch.Size([1, 2712])\n"
     ]
    }
   ],
   "source": [
    "outputs = trainer.predict(model, valid_loader)\n",
    "\n",
    "y_true, y_hat, sample_mask = zip(*outputs)\n",
    "y_true = t.cat(y_true).cpu()\n",
    "y_hat = t.cat(y_hat).cpu()\n",
    "sample_mask = t.cat(sample_mask).cpu()\n",
    "\n",
    "print(\"Original\")\n",
    "print(\"y_true.shape\", y_true.shape)\n",
    "print(\"y_hat.shape\", y_hat.shape)\n",
    "y_true = y_true.flatten(1,2)\n",
    "y_hat = y_hat.flatten(1,2)\n",
    "sample_mask = sample_mask.flatten(1,2)\n",
    "\n",
    "print(\"\\nFlatten\")\n",
    "print(\"y_true.shape\", y_true.shape)\n",
    "print(\"y_hat.shape\", y_hat.shape)\n",
    "print(\"sample_mask.shape\", sample_mask.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Returning level to Y_df and Y_hat_df\n",
    "y_true = y_true + Y_min - 20\n",
    "y_hat = y_hat + Y_min - 20"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "mae      13.490134\n",
      "rmae      0.742696\n",
      "smape    23.856249\n",
      "rmse     17.098967\n",
      "dtype: float64\n",
      "\n",
      "\n",
      "DescribeResult(nobs=720, minmax=(-54.182884, 39.697495), mean=-3.9000943, variance=277.5494, skewness=-0.416260302066803, kurtosis=-0.3610564153515341)\n",
      "model.training_percentile 51\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 5300 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "from scipy import stats\n",
    "\n",
    "import pandas as pd\n",
    "\n",
    "from nixtlats.losses.numpy import mae, rmae, smape, rmse\n",
    "from nixtlats.data.datasets.epf import epf_naive_forecast\n",
    "\n",
    "Y_naive_df = epf_naive_forecast(Y_df)\n",
    "\n",
    "# Filter test hours\n",
    "y_true = y_true[0, -30*24:]\n",
    "y_hat = y_hat[0, -30*24:]\n",
    "y_naive = Y_naive_df.y_hat.values[-30*24:]\n",
    "\n",
    "metrics = pd.Series({'mae' :  mae(y=y_true, y_hat=y_hat),\n",
    "                     'rmae':  rmae(y=y_true, y_hat1=y_hat, y_hat2=y_naive),\n",
    "                     'smape': smape(y=y_true, y_hat=y_hat),\n",
    "                     'rmse':  rmse(y=y_true, y_hat=y_hat)})\n",
    "\n",
    "print(metrics)\n",
    "print('\\n')\n",
    "print(stats.describe(y_true-y_hat))\n",
    "print(f'model.training_percentile {model.training_percentile}')\n",
    "\n",
    "plt.plot(sample_mask[0,:])\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 82080 characters]\n",
      "text/plain": [
       "<Figure size 1080x432 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "start = 0\n",
    "end = 7 * 24\n",
    "\n",
    "fig = plt.figure(figsize=(15, 6))\n",
    "plt.plot(y_true[start:end], color='#628793', linewidth=0.4, label='true')\n",
    "plt.plot(y_hat[start:end], color='peru', linewidth=0.4, label='forecast')\n",
    "plt.ylabel('Price [EUR/MWh]', fontsize=15)\n",
    "plt.xlabel('Date', fontsize=15)\n",
    "plt.legend()\n",
    "plt.grid()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## ES-RNN Multivariate Example"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Yearly\n"
     ]
    }
   ],
   "source": [
    "from nixtlats.data.datasets.tourism import Tourism, TourismInfo\n",
    "\n",
    "group = TourismInfo['Yearly']\n",
    "print(group.name)\n",
    "Y_df, *_ = Tourism.load(directory='./data', group=group.name)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fill_uids = Y_df.groupby('unique_id').size().sort_values() \\\n",
    "                .add(-group.horizon-1) \\\n",
    "                .loc[lambda x: x < 2 * group.horizon] \\\n",
    "                .index"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "new_ys = []\n",
    "for uid in fill_uids:\n",
    "    y = Y_df.query('unique_id == @uid')\n",
    "\n",
    "    ds_to_fill = 2 * group.horizon - y.shape[0] + group.horizon + 1\n",
    "\n",
    "    ds = pd.date_range(end = y.iloc[0].ds, periods = ds_to_fill + 1, freq = group.freq)[:ds_to_fill]\n",
    "    new_y = pd.DataFrame({'ds': ds})\n",
    "    new_y['unique_id'] = y.iloc[0].unique_id\n",
    "    new_y['y'] = y.iloc[0].y\n",
    "    \n",
    "    new_ys.append(new_y)\n",
    "new_ys = pd.concat(new_ys)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "Y_df = pd.concat([Y_df, new_ys]).sort_values(['unique_id', 'ds'], ignore_index=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "INFO:root:Train Validation splits\n",
      "\n",
      "INFO:root:                    ds           \n",
      "                   min        max\n",
      "sample_mask                      \n",
      "0           1990-12-31 2007-12-31\n",
      "1           1960-12-31 2003-12-31\n",
      "INFO:root:\n",
      "Total data \t\t\t12708 time stamps \n",
      "Available percentage=100.0, \t12708 time stamps \n",
      "Insample  percentage=83.7, \t10636 time stamps \n",
      "Outsample percentage=16.3, \t2072 time stamps \n",
      "\n",
      "INFO:root:Train Validation splits\n",
      "\n",
      "INFO:root:                    ds           \n",
      "                   min        max\n",
      "sample_mask                      \n",
      "0           1960-12-31 2007-12-31\n",
      "INFO:root:\n",
      "Total data \t\t\t12708 time stamps \n",
      "Available percentage=100.0, \t12708 time stamps \n",
      "Insample  percentage=0.0, \t0 time stamps \n",
      "Outsample percentage=100.0, \t12708 time stamps \n",
      "\n"
     ]
    }
   ],
   "source": [
    "train_dataset = TimeSeriesDataset(Y_df=Y_df,\n",
    "                                  ds_in_test=group.horizon,\n",
    "                                  is_test=False,\n",
    "                                  input_size=group.horizon,\n",
    "                                  output_size=group.horizon,\n",
    "                                  verbose=True)\n",
    "\n",
    "valid_dataset = TimeSeriesDataset(Y_df=Y_df,\n",
    "                                  ds_in_test=0,\n",
    "                                  is_test=True,\n",
    "                                  input_size=group.horizon,\n",
    "                                  output_size=group.horizon,\n",
    "                                  verbose=True)"
   ]
  },
  {
   "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=1,\n",
    "                                shuffle=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "S.shape torch.Size([32, 0])\n",
      "Y.shape torch.Size([32, 47])\n",
      "X.shape torch.Size([32, 0, 47])\n",
      "idxs.shape torch.Size([32])\n"
     ]
    }
   ],
   "source": [
    "dataloader = iter(train_loader)\n",
    "batch = next(dataloader)\n",
    "S, Y, X = batch['S'], batch['Y'], batch['X']\n",
    "available_mask = batch['available_mask']\n",
    "idxs = batch['idxs']\n",
    "\n",
    "print(\"S.shape\", S.shape)\n",
    "print(\"Y.shape\", Y.shape)\n",
    "print(\"X.shape\", X.shape)\n",
    "print(\"idxs.shape\", idxs.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = ESRNN(n_series=train_dataset.n_series,\n",
    "              n_s=train_dataset.n_s,\n",
    "              n_x=train_dataset.n_x,\n",
    "              input_size=train_dataset.input_size,\n",
    "              output_size=train_dataset.output_size,\n",
    "              sample_freq=1,\n",
    "              learning_rate=5e-3,\n",
    "              lr_scheduler_step_size=100,\n",
    "              lr_decay=0.9,\n",
    "              per_series_lr_multip=1.5,\n",
    "              gradient_eps=1e-8,\n",
    "              gradient_clipping_threshold=10,\n",
    "              rnn_weight_decay=0,\n",
    "              noise_std=0.001,\n",
    "              level_variability_penalty=10,\n",
    "              testing_percentile=50,\n",
    "              training_percentile=51,\n",
    "              es_component='multiplicative',\n",
    "              cell_type='GRU',\n",
    "              state_hsize=50,\n",
    "              dilations=[[24, 48], [168]],\n",
    "              add_nl_layer=False,\n",
    "              loss='SMYL',\n",
    "              val_loss='MAE',\n",
    "              seasonality=[1])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "early_stopping = EarlyStopping(monitor=\"val_loss\", \n",
    "                               min_delta=1e-4, \n",
    "                               patience=1, verbose=True, \n",
    "                               mode=\"min\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "GPU available: False, used: False\n",
      "TPU available: False, using: 0 TPU cores\n",
      "\n",
      "  | Name  | Type   | Params\n",
      "---------------------------------\n",
      "0 | model | _ESRNN | 40.8 K\n",
      "---------------------------------\n",
      "40.8 K    Trainable params\n",
      "0         Non-trainable params\n",
      "40.8 K    Total params\n",
      "0.163     Total estimated model params size (MB)\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "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/federicogarza2/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/federicogarza2/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": "a100181e955f408d8ad8428d73a99ae3",
       "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": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validating: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Metric val_loss improved. New best score: 68293.117\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validating: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Monitored metric val_loss did not improve in the last 1 records. Best score: 68293.117. Signaling Trainer to stop.\n"
     ]
    }
   ],
   "source": [
    "trainer = pl.Trainer(max_epochs=2, progress_bar_refresh_rate=1, \n",
    "                     log_every_n_steps=1, check_val_every_n_epoch=1,\n",
    "                     callbacks=[early_stopping])\n",
    "trainer.fit(model, train_loader, valid_loader)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/federicogarza2/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": "062daaa6c36443599d4b56c9fd29238b",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Predicting: 17it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "y_true.shape torch.Size([1, 9082, 4])\n",
      "y_hat.shape torch.Size([1, 9082, 4])\n",
      "mask.shape torch.Size([1, 9082, 4])\n"
     ]
    }
   ],
   "source": [
    "outputs = trainer.predict(model, valid_loader)\n",
    "y, y_hat, mask = zip(*outputs)\n",
    "y = t.cat(y, axis=1)\n",
    "y_hat = t.cat(y_hat, axis=1)\n",
    "mask = t.cat(mask, axis=1)\n",
    "print(\"y_true.shape\", y.shape)\n",
    "print(\"y_hat.shape\", y_hat.shape)\n",
    "print(\"mask.shape\", mask.shape)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "nixtla",
   "language": "python",
   "name": "nixtla"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
