From 44694f5ba69a0aeb3c44f7a705eb85981c2fcaef Mon Sep 17 00:00:00 2001 From: malfarhan7 Date: Sat, 26 Oct 2024 12:23:45 +0300 Subject: [PATCH] Implement snapshotting for the acoustic wave equation --- examples/seismic/acoustic/operators.py | 89 +++++++++++++++++++------ examples/seismic/acoustic/wavesolver.py | 66 +++++++++++++----- examples/seismic/utils.py | 2 + 3 files changed, 122 insertions(+), 35 deletions(-) diff --git a/examples/seismic/acoustic/operators.py b/examples/seismic/acoustic/operators.py index c8a2fdec5c..ac3210a6f3 100644 --- a/examples/seismic/acoustic/operators.py +++ b/examples/seismic/acoustic/operators.py @@ -1,4 +1,4 @@ -from devito import Eq, Operator, Function, TimeFunction, Inc, solve, sign +from devito import Eq, Operator, Function, TimeFunction, Inc, solve, sign, ConditionalDimension from devito.symbolics import retrieve_functions, INT, retrieve_derivatives @@ -108,7 +108,7 @@ def iso_stencil(field, model, kernel, **kwargs): def ForwardOperator(model, geometry, space_order=4, - save=False, kernel='OT2', **kwargs): + save=False, kernel='OT2', factor=None, **kwargs): """ Construct a forward modelling operator in an acoustic medium. @@ -126,6 +126,8 @@ def ForwardOperator(model, geometry, space_order=4, Defaults to False. kernel : str, optional Type of discretization, 'OT2' or 'OT4'. + factor : int, optional + Downsampling factor to save snapshots of the wavefield. """ m = model.m @@ -144,10 +146,28 @@ def ForwardOperator(model, geometry, space_order=4, # Create interpolation expression for receivers rec_term = rec.interpolate(expr=u) - + # Build operator equations + equations = eqn + src_term + rec_term + + if factor: + # Implement snapshotting + nsnaps = (geometry.nt + factor - 1) // factor + time_subsampled = ConditionalDimension( + 't_sub', parent=model.grid.time_dim, factor=factor) + usnaps = TimeFunction(name='usnaps', grid=model.grid, + time_order=2, space_order=space_order, + save=nsnaps, time_dim=time_subsampled) + # Add equation to save snapshots + snapshot_eq = Eq(usnaps, u) + equations += [snapshot_eq] + else: + usnaps = None # Substitute spacing terms to reduce flops - return Operator(eqn + src_term + rec_term, subs=model.spacing_map, - name='Forward', **kwargs) + op = Operator(equations, subs=model.spacing_map, name='Forward', **kwargs) + if usnaps is not None: + return op, usnaps + else: + return op def AdjointOperator(model, geometry, space_order=4, @@ -189,8 +209,8 @@ def AdjointOperator(model, geometry, space_order=4, def GradientOperator(model, geometry, space_order=4, save=True, - kernel='OT2', **kwargs): - """ + kernel='OT2', factor=None, **kwargs): + """ Construct a gradient operator in an acoustic media. Parameters @@ -206,30 +226,59 @@ def GradientOperator(model, geometry, space_order=4, save=True, Option to store the entire (unrolled) wavefield. kernel : str, optional Type of discretization, centered or shifted. + factor : int, optional + Downsampling factor to save snapshots of the wavefield. """ m = model.m - # Gradient symbol and wavefield symbols + # Gradient symbol grad = Function(name='grad', grid=model.grid) - u = TimeFunction(name='u', grid=model.grid, save=geometry.nt if save - else None, time_order=2, space_order=space_order) - v = TimeFunction(name='v', grid=model.grid, save=None, - time_order=2, space_order=space_order) - rec = geometry.rec + + # Create the adjoint wavefield + v = TimeFunction(name='v', grid=model.grid, time_order=2, space_order=space_order) s = model.grid.stepping_dim.spacing eqn = iso_stencil(v, model, kernel, forward=False) - if kernel == 'OT2': - gradient_update = Inc(grad, - u * v.dt2) - elif kernel == 'OT4': - gradient_update = Inc(grad, - u * v.dt2 - s**2 / 12.0 * u.biharmonic(m**(-2)) * v) # Add expression for receiver injection + rec = geometry.rec receivers = rec.inject(field=v.backward, expr=rec * s**2 / m) + time = model.grid.time_dim + + if factor is not None: + # Condition to apply gradient update only at snapshot times + condition = Eq(time % factor, 0) + # Create the ConditionalDimension for subsampling + time_subsampled = ConditionalDimension('t_sub', parent=time, factor=factor) + # Define usnaps with time_subsampled as its time dimension + nsnaps = (geometry.nt + factor - 1) // factor + usnaps = TimeFunction(name='usnaps', grid=model.grid, + time_order=2, space_order=space_order, + save=nsnaps, time_dim=time_subsampled) + # Gradient update without indexing usnaps + if kernel == 'OT2': + gradient_update = Inc(grad, - usnaps * v.dt2, implicit_dims=[time_subsampled], + condition=condition) + elif kernel == 'OT4': + gradient_update = Inc(grad, - usnaps * v.dt2 + - s**2 / 12.0 * usnaps.biharmonic(m**(-2)) * v, + implicit_dims=[time_subsampled], + condition=condition) + else: + u = TimeFunction(name='u', grid=model.grid, + save=geometry.nt if save else None, + time_order=2, space_order=space_order) + if kernel == 'OT2': + gradient_update = Inc(grad, - u * v.dt2) + elif kernel == 'OT4': + gradient_update = Inc(grad, - u * v.dt2 + - s**2 / 12.0 * u.biharmonic(m**(-2)) * v) + # Substitute spacing terms to reduce flops - return Operator(eqn + receivers + [gradient_update], subs=model.spacing_map, - name='Gradient', **kwargs) + op = Operator(eqn + receivers + [gradient_update], subs=model.spacing_map, + name='Gradient', **kwargs) + return op def BornOperator(model, geometry, space_order=4, @@ -274,4 +323,4 @@ def BornOperator(model, geometry, space_order=4, # Substitute spacing terms to reduce flops return Operator(eqn1 + source + eqn2 + receivers, subs=model.spacing_map, - name='Born', **kwargs) + name='Born', **kwargs) \ No newline at end of file diff --git a/examples/seismic/acoustic/wavesolver.py b/examples/seismic/acoustic/wavesolver.py index e191b34972..b1feaef9de 100644 --- a/examples/seismic/acoustic/wavesolver.py +++ b/examples/seismic/acoustic/wavesolver.py @@ -1,6 +1,6 @@ from devito import Function, TimeFunction, DevitoCheckpoint, CheckpointOperator, Revolver from devito.tools import memoized_meth -from examples.seismic.acoustic.operators import ( +from devitofwi.devito.acoustic.operators import ( ForwardOperator, AdjointOperator, GradientOperator, BornOperator ) @@ -23,6 +23,7 @@ class AcousticWaveSolver: space_order: int, optional Order of the spatial stencil discretisation. Defaults to 4. """ + def __init__(self, model, geometry, kernel='OT2', space_order=4, **kwargs): self.model = model self.model._initialize_bcs(bcs="damp") @@ -44,11 +45,11 @@ def dt(self): return self.model.critical_dt @memoized_meth - def op_fwd(self, save=None): + def op_fwd(self, save=None, factor=None): """Cached operator for forward runs with buffered wavefield""" return ForwardOperator(self.model, save=save, geometry=self.geometry, kernel=self.kernel, space_order=self.space_order, - **self._kwargs) + factor=factor, **self._kwargs) @memoized_meth def op_adj(self): @@ -58,11 +59,11 @@ def op_adj(self): **self._kwargs) @memoized_meth - def op_grad(self, save=True): + def op_grad(self, save=True, factor=None): """Cached operator for gradient runs""" return GradientOperator(self.model, save=save, geometry=self.geometry, kernel=self.kernel, space_order=self.space_order, - **self._kwargs) + factor=factor, **self._kwargs) @memoized_meth def op_born(self): @@ -71,7 +72,7 @@ def op_born(self): kernel=self.kernel, space_order=self.space_order, **self._kwargs) - def forward(self, src=None, rec=None, u=None, model=None, save=None, **kwargs): + def forward(self, src=None, rec=None, u=None, model=None, save=None, factor=None, **kwargs): """ Forward modelling function that creates the necessary data objects for running a forward modelling operator. @@ -90,6 +91,8 @@ def forward(self, src=None, rec=None, u=None, model=None, save=None, **kwargs): The time-constant velocity. save : bool, optional Whether or not to save the entire (unrolled) wavefield. + factor : int, optional + Downsampling factor to save snapshots of the wavefield. Returns ------- @@ -108,12 +111,24 @@ def forward(self, src=None, rec=None, u=None, model=None, save=None, **kwargs): model = model or self.model # Pick vp from model unless explicitly provided kwargs.update(model.physical_params(**kwargs)) + # Get the operator + op_fwd = self.op_fwd(save=save, factor=factor) + # Prepare parameters for operator apply + op_args = {'src': src, 'rec': rec, 'u': u, 'dt': kwargs.pop('dt', self.dt)} + op_args.update(kwargs) # Execute operator and return wavefield and receiver data - summary = self.op_fwd(save).apply(src=src, rec=rec, u=u, - dt=kwargs.pop('dt', self.dt), **kwargs) - - return rec, u, summary + if factor: + # Operator returned is op, usnaps + op, usnaps = op_fwd + op_args['usnaps'] = usnaps + summary = op.apply(**op_args) + + else: + op = op_fwd + usnaps = None + summary = op.apply(**op_args) + return rec, u, usnaps, summary def adjoint(self, rec, srca=None, v=None, model=None, **kwargs): """ @@ -155,8 +170,8 @@ def adjoint(self, rec, srca=None, v=None, model=None, **kwargs): dt=kwargs.pop('dt', self.dt), **kwargs) return srca, v, summary - def jacobian_adjoint(self, rec, u, src=None, v=None, grad=None, model=None, - checkpointing=False, **kwargs): + def jacobian_adjoint(self, rec, u=None, usnaps=None, src=None, v=None, grad=None, model=None, + factor=None, checkpointing=False, **kwargs): """ Gradient modelling function for computing the adjoint of the Linearized Born modelling function, ie. the action of the @@ -168,6 +183,8 @@ def jacobian_adjoint(self, rec, u, src=None, v=None, grad=None, model=None, Receiver data. u : TimeFunction Full wavefield `u` (created with save=True). + usnaps : TimeFunction + Snapshots of the wavefield `u`. v : TimeFunction, optional Stores the computed wavefield. grad : Function, optional @@ -176,12 +193,22 @@ def jacobian_adjoint(self, rec, u, src=None, v=None, grad=None, model=None, Object containing the physical parameters. vp : Function or float, optional The time-constant velocity. + checkpointing : boolean, optional + Flag to enable checkpointing (default False). + Cannot be used with snapshotting. + factor : int, optional + Downsampling factor for the saved snapshots of the wavefield `u`. + Cannot be used with checkpointing. Returns ------- Gradient field and performance summary. """ dt = kwargs.pop('dt', self.dt) + # Check that snapshotting and checkpointing are not used together + if factor is not None and checkpointing: + raise ValueError("Cannot use snapshotting (factor) and checkpointing simultaneously.") + # Gradient symbol grad = grad or Function(name='grad', grid=self.model.grid) @@ -209,8 +236,17 @@ def jacobian_adjoint(self, rec, u, src=None, v=None, grad=None, model=None, wrp.apply_forward() summary = wrp.apply_reverse() else: - summary = self.op_grad().apply(rec=rec, grad=grad, v=v, u=u, dt=dt, - **kwargs) + if factor is not None: + # Get the gradient operator + op = self.op_grad(save=False, factor=factor) + op_args = {'rec': rec, 'grad': grad, 'v': v, 'dt': dt, 'usnaps': usnaps} + else: + op = self.op_grad(save=True, factor=None) + op_args = {'rec': rec, 'grad': grad, 'v': v, 'dt': dt, 'u': u} + + op_args.update(kwargs) + summary = op.apply(**op_args) + return grad, summary def jacobian(self, dmin, src=None, rec=None, u=None, U=None, model=None, **kwargs): @@ -255,4 +291,4 @@ def jacobian(self, dmin, src=None, rec=None, u=None, U=None, model=None, **kwarg # Backward compatibility born = jacobian - gradient = jacobian_adjoint + gradient = jacobian_adjoint \ No newline at end of file diff --git a/examples/seismic/utils.py b/examples/seismic/utils.py index 6491d9ca5a..b9911381e8 100644 --- a/examples/seismic/utils.py +++ b/examples/seismic/utils.py @@ -264,4 +264,6 @@ def __call__(self, parser, args, values, option_string=None): choices=['float32', 'float64']) parser.add_argument("-interp", dest="interp", default="linear", choices=['linear', 'sinc']) + parser.add_argument("--factor", type=int, default=None, + help="Downsampling factor to use snapshotting, default is None") return parser