Merge pull request #640 from martinfleis/subset

ENH: add subgraph method to Graph to get subsets

libpysal/graph/_utils.py

@@ -246,7 +246,7 @@ def _resolve_islands(heads, tails, ids, weights):
     Induce self-loops for a collection of ids and links describing a
     contiguity graph. Induced self-loops will have zero weight.
     """
-    islands = np.setdiff1d(ids, heads)
+    islands = pd.Index(ids).difference(pd.Index(heads))
     if islands.shape != (0,):
         heads = np.hstack((heads, islands))
         tails = np.hstack((tails, islands))

libpysal/graph/base.py

@@ -20,7 +20,12 @@
 from ._set_ops import SetOpsMixin
 from ._spatial_lag import _lag_spatial
 from ._triangulation import _delaunay, _gabriel, _relative_neighborhood, _voronoi
-from ._utils import _evaluate_index, _neighbor_dict_to_edges, _sparse_to_arrays
+from ._utils import (
+    _evaluate_index,
+    _neighbor_dict_to_edges,
+    _resolve_islands,
+    _sparse_to_arrays,
+)
 
 ALLOWED_TRANSFORMATIONS = ("O", "B", "R", "D", "V")
 
@@ -1337,6 +1342,45 @@ def explore(
             **kwargs,
         )
 
+    def subgraph(self, ids):
+        """Returns a subset of Graph containing only nodes specified in ids
+
+        The resulting subgraph contains only the nodes in ``ids`` and the edges
+        between them or zero-weight self-loops in case of isolates.
+
+        The order of ``ids`` reflects a new canonical order of the resulting
+        subgraph. This means ``ids`` should be equal to the index of the DataFrame
+        containing data linked to the graph to ensure alignment of sparse representation
+        of subgraph.
+
+        Parameters
+        ----------
+        ids : array-like
+            An array of node IDs to be retained
+
+        Returns
+        -------
+        Graph
+            A new Graph that is a subset of the original
+
+        Notes
+        -----
+        Unlike the implementation in ``networkx``, this creates a copy since
+        Graphs in ``libpysal`` are immutable.
+        """
+        masked_adj = self._adjacency[ids]
+        filtered_adj = masked_adj[
+            masked_adj.index.get_level_values("neighbor").isin(ids)
+        ]
+        return Graph.from_arrays(
+            *_resolve_islands(
+                filtered_adj.index.get_level_values("focal"),
+                filtered_adj.index.get_level_values("neighbor"),
+                ids,
+                filtered_adj.values,
+            )
+        )
+
 
 def _arrange_arrays(heads, tails, weights, ids=None):
     """

libpysal/graph/tests/test_base.py

@@ -941,3 +941,20 @@ def test_component_labels(self):
         pd.testing.assert_series_equal(
             expected, nybb.component_labels, check_dtype=False
         )
+
+    def test_subgraph(self):
+        knn = graph.Graph.build_knn(self.nybb.set_geometry(self.nybb.centroid), k=2)
+        sub = knn.subgraph(["Staten Island", "Bronx", "Brooklyn"])
+        assert sub < knn
+        expected = pd.Series(
+            [1, 0, 0],
+            name="weight",
+            index=pd.MultiIndex.from_arrays(
+                [
+                    ["Staten Island", "Bronx", "Brooklyn"],
+                    ["Brooklyn", "Bronx", "Brooklyn"],
+                ],
+                names=["focal", "neighbor"],
+            ),
+        )
+        pd.testing.assert_series_equal(expected, sub._adjacency, check_dtype=False)