Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

api: Add Hicks (sinc) interpolation api #2342

Merged
merged 8 commits into from
Apr 3, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
16 changes: 15 additions & 1 deletion .github/workflows/pytest-core-mpi.yml
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand All @@ -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
Expand All @@ -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
56 changes: 39 additions & 17 deletions conftest.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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'

Expand All @@ -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):
Expand All @@ -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
Expand Down
111 changes: 103 additions & 8 deletions devito/operations/interpolators.py
Original file line number Diff line number Diff line change
Expand Up @@ -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):
Expand All @@ -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

Expand Down Expand Up @@ -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
Expand All @@ -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):
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

if we want this method here it would be cleaner to have this class inhereting from ArgProvider.

Perhaps we move it into a subclass?

Copy link
Contributor Author

@mloubout mloubout Apr 3, 2024

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

No, the interpolator isn't an arg provide, its sparse function is and is calling it only.

return {}


class WeightedInterpolator(GenericInterpolator):

Expand All @@ -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

Expand Down Expand Up @@ -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``.
"""
Expand All @@ -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
Expand Down Expand Up @@ -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),
Expand All @@ -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
Expand Down Expand Up @@ -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
Expand All @@ -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
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

What's Table 2? u mean Table 1 in the paper?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yes Table 1 in the paper

_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]
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Does this need a catch for r > 10?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

It will just throw a key-error but can add a check (I doubt anyone will run it this high that's a 20x20x20 cube in 3d super expensive)

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

agree to add a check. Users can be surprising...

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
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think this wants a NOTE: <comment about performance> ? if you measured it's slowish? or is it super quick since the loops are very short?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I haven't seen any drastic slow compute since it's numpified so not sure want to force user to be like "hey here is a warning should I nitpick here"

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
Loading
Loading