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 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 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 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