def nodes(self):
'''The list of nodes at which this boundary condition applies.'''
def hermite_stride(bcnodes):
if isinstance(self._function_space.finat_element, finat.Hermite) and \
self._function_space.mesh().topological_dimension() == 1:
return bcnodes[::2] # every second dof is the vertex value
else:
return bcnodes
sub_d = self.sub_domain
if isinstance(sub_d, str):
return hermite_stride(self._function_space.boundary_nodes(sub_d, self.method))
else:
sub_d = as_tuple(sub_d)
sub_d = [as_tuple(s) for s in sub_d]
bcnodes = []
for s in sub_d:
# s is of one of the following formats:
# facet: (i, )
# edge: (i, j)
# vertex: (i, j, k)
# take intersection of facet nodes, and add it to bcnodes
bcnodes1 = []
if len(s) > 1 and not isinstance(self._function_space.finat_element, finat.Lagrange):
raise TypeError("Currently, edge conditions have only been tested with Lagrange elements")
for ss in s:
# intersection of facets
# Edge conditions have only been tested with Lagrange elements.
# Need to expand the list.
bcnodes1.append(hermite_stride(self._function_space.boundary_nodes(ss, self.method)))
bcnodes1 = functools.reduce(np.intersect1d, bcnodes1)
bcnodes.append(bcnodes1)
return np.concatenate(tuple(bcnodes))
def nodes(self):
'''The list of nodes at which this boundary condition applies.'''
def hermite_stride(bcnodes):
if isinstance(self._function_space.finat_element, finat.Hermite) and \
self._function_space.mesh().topological_dimension() == 1:
return bcnodes[::2] # every second dof is the vertex value
else:
return bcnodes
sub_d = self.sub_domain
if isinstance(sub_d, str):
return hermite_stride(self._function_space.boundary_nodes(sub_d, self.method))
else:
sub_d = as_tuple(sub_d)
sub_d = [as_tuple(s) for s in sub_d]
bcnodes = []
for s in sub_d:
# s is of one of the following formats:
# facet: (i, )
# edge: (i, j)
# vertex: (i, j, k)
# take intersection of facet nodes, and add it to bcnodes
bcnodes1 = []
if len(s) > 1 and not isinstance(self._function_space.finat_element, finat.Lagrange):
raise TypeError("Currently, edge conditions have only been tested with Lagrange elements")
for ss in s:
# intersection of facets
# Edge conditions have only been tested with Lagrange elements.
# Need to expand the list.
bcnodes1.append(hermite_stride(self._function_space.boundary_nodes(ss, self.method)))
bcnodes1 = functools.reduce(np.intersect1d, bcnodes1)
bcnodes.append(bcnodes1)
return np.concatenate(tuple(bcnodes))
def __getitem__(self, key):
key = as_tuple(key)
# Make indexing with too few indices legal.
key = key + tuple(
slice(None) for i in range(self.tensor.rank - len(key)))
if len(key) > self.tensor.rank:
raise ValueError(
"Attempting to index a rank-%s tensor with %s indices." %
(self.tensor.rank, len(key)))
block_shape = tuple(len(V) for V in self.tensor.arg_function_spaces)
# Convert slice indices to tuple of indices.
blocks = tuple(
as_tuple(range(k.stop)[k] if isinstance(k, slice) else k)
for k, n in zip(key, block_shape))
if blocks == tuple(tuple(range(n)) for n in block_shape):
return self.tensor
# Avoid repeated instantiation of an equivalent block
try:
block = self.block_cache[blocks]
except KeyError:
block = Block(tensor=self.tensor, indices=blocks)
self.block_cache[blocks] = block
return block
def domain_args(self):
r"""The sub_domain the BC applies to."""
# Define facet, edge, vertex using tuples:
# Ex in 3D:
# user input returned keys
# facet = ((1, ), ) -> ((2, ((1, ), )), (1, ()), (0, ()))
# edge = ((1, 2), ) -> ((2, ()), (1, ((1, 2), )), (0, ()))
# vertex = ((1, 2, 4), ) -> ((2, ()), (1, ()), (0, ((1, 2, 4), ))
#
# Multiple facets:
# (1, 2, 4) := ((1, ), (2, ), (4,)) -> ((2, ((1, ), (2, ), (4, ))), (1, ()), (0, ()))
#
# One facet and two edges:
# ((1,), (1, 3), (1, 4)) -> ((2, ((1,),)), (1, ((1,3), (1, 4))), (0, ()))
#
sub_d = self.sub_domain
# if string, return
if isinstance(sub_d, str):
return (sub_d, )
# convert: i -> (i, )
sub_d = as_tuple(sub_d)
# convert: (i, j, (k, l)) -> ((i, ), (j, ), (k, l))
sub_d = [as_tuple(i) for i in sub_d]
ndim = self.function_space().mesh().topology_dm.getDimension()
sd = [[] for _ in range(ndim)]
for i in sub_d:
sd[ndim - len(i)].append(i)
s = []
for i in range(ndim):
s.append((ndim - 1 - i, as_tuple(sd[i])))
return as_tuple(s)
def domain_args(self):
r"""The sub_domain the BC applies to."""
# Define facet, edge, vertex using tuples:
# Ex in 3D:
# user input returned keys
# facet = ((1, ), ) -> ((2, ((1, ), )), (1, ()), (0, ()))
# edge = ((1, 2), ) -> ((2, ()), (1, ((1, 2), )), (0, ()))
# vertex = ((1, 2, 4), ) -> ((2, ()), (1, ()), (0, ((1, 2, 4), ))
#
# Multiple facets:
# (1, 2, 4) := ((1, ), (2, ), (4,)) -> ((2, ((1, ), (2, ), (4, ))), (1, ()), (0, ()))
#
# One facet and two edges:
# ((1,), (1, 3), (1, 4)) -> ((2, ((1,),)), (1, ((1,3), (1, 4))), (0, ()))
#
sub_d = self.sub_domain
# if string, return
if isinstance(sub_d, str):
return (sub_d, )
# convert: i -> (i, )
sub_d = as_tuple(sub_d)
# convert: (i, j, (k, l)) -> ((i, ), (j, ), (k, l))
sub_d = [as_tuple(i) for i in sub_d]
ndim = self.function_space().mesh()._plex.getDimension()
sd = [[] for _ in range(ndim)]
for i in sub_d:
sd[ndim - len(i)].append(i)
s = []
for i in range(ndim):
s.append((ndim - 1 - i, as_tuple(sd[i])))
return as_tuple(s)
def _needs_reassembly(self):
"""Does this :class:`Matrix` need reassembly.
The :class:`Matrix` needs reassembling if the subdomains over
which boundary conditions were applied the last time it was
assembled are different from the subdomains of the current set
of boundary conditions.
"""
old_subdomains = set(flatten(as_tuple(bc.sub_domain) for bc in self._bcs_at_point_of_assembly))
new_subdomains = set(flatten(as_tuple(bc.sub_domain) for bc in self.bcs))
return old_subdomains != new_subdomains
def _needs_reassembly(self):
"""Does this :class:`Matrix` need reassembly.
The :class:`Matrix` needs reassembling if the subdomains over
which boundary conditions were applied the last time it was
assembled are different from the subdomains of the current set
of boundary conditions.
"""
old_subdomains = set(flatten(as_tuple(bc.sub_domain)
for bc in self._bcs_at_point_of_assembly))
new_subdomains = set(flatten(as_tuple(bc.sub_domain)
for bc in self.bcs))
return old_subdomains != new_subdomains
def __init__(self, mesh, classes, kind, facet_cell, local_facet_number, markers=None,
unique_markers=None):
self.mesh = mesh
classes = as_tuple(classes, int, 4)
self.classes = classes
self.kind = kind
assert(kind in ["interior", "exterior"])
if kind == "interior":
self._rank = 2
else:
self._rank = 1
self.facet_cell = facet_cell
self.local_facet_number = local_facet_number
# assert that markers is a proper subset of unique_markers
if markers is not None:
for marker in markers:
assert (marker in unique_markers), \
"Every marker has to be contained in unique_markers"
self.markers = markers
self.unique_markers = [] if unique_markers is None else unique_markers
self._subsets = {}
def c_addto(self,
i,
j,
buf_name,
tmp_name,
tmp_decl,
extruded=None,
is_facet=False,
applied_blas=False):
# Override global c_addto to index the map locally rather than globally.
# Replaces MatSetValuesLocal with MatSetValues
from pyop2.utils import as_tuple
maps = as_tuple(self.map, op2.Map)
nrows = maps[0].split[i].arity
ncols = maps[1].split[j].arity
rows_str = "%s + n * %s" % (self.c_map_name(0, i), nrows)
cols_str = "%s + n * %s" % (self.c_map_name(1, j), ncols)
if extruded is not None:
rows_str = extruded + self.c_map_name(0, i)
cols_str = extruded + self.c_map_name(1, j)
if is_facet:
nrows *= 2
ncols *= 2
ret = []
rbs, cbs = self.data.sparsity[i, j].dims[0][0]
rdim = rbs * nrows
addto_name = buf_name
addto = 'MatSetValues'
if self.data._is_vector_field:
addto = 'MatSetValuesBlocked'
rmap, cmap = maps
rdim, cdim = self.data.dims[i][j]
if rmap.vector_index is not None or cmap.vector_index is not None:
raise NotImplementedError
ret.append(
"""%(addto)s(%(mat)s, %(nrows)s, %(rows)s,
%(ncols)s, %(cols)s,
(const PetscScalar *)%(vals)s,
%(insert)s);""" % {
'mat':
self.c_arg_name(i, j),
'vals':
addto_name,
'addto':
addto,
'nrows':
nrows,
'ncols':
ncols,
'rows':
rows_str,
'cols':
cols_str,
'insert':
"INSERT_VALUES" if self.access == op2.WRITE else "ADD_VALUES"
})
return "\n".join(ret)
def sanitise_input(v, V):
if isinstance(v, expression.Expression):
shape = v.value_shape()
# Build a function space that supports PointEvaluation so that
# we can interpolate into it.
deg = max(as_tuple(V.ufl_element().degree()))
if v.rank() == 0:
fs = functionspace.FunctionSpace(V.mesh(), 'DG', deg+1)
elif v.rank() == 1:
fs = functionspace.VectorFunctionSpace(V.mesh(), 'DG',
deg+1,
dim=shape[0])
else:
fs = functionspace.TensorFunctionSpace(V.mesh(), 'DG',
deg+1,
shape=shape)
f = function.Function(fs)
f.interpolate(v)
return f
elif isinstance(v, function.Function):
return v
elif isinstance(v, ufl.classes.Expr):
return v
else:
raise ValueError("Can't project from source object %r" % v)
def __init__(self, problems):
problems = as_tuple(problems)
self._problems = problems
# Build the jacobian with the correct sparsity pattern. Note
# that since matrix assembly is lazy this doesn't actually
# force an additional assembly of the matrix since in
# form_jacobian we call assemble again which drops this
# computation on the floor.
from firedrake.assemble import assemble
self._jacs = tuple(assemble(problem.J, bcs=problem.bcs,
form_compiler_parameters=problem.form_compiler_parameters,
nest=problem._nest)
for problem in problems)
if problems[-1].Jp is not None:
self._pjacs = tuple(assemble(problem.Jp, bcs=problem.bcs,
form_compiler_parameters=problem.form_compiler_parameters,
nest=problem._nest)
for problem in problems)
else:
self._pjacs = self._jacs
# Function to hold current guess
self._xs = tuple(function.Function(problem.u) for problem in problems)
self.Fs = tuple(ufl.replace(problem.F, {problem.u: x}) for problem, x in zip(problems,
self._xs))
self.Js = tuple(ufl.replace(problem.J, {problem.u: x}) for problem, x in zip(problems,
self._xs))
if problems[-1].Jp is not None:
self.Jps = tuple(ufl.replace(problem.Jp, {problem.u: x}) for problem, x in zip(problems,
self._xs))
else:
self.Jps = tuple(None for _ in problems)
self._Fs = tuple(function.Function(F.arguments()[0].function_space())
for F in self.Fs)
self._jacobians_assembled = [False for _ in problems]
def c_add_offset_map(self, is_facet=False):
if self._is_mat:
dsets = self.data.sparsity.dsets
else:
dsets = (self.data.dataset, )
val = []
for i, (map, dset) in enumerate(zip(as_tuple(self.map, Map), dsets)):
if not map.iterset._extruded:
continue
for j, (m, d) in enumerate(zip(map, dset)):
for idx in range(m.arity):
val.append(
"xtr_%(name)s[%(ind)s] += %(off)d;" % {
'name': self.c_map_name(i, j),
'off': m.offset[idx],
'ind': idx
})
if is_facet:
for idx in range(m.arity):
val.append(
"xtr_%(name)s[%(ind)s] += %(off)d;" % {
'name': self.c_map_name(i, j),
'off': m.offset[idx],
'ind': m.arity + idx
})
return '\n'.join(val) + '\n'
def c_map_init(self, is_top=False, is_facet=False):
if self._is_mat:
dsets = self.data.sparsity.dsets
else:
dsets = (self.data.dataset, )
val = []
for i, (map, dset) in enumerate(zip(as_tuple(self.map, Map), dsets)):
for j, (m, d) in enumerate(zip(map, dset)):
for idx in range(m.arity):
val.append(
"xtr_%(name)s[%(ind)s] = *(%(name)s + i * %(dim)s + %(ind)s)%(off_top)s;"
% {
'name':
self.c_map_name(i, j),
'dim':
m.arity,
'ind':
idx,
'off_top':
' + start_layer * ' +
str(m.offset[idx]) if is_top else ''
})
if is_facet:
for idx in range(m.arity):
val.append(
"xtr_%(name)s[%(ind)s] = *(%(name)s + i * %(dim)s + %(ind_zero)s)%(off_top)s%(off)s;"
% {
'name': self.c_map_name(i, j),
'dim': m.arity,
'ind': idx + m.arity,
'ind_zero': idx,
'off_top': ' + start_layer' if is_top else '',
'off': ' + ' + str(m.offset[idx])
})
return '\n'.join(val) + '\n'
def __init__(self, mesh, classes, kind, facet_cell, local_facet_number, markers=None,
unique_markers=None):
self.mesh = mesh
classes = as_tuple(classes, int, 4)
self.classes = classes
self.kind = kind
assert(kind in ["interior", "exterior"])
if kind == "interior":
self._rank = 2
else:
self._rank = 1
self.facet_cell = facet_cell
self.local_facet_number = local_facet_number
# assert that markers is a proper subset of unique_markers
if markers is not None:
for marker in markers:
assert (marker in unique_markers), \
"Every marker has to be contained in unique_markers"
self.markers = markers
self.unique_markers = [] if unique_markers is None else unique_markers
self._subsets = {}
def sanitise_input(v, V):
if isinstance(v, expression.Expression):
shape = v.value_shape()
# Build a function space that supports PointEvaluation so that
# we can interpolate into it.
deg = max(as_tuple(V.ufl_element().degree()))
if v.rank() == 0:
fs = functionspace.FunctionSpace(V.mesh(), 'DG', deg+1)
elif v.rank() == 1:
fs = functionspace.VectorFunctionSpace(V.mesh(), 'DG',
deg+1,
dim=shape[0])
else:
fs = functionspace.TensorFunctionSpace(V.mesh(), 'DG',
deg+1,
shape=shape)
f = function.Function(fs)
f.interpolate(v)
return f
elif isinstance(v, function.Function):
return v
elif isinstance(v, ufl.classes.Expr):
return v
else:
raise ValueError("Can't project from source object %r" % v)
def c_wrapper_arg(self):
if self._is_mat:
val = "Mat %s_" % self.c_arg_name()
else:
val = ', '.join(["%s *%s" % (self.ctype, self.c_arg_name(i))
for i in range(len(self.data))])
if self._is_indirect or self._is_mat:
for i, map in enumerate(as_tuple(self.map, Map)):
if map is not None:
for j, m in enumerate(map):
val += ", int *%s" % self.c_map_name(i, j)
return val
def bcs(self, bcs):
"""Attach some boundary conditions to this :class:`MatrixBase`.
:arg bcs: a boundary condition (of type
:class:`.DirichletBC`), or an iterable of boundary
conditions. If bcs is None, erase all boundary conditions
on the :class:`.MatrixBase`.
"""
if bcs is not None:
self._bcs = tuple(itertools.chain(*(as_tuple(bc) for bc in bcs)))
else:
self._bcs = ()
def c_map_decl(self, is_facet=False):
if self._is_mat:
dsets = self.data.sparsity.dsets
else:
dsets = (self.data.dataset,)
val = []
for i, (map, dset) in enumerate(zip(as_tuple(self.map, Map), dsets)):
for j, (m, d) in enumerate(zip(map, dset)):
dim = m.arity
if is_facet:
dim *= 2
val.append("int xtr_%(name)s[%(dim)s];" %
{'name': self.c_map_name(i, j), 'dim': dim})
return '\n'.join(val)+'\n'
def subset(self, markers):
"""Return the subset corresponding to a given marker value.
:param markers: integer marker id or an iterable of marker ids"""
if self.markers is None:
return self._null_subset
markers = as_tuple(markers, int)
try:
return self._subsets[markers]
except KeyError:
indices = np.concatenate([np.nonzero(self.markers == i)[0]
for i in markers])
self._subsets[markers] = op2.Subset(self.set, indices)
return self._subsets[markers]
def c_addto(self,
i,
j,
buf_name,
tmp_name,
tmp_decl,
extruded=None,
is_facet=False,
applied_blas=False):
# Override global c_addto to index the map locally rather than globally.
# Replaces MatSetValuesLocal with MatSetValues
from pyop2.utils import as_tuple
rmap, cmap = as_tuple(self.map, op2.Map)
rset, cset = self.data.sparsity.dsets
nrows = sum(m.arity * s.cdim for m, s in zip(rmap, rset))
ncols = sum(m.arity * s.cdim for m, s in zip(cmap, cset))
rows_str = "%s + n * %s" % (self.c_map_name(0, i), nrows)
cols_str = "%s + n * %s" % (self.c_map_name(1, j), ncols)
if extruded is not None:
raise NotImplementedError("Not for extruded right now")
if is_facet:
raise NotImplementedError("Not for interior facets and extruded")
ret = []
addto_name = buf_name
if rmap.vector_index is not None or cmap.vector_index is not None:
raise NotImplementedError
ret.append(
"""MatSetValues(%(mat)s, %(nrows)s, %(rows)s,
%(ncols)s, %(cols)s,
(const PetscScalar *)%(vals)s,
%(insert)s);""" % {
'mat':
self.c_arg_name(i, j),
'vals':
addto_name,
'nrows':
nrows,
'ncols':
ncols,
'rows':
rows_str,
'cols':
cols_str,
'insert':
"INSERT_VALUES" if self.access == op2.WRITE else "ADD_VALUES"
})
return "\n".join(ret)
def __new__(cls, tensor, indices):
if not isinstance(tensor, TensorBase):
raise TypeError("Can only extract blocks of Slate tensors.")
if len(indices) != tensor.rank:
raise ValueError("Length of indices must be equal to the tensor rank.")
if not all(0 <= i < len(arg.function_space())
for arg, idx in zip(tensor.arguments(), indices) for i in as_tuple(idx)):
raise ValueError("Indices out of range.")
if not tensor.is_mixed:
return tensor
return super().__new__(cls)
def __new__(cls, tensor, indices):
if not isinstance(tensor, TensorBase):
raise TypeError("Can only extract blocks of Slate tensors.")
if len(indices) != tensor.rank:
raise ValueError("Length of indices must be equal to the tensor rank.")
if not all(0 <= i < len(arg.function_space())
for arg, idx in zip(tensor.arguments(), indices) for i in as_tuple(idx)):
raise ValueError("Indices out of range.")
if not tensor.is_mixed:
return tensor
return super().__new__(cls)
def boundary_nodes(self, V, sub_domain):
if sub_domain in ["bottom", "top"]:
if not V.extruded:
raise ValueError("Invalid subdomain '%s' for non-extruded mesh",
sub_domain)
entity_dofs = eutils.flat_entity_dofs(V.finat_element.entity_dofs())
key = (entity_dofs_key(entity_dofs), sub_domain)
return get_top_bottom_boundary_nodes(V.mesh(), key, V)
else:
if sub_domain == "on_boundary":
sdkey = sub_domain
else:
sdkey = as_tuple(sub_domain)
key = (entity_dofs_key(V.finat_element.entity_dofs()), sdkey)
return get_facet_closure_nodes(V.mesh(), key, V)
def c_wrapper_arg(self):
if self._is_mat:
val = "Mat %s_" % self.c_arg_name()
else:
val = ', '.join([
"%s *%s" % (self.ctype, self.c_arg_name(i))
for i in range(len(self.data))
])
if self._is_indirect or self._is_mat:
for i, map in enumerate(as_tuple(self.map, Map)):
if map is not None:
for j, m in enumerate(map):
val += ", %s *%s" % (as_cstr(IntType),
self.c_map_name(i, j))
return val
def boundary_nodes(self, V, sub_domain, method):
if method not in {"topological", "geometric"}:
raise ValueError("Don't know how to extract nodes with method '%s'", method)
if sub_domain in ["bottom", "top"]:
if not V.extruded:
raise ValueError("Invalid subdomain '%s' for non-extruded mesh",
sub_domain)
entity_dofs = eutils.flat_entity_dofs(V.finat_element.entity_dofs())
key = (entity_dofs_key(entity_dofs), sub_domain, method)
return get_top_bottom_boundary_nodes(V.mesh(), key, V)
else:
if sub_domain == "on_boundary":
sdkey = sub_domain
else:
sdkey = as_tuple(sub_domain)
key = (entity_dofs_key(V.finat_element.entity_dofs()), sdkey, method)
return get_boundary_nodes(V.mesh(), key, V)
def c_map_decl(self, is_facet=False):
if self._is_mat:
dsets = self.data.sparsity.dsets
else:
dsets = (self.data.dataset, )
val = []
for i, (map, dset) in enumerate(zip(as_tuple(self.map, Map), dsets)):
for j, (m, d) in enumerate(zip(map, dset)):
dim = m.arity
if is_facet:
dim *= 2
val.append(
"%(IntType)s xtr_%(name)s[%(dim)s];" % {
'name': self.c_map_name(i, j),
'dim': dim,
'IntType': as_cstr(IntType)
})
return '\n'.join(val) + '\n'
def nodes(self):
dm = self.function_space().mesh().topology_dm
section = self.function_space().dm.getDefaultSection()
nodes = []
for sd in as_tuple(self.sub_domain):
nfaces = dm.getStratumSize(FACE_SETS_LABEL, sd)
faces = dm.getStratumIS(FACE_SETS_LABEL, sd)
if nfaces == 0:
continue
for face in faces.indices:
# if dm.getLabelValue("interior_facets", face) < 0:
# continue
closure, _ = dm.getTransitiveClosure(face)
for p in closure:
dof = section.getDof(p)
offset = section.getOffset(p)
nodes.extend((offset + d) for d in range(dof))
return np.unique(np.asarray(nodes, dtype=IntType))
def subset(self, markers):
"""Return the subset corresponding to a given marker value.
:param markers: integer marker id or an iterable of marker ids"""
if self.markers is None:
return self._null_subset
markers = as_tuple(markers, int)
try:
return self._subsets[markers]
except KeyError:
# check that the given markers are valid
for marker in markers:
if marker not in self.unique_markers:
raise LookupError("{0} is not a valid marker".format(marker))
# build a list of indices corresponding to the subsets selected by
# markers
indices = np.concatenate([np.nonzero(self.markers == i)[0] for i in markers])
self._subsets[markers] = op2.Subset(self.set, indices)
return self._subsets[markers]
def __init__(self, problems):
problems = as_tuple(problems)
self._problems = problems
# Build the jacobian with the correct sparsity pattern. Note
# that since matrix assembly is lazy this doesn't actually
# force an additional assembly of the matrix since in
# form_jacobian we call assemble again which drops this
# computation on the floor.
from firedrake.assemble import assemble
self._jacs = tuple(
assemble(problem.J,
bcs=problem.bcs,
form_compiler_parameters=problem.form_compiler_parameters,
nest=problem._nest) for problem in problems)
if problems[-1].Jp is not None:
self._pjacs = tuple(
assemble(
problem.Jp,
bcs=problem.bcs,
form_compiler_parameters=problem.form_compiler_parameters,
nest=problem._nest) for problem in problems)
else:
self._pjacs = self._jacs
# Function to hold current guess
self._xs = tuple(function.Function(problem.u) for problem in problems)
self.Fs = tuple(
ufl.replace(problem.F, {problem.u: x})
for problem, x in zip(problems, self._xs))
self.Js = tuple(
ufl.replace(problem.J, {problem.u: x})
for problem, x in zip(problems, self._xs))
if problems[-1].Jp is not None:
self.Jps = tuple(
ufl.replace(problem.Jp, {problem.u: x})
for problem, x in zip(problems, self._xs))
else:
self.Jps = tuple(None for _ in problems)
self._Fs = tuple(
function.Function(F.arguments()[0].function_space())
for F in self.Fs)
self._jacobians_assembled = [False for _ in problems]
def set_argtypes(self, iterset, *args):
argtypes = [slope.Executor.meta['py_ctype_exec']]
for itspace in self._all_itspaces:
if isinstance(itspace.iterset, base.Subset):
argtypes.append(itspace.iterset._argtype)
for arg in args:
if arg._is_mat:
argtypes.append(arg.data._argtype)
else:
for d in arg.data:
argtypes.append(d._argtype)
if arg._is_indirect or arg._is_mat:
maps = as_tuple(arg.map, base.Map)
for map in maps:
for m in map:
argtypes.append(m._argtype)
# MPI related stuff (rank, region)
argtypes.append(ctypes.c_int)
argtypes.append(ctypes.c_int)
self._argtypes = argtypes
def __init__(self, mesh, classes, kind, facet_cell, local_facet_number, markers=None,
unique_markers=None):
self.mesh = mesh
classes = as_tuple(classes, int, 4)
self.classes = classes
self.kind = kind
assert(kind in ["interior", "exterior"])
if kind == "interior":
self._rank = 2
else:
self._rank = 1
self.facet_cell = facet_cell
self.local_facet_number = local_facet_number
self.markers = markers
self.unique_markers = [] if unique_markers is None else unique_markers
self._subsets = {}
def subset(self, markers):
"""Return the subset corresponding to a given marker value.
:param markers: integer marker id or an iterable of marker ids"""
if self.markers is None:
return self._null_subset
markers = as_tuple(markers, int)
try:
return self._subsets[markers]
except KeyError:
# check that the given markers are valid
for marker in markers:
if marker not in self.unique_markers:
raise LookupError(
'{0} is not a valid marker'.format(marker))
# build a list of indices corresponding to the subsets selected by
# markers
indices = np.concatenate(
[np.nonzero(self.markers == i)[0] for i in markers])
self._subsets[markers] = op2.Subset(self.set, indices)
return self._subsets[markers]
def __init__(self, arg, loop_position, gtl_maps=None):
"""Initialize a :class:`TilingArg`.
:arg arg: a supertype of :class:`TilingArg`, from which this Arg is derived.
:arg loop_position: the position of the loop in the loop chain that this
object belongs to.
:arg gtl_maps: a dict associating global map names to local map names.
"""
super(TilingArg, self).__init__(arg)
self.position = arg.position
self.indirect_position = arg.indirect_position
self.loop_position = loop_position
c_local_maps = None
maps = as_tuple(arg.map, base.Map)
if gtl_maps:
c_local_maps = [None]*len(maps)
for i, map in enumerate(maps):
c_local_maps[i] = [None]*len(map)
for j, m in enumerate(map):
c_local_maps[i][j] = gtl_maps["%s%d_%d" % (m.name, i, j)]
self._c_local_maps = c_local_maps
def c_add_offset_map(self, is_facet=False):
if self._is_mat:
dsets = self.data.sparsity.dsets
else:
dsets = (self.data.dataset,)
val = []
for i, (map, dset) in enumerate(zip(as_tuple(self.map, Map), dsets)):
if not map.iterset._extruded:
continue
for j, (m, d) in enumerate(zip(map, dset)):
for idx in range(m.arity):
val.append("xtr_%(name)s[%(ind)s] += %(off)d;" %
{'name': self.c_map_name(i, j),
'off': m.offset[idx],
'ind': idx})
if is_facet:
for idx in range(m.arity):
val.append("xtr_%(name)s[%(ind)s] += %(off)d;" %
{'name': self.c_map_name(i, j),
'off': m.offset[idx],
'ind': m.arity + idx})
return '\n'.join(val)+'\n'
def set_argtypes(self, iterset, *args):
argtypes = [ctypes.c_int, ctypes.c_int]
if isinstance(iterset, Subset):
argtypes.append(iterset._argtype)
for arg in args:
if arg._is_mat:
argtypes.append(arg.data._argtype)
else:
for d in arg.data:
argtypes.append(d._argtype)
if arg._is_indirect or arg._is_mat:
maps = as_tuple(arg.map, Map)
for map in maps:
if map is not None:
for m in map:
argtypes.append(m._argtype)
if iterset._extruded:
argtypes.append(ctypes.c_int)
argtypes.append(ctypes.c_int)
self._argtypes = argtypes
def set_argtypes(self, iterset, *args):
index_type = as_ctypes(IntType)
argtypes = [index_type, index_type]
if isinstance(iterset, Subset):
argtypes.append(iterset._argtype)
for arg in args:
if arg._is_mat:
argtypes.append(arg.data._argtype)
else:
for d in arg.data:
argtypes.append(d._argtype)
if arg._is_indirect or arg._is_mat:
maps = as_tuple(arg.map, Map)
for map in maps:
if map is not None:
for m in map:
argtypes.append(m._argtype)
if iterset._extruded:
argtypes.append(index_type)
argtypes.append(index_type)
self._argtypes = argtypes
def __getitem__(self, key):
key = list(as_tuple(key))
# Make indexing with too few indices legal.
key += [slice(None) for i in range(self.tensor.rank - len(key))]
if len(key) > self.tensor.rank:
raise ValueError("Attempting to index a rank-%s tensor with %s indices."
% (self.tensor.rank, len(key)))
# Convert slice indices to tuple of indices.
blocks = tuple(range(n)[k] if isinstance(k, slice) else k
for k, n in zip(key, self.tensor.shape))
# Avoid repeated instantiation of an equivalent block
try:
block = self.block_cache[blocks]
except KeyError:
block = Block(tensor=self.tensor, indices=blocks)
self.block_cache[blocks] = block
return block
def __getitem__(self, key):
key = list(as_tuple(key))
# Make indexing with too few indices legal.
key += [slice(None) for i in range(self.tensor.rank - len(key))]
if len(key) > self.tensor.rank:
raise ValueError("Attempting to index a rank-%s tensor with %s indices."
% (self.tensor.rank, len(key)))
# Convert slice indices to tuple of indices.
blocks = tuple(range(n)[k] if isinstance(k, slice) else k
for k, n in zip(key, self.tensor.shape))
# Avoid repeated instantiation of an equivalent block
try:
block = self.block_cache[blocks]
except KeyError:
block = Block(tensor=self.tensor, indices=blocks)
self.block_cache[blocks] = block
return block
def _split_arguments(self):
"""Splits the function space and stores the component
spaces determined by the indices.
"""
from firedrake.functionspace import FunctionSpace, MixedFunctionSpace
from firedrake.ufl_expr import Argument
tensor, = self.operands
nargs = []
for i, arg in enumerate(tensor.arguments()):
V = arg.function_space()
V_is = V.split()
idx = as_tuple(self._blocks[i])
if len(idx) == 1:
fidx, = idx
W = V_is[fidx]
W = FunctionSpace(W.mesh(), W.ufl_element())
else:
W = MixedFunctionSpace([V_is[fidx] for fidx in idx])
nargs.append(Argument(W, arg.number(), part=arg.part()))
return tuple(nargs)
def _split_arguments(self):
"""Splits the function space and stores the component
spaces determined by the indices.
"""
from firedrake.functionspace import FunctionSpace, MixedFunctionSpace
from firedrake.ufl_expr import Argument
tensor, = self.operands
nargs = []
for i, arg in enumerate(tensor.arguments()):
V = arg.function_space()
V_is = V.split()
idx = as_tuple(self._blocks[i])
if len(idx) == 1:
fidx, = idx
W = V_is[fidx]
W = FunctionSpace(W.mesh(), W.ufl_element())
else:
W = MixedFunctionSpace([V_is[fidx] for fidx in idx])
nargs.append(Argument(W, arg.number(), part=arg.part()))
return tuple(nargs)
def prepare_arglist(self, part, *args):
arglist = [self._inspection]
for itspace in self._all_itspaces:
if isinstance(itspace._iterset, base.Subset):
arglist.append(itspace._iterset._indices.ctypes.data)
for arg in args:
if arg._is_mat:
arglist.append(arg.data.handle.handle)
else:
for d in arg.data:
# Cannot access a property of the Dat or we will force
# evaluation of the trace
arglist.append(d._data.ctypes.data)
if arg._is_indirect or arg._is_mat:
maps = as_tuple(arg.map, base.Map)
for map in maps:
for m in map:
arglist.append(m._values.ctypes.data)
arglist.append(self.it_space.comm.rank)
return arglist
def get_sup_element(*elements, continuous=False, max_degree=None):
"""Given ufl elements and a continuity flag, return
a new ufl element that contains all elements.
:arg elements: ufl elements.
:continous: A flag indicating if all elements are continous.
:returns: A ufl element containing all elements.
"""
try:
cell, = set(e.cell() for e in elements)
except ValueError:
raise ValueError("All cells must be identical")
degree = max(chain(*(as_tuple(e.degree()) for e in elements)))
if continuous:
family = "CG"
else:
if cell.cellname() in {"interval", "triangle", "tetrahedron"}:
family = "DG"
else:
family = "DQ"
return ufl.FiniteElement(family,
cell=cell,
degree=degree if max_degree is None else max_degree,
variant="equispaced")
def c_map_init(self, is_top=False, is_facet=False):
if self._is_mat:
dsets = self.data.sparsity.dsets
else:
dsets = (self.data.dataset,)
val = []
for i, (map, dset) in enumerate(zip(as_tuple(self.map, Map), dsets)):
for j, (m, d) in enumerate(zip(map, dset)):
for idx in range(m.arity):
val.append("xtr_%(name)s[%(ind)s] = *(%(name)s + i * %(dim)s + %(ind)s)%(off_top)s;" %
{'name': self.c_map_name(i, j),
'dim': m.arity,
'ind': idx,
'off_top': ' + start_layer * '+str(m.offset[idx]) if is_top else ''})
if is_facet:
for idx in range(m.arity):
val.append("xtr_%(name)s[%(ind)s] = *(%(name)s + i * %(dim)s + %(ind_zero)s)%(off_top)s%(off)s;" %
{'name': self.c_map_name(i, j),
'dim': m.arity,
'ind': idx + m.arity,
'ind_zero': idx,
'off_top': ' + start_layer' if is_top else '',
'off': ' + ' + str(m.offset[idx])})
return '\n'.join(val)+'\n'
def condense_and_forward_eliminate(A, b, elim_fields):
"""Returns Slate expressions for the operator and
right-hand side vector after eliminating specified
unknowns.
:arg A: a `slate.Tensor` corresponding to the
mixed UFL operator.
:arg b: a `firedrake.Function` corresponding
to the right-hand side.
:arg elim_fields: a `tuple` of indices denoting
which fields to eliminate.
"""
if not isinstance(A, slate.Tensor):
raise ValueError("Left-hand operator must be a Slate Tensor")
# Ensures field indices are in increasing order
elim_fields = list(as_tuple(elim_fields))
elim_fields.sort()
all_fields = list(range(len(A.arg_function_spaces[0])))
condensed_fields = list(set(all_fields) - set(elim_fields))
condensed_fields.sort()
_A = A.blocks
_b = slate.AssembledVector(b).blocks
# NOTE: Does not support non-contiguous field elimination
e_idx0 = elim_fields[0]
e_idx1 = elim_fields[-1]
f_idx0 = condensed_fields[0]
f_idx1 = condensed_fields[-1]
# Finite element systems where static condensation
# is possible have the general form:
#
# | A_ee A_ef || x_e | | b_e |
# | || | = | |
# | A_fe A_ff || x_f | | b_f |
#
# where subscript `e` denotes the coupling with fields
# that will be eliminated, and `f` denotes the condensed
# fields.
Aff = _A[f_idx0:f_idx1 + 1, f_idx0:f_idx1 + 1]
Aef = _A[e_idx0:e_idx1 + 1, f_idx0:f_idx1 + 1]
Afe = _A[f_idx0:f_idx1 + 1, e_idx0:e_idx1 + 1]
Aee = _A[e_idx0:e_idx1 + 1, e_idx0:e_idx1 + 1]
bf = _b[f_idx0:f_idx1 + 1]
be = _b[e_idx0:e_idx1 + 1]
# The reduced operator and right-hand side are:
# S = A_ff - A_fe * A_ee.inv * A_ef
# r = b_f - A_fe * A_ee.inv * b_e
# as show in Slate:
S = Aff - Afe * Aee.inv * Aef
r = bf - Afe * Aee.inv * be
field_idx = [idx for idx in range(f_idx0, f_idx1)]
return LAContext(lhs=S, rhs=r, field_idx=field_idx)
def _process_args(cls, *args, **kwargs):
"""Convert list of spaces to tuple (to make it hashable)"""
mesh = args[0][0].mesh()
pargs = tuple(as_tuple(arg) for arg in args)
return (mesh, ) + pargs, kwargs
def backward_solve(A, b, x, reconstruct_fields):
"""Returns a sequence of linear algebra contexts containing
Slate expressions for backwards substitution.
:arg A: a `slate.Tensor` corresponding to the
mixed UFL operator.
:arg b: a `firedrake.Function` corresponding
to the right-hand side.
:arg x: a `firedrake.Function` corresponding
to the solution.
:arg reconstruct_fields: a `tuple` of indices denoting
which fields to reconstruct.
"""
if not isinstance(A, slate.Tensor):
raise ValueError("Left-hand operator must be a Slate Tensor")
all_fields = list(range(len(A.arg_function_spaces[0])))
nfields = len(all_fields)
reconstruct_fields = as_tuple(reconstruct_fields)
_A = A.blocks
_b = b.split()
_x = x.split()
# Ordering matters
systems = []
# Reconstruct one unknown from one determined field:
#
# | A_ee A_ef || x_e | | b_e |
# | || | = | |
# | A_fe A_ff || x_f | | b_f |
#
# where x_f is known from a previous computation.
# Returns the system:
#
# A_ee x_e = b_e - A_ef * x_f.
if nfields == 2:
id_e, = reconstruct_fields
id_f, = [idx for idx in all_fields if idx != id_e]
A_ee = _A[id_e, id_e]
A_ef = _A[id_e, id_f]
b_e = slate.AssembledVector(_b[id_e])
x_f = slate.AssembledVector(_x[id_f])
r_e = b_e - A_ef * x_f
local_system = LAContext(lhs=A_ee, rhs=r_e, field_idx=(id_e,))
systems.append(local_system)
# Reconstruct two unknowns from one determined field:
#
# | A_e0e0 A_e0e1 A_e0f || x_e0 | | b_e0 |
# | A_e1e0 A_e1e1 A_e1f || x_e1 | = | b_e1 |
# | A_fe0 A_fe1 A_ff || x_f | | b_f |
#
# where x_f is the known field. Returns two systems to be
# solved in order (determined from the reverse order of indices
# e0 and e1):
#
# Solve for e1 first (obtained from eliminating x_e0):
#
# S_e1 x_e1 = r_e1
#
# where
#
# S_e1 = A_e1e1 - A_e1e0 * A_e0e0.inv * A_e0e1, and
# r_e1 = b_e1 - A_e1e0 * A_e0e0.inv * b_e0
# - (A_e1f - A_e1e0 * A_e0e0.inv * A_e0f) * x_f,
#
# And then solve for x_e0 given x_f and x_e1:
#
# A_e0e0 x_e0 = b_e0 - A_e0e1 * x_e1 - A_e0f * x_f.
elif nfields == 3:
if len(reconstruct_fields) != nfields - 1:
raise NotImplementedError("Implemented for 1 determined field")
# Order of reconstruction doesn't need to be in order
# of increasing indices
id_e0, id_e1 = reconstruct_fields
id_f, = [idx for idx in all_fields if idx not in reconstruct_fields]
A_e0e0 = _A[id_e0, id_e0]
A_e0e1 = _A[id_e0, id_e1]
A_e1e0 = _A[id_e1, id_e0]
A_e1e1 = _A[id_e1, id_e1]
A_e0f = _A[id_e0, id_f]
A_e1f = _A[id_e1, id_f]
x_e1 = slate.AssembledVector(_x[id_e1])
x_f = slate.AssembledVector(_x[id_f])
b_e0 = slate.AssembledVector(_b[id_e0])
b_e1 = slate.AssembledVector(_b[id_e1])
# Solve for e1
Sf = A_e1f - A_e1e0 * A_e0e0.inv * A_e0f
S_e1 = A_e1e1 - A_e1e0 * A_e0e0.inv * A_e0e1
r_e1 = b_e1 - A_e1e0 * A_e0e0.inv * b_e0 - Sf * x_f
systems.append(LAContext(lhs=S_e1, rhs=r_e1, field_idx=(id_e1,)))
# Solve for e0
r_e0 = b_e0 - A_e0e1 * x_e1 - A_e0f * x_f
systems.append(LAContext(lhs=A_e0e0, rhs=r_e0, field_idx=(id_e0,)))
else:
msg = "Not implemented for systems with %s fields" % nfields
raise NotImplementedError(msg)
return systems
def slate_to_cpp(expr, temps, prec=None):
"""Translates a Slate expression into its equivalent representation in
the Eigen C++ syntax.
:arg expr: a :class:`slate.TensorBase` expression.
:arg temps: a `dict` of temporaries which map a given expression to its
corresponding representation as a `coffee.Symbol` object.
:arg prec: an argument dictating the order of precedence in the linear
algebra operations. This ensures that parentheticals are placed
appropriately and the order in which linear algebra operations
are performed are correct.
Returns:
a `string` which represents the C/C++ code representation of the
`slate.TensorBase` expr.
"""
# If the tensor is terminal, it has already been declared.
# Coefficients defined as AssembledVectors will have been declared
# by now, as well as any other nodes with high reference count or
# matrix factorizations.
if expr in temps:
return temps[expr].gencode()
elif isinstance(expr, slate.Transpose):
tensor, = expr.operands
return "(%s).transpose()" % slate_to_cpp(tensor, temps)
elif isinstance(expr, slate.Inverse):
tensor, = expr.operands
return "(%s).inverse()" % slate_to_cpp(tensor, temps)
elif isinstance(expr, slate.Negative):
tensor, = expr.operands
result = "-%s" % slate_to_cpp(tensor, temps, expr.prec)
return parenthesize(result, expr.prec, prec)
elif isinstance(expr, (slate.Add, slate.Mul)):
op = {slate.Add: '+', slate.Mul: '*'}[type(expr)]
A, B = expr.operands
result = "%s %s %s" % (slate_to_cpp(
A, temps, expr.prec), op, slate_to_cpp(B, temps, expr.prec))
return parenthesize(result, expr.prec, prec)
elif isinstance(expr, slate.Block):
tensor, = expr.operands
indices = expr._indices
try:
ridx, cidx = indices
except ValueError:
ridx, = indices
cidx = 0
rids = as_tuple(ridx)
cids = as_tuple(cidx)
# Check if indices are non-contiguous
if not all(
all(ids[i] + 1 == ids[i + 1] for i in range(len(ids) - 1))
for ids in (rids, cids)):
raise NotImplementedError("Non-contiguous blocks not implemented")
rshape = expr.shape[0]
rstart = sum(tensor.shapes[0][:min(rids)])
if expr.rank == 1:
cshape = 1
cstart = 0
else:
cshape = expr.shape[1]
cstart = sum(tensor.shapes[1][:min(cids)])
result = "(%s).block<%d, %d>(%d, %d)" % (slate_to_cpp(
tensor, temps, expr.prec), rshape, cshape, rstart, cstart)
return parenthesize(result, expr.prec, prec)
elif isinstance(expr, slate.Solve):
A, B = expr.operands
result = "%s.solve(%s)" % (slate_to_cpp(
A, temps, expr.prec), slate_to_cpp(B, temps, expr.prec))
return parenthesize(result, expr.prec, prec)
else:
raise NotImplementedError("Type %s not supported.", type(expr))
def c_map_bcs(self, sign, is_facet):
maps = as_tuple(self.map, Map)
val = []
# To throw away boundary condition values, we subtract a large
# value from the map to make it negative then add it on later to
# get back to the original
max_int = 10000000
need_bottom = False
# Apply any bcs on the first (bottom) layer
for i, map in enumerate(maps):
if not map.iterset._extruded:
continue
for j, m in enumerate(map):
bottom_masks = None
for location, name in m.implicit_bcs:
if location == "bottom":
if bottom_masks is None:
bottom_masks = m.bottom_mask[name].copy()
else:
bottom_masks += m.bottom_mask[name]
need_bottom = True
if bottom_masks is not None:
for idx in range(m.arity):
if bottom_masks[idx] < 0:
val.append("xtr_%(name)s[%(ind)s] %(sign)s= %(val)s;" %
{'name': self.c_map_name(i, j),
'val': max_int,
'ind': idx,
'sign': sign})
if need_bottom:
val.insert(0, "if (j_0 == 0) {")
val.append("}")
need_top = False
pos = len(val)
# Apply any bcs on last (top) layer
for i, map in enumerate(maps):
if not map.iterset._extruded:
continue
for j, m in enumerate(map):
top_masks = None
for location, name in m.implicit_bcs:
if location == "top":
if top_masks is None:
top_masks = m.top_mask[name].copy()
else:
top_masks += m.top_mask[name]
need_top = True
if top_masks is not None:
facet_offset = m.arity if is_facet else 0
for idx in range(m.arity):
if top_masks[idx] < 0:
val.append("xtr_%(name)s[%(ind)s] %(sign)s= %(val)s;" %
{'name': self.c_map_name(i, j),
'val': max_int,
'ind': idx + facet_offset,
'sign': sign})
if need_top:
val.insert(pos, "if (j_0 == top_layer - 1) {")
val.append("}")
return '\n'.join(val)+'\n'
def initialize(self, pc):
"""Set up the problem context. Take the original
mixed problem and reformulate the problem as a
hybridized mixed system.
A KSP is created for the Lagrange multiplier system.
"""
from firedrake import (FunctionSpace, Function, Constant,
TrialFunction, TrialFunctions, TestFunction,
DirichletBC)
from firedrake.assemble import (allocate_matrix,
create_assembly_callable)
from firedrake.formmanipulation import split_form
from ufl.algorithms.replace import replace
# Extract the problem context
prefix = pc.getOptionsPrefix() + "hybridization_"
_, P = pc.getOperators()
self.ctx = P.getPythonContext()
if not isinstance(self.ctx, ImplicitMatrixContext):
raise ValueError("The python context must be an ImplicitMatrixContext")
test, trial = self.ctx.a.arguments()
V = test.function_space()
mesh = V.mesh()
if len(V) != 2:
raise ValueError("Expecting two function spaces.")
if all(Vi.ufl_element().value_shape() for Vi in V):
raise ValueError("Expecting an H(div) x L2 pair of spaces.")
# Automagically determine which spaces are vector and scalar
for i, Vi in enumerate(V):
if Vi.ufl_element().sobolev_space().name == "HDiv":
self.vidx = i
else:
assert Vi.ufl_element().sobolev_space().name == "L2"
self.pidx = i
# Create the space of approximate traces.
W = V[self.vidx]
if W.ufl_element().family() == "Brezzi-Douglas-Marini":
tdegree = W.ufl_element().degree()
else:
try:
# If we have a tensor product element
h_deg, v_deg = W.ufl_element().degree()
tdegree = (h_deg - 1, v_deg - 1)
except TypeError:
tdegree = W.ufl_element().degree() - 1
TraceSpace = FunctionSpace(mesh, "HDiv Trace", tdegree)
# Break the function spaces and define fully discontinuous spaces
broken_elements = ufl.MixedElement([ufl.BrokenElement(Vi.ufl_element()) for Vi in V])
V_d = FunctionSpace(mesh, broken_elements)
# Set up the functions for the original, hybridized
# and schur complement systems
self.broken_solution = Function(V_d)
self.broken_residual = Function(V_d)
self.trace_solution = Function(TraceSpace)
self.unbroken_solution = Function(V)
self.unbroken_residual = Function(V)
shapes = (V[self.vidx].finat_element.space_dimension(),
np.prod(V[self.vidx].shape))
domain = "{[i,j]: 0 <= i < %d and 0 <= j < %d}" % shapes
instructions = """
for i, j
w[i,j] = w[i,j] + 1
end
"""
self.weight = Function(V[self.vidx])
par_loop((domain, instructions), ufl.dx, {"w": (self.weight, INC)},
is_loopy_kernel=True)
instructions = """
for i, j
vec_out[i,j] = vec_out[i,j] + vec_in[i,j]/w[i,j]
end
"""
self.average_kernel = (domain, instructions)
# Create the symbolic Schur-reduction:
# Original mixed operator replaced with "broken"
# arguments
arg_map = {test: TestFunction(V_d),
trial: TrialFunction(V_d)}
Atilde = Tensor(replace(self.ctx.a, arg_map))
gammar = TestFunction(TraceSpace)
n = ufl.FacetNormal(mesh)
sigma = TrialFunctions(V_d)[self.vidx]
if mesh.cell_set._extruded:
Kform = (gammar('+') * ufl.jump(sigma, n=n) * ufl.dS_h
+ gammar('+') * ufl.jump(sigma, n=n) * ufl.dS_v)
else:
Kform = (gammar('+') * ufl.jump(sigma, n=n) * ufl.dS)
# Here we deal with boundaries. If there are Neumann
# conditions (which should be enforced strongly for
# H(div)xL^2) then we need to add jump terms on the exterior
# facets. If there are Dirichlet conditions (which should be
# enforced weakly) then we need to zero out the trace
# variables there as they are not active (otherwise the hybrid
# problem is not well-posed).
# If boundary conditions are contained in the ImplicitMatrixContext:
if self.ctx.row_bcs:
# Find all the subdomains with neumann BCS
# These are Dirichlet BCs on the vidx space
neumann_subdomains = set()
for bc in self.ctx.row_bcs:
if bc.function_space().index == self.pidx:
raise NotImplementedError("Dirichlet conditions for scalar variable not supported. Use a weak bc")
if bc.function_space().index != self.vidx:
raise NotImplementedError("Dirichlet bc set on unsupported space.")
# append the set of sub domains
subdom = bc.sub_domain
if isinstance(subdom, str):
neumann_subdomains |= set([subdom])
else:
neumann_subdomains |= set(as_tuple(subdom, int))
# separate out the top and bottom bcs
extruded_neumann_subdomains = neumann_subdomains & {"top", "bottom"}
neumann_subdomains = neumann_subdomains - extruded_neumann_subdomains
integrand = gammar * ufl.dot(sigma, n)
measures = []
trace_subdomains = []
if mesh.cell_set._extruded:
ds = ufl.ds_v
for subdomain in sorted(extruded_neumann_subdomains):
measures.append({"top": ufl.ds_t, "bottom": ufl.ds_b}[subdomain])
trace_subdomains.extend(sorted({"top", "bottom"} - extruded_neumann_subdomains))
else:
ds = ufl.ds
if "on_boundary" in neumann_subdomains:
measures.append(ds)
else:
measures.extend((ds(sd) for sd in sorted(neumann_subdomains)))
markers = [int(x) for x in mesh.exterior_facets.unique_markers]
dirichlet_subdomains = set(markers) - neumann_subdomains
trace_subdomains.extend(sorted(dirichlet_subdomains))
for measure in measures:
Kform += integrand*measure
trace_bcs = [DirichletBC(TraceSpace, Constant(0.0), subdomain) for subdomain in trace_subdomains]
else:
# No bcs were provided, we assume weak Dirichlet conditions.
# We zero out the contribution of the trace variables on
# the exterior boundary. Extruded cells will have both
# horizontal and vertical facets
trace_subdomains = ["on_boundary"]
if mesh.cell_set._extruded:
trace_subdomains.extend(["bottom", "top"])
trace_bcs = [DirichletBC(TraceSpace, Constant(0.0), subdomain) for subdomain in trace_subdomains]
# Make a SLATE tensor from Kform
K = Tensor(Kform)
# Assemble the Schur complement operator and right-hand side
self.schur_rhs = Function(TraceSpace)
self._assemble_Srhs = create_assembly_callable(
K * Atilde.inv * AssembledVector(self.broken_residual),
tensor=self.schur_rhs,
form_compiler_parameters=self.ctx.fc_params)
mat_type = PETSc.Options().getString(prefix + "mat_type", "aij")
schur_comp = K * Atilde.inv * K.T
self.S = allocate_matrix(schur_comp, bcs=trace_bcs,
form_compiler_parameters=self.ctx.fc_params,
mat_type=mat_type,
options_prefix=prefix)
self._assemble_S = create_assembly_callable(schur_comp,
tensor=self.S,
bcs=trace_bcs,
form_compiler_parameters=self.ctx.fc_params,
mat_type=mat_type)
self._assemble_S()
self.S.force_evaluation()
Smat = self.S.petscmat
nullspace = self.ctx.appctx.get("trace_nullspace", None)
if nullspace is not None:
nsp = nullspace(TraceSpace)
Smat.setNullSpace(nsp.nullspace(comm=pc.comm))
# Set up the KSP for the system of Lagrange multipliers
trace_ksp = PETSc.KSP().create(comm=pc.comm)
trace_ksp.setOptionsPrefix(prefix)
trace_ksp.setOperators(Smat)
trace_ksp.setUp()
trace_ksp.setFromOptions()
self.trace_ksp = trace_ksp
split_mixed_op = dict(split_form(Atilde.form))
split_trace_op = dict(split_form(K.form))
# Generate reconstruction calls
self._reconstruction_calls(split_mixed_op, split_trace_op)
def initialize(self, pc):
"""Set up the problem context. Take the original
mixed problem and reformulate the problem as a
hybridized mixed system.
A KSP is created for the Lagrange multiplier system.
"""
from firedrake import (FunctionSpace, Function, Constant,
TrialFunction, TrialFunctions, TestFunction,
DirichletBC)
from firedrake.assemble import (allocate_matrix,
create_assembly_callable)
from firedrake.formmanipulation import split_form
from ufl.algorithms.replace import replace
# Extract the problem context
prefix = pc.getOptionsPrefix() + "hybridization_"
_, P = pc.getOperators()
self.ctx = P.getPythonContext()
if not isinstance(self.ctx, ImplicitMatrixContext):
raise ValueError(
"The python context must be an ImplicitMatrixContext")
test, trial = self.ctx.a.arguments()
V = test.function_space()
mesh = V.mesh()
if len(V) != 2:
raise ValueError("Expecting two function spaces.")
if all(Vi.ufl_element().value_shape() for Vi in V):
raise ValueError("Expecting an H(div) x L2 pair of spaces.")
# Automagically determine which spaces are vector and scalar
for i, Vi in enumerate(V):
if Vi.ufl_element().sobolev_space().name == "HDiv":
self.vidx = i
else:
assert Vi.ufl_element().sobolev_space().name == "L2"
self.pidx = i
# Create the space of approximate traces.
W = V[self.vidx]
if W.ufl_element().family() == "Brezzi-Douglas-Marini":
tdegree = W.ufl_element().degree()
else:
try:
# If we have a tensor product element
h_deg, v_deg = W.ufl_element().degree()
tdegree = (h_deg - 1, v_deg - 1)
except TypeError:
tdegree = W.ufl_element().degree() - 1
TraceSpace = FunctionSpace(mesh, "HDiv Trace", tdegree)
# Break the function spaces and define fully discontinuous spaces
broken_elements = ufl.MixedElement(
[ufl.BrokenElement(Vi.ufl_element()) for Vi in V])
V_d = FunctionSpace(mesh, broken_elements)
# Set up the functions for the original, hybridized
# and schur complement systems
self.broken_solution = Function(V_d)
self.broken_residual = Function(V_d)
self.trace_solution = Function(TraceSpace)
self.unbroken_solution = Function(V)
self.unbroken_residual = Function(V)
shapes = (V[self.vidx].finat_element.space_dimension(),
np.prod(V[self.vidx].shape))
domain = "{[i,j]: 0 <= i < %d and 0 <= j < %d}" % shapes
instructions = """
for i, j
w[i,j] = w[i,j] + 1
end
"""
self.weight = Function(V[self.vidx])
par_loop((domain, instructions),
ufl.dx, {"w": (self.weight, INC)},
is_loopy_kernel=True)
instructions = """
for i, j
vec_out[i,j] = vec_out[i,j] + vec_in[i,j]/w[i,j]
end
"""
self.average_kernel = (domain, instructions)
# Create the symbolic Schur-reduction:
# Original mixed operator replaced with "broken"
# arguments
arg_map = {test: TestFunction(V_d), trial: TrialFunction(V_d)}
Atilde = Tensor(replace(self.ctx.a, arg_map))
gammar = TestFunction(TraceSpace)
n = ufl.FacetNormal(mesh)
sigma = TrialFunctions(V_d)[self.vidx]
if mesh.cell_set._extruded:
Kform = (gammar('+') * ufl.jump(sigma, n=n) * ufl.dS_h +
gammar('+') * ufl.jump(sigma, n=n) * ufl.dS_v)
else:
Kform = (gammar('+') * ufl.jump(sigma, n=n) * ufl.dS)
# Here we deal with boundaries. If there are Neumann
# conditions (which should be enforced strongly for
# H(div)xL^2) then we need to add jump terms on the exterior
# facets. If there are Dirichlet conditions (which should be
# enforced weakly) then we need to zero out the trace
# variables there as they are not active (otherwise the hybrid
# problem is not well-posed).
# If boundary conditions are contained in the ImplicitMatrixContext:
if self.ctx.row_bcs:
# Find all the subdomains with neumann BCS
# These are Dirichlet BCs on the vidx space
neumann_subdomains = set()
for bc in self.ctx.row_bcs:
if bc.function_space().index == self.pidx:
raise NotImplementedError(
"Dirichlet conditions for scalar variable not supported. Use a weak bc"
)
if bc.function_space().index != self.vidx:
raise NotImplementedError(
"Dirichlet bc set on unsupported space.")
# append the set of sub domains
subdom = bc.sub_domain
if isinstance(subdom, str):
neumann_subdomains |= set([subdom])
else:
neumann_subdomains |= set(
as_tuple(subdom, numbers.Integral))
# separate out the top and bottom bcs
extruded_neumann_subdomains = neumann_subdomains & {
"top", "bottom"
}
neumann_subdomains = neumann_subdomains - extruded_neumann_subdomains
integrand = gammar * ufl.dot(sigma, n)
measures = []
trace_subdomains = []
if mesh.cell_set._extruded:
ds = ufl.ds_v
for subdomain in sorted(extruded_neumann_subdomains):
measures.append({
"top": ufl.ds_t,
"bottom": ufl.ds_b
}[subdomain])
trace_subdomains.extend(
sorted({"top", "bottom"} - extruded_neumann_subdomains))
else:
ds = ufl.ds
if "on_boundary" in neumann_subdomains:
measures.append(ds)
else:
measures.extend((ds(sd) for sd in sorted(neumann_subdomains)))
markers = [int(x) for x in mesh.exterior_facets.unique_markers]
dirichlet_subdomains = set(markers) - neumann_subdomains
trace_subdomains.extend(sorted(dirichlet_subdomains))
for measure in measures:
Kform += integrand * measure
trace_bcs = [
DirichletBC(TraceSpace, Constant(0.0), subdomain)
for subdomain in trace_subdomains
]
else:
# No bcs were provided, we assume weak Dirichlet conditions.
# We zero out the contribution of the trace variables on
# the exterior boundary. Extruded cells will have both
# horizontal and vertical facets
trace_subdomains = ["on_boundary"]
if mesh.cell_set._extruded:
trace_subdomains.extend(["bottom", "top"])
trace_bcs = [
DirichletBC(TraceSpace, Constant(0.0), subdomain)
for subdomain in trace_subdomains
]
# Make a SLATE tensor from Kform
K = Tensor(Kform)
# Assemble the Schur complement operator and right-hand side
self.schur_rhs = Function(TraceSpace)
self._assemble_Srhs = create_assembly_callable(
K * Atilde.inv * AssembledVector(self.broken_residual),
tensor=self.schur_rhs,
form_compiler_parameters=self.ctx.fc_params)
mat_type = PETSc.Options().getString(prefix + "mat_type", "aij")
schur_comp = K * Atilde.inv * K.T
self.S = allocate_matrix(schur_comp,
bcs=trace_bcs,
form_compiler_parameters=self.ctx.fc_params,
mat_type=mat_type,
options_prefix=prefix)
self._assemble_S = create_assembly_callable(
schur_comp,
tensor=self.S,
bcs=trace_bcs,
form_compiler_parameters=self.ctx.fc_params,
mat_type=mat_type)
with timed_region("HybridOperatorAssembly"):
self._assemble_S()
Smat = self.S.petscmat
nullspace = self.ctx.appctx.get("trace_nullspace", None)
if nullspace is not None:
nsp = nullspace(TraceSpace)
Smat.setNullSpace(nsp.nullspace(comm=pc.comm))
# Set up the KSP for the system of Lagrange multipliers
trace_ksp = PETSc.KSP().create(comm=pc.comm)
trace_ksp.setOptionsPrefix(prefix)
trace_ksp.setOperators(Smat)
trace_ksp.setUp()
trace_ksp.setFromOptions()
self.trace_ksp = trace_ksp
split_mixed_op = dict(split_form(Atilde.form))
split_trace_op = dict(split_form(K.form))
# Generate reconstruction calls
self._reconstruction_calls(split_mixed_op, split_trace_op)
def _tile(self):
"""Tile consecutive loops over different iteration sets characterized
by RAW and WAR dependencies. This requires interfacing with the SLOPE
library."""
loop_chain = self._loop_chain
if len(loop_chain) == 1:
# Nothing more to try fusing after soft and hard fusion
return
tile_size = self._options.get('tile_size', 1)
seed_loop = self._options.get('seed_loop', 0)
extra_halo = self._options.get('extra_halo', False)
coloring = self._options.get('coloring', 'default')
use_prefetch = self._options.get('use_prefetch', 0)
ignore_war = self._options.get('ignore_war', False)
log = self._options.get('log', False)
rank = MPI.COMM_WORLD.rank
# The SLOPE inspector, which needs be populated with sets, maps,
# descriptors, and loop chain structure
inspector = slope.Inspector(self._name)
# Build inspector and argument types and values
# Note: we need ordered containers to be sure that SLOPE generates
# identical code for all ranks
arguments = []
insp_sets, insp_maps, insp_loops = OrderedDict(), OrderedDict(), []
for loop in loop_chain:
slope_desc = set()
# 1) Add sets
iterset = loop.it_space.iterset
iterset = iterset.subset if hasattr(iterset, 'subset') else iterset
slope_set = create_slope_set(iterset, extra_halo, insp_sets)
# If iterating over a subset, we fake an indirect parloop from the
# (iteration) subset to the superset. This allows the propagation of
# tiling across the hierarchy of sets (see SLOPE for further info)
if slope_set.superset:
create_slope_set(iterset.superset, extra_halo, insp_sets)
map_name = "%s_tosuperset" % slope_set.name
insp_maps[slope_set.name] = (map_name, slope_set.name,
iterset.superset.name, iterset.indices)
slope_desc.add((map_name, INC._mode))
for a in loop.args:
# 2) Add access descriptors
maps = as_tuple(a.map, Map)
if not maps:
# Simplest case: direct loop
slope_desc.add(('DIRECT', a.access._mode))
else:
# Add maps (there can be more than one per argument if the arg
# is actually a Mat - in which case there are two maps - or if
# a MixedMap) and relative descriptors
for i, map in enumerate(maps):
for j, m in enumerate(map):
map_name = "%s%d_%d" % (m.name, i, j)
insp_maps[m.name] = (map_name, m.iterset.name,
m.toset.name, m.values_with_halo)
slope_desc.add((map_name, a.access._mode))
create_slope_set(m.iterset, extra_halo, insp_sets)
create_slope_set(m.toset, extra_halo, insp_sets)
# 3) Add loop
insp_loops.append((loop.kernel.name, slope_set.name, list(slope_desc)))
# Provide structure of loop chain to SLOPE
arguments.extend([inspector.add_sets(insp_sets.keys())])
arguments.extend([inspector.add_maps(insp_maps.values())])
inspector.add_loops(insp_loops)
# Set a specific tile size
arguments.extend([inspector.set_tile_size(tile_size)])
# Tell SLOPE the rank of the MPI process
arguments.extend([inspector.set_mpi_rank(rank)])
# Get type and value of additional arguments that SLOPE can exploit
arguments.extend(inspector.add_extra_info())
# Add any available partitioning
partitionings = [(s[0], v) for s, v in insp_sets.items() if v is not None]
arguments.extend([inspector.add_partitionings(partitionings)])
# Arguments types and values
argtypes, argvalues = zip(*arguments)
# Set key tiling properties
inspector.drive_inspection(ignore_war=ignore_war,
seed_loop=seed_loop,
prefetch=use_prefetch,
coloring=coloring,
part_mode='chunk')
# Generate the C code
src = inspector.generate_code()
# Return type of the inspector
rettype = slope.Executor.meta['py_ctype_exec']
# Compiler and linker options
compiler = coffee.system.compiler.get('name')
cppargs = slope.get_compile_opts(compiler)
cppargs += ['-I%s/include/SLOPE' % sys.prefix]
ldargs = ['-L%s/lib' % sys.prefix, '-l%s' % slope.get_lib_name(),
'-Wl,-rpath,%s/lib' % sys.prefix, '-lrt']
# Compile and run inspector
fun = compilation.load(src, "cpp", "inspector", cppargs, ldargs,
argtypes, rettype, compiler)
inspection = fun(*argvalues)
# Log the inspector output
if log and rank == 0:
estimate_data_reuse(self._name, loop_chain)
# Finally, get the Executor representation, to be used at executor
# code generation time
executor = slope.Executor(inspector)
kernel = Kernel(tuple(loop.kernel for loop in loop_chain))
self._schedule = TilingSchedule(self._name, self._schedule, kernel, inspection,
executor, **self._options)
def c_addto(self, i, j, buf_name, tmp_name, tmp_decl,
extruded=None, is_facet=False):
maps = as_tuple(self.map, Map)
nrows = maps[0].split[i].arity
ncols = maps[1].split[j].arity
rows_str = "%s + i * %s" % (self.c_map_name(0, i), nrows)
cols_str = "%s + i * %s" % (self.c_map_name(1, j), ncols)
if extruded is not None:
rows_str = extruded + self.c_map_name(0, i)
cols_str = extruded + self.c_map_name(1, j)
if is_facet:
nrows *= 2
ncols *= 2
ret = []
rbs, cbs = self.data.sparsity[i, j].dims[0][0]
rdim = rbs * nrows
addto_name = buf_name
addto = 'MatSetValuesLocal'
if self.data._is_vector_field:
addto = 'MatSetValuesBlockedLocal'
rmap, cmap = maps
rdim, cdim = self.data.dims[i][j]
if rmap.vector_index is not None or cmap.vector_index is not None:
rows_str = "rowmap"
cols_str = "colmap"
addto = "MatSetValuesLocal"
fdict = {'nrows': nrows,
'ncols': ncols,
'rdim': rdim,
'cdim': cdim,
'rowmap': self.c_map_name(0, i),
'colmap': self.c_map_name(1, j),
'drop_full_row': 0 if rmap.vector_index is not None else 1,
'drop_full_col': 0 if cmap.vector_index is not None else 1}
# Horrible hack alert
# To apply BCs to a component of a Dat with cdim > 1
# we encode which components to apply things to in the
# high bits of the map value
# The value that comes in is:
# -(row + 1 + sum_i 2 ** (30 - i))
# where i are the components to zero
#
# So, the actual row (if it's negative) is:
# (~input) & ~0x70000000
# And we can determine which components to zero by
# inspecting the high bits (1 << 30 - i)
ret.append("""
PetscInt rowmap[%(nrows)d*%(rdim)d];
PetscInt colmap[%(ncols)d*%(cdim)d];
int discard, tmp, block_row, block_col;
for ( int j = 0; j < %(nrows)d; j++ ) {
block_row = %(rowmap)s[i*%(nrows)d + j];
discard = 0;
tmp = -(block_row + 1);
if ( block_row < 0 ) {
discard = 1;
block_row = tmp & ~0x70000000;
}
for ( int k = 0; k < %(rdim)d; k++ ) {
if ( discard && (!(tmp & 0x70000000) || %(drop_full_row)d || ((tmp & (1 << (30 - k))) != 0)) ) {
rowmap[j*%(rdim)d + k] = -1;
} else {
rowmap[j*%(rdim)d + k] = (block_row)*%(rdim)d + k;
}
}
}
for ( int j = 0; j < %(ncols)d; j++ ) {
discard = 0;
block_col = %(colmap)s[i*%(ncols)d + j];
tmp = -(block_col + 1);
if ( block_col < 0 ) {
discard = 1;
block_col = tmp & ~0x70000000;
}
for ( int k = 0; k < %(cdim)d; k++ ) {
if ( discard && (!(tmp & 0x70000000) || %(drop_full_col)d || ((tmp & (1 << (30 - k))) != 0)) ) {
colmap[j*%(cdim)d + k] = -1;
} else {
colmap[j*%(cdim)d + k] = (block_col)*%(cdim)d + k;
}
}
}
""" % fdict)
nrows *= rdim
ncols *= cdim
ret.append("""%(addto)s(%(mat)s, %(nrows)s, %(rows)s,
%(ncols)s, %(cols)s,
(const PetscScalar *)%(vals)s,
%(insert)s);""" %
{'mat': self.c_arg_name(i, j),
'vals': addto_name,
'addto': addto,
'nrows': nrows,
'ncols': ncols,
'rows': rows_str,
'cols': cols_str,
'insert': "INSERT_VALUES" if self.access == WRITE else "ADD_VALUES"})
ret = " "*16 + "{\n" + "\n".join(ret) + "\n" + " "*16 + "}"
return ret
def domain_args(self):
"""The sub_domain the BC applies to."""
if isinstance(self.sub_domain, str):
return (self.sub_domain, )
return (as_tuple(self.sub_domain), )
