examples: add source as optional input for checkpointing that runs fo…

…rward
devitocodes · Oct 5, 2021 · 624e3ba · 624e3ba
1 parent a7d3b38
commit 624e3ba
Showing 1 changed file with 3 additions and 2 deletions.
diff --git a/examples/seismic/acoustic/wavesolver.py b/examples/seismic/acoustic/wavesolver.py
@@ -157,7 +157,7 @@ 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, v=None, grad=None, model=None,
+    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
@@ -200,7 +200,8 @@ def jacobian_adjoint(self, rec, u, v=None, grad=None, model=None,
                              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,
+            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)