Added ray tracing to velocity modelling

app.py

@@ -9,6 +9,8 @@ import io
 
 from scipy.stats import gaussian_kde
 from scipy.signal import resample
+from scipy.interpolate import interp1d
+
 from bmi_topography import Topography
 import earthpy.spatial as es
 
@@ -435,16 +437,58 @@ def predict_on_section(client_name, timestamp, eq_lat, eq_lon, radius_km, source
 
 import numpy as np
 from matplotlib import colors, cm
+from scipy.interpolate import griddata
+
+def interpolate_vel_model(velocity_model, initial_velocity, lat_values, lon_values, depth_values, n_lat, n_lon, n_depth):
+    # Create a mask for points with the initial velocity
+    initial_velocity_mask = (velocity_model == initial_velocity)
+
+    # Find the indices of points with non-initial velocities
+    non_initial_velocity_indices = np.argwhere(~initial_velocity_mask)
+
+    # Extract the coordinates and corresponding velocities of the known points
+    known_points = np.column_stack([lat_values[non_initial_velocity_indices[:, 0]],
+                                    lon_values[non_initial_velocity_indices[:, 1]],
+                                    depth_values[non_initial_velocity_indices[:, 2]]])
+    
+    # Find the maximum depth in the known_points
+    max_known_depth = np.max(known_points[:, 2])
+
+    known_velocities = velocity_model[~initial_velocity_mask]
+
+    # Create a grid of points for the entire volume
+    grid_points = np.array(np.meshgrid(lat_values, lon_values, depth_values, indexing='ij')).reshape(3, -1).T
+
+    # Create a mask for grid points that are deeper than the maximum known depth
+    depth_mask = grid_points[:, 2] <= max_known_depth
+
+    # Interpolate the velocities at the grid points
+    interpolated_velocities = griddata(known_points, known_velocities, grid_points[depth_mask], method='linear')
+
+    # Fill nan values with the nearest known velocities
+    interpolated_velocities_filled = griddata(known_points, known_velocities, grid_points[depth_mask], method='nearest')
+    interpolated_velocities[np.isnan(interpolated_velocities)] = interpolated_velocities_filled[np.isnan(interpolated_velocities)]
+
+    # Initialize an array with the same length as grid_points and fill it with nan values
+    interpolated_velocities_with_depth_limit = np.full(grid_points.shape[0], np.nan)
+
+    # Update the array with the interpolated velocities for the masked grid points
+    interpolated_velocities_with_depth_limit[depth_mask] = interpolated_velocities
+
+    # Reshape the interpolated velocities to match the shape of the velocity_model
+    interpolated_velocity_model = interpolated_velocities_with_depth_limit.reshape(n_lat, n_lon, n_depth)
+
+    return interpolated_velocity_model
+
 
 # Function to find the closest index for a given value in an array
 def find_closest_index(array, value):
     return np.argmin(np.abs(array - value))
 
-def compute_velocity_model(azimuth, elevation):
+def compute_velocity_model(azimuth, elevation, interpolate, n_lat, n_lon, n_depth):
     filename = list(output_csv.temp_files)[0]
     
     df = pd.read_csv(filename)
-    print(df)
     filename = filename.split('/')[-1]
     
     # Current EQ location
@@ -458,11 +502,10 @@ def compute_velocity_model(azimuth, elevation):
     depth_range = (0, 50)
 
     # Define the number of nodes in each dimension
-    n_lat = 10
-    n_lon = 10
-    n_depth = 10
     num_points = 100
 
+    taup_model = TauPyModel(model='1066a')
+
     # Create the grid
     lat_values = np.linspace(lat_range[0], lat_range[1], n_lat)
     lon_values = np.linspace(lon_range[0], lon_range[1], n_lon)
@@ -475,23 +518,36 @@ def compute_velocity_model(azimuth, elevation):
     # Loop through the stations and update the velocity model
     for i in range(len(df)):
         if ~np.isnan(df['velocity_p, km/s'].iloc[i]):
-            # Interpolate coordinates along the great circle path between the earthquake and the station
-            lon_deg = np.linspace(df.st_lon.iloc[i], eq_lon, num_points)
-            lat_deg = np.linspace(df.st_lat.iloc[i], eq_lat, num_points)
-            depth_interpolated = np.interp(np.linspace(0, 1, num_points), [0, 1], [eq_depth, 0])
-
-            # Loop through the interpolated coordinates and update the grid cells with the average P-wave velocity
-            for lat, lon, depth in zip(lat_deg, lon_deg, depth_interpolated):
-                lat_index = find_closest_index(lat_values, lat)
-                lon_index = find_closest_index(lon_values, lon)
-                depth_index = find_closest_index(depth_values, depth)
-                
-                if velocity_model[lat_index, lon_index, depth_index] == initial_velocity:
-                    velocity_model[lat_index, lon_index, depth_index] = df['velocity_p, km/s'].iloc[i]
-                else:
-                    velocity_model[lat_index, lon_index, depth_index] = (velocity_model[lat_index, lon_index, depth_index] +
-                                                                        df['velocity_p, km/s'].iloc[i]) / 2
+
+                ray_path = taup_model.get_ray_paths_geo(source_depth_in_km=eq_depth,
+                                        source_latitude_in_deg=eq_lat,
+                                        source_longitude_in_deg=eq_lon,
+                                        receiver_latitude_in_deg=df.st_lat.iloc[i],
+                                        receiver_longitude_in_deg=df.st_lon.iloc[i],
+                                        phase_list=['P', 'S'])
+
+                # Create the interpolator objects for latitude, longitude, and depth
+                interp_latitude = interp1d(np.linspace(0, ray_path[0].path['lat'].max(), len(ray_path[0].path['lat'])), ray_path[0].path['lat'])
+                interp_longitude = interp1d(np.linspace(0, ray_path[0].path['lon'].max(), len(ray_path[0].path['lon'])), ray_path[0].path['lon'])
+                interp_depth = interp1d(np.linspace(0, ray_path[0].path['depth'].max(), len(ray_path[0].path['depth'])), ray_path[0].path['depth'])
+
+                # Resample the ray path to N points
+                lat_values_interp = interp_latitude(np.linspace(0, ray_path[0].path['lat'].max(), num_points))
+                lon_values_interp = interp_longitude(np.linspace(0, ray_path[0].path['lon'].max(), num_points))
+                depth_values_interp = interp_depth(np.linspace(0, ray_path[0].path['depth'].max(), num_points))
+
+                # Loop through the interpolated coordinates and update the grid cells with the average P-wave velocity
+                for lat, lon, depth in zip(lat_values_interp, lon_values_interp, depth_values_interp):
+                    lat_index = find_closest_index(lat_values, lat)
+                    lon_index = find_closest_index(lon_values, lon)
+                    depth_index = find_closest_index(depth_values, depth)
                     
+                    if velocity_model[lat_index, lon_index, depth_index] == initial_velocity:
+                        velocity_model[lat_index, lon_index, depth_index] = df['velocity_p, km/s'].iloc[i]
+                    else:
+                        velocity_model[lat_index, lon_index, depth_index] = (velocity_model[lat_index, lon_index, depth_index] +
+                                                                            df['velocity_p, km/s'].iloc[i]) / 2
+
     # Create the figure and axis
     fig = plt.figure(figsize=(8, 8))
     ax = fig.add_subplot(111, projection='3d')
@@ -515,11 +571,16 @@ def compute_velocity_model(azimuth, elevation):
     )
 
     # Create the color array
-    norm = plt.Normalize(vmin=velocity_model.min(), vmax=velocity_model.max())
-    colors = plt.cm.plasma(norm(velocity_model))
-
+    norm = plt.Normalize(vmin=2, vmax=8)
+    colors_vel = plt.cm.plasma(norm(velocity_model)) 
+    
     # Plot the voxels
-    ax.voxels(x, y, z, velocity_model > 0, facecolors=colors, alpha=0.5, edgecolor='k')
+    if interpolate:
+        interpolated_velocity_model = interpolate_vel_model(velocity_model, initial_velocity, lat_values, lon_values, depth_values, n_lat, n_lon, n_depth)
+        colors_interp = plt.cm.plasma(norm(interpolated_velocity_model))
+        ax.voxels(x, y, z, interpolated_velocity_model > 0, facecolors=colors_interp, alpha=0.5, edgecolor='k')
+    
+    ax.voxels(x, y, z, velocity_model > 0, facecolors=colors_vel, alpha=1, edgecolor='black')
 
     # Set the view angle
     ax.view_init(elev=elevation, azim=azimuth)
@@ -585,7 +646,7 @@ with gr.Blocks() as demo:
                 P_thres_inputs = gr.Slider(minimum=0.01,
                                     maximum=1,
                                     value=0.1,
-                                    label="P uncertainty threshold, s",
+                                    label="P uncertainty threshold (s)",
                                     step=0.01,
                                     info="Acceptable uncertainty for P picks expressed in std() seconds",
                                     interactive=True,
@@ -594,13 +655,13 @@ with gr.Blocks() as demo:
                 S_thres_inputs = gr.Slider(minimum=0.01,
                                     maximum=1,
                                     value=0.2,
-                                    label="S uncertainty threshold, s",
+                                    label="S uncertainty threshold (s)",
                                     step=0.01,
                                     info="Acceptable uncertainty for S picks expressed in std() seconds",
                                     interactive=True,
                                     )
             with gr.Column(scale=1):
-                upload = gr.File(label="Or upload your own waveform")
+                upload = gr.File(label="Upload your waveform")
                 with gr.Row():
                     sampling_rate_inputs = gr.Slider(minimum=10,
                             maximum=1000,
@@ -730,15 +791,22 @@ with gr.Blocks() as demo:
         
         with gr.Row():
             with gr.Column(scale=2):
-                inputs_vel_model = [
-                    gr.Slider(minimum=-180, maximum=180, value=0, step=5, label="Azimuth", interactive=True),
-                    gr.Slider(minimum=-90, maximum=90, value=30, step=5, label="Elevation", interactive=True)
-                ]
+                azimuth_input = gr.Slider(minimum=-180, maximum=180, value=0, step=5, label="Azimuth", interactive=True)
+                elevation_input = gr.Slider(minimum=-90, maximum=90, value=30, step=5, label="Elevation", interactive=True)
+
+                with gr.Row():
+                    interpolate_input = gr.Checkbox(label="Interpolate", info="Interpolate velocity model")
+                    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_lon_input = gr.Slider(minimum=5, maximum=100, value=50, step=5, label="N lon", info='Number of Lon grid points', interactive=True)
+                    n_depth_input = gr.Slider(minimum=5, maximum=100, value=50, step=5, label="N depth", info='Number of Depth grid points', interactive=True)
+                
                 button = gr.Button("Look at 3D Velocities")
             outputs_vel_model = gr.Image(label="3D Velocity Model")
 
-            button.click(compute_velocity_model, 
-                         inputs=inputs_vel_model, 
-                         outputs=outputs_vel_model)
+        button.click(compute_velocity_model, 
+                        inputs=[azimuth_input, elevation_input, 
+                                interpolate_input, n_lat_input, 
+                                n_lon_input, n_depth_input], 
+                        outputs=[outputs_vel_model])
 
 demo.launch()