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
