def facetarea(): from ufl import Measure assert integral_type != 'cell' integrand, degree = ufl_utils.one_times(Measure(integral_type, domain=domain)) integrand = ufl_utils.replace_coordinates(integrand, coordinate_coefficient) config = kernel_config.copy() config.update(quadrature_degree=degree) expr, = fem.compile_ufl(integrand, point_sum=True, **config) return expr
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 cellvolume(restriction): from ufl import dx integrand, degree = ufl_utils.one_times(dx(domain=domain)) integrand = ufl_utils.replace_coordinates(integrand, coordinate_coefficient) interface = CellVolumeKernelInterface(kernel_config["interface"], restriction) config = {k: v for k, v in kernel_config.items() if k in ["ufl_cell", "precision", "index_cache"]} config.update(interface=interface, quadrature_degree=degree) expr, = fem.compile_ufl(integrand, point_sum=True, **config) return expr
def facetarea(): from ufl import Measure assert integral_type != 'cell' integrand, degree = ufl_utils.one_times(Measure(integral_type, domain=domain)) integrand = ufl_utils.replace_coordinates(integrand, coordinate_coefficient) quadrature_index = gem.Index(name='q') config = kernel_config.copy() config.update(quadrature_degree=degree, point_index=quadrature_index) expr, = fem.compile_ufl(integrand, **config) if quadrature_index in expr.free_indices: expr = gem.IndexSum(expr, quadrature_index) return expr
def cellvolume(restriction): from ufl import dx integrand, degree = ufl_utils.one_times(dx(domain=domain)) integrand = ufl_utils.replace_coordinates(integrand, coordinate_coefficient) interface = CellVolumeKernelInterface(kernel_config["interface"], restriction) quadrature_index = gem.Index(name='q') config = {k: v for k, v in kernel_config.items() if k in ["ufl_cell", "precision", "index_cache"]} config.update(interface=interface, quadrature_degree=degree, point_index=quadrature_index) expr, = fem.compile_ufl(integrand, **config) if quadrature_index in expr.free_indices: expr = gem.IndexSum(expr, quadrature_index) return expr
def compile_ufl_kernel(expression, to_pts, to_element, fs): import collections from ufl.algorithms.apply_function_pullbacks import apply_function_pullbacks from ufl.algorithms.apply_algebra_lowering import apply_algebra_lowering from ufl.algorithms.apply_derivatives import apply_derivatives from ufl.algorithms.apply_geometry_lowering import apply_geometry_lowering from ufl.algorithms import extract_arguments, extract_coefficients from gem import gem, impero_utils from tsfc import fem, ufl_utils from tsfc.coffee import generate as generate_coffee from tsfc.kernel_interface import (KernelBuilderBase, needs_cell_orientations, cell_orientations_coffee_arg) # Imitate the compute_form_data processing pipeline # # Unfortunately, we cannot call compute_form_data here, since # we only have an expression, not a form expression = apply_algebra_lowering(expression) expression = apply_derivatives(expression) expression = apply_function_pullbacks(expression) expression = apply_geometry_lowering(expression) expression = apply_derivatives(expression) expression = apply_geometry_lowering(expression) expression = apply_derivatives(expression) # Replace coordinates (if any) if expression.ufl_domain(): assert fs.mesh() == expression.ufl_domain() expression = ufl_utils.replace_coordinates(expression, fs.mesh().coordinates) if extract_arguments(expression): return ValueError("Cannot interpolate UFL expression with Arguments!") builder = KernelBuilderBase() args = [] coefficients = extract_coefficients(expression) for i, coefficient in enumerate(coefficients): args.append(builder.coefficient(coefficient, "w_%d" % i)) point_index = gem.Index(name='p') ir = fem.process('cell', fs.mesh().ufl_cell(), to_pts, None, point_index, (), expression, builder.coefficient_mapper, collections.defaultdict(gem.Index)) assert len(ir) == 1 # Deal with non-scalar expressions tensor_indices = () if fs.shape: tensor_indices = tuple(gem.Index() for s in fs.shape) ir = [gem.Indexed(ir[0], tensor_indices)] # Build kernel body return_var = gem.Variable('A', (len(to_pts),) + fs.shape) return_expr = gem.Indexed(return_var, (point_index,) + tensor_indices) impero_c = impero_utils.compile_gem([return_expr], ir, [point_index]) body = generate_coffee(impero_c, index_names={point_index: 'p'}) oriented = needs_cell_orientations(ir) if oriented: args.insert(0, cell_orientations_coffee_arg) # Build kernel args.insert(0, ast.Decl("double", ast.Symbol('A', rank=(len(to_pts),) + fs.shape))) kernel_code = builder.construct_kernel("expression_kernel", args, body) return op2.Kernel(kernel_code, kernel_code.name), oriented, coefficients
def compile_ufl_kernel(expression, to_pts, to_element, fs): import collections from ufl.algorithms.apply_function_pullbacks import apply_function_pullbacks from ufl.algorithms.apply_algebra_lowering import apply_algebra_lowering from ufl.algorithms.apply_derivatives import apply_derivatives from ufl.algorithms.apply_geometry_lowering import apply_geometry_lowering from ufl.algorithms import extract_arguments, extract_coefficients from gem import gem, impero_utils from tsfc import fem, ufl_utils from tsfc.coffee import generate as generate_coffee from tsfc.kernel_interface import (KernelBuilderBase, needs_cell_orientations, cell_orientations_coffee_arg) # Imitate the compute_form_data processing pipeline # # Unfortunately, we cannot call compute_form_data here, since # we only have an expression, not a form expression = apply_algebra_lowering(expression) expression = apply_derivatives(expression) expression = apply_function_pullbacks(expression) expression = apply_geometry_lowering(expression) expression = apply_derivatives(expression) expression = apply_geometry_lowering(expression) expression = apply_derivatives(expression) # Replace coordinates (if any) if expression.ufl_domain(): assert fs.mesh() == expression.ufl_domain() expression = ufl_utils.replace_coordinates(expression, fs.mesh().coordinates) if extract_arguments(expression): return ValueError("Cannot interpolate UFL expression with Arguments!") builder = KernelBuilderBase() args = [] coefficients = extract_coefficients(expression) for i, coefficient in enumerate(coefficients): args.append(builder.coefficient(coefficient, "w_%d" % i)) point_index = gem.Index(name='p') ir = fem.compile_ufl(expression, cell=fs.mesh().ufl_cell(), points=to_pts, point_index=point_index, coefficient_mapper=builder.coefficient_mapper) assert len(ir) == 1 # Deal with non-scalar expressions tensor_indices = () if fs.shape: tensor_indices = tuple(gem.Index() for s in fs.shape) ir = [gem.Indexed(ir[0], tensor_indices)] # Build kernel body return_var = gem.Variable('A', (len(to_pts),) + fs.shape) return_expr = gem.Indexed(return_var, (point_index,) + tensor_indices) impero_c = impero_utils.compile_gem([return_expr], ir, [point_index]) body = generate_coffee(impero_c, index_names={point_index: 'p'}) oriented = needs_cell_orientations(ir) if oriented: args.insert(0, cell_orientations_coffee_arg) # Build kernel args.insert(0, ast.Decl("double", ast.Symbol('A', rank=(len(to_pts),) + fs.shape))) kernel_code = builder.construct_kernel("expression_kernel", args, body) return op2.Kernel(kernel_code, kernel_code.name), oriented, coefficients
def compile_integral(integral_data, form_data, prefix, parameters, interface=firedrake_interface): """Compiles a UFL integral into an assembly kernel. :arg integral_data: UFL integral data :arg form_data: UFL form data :arg prefix: kernel name will start with this string :arg parameters: parameters object :arg interface: backend module for the kernel interface :returns: a kernel constructed by the kernel interface """ if parameters is None: parameters = default_parameters() else: _ = default_parameters() _.update(parameters) parameters = _ # Remove these here, they're handled below. if parameters.get("quadrature_degree") in ["auto", "default", None, -1, "-1"]: del parameters["quadrature_degree"] if parameters.get("quadrature_rule") in ["auto", "default", None]: del parameters["quadrature_rule"] integral_type = integral_data.integral_type interior_facet = integral_type.startswith("interior_facet") mesh = integral_data.domain cell = integral_data.domain.ufl_cell() arguments = form_data.preprocessed_form.arguments() fiat_cell = as_fiat_cell(cell) integration_dim, entity_ids = lower_integral_type(fiat_cell, integral_type) argument_indices = tuple(tuple(gem.Index(extent=e) for e in create_element(arg.ufl_element()).index_shape) for arg in arguments) flat_argument_indices = tuple(chain(*argument_indices)) quadrature_indices = [] # Dict mapping domains to index in original_form.ufl_domains() domain_numbering = form_data.original_form.domain_numbering() builder = interface.KernelBuilder(integral_type, integral_data.subdomain_id, domain_numbering[integral_data.domain]) return_variables = builder.set_arguments(arguments, argument_indices) coordinates = ufl_utils.coordinate_coefficient(mesh) builder.set_coordinates(coordinates) builder.set_coefficients(integral_data, form_data) # Map from UFL FiniteElement objects to Index instances. This is # so we reuse Index instances when evaluating the same coefficient # multiple times with the same table. Occurs, for example, if we # have multiple integrals here (and the affine coordinate # evaluation can be hoisted). index_cache = collections.defaultdict(gem.Index) kernel_cfg = dict(interface=builder, ufl_cell=cell, precision=parameters["precision"], integration_dim=integration_dim, entity_ids=entity_ids, argument_indices=argument_indices, index_cache=index_cache) kernel_cfg["facetarea"] = facetarea_generator(mesh, coordinates, kernel_cfg, integral_type) kernel_cfg["cellvolume"] = cellvolume_generator(mesh, coordinates, kernel_cfg) irs = [] for integral in integral_data.integrals: params = {} # Record per-integral parameters params.update(integral.metadata()) if params.get("quadrature_rule") == "default": del params["quadrature_rule"] # parameters override per-integral metadata params.update(parameters) integrand = ufl_utils.replace_coordinates(integral.integrand(), coordinates) integrand = ufl_utils.split_coefficients(integrand, builder.coefficient_split) # Check if the integral has a quad degree attached, otherwise use # the estimated polynomial degree attached by compute_form_data quadrature_degree = params.get("quadrature_degree", params["estimated_polynomial_degree"]) try: quad_rule = params["quadrature_rule"] except KeyError: integration_cell = fiat_cell.construct_subelement(integration_dim) quad_rule = make_quadrature(integration_cell, quadrature_degree) if not isinstance(quad_rule, AbstractQuadratureRule): raise ValueError("Expected to find a QuadratureRule object, not a %s" % type(quad_rule)) quadrature_multiindex = quad_rule.point_set.indices quadrature_indices += quadrature_multiindex config = kernel_cfg.copy() config.update(quadrature_rule=quad_rule) ir = fem.compile_ufl(integrand, interior_facet=interior_facet, **config) if parameters["unroll_indexsum"]: def predicate(index): return index.extent <= parameters["unroll_indexsum"] ir = opt.unroll_indexsum(ir, predicate=predicate) ir = [gem.index_sum(expr, quadrature_multiindex) for expr in ir] irs.append(ir) # Sum the expressions that are part of the same restriction ir = list(reduce(gem.Sum, e, gem.Zero()) for e in zip(*irs)) # Need optimised roots for COFFEE ir = impero_utils.preprocess_gem(ir) # Look for cell orientations in the IR if builder.needs_cell_orientations(ir): builder.require_cell_orientations() impero_c = impero_utils.compile_gem(return_variables, ir, tuple(quadrature_indices) + flat_argument_indices, remove_zeros=True) # Generate COFFEE index_names = [(si, name + str(n)) for index, name in zip(argument_indices, ['j', 'k']) for n, si in enumerate(index)] if len(quadrature_indices) == 1: index_names.append((quadrature_indices[0], 'ip')) else: for i, quadrature_index in enumerate(quadrature_indices): index_names.append((quadrature_index, 'ip_%d' % i)) body = generate_coffee(impero_c, index_names, parameters["precision"], ir, flat_argument_indices) kernel_name = "%s_%s_integral_%s" % (prefix, integral_type, integral_data.subdomain_id) return builder.construct_kernel(kernel_name, body)
def compile_expression_at_points(expression, points, coordinates, parameters=None): """Compiles a UFL expression to be evaluated at compile-time known reference points. Useful for interpolating UFL expressions onto function spaces with only point evaluation nodes. :arg expression: UFL expression :arg points: reference coordinates of the evaluation points :arg coordinates: the coordinate function :arg parameters: parameters object """ import coffee.base as ast 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 = ufl_utils.preprocess_expression(expression) # Replace coordinates (if any) domain = expression.ufl_domain() if domain: assert coordinates.ufl_domain() == domain expression = ufl_utils.replace_coordinates(expression, coordinates) # Collect required coefficients coefficients = extract_coefficients(expression) if coordinates not in coefficients and has_type(expression, CellVolume): coefficients = [coordinates] + coefficients # Initialise kernel builder builder = firedrake_interface.ExpressionKernelBuilder() builder.set_coefficients(coefficients) # Split mixed coefficients expression = ufl_utils.split_coefficients(expression, builder.coefficient_split) # Translate to GEM point_set = PointSet(points) config = dict(interface=builder, ufl_cell=coordinates.ufl_domain().ufl_cell(), precision=parameters["precision"], point_set=point_set) config["cellvolume"] = cellvolume_generator(coordinates.ufl_domain(), coordinates, config) ir, = fem.compile_ufl(expression, point_sum=False, **config) # Deal with non-scalar expressions value_shape = ir.shape tensor_indices = tuple(gem.Index() for s in value_shape) if value_shape: ir = gem.Indexed(ir, tensor_indices) # Build kernel body return_shape = (len(points),) + value_shape return_indices = point_set.indices + tensor_indices return_var = gem.Variable('A', return_shape) return_arg = ast.Decl(SCALAR_TYPE, ast.Symbol('A', rank=return_shape)) return_expr = gem.Indexed(return_var, return_indices) ir, = impero_utils.preprocess_gem([ir]) impero_c = impero_utils.compile_gem([return_expr], [ir], return_indices) point_index, = point_set.indices body = generate_coffee(impero_c, {point_index: 'p'}, parameters["precision"]) # Handle cell orientations if builder.needs_cell_orientations([ir]): builder.require_cell_orientations() # Build kernel tuple return builder.construct_kernel(return_arg, body)
def compile_integral(integral_data, form_data, prefix, parameters, interface=firedrake_interface): """Compiles a UFL integral into an assembly kernel. :arg integral_data: UFL integral data :arg form_data: UFL form data :arg prefix: kernel name will start with this string :arg parameters: parameters object :arg interface: backend module for the kernel interface :returns: a kernel constructed by the kernel interface """ if parameters is None: parameters = default_parameters() else: _ = default_parameters() _.update(parameters) parameters = _ # Remove these here, they're handled below. if parameters.get("quadrature_degree") in [ "auto", "default", None, -1, "-1" ]: del parameters["quadrature_degree"] if parameters.get("quadrature_rule") in ["auto", "default", None]: del parameters["quadrature_rule"] integral_type = integral_data.integral_type interior_facet = integral_type.startswith("interior_facet") mesh = integral_data.domain cell = integral_data.domain.ufl_cell() arguments = form_data.preprocessed_form.arguments() kernel_name = "%s_%s_integral_%s" % (prefix, integral_type, integral_data.subdomain_id) fiat_cell = as_fiat_cell(cell) integration_dim, entity_ids = lower_integral_type(fiat_cell, integral_type) argument_multiindices = tuple( create_element(arg.ufl_element()).get_indices() for arg in arguments) argument_indices = tuple(chain(*argument_multiindices)) quadrature_indices = [] # Dict mapping domains to index in original_form.ufl_domains() domain_numbering = form_data.original_form.domain_numbering() builder = interface.KernelBuilder(integral_type, integral_data.subdomain_id, domain_numbering[integral_data.domain]) return_variables = builder.set_arguments(arguments, argument_multiindices) coordinates = ufl_utils.coordinate_coefficient(mesh) builder.set_coordinates(coordinates) builder.set_coefficients(integral_data, form_data) # Map from UFL FiniteElement objects to multiindices. This is # so we reuse Index instances when evaluating the same coefficient # multiple times with the same table. index_cache = {} kernel_cfg = dict(interface=builder, ufl_cell=cell, precision=parameters["precision"], integration_dim=integration_dim, entity_ids=entity_ids, argument_multiindices=argument_multiindices, index_cache=index_cache) kernel_cfg["facetarea"] = facetarea_generator(mesh, coordinates, kernel_cfg, integral_type) kernel_cfg["cellvolume"] = cellvolume_generator(mesh, coordinates, kernel_cfg) mode_irs = collections.OrderedDict() for integral in integral_data.integrals: params = parameters.copy() params.update(integral.metadata()) # integral metadata overrides if params.get("quadrature_rule") == "default": del params["quadrature_rule"] mode = pick_mode(params["mode"]) mode_irs.setdefault(mode, collections.OrderedDict()) integrand = ufl_utils.replace_coordinates(integral.integrand(), coordinates) integrand = ufl.replace(integrand, form_data.function_replace_map) integrand = ufl_utils.split_coefficients(integrand, builder.coefficient_split) # Check if the integral has a quad degree attached, otherwise use # the estimated polynomial degree attached by compute_form_data quadrature_degree = params.get("quadrature_degree", params["estimated_polynomial_degree"]) try: quad_rule = params["quadrature_rule"] except KeyError: integration_cell = fiat_cell.construct_subelement(integration_dim) quad_rule = make_quadrature(integration_cell, quadrature_degree) if not isinstance(quad_rule, AbstractQuadratureRule): raise ValueError( "Expected to find a QuadratureRule object, not a %s" % type(quad_rule)) quadrature_multiindex = quad_rule.point_set.indices quadrature_indices.extend(quadrature_multiindex) config = kernel_cfg.copy() config.update(quadrature_rule=quad_rule) expressions = fem.compile_ufl(integrand, interior_facet=interior_facet, **config) reps = mode.Integrals(expressions, quadrature_multiindex, argument_multiindices, params) for var, rep in zip(return_variables, reps): mode_irs[mode].setdefault(var, []).append(rep) # Finalise mode representations into a set of assignments assignments = [] for mode, var_reps in iteritems(mode_irs): assignments.extend(mode.flatten(viewitems(var_reps))) if assignments: return_variables, expressions = zip(*assignments) else: return_variables = [] expressions = [] # Need optimised roots for COFFEE options = dict( reduce( operator.and_, [viewitems(mode.finalise_options) for mode in iterkeys(mode_irs)])) expressions = impero_utils.preprocess_gem(expressions, **options) assignments = list(zip(return_variables, expressions)) # Look for cell orientations in the IR if builder.needs_cell_orientations(expressions): builder.require_cell_orientations() # Construct ImperoC index_ordering = tuple(quadrature_indices) + argument_indices try: impero_c = impero_utils.compile_gem(assignments, index_ordering, remove_zeros=True) except impero_utils.NoopError: # No operations, construct empty kernel return builder.construct_empty_kernel(kernel_name) # Generate COFFEE index_names = [(si, name + str(n)) for index, name in zip(argument_multiindices, ['j', 'k']) for n, si in enumerate(index)] if len(quadrature_indices) == 1: index_names.append((quadrature_indices[0], 'ip')) else: for i, quadrature_index in enumerate(quadrature_indices): index_names.append((quadrature_index, 'ip_%d' % i)) # Construct kernel body = generate_coffee(impero_c, index_names, parameters["precision"], expressions, argument_indices) return builder.construct_kernel(kernel_name, body)
def compile_integral(integral_data, form_data, prefix, parameters, interface=firedrake_interface): """Compiles a UFL integral into an assembly kernel. :arg integral_data: UFL integral data :arg form_data: UFL form data :arg prefix: kernel name will start with this string :arg parameters: parameters object :arg interface: backend module for the kernel interface :returns: a kernel constructed by the kernel interface """ if parameters is None: parameters = default_parameters() else: _ = default_parameters() _.update(parameters) parameters = _ # Remove these here, they're handled below. if parameters.get("quadrature_degree") in ["auto", "default", None, -1, "-1"]: del parameters["quadrature_degree"] if parameters.get("quadrature_rule") in ["auto", "default", None]: del parameters["quadrature_rule"] integral_type = integral_data.integral_type interior_facet = integral_type.startswith("interior_facet") mesh = integral_data.domain cell = integral_data.domain.ufl_cell() arguments = form_data.preprocessed_form.arguments() fiat_cell = as_fiat_cell(cell) integration_dim, entity_ids = lower_integral_type(fiat_cell, integral_type) argument_indices = tuple(gem.Index(name=name) for arg, name in zip(arguments, ['j', 'k'])) quadrature_indices = [] # Dict mapping domains to index in original_form.ufl_domains() domain_numbering = form_data.original_form.domain_numbering() builder = interface.KernelBuilder(integral_type, integral_data.subdomain_id, domain_numbering[integral_data.domain]) return_variables = builder.set_arguments(arguments, argument_indices) coordinates = ufl_utils.coordinate_coefficient(mesh) if ufl_utils.is_element_affine(mesh.ufl_coordinate_element()): # For affine mesh geometries we prefer code generation that # composes well with optimisations. builder.set_coordinates(coordinates, mode='list_tensor') else: # Otherwise we use the approach that might be faster (?) builder.set_coordinates(coordinates) builder.set_coefficients(integral_data, form_data) # Map from UFL FiniteElement objects to Index instances. This is # so we reuse Index instances when evaluating the same coefficient # multiple times with the same table. Occurs, for example, if we # have multiple integrals here (and the affine coordinate # evaluation can be hoisted). index_cache = collections.defaultdict(gem.Index) kernel_cfg = dict(interface=builder, ufl_cell=cell, precision=parameters["precision"], integration_dim=integration_dim, entity_ids=entity_ids, argument_indices=argument_indices, index_cache=index_cache) kernel_cfg["facetarea"] = facetarea_generator(mesh, coordinates, kernel_cfg, integral_type) kernel_cfg["cellvolume"] = cellvolume_generator(mesh, coordinates, kernel_cfg) irs = [] for integral in integral_data.integrals: params = {} # Record per-integral parameters params.update(integral.metadata()) if params.get("quadrature_rule") == "default": del params["quadrature_rule"] # parameters override per-integral metadata params.update(parameters) # Check if the integral has a quad degree attached, otherwise use # the estimated polynomial degree attached by compute_form_data quadrature_degree = params.get("quadrature_degree", params["estimated_polynomial_degree"]) integration_cell = fiat_cell.construct_subelement(integration_dim) quad_rule = params.get("quadrature_rule", create_quadrature(integration_cell, quadrature_degree)) if not isinstance(quad_rule, QuadratureRule): raise ValueError("Expected to find a QuadratureRule object, not a %s" % type(quad_rule)) integrand = ufl_utils.replace_coordinates(integral.integrand(), coordinates) integrand = ufl_utils.split_coefficients(integrand, builder.coefficient_split) quadrature_index = gem.Index(name='ip') quadrature_indices.append(quadrature_index) config = kernel_cfg.copy() config.update(quadrature_rule=quad_rule, point_index=quadrature_index) ir = fem.compile_ufl(integrand, interior_facet=interior_facet, **config) if parameters["unroll_indexsum"]: ir = opt.unroll_indexsum(ir, max_extent=parameters["unroll_indexsum"]) irs.append([(gem.IndexSum(expr, quadrature_index) if quadrature_index in expr.free_indices else expr) for expr in ir]) # Sum the expressions that are part of the same restriction ir = list(reduce(gem.Sum, e, gem.Zero()) for e in zip(*irs)) # Need optimised roots for COFFEE ir = opt.remove_componenttensors(ir) # Look for cell orientations in the IR if builder.needs_cell_orientations(ir): builder.require_cell_orientations() impero_c = impero_utils.compile_gem(return_variables, ir, tuple(quadrature_indices) + argument_indices, remove_zeros=True) # Generate COFFEE index_names = [(index, index.name) for index in argument_indices] if len(quadrature_indices) == 1: index_names.append((quadrature_indices[0], 'ip')) else: for i, quadrature_index in enumerate(quadrature_indices): index_names.append((quadrature_index, 'ip_%d' % i)) body = generate_coffee(impero_c, index_names, parameters["precision"], ir, argument_indices) kernel_name = "%s_%s_integral_%s" % (prefix, integral_type, integral_data.subdomain_id) return builder.construct_kernel(kernel_name, body)
def compile_expression_at_points(expression, points, coordinates, parameters=None): """Compiles a UFL expression to be evaluated at compile-time known reference points. Useful for interpolating UFL expressions onto function spaces with only point evaluation nodes. :arg expression: UFL expression :arg points: reference coordinates of the evaluation points :arg coordinates: the coordinate function :arg parameters: parameters object """ import coffee.base as ast 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 = ufl_utils.preprocess_expression(expression) # Replace coordinates (if any) domain = expression.ufl_domain() if domain: assert coordinates.ufl_domain() == domain expression = ufl_utils.replace_coordinates(expression, coordinates) # Collect required coefficients coefficients = extract_coefficients(expression) if coordinates not in coefficients and has_type(expression, CellVolume): coefficients = [coordinates] + coefficients # Initialise kernel builder builder = firedrake_interface.ExpressionKernelBuilder() builder.set_coefficients(coefficients) # Split mixed coefficients expression = ufl_utils.split_coefficients(expression, builder.coefficient_split) # Translate to GEM point_index = gem.Index(name='p') config = dict(interface=builder, ufl_cell=coordinates.ufl_domain().ufl_cell(), precision=parameters["precision"], points=points, point_index=point_index) config["cellvolume"] = cellvolume_generator(coordinates.ufl_domain(), coordinates, config) ir, = fem.compile_ufl(expression, **config) # Deal with non-scalar expressions value_shape = ir.shape tensor_indices = tuple(gem.Index() for s in value_shape) if value_shape: ir = gem.Indexed(ir, tensor_indices) # Build kernel body return_shape = (len(points),) + value_shape return_indices = (point_index,) + tensor_indices return_var = gem.Variable('A', return_shape) return_arg = ast.Decl(SCALAR_TYPE, ast.Symbol('A', rank=return_shape)) return_expr = gem.Indexed(return_var, return_indices) impero_c = impero_utils.compile_gem([return_expr], [ir], return_indices) body = generate_coffee(impero_c, {point_index: 'p'}, parameters["precision"]) # Handle cell orientations if builder.needs_cell_orientations([ir]): builder.require_cell_orientations() # Build kernel tuple return builder.construct_kernel(return_arg, body)