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 prepare_coefficient(coefficient, num, name, interior_facet=False):
"""Bridges the kernel interface and the GEM abstraction for
Coefficients.
:arg coefficient: UFL Coefficient
: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 coefficient.ufl_element().family() == 'Real':
size = numpy.prod(coefficient.ufl_shape, dtype=int)
data = gem.view(varexp, slice(num, num + 1), slice(size))
return gem.reshape(data, (), coefficient.ufl_shape)
element = create_element(coefficient.ufl_element())
size = numpy.prod(element.index_shape, dtype=int)
def expression(data):
result, = prune(
[gem.reshape(gem.view(data, slice(size)), element.index_shape)])
return result
if not interior_facet:
data = gem.view(varexp, slice(num, num + 1), slice(size))
return expression(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 (expression(gem.reshape(data_p, (), (size, ))),
expression(gem.reshape(data_m, (), (size, ))))
def prepare_coefficient(coefficient, num, name, interior_facet=False):
"""Bridges the kernel interface and the GEM abstraction for
Coefficients.
:arg coefficient: UFL Coefficient
: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 coefficient.ufl_element().family() == 'Real':
size = numpy.prod(coefficient.ufl_shape, dtype=int)
data = gem.view(varexp, slice(num, num + 1), slice(size))
return gem.reshape(data, (), coefficient.ufl_shape)
element = create_element(coefficient.ufl_element())
size = numpy.prod(element.index_shape, dtype=int)
def expression(data):
result, = prune([gem.reshape(gem.view(data, slice(size)), element.index_shape)])
return result
if not interior_facet:
data = gem.view(varexp, slice(num, num + 1), slice(size))
return expression(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 (expression(gem.reshape(data_p, (), (size,))),
expression(gem.reshape(data_m, (), (size,))))
def test_reshape_reshape(vector):
expression = gem.reshape(gem.reshape(vector, (4, 3)), (2, 2), (3,))
assert expression.shape == (2, 2, 3)
actual = [convert(expression, multiindex)
for multiindex in numpy.ndindex(expression.shape)]
assert [(i,) for i in range(12)] == actual
def test_reshape_reshape(vector):
expression = gem.reshape(gem.reshape(vector, (4, 3)), (2, 2), (3,))
assert expression.shape == (2, 2, 3)
actual = [convert(expression, multiindex)
for multiindex in numpy.ndindex(expression.shape)]
assert [(i,) for i in range(12)] == actual
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, 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)
shapes = [element.index_shape for element in elements]
indices = tuple(chain(*multiindices))
def expression(restricted):
return gem.Indexed(gem.reshape(restricted, *shapes), 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 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 prepare_coordinates(coefficient, name, scalar_type, 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)
assert isinstance(finat_element, TensorFiniteElement)
scalar_shape = finat_element.base_element.index_shape
tensor_shape = finat_element._shape
transposed_shape = scalar_shape + tensor_shape
scalar_rank = len(scalar_shape)
def transpose(expr):
indices = tuple(gem.Index(extent=extent) for extent in expr.shape)
transposed_indices = indices[scalar_rank:] + indices[:scalar_rank]
return gem.ComponentTensor(gem.Indexed(expr, indices),
transposed_indices)
if not interior_facet:
funargs = [
coffee.Decl(scalar_type,
coffee.Symbol(name),
pointers=[("", )],
qualifiers=["const"])
]
variable = gem.Variable(name, (size, ))
expression = transpose(gem.reshape(variable, transposed_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 = (transpose(gem.reshape(variable0, transposed_shape)),
transpose(gem.reshape(variable1, transposed_shape)))
return funargs, expression
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)
shapes = [element.index_shape for element in elements]
indices = tuple(chain(*multiindices))
def expression(restricted):
return gem.Indexed(gem.reshape(restricted, *shapes), 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 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)
assert isinstance(finat_element, TensorFiniteElement)
scalar_shape = finat_element.base_element.index_shape
tensor_shape = finat_element._shape
transposed_shape = scalar_shape + tensor_shape
scalar_rank = len(scalar_shape)
def transpose(expr):
indices = tuple(gem.Index(extent=extent) for extent in expr.shape)
transposed_indices = indices[scalar_rank:] + indices[:scalar_rank]
return gem.ComponentTensor(gem.Indexed(expr, indices),
transposed_indices)
if not interior_facet:
funargs = [coffee.Decl(SCALAR_TYPE, coffee.Symbol(name),
pointers=[("",)],
qualifiers=["const"])]
variable = gem.Variable(name, (size,))
expression = transpose(gem.reshape(variable, transposed_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 = (transpose(gem.reshape(variable0, transposed_shape)),
transpose(gem.reshape(variable1, transposed_shape)))
return funargs, expression
def test_view_reshape(vector):
expression = gem.view(gem.reshape(vector, (3, 4)), slice(2), slice(1, 3))
assert expression.shape == (2, 2)
actual = [convert(expression, multiindex)
for multiindex in numpy.ndindex(expression.shape)]
assert [(1,), (2,), (5,), (6,)] == actual
def test_view_reshape(vector):
expression = gem.view(gem.reshape(vector, (3, 4)), slice(2), slice(1, 3))
assert expression.shape == (2, 2)
actual = [convert(expression, multiindex)
for multiindex in numpy.ndindex(expression.shape)]
assert [(1,), (2,), (5,), (6,)] == actual
def test_reshape_shape(vector):
expression = gem.reshape(gem.view(vector, slice(5, 11)), (3, 2))
assert expression.shape == (3, 2)
actual = [convert(expression, multiindex)
for multiindex in numpy.ndindex(expression.shape)]
assert [(i,) for i in range(5, 11)] == actual
def test_reshape_shape(vector):
expression = gem.reshape(gem.view(vector, slice(5, 11)), (3, 2))
assert expression.shape == (3, 2)
actual = [convert(expression, multiindex)
for multiindex in numpy.ndindex(expression.shape)]
assert [(i,) for i in range(5, 11)] == actual
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 _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())
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 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 expression(data):
result, = prune(
[gem.reshape(gem.view(data, slice(size)), element.index_shape)])
return result
def expression(restricted):
return gem.Indexed(gem.reshape(restricted, *transposed_shapes),
transposed_indices)
def prepare_coefficient(coefficient, name, mode=None, interior_facet=False):
"""Bridges the kernel interface and the GEM abstraction for
Coefficients. Mixed element Coefficients are rearranged here for
interior facet integrals.
:arg coefficient: UFL Coefficient
:arg name: unique name to refer to the Coefficient in the kernel
:arg mode: 'manual_loop' or 'list_tensor'; two ways to deal with
interior facet integrals on mixed elements
:arg interior_facet: interior facet integral?
:returns: (funarg, prepare, expression)
funarg - :class:`coffee.Decl` function argument
prepare - list of COFFEE nodes to be prepended to the
kernel body
expression - GEM expression referring to the Coefficient
values
"""
if mode is None:
mode = 'manual_loop'
assert mode in ['manual_loop', 'list_tensor']
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
fiat_element = create_element(coefficient.ufl_element())
size = fiat_element.space_dimension()
if not interior_facet:
# Simple case
funarg = coffee.Decl(SCALAR_TYPE, coffee.Symbol(name),
pointers=[("const", "restrict"), ("restrict",)],
qualifiers=["const"])
expression = gem.reshape(gem.Variable(name, (size, 1)), (size,), ())
return funarg, [], expression
if not isinstance(fiat_element, MixedElement):
# Interior facet integral
funarg = coffee.Decl(SCALAR_TYPE, coffee.Symbol(name),
pointers=[("const", "restrict"), ("restrict",)],
qualifiers=["const"])
expression = gem.reshape(
gem.Variable(name, (2 * size, 1)), (2, size), ()
)
return funarg, [], expression
# Interior facet integral + mixed / vector element
# Here we need to reorder the coefficient values.
#
# Incoming ordering: E1+ E1- E2+ E2- E3+ E3-
# Required ordering: E1+ E2+ E3+ E1- E2- E3-
#
# Each of E[n]{+,-} is a vector of basis function coefficients for
# subelement E[n].
#
# There are two code generation method to reorder the values.
# We have not done extensive research yet as to which way yield
# faster code.
if mode == 'manual_loop':
# In this case we generate loops outside the GEM abstraction
# to reorder the values. A whole E[n]{+,-} block is copied by
# a single loop.
name_ = name + "_"
shape = (2, size)
funarg = coffee.Decl(SCALAR_TYPE, coffee.Symbol(name_),
pointers=[("const", "restrict"), ("restrict",)],
qualifiers=["const"])
prepare = [coffee.Decl(SCALAR_TYPE, coffee.Symbol(name, rank=shape))]
expression = gem.Variable(name, shape)
offset = 0
i = coffee.Symbol("i")
for element in fiat_element.elements():
space_dim = element.space_dimension()
loop_body = coffee.Assign(coffee.Symbol(name, rank=(0, "i"),
offset=((1, 0), (1, offset))),
coffee.Symbol(name_, rank=("i", 0),
offset=((1, 2 * offset), (1, 0))))
prepare.append(coffee_for(i, space_dim, loop_body))
loop_body = coffee.Assign(coffee.Symbol(name, rank=(1, "i"),
offset=((1, 0), (1, offset))),
coffee.Symbol(name_, rank=("i", 0),
offset=((1, 2 * offset + space_dim), (1, 0))))
prepare.append(coffee_for(i, space_dim, loop_body))
offset += space_dim
return funarg, prepare, expression
elif mode == 'list_tensor':
# In this case we generate a gem.ListTensor to do the
# reordering. Every single element in a E[n]{+,-} block is
# referenced separately.
funarg = coffee.Decl(SCALAR_TYPE, coffee.Symbol(name),
pointers=[("const", "restrict"), ("restrict",)],
qualifiers=["const"])
variable = gem.Variable(name, (2 * size, 1))
facet_0 = []
facet_1 = []
offset = 0
for element in fiat_element.elements():
space_dim = element.space_dimension()
for i in range(offset, offset + space_dim):
facet_0.append(gem.Indexed(variable, (i, 0)))
offset += space_dim
for i in range(offset, offset + space_dim):
facet_1.append(gem.Indexed(variable, (i, 0)))
offset += space_dim
expression = gem.ListTensor(numpy.array([facet_0, facet_1]))
return funarg, [], expression
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)])
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,)))))
def expression(restricted):
return gem.Indexed(gem.reshape(restricted, *shapes),
tuple(chain(*multiindices)))
def expression(restricted):
return gem.Indexed(gem.reshape(restricted, *shapes),
tuple(chain(*multiindices)))
def expression(data):
result, = prune([gem.reshape(gem.view(data, slice(size)), element.index_shape)])
return result
def expression(restricted):
return gem.Indexed(gem.reshape(restricted, *shapes), indices)
def expression(restricted):
return gem.Indexed(gem.reshape(restricted, *shapes), indices)
