Esempio n. 1
0
    def facetarea():
        from ufl import Measure
        assert integral_type != 'cell'
        integrand, degree = ufl_utils.one_times(Measure(integral_type, domain=domain))
        integrand = ufl_utils.replace_coordinates(integrand, coordinate_coefficient)

        config = kernel_config.copy()
        config.update(quadrature_degree=degree)
        expr, = fem.compile_ufl(integrand, point_sum=True, **config)
        return expr
Esempio n. 2
0
    def cellvolume(restriction):
        from ufl import dx
        integrand, degree = ufl_utils.one_times(dx(domain=domain))
        integrand = ufl_utils.replace_coordinates(integrand, coordinate_coefficient)
        interface = CellVolumeKernelInterface(kernel_config["interface"], restriction)

        config = {k: v for k, v in kernel_config.items()
                  if k in ["ufl_cell", "precision", "index_cache"]}
        config.update(interface=interface, quadrature_degree=degree)
        expr, = fem.compile_ufl(integrand, point_sum=True, **config)
        return expr
Esempio n. 3
0
    def facetarea():
        from ufl import Measure
        assert integral_type != 'cell'
        integrand, degree = ufl_utils.one_times(Measure(integral_type, domain=domain))
        integrand = ufl_utils.replace_coordinates(integrand, coordinate_coefficient)

        quadrature_index = gem.Index(name='q')

        config = kernel_config.copy()
        config.update(quadrature_degree=degree, point_index=quadrature_index)
        expr, = fem.compile_ufl(integrand, **config)
        if quadrature_index in expr.free_indices:
            expr = gem.IndexSum(expr, quadrature_index)
        return expr
Esempio n. 4
0
    def cellvolume(restriction):
        from ufl import dx
        integrand, degree = ufl_utils.one_times(dx(domain=domain))
        integrand = ufl_utils.replace_coordinates(integrand, coordinate_coefficient)

        interface = CellVolumeKernelInterface(kernel_config["interface"], restriction)
        quadrature_index = gem.Index(name='q')

        config = {k: v for k, v in kernel_config.items()
                  if k in ["ufl_cell", "precision", "index_cache"]}
        config.update(interface=interface,
                      quadrature_degree=degree,
                      point_index=quadrature_index)
        expr, = fem.compile_ufl(integrand, **config)
        if quadrature_index in expr.free_indices:
            expr = gem.IndexSum(expr, quadrature_index)
        return expr
Esempio n. 5
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)
Esempio n. 6
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
Esempio n. 7
0
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()
    kernel_name = "%s_%s_integral_%s" % (prefix, integral_type, integral_data.subdomain_id)
    # Handle negative subdomain_id
    kernel_name = kernel_name.replace("-", "_")

    fiat_cell = as_fiat_cell(cell)
    integration_dim, entity_ids = lower_integral_type(fiat_cell, integral_type)

    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])
    argument_multiindices = tuple(builder.create_element(arg.ufl_element()).get_indices()
                                  for arg in arguments)
    return_variables = builder.set_arguments(arguments, argument_multiindices)

    builder.set_coordinates(mesh)

    builder.set_coefficients(integral_data, form_data)

    # Map from UFL FiniteElement objects to multiindices.  This is
    # so we reuse Index instances when evaluating the same coefficient
    # multiple times with the same table.
    #
    # We also use the same dict for the unconcatenate index cache,
    # which maps index objects to tuples of multiindices.  These two
    # caches shall never conflict as their keys have different types
    # (UFL finite elements vs. GEM index objects).
    index_cache = {}

    kernel_cfg = dict(interface=builder,
                      ufl_cell=cell,
                      integral_type=integral_type,
                      precision=parameters["precision"],
                      integration_dim=integration_dim,
                      entity_ids=entity_ids,
                      argument_multiindices=argument_multiindices,
                      index_cache=index_cache)

    mode_irs = collections.OrderedDict()
    for integral in integral_data.integrals:
        params = parameters.copy()
        params.update(integral.metadata())  # integral metadata overrides
        if params.get("quadrature_rule") == "default":
            del params["quadrature_rule"]

        mode = pick_mode(params["mode"])
        mode_irs.setdefault(mode, collections.OrderedDict())

        integrand = ufl.replace(integral.integrand(), form_data.function_replace_map)
        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:
            quadrature_degree = params["quadrature_degree"]
        except KeyError:
            quadrature_degree = params["estimated_polynomial_degree"]
            functions = list(arguments) + [builder.coordinate(mesh)] + list(integral_data.integral_coefficients)
            function_degrees = [f.ufl_function_space().ufl_element().degree() for f in functions]
            if all((asarray(quadrature_degree) > 10 * asarray(degree)).all()
                   for degree in function_degrees):
                logger.warning("Estimated quadrature degree %s more "
                               "than tenfold greater than any "
                               "argument/coefficient degree (max %s)",
                               quadrature_degree, max_degree(function_degrees))

        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.extend(quadrature_multiindex)

        config = kernel_cfg.copy()
        config.update(quadrature_rule=quad_rule)
        expressions = fem.compile_ufl(integrand,
                                      interior_facet=interior_facet,
                                      **config)
        reps = mode.Integrals(expressions, quadrature_multiindex,
                              argument_multiindices, params)
        for var, rep in zip(return_variables, reps):
            mode_irs[mode].setdefault(var, []).append(rep)

    # Finalise mode representations into a set of assignments
    assignments = []
    for mode, var_reps in mode_irs.items():
        assignments.extend(mode.flatten(var_reps.items(), index_cache))

    if assignments:
        return_variables, expressions = zip(*assignments)
    else:
        return_variables = []
        expressions = []

    # Need optimised roots for COFFEE
    options = dict(reduce(operator.and_,
                          [mode.finalise_options.items()
                           for mode in mode_irs.keys()]))
    expressions = impero_utils.preprocess_gem(expressions, **options)
    assignments = list(zip(return_variables, expressions))

    # Look for cell orientations in the IR
    if builder.needs_cell_orientations(expressions):
        builder.require_cell_orientations()

    # Construct ImperoC
    split_argument_indices = tuple(chain(*[var.index_ordering()
                                           for var in return_variables]))
    index_ordering = tuple(quadrature_indices) + split_argument_indices
    try:
        impero_c = impero_utils.compile_gem(assignments, index_ordering, remove_zeros=True)
    except impero_utils.NoopError:
        # No operations, construct empty kernel
        return builder.construct_empty_kernel(kernel_name)

    # Generate COFFEE
    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)

    # Construct kernel
    body = generate_coffee(impero_c, index_names, parameters["precision"], expressions, split_argument_indices)

    return builder.construct_kernel(kernel_name, body)
Esempio 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)
Esempio n. 9
0
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)
Esempio n. 10
0
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()
    kernel_name = "%s_%s_integral_%s" % (prefix, integral_type,
                                         integral_data.subdomain_id)

    fiat_cell = as_fiat_cell(cell)
    integration_dim, entity_ids = lower_integral_type(fiat_cell, integral_type)

    argument_multiindices = tuple(
        create_element(arg.ufl_element()).get_indices() for arg in arguments)
    argument_indices = tuple(chain(*argument_multiindices))
    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_multiindices)

    coordinates = ufl_utils.coordinate_coefficient(mesh)
    builder.set_coordinates(coordinates)

    builder.set_coefficients(integral_data, form_data)

    # Map from UFL FiniteElement objects to multiindices.  This is
    # so we reuse Index instances when evaluating the same coefficient
    # multiple times with the same table.
    index_cache = {}

    kernel_cfg = dict(interface=builder,
                      ufl_cell=cell,
                      precision=parameters["precision"],
                      integration_dim=integration_dim,
                      entity_ids=entity_ids,
                      argument_multiindices=argument_multiindices,
                      index_cache=index_cache)

    kernel_cfg["facetarea"] = facetarea_generator(mesh, coordinates,
                                                  kernel_cfg, integral_type)
    kernel_cfg["cellvolume"] = cellvolume_generator(mesh, coordinates,
                                                    kernel_cfg)

    mode_irs = collections.OrderedDict()
    for integral in integral_data.integrals:
        params = parameters.copy()
        params.update(integral.metadata())  # integral metadata overrides
        if params.get("quadrature_rule") == "default":
            del params["quadrature_rule"]

        mode = pick_mode(params["mode"])
        mode_irs.setdefault(mode, collections.OrderedDict())

        integrand = ufl_utils.replace_coordinates(integral.integrand(),
                                                  coordinates)
        integrand = ufl.replace(integrand, form_data.function_replace_map)
        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.extend(quadrature_multiindex)

        config = kernel_cfg.copy()
        config.update(quadrature_rule=quad_rule)
        expressions = fem.compile_ufl(integrand,
                                      interior_facet=interior_facet,
                                      **config)
        reps = mode.Integrals(expressions, quadrature_multiindex,
                              argument_multiindices, params)
        for var, rep in zip(return_variables, reps):
            mode_irs[mode].setdefault(var, []).append(rep)

    # Finalise mode representations into a set of assignments
    assignments = []
    for mode, var_reps in iteritems(mode_irs):
        assignments.extend(mode.flatten(viewitems(var_reps)))

    if assignments:
        return_variables, expressions = zip(*assignments)
    else:
        return_variables = []
        expressions = []

    # Need optimised roots for COFFEE
    options = dict(
        reduce(
            operator.and_,
            [viewitems(mode.finalise_options) for mode in iterkeys(mode_irs)]))
    expressions = impero_utils.preprocess_gem(expressions, **options)
    assignments = list(zip(return_variables, expressions))

    # Look for cell orientations in the IR
    if builder.needs_cell_orientations(expressions):
        builder.require_cell_orientations()

    # Construct ImperoC
    index_ordering = tuple(quadrature_indices) + argument_indices
    try:
        impero_c = impero_utils.compile_gem(assignments,
                                            index_ordering,
                                            remove_zeros=True)
    except impero_utils.NoopError:
        # No operations, construct empty kernel
        return builder.construct_empty_kernel(kernel_name)

    # Generate COFFEE
    index_names = [(si, name + str(n))
                   for index, name in zip(argument_multiindices, ['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))

    # Construct kernel
    body = generate_coffee(impero_c, index_names, parameters["precision"],
                           expressions, argument_indices)

    return builder.construct_kernel(kernel_name, body)
Esempio n. 11
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'})
Esempio n. 12
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(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)
    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"])

    retarg = ast.Decl(SCALAR_TYPE, ast.Symbol("R", rank=(Vce.space_dimension(), )))
    local_tensor = coarse_builder.local_tensor
    local_tensor.init = ast.ArrayInit(numpy.zeros(Vce.space_dimension(), dtype=SCALAR_TYPE))
    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)
Esempio 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)
Esempio n. 14
0
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(gem.Index(name=name) for arg, name in zip(arguments, ['j', 'k']))
    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)
    if ufl_utils.is_element_affine(mesh.ufl_coordinate_element()):
        # For affine mesh geometries we prefer code generation that
        # composes well with optimisations.
        builder.set_coordinates(coordinates, mode='list_tensor')
    else:
        # Otherwise we use the approach that might be faster (?)
        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)

        # 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"])
        integration_cell = fiat_cell.construct_subelement(integration_dim)
        quad_rule = params.get("quadrature_rule",
                               create_quadrature(integration_cell,
                                                 quadrature_degree))

        if not isinstance(quad_rule, QuadratureRule):
            raise ValueError("Expected to find a QuadratureRule object, not a %s" %
                             type(quad_rule))

        integrand = ufl_utils.replace_coordinates(integral.integrand(), coordinates)
        integrand = ufl_utils.split_coefficients(integrand, builder.coefficient_split)
        quadrature_index = gem.Index(name='ip')
        quadrature_indices.append(quadrature_index)
        config = kernel_cfg.copy()
        config.update(quadrature_rule=quad_rule, point_index=quadrature_index)
        ir = fem.compile_ufl(integrand, interior_facet=interior_facet, **config)
        if parameters["unroll_indexsum"]:
            ir = opt.unroll_indexsum(ir, max_extent=parameters["unroll_indexsum"])
        irs.append([(gem.IndexSum(expr, quadrature_index)
                     if quadrature_index in expr.free_indices
                     else expr)
                    for expr in 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 = opt.remove_componenttensors(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) + argument_indices,
                                        remove_zeros=True)

    # Generate COFFEE
    index_names = [(index, index.name) for index in argument_indices]
    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, argument_indices)

    kernel_name = "%s_%s_integral_%s" % (prefix, integral_type, integral_data.subdomain_id)
    return builder.construct_kernel(kernel_name, body)