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 _calculate_values(function, points, dimension, cell_mask=None): """Calculate function values at given reference points :arg function: function to be sampled :arg points: points to be sampled in reference space :arg cell_mask: Masks for cell node list """ import numpy.ma as ma function_space = function.function_space() keys = {1: (0, ), 2: (0, 0)} # Such tabulation could be done with FInAT, but that would be more # complicated, and there is no clear benefit to changing. fiat_element = create_base_element( function_space.ufl_element()).fiat_equivalent elem = fiat_element.tabulate(0, points)[keys[dimension]] cell_node_list = function_space.cell_node_list if cell_mask is not None: cell_mask = np.tile(cell_mask.reshape(-1, 1), cell_node_list.shape[1]) cell_node_list = ma.compress_rows( ma.masked_array(cell_node_list, mask=cell_mask)) data = function.dat.data_ro[cell_node_list] if function.ufl_shape == (): vec_length = 1 else: vec_length = function.ufl_shape[0] if vec_length == 1: data = np.reshape(data, data.shape + (1, )) return np.einsum("ijk,jl->ilk", data, elem)
def compile_python_kernel(expression, to_pts, to_element, fs, coords): """Produce a :class:`PyOP2.Kernel` wrapping the eval method on the function provided.""" coords_space = coords.function_space() coords_element = create_base_element( coords_space.ufl_element()).fiat_equivalent X_remap = list(coords_element.tabulate(0, to_pts).values())[0] # The par_loop will just pass us arguments, since it doesn't # know about keyword arguments at all so unpack into a dict that we # can pass to the user's eval method. def kernel(output, x, *arguments): kwargs = {} for (slot, _), arg in zip(expression._user_args, arguments): kwargs[slot] = arg X = numpy.dot(X_remap.T, x) for i in range(len(output)): # Pass a slice for the scalar case but just the # current vector in the VFS case. This ensures the # eval method has a Dolfin compatible API. expression.eval( output[i:i + 1, ...] if numpy.ndim(output) == 1 else output[i, ...], X[i:i + 1, ...] if numpy.ndim(X) == 1 else X[i, ...], **kwargs) coefficients = [coords] for _, arg in expression._user_args: coefficients.append(GlobalWrapper(arg)) return kernel, False, False, tuple(coefficients)
def firedrake_local_to_cart(element): r"""Gets the list of nodes for an element (provided they exist.) :arg element: a ufl element. :returns: a list of arrays of floats where each array is a node. """ # TODO: Revise this when FInAT gets dual evaluation fiat_element = create_base_element(element).fiat_equivalent # TODO: Surely there is an easier way that I've forgotten? return [ np.array(list(phi.get_point_dict().keys())[0]) for phi in fiat_element.dual_basis() ]
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, _, name = compile_expression_dual_evaluation(expr, to_element, coffee=False) kernel = op2.Kernel(ast, name, requires_zeroed_output_arguments=True) 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_base_element from firedrake import TestFunction expr = TestFunction(P1) coords = Pk.ufl_domain().coordinates to_element = create_base_element(Pk.ufl_element()) ast, oriented, needs_cell_sizes, coefficients, _ = compile_expression_dual_evaluation(expr, to_element, coords, coffee=False) kernel = op2.Kernel(ast, ast.name) return kernel
def _bezier_calculate_points(function): """Calculate points values for a function used for bezier plotting :arg function: 1D Function with 1 < deg < 4 """ deg = function.function_space().ufl_element().degree() M = np.empty([deg + 1, deg + 1], dtype=float) finat_element = create_base_element( function.function_space().ufl_element()) # TODO: Revise this when FInAT gets dual evaluation basis = finat_element.fiat_equivalent.dual_basis() for i in range(deg + 1): for j in range(deg + 1): M[i, j] = _bernstein( list(basis[j].get_point_dict().keys())[0][0], i, deg) M_inv = np.linalg.inv(M) cell_node_list = function.function_space().cell_node_list return np.dot(function.dat.data_ro[cell_node_list], M_inv)
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