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 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 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_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 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 compile_integral(integral_data, form_data, prefix, parameters, interface=firedrake_interface): """Compiles a UFL integral into an assembly kernel. :arg integral_data: UFL integral data :arg form_data: UFL form data :arg prefix: kernel name will start with this string :arg parameters: parameters object :arg interface: backend module for the kernel interface :returns: a kernel constructed by the kernel interface """ if parameters is None: parameters = default_parameters() else: _ = default_parameters() _.update(parameters) parameters = _ # Remove these here, they're handled below. if parameters.get("quadrature_degree") in ["auto", "default", None, -1, "-1"]: del parameters["quadrature_degree"] if parameters.get("quadrature_rule") in ["auto", "default", None]: del parameters["quadrature_rule"] integral_type = integral_data.integral_type interior_facet = integral_type.startswith("interior_facet") mesh = integral_data.domain cell = integral_data.domain.ufl_cell() arguments = form_data.preprocessed_form.arguments() fiat_cell = as_fiat_cell(cell) integration_dim, entity_ids = lower_integral_type(fiat_cell, integral_type) argument_indices = tuple(tuple(gem.Index(extent=e) for e in create_element(arg.ufl_element()).index_shape) for arg in arguments) flat_argument_indices = tuple(chain(*argument_indices)) quadrature_indices = [] # Dict mapping domains to index in original_form.ufl_domains() domain_numbering = form_data.original_form.domain_numbering() builder = interface.KernelBuilder(integral_type, integral_data.subdomain_id, domain_numbering[integral_data.domain]) return_variables = builder.set_arguments(arguments, argument_indices) coordinates = ufl_utils.coordinate_coefficient(mesh) builder.set_coordinates(coordinates) builder.set_coefficients(integral_data, form_data) # Map from UFL FiniteElement objects to Index instances. This is # so we reuse Index instances when evaluating the same coefficient # multiple times with the same table. Occurs, for example, if we # have multiple integrals here (and the affine coordinate # evaluation can be hoisted). index_cache = collections.defaultdict(gem.Index) kernel_cfg = dict(interface=builder, ufl_cell=cell, precision=parameters["precision"], integration_dim=integration_dim, entity_ids=entity_ids, argument_indices=argument_indices, index_cache=index_cache) kernel_cfg["facetarea"] = facetarea_generator(mesh, coordinates, kernel_cfg, integral_type) kernel_cfg["cellvolume"] = cellvolume_generator(mesh, coordinates, kernel_cfg) irs = [] for integral in integral_data.integrals: params = {} # Record per-integral parameters params.update(integral.metadata()) if params.get("quadrature_rule") == "default": del params["quadrature_rule"] # parameters override per-integral metadata params.update(parameters) integrand = ufl_utils.replace_coordinates(integral.integrand(), coordinates) integrand = ufl_utils.split_coefficients(integrand, builder.coefficient_split) # Check if the integral has a quad degree attached, otherwise use # the estimated polynomial degree attached by compute_form_data quadrature_degree = params.get("quadrature_degree", params["estimated_polynomial_degree"]) try: quad_rule = params["quadrature_rule"] except KeyError: integration_cell = fiat_cell.construct_subelement(integration_dim) quad_rule = make_quadrature(integration_cell, quadrature_degree) if not isinstance(quad_rule, AbstractQuadratureRule): raise ValueError("Expected to find a QuadratureRule object, not a %s" % type(quad_rule)) quadrature_multiindex = quad_rule.point_set.indices quadrature_indices += quadrature_multiindex config = kernel_cfg.copy() config.update(quadrature_rule=quad_rule) ir = fem.compile_ufl(integrand, interior_facet=interior_facet, **config) if parameters["unroll_indexsum"]: def predicate(index): return index.extent <= parameters["unroll_indexsum"] ir = opt.unroll_indexsum(ir, predicate=predicate) ir = [gem.index_sum(expr, quadrature_multiindex) for expr in ir] irs.append(ir) # Sum the expressions that are part of the same restriction ir = list(reduce(gem.Sum, e, gem.Zero()) for e in zip(*irs)) # Need optimised roots for COFFEE ir = impero_utils.preprocess_gem(ir) # Look for cell orientations in the IR if builder.needs_cell_orientations(ir): builder.require_cell_orientations() impero_c = impero_utils.compile_gem(return_variables, ir, tuple(quadrature_indices) + flat_argument_indices, remove_zeros=True) # Generate COFFEE index_names = [(si, name + str(n)) for index, name in zip(argument_indices, ['j', 'k']) for n, si in enumerate(index)] if len(quadrature_indices) == 1: index_names.append((quadrature_indices[0], 'ip')) else: for i, quadrature_index in enumerate(quadrature_indices): index_names.append((quadrature_index, 'ip_%d' % i)) body = generate_coffee(impero_c, index_names, parameters["precision"], ir, flat_argument_indices) kernel_name = "%s_%s_integral_%s" % (prefix, integral_type, integral_data.subdomain_id) return builder.construct_kernel(kernel_name, body)
def _transform(self, v): u = [gem.Zero()] * self.size for j, zeta in enumerate(numpy.ndindex(self.element.value_shape)): u[self.offset + j] = gem.Indexed(v, zeta) return u
def compile_expression_dual_evaluation(expression, to_element, *, domain=None, interface=None, parameters=None, coffee=False): """Compile a UFL expression to be evaluated against a compile-time known reference element's dual basis. Useful for interpolating UFL expressions into e.g. N1curl spaces. :arg expression: UFL expression :arg to_element: A FInAT element for the target space :arg domain: optional UFL domain the expression is defined on (required when expression contains no domain). :arg interface: backend module for the kernel interface :arg parameters: parameters object :arg coffee: compile coffee kernel instead of loopy kernel """ import coffee.base as ast import loopy as lp # Just convert FInAT element to FIAT for now. # Dual evaluation in FInAT will bring a thorough revision. to_element = to_element.fiat_equivalent if any(len(dual.deriv_dict) != 0 for dual in to_element.dual_basis()): raise NotImplementedError( "Can only interpolate onto dual basis functionals without derivative evaluation, sorry!" ) if parameters is None: parameters = default_parameters() else: _ = default_parameters() _.update(parameters) parameters = _ # Determine whether in complex mode complex_mode = is_complex(parameters["scalar_type"]) # Find out which mapping to apply try: mapping, = set(to_element.mapping()) except ValueError: raise NotImplementedError( "Don't know how to interpolate onto zany spaces, sorry") expression = apply_mapping(expression, mapping, domain) # Apply UFL preprocessing expression = ufl_utils.preprocess_expression(expression, complex_mode=complex_mode) # Initialise kernel builder if interface is None: if coffee: import tsfc.kernel_interface.firedrake as firedrake_interface_coffee interface = firedrake_interface_coffee.ExpressionKernelBuilder else: # Delayed import, loopy is a runtime dependency import tsfc.kernel_interface.firedrake_loopy as firedrake_interface_loopy interface = firedrake_interface_loopy.ExpressionKernelBuilder builder = interface(parameters["scalar_type"]) arguments = extract_arguments(expression) argument_multiindices = tuple( builder.create_element(arg.ufl_element()).get_indices() for arg in arguments) # Replace coordinates (if any) unless otherwise specified by kwarg if domain is None: domain = expression.ufl_domain() assert domain is not None # Collect required coefficients first_coefficient_fake_coords = False coefficients = extract_coefficients(expression) if has_type(expression, GeometricQuantity) or any( fem.needs_coordinate_mapping(c.ufl_element()) for c in coefficients): # Create a fake coordinate coefficient for a domain. coords_coefficient = ufl.Coefficient( ufl.FunctionSpace(domain, domain.ufl_coordinate_element())) builder.domain_coordinate[domain] = coords_coefficient builder.set_cell_sizes(domain) coefficients = [coords_coefficient] + coefficients first_coefficient_fake_coords = True builder.set_coefficients(coefficients) # Split mixed coefficients expression = ufl_utils.split_coefficients(expression, builder.coefficient_split) # Translate to GEM kernel_cfg = dict( interface=builder, ufl_cell=domain.ufl_cell(), # FIXME: change if we ever implement # interpolation on facets. integral_type="cell", argument_multiindices=argument_multiindices, index_cache={}, scalar_type=parameters["scalar_type"]) if all( isinstance(dual, PointEvaluation) for dual in to_element.dual_basis()): # This is an optimisation for point-evaluation nodes which # should go away once FInAT offers the interface properly qpoints = [] # Everything is just a point evaluation. for dual in to_element.dual_basis(): ptdict = dual.get_point_dict() qpoint, = ptdict.keys() (qweight, component), = ptdict[qpoint] assert allclose(qweight, 1.0) assert component == () qpoints.append(qpoint) point_set = PointSet(qpoints) config = kernel_cfg.copy() config.update(point_set=point_set) # Allow interpolation onto QuadratureElements to refer to the quadrature # rule they represent if isinstance(to_element, FIAT.QuadratureElement): assert allclose(asarray(qpoints), asarray(to_element._points)) quad_rule = QuadratureRule(point_set, to_element._weights) config["quadrature_rule"] = quad_rule expr, = fem.compile_ufl(expression, **config, point_sum=False) # In some cases point_set.indices may be dropped from expr, but nothing # new should now appear assert set(expr.free_indices) <= set( chain(point_set.indices, *argument_multiindices)) shape_indices = tuple(gem.Index() for _ in expr.shape) basis_indices = point_set.indices ir = gem.Indexed(expr, shape_indices) else: # This is general code but is more unrolled than necssary. dual_expressions = [] # one for each functional broadcast_shape = len(expression.ufl_shape) - len( to_element.value_shape()) shape_indices = tuple(gem.Index() for _ in expression.ufl_shape[:broadcast_shape]) expr_cache = {} # Sharing of evaluation of the expression at points for dual in to_element.dual_basis(): pts = tuple(sorted(dual.get_point_dict().keys())) try: expr, point_set = expr_cache[pts] except KeyError: point_set = PointSet(pts) config = kernel_cfg.copy() config.update(point_set=point_set) expr, = fem.compile_ufl(expression, **config, point_sum=False) # In some cases point_set.indices may be dropped from expr, but # nothing new should now appear assert set(expr.free_indices) <= set( chain(point_set.indices, *argument_multiindices)) expr = gem.partial_indexed(expr, shape_indices) expr_cache[pts] = expr, point_set weights = collections.defaultdict(list) for p in pts: for (w, cmp) in dual.get_point_dict()[p]: weights[cmp].append(w) qexprs = gem.Zero() for cmp in sorted(weights): qweights = gem.Literal(weights[cmp]) qexpr = gem.Indexed(expr, cmp) qexpr = gem.index_sum( gem.Indexed(qweights, point_set.indices) * qexpr, point_set.indices) qexprs = gem.Sum(qexprs, qexpr) assert qexprs.shape == () assert set(qexprs.free_indices) == set( chain(shape_indices, *argument_multiindices)) dual_expressions.append(qexprs) basis_indices = (gem.Index(), ) ir = gem.Indexed(gem.ListTensor(dual_expressions), basis_indices) # Build kernel body return_indices = basis_indices + shape_indices + tuple( chain(*argument_multiindices)) return_shape = tuple(i.extent for i in return_indices) return_var = gem.Variable('A', return_shape) if coffee: return_arg = ast.Decl(parameters["scalar_type"], ast.Symbol('A', rank=return_shape)) else: return_arg = lp.GlobalArg("A", dtype=parameters["scalar_type"], shape=return_shape) return_expr = gem.Indexed(return_var, return_indices) # TODO: one should apply some GEM optimisations as in assembly, # but we don't for now. ir, = impero_utils.preprocess_gem([ir]) impero_c = impero_utils.compile_gem([(return_expr, ir)], return_indices) index_names = dict( (idx, "p%d" % i) for (i, idx) in enumerate(basis_indices)) # Handle kernel interface requirements builder.register_requirements([ir]) # Build kernel tuple return builder.construct_kernel(return_arg, impero_c, index_names, first_coefficient_fake_coords)
def test_listtensor(protocol): expr = gem.ListTensor([gem.Variable('x', ()), gem.Zero()]) unpickled = pickle.loads(pickle.dumps(expr, protocol)) assert expr == unpickled