diff --git a/devito/operations/interpolators.py b/devito/operations/interpolators.py index dcfadfbe372..411d41bc194 100644 --- a/devito/operations/interpolators.py +++ b/devito/operations/interpolators.py @@ -1,11 +1,12 @@ from abc import ABC, abstractmethod +from itertools import product import sympy -import numpy as np from cached_property import cached_property -from devito.symbolics import retrieve_function_carriers -from devito.tools import as_tuple, powerset, flatten, prod +from devito.finite_differences.elementary import floor +from devito.symbolics import retrieve_function_carriers, INT +from devito.tools import as_tuple, flatten, prod from devito.types import (ConditionalDimension, Eq, Inc, Evaluable, Symbol) __all__ = ['LinearInterpolator', 'PrecomputedInterpolator'] @@ -106,15 +107,12 @@ def interpolate(self, *args, **kwargs): pass -class LinearInterpolator(GenericInterpolator): +class WeightedInterpolation(GenericInterpolator): """ - Concrete implementation of GenericInterpolator implementing a Linear interpolation - scheme, i.e. Bilinear for 2D and Trilinear for 3D problems. - - Parameters - ---------- - sfunction: The SparseFunction that this Interpolator operates on. + Represent an Interpolation operation on a SparseFunction that is separable + in space, meaning hte coefficient are defined for each Dimension separately + and multiplied at a given point: `w[x, y] = wx[x] * wy[y]` """ def __init__(self, sfunction): @@ -124,86 +122,46 @@ def __init__(self, sfunction): def grid(self): return self.sfunction.grid - @cached_property - def _interpolation_coeffs(self): - """ - Symbolic expression for the coefficients for sparse point interpolation - according to: + @property + def _weights(self): + raise NotImplementedError - https://en.wikipedia.org/wiki/Bilinear_interpolation. + @property + def _psym(self): + return self.sfunction._point_symbols - Returns - ------- - Matrix of coefficient expressions. - """ - # Grid indices corresponding to the corners of the cell ie x1, y1, z1 - indices1 = tuple(sympy.symbols('%s1' % d) for d in self.grid.dimensions) - indices2 = tuple(sympy.symbols('%s2' % d) for d in self.grid.dimensions) - # 1, x1, y1, z1, x1*y1, ... - indices = list(powerset(indices1)) - indices[0] = (1,) - point_sym = list(powerset(self.sfunction._point_symbols)) - point_sym[0] = (1,) - # 1, px. py, pz, px*py, ... - A = [] - ref_A = [np.prod(ind) for ind in indices] - # Create the matrix with the same increment order as the point increment - for i in self.sfunction._point_increments: - # substitute x1 by x2 if increment in that dimension - subs = dict((indices1[d], indices2[d] if i[d] == 1 else indices1[d]) - for d in range(len(i))) - A += [[1] + [a.subs(subs) for a in ref_A[1:]]] - - A = sympy.Matrix(A) - # Coordinate values of the sparse point - p = sympy.Matrix([[np.prod(ind)] for ind in point_sym]) - - # reference cell x1:0, x2:h_x - left = dict((a, 0) for a in indices1) - right = dict((b, dim.spacing) for b, dim in zip(indices2, self.grid.dimensions)) - reference_cell = {**left, **right} - # Substitute in interpolation matrix - A = A.subs(reference_cell) - return A.inv().T * p + @property + def _gdim(self): + return self.grid.dimensions - def _interpolation_indices(self, variables, offset=0, field_offset=0, - implicit_dims=None): - """ - Generate interpolation indices for the DiscreteFunctions in ``variables``. - """ - index_matrix, points = self.sfunction._index_matrix(offset) + def implicit_dims(self, implicit_dims): + return as_tuple(implicit_dims) + self.sfunction.dimensions - idx_subs = [] - for i, idx in enumerate(index_matrix): - # Introduce ConditionalDimension so that we don't go OOB - mapper = {} - for j, d in zip(idx, self.grid.dimensions): - p = points[j] - lb = sympy.And(p >= d.symbolic_min - self.sfunction._radius, - evaluate=False) - ub = sympy.And(p <= d.symbolic_max + self.sfunction._radius, - evaluate=False) - condition = sympy.And(lb, ub, evaluate=False) - mapper[d] = ConditionalDimension(p.name, self.sfunction._sparse_dim, - condition=condition, indirect=True) + @property + def r(self): + return self.sfunction.r - # Track Indexed substitutions - idx_subs.append(mapper) + @property + def _interp_points(self): + return range(-self.r+1, self.r+1) - # Temporaries for the position - temps = [Eq(v, k, implicit_dims=implicit_dims) - for k, v in self.sfunction._position_map.items()] - # Temporaries for the indirection dimensions - temps.extend([Eq(v, k.subs(self.sfunction._position_map), - implicit_dims=implicit_dims) - for k, v in points.items()]) - # Temporaries for the coefficients - temps.extend([Eq(p, c.subs(self.sfunction._position_map), - implicit_dims=implicit_dims) - for p, c in zip(self.sfunction._point_symbols, - self.sfunction._coordinate_bases(field_offset))]) + @property + def _nd_points(self): + return product(self._interp_points, repeat=self.grid.dim) - return idx_subs, temps + @property + def _interpolation_coeffs(self): + coeffs = {} + for p in self._nd_points: + coeffs[p] = prod([self._weights[d][i] for (d, i) in zip(self._gdim, p)]) + return list(coeffs.values()) + + def _coeff_temps(self, implicit_dims): + return [] + + def _positions(self, implicit_dims): + return [Eq(v, k, implicit_dims=implicit_dims) + for k, v in self.sfunction._position_map.items()] def subs_coords(self, _expr, *idx_subs): return [_expr.xreplace(v_sub) * b.xreplace(v_sub) @@ -214,8 +172,47 @@ def subs_coords_eq(self, field, _expr, *idx_subs, implicit_dims=None): implicit_dims=implicit_dims) for b, vsub in zip(self._interpolation_coeffs, idx_subs)] - def implicit_dims(self, implicit_dims): - return as_tuple(implicit_dims) + self.sfunction.dimensions + def _interpolation_indices(self, variables, offset=0, field_offset=0, + implicit_dims=None): + """ + Generate interpolation indices for the DiscreteFunctions in ``variables``. + """ + idx_subs = [] + points = {d: [] for d in self._gdim} + mapper = {d: [] for d in self._gdim} + + # Positon map and temporaries for it + pmap = self.sfunction._coordinate_indices + + # Temporaries for the position + temps = self._positions(implicit_dims) + + # Coefficient symbol expression + temps.extend(self._coeff_temps(implicit_dims)) + + # Create positions and indices temporaries/indirections + for ((di, d), pos) in zip(enumerate(self._gdim), pmap): + for (ri, r) in enumerate(self._interp_points): + p = Symbol(name='ii_%s_%s_%d' % (self.sfunction.name, d.name, ri)) + points[d].append(p) + # Conditionals to avoid OOB + lb = sympy.And(p >= d.symbolic_min - self.r, evaluate=False) + ub = sympy.And(p <= d.symbolic_max + self.r, evaluate=False) + condition = sympy.And(lb, ub, evaluate=False) + mapper[d].append(ConditionalDimension(p.name, self.sfunction._sparse_dim, + condition=condition, indirect=True)) + temps.extend([Eq(p, pos + r, implicit_dims=implicit_dims)]) + + # Substitution mapper + for p in self._nd_points: + # Apply mapper to each variable with origin correction before the + # Dimensions get replaced + subs = {v: v.subs({k: c[pi] - v.origin.get(k, 0) + for ((k, c), pi) in zip(mapper.items(), p)}) + for v in variables} + idx_subs.append(subs) + + return idx_subs, temps def interpolate(self, expr, offset=0, increment=False, self_subs={}, implicit_dims=None): @@ -257,7 +254,6 @@ def callback(): # Substitute coordinate base symbols into the interpolation coefficients args = self.subs_coords(_expr, *idx_subs) - # Accumulate point-wise contributions into a temporary rhs = Symbol(name='sum', dtype=self.sfunction.dtype) summands = [Eq(rhs, 0., implicit_dims=implicit_dims)] @@ -317,66 +313,56 @@ def callback(): return Injection(field, expr, offset, self, callback) -class PrecomputedInterpolator(LinearInterpolator): +class LinearInterpolator(WeightedInterpolation): - def __init__(self, obj): - self.sfunction = obj + """ + Concrete implementation of GenericInterpolator implementing a Linear interpolation + scheme, i.e. Bilinear for 2D and Trilinear for 3D problems. - @property - def r(self): - return self.obj.r + Parameters + ---------- + sfunction: The SparseFunction that this Interpolator operates on. + """ - def _interpolation_indices(self, variables, offset=0, field_offset=0, - implicit_dims=None): - """ - Generate interpolation indices for the DiscreteFunctions in ``variables``. - """ - if self.sfunction.gridpoints is None: - return super()._interpolation_indices(variables, offset=offset, - field_offset=field_offset, - implicit_dims=implicit_dims) + @cached_property + def _weights(self): + return {d: [1 - p/d.spacing, p/d.spacing] + for (d, p) in zip(self._gdim, self._psym)} - index_matrix, points, shifts = self.sfunction._index_matrix(offset) + def _coeff_temps(self, implicit_dims): + pmap = self.sfunction._position_map.values() + return [Eq(self._psym[d], pos - d.spacing*INT(floor(pos/d.spacing)), + implicit_dims=implicit_dims) + for (d, pos) in zip(self._gdim, pmap)] - idx_subs = [] - coeffs = self._interpolation_coeffs - dt, it = coeffs.dimensions[1:] - for i, idx in enumerate(index_matrix): - # Introduce ConditionalDimension so that we don't go OOB - mapper = {} - for j, (di, d) in zip(idx, enumerate(self.grid.dimensions)): - p = points[j] - lb = sympy.And(p >= d.symbolic_min - self.sfunction.r // 2, - evaluate=False) - ub = sympy.And(p <= d.symbolic_max + self.sfunction.r // 2, - evaluate=False) - condition = sympy.And(lb, ub, evaluate=False) - mapper[d] = ConditionalDimension(p.name, self.sfunction._sparse_dim, - condition=condition, indirect=True) - mapper[coeffs._subs(dt, di)] = coeffs.subs({dt: di, it: shifts[i][di]}) - # Track Indexed substitutions - idx_subs.append(mapper) - # Temporaries for the indirection dimensions - temps = [Eq(v, k, implicit_dims=implicit_dims) for k, v in points.items()] +class PrecomputedInterpolator(WeightedInterpolation): - return idx_subs, temps + def _positions(self, implicit_dims): + if self.sfunction.gridpoints is None: + return [Eq(v, k, implicit_dims=implicit_dims) + for k, v in self.sfunction._position_map.items()] + # No position temp as we have directly the gridpoints + return [] @property - def _interpolation_coeffs(self): + def _interp_points(self): + return range(-self.r//2 + 1, self.r//2 + 1) + + @property + def _icoeffs(self): return self.sfunction.interpolation_coeffs @property - def _interpolation_coeffsp(self): - d = self.sfunction.interpolation_coeffs.dimensions[1] - return prod([self.sfunction.interpolation_coeffs._subs(d, i) - for (i, _) in enumerate(self.sfunction.grid.dimensions)]) + def _idim(self): + return self.sfunction.interpolation_coeffs.dimensions[-1] - def subs_coords(self, _expr, *idx_subs): - b = self._interpolation_coeffsp - return [_expr.xreplace(v_sub) * b.xreplace(v_sub) for v_sub in idx_subs] + @property + def _ddim(self): + return self.sfunction.interpolation_coeffs.dimensions[1] - def subs_coords_eq(self, field, _expr, *idx_subs, implicit_dims=None): - b = self._interpolation_coeffsp - return [Inc(field.xreplace(vsub), _expr.xreplace(vsub) * b.xreplace(vsub), - implicit_dims=implicit_dims) for vsub in idx_subs] + @cached_property + def _weights(self): + return {d: [self._icoeffs.subs({self._ddim: di, self._idim: k}) + for k in self._interp_points] + for (di, d) in enumerate(self._gdim)} diff --git a/devito/tools/algorithms.py b/devito/tools/algorithms.py index 0021a6f608e..f7edde22f55 100644 --- a/devito/tools/algorithms.py +++ b/devito/tools/algorithms.py @@ -72,4 +72,5 @@ def toposort(data): if item not in ordered]) if len(processed) != len(set(flatten(data) + flatten(data.values()))): raise ValueError("A cyclic dependency exists amongst %r" % data) + return processed diff --git a/devito/tools/data_structures.py b/devito/tools/data_structures.py index bd89b4ebe73..10a1d906723 100644 --- a/devito/tools/data_structures.py +++ b/devito/tools/data_structures.py @@ -66,6 +66,9 @@ def __getnewargs_ex__(self): # objects with varying number of attributes return (tuple(self), dict(self.__dict__)) + def get(self, key, val): + return self._getters.get(key, val) + class ReducerMap(MultiDict): diff --git a/devito/types/basic.py b/devito/types/basic.py index 97495234ea6..57463864e58 100644 --- a/devito/types/basic.py +++ b/devito/types/basic.py @@ -945,7 +945,8 @@ def origin(self): f(x) : origin = 0 f(x + hx/2) : origin = hx/2 """ - return tuple(r - d for d, r in zip(self.dimensions, self.indices_ref)) + return DimensionTuple(*(r - d for d, r in zip(self.dimensions, self.indices_ref)), + getters=self.grid.dimensions) @property def dimensions(self): @@ -1199,8 +1200,10 @@ def indexify(self, indices=None, subs=None): # Indices after substitutions indices = [sympy.sympify((a - o).xreplace(s)) for a, o, s in zip(self.args, self.origin, subs)] + indices = [i.xreplace({k: sympy.Integer(k) for k in i.atoms(sympy.Float)}) for i in indices] + return self.indexed[indices] def __getitem__(self, index): diff --git a/devito/types/dense.py b/devito/types/dense.py index e7d488e90c8..2a9579f4cfe 100644 --- a/devito/types/dense.py +++ b/devito/types/dense.py @@ -1471,6 +1471,13 @@ def __padding_setup__(self, **kwargs): def _halo_exchange(self): return + @property + def origin(self): + """ + SubFunction have zero origin + """ + return DimensionTuple(*(0 for _ in range(self.ndim)), getters=self.dimensions) + def _arg_values(self, **kwargs): if self.name in kwargs: raise RuntimeError("`%s` is a SubFunction, so it can't be assigned " diff --git a/devito/types/sparse.py b/devito/types/sparse.py index f2878cd90d5..7e7392626e0 100644 --- a/devito/types/sparse.py +++ b/devito/types/sparse.py @@ -24,7 +24,7 @@ DynamicDimension) from devito.types.basic import Symbol from devito.types.equation import Eq, Inc -from devito.types.utils import IgnoreDimSort +from devito.types.utils import IgnoreDimSort, DimensionTuple __all__ = ['SparseFunction', 'SparseTimeFunction', 'PrecomputedSparseFunction', @@ -101,7 +101,7 @@ def __subfunc_setup__(self, key, suffix): dimensions = (self.indices[self._sparse_position], Dimension(name='d'), Dimension(name='i')) - shape = (self.npoint, self.grid.dim, self.r) + shape = (self.npoint, self.grid.dim, self._radius) if key is None: # Fallback to default behaviour @@ -170,7 +170,8 @@ def _mpitype(self): @property def _smpitype(self): - sfuncs = [getattr(self, s) for s in self._sub_functions] + sfuncs = [getattr(self, s) for s in self._sub_functions + if getattr(self, s) is not None] return {s: MPI._typedict[np.dtype(s.dtype).char] for s in sfuncs} @property @@ -194,14 +195,26 @@ def _subfunc_names(self): @cached_property def _point_symbols(self): """Symbol for coordinate value in each dimension of the point.""" - return tuple(Symbol(name='p%s' % d, dtype=self.dtype) - for d in self.grid.dimensions) + return DimensionTuple(*(Symbol(name='p%s' % d, dtype=self.dtype) + for d in self.grid.dimensions), + getters=self.grid.dimensions) @cached_property def _position_map(self): """ - Symbols map for the position of the sparse points relative to the grid + Symbols map for the physical position of the sparse points relative to the grid origin. + """ + symbols = [Symbol(name='pos%s' % d, dtype=self.dtype) + for d in self.grid.dimensions] + return OrderedDict([(c - o, p) for p, c, o in zip(symbols, + self._coordinate_symbols, + self.grid.origin_symbols)]) + + @cached_property + def _coordinate_indices(self): + """ + Symbol for each grid index according to the coordinates. Notes ----- @@ -213,24 +226,9 @@ def _position_map(self): the position. We mitigate this problem by computing the positions individually (hence the need for a position map). """ - symbols = [Symbol(name='pos%s' % d, dtype=self.dtype) - for d in self.grid.dimensions] - return OrderedDict([(c - o, p) for p, c, o in zip(symbols, - self._coordinate_symbols, - self.grid.origin_symbols)]) - - @cached_property - def _point_increments(self): - """Index increments in each dimension for each point symbol.""" - return tuple(product(range(self.r+1), repeat=self.grid.dim)) - - @cached_property - def _coordinate_indices(self): - """Symbol for each grid index according to the coordinates.""" - return tuple([INT(floor((c - o) / i.spacing)) - for c, o, i in zip(self._coordinate_symbols, - self.grid.origin_symbols, - self.grid.dimensions[:self.grid.dim])]) + return tuple([INT(floor(p / i.spacing)) + for p, i in zip(self._position_map.values(), + self.grid.dimensions[:self.grid.dim])]) def _coordinate_bases(self, field_offset): """Symbol for the base coordinates of the reference grid point.""" @@ -265,10 +263,14 @@ def inject(self, *args, **kwargs): """ return self.interpolator.inject(*args, **kwargs) + @cached_property + def _point_increments(self): + """Index increments in each dimension for each point symbol.""" + return tuple(product(range(self.r+1), repeat=self.grid.dim)) + @cached_property def _point_support(self): - return np.array(tuple(product(range(-self._radius + 1, self._radius + 1), - repeat=self.grid.dim))) + return np.array(self._point_increments) @property def _support(self): @@ -278,7 +280,7 @@ def _support(self): """ max_shape = np.array(self.grid.shape).reshape(1, self.grid.dim) minmax = lambda arr: np.minimum(max_shape, np.maximum(0, arr)) - return np.stack([minmax(self.gridpoints_data + s) for s in self._point_support], + return np.stack([minmax(self._coords_indices + s) for s in self._point_support], axis=2) @property @@ -774,15 +776,13 @@ def dist_origin(self): return self._dist_origin @property - def gridpoints(self): + def _coords_indices(self): if self.coordinates.data is None: raise ValueError("No coordinates attached to this SparseFunction") return ( np.floor(self.coordinates.data._local - self.grid.origin) / self.grid.spacing ).astype(np.int32) - gridpoints_data = gridpoints - def guard(self, expr=None, offset=0): """ Generate guarded expressions, that is expressions that are evaluated @@ -1036,10 +1036,11 @@ class PrecomputedSparseFunction(AbstractSparseFunction): uses `*args` to (re-)create the dimension arguments of the symbolic object. """ - _sub_functions = ('gridpoints', 'interpolation_coeffs') + _sub_functions = ('gridpoints', 'coordinates', 'interpolation_coeffs') __rkwargs__ = (AbstractSparseFunction.__rkwargs__ + - ('r', 'gridpoints_data', 'interpolation_coeffs_data')) + ('r', 'gridpoints_data', 'coordinates_data', + 'interpolation_coeffs_data')) def __init_finalize__(self, *args, **kwargs): super().__init_finalize__(*args, **kwargs) @@ -1061,19 +1062,19 @@ def __init_finalize__(self, *args, **kwargs): # Specifying only `npoints` is acceptable; this will require the user # to setup the coordinates data later on + npoint = kwargs.get('npoint', None) if self.npoint and coordinates is None and gridpoints is None: - gridpoints = np.zeros((self.npoint, self.grid.dim)) + coordinates = np.zeros((npoint, self.grid.dim)) if coordinates is not None: - # Convert to gridpoints - if isinstance(coordinates, SubFunction): - raise ValueError("`coordinates` only accepted as array") - loc = np.floor((coordinates - self.grid.origin) / self.grid.spacing) - self._gridpoints = self.__subfunc_setup__(loc.astype(int), 'gridpoints') + self._coordinates = self.__subfunc_setup__(coordinates, 'coords') + self._gridpoints = None + self._dist_origin = {self._coordinates: self.grid.origin_offset} else: assert gridpoints is not None + self._coordinates = None self._gridpoints = self.__subfunc_setup__(gridpoints, 'gridpoints') - self._dist_origin = {self._gridpoints: self.grid.origin_ioffset} + self._dist_origin = {self._gridpoints: self.grid.origin_ioffset} # Setup the interpolation coefficients. These are compulsory interpolation_coeffs = kwargs.get('interpolation_coeffs', @@ -1093,6 +1094,10 @@ def __init_finalize__(self, *args, **kwargs): self.interpolator = PrecomputedInterpolator(self) + @property + def r(self): + return self._radius + @cached_property def _point_increments(self): """Index increments in each dimension for each point symbol.""" @@ -1100,8 +1105,16 @@ def _point_increments(self): @cached_property def _point_support(self): - return np.array(tuple(product(range(-self.r // 2 + 1, self.r // 2 + 1), - repeat=self.grid.dim))) + return np.array(self._point_increments) + + @property + def _coords_indices(self): + if self.gridpoints is not None: + return self.gridpoints.data._local + else: + return ( + np.floor(self.coordinates_data - self.grid.origin) / self.grid.spacing + ).astype(np.int32) @property def gridpoints(self): @@ -1109,7 +1122,21 @@ def gridpoints(self): @property def gridpoints_data(self): - return self.gridpoints.data._local + try: + return self.gridpoints.data._local + except AttributeError: + return None + + @property + def coordinates(self): + return self._coordinates + + @property + def coordinates_data(self): + try: + return self.coordinates.data._local + except AttributeError: + return None @property def dist_origin(self): @@ -1128,29 +1155,35 @@ def interpolation_coeffs_data(self): def _coordinate_symbols(self): """Symbol representing the coordinate values in each dimension.""" p_dim = self.indices[self._sparse_position] + if self.gridpoints is not None: + return tuple([self.gridpoints.indexify((p_dim, di)) * d.spacing + o + for ((di, d), o) in zip(enumerate(self.grid.dimensions), + self.grid.origin)]) + return tuple([self.coordinates.indexify((p_dim, i)) for i in range(self.grid.dim)]) - @memoized_meth - def _index_matrix(self, offset): - # Note about the use of *memoization* - # Since this method is called by `_interpolation_indices`, using - # memoization avoids a proliferation of symbolically identical - # ConditionalDimensions for a given set of indirection indices - - # List of indirection indices for all adjacent grid points - ddim = self._gridpoints.dimensions[1] - index_matrix = [tuple(self._gridpoints._subs(ddim, d) + ii + offset - for (ii, d) in zip(inc, range(self.grid.dim))) - for inc in self._point_increments] - shifts = [tuple(ii + offset for ii in inc) - for inc in self._point_increments] - # A unique symbol for each indirection index - indices = filter_ordered(flatten(index_matrix)) - points = OrderedDict([(p, Symbol(name='ii_%s_%d' % (self.name, i))) - for i, p in enumerate(indices)]) + @cached_property + def _coordinate_indices(self): + """ + Symbol for each grid index according to the coordinates. - return index_matrix, points, shifts + Notes + ----- + The expression `(coord - origin)/spacing` could also be computed in the + mathematically equivalent expanded form `coord/spacing - + origin/spacing`. This particular form is problematic when a sparse + point is in close proximity of the grid origin, since due to a larger + machine precision error it may cause a +-1 error in the computation of + the position. We mitigate this problem by computing the positions + individually (hence the need for a position map). + """ + if self.gridpoints is not None: + ddim = self.gridpoints.dimensions[-1] + return tuple([self.gridpoints._subs(ddim, di) for di in range(self.grid.dim)]) + return tuple([INT(floor(p / i.spacing)) + for p, i in zip(self._position_map.keys(), + self.grid.dimensions[:self.grid.dim])]) class PrecomputedSparseTimeFunction(AbstractSparseTimeFunction, @@ -2140,11 +2173,11 @@ def _dist_scatter(self, data=None): # The implementation in AbstractSparseFunction now relies on us # having a .coordinates property, which we don't have. - def _arg_apply(self, dataobj, alias=None): + def _arg_apply(self, *dataobj, alias=None): key = alias if alias is not None else self if isinstance(key, AbstractSparseFunction): # Gather into `self.data` - key._dist_gather(self._C_as_ndarray(dataobj)) + key._dist_gather(self._C_as_ndarray(dataobj[0])) elif self.grid.distributor.nprocs > 1: raise NotImplementedError("Don't know how to gather data from an " "object of type `%s`" % type(key)) diff --git a/tests/test_dle.py b/tests/test_dle.py index 23e23a6fe7c..e24b67941f8 100644 --- a/tests/test_dle.py +++ b/tests/test_dle.py @@ -937,16 +937,7 @@ def test_parallel_prec_inject(self): 'par-collapse-ncores': 1})) iterations = FindNodes(Iteration).visit(op0) assert not iterations[0].pragmas - assert 'omp for collapse(1) schedule(dynamic,chunk_size)'\ - in iterations[1].pragmas[0].value - - op1 = Operator(eqns, opt=('advanced', {'openmp': True, - 'par-collapse-ncores': 1, - 'par-collapse-work': 1})) - iterations = FindNodes(Iteration).visit(op1) - assert not iterations[0].pragmas - assert 'omp for collapse(1) schedule(dynamic,chunk_size)'\ - in iterations[1].pragmas[0].value + assert 'omp for' in iterations[1].pragmas[0].value class TestNestedParallelism(object): diff --git a/tests/test_mpi.py b/tests/test_mpi.py index 1337f947539..f1a93d363fc 100644 --- a/tests/test_mpi.py +++ b/tests/test_mpi.py @@ -523,16 +523,15 @@ def test_precomputed_sparse(self, r): assert sf1.npoint == 1 assert sf2.npoint == 1 - assert np.all(sf1.gridpoints.data.shape == (1, 2)) + assert np.all(sf1.coordinates.data.shape == (1, 2)) assert np.all(sf2.gridpoints.data.shape == (1, 2)) - assert np.all(sf1.gridpoints_data == sf2.gridpoints_data) + assert np.all(sf1._coords_indices == sf2.gridpoints_data) assert np.all(sf1.interpolation_coeffs.shape == (1, 2, r)) assert np.all(sf2.interpolation_coeffs.shape == (1, 2, r)) u = Function(name="u", grid=grid, space_order=r) u._data_with_outhalo[:] = 1 Operator(sf2.interpolate(u))() - print(sf2.data) assert np.all(sf2.data == 4) Operator(sf1.interpolate(u))() assert np.all(sf1.data == 4) diff --git a/tests/test_pickle.py b/tests/test_pickle.py index 1307ca3e9ab..31e99c69858 100644 --- a/tests/test_pickle.py +++ b/tests/test_pickle.py @@ -97,13 +97,19 @@ def test_sparse_function(): assert sf.npoint == new_sf.npoint -def test_precomputed_sparse_function(): - grid = Grid(shape=(10, 10)) +@pytest.mark.parametrize('mode', ['coordinates', 'gridpoints']) +def test_precomputed_sparse_function(mode): + grid = Grid(shape=(11, 11)) + + coords = [(0., 0.), (.5, .5), (.7, .2)] + gridpoints = [(0, 0), (6, 6), (8, 3)] + keys = {'coordinates': coords, 'gridpoints': gridpoints} + kw = {mode: keys[mode]} + othermode = 'coordinates' if mode == 'gridpoints' else 'gridpoints' sf = PrecomputedSparseTimeFunction( name='sf', grid=grid, r=2, npoint=3, nt=5, - coordinates=[(0., 0.), (1., 1.), (2., 2.)], - interpolation_coeffs=np.ndarray(shape=(3, 2, 2)), + interpolation_coeffs=np.ndarray(shape=(3, 2, 2)), **kw ) sf.data[2, 1] = 5. @@ -117,7 +123,9 @@ def test_precomputed_sparse_function(): assert np.all(sf.interpolation_coeffs.data == new_sf.interpolation_coeffs.data) # coordinates, since they were given, should also have been pickled - assert np.all(sf.coordinates.data == new_sf.coordinates.data) + assert np.all(getattr(sf, mode).data == getattr(new_sf, mode).data) + assert getattr(sf, othermode) is None + assert getattr(new_sf, othermode) is None assert sf._radius == new_sf._radius == 2 assert sf.space_order == new_sf.space_order