def strip_coordinate_derivatives(integrals):
if isinstance(integrals, list):
if len(integrals) == 0:
return integrals, None
stripped_integrals_and_cds = []
for integral in integrals:
(si, cd) = strip_coordinate_derivatives(integral)
stripped_integrals_and_cds.append((si, cd))
return stripped_integrals_and_cds
elif isinstance(integrals, Integral):
integral = integrals
integrand = integral.integrand()
checker = CoordinateDerivativeIsOutermostChecker()
map_expr_dags(checker, [integrand])
coordinate_derivatives = []
# grab all coordinate derivatives and store them, so that we can apply
# them later again
def take_top_coordinate_derivatives(o):
o_ = o.ufl_operands
if isinstance(o, CoordinateDerivative):
coordinate_derivatives.append((o_[1], o_[2], o_[3]))
return take_top_coordinate_derivatives(o_[0])
else:
return o
newintegrand = take_top_coordinate_derivatives(integrand)
return (integral.reconstruct(integrand=newintegrand), coordinate_derivatives)
else:
error("Invalid type %s" % (integrals.__class__.__name__,))
def strip_coordinate_derivatives(integrals):
if isinstance(integrals, list):
if len(integrals) == 0:
return integrals, None
stripped_integrals_and_cds = []
for integral in integrals:
(si, cd) = strip_coordinate_derivatives(integral)
stripped_integrals_and_cds.append((si, cd))
return stripped_integrals_and_cds
elif isinstance(integrals, Integral):
integral = integrals
integrand = integral.integrand()
checker = CoordinateDerivativeIsOutermostChecker()
map_expr_dags(checker, [integrand])
coordinate_derivatives = []
# grab all coordinate derivatives and store them, so that we can apply
# them later again
def take_top_coordinate_derivatives(o):
o_ = o.ufl_operands
if isinstance(o, CoordinateDerivative):
coordinate_derivatives.append((o_[1], o_[2], o_[3]))
return take_top_coordinate_derivatives(o_[0])
else:
return o
newintegrand = take_top_coordinate_derivatives(integrand)
return (integral.reconstruct(integrand=newintegrand), coordinate_derivatives)
else:
error("Invalid type %s" % (integrals.__class__.__name__,))
def compile_to_gem(expr, translator):
"""Compile a single pointwise expression to GEM.
:arg expr: The expression to compile.
:arg translator: a :class:`Translator` instance.
:returns: A (lvalue, rvalue) pair of preprocessed GEM."""
if not isinstance(expr, Assign):
raise ValueError(
f"Don't know how to assign expression of type {type(expr)}")
spaces = tuple(c.function_space() for c in expr.coefficients)
if any(
type(s.ufl_element()) is ufl.MixedElement for s in spaces
if s is not None):
raise ValueError("Not expecting a mixed space at this point, "
"did you forget to index a function with .sub(...)?")
if len(set(s.ufl_element() for s in spaces if s is not None)) != 1:
raise ValueError("All coefficients must be defined on the same space")
lvalue = expr.lvalue
rvalue = expr.rvalue
broadcast = all(
isinstance(c, firedrake.Constant)
for c in expr.rcoefficients) and rvalue.ufl_shape == ()
if not broadcast and lvalue.ufl_shape != rvalue.ufl_shape:
try:
rvalue = reshape(rvalue, lvalue.ufl_shape)
except ValueError:
raise ValueError(
"Mismatching shapes between lvalue and rvalue in pointwise assignment"
)
rvalue, = map_expr_dags(LowerCompoundAlgebra(), [rvalue])
try:
lvalue, rvalue = map_expr_dags(translator, [lvalue, rvalue])
except (AssertionError, ValueError):
raise ValueError("Mismatching shapes in pointwise assignment. "
"For intrinsically vector-/tensor-valued spaces make "
"sure you're not using shaped Constants or literals.")
indices = gem.indices(len(lvalue.shape))
if not broadcast:
if rvalue.shape != lvalue.shape:
raise ValueError(
"Mismatching shapes in pointwise assignment. "
"For intrinsically vector-/tensor-valued spaces make "
"sure you're not using shaped Constants or literals.")
rvalue = gem.Indexed(rvalue, indices)
lvalue = gem.Indexed(lvalue, indices)
if isinstance(expr, IAdd):
rvalue = gem.Sum(lvalue, rvalue)
elif isinstance(expr, ISub):
rvalue = gem.Sum(lvalue, gem.Product(gem.Literal(-1), rvalue))
elif isinstance(expr, IMul):
rvalue = gem.Product(lvalue, rvalue)
elif isinstance(expr, IDiv):
rvalue = gem.Division(lvalue, rvalue)
return preprocess_gem([lvalue, rvalue])
def gatherDerivatives(form, arguments=None):
if arguments is None:
arguments = form.arguments()
form = apply_derivatives(apply_algebra_lowering(form))
gatherer = DerivativeGatherer(arguments)
map_expr_dags(gatherer, [i.integrand() for i in form.integrals()])
return [
sorted(list(gatherer.derivatives[arg]), key=lambda d: len(d.ufl_shape))
for arg in arguments
]
def generateCode(predefined, expressions, coefficients=None, tempVars=True):
expressions = [applyRestrictions(expand_indices(apply_derivatives(apply_algebra_lowering(expr)))) for expr in expressions]
generator = CodeGenerator(predefined, coefficients, tempVars)
results = map_expr_dags(generator, expressions)
l = list(generator.using)
l.sort()
return l + generator.code, results
def compile_ufl(expression, interior_facet=False, **kwargs):
context = Context(**kwargs)
# Abs-simplification
expression = simplify_abs(expression)
# Collect modified terminals
modified_terminals = []
map_expr_dag(CollectModifiedTerminals(modified_terminals), expression)
# Collect maximal derivatives that needs tabulation
max_derivs = collections.defaultdict(int)
for mt in map(analyse_modified_terminal, modified_terminals):
if isinstance(mt.terminal, FormArgument):
ufl_element = mt.terminal.ufl_element()
max_derivs[ufl_element] = max(mt.local_derivatives, max_derivs[ufl_element])
# Collect tabulations for all components and derivatives
tabulation_manager = TabulationManager(context.points, context.integration_dim, context.entity_ids, context.epsilon)
for ufl_element, max_deriv in max_derivs.items():
if ufl_element.family() != 'Real':
tabulation_manager.tabulate(ufl_element, max_deriv)
if interior_facet:
expressions = []
for rs in itertools.product(("+", "-"), repeat=len(context.argument_indices)):
expressions.append(map_expr_dag(PickRestriction(*rs), expression))
else:
expressions = [expression]
# Translate UFL to GEM, lowering finite element specific nodes
translator = Translator(tabulation_manager, context)
return map_expr_dags(translator, expressions)
def strip_coordinate_derivatives(form):
if isinstance(form, Form):
if len(form.integrals()) == 0:
return form, None
stripped_integrals = []
coordinate_derivatives = []
for integral in form.integrals():
(si, cd) = strip_coordinate_derivatives(integral)
stripped_integrals.append(si)
coordinate_derivatives.append(cd)
assert_that_coordinate_derivatives_are_the_same(coordinate_derivatives)
# now that we have checked that the coordinate derivative that we apply is
# the same for all integrals, we only have to return them for the first integral
return (Form(stripped_integrals), coordinate_derivatives[0])
elif isinstance(form, Integral):
integral = form
integrand = integral.integrand()
checker = CoordinateDerivativeIsOutermostChecker()
map_expr_dags(checker, [integrand])
coordinate_derivatives = []
# grab all coordinate derivatives and store them, so that we can apply
# them later again
def take_top_coordinate_derivatives(o):
o_ = o.ufl_operands
if isinstance(o, CoordinateDerivative):
coordinate_derivatives.append((o_[1], o_[2], o_[3]))
return take_top_coordinate_derivatives(o_[0])
else:
return o
newintegrand = take_top_coordinate_derivatives(integrand)
return (integral.reconstruct(integrand=newintegrand), coordinate_derivatives)
else:
error("Invalid type %s" % (form.__class__.__name__,))
def estimate_total_polynomial_degree(e, default_degree=1,
element_replace_map={}):
"""Estimate total polynomial degree of integrand.
NB! Although some compound types are supported here,
some derivatives and compounds must be preprocessed
prior to degree estimation. In generic code, this algorithm
should only be applied after preprocessing.
For coefficients defined on an element with unspecified degree (None),
the degree is set to the given default degree.
"""
de = SumDegreeEstimator(default_degree, element_replace_map)
if isinstance(e, Form):
if not e.integrals():
error("Got form with no integrals!")
degrees = map_expr_dags(de, [it.integrand() for it in e.integrals()])
elif isinstance(e, Integral):
degrees = map_expr_dags(de, [e.integrand()])
else:
degrees = map_expr_dags(de, [e])
degree = max(degrees) if degrees else default_degree
return degree
def compile_ufl(expression, interior_facet=False, point_sum=False, **kwargs):
context = PointSetContext(**kwargs)
# Abs-simplification
expression = simplify_abs(expression)
if interior_facet:
expressions = []
for rs in itertools.product(("+", "-"), repeat=len(context.argument_multiindices)):
expressions.append(map_expr_dag(PickRestriction(*rs), expression))
else:
expressions = [expression]
# Translate UFL to GEM, lowering finite element specific nodes
result = map_expr_dags(context.translator, expressions)
if point_sum:
result = [gem.index_sum(expr, context.point_indices) for expr in result]
return result
def slow_key(self):
"""A slow lookup key for this expression (relabelling UFL indices)."""
self.relabeller._reset()
rvalue, = map_expr_dags(self.relabeller, [self.rvalue])
return (type(self), hash(self.lvalue), hash(rvalue))
def _compileUFL(integrands, form, *args, tempVars=True):
if isinstance(form, Equation):
form = form.lhs - form.rhs
if not isinstance(form, Form):
raise ValueError("ufl.Form or ufl.Equation expected.")
# added for dirichlet treatment same as conservationlaw model
dirichletBCs = [arg for arg in args if isinstance(arg, DirichletBC)]
uflExpr = [form] + [bc.ufl_value for bc in dirichletBCs]
if len(form.arguments()) < 2:
raise ValueError(
"Integrands model requires form with at least two arguments.")
x = SpatialCoordinate(form.ufl_cell())
n = FacetNormal(form.ufl_cell())
cellVolume = CellVolume(form.ufl_cell())
maxCellEdgeLength = MaxCellEdgeLength(form.ufl_cell())
minCellEdgeLength = MinCellEdgeLength(form.ufl_cell())
facetArea = FacetArea(form.ufl_cell())
maxFacetEdgeLength = MaxFacetEdgeLength(form.ufl_cell())
minFacetEdgeLength = MinFacetEdgeLength(form.ufl_cell())
phi, u = form.arguments()
ubar = Coefficient(u.ufl_function_space())
derivatives = gatherDerivatives(form, [phi, u])
derivatives_phi = derivatives[0]
derivatives_u = derivatives[1]
derivatives_ubar = map_expr_dags(Replacer({u: ubar}), derivatives_u)
try:
integrands.field = u.ufl_function_space().field
except AttributeError:
pass
integrals = splitForm(form, [phi])
dform = apply_derivatives(derivative(action(form, ubar), ubar, u))
linearizedIntegrals = splitForm(dform, [phi, u])
if not set(
integrals.keys()) <= {'cell', 'exterior_facet', 'interior_facet'}:
raise Exception('unknown integral encountered in ' +
str(set(integrals.keys())) + '.')
if 'cell' in integrals.keys():
arg = Variable(integrands.domainValueTuple, 'u')
predefined = {
derivatives_u[i]: arg[i]
for i in range(len(derivatives_u))
}
predefined[x] = integrands.spatialCoordinate('x')
predefined[cellVolume] = integrands.cellVolume()
predefined[maxCellEdgeLength] = maxEdgeLength(
integrands.cellGeometry())
predefined[minCellEdgeLength] = minEdgeLength(
integrands.cellGeometry())
integrands.predefineCoefficients(predefined, False)
integrands.interior = generateUnaryCode(predefined,
derivatives_phi,
integrals['cell'],
tempVars=tempVars)
predefined = {
derivatives_ubar[i]: arg[i]
for i in range(len(derivatives_u))
}
predefined[x] = integrands.spatialCoordinate('x')
predefined[cellVolume] = integrands.cellVolume()
predefined[maxCellEdgeLength] = maxEdgeLength(
integrands.cellGeometry())
predefined[minCellEdgeLength] = minEdgeLength(
integrands.cellGeometry())
integrands.predefineCoefficients(predefined, False)
integrands.linearizedInterior = generateUnaryLinearizedCode(
predefined,
derivatives_phi,
derivatives_u,
linearizedIntegrals.get('cell'),
tempVars=tempVars)
if 'exterior_facet' in integrals.keys():
arg = Variable(integrands.domainValueTuple, 'u')
predefined = {
derivatives_u[i]: arg[i]
for i in range(len(derivatives_u))
}
predefined[x] = integrands.spatialCoordinate('x')
predefined[n] = integrands.facetNormal('x')
predefined[cellVolume] = integrands.cellVolume()
predefined[maxCellEdgeLength] = maxEdgeLength(
integrands.cellGeometry())
predefined[minCellEdgeLength] = minEdgeLength(
integrands.cellGeometry())
predefined[facetArea] = integrands.facetArea()
predefined[maxFacetEdgeLength] = maxEdgeLength(
integrands.facetGeometry())
predefined[minFacetEdgeLength] = minEdgeLength(
integrands.facetGeometry())
integrands.predefineCoefficients(predefined, False)
integrands.boundary = generateUnaryCode(predefined,
derivatives_phi,
integrals['exterior_facet'],
tempVars=tempVars)
predefined = {
derivatives_ubar[i]: arg[i]
for i in range(len(derivatives_u))
}
predefined[x] = integrands.spatialCoordinate('x')
predefined[n] = integrands.facetNormal('x')
predefined[cellVolume] = integrands.cellVolume()
predefined[maxCellEdgeLength] = maxEdgeLength(
integrands.cellGeometry())
predefined[minCellEdgeLength] = minEdgeLength(
integrands.cellGeometry())
predefined[facetArea] = integrands.facetArea()
predefined[maxFacetEdgeLength] = maxEdgeLength(
integrands.facetGeometry())
predefined[minFacetEdgeLength] = minEdgeLength(
integrands.facetGeometry())
integrands.predefineCoefficients(predefined, False)
integrands.linearizedBoundary = generateUnaryLinearizedCode(
predefined,
derivatives_phi,
derivatives_u,
linearizedIntegrals.get('exterior_facet'),
tempVars=tempVars)
if 'interior_facet' in integrals.keys():
argIn = Variable(integrands.domainValueTuple, 'uIn')
argOut = Variable(integrands.domainValueTuple, 'uOut')
predefined = {
derivatives_u[i](s): arg[i]
for i in range(len(derivatives_u))
for s, arg in (('+', argIn), ('-', argOut))
}
predefined[x] = integrands.spatialCoordinate('xIn')
predefined[n('+')] = integrands.facetNormal('xIn')
predefined[cellVolume('+')] = integrands.cellVolume('Side::in')
predefined[cellVolume('-')] = integrands.cellVolume('Side::out')
predefined[maxCellEdgeLength('+')] = maxEdgeLength(
integrands.cellGeometry('Side::in'))
predefined[maxCellEdgeLength('-')] = maxEdgeLength(
integrands.cellGeometry('Side::out'))
predefined[minCellEdgeLength('+')] = minEdgeLength(
integrands.cellGeometry('Side::in'))
predefined[minCellEdgeLength('-')] = minEdgeLength(
integrands.cellGeometry('Side::out'))
predefined[facetArea] = integrands.facetArea()
predefined[maxFacetEdgeLength] = maxEdgeLength(
integrands.facetGeometry())
predefined[minFacetEdgeLength] = minEdgeLength(
integrands.facetGeometry())
integrands.predefineCoefficients(predefined, True)
integrands.skeleton = generateBinaryCode(predefined,
derivatives_phi,
integrals['interior_facet'],
tempVars=tempVars)
predefined = {
derivatives_ubar[i](s): arg[i]
for i in range(len(derivatives_u))
for s, arg in (('+', argIn), ('-', argOut))
}
predefined[x] = integrands.spatialCoordinate('xIn')
predefined[n('+')] = integrands.facetNormal('xIn')
predefined[cellVolume('+')] = integrands.cellVolume('Side::in')
predefined[cellVolume('-')] = integrands.cellVolume('Side::out')
predefined[maxCellEdgeLength('+')] = maxEdgeLength(
integrands.cellGeometry('Side::in'))
predefined[maxCellEdgeLength('-')] = maxEdgeLength(
integrands.cellGeometry('Side::out'))
predefined[minCellEdgeLength('+')] = minEdgeLength(
integrands.cellGeometry('Side::in'))
predefined[minCellEdgeLength('-')] = minEdgeLength(
integrands.cellGeometry('Side::out'))
predefined[facetArea] = integrands.facetArea()
predefined[maxFacetEdgeLength] = maxEdgeLength(
integrands.facetGeometry())
predefined[minFacetEdgeLength] = minEdgeLength(
integrands.facetGeometry())
integrands.predefineCoefficients(predefined, True)
integrands.linearizedSkeleton = generateBinaryLinearizedCode(
predefined,
derivatives_phi,
derivatives_u,
linearizedIntegrals.get('interior_facet'),
tempVars=tempVars)
if dirichletBCs:
integrands.hasDirichletBoundary = True
predefined = {}
# predefined[x] = UnformattedExpression('auto', 'entity().geometry().global( Dune::Fem::coordinate( ' + integrands.arg_x.name + ' ) )')
predefined[x] = UnformattedExpression(
'auto', 'intersection.geometry().center( )')
integrands.predefineCoefficients(predefined, False)
maxId = 0
codeDomains = []
bySubDomain = dict()
neuman = []
wholeDomain = None
for bc in dirichletBCs:
if bc.subDomain in bySubDomain:
raise Exception(
'Multiply defined Dirichlet boundary for subdomain ' +
str(bc.subDomain))
if not isinstance(bc.functionSpace,
(FunctionSpace, FiniteElementBase)):
raise Exception(
'Function space must either be a ufl.FunctionSpace or a ufl.FiniteElement'
)
if isinstance(bc.functionSpace, FunctionSpace) and (
bc.functionSpace != u.ufl_function_space()):
raise Exception(
'Space of trial function and dirichlet boundary function must be the same - note that boundary conditions on subspaces are not available, yet'
)
if isinstance(bc.functionSpace, FiniteElementBase) and (
bc.functionSpace != u.ufl_element()):
raise Exception(
'Cannot handle boundary conditions on subspaces, yet')
if isinstance(bc.value, list):
neuman = [i for i, x in enumerate(bc.value) if x == None]
else:
neuman = []
value = ExprTensor(u.ufl_shape)
for key in value.keys():
value[key] = Indexed(
bc.ufl_value,
MultiIndex(tuple(FixedIndex(k) for k in key)))
if bc.subDomain is None:
wholeDomain = value, neuman
elif isinstance(bc.subDomain, int):
bySubDomain[bc.subDomain] = value, neuman
maxId = max(maxId, bc.subDomain)
else:
domain = ExprTensor(())
for key in domain.keys():
domain[key] = Indexed(
bc.subDomain,
MultiIndex(tuple(FixedIndex(k) for k in key)))
codeDomains.append((value, neuman, domain))
defaultCode = []
if len(codeDomains) > 0:
defaultCode.append(Declaration(Variable('int', 'domainId')))
# defaultCode.append(Declaration(Variable('auto', 'x'),
# initializer=UnformattedExpression('auto','intersection.geometry().center()')))
for i, v in enumerate(codeDomains):
block = Block()
block.append(
generateDirichletDomainCode(predefined,
v[2],
tempVars=tempVars))
block.append('if (domainId)')
ifBlock = UnformattedBlock()
ifBlock.append(
'std::fill( dirichletComponent.begin(), dirichletComponent.end(), '
+ str(maxId + i + 2) + ' );')
if len(v[1]) > 0:
[
ifBlock.append('dirichletComponent[' + str(c) + '] = 0;')
for c in v[1]
]
ifBlock.append('return true;')
block.append(ifBlock)
defaultCode.append(block)
if wholeDomain is not None:
block = UnformattedBlock()
block.append(
'std::fill( dirichletComponent.begin(), dirichletComponent.end(), '
+ str(maxId + 1) + ' );')
if len(wholeDomain[1]) > 0:
[
block.append('dirichletComponent[' + str(c) + '] = 0;')
for c in wholeDomain[1]
]
block.append('return true;')
defaultCode.append(block)
defaultCode.append(return_(False))
bndId = Variable('const int', 'bndId')
getBndId = UnformattedExpression(
'int', 'BoundaryIdProviderType::boundaryId( ' +
integrands.arg_i.name + ' )')
# getBndId = UnformattedExpression('int', 'boundaryIdGetter_.boundaryId( ' + integrands.arg_i.name + ' )')
switch = SwitchStatement(bndId, default=defaultCode)
for i, v in bySubDomain.items():
code = []
if len(v[1]) > 0:
[
code.append('dirichletComponent[' + str(c) + '] = 0;')
for c in v[1]
]
code.append(return_(True))
switch.append(i, code)
integrands.isDirichletIntersection = [
Declaration(bndId, initializer=getBndId),
UnformattedBlock(
'std::fill( dirichletComponent.begin(), dirichletComponent.end(), '
+ bndId.name + ' );'), switch
]
predefined[x] = UnformattedExpression(
'auto', 'entity().geometry().global( Dune::Fem::coordinate( ' +
integrands.arg_x.name + ' ) )')
if wholeDomain is None:
defaultCode = assign(integrands.arg_r, construct("RRangeType", 0))
else:
defaultCode = generateDirichletCode(predefined,
wholeDomain[0],
tempVars=tempVars)
switch = SwitchStatement(integrands.arg_bndId, default=defaultCode)
for i, v in bySubDomain.items():
switch.append(
i, generateDirichletCode(predefined, v[0], tempVars=tempVars))
for i, v in enumerate(codeDomains):
switch.append(
i + maxId + 2,
generateDirichletCode(predefined, v[0], tempVars=tempVars))
integrands.dirichlet = [switch]
return integrands
def compile_element(expression, coordinates, parameters=None):
"""Generates C code for point evaluations.
:arg expression: UFL expression
:arg coordinates: coordinate field
:arg parameters: form compiler parameters
:returns: C code as string
"""
if parameters is None:
parameters = default_parameters()
else:
_ = default_parameters()
_.update(parameters)
parameters = _
# No arguments, please!
if extract_arguments(expression):
return ValueError("Cannot interpolate UFL expression with Arguments!")
# Apply UFL preprocessing
expression = tsfc.ufl_utils.preprocess_expression(expression)
# Collect required coefficients
coefficient, = extract_coefficients(expression)
# Point evaluation of mixed coefficients not supported here
if type(coefficient.ufl_element()) == MixedElement:
raise NotImplementedError("Cannot point evaluate mixed elements yet!")
# Replace coordinates (if any)
domain = expression.ufl_domain()
assert coordinates.ufl_domain() == domain
expression = tsfc.ufl_utils.replace_coordinates(expression, coordinates)
# Initialise kernel builder
builder = firedrake_interface.KernelBuilderBase()
x_arg = builder._coefficient(coordinates, "x")
f_arg = builder._coefficient(coefficient, "f")
# TODO: restore this for expression evaluation!
# expression = ufl_utils.split_coefficients(expression, builder.coefficient_split)
# 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=coordinates.ufl_domain().ufl_cell(),
precision=parameters["precision"],
point_indices=(),
point_expr=point)
# TODO: restore this for expression evaluation!
# config["cellvolume"] = cellvolume_generator(coordinates.ufl_domain(), coordinates, config)
context = tsfc.fem.GemPointContext(**config)
# Abs-simplification
expression = tsfc.ufl_utils.simplify_abs(expression)
# Translate UFL -> GEM
translator = tsfc.fem.Translator(context)
result, = map_expr_dags(translator, [expression])
tensor_indices = ()
if expression.ufl_shape:
tensor_indices = tuple(gem.Index() for s in 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, )))
# 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(
"evaluate_kernel", [result_arg, point_arg, x_arg, f_arg], body)
# Fill the code template
extruded = isinstance(cell, TensorProductCell)
#code = {
#"geometric_dimension": cell.geometric_dimension(),
#"extruded_arg": ", %s nlayers" % as_cstr(IntType) if extruded else "",
#"nlayers": ", f->n_layers" if extruded else "",
#"IntType": as_cstr(IntType),
#}
#evaluate_template_c = """static inline void wrap_evaluate(double *result, double *X, double *coords, %(IntType)s *coords_map, double *f, %(IntType)s *f_map%(extruded_arg)s, %(IntType)s cell);
#int evaluate(struct Function *f, double *x, double *result)
#{
#struct ReferenceCoords reference_coords;
#%(IntType)s cell = locate_cell(f, x, %(geometric_dimension)d, &to_reference_coords, &reference_coords);
#if (cell == -1) {
#return -1;
#}
#if (!result) {
#return 0;
#}
#wrap_evaluate(result, reference_coords.X, f->coords, f->coords_map, f->f, f->f_map%(nlayers)s, cell);
#return 0;
#}
#"""
# return (evaluate_template_c % code) + kernel_code.gencode()
return kernel_code.gencode()
def compile_element(expression, coordinates, parameters=None):
"""Generates C code for point evaluations.
:arg expression: UFL expression
:arg coordinates: coordinate field
:arg parameters: form compiler parameters
:returns: C code as string
"""
if parameters is None:
parameters = default_parameters()
else:
_ = default_parameters()
_.update(parameters)
parameters = _
# No arguments, please!
if extract_arguments(expression):
return ValueError("Cannot interpolate UFL expression with Arguments!")
# Apply UFL preprocessing
expression = tsfc.ufl_utils.preprocess_expression(
expression, complex_mode=utils.complex_mode)
# Collect required coefficients
coefficient, = extract_coefficients(expression)
# Point evaluation of mixed coefficients not supported here
if type(coefficient.ufl_element()) == MixedElement:
raise NotImplementedError("Cannot point evaluate mixed elements yet!")
# Replace coordinates (if any)
domain = expression.ufl_domain()
assert coordinates.ufl_domain() == domain
# Initialise kernel builder
builder = firedrake_interface.KernelBuilderBase(utils.ScalarType_c)
builder.domain_coordinate[domain] = coordinates
x_arg = builder._coefficient(coordinates, "x")
f_arg = builder._coefficient(coefficient, "f")
# TODO: restore this for expression evaluation!
# expression = ufl_utils.split_coefficients(expression, builder.coefficient_split)
# Translate to GEM
cell = domain.ufl_cell()
dim = cell.topological_dimension()
point = gem.Variable('X', (dim, ))
point_arg = ast.Decl(utils.ScalarType_c, ast.Symbol('X', rank=(dim, )))
config = dict(interface=builder,
ufl_cell=coordinates.ufl_domain().ufl_cell(),
precision=parameters["precision"],
point_indices=(),
point_expr=point,
complex_mode=utils.complex_mode)
# TODO: restore this for expression evaluation!
# config["cellvolume"] = cellvolume_generator(coordinates.ufl_domain(), coordinates, config)
context = tsfc.fem.GemPointContext(**config)
# Abs-simplification
expression = tsfc.ufl_utils.simplify_abs(expression, utils.complex_mode)
# Translate UFL -> GEM
translator = tsfc.fem.Translator(context)
result, = map_expr_dags(translator, [expression])
tensor_indices = ()
if expression.ufl_shape:
tensor_indices = tuple(gem.Index() for s in expression.ufl_shape)
return_variable = gem.Indexed(gem.Variable('R', expression.ufl_shape),
tensor_indices)
result_arg = ast.Decl(utils.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(utils.ScalarType_c, ast.Symbol('R', rank=(1, )))
# 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"],
utils.ScalarType_c)
# Build kernel tuple
kernel_code = builder.construct_kernel(
"evaluate_kernel", [result_arg, point_arg, x_arg, f_arg], body)
# Fill the code template
extruded = isinstance(cell, TensorProductCell)
code = {
"geometric_dimension": cell.geometric_dimension(),
"layers_arg": ", int const *__restrict__ layers" if extruded else "",
"layers": ", layers" if extruded else "",
"IntType": as_cstr(IntType),
"scalar_type": utils.ScalarType_c,
}
# if maps are the same, only need to pass one of them
if coordinates.cell_node_map() == coefficient.cell_node_map():
code[
"wrapper_map_args"] = "%(IntType)s const *__restrict__ coords_map" % code
code["map_args"] = "f->coords_map"
else:
code[
"wrapper_map_args"] = "%(IntType)s const *__restrict__ coords_map, %(IntType)s const *__restrict__ f_map" % code
code["map_args"] = "f->coords_map, f->f_map"
evaluate_template_c = """
static inline void wrap_evaluate(%(scalar_type)s* const result, %(scalar_type)s* const X, int const start, int const end%(layers_arg)s,
%(scalar_type)s const *__restrict__ coords, %(scalar_type)s const *__restrict__ f, %(wrapper_map_args)s);
int evaluate(struct Function *f, %(scalar_type)s *x, %(scalar_type)s *result)
{
struct ReferenceCoords reference_coords;
%(IntType)s cell = locate_cell(f, x, %(geometric_dimension)d, &to_reference_coords, &to_reference_coords_xtr, &reference_coords);
if (cell == -1) {
return -1;
}
if (!result) {
return 0;
}
int layers[2] = {0, 0};
if (f->extruded != 0) {
int nlayers = f->n_layers;
layers[1] = cell %% nlayers + 2;
cell = cell / nlayers;
}
wrap_evaluate(result, reference_coords.X, cell, cell+1%(layers)s, f->coords, f->f, %(map_args)s);
return 0;
}
"""
return (evaluate_template_c % code) + kernel_code.gencode()
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 expr_list(self, o, *operands):
# Inline list tensor indexing.
# This fixes a problem where we extract a subblock from
# derivative(foo, ...) and end up with the "Argument" looking like
# [v_0, v_2, v_3][1, 2]
return self.expr(o, *map_expr_dags(self.index_inliner, operands))
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 expr_list(self, o, *operands):
# Inline list tensor indexing.
# This fixes a problem where we extract a subblock from
# derivative(foo, ...) and end up with the "Argument" looking like
# [v_0, v_2, v_3][1, 2]
return self.expr(o, *map_expr_dags(self.index_inliner, operands))
def compile_element(expression, coordinates, parameters=None):
"""Generates C code for point evaluations.
:arg expression: UFL expression
:arg coordinates: coordinate field
:arg parameters: form compiler parameters
:returns: C code as string
"""
if parameters is None:
parameters = default_parameters()
else:
_ = default_parameters()
_.update(parameters)
parameters = _
# No arguments, please!
if extract_arguments(expression):
return ValueError("Cannot interpolate UFL expression with Arguments!")
# Apply UFL preprocessing
expression = tsfc.ufl_utils.preprocess_expression(expression)
# Collect required coefficients
coefficient, = extract_coefficients(expression)
# Point evaluation of mixed coefficients not supported here
if type(coefficient.ufl_element()) == MixedElement:
raise NotImplementedError("Cannot point evaluate mixed elements yet!")
# Replace coordinates (if any)
domain = expression.ufl_domain()
assert coordinates.ufl_domain() == domain
# Initialise kernel builder
builder = firedrake_interface.KernelBuilderBase()
builder.domain_coordinate[domain] = coordinates
x_arg = builder._coefficient(coordinates, "x")
f_arg = builder._coefficient(coefficient, "f")
# TODO: restore this for expression evaluation!
# expression = ufl_utils.split_coefficients(expression, builder.coefficient_split)
# 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=coordinates.ufl_domain().ufl_cell(),
precision=parameters["precision"],
point_indices=(),
point_expr=point)
# TODO: restore this for expression evaluation!
# config["cellvolume"] = cellvolume_generator(coordinates.ufl_domain(), coordinates, config)
context = tsfc.fem.GemPointContext(**config)
# Abs-simplification
expression = tsfc.ufl_utils.simplify_abs(expression)
# Translate UFL -> GEM
translator = tsfc.fem.Translator(context)
result, = map_expr_dags(translator, [expression])
tensor_indices = ()
if expression.ufl_shape:
tensor_indices = tuple(gem.Index() for s in 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,)))
# 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("evaluate_kernel", [result_arg, point_arg, x_arg, f_arg], body)
# Fill the code template
extruded = isinstance(cell, TensorProductCell)
code = {
"geometric_dimension": cell.geometric_dimension(),
"layers_arg": ", int const *__restrict__ layers" if extruded else "",
"layers": ", layers" if extruded else "",
"IntType": as_cstr(IntType),
}
# if maps are the same, only need to pass one of them
if coordinates.cell_node_map() == coefficient.cell_node_map():
code["wrapper_map_args"] = "%(IntType)s const *__restrict__ coords_map" % code
code["map_args"] = "f->coords_map"
else:
code["wrapper_map_args"] = "%(IntType)s const *__restrict__ coords_map, %(IntType)s const *__restrict__ f_map" % code
code["map_args"] = "f->coords_map, f->f_map"
evaluate_template_c = """
static inline void wrap_evaluate(double* const result, double* const X, int const start, int const end%(layers_arg)s,
double const *__restrict__ coords, double const *__restrict__ f, %(wrapper_map_args)s);
int evaluate(struct Function *f, double *x, double *result)
{
struct ReferenceCoords reference_coords;
%(IntType)s cell = locate_cell(f, x, %(geometric_dimension)d, &to_reference_coords, &to_reference_coords_xtr, &reference_coords);
if (cell == -1) {
return -1;
}
if (!result) {
return 0;
}
int layers[2] = {0, 0};
if (f->extruded != 0) {
int nlayers = f->n_layers;
layers[1] = cell %% nlayers + 2;
cell = cell / nlayers;
}
wrap_evaluate(result, reference_coords.X, cell, cell+1%(layers)s, f->coords, f->f, %(map_args)s);
return 0;
}
"""
return (evaluate_template_c % code) + kernel_code.gencode()
