def gradient(self, m0, maxmem=None):
cp = DevitoCheckpoint([self.forward_field])
n_checkpoints = None
if maxmem is not None:
n_checkpoints = int(
floor(maxmem * 10**6 /
(cp.size * self.forward_field.data.itemsize)))
wrap_fw = CheckpointOperator(self.forward_operator,
u=self.forward_field,
rec=self.rec,
m=m0,
src=self.src,
dt=self.dt)
wrap_rev = CheckpointOperator(self.gradient_operator,
u=self.forward_field,
v=self.adjoint_field,
m=m0,
rec=self.rec_g,
grad=self.grad,
dt=self.dt)
wrp = Revolver(cp, wrap_fw, wrap_rev, n_checkpoints,
self.nt - self.time_order)
wrp.apply_forward()
self.rec_g.data[:] = self.rec.data[:] - self.rec_t.data[:]
wrp.apply_reverse()
# The result is in grad
return self.grad.data, self.rec.data
def test_forward_with_breaks(shape, kernel, space_order):
""" Test running forward in one go and "with breaks"
and ensure they produce the same result
"""
spacing = tuple([15.0 for _ in shape])
tn = 500.
time_order = 2
nrec = shape[0]
solver = acoustic_setup(shape=shape, spacing=spacing, tn=tn,
space_order=space_order, kernel=kernel)
grid = solver.model.grid
rec = Receiver(name='rec', grid=grid, time_range=solver.geometry.time_axis,
npoint=nrec)
rec.coordinates.data[:, :] = solver.geometry.rec_positions
dt = solver.model.critical_dt
u = TimeFunction(name='u', grid=grid, time_order=2, space_order=space_order)
cp = DevitoCheckpoint([u])
wrap_fw = CheckpointOperator(solver.op_fwd(save=False), rec=rec,
src=solver.geometry.src, u=u, dt=dt)
wrap_rev = CheckpointOperator(solver.op_grad(save=False), u=u, dt=dt, rec=rec)
wrp = Revolver(cp, wrap_fw, wrap_rev, None, rec._time_range.num-time_order)
rec1, u1, summary = solver.forward()
wrp.apply_forward()
assert(np.allclose(u1.data, u.data))
assert(np.allclose(rec1.data, rec.data))
def test_forward_with_breaks(shape, time_order, space_order):
""" Test running forward in one go and "with breaks"
and ensure they produce the same result
"""
spacing = tuple([15.0 for _ in shape])
tn = 500.
example = CheckpointingExample(shape, spacing, tn, time_order, space_order)
m0, dm = example.initial_estimate()
cp = DevitoCheckpoint([example.forward_field])
wrap_fw = CheckpointOperator(example.forward_operator,
u=example.forward_field,
rec=example.rec,
m=m0,
src=example.src,
dt=example.dt)
wrap_rev = CheckpointOperator(example.gradient_operator,
u=example.forward_field,
v=example.adjoint_field,
m=m0,
rec=example.rec_g,
grad=example.grad,
dt=example.dt)
wrp = Revolver(cp, wrap_fw, wrap_rev, None, example.nt - time_order)
example.forward_operator.apply(u=example.forward_field,
rec=example.rec,
m=m0,
src=example.src,
dt=example.dt)
u_temp = np.copy(example.forward_field.data)
rec_temp = np.copy(example.rec.data)
example.forward_field.data[:] = 0
wrp.apply_forward()
assert (np.allclose(u_temp, example.forward_field.data))
assert (np.allclose(rec_temp, example.rec.data))
def jacobian_adjoint(self, rec, u, v=None, grad=None, vp=None,
checkpointing=False, **kwargs):
"""
Gradient modelling function for computing the adjoint of the
Linearized Born modelling function, ie. the action of the
Jacobian adjoint on an input data.
Parameters
----------
rec : SparseTimeFunction
Receiver data.
u : TimeFunction
Full wavefield `u` (created with save=True).
v : TimeFunction, optional
Stores the computed wavefield.
grad : Function, optional
Stores the gradient field.
vp : Function or float, optional
The time-constant velocity.
Returns
-------
Gradient field and performance summary.
"""
dt = kwargs.pop('dt', self.dt)
# Gradient symbol
grad = grad or Function(name='grad', grid=self.model.grid)
# Create the forward wavefield
v = v or TimeFunction(name='v', grid=self.model.grid,
time_order=2, space_order=self.space_order)
# Pick vp from model unless explicitly provided
vp = vp or self.model.vp
if checkpointing:
u = TimeFunction(name='u', grid=self.model.grid,
time_order=2, space_order=self.space_order)
cp = DevitoCheckpoint([u])
n_checkpoints = None
wrap_fw = CheckpointOperator(self.op_fwd(save=False), src=self.geometry.src,
u=u, vp=vp, dt=dt)
wrap_rev = CheckpointOperator(self.op_grad(save=False), u=u, v=v,
vp=vp, rec=rec, dt=dt, grad=grad)
# Run forward
wrp = Revolver(cp, wrap_fw, wrap_rev, n_checkpoints, rec.data.shape[0]-2)
wrp.apply_forward()
summary = wrp.apply_reverse()
else:
summary = self.op_grad().apply(rec=rec, grad=grad, v=v, u=u, vp=vp,
dt=dt, **kwargs)
return grad, summary
def verify(space_order=4, kernel='OT4', nbpml=40, filename='',
compression_params={}, **kwargs):
solver = acoustic_setup(shape=(10, 10), spacing=(10, 10), nbpml=10, tn=50,
space_order=space_order, kernel=kernel, **kwargs)
# solver = overthrust_setup(filename=filename, tn=50, nbpml=nbpml,
# space_order=space_order, kernel=kernel,
# **kwargs)
u = TimeFunction(name='u', grid=solver.model.grid, time_order=2,
space_order=solver.space_order)
rec = Receiver(name='rec', grid=solver.model.grid,
time_range=solver.geometry.time_axis,
coordinates=solver.geometry.rec_positions)
cp = DevitoCheckpoint([u])
n_checkpoints = None
dt = solver.dt
v = TimeFunction(name='v', grid=solver.model.grid, time_order=2,
space_order=solver.space_order)
grad = Function(name='grad', grid=solver.model.grid)
wrap_fw = CheckpointOperator(solver.op_fwd(save=False),
src=solver.geometry.src, u=u, rec=rec, dt=dt)
wrap_rev = CheckpointOperator(solver.op_grad(save=False), u=u, v=v,
rec=rec, dt=dt, grad=grad)
nt = rec.data.shape[0] - 2
print("Verifying for %d timesteps" % nt)
wrp = Revolver(cp, wrap_fw, wrap_rev, n_checkpoints, nt,
compression_params=compression_params)
wrp.apply_forward()
wrp.apply_reverse()
print(wrp.profiler.timings)
with Timer([]) as tf:
rec2, u2, _ = solver.forward(save=True)
with Timer([]) as tr:
grad2, _ = solver.gradient(rec=rec2, u=u2)
error = grad.data - grad2.data
# to_hdf5(error, 'zfp_grad_errors.h5')
print("Error norm", np.linalg.norm(error))
# assert(np.allclose(grad.data, grad2.data))
print("Checkpointing implementation is numerically verified")
print("Verification took %d ms for forward and %d ms for reverse" % (tf.elapsed, tr.elapsed))
def gradient(self, rec, u, v=None, grad=None, m=None, checkpointing=False, **kwargs):
"""
Gradient modelling function for computing the adjoint of the
Linearized Born modelling function, ie. the action of the
Jacobian adjoint on an input data.
:param recin: Receiver data as a numpy array
:param u: Symbol for full wavefield `u` (created with save=True)
:param v: (Optional) Symbol to store the computed wavefield
:param grad: (Optional) Symbol to store the gradient field
:returns: Gradient field and performance summary
"""
dt = kwargs.pop('dt', self.dt)
# Gradient symbol
grad = grad or Function(name='grad', grid=self.model.grid)
# Create the forward wavefield
v = v or TimeFunction(name='v', grid=self.model.grid,
time_order=2, space_order=self.space_order)
# Pick m from model unless explicitly provided
m = m or self.model.m
if checkpointing:
u = TimeFunction(name='u', grid=self.model.grid,
time_order=2, space_order=self.space_order)
cp = DevitoCheckpoint([u])
n_checkpoints = None
wrap_fw = CheckpointOperator(self.op_fwd(save=False), src=self.geometry.src,
u=u, m=m, dt=dt)
wrap_rev = CheckpointOperator(self.op_grad(save=False), u=u, v=v,
m=m, rec=rec, dt=dt, grad=grad)
# Run forward
wrp = Revolver(cp, wrap_fw, wrap_rev, n_checkpoints, rec.data.shape[0]-2)
wrp.apply_forward()
summary = wrp.apply_reverse()
else:
summary = self.op_grad().apply(rec=rec, grad=grad, v=v, u=u, m=m,
dt=dt, **kwargs)
return grad, summary
def calculate_lossy_gradient(i, solver, vp, grad, path_prefix, to=2, so=4,
n_checkpoints=1000, compression_params=None):
true_d, source_location = load_shot(i, path_prefix)
dt = solver.dt
# Update source location
solver.geometry.src_positions[0, :] = source_location[:]
if compression_params is None:
print("Using default compression params")
compression_params = {}
# Compute smooth data and full forward wavefield u0
u = TimeFunction(name='u', grid=solver.model.grid, time_order=to,
space_order=so, save=None)
v = TimeFunction(name='v', grid=solver.model.grid, time_order=to,
space_order=so)
residual = Receiver(name='rec', grid=solver.model.grid,
time_range=solver.geometry.time_axis,
coordinates=solver.geometry.rec_positions)
smooth_d = Receiver(name='rec', grid=solver.model.grid,
time_range=solver.geometry.time_axis,
coordinates=solver.geometry.rec_positions)
fwd_op = solver.op_fwd(save=False)
rev_op = solver.op_grad(save=False)
wrap_fw = CheckpointOperator(fwd_op, src=solver.geometry.src, u=u,
rec=smooth_d, vp=vp, dt=dt)
wrap_rev = CheckpointOperator(rev_op, vp=vp, u=u, v=v, rec=residual,
grad=grad, dt=dt)
cp = DevitoCheckpoint([u])
nt = smooth_d.data.shape[0] - 2
wrp = Revolver(cp, wrap_fw, wrap_rev, n_checkpoints, nt,
compression_params=compression_params)
wrp.apply_forward()
# Compute gradient from data residual
residual.data[:] = smooth_d.data[:] - true_d[:]
wrp.apply_reverse()
def checkpointed_run(space_order=4, ncp=None, kernel='OT4', nbpml=40,
filename='', compression_params={}, tn=1000, **kwargs):
solver = overthrust_setup(filename=filename, tn=tn, nbpml=nbpml,
space_order=space_order, kernel=kernel, **kwargs)
u = TimeFunction(name='u', grid=solver.model.grid, time_order=2,
space_order=solver.space_order)
rec = Receiver(name='rec', grid=solver.model.grid,
time_range=solver.geometry.time_axis,
coordinates=solver.geometry.rec_positions)
cp = DevitoCheckpoint([u])
n_checkpoints = ncp
dt = solver.dt
v = TimeFunction(name='v', grid=solver.model.grid, time_order=2,
space_order=solver.space_order)
grad = Function(name='grad', grid=solver.model.grid)
wrap_fw = CheckpointOperator(solver.op_fwd(save=False),
src=solver.geometry.src, u=u,
rec=rec, dt=dt)
wrap_rev = CheckpointOperator(solver.op_grad(save=False), u=u, v=v,
rec=rec, dt=dt, grad=grad)
fw_timings = []
rev_timings = []
nt = rec.data.shape[0] - 2
print("Running %d timesteps" % (nt))
print(compression_params)
wrp = Revolver(cp, wrap_fw, wrap_rev, n_checkpoints, nt,
compression_params=compression_params)
with Timer(fw_timings):
wrp.apply_forward()
with Timer(rev_timings):
wrp.apply_reverse()
return grad, wrp, fw_timings, rev_timings
def jacobian_adjoint(self, rec, p, pa=None, grad=None, r=None, model=None,
checkpointing=False, **kwargs):
"""
Gradient modelling function for computing the adjoint of the
Linearized Born modelling function, ie. the action of the
Jacobian adjoint on an input data.
Parameters
----------
rec : SparseTimeFunction
Receiver data.
p : TimeFunction
Full wavefield `p` (created with save=True).
pa : TimeFunction, optional
Stores the computed wavefield.
grad : Function, optional
Stores the gradient field.
r : TimeFunction, optional
The computed attenuation memory variable.
model : Model, optional
Object containing the physical parameters.
vp : Function or float, optional
The time-constant velocity.
qp : Function, optional
The P-wave quality factor.
b : Function, optional
The time-constant inverse density.
Returns
-------
Gradient field and performance summary.
"""
dt = kwargs.pop('dt', self.dt)
# Gradient symbol
grad = grad or Function(name='grad', grid=self.model.grid)
# Create the forward wavefield
pa = pa or TimeFunction(name='pa', grid=self.model.grid,
time_order=self.time_order, space_order=self.space_order,
staggered=NODE)
model = model or self.model
# Pick vp and physical parameters from model unless explicitly provided
kwargs.update(model.physical_params(**kwargs))
if checkpointing:
p = TimeFunction(name='p', grid=self.model.grid,
time_order=self.time_order, space_order=self.space_order,
staggered=NODE)
r = TimeFunction(name="r", grid=self.model.grid, time_order=self.time_order,
space_order=self.space_order, staggered=NODE)
cp = DevitoCheckpoint([p, r])
n_checkpoints = None
wrap_fw = CheckpointOperator(self.op_fwd(save=False),
src=self.geometry.src, p=p, r=r, dt=dt, **kwargs)
wrap_rev = CheckpointOperator(self.op_grad(save=False), p=p, pa=pa, rec=rec,
dt=dt, grad=grad, **kwargs)
# Run forward
wrp = Revolver(cp, wrap_fw, wrap_rev, n_checkpoints, rec.data.shape[0]-2)
wrp.apply_forward()
summary = wrp.apply_reverse()
else:
# Memory variable:
r = r or TimeFunction(name="r", grid=self.model.grid,
time_order=self.time_order,
space_order=self.space_order, staggered=NODE)
summary = self.op_grad().apply(rec=rec, grad=grad, pa=pa, p=p, r=r, dt=dt,
**kwargs)
return grad, summary
def test_index_alignment(const):
""" A much simpler test meant to ensure that the forward and reverse indices are
correctly aligned (i.e. u * v , where u is the forward field and v the reverse field
corresponds to the correct timesteps in u and v). The forward operator does u = u + 1
which means that the field a will be equal to nt (0 -> 1 -> 2 -> 3), the number of
timesteps this operator is run for. The field at the last time step of the forward is
used to initialise the field v for the reverse pass. The reverse operator does
v = v - 1, which means that if the reverse operator is run for the same number of
timesteps as the forward operator, v should be 0 at the last time step
(3 -> 2 -> 1 -> 0). There is also a grad = grad + u * v accumulator in the reverse
operator. If the alignment is correct, u and v should have the same value at every
time step:
0 -> 1 -> 2 -> 3 u
0 <- 1 <- 2 <- 3 v
and hence grad = 0*0 + 1*1 + 2*2 + 3*3 = sum(n^2) where n -> [0, nt]
If the test fails, the resulting number can tell you how the fields are misaligned
"""
n = 4
grid = Grid(shape=(2, 2))
order_of_eqn = 1
modulo_factor = order_of_eqn + 1
nt = n - order_of_eqn
u = TimeFunction(name='u', grid=grid, save=n)
# Increment one in the forward pass 0 -> 1 -> 2 -> 3
fwd_op = Operator(Eq(u.forward, u + 1. * const))
# Invocation 1
fwd_op(time=nt - 1, constant=1)
last_time_step_v = nt % modulo_factor
# Last time step should be equal to the number of timesteps we ran
assert (np.allclose(u.data[nt, :, :], nt))
v = TimeFunction(name='v', grid=grid, save=None)
v.data[last_time_step_v, :, :] = u.data[nt, :, :]
# Decrement one in the reverse pass 3 -> 2 -> 1 -> 0
adj_eqn = Eq(v, v.forward - 1. * const)
adj_op = Operator(adj_eqn)
# Invocation 2
adj_op(time=nt - 1, constant=1)
# Last time step should be back to 0
assert (np.allclose(v.data[0, :, :], 0))
# Reset v to run the backward again
v.data[last_time_step_v, :, :] = u.data[nt, :, :]
prod = Function(name="prod", grid=grid)
# Multiply u and v and add them
# = 3*3 + 2*2 + 1*1 + 0*0
prod_eqn = Eq(prod, prod + u * v)
comb_op = Operator([adj_eqn, prod_eqn])
# Invocation 3
comb_op(time=nt - 1, constant=1)
final_value = sum([n**2 for n in range(nt)])
# Final value should be sum of squares of first nt natural numbers
assert (np.allclose(prod.data, final_value))
# Now reset to repeat all the above tests with checkpointing
prod.data[:] = 0
v.data[last_time_step_v, :, :] = u.data[nt, :, :]
# Checkpointed version doesn't require to save u
u_nosave = TimeFunction(name='u_n', grid=grid)
# change equations to use new symbols
fwd_eqn_2 = Eq(u_nosave.forward, u_nosave + 1. * const)
fwd_op_2 = Operator(fwd_eqn_2)
cp = DevitoCheckpoint([u_nosave])
wrap_fw = CheckpointOperator(fwd_op_2, constant=1)
prod_eqn_2 = Eq(prod, prod + u_nosave * v)
comb_op_2 = Operator([adj_eqn, prod_eqn_2])
wrap_rev = CheckpointOperator(comb_op_2, constant=1)
wrp = Revolver(cp, wrap_fw, wrap_rev, None, nt)
# Invocation 4
wrp.apply_forward()
assert (np.allclose(u_nosave.data[last_time_step_v, :, :], nt))
# Invocation 5
wrp.apply_reverse()
assert (np.allclose(v.data[0, :, :], 0))
assert (np.allclose(prod.data, final_value))
def jacobian_adjoint(self,
rec,
u0,
v0,
du=None,
dv=None,
dm=None,
model=None,
checkpointing=False,
kernel='centered',
**kwargs):
"""
Gradient modelling function for computing the adjoint of the
Linearized Born modelling function, ie. the action of the
Jacobian adjoint on an input data.
Parameters
----------
rec : SparseTimeFunction
Receiver data.
u0 : TimeFunction
The computed background wavefield.
v0 : TimeFunction, optional
The computed background wavefield.
du : Function or float
The computed perturbed wavefield.
dv : Function or float
The computed perturbed wavefield.
dm : Function, optional
Stores the gradient field.
model : Model, optional
Object containing the physical parameters.
vp : Function or float, optional
The time-constant velocity.
epsilon : Function or float, optional
The time-constant first Thomsen parameter.
delta : Function or float, optional
The time-constant second Thomsen parameter.
theta : Function or float, optional
The time-constant Dip angle (radians).
phi : Function or float, optional
The time-constant Azimuth angle (radians).
Returns
-------
Gradient field and performance summary.
"""
if kernel != 'centered':
raise ValueError(
'Only centered kernel is supported for the jacobian_adj')
dt = kwargs.pop('dt', self.dt)
# Gradient symbol
dm = dm or Function(name='dm', grid=self.model.grid)
# Create the perturbation wavefields if not provided
du = du or TimeFunction(name='du',
grid=self.model.grid,
time_order=2,
space_order=self.space_order)
dv = dv or TimeFunction(name='dv',
grid=self.model.grid,
time_order=2,
space_order=self.space_order)
model = model or self.model
# Pick vp and Thomsen parameters from model unless explicitly provided
kwargs.update(model.physical_params(**kwargs))
if self.model.dim < 3:
kwargs.pop('phi', None)
if checkpointing:
u0 = TimeFunction(name='u0',
grid=self.model.grid,
time_order=2,
space_order=self.space_order)
v0 = TimeFunction(name='v0',
grid=self.model.grid,
time_order=2,
space_order=self.space_order)
cp = DevitoCheckpoint([u0, v0])
n_checkpoints = None
wrap_fw = CheckpointOperator(self.op_fwd(save=False),
src=self.geometry.src,
u=u0,
v=v0,
dt=dt,
**kwargs)
wrap_rev = CheckpointOperator(self.op_jacadj(save=False),
u0=u0,
v0=v0,
du=du,
dv=dv,
rec=rec,
dm=dm,
dt=dt,
**kwargs)
# Run forward
wrp = Revolver(cp, wrap_fw, wrap_rev, n_checkpoints,
rec.data.shape[0] - 2)
wrp.apply_forward()
summary = wrp.apply_reverse()
else:
summary = self.op_jacadj().apply(rec=rec,
dm=dm,
u0=u0,
v0=v0,
du=du,
dv=dv,
dt=dt,
**kwargs)
return dm, summary
def run(space_order=4,
kernel='OT2',
nbpml=40,
autotune=False,
filename='',
chunk=1000000,
algo=None,
shuffle="SHUFFLE",
**kwargs):
solver = overthrust_setup(filename=filename,
nbpml=nbpml,
space_order=space_order,
kernel=kernel,
**kwargs)
m = solver.model.m
dt = solver.dt
u = TimeFunction(name='u',
grid=solver.model.grid,
time_order=2,
space_order=solver.space_order)
v = TimeFunction(name='v',
grid=solver.model.grid,
time_order=2,
space_order=solver.space_order)
rec = Receiver(name='rec',
grid=solver.model.grid,
time_range=solver.receiver.time_range,
coordinates=solver.receiver.coordinates.data)
grad = Function(name='grad', grid=solver.model.grid)
cp = CompressionCheckpoint([u])
n_checkpoints = 60
wrap_fw = CheckpointOperator(solver.op_fwd(save=False),
src=solver.source,
u=u,
m=m,
dt=solver.dt,
rec=rec)
wrap_rev = CheckpointOperator(solver.op_grad(save=False),
u=u,
v=v,
m=m,
rec=rec,
dt=dt,
grad=grad)
# Run forward
wrp = Revolver(cp,
wrap_fw,
wrap_rev,
n_checkpoints,
rec.data.shape[0] - 2,
compression='blosc',
compression_params={
CHUNK_SIZE: chunk,
CNAME: algo,
SHUFFLE: shuffle
})
info("Applying Forward")
solver.forward(time=100)
#raw_fw(dt=dt)
info("Again")
wrp.apply_forward()
print(np.linalg.norm(u.data))
print(np.linalg.norm(rec.data))
info("Applying Gradient")
summary = wrp.apply_reverse()
print("Gradient is: %d" % np.linalg.norm(grad.data))
def forward_modeling_single_shot(record, table, par_files):
"Serial modeling function"
worker = get_worker() # The worker on which this task is running
strng = '{} : {} =>'.format(worker.address, str(record).zfill(5))
filename = 'logfile_{}.txt'.format(str(record).zfill(5))
g = open(filename, 'w')
g.write("This will show up in the worker logs")
# Read velocity model
f = segyio.open(par_files[-1], iline=segyio.tracefield.TraceField.FieldRecord,
xline=segyio.tracefield.TraceField.CDP)
xl, il, t = f.xlines, f.ilines, f.samples
dz = t[1] - t[0]
dx = f.header[1][segyio.TraceField.SourceX]-f.header[0][segyio.TraceField.SourceX]
if len(il) != 1:
dims = (len(xl), len(il), len(f.samples))
else:
dims = (len(xl), len(f.samples))
vp = f.trace.raw[:].reshape(dims)
vp *= 1./1000 # convert to km/sec
epsilon = np.empty(dims)
delta = np.empty(dims)
theta = np.empty(dims)
params = [epsilon, delta, theta]
# Read Thomsem parameters
for segyfile, par in zip(par_files, params):
f = segyio.open(segyfile, iline=segyio.tracefield.TraceField.FieldRecord,
xline=segyio.tracefield.TraceField.CDP)
par[:] = f.trace.raw[:].reshape(dims)
theta *= (np.pi/180.) # use radians
g.write('{} Parameter model dims: {}\n'.format(strng, vp.shape))
origin = (0., 0.)
shape = vp.shape
spacing = (dz, dz)
# Get a single shot as a numpy array
filename = table[record]['filename']
position = table[record]['Trace_Position']
traces_in_shot = table[record]['Num_Traces']
src_coord = np.array(table[record]['Source']).reshape((1, len(dims)))
rec_coord = np.array(table[record]['Receivers'])
start = time.time()
f = segyio.open(filename, ignore_geometry=True)
num_samples = len(f.samples)
samp_int = f.bin[segyio.BinField.Interval]/1000
retrieved_shot = np.zeros((traces_in_shot, num_samples))
shot_traces = f.trace[position:position+traces_in_shot]
for i, trace in enumerate(shot_traces):
retrieved_shot[i] = trace
g.write('{} Shot loaded in: {} seconds\n'.format(strng, time.time()-start))
# Only keep receivers within the model'
xmin = origin[0]
idx_xrec = np.where(rec_coord[:, 0] < xmin)[0]
is_empty = idx_xrec.size == 0
if not is_empty:
g.write('{} in {}\n'.format(strng, rec_coord.shape))
idx_tr = np.where(rec_coord[:, 0] >= xmin)[0]
rec_coord = np.delete(rec_coord, idx_xrec, axis=0)
# For 3D shot records, scan also y-receivers
if len(origin) == 3:
ymin = origin[1]
idx_yrec = np.where(rec_coord[:, 1] < ymin)[0]
is_empty = idx_yrec.size == 0
if not is_empty:
rec_coord = np.delete(rec_coord, idx_yrec, axis=0)
if rec_coord.size == 0:
g.write('all receivers outside of model\n')
return np.zeros(vp.shape)
space_order = 8
g.write('{} before: {} {} {}\n'.format(strng, params[0].shape, rec_coord.shape, src_coord.shape))
model = limit_model_to_receiver_area(rec_coord, src_coord, origin, spacing,
shape, vp, params, space_order=space_order, nbl=80)
g.write('{} shape_vp: {}\n'.format(strng, model.vp.shape))
model.smooth(('epsilon', 'delta', 'theta'))
# Geometry for current shot
geometry = AcquisitionGeometry(model, rec_coord, src_coord, 0, (num_samples-1)*samp_int, f0=0.018, src_type='Ricker')
g.write("{} Number of samples modelled data & dt: {} & {}\n".format(strng, geometry.nt, model.critical_dt))
g.write("{} Samples & dt: {} & {}\n".format(strng, num_samples, samp_int))
# Set up solver.
solver_tti = AnisotropicWaveSolver(model, geometry, space_order=space_order)
# Create image symbol and instantiate the previously defined imaging operator
image = Function(name='image', grid=model.grid)
itemsize = np.dtype(np.float32).itemsize
full_fld_mem = model.vp.size*itemsize*geometry.nt*2.
checkpointing = True
if checkpointing:
op_imaging = ImagingOperator(geometry, image, space_order, save=False)
n_checkpoints = 150
ckp_fld_mem = model.vp.size*itemsize*n_checkpoints*2.
g.write('Mem full fld: {} == {} use ckp instead\n'.format(full_fld_mem, humanbytes(full_fld_mem)))
g.write('Number of checkpoints/timesteps: {}/{}\n'.format(n_checkpoints, geometry.nt))
g.write('Memory saving: {}\n'.format(humanbytes(full_fld_mem-ckp_fld_mem)))
u = TimeFunction(name='u', grid=model.grid, staggered=None,
time_order=2, space_order=space_order)
v = TimeFunction(name='v', grid=model.grid, staggered=None,
time_order=2, space_order=space_order)
vv = TimeFunction(name='vv', grid=model.grid, staggered=None,
time_order=2, space_order=space_order)
uu = TimeFunction(name='uu', grid=model.grid, staggered=None,
time_order=2, space_order=space_order)
cp = DevitoCheckpoint([u, v])
op_fwd = solver_tti.op_fwd(save=False)
op_fwd.cfunction
op_imaging.cfunction
wrap_fw = CheckpointOperator(op_fwd, src=geometry.src,
u=u, v=v, vp=model.vp, epsilon=model.epsilon,
delta=model.delta, theta=model.theta, dt=model.critical_dt)
time_range = TimeAxis(start=0, stop=(num_samples-1)*samp_int, step=samp_int)
dobs = Receiver(name='dobs', grid=model.grid, time_range=time_range, coordinates=geometry.rec_positions)
if not is_empty:
dobs.data[:] = retrieved_shot[idx_tr, :].T
else:
dobs.data[:] = retrieved_shot[:].T
dobs_resam = dobs.resample(num=geometry.nt)
g.write('Shape of residual: {}\n'.format(dobs_resam.data.shape))
wrap_rev = CheckpointOperator(op_imaging, u=u, v=v, vv=vv, uu=uu, vp=model.vp,
epsilon=model.epsilon, delta=model.delta, theta=model.theta,
dt=model.critical_dt, residual=dobs_resam.data)
# Run forward
wrp = Revolver(cp, wrap_fw, wrap_rev, n_checkpoints, dobs_resam.shape[0]-2)
g.write('Revolver storage: {}\n'.format(humanbytes(cp.size*n_checkpoints*itemsize)))
wrp.apply_forward()
g.write('{} run finished\n'.format(strng))
summary = wrp.apply_reverse()
form = 'image_{}.bin'.format(str(record).zfill(5))
h = open(form, 'wb')
g.write('{}\n'.format(str(image.data.shape)))
np.transpose(image.data).astype('float32').tofile(h)
else:
# For illustrative purposes, assuming that there is enough memory
g.write('enough memory to save full fld: {} == {}\n'.format(full_fld_mem, humanbytes(full_fld_mem)))
op_imaging = ImagingOperator(geometry, image, space_order)
vv = TimeFunction(name='vv', grid=model.grid, staggered=None, time_order=2, space_order=space_order)
uu = TimeFunction(name='uu', grid=model.grid, staggered=None, time_order=2, space_order=space_order)
time_range = TimeAxis(start=0, stop=(num_samples-1)*samp_int, step=samp_int)
dobs = Receiver(name='dobs', grid=model.grid, time_range=time_range, coordinates=geometry.rec_positions)
if not is_empty:
dobs.data[:] = retrieved_shot[idx_tr, :].T
else:
dobs.data[:] = retrieved_shot[:].T
dobs_resam = dobs.resample(num=geometry.nt)
u, v = solver_tti.forward(vp=model.vp, epsilon=model.epsilon, delta=model.delta,
theta=model.theta, dt=model.critical_dt, save=True)[1:-1]
op_imaging(u=u, v=v, vv=vv, uu=uu, epsilon=model.epsilon, delta=model.delta,
theta=model0.theta, vp=model.vp, dt=model.critical_dt, residual=dobs_resam)
full_image = extend_image(origin, vp, model, image)
return full_image
