{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "05600e0a-8b4b-42ab-96d5-9d6eb2c72102",
   "metadata": {},
   "source": [
    "# Check model performance depending on station\n",
    "\n",
    "#### Samples generated with fixed window"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "1732f721-6b13-4fb3-8755-dc63cb255285",
   "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",
    "from pathlib import Path\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": "1a322214-ba56-4651-966d-f9afdcfecbc5",
   "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_110541-ir19n1xv</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/demo_scripts-notebooks/runs/ir19n1xv' target=\"_blank\">eager-forest-7</a></strong> to <a href='https://wandb.ai/epos/demo_scripts-notebooks' 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/demo_scripts-notebooks' target=\"_blank\">https://wandb.ai/epos/demo_scripts-notebooks</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/demo_scripts-notebooks/runs/ir19n1xv' target=\"_blank\">https://wandb.ai/epos/demo_scripts-notebooks/runs/ir19n1xv</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"
     ]
    },
    {
     "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": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model = train.load_model()\n",
    "\n",
    "run = wandb.init()\n",
    "artifact = run.use_artifact('epos/training_seisbench_models_on_igf_data/model:v113', type='model')\n",
    "artifact_dir = artifact.download()\n",
    "fname = artifact_dir + \"/\" + os.listdir(artifact_dir)[0]\n",
    "\n",
    "model.load_state_dict(torch.load(fname))\n",
    "model.eval()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "75cadd4b-7a0e-44cd-acc6-fef26f2cf551",
   "metadata": {},
   "outputs": [],
   "source": [
    "data_path = '../../../data/igf/seisbench_format'\n",
    "sampling_rate = 100\n",
    "data = sbd.WaveformDataset(data_path, sampling_rate=sampling_rate)\n",
    "data.filter(data.metadata.trace_Pg_arrival_sample.notna())\n",
    "\n",
    "pick_mae = train.PickMAE(sampling_rate)\n",
    "splits = ['train', 'dev', 'test']\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ccf454fe-e55a-47bc-882a-3e39395147a5",
   "metadata": {},
   "source": [
    "### Calculate overall model performance"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "4c7b4137-4599-4599-9e3e-064145bfccb4",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "\n",
      "Model resutls for train set\n",
      "\n",
      "Fixed window\n",
      "train (12444, 17) 100\n",
      "Test avg loss: 0.025157\n",
      "Test avg mae: 0.047488\n",
      "\n",
      "\n",
      "\n",
      "Model resutls for dev set\n",
      "\n",
      "Fixed window\n",
      "dev (2773, 17) 100\n",
      "Test avg loss: 0.025309\n",
      "Test avg mae: 0.051242\n",
      "\n",
      "\n",
      "\n",
      "Model resutls for test set\n",
      "\n",
      "Fixed window\n",
      "test (2785, 17) 100\n",
      "Test avg loss: 0.025555\n",
      "Test avg mae: 0.047198\n",
      "\n"
     ]
    }
   ],
   "source": [
    "for split in splits: \n",
    "    print(f\"\\n\\nModel resutls for {split} set\")\n",
    "    print(\"\\nFixed window\")\n",
    "    # the results depend on random selection of a window around a pick\n",
    "    gen = train.get_data_generator(split=split, station=None, sampling_rate=sampling_rate, path=data_path, window='fixed')\n",
    "\n",
    "    data_loader = DataLoader(gen, batch_size=256, shuffle=False, num_workers=0,\n",
    "                                    worker_init_fn=worker_seeding)\n",
    "        \n",
    "    test_loss, test_mae = train.test_one_epoch(model, data_loader, pick_mae, wandb_log=False)\n",
    "    "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a361bc6a-2dac-4e27-9417-3abe482161a5",
   "metadata": {},
   "source": [
    "## Check # frames per station in each set"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "b23b71b7-ba71-4e86-8978-5a70d4fbbdfb",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<Axes: title={'center': 'Frames per station'}, xlabel='station_code'>"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 46184 characters]",
      "text/plain": [
       "<Figure size 1500x300 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "frames_per_station = []\n",
    "for split in splits: \n",
    "    frames_per_station.append(data.get_split(split).metadata.groupby('station_code').count()['index'])\n",
    "    \n",
    "frames_per_station = pd.DataFrame(frames_per_station, index=splits).transpose()\n",
    "frames_per_station.plot(kind='bar', figsize=(15,3), title='Frames per station')\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "acd81c2b-318f-45ea-b639-bf1ed26b0c10",
   "metadata": {},
   "source": [
    "## Calculate MAE per station for train/dev/test set"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "66992bf3-c49a-4885-8ab1-5b372ad65202",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "train\n",
      "BRDW\n",
      "train (12444, 17) 100\n",
      "Test avg loss: 0.029592\n",
      "Test avg mae: 0.089062\n",
      "\n",
      "GROD\n",
      "train (12444, 17) 100\n",
      "Test avg loss: 0.025606\n",
      "Test avg mae: 0.057617\n",
      "\n",
      "GUZI\n",
      "train (12444, 17) 100\n",
      "Test avg loss: 0.024516\n",
      "Test avg mae: 0.040501\n",
      "\n",
      "JEDR\n",
      "train (12444, 17) 100\n",
      "Test avg loss: 0.024908\n",
      "Test avg mae: 0.044932\n",
      "\n",
      "MOSK2\n",
      "train (12444, 17) 100\n",
      "Test avg loss: 0.024886\n",
      "Test avg mae: 0.040648\n",
      "\n",
      "NWLU\n",
      "train (12444, 17) 100\n",
      "Test avg loss: 0.024519\n",
      "Test avg mae: 0.032965\n",
      "\n",
      "PCHB\n",
      "train (12444, 17) 100\n",
      "Test avg loss: 0.024587\n",
      "Test avg mae: 0.042854\n",
      "\n",
      "PPOL\n",
      "train (12444, 17) 100\n",
      "Test avg loss: 0.026397\n",
      "Test avg mae: 0.076008\n",
      "\n",
      "RUDN\n",
      "train (12444, 17) 100\n",
      "Test avg loss: 0.025373\n",
      "Test avg mae: 0.047411\n",
      "\n",
      "RYNR\n",
      "train (12444, 17) 100\n",
      "Test avg loss: 0.025934\n",
      "Test avg mae: 0.066802\n",
      "\n",
      "RZEC\n",
      "train (12444, 17) 100\n",
      "Test avg loss: 0.023816\n",
      "Test avg mae: 0.029310\n",
      "\n",
      "SGOR\n",
      "train (12444, 17) 100\n",
      "Test avg loss: 0.024461\n",
      "Test avg mae: 0.034385\n",
      "\n",
      "TRBC2\n",
      "train (12444, 17) 100\n",
      "Test avg loss: 0.025840\n",
      "Test avg mae: 0.046694\n",
      "\n",
      "TRN2\n",
      "train (12444, 17) 100\n",
      "Test avg loss: 0.025408\n",
      "Test avg mae: 0.051002\n",
      "\n",
      "TRZS\n",
      "train (12444, 17) 100\n",
      "Test avg loss: 0.025081\n",
      "Test avg mae: 0.043014\n",
      "\n",
      "ZMST\n",
      "train (12444, 17) 100\n",
      "Test avg loss: 0.025119\n",
      "Test avg mae: 0.049605\n",
      "\n",
      "LUBW\n",
      "train (12444, 17) 100\n",
      "Test avg loss: 0.029460\n",
      "Test avg mae: 0.095455\n",
      "\n",
      "DWOL\n",
      "train (12444, 17) 100\n",
      "Test avg loss: 0.024191\n",
      "Test avg mae: 0.026768\n",
      "\n",
      "LUBZ\n",
      "train (12444, 17) 100\n",
      "Test avg loss: 0.032006\n",
      "Test avg mae: 0.180000\n",
      "\n",
      "ZUKW2\n",
      "train (12444, 17) 100\n",
      "Test avg loss: 0.024766\n",
      "Test avg mae: 0.032847\n",
      "\n",
      "DABR\n",
      "train (12444, 17) 100\n",
      "Test avg loss: 0.024297\n",
      "Test avg mae: 0.030349\n",
      "\n",
      "PEKW2\n",
      "train (12444, 17) 100\n",
      "Test avg loss: 0.025126\n",
      "Test avg mae: 0.044390\n",
      "\n",
      "KRZY\n",
      "train (12444, 17) 100\n",
      "Test avg loss: 0.025487\n",
      "Test avg mae: 0.060000\n",
      "\n",
      "OBIS\n",
      "train (12444, 17) 100\n",
      "Test avg loss: 0.024071\n",
      "Test avg mae: 0.028000\n",
      "\n",
      "KAZI\n",
      "train (12444, 17) 100\n",
      "Test avg loss: 0.023955\n",
      "Test avg mae: 0.025048\n",
      "\n",
      "KWLC\n",
      "train (12444, 17) 100\n",
      "division by zero\n",
      "dev\n",
      "BRDW\n",
      "dev (2773, 17) 100\n",
      "Test avg loss: 0.029170\n",
      "Test avg mae: 0.093000\n",
      "\n",
      "GROD\n",
      "dev (2773, 17) 100\n",
      "Test avg loss: 0.025056\n",
      "Test avg mae: 0.045482\n",
      "\n",
      "GUZI\n",
      "dev (2773, 17) 100\n",
      "Test avg loss: 0.024413\n",
      "Test avg mae: 0.039008\n",
      "\n",
      "JEDR\n",
      "dev (2773, 17) 100\n",
      "Test avg loss: 0.024724\n",
      "Test avg mae: 0.017778\n",
      "\n",
      "MOSK2\n",
      "dev (2773, 17) 100\n",
      "Test avg loss: 0.024798\n",
      "Test avg mae: 0.039188\n",
      "\n",
      "NWLU\n",
      "dev (2773, 17) 100\n",
      "Test avg loss: 0.025681\n",
      "Test avg mae: 0.037319\n",
      "\n",
      "PCHB\n",
      "dev (2773, 17) 100\n",
      "Test avg loss: 0.024584\n",
      "Test avg mae: 0.044390\n",
      "\n",
      "PPOL\n",
      "dev (2773, 17) 100\n",
      "Test avg loss: 0.025954\n",
      "Test avg mae: 0.074627\n",
      "\n",
      "RUDN\n",
      "dev (2773, 17) 100\n",
      "Test avg loss: 0.025528\n",
      "Test avg mae: 0.049302\n",
      "\n",
      "RYNR\n",
      "dev (2773, 17) 100\n",
      "Test avg loss: 0.026649\n",
      "Test avg mae: 0.071854\n",
      "\n",
      "RZEC\n",
      "dev (2773, 17) 100\n",
      "division by zero\n",
      "SGOR\n",
      "dev (2773, 17) 100\n",
      "Test avg loss: 0.025600\n",
      "Test avg mae: 0.083871\n",
      "\n",
      "TRBC2\n",
      "dev (2773, 17) 100\n",
      "Test avg loss: 0.025220\n",
      "Test avg mae: 0.032712\n",
      "\n",
      "TRN2\n",
      "dev (2773, 17) 100\n",
      "Test avg loss: 0.024823\n",
      "Test avg mae: 0.039834\n",
      "\n",
      "TRZS\n",
      "dev (2773, 17) 100\n",
      "Test avg loss: 0.026505\n",
      "Test avg mae: 0.093793\n",
      "\n",
      "ZMST\n",
      "dev (2773, 17) 100\n",
      "Test avg loss: 0.024234\n",
      "Test avg mae: 0.030628\n",
      "\n",
      "LUBW\n",
      "dev (2773, 17) 100\n",
      "Test avg loss: 0.029821\n",
      "Test avg mae: 0.129583\n",
      "\n",
      "DWOL\n",
      "dev (2773, 17) 100\n",
      "Test avg loss: 0.024637\n",
      "Test avg mae: 0.048252\n",
      "\n",
      "LUBZ\n",
      "dev (2773, 17) 100\n",
      "Test avg loss: 0.023594\n",
      "Test avg mae: 0.010000\n",
      "\n",
      "ZUKW2\n",
      "dev (2773, 17) 100\n",
      "Test avg loss: 0.024708\n",
      "Test avg mae: 0.035825\n",
      "\n",
      "DABR\n",
      "dev (2773, 17) 100\n",
      "Test avg loss: 0.024622\n",
      "Test avg mae: 0.036894\n",
      "\n",
      "PEKW2\n",
      "dev (2773, 17) 100\n",
      "Test avg loss: 0.026771\n",
      "Test avg mae: 0.046392\n",
      "\n",
      "KRZY\n",
      "dev (2773, 17) 100\n",
      "Test avg loss: 0.038841\n",
      "Test avg mae: 0.099091\n",
      "\n",
      "OBIS\n",
      "dev (2773, 17) 100\n",
      "Test avg loss: 0.025738\n",
      "Test avg mae: 0.112784\n",
      "\n",
      "KAZI\n",
      "dev (2773, 17) 100\n",
      "Test avg loss: 0.024498\n",
      "Test avg mae: 0.029636\n",
      "\n",
      "KWLC\n",
      "dev (2773, 17) 100\n",
      "Test avg loss: 0.023677\n",
      "Test avg mae: 0.030000\n",
      "\n",
      "test\n",
      "BRDW\n",
      "test (2785, 17) 100\n",
      "Test avg loss: 0.026271\n",
      "Test avg mae: 0.062836\n",
      "\n",
      "GROD\n",
      "test (2785, 17) 100\n",
      "Test avg loss: 0.025284\n",
      "Test avg mae: 0.057771\n",
      "\n",
      "GUZI\n",
      "test (2785, 17) 100\n",
      "Test avg loss: 0.025403\n",
      "Test avg mae: 0.050159\n",
      "\n",
      "JEDR\n",
      "test (2785, 17) 100\n",
      "Test avg loss: 0.027117\n",
      "Test avg mae: 0.039194\n",
      "\n",
      "MOSK2\n",
      "test (2785, 17) 100\n",
      "Test avg loss: 0.024491\n",
      "Test avg mae: 0.035829\n",
      "\n",
      "NWLU\n",
      "test (2785, 17) 100\n",
      "Test avg loss: 0.024571\n",
      "Test avg mae: 0.030738\n",
      "\n",
      "PCHB\n",
      "test (2785, 17) 100\n",
      "Test avg loss: 0.024751\n",
      "Test avg mae: 0.039241\n",
      "\n",
      "PPOL\n",
      "test (2785, 17) 100\n",
      "Test avg loss: 0.025289\n",
      "Test avg mae: 0.047288\n",
      "\n",
      "RUDN\n",
      "test (2785, 17) 100\n",
      "Test avg loss: 0.025869\n",
      "Test avg mae: 0.056053\n",
      "\n",
      "RYNR\n",
      "test (2785, 17) 100\n",
      "Test avg loss: 0.027814\n",
      "Test avg mae: 0.101556\n",
      "\n",
      "RZEC\n",
      "test (2785, 17) 100\n",
      "division by zero\n",
      "SGOR\n",
      "test (2785, 17) 100\n",
      "Test avg loss: 0.025836\n",
      "Test avg mae: 0.043481\n",
      "\n",
      "TRBC2\n",
      "test (2785, 17) 100\n",
      "Test avg loss: 0.024740\n",
      "Test avg mae: 0.031455\n",
      "\n",
      "TRN2\n",
      "test (2785, 17) 100\n",
      "Test avg loss: 0.027810\n",
      "Test avg mae: 0.063457\n",
      "\n",
      "TRZS\n",
      "test (2785, 17) 100\n",
      "Test avg loss: 0.024783\n",
      "Test avg mae: 0.036907\n",
      "\n",
      "ZMST\n",
      "test (2785, 17) 100\n",
      "Test avg loss: 0.025570\n",
      "Test avg mae: 0.050635\n",
      "\n",
      "LUBW\n",
      "test (2785, 17) 100\n",
      "Test avg loss: 0.024866\n",
      "Test avg mae: 0.038000\n",
      "\n",
      "DWOL\n",
      "test (2785, 17) 100\n",
      "Test avg loss: 0.024782\n",
      "Test avg mae: 0.035500\n",
      "\n",
      "LUBZ\n",
      "test (2785, 17) 100\n",
      "Test avg loss: 0.023713\n",
      "Test avg mae: 0.020000\n",
      "\n",
      "ZUKW2\n",
      "test (2785, 17) 100\n",
      "Test avg loss: 0.024829\n",
      "Test avg mae: 0.034684\n",
      "\n",
      "DABR\n",
      "test (2785, 17) 100\n",
      "Test avg loss: 0.025824\n",
      "Test avg mae: 0.040625\n",
      "\n",
      "PEKW2\n",
      "test (2785, 17) 100\n",
      "Test avg loss: 0.024243\n",
      "Test avg mae: 0.036092\n",
      "\n",
      "KRZY\n",
      "test (2785, 17) 100\n",
      "Test avg loss: 0.025173\n",
      "Test avg mae: 0.060000\n",
      "\n",
      "OBIS\n",
      "test (2785, 17) 100\n",
      "Test avg loss: 0.026633\n",
      "Test avg mae: 0.056509\n",
      "\n",
      "KAZI\n",
      "test (2785, 17) 100\n",
      "Test avg loss: 0.025799\n",
      "Test avg mae: 0.038860\n",
      "\n",
      "KWLC\n",
      "test (2785, 17) 100\n",
      "Test avg loss: 0.025239\n",
      "Test avg mae: 0.052885\n",
      "\n"
     ]
    }
   ],
   "source": [
    "stations = data.metadata.station_code.unique()\n",
    "\n",
    "results = []\n",
    "# highest_dev_mae = 0\n",
    "# highest_test_mae = 0\n",
    "\n",
    "for split in splits: \n",
    "    split_results = {}\n",
    "    print(split)\n",
    "    for station in stations: \n",
    "        print(station)\n",
    "        \n",
    "        split_results[station] = {}\n",
    "        gen = train.get_data_generator(split=split, station=station, sampling_rate=sampling_rate, path=data_path, window='fixed')\n",
    "        data_loader = DataLoader(gen, batch_size=256, shuffle=False, num_workers=0,\n",
    "                                worker_init_fn=worker_seeding)\n",
    "    \n",
    "        test_loss, test_mae = None, None\n",
    "        try: \n",
    "            test_loss, test_mae = train.test_one_epoch(model, data_loader, pick_mae, wandb_log=False)\n",
    "            test_mae = float(test_mae)\n",
    "            \n",
    "        except Exception as e: \n",
    "            print(e)\n",
    "    \n",
    "        split_results[station]['mae']=test_mae\n",
    "        split_results[station]['loss']=test_loss\n",
    "    results.append(split_results)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "206e6973-61bb-4f84-95a5-f910c7c7dc2b",
   "metadata": {},
   "source": [
    "#### Plot MAE per station"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "f0366dc6-d1e1-44cb-b748-d0af44a77219",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 49748 characters]",
      "text/plain": [
       "<Figure size 1500x400 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 47592 characters]",
      "text/plain": [
       "<Figure size 1500x400 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 46548 characters]",
      "text/plain": [
       "<Figure size 1500x400 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "for i, split in enumerate(splits): \n",
    "    results_df = pd.DataFrame(results[i]).transpose()\n",
    "\n",
    "    # for station, values in results_df[['mae']].itertuples():\n",
    "    #     results_df.loc[station, 'mae'] =  np.mean([v for v in values if v is not None]) \n",
    "    # for station, values in results_df[['loss']].itertuples():\n",
    "    #     results_df.loc[station, 'loss'] =  np.mean([v for v in values if v is not None]) \n",
    "        \n",
    "    results_df.plot(kind='bar', figsize=(15,4), title=f\"Mean results per station in {split} set\")\n",
    "\n",
    "    results[i] = results_df"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f02c9173-986c-45d5-973c-a68406ff20df",
   "metadata": {},
   "source": [
    "### Check predictions for stations with highest MAE"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "66a99163-04cd-4ecb-a3c8-fa8dc112c435",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "highest mean MAE in dev set: 0.13 for station: LUBW\n",
      "highest mean MAE in test set: 0.10 for station: RYNR\n"
     ]
    }
   ],
   "source": [
    "dev_res = results[1]\n",
    "station_with_worst_res_dev_set = dev_res[dev_res.mae == dev_res.mae.max()].index[0]\n",
    "highest_dev_mae = dev_res.loc[station_with_worst_res_dev_set, 'mae']\n",
    "\n",
    "test_res = results[2]\n",
    "station_with_worst_res_test_set = test_res[test_res.mae == test_res.mae.max()].index[0]\n",
    "highest_test_mae = test_res.loc[station_with_worst_res_test_set, 'mae']\n",
    "print(f\"highest mean MAE in dev set: {highest_dev_mae:.2f} for station: {station_with_worst_res_dev_set}\")\n",
    "print(f\"highest mean MAE in test set: {highest_test_mae:.2f} for station: {station_with_worst_res_test_set}\")\n",
    "\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "73a07433-1663-440c-bce3-d73988e9685b",
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot_sample(sample, model, i, desc=None): \n",
    "    fig = plt.figure(figsize=(15, 10))\n",
    "    \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",
    "        \n",
    "        axs[1].plot(pred[0], label='pred', color='orange')\n",
    "        plt.legend()\n",
    "\n",
    "        pred_pick_idx  = np.argmax(pred[0])\n",
    "        true_pick_idx = np.argmax(sample['y'][0])\n",
    "\n",
    "        \n",
    "        \n",
    "        mae_error = np.abs(pred_pick_idx - true_pick_idx) /100 #mae in seconds\n",
    "\n",
    "        fig.suptitle(f\"Predictions for sample: {i} {desc}, mae: {mae_error}s\")\n",
    "        \n",
    "        plt.show()\n",
    "        \n",
    "    "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d28d456a-4034-4eea-ac28-ed29ddb98d75",
   "metadata": {},
   "source": [
    "### Find random samples that reproduce obtained results \n",
    "\n",
    "the model makes at least as bad predictions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "235ae13e-3250-49e6-8bb2-765c037b0870",
   "metadata": {},
   "outputs": [],
   "source": [
    "##### dev set"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "5e3829f5-360a-431f-9fdd-bdeff6a5a1e4",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "dev (2773, 17) 100\n"
     ]
    },
    {
     "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": [
      "0.12958333333333336 LUBW\n",
      "[1.06 0.4  0.33 0.28 0.23 0.1  0.08 0.07 0.07 0.07 0.07 0.06 0.05 0.05\n",
      " 0.05 0.04 0.03 0.03 0.01 0.01 0.01 0.01 0.   0.  ]\n",
      "2 1.06\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 144208 characters]",
      "text/plain": [
       "<Figure size 1500x1000 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "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": [
      "3 0.4\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 149900 characters]",
      "text/plain": [
       "<Figure size 1500x1000 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "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": [
      "13 0.33\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 129660 characters]",
      "text/plain": [
       "<Figure size 1500x1000 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "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": [
      "22 0.28\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 149768 characters]",
      "text/plain": [
       "<Figure size 1500x1000 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "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": [
      "17 0.23\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 145056 characters]",
      "text/plain": [
       "<Figure size 1500x1000 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "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": [
      "18 0.1\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 116608 characters]",
      "text/plain": [
       "<Figure size 1500x1000 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "\n",
    "mean_mae = 0\n",
    "samples = []\n",
    "split = 'dev'\n",
    "station = station_with_worst_res_dev_set\n",
    "window='fixed'\n",
    "\n",
    "gen = train.get_data_generator(split=split, station=station , sampling_rate=sampling_rate, path=data_path, \n",
    "                                window='fixed')\n",
    "\n",
    "\n",
    "station_mae = []\n",
    "with torch.no_grad():\n",
    "    for i in range(len(gen)): \n",
    "        # idx = np.random.randint(len(gen))\n",
    "        idx = i\n",
    "        sample = gen[idx]\n",
    "        samples.append(sample)\n",
    "        pred = model(torch.tensor(sample[\"X\"], device=model.device).unsqueeze(0))  \n",
    "        pred = pred[0].cpu().numpy()\n",
    "\n",
    "        pred_pick_idx  = np.argmax(pred[0])\n",
    "        true_pick_idx = np.argmax(sample['y'][0])            \n",
    "            \n",
    "        mae_error = np.abs(pred_pick_idx - true_pick_idx) /100 #mae in seconds\n",
    "        station_mae.append(mae_error)\n",
    "\n",
    "sorted = np.argsort(station_mae)[::-1]\n",
    "mean_mae = np.mean(station_mae)\n",
    "print(mean_mae, station_with_worst_res_dev_set)\n",
    "print(np.array(station_mae)[sorted])\n",
    "\n",
    "\n",
    "## plot samples with mae error at leas 0.2s\n",
    "for idx in sorted:\n",
    "    if station_mae[idx] < 0.1: \n",
    "        break\n",
    "    print(idx, station_mae[idx])\n",
    "    plot_sample(samples[idx], model, idx, desc=f\" from station {station} {split} set\")\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "fc074059-8faf-4d7b-8b2c-49995df6fb37",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "test (2785, 17) 100\n",
      "test (2785, 17) 100\n",
      "test (2785, 17) 100\n"
     ]
    },
    {
     "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": [
      "[3.11 2.33 1.23 1.21 0.59 0.53 0.39 0.23 0.22 0.12 0.11 0.11 0.1  0.1\n",
      " 0.1  0.09 0.09 0.08 0.08 0.08 0.07 0.07 0.07 0.06 0.06 0.06 0.06 0.06\n",
      " 0.05 0.05 0.05 0.04 0.04 0.04 0.04 0.04 0.04 0.04 0.04 0.04 0.04 0.04\n",
      " 0.04 0.04 0.04 0.04 0.04 0.03 0.03 0.03 0.03 0.03 0.03 0.03 0.03 0.03\n",
      " 0.03 0.03 0.03 0.03 0.03 0.03 0.03 0.03 0.03 0.03 0.03 0.03 0.03 0.03\n",
      " 0.03 0.02 0.02 0.02 0.02 0.02 0.02 0.02 0.02 0.02 0.02 0.02 0.02 0.02\n",
      " 0.02 0.02 0.02 0.02 0.02 0.02 0.02 0.02 0.02 0.02 0.02 0.02 0.02 0.02\n",
      " 0.01 0.01 0.01 0.01 0.01 0.01 0.01 0.01 0.01 0.01 0.01 0.01 0.01 0.01\n",
      " 0.01 0.01 0.01 0.01 0.01 0.01 0.01 0.01 0.01 0.01 0.01 0.   0.   0.\n",
      " 0.   0.   0.   0.   0.   0.   0.   0.   0.  ]\n",
      "91 3.11\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 118924 characters]",
      "text/plain": [
       "<Figure size 1500x1000 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "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": [
      "46 2.33\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 137804 characters]",
      "text/plain": [
       "<Figure size 1500x1000 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "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": [
      "63 1.23\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 130116 characters]",
      "text/plain": [
       "<Figure size 1500x1000 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "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": [
      "40 1.21\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 130812 characters]",
      "text/plain": [
       "<Figure size 1500x1000 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "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": [
      "28 0.59\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 146148 characters]",
      "text/plain": [
       "<Figure size 1500x1000 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "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": [
      "106 0.53\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 142356 characters]",
      "text/plain": [
       "<Figure size 1500x1000 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "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": [
      "24 0.39\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 123044 characters]",
      "text/plain": [
       "<Figure size 1500x1000 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "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": [
      "100 0.23\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 121268 characters]",
      "text/plain": [
       "<Figure size 1500x1000 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "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": [
      "119 0.22\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 94624 characters]",
      "text/plain": [
       "<Figure size 1500x1000 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "\n",
    "mean_mae = 0\n",
    "samples = []\n",
    "split = 'test'\n",
    "station = station_with_worst_res_test_set\n",
    "\n",
    "for i in range(3): \n",
    "    \n",
    "    gen = train.get_data_generator(split=split, station=station , sampling_rate=sampling_rate, path=data_path, \n",
    "                                  window='fixed')\n",
    "    station_mae = []\n",
    "    with torch.no_grad():\n",
    "        for i in range(len(gen)): \n",
    "            # idx = np.random.randint(len(gen))\n",
    "            idx = i\n",
    "            sample = gen[idx]\n",
    "            samples.append(sample)\n",
    "            pred = model(torch.tensor(sample[\"X\"], device=model.device).unsqueeze(0))  \n",
    "            pred = pred[0].cpu().numpy()\n",
    "    \n",
    "            pred_pick_idx  = np.argmax(pred[0])\n",
    "            true_pick_idx = np.argmax(sample['y'][0])            \n",
    "                \n",
    "            mae_error = np.abs(pred_pick_idx - true_pick_idx) /100 #mae in seconds\n",
    "            station_mae.append(mae_error)\n",
    "    \n",
    "    sorted = np.argsort(station_mae)[::-1]\n",
    "    mean_mae = np.mean(station_mae)\n",
    "\n",
    "print(np.array(station_mae)[sorted])\n",
    "\n",
    "\n",
    "## plot samples with mae error at leas 0.2s\n",
    "for idx in sorted:\n",
    "    if station_mae[idx] < 0.2: \n",
    "        break\n",
    "    print(idx, station_mae[idx])\n",
    "    plot_sample(samples[idx], model, idx, desc=f\" from station {station} {split} set\")\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7e599f32-b888-4235-8d15-f9d2c98a59cf",
   "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.9.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
