clean up progress bar

Author: jsalvatier

pymc/progressbar.py

@@ -9,70 +9,104 @@
 import uuid
 try:
     from IPython.core.display import HTML, Javascript, display
-
-    have_ipython = True
 except ImportError:
-    have_ipython = False
+    pass
 
-class ProgressBar:
-    def __init__(self, iterations):
-        self.iterations = iterations
-        self.prog_bar = '[]'
-        self.fill_char = '*'
-        self.width = 40
-        self.__update_amount(0)
+__all__ = ['progress_bar']
 
+class ProgressBar(object): 
+    def __init__(self, iterations, animation_interval = .5):
+        self.iterations = iterations
         self.start = time.time()
         self.last = 0
-        if have_ipython:
-            self.animate = self.animate_ipython
-            self.divid = str(uuid.uuid4())
-            self.sec_id = str(uuid.uuid4())
-
-            pb = HTML(
-                """
-                <div style="float: left; border: 1px solid black; width:500px">
-                  <div id="%s" style="background-color:blue; width:0%%">&nbsp;</div>
-                </div> 
-                <label id="%s" style="padding-left: 10px;" text = ""/>
-                """ % (self.divid,self.sec_id))
-            display(pb)
-        else:
-            self.animate = self.animate_noipython
-
-    def animate_noipython(self, iter):
-        if sys.platform.lower().startswith('win'):
-            print(self, '\r', end='')
-        else:
-            print(self)
-        self.update_iteration(iter)
-
-    def animate_ipython(self, iter):
-        elapsed = time.time() - self.start
-        iter = iter + 1
-        if elapsed - self.last > .5 or iter == self.iterations:
-            self.last = elapsed
+        self.animation_interval = animation_interval
 
-            self.update_iteration(iter)
-            fraction = int(100*iter/float(self.iterations))
+    def percentage(self, i): 
+        return 100*i/float(self.iterations)
 
-            display(Javascript("$('div#%s').width('%i%%')" % (self.divid, fraction)))
-            display(Javascript("$('label#%s').text('%i%% in %.1f sec')" % (self.sec_id, fraction, round(elapsed, 1))))
+    def update(self, i):
+        elapsed = time.time() - self.start
+        i = i + 1
+        
+        if elapsed - self.last > self.animation_interval or i == self.iterations:
+             self.animate(i+1, elapsed)
+             self.last = elapsed
+        
+class TextProgressBar(ProgressBar):
+    def __init__(self, iterations, printer):
+        self.fill_char = '-'
+        self.width = 40
+        self.printer = printer
 
-    def update_iteration(self, elapsed_iter):
-        self.__update_amount((elapsed_iter / float(self.iterations)) * 100.0)
-        self.prog_bar += '  %d of %s complete' % (elapsed_iter, self.iterations)
+        ProgressBar.__init__(self, iterations)
+        self.update(0)
 
+    def animate(self, i, elapsed):
+        self.printer(self.progbar(i, elapsed))
 
-    def __update_amount(self, new_amount):
-        percent_done = int(round((new_amount / 100.0) * 100.0))
+
+    def progbar(self,i, elapsed): 
+        bar = self.bar(self.percentage(i))
+        return "[%s] %i of %i complete in %.1f sec" % (bar, i, self.iterations, round(elapsed,1))
+ 
+    def bar(self, percent):
         all_full = self.width - 2
-        num_hashes = int(round((percent_done / 100.0) * all_full))
-        self.prog_bar = '[' + self.fill_char * num_hashes + ' ' * (all_full - num_hashes) + ']'
-        pct_place = (len(self.prog_bar) / 2) - len(str(percent_done))
-        pct_string = '%d%%' % percent_done
-        self.prog_bar = self.prog_bar[0:pct_place] + \
-            (pct_string + self.prog_bar[pct_place + len(pct_string):])
-
-    def __str__(self):
-        return str(self.prog_bar)
+        num_hashes = int(percent/100 * all_full)
+        
+        bar = self.fill_char * num_hashes + ' ' * (all_full - num_hashes)
+        
+        info = '%d%%' % percent
+        loc = (  len(bar) - len(info)  )/2
+        return replace_at(bar, info, loc, loc + len(info))
+
+
+def replace_at(str, new, start, stop):
+    return str[:start] + new + str[stop:]
+
+def consoleprint(s):
+    if sys.platform.lower().startswith('win'):
+        print(s, '\r', end='')
+    else:
+        print(s)
+    
+def ipythonprint(s):
+    print('\r', s, end='')
+    sys.stdout.flush()
+
+class IPythonNotebookPB(ProgressBar):
+    def __init__(self, iterations):
+        self.divid = str(uuid.uuid4())
+        self.sec_id = str(uuid.uuid4())
+
+        pb = HTML(
+            """
+            <div style="float: left; border: 1px solid black; width:500px">
+              <div id="%s" style="background-color:blue; width:0%%">&nbsp;</div>
+            </div> 
+            <label id="%s" style="padding-left: 10px;" text = ""/>
+            """ % (self.divid,self.sec_id))
+        display(pb)
+
+        ProgressBar.__init__(self, iterations)
+
+    def animate(self, i, elapsed):
+        percentage = int(self.fraction(i))
+
+        display(Javascript("$('div#%s').width('%i%%')" % (self.divid, percentage)))
+        display(Javascript("$('label#%s').text('%i%% in %.1f sec')" % (self.sec_id, fraction, round(elapsed, 1))))
+
+def run_from_ipython():
+        try:
+            __IPYTHON__
+            return True
+        except NameError:
+            return False
+
+def progress_bar(iters):
+    if run_from_ipython(): 
+        if in_notebook():
+            return NotebookProgressBar(iters)
+        else: 
+            return IPythonConsolePB(iters, ipythonprint)
+    else:
+        return TextProgressBar(iters, consoleprint)

pymc/sample.py

@@ -4,12 +4,12 @@
 from time import time
 from core import *
 import step_methods
-from progressbar import ProgressBar
+from progressbar import progress_bar
 
 __all__ = ['sample', 'psample']
 
 @withmodel
-def sample(model, draws, step, start = None, trace = None, progress_bar = True): 
+def sample(model, draws, step, start = None, trace = None, track_progress = True): 
     """
     Draw a number of samples using the given step method. 
     Multiple step methods supported via compound step method 
@@ -51,15 +51,13 @@ def sample(model, draws, step, start = None, trace = None, progress_bar = True):
     except TypeError:
         pass
 
-
-    
-    progress = ProgressBar(draws)
+    progress = progress_bar(draws)
 
     for i in xrange(draws):
         point = step.step(point)
         trace = trace.record(point)
-        if progress_bar:
-            progress.animate(i)
+        if track_progress:
+            progress.update(i)
 
     return trace