Updated layout

Gradio_app.ipynb

@@ -2,29 +2,14 @@
  "cells": [
   {
    "cell_type": "code",
-   "execution_count": 5,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "model = Onset_picker.load_from_checkpoint(\"./weights.ckpt\",\n",
-    "                                 picker=Updated_onset_picker(),\n",
-    "                                    learning_rate=3e-4)\n",
-    "model.eval()\n",
-    "model.freeze()\n",
-    "script = model.to_torchscript()\n",
-    "torch.jit.save(script, \"model.pt\")"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 32,
+   "execution_count": 51,
    "metadata": {},
    "outputs": [
     {
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "Running on local URL:  http://127.0.0.1:7878\n",
+      "Running on local URL:  http://127.0.0.1:7897\n",
       "\n",
       "To create a public link, set `share=True` in `launch()`.\n"
      ]
@@ -32,7 +17,7 @@
     {
      "data": {
       "text/html": [
-       "<div><iframe src=\"http://127.0.0.1:7878/\" width=\"100%\" height=\"500\" allow=\"autoplay; camera; microphone; clipboard-read; clipboard-write;\" frameborder=\"0\" allowfullscreen></iframe></div>"
+       "<div><iframe src=\"http://127.0.0.1:7897/\" width=\"100%\" height=\"500\" allow=\"autoplay; camera; microphone; clipboard-read; clipboard-write;\" frameborder=\"0\" allowfullscreen></iframe></div>"
       ],
       "text/plain": [
        "<IPython.core.display.HTML object>"
@@ -45,7 +30,7 @@
      "data": {
       "text/plain": []
      },
-     "execution_count": 32,
+     "execution_count": 51,
      "metadata": {},
      "output_type": "execute_result"
     },
@@ -116,13 +101,69 @@
      "name": "stderr",
      "output_type": "stream",
      "text": [
-      "/var/folders/_g/3q5q8_dj0ydcpktxlwxb5vrh0000gq/T/ipykernel_27324/2440224661.py:224: FutureWarning: The input object of type 'Tensor' is an array-like implementing one of the corresponding protocols (`__array__`, `__array_interface__` or `__array_struct__`); but not a sequence (or 0-D). In the future, this object will be coerced as if it was first converted using `np.array(obj)`. To retain the old behaviour, you have to either modify the type 'Tensor', or assign to an empty array created with `np.empty(correct_shape, dtype=object)`.\n",
+      "/var/folders/_g/3q5q8_dj0ydcpktxlwxb5vrh0000gq/T/ipykernel_27324/1938231065.py:224: FutureWarning: The input object of type 'Tensor' is an array-like implementing one of the corresponding protocols (`__array__`, `__array_interface__` or `__array_struct__`); but not a sequence (or 0-D). In the future, this object will be coerced as if it was first converted using `np.array(obj)`. To retain the old behaviour, you have to either modify the type 'Tensor', or assign to an empty array created with `np.empty(correct_shape, dtype=object)`.\n",
       "  waveforms = np.array(waveforms)[selection_indexes]\n",
-      "/var/folders/_g/3q5q8_dj0ydcpktxlwxb5vrh0000gq/T/ipykernel_27324/2440224661.py:224: VisibleDeprecationWarning: Creating an ndarray from ragged nested sequences (which is a list-or-tuple of lists-or-tuples-or ndarrays with different lengths or shapes) is deprecated. If you meant to do this, you must specify 'dtype=object' when creating the ndarray.\n",
+      "/var/folders/_g/3q5q8_dj0ydcpktxlwxb5vrh0000gq/T/ipykernel_27324/1938231065.py:224: VisibleDeprecationWarning: Creating an ndarray from ragged nested sequences (which is a list-or-tuple of lists-or-tuples-or ndarrays with different lengths or shapes) is deprecated. If you meant to do this, you must specify 'dtype=object' when creating the ndarray.\n",
       "  waveforms = np.array(waveforms)[selection_indexes]\n",
-      "/var/folders/_g/3q5q8_dj0ydcpktxlwxb5vrh0000gq/T/ipykernel_27324/2440224661.py:231: UserWarning: To copy construct from a tensor, it is recommended to use sourceTensor.clone().detach() or sourceTensor.clone().detach().requires_grad_(True), rather than torch.tensor(sourceTensor).\n",
+      "/var/folders/_g/3q5q8_dj0ydcpktxlwxb5vrh0000gq/T/ipykernel_27324/1938231065.py:231: UserWarning: To copy construct from a tensor, it is recommended to use sourceTensor.clone().detach() or sourceTensor.clone().detach().requires_grad_(True), rather than torch.tensor(sourceTensor).\n",
       "  waveforms = [torch.tensor(waveform) for waveform in waveforms]\n"
      ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Starting plotting 3 waveforms\n",
+      "Fetching topography\n",
+      "Plotting topo\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "/Users/anovosel/miniconda3/envs/phasehunter/lib/python3.11/site-packages/bmi_topography/api_key.py:49: UserWarning: You are using a demo key to fetch data from OpenTopography, functionality will be limited. See https://bmi-topography.readthedocs.io/en/latest/#api-key for more information.\n",
+      "  warnings.warn(\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Plotting waveform 1/3\n",
+      "Station 36.11758, -117.85486 has P velocity 4.987805380766392 and S velocity 2.9782985042350987\n",
+      "Plotting waveform 2/3\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "/var/folders/_g/3q5q8_dj0ydcpktxlwxb5vrh0000gq/T/ipykernel_27324/1938231065.py:299: FutureWarning: The frame.append method is deprecated and will be removed from pandas in a future version. Use pandas.concat instead.\n",
+      "  output_picks = output_picks.append(pd.DataFrame({'station_name': [names[i]], 'starttime' : [str(t0s[i])],\n",
+      "/var/folders/_g/3q5q8_dj0ydcpktxlwxb5vrh0000gq/T/ipykernel_27324/1938231065.py:299: FutureWarning: The frame.append method is deprecated and will be removed from pandas in a future version. Use pandas.concat instead.\n",
+      "  output_picks = output_picks.append(pd.DataFrame({'station_name': [names[i]], 'starttime' : [str(t0s[i])],\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Station 35.98249, -117.80885 has P velocity 4.255522557803516 and S velocity 2.2929437916670583\n",
+      "Plotting waveform 3/3\n",
+      "Station 35.69235, -117.75051 has P velocity 2.979034174961547 and S velocity 1.3728192788753049\n",
+      "Plotting stations\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "/var/folders/_g/3q5q8_dj0ydcpktxlwxb5vrh0000gq/T/ipykernel_27324/1938231065.py:299: FutureWarning: The frame.append method is deprecated and will be removed from pandas in a future version. Use pandas.concat instead.\n",
+      "  output_picks = output_picks.append(pd.DataFrame({'station_name': [names[i]], 'starttime' : [str(t0s[i])],\n",
+      "/var/folders/_g/3q5q8_dj0ydcpktxlwxb5vrh0000gq/T/ipykernel_27324/1938231065.py:324: UserWarning: FixedFormatter should only be used together with FixedLocator\n",
+      "  ax[i].set_xticklabels(ax[i].get_xticks(), rotation = 50)\n"
+     ]
     }
    ],
    "source": [
@@ -149,7 +190,7 @@
     "\n",
     "import matplotlib.pyplot as plt\n",
     "import matplotlib.dates as mdates\n",
-    "from matplotlib.colors import LightSource\n",
+    "from mpl_toolkits.axes_grid1 import ImageGrid\n",
     "\n",
     "from glob import glob\n",
     "\n",
@@ -309,8 +350,8 @@
     "            \n",
     "                waveform = waveform.select(channel=\"H[BH][ZNE]\")\n",
     "                waveform = waveform.merge(fill_value=0)\n",
-    "                waveform = waveform[:3]\n",
-    "                \n",
+    "                waveform = waveform[:3].sort(keys=['channel'], reverse=True)\n",
+    "\n",
     "                len_check = [len(x.data) for x in waveform]\n",
     "                if len(set(len_check)) > 1:\n",
     "                    continue\n",
@@ -371,8 +412,8 @@
     "    s_max_confidence = np.min([s_phases[i::len(waveforms)].std() for i in range(len(waveforms))])\n",
     "\n",
     "    print(f\"Starting plotting {len(waveforms)} waveforms\")\n",
-    "    fig, ax = plt.subplots(nrows=1, ncols=3, figsize=(10, 3))\n",
-    "\n",
+    "    fig, ax = plt.subplots(ncols=3, figsize=(10, 3))\n",
+    "    \n",
     "    # Plot topography\n",
     "    print('Fetching topography')\n",
     "    params = Topography.DEFAULT.copy()\n",
@@ -417,9 +458,6 @@
     "        ax[0].scatter(x[int(current_S.mean()*waveforms[i][0].shape[-1])], waveforms[i][0, 0].mean()+distances[i]*111.2, color='b', alpha=s_conf, marker='|')\n",
     "        ax[0].set_ylabel('Z')\n",
     "\n",
-    "        ax[0].xaxis.set_major_formatter(mdates.DateFormatter('%H:%M:%S'))\n",
-    "        ax[0].xaxis.set_major_locator(mdates.SecondLocator(interval=20))\n",
-    "\n",
     "        delta_t = t0s[i].timestamp - obspy.UTCDateTime(timestamp).timestamp\n",
     "\n",
     "        velocity_p = (distances[i]*111.2)/(delta_t+current_P.mean()*60).item()\n",
@@ -437,30 +475,37 @@
     "        y = np.linspace(st_lats[i], eq_lat, 50)\n",
     "        \n",
     "        # Plot the array\n",
-    "        ax[1].scatter(x, y, c=np.zeros_like(x)+velocity_p, alpha=0.5, vmin=0, vmax=8)\n",
-    "        ax[2].scatter(x, y, c=np.zeros_like(x)+velocity_s, alpha=0.5, vmin=0, vmax=8)\n",
+    "        ax[1].scatter(x, y, c=np.zeros_like(x)+velocity_p, alpha=0.1, vmin=0, vmax=8)\n",
+    "        ax[2].scatter(x, y, c=np.zeros_like(x)+velocity_s, alpha=0.1, vmin=0, vmax=8)\n",
     "\n",
     "    # Add legend\n",
     "    ax[0].scatter(None, None, color='r', marker='|', label='P')\n",
     "    ax[0].scatter(None, None, color='b', marker='|', label='S')\n",
+    "    ax[0].xaxis.set_major_formatter(mdates.DateFormatter('%H:%M:%S'))\n",
+    "    ax[0].xaxis.set_major_locator(mdates.SecondLocator(interval=20))\n",
     "    ax[0].legend()\n",
     "\n",
     "    print('Plotting stations')\n",
     "    for i in range(1,3):\n",
     "        ax[i].scatter(st_lons, st_lats, color='b', label='Stations')\n",
     "        ax[i].scatter(eq_lon, eq_lat, color='r', marker='*', label='Earthquake')\n",
+    "        ax[i].set_aspect('equal')\n",
+    "        ax[i].set_xticklabels(ax[i].get_xticks(), rotation = 50)\n",
     "\n",
-    "    # Generate colorbar for the velocity plot\n",
-    "    cbar = plt.colorbar(ax[1].scatter(None, None, c=velocity_p, alpha=0.5, vmin=0, vmax=8), ax=ax[1])\n",
-    "    cbar.set_label('P Velocity (km/s)')\n",
-    "    ax[1].set_title('P Velocity')\n",
+    "    fig.subplots_adjust(bottom=0.1, top=0.9, left=0.1, right=0.8,\n",
+    "                    wspace=0.02, hspace=0.02)\n",
+    "    \n",
+    "    cb_ax = fig.add_axes([0.83, 0.1, 0.02, 0.8])\n",
+    "    cbar = fig.colorbar(ax[2].scatter(None, None, c=velocity_p, alpha=0.5, vmin=0, vmax=8), cax=cb_ax)\n",
     "\n",
-    "    cbar = plt.colorbar(ax[2].scatter(None, None, c=velocity_s, alpha=0.5, vmin=0, vmax=8), ax=ax[2])\n",
-    "    cbar.set_label('S Velocity (km/s)')\n",
+    "    cbar.set_label('Velocity (km/s)')\n",
+    "    ax[1].set_title('P Velocity')\n",
     "    ax[2].set_title('S Velocity')\n",
     "\n",
+    "    for a in ax:\n",
+    "        a.tick_params(axis='both', which='major', labelsize=8)\n",
+    "        \n",
     "    plt.subplots_adjust(hspace=0., wspace=0.5)\n",
-    "\n",
     "    fig.canvas.draw();\n",
     "    image = np.array(fig.canvas.renderer.buffer_rgba())\n",
     "    plt.close(fig)\n",
@@ -482,7 +527,6 @@
     "    }\n",
     "</style></h1> \n",
     "    \n",
-    "\n",
     "    <p style=\"font-size: 16px; margin-bottom: 20px;\">Detect <span style=\"background-image: linear-gradient(to right, #ED213A, #93291E); \n",
     "    -webkit-background-clip: text;\n",
     "    -webkit-text-fill-color: transparent;\n",
@@ -531,68 +575,66 @@
     "        </div>\n",
     "        \"\"\")\n",
     "        with gr.Row(): \n",
-    "            client_inputs = gr.Dropdown(\n",
-    "                choices = list(URL_MAPPINGS.keys()), \n",
-    "                label=\"FDSN Client\", \n",
-    "                info=\"Select one of the available FDSN clients\",\n",
-    "                value = \"IRIS\",\n",
-    "                interactive=True\n",
-    "            )\n",
-    "\n",
-    "            velocity_inputs = gr.Dropdown(\n",
-    "                choices = ['1066a', '1066b', 'ak135', \n",
-    "                        'ak135f', 'herrin', 'iasp91', \n",
-    "                        'jb', 'prem', 'pwdk'], \n",
-    "                label=\"1D velocity model\", \n",
-    "                info=\"Velocity model for station selection\",\n",
-    "                value = \"1066a\",\n",
-    "                interactive=True\n",
-    "            )\n",
+    "            with gr.Column(scale=2):\n",
+    "                client_inputs = gr.Dropdown(\n",
+    "                    choices = list(URL_MAPPINGS.keys()), \n",
+    "                    label=\"FDSN Client\", \n",
+    "                    info=\"Select one of the available FDSN clients\",\n",
+    "                    value = \"IRIS\",\n",
+    "                    interactive=True\n",
+    "                )\n",
+    "\n",
+    "                velocity_inputs = gr.Dropdown(\n",
+    "                    choices = ['1066a', '1066b', 'ak135', \n",
+    "                            'ak135f', 'herrin', 'iasp91', \n",
+    "                            'jb', 'prem', 'pwdk'], \n",
+    "                    label=\"1D velocity model\", \n",
+    "                    info=\"Velocity model for station selection\",\n",
+    "                    value = \"1066a\",\n",
+    "                    interactive=True\n",
+    "                )\n",
     "\n",
-    "            with gr.Column(scale=4):\n",
-    "                with gr.Row():   \n",
-    "                    timestamp_inputs = gr.Textbox(value='2019-07-04 17:33:49',\n",
-    "                                        placeholder='YYYY-MM-DD HH:MM:SS',\n",
-    "                                        label=\"Timestamp\",\n",
-    "                                        info=\"Timestamp of the earthquake\",\n",
-    "                                        max_lines=1,\n",
-    "                                        interactive=True)\n",
-    "                            \n",
-    "                    eq_lat_inputs = gr.Number(value=35.766, \n",
-    "                                    label=\"Latitude\", \n",
-    "                                    info=\"Latitude of the earthquake\",\n",
+    "            with gr.Column(scale=2):\n",
+    "                timestamp_inputs = gr.Textbox(value='2019-07-04 17:33:49',\n",
+    "                                    placeholder='YYYY-MM-DD HH:MM:SS',\n",
+    "                                    label=\"Timestamp\",\n",
+    "                                    info=\"Timestamp of the earthquake\",\n",
+    "                                    max_lines=1,\n",
+    "                                    interactive=True)\n",
+    "                \n",
+    "                source_depth_inputs = gr.Number(value=10,\n",
+    "                    label=\"Source depth (km)\",\n",
+    "                    info=\"Depth of the earthquake\",\n",
+    "                    interactive=True)\n",
+    "                \n",
+    "            with gr.Column(scale=2):\n",
+    "                eq_lat_inputs = gr.Number(value=35.766, \n",
+    "                                label=\"Latitude\", \n",
+    "                                info=\"Latitude of the earthquake\",\n",
+    "                                interactive=True)\n",
+    "                \n",
+    "                eq_lon_inputs = gr.Number(value=-117.605,\n",
+    "                                    label=\"Longitude\",\n",
+    "                                    info=\"Longitude of the earthquake\",\n",
     "                                    interactive=True)\n",
-    "                    \n",
-    "                    eq_lon_inputs = gr.Number(value=-117.605,\n",
-    "                                        label=\"Longitude\",\n",
-    "                                        info=\"Longitude of the earthquake\",\n",
-    "                                        interactive=True)\n",
-    "                    \n",
-    "                    source_depth_inputs = gr.Number(value=10,\n",
-    "                        label=\"Source depth (km)\",\n",
-    "                        info=\"Depth of the earthquake\",\n",
-    "                        interactive=True)\n",
     "                \n",
-    "\n",
-    "            \n",
     "            with gr.Column(scale=2):\n",
-    "                with gr.Row():   \n",
-    "                    radius_inputs = gr.Slider(minimum=1, \n",
-    "                                            maximum=150, \n",
-    "                                            value=50, label=\"Radius (km)\", \n",
-    "                                            step=10,\n",
-    "                                            info=\"\"\"Select the radius around the earthquake to download data from.\\n \n",
-    "                                            Note that the larger the radius, the longer the app will take to run.\"\"\",\n",
-    "                                            interactive=True)\n",
-    "                    \n",
-    "                    max_waveforms_inputs = gr.Slider(minimum=1,\n",
-    "                                    maximum=100,\n",
-    "                                    value=10,\n",
-    "                                    label=\"Max waveforms per section\",\n",
-    "                                    step=1,\n",
-    "                                    info=\"Maximum number of waveforms to show per section\\n (to avoid long prediction times)\",\n",
-    "                                    interactive=True,\n",
-    "                                    )\n",
+    "                radius_inputs = gr.Slider(minimum=1, \n",
+    "                                        maximum=200, \n",
+    "                                        value=50, label=\"Radius (km)\", \n",
+    "                                        step=10,\n",
+    "                                        info=\"\"\"Select the radius around the earthquake to download data from.\\n \n",
+    "                                        Note that the larger the radius, the longer the app will take to run.\"\"\",\n",
+    "                                        interactive=True)\n",
+    "                \n",
+    "                max_waveforms_inputs = gr.Slider(minimum=1,\n",
+    "                                maximum=100,\n",
+    "                                value=10,\n",
+    "                                label=\"Max waveforms per section\",\n",
+    "                                step=1,\n",
+    "                                info=\"Maximum number of waveforms to show per section\\n (to avoid long prediction times)\",\n",
+    "                                interactive=True,\n",
+    "                                )\n",
     "            \n",
     "        button = gr.Button(\"Predict phases\")\n",
     "        output_image = gr.Image(label='Waveforms with Phases Marked', type='numpy', interactive=False)\n",

app.py

@@ -21,7 +21,7 @@ from obspy.clients.fdsn.header import URL_MAPPINGS
 
 import matplotlib.pyplot as plt
 import matplotlib.dates as mdates
-from matplotlib.colors import LightSource
+from mpl_toolkits.axes_grid1 import ImageGrid
 
 from glob import glob
 
@@ -181,8 +181,8 @@ def predict_on_section(client_name, timestamp, eq_lat, eq_lon, radius_km, source
             
                 waveform = waveform.select(channel="H[BH][ZNE]")
                 waveform = waveform.merge(fill_value=0)
-                waveform = waveform[:3]
-                
+                waveform = waveform[:3].sort(keys=['channel'], reverse=True)
+
                 len_check = [len(x.data) for x in waveform]
                 if len(set(len_check)) > 1:
                     continue
@@ -243,8 +243,8 @@ def predict_on_section(client_name, timestamp, eq_lat, eq_lon, radius_km, source
     s_max_confidence = np.min([s_phases[i::len(waveforms)].std() for i in range(len(waveforms))])
 
     print(f"Starting plotting {len(waveforms)} waveforms")
-    fig, ax = plt.subplots(nrows=1, ncols=3, figsize=(10, 3))
-
+    fig, ax = plt.subplots(ncols=3, figsize=(10, 3))
+    
     # Plot topography
     print('Fetching topography')
     params = Topography.DEFAULT.copy()
@@ -289,9 +289,6 @@ def predict_on_section(client_name, timestamp, eq_lat, eq_lon, radius_km, source
         ax[0].scatter(x[int(current_S.mean()*waveforms[i][0].shape[-1])], waveforms[i][0, 0].mean()+distances[i]*111.2, color='b', alpha=s_conf, marker='|')
         ax[0].set_ylabel('Z')
 
-        ax[0].xaxis.set_major_formatter(mdates.DateFormatter('%H:%M:%S'))
-        ax[0].xaxis.set_major_locator(mdates.SecondLocator(interval=20))
-
         delta_t = t0s[i].timestamp - obspy.UTCDateTime(timestamp).timestamp
 
         velocity_p = (distances[i]*111.2)/(delta_t+current_P.mean()*60).item()
@@ -309,30 +306,37 @@ def predict_on_section(client_name, timestamp, eq_lat, eq_lon, radius_km, source
         y = np.linspace(st_lats[i], eq_lat, 50)
         
         # Plot the array
-        ax[1].scatter(x, y, c=np.zeros_like(x)+velocity_p, alpha=0.5, vmin=0, vmax=8)
-        ax[2].scatter(x, y, c=np.zeros_like(x)+velocity_s, alpha=0.5, vmin=0, vmax=8)
+        ax[1].scatter(x, y, c=np.zeros_like(x)+velocity_p, alpha=0.1, vmin=0, vmax=8)
+        ax[2].scatter(x, y, c=np.zeros_like(x)+velocity_s, alpha=0.1, vmin=0, vmax=8)
 
     # Add legend
     ax[0].scatter(None, None, color='r', marker='|', label='P')
     ax[0].scatter(None, None, color='b', marker='|', label='S')
+    ax[0].xaxis.set_major_formatter(mdates.DateFormatter('%H:%M:%S'))
+    ax[0].xaxis.set_major_locator(mdates.SecondLocator(interval=20))
     ax[0].legend()
 
     print('Plotting stations')
     for i in range(1,3):
         ax[i].scatter(st_lons, st_lats, color='b', label='Stations')
         ax[i].scatter(eq_lon, eq_lat, color='r', marker='*', label='Earthquake')
+        ax[i].set_aspect('equal')
+        ax[i].set_xticklabels(ax[i].get_xticks(), rotation = 50)
 
-    # Generate colorbar for the velocity plot
-    cbar = plt.colorbar(ax[1].scatter(None, None, c=velocity_p, alpha=0.5, vmin=0, vmax=8), ax=ax[1])
-    cbar.set_label('P Velocity (km/s)')
-    ax[1].set_title('P Velocity')
+    fig.subplots_adjust(bottom=0.1, top=0.9, left=0.1, right=0.8,
+                    wspace=0.02, hspace=0.02)
+    
+    cb_ax = fig.add_axes([0.83, 0.1, 0.02, 0.8])
+    cbar = fig.colorbar(ax[2].scatter(None, None, c=velocity_p, alpha=0.5, vmin=0, vmax=8), cax=cb_ax)
 
-    cbar = plt.colorbar(ax[2].scatter(None, None, c=velocity_s, alpha=0.5, vmin=0, vmax=8), ax=ax[2])
-    cbar.set_label('S Velocity (km/s)')
+    cbar.set_label('Velocity (km/s)')
+    ax[1].set_title('P Velocity')
     ax[2].set_title('S Velocity')
 
+    for a in ax:
+        a.tick_params(axis='both', which='major', labelsize=8)
+        
     plt.subplots_adjust(hspace=0., wspace=0.5)
-
     fig.canvas.draw();
     image = np.array(fig.canvas.renderer.buffer_rgba())
     plt.close(fig)
@@ -354,7 +358,6 @@ with gr.Blocks() as demo:
     }
 </style></h1> 
     
-
     <p style="font-size: 16px; margin-bottom: 20px;">Detect <span style="background-image: linear-gradient(to right, #ED213A, #93291E); 
     -webkit-background-clip: text;
     -webkit-text-fill-color: transparent;
@@ -393,77 +396,76 @@ with gr.Blocks() as demo:
         button.click(mark_phases, inputs=[inputs, upload], outputs=outputs)
         
     with gr.Tab("Select earthquake from catalogue"):
-        gr.Markdown("""
-        Select an earthquake from the global earthquake catalogue and the app will download the waveform from the FDSN client of your choice. 
-        The app will use a velocity model of your choice to select appropriate time windows for each station within specify radius of the earthquake.
-        The app will then analyze the waveforms and mark the detected phases on the waveform.
-        Pick data for each waveform is reported in seconds from the start of the waveform.
-        Velocities are derived from distance and travel time determined by PhaseHunter picks ($v = \mathrm{distance}/\mathrm{predicted_pick_time}$).
-        Backround of velocity plot is colored by DEM.
+
+        gr.HTML("""
+        <div style="padding: 20px; border-radius: 10px; font-size: 16px;">
+        <p style="font-weight: bold; font-size: 24px; margin-bottom: 20px;">Using PhaseHunter to Analyze Seismic Waveforms</p>
+        <p>Select an earthquake from the global earthquake catalogue and the app will download the waveform from the FDSN client of your choice. The app will use a velocity model of your choice to select appropriate time windows for each station within a specified radius of the earthquake.</p>
+        <p>The app will then analyze the waveforms and mark the detected phases on the waveform. Pick data for each waveform is reported in seconds from the start of the waveform.</p>
+        <p>Velocities are derived from distance and travel time determined by PhaseHunter picks (<span style="font-style: italic;">v = distance/predicted_pick_time</span>). The background of the velocity plot is colored by DEM.</p>
+        </div>
         """)
         with gr.Row(): 
-            client_inputs = gr.Dropdown(
-                choices = list(URL_MAPPINGS.keys()), 
-                label="FDSN Client", 
-                info="Select one of the available FDSN clients",
-                value = "IRIS",
-                interactive=True
-            )
-
-            velocity_inputs = gr.Dropdown(
-                choices = ['1066a', '1066b', 'ak135', 
-                        'ak135f', 'herrin', 'iasp91', 
-                        'jb', 'prem', 'pwdk'], 
-                label="1D velocity model", 
-                info="Velocity model for station selection",
-                value = "1066a",
-                interactive=True
-            )
+            with gr.Column(scale=2):
+                client_inputs = gr.Dropdown(
+                    choices = list(URL_MAPPINGS.keys()), 
+                    label="FDSN Client", 
+                    info="Select one of the available FDSN clients",
+                    value = "IRIS",
+                    interactive=True
+                )
+
+                velocity_inputs = gr.Dropdown(
+                    choices = ['1066a', '1066b', 'ak135', 
+                            'ak135f', 'herrin', 'iasp91', 
+                            'jb', 'prem', 'pwdk'], 
+                    label="1D velocity model", 
+                    info="Velocity model for station selection",
+                    value = "1066a",
+                    interactive=True
+                )
 
-            with gr.Column(scale=4):
-                with gr.Row():   
-                    timestamp_inputs = gr.Textbox(value='2019-07-04 17:33:49',
-                                        placeholder='YYYY-MM-DD HH:MM:SS',
-                                        label="Timestamp",
-                                        info="Timestamp of the earthquake",
-                                        max_lines=1,
-                                        interactive=True)
-                            
-                    eq_lat_inputs = gr.Number(value=35.766, 
-                                    label="Latitude", 
-                                    info="Latitude of the earthquake",
+            with gr.Column(scale=2):
+                timestamp_inputs = gr.Textbox(value='2019-07-04 17:33:49',
+                                    placeholder='YYYY-MM-DD HH:MM:SS',
+                                    label="Timestamp",
+                                    info="Timestamp of the earthquake",
+                                    max_lines=1,
                                     interactive=True)
-                    
-                    eq_lon_inputs = gr.Number(value=-117.605,
-                                        label="Longitude",
-                                        info="Longitude of the earthquake",
-                                        interactive=True)
-                    
-                    source_depth_inputs = gr.Number(value=10,
-                        label="Source depth (km)",
-                        info="Depth of the earthquake",
-                        interactive=True)
                 
-
-            
+                source_depth_inputs = gr.Number(value=10,
+                    label="Source depth (km)",
+                    info="Depth of the earthquake",
+                    interactive=True)
+                
             with gr.Column(scale=2):
-                with gr.Row():   
-                    radius_inputs = gr.Slider(minimum=1, 
-                                            maximum=150, 
-                                            value=50, label="Radius (km)", 
-                                            step=10,
-                                            info="""Select the radius around the earthquake to download data from.\n 
-                                            Note that the larger the radius, the longer the app will take to run.""",
-                                            interactive=True)
-                    
-                    max_waveforms_inputs = gr.Slider(minimum=1,
-                                    maximum=100,
-                                    value=10,
-                                    label="Max waveforms per section",
-                                    step=1,
-                                    info="Maximum number of waveforms to show per section\n (to avoid long prediction times)",
-                                    interactive=True,
-                                    )
+                eq_lat_inputs = gr.Number(value=35.766, 
+                                label="Latitude", 
+                                info="Latitude of the earthquake",
+                                interactive=True)
+                
+                eq_lon_inputs = gr.Number(value=-117.605,
+                                    label="Longitude",
+                                    info="Longitude of the earthquake",
+                                    interactive=True)
+                
+            with gr.Column(scale=2):
+                radius_inputs = gr.Slider(minimum=1, 
+                                        maximum=200, 
+                                        value=50, label="Radius (km)", 
+                                        step=10,
+                                        info="""Select the radius around the earthquake to download data from.\n 
+                                        Note that the larger the radius, the longer the app will take to run.""",
+                                        interactive=True)
+                
+                max_waveforms_inputs = gr.Slider(minimum=1,
+                                maximum=100,
+                                value=10,
+                                label="Max waveforms per section",
+                                step=1,
+                                info="Maximum number of waveforms to show per section\n (to avoid long prediction times)",
+                                interactive=True,
+                                )
             
         button = gr.Button("Predict phases")
         output_image = gr.Image(label='Waveforms with Phases Marked', type='numpy', interactive=False)

phasehunter/model.py

@@ -1,313 +0,0 @@
-import numpy as np
-import torch
-import torch.nn.functional as F
-from torch import nn
-from torchmetrics import MeanAbsoluteError
-from torch.optim.lr_scheduler import ReduceLROnPlateau
-
-import lightning as pl
-
-class BlurPool1D(nn.Module):
-    def __init__(self, channels, pad_type="reflect", filt_size=3, stride=2, pad_off=0):
-        super(BlurPool1D, self).__init__()
-        self.filt_size = filt_size
-        self.pad_off = pad_off
-        self.pad_sizes = [
-            int(1.0 * (filt_size - 1) / 2),
-            int(np.ceil(1.0 * (filt_size - 1) / 2)),
-        ]
-        self.pad_sizes = [pad_size + pad_off for pad_size in self.pad_sizes]
-        self.stride = stride
-        self.off = int((self.stride - 1) / 2.0)
-        self.channels = channels
-
-        # print('Filter size [%i]' % filt_size)
-        if self.filt_size == 1:
-            a = np.array(
-                [
-                    1.0,
-                ]
-            )
-        elif self.filt_size == 2:
-            a = np.array([1.0, 1.0])
-        elif self.filt_size == 3:
-            a = np.array([1.0, 2.0, 1.0])
-        elif self.filt_size == 4:
-            a = np.array([1.0, 3.0, 3.0, 1.0])
-        elif self.filt_size == 5:
-            a = np.array([1.0, 4.0, 6.0, 4.0, 1.0])
-        elif self.filt_size == 6:
-            a = np.array([1.0, 5.0, 10.0, 10.0, 5.0, 1.0])
-        elif self.filt_size == 7:
-            a = np.array([1.0, 6.0, 15.0, 20.0, 15.0, 6.0, 1.0])
-
-        filt = torch.Tensor(a)
-        filt = filt / torch.sum(filt)
-        self.register_buffer("filt", filt[None, None, :].repeat((self.channels, 1, 1)))
-
-        self.pad = get_pad_layer_1d(pad_type)(self.pad_sizes)
-
-    def forward(self, inp):
-        if self.filt_size == 1:
-            if self.pad_off == 0:
-                return inp[:, :, :: self.stride]
-            else:
-                return self.pad(inp)[:, :, :: self.stride]
-        else:
-            return F.conv1d(
-                self.pad(inp), self.filt, stride=self.stride, groups=inp.shape[1]
-            )
-
-
-def get_pad_layer_1d(pad_type):
-    if pad_type in ["refl", "reflect"]:
-        PadLayer = nn.ReflectionPad1d
-    elif pad_type in ["repl", "replicate"]:
-        PadLayer = nn.ReplicationPad1d
-    elif pad_type == "zero":
-        PadLayer = nn.ZeroPad1d
-    else:
-        print("Pad type [%s] not recognized" % pad_type)
-    return PadLayer
-
-
-from masksembles import common
-
-
-class Masksembles1D(nn.Module):
-    def __init__(self, channels: int, n: int, scale: float):
-        super().__init__()
-
-        self.channels = channels
-        self.n = n
-        self.scale = scale
-
-        masks = common.generation_wrapper(channels, n, scale)
-        masks = torch.from_numpy(masks)
-
-        self.masks = torch.nn.Parameter(masks, requires_grad=False)
-
-    def forward(self, inputs):
-        batch = inputs.shape[0]
-        x = torch.split(inputs.unsqueeze(1), batch // self.n, dim=0)
-        x = torch.cat(x, dim=1).permute([1, 0, 2, 3])
-        x = x * self.masks.unsqueeze(1).unsqueeze(-1)
-        x = torch.cat(torch.split(x, 1, dim=0), dim=1)
-
-        return x.squeeze(0).type(inputs.dtype)
-
-
-class BasicBlock(nn.Module):
-    expansion = 1
-
-    def __init__(self, in_planes, planes, stride=1, kernel_size=7, groups=1):
-        super(BasicBlock, self).__init__()
-        self.conv1 = nn.Conv1d(
-            in_planes,
-            planes,
-            kernel_size=kernel_size,
-            stride=stride,
-            padding="same",
-            bias=False,
-        )
-        self.bn1 = nn.BatchNorm1d(planes)
-        self.conv2 = nn.Conv1d(
-            planes,
-            planes,
-            kernel_size=kernel_size,
-            stride=1,
-            padding="same",
-            bias=False,
-        )
-        self.bn2 = nn.BatchNorm1d(planes)
-
-        self.shortcut = nn.Sequential(
-            nn.Conv1d(
-                in_planes,
-                self.expansion * planes,
-                kernel_size=1,
-                stride=stride,
-                padding="same",
-                bias=False,
-            ),
-            nn.BatchNorm1d(self.expansion * planes),
-        )
-
-    def forward(self, x):
-        out = F.relu(self.bn1(self.conv1(x)))
-        out = self.bn2(self.conv2(out))
-        out += self.shortcut(x)
-        out = F.relu(out)
-        return out
-
-
-class Updated_onset_picker(nn.Module):
-    def __init__(
-        self,
-    ):
-        super().__init__()
-
-        # self.activation = nn.ReLU()
-        # self.maxpool = nn.MaxPool1d(2)
-
-        self.n_masks = 128
-
-        self.block1 = nn.Sequential(
-            BasicBlock(3, 8, kernel_size=7, groups=1),
-            nn.GELU(),
-            BlurPool1D(8, filt_size=3, stride=2),
-            nn.GroupNorm(2, 8),
-        )
-
-        self.block2 = nn.Sequential(
-            BasicBlock(8, 16, kernel_size=7, groups=8),
-            nn.GELU(),
-            BlurPool1D(16, filt_size=3, stride=2),
-            nn.GroupNorm(2, 16),
-        )
-
-        self.block3 = nn.Sequential(
-            BasicBlock(16, 32, kernel_size=7, groups=16),
-            nn.GELU(),
-            BlurPool1D(32, filt_size=3, stride=2),
-            nn.GroupNorm(2, 32),
-        )
-
-        self.block4 = nn.Sequential(
-            BasicBlock(32, 64, kernel_size=7, groups=32),
-            nn.GELU(),
-            BlurPool1D(64, filt_size=3, stride=2),
-            nn.GroupNorm(2, 64),
-        )
-
-        self.block5 = nn.Sequential(
-            BasicBlock(64, 128, kernel_size=7, groups=64),
-            nn.GELU(),
-            BlurPool1D(128, filt_size=3, stride=2),
-            nn.GroupNorm(2, 128),
-        )
-
-        self.block6 = nn.Sequential(
-            Masksembles1D(128, self.n_masks, 2.0),
-            BasicBlock(128, 256, kernel_size=7, groups=128),
-            nn.GELU(),
-            BlurPool1D(256, filt_size=3, stride=2),
-            nn.GroupNorm(2, 256),
-        )
-
-        self.block7 = nn.Sequential(
-            Masksembles1D(256, self.n_masks, 2.0),
-            BasicBlock(256, 512, kernel_size=7, groups=256),
-            BlurPool1D(512, filt_size=3, stride=2),
-            nn.GELU(),
-            nn.GroupNorm(2, 512),
-        )
-
-        self.block8 = nn.Sequential(
-            Masksembles1D(512, self.n_masks, 2.0),
-            BasicBlock(512, 1024, kernel_size=7, groups=512),
-            BlurPool1D(1024, filt_size=3, stride=2),
-            nn.GELU(),
-            nn.GroupNorm(2, 1024),
-        )
-
-        self.block9 = nn.Sequential(
-            Masksembles1D(1024, self.n_masks, 2.0),
-            BasicBlock(1024, 128, kernel_size=7, groups=128),
-            # BlurPool1D(512, filt_size=3, stride=2),
-            # nn.GELU(),
-            # nn.GroupNorm(2,512),
-        )
-
-        self.out = nn.Sequential(nn.Linear(3072, 2), nn.Sigmoid())
-
-    def forward(self, x):
-        # Feature extraction
-
-        x = self.block1(x)
-        x = self.block2(x)
-
-        x = self.block3(x)
-        x = self.block4(x)
-
-        x = self.block5(x)
-        x = self.block6(x)
-
-        x = self.block7(x)
-        x = self.block8(x)
-
-        x = self.block9(x)
-
-        # Regressor
-        x = x.flatten(start_dim=1)
-        x = self.out(x)
-
-        return x
-
-class Onset_picker(pl.LightningModule):
-    def __init__(self, picker, learning_rate):
-        super().__init__()
-        self.picker = picker
-        self.learning_rate = learning_rate
-        self.save_hyperparameters(ignore=['picker'])
-        self.mae = MeanAbsoluteError()
-    
-    def compute_loss(self, y, pick, mae_name=False):
-        y_filt = y[y != 0]
-        pick_filt = pick[y != 0]
-        if len(y_filt) > 0:
-            loss = F.l1_loss(y_filt, pick_filt.flatten())
-            if mae_name != False:
-                mae_phase = self.mae(y_filt, pick_filt.flatten())*60
-                self.log(f'MAE/{mae_name}_val', mae_phase,  on_step=False, on_epoch=True, prog_bar=False, sync_dist=True)
-        else:
-            loss = 0
-        return loss
-            
-    def training_step(self, batch, batch_idx):
-        # training_step defines the train loop.
-        x, y_p, y_s = batch
-        # x, y_p, y_s, y_pg, y_sg, y_pn, y_sn = batch
-        
-        picks = self.picker(x)
-        
-        p_pick  = picks[:,0]
-        s_pick  = picks[:,1]
-        
-        p_loss = self.compute_loss(y_p, p_pick)
-        s_loss = self.compute_loss(y_s, s_pick)
-        
-        loss = (p_loss+s_loss)/2
-        
-        self.log('Loss/train', loss, on_step=True, on_epoch=False, prog_bar=True, sync_dist=True)
-        
-        return loss
-    
-    def validation_step(self, batch, batch_idx):
-        
-        x, y_p, y_s = batch
-        
-        picks = self.picker(x)
-        
-        p_pick  = picks[:,0]
-        s_pick  = picks[:,1]
-        
-        p_loss = self.compute_loss(y_p, p_pick, mae_name='P')
-        s_loss = self.compute_loss(y_s, s_pick, mae_name='S')
-        
-        loss = (p_loss+s_loss)/2
-            
-        self.log('Loss/val',  loss, on_step=False, on_epoch=True, prog_bar=False, sync_dist=True)
-        
-        return loss
-
-    def configure_optimizers(self):
-        optimizer = torch.optim.Adam(self.parameters(), lr=self.learning_rate)
-        scheduler = ReduceLROnPlateau(optimizer, mode='min', factor=0.5, patience=10, cooldown=10, threshold=1e-3)
-        # scheduler = torch.optim.lr_scheduler.OneCycleLR(optimizer, 3e-4,  epochs=300, steps_per_epoch=len(train_loader))
-        monitor = 'Loss/train'
-        return {"optimizer": optimizer,  "lr_scheduler": scheduler, 'monitor': monitor}
-    
-    def forward(self, x):
-        picks = self.picker(x)
-        return picks

phasehunter/training.py

@@ -1,104 +0,0 @@
-
-import torch
-
-from data_preparation import augment, collation_fn, my_split_by_node
-from model import Onset_picker, Updated_onset_picker
-
-import webdataset as wds
-
-from lightning.pytorch.callbacks import LearningRateMonitor, ModelCheckpoint
-from lightning.pytorch.loggers.tensorboard import TensorBoardLogger
-from lightning.pytorch.strategies import DDPStrategy
-from lightning import seed_everything
-import lightning as pl
-
-seed_everything(42, workers=False)
-torch.set_float32_matmul_precision('medium')
-
-batch_size = 256
-num_workers = 16 #int(os.cpu_count())
-n_iters_in_epoch = 5000
-
-train_dataset = (
-      wds.WebDataset("data/sample/shard-00{0000..0001}.tar", 
-                     # splitter=my_split_by_worker, 
-                     nodesplitter=my_split_by_node)
-      .decode()
-      .map(augment)
-      .shuffle(5000)
-      .batched(batchsize=batch_size,
-               collation_fn=collation_fn,
-               partial=False
-              )
-).with_epoch(n_iters_in_epoch//num_workers)
-
-
-val_dataset = (
-      wds.WebDataset("data/sample/shard-00{0000..0000}.tar", 
-                     # splitter=my_split_by_worker, 
-                     nodesplitter=my_split_by_node)
-      .decode()
-      .map(augment)
-      .repeat()
-      .batched(batchsize=batch_size,
-               collation_fn=collation_fn,
-               partial=False
-              )
-).with_epoch(100)
-
-
-train_loader = wds.WebLoader(train_dataset, 
-                             num_workers=num_workers, 
-                             shuffle=False,
-                             pin_memory=True, 
-                             batch_size=None)
-
-val_loader = wds.WebLoader(val_dataset, 
-                           num_workers=0,
-                           shuffle=False,
-                           pin_memory=True, 
-                           batch_size=None) 
-
-
-
-# model
-model = Onset_picker(picker=Updated_onset_picker(), 
-                     learning_rate=3e-4)
-# model = torch.compile(model, mode="reduce-overhead")
-
-logger = TensorBoardLogger("tensorboard_logdir", name="FAST")
-
-checkpoint_callback = ModelCheckpoint(save_top_k=1, monitor="Loss/val", filename="chkp-{epoch:02d}")
-lr_callback = LearningRateMonitor(logging_interval='epoch')
-# swa_callback = StochasticWeightAveraging(swa_lrs=0.05)
-    
-# # train model
-trainer = pl.Trainer(
-            precision='16-mixed',
-            
-            callbacks=[checkpoint_callback, lr_callback],
-    
-            devices='auto', 
-            accelerator='auto', 
-    
-            strategy=DDPStrategy(find_unused_parameters=False,
-                                 static_graph=True,
-                                 gradient_as_bucket_view=True),
-            benchmark=True,
-            
-            gradient_clip_val=0.5,
-            # ckpt_path='path/to/saved/checkpoints/chkp.ckpt',
-
-            # fast_dev_run=True,
-
-            logger=logger,
-            log_every_n_steps=50,
-            enable_progress_bar=True,
-    
-            max_epochs=300,
-        )
-
-trainer.fit(model=model, 
-            train_dataloaders=train_loader, 
-            val_dataloaders=val_loader,
-            )