{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# GraphBEAN: A Powerful Tool for Anomaly Detection on Bipartite Graphs\n",
    "<a target=\"_blank\" href=\"https://colab.research.google.com/github/AI4FinTech/FinTorch/blob/main/docs/tutorials/graphbean/graphbean.ipynb\">\n",
    "  <img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/>\n",
    "</a>\n",
    "\n",
    "## Introduction\n",
    "This tutorial provides a hands-on introduction to the GraphBEAN model, a novel graph neural network architecture for unsupervised anomaly detection on bipartite node-and-edge-attributed graphs. This model was originally presented in the paper \"Interaction-Focused Anomaly Detection on Bipartite Node-and-Edge-Attributed Graphs\" by Fathony et al. (2023), which we have implemented as part of our FinTorch project. Note that we generalized the concepts of GraphBEAN from bipartite networks to k-partite networks.\n",
    "\n",
    "GraphBEAN addresses the limitations of existing anomaly detection models, which typically focus on homogeneous graphs or neglect rich edge information. It leverages an autoencoder-like approach, employing a customized encoder-decoder structure to effectively encode both node and edge attributes, as well as the underlying graph structure, into low-dimensional latent representations. These representations are then used to reconstruct the original graph, and reconstruction errors are used to identify anomalous edges and nodes.\n",
    "\n",
    "This tutorial will guide you through the core concepts of GraphBEAN, demonstrating its usage with a practical example using the Elliptic dataset. You will learn how to:\n",
    "* Load and explore bipartite node-and-edge-attributed graph data.\n",
    "* Define and train a GraphBEAN model using PyTorch Lightning.\n",
    "* Analyze and interpret anomaly detection results.\n",
    "\n",
    "This tutorial will enable you to effectively apply GraphBEAN to diverse applications involving bipartite graphs, such as fraud detection in financial transactions, malicious activity detection in network security, or anomaly detection in user-item interaction networks."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Install FinTorch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install fintorch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "\n",
    "# Installation of PyTorch Geometric and dependencies based on detected versions\n",
    "def install_pyg_and_dependencies():\n",
    "  !pip install pyg-lib -f https://data.pyg.org/whl/torch-{torch.__version__}.html\n",
    "  !pip install torch-scatter torch-sparse -f https://data.pyg.org/whl/torch-{torch.__version__}.html\n",
    "\n",
    "# Detect PyTorch version\n",
    "if torch.__version__ >= \"1.13.0\":\n",
    "  print(\"PyTorch version 1.13.0 or newer detected. Installing PyG and dependencies...\")\n",
    "  install_pyg_and_dependencies()\n",
    "else:\n",
    "  print(\"PyTorch version is older than 1.13.0. PyG might not work correctly. Please upgrade PyTorch or use the pip install torch_geometric method.\")\n",
    "  \n",
    "\n",
    "# Verify installation\n",
    "try:\n",
    "  import torch_geometric\n",
    "  print(f\"PyTorch Geometric successfully installed. Version: {torch_geometric.__version__}\")\n",
    "except ImportError:\n",
    "  print(\"PyTorch Geometric not found. Installation might have failed.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Code"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The tutorial kicks off by importing the necessary libraries: PyTorch Lightning for streamlined training, PyTorch Geometric for powerful graph convolution operations, and FinTorch modules for loading the Elliptic dataset and utilizing the GraphBEAN model. We then create an instance of the EllipticppDataModule, clearly defining the dataset's bipartite structure with \"wallets\" and \"transactions\" as node types and \"to\" as the edge type. This module takes care of data loading, splitting, and generating data loaders for efficient training."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import lightning as L\n",
    "from torch_geometric.nn.conv import TransformerConv\n",
    "\n",
    "from fintorch.datasets.ellipticpp import EllipticppDataModule\n",
    "from fintorch.models.graph.graphbean.graphBEAN import GraphBEANModule"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next, we prepare the dataset by initializing the data module and displaying its structure, revealing the node types, their attributes, and edge connections. We then delve deeper into the dataset's structure by retrieving its metadata, gaining a high-level understanding of the relationships within the bipartite graph."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "# We use an example data module from the elliptic dataset which is bipartite\n",
    "data_module = EllipticppDataModule((\"wallets\", \"to\", \"transactions\"))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Start download from HuggingFace...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "  0%|          | 0/7 [00:00<?, ?it/s]"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 7/7 [00:01<00:00,  6.68it/s]\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "HeteroData(\n",
       "  wallets={\n",
       "    x=[1268260, 55],\n",
       "    y=[1268260],\n",
       "    train_mask=[1268260],\n",
       "    val_mask=[1268260],\n",
       "    test_mask=[1268260],\n",
       "  },\n",
       "  transactions={\n",
       "    x=[203769, 182],\n",
       "    y=[203769],\n",
       "    train_mask=[203769],\n",
       "    val_mask=[203769],\n",
       "    test_mask=[203769],\n",
       "  },\n",
       "  (transactions, to, transactions)={ edge_index=[2, 234355] },\n",
       "  (transactions, to, wallets)={ edge_index=[2, 837124] },\n",
       "  (wallets, to, transactions)={ edge_index=[2, 477117] },\n",
       "  (wallets, to, wallets)={ edge_index=[2, 2868964] }\n",
       ")"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "data_module.setup()\n",
    "\n",
    "data_module.dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(['wallets', 'transactions'],\n",
       " [('transactions', 'to', 'transactions'),\n",
       "  ('transactions', 'to', 'wallets'),\n",
       "  ('wallets', 'to', 'transactions'),\n",
       "  ('wallets', 'to', 'wallets')])"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "data_module.dataset.metadata()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We proceed to create a GraphBEAN model instance, specifying the node and edge types we use in the convolution, learning rate, convolution type, and the number of layers in the encoder, decoder, and hidden layers. We also set up a PyTorch Lightning trainer, defining the maximum number of training epochs and enabling GPU acceleration for faster training."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "GPU available: True (cuda), used: True\n",
      "TPU available: False, using: 0 TPU cores\n",
      "IPU available: False, using: 0 IPUs\n",
      "HPU available: False, using: 0 HPUs\n"
     ]
    }
   ],
   "source": [
    "mapping = dict()\n",
    "for key in data_module.dataset.metadata()[0]:\n",
    "    mapping[key] = data_module.dataset[key].x.shape[1]\n",
    "\n",
    "# Create an instance of the GraphBEANModule\n",
    "module = GraphBEANModule(\n",
    "    (\"wallets\", \"to\", \"transactions\"),\n",
    "    edge_types=[(\"wallets\", \"to\", \"transactions\"),\n",
    "                (\"transactions\", \"to\", \"wallets\")],\n",
    "    mapping=mapping,\n",
    "    learning_rate=0.001,\n",
    "    encoder_layers=5,\n",
    "    decoder_layers=5,\n",
    "    hidden_layers=50,\n",
    ")\n",
    "\n",
    "# Create a PyTorch Lightning Trainer and train the module\n",
    "trainer = L.Trainer(max_epochs=1, accelerator=\"auto\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Finally, we train the model by using the PyTorch Lightning trainer with the data loaders provided by the EllipticDataModule. The trainer automatically manages the training loop, logging, and progress tracking, ensuring a smooth and efficient training experience."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "You are using a CUDA device ('NVIDIA GeForce RTX 3090') that has Tensor Cores. To properly utilize them, you should set `torch.set_float32_matmul_precision('medium' | 'high')` which will trade-off precision for performance. For more details, read https://pytorch.org/docs/stable/generated/torch.set_float32_matmul_precision.html#torch.set_float32_matmul_precision\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Start download from HuggingFace...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 7/7 [00:00<00:00,  9.29it/s]\n",
      "LOCAL_RANK: 0 - CUDA_VISIBLE_DEVICES: [0]\n",
      "/home/marcel/Documents/research/FinTorch/.conda/lib/python3.11/site-packages/lightning/pytorch/utilities/model_summary/model_summary.py:454: A layer with UninitializedParameter was found. Thus, the total number of parameters detected may be inaccurate.\n",
      "\n",
      "  | Name      | Type                      | Params\n",
      "--------------------------------------------------------\n",
      "0 | accuracy  | MulticlassAccuracy        | 0     \n",
      "1 | f1        | MulticlassF1Score         | 0     \n",
      "2 | recall    | MulticlassRecall          | 0     \n",
      "3 | precision | MulticlassPrecision       | 0     \n",
      "4 | confmat   | MulticlassConfusionMatrix | 0     \n",
      "5 | aucroc    | MulticlassAUROC           | 0     \n",
      "6 | model     | GraphBEAN                 | 165 K \n",
      "--------------------------------------------------------\n",
      "165 K     Trainable params\n",
      "0         Non-trainable params\n",
      "165 K     Total params\n",
      "0.664     Total estimated model params size (MB)\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "a8a217e8bce44e34bab185f5cd479324",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Sanity Checking: |          | 0/? [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "0088d74e26d144da9d0f475d0c084bdf",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Training: |          | 0/? [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "e6e4122fe80d40fe95bcfcc82600d9b1",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: |          | 0/? [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "`Trainer.fit` stopped: `max_epochs=1` reached.\n"
     ]
    }
   ],
   "source": [
    "# Train the module using the dataloaders\n",
    "trainer.fit(module, datamodule=data_module)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Conclusion\n",
    "This concludes our tutorial on GraphBEAN, a powerful tool for anomaly detection on bipartite graphs. You've learned how to load and explore bipartite node-and-edge-attributed graphs using the Elliptic dataset and FinTorch's data module. You've also gained experience in defining and configuring a GraphBEAN model, customizing its parameters to suit your specific needs. Finally, you've trained the GraphBEAN model using PyTorch Lightning.\n",
    "\n",
    "This tutorial provided a foundation for applying GraphBEAN to various real-world applications involving bipartite graphs. You can now adapt these concepts and further explore GraphBEAN's capabilities. Experiment with different datasets, explore hyperparameter tuning, and delve into advanced anomaly detection techniques.\n",
    "\n",
    "Remember, GraphBEAN is a powerful tool for uncovering hidden anomalies in complex bipartite data, potentially leading to valuable insights in domains like fraud detection, network security, or user behavior analysis. We encourage you to continue exploring and implementing this model, pushing the boundaries of anomaly detection in bipartite graphs."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.11.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}