Fixed marmottant inconsistencies

stride/physics/common/devito.py

@@ -238,7 +238,7 @@ def __init__(self, space_order, time_order, time_dim=None, **kwargs):
 
         if space is None:
             origin = (0,)
-            extended_shape = (1,)
+            extended_shape = (2,)
             extended_extent = (1,)
         else:
             extra = space.absorbing
@@ -286,7 +286,7 @@ def __init__(self, space_order, time_order, time_dim=None, **kwargs):
         if parent_grid is not None:
             dimensions = parent_grid.dimensions
             time_dimension = devito.TimeDimension(name='time_inner',
-                                                  spacing=devito.types.Scalar(name='dt_inner', is_const=True))
+                                                  spacing=devito.Scalar(name='dt_inner', is_const=True))
             self.num_inner = kwargs.pop('num_inner', 1)
         else:
             self.num_inner = None
@@ -691,7 +691,7 @@ def sparse_function(self, name, num=1, space_order=None,
 
             reference_gridpoints, coefficients = self._calculate_hicks(coordinates)
 
-            fun = devito.PrecomputedSparseFunction(r=r,
+            fun = devito.PrecomputedSparseFunction(r=r+1,
                                                    gridpoints=reference_gridpoints,
                                                    interpolation_coeffs=coefficients,
                                                    **sparse_kwargs)

stride/physics/common/import_devito.py

@@ -1,7 +1,8 @@
 
 from devito import *  # noqa: F401
-from devito.types import Symbol  # noqa: F401
+from devito.types import Symbol, Scalar  # noqa: F401
 from devito.symbolics import INT, IntDiv  # noqa: F401
+from devito import TimeFunction as TimeFunctionOSS  # noqa: F401
 
 try:
     from devitopro import *  # noqa: F401

stride/physics/iso_acoustic/devito.py

@@ -381,7 +381,7 @@ def before_forward(self, wavelets, vp, rho=None, alpha=None, **kwargs):
                     # ValueError: Cannot access `shape_allocated` as unfinalized - so no size estimate
                     pass
 
-                if self._needs_grad(wavelets, rho, alpha):
+                if self._needs_grad(wavelets, rho, alpha, **kwargs):
                     p_saved_expr = p
                 else:
                     p_saved_expr = self._forward_save(p)
@@ -1595,8 +1595,9 @@ def _weights(self):
     def _dt_max(self, k, h, vp_max):
         return k * h / vp_max * 1 / np.sqrt(self.space.dim)
 
-    def _needs_grad(self, *wrt):
-        return any(v is not None and v.needs_grad for v in wrt)
+    def _needs_grad(self, *wrt, **kwargs):
+        force_raw_wavefield = kwargs.pop('force_raw_wavefield', False)
+        return any(v is not None and v.needs_grad for v in wrt) or force_raw_wavefield
 
     def _forward_save(self, field):
         return field.dt2

stride/physics/marmottant/devito.py

@@ -174,7 +174,10 @@ def run_forward(self, *args, **kwargs):
         """
         functions = dict()
 
+        time_bounds = kwargs.get('time_bounds', (0, self.time.extended_num))
         self.state_operator.run(dt=self.time.step,
+                                time_m=1,
+                                time_M=time_bounds[1]-1,
                                 **functions,
                                 **kwargs.pop('devito_args', {}))
 
@@ -286,7 +289,7 @@ def sub_stencil(self, **kwargs):
             rho = parent_grid.vars.rho
         except AttributeError:
             rho = self.dev_grid.vars.rho_sparse
-        vp2 = parent_grid.vars.vp2
+        vp2 = parent_grid.vars.vp**2
 
         inject_term = v_inject.inject(field=p_out.forward, expr=vp2 * self.time.step**2 * rho * inject_scale * v_inject)
 
@@ -481,6 +484,7 @@ def _make_saved_time_function(self, name, num, **kwargs):
                                            shape=(self.time.num, num),
                                            space_order=self.space_order,
                                            time_order=self.time_order,
+                                           layers=devito.NoLayers,
                                            **kwargs)
 
     def _make_interp_function(self, name, value, x_0, num, **kwargs):
@@ -519,7 +523,7 @@ def _make_interp_function(self, name, value, x_0, num, **kwargs):
         return fun, dense_fun, interp_term
 
     def _make_interp_time_function(self, name, value, x_0, num, **kwargs):
-        if not isinstance(value, devito.TimeFunction):
+        if not isinstance(value, devito.TimeFunctionOSS):
             fun = self._make_saved_time_function(name, num=num, save=self.time.num)
             fun.data[:] = value.data.T
 
@@ -541,7 +545,7 @@ def _make_interp_time_function(self, name, value, x_0, num, **kwargs):
 
             dense_fun = value
 
-            interp_term = fun.interpolate(expr=dense_fun.forward)
+            interp_term = fun.interpolate(expr=dense_fun)
 
             if x_0 is None:
                 raise ValueError('Bubble location x_0 needs to be provided when'