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 = 2
  22 ytfreq = 1
  23 xpad = 0.01
  24 ypad = 0.01
  25 ms = 750
  26 ystep = 0.01
  27 xstep = 0.1
  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,vel,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                 tmpv = []
  65                 for j in range(samples):
  66                         assert i+j <= len(xar)
  67                         #assert yar[i+j] == yar[i] and xar[i+j] == xar[i]
  68                         tmp.append(colors[i+j])
  69                         tmpv.append(vel[i+j])
  70 #               print len(tmp), samples
  71                 assert len(tmp) == samples
  72                 ccount = len(unique(tmp))
  73                 if ccount == 2 and 'green' in tmp and 'orange' in tmp:
  74                         ret_xcolors[6].append(xar[i])
  75                         ret_ycolors[6].append(yar[i])
  76                 elif ccount == 2 and 'blue' in tmp and 'red' in tmp:
  77                         ret_xcolors[7].append(xar[i])
  78                         ret_ycolors[7].append(yar[i])
  79                 elif ccount == 3 and 'green' in tmp and 'orange' in tmp and 'grey' in tmp:
  80                         ret_xcolors[12].append(xar[i])
  81                         ret_ycolors[12].append(yar[i])
  82                 elif ccount == 3 and 'blue' in tmp and 'red' in tmp and 'grey' in tmp:
  83                         ret_xcolors[13].append(xar[i])
  84                         ret_ycolors[13].append(yar[i])
  85                 elif ccount == 4 and 'blue' in tmp and 'red' in tmp and 'green' in tmp and 'orange' in tmp:
  86                         print tmp
  87                         ret_xcolors[14].append(xar[i])
  88                         ret_ycolors[14].append(yar[i])
  89                 elif ccount == 1 and 'grey' in tmp:
  90                         ret_xcolors[1].append(xar[i])
  91                         ret_ycolors[1].append(yar[i])
  92                 elif ccount == 1 and 'blue' in tmp:
  93                         ret_xcolors[2].append(xar[i])
  94                         ret_ycolors[2].append(yar[i])
  95                 elif ccount == 1 and 'green' in tmp:
  96                         ret_xcolors[3].append(xar[i])
  97                         ret_ycolors[3].append(yar[i])
  98                 elif ccount == 1 and 'orange' in tmp:
  99                         ret_xcolors[4].append(xar[i])
 100                         ret_ycolors[4].append(yar[i])
 101                 elif ccount == 1 and 'red' in tmp:
 102                         ret_xcolors[5].append(xar[i])
 103                         ret_ycolors[5].append(yar[i])
 104                 elif ('green' in tmp or 'blue' in tmp) and ('black' in tmp or 'grey' in tmp):
 105                         ret_xcolors[0].append(xar[i])
 106                         ret_ycolors[0].append(yar[i])
 107                 elif ('orange' in tmp or 'red' in tmp) and ('black' in tmp or 'grey' in tmp):
 108                         ret_xcolors[0].append(xar[i])
 109                         ret_ycolors[0].append(yar[i])
 110                 elif ccount == 1 and 'black' in tmp:
 111                         ret_xcolors[0].append(xar[i])
 112                         ret_ycolors[0].append(yar[i])
 113                 elif ccount == 2 and 'black' in tmp and 'grey' in tmp:
 114                         if tmp.count ('black') >= samples/2:
 115                                 ret_xcolors[15].append(xar[i])
 116                                 ret_ycolors[15].append(yar[i])
 117                         if tmp.count ('black') < samples/2:
 118                                 ret_xcolors[11].append(xar[i])
 119                                 ret_ycolors[11].append(yar[i])
 120                 elif ccount >= 5:
 121                         ret_xcolors[0].append(xar[i])
 122                         ret_ycolors[0].append(yar[i])
 123                 else:
 124                         print tmp
 125                         ret_xcolors[10].append(xar[i])
 126                         ret_ycolors[10].append(yar[i])
 127                 i = i+samples
 128                 #print "increment:",i
 129         for j  in range(len(clist)):
 130                 print clist[j],":",len(ret_xcolors[j])
 131         #print "unmatched velocities:"
 132         print ret_xcolors[13], ret_ycolors[13]
 133         print "---------------------------------------------------------"
 134         return (ret_xcolors, ret_ycolors)
 135
 136 def assign_color(colors,xar,yar):
 137         print "---------------------------------------------------------"
 138         unique_colors = unique(colors)
 139         ret_xcolors = [ [] for DUMMYVAR in range(len(unique_colors)) ]
 140         ret_ycolors = [ [] for DUMMYVAR in range(len(unique_colors)) ]
 141         for j in range(len(unique_colors)):
 142                 tmpx = []
 143                 tmpy = []
 144                 k = 0
 145                 for i in range(len(colors)):
 146                         if (colors[i] == unique_colors[j]):
 147                                 k = k+1
 148                                 tmpxcolors = ret_xcolors[j]
 149                                 tmpycolors = ret_ycolors[j]
 150                                 #print tmpxcolors
 151                                 if len(tmpxcolors) > 0:
 152                                         if (tmpxcolors[-1] != xar[i]) or (tmpycolors[-1] != yar[i]):
 153                                                 #print tmpxcolors[-1],tmpycolors[-1],"is not", xar[i],yar[i],"!"
 154                                                 ret_xcolors[j].append(xar[i])
 155                                                 ret_ycolors[j].append(yar[i])
 156                                 else:
 157                                         ret_xcolors[j].append(xar[i])
 158                                         ret_ycolors[j].append(yar[i])
 159                 lenc = float(len(colors))
 160                 per = 100*k/lenc
 161                 print unique_colors[j],":",len(ret_xcolors[j]),k,"(",'%.2f' % per ,"%)"
 162         print "---------------------------------------------------------"
 163         return (ret_xcolors, ret_ycolors)
 164
 165 print "reading data..."
 166 r  = mlab.csv2rec(FIN, delimiter='\t', comments='#')
 167
 168 x = r.x
 169 x_color = r.vx_color
 170 y_color = r.vy_color
 171 y = r.y
 172 vx = r.vx
 173 vy = r.vy
 174
 175 assert len(x) == len(y) == len(vx) == len(vy)
 176 print "done."
 177
 178 xmin = np.min(unique(x))
 179 xmax = np.max(unique(x))
 180 ymin = np.min(unique(y))
 181 ymax = np.max(unique(y))
 182
 183 # assumes that the step distance doesn't change over the course of time
 184 rat = xstep/ystep
 185 print "size of xstep:", xstep
 186 print "size of ystep:", ystep
 187 print "aspect ratio:",rat
 188
 189 print "different values of vx:"
 190 print unique(vx),"(",len(vx),"total)"
 191
 192 print "different values of vy:"
 193 print unique(vy),"(",len(vy),"total)"
 194
 195 # needed for assign_color2
 196 color_list = [
 197                 'black',        # 0     everything else
 198                 'gray',         # 1     v = 0
 199                 'blue',         # 2     v = -1.0
 200                 'green',        # 3     v = -0.5
 201                 'orange',       # 4     v = 0.5
 202                 'red',          # 5     v = 1.0
 203                 'khaki',        # 6     v = -0.5 && v = 0.5
 204                 'purple',       # 7     v = -1.0 && v = 1.0
 205                 'salmon',       # 8     v > 0
 206                 'palegreen',    # 9     v < 0
 207                 'white',        # 10    not understood
 208                 'dimgray',      # 11    grey and black
 209                 'darkgoldenrod',# 12    v = +/- 0.5 & 0
 210                 'mediumpurple', # 13    v = +/- 1 1 & -0
 211                 'brown',        # 14    v = +/1 1 & +/- 0.5
 212                 'darkslategray',# 15    black and grey
 213         ]
 214
 215 label_list = [
 216                 r'other rational $v$',  # 0 black
 217                 r'$v = 0$',                     # 1 grey
 218                 r'$v = -1$',                    # 2 blue
 219                 r'$v = -0.5$',                  # 3 green
 220                 r'$v = 0.5$',                   # 4 orange
 221                 r'$v = 1.0$',                   # 5 red
 222                 r'$v = \pm 0.5$',               # 6 brown
 223                 r'$v = \pm 1.0$',               # 7 purple
 224                 r'$v > 0$',                     # 8 magenta
 225                 r'$v < 0$',                     # 9 cyan
 226                 r'???',                         # 10 white
 227                 r'small',                       # 11 darkgrey
 228                 r'$v = \pm 0.5, 0$',            # 12 darkgoldenrod
 229                 r'$v = \pm 1.0, 0$',            # 13 mediumpurple
 230                 r'$v = \pm 1.0,\pm 0.5$',       # 14 brown
 231                 r'bigger',                      # 15 darkslategray
 232         ]
 233
 234 xpadmin = xmin-xstep
 235 xpadmax = xmax+xstep
 236 ypadmin = ymin-ystep
 237 ypadmax = ymax+ystep
 238 xlim = (xpadmin, xpadmax)               # small padding between plot and
 239 ylim = (ypadmin, ypadmax)               # axis so they don't overlap
 240 xticks = np.arange(xmin, xmax+xpad, xtfreq)
 241 yticks = np.arange(ymin, ymax+ypad, ytfreq)
 242
 243
 244 fig = plt.figure(figsize=res)
 245 #fig.suptitle(title, size=labelsize)
 246
 247 max = 0.84
 248 c = max/2
 249 rect1 = [0.1, 0.1, c, 0.9]
 250 rect2 = [c+0.1, 0.1, c, 0.9]
 251
 252 ax1 = fig.add_axes(rect1, aspect=rat)
 253 ax2 = fig.add_axes(rect2, aspect=rat)
 254
 255 if (method == 'mean'):
 256
 257         print "assigning colors...."
 258         xcolors_x,xcolors_y = assign_color(x_color,x,y)
 259         ycolors_x,ycolors_y = assign_color(y_color,x,y)
 260         print "done."
 261
 262         print "plotting..."
 263         for k in range(len(unique(x_color))):
 264                 x_xtli = xcolors_x[k]
 265                 x_ytli = xcolors_y[k]
 266                 labelx = assign_label(unique(x_color)[k])
 267                 ax1.scatter(x_xtli, x_ytli, marker='s', lw=0, s=ms, alpha=0.5,
 268                                 c = unique(x_color)[k], label=labelx,
 269                                 edgecolor='face')
 270
 271         for k in range(len(unique(y_color))):
 272                 y_xtli = ycolors_x[k]
 273                 y_ytli = ycolors_y[k]
 274                 labely = assign_label(unique(y_color)[k])
 275                 ax2.scatter(y_xtli, y_ytli, marker='s', lw=0, s=ms, alpha=0.5,
 276                                 c = unique(y_color)[k], label=labely,
 277                                 edgecolor='none')
 278
 279         print "done."
 280
 281 elif (method == 'direct'):
 282
 283         print "assigning colors..."
 284         xcolors_x,xcolors_y = assign_color2(x_color,x,y,vx,color_list)
 285         ycolors_x,ycolors_y = assign_color2(y_color,x,y,vy,color_list)
 286         print "done."
 287
 288         #for i in range(8):
 289         #       print i, xcolors_x[i],xcolors_y[i];
 290
 291         print "plotting..."
 292         for k in range(len(xcolors_x)):
 293                 if len(xcolors_x[k]) > 0:
 294                         x_xtli = xcolors_x[k]
 295                         x_ytli = xcolors_y[k]
 296                         ax1.scatter(x_xtli, x_ytli, marker='s', lw=0,
 297                                         s=ms, alpha=1, c=color_list[k],
 298                                         label=label_list[k])
 299
 300         for k in range(len(ycolors_x)):
 301                 if len(ycolors_x[k]) > 0:
 302                         y_xtli = ycolors_x[k]
 303                         y_ytli = ycolors_y[k]
 304                         ax2.scatter(y_xtli, y_ytli, marker='s', lw=0,
 305                                         s=ms, alpha=1, c=color_list[k],
 306                                         label=label_list[k])
 307
 308         print "done."
 309 else:
 310         print "fatal error: specify a valid method"
 311
 312 leg1 = ax1.legend(bbox_to_anchor=(0.05, 2), loc=2, shadow=True,
 313                 fancybox=True, scatterpoints=1, markerscale=1,
 314                 borderaxespad=0.)
 315 leg2 = ax2.legend(bbox_to_anchor=(0.85, 2), loc=2, shadow=True,
 316                 fancybox=True, scatterpoints=1, markerscale=1,
 317                 borderaxespad=0.)
 318 for t in leg1.get_texts():
 319         t.set_fontsize('small')
 320 for t in leg2.get_texts():
 321         t.set_fontsize('small')
 322
 323 ax1.set_title(r'$v_x$', size=labelsize)
 324 ax1.set_xlim(xlim)
 325 ax1.set_ylim(ylim)
 326 ax1.set_xlabel(xlabel, size=labelsize)
 327 ax1.set_ylabel(ylabel, size=labelsize)
 328 ax1.set_xticks(xticks, minor=False)
 329 ax1.set_yticks(yticks, minor=False)
 330 ax1.set_xticklabels(ax1.get_xticks(), size=ticksize)
 331 ax1.set_yticklabels(ax1.get_yticks(), size=ticksize)
 332
 333 ax2.set_title(r'$v_y$', size=labelsize)
 334 ax2.set_xlim(xlim)
 335 ax2.set_ylim(ylim)
 336 ax2.set_xlabel(xlabel, size=labelsize)
 337 ax2.set_ylabel('', size=labelsize)
 338 ax2.set_xticks(xticks, minor=False)
 339 ax2.set_yticks(yticks, minor=False)
 340 ax2.set_xticklabels(ax2.get_xticks(), size=ticksize)
 341 ax2.set_yticklabels('', size=ticksize)
 342
 343 leg1.get_frame().set_alpha(0.75)
 344 leg2.get_frame().set_alpha(0.75)
 345
 346 plt.subplots_adjust(hspace=0)
 347 print "flushing into a file..."
 348 plt.savefig(FOUT)
 349 print "done."
 350 #plt.show()
 351
 352
 353
 354