def test_decoratedmap_cache_hit(self, base_map):
sm = op2.DecoratedMap(base_map, [op2.ALL])
sm2 = op2.DecoratedMap(base_map, [op2.ALL])
assert sm is sm2
assert not sm != sm2
assert sm == sm2
def test_decoratedmap_change_bcs(self, base_map):
sm = op2.DecoratedMap(base_map, [op2.ALL])
smbc = op2.DecoratedMap(base_map, [op2.ALL], implicit_bcs=["top"])
assert "top" in smbc.implicit_bcs
assert "top" not in sm.implicit_bcs
smbc = op2.DecoratedMap(sm, implicit_bcs=["top"])
assert "top" in smbc.implicit_bcs
assert op2.ALL in smbc.iteration_region
assert len(sm.implicit_bcs) == 0
assert op2.ALL in smbc.iteration_region
def test_decoratedmap_le(self, base_map):
sm = op2.DecoratedMap(base_map, [op2.ALL])
assert base_map <= sm
assert sm <= base_map
smbc = op2.DecoratedMap(base_map, [op2.ALL], implicit_bcs=["top"])
assert base_map <= smbc
assert smbc <= base_map
sm2 = op2.DecoratedMap(base_map, [op2.ON_BOTTOM])
assert not base_map <= sm2
assert not sm2 <= base_map
def test_decoratedmap_cache_miss(self, base_map, base_map2):
sm = op2.DecoratedMap(base_map, [op2.ALL])
sm2 = op2.DecoratedMap(base_map2, [op2.ALL])
assert sm is not sm2
assert sm != sm2
assert not sm == sm2
sm3 = op2.DecoratedMap(base_map, [op2.ON_BOTTOM])
assert sm is not sm3
assert sm != sm3
assert not sm == sm3
assert sm2 is not sm3
assert sm2 != sm3
assert not sm2 == sm3
def test_sparsity_cache_miss(self, base_set, base_set2, base_map,
base_map2):
dsets = (base_set, base_set)
maps = (base_map, base_map)
sp = op2.Sparsity(dsets, maps)
dsets2 = op2.MixedSet([base_set, base_set])
maps2 = op2.MixedMap([base_map, base_map])
maps2 = op2.DecoratedMap(maps2, [op2.ALL])
sp2 = op2.Sparsity(dsets2, maps2)
assert sp is not sp2
assert sp != sp2
assert not sp == sp2
dsets2 = (base_set, base_set2)
maps2 = (base_map, base_map2)
sp2 = op2.Sparsity(dsets2, maps2)
assert sp is not sp2
assert sp != sp2
assert not sp == sp2
def get_map(self,
V,
entity_set,
map_arity,
bcs,
name,
offset,
parent,
kind=None):
"""Return a :class:`pyop2.Map` from some topological entity to
degrees of freedom.
:arg V: The :class:`FunctionSpace` to create the map for.
:arg entity_set: The :class:`pyop2.Set` of entities to map from.
:arg map_arity: The arity of the resulting map.
:arg bcs: An iterable of :class:`~.DirichletBC` objects (may
be ``None``.
:arg name: A name for the resulting map.
:arg offset: Map offset (for extruded).
:arg parent: The parent map (used when bcs are provided)."""
# V is only really used for error checking and "name".
assert len(
V) == 1, "get_map should not be called on MixedFunctionSpace"
entity_node_list = self.entity_node_lists[entity_set]
cache = self.map_caches[entity_set]
if bcs is not None:
# Separate explicit bcs (we just place negative entries in
# the appropriate map values) from implicit ones (extruded
# top and bottom) that require PyOP2 code gen.
explicit_bcs = [
bc for bc in bcs if bc.sub_domain not in ['top', 'bottom']
]
implicit_bcs = [(bc.sub_domain, bc.method) for bc in bcs
if bc.sub_domain in ['top', 'bottom']]
if len(explicit_bcs) == 0:
# Implicit bcs are not part of the cache key for the
# map (they only change the generated PyOP2 code),
# hence rewrite bcs here.
bcs = ()
if len(implicit_bcs) == 0:
implicit_bcs = None
else:
# Empty tuple if no bcs found. This is so that matrix
# assembly, which uses a set to keep track of the bcs
# applied to matrix hits the cache when that set is
# empty. tuple(set([])) == tuple().
bcs = ()
implicit_bcs = None
for bc in bcs:
fs = bc.function_space()
# Unwind proxies for ComponentFunctionSpace, but not
# IndexedFunctionSpace.
while fs.component is not None and fs.parent is not None:
fs = fs.parent
if fs.topological != V:
raise RuntimeError(
"DirichletBC defined on a different FunctionSpace!")
# Ensure bcs is a tuple in a canonical order for the hash key.
lbcs = tuple(sorted(bcs, key=lambda bc: bc.__hash__()))
cache = self.map_caches[entity_set]
try:
# Cache hit
val = cache[lbcs]
# In the implicit bc case, we decorate the cached map with
# the list of implicit boundary conditions so PyOP2 knows
# what to do.
if implicit_bcs:
val = op2.DecoratedMap(val, implicit_bcs=implicit_bcs)
return val
except KeyError:
# Cache miss.
# Any top and bottom bcs (for the extruded case) are handled elsewhere.
nodes = [
bc.nodes for bc in lbcs
if bc.sub_domain not in ['top', 'bottom']
]
decorate = any(bc.function_space().component is not None
for bc in lbcs)
if nodes:
bcids = reduce(numpy.union1d, nodes)
negids = numpy.copy(bcids)
for bc in lbcs:
if bc.sub_domain in ["top", "bottom"]:
continue
nbits = IntType.itemsize * 8 - 2
if decorate and bc.function_space().component is None:
# Some of the other entries will be marked
# with high bits, so we need to set all the
# high bits for these bcs
idx = numpy.searchsorted(bcids, bc.nodes)
if bc.function_space().dim > 3:
raise ValueError(
"Can't have component BCs with more than three components (have %d)",
bc.function_space().dim)
for cmp in range(bc.function_space().dim):
negids[idx] |= (1 << (nbits - cmp))
# FunctionSpace with component is IndexedVFS
if bc.function_space().component is not None:
# For indexed VFS bcs, we encode the component
# in the high bits of the map value.
# That value is then negated to indicate to
# the generated code to discard the values
#
# So here we do:
#
# node = -(node + 2**(nbits-cmpt) + 1)
#
# And in the generated code we can then
# extract the information to discard the
# correct entries.
# bcids is sorted, so use searchsorted to find indices
idx = numpy.searchsorted(bcids, bc.nodes)
# Set appropriate bit
negids[idx] |= (
1 << (nbits - bc.function_space().component))
node_list_bc = numpy.arange(self.node_set.total_size,
dtype=IntType)
# Fix up for extruded, doesn't commute with indexedvfs
# for now
if self.extruded:
node_list_bc[bcids] = -10000000
else:
node_list_bc[bcids] = -(negids + 1)
new_entity_node_list = node_list_bc.take(entity_node_list)
else:
new_entity_node_list = entity_node_list
if kind == "interior_facet" and self.bt_masks is not None:
bt_masks = {}
off = map_arity // 2
for method, (bottom, top) in iteritems(self.bt_masks):
b = []
t = []
for i in bottom:
b.append(i)
b.append(i + off)
for i in top:
t.append(i)
t.append(i + off)
bt_masks[method] = (b, t)
else:
bt_masks = self.bt_masks
val = op2.Map(entity_set,
self.node_set,
map_arity,
new_entity_node_list, ("%s_" + name) % (V.name),
offset=offset,
parent=parent,
bt_masks=bt_masks)
if decorate:
val = op2.DecoratedMap(val, vector_index=True)
cache[lbcs] = val
if implicit_bcs:
return op2.DecoratedMap(val, implicit_bcs=implicit_bcs)
return val
def get_map(x, bcs=None, decoration=None):
map_ = x.cell_node_map(bcs)
if decoration is not None:
return op2.DecoratedMap(map_, decoration)
return map_
def _assemble(f,
tensor=None,
bcs=None,
form_compiler_parameters=None,
inverse=False,
mat_type=None,
sub_mat_type=None,
appctx={},
options_prefix=None,
collect_loops=False,
allocate_only=False):
"""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 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):
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)
integral_types = [
integral.integral_type() for integral in f.integrals()
]
rank = len(f.arguments())
is_mat = rank == 2
is_vec = rank == 1
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 = 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"
if matfree: # intercept matrix-free matrices here
if inverse:
raise NotImplementedError(
"Inverse not implemented with matfree")
if collect_loops:
raise NotImplementedError("Can't collect loops with matfree")
if tensor is None:
return matrix.ImplicitMatrix(
f,
bcs,
fc_params=form_compiler_parameters,
appctx=appctx,
options_prefix=options_prefix)
if not isinstance(tensor, matrix.ImplicitMatrix):
raise ValueError("Expecting implicit matrix with matfree")
tensor.assemble()
return tensor
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)
# To avoid an extra check for extruded domains, the maps that are being passed in
# are DecoratedMaps. For the non-extruded case the DecoratedMaps don't restrict the
# space over which we iterate as the domains are dropped at Sparsity construction
# time. In the extruded case the cell domains are used to identify the regions of the
# mesh which require allocation in the sparsity.
if cell_domains:
map_pairs.append(
(op2.DecoratedMap(test.cell_node_map(), cell_domains),
op2.DecoratedMap(trial.cell_node_map(), cell_domains)))
if exterior_facet_domains:
map_pairs.append(
(op2.DecoratedMap(test.exterior_facet_node_map(),
exterior_facet_domains),
op2.DecoratedMap(trial.exterior_facet_node_map(),
exterior_facet_domains)))
if interior_facet_domains:
map_pairs.append(
(op2.DecoratedMap(test.interior_facet_node_map(),
interior_facet_domains),
op2.DecoratedMap(trial.interior_facet_node_map(),
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,
"%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,
numpy.float64,
"%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
# Replace any bcs on the tensor we passed in
result_matrix.bcs = bcs
tensor = tensor._M
zero_tensor = 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, i, j):
m = testmap(test.function_space()[i])
n = trialmap(trial.function_space()[j])
maps = (m[op2.i[0]] if m else None,
n[op2.i[1 if m else 0]] if n else None)
return tensor[i, j](op2.INC, maps)
result = lambda: result_matrix
if allocate_only:
result_matrix._assembly_callback = None
return result_matrix
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 = tensor.zero
def vec(testmap, i):
_testmap = testmap(test.function_space()[i])
return tensor[i](op2.INC, _testmap[op2.i[0]] 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))
# Since applying boundary conditions to a matrix changes the
# initial assembly, to support:
# A = assemble(a)
# bc.apply(A)
# solve(A, ...)
# we need to defer actually assembling the matrix until just
# before we need it (when we know if there are any bcs to be
# applied). To do so, we build a closure that carries out the
# assembly and stash that on the Matrix object. When we hit a
# solve, we funcall the closure with any bcs the Matrix now has to
# assemble it.
# 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.
loops = []
def thunk(bcs):
if collect_loops:
loops.append(zero_tensor)
else:
zero_tensor()
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 and result_matrix.block_shape > (1, 1):
tsbc = []
trbc = []
# Unwind ComponentFunctionSpace to check for matching BCs
for bc in bcs:
fs = bc.function_space()
if fs.component is not None:
fs = fs.parent
if fs.index == i:
tsbc.append(bc)
if fs.index == j:
trbc.append(bc)
elif is_mat:
tsbc, trbc = bcs, bcs
# 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, bcs=None, decoration=None):
return x.cell_node_map(bcs)
elif integral_type in ("exterior_facet", "exterior_facet_vert"):
extra_args.append(m.exterior_facets.local_facet_dat(op2.READ))
def get_map(x, bcs=None, decoration=None):
return x.exterior_facet_node_map(bcs)
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, bcs=None, decoration=None):
map_ = x.cell_node_map(bcs)
if decoration is not None:
return op2.DecoratedMap(map_, decoration)
return 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, bcs=None, decoration=None):
return x.interior_facet_node_map(bcs)
elif integral_type == "interior_facet_horiz":
decoration = op2.ON_INTERIOR_FACETS
kwargs["iterate"] = decoration
def get_map(x, bcs=None, decoration=None):
map_ = x.cell_node_map(bcs)
if decoration is not None:
return op2.DecoratedMap(map_, decoration)
return map_
else:
raise ValueError("Unknown integral type '%s'" % integral_type)
# Output argument
if is_mat:
tensor_arg = mat(lambda s: get_map(s, tsbc, decoration),
lambda s: get_map(s, trbc, decoration), 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)[op2.i[0]])
]
if needs_orientations:
o = m.cell_orientations()
args.append(o.dat(op2.READ, get_map(o)[op2.i[0]]))
if needs_cell_sizes:
o = m.cell_sizes
args.append(o.dat(op2.READ, get_map(o)[op2.i[0]]))
for n in coeff_map:
c = coefficients[n]
for c_ in c.split():
m_ = get_map(c_)
args.append(c_.dat(op2.READ, m_ and m_[op2.i[0]]))
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:
with collecting_loops(collect_loops):
loops.append(op2.par_loop(*args, **kwargs))
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:
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(
"""Cannot apply boundary conditions to full mixed space. Did you forget to index it?"""
)
shape = result_matrix.block_shape
with collecting_loops(collect_loops):
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:
loops.append(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
loops.append(
tensor[i, j].set_local_diagonal_entries(
nodes, idx=fs.component))
elif fs.index is None:
loops.append(tensor[
i, j].set_local_diagonal_entries(nodes))
else:
raise RuntimeError("Unhandled BC case")
if bcs is not None and is_vec:
if len(bcs) > 0 and collect_loops:
raise NotImplementedError(
"Loop collection not handled in this case")
for bc in bcs:
bc.apply(result_function)
if is_mat:
# Queue up matrix assembly (after we've done all the other operations)
loops.append(tensor.assemble())
return result()
if collect_loops:
thunk(bcs)
return loops
if is_mat:
result_matrix._assembly_callback = thunk
return result()
else:
return thunk(bcs)
def _assemble(f, tensor=None, bcs=None, form_compiler_parameters=None,
inverse=False, nest=None):
"""Assemble the form f and return a Firedrake object representing the
result. This will be a :class:`float` for 0-forms, a
:class:`.Function` for 1-forms and a :class:`.Matrix` for 2-forms.
:arg bcs: A tuple of :class`.DirichletBC`\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 nest: (optional) flag indicating if matrices on mixed spaces
should be built in blocks or monolithically.
"""
if form_compiler_parameters:
form_compiler_parameters = form_compiler_parameters.copy()
else:
form_compiler_parameters = {}
form_compiler_parameters["assemble_inverse"] = inverse
kernels = tsfc_interface.compile_form(f, "form", parameters=form_compiler_parameters,
inverse=inverse)
rank = len(f.arguments())
is_mat = rank == 2
is_vec = rank == 1
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")
integrals = f.integrals()
if nest is None:
nest = parameters.parameters["matnest"]
# Pass this through for assembly caching purposes
form_compiler_parameters["matnest"] = nest
zero_tensor = lambda: None
if is_mat:
test, trial = f.arguments()
map_pairs = []
cell_domains = []
exterior_facet_domains = []
interior_facet_domains = []
if tensor is None:
# For horizontal facets of extrded 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 in integrals:
integral_type = integral.integral_type()
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)
# To avoid an extra check for extruded domains, the maps that are being passed in
# are DecoratedMaps. For the non-extruded case the DecoratedMaps don't restrict the
# space over which we iterate as the domains are dropped at Sparsity construction
# time. In the extruded case the cell domains are used to identify the regions of the
# mesh which require allocation in the sparsity.
if cell_domains:
map_pairs.append((op2.DecoratedMap(test.cell_node_map(), cell_domains),
op2.DecoratedMap(trial.cell_node_map(), cell_domains)))
if exterior_facet_domains:
map_pairs.append((op2.DecoratedMap(test.exterior_facet_node_map(), exterior_facet_domains),
op2.DecoratedMap(trial.exterior_facet_node_map(), exterior_facet_domains)))
if interior_facet_domains:
map_pairs.append((op2.DecoratedMap(test.interior_facet_node_map(), interior_facet_domains),
op2.DecoratedMap(trial.interior_facet_node_map(), interior_facet_domains)))
map_pairs = tuple(map_pairs)
# Construct OP2 Mat to assemble into
fs_names = (test.function_space().name, trial.function_space().name)
sparsity = op2.Sparsity((test.function_space().dof_dset,
trial.function_space().dof_dset),
map_pairs,
"%s_%s_sparsity" % fs_names,
nest=nest)
result_matrix = matrix.Matrix(f, bcs, sparsity, numpy.float64,
"%s_%s_matrix" % fs_names)
tensor = result_matrix._M
else:
result_matrix = tensor
# Replace any bcs on the tensor we passed in
result_matrix.bcs = bcs
tensor = tensor._M
zero_tensor = lambda: tensor.zero()
def mat(testmap, trialmap, i, j):
return tensor[i, j](op2.INC,
(testmap(test.function_space()[i])[op2.i[0]],
trialmap(trial.function_space()[j])[op2.i[1]]),
flatten=True)
result = lambda: result_matrix
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 = lambda: tensor.zero()
def vec(testmap, i):
return tensor[i](op2.INC,
testmap(test.function_space()[i])[op2.i[0]],
flatten=True)
result = lambda: result_function
else:
# 0-forms are always scalar
if tensor is None:
tensor = op2.Global(1, [0.0])
result = lambda: tensor.data[0]
subdomain_data = f.subdomain_data()
coefficients = f.coefficients()
domains = f.ufl_domains()
if len(domains) != 1:
raise NotImplementedError("Assembly of forms with more than one domain not supported")
m = domains[0]
# Ensure mesh is "initialised" (could have got here without
# building a functionspace (e.g. if integrating a constant)).
m.init()
subdomain_data = subdomain_data[m]
# Since applying boundary conditions to a matrix changes the
# initial assembly, to support:
# A = assemble(a)
# bc.apply(A)
# solve(A, ...)
# we need to defer actually assembling the matrix until just
# before we need it (when we know if there are any bcs to be
# applied). To do so, we build a closure that carries out the
# assembly and stash that on the Matrix object. When we hit a
# solve, we funcall the closure with any bcs the Matrix now has to
# assemble it.
def thunk(bcs):
zero_tensor()
for indices, (kernel, integral_type, needs_orientations, subdomain_id, coeff_map) in kernels:
# 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 and tensor.sparsity.shape > (1, 1):
tsbc = []
trbc = []
# Unwind ComponentFunctionSpace to check for matching BCs
for bc in bcs:
fs = bc.function_space()
if fs.component is not None:
fs = fs.parent
if fs.index == i:
tsbc.append(bc)
if fs.index == j:
trbc.append(bc)
elif is_mat:
tsbc, trbc = bcs, bcs
# 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
if integral_type == "cell":
itspace = sdata or m.cell_set
def get_map(x, bcs=None, decoration=None):
return x.cell_node_map(bcs)
elif integral_type in ("exterior_facet", "exterior_facet_vert"):
itspace = m.exterior_facets.measure_set(integral_type, subdomain_id)
extra_args.append(m.exterior_facets.local_facet_dat(op2.READ))
def get_map(x, bcs=None, decoration=None):
return x.exterior_facet_node_map(bcs)
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.
index, itspace = m.exterior_facets.measure_set(integral_type, subdomain_id)
decoration = index
kwargs["iterate"] = index
def get_map(x, bcs=None, decoration=None):
map_ = x.cell_node_map(bcs)
if decoration is not None:
return op2.DecoratedMap(map_, decoration)
return map_
elif integral_type in ("interior_facet", "interior_facet_vert"):
itspace = m.interior_facets.set
extra_args.append(m.interior_facets.local_facet_dat(op2.READ))
def get_map(x, bcs=None, decoration=None):
return x.interior_facet_node_map(bcs)
elif integral_type == "interior_facet_horiz":
itspace = m.interior_facets.measure_set(integral_type, subdomain_id)
decoration = op2.ON_INTERIOR_FACETS
kwargs["iterate"] = op2.ON_INTERIOR_FACETS
def get_map(x, bcs=None, decoration=None):
map_ = x.cell_node_map(bcs)
if decoration is not None:
return op2.DecoratedMap(map_, decoration)
return map_
else:
raise ValueError("Unknown integral type '%s'" % integral_type)
# Output argument
if is_mat:
tensor_arg = mat(lambda s: get_map(s, tsbc, decoration),
lambda s: get_map(s, trbc, decoration),
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), flatten=True)]
if needs_orientations:
o = m.cell_orientations()
args.append(o.dat(op2.READ, get_map(o), flatten=True))
for n in coeff_map:
c = coefficients[n]
args.append(c.dat(op2.READ, get_map(c), flatten=True))
args.extend(extra_args)
try:
op2.par_loop(*args, **kwargs)
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:
fs = bc.function_space()
if len(fs) > 1:
raise RuntimeError("""Cannot apply boundary conditions to full mixed space. Did you forget to index it?""")
shape = tensor.sparsity.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:
tensor[i, j].set_local_diagonal_entries(bc.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
tensor[i, j].set_local_diagonal_entries(bc.nodes, idx=fs.component)
elif fs.index is None:
tensor[i, j].set_local_diagonal_entries(bc.nodes)
else:
raise RuntimeError("Unhandled BC case")
if bcs is not None and is_vec:
for bc in bcs:
bc.apply(result_function)
if is_mat:
# Queue up matrix assembly (after we've done all the other operations)
tensor.assemble()
return result()
thunk = assembly_cache._cache_thunk(thunk, f, result(), form_compiler_parameters)
if is_mat:
result_matrix._assembly_callback = thunk
return result()
else:
return thunk(bcs)
def _map_cache(self,
cache,
entity_set,
entity_node_list,
map_arity,
bcs,
name,
offset=None,
parent=None):
if bcs is not None:
# Separate explicit bcs (we just place negative entries in
# the appropriate map values) from implicit ones (extruded
# top and bottom) that require PyOP2 code gen.
explicit_bcs = [
bc for bc in bcs if bc.sub_domain not in ['top', 'bottom']
]
implicit_bcs = [
bc.sub_domain for bc in bcs
if bc.sub_domain in ['top', 'bottom']
]
if len(explicit_bcs) == 0:
# Implicit bcs are not part of the cache key for the
# map (they only change the generated PyOP2 code),
# hence rewrite bcs here.
bcs = None
if len(implicit_bcs) == 0:
implicit_bcs = None
else:
implicit_bcs = None
if bcs is None:
# Empty tuple if no bcs found. This is so that matrix
# assembly, which uses a set to keep track of the bcs
# applied to matrix hits the cache when that set is
# empty. tuple(set([])) == tuple().
lbcs = tuple()
else:
for bc in bcs:
fs = bc.function_space()
if isinstance(fs, IndexedVFS):
fs = fs._parent
if fs != self:
raise RuntimeError(
"DirichletBC defined on a different FunctionSpace!")
# Ensure bcs is a tuple in a canonical order for the hash key.
lbcs = tuple(sorted(bcs, key=lambda bc: bc.__hash__()))
try:
# Cache hit
val = cache[lbcs]
# In the implicit bc case, we decorate the cached map with
# the list of implicit boundary conditions so PyOP2 knows
# what to do.
if implicit_bcs:
val = op2.DecoratedMap(val, implicit_bcs=implicit_bcs)
return val
except KeyError:
# Cache miss.
# Any top and bottom bcs (for the extruded case) are handled elsewhere.
nodes = [
bc.nodes for bc in lbcs
if bc.sub_domain not in ['top', 'bottom']
]
decorate = False
if nodes:
bcids = reduce(np.union1d, nodes)
negids = np.copy(bcids)
for bc in lbcs:
if bc.sub_domain in ["top", "bottom"]:
continue
if isinstance(bc.function_space(), IndexedVFS):
# For indexed VFS bcs, we encode the component
# in the high bits of the map value.
# That value is then negated to indicate to
# the generated code to discard the values
#
# So here we do:
#
# node = -(node + 2**(30-cmpt) + 1)
#
# And in the generated code we can then
# extract the information to discard the
# correct entries.
val = 2**(30 - bc.function_space().index)
# bcids is sorted, so use searchsorted to find indices
idx = np.searchsorted(bcids, bc.nodes)
negids[idx] += val
decorate = True
node_list_bc = np.arange(self.node_count, dtype=np.int32)
# Fix up for extruded, doesn't commute with indexedvfs for now
if isinstance(self.mesh(), mesh_t.ExtrudedMesh):
node_list_bc[bcids] = -10000000
else:
node_list_bc[bcids] = -(negids + 1)
new_entity_node_list = node_list_bc.take(entity_node_list)
else:
new_entity_node_list = entity_node_list
val = op2.Map(entity_set, self.node_set, map_arity,
new_entity_node_list, ("%s_" + name) % (self.name),
offset, parent, self.bt_masks)
if decorate:
val = op2.DecoratedMap(val, vector_index=True)
cache[lbcs] = val
if implicit_bcs:
return op2.DecoratedMap(val, implicit_bcs=implicit_bcs)
return val
def thunk(bcs):
zero_tensor()
for (
i, j
), integral_type, subdomain_id, coords, coefficients, needs_orientations, kernel in kernels:
m = coords.function_space().mesh()
if needs_orientations:
cell_orientations = m.cell_orientations()
# 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 and tensor.sparsity.shape > (1, 1):
tsbc = []
trbc = []
# Unwind IndexedVFS to check for matching BCs
for bc in bcs:
fs = bc.function_space()
if isinstance(fs, functionspace.IndexedVFS):
fs = fs._parent
if fs.index == i:
tsbc.append(bc)
if fs.index == j:
trbc.append(bc)
elif is_mat:
tsbc, trbc = bcs, bcs
if integral_type == 'cell':
with timed_region("Assemble cells"):
if is_mat:
tensor_arg = mat(lambda s: s.cell_node_map(tsbc),
lambda s: s.cell_node_map(trbc), i, j)
elif is_vec:
tensor_arg = vec(lambda s: s.cell_node_map(), i)
else:
tensor_arg = tensor(op2.INC)
itspace = m.cell_set
args = [
kernel, itspace, tensor_arg,
coords.dat(op2.READ,
coords.cell_node_map(),
flatten=True)
]
if needs_orientations:
args.append(
cell_orientations.dat(
op2.READ,
cell_orientations.cell_node_map(),
flatten=True))
for c in coefficients:
args.append(
c.dat(op2.READ, c.cell_node_map(), flatten=True))
try:
op2.par_loop(*args)
except MapValueError:
raise RuntimeError(
"Integral measure does not match measure of all coefficients/arguments"
)
elif integral_type in ['exterior_facet', 'exterior_facet_vert']:
with timed_region("Assemble exterior facets"):
if is_mat:
tensor_arg = mat(
lambda s: s.exterior_facet_node_map(tsbc),
lambda s: s.exterior_facet_node_map(trbc), i, j)
elif is_vec:
tensor_arg = vec(lambda s: s.exterior_facet_node_map(),
i)
else:
tensor_arg = tensor(op2.INC)
args = [
kernel,
m.exterior_facets.measure_set(integral_type,
subdomain_id),
tensor_arg,
coords.dat(op2.READ,
coords.exterior_facet_node_map(),
flatten=True)
]
if needs_orientations:
args.append(
cell_orientations.dat(
op2.READ,
cell_orientations.exterior_facet_node_map(),
flatten=True))
for c in coefficients:
args.append(
c.dat(op2.READ,
c.exterior_facet_node_map(),
flatten=True))
args.append(m.exterior_facets.local_facet_dat(op2.READ))
try:
op2.par_loop(*args)
except MapValueError:
raise RuntimeError(
"Integral measure does not match measure of all coefficients/arguments"
)
elif integral_type in [
'exterior_facet_top', 'exterior_facet_bottom'
]:
with timed_region("Assemble exterior facets"):
# 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.
# Get the list of sets and globals required for parallel loop construction.
set_global_list = m.exterior_facets.measure_set(
integral_type, subdomain_id)
# Iterate over the list and assemble all the args of the parallel loop
for (index, set) in set_global_list:
if is_mat:
tensor_arg = mat(
lambda s: op2.DecoratedMap(
s.cell_node_map(tsbc), index),
lambda s: op2.DecoratedMap(
s.cell_node_map(trbc), index), i, j)
elif is_vec:
tensor_arg = vec(lambda s: s.cell_node_map(), i)
else:
tensor_arg = tensor(op2.INC)
# Add the kernel, iteration set and coordinate fields to the loop args
args = [
kernel, set, tensor_arg,
coords.dat(op2.READ,
coords.cell_node_map(),
flatten=True)
]
if needs_orientations:
args.append(
cell_orientations.dat(
op2.READ,
cell_orientations.cell_node_map(),
flatten=True))
for c in coefficients:
args.append(
c.dat(op2.READ,
c.cell_node_map(),
flatten=True))
try:
op2.par_loop(*args, iterate=index)
except MapValueError:
raise RuntimeError(
"Integral measure does not match measure of all coefficients/arguments"
)
elif integral_type in ['interior_facet', 'interior_facet_vert']:
with timed_region("Assemble interior facets"):
if is_mat:
tensor_arg = mat(
lambda s: s.interior_facet_node_map(tsbc),
lambda s: s.interior_facet_node_map(trbc), i, j)
elif is_vec:
tensor_arg = vec(lambda s: s.interior_facet_node_map(),
i)
else:
tensor_arg = tensor(op2.INC)
args = [
kernel, m.interior_facets.set, tensor_arg,
coords.dat(op2.READ,
coords.interior_facet_node_map(),
flatten=True)
]
if needs_orientations:
args.append(
cell_orientations.dat(
op2.READ,
cell_orientations.interior_facet_node_map(),
flatten=True))
for c in coefficients:
args.append(
c.dat(op2.READ,
c.interior_facet_node_map(),
flatten=True))
args.append(m.interior_facets.local_facet_dat(op2.READ))
try:
op2.par_loop(*args)
except MapValueError:
raise RuntimeError(
"Integral measure does not match measure of all coefficients/arguments"
)
elif integral_type == 'interior_facet_horiz':
with timed_region("Assemble interior facets"):
if is_mat:
tensor_arg = mat(
lambda s: op2.DecoratedMap(s.cell_node_map(tsbc),
op2.ON_INTERIOR_FACETS),
lambda s: op2.DecoratedMap(s.cell_node_map(
trbc), op2.ON_INTERIOR_FACETS), i, j)
elif is_vec:
tensor_arg = vec(lambda s: s.cell_node_map(), i)
else:
tensor_arg = tensor(op2.INC)
args = [
kernel,
m.interior_facets.measure_set(integral_type,
subdomain_id),
tensor_arg,
coords.dat(op2.READ,
coords.cell_node_map(),
flatten=True)
]
if needs_orientations:
args.append(
cell_orientations.dat(
op2.READ,
cell_orientations.cell_node_map(),
flatten=True))
for c in coefficients:
args.append(
c.dat(op2.READ, c.cell_node_map(), flatten=True))
try:
op2.par_loop(*args, iterate=op2.ON_INTERIOR_FACETS)
except MapValueError:
raise RuntimeError(
"Integral measure does not match measure of all coefficients/arguments"
)
else:
raise RuntimeError('Unknown integral type "%s"' %
integral_type)
# 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:
with timed_region('DirichletBC apply'):
for bc in bcs:
fs = bc.function_space()
if isinstance(fs, functionspace.MixedFunctionSpace):
raise RuntimeError(
"""Cannot apply boundary conditions to full mixed space. Did you forget to index it?"""
)
shape = tensor.sparsity.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 isinstance(fs,
functionspace.IndexedFunctionSpace):
# Mixed, index
if fs.index == i:
tensor[i, j].set_local_diagonal_entries(
bc.nodes)
elif isinstance(fs, functionspace.IndexedVFS):
if isinstance(
fs._parent,
functionspace.IndexedFunctionSpace):
if fs._parent.index != i:
continue
tensor[i, j].set_local_diagonal_entries(
bc.nodes, idx=fs.index)
elif fs.index is None:
tensor[i,
j].set_local_diagonal_entries(bc.nodes)
else:
raise RuntimeError("Unhandled BC case")
if bcs is not None and is_vec:
for bc in bcs:
bc.apply(result_function)
if is_mat:
# Queue up matrix assembly (after we've done all the other operations)
tensor.assemble()
return result()