N-BEATS-S.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import gc\n",
    "import pandas as pd\n",
    "import numpy as np\n",
    "import math\n",
    "import plotly.graph_objects as go\n",
    "import random\n",
    "import os\n",
    "import copy\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "from time import time\n",
    "from datetime import datetime\n",
    "\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "\n",
    "from pytorchtools import EarlyStopping\n",
    "\n",
    "#!pip install wandb\n",
    "import wandb"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Data loading & data preprocessing"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load monthly M4 data.\n",
    "# Transform the data into a lists of arrays. Each inner array represents a timeseries.\n",
    "# Remove all the NaN values from the datasets.\n",
    "\n",
    "# M4\n",
    "trainset = pd.read_csv('https://raw.githubusercontent.com/Mcompetitions/M4-methods/master/Dataset/Train/Monthly-train.csv')\n",
    "testset = pd.read_csv('https://raw.githubusercontent.com/Mcompetitions/M4-methods/master/Dataset/Test/Monthly-test.csv')\n",
    "trainset.set_index('V1', inplace = True)\n",
    "testset.set_index('V1', inplace = True)\n",
    "# Add the testset columns behind the trainset columns\n",
    "testset_merge = trainset.merge(testset, on = 'V1', how = 'inner')\n",
    "# Get the data in numpy representation\n",
    "trainset_np = trainset.values\n",
    "testset_np = testset_merge.values\n",
    "# Select all non NaN values from the trainset\n",
    "trainset_clean = [x[x == x] for x in trainset_np]\n",
    "# Train/validation/test --------------------------------- NBeats paper validation strategy\n",
    "testset_m4m = [x[x == x] for x in testset_np]\n",
    "valset_m4m = trainset_clean.copy()\n",
    "trainset_m4m = [x[:-18] for x in trainset_clean]\n",
    "\n",
    "del(trainset, testset, testset_merge, trainset_np, testset_np, trainset_clean)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Scale independent loss functions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# output = batch x fl\n",
    "# target = batch x fl\n",
    "# actuals_train = batch x bl\n",
    "\n",
    "def SMAPE(output, target, actuals_train = None):\n",
    "    \n",
    "    abs_errors = torch.abs(target - output)\n",
    "    abs_output = torch.abs(output)\n",
    "    abs_target = torch.abs(target)\n",
    "    loss = 200 * torch.mean(abs_errors / (abs_output.detach() + abs_target + 1e-5))\n",
    "    # possibly nan values in training networks if no offset in denominator is used\n",
    "    \n",
    "    return loss\n",
    "\n",
    "\n",
    "def MAPE(output, target, actuals_train = None):\n",
    "\n",
    "    abs_errors = torch.abs(target - output)\n",
    "    abs_target = torch.abs(target)\n",
    "    loss = 100 * torch.mean(abs_errors / (abs_target + 1e-5))\n",
    "    # possibly nan values in training networks if no offset in denominator is used\n",
    "\n",
    "    return loss\n",
    "\n",
    "\n",
    "def MASE(output, target, actuals_train):\n",
    "    \n",
    "    mask = torch.abs(actuals_train)>1e-6\n",
    "    mad = torch.sum(torch.abs(actuals_train[:, 1:] - actuals_train[:, :-1]), dim = -1) / (torch.sum(mask, dim = -1) - 1)\n",
    "    mad_reshaped = mad.unsqueeze(-1).repeat_interleave(target.shape[-1], dim = -1)\n",
    "    loss_items = torch.mean((torch.abs(target - output)) / (mad_reshaped + 1e-5), dim = -1)\n",
    "    loss_items_clamped = torch.clamp(loss_items, 0, 5)\n",
    "    loss = torch.mean(loss_items_clamped)\n",
    "\n",
    "    return loss\n",
    "\n",
    "\n",
    "def MASE_m(output, target, actuals_train):\n",
    "    \n",
    "    mask = torch.abs(actuals_train)>1e-6\n",
    "    mad = torch.sum(torch.abs(actuals_train[:, 12:] - actuals_train[:, :-12]), dim = -1) / (torch.sum(mask, dim = -1) - 12)\n",
    "    mad_reshaped = mad.unsqueeze(-1).repeat_interleave(target.shape[-1], dim = -1)\n",
    "    loss_items = torch.mean((torch.abs(target - output)) / (mad_reshaped + 1e-5), dim = -1)\n",
    "    loss_items_clamped = torch.clamp(loss_items, 0, 5)\n",
    "    loss = torch.mean(loss_items_clamped) \n",
    "\n",
    "    return loss\n",
    "\n",
    "\n",
    "def RMSSE(output, target, actuals_train):\n",
    "    \n",
    "    mask = torch.abs(actuals_train)>1e-6\n",
    "    msd = torch.sum((actuals_train[:, 1:] - actuals_train[:, :-1])**2, dim = -1) / (torch.sum(mask, dim = -1) - 1)\n",
    "    msd_reshaped = msd.unsqueeze(-1).repeat_interleave(target.shape[-1], dim = -1)\n",
    "    loss_items = torch.sqrt(torch.mean((target - output)**2 / (msd_reshaped + 1e-5), dim = -1))\n",
    "    loss_items_clamped = torch.clamp(loss_items, 0, 5)\n",
    "    loss = torch.mean(loss_items_clamped)\n",
    "    #loss = torch.sqrt(torch.mean((target - output)**2 / msd_reshaped))\n",
    "\n",
    "    return loss\n",
    "\n",
    "\n",
    "def RMSSE_m(output, target, actuals_train):\n",
    "    \n",
    "    mask = torch.abs(actuals_train)>1e-6\n",
    "    msd = torch.sum((actuals_train[:, 12:] - actuals_train[:, :-12])**2, dim = -1) / (torch.sum(mask, dim = -1) - 12)\n",
    "    msd_reshaped = msd.unsqueeze(-1).repeat_interleave(target.shape[-1], dim = -1)\n",
    "    loss_items = torch.sqrt(torch.mean((target - output)**2 / (msd_reshaped + 1e-5), dim = -1))\n",
    "    loss_items_clamped = torch.clamp(loss_items, 0, 5)\n",
    "    loss = torch.mean(loss_items_clamped)\n",
    "    \n",
    "    return loss"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# NBeats models"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Generic building block:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class GenericNBeatsBlock(nn.Module):\n",
    "    \n",
    "    def __init__(self,\n",
    "                 device,\n",
    "                 backcast_length,\n",
    "                 forecast_length,\n",
    "                 hidden_layer_units, thetas_dims, \n",
    "                 share_thetas,\n",
    "                 dropout = False, dropout_p = 0.0, \n",
    "                 neg_slope = 0.00):\n",
    "        \n",
    "        super().__init__()\n",
    "        self.device = device\n",
    "        self.backcast_length = backcast_length\n",
    "        self.forecast_length = forecast_length        \n",
    "        if isinstance(hidden_layer_units, int):\n",
    "            self.hidden_layer_units = [hidden_layer_units for FC_layer in range(4)]\n",
    "        else:\n",
    "            #assert(len(hidden_layer_units) == 4)\n",
    "            self.hidden_layer_units = hidden_layer_units\n",
    "        self.thetas_dims = thetas_dims\n",
    "        self.share_thetas = share_thetas\n",
    "        self.dropout = dropout\n",
    "        self.dropout_p = dropout_p\n",
    "        self.neg_slope = neg_slope\n",
    "        \n",
    "        # shared layers in block\n",
    "        self.fc1 = nn.Linear(self.backcast_length,\n",
    "                             self.hidden_layer_units[0])#, bias = False)\n",
    "        self.fc2 = nn.Linear(self.hidden_layer_units[0], self.hidden_layer_units[1])#, bias = False)\n",
    "        self.fc3 = nn.Linear(self.hidden_layer_units[1], self.hidden_layer_units[2])#, bias = False)\n",
    "        self.fc4 = nn.Linear(self.hidden_layer_units[2], self.hidden_layer_units[3])#, bias = False)\n",
    "        \n",
    "        # do not use F.dropout as you want dropout to only affect training (not evaluation mode)\n",
    "        # nn.Dropout handles this automatically\n",
    "        if self.dropout:\n",
    "            self.dropoutlayer = nn.Dropout(p = self.dropout_p)\n",
    "        \n",
    "        # task specific (backcast & forecast) layers in block\n",
    "        # do not include bias - see section 3.1 - Ruben does include bias for generic blocks\n",
    "        if self.share_thetas:\n",
    "            self.theta_b_fc = self.theta_f_fc = nn.Linear(self.hidden_layer_units[3], self.thetas_dims)#, bias = False)\n",
    "        else:\n",
    "            self.theta_b_fc = nn.Linear(self.hidden_layer_units[3], self.thetas_dims)#, bias = False)\n",
    "            self.theta_f_fc = nn.Linear(self.hidden_layer_units[3], self.thetas_dims)#, bias = False)\n",
    "        \n",
    "        # block output layers\n",
    "        self.backcast_out = nn.Linear(self.thetas_dims, self.backcast_length)#, bias = False) # include bias - see section 3.3\n",
    "        self.forecast_out = nn.Linear(self.thetas_dims, self.forecast_length)#, bias = False) # include bias - see section 3.3\n",
    "        \n",
    "        \n",
    "    def forward(self, x):\n",
    "        \n",
    "        if self.dropout:\n",
    "            h1 = F.leaky_relu(self.fc1(x.to(self.device)), negative_slope = self.neg_slope)\n",
    "            h1 = self.dropoutlayer(h1)\n",
    "            h2 = F.leaky_relu(self.fc2(h1), negative_slope = self.neg_slope)\n",
    "            h2 = self.dropoutlayer(h2)\n",
    "            h3 = F.leaky_relu(self.fc3(h2), negative_slope = self.neg_slope)\n",
    "            h3 = self.dropoutlayer(h3)\n",
    "            h4 = F.leaky_relu(self.fc4(h3), negative_slope = self.neg_slope)\n",
    "            theta_b = F.leaky_relu(self.theta_b_fc(h4), negative_slope = self.neg_slope)\n",
    "            #theta_b = self.theta_b_fc(h4)\n",
    "            theta_f = F.leaky_relu(self.theta_f_fc(h4), negative_slope = self.neg_slope)\n",
    "            #theta_f = self.theta_f_fc(h4)\n",
    "            backcast = self.backcast_out(theta_b)\n",
    "            forecast = self.forecast_out(theta_f)\n",
    "        else:\n",
    "            h1 = F.leaky_relu(self.fc1(x.to(self.device)), negative_slope = self.neg_slope)\n",
    "            h2 = F.leaky_relu(self.fc2(h1), negative_slope = self.neg_slope)\n",
    "            h3 = F.leaky_relu(self.fc3(h2), negative_slope = self.neg_slope)\n",
    "            h4 = F.leaky_relu(self.fc4(h3), negative_slope = self.neg_slope)\n",
    "            theta_b = F.leaky_relu(self.theta_b_fc(h4), negative_slope = self.neg_slope)\n",
    "            #theta_b = self.theta_b_fc(h4)\n",
    "            theta_f = F.leaky_relu(self.theta_f_fc(h4), negative_slope = self.neg_slope)\n",
    "            #theta_f = self.theta_f_fc(h4)\n",
    "            backcast = self.backcast_out(theta_b)\n",
    "            forecast = self.forecast_out(theta_f)\n",
    "            \n",
    "        return backcast, forecast\n",
    "    \n",
    "    \n",
    "    def __str__(self):\n",
    "        \n",
    "        block_type = type(self).__name__\n",
    "        \n",
    "        return f'{block_type}(units={self.hidden_layer_units}, thetas_dims={self.thetas_dims}, ' \\\n",
    "            f'backcast_length={self.backcast_length}, ' \\\n",
    "            f'forecast_length={self.forecast_length}, share_thetas={self.share_thetas}, ' \\\n",
    "            f'dropout={self.dropout}, dropout_p={self.dropout_p}, neg_slope={self.neg_slope}) at @{id(self)}'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "StableNBeatsNet:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Only the forward method is changed compared to standard NBeatsNet\n",
    "class StableNBeatsNet(nn.Module): \n",
    "    \n",
    "    def __init__(self, \n",
    "                 device,\n",
    "                 backcast_length_multiplier,\n",
    "                 forecast_length,\n",
    "                 hidden_layer_units, thetas_dims, \n",
    "                 share_thetas,\n",
    "                 nb_blocks_per_stack, n_stacks, share_weights_in_stack,\n",
    "                 dropout = False, dropout_p = 0.0, \n",
    "                 neg_slope = 0.00):\n",
    "        \n",
    "        super().__init__()\n",
    "        self.device = device\n",
    "        self.backcast_length = backcast_length_multiplier * forecast_length\n",
    "        self.forecast_length = forecast_length\n",
    "        self.hidden_layer_units = hidden_layer_units\n",
    "        self.thetas_dims = thetas_dims\n",
    "        self.share_thetas = share_thetas\n",
    "        self.nb_blocks_per_stack = nb_blocks_per_stack\n",
    "        self.n_stacks = n_stacks\n",
    "        self.share_weights_in_stack = share_weights_in_stack\n",
    "        self.dropout = dropout\n",
    "        self.dropout_p = dropout_p\n",
    "        self.neg_slope = neg_slope\n",
    "        \n",
    "        self.stacks = []\n",
    "        self.parameters = []\n",
    "        \n",
    "        print(f'| N-Beats')\n",
    "        for stack_id in range(self.n_stacks):\n",
    "            self.stacks.append(self.create_stack(stack_id))\n",
    "        self.parameters = nn.ParameterList(self.parameters)\n",
    "        \n",
    "        \n",
    "    def create_stack(self, stack_id):\n",
    "        \n",
    "        print(f'| --  Stack Generic (#{stack_id}) (share_weights_in_stack={self.share_weights_in_stack})')\n",
    "        blocks = []\n",
    "        for block_id in range(self.nb_blocks_per_stack):\n",
    "            if self.share_weights_in_stack and block_id != 0:\n",
    "                block = blocks[-1]  # pick up the last one when we share weights\n",
    "            else:\n",
    "                block = GenericNBeatsBlock(self.device,\n",
    "                                           self.backcast_length,\n",
    "                                           self.forecast_length,\n",
    "                                           self.hidden_layer_units, self.thetas_dims, \n",
    "                                           self.share_thetas,\n",
    "                                           self.dropout, self.dropout_p, \n",
    "                                           self.neg_slope)\n",
    "                self.parameters.extend(block.parameters())\n",
    "                print(f'     | -- {block}')\n",
    "                blocks.append(block)\n",
    "                \n",
    "        return blocks\n",
    "\n",
    "    \n",
    "    def forward(self, backcast_arr):\n",
    "        \n",
    "        # dim backcast_arr = batch_size x shifts x backcast_length\n",
    "        # shifts == 0 is standard input window, others are shifted lookback windows \n",
    "        # higher index = further back in time\n",
    "        # feed different input windows (per batch) through the SAME network (check via list of learnable parameters)\n",
    "        # see https://stackoverflow.com/questions/54444630/application-of-nn-linear-layer-in-pytorch-on-additional-dimentions\n",
    "        \n",
    "        forecast_arr = torch.zeros((backcast_arr.shape[0], # take batch size from backcast\n",
    "                                    backcast_arr.shape[1], # take n of shifts from backcast\n",
    "                                    self.forecast_length), dtype = torch.float).to(self.device)\n",
    "        backcast_arr = backcast_arr.to(self.device)\n",
    "        \n",
    "        # loop through stacks (and blocks)\n",
    "        for stack_id in range(len(self.stacks)):\n",
    "            for block_id in range(len(self.stacks[stack_id])):\n",
    "                b, f = self.stacks[stack_id][block_id](backcast_arr)\n",
    "                backcast_arr = backcast_arr - b\n",
    "                forecast_arr = forecast_arr + f  \n",
    "                \n",
    "        return backcast_arr, forecast_arr"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "\n",
    "# Training and evaluation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def seed_torch(seed = 5101992):\n",
    "    \n",
    "    random.seed(seed)\n",
    "    os.environ[\"PYTHONHASHSEED\"] = str(seed)\n",
    "    np.random.seed(seed)\n",
    "    torch.manual_seed(seed) "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class StableNBeatsLearner:\n",
    "    \n",
    "    def __init__(self,\n",
    "                 device,\n",
    "                 forecast_length,\n",
    "                 configNBeats):\n",
    "        \n",
    "        gc.collect()\n",
    "                \n",
    "        self.device = device \n",
    "        self.forecast_length = forecast_length\n",
    "        self.configNBeats = configNBeats\n",
    "        \n",
    "        if self.configNBeats[\"loss_function\"] == 1:\n",
    "            self.loss = RMSSE\n",
    "        elif self.configNBeats[\"loss_function\"] == 2:\n",
    "            self.loss = RMSSE_m\n",
    "        elif self.configNBeats[\"loss_function\"] == 3:\n",
    "            self.loss = SMAPE\n",
    "        elif self.configNBeats[\"loss_function\"] == 4:\n",
    "            self.loss = MAPE\n",
    "            \n",
    "        self.rndseed = self.configNBeats[\"rndseed\"]\n",
    "        \n",
    "        seed_torch(self.rndseed)\n",
    "        \n",
    "        print('--- Model ---')    \n",
    "        self.model = StableNBeatsNet(self.device,\n",
    "                                     self.configNBeats[\"backcast_length_multiplier\"],\n",
    "                                     self.forecast_length,\n",
    "                                     self.configNBeats[\"hidden_layer_units\"],\n",
    "                                     self.configNBeats[\"thetas_dims\"],\n",
    "                                     self.configNBeats[\"share_thetas\"],\n",
    "                                     self.configNBeats[\"nb_blocks_per_stack\"],\n",
    "                                     self.configNBeats[\"n_stacks\"],\n",
    "                                     self.configNBeats[\"share_weights_in_stack\"],                               \n",
    "                                     self.configNBeats[\"dropout\"],\n",
    "                                     self.configNBeats[\"dropout_p\"],\n",
    "                                     self.configNBeats[\"neg_slope\"])\n",
    "        \n",
    "        self.model = self.model.to(self.device)\n",
    "        \n",
    "        self.optim = torch.optim.Adam(self.model.parameters(), \n",
    "                                      lr = self.configNBeats[\"learning_rate\"],\n",
    "                                      weight_decay = self.configNBeats[\"weight_decay\"])\n",
    "        \n",
    "        self.init_state = copy.deepcopy(self.model.state_dict())\n",
    "        self.init_state_opt = copy.deepcopy(self.optim.state_dict())\n",
    "        \n",
    "        wandb.watch(self.model)\n",
    "    \n",
    "    \n",
    "    def ts_padding(self, ts_train_data, ts_eval_data):\n",
    "        \n",
    "        # Some time series in the dataset are not long enough to support the specified:\n",
    "        # forecast_length and backcast_length + backcast input shifts\n",
    "        # we use zero padding for the time series that are too short (neutral effect on loss calculations)\n",
    "        # + self.shifts comes from the number of extra observations needed to create the shifted inputs/targets\n",
    "        # + (self.forgins - 1) comes from rolling origin evaluation\n",
    "        \n",
    "        length_train = (self.configNBeats[\"backcast_length_multiplier\"] * self.forecast_length + \n",
    "                        self.forecast_length +\n",
    "                        self.shifts)\n",
    "        length_eval = (self.configNBeats[\"backcast_length_multiplier\"] * self.forecast_length +\n",
    "                       self.forecast_length +\n",
    "                       self.shifts +\n",
    "                       (self.forigins - 1))\n",
    "        \n",
    "        ts_train_pad = [x if x.size >= length_train else np.pad(x,\n",
    "                                                                (int(length_train - x.size), 0), \n",
    "                                                                'constant', \n",
    "                                                                constant_values = 0) for x in ts_train_data]\n",
    "        ts_eval_pad = [x if x.size >= length_eval else np.pad(x,\n",
    "                                                              (int(length_eval - x.size), 0), \n",
    "                                                              'constant', \n",
    "                                                              constant_values = 0) for x in ts_eval_data]\n",
    "        \n",
    "        return ts_train_pad, ts_eval_pad\n",
    "        \n",
    "        \n",
    "    def make_batch(self, batch_data, shuffle_origin = True):\n",
    "        \n",
    "        # If shuffle_origin = True --> batch for training --> random forecast origin based on LH \n",
    "        # If shuffle_origin = False --> batch for evaluation --> fixed forecast origin\n",
    "        \n",
    "        # Split the batch into input_list and target_list\n",
    "        # In x_arr and target_arr: batch x shift x backcats_length/forecast_length\n",
    "        x_arr = np.empty(shape = (len(batch_data), \n",
    "                                  self.shifts + 1,\n",
    "                                  self.configNBeats[\"backcast_length_multiplier\"] * self.forecast_length))\n",
    "        target_arr = np.empty(shape = (len(batch_data), \n",
    "                                       self.shifts + 1, \n",
    "                                       self.forecast_length))\n",
    "        \n",
    "        # For every time series in the batch:\n",
    "        # (1) slice the time series according to specific forecasting origin (depending on shuffle_origin)\n",
    "        # (2) make shifted inputs/targets --> max number of shifts = forecast_length - 1\n",
    "        # (3) fill x_arr and target_arr\n",
    "        for j in range(len(batch_data)):\n",
    "            i = batch_data[j]\n",
    "            \n",
    "            if shuffle_origin: \n",
    "                # suffle_origin --> only in training \n",
    "                \n",
    "                ### --> also pick random scale --> does not result in improved results\n",
    "                ### to remain as close as possible to nbeats paper: do not pick random scale\n",
    "                ### i = i + i * np.random.default_rng().uniform(-0.95, 0.95, 1)\n",
    "                \n",
    "                # pick origin\n",
    "                LH_max_offset = int(self.configNBeats[\"LH\"] * self.forecast_length)\n",
    "                ts_max_offset = int(len(i) -\n",
    "                                    (self.configNBeats[\"backcast_length_multiplier\"] * self.forecast_length + \n",
    "                                     self.forecast_length +\n",
    "                                     self.shifts))\n",
    "                max_offset = min(LH_max_offset, ts_max_offset)\n",
    "                if max_offset < 1:\n",
    "                    offset = np.zeros(1)\n",
    "                else:\n",
    "                    offset = np.random.randint(low = 0, high = max_offset)   \n",
    "            else:\n",
    "                offset = np.zeros(1)\n",
    "            \n",
    "            if offset == 0:\n",
    "                for shift in range(self.shifts + 1):\n",
    "                    if shift == 0:\n",
    "                        x_arr[j, shift, :] = i[-self.forecast_length-self.configNBeats[\"backcast_length_multiplier\"]*self.forecast_length:-self.forecast_length]\n",
    "                        target_arr[j, shift, :] = i[-self.forecast_length:]\n",
    "                    else:\n",
    "                        x_arr[j, shift, :] = i[-self.forecast_length-self.configNBeats[\"backcast_length_multiplier\"]*self.forecast_length-shift:-self.forecast_length-shift]\n",
    "                        target_arr[j, shift, :] = i[-self.forecast_length-shift:-shift]\n",
    "            else:\n",
    "                for shift in range(self.shifts + 1):\n",
    "                    x_arr[j, shift, :] = i[-self.forecast_length-self.configNBeats[\"backcast_length_multiplier\"]*self.forecast_length-offset-shift:-self.forecast_length-offset-shift]\n",
    "                    target_arr[j, shift, :] = i[-self.forecast_length-offset-shift:-offset-shift]\n",
    "                    \n",
    "        return x_arr, target_arr\n",
    "                                        \n",
    "                    \n",
    "    def create_example_plots(self, output, target, actuals_train, final_evaluation = False):\n",
    "        \n",
    "        plot_forecasts = torch.cat((actuals_train, output))\n",
    "        plot_actuals = torch.cat((actuals_train, target))\n",
    "        random_sample_forecasts = plot_forecasts.squeeze()\n",
    "        random_sample_actuals = plot_actuals.squeeze()\n",
    "        x_axis = torch.arange(1, random_sample_forecasts.shape[0]+1)\n",
    "        fig = go.Figure()\n",
    "        fig.add_trace(go.Scatter(x = x_axis.numpy(), y = random_sample_forecasts.numpy(),\n",
    "                                 mode = 'lines+markers', name = 'forecasts'))\n",
    "        fig.add_trace(go.Scatter(x = x_axis.numpy(), y = random_sample_actuals.numpy(),\n",
    "                                 mode = 'lines+markers', name = 'actuals'))\n",
    "        \n",
    "        # We only visualize examples for last epoch\n",
    "        if not final_evaluation:\n",
    "            wandb.log({\"example_plots_evaluation\": fig})\n",
    "        else:\n",
    "            wandb.log({\"example_plots_final_evaluation\": fig})\n",
    "            \n",
    "            \n",
    "    def evaluate(self, x_arr, target_arr,\n",
    "                 epoch = None,\n",
    "                 need_grad = True,\n",
    "                 early_stop = False):\n",
    "        \n",
    "        losses = dict()\n",
    "        \n",
    "        # Inputs must be converted to np.array of Tensors (float)\n",
    "        x_arr = torch.from_numpy(x_arr).float().to(self.device)\n",
    "        target_arr = torch.from_numpy(target_arr).float().to(self.device)\n",
    "        \n",
    "        if need_grad:\n",
    "            self.model.train()\n",
    "            self.model.to(self.device)\n",
    "            _, forecast_arr = self.model(x_arr)\n",
    "            \n",
    "            losses_forecast_shifts = 0.0\n",
    "            for shift in range(self.shifts + 1):\n",
    "                losses_forecast_shifts += self.loss(forecast_arr[:, shift, :], \n",
    "                                                    target_arr[:, shift, :], \n",
    "                                                    x_arr[:, shift, :])\n",
    "            losses[\"forecast_accuracy\"] = losses_forecast_shifts / (self.shifts + 1)\n",
    "                        \n",
    "            if self.shifts > 0:\n",
    "                # dimensions = batch_size x shifted forecasts for stability computations\n",
    "                forecast_base_arr = torch.zeros((forecast_arr.shape[0],\n",
    "                                                 sum(range(self.forecast_length - self.shifts, self.forecast_length))),\n",
    "                                                dtype = torch.float).to(self.device)\n",
    "                forecast_shift_arr = torch.zeros((forecast_arr.shape[0],\n",
    "                                                  sum(range(self.forecast_length - self.shifts, self.forecast_length))),\n",
    "                                                 dtype = torch.float).to(self.device)\n",
    "                col = 0\n",
    "                for shift in range(1, self.shifts + 1):\n",
    "                    for horizon_m1 in range(self.forecast_length - shift):\n",
    "                        forecast_base_arr[:, col] = forecast_arr[:, 0, horizon_m1]\n",
    "                        forecast_shift_arr[:, col] = forecast_arr[:, shift, horizon_m1 + shift]\n",
    "                        col = col + 1\n",
    "                losses[\"forecast_stability\"] = self.loss(forecast_shift_arr, forecast_base_arr, x_arr[:, 0, :])\n",
    "            else:\n",
    "                losses[\"forecast_stability\"] = torch.zeros(1)\n",
    "                \n",
    "        else:\n",
    "            with torch.no_grad():\n",
    "                self.model.eval()\n",
    "                self.model.to(self.device)\n",
    "                _, forecast_arr = self.model(x_arr)\n",
    "                \n",
    "                losses_forecast_shifts = 0.0\n",
    "                for shift in range(self.shifts + 1):\n",
    "                    losses_forecast_shifts += self.loss(forecast_arr[:, shift, :], target_arr[:, shift, :], x_arr[:, shift, :])\n",
    "                losses[\"forecast_accuracy\"] = losses_forecast_shifts / (self.shifts + 1)\n",
    "                \n",
    "                if self.shifts > 0:\n",
    "                    forecast_base_arr = torch.zeros((forecast_arr.shape[0],\n",
    "                                                     sum(range(self.forecast_length - self.shifts, self.forecast_length))),\n",
    "                                                    dtype = torch.float).to(self.device)\n",
    "                    forecast_shift_arr = torch.zeros((forecast_arr.shape[0],\n",
    "                                                      sum(range(self.forecast_length - self.shifts, self.forecast_length))),\n",
    "                                                     dtype = torch.float).to(self.device)\n",
    "                    col = 0\n",
    "                    for shift in range(1, self.shifts + 1):\n",
    "                        for horizon_m1 in range(self.forecast_length - shift):\n",
    "                            forecast_base_arr[:, col] = forecast_arr[:, 0, horizon_m1]\n",
    "                            forecast_shift_arr[:, col] = forecast_arr[:, shift, horizon_m1 + shift]\n",
    "                            col = col + 1\n",
    "                    losses[\"forecast_stability\"] = self.loss(forecast_base_arr, forecast_shift_arr, x_arr[:, 0, :])\n",
    "                else:\n",
    "                    losses[\"forecast_stability\"] = torch.zeros(1)\n",
    "                    \n",
    "                if not self.disable_plot:\n",
    "                    if early_stop:\n",
    "                        # Plot validation examples - of standard/unshifted input - for last epoch before break\n",
    "                        # This part of the evaluation function is only called after training has been forced to stop\n",
    "                        self.create_example_plots(forecast_arr[0, 0, :], target_arr[0, 0, :], x_arr[0, 0, :])\n",
    "                    else:\n",
    "                        # Plot validation examples - of standard/unshifted input - for last epoch\n",
    "                        # This part of the evaluation function is called after training has been completed\n",
    "                        if (epoch == self.configNBeats[\"epochs\"]):\n",
    "                            self.create_example_plots(forecast_arr[0, 0, :], target_arr[0, 0, :], x_arr[0, 0, :])\n",
    "        \n",
    "        return losses\n",
    "    \n",
    "    \n",
    "    # Training of net (training data can include validation data) + validation or testing\n",
    "    def train_net(self,\n",
    "                  ts_train_m4m,\n",
    "                  ts_eval_m4m,\n",
    "                  forigins,\n",
    "                  validation = True,\n",
    "                  validation_earlystop = False,\n",
    "                  disable_plot = True):\n",
    "        \n",
    "        self.forigins = forigins\n",
    "        self.shifts = self.configNBeats[\"shifts\"]\n",
    "        self.validation = validation\n",
    "        self.validation_earlystop = validation_earlystop\n",
    "        self.disable_plot = disable_plot\n",
    "        #assert self.shifts < self.forecast_length # max allowed number of shifts is forecast_length - 1\n",
    "        \n",
    "        # Data preprocessing depends on backcast_length_multiplier\n",
    "        ts_train_pad, ts_eval_pad = self.ts_padding(ts_train_m4m, ts_eval_m4m)\n",
    "        ts_train_pad = np.array(ts_train_pad, dtype = object)\n",
    "        ts_eval_pad = np.array(ts_eval_pad, dtype = object)\n",
    "        \n",
    "        print('--- Training ---')\n",
    "        \n",
    "        # Containers to save train/evaluation losses and parameters\n",
    "        tloss_combined, tloss_forecast_accuracy, tloss_forecast_stability = [], [], []\n",
    "        eloss_combined, eloss_forecast_accuracy, eloss_forecast_stability = [], [], []\n",
    "        #params = []\n",
    "        \n",
    "        # Main training loop\n",
    "        self.model.load_state_dict(self.init_state)\n",
    "        self.optim.load_state_dict(self.init_state_opt)\n",
    "            \n",
    "        seed_torch(self.rndseed)\n",
    "        # Initialize early stopping object\n",
    "        if self.validation_earlystop:\n",
    "            early_stopping = EarlyStopping(patience = self.configNBeats[\"patience\"], verbose = True)\n",
    "            \n",
    "        for epoch in range(1, self.configNBeats[\"epochs\"]+1):\n",
    "            \n",
    "            start_time = time()\n",
    "            # Shuffle train data\n",
    "            np.random.shuffle(ts_train_pad)\n",
    "            # Determine number of batches per epoch\n",
    "            num_batches = int(ts_train_pad.shape[0] / self.configNBeats[\"batch_size\"])\n",
    "            \n",
    "            # Training per epoch\n",
    "            avg_tloss_combined_epoch = 0.0\n",
    "            avg_tloss_forecast_accuracy_epoch = 0.0\n",
    "            avg_tloss_forecast_stability_epoch = 0.0\n",
    "            \n",
    "            for k in range(num_batches):\n",
    "                \n",
    "                batch = np.array(ts_train_pad[k*self.configNBeats[\"batch_size\"]:(k+1)*self.configNBeats[\"batch_size\"]])\n",
    "                x_arr, target_arr = self.make_batch(batch, shuffle_origin = True)\n",
    "                \n",
    "                self.optim.zero_grad()\n",
    "                losses_batch = self.evaluate(x_arr, target_arr,\n",
    "                                             epoch, need_grad = True, \n",
    "                                             early_stop = False)\n",
    "                if self.shifts > 0:\n",
    "                    loss_combined = ((self.configNBeats[\"lambda\"] * losses_batch[\"forecast_stability\"]) +\n",
    "                                     ((1 - self.configNBeats[\"lambda\"]) * losses_batch[\"forecast_accuracy\"]))\n",
    "                else:\n",
    "                    loss_combined = losses_batch[\"forecast_accuracy\"]\n",
    "                loss_combined.backward()\n",
    "                self.optim.step()\n",
    "                \n",
    "                #params = self.model.parameters()\n",
    "                #total_params = sum(p.numel() for p in self.model.parameters() if p.requires_grad)\n",
    "                #if (epoch == 1 or epoch == self.configNBeats[\"epochs\"]) and k == 0:\n",
    "                #    print('Epoch {}/{} \\t n_learnable_pars={:.4f}'.format(\n",
    "                #        epoch,\n",
    "                #        self.configNBeats[\"epochs\"],\n",
    "                #        total_params))\n",
    "                \n",
    "                wandb.log({\"tloss_comb_step\": loss_combined,\n",
    "                           \"tloss_fcacc_step\": losses_batch[\"forecast_accuracy\"],\n",
    "                           \"tloss_fcstab_step\": losses_batch[\"forecast_stability\"]})\n",
    "                \n",
    "                avg_tloss_combined_epoch += (loss_combined / num_batches)\n",
    "                avg_tloss_forecast_accuracy_epoch += (losses_batch[\"forecast_accuracy\"] / num_batches)\n",
    "                avg_tloss_forecast_stability_epoch += (losses_batch[\"forecast_stability\"] / num_batches)\n",
    "                \n",
    "            if self.validation: # validation_full and validation_earlystop\n",
    "                \n",
    "                # Evaluation per epoch\n",
    "                avg_eloss_combined_epoch = 0.0\n",
    "                avg_eloss_forecast_accuracy_epoch = 0.0\n",
    "                avg_eloss_forecast_stability_epoch = 0.0\n",
    "\n",
    "                for forigin in range(self.forigins):\n",
    "\n",
    "                    # Only one batch, but one batch per forecast origin\n",
    "                    if forigin < self.forigins-1:\n",
    "                        eval_data_subset = np.array([x[:(-18 + forigin + self.forecast_length)] for x in ts_eval_pad],\n",
    "                                                   dtype = object)\n",
    "                    else:\n",
    "                        eval_data_subset = np.array([x for x in ts_eval_pad], dtype = object)\n",
    "                    x_arr, target_arr = self.make_batch(eval_data_subset, shuffle_origin = False)\n",
    "                    \n",
    "                    losses_evaluation = self.evaluate(x_arr, target_arr,\n",
    "                                                      epoch, need_grad = False,\n",
    "                                                      early_stop = False)\n",
    "                    if self.shifts > 0:\n",
    "                        loss_combined = ((self.configNBeats[\"lambda\"] * losses_evaluation[\"forecast_stability\"]) +\n",
    "                                         ((1 - self.configNBeats[\"lambda\"]) * losses_evaluation[\"forecast_accuracy\"]))\n",
    "                    else:\n",
    "                        loss_combined = losses_evaluation[\"forecast_accuracy\"]\n",
    "\n",
    "                    avg_eloss_combined_epoch += (loss_combined / self.forigins)\n",
    "                    avg_eloss_forecast_accuracy_epoch += (losses_evaluation[\"forecast_accuracy\"] / self.forigins)\n",
    "                    avg_eloss_forecast_stability_epoch += (losses_evaluation[\"forecast_stability\"] / self.forigins)\n",
    "                \n",
    "                elapsed_time = time() - start_time\n",
    "\n",
    "                print('Epoch {}/{} \\t tloss_combined={:.4f} \\t eloss_combined={:.4f} \\t time={:.2f}s'.format(\n",
    "                    epoch,\n",
    "                    self.configNBeats[\"epochs\"],\n",
    "                    avg_tloss_combined_epoch,\n",
    "                    avg_eloss_combined_epoch,\n",
    "                    elapsed_time))\n",
    "                \n",
    "                wandb.log({\"epoch\": epoch,\n",
    "                           \"tloss_comb_evol\": avg_tloss_combined_epoch,\n",
    "                           \"tloss_fcacc_evol\": avg_tloss_forecast_accuracy_epoch,\n",
    "                           \"tloss_fcstab_evol\": avg_tloss_forecast_stability_epoch,\n",
    "                           \"eloss_comb_evol\": avg_eloss_combined_epoch,\n",
    "                           \"eloss_fcacc_evol\": avg_eloss_forecast_accuracy_epoch,\n",
    "                           \"eloss_fcstab_evol\": avg_eloss_forecast_stability_epoch})\n",
    "                \n",
    "                if self.validation_earlystop: # validation_earlystop \n",
    "\n",
    "                    # early_stopping needs the average epoch validation loss to check if it has decreased, \n",
    "                    # and if it has, it will make a checkpoint of the current model\n",
    "                    early_stopping(avg_eloss_combined_epoch, self.model)\n",
    "\n",
    "                    if early_stopping.early_stop:\n",
    "                        print(\"Early stopping\")\n",
    "                        # Load the last checkpoint with the best model\n",
    "                        self.model.load_state_dict(torch.load('checkpoint.pt'))\n",
    "                        # Produce plots for final epoch before break\n",
    "                        if not self.disable_plot:\n",
    "                            for forigin in range(self.forigins):\n",
    "                                # Only one batch, but one batch per forecast origin\n",
    "                                if forigin < self.forigins-1:\n",
    "                                    eval_data_subset = np.array([x[:(-18 + forigin + self.forecast_length)] for x in ts_eval_pad],\n",
    "                                                                dtype = object)\n",
    "                                else:\n",
    "                                    eval_data_subset = np.array([x for x in ts_eval_pad], dtype = object)\n",
    "                                x_arr, target_arr = self.make_batch(eval_data_subset, shuffle_origin = False)\n",
    "\n",
    "                                losses_evaluation = self.evaluate(x_arr, target_arr,\n",
    "                                                                  epoch, need_grad = False,\n",
    "                                                                  early_stop = True)\n",
    "                        # Break loop over epochs\n",
    "                        break\n",
    "                    \n",
    "            else: # testing\n",
    "                \n",
    "                elapsed_time = time() - start_time\n",
    "\n",
    "                print('Epoch {}/{} \\t tloss_combined={:.4f} \\t time={:.2f}s'.format(\n",
    "                    epoch,\n",
    "                    self.configNBeats[\"epochs\"],\n",
    "                    avg_tloss_combined_epoch,\n",
    "                    elapsed_time))\n",
    "                \n",
    "                wandb.log({\"epoch\": epoch,\n",
    "                           \"tloss_comb_evol\": avg_tloss_combined_epoch,\n",
    "                           \"tloss_fcacc_evol\": avg_tloss_forecast_accuracy_epoch,\n",
    "                           \"tloss_fcstab_evol\": avg_tloss_forecast_stability_epoch})\n",
    "\n",
    "        wandb.log({\"tloss_comb\": avg_tloss_combined_epoch,\n",
    "                   \"tloss_fcacc\": avg_tloss_forecast_accuracy_epoch,\n",
    "                   \"tloss_fcstab\": avg_tloss_forecast_stability_epoch})\n",
    "        \n",
    "        print('--- Training done ---')\n",
    "        print('--- Final evaluation ---')\n",
    "        \n",
    "        print('--- M4 evaluation ---')\n",
    "        \n",
    "        # Containers to save actuals and forecasts\n",
    "        actuals = np.empty(shape = (len(ts_eval_pad), self.forigins, self.forecast_length)) # n_series, forigin, forecast_length\n",
    "        forecasts = np.empty(shape = (len(ts_eval_pad), self.forigins, self.forecast_length)) # n_series, forigin, forecast_length\n",
    "        \n",
    "        # Forecasts for each origin in rolling_window\n",
    "        for forigin in range(self.forigins):\n",
    "            \n",
    "            # Only one batch, but one batch per forecast origin\n",
    "            if forigin < self.forigins-1:\n",
    "                eval_data_subset = np.array([x[:(-18 + forigin + self.forecast_length)] for x in ts_eval_pad],\n",
    "                                           dtype = object)\n",
    "            else:\n",
    "                eval_data_subset = np.array([x for x in ts_eval_pad], dtype = object)\n",
    "            x_arr, target_arr = self.make_batch(eval_data_subset, shuffle_origin = False)\n",
    "            \n",
    "            # Produce forecasts for subset of test data\n",
    "            x_arr = torch.from_numpy(x_arr).float().to(self.device)\n",
    "            target_arr = torch.from_numpy(target_arr).float().to(self.device)\n",
    "            with torch.no_grad():\n",
    "                self.model.eval()\n",
    "                self.model.to(self.device)\n",
    "                _, forecast_arr = self.model(x_arr)\n",
    "                \n",
    "            x_arr = x_arr.cpu() \n",
    "            target_arr = target_arr.cpu()\n",
    "            forecast_arr = forecast_arr.cpu()\n",
    "                \n",
    "            # Plot 10 random examples per origin - of standard/unshifted input\n",
    "            sample_ids = np.random.randint(low = 0, high = int(x_arr.shape[0]), size = 10)\n",
    "            for sample_id in sample_ids:\n",
    "                self.create_example_plots(forecast_arr[sample_id, 0, :], \n",
    "                                          target_arr[sample_id, 0, :], \n",
    "                                          x_arr[sample_id, 0, :],\n",
    "                                          final_evaluation = True)\n",
    "                \n",
    "            # Save to containers\n",
    "            forecasts[:, forigin, :] = forecast_arr[:, 0, :]\n",
    "            actuals[:, forigin, :] = target_arr[:, 0, :]\n",
    "            \n",
    "        # Compute accuracy sMAPE\n",
    "        sMAPE = 200 * np.mean(np.abs(actuals - forecasts) / (np.abs(forecasts) + np.abs(actuals)))\n",
    "        \n",
    "        # Compute stability Total MAC\n",
    "        if self.forecast_length == 6:\n",
    "            weight = np.mean(actuals[:, [0, 6, 12], :], axis = (1, 2))\n",
    "        elif self.forecast_length == 18:\n",
    "            weight = np.mean(actuals[:, 0, :], axis = -1)\n",
    "        forecasts_helper = np.full((actuals.shape[0], \n",
    "                                    self.forigins,\n",
    "                                    (self.forecast_length - 1) + self.forigins), np.nan)\n",
    "        # n_series x self.forigins x ((forecast_length - 1) + forigins)\n",
    "        for forigin in range(self.forigins):\n",
    "            forecasts_helper[:, forigin, forigin:(forigin + self.forecast_length)] = forecasts[:, forigin, :]\n",
    "        MAC_mat = np.abs(np.diff(forecasts_helper, axis = 1))\n",
    "        # n_series x (self.forigins - 1) x ((forecast_length - 1) + forigins)\n",
    "        MAC_mat_adjust = np.delete(MAC_mat, [0, (self.forecast_length - 1) + self.forigins - 1], 2)\n",
    "        # n_series x (self.forigins - 1) x ((forecast_length - 1) + forigins - 2)\n",
    "        MAC = np.nanmean(MAC_mat_adjust, axis = 1)\n",
    "        # n_series x ((forecast_length - 1) + forigins - 2)\n",
    "        ItemMAC = np.mean(MAC, axis = 1) / weight\n",
    "        TotalMAC = np.mean(ItemMAC) * 100\n",
    "        \n",
    "        print('sMAPE_m4m={:.4f} \\t TotalMAC_m4m={:.4f}'.format(sMAPE, TotalMAC))\n",
    "        \n",
    "        wandb.log({\"sMAPE_m4m\": sMAPE,\n",
    "                   \"TotalMAC_m4m\": TotalMAC})\n",
    "        \n",
    "        # n_series, forigin, forecast_length\n",
    "        fc_colnames = [str(i) for i in range(1, self.forecast_length + 1)]\n",
    "        \n",
    "        actuals_np = actuals#.numpy()\n",
    "        m,n,r = actuals_np.shape\n",
    "        actuals_arr = np.column_stack((np.repeat(np.arange(m) + 1, n), \n",
    "                                       np.tile(np.arange(n) + 1, m),\n",
    "                                       actuals_np.reshape(m*n, -1)))\n",
    "        actuals_df = pd.DataFrame(actuals_arr, columns = ['item_id', 'fc_origin'] + fc_colnames)\n",
    "        helper_col = ['actual'] * len(actuals_df)\n",
    "        actuals_df['type'] = helper_col\n",
    "        \n",
    "        forecasts_np = forecasts#.numpy()\n",
    "        m,n,r = forecasts_np.shape\n",
    "        forecasts_arr = np.column_stack((np.repeat(np.arange(m) + 1, n), \n",
    "                                         np.tile(np.arange(n) + 1, m),\n",
    "                                         forecasts_np.reshape(m*n, -1)))\n",
    "        forecasts_df = pd.DataFrame(forecasts_arr, columns = ['item_id', 'fc_origin'] + fc_colnames)\n",
    "        helper_col = ['forecast'] * len(forecasts_df)\n",
    "        forecasts_df['type'] = helper_col\n",
    "        \n",
    "        output_df_m4m = pd.concat([actuals_df, forecasts_df])\n",
    "         \n",
    "        wandb.join()\n",
    "        \n",
    "        return output_df_m4m"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Run experiments"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Setup:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if torch.cuda.is_available():\n",
    "    device = torch.device(\"cuda:0\")\n",
    "else:\n",
    "    device = torch.device(\"cpu\")\n",
    "print(device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "wandb_project_name = 'nbeats_stability_m4_monthly'\n",
    "job_type_name = 'test'\n",
    "# one of:\n",
    "# - 'test', \n",
    "# - 'validation_full' --> e.g., for lambda value tuning\n",
    "# - 'validation_earlystop' --> for specifying number of epochs\n",
    "\n",
    "hyperparameter_defaults = dict()\n",
    "hyperparameter_defaults['epochs'] = 155 \n",
    "hyperparameter_defaults['batch_size'] = 512\n",
    "hyperparameter_defaults['nb_blocks_per_stack'] = 1\n",
    "hyperparameter_defaults['thetas_dims'] = 256\n",
    "hyperparameter_defaults['n_stacks'] = 20\n",
    "hyperparameter_defaults['share_weights_in_stack'] = False\n",
    "hyperparameter_defaults[\"backcast_length_multiplier\"] = 4\n",
    "hyperparameter_defaults['hidden_layer_units'] = 256\n",
    "hyperparameter_defaults['share_thetas'] = False\n",
    "hyperparameter_defaults[\"dropout\"] = False\n",
    "hyperparameter_defaults[\"dropout_p\"] = 0.0\n",
    "hyperparameter_defaults[\"neg_slope\"] = 0.00\n",
    "hyperparameter_defaults['learning_rate'] = 0.001\n",
    "hyperparameter_defaults[\"weight_decay\"] = 0.00\n",
    "hyperparameter_defaults[\"LH\"] = 10\n",
    "hyperparameter_defaults[\"rndseed\"] = 2000\n",
    "hyperparameter_defaults[\"loss_function\"] = 1 # 1 == RMSSE / 2 == RMSSE_m / 3 == SMAPE / 4 == MAPE\n",
    "hyperparameter_defaults[\"shifts\"] = 1\n",
    "hyperparameter_defaults['patience'] = 2000 # Only affects 'validation_earlystop' runs\n",
    "hyperparameter_defaults['lambda'] = 0.15 # Weight of forecast stability loss in loss_combined\n",
    "# Note that lambda is defined in the code as the proportion of stability in total loss (relative terms)\n",
    "# In the paper, lambda is defined in absolute terms\n",
    "# So: lambda_paper = lambda_code/(1-lambda_code)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if job_type_name == 'test':\n",
    "    is_val = False\n",
    "    do_earlystop = False\n",
    "    m4_train, m4_eval = valset_m4m, testset_m4m\n",
    "elif job_type_name == 'validation_full':\n",
    "    is_val = True\n",
    "    do_earlystop = False\n",
    "    m4_train, m4_eval = trainset_m4m, valset_m4m\n",
    "elif job_type_name == 'validation_earlystop':\n",
    "    is_val = True\n",
    "    do_earlystop = True\n",
    "    m4_train, m4_eval = trainset_m4m, valset_m4m"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#from google.colab import drive\n",
    "#drive.mount('/content/drive')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Run:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def sweep_function():\n",
    "    wandb.init(config = hyperparameter_defaults,\n",
    "               project = wandb_project_name,\n",
    "               job_type = job_type_name)\n",
    "    config = wandb.config\n",
    "    run_name = wandb.run.name\n",
    "    \n",
    "    # Initialize model\n",
    "    StableNBeats_model = StableNBeatsLearner(device, 6, config)\n",
    "    # Train & evaluate\n",
    "    forecasts_df_m4m = StableNBeats_model.train_net(m4_train, m4_eval, 13, is_val, do_earlystop)\n",
    "    # Save forecasts\n",
    "    forecasts_df_m4m.to_csv('m4m_nbeats_stability_' + job_type_name + '_' + run_name + '.csv', index = False)\n",
    "    #forecasts_df_m4m.to_csv('/content/drive/My Drive/Colab Notebooks/Sweeps/m4m_nbeats_stability_' + job_type_name + '_' + run_name + '.csv', index = False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sweep_config = {\n",
    "    \"name\": \"sweep\",\n",
    "    \"method\": \"grid\",\n",
    "    \"parameters\": {\n",
    "        \"rndseed\": {\n",
    "            \"values\": [2000]\n",
    "        },\n",
    "        \"weight_decay\": {\n",
    "            \"values\": [0.0]\n",
    "        },\n",
    "        \"dropout_p\": {\n",
    "            \"values\": [0.0]\n",
    "        },\n",
    "        \"lambda\": {\n",
    "            \"values\": [0.0]\n",
    "        }\n",
    "    }\n",
    "}\n",
    "sweep_id = wandb.sweep(sweep_config, project = wandb_project_name)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "wandb.agent(sweep_id, function = sweep_function)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}