def _find_linear_terms(rhs_exprs, u):
    """Help function that takes a list of rhs expressions and return a
    list of bools determining what component, rhs_exprs[i], is linear
    wrt u[i].

    """

    uu = [Constant(1.0) for _ in rhs_exprs]
    if len(rhs_exprs) > 1:
        repl = {u: ufl.as_vector(uu)}
    else:
        repl = {u: uu[0]}

    linear_terms = []
    for i, ui in enumerate(uu):
        comp_i_s = expand_indices(ufl.replace(rhs_exprs[i], repl))
        linear_terms.append(ui in extract_coefficients(comp_i_s) and
                            ui not in extract_coefficients(
                                expand_derivatives(ufl.diff(comp_i_s, ui))))
    return linear_terms
Ejemplo n.º 2
0
def derivative(form, u, du=None, coefficient_derivatives=None):
    """Compute the derivative of a form.

    Given a form, this computes its linearization with respect to the
    provided :class:`.Function`.  The resulting form has one
    additional :class:`Argument` in the same finite element space as
    the Function.

    :arg form: a :class:`~ufl.classes.Form` to compute the derivative of.
    :arg u: a :class:`.Function` to compute the derivative with
         respect to.
    :arg du: an optional :class:`Argument` to use as the replacement
         in the new form (constructed automatically if not provided).
    :arg coefficient_derivatives: an optional :class:`dict` to
         provide the derivative of a coefficient function.

    :raises ValueError: If any of the coefficients in ``form`` were
        obtained from ``u.split()``.  UFL doesn't notice that these
        are related to ``u`` and so therefore the derivative is
        wrong (instead one should have written ``split(u)``).

    See also :func:`ufl.derivative`.
    """
    if set(extract_coefficients(form)) & set(u.split()):
        raise ValueError(
            "Taking derivative of form wrt u, but form contains coefficients from u.split()."
            "\nYou probably meant to write split(u) when defining your form.")
    if du is None:
        if isinstance(u, function.Function):
            V = u.function_space()
            args = form.arguments()
            number = max(a.number() for a in args) if args else -1
            du = Argument(V, number + 1)
        else:
            raise RuntimeError("Can't compute derivative for form")
    return ufl.derivative(form, u, du, coefficient_derivatives)
Ejemplo n.º 3
0
def test_extract_coefficients_vs_fixture(coefficients, forms):
    assert coefficients == tuple(extract_coefficients(forms[2]))
Ejemplo n.º 4
0
def compile_ufl_kernel(expression, to_pts, to_element, fs):
    import collections
    from ufl.algorithms.apply_function_pullbacks import apply_function_pullbacks
    from ufl.algorithms.apply_algebra_lowering import apply_algebra_lowering
    from ufl.algorithms.apply_derivatives import apply_derivatives
    from ufl.algorithms.apply_geometry_lowering import apply_geometry_lowering
    from ufl.algorithms import extract_arguments, extract_coefficients
    from gem import gem, impero_utils
    from tsfc import fem, ufl_utils
    from tsfc.coffee import generate as generate_coffee
    from tsfc.kernel_interface import (KernelBuilderBase,
                                       needs_cell_orientations,
                                       cell_orientations_coffee_arg)

    # Imitate the compute_form_data processing pipeline
    #
    # Unfortunately, we cannot call compute_form_data here, since
    # we only have an expression, not a form
    expression = apply_algebra_lowering(expression)
    expression = apply_derivatives(expression)
    expression = apply_function_pullbacks(expression)
    expression = apply_geometry_lowering(expression)
    expression = apply_derivatives(expression)
    expression = apply_geometry_lowering(expression)
    expression = apply_derivatives(expression)

    # Replace coordinates (if any)
    if expression.ufl_domain():
        assert fs.mesh() == expression.ufl_domain()
        expression = ufl_utils.replace_coordinates(expression, fs.mesh().coordinates)

    if extract_arguments(expression):
        return ValueError("Cannot interpolate UFL expression with Arguments!")

    builder = KernelBuilderBase()
    args = []

    coefficients = extract_coefficients(expression)
    for i, coefficient in enumerate(coefficients):
        args.append(builder.coefficient(coefficient, "w_%d" % i))

    point_index = gem.Index(name='p')
    ir = fem.process('cell', fs.mesh().ufl_cell(), to_pts, None,
                     point_index, (), expression,
                     builder.coefficient_mapper,
                     collections.defaultdict(gem.Index))
    assert len(ir) == 1

    # Deal with non-scalar expressions
    tensor_indices = ()
    if fs.shape:
        tensor_indices = tuple(gem.Index() for s in fs.shape)
        ir = [gem.Indexed(ir[0], tensor_indices)]

    # Build kernel body
    return_var = gem.Variable('A', (len(to_pts),) + fs.shape)
    return_expr = gem.Indexed(return_var, (point_index,) + tensor_indices)
    impero_c = impero_utils.compile_gem([return_expr], ir, [point_index])
    body = generate_coffee(impero_c, index_names={point_index: 'p'})

    oriented = needs_cell_orientations(ir)
    if oriented:
        args.insert(0, cell_orientations_coffee_arg)

    # Build kernel
    args.insert(0, ast.Decl("double", ast.Symbol('A', rank=(len(to_pts),) + fs.shape)))
    kernel_code = builder.construct_kernel("expression_kernel", args, body)

    return op2.Kernel(kernel_code, kernel_code.name), oriented, coefficients
Ejemplo n.º 5
0
def compile_ufl_kernel(expression, to_pts, to_element, fs):
    import collections
    from ufl.algorithms.apply_function_pullbacks import apply_function_pullbacks
    from ufl.algorithms.apply_algebra_lowering import apply_algebra_lowering
    from ufl.algorithms.apply_derivatives import apply_derivatives
    from ufl.algorithms.apply_geometry_lowering import apply_geometry_lowering
    from ufl.algorithms import extract_arguments, extract_coefficients
    from gem import gem, impero_utils
    from tsfc import fem, ufl_utils
    from tsfc.coffee import generate as generate_coffee
    from tsfc.kernel_interface import (KernelBuilderBase,
                                       needs_cell_orientations,
                                       cell_orientations_coffee_arg)

    # Imitate the compute_form_data processing pipeline
    #
    # Unfortunately, we cannot call compute_form_data here, since
    # we only have an expression, not a form
    expression = apply_algebra_lowering(expression)
    expression = apply_derivatives(expression)
    expression = apply_function_pullbacks(expression)
    expression = apply_geometry_lowering(expression)
    expression = apply_derivatives(expression)
    expression = apply_geometry_lowering(expression)
    expression = apply_derivatives(expression)

    # Replace coordinates (if any)
    if expression.ufl_domain():
        assert fs.mesh() == expression.ufl_domain()
        expression = ufl_utils.replace_coordinates(expression, fs.mesh().coordinates)

    if extract_arguments(expression):
        return ValueError("Cannot interpolate UFL expression with Arguments!")

    builder = KernelBuilderBase()
    args = []

    coefficients = extract_coefficients(expression)
    for i, coefficient in enumerate(coefficients):
        args.append(builder.coefficient(coefficient, "w_%d" % i))

    point_index = gem.Index(name='p')
    ir = fem.compile_ufl(expression,
                         cell=fs.mesh().ufl_cell(),
                         points=to_pts,
                         point_index=point_index,
                         coefficient_mapper=builder.coefficient_mapper)
    assert len(ir) == 1

    # Deal with non-scalar expressions
    tensor_indices = ()
    if fs.shape:
        tensor_indices = tuple(gem.Index() for s in fs.shape)
        ir = [gem.Indexed(ir[0], tensor_indices)]

    # Build kernel body
    return_var = gem.Variable('A', (len(to_pts),) + fs.shape)
    return_expr = gem.Indexed(return_var, (point_index,) + tensor_indices)
    impero_c = impero_utils.compile_gem([return_expr], ir, [point_index])
    body = generate_coffee(impero_c, index_names={point_index: 'p'})

    oriented = needs_cell_orientations(ir)
    if oriented:
        args.insert(0, cell_orientations_coffee_arg)

    # Build kernel
    args.insert(0, ast.Decl("double", ast.Symbol('A', rank=(len(to_pts),) + fs.shape)))
    kernel_code = builder.construct_kernel("expression_kernel", args, body)

    return op2.Kernel(kernel_code, kernel_code.name), oriented, coefficients
Ejemplo n.º 6
0
def derivative(form, u, du=None, coefficient_derivatives=None):
    """Compute the derivative of a form.

    Given a form, this computes its linearization with respect to the
    provided :class:`.Function`.  The resulting form has one
    additional :class:`Argument` in the same finite element space as
    the Function.

    :arg form: a :class:`~ufl.classes.Form` to compute the derivative of.
    :arg u: a :class:`.Function` to compute the derivative with
         respect to.
    :arg du: an optional :class:`Argument` to use as the replacement
         in the new form (constructed automatically if not provided).
    :arg coefficient_derivatives: an optional :class:`dict` to
         provide the derivative of a coefficient function.

    :raises ValueError: If any of the coefficients in ``form`` were
        obtained from ``u.split()``.  UFL doesn't notice that these
        are related to ``u`` and so therefore the derivative is
        wrong (instead one should have written ``split(u)``).

    See also :func:`ufl.derivative`.
    """
    # TODO: What about Constant?
    u_is_x = isinstance(u, ufl.SpatialCoordinate)
    if not u_is_x and len(u.split()) > 1 and set(
            extract_coefficients(form)) & set(u.split()):
        raise ValueError(
            "Taking derivative of form wrt u, but form contains coefficients from u.split()."
            "\nYou probably meant to write split(u) when defining your form.")

    mesh = form.ufl_domain()
    is_dX = u_is_x or u is mesh.coordinates
    args = form.arguments()

    def argument(V):
        if du is None:
            n = max(a.number() for a in args) if args else -1
            return Argument(V, n + 1)
        else:
            return du

    if is_dX:
        coords = mesh.coordinates
        u = ufl.SpatialCoordinate(mesh)
        V = coords.function_space()
        du = argument(V)
        cds = {coords: du}
        if coefficient_derivatives is not None:
            cds.update(coefficient_derivatives)
        coefficient_derivatives = cds
    elif isinstance(u, firedrake.Function):
        V = u.function_space()
        du = argument(V)
    elif isinstance(u, firedrake.Constant):
        if u.ufl_shape != ():
            raise ValueError(
                "Real function space of vector elements not supported")
        V = firedrake.FunctionSpace(mesh, "Real", 0)
        du = argument(V)
    else:
        raise RuntimeError("Can't compute derivative for form")

    return ufl.derivative(form, u, du, coefficient_derivatives)
Ejemplo n.º 7
0
def derivative(form, u, du=None, coefficient_derivatives=None):
    """Compute the derivative of a form.

    Given a form, this computes its linearization with respect to the
    provided :class:`.Function`.  The resulting form has one
    additional :class:`Argument` in the same finite element space as
    the Function.

    :arg form: a :class:`~ufl.classes.Form` to compute the derivative of.
    :arg u: a :class:`.Function` to compute the derivative with
         respect to.
    :arg du: an optional :class:`Argument` to use as the replacement
         in the new form (constructed automatically if not provided).
    :arg coefficient_derivatives: an optional :class:`dict` to
         provide the derivative of a coefficient function.

    :raises ValueError: If any of the coefficients in ``form`` were
        obtained from ``u.split()``.  UFL doesn't notice that these
        are related to ``u`` and so therefore the derivative is
        wrong (instead one should have written ``split(u)``).

    See also :func:`ufl.derivative`.
    """
    if isinstance(form, firedrake.slate.TensorBase):
        raise TypeError(
            f"Cannot take the derivative of a {type(form).__name__}"
        )
    # TODO: What about Constant?
    u_is_x = isinstance(u, ufl.SpatialCoordinate)
    uc, = (u,) if u_is_x else extract_coefficients(u)
    if not u_is_x and len(uc.split()) > 1 and set(extract_coefficients(form)) & set(uc.split()):
        raise ValueError("Taking derivative of form wrt u, but form contains coefficients from u.split()."
                         "\nYou probably meant to write split(u) when defining your form.")

    mesh = form.ufl_domain()
    if not mesh:
        raise ValueError("Expression to be differentiated has no ufl domain."
                         "\nDo you need to add a domain to your Constant?")
    is_dX = u_is_x or u is mesh.coordinates
    try:
        args = form.arguments()
    except AttributeError:
        args = extract_arguments(form)

    def argument(V):
        if du is None:
            n = max(a.number() for a in args) if args else -1
            return Argument(V, n + 1)
        else:
            return du

    if is_dX:
        coords = mesh.coordinates
        u = ufl.SpatialCoordinate(mesh)
        V = coords.function_space()
        du = argument(V)
        cds = {coords: du}
        if coefficient_derivatives is not None:
            cds.update(coefficient_derivatives)
        coefficient_derivatives = cds
    elif isinstance(uc, firedrake.Function):
        V = uc.function_space()
        du = argument(V)
    elif isinstance(uc, firedrake.Constant):
        if uc.ufl_shape != ():
            raise ValueError("Real function space of vector elements not supported")
        V = firedrake.FunctionSpace(mesh, "Real", 0)
        du = argument(V)
    else:
        raise RuntimeError("Can't compute derivative for form")

    if u.ufl_shape != du.ufl_shape:
        raise ValueError("Shapes of u and du do not match.\n"
                         "If you passed an indexed part of split(u) into "
                         "derivative, you need to provide an appropriate du as well.")
    return ufl.derivative(form, u, du, coefficient_derivatives)
Ejemplo n.º 8
0
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)
Ejemplo n.º 9
0
def derivative(form, u, du=None, coefficient_derivatives=None):
    """Compute the derivative of a form.

    Given a form, this computes its linearization with respect to the
    provided :class:`.Function`.  The resulting form has one
    additional :class:`Argument` in the same finite element space as
    the Function.

    :arg form: a :class:`~ufl.classes.Form` to compute the derivative of.
    :arg u: a :class:`.Function` to compute the derivative with
         respect to.
    :arg du: an optional :class:`Argument` to use as the replacement
         in the new form (constructed automatically if not provided).
    :arg coefficient_derivatives: an optional :class:`dict` to
         provide the derivative of a coefficient function.

    :raises ValueError: If any of the coefficients in ``form`` were
        obtained from ``u.split()``.  UFL doesn't notice that these
        are related to ``u`` and so therefore the derivative is
        wrong (instead one should have written ``split(u)``).

    See also :func:`ufl.derivative`.
    """
    # TODO: What about Constant?
    u_is_x = isinstance(u, ufl.SpatialCoordinate)
    if not u_is_x and len(u.split()) > 1 and set(extract_coefficients(form)) & set(u.split()):
        raise ValueError("Taking derivative of form wrt u, but form contains coefficients from u.split()."
                         "\nYou probably meant to write split(u) when defining your form.")

    mesh = form.ufl_domain()
    is_dX = u_is_x or u is mesh.coordinates
    args = form.arguments()

    def argument(V):
        if du is None:
            n = max(a.number() for a in args) if args else -1
            return Argument(V, n + 1)
        else:
            return du

    if is_dX:
        coords = mesh.coordinates
        u = ufl.SpatialCoordinate(mesh)
        V = coords.function_space()
        du = argument(V)
        cds = {coords: du}
        if coefficient_derivatives is not None:
            cds.update(coefficient_derivatives)
        coefficient_derivatives = cds
    elif isinstance(u, firedrake.Function):
        V = u.function_space()
        du = argument(V)
    elif isinstance(u, firedrake.Constant):
        if u.ufl_shape != ():
            raise ValueError("Real function space of vector elements not supported")
        V = firedrake.FunctionSpace(mesh, "Real", 0)
        du = argument(V)
    else:
        raise RuntimeError("Can't compute derivative for form")

    return ufl.derivative(form, u, du, coefficient_derivatives)
Ejemplo n.º 10
0
def compile_expression_at_points(expression,
                                 points,
                                 coordinates,
                                 interface=None,
                                 parameters=None,
                                 coffee=True):
    """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 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

    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"])

    # 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)
    domain = expression.ufl_domain()
    if domain:
        assert coordinates.ufl_domain() == domain
        builder.domain_coordinate[domain] = coordinates
        builder.set_cell_sizes(domain)

    # Collect required coefficients
    coefficients = extract_coefficients(expression)
    if has_type(expression, GeometricQuantity) or any(
            fem.needs_coordinate_mapping(c.ufl_element())
            for c in coefficients):
        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,
                  argument_multiindices=argument_multiindices)
    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_indices = point_set.indices + tensor_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)
    ir, = impero_utils.preprocess_gem([ir])
    impero_c = impero_utils.compile_gem([(return_expr, ir)], return_indices)
    point_index, = point_set.indices

    # Handle kernel interface requirements
    builder.register_requirements([ir])
    # Build kernel tuple
    return builder.construct_kernel(return_arg, impero_c,
                                    parameters["precision"],
                                    {point_index: 'p'})
Ejemplo n.º 11
0
def test_extract_coefficients_vs_fixture(coefficients, forms):
    assert coefficients == tuple(extract_coefficients(forms[2]))
Ejemplo n.º 12
0
def compile_element(expression, coordinates, parameters=None):
    """Generates C code for point evaluations.

    :arg expression: UFL expression
    :arg coordinates: coordinate field
    :arg parameters: form compiler parameters
    :returns: C code as string
    """
    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 = tsfc.ufl_utils.preprocess_expression(
        expression, complex_mode=utils.complex_mode)

    # Collect required coefficients
    coefficient, = extract_coefficients(expression)

    # Point evaluation of mixed coefficients not supported here
    if type(coefficient.ufl_element()) == MixedElement:
        raise NotImplementedError("Cannot point evaluate mixed elements yet!")

    # Replace coordinates (if any)
    domain = expression.ufl_domain()
    assert coordinates.ufl_domain() == domain

    # Initialise kernel builder
    builder = firedrake_interface.KernelBuilderBase(utils.ScalarType_c)
    builder.domain_coordinate[domain] = coordinates
    x_arg = builder._coefficient(coordinates, "x")
    f_arg = builder._coefficient(coefficient, "f")

    # TODO: restore this for expression evaluation!
    # expression = ufl_utils.split_coefficients(expression, builder.coefficient_split)

    # Translate to GEM
    cell = domain.ufl_cell()
    dim = cell.topological_dimension()
    point = gem.Variable('X', (dim, ))
    point_arg = ast.Decl(utils.ScalarType_c, ast.Symbol('X', rank=(dim, )))

    config = dict(interface=builder,
                  ufl_cell=coordinates.ufl_domain().ufl_cell(),
                  precision=parameters["precision"],
                  point_indices=(),
                  point_expr=point,
                  complex_mode=utils.complex_mode)
    # TODO: restore this for expression evaluation!
    # config["cellvolume"] = cellvolume_generator(coordinates.ufl_domain(), coordinates, config)
    context = tsfc.fem.GemPointContext(**config)

    # Abs-simplification
    expression = tsfc.ufl_utils.simplify_abs(expression, utils.complex_mode)

    # Translate UFL -> GEM
    translator = tsfc.fem.Translator(context)
    result, = map_expr_dags(translator, [expression])

    tensor_indices = ()
    if expression.ufl_shape:
        tensor_indices = tuple(gem.Index() for s in expression.ufl_shape)
        return_variable = gem.Indexed(gem.Variable('R', expression.ufl_shape),
                                      tensor_indices)
        result_arg = ast.Decl(utils.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(utils.ScalarType_c, ast.Symbol('R', rank=(1, )))

    # 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"],
                           utils.ScalarType_c)

    # Build kernel tuple
    kernel_code = builder.construct_kernel(
        "evaluate_kernel", [result_arg, point_arg, x_arg, f_arg], body)

    # Fill the code template
    extruded = isinstance(cell, TensorProductCell)

    code = {
        "geometric_dimension": cell.geometric_dimension(),
        "layers_arg": ", int const *__restrict__ layers" if extruded else "",
        "layers": ", layers" if extruded else "",
        "IntType": as_cstr(IntType),
        "scalar_type": utils.ScalarType_c,
    }
    # if maps are the same, only need to pass one of them
    if coordinates.cell_node_map() == coefficient.cell_node_map():
        code[
            "wrapper_map_args"] = "%(IntType)s const *__restrict__ coords_map" % code
        code["map_args"] = "f->coords_map"
    else:
        code[
            "wrapper_map_args"] = "%(IntType)s const *__restrict__ coords_map, %(IntType)s const *__restrict__ f_map" % code
        code["map_args"] = "f->coords_map, f->f_map"

    evaluate_template_c = """
static inline void wrap_evaluate(%(scalar_type)s* const result, %(scalar_type)s* const X, int const start, int const end%(layers_arg)s,
    %(scalar_type)s const *__restrict__ coords, %(scalar_type)s const *__restrict__ f, %(wrapper_map_args)s);

int evaluate(struct Function *f, %(scalar_type)s *x, %(scalar_type)s *result)
{
    struct ReferenceCoords reference_coords;
    %(IntType)s cell = locate_cell(f, x, %(geometric_dimension)d, &to_reference_coords, &to_reference_coords_xtr, &reference_coords);
    if (cell == -1) {
        return -1;
    }

    if (!result) {
        return 0;
    }
    int layers[2] = {0, 0};
    if (f->extruded != 0) {
        int nlayers = f->n_layers;
        layers[1] = cell %% nlayers + 2;
        cell = cell / nlayers;
    }

    wrap_evaluate(result, reference_coords.X, cell, cell+1%(layers)s, f->coords, f->f, %(map_args)s);
    return 0;
}
"""

    return (evaluate_template_c % code) + kernel_code.gencode()
Ejemplo n.º 13
0
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)
Ejemplo n.º 14
0
def compile_element(expression, coordinates, parameters=None):
    """Generates C code for point evaluations.

    :arg expression: UFL expression
    :arg coordinates: coordinate field
    :arg parameters: form compiler parameters
    :returns: C code as string
    """
    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 = tsfc.ufl_utils.preprocess_expression(expression)

    # Collect required coefficients
    coefficient, = extract_coefficients(expression)

    # Point evaluation of mixed coefficients not supported here
    if type(coefficient.ufl_element()) == MixedElement:
        raise NotImplementedError("Cannot point evaluate mixed elements yet!")

    # Replace coordinates (if any)
    domain = expression.ufl_domain()
    assert coordinates.ufl_domain() == domain
    expression = tsfc.ufl_utils.replace_coordinates(expression, coordinates)

    # Initialise kernel builder
    builder = firedrake_interface.KernelBuilderBase()
    x_arg = builder._coefficient(coordinates, "x")
    f_arg = builder._coefficient(coefficient, "f")

    # TODO: restore this for expression evaluation!
    # expression = ufl_utils.split_coefficients(expression, builder.coefficient_split)

    # Translate to GEM
    cell = domain.ufl_cell()
    dim = cell.topological_dimension()
    point = gem.Variable('X', (dim, ))
    point_arg = ast.Decl(SCALAR_TYPE, ast.Symbol('X', rank=(dim, )))

    config = dict(interface=builder,
                  ufl_cell=coordinates.ufl_domain().ufl_cell(),
                  precision=parameters["precision"],
                  point_indices=(),
                  point_expr=point)
    # TODO: restore this for expression evaluation!
    # config["cellvolume"] = cellvolume_generator(coordinates.ufl_domain(), coordinates, config)
    context = tsfc.fem.GemPointContext(**config)

    # Abs-simplification
    expression = tsfc.ufl_utils.simplify_abs(expression)

    # Translate UFL -> GEM
    translator = tsfc.fem.Translator(context)
    result, = map_expr_dags(translator, [expression])

    tensor_indices = ()
    if expression.ufl_shape:
        tensor_indices = tuple(gem.Index() for s in expression.ufl_shape)
        return_variable = gem.Indexed(gem.Variable('R', expression.ufl_shape),
                                      tensor_indices)
        result_arg = ast.Decl(SCALAR_TYPE,
                              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(SCALAR_TYPE, ast.Symbol('R', rank=(1, )))

    # 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"])

    # Build kernel tuple
    kernel_code = builder.construct_kernel(
        "evaluate_kernel", [result_arg, point_arg, x_arg, f_arg], body)

    # Fill the code template
    extruded = isinstance(cell, TensorProductCell)

    #code = {
    #"geometric_dimension": cell.geometric_dimension(),
    #"extruded_arg": ", %s nlayers" % as_cstr(IntType) if extruded else "",
    #"nlayers": ", f->n_layers" if extruded else "",
    #"IntType": as_cstr(IntType),
    #}

    #evaluate_template_c = """static inline void wrap_evaluate(double *result, double *X, double *coords, %(IntType)s *coords_map, double *f, %(IntType)s *f_map%(extruded_arg)s, %(IntType)s cell);

    #int evaluate(struct Function *f, double *x, double *result)
    #{
    #struct ReferenceCoords reference_coords;
    #%(IntType)s cell = locate_cell(f, x, %(geometric_dimension)d, &to_reference_coords, &reference_coords);
    #if (cell == -1) {
    #return -1;
    #}

    #if (!result) {
    #return 0;
    #}

    #wrap_evaluate(result, reference_coords.X, f->coords, f->coords_map, f->f, f->f_map%(nlayers)s, cell);
    #return 0;
    #}
    #"""

    #    return (evaluate_template_c % code) + kernel_code.gencode()
    return kernel_code.gencode()
Ejemplo n.º 15
0
def compute_expression_ir(expression, analysis, parameters, visualise):
    """Compute IR for expression.

    Parameters
    ----------
    expression
        Triple of (UFL expression, array of evaluation points, original UFL expression).

    Note
    ----
    Original UFL expression is needed to compute original positions of coefficients.

    """
    logger.info("Computing uflacs representation of expression")

    original_expression = expression[2]
    points = expression[1]
    expression = expression[0]

    num_points = points.shape[0]
    weights = numpy.array([1.0] * num_points)

    cell = expression.ufl_domain().ufl_cell()

    ir = {}

    # Prepare dimensions of all unique element in expression,
    # including elements for arguments, coefficients and coordinate mappings
    ir["element_dimensions"] = {
        ufl_element: create_element(ufl_element).space_dimension()
        for ufl_element in analysis.unique_elements
    }

    # Extract dimensions for elements of arguments only
    arguments = extract_arguments(expression)
    argument_elements = tuple(f.ufl_element() for f in arguments)
    argument_dimensions = [
        ir["element_dimensions"][ufl_element]
        for ufl_element in argument_elements
    ]

    tensor_shape = argument_dimensions
    ir["tensor_shape"] = tensor_shape

    ir["expression_shape"] = list(expression.ufl_shape)

    coefficients = extract_coefficients(expression)
    coefficient_numbering = {}
    for i, coeff in enumerate(coefficients):
        coefficient_numbering[coeff] = i

    # Add coefficient numbering to IR
    ir["coefficient_numbering"] = coefficient_numbering

    original_coefficient_positions = []
    original_coefficients = extract_coefficients(original_expression)
    for coeff in coefficients:
        original_coefficient_positions.append(
            original_coefficients.index(coeff))

    ir["original_coefficient_positions"] = original_coefficient_positions

    coefficient_elements = tuple(f.ufl_element() for f in coefficients)

    offsets = {}
    _offset = 0
    for i, el in enumerate(coefficient_elements):
        offsets[coefficients[i]] = _offset
        _offset += ir["element_dimensions"][el]

    # Copy offsets also into IR
    ir["coefficient_offsets"] = offsets

    ir["integral_type"] = "expression"
    ir["entitytype"] = "cell"

    # Build offsets for Constants
    original_constant_offsets = {}
    _offset = 0
    for constant in extract_constants(expression):
        original_constant_offsets[constant] = _offset
        _offset += numpy.product(constant.ufl_shape, dtype=numpy.int)

    ir["original_constant_offsets"] = original_constant_offsets

    ir["points"] = points

    integrands = {num_points: expression}
    quadrature_rules = {num_points: (points, weights)}

    uflacs_ir = build_uflacs_ir(cell, ir["integral_type"], ir["entitytype"],
                                integrands, tensor_shape, quadrature_rules,
                                parameters, visualise)

    ir.update(uflacs_ir)

    return ir
Ejemplo n.º 16
0
def compile_element(expression, coordinates, parameters=None):
    """Generates C code for point evaluations.

    :arg expression: UFL expression
    :arg coordinates: coordinate field
    :arg parameters: form compiler parameters
    :returns: C code as string
    """
    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 = tsfc.ufl_utils.preprocess_expression(expression)

    # Collect required coefficients
    coefficient, = extract_coefficients(expression)

    # Point evaluation of mixed coefficients not supported here
    if type(coefficient.ufl_element()) == MixedElement:
        raise NotImplementedError("Cannot point evaluate mixed elements yet!")

    # Replace coordinates (if any)
    domain = expression.ufl_domain()
    assert coordinates.ufl_domain() == domain

    # Initialise kernel builder
    builder = firedrake_interface.KernelBuilderBase()
    builder.domain_coordinate[domain] = coordinates
    x_arg = builder._coefficient(coordinates, "x")
    f_arg = builder._coefficient(coefficient, "f")

    # TODO: restore this for expression evaluation!
    # expression = ufl_utils.split_coefficients(expression, builder.coefficient_split)

    # Translate to GEM
    cell = domain.ufl_cell()
    dim = cell.topological_dimension()
    point = gem.Variable('X', (dim,))
    point_arg = ast.Decl(SCALAR_TYPE, ast.Symbol('X', rank=(dim,)))

    config = dict(interface=builder,
                  ufl_cell=coordinates.ufl_domain().ufl_cell(),
                  precision=parameters["precision"],
                  point_indices=(),
                  point_expr=point)
    # TODO: restore this for expression evaluation!
    # config["cellvolume"] = cellvolume_generator(coordinates.ufl_domain(), coordinates, config)
    context = tsfc.fem.GemPointContext(**config)

    # Abs-simplification
    expression = tsfc.ufl_utils.simplify_abs(expression)

    # Translate UFL -> GEM
    translator = tsfc.fem.Translator(context)
    result, = map_expr_dags(translator, [expression])

    tensor_indices = ()
    if expression.ufl_shape:
        tensor_indices = tuple(gem.Index() for s in expression.ufl_shape)
        return_variable = gem.Indexed(gem.Variable('R', expression.ufl_shape), tensor_indices)
        result_arg = ast.Decl(SCALAR_TYPE, 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(SCALAR_TYPE, ast.Symbol('R', rank=(1,)))

    # 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"])

    # Build kernel tuple
    kernel_code = builder.construct_kernel("evaluate_kernel", [result_arg, point_arg, x_arg, f_arg], body)

    # Fill the code template
    extruded = isinstance(cell, TensorProductCell)

    code = {
        "geometric_dimension": cell.geometric_dimension(),
        "layers_arg": ", int const *__restrict__ layers" if extruded else "",
        "layers": ", layers" if extruded else "",
        "IntType": as_cstr(IntType),
    }
    # if maps are the same, only need to pass one of them
    if coordinates.cell_node_map() == coefficient.cell_node_map():
        code["wrapper_map_args"] = "%(IntType)s const *__restrict__ coords_map" % code
        code["map_args"] = "f->coords_map"
    else:
        code["wrapper_map_args"] = "%(IntType)s const *__restrict__ coords_map, %(IntType)s const *__restrict__ f_map" % code
        code["map_args"] = "f->coords_map, f->f_map"

    evaluate_template_c = """
static inline void wrap_evaluate(double* const result, double* const X, int const start, int const end%(layers_arg)s,
    double const *__restrict__ coords, double const *__restrict__ f, %(wrapper_map_args)s);

int evaluate(struct Function *f, double *x, double *result)
{
    struct ReferenceCoords reference_coords;
    %(IntType)s cell = locate_cell(f, x, %(geometric_dimension)d, &to_reference_coords, &to_reference_coords_xtr, &reference_coords);
    if (cell == -1) {
        return -1;
    }

    if (!result) {
        return 0;
    }
    int layers[2] = {0, 0};
    if (f->extruded != 0) {
        int nlayers = f->n_layers;
        layers[1] = cell %% nlayers + 2;
        cell = cell / nlayers;
    }

    wrap_evaluate(result, reference_coords.X, cell, cell+1%(layers)s, f->coords, f->f, %(map_args)s);
    return 0;
}
"""

    return (evaluate_template_c % code) + kernel_code.gencode()