refactor visualization function

Author: lsampras

Visualization/MapVisualizer.py

@@ -0,0 +1,20 @@
+import matplotlib.pyplot as plt
+from matplotlib import colors
+import matplotlib as mpl
+import numpy as np
+
+
+def saveColorGradedMap(name,map,color='Reds',maxRange=1.0,minRange=0.0,underRange=(0.,0.,0.),overRange=(1.,1.,1.)):
+    norm = mpl.colors.Normalize(vmin=minRange, vmax=maxRange)
+    cmap = mpl.cm.get_cmap(color)
+    cmap.set_under(underRange)
+    cmap.set_over(overRange)
+
+    fig, ax = plt.subplots()
+    ax.imshow(map, cmap=cmap, norm=norm)
+    # draw gridlines
+    ax.grid(which='major', axis='both', linestyle='-', color='k', linewidth=2)
+    ax.set_xticks(np.arange(-0.5, map.shape[0], 1))
+    ax.set_yticks(np.arange(-0.5, map.shape[1], 1))
+    plt.savefig(name + ".png")
+    plt.close(fig)

pheromone-driver.py

@@ -2,6 +2,8 @@
 import numpy as np
 import moves
 import random
+from Visualization.MapVisualizer import saveColorGradedMap
+
 np.set_printoptions(formatter={'float': lambda x: "{0:0.3f}".format(x)})
 
 def getObstacleData(size,prob):
@@ -19,15 +21,13 @@ def addEnemy():
 for i in range(40):
         addEnemy()
 
-# make sure directory exists
+
 
 for turn in range(10):
     if(turn%1 == 0):
-        # print("\n\n\n")
-        # print(ph.map)
         addEnemy()
-        ph.save_map("pheromone-sample/{}".format(turn))
+        # TODO: make sure directory exists
+        saveColorGradedMap("pheromone-sample/{}".format(turn),ph.map,underRange=(0.,0.,1.))
     ph.propogate()
 
 
-ph.save_map("pheromone-sample/{}".format(turn))

pheromone.py

@@ -1,7 +1,4 @@
 import numpy as np
-import matplotlib.pyplot as plt
-from matplotlib import colors
-import matplotlib as mpl
 from point import Point
 class pheromone(object):
     def __init__(self,obstacleData,possible_moves,decayRate = 0.7,dropOff = 0.1):
@@ -13,15 +10,11 @@ def __init__(self,obstacleData,possible_moves,decayRate = 0.7,dropOff = 0.1):
         self.moves.append(Point(0,0))
         self.multiplier = (1 - decayRate) #normalize
         self.temp = np.zeros(shape = self.map.shape,dtype = np.float32)
-        self.cmap = mpl.cm.Reds
-        self.cmap.set_under((0.,0.,1.))
-        self.norm = mpl.colors.Normalize(vmin=0, vmax=1)
         return
 
     def addEnemy(self,x,y):
         self.map[x][y] = 1.0
 
-
     def decay(self):
         self.map = self.map*self.multiplier
 
@@ -49,14 +42,3 @@ def propogate_point(self,value,x,y):
         delta = (value * self.multiplier)/len(changes)
         for pos in changes:
             self.temp[pos[0]][pos[1]] +=delta
-
-
-    def save_map(self,name):
-        fig, ax = plt.subplots()
-        ax.imshow(self.map, cmap=self.cmap, norm=self.norm)
-        # draw gridlines
-        ax.grid(which='major', axis='both', linestyle='-', color='k', linewidth=2)
-        ax.set_xticks(np.arange(-0.5, self.map.shape[0], 1));
-        ax.set_yticks(np.arange(-0.5, self.map.shape[1], 1));
-        plt.savefig(name + ".png")
-        plt.close(fig)