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 __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 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 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 _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), )