{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "55ef77b9-9320-44db-862b-088d7af03112",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\u001b[34m\u001b[1mwandb\u001b[0m: Currently logged in as: \u001b[33mkmilian\u001b[0m (\u001b[33mepos\u001b[0m). Use \u001b[1m`wandb login --relogin`\u001b[0m to force relogin\n",
      "\u001b[34m\u001b[1mwandb\u001b[0m: \u001b[33mWARNING\u001b[0m If you're specifying your api key in code, ensure this code is not shared publicly.\n",
      "\u001b[34m\u001b[1mwandb\u001b[0m: \u001b[33mWARNING\u001b[0m Consider setting the WANDB_API_KEY environment variable, or running `wandb login` from the command line.\n",
      "\u001b[34m\u001b[1mwandb\u001b[0m: Appending key for api.wandb.ai to your netrc file: /Users/krystynamilian/.netrc\n"
     ]
    }
   ],
   "source": [
    "import pandas as pd\n",
    "from obspy.core.event import read_events\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "import seisbench.models as sbm\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "\n",
    "import seisbench.data as sbd\n",
    "import seisbench.generate as sbg\n",
    "import seisbench.models as sbm\n",
    "from seisbench.util import worker_seeding\n",
    "import numpy as np\n",
    "from torch.utils.data import DataLoader\n",
    "\n",
    "import wandb\n",
    "import os\n",
    "import sys\n",
    "\n",
    "from pathlib import Path\n",
    "project_path = str(Path.cwd().parent)\n",
    "sys.path.append(project_path)\n",
    "from scripts import train"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "694ceb35-d2f3-4654-a6af-d84e494a4660",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "wandb version 0.15.4 is available!  To upgrade, please run:\n",
       " $ pip install wandb --upgrade"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "Tracking run with wandb version 0.15.3"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "Run data is saved locally in <code>/Users/krystynamilian/Documents/praca/Cyfronet/epos/ai/repo/demo_scripts/notebooks/wandb/run-20230704_110602-ufltoqra</code>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "Syncing run <strong><a href='https://wandb.ai/epos/training_seisbench_models_on_igf_data/runs/ufltoqra' target=\"_blank\">polished-totem-255</a></strong> to <a href='https://wandb.ai/epos/training_seisbench_models_on_igf_data' target=\"_blank\">Weights & Biases</a> (<a href='https://wandb.me/run' target=\"_blank\">docs</a>)<br/>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       " View project at <a href='https://wandb.ai/epos/training_seisbench_models_on_igf_data' target=\"_blank\">https://wandb.ai/epos/training_seisbench_models_on_igf_data</a>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       " View run at <a href='https://wandb.ai/epos/training_seisbench_models_on_igf_data/runs/ufltoqra' target=\"_blank\">https://wandb.ai/epos/training_seisbench_models_on_igf_data/runs/ufltoqra</a>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\u001b[34m\u001b[1mwandb\u001b[0m:   1 of 1 files downloaded.  \n"
     ]
    }
   ],
   "source": [
    "run = wandb.init(project=\"training_seisbench_models_on_igf_data\", entity=\"epos\", mode=\"online\")\n",
    "artifact = run.use_artifact('epos/model-registry/phasenet_p:v0', type='model')\n",
    "artifact_dir = artifact.download()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0589f4aa-e5a4-485e-9213-a4696c38a60c",
   "metadata": {},
   "source": [
    "# Load model, get data loaders"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "9935a11d-e8b5-4019-aafa-f031f1024a71",
   "metadata": {},
   "outputs": [],
   "source": [
    "model = train.load_model()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "922e8062-e503-4958-b192-b274453d64c3",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'./artifacts/model:v4'"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "artifact_dir"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "7e100ac5-af2b-4d20-95f3-4f0a6a4fc0ad",
   "metadata": {},
   "outputs": [],
   "source": [
    "fname = artifact_dir + \"/\" + os.listdir(artifact_dir)[0]\n",
    "# fname \n",
    "# fname = \"../models/PhaseNet_pretrained_on_iquique_finetuned_on_igf_1.pt\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "ffdf1280-2b25-49a3-9e6b-38a85a9c5501",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "PhaseNet(\n",
       "  (inc): Conv1d(3, 8, kernel_size=(7,), stride=(1,), padding=same)\n",
       "  (in_bn): BatchNorm1d(8, eps=0.001, momentum=0.1, affine=True, track_running_stats=True)\n",
       "  (down_branch): ModuleList(\n",
       "    (0): ModuleList(\n",
       "      (0): Conv1d(8, 8, kernel_size=(7,), stride=(1,), padding=same, bias=False)\n",
       "      (1): BatchNorm1d(8, eps=0.001, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (2): Conv1d(8, 8, kernel_size=(7,), stride=(4,), padding=(3,), bias=False)\n",
       "      (3): BatchNorm1d(8, eps=0.001, momentum=0.1, affine=True, track_running_stats=True)\n",
       "    )\n",
       "    (1): ModuleList(\n",
       "      (0): Conv1d(8, 16, kernel_size=(7,), stride=(1,), padding=same, bias=False)\n",
       "      (1): BatchNorm1d(16, eps=0.001, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (2): Conv1d(16, 16, kernel_size=(7,), stride=(4,), bias=False)\n",
       "      (3): BatchNorm1d(16, eps=0.001, momentum=0.1, affine=True, track_running_stats=True)\n",
       "    )\n",
       "    (2): ModuleList(\n",
       "      (0): Conv1d(16, 32, kernel_size=(7,), stride=(1,), padding=same, bias=False)\n",
       "      (1): BatchNorm1d(32, eps=0.001, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (2): Conv1d(32, 32, kernel_size=(7,), stride=(4,), bias=False)\n",
       "      (3): BatchNorm1d(32, eps=0.001, momentum=0.1, affine=True, track_running_stats=True)\n",
       "    )\n",
       "    (3): ModuleList(\n",
       "      (0): Conv1d(32, 64, kernel_size=(7,), stride=(1,), padding=same, bias=False)\n",
       "      (1): BatchNorm1d(64, eps=0.001, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (2): Conv1d(64, 64, kernel_size=(7,), stride=(4,), bias=False)\n",
       "      (3): BatchNorm1d(64, eps=0.001, momentum=0.1, affine=True, track_running_stats=True)\n",
       "    )\n",
       "    (4): ModuleList(\n",
       "      (0): Conv1d(64, 128, kernel_size=(7,), stride=(1,), padding=same, bias=False)\n",
       "      (1): BatchNorm1d(128, eps=0.001, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (2-3): 2 x None\n",
       "    )\n",
       "  )\n",
       "  (up_branch): ModuleList(\n",
       "    (0): ModuleList(\n",
       "      (0): ConvTranspose1d(128, 64, kernel_size=(7,), stride=(4,), bias=False)\n",
       "      (1): BatchNorm1d(64, eps=0.001, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (2): Conv1d(128, 64, kernel_size=(7,), stride=(1,), padding=same, bias=False)\n",
       "      (3): BatchNorm1d(64, eps=0.001, momentum=0.1, affine=True, track_running_stats=True)\n",
       "    )\n",
       "    (1): ModuleList(\n",
       "      (0): ConvTranspose1d(64, 32, kernel_size=(7,), stride=(4,), bias=False)\n",
       "      (1): BatchNorm1d(32, eps=0.001, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (2): Conv1d(64, 32, kernel_size=(7,), stride=(1,), padding=same, bias=False)\n",
       "      (3): BatchNorm1d(32, eps=0.001, momentum=0.1, affine=True, track_running_stats=True)\n",
       "    )\n",
       "    (2): ModuleList(\n",
       "      (0): ConvTranspose1d(32, 16, kernel_size=(7,), stride=(4,), bias=False)\n",
       "      (1): BatchNorm1d(16, eps=0.001, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (2): Conv1d(32, 16, kernel_size=(7,), stride=(1,), padding=same, bias=False)\n",
       "      (3): BatchNorm1d(16, eps=0.001, momentum=0.1, affine=True, track_running_stats=True)\n",
       "    )\n",
       "    (3): ModuleList(\n",
       "      (0): ConvTranspose1d(16, 8, kernel_size=(7,), stride=(4,), bias=False)\n",
       "      (1): BatchNorm1d(8, eps=0.001, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (2): Conv1d(16, 8, kernel_size=(7,), stride=(1,), padding=same, bias=False)\n",
       "      (3): BatchNorm1d(8, eps=0.001, momentum=0.1, affine=True, track_running_stats=True)\n",
       "    )\n",
       "  )\n",
       "  (out): Conv1d(8, 2, kernel_size=(1,), stride=(1,), padding=same)\n",
       "  (softmax): Softmax(dim=1)\n",
       ")"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model.load_state_dict(torch.load(fname))\n",
    "model.eval()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "4656e5fe-bb7b-4564-923a-8af385eb312d",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "train (12444, 17) 100\n",
      "using random window\n",
      "dev (2773, 17) 100\n",
      "using random window\n",
      "test (2785, 17) 100\n",
      "using random window\n"
     ]
    }
   ],
   "source": [
    "train_gen, dev_gen, test_gen = train.get_data_generators()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "dfdaa857-9f2b-419f-8f28-9ce8f4b11a0f",
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot_sample(sample, model, i): \n",
    "    fig = plt.figure(figsize=(15, 10))\n",
    "    fig.suptitle(\"Predictions for test sample: \" + str(i))\n",
    "    axs = fig.subplots(2, 1, sharex=True, gridspec_kw={\"hspace\": 0, \"height_ratios\": [3,  2]})\n",
    "    axs[0].plot(sample[\"X\"][0].T, label='x')\n",
    "    plt.legend()\n",
    "    axs[1].plot(sample[\"y\"][0].T, label='y')\n",
    "    \n",
    "    model.eval()  # close the model for evaluation\n",
    "    \n",
    "    with torch.no_grad():\n",
    "        pred = model(torch.tensor(sample[\"X\"], device=model.device).unsqueeze(0))  # Add a fake batch dimension\n",
    "        pred = pred[0].cpu().numpy()\n",
    "        pick_idx = np.argmax(pred[0])\n",
    "        print(pred.shape)\n",
    "    \n",
    "        axs[1].plot(pred[0], label='pred', color='orange')\n",
    "        plt.legend()\n",
    "        \n",
    "        plt.show()\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "26ee3888-8138-4d02-a5f8-a899b8ca6436",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "No artists with labels found to put in legend.  Note that artists whose label start with an underscore are ignored when legend() is called with no argument.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(2, 3001)\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 120492 characters]",
      "text/plain": [
       "<Figure size 1500x1000 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "idx = np.random.randint(len(test_gen))\n",
    "sample = test_gen[idx]\n",
    "plot_sample(sample, model, idx)\n",
    "\n"
   ]
  }
 ],
 "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.9.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
