def inject_kernel(Vf, Vc):
hierarchy, level = utils.get_level(Vc.ufl_domain())
cache = hierarchy._shared_data_cache["transfer_kernels"]
coordinates = Vf.ufl_domain().coordinates
key = (("inject", )
+ Vf.ufl_element().value_shape()
+ entity_dofs_key(Vc.finat_element.entity_dofs())
+ entity_dofs_key(Vf.finat_element.entity_dofs())
+ entity_dofs_key(Vc.mesh().coordinates.function_space().finat_element.entity_dofs())
+ entity_dofs_key(coordinates.function_space().finat_element.entity_dofs()))
try:
return cache[key]
except KeyError:
ncandidate = hierarchy.coarse_to_fine_cells[level].shape[1]
if Vc.finat_element.entity_dofs() == Vc.finat_element.entity_closure_dofs():
return cache.setdefault(key, (dg_injection_kernel(Vf, Vc, ncandidate), True))
coordinates = Vf.ufl_domain().coordinates
evaluate_kernel = compile_element(ufl.Coefficient(Vf))
to_reference_kernel = to_reference_coordinates(coordinates.ufl_element())
coords_element = create_element(coordinates.ufl_element())
Vf_element = create_element(Vf.ufl_element())
kernel = """
%(to_reference)s
%(evaluate)s
__attribute__((noinline)) /* Clang bug */
static void pyop2_kernel_inject(double *R, const double *X, const double *f, const double *Xf)
{
double Xref[%(tdim)d];
int cell = -1;
for (int i = 0; i < %(ncandidate)d; i++) {
const double *Xfi = Xf + i*%(Xf_cell_inc)d;
to_reference_coords_kernel(Xref, X, Xfi);
if (%(inside_cell)s) {
cell = i;
break;
}
}
if (cell == -1) {
abort();
}
const double *fi = f + cell*%(f_cell_inc)d;
for ( int i = 0; i < %(Rdim)d; i++ ) {
R[i] = 0;
}
pyop2_kernel_evaluate(R, fi, Xref);
}
""" % {
"to_reference": str(to_reference_kernel),
"evaluate": str(evaluate_kernel),
"inside_cell": inside_check(Vc.finat_element.cell, eps=1e-8, X="Xref"),
"tdim": Vc.ufl_domain().topological_dimension(),
"ncandidate": ncandidate,
"Rdim": numpy.prod(Vf_element.value_shape),
"Xf_cell_inc": coords_element.space_dimension(),
"f_cell_inc": Vf_element.space_dimension()
}
return cache.setdefault(key, (op2.Kernel(kernel, name="pyop2_kernel_inject"), False))
def inject_kernel(Vf, Vc):
hierarchy, level = utils.get_level(Vc.ufl_domain())
cache = hierarchy._shared_data_cache["transfer_kernels"]
coordinates = Vf.ufl_domain().coordinates
key = (("inject", ) +
Vf.ufl_element().value_shape() +
entity_dofs_key(Vc.finat_element.entity_dofs()) +
entity_dofs_key(Vf.finat_element.entity_dofs()) +
entity_dofs_key(Vc.mesh().coordinates.function_space().finat_element.entity_dofs()) +
entity_dofs_key(coordinates.function_space().finat_element.entity_dofs()))
try:
return cache[key]
except KeyError:
ncandidate = hierarchy.coarse_to_fine_cells[level].shape[1]
if Vc.finat_element.entity_dofs() == Vc.finat_element.entity_closure_dofs():
return cache.setdefault(key, (dg_injection_kernel(Vf, Vc, ncandidate), True))
coordinates = Vf.ufl_domain().coordinates
evaluate_kernel = compile_element(ufl.Coefficient(Vf))
to_reference_kernel = to_reference_coordinates(coordinates.ufl_element())
coords_element = create_element(coordinates.ufl_element())
Vf_element = create_element(Vf.ufl_element())
kernel = """
%(to_reference)s
%(evaluate)s
void inject_kernel(double *R, const double *X, const double *f, const double *Xf)
{
double Xref[%(tdim)d];
int cell = -1;
for (int i = 0; i < %(ncandidate)d; i++) {
const double *Xfi = Xf + i*%(Xf_cell_inc)d;
to_reference_coords_kernel(Xref, X, Xfi);
if (%(inside_cell)s) {
cell = i;
break;
}
}
if (cell == -1) {
abort();
}
const double *fi = f + cell*%(f_cell_inc)d;
for ( int i = 0; i < %(Rdim)d; i++ ) {
R[i] = 0;
}
evaluate_kernel(R, fi, Xref);
}
""" % {
"to_reference": str(to_reference_kernel),
"evaluate": str(evaluate_kernel),
"inside_cell": inside_check(Vc.finat_element.cell, eps=1e-8, X="Xref"),
"tdim": Vc.ufl_domain().topological_dimension(),
"ncandidate": ncandidate,
"Rdim": numpy.prod(Vf_element.value_shape),
"Xf_cell_inc": coords_element.space_dimension(),
"f_cell_inc": Vf_element.space_dimension()
}
return cache.setdefault(key, (op2.Kernel(kernel, name="inject_kernel"), False))
def prolong_kernel(expression):
hierarchy, level = utils.get_level(expression.ufl_domain())
levelf = level + Fraction(1 / hierarchy.refinements_per_level)
cache = hierarchy._shared_data_cache["transfer_kernels"]
coordinates = expression.ufl_domain().coordinates
key = (("prolong", )
+ expression.ufl_element().value_shape()
+ entity_dofs_key(expression.function_space().finat_element.entity_dofs())
+ entity_dofs_key(coordinates.function_space().finat_element.entity_dofs()))
try:
return cache[key]
except KeyError:
mesh = coordinates.ufl_domain()
evaluate_kernel = compile_element(expression)
to_reference_kernel = to_reference_coordinates(coordinates.ufl_element())
element = create_element(expression.ufl_element())
eval_args = evaluate_kernel.args[:-1]
coords_element = create_element(coordinates.ufl_element())
args = eval_args[-1].gencode(not_scope=True)
R, coarse = (a.sym.symbol for a in eval_args)
my_kernel = """
%(to_reference)s
%(evaluate)s
__attribute__((noinline)) /* Clang bug */
static void pyop2_kernel_prolong(double *R, %(args)s, const double *X, const double *Xc)
{
double Xref[%(tdim)d];
int cell = -1;
for (int i = 0; i < %(ncandidate)d; i++) {
const double *Xci = Xc + i*%(Xc_cell_inc)d;
to_reference_coords_kernel(Xref, X, Xci);
if (%(inside_cell)s) {
cell = i;
break;
}
}
if (cell == -1) abort();
const double *coarsei = %(coarse)s + cell*%(coarse_cell_inc)d;
for ( int i = 0; i < %(Rdim)d; i++ ) {
%(R)s[i] = 0;
}
pyop2_kernel_evaluate(%(R)s, coarsei, Xref);
}
""" % {"to_reference": str(to_reference_kernel),
"evaluate": str(evaluate_kernel),
"args": args,
"R": R,
"coarse": coarse,
"ncandidate": hierarchy.fine_to_coarse_cells[levelf].shape[1],
"Rdim": numpy.prod(element.value_shape),
"inside_cell": inside_check(element.cell, eps=1e-8, X="Xref"),
"Xc_cell_inc": coords_element.space_dimension(),
"coarse_cell_inc": element.space_dimension(),
"tdim": mesh.topological_dimension()}
return cache.setdefault(key, op2.Kernel(my_kernel, name="pyop2_kernel_prolong"))
def restrict_kernel(Vf, Vc):
hierarchy, level = utils.get_level(Vc.ufl_domain())
levelf = level + Fraction(1 / hierarchy.refinements_per_level)
cache = hierarchy._shared_data_cache["transfer_kernels"]
coordinates = Vc.ufl_domain().coordinates
key = (("restrict", ) + Vf.ufl_element().value_shape() +
entity_dofs_key(Vf.finat_element.entity_dofs()) +
entity_dofs_key(Vc.finat_element.entity_dofs()) + entity_dofs_key(
coordinates.function_space().finat_element.entity_dofs()))
try:
return cache[key]
except KeyError:
mesh = coordinates.ufl_domain()
evaluate_kernel = compile_element(firedrake.TestFunction(Vc), Vf)
to_reference_kernel = to_reference_coordinates(
coordinates.ufl_element())
coords_element = create_element(coordinates.ufl_element())
element = create_element(Vc.ufl_element())
eval_args = evaluate_kernel.args[:-1]
args = eval_args[-1].gencode(not_scope=True)
R, fine = (a.sym.symbol for a in eval_args)
my_kernel = """
%(to_reference)s
%(evaluate)s
__attribute__((noinline)) /* Clang bug */
static void pyop2_kernel_restrict(double *R, %(args)s, const double *X, const double *Xc)
{
double Xref[%(tdim)d];
int cell = -1;
for (int i = 0; i < %(ncandidate)d; i++) {
const double *Xci = Xc + i*%(Xc_cell_inc)d;
to_reference_coords_kernel(Xref, X, Xci);
if (%(inside_cell)s) {
cell = i;
const double *Ri = %(R)s + cell*%(coarse_cell_inc)d;
pyop2_kernel_evaluate(Ri, %(fine)s, Xref);
break;
}
}
}
""" % {
"to_reference": str(to_reference_kernel),
"evaluate": str(evaluate_kernel),
"ncandidate": hierarchy.fine_to_coarse_cells[levelf].shape[1],
"inside_cell": inside_check(element.cell, eps=1e-8, X="Xref"),
"Xc_cell_inc": coords_element.space_dimension(),
"coarse_cell_inc": element.space_dimension(),
"args": args,
"R": R,
"fine": fine,
"tdim": mesh.topological_dimension()
}
return cache.setdefault(
key, op2.Kernel(my_kernel, name="pyop2_kernel_restrict"))
def test_tensor_prod_simple(ufl_A, ufl_B):
tensor_ufl = ufl.TensorProductElement(ufl_A, ufl_B)
tensor = create_element(tensor_ufl)
A = create_element(ufl_A)
B = create_element(ufl_B)
assert isinstance(tensor, finat.TensorProductElement)
assert tensor.factors == (A, B)
def needs_coordinate_mapping(element):
"""Does this UFL element require a CoordinateMapping for translation?"""
if element.family() == 'Real':
return False
else:
return isinstance(create_element(element),
NeedsCoordinateMappingElement)
def extent(self, coefficient):
""" Calculation of the range of a coefficient."""
element = coefficient.ufl_element()
if element.family() == "Real":
return (coefficient.dat.cdim, )
else:
return (create_element(element).space_dimension(), )
def prepare_coordinates(coefficient, name, interior_facet=False):
"""Bridges the kernel interface and the GEM abstraction for
coordinates.
:arg coefficient: UFL Coefficient
:arg name: unique name to refer to the Coefficient in the kernel
:arg interior_facet: interior facet integral?
:returns: (funarg, expression)
funarg - :class:`coffee.Decl` function argument
expression - GEM expression referring to the Coefficient
values
"""
finat_element = create_element(coefficient.ufl_element())
shape = finat_element.index_shape
size = numpy.prod(shape, dtype=int)
if not interior_facet:
funargs = [coffee.Decl(SCALAR_TYPE, coffee.Symbol(name),
pointers=[("",)],
qualifiers=["const"])]
variable = gem.Variable(name, (size,))
expression = gem.reshape(variable, shape)
else:
funargs = [coffee.Decl(SCALAR_TYPE, coffee.Symbol(name+"_0"),
pointers=[("",)],
qualifiers=["const"]),
coffee.Decl(SCALAR_TYPE, coffee.Symbol(name+"_1"),
pointers=[("",)],
qualifiers=["const"])]
variable0 = gem.Variable(name+"_0", (size,))
variable1 = gem.Variable(name+"_1", (size,))
expression = (gem.reshape(variable0, shape),
gem.reshape(variable1, shape))
return funargs, expression
def translate_coefficient(terminal, mt, ctx):
vec = ctx.coefficient(terminal, mt.restriction)
if terminal.ufl_element().family() == 'Real':
assert mt.local_derivatives == 0
return vec
element = create_element(terminal.ufl_element())
# Collect FInAT tabulation for all entities
per_derivative = collections.defaultdict(list)
for entity_id in ctx.entity_ids:
finat_dict = element.basis_evaluation(mt.local_derivatives,
ctx.point_set,
(ctx.integration_dim, entity_id))
for alpha, table in iteritems(finat_dict):
# Filter out irrelevant derivatives
if sum(alpha) == mt.local_derivatives:
# A numerical hack that FFC used to apply on FIAT
# tables still lives on after ditching FFC and
# switching to FInAT.
table = ffc_rounding(table, ctx.epsilon)
per_derivative[alpha].append(table)
# Merge entity tabulations for each derivative
if len(ctx.entity_ids) == 1:
def take_singleton(xs):
x, = xs # asserts singleton
return x
per_derivative = {alpha: take_singleton(tables)
for alpha, tables in iteritems(per_derivative)}
else:
f = ctx.facet_number(mt.restriction)
per_derivative = {alpha: gem.select_expression(tables, f)
for alpha, tables in iteritems(per_derivative)}
# Coefficient evaluation
beta = element.get_indices()
zeta = element.get_value_indices()
value_dict = {}
for alpha, table in iteritems(per_derivative):
value = gem.IndexSum(gem.Product(gem.Indexed(table, beta + zeta),
gem.Indexed(vec, beta)),
beta)
optimised_value = gem.optimise.contraction(value)
value_dict[alpha] = gem.ComponentTensor(optimised_value, zeta)
# Change from FIAT to UFL arrangement
result = fiat_to_ufl(value_dict, mt.local_derivatives)
assert result.shape == mt.expr.ufl_shape
assert set(result.free_indices) <= set(ctx.point_set.indices)
# Detect Jacobian of affine cells
if not result.free_indices and all(numpy.count_nonzero(node.array) <= 2
for node in traversal((result,))
if isinstance(node, gem.Literal)):
result = gem.optimise.aggressive_unroll(result)
return result
def __init__(self, mesh, element, name=None, real_tensorproduct=False):
super(FunctionSpace, self).__init__()
if type(element) is ufl.MixedElement:
raise ValueError("Can't create FunctionSpace for MixedElement")
finat_element = create_element(element)
if isinstance(finat_element, finat.TensorFiniteElement):
# Retrieve scalar element
finat_element = finat_element.base_element
# Used for reconstruction of mixed/component spaces
self.real_tensorproduct = real_tensorproduct
sdata = get_shared_data(mesh,
finat_element,
real_tensorproduct=real_tensorproduct)
# The function space shape is the number of dofs per node,
# hence it is not always the value_shape. Vector and Tensor
# element modifiers *must* live on the outside!
if type(element) is ufl.TensorElement:
# UFL enforces value_shape of the subelement to be empty
# on a TensorElement.
self.shape = element.value_shape()
elif type(element) is ufl.VectorElement:
# First dimension of the value_shape is the VectorElement
# shape.
self.shape = element.value_shape()[:1]
else:
self.shape = ()
self._ufl_function_space = ufl.FunctionSpace(mesh.ufl_mesh(), element)
self._shared_data = sdata
self._mesh = mesh
self.rank = len(self.shape)
r"""The rank of this :class:`FunctionSpace`. Spaces where the
element is scalar-valued (or intrinsically vector-valued) have
rank zero. Spaces built on :class:`~ufl.classes.VectorElement` or
:class:`~ufl.classes.TensorElement` instances have rank equivalent to
the number of components of their
:meth:`~ufl.classes.FiniteElementBase.value_shape`."""
self.value_size = int(numpy.prod(self.shape, dtype=int))
r"""The total number of degrees of freedom at each function
space node."""
self.name = name
r"""The (optional) descriptive name for this space."""
self.node_set = sdata.node_set
r"""A :class:`pyop2.Set` representing the function space nodes."""
self.dof_dset = op2.DataSet(self.node_set,
self.shape or 1,
name="%s_nodes_dset" % self.name)
r"""A :class:`pyop2.DataSet` representing the function space
degrees of freedom."""
self.comm = self.node_set.comm
self.finat_element = finat_element
self.extruded = sdata.extruded
self.offset = sdata.offset
self.cell_boundary_masks = sdata.cell_boundary_masks
self.interior_facet_boundary_masks = sdata.interior_facet_boundary_masks
def prepare_arguments(arguments, multiindices, scalar_type, interior_facet=False, diagonal=False):
"""Bridges the kernel interface and the GEM abstraction for
Arguments. Vector Arguments are rearranged here for interior
facet integrals.
:arg arguments: UFL Arguments
:arg multiindices: Argument multiindices
:arg interior_facet: interior facet integral?
:arg diagonal: Are we assembling the diagonal of a rank-2 element tensor?
:returns: (funarg, expression)
funarg - :class:`loopy.GlobalArg` function argument
expressions - GEM expressions referring to the argument
tensor
"""
assert isinstance(interior_facet, bool)
if len(arguments) == 0:
# No arguments
funarg = lp.GlobalArg("A", dtype=scalar_type, shape=(1,))
expression = gem.Indexed(gem.Variable("A", (1,)), (0,))
return funarg, [expression]
elements = tuple(create_element(arg.ufl_element()) for arg in arguments)
shapes = tuple(element.index_shape for element in elements)
if diagonal:
if len(arguments) != 2:
raise ValueError("Diagonal only for 2-forms")
try:
element, = set(elements)
except ValueError:
raise ValueError("Diagonal only for diagonal blocks (test and trial spaces the same)")
elements = (element, )
shapes = tuple(element.index_shape for element in elements)
multiindices = multiindices[:1]
def expression(restricted):
return gem.Indexed(gem.reshape(restricted, *shapes),
tuple(chain(*multiindices)))
u_shape = numpy.array([numpy.prod(shape, dtype=int) for shape in shapes])
if interior_facet:
c_shape = tuple(2 * u_shape)
slicez = [[slice(r * s, (r + 1) * s)
for r, s in zip(restrictions, u_shape)]
for restrictions in product((0, 1), repeat=len(arguments))]
else:
c_shape = tuple(u_shape)
slicez = [[slice(s) for s in u_shape]]
funarg = lp.GlobalArg("A", dtype=scalar_type, shape=c_shape)
varexp = gem.Variable("A", c_shape)
expressions = [expression(gem.view(varexp, *slices)) for slices in slicez]
return funarg, prune(expressions)
def test_ufl_only_to_contravariant_piola():
mesh = ufl.Mesh(ufl.VectorElement("P", ufl.triangle, 1))
V = ufl.FunctionSpace(mesh, ufl.FiniteElement("P", ufl.triangle, 2))
v = ufl.Coefficient(V)
expr = ufl.as_vector([v, v])
W = ufl.FunctionSpace(mesh, ufl.FiniteElement("RT", ufl.triangle, 1))
to_element = create_element(W.ufl_element())
ast, oriented, needs_cell_sizes, coefficients, first_coeff_fake_coords, *_ = compile_expression_dual_evaluation(
expr, to_element, coffee=False)
assert first_coeff_fake_coords is True
def test_ufl_only_spatialcoordinate():
mesh = ufl.Mesh(ufl.VectorElement("P", ufl.triangle, 1))
V = ufl.FunctionSpace(mesh, ufl.FiniteElement("P", ufl.triangle, 2))
x, y = ufl.SpatialCoordinate(mesh)
expr = x * y - y**2 + x
W = V
to_element = create_element(W.ufl_element())
ast, oriented, needs_cell_sizes, coefficients, first_coeff_fake_coords, *_ = compile_expression_dual_evaluation(
expr, to_element, coffee=False)
assert first_coeff_fake_coords is True
def _coefficient(self, coefficient, name):
element = create_element(coefficient.ufl_element())
shape = self.shape + element.index_shape
size = numpy.prod(shape, dtype=int)
funarg = ast.Decl(SCALAR_TYPE, ast.Symbol(name), pointers=[("restrict", )],
qualifiers=["const"])
expression = gem.reshape(gem.Variable(name, (size, )), shape)
expression = gem.partial_indexed(expression, self.indices)
self.coefficient_map[coefficient] = expression
return funarg
def test_ufl_only_simple():
mesh = ufl.Mesh(ufl.VectorElement("P", ufl.triangle, 1))
V = ufl.FunctionSpace(mesh, ufl.FiniteElement("P", ufl.triangle, 2))
v = ufl.Coefficient(V)
expr = ufl.inner(v, v)
W = V
to_element = create_element(W.ufl_element())
ast, oriented, needs_cell_sizes, coefficients, first_coeff_fake_coords, *_ = compile_expression_dual_evaluation(
expr, to_element, coffee=False)
assert first_coeff_fake_coords is False
def _coefficient(self, coefficient, name):
element = create_element(coefficient.ufl_element())
shape = self.shape + element.index_shape
size = numpy.prod(shape, dtype=int)
funarg = ast.Decl(SCALAR_TYPE, ast.Symbol(name), pointers=[("restrict", )],
qualifiers=["const"])
expression = gem.reshape(gem.Variable(name, (size, )), shape)
expression = gem.partial_indexed(expression, self.indices)
self.coefficient_map[coefficient] = expression
return funarg
def prepare_coefficient(coefficient, name, interior_facet=False):
"""Bridges the kernel interface and the GEM abstraction for
Coefficients.
:arg coefficient: UFL Coefficient
:arg name: unique name to refer to the Coefficient in the kernel
:arg interior_facet: interior facet integral?
:returns: (funarg, expression)
funarg - :class:`coffee.Decl` function argument
expression - GEM expression referring to the Coefficient
values
"""
assert isinstance(interior_facet, bool)
if coefficient.ufl_element().family() == 'Real':
# Constant
funarg = coffee.Decl(SCALAR_TYPE,
coffee.Symbol(name),
pointers=[("restrict", )],
qualifiers=["const"])
expression = gem.reshape(gem.Variable(name, (None, )),
coefficient.ufl_shape)
return funarg, expression
finat_element = create_element(coefficient.ufl_element())
if isinstance(finat_element, TensorFiniteElement):
scalar_shape = finat_element.base_element.index_shape
tensor_shape = finat_element.index_shape[len(scalar_shape):]
else:
scalar_shape = finat_element.index_shape
tensor_shape = ()
scalar_size = numpy.prod(scalar_shape, dtype=int)
tensor_size = numpy.prod(tensor_shape, dtype=int)
funarg = coffee.Decl(SCALAR_TYPE,
coffee.Symbol(name),
pointers=[("const", "restrict"), ("restrict", )],
qualifiers=["const"])
if not interior_facet:
expression = gem.reshape(
gem.Variable(name, (scalar_size, tensor_size)), scalar_shape,
tensor_shape)
else:
varexp = gem.Variable(name, (2 * scalar_size, tensor_size))
plus = gem.view(varexp, slice(scalar_size), slice(tensor_size))
minus = gem.view(varexp, slice(scalar_size, 2 * scalar_size),
slice(tensor_size))
expression = (gem.reshape(plus, scalar_shape, tensor_shape),
gem.reshape(minus, scalar_shape, tensor_shape))
return funarg, expression
def test_ufl_only_shape_mismatch():
mesh = ufl.Mesh(ufl.VectorElement("P", ufl.triangle, 1))
V = ufl.FunctionSpace(mesh, ufl.FiniteElement("RT", ufl.triangle, 1))
v = ufl.Coefficient(V)
expr = ufl.inner(v, v)
assert expr.ufl_shape == ()
W = V
to_element = create_element(W.ufl_element())
assert to_element.value_shape == (2, )
with pytest.raises(ValueError):
ast, oriented, needs_cell_sizes, coefficients, first_coeff_fake_coords, *_ = compile_expression_dual_evaluation(
expr, to_element, coffee=False)
def prolongation_transfer_kernel_action(Vf, expr):
from tsfc import compile_expression_dual_evaluation
from tsfc.finatinterface import create_element
to_element = create_element(Vf.ufl_element())
kernel = compile_expression_dual_evaluation(expr, to_element,
Vf.ufl_element())
ast = kernel.ast
name = kernel.name
flop_count = kernel.flop_count
return op2.Kernel(ast,
name,
requires_zeroed_output_arguments=True,
flop_count=flop_count)
def prepare_arguments(arguments,
multiindices,
scalar_type,
interior_facet=False):
"""Bridges the kernel interface and the GEM abstraction for
Arguments. Vector Arguments are rearranged here for interior
facet integrals.
:arg arguments: UFL Arguments
:arg multiindices: Argument multiindices
:arg interior_facet: interior facet integral?
:returns: (funarg, expression)
funarg - :class:`coffee.Decl` function argument
expressions - GEM expressions referring to the argument
tensor
"""
assert isinstance(interior_facet, bool)
if len(arguments) == 0:
# No arguments
funarg = coffee.Decl(scalar_type, coffee.Symbol("A", rank=(1, )))
expression = gem.Indexed(gem.Variable("A", (1, )), (0, ))
return funarg, [expression]
elements = tuple(create_element(arg.ufl_element()) for arg in arguments)
shapes = tuple(element.index_shape for element in elements)
def expression(restricted):
return gem.Indexed(gem.reshape(restricted, *shapes),
tuple(chain(*multiindices)))
u_shape = numpy.array([numpy.prod(shape, dtype=int) for shape in shapes])
if interior_facet:
c_shape = tuple(2 * u_shape)
slicez = [[
slice(r * s, (r + 1) * s) for r, s in zip(restrictions, u_shape)
] for restrictions in product((0, 1), repeat=len(arguments))]
else:
c_shape = tuple(u_shape)
slicez = [[slice(s) for s in u_shape]]
funarg = coffee.Decl(scalar_type, coffee.Symbol("A", rank=c_shape))
varexp = gem.Variable("A", c_shape)
expressions = [expression(gem.view(varexp, *slices)) for slices in slicez]
return funarg, prune(expressions)
def prepare_arguments(arguments, multiindices, interior_facet=False):
"""Bridges the kernel interface and the GEM abstraction for
Arguments. Vector Arguments are rearranged here for interior
facet integrals.
:arg arguments: UFL Arguments
:arg multiindices: Argument multiindices
:arg interior_facet: interior facet integral?
:returns: (funarg, expression)
funarg - :class:`coffee.Decl` function argument
expressions - GEM expressions referring to the argument
tensor
"""
assert isinstance(interior_facet, bool)
if len(arguments) == 0:
# No arguments
funarg = coffee.Decl(SCALAR_TYPE, coffee.Symbol("A", rank=(1,)))
expression = gem.Indexed(gem.Variable("A", (1,)), (0,))
return funarg, [expression]
elements = tuple(create_element(arg.ufl_element()) for arg in arguments)
shapes = tuple(element.index_shape for element in elements)
def expression(restricted):
return gem.Indexed(gem.reshape(restricted, *shapes),
tuple(chain(*multiindices)))
u_shape = numpy.array([numpy.prod(shape, dtype=int) for shape in shapes])
if interior_facet:
c_shape = tuple(2 * u_shape)
slicez = [[slice(r * s, (r + 1) * s)
for r, s in zip(restrictions, u_shape)]
for restrictions in product((0, 1), repeat=len(arguments))]
else:
c_shape = tuple(u_shape)
slicez = [[slice(s) for s in u_shape]]
funarg = coffee.Decl(SCALAR_TYPE, coffee.Symbol("A", rank=c_shape))
varexp = gem.Variable("A", c_shape)
expressions = [expression(gem.view(varexp, *slices)) for slices in slicez]
return funarg, prune(expressions)
def prepare_coefficient(coefficient, name, interior_facet=False):
"""Bridges the kernel interface and the GEM abstraction for
Coefficients.
:arg coefficient: UFL Coefficient
:arg name: unique name to refer to the Coefficient in the kernel
:arg interior_facet: interior facet integral?
:returns: (funarg, expression)
funarg - :class:`coffee.Decl` function argument
expression - GEM expression referring to the Coefficient
values
"""
assert isinstance(interior_facet, bool)
if coefficient.ufl_element().family() == 'Real':
# Constant
funarg = coffee.Decl(SCALAR_TYPE, coffee.Symbol(name),
pointers=[("restrict",)],
qualifiers=["const"])
expression = gem.reshape(gem.Variable(name, (None,)),
coefficient.ufl_shape)
return funarg, expression
finat_element = create_element(coefficient.ufl_element())
shape = finat_element.index_shape
size = numpy.prod(shape, dtype=int)
funarg = coffee.Decl(SCALAR_TYPE, coffee.Symbol(name),
pointers=[("restrict",)],
qualifiers=["const"])
if not interior_facet:
expression = gem.reshape(gem.Variable(name, (size,)), shape)
else:
varexp = gem.Variable(name, (2 * size,))
plus = gem.view(varexp, slice(size))
minus = gem.view(varexp, slice(size, 2 * size))
expression = (gem.reshape(plus, shape),
gem.reshape(minus, shape))
return funarg, expression
def prolongation_transfer_kernel_aij(Pk, P1):
# Works for Pk, Pm; I just retain the notation
# P1 to remind you that P1 is of lower degree
# than Pk
from tsfc import compile_expression_dual_evaluation
from tsfc.finatinterface import create_element
from firedrake import TestFunction
expr = TestFunction(P1)
to_element = create_element(Pk.ufl_element())
kernel = compile_expression_dual_evaluation(expr, to_element,
Pk.ufl_element())
ast = kernel.ast
name = kernel.name
flop_count = kernel.flop_count
return op2.Kernel(ast,
name,
requires_zeroed_output_arguments=True,
flop_count=flop_count)
def translate_argument(terminal, mt, ctx):
argument_multiindex = ctx.argument_multiindices[terminal.number()]
sigma = tuple(gem.Index(extent=d) for d in mt.expr.ufl_shape)
element = create_element(terminal.ufl_element())
def callback(entity_id):
finat_dict = ctx.basis_evaluation(element, mt.local_derivatives, entity_id)
# Filter out irrelevant derivatives
filtered_dict = {alpha: table
for alpha, table in iteritems(finat_dict)
if sum(alpha) == mt.local_derivatives}
# Change from FIAT to UFL arrangement
square = fiat_to_ufl(filtered_dict, mt.local_derivatives)
# A numerical hack that FFC used to apply on FIAT tables still
# lives on after ditching FFC and switching to FInAT.
return ffc_rounding(square, ctx.epsilon)
table = ctx.entity_selector(callback, mt.restriction)
return gem.ComponentTensor(gem.Indexed(table, argument_multiindex + sigma), sigma)
def prolong_kernel(expression):
hierarchy, _ = utils.get_level(expression.ufl_domain())
cache = hierarchy._shared_data_cache["transfer_kernels"]
coordinates = expression.ufl_domain().coordinates
key = (("prolong", ) +
expression.ufl_element().value_shape() +
entity_dofs_key(expression.function_space().finat_element.entity_dofs()) +
entity_dofs_key(coordinates.function_space().finat_element.entity_dofs()))
try:
return cache[key]
except KeyError:
mesh = coordinates.ufl_domain()
evaluate_kernel = compile_element(expression)
to_reference_kernel = to_reference_coordinates(coordinates.ufl_element())
element = create_element(expression.ufl_element())
eval_args = evaluate_kernel.args[:-1]
args = ", ".join(a.gencode(not_scope=True) for a in eval_args)
arg_names = ", ".join(a.sym.symbol for a in eval_args)
my_kernel = """
%(to_reference)s
%(evaluate)s
void prolong_kernel(%(args)s, const double *X, const double *Xc)
{
double Xref[%(tdim)d];
to_reference_coords_kernel(Xref, X, Xc);
for ( int i = 0; i < %(Rdim)d; i++ ) {
%(R)s[i] = 0;
}
evaluate_kernel(%(arg_names)s, Xref);
}
""" % {"to_reference": str(to_reference_kernel),
"evaluate": str(evaluate_kernel),
"args": args,
"R": arg_names[0],
"Rdim": numpy.prod(element.value_shape),
"arg_names": arg_names,
"tdim": mesh.topological_dimension()}
return cache.setdefault(key, op2.Kernel(my_kernel, name="prolong_kernel"))
def prepare_coefficient(coefficient, name, scalar_type, interior_facet=False):
"""Bridges the kernel interface and the GEM abstraction for
Coefficients.
:arg coefficient: UFL Coefficient
:arg name: unique name to refer to the Coefficient in the kernel
:arg interior_facet: interior facet integral?
:returns: (funarg, expression)
funarg - :class:`loopy.GlobalArg` function argument
expression - GEM expression referring to the Coefficient
values
"""
assert isinstance(interior_facet, bool)
if coefficient.ufl_element().family() == 'Real':
# Constant
funarg = lp.GlobalArg(name,
dtype=scalar_type,
shape=(coefficient.ufl_element().value_size(), ))
expression = gem.reshape(gem.Variable(name, (None, )),
coefficient.ufl_shape)
return funarg, expression
finat_element = create_element(coefficient.ufl_element())
shape = finat_element.index_shape
size = numpy.prod(shape, dtype=int)
if not interior_facet:
expression = gem.reshape(gem.Variable(name, (size, )), shape)
else:
varexp = gem.Variable(name, (2 * size, ))
plus = gem.view(varexp, slice(size))
minus = gem.view(varexp, slice(size, 2 * size))
expression = (gem.reshape(plus, shape), gem.reshape(minus, shape))
size = size * 2
funarg = lp.GlobalArg(name, dtype=scalar_type, shape=(size, ))
return funarg, expression
def dg_injection_kernel(Vf, Vc, ncell):
from firedrake import Tensor, AssembledVector, TestFunction, TrialFunction
from firedrake.slate.slac import compile_expression
macro_builder = MacroKernelBuilder(ScalarType_c, ncell)
f = ufl.Coefficient(Vf)
macro_builder.set_coefficients([f])
macro_builder.set_coordinates(Vf.mesh())
Vfe = create_element(Vf.ufl_element())
macro_quadrature_rule = make_quadrature(
Vfe.cell, estimate_total_polynomial_degree(ufl.inner(f, f)))
index_cache = {}
parameters = default_parameters()
integration_dim, entity_ids = lower_integral_type(Vfe.cell, "cell")
macro_cfg = dict(interface=macro_builder,
ufl_cell=Vf.ufl_cell(),
precision=parameters["precision"],
integration_dim=integration_dim,
entity_ids=entity_ids,
index_cache=index_cache,
quadrature_rule=macro_quadrature_rule)
fexpr, = fem.compile_ufl(f, **macro_cfg)
X = ufl.SpatialCoordinate(Vf.mesh())
C_a, = fem.compile_ufl(X, **macro_cfg)
detJ = ufl_utils.preprocess_expression(
abs(ufl.JacobianDeterminant(f.ufl_domain())))
macro_detJ, = fem.compile_ufl(detJ, **macro_cfg)
Vce = create_element(Vc.ufl_element())
coarse_builder = firedrake_interface.KernelBuilder("cell", "otherwise", 0,
ScalarType_c)
coarse_builder.set_coordinates(Vc.mesh())
argument_multiindices = (Vce.get_indices(), )
argument_multiindex, = argument_multiindices
return_variable, = coarse_builder.set_arguments((ufl.TestFunction(Vc), ),
argument_multiindices)
integration_dim, entity_ids = lower_integral_type(Vce.cell, "cell")
# Midpoint quadrature for jacobian on coarse cell.
quadrature_rule = make_quadrature(Vce.cell, 0)
coarse_cfg = dict(interface=coarse_builder,
ufl_cell=Vc.ufl_cell(),
precision=parameters["precision"],
integration_dim=integration_dim,
entity_ids=entity_ids,
index_cache=index_cache,
quadrature_rule=quadrature_rule)
X = ufl.SpatialCoordinate(Vc.mesh())
K = ufl_utils.preprocess_expression(ufl.JacobianInverse(Vc.mesh()))
C_0, = fem.compile_ufl(X, **coarse_cfg)
K, = fem.compile_ufl(K, **coarse_cfg)
i = gem.Index()
j = gem.Index()
C_0 = gem.Indexed(C_0, (j, ))
C_0 = gem.index_sum(C_0, quadrature_rule.point_set.indices)
C_a = gem.Indexed(C_a, (j, ))
X_a = gem.Sum(C_0, gem.Product(gem.Literal(-1), C_a))
K_ij = gem.Indexed(K, (i, j))
K_ij = gem.index_sum(K_ij, quadrature_rule.point_set.indices)
X_a = gem.index_sum(gem.Product(K_ij, X_a), (j, ))
C_0, = quadrature_rule.point_set.points
C_0 = gem.Indexed(gem.Literal(C_0), (i, ))
# fine quad points in coarse reference space.
X_a = gem.Sum(C_0, gem.Product(gem.Literal(-1), X_a))
X_a = gem.ComponentTensor(X_a, (i, ))
# Coarse basis function evaluated at fine quadrature points
phi_c = fem.fiat_to_ufl(
Vce.point_evaluation(0, X_a, (Vce.cell.get_dimension(), 0)), 0)
tensor_indices = tuple(gem.Index(extent=d) for d in f.ufl_shape)
phi_c = gem.Indexed(phi_c, argument_multiindex + tensor_indices)
fexpr = gem.Indexed(fexpr, tensor_indices)
quadrature_weight = macro_quadrature_rule.weight_expression
expr = gem.Product(gem.IndexSum(gem.Product(phi_c, fexpr), tensor_indices),
gem.Product(macro_detJ, quadrature_weight))
quadrature_indices = macro_builder.indices + macro_quadrature_rule.point_set.indices
reps = spectral.Integrals([expr], quadrature_indices,
argument_multiindices, parameters)
assignments = spectral.flatten([(return_variable, reps)], index_cache)
return_variables, expressions = zip(*assignments)
expressions = impero_utils.preprocess_gem(expressions,
**spectral.finalise_options)
assignments = list(zip(return_variables, expressions))
impero_c = impero_utils.compile_gem(assignments,
quadrature_indices +
argument_multiindex,
remove_zeros=True)
index_names = []
def name_index(index, name):
index_names.append((index, name))
if index in index_cache:
for multiindex, suffix in zip(index_cache[index],
string.ascii_lowercase):
name_multiindex(multiindex, name + suffix)
def name_multiindex(multiindex, name):
if len(multiindex) == 1:
name_index(multiindex[0], name)
else:
for i, index in enumerate(multiindex):
name_index(index, name + str(i))
name_multiindex(quadrature_indices, 'ip')
for multiindex, name in zip(argument_multiindices, ['j', 'k']):
name_multiindex(multiindex, name)
index_names.extend(zip(macro_builder.indices, ["entity"]))
body = generate_coffee(impero_c, index_names, parameters["precision"],
ScalarType_c)
retarg = ast.Decl(ScalarType_c,
ast.Symbol("R", rank=(Vce.space_dimension(), )))
local_tensor = coarse_builder.local_tensor
local_tensor.init = ast.ArrayInit(
numpy.zeros(Vce.space_dimension(), dtype=ScalarType_c))
body.children.insert(0, local_tensor)
args = [retarg] + macro_builder.kernel_args + [
macro_builder.coordinates_arg, coarse_builder.coordinates_arg
]
# Now we have the kernel that computes <f, phi_c>dx_c
# So now we need to hit it with the inverse mass matrix on dx_c
u = TrialFunction(Vc)
v = TestFunction(Vc)
expr = Tensor(ufl.inner(u, v) * ufl.dx).inv * AssembledVector(
ufl.Coefficient(Vc))
Ainv, = compile_expression(expr)
Ainv = Ainv.kinfo.kernel
A = ast.Symbol(local_tensor.sym.symbol)
R = ast.Symbol("R")
body.children.append(
ast.FunCall(Ainv.name, R, coarse_builder.coordinates_arg.sym, A))
from coffee.base import Node
assert isinstance(Ainv._code, Node)
return op2.Kernel(ast.Node([
Ainv._code,
ast.FunDecl("void",
"pyop2_kernel_injection_dg",
args,
body,
pred=["static", "inline"])
]),
name="pyop2_kernel_injection_dg",
cpp=True,
include_dirs=Ainv._include_dirs,
headers=Ainv._headers)
def compile_element(expression,
dual_space=None,
parameters=None,
name="evaluate"):
"""Generate code for point evaluations.
:arg expression: A UFL expression (may contain up to one coefficient, or one argument)
:arg dual_space: if the expression has an argument, should we also distribute residual data?
:returns: Some coffee AST
"""
if parameters is None:
parameters = default_parameters()
else:
_ = default_parameters()
_.update(parameters)
parameters = _
expression = tsfc.ufl_utils.preprocess_expression(expression)
# # Collect required coefficients
try:
arg, = extract_coefficients(expression)
argument_multiindices = ()
coefficient = True
if expression.ufl_shape:
tensor_indices = tuple(gem.Index() for s in expression.ufl_shape)
else:
tensor_indices = ()
except ValueError:
arg, = extract_arguments(expression)
finat_elem = create_element(arg.ufl_element())
argument_multiindices = (finat_elem.get_indices(), )
argument_multiindex, = argument_multiindices
value_shape = finat_elem.value_shape
if value_shape:
tensor_indices = argument_multiindex[-len(value_shape):]
else:
tensor_indices = ()
coefficient = False
# Replace coordinates (if any)
builder = firedrake_interface.KernelBuilderBase(scalar_type=ScalarType_c)
domain = expression.ufl_domain()
# Translate to GEM
cell = domain.ufl_cell()
dim = cell.topological_dimension()
point = gem.Variable('X', (dim, ))
point_arg = ast.Decl(ScalarType_c, ast.Symbol('X', rank=(dim, )))
config = dict(interface=builder,
ufl_cell=cell,
precision=parameters["precision"],
point_indices=(),
point_expr=point,
argument_multiindices=argument_multiindices)
context = tsfc.fem.GemPointContext(**config)
# Abs-simplification
expression = tsfc.ufl_utils.simplify_abs(expression)
# Translate UFL -> GEM
if coefficient:
assert dual_space is None
f_arg = [builder._coefficient(arg, "f")]
else:
f_arg = []
translator = tsfc.fem.Translator(context)
result, = map_expr_dags(translator, [expression])
b_arg = []
if coefficient:
if expression.ufl_shape:
return_variable = gem.Indexed(
gem.Variable('R', expression.ufl_shape), tensor_indices)
result_arg = ast.Decl(ScalarType_c,
ast.Symbol('R', rank=expression.ufl_shape))
result = gem.Indexed(result, tensor_indices)
else:
return_variable = gem.Indexed(gem.Variable('R', (1, )), (0, ))
result_arg = ast.Decl(ScalarType_c, ast.Symbol('R', rank=(1, )))
else:
return_variable = gem.Indexed(
gem.Variable('R', finat_elem.index_shape), argument_multiindex)
result = gem.Indexed(result, tensor_indices)
if dual_space:
elem = create_element(dual_space.ufl_element())
if elem.value_shape:
var = gem.Indexed(gem.Variable("b", elem.value_shape),
tensor_indices)
b_arg = [
ast.Decl(ScalarType_c,
ast.Symbol("b", rank=elem.value_shape))
]
else:
var = gem.Indexed(gem.Variable("b", (1, )), (0, ))
b_arg = [ast.Decl(ScalarType_c, ast.Symbol("b", rank=(1, )))]
result = gem.Product(result, var)
result_arg = ast.Decl(ScalarType_c,
ast.Symbol('R', rank=finat_elem.index_shape))
# Unroll
max_extent = parameters["unroll_indexsum"]
if max_extent:
def predicate(index):
return index.extent <= max_extent
result, = gem.optimise.unroll_indexsum([result], predicate=predicate)
# Translate GEM -> COFFEE
result, = gem.impero_utils.preprocess_gem([result])
impero_c = gem.impero_utils.compile_gem([(return_variable, result)],
tensor_indices)
body = generate_coffee(impero_c, {}, parameters["precision"], ScalarType_c)
# Build kernel tuple
kernel_code = builder.construct_kernel(
"pyop2_kernel_" + name, [result_arg] + b_arg + f_arg + [point_arg],
body)
return kernel_code
def restrict_kernel(Vf, Vc):
hierarchy, level = utils.get_level(Vc.ufl_domain())
levelf = level + Fraction(1 / hierarchy.refinements_per_level)
cache = hierarchy._shared_data_cache["transfer_kernels"]
coordinates = Vc.ufl_domain().coordinates
key = (("restrict", ) + Vf.ufl_element().value_shape() +
entity_dofs_key(Vf.finat_element.entity_dofs()) +
entity_dofs_key(Vc.finat_element.entity_dofs()) + entity_dofs_key(
coordinates.function_space().finat_element.entity_dofs()))
try:
return cache[key]
except KeyError:
mesh = coordinates.ufl_domain()
evaluate_kernel = compile_element(firedrake.TestFunction(Vc), Vf)
to_reference_kernel = to_reference_coordinates(
coordinates.ufl_element())
coords_element = create_element(coordinates.ufl_element())
element = create_element(Vc.ufl_element())
eval_args = evaluate_kernel.args[:-1]
args = eval_args[-1].gencode(not_scope=True)
R, fine = (a.sym.symbol for a in eval_args)
my_kernel = """
%(to_reference)s
%(evaluate)s
__attribute__((noinline)) /* Clang bug */
static void pyop2_kernel_restrict(double *R, %(args)s, const double *X, const double *Xc)
{
double Xref[%(tdim)d];
int cell = -1;
int bestcell = -1;
double bestdist = 1e10;
for (int i = 0; i < %(ncandidate)d; i++) {
const double *Xci = Xc + i*%(Xc_cell_inc)d;
double celldist = 2*bestdist;
to_reference_coords_kernel(Xref, X, Xci);
if (%(inside_cell)s) {
cell = i;
break;
}
%(compute_celldist)s
/* fprintf(stderr, "cell %%d celldist: %%.14e\\n", i, celldist);
fprintf(stderr, "Xref: %%.14e %%.14e %%.14e\\n", Xref[0], Xref[1], Xref[2]); */
if (celldist < bestdist) {
bestdist = celldist;
bestcell = i;
}
}
if (cell == -1) {
/* We didn't find a cell that contained this point exactly.
Did we find one that was close enough? */
if (bestdist < 10) {
cell = bestcell;
} else {
fprintf(stderr, "Could not identify cell in transfer operator. Point: ");
for (int coord = 0; coord < %(spacedim)s; coord++) {
fprintf(stderr, "%%.14e ", X[coord]);
}
fprintf(stderr, "\\n");
fprintf(stderr, "Number of candidates: %%d. Best distance located: %%14e", %(ncandidate)d, bestdist);
abort();
}
}
{
const double *Ri = %(R)s + cell*%(coarse_cell_inc)d;
pyop2_kernel_evaluate(Ri, %(fine)s, Xref);
}
}
""" % {
"to_reference":
str(to_reference_kernel),
"evaluate":
str(evaluate_kernel),
"ncandidate":
hierarchy.fine_to_coarse_cells[levelf].shape[1],
"inside_cell":
inside_check(element.cell, eps=1e-8, X="Xref"),
"compute_celldist":
compute_celldist(element.cell, X="Xref", celldist="celldist"),
"Xc_cell_inc":
coords_element.space_dimension(),
"coarse_cell_inc":
element.space_dimension(),
"args":
args,
"spacedim":
element.cell.get_spatial_dimension(),
"R":
R,
"fine":
fine,
"tdim":
mesh.topological_dimension()
}
return cache.setdefault(
key, op2.Kernel(my_kernel, name="pyop2_kernel_restrict"))
def inject_kernel(Vf, Vc):
hierarchy, level = utils.get_level(Vc.ufl_domain())
cache = hierarchy._shared_data_cache["transfer_kernels"]
coordinates = Vf.ufl_domain().coordinates
key = (
("inject", ) + Vf.ufl_element().value_shape() +
entity_dofs_key(Vc.finat_element.entity_dofs()) +
entity_dofs_key(Vf.finat_element.entity_dofs()) + entity_dofs_key(
Vc.mesh().coordinates.function_space().finat_element.entity_dofs())
+ entity_dofs_key(
coordinates.function_space().finat_element.entity_dofs()))
try:
return cache[key]
except KeyError:
ncandidate = hierarchy.coarse_to_fine_cells[level].shape[1]
if Vc.finat_element.entity_dofs(
) == Vc.finat_element.entity_closure_dofs():
return cache.setdefault(
key, (dg_injection_kernel(Vf, Vc, ncandidate), True))
coordinates = Vf.ufl_domain().coordinates
evaluate_kernel = compile_element(ufl.Coefficient(Vf))
to_reference_kernel = to_reference_coordinates(
coordinates.ufl_element())
coords_element = create_element(coordinates.ufl_element())
Vf_element = create_element(Vf.ufl_element())
kernel = """
%(to_reference)s
%(evaluate)s
__attribute__((noinline)) /* Clang bug */
static void pyop2_kernel_inject(double *R, const double *X, const double *f, const double *Xf)
{
double Xref[%(tdim)d];
int cell = -1;
int bestcell = -1;
double bestdist = 1e10;
for (int i = 0; i < %(ncandidate)d; i++) {
const double *Xfi = Xf + i*%(Xf_cell_inc)d;
double celldist = 2*bestdist;
to_reference_coords_kernel(Xref, X, Xfi);
if (%(inside_cell)s) {
cell = i;
break;
}
%(compute_celldist)s
if (celldist < bestdist) {
bestdist = celldist;
bestcell = i;
}
}
if (cell == -1) {
/* We didn't find a cell that contained this point exactly.
Did we find one that was close enough? */
if (bestdist < 10) {
cell = bestcell;
} else {
fprintf(stderr, "Could not identify cell in transfer operator. Point: ");
for (int coord = 0; coord < %(spacedim)s; coord++) {
fprintf(stderr, "%%.14e ", X[coord]);
}
fprintf(stderr, "\\n");
fprintf(stderr, "Number of candidates: %%d. Best distance located: %%14e", %(ncandidate)d, bestdist);
abort();
}
}
const double *fi = f + cell*%(f_cell_inc)d;
for ( int i = 0; i < %(Rdim)d; i++ ) {
R[i] = 0;
}
pyop2_kernel_evaluate(R, fi, Xref);
}
""" % {
"to_reference":
str(to_reference_kernel),
"evaluate":
str(evaluate_kernel),
"inside_cell":
inside_check(Vc.finat_element.cell, eps=1e-8, X="Xref"),
"spacedim":
Vc.finat_element.cell.get_spatial_dimension(),
"compute_celldist":
compute_celldist(
Vc.finat_element.cell, X="Xref", celldist="celldist"),
"tdim":
Vc.ufl_domain().topological_dimension(),
"ncandidate":
ncandidate,
"Rdim":
numpy.prod(Vf_element.value_shape),
"Xf_cell_inc":
coords_element.space_dimension(),
"f_cell_inc":
Vf_element.space_dimension()
}
return cache.setdefault(
key, (op2.Kernel(kernel, name="pyop2_kernel_inject"), False))
def compile_element(expression, dual_space=None, parameters=None,
name="evaluate"):
"""Generate code for point evaluations.
:arg expression: A UFL expression (may contain up to one coefficient, or one argument)
:arg dual_space: if the expression has an argument, should we also distribute residual data?
:returns: Some coffee AST
"""
if parameters is None:
parameters = default_parameters()
else:
_ = default_parameters()
_.update(parameters)
parameters = _
expression = tsfc.ufl_utils.preprocess_expression(expression)
# # Collect required coefficients
try:
arg, = extract_coefficients(expression)
argument_multiindices = ()
coefficient = True
if expression.ufl_shape:
tensor_indices = tuple(gem.Index() for s in expression.ufl_shape)
else:
tensor_indices = ()
except ValueError:
arg, = extract_arguments(expression)
finat_elem = create_element(arg.ufl_element())
argument_multiindices = (finat_elem.get_indices(), )
argument_multiindex, = argument_multiindices
value_shape = finat_elem.value_shape
if value_shape:
tensor_indices = argument_multiindex[-len(value_shape):]
else:
tensor_indices = ()
coefficient = False
# Replace coordinates (if any)
builder = firedrake_interface.KernelBuilderBase()
domain = expression.ufl_domain()
# Translate to GEM
cell = domain.ufl_cell()
dim = cell.topological_dimension()
point = gem.Variable('X', (dim,))
point_arg = ast.Decl(SCALAR_TYPE, ast.Symbol('X', rank=(dim,)))
config = dict(interface=builder,
ufl_cell=cell,
precision=parameters["precision"],
point_indices=(),
point_expr=point,
argument_multiindices=argument_multiindices)
context = tsfc.fem.GemPointContext(**config)
# Abs-simplification
expression = tsfc.ufl_utils.simplify_abs(expression)
# Translate UFL -> GEM
if coefficient:
assert dual_space is None
f_arg = [builder._coefficient(arg, "f")]
else:
f_arg = []
translator = tsfc.fem.Translator(context)
result, = map_expr_dags(translator, [expression])
b_arg = []
if coefficient:
if expression.ufl_shape:
return_variable = gem.Indexed(gem.Variable('R', expression.ufl_shape), tensor_indices)
result_arg = ast.Decl(SCALAR_TYPE, ast.Symbol('R', rank=expression.ufl_shape))
result = gem.Indexed(result, tensor_indices)
else:
return_variable = gem.Indexed(gem.Variable('R', (1,)), (0,))
result_arg = ast.Decl(SCALAR_TYPE, ast.Symbol('R', rank=(1,)))
else:
return_variable = gem.Indexed(gem.Variable('R', finat_elem.index_shape), argument_multiindex)
result = gem.Indexed(result, tensor_indices)
if dual_space:
elem = create_element(dual_space.ufl_element())
if elem.value_shape:
var = gem.Indexed(gem.Variable("b", elem.value_shape),
tensor_indices)
b_arg = [ast.Decl(SCALAR_TYPE, ast.Symbol("b", rank=elem.value_shape))]
else:
var = gem.Indexed(gem.Variable("b", (1, )), (0, ))
b_arg = [ast.Decl(SCALAR_TYPE, ast.Symbol("b", rank=(1, )))]
result = gem.Product(result, var)
result_arg = ast.Decl(SCALAR_TYPE, ast.Symbol('R', rank=finat_elem.index_shape))
# Unroll
max_extent = parameters["unroll_indexsum"]
if max_extent:
def predicate(index):
return index.extent <= max_extent
result, = gem.optimise.unroll_indexsum([result], predicate=predicate)
# Translate GEM -> COFFEE
result, = gem.impero_utils.preprocess_gem([result])
impero_c = gem.impero_utils.compile_gem([(return_variable, result)], tensor_indices)
body = generate_coffee(impero_c, {}, parameters["precision"])
# Build kernel tuple
kernel_code = builder.construct_kernel("pyop2_kernel_" + name, [result_arg] + b_arg + f_arg + [point_arg], body)
return kernel_code
def dg_injection_kernel(Vf, Vc, ncell):
from firedrake import Tensor, AssembledVector, TestFunction, TrialFunction
from firedrake.slate.slac import compile_expression
macro_builder = MacroKernelBuilder(ncell)
f = ufl.Coefficient(Vf)
macro_builder.set_coefficients([f])
macro_builder.set_coordinates(Vf.mesh())
Vfe = create_element(Vf.ufl_element())
macro_quadrature_rule = make_quadrature(Vfe.cell, estimate_total_polynomial_degree(ufl.inner(f, f)))
index_cache = {}
parameters = default_parameters()
integration_dim, entity_ids = lower_integral_type(Vfe.cell, "cell")
macro_cfg = dict(interface=macro_builder,
ufl_cell=Vf.ufl_cell(),
precision=parameters["precision"],
integration_dim=integration_dim,
entity_ids=entity_ids,
index_cache=index_cache,
quadrature_rule=macro_quadrature_rule)
fexpr, = fem.compile_ufl(f, **macro_cfg)
X = ufl.SpatialCoordinate(Vf.mesh())
C_a, = fem.compile_ufl(X, **macro_cfg)
detJ = ufl_utils.preprocess_expression(abs(ufl.JacobianDeterminant(f.ufl_domain())))
macro_detJ, = fem.compile_ufl(detJ, **macro_cfg)
Vce = create_element(Vc.ufl_element())
coarse_builder = firedrake_interface.KernelBuilder("cell", "otherwise", 0)
coarse_builder.set_coordinates(Vc.mesh())
argument_multiindices = (Vce.get_indices(), )
argument_multiindex, = argument_multiindices
return_variable, = coarse_builder.set_arguments((ufl.TestFunction(Vc), ), argument_multiindices)
integration_dim, entity_ids = lower_integral_type(Vce.cell, "cell")
# Midpoint quadrature for jacobian on coarse cell.
quadrature_rule = make_quadrature(Vce.cell, 0)
coarse_cfg = dict(interface=coarse_builder,
ufl_cell=Vc.ufl_cell(),
precision=parameters["precision"],
integration_dim=integration_dim,
entity_ids=entity_ids,
index_cache=index_cache,
quadrature_rule=quadrature_rule)
X = ufl.SpatialCoordinate(Vc.mesh())
K = ufl_utils.preprocess_expression(ufl.JacobianInverse(Vc.mesh()))
C_0, = fem.compile_ufl(X, **coarse_cfg)
K, = fem.compile_ufl(K, **coarse_cfg)
i = gem.Index()
j = gem.Index()
C_0 = gem.Indexed(C_0, (j, ))
C_0 = gem.index_sum(C_0, quadrature_rule.point_set.indices)
C_a = gem.Indexed(C_a, (j, ))
X_a = gem.Sum(C_0, gem.Product(gem.Literal(-1), C_a))
K_ij = gem.Indexed(K, (i, j))
K_ij = gem.index_sum(K_ij, quadrature_rule.point_set.indices)
X_a = gem.index_sum(gem.Product(K_ij, X_a), (j, ))
C_0, = quadrature_rule.point_set.points
C_0 = gem.Indexed(gem.Literal(C_0), (i, ))
# fine quad points in coarse reference space.
X_a = gem.Sum(C_0, gem.Product(gem.Literal(-1), X_a))
X_a = gem.ComponentTensor(X_a, (i, ))
# Coarse basis function evaluated at fine quadrature points
phi_c = fem.fiat_to_ufl(Vce.point_evaluation(0, X_a, (Vce.cell.get_dimension(), 0)), 0)
tensor_indices = tuple(gem.Index(extent=d) for d in f.ufl_shape)
phi_c = gem.Indexed(phi_c, argument_multiindex + tensor_indices)
fexpr = gem.Indexed(fexpr, tensor_indices)
quadrature_weight = macro_quadrature_rule.weight_expression
expr = gem.Product(gem.IndexSum(gem.Product(phi_c, fexpr), tensor_indices),
gem.Product(macro_detJ, quadrature_weight))
quadrature_indices = macro_builder.indices + macro_quadrature_rule.point_set.indices
reps = spectral.Integrals([expr], quadrature_indices, argument_multiindices, parameters)
assignments = spectral.flatten([(return_variable, reps)], index_cache)
return_variables, expressions = zip(*assignments)
expressions = impero_utils.preprocess_gem(expressions, **spectral.finalise_options)
assignments = list(zip(return_variables, expressions))
impero_c = impero_utils.compile_gem(assignments, quadrature_indices + argument_multiindex,
remove_zeros=True)
index_names = []
def name_index(index, name):
index_names.append((index, name))
if index in index_cache:
for multiindex, suffix in zip(index_cache[index],
string.ascii_lowercase):
name_multiindex(multiindex, name + suffix)
def name_multiindex(multiindex, name):
if len(multiindex) == 1:
name_index(multiindex[0], name)
else:
for i, index in enumerate(multiindex):
name_index(index, name + str(i))
name_multiindex(quadrature_indices, 'ip')
for multiindex, name in zip(argument_multiindices, ['j', 'k']):
name_multiindex(multiindex, name)
index_names.extend(zip(macro_builder.indices, ["entity"]))
body = generate_coffee(impero_c, index_names, parameters["precision"])
retarg = ast.Decl(SCALAR_TYPE, ast.Symbol("R", rank=(Vce.space_dimension(), )))
local_tensor = coarse_builder.local_tensor
local_tensor.init = ast.ArrayInit(numpy.zeros(Vce.space_dimension(), dtype=SCALAR_TYPE))
body.children.insert(0, local_tensor)
args = [retarg] + macro_builder.kernel_args + [macro_builder.coordinates_arg,
coarse_builder.coordinates_arg]
# Now we have the kernel that computes <f, phi_c>dx_c
# So now we need to hit it with the inverse mass matrix on dx_c
u = TrialFunction(Vc)
v = TestFunction(Vc)
expr = Tensor(ufl.inner(u, v)*ufl.dx).inv * AssembledVector(ufl.Coefficient(Vc))
Ainv, = compile_expression(expr)
Ainv = Ainv.kinfo.kernel
A = ast.Symbol(local_tensor.sym.symbol)
R = ast.Symbol("R")
body.children.append(ast.FunCall(Ainv.name, R, coarse_builder.coordinates_arg.sym, A))
from coffee.base import Node
assert isinstance(Ainv._code, Node)
return op2.Kernel(ast.Node([Ainv._code,
ast.FunDecl("void", "pyop2_kernel_injection_dg", args, body,
pred=["static", "inline"])]),
name="pyop2_kernel_injection_dg",
cpp=True,
include_dirs=Ainv._include_dirs,
headers=Ainv._headers)
def test_quadrilateral_variant_spectral_dq():
element = create_element(
ufl.FiniteElement('DQ', ufl.quadrilateral, 1, variant='spectral'))
assert isinstance(element.product.factors[0], finat.GaussLegendre)
assert isinstance(element.product.factors[1], finat.GaussLegendre)
def dimension(e):
return create_element(e).space_dimension()
def make_extruded_coords(extruded_topology,
base_coords,
ext_coords,
layer_height,
extrusion_type='uniform',
kernel=None):
"""
Given either a kernel or a (fixed) layer_height, compute an
extruded coordinate field for an extruded mesh.
:arg extruded_topology: an :class:`~.ExtrudedMeshTopology` to extrude
a coordinate field for.
:arg base_coords: a :class:`~.Function` to read the base
coordinates from.
:arg ext_coords: a :class:`~.Function` to write the extruded
coordinates into.
:arg layer_height: the height for each layer. Either a scalar,
where layers will be equi-spaced at the specified height, or a
1D array of variable layer heights to use through the extrusion.
:arg extrusion_type: the type of extrusion to use. Predefined
options are either "uniform" (creating equi-spaced layers by
extruding in the (n+1)dth direction), "radial" (creating
equi-spaced layers by extruding in the outward direction from
the origin) or "radial_hedgehog" (creating equi-spaced layers
by extruding coordinates in the outward cell-normal
direction, needs a P1dgxP1 coordinate field).
:arg kernel: an optional kernel to carry out coordinate extrusion.
The kernel signature (if provided) is::
void kernel(double **base_coords, double **ext_coords,
double *layer_height, int layer)
The kernel iterates over the cells of the mesh and receives as
arguments the coordinates of the base cell (to read), the
coordinates on the extruded cell (to write to), the fixed layer
height, and the current cell layer.
"""
_, vert_space = ext_coords.function_space().ufl_element().sub_elements(
)[0].sub_elements()
if kernel is None and not (vert_space.degree() == 1
and vert_space.family()
in ['Lagrange', 'Discontinuous Lagrange']):
raise RuntimeError(
'Extrusion of coordinates is only possible for a P1 or P1dg interval unless a custom kernel is provided'
)
layer_height = numpy.atleast_1d(numpy.array(layer_height, dtype=RealType))
if layer_height.ndim > 1:
raise RuntimeError('Extrusion layer height should be 1d or scalar')
if layer_height.size > 1:
layer_height = numpy.cumsum(numpy.concatenate(([0], layer_height)))
layer_heights = layer_height.size
layer_height = op2.Global(layer_heights, layer_height, dtype=RealType)
if kernel is not None:
op2.ParLoop(kernel,
ext_coords.cell_set,
ext_coords.dat(op2.WRITE, ext_coords.cell_node_map()),
base_coords.dat(op2.READ, base_coords.cell_node_map()),
layer_height(op2.READ),
pass_layer_arg=True,
is_loopy_kernel=True).compute()
return
ext_fe = create_element(ext_coords.ufl_element())
ext_shape = ext_fe.index_shape
base_fe = create_element(base_coords.ufl_element())
base_shape = base_fe.index_shape
data = []
data.append(lp.GlobalArg("ext_coords", dtype=ScalarType, shape=ext_shape))
data.append(lp.GlobalArg("base_coords", dtype=ScalarType,
shape=base_shape))
data.append(
lp.GlobalArg("layer_height", dtype=RealType, shape=(layer_heights, )))
data.append(lp.ValueArg('layer'))
base_coord_dim = base_coords.function_space().value_size
# Deal with tensor product cells
adim = len(ext_shape) - 2
# handle single or variable layer heights
if layer_heights == 1:
height_var = "layer_height[0] * (layer + l)"
else:
height_var = "layer_height[layer + l]"
def _get_arity_axis_inames(_base):
return tuple(_base + str(i) for i in range(adim))
def _get_lp_domains(_inames, _extents):
domains = []
for idx, extent in zip(_inames, _extents):
inames = isl.make_zero_and_vars([idx])
domains.append(((inames[0].le_set(inames[idx])) &
(inames[idx].lt_set(inames[0] + extent))))
return domains
if extrusion_type == 'uniform':
domains = []
dd = _get_arity_axis_inames('d')
domains.extend(_get_lp_domains(dd, ext_shape[:adim]))
domains.extend(_get_lp_domains(('c', ), (base_coord_dim, )))
if layer_heights == 1:
domains.extend(_get_lp_domains(('l', ), (2, )))
else:
domains.append(
"[layer] -> { [l] : 0 <= l <= 1 & 0 <= l + layer < %d}" %
layer_heights)
instructions = """
ext_coords[{dd}, l, c] = base_coords[{dd}, c]
ext_coords[{dd}, l, {base_coord_dim}] = ({hv})
""".format(dd=', '.join(dd),
base_coord_dim=base_coord_dim,
hv=height_var)
name = "pyop2_kernel_uniform_extrusion"
elif extrusion_type == 'radial':
domains = []
dd = _get_arity_axis_inames('d')
domains.extend(_get_lp_domains(dd, ext_shape[:adim]))
domains.extend(_get_lp_domains(('c', 'k'), (base_coord_dim, ) * 2))
if layer_heights == 1:
domains.extend(_get_lp_domains(('l', ), (2, )))
else:
domains.append(
"[layer] -> { [l] : 0 <= l <= 1 & 0 <= l + layer < %d}" %
layer_heights)
instructions = """
<{RealType}> tt[{dd}] = 0
<{RealType}> bc[{dd}] = 0
for k
bc[{dd}] = real(base_coords[{dd}, k])
tt[{dd}] = tt[{dd}] + bc[{dd}] * bc[{dd}]
end
tt[{dd}] = sqrt(tt[{dd}])
ext_coords[{dd}, l, c] = base_coords[{dd}, c] + base_coords[{dd}, c] * ({hv}) / tt[{dd}]
""".format(RealType=RealType, dd=', '.join(dd), hv=height_var)
name = "pyop2_kernel_radial_extrusion"
elif extrusion_type == 'radial_hedgehog':
# Only implemented for interval in 2D and triangle in 3D.
# gdim != tdim already checked in ExtrudedMesh constructor.
tdim = base_coords.ufl_domain().ufl_cell().topological_dimension()
if tdim not in [1, 2]:
raise NotImplementedError(
"Hedgehog extrusion not implemented for %s" %
base_coords.ufl_domain().ufl_cell())
# tdim == 1:
#
# normal is:
# (0 -1) (x2 - x1)
# (1 0) (y2 - y1)
#
# tdim == 2:
# normal is
# v0 x v1
#
# /\
# v0/ \
# / \
# /------\
# v1
domains = []
dd = _get_arity_axis_inames('d')
_dd = _get_arity_axis_inames('_d')
domains.extend(_get_lp_domains(dd, ext_shape[:adim]))
domains.extend(_get_lp_domains(_dd, ext_shape[:adim]))
domains.extend(
_get_lp_domains(('c0', 'c1', 'c2', 'c3', 'k', 'l'),
(base_coord_dim, ) * 5 + (2, )))
# Formula for normal, n
n_1_1 = """
n[0] = -bc[1, 1] + bc[0, 1]
n[1] = bc[1, 0] - bc[0, 0]
"""
n_2_1 = """
v0[c3] = bc[1, c3] - bc[0, c3]
v1[c3] = bc[2, c3] - bc[0, c3]
n[0] = v0[1] * v1[2] - v0[2] * v1[1]
n[1] = v0[2] * v1[0] - v0[0] * v1[2]
n[2] = v0[0] * v1[1] - v0[1] * v1[0]
"""
n_2_2 = """
v0[c3] = bc[0, 1, c3] - bc[0, 0, c3]
v1[c3] = bc[1, 0, c3] - bc[0, 0, c3]
n[0] = v0[1] * v1[2] - v0[2] * v1[1]
n[1] = v0[2] * v1[0] - v0[0] * v1[2]
n[2] = v0[0] * v1[1] - v0[1] * v1[0]
"""
n_dict = {1: {1: n_1_1}, 2: {1: n_2_1, 2: n_2_2}}
instructions = """
<{RealType}> dot = 0
<{RealType}> norm = 0
<{RealType}> v0[c2] = 0
<{RealType}> v1[c2] = 0
<{RealType}> n[c2] = 0
<{RealType}> x[c2] = 0
<{RealType}> bc[{_dd}, c1] = real(base_coords[{_dd}, c1])
for {_dd}
x[c1] = x[c1] + bc[{_dd}, c1]
end
{ninst}
for k
dot = dot + x[k] * n[k]
norm = norm + n[k] * n[k]
end
norm = sqrt(norm)
norm = -norm if dot < 0 else norm
ext_coords[{dd}, l, c0] = base_coords[{dd}, c0] + n[c0] * ({hv}) / norm
""".format(RealType=RealType,
dd=', '.join(dd),
_dd=', '.join(_dd),
ninst=n_dict[tdim][adim],
hv=height_var)
name = "pyop2_kernel_radial_hedgehog_extrusion"
else:
raise NotImplementedError('Unsupported extrusion type "%s"' %
extrusion_type)
ast = lp.make_function(domains,
instructions,
data,
name=name,
target=lp.CTarget(),
seq_dependencies=True,
silenced_warnings=["summing_if_branches_ops"])
kernel = op2.Kernel(ast, name)
op2.ParLoop(kernel,
ext_coords.cell_set,
ext_coords.dat(op2.WRITE, ext_coords.cell_node_map()),
base_coords.dat(op2.READ, base_coords.cell_node_map()),
layer_height(op2.READ),
pass_layer_arg=True,
is_loopy_kernel=True).compute()
def create_element(self, element, **kwargs):
"""Create a FInAT element (suitable for tabulating with) given
a UFL element."""
return create_element(element, **kwargs)
def prepare_arguments(arguments, multiindices, interior_facet=False):
"""Bridges the kernel interface and the GEM abstraction for
Arguments. Vector Arguments are rearranged here for interior
facet integrals.
:arg arguments: UFL Arguments
:arg multiindices: Argument multiindices
:arg interior_facet: interior facet integral?
:returns: (funarg, prepare, expressions)
funarg - :class:`coffee.Decl` function argument
prepare - list of COFFEE nodes to be prepended to the
kernel body
expressions - GEM expressions referring to the argument
tensor
"""
funarg = coffee.Decl(SCALAR_TYPE, coffee.Symbol("A"), pointers=[()])
varexp = gem.Variable("A", (None,))
if len(arguments) == 0:
# No arguments
zero = coffee.FlatBlock(
"memset({name}, 0, sizeof(*{name}));\n".format(name=funarg.sym.gencode())
)
return funarg, [zero], [gem.reshape(varexp, ())]
elements = tuple(create_element(arg.ufl_element()) for arg in arguments)
transposed_shapes = []
transposed_indices = []
for element, multiindex in zip(elements, multiindices):
if isinstance(element, TensorFiniteElement):
scalar_shape = element.base_element.index_shape
tensor_shape = element.index_shape[len(scalar_shape):]
else:
scalar_shape = element.index_shape
tensor_shape = ()
transposed_shapes.append(tensor_shape + scalar_shape)
scalar_rank = len(scalar_shape)
transposed_indices.extend(multiindex[scalar_rank:] + multiindex[:scalar_rank])
transposed_indices = tuple(transposed_indices)
def expression(restricted):
return gem.Indexed(gem.reshape(restricted, *transposed_shapes),
transposed_indices)
u_shape = numpy.array([numpy.prod(element.index_shape, dtype=int)
for element in elements])
if interior_facet:
c_shape = tuple(2 * u_shape)
slicez = [[slice(r * s, (r + 1) * s)
for r, s in zip(restrictions, u_shape)]
for restrictions in product((0, 1), repeat=len(arguments))]
else:
c_shape = tuple(u_shape)
slicez = [[slice(s) for s in u_shape]]
expressions = [expression(gem.view(gem.reshape(varexp, c_shape), *slices))
for slices in slicez]
zero = coffee.FlatBlock(
str.format("memset({name}, 0, {size} * sizeof(*{name}));\n",
name=funarg.sym.gencode(), size=numpy.product(c_shape, dtype=int))
)
return funarg, [zero], prune(expressions)
def prolong_kernel(expression):
hierarchy, level = utils.get_level(expression.ufl_domain())
levelf = level + Fraction(1 / hierarchy.refinements_per_level)
cache = hierarchy._shared_data_cache["transfer_kernels"]
coordinates = expression.ufl_domain().coordinates
key = (("prolong", ) + expression.ufl_element().value_shape() +
entity_dofs_key(
expression.function_space().finat_element.entity_dofs()) +
entity_dofs_key(
coordinates.function_space().finat_element.entity_dofs()))
try:
return cache[key]
except KeyError:
mesh = coordinates.ufl_domain()
evaluate_kernel = compile_element(expression)
to_reference_kernel = to_reference_coordinates(
coordinates.ufl_element())
element = create_element(expression.ufl_element())
eval_args = evaluate_kernel.args[:-1]
coords_element = create_element(coordinates.ufl_element())
args = eval_args[-1].gencode(not_scope=True)
R, coarse = (a.sym.symbol for a in eval_args)
my_kernel = """
%(to_reference)s
%(evaluate)s
void prolong_kernel(double *R, %(args)s, const double *X, const double *Xc)
{
double Xref[%(tdim)d];
int cell = -1;
for (int i = 0; i < %(ncandidate)d; i++) {
const double *Xci = Xc + i*%(Xc_cell_inc)d;
to_reference_coords_kernel(Xref, X, Xci);
if (%(inside_cell)s) {
cell = i;
break;
}
}
if (cell == -1) abort();
const double *coarsei = %(coarse)s + cell*%(coarse_cell_inc)d;
for ( int i = 0; i < %(Rdim)d; i++ ) {
%(R)s[i] = 0;
}
evaluate_kernel(%(R)s, coarsei, Xref);
}
""" % {
"to_reference": str(to_reference_kernel),
"evaluate": str(evaluate_kernel),
"args": args,
"R": R,
"coarse": coarse,
"ncandidate": hierarchy.fine_to_coarse_cells[levelf].shape[1],
"Rdim": numpy.prod(element.value_shape),
"inside_cell": inside_check(element.cell, eps=1e-8, X="Xref"),
"Xc_cell_inc": coords_element.space_dimension(),
"coarse_cell_inc": element.space_dimension(),
"tdim": mesh.topological_dimension()
}
return cache.setdefault(key,
op2.Kernel(my_kernel, name="prolong_kernel"))
def prepare_coefficients(coefficients, num, name, interior_facet=False):
"""Bridges the kernel interface and the GEM abstraction for
Coefficients.
:arg coefficients: split UFL Coefficients
:arg num: coefficient index in the original form
:arg name: unique name to refer to the Coefficient in the kernel
:arg interior_facet: interior facet integral?
:returns: GEM expression referring to the Coefficient value
"""
varexp = gem.Variable(name, (None, None))
if len(coefficients) == 1 and coefficients[0].ufl_element().family() == 'Real':
coefficient, = coefficients
size = numpy.prod(coefficient.ufl_shape, dtype=int)
data = gem.view(varexp, slice(num, num + 1), slice(size))
expression = gem.reshape(data, (), coefficient.ufl_shape)
return [expression]
elements = [create_element(coeff.ufl_element()) for coeff in coefficients]
space_dimensions = [numpy.prod(element.index_shape, dtype=int)
for element in elements]
ends = list(numpy.cumsum(space_dimensions))
starts = [0] + ends[:-1]
slices = [slice(start, end) for start, end in zip(starts, ends)]
transposed_shapes = []
tensor_ranks = []
for element in elements:
if isinstance(element, TensorFiniteElement):
scalar_shape = element.base_element.index_shape
tensor_shape = element.index_shape[len(scalar_shape):]
else:
scalar_shape = element.index_shape
tensor_shape = ()
transposed_shapes.append(tensor_shape + scalar_shape)
tensor_ranks.append(len(tensor_shape))
def transpose(expr, rank):
assert not expr.free_indices
assert 0 <= rank < len(expr.shape)
if rank == 0:
return expr
else:
indices = tuple(gem.Index(extent=extent) for extent in expr.shape)
transposed_indices = indices[rank:] + indices[:rank]
return gem.ComponentTensor(gem.Indexed(expr, indices),
transposed_indices)
def expressions(data):
return prune([transpose(gem.reshape(gem.view(data, slice_), shape), rank)
for slice_, shape, rank in zip(slices, transposed_shapes, tensor_ranks)])
size = sum(space_dimensions)
if not interior_facet:
data = gem.view(varexp, slice(num, num + 1), slice(size))
return expressions(gem.reshape(data, (), (size,)))
else:
data_p = gem.view(varexp, slice(num, num + 1), slice(size))
data_m = gem.view(varexp, slice(num, num + 1), slice(size, 2 * size))
return list(zip(expressions(gem.reshape(data_p, (), (size,))),
expressions(gem.reshape(data_m, (), (size,)))))
