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
