WIP on section plot

app.py

@@ -78,32 +78,89 @@ def mark_phases(waveform):
     plt.close(fig)
     return image
 
-#??
-
-def download_data(timestamp, eq_lat, eq_lon, client_name, radius_km):
+def predict_on_section(client_name, timestamp, eq_lat, eq_lon, radius_km, source_depth_km, velocity_model):
+    distances, t0s, st_lats, st_lons, waveforms = [], [], [], [], []
+    
+    taup_model = TauPyModel(model=velocity_model)
     client = Client(client_name)
+
     window = radius_km / 111.2
 
     assert eq_lat - window > -90 and eq_lat + window < 90, "Latitude out of bounds"
     assert eq_lon - window > -180 and eq_lon + window < 180, "Longitude out of bounds"
 
     starttime = obspy.UTCDateTime(timestamp)
-    endtime = startime + 120
+    endtime = starttime + 120
 
     inv = client.get_stations(network="*", station="*", location="*", channel="*H*", 
-                          starttime=obspy.UTCDateTime(starttime), endtime=endtime, 
-                          minlatitude=eq_lat-window, maxlatitude=eq_lat+window,
-                          minlongitude=eq_lon-window, maxlongitude=eq_lon+window, 
-                          level='channel')
+                          starttime=starttime, endtime=endtime, 
+                          minlatitude=(eq_lat-window), maxlatitude=(eq_lat+window),
+                          minlongitude=(eq_lon-window), maxlongitude=(eq_lon+window), 
+                          level='station')
     
+    waveforms = []
     for network in inv:
         for station in network:
-            print(station)
+            try:
+                distance = locations2degrees(eq_lat, eq_lon, station.latitude, station.longitude)
+
+                arrivals = taup_model.get_travel_times(source_depth_in_km=source_depth_km, 
+                                                       distance_in_degree=distance, 
+                                                       phase_list=["P", "S"])
+
+                if len(arrivals) > 0:
+
+                    starttime = obspy.UTCDateTime(timestamp) + arrivals[0].time - 15
+                    endtime = starttime + 60
+
+                    waveform = client.get_waveforms(network=network.code, station=station.code, location="*", channel="*", 
+                                                starttime=starttime, endtime=endtime)
+                    
+                    waveform = waveform.select(channel="H[BH][ZNE]")
+                    waveform = waveform.merge(fill_value=0)
+                    waveform = waveform[:3]
+                    
+                    len_check = [len(x.data) for x in waveform]
+                    if len(set(len_check)) > 1:
+                        continue
+
+                    if len(waveform) == 3:
+                        waveform = prepare_waveform(np.stack([x.data for x in waveform]))
+                    
+                        distances.append(distance)
+                        t0s.append(starttime)
+                        st_lats.append(station.latitude)
+                        st_lons.append(station.longitude)
+                        waveforms.append(waveform)
+
+            except (IndexError, FDSNNoDataException, FDSNTimeoutException):
+                continue
+
+    with torch.no_grad():
+        waveforms_torch = torch.vstack(waveforms)
+        output = model(waveforms_torch)
+
+    p_phases = output[:, 0]
+    s_phases = output[:, 1]
+
+
+    print(p_phases.shape)
+    # for i in range(len(waveforms)):
+    #     current_P = P_batch[i::len(waveforms)].cpu()
+    #     current_S_batch = S_batch[i::len(waveforms)].cpu()
+    #     current_Pg_batch = Pg_batch[i::len(waveforms)].cpu()
+    #     current_Sg_batch = Sg_batch[i::len(waveforms)].cpu()
+    #     current_Pn_batch = Pn_batch[i::len(waveforms)].cpu()
+    #     current_Sn_batch = Sn_batch[i::len(waveforms)].cpu()
+    
+    fig,ax = plt.subplots()
+    ax.scatter(st_lats, st_lons)
+    fig.canvas.draw()
+    image = np.array(fig.canvas.renderer.buffer_rgba())
+    plt.close(fig)
 
-            # waveform = client.get_waveforms(network=network.code, station=station.code, location="*", channel="*", 
-            #                             starttime=obspy.UTCDateTime(start_date), endtime=obspy.UTCDateTime(end_date))
+    return image
 
-    return 0
 
 model = Onset_picker.load_from_checkpoint("./weights.ckpt",
                                  picker=Updated_onset_picker(),
@@ -119,10 +176,10 @@ model.eval()
 with gr.Blocks() as demo:
     gr.Markdown("# PhaseHunter")
     gr.Markdown("""This app allows one to detect P and S seismic phases along with uncertainty of the detection. 
-                The app can be used in three ways: either by selecting one of the sample waveforms;
-                                                   or by selecting an earthquake from the global earthquake catalogue;
-                                                   or by uploading a waveform of interest.
-                """)
+            The app can be used in three ways: either by selecting one of the sample waveforms;
+                                                or by selecting an earthquake from the global earthquake catalogue;
+                                                or by uploading a waveform of interest.
+            """)
     with gr.Tab("Default example"):
         # Define the input and output types for Gradio
         inputs = gr.Dropdown(
@@ -163,19 +220,39 @@ with gr.Blocks() as demo:
                             info="Latitude of the earthquake",
                             interactive=True)
             
-            eq_lo_inputs = gr.Number(value=117.605,
+            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)
+            
             radius_inputs = gr.Slider(minimum=1, 
                                     maximum=150, 
                                     value=50, label="Radius (km)", 
                                     info="Select the radius around the earthquake to download data from",
                                     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
+            )
+            
+            
         button = gr.Button("Predict phases")
-        button.click(mark_phases, inputs=inputs, outputs=outputs)
+        outputs_section = gr.outputs.Image(label='Waveforms with Phases Marked', type='numpy')
+        
+        button.click(predict_on_section, 
+                 inputs=[client_inputs, timestamp_inputs, 
+                         eq_lat_inputs, eq_lon_inputs, 
+                         radius_inputs, source_depth_inputs, velocity_inputs],
+                 outputs=outputs_section)
 
     with gr.Tab("Predict on your own waveform"):
         gr.Markdown("""
@@ -183,6 +260,6 @@ with gr.Blocks() as demo:
         Your waveform should be sampled at 100 sps and have 3 (Z, N, E) or 1 (Z) channels.
         """)
 
-    button.click(download_data, inputs=[timestamp_inputs, eq_lat_inputs,eq_lo_inputs, radius_inputs], outputs=outputs)
+
 
 demo.launch()