def prepare_coordinates(coefficient, name, interior_facet=False): """Bridges the kernel interface and the GEM abstraction for coordinates. :arg coefficient: UFL Coefficient :arg name: unique name to refer to the Coefficient in the kernel :arg interior_facet: interior facet integral? :returns: (funarg, expression) funarg - :class:`coffee.Decl` function argument expression - GEM expression referring to the Coefficient values """ finat_element = create_element(coefficient.ufl_element()) shape = finat_element.index_shape size = numpy.prod(shape, dtype=int) if not interior_facet: funargs = [coffee.Decl(SCALAR_TYPE, coffee.Symbol(name), pointers=[("",)], qualifiers=["const"])] variable = gem.Variable(name, (size,)) expression = gem.reshape(variable, shape) else: funargs = [coffee.Decl(SCALAR_TYPE, coffee.Symbol(name+"_0"), pointers=[("",)], qualifiers=["const"]), coffee.Decl(SCALAR_TYPE, coffee.Symbol(name+"_1"), pointers=[("",)], qualifiers=["const"])] variable0 = gem.Variable(name+"_0", (size,)) variable1 = gem.Variable(name+"_1", (size,)) expression = (gem.reshape(variable0, shape), gem.reshape(variable1, shape)) return funargs, expression
def __init__(self, integral_type, subdomain_id, domain_number): """Initialise a kernel builder.""" super(KernelBuilder, self).__init__(integral_type.startswith("interior_facet")) self.kernel = Kernel(integral_type=integral_type, subdomain_id=subdomain_id, domain_number=domain_number) self.local_tensor = None self.coordinates_arg = None self.coefficient_args = [] self.coefficient_split = {} # Facet number if integral_type in ['exterior_facet', 'exterior_facet_vert']: facet = gem.Variable('facet', (1, )) self._facet_number = { None: gem.VariableIndex(gem.Indexed(facet, (0, ))) } elif integral_type in ['interior_facet', 'interior_facet_vert']: facet = gem.Variable('facet', (2, )) self._facet_number = { '+': gem.VariableIndex(gem.Indexed(facet, (0, ))), '-': gem.VariableIndex(gem.Indexed(facet, (1, ))) } elif integral_type == 'interior_facet_horiz': self._facet_number = {'+': 1, '-': 0}
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 prepare_arguments(arguments, multiindices, scalar_type, interior_facet=False, diagonal=False): """Bridges the kernel interface and the GEM abstraction for Arguments. Vector Arguments are rearranged here for interior facet integrals. :arg arguments: UFL Arguments :arg multiindices: Argument multiindices :arg interior_facet: interior facet integral? :arg diagonal: Are we assembling the diagonal of a rank-2 element tensor? :returns: (funarg, expression) funarg - :class:`loopy.GlobalArg` function argument expressions - GEM expressions referring to the argument tensor """ assert isinstance(interior_facet, bool) if len(arguments) == 0: # No arguments funarg = lp.GlobalArg("A", dtype=scalar_type, shape=(1,)) expression = gem.Indexed(gem.Variable("A", (1,)), (0,)) return funarg, [expression] elements = tuple(create_element(arg.ufl_element()) for arg in arguments) shapes = tuple(element.index_shape for element in elements) if diagonal: if len(arguments) != 2: raise ValueError("Diagonal only for 2-forms") try: element, = set(elements) except ValueError: raise ValueError("Diagonal only for diagonal blocks (test and trial spaces the same)") elements = (element, ) shapes = tuple(element.index_shape for element in elements) multiindices = multiindices[:1] def expression(restricted): return gem.Indexed(gem.reshape(restricted, *shapes), tuple(chain(*multiindices))) u_shape = numpy.array([numpy.prod(shape, dtype=int) for shape in shapes]) if interior_facet: c_shape = tuple(2 * u_shape) slicez = [[slice(r * s, (r + 1) * s) for r, s in zip(restrictions, u_shape)] for restrictions in product((0, 1), repeat=len(arguments))] else: c_shape = tuple(u_shape) slicez = [[slice(s) for s in u_shape]] funarg = lp.GlobalArg("A", dtype=scalar_type, shape=c_shape) varexp = gem.Variable("A", c_shape) expressions = [expression(gem.view(varexp, *slices)) for slices in slicez] return funarg, prune(expressions)
def prepare_coefficient(coefficient, name, interior_facet=False): """Bridges the kernel interface and the GEM abstraction for Coefficients. :arg coefficient: UFL Coefficient :arg name: unique name to refer to the Coefficient in the kernel :arg interior_facet: interior facet integral? :returns: (funarg, expression) funarg - :class:`coffee.Decl` function argument expression - GEM expression referring to the Coefficient values """ assert isinstance(interior_facet, bool) if coefficient.ufl_element().family() == 'Real': # Constant funarg = coffee.Decl(SCALAR_TYPE, coffee.Symbol(name), pointers=[("restrict", )], qualifiers=["const"]) expression = gem.reshape(gem.Variable(name, (None, )), coefficient.ufl_shape) return funarg, expression finat_element = create_element(coefficient.ufl_element()) if isinstance(finat_element, TensorFiniteElement): scalar_shape = finat_element.base_element.index_shape tensor_shape = finat_element.index_shape[len(scalar_shape):] else: scalar_shape = finat_element.index_shape tensor_shape = () scalar_size = numpy.prod(scalar_shape, dtype=int) tensor_size = numpy.prod(tensor_shape, dtype=int) funarg = coffee.Decl(SCALAR_TYPE, coffee.Symbol(name), pointers=[("const", "restrict"), ("restrict", )], qualifiers=["const"]) if not interior_facet: expression = gem.reshape( gem.Variable(name, (scalar_size, tensor_size)), scalar_shape, tensor_shape) else: varexp = gem.Variable(name, (2 * scalar_size, tensor_size)) plus = gem.view(varexp, slice(scalar_size), slice(tensor_size)) minus = gem.view(varexp, slice(scalar_size, 2 * scalar_size), slice(tensor_size)) expression = (gem.reshape(plus, scalar_shape, tensor_shape), gem.reshape(minus, scalar_shape, tensor_shape)) return funarg, expression
def to_reference_coordinates(ufl_coordinate_element): # Set up UFL form cell = ufl_coordinate_element.cell() domain = ufl.Mesh(ufl_coordinate_element) K = ufl.JacobianInverse(domain) x = ufl.SpatialCoordinate(domain) x0_element = ufl.VectorElement("Real", cell, 0) x0 = ufl.Coefficient(ufl.FunctionSpace(domain, x0_element)) expr = ufl.dot(K, x - x0) # Translation to GEM C = ufl_utils.coordinate_coefficient(domain) expr = ufl_utils.preprocess_expression(expr) expr = ufl_utils.replace_coordinates(expr, C) expr = ufl_utils.simplify_abs(expr) builder = firedrake_interface.KernelBuilderBase() builder._coefficient(C, "C") builder._coefficient(x0, "x0") dim = cell.topological_dimension() point = gem.Variable('X', (dim, )) context = tsfc.fem.GemPointContext( interface=builder, ufl_cell=cell, precision=parameters["precision"], point_indices=(), point_expr=point, ) translator = tsfc.fem.Translator(context) ir = map_expr_dag(translator, expr) # Unroll result ir = [gem.Indexed(ir, alpha) for alpha in numpy.ndindex(ir.shape)] # Unroll IndexSums max_extent = parameters["unroll_indexsum"] if max_extent: def predicate(index): return index.extent <= max_extent ir = gem.optimise.unroll_indexsum(ir, predicate=predicate) # Translate to COFFEE ir = impero_utils.preprocess_gem(ir) return_variable = gem.Variable('dX', (dim, )) assignments = [(gem.Indexed(return_variable, (i, )), e) for i, e in enumerate(ir)] impero_c = impero_utils.compile_gem(assignments, ()) body = tsfc.coffee.generate(impero_c, {}, parameters["precision"]) body.open_scope = False return body
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 prepare_coordinates(coefficient, name, scalar_type, interior_facet=False): """Bridges the kernel interface and the GEM abstraction for coordinates. :arg coefficient: UFL Coefficient :arg name: unique name to refer to the Coefficient in the kernel :arg interior_facet: interior facet integral? :returns: (funarg, expression) funarg - :class:`coffee.Decl` function argument expression - GEM expression referring to the Coefficient values """ finat_element = create_element(coefficient.ufl_element()) shape = finat_element.index_shape size = numpy.prod(shape, dtype=int) assert isinstance(finat_element, TensorFiniteElement) scalar_shape = finat_element.base_element.index_shape tensor_shape = finat_element._shape transposed_shape = scalar_shape + tensor_shape scalar_rank = len(scalar_shape) 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) if not interior_facet: funargs = [ coffee.Decl(scalar_type, coffee.Symbol(name), pointers=[("", )], qualifiers=["const"]) ] variable = gem.Variable(name, (size, )) expression = transpose(gem.reshape(variable, transposed_shape)) else: funargs = [ coffee.Decl(scalar_type, coffee.Symbol(name + "_0"), pointers=[("", )], qualifiers=["const"]), coffee.Decl(scalar_type, coffee.Symbol(name + "_1"), pointers=[("", )], qualifiers=["const"]) ] variable0 = gem.Variable(name + "_0", (size, )) variable1 = gem.Variable(name + "_1", (size, )) expression = (transpose(gem.reshape(variable0, transposed_shape)), transpose(gem.reshape(variable1, transposed_shape))) return funargs, expression
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 __init__(self, interior_facet=False): """Initialise a kernel builder. :arg interior_facet: kernel accesses two cells """ super(KernelBuilderBase, self).__init__(interior_facet=interior_facet) # Cell orientation if self.interior_facet: cell_orientations = gem.Variable("cell_orientations", (2, )) self._cell_orientations = (gem.Indexed(cell_orientations, (0, )), gem.Indexed(cell_orientations, (1, ))) else: cell_orientations = gem.Variable("cell_orientations", (1, )) self._cell_orientations = (gem.Indexed(cell_orientations, (0, )), )
def prepare_coefficient(coefficient, num, name, interior_facet=False): """Bridges the kernel interface and the GEM abstraction for Coefficients. :arg coefficient: UFL Coefficient :arg num: coefficient index in the original form :arg name: unique name to refer to the Coefficient in the kernel :arg interior_facet: interior facet integral? :returns: GEM expression referring to the Coefficient value """ varexp = gem.Variable(name, (None, None)) if coefficient.ufl_element().family() == 'Real': size = numpy.prod(coefficient.ufl_shape, dtype=int) data = gem.view(varexp, slice(num, num + 1), slice(size)) return gem.reshape(data, (), coefficient.ufl_shape) element = create_element(coefficient.ufl_element()) size = numpy.prod(element.index_shape, dtype=int) def expression(data): result, = prune( [gem.reshape(gem.view(data, slice(size)), element.index_shape)]) return result if not interior_facet: data = gem.view(varexp, slice(num, num + 1), slice(size)) return expression(gem.reshape(data, (), (size, ))) else: data_p = gem.view(varexp, slice(num, num + 1), slice(size)) data_m = gem.view(varexp, slice(num, num + 1), slice(size, 2 * size)) return (expression(gem.reshape(data_p, (), (size, ))), expression(gem.reshape(data_m, (), (size, ))))
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 object. :param entity: the cell entity on which to tabulate. """ # Spatial dimension dimension = self.cell.get_spatial_dimension() # Shape of the tabulation matrix shape = tuple( index.extent for index in ps.indices) + self.index_shape + self.value_shape result = {} for derivative in range(order + 1): for alpha in mis(dimension, derivative): name = str.format("rt_{}_{}_{}_{}_{}_{}", self.variant, self.degree, ''.join(map(str, alpha)), self.shift_axes, 'c' if self.continuous else 'd', { None: "", '+': "p", '-': "m" }[self.restriction]) result[alpha] = gem.partial_indexed(gem.Variable(name, shape), ps.indices) return result
def gem_to_loopy(gem_expr, var2terminal, scalar_type): """ Method encapsulating stage 2. Converts the gem expression dag into imperoc first, and then further into loopy. :return slate_loopy: 2-tuple of loopy kernel for slate operations and loopy GlobalArg for the output variable. """ # Creation of return variables for outer loopy shape = gem_expr.shape if len(gem_expr.shape) != 0 else (1, ) idx = make_indices(len(shape)) indexed_gem_expr = gem.Indexed(gem_expr, idx) args = ([ loopy.GlobalArg("output", shape=shape, dtype=scalar_type, is_output=True, is_input=True) ] + [ loopy.GlobalArg(var.name, shape=var.shape, dtype=scalar_type) for var in var2terminal.keys() ]) ret_vars = [gem.Indexed(gem.Variable("output", shape), idx)] preprocessed_gem_expr = impero_utils.preprocess_gem([indexed_gem_expr]) # glue assignments to return variable assignments = list(zip(ret_vars, preprocessed_gem_expr)) # Part A: slate to impero_c impero_c = impero_utils.compile_gem(assignments, (), remove_zeros=False) # Part B: impero_c to loopy return generate_loopy(impero_c, args, scalar_type, "slate_loopy", []), args[0].copy()
def prepare_arguments(arguments, multiindices, scalar_type, interior_facet=False): """Bridges the kernel interface and the GEM abstraction for Arguments. Vector Arguments are rearranged here for interior facet integrals. :arg arguments: UFL Arguments :arg multiindices: Argument multiindices :arg interior_facet: interior facet integral? :returns: (funarg, expression) funarg - :class:`coffee.Decl` function argument expressions - GEM expressions referring to the argument tensor """ assert isinstance(interior_facet, bool) if len(arguments) == 0: # No arguments funarg = coffee.Decl(scalar_type, coffee.Symbol("A", rank=(1, ))) expression = gem.Indexed(gem.Variable("A", (1, )), (0, )) return funarg, [expression] elements = tuple(create_element(arg.ufl_element()) for arg in arguments) shapes = tuple(element.index_shape for element in elements) def expression(restricted): return gem.Indexed(gem.reshape(restricted, *shapes), tuple(chain(*multiindices))) u_shape = numpy.array([numpy.prod(shape, dtype=int) for shape in shapes]) if interior_facet: c_shape = tuple(2 * u_shape) slicez = [[ slice(r * s, (r + 1) * s) for r, s in zip(restrictions, u_shape) ] for restrictions in product((0, 1), repeat=len(arguments))] else: c_shape = tuple(u_shape) slicez = [[slice(s) for s in u_shape]] funarg = coffee.Decl(scalar_type, coffee.Symbol("A", rank=c_shape)) varexp = gem.Variable("A", c_shape) expressions = [expression(gem.view(varexp, *slices)) for slices in slicez] return funarg, prune(expressions)
def __init__(self, integral_type, subdomain_id, domain_number): """Initialise a kernel builder.""" super(KernelBuilder, self).__init__(integral_type.startswith("interior_facet")) self.integral_type = integral_type self.local_tensor = None self.coordinates_args = None self.coefficient_args = None self.coefficient_split = None if self.interior_facet: self._cell_orientations = (gem.Variable("cell_orientation_0", ()), gem.Variable("cell_orientation_1", ())) else: self._cell_orientations = (gem.Variable("cell_orientation", ()), ) if integral_type == "exterior_facet": self._entity_number = { None: gem.VariableIndex(gem.Variable("facet", ())) } elif integral_type == "interior_facet": self._entity_number = { '+': gem.VariableIndex(gem.Variable("facet_0", ())), '-': gem.VariableIndex(gem.Variable("facet_1", ())) } elif integral_type == "vertex": self._entity_number = { None: gem.VariableIndex(gem.Variable("vertex", ())) }
def __init__(self, scalar_type=None, interior_facet=False): """Initialise a kernel builder. :arg interior_facet: kernel accesses two cells """ if scalar_type is None: from tsfc.parameters import SCALAR_TYPE scalar_type = SCALAR_TYPE super(KernelBuilderBase, self).__init__(scalar_type=scalar_type, interior_facet=interior_facet) # Cell orientation if self.interior_facet: cell_orientations = gem.Variable("cell_orientations", (2, )) self._cell_orientations = (gem.Indexed(cell_orientations, (0, )), gem.Indexed(cell_orientations, (1, ))) else: cell_orientations = gem.Variable("cell_orientations", (1, )) self._cell_orientations = (gem.Indexed(cell_orientations, (0, )), )
def _coefficient(self, coefficient, name): element = create_element(coefficient.ufl_element()) shape = self.shape + element.index_shape size = numpy.prod(shape, dtype=int) funarg = ast.Decl(SCALAR_TYPE, ast.Symbol(name), pointers=[("restrict", )], qualifiers=["const"]) expression = gem.reshape(gem.Variable(name, (size, )), shape) expression = gem.partial_indexed(expression, self.indices) self.coefficient_map[coefficient] = expression return funarg
def __init__(self, scalar_type, interior_facet=False): """Initialise a kernel builder. :arg interior_facet: kernel accesses two cells """ super(KernelBuilderBase, self).__init__(scalar_type=scalar_type, interior_facet=interior_facet) # Cell orientation if self.interior_facet: shape = (2,) cell_orientations = gem.Variable("cell_orientations", shape) self._cell_orientations = (gem.Indexed(cell_orientations, (0,)), gem.Indexed(cell_orientations, (1,))) else: shape = (1,) cell_orientations = gem.Variable("cell_orientations", shape) self._cell_orientations = (gem.Indexed(cell_orientations, (0,)),) self.cell_orientations_loopy_arg = lp.GlobalArg("cell_orientations", dtype=numpy.int32, shape=shape)
def prepare_arguments(arguments, multiindices, scalar_type, interior_facet=False): """Bridges the kernel interface and the GEM abstraction for Arguments. Vector Arguments are rearranged here for interior facet integrals. :arg arguments: UFL Arguments :arg multiindices: Argument multiindices :arg interior_facet: interior facet integral? :returns: (funarg, prepare, expressions) funarg - :class:`coffee.Decl` function argument prepare - list of COFFEE nodes to be prepended to the kernel body expressions - GEM expressions referring to the argument tensor """ funarg = coffee.Decl(scalar_type, coffee.Symbol("A"), pointers=[()]) varexp = gem.Variable("A", (None, )) if len(arguments) == 0: # No arguments zero = coffee.FlatBlock("memset({name}, 0, sizeof(*{name}));\n".format( name=funarg.sym.gencode())) return funarg, [zero], [gem.reshape(varexp, ())] elements = tuple(create_element(arg.ufl_element()) for arg in arguments) shapes = [element.index_shape for element in elements] indices = tuple(chain(*multiindices)) def expression(restricted): return gem.Indexed(gem.reshape(restricted, *shapes), indices) u_shape = numpy.array( [numpy.prod(element.index_shape, dtype=int) for element in elements]) if interior_facet: c_shape = tuple(2 * u_shape) slicez = [[ slice(r * s, (r + 1) * s) for r, s in zip(restrictions, u_shape) ] for restrictions in product((0, 1), repeat=len(arguments))] else: c_shape = tuple(u_shape) slicez = [[slice(s) for s in u_shape]] expressions = [ expression(gem.view(gem.reshape(varexp, c_shape), *slices)) for slices in slicez ] zero = coffee.FlatBlock( str.format("memset({name}, 0, {size} * sizeof(*{name}));\n", name=funarg.sym.gencode(), size=numpy.product(c_shape, dtype=int))) return funarg, [zero], prune(expressions)
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 prepare_coefficient(coefficient, name, scalar_type, interior_facet=False): """Bridges the kernel interface and the GEM abstraction for Coefficients. :arg coefficient: UFL Coefficient :arg name: unique name to refer to the Coefficient in the kernel :arg interior_facet: interior facet integral? :returns: (funarg, expression) funarg - :class:`loopy.GlobalArg` function argument expression - GEM expression referring to the Coefficient values """ assert isinstance(interior_facet, bool) if coefficient.ufl_element().family() == 'Real': # Constant funarg = lp.GlobalArg(name, dtype=scalar_type, shape=(coefficient.ufl_element().value_size(), )) expression = gem.reshape(gem.Variable(name, (None, )), coefficient.ufl_shape) return funarg, expression finat_element = create_element(coefficient.ufl_element()) shape = finat_element.index_shape size = numpy.prod(shape, dtype=int) if not interior_facet: expression = gem.reshape(gem.Variable(name, (size, )), shape) else: varexp = gem.Variable(name, (2 * size, )) plus = gem.view(varexp, slice(size)) minus = gem.view(varexp, slice(size, 2 * size)) expression = (gem.reshape(plus, shape), gem.reshape(minus, shape)) size = size * 2 funarg = lp.GlobalArg(name, dtype=scalar_type, shape=(size, )) return funarg, expression
def coefficient(self, o): # Because we act on dofs, the ufl_shape is not the right thing to check shape = o.dat.dim try: var = self.varmapping[o] except KeyError: name = f"C{len(self.varmapping)}" var = gem.Variable(name, shape) self.varmapping[o] = var if o.ufl_shape == (): assert shape == (1, ) return gem.Indexed(var, (0, )) else: return var
def gem_to_loopy(gem_expr): """ Method encapsulating stage 2. Converts the gem expression dag into imperoc first, and then further into loopy. :return slate_loopy: loopy kernel for slate operations. """ # Creation of return variables for outer loopy shape = gem_expr.shape if len(gem_expr.shape) != 0 else (1,) idx = make_indices(len(shape)) indexed_gem_expr = gem.Indexed(gem_expr, idx) arg = [loopy.GlobalArg("output", shape=shape)] ret_vars = [gem.Indexed(gem.Variable("output", shape), idx)] preprocessed_gem_expr = impero_utils.preprocess_gem([indexed_gem_expr]) # glue assignments to return variable assignments = list(zip(ret_vars, preprocessed_gem_expr)) # Part A: slate to impero_c impero_c = impero_utils.compile_gem(assignments, (), remove_zeros=False) # Part B: impero_c to loopy return generate_loopy(impero_c, arg, parameters["form_compiler"]["scalar_type"], "slate_loopy", [])
def rebuild_dg(element, expr, rt_var_name): # To tabulate on the given element (which is on a different mesh to the # expression) we must do so at runtime. We therefore create a quadrature # element with runtime points to evaluate for each point in the element's # dual basis. This exists on the same reference cell as the input element # and we can interpolate onto it before mapping the result back onto the # target space. expr_tdim = expr.ufl_domain().topological_dimension() # Need point evaluations and matching weights from dual basis. # This could use FIAT's dual basis as below: # num_points = sum(len(dual.get_point_dict()) for dual in element.fiat_equivalent.dual_basis()) # weights = [] # for dual in element.fiat_equivalent.dual_basis(): # pts = dual.get_point_dict().keys() # for p in pts: # for w, _ in dual.get_point_dict()[p]: # weights.append(w) # assert len(weights) == num_points # but for now we just fix the values to what we know works: if element.degree != 0 or not isinstance(element.cell, FIAT.reference_element.Point): raise NotImplementedError( "Cross mesh interpolation only implemented for P0DG on vertex cells." ) num_points = 1 weights = [1.] * num_points # gem.Variable name starting with rt_ forces TSFC runtime tabulation assert rt_var_name.startswith("rt_") runtime_points_expr = gem.Variable(rt_var_name, (num_points, expr_tdim)) rule_pointset = finat.point_set.UnknownPointSet(runtime_points_expr) try: expr_fiat_cell = as_fiat_cell(expr.ufl_element().cell()) except AttributeError: # expression must be pure function of spatial coordinates so # domain has correct ufl cell expr_fiat_cell = as_fiat_cell(expr.ufl_domain().ufl_cell()) rule = finat.quadrature.QuadratureRule(rule_pointset, weights=weights) return finat.QuadratureElement(expr_fiat_cell, rule)
def __init__(self, integral_type, subdomain_id, domain_number, scalar_type=None, diagonal=False): """Initialise a kernel builder.""" if diagonal: raise NotImplementedError( "Assembly of diagonal not implemented yet, sorry") super(KernelBuilder, self).__init__(scalar_type, integral_type.startswith("interior_facet")) self.integral_type = integral_type self.local_tensor = None self.coordinates_args = None self.coefficient_args = None self.coefficient_split = None if self.interior_facet: self._cell_orientations = (gem.Variable("cell_orientation_0", ()), gem.Variable("cell_orientation_1", ())) else: self._cell_orientations = (gem.Variable("cell_orientation", ()), ) if integral_type == "exterior_facet": self._entity_number = { None: gem.VariableIndex(gem.Variable("facet", ())) } elif integral_type == "interior_facet": self._entity_number = { '+': gem.VariableIndex(gem.Variable("facet_0", ())), '-': gem.VariableIndex(gem.Variable("facet_1", ())) } elif integral_type == "vertex": self._entity_number = { None: gem.VariableIndex(gem.Variable("vertex", ())) }
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_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 prepare_coefficients(coefficients, num, name, interior_facet=False): """Bridges the kernel interface and the GEM abstraction for Coefficients. :arg coefficients: split UFL Coefficients :arg num: coefficient index in the original form :arg name: unique name to refer to the Coefficient in the kernel :arg interior_facet: interior facet integral? :returns: GEM expression referring to the Coefficient value """ varexp = gem.Variable(name, (None, None)) if len(coefficients) == 1 and coefficients[0].ufl_element().family() == 'Real': coefficient, = coefficients size = numpy.prod(coefficient.ufl_shape, dtype=int) data = gem.view(varexp, slice(num, num + 1), slice(size)) expression = gem.reshape(data, (), coefficient.ufl_shape) return [expression] elements = [create_element(coeff.ufl_element()) for coeff in coefficients] space_dimensions = [numpy.prod(element.index_shape, dtype=int) for element in elements] ends = list(numpy.cumsum(space_dimensions)) starts = [0] + ends[:-1] slices = [slice(start, end) for start, end in zip(starts, ends)] transposed_shapes = [] tensor_ranks = [] for element in elements: if isinstance(element, TensorFiniteElement): scalar_shape = element.base_element.index_shape tensor_shape = element.index_shape[len(scalar_shape):] else: scalar_shape = element.index_shape tensor_shape = () transposed_shapes.append(tensor_shape + scalar_shape) tensor_ranks.append(len(tensor_shape)) 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 expressions(data): return prune([transpose(gem.reshape(gem.view(data, slice_), shape), rank) for slice_, shape, rank in zip(slices, transposed_shapes, tensor_ranks)]) size = sum(space_dimensions) if not interior_facet: data = gem.view(varexp, slice(num, num + 1), slice(size)) return expressions(gem.reshape(data, (), (size,))) else: data_p = gem.view(varexp, slice(num, num + 1), slice(size)) data_m = gem.view(varexp, slice(num, num + 1), slice(size, 2 * size)) return list(zip(expressions(gem.reshape(data_p, (), (size,))), expressions(gem.reshape(data_m, (), (size,)))))
def vector(): return gem.Variable('u', (12, ))
def matrix(): return gem.Variable('A', (10, 12))