Merge pull request #2456 from devitocodes/avg-mode

API: Support harmonic averaging for parameter Function

devito/__init__.py

@@ -1,7 +1,6 @@
 import atexit
 from itertools import product
 import os
-from . import _version
 
 import numpy as np
 
@@ -195,9 +194,5 @@ def mode_performance():
 # job properly and thus we may end up missing some custom __del__'s
 atexit.register(clear_cache)
 
-
-__version__ = _version.get_versions()['version']
-
-
 # Clean up namespace
 del atexit, product

devito/finite_differences/differentiable.py

@@ -1033,7 +1033,7 @@ def _(expr, x0, **kwargs):
 def _(expr, x0, **kwargs):
     from devito.finite_differences.derivative import Derivative
     x0_expr = {d: v for d, v in x0.items() if v is not expr.indices_ref[d]}
-    if x0_expr and not expr.is_parameter:
+    if x0_expr:
         dims = tuple((d, 0) for d in x0_expr)
         fd_o = tuple([2]*len(dims))
         return Derivative(expr, *dims, fd_order=fd_o, x0=x0_expr)._evaluate(**kwargs)

devito/types/basic.py

@@ -852,7 +852,7 @@ class AbstractFunction(sympy.Function, Basic, Pickable, Evaluable):
     """
 
     __rkwargs__ = ('name', 'dtype', 'grid', 'halo', 'padding', 'ghost',
-                   'alias', 'space', 'function', 'is_transient')
+                   'alias', 'space', 'function', 'is_transient', 'avg_mode')
 
     def __new__(cls, *args, **kwargs):
         # Preprocess arguments
@@ -983,6 +983,12 @@ def __init_finalize__(self, *args, **kwargs):
         # certain optimizations, such as avoiding memory copies
         self._is_transient = kwargs.get('is_transient', False)
 
+        # Averaging mode for off the grid evaluation
+        self._avg_mode = kwargs.get('avg_mode', 'arithmetic')
+        if self._avg_mode not in ['arithmetic', 'harmonic']:
+            raise ValueError("Invalid averaging mode_mode %s, accepted values are"
+                             " arithmetic or harmonic" % self._avg_mode)
+
     @classmethod
     def __args_setup__(cls, *args, **kwargs):
         """
@@ -1147,13 +1153,19 @@ def _evaluate(self, **kwargs):
             return self
 
         # Base function
-        retval = self.function
+        if self._avg_mode == 'harmonic':
+            retval = 1 / self.function
+        else:
+            retval = self.function
         # Apply interpolation from inner most dim
         for d, i in self._grid_map.items():
-            retval = retval.diff(d, 0, fd_order=2, x0={d: i})
+            retval = retval.diff(d, deriv_order=0, fd_order=2, x0={d: i})
+        if self._avg_mode == 'harmonic':
+            retval = 1 / retval
+
         # Evaluate. Since we used `self.function` it will be on the grid when evaluate
         # is called again within FD
-        return retval.evaluate
+        return retval.evaluate.expand()
 
     @property
     def shape(self):
@@ -1255,6 +1267,10 @@ def is_const(self):
     def is_transient(self):
         return self._is_transient
 
+    @property
+    def avg_mode(self):
+        return self._avg_mode
+
     @property
     def alias(self):
         return self._alias

examples/seismic/acoustic/acoustic_example.py

@@ -1,5 +1,8 @@
 import numpy as np
-import pytest
+try:
+    import pytest
+except ImportError:
+    pass
 
 from devito.logger import info
 from devito import Constant, Function, smooth, norm

examples/seismic/elastic/elastic_example.py

@@ -1,5 +1,8 @@
 import numpy as np
-import pytest
+try:
+    import pytest
+except ImportError:
+    pass
 from devito import norm
 from devito.logger import info
 from examples.seismic.elastic import ElasticWaveSolver
@@ -36,8 +39,8 @@ def run(shape=(50, 50), spacing=(20.0, 20.0), tn=1000.0,
 @pytest.mark.parametrize("dtype", [np.float32, np.float64])
 def test_elastic(dtype):
     _, _, _, [rec1, rec2, v, tau] = run(dtype=dtype)
-    assert np.isclose(norm(rec1), 19.25636, atol=1e-3, rtol=0)
-    assert np.isclose(norm(rec2), 0.6276, atol=1e-3, rtol=0)
+    assert np.isclose(norm(rec1), 19.43005, atol=1e-3, rtol=0)
+    assert np.isclose(norm(rec2), 0.6345, atol=1e-3, rtol=0)
 
 
 @pytest.mark.parametrize('shape', [(51, 51), (16, 16, 16)])

examples/seismic/model.py

@@ -1,6 +1,9 @@
 import numpy as np
 from sympy import finite_diff_weights as fd_w
-import pytest
+try:
+    import pytest
+except:
+    pass
 
 from devito import (Grid, SubDomain, Function, Constant, warning,
                     SubDimension, Eq, Inc, Operator, div, sin, Abs)
@@ -167,12 +170,12 @@ def physical_params(self, **kwargs):
         return {i.name: kwargs.get(i.name, i) or i for i in known}
 
     def _gen_phys_param(self, field, name, space_order, is_param=True,
-                        default_value=0, **kwargs):
+                        default_value=0, avg_mode='arithmetic', **kwargs):
         if field is None:
             return default_value
         if isinstance(field, np.ndarray):
             function = Function(name=name, grid=self.grid, space_order=space_order,
-                                parameter=is_param)
+                                parameter=is_param, avg_mode=avg_mode)
             initialize_function(function, field, self.padsizes)
         else:
             function = Constant(name=name, value=field, dtype=self.grid.dtype)
@@ -304,7 +307,11 @@ def _initialize_physics(self, vp, space_order, **kwargs):
             vs = kwargs.pop('vs')
             self.lam = self._gen_phys_param((vp**2 - 2. * vs**2)/b, 'lam', space_order,
                                             is_param=True)
-            self.mu = self._gen_phys_param(vs**2 / b, 'mu', space_order, is_param=True)
+            # Need to add small value to avoid division by zero
+            if isinstance(vs, np.ndarray):
+                vs = vs + 1e-12
+            self.mu = self._gen_phys_param(vs**2 / b, 'mu', space_order, is_param=True,
+                                           avg_mode='harmonic')
         else:
             # All other seismic models have at least a velocity
             self.vp = self._gen_phys_param(vp, 'vp', space_order)

examples/seismic/self_adjoint/test_utils.py

@@ -1,5 +1,8 @@
 import numpy as np
-import pytest
+try:
+    import pytest
+except:
+    pass
 from devito import Grid, Function
 from examples.seismic.self_adjoint import setup_w_over_q
 

examples/seismic/self_adjoint/test_wavesolver_iso.py

@@ -1,6 +1,9 @@
 from scipy.special import hankel2
 import numpy as np
-import pytest
+try:
+    import pytest
+except:
+    pass
 from devito import Grid, Function, Eq, Operator, info
 from examples.seismic import RickerSource, TimeAxis, Model, AcquisitionGeometry
 from examples.seismic.self_adjoint import (acoustic_sa_setup, setup_w_over_q,

examples/seismic/test_seismic_utils.py

@@ -1,4 +1,7 @@
-import pytest
+try:
+    import pytest
+except:
+    pass
 import numpy as np
 
 from devito import norm

examples/seismic/tti/tti_example.py

@@ -1,5 +1,8 @@
 import numpy as np
-import pytest
+try:
+    import pytest
+except ImportError:
+    pass
 
 from devito import Function, smooth, norm, info, Constant
 

examples/seismic/viscoacoustic/viscoacoustic_example.py

@@ -1,5 +1,8 @@
 import numpy as np
-import pytest
+try:
+    import pytest
+except ImportError:
+    pass
 
 from devito.logger import info
 from devito import norm

examples/seismic/viscoelastic/viscoelastic_example.py

@@ -1,5 +1,8 @@
 import numpy as np
-import pytest
+try:
+    import pytest
+except ImportError:
+    pass
 
 from devito import norm
 from devito.logger import info
@@ -38,7 +41,7 @@ def run(shape=(50, 50), spacing=(20.0, 20.0), tn=1000.0,
 @pytest.mark.parametrize("dtype", [np.float32, np.float64])
 def test_viscoelastic(dtype):
     _, _, _, [rec1, rec2, v, tau] = run(dtype=dtype)
-    assert np.isclose(norm(rec1), 12.28040, atol=1e-3, rtol=0)
+    assert np.isclose(norm(rec1), 12.30114, atol=1e-3, rtol=0)
     assert np.isclose(norm(rec2), 0.312462, atol=1e-3, rtol=0)
 
 

examples/userapi/03_subdomains.ipynb

@@ -941,7 +941,7 @@
    "outputs": [],
    "source": [
     "from devito import norm\n",
-    "assert np.isclose(norm(v[0]), 0.10312, rtol=1e-4)"
+    "assert np.isclose(norm(v[0]), 0.10301, rtol=1e-4)"
    ]
   },
   {

tests/test_differentiable.py

@@ -1,4 +1,8 @@
+from itertools import product
+
 import sympy
+import pytest
+
 from devito import Function, Grid, Differentiable, NODE
 from devito.finite_differences.differentiable import Add, Mul, Pow, diffify, interp_for_fd
 
@@ -82,3 +86,28 @@ def test_interp():
                    sa + interp_for_fd(a, {x: x + x.spacing/2}, expand=True))
     assert sp_diff(interp_for_fd(a + sa, {x: x}, expand=True),
                    a + interp_for_fd(sa, {x: x}, expand=True))
+
+
+@pytest.mark.parametrize('ndim', [1, 2, 3])
+def test_avg_mode(ndim):
+    grid = Grid([11]*ndim)
+    v = Function(name='v', grid=grid, staggered=grid.dimensions)
+    a0 = Function(name="a0", grid=grid)
+    a = Function(name="a", grid=grid, parameter=True)
+    b = Function(name="b", grid=grid, parameter=True, avg_mode='harmonic')
+
+    a0_avg = a0._eval_at(v)
+    a_avg = a._eval_at(v).evaluate
+    b_avg = b._eval_at(v).evaluate
+
+    assert a0_avg == a0
+
+    # Indices around the point at the center of a cell
+    all_shift = tuple(product(*[[0, 1] for _ in range(ndim)]))
+    args = [{d: d + i * d.spacing for d, i in zip(grid.dimensions, s)} for s in all_shift]
+
+    # Default is arithmetic average
+    assert sympy.simplify(a_avg - 0.5**ndim * sum(a.subs(arg) for arg in args)) == 0
+
+    # Harmonic average, h(a[.5]) = 1/(.5/a[0] + .5/a[1])
+    assert sympy.simplify(b_avg - 1/(0.5**ndim * sum(1/b.subs(arg) for arg in args))) == 0