diff --git a/devito/finite_differences/derivative.py b/devito/finite_differences/derivative.py index 22984af6fb..a16843d858 100644 --- a/devito/finite_differences/derivative.py +++ b/devito/finite_differences/derivative.py @@ -489,6 +489,9 @@ def _eval_at(self, func): # compare staggering if self.expr.staggered == func.staggered and self.expr.is_Function: return self + # Time derivatives are not affected by space staggering + if all(d.is_Time for d in self.dims): + return self # Check if x0's keys come from a DerivedDimension x0 = func.indices_ref.getters diff --git a/devito/operations/interpolators.py b/devito/operations/interpolators.py index defc1b5656..6467e65df4 100644 --- a/devito/operations/interpolators.py +++ b/devito/operations/interpolators.py @@ -1,5 +1,7 @@ from abc import ABC, abstractmethod +from contextlib import suppress from functools import cached_property, wraps +from itertools import groupby import numpy as np import sympy @@ -67,12 +69,9 @@ def _extract_subdomain(variables): """ sdms = set() for v in variables: - try: + with suppress(AttributeError): if v.grid.is_SubDomain: sdms.add(v.grid) - except AttributeError: - # Variable not on a grid (Indexed for example) - pass if len(sdms) > 1: raise NotImplementedError("Sparse operation on multiple Functions defined on" @@ -230,7 +229,7 @@ def r(self): return self.sfunction.r @memoized_meth - def _weights(self, subdomain=None): + def _weights(self, subdomain=None, shifts=None): raise NotImplementedError @property @@ -243,8 +242,22 @@ def _cdim(self): dims = [self.sfunction._crdim(d) for d in self._gdims] return dims + def _field_shifts(self, field): + """ + Per-grid-Dimension half-cell shift induced by `field`'s staggering + (e.g. `h_x/2` for a field staggered in `x`). Returns None for + unstaggered fields. SubDomain-induced origin offsets are deliberately + ignored — they are not staggering. + """ + staggered = field.staggered + if not staggered or staggered.on_node: + return () + return tuple((d.spacing / 2) if s else 0 + for d, s in zip(field.dimensions, staggered, strict=True) + if d.is_Space) + @memoized_meth - def _rdim(self, subdomain=None): + def _rdim(self, subdomain=None, shifts=None): # If the interpolation operation is limited to a SubDomain, # use the SubDimensions of that SubDomain if subdomain: @@ -254,7 +267,7 @@ def _rdim(self, subdomain=None): # Make radius dimension conditional to avoid OOB rdims = [] - pos = self.sfunction._position_map.values() + pos = self.sfunction._position_map(shifts=shifts).values() for (d, rd, p) in zip(gdims, self._cdim, pos, strict=True): # Add conditional to avoid OOB @@ -279,12 +292,10 @@ def _augment_implicit_dims(self, implicit_dims, extras=None): # dimensions of that SubDomain from any extra dimensions found edims = [] for v in extras: - try: + with suppress(AttributeError): if v.grid.is_SubDomain: edims.extend([d for d in v.grid.dimensions if d.is_Sub and d.root in self._gdims]) - except AttributeError: - pass gdims = filter_ordered(edims + list(self._gdims)) extra = filter_ordered([i for v in extras for i in v.dimensions @@ -300,27 +311,34 @@ def _augment_implicit_dims(self, implicit_dims, extras=None): idims = extra + as_tuple(implicit_dims) + self.sfunction.dimensions return tuple(idims) - def _coeff_temps(self, implicit_dims): + def _coeff_temps(self, implicit_dims, shifts=None): return [] - def _positions(self, implicit_dims): + def _positions(self, implicit_dims, shifts=None): return [Eq(v, INT(floor(k)), implicit_dims=implicit_dims) - for k, v in self.sfunction._position_map.items()] + for k, v in self.sfunction._position_map(shifts=shifts).items()] - def _interp_idx(self, variables, implicit_dims=None, subdomain=None): + def _interp_idx(self, variables, implicit_dims=None, subdomain=None, + shifts=None): """ - Generate interpolation indices for the DiscreteFunctions in ``variables``. + Generate interpolation indices for the DiscreteFunctions in `variables`. + + `shifts` is a per-Dimension physical offset for the target field's + origin: it only affects the integer position symbol via the position + map (`pos = floor((c - o - shift)/h)`). The index substitution itself + is unchanged — any staggered offset in a field's own symbolic access is + absorbed by Devito's normal indexing. """ - pos = self.sfunction._position_map.values() + pos = self.sfunction._position_map(shifts=shifts).values() # Temporaries for the position - temps = self._positions(implicit_dims) + temps = self._positions(implicit_dims, shifts=shifts) # Coefficient symbol expression - temps.extend(self._coeff_temps(implicit_dims)) + temps.extend(self._coeff_temps(implicit_dims, shifts=shifts)) # Substitution mapper for variables - mapper = self._rdim(subdomain=subdomain).getters + mapper = self._rdim(subdomain=subdomain, shifts=shifts).getters # Index substitution to make in variables subs = { @@ -337,7 +355,7 @@ def _interp_idx(self, variables, implicit_dims=None, subdomain=None): @check_coords def interpolate(self, expr, increment=False, self_subs=None, implicit_dims=None): """ - Generate equations interpolating an arbitrary expression into ``self``. + Generate equations interpolating an arbitrary expression into `self`. Parameters ---------- @@ -375,7 +393,7 @@ def inject(self, field, expr, implicit_dims=None): def _interpolate(self, expr, increment=False, self_subs=None, implicit_dims=None): """ - Generate equations interpolating an arbitrary expression into ``self``. + Generate equations interpolating an arbitrary expression into `self`. Parameters ---------- @@ -389,16 +407,13 @@ def _interpolate(self, expr, increment=False, self_subs=None, implicit_dims=None the operator. """ # Derivatives must be evaluated before the introduction of indirect accesses - try: - _expr = expr.evaluate - except AttributeError: - # E.g., a generic SymPy expression or a number - _expr = expr + with suppress(AttributeError): + expr = expr._eval_at(self.sfunction).evaluate if self_subs is None: self_subs = {} - variables = list(retrieve_function_carriers(_expr)) + variables = list(retrieve_function_carriers(expr)) subdomain = _extract_subdomain(variables) # Implicit dimensions @@ -413,7 +428,7 @@ def _interpolate(self, expr, increment=False, self_subs=None, implicit_dims=None summands = [Eq(rhs, 0., implicit_dims=implicit_dims)] # Substitute coordinate base symbols into the interpolation coefficients weights = self._weights(subdomain=subdomain) - summands.extend([Inc(rhs, (weights * _expr).xreplace(idx_subs), + summands.extend([Inc(rhs, (weights * expr).xreplace(idx_subs), implicit_dims=implicit_dims)]) # Write/Incr `self` @@ -451,35 +466,48 @@ def _inject(self, field, expr, implicit_dims=None): subdomain = _extract_subdomain(fields) - # Derivatives must be evaluated before the introduction of indirect accesses - try: - _exprs = tuple(e.evaluate for e in exprs) - except AttributeError: - # E.g., a generic SymPy expression or a number - _exprs = exprs - - variables = list(v for e in _exprs for v in retrieve_function_carriers(e)) - - # Implicit dimensions - implicit_dims = self._augment_implicit_dims(implicit_dims, variables) - # Move all temporaries inside inner loop to improve parallelism - # Can only be done for inject as interpolation need a temporary - # summing temp that wouldn't allow collapsing - implicit_dims = implicit_dims + tuple(r.parent for r in - self._rdim(subdomain=subdomain)) - - # List of indirection indices for all adjacent grid points - finterp = fields + as_tuple(variables) - idx_subs, temps = self._interp_idx(finterp, implicit_dims=implicit_dims, - subdomain=subdomain) - - # Substitute coordinate base symbols into the interpolation coefficients - eqns = [Inc(_field.xreplace(idx_subs), - (self._weights(subdomain=subdomain) * _expr).xreplace(idx_subs), - implicit_dims=implicit_dims) - for (_field, _expr) in zip(fields, _exprs, strict=True)] - - return temps + eqns + # Derivatives must be evaluated before the introduction of indirect + # accesses. Variables are sampled at their own grid location; the + # position map for the target field carries the staggering so the + # field's stencil neighbors land on the right indices. + with suppress(AttributeError): + exprs = tuple(e._eval_at(f).evaluate + for e, f in zip(exprs, fields, strict=True)) + + eqns = [] + temps = [] + # We need to create one set of equations (temps and and coeffs) per staggering + # field in which we inject as the reference index depends on the field's origin + for _, g in groupby(zip(fields, exprs, strict=True), lambda f: f[0].staggered): + g_fields, g_exprs = zip(*g, strict=True) + variables = list(v for e in g_exprs for v in retrieve_function_carriers(e)) + + implicit_dims = self._augment_implicit_dims(implicit_dims, variables) + + # All fields in a single injection share the same staggering by + # construction (they are written together at the same indices), so + # derive shifts from the first field. + shifts = self._field_shifts(g_fields[0]) + + # Move all temporaries inside inner loop to improve parallelism + # Can only be done for inject as interpolation needs a summing temp + # that wouldn't allow collapsing + implicit_dims = implicit_dims + tuple(r.parent for r in + self._rdim(subdomain=subdomain, + shifts=shifts)) + + # List of indirection indices for all adjacent grid points + idx_subs, _temps = self._interp_idx(list(g_fields) + variables, + implicit_dims=implicit_dims, + subdomain=subdomain, shifts=shifts) + + w = self._weights(subdomain=subdomain, shifts=shifts) + temps.extend(_temps) + eqns.extend([Inc(f.xreplace(idx_subs), (w * e).xreplace(idx_subs), + implicit_dims=implicit_dims) + for f, e in zip(g_fields, g_exprs, strict=True)]) + + return filter_ordered(temps) + eqns class LinearInterpolator(WeightedInterpolator): @@ -495,24 +523,30 @@ class LinearInterpolator(WeightedInterpolator): _name = 'linear' @memoized_meth - def _weights(self, subdomain=None): - rdim = self._rdim(subdomain=subdomain) + def _weights(self, subdomain=None, shifts=None): + rdim = self._rdim(subdomain=subdomain, shifts=shifts) c = [(1 - p) * (1 - r) + p * r - for (p, d, r) in zip(self._point_symbols, self._gdims, rdim, strict=True)] + for (p, d, r) in zip(self._point_symbols(shifts), self._gdims, rdim, + strict=True)] return Mul(*c) - @cached_property - def _point_symbols(self): + @memoized_meth + def _point_symbols(self, shifts=None): """Symbol for coordinate value in each Dimension of the point.""" dtype = self.sfunction.coordinates.dtype - return DimensionTuple(*(Symbol(name=f'p{d}', dtype=dtype) - for d in self.grid.dimensions), - getters=self.grid.dimensions) + symbols = [] + for d in self.grid.dimensions: + if shifts and shifts[self.grid.dimensions.index(d)] != 0: + symbols.append(Symbol(name=f'p{d}_s1', dtype=dtype)) + else: + symbols.append(Symbol(name=f'p{d}', dtype=dtype)) + return DimensionTuple(*symbols, getters=self.grid.dimensions) - def _coeff_temps(self, implicit_dims): + def _coeff_temps(self, implicit_dims, shifts=None): # Positions - pmap = self.sfunction._position_map - poseq = [Eq(self._point_symbols[d], pos - floor(pos), + pmap = self.sfunction._position_map(shifts=shifts) + psyms = self._point_symbols(shifts) + poseq = [Eq(psyms[d], pos - floor(pos), implicit_dims=implicit_dims) for (d, pos) in zip(self._gdims, pmap.keys(), strict=True)] return poseq @@ -531,23 +565,24 @@ class PrecomputedInterpolator(WeightedInterpolator): _name = 'precomp' - def _positions(self, implicit_dims): + def _positions(self, implicit_dims, shifts=None): if self.sfunction.gridpoints_data is None: - return super()._positions(implicit_dims) + return super()._positions(implicit_dims, shifts=shifts) else: # No position temp as we have directly the gridpoints return[Eq(p, k, implicit_dims=implicit_dims) - for (k, p) in self.sfunction._position_map.items()] + for (k, p) in self.sfunction._position_map(shifts=shifts).items()] @property def interpolation_coeffs(self): return self.sfunction.interpolation_coeffs @memoized_meth - def _weights(self, subdomain=None): + def _weights(self, subdomain=None, shifts=None): ddim, cdim = self.interpolation_coeffs.dimensions[1:] mappers = [{ddim: ri, cdim: rd-rd.parent.symbolic_min} - for (ri, rd) in enumerate(self._rdim(subdomain=subdomain))] + for (ri, rd) in enumerate(self._rdim(subdomain=subdomain, + shifts=shifts))] return Mul(*[self.interpolation_coeffs.subs(mapper) for mapper in mappers]) @@ -592,8 +627,8 @@ def interpolation_coeffs(self): return tuple(coeffs) @memoized_meth - def _weights(self, subdomain=None): - rdims = self._rdim(subdomain=subdomain) + def _weights(self, subdomain=None, shifts=None): + rdims = self._rdim(subdomain=subdomain, shifts=shifts) return Mul(*[ w._subs(rd, rd-rd.parent.symbolic_min) for (rd, w) in zip(rdims, self.interpolation_coeffs, strict=True) diff --git a/devito/types/basic.py b/devito/types/basic.py index cc0477bea9..22a87e848f 100644 --- a/devito/types/basic.py +++ b/devito/types/basic.py @@ -960,12 +960,12 @@ def indices(self): """The indices of the object.""" return DimensionTuple(*self.args, getters=self.dimensions) - @property + @cached_property def indices_ref(self): """The reference indices of the object (indices at first creation).""" return DimensionTuple(*self.function.indices, getters=self.dimensions) - @property + @cached_property def origin(self): """ Origin of the AbstractFunction in term of Dimension diff --git a/devito/types/dense.py b/devito/types/dense.py index 30205a9712..a1dc95d75c 100644 --- a/devito/types/dense.py +++ b/devito/types/dense.py @@ -1126,7 +1126,8 @@ def _eval_at(self, func): for d in self.dimensions: try: if self.indices_ref[d] is not func.indices_ref[d]: - f_idx = func.indices_ref[d]._subs(func.dimensions[d], d) + d0 = func.dimensions.get(d, d) + f_idx = func.indices_ref[d]._subs(d0, d) mapper[self.indices_ref[d]] = f_idx except KeyError: pass diff --git a/devito/types/sparse.py b/devito/types/sparse.py index aeb3c6c544..e7b528b370 100644 --- a/devito/types/sparse.py +++ b/devito/types/sparse.py @@ -23,7 +23,7 @@ ) from devito.types.dimension import dimensions as mkdims from devito.types.equation import Eq, Inc -from devito.types.utils import IgnoreDimSort +from devito.types.utils import DimensionTuple, IgnoreDimSort __all__ = [ 'MatrixSparseTimeFunction', @@ -257,6 +257,21 @@ def sparse_position(self): def _sparse_dim(self): return self.dimensions[self.sparse_position] + @cached_property + def indices_ref(self): + getters = (*self.dimensions, *self.grid.dimensions) + indices = (*self.dimensions, *self.grid.dimensions) + return DimensionTuple(*indices, getters=getters) + + @property + def _grid_map(self): + return {} + + @cached_property + def origin(self): + return DimensionTuple(*[0]*len(self.dimensions), + getters=self.dimensions) + @property def _mpitype(self): return dtype_to_mpidtype(self.dtype) @@ -352,10 +367,15 @@ def coordinates_data(self): except AttributeError: return None - @cached_property - def _pos_symbols(self): - return [Symbol(name=f'pos{d}', dtype=np.int32) - for d in self.grid.dimensions] + @memoized_meth + def _pos_symbols(self, shifts=None): + symbols = [] + for d in self.grid.dimensions: + if shifts and shifts[self.grid.dimensions.index(d)] != 0: + symbols.append(Symbol(name=f'pos{d}_s1', dtype=np.int32)) + else: + symbols.append(Symbol(name=f'pos{d}', dtype=np.int32)) + return symbols @cached_property def _point_increments(self): @@ -366,19 +386,24 @@ def _point_increments(self): def _point_support(self): return np.array(self._point_increments) - @cached_property - def _position_map(self): + @memoized_meth + def _position_map(self, shifts=None): """ - Symbols map for the physical position of the sparse points relative to the grid - origin. + Symbols map for the physical position of the sparse points relative to the + origin of the target field. ``shifts`` is a tuple, one entry per grid + Dimension, of additional physical offsets to subtract from the sparse + coordinates (e.g. ``h_x/2`` for a field staggered in ``x``). If ``shifts`` + is None, only the grid origin is subtracted. """ + shifts = shifts or (0,) * len(self.grid.dimensions) return OrderedDict([ - ((c - o)/d.spacing, p) - for p, c, d, o in zip( - self._pos_symbols, + ((c - o - s)/d.spacing, p) + for p, c, d, o, s in zip( + self._pos_symbols(shifts=shifts), self._coordinate_symbols, self.grid.dimensions, self.grid.origin_symbols, + shifts, strict=True ) ]) @@ -405,6 +430,9 @@ def _cond_rdim(self, dim, cond): parent = self._crdim(dim) return ConditionalDimension(parent.name, parent, condition=cond, indirect=True) + def _eval_at(self, func): + return self + def interpolate(self, *args, **kwargs): """ Implement an interpolation operation from the grid onto the given sparse points @@ -433,7 +461,7 @@ def guard(self, expr=None): conditions = {} # Position map and temporaries for it - pmap = self._position_map + pmap = self._position_map() # Temporaries for the position temps = self.interpolator._positions(self.dimensions) @@ -695,9 +723,6 @@ def _dist_scatter(self, alias=None, data=None): mapper.update(self._dist_subfunc_scatter(sf)) return mapper - def _eval_at(self, func): - return self - def _halo_exchange(self): # no-op for SparseFunctions return @@ -1268,8 +1293,8 @@ def _coordinate_symbols(self): return tuple([self.coordinates._subs(d_dim, i) for i in range(self.grid.dim)]) - @cached_property - def _position_map(self): + @memoized_meth + def _position_map(self, shifts=None): """ Symbol for each grid index according to the coordinates. @@ -1289,12 +1314,12 @@ def _position_map(self): (self.gridpoints._subs(ddim, di), p) for (di, p) in zip( range(self.grid.dim), - self._pos_symbols, + self._pos_symbols(shifts=shifts), strict=True ) ) else: - return super()._position_map + return super()._position_map(shifts=shifts) class PrecomputedSparseTimeFunction(AbstractSparseTimeFunction, diff --git a/examples/seismic/tutorials/06_elastic_varying_parameters.ipynb b/examples/seismic/tutorials/06_elastic_varying_parameters.ipynb index 72161dbd30..2ef27a6c78 100644 --- a/examples/seismic/tutorials/06_elastic_varying_parameters.ipynb +++ b/examples/seismic/tutorials/06_elastic_varying_parameters.ipynb @@ -604,8 +604,8 @@ "outputs": [], "source": [ "assert np.isclose(norm(rec), 23.504, atol=0, rtol=1e-3)\n", - "assert np.isclose(norm(rec2), 2.425, atol=0, rtol=1e-3)\n", - "assert np.isclose(norm(rec3), 2.889, atol=0, rtol=1e-3)" + "assert np.isclose(norm(rec2), 2.4298, atol=0, rtol=1e-3)\n", + "assert np.isclose(norm(rec3), 2.7481, atol=0, rtol=1e-3)" ] }, { @@ -838,8 +838,8 @@ "metadata": {}, "outputs": [], "source": [ - "assert np.isclose(norm(rec2), .3250, atol=0, rtol=1e-3)\n", - "assert np.isclose(norm(rec3), .26745, atol=0, rtol=1e-3)" + "assert np.isclose(norm(rec2), .30388, atol=0, rtol=1e-3)\n", + "assert np.isclose(norm(rec3), .26633, atol=0, rtol=1e-3)" ] }, { @@ -1102,9 +1102,9 @@ "metadata": {}, "outputs": [], "source": [ - "assert np.isclose(norm(rec), 31.23, atol=0, rtol=1e-3)\n", - "assert np.isclose(norm(rec2), 3.5482, atol=0, rtol=1e-3)\n", - "assert np.isclose(norm(rec3), 4.7007, atol=0, rtol=1e-3)" + "assert np.isclose(norm(rec), 29.538, atol=0, rtol=1e-3)\n", + "assert np.isclose(norm(rec2), 1.9116, atol=0, rtol=1e-3)\n", + "assert np.isclose(norm(rec3), 3.4919, atol=0, rtol=1e-3)" ] } ], diff --git a/tests/test_dimension.py b/tests/test_dimension.py index aae7e4e0cb..19d807466b 100644 --- a/tests/test_dimension.py +++ b/tests/test_dimension.py @@ -1274,10 +1274,10 @@ def test_no_index_sparse(self): radius = 1 indices = [(INT(floor(i)), INT(floor(i))+radius) - for i in sf._position_map] + for i in sf._position_map()] bounds = [i.symbolic_size - radius for i in grid.dimensions] - eqs = [Eq(p, v) for (v, p) in sf._position_map.items()] + eqs = [Eq(p, v) for (v, p) in sf._position_map().items()] for e, i in enumerate(product(*indices)): args = [j > 0 for j in i] args.extend([j < k for j, k in zip(i, bounds, strict=True)]) diff --git a/tests/test_interpolation.py b/tests/test_interpolation.py index 8ee0f6c8dd..eda7351bb4 100644 --- a/tests/test_interpolation.py +++ b/tests/test_interpolation.py @@ -6,11 +6,13 @@ from conftest import assert_structure from devito import ( - DefaultDimension, Dimension, Eq, Function, Grid, MatrixSparseTimeFunction, Operator, - PrecomputedSparseFunction, PrecomputedSparseTimeFunction, Real, SparseFunction, - SparseTimeFunction, SubDomain, TimeFunction, switchconfig + NODE, DefaultDimension, Dimension, Eq, Function, Grid, MatrixSparseTimeFunction, + Operator, PrecomputedSparseFunction, PrecomputedSparseTimeFunction, Real, + SparseFunction, SparseTimeFunction, SubDomain, TimeFunction, VectorFunction, + switchconfig ) from devito.operations.interpolators import LinearInterpolator, SincInterpolator +from devito.tools import as_tuple from examples.seismic import ( AcquisitionGeometry, Receiver, RickerSource, TimeAxis, demo_model ) @@ -22,6 +24,11 @@ class SparseFirst(SparseFunction): _sparse_position = 0 +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + + def unit_box(name='a', shape=(11, 11), grid=None, space_order=1): """Create a field with value 0. to 1. in each dimension""" grid = grid or Grid(shape=shape) @@ -61,24 +68,6 @@ def time_points(grid, ranges, npoints, name='points', nt=10): return points -def a(shape=(11, 11)): - grid = Grid(shape=shape) - a = Function(name='a', grid=grid) - xarr = np.linspace(0., 1., shape[0]) - yarr = np.linspace(0., 1., shape[1]) - a.data[:] = np.meshgrid(xarr, yarr)[1] - return a - - -def at(shape=(11, 11)): - grid = Grid(shape=shape) - a = TimeFunction(name='a', grid=grid) - xarr = np.linspace(0., 1., shape[0]) - yarr = np.linspace(0., 1., shape[1]) - a.data[:] = np.meshgrid(xarr, yarr)[1] - return a - - def custom_points(grid, ranges, npoints, name='points'): """Create a set of sparse points from a set of coordinate ranges for each spatial dimension. @@ -115,781 +104,934 @@ def precompute_linear_interpolation(points, grid, origin, r=2): return gridpoints, interpolation_coeffs -@pytest.mark.parametrize('r', [2, 4, 6]) -def test_precomputed_interpolation(r): - """ Test interpolation with PrecomputedSparseFunction which accepts - precomputed values for interpolation coefficients - """ - shape = (101, 101) - points = np.array([(.05, .9), (.01, .8), (0.07, 0.84)]) - origin = (0, 0) - - grid = Grid(shape=shape, origin=origin) - - def init(data): - # This is data with halo so need to shift to match the m.data expectations - print(grid.spacing) - for i in range(data.shape[0]): - for j in range(data.shape[1]): - data[i, j] = sin(grid.spacing[0]*(i-r)) + sin(grid.spacing[1]*(j-r)) - return data - - m = Function(name='m', grid=grid, initializer=init, space_order=r) - - gridpoints, interpolation_coeffs = precompute_linear_interpolation(points, - grid, origin, - r=r) - - sf = PrecomputedSparseFunction(name='s', grid=grid, r=r, npoint=len(points), - gridpoints=gridpoints, - interpolation_coeffs=interpolation_coeffs) - eqn = sf.interpolate(m) - op = Operator(eqn) - - op() - expected_values = [sin(point[0]) + sin(point[1]) for point in points] - assert(all(np.isclose(sf.data, expected_values, rtol=1e-6))) - - -@pytest.mark.parametrize('r', [2, 4, 6]) -def test_precomputed_interpolation_time(r): - """ Test interpolation with PrecomputedSparseFunction which accepts - precomputed values for interpolation coefficients, but this time - with a TimeFunction - """ - shape = (101, 101) - points = [(.05, .9), (.01, .8), (0.07, 0.84)] - origin = (0, 0) +# Parametrization shared across linear interpolation/injection tests +SHAPE_COORDS = [ + ((11, 11), [(.05, .9), (.01, .8)]), + ((11, 11, 11), [(.05, .9), (.01, .8), (0.07, 0.84)]) +] - grid = Grid(shape=shape, origin=origin) +SHAPE_COORDS_INJECT = [ + ((11, 11), [(.05, .95), (.45, .45)], 1.), + ((11, 11, 11), [(.05, .95), (.45, .45), (.45, .45)], 0.5) +] - u = TimeFunction(name='u', grid=grid, space_order=r, save=5) - for it in range(5): - u.data[it, :] = it - gridpoints, interpolation_coeffs = precompute_linear_interpolation(points, - grid, origin, - r=r) +# --------------------------------------------------------------------------- +# Linear interpolation / injection +# --------------------------------------------------------------------------- - sf = PrecomputedSparseTimeFunction(name='s', grid=grid, r=r, npoint=len(points), - nt=5, gridpoints=gridpoints, - interpolation_coeffs=interpolation_coeffs) - assert sf.data.shape == (5, 3) +class TestLinear: + """Tests for the default linear interpolation and injection.""" - eqn = sf.interpolate(u) - op = Operator(eqn) + @pytest.mark.parametrize('shape, coords', SHAPE_COORDS) + def test_interpolate(self, shape, coords, npoints=20): + """Test generic point interpolation testing the x-coordinate of an + arbitrary set of points going across the grid. + """ + a = unit_box(shape=shape) + p = points(a.grid, coords, npoints=npoints) + xcoords = p.coordinates.data[:, 0] - op(time_m=0, time_M=4) + expr = p.interpolate(a) + op = Operator(expr) - for it in range(5): - assert np.allclose(sf.data[it, :], it) + op(a=a) + assert np.allclose(p.data[:], xcoords, rtol=1e-6) + @pytest.mark.parametrize('shape, coords', SHAPE_COORDS) + def test_interpolate_cumm(self, shape, coords, npoints=20): + """Test generic point interpolation testing the x-coordinate of an + arbitrary set of points going across the grid. + """ + a = unit_box(shape=shape) + p = points(a.grid, coords, npoints=npoints) + xcoords = p.coordinates.data[:, 0] -@pytest.mark.parametrize('r', [2, 4, 6]) -def test_precomputed_injection(r): - """Test injection with PrecomputedSparseFunction which accepts - precomputed values for interpolation coefficients - """ - shape = (11, 11) - coords = [(.05, .95), (.45, .45)] - origin = (0, 0) - result = 0.25 + p.data[:] = 1. + expr = p.interpolate(a, increment=True) + op = Operator(expr) - m = unit_box(shape=shape, space_order=r) - m.data[:] = 0. + op(a=a) - gridpoints, interpolation_coeffs = precompute_linear_interpolation(coords, - m.grid, origin, - r=r) + assert np.allclose(p.data[:], xcoords + 1., rtol=1e-6) - sf = PrecomputedSparseFunction(name='s', grid=m.grid, r=r, npoint=len(coords), - gridpoints=gridpoints, - interpolation_coeffs=interpolation_coeffs) + @pytest.mark.parametrize('shape, coords', SHAPE_COORDS) + def test_interpolate_time_shift(self, shape, coords, npoints=20): + """Test generic point interpolation testing the x-coordinate of an + arbitrary set of points going across the grid. + This test verifies the optional time shifting for SparseTimeFunctions + """ + a = unit_box_time(shape=shape) + p = time_points(a.grid, coords, npoints=npoints, nt=10) + xcoords = p.coordinates.data[:, 0] - expr = sf.inject(m, Float(1.)) + p.data[:] = 1. + expr = p.interpolate(a, u_t=a.indices[0]+1) + op = Operator(expr) - op = Operator(expr) + op(a=a) - op() - indices = [slice(0, 2, 1), slice(9, 11, 1)] - assert np.allclose(m.data[indices], result, rtol=1.e-5) + assert np.allclose(p.data[0, :], xcoords, rtol=1e-6) - indices = [slice(4, 6, 1) for _ in coords] - assert np.allclose(m.data[indices], result, rtol=1.e-5) + p.data[:] = 1. + expr = p.interpolate(a, p_t=p.indices[0]+1) + op = Operator(expr) + op(a=a) -@pytest.mark.parametrize('r', [2, 4, 6]) -def test_precomputed_injection_time(r): - """Test injection with PrecomputedSparseFunction which accepts - precomputed values for interpolation coefficients - """ - shape = (11, 11) - coords = [(.05, .95), (.45, .45)] - origin = (0, 0) - result = 0.25 - nt = 20 + assert np.allclose(p.data[1, :], xcoords, rtol=1e-6) - m = unit_box_time(shape=shape, space_order=r) - m.data[:] = 0. + p.data[:] = 1. + expr = p.interpolate(a, u_t=a.indices[0]+1, + p_t=p.indices[0]+1) + op = Operator(expr) - gridpoints, interpolation_coeffs = precompute_linear_interpolation(coords, - m.grid, origin, - r=r) + op(a=a) - sf = PrecomputedSparseTimeFunction(name='s', grid=m.grid, r=r, npoint=len(coords), - gridpoints=gridpoints, nt=nt, - interpolation_coeffs=interpolation_coeffs) - sf.data.fill(1.) - expr = sf.inject(m, sf) + assert np.allclose(p.data[1, :], xcoords, rtol=1e-6) - op = Operator(expr) + @pytest.mark.parametrize('shape, coords', SHAPE_COORDS) + def test_interpolate_array(self, shape, coords, npoints=20): + """Test generic point interpolation testing the x-coordinate of an + arbitrary set of points going across the grid. + """ + a = unit_box(shape=shape) + p = points(a.grid, coords, npoints=npoints) + xcoords = p.coordinates.data[:, 0] - op() - for ti in range(2): - indices = [slice(0, 2, 1), slice(9, 11, 1)] - assert np.allclose(m.data[ti][indices], nt*result/2, rtol=1.e-5) + expr = p.interpolate(a) + op = Operator(expr) - indices = [slice(4, 6, 1) for _ in coords] - assert np.allclose(m.data[ti][indices], nt*result/2, rtol=1.e-5) + op(a=a, points=p.data[:]) + assert np.allclose(p.data[:], xcoords, rtol=1e-6) -@pytest.mark.parametrize('shape, coords', [ - ((11, 11), [(.05, .9), (.01, .8)]), - ((11, 11, 11), [(.05, .9), (.01, .8), (0.07, 0.84)]) -]) -def test_interpolate(shape, coords, npoints=20): - """Test generic point interpolation testing the x-coordinate of an - arbitrary set of points going across the grid. - """ - a = unit_box(shape=shape) - p = points(a.grid, coords, npoints=npoints) - xcoords = p.coordinates.data[:, 0] + @pytest.mark.parametrize('shape, coords', SHAPE_COORDS) + def test_interpolate_custom(self, shape, coords, npoints=20): + """Test generic point interpolation testing the x-coordinate of an + arbitrary set of points going across the grid. + """ + a = unit_box(shape=shape) + p = custom_points(a.grid, coords, npoints=npoints) + xcoords = p.coordinates.data[:, 0] - expr = p.interpolate(a) - op = Operator(expr) + p.data[:] = 1. + expr = p.interpolate(a * p.indices[0]) + op = Operator(expr) - op(a=a) - assert np.allclose(p.data[:], xcoords, rtol=1e-6) + op(a=a) + assert np.allclose(p.data[0, :], 0.0 * xcoords, rtol=1e-6) + assert np.allclose(p.data[1, :], 1.0 * xcoords, rtol=1e-6) + assert np.allclose(p.data[2, :], 2.0 * xcoords, rtol=1e-6) -@pytest.mark.parametrize('shape, coords', [ - ((11, 11), [(.05, .9), (.01, .8)]), - ((11, 11, 11), [(.05, .9), (.01, .8), (0.07, 0.84)]) -]) -def test_interpolate_cumm(shape, coords, npoints=20): - """Test generic point interpolation testing the x-coordinate of an - arbitrary set of points going across the grid. - """ - a = unit_box(shape=shape) - p = points(a.grid, coords, npoints=npoints) - xcoords = p.coordinates.data[:, 0] + @pytest.mark.parametrize('shape, coords', SHAPE_COORDS) + def test_interpolate_indexed(self, shape, coords, npoints=20): + """Test generic point interpolation testing the x-coordinate of an + arbitrary set of points going across the grid. Unlike other tests, + here we interpolate an expression built using the indexed notation. + """ + a = unit_box(shape=shape) + p = custom_points(a.grid, coords, npoints=npoints) + xcoords = p.coordinates.data[:, 0] - p.data[:] = 1. - expr = p.interpolate(a, increment=True) - op = Operator(expr) + p.data[:] = 1. + expr = p.interpolate(a[a.grid.dimensions] * p.indices[0]) + op = Operator(expr) - op(a=a) + op(a=a) - assert np.allclose(p.data[:], xcoords + 1., rtol=1e-6) + assert np.allclose(p.data[0, :], 0.0 * xcoords, rtol=1e-6) + assert np.allclose(p.data[1, :], 1.0 * xcoords, rtol=1e-6) + assert np.allclose(p.data[2, :], 2.0 * xcoords, rtol=1e-6) + def test_interpolation_dx(self): + """ + Test interpolation of a SparseFunction from a Derivative of + a Function. + """ + u = unit_box(shape=(11, 11)) + sf1 = SparseFunction(name='s', grid=u.grid, npoint=1) + sf1.coordinates.data[0, :] = (0.5, 0.5) -@pytest.mark.parametrize('shape, coords', [ - ((11, 11), [(.05, .9), (.01, .8)]), - ((11, 11, 11), [(.05, .9), (.01, .8), (0.07, 0.84)]) -]) -def test_interpolate_time_shift(shape, coords, npoints=20): - """Test generic point interpolation testing the x-coordinate of an - arbitrary set of points going across the grid. - This test verifies the optional time shifting for SparseTimeFunctions - """ - a = unit_box_time(shape=shape) - p = time_points(a.grid, coords, npoints=npoints, nt=10) - xcoords = p.coordinates.data[:, 0] + op = Operator(sf1.interpolate(u.dx)) - p.data[:] = 1. - expr = p.interpolate(a, u_t=a.indices[0]+1) - op = Operator(expr) + assert sf1.data.shape == (1,) + u.data[:] = 0.0 + u.data[5, 5] = 4.0 + u.data[4, 5] = 2.0 + u.data[6, 5] = 2.0 - op(a=a) + op.apply() + # Exactly in the middle of 4 points, only 1 nonzero is 4 + assert sf1.data[0] == pytest.approx(-20.0) - assert np.allclose(p.data[0, :], xcoords, rtol=1e-6) + @pytest.mark.parametrize('shape, coords, result', SHAPE_COORDS_INJECT) + def test_inject(self, shape, coords, result, npoints=19): + """Test point injection with a set of points forming a line + through the middle of the grid. + """ + a = unit_box(shape=shape) + a.data[:] = 0. + p = points(a.grid, ranges=coords, npoints=npoints) - p.data[:] = 1. - expr = p.interpolate(a, p_t=p.indices[0]+1) - op = Operator(expr) + expr = p.inject(a, Float(1.)) - op(a=a) + op = Operator(expr) - assert np.allclose(p.data[1, :], xcoords, rtol=1e-6) + op(a=a) - p.data[:] = 1. - expr = p.interpolate(a, u_t=a.indices[0]+1, - p_t=p.indices[0]+1) - op = Operator(expr) + indices = [slice(4, 6, 1) for _ in coords] + indices[0] = slice(1, -1, 1) + assert np.allclose(a.data[indices], result, rtol=1.e-5) + + @pytest.mark.parametrize('shape, coords, nexpr, result', [ + ((11, 11), [(.05, .95), (.45, .45)], 1, 1.), + ((11, 11), [(.05, .95), (.45, .45)], 2, 1.), + ((11, 11, 11), [(.05, .95), (.45, .45), (.45, .45)], 1, 0.5), + ((11, 11, 11), [(.05, .95), (.45, .45), (.45, .45)], 2, 0.5) + ]) + def test_multi_inject(self, shape, coords, nexpr, result, npoints=19): + """Test point injection with a set of points forming a line + through the middle of the grid. + """ + a1 = unit_box(name='a1', shape=shape) + a2 = unit_box(name='a2', shape=shape, grid=a1.grid) + a1.data[:] = 0. + a2.data[:] = 0. + p = points(a1.grid, ranges=coords, npoints=npoints) - op(a=a) + iexpr = Float(1.) if nexpr == 1 else (Float(1.), Float(2.)) + expr = p.inject((a1, a2), iexpr) - assert np.allclose(p.data[1, :], xcoords, rtol=1e-6) + op = Operator(expr) + op(a1=a1, a2=a2) -@pytest.mark.parametrize('shape, coords', [ - ((11, 11), [(.05, .9), (.01, .8)]), - ((11, 11, 11), [(.05, .9), (.01, .8), (0.07, 0.84)]) -]) -def test_interpolate_array(shape, coords, npoints=20): - """Test generic point interpolation testing the x-coordinate of an - arbitrary set of points going across the grid. - """ - a = unit_box(shape=shape) - p = points(a.grid, coords, npoints=npoints) - xcoords = p.coordinates.data[:, 0] + indices = [slice(4, 6, 1) for _ in coords] + indices[0] = slice(1, -1, 1) + result = (result, result) if nexpr == 1 else (result, 2 * result) + for r, a in zip(result, (a1, a2), strict=True): + assert np.allclose(a.data[indices], r, rtol=1.e-5) + + @pytest.mark.parametrize('shape, coords, result', SHAPE_COORDS_INJECT) + def test_inject_time_shift(self, shape, coords, result, npoints=19): + """Test generic point injection testing the x-coordinate of an + arbitrary set of points going across the grid. + This test verifies the optional time shifting for SparseTimeFunctions + """ + a = unit_box_time(shape=shape) + a.data[:] = 0. + p = time_points(a.grid, ranges=coords, npoints=npoints) - expr = p.interpolate(a) - op = Operator(expr) + expr = p.inject(a, Float(1.), u_t=a.indices[0]+1) - op(a=a, points=p.data[:]) + op = Operator(expr) - assert np.allclose(p.data[:], xcoords, rtol=1e-6) + op(a=a, time=1) + indices = [slice(1, 1, 1)] + [slice(4, 6, 1) for _ in coords] + indices[1] = slice(1, -1, 1) + assert np.allclose(a.data[indices], result, rtol=1.e-5) -@pytest.mark.parametrize('shape, coords', [ - ((11, 11), [(.05, .9), (.01, .8)]), - ((11, 11, 11), [(.05, .9), (.01, .8), (0.07, 0.84)]) -]) -def test_interpolate_custom(shape, coords, npoints=20): - """Test generic point interpolation testing the x-coordinate of an - arbitrary set of points going across the grid. - """ - a = unit_box(shape=shape) - p = custom_points(a.grid, coords, npoints=npoints) - xcoords = p.coordinates.data[:, 0] + a.data[:] = 0. + expr = p.inject(a, Float(1.), p_t=p.indices[0]+1) - p.data[:] = 1. - expr = p.interpolate(a * p.indices[0]) - op = Operator(expr) + op = Operator(expr) - op(a=a) + op(a=a, time=1) - assert np.allclose(p.data[0, :], 0.0 * xcoords, rtol=1e-6) - assert np.allclose(p.data[1, :], 1.0 * xcoords, rtol=1e-6) - assert np.allclose(p.data[2, :], 2.0 * xcoords, rtol=1e-6) + indices = [slice(0, 0, 1)] + [slice(4, 6, 1) for _ in coords] + indices[1] = slice(1, -1, 1) + assert np.allclose(a.data[indices], result, rtol=1.e-5) + a.data[:] = 0. + expr = p.inject(a, Float(1.), u_t=a.indices[0]+1, p_t=p.indices[0]+1) -def test_interpolation_dx(): - """ - Test interpolation of a SparseFunction from a Derivative of - a Function. - """ - u = unit_box(shape=(11, 11)) - sf1 = SparseFunction(name='s', grid=u.grid, npoint=1) - sf1.coordinates.data[0, :] = (0.5, 0.5) + op = Operator(expr) - op = Operator(sf1.interpolate(u.dx)) + op(a=a, time=1) - assert sf1.data.shape == (1,) - u.data[:] = 0.0 - u.data[5, 5] = 4.0 - u.data[4, 5] = 2.0 - u.data[6, 5] = 2.0 + indices = [slice(1, 1, 1)] + [slice(4, 6, 1) for _ in coords] + indices[1] = slice(1, -1, 1) + assert np.allclose(a.data[indices], result, rtol=1.e-5) - op.apply() - # Exactly in the middle of 4 points, only 1 nonzero is 4 - assert sf1.data[0] == pytest.approx(-20.0) + @pytest.mark.parametrize('shape, coords, result', SHAPE_COORDS_INJECT) + def test_inject_array(self, shape, coords, result, npoints=19): + """Test point injection with a set of points forming a line + through the middle of the grid. + """ + a = unit_box(shape=shape) + a.data[:] = 0. + p = points(a.grid, ranges=coords, npoints=npoints) + p2 = points(a.grid, ranges=coords, npoints=npoints, name='p2') + p2.data[:] = 1. + expr = p.inject(a, p) + op = Operator(expr) -@pytest.mark.parametrize('shape, coords', [ - ((11, 11), [(.05, .9), (.01, .8)]), - ((11, 11, 11), [(.05, .9), (.01, .8), (0.07, 0.84)]) -]) -def test_interpolate_indexed(shape, coords, npoints=20): - """Test generic point interpolation testing the x-coordinate of an - arbitrary set of points going across the grid. Unlike other tests, - here we interpolate an expression built using the indexed notation. - """ - a = unit_box(shape=shape) - p = custom_points(a.grid, coords, npoints=npoints) - xcoords = p.coordinates.data[:, 0] + op(a=a, points=p2.data[:]) + + indices = [slice(4, 6, 1) for _ in coords] + indices[0] = slice(1, -1, 1) + assert np.allclose(a.data[indices], result, rtol=1.e-5) - p.data[:] = 1. - expr = p.interpolate(a[a.grid.dimensions] * p.indices[0]) - op = Operator(expr) + @pytest.mark.parametrize('shape, coords, result', SHAPE_COORDS_INJECT) + def test_inject_from_field(self, shape, coords, result, npoints=19): + """Test point injection from a second field along a line + through the middle of the grid. + """ + a = unit_box(shape=shape) + a.data[:] = 0. + b = Function(name='b', grid=a.grid) + b.data[:] = 1. + p = points(a.grid, ranges=coords, npoints=npoints) - op(a=a) + expr = p.inject(field=a, expr=b) + op = Operator(expr) - assert np.allclose(p.data[0, :], 0.0 * xcoords, rtol=1e-6) - assert np.allclose(p.data[1, :], 1.0 * xcoords, rtol=1e-6) - assert np.allclose(p.data[2, :], 2.0 * xcoords, rtol=1e-6) + op(a=a, b=b) + indices = [slice(4, 6, 1) for _ in coords] + indices[0] = slice(1, -1, 1) + assert np.allclose(a.data[indices], result, rtol=1.e-5) -@pytest.mark.parametrize('shape, coords, result', [ - ((11, 11), [(.05, .95), (.45, .45)], 1.), - ((11, 11, 11), [(.05, .95), (.45, .45), (.45, .45)], 0.5) -]) -def test_inject(shape, coords, result, npoints=19): - """Test point injection with a set of points forming a line - through the middle of the grid. - """ - a = unit_box(shape=shape) - a.data[:] = 0. - p = points(a.grid, ranges=coords, npoints=npoints) + def test_edge_sparse(self): + """ + Test that interpolation uses the correct point for the edge case + where the sparse point is at the origin with non rational grid spacing. + Due to round up error the interpolation would use the halo point instead of + the point (0, 0) without the factorizaion of the expressions. + """ + grid = Grid(shape=(16, 16), extent=(225., 225.), origin=(25., 35.)) + u = unit_box(shape=(16, 16), grid=grid) + u._data_with_outhalo[:u.space_order, :] = -1 + u._data_with_outhalo[:, :u.space_order] = -1 + sf1 = SparseFunction(name='s', grid=u.grid, npoint=1) + sf1.coordinates.data[0, :] = (25.0, 35.0) + + expr = sf1.interpolate(u) + subs = {d.spacing: v for d, v in zip(u.grid.dimensions, u.grid.spacing, + strict=True)} + op = Operator(expr, subs=subs) + + op() + assert sf1.data[0] == 0 + + @pytest.mark.parametrize('stagg', [ + 'NODE', 'x', 'y', 'z', + '(x, y)', '(x, z)', '(y, z)', + '(x, y, z)' + ]) + def test_interpolate_staggered(self, stagg): + grid = Grid((11, 11, 11)) + x, y, z = grid.dimensions # noqa + staggered = eval(stagg) + a = Function(name='a', grid=grid, space_order=2, + staggered=staggered) + a.data.fill(1) + a.data[5, 5, 5] = 2 + + p = SparseFunction(name="p", grid=a.grid, nt=10, npoint=1) + p.coordinates.data[0, :] = (0.5, 0.5, 0.5) + + expr = p.interpolate(a) + op = Operator(expr) + + op() + # The expected value is 0.5 for the NODE case. + # For the staggered case, the center point is + # shifted left compared to the staggered function + if stagg == 'NODE': + expected = 2 + elif len(as_tuple(staggered)) == 1: + # Staggered in one dimension leads to single + # two point interpolation along that dim. + # Data is 1 everywhere and 2 at the center so + # (.5) * 1 + (.5) * 2 = 1.5 + expected = 1.5 + elif len(as_tuple(staggered)) == 2: + # Square interpolation with three 1 and one 2 + expected = 1.25 + elif len(as_tuple(staggered)) == 3: + # Cube interpolation with seven 1 and one 2 + expected = 1.125 + + assert np.allclose(p.data[:], expected, rtol=1e-6) + + @pytest.mark.parametrize('stagg', [ + 'NODE', 'x', 'y', 'z', + '(x, y)', '(x, z)', '(y, z)', + '(x, y, z)' + ]) + def test_inject_staggered(self, stagg): + grid = Grid((11, 11, 11)) + x, y, z = grid.dimensions # noqa + staggered = eval(stagg) + a = Function(name='a', grid=grid, space_order=2, + staggered=staggered) + a.data.fill(0) + + b = Function(name='b', grid=a.grid, space_order=8, + staggered=NODE) + b.data.fill(1) + b.data[5, 5, 5] = 2 + + p = SparseFunction(name="p", grid=a.grid, nt=10, npoint=1) + p.coordinates.data[0, :] = (0.5, 0.5, 0.5) + p.data[:] = 1. + + expr = p.inject(a, expr=b * p) + op = Operator(expr) + + op() + if stagg == 'NODE': + assert np.isclose(a.data[5, 5, 5], 2, rtol=1e-6) + # all other should be zero + assert np.sum(a.data) == 2 + else: + # Bottom corner since the source position is left of the staggered field + corner = [5, 5, 5] - np.array(a.staggered) + # Indices touched by the interpolation based on staggering + slices = [slice(corner[i], 6) for i in range(3)] + # Number of points for the interpolation. + # Single dim only interpolates between two points, + # so 2**(number of staggered dims) is the number of points. + npoints = 2**(np.sum(np.array(a.staggered, dtype=np.int32))) + # b value at the staggered field from b._eval_at(field) + b_val = (1 * (npoints - 1) + 2) / npoints + # Then source interpolation. Source at the center of the staggered field, + # so all points have the same weight 1/npoints. + interp_val = b_val / npoints + assert np.allclose(a.data[slices], interp_val, rtol=1e-6) + # All other should be zero so should sum to the interp_val * number of points. + # Use abs to make sure there is no +- cancellations + assert np.sum(np.abs(a.data)) == interp_val * 2**(sum(np.array(a.staggered))) + + def test_inject_staggered_mixed(self): + grid = Grid((11, 11, 11)) + v = VectorFunction(name='v', grid=grid, space_order=2) + b = Function(name='b', grid=grid, space_order=2, staggered=NODE) + p = SparseFunction(name="p", grid=grid, nt=10, npoint=1) + + eq = p.inject(v, expr=b * p).evaluate + + # We should have + # - 3 injection equations v_x, v_y, v_z + # The standard 6 on node temps posx, posy, posz, px, py, pz + # 2 temps for the staggered in x vx posz_s1, px_s1 + # 2 temps for the staggered in y vy posz_s1, py_s1 + # 2 temps for the staggered in z vz posz_s1, pz_s1 + assert len(eq) == 3 + 6 + 2 + 2 + 2 + + op = Operator(eq) + # Should be a single loop nest with 3 injections + assert_structure(op, ['p_p,rp_px,rp_py,rp_pz'], 'p_prp_pxrp_py,rp_pz') + +# --------------------------------------------------------------------------- +# Precomputed interpolation / injection +# --------------------------------------------------------------------------- + + +class TestPrecomputed: + """Tests for PrecomputedSparseFunction / PrecomputedSparseTimeFunction.""" + + @pytest.fixture + def coords(self): + return np.array([(.05, .9), (.01, .8), (0.07, 0.84)]) + + @pytest.fixture + def inject_coords(self): + return [(.05, .95), (.45, .45)] + + @pytest.fixture + def origin(self): + return (0, 0) + + @pytest.mark.parametrize('r', [2, 4, 6]) + def test_precomputed_interpolation(self, r, coords, origin): + """ Test interpolation with PrecomputedSparseFunction which accepts + precomputed values for interpolation coefficients + """ + shape = (101, 101) + grid = Grid(shape=shape, origin=origin) - expr = p.inject(a, Float(1.)) + def init(data): + # This is data with halo so need to shift to match the m.data expectations + print(grid.spacing) + for i in range(data.shape[0]): + for j in range(data.shape[1]): + data[i, j] = (sin(grid.spacing[0]*(i-r)) + + sin(grid.spacing[1]*(j-r))) + return data - op = Operator(expr) + m = Function(name='m', grid=grid, initializer=init, space_order=r) - op(a=a) + gridpoints, interpolation_coeffs = precompute_linear_interpolation( + coords, grid, origin, r=r) - indices = [slice(4, 6, 1) for _ in coords] - indices[0] = slice(1, -1, 1) - assert np.allclose(a.data[indices], result, rtol=1.e-5) + sf = PrecomputedSparseFunction(name='s', grid=grid, r=r, npoint=len(coords), + gridpoints=gridpoints, + interpolation_coeffs=interpolation_coeffs) + eqn = sf.interpolate(m) + op = Operator(eqn) + + op() + expected_values = [sin(point[0]) + sin(point[1]) for point in coords] + assert(all(np.isclose(sf.data, expected_values, rtol=1e-6))) + + @pytest.mark.parametrize('r', [2, 4, 6]) + def test_precomputed_interpolation_time(self, r, coords, origin): + """ Test interpolation with PrecomputedSparseFunction which accepts + precomputed values for interpolation coefficients, but this time + with a TimeFunction + """ + shape = (101, 101) + grid = Grid(shape=shape, origin=origin) + u = TimeFunction(name='u', grid=grid, space_order=r, save=5) + for it in range(5): + u.data[it, :] = it -@pytest.mark.parametrize('shape, coords, nexpr, result', [ - ((11, 11), [(.05, .95), (.45, .45)], 1, 1.), - ((11, 11), [(.05, .95), (.45, .45)], 2, 1.), - ((11, 11, 11), [(.05, .95), (.45, .45), (.45, .45)], 1, 0.5), - ((11, 11, 11), [(.05, .95), (.45, .45), (.45, .45)], 2, 0.5) -]) -def test_multi_inject(shape, coords, nexpr, result, npoints=19): - """Test point injection with a set of points forming a line - through the middle of the grid. - """ - a1 = unit_box(name='a1', shape=shape) - a2 = unit_box(name='a2', shape=shape, grid=a1.grid) - a1.data[:] = 0. - a2.data[:] = 0. - p = points(a1.grid, ranges=coords, npoints=npoints) + gridpoints, interpolation_coeffs = precompute_linear_interpolation( + coords, grid, origin, r=r) - iexpr = Float(1.) if nexpr == 1 else (Float(1.), Float(2.)) - expr = p.inject((a1, a2), iexpr) + sf = PrecomputedSparseTimeFunction(name='s', grid=grid, r=r, + npoint=len(coords), nt=5, + gridpoints=gridpoints, + interpolation_coeffs=interpolation_coeffs) - op = Operator(expr) + assert sf.data.shape == (5, 3) - op(a1=a1, a2=a2) + eqn = sf.interpolate(u) + op = Operator(eqn) - indices = [slice(4, 6, 1) for _ in coords] - indices[0] = slice(1, -1, 1) - result = (result, result) if nexpr == 1 else (result, 2 * result) - for r, a in zip(result, (a1, a2), strict=True): - assert np.allclose(a.data[indices], r, rtol=1.e-5) + op(time_m=0, time_M=4) + for it in range(5): + assert np.allclose(sf.data[it, :], it) -@pytest.mark.parametrize('shape, coords, result', [ - ((11, 11), [(.05, .95), (.45, .45)], 1.), - ((11, 11, 11), [(.05, .95), (.45, .45), (.45, .45)], 0.5) -]) -def test_inject_time_shift(shape, coords, result, npoints=19): - """Test generic point injection testing the x-coordinate of an - arbitrary set of points going across the grid. - This test verifies the optional time shifting for SparseTimeFunctions - """ - a = unit_box_time(shape=shape) - a.data[:] = 0. - p = time_points(a.grid, ranges=coords, npoints=npoints) + @pytest.mark.parametrize('r', [2, 4, 6]) + def test_precomputed_injection(self, r, inject_coords, origin): + """Test injection with PrecomputedSparseFunction which accepts + precomputed values for interpolation coefficients + """ + shape = (11, 11) + result = 0.25 - expr = p.inject(a, Float(1.), u_t=a.indices[0]+1) + m = unit_box(shape=shape, space_order=r) + m.data[:] = 0. - op = Operator(expr) + gridpoints, interpolation_coeffs = precompute_linear_interpolation( + inject_coords, m.grid, origin, r=r) - op(a=a, time=1) + sf = PrecomputedSparseFunction(name='s', grid=m.grid, r=r, + npoint=len(inject_coords), + gridpoints=gridpoints, + interpolation_coeffs=interpolation_coeffs) - indices = [slice(1, 1, 1)] + [slice(4, 6, 1) for _ in coords] - indices[1] = slice(1, -1, 1) - assert np.allclose(a.data[indices], result, rtol=1.e-5) + expr = sf.inject(m, Float(1.)) - a.data[:] = 0. - expr = p.inject(a, Float(1.), p_t=p.indices[0]+1) + op = Operator(expr) - op = Operator(expr) + op() + indices = [slice(0, 2, 1), slice(9, 11, 1)] + assert np.allclose(m.data[indices], result, rtol=1.e-5) - op(a=a, time=1) + indices = [slice(4, 6, 1) for _ in inject_coords] + assert np.allclose(m.data[indices], result, rtol=1.e-5) - indices = [slice(0, 0, 1)] + [slice(4, 6, 1) for _ in coords] - indices[1] = slice(1, -1, 1) - assert np.allclose(a.data[indices], result, rtol=1.e-5) + @pytest.mark.parametrize('r', [2, 4, 6]) + def test_precomputed_injection_time(self, r, inject_coords, origin): + """Test injection with PrecomputedSparseFunction which accepts + precomputed values for interpolation coefficients + """ + shape = (11, 11) + result = 0.25 + nt = 20 + + m = unit_box_time(shape=shape, space_order=r) + m.data[:] = 0. + + gridpoints, interpolation_coeffs = precompute_linear_interpolation( + inject_coords, m.grid, origin, r=r) + + sf = PrecomputedSparseTimeFunction(name='s', grid=m.grid, r=r, + npoint=len(inject_coords), + gridpoints=gridpoints, nt=nt, + interpolation_coeffs=interpolation_coeffs) + sf.data.fill(1.) + expr = sf.inject(m, sf) + + op = Operator(expr) + + op() + for ti in range(2): + indices = [slice(0, 2, 1), slice(9, 11, 1)] + assert np.allclose(m.data[ti][indices], nt*result/2, rtol=1.e-5) + + indices = [slice(4, 6, 1) for _ in inject_coords] + assert np.allclose(m.data[ti][indices], nt*result/2, rtol=1.e-5) + + +# --------------------------------------------------------------------------- +# Sinc interpolation +# --------------------------------------------------------------------------- + + +class TestSinc: + """Tests for sinc interpolation.""" + + @pytest.mark.parametrize('r, tol', [(2, 0.051), (3, 0.003), (4, 0.008), + (5, 0.002), (6, 0.0005), (7, 8e-5), + (8, 6e-5), (9, 5e-5), (10, 4.2e-5)]) + def test_sinc_accuracy(self, r, tol): + so = max(2, r) + solver_lin = acoustic_setup(preset='constant-isotropic', shape=(101, 101), + spacing=(10, 10), interpolation='linear', + space_order=so) + solver_sinc = acoustic_setup(preset='constant-isotropic', shape=(101, 101), + spacing=(10, 10), interpolation='sinc', r=r, + space_order=so) + + # On node source + s_node = [500, 500] + src_n = solver_lin.geometry.src + src_n.coordinates.data[:] = s_node + + # Half node src + s_mid = [505, 505] + src_h = solver_lin.geometry.src + src_h.coordinates.data[:] = s_mid + + # On node rec + r_node = [750, 750] + rec_n = solver_lin.geometry.new_src(name='rec', src_type=None) + rec_n.coordinates.data[:] = r_node + + # Half node rec for linear + r_mid = [755, 755] + rec_hl = solver_lin.geometry.new_src(name='recl', src_type=None) + rec_hl.coordinates.data[:] = r_mid + + # Half node rec for sinc + r_mid = [755, 755] + rec_hs = solver_lin.geometry.new_src(name='recs', src_type=None) + rec_hs.coordinates.data[:] = r_mid + + # Reference solution, on node + _, un, _ = solver_lin.forward(src=src_n, rec=rec_n) + # Linear interp on half node + _, ul, _ = solver_lin.forward(src=src_h, rec=rec_hl) + # Sinc interp on half node + _, us, _ = solver_sinc.forward(src=src_h, rec=rec_hs) + + # Check sinc is more accuracte + nref = np.linalg.norm(rec_n.data) + err_lin = np.linalg.norm(rec_n.data - rec_hl.data)/nref + err_sinc = np.linalg.norm(rec_n.data - rec_hs.data)/nref + + print(f"Error linear: {err_lin}, Error sinc: {err_sinc}") + assert np.isclose(err_sinc, 0, rtol=0, atol=tol) + assert err_sinc < err_lin + assert err_lin > 0.01 + + +# --------------------------------------------------------------------------- +# Matrix sparse function interpolation / injection +# --------------------------------------------------------------------------- + + +class TestMatrixSparse: + """Tests for MatrixSparseTimeFunction interpolation / injection.""" + + def test_msf_interpolate(self): + """ Test interpolation with MatrixSparseTimeFunction which accepts + precomputed values for interpolation coefficients, but this time + with a TimeFunction + """ + shape = (101, 101) + coords = np.array([(.05, .9), (.01, .8), (0.07, 0.84)]) + origin = (0, 0) + + grid = Grid(shape=shape, origin=origin) + r = 2 # Constant for linear interpolation + # because we interpolate across 2 neighbouring points in each dimension + + u = TimeFunction(name='u', grid=grid, space_order=0, save=5) + for it in range(5): + u.data[it, :] = it + + gridpoints, interpolation_coeffs = precompute_linear_interpolation( + coords, grid, origin) - a.data[:] = 0. - expr = p.inject(a, Float(1.), u_t=a.indices[0]+1, p_t=p.indices[0]+1) + matrix = scipy.sparse.eye(len(coords)) + + sf = MatrixSparseTimeFunction( + name='s', grid=grid, r=r, matrix=matrix, nt=5 + ) - op = Operator(expr) + sf.gridpoints.data[:] = gridpoints + sf.coefficients_x.data[:] = interpolation_coeffs[:, 0, :] + sf.coefficients_y.data[:] = interpolation_coeffs[:, 0, :] - op(a=a, time=1) + assert sf.data.shape == (5, 3) - indices = [slice(1, 1, 1)] + [slice(4, 6, 1) for _ in coords] - indices[1] = slice(1, -1, 1) - assert np.allclose(a.data[indices], result, rtol=1.e-5) + eqn = sf.interpolate(u) + op = Operator(eqn) + sf.manual_scatter() + op(time_m=0, time_M=4) + sf.manual_gather() -@pytest.mark.parametrize('shape, coords, result', [ - ((11, 11), [(.05, .95), (.45, .45)], 1.), - ((11, 11, 11), [(.05, .95), (.45, .45), (.45, .45)], 0.5) -]) -def test_inject_array(shape, coords, result, npoints=19): - """Test point injection with a set of points forming a line - through the middle of the grid. - """ - a = unit_box(shape=shape) - a.data[:] = 0. - p = points(a.grid, ranges=coords, npoints=npoints) - p2 = points(a.grid, ranges=coords, npoints=npoints, name='p2') - p2.data[:] = 1. - expr = p.inject(a, p) + for it in range(5): + assert np.allclose(sf.data[it, :], it) - op = Operator(expr) + # Now test injection + u.data[:] = 0 - op(a=a, points=p2.data[:]) + eqn_inject = sf.inject(field=u, expr=sf) + op2 = Operator(eqn_inject) + op2(time_m=0, time_M=4) - indices = [slice(4, 6, 1) for _ in coords] - indices[0] = slice(1, -1, 1) - assert np.allclose(a.data[indices], result, rtol=1.e-5) + # There should be 4 points touched for each source point + # (5, 90), (1, 80), (7, 84) and x+1, y+1 for each + nzt, nzx, nzy = np.nonzero(u.data) + assert np.all(np.unique(nzx) == np.array([1, 2, 5, 6, 7, 8])) + assert np.all(np.unique(nzy) == np.array([80, 81, 84, 85, 90, 91])) + assert np.all(np.unique(nzt) == np.array([1, 2, 3, 4])) + # 12 points x 4 timesteps + assert nzt.size == 48 -@pytest.mark.parametrize('shape, coords, result', [ - ((11, 11), [(.05, .95), (.45, .45)], 1.), - ((11, 11, 11), [(.05, .95), (.45, .45), (.45, .45)], 0.5) -]) -def test_inject_from_field(shape, coords, result, npoints=19): - """Test point injection from a second field along a line - through the middle of the grid. - """ - a = unit_box(shape=shape) - a.data[:] = 0. - b = Function(name='b', grid=a.grid) - b.data[:] = 1. - p = points(a.grid, ranges=coords, npoints=npoints) +# --------------------------------------------------------------------------- +# Interpolator selection / configuration +# --------------------------------------------------------------------------- - expr = p.inject(field=a, expr=b) - op = Operator(expr) - op(a=a, b=b) +class TestInterpolator: + """Tests for interpolator selection, radius and dtype handling.""" - indices = [slice(4, 6, 1) for _ in coords] - indices[0] = slice(1, -1, 1) - assert np.allclose(a.data[indices], result, rtol=1.e-5) + @pytest.mark.parametrize('r, interp', [(2, 'linear'), (4, 'sinc')]) + def test_interpolation_radius(self, r, interp): + nt = 11 + grid = Grid(shape=(5, 5)) + u = TimeFunction(name="u", grid=grid, space_order=0) + src = SparseTimeFunction(name="src", grid=grid, nt=nt, npoint=1, + r=r, interpolation=interp) + try: + src.interpolate(u) + raise AssertionError('Assert False') + except ValueError: + assert True -@pytest.mark.parametrize('shape', [(50, 50, 50)]) -def test_position(shape): - t0 = 0.0 # Start time - tn = 500. # Final time - nrec = 130 # Number of receivers + def test_interp_default(self): + nt = 3 + grid = Grid(shape=(5, 5)) - # Create model from preset - model = demo_model('constant-isotropic', spacing=[15. for _ in shape], - shape=shape, nbl=10) + src = SparseTimeFunction(name="src", grid=grid, nt=nt, npoint=1) + assert isinstance(src.interpolator, LinearInterpolator) + assert src.r == 1 - # Derive timestepping from model spacing - dt = model.critical_dt - time_range = TimeAxis(start=t0, stop=tn, step=dt) + src = SparseTimeFunction(name="src", grid=grid, nt=nt, npoint=1, + interpolation='sinc') + assert isinstance(src.interpolator, SincInterpolator) + assert src.r == 4 - # Source and receiver geometries - src_coordinates = np.empty((1, len(shape))) - src_coordinates[0, :] = np.array(model.domain_size) * .5 - src_coordinates[0, -1] = 30. + src = SparseTimeFunction(name="src", grid=grid, nt=nt, npoint=1, + interpolation='sinc', r=6) + assert isinstance(src.interpolator, SincInterpolator) + assert src.r == 6 - rec_coordinates = np.empty((nrec, len(shape))) - rec_coordinates[:, 0] = np.linspace(0., model.domain_size[0], num=nrec) - rec_coordinates[:, 1:] = src_coordinates[0, 1:] + @pytest.mark.parametrize('dtype, expected', [(np.complex64, np.float32), + (np.complex128, np.float64)]) + def test_point_symbol_types(self, dtype, expected): + """Test that positions are always real""" + grid = Grid(shape=(11,)) + s = SparseFunction(name='src', npoint=1, + grid=grid, dtype=dtype) + point_symbol = s.interpolator._point_symbols()[0] - geometry = AcquisitionGeometry(model, rec_coordinates, src_coordinates, - t0=t0, tn=tn, src_type='Ricker', f0=0.010) - # Create solver object to provide relevant operators - solver = AcousticWaveSolver(model, geometry, time_order=2, space_order=4) + assert point_symbol.dtype is expected - rec, u, _ = solver.forward(save=False) + def test_wrong_coords(self): + grid = Grid(shape=(11, 11)) + s = SparseFunction(name='src', npoint=1, grid=grid) + s2 = SparseFunction(name='src2', npoint=1, grid=grid) + u = Function(name='u', grid=grid) - # Define source geometry (center of domain, just below surface) with 100. origin - src = RickerSource(name='src', grid=model.grid, f0=0.01, time_range=time_range) - src.coordinates.data[0, :] = np.array(model.domain_size) * .5 + 100. - src.coordinates.data[0, -1] = 130. + with pytest.raises(ValueError) as vinfo: + s.inject(u, expr=s2) + assert "Interpolation/injection with" in str(vinfo.value) - # Define receiver geometry (same as source, but spread across `x, y`) - rec2 = Receiver(name='rec', grid=model.grid, time_range=time_range, npoint=nrec) - rec2.coordinates.data[:, 0] = np.linspace(100., 100. + model.domain_size[0], - num=nrec) - rec2.coordinates.data[:, 1:] = src.coordinates.data[0, 1:] + with pytest.raises(ValueError) as vinfo: + s.interpolate(u + s2) + assert "Interpolation/injection with" in str(vinfo.value) - ox_g, oy_g, oz_g = tuple(o + 100. for o in model.grid.origin) - rec1, u1, _ = solver.forward(save=False, src=src, rec=rec2, - o_x=ox_g, o_y=oy_g, o_z=oz_g) +# --------------------------------------------------------------------------- +# Complex-valued interpolation +# --------------------------------------------------------------------------- - assert(np.allclose(rec.data, rec1.data, atol=1e-5)) +class TestComplex: + """Tests for interpolation with complex-valued Functions.""" -def test_edge_sparse(): - """ - Test that interpolation uses the correct point for the edge case - where the sparse point is at the origin with non rational grid spacing. - Due to round up error the interpolation would use the halo point instead of - the point (0, 0) without the factorizaion of the expressions. - """ - grid = Grid(shape=(16, 16), extent=(225., 225.), origin=(25., 35.)) - u = unit_box(shape=(16, 16), grid=grid) - u._data_with_outhalo[:u.space_order, :] = -1 - u._data_with_outhalo[:, :u.space_order] = -1 - sf1 = SparseFunction(name='s', grid=u.grid, npoint=1) - sf1.coordinates.data[0, :] = (25.0, 35.0) + @pytest.mark.parametrize('dtype', [np.complex64, np.complex128]) + def test_interp_complex(self, dtype): + grid = Grid((11, 11, 11)) - expr = sf1.interpolate(u) - subs = {d.spacing: v for d, v in zip(u.grid.dimensions, u.grid.spacing, strict=True)} - op = Operator(expr, subs=subs) + sc = SparseFunction(name="sc", grid=grid, npoint=1, dtype=dtype) + sc.coordinates.data[:] = [.5, .5, .5] - op() - assert sf1.data[0] == 0 + fc = Function(name="fc", grid=grid, npoint=2, dtype=dtype) + fc.data[:] = np.random.randn(*grid.shape) + 1j * np.random.randn(*grid.shape) + opC = Operator([sc.interpolate(expr=fc)], name="OpC") + opC() + assert np.isclose(sc.data[0], fc.data[5, 5, 5]) -def test_msf_interpolate(): - """ Test interpolation with MatrixSparseTimeFunction which accepts - precomputed values for interpolation coefficients, but this time - with a TimeFunction - """ - shape = (101, 101) - points = np.array([(.05, .9), (.01, .8), (0.07, 0.84)]) - origin = (0, 0) + @pytest.mark.parametrize('dtype', [np.complex64, np.complex128]) + def test_interp_complex_and_real(self, dtype): + grid = Grid((11, 11, 11)) - grid = Grid(shape=shape, origin=origin) - r = 2 # Constant for linear interpolation - # because we interpolate across 2 neighbouring points in each dimension + sc = SparseFunction(name="sc", grid=grid, npoint=1, dtype=dtype) + sc.coordinates.data[:] = [.5, .5, .5] + scre = SparseFunction(name="sce", grid=grid, npoint=1, + coordinates=sc.coordinates) - u = TimeFunction(name='u', grid=grid, space_order=0, save=5) - for it in range(5): - u.data[it, :] = it + fc = Function(name="fc", grid=grid, npoint=2, dtype=dtype) + fc.data[:] = np.random.randn(*grid.shape) + 1j * np.random.randn(*grid.shape) + exprs = sc.interpolate(expr=fc) + scre.interpolate(expr=Real(fc)) + opC = Operator(exprs, name="OpC") + opC() - gridpoints, interpolation_coeffs = precompute_linear_interpolation(points, - grid, origin) + assert np.isclose(sc.data[0], fc.data[5, 5, 5]) + assert np.isclose(scre.data[0], fc.data[5, 5, 5].real) - matrix = scipy.sparse.eye(len(points)) + assert_structure(opC, ['p_sc', 'p_sc,rp_scx,rp_scy,rp_scz', + 'p_sc,rp_scx,rp_scy,rp_scz'], + 'p_sc,rp_scx,rp_scy,rp_scz,rp_scx,rp_scy,rp_scz') - sf = MatrixSparseTimeFunction( - name='s', grid=grid, r=r, matrix=matrix, nt=5 - ) - sf.gridpoints.data[:] = gridpoints - sf.coefficients_x.data[:] = interpolation_coeffs[:, 0, :] - sf.coefficients_y.data[:] = interpolation_coeffs[:, 0, :] +# --------------------------------------------------------------------------- +# Custom sparse function layouts / acoustic-driven scenarios +# --------------------------------------------------------------------------- - assert sf.data.shape == (5, 3) - eqn = sf.interpolate(u) - op = Operator(eqn) +class TestCustom: + """Tests for custom SparseFunction layouts and acoustic-driven scenarios.""" - sf.manual_scatter() - op(time_m=0, time_M=4) - sf.manual_gather() + @pytest.mark.parametrize('shape', [(50, 50, 50)]) + def test_position(self, shape): + t0 = 0.0 # Start time + tn = 500. # Final time + nrec = 130 # Number of receivers - for it in range(5): - assert np.allclose(sf.data[it, :], it) + # Create model from preset + model = demo_model('constant-isotropic', spacing=[15. for _ in shape], + shape=shape, nbl=10) - # Now test injection - u.data[:] = 0 + # Derive timestepping from model spacing + dt = model.critical_dt + time_range = TimeAxis(start=t0, stop=tn, step=dt) - eqn_inject = sf.inject(field=u, expr=sf) - op2 = Operator(eqn_inject) - op2(time_m=0, time_M=4) + # Source and receiver geometries + src_coordinates = np.empty((1, len(shape))) + src_coordinates[0, :] = np.array(model.domain_size) * .5 + src_coordinates[0, -1] = 30. - # There should be 4 points touched for each source point - # (5, 90), (1, 80), (7, 84) and x+1, y+1 for each - nzt, nzx, nzy = np.nonzero(u.data) - assert np.all(np.unique(nzx) == np.array([1, 2, 5, 6, 7, 8])) - assert np.all(np.unique(nzy) == np.array([80, 81, 84, 85, 90, 91])) - assert np.all(np.unique(nzt) == np.array([1, 2, 3, 4])) - # 12 points x 4 timesteps - assert nzt.size == 48 + rec_coordinates = np.empty((nrec, len(shape))) + rec_coordinates[:, 0] = np.linspace(0., model.domain_size[0], num=nrec) + rec_coordinates[:, 1:] = src_coordinates[0, 1:] + geometry = AcquisitionGeometry(model, rec_coordinates, src_coordinates, + t0=t0, tn=tn, src_type='Ricker', f0=0.010) + # Create solver object to provide relevant operators + solver = AcousticWaveSolver(model, geometry, time_order=2, space_order=4) -def test_sparse_first(): - """ - Tests custom sprase function with sparse dimension as first index. - """ - dr = Dimension("cd") - ds = DefaultDimension("ps", default_value=3) - grid = Grid((11, 11)) - dims = grid.dimensions - s = SparseFirst(name="s", grid=grid, npoint=2, dimensions=(dr, ds), shape=(2, 3), - coordinates=[[.5, .5], [.2, .2]]) - - # Check dimensions and shape are correctly initialized - assert s.indices[s._sparse_position] == dr - assert s.shape == (2, 3) - assert s.coordinates.indices[0] == dr - - # Operator - u = TimeFunction(name="u", grid=grid, time_order=1) - fs = Function(name="fs", grid=grid, dimensions=(*dims, ds), shape=(11, 11, 3)) - - eqs = [Eq(u.forward, u+1), Eq(fs, u)] - # No time dependence so need the implicit dim - rec = s.interpolate(expr=s+fs, implicit_dims=grid.stepping_dim) - op = Operator(eqs + rec) - - op(time_M=10) - expected = 10*11/2 # n (n+1)/2 - assert np.allclose(s.data, expected) - - -@switchconfig(safe_math=True) -def test_inject_function(): - nt = 11 - - grid = Grid(shape=(5, 5)) - u = TimeFunction(name="u", grid=grid, time_order=2) - src = SparseTimeFunction(name="src", grid=grid, nt=nt, npoint=1, - coordinates=[[0.5, 0.5]]) - - nfreq = 5 - freq_dim = DefaultDimension(name="freq", default_value=nfreq) - omega = Function(name="omega", dimensions=(freq_dim,), shape=(nfreq,), grid=grid) - omega.data.fill(1.) - - inj = src.inject(field=u.forward, expr=omega) - - op = Operator([inj]) - - op(time_M=0) - assert u.data[1, 2, 2] == nfreq - assert np.all(u.data[0] == 0) - assert np.all(u.data[2] == 0) - for i in [0, 1, 3, 4]: - for j in [0, 1, 3, 4]: - assert u.data[1, i, j] == 0 - - -@pytest.mark.parametrize('r, interp', [(2, 'linear'), (4, 'sinc')]) -def test_interpolation_radius(r, interp): - nt = 11 - - grid = Grid(shape=(5, 5)) - u = TimeFunction(name="u", grid=grid, space_order=0) - src = SparseTimeFunction(name="src", grid=grid, nt=nt, npoint=1, - r=r, interpolation=interp) - try: - src.interpolate(u) - raise AssertionError('Assert False') - except ValueError: - assert True - - -def test_interp_default(): - nt = 3 - grid = Grid(shape=(5, 5)) - - src = SparseTimeFunction(name="src", grid=grid, nt=nt, npoint=1) - assert isinstance(src.interpolator, LinearInterpolator) - assert src.r == 1 - - src = SparseTimeFunction(name="src", grid=grid, nt=nt, npoint=1, interpolation='sinc') - assert isinstance(src.interpolator, SincInterpolator) - assert src.r == 4 - - src = SparseTimeFunction(name="src", grid=grid, nt=nt, npoint=1, - interpolation='sinc', r=6) - assert isinstance(src.interpolator, SincInterpolator) - assert src.r == 6 - - -@pytest.mark.parametrize('r, tol', [(2, 0.051), (3, 0.003), (4, 0.008), - (5, 0.002), (6, 0.0005), (7, 8e-5), - (8, 6e-5), (9, 5e-5), (10, 4.2e-5)]) -def test_sinc_accuracy(r, tol): - so = max(2, r) - solver_lin = acoustic_setup(preset='constant-isotropic', shape=(101, 101), - spacing=(10, 10), interpolation='linear', space_order=so) - solver_sinc = acoustic_setup(preset='constant-isotropic', shape=(101, 101), - spacing=(10, 10), interpolation='sinc', r=r, - space_order=so) - - # On node source - s_node = [500, 500] - src_n = solver_lin.geometry.src - src_n.coordinates.data[:] = s_node - - # Half node src - s_mid = [505, 505] - src_h = solver_lin.geometry.src - src_h.coordinates.data[:] = s_mid + rec, u, _ = solver.forward(save=False) + + # Define source geometry (center of domain, just below surface) with 100. origin + src = RickerSource(name='src', grid=model.grid, f0=0.01, time_range=time_range) + src.coordinates.data[0, :] = np.array(model.domain_size) * .5 + 100. + src.coordinates.data[0, -1] = 130. + + # Define receiver geometry (same as source, but spread across `x, y`) + rec2 = Receiver(name='rec', grid=model.grid, time_range=time_range, npoint=nrec) + rec2.coordinates.data[:, 0] = np.linspace(100., 100. + model.domain_size[0], + num=nrec) + rec2.coordinates.data[:, 1:] = src.coordinates.data[0, 1:] + + ox_g, oy_g, oz_g = tuple(o + 100. for o in model.grid.origin) - # On node rec - r_node = [750, 750] - rec_n = solver_lin.geometry.new_src(name='rec', src_type=None) - rec_n.coordinates.data[:] = r_node + rec1, u1, _ = solver.forward(save=False, src=src, rec=rec2, + o_x=ox_g, o_y=oy_g, o_z=oz_g) + + assert(np.allclose(rec.data, rec1.data, atol=1e-5)) + + def test_sparse_first(self): + """ + Tests custom sprase function with sparse dimension as first index. + """ + dr = Dimension("cd") + ds = DefaultDimension("ps", default_value=3) + grid = Grid((11, 11)) + dims = grid.dimensions + s = SparseFirst(name="s", grid=grid, npoint=2, dimensions=(dr, ds), + shape=(2, 3), coordinates=[[.5, .5], [.2, .2]]) - # Half node rec for linear - r_mid = [755, 755] - rec_hl = solver_lin.geometry.new_src(name='recl', src_type=None) - rec_hl.coordinates.data[:] = r_mid - - # Half node rec for sinc - r_mid = [755, 755] - rec_hs = solver_lin.geometry.new_src(name='recs', src_type=None) - rec_hs.coordinates.data[:] = r_mid - - # Reference solution, on node - _, un, _ = solver_lin.forward(src=src_n, rec=rec_n) - # Linear interp on half node - _, ul, _ = solver_lin.forward(src=src_h, rec=rec_hl) - # Sinc interp on half node - _, us, _ = solver_sinc.forward(src=src_h, rec=rec_hs) - - # Check sinc is more accuracte - nref = np.linalg.norm(rec_n.data) - err_lin = np.linalg.norm(rec_n.data - rec_hl.data)/nref - err_sinc = np.linalg.norm(rec_n.data - rec_hs.data)/nref - - print(f"Error linear: {err_lin}, Error sinc: {err_sinc}") - assert np.isclose(err_sinc, 0, rtol=0, atol=tol) - assert err_sinc < err_lin - assert err_lin > 0.01 + # Check dimensions and shape are correctly initialized + assert s.indices[s._sparse_position] == dr + assert s.shape == (2, 3) + assert s.coordinates.indices[0] == dr + # Operator + u = TimeFunction(name="u", grid=grid, time_order=1) + fs = Function(name="fs", grid=grid, dimensions=(*dims, ds), shape=(11, 11, 3)) -@pytest.mark.parametrize('dtype, expected', [(np.complex64, np.float32), - (np.complex128, np.float64)]) -def test_point_symbol_types(dtype, expected): - """Test that positions are always real""" - grid = Grid(shape=(11,)) - s = SparseFunction(name='src', npoint=1, - grid=grid, dtype=dtype) - point_symbol = s.interpolator._point_symbols[0] - - assert point_symbol.dtype is expected + eqs = [Eq(u.forward, u+1), Eq(fs, u)] + # No time dependence so need the implicit dim + rec = s.interpolate(expr=s+fs, implicit_dims=grid.stepping_dim) + op = Operator(eqs + rec) + op(time_M=10) + expected = 10*11/2 # n (n+1)/2 + assert np.allclose(s.data, expected) -@pytest.mark.parametrize('dtype', [np.complex64, np.complex128]) -def test_interp_complex(dtype): - grid = Grid((11, 11, 11)) - - sc = SparseFunction(name="sc", grid=grid, npoint=1, dtype=dtype) - sc.coordinates.data[:] = [.5, .5, .5] - - fc = Function(name="fc", grid=grid, npoint=2, dtype=dtype) - fc.data[:] = np.random.randn(*grid.shape) + 1j * np.random.randn(*grid.shape) - opC = Operator([sc.interpolate(expr=fc)], name="OpC") - opC() + @switchconfig(safe_math=True) + def test_inject_function(self): + nt = 11 - assert np.isclose(sc.data[0], fc.data[5, 5, 5]) + grid = Grid(shape=(5, 5)) + u = TimeFunction(name="u", grid=grid, time_order=2) + src = SparseTimeFunction(name="src", grid=grid, nt=nt, npoint=1, + coordinates=[[0.5, 0.5]]) + nfreq = 5 + freq_dim = DefaultDimension(name="freq", default_value=nfreq) + omega = Function(name="omega", dimensions=(freq_dim,), shape=(nfreq,), + grid=grid) + omega.data.fill(1.) -@pytest.mark.parametrize('dtype', [np.complex64, np.complex128]) -def test_interp_complex_and_real(dtype): - grid = Grid((11, 11, 11)) + inj = src.inject(field=u.forward, expr=omega) - sc = SparseFunction(name="sc", grid=grid, npoint=1, dtype=dtype) - sc.coordinates.data[:] = [.5, .5, .5] - scre = SparseFunction(name="sce", grid=grid, npoint=1, coordinates=sc.coordinates) + op = Operator([inj]) - fc = Function(name="fc", grid=grid, npoint=2, dtype=dtype) - fc.data[:] = np.random.randn(*grid.shape) + 1j * np.random.randn(*grid.shape) - exprs = sc.interpolate(expr=fc) + scre.interpolate(expr=Real(fc)) - opC = Operator(exprs, name="OpC") - opC() - - assert np.isclose(sc.data[0], fc.data[5, 5, 5]) - assert np.isclose(scre.data[0], fc.data[5, 5, 5].real) + op(time_M=0) + assert u.data[1, 2, 2] == nfreq + assert np.all(u.data[0] == 0) + assert np.all(u.data[2] == 0) + for i in [0, 1, 3, 4]: + for j in [0, 1, 3, 4]: + assert u.data[1, i, j] == 0 - assert_structure(opC, ['p_sc', 'p_sc,rp_scx,rp_scy,rp_scz', - 'p_sc,rp_scx,rp_scy,rp_scz'], - 'p_sc,rp_scx,rp_scy,rp_scz,rp_scx,rp_scy,rp_scz') + +# --------------------------------------------------------------------------- +# SubDomain interpolation / injection +# --------------------------------------------------------------------------- class SD0(SubDomain): @@ -962,12 +1104,25 @@ class TestSubDomainInterpolation: SubDomains. """ - def test_interpolate_subdomain(self): + @pytest.fixture + def grid(self): + return Grid(shape=(11, 11), extent=(10., 10.)) + + @pytest.fixture + def coords(self): + return np.array([[2.5, 1.5], [4.5, 2.], [8.5, 4.], + [0.5, 6.], [7.5, 4.], [5.5, 5.5], + [1.5, 4.5], [7.5, 8.5]]) + + @pytest.fixture + def sinc_coords(self): + return np.array([[2.5, 6.5], [3.5, 4.5], [6., 6.], + [5.5, 4.5], [4.5, 6.]]) + + def test_interpolate_subdomain(self, grid, coords): """ Test interpolation off of a Function defined on a SubDomain. """ - - grid = Grid(shape=(11, 11), extent=(10., 10.)) sd0 = SD0(grid=grid) sd1 = SD1(grid=grid) @@ -985,10 +1140,6 @@ def test_interpolate_subdomain(self): sr1 = SparseFunction(name='sr1', grid=grid, npoint=8) sr2 = SparseFunction(name='sr2', grid=grid, npoint=8) - coords = np.array([[2.5, 1.5], [4.5, 2.], [8.5, 4.], - [0.5, 6.], [7.5, 4.], [5.5, 5.5], - [1.5, 4.5], [7.5, 8.5]]) - sr0.coordinates.data[:] = coords sr1.coordinates.data[:] = coords sr2.coordinates.data[:] = coords @@ -1015,11 +1166,10 @@ def test_interpolate_subdomain(self): 'p_sr1rp_sr1xrp_sr1y', 'p_sr2', 'p_sr2rp_sr2xrp_sr2y'], 'p_sr0rp_sr0xrp_sr0yp_sr1rp_sr1xrp_sr1yp_sr2rp_sr2xrp_sr2y') - def test_interpolate_subdomain_sinc(self): + def test_interpolate_subdomain_sinc(self, grid, sinc_coords): """ Check that sinc interpolation off a SubDomain works as expected. """ - grid = Grid(shape=(11, 11), extent=(10., 10.)) sd0 = SD6(grid=grid) sd1 = SD7(grid=grid) @@ -1037,11 +1187,9 @@ def test_interpolate_subdomain_sinc(self): sr1 = SparseFunction(name='sr1', grid=grid, npoint=5, interpolation='sinc', r=2) sr2 = SparseFunction(name='sr2', grid=grid, npoint=5, interpolation='sinc', r=2) - coords = np.array([[2.5, 6.5], [3.5, 4.5], [6., 6.], [5.5, 4.5], [4.5, 6.]]) - - sr0.coordinates.data[:] = coords - sr1.coordinates.data[:] = coords - sr2.coordinates.data[:] = coords + sr0.coordinates.data[:] = sinc_coords + sr1.coordinates.data[:] = sinc_coords + sr2.coordinates.data[:] = sinc_coords rec0 = sr0.interpolate(f0) rec1 = sr1.interpolate(f1) @@ -1058,11 +1206,10 @@ def test_interpolate_subdomain_sinc(self): 'p_sr1rp_sr1xrp_sr1y', 'p_sr2', 'p_sr2rp_sr2xrp_sr2y'], 'p_sr0rp_sr0xrp_sr0yp_sr1rp_sr1xrp_sr1yp_sr2rp_sr2xrp_sr2y') - def test_inject_subdomain(self): + def test_inject_subdomain(self, grid, coords): """ Test injection into a Function defined on a SubDomain. """ - grid = Grid(shape=(11, 11), extent=(10., 10.)) sd0 = SD0(grid=grid) sd1 = SD1(grid=grid) @@ -1071,10 +1218,6 @@ def test_inject_subdomain(self): sr0 = SparseFunction(name='sr0', grid=grid, npoint=8) - coords = np.array([[2.5, 1.5], [4.5, 2.], [8.5, 4.], - [0.5, 6.], [7.5, 4.], [5.5, 5.5], - [1.5, 4.5], [7.5, 8.5]]) - sr0.coordinates.data[:] = coords src0 = sr0.inject(f0, Float(1.)) @@ -1105,12 +1248,11 @@ def test_inject_subdomain(self): ['p_sr0rp_sr0xrp_sr0y'], 'p_sr0rp_sr0xrp_sr0y') - def test_inject_subdomain_sinc(self): + def test_inject_subdomain_sinc(self, grid, sinc_coords): """ Check sinc injection into a Function defined on a SubDomain functions as expected. """ - grid = Grid(shape=(11, 11), extent=(10., 10.)) sd0 = SD6(grid=grid) sd1 = SD7(grid=grid) @@ -1120,9 +1262,7 @@ def test_inject_subdomain_sinc(self): sr0 = SparseFunction(name='sr0', grid=grid, npoint=5, interpolation='sinc', r=2) - coords = np.array([[2.5, 6.5], [3.5, 4.5], [6.0, 6.], [5.5, 4.5], [4.5, 6.]]) - - sr0.coordinates.data[:] = coords + sr0.coordinates.data[:] = sinc_coords src0 = sr0.inject(f0, Float(1.)) src1 = sr0.inject(f1, Float(1.)) @@ -1139,12 +1279,10 @@ def test_inject_subdomain_sinc(self): @pytest.mark.xfail(reason="OOB issue") @pytest.mark.parallel(mode=4) - def test_interpolate_subdomain_mpi(self, mode): + def test_interpolate_subdomain_mpi(self, mode, grid, coords): """ Test interpolation off of a Function defined on a SubDomain with MPI. """ - - grid = Grid(shape=(11, 11), extent=(10., 10.)) sd2 = SD2(grid=grid) sd3 = SD3(grid=grid) sd4 = SD4(grid=grid) @@ -1166,10 +1304,6 @@ def test_interpolate_subdomain_mpi(self, mode): sr2 = SparseFunction(name='sr2', grid=grid, npoint=8) sr3 = SparseFunction(name='sr3', grid=grid, npoint=8) - coords = np.array([[2.5, 1.5], [4.5, 2.], [8.5, 4.], - [0.5, 6.], [7.5, 4.], [5.5, 5.5], - [1.5, 4.5], [7.5, 8.5]]) - sr0.coordinates.data[:] = coords sr1.coordinates.data[:] = coords sr2.coordinates.data[:] = coords @@ -1206,12 +1340,10 @@ def test_interpolate_subdomain_mpi(self, mode): assert np.all(np.isclose(sr3.data, [0., 0.])) @pytest.mark.parallel(mode=4) - def test_inject_subdomain_mpi(self, mode): + def test_inject_subdomain_mpi(self, mode, grid, coords): """ Test injection into a Function defined on a SubDomain with MPI. """ - - grid = Grid(shape=(11, 11), extent=(10., 10.)) sd2 = SD2(grid=grid) sd3 = SD3(grid=grid) sd4 = SD4(grid=grid) @@ -1224,10 +1356,6 @@ def test_inject_subdomain_mpi(self, mode): sr0 = SparseFunction(name='sr0', grid=grid, npoint=8) - coords = np.array([[2.5, 1.5], [4.5, 2.], [8.5, 4.], - [0.5, 6.], [7.5, 4.], [5.5, 5.5], - [1.5, 4.5], [7.5, 8.5]]) - sr0.coordinates.data[:] = coords src0 = sr0.inject(f0, Float(1.)) @@ -1277,18 +1405,3 @@ def test_inject_subdomain_mpi(self, mode): assert data1 == None # noqa assert data2 == None # noqa assert data3 == None # noqa - - -def test_wrong_coords(): - grid = Grid(shape=(11, 11)) - s = SparseFunction(name='src', npoint=1, grid=grid) - s2 = SparseFunction(name='src2', npoint=1, grid=grid) - u = Function(name='u', grid=grid) - - with pytest.raises(ValueError) as vinfo: - s.inject(u, expr=s2) - assert "Interpolation/injection with" in str(vinfo.value) - - with pytest.raises(ValueError) as vinfo: - s.interpolate(u + s2) - assert "Interpolation/injection with" in str(vinfo.value)