Extend UST class to handle weighted graphs

dppy/exotic_dpps.py

@@ -487,9 +487,9 @@ def __init__(self, graph):
         self.edge_labels = {edge: r'$e_{}$'.format(i)
                             for i, edge in enumerate(self.edges)}
 
-        self.neighbors = [list(graph.neighbors(v))
-                          for v in range(self.nb_nodes)]
-
+        #self.neighbors = [list(graph.neighbors(v)) for v in range(self.nb_nodes)]
+        self.adjacency_matrix = nx.adjacency_matrix(self.graph)
+        
         self.sampling_mode = 'Wilson'  # Default (avoid eig_vecs computation)
         self._sampling_modes = {'markov-chain': ['Wilson', 'Aldous-Broder', 'Wilson_node'],
                                 'spectral-method': ['GS'],
@@ -559,16 +559,15 @@ def sample(self, mode='Wilson', root=None, random_state=None):
 
         if self.sampling_mode in self._sampling_modes['markov-chain']:
             if self.sampling_mode == 'Wilson':
-                sampl = ust_sampler_wilson(self.neighbors,
+                sampl = ust_sampler_wilson(self.adjacency_matrix,
                                            random_state=rng)
 
             elif self.sampling_mode == 'Aldous-Broder':
-                sampl = ust_sampler_aldous_broder(self.neighbors,
+                sampl = ust_sampler_aldous_broder(self.adjacency_matrix,
                                                   random_state=rng)
 
             elif self.sampling_mode == 'Wilson_node':
-                W = nx.adjacency_matrix(self.graph).todense()
-                Y, P, sampl = ust_sampler_wilson_nodes(W, absorbing_weight=1, random_state=rng)
+                Y, P, sampl = ust_sampler_wilson_nodes(self.adjacency_matrix, absorbing_weight=1, random_state=rng)
 
 
         elif self.sampling_mode in self._sampling_modes['spectral-method']:

dppy/exotic_dpps_core.py

@@ -36,14 +36,22 @@
 from dppy.utils import check_random_state
 
 
-def ust_sampler_wilson(list_of_neighbors, root=None,
+def ust_sampler_wilson(W, root=None,
                        random_state=None):
-
+    """
+    Compute a random spanning tree of a graph G given by his adjacency matrix
+    
+    param W:
+        Adjacency matrix of the Graph
+    param W:
+        scipy.sparse.csr_matrix
+    """
     rng = check_random_state(random_state)
 
     # Initialize the tree
     wilson_tree_graph = nx.Graph()
-    nb_nodes = len(list_of_neighbors)
+    #nb_nodes = len(list_of_neighbors)
+    nb_nodes = W.shape[0]
 
     # Initialize the root, if root not specified start from any node
     n0 = root if root else rng.choice(nb_nodes)  # size=1)[0]
@@ -57,7 +65,10 @@ def ust_sampler_wilson(list_of_neighbors, root=None,
     while nb_nodes_in_tree < nb_nodes:  # |Tree| = |V| - 1
 
         # visit a neighbor of n0 uniformly at random
-        n1 = rng.choice(list_of_neighbors[n0])  # size=1)[0]
+        #n1 = rng.choice(list_of_neighbors[n0])  # size=1)[0]
+        weights = (W.getrow(n0).toarray())[0].astype('float')
+        weights /= np.sum(weights)
+        n1 = rng.choice(np.arange(nb_nodes), p=weights)
 
         if state[n1] == -1:  # not visited => continue the walk
 
@@ -105,7 +116,7 @@ def ust_sampler_wilson_nodes(W, absorbing_weight=0, random_state=None):
     :param W:
         Adjacency matrix of the graph
     :type W:
-        array_like
+        scipy.sparse.csr_matrix
 
     :param absorbing_weight:
         Weight of the node Delta added to the graph
@@ -138,7 +149,7 @@ def ust_sampler_wilson_nodes(W, absorbing_weight=0, random_state=None):
     all_path = []
 
     # Compute the probabilities of transition
-    transition_probabilities = np.pad(W, [(0, 0), (0, 1)], mode='constant', constant_values=absorbing_weight).astype('float')
+    transition_probabilities = np.pad(W.toarray(), [(0, 0), (0, 1)], mode='constant', constant_values=absorbing_weight).astype('float')
     norm = np.sum(transition_probabilities, axis=1)
     transition_probabilities[np.nonzero(norm), :] /= norm[np.nonzero(norm), None]
 
@@ -211,14 +222,15 @@ def ust_sampler_wilson_nodes(W, absorbing_weight=0, random_state=None):
     #print("Nu=", Nu)
     return Y, all_path, wilson_tree_from_path
 
-def ust_sampler_aldous_broder(list_of_neighbors, root=None,
+def ust_sampler_aldous_broder(W, root=None,
                               random_state=None):
 
     rng = check_random_state(random_state)
 
     # Initialize the tree
     aldous_tree_graph = nx.Graph()
-    nb_nodes = len(list_of_neighbors)
+    #nb_nodes = len(list_of_neighbors)
+    nb_nodes = W.shape[0]
 
     # Initialize the root, if root not specified start from any node
     n0 = root if root else rng.choice(nb_nodes)  # size=1)[0]
@@ -231,7 +243,10 @@ def ust_sampler_aldous_broder(list_of_neighbors, root=None,
     while nb_nodes_in_tree < nb_nodes:
 
         # visit a neighbor of n0 uniformly at random
-        n1 = rng.choice(list_of_neighbors[n0])  # size=1)[0]
+        #n1 = rng.choice(list_of_neighbors[n0])  # size=1)[0]
+        weights = (W.getrow(n0).toarray())[0].astype('float')
+        weights /= np.sum(weights)
+        n1 = rng.choice(np.arange(nb_nodes), p=weights)
 
         if visited[n1]:
             pass  # continue the walk