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 to_reference_coordinates(ufl_coordinate_element): # Set up UFL form cell = ufl_coordinate_element.cell() domain = ufl.Mesh(ufl_coordinate_element) K = ufl.JacobianInverse(domain) x = ufl.SpatialCoordinate(domain) x0_element = ufl.VectorElement("Real", cell, 0) x0 = ufl.Coefficient(ufl.FunctionSpace(domain, x0_element)) expr = ufl.dot(K, x - x0) # Translation to GEM C = ufl_utils.coordinate_coefficient(domain) expr = ufl_utils.preprocess_expression(expr) expr = ufl_utils.replace_coordinates(expr, C) expr = ufl_utils.simplify_abs(expr) builder = firedrake_interface.KernelBuilderBase() builder._coefficient(C, "C") builder._coefficient(x0, "x0") dim = cell.topological_dimension() point = gem.Variable('X', (dim, )) context = tsfc.fem.GemPointContext( interface=builder, ufl_cell=cell, precision=parameters["precision"], point_indices=(), point_expr=point, ) translator = tsfc.fem.Translator(context) ir = map_expr_dag(translator, expr) # Unroll result ir = [gem.Indexed(ir, alpha) for alpha in numpy.ndindex(ir.shape)] # Unroll IndexSums max_extent = parameters["unroll_indexsum"] if max_extent: def predicate(index): return index.extent <= max_extent ir = gem.optimise.unroll_indexsum(ir, predicate=predicate) # Translate to COFFEE ir = impero_utils.preprocess_gem(ir) return_variable = gem.Variable('dX', (dim, )) assignments = [(gem.Indexed(return_variable, (i, )), e) for i, e in enumerate(ir)] impero_c = impero_utils.compile_gem(assignments, ()) body = tsfc.coffee.generate(impero_c, {}, parameters["precision"]) body.open_scope = False return body
def to_reference_coordinates(ufl_coordinate_element, parameters): # Set up UFL form cell = ufl_coordinate_element.cell() domain = ufl.Mesh(ufl_coordinate_element) K = ufl.JacobianInverse(domain) x = ufl.SpatialCoordinate(domain) x0_element = ufl.VectorElement("Real", cell, 0) x0 = ufl.Coefficient(ufl.FunctionSpace(domain, x0_element)) expr = ufl.dot(K, x - x0) # Translation to GEM C = ufl.Coefficient(ufl.FunctionSpace(domain, ufl_coordinate_element)) expr = ufl_utils.preprocess_expression(expr) expr = ufl_utils.simplify_abs(expr) builder = firedrake_interface.KernelBuilderBase() builder.domain_coordinate[domain] = C builder._coefficient(C, "C") builder._coefficient(x0, "x0") dim = cell.topological_dimension() point = gem.Variable('X', (dim,)) context = tsfc.fem.GemPointContext( interface=builder, ufl_cell=cell, precision=parameters["precision"], point_indices=(), point_expr=point, ) translator = tsfc.fem.Translator(context) ir = map_expr_dag(translator, expr) # Unroll result ir = [gem.Indexed(ir, alpha) for alpha in numpy.ndindex(ir.shape)] # Unroll IndexSums max_extent = parameters["unroll_indexsum"] if max_extent: def predicate(index): return index.extent <= max_extent ir = gem.optimise.unroll_indexsum(ir, predicate=predicate) # Translate to COFFEE ir = impero_utils.preprocess_gem(ir) return_variable = gem.Variable('dX', (dim,)) assignments = [(gem.Indexed(return_variable, (i,)), e) for i, e in enumerate(ir)] impero_c = impero_utils.compile_gem(assignments, ()) body = tsfc.coffee.generate(impero_c, {}, parameters["precision"]) body.open_scope = False return body
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