diff --git a/devito/operations/interpolators.py b/devito/operations/interpolators.py index 8bc8535e59..a7fa4d76be 100644 --- a/devito/operations/interpolators.py +++ b/devito/operations/interpolators.py @@ -219,12 +219,22 @@ class WeightedInterpolator(GenericInterpolator): Represent an Interpolation operation on a SparseFunction that is separable in space, meaning the coefficients are defined for each Dimension separately and multiplied at a given point: `w[x, y] = wx[x] * wy[y]` + + The base cell index of each sparse point, and (for the coordinate-based + schemes) the interpolation weights, are tabulated on the host in fp64 by + `_arg_defaults` and handed to the kernel as int32/fp SubFunctions, so the + generated C never computes `floor((c - o)/h)` on fp32. """ _name = 'weighted' - def __init__(self, sfunction): + def __init__(self, sfunction, shifts=()): self.sfunction = sfunction + # Every shift set the interpolator has been asked to produce tables + # for. Persisted with the parent SparseFunction via `__rkwargs__` so + # a pickled/rebuilt interpolator (decoupled workers) knows which + # tables to emit in `_arg_defaults`. + self._shifts_used = set(tuple(s) if s else None for s in shifts) @property def grid(self): @@ -234,6 +244,14 @@ def grid(self): def r(self): return self.sfunction.r + @cached_property + def _coeff_dtype(self): + # Weights are real even for complex-valued sparse fields. + dtype = np.dtype(self.sfunction.dtype) + if np.issubdtype(dtype, np.complexfloating): + return np.finfo(dtype).dtype.type + return dtype.type + @memoized_meth def _weights(self, subdomain=None, shifts=None): raise NotImplementedError @@ -320,10 +338,88 @@ def _augment_implicit_dims(self, implicit_dims, extras=None): def _coeff_temps(self, implicit_dims, shifts=None): return [] - def _positions(self, implicit_dims, shifts=None): + @memoized_meth + def _generate_gridpoints(self, key): + """ + Create the `(npoint, ndim)` int32 table of base cell indices for a + given shift set. `key` is either None (plain, non-staggered) or a + tuple of per-dim shifts. + """ + # Record that the caller has emitted a table for this shift set, so + # `_arg_defaults` can regenerate it on the fly. + self._shifts_used.add(key) + + name = f'{self.sfunction.name}_gp{_shift_tag(as_list(key))}' + sfdim = self.sfunction._sparse_dim + ddim = CustomDimension(f'{name}d', 0, self.grid.dim - 1, + self.grid.dim, sfdim) + return Gridpoints(name=name, dtype=np.int32, + shape=(self.sfunction.npoint, self.grid.dim), + dimensions=(sfdim, ddim), space_order=0, + alias=self.sfunction.alias, + parent=self.sfunction) + + def _gridpoints(self, shifts=None): + return self._generate_gridpoints(tuple(shifts) if shifts else None) + + def _coeffs(self, shifts=None): + """The per-Dimension weight SubFunctions for a given shift set.""" + return () + + def _floor_positions(self, implicit_dims, shifts=None): + """ + Position temporaries computed in the kernel as `floor((c - o - s)/h)`. + Only for schemes whose tables aren't tabulated on the host. + """ return [Eq(v, INT(floor(k)), implicit_dims=implicit_dims) for k, v in self.sfunction._position_map(shifts=shifts).items()] + def _positions(self, implicit_dims, shifts=None): + gp = self._gridpoints(shifts=shifts) + ddim = gp.dimensions[-1] + return [Eq(p, gp._subs(ddim, di), implicit_dims=implicit_dims) + for (di, p) in enumerate( + self.sfunction._pos_symbols(shifts=shifts))] + + def _coeff_data(self, coords, grid, shifts, spacing, origin): + """ + Mapper `M : weight SubFunction -> tabulated data` for one shift set. + """ + return {} + + def _arg_defaults(self, coords=None, sfunc=None, origin=None): + """ + Fill the gridpoints/weights tables from the already-scattered `coords` + handed in by `SparseFunction._arg_defaults`. The tables are regenerated + from the persisted shift set, so a pickled/rebuilt interpolator + (decoupled workers) still emits data for every table the operator was + compiled with. + """ + if coords is None or sfunc is None: + raise ValueError("No coordinates or sparse function provided") + + # Fp64 grid geometry -- avoids fp32 rounding on cell boundaries. + grid = sfunc.grid + spacing = np.array([as_fp64_decimal(h) for h in grid.spacing]) + origin = np.array([as_fp64_decimal(o) + for o in (origin or grid.origin)]) + + args = {} + for key in self._shifts_used or {None}: + shifts = as_list(key) + gp = self._gridpoints(shifts=key) + + tables = {gp: _cell_indices(coords, grid, shifts, spacing, origin)} + tables.update(self._coeff_data(coords, grid, shifts, spacing, origin)) + + for f, data in tables.items(): + args[f.name] = data + # Bounds for each table's dimensions, matching the computed data + for d, s in zip(f.dimensions, data.shape, strict=True): + args.update(d._arg_defaults(_min=0, size=s)) + + return args + def _interp_idx(self, variables, implicit_dims=None, subdomain=None, shifts=None): """ @@ -546,12 +642,11 @@ def _shift_values(shifts, grid, spacing): class _HostTable(SubFunction): """SubFunction populated on the host by the parent interpolator's - `_arg_defaults` from the already-scattered coordinates, exactly the way - sinc's precomputed weights are populated. + `_arg_defaults` from the already-scattered coordinates. `parent` links the table to its SparseFunction so the base `SubFunction._arg_values` routing triggers `SparseFunction._arg_defaults` - -> `LinearInterpolator._arg_defaults`. The link is *not* preserved by + -> `WeightedInterpolator._arg_defaults`. The link is *not* preserved by pickle (parent is not in `__rkwargs__`) so decoupled workers don't ship the sfunction back through every table.""" @@ -564,7 +659,7 @@ class Gridpoints(_HostTable): class Coeffs(_HostTable): - """Per-dim `(1 - frac, frac)` interpolation weights, shape ``(npoint, 2)``.""" + """Per-dim interpolation weights, shape ``(npoint, 2*r)``.""" def _resolved_geometry(grid, kwargs): @@ -591,18 +686,38 @@ def _cell_indices(coords, grid, shifts, spacing, origin): origin)).astype(np.int32) +def _cell_fractions(coords, grid, shifts, j, spacing, origin): + """Fractional cell offset along dim `j` from the base index returned by + `_cell_indices`, in fp64.""" + pos = _positions_fp64(coords, grid, shifts, spacing, origin)[:, j] + return pos - np.floor(pos) + + def _linear_weights(coords, grid, shifts, j, dtype, spacing, origin): """`(1 - frac, frac)` linear interpolation weights along dim `j`; `frac` is the fractional cell offset from the base index returned by `_cell_indices`.""" - pos = _positions_fp64(coords, grid, shifts, spacing, origin) - frac = pos[:, j] - np.floor(pos[:, j]) - data = np.empty((pos.shape[0], 2), dtype=dtype) + frac = _cell_fractions(coords, grid, shifts, j, spacing, origin) + data = np.empty((frac.size, 2), dtype=dtype) data[:, 0] = 1.0 - frac data[:, 1] = frac return data +def _sinc_weights(coords, grid, shifts, j, dtype, spacing, origin, r, b): + """Kaiser windowed sinc weights along dim `j` for the `2*r` grid points + surrounding the base index returned by `_cell_indices`. `b` is the Kaiser + window parameter for radius `r`.""" + frac = _cell_fractions(coords, grid, shifts, j, spacing, origin) + b0 = i0(b) + data = np.zeros((frac.size, 2*r), dtype=dtype) + for ri in range(2*r): + rpos = ri - r + 1 - frac + num = i0(b*np.sqrt(1 - (rpos/r)**2)) + data[:, ri] = num / b0 * np.sinc(rpos) + return data + + class LinearInterpolator(WeightedInterpolator): """ Linear (bilinear/trilinear) interpolator. @@ -615,49 +730,21 @@ class LinearInterpolator(WeightedInterpolator): _name = 'linear' - def __init__(self, sfunction, shifts=()): - super().__init__(sfunction) - # Every shift set the interpolator has been asked to produce tables - # for. Persisted with the parent SparseFunction via `__rkwargs__` so - # a pickled/rebuilt interpolator (decoupled workers) knows which - # tables to emit in `_arg_defaults`. - self._shifts_used = set(tuple(s) if s else None for s in shifts) - - @cached_property - def _coeff_dtype(self): - # Weights are real even for complex-valued sparse fields. - dtype = np.dtype(self.sfunction.dtype) - if np.issubdtype(dtype, np.complexfloating): - return np.finfo(dtype).dtype.type - return dtype.type - @memoized_meth def _generate_coeffs(self, key): - """Create the ``(gridpoints, coeffs_per_dim)`` SubFunction tuple for - a given shift set. ``key`` is either ``None`` (plain, non-staggered) - or a tuple of per-dim shifts. Mirrors sinc's ``interpolation_coeffs`` - cached_property but keyed on ``shifts``.""" + """Create the per-Dimension weight SubFunctions for a given shift set. + ``key`` is either ``None`` (plain, non-staggered) or a tuple of per-dim + shifts.""" # Record that the caller has emitted tables for this shift set, so # `_arg_defaults` can regenerate them on the fly. self._shifts_used.add(key) - shifts = as_list(key) - tag = _shift_tag(shifts) + tag = _shift_tag(as_list(key)) sfname = self.sfunction.name sfdim = self.sfunction._sparse_dim - # Gridpoints: `(npoint, ndim)` int32 base cell index per sparse point. - gp_name = f'{sfname}_gp{tag}' - ddim = CustomDimension(f'{gp_name}d', 0, self.grid.dim - 1, - self.grid.dim, sfdim) - gp = Gridpoints(name=gp_name, dtype=np.int32, - shape=(self.sfunction.npoint, self.grid.dim), - dimensions=(sfdim, ddim), space_order=0, - alias=self.sfunction.alias, - parent=self.sfunction) - # Per-dim linear weights: `(npoint, 2)` holding `(1 - frac, frac)`. - coeffs = tuple( + return tuple( Coeffs(name=f'{sfname}_w{d.name}{tag}', dtype=self._coeff_dtype, shape=(self.sfunction.npoint, 2), @@ -667,23 +754,8 @@ def _generate_coeffs(self, key): for d, r in zip(self._gdims, self._cdim, strict=True) ) - return gp, coeffs - - def _gridpoints(self, shifts=None): - return self._generate_coeffs(tuple(shifts) if shifts else None)[0] - def _coeffs(self, shifts=None): - return self._generate_coeffs(tuple(shifts) if shifts else None)[1] - - def _positions(self, implicit_dims, shifts=None): - gp = self._gridpoints(shifts=shifts) - ddim = gp.dimensions[-1] - return [Eq(p, gp._subs(ddim, di), implicit_dims=implicit_dims) - for (di, p) in enumerate( - self.sfunction._pos_symbols(shifts=shifts))] - - def _coeff_temps(self, implicit_dims, shifts=None): - return [] + return self._generate_coeffs(tuple(shifts) if shifts else None) @memoized_meth def _weights(self, subdomain=None, shifts=None): @@ -694,39 +766,11 @@ def _weights(self, subdomain=None, shifts=None): for (rd, w) in zip(rdims, coeffs, strict=True) ]) - def _arg_defaults(self, coords=None, sfunc=None, origin=None): - """Fill the gridpoints/coeffs tables from the already-scattered - ``coords`` handed in by ``SparseFunction._arg_defaults``. Mirrors - sinc's `_arg_defaults`: regenerates the tables from the persisted - shift set, so a pickled/rebuilt interpolator (decoupled workers) - still emits data for every table the operator was compiled with.""" - if coords is None or sfunc is None: - raise ValueError("No coordinates or sparse function provided") - - # Fp64 grid geometry -- avoids fp32 rounding on cell boundaries. - grid = sfunc.grid - spacing = np.array([as_fp64_decimal(h) for h in grid.spacing]) - origin = np.array([as_fp64_decimal(o) - for o in (origin or grid.origin)]) - - args = {} - for key in self._shifts_used or {None}: - shifts = as_list(key) - gp, coeffs = self._generate_coeffs(key) - - # Tabulated data (int32 cell indices + fp linear weights per dim). - args[gp.name] = _cell_indices(coords, grid, shifts, spacing, origin) - for i, w in enumerate(coeffs): - args[w.name] = _linear_weights( - coords, grid, shifts, i, w.dtype, spacing, origin - ) - - # Bounds for each table's dimensions, matching the computed data. - for f in (gp, *coeffs): - for d, s in zip(f.dimensions, args[f.name].shape, strict=True): - args.update(d._arg_defaults(_min=0, size=s)) - - return args + def _coeff_data(self, coords, grid, shifts, spacing, origin): + return { + w: _linear_weights(coords, grid, shifts, i, w.dtype, spacing, origin) + for i, w in enumerate(self._coeffs(shifts=shifts)) + } class NearestInterpolator(LinearInterpolator): @@ -752,11 +796,8 @@ def _rdim(self, subdomain=None, shifts=None): def _weights(self, subdomain=None, shifts=None): return sympy.S.One - def _gridpoints(self, shifts=None): - return self._generate_coeffs(tuple(shifts) if shifts else None)[0] - def _coeffs(self, shifts=None): - return [] + return () class PrecomputedInterpolator(WeightedInterpolator): @@ -774,12 +815,19 @@ class PrecomputedInterpolator(WeightedInterpolator): def _positions(self, implicit_dims, shifts=None): if self.sfunction.gridpoints_data is None: - return super()._positions(implicit_dims, shifts=shifts) + # Only the coordinates are known, and the user-provided coefficients + # are tied to the cell index the kernel derives from them + return self._floor_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(shifts=shifts).items()] + def _arg_defaults(self, **kwargs): + # Gridpoints and coefficients are user-provided SubFunctions of the + # PrecomputedSparseFunction, so there is nothing to tabulate + return {} + @property def interpolation_coeffs(self): return self.sfunction.interpolation_coeffs @@ -794,7 +842,7 @@ def _weights(self, subdomain=None, shifts=None): for mapper in mappers]) -class SincInterpolator(PrecomputedInterpolator): +class SincInterpolator(WeightedInterpolator): """ Hicks windowed sinc interpolation scheme. @@ -803,6 +851,9 @@ class SincInterpolator(PrecomputedInterpolator): https://library.seg.org/doi/10.1190/1.1451454 + Like the linear scheme, the gridpoints and the `2*r` windowed sinc weights + per Dimension are tabulated on the host in fp64 (see `_arg_defaults`), so + that the weights are always those of the cell the kernel indexes into. """ _name = 'sinc' @@ -812,52 +863,49 @@ class SincInterpolator(PrecomputedInterpolator): 4: 4.14, 5: 5.26, 6: 6.40, 7: 7.51, 8: 8.56, 9: 9.56, 10: 10.64} - def __init__(self, sfunction): + def __init__(self, sfunction, shifts=()): 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) + super().__init__(sfunction, shifts=shifts) - @cached_property - def interpolation_coeffs(self): - coeffs = [] + @memoized_meth + def _generate_coeffs(self, key): + """Create the per-Dimension windowed sinc weight SubFunctions for a + given shift set. ``key`` is either ``None`` (plain, non-staggered) or + a tuple of per-dim shifts.""" + # Record that the caller has emitted tables for this shift set, so + # `_arg_defaults` can regenerate them on the fly. + self._shifts_used.add(key) + + tag = _shift_tag(as_list(key)) shape = (self.sfunction.npoint, 2 * self.r) - for r in self._cdim: - dimensions = (self.sfunction._sparse_dim, r) - sf = SubFunction(name=f"wsinc{r.name}", dtype=self.sfunction.dtype, - shape=shape, dimensions=dimensions, - space_order=0, alias=self.sfunction.alias, - parent=None) - coeffs.append(sf) - return tuple(coeffs) + + return tuple( + Coeffs(name=f'wsinc{r.name}{tag}', dtype=self._coeff_dtype, + shape=shape, dimensions=(self.sfunction._sparse_dim, r), + space_order=0, alias=self.sfunction.alias, + parent=self.sfunction) + for r in self._cdim + ) + + def _coeffs(self, shifts=None): + return self._generate_coeffs(tuple(shifts) if shifts else None) @memoized_meth 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) + for (rd, w) in zip(rdims, self._coeffs(shifts=shifts), strict=True) ]) - def _arg_defaults(self, coords=None, sfunc=None, origin=None): - args = {} + def _coeff_data(self, coords, grid, shifts, spacing, origin): b = self._b_table[self.r] - b0 = i0(b) - if coords is None or sfunc is None: - raise ValueError("No coordinates or sparse function provided") - # Coords to indices - coords = coords / np.array(sfunc.grid.spacing) - coords = coords - np.floor(coords) - - # Precompute sinc - for j in range(len(self._gdims)): - 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[j].name] = data - - return args + return { + w: _sinc_weights(coords, grid, shifts, i, w.dtype, spacing, + origin, self.r, b) + for i, w in enumerate(self._coeffs(shifts=shifts)) + } diff --git a/examples/userapi/06_sparse_operations.ipynb b/examples/userapi/06_sparse_operations.ipynb index 4e2d08417d..286ad0297b 100644 --- a/examples/userapi/06_sparse_operations.ipynb +++ b/examples/userapi/06_sparse_operations.ipynb @@ -482,8 +482,8 @@ "name": "stdout", "output_type": "stream", "text": [ - "Eq(posx, (int)floor((-o_x + s_coords(p_s, 0))/h_x))\n", - "Eq(posy, (int)floor((-o_y + s_coords(p_s, 1))/h_y))\n", + "Eq(posx, s_gp(p_s, 0))\n", + "Eq(posy, s_gp(p_s, 1))\n", "Eq(sums, 0.0)\n", "Inc(sums, wsincrp_sx(p_s, rp_sx + 3)*wsincrp_sy(p_s, rp_sy + 3)*f(t, rp_sx + posx, rp_sy + posy))\n", "Eq(s(time, p_s), sums)\n" diff --git a/tests/test_interpolation.py b/tests/test_interpolation.py index 635857bdea..af1f1954a3 100644 --- a/tests/test_interpolation.py +++ b/tests/test_interpolation.py @@ -818,6 +818,81 @@ def test_sinc_accuracy(self, r, tol): assert err_sinc < err_lin assert err_lin > 0.01 + def test_sinc_position_tabulated(self): + """ + Sinc must read its gridpoints from the fp64 host-tabulated table rather + than recomputing `floor((c - o)/h)` in the kernel, so that the cell it + indexes is the one its weights were computed for. + """ + grid = Grid(shape=(30, 30), extent=(2.9, 2.9)) + assert grid.dtype is np.float32 + + sf = SparseTimeFunction(name='sf', grid=grid, npoint=1, nt=2, + interpolation='sinc', r=2) + u = TimeFunction(name='u', grid=grid, space_order=4, time_order=1) + sf.coordinates.data[0, :] = 0.6999999990000001 + sf.data[:] = 1.0 + + op = Operator(sf.inject(field=u.forward, expr=sf)) + code = str(op) + assert 'floor' not in code + assert 'o_x' not in code and 'o_y' not in code + + op.apply(time_M=0) + # fp64 `floor(0.6999999990000001/0.1)` is 6, so the stencil spans 5:9 + assert np.all(u.data[1, :5, :] == 0.0) + assert np.all(u.data[1, 9:, :] == 0.0) + assert np.argmax(np.abs(u.data[1])) == 7 * u.data.shape[2] + 7 + + def test_sinc_origin(self): + """ + The sinc weights must be tabulated from the same origin-corrected + position as the gridpoints, so injecting at `c` on a grid with origin + `0` matches injecting at `c + shift` on a grid with origin `shift`. + """ + shape, spacing, coord = (41, 41), (10., 10.), 133. + extent = tuple((s - 1) * h for s, h in zip(shape, spacing, strict=True)) + + def run(origin): + grid = Grid(shape=shape, extent=extent, origin=origin) + u = TimeFunction(name='u', grid=grid, space_order=8) + src = SparseTimeFunction(name='src', grid=grid, npoint=1, nt=2, + interpolation='sinc', r=4) + src.coordinates.data[0, :] = coord + origin[0] + src.data[:] = 1. + Operator(src.inject(field=u.forward, expr=src)).apply(time_M=0) + return u.data[1] + + assert np.allclose(run((0., 0.)), run((-25., -25.)), rtol=1e-6) + + @pytest.mark.parametrize('stagg', ['x', 'y', '(x, y)']) + def test_sinc_staggered_weights(self, stagg): + """ + Injection into a staggered field must use the weights of the half-cell + shifted position, i.e. be equivalent to injecting into an unstaggered + field with the coordinates shifted by the same half cell. + """ + grid = Grid(shape=(21, 21), extent=(20., 20.)) + x, y = grid.dimensions # noqa + staggered = as_tuple(eval(stagg)) + coord = (9.3, 11.7) + + def run(stagg, shift): + a = Function(name='a', grid=grid, space_order=8, staggered=stagg) + p = SparseFunction(name='p', grid=grid, npoint=1, + interpolation='sinc', r=3) + p.coordinates.data[0, :] = [c - s + for c, s in zip(coord, shift, strict=True)] + p.data[:] = 1. + Operator(p.inject(a, expr=p)).apply() + return np.array(a.data) + + shift = [h/2 if d in staggered else 0 + for d, h in zip(grid.dimensions, grid.spacing, strict=True)] + assert np.allclose(run(staggered, [0, 0]), run(NODE, shift), rtol=1e-6) + # Sanity check: the shift is what makes the two match + assert not np.allclose(run(staggered, [0, 0]), run(NODE, [0, 0])) + # --------------------------------------------------------------------------- # Matrix sparse function interpolation / injection diff --git a/tests/test_mpi.py b/tests/test_mpi.py index d3fc001146..d5ad76323a 100644 --- a/tests/test_mpi.py +++ b/tests/test_mpi.py @@ -2541,6 +2541,30 @@ def test_interpolation_dup(self, mode): assert np.all(sf.data == [1.5, 2.5, 2.5, 3.5][grid.distributor.myrank]) + @pytest.mark.parallel(mode=4) + def test_interpolation_sinc(self, mode): + """ + Sinc interpolation reads both its gridpoints and its weights from + host-tabulated SubFunctions, whose local size is that of the scattered + coordinates rather than the global `npoint`. + """ + grid = Grid(shape=(4, 4), extent=(3.0, 3.0)) + + f = Function(name='f', grid=grid, space_order=2) + f.data[:] = np.array([[1, 1, 1, 1], [2, 2, 2, 2], + [3, 3, 3, 3], [4, 4, 4, 4]]) + + # On-node points, for which the windowed sinc weights are a delta + coords = np.array([(1.0, 1.0), (1.0, 2.0), (2.0, 1.0), (2.0, 2.0)]) + sf = SparseFunction(name='sf', grid=grid, npoint=len(coords), + coordinates=coords, interpolation='sinc', r=2) + sf.data[:] = 0. + + op = Operator(sf.interpolate(expr=f)) + op.apply() + + assert np.allclose(sf.data, [2., 2., 3., 3.][grid.distributor.myrank]) + @pytest.mark.parallel(mode=2) def test_subsampling(self, mode): grid = Grid(shape=(40,))