def _expression_mathfunction(expr, ctx): if expr.name.startswith('cyl_bessel_'): # Bessel functions if is_complex(ctx.scalar_type): raise NotImplementedError("Bessel functions for complex numbers: " "missing implementation") nu, arg = expr.children nu_ = expression(nu, ctx) arg_ = expression(arg, ctx) # Modified Bessel functions (C++ only) # # These mappings work for FEniCS only, and fail with Firedrake # since no Boost available. if expr.name in {'cyl_bessel_i', 'cyl_bessel_k'}: name = 'boost::math::' + expr.name return p.Variable(name)(nu_, arg_) else: # cyl_bessel_{jy} -> {jy} name = expr.name[-1:] if nu == gem.Zero(): return p.Variable(f"{name}0")(arg_) elif nu == gem.one: return p.Variable(f"{name}1")(arg_) else: return p.Variable(f"{name}n")(nu_, arg_) else: if expr.name == "ln": name = "log" else: name = expr.name # Not all mathfunctions apply to complex numbers, but this # will be picked up in loopy. This way we allow erf(real(...)) # in complex mode (say). return p.Variable(name)(*(expression(c, ctx) for c in expr.children))
def compile_form(form, prefix="form", parameters=None, interface=None, coffee=True, diagonal=False): """Compiles a UFL form into a set of assembly kernels. :arg form: UFL form :arg prefix: kernel name will start with this string :arg parameters: parameters object :arg coffee: compile coffee kernel instead of loopy kernel :arg diagonal: Are we building a kernel for the diagonal of a rank-2 element tensor? :returns: list of kernels """ cpu_time = time.time() assert isinstance(form, Form) # Determine whether in complex mode: # complex nodes would break the refactoriser. complex_mode = parameters and is_complex(parameters.get("scalar_type")) if complex_mode: logger.warning("Disabling whole expression optimisations" " in GEM for supporting complex mode.") parameters = parameters.copy() parameters["mode"] = 'vanilla' fd = ufl_utils.compute_form_data(form, complex_mode=complex_mode) logger.info(GREEN % "compute_form_data finished in %g seconds.", time.time() - cpu_time) kernels = [] for integral_data in fd.integral_data: start = time.time() kernel = compile_integral(integral_data, fd, prefix, parameters, interface=interface, coffee=coffee, diagonal=diagonal) if kernel is not None: kernels.append(kernel) logger.info(GREEN % "compile_integral finished in %g seconds.", time.time() - start) logger.info(GREEN % "TSFC finished in %g seconds.", time.time() - cpu_time) return kernels
def _expression_mathfunction(expr, ctx): from tsfc.coffee import math_table math_table = math_table.copy() math_table['abs'] = ('abs', 'cabs') complex_mode = int(is_complex(ctx.scalar_type)) # Bessel functions if expr.name.startswith('cyl_bessel_'): if complex_mode: msg = "Bessel functions for complex numbers: missing implementation" raise NotImplementedError(msg) nu, arg = expr.children nu_thunk = lambda: expression(nu, ctx) arg_loopy = expression(arg, ctx) if expr.name == 'cyl_bessel_j': if nu == gem.Zero(): return p.Variable("j0")(arg_loopy) elif nu == gem.one: return p.Variable("j1")(arg_loopy) else: return p.Variable("jn")(nu_thunk(), arg_loopy) if expr.name == 'cyl_bessel_y': if nu == gem.Zero(): return p.Variable("y0")(arg_loopy) elif nu == gem.one: return p.Variable("y1")(arg_loopy) else: return p.Variable("yn")(nu_thunk(), arg_loopy) # Modified Bessel functions (C++ only) # # These mappings work for FEniCS only, and fail with Firedrake # since no Boost available. if expr.name in ['cyl_bessel_i', 'cyl_bessel_k']: name = 'boost::math::' + expr.name return p.Variable(name)(nu_thunk(), arg_loopy) assert False, "Unknown Bessel function: {}".format(expr.name) # Other math functions name = math_table[expr.name][complex_mode] if name is None: raise RuntimeError("{} not supported in complex mode".format(expr.name)) return p.Variable(name)(*[expression(c, ctx) for c in expr.children])
def _expression_mathfunction(expr, parameters): complex_mode = int(is_complex(parameters.scalar_type)) # Bessel functions if expr.name.startswith('cyl_bessel_'): if complex_mode: msg = "Bessel functions for complex numbers: missing implementation" raise NotImplementedError(msg) nu, arg = expr.children nu_thunk = lambda: expression(nu, parameters) arg_coffee = expression(arg, parameters) if expr.name == 'cyl_bessel_j': if nu == gem.Zero(): return coffee.FunCall('j0', arg_coffee) elif nu == gem.one: return coffee.FunCall('j1', arg_coffee) else: return coffee.FunCall('jn', nu_thunk(), arg_coffee) if expr.name == 'cyl_bessel_y': if nu == gem.Zero(): return coffee.FunCall('y0', arg_coffee) elif nu == gem.one: return coffee.FunCall('y1', arg_coffee) else: return coffee.FunCall('yn', nu_thunk(), arg_coffee) # Modified Bessel functions (C++ only) # # These mappings work for FEniCS only, and fail with Firedrake # since no Boost available. if expr.name in ['cyl_bessel_i', 'cyl_bessel_k']: name = 'boost::math::' + expr.name return coffee.FunCall(name, nu_thunk(), arg_coffee) assert False, "Unknown Bessel function: {}".format(expr.name) # Other math functions name = math_table[expr.name][complex_mode] if name is None: raise RuntimeError("{} not supported in complex mode".format( expr.name)) return coffee.FunCall(name, *[expression(c, parameters) for c in expr.children])
def _expression_power(expr, parameters): base, exponent = expr.children complex_mode = int(is_complex(parameters.scalar_type)) return coffee.FunCall(math_table['power'][complex_mode], expression(base, parameters), expression(exponent, parameters))
def complex_mode(self): return is_complex(self.scalar_type)
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'})
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!") # 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 builder = firedrake_interface.ExpressionKernelBuilder( parameters["scalar_type"]) # 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) 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(parameters["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"], parameters["scalar_type"]) # Handle cell orientations if builder.needs_cell_orientations([ir]): builder.require_cell_orientations() # Handle cell sizes (physically mapped elements) if builder.needs_cell_sizes([ir]): builder.require_cell_sizes() # Build kernel tuple return builder.construct_kernel(return_arg, body)
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)