GAD-EBM_disney.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "065f66a4-5f68-4e14-acc7-0aa91959aa75",
   "metadata": {
    "tags": []
   },
   "source": [
    "# Import"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "e1bcddab-f939-4f06-9009-da4a425cde65",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/roy206/.conda/envs/cent7/2020.11-py38/NWRGAE/lib/python3.8/site-packages/tqdm/auto.py:22: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
      "  from .autonotebook import tqdm as notebook_tqdm\n"
     ]
    }
   ],
   "source": [
    "import numpy as np\n",
    "import scipy.sparse as sp\n",
    "import time\n",
    "import torch\n",
    "import scipy\n",
    "import random\n",
    "import pdb\n",
    "import copy\n",
    "import os\n",
    "import argparse\n",
    "\n",
    "\n",
    "import numpy as np\n",
    "import pickle as pkl\n",
    "import networkx as nx\n",
    "import scipy.sparse as sp\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "# from scipy.sparse.linalg.eigen.arpack import eigsh\n",
    "import sys\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "\n",
    "from torch_geometric.nn import GCNConv,GINConv,SAGEConv,GATConv,PNAConv, GraphSAGE\n",
    "from torch_geometric.utils import to_dense_adj\n",
    "from pygod.utils import load_data\n",
    "from pygod.metrics import eval_roc_auc"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0fc66f2b-ba3a-46b4-9b1e-4ae84087280f",
   "metadata": {},
   "source": [
    "# Helper Function"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "0c1b89a7-31e7-462b-b118-b961675b58f4",
   "metadata": {},
   "outputs": [],
   "source": [
    "def aug_random_edge(input_adj, perturb_percent=0.2, drop_edge = True, add_edge = True, self_loop = True):\n",
    "\n",
    "    aug_adj = copy.deepcopy(input_adj)\n",
    "    nb_nodes = input_adj.shape[0]\n",
    "    edge_index = (input_adj>0).nonzero().t()\n",
    "    \n",
    "    edge_dict = {}\n",
    "    for i in range(nb_nodes):\n",
    "        edge_dict[i] = set()\n",
    "    \n",
    "    for edge in edge_index:\n",
    "        i,j = edge[0],edge[1]\n",
    "        i = i.item()\n",
    "        j = j.item()\n",
    "        edge_dict[i].add(j)\n",
    "        edge_dict[j].add(i)\n",
    "    \n",
    "    if drop_edge: \n",
    "        for i in range(nb_nodes):\n",
    "            d = len(edge_dict[i])\n",
    "            node_list = list(edge_dict[i])\n",
    "            num_edge_to_drop = int(d * perturb_percent)\n",
    "\n",
    "            sampled_nodes = random.sample(node_list, num_edge_to_drop)\n",
    "\n",
    "            for j in sampled_nodes:\n",
    "                aug_adj[i][j] = 0\n",
    "                aug_adj[j][i] = 0\n",
    "\n",
    "            \n",
    "    node_list = [i for i in range(nb_nodes)]\n",
    "    # num_edge_to_add = int(nb_nodes * perturb_percent)\n",
    "    \n",
    "    add_list = []\n",
    "    for i in range(nb_nodes):\n",
    "        d = len(edge_dict[i])\n",
    "        num_edge_to_add = int(d * perturb_percent)\n",
    "        sampled_nodes = random.sample(node_list, num_edge_to_add)\n",
    "        for j in sampled_nodes:\n",
    "            add_list.append((i,j))\n",
    "            \n",
    "    if add_edge:\n",
    "        for i in add_list:\n",
    "            aug_adj[i[0]][i[1]] = 1\n",
    "            aug_adj[i[1]][i[0]] = 1\n",
    "    \n",
    "    if self_loop: \n",
    "        for i in range(nb_nodes):\n",
    "            aug_adj[i][i] = 1\n",
    "            aug_adj[i][i] = 1\n",
    "    \n",
    "    \n",
    "    return aug_adj\n",
    "\n",
    "\n",
    "def _aug_random_edge(nb_nodes, edge_index, perturb_percent=0.2, drop_edge = True, add_edge = True, self_loop = True, use_avg_deg = True):\n",
    "\n",
    "    \n",
    "    total_edges = edge_index.shape[1]\n",
    "    avg_degree = int(total_edges/nb_nodes)\n",
    "    \n",
    "    \n",
    "    edge_dict = {}\n",
    "    for i in range(nb_nodes):\n",
    "        edge_dict[i] = set()\n",
    "    \n",
    "    for edge in edge_index:\n",
    "        i,j = edge[0],edge[1]\n",
    "        i = i.item()\n",
    "        j = j.item()\n",
    "        edge_dict[i].add(j)\n",
    "        edge_dict[j].add(i)\n",
    "    \n",
    "    if drop_edge: \n",
    "        for i in range(nb_nodes):\n",
    "            \n",
    "            d = len(edge_dict[i])\n",
    "            if use_avg_deg:\n",
    "                num_edge_to_drop = avg_degree\n",
    "            else:\n",
    "                num_edge_to_drop = int(d * perturb_percent)\n",
    "\n",
    "            node_list = list(edge_dict[i])\n",
    "            num_edge_to_drop = min(num_edge_to_drop, d)\n",
    "            sampled_nodes = random.sample(node_list, num_edge_to_drop)\n",
    "\n",
    "            for j in sampled_nodes:\n",
    "                edge_dict[i].discard(j)\n",
    "                edge_dict[j].discard(i)\n",
    "            \n",
    "    node_list = [i for i in range(nb_nodes)]\n",
    "    \n",
    "    add_list = []\n",
    "    for i in range(nb_nodes):\n",
    "        \n",
    "        if use_avg_deg:\n",
    "            num_edge_to_add =  avg_degree\n",
    "        else:\n",
    "            d = len(edge_dict[i])\n",
    "            num_edge_to_add = int(d * perturb_percent)\n",
    "        \n",
    "        sampled_nodes = random.sample(node_list, num_edge_to_add)\n",
    "        for j in sampled_nodes:\n",
    "            add_list.append((i,j))\n",
    "            \n",
    "    if add_edge:\n",
    "        for edge in add_list:\n",
    "            u = edge[0]\n",
    "            v = edge[1]\n",
    "            \n",
    "            edge_dict[u].add(v)\n",
    "            edge_dict[v].add(u)\n",
    "    \n",
    "    if self_loop: \n",
    "        for i in range(nb_nodes):\n",
    "            edge_dict[i].add(i)\n",
    "            \n",
    "    updated_edges = set()\n",
    "    for i in range(nb_nodes):\n",
    "        for j in edge_dict[i]:\n",
    "            updated_edges.add((i,j))\n",
    "            updated_edges.add((j,i))\n",
    "    \n",
    "    row = []\n",
    "    col = []\n",
    "    for edge in updated_edges:\n",
    "        u = edge[0]\n",
    "        v = edge[1]\n",
    "        row.append(u)\n",
    "        col.append(v)\n",
    "    \n",
    "    aug_edge_index = [row,col]\n",
    "    aug_edge_index = torch.tensor(aug_edge_index)\n",
    "    \n",
    "    return aug_edge_index\n",
    "\n",
    "\n",
    "def preprocess_features(features):\n",
    "    \"\"\"Row-normalize feature matrix and convert to tuple representation\"\"\"\n",
    "    features = features.squeeze()\n",
    "    rowsum = np.array(features.sum(1))\n",
    "    r_inv = np.power(rowsum, -1).flatten()\n",
    "    r_inv[np.isinf(r_inv)] = 0.\n",
    "    r_mat_inv = sp.diags(r_inv)\n",
    "    features = r_mat_inv.dot(features)\n",
    "    return features\n",
    "    # return features, sparse_to_tuple(features)\n",
    "\n",
    "def normalize_adj(adj):\n",
    "    \"\"\"Symmetrically normalize adjacency matrix.\"\"\"\n",
    "    adj = sp.coo_matrix(adj)\n",
    "    rowsum = np.array(adj.sum(1))\n",
    "    d_inv_sqrt = np.power(rowsum, -0.5).flatten()\n",
    "    d_inv_sqrt[np.isinf(d_inv_sqrt)] = 0.\n",
    "    d_mat_inv_sqrt = sp.diags(d_inv_sqrt)\n",
    "    return adj.dot(d_mat_inv_sqrt).transpose().dot(d_mat_inv_sqrt).tocoo()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ea0866d6-d178-4377-93a3-6d6cc8de664a",
   "metadata": {},
   "source": [
    "# MLP"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "f4bd95b9-2f6a-4331-bb74-bd175d2a8135",
   "metadata": {},
   "outputs": [],
   "source": [
    "class MLP(nn.Module):\n",
    "    def __init__(self, num_layers, input_dim, hidden_dim, output_dim):\n",
    "        \n",
    "        super(MLP, self).__init__()\n",
    "\n",
    "        self.linear_or_not = True  # default is linear model\n",
    "        self.num_layers = num_layers\n",
    "\n",
    "        if num_layers < 1:\n",
    "            raise ValueError(\"number of layers should be positive!\")\n",
    "        elif num_layers == 1:\n",
    "            # Linear model\n",
    "            self.linear = nn.Linear(input_dim, output_dim)\n",
    "        else:\n",
    "            # Multi-layer model\n",
    "            self.linear_or_not = False\n",
    "            self.linears = torch.nn.ModuleList()\n",
    "            self.batch_norms = torch.nn.ModuleList()\n",
    "\n",
    "            self.linears.append(nn.Linear(input_dim, hidden_dim))\n",
    "            for layer in range(num_layers - 2):\n",
    "                self.linears.append(nn.Linear(hidden_dim, hidden_dim))\n",
    "            self.linears.append(nn.Linear(hidden_dim, output_dim))\n",
    "\n",
    "            for layer in range(num_layers - 1):\n",
    "                self.batch_norms.append(nn.BatchNorm1d((hidden_dim)))\n",
    "\n",
    "    def forward(self, x):\n",
    "        if self.linear_or_not:\n",
    "            # If linear model\n",
    "            return self.linear(x)\n",
    "        else:\n",
    "            # If MLP\n",
    "            h = x\n",
    "            for layer in range(self.num_layers - 1):\n",
    "                h = self.linears[layer](h)\n",
    "                \n",
    "                if len(h.shape) > 2:\n",
    "                    h = torch.transpose(h, 0, 1)\n",
    "                    h = torch.transpose(h, 1, 2)\n",
    "                    \n",
    "                \n",
    "                if len(h.shape) > 2:\n",
    "                    h = torch.transpose(h, 1, 2)\n",
    "                    h = torch.transpose(h, 0, 1)\n",
    "\n",
    "                h = F.relu(h)\n",
    "                \n",
    "            return self.linears[self.num_layers - 1](h)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "64572324-7c46-4363-b541-8c9194e22e6e",
   "metadata": {},
   "source": [
    "# GNN Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "6d7d67c6-672e-428d-91d5-dd0dc6d82fe1",
   "metadata": {},
   "outputs": [],
   "source": [
    "class GNN(nn.Module):\n",
    "    def __init__(self, in_dim, out_dim, GNN_name = \"GCN\"):\n",
    "    \n",
    "        super(GNN, self).__init__()    \n",
    "        \n",
    "        self.mlp0 = MLP(3, in_dim, out_dim, out_dim)\n",
    "        \n",
    "        if GNN_name == \"GIN\":\n",
    "            self.linear1 = MLP(4, out_dim, out_dim, out_dim)\n",
    "            self.graphconv1 = GINConv(self.linear1)\n",
    "        elif GNN_name == \"GCN\":\n",
    "            self.graphconv1 = GCNConv(out_dim, out_dim, aggr='mean')\n",
    "        elif GNN_name == \"GAT\":\n",
    "            self.graphconv1 = GATConv(out_dim, out_dim, aggr='mean')\n",
    "        elif GNN_name == \"SAGE\":\n",
    "            self.graphconv1 = SAGEConv(out_dim, out_dim, aggr='mean')\n",
    "            \n",
    "        self.mlp1 = nn.Linear(out_dim,1)\n",
    "        self.sigmoid = nn.Sigmoid()\n",
    "        self.relu = nn.ReLU()\n",
    "    \n",
    "    def forward(self, x, edge_index):\n",
    "        h0 = self.mlp0(x)\n",
    "        h1 = self.graphconv1(h0,edge_index)\n",
    "        h2 = self.mlp1(h1)\n",
    "        h2 = self.relu(h2)\n",
    "        p = torch.exp(h2)\n",
    "        \n",
    "        # h0 = self.mlp0(x)\n",
    "        # h1 = self.mlp1(h0)\n",
    "        # p = torch.exp(h1)\n",
    "        return p\n",
    "        \n",
    "    "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7aaafc73-dc35-4c78-80eb-779cb790fa14",
   "metadata": {},
   "source": [
    "# Execution"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "ceafb622-a294-4cd9-b76a-dd8ae8539be1",
   "metadata": {
    "scrolled": true,
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "----------------------------------------------------------------------------------------------------\n",
      "Namespace(GNN_name='GCN', add_edge=False, dataset='disney', drop_edge=True, f='/home/roy206/.local/share/jupyter/runtime/kernel-fc3e89bb-cf5d-4676-bfed-54100c3c87fa.json', gpu=0, hidden_dim=16, l2_coef=0, lr=0.01, nb_epochs=200, num_neigh=1, perturb_percent=0.05, preprocess_feat=True, save_name='try.pkl', seed=42, self_loop=True, use_avg_deg=False)\n",
      "----------------------------------------------------------------------------------------------------\n",
      "2\n",
      "Loss:  1.0129443407058716 AUC Score:  75.28248587570621\n",
      "Loss:  0.9694161415100098 AUC Score:  35.3813559322034\n",
      "Loss:  0.938948392868042 AUC Score:  31.85028248587571\n",
      "Loss:  0.8653993606567383 AUC Score:  26.906779661016948\n",
      "Loss:  0.7981433868408203 AUC Score:  26.412429378531066\n",
      "Loss:  0.7775387763977051 AUC Score:  26.34180790960452\n",
      "Loss:  0.7239422798156738 AUC Score:  26.271186440677962\n",
      "Loss:  0.6635088920593262 AUC Score:  26.906779661016948\n",
      "Loss:  0.6232538223266602 AUC Score:  26.906779661016948\n",
      "Loss:  0.5720295906066895 AUC Score:  27.54237288135593\n",
      "Loss:  0.5147764682769775 AUC Score:  25.918079096045197\n",
      "Loss:  0.4748563766479492 AUC Score:  27.04802259887006\n",
      "Loss:  0.40635228157043457 AUC Score:  27.04802259887006\n",
      "Loss:  0.38480162620544434 AUC Score:  26.55367231638418\n",
      "Loss:  0.32447075843811035 AUC Score:  26.55367231638418\n",
      "Loss:  0.22003674507141113 AUC Score:  26.9774011299435\n",
      "Loss:  0.1939702033996582 AUC Score:  27.04802259887006\n",
      "Loss:  0.14691948890686035 AUC Score:  27.04802259887006\n",
      "Loss:  0.05830574035644531 AUC Score:  27.04802259887006\n",
      "Loss:  -0.07116389274597168 AUC Score:  27.04802259887006\n",
      "Loss:  -0.11998510360717773 AUC Score:  26.836158192090398\n",
      "Loss:  -0.18028020858764648 AUC Score:  26.55367231638418\n",
      "Loss:  -0.3119211196899414 AUC Score:  26.55367231638418\n",
      "Loss:  -0.46870970726013184 AUC Score:  27.04802259887006\n",
      "Loss:  -0.5536837577819824 AUC Score:  27.04802259887006\n",
      "Loss:  -0.69537353515625 AUC Score:  27.04802259887006\n",
      "Loss:  -0.7887451648712158 AUC Score:  26.271186440677962\n",
      "Loss:  -0.9347105026245117 AUC Score:  26.271186440677962\n",
      "Loss:  -1.0940830707550049 AUC Score:  27.04802259887006\n",
      "Loss:  -1.233670949935913 AUC Score:  27.04802259887006\n",
      "Loss:  -1.45541512966156 AUC Score:  25.635593220338983\n",
      "Loss:  -1.4947799444198608 AUC Score:  26.271186440677962\n",
      "Loss:  -1.6996363401412964 AUC Score:  26.271186440677962\n",
      "Loss:  -2.054671287536621 AUC Score:  27.04802259887006\n",
      "Loss:  -2.2650578022003174 AUC Score:  27.04802259887006\n",
      "Loss:  -2.551859140396118 AUC Score:  26.129943502824858\n",
      "Loss:  -2.81282639503479 AUC Score:  27.04802259887006\n",
      "Loss:  -3.1877527236938477 AUC Score:  26.69491525423729\n",
      "Loss:  -3.3301219940185547 AUC Score:  26.129943502824858\n",
      "Loss:  -3.8915979862213135 AUC Score:  26.129943502824858\n",
      "Loss:  -4.187491416931152 AUC Score:  27.04802259887006\n",
      "Loss:  -4.538972854614258 AUC Score:  27.04802259887006\n",
      "Loss:  -5.246111869812012 AUC Score:  27.04802259887006\n",
      "Loss:  -6.370660781860352 AUC Score:  26.76553672316384\n",
      "Loss:  -6.475280284881592 AUC Score:  26.271186440677962\n",
      "Loss:  -8.114704132080078 AUC Score:  25.98870056497175\n",
      "Loss:  -8.469176292419434 AUC Score:  27.04802259887006\n",
      "Loss:  -9.139179229736328 AUC Score:  27.04802259887006\n",
      "Loss:  -10.74445915222168 AUC Score:  27.04802259887006\n",
      "Loss:  -13.732972145080566 AUC Score:  27.04802259887006\n",
      "Loss:  -14.27975082397461 AUC Score:  26.129943502824858\n",
      "Loss:  -17.45887565612793 AUC Score:  26.271186440677962\n",
      "Loss:  -22.4130916595459 AUC Score:  27.04802259887006\n",
      "Loss:  -25.419429779052734 AUC Score:  27.04802259887006\n",
      "Loss:  -32.02189254760742 AUC Score:  27.04802259887006\n",
      "Loss:  -33.383907318115234 AUC Score:  26.412429378531073\n",
      "Loss:  -47.61634826660156 AUC Score:  26.694915254237284\n",
      "Loss:  -49.516151428222656 AUC Score:  27.04802259887006\n",
      "Loss:  -53.969078063964844 AUC Score:  27.04802259887006\n",
      "Loss:  -66.60535430908203 AUC Score:  27.04802259887006\n",
      "Loss:  -92.87590789794922 AUC Score:  27.04802259887006\n",
      "Loss:  -131.50515747070312 AUC Score:  26.271186440677962\n",
      "Loss:  -162.4398651123047 AUC Score:  27.04802259887006\n",
      "Loss:  -203.77369689941406 AUC Score:  27.04802259887006\n",
      "Loss:  -274.1226501464844 AUC Score:  26.694915254237284\n",
      "Loss:  -322.9549255371094 AUC Score:  27.04802259887006\n",
      "Loss:  -406.8162536621094 AUC Score:  27.04802259887006\n",
      "Loss:  -588.5031127929688 AUC Score:  27.04802259887006\n",
      "Loss:  -470.78533935546875 AUC Score:  26.694915254237284\n",
      "Loss:  -850.4251708984375 AUC Score:  26.694915254237284\n",
      "Loss:  -1154.9154052734375 AUC Score:  27.04802259887006\n",
      "Loss:  -1303.9014892578125 AUC Score:  27.04802259887006\n",
      "Loss:  -1778.9193115234375 AUC Score:  27.04802259887006\n",
      "Loss:  -2945.131103515625 AUC Score:  27.04802259887006\n",
      "Loss:  -2987.5947265625 AUC Score:  26.271186440677962\n",
      "Loss:  -6242.1484375 AUC Score:  27.04802259887006\n",
      "Loss:  -8096.017578125 AUC Score:  27.04802259887006\n",
      "Loss:  -11927.46484375 AUC Score:  27.04802259887006\n",
      "Loss:  -12014.3291015625 AUC Score:  26.412429378531073\n",
      "Loss:  -22049.515625 AUC Score:  27.04802259887006\n",
      "Loss:  -28015.443359375 AUC Score:  27.04802259887006\n",
      "Loss:  -43407.95703125 AUC Score:  27.04802259887006\n",
      "Loss:  -61753.53125 AUC Score:  25.918079096045197\n",
      "Loss:  -27609.25 AUC Score:  26.9774011299435\n",
      "Loss:  -80695.796875 AUC Score:  26.9774011299435\n",
      "Loss:  -209541.1875 AUC Score:  26.906779661016948\n",
      "Loss:  -237413.03125 AUC Score:  26.906779661016948\n",
      "Loss:  -390320.3125 AUC Score:  26.906779661016948\n",
      "Loss:  -502659.0625 AUC Score:  26.906779661016948\n",
      "Loss:  -938378.25 AUC Score:  26.906779661016948\n",
      "Loss:  -1283965.75 AUC Score:  26.906779661016948\n",
      "Loss:  -2252255.0 AUC Score:  27.04802259887006\n",
      "Loss:  -1457679.0 AUC Score:  27.04802259887006\n",
      "Loss:  -4065704.0 AUC Score:  27.04802259887006\n",
      "Loss:  -3894451.5 AUC Score:  27.04802259887006\n",
      "Loss:  -4672488.5 AUC Score:  27.04802259887006\n",
      "Loss:  -7393828.5 AUC Score:  27.04802259887006\n",
      "Loss:  -16213099.0 AUC Score:  27.04802259887006\n",
      "Loss:  -41016616.0 AUC Score:  27.04802259887006\n",
      "Loss:  44811520.0 AUC Score:  25.282485875706218\n",
      "Loss:  -48210336.0 AUC Score:  27.04802259887006\n",
      "Loss:  -26313944.0 AUC Score:  27.04802259887006\n",
      "Loss:  -15963240.0 AUC Score:  27.04802259887006\n",
      "Loss:  -10500565.0 AUC Score:  28.460451977401135\n",
      "Loss:  -7340480.5 AUC Score:  31.285310734463273\n",
      "Loss:  -5382371.5 AUC Score:  31.285310734463273\n",
      "Loss:  -4103458.5 AUC Score:  32.344632768361585\n",
      "Loss:  -3232617.25 AUC Score:  33.40395480225989\n",
      "Loss:  -2619271.0 AUC Score:  33.40395480225989\n",
      "Loss:  -2175105.75 AUC Score:  35.16949152542374\n",
      "Loss:  -1845937.75 AUC Score:  35.87570621468927\n",
      "Loss:  -1597236.25 AUC Score:  35.87570621468927\n",
      "Loss:  -1406283.75 AUC Score:  36.228813559322035\n",
      "Loss:  -1257723.875 AUC Score:  36.5819209039548\n",
      "Loss:  -1140917.25 AUC Score:  37.64124293785311\n",
      "Loss:  -1048330.125 AUC Score:  37.994350282485875\n",
      "Loss:  -974541.6875 AUC Score:  38.70056497175142\n",
      "Loss:  -915570.75 AUC Score:  38.70056497175142\n",
      "Loss:  -868463.9375 AUC Score:  38.70056497175142\n",
      "Loss:  -830995.3125 AUC Score:  39.05367231638419\n",
      "Loss:  -801467.9375 AUC Score:  39.05367231638419\n",
      "Loss:  -778582.0625 AUC Score:  39.05367231638419\n",
      "Loss:  -761323.3125 AUC Score:  39.05367231638419\n",
      "Loss:  -748908.625 AUC Score:  39.05367231638419\n",
      "Loss:  -740717.875 AUC Score:  39.05367231638419\n",
      "Loss:  -736273.25 AUC Score:  39.40677966101695\n",
      "Loss:  -735198.4375 AUC Score:  39.40677966101695\n",
      "Loss:  -737201.75 AUC Score:  39.40677966101695\n",
      "Loss:  -742063.9375 AUC Score:  39.40677966101695\n",
      "Loss:  -749626.625 AUC Score:  39.40677966101695\n",
      "Loss:  -759777.3125 AUC Score:  39.40677966101695\n",
      "Loss:  -772449.5 AUC Score:  39.40677966101695\n",
      "Loss:  -787616.875 AUC Score:  39.40677966101695\n",
      "Loss:  -805290.8125 AUC Score:  39.40677966101695\n",
      "Loss:  -825519.375 AUC Score:  39.40677966101695\n",
      "Loss:  -848387.0625 AUC Score:  39.40677966101695\n",
      "Loss:  -874009.0625 AUC Score:  39.40677966101695\n",
      "Loss:  -902548.25 AUC Score:  39.40677966101695\n",
      "Loss:  -934207.1875 AUC Score:  39.40677966101695\n",
      "Loss:  -969235.875 AUC Score:  39.40677966101695\n",
      "Loss:  -1007934.8125 AUC Score:  39.40677966101695\n",
      "Loss:  -1050670.25 AUC Score:  39.40677966101695\n",
      "Loss:  -1097886.125 AUC Score:  39.40677966101695\n",
      "Loss:  -1150102.875 AUC Score:  39.05367231638419\n",
      "Loss:  -1207955.375 AUC Score:  39.05367231638419\n",
      "Loss:  -1272197.5 AUC Score:  39.05367231638419\n",
      "Loss:  -1343751.625 AUC Score:  39.05367231638419\n",
      "Loss:  -1423716.375 AUC Score:  39.05367231638419\n",
      "Loss:  -1513459.625 AUC Score:  39.05367231638419\n",
      "Loss:  -1614643.0 AUC Score:  39.05367231638419\n",
      "Loss:  -1729330.25 AUC Score:  39.05367231638419\n",
      "Loss:  -1860119.75 AUC Score:  39.05367231638419\n",
      "Loss:  -2010288.875 AUC Score:  38.70056497175142\n",
      "Loss:  -2184045.75 AUC Score:  38.70056497175142\n",
      "Loss:  -2386866.0 AUC Score:  38.70056497175142\n",
      "Loss:  -2625979.0 AUC Score:  38.70056497175142\n",
      "Loss:  -2911103.5 AUC Score:  38.70056497175142\n",
      "Loss:  -3255576.75 AUC Score:  38.34745762711865\n",
      "Loss:  -3678107.75 AUC Score:  37.994350282485875\n",
      "Loss:  -4205719.0 AUC Score:  37.994350282485875\n",
      "Loss:  -4878599.5 AUC Score:  37.64124293785311\n",
      "Loss:  -5758782.5 AUC Score:  37.28813559322034\n",
      "Loss:  -6946533.5 AUC Score:  36.5819209039548\n",
      "Loss:  -8612748.0 AUC Score:  36.228813559322035\n",
      "Loss:  -11069604.0 AUC Score:  35.87570621468927\n",
      "Loss:  -14939157.0 AUC Score:  35.87570621468927\n",
      "Loss:  -21613478.0 AUC Score:  35.5225988700565\n",
      "Loss:  -34743184.0 AUC Score:  34.81638418079096\n",
      "Loss:  -66322832.0 AUC Score:  33.40395480225989\n",
      "Loss:  -170986592.0 AUC Score:  31.285310734463273\n",
      "Loss:  -741666048.0 AUC Score:  31.285310734463273\n",
      "Loss:  -5744830976.0 AUC Score:  27.04802259887006\n",
      "Loss:  -9488601088.0 AUC Score:  27.04802259887006\n",
      "Loss:  -12251148288.0 AUC Score:  27.04802259887006\n",
      "Loss:  -5850822144.0 AUC Score:  27.04802259887006\n",
      "Loss:  -3469513216.0 AUC Score:  31.285310734463273\n",
      "Loss:  -2354590976.0 AUC Score:  31.285310734463273\n",
      "Loss:  -1761553408.0 AUC Score:  33.40395480225989\n",
      "Loss:  -1425746048.0 AUC Score:  33.40395480225989\n",
      "Loss:  -1236676864.0 AUC Score:  35.5225988700565\n",
      "Loss:  -1144851200.0 AUC Score:  35.87570621468927\n",
      "Loss:  -1130519680.0 AUC Score:  36.228813559322035\n",
      "Loss:  -1193678464.0 AUC Score:  36.5819209039548\n",
      "Loss:  -1355056512.0 AUC Score:  37.28813559322034\n",
      "Loss:  -1669235968.0 AUC Score:  37.64124293785311\n",
      "Loss:  -2263991808.0 AUC Score:  37.28813559322034\n",
      "Loss:  -3459106048.0 AUC Score:  36.5819209039548\n",
      "Loss:  -6200788992.0 AUC Score:  36.228813559322035\n",
      "Loss:  -14264604672.0 AUC Score:  35.87570621468927\n",
      "Loss:  -52354469888.0 AUC Score:  34.81638418079096\n",
      "Loss:  -452189945856.0 AUC Score:  31.285310734463273\n",
      "Loss:  -6802426036224.0 AUC Score:  27.04802259887006\n",
      "Loss:  -74164346028032.0 AUC Score:  27.04802259887006\n",
      "Loss:  11770962903040.0 AUC Score:  26.271186440677962\n",
      "Loss:  36335404.0 AUC Score:  26.271186440677962\n",
      "Loss:  -17494.99609375 AUC Score:  26.129943502824858\n",
      "Loss:  -95.56375885009766 AUC Score:  27.47175141242938\n",
      "Loss:  -0.5479435920715332 AUC Score:  27.401129943502827\n",
      "Loss:  130.54052734375 AUC Score:  73.44632768361582\n",
      "Loss:  31117.021484375 AUC Score:  72.17514124293785\n",
      "Maximum AUC:  75.28248587570621\n",
      "Required Time:  1.9168684482574463\n"
     ]
    }
   ],
   "source": [
    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "\n",
    "parser = argparse.ArgumentParser(\"GAD-EBM: Graph Anomaly Detection via Energy-based Model\")\n",
    "\n",
    "parser.add_argument('-f')\n",
    "parser.add_argument('--dataset',          type=str,           default=\"disney\",         help='dataset name')\n",
    "parser.add_argument('--perturb_percent',  type=float,         default=0.05,             help='perturb percent')\n",
    "parser.add_argument('--seed',             type=int,           default=42,               help='seed')\n",
    "parser.add_argument('--nb_epochs',        type=int,           default=200,              help='total epochs')\n",
    "parser.add_argument('--hidden_dim',       type=int,           default=16,               help='hidden dimension')\n",
    "parser.add_argument('--lr',               type=float,         default=0.01,             help='learning rate')\n",
    "parser.add_argument('--l2_coef',          type=float,         default=0,                help='regularization coefficeint')\n",
    "parser.add_argument('--gpu',              type=int,           default=0,                help='gpu')\n",
    "parser.add_argument('--save_name',        type=str,           default='try.pkl',        help='save ckpt name')\n",
    "parser.add_argument('--drop_edge',        type=bool,          default=True,             help='drop edge flag to produce state space neighbor')\n",
    "parser.add_argument('--add_edge',         type=bool,          default=False,            help='add edge flag to produce state space neighbor')\n",
    "parser.add_argument('--self_loop',        type=bool,          default=True,             help='self loop in state space neighbor')\n",
    "parser.add_argument('--preprocess_feat',  type=bool,          default=True,             help='preprocess feature flag')\n",
    "parser.add_argument('--use_avg_deg',      type=bool,          default=False,            help='use avg_deg to add/drop edge in state space neighbor')\n",
    "parser.add_argument('--GNN_name',         type=str,           default=\"GCN\",            help='gnn encoder')\n",
    "parser.add_argument('--num_neigh',        type=int,           default=1,                help='state space graph number of neighbors')\n",
    "\n",
    "args = parser.parse_args()\n",
    "\n",
    "print('-' * 100)\n",
    "print(args)\n",
    "print('-' * 100)\n",
    "\n",
    "dataset_str = args.dataset\n",
    "perturb_percent = args.perturb_percent\n",
    "os.environ[\"CUDA_VISIBLE_DEVICES\"] = str(args.gpu) \n",
    "seed = args.seed\n",
    "random.seed(seed)\n",
    "np.random.seed(seed)\n",
    "torch.manual_seed(seed)\n",
    "torch.cuda.manual_seed(seed)\n",
    "\n",
    "nb_epochs = args.nb_epochs\n",
    "lr = args.lr\n",
    "l2_coef = args.l2_coef\n",
    "hidden_dim = args.hidden_dim\n",
    "k = args.num_neigh\n",
    "\n",
    "\n",
    "\n",
    "data = load_data(dataset_str)\n",
    "edge_index = data.edge_index\n",
    "\n",
    "\n",
    "adj = to_dense_adj(edge_index).squeeze()\n",
    "features = data.x\n",
    "labels = data.y\n",
    "y = labels.bool()\n",
    "\n",
    "anomaly_nodes = np.nonzero(y)\n",
    "\n",
    "nb_nodes = features.shape[0]  # total node\n",
    "input_dim = features.shape[1]   # total features\n",
    "\n",
    "\n",
    "\n",
    "model = GNN(input_dim, hidden_dim, args.GNN_name)\n",
    "optimiser = torch.optim.Adam(model.parameters(), lr=lr, weight_decay=l2_coef)\n",
    "\n",
    "\n",
    "if args.preprocess_feat:\n",
    "    features = preprocess_features(features)\n",
    "\n",
    "    \n",
    "# if args.preprocess_adj:\n",
    "#     adj = normalize_adj(adj)\n",
    "\n",
    "    \n",
    "\n",
    "\n",
    "st = time.time()\n",
    "mx_auc = 0\n",
    "\n",
    "aug_edge_indexes = []\n",
    "\n",
    "for i in range(k):\n",
    "    aug_edge_index = _aug_random_edge(nb_nodes,edge_index, perturb_percent=perturb_percent,drop_edge=args.drop_edge,\n",
    "                                  add_edge=args.add_edge,self_loop=args.self_loop, use_avg_deg = args.use_avg_deg) # add/drop perturb percentage of edges\n",
    "    \n",
    "    \n",
    "    aug_edge_index = aug_edge_index.to(device)\n",
    "    \n",
    "    aug_edge_indexes.append(aug_edge_index)\n",
    "    \n",
    "    \n",
    "model = model.to(device)\n",
    "features = torch.FloatTensor(features[np.newaxis])\n",
    "features = features.to(device)\n",
    "edge_index = edge_index.to(device)\n",
    "\n",
    "losses = []\n",
    "    \n",
    "for epoch in range(nb_epochs):\n",
    "\n",
    "    model.train()\n",
    "    optimiser.zero_grad()\n",
    "    \n",
    "    p_data = model(features, edge_index)\n",
    "    loss = 0\n",
    "    \n",
    "    for i in range(k):\n",
    "        \n",
    "        aug_edge_index = aug_edge_indexes[i]\n",
    "        \n",
    "        shuf_fts = features\n",
    "        idx = np.random.permutation(nb_nodes)\n",
    "        shuf_fts = features[:, idx, :]\n",
    "        \n",
    "        \n",
    "        p_neigh = model(shuf_fts, aug_edge_index)\n",
    "    \n",
    "        c_theta_j1 = p_neigh/p_data\n",
    "        c_theta_j2 = p_data/p_neigh\n",
    "            \n",
    "        j1 = (c_theta_j1**2 + 2 * c_theta_j1).mean()\n",
    "        j2 = (2 * c_theta_j2).mean()\n",
    "        \n",
    "        \n",
    "        \n",
    "        neigh_loss = j1 - j2\n",
    "        neigh_loss = neigh_loss.mean()\n",
    "        loss = loss + neigh_loss\n",
    "    \n",
    "    loss = loss / k\n",
    "    \n",
    "    losses.append(loss)\n",
    "    \n",
    "\n",
    "    \n",
    "    logits = p_data.squeeze().detach().cpu() \n",
    "    auc_score = eval_roc_auc(y.numpy(), logits.numpy()) * 100\n",
    "    \n",
    "    print(\"Loss: \", loss.item(), \"AUC Score: \", auc_score)\n",
    "    mx_auc = max(mx_auc, auc_score)\n",
    "    \n",
    "    loss.backward()\n",
    "    optimiser.step()\n",
    "\n",
    "en = time.time()\n",
    "\n",
    "\n",
    "print(\"Maximum AUC: \", mx_auc)\n",
    "print(\"Required Time: \", en-st)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a7cf09fc-956d-48ad-9cb3-198e8d56e235",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "NWRGAE",
   "language": "python",
   "name": "nwrgae"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}