ckp: Checkpointing update

examples/checkpointing/checkpoint.py → devito/checkpointing/checkpoint.py

@@ -6,8 +6,8 @@
 class CheckpointOperator(Operator):
     """Devito's concrete implementation of the ABC pyrevolve.Operator. This class wraps
        devito.Operator so it conforms to the pyRevolve API. pyRevolve will call apply
-       with arguments t_start and t_end. Devito calls these arguments t_s and t_e so
-       the following dict is used to perform the translations between different names.
+       with arguments t_start and t_end. Devito calls these arguments time_m and time_M
+       so the following dict is used to perform the translations between different names.
 
        Parameters
        ----------
@@ -23,9 +23,8 @@ class CheckpointOperator(Operator):
 
     def __init__(self, op, **kwargs):
         self.op = op
-        self.args = kwargs
-        op_default_args = self.op._prepare_arguments(**kwargs)
-        self.start_offset = op_default_args[self.t_arg_names['t_start']]
+        self.args = self.op._prepare_arguments(**kwargs)
+        self.start_offset = self.args[self.t_arg_names['t_start']]
 
     def _prepare_args(self, t_start, t_end):
         args = self.args.copy()
@@ -40,7 +39,7 @@ def apply(self, t_start, t_end):
             this method passes them on correctly to devito.Operator
         """
         # Build the arguments list to invoke the kernel function
-        args = self.op.arguments(**self._prepare_args(t_start, t_end))
+        args = self._prepare_args(t_start, t_end)
         # Invoke kernel function with args
         arg_values = [args[p.name] for p in self.op.parameters]
         self.op.cfunction(*arg_values)

examples/seismic/acoustic/wavesolver.py

@@ -1,9 +1,9 @@
 from devito import Function, TimeFunction
 from devito.tools import memoized_meth
+from devito.checkpointing import DevitoCheckpoint, CheckpointOperator
 from examples.seismic.acoustic.operators import (
     ForwardOperator, AdjointOperator, GradientOperator, BornOperator
 )
-from examples.checkpointing.checkpoint import DevitoCheckpoint, CheckpointOperator
 from pyrevolve import Revolver
 
 

examples/seismic/tti/wavesolver.py

@@ -1,10 +1,10 @@
 # coding: utf-8
 from devito import Function, TimeFunction, warning
 from devito.tools import memoized_meth
+from devito.checkpointing import DevitoCheckpoint, CheckpointOperator
 from examples.seismic.tti.operators import ForwardOperator, AdjointOperator
 from examples.seismic.tti.operators import JacobianOperator, JacobianAdjOperator
 from examples.seismic.tti.operators import particle_velocity_fields
-from examples.checkpointing.checkpoint import DevitoCheckpoint, CheckpointOperator
 from pyrevolve import Revolver
 
 

examples/seismic/viscoacoustic/wavesolver.py

@@ -1,10 +1,10 @@
 from devito import VectorTimeFunction, TimeFunction, Function, NODE
+from devito.checkpointing import DevitoCheckpoint, CheckpointOperator
 from devito.tools import memoized_meth
 from examples.seismic import PointSource
 from examples.seismic.viscoacoustic.operators import (
     ForwardOperator, AdjointOperator, GradientOperator, BornOperator
 )
-from examples.checkpointing.checkpoint import DevitoCheckpoint, CheckpointOperator
 from pyrevolve import Revolver
 
 

tests/test_checkpointing.py

@@ -5,7 +5,7 @@
 import numpy as np
 
 from devito import Grid, TimeFunction, Operator, Function, Eq, switchconfig, Constant
-from examples.checkpointing.checkpoint import DevitoCheckpoint, CheckpointOperator
+from devito.checkpointing import DevitoCheckpoint, CheckpointOperator
 from examples.seismic.acoustic.acoustic_example import acoustic_setup