From 2705f2683683bd0eb0a1327828721e0184f10520 Mon Sep 17 00:00:00 2001
From: mloubout <mathias.louboutin@gmail.com>
Date: Wed, 3 Apr 2024 09:15:43 -0400
Subject: [PATCH] api: switch to recommended scipy i0 when available

---
 .github/workflows/pytest-core-mpi.yml         |  5 +++++
 devito/operations/interpolators.py            | 18 ++++++++++++++++--
 devito/types/sparse.py                        | 12 ++++++------
 examples/seismic/acoustic/acoustic_example.py |  2 +-
 examples/seismic/utils.py                     |  2 ++
 5 files changed, 30 insertions(+), 9 deletions(-)

diff --git a/.github/workflows/pytest-core-mpi.yml b/.github/workflows/pytest-core-mpi.yml
index a3963fe866e..250e48d9999 100644
--- a/.github/workflows/pytest-core-mpi.yml
+++ b/.github/workflows/pytest-core-mpi.yml
@@ -85,3 +85,8 @@ 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
+
+      - name: Test examples with MPI
+        run: |
+          docker run --rm -e DEVITO_SAFE_MATH=1 -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
\ No newline at end of file
diff --git a/devito/operations/interpolators.py b/devito/operations/interpolators.py
index 38a3571b17f..3469e9539cf 100644
--- a/devito/operations/interpolators.py
+++ b/devito/operations/interpolators.py
@@ -5,8 +5,14 @@
 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,
@@ -465,6 +471,14 @@ class SincInterpolator(PrecomputedInterpolator):
                 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 = {}
@@ -487,7 +501,7 @@ def _weights(self):
     def _arg_defaults(self, coords=None, sfunc=None):
         args = {}
         b = self._b_table[self.r]
-        b0 = np.i0(b)
+        b0 = i0(b)
         if coords is None or sfunc is None:
             raise ValueError("No coordinates or sparse function provided")
         # Coords to indices
@@ -499,7 +513,7 @@ def _arg_defaults(self, coords=None, sfunc=None):
             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 = np.i0(b*np.sqrt(1 - (rpos/self.r)**2))
+                num = i0(b*np.sqrt(1 - (rpos/self.r)**2))
                 data[:, ri] = num / b0 * np.sinc(rpos)
             args[self.interpolation_coeffs[r].name] = data
 
diff --git a/devito/types/sparse.py b/devito/types/sparse.py
index a5f711f7cf2..85910ff9550 100644
--- a/devito/types/sparse.py
+++ b/devito/types/sparse.py
@@ -822,16 +822,16 @@ def __init_finalize__(self, *args, **kwargs):
         self._coordinates = self.__subfunc_setup__(coordinates, 'coords')
         self._dist_origin = {self._coordinates: self.grid.origin_offset}
 
-    def __interp_setup__(self, interp='linear', r=None, **kwargs):
-        self.interpolation = interp
-        self.interpolator = _interpolators[interp](self)
-        self._radius = r or _default_radius[interp]
-        if interp == 'sinc':
+    def __interp_setup__(self, interpolation='linear', r=None, **kwargs):
+        self.interpolation = interpolation
+        self.interpolator = _interpolators[interpolation](self)
+        self._radius = r or _default_radius[interpolation]
+        if interpolation == 'sinc':
             if self._radius < 2:
                 raise ValueError("'sinc' interpolator requires a radius of at least 2")
             elif self._radius > 10:
                 raise ValueError("'sinc' interpolator requires a radius of at most 10")
-        elif interp == 'linear' and self._radius != 1:
+        elif interpolation == 'linear' and self._radius != 1:
             self._radius = 1
 
     @cached_property
diff --git a/examples/seismic/acoustic/acoustic_example.py b/examples/seismic/acoustic/acoustic_example.py
index aced327cdca..bbcc4a7c41d 100644
--- a/examples/seismic/acoustic/acoustic_example.py
+++ b/examples/seismic/acoustic/acoustic_example.py
@@ -103,4 +103,4 @@ def test_isoacoustic(fs, normrec, dtype, interp):
     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, dtype=args.dtype)
+        checkpointing=args.checkpointing, dtype=args.dtype, interpolation=args.interp)
diff --git a/examples/seismic/utils.py b/examples/seismic/utils.py
index 42ca0ced721..835961fa028 100644
--- a/examples/seismic/utils.py
+++ b/examples/seismic/utils.py
@@ -253,4 +253,6 @@ def __call__(self, parser, args, values, option_string=None):
                         type=float, help="Simulation time in millisecond")
     parser.add_argument("-dtype", action=_dtype_store, dest="dtype", default=np.float32,
                         choices=['float32', 'float64'])
+    parser.add_argument("-interp", dest="interp", default="linear",
+                        choices=['linear', 'sinc'])
     return parser