This repository has been archived by the owner on Jun 13, 2022. It is now read-only.
-
Notifications
You must be signed in to change notification settings - Fork 2
/
example_born_fully_connected.jl
109 lines (79 loc) · 2.44 KB
/
example_born_fully_connected.jl
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
# Example for linearized modeling with Flux networks
# Author: Philipp Witte, pwitte3@gatech.edu
# Date: April 2020
#
using JUDI, SegyIO, JOLI, Flux, JUDI4Flux
# Set up model structure
n = (120, 100) # (x,y,z) or (x,z)
d = (10., 10.)
o = (0., 0.)
# Velocity [km/s]
v = ones(Float32,n) .+ 0.4f0
v0 = ones(Float32,n) .+ 0.4f0
v[:, Int(round(end/2)):end] .= 4f0
# Slowness squared [s^2/km^2]
m = (1f0 ./ v).^2
m0 = (1f0 ./ v0).^2
dm = vec(m - m0)
# Setup info and model structure
nsrc = 1 # number of sources
model0 = Model(n, d, o, m0)
# Set up receiver geometry
nxrec = 120
xrec = range(50f0, stop=1150f0, length=nxrec)
yrec = 0f0
zrec = range(50f0, stop=50f0, length=nxrec)
# receiver sampling and recording time
time = 1000f0 # receiver recording time [ms]
dt = 1f0 # receiver sampling interval [ms]
# Set up receiver structure
recGeometry = Geometry(xrec, yrec, zrec; dt=dt, t=time, nsrc=nsrc)
## Set up source geometry (cell array with source locations for each shot)
xsrc = convertToCell([600f0])
ysrc = convertToCell([0f0])
zsrc = convertToCell([20f0])
# Set up source structure
srcGeometry = Geometry(xsrc, ysrc, zsrc; dt=dt, t=time)
# setup wavelet
f0 = 0.01f0 # MHz
wavelet = ricker_wavelet(time, dt, f0)
q = judiVector(srcGeometry, wavelet)
# Set up info structure for linear operators
ntComp = get_computational_nt(srcGeometry, recGeometry, model0)
info = Info(prod(n), nsrc, ntComp)
# return linearized data as Julia array
opt = Options(return_array=true)
# Non-linear forward modeling operator
F0 = judiModeling(info, model0, srcGeometry, recGeometry; options=opt)
num_samples = recGeometry.nt[1] * nxrec
##################################################################################
# Fully connected neural network with linearized modeling operator
n_in = 100
n_out = 10
W1 = randn(Float32, info.n, n_in)
b1 = randn(Float32, info.n)
W2 = judiJacobian(F0, q)
b2 = randn(Float32, num_samples)
W3 = randn(Float32, n_out, num_samples)
b3 = randn(Float32, n_out)
function network(x)
x = W1*x .+ b1
x = W2*x .+ b2
x = W3*x .+ b3
return x
end
# Gradient
x = zeros(Float32, n_in)
y = randn(Float32, n_out)
# Evaluate MSE loss
loss(x, y) = Flux.mse(network(x), y)
# Compute gradient w.r.t. x and y
Δx, Δy = gradient(loss, x, y)
# Compute gradient for x, y and weights (except for W2)
p = params(x, y, W1, b1, b2, W3, b3)
gs = gradient(() -> loss(x, y), p)
# Access gradients
Δx = gs[x]
ΔW1 = gs[W1]
Δb1 = gs[b1]
# and so on...