def coordinate_indices(self):
"""Symbol for each grid index according to the coordinates"""
indices = (x, y, z)
return tuple([
INT(Function('floor')(c / i.spacing))
for c, i in zip(self.coordinate_symbols, indices[:self.ndim])
])
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.
"""
expr = indexify(expr)
field = indexify(field)
p, _ = self.gridpoints.indices
dim_subs = []
coeffs = []
for i, d in enumerate(self.grid.dimensions):
rd = DefaultDimension(name="r%s" % d.name, default_value=self.r)
dim_subs.append((d, INT(rd + self.gridpoints[p, i])))
coeffs.append(self.interpolation_coeffs[p, i, rd])
rhs = prod(coeffs) * expr
field = field.subs(dim_subs)
return [Eq(field, field + rhs.subs(dim_subs))]
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).
"""
expr = indexify(expr)
p, _, _ = self.interpolation_coeffs.indices
dim_subs = []
coeffs = []
for i, d in enumerate(self.grid.dimensions):
rd = DefaultDimension(name="r%s" % d.name, default_value=self.r)
dim_subs.append((d, INT(rd + self.gridpoints[p, i])))
coeffs.append(self.interpolation_coeffs[p, i, rd])
# Apply optional time symbol substitutions to lhs of assignment
lhs = self.subs(self_subs)
rhs = prod(coeffs) * expr.subs(dim_subs)
return [Eq(lhs, lhs + rhs)]
def coordinate_indices(self):
"""Symbol for each grid index according to the coordinates"""
indices = self.grid.dimensions
return tuple([
INT(sympy.Function('floor')((c - o) / i.spacing))
for c, o, i in zip(self.coordinate_symbols, self.grid.origin,
indices[:self.grid.dim])
])
def make_grid_gets(expr):
mapper = {}
indexeds = retrieve_indexed(expr)
data_carriers = [i for i in indexeds if i.base.function.from_YASK]
for i in data_carriers:
name = namespace['code-grid-name'](i.base.function.name)
args = [ListInitializer([INT(make_grid_gets(j)) for j in i.indices])]
mapper[i] = make_sharedptr_funcall(namespace['code-grid-get'], args, name)
return expr.xreplace(mapper)
def make_grid_accesses(node, yk_grid_objs):
"""
Construct a new Iteration/Expression based on ``node``, in which all
:class:`types.Indexed` accesses have been converted into YASK grid
accesses.
"""
def make_grid_gets(expr):
mapper = {}
indexeds = retrieve_indexed(expr)
data_carriers = [i for i in indexeds if i.base.function.from_YASK]
for i in data_carriers:
args = [
ListInitializer([INT(make_grid_gets(j)) for j in i.indices])
]
mapper[i] = make_sharedptr_funcall(
namespace['code-grid-get'], args,
yk_grid_objs[i.base.function.name])
return expr.xreplace(mapper)
mapper = {}
for i, e in enumerate(FindNodes(Expression).visit(node)):
if e.is_ForeignExpression:
continue
lhs, rhs = e.expr.args
# RHS translation
rhs = make_grid_gets(rhs)
# LHS translation
if e.write.from_YASK:
args = [rhs]
args += [
ListInitializer([INT(make_grid_gets(i)) for i in lhs.indices])
]
call = namespace['code-grid-add' if e.
is_Increment else 'code-grid-put']
handle = make_sharedptr_funcall(call, args,
yk_grid_objs[e.write.name])
processed = ForeignExpression(handle,
e.dtype,
is_Increment=e.is_Increment)
else:
# Writing to a scalar temporary
processed = e._rebuild(expr=e.expr.func(lhs, rhs))
mapper.update({e: processed})
return Transformer(mapper).visit(node)
def make_grid_accesses(node):
"""
Construct a new Iteration/Expression based on ``node``, in which all
:class:`types.Indexed` accesses have been converted into YASK grid
accesses.
"""
def make_grid_gets(expr):
mapper = {}
indexeds = retrieve_indexed(expr)
data_carriers = [i for i in indexeds if i.base.function.from_YASK]
for i in data_carriers:
name = namespace['code-grid-name'](i.base.function.name)
args = [ListInitializer([INT(make_grid_gets(j)) for j in i.indices])]
mapper[i] = make_sharedptr_funcall(namespace['code-grid-get'], args, name)
return expr.xreplace(mapper)
mapper = {}
for i, e in enumerate(FindNodes(Expression).visit(node)):
lhs, rhs = e.expr.args
# RHS translation
rhs = make_grid_gets(rhs)
# LHS translation
if e.write.from_YASK:
name = namespace['code-grid-name'](e.write.name)
args = [rhs]
args += [ListInitializer([INT(make_grid_gets(i)) for i in lhs.indices])]
handle = make_sharedptr_funcall(namespace['code-grid-put'], args, name)
processed = Element(c.Statement(ccode(handle)))
else:
# Writing to a scalar temporary
processed = Expression(e.expr.func(lhs, rhs))
mapper.update({e: processed})
return Transformer(mapper).visit(node)
