Gradio_app.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Running on local URL:  http://127.0.0.1:7865\n",
      "\n",
      "To create a public link, set `share=True` in `launch()`.\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<div><iframe src=\"http://127.0.0.1:7865/\" width=\"100%\" height=\"500\" allow=\"autoplay; camera; microphone; clipboard-read; clipboard-write;\" frameborder=\"0\" allowfullscreen></iframe></div>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": []
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Starting to download inventory\n",
      "Finished downloading inventory\n",
      "Processing CI.CCC...\n",
      "Downloading waveform for CI_CCC_2019-07-04T17:33:40.494920Z\n",
      "Skipping CI_CCC_2019-07-04T17:33:40.494920Z\n",
      "Processing CI.CLC...\n",
      "Processing CI.JRC2...\n",
      "Reading cached waveform\n",
      "Added CI.JRC2 to the list of waveforms\n",
      "Processing CI.LRL...\n",
      "Reading cached waveform\n",
      "Added CI.LRL to the list of waveforms\n",
      "Processing CI.MPM...\n",
      "Reading cached waveform\n",
      "Processing CI.Q0072...\n",
      "Reading cached waveform\n",
      "Processing CI.SLA...\n",
      "Reading cached waveform\n",
      "Added CI.SLA to the list of waveforms\n",
      "Processing CI.SRT...\n",
      "Reading cached waveform\n",
      "Added CI.SRT to the list of waveforms\n",
      "Processing CI.TOW2...\n",
      "Reading cached waveform\n",
      "Added CI.TOW2 to the list of waveforms\n",
      "Processing CI.WBM...\n",
      "Downloading waveform for CI_WBM_2019-07-04T17:33:40.063616Z\n",
      "Skipping CI_WBM_2019-07-04T17:33:40.063616Z\n",
      "Processing CI.WCS2...\n",
      "Downloading waveform for CI_WCS2_2019-07-04T17:33:40.200958Z\n"
     ]
    }
   ],
   "source": [
    "# Gradio app that takes seismic waveform as input and marks 2 phases on the waveform as output.\n",
    "\n",
    "import gradio as gr\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "from phasehunter.data_preparation import prepare_waveform\n",
    "import torch\n",
    "import io\n",
    "\n",
    "from scipy.stats import gaussian_kde\n",
    "from scipy.signal import resample\n",
    "from scipy.interpolate import interp1d\n",
    "\n",
    "from bmi_topography import Topography\n",
    "import earthpy.spatial as es\n",
    "\n",
    "import obspy\n",
    "from obspy.clients.fdsn import Client\n",
    "from obspy.clients.fdsn.header import FDSNNoDataException, FDSNTimeoutException, FDSNInternalServerException\n",
    "from obspy.geodetics.base import locations2degrees\n",
    "from obspy.taup import TauPyModel\n",
    "from obspy.taup.helper_classes import SlownessModelError\n",
    "\n",
    "from obspy.clients.fdsn.header import URL_MAPPINGS\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "import matplotlib.dates as mdates\n",
    "from mpl_toolkits.axes_grid1 import ImageGrid\n",
    "\n",
    "from glob import glob\n",
    "\n",
    "import numpy as np\n",
    "from matplotlib import colors, cm\n",
    "from scipy.interpolate import griddata\n",
    "\n",
    "def resample_waveform(waveform, original_freq, target_freq):\n",
    "    \"\"\"\n",
    "    Resample a waveform from original frequency to target frequency using SciPy's resample function.\n",
    "    \n",
    "    Args:\n",
    "    waveform (numpy.ndarray): The input waveform as a 1D array.\n",
    "    original_freq (float): The original sampling frequency of the waveform.\n",
    "    target_freq (float): The target sampling frequency of the waveform.\n",
    "    \n",
    "    Returns:\n",
    "    resampled_waveform (numpy.ndarray): The resampled waveform as a 1D array.\n",
    "    \"\"\"\n",
    "    # Calculate the resampling ratio\n",
    "    resampling_ratio = target_freq / original_freq\n",
    "    # Calculate the new length of the resampled waveform\n",
    "    resampled_length = int(waveform.shape[-1] * resampling_ratio)\n",
    "    # Resample the waveform using SciPy's resample function\n",
    "    resampled_waveform = resample(waveform, resampled_length, axis=-1)\n",
    "    \n",
    "    return resampled_waveform\n",
    "\n",
    "def sort_channels_to_ZNE(waveform, channels):\n",
    "    # Input:\n",
    "    # waveform: a 2D numpy array with shape (3, n), where n is the number of samples\n",
    "    # channels: a list or tuple of 3 strings representing the channel order, e.g. ('N', 'Z', 'E')\n",
    "    channels = list(channels)\n",
    "\n",
    "    if len(channels) != 3 or set(channels) != {'Z', 'N', 'E'}:\n",
    "        raise ValueError(\"Invalid channel input. It should be a permutation of 'Z', 'N', and 'E'.\")\n",
    "\n",
    "    # Find the indices of the Z, N, and E channels\n",
    "    z_index = channels.index('Z')\n",
    "    n_index = channels.index('N')\n",
    "    e_index = channels.index('E')\n",
    "    \n",
    "    print(z_index, n_index, e_index)\n",
    "    # Sort the channels to ZNE\n",
    "    sorted_waveform = waveform[[z_index, n_index, e_index], :]\n",
    "    \n",
    "    return sorted_waveform\n",
    "\n",
    "def make_prediction(waveform, sampling_rate, order):\n",
    "    waveform = np.load(waveform)\n",
    "    print('Loaded', waveform.shape)\n",
    "\n",
    "    if len(waveform.shape) == 1:\n",
    "        waveform = waveform.reshape(1, waveform.shape[0])\n",
    "\n",
    "    elif waveform.shape[0] == 3:\n",
    "        waveform = sort_channels_to_ZNE(waveform, order)\n",
    "\n",
    "    if sampling_rate != 100:\n",
    "        waveform = resample_waveform(waveform, sampling_rate, 100)\n",
    "        print('Resampled', waveform.shape)\n",
    "\n",
    "\n",
    "    orig_waveform = waveform[:, :6000].copy()\n",
    "    processed_input = prepare_waveform(waveform)\n",
    "\n",
    "    # Make prediction\n",
    "    with torch.inference_mode():\n",
    "        output = model(processed_input)\n",
    "\n",
    "    p_phase = output[:, 0]\n",
    "    s_phase = output[:, 1]\n",
    "\n",
    "    return processed_input, p_phase, s_phase, orig_waveform\n",
    "\n",
    "\n",
    "def mark_phases(waveform, uploaded_file, p_thres, s_thres, sampling_rate, order):\n",
    "\n",
    "    if uploaded_file is not None:\n",
    "        waveform = uploaded_file.name\n",
    "\n",
    "    processed_input, p_phase, s_phase, orig_waveform = make_prediction(waveform, sampling_rate, order)\n",
    "\n",
    "    # Create a plot of the waveform with the phases marked\n",
    "    if sum(processed_input[0][2] == 0): #if input is 1C\n",
    "        fig, ax = plt.subplots(nrows=2, figsize=(10, 2), sharex=True)\n",
    "\n",
    "        ax[0].plot(orig_waveform[0], color='black', lw=1)\n",
    "        ax[0].set_ylabel('Norm. Ampl.')\n",
    "\n",
    "    else: #if input is 3C\n",
    "        fig, ax = plt.subplots(nrows=4, figsize=(10, 6), sharex=True)\n",
    "        ax[0].plot(orig_waveform[0], color='black', lw=1)\n",
    "        ax[1].plot(orig_waveform[1], color='black', lw=1)\n",
    "        ax[2].plot(orig_waveform[2], color='black', lw=1)\n",
    "\n",
    "        ax[0].set_ylabel('Z')\n",
    "        ax[1].set_ylabel('N')\n",
    "        ax[2].set_ylabel('E')\n",
    "\n",
    "\n",
    "    do_we_have_p = (p_phase.std().item()*60 < p_thres)\n",
    "    if do_we_have_p:\n",
    "        p_phase_plot = p_phase*processed_input.shape[-1]\n",
    "        p_kde = gaussian_kde(p_phase_plot)\n",
    "        p_dist_space = np.linspace( min(p_phase_plot)-10, max(p_phase_plot)+10, 500 )\n",
    "        ax[-1].plot( p_dist_space, p_kde(p_dist_space), color='r')\n",
    "    else:\n",
    "        ax[-1].text(0.5, 0.75, 'No P phase detected', horizontalalignment='center', verticalalignment='center', transform=ax[-1].transAxes)\n",
    "\n",
    "    do_we_have_s = (s_phase.std().item()*60 < s_thres)\n",
    "    if do_we_have_s:\n",
    "        s_phase_plot = s_phase*processed_input.shape[-1]\n",
    "        s_kde = gaussian_kde(s_phase_plot)\n",
    "        s_dist_space = np.linspace( min(s_phase_plot)-10, max(s_phase_plot)+10, 500 )\n",
    "        ax[-1].plot( s_dist_space, s_kde(s_dist_space), color='b')\n",
    "\n",
    "        for a in ax:\n",
    "            a.axvline(p_phase.mean()*processed_input.shape[-1], color='r', linestyle='--', label='P', alpha=do_we_have_p)\n",
    "            a.axvline(s_phase.mean()*processed_input.shape[-1], color='b', linestyle='--', label='S', alpha=do_we_have_s)\n",
    "    else:\n",
    "        ax[-1].text(0.5, 0.25, 'No S phase detected', horizontalalignment='center', verticalalignment='center', transform=ax[-1].transAxes)\n",
    "\n",
    "    ax[-1].set_xlabel('Time, samples')\n",
    "    ax[-1].set_ylabel('Uncert., samples')\n",
    "    ax[-1].legend()\n",
    "\n",
    "    plt.subplots_adjust(hspace=0., wspace=0.)\n",
    "\n",
    "    # Convert the plot to an image and return it\n",
    "    fig.canvas.draw()\n",
    "    image = np.array(fig.canvas.renderer.buffer_rgba())\n",
    "    plt.close(fig)\n",
    "    return image\n",
    "\n",
    "def bin_distances(distances, bin_size=10):\n",
    "    # Bin the distances into groups of `bin_size` kilometers\n",
    "    binned_distances = {}\n",
    "    for i, distance in enumerate(distances):\n",
    "        bin_index = distance // bin_size\n",
    "        if bin_index not in binned_distances:\n",
    "            binned_distances[bin_index] = (distance, i)\n",
    "        elif i < binned_distances[bin_index][1]:\n",
    "            binned_distances[bin_index] = (distance, i)\n",
    "\n",
    "    # Select the first distance in each bin and its index\n",
    "    first_distances = []\n",
    "    for bin_index in binned_distances:\n",
    "        first_distance, first_distance_index = binned_distances[bin_index]\n",
    "        first_distances.append(first_distance_index)\n",
    "    \n",
    "    return first_distances\n",
    "\n",
    "def variance_coefficient(residuals):\n",
    "    # calculate the variance of the residuals\n",
    "    var = residuals.var()\n",
    "    # scale the variance to a coefficient between 0 and 1\n",
    "    coeff = 1 - (var / (residuals.max() - residuals.min()))\n",
    "    return coeff\n",
    "\n",
    "def predict_on_section(client_name, timestamp, eq_lat, eq_lon, radius_km, source_depth_km, velocity_model, max_waveforms, conf_thres_P, conf_thres_S):\n",
    "    distances, t0s, st_lats, st_lons, waveforms, names = [], [], [], [], [], []\n",
    "    \n",
    "    taup_model = TauPyModel(model=velocity_model)\n",
    "    client = Client(client_name)\n",
    "\n",
    "    window = radius_km / 111.2\n",
    "    max_waveforms = int(max_waveforms)\n",
    "\n",
    "    assert eq_lat - window > -90 and eq_lat + window < 90, \"Latitude out of bounds\"\n",
    "    assert eq_lon - window > -180 and eq_lon + window < 180, \"Longitude out of bounds\"\n",
    "\n",
    "    starttime = obspy.UTCDateTime(timestamp)\n",
    "    endtime = starttime + 120\n",
    "\n",
    "    try:\n",
    "        print('Starting to download inventory')\n",
    "        inv = client.get_stations(network=\"*\", station=\"*\", location=\"*\", channel=\"*H*\", \n",
    "                            starttime=starttime, endtime=endtime, \n",
    "                            minlatitude=(eq_lat-window), maxlatitude=(eq_lat+window),\n",
    "                            minlongitude=(eq_lon-window), maxlongitude=(eq_lon+window), \n",
    "                            level='station')\n",
    "        print('Finished downloading inventory')\n",
    "        \n",
    "    except (IndexError, FDSNNoDataException, FDSNTimeoutException, FDSNInternalServerException):\n",
    "        fig, ax = plt.subplots()\n",
    "        ax.text(0.5,0.5,'Something is wrong with the data provider, try another')\n",
    "        fig.canvas.draw();\n",
    "        image = np.array(fig.canvas.renderer.buffer_rgba())\n",
    "        plt.close(fig)\n",
    "        return image\n",
    "    \n",
    "    waveforms = []\n",
    "    cached_waveforms = glob(\"data/cached/*.mseed\")\n",
    "\n",
    "    for network in inv:\n",
    "        if network.code == 'SY':\n",
    "            continue\n",
    "        for station in network:\n",
    "            print(f\"Processing {network.code}.{station.code}...\")\n",
    "            distance = locations2degrees(eq_lat, eq_lon, station.latitude, station.longitude)\n",
    "\n",
    "            arrivals = taup_model.get_travel_times(source_depth_in_km=source_depth_km, \n",
    "                                                    distance_in_degree=distance, \n",
    "                                                    phase_list=[\"P\", \"S\"])\n",
    "\n",
    "            if len(arrivals) > 0:\n",
    "\n",
    "                starttime = obspy.UTCDateTime(timestamp) + arrivals[0].time - 15\n",
    "                endtime = starttime + 60\n",
    "                try:\n",
    "                    filename=f'{network.code}_{station.code}_{starttime}'\n",
    "                    if f\"data/cached/{filename}.mseed\" not in cached_waveforms:\n",
    "                        print(f'Downloading waveform for {filename}')\n",
    "                        waveform = client.get_waveforms(network=network.code, station=station.code, location=\"*\", channel=\"*\", \n",
    "                                                    starttime=starttime, endtime=endtime)\n",
    "                        waveform.write(f\"data/cached/{network.code}_{station.code}_{starttime}.mseed\", format=\"MSEED\")\n",
    "                        print('Finished downloading and caching waveform')\n",
    "                    else:\n",
    "                        print('Reading cached waveform')\n",
    "                        waveform = obspy.read(f\"data/cached/{network.code}_{station.code}_{starttime}.mseed\")\n",
    "                        \n",
    "\n",
    "                except (IndexError, FDSNNoDataException, FDSNTimeoutException, FDSNInternalServerException):\n",
    "                    print(f'Skipping {network.code}_{station.code}_{starttime}')\n",
    "                    continue\n",
    "            \n",
    "                waveform = waveform.select(channel=\"H[BH][ZNE]\")\n",
    "                waveform = waveform.merge(fill_value=0)\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",
    "\n",
    "                if len(waveform) == 3:\n",
    "                    try:\n",
    "                        waveform = prepare_waveform(np.stack([x.data for x in waveform]))\n",
    "\n",
    "                        distances.append(distance)\n",
    "                        t0s.append(starttime)\n",
    "                        st_lats.append(station.latitude)\n",
    "                        st_lons.append(station.longitude)\n",
    "                        waveforms.append(waveform)\n",
    "                        names.append(f\"{network.code}.{station.code}\")\n",
    "\n",
    "                        print(f\"Added {network.code}.{station.code} to the list of waveforms\")\n",
    "\n",
    "                    except:\n",
    "                        continue\n",
    "                \n",
    "    \n",
    "    # If there are no waveforms, return an empty plot\n",
    "    if len(waveforms) == 0:\n",
    "        print('No waveforms found')\n",
    "        fig, ax = plt.subplots()\n",
    "        # prints \"No waveforms found\" on the plot aligned at center and vertically\n",
    "        ax.text(0.5,0.5,'No waveforms found', horizontalalignment='center', verticalalignment='center', transform=ax.transAxes)\n",
    "        fig.canvas.draw();\n",
    "        image = np.array(fig.canvas.renderer.buffer_rgba())\n",
    "        plt.close(fig)\n",
    "\n",
    "        output_picks = pd.DataFrame()\n",
    "        output_picks.to_csv('data/picks.csv', index=False)\n",
    "        output_csv = 'data/picks.csv'\n",
    "        return image, output_picks, output_csv\n",
    "    \n",
    "\n",
    "    first_distances = bin_distances(distances, bin_size=10/111.2)\n",
    "\n",
    "    # Edge case when there are way too many waveforms to process\n",
    "    selection_indexes = np.random.choice(first_distances, \n",
    "                                         np.min([len(first_distances), max_waveforms]),\n",
    "                                         replace=False)\n",
    "\n",
    "    waveforms = np.array(waveforms)[selection_indexes]\n",
    "    distances = np.array(distances)[selection_indexes]\n",
    "    t0s = np.array(t0s)[selection_indexes]\n",
    "    st_lats = np.array(st_lats)[selection_indexes]\n",
    "    st_lons = np.array(st_lons)[selection_indexes]\n",
    "    names = np.array(names)[selection_indexes]\n",
    "\n",
    "    waveforms = [torch.tensor(waveform) for waveform in waveforms]\n",
    "\n",
    "    print('Starting to run predictions')\n",
    "    with torch.no_grad():\n",
    "        waveforms_torch = torch.vstack(waveforms)\n",
    "        output = model(waveforms_torch)\n",
    "\n",
    "    p_phases = output[:, 0]\n",
    "    s_phases = output[:, 1]\n",
    "\n",
    "    p_phases = p_phases.reshape(len(waveforms),-1)\n",
    "    s_phases = s_phases.reshape(len(waveforms),-1)\n",
    "\n",
    "    # Max confidence - min variance    \n",
    "    p_max_confidence = p_phases.std(axis=-1).min()\n",
    "    s_max_confidence = s_phases.std(axis=-1).min()\n",
    "\n",
    "    print(f\"Starting plotting {len(waveforms)} waveforms\")\n",
    "    fig, ax = plt.subplots(ncols=3, figsize=(10, 3))\n",
    "    \n",
    "    # Plot topography\n",
    "    print('Fetching topography')\n",
    "    params = Topography.DEFAULT.copy()\n",
    "    extra_window = 0.5\n",
    "    params[\"south\"] = np.min([st_lats.min(), eq_lat])-extra_window\n",
    "    params[\"north\"] = np.max([st_lats.max(), eq_lat])+extra_window\n",
    "    params[\"west\"] = np.min([st_lons.min(), eq_lon])-extra_window\n",
    "    params[\"east\"] = np.max([st_lons.max(), eq_lon])+extra_window\n",
    "\n",
    "    topo_map = Topography(**params)\n",
    "    topo_map.fetch()\n",
    "    topo_map.load()\n",
    "\n",
    "    print('Plotting topo')\n",
    "    hillshade = es.hillshade(topo_map.da[0], altitude=10)\n",
    "    \n",
    "    topo_map.da.plot(ax = ax[1], cmap='Greys', add_colorbar=False, add_labels=False)\n",
    "    topo_map.da.plot(ax = ax[2], cmap='Greys', add_colorbar=False, add_labels=False)\n",
    "    ax[1].imshow(hillshade, cmap=\"Greys\", alpha=0.5)\n",
    "\n",
    "    output_picks = pd.DataFrame({'station_name' : [], \n",
    "                                'st_lat' : [], 'st_lon' : [],\n",
    "                                 'starttime' : [], \n",
    "                                 'p_phase, s' : [], 'p_uncertainty, s' : [], \n",
    "                                 's_phase, s' : [], 's_uncertainty, s' : [],\n",
    "                                 'velocity_p, km/s' : [], 'velocity_s, km/s' : []})\n",
    "                        \n",
    "    for i in range(len(waveforms)):\n",
    "        print(f\"Plotting waveform {i+1}/{len(waveforms)}\")\n",
    "        current_P = p_phases[i]\n",
    "        current_S = s_phases[i]\n",
    "        \n",
    "        x = [t0s[i] + pd.Timedelta(seconds=k/100) for k in np.linspace(0,6000,6000)]\n",
    "        x = mdates.date2num(x)\n",
    "\n",
    "        # Normalize confidence for the plot\n",
    "        p_conf = 1/(current_P.std()/p_max_confidence).item()\n",
    "        s_conf = 1/(current_S.std()/s_max_confidence).item()\n",
    "\n",
    "        delta_t = t0s[i].timestamp - obspy.UTCDateTime(timestamp).timestamp\n",
    "\n",
    "        ax[0].plot(x, waveforms[i][0, 0]*10+distances[i]*111.2, color='black', alpha=0.5, lw=1)\n",
    "\n",
    "        if (current_P.std().item()*60 < conf_thres_P) or (current_S.std().item()*60 < conf_thres_S):\n",
    "            ax[0].scatter(x[int(current_P.mean()*waveforms[i][0].shape[-1])], waveforms[i][0, 0].mean()+distances[i]*111.2, color='r', alpha=p_conf, marker='|')\n",
    "            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",
    "        \n",
    "            velocity_p = (distances[i]*111.2)/(delta_t+current_P.mean()*60).item()\n",
    "            velocity_s = (distances[i]*111.2)/(delta_t+current_S.mean()*60).item()\n",
    "\n",
    "            # Generate an array from st_lat to eq_lat and from st_lon to eq_lon\n",
    "            x = np.linspace(st_lons[i], eq_lon, 50)\n",
    "            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.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",
    "        else:\n",
    "            velocity_p = np.nan\n",
    "            velocity_s = np.nan\n",
    "        \n",
    "        ax[0].set_ylabel('Z')\n",
    "        print(f\"Station {st_lats[i]}, {st_lons[i]} has P velocity {velocity_p} and S velocity {velocity_s}\")\n",
    "\n",
    "        output_picks = output_picks.append(pd.DataFrame({'station_name': [names[i]], \n",
    "                                                        'st_lat' : [st_lats[i]], 'st_lon' : [st_lons[i]],\n",
    "                                                        'starttime' : [str(t0s[i])], \n",
    "                                                        'p_phase, s' : [(delta_t+current_P.mean()*60).item()], 'p_uncertainty, s' : [current_P.std().item()*60], \n",
    "                                                        's_phase, s' : [(delta_t+current_S.mean()*60).item()], 's_uncertainty, s' : [current_S.std().item()*60],\n",
    "                                                        'velocity_p, km/s' : [velocity_p], 'velocity_s, km/s' : [velocity_s]}))\n",
    "        \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",
    "    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.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",
    "    fig.canvas.draw();\n",
    "    image = np.array(fig.canvas.renderer.buffer_rgba())\n",
    "    plt.close(fig)\n",
    "\n",
    "    output_csv = f'data/velocity/{eq_lat}_{eq_lon}_{source_depth_km}_{timestamp}_{len(waveforms)}.csv'\n",
    "    output_picks.to_csv(output_csv, index=False)\n",
    "    \n",
    "    return image, output_picks, output_csv\n",
    "\n",
    "def interpolate_vel_model(velocity_model, initial_velocity, lat_values, lon_values, depth_values, n_lat, n_lon, n_depth):\n",
    "    # Create a mask for points with the initial velocity\n",
    "    initial_velocity_mask = (velocity_model == initial_velocity)\n",
    "\n",
    "    # Find the indices of points with non-initial velocities\n",
    "    non_initial_velocity_indices = np.argwhere(~initial_velocity_mask)\n",
    "\n",
    "    # Extract the coordinates and corresponding velocities of the known points\n",
    "    known_points = np.column_stack([lat_values[non_initial_velocity_indices[:, 0]],\n",
    "                                    lon_values[non_initial_velocity_indices[:, 1]],\n",
    "                                    depth_values[non_initial_velocity_indices[:, 2]]])\n",
    "    \n",
    "    # Find the maximum depth in the known_points\n",
    "    max_known_depth = np.max(known_points[:, 2])\n",
    "\n",
    "    known_velocities = velocity_model[~initial_velocity_mask]\n",
    "\n",
    "    # Create a grid of points for the entire volume\n",
    "    grid_points = np.array(np.meshgrid(lat_values, lon_values, depth_values, indexing='ij')).reshape(3, -1).T\n",
    "\n",
    "    # Create a mask for grid points that are deeper than the maximum known depth\n",
    "    depth_mask = grid_points[:, 2] <= max_known_depth\n",
    "\n",
    "    # Interpolate the velocities at the grid points\n",
    "    interpolated_velocities = griddata(known_points, known_velocities, grid_points[depth_mask], method='linear')\n",
    "\n",
    "    # Fill nan values with the nearest known velocities\n",
    "    interpolated_velocities_filled = griddata(known_points, known_velocities, grid_points[depth_mask], method='nearest')\n",
    "    interpolated_velocities[np.isnan(interpolated_velocities)] = interpolated_velocities_filled[np.isnan(interpolated_velocities)]\n",
    "\n",
    "    # Initialize an array with the same length as grid_points and fill it with nan values\n",
    "    interpolated_velocities_with_depth_limit = np.full(grid_points.shape[0], np.nan)\n",
    "\n",
    "    # Update the array with the interpolated velocities for the masked grid points\n",
    "    interpolated_velocities_with_depth_limit[depth_mask] = interpolated_velocities\n",
    "\n",
    "    # Reshape the interpolated velocities to match the shape of the velocity_model\n",
    "    interpolated_velocity_model = interpolated_velocities_with_depth_limit.reshape(n_lat, n_lon, n_depth)\n",
    "\n",
    "    return interpolated_velocity_model\n",
    "\n",
    "\n",
    "# Function to find the closest index for a given value in an array\n",
    "def find_closest_index(array, value):\n",
    "    return np.argmin(np.abs(array - value))\n",
    "\n",
    "# FIX AFTER CONFERENCE\n",
    "# def compute_velocity_model(azimuth, elevation, interpolate, n_lat, n_lon, n_depth):\n",
    "#     filename = list(output_csv.temp_files)[0]\n",
    "    \n",
    "#     df = pd.read_csv(filename)\n",
    "#     filename = filename.split('/')[-1]\n",
    "    \n",
    "#     # Current EQ location\n",
    "#     eq_lat = float(filename.split(\"_\")[0])\n",
    "#     eq_lon = float(filename.split(\"_\")[1])\n",
    "#     eq_depth = float(filename.split(\"_\")[2])\n",
    "\n",
    "#     # Define the region of interest (latitude, longitude, and depth ranges)\n",
    "#     lat_range = (np.min([df.st_lat.min(), eq_lat]), np.max([df.st_lat.max(), eq_lat]))\n",
    "#     lon_range = (np.min([df.st_lon.min(), eq_lon]),  np.max([df.st_lon.max(), eq_lon]))\n",
    "#     depth_range = (0, 50)\n",
    "\n",
    "#     # Define the number of nodes in each dimension\n",
    "#     num_points = 100\n",
    "\n",
    "#     taup_model = TauPyModel(model='1066a')\n",
    "\n",
    "#     # Create the grid\n",
    "#     lat_values = np.linspace(lat_range[0], lat_range[1], n_lat)\n",
    "#     lon_values = np.linspace(lon_range[0], lon_range[1], n_lon)\n",
    "#     depth_values = np.linspace(depth_range[0], depth_range[1], n_depth)\n",
    "\n",
    "#     # Initialize the velocity model with constant values\n",
    "#     initial_velocity = 0  # km/s, this can be P-wave or S-wave velocity\n",
    "#     velocity_model = np.full((n_lat, n_lon, n_depth), initial_velocity, dtype=float)\n",
    "\n",
    "#     # Loop through the stations and update the velocity model\n",
    "#     for i in range(len(df)):\n",
    "#         if ~np.isnan(df['velocity_p, km/s'].iloc[i]):\n",
    "\n",
    "#                 ray_path = taup_model.get_ray_paths_geo(source_depth_in_km=eq_depth,\n",
    "#                                         source_latitude_in_deg=eq_lat,\n",
    "#                                         source_longitude_in_deg=eq_lon,\n",
    "#                                         receiver_latitude_in_deg=df.st_lat.iloc[i],\n",
    "#                                         receiver_longitude_in_deg=df.st_lon.iloc[i],\n",
    "#                                         phase_list=['P', 'S'])\n",
    "                \n",
    "#                 # THERE IS A PROBLEM WITH THE RAY PATHS. APPARENTLY LAT AND LON DON'T EXIST (HOW DID IT WORK BEFORE?)\n",
    "#                 print(ray_path[0].path)\n",
    "\n",
    "#                 # Create the interpolator objects for latitude, longitude, and depth\n",
    "#                 interp_latitude = interp1d(np.linspace(0, ray_path[0].path['lat'].max(), len(ray_path[0].path['lat'])), ray_path[0].path['lat'])\n",
    "#                 interp_longitude = interp1d(np.linspace(0, ray_path[0].path['lon'].max(), len(ray_path[0].path['lon'])), ray_path[0].path['lon'])\n",
    "#                 interp_depth = interp1d(np.linspace(0, ray_path[0].path['depth'].max(), len(ray_path[0].path['depth'])), ray_path[0].path['depth'])\n",
    "\n",
    "#                 # Resample the ray path to N points\n",
    "#                 lat_values_interp = interp_latitude(np.linspace(0, ray_path[0].path['lat'].max(), num_points))\n",
    "#                 lon_values_interp = interp_longitude(np.linspace(0, ray_path[0].path['lon'].max(), num_points))\n",
    "#                 depth_values_interp = interp_depth(np.linspace(0, ray_path[0].path['depth'].max(), num_points))\n",
    "\n",
    "#                 # Loop through the interpolated coordinates and update the grid cells with the average P-wave velocity\n",
    "#                 for lat, lon, depth in zip(lat_values_interp, lon_values_interp, depth_values_interp):\n",
    "#                     lat_index = find_closest_index(lat_values, lat)\n",
    "#                     lon_index = find_closest_index(lon_values, lon)\n",
    "#                     depth_index = find_closest_index(depth_values, depth)\n",
    "                    \n",
    "#                     if velocity_model[lat_index, lon_index, depth_index] == initial_velocity:\n",
    "#                         velocity_model[lat_index, lon_index, depth_index] = df['velocity_p, km/s'].iloc[i]\n",
    "#                     else:\n",
    "#                         velocity_model[lat_index, lon_index, depth_index] = (velocity_model[lat_index, lon_index, depth_index] +\n",
    "#                                                                             df['velocity_p, km/s'].iloc[i]) / 2\n",
    "\n",
    "#     # Create the figure and axis\n",
    "#     fig = plt.figure(figsize=(8, 8))\n",
    "#     ax = fig.add_subplot(111, projection='3d')\n",
    "\n",
    "#     # Set the plot limits\n",
    "#     ax.set_xlim3d(lat_range[0], lat_range[1])\n",
    "#     ax.set_ylim3d(lon_range[0], lon_range[1])\n",
    "#     ax.set_zlim3d(depth_range[1], depth_range[0])\n",
    "\n",
    "#     ax.set_xlabel('Latitude')\n",
    "#     ax.set_ylabel('Longitude')\n",
    "#     ax.set_zlabel('Depth (km)')\n",
    "#     ax.set_title('Velocity Model')\n",
    "    \n",
    "#     # Create the meshgrid\n",
    "#     x, y, z = np.meshgrid(\n",
    "#         np.linspace(lat_range[0], lat_range[1], velocity_model.shape[0]+1),\n",
    "#         np.linspace(lon_range[0], lon_range[1], velocity_model.shape[1]+1),\n",
    "#         np.linspace(depth_range[0], depth_range[1], velocity_model.shape[2]+1),\n",
    "#         indexing='ij'\n",
    "#     )\n",
    "\n",
    "#     # Create the color array\n",
    "#     norm = plt.Normalize(vmin=2, vmax=8)\n",
    "#     colors_vel = plt.cm.plasma(norm(velocity_model)) \n",
    "    \n",
    "#     # Plot the voxels\n",
    "#     if interpolate:\n",
    "#         interpolated_velocity_model = interpolate_vel_model(velocity_model, initial_velocity, lat_values, lon_values, depth_values, n_lat, n_lon, n_depth)\n",
    "#         colors_interp = plt.cm.plasma(norm(interpolated_velocity_model))\n",
    "#         ax.voxels(x, y, z, interpolated_velocity_model > 0, facecolors=colors_interp, alpha=0.5, edgecolor='k')\n",
    "    \n",
    "#     ax.voxels(x, y, z, velocity_model > 0, facecolors=colors_vel, alpha=1, edgecolor='black')\n",
    "\n",
    "#     # Set the view angle\n",
    "#     ax.view_init(elev=elevation, azim=azimuth)\n",
    "\n",
    "#     m = cm.ScalarMappable(cmap=plt.cm.plasma, norm=norm)\n",
    "#     m.set_array([])\n",
    "#     plt.colorbar(m)\n",
    "\n",
    "#     # Show the plot\n",
    "#     fig.canvas.draw();\n",
    "#     image = np.array(fig.canvas.renderer.buffer_rgba())\n",
    "#     plt.close(fig)\n",
    "\n",
    "#     return image\n",
    "\n",
    "# model = torch.jit.load(\"model.pt\")\n",
    "model = torch.jit.load(\"model.pt\")\n",
    "model.eval()\n",
    "\n",
    "with gr.Blocks() as demo:\n",
    "    gr.HTML(\"\"\"\n",
    "<div style=\"padding: 20px; border-radius: 10px;\">\n",
    "    <h1 style=\"font-size: 30px; text-align: center; margin-bottom: 20px;\">PhaseHunter <span style=\"animation: arrow-anim 10s linear infinite; display: inline-block; transform: rotate(45deg) translateX(-20px);\">🏹</span>\n",
    "\n",
    "<style>\n",
    "    @keyframes arrow-anim {\n",
    "        0% { transform: translateX(-20px); }\n",
    "        50% { transform: translateX(20px); }\n",
    "        100% { transform: translateX(-20px); }\n",
    "    }\n",
    "</style></h1> \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",
    "    background-clip: text;\">P</span> and <span style=\"background-image: linear-gradient(to right, #00B4DB, #0083B0); \n",
    "    -webkit-background-clip: text;\n",
    "    -webkit-text-fill-color: transparent;\n",
    "    background-clip: text;\">S</span> seismic phases with <span style=\"background-image: linear-gradient(to right, #f12711, #f5af19); \n",
    "    -webkit-background-clip: text;\n",
    "    -webkit-text-fill-color: transparent;\n",
    "    background-clip: text;\">uncertainty</span></p>\n",
    "    <ul style=\"font-size: 16px; margin-bottom: 40px;\">\n",
    "        <li>Tab 1: Detect seismic phases by selecting a sample waveform or uploading your own waveform in <code>.npy</code> format.</li>\n",
    "        <li>Tab 2: Select an earthquake from the global earthquake catalogue and PhaseHunter will analyze seismic stations in the given radius.</li>\n",
    "        <li>Waveforms should be sampled at 100 samples/sec and have 3 (Z, N, E) or 1 (Z) channels. PhaseHunter analyzes the first 6000 samples of your file.</li>\n",
    "    </ul>\n",
    "    <p style=\"font-size: 16px; margin-bottom: 20px;\">Please contact me at [email protected] with questions and feedback</p>\n",
    "</div>\n",
    "\"\"\")\n",
    "    with gr.Tab(\"Try on a single station\"):\n",
    "        with gr.Row(): \n",
    "            # Define the input and output types for Gradio\n",
    "            inputs = gr.Dropdown(\n",
    "                [\"data/sample/sample_0.npy\", \n",
    "                \"data/sample/sample_1.npy\", \n",
    "                \"data/sample/sample_2.npy\"], \n",
    "                label=\"Sample waveform\", \n",
    "                info=\"Select one of the samples\",\n",
    "                value = \"data/sample/sample_0.npy\"\n",
    "            )\n",
    "            with gr.Column(scale=1):\n",
    "                P_thres_inputs = gr.Slider(minimum=0.01,\n",
    "                                    maximum=1,\n",
    "                                    value=0.1,\n",
    "                                    label=\"P uncertainty threshold (s)\",\n",
    "                                    step=0.01,\n",
    "                                    info=\"Acceptable uncertainty for P picks expressed in std() seconds\",\n",
    "                                    interactive=True,\n",
    "                                    )\n",
    "                \n",
    "                S_thres_inputs = gr.Slider(minimum=0.01,\n",
    "                                    maximum=1,\n",
    "                                    value=0.2,\n",
    "                                    label=\"S uncertainty threshold (s)\",\n",
    "                                    step=0.01,\n",
    "                                    info=\"Acceptable uncertainty for S picks expressed in std() seconds\",\n",
    "                                    interactive=True,\n",
    "                                    )\n",
    "            with gr.Column(scale=1):\n",
    "                upload = gr.File(label=\"Upload your waveform\")\n",
    "                with gr.Row():\n",
    "                    sampling_rate_inputs = gr.Slider(minimum=10,\n",
    "                            maximum=1000,\n",
    "                            value=100,\n",
    "                            label=\"Samlping rate, Hz\",\n",
    "                            step=10,\n",
    "                            info=\"Sampling rate of the waveform\",\n",
    "                            interactive=True,\n",
    "                            )\n",
    "                    order_input = gr.Text(value='ZNE', \n",
    "                                        label='Channel order', \n",
    "                                        info='Order of the channels in the waveform file (e.g. ZNE)')\n",
    "\n",
    "        button = gr.Button(\"Predict phases\")\n",
    "        outputs = gr.Image(label='Waveform with Phases Marked', type='numpy', interactive=False)\n",
    "    \n",
    "        button.click(mark_phases, inputs=[inputs, upload, \n",
    "                                          P_thres_inputs, S_thres_inputs,\n",
    "                                          sampling_rate_inputs, order_input], \n",
    "                                          outputs=outputs)    \n",
    "    with gr.Tab(\"Select earthquake from catalogue\"):\n",
    "\n",
    "        gr.HTML(\"\"\"\n",
    "        <div style=\"padding: 20px; border-radius: 10px; font-size: 16px;\">\n",
    "        <p style=\"font-weight: bold; font-size: 24px; margin-bottom: 20px;\">Using PhaseHunter to Analyze Seismic Waveforms</p>\n",
    "        <p>Select an earthquake from the global earthquake catalogue (e.g. <a href=\"https://earthquake.usgs.gov/earthquakes/map\">USGS</a>) 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>\n",
    "        <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>\n",
    "        <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>\n",
    "        </div>\n",
    "        \"\"\")\n",
    "        with gr.Row(): \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=2):\n",
    "                timestamp_inputs = gr.Textbox(value='2019-07-04T17:33:49-00',\n",
    "                                    placeholder='YYYY-MM-DDTHH:MM:SS-TZ',\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",
    "            with gr.Column(scale=2):\n",
    "                radius_inputs = gr.Slider(minimum=1, \n",
    "                                        maximum=200, \n",
    "                                        value=50, \n",
    "                                        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",
    "            with gr.Column(scale=2):\n",
    "                P_thres_inputs = gr.Slider(minimum=0.01,\n",
    "                                maximum=1,\n",
    "                                value=0.1,\n",
    "                                label=\"P uncertainty threshold, s\",\n",
    "                                step=0.01,\n",
    "                                info=\"Acceptable uncertainty for P picks expressed in std() seconds\",\n",
    "                                interactive=True,\n",
    "                                )\n",
    "                S_thres_inputs = gr.Slider(minimum=0.01,\n",
    "                                maximum=1,\n",
    "                                value=0.2,\n",
    "                                label=\"S uncertainty threshold, s\",\n",
    "                                step=0.01,\n",
    "                                info=\"Acceptable uncertainty for S picks expressed in std() seconds\",\n",
    "                                interactive=True,\n",
    "                                )\n",
    "            \n",
    "        button_phases = gr.Button(\"Predict phases\")\n",
    "        output_image = gr.Image(label='Waveforms with Phases Marked', type='numpy', interactive=False)\n",
    "        \n",
    "        # with gr.Row():\n",
    "        #     with gr.Column(scale=2):\n",
    "        #         azimuth_input = gr.Slider(minimum=-180, maximum=180, value=0, step=5, label=\"Azimuth\", interactive=True)\n",
    "        #         elevation_input = gr.Slider(minimum=-90, maximum=90, value=30, step=5, label=\"Elevation\", interactive=True)\n",
    "\n",
    "        #         with gr.Row():\n",
    "        #             interpolate_input = gr.Checkbox(label=\"Interpolate\", info=\"Interpolate velocity model\")\n",
    "        #             n_lat_input = gr.Slider(minimum=5, maximum=100, value=50, step=5, label=\"N lat\", info='Number of Lat grid points', interactive=True)\n",
    "        #             n_lon_input = gr.Slider(minimum=5, maximum=100, value=50, step=5, label=\"N lon\", info='Number of Lon grid points', interactive=True)\n",
    "        #             n_depth_input = gr.Slider(minimum=5, maximum=100, value=50, step=5, label=\"N depth\", info='Number of Depth grid points', interactive=True)\n",
    "                \n",
    "        #         button = gr.Button(\"Look at 3D Velocities\")\n",
    "        #     outputs_vel_model = gr.Image(label=\"3D Velocity Model\")\n",
    "\n",
    "        # button.click(compute_velocity_model, \n",
    "        #                 inputs=[azimuth_input, elevation_input, \n",
    "        #                         interpolate_input, n_lat_input, \n",
    "        #                         n_lon_input, n_depth_input], \n",
    "        #                 outputs=[outputs_vel_model])\n",
    "        \n",
    "        with gr.Row():\n",
    "            output_picks = gr.Dataframe(label='Pick data', \n",
    "                                        type='pandas', \n",
    "                                        interactive=False)\n",
    "            output_csv = gr.File(label=\"Output File\", file_types=[\".csv\"])\n",
    "\n",
    "        button_phases.click(predict_on_section, \n",
    "                 inputs=[client_inputs, timestamp_inputs, \n",
    "                         eq_lat_inputs, eq_lon_inputs, \n",
    "                         radius_inputs, source_depth_inputs, \n",
    "                         velocity_inputs, max_waveforms_inputs,\n",
    "                         P_thres_inputs, S_thres_inputs],\n",
    "                 outputs=[output_image, output_picks, output_csv])\n",
    "\n",
    "demo.launch()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def compute_velocity_model(azimuth, elevation, interpolate, n_lat, n_lon, n_depth):\n",
    "    filename = list(output_csv.temp_files)[0]\n",
    "    \n",
    "    df = pd.read_csv(filename)\n",
    "    filename = filename.split('/')[-1]\n",
    "    \n",
    "    # Current EQ location\n",
    "    eq_lat = float(filename.split(\"_\")[0])\n",
    "    eq_lon = float(filename.split(\"_\")[1])\n",
    "    eq_depth = float(filename.split(\"_\")[2])\n",
    "\n",
    "    # Define the region of interest (latitude, longitude, and depth ranges)\n",
    "    lat_range = (np.min([df.st_lat.min(), eq_lat]), np.max([df.st_lat.max(), eq_lat]))\n",
    "    lon_range = (np.min([df.st_lon.min(), eq_lon]),  np.max([df.st_lon.max(), eq_lon]))\n",
    "    depth_range = (0, 50)\n",
    "\n",
    "    # Define the number of nodes in each dimension\n",
    "    num_points = 100\n",
    "\n",
    "    taup_model = TauPyModel(model='1066a')\n",
    "\n",
    "    # Create the grid\n",
    "    lat_values = np.linspace(lat_range[0], lat_range[1], n_lat)\n",
    "    lon_values = np.linspace(lon_range[0], lon_range[1], n_lon)\n",
    "    depth_values = np.linspace(depth_range[0], depth_range[1], n_depth)\n",
    "\n",
    "    # Initialize the velocity model with constant values\n",
    "    initial_velocity = 0  # km/s, this can be P-wave or S-wave velocity\n",
    "    velocity_model = np.full((n_lat, n_lon, n_depth), initial_velocity, dtype=float)\n",
    "\n",
    "    # Loop through the stations and update the velocity model\n",
    "    for i in range(len(df)):\n",
    "        if ~np.isnan(df['velocity_p, km/s'].iloc[i]):\n",
    "\n",
    "                ray_path = taup_model.get_ray_paths_geo(source_depth_in_km=eq_depth,\n",
    "                                        source_latitude_in_deg=eq_lat,\n",
    "                                        source_longitude_in_deg=eq_lon,\n",
    "                                        receiver_latitude_in_deg=df.st_lat.iloc[i],\n",
    "                                        receiver_longitude_in_deg=df.st_lon.iloc[i],\n",
    "                                        phase_list=['P', 'S'])\n",
    "                \n",
    "                # HERE IS THE PROBLEM\n",
    "                print(ray_path[0].path)\n",
    "\n",
    "                # Create the interpolator objects for latitude, longitude, and depth\n",
    "                interp_latitude = interp1d(np.linspace(0, ray_path[0].path['lat'].max(), len(ray_path[0].path['lat'])), ray_path[0].path['lat'])\n",
    "                interp_longitude = interp1d(np.linspace(0, ray_path[0].path['lon'].max(), len(ray_path[0].path['lon'])), ray_path[0].path['lon'])\n",
    "                interp_depth = interp1d(np.linspace(0, ray_path[0].path['depth'].max(), len(ray_path[0].path['depth'])), ray_path[0].path['depth'])\n",
    "\n",
    "                # Resample the ray path to N points\n",
    "                lat_values_interp = interp_latitude(np.linspace(0, ray_path[0].path['lat'].max(), num_points))\n",
    "                lon_values_interp = interp_longitude(np.linspace(0, ray_path[0].path['lon'].max(), num_points))\n",
    "                depth_values_interp = interp_depth(np.linspace(0, ray_path[0].path['depth'].max(), num_points))\n",
    "\n",
    "                # Loop through the interpolated coordinates and update the grid cells with the average P-wave velocity\n",
    "                for lat, lon, depth in zip(lat_values_interp, lon_values_interp, depth_values_interp):\n",
    "                    lat_index = find_closest_index(lat_values, lat)\n",
    "                    lon_index = find_closest_index(lon_values, lon)\n",
    "                    depth_index = find_closest_index(depth_values, depth)\n",
    "                    \n",
    "                    if velocity_model[lat_index, lon_index, depth_index] == initial_velocity:\n",
    "                        velocity_model[lat_index, lon_index, depth_index] = df['velocity_p, km/s'].iloc[i]\n",
    "                    else:\n",
    "                        velocity_model[lat_index, lon_index, depth_index] = (velocity_model[lat_index, lon_index, depth_index] +\n",
    "                                                                            df['velocity_p, km/s'].iloc[i]) / 2\n",
    "\n",
    "    # Create the figure and axis\n",
    "    fig = plt.figure(figsize=(8, 8))\n",
    "    ax = fig.add_subplot(111, projection='3d')\n",
    "\n",
    "    # Set the plot limits\n",
    "    ax.set_xlim3d(lat_range[0], lat_range[1])\n",
    "    ax.set_ylim3d(lon_range[0], lon_range[1])\n",
    "    ax.set_zlim3d(depth_range[1], depth_range[0])\n",
    "\n",
    "    ax.set_xlabel('Latitude')\n",
    "    ax.set_ylabel('Longitude')\n",
    "    ax.set_zlabel('Depth (km)')\n",
    "    ax.set_title('Velocity Model')\n",
    "    \n",
    "    # Create the meshgrid\n",
    "    x, y, z = np.meshgrid(\n",
    "        np.linspace(lat_range[0], lat_range[1], velocity_model.shape[0]+1),\n",
    "        np.linspace(lon_range[0], lon_range[1], velocity_model.shape[1]+1),\n",
    "        np.linspace(depth_range[0], depth_range[1], velocity_model.shape[2]+1),\n",
    "        indexing='ij'\n",
    "    )\n",
    "\n",
    "    # Create the color array\n",
    "    norm = plt.Normalize(vmin=2, vmax=8)\n",
    "    colors_vel = plt.cm.plasma(norm(velocity_model)) \n",
    "    \n",
    "    # Plot the voxels\n",
    "    if interpolate:\n",
    "        interpolated_velocity_model = interpolate_vel_model(velocity_model, initial_velocity, lat_values, lon_values, depth_values, n_lat, n_lon, n_depth)\n",
    "        colors_interp = plt.cm.plasma(norm(interpolated_velocity_model))\n",
    "        ax.voxels(x, y, z, interpolated_velocity_model > 0, facecolors=colors_interp, alpha=0.5, edgecolor='k')\n",
    "    \n",
    "    ax.voxels(x, y, z, velocity_model > 0, facecolors=colors_vel, alpha=1, edgecolor='black')\n",
    "\n",
    "    # Set the view angle\n",
    "    ax.view_init(elev=elevation, azim=azimuth)\n",
    "\n",
    "    m = cm.ScalarMappable(cmap=plt.cm.plasma, norm=norm)\n",
    "    m.set_array([])\n",
    "    plt.colorbar(m)\n",
    "\n",
    "    # Show the plot\n",
    "    fig.canvas.draw();\n",
    "    image = np.array(fig.canvas.renderer.buffer_rgba())\n",
    "    plt.close(fig)\n",
    "\n",
    "    return image"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.8"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}