-
Notifications
You must be signed in to change notification settings - Fork 46
/
example_location_interpolation.py
113 lines (102 loc) · 3.7 KB
/
example_location_interpolation.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
import torch
import matplotlib.pyplot as plt
import deepwave
from deepwave import scalar
device = torch.device('cuda' if torch.cuda.is_available()
else 'cpu')
ny = 100
nx = 300
dx = 4.0
v = torch.ones(ny, nx, device=device) * 1500
freq = 25
nt = 40
dt = 0.004
peak_time = 1.5 / freq
# source_amplitudes
source_amplitudes = (
deepwave.wavelets.ricker(freq, nt, dt, peak_time)
.repeat(1, 3, 1)
.to(device)
)
# Location interpolation
source_locations = torch.zeros(1, 3, 2, device=device)
source_locations[..., 0] = ny / 2 + 0.5
source_locations[..., 1] = torch.arange(3) * nx / 3 + nx / 6 + 0.5
source_monopole = torch.zeros(1, 3, dtype=torch.bool, device=device)
source_monopole[0, 0] = True
source_monopole[0, 1] = False
source_monopole[0, 2] = False
source_dipole_dim = torch.zeros(1, 3, dtype=torch.int,
device=device)
source_dipole_dim[0, 2] = 1
hicks_source = deepwave.location_interpolation.Hicks(
source_locations,
monopole=source_monopole,
dipole_dim=source_dipole_dim
)
hicks_source_locations = hicks_source.get_locations()
hicks_source_amplitudes = hicks_source.source(source_amplitudes)
receiver_locations = torch.zeros(1, 3, 2, device=device)
receiver_locations[..., 0] = source_locations[0, 0, 0] - ny / 8
receiver_locations[..., 1] = source_locations[0, 0, 1]
receiver_monopole = torch.zeros(1, 3, dtype=torch.bool,
device=device)
receiver_monopole[0, 0] = True
receiver_monopole[0, 1] = False
receiver_monopole[0, 2] = False
receiver_dipole_dim = torch.zeros(1, 3, dtype=torch.int,
device=device)
receiver_dipole_dim[0, 2] = 1
hicks_receiver = deepwave.location_interpolation.Hicks(
receiver_locations,
monopole=receiver_monopole,
dipole_dim=receiver_dipole_dim
)
hicks_receiver_locations = hicks_receiver.get_locations()
# Propagate
out = scalar(v, dx, dt, source_amplitudes=hicks_source_amplitudes,
source_locations=hicks_source_locations,
receiver_locations=hicks_receiver_locations,
pml_freq=freq)
# Plot wavefield
wavefield = out[0]
plt.figure(figsize=(10.5, 7))
plt.imshow(wavefield[0].cpu(), cmap='gray')
plt.savefig('example_location_interpolation_wavefield.jpg')
# Plot receiver amplitudes
receiver_hicks_amplitudes = out[-1]
receiver_amplitudes = hicks_receiver.receiver(
receiver_hicks_amplitudes
)
plt.figure(figsize=(10.5, 7))
plt.plot(receiver_amplitudes[0, 0].cpu(), label='Monopole')
plt.plot(receiver_amplitudes[0, 1].cpu(), label='Dipole 0')
plt.plot(receiver_amplitudes[0, 2].cpu(), label='Dipole 1')
plt.legend()
plt.xlabel('Time sample')
plt.ylabel('Receiver amplitude')
plt.savefig('example_location_interpolation_receiver_data.jpg')
# Plot comparison with grid-centred points
out2 = scalar(v, dx, dt,
source_amplitudes=source_amplitudes,
source_locations=source_locations.long(),
receiver_locations=receiver_locations.long(),
pml_freq=freq)
v = torch.ones(2 * ny, 2 * nx, device=device) * 1500
out3 = scalar(v, dx / 2, dt,
source_amplitudes=source_amplitudes*4,
source_locations=(source_locations*2).long(),
receiver_locations=(receiver_locations*2).long(),
pml_freq=freq)
receiver_amplitudes2 = out2[-1]
receiver_amplitudes3 = out3[-1]
plt.figure(figsize=(10.5, 7))
plt.plot(receiver_amplitudes[0, 0].cpu(), label='Hicks',
linewidth=6)
plt.plot(receiver_amplitudes2[0, 0].cpu(), 'g',
label='Without Hicks')
plt.plot(receiver_amplitudes3[0, 0].cpu(), 'r', label='Denser grid')
plt.legend()
plt.xlabel('Time sample')
plt.ylabel('Receiver amplitude')
plt.savefig('example_location_interpolation_comparison.jpg')