Skip to content

Commit

Permalink
Merge pull request #1948 from devitocodes/patch-conddim-sparse
Browse files Browse the repository at this point in the history
compiler: Support interpolation with user-provided implicit dims
  • Loading branch information
mloubout authored Jul 17, 2022
2 parents 2f96bfb + ea8b9c3 commit d9df3b1
Show file tree
Hide file tree
Showing 3 changed files with 75 additions and 20 deletions.
45 changes: 30 additions & 15 deletions devito/operations/interpolators.py
Original file line number Diff line number Diff line change
Expand Up @@ -6,7 +6,7 @@

from devito.logger import warning
from devito.symbolics import retrieve_function_carriers, indexify, INT
from devito.tools import powerset, flatten, prod
from devito.tools import as_tuple, powerset, flatten, prod
from devito.types import (ConditionalDimension, Dimension, DefaultDimension, Eq, Inc,
Evaluable, Symbol, SubFunction)

Expand Down Expand Up @@ -168,7 +168,8 @@ def _interpolation_coeffs(self):
A = A.subs(reference_cell)
return A.inv().T * p

def _interpolation_indices(self, variables, offset=0, field_offset=0):
def _interpolation_indices(self, variables, offset=0, field_offset=0,
implicit_dims=None):
"""
Generate interpolation indices for the DiscreteFunctions in ``variables``.
"""
Expand All @@ -192,21 +193,22 @@ def _interpolation_indices(self, variables, offset=0, field_offset=0):
idx_subs.append(mapper)

# Temporaries for the position
temps = [Eq(v, k, implicit_dims=self.sfunction.dimensions)
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=self.sfunction.dimensions)
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=self.sfunction.dimensions)
implicit_dims=implicit_dims)
for p, c in zip(self.sfunction._point_symbols,
self.sfunction._coordinate_bases(field_offset))])

return idx_subs, temps

def interpolate(self, expr, offset=0, increment=False, self_subs={}):
def interpolate(self, expr, offset=0, increment=False, self_subs={},
implicit_dims=None):
"""
Generate equations interpolating an arbitrary expression into ``self``.
Expand All @@ -218,7 +220,13 @@ def interpolate(self, expr, offset=0, increment=False, self_subs={}):
Additional offset from the boundary.
increment: bool, optional
If True, generate increments (Inc) rather than assignments (Eq).
implicit_dims : Dimension or list of Dimension, optional
An ordered list of Dimensions that do not explicitly appear in the
interpolation expression, but that should be honored when constructing
the operator.
"""
implicit_dims = as_tuple(implicit_dims) + self.sfunction.dimensions

def callback():
# Derivatives must be evaluated before the introduction of indirect accesses
try:
Expand All @@ -233,18 +241,18 @@ def callback():
# TODO: handle each variable staggereing spearately
field_offset = variables[0].origin
# List of indirection indices for all adjacent grid points
idx_subs, temps = self._interpolation_indices(variables, offset,
field_offset=field_offset)
idx_subs, temps = self._interpolation_indices(
variables, offset, field_offset=field_offset, implicit_dims=implicit_dims
)

# Substitute coordinate base symbols into the interpolation coefficients
args = [_expr.xreplace(v_sub) * b.xreplace(v_sub)
for b, v_sub in zip(self._interpolation_coeffs, idx_subs)]

# Accumulate point-wise contributions into a temporary
rhs = Symbol(name='sum', dtype=self.sfunction.dtype)
summands = [Eq(rhs, 0., implicit_dims=self.sfunction.dimensions)]
summands.extend([Inc(rhs, i, implicit_dims=self.sfunction.dimensions)
for i in args])
summands = [Eq(rhs, 0., implicit_dims=implicit_dims)]
summands.extend([Inc(rhs, i, implicit_dims=implicit_dims) for i in args])

# Write/Incr `self`
lhs = self.sfunction.subs(self_subs)
Expand All @@ -254,7 +262,7 @@ def callback():

return Interpolation(expr, offset, increment, self_subs, self, callback)

def inject(self, field, expr, offset=0):
def inject(self, field, expr, offset=0, implicit_dims=None):
"""
Generate equations injecting an arbitrary expression into a field.
Expand All @@ -266,7 +274,13 @@ def inject(self, field, expr, offset=0):
Injected expression.
offset : int, optional
Additional offset from the boundary.
implicit_dims : Dimension or list of Dimension, optional
An ordered list of Dimensions that do not explicitly appear in the
injection expression, but that should be honored when constructing
the operator.
"""
implicit_dims = as_tuple(implicit_dims) + self.sfunction.dimensions

def callback():
# Derivatives must be evaluated before the introduction of indirect accesses
try:
Expand All @@ -280,12 +294,13 @@ def callback():
# Need to get origin of the field in case it is staggered
field_offset = field.origin
# List of indirection indices for all adjacent grid points
idx_subs, temps = self._interpolation_indices(variables, offset,
field_offset=field_offset)
idx_subs, temps = self._interpolation_indices(
variables, offset, field_offset=field_offset, implicit_dims=implicit_dims
)

# Substitute coordinate base symbols into the interpolation coefficients
eqns = [Inc(field.xreplace(vsub), _expr.xreplace(vsub) * b,
implicit_dims=self.sfunction.dimensions)
implicit_dims=implicit_dims)
for b, vsub in zip(self._interpolation_coeffs, idx_subs)]

return temps + eqns
Expand Down
8 changes: 6 additions & 2 deletions devito/types/sparse.py
Original file line number Diff line number Diff line change
Expand Up @@ -848,7 +848,7 @@ def interpolate(self, expr, offset=0, u_t=None, p_t=None, increment=False):
increment=increment,
self_subs=subs)

def inject(self, field, expr, offset=0, u_t=None, p_t=None):
def inject(self, field, expr, offset=0, u_t=None, p_t=None, implicit_dims=None):
"""
Generate equations injecting an arbitrary expression into a field.
Expand All @@ -864,14 +864,18 @@ def inject(self, field, expr, offset=0, u_t=None, p_t=None):
Time index at which the interpolation is performed.
p_t : expr-like, optional
Time index at which the result of the interpolation is stored.
implicit_dims : Dimension or list of Dimension, optional
An ordered list of Dimensions that do not explicitly appear in the
injection expression, but that should be honored when constructing
the operator.
"""
# Apply optional time symbol substitutions to field and expr
if u_t is not None:
field = field.subs({field.time_dim: u_t})
if p_t is not None:
expr = expr.subs({self.time_dim: p_t})

return super(SparseTimeFunction, self).inject(field, expr, offset=offset)
return super().inject(field, expr, offset=offset, implicit_dims=implicit_dims)

# Pickling support
_pickle_kwargs = AbstractSparseTimeFunction._pickle_kwargs +\
Expand Down
42 changes: 39 additions & 3 deletions tests/test_dimension.py
Original file line number Diff line number Diff line change
Expand Up @@ -5,9 +5,10 @@
import pytest

from conftest import assert_blocking, skipif, opts_tiling
from devito import (ConditionalDimension, Grid, Function, TimeFunction, SparseFunction, # noqa
Eq, Operator, Constant, Dimension, SubDimension, switchconfig,
SubDomain, Lt, Le, Gt, Ge, Ne, Buffer)
from devito import (ConditionalDimension, Grid, Function, TimeFunction, # noqa
SparseFunction, SparseTimeFunction, Eq, Operator, Constant,
Dimension, SubDimension, switchconfig, SubDomain, Lt, Le,
Gt, Ge, Ne, Buffer, sin)
from devito.ir.iet import (Conditional, Expression, Iteration, FindNodes,
retrieve_iteration_tree)
from devito.symbolics import indexify, retrieve_functions, IntDiv
Expand Down Expand Up @@ -1178,6 +1179,41 @@ def test_affiness(self):
iterations = [i for i in FindNodes(Iteration).visit(op) if i.dim is not time]
assert all(i.is_Affine for i in iterations)

def test_sparse_time_function(self):
nt = 20

shape = (21, 21, 21)
origin = (0., 0., 0.)
spacing = (1., 1., 1.)
extent = tuple([d * (s - 1) for s, d in zip(shape, spacing)])
grid = Grid(shape=shape, extent=extent, origin=origin)
time = grid.time_dim
x, y, z = grid.dimensions

p = TimeFunction(name="p", grid=grid, time_order=2, space_order=4, save=nt)

# Place source in the middle of the grid
src_coords = np.empty((1, len(shape)), dtype=np.float32)
src_coords[0, :] = [o + d * (s-1)//2 for o, d, s in zip(origin, spacing, shape)]
src = SparseTimeFunction(name='src', grid=grid, npoint=1, nt=nt)
src.data[:] = 1.
src.coordinates.data[:] = src_coords[:]

cd = ConditionalDimension(name="cd", parent=time,
condition=And(Ge(time, 1), Le(time, 10)))

src_term = src.inject(field=p.forward, expr=src*time, implicit_dims=cd)

op = Operator(src_term)

op.apply(time_m=1, time_M=nt-2, dt=1.0)

assert np.all(p.data[0] == 0)
# Note the endpoint of the range is 12 because we inject at p.forward
assert all(p.data[i].sum() == i - 1 for i in range(1, 12))
assert all(p.data[i, 10, 10, 10] == i - 1 for i in range(1, 12))
assert all(np.all(p.data[i] == 0) for i in range(12, 20))


class TestMashup(object):

Expand Down

0 comments on commit d9df3b1

Please sign in to comment.