Source code for firedrake.functionspacedata

"""This module provides an object that encapsulates data that can be
shared between different :class:`~.FunctionSpace` objects.

The sharing is based on the idea of compatibility of function space
node layout.  The shared data is stored on the :func:`~.Mesh` the
function space is created on, since the created objects are
mesh-specific.  The sharing is done on an individual key basis.  So,
for example, Sets can be shared between all function spaces with the
same number of nodes per topological entity.  However, maps are
specific to the node *ordering*.

This means, for example, that function spaces with the same *node*
ordering, but different numbers of dofs per node (e.g. FiniteElement
vs VectorElement) can share the PyOP2 Set and Map data.
"""

import numpy
import finat.ufl
import finat
from decorator import decorator
from functools import partial

from tsfc.finatinterface import create_element as _create_element

from pyop2 import op2
from firedrake.utils import IntType
from pyop2.utils import as_tuple

from firedrake.cython import extrusion_numbering as extnum
from firedrake.cython import dmcommon
from firedrake import halo as halo_mod
from firedrake import mesh as mesh_mod
from firedrake import extrusion_utils as eutils
from firedrake.petsc import PETSc


__all__ = ("get_shared_data", )


@PETSc.Log.EventDecorator("FunctionSpaceData: CreateElement")
def create_element(ufl_element):
    finat_element = _create_element(ufl_element)
    if isinstance(finat_element, finat.TensorFiniteElement):
        # Retrieve scalar element
        finat_element = finat_element.base_element
    return finat_element


@decorator
def cached(f, mesh, key, *args, **kwargs):
    """Sui generis caching for a function whose data is
    associated with a mesh.

    :arg f: The function to cache.
    :arg mesh: The mesh to cache on (should have a
        ``_shared_data_cache`` object).
    :arg key: The key to the cache.
    :args args: Additional arguments to ``f``.
    :kwargs kwargs:  Additional keyword arguments to ``f``."""
    assert hasattr(mesh, "_shared_data_cache")
    cache = mesh._shared_data_cache[f.__name__]
    try:
        return cache[key]
    except KeyError:
        result = f(mesh, key, *args, **kwargs)
        cache[key] = result
        return result


@cached
def get_global_numbering(mesh, key, global_numbering=None):
    """Get a PETSc Section describing the global numbering.

    This numbering associates function space nodes with topological
    entities.

    :arg mesh: The mesh to use.
    :arg key: a (nodes_per_entity, real_tensorproduct, boundary_set) tuple where
        nodes_per_entity is a tuple of the number of nodes per topological
        entity; real_tensorproduct is True if the function space is a
        degenerate fs x Real tensorproduct; boundary_set is a set of boundary
        markers, indicating sub-domains a boundary condition is specified on.
    :returns: A new PETSc Section.
    """
    if global_numbering:
        return global_numbering
    nodes_per_entity, real_tensorproduct, boundary_set = key
    return mesh.create_section(nodes_per_entity, real_tensorproduct, boundary_set=boundary_set)


@cached
def get_node_set(mesh, key):
    """Get the :class:`node set <pyop2.Set>`.

    :arg mesh: The mesh to use.
    :arg key: a (nodes_per_entity, real_tensorproduct, boundary_set) tuple
        where nodes_per_entity is a tuple of the number of nodes per
        topological entity; real_tensorproduct is True if the function space is
        a degenerate fs x Real tensorproduct; boundary_set is a set of boundary
        markers, indicating sub-domains a boundary condition is specified on.
    :returns: A :class:`pyop2.Set` for the function space nodes.
    """
    nodes_per_entity, real_tensorproduct, _ = key
    global_numbering, constrained_size = get_global_numbering(mesh, key)
    node_classes = mesh.node_classes(nodes_per_entity, real_tensorproduct=real_tensorproduct)
    halo = halo_mod.Halo(mesh.topology_dm, global_numbering, comm=mesh.comm)
    node_set = op2.Set(node_classes, halo=halo, comm=mesh.comm, constrained_size=constrained_size)
    extruded = mesh.cell_set._extruded

    assert global_numbering.getStorageSize() == node_set.total_size
    if not extruded and node_set.total_size >= (1 << (IntType.itemsize * 8 - 4)):
        raise RuntimeError("Problems with more than %d nodes per process unsupported", (1 << (IntType.itemsize * 8 - 4)))
    return node_set


def get_cell_node_list(mesh, entity_dofs, entity_permutations, global_numbering, offsets):
    """Get the cell->node list for specified dof layout.

    :arg mesh: The mesh to use.
    :arg entity_dofs: The FInAT entity_dofs dict.
    :arg entity_permutations: The FInAT entity_permutations dict.
    :arg global_numbering: The PETSc Section describing node layout
        (see :func:`get_global_numbering`).
    :arg offsets: layer offsets for each entity (maybe ignored).
    :returns: A numpy array mapping mesh cells to function space
        nodes.
    """
    return mesh.make_cell_node_list(global_numbering, entity_dofs, entity_permutations, offsets)


def get_facet_node_list(mesh, kind, cell_node_list, offsets):
    """Get the facet->node list for specified dof layout.

    :arg mesh: The mesh to use.
    :arg kind: The facet kind (one of ``"interior_facets"`` or
        ``"exterior_facets"``).
    :arg cell_node_list: The map from mesh cells to function space
        nodes, see :func:`get_cell_node_list`.
    :arg offsets: layer offsets for each entity (maybe ignored).
    :returns: A numpy array mapping mesh facets to function space
        nodes.
    """
    assert kind in ["interior_facets", "exterior_facets"]
    if mesh.topology_dm.getStratumSize(kind, 1) > 0:
        return dmcommon.get_facet_nodes(mesh, cell_node_list, kind, offsets)
    else:
        return numpy.array([], dtype=IntType)


@cached
def get_entity_node_lists(mesh, key, entity_dofs, entity_permutations, global_numbering, offsets):
    """Get the map from mesh entity sets to function space nodes.

    :arg mesh: The mesh to use.
    :arg key: a (entity_dofs_key, real_tensorproduct, entity_permutations_key,
        boundary_set) tuple.
    :arg entity_dofs: FInAT entity dofs.
    :arg entity_permutations: FInAT entity permutations.
    :arg global_numbering: The PETSc Section describing node layout
        (see :func:`get_global_numbering`).
    :arg offsets: layer offsets for each entity (maybe ignored).
    :returns: A dict mapping mesh entity sets to numpy arrays of
        function space nodes.
    """
    # set->node lists are specific to the sorted entity_dofs.
    cell_node_list = get_cell_node_list(mesh, entity_dofs, entity_permutations, global_numbering, offsets)
    interior_facet_node_list = partial(get_facet_node_list, mesh, "interior_facets", cell_node_list, offsets)
    exterior_facet_node_list = partial(get_facet_node_list, mesh, "exterior_facets", cell_node_list, offsets)

    class magic(dict):
        def __missing__(self, key):
            if type(mesh.topology) is mesh_mod.VertexOnlyMeshTopology:
                return self.setdefault(key,
                                       {mesh.cell_set: lambda: cell_node_list}[key]())
            else:
                return self.setdefault(key,
                                       {mesh.cell_set: lambda: cell_node_list,
                                        mesh.interior_facets.set: interior_facet_node_list,
                                        mesh.exterior_facets.set: exterior_facet_node_list}[key]())

    return magic()


@cached
def get_map_cache(mesh, key):
    """Get the map cache for this mesh.

    :arg mesh: The mesh to use.
    :arg key: a (entity_dofs_key, real_tensorproduct, entity_permutations_key,
        boundary_set) tuple where entity_dofs is Canonicalised entity_dofs
        (see :func:`entity_dofs_key`); real_tensorproduct is True if the
        function space is a degenerate fs x Real tensorproduct; boundary_set is
        the set of subdomains a restricted function space is applied to, or
        None if using a regular function space.
    """
    if type(mesh.topology) is mesh_mod.VertexOnlyMeshTopology:
        return {mesh.cell_set: None}
    else:
        return {mesh.cell_set: None,
                mesh.interior_facets.set: None,
                mesh.exterior_facets.set: None,
                "boundary_node": None}


@cached
def get_boundary_masks(mesh, key, finat_element):
    """Get masks for facet dofs.

    :arg mesh: The mesh to use.
    :arg key: Canonicalised entity_dofs (see :func:`entity_dofs_key`).
    :arg finat_element: The FInAT element.
    :returns: ``None`` or a 3-tuple of a Section, an array of indices, and
        an array indicating which points in the Section correspond to
        the facets of the cell.  If section.getDof(p) is non-zero,
        then there are ndof basis functions topologically associated
        with points in the closure of point p.  The basis function
        indices are in the index array, starting at section.getOffset(p).
    """
    if not mesh.cell_set._extruded:
        return None
    _, kind = key
    assert kind in {"cell", "interior_facet"}
    dim = finat_element.cell.get_spatial_dimension()
    ecd = finat_element.entity_closure_dofs()
    # Number of entities on cell excepting the cell itself.
    chart = sum(map(len, ecd.values())) - 1
    closure_section = PETSc.Section().create(comm=PETSc.COMM_SELF)
    # Double up for interior facets.
    if kind == "cell":
        ncell = 1
    else:
        ncell = 2
    closure_section.setChart(0, ncell*chart)
    closure_indices = []
    facet_points = []
    p = 0

    offset = finat_element.space_dimension()
    for cell in range(ncell):
        for ent in sorted(ecd.keys()):
            # Never need closure of cell
            if sum(ent) == dim:
                continue
            for key in sorted(ecd[ent].keys()):
                closure_section.setDof(p, len(ecd[ent][key]))
                vals = numpy.asarray(sorted(ecd[ent][key]), dtype=IntType)
                closure_indices.extend(vals + cell*offset)
                if sum(ent) == dim - 1:
                    facet_points.append(p)
                p += 1
    closure_section.setUp()
    closure_indices = numpy.asarray(closure_indices, dtype=IntType)
    facet_points = numpy.asarray(facet_points, dtype=IntType)
    return (closure_section, closure_indices, facet_points)


@cached
def get_work_function_cache(mesh, ufl_element):
    """Get the cache for work functions.

    :arg mesh: The mesh to use.
    :arg ufl_element: The ufl element, used as a key.
    :returns: A dict.

    :class:`.FunctionSpace` objects sharing the same UFL element (and
    therefore comparing equal) share a work function cache.
    """
    return {}


@cached
def get_top_bottom_boundary_nodes(mesh, key, V):
    """Get top or bottom boundary nodes of an extruded function space.

    :arg mesh: The mesh to cache on.
    :arg key: A 3-tuple of ``(entity_dofs_key, sub_domain, boundary_set)`` key.
        Where sub_domain indicates top or bottom.
    :arg V: The FunctionSpace to select from.
    :arg entity_dofs: The flattened entity dofs.
    :returnsL: A numpy array of the (unique) boundary nodes.
    """
    _, sub_domain, boundary_set = key
    cell_node_list = V.cell_node_list
    offset = V.offset
    if mesh.variable_layers:
        return extnum.top_bottom_boundary_nodes(mesh, cell_node_list,
                                                V.cell_boundary_masks,
                                                offset,
                                                sub_domain)
    else:
        if mesh.extruded_periodic and sub_domain == "top":
            raise ValueError("Invalid subdomain 'top': 'top' boundary is identified as 'bottom' boundary in periodic extrusion")
        idx = {"bottom": -2, "top": -1}[sub_domain]
        section, indices, facet_points = V.cell_boundary_masks
        facet = facet_points[idx]
        dof = section.getDof(facet)
        off = section.getOffset(facet)
        mask = indices[off:off+dof]
        nodes = cell_node_list[..., mask]
        if sub_domain == "top":
            nodes = nodes + offset[mask]*(mesh.cell_set.layers - 2)
        return numpy.unique(nodes)


@cached
def get_facet_closure_nodes(mesh, key, V):
    """Function space nodes in the closure of facets with a given
    marker.
    :arg mesh: Mesh to cache on
    :arg key: (edofs, sub_domain, boundary_set) tuple
    :arg V: function space.
    :returns: numpy array of unique nodes in the closure of facets
       with provided markers (both interior and exterior)."""
    _, sub_domain, boundary_set = key
    if sub_domain not in {"on_boundary", "top", "bottom"}:
        valid = set(mesh.interior_facets.unique_markers)
        valid |= set(mesh.exterior_facets.unique_markers)
        invalid = set(sub_domain) - valid
        if invalid:
            raise LookupError(f"BC construction got invalid markers {invalid}. "
                              f"Valid markers are '{valid}'")
    return dmcommon.facet_closure_nodes(V, sub_domain)


def get_max_work_functions(V):
    """Get the maximum number of work functions.

    :arg V: The function space to get the number of work functions for.
    :returns: The maximum number of work functions.

    This number is shared between all function spaces with the same
    :meth:`~.FunctionSpace.ufl_element` and
    :meth:`~FunctionSpace.mesh`.

    The default is 25 work functions per function space.  This can be
    set using :func:`set_max_work_functions`.
    """
    mesh = V.mesh()
    assert hasattr(mesh, "_shared_data_cache")
    cache = mesh._shared_data_cache["max_work_functions"]
    return cache.get(V.ufl_element(), 25)


def set_max_work_functions(V, val):
    """Set the maximum number of work functions.

    :arg V: The function space to set the number of work functions
        for.
    :arg val: The new maximum number of work functions.

    This number is shared between all function spaces with the same
    :meth:`~.FunctionSpace.ufl_element` and
    :meth:`~FunctionSpace.mesh`.
    """
    mesh = V.mesh()
    assert hasattr(mesh, "_shared_data_cache")
    cache = mesh._shared_data_cache["max_work_functions"]
    cache[V.ufl_element()] = val


def entity_dofs_key(entity_dofs):
    """Provide a canonical key for an entity_dofs dict.

    :arg entity_dofs: The FInAT entity_dofs.
    :returns: A tuple of canonicalised entity_dofs (suitable for
        caching).
    """
    key = []
    for k in sorted(entity_dofs.keys()):
        sub_key = [k]
        for sk in sorted(entity_dofs[k]):
            sub_key.append(tuple(entity_dofs[k][sk]))
        key.append(tuple(sub_key))
    key = tuple(key)
    return key


def entity_permutations_key(entity_permutations):
    """Provide a canonical key for an entity_permutations dict.

    :arg entity_permutations: The FInAT entity_permutations.
    :returns: A tuple of canonicalised entity_permutations (suitable for
        caching).
    """
    key = []
    for k in sorted(entity_permutations.keys()):
        sub_key = [k]
        for sk in sorted(entity_permutations[k]):
            subsub_key = [sk]
            for ssk in sorted(entity_permutations[k][sk]):
                subsub_key.append((ssk, tuple(entity_permutations[k][sk][ssk])))
            sub_key.append(tuple(subsub_key))
        key.append(tuple(sub_key))
    key = tuple(key)
    return key


class FunctionSpaceData(object):
    """Function spaces with the same entity dofs share data.  This class
    stores that shared data.  It is cached on the mesh.

    :arg mesh: The mesh to share the data on.
    :arg ufl_element: The UFL element.
    :arg boundary_set: The set of subdomains that a Dirichlet boundary condition
        will act on. This is None if the function space is not a
        :class:`.RestrictedFunctionSpace`.
    """
    __slots__ = ("real_tensorproduct", "map_cache", "entity_node_lists",
                 "node_set", "cell_boundary_masks",
                 "interior_facet_boundary_masks", "offset", "offset_quotient",
                 "extruded", "mesh", "global_numbering", "boundary_set")

    @PETSc.Log.EventDecorator()
    def __init__(self, mesh, ufl_element, boundary_set=None):
        if type(ufl_element) is finat.ufl.MixedElement:
            raise ValueError("Can't create FunctionSpace for MixedElement")

        self.boundary_set = boundary_set

        finat_element = create_element(ufl_element)
        real_tensorproduct = eutils.is_real_tensor_product_element(finat_element)
        entity_dofs = finat_element.entity_dofs()
        nodes_per_entity = tuple(mesh.make_dofs_per_plex_entity(entity_dofs))
        try:
            entity_permutations = finat_element.entity_permutations
        except NotImplementedError:
            entity_permutations = None

        # Create the PetscSection mapping topological entities to functionspace nodes
        # For non-scalar valued function spaces, there are multiple dofs per node.
        key = (nodes_per_entity, real_tensorproduct, boundary_set)
        # These are keyed only on nodes per topological entity.
        global_numbering, constrained_size = get_global_numbering(mesh, key)
        node_set = get_node_set(mesh, key)

        edofs_key = entity_dofs_key(entity_dofs)
        # entity_permutations is None if not yet implemented
        eperm_key = entity_permutations_key(entity_permutations) if entity_permutations else None

        self.real_tensorproduct = real_tensorproduct
        # Empty map caches. This is a sui generis cache
        # implementation because of the need to support boundary
        # conditions.
        # Map caches are specific to a cell_node_list, which is keyed by entity_dof
        self.map_cache = get_map_cache(mesh, (edofs_key, real_tensorproduct, eperm_key, boundary_set))

        if isinstance(mesh, mesh_mod.ExtrudedMeshTopology):
            self.offset = eutils.calculate_dof_offset(finat_element)
        else:
            self.offset = None
        if isinstance(mesh, mesh_mod.ExtrudedMeshTopology) and mesh.extruded_periodic:
            self.offset_quotient = eutils.calculate_dof_offset_quotient(finat_element)
        else:
            self.offset_quotient = None

        self.entity_node_lists = get_entity_node_lists(mesh, (edofs_key, real_tensorproduct, eperm_key, boundary_set), entity_dofs, entity_permutations, global_numbering, self.offset)
        self.node_set = node_set
        self.cell_boundary_masks = get_boundary_masks(mesh, (edofs_key, "cell"), finat_element)
        self.interior_facet_boundary_masks = get_boundary_masks(mesh, (edofs_key, "interior_facet"), finat_element)
        self.extruded = mesh.cell_set._extruded
        self.mesh = mesh
        self.global_numbering = global_numbering

    def __eq__(self, other):
        if type(self) is not type(other):
            return False
        return all(getattr(self, s) is getattr(other, s) for s in
                   FunctionSpaceData.__slots__)

    def __ne__(self, other):
        return not self.__eq__(other)

    def __repr__(self):
        return "FunctionSpaceData(%r, %r)" % (self.mesh, self.node_set)

    def __str__(self):
        return "FunctionSpaceData(%s, %s)" % (self.mesh, self.node_set)

    @PETSc.Log.EventDecorator()
    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, V.boundary_set)
            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, V.boundary_set)
            return get_facet_closure_nodes(V.mesh(), key, V)

    @PETSc.Log.EventDecorator()
    def get_map(self, V, entity_set, map_arity, name, offset, offset_quotient):
        """Return a :class:`pyop2.Map` from some topological entity to
        degrees of freedom.

        :arg V: The :class:`FunctionSpace` to create the map for.
        :arg entity_set: The :class:`pyop2.Set` of entities to map from.
        :arg map_arity: The arity of the resulting map.
        :arg name: A name for the resulting map.
        :arg offset: Map offset (for extruded).
        :arg offset_quotient: Map offset_quotient (for extruded)."""
        # V is only really used for error checking and "name".
        assert len(V) == 1, "get_map should not be called on MixedFunctionSpace"
        entity_node_list = self.entity_node_lists[entity_set]
        val = self.map_cache[entity_set]
        if val is None:
            val = op2.Map(entity_set, self.node_set,
                          map_arity,
                          entity_node_list,
                          ("%s_"+name) % (V.name),
                          offset=offset,
                          offset_quotient=offset_quotient)

            self.map_cache[entity_set] = val
        return val


[docs] @PETSc.Log.EventDecorator() def get_shared_data(mesh, ufl_element, boundary_set=None): """Return the ``FunctionSpaceData`` for the given element. :arg mesh: The mesh to build the function space data on. :arg ufl_element: A UFL element. :arg boundary_set: A set of boundary markers, indicating the subdomains a boundary condition is specified on. :raises ValueError: if mesh or ufl_element are invalid. :returns: a ``FunctionSpaceData`` object with the shared data. """ if not isinstance(mesh, mesh_mod.AbstractMeshTopology): raise ValueError("%s is not an AbstractMeshTopology" % mesh) if not isinstance(ufl_element, finat.ufl.finiteelement.FiniteElementBase): raise ValueError("Can't create function space data from a %s" % type(ufl_element)) return FunctionSpaceData(mesh, ufl_element, boundary_set)