{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# default_exp losses.utils"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#hide\n",
    "%load_ext autoreload\n",
    "%autoreload 2"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Losses utils \n",
    "> Generic loss class."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "from typing import Union, List, Optional\n",
    "\n",
    "import torch as t\n",
    "from fastcore.foundation import patch\n",
    "\n",
    "from nixtlats.losses.pytorch import (\n",
    "    MAPELoss, MASELoss, SMAPELoss, \n",
    "    MSELoss, MAELoss, SmylLoss, \n",
    "    PinballLoss, MQLoss, wMQLoss \n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "class LossFunction:\n",
    "    def __init__(self, loss_name: str, seasonality: Optional[int] = None,\n",
    "                 percentile: Optional[Union[List[int], int]] = None,\n",
    "                 level_variability_penalty: Optional[int] = None) -> 'LossFunction':\n",
    "        \"\"\"Instantiates a callable class of the `loss_name` loss.\n",
    "        \n",
    "        Parameters\n",
    "        ----------\n",
    "        loss_name: str\n",
    "            Name of the loss.\n",
    "        seasonality: int\n",
    "            main frequency of the time series\n",
    "            Hourly 24,  Daily 7, Weekly 52,\n",
    "            Monthly 12, Quarterly 4, Yearly.\n",
    "            Default `None`. \n",
    "            Mandatory for MASE loss.\n",
    "        percentile: Union[List[int], int]\n",
    "            Target percentile. \n",
    "            For SMYL and PINBALL losses an int \n",
    "            is expected.\n",
    "            For MQ and wMQ losses a list of ints\n",
    "            is expected.\n",
    "            Default `None`.\n",
    "        level_variability_penalty: int\n",
    "            Only used for SMYL loss.\n",
    "        \"\"\"\n",
    "        if loss_name in ['SMYL', 'PINBALL'] and not isinstance(percentile, int):\n",
    "            raise Exception(f'Percentile should be integer for {loss_name} loss.')\n",
    "        elif loss_name in ['MQ', 'wMQ'] and not isinstance(percentile, list):\n",
    "            raise Exception(f'Percentile should be list for {loss_name} loss')\n",
    "        elif loss_name == 'MASE' and seasonality is None:\n",
    "            raise Exception(f'Seasonality should be a list of integers for {loss_name} loss')\n",
    "            \n",
    "\n",
    "        self.loss_name = loss_name\n",
    "        self.seasonality = seasonality\n",
    "        self.percentile = percentile\n",
    "        self.level_variability_penalty = level_variability_penalty\n",
    "        \n",
    "        self.tau = self.percentile / 100 if isinstance(percentile, int) else None\n",
    "        self.quantiles = [tau / 100 for tau in percentile] if isinstance(percentile, list) else None"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "@patch\n",
    "def __call__(self: LossFunction, \n",
    "             y: t.Tensor,\n",
    "             y_hat: t.Tensor, \n",
    "             mask: Optional[t.Tensor] = None, \n",
    "             y_insample: Optional[t.Tensor] = None, \n",
    "             levels: Optional[t.Tensor] = None) -> t.Tensor:\n",
    "    \"\"\"Returns loss according to `loss_name`.\"\"\"\n",
    "    if self.loss_name == 'SMYL':\n",
    "        return SmylLoss(y=y, y_hat=y_hat, levels=levels, mask=mask,\n",
    "                        tau=self.tau,\n",
    "                        level_variability_penalty=self.level_variability_penalty)\n",
    "\n",
    "    elif self.loss_name == 'PINBALL':\n",
    "        return PinballLoss(y=y, y_hat=y_hat, mask=mask, \n",
    "                           tau=self.tau)\n",
    "\n",
    "    elif self.loss_name == 'MQ':\n",
    "        quantiles = t.Tensor(self.quantiles, device=y.device)\n",
    "        return MQLoss(y=y, y_hat=y_hat, quantiles=quantiles, mask=mask)\n",
    "\n",
    "    elif self.loss_name == 'wMQ':\n",
    "        quantiles = t.Tensor(self.quantiles, device=y.device)\n",
    "        return wMQLoss(y=y, y_hat=y_hat, quantiles=quantiles, mask=mask)\n",
    "\n",
    "    elif self.loss_name == 'MAPE':\n",
    "        return MAPELoss(y=y, y_hat=y_hat, mask=mask)\n",
    "\n",
    "    elif self.loss_name == 'MASE':\n",
    "        return MASELoss(y=y, y_hat=y_hat, y_insample=y_insample, \n",
    "                        seasonality=self.seasonality, mask=mask)\n",
    "\n",
    "    elif self.loss_name == 'SMAPE':\n",
    "        return SMAPELoss(y=y, y_hat=y_hat, mask=mask)\n",
    "\n",
    "    elif self.loss_name == 'MSE':\n",
    "        return MSELoss(y=y, y_hat=y_hat, mask=mask)\n",
    "\n",
    "    elif self.loss_name == 'MAE':\n",
    "        return MAELoss(y=y, y_hat=y_hat, mask=mask)\n",
    "\n",
    "    raise Exception(f'Unknown loss function: {loss_name}')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Usage"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Multi-Quantile losses"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor(0.5067)"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "loss = LossFunction('MQ', percentile=[30, 60])\n",
    "y = t.normal(0, 1, size=(10, 10))\n",
    "y_hat = t.normal(0, 1, size=(10, 10, 2))\n",
    "\n",
    "loss(y, y_hat)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Mean losses"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor(1.1782)"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "loss = LossFunction('MAE')\n",
    "y = t.normal(0, 1, size=(10, 10))\n",
    "y_hat = t.normal(0, 1, size=(10, 10))\n",
    "\n",
    "loss(y, y_hat)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "nixtla",
   "language": "python",
   "name": "nixtla"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
