Merge pull request #2342 from devitocodes/sinc-interp-final

api: Add Hicks (sinc) interpolation api

.github/workflows/pytest-core-mpi.yml

@@ -48,6 +48,12 @@ jobs:
         python3 scripts/clear_devito_cache.py
         python3 -m pytest --cov --cov-config=.coveragerc --cov-report=xml -m parallel tests/
 
+    - name: Test examples with MPI
+      run: |
+        python3 scripts/clear_devito_cache.py
+        DEVITO_MPI=1 mpirun -n 2 python3 -m pytest --cov --cov-config=.coveragerc --cov-report=xml examples/seismic/acoustic
+        DEVITO_MPI=1 mpirun -n 2 python3 -m pytest --cov --cov-config=.coveragerc --cov-report=xml examples/seismic/tti
+
     - name: Upload coverage to Codecov
       uses: codecov/codecov-action@v4
       with:
@@ -64,9 +70,12 @@ jobs:
             - name: gcc
               arch: gcc
               os: ubuntu-latest
+              mpiflag: ""
             - name: icx
               arch: icx
               os: ubuntu-latest
+              # Need safe math for icx due to inaccuracy with mpi+sinc interpolation
+              mpiflag: "-e DEVITO_SAFE_MATH=1"
 
       steps:
       - name: Checkout devito
@@ -78,4 +87,9 @@ jobs:
       
       - name: Test with pytest
         run: |
-          docker run --rm -e CODECOV_TOKEN=${{ secrets.CODECOV_TOKEN }} -e OMP_NUM_THREADS=1 --name testrun devito_img pytest tests/test_mpi.py
+          docker run --rm -e CODECOV_TOKEN=${{ secrets.CODECOV_TOKEN }} -e OMP_NUM_THREADS=1 --name testrun devito_img pytest tests/test_mpi.py
+
+      - name: Test examples with MPI
+        run: |
+          docker run --rm ${{ matrix.mpiflag }} -e DEVITO_MPI=1 -e OMP_NUM_THREADS=1 --name examplerun devito_img mpiexec -n 2 pytest examples/seismic/acoustic
+          docker run --rm -e DEVITO_MPI=1 -e OMP_NUM_THREADS=1 --name examplerun devito_img mpiexec -n 2 pytest examples/seismic/tti

conftest.py

@@ -34,16 +34,12 @@ def skipif(items, whole_module=False):
     accepted.update({'device', 'device-C', 'device-openmp', 'device-openacc',
                      'device-aomp', 'cpu64-icc', 'cpu64-icx', 'cpu64-nvc', 'cpu64-arm',
                      'cpu64-icpx', 'chkpnt'})
-    accepted.update({'nompi', 'nodevice'})
+    accepted.update({'nodevice'})
     unknown = sorted(set(items) - accepted)
     if unknown:
         raise ValueError("Illegal skipif argument(s) `%s`" % unknown)
     skipit = False
     for i in items:
-        # Skip if no MPI
-        if i == 'nompi' and MPI is None:
-            skipit = "mpi4py/MPI not installed"
-            break
         # Skip if won't run on GPUs
         if i == 'device' and isinstance(configuration['platform'], Device):
             skipit = "device `%s` unsupported" % configuration['platform'].name
@@ -171,6 +167,9 @@ def parallel(item):
         os.environ['DEVITO_MPI'] = scheme
         try:
             check_call(call)
+            return True
+        except:
+            return False
         finally:
             os.environ['DEVITO_MPI'] = '0'
 
@@ -189,18 +188,21 @@ def pytest_runtest_setup(item):
         partest = int(partest)
     except ValueError:
         pass
-    if item.get_closest_marker("parallel") and not partest:
-        # Blow away function arg in "master" process, to ensure
-        # this test isn't run on only one process
-        dummy_test = lambda *args, **kwargs: True
-        # For pytest <7
-        if item.cls is not None:
-            attr = item.originalname or item.name
-            setattr(item.cls, attr, dummy_test)
+    if item.get_closest_marker("parallel"):
+        if MPI is None:
+            pytest.skip("mpi4py/MPI not installed")
         else:
-            item.obj = dummy_test
-        # For pytest >= 7
-        setattr(item, '_obj', dummy_test)
+            # Blow away function arg in "master" process, to ensure
+            # this test isn't run on only one process
+            dummy_test = lambda *args, **kwargs: True
+            # For pytest <7
+            if item.cls is not None:
+                attr = item.originalname or item.name
+                setattr(item.cls, attr, dummy_test)
+            else:
+                item.obj = dummy_test
+            # For pytest >= 7
+            setattr(item, '_obj', dummy_test)
 
 
 def pytest_runtest_call(item):
@@ -211,7 +213,27 @@ def pytest_runtest_call(item):
         pass
     if item.get_closest_marker("parallel") and not partest:
         # Spawn parallel processes to run test
-        parallel(item)
+        passed = parallel(item)
+        if not passed:
+            pytest.fail(f"{item} failed in parallel execution")
+        else:
+            pytest.skip(f"{item}t passed in parallel execution")
+
+
+@pytest.hookimpl(tryfirst=True, hookwrapper=True)
+def pytest_runtest_makereport(item, call):
+    outcome = yield
+    result = outcome.get_result()
+
+    partest = os.environ.get('DEVITO_MPI', 0)
+    try:
+        partest = int(partest)
+    except ValueError:
+        pass
+
+    if item.get_closest_marker("parallel") and not partest:
+        if call.when == 'call' and result.outcome == 'skipped':
+            result.outcome = 'passed'
 
 
 # A list of optimization options/pipelines to be used in testing

devito/operations/interpolators.py

@@ -2,17 +2,24 @@
 from functools import wraps
 
 import sympy
+import numpy as np
 from cached_property import cached_property
 
+try:
+    from scipy.special import i0
+except ImportError:
+    from numpy import i0
+
 from devito.finite_differences.differentiable import Mul
 from devito.finite_differences.elementary import floor
+from devito.logger import warning
 from devito.symbolics import retrieve_function_carriers, retrieve_functions, INT
 from devito.tools import as_tuple, flatten, filter_ordered
 from devito.types import (ConditionalDimension, Eq, Inc, Evaluable, Symbol,
-                          CustomDimension)
+                          CustomDimension, SubFunction)
 from devito.types.utils import DimensionTuple
 
-__all__ = ['LinearInterpolator', 'PrecomputedInterpolator']
+__all__ = ['LinearInterpolator', 'PrecomputedInterpolator', 'SincInterpolator']
 
 
 def check_radius(func):
@@ -22,7 +29,7 @@ def wrapper(interp, *args, **kwargs):
         funcs = set().union(*[retrieve_functions(a) for a in args])
         so = min({f.space_order for f in funcs if not f.is_SparseFunction} or {r})
         if so < r:
-            raise ValueError("Space order %d smaller than interpolation r %d" % (so, r))
+            raise ValueError("Space order %d too small for interpolation r %d" % (so, r))
         return func(interp, *args, **kwargs)
     return wrapper
 
@@ -125,6 +132,8 @@ class GenericInterpolator(ABC):
     Abstract base class defining the interface for an interpolator.
     """
 
+    _name = "generic"
+
     @abstractmethod
     def inject(self, *args, **kwargs):
         pass
@@ -133,6 +142,13 @@ def inject(self, *args, **kwargs):
     def interpolate(self, *args, **kwargs):
         pass
 
+    @property
+    def name(self):
+        return self._name
+
+    def _arg_defaults(self, **args):
+        return {}
+
 
 class WeightedInterpolator(GenericInterpolator):
 
@@ -142,6 +158,8 @@ class WeightedInterpolator(GenericInterpolator):
     and multiplied at a given point: `w[x, y] = wx[x] * wy[y]`
     """
 
+    _name = 'weighted'
+
     def __init__(self, sfunction):
         self.sfunction = sfunction
 
@@ -204,7 +222,7 @@ def _positions(self, implicit_dims):
         return [Eq(v, INT(floor(k)), implicit_dims=implicit_dims)
                 for k, v in self.sfunction._position_map.items()]
 
-    def _interp_idx(self, variables, implicit_dims=None):
+    def _interp_idx(self, variables, implicit_dims=None, pos_only=()):
         """
         Generate interpolation indices for the DiscreteFunctions in ``variables``.
         """
@@ -223,6 +241,12 @@ def _interp_idx(self, variables, implicit_dims=None):
                     for ((k, c), p) in zip(mapper.items(), pos)})
                     for v in variables}
 
+        # Position only replacement, not radius dependent.
+        # E.g src.inject(vp(x)*src) needs to use vp[posx] at all points
+        # not vp[posx + rx]
+        idx_subs.update({v: v.subs({k: p for (k, p) in zip(mapper, pos)})
+                         for v in pos_only})
+
         return idx_subs, temps
 
     @check_radius
@@ -347,10 +371,9 @@ def _inject(self, field, expr, implicit_dims=None):
         # summing temp that wouldn't allow collapsing
         implicit_dims = implicit_dims + tuple(r.parent for r in self._rdim)
 
-        variables = variables + list(fields)
-
         # List of indirection indices for all adjacent grid points
-        idx_subs, temps = self._interp_idx(variables, implicit_dims=implicit_dims)
+        idx_subs, temps = self._interp_idx(fields, implicit_dims=implicit_dims,
+                                           pos_only=variables)
 
         # Substitute coordinate base symbols into the interpolation coefficients
         eqns = [Inc(_field.xreplace(idx_subs),
@@ -370,6 +393,9 @@ class LinearInterpolator(WeightedInterpolator):
     ----------
     sfunction: The SparseFunction that this Interpolator operates on.
     """
+
+    _name = 'linear'
+
     @property
     def _weights(self):
         c = [(1 - p) * (1 - r) + p * r
@@ -403,11 +429,13 @@ class PrecomputedInterpolator(WeightedInterpolator):
     sfunction: The SparseFunction that this Interpolator operates on.
     """
 
+    _name = 'precomp'
+
     def _positions(self, implicit_dims):
         if self.sfunction.gridpoints is None:
             return super()._positions(implicit_dims)
         # No position temp as we have directly the gridpoints
-        return [Eq(p, k, implicit_dims=implicit_dims)
+        return [Eq(p, floor(k), implicit_dims=implicit_dims)
                 for (k, p) in self.sfunction._position_map.items()]
 
     @property
@@ -421,3 +449,70 @@ def _weights(self):
                    for (ri, rd) in enumerate(self._rdim)]
         return Mul(*[self.interpolation_coeffs.subs(mapper)
                      for mapper in mappers])
+
+
+class SincInterpolator(PrecomputedInterpolator):
+    """
+    Hicks windowed sinc interpolation scheme.
+
+    Arbitrary source and receiver positioning in finite‐difference schemes
+    using Kaiser windowed sinc functions
+
+    https://library.seg.org/doi/10.1190/1.1451454
+
+    """
+
+    _name = 'sinc'
+
+    # Table 1
+    _b_table = {2: 2.94, 3: 4.53,
+                4: 4.14, 5: 5.26, 6: 6.40,
+                7: 7.51, 8: 8.56, 9: 9.56, 10: 10.64}
+
+    def __init__(self, sfunction):
+        if i0 is np.i0:
+            warning("""
+Using `numpy.i0`. We (and numpy) recommend to install scipy to improve the performance
+of the SincInterpolator that uses i0 (Bessel function).
+""")
+        super().__init__(sfunction)
+
+    @cached_property
+    def interpolation_coeffs(self):
+        coeffs = {}
+        shape = (self.sfunction.npoint, 2 * self.r)
+        for r in self._rdim:
+            dimensions = (self.sfunction._sparse_dim, r.parent)
+            sf = SubFunction(name="wsinc%s" % r.name, dtype=self.sfunction.dtype,
+                             shape=shape, dimensions=dimensions,
+                             space_order=0, alias=self.sfunction.alias,
+                             distributor=self.sfunction._distributor,
+                             parent=self.sfunction)
+            coeffs[r] = sf
+        return coeffs
+
+    @property
+    def _weights(self):
+        return Mul(*[w._subs(rd, rd-rd.parent.symbolic_min)
+                     for (rd, w) in self.interpolation_coeffs.items()])
+
+    def _arg_defaults(self, coords=None, sfunc=None):
+        args = {}
+        b = self._b_table[self.r]
+        b0 = i0(b)
+        if coords is None or sfunc is None:
+            raise ValueError("No coordinates or sparse function provided")
+        # Coords to indices
+        coords = coords / np.array(sfunc.grid.spacing)
+        coords = coords - np.floor(coords)
+
+        # Precompute sinc
+        for j, r in enumerate(self._rdim):
+            data = np.zeros((coords.shape[0], 2*self.r), dtype=sfunc.dtype)
+            for ri in range(2*self.r):
+                rpos = ri - self.r + 1 - coords[:, j]
+                num = i0(b*np.sqrt(1 - (rpos/self.r)**2))
+                data[:, ri] = num / b0 * np.sinc(rpos)
+            args[self.interpolation_coeffs[r].name] = data
+
+        return args