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