def prolongation_transfer_kernel_action(Vf, expr): from tsfc import compile_expression_dual_evaluation from tsfc.finatinterface import create_base_element to_element = create_base_element(Vf.ufl_element()) ast, oriented, needs_cell_sizes, coefficients, first_coeff_fake_coords, _ = compile_expression_dual_evaluation( expr, to_element, coffee=False) return op2.Kernel(ast, ast.name)
def prolongation_transfer_kernel_action(Vf, expr): from tsfc import compile_expression_dual_evaluation from tsfc.fiatinterface import create_element coords = Vf.ufl_domain().coordinates to_element = create_element(Vf.ufl_element(), vector_is_mixed=False) ast, oriented, needs_cell_sizes, coefficients, _ = compile_expression_dual_evaluation( expr, to_element, coords, coffee=False) return op2.Kernel(ast, ast.name)
def test_ufl_only_simple(): mesh = ufl.Mesh(ufl.VectorElement("P", ufl.triangle, 1)) V = ufl.FunctionSpace(mesh, ufl.FiniteElement("P", ufl.triangle, 2)) v = ufl.Coefficient(V) expr = ufl.inner(v, v) W = V to_element = create_element(W.ufl_element()) ast, oriented, needs_cell_sizes, coefficients, first_coeff_fake_coords, *_ = compile_expression_dual_evaluation( expr, to_element, coffee=False) assert first_coeff_fake_coords is False
def test_ufl_only_to_contravariant_piola(): mesh = ufl.Mesh(ufl.VectorElement("P", ufl.triangle, 1)) V = ufl.FunctionSpace(mesh, ufl.FiniteElement("P", ufl.triangle, 2)) v = ufl.Coefficient(V) expr = ufl.as_vector([v, v]) W = ufl.FunctionSpace(mesh, ufl.FiniteElement("RT", ufl.triangle, 1)) to_element = create_element(W.ufl_element()) ast, oriented, needs_cell_sizes, coefficients, first_coeff_fake_coords, *_ = compile_expression_dual_evaluation( expr, to_element, coffee=False) assert first_coeff_fake_coords is True
def test_ufl_only_spatialcoordinate(): mesh = ufl.Mesh(ufl.VectorElement("P", ufl.triangle, 1)) V = ufl.FunctionSpace(mesh, ufl.FiniteElement("P", ufl.triangle, 2)) x, y = ufl.SpatialCoordinate(mesh) expr = x * y - y**2 + x W = V to_element = create_element(W.ufl_element()) ast, oriented, needs_cell_sizes, coefficients, first_coeff_fake_coords, *_ = compile_expression_dual_evaluation( expr, to_element, coffee=False) assert first_coeff_fake_coords is True
def test_ufl_only_shape_mismatch(): mesh = ufl.Mesh(ufl.VectorElement("P", ufl.triangle, 1)) V = ufl.FunctionSpace(mesh, ufl.FiniteElement("RT", ufl.triangle, 1)) v = ufl.Coefficient(V) expr = ufl.inner(v, v) assert expr.ufl_shape == () W = V to_element = create_element(W.ufl_element()) assert to_element.value_shape == (2, ) with pytest.raises(ValueError): ast, oriented, needs_cell_sizes, coefficients, first_coeff_fake_coords, *_ = compile_expression_dual_evaluation( expr, to_element, coffee=False)
def prolongation_transfer_kernel_action(Vf, expr): from tsfc import compile_expression_dual_evaluation from tsfc.finatinterface import create_element to_element = create_element(Vf.ufl_element()) kernel = compile_expression_dual_evaluation(expr, to_element, Vf.ufl_element()) ast = kernel.ast name = kernel.name flop_count = kernel.flop_count return op2.Kernel(ast, name, requires_zeroed_output_arguments=True, flop_count=flop_count)
def prolongation_transfer_kernel_aij(Pk, P1): # Works for Pk, Pm; I just retain the notation # P1 to remind you that P1 is of lower degree # than Pk from tsfc import compile_expression_dual_evaluation from tsfc.finatinterface import create_base_element from firedrake import TestFunction expr = TestFunction(P1) to_element = create_base_element(Pk.ufl_element()) ast, oriented, needs_cell_sizes, coefficients, first_coeff_fake_coords, _ = compile_expression_dual_evaluation( expr, to_element, coffee=False) kernel = op2.Kernel(ast, ast.name) return kernel
def prolongation_transfer_kernel_aij(Pk, P1): # Works for Pk, Pm; I just retain the notation # P1 to remind you that P1 is of lower degree # than Pk from tsfc import compile_expression_dual_evaluation from tsfc.finatinterface import create_element from firedrake import TestFunction expr = TestFunction(P1) to_element = create_element(Pk.ufl_element()) kernel = compile_expression_dual_evaluation(expr, to_element, Pk.ufl_element()) ast = kernel.ast name = kernel.name flop_count = kernel.flop_count return op2.Kernel(ast, name, requires_zeroed_output_arguments=True, flop_count=flop_count)
def _interpolator(V, tensor, expr, subset, arguments, access): try: to_element = create_base_element(V.ufl_element()) except KeyError: # FInAT only elements raise NotImplementedError( "Don't know how to create FIAT element for %s" % V.ufl_element()) if access is op2.READ: raise ValueError("Can't have READ access for output function") if len(expr.ufl_shape) != len(V.ufl_element().value_shape()): raise RuntimeError( 'Rank mismatch: Expression rank %d, FunctionSpace rank %d' % (len(expr.ufl_shape), len(V.ufl_element().value_shape()))) if expr.ufl_shape != V.ufl_element().value_shape(): raise RuntimeError( 'Shape mismatch: Expression shape %r, FunctionSpace shape %r' % (expr.ufl_shape, V.ufl_element().value_shape())) mesh = V.ufl_domain() coords = mesh.coordinates if not isinstance(expr, firedrake.Expression): if expr.ufl_domain() and expr.ufl_domain() != V.mesh(): raise NotImplementedError( "Interpolation onto another mesh not supported.") ast, oriented, needs_cell_sizes, coefficients, _ = compile_expression_dual_evaluation( expr, to_element, coords, domain=V.mesh(), coffee=False) kernel = op2.Kernel(ast, ast.name, requires_zeroed_output_arguments=True) elif hasattr(expr, "eval"): to_pts = [] for dual in to_element.fiat_equivalent.dual_basis(): if not isinstance(dual, FIAT.functional.PointEvaluation): raise NotImplementedError( "Can only interpolate Python kernels with Lagrange elements" ) pts, = dual.pt_dict.keys() to_pts.append(pts) kernel, oriented, needs_cell_sizes, coefficients = compile_python_kernel( expr, to_pts, to_element, V, coords) else: raise RuntimeError( "Attempting to evaluate an Expression which has no value.") cell_set = coords.cell_set if subset is not None: assert subset.superset == cell_set cell_set = subset parloop_args = [kernel, cell_set] if tensor in set((c.dat for c in coefficients)): output = tensor tensor = op2.Dat(tensor.dataset) if access is not op2.WRITE: copyin = (partial(output.copy, tensor), ) else: copyin = () copyout = (partial(tensor.copy, output), ) else: copyin = () copyout = () if isinstance(tensor, op2.Global): parloop_args.append(tensor(access)) elif isinstance(tensor, op2.Dat): parloop_args.append(tensor(access, V.cell_node_map())) else: assert access == op2.WRITE # Other access descriptors not done for Matrices. parloop_args.append( tensor(op2.WRITE, (V.cell_node_map(), arguments[0].function_space().cell_node_map()))) if oriented: co = mesh.cell_orientations() parloop_args.append(co.dat(op2.READ, co.cell_node_map())) if needs_cell_sizes: cs = mesh.cell_sizes parloop_args.append(cs.dat(op2.READ, cs.cell_node_map())) for coefficient in coefficients: m_ = coefficient.cell_node_map() parloop_args.append(coefficient.dat(op2.READ, m_)) for o in coefficients: domain = o.ufl_domain() if domain is not None and domain.topology != mesh.topology: raise NotImplementedError( "Interpolation onto another mesh not supported.") parloop = op2.ParLoop(*parloop_args).compute if isinstance(tensor, op2.Mat): return parloop, tensor.assemble else: return copyin + (parloop, ) + copyout
def _interpolator(V, tensor, expr, subset, arguments, access): try: expr = ufl.as_ufl(expr) except ufl.UFLException: raise ValueError("Expecting to interpolate a UFL expression") try: to_element = create_element(V.ufl_element()) except KeyError: # FInAT only elements raise NotImplementedError( "Don't know how to create FIAT element for %s" % V.ufl_element()) if access is op2.READ: raise ValueError("Can't have READ access for output function") if len(expr.ufl_shape) != len(V.ufl_element().value_shape()): raise RuntimeError( 'Rank mismatch: Expression rank %d, FunctionSpace rank %d' % (len(expr.ufl_shape), len(V.ufl_element().value_shape()))) if expr.ufl_shape != V.ufl_element().value_shape(): raise RuntimeError( 'Shape mismatch: Expression shape %r, FunctionSpace shape %r' % (expr.ufl_shape, V.ufl_element().value_shape())) # NOTE: The par_loop is always over the target mesh cells. target_mesh = V.ufl_domain() source_mesh = expr.ufl_domain() or target_mesh if target_mesh is not source_mesh: if not isinstance(target_mesh.topology, firedrake.mesh.VertexOnlyMeshTopology): raise NotImplementedError( "Can only interpolate onto a Vertex Only Mesh") if target_mesh.geometric_dimension( ) != source_mesh.geometric_dimension(): raise ValueError( "Cannot interpolate onto a mesh of a different geometric dimension" ) if not hasattr( target_mesh, "_parent_mesh") or target_mesh._parent_mesh is not source_mesh: raise ValueError( "Can only interpolate across meshes where the source mesh is the parent of the target" ) # For trans-mesh interpolation we use a FInAT QuadratureElement as the # (base) target element with runtime point set expressions as their # quadrature rule point set and weights from their dual basis. # NOTE: This setup is useful for thinking about future design - in the # future this `rebuild` function can be absorbed into FInAT as a # transformer that eats an element and gives you an equivalent (which # may or may not be a QuadratureElement) that lets you do run time # tabulation. Alternatively (and this all depends on future design # decision about FInAT how dual evaluation should work) the # to_element's dual basis (which look rather like quadrature rules) can # have their pointset(s) directly replaced with run-time tabulated # equivalent(s) (i.e. finat.point_set.UnknownPointSet(s)) rt_var_name = 'rt_X' to_element = rebuild(to_element, expr, rt_var_name) parameters = {} parameters['scalar_type'] = utils.ScalarType # We need to pass both the ufl element and the finat element # because the finat elements might not have the right mapping # (e.g. L2 Piola, or tensor element with symmetries) # FIXME: for the runtime unknown point set (for cross-mesh # interpolation) we have to pass the finat element we construct # here. Ideally we would only pass the UFL element through. kernel = compile_expression_dual_evaluation(expr, to_element, V.ufl_element(), domain=source_mesh, parameters=parameters) ast = kernel.ast oriented = kernel.oriented needs_cell_sizes = kernel.needs_cell_sizes coefficients = kernel.coefficients first_coeff_fake_coords = kernel.first_coefficient_fake_coords name = kernel.name kernel = op2.Kernel(ast, name, requires_zeroed_output_arguments=True, flop_count=kernel.flop_count) cell_set = target_mesh.cell_set if subset is not None: assert subset.superset == cell_set cell_set = subset parloop_args = [kernel, cell_set] if first_coeff_fake_coords: # Replace with real source mesh coordinates coefficients[0] = source_mesh.coordinates if target_mesh is not source_mesh: # NOTE: TSFC will sometimes drop run-time arguments in generated # kernels if they are deemed not-necessary. # FIXME: Checking for argument name in the inner kernel to decide # whether to add an extra coefficient is a stopgap until # compile_expression_dual_evaluation # (a) outputs a coefficient map to indicate argument ordering in # parloops as `compile_form` does and # (b) allows the dual evaluation related coefficients to be supplied to # them rather than having to be added post-hoc (likely by # replacing `to_element` with a CoFunction/CoArgument as the # target `dual` which would contain `dual` related # coefficient(s)) if rt_var_name in [arg.name for arg in kernel.code[name].args]: # Add the coordinates of the target mesh quadrature points in the # source mesh's reference cell as an extra argument for the inner # loop. (With a vertex only mesh this is a single point for each # vertex cell.) coefficients.append(target_mesh.reference_coordinates) if tensor in set((c.dat for c in coefficients)): output = tensor tensor = op2.Dat(tensor.dataset) if access is not op2.WRITE: copyin = (partial(output.copy, tensor), ) else: copyin = () copyout = (partial(tensor.copy, output), ) else: copyin = () copyout = () if isinstance(tensor, op2.Global): parloop_args.append(tensor(access)) elif isinstance(tensor, op2.Dat): parloop_args.append(tensor(access, V.cell_node_map())) else: assert access == op2.WRITE # Other access descriptors not done for Matrices. rows_map = V.cell_node_map() columns_map = arguments[0].function_space().cell_node_map() if target_mesh is not source_mesh: # Since the par_loop is over the target mesh cells we need to # compose a map that takes us from target mesh cells to the # function space nodes on the source mesh. columns_map = compose_map_and_cache( target_mesh.cell_parent_cell_map, columns_map) parloop_args.append(tensor(op2.WRITE, (rows_map, columns_map))) if oriented: co = target_mesh.cell_orientations() parloop_args.append(co.dat(op2.READ, co.cell_node_map())) if needs_cell_sizes: cs = target_mesh.cell_sizes parloop_args.append(cs.dat(op2.READ, cs.cell_node_map())) for coefficient in coefficients: coeff_mesh = coefficient.ufl_domain() if coeff_mesh is target_mesh or not coeff_mesh: # NOTE: coeff_mesh is None is allowed e.g. when interpolating from # a Real space m_ = coefficient.cell_node_map() elif coeff_mesh is source_mesh: if coefficient.cell_node_map(): # Since the par_loop is over the target mesh cells we need to # compose a map that takes us from target mesh cells to the # function space nodes on the source mesh. m_ = compose_map_and_cache(target_mesh.cell_parent_cell_map, coefficient.cell_node_map()) else: # m_ is allowed to be None when interpolating from a Real space, # even in the trans-mesh case. m_ = coefficient.cell_node_map() else: raise ValueError("Have coefficient with unexpected mesh") parloop_args.append(coefficient.dat(op2.READ, m_)) parloop = op2.ParLoop(*parloop_args) parloop_compute_callable = parloop.compute if isinstance(tensor, op2.Mat): return parloop_compute_callable, tensor.assemble else: return copyin + (parloop_compute_callable, ) + copyout