def test_expressions(): x = gem.Variable("x", (3, 4)) y = gem.Variable("y", (4, )) i, j = gem.indices(2) xij = x[i, j] yj = y[j] assert xij == gem.Indexed(x, (i, j)) assert yj == gem.Indexed(y, (j, )) assert xij + yj == gem.Sum(xij, yj) assert xij * yj == gem.Product(xij, yj) assert xij - yj == gem.Sum(xij, gem.Product(gem.Literal(-1), yj)) assert xij / yj == gem.Division(xij, yj) assert xij + 1 == gem.Sum(xij, gem.Literal(1)) assert 1 + xij == gem.Sum(gem.Literal(1), xij) assert (xij + y).shape == (4, ) assert (x @ y).shape == (3, ) assert x.T.shape == (4, 3) with pytest.raises(ValueError): xij.T @ y with pytest.raises(ValueError): xij + "foo"
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 cell_orientation(self, restriction): """Cell orientation as a GEM expression.""" f = {None: 0, '+': 0, '-': 1}[restriction] # Assume self._cell_orientations tuple is set up at this point. co_int = self._cell_orientations[f] return gem.Conditional( gem.Comparison("==", co_int, gem.Literal(1)), gem.Literal(-1), gem.Conditional(gem.Comparison("==", co_int, gem.Zero()), gem.Literal(1), gem.Literal(numpy.nan)))
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 einsum(factors, sum_indices): """Evaluates a tensor product at compile time. :arg factors: iterable of indexed GEM literals :arg sum_indices: indices to sum over :returns: a single indexed GEM literal """ # Maps the problem onto numpy.einsum index2letter = defaultdict(partial(lambda c: chr(ord('i') + next(c)), count())) operands = [] subscript_parts = [] for factor in factors: literal, = factor.children selectors = [] letters = [] for index in factor.multiindex: if isinstance(index, int): selectors.append(index) else: selectors.append(slice(None)) letters.append(index2letter[index]) operands.append(literal.array.__getitem__(tuple(selectors))) subscript_parts.append(''.join(letters)) result_pairs = sorted((letter, index) for index, letter in index2letter.items() if index not in sum_indices) subscripts = ','.join(subscript_parts) + '->' + ''.join(l for l, i in result_pairs) tensor = numpy.einsum(subscripts, *operands) return gem.Indexed(gem.Literal(tensor), tuple(i for l, i in result_pairs))
def reference_normals(self): if not (isinstance(self.interface.fiat_cell, UFCSimplex) and self.interface.fiat_cell.get_spatial_dimension() == 2): raise NotImplementedError("Only works for triangles for now") return gem.Literal( numpy.asarray([ self.interface.fiat_cell.compute_normal(i) for i in range(3) ]))
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_edge_vectors(terminal, mt, ctx): from FIAT.reference_element import TensorProductCell as fiat_TensorProductCell fiat_cell = ctx.fiat_cell if isinstance(fiat_cell, fiat_TensorProductCell): raise NotImplementedError("CellEdgeVectors not implemented on TensorProductElements yet") nedges = len(fiat_cell.get_topology()[1]) vecs = numpy.vstack(map(fiat_cell.compute_edge_tangent, range(nedges))).astype(NUMPY_TYPE) assert vecs.shape == terminal.ufl_shape return gem.Literal(vecs)
def compile_to_gem(expr, translator): """Compile a single pointwise expression to GEM. :arg expr: The expression to compile. :arg translator: a :class:`Translator` instance. :returns: A (lvalue, rvalue) pair of preprocessed GEM.""" if not isinstance(expr, Assign): raise ValueError( f"Don't know how to assign expression of type {type(expr)}") spaces = tuple(c.function_space() for c in expr.coefficients) if any( type(s.ufl_element()) is ufl.MixedElement for s in spaces if s is not None): raise ValueError("Not expecting a mixed space at this point, " "did you forget to index a function with .sub(...)?") if len(set(s.ufl_element() for s in spaces if s is not None)) != 1: raise ValueError("All coefficients must be defined on the same space") lvalue = expr.lvalue rvalue = expr.rvalue broadcast = all( isinstance(c, firedrake.Constant) for c in expr.rcoefficients) and rvalue.ufl_shape == () if not broadcast and lvalue.ufl_shape != rvalue.ufl_shape: try: rvalue = reshape(rvalue, lvalue.ufl_shape) except ValueError: raise ValueError( "Mismatching shapes between lvalue and rvalue in pointwise assignment" ) rvalue, = map_expr_dags(LowerCompoundAlgebra(), [rvalue]) try: lvalue, rvalue = map_expr_dags(translator, [lvalue, rvalue]) except (AssertionError, ValueError): raise ValueError("Mismatching shapes in pointwise assignment. " "For intrinsically vector-/tensor-valued spaces make " "sure you're not using shaped Constants or literals.") indices = gem.indices(len(lvalue.shape)) if not broadcast: if rvalue.shape != lvalue.shape: raise ValueError( "Mismatching shapes in pointwise assignment. " "For intrinsically vector-/tensor-valued spaces make " "sure you're not using shaped Constants or literals.") rvalue = gem.Indexed(rvalue, indices) lvalue = gem.Indexed(lvalue, indices) if isinstance(expr, IAdd): rvalue = gem.Sum(lvalue, rvalue) elif isinstance(expr, ISub): rvalue = gem.Sum(lvalue, gem.Product(gem.Literal(-1), rvalue)) elif isinstance(expr, IMul): rvalue = gem.Product(lvalue, rvalue) elif isinstance(expr, IDiv): rvalue = gem.Division(lvalue, rvalue) return preprocess_gem([lvalue, rvalue])
def select_hdiv_transformer(element): # Assume: something x interval assert len(element.factors) == 2 assert element.factors[1].cell.get_shape() == LINE # Globally consistent edge orientations of the reference # quadrilateral: rightward horizontally, upward vertically. # Their rotation by 90 degrees anticlockwise is interpreted as the # positive direction for normal vectors. ks = tuple(fe.formdegree for fe in element.factors) if ks == (0, 1): # Make the scalar value the leftward-pointing normal on the # y-aligned edges. return lambda v: [gem.Product(gem.Literal(-1), v), gem.Zero()] elif ks == (1, 0): # Make the scalar value the upward-pointing normal on the # x-aligned edges. return lambda v: [gem.Zero(), v] elif ks == (2, 0): # Same for 3D, so z-plane. return lambda v: [gem.Zero(), gem.Zero(), v] elif ks == (1, 1): if element.mapping == "contravariant piola": # Pad the 2-vector normal on the "base" cell into a # 3-vector, maintaining direction. return lambda v: [ gem.Indexed(v, (0, )), gem.Indexed(v, (1, )), gem.Zero() ] elif element.mapping == "covariant piola": # Rotate the 2-vector tangential component on the "base" # cell 90 degrees anticlockwise into a 3-vector and pad. return lambda v: [ gem.Indexed(v, (1, )), gem.Product(gem.Literal(-1), gem.Indexed(v, (0, ))), gem.Zero() ] else: assert False, "Unexpected original mapping!" else: assert False, "Unexpected form degree combination!"
def physical_points(self, ps, entity=None): assert entity is None prefs = ps.points pvs = self.verts pps = np.zeros(prefs.shape, dtype=float) for i in range(pps.shape[0]): pps[i, :] = (pvs[0, :] * (1 - prefs[i, 0]) * (1 - prefs[i, 1]) + pvs[1, :] * (1 - prefs[i, 0]) * prefs[i, 1] + pvs[2, :] * prefs[i, 0] * (1 - prefs[i, 1]) + pvs[3, :] * prefs[i, 0] * prefs[i, 1]) return gem.Literal(pps)
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 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 select_hcurl_transformer(element): # Assume: something x interval assert len(element.factors) == 2 assert element.factors[1].cell.get_shape() == LINE # Globally consistent edge orientations of the reference # quadrilateral: rightward horizontally, upward vertically. # Tangential vectors interpret these as the positive direction. dim = element.cell.get_spatial_dimension() ks = tuple(fe.formdegree for fe in element.factors) if element.mapping == "affine": if ks == (1, 0): # Can only be 2D. Make the scalar value the # rightward-pointing tangential on the x-aligned edges. return lambda v: [v, gem.Zero()] elif ks == (0, 1): # Can be any spatial dimension. Make the scalar value the # upward-pointing tangential. return lambda v: [gem.Zero()] * (dim - 1) + [v] else: assert False elif element.mapping == "covariant piola": # Second factor must be continuous interval. Just padding. return lambda v: [ gem.Indexed(v, (0, )), gem.Indexed(v, (1, )), gem.Zero() ] elif element.mapping == "contravariant piola": # Second factor must be continuous interval. Rotate the # 2-vector tangential component on the "base" cell 90 degrees # clockwise into a 3-vector and pad. return lambda v: [ gem.Product(gem.Literal(-1), gem.Indexed(v, (1, ))), gem.Indexed(v, (0, )), gem.Zero() ] else: assert False, "Unexpected original mapping!"
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 callback(entity_id): t = ctx.fiat_cell.get_entity_transform(ctx.integration_dim, entity_id) data = numpy.asarray(list(map(t, ps.points))) return gem.Literal(data.reshape(point_shape + data.shape[1:]))
def callback(facet_i): n = ctx.fiat_cell.compute_reference_normal(ctx.integration_dim, facet_i) return gem.Literal(n)
def callback(entity_id): return gem.Literal(make_cell_facet_jacobian(cell, facet_dim, entity_id))
def expression(self): return gem.partial_indexed(gem.Literal(self.points), self.indices)
def physical_vertices(self): return gem.Literal(self.phys_cell.verts)
def expression(self): return gem.Literal(self.point)
def translate_reference_facet_volume(terminal, mt, ctx): # FIXME: simplex only code path dim = ctx.fiat_cell.get_spatial_dimension() facet_cell = ctx.fiat_cell.construct_subelement(dim - 1) return gem.Literal(facet_cell.volume())
def physical_points(self, ps, entity=None): prefs = ps.points A, b = self.A, self.b return gem.Literal(np.asarray([A @ x + b for x in prefs]))
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 reference_normals(self): return gem.Literal( numpy.asarray([ self.interface.fiat_cell.compute_normal(i) for i in range(3) ]))
def translate_reference_cell_volume(terminal, mt, ctx): return gem.Literal(ctx.fiat_cell.volume())
def weight_expression(self): return gem.Indexed(gem.Literal(self.weights), self.point_set.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 test_as_gem(): with pytest.raises(ValueError): gem.as_gem([1, 2]) assert gem.as_gem(1) == gem.Literal(1)