map.py

   1 import matplotlib as mpl
   2 import matplotlib.mlab as mlab
   3 import matplotlib.pyplot as plt
   4 import matplotlib.axes as axe
   5 import matplotlib.collections as collections
   6 import numpy as np
   7 import random
   8 from pylab import *
   9 from matplotlib.pyplot import *
  10
  11 FIN = 'tmp.dat'
  12 FOUT = '../map_test.png'
  13
  14 samples = 10
  15 ticksize = 14
  16 labelsize = 24
  17 title = "a = 1.15, w = 0.1, F = 0.00, eta1 = 1.0, eta2 = 1.5, theta = 0"
  18 xlabel = r'$a$'
  19 ylabel = r'$l$'
  20
  21 xtfreq = 0.01
  22 ytfreq = 2
  23 xpad = 0.001
  24 ypad = 0.001
  25 ms = 40
  26 ystep = 0.1
  27 xstep = 0.001
  28 res = (1440/80,900/80) # default dpi is 80
  29 method = 'direct' # mean, direct
  30
  31 ################################################################################
  32
  33 def assign_label(color):
  34         ret_label = 'DUMMY'
  35         if (color == 'grey'):
  36                 ret_label = r'$|v| = 0$'
  37         if (color == 'blue'):
  38                 ret_label = r'$|v| = -1$'
  39         if (color == 'red'):
  40                 ret_label = r'$|v| = 1$'
  41         if (color == 'orange'):
  42                 ret_label = r'$|v| = 0.5$'
  43         if (color == 'green'):
  44                 ret_label = r'$|v| = -0.5$'
  45         if (color == 'black'):
  46                 ret_label = r'other rational $v$'
  47         if (color == 'cyan'):
  48                 ret_label = r'$v < 0$'
  49         if (color == 'magenta'):
  50                 ret_label = r'$v > 0$'
  51         return (ret_label)
  52
  53 def assign_color2(colors,xar,yar,clist):
  54         print "---------------------------------------------------------"
  55         unique_colors = unique(colors)
  56         #print range(len(clist))
  57         ret_xcolors = [ [] for DUMMYVAR in range(len(clist)) ]
  58         ret_ycolors = [ [] for DUMMYVAR in range(len(clist)) ]
  59         i = 0
  60         assert len(xar) == len(yar) == len(colors)
  61         xarlen_trunc = len(xar)
  62         while i < xarlen_trunc:
  63                 tmp = []
  64                 for j in range(samples):
  65                         assert i+j <= len(xar)
  66                         #assert yar[i+j] == yar[i] and xar[i+j] == xar[i]
  67                         tmp.append(colors[i+j])
  68 #               print len(tmp), samples
  69                 assert len(tmp) == samples
  70                 if 'green' in tmp and 'orange' in tmp:
  71                         ret_xcolors[6].append(xar[i])
  72                         ret_ycolors[6].append(yar[i])
  73                 elif 'blue' in tmp and 'red' in tmp:
  74                         ret_xcolors[7].append(xar[i])
  75                         ret_ycolors[7].append(yar[i])
  76                 elif len(unique(tmp)) == 1 and 'grey' in tmp:
  77                         ret_xcolors[1].append(xar[i])
  78                         ret_ycolors[1].append(yar[i])
  79                 elif len(unique(tmp)) == 1 and 'blue' in tmp:
  80                         ret_xcolors[2].append(xar[i])
  81                         ret_ycolors[2].append(yar[i])
  82                 elif len(unique(tmp)) == 1 and 'green' in tmp:
  83                         ret_xcolors[3].append(xar[i])
  84                         ret_ycolors[3].append(yar[i])
  85                 elif len(unique(tmp)) == 1 and 'orange' in tmp:
  86                         ret_xcolors[4].append(xar[i])
  87                         ret_ycolors[4].append(yar[i])
  88                 elif len(unique(tmp)) == 1 and 'red' in tmp:
  89                         ret_xcolors[5].append(xar[i])
  90                         ret_ycolors[5].append(yar[i])
  91                 elif len(unique(tmp)) == 1 and 'black' in tmp:
  92                         if (xar[i] > 0):
  93                                 ret_xcolors[8].append(xar[i])
  94                         if (xar[i] < 0):
  95                                 ret_xcolors[9].append(xar[i])
  96                         if (yar[i] > 0):
  97                                 ret_ycolors[8].append(yar[i])
  98                         if (yar[i] < 0):
  99                                 ret_ycolors[9].append(yar[i])
 100                 else:
 101                         ret_xcolors[0].append(xar[i])
 102                         ret_ycolors[0].append(yar[i])
 103                 i = i+samples
 104                 #print "increment:",i
 105         for j  in range(10):
 106                 print clist[j],":",len(ret_xcolors[j])
 107         print "---------------------------------------------------------"
 108         return (ret_xcolors, ret_ycolors)
 109
 110 def assign_color(colors,xar,yar):
 111         print "---------------------------------------------------------"
 112         unique_colors = unique(colors)
 113         ret_xcolors = [ [] for DUMMYVAR in range(len(unique_colors)) ]
 114         ret_ycolors = [ [] for DUMMYVAR in range(len(unique_colors)) ]
 115         for j in range(len(unique_colors)):
 116                 tmpx = []
 117                 tmpy = []
 118                 k = 0
 119                 for i in range(len(colors)):
 120                         if (colors[i] == unique_colors[j]):
 121                                 k = k+1
 122                                 tmpxcolors = ret_xcolors[j]
 123                                 tmpycolors = ret_ycolors[j]
 124                                 #print tmpxcolors
 125                                 if len(tmpxcolors) > 0:
 126                                         if (tmpxcolors[-1] != xar[i]) or (tmpycolors[-1] != yar[i]):
 127                                                 #print tmpxcolors[-1],tmpycolors[-1],"is not", xar[i],yar[i],"!"
 128                                                 ret_xcolors[j].append(xar[i])
 129                                                 ret_ycolors[j].append(yar[i])
 130                                 else:
 131                                         ret_xcolors[j].append(xar[i])
 132                                         ret_ycolors[j].append(yar[i])
 133                 lenc = float(len(colors))
 134                 per = 100*k/lenc
 135                 print unique_colors[j],":",len(ret_xcolors[j]),k,"(",'%.2f' % per ,"%)"
 136         print "---------------------------------------------------------"
 137         return (ret_xcolors, ret_ycolors)
 138
 139 print "reading data..."
 140 r  = mlab.csv2rec(FIN, delimiter='\t', comments='#')
 141
 142 x = r.x
 143 x_color = r.vx_color
 144 y_color = r.vy_color
 145 y = r.y
 146 vx = r.vx
 147 vy = r.vy
 148
 149 assert len(x) == len(y) == len(vx) == len(vy)
 150 print "done."
 151
 152 xmin = np.min(unique(x))
 153 xmax = np.max(unique(x))
 154 ymin = np.min(unique(y))
 155 ymax = np.max(unique(y))
 156
 157 # assumes that the step distance doesn't change over the course of time
 158 rat = xstep/ystep
 159 print "size of xstep:", xstep
 160 print "size of ystep:", ystep
 161 print "aspect ratio:",rat
 162
 163 print "different values of vx:"
 164 print unique(vx),"(",len(vx),"total)"
 165
 166 print "different values of vy:"
 167 print unique(vy),"(",len(vy),"total)"
 168
 169 # needed for assign_color2
 170 color_list = [
 171                 'black',        # 0     everything else
 172                 'grey',         # 1     v = 0
 173                 'blue',         # 2     v = -1.0
 174                 'green',        # 3     v = -0.5
 175                 'orange',       # 4     v = 0.5
 176                 'red',          # 5     v = 1.0
 177                 'teal',         # 6     v = -0.5 && v = 0.5
 178                 'purple',       # 7     v = -1.0 && v = 1.0
 179                 'magenta',      # 8     v > 0
 180                 'cyan'          # 9     v < 0
 181         ]
 182
 183 label_list = [
 184                 r'other rational $v$',  # 0 black
 185                 r'$v = 0$',                             # 1 grey
 186                 r'$v = -1$',                    # 2 blue
 187                 r'$v = -0.5$',                  # 3 green
 188                 r'$v = 0.5$',                   # 4 orange
 189                 r'$v = 1.0$',                   # 5 red
 190                 r'$v = \pm 0.5$',               # 6 teal
 191                 r'$v = \pm 1.0$',               # 7 purple
 192                 r'$v > 0$',                     # 8 magenta
 193                 r'$v < 0$',                     # 9 cyan
 194         ]
 195
 196 xpadmin = xmin-xstep
 197 xpadmax = xmax+xstep
 198 ypadmin = ymin-ystep
 199 ypadmax = ymax+ystep
 200 xlim = (xpadmin, xpadmax)               # small padding between plot and
 201 ylim = (ypadmin, ypadmax)               # axis so they don't overlap
 202 xticks = np.arange(xmin, xmax+xpad, xtfreq)
 203 yticks = np.arange(ymin, ymax+ypad, ytfreq)
 204
 205
 206 fig = plt.figure(figsize=res)
 207 #fig.suptitle(title, size=labelsize)
 208
 209 max = 0.84
 210 c = max/2
 211 rect1 = [0.1, 0.1, c, 0.9]
 212 rect2 = [c+0.1, 0.1, c, 0.9]
 213
 214 ax1 = fig.add_axes(rect1, aspect=rat)
 215 ax2 = fig.add_axes(rect2, aspect=rat)
 216
 217 if (method == 'mean'):
 218
 219         print "assigning colors...."
 220         xcolors_x,xcolors_y = assign_color(x_color,x,y)
 221         ycolors_x,ycolors_y = assign_color(y_color,x,y)
 222         print "done."
 223
 224         print "plotting..."
 225         for k in range(len(unique(x_color))):
 226                 x_xtli = xcolors_x[k]
 227                 x_ytli = xcolors_y[k]
 228                 labelx = assign_label(unique(x_color)[k])
 229                 ax1.scatter(x_xtli, x_ytli, marker='s', lw=0, s=ms, alpha=0.5,
 230                                 c = unique(x_color)[k], label=labelx,
 231                                 edgecolor='face')
 232
 233         for k in range(len(unique(y_color))):
 234                 y_xtli = ycolors_x[k]
 235                 y_ytli = ycolors_y[k]
 236                 labely = assign_label(unique(y_color)[k])
 237                 ax2.scatter(y_xtli, y_ytli, marker='s', lw=0, s=ms, alpha=0.5,
 238                                 c = unique(y_color)[k], label=labely,
 239                                 edgecolor='none')
 240
 241         print "done."
 242
 243 elif (method == 'direct'):
 244
 245         print "assigning colors..."
 246         xcolors_x,xcolors_y = assign_color2(x_color,x,y,color_list)
 247         ycolors_x,ycolors_y = assign_color2(y_color,x,y,color_list)
 248         print "done."
 249
 250         #for i in range(8):
 251         #       print i, xcolors_x[i],xcolors_y[i];
 252
 253         print "plotting..."
 254         for k in range(len(xcolors_x)):
 255                 if len(xcolors_x[k]) > 0:
 256                         x_xtli = xcolors_x[k]
 257                         x_ytli = xcolors_y[k]
 258                         ax1.scatter(x_xtli, x_ytli, marker='s', lw=0,
 259                                         s=ms, alpha=1, c=color_list[k],
 260                                         label=label_list[k])
 261
 262         for k in range(len(ycolors_x)):
 263                 if len(ycolors_x[k]) > 0:
 264                         y_xtli = ycolors_x[k]
 265                         y_ytli = ycolors_y[k]
 266                         ax2.scatter(y_xtli, y_ytli, marker='s', lw=0,
 267                                         s=ms, alpha=1, c=color_list[k],
 268                                         label=label_list[k])
 269
 270         print "done."
 271 else:
 272         print "fatal error: specify a valid method"
 273
 274 leg1 = ax1.legend(bbox_to_anchor=(0.05, 1.05), loc=2, shadow=True,
 275                 fancybox=True, scatterpoints=1, markerscale=1,
 276                 borderaxespad=0.)
 277 leg2 = ax2.legend(bbox_to_anchor=(0.85, 1.05), loc=2, shadow=True,
 278                 fancybox=True, scatterpoints=1, markerscale=1,
 279                 borderaxespad=0.)
 280 for t in leg1.get_texts():
 281         t.set_fontsize('small')
 282 for t in leg2.get_texts():
 283         t.set_fontsize('small')
 284
 285 ax1.set_title(r'$v_x$', size=labelsize)
 286 ax1.set_xlim(xlim)
 287 ax1.set_ylim(ylim)
 288 ax1.set_xlabel(xlabel, size=labelsize)
 289 ax1.set_ylabel(ylabel, size=labelsize)
 290 ax1.set_xticks(xticks, minor=False)
 291 ax1.set_yticks(yticks, minor=False)
 292 ax1.set_xticklabels(ax1.get_xticks(), size=ticksize)
 293 ax1.set_yticklabels(ax1.get_yticks(), size=ticksize)
 294
 295 ax2.set_title(r'$v_y$', size=labelsize)
 296 ax2.set_xlim(xlim)
 297 ax2.set_ylim(ylim)
 298 ax2.set_xlabel(xlabel, size=labelsize)
 299 ax2.set_ylabel('', size=labelsize)
 300 ax2.set_xticks(xticks, minor=False)
 301 ax2.set_yticks(yticks, minor=False)
 302 ax2.set_xticklabels(ax2.get_xticks(), size=ticksize)
 303 ax2.set_yticklabels('', size=ticksize)
 304
 305 leg1.get_frame().set_alpha(0.75)
 306 leg2.get_frame().set_alpha(0.75)
 307
 308 plt.subplots_adjust(hspace=0)
 309 print "flushing into a file..."
 310 plt.savefig(FOUT)
 311 print "done."
 312 #plt.show()
 313
 314
 315
 316