Added progress bars to option generators

src/Environment.py

@@ -8,8 +8,11 @@
 import numpy as np
 import networkx as nx
 import util
+import sys
 import pdb
 
+from ProgressBar import ProgressBar
+
 class Environment:
     """Environment represented as an MDP"""
     S = 0
@@ -222,8 +225,27 @@ def make_options_from_random_nodes( g, gr, count ):
         random.shuffle( nodes )
         nodes = nodes[ :count ]
 
-        # Get paths to node
-        return [ OptionEnvironment.make_point_option( g, gr, n ) for n in nodes ]
+        print "Building options...\n\n"
+        progress = ProgressBar( 0, count, mode='fixed' )
+        # Needed to prevent glitches
+        print progress, "\r",
+        oldprog = str(progress)
+
+        options = []
+        for n in nodes:
+            options.append( OptionEnvironment.make_point_option( g, gr, n ) )
+
+
+            # print progress
+            progress.increment_amount()
+            if oldprog != str(progress):
+                print progress, "\r",
+                sys.stdout.flush()
+                oldprog=str(progress)
+
+        print "\n\n"
+
+        return options
 
     @staticmethod
     def make_options_from_betweenness( g, gr, count ):
@@ -241,8 +263,27 @@ def make_options_from_betweenness( g, gr, count ):
         local_maximas.sort( key = lambda (x,v): -v )
         local_maximas = [ x for (x,v) in local_maximas[ :count ] ]
 
-        # Get paths to node
-        return [ OptionEnvironment.make_point_option( g, gr, n ) for n in local_maximas ]
+        print "Building options...\n\n"
+        progress = ProgressBar( 0, count, mode='fixed' )
+        # Needed to prevent glitches
+        print progress, "\r",
+        oldprog = str(progress)
+
+        options = []
+        for n in local_maximas:
+            options.append( OptionEnvironment.make_point_option( g, gr, n ) )
+
+
+            # print progress
+            progress.increment_amount()
+            if oldprog != str(progress):
+                print progress, "\r",
+                sys.stdout.flush()
+                oldprog=str(progress)
+
+        print "\n\n"
+
+        return options
 
     @staticmethod
     def make_options_from_random_paths( g, gr, count, markov ):
@@ -252,6 +293,12 @@ def make_options_from_random_paths( g, gr, count, markov ):
         nodes = g.nodes()
         random.shuffle( nodes )
 
+        print "Building options...\n\n"
+        progress = ProgressBar( 0, count, mode='fixed' )
+        # Needed to prevent glitches
+        print progress, "\r",
+        oldprog = str(progress)
+
         options = []
         for node in nodes:
             if len( options ) > count: 
@@ -269,6 +316,15 @@ def make_options_from_random_paths( g, gr, count, markov ):
             if o:
                 options.append( o )
 
+            # print progress
+            progress.increment_amount()
+            if oldprog != str(progress):
+                print progress, "\r",
+                sys.stdout.flush()
+                oldprog=str(progress)
+
+        print "\n\n"
+
         # Get paths to node
         return options
 
@@ -283,6 +339,11 @@ def make_options_from_small_world( g, gr, count, markov, r = None ):
         # Get all the edges in the graph
         path_lengths = nx.shortest_path_length( g ).items()
         random.shuffle( path_lengths )
+        print "Building options...\n\n"
+        progress = ProgressBar( 0, count, mode='fixed' )
+        # Needed to prevent glitches
+        print progress, "\r",
+        oldprog = str(progress)
 
         options = []
         for node, dists in path_lengths:
@@ -309,6 +370,15 @@ def make_options_from_small_world( g, gr, count, markov, r = None ):
             if o:
                 options.append( o )
 
+            # print progress
+            progress.increment_amount()
+            if oldprog != str(progress):
+                print progress, "\r",
+                sys.stdout.flush()
+                oldprog=str(progress)
+
+        print "\n\n"
+
         # Get paths to node
         return options