Exemple #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"
Exemple #2
0
    def _entity_support_dofs(self):
        esd = {}
        for entity_dim in self.cell.sub_entities.keys():
            beta = self.get_indices()
            zeta = self.get_value_indices()

            entity_cell = self.cell.construct_subelement(entity_dim)
            quad = make_quadrature(entity_cell,
                                   (2 * numpy.array(self.degree)).tolist())

            eps = 1.e-8  # Is this a safe value?

            result = {}
            for f in self.entity_dofs()[entity_dim].keys():
                # Tabulate basis functions on the facet
                vals, = self.basis_evaluation(0,
                                              quad.point_set,
                                              entity=(entity_dim, f)).values()
                # Integrate the square of the basis functions on the facet.
                ints = gem.IndexSum(
                    gem.Product(
                        gem.IndexSum(
                            gem.Product(gem.Indexed(vals, beta + zeta),
                                        gem.Indexed(vals, beta + zeta)), zeta),
                        quad.weight_expression), quad.point_set.indices)
                evaluation, = evaluate([gem.ComponentTensor(ints, beta)])
                ints = evaluation.arr.flatten()
                assert evaluation.fids == ()
                result[f] = [dof for dof, i in enumerate(ints) if i > eps]

            esd[entity_dim] = result
        return esd
Exemple #3
0
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])
Exemple #4
0
    def basis_evaluation(self, order, ps, entity=None):
        r"""Produce the recipe for basis function evaluation at a set of points :math:`q`:

        .. math::
            \boldsymbol\phi_{(\gamma \epsilon) (i \alpha \beta) q} = \delta_{\alpha \gamma}\delta{\beta \epsilon}\phi_{i q}

            \nabla\boldsymbol\phi_{(\epsilon \gamma \zeta) (i \alpha \beta) q} = \delta_{\alpha \epsilon} \deta{\beta \gamma}\nabla\phi_{\zeta i q}
        """
        # 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)))

        scalar_result = self._base_element.basis_evaluation(order, ps, entity)
        result = {}
        for alpha, expr in iteritems(scalar_result):
            result[alpha] = gem.ComponentTensor(
                gem.Product(deltas, gem.Indexed(expr, scalar_i + scalar_vi)),
                scalar_i + tensor_i + scalar_vi + tensor_vi)
        return result
Exemple #5
0
def translate_coefficient(terminal, mt, ctx):
    vec = ctx.coefficient(terminal, mt.restriction)

    if terminal.ufl_element().family() == 'Real':
        assert mt.local_derivatives == 0
        return vec

    element = ctx.create_element(terminal.ufl_element())

    # Collect FInAT tabulation for all entities
    per_derivative = collections.defaultdict(list)
    for entity_id in ctx.entity_ids:
        finat_dict = ctx.basis_evaluation(element, mt.local_derivatives, entity_id)
        for alpha, table in finat_dict.items():
            # Filter out irrelevant derivatives
            if sum(alpha) == mt.local_derivatives:
                # A numerical hack that FFC used to apply on FIAT
                # tables still lives on after ditching FFC and
                # switching to FInAT.
                table = ffc_rounding(table, ctx.epsilon)
                per_derivative[alpha].append(table)

    # Merge entity tabulations for each derivative
    if len(ctx.entity_ids) == 1:
        def take_singleton(xs):
            x, = xs  # asserts singleton
            return x
        per_derivative = {alpha: take_singleton(tables)
                          for alpha, tables in per_derivative.items()}
    else:
        f = ctx.entity_number(mt.restriction)
        per_derivative = {alpha: gem.select_expression(tables, f)
                          for alpha, tables in per_derivative.items()}

    # Coefficient evaluation
    ctx.index_cache.setdefault(terminal.ufl_element(), element.get_indices())
    beta = ctx.index_cache[terminal.ufl_element()]
    zeta = element.get_value_indices()
    vec_beta, = gem.optimise.remove_componenttensors([gem.Indexed(vec, beta)])
    value_dict = {}
    for alpha, table in per_derivative.items():
        table_qi = gem.Indexed(table, beta + zeta)
        summands = []
        for var, expr in unconcatenate([(vec_beta, table_qi)], ctx.index_cache):
            value = gem.IndexSum(gem.Product(expr, var), var.index_ordering())
            summands.append(gem.optimise.contraction(value))
        optimised_value = gem.optimise.make_sum(summands)
        value_dict[alpha] = gem.ComponentTensor(optimised_value, zeta)

    # Change from FIAT to UFL arrangement
    result = fiat_to_ufl(value_dict, mt.local_derivatives)
    assert result.shape == mt.expr.ufl_shape
    assert set(result.free_indices) <= set(ctx.point_indices)

    # Detect Jacobian of affine cells
    if not result.free_indices and all(numpy.count_nonzero(node.array) <= 2
                                       for node in traversal((result,))
                                       if isinstance(node, gem.Literal)):
        result = gem.optimise.aggressive_unroll(result)
    return result
Exemple #6
0
    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
Exemple #7
0
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)
Exemple #8
0
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
Exemple #9
0
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!"
Exemple #10
0
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
Exemple #11
0
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)
Exemple #12
0
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))
Exemple #13
0
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
Exemple #14
0
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!"
Exemple #15
0
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
Exemple #16
0
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)