def make_extruded_coords(extruded_topology,
base_coords,
ext_coords,
layer_height,
extrusion_type='uniform',
kernel=None):
"""
Given either a kernel or a (fixed) layer_height, compute an
extruded coordinate field for an extruded mesh.
:arg extruded_topology: an :class:`~.ExtrudedMeshTopology` to extrude
a coordinate field for.
:arg base_coords: a :class:`~.Function` to read the base
coordinates from.
:arg ext_coords: a :class:`~.Function` to write the extruded
coordinates into.
:arg layer_height: the height for each layer. Either a scalar,
where layers will be equi-spaced at the specified height, or a
1D array of variable layer heights to use through the extrusion.
:arg extrusion_type: the type of extrusion to use. Predefined
options are either "uniform" (creating equi-spaced layers by
extruding in the (n+1)dth direction), "radial" (creating
equi-spaced layers by extruding in the outward direction from
the origin) or "radial_hedgehog" (creating equi-spaced layers
by extruding coordinates in the outward cell-normal
direction, needs a P1dgxP1 coordinate field).
:arg kernel: an optional kernel to carry out coordinate extrusion.
The kernel signature (if provided) is::
void kernel(double **base_coords, double **ext_coords,
double *layer_height, int layer)
The kernel iterates over the cells of the mesh and receives as
arguments the coordinates of the base cell (to read), the
coordinates on the extruded cell (to write to), the fixed layer
height, and the current cell layer.
"""
_, vert_space = ext_coords.function_space().ufl_element().sub_elements(
)[0].sub_elements()
if kernel is None and not (vert_space.degree() == 1
and vert_space.family()
in ['Lagrange', 'Discontinuous Lagrange']):
raise RuntimeError(
'Extrusion of coordinates is only possible for a P1 or P1dg interval unless a custom kernel is provided'
)
layer_height = numpy.atleast_1d(numpy.array(layer_height, dtype=RealType))
if layer_height.ndim > 1:
raise RuntimeError('Extrusion layer height should be 1d or scalar')
if layer_height.size > 1:
layer_height = numpy.cumsum(numpy.concatenate(([0], layer_height)))
layer_heights = layer_height.size
layer_height = op2.Global(layer_heights, layer_height, dtype=RealType)
if kernel is not None:
op2.ParLoop(kernel,
ext_coords.cell_set,
ext_coords.dat(op2.WRITE, ext_coords.cell_node_map()),
base_coords.dat(op2.READ, base_coords.cell_node_map()),
layer_height(op2.READ),
pass_layer_arg=True,
is_loopy_kernel=True).compute()
return
ext_fe = create_element(ext_coords.ufl_element())
ext_shape = ext_fe.index_shape
base_fe = create_element(base_coords.ufl_element())
base_shape = base_fe.index_shape
data = []
data.append(lp.GlobalArg("ext_coords", dtype=ScalarType, shape=ext_shape))
data.append(lp.GlobalArg("base_coords", dtype=ScalarType,
shape=base_shape))
data.append(
lp.GlobalArg("layer_height", dtype=RealType, shape=(layer_heights, )))
data.append(lp.ValueArg('layer'))
base_coord_dim = base_coords.function_space().value_size
# Deal with tensor product cells
adim = len(ext_shape) - 2
# handle single or variable layer heights
if layer_heights == 1:
height_var = "layer_height[0] * (layer + l)"
else:
height_var = "layer_height[layer + l]"
def _get_arity_axis_inames(_base):
return tuple(_base + str(i) for i in range(adim))
def _get_lp_domains(_inames, _extents):
domains = []
for idx, extent in zip(_inames, _extents):
inames = isl.make_zero_and_vars([idx])
domains.append(((inames[0].le_set(inames[idx])) &
(inames[idx].lt_set(inames[0] + extent))))
return domains
if extrusion_type == 'uniform':
domains = []
dd = _get_arity_axis_inames('d')
domains.extend(_get_lp_domains(dd, ext_shape[:adim]))
domains.extend(_get_lp_domains(('c', ), (base_coord_dim, )))
if layer_heights == 1:
domains.extend(_get_lp_domains(('l', ), (2, )))
else:
domains.append(
"[layer] -> { [l] : 0 <= l <= 1 & 0 <= l + layer < %d}" %
layer_heights)
instructions = """
ext_coords[{dd}, l, c] = base_coords[{dd}, c]
ext_coords[{dd}, l, {base_coord_dim}] = ({hv})
""".format(dd=', '.join(dd),
base_coord_dim=base_coord_dim,
hv=height_var)
name = "pyop2_kernel_uniform_extrusion"
elif extrusion_type == 'radial':
domains = []
dd = _get_arity_axis_inames('d')
domains.extend(_get_lp_domains(dd, ext_shape[:adim]))
domains.extend(_get_lp_domains(('c', 'k'), (base_coord_dim, ) * 2))
if layer_heights == 1:
domains.extend(_get_lp_domains(('l', ), (2, )))
else:
domains.append(
"[layer] -> { [l] : 0 <= l <= 1 & 0 <= l + layer < %d}" %
layer_heights)
instructions = """
<{RealType}> tt[{dd}] = 0
<{RealType}> bc[{dd}] = 0
for k
bc[{dd}] = real(base_coords[{dd}, k])
tt[{dd}] = tt[{dd}] + bc[{dd}] * bc[{dd}]
end
tt[{dd}] = sqrt(tt[{dd}])
ext_coords[{dd}, l, c] = base_coords[{dd}, c] + base_coords[{dd}, c] * ({hv}) / tt[{dd}]
""".format(RealType=RealType, dd=', '.join(dd), hv=height_var)
name = "pyop2_kernel_radial_extrusion"
elif extrusion_type == 'radial_hedgehog':
# Only implemented for interval in 2D and triangle in 3D.
# gdim != tdim already checked in ExtrudedMesh constructor.
tdim = base_coords.ufl_domain().ufl_cell().topological_dimension()
if tdim not in [1, 2]:
raise NotImplementedError(
"Hedgehog extrusion not implemented for %s" %
base_coords.ufl_domain().ufl_cell())
# tdim == 1:
#
# normal is:
# (0 -1) (x2 - x1)
# (1 0) (y2 - y1)
#
# tdim == 2:
# normal is
# v0 x v1
#
# /\
# v0/ \
# / \
# /------\
# v1
domains = []
dd = _get_arity_axis_inames('d')
_dd = _get_arity_axis_inames('_d')
domains.extend(_get_lp_domains(dd, ext_shape[:adim]))
domains.extend(_get_lp_domains(_dd, ext_shape[:adim]))
domains.extend(
_get_lp_domains(('c0', 'c1', 'c2', 'c3', 'k', 'l'),
(base_coord_dim, ) * 5 + (2, )))
# Formula for normal, n
n_1_1 = """
n[0] = -bc[1, 1] + bc[0, 1]
n[1] = bc[1, 0] - bc[0, 0]
"""
n_2_1 = """
v0[c3] = bc[1, c3] - bc[0, c3]
v1[c3] = bc[2, c3] - bc[0, c3]
n[0] = v0[1] * v1[2] - v0[2] * v1[1]
n[1] = v0[2] * v1[0] - v0[0] * v1[2]
n[2] = v0[0] * v1[1] - v0[1] * v1[0]
"""
n_2_2 = """
v0[c3] = bc[0, 1, c3] - bc[0, 0, c3]
v1[c3] = bc[1, 0, c3] - bc[0, 0, c3]
n[0] = v0[1] * v1[2] - v0[2] * v1[1]
n[1] = v0[2] * v1[0] - v0[0] * v1[2]
n[2] = v0[0] * v1[1] - v0[1] * v1[0]
"""
n_dict = {1: {1: n_1_1}, 2: {1: n_2_1, 2: n_2_2}}
instructions = """
<{RealType}> dot = 0
<{RealType}> norm = 0
<{RealType}> v0[c2] = 0
<{RealType}> v1[c2] = 0
<{RealType}> n[c2] = 0
<{RealType}> x[c2] = 0
<{RealType}> bc[{_dd}, c1] = real(base_coords[{_dd}, c1])
for {_dd}
x[c1] = x[c1] + bc[{_dd}, c1]
end
{ninst}
for k
dot = dot + x[k] * n[k]
norm = norm + n[k] * n[k]
end
norm = sqrt(norm)
norm = -norm if dot < 0 else norm
ext_coords[{dd}, l, c0] = base_coords[{dd}, c0] + n[c0] * ({hv}) / norm
""".format(RealType=RealType,
dd=', '.join(dd),
_dd=', '.join(_dd),
ninst=n_dict[tdim][adim],
hv=height_var)
name = "pyop2_kernel_radial_hedgehog_extrusion"
else:
raise NotImplementedError('Unsupported extrusion type "%s"' %
extrusion_type)
ast = lp.make_function(domains,
instructions,
data,
name=name,
target=lp.CTarget(),
seq_dependencies=True,
silenced_warnings=["summing_if_branches_ops"])
kernel = op2.Kernel(ast, name)
op2.ParLoop(kernel,
ext_coords.cell_set,
ext_coords.dat(op2.WRITE, ext_coords.cell_node_map()),
base_coords.dat(op2.READ, base_coords.cell_node_map()),
layer_height(op2.READ),
pass_layer_arg=True,
is_loopy_kernel=True).compute()
def _interpolator(V, dat, expr, subset, access):
to_element = create_element(V.ufl_element(), vector_is_mixed=False)
to_pts = []
if access is op2.READ:
raise ValueError("Can't have READ access for output function")
if V.ufl_element().mapping() != "identity":
raise NotImplementedError("Can only interpolate onto elements "
"with affine mapping. Try projecting instead")
for dual in to_element.dual_basis():
if not isinstance(dual, FIAT.functional.PointEvaluation):
raise NotImplementedError("Can only interpolate onto point "
"evaluation operators. Try projecting instead")
pts, = dual.pt_dict.keys()
to_pts.append(pts)
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.")
if expr.ufl_shape != V.shape:
raise ValueError("UFL expression has incorrect shape for interpolation.")
ast, oriented, needs_cell_sizes, coefficients, _ = compile_ufl_kernel(expr, to_pts, coords, coffee=False)
kernel = op2.Kernel(ast, ast.name)
elif hasattr(expr, "eval"):
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
args = [kernel, cell_set]
if dat in set((c.dat for c in coefficients)):
output = dat
dat = op2.Dat(dat.dataset)
if access is not op2.WRITE:
copyin = (partial(output.copy, dat), )
else:
copyin = ()
copyout = (partial(dat.copy, output), )
else:
copyin = ()
copyout = ()
args.append(dat(access, V.cell_node_map()))
if oriented:
co = mesh.cell_orientations()
args.append(co.dat(op2.READ, co.cell_node_map()))
if needs_cell_sizes:
cs = mesh.cell_sizes
args.append(cs.dat(op2.READ, cs.cell_node_map()))
for coefficient in coefficients:
m_ = coefficient.cell_node_map()
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.")
return copyin + (op2.ParLoop(*args).compute, ) + copyout
def _assemble(f, tensor=None, bcs=None, form_compiler_parameters=None,
inverse=False, mat_type=None, sub_mat_type=None,
appctx={},
options_prefix=None,
assemble_now=False,
allocate_only=False,
zero_tensor=True,
diagonal=False):
r"""Assemble the form or Slate expression f and return a Firedrake object
representing the result. This will be a :class:`float` for 0-forms/rank-0
Slate tensors, a :class:`.Function` for 1-forms/rank-1 Slate tensors and
a :class:`.Matrix` for 2-forms/rank-2 Slate tensors.
:arg bcs: A tuple of :class`.DirichletBC`\s and/or :class`.EquationBCSplit`\s to be applied.
:arg tensor: An existing tensor object into which the form should be
assembled. If this is not supplied, a new tensor will be created for
the purpose.
:arg form_compiler_parameters: (optional) dict of parameters to pass to
the form compiler.
:arg inverse: (optional) if f is a 2-form, then assemble the inverse
of the local matrices.
:arg mat_type: (optional) type for assembled matrices, one of
"nest", "aij", "baij", or "matfree".
:arg sub_mat_type: (optional) type for assembled sub matrices
inside a "nest" matrix. One of "aij" or "baij".
:arg appctx: Additional information to hang on the assembled
matrix if an implicit matrix is requested (mat_type "matfree").
:arg options_prefix: An options prefix for the PETSc matrix
(ignored if not assembling a bilinear form).
"""
if mat_type is None:
mat_type = parameters.parameters["default_matrix_type"]
if mat_type not in ["matfree", "aij", "baij", "nest"]:
raise ValueError("Unrecognised matrix type, '%s'" % mat_type)
if sub_mat_type is None:
sub_mat_type = parameters.parameters["default_sub_matrix_type"]
if sub_mat_type not in ["aij", "baij"]:
raise ValueError("Invalid submatrix type, '%s' (not 'aij' or 'baij')", sub_mat_type)
if form_compiler_parameters:
form_compiler_parameters = form_compiler_parameters.copy()
else:
form_compiler_parameters = {}
form_compiler_parameters["assemble_inverse"] = inverse
topology = f.ufl_domains()[0].topology
for m in f.ufl_domains():
# Ensure mesh is "initialised" (could have got here without
# building a functionspace (e.g. if integrating a constant)).
m.init()
if m.topology != topology:
raise NotImplementedError("All integration domains must share a mesh topology.")
for o in chain(f.arguments(), f.coefficients()):
domain = o.ufl_domain()
if domain is not None and domain.topology != topology:
raise NotImplementedError("Assembly with multiple meshes not supported.")
if isinstance(f, slate.TensorBase):
if diagonal:
raise NotImplementedError("Diagonal + slate not supported")
kernels = slac.compile_expression(f, tsfc_parameters=form_compiler_parameters)
integral_types = [kernel.kinfo.integral_type for kernel in kernels]
else:
kernels = tsfc_interface.compile_form(f, "form", parameters=form_compiler_parameters, inverse=inverse, diagonal=diagonal)
integral_types = [integral.integral_type() for integral in f.integrals()]
if bcs is not None:
for bc in bcs:
integral_types += [integral.integral_type() for integral in bc.integrals()]
rank = len(f.arguments())
if diagonal:
assert rank == 2
is_mat = rank == 2 and not diagonal
is_vec = rank == 1 or diagonal
if any((coeff.function_space() and coeff.function_space().component is not None)
for coeff in f.coefficients()):
raise NotImplementedError("Integration of subscripted VFS not yet implemented")
if inverse and rank != 2:
raise ValueError("Can only assemble the inverse of a 2-form")
zero_tensor_parloop = lambda: None
if is_mat:
matfree = mat_type == "matfree"
nest = mat_type == "nest"
if nest:
baij = sub_mat_type == "baij"
else:
baij = mat_type == "baij"
# intercept matrix-free matrices here
if matfree:
if inverse:
raise NotImplementedError("Inverse not implemented with matfree")
if diagonal:
raise NotImplementedError("Diagonal not implemented with matfree")
if tensor is None:
result_matrix = matrix.ImplicitMatrix(f, bcs,
fc_params=form_compiler_parameters,
appctx=appctx,
options_prefix=options_prefix)
yield lambda: result_matrix
return
if not isinstance(tensor, matrix.ImplicitMatrix):
raise ValueError("Expecting implicit matrix with matfree")
tensor.assemble()
yield lambda: tensor
return
test, trial = f.arguments()
map_pairs = []
cell_domains = []
exterior_facet_domains = []
interior_facet_domains = []
if tensor is None:
# For horizontal facets of extruded meshes, the corresponding domain
# in the base mesh is the cell domain. Hence all the maps used for top
# bottom and interior horizontal facets will use the cell to dofs map
# coming from the base mesh as a starting point for the actual dynamic map
# computation.
for integral_type in integral_types:
if integral_type == "cell":
cell_domains.append(op2.ALL)
elif integral_type == "exterior_facet":
exterior_facet_domains.append(op2.ALL)
elif integral_type == "interior_facet":
interior_facet_domains.append(op2.ALL)
elif integral_type == "exterior_facet_bottom":
cell_domains.append(op2.ON_BOTTOM)
elif integral_type == "exterior_facet_top":
cell_domains.append(op2.ON_TOP)
elif integral_type == "exterior_facet_vert":
exterior_facet_domains.append(op2.ALL)
elif integral_type == "interior_facet_horiz":
cell_domains.append(op2.ON_INTERIOR_FACETS)
elif integral_type == "interior_facet_vert":
interior_facet_domains.append(op2.ALL)
else:
raise ValueError('Unknown integral type "%s"' % integral_type)
# Used for the sparsity construction
iteration_regions = []
if cell_domains:
map_pairs.append((test.cell_node_map(), trial.cell_node_map()))
iteration_regions.append(tuple(cell_domains))
if exterior_facet_domains:
map_pairs.append((test.exterior_facet_node_map(), trial.exterior_facet_node_map()))
iteration_regions.append(tuple(exterior_facet_domains))
if interior_facet_domains:
map_pairs.append((test.interior_facet_node_map(), trial.interior_facet_node_map()))
iteration_regions.append(tuple(interior_facet_domains))
map_pairs = tuple(map_pairs)
# Construct OP2 Mat to assemble into
fs_names = (test.function_space().name, trial.function_space().name)
try:
sparsity = op2.Sparsity((test.function_space().dof_dset,
trial.function_space().dof_dset),
map_pairs,
iteration_regions=iteration_regions,
name="%s_%s_sparsity" % fs_names,
nest=nest,
block_sparse=baij)
except SparsityFormatError:
raise ValueError("Monolithic matrix assembly is not supported for systems with R-space blocks.")
result_matrix = matrix.Matrix(f, bcs, mat_type, sparsity, ScalarType,
"%s_%s_matrix" % fs_names,
options_prefix=options_prefix)
tensor = result_matrix.M
else:
if isinstance(tensor, matrix.ImplicitMatrix):
raise ValueError("Expecting matfree with implicit matrix")
result_matrix = tensor
tensor = tensor.M
zero_tensor_parloop = tensor.zero
if result_matrix.block_shape != (1, 1) and mat_type == "baij":
raise ValueError("BAIJ matrix type makes no sense for mixed spaces, use 'aij'")
def mat(testmap, trialmap, rowbc, colbc, i, j):
m = testmap(test.function_space()[i])
n = trialmap(trial.function_space()[j])
maps = (m if m else None, n if n else None)
rlgmap, clgmap = tensor[i, j].local_to_global_maps
V = test.function_space()[i]
rlgmap = V.local_to_global_map(rowbc, lgmap=rlgmap)
V = trial.function_space()[j]
clgmap = V.local_to_global_map(colbc, lgmap=clgmap)
if rowbc is None:
rowbc = []
if colbc is None:
colbc = []
unroll = any(bc.function_space().component is not None
for bc in chain(rowbc, colbc) if bc is not None)
return tensor[i, j](op2.INC, maps, lgmaps=(rlgmap, clgmap), unroll_map=unroll)
result = lambda: result_matrix
if allocate_only:
yield result
return
elif is_vec:
test = f.arguments()[0]
if tensor is None:
result_function = function.Function(test.function_space())
tensor = result_function.dat
else:
result_function = tensor
tensor = result_function.dat
zero_tensor_parloop = tensor.zero
def vec(testmap, i):
_testmap = testmap(test.function_space()[i])
return tensor[i](op2.INC, _testmap if _testmap else None)
result = lambda: result_function
else:
# 0-forms are always scalar
if tensor is None:
tensor = op2.Global(1, [0.0])
else:
raise ValueError("Can't assemble 0-form into existing tensor")
result = lambda: tensor.data[0]
coefficients = f.coefficients()
domains = f.ufl_domains()
# These will be used to correctly interpret the "otherwise"
# subdomain
all_integer_subdomain_ids = defaultdict(list)
for k in kernels:
if k.kinfo.subdomain_id != "otherwise":
all_integer_subdomain_ids[k.kinfo.integral_type].append(k.kinfo.subdomain_id)
for k, v in all_integer_subdomain_ids.items():
all_integer_subdomain_ids[k] = tuple(sorted(v))
# In collecting loops mode, we collect the loops, and assume the
# boundary conditions provided are the ones we want. It therefore
# is only used inside residual and jacobian assembly.
if zero_tensor:
yield zero_tensor_parloop
for indices, kinfo in kernels:
kernel = kinfo.kernel
integral_type = kinfo.integral_type
domain_number = kinfo.domain_number
subdomain_id = kinfo.subdomain_id
coeff_map = kinfo.coefficient_map
pass_layer_arg = kinfo.pass_layer_arg
needs_orientations = kinfo.oriented
needs_cell_facets = kinfo.needs_cell_facets
needs_cell_sizes = kinfo.needs_cell_sizes
m = domains[domain_number]
subdomain_data = f.subdomain_data()[m]
# Find argument space indices
if is_mat:
i, j = indices
elif is_vec:
i, = indices
else:
assert len(indices) == 0
sdata = subdomain_data.get(integral_type, None)
if integral_type != 'cell' and sdata is not None:
raise NotImplementedError("subdomain_data only supported with cell integrals.")
# Extract block from tensor and test/trial spaces
# FIXME Ugly variable renaming required because functions are not
# lexical closures in Python and we're writing to these variables
if is_mat:
if bcs is not None:
tsbc = list(bc for bc in chain(*bcs))
if result_matrix.block_shape > (1, 1):
trbc = [bc for bc in tsbc if bc.function_space_index() == j and isinstance(bc, DirichletBC)]
tsbc = [bc for bc in tsbc if bc.function_space_index() == i]
else:
trbc = [bc for bc in tsbc if isinstance(bc, DirichletBC)]
else:
tsbc = []
trbc = []
# Now build arguments for the par_loop
kwargs = {}
# Some integrals require non-coefficient arguments at the
# end (facet number information).
extra_args = []
# Decoration for applying to matrix maps in extruded case
decoration = None
itspace = m.measure_set(integral_type, subdomain_id,
all_integer_subdomain_ids)
if integral_type == "cell":
itspace = sdata or itspace
if subdomain_id not in ["otherwise", "everywhere"] and sdata is not None:
raise ValueError("Cannot use subdomain data and subdomain_id")
def get_map(x):
return x.cell_node_map()
elif integral_type in ("exterior_facet", "exterior_facet_vert"):
extra_args.append(m.exterior_facets.local_facet_dat(op2.READ))
def get_map(x):
return x.exterior_facet_node_map()
elif integral_type in ("exterior_facet_top", "exterior_facet_bottom"):
# In the case of extruded meshes with horizontal facet integrals, two
# parallel loops will (potentially) get created and called based on the
# domain id: interior horizontal, bottom or top.
decoration = {"exterior_facet_top": op2.ON_TOP,
"exterior_facet_bottom": op2.ON_BOTTOM}[integral_type]
kwargs["iterate"] = decoration
def get_map(x):
return x.cell_node_map()
elif integral_type in ("interior_facet", "interior_facet_vert"):
extra_args.append(m.interior_facets.local_facet_dat(op2.READ))
def get_map(x):
return x.interior_facet_node_map()
elif integral_type == "interior_facet_horiz":
decoration = op2.ON_INTERIOR_FACETS
kwargs["iterate"] = decoration
def get_map(x):
return x.cell_node_map()
else:
raise ValueError("Unknown integral type '%s'" % integral_type)
# Output argument
if is_mat:
tensor_arg = mat(lambda s: get_map(s),
lambda s: get_map(s),
tsbc, trbc,
i, j)
elif is_vec:
tensor_arg = vec(lambda s: get_map(s), i)
else:
tensor_arg = tensor(op2.INC)
coords = m.coordinates
args = [kernel, itspace, tensor_arg,
coords.dat(op2.READ, get_map(coords))]
if needs_orientations:
o = m.cell_orientations()
args.append(o.dat(op2.READ, get_map(o)))
if needs_cell_sizes:
o = m.cell_sizes
args.append(o.dat(op2.READ, get_map(o)))
for n in coeff_map:
c = coefficients[n]
for c_ in c.split():
m_ = get_map(c_)
args.append(c_.dat(op2.READ, m_))
if needs_cell_facets:
assert integral_type == "cell"
extra_args.append(m.cell_to_facets(op2.READ))
args.extend(extra_args)
kwargs["pass_layer_arg"] = pass_layer_arg
try:
yield op2.ParLoop(*args, **kwargs).compute
except MapValueError:
raise RuntimeError("Integral measure does not match measure of all coefficients/arguments")
# Must apply bcs outside loop over kernels because we may wish
# to apply bcs to a block which is otherwise zero, and
# therefore does not have an associated kernel.
if bcs is not None and is_mat:
for bc in bcs:
if isinstance(bc, DirichletBC):
fs = bc.function_space()
# Evaluate this outwith a "collecting_loops" block,
# since creation of the bc nodes actually can create a
# par_loop.
nodes = bc.nodes
if len(fs) > 1:
raise RuntimeError(r"""Cannot apply boundary conditions to full mixed space. Did you forget to index it?""")
shape = result_matrix.block_shape
for i in range(shape[0]):
for j in range(shape[1]):
# Set diagonal entries on bc nodes to 1 if the current
# block is on the matrix diagonal and its index matches the
# index of the function space the bc is defined on.
if i != j:
continue
if fs.component is None and fs.index is not None:
# Mixed, index (no ComponentFunctionSpace)
if fs.index == i:
yield functools.partial(tensor[i, j].set_local_diagonal_entries, nodes)
elif fs.component is not None:
# ComponentFunctionSpace, check parent index
if fs.parent.index is not None:
# Mixed, index doesn't match
if fs.parent.index != i:
continue
# Index matches
yield functools.partial(tensor[i, j].set_local_diagonal_entries, nodes, idx=fs.component)
elif fs.index is None:
yield functools.partial(tensor[i, j].set_local_diagonal_entries, nodes)
else:
raise RuntimeError("Unhandled BC case")
elif isinstance(bc, EquationBCSplit):
yield from _assemble(bc.f, tensor=result_matrix, bcs=bc.bcs,
form_compiler_parameters=form_compiler_parameters,
inverse=inverse, mat_type=mat_type,
sub_mat_type=sub_mat_type,
appctx=appctx,
assemble_now=assemble_now,
allocate_only=False,
zero_tensor=False)
else:
raise NotImplementedError("Undefined type of bcs class provided.")
if bcs is not None and is_vec:
for bc in bcs:
if isinstance(bc, DirichletBC):
if assemble_now:
if diagonal:
yield functools.partial(bc.set, result_function, 1)
else:
yield functools.partial(bc.apply, result_function)
else:
yield functools.partial(bc.zero, result_function)
elif isinstance(bc, EquationBCSplit):
if diagonal:
raise NotImplementedError("diagonal assembly and EquationBC not supported")
yield functools.partial(bc.zero, result_function)
yield from _assemble(bc.f, tensor=result_function, bcs=bc.bcs,
form_compiler_parameters=form_compiler_parameters,
inverse=inverse, mat_type=mat_type,
sub_mat_type=sub_mat_type,
appctx=appctx,
assemble_now=assemble_now,
allocate_only=False,
zero_tensor=False)
if zero_tensor:
if is_mat:
# Queue up matrix assembly (after we've done all the other operations)
yield tensor.assemble
if assemble_now:
yield result
def _make_parloops(expr, tensor, bcs, diagonal, fc_params, assembly_rank):
"""Create parloops for the assembly of the expression.
:arg expr: The expression to be assembled.
:arg tensor: The tensor to write to. Depending on ``expr`` and ``diagonal``
this will either be a scalar (:class:`~pyop2.op2.Global`),
vector/cofunction (masquerading as a :class:`.Function`) or :class:`.Matrix`.
:arg bcs: Iterable of boundary conditions.
:arg diagonal: (:class:`bool`) If assembling a matrix is it diagonal?
:arg fc_params: Dictionary of parameters to pass to the form compiler.
:arg assembly_rank: The appropriate :class:`_AssemblyRank`.
:returns: A tuple of the generated :class:`~pyop2..op2.ParLoop` objects.
"""
if fc_params:
form_compiler_parameters = fc_params.copy()
else:
form_compiler_parameters = {}
try:
topology, = set(d.topology for d in expr.ufl_domains())
except ValueError:
raise NotImplementedError(
"All integration domains must share a mesh topology")
for m in expr.ufl_domains():
# Ensure mesh is "initialised" (could have got here without
# building a functionspace (e.g. if integrating a constant)).
m.init()
for o in chain(expr.arguments(), expr.coefficients()):
domain = o.ufl_domain()
if domain is not None and domain.topology != topology:
raise NotImplementedError(
"Assembly with multiple meshes not supported.")
if assembly_rank == _AssemblyRank.MATRIX:
test, trial = expr.arguments()
create_op2arg = functools.partial(_matrix_arg,
all_bcs=tuple(chain(*bcs)),
matrix=tensor,
Vrow=test.function_space(),
Vcol=trial.function_space())
elif assembly_rank == _AssemblyRank.VECTOR:
if diagonal:
# actually a 2-form but throw away the trial space
test, _ = expr.arguments()
else:
test, = expr.arguments()
create_op2arg = functools.partial(_vector_arg,
function=tensor,
V=test.function_space())
else:
create_op2arg = tensor
coefficients = expr.coefficients()
domains = expr.ufl_domains()
if isinstance(expr, slate.TensorBase):
kernels = slac.compile_expression(
expr, compiler_parameters=form_compiler_parameters)
else:
kernels = tsfc_interface.compile_form(
expr,
"form",
parameters=form_compiler_parameters,
diagonal=diagonal)
# These will be used to correctly interpret the "otherwise"
# subdomain
all_integer_subdomain_ids = defaultdict(list)
for k in kernels:
if k.kinfo.subdomain_id != "otherwise":
all_integer_subdomain_ids[k.kinfo.integral_type].append(
k.kinfo.subdomain_id)
for k, v in all_integer_subdomain_ids.items():
all_integer_subdomain_ids[k] = tuple(sorted(v))
parloops = []
for indices, kinfo in kernels:
kernel = kinfo.kernel
integral_type = kinfo.integral_type
domain_number = kinfo.domain_number
subdomain_id = kinfo.subdomain_id
coeff_map = kinfo.coefficient_map
pass_layer_arg = kinfo.pass_layer_arg
needs_orientations = kinfo.oriented
needs_cell_facets = kinfo.needs_cell_facets
needs_cell_sizes = kinfo.needs_cell_sizes
m = domains[domain_number]
subdomain_data = expr.subdomain_data()[m]
# Find argument space indices
if assembly_rank == _AssemblyRank.MATRIX:
i, j = indices
elif assembly_rank == _AssemblyRank.VECTOR:
i, = indices
else:
assert len(indices) == 0
sdata = subdomain_data.get(integral_type, None)
if integral_type != 'cell' and sdata is not None:
raise NotImplementedError(
"subdomain_data only supported with cell integrals.")
# Now build arguments for the par_loop
kwargs = {}
# Some integrals require non-coefficient arguments at the
# end (facet number information).
extra_args = []
itspace = m.measure_set(integral_type, subdomain_id,
all_integer_subdomain_ids)
if integral_type == "cell":
itspace = sdata or itspace
if subdomain_id not in ["otherwise", "everywhere"
] and sdata is not None:
raise ValueError("Cannot use subdomain data and subdomain_id")
def get_map(x):
return x.cell_node_map()
elif integral_type in ("exterior_facet", "exterior_facet_vert"):
extra_args.append(m.exterior_facets.local_facet_dat(op2.READ))
def get_map(x):
return x.exterior_facet_node_map()
elif integral_type in ("exterior_facet_top", "exterior_facet_bottom"):
# In the case of extruded meshes with horizontal facet integrals, two
# parallel loops will (potentially) get created and called based on the
# domain id: interior horizontal, bottom or top.
kwargs["iterate"] = {
"exterior_facet_top": op2.ON_TOP,
"exterior_facet_bottom": op2.ON_BOTTOM
}[integral_type]
def get_map(x):
return x.cell_node_map()
elif integral_type in ("interior_facet", "interior_facet_vert"):
extra_args.append(m.interior_facets.local_facet_dat(op2.READ))
def get_map(x):
return x.interior_facet_node_map()
elif integral_type == "interior_facet_horiz":
kwargs["iterate"] = op2.ON_INTERIOR_FACETS
def get_map(x):
return x.cell_node_map()
else:
raise ValueError("Unknown integral type '%s'" % integral_type)
# Output argument
if assembly_rank == _AssemblyRank.MATRIX:
tensor_arg = create_op2arg(op2.INC, get_map, i, j)
elif assembly_rank == _AssemblyRank.VECTOR:
tensor_arg = create_op2arg(op2.INC, get_map, i)
else:
tensor_arg = create_op2arg(op2.INC)
coords = m.coordinates
args = [
kernel, itspace, tensor_arg,
coords.dat(op2.READ, get_map(coords))
]
if needs_orientations:
o = m.cell_orientations()
args.append(o.dat(op2.READ, get_map(o)))
if needs_cell_sizes:
o = m.cell_sizes
args.append(o.dat(op2.READ, get_map(o)))
for n, split_map in coeff_map:
c = coefficients[n]
split_c = c.split()
for c_ in (split_c[i] for i in split_map):
m_ = get_map(c_)
args.append(c_.dat(op2.READ, m_))
if needs_cell_facets:
assert integral_type == "cell"
extra_args.append(m.cell_to_facets(op2.READ))
if pass_layer_arg:
c = op2.Global(1,
itspace.layers - 2,
dtype=numpy.dtype(numpy.int32))
o = c(op2.READ)
extra_args.append(o)
args.extend(extra_args)
kwargs["pass_layer_arg"] = pass_layer_arg
try:
parloops.append(op2.ParLoop(*args, **kwargs))
except MapValueError:
raise RuntimeError(
"Integral measure does not match measure of all coefficients/arguments"
)
return tuple(parloops)
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 create_parloops(expr,
create_op2arg,
*,
assembly_rank=None,
diagonal=False,
form_compiler_parameters=None):
"""Create parallel loops for assembly of expr.
:arg expr: The expression to assemble.
:arg create_op2arg: callable that creates the Arg corresponding to
the output tensor.
:arg assembly_rank: are we assembling a scalar, vector, or matrix?
:arg diagonal: For matrices are we actually assembling the
diagonal into a vector?
:arg form_compiler_parameters: parameters to pass to the form
compiler.
:returns: a generator of op2.ParLoop objects."""
coefficients = expr.coefficients()
domains = expr.ufl_domains()
if isinstance(expr, slate.TensorBase):
if diagonal:
raise NotImplementedError("Diagonal + slate not supported")
kernels = slac.compile_expression(
expr, tsfc_parameters=form_compiler_parameters)
else:
kernels = tsfc_interface.compile_form(
expr,
"form",
parameters=form_compiler_parameters,
diagonal=diagonal)
# These will be used to correctly interpret the "otherwise"
# subdomain
all_integer_subdomain_ids = defaultdict(list)
for k in kernels:
if k.kinfo.subdomain_id != "otherwise":
all_integer_subdomain_ids[k.kinfo.integral_type].append(
k.kinfo.subdomain_id)
for k, v in all_integer_subdomain_ids.items():
all_integer_subdomain_ids[k] = tuple(sorted(v))
for indices, kinfo in kernels:
kernel = kinfo.kernel
integral_type = kinfo.integral_type
domain_number = kinfo.domain_number
subdomain_id = kinfo.subdomain_id
coeff_map = kinfo.coefficient_map
pass_layer_arg = kinfo.pass_layer_arg
needs_orientations = kinfo.oriented
needs_cell_facets = kinfo.needs_cell_facets
needs_cell_sizes = kinfo.needs_cell_sizes
m = domains[domain_number]
subdomain_data = expr.subdomain_data()[m]
# Find argument space indices
if assembly_rank == AssemblyRank.MATRIX:
i, j = indices
elif assembly_rank == AssemblyRank.VECTOR:
i, = indices
else:
assert len(indices) == 0
sdata = subdomain_data.get(integral_type, None)
if integral_type != 'cell' and sdata is not None:
raise NotImplementedError(
"subdomain_data only supported with cell integrals.")
# Now build arguments for the par_loop
kwargs = {}
# Some integrals require non-coefficient arguments at the
# end (facet number information).
extra_args = []
itspace = m.measure_set(integral_type, subdomain_id,
all_integer_subdomain_ids)
if integral_type == "cell":
itspace = sdata or itspace
if subdomain_id not in ["otherwise", "everywhere"
] and sdata is not None:
raise ValueError("Cannot use subdomain data and subdomain_id")
def get_map(x):
return x.cell_node_map()
elif integral_type in ("exterior_facet", "exterior_facet_vert"):
extra_args.append(m.exterior_facets.local_facet_dat(op2.READ))
def get_map(x):
return x.exterior_facet_node_map()
elif integral_type in ("exterior_facet_top", "exterior_facet_bottom"):
# In the case of extruded meshes with horizontal facet integrals, two
# parallel loops will (potentially) get created and called based on the
# domain id: interior horizontal, bottom or top.
kwargs["iterate"] = {
"exterior_facet_top": op2.ON_TOP,
"exterior_facet_bottom": op2.ON_BOTTOM
}[integral_type]
def get_map(x):
return x.cell_node_map()
elif integral_type in ("interior_facet", "interior_facet_vert"):
extra_args.append(m.interior_facets.local_facet_dat(op2.READ))
def get_map(x):
return x.interior_facet_node_map()
elif integral_type == "interior_facet_horiz":
kwargs["iterate"] = op2.ON_INTERIOR_FACETS
def get_map(x):
return x.cell_node_map()
else:
raise ValueError("Unknown integral type '%s'" % integral_type)
# Output argument
if assembly_rank == AssemblyRank.MATRIX:
tensor_arg = create_op2arg(op2.INC, get_map, i, j)
elif assembly_rank == AssemblyRank.VECTOR:
tensor_arg = create_op2arg(op2.INC, get_map, i)
else:
tensor_arg = create_op2arg(op2.INC)
coords = m.coordinates
args = [
kernel, itspace, tensor_arg,
coords.dat(op2.READ, get_map(coords))
]
if needs_orientations:
o = m.cell_orientations()
args.append(o.dat(op2.READ, get_map(o)))
if needs_cell_sizes:
o = m.cell_sizes
args.append(o.dat(op2.READ, get_map(o)))
for n in coeff_map:
c = coefficients[n]
for c_ in c.split():
m_ = get_map(c_)
args.append(c_.dat(op2.READ, m_))
if needs_cell_facets:
assert integral_type == "cell"
extra_args.append(m.cell_to_facets(op2.READ))
if pass_layer_arg:
c = op2.Global(1,
itspace.layers - 2,
dtype=numpy.dtype(numpy.int32))
o = c(op2.READ)
extra_args.append(o)
args.extend(extra_args)
kwargs["pass_layer_arg"] = pass_layer_arg
try:
yield op2.ParLoop(*args, **kwargs).compute
except MapValueError:
raise RuntimeError(
"Integral measure does not match measure of all coefficients/arguments"
)
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
