GAD-EBM_enron.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "065f66a4-5f68-4e14-acc7-0aa91959aa75",
   "metadata": {
    "tags": []
   },
   "source": [
    "# Import"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "e1bcddab-f939-4f06-9009-da4a425cde65",
   "metadata": {},
   "outputs": [],
   "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",
    "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.metric import eval_roc_auc\n",
    "from pygod.generator import gen_contextual_outlier, gen_structural_outlier"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0fc66f2b-ba3a-46b4-9b1e-4ae84087280f",
   "metadata": {},
   "source": [
    "# Helper Function"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "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",
    "        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 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": 8,
   "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",
    "                h = self.batch_norms[layer](h)\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",
    "                # h = F.relu(self.linears[layer](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": 9,
   "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(1, in_dim, out_dim, out_dim)\n",
    "        \n",
    "        if GNN_name == \"GIN\":\n",
    "            self.linear1 = MLP(1, 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='sum')\n",
    "        elif GNN_name == \"GAT\":\n",
    "            self.graphconv1 = GATConv(out_dim, out_dim, aggr='sum')\n",
    "        elif GNN_name == \"SAGE\":\n",
    "            self.graphconv1 = SAGEConv(out_dim, out_dim, aggr='sum')\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",
    "        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": null,
   "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='enron', drop_edge=True, f='/home/roy206/.local/share/jupyter/runtime/kernel-9959e154-abd2-481d-97fc-85cb71b2a51e.json', gpu=0, hidden_dim=16, l2_coef=0.1, lr=0.005, nb_epochs=500, num_neigh=1, perturb_percent=0.2, preprocess_feat=False, save_name='try.pkl', seed=42, self_loop=True)\n",
      "----------------------------------------------------------------------------------------------------\n",
      "Loss:  1.0158624649047852 AUC Score:  80.06948551153164\n",
      "Loss:  0.9975061416625977 AUC Score:  80.07318154937909\n",
      "Loss:  1.0095503330230713 AUC Score:  80.03991720875221\n",
      "Loss:  0.9879255294799805 AUC Score:  80.01330573625074\n",
      "Loss:  1.0092661380767822 AUC Score:  80.0510053222945\n",
      "Loss:  1.0075030326843262 AUC Score:  80.04878769958604\n",
      "Loss:  0.9965310096740723 AUC Score:  80.10053222945002\n",
      "Loss:  0.9977507591247559 AUC Score:  80.15745121230043\n",
      "Loss:  0.996845006942749 AUC Score:  80.1434062684802\n",
      "Loss:  0.9974918365478516 AUC Score:  80.1323181549379\n",
      "Loss:  0.9999473094940186 AUC Score:  80.13157894736842\n",
      "Loss:  1.0014808177947998 AUC Score:  80.18554109994086\n",
      "Loss:  1.000932216644287 AUC Score:  80.19293317563572\n",
      "Loss:  0.9992256164550781 AUC Score:  80.24393849793022\n",
      "Loss:  0.9999203681945801 AUC Score:  80.31712004730929\n",
      "Loss:  1.000619649887085 AUC Score:  80.37182140745122\n",
      "Loss:  1.0027955770492554 AUC Score:  80.33486102897693\n",
      "Loss:  0.9980025291442871 AUC Score:  80.39473684210526\n",
      "Loss:  0.9987034797668457 AUC Score:  80.58471318746304\n",
      "Loss:  0.9996404647827148 AUC Score:  80.43243642814902\n",
      "Loss:  1.0014795064926147 AUC Score:  80.47161442933177\n",
      "Loss:  1.000112771987915 AUC Score:  79.67696629213484\n",
      "Loss:  1.0018584728240967 AUC Score:  79.8115020697812\n",
      "Loss:  1.0001866817474365 AUC Score:  79.70505617977528\n",
      "Loss:  0.9996883869171143 AUC Score:  79.8573329390893\n",
      "Loss:  0.9991137981414795 AUC Score:  79.81963335304553\n",
      "Loss:  0.9992117881774902 AUC Score:  79.6274393849793\n",
      "Loss:  0.9997658729553223 AUC Score:  79.61487285629805\n",
      "Loss:  1.000315546989441 AUC Score:  79.48625073920758\n",
      "Loss:  0.9998493194580078 AUC Score:  79.63483146067416\n",
      "Loss:  0.9992716312408447 AUC Score:  79.60674157303372\n",
      "Loss:  0.9997868537902832 AUC Score:  79.67474866942638\n",
      "Loss:  1.0000770092010498 AUC Score:  79.7397989355411\n",
      "Loss:  0.9997832775115967 AUC Score:  79.7146658781786\n",
      "Loss:  1.0001651048660278 AUC Score:  79.68214074512123\n",
      "Loss:  1.0000618696212769 AUC Score:  79.68066232998225\n",
      "Loss:  0.9998233318328857 AUC Score:  79.668095801301\n",
      "Loss:  0.9999938011169434 AUC Score:  79.50694855115317\n",
      "Loss:  0.9999332427978516 AUC Score:  79.67327025428742\n",
      "Loss:  0.9999582767486572 AUC Score:  79.67400946185688\n",
      "Loss:  0.9996838569641113 AUC Score:  79.59417504435245\n",
      "Loss:  1.000018835067749 AUC Score:  79.56682436428149\n",
      "Loss:  0.9999425411224365 AUC Score:  79.49142519219396\n",
      "Loss:  0.9999427795410156 AUC Score:  78.93849793021882\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=\"enron\",          help='dataset name')\n",
    "parser.add_argument('--perturb_percent',  type=float,         default=0.2,              help='perturb percent')\n",
    "parser.add_argument('--seed',             type=int,           default=42,               help='seed')\n",
    "parser.add_argument('--nb_epochs',        type=int,           default=500,              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.005,            help='learning rate')\n",
    "parser.add_argument('--l2_coef',          type=float,         default=0.1,              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 for state space neighbor')\n",
    "parser.add_argument('--add_edge',         type=bool,          default=False,            help='add edge flag for state space neighbor')\n",
    "parser.add_argument('--self_loop',        type=bool,          default=True,             help='self loop for state space neighbor')\n",
    "parser.add_argument('--preprocess_feat',  type=bool,          default=False,            help='preprocess feature flag')\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",
    "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",
    "for i in range(nb_nodes):\n",
    "    adj[i][i] = 1\n",
    "edge_index = adj.nonzero().t()\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",
    "    \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",
    "st = time.time()\n",
    "mx_auc = 0\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_adj = aug_random_edge(adj, perturb_percent=perturb_percent,drop_edge=args.drop_edge,\n",
    "                                  add_edge=args.add_edge,self_loop=args.self_loop) # add/drop perturb percentage of edges\n",
    "        aug_edge_index = aug_adj.nonzero().t()\n",
    "        aug_edge_index = aug_edge_index.to(device)\n",
    "        \n",
    "        shuf_fts = features\n",
    "        idx = np.random.permutation(nb_nodes)\n",
    "        shuf_fts = features[:, idx, :]\n",
    "        \n",
    "        p_neigh = model(shuf_fts, aug_edge_index)\n",
    "        \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",
    "        neigh_loss = j1 - j2\n",
    "        neigh_loss = neigh_loss.mean()\n",
    "        loss = loss + neigh_loss\n",
    "    \n",
    "    loss = loss / k\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)"
   ]
  }
 ],
 "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
}