def _find_linear_terms(rhs_exprs, u): """Help function that takes a list of rhs expressions and return a list of bools determining what component, rhs_exprs[i], is linear wrt u[i]. """ uu = [Constant(1.0) for _ in rhs_exprs] if len(rhs_exprs) > 1: repl = {u: ufl.as_vector(uu)} else: repl = {u: uu[0]} linear_terms = [] for i, ui in enumerate(uu): comp_i_s = expand_indices(ufl.replace(rhs_exprs[i], repl)) linear_terms.append(ui in extract_coefficients(comp_i_s) and ui not in extract_coefficients( expand_derivatives(ufl.diff(comp_i_s, ui)))) return linear_terms
def derivative(form, u, du=None, coefficient_derivatives=None): """Compute the derivative of a form. Given a form, this computes its linearization with respect to the provided :class:`.Function`. The resulting form has one additional :class:`Argument` in the same finite element space as the Function. :arg form: a :class:`~ufl.classes.Form` to compute the derivative of. :arg u: a :class:`.Function` to compute the derivative with respect to. :arg du: an optional :class:`Argument` to use as the replacement in the new form (constructed automatically if not provided). :arg coefficient_derivatives: an optional :class:`dict` to provide the derivative of a coefficient function. :raises ValueError: If any of the coefficients in ``form`` were obtained from ``u.split()``. UFL doesn't notice that these are related to ``u`` and so therefore the derivative is wrong (instead one should have written ``split(u)``). See also :func:`ufl.derivative`. """ if set(extract_coefficients(form)) & set(u.split()): raise ValueError( "Taking derivative of form wrt u, but form contains coefficients from u.split()." "\nYou probably meant to write split(u) when defining your form.") if du is None: if isinstance(u, function.Function): V = u.function_space() args = form.arguments() number = max(a.number() for a in args) if args else -1 du = Argument(V, number + 1) else: raise RuntimeError("Can't compute derivative for form") return ufl.derivative(form, u, du, coefficient_derivatives)
def test_extract_coefficients_vs_fixture(coefficients, forms): assert coefficients == tuple(extract_coefficients(forms[2]))
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 derivative(form, u, du=None, coefficient_derivatives=None): """Compute the derivative of a form. Given a form, this computes its linearization with respect to the provided :class:`.Function`. The resulting form has one additional :class:`Argument` in the same finite element space as the Function. :arg form: a :class:`~ufl.classes.Form` to compute the derivative of. :arg u: a :class:`.Function` to compute the derivative with respect to. :arg du: an optional :class:`Argument` to use as the replacement in the new form (constructed automatically if not provided). :arg coefficient_derivatives: an optional :class:`dict` to provide the derivative of a coefficient function. :raises ValueError: If any of the coefficients in ``form`` were obtained from ``u.split()``. UFL doesn't notice that these are related to ``u`` and so therefore the derivative is wrong (instead one should have written ``split(u)``). See also :func:`ufl.derivative`. """ # TODO: What about Constant? u_is_x = isinstance(u, ufl.SpatialCoordinate) if not u_is_x and len(u.split()) > 1 and set( extract_coefficients(form)) & set(u.split()): raise ValueError( "Taking derivative of form wrt u, but form contains coefficients from u.split()." "\nYou probably meant to write split(u) when defining your form.") mesh = form.ufl_domain() is_dX = u_is_x or u is mesh.coordinates args = form.arguments() def argument(V): if du is None: n = max(a.number() for a in args) if args else -1 return Argument(V, n + 1) else: return du if is_dX: coords = mesh.coordinates u = ufl.SpatialCoordinate(mesh) V = coords.function_space() du = argument(V) cds = {coords: du} if coefficient_derivatives is not None: cds.update(coefficient_derivatives) coefficient_derivatives = cds elif isinstance(u, firedrake.Function): V = u.function_space() du = argument(V) elif isinstance(u, firedrake.Constant): if u.ufl_shape != (): raise ValueError( "Real function space of vector elements not supported") V = firedrake.FunctionSpace(mesh, "Real", 0) du = argument(V) else: raise RuntimeError("Can't compute derivative for form") return ufl.derivative(form, u, du, coefficient_derivatives)
def derivative(form, u, du=None, coefficient_derivatives=None): """Compute the derivative of a form. Given a form, this computes its linearization with respect to the provided :class:`.Function`. The resulting form has one additional :class:`Argument` in the same finite element space as the Function. :arg form: a :class:`~ufl.classes.Form` to compute the derivative of. :arg u: a :class:`.Function` to compute the derivative with respect to. :arg du: an optional :class:`Argument` to use as the replacement in the new form (constructed automatically if not provided). :arg coefficient_derivatives: an optional :class:`dict` to provide the derivative of a coefficient function. :raises ValueError: If any of the coefficients in ``form`` were obtained from ``u.split()``. UFL doesn't notice that these are related to ``u`` and so therefore the derivative is wrong (instead one should have written ``split(u)``). See also :func:`ufl.derivative`. """ if isinstance(form, firedrake.slate.TensorBase): raise TypeError( f"Cannot take the derivative of a {type(form).__name__}" ) # TODO: What about Constant? u_is_x = isinstance(u, ufl.SpatialCoordinate) uc, = (u,) if u_is_x else extract_coefficients(u) if not u_is_x and len(uc.split()) > 1 and set(extract_coefficients(form)) & set(uc.split()): raise ValueError("Taking derivative of form wrt u, but form contains coefficients from u.split()." "\nYou probably meant to write split(u) when defining your form.") mesh = form.ufl_domain() if not mesh: raise ValueError("Expression to be differentiated has no ufl domain." "\nDo you need to add a domain to your Constant?") is_dX = u_is_x or u is mesh.coordinates try: args = form.arguments() except AttributeError: args = extract_arguments(form) def argument(V): if du is None: n = max(a.number() for a in args) if args else -1 return Argument(V, n + 1) else: return du if is_dX: coords = mesh.coordinates u = ufl.SpatialCoordinate(mesh) V = coords.function_space() du = argument(V) cds = {coords: du} if coefficient_derivatives is not None: cds.update(coefficient_derivatives) coefficient_derivatives = cds elif isinstance(uc, firedrake.Function): V = uc.function_space() du = argument(V) elif isinstance(uc, firedrake.Constant): if uc.ufl_shape != (): raise ValueError("Real function space of vector elements not supported") V = firedrake.FunctionSpace(mesh, "Real", 0) du = argument(V) else: raise RuntimeError("Can't compute derivative for form") if u.ufl_shape != du.ufl_shape: raise ValueError("Shapes of u and du do not match.\n" "If you passed an indexed part of split(u) into " "derivative, you need to provide an appropriate du as well.") return ufl.derivative(form, u, du, coefficient_derivatives)
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) # Initialise kernel builder builder = firedrake_interface.ExpressionKernelBuilder() # Replace coordinates (if any) domain = expression.ufl_domain() if domain: assert coordinates.ufl_domain() == domain builder.domain_coordinate[domain] = coordinates # Collect required coefficients coefficients = extract_coefficients(expression) if has_type(expression, GeometricQuantity): coefficients = [coordinates] + coefficients 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) 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 derivative(form, u, du=None, coefficient_derivatives=None): """Compute the derivative of a form. Given a form, this computes its linearization with respect to the provided :class:`.Function`. The resulting form has one additional :class:`Argument` in the same finite element space as the Function. :arg form: a :class:`~ufl.classes.Form` to compute the derivative of. :arg u: a :class:`.Function` to compute the derivative with respect to. :arg du: an optional :class:`Argument` to use as the replacement in the new form (constructed automatically if not provided). :arg coefficient_derivatives: an optional :class:`dict` to provide the derivative of a coefficient function. :raises ValueError: If any of the coefficients in ``form`` were obtained from ``u.split()``. UFL doesn't notice that these are related to ``u`` and so therefore the derivative is wrong (instead one should have written ``split(u)``). See also :func:`ufl.derivative`. """ # TODO: What about Constant? u_is_x = isinstance(u, ufl.SpatialCoordinate) if not u_is_x and len(u.split()) > 1 and set(extract_coefficients(form)) & set(u.split()): raise ValueError("Taking derivative of form wrt u, but form contains coefficients from u.split()." "\nYou probably meant to write split(u) when defining your form.") mesh = form.ufl_domain() is_dX = u_is_x or u is mesh.coordinates args = form.arguments() def argument(V): if du is None: n = max(a.number() for a in args) if args else -1 return Argument(V, n + 1) else: return du if is_dX: coords = mesh.coordinates u = ufl.SpatialCoordinate(mesh) V = coords.function_space() du = argument(V) cds = {coords: du} if coefficient_derivatives is not None: cds.update(coefficient_derivatives) coefficient_derivatives = cds elif isinstance(u, firedrake.Function): V = u.function_space() du = argument(V) elif isinstance(u, firedrake.Constant): if u.ufl_shape != (): raise ValueError("Real function space of vector elements not supported") V = firedrake.FunctionSpace(mesh, "Real", 0) du = argument(V) else: raise RuntimeError("Can't compute derivative for form") return ufl.derivative(form, u, du, coefficient_derivatives)
def compile_expression_at_points(expression, points, coordinates, interface=None, parameters=None, coffee=True): """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 interface: backend module for the kernel interface :arg parameters: parameters object :arg coffee: compile coffee kernel instead of loopy kernel """ import coffee.base as ast import loopy as lp if parameters is None: parameters = default_parameters() else: _ = default_parameters() _.update(parameters) parameters = _ # Determine whether in complex mode complex_mode = is_complex(parameters["scalar_type"]) # Apply UFL preprocessing expression = ufl_utils.preprocess_expression(expression, complex_mode=complex_mode) # Initialise kernel builder if interface is None: if coffee: import tsfc.kernel_interface.firedrake as firedrake_interface_coffee interface = firedrake_interface_coffee.ExpressionKernelBuilder else: # Delayed import, loopy is a runtime dependency import tsfc.kernel_interface.firedrake_loopy as firedrake_interface_loopy interface = firedrake_interface_loopy.ExpressionKernelBuilder builder = interface(parameters["scalar_type"]) arguments = extract_arguments(expression) argument_multiindices = tuple( builder.create_element(arg.ufl_element()).get_indices() for arg in arguments) # Replace coordinates (if any) domain = expression.ufl_domain() if domain: assert coordinates.ufl_domain() == domain builder.domain_coordinate[domain] = coordinates builder.set_cell_sizes(domain) # Collect required coefficients coefficients = extract_coefficients(expression) if has_type(expression, GeometricQuantity) or any( fem.needs_coordinate_mapping(c.ufl_element()) for c in coefficients): coefficients = [coordinates] + coefficients 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, argument_multiindices=argument_multiindices) 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_indices = point_set.indices + tensor_indices + tuple( chain(*argument_multiindices)) return_shape = tuple(i.extent for i in return_indices) return_var = gem.Variable('A', return_shape) if coffee: return_arg = ast.Decl(parameters["scalar_type"], ast.Symbol('A', rank=return_shape)) else: return_arg = lp.GlobalArg("A", dtype=parameters["scalar_type"], shape=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 # Handle kernel interface requirements builder.register_requirements([ir]) # Build kernel tuple return builder.construct_kernel(return_arg, impero_c, parameters["precision"], {point_index: 'p'})
def test_extract_coefficients_vs_fixture(coefficients, forms): assert coefficients == tuple(extract_coefficients(forms[2]))
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_expression_dual_evaluation(expression, to_element, *, domain=None, interface=None, parameters=None, coffee=False): """Compile a UFL expression to be evaluated against a compile-time known reference element's dual basis. Useful for interpolating UFL expressions into e.g. N1curl spaces. :arg expression: UFL expression :arg to_element: A FInAT element for the target space :arg domain: optional UFL domain the expression is defined on (required when expression contains no domain). :arg interface: backend module for the kernel interface :arg parameters: parameters object :arg coffee: compile coffee kernel instead of loopy kernel """ import coffee.base as ast import loopy as lp # Just convert FInAT element to FIAT for now. # Dual evaluation in FInAT will bring a thorough revision. to_element = to_element.fiat_equivalent if any(len(dual.deriv_dict) != 0 for dual in to_element.dual_basis()): raise NotImplementedError( "Can only interpolate onto dual basis functionals without derivative evaluation, sorry!" ) if parameters is None: parameters = default_parameters() else: _ = default_parameters() _.update(parameters) parameters = _ # Determine whether in complex mode complex_mode = is_complex(parameters["scalar_type"]) # Find out which mapping to apply try: mapping, = set(to_element.mapping()) except ValueError: raise NotImplementedError( "Don't know how to interpolate onto zany spaces, sorry") expression = apply_mapping(expression, mapping, domain) # Apply UFL preprocessing expression = ufl_utils.preprocess_expression(expression, complex_mode=complex_mode) # Initialise kernel builder if interface is None: if coffee: import tsfc.kernel_interface.firedrake as firedrake_interface_coffee interface = firedrake_interface_coffee.ExpressionKernelBuilder else: # Delayed import, loopy is a runtime dependency import tsfc.kernel_interface.firedrake_loopy as firedrake_interface_loopy interface = firedrake_interface_loopy.ExpressionKernelBuilder builder = interface(parameters["scalar_type"]) arguments = extract_arguments(expression) argument_multiindices = tuple( builder.create_element(arg.ufl_element()).get_indices() for arg in arguments) # Replace coordinates (if any) unless otherwise specified by kwarg if domain is None: domain = expression.ufl_domain() assert domain is not None # Collect required coefficients first_coefficient_fake_coords = False coefficients = extract_coefficients(expression) if has_type(expression, GeometricQuantity) or any( fem.needs_coordinate_mapping(c.ufl_element()) for c in coefficients): # Create a fake coordinate coefficient for a domain. coords_coefficient = ufl.Coefficient( ufl.FunctionSpace(domain, domain.ufl_coordinate_element())) builder.domain_coordinate[domain] = coords_coefficient builder.set_cell_sizes(domain) coefficients = [coords_coefficient] + coefficients first_coefficient_fake_coords = True builder.set_coefficients(coefficients) # Split mixed coefficients expression = ufl_utils.split_coefficients(expression, builder.coefficient_split) # Translate to GEM kernel_cfg = dict( interface=builder, ufl_cell=domain.ufl_cell(), # FIXME: change if we ever implement # interpolation on facets. integral_type="cell", argument_multiindices=argument_multiindices, index_cache={}, scalar_type=parameters["scalar_type"]) if all( isinstance(dual, PointEvaluation) for dual in to_element.dual_basis()): # This is an optimisation for point-evaluation nodes which # should go away once FInAT offers the interface properly qpoints = [] # Everything is just a point evaluation. for dual in to_element.dual_basis(): ptdict = dual.get_point_dict() qpoint, = ptdict.keys() (qweight, component), = ptdict[qpoint] assert allclose(qweight, 1.0) assert component == () qpoints.append(qpoint) point_set = PointSet(qpoints) config = kernel_cfg.copy() config.update(point_set=point_set) # Allow interpolation onto QuadratureElements to refer to the quadrature # rule they represent if isinstance(to_element, FIAT.QuadratureElement): assert allclose(asarray(qpoints), asarray(to_element._points)) quad_rule = QuadratureRule(point_set, to_element._weights) config["quadrature_rule"] = quad_rule expr, = fem.compile_ufl(expression, **config, point_sum=False) # In some cases point_set.indices may be dropped from expr, but nothing # new should now appear assert set(expr.free_indices) <= set( chain(point_set.indices, *argument_multiindices)) shape_indices = tuple(gem.Index() for _ in expr.shape) basis_indices = point_set.indices ir = gem.Indexed(expr, shape_indices) else: # This is general code but is more unrolled than necssary. dual_expressions = [] # one for each functional broadcast_shape = len(expression.ufl_shape) - len( to_element.value_shape()) shape_indices = tuple(gem.Index() for _ in expression.ufl_shape[:broadcast_shape]) expr_cache = {} # Sharing of evaluation of the expression at points for dual in to_element.dual_basis(): pts = tuple(sorted(dual.get_point_dict().keys())) try: expr, point_set = expr_cache[pts] except KeyError: point_set = PointSet(pts) config = kernel_cfg.copy() config.update(point_set=point_set) expr, = fem.compile_ufl(expression, **config, point_sum=False) # In some cases point_set.indices may be dropped from expr, but # nothing new should now appear assert set(expr.free_indices) <= set( chain(point_set.indices, *argument_multiindices)) expr = gem.partial_indexed(expr, shape_indices) expr_cache[pts] = expr, point_set weights = collections.defaultdict(list) for p in pts: for (w, cmp) in dual.get_point_dict()[p]: weights[cmp].append(w) qexprs = gem.Zero() for cmp in sorted(weights): qweights = gem.Literal(weights[cmp]) qexpr = gem.Indexed(expr, cmp) qexpr = gem.index_sum( gem.Indexed(qweights, point_set.indices) * qexpr, point_set.indices) qexprs = gem.Sum(qexprs, qexpr) assert qexprs.shape == () assert set(qexprs.free_indices) == set( chain(shape_indices, *argument_multiindices)) dual_expressions.append(qexprs) basis_indices = (gem.Index(), ) ir = gem.Indexed(gem.ListTensor(dual_expressions), basis_indices) # Build kernel body return_indices = basis_indices + shape_indices + tuple( chain(*argument_multiindices)) return_shape = tuple(i.extent for i in return_indices) return_var = gem.Variable('A', return_shape) if coffee: return_arg = ast.Decl(parameters["scalar_type"], ast.Symbol('A', rank=return_shape)) else: return_arg = lp.GlobalArg("A", dtype=parameters["scalar_type"], shape=return_shape) return_expr = gem.Indexed(return_var, return_indices) # TODO: one should apply some GEM optimisations as in assembly, # but we don't for now. ir, = impero_utils.preprocess_gem([ir]) impero_c = impero_utils.compile_gem([(return_expr, ir)], return_indices) index_names = dict( (idx, "p%d" % i) for (i, idx) in enumerate(basis_indices)) # Handle kernel interface requirements builder.register_requirements([ir]) # Build kernel tuple return builder.construct_kernel(return_arg, impero_c, index_names, first_coefficient_fake_coords)
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 compute_expression_ir(expression, analysis, parameters, visualise): """Compute IR for expression. Parameters ---------- expression Triple of (UFL expression, array of evaluation points, original UFL expression). Note ---- Original UFL expression is needed to compute original positions of coefficients. """ logger.info("Computing uflacs representation of expression") original_expression = expression[2] points = expression[1] expression = expression[0] num_points = points.shape[0] weights = numpy.array([1.0] * num_points) cell = expression.ufl_domain().ufl_cell() ir = {} # Prepare dimensions of all unique element in expression, # including elements for arguments, coefficients and coordinate mappings ir["element_dimensions"] = { ufl_element: create_element(ufl_element).space_dimension() for ufl_element in analysis.unique_elements } # Extract dimensions for elements of arguments only arguments = extract_arguments(expression) argument_elements = tuple(f.ufl_element() for f in arguments) argument_dimensions = [ ir["element_dimensions"][ufl_element] for ufl_element in argument_elements ] tensor_shape = argument_dimensions ir["tensor_shape"] = tensor_shape ir["expression_shape"] = list(expression.ufl_shape) coefficients = extract_coefficients(expression) coefficient_numbering = {} for i, coeff in enumerate(coefficients): coefficient_numbering[coeff] = i # Add coefficient numbering to IR ir["coefficient_numbering"] = coefficient_numbering original_coefficient_positions = [] original_coefficients = extract_coefficients(original_expression) for coeff in coefficients: original_coefficient_positions.append( original_coefficients.index(coeff)) ir["original_coefficient_positions"] = original_coefficient_positions coefficient_elements = tuple(f.ufl_element() for f in coefficients) offsets = {} _offset = 0 for i, el in enumerate(coefficient_elements): offsets[coefficients[i]] = _offset _offset += ir["element_dimensions"][el] # Copy offsets also into IR ir["coefficient_offsets"] = offsets ir["integral_type"] = "expression" ir["entitytype"] = "cell" # Build offsets for Constants original_constant_offsets = {} _offset = 0 for constant in extract_constants(expression): original_constant_offsets[constant] = _offset _offset += numpy.product(constant.ufl_shape, dtype=numpy.int) ir["original_constant_offsets"] = original_constant_offsets ir["points"] = points integrands = {num_points: expression} quadrature_rules = {num_points: (points, weights)} uflacs_ir = build_uflacs_ir(cell, ir["integral_type"], ir["entitytype"], integrands, tensor_shape, quadrature_rules, parameters, visualise) ir.update(uflacs_ir) return ir
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()