experiments/GIB_node_model.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "from collections import OrderedDict\n",
    "from copy import deepcopy\n",
    "import itertools\n",
    "import matplotlib.pylab as plt\n",
    "import numpy as np\n",
    "import os.path as osp\n",
    "import pickle\n",
    "import scipy.sparse as sp\n",
    "from scipy.sparse.csgraph import connected_components\n",
    "import sklearn\n",
    "from sklearn.manifold import TSNE\n",
    "import torch\n",
    "from torch.nn import Parameter, Linear\n",
    "import torch.nn.functional as F\n",
    "from torch.distributions.normal import Normal\n",
    "\n",
    "import sys, os\n",
    "sys.path.append(os.path.join(os.path.dirname(\"__file__\"), '..', '..'))\n",
    "from torch_scatter import scatter_add\n",
    "import torch_geometric.transforms as T\n",
    "from torch_geometric.utils import add_remaining_self_loops, remove_self_loops, add_self_loops, softmax, degree, to_undirected\n",
    "from torch_geometric.nn.inits import glorot, zeros\n",
    "from torch_geometric.nn.conv import MessagePassing\n",
    "from torch_geometric.datasets import Planetoid\n",
    "from torch_geometric.data import Data\n",
    "from GIB.pytorch_net.net import reparameterize, Mixture_Gaussian_reparam\n",
    "from GIB.pytorch_net.util import sample, to_cpu_recur, to_np_array, to_Variable, record_data, make_dir, remove_duplicates, update_dict, get_list_elements, to_string, filter_filename\n",
    "from GIB.util import get_reparam_num_neurons, sample_lognormal, scatter_sample, uniform_prior, compose_log, edge_index_2_csr, COLOR_LIST, LINESTYLE_LIST, process_data_for_nettack, parse_filename, add_distant_neighbors\n",
    "from GIB.DeepRobust.deeprobust.graph.targeted_attack import Nettack"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## GCNConv:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "class GCNConv(MessagePassing):\n",
    "    r\"\"\"The graph convolutional operator from the `\"Semi-supervised\n",
    "    Classification with Graph Convolutional Networks\"\n",
    "    <https://arxiv.org/abs/1609.02907>`_ paper\n",
    "\n",
    "    .. math::\n",
    "        \\mathbf{X}^{\\prime} = \\mathbf{\\hat{D}}^{-1/2} \\mathbf{\\hat{A}}\n",
    "        \\mathbf{\\hat{D}}^{-1/2} \\mathbf{X} \\mathbf{\\Theta},\n",
    "\n",
    "    where :math:`\\mathbf{\\hat{A}} = \\mathbf{A} + \\mathbf{I}` denotes the\n",
    "    adjacency matrix with inserted self-loops and\n",
    "    :math:`\\hat{D}_{ii} = \\sum_{j=0} \\hat{A}_{ij}` its diagonal degree matrix.\n",
    "\n",
    "    Args:\n",
    "        in_channels (int): Size of each input sample.\n",
    "        out_channels (int): Size of each output sample.\n",
    "        improved (bool, optional): If set to :obj:`True`, the layer computes\n",
    "            :math:`\\mathbf{\\hat{A}}` as :math:`\\mathbf{A} + 2\\mathbf{I}`.\n",
    "            (default: :obj:`False`)\n",
    "        cached (bool, optional): If set to :obj:`True`, the layer will cache\n",
    "            the computation of :math:`\\mathbf{\\hat{D}}^{-1/2} \\mathbf{\\hat{A}}\n",
    "            \\mathbf{\\hat{D}}^{-1/2}` on first execution, and will use the\n",
    "            cached version for further executions.\n",
    "            This parameter should only be set to :obj:`True` in transductive\n",
    "            learning scenarios. (default: :obj:`False`)\n",
    "        bias (bool, optional): If set to :obj:`False`, the layer will not learn\n",
    "            an additive bias. (default: :obj:`True`)\n",
    "        normalize (bool, optional): Whether to add self-loops and apply\n",
    "            symmetric normalization. (default: :obj:`True`)\n",
    "        **kwargs (optional): Additional arguments of\n",
    "            :class:`torch_geometric.nn.conv.MessagePassing`.\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self, in_channels, out_channels, improved=False, cached=False,\n",
    "                 bias=True, normalize=True,\n",
    "                 reparam_mode=None, prior_mode=None, sample_size=1, val_use_mean=True,\n",
    "                 **kwargs):\n",
    "        super(GCNConv, self).__init__(aggr='add', **kwargs)\n",
    "\n",
    "        self.reparam_mode = None if reparam_mode == \"None\" else reparam_mode\n",
    "        self.prior_mode = prior_mode\n",
    "        self.val_use_mean = val_use_mean\n",
    "        self.sample_size = sample_size\n",
    "        self.in_channels = in_channels\n",
    "        self.out_channels = out_channels\n",
    "        self.out_neurons = get_reparam_num_neurons(out_channels, self.reparam_mode)\n",
    "        self.improved = improved\n",
    "        self.cached = cached\n",
    "        self.normalize = normalize\n",
    "\n",
    "        self.weight = Parameter(torch.Tensor(in_channels, self.out_neurons))\n",
    "\n",
    "        if bias:\n",
    "            self.bias = Parameter(torch.Tensor(self.out_neurons))\n",
    "        else:\n",
    "            self.register_parameter('bias', None)\n",
    "\n",
    "        if self.reparam_mode is not None:\n",
    "            if self.prior_mode.startswith(\"mixGau\"):\n",
    "                n_components = eval(self.prior_mode.split(\"-\")[1])\n",
    "                self.feature_prior = Mixture_Gaussian_reparam(is_reparam=False, Z_size=self.out_channels, n_components=n_components)\n",
    "\n",
    "        self.reset_parameters()\n",
    "\n",
    "    def reset_parameters(self):\n",
    "        glorot(self.weight)\n",
    "        zeros(self.bias)\n",
    "        self.cached_result = None\n",
    "        self.cached_num_edges = None\n",
    "    \n",
    "\n",
    "    def set_cache(self, cached):\n",
    "        self.cached = cached\n",
    "\n",
    "\n",
    "    def to_device(self, device):\n",
    "        self.to(device)\n",
    "        if self.cached and self.cached_result is not None:\n",
    "            edge_index, norm = self.cached_result\n",
    "            self.cached_result = edge_index.to(device), norm.to(device)\n",
    "        return self\n",
    "\n",
    "\n",
    "    @staticmethod\n",
    "    def norm(edge_index, num_nodes, edge_weight=None, improved=False,\n",
    "             dtype=None):\n",
    "        if edge_weight is None:\n",
    "            edge_weight = torch.ones((edge_index.size(1), ), dtype=dtype,\n",
    "                                     device=edge_index.device)\n",
    "\n",
    "        fill_value = 1 if not improved else 2\n",
    "        edge_index, edge_weight = add_remaining_self_loops(\n",
    "            edge_index, edge_weight, fill_value, num_nodes)\n",
    "\n",
    "        row, col = edge_index\n",
    "        deg = scatter_add(edge_weight, row, dim=0, dim_size=num_nodes)\n",
    "        deg_inv_sqrt = deg.pow(-0.5)\n",
    "        deg_inv_sqrt[deg_inv_sqrt == float('inf')] = 0\n",
    "\n",
    "        return edge_index, deg_inv_sqrt[row] * edge_weight * deg_inv_sqrt[col]\n",
    "\n",
    "    def forward(self, x, edge_index, edge_weight=None):\n",
    "        \"\"\"\"\"\"\n",
    "        x = torch.matmul(x, self.weight)\n",
    "\n",
    "        if self.cached and self.cached_result is not None:\n",
    "            if edge_index.size(1) != self.cached_num_edges:\n",
    "                raise RuntimeError(\n",
    "                    'Cached {} number of edges, but found {}. Please '\n",
    "                    'disable the caching behavior of this layer by removing '\n",
    "                    'the `cached=True` argument in its constructor.'.format(\n",
    "                        self.cached_num_edges, edge_index.size(1)))\n",
    "\n",
    "        if not self.cached or self.cached_result is None:\n",
    "            self.cached_num_edges = edge_index.size(1)\n",
    "            if self.normalize:\n",
    "                edge_index, norm = self.norm(edge_index, x.size(\n",
    "                    self.node_dim), edge_weight, self.improved, x.dtype)\n",
    "            else:\n",
    "                norm = edge_weight\n",
    "            self.cached_result = edge_index, norm\n",
    "\n",
    "        edge_index, norm = self.cached_result\n",
    "        out = self.propagate(edge_index, x=x, norm=norm)\n",
    "\n",
    "        if self.reparam_mode is not None:\n",
    "            # Reparameterize:\n",
    "            self.dist, _ = reparameterize(model=None, input=out, \n",
    "                                          mode=self.reparam_mode, \n",
    "                                          size=self.out_channels\n",
    "                                         )  # [B, Z]\n",
    "            Z = sample(self.dist, self.sample_size)  # [S, B, Z]\n",
    "\n",
    "            if self.prior_mode == \"Gaussian\":\n",
    "                self.feature_prior = Normal(loc=torch.zeros(x.size(0), self.out_channels).to(x.device),\n",
    "                                            scale=torch.ones(x.size(0), self.out_channels).to(x.device),\n",
    "                                           )  # [B, Z]\n",
    "\n",
    "            # Calculate prior loss:\n",
    "            if self.reparam_mode == \"diag\" and self.prior_mode == \"Gaussian\":\n",
    "                ixz = torch.distributions.kl.kl_divergence(self.dist, self.feature_prior).sum(-1)\n",
    "            else:\n",
    "                Z_logit = self.dist.log_prob(Z).sum(-1) if self.reparam_mode.startswith(\"diag\") else self.dist.log_prob(Z)  # [S, B]\n",
    "                prior_logit = self.feature_prior.log_prob(Z).sum(-1)  # [S, B]\n",
    "                # upper bound of I(X; Z):\n",
    "                ixz = (Z_logit - prior_logit).mean(0)  # [B]\n",
    "\n",
    "            self.Z_std = to_np_array(Z.std((0, 1)).mean())\n",
    "            if self.val_use_mean is False or self.training:\n",
    "                out = Z.mean(0)  # [B, Z]\n",
    "            else:\n",
    "                out = out[:, :self.out_channels]  # [B, Z]\n",
    "        else:\n",
    "            ixz = torch.zeros(x.size(0)).to(x.device)  # [B]\n",
    "\n",
    "        structure_kl_loss = torch.zeros([]).to(x.device)\n",
    "        return out, ixz, structure_kl_loss\n",
    "            \n",
    "\n",
    "    def message(self, x_j, norm):\n",
    "        return norm.view(-1, 1) * x_j if norm is not None else x_j\n",
    "\n",
    "    def update(self, aggr_out):\n",
    "        if self.bias is not None:\n",
    "            aggr_out = aggr_out + self.bias\n",
    "        return aggr_out\n",
    "\n",
    "    def __repr__(self):\n",
    "        return '{}({}, {})'.format(self.__class__.__name__, self.in_channels,\n",
    "                                   self.out_channels)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## GATConv:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "class GATConv(MessagePassing):\n",
    "    r\"\"\"The graph attentional operator from the `\"Graph Attention Networks\"\n",
    "    <https://arxiv.org/abs/1710.10903>`_ paper\n",
    "\n",
    "    .. math::\n",
    "        \\mathbf{x}^{\\prime}_i = \\alpha_{i,i}\\mathbf{\\Theta}\\mathbf{x}_{i} +\n",
    "        \\sum_{j \\in \\mathcal{N}(i)} \\alpha_{i,j}\\mathbf{\\Theta}\\mathbf{x}_{j},\n",
    "\n",
    "    where the attention coefficients :math:`\\alpha_{i,j}` are computed as\n",
    "\n",
    "    .. math::\n",
    "        \\alpha_{i,j} =\n",
    "        \\frac{\n",
    "        \\exp\\left(\\mathrm{LeakyReLU}\\left(\\mathbf{a}^{\\top}\n",
    "        [\\mathbf{\\Theta}\\mathbf{x}_i \\, \\Vert \\, \\mathbf{\\Theta}\\mathbf{x}_j]\n",
    "        \\right)\\right)}\n",
    "        {\\sum_{k \\in \\mathcal{N}(i) \\cup \\{ i \\}}\n",
    "        \\exp\\left(\\mathrm{LeakyReLU}\\left(\\mathbf{a}^{\\top}\n",
    "        [\\mathbf{\\Theta}\\mathbf{x}_i \\, \\Vert \\, \\mathbf{\\Theta}\\mathbf{x}_k]\n",
    "        \\right)\\right)}.\n",
    "\n",
    "    Args:\n",
    "        in_channels (int): Size of each input sample.\n",
    "        out_channels (int): Size of each output sample.\n",
    "        heads (int, optional): Number of multi-head-attentions.\n",
    "            (default: :obj:`1`)\n",
    "        concat (bool, optional): If set to :obj:`False`, the multi-head\n",
    "            attentions are averaged instead of concatenated.\n",
    "            (default: :obj:`True`)\n",
    "        negative_slope (float, optional): LeakyReLU angle of the negative\n",
    "            slope. (default: :obj:`0.2`)\n",
    "        struct_dropout_mode (tuple, optional): Choose from: None, (\"standard\", prob), (\"info\", ${MODE}),\n",
    "            where ${MODE} chooses from \"subset\", \"lognormal\", \"loguniform\".\n",
    "        bias (bool, optional): If set to :obj:`False`, the layer will not learn\n",
    "            an additive bias. (default: :obj:`True`)\n",
    "        **kwargs (optional): Additional arguments of\n",
    "            :class:`torch_geometric.nn.conv.MessagePassing`.\n",
    "    \"\"\"\n",
    "    def __init__(self, in_channels, out_channels, heads=1, concat=True,\n",
    "                 negative_slope=0.2, reparam_mode=None, prior_mode=None,\n",
    "                 struct_dropout_mode=None, sample_size=1,\n",
    "                 val_use_mean=True,\n",
    "                 bias=True,\n",
    "                 **kwargs):\n",
    "        super(GATConv, self).__init__(aggr='add', **kwargs)\n",
    "\n",
    "        self.in_channels = in_channels\n",
    "        self.out_channels = out_channels\n",
    "        self.heads = heads\n",
    "        self.concat = concat\n",
    "        self.negative_slope = negative_slope\n",
    "        self.reparam_mode = reparam_mode if reparam_mode != \"None\" else None\n",
    "        self.prior_mode = prior_mode\n",
    "        self.out_neurons = get_reparam_num_neurons(out_channels, self.reparam_mode)\n",
    "        self.struct_dropout_mode = struct_dropout_mode\n",
    "        self.sample_size = sample_size\n",
    "        self.val_use_mean = val_use_mean\n",
    "\n",
    "        self.weight = Parameter(torch.Tensor(in_channels,\n",
    "                                             heads * self.out_neurons))\n",
    "        self.att = Parameter(torch.Tensor(1, heads, 2 * self.out_neurons))\n",
    "\n",
    "        if bias and concat:\n",
    "            self.bias = Parameter(torch.Tensor(heads * self.out_neurons))\n",
    "        elif bias and not concat:\n",
    "            self.bias = Parameter(torch.Tensor(self.out_neurons))\n",
    "        else:\n",
    "            self.register_parameter('bias', None)\n",
    "            \n",
    "        if self.reparam_mode is not None:\n",
    "            if self.prior_mode.startswith(\"mixGau\"):\n",
    "                n_components = eval(self.prior_mode.split(\"-\")[1])\n",
    "                self.feature_prior = Mixture_Gaussian_reparam(is_reparam=False, Z_size=self.out_channels, n_components=n_components)\n",
    "\n",
    "        self.skip_editing_edge_index = struct_dropout_mode[0] == 'DNsampling'\n",
    "        self.reset_parameters()\n",
    "\n",
    "    def reset_parameters(self):\n",
    "        glorot(self.weight)\n",
    "        glorot(self.att)\n",
    "        zeros(self.bias)\n",
    "\n",
    "    def forward(self, x, edge_index, size=None):\n",
    "        \"\"\"\"\"\"\n",
    "        if size is None and torch.is_tensor(x) and not self.skip_editing_edge_index:\n",
    "            edge_index, _ = remove_self_loops(edge_index)\n",
    "            edge_index, _ = add_self_loops(edge_index,\n",
    "                                           num_nodes=x.size(self.node_dim))\n",
    "\n",
    "        if torch.is_tensor(x):\n",
    "            x = torch.matmul(x, self.weight)\n",
    "        else:\n",
    "            x = (None if x[0] is None else torch.matmul(x[0], self.weight),\n",
    "                 None if x[1] is None else torch.matmul(x[1], self.weight))\n",
    "\n",
    "        out = self.propagate(edge_index, size=size, x=x)\n",
    "\n",
    "        if self.reparam_mode is not None:\n",
    "            # Reparameterize:\n",
    "            out = out.view(-1, self.out_neurons)\n",
    "            self.dist, _ = reparameterize(model=None, input=out,\n",
    "                                          mode=self.reparam_mode,\n",
    "                                          size=self.out_channels,\n",
    "                                         )  # dist: [B * head, Z]\n",
    "            Z_core = sample(self.dist, self.sample_size)  # [S, B * head, Z]\n",
    "            Z = Z_core.view(self.sample_size, -1, self.heads * self.out_channels)  # [S, B, head * Z]\n",
    "\n",
    "            if self.prior_mode == \"Gaussian\":\n",
    "                self.feature_prior = Normal(loc=torch.zeros(out.size(0), self.out_channels).to(x.device),\n",
    "                                            scale=torch.ones(out.size(0), self.out_channels).to(x.device),\n",
    "                                           )  # feature_prior: [B * head, Z]\n",
    "\n",
    "            if self.reparam_mode == \"diag\" and self.prior_mode == \"Gaussian\":\n",
    "                ixz = torch.distributions.kl.kl_divergence(self.dist, self.feature_prior).sum(-1).view(-1, self.heads).mean(-1)\n",
    "            else:\n",
    "                Z_logit = self.dist.log_prob(Z_core).sum(-1) if self.reparam_mode.startswith(\"diag\") else self.dist.log_prob(Z_core)  # [S, B * head]\n",
    "                prior_logit = self.feature_prior.log_prob(Z_core).sum(-1)  # [S, B * head]\n",
    "                # upper bound of I(X; Z):\n",
    "                ixz = (Z_logit - prior_logit).mean(0).view(-1, self.heads).mean(-1)  # [B]\n",
    "\n",
    "            self.Z_std = to_np_array(Z.std((0, 1)).mean())\n",
    "            if self.val_use_mean is False or self.training:\n",
    "                out = Z.mean(0)\n",
    "            else:\n",
    "                out = out[:, :self.out_channels].contiguous().view(-1, self.heads * self.out_channels)\n",
    "        else:\n",
    "            ixz = torch.zeros(x.size(0)).to(x.device)\n",
    "\n",
    "        if \"Nsampling\" in self.struct_dropout_mode[0]:\n",
    "            if 'categorical' in self.struct_dropout_mode[1]:\n",
    "                structure_kl_loss = torch.sum(self.alpha*torch.log((self.alpha+1e-16)/self.prior))\n",
    "            elif 'Bernoulli' in self.struct_dropout_mode[1]:\n",
    "                posterior = torch.distributions.bernoulli.Bernoulli(self.alpha)\n",
    "                prior = torch.distributions.bernoulli.Bernoulli(self.prior) \n",
    "                structure_kl_loss = torch.distributions.kl.kl_divergence(posterior, prior).sum(-1).mean()\n",
    "            else:\n",
    "                raise Exception(\"I think this belongs to the diff subset sampling that is not implemented\")\n",
    "        else:\n",
    "            structure_kl_loss = torch.zeros([]).to(x.device)\n",
    "\n",
    "        return out, ixz, structure_kl_loss\n",
    "\n",
    "    def message(self, edge_index_i, x_i, x_j, size_i):\n",
    "        # Compute attention coefficients.\n",
    "        x_j = x_j.view(-1, self.heads, self.out_neurons)  # [N_edge, heads, out_channels]\n",
    "        if x_i is None:\n",
    "            alpha = (x_j * self.att[:, :, self.out_neurons:]).sum(dim=-1)\n",
    "        else:\n",
    "            x_i = x_i.view(-1, self.heads, self.out_neurons)\n",
    "            alpha = (torch.cat([x_i, x_j], dim=-1) * self.att).sum(dim=-1)  # [N_edge, heads]\n",
    "\n",
    "        alpha = F.leaky_relu(alpha, self.negative_slope)\n",
    "\n",
    "        # Sample attention coefficients stochastically.\n",
    "        if self.struct_dropout_mode[0] == \"None\":\n",
    "            alpha = softmax(alpha, edge_index_i, num_nodes=size_i)\n",
    "        else:\n",
    "            if self.struct_dropout_mode[0] == \"standard\":\n",
    "                alpha = softmax(alpha, edge_index_i, num_nodes=size_i)\n",
    "                prob_dropout = self.struct_dropout_mode[1]\n",
    "                alpha = F.dropout(alpha, p=prob_dropout, training=self.training)\n",
    "            elif self.struct_dropout_mode[0] == \"identity\":\n",
    "                alpha = torch.ones_like(alpha)\n",
    "                alpha = softmax(alpha, edge_index_i, num_nodes=size_i)\n",
    "            elif self.struct_dropout_mode[0] == \"info\":\n",
    "                mode = self.struct_dropout_mode[1]\n",
    "                if mode == \"lognormal\":\n",
    "                    max_alpha = self.struct_dropout_mode[2] if len(self.struct_dropout_mode) > 2 else 0.7\n",
    "                    alpha = 0.001 + max_alpha * alpha\n",
    "                    self.kl = -torch.log(alpha/(max_alpha + 0.001))\n",
    "                    sigma0 = 1. if self.training else 0.\n",
    "                    alpha = sample_lognormal(mean=torch.zeros_like(alpha), sigma=alpha, sigma0=sigma0)\n",
    "                else:\n",
    "                    raise Exception(\"Mode {} for the InfoDropout is invalid!\".format(mode))\n",
    "            elif \"Nsampling\" in self.struct_dropout_mode[0]:\n",
    "                neighbor_sampling_mode = self.struct_dropout_mode[1]\n",
    "                if 'categorical' in neighbor_sampling_mode:\n",
    "                    alpha = softmax(alpha, edge_index_i, num_nodes=size_i)\n",
    "                    self.alpha = alpha\n",
    "                    self.prior = uniform_prior(edge_index_i)\n",
    "                    if self.val_use_mean is False or self.training:\n",
    "                        temperature = self.struct_dropout_mode[2]\n",
    "                        sample_neighbor_size = self.struct_dropout_mode[3]\n",
    "                        if neighbor_sampling_mode == 'categorical':\n",
    "                            alpha = scatter_sample(alpha, edge_index_i, temperature, size_i)\n",
    "                        elif 'multi-categorical' in neighbor_sampling_mode:\n",
    "                            alphas = []\n",
    "                            for _ in range(sample_neighbor_size): #! this can be improved by parallel sampling\n",
    "                                alphas.append(scatter_sample(alpha, edge_index_i, temperature, size_i))\n",
    "                            alphas = torch.stack(alphas, dim=0)\n",
    "                            if 'sum' in neighbor_sampling_mode:\n",
    "                                alpha = alphas.sum(dim=0)\n",
    "                            elif 'max' in neighbor_sampling_mode:\n",
    "                                alpha, _ = torch.max(alphas, dim=0)\n",
    "                            else:\n",
    "                                raise\n",
    "                        else:\n",
    "                            raise\n",
    "                elif neighbor_sampling_mode == 'Bernoulli':\n",
    "                    if self.struct_dropout_mode[4] == 'norm':\n",
    "                        alpha_normalization = torch.ones_like(alpha)\n",
    "                        alpha_normalization = softmax(alpha_normalization, edge_index_i, num_nodes=size_i)\n",
    "                    alpha = torch.clamp(torch.sigmoid(alpha), 0.01, 0.99)\n",
    "                    self.alpha = alpha\n",
    "                    self.prior = (torch.ones_like(self.alpha)*self.struct_dropout_mode[3]).to(alpha.device)\n",
    "                    if not self.val_use_mean or self.training:\n",
    "                        temperature = self.struct_dropout_mode[2]\n",
    "                        alpha = torch.distributions.relaxed_bernoulli.RelaxedBernoulli(torch.Tensor([temperature]).to(alpha.device),\n",
    "                            probs=alpha).rsample()\n",
    "                    if self.struct_dropout_mode[4] == 'norm':\n",
    "                        alpha = alpha*alpha_normalization\n",
    "                else:\n",
    "                    raise\n",
    "            else:\n",
    "                raise\n",
    "\n",
    "        return x_j * alpha.view(-1, self.heads, 1)\n",
    "\n",
    "    def update(self, aggr_out):\n",
    "        if self.concat is True:\n",
    "            aggr_out = aggr_out.view(-1, self.heads * self.out_neurons)\n",
    "        else:\n",
    "            aggr_out = aggr_out.mean(dim=1)\n",
    "\n",
    "        if self.bias is not None:\n",
    "            aggr_out = aggr_out + self.bias\n",
    "        return aggr_out\n",
    "\n",
    "\n",
    "    def to_device(self, device):\n",
    "        self.to(device)\n",
    "        return self\n",
    "\n",
    "    def __repr__(self):\n",
    "        return '{}({}, {}, heads={})'.format(self.__class__.__name__,\n",
    "                                             self.in_channels,\n",
    "                                             self.out_channels, self.heads)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## GNN:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "class GNN(torch.nn.Module):\n",
    "    def __init__(\n",
    "        self,\n",
    "        model_type,\n",
    "        num_features,\n",
    "        num_classes,\n",
    "        reparam_mode,\n",
    "        prior_mode,\n",
    "        latent_size,\n",
    "        sample_size=1,\n",
    "        num_layers=2,\n",
    "        struct_dropout_mode=(\"standard\", 0.6),\n",
    "        dropout=True,\n",
    "        with_relu=True,\n",
    "        val_use_mean=True,\n",
    "        reparam_all_layers=True,\n",
    "        normalize=True,\n",
    "        is_cuda=False,\n",
    "    ):\n",
    "        \"\"\"Class implementing a general GNN, which can realize GAT, GIB-GAT, GCN.\n",
    "        \n",
    "        Args:\n",
    "            model_type:   name of the base model. Choose from \"GAT\", \"GCN\".\n",
    "            num_features: number of features of the data.x.\n",
    "            num_classes:  number of classes for data.y.\n",
    "            reparam_mode: reparameterization mode for XIB. Choose from \"diag\" and \"full\". Default \"diag\" that parameterizes the mean and diagonal element of the Gaussian\n",
    "            prior_mode:   distribution type for the prior. Choose from \"Gaussian\" or \"mixGau-{Number}\", where {Number} is the number of components for mixture of Gaussian.\n",
    "            latent_size:  latent size for each layer of GNN. If model_type=\"GAT\", the true latent size is int(latent_size/2)\n",
    "            sample_size=1:how many Z to sample for each feature X.\n",
    "            num_layers=2: number of layers for the GNN\n",
    "            struct_dropout_mode: Mode for how the structural representation is generated. Only effective for model_type==\"GAT\"\n",
    "                          Choose from (\"Nsampling\", 'multi-categorical-sum', 0.1, 3) (here 0.1 is temperature, k=3 is the number of sampled edges with replacement), \n",
    "                          (\"DNsampling\", 'multi-categorical-sum', 0.1, 3, 2) (similar as above, with the local dependence range T=2) \n",
    "                          (\"standard\", 0.6) (standard dropout used on the attention weights in GAT)\n",
    "            dropout:      whether to use dropout on features.\n",
    "            with_relu:    whether to use nonlinearity for GCN.\n",
    "            val_use_mean: Whether during evaluation use the parameter value instead of sampling. If True, during evaluation,\n",
    "                          XIB will use mean for prediction, and AIB will use the parameter of the categorical distribution for prediction.\n",
    "            reparam_all_layers: Which layers to use XIB, e.g. (1,2,4). Default (-2,), meaning the second last layer. If True, use XIB for all layers.\n",
    "            normalize:    whether to normalize for GCN (only effective for GCN)\n",
    "            is_cuda:      whether to use CUDA, and if so, which GPU to use. Choose from False, True, \"CUDA:{GPU_ID}\", where {GPU_ID} is the ID for the CUDA.\n",
    "        \"\"\"\n",
    "        super(GNN, self).__init__()\n",
    "        self.model_type = model_type\n",
    "        self.num_features = num_features\n",
    "        self.num_classes = num_classes\n",
    "        self.normalize = normalize\n",
    "        self.reparam_mode = reparam_mode\n",
    "        self.prior_mode = prior_mode\n",
    "        self.struct_dropout_mode = struct_dropout_mode\n",
    "        self.dropout = dropout\n",
    "        self.latent_size = latent_size\n",
    "        self.sample_size = sample_size\n",
    "        self.num_layers = num_layers\n",
    "        self.with_relu = with_relu\n",
    "        self.val_use_mean = val_use_mean\n",
    "        self.reparam_all_layers = reparam_all_layers\n",
    "        self.is_cuda = is_cuda\n",
    "        self.device = torch.device(self.is_cuda if isinstance(self.is_cuda, str) else \"cuda\" if self.is_cuda else \"cpu\")\n",
    "\n",
    "        self.init()\n",
    "\n",
    "\n",
    "    def init(self):\n",
    "        \"\"\"Initialize the layers for the GNN.\"\"\"\n",
    "        self.reparam_layers = []\n",
    "        if self.model_type == \"GCN\":\n",
    "            for i in range(self.num_layers):\n",
    "                if self.reparam_all_layers is True:\n",
    "                    is_reparam = True\n",
    "                elif isinstance(self.reparam_all_layers, tuple):\n",
    "                    reparam_all_layers = tuple([kk + self.num_layers if kk < 0 else kk for kk in self.reparam_all_layers])\n",
    "                    is_reparam = i in reparam_all_layers\n",
    "                else:\n",
    "                    raise\n",
    "                if is_reparam:\n",
    "                    self.reparam_layers.append(i)\n",
    "                setattr(self, \"conv{}\".format(i + 1),\n",
    "                        GCNConv(self.num_features if i == 0 else self.latent_size,\n",
    "                                self.latent_size if i != self.num_layers - 1 else self.num_classes,\n",
    "                                cached=True,\n",
    "                                reparam_mode=self.reparam_mode if is_reparam else None,\n",
    "                                prior_mode=self.prior_mode if is_reparam else None,\n",
    "                                sample_size=self.sample_size,\n",
    "                                bias=True if self.with_relu else False,\n",
    "                                val_use_mean=self.val_use_mean,\n",
    "                                normalize=self.normalize,\n",
    "                ))\n",
    "            # self.conv1 = ChebConv(self.num_features, 16, K=2)\n",
    "            # self.conv2 = ChebConv(16, self.num_features, K=2)\n",
    "\n",
    "        elif self.model_type == \"GAT\":\n",
    "            latent_size = int(self.latent_size / 2)  # Under the default setting, latent_size = 8\n",
    "            for i in range(self.num_layers):\n",
    "                if i == 0:\n",
    "                    input_size = self.num_features\n",
    "                else:\n",
    "                    if self.struct_dropout_mode[0] == 'DNsampling' or (self.struct_dropout_mode[0] == 'standard' and len(self.struct_dropout_mode) == 3):\n",
    "                        input_size = latent_size * 8 * 2\n",
    "                    else:\n",
    "                        input_size = latent_size * 8\n",
    "                if self.reparam_all_layers is True:\n",
    "                    is_reparam = True\n",
    "                elif isinstance(self.reparam_all_layers, tuple):\n",
    "                    reparam_all_layers = tuple([kk + self.num_layers if kk < 0 else kk for kk in self.reparam_all_layers])\n",
    "                    is_reparam = i in reparam_all_layers\n",
    "                else:\n",
    "                    raise\n",
    "                if is_reparam:\n",
    "                    self.reparam_layers.append(i)\n",
    "                setattr(self, \"conv{}\".format(i + 1), GATConv(\n",
    "                    input_size,\n",
    "                    latent_size if i != self.num_layers - 1 else self.num_classes,\n",
    "                    heads=8 if i != self.num_layers - 1 else 1, concat=True,\n",
    "                    reparam_mode=self.reparam_mode if is_reparam else None,\n",
    "                    prior_mode=self.prior_mode if is_reparam else None,\n",
    "                    val_use_mean=self.val_use_mean,\n",
    "                    struct_dropout_mode=self.struct_dropout_mode,\n",
    "                    sample_size=self.sample_size,\n",
    "                ))\n",
    "                if self.struct_dropout_mode[0] == 'DNsampling' or (self.struct_dropout_mode[0] == 'standard' and len(self.struct_dropout_mode) == 3):\n",
    "                    setattr(self, \"conv{}_1\".format(i + 1), GATConv(\n",
    "                        input_size,\n",
    "                        latent_size if i != self.num_layers - 1 else self.num_classes,\n",
    "                        heads=8 if i != self.num_layers - 1 else 1, concat=True,\n",
    "                        reparam_mode=self.reparam_mode if is_reparam else None,\n",
    "                        prior_mode=self.prior_mode if is_reparam  else None,\n",
    "                        val_use_mean=self.val_use_mean,\n",
    "                        struct_dropout_mode=self.struct_dropout_mode,\n",
    "                        sample_size=self.sample_size,\n",
    "                    ))\n",
    "            # On the Pubmed dataset, use heads=8 in conv2.\n",
    "        \n",
    "        else:\n",
    "            raise Exception(\"Model_type {} is not valid!\".format(self.model_type))\n",
    "\n",
    "        self.reparam_layers = sorted(self.reparam_layers)\n",
    "   \n",
    "        if self.model_type == \"GCN\":\n",
    "            if self.with_relu:\n",
    "                reg_params = [getattr(self, \"conv{}\".format(i+1)).parameters() for i in range(self.num_layers - 1)]\n",
    "                self.reg_params = itertools.chain(*reg_params)\n",
    "                self.non_reg_params = getattr(self, \"conv{}\".format(self.num_layers)).parameters()\n",
    "            else:\n",
    "                self.reg_params = OrderedDict()\n",
    "                self.non_reg_params = self.parameters()\n",
    "        else:\n",
    "            self.reg_params = self.parameters()\n",
    "            self.non_reg_params = OrderedDict()\n",
    "        self.to(self.device)\n",
    "\n",
    "\n",
    "    def set_cache(self, cached):\n",
    "        \"\"\"Set cache for GCN.\"\"\"\n",
    "        for i in range(self.num_layers):\n",
    "            if hasattr(getattr(self, \"conv{}\".format(i+1)), \"set_cache\"):\n",
    "                getattr(self, \"conv{}\".format(i+1)).set_cache(cached)\n",
    "\n",
    "\n",
    "    def to_device(self, device):\n",
    "        \"\"\"Send all the layers to the specified device.\"\"\"\n",
    "        for i in range(self.num_layers):\n",
    "            getattr(self, \"conv{}\".format(i+1)).to_device(device)\n",
    "        self.to(device)\n",
    "        return self\n",
    "\n",
    "\n",
    "    def forward(self, data, record_Z=False, isplot=False):\n",
    "        \"\"\"Main forward function.\n",
    "        \n",
    "        Args:\n",
    "            data: the pytorch-geometric data class.\n",
    "            record_Z: whether to record the standard deviation for the representation Z.\n",
    "            isplot:   whether to plot.\n",
    "        \n",
    "        Returns:\n",
    "            x: output\n",
    "            reg_info: other information or metrics.\n",
    "        \"\"\"\n",
    "        reg_info = {}\n",
    "        if self.model_type == \"GCN\":\n",
    "            x, edge_index, edge_weight = data.x, data.edge_index, data.edge_attr\n",
    "            for i in range(self.num_layers - 1):\n",
    "                layer = getattr(self, \"conv{}\".format(i + 1))\n",
    "                x, ixz, structure_kl_loss = layer(x, edge_index, edge_weight)\n",
    "                # Record:\n",
    "                record_data(reg_info, [ixz, structure_kl_loss], [\"ixz_list\", \"structure_kl_list\"])\n",
    "                if layer.reparam_mode is not None:\n",
    "                    record_data(reg_info, [layer.Z_std], [\"Z_std\"])\n",
    "                if record_Z:\n",
    "                    record_data(reg_info, [to_np_array(x)], [\"Z_{}\".format(i)], nolist=True)\n",
    "                if self.with_relu:\n",
    "                    x = F.relu(x)\n",
    "                    self.plot(x, data.y, titles=\"Layer{}\".format(i + 1), isplot=isplot)\n",
    "                    if self.dropout is True:\n",
    "                        x = F.dropout(x, training=self.training)\n",
    "            layer = getattr(self, \"conv{}\".format(self.num_layers))\n",
    "            x, ixz, structure_kl_loss = layer(x, edge_index, edge_weight)\n",
    "            # Record:\n",
    "            record_data(reg_info, [ixz, structure_kl_loss], [\"ixz_list\", \"structure_kl_list\"])\n",
    "            if layer.reparam_mode is not None:\n",
    "                record_data(reg_info, [layer.Z_std], [\"Z_std\"])\n",
    "            if record_Z:\n",
    "                record_data(reg_info, [to_np_array(x)], [\"Z_{}\".format(self.num_layers - 1)], nolist=True)\n",
    "            self.plot(x, data.y, titles=\"Layer{}\".format(self.num_layers), isplot=isplot)\n",
    "\n",
    "        elif self.model_type == \"GAT\":\n",
    "            x = F.dropout(data.x, p=0.6, training=self.training)\n",
    "\n",
    "            for i in range(self.num_layers - 1):\n",
    "                if self.struct_dropout_mode[0] == 'DNsampling' or (self.struct_dropout_mode[0] == 'standard' and len(self.struct_dropout_mode) == 3):\n",
    "                    x_1, ixz_1, structure_kl_loss_1 = getattr(self, \"conv{}_1\".format(i + 1))(x, data.multi_edge_index)\n",
    "                layer = getattr(self, \"conv{}\".format(i + 1))\n",
    "                x, ixz, structure_kl_loss = layer(x, data.edge_index)\n",
    "                # Record:\n",
    "                record_data(reg_info, [ixz, structure_kl_loss], [\"ixz_list\", \"structure_kl_list\"])\n",
    "                if layer.reparam_mode is not None:\n",
    "                    record_data(reg_info, [layer.Z_std], [\"Z_std\"])\n",
    "                if record_Z:\n",
    "                    record_data(reg_info, [to_np_array(x)], [\"Z_{}\".format(i)], nolist=True)\n",
    "                # Multi-hop:\n",
    "                if self.struct_dropout_mode[0] == 'DNsampling' or (self.struct_dropout_mode[0] == 'standard' and len(self.struct_dropout_mode) == 3):\n",
    "                    x = torch.cat([x, x_1], dim=-1)\n",
    "                    record_data(reg_info, [ixz_1, structure_kl_loss_1], [\"ixz_DN_list\", \"structure_kl_DN_list\"])\n",
    "                x = F.elu(x)\n",
    "                self.plot(x, data.y, titles=\"Layer{}\".format(i + 1), isplot=isplot)\n",
    "                x = F.dropout(x, p=0.6, training=self.training)\n",
    "\n",
    "            if self.struct_dropout_mode[0] == 'DNsampling' or (self.struct_dropout_mode[0] == 'standard' and len(self.struct_dropout_mode) == 3):\n",
    "                x_1, ixz_1, structure_kl_loss_1 = getattr(self, \"conv{}_1\".format(self.num_layers))(x, data.multi_edge_index)\n",
    "            layer = getattr(self, \"conv{}\".format(self.num_layers))\n",
    "            x, ixz, structure_kl_loss = layer(x, data.edge_index)\n",
    "            # Record:\n",
    "            record_data(reg_info, [ixz, structure_kl_loss], [\"ixz_list\", \"structure_kl_list\"])\n",
    "            if layer.reparam_mode is not None:\n",
    "                record_data(reg_info, [layer.Z_std], [\"Z_std\"])\n",
    "            if record_Z:\n",
    "                record_data(reg_info, [to_np_array(x)], [\"Z_{}\".format(self.num_layers - 1)], nolist=True)\n",
    "            # Multi-hop:\n",
    "            if self.struct_dropout_mode[0] == 'DNsampling' or (self.struct_dropout_mode[0] == 'standard' and len(self.struct_dropout_mode) == 3):\n",
    "                x = x + x_1\n",
    "                record_data(reg_info, [ixz_1, structure_kl_loss_1], [\"ixz_DN_list\", \"structure_kl_DN_list\"])\n",
    "            self.plot(x, data.y, titles=\"Layer{}\".format(self.num_layers), isplot=isplot)\n",
    "\n",
    "        return x, reg_info\n",
    "\n",
    "\n",
    "    def compute_metrics_fun(self, data, metrics, mask=None, mask_id=None):\n",
    "        \"\"\"Compute metrics for measuring clustering performance.\n",
    "        Choices: \"Silu\", \"CH\", \"DB\".\n",
    "        \"\"\"\n",
    "        _, info_dict = self(data, record_Z=True)\n",
    "        y = to_np_array(data.y)\n",
    "        info_metrics = {}\n",
    "        if mask is not None:\n",
    "            mask = to_np_array(mask)\n",
    "            mask_id += \"_\"\n",
    "        else:\n",
    "            mask_id = \"\"\n",
    "        for k in range(self.num_layers):\n",
    "            if mask is not None:\n",
    "                Z_i = info_dict[\"Z_{}\".format(k)][mask]\n",
    "                y_i = y[mask]\n",
    "            else:\n",
    "                Z_i = info_dict[\"Z_{}\".format(k)]\n",
    "                y_i = y\n",
    "            for metric in metrics:\n",
    "                if metric == \"Silu\":\n",
    "                    score = sklearn.metrics.silhouette_score(Z_i, y_i, metric='euclidean')\n",
    "                elif metric == \"DB\":\n",
    "                    score = sklearn.metrics.davies_bouldin_score(Z_i, y_i)\n",
    "                elif metric == \"CH\":\n",
    "                    score = sklearn.metrics.calinski_harabasz_score(Z_i, y_i)\n",
    "                info_metrics[\"{}{}_{}\".format(mask_id, metric, k)] = score\n",
    "        return info_metrics\n",
    "\n",
    "\n",
    "    def plot(self, Z_list, y, titles=None, isplot=False):\n",
    "        \"\"\"Plot the intermediate representation Z.\"\"\"\n",
    "        if isplot:\n",
    "            if not isinstance(Z_list, list):\n",
    "                Z_list = [Z_list]\n",
    "            if titles is not None and not isinstance(titles, list):\n",
    "                titles = [titles]\n",
    "            length = len(Z_list)\n",
    "            tsne = TSNE(n_components=2, init='pca')\n",
    "            plt.figure(figsize=(8 * length, 6))\n",
    "            for k, Z in enumerate(Z_list):\n",
    "                plt.subplot(1, length, k+1)\n",
    "                for i in range(self.num_classes):\n",
    "                    idx = y == i\n",
    "                    Z_embed_i = tsne.fit_transform(to_np_array(Z[idx]))\n",
    "                    plt.plot(Z_embed_i[:, 0], Z_embed_i[:, 1], \".\", c=COLOR_LIST[i], label=str(i))\n",
    "                if titles is not None:\n",
    "                    plt.title(titles[k], fontsize=18)\n",
    "            plt.legend(bbox_to_anchor=[1,1])\n",
    "            plt.show()\n",
    "\n",
    "\n",
    "    @property\n",
    "    def model_dict(self):\n",
    "        \"\"\"Record model_dict for saving.\"\"\"\n",
    "        model_dict = {}\n",
    "        model_dict[\"model_type\"] = self.model_type\n",
    "        model_dict[\"num_features\"] = self.num_features\n",
    "        model_dict[\"num_classes\"] = self.num_classes\n",
    "        model_dict[\"normalize\"] = self.normalize\n",
    "        model_dict[\"reparam_mode\"] = self.reparam_mode\n",
    "        model_dict[\"prior_mode\"] = self.prior_mode\n",
    "        model_dict[\"struct_dropout_mode\"] = self.struct_dropout_mode\n",
    "        model_dict[\"dropout\"] = self.dropout\n",
    "        model_dict[\"latent_size\"] = self.latent_size\n",
    "        model_dict[\"sample_size\"] = self.sample_size\n",
    "        model_dict[\"num_layers\"] = self.num_layers\n",
    "        model_dict[\"with_relu\"] = self.with_relu\n",
    "        model_dict[\"val_use_mean\"] = self.val_use_mean\n",
    "        model_dict[\"reparam_all_layers\"] = self.reparam_all_layers\n",
    "        model_dict[\"state_dict\"] = to_cpu_recur(self.state_dict())\n",
    "        return model_dict\n",
    "\n",
    "\n",
    "def load_model_dict_GNN(model_dict, is_cuda=False):\n",
    "    \"\"\"Load the GNN model.\"\"\"\n",
    "    model = GNN(\n",
    "        model_type=model_dict[\"model_type\"],\n",
    "        num_features=model_dict[\"num_features\"],\n",
    "        num_classes=model_dict[\"num_classes\"],\n",
    "        normalize=model_dict[\"normalize\"],\n",
    "        reparam_mode=model_dict[\"reparam_mode\"],\n",
    "        prior_mode=model_dict[\"prior_mode\"],\n",
    "        struct_dropout_mode=model_dict[\"struct_dropout_mode\"],\n",
    "        dropout=model_dict[\"dropout\"],\n",
    "        latent_size=model_dict[\"latent_size\"],\n",
    "        sample_size=model_dict[\"sample_size\"],\n",
    "        num_layers=model_dict[\"num_layers\"],\n",
    "        with_relu=model_dict[\"with_relu\"],\n",
    "        val_use_mean=model_dict[\"val_use_mean\"],\n",
    "        reparam_all_layers=model_dict[\"reparam_all_layers\"],\n",
    "        is_cuda=is_cuda,\n",
    "    )\n",
    "    if \"state_dict\" in model_dict:\n",
    "        model.load_state_dict(model_dict[\"state_dict\"])\n",
    "    return model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Training and testing:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Train and test functions:\n",
    "def train_model(model, data, optimizer, loss_type, beta1=None, beta2=None):\n",
    "    \"\"\"Train the model for one epoch.\"\"\"\n",
    "    model.train()\n",
    "    optimizer.zero_grad()\n",
    "    logits, reg_info = model(data)\n",
    "    if loss_type == 'sigmoid':\n",
    "        loss = torch.nn.BCEWithLogitsLoss(reduction='mean')(logits[data.train_mask], data.y[data.train_mask])\n",
    "    elif loss_type == 'softmax':\n",
    "        loss = torch.nn.CrossEntropyLoss(reduction='mean')(logits[data.train_mask], data.y[data.train_mask])\n",
    "    else:\n",
    "        raise\n",
    "    # Add IB loss:\n",
    "    if beta1 is not None and beta1 != 0:\n",
    "        ixz = torch.stack(reg_info[\"ixz_list\"], 1).mean(0).sum()\n",
    "        if model.struct_dropout_mode[0] == 'DNsampling' or (model.struct_dropout_mode[0] == 'standard' and len(model.struct_dropout_mode) == 3):\n",
    "            ixz = ixz + torch.stack(reg_info[\"ixz_DN_list\"], 1).mean(0).sum()\n",
    "        loss = loss + ixz * beta1\n",
    "    if beta2 is not None and beta2 != 0:\n",
    "        structure_kl_loss = torch.stack(reg_info[\"structure_kl_list\"]).mean()\n",
    "        if model.struct_dropout_mode[0] == 'DNsampling' or (model.struct_dropout_mode[0] == 'standard' and len(model.struct_dropout_mode) == 3):\n",
    "            structure_kl_loss = structure_kl_loss + torch.stack(reg_info[\"structure_kl_DN_list\"]).mean()\n",
    "        loss = loss + structure_kl_loss * beta2\n",
    "    loss.backward()\n",
    "    optimizer.step()\n",
    "\n",
    "\n",
    "def get_test_metrics(model, data, loss_type, beta1=None, beta2=None, compute_metrics=None, isplot=False):\n",
    "    \"\"\"Obtain the metrics of the data evaluated by the model.\"\"\"\n",
    "    logits, info_dict = model(data, isplot=isplot)\n",
    "    metrics = OrderedDict()\n",
    "\n",
    "    # Record structure kl:\n",
    "    structure_kl_list = [value.item() for value in info_dict[\"structure_kl_list\"]]\n",
    "    record_data(metrics, structure_kl_list, [\"structure_kl_{}\".format(i+1) for i in range(model.num_layers)], nolist=True)\n",
    "    record_data(metrics, [np.sum(structure_kl_list)], [\"structure_kl\"], nolist=True)\n",
    "    if model.struct_dropout_mode[0] == 'DNsampling' or (model.struct_dropout_mode[0] == 'standard' and len(model.struct_dropout_mode) == 3):\n",
    "        structure_kl_DN_list = [value.item() for value in info_dict[\"structure_kl_DN_list\"]]\n",
    "        record_data(metrics, structure_kl_DN_list, \n",
    "                    [\"structure_kl_DN_{}\".format(i+1) for i in range(model.num_layers)], nolist=True)\n",
    "        record_data(metrics, [np.sum(structure_kl_DN_list)], [\"structure_kl_DN\"], nolist=True)\n",
    "    if compute_metrics is not None:\n",
    "        info_metrics = model.compute_metrics_fun(data, compute_metrics)\n",
    "        record_data(metrics, list(info_metrics.values()), list(info_metrics.keys()), nolist=True)\n",
    "    for mask_id, mask in data('train_mask', 'val_mask', 'test_mask'):\n",
    "        mask_id = mask_id.split(\"_\")[0]\n",
    "        # Record Ixz:\n",
    "        ixz_list = to_np_array(torch.stack(info_dict[\"ixz_list\"], 1)[mask].mean(0)).tolist()\n",
    "        record_data(metrics, ixz_list, [\"{}_ixz{}\".format(mask_id, i+1) for i in range(model.num_layers)], nolist=True)\n",
    "        record_data(metrics, [np.sum(ixz_list)], [\"{}_ixz\".format(mask_id)], nolist=True)\n",
    "        if model.struct_dropout_mode[0] == 'DNsampling' or (model.struct_dropout_mode[0] == 'standard' and len(model.struct_dropout_mode) == 3):\n",
    "            ixz_DN_list = to_np_array(torch.stack(info_dict[\"ixz_DN_list\"], 1)[mask].mean(0)).tolist()\n",
    "            record_data(metrics, ixz_DN_list, [\"{}_ixz{}_DN\".format(mask_id, i+1) for i in range(model.num_layers)], nolist=True)\n",
    "            record_data(metrics, [np.sum(ixz_DN_list)], [\"{}_ixz_DN\".format(mask_id)], nolist=True)\n",
    "        if \"Z_std\" in info_dict and \"Z_std\" not in metrics:\n",
    "            record_data(metrics, info_dict[\"Z_std\"], [\"Z_std_{}\".format(kk) for kk in model.reparam_layers])\n",
    "        # Record contrastive loss:\n",
    "        if compute_metrics is not None:\n",
    "            info_metrics = model.compute_metrics_fun(data, compute_metrics, mask=mask, mask_id=mask_id)\n",
    "            record_data(metrics, list(info_metrics.values()), list(info_metrics.keys()), nolist=True)\n",
    "        if loss_type == 'softmax':\n",
    "            pred = logits[mask].max(1)[1]\n",
    "            loss = torch.nn.CrossEntropyLoss(reduction='mean')(logits[mask], data.y[mask]).item()\n",
    "            acc = pred.eq(data.y[mask]).sum().item() / mask.sum().item()\n",
    "            f1_micro = sklearn.metrics.f1_score(data.y[mask].tolist(), pred.tolist(), average='micro')\n",
    "            f1_macro = sklearn.metrics.f1_score(data.y[mask].tolist(), pred.tolist(), average='macro')\n",
    "        elif loss_type == 'sigmoid':\n",
    "            pred = torch.sigmoid(logits[mask])\n",
    "            pred[pred>0.5] = 1\n",
    "            pred[pred<=0.5] = 0\n",
    "            loss = torch.nn.BCEWithLogitsLoss(reduction='mean')(logits[mask], data.y[mask]).item()\n",
    "            acc = 0\n",
    "            f1_micro = sklearn.metrics.f1_score(data.y[mask].tolist(), pred.tolist(), average='micro')\n",
    "            f1_macro = sklearn.metrics.f1_score(data.y[mask].tolist(), pred.tolist(), average='macro')\n",
    "\n",
    "        record_data(metrics, [acc, loss, f1_micro, f1_macro],\n",
    "                    [\"{}_acc\".format(mask_id), \"{}_loss\".format(mask_id), \"{}_f1_micro\".format(mask_id), \"{}_f1_macro\".format(mask_id)], nolist=True)\n",
    "    return metrics\n",
    "\n",
    "\n",
    "@torch.no_grad()\n",
    "def test_model(\n",
    "    model,\n",
    "    data,\n",
    "    loss_type,\n",
    "    beta1=None,\n",
    "    beta2=None,\n",
    "    compute_metrics=None,\n",
    "    isplot=False,\n",
    "):\n",
    "    model.eval()\n",
    "    metrics_all = get_test_metrics(model, data, loss_type=loss_type, beta1=beta1, beta2=beta2,\n",
    "                                   compute_metrics=compute_metrics, isplot=isplot,\n",
    "                                  )\n",
    "    return metrics_all\n",
    "\n",
    "\n",
    "def train_baseline(model, model_type, data, device, threshold, filename, epochs, save_best_model=False, verbose=True):\n",
    "    \"\"\"Train the baseline model for the specified number of epochs.\"\"\"\n",
    "    data_record = {}\n",
    "    adj, features, labels = data.adj, data.features, data.labels\n",
    "    idx_train, idx_val, idx_test = data.idx_train, data.idx_val, data.idx_test\n",
    "\n",
    "    if model_type == 'GCNJaccard':\n",
    "        model = model.to(device)\n",
    "        model.fit(features, adj, labels, idx_train, idx_val, train_iters=epochs, threshold=threshold, verbose=verbose)\n",
    "    elif model_type == 'RGCN':\n",
    "        model = model.to(device)\n",
    "        model.fit(features, adj, labels, idx_train, idx_val, train_iters=epochs, verbose=verbose)\n",
    "    model.eval()\n",
    "\n",
    "    output = model.predict()\n",
    "    output_train = output[idx_train].max(1)[1]\n",
    "    output_val = output[idx_val].max(1)[1]\n",
    "    output_test = output[idx_test].max(1)[1]\n",
    "    train_f1_micro = sklearn.metrics.f1_score(data.y[idx_train].tolist(), output_train.tolist(), average='micro')\n",
    "    train_f1_macro = sklearn.metrics.f1_score(data.y[idx_train].tolist(), output_train.tolist(), average='macro')\n",
    "    val_f1_micro   = sklearn.metrics.f1_score(data.y[idx_val].tolist(), output_val.tolist(), average='micro')\n",
    "    val_f1_macro   = sklearn.metrics.f1_score(data.y[idx_val].tolist(), output_val.tolist(), average='macro')\n",
    "    test_f1_micro  = sklearn.metrics.f1_score(data.y[idx_test].tolist(), output_test.tolist(), average='micro')\n",
    "    test_f1_macro  = sklearn.metrics.f1_score(data.y[idx_test].tolist(), output_test.tolist(), average='macro')\n",
    "    record_data(data_record, [train_f1_micro, train_f1_macro, val_f1_micro, val_f1_macro, test_f1_micro, test_f1_macro, test_f1_micro, test_f1_macro, epochs],\n",
    "                ['train_f1_micro', 'train_f1_macro', 'val_f1_micro', 'val_f1_macro', 'test_f1_micro', 'test_f1_macro', 'b_test_f1_micro', 'b_test_f1_macro', 'epochs'])\n",
    "    if save_best_model:\n",
    "        data_record[\"best_model_dict\"] = model.state_dict()\n",
    "    if filename is not None:\n",
    "        make_dir(filename)\n",
    "        pickle.dump(data_record, open(filename + \".p\", \"wb\"))\n",
    "\n",
    "    return data_record\n",
    "\n",
    "\n",
    "def train(\n",
    "    model,\n",
    "    data,\n",
    "    data_type,\n",
    "    model_type,\n",
    "    loss_type,\n",
    "    beta1_list,\n",
    "    beta2_list,\n",
    "    epochs,\n",
    "    verbose=True,\n",
    "    inspect_interval=10,\n",
    "    isplot=True,\n",
    "    filename=None,\n",
    "    compute_metrics=None, # \"Silu\", \"DB\", \"CH\"\n",
    "    lr=None,\n",
    "    weight_decay=None,\n",
    "    save_best_model=False,\n",
    "):\n",
    "    \"\"\"Training multiple epochs.\"\"\"\n",
    "    if lr is None:\n",
    "        if model_type == \"GCN\":\n",
    "            lr = 0.01\n",
    "        elif model_type == \"GAT\":\n",
    "            lr = 0.01 if data_type.startswith(\"Pubmed\") else 0.005\n",
    "        else:\n",
    "            lr = 0.01\n",
    "    if weight_decay is None:\n",
    "        if model_type == \"GCN\":\n",
    "            weight_decay = 5e-4\n",
    "        elif model_type == \"GAT\":\n",
    "            weight_decay = 1e-3 if data_type.startswith(\"Pubmed\") else 5e-4\n",
    "        else:\n",
    "            weight_decay = 5e-4 \n",
    "\n",
    "    # Training:\n",
    "    optimizer = torch.optim.Adam([\n",
    "        dict(params=model.reg_params, weight_decay=weight_decay),\n",
    "        dict(params=model.non_reg_params, weight_decay=0)\n",
    "    ], lr=lr)\n",
    "    best_val_f1_micro = 0\n",
    "    b_test_f1_micro = 0\n",
    "    data_record = {\"num_layers\": model.num_layers}\n",
    "\n",
    "    # Train:\n",
    "    for epoch in range(1, epochs + 1):\n",
    "        beta1 = beta1_list[epoch] if beta1_list is not None else None\n",
    "        beta2 = beta2_list[epoch] if beta2_list is not None else None\n",
    "        train_model(\n",
    "            model,\n",
    "            data,\n",
    "            optimizer,\n",
    "            loss_type,\n",
    "            beta1=beta1,\n",
    "            beta2=beta2,\n",
    "        )\n",
    "        metrics = test_model(model, data, loss_type, beta1=beta1, beta2=beta2,\n",
    "                             compute_metrics=compute_metrics if epoch % inspect_interval == 0 else None,\n",
    "                             isplot=isplot >= 2 if epoch % inspect_interval == 0 else False,\n",
    "                            )\n",
    "        if metrics[\"val_f1_micro\"] > best_val_f1_micro:\n",
    "            best_val_f1_micro = metrics[\"val_f1_micro\"]\n",
    "            b_test_f1_micro = metrics[\"test_f1_micro\"]\n",
    "            data_record[\"best_epoch\"] = epoch\n",
    "            if save_best_model:\n",
    "                data_record[\"best_model_dict\"] = deepcopy(model.model_dict)\n",
    "        record_data(data_record, [epoch, best_val_f1_micro, b_test_f1_micro], [\"epoch\", \"best_val_f1_micro\", \"b_test_f1_micro\"])\n",
    "        record_data(data_record, list(metrics.values()), list(metrics.keys()))\n",
    "        if verbose and epoch % inspect_interval == 0:\n",
    "            record_data(data_record, [epoch], [\"inspect_epoch\"])\n",
    "            log = 'Epoch: {:03d}:'.format(epoch) + '\\tF1 micro: ({:.4f}, {:.4f}, {:.4f})'.format(metrics[\"train_f1_micro\"], best_val_f1_micro, b_test_f1_micro)\n",
    "            log += compose_log(metrics, \"f1_macro\", 2)\n",
    "            log += compose_log(metrics, \"acc\", tabs=2, newline=True) + compose_log(metrics, \"loss\", 7)\n",
    "            if beta1 is not None:\n",
    "                log += \"\\n\\t\\tixz: ({:.4f}, {:.4f}, {:.4f})\".format(metrics[\"train_ixz\"], metrics[\"val_ixz\"], metrics[\"test_ixz\"])\n",
    "                if model.struct_dropout_mode[0] == 'DNsampling' or (model.struct_dropout_mode[0] == 'standard' and len(model.struct_dropout_mode) == 3):\n",
    "                    log += \" \" * 7 + \"ixz_DN: ({:.4f}, {:.4f}, {:.4f})\".format(metrics[\"train_ixz_DN\"], metrics[\"val_ixz_DN\"], metrics[\"test_ixz_DN\"])\n",
    "                if \"Z_std\" in metrics:\n",
    "                    log += \"\\n\\t\\tZ_std: {}\".format(to_string(metrics[\"Z_std\"], connect=\", \", num_digits=4))\n",
    "            if beta2 is not None:\n",
    "                log += \"\\n\\t\\tstruct_kl: {:.4f}\".format(metrics[\"structure_kl\"])\n",
    "            if compute_metrics is not None:\n",
    "                for metric in compute_metrics:\n",
    "                    log += \"\\n\\t\"\n",
    "                    for kk in range(model.num_layers):\n",
    "                        List = [metrics[\"{}{}_{}\".format(id, metric, kk)] for id in [\"\", \"train_\", \"val_\", \"test_\"]]\n",
    "                        log += \"\\t{}_{}:\\t({})\".format(metric, kk, \"{:.4f}; \".format(List[0]) + to_string(List[1:], connect=\", \", num_digits=4))\n",
    "            log += \"\\n\"\n",
    "            print(log)\n",
    "            try:\n",
    "                sys.stdout.flush()\n",
    "            except:\n",
    "                pass\n",
    "\n",
    "        # Saving:\n",
    "        if epoch % 200 == 0:\n",
    "            data_record[\"model_dict\"] = model.model_dict\n",
    "            if filename is not None:\n",
    "                make_dir(filename)\n",
    "                pickle.dump(data_record, open(filename + \".p\", \"wb\"))\n",
    "\n",
    "    # Plotting:\n",
    "    if isplot:\n",
    "        plot(data_record, compute_metrics=compute_metrics)\n",
    "    return data_record\n",
    "\n",
    "\n",
    "def plot(data_record, compute_metrics=None):\n",
    "    def plot_metrics(\n",
    "        metric_list,\n",
    "        title=None,\n",
    "        X_source=\"epoch\",\n",
    "        linestyle_list=None,\n",
    "        color_list=None,\n",
    "        y_scale=\"standard\",\n",
    "        ax=None,\n",
    "        figsize=(8, 6),\n",
    "        is_legend=True,\n",
    "    ):\n",
    "        import matplotlib.pylab as plt\n",
    "        if ax is None:\n",
    "            plt.figure(figsize=figsize)\n",
    "            plt_save = plt\n",
    "        else:\n",
    "            plt_save = plt\n",
    "            plt = ax\n",
    "        fontsize = 15\n",
    "        for i, metric in enumerate(metric_list):\n",
    "            if y_scale == \"standard\":\n",
    "                plt.plot(data_record[X_source], data_record[metric], label=metric, alpha=0.8,\n",
    "                    linestyle=linestyle_list[i] if linestyle_list is not None else \"-\",\n",
    "                    c=color_list[i] if color_list is not None else None,\n",
    "                )\n",
    "            elif y_scale == \"log\":\n",
    "                plt.semilogy(data_record[X_source], data_record[metric], label=metric, alpha=0.8,\n",
    "                    linestyle=linestyle_list[i] if linestyle_list is not None else \"-\",\n",
    "                    c=color_list[i] if color_list is not None else None,\n",
    "                )\n",
    "            else:\n",
    "                raise\n",
    "        if is_legend:\n",
    "            plt.legend(bbox_to_anchor=[1,1])\n",
    "        plt.tick_params(labelsize=fontsize)\n",
    "        if ax is None:\n",
    "            plt.xlabel(\"epoch\", fontsize=fontsize)\n",
    "        else:\n",
    "            plt.set_xlabel(\"epoch\", fontsize=fontsize)\n",
    "        if title is not None:\n",
    "            if ax is None:\n",
    "                plt.title(title, fontsize=fontsize)\n",
    "            else:\n",
    "                plt.set_title(title, fontsize=fontsize)\n",
    "        if ax is None:\n",
    "            plt_save.show()\n",
    "\n",
    "    plot_metrics([\"train_f1_micro\", \"best_val_f1_micro\", \"b_test_f1_micro\", \"test_f1_micro\"], \"f1_micro\")\n",
    "    plot_metrics([\"train_acc\", \"val_acc\", \"test_acc\"], \"acc\")\n",
    "    plot_metrics([\"train_loss\", \"val_loss\", \"test_loss\"], \"loss\")\n",
    "    if \"train_ixz\" in data_record:\n",
    "        plot_metrics([\"train_ixz\", \"val_ixz\", \"test_ixz\"], \"I(X;Z)\")\n",
    "        if \"train_ixz_DN\" in data_record:\n",
    "            plot_metrics([\"train_ixz_DN\", \"val_ixz_DN\", \"test_ixz_DN\"], \"I_DN(X;Z)\")\n",
    "        if \"Z_std\" in data_record:\n",
    "            plot_metrics([\"Z_std\"])\n",
    "    if \"structure_kl\" in data_record:\n",
    "        plot_metrics([\"structure_kl\"], \"structure_IB\", y_scale=\"log\")\n",
    "        if \"structure_kl_DN\" in data_record:\n",
    "            plot_metrics([\"structure_kl_DN\"], \"structure_IB_DN\", y_scale=\"log\")\n",
    "    if compute_metrics is not None:\n",
    "        for metric in compute_metrics:\n",
    "            fig, axs = plt.subplots(1, 3, sharey=True, figsize=(24, 6))\n",
    "            for i, mask_id in enumerate([\"train_\", \"val_\", \"test_\"]):\n",
    "                color_list = [COLOR_LIST[kk] for kk in range(data_record[\"num_layers\"])]\n",
    "                plot_metrics([\"{}{}_{}\".format(mask_id, metric, kk) for kk in range(data_record[\"num_layers\"])],\n",
    "                             X_source=\"inspect_epoch\",\n",
    "                             color_list=color_list,\n",
    "                             title=\"{}{}\".format(mask_id, metric),\n",
    "                             ax=axs[i],\n",
    "                             is_legend=i==2,\n",
    "                            )\n",
    "            plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Dataset:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_data(\n",
    "    data_type,\n",
    "    train_fraction=1,\n",
    "    added_edge_fraction=0,\n",
    "    feature_noise_ratio=0,\n",
    "    **kwargs):\n",
    "    \"\"\"Get the pytorch-geometric data object.\n",
    "    \n",
    "    Args:\n",
    "        data_type: Data type. Choose from \"Cora\", \"Pubmed\", \"citeseer\". If want the feature to be binarized, include \"-bool\" in data_type string.\n",
    "                   if want to use largest connected components, include \"-lcc\" in data_type. If use random splitting with train:val:test=0.1:0.1:0.8,\n",
    "                   include \"-rand\" in the data_type string.\n",
    "        train_fraction: Fraction of training labels preserved for the training set.\n",
    "        added_edge_fraction: Fraction of added (or deleted) random edges. Use positive (negative) number for randomly adding (deleting) edges.\n",
    "        feature_noise_ratio: Noise ratio for the additive independent Gaussian noise on the features.\n",
    "\n",
    "    Returns:\n",
    "        A pytorch-geometric data object containing the specified dataset.\n",
    "    \"\"\"\n",
    "    def to_mask(idx, size):\n",
    "        mask = torch.zeros(size).bool()\n",
    "        mask[idx] = True\n",
    "        return mask\n",
    "    path = osp.join(osp.dirname(osp.realpath(\"__file__\")), '..', 'data', data_type)\n",
    "    # Obtain the mode if given:\n",
    "    data_type_split = data_type.split(\"-\")\n",
    "    \n",
    "    data_type_full = data_type\n",
    "    data_type = data_type_split[0]\n",
    "    mode = \"lcc\" if \"lcc\" in data_type_split else None\n",
    "    boolean = True if \"bool\" in data_type_split else False\n",
    "    split = \"rand\" if \"rand\" in data_type_split else None\n",
    "    \n",
    "    # Load data:\n",
    "    info = {}\n",
    "    if data_type in [\"Cora\", \"Pubmed\", \"citeseer\"]:\n",
    "        dataset = Planetoid(path, data_type, transform=T.NormalizeFeatures())\n",
    "        data = dataset[0]\n",
    "        info[\"num_features\"] = dataset.num_features\n",
    "        info[\"num_classes\"] = dataset.num_classes\n",
    "        info['loss'] = 'softmax'\n",
    "    else:\n",
    "        raise Exception(\"data_type {} is not valid!\".format(data_type))\n",
    "\n",
    "    # Process the dataset according to the mode given:\n",
    "    if mode is not None:\n",
    "        if mode == \"lcc\":\n",
    "            data = get_data_lcc(dataset.data)\n",
    "        else:\n",
    "            raise\n",
    "\n",
    "    if boolean:\n",
    "        data.x = data.x.bool().float()\n",
    "    \n",
    "    if split == \"rand\":\n",
    "        unlabeled_share = 0.8\n",
    "        val_share = 0.1\n",
    "        train_share = 1 - unlabeled_share - val_share\n",
    "\n",
    "        split_train, split_val, split_unlabeled = train_val_test_split_tabular(np.arange(data.x.shape[0]),\n",
    "                                                                               train_size=train_share,\n",
    "                                                                               val_size=val_share,\n",
    "                                                                               test_size=unlabeled_share,\n",
    "                                                                               stratify=to_np_array(data.y),\n",
    "                                                                               random_state=kwargs[\"seed\"] if \"seed\" in kwargs else None,\n",
    "                                                                              )\n",
    "        data.train_mask = to_mask(split_train, data.x.shape[0])\n",
    "        data.val_mask = to_mask(split_val, data.x.shape[0])\n",
    "        data.test_mask = to_mask(split_unlabeled, data.x.shape[0])\n",
    "\n",
    "    # Reduce the number of training examples by randomly choosing some of the original training examples:\n",
    "    if train_fraction != 1:\n",
    "        try:\n",
    "            train_mask_file = \"../attack_data/{}/train_mask_tr_{}_seed_{}.p\".format(data_type_full, train_fraction, kwargs[\"seed\"] % 10)\n",
    "            new_train_mask = pickle.load(open(train_mask_file, \"rb\"))\n",
    "            data.train_mask = torch.BoolTensor(new_train_mask).to(data.y.device)\n",
    "            print(\"Load train_mask at {}\".format(train_mask_file))\n",
    "        except:\n",
    "            raise\n",
    "            ids_chosen = []\n",
    "            n_per_class = int(to_np_array(data.train_mask.sum()) * train_fraction / info[\"num_classes\"])\n",
    "            train_ids = torch.where(data.train_mask)[0]\n",
    "            for i in range(info[\"num_classes\"]):\n",
    "                class_id_train = to_np_array(torch.where(((data.y == i) & data.train_mask))[0])\n",
    "                ids_chosen = ids_chosen + np.random.choice(class_id_train, size=n_per_class, replace=False).tolist()\n",
    "            new_train_mask = torch.zeros(data.train_mask.shape[0]).bool().to(data.y.device)\n",
    "            new_train_mask[ids_chosen] = True\n",
    "            data.train_mask = new_train_mask\n",
    "            make_dir(\"../attack_data/{}/\".format(data_type_full))\n",
    "            pickle.dump(to_np_array(new_train_mask), open(\"../attack_data/{}/train_mask_tr_{}_seed_{}.p\".format(data_type_full, train_fraction, kwargs[\"seed\"] % 10), \"wb\"))\n",
    "\n",
    "    # Add random edges for untargeted attacks:\n",
    "    if added_edge_fraction > 0:\n",
    "        data = add_random_edge(data, added_edge_fraction=added_edge_fraction)\n",
    "    elif added_edge_fraction < 0:\n",
    "        data = remove_edge_random(data, remove_edge_fraction=-added_edge_fraction)\n",
    "\n",
    "    # Perturb features for untargeted attacks:\n",
    "    if feature_noise_ratio > 0:\n",
    "        x_max_mean = data.x.max(1)[0].mean()\n",
    "        data.x = data.x + torch.randn(data.x.shape) * x_max_mean * feature_noise_ratio\n",
    "\n",
    "    # For adversarial attacks:\n",
    "    data.data_type = data_type\n",
    "    if \"attacked_nodes\" in kwargs:\n",
    "        attack_path = osp.join(osp.dirname(osp.realpath(\"__file__\")), '..', 'attack_data', data_type_full) \n",
    "        if not os.path.exists(attack_path):\n",
    "            os.makedirs(attack_path)\n",
    "        try:\n",
    "            with open(os.path.join(attack_path, \"test-node.pkl\"), 'rb') as f:\n",
    "                node_ids = pickle.load(f)\n",
    "                info['node_ids'] = node_ids\n",
    "                print(\"Load previous attacked node_ids saved in {}.\".format(attack_path))\n",
    "        except:\n",
    "            test_ids = to_np_array(torch.where(data.test_mask)[0])\n",
    "            node_ids = get_list_elements(test_ids, kwargs['attacked_nodes'])\n",
    "            with open(os.path.join(attack_path, \"test-node.pkl\"), 'wb') as f:\n",
    "                pickle.dump(node_ids, f)\n",
    "            info['node_ids'] = node_ids\n",
    "            print(\"Save attacked node_ids into {}.\".format(attack_path))\n",
    "    return data, info\n",
    "\n",
    "\n",
    "def remove_edge_random(data, remove_edge_fraction):\n",
    "    \"\"\"Randomly remove a certain fraction of edges.\"\"\"\n",
    "    data_c = deepcopy(data)\n",
    "    num_edges = int(data_c.edge_index.shape[1] / 2)\n",
    "    num_removed_edges = int(num_edges * remove_edge_fraction)\n",
    "    edges = [tuple(ele) for ele in to_np_array(data_c.edge_index.T)]\n",
    "    for i in range(num_removed_edges):\n",
    "        idx = np.random.choice(len(edges))\n",
    "        edge = edges[idx]\n",
    "        edge_r = (edge[1], edge[0])\n",
    "        edges.pop(idx)\n",
    "        try:\n",
    "            edges.remove(edge_r)\n",
    "        except:\n",
    "            pass\n",
    "    data_c.edge_index = torch.LongTensor(np.array(edges).T).to(data.edge_index.device)\n",
    "    return data_c\n",
    "\n",
    "\n",
    "def add_random_edge(data, added_edge_fraction=0):\n",
    "    \"\"\"Add random edges to the original data's edge_index.\"\"\"\n",
    "    if added_edge_fraction == 0:\n",
    "        return data\n",
    "    data_c = deepcopy(data)\n",
    "    num_edges = int(data.edge_index.shape[1] / 2)\n",
    "    num_added_edges = int(num_edges * added_edge_fraction)\n",
    "    edges = [tuple(ele) for ele in to_np_array(data.edge_index.T)]\n",
    "    added_edges = []\n",
    "    for i in range(num_added_edges):\n",
    "        while True:\n",
    "            added_edge_cand = tuple(np.random.choice(data.x.shape[0], size=2, replace=False))\n",
    "            added_edge_r_cand = (added_edge_cand[1], added_edge_cand[0])\n",
    "            if added_edge_cand in edges or added_edge_cand in added_edges:\n",
    "                if added_edge_cand in edges:\n",
    "                    assert added_edge_r_cand in edges\n",
    "                if added_edge_cand in added_edges:\n",
    "                    assert added_edge_r_cand in added_edges\n",
    "                continue\n",
    "            else:\n",
    "                added_edges.append(added_edge_cand)\n",
    "                added_edges.append(added_edge_r_cand)\n",
    "                break\n",
    "\n",
    "    added_edge_index = torch.LongTensor(np.array(added_edges).T).to(data.edge_index.device)\n",
    "    data_c.edge_index = torch.cat([data.edge_index, added_edge_index], 1)\n",
    "    return data_c\n",
    "\n",
    "\n",
    "def get_edge_corrupted_data(data, corrupt_fraction, is_original_included=True):\n",
    "    \"\"\"Add random edges to the original data's edge_index.\n",
    "\n",
    "    Args:\n",
    "        data: PyG data instance\n",
    "        corrupt_fraction: fraction of edges being removed and then the corresponding random edge added.\n",
    "        is_original_included: if True, the original edges may be included in the random edges.\n",
    "\n",
    "    Returns:\n",
    "        data_edge_corrupted: new data instance where the edge is replaced by random edges.\n",
    "    \"\"\"\n",
    "    data_edge_corrupted = deepcopy(data)\n",
    "    num_edges = int(data.edge_index.shape[1] / 2)\n",
    "    num_corrupted_edges = int(num_edges * corrupt_fraction)\n",
    "    edges = [tuple(item) for item in to_np_array(data.edge_index.T)]\n",
    "    removed_edges = []\n",
    "    num_nodes = data.x.shape[0]\n",
    "\n",
    "    # Remove edges:\n",
    "    for i in range(num_corrupted_edges):\n",
    "        id = np.random.choice(range(len(edges)))\n",
    "        edge = edges.pop(id)\n",
    "        try:\n",
    "            edge_r = edges.remove((edge[1], edge[0]))\n",
    "        except:\n",
    "            pass\n",
    "        removed_edges.append(edge)\n",
    "        removed_edges.append((edge[1], edge[0]))\n",
    "\n",
    "    # Setting up excluded edges when adding:\n",
    "    remaining_edges = list(set(edges).difference(set(removed_edges)))\n",
    "    if is_original_included:\n",
    "        edges_exclude = remaining_edges\n",
    "    else:\n",
    "        edges_exclude = edges\n",
    "\n",
    "    # Add edges:\n",
    "    added_edges = []\n",
    "    for i in range(num_corrupted_edges):\n",
    "        while True:\n",
    "            added_edge_cand = tuple(np.random.choice(num_nodes, size=2, replace=False))\n",
    "            added_edge_r_cand = (added_edge_cand[1], added_edge_cand[0])\n",
    "            if added_edge_cand in edges_exclude or added_edge_cand in added_edges:\n",
    "                continue\n",
    "            else:\n",
    "                added_edges.append(added_edge_cand)\n",
    "                added_edges.append(added_edge_r_cand)\n",
    "                break\n",
    "\n",
    "    added_edge_index = torch.LongTensor(np.array(added_edges + remaining_edges).T).to(data.edge_index.device)\n",
    "    data_edge_corrupted.edge_index = added_edge_index\n",
    "    return data_edge_corrupted"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Nettack:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_attacked_data_deeprobust(\n",
    "    data,\n",
    "    surrogate_model,\n",
    "    target_node,\n",
    "    direct_attack,\n",
    "    n_perturbations,\n",
    "    verbose=False,\n",
    "):\n",
    "    \"\"\"Nettack implemented by DeepRobust.\"\"\"\n",
    "    data_attacked = deepcopy(data)\n",
    "    features, adj, labels = data.features, data.adj, data.labels\n",
    "    device = data.x.device\n",
    "    nettack = Nettack(surrogate_model, nnodes=adj.shape[0], attack_structure=True, attack_features=True, device=device).to(device)\n",
    "    if direct_attack:\n",
    "        nettack.attack(features, adj, labels, target_node, n_perturbations, direct=True, verbose=verbose)\n",
    "    else:\n",
    "        nettack.attack(features, adj, labels, target_node, n_perturbations, direct=False, n_influencers=5, verbose=verbose)\n",
    "    modified_adj = nettack.modified_adj\n",
    "    modified_features = nettack.modified_features\n",
    "\n",
    "    adj_coo = nettack.modified_adj.tocoo()\n",
    "    data_attacked.edge_index = torch.stack(\n",
    "        list(to_Variable(adj_coo.row, adj_coo.col))).long().to(data.x.device)\n",
    "    data_attacked.x = to_Variable(nettack.modified_features.toarray()).to(data.x.device)\n",
    "    data_attacked = process_data_for_nettack(data_attacked)\n",
    "    info = {\"structure_perturbations\": nettack.structure_perturbations,\n",
    "            \"feature_perturbations\": nettack.feature_perturbations,\n",
    "           }\n",
    "    return data_attacked, info"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Evasive edge attack:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_evasive_dict(\n",
    "    dirname,\n",
    "    filename,\n",
    "    perturb_mode=\"edge\",\n",
    "    best_model=None,\n",
    "    feature_perturb_target=None,\n",
    "    structure_perturb_target=None,\n",
    "    verbose=False,\n",
    "    n_repeats=1,\n",
    "    device_name=\"cpu\",\n",
    "):\n",
    "    \"\"\"Get evasive attack metrics on feature or structure attacks.\"\"\"\n",
    "    from GIB.DeepRobust.deeprobust.graph.defense import GCNJaccard, RGCN\n",
    "    data_record = pickle.load(open(dirname + filename, \"rb\"))\n",
    "    parse_dict = parse_filename(filename)\n",
    "    data, _ = get_data(parse_dict[\"data_type\"])\n",
    "    baseline = parse_dict[\"model_type\"] in [\"GCNJaccard\", \"RGCN\"]\n",
    "    # Load model:\n",
    "    if best_model is None:\n",
    "        if parse_dict[\"model_type\"] == 'GCNJaccard':\n",
    "            data = process_data_for_nettack(data)\n",
    "            best_model = GCNJaccard(nfeat=data.features.shape[1], nclass=data.labels.max()+1,\n",
    "                               nhid=parse_dict[\"latent_size\"], #device=device,\n",
    "                               weight_decay=parse_dict[\"weight_decay\"],\n",
    "                               lr=parse_dict[\"lr\"],\n",
    "                              )\n",
    "            best_model.load_state_dict(data_record['best_model_dict'])\n",
    "        elif parse_dict[\"model_type\"] == 'RGCN':\n",
    "            data = process_data_for_nettack(data)\n",
    "            best_model = RGCN(nnodes=data.adj.shape[0], nfeat=data.features.shape[1], nclass=data.labels.max()+1,\n",
    "                         nhid=parse_dict[\"latent_size\"], #device=device,\n",
    "                         lr=parse_dict[\"lr\"],\n",
    "                         gamma=parse_dict[\"gamma\"],\n",
    "                         beta1=parse_dict[\"beta1\"],\n",
    "                         beta2=parse_dict[\"weight_decay\"],\n",
    "                        )\n",
    "            best_model.load_state_dict(data_record['best_model_dict'])\n",
    "        else:\n",
    "            best_model = load_model_dict_GNN(data_record[\"best_model_dict\"])\n",
    "            best_model.set_cache(False)\n",
    "    \n",
    "    if device_name != \"cpu\":\n",
    "        best_model.to(torch.device(device_name))\n",
    "\n",
    "    assert parse_dict[\"added_edge_fraction\"] == 0\n",
    "\n",
    "    best_model.eval()\n",
    "    df_dict_list = []\n",
    "    if perturb_mode == \"edge\":\n",
    "        if structure_perturb_target is None:\n",
    "            perturb_target = [-0.75, -0.5, -0.2, 0, 0.2, 0.5, 1, 2]\n",
    "        else:\n",
    "            perturb_target = structure_perturb_target\n",
    "    elif perturb_mode == \"feature\":\n",
    "        if feature_perturb_target is None:\n",
    "            perturb_target = [0.5, 1., 1.5]\n",
    "        else:\n",
    "            perturb_target = feature_perturb_target\n",
    "    else:\n",
    "        raise\n",
    "\n",
    "    for perturb_ratio in perturb_target:\n",
    "        if verbose:\n",
    "            print(\"seed: {}    {}: {}\".format(parse_dict[\"seed\"], perturb_mode, perturb_ratio))\n",
    "        for k in range(n_repeats):\n",
    "            df_dict = deepcopy(parse_dict)\n",
    "            df_dict[\"seed_evasive\"] = parse_dict[\"seed\"]\n",
    "            df_dict[\"best_epoch\"] = data_record[\"best_epoch\"] if \"best_epoch\" in data_record else np.NaN\n",
    "            df_dict[\"b_test_f1_micro\"] = data_record[\"b_test_f1_micro\"][-1]\n",
    "            df_dict[\"repeat_id\"] = k\n",
    "            # Get perturbed data:\n",
    "            if perturb_mode == \"edge\":\n",
    "                df_dict[\"added_edge_fraction_evasive\"] = perturb_ratio\n",
    "                if perturb_ratio != 0:\n",
    "                    data_c = deepcopy(data)\n",
    "                    edge_index_c = torch.LongTensor(pickle.load(open(\"data_evasive/{}_{}_seed_{}.p\".format(parse_dict[\"data_type\"], perturb_ratio, parse_dict[\"seed\"]), \"rb\"))).to(data.x.device)\n",
    "                    data_c.edge_index = edge_index_c\n",
    "                else:\n",
    "                    data_c = deepcopy(data)\n",
    "            elif perturb_mode == \"feature\":\n",
    "                df_dict[\"feature_noise_ratio_evasive\"] = perturb_ratio\n",
    "                if perturb_ratio != 0:\n",
    "                    data_c = deepcopy(data)\n",
    "                    x_max_mean = data_c.x.max(1)[0].mean()\n",
    "                    data_c.x = data.x + torch.randn(data.x.shape) * x_max_mean * perturb_ratio\n",
    "                else:\n",
    "                    data_c = deepcopy(data)\n",
    "\n",
    "            if \"struct_dropout_mode\" in parse_dict:\n",
    "                struct_dropout_mode = parse_dict[\"struct_dropout_mode\"]\n",
    "                if 'DNsampling' in struct_dropout_mode or (\"standard\" in struct_dropout_mode and len(struct_dropout_mode.split(\"-\")) == 3):\n",
    "                    add_distant_neighbors(data_c, int(struct_dropout_mode.split(\"-\")[-1]))\n",
    "\n",
    "            # Evaluate:\n",
    "            data_c = process_data_for_nettack(data_c)\n",
    "            if device_name != \"cpu\":\n",
    "                data_c.to(torch.device(device_name))\n",
    "            if baseline:\n",
    "                output = best_model.predict(features=data_c.features, adj=data_c.adj)\n",
    "                output_test = output[data_c.idx_test].max(1)[1]\n",
    "                test_f1_micro_best = sklearn.metrics.f1_score(data_c.y[data_c.idx_test].tolist(), output_test.tolist(), average='micro')\n",
    "            else:\n",
    "                test_f1_micro_best = get_test_metrics(best_model, data_c, loss_type=\"softmax\")['test_f1_micro']\n",
    "            df_dict[\"test_f1_micro_evasive_best\"] = test_f1_micro_best\n",
    "            df_dict_list.append(df_dict)\n",
    "    return df_dict_list"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Helper functions:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "def largest_connected_components(adj, n_components=1):\n",
    "    \"\"\"Select the largest connected components in the graph.\n",
    "    \n",
    "    Adapted from https://github.com/danielzuegner/nettack/blob/master/nettack/utils.py\n",
    "\n",
    "    Parameters\n",
    "    ----------\n",
    "    sparse_graph : gust.SparseGraph\n",
    "        Input graph.\n",
    "    n_components : int, default 1\n",
    "        Number of largest connected components to keep.\n",
    "\n",
    "    Returns\n",
    "    -------\n",
    "    sparse_graph : gust.SparseGraph\n",
    "        Subgraph of the input graph where only the nodes in largest n_components are kept.\n",
    "\n",
    "    \"\"\"\n",
    "    _, component_indices = connected_components(adj)\n",
    "    component_sizes = np.bincount(component_indices)\n",
    "    components_to_keep = np.argsort(component_sizes)[::-1][:n_components]  # reverse order to sort descending\n",
    "    nodes_to_keep = [\n",
    "        idx for (idx, component) in enumerate(component_indices) if component in components_to_keep\n",
    "    ]\n",
    "    return nodes_to_keep\n",
    "\n",
    "\n",
    "def get_data_lcc(data):\n",
    "    \"\"\"Return a new data object consisting of the largest connected component.\"\"\"\n",
    "    data_lcc = deepcopy(data)\n",
    "    edge_index_sparse = edge_index_2_csr(data.edge_index, data.num_nodes)\n",
    "    lcc = largest_connected_components(edge_index_sparse)\n",
    "    edge_index_lcc_sparse = edge_index_sparse[lcc][:, lcc].tocoo()\n",
    "    data_lcc.edge_index = torch.stack(list(to_Variable(edge_index_lcc_sparse.row, edge_index_lcc_sparse.col))).long()\n",
    "\n",
    "    data_lcc.x = data.x[lcc]\n",
    "    data_lcc.y = data.y[lcc]\n",
    "    data_lcc.train_mask = data.train_mask[lcc]\n",
    "    data_lcc.val_mask = data.val_mask[lcc]\n",
    "    data_lcc.test_mask = data.test_mask[lcc]\n",
    "    return data_lcc"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if __name__ == \"__main__\":\n",
    "    # Test dataset:\n",
    "    data, info = get_data(\"Cora\")\n",
    "    A_sparse = edge_index_2_csr(data.edge_index, data.num_nodes)\n",
    "    lcc = largest_connected_components(A_sparse)\n",
    "    data, info = get_data(\"Cora-lcc\")\n",
    "    data, info = get_data(\"Cora-lcc-bool\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}