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 = ffc_interface.compile_form(f, "form", parameters=form_compiler_parameters,
inverse=inverse)
rank = len(f.arguments())
is_mat = rank == 2
is_vec = rank == 1
if inverse and rank != 2:
raise ValueError("Can only assemble the inverse of a 2-form")
integrals = f.integrals()
def get_rank(arg):
return arg.function_space().rank
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 RuntimeError('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]
# 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 (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 = [bc for bc in bcs if bc.function_space().index == i]
trbc = [bc for bc in bcs if bc.function_space().index == j]
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 and (fs.index is None or fs.index == i):
tensor[i, j].inc_local_diagonal_entries(bc.nodes)
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)
