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_ptr_loads_and_saves(nt, ncp): cp = SimpleCheckpoint() f = SimpleOperator() b = SimpleOperator() rev = Revolver(cp, f, b, ncp, nt) rev.apply_forward() rev.apply_reverse() assert(cp.save_pointers == cp.load_pointers) assert(len(cp.save_pointers) == min(ncp, nt - 1))
def test_reverse_nt(nt, ncp): cp = SimpleCheckpoint() f = SimpleOperator() b = SimpleOperator() rev = Revolver(cp, f, b, ncp, nt) rev.apply_forward() assert(b.counter == 0) rev.apply_reverse() assert(b.counter == nt)
def test_num_loads_and_saves(nt, ncp): cp = SimpleCheckpoint() f = SimpleOperator() b = SimpleOperator() rev = Revolver(cp, f, b, ncp, nt) rev.apply_forward() assert(cp.load_counter == 0) rev.apply_reverse() assert(cp.load_counter >= cp.save_counter)
def revolve_factors(nt, ncp): fwd = CounterOperator() rev = CounterOperator() cp = CounterCheckpoint() revolver = Revolver(cp, fwd, rev, ncp, nt) #, compression_params={'scheme': None}) revolver.apply_forward() assert (fwd.counter == nt) revolver.apply_reverse() assert (rev.counter == nt) return fwd.counter, cp.savecount, cp.loadcount
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 jacobian_adjoint(self, rec, u, src=None, v=None, grad=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. u : TimeFunction Full wavefield `u` (created with save=True). v : TimeFunction, optional Stores the computed wavefield. grad : Function, optional Stores the gradient field. model : Model, optional Object containing the physical parameters. 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) model = model or self.model # Pick vp from model unless explicitly provided kwargs.update(model.physical_params(**kwargs)) 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=src or self.geometry.src, u=u, dt=dt, **kwargs) wrap_rev = CheckpointOperator(self.op_grad(save=False), u=u, v=v, 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: summary = self.op_grad().apply(rec=rec, grad=grad, v=v, u=u, dt=dt, **kwargs) return grad, summary
def test_complete(scheme): nt = 100 ncp = 10 shape = (10, 10) a = np.zeros(shape) fwd = IncrementOperator(1, a) rev = IncrementOperator(-1, a) cp = YoCheckpoint(a) compression_params = {'scheme': scheme} revolver = Revolver(cp, fwd, rev, ncp, nt, compression_params=compression_params) revolver.apply_forward() assert (np.all(np.isclose(a, np.zeros(shape) + nt))) revolver.apply_reverse() assert (np.all(np.isclose(a, np.zeros(shape))))
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 test_compression_is_used(): nt = 100 ncp = 10 shape = (10, 10) a = np.zeros(shape) fwd = IncrementOperator(1, a) rev = IncrementOperator(-1, a) cp = YoCheckpoint(a) counters = [0, 0] compressor, decompressor = compressors[None], decompressors[None] def this_compressor(params, data): counters[0] += 1 return compressor(params, data) def this_decompressor(params, data): counters[1] += 1 return decompressor(params, data) compression_params = { 'scheme': 'custom', 'compressor': this_compressor, 'decompressor': this_decompressor } revolver = Revolver(cp, fwd, rev, ncp, nt, compression_params=compression_params) revolver.apply_forward() assert (counters[0] >= ncp) revolver.apply_reverse() assert (counters[1] >= counters[0])
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 adjoint_born(model, rec_coords, rec_data, u=None, op_forward=None, is_residual=False, space_order=8, nb=40, isic=False, dt=None, n_checkpoints=None, maxmem=None): clear_cache() # Parameters nt = rec_data.shape[0] if dt is None: dt = model.critical_dt m, rho, damp = model.m, model.rho, model.damp # Create adjoint wavefield and gradient v = TimeFunction(name='v', grid=model.grid, time_order=2, space_order=space_order) gradient = Function(name='gradient', grid=model.grid) # Set up PDE and rearrange vlaplace, rho = acoustic_laplacian(v, rho) H = symbols('H') eqn = m / rho * v.dt2 - H - damp * v.dt stencil = solve(eqn, v.backward, simplify=False, rational=False)[0] expression = [Eq(v.backward, stencil.subs({H: vlaplace}))] # Data at receiver locations as adjoint source rec_g = Receiver(name='rec_g', grid=model.grid, ntime=nt, coordinates=rec_coords) if op_forward is None: rec_g.data[:] = rec_data[:] adj_src = rec_g.inject(field=v.backward, offset=model.nbpml, expr=rec_g * rho * dt**2 / m) # Gradient update if u is None: u = TimeFunction(name='u', grid=model.grid, time_order=2, space_order=space_order) if isic is not True: gradient_update = [Eq(gradient, gradient - dt * u.dt2 / rho * v)] else: # sum u.dx * v.dx fo x in dimensions. # space_order//2 diff_u_v = sum([ first_derivative(u, dim=d, order=space_order // 2) * first_derivative(v, dim=d, order=space_order // 2) for d in u.space_dimensions ]) gradient_update = [ Eq(gradient, gradient - dt * (u * v.dt2 * m + diff_u_v) / rho) ] # Create operator and run set_log_level('ERROR') expression += adj_src + gradient_update subs = model.spacing_map subs[u.grid.time_dim.spacing] = dt op = Operator(expression, subs=subs, dse='advanced', dle='advanced', name="Gradient%s" % randint(1e5)) # Optimal checkpointing if op_forward is not None: rec = Receiver(name='rec', grid=model.grid, ntime=nt, coordinates=rec_coords) cp = DevitoCheckpoint([u]) if maxmem is not None: n_checkpoints = int( np.floor(maxmem * 10**6 / (cp.size * u.data.itemsize))) wrap_fw = CheckpointOperator(op_forward, u=u, m=model.m, rec=rec) wrap_rev = CheckpointOperator(op, u=u, v=v, m=model.m, rec_g=rec_g) # Run forward wrp = Revolver(cp, wrap_fw, wrap_rev, n_checkpoints, nt - 2) wrp.apply_forward() # Residual and gradient if is_residual is True: # input data is already the residual rec_g.data[:] = rec_data[:] else: rec_g.data[:] = rec.data[:] - rec_data[:] # input is observed data fval = .5 * np.dot(rec_g.data[:].flatten(), rec_g.data[:].flatten()) * dt wrp.apply_reverse() else: op() clear_cache() if op_forward is not None and is_residual is not True: return fval, gradient.data else: return gradient.data
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 fwi_gradient_checkpointed(vp_in, model, geometry, nshots, client, solver_params, n_checkpoints=1000, compression_params=None): # Create symbols to hold the gradient and residual grad = Function(name="grad", grid=model.grid) vp = Function(name="vp", grid=model.grid) smooth_d = Receiver(name='rec', grid=model.grid, time_range=geometry.time_axis, coordinates=geometry.rec_positions) residual = Receiver(name='rec', grid=model.grid, time_range=geometry.time_axis, coordinates=geometry.rec_positions) objective = 0. time_order = 2 vp_in = vec2mat(vp_in, model.shape) assert (model.vp.shape == vp_in.shape) vp.data[:] = vp_in[:] filename = solver_params['filename'] solver = overthrust_solver_iso(filename, datakey="m0") dt = solver.dt nt = smooth_d.data.shape[0] - 2 u = TimeFunction(name='u', grid=model.grid, time_order=time_order, space_order=4) v = TimeFunction(name='v', grid=model.grid, time_order=time_order, space_order=4) fwd_op = solver.op_fwd(save=False) rev_op = solver.op_grad(save=False) cp = DevitoCheckpoint([u]) for i in range(nshots): true_d, source_location = load_shot(i) # Update source location solver.geometry.src_positions[0, :] = source_location[:] # Compute smooth data and full forward wavefield u0 u.data[:] = 0. residual.data[:] = 0. v.data[:] = 0. smooth_d.data[:] = 0. 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) wrp = Revolver(cp, wrap_fw, wrap_rev, n_checkpoints, nt, compression_params=compression_params) wrp.apply_forward() # Compute gradient from data residual and update objective function residual.data[:] = smooth_d.data[:] - true_d[:] objective += .5 * np.linalg.norm(residual.data.ravel())**2 wrp.apply_reverse() print(wrp.profiler.summary()) return objective, -np.ravel(grad.data).astype(np.float64)
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 adjoint_born(model, rec_coords, rec_data, u=None, op_forward=None, is_residual=False, space_order=8, isic=False, dt=None, n_checkpoints=None, maxmem=None, free_surface=False, tsub_factor=1, checkpointing=False): clear_cache() # Parameters nt = rec_data.shape[0] if dt is None: dt = model.critical_dt m, rho, damp = model.m, model.rho, model.damp # Create adjoint wavefield and gradient v = TimeFunction(name='v', grid=model.grid, time_order=2, space_order=space_order) gradient = Function(name='gradient', grid=model.grid) # Set up PDE and rearrange vlaplace, rho = acoustic_laplacian(v, rho) stencil = damp * (2.0 * v - damp * v.forward + dt**2 * rho / m * vlaplace) expression = [Eq(v.backward, stencil)] # Data at receiver locations as adjoint source rec_g = Receiver(name='rec_g', grid=model.grid, ntime=nt, coordinates=rec_coords) if op_forward is None: rec_g.data[:] = rec_data[:] adj_src = rec_g.inject(field=v.backward, expr=rec_g * rho * dt**2 / m) # Gradient update if u is None: u = TimeFunction(name='u', grid=model.grid, time_order=2, space_order=space_order) if isic is not True: gradient_update = [Inc(gradient, - dt * u.dt2 / rho * v)] else: # sum u.dx * v.dx fo x in dimensions. # space_order//2 diff_u_v = sum([first_derivative(u, dim=d, fd_order=space_order//2)* first_derivative(v, dim=d, fd_order=space_order//2) for d in u.space_dimensions]) gradient_update = [Inc(gradient, - tsub_factor * dt * (u * v.dt2 * m + diff_u_v) / rho)] # Create operator and run # Free surface if free_surface is True: expression += freesurface(v, space_order//2, model.nbpml, forward=False) expression += adj_src + gradient_update subs = model.spacing_map subs[u.grid.time_dim.spacing] = dt op = Operator(expression, subs=subs, dse='advanced', dle='advanced') # Optimal checkpointing summary1 = None summary2 = None if op_forward is not None and checkpointing is True: rec = Receiver(name='rec', grid=model.grid, ntime=nt, coordinates=rec_coords) cp = DevitoCheckpoint([u]) if maxmem is not None: n_checkpoints = int(np.floor(maxmem * 10**6 / (cp.size * u.data.itemsize))) wrap_fw = CheckpointOperator(op_forward, u=u, m=model.m, rec=rec) wrap_rev = CheckpointOperator(op, u=u, v=v, m=model.m, rec_g=rec_g) # Run forward wrp = Revolver(cp, wrap_fw, wrap_rev, n_checkpoints, nt-2) wrp.apply_forward() # Residual and gradient if is_residual is True: # input data is already the residual rec_g.data[:] = rec_data[:] else: rec_g.data[:] = rec.data[:] - rec_data[:] # input is observed data fval = .5*np.dot(rec_g.data[:].flatten(), rec_g.data[:].flatten()) * dt wrp.apply_reverse() elif op_forward is not None and checkpointing is False: # Compile first cf1 = op_forward.cfunction cf2 = op.cfunction # Run forward and adjoint summary1 = op_forward.apply() if is_residual is True: rec_g.data[:] = rec_data[:] else: rec = Receiver(name='rec', grid=model.grid, ntime=nt, coordinates=rec_coords) rec_g.data[:] = rec.data[:] - rec_data[:] fval = .5*np.dot(rec_g.data[:].flatten(), rec_g.data[:].flatten()) * dt summary2 = op.apply() else: cf = op.cfunction summary1 = op.apply() clear_cache() if op_forward is not None and is_residual is not True: if summary2 is not None: return fval, gradient.data, summary1, summary2 elif summary1 is not None: return fval, gradient.data, summary1 else: return fval, gradient.data else: if summary2 is not None: return gradient.data, summary1, summary2 elif summary1 is not None: return gradient.data, summary1 else: return gradient.data
def adjoint_born(model, rec_coords, rec_data, u=None, op_forward=None, is_residual=False, space_order=8, nb=40, isic=False, dt=None): clear_cache() # Parameters nt = rec_data.shape[0] if dt is None: dt = model.critical_dt m, damp = model.m, model.damp # Create adjoint wavefield and gradient v = TimeFunction(name='v', grid=model.grid, time_order=2, space_order=space_order) gradient = Function(name='gradient', grid=model.grid) # Set up PDE and rearrange eqn = m * v.dt2 - v.laplace - damp * v.dt stencil = solve(eqn, v.backward)[0] expression = [Eq(v.backward, stencil)] # Data at receiver locations as adjoint source rec_g = Receiver(name='rec_g', grid=model.grid, ntime=nt, coordinates=rec_coords) if op_forward is None: rec_g.data[:] = rec_data[:] adj_src = rec_g.inject(field=v.backward, offset=model.nbpml, expr=rec_g * dt**2 / m) # Gradient update if u is None: u = TimeFunction(name='u', grid=model.grid, time_order=2, space_order=space_order) if isic is not True: gradient_update = [Eq(gradient, gradient - u * v.dt2) ] # zero-lag cross-correlation imaging condition else: # linearized inverse scattering imaging condition (Op't Root et al. 2010; Whitmore and Crawley 2012) if len(model.shape) == 2: gradient_update = [ Eq(gradient, gradient - (u * v.dt2 * m + u.dx * v.dx + u.dy * v.dy)) ] else: gradient_update = [ Eq( gradient, gradient - (u * v.dt2 * m + u.dx * v.dx + u.dy * v.dy + u.dz * v.dz)) ] # Create operator and run set_log_level('ERROR') expression += adj_src + gradient_update op = Operator(expression, subs=model.spacing_map, dse='advanced', dle='advanced', name="Gradient%s" % randint(1e5)) # Optimal checkpointing if op_forward is not None: rec = Receiver(name='rec', grid=model.grid, ntime=nt, coordinates=rec_coords) cp = DevitoCheckpoint([u]) n_checkpoints = None wrap_fw = CheckpointOperator(op_forward, u=u, m=model.m.data, rec=rec, dt=dt) wrap_rev = CheckpointOperator(op, u=u, v=v, m=model.m.data, rec_g=rec_g, dt=dt) # Run forward wrp = Revolver(cp, wrap_fw, wrap_rev, n_checkpoints, nt - 2) wrp.apply_forward() # Residual and gradient if is_residual is True: # input data is already the residual rec_g.data[:] = rec_data[:] else: rec_g.data[:] = rec.data[:] - rec_data[:] # input is observed data fval = .5 * np.linalg.norm(rec_g.data[:])**2 wrp.apply_reverse() else: op(dt=dt) clear_cache() if op_forward is not None and is_residual is not True: return fval, gradient.data else: return gradient.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
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 J_adjoint_checkpointing(model, src_coords, wavelet, rec_coords, recin, space_order=8, is_residual=False, n_checkpoints=None, maxmem=None, return_obj=False, isic=False, ws=None, t_sub=1): """ Jacobian (adjoint fo born modeling operator) operator on a shot record as a source (i.e data residual). Outputs the gradient with Checkpointing. Parameters ---------- model: Model Physical model src_coords: Array Coordiantes of the source(s) wavelet: Array Source signature rec_coords: Array Coordiantes of the receiver(s) recin: Array Receiver data space_order: Int (optional) Spatial discretization order, defaults to 8 checkpointing: Bool Whether or not to use checkpointing n_checkpoints: Int Number of checkpoints for checkpointing maxmem: Float Maximum memory to use for checkpointing isic : Bool Whether or not to use ISIC imaging condition ws : Array Extended source spatial distribution is_residual: Bool Whether to treat the input as the residual or as the observed data Returns ---------- Array Adjoint jacobian on the input data (gradient) """ # Optimal checkpointing op_f, u, rec_g = forward(model, src_coords, rec_coords, wavelet, space_order=space_order, return_op=True, ws=ws) op, g, v = gradient(model, recin, rec_coords, u, space_order=space_order, return_op=True, isic=isic) nt = wavelet.shape[0] rec = Receiver(name='rec', grid=model.grid, ntime=nt, coordinates=rec_coords) cp = DevitoCheckpoint([uu for uu in as_tuple(u)]) if maxmem is not None: memsize = (cp.size * u.data.itemsize) n_checkpoints = int(np.floor(maxmem * 10**6 / memsize)) # Op arguments uk = {uu.name: uu for uu in as_tuple(u)} vk = {**uk, **{vv.name: vv for vv in as_tuple(v)}} uk.update({'rcv%s' % as_tuple(u)[0].name: rec_g}) vk.update({'src%s' % as_tuple(v)[0].name: rec}) # Wrapped ops wrap_fw = CheckpointOperator(op_f, vp=model.vp, **uk) wrap_rev = CheckpointOperator(op, vp=model.vp, **vk) # Run forward wrp = Revolver(cp, wrap_fw, wrap_rev, n_checkpoints, nt - 2) wrp.apply_forward() # Residual and gradient if is_residual is True: # input data is already the residual rec.data[:] = recin[:] else: rec.data[:] = rec.data[:] - recin[:] # input is observed data wrp.apply_reverse() if return_obj: return .5 * model.critical_dt * norm(rec)**2, g.data return g.data
def process_shot_checkpointed(i, solver, shots_container, auth, exclude_boundaries=True, checkpoint_params=None): model = solver.model so = solver.space_order geometry = solver.geometry time_order = 2 dt = solver.model.critical_dt # 1.75 nt = solver.geometry.time_axis.num - time_order if checkpoint_params is not None: n_checkpoints = checkpoint_params.pop('n_checkpoints', 1000) compression_params = checkpoint_params else: n_checkpoints = 1000 compression_params = None # Create symbols to hold the gradient and residual grad = Function(name="grad", grid=model.grid) smooth_d = Receiver(name='rec', grid=model.grid, time_range=geometry.time_axis, coordinates=geometry.rec_positions) residual = Receiver(name='rec', grid=model.grid, time_range=geometry.time_axis, coordinates=geometry.rec_positions) u = TimeFunction(name='u', grid=model.grid, time_order=time_order, space_order=so) v = TimeFunction(name='v', grid=model.grid, time_order=time_order, space_order=so) fwd_op = solver.op_fwd(save=False) rev_op = solver.op_grad(save=False) cp = DevitoCheckpoint([u]) true_d, source_location, old_dt = load_shot(i, auth=auth, container=shots_container) true_d = reinterpolate(true_d, solver.geometry.nt, old_dt) # Update source location solver.geometry.src_positions[0, :] = source_location[:] wrap_fw = CheckpointOperator(fwd_op, src=solver.geometry.src, u=u, rec=smooth_d, dt=dt, vp=solver.model.vp) wrap_rev = CheckpointOperator(rev_op, u=u, v=v, rec=residual, grad=grad, dt=dt, vp=solver.model.vp) wrp = Revolver(cp, wrap_fw, wrap_rev, n_checkpoints, nt, compression_params=compression_params) wrp.apply_forward() # Compute gradient from data residual and update objective function residual.data[:] = smooth_d.data[:] - true_d[:] objective = .5 * np.linalg.norm(residual.data.ravel())**2 wrp.apply_reverse() # Prepare for serialization before returning if exclude_boundaries: grad = trim_boundary(grad, solver.model.nbl) else: grad = grad.data grad = np.array(grad, dtype=solver.model.dtype) return objective, grad
def jacobian_adjoint(self, rec, u, v=None, grad=None, precon=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: if precon is not None: precon = precon #don't know if this is needded but I'll just follow style summary = self.op_grad_precon().apply(rec=rec, grad=grad, precon=precon, v=v, u=u, vp=vp, dt=dt, **kwargs) else: summary = self.op_grad().apply(rec=rec, grad=grad, v=v, u=u, vp=vp, dt=dt, **kwargs) return grad, summary