{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "0e78e9a7",
   "metadata": {},
   "source": [
    "# Results of anomaly detection using STGAT-MAD"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "21763f12-3752-43b5-8f89-e9e131d3c979",
   "metadata": {},
   "outputs": [],
   "source": [
    "from ice.anomaly_detection.datasets import AnomalyDetectionRiethTEP\n",
    "from ice.anomaly_detection.models import STGAT_MAD\n",
    "from sklearn.preprocessing import StandardScaler\n",
    "import numpy as np\n",
    "import pandas as pd"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "67244937-7e66-4f90-82e4-a9d4e634cd08",
   "metadata": {},
   "source": [
    "Download the dataset."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "c87e9ba4-4dc4-447b-91e1-df591298e756",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "16ec0e4a46ba4b91b67126185a088af7",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Reading data/rieth_tep/df.csv:   0%|          | 0/15330000 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "5b095d76964347f0b20c92658896fac1",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Reading data/rieth_tep/target.csv:   0%|          | 0/15330000 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "386ff939a156425a9d9323751e34de45",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Reading data/rieth_tep/train_mask.csv:   0%|          | 0/15330000 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "dataset = AnomalyDetectionRiethTEP()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f7a5479d-ac9e-47fc-819f-509cd0f1b960",
   "metadata": {},
   "source": [
    "Normalize the data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "355b8279-7b82-41ec-9ad0-5b45c0630fb8",
   "metadata": {},
   "outputs": [],
   "source": [
    "scaler = StandardScaler()\n",
    "dataset.df[dataset.train_mask] = scaler.fit_transform(dataset.df[dataset.train_mask])\n",
    "dataset.df[dataset.test_mask] = scaler.transform(dataset.df[dataset.test_mask])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a149d31c-e492-4c2b-b854-e83cdc752b87",
   "metadata": {},
   "source": [
    "Create the GNN model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "a975476c-1428-486e-b6c6-f947a096a0d9",
   "metadata": {},
   "outputs": [],
   "source": [
    "model = STGAT_MAD(\n",
    "    window_size=32, \n",
    "    num_epochs=30, \n",
    "    device='cuda',\n",
    "    verbose=True,\n",
    "    val_ratio=0.1,\n",
    "    save_checkpoints=True,\n",
    "    threshold_level=0.98\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "111cff61-98ff-4bff-862f-fba252f4f803",
   "metadata": {},
   "source": [
    "Load the checkpoint."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "7dedc39e-0157-47c4-b88b-4a00562da48f",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "model.load_checkpoint('stgat_anomaly_detection_epoch_30.tar')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "90f3470f-306e-4669-a1f8-ec5c15fccbe4",
   "metadata": {},
   "source": [
    "Evaluate the model on the test data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "4aad154d-1a4d-4a2a-be1e-08ba4f893ea2",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "e7cf8b9b2e6b42f3a337b33785e95c6b",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Creating sequence of samples:   0%|          | 0/10500 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Steps ...:   0%|          | 0/76125 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "{'accuracy': 0.861279659277504,\n",
       " 'true_positive_rate': [0.8352755],\n",
       " 'false_positive_rate': [0.019435206422018347]}"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "metrics = model.evaluate(dataset.df[dataset.test_mask], dataset.target[dataset.test_mask])\n",
    "metrics"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "395517a9-16b3-42fb-987b-adeb11e13a5e",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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": 5
}