examples/xgraph/flowx.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "source": [
    "# GradCAM on ClinTox dataset for GIN"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%% md\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "collapsed": true
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/data/shurui.gui/anaconda3/envs/torch_110/lib/python3.8/site-packages/tqdm/auto.py:22: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
      "  from .autonotebook import tqdm as notebook_tqdm\n"
     ]
    }
   ],
   "source": [
    "import os\n",
    "import os.path as osp\n",
    "\n",
    "import torch\n",
    "from torch.utils.data import random_split\n",
    "from torch_geometric.data import download_url, extract_zip\n",
    "from torch_geometric.loader import DataLoader\n",
    "\n",
    "from dig.xgraph.dataset import MoleculeDataset\n",
    "from dig.xgraph.utils.compatibility import compatible_state_dict\n",
    "from dig.xgraph.utils.init import fix_random_seed\n",
    "\n",
    "device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')"
   ]
  },
  {
   "cell_type": "markdown",
   "source": [
    "### Load dataset"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%% md\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Downloading https://deepchemdata.s3-us-west-1.amazonaws.com/datasets/clintox.csv.gz\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Extracting datasets/clintox/raw/clintox.csv.gz\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Processing...\n",
      "[16:00:42] Explicit valence for atom # 0 N, 5, is greater than permitted\n",
      "[16:00:43] Can't kekulize mol.  Unkekulized atoms: 9\n",
      "[16:00:43] Explicit valence for atom # 10 N, 4, is greater than permitted\n",
      "[16:00:43] Explicit valence for atom # 10 N, 4, is greater than permitted\n",
      "[16:00:43] Can't kekulize mol.  Unkekulized atoms: 4\n",
      "[16:00:43] Can't kekulize mol.  Unkekulized atoms: 4\n",
      "Done!\n"
     ]
    }
   ],
   "source": [
    "def split_dataset(dataset, dataset_split=[0.8, 0.1, 0.1]):\n",
    "    dataset_len = len(dataset)\n",
    "    dataset_split = [int(dataset_len * dataset_split[0]),\n",
    "                     int(dataset_len * dataset_split[1]),\n",
    "                     0]\n",
    "    dataset_split[2] = dataset_len - dataset_split[0] - dataset_split[1]\n",
    "    train_set, val_set, test_set = random_split(dataset, dataset_split)\n",
    "\n",
    "    return {'train': train_set, 'val': val_set, 'test': test_set}\n",
    "\n",
    "\n",
    "fix_random_seed(123)\n",
    "dataset = MoleculeDataset('datasets', 'clintox')\n",
    "dataset.data.x = dataset.data.x.to(torch.float32)\n",
    "dataset.data.y = dataset.data.y[:, 0]\n",
    "dim_node = dataset.num_node_features\n",
    "dim_edge = dataset.num_edge_features\n",
    "num_targets = dataset.num_classes\n",
    "num_classes = 2\n",
    "\n",
    "splitted_dataset = split_dataset(dataset)\n",
    "dataloader = DataLoader(splitted_dataset['test'], batch_size=1, shuffle=False)"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "### Load model and checkpoints"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%% md\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Downloading https://github.com/divelab/DIG_storage/raw/main/xgraph/checkpoints.zip\n",
      "Extracting ./checkpoints.zip\n"
     ]
    },
    {
     "data": {
      "text/plain": "<All keys matched successfully>"
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from dig.xgraph.models import GIN_3l\n",
    "\n",
    "\n",
    "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",
    "\n",
    "model = GIN_3l(model_level='graph', dim_node=dim_node, dim_hidden=300, num_classes=num_classes)\n",
    "model.to(device)\n",
    "check_checkpoints()\n",
    "ckpt_path = osp.join('checkpoints', 'clintox', 'GIN_3l', '0', 'GIN_3l_best.ckpt')\n",
    "state_dict = compatible_state_dict(torch.load(ckpt_path, map_location='cpu')['state_dict'])\n",
    "model.load_state_dict(state_dict)"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "### Display example output"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%% md\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([[-1.0756,  1.0493]], device='cuda:0', grad_fn=<AddmmBackward0>)\n"
     ]
    }
   ],
   "source": [
    "data = list(dataloader)[0].to(device)\n",
    "out = model(data.x, data.edge_index)\n",
    "print(out)"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "### Load the explainer"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%% md\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "outputs": [],
   "source": [
    "from dig.xgraph.method import FlowX\n",
    "\n",
    "explainer = FlowX(model, explain_graph=True)"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "### Setup for evaluation"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%% md\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "outputs": [],
   "source": [
    "# --- Set the Sparsity to 0.5 ---\n",
    "sparsity = 0.5\n",
    "\n",
    "# --- Create data collector and explanation processor ---\n",
    "from dig.xgraph.evaluation import XCollector\n",
    "\n",
    "x_collector = XCollector(sparsity)\n",
    "# x_processor = ExplanationProcessor(model=model, device=device)"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "### Run explainer on the given model and dataset"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%% md\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "explain graph line 1278\n",
      "explain graph line 1153\n",
      "explain graph line 25\n",
      "explain graph line 721\n",
      "explain graph line 292\n",
      "explain graph line 186\n",
      "explain graph line 1402\n",
      "explain graph line 1095\n",
      "explain graph line 1093\n",
      "explain graph line 221\n",
      "explain graph line 404\n",
      "explain graph line 472\n",
      "explain graph line 1180\n",
      "explain graph line 419\n",
      "explain graph line 1117\n",
      "explain graph line 467\n",
      "explain graph line 942\n",
      "explain graph line 401\n",
      "explain graph line 760\n",
      "explain graph line 257\n",
      "explain graph line 161\n",
      "explain graph line 656\n",
      "explain graph line 1350\n",
      "explain graph line 744\n",
      "explain graph line 56\n",
      "explain graph line 835\n",
      "explain graph line 1383\n",
      "explain graph line 1063\n",
      "explain graph line 18\n",
      "explain graph line 174\n",
      "explain graph line 1261\n",
      "explain graph line 1341\n",
      "explain graph line 973\n",
      "explain graph line 1203\n",
      "explain graph line 1280\n",
      "explain graph line 671\n",
      "explain graph line 1303\n",
      "explain graph line 1311\n",
      "explain graph line 1214\n",
      "explain graph line 141\n",
      "explain graph line 952\n",
      "explain graph line 881\n",
      "explain graph line 1213\n",
      "explain graph line 7\n",
      "explain graph line 130\n",
      "explain graph line 1451\n",
      "explain graph line 293\n",
      "explain graph line 73\n",
      "explain graph line 677\n",
      "explain graph line 892\n",
      "explain graph line 868\n",
      "explain graph line 1138\n",
      "explain graph line 134\n",
      "explain graph line 1328\n",
      "explain graph line 1370\n",
      "explain graph line 251\n",
      "explain graph line 585\n",
      "explain graph line 580\n",
      "explain graph line 665\n",
      "explain graph line 302\n",
      "explain graph line 731\n",
      "explain graph line 139\n",
      "explain graph line 1025\n",
      "explain graph line 851\n",
      "explain graph line 802\n",
      "explain graph line 1225\n",
      "explain graph line 910\n",
      "explain graph line 1306\n",
      "explain graph line 521\n",
      "explain graph line 458\n",
      "explain graph line 258\n",
      "explain graph line 625\n",
      "explain graph line 395\n",
      "explain graph line 111\n",
      "explain graph line 1296\n",
      "explain graph line 345\n",
      "explain graph line 1355\n",
      "explain graph line 382\n",
      "explain graph line 1240\n",
      "explain graph line 1134\n",
      "explain graph line 208\n",
      "explain graph line 336\n",
      "explain graph line 1034\n",
      "explain graph line 563\n",
      "explain graph line 396\n",
      "explain graph line 735\n",
      "explain graph line 1228\n",
      "explain graph line 555\n",
      "explain graph line 703\n",
      "explain graph line 370\n",
      "explain graph line 578\n",
      "explain graph line 6\n",
      "explain graph line 1452\n",
      "explain graph line 1135\n",
      "explain graph line 272\n",
      "explain graph line 394\n",
      "explain graph line 1252\n",
      "explain graph line 1065\n",
      "explain graph line 329\n",
      "explain graph line 1424\n"
     ]
    }
   ],
   "source": [
    "for index, data in enumerate(dataloader):\n",
    "    print(f'explain graph line {dataloader.dataset.indices[index] + 2}')\n",
    "    data.to(device)\n",
    "\n",
    "    if torch.isnan(data.y[0].squeeze()):\n",
    "        continue\n",
    "\n",
    "    walks, masks, related_preds = explainer(data.x, data.edge_index, sparsity=sparsity, num_classes=num_classes)\n",
    "\n",
    "    x_collector.collect_data(masks, related_preds, data.y[0].squeeze().long().item())\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",
    "\n",
    "    if index >= 99:\n",
    "        break"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "### Output metrics evaluation results"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%% md\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Fidelity: 0.5917\n",
      "Fidelity_inv: -0.0093\n",
      "Sparsity: 0.5000\n"
     ]
    }
   ],
   "source": [
    "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}')"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "markdown",
   "source": [],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%% md\n"
    }
   }
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}