def tensor_assembly_calls(builder):
"""Generates a block of statements for assembling the local
finite element tensors.
:arg builder: The :class:`LocalKernelBuilder` containing
all relevant expression information and assembly calls.
"""
assembly_calls = builder.assembly_calls
statements = [ast.FlatBlock("/* Assemble local tensors */\n")]
# Cell integrals are straightforward. Just splat them out.
statements.extend(assembly_calls["cell"])
if builder.needs_cell_facets:
# The for-loop will have the general structure:
#
# FOR (facet=0; facet<num_facets; facet++):
# IF (facet is interior):
# *interior calls
# ELSE IF (facet is exterior):
# *exterior calls
#
# If only interior (exterior) facets are present,
# then only a single IF-statement checking for interior
# (exterior) facets will be present within the loop. The
# cell facets are labelled `1` for interior, and `0` for
# exterior.
statements.append(ast.FlatBlock("/* Loop over cell facets */\n"))
int_calls = list(
chain(*[
assembly_calls[it_type]
for it_type in ("interior_facet", "interior_facet_vert")
]))
ext_calls = list(
chain(*[
assembly_calls[it_type]
for it_type in ("exterior_facet", "exterior_facet_vert")
]))
# Generate logical statements for handling exterior/interior facet
# integrals on subdomains.
# Currently only facet integrals are supported.
for sd_type in ("subdomains_exterior_facet",
"subdomains_interior_facet"):
stmts = []
for sd, sd_calls in groupby(assembly_calls[sd_type],
lambda x: x[0]):
_, calls = zip(*sd_calls)
if_sd = ast.Eq(
ast.Symbol(builder.cell_facet_sym,
rank=(builder.it_sym, 1)), sd)
stmts.append(
ast.If(if_sd, (ast.Block(calls, open_scope=True), )))
if sd_type == "subdomains_exterior_facet":
ext_calls.extend(stmts)
if sd_type == "subdomains_interior_facet":
int_calls.extend(stmts)
# Compute the number of facets to loop over
domain = builder.expression.ufl_domain()
if domain.cell_set._extruded:
num_facets = domain.ufl_cell()._cells[0].num_facets()
else:
num_facets = domain.ufl_cell().num_facets()
if_ext = ast.Eq(
ast.Symbol(builder.cell_facet_sym, rank=(builder.it_sym, 0)), 0)
if_int = ast.Eq(
ast.Symbol(builder.cell_facet_sym, rank=(builder.it_sym, 0)), 1)
body = []
if ext_calls:
body.append(
ast.If(if_ext, (ast.Block(ext_calls, open_scope=True), )))
if int_calls:
body.append(
ast.If(if_int, (ast.Block(int_calls, open_scope=True), )))
statements.append(
ast.For(ast.Decl("unsigned int", builder.it_sym, init=0),
ast.Less(builder.it_sym, num_facets),
ast.Incr(builder.it_sym, 1), body))
if builder.needs_mesh_layers:
# In the presence of interior horizontal facet calls, an
# IF-ELIF-ELSE block is generated using the mesh levels
# as conditions for which calls are needed:
#
# IF (layer == bottom_layer):
# *bottom calls
# ELSE IF (layer == top_layer):
# *top calls
# ELSE:
# *top calls
# *bottom calls
#
# Any extruded top or bottom calls for extruded facets are
# included within the appropriate mesh-level IF-blocks. If
# no interior horizontal facet calls are present, then
# standard IF-blocks are generated for exterior top/bottom
# facet calls when appropriate:
#
# IF (layer == bottom_layer):
# *bottom calls
#
# IF (layer == top_layer):
# *top calls
#
# The mesh level is an integer provided as a macro kernel
# argument.
# FIXME: No variable layers assumption
statements.append(ast.FlatBlock("/* Mesh levels: */\n"))
num_layers = ast.Symbol(builder.mesh_layer_count_sym, rank=(0, ))
layer = builder.mesh_layer_sym
types = [
"interior_facet_horiz_top", "interior_facet_horiz_bottom",
"exterior_facet_top", "exterior_facet_bottom"
]
decide = [
ast.Less(layer, num_layers),
ast.Greater(layer, 0),
ast.Eq(layer, num_layers),
ast.Eq(layer, 0)
]
for (integral_type, which) in zip(types, decide):
statements.append(
ast.If(which, (ast.Block(assembly_calls[integral_type],
open_scope=True), )))
return statements
_to_sum = lambda o: ast.Sum(_ast(o[0]), _to_sum(o[1:])) if len(
o) > 1 else _ast(o[0])
_to_prod = lambda o: ast.Prod(_ast(o[0]), _to_sum(o[1:])) if len(
o) > 1 else _ast(o[0])
_to_aug_assign = lambda op, o: op(_ast(o[0]), _ast(o[1]))
_ast_map = {
MathFunction:
(lambda e: ast.FunCall(e._name, *[_ast(o) for o in e.ufl_operands])),
ufl.algebra.Sum: (lambda e: ast.Par(_to_sum(e.ufl_operands))),
ufl.algebra.Product: (lambda e: ast.Par(_to_prod(e.ufl_operands))),
ufl.algebra.Division:
(lambda e: ast.Par(ast.Div(*[_ast(o) for o in e.ufl_operands]))),
ufl.algebra.Abs: (lambda e: ast.FunCall("abs", _ast(e.ufl_operands[0]))),
Assign: (lambda e: _to_aug_assign(e._ast, e.ufl_operands)),
AugmentedAssignment: (lambda e: _to_aug_assign(e._ast, e.ufl_operands)),
ufl.constantvalue.ScalarValue: (lambda e: ast.Symbol(e._value)),
ufl.constantvalue.Zero: (lambda e: ast.Symbol(0)),
ufl.classes.Conditional:
(lambda e: ast.Ternary(*[_ast(o) for o in e.ufl_operands])),
ufl.classes.EQ: (lambda e: ast.Eq(*[_ast(o) for o in e.ufl_operands])),
ufl.classes.NE: (lambda e: ast.NEq(*[_ast(o) for o in e.ufl_operands])),
ufl.classes.LT: (lambda e: ast.Less(*[_ast(o) for o in e.ufl_operands])),
ufl.classes.LE: (lambda e: ast.LessEq(*[_ast(o) for o in e.ufl_operands])),
ufl.classes.GT: (lambda e: ast.Greater(*[_ast(o)
for o in e.ufl_operands])),
ufl.classes.GE:
(lambda e: ast.GreaterEq(*[_ast(o) for o in e.ufl_operands]))
}