def dual_evaluation(self, fn): tQ, x = self.dual_basis expr = fn(x) # Apply targeted sum factorisation and delta elimination to # the expression sum_indices, factors = delta_elimination(*traverse_product(expr)) expr = sum_factorise(sum_indices, factors) # NOTE: any shape indices in the expression are because the # expression is tensor valued. assert expr.shape == self.value_shape scalar_i = self.base_element.get_indices() scalar_vi = self.base_element.get_value_indices() tensor_i = tuple(gem.Index(extent=d) for d in self._shape) tensor_vi = tuple(gem.Index(extent=d) for d in self._shape) if self._transpose: index_ordering = tensor_i + scalar_i + tensor_vi + scalar_vi else: index_ordering = scalar_i + tensor_i + tensor_vi + scalar_vi tQi = tQ[index_ordering] expri = expr[tensor_i + scalar_vi] evaluation = gem.IndexSum(tQi * expri, x.indices + scalar_vi + tensor_i) # This doesn't work perfectly, the resulting code doesn't have # a minimal memory footprint, although the operation count # does appear to be minimal. evaluation = gem.optimise.contraction(evaluation) return evaluation, scalar_i + tensor_vi
def _tensorise(self, scalar_evaluation): # Old basis function and value indices scalar_i = self._base_element.get_indices() scalar_vi = self._base_element.get_value_indices() # New basis function and value indices tensor_i = tuple(gem.Index(extent=d) for d in self._shape) tensor_vi = tuple(gem.Index(extent=d) for d in self._shape) # Couple new basis function and value indices deltas = reduce(gem.Product, (gem.Delta(j, k) for j, k in zip(tensor_i, tensor_vi))) if self._transpose: index_ordering = tensor_i + scalar_i + tensor_vi + scalar_vi else: index_ordering = scalar_i + tensor_i + tensor_vi + scalar_vi result = {} for alpha, expr in scalar_evaluation.items(): result[alpha] = gem.ComponentTensor( gem.Product(deltas, gem.Indexed(expr, scalar_i + scalar_vi)), index_ordering ) return result
def basis_evaluation(self, order, ps, entity=None): r"""Produce the recipe for basis function evaluation at a set of points :math:`q`: .. math:: \boldsymbol\phi_{(\gamma \epsilon) (i \alpha \beta) q} = \delta_{\alpha \gamma}\delta{\beta \epsilon}\phi_{i q} \nabla\boldsymbol\phi_{(\epsilon \gamma \zeta) (i \alpha \beta) q} = \delta_{\alpha \epsilon} \deta{\beta \gamma}\nabla\phi_{\zeta i q} """ # Old basis function and value indices scalar_i = self._base_element.get_indices() scalar_vi = self._base_element.get_value_indices() # New basis function and value indices tensor_i = tuple(gem.Index(extent=d) for d in self._shape) tensor_vi = tuple(gem.Index(extent=d) for d in self._shape) # Couple new basis function and value indices deltas = reduce(gem.Product, (gem.Delta(j, k) for j, k in zip(tensor_i, tensor_vi))) scalar_result = self._base_element.basis_evaluation(order, ps, entity) result = {} for alpha, expr in iteritems(scalar_result): result[alpha] = gem.ComponentTensor( gem.Product(deltas, gem.Indexed(expr, scalar_i + scalar_vi)), scalar_i + tensor_i + scalar_vi + tensor_vi) return result
def point_evaluation_ciarlet(fiat_element, order, refcoords, entity): # Coordinates on the reference entity (SymPy) esd, = refcoords.shape Xi = sp.symbols('X Y Z')[:esd] # Coordinates on the reference cell cell = fiat_element.get_reference_element() X = cell.get_entity_transform(*entity)(Xi) # Evaluate expansion set at SymPy point poly_set = fiat_element.get_nodal_basis() degree = poly_set.get_embedded_degree() base_values = poly_set.get_expansion_set().tabulate(degree, [X]) m = len(base_values) assert base_values.shape == (m, 1) base_values_sympy = np.array(list(base_values.flat)) # Find constant polynomials def is_const(expr): try: float(expr) return True except TypeError: return False const_mask = np.array(list(map(is_const, base_values_sympy))) # Convert SymPy expression to GEM mapper = gem.node.Memoizer(sympy2gem) mapper.bindings = { s: gem.Indexed(refcoords, (i, )) for i, s in enumerate(Xi) } base_values = gem.ListTensor(list(map(mapper, base_values.flat))) # Populate result dict, creating precomputed coefficient # matrices for each derivative tuple. result = {} for i in range(order + 1): for alpha in mis(cell.get_spatial_dimension(), i): D = form_matrix_product(poly_set.get_dmats(), alpha) table = np.dot(poly_set.get_coeffs(), np.transpose(D)) assert table.shape[-1] == m zerocols = np.isclose( abs(table).max(axis=tuple(range(table.ndim - 1))), 0.0) if all(np.logical_or(const_mask, zerocols)): # Casting is safe by assertion of is_const vals = base_values_sympy[const_mask].astype(np.float64) result[alpha] = gem.Literal(table[..., const_mask].dot(vals)) else: beta = tuple(gem.Index() for s in table.shape[:-1]) k = gem.Index() result[alpha] = gem.ComponentTensor( gem.IndexSum( gem.Product( gem.Indexed(gem.Literal(table), beta + (k, )), gem.Indexed(base_values, (k, ))), (k, )), beta) return result
def test_pickle_gem(protocol): f = gem.VariableIndex(gem.Indexed(gem.Variable('facet', (2, )), (1, ))) q = gem.Index() r = gem.Index() _1 = gem.Indexed(gem.Literal(numpy.random.rand(3, 6, 8)), (f, q, r)) _2 = gem.Indexed( gem.view(gem.Variable('w', (None, None)), slice(8), slice(1)), (r, 0)) expr = gem.ComponentTensor(gem.IndexSum(gem.Product(_1, _2), (r, )), (q, )) unpickled = pickle.loads(pickle.dumps(expr, protocol)) assert repr(expr) == repr(unpickled)
def translate_cell_edge_vectors(terminal, mt, ctx): # WARNING: Assumes straight edges! coords = CellVertices(terminal.ufl_domain()) ufl_expr = construct_modified_terminal(mt, coords) cell_vertices = ctx.translator(ufl_expr) e = gem.Index() c = gem.Index() expr = gem.ListTensor([ gem.Sum( gem.Indexed(cell_vertices, (u, c)), gem.Product(gem.Literal(-1), gem.Indexed(cell_vertices, (v, c)))) for _, (u, v) in sorted(ctx.fiat_cell.get_topology()[1].items()) ]) return gem.ComponentTensor(gem.Indexed(expr, (e, )), (e, c))
def physical_points(self, point_set, entity=None): """Converts point_set from reference to physical space""" expr = SpatialCoordinate(self.mt.terminal.ufl_domain()) point_shape, = point_set.expression.shape if entity is not None: e, _ = entity assert point_shape == e else: assert point_shape == expr.ufl_domain().topological_dimension() if self.mt.restriction == '+': expr = PositiveRestricted(expr) elif self.mt.restriction == '-': expr = NegativeRestricted(expr) config = {"point_set": point_set} config.update(self.config) if entity is not None: config.update({ name: getattr(self.interface, name) for name in ["integration_dim", "entity_ids"] }) context = PointSetContext(**config) expr = self.preprocess(expr, context) mapped = map_expr_dag(context.translator, expr) indices = tuple(gem.Index() for _ in mapped.shape) return gem.ComponentTensor(gem.Indexed(mapped, indices), point_set.indices + indices)
def translate_argument(terminal, mt, ctx): argument_multiindex = ctx.argument_multiindices[terminal.number()] sigma = tuple(gem.Index(extent=d) for d in mt.expr.ufl_shape) element = ctx.create_element(terminal.ufl_element(), restriction=mt.restriction) def callback(entity_id): finat_dict = ctx.basis_evaluation(element, mt, entity_id) # Filter out irrelevant derivatives filtered_dict = { alpha: table for alpha, table in finat_dict.items() if sum(alpha) == mt.local_derivatives } # Change from FIAT to UFL arrangement square = fiat_to_ufl(filtered_dict, mt.local_derivatives) # A numerical hack that FFC used to apply on FIAT tables still # lives on after ditching FFC and switching to FInAT. return ffc_rounding(square, ctx.epsilon) table = ctx.entity_selector(callback, mt.restriction) return gem.ComponentTensor(gem.Indexed(table, argument_multiindex + sigma), sigma)
def transpose(expr, rank): assert not expr.free_indices assert 0 <= rank < len(expr.shape) if rank == 0: return expr else: indices = tuple(gem.Index(extent=extent) for extent in expr.shape) transposed_indices = indices[rank:] + indices[:rank] return gem.ComponentTensor(gem.Indexed(expr, indices), transposed_indices)
def point_evaluation_generic(fiat_element, order, refcoords, entity): # Coordinates on the reference entity (SymPy) esd, = refcoords.shape Xi = sp.symbols('X Y Z')[:esd] space_dimension = fiat_element.space_dimension() value_size = np.prod(fiat_element.value_shape(), dtype=int) fiat_result = fiat_element.tabulate(order, [Xi], entity) result = {} for alpha, fiat_table in fiat_result.items(): if isinstance(fiat_table, Exception): result[alpha] = gem.Failure( (space_dimension, ) + fiat_element.value_shape(), fiat_table) continue # Convert SymPy expression to GEM mapper = gem.node.Memoizer(sympy2gem) mapper.bindings = { s: gem.Indexed(refcoords, (i, )) for i, s in enumerate(Xi) } gem_table = np.vectorize(mapper)(fiat_table) table_roll = gem_table.reshape(space_dimension, value_size).transpose() exprs = [] for table in table_roll: exprs.append(gem.ListTensor(table.reshape(space_dimension))) if fiat_element.value_shape(): beta = (gem.Index(extent=space_dimension), ) zeta = tuple( gem.Index(extent=d) for d in fiat_element.value_shape()) result[alpha] = gem.ComponentTensor( gem.Indexed( gem.ListTensor( np.array([gem.Indexed(expr, beta) for expr in exprs ]).reshape(fiat_element.value_shape())), zeta), beta + zeta) else: expr, = exprs result[alpha] = expr return result
def fiat_to_ufl(fiat_dict, order): # All derivative multiindices must be of the same dimension. dimension, = set(len(alpha) for alpha in fiat_dict.keys()) # All derivative tables must have the same shape. shape, = set(table.shape for table in fiat_dict.values()) sigma = tuple(gem.Index(extent=extent) for extent in shape) # Convert from FIAT to UFL format eye = numpy.eye(dimension, dtype=int) tensor = numpy.empty((dimension, ) * order, dtype=object) for multiindex in numpy.ndindex(tensor.shape): alpha = tuple(eye[multiindex, :].sum(axis=0)) tensor[multiindex] = gem.Indexed(fiat_dict[alpha], sigma) delta = tuple(gem.Index(extent=dimension) for _ in range(order)) if order > 0: tensor = gem.Indexed(gem.ListTensor(tensor), delta) else: tensor = tensor[()] return gem.ComponentTensor(tensor, sigma + delta)
def test_cellwise_constant(cell, degree): dim = cell.get_spatial_dimension() element = finat.Lagrange(cell, degree) index = gem.Index() point = gem.partial_indexed(gem.Variable('X', (17, dim)), (index, )) order = 2 for alpha, table in element.point_evaluation(order, point).items(): if sum(alpha) < degree: assert table.free_indices == (index, ) else: assert table.free_indices == ()
def dual_basis(self): base = self.base_element Q, points = base.dual_basis # Suppose the tensor element has shape (2, 4) # These identity matrices may have difference sizes depending the shapes # tQ = Q ⊗ 𝟙₂ ⊗ 𝟙₄ scalar_i = self.base_element.get_indices() scalar_vi = self.base_element.get_value_indices() tensor_i = tuple(gem.Index(extent=d) for d in self._shape) tensor_vi = tuple(gem.Index(extent=d) for d in self._shape) # Couple new basis function and value indices deltas = reduce(gem.Product, (gem.Delta(j, k) for j, k in zip(tensor_i, tensor_vi))) if self._transpose: index_ordering = tensor_i + scalar_i + tensor_vi + scalar_vi else: index_ordering = scalar_i + tensor_i + tensor_vi + scalar_vi Qi = Q[scalar_i + scalar_vi] tQ = gem.ComponentTensor(Qi*deltas, index_ordering) return tQ, points
def test_refactorise(): f = gem.Variable('f', (3,)) u = gem.Variable('u', (3,)) v = gem.Variable('v', ()) i = gem.Index() f_i = gem.Indexed(f, (i,)) u_i = gem.Indexed(u, (i,)) def classify(atomics_set, expression): if expression in atomics_set: return ATOMIC for node in traversal([expression]): if node in atomics_set: return COMPOUND return OTHER classifier = partial(classify, {u_i, v}) # \sum_i 5*(2*u_i + -1*v)*(u_i + v*f) expr = gem.IndexSum( gem.Product( gem.Literal(5), gem.Product( gem.Sum(gem.Product(gem.Literal(2), u_i), gem.Product(gem.Literal(-1), v)), gem.Sum(u_i, gem.Product(v, f_i)) ) ), (i,) ) expected = [ Monomial((i,), (u_i, u_i), gem.Literal(10)), Monomial((i,), (u_i, v), gem.Product(gem.Literal(5), gem.Sum(gem.Product(f_i, gem.Literal(2)), gem.Literal(-1)))), Monomial((), (v, v), gem.Product(gem.Literal(5), gem.IndexSum(gem.Product(f_i, gem.Literal(-1)), (i,)))), ] actual, = collect_monomials([expr], classifier) assert expected == list(actual)
def translate_cell_vertices(terminal, mt, ctx): coords = SpatialCoordinate(terminal.ufl_domain()) ufl_expr = construct_modified_terminal(mt, coords) ps = PointSet(numpy.array(ctx.fiat_cell.get_vertices())) config = {name: getattr(ctx, name) for name in ["ufl_cell", "precision", "index_cache"]} config.update(interface=ctx, point_set=ps) context = PointSetContext(**config) expr = context.translator(ufl_expr) # Wrap up point (vertex) index c = gem.Index() return gem.ComponentTensor(gem.Indexed(expr, (c,)), ps.indices + (c,))
def __init__(self, scalar_type, num_entities): """:arg num_entities: the number of micro-entities to integrate over.""" super().__init__(scalar_type) self.indices = (gem.Index("entity", extent=num_entities), ) self.shape = tuple(i.extent for i in self.indices) self.unsummed_coefficient_indices = frozenset(self.indices)
def transpose(expr): indices = tuple(gem.Index(extent=extent) for extent in expr.shape) transposed_indices = indices[scalar_rank:] + indices[:scalar_rank] return gem.ComponentTensor(gem.Indexed(expr, indices), transposed_indices)
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 get_value_indices(self): '''A tuple of GEM :class:`~gem.Index` of the correct extents to loop over the value shape of this element.''' return tuple(gem.Index(extent=d) for d in self.value_shape)
def get_indices(self): '''A tuple of GEM :class:`Index` of the correct extents to loop over the basis functions of this element.''' return tuple(gem.Index(extent=d) for d in self.index_shape)
def basis_evaluation(self, order, ps, entity=None, coordinate_mapping=None): '''Return code for evaluating the element at known points on the reference element. :param order: return derivatives up to this order. :param ps: the point set. :param entity: the cell entity on which to tabulate. ''' space_dimension = self._element.space_dimension() value_size = np.prod(self._element.value_shape(), dtype=int) fiat_result = self._element.tabulate(order, ps.points, entity) result = {} # In almost all cases, we have # self.space_dimension() == self._element.space_dimension() # But for Bell, FIAT reports 21 basis functions, # but FInAT only 18 (because there are actually 18 # basis functions, and the additional 3 are for # dealing with transformations between physical # and reference space). index_shape = (self._element.space_dimension(), ) for alpha, fiat_table in fiat_result.items(): if isinstance(fiat_table, Exception): result[alpha] = gem.Failure( self.index_shape + self.value_shape, fiat_table) continue derivative = sum(alpha) table_roll = fiat_table.reshape(space_dimension, value_size, len(ps.points)).transpose(1, 2, 0) exprs = [] for table in table_roll: if derivative < self.degree: point_indices = ps.indices point_shape = tuple(index.extent for index in point_indices) exprs.append( gem.partial_indexed( gem.Literal( table.reshape(point_shape + index_shape)), point_indices)) elif derivative == self.degree: # Make sure numerics satisfies theory exprs.append(gem.Literal(table[0])) else: # Make sure numerics satisfies theory assert np.allclose(table, 0.0) exprs.append(gem.Zero(self.index_shape)) if self.value_shape: # As above, this extent may be different from that # advertised by the finat element. beta = tuple(gem.Index(extent=i) for i in index_shape) assert len(beta) == len(self.get_indices()) zeta = self.get_value_indices() result[alpha] = gem.ComponentTensor( gem.Indexed( gem.ListTensor( np.array([ gem.Indexed(expr, beta) for expr in exprs ]).reshape(self.value_shape)), zeta), beta + zeta) else: expr, = exprs result[alpha] = expr return result
def indices(self): N, _ = self._points_expr.shape return (gem.Index(extent=N), )
def indices(self): return (gem.Index(extent=len(self.points)),)
def _dual_basis(self): # Return the numerical part of the dual basis, this split is # needed because the dual_basis itself can't produce the same # point set over and over in case it is used multiple times # (in for example a tensorproductelement). fiat_dual_basis = self._element.dual_basis() seen = dict() allpts = [] # Find the unique points to evaluate at. # We might be able to make this a smaller set by treating each # point one by one, but most of the redundancy comes from # multiple functionals using the same quadrature rule. for dual in fiat_dual_basis: if len(dual.deriv_dict) != 0: raise NotImplementedError( "FIAT dual bases with derivative nodes represented via a ``Functional.deriv_dict`` property do not currently have a FInAT dual basis" ) pts = dual.get_point_dict().keys() pts = tuple(sorted(pts)) # need this for determinism if pts not in seen: # k are indices into Q (see below) for the seen points kstart = len(allpts) kend = kstart + len(pts) seen[pts] = kstart, kend allpts.extend(pts) # Build Q. # Q is a tensor of weights (of total rank R) to contract with a unique # vector of points to evaluate at, giving a tensor (of total rank R-1) # where the first indices (rows) correspond to a basis functional # (node). # Q is a DOK Sparse matrix in (row, col, higher,..)=>value pairs (to # become a gem.SparseLiteral when implemented). # Rows (i) are number of nodes/dual functionals. # Columns (k) are unique points to evaluate. # Higher indices (*cmp) are tensor indices of the weights when weights # are tensor valued. Q = {} for i, dual in enumerate(fiat_dual_basis): point_dict = dual.get_point_dict() pts = tuple(sorted(point_dict.keys())) kstart, kend = seen[pts] for p, k in zip(pts, range(kstart, kend)): for weight, cmp in point_dict[p]: Q[(i, k, *cmp)] = weight if all( len(set(key)) == 1 and np.isclose(weight, 1) and len(key) == 2 for key, weight in Q.items()): # Identity matrix Q can be expressed symbolically extents = tuple(map(max, zip(*Q.keys()))) js = tuple(gem.Index(extent=e + 1) for e in extents) assert len(js) == 2 Q = gem.ComponentTensor(gem.Delta(*js), js) else: # temporary until sparse literals are implemented in GEM which will # automatically convert a dictionary of keys internally. # TODO the below is unnecessarily slow and would be sped up # significantly by building Q in a COO format rather than DOK (i.e. # storing coords and associated data in (nonzeros, entries) shaped # numpy arrays) to take advantage of numpy multiindexing Qshape = tuple(s + 1 for s in map(max, *Q)) Qdense = np.zeros(Qshape, dtype=np.float64) for idx, value in Q.items(): Qdense[idx] = value Q = gem.Literal(Qdense) return Q, np.asarray(allpts)
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 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) # 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 dg_injection_kernel(Vf, Vc, ncell): from firedrake import Tensor, AssembledVector, TestFunction, TrialFunction from firedrake.slate.slac import compile_expression macro_builder = MacroKernelBuilder(ScalarType_c, ncell) f = ufl.Coefficient(Vf) macro_builder.set_coefficients([f]) macro_builder.set_coordinates(Vf.mesh()) Vfe = create_element(Vf.ufl_element()) macro_quadrature_rule = make_quadrature( Vfe.cell, estimate_total_polynomial_degree(ufl.inner(f, f))) index_cache = {} parameters = default_parameters() integration_dim, entity_ids = lower_integral_type(Vfe.cell, "cell") macro_cfg = dict(interface=macro_builder, ufl_cell=Vf.ufl_cell(), precision=parameters["precision"], integration_dim=integration_dim, entity_ids=entity_ids, index_cache=index_cache, quadrature_rule=macro_quadrature_rule) fexpr, = fem.compile_ufl(f, **macro_cfg) X = ufl.SpatialCoordinate(Vf.mesh()) C_a, = fem.compile_ufl(X, **macro_cfg) detJ = ufl_utils.preprocess_expression( abs(ufl.JacobianDeterminant(f.ufl_domain()))) macro_detJ, = fem.compile_ufl(detJ, **macro_cfg) Vce = create_element(Vc.ufl_element()) coarse_builder = firedrake_interface.KernelBuilder("cell", "otherwise", 0, ScalarType_c) coarse_builder.set_coordinates(Vc.mesh()) argument_multiindices = (Vce.get_indices(), ) argument_multiindex, = argument_multiindices return_variable, = coarse_builder.set_arguments((ufl.TestFunction(Vc), ), argument_multiindices) integration_dim, entity_ids = lower_integral_type(Vce.cell, "cell") # Midpoint quadrature for jacobian on coarse cell. quadrature_rule = make_quadrature(Vce.cell, 0) coarse_cfg = dict(interface=coarse_builder, ufl_cell=Vc.ufl_cell(), precision=parameters["precision"], integration_dim=integration_dim, entity_ids=entity_ids, index_cache=index_cache, quadrature_rule=quadrature_rule) X = ufl.SpatialCoordinate(Vc.mesh()) K = ufl_utils.preprocess_expression(ufl.JacobianInverse(Vc.mesh())) C_0, = fem.compile_ufl(X, **coarse_cfg) K, = fem.compile_ufl(K, **coarse_cfg) i = gem.Index() j = gem.Index() C_0 = gem.Indexed(C_0, (j, )) C_0 = gem.index_sum(C_0, quadrature_rule.point_set.indices) C_a = gem.Indexed(C_a, (j, )) X_a = gem.Sum(C_0, gem.Product(gem.Literal(-1), C_a)) K_ij = gem.Indexed(K, (i, j)) K_ij = gem.index_sum(K_ij, quadrature_rule.point_set.indices) X_a = gem.index_sum(gem.Product(K_ij, X_a), (j, )) C_0, = quadrature_rule.point_set.points C_0 = gem.Indexed(gem.Literal(C_0), (i, )) # fine quad points in coarse reference space. X_a = gem.Sum(C_0, gem.Product(gem.Literal(-1), X_a)) X_a = gem.ComponentTensor(X_a, (i, )) # Coarse basis function evaluated at fine quadrature points phi_c = fem.fiat_to_ufl( Vce.point_evaluation(0, X_a, (Vce.cell.get_dimension(), 0)), 0) tensor_indices = tuple(gem.Index(extent=d) for d in f.ufl_shape) phi_c = gem.Indexed(phi_c, argument_multiindex + tensor_indices) fexpr = gem.Indexed(fexpr, tensor_indices) quadrature_weight = macro_quadrature_rule.weight_expression expr = gem.Product(gem.IndexSum(gem.Product(phi_c, fexpr), tensor_indices), gem.Product(macro_detJ, quadrature_weight)) quadrature_indices = macro_builder.indices + macro_quadrature_rule.point_set.indices reps = spectral.Integrals([expr], quadrature_indices, argument_multiindices, parameters) assignments = spectral.flatten([(return_variable, reps)], index_cache) return_variables, expressions = zip(*assignments) expressions = impero_utils.preprocess_gem(expressions, **spectral.finalise_options) assignments = list(zip(return_variables, expressions)) impero_c = impero_utils.compile_gem(assignments, quadrature_indices + argument_multiindex, remove_zeros=True) index_names = [] def name_index(index, name): index_names.append((index, name)) if index in index_cache: for multiindex, suffix in zip(index_cache[index], string.ascii_lowercase): name_multiindex(multiindex, name + suffix) def name_multiindex(multiindex, name): if len(multiindex) == 1: name_index(multiindex[0], name) else: for i, index in enumerate(multiindex): name_index(index, name + str(i)) name_multiindex(quadrature_indices, 'ip') for multiindex, name in zip(argument_multiindices, ['j', 'k']): name_multiindex(multiindex, name) index_names.extend(zip(macro_builder.indices, ["entity"])) body = generate_coffee(impero_c, index_names, parameters["precision"], ScalarType_c) retarg = ast.Decl(ScalarType_c, ast.Symbol("R", rank=(Vce.space_dimension(), ))) local_tensor = coarse_builder.local_tensor local_tensor.init = ast.ArrayInit( numpy.zeros(Vce.space_dimension(), dtype=ScalarType_c)) body.children.insert(0, local_tensor) args = [retarg] + macro_builder.kernel_args + [ macro_builder.coordinates_arg, coarse_builder.coordinates_arg ] # Now we have the kernel that computes <f, phi_c>dx_c # So now we need to hit it with the inverse mass matrix on dx_c u = TrialFunction(Vc) v = TestFunction(Vc) expr = Tensor(ufl.inner(u, v) * ufl.dx).inv * AssembledVector( ufl.Coefficient(Vc)) Ainv, = compile_expression(expr) Ainv = Ainv.kinfo.kernel A = ast.Symbol(local_tensor.sym.symbol) R = ast.Symbol("R") body.children.append( ast.FunCall(Ainv.name, R, coarse_builder.coordinates_arg.sym, A)) from coffee.base import Node assert isinstance(Ainv._code, Node) return op2.Kernel(ast.Node([ Ainv._code, ast.FunDecl("void", "pyop2_kernel_injection_dg", args, body, pred=["static", "inline"]) ]), name="pyop2_kernel_injection_dg", cpp=True, include_dirs=Ainv._include_dirs, headers=Ainv._headers)
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 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)