def basis_evaluation(self, order, ps, entity=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 = {} for alpha, fiat_table in iteritems(fiat_result): 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 + self.index_shape)), point_indices)) elif derivative == self.degree: # Make sure numerics satisfies theory assert np.allclose(table, table.mean(axis=0, keepdims=True)) 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: beta = 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 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 dual_basis(self): Q, x = self.wrappee.dual_basis beta = self.get_indices() zeta = self.get_value_indices() # Index out the basis indices from wrapee's Q, to get # something of wrappee.value_shape, then promote to new shape # with the same transform as done for basis evaluation Q = gem.ListTensor(self.transform(gem.partial_indexed(Q, beta))) # Finally wrap up Q in shape again (now with some extra # value_shape indices) return gem.ComponentTensor(Q[zeta], beta + zeta), x
def basis_evaluation(self, order, ps, entity=None): # Default entity if entity is None: entity = (self.cell.get_dimension(), 0) entity_dim, entity_id = entity # Factor entity assert isinstance(entity_dim, tuple) assert len(entity_dim) == len(self.factors) shape = tuple( len(c.get_topology()[d]) for c, d in zip(self.cell.cells, entity_dim)) entities = list(zip(entity_dim, numpy.unravel_index(entity_id, shape))) # Factor point set ps_factors = factor_point_set(self.cell, entity_dim, ps) # Subelement results factor_results = [ fe.basis_evaluation(order, ps_, e) for fe, ps_, e in zip(self.factors, ps_factors, entities) ] # Spatial dimension dimension = self.cell.get_spatial_dimension() # A list of slices that are used to select dimensions # corresponding to each subelement. dim_slices = TensorProductCell._split_slices( [c.get_spatial_dimension() for c in self.cell.cells]) # A list of multiindices, one multiindex per subelement, each # multiindex describing the shape of basis functions of the # subelement. alphas = [fe.get_indices() for fe in self.factors] result = {} for derivative in range(order + 1): for Delta in mis(dimension, derivative): # Split the multiindex for the subelements deltas = [Delta[s] for s in dim_slices] # GEM scalars (can have free indices) for collecting # the contributions from the subelements. scalars = [] for fr, delta, alpha in zip(factor_results, deltas, alphas): # Turn basis shape to free indices, select the # right derivative entry, and collect the result. scalars.append(gem.Indexed(fr[delta], alpha)) # Multiply the values from the subelements and wrap up # non-point indices into shape. result[Delta] = gem.ComponentTensor( reduce(gem.Product, scalars), tuple(chain(*alphas))) return result
def dual_basis(self): # Return Q with x.indices already a free index for the # consumer to use # expensive numerical extraction is done once per element # instance, but the point set must be created every time we # build the dual. Q, pts = self._dual_basis x = PointSet(pts) assert len(x.indices) == 1 assert Q.shape[1] == x.indices[0].extent i, *js = gem.indices(len(Q.shape) - 1) Q = gem.ComponentTensor(gem.Indexed(Q, (i, *x.indices, *js)), (i, *js)) return Q, x
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 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 dual_basis(self): # Outer product the dual bases of the factors qs, pss = zip(*(factor.dual_basis for factor in self.factors)) ps = TensorPointSet(pss) # Naming as _merge_evaluations above alphas = [factor.get_indices() for factor in self.factors] zetas = [factor.get_value_indices() for factor in self.factors] # Index the factors by so that we can reshape into index-shape # followed by value-shape qis = [q[alpha + zeta] for q, alpha, zeta in zip(qs, alphas, zetas)] Q = gem.ComponentTensor( reduce(gem.Product, qis), tuple(chain(*(alphas + zetas))) ) return Q, ps
def basis_evaluation(self, order, ps, entity=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. ''' result = super(GaussLegendre, self).basis_evaluation(order, ps, entity) cell_dimension = self.cell.get_dimension() if entity is None or entity == (cell_dimension, 0): # on cell interior space_dim = self.space_dimension() if isinstance(ps, GaussLegendrePointSet) and len(ps.points) == space_dim: # Bingo: evaluation points match node locations! spatial_dim = self.cell.get_spatial_dimension() q, = ps.indices r, = self.get_indices() result[(0,) * spatial_dim] = gem.ComponentTensor(gem.Delta(q, r), (r,)) 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 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 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 = create_element(terminal.ufl_element()) def callback(entity_id): finat_dict = ctx.basis_evaluation(element, mt.local_derivatives, entity_id) # Filter out irrelevant derivatives filtered_dict = {alpha: table for alpha, table in iteritems(finat_dict) 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 _merge_evaluations(self, factor_results): # Spatial dimension dimension = self.cell.get_spatial_dimension() # Derivative order order = max(map(sum, chain(*factor_results))) # A list of slices that are used to select dimensions # corresponding to each subelement. dim_slices = TensorProductCell._split_slices([c.get_spatial_dimension() for c in self.cell.cells]) # A list of multiindices, one multiindex per subelement, each # multiindex describing the shape of basis functions of the # subelement. alphas = [fe.get_indices() for fe in self.factors] # A list of multiindices, one multiindex per subelement, each # multiindex describing the value shape of the subelement. zetas = [fe.get_value_indices() for fe in self.factors] result = {} for derivative in range(order + 1): for Delta in mis(dimension, derivative): # Split the multiindex for the subelements deltas = [Delta[s] for s in dim_slices] # GEM scalars (can have free indices) for collecting # the contributions from the subelements. scalars = [] for fr, delta, alpha, zeta in zip(factor_results, deltas, alphas, zetas): # Turn basis shape to free indices, select the # right derivative entry, and collect the result. scalars.append(gem.Indexed(fr[delta], alpha + zeta)) # Multiply the values from the subelements and wrap up # non-point indices into shape. result[Delta] = gem.ComponentTensor( reduce(gem.Product, scalars), tuple(chain(*(alphas + zetas))) ) return result
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 _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, 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 sum(self, o, *ops): shape, = set(o.shape for o in ops) indices = gem.indices(len(shape)) return gem.ComponentTensor( gem.Sum(*[gem.Indexed(op, indices) for op in ops]), indices)
def promote(table): v = gem.partial_indexed(table, beta) u = gem.ListTensor(self.transform(v)) return gem.ComponentTensor(gem.Indexed(u, zeta), beta + zeta)
def abs(self, o, expr): indices = gem.indices(len(expr.shape)) return gem.ComponentTensor( gem.MathFunction('abs', gem.Indexed(expr, indices)), indices)
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)
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(), restriction=mt.restriction) # 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, 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): indices = tuple(i for i in var.index_ordering() if i not in ctx.unsummed_coefficient_indices) value = gem.IndexSum(gem.Product(expr, var), indices) 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) - ctx.unsummed_coefficient_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 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 _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 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 component_tensor(self, o, expression, index): index = tuple(i for i in index if i in expression.free_indices) return gem.ComponentTensor(expression, index)