def transfer_kernel(Pk, P1): """Compile a kernel that will map between Pk and P1. :returns: a PyOP2 kernel. The prolongation maps a solution in P1 into Pk using the natural embedding. The restriction maps a residual in the dual of Pk into the dual of P1 (it is the dual of the prolongation), computed using linearity of the test function. """ # Mapping of a residual in Pk into a residual in P1 from coffee import base as coffee from tsfc.coffee import generate as generate_coffee, SCALAR_TYPE from tsfc.parameters import default_parameters from gem import gem, impero_utils as imp # Pk should be at least the same size as P1 assert Pk.finat_element.space_dimension() >= P1.finat_element.space_dimension() # In the general case we should compute this by doing: # numpy.linalg.solve(Pkmass, PkP1mass) Pke = Pk.finat_element._element P1e = P1.finat_element._element # TODO, rework to use finat. matrix = numpy.dot(Pke.dual.to_riesz(P1e.get_nodal_basis()), P1e.get_coeffs().T).T Vout, Vin = P1, Pk weights = gem.Literal(matrix) name = "Pk_P1_mapper" funargs = [] assert Vin.shape == Vout.shape shape = (P1e.space_dimension() * Vout.value_size, Pke.space_dimension() * Vin.value_size) outarg = coffee.Decl(SCALAR_TYPE, coffee.Symbol("A", rank=shape)) i = gem.Index() j = gem.Index() k = gem.Index() indices = i, j, k A = gem.Variable("A", shape) outgem = [gem.Indexed(gem.reshape(A, (P1e.space_dimension(), Vout.value_size), (Pke.space_dimension(), Vin.value_size)), (i, k, j, k))] funargs.append(outarg) expr = gem.Indexed(weights, (i, j)) outgem, = imp.preprocess_gem(outgem) ir = imp.compile_gem([(outgem, expr)], indices) index_names = [(i, "i"), (j, "j"), (k, "k")] body = generate_coffee(ir, index_names, default_parameters()["precision"]) function = coffee.FunDecl("void", name, funargs, body, pred=["static", "inline"]) return op2.Kernel(function, name=function.name)
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