{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "3a96a25b-b38a-4485-b7ff-2509236757cb",
   "metadata": {},
   "outputs": [],
   "source": [
    "from pathlib import Path\n",
    "import pandas as pd\n",
    "import os\n",
    "import matplotlib.pyplot as plt\n",
    "import obspy\n",
    "from obspy.core.event import read_events\n",
    "\n",
    "import seisbench\n",
    "import seisbench.data as sbd\n",
    "import seisbench.util as sbu\n",
    "import numpy as np\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b8d8aeaf-56e7-4bb0-9e3d-a41f2197d0f4",
   "metadata": {},
   "source": [
    "# Creating a dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "53e5a331-6b24-4da7-a0ba-7931a9d79dd0",
   "metadata": {},
   "outputs": [],
   "source": [
    "input_path = \"/net/pr2/projects/plgrid/plggeposai/datasets/bogdanka\"\n",
    "catalog_path = \"/net/pr2/projects/plgrid/plggeposai/datasets/bogdanka/BOIS_all.xml\"\n",
    "output_path = \"/net/pr2/projects/plgrid/plggeposai/kmilian/platform-demo-scripts/datasets/bogdanka/seisbench_format\"\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "426d4428-2bfa-492a-93d0-01c143d77ddc",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/net/pr2/projects/plgrid/plggeposai/mambaforge/envs/eposai/lib/python3.11/site-packages/obspy/io/quakeml/core.py:193: UserWarning: Setting attribute \"preferred_description\" failed. Value \"uncertainty elipse\" could not be converted to type \"Enum([\"horizontal uncertainty\", \"uncertainty ellipse\", \"confidence ellipsoid\"])\". The attribute \"preferred_description\" will not be set and will be missing in the resulting object.\n",
      "  warnings.warn(msg)\n"
     ]
    }
   ],
   "source": [
    "events = read_events(catalog_path)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "281a13f7-d8d9-4d6c-a97c-89b1044671ad",
   "metadata": {},
   "source": [
    "### Define train/dev/test split"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "28f666f4-16c8-446f-9428-d28d9ce6f8f2",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_event_params(event): \n",
    "\n",
    "    \n",
    "    origin = event.preferred_origin()\n",
    "    if origin is None:\n",
    "        return {}\n",
    "\n",
    "    mag = event.preferred_magnitude()\n",
    "\n",
    "    source_id = str(event.resource_id)\n",
    "\n",
    "    event_params = {\n",
    "        \"source_id\": source_id,\n",
    "        \"source_origin_uncertainty_sec\": origin.time_errors[\"uncertainty\"],\n",
    "        \"source_latitude_deg\": origin.latitude,\n",
    "        \"source_latitude_uncertainty_km\": origin.latitude_errors[\"uncertainty\"],\n",
    "        \"source_longitude_deg\": origin.longitude,\n",
    "        \"source_longitude_uncertainty_km\": origin.longitude_errors[\"uncertainty\"],\n",
    "        \"source_depth_km\": origin.depth / 1e3,\n",
    "        \"source_depth_uncertainty_km\": origin.depth_errors[\"uncertainty\"] / 1e3 if origin.depth_errors[\"uncertainty\"] is not None else None,\n",
    "    }\n",
    "\n",
    "    if mag is not None:\n",
    "        event_params[\"source_magnitude\"] = mag.mag\n",
    "        event_params[\"source_magnitude_uncertainty\"] = mag.mag_errors[\"uncertainty\"]\n",
    "        event_params[\"source_magnitude_type\"] = mag.magnitude_type\n",
    "        event_params[\"source_magnitude_author\"] = mag.creation_info.agency_id if mag.creation_info is not None else None\n",
    "        \n",
    "    return event_params\n",
    "\n",
    "\n",
    "def get_trace_params(pick):\n",
    "    net = pick.waveform_id.network_code\n",
    "    sta = pick.waveform_id.station_code\n",
    "\n",
    "    trace_params = {\n",
    "        \"station_network_code\": net,\n",
    "        \"station_code\": sta,\n",
    "        \"trace_channel\": pick.waveform_id.channel_code,\n",
    "        \"station_location_code\": pick.waveform_id.location_code,\n",
    "        \"time\": pick.time\n",
    "    }\n",
    "\n",
    "    return trace_params\n",
    "\n",
    "\n",
    "def get_trace_path(input_path, trace_params):\n",
    "\n",
    "    year = trace_params[\"time\"].year\n",
    "    day_of_year = pd.Timestamp(str(trace_params[\"time\"])).day_of_year\n",
    "    net = trace_params[\"station_network_code\"]\n",
    "    station = trace_params[\"station_code\"]\n",
    "    tr_channel = trace_params[\"trace_channel\"]\n",
    "\n",
    "    path = f\"{input_path}/{year}/{net}/{station}/{tr_channel}.D/{net}.{station}..{tr_channel}.D.{year}.{day_of_year}\"\n",
    "    return path"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "8415f6c1-0dfc-41e2-a06e-2e9bc94f0d4c",
   "metadata": {},
   "outputs": [],
   "source": [
    "def split_events(events, input_path):\n",
    "\n",
    "    print(\"Splitting available events into train, dev and test sets ...\")\n",
    "    events_stats = pd.DataFrame()\n",
    "    events_stats.index.name = \"event\"\n",
    "    missing_mseeds = []\n",
    "\n",
    "    for i, event in enumerate(events):\n",
    "        #check if mseed exists\n",
    "        actual_picks = 0\n",
    "        for pick in event.picks:\n",
    "            trace_params = get_trace_params(pick)\n",
    "            trace_path = get_trace_path(input_path, trace_params)\n",
    "            if os.path.isfile(trace_path):\n",
    "                actual_picks += 1\n",
    "            else: \n",
    "                missing_mseeds.append(trace_path.split('/'))\n",
    "                \n",
    "\n",
    "        events_stats.loc[i, \"pick_count\"] = len(event.picks)\n",
    "        events_stats.loc[i, \"available_picks\"] = actual_picks\n",
    "        if len(event.picks):\n",
    "            events_stats.loc[i, \"first_peak_time\"] = trace_params[\"time\"]\n",
    "\n",
    "    events_stats['pick_count_cumsum'] = events_stats.pick_count.cumsum()\n",
    "\n",
    "    train_th = 0.7 * events_stats.pick_count_cumsum.values[-1]\n",
    "    dev_th = 0.85 * events_stats.pick_count_cumsum.values[-1]\n",
    "\n",
    "    events_stats['split'] = 'test'\n",
    "    for i, event in events_stats.iterrows():\n",
    "        if event['pick_count_cumsum'] < train_th:\n",
    "            events_stats.loc[i, 'split'] = 'train'\n",
    "        elif event['pick_count_cumsum'] < dev_th:\n",
    "            events_stats.loc[i, 'split'] = 'dev'\n",
    "        else:\n",
    "            break\n",
    "\n",
    "    missing_mseeds = pd.DataFrame(np.array(missing_mseeds)[:, -5:], columns=['y', 'n','s','tr','f'])\n",
    "    missing_mseeds['last_dot'] = missing_mseeds['f'].apply(lambda x: x.rindex('.'))\n",
    "\n",
    "    for i, row in missing_mseeds.iterrows(): \n",
    "        missing_mseeds.loc[i, 'd'] = int(row['f'][row['last_dot']+1:])\n",
    "\n",
    "\n",
    "    return events_stats, missing_mseeds"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "96b8945a-06d6-498c-81d9-ba76788fdf37",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Splitting available events into train, dev and test sets ...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/net/ascratch/people/plgkmilian/slurm_jobdir/5092284/tmp/ipykernel_248521/649615512.py:23: FutureWarning: Setting an item of incompatible dtype is deprecated and will raise in a future error of pandas. Value '2018-09-24T14:20:07.521000Z' has dtype incompatible with float64, please explicitly cast to a compatible dtype first.\n",
      "  events_stats.loc[i, \"first_peak_time\"] = trace_params[\"time\"]\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>pick_count</th>\n",
       "      <th>available_picks</th>\n",
       "      <th>first_peak_time</th>\n",
       "      <th>pick_count_cumsum</th>\n",
       "      <th>split</th>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>event</th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>NaN</td>\n",
       "      <td>0.0</td>\n",
       "      <td>train</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>10.0</td>\n",
       "      <td>8.0</td>\n",
       "      <td>2018-09-24T14:20:07.521000Z</td>\n",
       "      <td>10.0</td>\n",
       "      <td>train</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>10.0</td>\n",
       "      <td>10.0</td>\n",
       "      <td>2018-10-04T21:44:08.474000Z</td>\n",
       "      <td>20.0</td>\n",
       "      <td>train</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>12.0</td>\n",
       "      <td>10.0</td>\n",
       "      <td>2018-10-05T23:25:45.835000Z</td>\n",
       "      <td>32.0</td>\n",
       "      <td>train</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>10.0</td>\n",
       "      <td>10.0</td>\n",
       "      <td>2018-10-07T22:02:11.773000Z</td>\n",
       "      <td>42.0</td>\n",
       "      <td>train</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "       pick_count  available_picks              first_peak_time  \\\n",
       "event                                                             \n",
       "0             0.0              0.0                          NaN   \n",
       "1            10.0              8.0  2018-09-24T14:20:07.521000Z   \n",
       "2            10.0             10.0  2018-10-04T21:44:08.474000Z   \n",
       "3            12.0             10.0  2018-10-05T23:25:45.835000Z   \n",
       "4            10.0             10.0  2018-10-07T22:02:11.773000Z   \n",
       "\n",
       "       pick_count_cumsum  split  \n",
       "event                            \n",
       "0                    0.0  train  \n",
       "1                   10.0  train  \n",
       "2                   20.0  train  \n",
       "3                   32.0  train  \n",
       "4                   42.0  train  "
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "events_stats, missing_mseeds = split_events(events, input_path)\n",
    "events_stats.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "4f2e3d49-03da-4086-8d42-a77c2aa29c5f",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>y</th>\n",
       "      <th>n</th>\n",
       "      <th>s</th>\n",
       "      <th>tr</th>\n",
       "      <th>f</th>\n",
       "      <th>last_dot</th>\n",
       "      <th>d</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>2018</td>\n",
       "      <td>PL</td>\n",
       "      <td>BOG5</td>\n",
       "      <td>EHN.D</td>\n",
       "      <td>PL.BOG5..EHN.D.2018.267</td>\n",
       "      <td>19</td>\n",
       "      <td>267.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>2018</td>\n",
       "      <td>PL</td>\n",
       "      <td>BOG5</td>\n",
       "      <td>EHZ.D</td>\n",
       "      <td>PL.BOG5..EHZ.D.2018.267</td>\n",
       "      <td>19</td>\n",
       "      <td>267.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>2018</td>\n",
       "      <td>PL</td>\n",
       "      <td>BOG5</td>\n",
       "      <td>EHE.D</td>\n",
       "      <td>PL.BOG5..EHE.D.2018.278</td>\n",
       "      <td>19</td>\n",
       "      <td>278.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>2018</td>\n",
       "      <td>PL</td>\n",
       "      <td>BOG5</td>\n",
       "      <td>EHZ.D</td>\n",
       "      <td>PL.BOG5..EHZ.D.2018.278</td>\n",
       "      <td>19</td>\n",
       "      <td>278.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>2018</td>\n",
       "      <td>PL</td>\n",
       "      <td>BOG5</td>\n",
       "      <td>EHN.D</td>\n",
       "      <td>PL.BOG5..EHN.D.2018.281</td>\n",
       "      <td>19</td>\n",
       "      <td>281.0</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "      y   n     s     tr                        f  last_dot      d\n",
       "0  2018  PL  BOG5  EHN.D  PL.BOG5..EHN.D.2018.267        19  267.0\n",
       "1  2018  PL  BOG5  EHZ.D  PL.BOG5..EHZ.D.2018.267        19  267.0\n",
       "2  2018  PL  BOG5  EHE.D  PL.BOG5..EHE.D.2018.278        19  278.0\n",
       "3  2018  PL  BOG5  EHZ.D  PL.BOG5..EHZ.D.2018.278        19  278.0\n",
       "4  2018  PL  BOG5  EHN.D  PL.BOG5..EHN.D.2018.281        19  281.0"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "missing_mseeds.head()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fadeaa9e-b7d7-48c7-a09c-872a545f000f",
   "metadata": {},
   "source": [
    "#### Check missing files"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "ffd246b8-118b-406a-bb87-4cd504970711",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<Axes: title={'center': 'Missing mseeds grouped by year and station'}, xlabel='y,s'>"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 48996 characters]",
      "text/plain": [
       "<Figure size 2000x400 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "missing_mseeds.groupby(['y','s']).size().plot(kind='bar', figsize=(20,4), title='Missing mseeds grouped by year and station')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "24ee2524-6d1a-41a2-bc29-bb64905e8201",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th>min</th>\n",
       "      <th>max</th>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>y</th>\n",
       "      <th>s</th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>2018</th>\n",
       "      <th>BOG5</th>\n",
       "      <td>267.0</td>\n",
       "      <td>281.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th rowspan=\"4\" valign=\"top\">2019</th>\n",
       "      <th>ALBE</th>\n",
       "      <td>2.0</td>\n",
       "      <td>99.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>BARK</th>\n",
       "      <td>11.0</td>\n",
       "      <td>99.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>CYCO</th>\n",
       "      <td>2.0</td>\n",
       "      <td>99.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>DRAT</th>\n",
       "      <td>2.0</td>\n",
       "      <td>99.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>...</th>\n",
       "      <th>...</th>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th rowspan=\"5\" valign=\"top\">2023</th>\n",
       "      <th>OSTR</th>\n",
       "      <td>10.0</td>\n",
       "      <td>149.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>PIAS</th>\n",
       "      <td>12.0</td>\n",
       "      <td>124.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>PUCH</th>\n",
       "      <td>10.0</td>\n",
       "      <td>149.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>ROGO</th>\n",
       "      <td>12.0</td>\n",
       "      <td>149.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>WESO</th>\n",
       "      <td>10.0</td>\n",
       "      <td>149.0</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "<p>62 rows × 2 columns</p>\n",
       "</div>"
      ],
      "text/plain": [
       "             min    max\n",
       "y    s                 \n",
       "2018 BOG5  267.0  281.0\n",
       "2019 ALBE    2.0   99.0\n",
       "     BARK   11.0   99.0\n",
       "     CYCO    2.0   99.0\n",
       "     DRAT    2.0   99.0\n",
       "...          ...    ...\n",
       "2023 OSTR   10.0  149.0\n",
       "     PIAS   12.0  124.0\n",
       "     PUCH   10.0  149.0\n",
       "     ROGO   12.0  149.0\n",
       "     WESO   10.0  149.0\n",
       "\n",
       "[62 rows x 2 columns]"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "missing_mseeds.groupby(['y', 's'])['d'].agg(['min', 'max'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 106,
   "id": "211c37be-e5eb-40ed-bbe8-0d7a5e92ae77",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "s\n",
       "ALBE    Axes(0.125,0.11;0.775x0.77)\n",
       "BARK    Axes(0.125,0.11;0.775x0.77)\n",
       "CYCO    Axes(0.125,0.11;0.775x0.77)\n",
       "DRAT    Axes(0.125,0.11;0.775x0.77)\n",
       "GARB    Axes(0.125,0.11;0.775x0.77)\n",
       "KANI    Axes(0.125,0.11;0.775x0.77)\n",
       "KOPI    Axes(0.125,0.11;0.775x0.77)\n",
       "OSTR    Axes(0.125,0.11;0.775x0.77)\n",
       "PIAS    Axes(0.125,0.11;0.775x0.77)\n",
       "PUCH    Axes(0.125,0.11;0.775x0.77)\n",
       "ROGO    Axes(0.125,0.11;0.775x0.77)\n",
       "WESO    Axes(0.125,0.11;0.775x0.77)\n",
       "Name: d, dtype: object"
      ]
     },
     "execution_count": 106,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 26460 characters]",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "missing_mseeds[missing_mseeds.y=='2019'].groupby(['s'])['d'].plot(kind='hist', legend=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 64,
   "id": "1ead7261-da65-4832-af20-f841a798febc",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "y\n",
       "2018       6\n",
       "2019    2540\n",
       "2020    4486\n",
       "2021    6507\n",
       "2022    6111\n",
       "2023    1379\n",
       "dtype: int64"
      ]
     },
     "execution_count": 64,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "missing_mseeds.groupby('y').size()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "327f21de-64bc-49b1-849b-4d2d79b505ba",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(2099, 3)"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "events_stats.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "a7407b0a-2efc-4c82-bf6a-8509dbdbc07c",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<Axes: xlabel='event'>"
      ]
     },
     "execution_count": 18,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 22588 characters]",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "events_stats.pick_count_cumsum.plot()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "12dbaa70-a178-4867-a710-0d8147ce4ecb",
   "metadata": {},
   "source": [
    "### Transform mseeds to Seisbench format"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "80fc0198-6a1a-48e1-ada9-4faa048c6e69",
   "metadata": {},
   "outputs": [],
   "source": [
    "def find_trace(pick_time, traces): \n",
    "    for tr in traces: \n",
    "        if pick_time>tr.stats.endtime: \n",
    "            continue\n",
    "        if pick_time>=tr.stats.starttime: \n",
    "            # print(pick_time, \" - selected trace: \", tr)\n",
    "            return tr\n",
    "\n",
    "    print(pick_time, \" no matching trace \")\n",
    "    return None\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "60ee1d41-d3db-4661-ade7-882bd2371f62",
   "metadata": {},
   "outputs": [],
   "source": [
    "def load_trace(input_path, trace_params):\n",
    "    trace_path = get_trace_path(input_path, trace_params)\n",
    "    trace, stream = None, None\n",
    "\n",
    "    if not os.path.isfile(trace_path):\n",
    "        print(trace_path + \" not found\")\n",
    "    else:\n",
    "        stream = obspy.read(trace_path)\n",
    "        if len(stream.traces) > 1:\n",
    "            trace = find_trace(trace_params[\"time\"], stream.traces)\n",
    "        elif len(stream.traces) == 0:\n",
    "            print(\"no data in:\", trace_path)\n",
    "        else:\n",
    "            trace = stream.traces[0]\n",
    "\n",
    "    return trace, stream\n",
    "\n",
    "\n",
    "def load_stream(input_path, trace_params, time_before=60, time_after=60):\n",
    "    trace_path = get_trace_path(input_path, trace_params)\n",
    "    sampling_rate, stream = None, None\n",
    "    pick_time = trace_params[\"time\"]\n",
    "\n",
    "    if not os.path.isfile(trace_path):\n",
    "        print(trace_path + \" not found\")\n",
    "    else:\n",
    "        stream = obspy.read(trace_path)\n",
    "        stream = stream.slice(pick_time - time_before, pick_time + time_after)\n",
    "        if len(stream.traces) == 0:\n",
    "            print(f\"no data in: {trace_path}\")\n",
    "        else:\n",
    "            sampling_rate = stream.traces[0].stats.sampling_rate\n",
    "\n",
    "    return sampling_rate, stream"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "7426f90a-4779-4c87-b445-4c41635a04f9",
   "metadata": {},
   "outputs": [],
   "source": [
    "metadata_path = output_path + \"/metadata.csv\"\n",
    "waveforms_path = output_path + \"/waveforms.hdf5\"\n",
    "train = 0.7\n",
    "dev = 0.15\n",
    "test = 0.15"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cc393ac1-791e-414a-892f-e8cf0d8e2a12",
   "metadata": {},
   "outputs": [],
   "source": [
    "with sbd.WaveformDataWriter(metadata_path, waveforms_path) as writer:\n",
    "    writer.data_format = {\n",
    "        \"dimension_order\": \"CW\",\n",
    "        \"component_order\": \"ZNE\",\n",
    "    }\n",
    "    for i, event in enumerate(events):\n",
    "        logger.debug(f\"Converting {i} event\")\n",
    "        event_params = get_event_params(event)\n",
    "        event_params[\"split\"] = events_stats.loc[i, \"split\"]\n",
    "        #             b = False\n",
    "\n",
    "        for pick in event.picks:\n",
    "            trace_params = get_trace_params(pick)\n",
    "            sampling_rate, stream = load_stream(input_path, trace_params)\n",
    "            if stream is None:\n",
    "                continue\n",
    "\n",
    "            actual_t_start, data, _ = sbu.stream_to_array(\n",
    "                stream,\n",
    "                component_order=writer.data_format[\"component_order\"],\n",
    "            )\n",
    "\n",
    "            trace_params[\"trace_sampling_rate_hz\"] = sampling_rate\n",
    "            trace_params[\"trace_start_time\"] = str(actual_t_start)\n",
    "\n",
    "            pick_time = obspy.core.utcdatetime.UTCDateTime(trace_params[\"time\"])\n",
    "            pick_idx = (pick_time - actual_t_start) * sampling_rate\n",
    "\n",
    "            trace_params[f\"trace_{pick.phase_hint}_arrival_sample\"] = int(pick_idx)\n",
    "\n",
    "            writer.add_trace({**event_params, **trace_params}, data)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5cf46368-2e92-4aea-8fe8-298f3200f509",
   "metadata": {},
   "source": [
    "### Load created dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "bcbefdd8-91b0-461b-8677-5b10325822d0",
   "metadata": {},
   "outputs": [],
   "source": [
    "data = sbd.WaveformDataset(output_path, sampling_rate=100)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "77f8060d-0703-41ce-ade8-c738e7a8e540",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<seisbench.data.base.WaveformDataset at 0x14b58e696850>"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "data"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ebc4f1ea-5e33-458d-82eb-ab6e5982ef63",
   "metadata": {},
   "source": [
    "#### Check train/dev/test proportions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "98f2ea18-18d8-4d7c-8e09-8c242842bda4",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "All samples:  6743\n",
      "Training examples: 4720 70.0%\n",
      "Development examples: 1009 15.0%\n",
      "Test examples: 1014 15.0 %\n"
     ]
    }
   ],
   "source": [
    "all_samples = len(data.train()) + len(data.dev()) + len(data.test())\n",
    "print(\"All samples: \", all_samples)\n",
    "print(f\"Training examples: {len(data.train())} {len(data.train())/all_samples * 100:.1f}%\" )\n",
    "print(f\"Development examples: {len(data.dev())} {len(data.dev())/all_samples * 100:.1f}%\")\n",
    "print(f\"Test examples: {len(data.test())} {len(data.test())/all_samples * 100:.1f} %\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5ffd0f88-5a5c-4f21-8688-b34d8540ad30",
   "metadata": {},
   "source": [
    "#### Plot converted sample"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "c3af6da7-864f-490b-b260-733e49bd6d12",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 64212 characters]",
      "text/plain": [
       "<Figure size 1500x500 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "fig = plt.figure(figsize=(15, 5))\n",
    "ax = fig.add_subplot(111)\n",
    "\n",
    "idx = 0\n",
    "ax.plot(data.get_waveforms(idx).T)\n",
    "ax.axvline(data.metadata[\"trace_Pg_arrival_sample\"].iloc[idx], color=\"green\", lw=1)\n",
    "\n",
    "plt.show()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "12d829c7-1478-4fc4-9b8f-565d4fdae824",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 72116 characters]",
      "text/plain": [
       "<Figure size 1500x500 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "fig = plt.figure(figsize=(15, 5))\n",
    "ax = fig.add_subplot(111)\n",
    "idx = 1\n",
    "ax.plot(data.get_waveforms(idx).T)\n",
    "# ax.axvline(data.metadata[\"trace_Pg_arrival_sample\"].iloc[idx], color=\"green\", lw=1)\n",
    "ax.axvline(data.metadata[\"trace_Sg_arrival_sample\"].iloc[idx], color=\"black\", lw=1)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "d3534d4e-13a4-4780-8f7d-7490ae3608c6",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>trace_Pg_arrival_sample</th>\n",
       "      <th>trace_Sg_arrival_sample</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>5999.0</td>\n",
       "      <td>NaN</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>NaN</td>\n",
       "      <td>5999.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>NaN</td>\n",
       "      <td>6000.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>5999.0</td>\n",
       "      <td>NaN</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>6000.0</td>\n",
       "      <td>NaN</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>...</th>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>6738</th>\n",
       "      <td>NaN</td>\n",
       "      <td>5999.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>6739</th>\n",
       "      <td>5999.0</td>\n",
       "      <td>NaN</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>6740</th>\n",
       "      <td>6000.0</td>\n",
       "      <td>NaN</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>6741</th>\n",
       "      <td>NaN</td>\n",
       "      <td>5999.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>6742</th>\n",
       "      <td>5999.0</td>\n",
       "      <td>NaN</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "<p>6743 rows × 2 columns</p>\n",
       "</div>"
      ],
      "text/plain": [
       "      trace_Pg_arrival_sample  trace_Sg_arrival_sample\n",
       "0                      5999.0                      NaN\n",
       "1                         NaN                   5999.0\n",
       "2                         NaN                   6000.0\n",
       "3                      5999.0                      NaN\n",
       "4                      6000.0                      NaN\n",
       "...                       ...                      ...\n",
       "6738                      NaN                   5999.0\n",
       "6739                   5999.0                      NaN\n",
       "6740                   6000.0                      NaN\n",
       "6741                      NaN                   5999.0\n",
       "6742                   5999.0                      NaN\n",
       "\n",
       "[6743 rows x 2 columns]"
      ]
     },
     "execution_count": 19,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "data.metadata[[\"trace_Pg_arrival_sample\", \"trace_Sg_arrival_sample\"]]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5d023f5f-120e-46ea-bfed-ead098fb247f",
   "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.10.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
