dense: memoize on_grid for speed

devito/types/basic.py

@@ -16,7 +16,7 @@
 from devito.parameters import configuration
 from devito.symbolics import aligned_indices
 from devito.tools import (Pickable, ctypes_to_cstr, dtype_to_cstr, dtype_to_ctype,
-                          frozendict)
+                          frozendict, memoized_meth)
 from devito.types.args import ArgProvider
 from devito.types.caching import Cached
 from devito.types.utils import DimensionTuple
@@ -736,8 +736,16 @@ def _is_on_grid(self):
         For example, if the original non-staggered function is f(x)
         then f(x) is on the grid and f(x + h_x/2) is off the grid.
         """
+        return self._check_indices(inds=self.indices)
+
+    @memoized_meth
+    def _check_indices(self, inds=None):
+        """
+        Check if the function indices are aligned with the dimensions.
+        """
+        inds = inds or self.indices
         return all([aligned_indices(i, j, d.spacing) for i, j, d in
-                    zip(self.indices, self.indices_ref, self.dimensions)])
+                    zip(inds, self.indices_ref, self.dimensions)])
 
     @property
     def evaluate(self):

examples/seismic/acoustic/acoustic_example.py

@@ -2,17 +2,17 @@
 import pytest
 
 from devito.logger import info
-from devito import Constant, Function, smooth
+from devito import Constant, Function, smooth, norm
 from examples.seismic.acoustic import AcousticWaveSolver
 from examples.seismic import demo_model, setup_geometry, seismic_args
 
 
 def acoustic_setup(shape=(50, 50, 50), spacing=(15.0, 15.0, 15.0),
                    tn=500., kernel='OT2', space_order=4, nbl=10,
-                   preset='layers-isotropic', **kwargs):
+                   preset='layers-isotropic', fs=False, **kwargs):
     model = demo_model(preset, space_order=space_order, shape=shape, nbl=nbl,
                        dtype=kwargs.pop('dtype', np.float32), spacing=spacing,
-                       **kwargs)
+                       fs=fs, **kwargs)
 
     # Source and receiver geometries
     geometry = setup_geometry(model, tn)
@@ -24,11 +24,11 @@ def acoustic_setup(shape=(50, 50, 50), spacing=(15.0, 15.0, 15.0),
 
 
 def run(shape=(50, 50, 50), spacing=(20.0, 20.0, 20.0), tn=1000.0,
-        space_order=4, kernel='OT2', nbl=40, full_run=False,
+        space_order=4, kernel='OT2', nbl=40, full_run=False, fs=False,
         autotune=False, preset='layers-isotropic', checkpointing=False, **kwargs):
 
     solver = acoustic_setup(shape=shape, spacing=spacing, nbl=nbl, tn=tn,
-                            space_order=space_order, kernel=kernel,
+                            space_order=space_order, kernel=kernel, fs=fs,
                             preset=preset, **kwargs)
 
     info("Applying Forward")
@@ -68,10 +68,15 @@ def test_isoacoustic_stability(ndim):
     shape = tuple([11]*ndim)
     spacing = tuple([20]*ndim)
     _, _, _, [rec, _] = run(shape=shape, spacing=spacing, tn=20000.0, nbl=0)
-    norm = lambda x: np.linalg.norm(x.data.reshape(-1))
     assert np.isfinite(norm(rec))
 
 
+@pytest.mark.parametrize('fs, normrec', [(True, 369.9544), (False, 459.1678)])
+def test_isoacoustic(fs, normrec):
+    _, _, _, [rec, _] = run(fs=fs)
+    assert np.isclose(norm(rec), normrec, atol=1e-3, rtol=0)
+
+
 if __name__ == "__main__":
     description = ("Example script for a set of acoustic operators.")
     parser = seismic_args(description)
@@ -86,7 +91,7 @@ def test_isoacoustic_stability(ndim):
     tn = args.tn if args.tn > 0 else (750. if ndim < 3 else 1250.)
 
     preset = 'constant-isotropic' if args.constant else 'layers-isotropic'
-    run(shape=shape, spacing=spacing, nbl=args.nbl, tn=tn,
+    run(shape=shape, spacing=spacing, nbl=args.nbl, tn=tn, fs=args.fs,
         space_order=args.space_order, preset=preset, kernel=args.kernel,
         autotune=args.autotune, opt=args.opt, full_run=args.full,
         checkpointing=args.checkpointing)

examples/seismic/acoustic/operators.py

@@ -17,7 +17,8 @@ def freesurface(model, eq):
     eq : Eq
         Time-stepping stencil (time update) to mirror at the freesurface.
     """
-    lhs, rhs = eq.evaluate
+    eq = eq.evaluate
+    lhs, rhs = eq.lhs, eq.rhs
     # Get vertical dimension and corresponding subdimension
     zfs = model.grid.subdomains['fsdomain'].dimensions[-1]
     z = zfs.parent

examples/seismic/elastic/elastic_example.py

@@ -2,6 +2,7 @@
 import pytest
 
 from devito.logger import info
+from devito import norm
 from examples.seismic.elastic import ElasticWaveSolver
 from examples.seismic import demo_model, setup_geometry, seismic_args
 
@@ -36,7 +37,6 @@ def run(shape=(50, 50), spacing=(20.0, 20.0), tn=1000.0,
 
 def test_elastic():
     _, _, _, [rec1, rec2, v, tau] = run()
-    norm = lambda x: np.linalg.norm(x.data.reshape(-1))
     assert np.isclose(norm(rec1), 19.25636, atol=1e-3, rtol=0)
     assert np.isclose(norm(rec2), 0.627606, atol=1e-3, rtol=0)
 
@@ -46,7 +46,6 @@ def test_elastic_stability(ndim):
     shape = tuple([11]*ndim)
     spacing = tuple([20]*ndim)
     _, _, _, [rec1, rec2, v, tau] = run(shape=shape, spacing=spacing, tn=20000.0, nbl=0)
-    norm = lambda x: np.linalg.norm(x.data.reshape(-1))
     assert np.isfinite(norm(rec1))
 
 

examples/seismic/model.py

@@ -76,7 +76,7 @@ def __init__(self, so):
 
     def define(self, dimensions):
         """
-        Definition of the top part of the domain for wrapped indices FS
+        Definition of the top part of the domain for wrapped indices FS.
         """
 
         return {d: (d if not d == dimensions[-1] else ('left', self.size))

examples/seismic/viscoacoustic/viscoacoustic_example.py

@@ -1,5 +1,6 @@
 import numpy as np
 from devito.logger import info
+from devito import norm
 from examples.seismic.viscoacoustic import ViscoacousticWaveSolver
 from examples.seismic import demo_model, setup_geometry, seismic_args
 
@@ -37,7 +38,6 @@ def run(shape=(50, 50), spacing=(20.0, 20.0), tn=1000.0,
 
 def test_viscoacoustic():
     _, _, _, [rec] = run()
-    norm = lambda x: np.linalg.norm(x.data.reshape(-1))
     assert np.isclose(norm(rec), 18.7749, atol=1e-3, rtol=0)
 
 

examples/seismic/viscoelastic/viscoelastic_example.py

@@ -2,6 +2,7 @@
 import pytest
 
 from devito.logger import info
+from devito import norm
 from examples.seismic.viscoelastic import ViscoelasticWaveSolver
 from examples.seismic import demo_model, setup_geometry, seismic_args
 
@@ -36,7 +37,6 @@ def run(shape=(50, 50), spacing=(20.0, 20.0), tn=1000.0,
 
 def test_viscoelastic():
     _, _, _, [rec1, rec2, v, tau] = run()
-    norm = lambda x: np.linalg.norm(x.data.reshape(-1))
     assert np.isclose(norm(rec1), 12.28040, atol=1e-3, rtol=0)
     assert np.isclose(norm(rec2), 0.312461, atol=1e-3, rtol=0)
 
@@ -46,7 +46,6 @@ def test_viscoelastic_stability(ndim):
     shape = tuple([11]*ndim)
     spacing = tuple([20]*ndim)
     _, _, _, [rec1, rec2, v, tau] = run(shape=shape, spacing=spacing, tn=20000.0, nbl=0)
-    norm = lambda x: np.linalg.norm(x.data.reshape(-1))
     assert np.isfinite(norm(rec1))