visualisation added

simanneal/anneal.py

@@ -11,21 +11,21 @@
 import signal
 import sys
 import time
+import matplotlib.pyplot as plt
+
 
 
 def round_figures(x, n):
     """Returns x rounded to n significant figures."""
     return round(x, int(n - math.ceil(math.log10(abs(x)))))
 
-
 def time_string(seconds):
     """Returns time in seconds as a string formatted HHHH:MM:SS."""
     s = int(round(seconds))  # round to nearest second
     h, s = divmod(s, 3600)   # get hours and remainder
     m, s = divmod(s, 60)     # split remainder into minutes and seconds
     return '%4i:%02i:%02i' % (h, m, s)
 
-
 class Annealer(object):
 
     """Performs simulated annealing by calling functions to calculate
@@ -50,14 +50,19 @@ class Annealer(object):
     start = None
 
     def __init__(self, initial_state=None, load_state=None):
+
+        self.energy_list = []
+        self.initial_energy = None
+        self.vis = True
+        self.save = True
+
         if initial_state is not None:
             self.state = self.copy_state(initial_state)
         elif load_state:
             self.load_state(load_state)
         else:
             raise ValueError('No valid values supplied for neither \
             initial_state nor load_state')
-
         signal.signal(signal.SIGINT, self.set_user_exit)
 
     def save_state(self, fname=None):
@@ -73,6 +78,27 @@ def load_state(self, fname=None):
         with open(fname, 'rb') as fh:
             self.state = pickle.load(fh)
 
+    def save_plot(self, fname=None):
+        """Saves progress plot"""
+
+        if not fname:
+            date = datetime.datetime.now().strftime("%Y-%m-%dT%Hh%Mm%Ss")
+            fname = date 
+
+        print(len(self.energy_list))
+        fig = plt.figure(figsize=(10,8))
+        plt.plot([i for i in range(len(self.energy_list))], self.energy_list)
+        initial = plt.axhline(y=self.initial_energy, color='r', linestyle='--')
+        final = plt.axhline(y=self.initial_energy, color='g', linestyle='--')
+        plt.legend( [ initial , final ] , ['Initial Energy', 'Final Energy'])
+        plt.ylabel('State Energy')
+        plt.xlabel('Iteration')
+
+        if self.save:
+            fig.savefig( 'progress.png', bbox_inches='tight')
+
+        plt.show()
+
     @abc.abstractmethod
     def move(self):
         """Create a state change"""
@@ -123,7 +149,7 @@ def update(self, *args, **kwargs):
         from your own Annealer.
         """
         self.default_update(*args, **kwargs)
-
+        
     def default_update(self, step, T, E, acceptance, improvement):
         """Default update, outputs to stderr.
 
@@ -145,6 +171,8 @@ def default_update(self, step, T, E, acceptance, improvement):
         are exhausted and moves that would increase the energy are no longer
         thermally accessible."""
 
+        self.energy_list.append(E)
+
         elapsed = time.time() - self.start
         if step == 0:
             print('\n Temperature        Energy    Accept   Improve     Elapsed   Remaining',
@@ -188,6 +216,8 @@ def anneal(self):
         # Note initial state
         T = self.Tmax
         E = self.energy()
+        self.initial_energy = E
+
         prevState = self.copy_state(self.state)
         prevEnergy = E
         self.best_state = self.copy_state(self.state)
@@ -229,9 +259,13 @@ def anneal(self):
                     trials = accepts = improves = 0
 
         self.state = self.copy_state(self.best_state)
+
         if self.save_state_on_exit:
             self.save_state()
 
+        if self.vis:
+            self.save_plot()
+
         # Return best state and energy
         return self.best_state, self.best_energy