examples/xgraph/pgexplainer.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "# PGExplainer on BA-Shapes dataset for 2-layer GCN"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "from dig.xgraph.dataset import SynGraphDataset\n",
    "from dig.xgraph.models import *\n",
    "import torch\n",
    "from torch_geometric.data import DataLoader\n",
    "from torch_geometric.data import Data, InMemoryDataset, download_url, extract_zip\n",
    "import os.path as osp\n",
    "import os\n",
    "\n",
    "device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "### Load dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "def index_to_mask(index, size):\n",
    "    mask = torch.zeros(size, dtype=torch.bool, device=index.device)\n",
    "    mask[index] = 1\n",
    "    return mask\n",
    "\n",
    "def split_dataset(dataset):\n",
    "    indices = []\n",
    "    num_classes = 4\n",
    "    train_percent = 0.7\n",
    "    for i in range(num_classes):\n",
    "        index = (dataset.data.y == i).nonzero().view(-1)\n",
    "        index = index[torch.randperm(index.size(0))]\n",
    "        indices.append(index)\n",
    "\n",
    "    train_index = torch.cat([i[:int(len(i) * train_percent)] for i in indices], dim=0)\n",
    "\n",
    "    rest_index = torch.cat([i[int(len(i) * train_percent):] for i in indices], dim=0)\n",
    "    rest_index = rest_index[torch.randperm(rest_index.size(0))]\n",
    "\n",
    "    dataset.data.train_mask = index_to_mask(train_index, size=dataset.data.num_nodes)\n",
    "    dataset.data.val_mask = index_to_mask(rest_index[:len(rest_index) // 2], size=dataset.data.num_nodes)\n",
    "    dataset.data.test_mask = index_to_mask(rest_index[len(rest_index) // 2:], size=dataset.data.num_nodes)\n",
    "\n",
    "    dataset.data, dataset.slices = dataset.collate([dataset.data])\n",
    "\n",
    "    return dataset\n",
    "\n",
    "dataset = SynGraphDataset('./datasets', 'BA_shapes')\n",
    "dataset.data.x = dataset.data.x.to(torch.float32)\n",
    "dataset.data.x = dataset.data.x[:, :1]\n",
    "# dataset.data.y = dataset.data.y[:, 2]\n",
    "dim_node = dataset.num_node_features\n",
    "dim_edge = dataset.num_edge_features\n",
    "# num_targets = dataset.num_classes\n",
    "num_classes = dataset.num_classes\n",
    "\n",
    "splitted_dataset = split_dataset(dataset)\n",
    "splitted_dataset.data.mask = splitted_dataset.data.test_mask\n",
    "splitted_dataset.slices['mask'] = splitted_dataset.slices['train_mask']\n",
    "dataloader = DataLoader(splitted_dataset, batch_size=1, shuffle=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "### Load model and checkpoints"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": "<All keys matched successfully>"
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "def check_checkpoints(root='./'):\n",
    "    if osp.exists(osp.join(root, 'checkpoints')):\n",
    "        return\n",
    "    url = ('https://github.com/divelab/DIG_storage/raw/main/xgraph/checkpoints.zip')\n",
    "    path = download_url(url, root)\n",
    "    extract_zip(path, root)\n",
    "    os.unlink(path)\n",
    "\n",
    "model = GCN_2l(model_level='node', dim_node=dim_node, dim_hidden=300, num_classes=num_classes)\n",
    "model.to(device)\n",
    "check_checkpoints()\n",
    "ckpt_path = osp.join('checkpoints', 'ba_shapes', 'GCN_2l', '0', 'GCN_2l_best.ckpt')\n",
    "model.load_state_dict(torch.load(ckpt_path)['state_dict'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 560/560 [00:00<00:00, 10215.51it/s]\n",
      "560it [00:15, 36.69it/s]\n",
      "560it [00:13, 40.76it/s]\n",
      "560it [00:12, 43.24it/s]\n",
      "560it [00:17, 31.17it/s]\n",
      "560it [00:22, 24.45it/s]\n",
      "560it [00:25, 21.85it/s]\n",
      "560it [00:24, 23.21it/s]\n",
      "560it [00:22, 25.10it/s]\n",
      "560it [00:22, 25.44it/s]\n",
      "560it [00:21, 25.67it/s]\n",
      "560it [00:23, 24.13it/s]\n",
      "560it [00:24, 23.08it/s]\n",
      "560it [00:23, 23.52it/s]\n",
      "560it [00:18, 30.19it/s]\n",
      "560it [00:17, 32.79it/s]\n",
      "560it [00:16, 34.47it/s]\n",
      "560it [00:16, 33.35it/s]\n",
      "560it [00:15, 35.07it/s]\n",
      "560it [00:16, 34.81it/s]\n",
      "560it [00:16, 33.43it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: 0 | Loss: 8.894268180855683\n",
      "Epoch: 1 | Loss: 8.846726562508515\n",
      "Epoch: 2 | Loss: 8.56760711020657\n",
      "Epoch: 3 | Loss: 8.010523775539228\n",
      "Epoch: 4 | Loss: 7.263117306360177\n",
      "Epoch: 5 | Loss: 6.238618776787605\n",
      "Epoch: 6 | Loss: 5.265949377949749\n",
      "Epoch: 7 | Loss: 4.691092160317514\n",
      "Epoch: 8 | Loss: 4.530065295419523\n",
      "Epoch: 9 | Loss: 4.613764908297786\n",
      "Epoch: 10 | Loss: 4.71995146330446\n",
      "Epoch: 11 | Loss: 4.820705992249506\n",
      "Epoch: 12 | Loss: 4.902427153342536\n",
      "Epoch: 13 | Loss: 4.950923837881003\n",
      "Epoch: 14 | Loss: 4.986500304724489\n",
      "Epoch: 15 | Loss: 5.010961620828935\n",
      "Epoch: 16 | Loss: 5.024333437266094\n",
      "Epoch: 17 | Loss: 5.03567554785737\n",
      "Epoch: 18 | Loss: 5.043968493065663\n",
      "Epoch: 19 | Loss: 5.047438699007034\n",
      "training time is 387.54s\n"
     ]
    },
    {
     "data": {
      "text/plain": "<All keys matched successfully>"
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from dig.xgraph.method import PGExplainer\n",
    "explainer = PGExplainer(model, in_channels=900, device=device, explain_graph=False)\n",
    "\n",
    "explainer.train_explanation_network(splitted_dataset)\n",
    "torch.save(explainer.state_dict(), 'tmp.pt')\n",
    "state_dict = torch.load('tmp.pt')\n",
    "explainer.load_state_dict(state_dict)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "### Visualization"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": "<Figure size 432x288 with 1 Axes>",
      "image/png": "iVBORw0KGgoAAAANSUhEUgAAAV0AAADnCAYAAAC9roUQAAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjQuMSwgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy/Z1A+gAAAACXBIWXMAAAsTAAALEwEAmpwYAAAowklEQVR4nO3de3xcZZ348c+Zya1JZibTtLTQlkwK5VIolCa9TDsod6WiIK7o1uW1rvsTdF3Eqgj7c+Xw6LpeAOOuuKiwP12VoKIIrkS5KKgTUtKmQFNpaaGZ9EJ6zWQmt+Y25/fHmWmnIZe5XzLf9+vly2fOnPOcr7b59sk5z/N9NMMwEEIIkRmWbAcghBCFRJKuEEJkkCRdIYTIIEm6QgiRQZJ0hRAigyTpCiFEBhWlrKdGrRS4EbgdcAFlwADQDnwLeJYNRihl9xNCiDykJT1Pt1ErBu4GPg1ogG3cGQbQB/QD9wA/YINMDhZCFKbkkm6jVgE0AfVAeQxX9AM/A26RUa8QohAl/ky3UbMCTwKriC3hAlQAH8Z83CCEEAUnmRdptwJrMJ/dxqMC+D80au9M4t5CCJGXEku6jZoGfAEzgSaiHLgjwWuFECJvJTrSvRSoTuK+GnAljdrpSfQhhBB5J9Gk+48kPsqNdlMK+hBCiLyRaNKtwRytJqMMWJBkH0IIkVcSTbqzUnT/VIyWhRAibySadI+l4N4GcDgF/QghRN5INOn+AXOJbzL6gE1J9iGEEHkl0aT7wySuBcAw6AeeTaYPIYTIN4klzg1GN/AEMJbI5cOhIl44dlmJ2n3PvITuL4QQeSqZ0aoCjsd7UciA4VAprYHVs4FOpdSGJGIQQoi8kmzBm2uAXxNj7QXDIDRslPDw3o9bjo7Mjf7ql8BNuq5L9TEhxIyWXBHzDcYzwLVAELOC2FR6NY2DpZbh84+OzH1g3Hd/AxxQSrmSikcIIXJc8vV0ARo1O3AzZj2FaiAEWDGf+RYBu4BvAI+zwRgGUEpdAfwvp46Sx4BP6rr+UPJBCSFE7klN0o0wC+GsAmoxFz4EgdfYYPx1otOVUjbgz8DycV/9DnivrusJvagTQohcldqkmyCl1NeAOzl1afFR4J26rr+WnaiEECL1cmJjSl3X/wVYC/RGHZ4DtCulPpudqIQQIvVyYqQboZSaBTyHmYCj/Rm4Rtf1ocxHJYQQqZNTSTdCKXUX8FVOHYkHgKt0Xd+SnaiEECJ5OfF4YTxd178OrAD8UYcdQKtS6p6sBCWEECmQkyPdCKVUMeZuw1eN+6oN8yXbdHODhRAip+R00o1QSv0z8G3Mub8R/cB7dF3/U1aCEkKIBORF0gVQSp2D+UJtfJGcb+m6/rkshCSEEHHLm6QLoJTSMGs9XD/uq9cAj67r/rdfJYQQuSOvkm6EUurvgYcxlxhHHAdu1HX9d9mJSgghppeXSRdAKbUIaAYWjfvqIV3Xb8lCSEIIMa28Tbpw4nHDT4CPjPtqD7BW1/VDmY9KCCEml9dJN0IpdSPwKFASdXgEuFnX9Z9nJyohhHi7GZF0AZRSc4EXgbPHffVz4G+lQLoQIhfMmKQboZT6HnDruMMHgHW6rndmISQhhDghJ5cBJ0PX9U9g7mYRvX/bAuANpdTHshOVEEKYZtxIN0Ip5QC8wIXjvnoKuF4KpAshsmHGJt0IpdS9wOfHHT4CvEPX9Z1ZCEkIUcBm3OOF8XRdvwO4FOiLOjwX+KtS6jNZCUoIUbBm/Eg3Ilwg/Xlg9bivngfepev6SOajEkIUmoJJuhFKqS8BilP3Y+sBrtB1/eWsBCWEKBgz/vHCeLqufwWow0y0EVVAm1Lqi9mISQhROApupBsRLpD+NHD5uK9agct0XR/MfFRCiJmuYJNuRPhl2v2cOurvA67Vdd2blaCEEDNWwSddAKXUeZgF0ueO++q+8OwHIYRICUm6YUopK/Ak8J5xX23HLJAeyHxUQoiZRpLuOEqpfwS+z6n7sR3HXMX2THaiEkLMFJJ0J6CUcmEWSD9j3FffD9d2EEKIhEjSnUS4QPqjwIfGffUGZoH0I5mPSgiR7yTpTkMp9SHM3SmKow4PAxt0Xf9VdqISQuQrSboxUErNwyyQvnjcV43A30mBdCFErCTpxkEp9QPg4+MO78d83LAvCyEJIfJMwS0DTkZ4l+HrOLVA+kJgT3hbeCGEmJKMdBOglHICfwEuGPfVb4Ab5HGDEGIyknSToJS6H/jsuMOHgUt1Xd810TWGplmAJcDs8KGjwBvyByFEYZCf9SQppd4JNAHlUYfHgNt1Xf9u5IChadXAxzCTtA0YDX9VDHQD9wL/oxmGrHwTYgaTpJsCSqkK4E+YJSOjPTf38OH1n/yv//oM8GUgxKnJOVo/5jP2z2mG8WC6YhVCZJck3RRSSingS0QKpBsG1z711FBdW1vIYhizYuxmALhfM4y70xSmECKLJOmmmFKqHngOcKx+8UUuf/55Skbi3gloALhNM4z/l/IAhRBZJVPGUkzX9S3AvLKBgRev+OMfE0m4YD6C+LahaaWpjU4IkW0y0k0TQ9P+acxiabCGQiUJdtELfFIzjEdSGZcQIrsk6aaBoWka0AHUJNlVu2YYF6UgJCFEjpDHC+kxD5ifgn7ONzStIgX9CCFyhCTd9HBiViJLyhiEhsy+hBAzRFG2A5ihQqnoZAhKzoZ9/rIyo7Ky8rjD4Thgt9tfqKiouN/r9e5MxT2EEJklz3TTILz67C0g0ZdogLlkzQYMTfCdxWKhvLw8ZLPZeh0Ox06bzfa4xWL5zqZNm2TreCFymCTdNDE0bStwSTJ9/BG4Ks5riouLqaioGLbb7UccDscWm832X83NzQW1t9v6jS1OzPoWpZj/Zu1uanD7sxuVECZJumliaNqHgIcwB6uJ6B2DmzyrV9cGg8ENgUDg/GAwWDU4OGgdGxuLqyNN0ygtLTXsdvuAzWbb63A4np41a9Y3vV5vV4Kx5Zz1G1uWAxuBd2M+Bx+I+roc8AO/BxqaGtyvZDo+ISIk6aaJoWklwCGgKsEuDgILNMN42/Nhj8dTOzAw8LlgMHhlMBg8s7e3d9bQ0FDcf5ZWq5Xy8vIxu93eY7fbt9vt9saDBw8+5PP58uYvxfqNLecAP8Uss1nKqbs4jzeGOfLdDtzc1OCesBKcEOkkSTeNDE27DvgFEGvdhYgB4DrNMJ6P56K1a9de19vbe2swGKwLBoNz+vv7i0cSWBFXUlJCZWXlkN1uP+hwOJorKyv/w+v1tsbdUZqt39hyG/ANzGQbz0ycMczZJXc2Nbi/k47YhJiMJN00MzTtH4DvEnviHQQ+qhnGL1Jx/zVr1swKhUK3B4PBDwQCgSW9vb22wcFBSygU3wQLTdOYNWtWyGaz9Tscjj12u/03RUVF97a0tPSmIs54rd/Y8g3gU0Ay85gHgAeaGtx3piYqIaYnSTcDDE17N/DfgB0zSWjR34cw6zr6gTPhMs0w/pTumDwezwX9/f2fDQaDlwUCgdP7+vpmDQ1NNE9iakVFRZSXl486HI5jdrv9VZvN9nBXV9cv0/mIIjzC/RrJJdyIAeAuGfGKTJGkmyHhHSOuAO4ArgQMQAsBT4H1Psz9f1asWHFfW1vbHdmI0eVyaaeffvqG3t7ejwYCgYuCweDsgYGBotHR0ekvHqe0tDQtc4vDz3BfIf5HNlMZBC5uanDvTmGfQkxIkm4WhGszlGMm3kF7ZeVIX1+fFcDpdA51d3eXZTXAcTwez9zjx49/NhgMrg8Gg7W9vb0Vg4ODlnj/7kTNLQ44HI5dNpvtsbGxsf9sa2uL+cHz+o0tLwH1pHY15RjQ1tTgXp3CPoWYkCTdHHDRRRf9or29/YORz+vWrVudiy+uxlu3bt26vr6+TweDQXcgEDitv7+/dHg4/tXPkbnFDofjiN1un3Ru8fqNLZcAXibffSMZA8A6mU4m0k2Sbg5wuVza/v37Q5H5tzU1NQd8Pt/CLIeVkLq6uuKioqJbent7PxwIBJYGg0FHonOLy8rKDJvNNmC32zvtdvszp627b4FmLbuRqaeFJWoUeKSpwf3RNPQtxAmSdHNEbW2tz+fz1YD5a/iiRYss+TRfdjpRc4uvCgaDixKZW3zFLb+ltHz29Ccm7lBTgzsV1eGEmJRUGcsRCxYsuD7SDoVC2Gy232QznlTzer0dW7du/ec33njjvMOHD1eEp61pbrf7vRdeeOFvzzzzzK6qqqqR4uLiCa8vKrVRXJro4r6YzQ4vIRYibWSkm0PmzJkzeOzYsTKAioqKUOTlWqEJzy3+bDAYvCEyt7jE7rLU3dBAcWllOm8dAK5panDn/PN0kb9kpJtDamtr7420+/v7LWvWrLk5m/Fky6ZNmwZbW1u/unPnzpVdXV1VfX191jUfeugdRSXlmViIIfvSibSSpJtDNm/efHdpaemJXz06Oztlwv5JQ5pmSUmd4unuk4F7iAImSTfHuFyuzZH2oUOHHB6PZ24248khu0nPVLFo5eH7CJE2knRzTHV19bvNtRNgGAZHjhx5Nssh5YRwPdx018Ttlrq7It0k6eaY5uZm//z580/84Pt8vouzGU+O+T3m6rF0GA33L0RaSdLNQTU1NbdG2sPDw9TV1d071fkFpIH0PXMdBr6dpr6FOEGmjOUom802msv1GLJFai+IfCcj3RxVW1v7eKTt9/tLPR7PqmzGk0NuJvWj3WHg71LcpxATkqSbo4LB4Ies1pNrI/bv3/+rLIaTM8Jb7NzJqXugJWx0ZJC+vc9tk7KOIlMk6eYon89nLFq0qDPyef/+/QtdLpc21TWFIlxw/AGSTLyjI4N0vvJL/vL43auXLVv2ZGqiE2JqknRzWHQ9hrGxsRlXjyEZ4S127sIsQB7vjIYxIzQ6uqv5QXY1PwjA9u3b37d06dK49qQTIhGSdHOY1+t9dfbs2ccjn30+3/psxpNrwiPe5UAb5qh3ui0uRsPntWmWoqVzi3z3ROZEA+zYseOy8847b0uawhUCkKSb8xYvXnxiulhfX1/B1mOYTFODe1d41sE64BHMbe9HgICF0ZCVETRCWBk5Hv5+XVODe3VTg3v35s2bVX19/e0Wy8kfg9dff73u7LPPfj0b/1tEYZApY3mgrKwsNDQ0pAHMnz8/0NXVVZXlkHJauDzjkiUVLzeUaENryywDFFlGX9d1/byJzne73R/ZvHnzT6MLrdfU1Oz3+XyLMhWzKBwy0s0DLperLdKWegzTa2pw+5sa3K2zi4/8srIoSJFlFGDSBNrS0vLIqlWr1hcVFZ041tnZuXDhwoVHMxCuKDCSdPNAdXX1NVKPISGPR7XLlVKTVkF/8cUXf7d69ep1JSUlJ44dOHCgev78+UGZNSJSSZJuHpB6DInRdb0T8/luxHunOt/r9b64cuXKC6PLax46dMjW19c3UFdXN/GWFkLESZJunpB6DAk7FNV+93Qne73ev9bX19eUl5efqN177NixsjfffLPf7Xanfb8gMfNJ0s0TLS0tj1VWVp5409PR0XFbNuPJI9uj2itiucDr9e5bvnz5HJvNduL/70AgULx9+3a/PE8XyZKkm0ekHkNCXohqu2K9qLm52b906VJbVVXVcORYX1+fdevWrQc9Hk9tCuMTBUaSbh6RegwJ+XVUu0IpNSvWCzdt2jS4ePHiyjlz5gxGjg0ODlq2bNnypsfjkefqIiGSdPOI1GOIn67ruzh1mXBcq/ra2tpGKioqKubPnx+IHBsaGtJaW1tfWbdu3TtTFacoHJJ088wE9RieyGI4+SL6Zdq18V7s8/mMrq6uqoULFx6OHBsZGaG1tfUFt9t9/VTXCjGeJN08M0E9huuyGU+eeC2qvTLRTvbt2zfP5XLtjXweHR2ltbX1iTVr1nw0meBEYZGkm4ekHkPc/hLVTuolWEdHR81ZZ521I/I5FArR2tr6w1WrVm1Mpl9ROCTp5qHNmzffHT2Bv7Oz8zvZjCcPRL9wtCmlklro8MYbbyw999xzWyOfDcNgy5Yt36qvr/+3ZPoVhUGSbp6Segyx03X9r0Ao6tC0iySms3PnztUXXHDBc5HPhmHQ1tb2xRUrVnw32b7FzCZJN09JPYa4HYlqx/0ybSLbt2+/etmyZY9FH3v55Zf/6eKLL340Ff2LmUmSbp6Segxx2xnVTtmuv9u2bbvpkksueWjcsQ8vW7asKVX3EDOLJN08Nr4eQ319/TeyGU+O80a1z0plx1u3br2lrq7um9G7UGzfvv3a888/3zvFZaJASRHzPGez2Ub7+vqsAE6nc6i7u7ss2zHlIqVUPbA56lCRruvx7q02pVWrVt21ZcuWr0X/TC1ZsmTbrl275LeQfNGonQu4gSpgCDgIPM0GIyW7T4OMdPOe1GOIWRunvky7MtU3aG1t/fqqVatujd7+Z/fu3RctXrx4T6rvJVKoUSuiUbuRRu0l4GXgO8DXgfuBHwGHadQeoFE7JxW3k6Sb56QeQ2x0XTeA7qhD70nHfTZt2vSDVatW3RT9Z9LR0VF75plndqXjfiJJjdppmIn2f4BVwCygEigNt+1ABfBx4BUatc8ne0tJunlO6jHEJXrDyTXpuklLS8tjq1evvqq4+OR04H379s0/44wzuuXPJoc0anOBrcC5mIl2KiWYSVjRqH0lmdtK0p0BpB5DzJqj2kvSeqPm5j+sWrWqLnr7n66uLufAwECvJN4c0KhZgGeBuUA8i2XKgc/SqH040VtL0p0BpB5DzP43ql2llEpr8vN6vVtXrly5pKys7MSbtSNHjlQEg8Hja9asibnEpEiLqzBnsZRMd+IEyoFv0Kgl9PdHku4MIfUYYtIMRBKgBnjSfUOv1/tGXV3dgoqKihMv8fx+f8mOHTuCsv1PVn0B81ltomYD70jkQkm6M4TUY5he+GWaP+pQRsoyer3ermXLllXZ7fbRyLFgMFjU3t7e4/F4Ts9EDCJKo7YIWIf5D2+iKoA7ErlQku4MIvUYYrI7qu3O1E1bWlp6zz777PLZs2cPRY719/db2traDng8nvMyFYcAYDkwPN1J09BIsEyoJN0ZROoxxKQlqp2SeZexamtrG7HZbLNOO+20/six48ePa5s3b94h86vTTylVrJS63Nvt+dhIqCiZRwsRCfUhK9JmmDPOOKO7q6vLCVBSUsLQ0JC8KY+ilLoC+EP4owFYw48dMsblcmkjIyPH3nrrLWfkWHFxMStXrnxXc3PzM5mMZSYK74N3Wfg/KzBnqpyGOeWL8yv/yvtOe5Iya7KDXfxsMGbHe1FRsncVuaWmpubWrq6uX8DJegxbtmy5M9tx5ZA/RbUjvyK2TnJuWvh8PgOYXVNT89bevXtPB3P7n5deeulpt9t9U0tLy2PTdCEApZQNuBq4FLgEczbCXMyFDZPqHbUn9TA3yqHpT3k7GenOQFKPYWpKKT/m2nqAr+u6/i/ZiuWss856c8+ePYsjny0WCytXrrx106ZNP8hWTLlGKVXNyeR6MbAYmEN882sjQhqhwOdq77NVFA0kM+jsB+5ig/FAvBfKSHcGqq2tfaK9vf0DcLIeg9frzehoLse9CdSF22uzGsibb551zjnnvLp79+6L4MT2P99ftWrV7NbW1q9nM7ZMU0qdDrwLcyrfMsCFOTUrkTw1BvQAe4HtwIvAM7qum3UwGr/8OeDLmHNuE2EBfpzIhZJ0Z6BgMPhBq9UaGhszi2iF6zEsym5UOeUlTibdrM8c2LVr18VLly79y44dOzxgvgTdvHnz1+rq6qrb2treNi1p/cYWC+AIfww0NbhD48/JZUopF+buHWuBCzH/bjoB6xSXTWYUs6aGD2jHLOH5rK7rB6a57oeYSTcRx4FH2WAEE7lYHi/MULW1tT6fz1cDYLVaWbhwoSX8LLHgKaWuBSJFxkO6rifyw55yy5Yt++327dtPKcSzfPnyh19++eWPr9/YYsUcBX4B89fsyJxfK/A8cC/wXK4k4PBqv/OBazALyVyAmVwdJDZragQ4CuwBXsXcbPRZXdePJRxko3Y98CjhF2wxGg3HUM8GozeR20rSnaE8Hs/Fzc3Nr0Q+X3jhhb9pb2/PyGKAXBfemDL61fVyXddfzVY80ZYvX/7Iq6++uiH62Oprbm2evfTvl2Amh4lWsRlAH+Zzxr9vanBnbAZEOLkux1xWuxoz0S7ArM6VyPuqIcytld4EXgH+jJlcE0pw02rU/g74AbEl3iGgA7icDcbBRG8pSXcGq66uHoy8RKusrAz19vbmxIguFyilApiJAUDpun5PFsM5xYoVKx54+eWXPwWw8ML3sfSdn8FaHPO70EHg1qYG909SGZNSyoqZVK/EnPFxLnA6E/8jEIvjmG//d2NW+noBeEHX9cGkg41Xo7YO8zeF5ZiPXMe/oOsL//d/A//KBqOPJMgz3Rls8eLF93Z3d38JTtZj2LRpU0p/GPPYHswfMjB/Xc8ZW7du/eeVK1d27+2e9aU4Ey6YI7bvr9/Ycripwf10vPcO/xZwKXAF5hzXczCTa6IvnAaALmAXsAXzUYhX1/WRBPtLvQ1GM7A2XKT8dsz/7XbM34YOAd8Dfs4GIyX/IMhId4YrKysLRRZIzJ8/P9DV1VWV5ZByglLqB5iFqQEO6rqeUzUQ1m9sKTFCIz2apTjRamRHgflNDe4JtySaZAHBPCDR6YW9wFuYNYu3YC5AeSnVWyLNBDLSneFcLlfb66+/Xg8n6zF4vd4j011XAH7PyaSbizUqbtAsxaPTnzapUmC9UuoFzDmu78Ac2ce0gGASBhAEDgA7MGeBPAe8kulVfflMku4MF67H0G0YRqQewzOYq3cK3VNRbatS6jxd13dOenbm3Uniz0sBbJVW/5Mk9jIrBASAfcBfgU2Yc1xz6f+fvCWPFwqA1GOYmFKql5PbtHxR1/V/z2Y8Ees3tjgxnyUmsuIqSoh6xx+xapPOIhvDLHW5F3OOa2QBgS+5+4qpyEi3AEg9hkl1Ys4fBfPX75xIukA15vSkpJKuBYNRoxirNjQKHMNcQLCd2BcQiDSQpFsAWlpaHrPZbGORegwdHR23Y/76Wug2czLpXpjNQAA8Hs8FwWBQDYyWX7b48q9WFpUkOmHAFEIb2je4ZNVPvvmBbSkKUaSA1NMtELW1tU9E2t3d3aVSvxWA6ClV8zJ5Y7fbbaurq7v3rLPOenP27NlDRUVFRnNz8/b29vYP7O3YWW2xpGI8ZLEeGzmjIwUdiRSSkW6BkHoME4reqLJIKbX4REGUFHK5XNr8+fM/1tPT80m/339eT09PxfDw5LVcR44HGOw9RIUz6T+eHU0N7vSs5BIJk5FugfD5fMaiRYs6I5/379+/sNC3Atd1vR9z8n7Ejano1+PxrF22bFnTggULjlZUVIzt3bs39NJLLz38+uuv1x0+fHjKhGu1Wqmqqhru2/fHo0ZoJJkFBL3AN5K4XqSJjHQLyIIFC673+XyvAIyNjWGz2Z4gQ5sz5rC9nKw0dhlwXzwXezyeuQMDA1/1+/3v7u7uPr2vr68oFIqt5oymaVRUVISqqqqOOp3OF+12+z1er/dVgPUbWyqAwyT3Mu2XSVwr0kSSbgHxer2vVldXH4/UY/D5fNdlO6Yc0MbJpLtsqhNdLpd22mmn3dHT03Nzd3f3WcFgcNZIHIPR0tJSo6qqqs/pdG6vqqr6z5aWlp9Ndm5Tg7t//caWz2P+IxDvG7UB4NNNDe6hac8UGSdJt8BIPYa3eRb4SLh9ylLgtWvXXtvT0/MFv99/SSAQsA8ODmqdnZ1v72ECxcXF2Gy2IafT6XM6nY+GQqF/b2tri+txQVOD+8H1G1vOBD5N7Il3APhmU4P7R/HcS2SOLI4oQFKP4SSllAPo6e/v57XXXqO9vf2w3++v7u/vt8b6mMBisVBRUTHqdDqPOJ3OZysrK+/xer0pmzWwfmPLp8ZGhx4wjBBFk5di6MdcfXZ7U4P74VTdW6SejHQLUKHXY1izZs2s0dHRL/r9/pv8fv+Zvb29jI6eKHNw2nTXz5o1K1RVVdXrdDq3VFVVfTPdO/ge8X563ivbXuOM865lcf1HKLPNHdY06yBmki0GDmK+NGuU2Qq5T5JuASq0egxut/sjPT09t3V3d18QCAQq4lkGXVJSgt1uH3Q6nburqqp+dPjw4W9negeOjo6OO0aHB9i77Vcce/P3oXd8/NlzMfcOA7Oa2L6mBrf8ypon5PFCgZqp9Rg8Hs+KQCCg+/3+tT09PbMHBgYssf4dt1qt2O12w+Fw7HU6nU+Vl5ffk+3fACbYAeS37e3t781iSCJJMtItUDOhHsO6deucg4ODX+np6XlPd3f3gr6+vuLI4o/paJpGeXl5qKqqqnvRokVvrVmz5qKqqiqAEV3XXWkMOy4HDhx4MtK2Wq309va+L5vxiOTJSLeA2Wy20Ug9htmzZw8dO3Ys0QLWaReernVbIBD4WHd395JgMFg+1SKD8UpLS3E4HH1Op/O1qqqqBw8ePPg/kccESqm5mHNiI07TdT3rz7hdLpe2f//+E6sIa2pq9vt8vkJfRZj3ZKRbwGpra59ob2//AEB1d3fp7oqKB88eGACzwPVBzH2rntMMI+M7zK5bt+7KQCBwV3d3d10gEHAMDAxYYp2uVVRUhM1mG3Y6nXudTufjRUVF92zatGnSrVZ0XT+ilBriZGHv92NuVphVdrv959Ej94ULF34gi+GIFJGkW8D+Ye/eD/7FYgndEQpxEVA8MPCJqK8N4Dag19C0+4GHNcMIpCMOj8ezqL+//8vd3d1X9/T0zEtgVdeY0+k86nQ6n7fZbMrr9SZSbPsAsDjcvpocSLodHR0nliU7nc4hr9fbms14RGpI0i1QhqaVfgYab9E0o3zi3QU0zALflcBXgNsNTbtMM4ykCsLU1dUVWyyWu/x+/9/6/f7Fvb29pfGs6iorKzOqqqqCTqfzFYfD0dDS0vLk9FfFZBsnk+7yFPWZMLfb/cHIox+A2tra72QzHpE6knQLkKFpVuA3wKXlhhHLrIVZwAKg1dC0SzTD2Bfrvdxu942BQGBjd3f3xYFAoPL48eMxz5IoLi7Gbrcfdzqde5xO50/Gxsbuj3dVVxz+ANwQbi9M0z1i1tnZ+f1Iu6SkhLa2tjuyGY9IHUm6hekrgAczmcbKAjiAPxiadt5Ez3kjRbj9fv+lPT091f39/dZYX9RaLBYqKytHnU7nQafT+fuKioq7vV5vVxzxJetXQGQ0WaaUcuq67s/g/U9Yt26d8+DBg87IZ5fL9Wo24hDpIUm3wBiaNov41vJHK8KsT3CV2+1uGR4evqenp+cGv9+/MBgMlsQ6XQuITNcKOJ3OVofD8bXm5uY/JRBPyui63qWUGuFkVa/rgR9lI5Zjx449E/nHStM05s6de3U24hDpIUm38HwI8yVZQkJQ+bymPb1p06aYrykpKcHhcAw4nc6dTqfz4YMHD34v06u6YtQFnBluX0OWkq7P56uLtOfNmxfI9gINkVqSdAvP5zm5A27cLMBaw+B0zAw1ntVqjUzXOuB0Op8sKSm5u6WlJV/qAWzjZNJdkY0AVq5c+eXo1YE1NTW3ZSMOkT6SdAvPWcl2MAQsAQ6Gp2tVVVUdG1+EO089D0RqDGdlEcKePXtOvDCrrKwMFXjZzRlJtuspIIamaZxcAJCwYk0b+8T8+feGQiGtt7e3aN++ffO2bdv2/jxPuGC+TIsoV0rZMnlzj8dzcaTAPIDL5fptJu8vMkOSbgEJr/mOfe3sJCoMo+/DBw/+LgUh5RRd1zuB0ahDGS0sM0GdhRsyeX+RGZJ0C8+BFPRRAszUrb0PRrWvydRNXS6Xtm/fvprI54ULF+7P0ZeNIkmSdAvPtzF3GUjGNs0wfMmHkpO2R7XrM3VTu93+mNRZKAySdAvPj0nuz32mb+0dPV/YlambdnR03BBpz549W+oszGCSdAtMuGjNT4BJq25NwQD6gP9NaVC55fGodoVSKp5VewmZoM7Cf6T7niJ7JOkWptuBHZizv+LRB1ylGcbotGfmKV3XdwHRS+vWp/ue4+ss5FsxeREfSboFSDOM48AVwFbMLbunMwL4gcs1w3gtnbHliOiC5tem80YT1Fl4JZ33E9knSbdAhR8zXAbcBezDHMWOf1veh/nS7XvAMs0w2jIZYxZF/8OyMp03mqDOQsZmTIjskBVpBUwzjGHgO4amPQC8E7gZs6xhCeYus03AzzTDSOT5bz77M3BluF2bzhtJnYXCIyNdgWYYhmYYL2iG8Y+aYbxLM4zLNcP4oGYYPyzAhAvw66i2TSlVPOmZSaivr/+36DoLLpfrU+m4j8gtknSFGEfX9XYgul7wu9Nxn46Ojs9H2pWVlWMtLS2PpOM+IrdI0hViYtG/5qf8ZVq4zsKJOhi1tbVSZ6FASNIVYmLRm1uuTnXnBw4c+E2kbbVaCQaD70/1PURukqQrxMSao9pJl8OMFq6zEKnbK3UWCowkXSEmFv0yzaGUsk56ZpwmqLNwfar6FrlPkq4QE2vj1HnLl6eq4wnqLGxNVd8i90nSFWICuq4bwLGoQympret2uz88rs5CQyr6FflDkq4Qk3s9qr0mFR12dnZ+L9IuLS1ly5Yt/5KKfkX+kKQrxORaotpLku0sXGfBEflcU1MjjxUKkCRdISb3ZFS7SimlTXpmDI4dO/ZsdJ2FOXPmXJVMfyI/SdIVYnLNnHyZpgGeZDobX2ehubnZn0x/Ij9J0hViEuGXadGJMeGpXeE6Cyc+S52FwiVJV4ip7Y5quxPtROosiAhJukJMbVNU+5xEOvB4PCvG1VmYydsdiWlI0hViatGFaKoTeZl24MCBJyLtcJ2FG1MRmMhPknSFmNrzUW2NOHeSCNdZWBT5vGjRon1SZ6GwSdIVYgq6ro8BgahDcb1Ms9vtj0fXWViwYMENqYlM5CtJukJM782odlzTxjo6Ok4sH5Y6CwIk6QoRi+iXaefFetEEdRbuS2lUIi9J0hViek1R7TmxXjRBnYV/TWlUIi9J0hVies9EtS1KqYunu2B8nQWXy1Uo29eLaUjSFWIauq6PAMGoQ9NurXP06NHnoussVFdXX52m8ESekaQrRGw6otqXTndyZ2fnikh7/vz5UmdBnCBJV4jYbIlqL53qxPr6+q9F11moqan5RLqCEvlHkq4QsYl+mTZ3qhM7Ojo2RtrhOgs/S1tUIu9I0hUiNk9Fta1KqQmnjkmdBTEdSbpCxEDX9SGgP+rQhPUTpM6CmI4kXSFi54tqv2P8lxPUWdgrdRbEeJJ0hYhd9Mu0C8d/abfbfz2uzsL7MhGUyC+SdIWI3e+j2vPGf9nR0XFdpB2us/BqRqISeUWSrhCxi34pVqSUckU+uN3uj0idBRELSbpCxEjX9X5gIOrQ30QaPp/vu5F2aWmpIXUWxGSKsh2AEHlm7yij53XTTQ8979eU9vyct+YMHjt0TOosiJhI0hUiRprSapaxLLSLXRjmzuxu4A9H5x21cwvgBXZCdXX1NVkNVOQ0LVKUQwgxMU1pxcD3gb/V0IoMjIkHK0OgoRlGqXGloRvPT3iOKHiSdIWYgqa0EuBpYBVQHuNlg8AGQzeeSFdcIn/JizQhpvYj4ku4ALOARzSlrU5LRCKvSdIVYhKa0i7B3IgynoQbUQ58d9qzRMGRpCvE5DYCpdOeNbmlmtKmLAMpCo8kXSEmoCmtCvggYJ3m1KkUA59JRTxi5pCkK8TE1gJD0541tSLgumnPEgVFkq4QE3OSmp8PWwr6EDOIJF0hJhbKsX7EDCFJV4iJHQFSMYm9JwV9iBlEkq4QE/MCWpJ9HAd+moJYxAwiSVeICRi6cRx4CBhOsqsHUxCOmEEk6QoxuQdI/JnsGPBnQzf2pzAeMQNI0hViEoZudABf59QNKWPVC3wytRGJmUCSrhBT+zLwY04tXj6VEBAArjZ0Y0/aohJ5S5KuEFMwdMMAPgX8X8zRa+8kp45iJuZXgXpDN7ZMcp4ocFLaUYgYaUorxdyi507gAszntlbMlWuPAt82dKM9exGKfCBJV4gEaEqzAJXAcHimgxAxkaQrhBAZJM90hRAigyTpCiFEBknSFUKIDJKkK4QQGSRJVwghMkiSrhBCZND/B/qhqFludVQHAAAAAElFTkSuQmCC\n"
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "node_indices = torch.where(dataset[0].test_mask * dataset[0].y != 0)[0].tolist()\n",
    "from dig.xgraph.method.pgexplainer import PlotUtils\n",
    "plotutils = PlotUtils(dataset_name='ba_shapes', is_show=True)\n",
    "data = dataset[0]\n",
    "node_idx = node_indices[6]\n",
    "with torch.no_grad():\n",
    "    walks, masks, related_preds = \\\n",
    "        explainer(data.x, data.edge_index, node_idx=node_idx, y=data.y, top_k=5)\n",
    "\n",
    "explainer.visualization(data, edge_mask=masks[0], top_k=5, plot_utils=plotutils, node_idx=node_idx)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "### Metric results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Fidelity: 0.8181\n",
      "Fidelity_inv: 0.7251\n",
      "Sparsity: 0.2311\n"
     ]
    }
   ],
   "source": [
    "top_k = 5\n",
    "undirected_graph = True\n",
    "# --- Create data collector and explanation processor ---\n",
    "from dig.xgraph.evaluation import XCollector\n",
    "x_collector = XCollector()\n",
    "\n",
    "### Run explainer on the given model and dataset\n",
    "index = -1\n",
    "node_indices = torch.where(dataset[0].test_mask * dataset[0].y != 0)[0].tolist()\n",
    "top_k = top_k if not undirected_graph else top_k * 2\n",
    "\n",
    "for i, data in enumerate(dataloader):\n",
    "    for j, node_idx in enumerate(node_indices):\n",
    "        index += 1\n",
    "        data.to(device)\n",
    "\n",
    "        if torch.isnan(data.y[0].squeeze()):\n",
    "            continue\n",
    "\n",
    "        with torch.no_grad():\n",
    "            walks, masks, related_preds = \\\n",
    "                explainer(data.x, data.edge_index, node_idx=node_idx, y=data.y, top_k=top_k)\n",
    "            masks = [mask.detach() for mask in masks]\n",
    "        x_collector.collect_data(masks, related_preds)\n",
    "\n",
    "        # if you only have the edge masks without related_pred, please feed sparsity controlled mask to\n",
    "        # obtain the result: x_processor(data, masks, x_collector)\n",
    "        if index >= 99:\n",
    "            break\n",
    "\n",
    "    if index >= 20:\n",
    "        break\n",
    "\n",
    "\n",
    "print(f'Fidelity: {x_collector.fidelity:.4f}\\n'\n",
    "      f'Fidelity_inv: {x_collector.fidelity_inv:.4f}\\n'\n",
    "      f'Sparsity: {x_collector.sparsity:.4f}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "outputs": [],
   "source": [],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "outputs": [],
   "source": [],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  }
 ],
 "metadata": {
  "kernelspec": {
   "name": "python3",
   "language": "python",
   "display_name": "Python 3"
  },
  "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": 1
}