refactor visualization functions

Author: lsampras

Visualization/MapVisualizer.py

@@ -2,6 +2,7 @@
 
 import matplotlib.pyplot as plt
 import matplotlib.cm as cm
+import matplotlib.colors as colors
 import numpy as np
 import itertools
 
@@ -204,3 +205,18 @@ def mouse_move(event):
     x, y = int(round(event.xdata)), int(round(event.ydata))
     print(x, y)
     return
+
+def saveColorGradedMap(name,map,color='Reds',maxRange=1.0,minRange=0.0,underRange=(0.,0.,0.),overRange=(1.,1.,1.)):
+    norm = colors.Normalize(vmin=minRange, vmax=maxRange)
+    cmap = 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)

mapworks/visualization.py

@@ -2,7 +2,7 @@
 import numpy as np
 import moves
 import random
-from Visualization.MapVisualizer import saveColorGradedMap
+from mapworks import visualization
 
 np.set_printoptions(formatter={'float': lambda x: "{0:0.3f}".format(x)})
 
@@ -27,7 +27,7 @@ def addEnemy():
     if(turn%1 == 0):
         addEnemy()
         # TODO: make sure directory exists
-        saveColorGradedMap("pheromone-sample/{}".format(turn),ph.map,underRange=(0.,0.,1.))
+        visualization.saveColorGradedMap("pheromone-sample/{}".format(turn),ph.map,underRange=(0.,0.,1.))
     ph.propogate()