Allow referencing named boundaries or subdomains in get_dofs

docs/examples/performance.py

@@ -21,8 +21,8 @@ def pre(N=3):
 def assembler(m):
     basis = Basis(m, ElementTetP1())
     return (
-        asm(laplace, basis),
-        asm(unit_load, basis),
+        laplace.assemble(basis),
+        unit_load.assemble(basis),
     )
 
 

skfem/assembly/basis/abstract_basis.py

@@ -117,71 +117,121 @@ def complement_dofs(self, *D):
     def find_dofs(self,
                   facets: Dict[str, ndarray] = None,
                   skip: List[str] = None) -> Dict[str, DofsView]:
-        """Return global DOF numbers corresponding to a dictionary of facets.
+        """Deprecated in favor of :meth:`~skfem.AbstractBasis.get_dofs`."""
+        if facets is None:
+            if self.mesh.boundaries is None:
+                facets = {'all': self.mesh.boundary_facets()}
+            else:
+                facets = self.mesh.boundaries
+
+        return {k: self.dofs.get_facet_dofs(facets[k], skip_dofnames=skip)
+                for k in facets}
+
+    def get_dofs(self,
+                 facets: Optional[Any] = None,
+                 elements: Optional[Any] = None,
+                 skip: List[str] = None) -> Any:
+        """Find global DOF numbers.
+
+        Accepts an array of facet/element indices.  However, various argument
+        types can be turned into an array of facet/element indices.
 
-        Facets can be queried from :class:`~skfem.mesh.Mesh` objects:
+        Get all boundary DOFs:
 
-        >>> from skfem import MeshTri
+        >>> import numpy as np
+        >>> from skfem import MeshTri, Basis, ElementTriP1
         >>> m = MeshTri().refined()
-        >>> m.facets_satisfying(lambda x: x[0] == 0)
-        array([1, 5])
+        >>> basis = Basis(m, ElementTriP1())
+        >>> basis.get_dofs().flatten()
+        array([0, 1, 2, 3, 4, 5, 7, 8])
 
-        This corresponds to a list of facet indices that can be passed over:
+        Get DOFs via a function query:
 
         >>> import numpy as np
-        >>> from skfem import Basis, ElementTriP1
+        >>> from skfem import MeshTri, Basis, ElementTriP1
+        >>> m = MeshTri().refined()
         >>> basis = Basis(m, ElementTriP1())
-        >>> basis.find_dofs({'left': np.array([1, 5])})['left'].all()
+        >>> basis.get_dofs(lambda x: np.isclose(x[0], 0)).flatten()
         array([0, 2, 5])
 
-        Parameters
-        ----------
-        facets
-            A dictionary of facets. If ``None``, use ``self.mesh.boundaries``
-            if set or otherwise use ``{'all': self.mesh.boundary_facets()}``.
-        skip
-            List of dofnames to skip.
+        Get DOFs via named boundaries:
 
-        """
-        if facets is None:
-            if self.mesh.boundaries is None:
-                facets = {'all': self.mesh.boundary_facets()}
-            else:
-                facets = self.mesh.boundaries
+        >>> import numpy as np
+        >>> from skfem import MeshTri, Basis, ElementTriP1
+        >>> m = (MeshTri()
+        ...      .refined()
+        ...      .with_boundaries({'left': lambda x: np.isclose(x[0], 0)}))
+        >>> basis = Basis(m, ElementTriP1())
+        >>> basis.get_dofs('left').flatten()
+        array([0, 2, 5])
 
-        return {k: self.dofs.get_facet_dofs(facets[k], skip_dofnames=skip)
-                for k in facets}
+        Get DOFs via named subdomains:
 
-    def get_dofs(self, facets: Optional[Any] = None) -> Any:
-        """Find global DOF numbers.
+        >>> from skfem import MeshTri, Basis, ElementTriP1
+        >>> m = (MeshTri()
+        ...      .refined()
+        ...      .with_subdomains({'left': lambda x: x[0] < .5}))
+        >>> basis = Basis(m, ElementTriP1())
+        >>> basis.get_dofs(elements='left').flatten()
+        array([0, 2, 4, 5, 6, 8])
+
+        Further reduce the set of DOFs:
+
+        >>> import numpy as np
+        >>> from skfem import MeshTri, Basis, ElementTriArgyris
+        >>> m = MeshTri().refined()
+        >>> basis = Basis(m, ElementTriArgyris())
+        >>> basis.get_dofs(lambda x: np.isclose(x[0], 0)).nodal.keys()
+        dict_keys(['u', 'u_x', 'u_y', 'u_xx', 'u_xy', 'u_yy'])
+        >>> basis.get_dofs(lambda x: np.isclose(x[0], 0)).all(['u', 'u_x'])
+        array([ 0,  1, 12, 13, 30, 31])
 
-        Accepts a richer set of types than
-        :meth:`skfem.assembly.Basis.find_dofs`.
+        Skip some DOF names altogether:
+
+        >>> import numpy as np
+        >>> from skfem import MeshTri, Basis, ElementTriArgyris
+        >>> m = MeshTri().refined()
+        >>> basis = Basis(m, ElementTriArgyris())
+        >>> basis.get_dofs(lambda x: np.isclose(x[0], 0),
+        ...                skip=['u_x', 'u_y']).nodal.keys()
+        dict_keys(['u', 'u_xx', 'u_xy', 'u_yy'])
 
         Parameters
         ----------
         facets
-            A list of facet indices. If ``None``, find facets by
+            An array of facet indices.  If ``None``, find facets by
             ``self.mesh.boundary_facets()``.  If callable, call
-            ``self.mesh.facets_satisfying(facets)`` to get the facets.
-            If array, simply find the corresponding DOF's. If a dictionary
-            of arrays, find DOF's for each entry. If a dictionary of
-            callables, call ``self.mesh.facets_satisfying`` for each entry to
-            get facets and then find DOF's for those.
+            ``self.mesh.facets_satisfying(facets)`` to get the facets.  If
+            string, use ``self.mesh.boundaries[facets]`` to get the facets.
+        elements
+            An array of element indices.  See above.
+        skip
+            List of dofnames to skip.
 
         """
+        if elements is not None:
+            if callable(elements):
+                elements = self.mesh.elements_satisfying(elements)
+            elif isinstance(elements, str):
+                elements = self.mesh.subdomains[elements]
+            return self.dofs.get_element_dofs(elements,
+                                              skip_dofnames=skip)
         if facets is None:
             facets = self.mesh.boundary_facets()
         elif callable(facets):
             facets = self.mesh.facets_satisfying(facets)
-        if isinstance(facets, dict):
+        elif isinstance(facets, dict):
             def to_indices(f):
                 if callable(f):
                     return self.mesh.facets_satisfying(f)
                 return f
-            return {k: self.dofs.get_facet_dofs(to_indices(facets[k]))
+            return {k: self.dofs.get_facet_dofs(to_indices(facets[k]),
+                                                skip_dofnames=skip)
                     for k in facets}
-        return self.dofs.get_facet_dofs(facets)
+        elif isinstance(facets, str):
+            facets = self.mesh.boundaries[facets]
+        return self.dofs.get_facet_dofs(facets,
+                                        skip_dofnames=skip)
 
     def default_parameters(self):
         """This is used by :func:`skfem.assembly.asm` to get the default

skfem/assembly/dofs.py

@@ -242,6 +242,42 @@ def __init__(self, topo, element, offset=0):
         # total dofs
         self.N = np.max(self.element_dofs) + 1
 
+    def get_element_dofs(self,
+                         elements: ndarray,
+                         skip_dofnames: List[str] = None) -> DofsView:
+        """Return a subset of DOFs corresponding to the given elements.
+
+        Parameters
+        ----------
+        elements
+            An array of element indices.
+        skip_dofnames
+            An array of dofnames to skip.
+
+        """
+        nodal_ix = (np.empty((0,), dtype=np.int64)
+                    if self.element.nodal_dofs == 0
+                    else np.unique(self.topo.t[:, elements]))
+        edge_ix = (np.empty((0,), dtype=np.int64)
+                   if self.element.edge_dofs == 0
+                   else np.unique(self.topo.t2e[:, elements]))
+        facet_ix = (np.empty((0,), dtype=np.int64)
+                    if self.element.facet_dofs == 0
+                    else np.unique(self.topo.t2f[:, elements]))
+        interior_ix = elements
+
+        if skip_dofnames is None:
+            skip_dofnames = []
+
+        return DofsView(
+            self,
+            nodal_ix,
+            facet_ix,
+            edge_ix,
+            interior_ix,
+            *self._dofnames_to_rows(skip_dofnames, skip=True)
+        )
+
     def get_facet_dofs(self,
                        facets: ndarray,
                        skip_dofnames: List[str] = None) -> DofsView: