def test_sparsity_has_diagonal_space(self):
# A sparsity should have space for diagonal entries if rmap==cmap
s = op2.Set(1)
d = op2.Set(4)
m = op2.Map(s, d, 2, [1, 3])
d2 = op2.Set(4)
m2 = op2.Map(s, d2, 3, [1, 2, 3])
sparsity = op2.Sparsity((d, d), (m, m))
sparsity2 = op2.Sparsity((d, d2), (m, m2))
assert all(sparsity.nnz == [1, 2, 1, 2])
assert all(sparsity2.nnz == [0, 3, 0, 3])
def test_mat_set_diagonal(self, nodes, elem_node, n):
"Set the diagonal of the entire matrix to 1.0"
mat = op2.Mat(op2.Sparsity(nodes**n, elem_node), valuetype)
nrows = mat.sparsity.nrows
mat.set_local_diagonal_entries(list(range(nrows)))
mat.assemble()
assert (mat.values == np.identity(nrows * n)).all()
def make_interpolator(expr, V, subset, access):
assert isinstance(expr, ufl.classes.Expr)
if isinstance(expr, firedrake.Expression):
arguments = ()
else:
arguments = extract_arguments(expr)
if len(arguments) == 0:
if isinstance(V, firedrake.Function):
f = V
V = f.function_space()
else:
f = firedrake.Function(V)
tensor = f.dat
elif len(arguments) == 1:
if isinstance(V, firedrake.Function):
raise ValueError(
"Cannot interpolate an expression with an argument into a Function"
)
argfs = arguments[0].function_space()
sparsity = op2.Sparsity((V.dof_dset, argfs.dof_dset),
((V.cell_node_map(), argfs.cell_node_map()), ),
name="%s_%s_sparsity" % (V.name, argfs.name),
nest=False,
block_sparse=True)
tensor = op2.Mat(sparsity)
f = tensor
else:
raise ValueError("Cannot interpolate an expression with %d arguments" %
len(arguments))
# Make sure we have an expression of the right length i.e. a value for
# each component in the value shape of each function space
dims = [numpy.prod(fs.ufl_element().value_shape(), dtype=int) for fs in V]
loops = []
if numpy.prod(expr.ufl_shape, dtype=int) != sum(dims):
raise RuntimeError('Expression of length %d required, got length %d' %
(sum(dims), numpy.prod(expr.ufl_shape, dtype=int)))
if not isinstance(expr, firedrake.Expression):
if len(V) > 1:
raise NotImplementedError(
"UFL expressions for mixed functions are not yet supported.")
loops.extend(_interpolator(V, tensor, expr, subset, arguments, access))
elif hasattr(expr, 'eval'):
if len(V) > 1:
raise NotImplementedError(
"Python expressions for mixed functions are not yet supported."
)
loops.extend(_interpolator(V, tensor, expr, subset, arguments, access))
else:
raise ValueError("Don't know how to interpolate a %r" % expr)
def callable(loops, f):
for l in loops:
l()
return f
return partial(callable, loops, f), arguments
def time_sparsity(n):
nodes = op2.Set(n)
cells = op2.Set(n - 1)
m = op2.Map(cells, nodes, 2,
np.concatenate((np.arange(n - 1), np.arange(1, n))))
t = clock()
s = op2.Sparsity((m, m), 1)
return clock() - t
def setupHessian(poses, constraints_to_poses):
""" define the sparsity pattern of the hessian and return the Mat """
hess_sparsity = op2.Sparsity((poses ** POSES_DIM, poses ** POSES_DIM),
(constraints_to_poses, constraints_to_poses),
'hess_sparsity')
if _VERBOSE:
print 'hessian has dimension %d by %d' % (NUM_POSES * POSES_DIM, NUM_POSES * POSES_DIM)
return petsc_base.Mat(hess_sparsity, np.float64, 'hess_matrix')
def test_two_mats_on_same_sparsity_share_data(self, backend, skip_opencl,
m1, skip_sequential,
skip_openmp, ds2):
"""Sparsity data should be shared between Mat objects.
Even on the device."""
sp = op2.Sparsity((ds2, ds2), (m1, m1))
mat1 = op2.Mat(sp, 'float64')
mat2 = op2.Mat(sp, 'float64')
assert mat1._colidx is mat2._colidx
assert mat1._rowptr is mat2._rowptr
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, iteration_regions=[(op2.ALL, )])
dsets2 = op2.MixedSet([base_set, base_set])
maps2 = op2.MixedMap([base_map, base_map])
sp2 = op2.Sparsity(dsets2, maps2, iteration_regions=[(op2.ALL, )])
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, iteration_regions=[(op2.ALL, )])
assert sp is not sp2
assert sp != sp2
assert not sp == sp2
def test_build_sparsity(self):
"""Building a sparsity from a pair of maps should give the expected
rowptr and colidx."""
elements = op2.Set(4)
nodes = op2.Set(5)
elem_node = op2.Map(elements, nodes, 3, [0, 4, 3, 0, 1, 4,
1, 2, 4, 2, 3, 4])
sparsity = op2.Sparsity((nodes, nodes), (elem_node, elem_node))
assert all(sparsity._rowptr == [0, 4, 8, 12, 16, 21])
assert all(sparsity._colidx == [0, 1, 3, 4, 0, 1, 2, 4, 1, 2,
3, 4, 0, 2, 3, 4, 0, 1, 2, 3, 4])
def test_sparsity_cache_hit(self, base_set, base_map):
dsets = (base_set, base_set)
maps = (base_map, base_map)
sp = op2.Sparsity(dsets, maps)
sp2 = op2.Sparsity(dsets, maps)
assert sp is sp2
assert not sp != sp2
assert sp == sp2
dsets = op2.MixedSet([base_set, base_set])
maps = op2.MixedMap([base_map, base_map])
sp = op2.Sparsity(dsets, maps)
dsets2 = op2.MixedSet([base_set, base_set])
maps2 = op2.MixedMap([base_map, base_map])
sp2 = op2.Sparsity(dsets2, maps2)
assert sp is sp2
assert not sp != sp2
assert sp == sp2
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 restriction_matrix(Pk, P1, Pk_bcs, P1_bcs):
sp = op2.Sparsity((P1.dof_dset, Pk.dof_dset),
(P1.cell_node_map(), Pk.cell_node_map()))
mat = op2.Mat(sp, PETSc.ScalarType)
rlgmap, clgmap = mat.local_to_global_maps
rlgmap = P1.local_to_global_map(P1_bcs, lgmap=rlgmap)
clgmap = Pk.local_to_global_map(Pk_bcs, lgmap=clgmap)
unroll = any(bc.function_space().component is not None
for bc in chain(P1_bcs, Pk_bcs) if bc is not None)
matarg = mat(op2.WRITE, (P1.cell_node_map(), Pk.cell_node_map()),
lgmaps=(rlgmap, clgmap),
unroll_map=unroll)
mesh = Pk.ufl_domain()
op2.par_loop(transfer_kernel(Pk, P1), mesh.cell_set, matarg)
mat.assemble()
return mat.handle
def restriction_matrix(Pk, P1, Pk_bcs, P1_bcs):
sp = op2.Sparsity((P1.dof_dset, Pk.dof_dset),
(P1.cell_node_map(), Pk.cell_node_map()))
mat = op2.Mat(sp, PETSc.ScalarType)
matarg = mat(op2.WRITE, (P1.cell_node_map(P1_bcs)[op2.i[0]],
Pk.cell_node_map(Pk_bcs)[op2.i[1]]))
# # HACK HACK HACK, this seems like it might be a pyop2 bug
# sh = matarg._block_shape
# assert len(sh) == 1 and len(sh[0]) == 1 and len(sh[0][0]) == 2
# a, b = sh[0][0]
# nsh = (((a*self.P1.dof_dset.cdim, b*self.V.dof_dset.cdim), ), )
# matarg._block_shape = nsh
mesh = Pk.ufl_domain()
op2.par_loop(transfer_kernel(Pk, P1), mesh.cell_set, matarg)
mat.assemble()
mat._force_evaluation()
return mat.handle
def test_mat_always_has_diagonal_space(self):
# A sparsity should always have space for diagonal entries
s = op2.Set(1)
d = op2.Set(4)
m = op2.Map(s, d, 1, [2])
d2 = op2.Set(3)
m2 = op2.Map(s, d2, 1, [1])
sparsity = op2.Sparsity((d, d2), (m, m2))
from petsc4py import PETSc
# petsc4py default error handler swallows SETERRQ, so just
# install the abort handler to notice an error.
PETSc.Sys.pushErrorHandler("abort")
mat = op2.Mat(sparsity)
PETSc.Sys.popErrorHandler()
assert np.allclose(mat.handle.getDiagonal().array, 0.0)
def test_matrix(self, backend):
"""Test a indirect par_loop with a matrix argument"""
iterset = op2.Set(2)
idset = op2.Set(2)
ss01 = op2.Subset(iterset, [0, 1])
ss10 = op2.Subset(iterset, [1, 0])
indset = op2.Set(4)
dat = op2.Dat(idset ** 1, data=[0, 1], dtype=np.float)
map = op2.Map(iterset, indset, 4, [0, 1, 2, 3, 0, 1, 2, 3])
idmap = op2.Map(iterset, idset, 1, [0, 1])
sparsity = op2.Sparsity((indset, indset), (map, map))
mat = op2.Mat(sparsity, np.float64)
mat01 = op2.Mat(sparsity, np.float64)
mat10 = op2.Mat(sparsity, np.float64)
assembly = c_for("i", 4,
c_for("j", 4,
Incr(Symbol("mat", ("i", "j")), FlatBlock("(*dat)*16+i*4+j"))))
kernel_code = FunDecl("void", "unique_id",
[Decl("double*", c_sym("dat")),
Decl("double", Symbol("mat", (4, 4)))],
Block([assembly], open_scope=False))
k = op2.Kernel(kernel_code, "unique_id")
mat.zero()
mat01.zero()
mat10.zero()
op2.par_loop(k, iterset,
dat(op2.READ, idmap[0]),
mat(op2.INC, (map[op2.i[0]], map[op2.i[1]])))
mat.assemble()
op2.par_loop(k, ss01,
dat(op2.READ, idmap[0]),
mat01(op2.INC, (map[op2.i[0]], map[op2.i[1]])))
mat01.assemble()
op2.par_loop(k, ss10,
dat(op2.READ, idmap[0]),
mat10(op2.INC, (map[op2.i[0]], map[op2.i[1]])))
mat10.assemble()
assert (mat01.values == mat.values).all()
assert (mat10.values == mat.values).all()
def prolongation_matrix_aij(Pk, P1, Pk_bcs, P1_bcs):
sp = op2.Sparsity((Pk.dof_dset, P1.dof_dset),
(Pk.cell_node_map(), P1.cell_node_map()))
mat = op2.Mat(sp, PETSc.ScalarType)
mesh = Pk.ufl_domain()
fele = Pk.ufl_element()
if isinstance(fele, MixedElement) and not isinstance(
fele, (VectorElement, TensorElement)):
for i in range(fele.num_sub_elements()):
Pk_bcs_i = [bc for bc in Pk_bcs if bc.function_space().index == i]
P1_bcs_i = [bc for bc in P1_bcs if bc.function_space().index == i]
rlgmap, clgmap = mat[i, i].local_to_global_maps
rlgmap = Pk.sub(i).local_to_global_map(Pk_bcs_i, lgmap=rlgmap)
clgmap = P1.sub(i).local_to_global_map(P1_bcs_i, lgmap=clgmap)
unroll = any(bc.function_space().component is not None
for bc in chain(Pk_bcs_i, P1_bcs_i) if bc is not None)
matarg = mat[i, i](
op2.WRITE,
(Pk.sub(i).cell_node_map(), P1.sub(i).cell_node_map()),
lgmaps=((rlgmap, clgmap), ),
unroll_map=unroll)
op2.par_loop(
prolongation_transfer_kernel_aij(Pk.sub(i), P1.sub(i)),
mesh.cell_set, matarg)
else:
rlgmap, clgmap = mat.local_to_global_maps
rlgmap = Pk.local_to_global_map(Pk_bcs, lgmap=rlgmap)
clgmap = P1.local_to_global_map(P1_bcs, lgmap=clgmap)
unroll = any(bc.function_space().component is not None
for bc in chain(Pk_bcs, P1_bcs) if bc is not None)
matarg = mat(op2.WRITE, (Pk.cell_node_map(), P1.cell_node_map()),
lgmaps=((rlgmap, clgmap), ),
unroll_map=unroll)
op2.par_loop(prolongation_transfer_kernel_aij(Pk, P1), mesh.cell_set,
matarg)
mat.assemble()
return mat.handle
def test_minimal_zero_mat(self):
"""Assemble a matrix that is all zeros."""
code = c_for("i", 1,
c_for("j", 1,
Assign(Symbol("local_mat", ("i", "j")), c_sym("0.0"))))
zero_mat_code = FunDecl("void", "zero_mat",
[Decl("double", Symbol("local_mat", (1, 1)))],
Block([code], open_scope=False))
nelems = 128
set = op2.Set(nelems)
map = op2.Map(set, set, 1, np.array(list(range(nelems)), np.uint32))
sparsity = op2.Sparsity((set, set), (map, map))
mat = op2.Mat(sparsity, np.float64)
kernel = op2.Kernel(zero_mat_code, "zero_mat")
op2.par_loop(kernel, set,
mat(op2.WRITE, (map[op2.i[0]], map[op2.i[1]])))
mat.assemble()
expected_matrix = np.zeros((nelems, nelems), dtype=np.float64)
eps = 1.e-12
assert_allclose(mat.values, expected_matrix, eps)
def test_sparsities_differing_map_tuples_not_cached(self, m1, m2, ds2):
"""Sparsities with different maps should not share a C handle."""
sp1 = op2.Sparsity((ds2, ds2), ((m1, m1), (m2, m2)))
sp2 = op2.Sparsity((ds2, ds2), ((m2, m2), (m2, m2)))
assert sp1 is not sp2
def test_sparsity_null_maps(self):
"""Building sparsity from a pair of non-initialized maps should fail."""
s = op2.Set(5)
with pytest.raises(MapValueError):
m = op2.Map(s, s, 1)
op2.Sparsity((s, s), (m, m))
def mat(s2, m2):
return op2.Mat(op2.Sparsity((s2, s2), (m2, m2)))
def mvsparsity(mset, mmap):
return op2.Sparsity(mset ** 2, mmap)
def non_nest_mixed_sparsity(mset, mmap):
return op2.Sparsity(mset, mmap, nest=False)
def msparsity(mset, mmap):
return op2.Sparsity(mset, mmap)
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 _make_matrix(expr, bcs, opts):
"""Make an empty matrix.
:arg expr: The expression being assembled.
:arg bcs: Iterable of boundary conditions.
:arg opts: :class:`_AssemblyOpts` containing the assembly options.
:returns: An empty :class:`.Matrix` or :class:`.ImplicitMatrix`.
"""
matfree = opts.mat_type == "matfree"
arguments = expr.arguments()
if bcs is None:
bcs = ()
else:
if any(isinstance(bc, EquationBC) for bc in bcs):
raise TypeError(
"EquationBC objects not expected here. "
"Preprocess by extracting the appropriate form with bc.extract_form('Jp') or bc.extract_form('J')"
)
if matfree:
return matrix.ImplicitMatrix(expr,
bcs,
fc_params=opts.fc_params,
appctx=opts.appctx,
options_prefix=opts.options_prefix)
integral_types = set(i.integral_type() for i in expr.integrals())
for bc in bcs:
integral_types.update(integral.integral_type()
for integral in bc.integrals())
nest = opts.mat_type == "nest"
if nest:
baij = opts.sub_mat_type == "baij"
else:
baij = opts.mat_type == "baij"
if any(len(a.function_space()) > 1
for a in arguments) and opts.mat_type == "baij":
raise ValueError(
"BAIJ matrix type makes no sense for mixed spaces, use 'aij'")
get_cell_map = operator.methodcaller("cell_node_map")
get_extf_map = operator.methodcaller("exterior_facet_node_map")
get_intf_map = operator.methodcaller("interior_facet_node_map")
domains = OrderedDict(
(k, set()) for k in (get_cell_map, get_extf_map, get_intf_map))
mapping = {
"cell": (get_cell_map, op2.ALL),
"exterior_facet_bottom": (get_cell_map, op2.ON_BOTTOM),
"exterior_facet_top": (get_cell_map, op2.ON_TOP),
"interior_facet_horiz": (get_cell_map, op2.ON_INTERIOR_FACETS),
"exterior_facet": (get_extf_map, op2.ALL),
"exterior_facet_vert": (get_extf_map, op2.ALL),
"interior_facet": (get_intf_map, op2.ALL),
"interior_facet_vert": (get_intf_map, op2.ALL)
}
for integral_type in integral_types:
try:
get_map, region = mapping[integral_type]
except KeyError:
raise ValueError(f"Unknown integral type '{integral_type}'")
domains[get_map].add(region)
test, trial = arguments
map_pairs, iteration_regions = zip(
*(((get_map(test), get_map(trial)), tuple(sorted(regions)))
for get_map, regions in domains.items() if regions))
try:
sparsity = op2.Sparsity(
(test.function_space().dof_dset, trial.function_space().dof_dset),
tuple(map_pairs),
iteration_regions=tuple(iteration_regions),
nest=nest,
block_sparse=baij)
except SparsityFormatError:
raise ValueError(
"Monolithic matrix assembly not supported for systems "
"with R-space blocks")
return matrix.Matrix(expr,
bcs,
opts.mat_type,
sparsity,
ScalarType,
options_prefix=opts.options_prefix)
def run(diffusivity, current_time, dt, endtime, **kwargs):
op2.init(**kwargs)
# Set up finite element problem
T = FiniteElement("Lagrange", "triangle", 1)
V = VectorElement("Lagrange", "triangle", 1)
p = TrialFunction(T)
q = TestFunction(T)
t = Coefficient(T)
u = Coefficient(V)
diffusivity = 0.1
M = p * q * dx
adv_rhs = (q * t + dt * dot(grad(q), u) * t) * dx
d = -dt * diffusivity * dot(grad(q), grad(p)) * dx
diff_matrix = M - 0.5 * d
diff_rhs = action(M + 0.5 * d, t)
# Generate code for mass and rhs assembly.
mass = compile_form(M, "mass")[0]
adv_rhs = compile_form(adv_rhs, "adv_rhs")[0]
diff_matrix = compile_form(diff_matrix, "diff_matrix")[0]
diff_rhs = compile_form(diff_rhs, "diff_rhs")[0]
# Set up simulation data structures
valuetype = np.float64
nodes, coords, elements, elem_node = read_triangle(kwargs['mesh'])
num_nodes = nodes.size
sparsity = op2.Sparsity((elem_node, elem_node), 1, "sparsity")
mat = op2.Mat(sparsity, valuetype, "mat")
tracer_vals = np.asarray([0.0] * num_nodes, dtype=valuetype)
tracer = op2.Dat(nodes, 1, tracer_vals, valuetype, "tracer")
b_vals = np.asarray([0.0] * num_nodes, dtype=valuetype)
b = op2.Dat(nodes, 1, b_vals, valuetype, "b")
velocity_vals = np.asarray([1.0, 0.0] * num_nodes, dtype=valuetype)
velocity = op2.Dat(nodes, 2, velocity_vals, valuetype, "velocity")
# Set initial condition
i_cond_code = """
void i_cond(double *c, double *t)
{
double i_t = 0.01; // Initial time
double A = 0.1; // Normalisation
double D = 0.1; // Diffusivity
double pi = 3.141459265358979;
double x = c[0]-0.5;
double y = c[1]-0.5;
double r = sqrt(x*x+y*y);
if (r<0.25)
*t = A*(exp((-(r*r))/(4*D*i_t))/(4*pi*D*i_t));
else
*t = 0.0;
}
"""
i_cond = op2.Kernel(i_cond_code, "i_cond")
op2.par_loop(i_cond, nodes, coords(op2.IdentityMap, op2.READ),
tracer(op2.IdentityMap, op2.WRITE))
zero_dat_code = """
void zero_dat(double *dat)
{
*dat = 0.0;
}
"""
zero_dat = op2.Kernel(zero_dat_code, "zero_dat")
# Assemble and solve
have_advection = True
have_diffusion = True
def timestep_iteration():
# Advection
if have_advection:
tic('advection')
tic('assembly')
mat.zero()
op2.par_loop(
mass, elements(3, 3),
mat((elem_node[op2.i[0]], elem_node[op2.i[1]]), op2.INC),
coords(elem_node, op2.READ))
op2.par_loop(zero_dat, nodes, b(op2.IdentityMap, op2.WRITE))
op2.par_loop(adv_rhs, elements(3), b(elem_node[op2.i[0]], op2.INC),
coords(elem_node, op2.READ),
tracer(elem_node, op2.READ),
velocity(elem_node, op2.READ))
toc('assembly')
tic('solve')
op2.solve(mat, tracer, b)
toc('solve')
toc('advection')
# Diffusion
if have_diffusion:
tic('diffusion')
tic('assembly')
mat.zero()
op2.par_loop(
diff_matrix, elements(3, 3),
mat((elem_node[op2.i[0]], elem_node[op2.i[1]]), op2.INC),
coords(elem_node, op2.READ))
op2.par_loop(zero_dat, nodes, b(op2.IdentityMap, op2.WRITE))
op2.par_loop(diff_rhs, elements(3), b(elem_node[op2.i[0]],
op2.INC),
coords(elem_node, op2.READ),
tracer(elem_node, op2.READ))
toc('assembly')
tic('solve')
op2.solve(mat, tracer, b)
toc('solve')
toc('diffusion')
# Perform 1 iteration to warm up plan cache then reset initial condition
timestep_iteration()
op2.par_loop(i_cond, nodes, coords(op2.IdentityMap, op2.READ),
tracer(op2.IdentityMap, op2.WRITE))
reset()
# Timed iteration
t1 = clock()
while current_time < endtime:
timestep_iteration()
current_time += dt
runtime = clock() - t1
print "/fluidity :: %f" % runtime
summary('profile_pyop2_%s_%s.csv' %
(opt['mesh'].split('/')[-1], opt['backend']))
def xtr_mat(xtr_elem_node, xtr_dnodes):
sparsity = op2.Sparsity((xtr_dnodes, xtr_dnodes),
(xtr_elem_node, xtr_elem_node), "xtr_sparsity")
return op2.Mat(sparsity, valuetype, "xtr_mat")
def test_sparsities_same_map_pair_cached(self, m1, ds2):
"""Sparsities with the same map pair should share a C handle."""
sp1 = op2.Sparsity((ds2, ds2), (m1, m1))
sp2 = op2.Sparsity((ds2, ds2), (m1, m1))
assert sp1 is sp2
def test_sparsities_different_ordered_map_tuple_cached(self, m1, m2, ds2):
"Sparsities with the same tuple of map pairs should share a C handle."
sp1 = op2.Sparsity((ds2, ds2), ((m1, m1), (m2, m2)))
sp2 = op2.Sparsity((ds2, ds2), ((m2, m2), (m1, m1)))
assert sp1 is sp2
def mat(elem_node, dnodes):
sparsity = op2.Sparsity((dnodes, dnodes), (elem_node, elem_node), "sparsity")
return op2.Mat(sparsity, valuetype, "mat")
def mat(cls, iter2ind1):
sparsity = op2.Sparsity(iter2ind1.toset, iter2ind1, "sparsity")
return op2.Mat(sparsity, 'float64', "mat")
