Model - September 22.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "37c5b30c",
   "metadata": {},
   "outputs": [],
   "source": [
    "CONVERTED_DIRS = {'train':'./dataset/converted_train',\n",
    "                 'val':'./dataset/converted_validation',\n",
    "                 'test':'./dataset/converted_test'\n",
    "                 }\n",
    "\n",
    "RAW_DIRS = {'train':f'./dataset/gnnet-ch21-dataset-{\"train\"}',\n",
    "           'val':f'./dataset/gnnet-ch21-dataset-{\"validation\"}',\n",
    "           'test':f'./dataset/gnnet-ch21-dataset-{\"test\"}'\n",
    "           }\n",
    "\n",
    "#Debug: Optionally, you can change the directory to some location on an SSD\n",
    "#CONVERTED_DIRS['train'] = '/usr/converted_train/converted_train'"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "50c09974",
   "metadata": {},
   "source": [
    "## Convert dataset to pytorch data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "fd15c536",
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "import os\n",
    "import os.path as osp\n",
    "import torch\n",
    "import re\n",
    "from torch_geometric.data import DataLoader, Dataset\n",
    "\n",
    "import numpy as np\n",
    "import torch.nn.functional as F\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "class ChallengeDataset(Dataset):\n",
    "    \"\"\"\n",
    "    Base class representing a dataset for the challenge.\n",
    "    \n",
    "    We assume that the conversion process is already done, i.e. we \n",
    "    work with a list of pytorch Data objects stored in .pt files,\n",
    "    all in the same folder.\n",
    "    \"\"\"\n",
    "    \n",
    "\n",
    "    def challenge_transform(self,data,converted_path=None,debug=False):\n",
    "        all_timeparams = ['EqLambda', 'AvgPktsLambda', 'ExpMaxFactor',\n",
    "                     'MinPktLambda','MaxPktLambda','StdDev','PktsLambdaOn',\n",
    "                      'AvgTOff','AvgTOn','BurstGenLambda','Bitrate',\n",
    "                      'ParetoMinSize','ParetoMaxSize','ParetoAlfa'\n",
    "                     ]\n",
    "\n",
    "        all_sizeparams = ['MinSize','MaxSize','AvgPktSize','PktSize1',\n",
    "                         'PktSize2','NumCandidates','Size_i','Prob_i']\n",
    "\n",
    "\n",
    "        \"\"\" 1a. Assert that some stuff remains constant...\"\"\"\n",
    "        if debug:\n",
    "            assert all(data.p_SizeDist == 2)\n",
    "            assert all(data.p_TimeDist == 0)\n",
    "            assert all(data.n_levelsQoS == 1)\n",
    "            assert all(data.p_ToS == 0.0)\n",
    "            assert all(data.n_queueSizes == 32)\n",
    "            assert data.n_schedulingPolicy[0] == \"FIFO\"\n",
    "            for a,v in zip(['p_size_AvgPktSize','p_size_PktSize1',\n",
    "                            'p_size_PktSize2', 'p_time_ExpMaxFactor'],\n",
    "                        [1000.0,300.0,1700.0,10.0]):\n",
    "                if not torch.allclose(getattr(data,a), \n",
    "                                      v*torch.ones_like(getattr(data,a)),rtol=1e-05):\n",
    "                    raise Exception(f\"{a} was expected to have the value close to {v}\")\n",
    "\n",
    "        del data.p_SizeDist, data.p_TimeDist, data.p_ToS\n",
    "        del data.n_queueSizes, data.n_levelsQoS, data.n_schedulingPolicy\n",
    "\n",
    "        \"\"\" 1b. Transform p_SizeDist and p_TimeDist into one-hot. We skip it because it \n",
    "        does not change in the training dataset.\n",
    "        \"\"\"\n",
    "        #data.p_SizeDist= (F.one_hot(data.p_SizeDist,4))\n",
    "        #data.p_TimeDist= (F.one_hot(data.p_SizeDist,6))\n",
    "\n",
    "        \"\"\" \n",
    "        2. Path attributes; Concatenate Size/Time distribution parameters.\n",
    "\n",
    "         It turns out all sizeparams have the same value. Useless.\n",
    "            Otherwise, we'd have:\n",
    "                data.p_sizeparams = torch.cat([getattr(data,a).view(-1,1) for a in sizeparams],axis=1)\n",
    "\n",
    "         Also, p_time_ExpMaxFactor is always equal to 10.0, so we delete it\n",
    "        \"\"\"\n",
    "        delattr(data,'p_time_ExpMaxFactor')\n",
    "\n",
    "        sizeparams = [f'p_size_{a}' for a in ['AvgPktSize','PktSize1',\n",
    "                     'PktSize2']]\n",
    "        timeparams = [f'p_time_{a}' for a in ['EqLambda', 'AvgPktsLambda']]\n",
    "    \n",
    "        p_params = timeparams + ['p_TotalPktsGen','p_PktsGen','p_AvgBw']\n",
    "        \n",
    "        mean_pkts_rate = data.p_time_AvgPktsLambda.mean().item()\n",
    "        \n",
    "        assert mean_pkts_rate > 0\n",
    "        for p in p_params:\n",
    "            setattr(data,p,getattr(data,p)/mean_pkts_rate)\n",
    "            \n",
    "        data.p_time_EqLambda /= 1000.0\n",
    "        data.p_AvgBw /= 1000.0\n",
    "        #data.p_time_EqLambda *= 0.0\n",
    "        data.p_TotalPktsGen *= 0.0\n",
    "        \n",
    "        \"\"\"\n",
    "        Time parameters: Total Packets Generated (unused)\n",
    "                         EqLambda (we divide by 1000)\n",
    "                         Packets Generated\n",
    "                         Average Packets Generated\n",
    "                         Average Bandwidth (we divide by 1000)\n",
    "        \n",
    "        \"\"\"\n",
    "        \n",
    "        \n",
    "        data.P = torch.cat([getattr(data,a).view(-1,1) for a in p_params],axis=1)\n",
    "        \n",
    "        \n",
    "        for p in p_params + sizeparams:\n",
    "            delattr(data,p)\n",
    "        \n",
    "        \"\"\"\n",
    "            3. Global Attributes\n",
    "        \"\"\"\n",
    "        global_attrs = ['g_delay','g_packets','g_losses','g_AvgPktsLambda']\n",
    "        data.G = torch.as_tensor([getattr(data,a) for a in global_attrs],device=data.P.device)\n",
    "        data.G = torch.tile(data.G.view(1,-1),(data.type.shape[0],1))\n",
    "        for a in global_attrs:\n",
    "            delattr(data,a)\n",
    "\n",
    "        \"\"\"\n",
    "            4. Link attributes\n",
    "        \"\"\"\n",
    "        data.L = data.l_capacity.clone().view(-1,1) / mean_pkts_rate\n",
    "        data.mean_pkts_rate = mean_pkts_rate*torch.ones_like(data.type)\n",
    "        \n",
    "        data.n_paths = data.P.shape[0]*torch.ones_like(data.type)\n",
    "        data.n_links = data.L.shape[0]*torch.ones_like(data.type)\n",
    "        \n",
    "        \n",
    "        delattr(data,'l_capacity')\n",
    "        \n",
    "        \n",
    "        if not converted_path is None:\n",
    "            torch.save(data,converted_path)\n",
    "\n",
    "        return data\n",
    "    \n",
    "    def __init__(self, root_dir,filenames=None):\n",
    "        \"\"\"\n",
    "        Args:\n",
    "            csv_file (string): Path to the csv file with annotations.\n",
    "            root_dir (string): Directory with all the .pt files.\n",
    "            transform (callable, optional): Optional transform to be applied\n",
    "                on a sample.\n",
    "        \"\"\"\n",
    "        self.root_dir = root_dir\n",
    "        if filenames is None:\n",
    "            onlyfiles = [f for f in os.listdir(self.root_dir) if osp.isfile(osp.join(self.root_dir, f))]\n",
    "            self.filenames = [f for f in onlyfiles if f.endswith('.pt')]\n",
    "        else:\n",
    "            self.filenames = filenames\n",
    "        \n",
    "    def __len__(self):\n",
    "        return len(self.filenames)\n",
    "\n",
    "    def __getitem__(self, idx):\n",
    "        if torch.is_tensor(idx):\n",
    "            idx = idx.tolist()\n",
    "            \n",
    "        filename = self.filenames[idx]\n",
    "        pt_path = osp.join(self.root_dir, filename) \n",
    "        \n",
    "        converted_dir = self.root_dir+'_2'\n",
    "        #os.makedirs(converted_dir,exist_ok=True)\n",
    "        converted_path = osp.join(converted_dir, filename) \n",
    "        try:\n",
    "            sample = torch.load(pt_path,map_location='cuda')\n",
    "        except KeyboardInterrupt:\n",
    "            raise KeyboardInterrupt\n",
    "        except:\n",
    "            print(f\"Couldn't load {pt_path}\")\n",
    "        sample = self.challenge_transform(sample,converted_path=None)\n",
    "\n",
    "        return sample\n",
    "\n",
    "    "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f535a96b",
   "metadata": {},
   "source": [
    "\n",
    "## Divide validation datasets into 3. Initialize datasets/dataloaders\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "3eeec346",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "results_100_400-2000_0_9.tar.gz\n",
      "1    1040\n",
      "2    1040\n",
      "3    1040\n",
      "Name: validation_setting, dtype: int64\n",
      "results_100_400-2000_0_9.tar.gz\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>full_path</th>\n",
       "      <th>num_nodes</th>\n",
       "      <th>validation_setting</th>\n",
       "      <th>sample_num</th>\n",
       "      <th>file_num</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>test_42_0.pt</th>\n",
       "      <td>./dataset/gnnet-ch21-dataset-test/ch21-test-se...</td>\n",
       "      <td>50</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>42</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>test_42_1.pt</th>\n",
       "      <td>./dataset/gnnet-ch21-dataset-test/ch21-test-se...</td>\n",
       "      <td>50</td>\n",
       "      <td>1</td>\n",
       "      <td>1</td>\n",
       "      <td>42</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>test_42_2.pt</th>\n",
       "      <td>./dataset/gnnet-ch21-dataset-test/ch21-test-se...</td>\n",
       "      <td>50</td>\n",
       "      <td>1</td>\n",
       "      <td>2</td>\n",
       "      <td>42</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>test_42_3.pt</th>\n",
       "      <td>./dataset/gnnet-ch21-dataset-test/ch21-test-se...</td>\n",
       "      <td>50</td>\n",
       "      <td>1</td>\n",
       "      <td>3</td>\n",
       "      <td>42</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>test_42_4.pt</th>\n",
       "      <td>./dataset/gnnet-ch21-dataset-test/ch21-test-se...</td>\n",
       "      <td>50</td>\n",
       "      <td>1</td>\n",
       "      <td>4</td>\n",
       "      <td>42</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>...</th>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>test_61_5.pt</th>\n",
       "      <td>./dataset/gnnet-ch21-dataset-test/ch21-test-se...</td>\n",
       "      <td>300</td>\n",
       "      <td>3</td>\n",
       "      <td>5</td>\n",
       "      <td>61</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>test_61_6.pt</th>\n",
       "      <td>./dataset/gnnet-ch21-dataset-test/ch21-test-se...</td>\n",
       "      <td>300</td>\n",
       "      <td>3</td>\n",
       "      <td>6</td>\n",
       "      <td>61</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>test_61_7.pt</th>\n",
       "      <td>./dataset/gnnet-ch21-dataset-test/ch21-test-se...</td>\n",
       "      <td>300</td>\n",
       "      <td>3</td>\n",
       "      <td>7</td>\n",
       "      <td>61</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>test_61_8.pt</th>\n",
       "      <td>./dataset/gnnet-ch21-dataset-test/ch21-test-se...</td>\n",
       "      <td>300</td>\n",
       "      <td>3</td>\n",
       "      <td>8</td>\n",
       "      <td>61</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>test_61_9.pt</th>\n",
       "      <td>./dataset/gnnet-ch21-dataset-test/ch21-test-se...</td>\n",
       "      <td>300</td>\n",
       "      <td>3</td>\n",
       "      <td>9</td>\n",
       "      <td>61</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "<p>1560 rows × 5 columns</p>\n",
       "</div>"
      ],
      "text/plain": [
       "                                                      full_path  num_nodes  \\\n",
       "test_42_0.pt  ./dataset/gnnet-ch21-dataset-test/ch21-test-se...         50   \n",
       "test_42_1.pt  ./dataset/gnnet-ch21-dataset-test/ch21-test-se...         50   \n",
       "test_42_2.pt  ./dataset/gnnet-ch21-dataset-test/ch21-test-se...         50   \n",
       "test_42_3.pt  ./dataset/gnnet-ch21-dataset-test/ch21-test-se...         50   \n",
       "test_42_4.pt  ./dataset/gnnet-ch21-dataset-test/ch21-test-se...         50   \n",
       "...                                                         ...        ...   \n",
       "test_61_5.pt  ./dataset/gnnet-ch21-dataset-test/ch21-test-se...        300   \n",
       "test_61_6.pt  ./dataset/gnnet-ch21-dataset-test/ch21-test-se...        300   \n",
       "test_61_7.pt  ./dataset/gnnet-ch21-dataset-test/ch21-test-se...        300   \n",
       "test_61_8.pt  ./dataset/gnnet-ch21-dataset-test/ch21-test-se...        300   \n",
       "test_61_9.pt  ./dataset/gnnet-ch21-dataset-test/ch21-test-se...        300   \n",
       "\n",
       "              validation_setting  sample_num  file_num  \n",
       "test_42_0.pt                   1           0        42  \n",
       "test_42_1.pt                   1           1        42  \n",
       "test_42_2.pt                   1           2        42  \n",
       "test_42_3.pt                   1           3        42  \n",
       "test_42_4.pt                   1           4        42  \n",
       "...                          ...         ...       ...  \n",
       "test_61_5.pt                   3           5        61  \n",
       "test_61_6.pt                   3           6        61  \n",
       "test_61_7.pt                   3           7        61  \n",
       "test_61_8.pt                   3           8        61  \n",
       "test_61_9.pt                   3           9        61  \n",
       "\n",
       "[1560 rows x 5 columns]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "ds_val =  ChallengeDataset(root_dir=CONVERTED_DIRS['val'])\n",
    "ds_test =  ChallengeDataset(root_dir=CONVERTED_DIRS['test'])\n",
    "\n",
    "filenames_val = ds_val.filenames\n",
    "\n",
    "import re\n",
    "import pandas as pd\n",
    "import numpy as np\n",
    "import datanetAPI\n",
    "import os.path as osp\n",
    "def converted_filenames_metadata(filenames,path_to_original_dataset):\n",
    "        import re\n",
    "        \n",
    "        def m(f):\n",
    "            g = re.match(\"(validation|train|test)\\_(\\d+)\\_(\\d+).*\",f).groups()\n",
    "            g  = [g[0], int(g[1]), int(g[2])]\n",
    "            return g\n",
    "        \n",
    "        matches =  [m(f) for f in filenames]\n",
    "        reader = datanetAPI.DatanetAPI(path_to_original_dataset)\n",
    "        files_num = np.array([m[1] for m in matches],dtype=np.int32)\n",
    "        samples_num = np.array([m[2] for m in matches],dtype=np.int32)\n",
    "        \n",
    "        all_paths = np.array(reader.get_available_files())\n",
    "        print(all_paths[0,1])\n",
    "        df = pd.DataFrame(index=filenames,columns=['full_path','num_nodes','validation_setting'])\n",
    "        df['full_path'] = all_paths[files_num,0]\n",
    "        df['sample_num'] = samples_num\n",
    "        df['file_num'] = files_num\n",
    "        \n",
    "        df['num_nodes'] = np.array([osp.split(f)[-1] for f in df['full_path'].values],dtype=np.int32)\n",
    "        \n",
    "        if matches[0][0] in ['validation','test']:\n",
    "            df['validation_setting'] = np.array([osp.split(f)[-2][-1] for f in df['full_path'].values],dtype=np.int32)\n",
    "        else:\n",
    "            df['validation_setting'] = -1\n",
    "            \n",
    "        \"\"\"\n",
    "            Put it in correct order\n",
    "        \"\"\"\n",
    "        df = df.sort_values(by=['validation_setting','num_nodes','file_num','sample_num'])\n",
    "        return df\n",
    "        \n",
    "        \n",
    "df_val = converted_filenames_metadata(ds_val.filenames,RAW_DIRS['val'])\n",
    "\"\"\"\n",
    "    We opt to make the validation set smaller (as it is more time consuming to run)\n",
    "\"\"\"\n",
    "print(df_val['validation_setting'].value_counts())\n",
    "df_val['filenames'] = df_val.index.values\n",
    "df_val = df_val.groupby('full_path').head(10)\n",
    "df_test =  converted_filenames_metadata(ds_test.filenames,RAW_DIRS['test'])\n",
    "display(df_test)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "c30c4e73",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Dataset Len (train): 120000\n",
      "Dataloader Len (train): 7500\n",
      "Dataset Len (val): 3120\n",
      "Dataloader Len (val): 780\n",
      "Dataset Len (test): 1560\n",
      "Dataloader Len (test): 390\n",
      "Dataset Len (val_1): 260\n",
      "Dataloader Len (val_1): 65\n",
      "Dataset Len (val_2): 260\n",
      "Dataloader Len (val_2): 65\n",
      "Dataset Len (val_3): 260\n",
      "Dataloader Len (val_3): 65\n"
     ]
    }
   ],
   "source": [
    "BATCH_SIZE = {'train':16,'val':4}\n",
    "datasets = {\"train\": ChallengeDataset(root_dir=CONVERTED_DIRS['train']),\n",
    "            \"val\":ChallengeDataset(root_dir=CONVERTED_DIRS['val']),\n",
    "            \"test\":ChallengeDataset(root_dir=CONVERTED_DIRS['test'],\n",
    "                                   filenames=list(df_test.index) )\n",
    "           }\n",
    "for i in range(3):\n",
    "    which_files = list(df_val[df_val['validation_setting']==i+1]['filenames'].values)\n",
    "    ds = ChallengeDataset(root_dir=CONVERTED_DIRS['val'],\n",
    "                         filenames=which_files)\n",
    "    datasets[f'val_{i+1}'] = ds\n",
    "    \n",
    "    \n",
    "    \n",
    "    \n",
    "dataloaders = {}\n",
    "for k in datasets.keys():\n",
    "    if k.startswith('train'):\n",
    "        batch_size = BATCH_SIZE['train']\n",
    "    else:\n",
    "        batch_size = BATCH_SIZE['val']\n",
    "    dataloaders[k] = DataLoader(datasets[k],batch_size=batch_size,shuffle=False)\n",
    "    dataloaders[k+\"_s\"] = DataLoader(datasets[k],batch_size=batch_size,shuffle=True)\n",
    "    print(f\"Dataset Len ({k}): {len(datasets[k])}\")\n",
    "    print(f\"Dataloader Len ({k}): {len(dataloaders[k])}\")\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "146e0c01",
   "metadata": {},
   "source": [
    "<h2><b>(Debug)</b></h2> Optionally, run the following to ensure that the training, validation and test datasets are correctly loaded"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 250,
   "id": "9d058ca2",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 17%|█▋        | 259/1560 [00:02<00:14, 87.76it/s] \n"
     ]
    },
    {
     "ename": "KeyboardInterrupt",
     "evalue": "",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mKeyboardInterrupt\u001b[0m                         Traceback (most recent call last)",
      "\u001b[0;32m<ipython-input-1-15dfa748da3c>\u001b[0m in \u001b[0;36m__getitem__\u001b[0;34m(self, idx)\u001b[0m\n\u001b[1;32m    146\u001b[0m         \u001b[0;32mtry\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 147\u001b[0;31m             \u001b[0msample\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mtorch\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mload\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mpt_path\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0mmap_location\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;34m'cuda'\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    148\u001b[0m         \u001b[0;32mexcept\u001b[0m \u001b[0mKeyboardInterrupt\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/anaconda3/envs/gcn2/lib/python3.9/site-packages/torch/serialization.py\u001b[0m in \u001b[0;36mload\u001b[0;34m(f, map_location, pickle_module, **pickle_load_args)\u001b[0m\n\u001b[1;32m    579\u001b[0m     \u001b[0;32mwith\u001b[0m \u001b[0m_open_file_like\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mf\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m'rb'\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;32mas\u001b[0m \u001b[0mopened_file\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 580\u001b[0;31m         \u001b[0;32mif\u001b[0m \u001b[0m_is_zipfile\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mopened_file\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    581\u001b[0m             \u001b[0;31m# The zipfile reader is going to advance the current file position.\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/anaconda3/envs/gcn2/lib/python3.9/site-packages/torch/serialization.py\u001b[0m in \u001b[0;36m_is_zipfile\u001b[0;34m(f)\u001b[0m\n\u001b[1;32m     56\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 57\u001b[0;31m     \u001b[0mbyte\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mf\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mread\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;36m1\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     58\u001b[0m     \u001b[0;32mwhile\u001b[0m \u001b[0mbyte\u001b[0m \u001b[0;34m!=\u001b[0m \u001b[0;34m\"\"\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;31mKeyboardInterrupt\u001b[0m: ",
      "\nDuring handling of the above exception, another exception occurred:\n",
      "\u001b[0;31mKeyboardInterrupt\u001b[0m                         Traceback (most recent call last)",
      "\u001b[0;32m<ipython-input-250-6ba3a926d8c5>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[1;32m      4\u001b[0m     \u001b[0;32mfor\u001b[0m \u001b[0mi\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mtqdm\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mrange\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mlen\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mdatasets\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mmode\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      5\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m----> 6\u001b[0;31m         \u001b[0mX\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mdatasets\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mmode\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mi\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m      7\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m<ipython-input-1-15dfa748da3c>\u001b[0m in \u001b[0;36m__getitem__\u001b[0;34m(self, idx)\u001b[0m\n\u001b[1;32m    147\u001b[0m             \u001b[0msample\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mtorch\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mload\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mpt_path\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0mmap_location\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;34m'cuda'\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    148\u001b[0m         \u001b[0;32mexcept\u001b[0m \u001b[0mKeyboardInterrupt\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 149\u001b[0;31m             \u001b[0;32mraise\u001b[0m \u001b[0mKeyboardInterrupt\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    150\u001b[0m         \u001b[0;32mexcept\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    151\u001b[0m             \u001b[0mprint\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34mf\"Couldn't load {pt_path}\"\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;31mKeyboardInterrupt\u001b[0m: "
     ]
    }
   ],
   "source": [
    "from tqdm import tqdm\n",
    "import numpy as np\n",
    "\n",
    "DEBUG_MODES  = ['train','val','test']\n",
    "for mode in DEBUG_MODES:\n",
    "    for i in tqdm(range(len(datasets[mode]))):\n",
    "        X = datasets[mode][i]\n",
    "       \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "327c5b32",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Batch(G=[836909, 4], L=[92884, 1], P=[725230, 5], batch=[836909], edge_index=[2, 4427010], edge_type=[4427010], mean_pkts_rate=[836909], n_links=[836909], n_paths=[836909], out_delay=[725230], out_occupancy=[92884], ptr=[501], type=[836909])\n"
     ]
    }
   ],
   "source": [
    "%matplotlib inline\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "import warnings\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import scipy.stats as st\n",
    "from scipy.stats._continuous_distns import _distn_names\n",
    "import matplotlib\n",
    "matplotlib.rcParams['figure.figsize'] = (16.0, 12.0)\n",
    "matplotlib.style.use('ggplot')\n",
    "\n",
    "def get_train_distribution_statistics(ds_train,attrs_to_normalize,calc_distr=False):\n",
    "    torch.manual_seed(42)\n",
    "    dl_train_oneshot = DataLoader(ds_train,batch_size=500,shuffle=True)\n",
    "    for sample in dl_train_oneshot:\n",
    "        print(sample)\n",
    "        means=dict([(k, torch.mean(getattr(sample,k).float(),axis=0))\\\n",
    "                   for k in attrs_to_normalize])\n",
    "        std=dict([(k, torch.std(getattr(sample,k).float(),axis=0))\\\n",
    "                   for k in attrs_to_normalize])\n",
    "        sample.out_occupancy.cpu().numpy()\n",
    "        break\n",
    "    return means, std\n",
    "\n",
    "means, stds = get_train_distribution_statistics(datasets['train'],['P','G','L','out_occupancy','out_delay'])\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "8259a004",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import pandas as pd\n",
    "from torch_geometric.nn import MessagePassing\n",
    "from torch_geometric.utils import add_self_loops, degree\n",
    "\n",
    "from torch.nn import Sequential, Linear, BatchNorm1d, ReLU, Sigmoid\n",
    "from torch_geometric.nn.conv import GATConv,TransformerConv,GCNConv,GINConv\n",
    "\n",
    "from torch_geometric_temporal.nn.recurrent import GConvGRU\n",
    "import torch.nn.functional as F\n",
    "from torch_scatter import scatter\n",
    "\n",
    "def separate_edge_timesteps(edge_index,edge_type):\n",
    "    all_edges= [[] for _ in range(3)]\n",
    "    for et in [0,1,2]:\n",
    "        et_edges = edge_index[:,edge_type==et]\n",
    "\n",
    "        init_tensor = torch.cat([torch.ones(1,device=et_edges.device).long(),torch.diff(et_edges[0,:])],dim=0)\n",
    "        init_tensor = torch.clip(torch.abs(init_tensor),0.,1.)\n",
    "        # [1, 0, 0, 0, 1, 0, 0, 1] where 0 iff edge source equal\n",
    "        \"\"\"  Debug: \n",
    "        init_tensor = torch.as_tensor([1,0, 0, 0, 1, 0, 0, 1] ) \n",
    "        sol =  [0,1,2,3,0,1,2,0]\n",
    "        \"\"\"\n",
    "        # [0, 1, 1, 1, 0, 1, 1, 0] where 1 iff edge source equal\n",
    "        init_tensor = 1 - init_tensor\n",
    "        # [0, 1, 1, 1, -4, 1, 1, -3] where 1 iff edge source equal\n",
    "        count_tensor = torch.nonzero(1-init_tensor).view(-1)\n",
    "        init_tensor[count_tensor[1:]] = -torch.diff(count_tensor) +1\n",
    "        # [0, 1, 2, 3, 0, 1, 2, 0] where 1 iff edge source equal\n",
    "        init_tensor = init_tensor.cumsum(axis=0)\n",
    "\n",
    "        # Will list all dsts that were the first linked to some src, then all second, etc..\n",
    "        ensure_stable = torch.linspace(start=0.0,end=0.5,steps=init_tensor.shape[0],device=init_tensor.device)\n",
    "        encountered_order = torch.sort(init_tensor+ensure_stable)[1]\n",
    "        et_edges = et_edges[:,encountered_order]\n",
    "        \n",
    "        \n",
    "        #vals[i] == number of edges that belong to time step i\n",
    "        idxs, vals = torch.unique(init_tensor,return_counts=True)\n",
    "        vs = [x for x in torch.split_with_sizes(et_edges,tuple(vals),dim=1)]\n",
    "         \n",
    "        \n",
    "        #if not torch.as_tensor([v.shape[1] for v in vs]).sum().item() == et_edges.shape[1]:\n",
    "        #    raise f\"Sum of disjoint timesteps is {torch.as_tensor([v.shape[0] for v in vs]).sum()} but should be {et_edges.shape[1]}\"\n",
    "\n",
    "        all_edges[et] = vs\n",
    "        \n",
    "    return all_edges\n",
    "\n",
    "import torch\n",
    "from torch_geometric.nn import MessagePassing\n",
    "from torch_geometric.utils import add_self_loops, degree\n",
    "\n",
    "\n",
    "class Baseline(torch.nn.Module):\n",
    "\n",
    "    \n",
    "    def __init__(self,num_iterations=5,G_dim = 4,P_dim =5,L_dim = 1,**kwargs):\n",
    "        super(Baseline, self).__init__(**kwargs)\n",
    "        self.num_iterations = num_iterations\n",
    "        self.G_dim, self.P_dim, self.L_dim = G_dim, P_dim, L_dim\n",
    "        self.H_n, self.H_p, self.H_l = 2,2,2\n",
    "    \n",
    "    def forward(self,data,means,stds,naive=True,mode=None):\n",
    "        edge_index = data.edge_index.long()\n",
    "        edge_type = data.edge_type.clone()\n",
    "        \n",
    "        is_p = data.type==0\n",
    "        is_l = data.type==1\n",
    "        is_n = data.type==2       \n",
    "        \n",
    "        \n",
    "        all_edges = separate_edge_timesteps(edge_index,edge_type)\n",
    "        \n",
    "        \"\"\" Each element $i$ of pl_by_time contains all edges that occur at  \n",
    "            position $i$ in some path.\n",
    "        \"\"\"\n",
    "        pl_at_time = all_edges[0]\n",
    "        \n",
    "        edges_pl = edge_index[:,edge_type==0]\n",
    "        edges_pn = edge_index[:,edge_type==1]\n",
    "        edges_ln = edge_index[:,edge_type==2]    \n",
    "        \n",
    "\n",
    "        G_dim, P_dim, L_dim = self.G_dim, self.P_dim, self.L_dim\n",
    "        H_n, H_p, H_l = self.H_n, self.H_p, self.H_l\n",
    "        \n",
    "        n_p = torch.sum(is_p)\n",
    "        n_l = torch.sum(is_l)\n",
    "        n_n = torch.sum(is_n)\n",
    "        \n",
    "        \"\"\" Get true value of P,L,G\"\"\"\n",
    "        P = data.P * data.mean_pkts_rate[is_p].view(-1,1)\n",
    "        L = data.L * data.mean_pkts_rate[is_l].view(-1,1)\n",
    "        L = L / 1000\n",
    "        G = data.G \n",
    "        \n",
    "          \n",
    "        \n",
    "        cnt = 0\n",
    "        cnt, node_hidden = cnt+H_n,   slice(cnt,cnt+H_n)\n",
    "        cnt, node_og     = cnt+G_dim, slice(cnt,cnt+G_dim)\n",
    "        cnt, link_hidden = cnt+H_l,   slice(cnt,cnt+H_l)\n",
    "        cnt, link_og     = cnt+L_dim, slice(cnt,cnt+L_dim)\n",
    "        cnt, path_hidden = cnt+H_p,   slice(cnt,cnt+H_p)\n",
    "        cnt, path_og     = cnt+P_dim, slice(cnt,cnt+P_dim)\n",
    "        \n",
    "        cnt = 0\n",
    "        cnt, node_all = cnt+H_n+G_dim, slice(cnt,cnt+H_n+G_dim)\n",
    "        cnt, link_all = cnt+H_l+L_dim, slice(cnt,cnt+H_l+L_dim)\n",
    "        cnt, path_all = cnt+H_p+P_dim, slice(cnt,cnt+H_p+P_dim)\n",
    "        \n",
    "        X = torch.zeros(data.G.size(0),H_n+G_dim+H_p+P_dim+H_l+L_dim,device='cuda')\n",
    "        X[:,node_og] = G[:,:]\n",
    "        X[is_l,link_og] = L\n",
    "        X[is_p,path_og] = P\n",
    "        \n",
    "        \"\"\"\n",
    "            Get Average bandwidth\n",
    "        \"\"\"\n",
    "        A = X[:,path_og.stop-2].view(-1,).clone() #Avg pkts sent\n",
    "        blocking_probs =  0.3*torch.ones_like(A)\n",
    "        \n",
    "        max_numpaths = len(pl_at_time) \n",
    "        T =  torch.zeros(X.size(0),device=A.device)\n",
    "        rhos =  torch.zeros(X.size(0),device=A.device)      \n",
    "        \n",
    "        \"\"\"\n",
    "            \\trafic[k]_{i}: traffic passing on some edge that appears in order k at path\n",
    "        \"\"\"\n",
    "        def update_traffic(L,T,A,pl_at_time,blocking_probs):\n",
    "            multiplier = 1.0\n",
    "            T = torch.zeros_like(T)\n",
    "            N =  torch.zeros_like(T)\n",
    "            for k in range(max_numpaths):\n",
    "                \n",
    "                if k == 0:\n",
    "                    \"\"\" Just map the demand on the respective path\"\"\"\n",
    "                    traffic = A.clone()\n",
    "                else:\n",
    "                    prev_paths = pl_at_time[k-1][0,:]\n",
    "                    prev_edges = pl_at_time[k-1][1,:]\n",
    "                    prev_edges_block_probs = torch.gather(blocking_probs,dim=0,\n",
    "                                                         index=prev_edges)\n",
    "\n",
    "                    traffic[prev_paths] *= (1.0 - prev_edges_block_probs)\n",
    "                    \n",
    "                which_paths = pl_at_time[k][0,:]\n",
    "                which_edges = pl_at_time[k][1,:]\n",
    "                T += scatter(src=torch.gather(traffic,0,which_paths),\n",
    "                             index=which_edges,\n",
    "                            dim=0,dim_size=X.size(0),reduce='sum')\n",
    "                N += scatter(src=torch.ones_like(torch.gather(traffic,0,which_paths)),\n",
    "                             index=which_edges,\n",
    "                            dim=0,dim_size=X.size(0),reduce='sum')\n",
    "                #print(T[is_l].mean())    \n",
    "            #T = T/torch.maximum(N,torch.ones_like(N))\n",
    "            #T /= max_numpaths\n",
    "            return T,N\n",
    "        B = buffer_size = 32\n",
    "        def update_blocking_probs(L,T,A,pl_at_time,blocking_probs):\n",
    "            blocking_probs = 0.0*blocking_probs\n",
    "            rhos = 0.0*blocking_probs      \n",
    "            rhos[is_l] = T[is_l] / X[is_l,link_og.start]\n",
    "            #print(rhos[is_l].mean())\n",
    "            \n",
    "            blocking_probs_num = (1.0 - rhos) * torch.pow(rhos,buffer_size)\n",
    "            blocking_probs_den = 1.0 - torch.pow(rhos,buffer_size+1)\n",
    "            return blocking_probs_num/(blocking_probs_den+1e-08)\n",
    "        \n",
    "        for t in range(self.num_iterations):\n",
    "            T, N = update_traffic(L,T,A,pl_at_time,blocking_probs)\n",
    "            #print(\"mean traffic: \",T[is_l].mean().item())\n",
    "            \n",
    "            blocking_probs = update_blocking_probs(L,T,A,pl_at_time,blocking_probs)\n",
    "\n",
    "            #print(\"mean block p.: \",blocking_probs[is_l].mean().item())\n",
    "            rhos = T[is_l] / (X[is_l,link_og.start])\n",
    "            pi_0 = (1 - rhos)/(1-torch.pow(rhos,B+1))\n",
    "            res = 1*pi_0\n",
    "            for j in range(32):\n",
    "                pi_0 = pi_0*rhos\n",
    "                res += (j+1)*pi_0\n",
    "                \n",
    "            res = res/32\n",
    "\n",
    "        L = res \n",
    "        \n",
    "        \"\"\" To predict the node, we use the formula:\n",
    "        \n",
    "            path delay ~= \\sum_{i=0}^{n_links} delay_link(i)\n",
    "            where \n",
    "                delay_link(i) := avg_utilization_{i} * (queue_size_{i}/link_capacity_{i})\n",
    "                \n",
    "            Our NN predicts avg_utilization_{i} \\in [0,1], \\forall i. \n",
    "            For this dataset, we have \\forall i: queue_size_{i} =32000\n",
    "        \"\"\"\n",
    "        X = torch.zeros(X.size(0), device=X.device)\n",
    "        data_L = data.L.squeeze(-1)\n",
    "        link_capacity = data.L.squeeze(-1) * data.mean_pkts_rate[is_l]\n",
    "        X[is_l] = L.squeeze(-1)  * 32000.0  / link_capacity\n",
    "        E   = torch.gather(X,index=edges_pl[1,:],dim=0)\n",
    "        res = scatter(src=E,index=edges_pl[0,:],dim=0,dim_size=X.size(0),reduce='sum')\n",
    "        res = res[is_p]\n",
    "        return res, L\n",
    "        \n",
    "        \n",
    "\n",
    "class ChallengeModel(torch.nn.Module):\n",
    "\n",
    "    \n",
    "    def __init__(self,num_iterations=3,G_dim = 4,P_dim =5,L_dim = 1,**kwargs):\n",
    "        super(ChallengeModel, self).__init__(**kwargs)\n",
    "        self.num_iterations = num_iterations\n",
    "        self.G_dim = G_dim\n",
    "        self.P_dim = P_dim\n",
    "        self.L_dim = L_dim\n",
    "        \n",
    "        self.H = 64\n",
    "        self.H_p = 64\n",
    "        self.H_l = 64\n",
    "        self.H_n = 64\n",
    "        \n",
    "        self.conv_pn_1 = []\n",
    "        self.conv_ln_1 = []\n",
    "        self.conv_pn_r_1 = []\n",
    "        self.conv_ln_r_1 = []\n",
    "        self.conv_pl_1 = []\n",
    "        self.conv_pl_r_1 = []\n",
    "        \n",
    "        for i in range(self.num_iterations):\n",
    "            self.conv_pn_1.append(GATConv(self.H_p+self.P_dim,self.H_n).cuda())\n",
    "            self.conv_ln_1.append(GATConv(self.H_l+L_dim,self.H_n).cuda())\n",
    "            self.conv_pl_1.append(GATConv(self.H_p+P_dim,self.H_l).cuda())\n",
    "\n",
    "            self.conv_pn_r_1.append(GATConv(self.H_n+G_dim,self.H_p,flow=\"target_to_source\").cuda())\n",
    "            self.conv_ln_r_1 .append(GATConv(self.H_n+G_dim,self.H_l,flow=\"target_to_source\").cuda())\n",
    "        \n",
    "            self.conv_pl_r_1.append(GConvGRU(self.H_l+L_dim,self.H_p,K=2).cuda())\n",
    "        \n",
    "        self.finalconv = GCNConv(self.H_p+P_dim,self.H_l,normalize=True).cuda()\n",
    "        \n",
    "        for c in ['conv_pn','conv_ln','conv_pl']:\n",
    "            setattr(self,c+'_1',torch.nn.ModuleList(getattr(self,c+'_1')))\n",
    "            setattr(self,c+'_r_1',torch.nn.ModuleList(getattr(self,c+'_r_1')))\n",
    "        \n",
    "        self.lin1 = torch.nn.Linear(L_dim+self.H_l,512)\n",
    "        self.lin2 = torch.nn.Linear(512,512)\n",
    "        self.lin3 = torch.nn.Linear(512,1)\n",
    "        self.xlin1 = torch.nn.Linear(self.H_n+self.G_dim+self.H_p+self.P_dim+self.H_l+self.L_dim,128)\n",
    "        self.xlin2 = torch.nn.Linear(128,self.H_n+self.G_dim+self.H_p+self.P_dim+self.H_l+self.L_dim)\n",
    "        \n",
    "    \n",
    "    def forward(self,data,means,stds,naive=True,mode=None,baseline_occup=None,baseline_out=None):\n",
    "        edge_index = data.edge_index.long()\n",
    "        edge_type = data.edge_type.clone()\n",
    "        \n",
    "        is_p = data.type==0\n",
    "        is_l = data.type==1\n",
    "        is_n = data.type==2       \n",
    "        \n",
    "        \n",
    "        all_edges = separate_edge_timesteps(edge_index,edge_type)\n",
    "        vs = all_edges[0]\n",
    "        \n",
    "        edges_pl = edge_index[:,edge_type==0]\n",
    "        edges_pn = edge_index[:,edge_type==1]\n",
    "        edges_ln = edge_index[:,edge_type==2]    \n",
    "        \n",
    "        #print(f\"Edges between path and link: {edges_pl.shape[1]}\")\n",
    "        #print(f\"Edges between path and node: {edges_pn.shape[1]}\")\n",
    "        #print(f\"Edges between link and node: {edges_ln.shape[1]}\")\n",
    "    \n",
    "        \n",
    "        G_dim, P_dim, L_dim = self.G_dim, self.P_dim, self.L_dim\n",
    "        H_n, H_p, H_l = self.H_n, self.H_p, self.H_l\n",
    "        \n",
    "        n_p = torch.sum(is_p)\n",
    "        n_l = torch.sum(is_l)\n",
    "        n_n = torch.sum(is_n)\n",
    "        \n",
    "        P = (data.P - means['P'])/(1e-08+stds['P'])\n",
    "        #P =  data.P\n",
    "        G = 0.0*(data.G - means['G'])/stds['G']\n",
    "        G[:,2] = 0*data.G[:,2]\n",
    "        L = (data.L - means['L'])/stds['L']\n",
    " \n",
    "        \n",
    "        \n",
    "        \n",
    "        X = torch.zeros(data.G.size(0),H_n+G_dim+H_p+P_dim+H_l+L_dim,device='cuda')\n",
    "        cnt = 0\n",
    "        cnt, node_hidden = cnt+H_n,   slice(cnt,cnt+H_n)\n",
    "        cnt, node_og     = cnt+G_dim, slice(cnt,cnt+G_dim)\n",
    "        cnt, link_hidden = cnt+H_l,   slice(cnt,cnt+H_l)\n",
    "        cnt, link_og     = cnt+L_dim, slice(cnt,cnt+L_dim)\n",
    "        cnt, path_hidden = cnt+H_p,   slice(cnt,cnt+H_p)\n",
    "        cnt, path_og     = cnt+P_dim, slice(cnt,cnt+P_dim)\n",
    "        \n",
    "        cnt = 0\n",
    "        cnt, node_all = cnt+H_n+G_dim, slice(cnt,cnt+H_n+G_dim)\n",
    "        cnt, link_all = cnt+H_l+L_dim, slice(cnt,cnt+H_l+L_dim)\n",
    "        cnt, path_all = cnt+H_p+P_dim, slice(cnt,cnt+H_p+P_dim)\n",
    "        \n",
    "        \n",
    "        X[:,node_og] = G[:,:]\n",
    "        X[is_l,link_og] = L\n",
    "        X[is_p,path_og] = P\n",
    "        \n",
    "        X[is_l,link_hidden.start] = baseline_occup\n",
    "        X[is_p,path_hidden.start] = baseline_out\n",
    "        \n",
    "        \n",
    "        X = F.leaky_relu(self.xlin1(X))\n",
    "        X = F.leaky_relu(self.xlin2(X))\n",
    "        \n",
    "\n",
    "        def act(x):\n",
    "            return F.leaky_relu(x)\n",
    "        for i in range(self.num_iterations):\n",
    "            X[is_p,path_hidden] =  act(self.conv_pn_r_1[i](X[:,node_all].clone(),edges_pn)[is_p,:])\n",
    "            \n",
    "            x = X[:,link_all].clone()\n",
    "            H = None\n",
    "            max_numpaths = len(vs)\n",
    "            for k in range(max_numpaths):\n",
    "                e = torch.cat([vs[k][1,:].unsqueeze(0),vs[k][0,:].unsqueeze(0)],axis=0)\n",
    "                H  = self.conv_pl_r_1[0](X=x,H=H,edge_index=e)    \n",
    "            X[is_p,path_hidden] = act(H[is_p,:]/max_numpaths)\n",
    "            X[is_p,path_hidden.start] = baseline_out\n",
    "            \n",
    "            X[is_n,node_hidden] = \\\n",
    "                act(self.conv_pn_1[i](X[:,path_all].clone(), edges_pn)[is_n,:] +\\\n",
    "                              self.conv_ln_1[i](X[:,link_all].clone(),edges_ln)[is_n,:])\n",
    "\n",
    "            \n",
    "            X[is_l,link_hidden] =  act(self.conv_ln_r_1[i](X[:,node_all].clone(),edges_ln)[is_l,:])\n",
    "            X[is_l,link_hidden] =  act(self.conv_pl_1[i](X[:,path_all].clone(),\n",
    "                                                         edges_pl)[is_l,:])\n",
    "            X[is_l,link_hidden.start] = baseline_occup\n",
    "        \n",
    "            \n",
    "        L = X[is_l,link_all]\n",
    "        L = self.lin1(L)\n",
    "        \n",
    "        L = F.leaky_relu(L)\n",
    "        L = F.leaky_relu(self.lin2(L))\n",
    "        #L = F.leaky_relu(L)\n",
    "        L = torch.sigmoid(self.lin3(L)) \n",
    "        #lamb = (1/0.05)\n",
    "        #L = -(1/lamb)* torch.log(1-0.99*L)\n",
    "        \n",
    "        \n",
    "        X = torch.zeros(X.size(0), device=X.device)\n",
    "        \n",
    "        \n",
    "        \"\"\" To predict the node, we use the formula:\n",
    "        \n",
    "            path delay ~= \\sum_{i=0}^{n_links} delay_link(i)\n",
    "            where \n",
    "                delay_link(i) := avg_utilization_{i} * (queue_size_{i}/link_capacity_{i})\n",
    "                \n",
    "            Our NN predicts avg_utilization_{i} \\in [0,1], \\forall i. \n",
    "            For this dataset, we have \\forall i: queue_size_{i} =32000\n",
    "        \"\"\"\n",
    "        link_capacity = data.L.squeeze(-1) * data.mean_pkts_rate[is_l]\n",
    "        \n",
    "        \n",
    "        \n",
    "        X[is_l] = L.squeeze(-1)  * 32000.0  / link_capacity\n",
    "        E   = torch.gather(X,index=edges_pl[1,:],dim=0)\n",
    "        \n",
    "        #print(f\"Shape after gather {E.shape}\")\n",
    "        res = scatter(src=E,index=edges_pl[0,:],dim=0,dim_size=X.size(0),reduce='sum')\n",
    "        res = res[is_p]\n",
    "        \n",
    "        return res, L\n",
    "               \n",
    "model = ChallengeModel().cuda()\n",
    "baseline = Baseline().cuda()\n",
    "       \n",
    "       "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2e6cc9d6",
   "metadata": {},
   "source": [
    "<h2> Use the cell below to train from scratch.</h2> <br/>\n",
    "Warning: Unlike the other model, training this from scratch may take multiple hours!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "2f70cca7",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "  3%|▎         | 2/65 [00:01<00:41,  1.52it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([[0.1144, 0.1163],\n",
      "        [0.1166, 0.1196],\n",
      "        [0.1885, 0.1938],\n",
      "        ...,\n",
      "        [0.1437, 0.1463],\n",
      "        [0.0993, 0.0999],\n",
      "        [0.3237, 0.3315]], device='cuda:0')\n",
      "tensor(1.9805, device='cuda:0')\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 65/65 [00:18<00:00,  3.61it/s]\n",
      "  2%|▏         | 1/65 [00:00<00:07,  8.50it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "======================================\n",
      "Epoch 0 - Avg stats (val_1)\n",
      "Mean loss: 2.7116544760190524\n",
      "Mean out: 0.44985095354226917\n",
      "Mean occup: 0.03714364377351908\n",
      "Mean actual_out: 0.4635581649266757\n",
      "Mean actual_occup: 0.035495109111070636\n",
      "tensor([[0.0517, 0.0517],\n",
      "        [0.2008, 0.2010],\n",
      "        [0.0314, 0.0312],\n",
      "        ...,\n",
      "        [0.0375, 0.0376],\n",
      "        [0.7551, 0.7673],\n",
      "        [0.0177, 0.0178]], device='cuda:0')\n",
      "tensor(1.0508, device='cuda:0')\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 65/65 [00:22<00:00,  2.85it/s]\n",
      "  2%|▏         | 1/65 [00:00<00:10,  6.29it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "======================================\n",
      "Epoch 0 - Avg stats (val_2)\n",
      "Mean loss: 1.332254853615394\n",
      "Mean out: 0.07160714635482202\n",
      "Mean occup: 0.08163771801269971\n",
      "Mean actual_out: 0.07174499582212705\n",
      "Mean actual_occup: 0.08150574467503108\n",
      "tensor([[0.0438, 0.0444],\n",
      "        [0.0484, 0.0489],\n",
      "        [0.0819, 0.0816],\n",
      "        ...,\n",
      "        [0.0562, 0.0556],\n",
      "        [0.0712, 0.0712],\n",
      "        [0.0115, 0.0115]], device='cuda:0')\n",
      "tensor(1.4480, device='cuda:0')\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 65/65 [00:37<00:00,  1.73it/s]\n",
      "  0%|          | 0/65 [00:00<?, ?it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "======================================\n",
      "Epoch 0 - Avg stats (val_3)\n",
      "Mean loss: 1.6559356038387005\n",
      "Mean out: 0.05248686512215779\n",
      "Mean occup: 0.062194769027141425\n",
      "Mean actual_out: 0.05243281760754494\n",
      "Mean actual_occup: 0.06198185212337054\n",
      "Flushed\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "  3%|▎         | 2/65 [00:00<00:10,  5.98it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([[0.1144, 0.1163],\n",
      "        [0.1166, 0.1196],\n",
      "        [0.1885, 0.1938],\n",
      "        ...,\n",
      "        [0.1437, 0.1463],\n",
      "        [0.0993, 0.0999],\n",
      "        [0.3237, 0.3315]], device='cuda:0')\n",
      "tensor(1.9805, device='cuda:0')\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 66%|██████▌   | 43/65 [00:08<00:04,  5.05it/s]\n"
     ]
    },
    {
     "ename": "KeyboardInterrupt",
     "evalue": "",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mKeyboardInterrupt\u001b[0m                         Traceback (most recent call last)",
      "\u001b[0;32m<ipython-input-11-0127940f8bec>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[1;32m     63\u001b[0m                         \u001b[0mopt\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mzero_grad\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     64\u001b[0m                     \u001b[0mcnt\u001b[0m \u001b[0;34m+=\u001b[0m \u001b[0;36m1\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 65\u001b[0;31m                     out, occup = model(sample,means,stds,mode,baseline_occup=b_occup,\n\u001b[0m\u001b[1;32m     66\u001b[0m                                       baseline_out=b_out)\n\u001b[1;32m     67\u001b[0m                     \u001b[0;32mif\u001b[0m \u001b[0;32mFalse\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/anaconda3/envs/gcn2/lib/python3.9/site-packages/torch/nn/modules/module.py\u001b[0m in \u001b[0;36m_call_impl\u001b[0;34m(self, *input, **kwargs)\u001b[0m\n\u001b[1;32m    887\u001b[0m             \u001b[0mresult\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_slow_forward\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0minput\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    888\u001b[0m         \u001b[0;32melse\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 889\u001b[0;31m             \u001b[0mresult\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mforward\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0minput\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    890\u001b[0m         for hook in itertools.chain(\n\u001b[1;32m    891\u001b[0m                 \u001b[0m_global_forward_hooks\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mvalues\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m<ipython-input-10-938fbc699c0b>\u001b[0m in \u001b[0;36mforward\u001b[0;34m(self, data, means, stds, naive, mode, baseline_occup, baseline_out)\u001b[0m\n\u001b[1;32m    327\u001b[0m             \u001b[0;32mfor\u001b[0m \u001b[0mk\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mrange\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mmax_numpaths\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    328\u001b[0m                 \u001b[0me\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mtorch\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mcat\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mvs\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mk\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;36m1\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0munsqueeze\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;36m0\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0mvs\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mk\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;36m0\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0munsqueeze\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;36m0\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0maxis\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;36m0\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 329\u001b[0;31m                 \u001b[0mH\u001b[0m  \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mconv_pl_r_1\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;36m0\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mX\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mx\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0mH\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mH\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0medge_index\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0me\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    330\u001b[0m             \u001b[0mX\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mis_p\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0mpath_hidden\u001b[0m\u001b[0;34m]\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mact\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mH\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mis_p\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m/\u001b[0m\u001b[0mmax_numpaths\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    331\u001b[0m             \u001b[0mX\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mis_p\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0mpath_hidden\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mstart\u001b[0m\u001b[0;34m]\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mbaseline_out\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/anaconda3/envs/gcn2/lib/python3.9/site-packages/torch/nn/modules/module.py\u001b[0m in \u001b[0;36m_call_impl\u001b[0;34m(self, *input, **kwargs)\u001b[0m\n\u001b[1;32m    887\u001b[0m             \u001b[0mresult\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_slow_forward\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0minput\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    888\u001b[0m         \u001b[0;32melse\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 889\u001b[0;31m             \u001b[0mresult\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mforward\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0minput\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    890\u001b[0m         for hook in itertools.chain(\n\u001b[1;32m    891\u001b[0m                 \u001b[0m_global_forward_hooks\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mvalues\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/anaconda3/envs/gcn2/lib/python3.9/site-packages/torch_geometric_temporal/nn/recurrent/gconv_gru.py\u001b[0m in \u001b[0;36mforward\u001b[0;34m(self, X, edge_index, edge_weight, H)\u001b[0m\n\u001b[1;32m    149\u001b[0m         \u001b[0mZ\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_calculate_update_gate\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mX\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0medge_index\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0medge_weight\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mH\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    150\u001b[0m         \u001b[0mR\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_calculate_reset_gate\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mX\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0medge_index\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0medge_weight\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mH\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 151\u001b[0;31m         \u001b[0mH_tilde\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_calculate_candidate_state\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mX\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0medge_index\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0medge_weight\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mH\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mR\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    152\u001b[0m         \u001b[0mH\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_calculate_hidden_state\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mZ\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mH\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mH_tilde\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    153\u001b[0m         \u001b[0;32mreturn\u001b[0m \u001b[0mH\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/anaconda3/envs/gcn2/lib/python3.9/site-packages/torch_geometric_temporal/nn/recurrent/gconv_gru.py\u001b[0m in \u001b[0;36m_calculate_candidate_state\u001b[0;34m(self, X, edge_index, edge_weight, H, R)\u001b[0m\n\u001b[1;32m    120\u001b[0m     \u001b[0;32mdef\u001b[0m \u001b[0m_calculate_candidate_state\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mX\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0medge_index\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0medge_weight\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mH\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mR\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    121\u001b[0m         \u001b[0mH_tilde\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mconv_x_h\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mX\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0medge_index\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0medge_weight\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 122\u001b[0;31m         \u001b[0mH_tilde\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mH_tilde\u001b[0m \u001b[0;34m+\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mconv_h_h\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mH\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0mR\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0medge_index\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0medge_weight\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    123\u001b[0m         \u001b[0mH_tilde\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mtorch\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mtanh\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mH_tilde\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    124\u001b[0m         \u001b[0;32mreturn\u001b[0m \u001b[0mH_tilde\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/anaconda3/envs/gcn2/lib/python3.9/site-packages/torch/nn/modules/module.py\u001b[0m in \u001b[0;36m_call_impl\u001b[0;34m(self, *input, **kwargs)\u001b[0m\n\u001b[1;32m    887\u001b[0m             \u001b[0mresult\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_slow_forward\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0minput\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    888\u001b[0m         \u001b[0;32melse\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 889\u001b[0;31m             \u001b[0mresult\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mforward\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0minput\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    890\u001b[0m         for hook in itertools.chain(\n\u001b[1;32m    891\u001b[0m                 \u001b[0m_global_forward_hooks\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mvalues\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/anaconda3/envs/gcn2/lib/python3.9/site-packages/torch_geometric/nn/conv/cheb_conv.py\u001b[0m in \u001b[0;36mforward\u001b[0;34m(self, x, edge_index, edge_weight, batch, lambda_max)\u001b[0m\n\u001b[1;32m    124\u001b[0m         \u001b[0;32massert\u001b[0m \u001b[0mlambda_max\u001b[0m \u001b[0;32mis\u001b[0m \u001b[0;32mnot\u001b[0m \u001b[0;32mNone\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    125\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 126\u001b[0;31m         edge_index, norm = self.__norm__(edge_index, x.size(self.node_dim),\n\u001b[0m\u001b[1;32m    127\u001b[0m                                          \u001b[0medge_weight\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mnormalization\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    128\u001b[0m                                          \u001b[0mlambda_max\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mdtype\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mx\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mdtype\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/anaconda3/envs/gcn2/lib/python3.9/site-packages/torch_geometric/nn/conv/cheb_conv.py\u001b[0m in \u001b[0;36m__norm__\u001b[0;34m(self, edge_index, num_nodes, edge_weight, normalization, lambda_max, dtype, batch)\u001b[0m\n\u001b[1;32m     93\u001b[0m         \u001b[0medge_index\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0medge_weight\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mremove_self_loops\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0medge_index\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0medge_weight\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     94\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 95\u001b[0;31m         edge_index, edge_weight = get_laplacian(edge_index, edge_weight,\n\u001b[0m\u001b[1;32m     96\u001b[0m                                                 \u001b[0mnormalization\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mdtype\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     97\u001b[0m                                                 num_nodes)\n",
      "\u001b[0;32m~/anaconda3/envs/gcn2/lib/python3.9/site-packages/torch_geometric/utils/get_laplacian.py\u001b[0m in \u001b[0;36mget_laplacian\u001b[0;34m(edge_index, edge_weight, normalization, dtype, num_nodes)\u001b[0m\n\u001b[1;32m     62\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     63\u001b[0m         \u001b[0;31m# L = I - A_norm.\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 64\u001b[0;31m         edge_index, tmp = add_self_loops(edge_index, -edge_weight,\n\u001b[0m\u001b[1;32m     65\u001b[0m                                          fill_value=1., num_nodes=num_nodes)\n\u001b[1;32m     66\u001b[0m         \u001b[0;32massert\u001b[0m \u001b[0mtmp\u001b[0m \u001b[0;32mis\u001b[0m \u001b[0;32mnot\u001b[0m \u001b[0;32mNone\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/anaconda3/envs/gcn2/lib/python3.9/site-packages/torch_geometric/utils/loop.py\u001b[0m in \u001b[0;36madd_self_loops\u001b[0;34m(edge_index, edge_weight, fill_value, num_nodes)\u001b[0m\n\u001b[1;32m     87\u001b[0m     \u001b[0;32mif\u001b[0m \u001b[0medge_weight\u001b[0m \u001b[0;32mis\u001b[0m \u001b[0;32mnot\u001b[0m \u001b[0;32mNone\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     88\u001b[0m         \u001b[0;32massert\u001b[0m \u001b[0medge_weight\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mnumel\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;34m==\u001b[0m \u001b[0medge_index\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0msize\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;36m1\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 89\u001b[0;31m         \u001b[0mloop_weight\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0medge_weight\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mnew_full\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mN\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m)\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mfill_value\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     90\u001b[0m         \u001b[0medge_weight\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mtorch\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mcat\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0medge_weight\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mloop_weight\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mdim\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;36m0\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     91\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;31mKeyboardInterrupt\u001b[0m: "
     ]
    }
   ],
   "source": [
    "from tqdm import tqdm,trange\n",
    "from convertDataset import total_samples\n",
    "from torch.nn import MSELoss\n",
    "def MAPE(preds,actuals):\n",
    "    return 100.0*torch.mean(torch.abs((preds-actuals)/actuals))\n",
    "\n",
    "def mape_all(preds,actuals):\n",
    "    return 100.0*torch.abs((preds-actuals)/actuals)\n",
    "def lMAPE(preds,actuals):\n",
    "    return 100.0*torch.mean(torch.abs((torch.log(preds)-torch.log(actuals))/actuals))\n",
    "\n",
    "def MSE(preds,actuals):\n",
    "    return torch.sqrt(torch.mean(torch.square(preds-actuals)))\n",
    "import numpy as np\n",
    "\n",
    "\n",
    "from torch.utils.tensorboard import SummaryWriter\n",
    "writer = SummaryWriter()\n",
    "num_epochs = 300\n",
    "opt = torch.optim.Adam(lr=1e-3,params=model.parameters())\n",
    "step = 0\n",
    "\n",
    "torch.manual_seed(420)\n",
    "\n",
    "\n",
    "#model.load_state_dict(torch.load(f'./22_setembro_modelo.pt'))\n",
    "\n",
    "for epoch in range(num_epochs):\n",
    "    \n",
    "        \n",
    "    for mode in ['train','val_1','val_2','val_3']:\n",
    "        if mode == 'train':\n",
    "                model.train()\n",
    "        else:\n",
    "                model.eval()\n",
    "        stats = {'loss':[],\n",
    "                 'out':[],\n",
    "                 'occup':[],\n",
    "                 'actual_out':[],\n",
    "                 'actual_occup':[]\n",
    "                }\n",
    "                \n",
    "        running_loss = []\n",
    "        out_mean = []\n",
    "        occup_mean = []\n",
    "        actual_out_mean = []\n",
    "        cnt = 0\n",
    "        total=len(dataloaders[mode])//30 if mode == 'train' else len(dataloaders[mode])\n",
    "        mode_maybe_shuffle = 'train_s' if mode == 'train' else mode\n",
    "        for i,sample in tqdm(enumerate(dataloaders[mode_maybe_shuffle]),\n",
    "                             total=total):\n",
    "            \n",
    "            #with torch.autograd.detect_anomaly():\n",
    "\n",
    "                with torch.set_grad_enabled(False):\n",
    "                    b_out, b_occup = baseline(sample,means,stds,mode)\n",
    "\n",
    "\n",
    "                with torch.set_grad_enabled(mode == 'train'):\n",
    "                    if i == total:\n",
    "                        break\n",
    "                    if mode == 'train':\n",
    "                        opt.zero_grad()\n",
    "                    cnt += 1\n",
    "                    out, occup = model(sample,means,stds,mode,baseline_occup=b_occup,\n",
    "                                      baseline_out=b_out)\n",
    "                    if False:\n",
    "                        loss = mape_all(out,sample.out_delay)\n",
    "                        loss = scatter(src=loss,index=sample.batch[sample.type==0],\n",
    "                                       dim=0,dim_size=(sample.batch.max()+1),reduce='mean')\n",
    "                        loss = loss.mean()\n",
    "                    else:\n",
    "                        loss = MAPE(out,sample.out_delay)\n",
    "                    if mode == 'train':\n",
    "                        #MSE(occup,sample.out_occupancy).backward()\n",
    "                        lMAPE(out,sample.out_delay).backward()\n",
    "                        opt.step()\n",
    "                    elif i == 0:\n",
    "                        print(torch.cat([out.view(-1,1),sample.out_delay.view(-1,1)],axis=1))\n",
    "                        print(loss)\n",
    "\n",
    "                    _stats = {'loss':loss,\n",
    "                              'out':out.mean(),\n",
    "                              'actual_out':sample.out_delay.mean() if not mode == 'test' else -1.0,\n",
    "                              'occup':occup.mean(),\n",
    "                              'actual_occup':sample.out_occupancy.mean() if not mode == 'test' else -1.0,\n",
    "                             }\n",
    "                    for k in _stats.keys():\n",
    "                        stats[k].append(_stats[k].cpu().item())\n",
    "                    del _stats\n",
    "        print(\"======================================\")\n",
    "        print(f\"Epoch {epoch} - Avg stats ({mode})\")\n",
    "        for k in stats.keys():\n",
    "            print(f'Mean {k}: {np.array(stats[k]).mean()}')\n",
    "            writer.add_scalar(f\"{k}/{mode}\", np.array(stats[k]).mean(), step)\n",
    "    writer.flush()\n",
    "    print(\"Flushed\")\n",
    "    step += 1\n",
    "\n",
    "    torch.cuda.empty_cache()\n",
    "\n",
    "    import os\n",
    "    os.makedirs('./model',exist_ok=True)\n",
    "    if i%1 == 0:\n",
    "        torch.save(model.state_dict(),f'./model/model_{epoch}.pt')\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d327c269",
   "metadata": {},
   "source": [
    "## Use the cell below to load the pre-trained model.\n",
    "\n",
    "You should see the message: \\<All keys matched successfully\\>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "83d7d3cb",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<All keys matched successfully>"
      ]
     },
     "execution_count": 24,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "\n",
    "model.load_state_dict(torch.load(f'./22_setembro_modelo.pt'))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "54de317e",
   "metadata": {},
   "source": [
    "## Create Test prediction"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "97dcc33c",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 390/390 [04:42<00:00,  1.38it/s]\n"
     ]
    }
   ],
   "source": [
    "model.eval()\n",
    "import os\n",
    "os.makedirs('./predictions',exist_ok=True)\n",
    "FILENAME = '22_setembro_NEW.csv'\n",
    "upload_file = open(f'./predictions/{FILENAME}', \"w\")\n",
    "from tqdm import tqdm\n",
    "\n",
    "cnt = 0\n",
    "for i,sample in tqdm(enumerate(dataloaders['test']),total=len(dataloaders['test'])):\n",
    "    with torch.set_grad_enabled(False):\n",
    "        b_out, b_occup = baseline(sample,means,stds,'test')\n",
    "    \n",
    "    with torch.set_grad_enabled(False):\n",
    "        out, occup = model(sample,means,stds,'test',baseline_occup=b_occup,\n",
    "                                      baseline_out=b_out)\n",
    "        batch = sample.batch[sample.type==0]\n",
    "\n",
    "        for b in range(sample.batch.max() + 1):\n",
    "            if cnt > 0:\n",
    "                upload_file.write(\"\\n\")\n",
    "            cnt += 1\n",
    "            out_batch = out[batch==b].cpu().numpy().round(5)            \n",
    "            upload_file.write(\"{}\".format(';'.join([str(i) for i in np.squeeze(out_batch)])))\n",
    "        \n",
    "upload_file.close()        "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0d258df9",
   "metadata": {},
   "source": [
    "### Check if test file is OK"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 459,
   "id": "3c258536",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Checking the file...\n",
      "Congratulations! The submission file has passed all the tests!\n"
     ]
    }
   ],
   "source": [
    "\n",
    "from itertools import zip_longest\n",
    "def check_submission(FILENAME,PATHS_PER_SAMPLE):\n",
    "    sample_num = 0\n",
    "    error = False\n",
    "    with open(FILENAME, \"r\") as uploaded_file, open(PATHS_PER_SAMPLE, \"r\") as path_per_sample:\n",
    "        # Load all files line by line (not at once)\n",
    "        for prediction, n_paths in zip_longest(uploaded_file, path_per_sample):\n",
    "            # Case 1: Line Count does not match.\n",
    "            if n_paths is None:\n",
    "                print(\"WARNING: File must contain 1560 lines in total for the final test datset (90 for the toy dataset). \"\n",
    "                      \"Looks like the uploaded file has {} lines\".format(sample_num))\n",
    "                error = True\n",
    "            if prediction is None:\n",
    "                print(\"WARNING: File must have 1560 lines in total for the final test datset (90 for the toy dataset). \"\n",
    "                      \"Looks like the uploaded file has {} lines\".format(sample_num))\n",
    "                error = True\n",
    "\n",
    "            # Remove the \\n at the end of lines\n",
    "            prediction = prediction.rstrip()\n",
    "            n_paths = n_paths.rstrip()\n",
    "\n",
    "            # Split the line, convert to float and then, to list\n",
    "            prediction = list(map(float, prediction.split(\";\")))\n",
    "\n",
    "            # Case 2: Wrong number of predictions in a sample\n",
    "            if int(len(prediction)) != int(n_paths):\n",
    "                print(\"WARNING in line {}: The line should have size {} but it has size {}\".format(sample_num, n_paths,\n",
    "                                                                                                   len(prediction)))\n",
    "                error = True\n",
    "\n",
    "            sample_num += 1\n",
    "\n",
    "    if not error:\n",
    "        print(\"Congratulations! The submission file has passed all the tests!\")\n",
    "    else:\n",
    "        print(\"Error\")\n",
    "        \n",
    "print(\"Checking the file...\")\n",
    "\n",
    "PATHS_PER_SAMPLE = './paths_per_sample_test_dataset.txt'\n",
    "FILEPATH= f'./predictions/{FILENAME}'\n",
    "check_submission(FILEPATH,PATHS_PER_SAMPLE)\n"
   ]
  }
 ],
 "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.9.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}