SST_GNN.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "SST-GNN.ipynb",
      "provenance": [],
      "collapsed_sections": [],
      "toc_visible": true
    },
    "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.5.2"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ILVaFChpXnQg"
      },
      "source": [
        "\n",
        "# Description\n",
        "\n",
        "\n",
        "---\n",
        "\n",
        "*Model: SST-GNN Model*\n",
        "\n",
        "\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "cNdyMxnv5nl1"
      },
      "source": [
        "# Packages"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "CMYQU_A33c8a",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "75c2f48a-7f57-4e87-8c05-0b07b59c3eb6"
      },
      "source": [
        "# !pip install numpy\n",
        "# !pip install Pillow\n",
        "# !pip install pyparsing\n",
        "# !pip install scipy\n",
        "# !pip install six\n",
        "# !pip install sklearn\n",
        "# !pip install torch\n",
        "# !pip install torchvision\n",
        "# !pip install matplotlib\n",
        "# !pip install pandas\n",
        "\n",
        "!pip install pyhocon"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Collecting pyhocon\n",
            "\u001b[?25l  Downloading https://files.pythonhosted.org/packages/af/c4/39cfc164ed77ddfe7dd80ff31b065c2b18583c4ee0a4336979c86c4b5d8f/pyhocon-0.3.57.tar.gz (110kB)\n",
            "\r\u001b[K     |███                             | 10kB 16.3MB/s eta 0:00:01\r\u001b[K     |██████                          | 20kB 10.3MB/s eta 0:00:01\r\u001b[K     |████████▉                       | 30kB 8.3MB/s eta 0:00:01\r\u001b[K     |███████████▉                    | 40kB 7.5MB/s eta 0:00:01\r\u001b[K     |██████████████▉                 | 51kB 4.4MB/s eta 0:00:01\r\u001b[K     |█████████████████▊              | 61kB 4.9MB/s eta 0:00:01\r\u001b[K     |████████████████████▊           | 71kB 5.0MB/s eta 0:00:01\r\u001b[K     |███████████████████████▋        | 81kB 5.3MB/s eta 0:00:01\r\u001b[K     |██████████████████████████▋     | 92kB 5.6MB/s eta 0:00:01\r\u001b[K     |█████████████████████████████▋  | 102kB 5.8MB/s eta 0:00:01\r\u001b[K     |████████████████████████████████| 112kB 5.8MB/s \n",
            "\u001b[?25hRequirement already satisfied: pyparsing>=2.0.3 in /usr/local/lib/python3.6/dist-packages (from pyhocon) (2.4.7)\n",
            "Building wheels for collected packages: pyhocon\n",
            "  Building wheel for pyhocon (setup.py) ... \u001b[?25l\u001b[?25hdone\n",
            "  Created wheel for pyhocon: filename=pyhocon-0.3.57-cp36-none-any.whl size=18542 sha256=a3ad56cfdcbcd3a2c530a2657ead9d4f93b471974df243fb8f6b0b77a11bc181\n",
            "  Stored in directory: /root/.cache/pip/wheels/43/ae/39/250b80e3aa8082c5fbb3f14782b9cad47de0b885f22e927a8f\n",
            "Successfully built pyhocon\n",
            "Installing collected packages: pyhocon\n",
            "Successfully installed pyhocon-0.3.57\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "hVfiZx4M6t-D"
      },
      "source": [
        "import sys\n",
        "import os\n",
        "import torch\n",
        "import argparse\n",
        "import pyhocon\n",
        "import random\n",
        "import math\n",
        "import copy\n",
        "from sklearn.utils import shuffle\n",
        "from sklearn.metrics import f1_score\n",
        "from sklearn.metrics import r2_score\n",
        "from sklearn.metrics import mean_absolute_error\n",
        "import torch.nn as nn\n",
        "import numpy as np\n",
        "import torch.nn.functional as F\n",
        "import time\n",
        "from collections import defaultdict\n",
        "import matplotlib.pyplot as plt\n",
        "import pandas as pd\n",
        "import pickle\n",
        "from datetime import datetime"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "r8zPeAdh30ex"
      },
      "source": [
        "# Data Center"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "A8OWTg9Ov3kj"
      },
      "source": [
        "class DataCenter(object):\n",
        "\n",
        "    def __init__(self, config):\n",
        "        super(DataCenter, self).__init__()\n",
        "        self.config = config\n",
        "\n",
        "    def getPositionEmbedding(self,pos):\n",
        "      input = np.arange(0,pos+1,1)\n",
        "      a = input * 360\n",
        "      day = a / 288\n",
        "      week = a / 2016\n",
        "      month = a / 8640\n",
        "      day = np.deg2rad(day)\n",
        "      week = np.deg2rad(week)\n",
        "      day = np.sin(day)\n",
        "      week = np.sin(week)\n",
        "      combined = day+week\n",
        "      return combined\n",
        "\n",
        "  \n",
        "\n",
        "    def load_data(self,ds,st_day,en_day,hr_sample,day,pred_len):\n",
        "\n",
        "        content_file = self.config['file_path.' + ds + '_content']\n",
        "\n",
        "        if ds==\"PeMSD8\" or ds == \"PeMSD4\":\n",
        "          timestamp_data = np.load(content_file)\n",
        "          timestamp_data = timestamp_data[:,:,2]\n",
        "        else:\n",
        "          timestamp_data = []\n",
        "\n",
        "          with open(content_file) as fp:\n",
        "            for i, line in enumerate(fp):\n",
        "              info = line.strip().split(\",\")\n",
        "              info = [float(x) for x in info]\n",
        "\n",
        "              timestamp_data.append(info)\n",
        "\n",
        "        timestamp_data = np.asarray(timestamp_data)\n",
        "        timestamp_data = timestamp_data.transpose()\n",
        "        tot_node = timestamp_data.shape[0]\n",
        "\n",
        "        pos = float(timestamp_data.shape[1])\n",
        "        pos_embd = self.getPositionEmbedding(pos)\n",
        "\n",
        "        \n",
        "\n",
        "        st_day -= 1\n",
        "\n",
        "        timestamp = 24 * hr_sample\n",
        "\n",
        "        ts_data = []\n",
        "        ps_data = []\n",
        "        for idx in range(st_day,en_day+1-day,1):\n",
        "\n",
        "            st_point = idx*timestamp\n",
        "            en_point = (idx+1)*timestamp\n",
        "\n",
        "            last_hour = False\n",
        "            for st in range(st_point, en_point):\n",
        "                ts = []\n",
        "                for nd in range(tot_node):\n",
        "                    a = timestamp_data[nd][st: st + (day  * timestamp) :timestamp]\n",
        "\n",
        "                    assert len(a) == day\n",
        "                    if (st + (day-1) * timestamp + pred_len) >= len(timestamp_data[nd]):\n",
        "                        last_hour = True\n",
        "                        break\n",
        "\n",
        "                    for pred in range(1,pred_len+1):\n",
        "                      gt = timestamp_data[nd][st + (day-1) * timestamp + pred]\n",
        "                      a = np.append(a, gt)\n",
        "\n",
        "                    assert len(a) == (day+pred_len)\n",
        "                    ts.append(a)\n",
        "\n",
        "                if last_hour:\n",
        "                    break\n",
        "                ts = np.asarray(ts)\n",
        "                pos_a = pos_embd[st: st + (day  * timestamp) :timestamp]\n",
        "                pos_a = np.expand_dims(pos_a,axis=0)\n",
        "                pos_a = np.repeat(pos_a,tot_node,axis=0)\n",
        "                ps_data.append(pos_a)\n",
        "                ts_data.append(ts)\n",
        "\n",
        "\n",
        "        return ts_data,ps_data\n",
        "    \n",
        "    def load_adj(self,ds):\n",
        "      W = self.load_PeMSD(self.config['file_path.'+ ds +'_cites'])\n",
        "      adj_lists = defaultdict(set)\n",
        "      for row in range(len(W)):\n",
        "        adj_lists[row] = set()\n",
        "        for col in range(len(W)):\n",
        "          if float(W[row][col]) >0 :\n",
        "            adj_lists[row].add(col)\n",
        "            adj_lists[col].add(row)\n",
        "      \n",
        "      adj = torch.zeros((len(adj_lists),len(adj_lists)))\n",
        "      for u in adj_lists:\n",
        "        for v in adj_lists[u]:\n",
        "            adj[u][v] = 1\n",
        "            adj[v][u] = 1\n",
        "      return adj\n",
        "        \n",
        "\n",
        "    def load_PeMSD(self,file_path, sigma2=0.1, epsilon=0.5, scaling=True):\n",
        "\n",
        "      try:\n",
        "          W = pd.read_csv(file_path, header=None).values\n",
        "      except FileNotFoundError:\n",
        "          print('ERROR: No File Found.')\n",
        "      \n",
        "      n = W.shape[0]\n",
        "      W = W / 10000.\n",
        "      W2, W_mask = W * W, np.ones([n, n]) - np.identity(n)\n",
        "\n",
        "      return np.exp(-W2 / sigma2) * (np.exp(-W2 / sigma2) >= epsilon) * W_mask\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "OtWLiXkR4DtC"
      },
      "source": [
        "# Utility Functions\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "sVSlgvcV4GbD"
      },
      "source": [
        "def evaluate(test_nodes,labels, graphSage, regression, device,test_loss):\n",
        "\n",
        "\n",
        "    models = [graphSage, regression]\n",
        "\n",
        "    params = []\n",
        "    for model in models:\n",
        "        for param in model.parameters():\n",
        "            if param.requires_grad:\n",
        "                param.requires_grad = False\n",
        "                params.append(param)\n",
        "\n",
        "    \n",
        "    val_nodes = test_nodes\n",
        "    embs = graphSage(val_nodes,False)\n",
        "    predicts = regression(embs)\n",
        "    loss_sup = torch.nn.MSELoss()(predicts, labels)\n",
        "    loss_sup /= len(val_nodes)\n",
        "    test_loss += loss_sup.item()\n",
        "\n",
        "    for param in params:\n",
        "        param.requires_grad = True\n",
        "\n",
        "    return predicts,test_loss\n",
        "\n",
        "def RMSELoss(yhat,y):\n",
        "    yhat = torch.FloatTensor(yhat)\n",
        "    y = torch.FloatTensor(y)\n",
        "    return torch.sqrt(torch.mean((yhat-y)**2)).item()\n",
        "\n",
        "\n",
        "def mean_absolute_percentage_error(y_true, y_pred):\n",
        "\n",
        "  y_true = np.asarray(y_true)\n",
        "  y_pred = np.asarray(y_pred)\n",
        "\n",
        "  return np.mean(np.abs((y_true - y_pred) / y_true)) * 100\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "KKHVm5cL4V16"
      },
      "source": [
        "# Downstream Task"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "lXMH_9cr4Yj0"
      },
      "source": [
        "\n",
        "class Regression(nn.Module):\n",
        "\n",
        "    def __init__(self, emb_size, out_size):\n",
        "        super(Regression, self).__init__()\n",
        "\n",
        "        self.layer = nn.Sequential(nn.Linear(emb_size, emb_size),\n",
        "                                nn.ReLU(),\n",
        "                                nn.Linear(emb_size, out_size),\n",
        "                                nn.ReLU()\n",
        "                                )\n",
        "            \n",
        "\n",
        "        self.init_params()\n",
        "\n",
        "    def init_params(self):\n",
        "        for param in self.parameters():\n",
        "            if len(param.size()) == 2:\n",
        "                nn.init.xavier_uniform_(param)\n",
        "\n",
        "    def forward(self, embds):\n",
        "        logists = self.layer(embds)\n",
        "        return logists"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "EAJ_FLV64l8I"
      },
      "source": [
        "# TimeStamp Layer Skeleton (SAGE Layer)"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "_TtR2n3D4u5S"
      },
      "source": [
        "class GNNLayer(nn.Module):\n",
        "    \n",
        "    def __init__(self, out_size,num_layers): \n",
        "        super(GNNLayer, self).__init__()\n",
        "\n",
        "        self.out_size = out_size\n",
        "        \n",
        "        dim = (num_layers+2) * self.out_size\n",
        "        self.weight = nn.Parameter(torch.FloatTensor(out_size, dim))\n",
        "\n",
        "        self.init_params()\n",
        "\n",
        "    def init_params(self):\n",
        "        for param in self.parameters():\n",
        "            nn.init.xavier_uniform_(param)\n",
        "\n",
        "    def forward(self, self_feats, aggregate_feats, his_feats, neighs=None):\n",
        "        combined = torch.cat([self_feats, aggregate_feats, his_feats], dim=1)\n",
        "        combined = F.relu(self.weight.mm(combined.t())).t()\n",
        "        return combined"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "NICDPOuv5AU6"
      },
      "source": [
        "# TimeStamp Model Skeleton\\(GraphSAGE Model)\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "dp5yw_sD5Bty"
      },
      "source": [
        "class GNN(nn.Module):\n",
        "  def __init__(self, num_layers,input_size,out_size, adj_lists, device):\n",
        "    super(GNN, self).__init__()\n",
        "\n",
        "    self.num_layers = num_layers\n",
        "    self.input_size = input_size\n",
        "    self.out_size = out_size\n",
        "    self.adj_lists = []\n",
        "    self.device = device\n",
        "    \n",
        "    _ones = torch.ones(adj_lists.shape).to(device)\n",
        "    _zeros = torch.zeros(adj_lists.shape).to(device)\n",
        "\n",
        "    setattr(self, 'layer_adj1', adj_lists)\n",
        "    \n",
        "    for index in range(2, num_layers+1):\n",
        "      cur_adj = torch.pow(adj_lists,index)\n",
        "      cur_adj = torch.where(cur_adj>0, _ones, _zeros)\n",
        "\n",
        "      prev_adj = torch.pow(adj_lists,index-1)\n",
        "      prev_adj = torch.where(prev_adj>0, _ones, _zeros)\n",
        "\n",
        "      layer_adj = cur_adj - prev_adj\n",
        "      setattr(self, 'layer_adj'+str(index), layer_adj)\n",
        "      \n",
        "    self.GNN_Layer = GNNLayer(out_size, num_layers)\n",
        "\n",
        "  def forward(self, nodes_batch,ts):\n",
        "\n",
        "    pre_hidden_embs = self.raw_features\n",
        "    nb = nodes_batch\n",
        "\n",
        "    aggregated_feats = []\n",
        "    for index in range(1, self.num_layers+1):\n",
        "      neigh_feats = self.aggregate(nb, pre_hidden_embs,index)        \n",
        "      aggregated_feats.append(neigh_feats)\n",
        "    \n",
        "    aggregated_feats = torch.cat(aggregated_feats,dim=1)\n",
        "\n",
        "    cur_hidden_embs = self.GNN_Layer(pre_hidden_embs,aggregated_feats, self.pre_latent_feats)\n",
        "\n",
        "    pre_hidden_embs = cur_hidden_embs\n",
        "    \n",
        "    return pre_hidden_embs\n",
        "\n",
        "  def aggregate(self, nodes,pre_hidden_embs,layer):\n",
        "      \n",
        "    embed_matrix = pre_hidden_embs\n",
        "    mask = getattr(self, 'layer_adj'+str(layer))\n",
        "\n",
        "    num_neigh = mask.sum(1, keepdim=True)\n",
        "    _ones = torch.ones(num_neigh.shape).to(self.device)\n",
        "    num_neigh = torch.where(num_neigh>0,num_neigh,_ones)\n",
        "    mask = mask.div(num_neigh)\n",
        "\n",
        "    aggregate_feats = mask.mm(embed_matrix) \n",
        "          \n",
        "    return aggregate_feats"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "7-ubtJTIybQs"
      },
      "source": [
        "# Data Loader"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "pTiG6tUZyeJ0"
      },
      "source": [
        "\n",
        "class DataLoader:\n",
        "\n",
        "  def __init__(self, config,ds,pred_len):\n",
        "    \n",
        "    super(DataLoader, self).__init__()\n",
        "\n",
        "    self.ds = ds\n",
        "    self.dataCenter = DataCenter(config)\n",
        "\n",
        "    if ds == \"PeMSD7\":\n",
        "      train_st = 1\n",
        "      train_en = 22\n",
        "\n",
        "      test_st = 23\n",
        "      test_en = 44 \n",
        "      \n",
        "    elif ds == \"PeMSD8\":\n",
        "      train_st = 1\n",
        "      train_en = 50\n",
        "\n",
        "      test_st = 51\n",
        "      test_en = 62 \n",
        "    \n",
        "    elif ds == \"PeMSD4\" :\n",
        "      train_st = 1\n",
        "      train_en = 47\n",
        "\n",
        "      test_st = 48\n",
        "      test_en = 58\n",
        "  \n",
        "    self.train_st = train_st\n",
        "    self.train_en = train_en\n",
        "    self.test_st = test_st\n",
        "    self.test_en = test_en\n",
        "\n",
        "    self.hr_sample = 12\n",
        "    self.day = 8\n",
        "    self.pred_len = pred_len\n",
        "    \n",
        "  def load_data(self):\n",
        "    print(\"Loading Data...\")\n",
        "    train_data,train_pos = self.dataCenter.load_data(self.ds,self.train_st,self.train_en,self.hr_sample,self.day,self.pred_len)\n",
        "    test_data,test_pos = self.dataCenter.load_data(self.ds,self.test_st,self.test_en,self.hr_sample,self.day,self.pred_len)\n",
        "    adj = self.dataCenter.load_adj(self.ds)\n",
        "    print(\"Data Loaded\")\n",
        "    print(\"Dataset: \", self.ds)\n",
        "    print(\"Total Nodes: \",adj.shape[0])\n",
        "    print(\"Train timestamps: \",len(train_data))\n",
        "    print(\"Test timestamps: \",len(test_data))\n",
        "    print(\"Predicting After: \",self.pred_len*5,\"minutes\")\n",
        "\n",
        "    return train_data,train_pos,test_data,test_pos,adj"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "taPueZtQg39k"
      },
      "source": [
        "# Traffic Model"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "HChymcOgg65S"
      },
      "source": [
        "\n",
        "class TrafficModel:\n",
        "\n",
        "    def __init__(self, train_data,train_pos,test_data,test_pos,adj, \n",
        "                 config, ds, input_size, out_size,GNN_layers,\n",
        "                epochs, device,num_timestamps, pred_len,save_flag,PATH,b_debug,t_debug):\n",
        "      \n",
        "      super(TrafficModel, self).__init__()\n",
        "      \n",
        "      \n",
        "      self.train_data,self.train_pos,self.test_data,self.test_pos,self.adj = train_data,train_pos,test_data,test_pos,adj\n",
        "      self.all_nodes = [i for i in range(self.adj.shape[0])]\n",
        "\n",
        "      self.ds = ds\n",
        "      self.input_size = input_size\n",
        "      self.out_size = out_size\n",
        "      self.GNN_layers = GNN_layers\n",
        "      self.day = input_size \n",
        "      self.device = device\n",
        "      self.epochs = epochs\n",
        "      self.regression = Regression(input_size * num_timestamps, pred_len)\n",
        "      self.num_timestamps = num_timestamps\n",
        "      self.pred_len = pred_len\n",
        "\n",
        "      self.node_bsz = 512\n",
        "      self.PATH = PATH\n",
        "      self.save_flag = save_flag\n",
        "\n",
        "      self.train_data = torch.FloatTensor(self.train_data).to(device)\n",
        "      self.test_data = torch.FloatTensor(self.test_data).to(device)\n",
        "      self.train_pos = torch.FloatTensor(self.train_pos).to(device)\n",
        "      self.test_pos = torch.FloatTensor(self.test_pos).to(device)\n",
        "      self.all_nodes = torch.LongTensor(self.all_nodes).to(device)\n",
        "      self.adj = torch.FloatTensor(self.adj).to(device)\n",
        "\n",
        "      self.b_debug = b_debug\n",
        "      self.t_debug = t_debug\n",
        "      \n",
        "    def run_model(self):\n",
        "\n",
        "\n",
        "      timeStampModel = CombinedGNN(self.input_size,self.out_size,self.adj, \n",
        "      self.device,self.train_data,self.train_pos,self.test_data,self.test_pos,1,self.GNN_layers,self.num_timestamps,self.day)\n",
        "      timeStampModel.to(self.device)\n",
        "\n",
        "      regression = self.regression\n",
        "      regression.to(self.device)\n",
        "\n",
        "      min_RMSE = float(\"Inf\") \n",
        "      min_MAE = float(\"Inf\") \n",
        "      min_MAPE = float(\"Inf\")\n",
        "      best_test = float(\"Inf\")\n",
        "      \n",
        "      lr = 0.001\n",
        "      # if self.ds == \"PeMSD7\":\n",
        "        # lr = 0.001\n",
        "      # elif self.ds == \"PeMSD8\":\n",
        "      #   lr = 0.0001\n",
        "        \n",
        "      train_loss = torch.tensor(0.).to(self.device)  \n",
        "      \n",
        "      for epoch in range(1,epochs):\n",
        "\n",
        "        print(\"Epoch: \",epoch,\" running...\")\n",
        "\n",
        "        tot_timestamp = len(self.train_data)\n",
        "        if self.t_debug:\n",
        "            tot_timestamp = 120\n",
        "        idx = np.random.permutation(tot_timestamp+1-self.num_timestamps)\n",
        "\n",
        "        for data_timestamp in idx:\n",
        "\n",
        "          timeStampModel, regression, train_loss = apply_model(self.all_nodes,timeStampModel, \n",
        "          regression, data_timestamp,self.node_bsz, self.device,\n",
        "          self.pred_len,self.train_data,self.num_timestamps,self.day,train_loss,lr)\n",
        "        \n",
        "          if self.b_debug:\n",
        "            break\n",
        "\n",
        "        train_loss /= len(idx)\n",
        "        # if self.ds==\"PeMSD7\":\n",
        "        if epoch<= 24 and epoch%8==0:\n",
        "          lr *= 0.5\n",
        "        else:\n",
        "          lr = 0.0001\n",
        "\n",
        "\n",
        "        print(\"Train avg loss: \",train_loss)\n",
        "\n",
        "\n",
        "        pred = []\n",
        "        label = []\n",
        "        tot_timestamp = len(self.test_data)\n",
        "        if self.t_debug:\n",
        "            tot_timestamp = 120\n",
        "        idx = np.random.permutation(tot_timestamp+1-self.num_timestamps)\n",
        "        test_loss = torch.tensor(0.).to(self.device)\n",
        "        for data_timestamp in idx:\n",
        "\n",
        "          #window slide\n",
        "          timeStampModel.st = data_timestamp\n",
        "\n",
        "          #test_label\n",
        "          raw_features = self.test_data[timeStampModel.st+self.num_timestamps-1]\n",
        "          test_label = raw_features[:,self.day:]\n",
        "          \n",
        "          #evaluate\n",
        "          temp_predicts,test_loss = evaluate(self.all_nodes,test_label, timeStampModel, regression, \n",
        "                  self.device,test_loss)\n",
        "\n",
        "          label = label + test_label.detach().tolist()\n",
        "          pred = pred + temp_predicts.detach().tolist()\n",
        "            \n",
        "          if self.b_debug:\n",
        "            break\n",
        "\n",
        "        \n",
        "        test_loss /= len(idx)\n",
        "        print(\"Average Test Loss: \",test_loss)\n",
        "\n",
        "        \n",
        "\n",
        "        RMSE = torch.nn.MSELoss()(torch.FloatTensor(pred), torch.FloatTensor(label))\n",
        "        RMSE = torch.sqrt(RMSE).item()\n",
        "        MAE = mean_absolute_error(pred,label)\n",
        "        MAPE = mean_absolute_percentage_error(label,pred)\n",
        "\n",
        "        if test_loss <= best_test:\n",
        "          best_test = test_loss\n",
        "          pred_after = self.pred_len * 5\n",
        "          min_RMSE = RMSE\n",
        "          min_MAE = MAE\n",
        "          min_MAPE = MAPE\n",
        "          if self.save_flag:\n",
        "              torch.save(timeStampModel, self.PATH + \"/\" + self.ds + \"/bestTmodel_\" + str(pred_after) +\"minutes.pth\")\n",
        "              torch.save(regression, self.PATH + \"/\" + self.ds + \"/bestRegression_\" + str(pred_after) +\"minutes.pth\")\n",
        "          \n",
        "        \n",
        "        \n",
        "        print(\"Epoch:\", epoch)\n",
        "        print(\"RMSE: \", RMSE)\n",
        "        print(\"MAE: \", MAE)\n",
        "        print(\"MAPE: \", MAPE)\n",
        "        print(\"===============================================\")\n",
        "\n",
        "        \n",
        "        print(\"Min RMSE: \", min_RMSE)\n",
        "        print(\"Min MAE: \", min_MAE)\n",
        "        print(\"Min MAPE: \", min_MAPE)\n",
        "        \n",
        "        print(\"===============================================\")\n",
        "\n",
        "\n",
        "\n",
        "      return\n",
        "    \n",
        "    def run_Trained_Model(self):\n",
        "      pred_after = self.pred_len * 5\n",
        "      timeStampModel = torch.load(self.PATH + \"/saved_model/\" + self.ds +  \"/bestTmodel_\" + str(pred_after) +\"minutes.pth\")\n",
        "      regression = torch.load(self.PATH + \"/saved_model/\" + self.ds +  \"/bestRegression_\" + str(pred_after) +\"minutes.pth\")\n",
        "      pred = []\n",
        "      label = []\n",
        "      tot_timestamp = len(self.test_data)\n",
        "      idx = np.random.permutation(tot_timestamp+1-self.num_timestamps)\n",
        "      test_loss = torch.tensor(0.).to(self.device)\n",
        "      for data_timestamp in idx:\n",
        "\n",
        "        #window slide\n",
        "        timeStampModel.st = data_timestamp\n",
        "\n",
        "        #test_label\n",
        "        raw_features = self.test_data[timeStampModel.st+self.num_timestamps-1]\n",
        "        test_label = raw_features[:,self.day:]\n",
        "        \n",
        "        #evaluate\n",
        "        temp_predicts,test_loss = evaluate(self.all_nodes,test_label, timeStampModel, regression, \n",
        "                self.device,test_loss)\n",
        "\n",
        "        label = label + test_label.detach().tolist()\n",
        "        pred = pred + temp_predicts.detach().tolist()\n",
        "\n",
        "      \n",
        "      test_loss /= len(idx)\n",
        "      print(\"Average Test Loss: \",test_loss)\n",
        "\n",
        "      \n",
        "      RMSE = torch.nn.MSELoss()(torch.FloatTensor(pred), torch.FloatTensor(label))\n",
        "      RMSE = torch.sqrt(RMSE).item()\n",
        "      MAE = mean_absolute_error(pred,label)\n",
        "      MAPE = mean_absolute_percentage_error(label,pred)\n",
        "      \n",
        "      \n",
        "      print(\"RMSE: \", RMSE)\n",
        "      print(\"MAE: \", MAE)\n",
        "      print(\"MAPE: \", MAPE)\n",
        "      print(\"===============================================\")"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "z4RY89h8UizM"
      },
      "source": [
        ""
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "EquCB81-NUyL"
      },
      "source": [
        "#Different Loss "
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "3hFnoaO-NX36"
      },
      "source": [
        "Window Length, T = 12\n",
        "Y_{T+1} = f(Y_1,...,Y_T)\n",
        "\n",
        "Loss1 =  MSE( Y_{T+1} -  f(Y_1,...,Y_T))\n",
        "\n",
        "\n",
        "Y_{t+1} = g(Y_t) where t= 1,2,...T\n",
        "\n",
        "Loss2 = Sum( MSE( Y_{i+1} -  g(Y_i) ) for all i in (1,2,...T)\n",
        "\n",
        "\n",
        "Loss3 = Sum( MSE( Y_{T+1} -  g(Y_i) ) for all i in (1,2,...T)\n",
        "\n",
        "Model1:"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "f0t4OExU9huJ"
      },
      "source": [
        "# Combined GraphSAGE\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Tkfh2IOB9khw"
      },
      "source": [
        "\n",
        "class CombinedGNN(nn.Module):\n",
        "    def __init__(self,input_size,out_size, adj_lists,\n",
        "                 device,train_data,train_pos,test_data,test_pos,\n",
        "                 st,GNN_layers,num_timestamps,day):\n",
        "        super(CombinedGNN, self).__init__()\n",
        "\n",
        "        self.train_data = train_data\n",
        "        self.train_pos = train_pos\n",
        "        self.test_data = test_data\n",
        "        self.test_pos = test_pos\n",
        "        self.st = st\n",
        "        self.num_timestamps = num_timestamps\n",
        "        self.out_size = out_size\n",
        "        self.tot_nodes = adj_lists.shape[0]\n",
        "        self.device = device\n",
        "        self.adj_lists = adj_lists\n",
        "\n",
        "        \n",
        "        self.day = day\n",
        "  \n",
        "        for timestamp in range(0, self.num_timestamps):\n",
        "\n",
        "            setattr(self, 'his_model' + str(timestamp),GNN(GNN_layers,input_size-1,\n",
        "            out_size-1,adj_lists,device))\n",
        "\n",
        "            setattr(self, 'cur_model' + str(timestamp),GNN(GNN_layers,1,\n",
        "            1,adj_lists,device))\n",
        "\n",
        "                \n",
        "        self.his_weight = nn.Parameter(torch.FloatTensor(out_size-1, self.num_timestamps*out_size-1))\n",
        "        self.cur_weight = nn.Parameter(torch.FloatTensor(1, self.num_timestamps*1))\n",
        "\n",
        "        dim = self.num_timestamps*out_size\n",
        "        self.final_weight = nn.Parameter(torch.FloatTensor(dim,dim))\n",
        "\n",
        "        self.init_params()\n",
        "\n",
        "    def init_params(self):\n",
        "      for param in self.parameters():\n",
        "\n",
        "          if(len(param.shape)>1):\n",
        "            nn.init.xavier_uniform_(param)\n",
        "\n",
        "\n",
        "    def forward(self,nodes_batch,isTrain):\n",
        "\n",
        "        his_timestamp_embds = torch.zeros((nodes_batch.shape[0],self.out_size-1)).to(self.device)\n",
        "        cur_timestamp_embds = torch.zeros((nodes_batch.shape[0],1)).to(self.device)\n",
        "\n",
        "\n",
        "        historical_embds = []\n",
        "        current_embds = []\n",
        "\n",
        "        for timestamp in range(0, self.num_timestamps):\n",
        "\n",
        "\n",
        "            historicalModel = getattr(self, 'his_model' + str(timestamp))\n",
        "            historicalModel.adj_lists = self.adj_lists\n",
        "            setattr(historicalModel, 'timestamp_no', timestamp)\n",
        "            setattr(historicalModel, 'pre_latent_feats', his_timestamp_embds)\n",
        "\n",
        "            if isTrain:\n",
        "                his_raw_features = self.train_data[self.st+timestamp]\n",
        "                his_pos = self.train_pos[self.st+timestamp]\n",
        "            else:\n",
        "                his_raw_features = self.test_data[self.st+timestamp]\n",
        "                his_pos = self.test_pos[self.st+timestamp]\n",
        "\n",
        "            his_raw_features = his_raw_features[:,:self.day-1]\n",
        "            his_pos = his_pos[:,:self.day-1]\n",
        "            setattr(historicalModel,'raw_features',his_raw_features)\n",
        "            \n",
        "\n",
        "            his_timestamp_embds = historicalModel(nodes_batch,timestamp) + his_pos\n",
        "            \n",
        "            historical_embds.append(his_timestamp_embds)\n",
        "            upto_current_timestamp = torch.cat(historical_embds,dim=1)\n",
        "            weight = self.his_weight[:,:(timestamp+1)*(self.out_size-1)]\n",
        "            his_timestamp_embds = F.relu(weight.mm(upto_current_timestamp.t())).t()\n",
        "\n",
        "            currentModel = getattr(self, 'cur_model' + str(timestamp))\n",
        "            currentModel.adj_lists = self.adj_lists\n",
        "            setattr(currentModel, 'timestamp_no', timestamp)\n",
        "            setattr(currentModel, 'pre_latent_feats', cur_timestamp_embds)\n",
        "\n",
        "            if isTrain:\n",
        "                cur_raw_features = self.train_data[self.st+timestamp]\n",
        "                cur_pos = self.train_pos[self.st+timestamp]\n",
        "            else:\n",
        "                cur_raw_features = self.test_data[self.st+timestamp]\n",
        "                cur_pos = self.test_pos[self.st+timestamp]\n",
        "\n",
        "            cur_raw_features = cur_raw_features[:,self.day-1:self.day]\n",
        "            cur_pos = cur_pos[:,self.day-1:self.day]\n",
        "            setattr(currentModel,'raw_features',cur_raw_features)\n",
        "\n",
        "            cur_timestamp_embds = currentModel(nodes_batch,timestamp) + cur_pos\n",
        "\n",
        "            current_embds.append(cur_timestamp_embds)\n",
        "            upto_current_timestamp = torch.cat(current_embds,dim=1)\n",
        "            weight = self.cur_weight[:,:(timestamp+1)*1]\n",
        "            cur_timestamp_embds = F.relu(weight.mm(upto_current_timestamp.t())).t()\n",
        "\n",
        "\n",
        "        his_final_embds = torch.cat(historical_embds,dim=1)\n",
        "        cur_final_embds = torch.cat(current_embds,dim=1)\n",
        "\n",
        "        final_embds = [his_final_embds,cur_final_embds]\n",
        "        final_embds = torch.cat(final_embds,dim=1)\n",
        "        final_embds = F.relu(self.final_weight.mm(final_embds.t()).t())\n",
        "        \n",
        "        return final_embds\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "UbmH7bBv5bql"
      },
      "source": [
        "# Applying Model"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "OE7wtDqj5gBy"
      },
      "source": [
        "\n",
        "def apply_model(train_nodes, CombinedGNN, regression, data_timestamp, \n",
        "                node_batch_sz, device,pred_len,train_data,num_timestamps,day,avg_loss,lr):\n",
        "\n",
        "\n",
        "    models = [CombinedGNN, regression]\n",
        "    params = []\n",
        "    for model in models:\n",
        "      for param in model.parameters():\n",
        "          if param.requires_grad:\n",
        "              params.append(param)\n",
        "\n",
        "\n",
        "\n",
        "    optimizer = torch.optim.Adam(params, lr=lr, weight_decay=0)\n",
        "\n",
        "    optimizer.zero_grad()  # set gradients in zero...\n",
        "    for model in models:\n",
        "      model.zero_grad()  # set gradients in zero\n",
        "\n",
        "    node_batches = math.ceil(len(train_nodes) / node_batch_sz)\n",
        "\n",
        "    loss = torch.tensor(0.).to(device)\n",
        "    #window slide\n",
        "    CombinedGNN.st = data_timestamp\n",
        "    #test_label\n",
        "    raw_features = train_data[CombinedGNN.st+num_timestamps-1]\n",
        "    labels = raw_features[:,day:]\n",
        "    for index in range(node_batches):\n",
        "\n",
        "      nodes_batch = train_nodes[index * node_batch_sz:(index + 1) * node_batch_sz]\n",
        "      nodes_batch = nodes_batch.view(nodes_batch.shape[0],1)\n",
        "      labels_batch = labels[nodes_batch]      \n",
        "      labels_batch = labels_batch.view(len(labels_batch),pred_len)\n",
        "      embs_batch = CombinedGNN(nodes_batch,True)  # Finds embeddings for all the ndoes in nodes_batch\n",
        "\n",
        "      logists = regression(embs_batch)\n",
        "\n",
        "\n",
        "      loss_sup = torch.nn.MSELoss()(logists, labels_batch)\n",
        "\n",
        "      loss_sup /= len(nodes_batch)\n",
        "      loss += loss_sup\n",
        "\n",
        "\n",
        "\n",
        "    avg_loss += loss.item()\n",
        "\n",
        "    loss.backward()\n",
        "    for model in models:\n",
        "      nn.utils.clip_grad_norm_(model.parameters(), 5) \n",
        "    optimizer.step()\n",
        "\n",
        "    optimizer.zero_grad()\n",
        "    for model in models:\n",
        "      model.zero_grad()\n",
        "\n",
        "    return CombinedGNN, regression,avg_loss"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Z9eSCiDc6BQY"
      },
      "source": [
        "# Main Function"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "m7cTd5nRaHC1",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "25d166fd-b7fd-4925-c344-990f8b9a5075"
      },
      "source": [
        "\n",
        "\n",
        "parser = argparse.ArgumentParser(description='pytorch version of Traffic Forecasting GNN')\n",
        "parser.add_argument('-f')  \n",
        "\n",
        "\n",
        "parser.add_argument('--dataset', type=str, default='PeMSD4')\n",
        "parser.add_argument('--input_size', type=int, default=8)\n",
        "parser.add_argument('--out_size', type=int, default=8)\n",
        "parser.add_argument('--epochs', type=int, default=200)\n",
        "parser.add_argument('--seed', type=int, default=824)\n",
        "parser.add_argument('--cuda', action='store_true',help='use CUDA')\n",
        "parser.add_argument('--trained_model', action='store_true')\n",
        "parser.add_argument('--save_model', action='store_true')\n",
        "parser.add_argument('--num_timestamps', type=int, default=12)\n",
        "parser.add_argument('--GNN_layers', type=int, default=4)\n",
        "parser.add_argument('--pred_len', type=int, default=9)\n",
        "args = parser.parse_args()\n",
        "\n",
        "args.cuda = False\n",
        "\n",
        "if torch.cuda.is_available():\n",
        "\tif not args.cuda:\n",
        "\t\tprint(\"WARNING: You have a CUDA device, so you should run with --cuda\")\n",
        "\telse:\n",
        "\t\tdevice_id = torch.cuda.current_device()\n",
        "\t\tprint('using device', device_id, torch.cuda.get_device_name(device_id))\n",
        "\n",
        "device = torch.device(\"cuda\" if args.cuda else \"cpu\")\n",
        "print('DEVICE:', device)\n",
        "\n",
        " \n",
        "if __name__ == '__main__':\n",
        "\n",
        "    print('Traffic Forecasting GNN with Historical and Current Model')\n",
        "\n",
        "    #set user given seed to every random generator\n",
        "    random.seed(args.seed)\n",
        "    np.random.seed(args.seed)\n",
        "    torch.manual_seed(args.seed)\n",
        "    torch.cuda.manual_seed_all(args.seed)\n",
        "\n",
        "    PATH = os.getcwd() + \"/\"\n",
        "    config_file = PATH + \"experiments.conf\"\n",
        " \n",
        "    config = pyhocon.ConfigFactory.parse_file(config_file)\n",
        "    ds = args.dataset\n",
        "    pred_len = args.pred_len\n",
        "    data_loader = DataLoader(config,ds,pred_len)\n",
        "    train_data,train_pos,test_data,test_pos,adj = data_loader.load_data()\n",
        "\n",
        "    num_timestamps = args.num_timestamps \n",
        "    GNN_layers = args.GNN_layers\n",
        "    input_size = args.input_size\n",
        "    out_size = args.input_size\n",
        "    epochs = args.epochs\n",
        "    save_flag = False\n",
        "\n",
        "\n",
        "\n"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "DEVICE: cpu\n",
            "Traffic Forecasting GNN with Historical and Current Model\n",
            "Loading Data...\n",
            "Data Loaded\n",
            "Dataset:  PeMSD4\n",
            "Total Nodes:  307\n",
            "Train timestamps:  11520\n",
            "Test timestamps:  1152\n",
            "Predicting After:  45 minutes\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "jiULj4g1zBCp",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "ec968399-6155-420b-a592-546f34d3f176"
      },
      "source": [
        "b_debug = False\n",
        "t_debug = False\n",
        "\n",
        "hModel = TrafficModel (train_data,train_pos,test_data,test_pos,adj,config, ds, input_size, out_size,GNN_layers,\n",
        "                epochs, device,num_timestamps,pred_len,save_flag,PATH,b_debug,t_debug)\n",
        "if not args.trained_model: #train model and evaluate\n",
        "    hModel.run_model()\n",
        "else:\n",
        "    print(\"Running Trained Model...\")\n",
        "    hModel.run_Trained_Model() #run trained model\n"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Epoch:  1  running...\n",
            "Train avg loss:  tensor(0.1460)\n",
            "Average Test Loss:  tensor(0.1121)\n",
            "Epoch: 1\n",
            "RMSE:  5.8674540519714355\n",
            "MAE:  4.485428726380531\n",
            "MAPE:  8.353322211540146\n",
            "===============================================\n",
            "Min RMSE:  5.8674540519714355\n",
            "Min MAE:  4.485428726380531\n",
            "Min MAPE:  8.353322211540146\n",
            "===============================================\n",
            "Epoch:  2  running...\n",
            "Train avg loss:  tensor(0.0565)\n",
            "Average Test Loss:  tensor(0.0533)\n",
            "Epoch: 2\n",
            "RMSE:  4.045316696166992\n",
            "MAE:  2.3446808339508025\n",
            "MAPE:  4.63821398079797\n",
            "===============================================\n",
            "Min RMSE:  4.045316696166992\n",
            "Min MAE:  2.3446808339508025\n",
            "Min MAPE:  4.63821398079797\n",
            "===============================================\n",
            "Epoch:  3  running...\n",
            "Train avg loss:  tensor(0.0570)\n",
            "Average Test Loss:  tensor(0.0512)\n",
            "Epoch: 3\n",
            "RMSE:  3.963270425796509\n",
            "MAE:  1.940579831297757\n",
            "MAPE:  4.1914123314848\n",
            "===============================================\n",
            "Min RMSE:  3.963270425796509\n",
            "Min MAE:  1.940579831297757\n",
            "Min MAPE:  4.1914123314848\n",
            "===============================================\n",
            "Epoch:  4  running...\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "_KpfK8S50XZt"
      },
      "source": [
        "# Mount Drive"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "mBRrCs3t0Y4E"
      },
      "source": [
        "from google.colab import drive\n",
        "drive.mount('/content/drive')"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "fGzIxxo1FOTb"
      },
      "source": [
        ""
      ],
      "execution_count": null,
      "outputs": []
    }
  ]
}