{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# default_exp losses.pytorch"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# PyTorch losses\n",
    "> Training losses."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#hide\n",
    "%load_ext autoreload\n",
    "%autoreload 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "import torch as t\n",
    "import torch.nn as nn"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "def divide_no_nan(a, b):\n",
    "    \"\"\"\n",
    "    Auxiliary funtion to handle divide by 0\n",
    "    \"\"\"\n",
    "    div = a / b\n",
    "    div[div != div] = 0.0\n",
    "    div[div == float('inf')] = 0.0\n",
    "    return div"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# PyTorch Train Losses "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "def MAPELoss(y, y_hat, mask=None):\n",
    "    \"\"\"MAPE Loss\n",
    "\n",
    "    Calculates Mean Absolute Percentage Error between\n",
    "    y and y_hat. MAPE measures the relative prediction\n",
    "    accuracy of a forecasting method by calculating the\n",
    "    percentual deviation of the prediction and the true\n",
    "    value at a given time and averages these devations\n",
    "    over the length of the series.\n",
    "    As defined in: https://en.wikipedia.org/wiki/Mean_absolute_percentage_error\n",
    "\n",
    "    Parameters\n",
    "    ----------\n",
    "    y: tensor (batch_size, output_size)\n",
    "        actual values in torch tensor.\n",
    "    y_hat: tensor (batch_size, output_size)\n",
    "        predicted values in torch tensor.\n",
    "    mask: tensor (batch_size, output_size)\n",
    "        specifies date stamps per serie\n",
    "        to consider in loss\n",
    "\n",
    "    Returns\n",
    "    -------\n",
    "    mape:\n",
    "    Mean absolute percentage error.\n",
    "    \"\"\"\n",
    "    if mask is None: mask = t.ones_like(y_hat)\n",
    "        \n",
    "    mask = divide_no_nan(mask, t.abs(y))\n",
    "    mape = t.abs(y - y_hat) * mask\n",
    "    mape = t.mean(mape)\n",
    "    return mape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# export\n",
    "def MSELoss(y, y_hat, mask=None):\n",
    "    \"\"\"MSE Loss\n",
    "\n",
    "    Calculates Mean Squared Error between\n",
    "    y and y_hat. MAPE measures the relative prediction\n",
    "    accuracy of a forecasting method by calculating the\n",
    "    percentual deviation of the prediction and the true\n",
    "    value at a given time and averages these devations\n",
    "    over the length of the series.\n",
    "\n",
    "    Parameters\n",
    "    ----------\n",
    "    y: tensor (batch_size, output_size)\n",
    "        actual values in torch tensor.\n",
    "    y_hat: tensor (batch_size, output_size)\n",
    "        predicted values in torch tensor.\n",
    "    mask: tensor (batch_size, output_size)\n",
    "        specifies date stamps per serie\n",
    "        to consider in loss\n",
    "\n",
    "    Returns\n",
    "    -------\n",
    "    mse:\n",
    "    Mean Squared Error.\n",
    "    \"\"\"\n",
    "    if mask is None: mask = t.ones_like(y_hat)\n",
    "        \n",
    "    mse = (y - y_hat)**2\n",
    "    mse = mask * mse\n",
    "    mse = t.mean(mse)\n",
    "    return mse"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# export\n",
    "def RMSELoss(y, y_hat, mask=None):\n",
    "    \"\"\"RMSE Loss\n",
    "\n",
    "    Calculates Mean Squared Error between\n",
    "    y and y_hat. MAPE measures the relative prediction\n",
    "    accuracy of a forecasting method by calculating the\n",
    "    percentual deviation of the prediction and the true\n",
    "    value at a given time and averages these devations\n",
    "    over the length of the series.\n",
    "\n",
    "    Parameters\n",
    "    ----------\n",
    "    y: tensor (batch_size, output_size)\n",
    "        actual values in torch tensor.\n",
    "    y_hat: tensor (batch_size, output_size)\n",
    "        predicted values in torch tensor.\n",
    "    mask: tensor (batch_size, output_size)\n",
    "        specifies date stamps per serie\n",
    "        to consider in loss\n",
    "\n",
    "    Returns\n",
    "    -------\n",
    "    rmse:\n",
    "    Root Mean Squared Error.\n",
    "    \"\"\"\n",
    "    if mask is None: mask = t.ones_like(y_hat)\n",
    "        \n",
    "    rmse = (y - y_hat)**2\n",
    "    rmse = mask * rmse\n",
    "    rmse = t.sqrt(t.mean(rmse))\n",
    "    return rmse"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# export\n",
    "def SMAPELoss(y, y_hat, mask=None):\n",
    "    \"\"\"SMAPE2 Loss\n",
    "\n",
    "    Calculates Symmetric Mean Absolute Percentage Error.\n",
    "    SMAPE measures the relative prediction accuracy of a\n",
    "    forecasting method by calculating the relative deviation\n",
    "    of the prediction and the true value scaled by the sum of the\n",
    "    absolute values for the prediction and true value at a\n",
    "    given time, then averages these devations over the length\n",
    "    of the series. This allows the SMAPE to have bounds between\n",
    "    0% and 200% which is desireble compared to normal MAPE that\n",
    "    may be undetermined.\n",
    "\n",
    "    Parameters\n",
    "    ----------\n",
    "    y: tensor (batch_size, output_size)\n",
    "        actual values in torch tensor.\n",
    "    y_hat: tensor (batch_size, output_size)\n",
    "        predicted values in torch tensor.\n",
    "\n",
    "    Returns\n",
    "    -------\n",
    "    smape:\n",
    "        symmetric mean absolute percentage error\n",
    "\n",
    "    References\n",
    "    ----------\n",
    "    [1] https://robjhyndman.com/hyndsight/smape/ (Makridakis 1993)\n",
    "    \"\"\"\n",
    "    if mask is None: mask = t.ones_like(y_hat)\n",
    "        \n",
    "    delta_y = t.abs((y - y_hat))\n",
    "    scale = t.abs(y) + t.abs(y_hat)\n",
    "    smape = divide_no_nan(delta_y, scale)\n",
    "    smape = smape * mask\n",
    "    smape = 2 * t.mean(smape)\n",
    "    return smape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# export\n",
    "def MASELoss(y, y_hat, y_insample, seasonality, mask=None) :\n",
    "    \"\"\" Calculates the M4 Mean Absolute Scaled Error.\n",
    "\n",
    "    MASE measures the relative prediction accuracy of a\n",
    "    forecasting method by comparinng the mean absolute errors\n",
    "    of the prediction and the true value against the mean\n",
    "    absolute errors of the seasonal naive model.\n",
    "\n",
    "    Parameters\n",
    "    ----------\n",
    "    seasonality: int\n",
    "        main frequency of the time series\n",
    "        Hourly 24,  Daily 7, Weekly 52,\n",
    "        Monthly 12, Quarterly 4, Yearly 1\n",
    "    y: tensor (batch_size, output_size)\n",
    "        actual test values\n",
    "    y_hat: tensor (batch_size, output_size)\n",
    "        predicted values\n",
    "    y_train: tensor (batch_size, input_size)\n",
    "        actual insample values for Seasonal Naive predictions\n",
    "\n",
    "    Returns\n",
    "    -------\n",
    "    mase:\n",
    "        mean absolute scaled error\n",
    "\n",
    "    References\n",
    "    ----------\n",
    "    [1] https://robjhyndman.com/papers/mase.pdf\n",
    "    \"\"\"\n",
    "    if mask is None: mask = t.ones_like(y_hat)\n",
    "        \n",
    "    delta_y = t.abs(y - y_hat)\n",
    "    scale = t.mean(t.abs(y_insample[:, seasonality:] - \\\n",
    "                            y_insample[:, :-seasonality]), axis=1)\n",
    "    mase = divide_no_nan(delta_y, scale[:, None])\n",
    "    mase = mase * mask\n",
    "    mase = t.mean(mase)\n",
    "    return mase"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# export\n",
    "def MAELoss(y, y_hat, mask=None):\n",
    "    \"\"\"MAE Loss\n",
    "\n",
    "    Calculates Mean Absolute Error between\n",
    "    y and y_hat. MAE measures the relative prediction\n",
    "    accuracy of a forecasting method by calculating the\n",
    "    deviation of the prediction and the true\n",
    "    value at a given time and averages these devations\n",
    "    over the length of the series.\n",
    "\n",
    "    Parameters\n",
    "    ----------\n",
    "    y: tensor (batch_size, output_size)\n",
    "        actual values in torch tensor.\n",
    "    y_hat: tensor (batch_size, output_size)\n",
    "        predicted values in torch tensor.\n",
    "    mask: tensor (batch_size, output_size)\n",
    "        specifies date stamps per serie\n",
    "        to consider in loss\n",
    "\n",
    "    Returns\n",
    "    -------\n",
    "    mae:\n",
    "    Mean absolute error.\n",
    "    \"\"\"\n",
    "    if mask is None: mask = t.ones_like(y_hat)\n",
    "        \n",
    "    mae = t.abs(y - y_hat) * mask\n",
    "    mae = t.mean(mae)\n",
    "    return mae"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# export\n",
    "def PinballLoss(y, y_hat, mask=None, tau=0.5):\n",
    "    \"\"\"Pinball Loss\n",
    "    Computes the pinball loss between y and y_hat.\n",
    "\n",
    "    Parameters\n",
    "    ----------\n",
    "    y: tensor (batch_size, output_size)\n",
    "        actual values in torch tensor.\n",
    "    y_hat: tensor (batch_size, output_size)\n",
    "        predicted values in torch tensor.\n",
    "    tau: float, between 0 and 1\n",
    "        the slope of the pinball loss, in the context of\n",
    "        quantile regression, the value of tau determines the\n",
    "        conditional quantile level.\n",
    "\n",
    "    Returns\n",
    "    -------\n",
    "    pinball:\n",
    "        average accuracy for the predicted quantile\n",
    "    \"\"\"\n",
    "    if mask is None: mask = t.ones_like(y_hat)\n",
    "        \n",
    "    delta_y = t.sub(y, y_hat)\n",
    "    pinball = t.max(t.mul(tau, delta_y), t.mul((tau - 1), delta_y))\n",
    "    pinball = pinball * mask\n",
    "    pinball = t.mean(pinball)\n",
    "    return pinball"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# ES-RNN PyTorch loss"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "def LevelVariabilityLoss(levels, level_variability_penalty):\n",
    "    \"\"\" Level Variability Loss\n",
    "    Computes the variability penalty for the level.\n",
    "\n",
    "    Parameters\n",
    "    ----------\n",
    "    levels: tensor with shape (batch, n_time)\n",
    "        levels obtained from exponential smoothing component of ESRNN\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",
    "    \n",
    "    Returns\n",
    "    ----------\n",
    "    level_var_loss:\n",
    "        wiggliness loss for the level vector\n",
    "    \"\"\"\n",
    "    assert levels.shape[1] > 2\n",
    "    level_prev = t.log(levels[:, :-1])\n",
    "    level_next = t.log(levels[:, 1:])\n",
    "    log_diff_of_levels = t.sub(level_prev, level_next)\n",
    "\n",
    "    log_diff_prev = log_diff_of_levels[:, :-1]\n",
    "    log_diff_next = log_diff_of_levels[:, 1:]\n",
    "    diff = t.sub(log_diff_prev, log_diff_next)\n",
    "    level_var_loss = diff**2\n",
    "    level_var_loss = level_var_loss.mean() * level_variability_penalty\n",
    "    \n",
    "    return level_var_loss"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "def SmylLoss(y, y_hat, levels, mask, tau, level_variability_penalty=0.0):\n",
    "    \"\"\"Computes the Smyl Loss that combines level variability with\n",
    "    with Pinball loss.\n",
    "    windows_y: tensor of actual values,\n",
    "                            shape (n_windows, batch_size, window_size).\n",
    "    windows_y_hat: tensor of predicted values,\n",
    "                                    shape (n_windows, batch_size, window_size).\n",
    "    levels: levels obtained from exponential smoothing component of ESRNN.\n",
    "                    tensor with shape (batch, n_time).\n",
    "    return: smyl_loss.\n",
    "    \"\"\"\n",
    "    \n",
    "    if mask is None: mask = t.ones_like(y_hat)\n",
    "        \n",
    "    smyl_loss = PinballLoss(y, y_hat, mask, tau)\n",
    "    \n",
    "    if level_variability_penalty > 0:\n",
    "        log_diff_of_levels = LevelVariabilityLoss(levels, level_variability_penalty) \n",
    "        smyl_loss += log_diff_of_levels\n",
    "    \n",
    "    return smyl_loss"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Multi-quantile PyTorch loss\n",
    "\n",
    "MQLoss definition and testing"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "def MQLoss(y, y_hat, quantiles, mask=None): \n",
    "    \"\"\"MQLoss\n",
    "\n",
    "    Calculates Average Multi-quantile Loss function, for\n",
    "    a given set of quantiles, based on the absolute \n",
    "    difference between predicted and true values.\n",
    "\n",
    "    Parameters\n",
    "    ----------\n",
    "    y: tensor (batch_size, output_size) actual values in torch tensor.\n",
    "    y_hat: tensor (batch_size, output_size, n_quantiles) predicted values in torch tensor.\n",
    "    mask: tensor (batch_size, output_size, n_quantiles) specifies date stamps per serie\n",
    "          to consider in loss\n",
    "    quantiles: tensor(n_quantiles) quantiles to estimate from the distribution of y.\n",
    "\n",
    "    Returns\n",
    "    -------\n",
    "    lq: tensor(n_quantiles) average multi-quantile loss.\n",
    "    \"\"\"    \n",
    "    assert len(quantiles) > 1, f'your quantiles are of len: {len(quantiles)}'\n",
    "    \n",
    "    if mask is None: mask = t.ones_like(y_hat)\n",
    "        \n",
    "    n_q = len(quantiles)\n",
    "    \n",
    "    error = y_hat - y.unsqueeze(-1)\n",
    "    sq = t.maximum(-error, t.zeros_like(error))\n",
    "    s1_q = t.maximum(error, t.zeros_like(error))\n",
    "    loss = (quantiles * sq + (1 - quantiles) * s1_q)\n",
    "   \n",
    "    return t.mean(t.mean(loss, axis=1))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "def wMQLoss(y, y_hat, quantiles, mask=None): \n",
    "    \"\"\"wMQLoss\n",
    "\n",
    "    Calculates Average Multi-quantile Loss function, for\n",
    "    a given set of quantiles, based on the absolute \n",
    "    difference between predicted and true values.\n",
    "\n",
    "    Parameters\n",
    "    ----------\n",
    "    y: tensor (batch_size, output_size) actual values in torch tensor.\n",
    "    y_hat: tensor (batch_size, output_size, n_quantiles) predicted values in torch tensor.\n",
    "    mask: tensor (batch_size, output_size, n_quantiles) specifies date stamps per serie\n",
    "          to consider in loss\n",
    "    quantiles: tensor(n_quantiles) quantiles to estimate from the distribution of y.\n",
    "\n",
    "    Returns\n",
    "    -------\n",
    "    lq: tensor(n_quantiles) average multi-quantile loss.\n",
    "    \"\"\"    \n",
    "    assert len(quantiles) > 1, f'your quantiles are of len: {len(quantiles)}'\n",
    "    \n",
    "    if mask is None: mask = t.ones_like(y_hat)\n",
    "    \n",
    "    n_q = len(quantiles)\n",
    "    \n",
    "    error = y_hat - y.unsqueeze(-1)\n",
    "    \n",
    "    sq = t.maximum(-error, t.zeros_like(error))\n",
    "    s1_q = t.maximum(error, t.zeros_like(error))\n",
    "    loss = (quantiles * sq + (1 - quantiles) * s1_q)\n",
    "    \n",
    "    loss = divide_no_nan(t.sum(loss * mask, axis=-2), \n",
    "                         t.sum(t.abs(y.unsqueeze(-1)) * mask, axis=-2))\n",
    "    \n",
    "    return t.mean(loss)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Checks for PyTorch train losses"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from torch import nn, optim\n",
    "from torch.utils.data import DataLoader, Dataset\n",
    "\n",
    "import numpy as np\n",
    "import time\n",
    "from scipy.stats import hmean\n",
    "import matplotlib.pyplot as plt\n",
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Model(nn.Module):  \n",
    "\n",
    "    def __init__(self, horizon, n_quantiles):\n",
    "        super(Model, self).__init__()\n",
    "        self.horizon = horizon\n",
    "        self.n_quantiles = n_quantiles\n",
    "        self.linear_layer = nn.Linear(in_features=n_obs, \n",
    "                                      out_features=horizon * n_quantiles, \n",
    "                                      bias=False)\n",
    "\n",
    "    def forward(self, x):\n",
    "        y_hat = self.linear_layer(x)\n",
    "        y_hat = y_hat.view(-1, self.horizon, self.n_quantiles)\n",
    "        return y_hat\n",
    "    \n",
    "class Data(Dataset):\n",
    "    \n",
    "    # Constructor\n",
    "    def __init__(self, Y, X):\n",
    "        self.X = X\n",
    "        self.Y = Y\n",
    "        self.len = Y.shape[0]\n",
    "\n",
    "    # Getter\n",
    "    def __getitem__(self, index):\n",
    "        return self.X[index], self.Y[index]\n",
    "    \n",
    "    # Get Length\n",
    "    def __len__(self):\n",
    "        return self.len"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "quantiles:\n",
      "tensor([0.0500, 0.3500, 0.6500, 0.9500])\n",
      "Y.shape: torch.Size([1000, 10]), X.shape: torch.Size([1000, 10])\n",
      "Y_test.shape: torch.Size([1000, 10]), X_test.shape: torch.Size([1000, 10])\n"
     ]
    }
   ],
   "source": [
    "# Hyperparameters and sample data parameters\n",
    "t.cuda.manual_seed(7)\n",
    "\n",
    "# Sample data\n",
    "n_ts = 1000\n",
    "n_obs = horizon = 10\n",
    "mean = 0.0 # to generate random numbers from N(mean, std)\n",
    "std = 7.0 # to generate random numbers from N(mean, std)\n",
    "start = 0.05 # First quantile\n",
    "end = 0.95 # Last quantiles\n",
    "steps = 4 # Number of quantiles\n",
    "\n",
    "# Hyperparameters\n",
    "batch_size = 500\n",
    "lr = 0.08\n",
    "epochs = 100\n",
    "\n",
    "# Sample data\n",
    "quantiles = t.Tensor([0.0500, 0.3500, 0.6500, 0.9500])\n",
    "print(f'quantiles:\\n{quantiles}')\n",
    "Y = t.normal(mean=mean, std=std, size=(n_ts, n_obs))\n",
    "X = t.ones(size=(n_ts, n_obs))\n",
    "\n",
    "Y_test = t.normal(mean=mean, std=std, size=(n_ts, horizon))\n",
    "X_test = t.ones(size=(n_ts, horizon))\n",
    "print(f'Y.shape: {Y.shape}, X.shape: {X.shape}')\n",
    "print(f'Y_test.shape: {Y_test.shape}, X_test.shape: {X_test.shape}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Model training \n",
    "model = Model(horizon=horizon, n_quantiles=len(quantiles))\n",
    "dataset = Data(X=X, Y=Y)\n",
    "dataloader = DataLoader(dataset=dataset, batch_size=batch_size)\n",
    "optimizer = optim.Adam(model.parameters(), lr=lr)\n",
    "\n",
    "def train_model(model, epochs, print_progress=False):\n",
    "\n",
    "    start = time.time()\n",
    "    i = 0 \n",
    "    training_trajectory = {'epoch': [],\n",
    "                           'train_loss': []}\n",
    "    \n",
    "    for epoch in range(epochs):\n",
    "        for x, y in dataloader:\n",
    "            \n",
    "            i += 1\n",
    "            y_hat = model(x)\n",
    "            #training_loss = wMQLoss(y=y, y_hat=y_hat, quantiles=quantiles)\n",
    "            training_loss = MQLoss(y=y, y_hat=y_hat, quantiles=quantiles)\n",
    "            if i % (epoch + 1) == 0: \n",
    "                training_trajectory['epoch'].append(i)\n",
    "                training_trajectory['train_loss'].append(training_loss.detach().numpy())\n",
    "            optimizer.zero_grad()\n",
    "            training_loss.backward()\n",
    "            optimizer.step()\n",
    "\n",
    "            display_string = 'Step: {}, Time: {:03.3f}, Insample {}: {:.5f}'.format(i, \n",
    "                                                                                    time.time()-start, \n",
    "                                                                                    \"MQLoss\", \n",
    "                                                                                    training_loss.cpu().data.numpy())\n",
    "            if print_progress: print(display_string)\n",
    "\n",
    "    return model, training_trajectory"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "nixtla",
   "language": "python",
   "name": "nixtla"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
