subgraph: sub-sample signals as well

pygsp/graphs/graph.py

@@ -222,7 +222,8 @@ def subgraph(self, vertices):
         Parameters
         ----------
         vertices : list
-            List of vertices to keep.
+            Vertices to keep.
+            Either a list of indices or an indicator function.
 
         Returns
         -------
@@ -249,7 +250,10 @@ def subgraph(self, vertices):
             coords = self.coords[vertices]
         except AttributeError:
             coords = None
-        return Graph(adjacency, self.lap_type, coords, self.plotting)
+        graph = Graph(adjacency, self.lap_type, coords, self.plotting)
+        for name, signal in self.signals.items():
+            graph.set_signal(signal[vertices], name)
+        return graph
 
     def is_connected(self):
         r"""Check if the graph is connected (cached).

pygsp/tests/test_graphs.py

@@ -449,9 +449,11 @@ def test_set_coordinates(self):
         self.assertRaises(ValueError, G.set_coordinates, 'invalid')
 
     def test_subgraph(self, n_vertices=100):
+        self._G.set_signal(self._G.coords, 'coords')
         graph = self._G.subgraph(range(n_vertices))
         self.assertEqual(graph.n_vertices, n_vertices)
         self.assertEqual(graph.coords.shape, (n_vertices, 2))
+        self.assertEqual(graph.signals['coords'].shape, (n_vertices, 2))
         self.assertIs(graph.lap_type, self._G.lap_type)
         self.assertEqual(graph.plotting, self._G.plotting)