Examples: extended 3D example to full run and added a PyVista plotter

examples/infrasound/3d_infrasound_example.py

@@ -9,19 +9,20 @@
 
 from model import InfrasoundModel
 from propagator import InfrasoundPropagator
+from plotting import plot_st_helens
 
-src_coords = np.array([4800., 4800., 2550.])[np.newaxis, :]
-rec_coords = np.array([[2400., 2400., 1275.],
-                       [7200., 2400., 1275.],
-                       [2400., 7200., 1275.],
-                       [7200., 7200., 1275.],
-                       [2400., 2400., 3825.],
-                       [7200., 2400., 3825.],
-                       [2400., 7200., 3825.],
-                       [7200., 7200., 3825.]])
+src_coords = np.array([4800., 4800., 2250.])[np.newaxis, :]
+rec_coords = np.array([[4800., 1400., 1400.],
+                       [8200., 4800., 1500.],
+                       [1400., 4800., 1500.],
+                       [4800., 8200., 1650.],
+                       [2400., 2400., 1400.],
+                       [7200., 2400., 1500.],
+                       [2400., 7200., 1550.],
+                       [7200., 7200., 1500.]])
 
 # Original time is 13s
-t0, tn, dt = 0., 6., 0.021  # Courant number ~0.25
+t0, tn, dt = 0., 14., 0.021  # Courant number ~0.25
 src_f = 1.
 
 sdf_data = -np.load('surface_files/mt_st_helens_3d.npy')
@@ -48,11 +49,18 @@
 plt.colorbar()
 plt.show()
 
-"""
-print(np.linalg.norm(model.rec.data[:, 0]))
-plt.plot(model.rec.data[:, 0])
+plt.imshow(model.p.data[-1, :, xmid].T, origin='lower', extent=plt_ext)
+plt.colorbar()
 plt.show()
 
+plot_st_helens(model.p.data[-1], src_coords, rec_coords,
+               np.array([-4800., -4800., 0.]), (30, 30, 30))
+
+for i in range(rec_coords.shape[0]):
+    print(np.linalg.norm(model.rec.data[:, i]))
+    plt.plot(model.rec.data[:, i])
+    plt.show()
+
 # Next step is backpropagation
 # Reset the fields
 model.p.data[:] = 0
@@ -74,7 +82,17 @@
 plt.colorbar()
 plt.show()
 
+plt.imshow(model.p.data[-1, :, xmid].T, origin='lower', extent=plt_ext)
+plt.colorbar()
+plt.show()
+
 plt.imshow(model.zsc.data[-1, xmid].T, origin='lower', extent=plt_ext)
 plt.colorbar()
 plt.show()
-"""
+
+plt.imshow(model.zsc.data[-1, :, xmid].T, origin='lower', extent=plt_ext)
+plt.colorbar()
+plt.show()
+
+plot_st_helens(model.zsc.data[-1], src_coords, rec_coords,
+               np.array([-4800., -4800., 0.]), (30, 30, 30))

examples/infrasound/plotting.py

@@ -0,0 +1,67 @@
+"""Useful pyvista plotters for visualising the 3D data"""
+
+import pyvista as pv
+import numpy as np
+
+
+def plot_st_helens(data, src_loc, rec_loc, origin, spacing):
+    """Plot a 3D render of slices through a dataset with topography"""
+    # Create the spatial reference
+    mesh = pv.UniformGrid()
+
+    # Set the grid dimensions: shape + 1 because we want to inject our values on
+    #   the CELL data
+    mesh.dimensions = np.array(data.shape) + 1
+
+    # Edit the spatial reference
+    mesh.origin = origin  # The bottom left corner of the data set
+    mesh.spacing = spacing  # These are the cell sizes along each axis
+
+    # Add the data values to the cell data
+    mesh.cell_data["values"] = data.flatten(order="F")  # Flatten the array!
+
+    # Now plot the grid!
+    # slices = mesh.slice_orthogonal(z=1800)
+    slicex = mesh.slice(normal=[1, 0, 0])
+    slicey = mesh.slice(normal=[0, 1, 0])
+
+    surface = pv.read("surface_files/mt_st_helens.ply")
+
+    plotter = pv.Plotter()
+    plotter.add_mesh(slicex, opacity=0.95)
+    plotter.add_mesh(slicey, opacity=0.95)
+    plotter.add_mesh(surface)
+
+    for src in src_loc:
+        center = src + origin
+        sphere = pv.Sphere(radius=90, center=center)
+        plotter.add_mesh(sphere, color='red')
+
+    for rec in rec_loc:
+        center = rec + origin
+        sphere = pv.Sphere(radius=90, center=center)
+        plotter.add_mesh(sphere, color='blue')
+
+    plotter.show()
+
+
+def main():
+    data = np.random.rand(321, 321, 171)
+
+    origin = np.array([-4800., -4800., 0.])
+    src_loc = np.array([4800., 4800., 2250.])[np.newaxis, :]
+    rec_loc = np.array([[4800., 1400., 1400.],
+                       [8200., 4800., 1500.],
+                       [1400., 4800., 1500.],
+                       [4800., 8200., 1650.],
+                       [2400., 2400., 1400.],
+                       [7200., 2400., 1500.],
+                       [2400., 7200., 1550.],
+                       [7200., 7200., 1500.]])
+
+    spacing = (30, 30, 30)
+    plot_st_helens(data, src_loc, rec_loc, origin, spacing)
+
+
+if __name__ == "__main__":
+    main()