def _entity_support_dofs(self): esd = {} for entity_dim in self.cell.sub_entities.keys(): beta = self.get_indices() zeta = self.get_value_indices() entity_cell = self.cell.construct_subelement(entity_dim) quad = make_quadrature(entity_cell, (2 * numpy.array(self.degree)).tolist()) eps = 1.e-8 # Is this a safe value? result = {} for f in self.entity_dofs()[entity_dim].keys(): # Tabulate basis functions on the facet vals, = self.basis_evaluation(0, quad.point_set, entity=(entity_dim, f)).values() # Integrate the square of the basis functions on the facet. ints = gem.IndexSum( gem.Product( gem.IndexSum( gem.Product(gem.Indexed(vals, beta + zeta), gem.Indexed(vals, beta + zeta)), zeta), quad.weight_expression), quad.point_set.indices) evaluation, = evaluate([gem.ComponentTensor(ints, beta)]) ints = evaluation.arr.flatten() assert evaluation.fids == () result[f] = [dof for dof, i in enumerate(ints) if i > eps] esd[entity_dim] = result return esd
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_coefficient(terminal, mt, ctx): vec = ctx.coefficient(terminal, mt.restriction) if terminal.ufl_element().family() == 'Real': assert mt.local_derivatives == 0 return vec element = ctx.create_element(terminal.ufl_element()) # Collect FInAT tabulation for all entities per_derivative = collections.defaultdict(list) for entity_id in ctx.entity_ids: finat_dict = ctx.basis_evaluation(element, mt.local_derivatives, entity_id) for alpha, table in finat_dict.items(): # Filter out irrelevant derivatives if sum(alpha) == mt.local_derivatives: # A numerical hack that FFC used to apply on FIAT # tables still lives on after ditching FFC and # switching to FInAT. table = ffc_rounding(table, ctx.epsilon) per_derivative[alpha].append(table) # Merge entity tabulations for each derivative if len(ctx.entity_ids) == 1: def take_singleton(xs): x, = xs # asserts singleton return x per_derivative = {alpha: take_singleton(tables) for alpha, tables in per_derivative.items()} else: f = ctx.entity_number(mt.restriction) per_derivative = {alpha: gem.select_expression(tables, f) for alpha, tables in per_derivative.items()} # Coefficient evaluation ctx.index_cache.setdefault(terminal.ufl_element(), element.get_indices()) beta = ctx.index_cache[terminal.ufl_element()] zeta = element.get_value_indices() vec_beta, = gem.optimise.remove_componenttensors([gem.Indexed(vec, beta)]) value_dict = {} for alpha, table in per_derivative.items(): table_qi = gem.Indexed(table, beta + zeta) summands = [] for var, expr in unconcatenate([(vec_beta, table_qi)], ctx.index_cache): value = gem.IndexSum(gem.Product(expr, var), var.index_ordering()) summands.append(gem.optimise.contraction(value)) optimised_value = gem.optimise.make_sum(summands) value_dict[alpha] = gem.ComponentTensor(optimised_value, zeta) # Change from FIAT to UFL arrangement result = fiat_to_ufl(value_dict, mt.local_derivatives) assert result.shape == mt.expr.ufl_shape assert set(result.free_indices) <= set(ctx.point_indices) # Detect Jacobian of affine cells if not result.free_indices and all(numpy.count_nonzero(node.array) <= 2 for node in traversal((result,)) if isinstance(node, gem.Literal)): result = gem.optimise.aggressive_unroll(result) return result
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 matvec(table): i, j = gem.indices(2) value_indices = self.get_value_indices() table = gem.Indexed(table, (j, ) + value_indices) val = gem.ComponentTensor(gem.IndexSum(M[i, j] * table, (j, )), (i, ) + value_indices) # Eliminate zeros return gem.optimise.aggressive_unroll(val)
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 entity_support_dofs(elem, entity_dim): """Return the map of entity id to the degrees of freedom for which the corresponding basis functions take non-zero values. :arg elem: FInAT finite element :arg entity_dim: Dimension of the cell subentity. """ if not hasattr(elem, "_entity_support_dofs"): elem._entity_support_dofs = {} cache = elem._entity_support_dofs try: return cache[entity_dim] except KeyError: pass beta = elem.get_indices() zeta = elem.get_value_indices() entity_cell = elem.cell.construct_subelement(entity_dim) quad = make_quadrature(entity_cell, (2 * numpy.array(elem.degree)).tolist()) eps = 1.e-8 # Is this a safe value? result = {} for f in elem.entity_dofs()[entity_dim].keys(): # Tabulate basis functions on the facet vals, = itervalues( elem.basis_evaluation(0, quad.point_set, entity=(entity_dim, f))) # Integrate the square of the basis functions on the facet. ints = gem.IndexSum( gem.Product( gem.IndexSum( gem.Product(gem.Indexed(vals, beta + zeta), gem.Indexed(vals, beta + zeta)), zeta), quad.weight_expression), quad.point_set.indices) ints = aggressive_unroll(gem.ComponentTensor(ints, beta)).array.flatten() result[f] = [dof for dof, i in enumerate(ints) if i > eps] cache[entity_dim] = result 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 Integrals(expressions, quadrature_multiindex, argument_multiindices, parameters): # Concatenate expressions = concatenate(expressions) # Unroll max_extent = parameters["unroll_indexsum"] if max_extent: def predicate(index): return index.extent <= max_extent expressions = unroll_indexsum(expressions, predicate=predicate) # Refactorise def classify(quadrature_indices, expression): if not quadrature_indices.intersection(expression.free_indices): return OTHER elif isinstance(expression, gem.Indexed) and isinstance( expression.children[0], gem.Literal): return ATOMIC else: return COMPOUND classifier = partial(classify, set(quadrature_multiindex)) result = [] for expr, monomial_sum in zip(expressions, collect_monomials(expressions, classifier)): # Select quadrature indices that are present quadrature_indices = set(index for index in quadrature_multiindex if index in expr.free_indices) products = [] for sum_indices, factors, rest in monomial_sum: # Collapse quadrature literals for each monomial if factors or quadrature_indices: replacement = einsum(remove_componenttensors(factors), quadrature_indices) else: replacement = gem.Literal(1) # Rebuild expression products.append( gem.IndexSum(gem.Product(replacement, rest), sum_indices)) result.append(reduce(gem.Sum, products, gem.Zero())) return result
def dual_evaluation(self, fn): '''Get a GEM expression for performing the dual basis evaluation at the nodes of the reference element. Currently only works for flat elements: tensor elements are implemented in :class:`TensorFiniteElement`. :param fn: Callable representing the function to dual evaluate. Callable should take in an :class:`AbstractPointSet` and return a GEM expression for evaluation of the function at those points. :returns: A tuple ``(dual_evaluation_gem_expression, basis_indices)`` where the given ``basis_indices`` are those needed to form a return expression for the code which is compiled from ``dual_evaluation_gem_expression`` (alongside any argument multiindices already encoded within ``fn``) ''' Q, 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 == Q.shape[len(Q.shape) - len(expr.shape):] shape_indices = gem.indices(len(expr.shape)) basis_indices = gem.indices(len(Q.shape) - len(expr.shape)) Qi = Q[basis_indices + shape_indices] expri = expr[shape_indices] evaluation = gem.IndexSum(Qi * expri, x.indices + shape_indices) # Now we want to factorise over the new contraction with x, # ignoring any shape indices to avoid hitting the sum- # factorisation index limit (this is a bit of a hack). # Really need to do a more targeted job here. evaluation = gem.optimise.contraction(evaluation, shape_indices) return evaluation, basis_indices
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 index_sum(self, o, summand, indices): index, = indices indices = gem.indices(len(summand.shape)) return gem.ComponentTensor( gem.IndexSum(gem.Indexed(summand, indices), (index, )), indices)