def interpolate(self, expr, offset=0, increment=False, self_subs={}):
"""
Generate equations interpolating an arbitrary expression into ``self``.
Parameters
----------
expr : expr-like
Input expression to interpolate.
offset : int, optional
Additional offset from the boundary.
increment: bool, optional
If True, generate increments (Inc) rather than assignments (Eq).
"""
variables = list(retrieve_function_carriers(expr))
# List of indirection indices for all adjacent grid points
idx_subs, eqns = self._interpolation_indices(variables, offset)
# Substitute coordinate base symbols into the interpolation coefficients
args = [
expr.subs(v_sub) * b.subs(v_sub)
for b, v_sub in zip(self._interpolation_coeffs, idx_subs)
]
# Accumulate point-wise contributions into a temporary
rhs = Scalar(name='sum', dtype=self.dtype)
summands = [Eq(rhs, 0.)] + [Inc(rhs, i) for i in args]
# Write/Incr `self`
lhs = self.subs(self_subs)
last = [Inc(lhs, rhs)] if increment else [Eq(lhs, rhs)]
return eqns + summands + last
def inject(self, field, expr, offset=0):
"""
Generate equations injecting an arbitrary expression into a field.
Parameters
----------
field : Function
Input field into which the injection is performed.
expr : expr-like
Injected expression.
offset : int, optional
Additional offset from the boundary.
"""
# 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
pass
variables = list(retrieve_function_carriers(expr)) + [field]
# List of indirection indices for all adjacent grid points
idx_subs, temps = self._interpolation_indices(variables, offset)
# Substitute coordinate base symbols into the interpolation coefficients
eqns = [
Inc(field.xreplace(vsub),
expr.xreplace(vsub) * b,
implicit_dims=self.dimensions)
for b, vsub in zip(self._interpolation_coeffs, idx_subs)
]
return temps + eqns
def callback():
# 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
variables = list(retrieve_function_carriers(_expr)) + [field]
# 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)
# Substitute coordinate base symbols into the interpolation coefficients
eqns = [
Inc(field.xreplace(vsub),
_expr.xreplace(vsub) * b,
implicit_dims=self.sfunction.dimensions)
for b, vsub in zip(self._interpolation_coeffs, idx_subs)
]
return temps + eqns
def inject(self, field, expr, offset=0):
"""
Generate equations injecting an arbitrary expression into a field.
Parameters
----------
field : Function
Input field into which the injection is performed.
expr : expr-like
Injected expression.
offset : int, optional
Additional offset from the boundary.
"""
variables = list(retrieve_function_carriers(expr)) + [field]
# List of indirection indices for all adjacent grid points
idx_subs, eqns = self._interpolation_indices(variables, offset)
# Substitute coordinate base symbols into the interpolation coefficients
eqns.extend([
Inc(field.subs(vsub),
expr.subs(vsub) * b)
for b, vsub in zip(self._interpolation_coeffs, idx_subs)
])
return eqns
def interpolate(self, expr, offset=0, increment=False, self_subs={}):
"""
Generate equations interpolating an arbitrary expression into ``self``.
Parameters
----------
expr : expr-like
Input expression to interpolate.
offset : int, optional
Additional offset from the boundary.
increment: bool, optional
If True, generate increments (Inc) rather than assignments (Eq).
"""
variables = list(retrieve_function_carriers(expr))
# List of indirection indices for all adjacent grid points
idx_subs, temps = self._interpolation_indices(variables, offset)
# Substitute coordinate base symbols into the interpolation coefficients
args = [expr.subs(v_sub) * b.subs(v_sub)
for b, v_sub in zip(self._interpolation_coeffs, idx_subs)]
# Accumulate point-wise contributions into a temporary
rhs = Scalar(name='sum', dtype=self.dtype)
summands = [Eq(rhs, 0., implicit_dims=self.dimensions)]
summands.extend([Inc(rhs, i, implicit_dims=self.dimensions) for i in args])
# Write/Incr `self`
lhs = self.subs(self_subs)
last = [Inc(lhs, rhs)] if increment else [Eq(lhs, rhs)]
return temps + summands + last
def callback():
# 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
variables = list(retrieve_function_carriers(_expr))
# Need to get origin of the field in case it is staggered
# 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)
# 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])
# Write/Incr `self`
lhs = self.sfunction.subs(self_subs)
last = [Inc(lhs, rhs)] if increment else [Eq(lhs, rhs)]
return temps + summands + last
def __init__(self, expressions):
mapper = DefaultOrderedDict(lambda: DefaultOrderedDict(AccessTuple))
for e in expressions:
for i in retrieve_function_carriers(e.rhs):
mapper[i.function][e].reads.append(i)
mapper[e.lhs.function][e].write = e.lhs
super().__init__([(f, AccessValue(f, mapper[f])) for f in mapper])
def guard(self, expr=None, offset=0):
"""
Generate guarded expressions, that is expressions that are evaluated
by an Operator only if certain conditions are met. The introduced
condition, here, is that all grid points in the support of a sparse
value must fall within the grid domain (i.e., *not* on the halo).
Parameters
----------
expr : expr-like, optional
Input expression, from which the guarded expression is derived.
If not specified, defaults to ``self``.
offset : int, optional
Relax the guard condition by introducing a tolerance offset.
"""
_, points = self._index_matrix(offset)
# Guard through ConditionalDimension
conditions = {}
for d, idx in zip(self.grid.dimensions, self._coordinate_indices):
p = points[idx]
lb = sympy.And(p >= d.symbolic_min - offset, evaluate=False)
ub = sympy.And(p <= d.symbolic_max + offset, evaluate=False)
conditions[p] = sympy.And(lb, ub, evaluate=False)
condition = sympy.And(*conditions.values(), evaluate=False)
cd = ConditionalDimension("%s_g" % self._sparse_dim,
self._sparse_dim,
condition=condition)
if expr is None:
out = self.indexify().xreplace({self._sparse_dim: cd})
else:
functions = {
f
for f in retrieve_function_carriers(expr)
if f.is_SparseFunction
}
out = indexify(expr).xreplace(
{f._sparse_dim: cd
for f in functions})
# Temporaries for the position
temps = [
Eq(v, k, implicit_dims=self.dimensions)
for k, v in self._position_map.items()
]
# Temporaries for the indirection dimensions
temps.extend([
Eq(v, k.subs(self._position_map), implicit_dims=self.dimensions)
for k, v in points.items() if v in conditions
])
return out, temps
def interpolate(self, expr, offset=0, increment=False, self_subs={}):
"""
Generate equations interpolating an arbitrary expression into ``self``.
Parameters
----------
expr : expr-like
Input expression to interpolate.
offset : int, optional
Additional offset from the boundary.
increment: bool, optional
If True, generate increments (Inc) rather than assignments (Eq).
"""
# 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
pass
variables = list(retrieve_function_carriers(expr))
# List of indirection indices for all adjacent grid points
idx_subs, temps = self._interpolation_indices(variables, offset)
# 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 = Scalar(name='sum', dtype=self.dtype)
summands = [Eq(rhs, 0., implicit_dims=self.dimensions)]
summands.extend(
[Inc(rhs, i, implicit_dims=self.dimensions) for i in args])
# Write/Incr `self`
lhs = self.subs(self_subs)
last = [Inc(lhs, rhs)] if increment else [Eq(lhs, rhs)]
return temps + summands + last
def guard(self, expr=None, offset=0):
"""
Generate guarded expressions, that is expressions that are evaluated
by an Operator only if certain conditions are met. The introduced
condition, here, is that all grid points in the support of a sparse
value must fall within the grid domain (i.e., *not* on the halo).
Parameters
----------
expr : expr-like, optional
Input expression, from which the guarded expression is derived.
If not specified, defaults to ``self``.
offset : int, optional
Relax the guard condition by introducing a tolerance offset.
"""
_, points = self._index_matrix(offset)
# Guard through ConditionalDimension
conditions = {}
for d, idx in zip(self.grid.dimensions, self._coordinate_indices):
p = points[idx]
lb = sympy.And(p >= d.symbolic_min - offset, evaluate=False)
ub = sympy.And(p <= d.symbolic_max + offset, evaluate=False)
conditions[p] = sympy.And(lb, ub, evaluate=False)
condition = sympy.And(*conditions.values(), evaluate=False)
cd = ConditionalDimension("%s_g" % self._sparse_dim, self._sparse_dim,
condition=condition)
if expr is None:
out = self.indexify().xreplace({self._sparse_dim: cd})
else:
functions = {f for f in retrieve_function_carriers(expr)
if f.is_SparseFunction}
out = indexify(expr).xreplace({f._sparse_dim: cd for f in functions})
# Temporaries for the indirection dimensions
temps = [Eq(v, k, implicit_dims=self.dimensions)
for k, v in points.items() if v in conditions]
return out, temps
def inject(self, field, expr, offset=0):
"""
Generate equations injecting an arbitrary expression into a field.
Parameters
----------
field : Function
Input field into which the injection is performed.
expr : expr-like
Injected expression.
offset : int, optional
Additional offset from the boundary.
"""
variables = list(retrieve_function_carriers(expr)) + [field]
# List of indirection indices for all adjacent grid points
idx_subs, temps = self._interpolation_indices(variables, offset)
# Substitute coordinate base symbols into the interpolation coefficients
eqns = [Inc(field.subs(vsub), expr.subs(vsub) * b, implicit_dims=self.dimensions)
for b, vsub in zip(self._interpolation_coeffs, idx_subs)]
return temps + eqns
