www/png/example.py

   1 # this code is really a mess ...
   2
   3 import sys
   4 from Numeric import *
   5 from RandomArray import *
   6 from LinearAlgebra import *
   7 from FFT import *
   8
   9 from pyx import *
  10
  11
  12 sqrt2 = math.sqrt(2)
  13
  14
  15 def norm(v):
  16     return math.sqrt(reduce(lambda x, y: x+y*y, v, 0))
  17
  18 def blochvector(l, k):
  19     assert k >= 0 and k <= l/2
  20     v = cos(2*math.pi/l*(k)*arange(l))
  21     return v/norm(v)
  22
  23 def localvector(l, pos):
  24     v = zeros(l)
  25     v[pos] = 1
  26     return v/norm(v)
  27
  28 def local2vector(l, pos1, pos2):
  29     v = zeros(l)
  30     v[pos1] = 1
  31     v[pos2] = 1
  32     return v/norm(v)
  33
  34
  35
  36 def diag_anderson(l, w):
  37
  38     """diagonalize the anderson model"""
  39
  40     seed(x=1705, y=1111)
  41     pot=w*random(l)-0.5*w
  42     ham = zeros((l, l), Float)
  43     for i in xrange(l):
  44         ham[i, i] = pot[i]
  45         ham[i, (i+1)%l] = -1
  46         ham[i, (i-1)%l] = -1
  47     evalues, evectors = eigenvectors(ham)
  48     sorter = argsort(evalues)
  49     return take(evalues, sorter), take(evectors, sorter)
  50
  51
  52
  53 def diag_harper_full(l, lam):
  54
  55     """diagonalize the harper model by full diagonalization"""
  56
  57     def fibs(tonumber):
  58         result = [1, 1]
  59         while result[-1] < tonumber:
  60             result.append(result[-2]+result[-1])
  61         return result
  62
  63     l2, ltest = fibs(l)[-2:]
  64     if ltest != l: raise ValueError("l isn't a fibonacci number")
  65
  66     pot=lam*cos((2*pi*l2*arange(l))/l)
  67     ham = zeros((l, l), Float)
  68     for i in xrange(l):
  69         ham[i, i] = pot[i]
  70         ham[i, (i+1)%l] = 1
  71         ham[i, (i-1)%l] = 1
  72     evalues, evectors = eigenvectors(ham)
  73     sorter = argsort(evalues)
  74     return take(evalues, sorter), take(evectors, sorter)
  75
  76
  77
  78 def diag_harper(l, lam):
  79
  80     """diagonalize the harper model by diagonalizing the symmetric and antisymmetic parts"""
  81
  82     def fibs(tonumber):
  83         result = [1, 1]
  84         while result[-1] < tonumber:
  85             result.append(result[-2]+result[-1])
  86         return result
  87
  88     def createsymm(dest, source):
  89         l = len(dest)
  90         if l % 2:
  91             dest[0] = source[0]
  92             for i in xrange(1, (l+1)/2):
  93                 dest[i] = dest[l-i] = source[i] / sqrt2
  94         else:
  95             dest[0] = source[0]
  96             dest[l/2] = source[l/2]
  97             for i in xrange(1, l/2):
  98                 dest[i] = dest[l-i] = source[i] / sqrt2
  99
 100     def createasymm(dest, source):
 101         l = len(dest)
 102         if l % 2:
 103             dest[0] = 0
 104             for i in xrange(1, (l+1)/2):
 105                 dest[i] = source[i-1] / sqrt2
 106                 dest[l-i] = -source[i-1] / sqrt2
 107         else:
 108             dest[0] = dest[l/2] = 0
 109             for i in xrange(1, l/2):
 110                 dest[i] = -source[i-1] / sqrt2
 111                 dest[l-i] = source[i-1] / sqrt2
 112
 113     l2, ltest = fibs(l)[-2:]
 114     if ltest != l: raise ValueError("l isn't a fibonacci number")
 115
 116     # symmetric and antisymmetric diagonalization
 117     if l%2:
 118         # odd
 119         symmsize = (l-1)/2 + 1
 120         pot = lam*cos((2*pi*l2*arange(symmsize))/l)
 121         pot[symmsize - 1] += 1
 122         symmham = zeros((symmsize, symmsize), Float)
 123         for i in xrange(symmsize):
 124             symmham[i, i] = pot[i]
 125             if i > 1:
 126                 symmham[i, i-1] = symmham[i-1, i] = 1
 127             elif i:
 128                 symmham[i, i-1] = symmham[i-1, i] = sqrt2
 129     else:
 130         # even
 131         symmsize = l/2 + 1
 132         pot = lam*cos((2*pi*l2*arange(symmsize))/l)
 133         symmham = zeros((symmsize, symmsize), Float)
 134         for i in xrange(symmsize):
 135             symmham[i, i] = pot[i]
 136             if i > 1 and i < symmsize - 1:
 137                 symmham[i, i-1] = symmham[i-1, i] = 1
 138             elif i:
 139                 symmham[i, i-1] = symmham[i-1, i] = sqrt2
 140     symmevalues, symmevectors = eigenvectors(symmham)
 141
 142     if l%2:
 143         # odd
 144         asymmsize = (l-1)/2
 145         pot = lam*cos((2*pi*l2*(arange(asymmsize)+1))/l)
 146         pot[asymmsize-1] -= 1
 147         asymmham = zeros((asymmsize, asymmsize), Float)
 148         for i in xrange(asymmsize):
 149             asymmham[i, i] = pot[i]
 150             if i:
 151                 asymmham[i, i-1] = asymmham[i-1, i] = 1
 152     else:
 153         # even
 154         asymmsize = l/2 - 1
 155         pot = lam*cos((2*pi*l2*(arange(asymmsize)+1))/l)
 156         asymmham = zeros((asymmsize, asymmsize), Float)
 157         for i in xrange(asymmsize):
 158             asymmham[i, i] = pot[i]
 159             if i:
 160                 asymmham[i, i-1] = asymmham[i-1, i] = 1
 161     asymmevalues, asymmevectors = eigenvectors(asymmham)
 162
 163     # build complete solution
 164     symmsort = argsort(symmevalues)
 165     asymmsort = argsort(asymmevalues)
 166     j = k = 0
 167     evalues = zeros((l,), Float)
 168     evectors = zeros((l, l), Float)
 169     for i in xrange(l):
 170         if j != symmsize and (k == asymmsize or
 171                               symmevalues[symmsort[j]] < asymmevalues[asymmsort[k]]):
 172             evalues[i] = symmevalues[symmsort[j]]
 173             createsymm(evectors[i], symmevectors[symmsort[j]])
 174             j += 1
 175         else:
 176             evalues[i] = asymmevalues[asymmsort[k]]
 177             createasymm(evectors[i], asymmevectors[asymmsort[k]])
 178             k += 1
 179     return evalues, evectors
 180
 181
 182
 183 def test_harper(l, lam):
 184
 185     """compare diag_harper and diag_harper_full"""
 186
 187     res1 = diag_harper(l, lam)
 188     res2 = diag_harper_full(l, lam)
 189     for x1, x2 in zip(res1[0], res2[0]):
 190         assert math.fabs(x1-x2) < 1e-8, "eigenvalues differ"
 191     for v1, v2 in zip(res1[1], res2[1]):
 192         sum = 0
 193         for x1, x2 in zip(v1, v2):
 194             sum += x1 * x2
 195         if sum > 0:
 196             for x1, x2 in zip(v1, v2):
 197                 assert math.fabs(x1-x2) < 1e-8, "eigenvectors differ\n%s\n%s" % (v1, v2)
 198         else:
 199             for x1, x2 in zip(v1, v2):
 200                 assert math.fabs(x1+x2) < 1e-8, "eigenvectors differ\n%s\n%s" % (v1, -v2)
 201
 202 # test_harper(13, 5)
 203
 204
 205
 206 def wigner_slow(vector):
 207
 208     """create wigner function of a state (direct way)"""
 209
 210     l = len(vector)
 211     wig = zeros((l, l), Float)
 212
 213     # wigner function (direct way)
 214     for x0loop in xrange(l):
 215         x0 = x0loop
 216         for k0loop in xrange(l):
 217             if l%2:
 218                 k0 = (k0loop-0.5*l+0.5)*2*pi/l
 219             else:
 220                 k0 = (k0loop-0.5*l+1)*2*pi/l
 221             sum = 0j
 222             for yloop in xrange(l):
 223                 y = yloop - l/2 + 1 - l%2
 224                 v = lambda x: vector[(x+10*l)%l]
 225                 if y%2:
 226                     v1 = 0.5*(v(x0-(y-1)/2) + v(x0-(y+1)/2))
 227                     v2 = 0.5*(v(x0+(y-1)/2) + v(x0+(y+1)/2))
 228                 #    v1 = 0
 229                 #    v2 = 0
 230                 else:
 231                     v1 = v(x0-y/2)
 232                     v2 = v(x0+y/2)
 233                 sum += v1 * v2 * exp(1j*k0*y)
 234             assert abs(sum.imag) < 1e-10, "imaginary part should be zero"
 235             wig[x0loop, k0loop] = sum.real
 236
 237     return wig
 238
 239
 240
 241 def wigner(vector):
 242
 243     """create wigner function of a state (fft)"""
 244
 245     l = len(vector)
 246     wig = zeros((l, l), Float)
 247
 248     fftarray = zeros(l, Float)
 249     fftresult = zeros(l, Complex)
 250     wig = zeros((l, l), Float)
 251     for x0loop in xrange(l):
 252         x0 = x0loop
 253         for yloop in xrange(l):
 254             y = yloop - l/2 + 1 - l%2
 255             v = lambda x: vector[(x+10*l)%l]
 256             if y%2:
 257                 v1 = 0.5*(v(x0-(y-1)/2) + v(x0-(y+1)/2))
 258                 v2 = 0.5*(v(x0+(y-1)/2) + v(x0+(y+1)/2))
 259             #    v1 = 0
 260             #    v2 = 0
 261             else:
 262                 v1 = v(x0-y/2)
 263                 v2 = v(x0+y/2)
 264             fftarray[(y+10*l)%l] = v1 * v2
 265         fftresult = real_fft(fftarray)
 266         for k0loop in xrange(l):
 267             if l%2:
 268                 index = int(abs(k0loop-0.5*l+0.5))
 269             else:
 270                 index = int(abs(k0loop-0.5*l+1))
 271             wig[x0loop, k0loop] = fftresult[index].real
 272
 273     return wig
 274
 275
 276
 277 def test_wigner(vector):
 278
 279     """compare wigner_slow and wigner"""
 280
 281     res1 = wigner_slow(vector)
 282     res2 = wigner(vector)
 283     for v1, v2 in zip(res1, res2):
 284         for x1, x2 in zip(v1, v2):
 285             assert math.fabs(x1-x2) < 1e-8, "wigner function differs\n%s\n%s" % (res1, res2)
 286
 287 # test_wigner(diag_anderson(10, 1)[1][5])
 288 # test_wigner(diag_anderson(11, 1)[1][5])
 289
 290
 291 def husimi_from_wigner(wig):
 292
 293     """calculate the husimi function out of the wigner function"""
 294
 295     l = len(wig)
 296     sigma = sqrt(l/pi/4)
 297
 298     hus = zeros((l, l), Float)
 299     for x1 in xrange(l):
 300         for k1 in xrange(l):
 301             sys.stderr.write("wigner -> husimi (very slow) ...%6.2f%%\r" % ((100.0*(x1*l+k1))/l/l))
 302             hus[x1, k1] = 0
 303             for x2 in xrange(l):
 304                 for k2 in xrange(l):
 305                     dx = x1-x2
 306                     if dx < -l/2: dx += l
 307                     if dx > l/2: dx -= l
 308                     dk = k1-k2
 309                     if dk < -l/2: dk += l
 310                     if dk > l/2: dk -= l
 311                     dk *= 2*pi/l
 312                     hus[x1, k1] += 2.0/l/l * wig[x2, k2] * exp(-0.5*dx*dx/sigma/sigma-2*sigma*sigma*dk*dk)
 313     sys.stderr.write("wigner -> husimi (very slow) ... done.   \n")
 314
 315     return hus
 316
 317
 318
 319 def husimi_slow(vector):
 320
 321     l = len(vector)
 322     hus = zeros((l, l), Float)
 323     phases = zeros((l, l), Float)
 324     sigma=sqrt(l/pi/4)
 325     factor=1/sqrt(sqrt(2*pi)*sigma*l)
 326
 327     for x0loop in xrange(l):
 328         x0 = x0loop
 329         for k0loop in xrange(l):
 330             if l%2:
 331                 k0 = (k0loop-0.5*l+0.5)*2*pi/l
 332             else:
 333                 k0 = (k0loop-0.5*l+1)*2*pi/l
 334             sum = 0j
 335             for x in xrange(l):
 336                 phase = exp(1j*k0*x)
 337                 sum += vector[x] * factor * exp(-(x-x0-l)*(x-x0-l)/(4*sigma*sigma)) * phase
 338                 sum += vector[x] * factor * exp(-(x-x0  )*(x-x0  )/(4*sigma*sigma)) * phase
 339                 sum += vector[x] * factor * exp(-(x-x0+l)*(x-x0+l)/(4*sigma*sigma)) * phase
 340             hus[x0loop, k0loop] = abs(sum)*abs(sum)
 341             phases[x0loop, k0loop] = math.atan2(sum.imag, sum.real)
 342
 343     return hus, phases
 344
 345
 346
 347 def husimi(vector):
 348
 349     l = len(vector)
 350     hus = zeros((l, l), Float)
 351     sigma=sqrt(l/pi/4)
 352     factor=1/sqrt(sqrt(2*pi)*sigma*l)
 353
 354     fftarray = zeros(l, Complex)
 355     fftresult = zeros(l, Complex)
 356     heights = zeros((l, l), Float)
 357     phases = zeros((l, l), Float)
 358     for x0loop in xrange(l):
 359         x0 = x0loop
 360         for xloop in xrange(l):
 361             x = xloop
 362             while (x-x0 < -0.5*l): x += l
 363             while (x-x0 > 0.5*l): x -=l
 364             fftarray[xloop] = vector[xloop] * factor * exp(-(x-x0)*(x-x0)/(4*sigma*sigma))
 365         #fftresult = real_fft(fftarray)
 366         fftresult = fft(fftarray)
 367         for k0loop in xrange(l):
 368             if l%2:
 369                 index = (int(k0loop-0.5*l+0.5) + 10*l)%l
 370                 #index = int(abs(k0loop-0.5*l+0.5))
 371             else:
 372                 raise
 373                 index = int(abs(k0loop-0.5*l+1))
 374             heights[x0loop, k0loop] = abs(fftresult[index])*abs(fftresult[index])
 375             phases[x0loop, k0loop] = math.atan2(fftresult[index].imag, fftresult[index].real)
 376             # if 2*x0loop > l: phases[x0loop, k0loop] = -phases[x0loop, k0loop] + math.pi
 377             # if 2*k0loop > l: phases[x0loop, k0loop] = -phases[x0loop, k0loop]
 378
 379     return heights, phases
 380
 381
 382
 383 def test_husimi(vector):
 384
 385     """compare husimi_slow and husimi"""
 386
 387     res1 = husimi_slow(vector)
 388     res2 = husimi(vector)
 389     for v1, v2 in zip(res1, res2):
 390         for x1, x2 in zip(v1, v2):
 391             assert math.fabs(x1-x2) < 1e-8, "husimi function differs\n%s\n%s" % (res1, res2)
 392
 393 # test_husimi(diag_anderson(10, 1)[1][5])
 394 # test_husimi(diag_anderson(11, 1)[1][5])
 395
 396
 397
 398
 399 def plot_grid(grid, filename):
 400
 401     class griddata(datafile._datafile):
 402
 403         def __init__(self, grid):
 404             data = []
 405             for i in range(len(grid)):
 406                 for j in range(len(grid[i])):
 407                     if l%2:
 408                         k = (j-0.5*l+0.5)*2*pi/l
 409                     else:
 410                         k = (j-0.5*l+1)*2*pi/l
 411                     if l%2 or j < len(grid[i]) - 1:
 412                         data.append([grid[i, j], i+0.5, i+1.5, k-math.pi/l, k+math.pi/l])
 413                     else:
 414                         data.append([grid[i, j], i+0.5, i+1.5, -math.pi, -math.pi+math.pi/l])
 415                         data.append([grid[i, j], i+0.5, i+1.5, math.pi-math.pi/l, math.pi])
 416             datafile._datafile.__init__(self, ["color", "xmin", "xmax", "kmin", "kmax"], data)
 417
 418     hd = griddata(grid)
 419
 420     mypainter = graph.axispainter(baselineattrs=None, zerolineattrs=None, innerticklengths=0, outerticklengths="0.2")
 421
 422     c = canvas.canvas()
 423     t = c.insert(tex.tex())
 424     g = c.insert(graph.graphxy(t, width=10, height=10, x2=None, y2=None, backgroundattrs=canvas.stroked(),
 425                                x=graph.linaxis(min=0.5, max=len(grid)+0.5, title="$x$",
 426                                                part=graph.manualpart(ticks=("1", str(len(grid))), texts=("1", "$L$")),
 427                                                painter=mypainter),
 428                                y=graph.linaxis(min=-pi, max=pi, divisor=pi, suffix=r"\pi", title="$k$",
 429                                                part=graph.linpart("1"), painter=mypainter)))
 430     g.plot(graph.data(hd, color="color", xmin="xmin", xmax="xmax", ymin="kmin", ymax="kmax"),
 431            graph.rect(color.gradient.ReverseRainbow))
 432     g.dodata()
 433     g.finish()
 434     c.writetofile(filename)
 435
 436
 437
 438
 439 theta = 15 * math.pi/180
 440 phi = -25 * math.pi/180
 441
 442 allscale = 0.07
 443 distance = 150*allscale
 444 gridscale = 1*allscale
 445 zscale = 15*allscale
 446
 447
 448 class gridpoint:
 449
 450     def __init__(self, x_3d, y_3d, z_3d, x_2d, y_2d, depth, phase):
 451         self.x_3d = x_3d
 452         self.y_3d = y_3d
 453         self.z_3d = z_3d
 454         self.x_2d = x_2d
 455         self.y_2d = y_2d
 456         self.depth = depth
 457         self.phase = phase
 458
 459     def __repr__(self):
 460         return "(%s, %s, %s), (%s, %s), %s, %s " % (self.x_3d, self.y_3d, self.z_3d, self.x_2d, self.y_2d, self.depth, self.phase)
 461
 462 def norm(*list):
 463     f = 1/math.sqrt(reduce(lambda x, y: x + y*y, list, 0))
 464     return map(lambda x, f=f: x*f, list)
 465
 466
 467 a = (math.sin(phi), -math.cos(phi), 0)
 468 b = (-math.cos(phi)*math.sin(theta), -math.sin(phi)*math.sin(theta), math.cos(theta))
 469 eye = (distance*math.cos(phi)*math.cos(theta), distance*math.sin(phi)*math.cos(theta), distance*math.sin(theta))
 470
 471 class grid:
 472
 473     def point_3d(self, x, y, z, phase=0):
 474         d0 = float(a[0]*b[1]*(z-eye[2]) + a[2]*b[0]*(y-eye[1]) + a[1]*b[2]*(x-eye[0])
 475                   -a[2]*b[1]*(x-eye[0]) - a[0]*b[2]*(y-eye[1]) - a[1]*b[0]*(z-eye[2]))
 476         da = (eye[0]*b[1]*(z-eye[2]) + eye[2]*b[0]*(y-eye[1]) + eye[1]*b[2]*(x-eye[0])
 477              -eye[2]*b[1]*(x-eye[0]) - eye[0]*b[2]*(y-eye[1]) - eye[1]*b[0]*(z-eye[2]))
 478         db = (a[0]*eye[1]*(z-eye[2]) + a[2]*eye[0]*(y-eye[1]) + a[1]*eye[2]*(x-eye[0])
 479              -a[2]*eye[1]*(x-eye[0]) - a[0]*eye[2]*(y-eye[1]) - a[1]*eye[0]*(z-eye[2]))
 480         dc = (a[0]*b[1]*eye[2] + a[2]*b[0]*eye[1] + a[1]*b[2]*eye[0]
 481              -a[2]*b[1]*eye[0] - a[0]*b[2]*eye[1] - a[1]*b[0]*eye[2])
 482         return gridpoint(x, y, z, da/d0, db/d0, -dc/d0, phase)
 483
 484     def point_grid(self, x, y, z, phase):
 485         x, y, z = x - planecenter[0], y - planecenter[1], z - planecenter[2]
 486         return self.point_3d(x * gridscale, y * gridscale, z * zscale, phase)
 487
 488
 489 l=101
 490 values, vectors = diag_anderson(l, 0.5)
 491 vector = vectors[24]
 492 heights, phases = husimi(vector)
 493 planecenter = (0.5*(len(heights) - 1), 0.5*(len(heights[0]) - 1), 0)
 494 datagrid = [[None for j in range(len(heights[i]))] for i in range(len(heights))]
 495
 496 mygrid = grid()
 497
 498 for i in range(len(heights)):
 499     for j in range(len(heights[i])):
 500         datagrid[i][j] = mygrid.point_grid(i, j, heights[i][j]*1000, phases[i][j])
 501
 502 def triangle(p1, p2, p3, hue):
 503     ax, ay, az = p2.x_3d - p1.x_3d, p2.y_3d - p1.y_3d, p2.z_3d - p1.z_3d
 504     bx, by, bz = p3.x_3d - p1.x_3d, p3.y_3d - p1.y_3d, p3.z_3d - p1.z_3d
 505     axb = norm(ay*bz - az*by, az*bx - ax*bz, ax*by-ay*bx)
 506     e = norm(eye[0]-p1.x_3d, eye[1]-p1.y_3d, eye[2]-p1.z_3d)
 507     axb_r = axb[0]*e[0] + axb[1]*e[1] + axb[2]*e[2]
 508     depthpoint = mygrid.point_3d((p1.x_3d + p2.x_3d + p3.x_3d)/3.0,
 509                                  (p1.y_3d + p2.y_3d + p3.y_3d)/3.0, 0)
 510     return [depthpoint.depth,
 511             axb_r,
 512             path.path(path.moveto(p1.x_2d, p1.y_2d),
 513                       path.lineto(p2.x_2d, p2.y_2d),
 514                       path.lineto(p3.x_2d, p3.y_2d),
 515                       path.closepath()),
 516             hue]
 517
 518 def gethue(phase1, phase2):
 519     x1 = 0.5*phase1/math.pi
 520     x2 = 0.5*phase2/math.pi
 521     while x1-x2 > 0.5: x1 -= 1
 522     while x1-x2 < -0.5: x1 += 1
 523     dx = 0.5*(x1+x2)
 524     while dx > 1: dx -= 1
 525     while dx < 0: dx += 1
 526     return dx
 527
 528 triangles = []
 529 for i in range(len(heights)-1):
 530     for j in range(len(heights[i])-1):
 531         midpoint = mygrid.point_3d(0.25 * (datagrid[i][j].x_3d + datagrid[i+1][j].x_3d + datagrid[i][j+1].x_3d + datagrid[i+1][j+1].x_3d),
 532                                    0.25 * (datagrid[i][j].y_3d + datagrid[i+1][j].y_3d + datagrid[i][j+1].y_3d + datagrid[i+1][j+1].y_3d),
 533                                    0.25 * (datagrid[i][j].z_3d + datagrid[i+1][j].z_3d + datagrid[i][j+1].z_3d + datagrid[i+1][j+1].z_3d))
 534         triangles.append(triangle(datagrid[i][j], datagrid[i+1][j], midpoint, gethue(datagrid[i][j].phase, datagrid[i+1][j].phase)))
 535         triangles.append(triangle(datagrid[i+1][j], datagrid[i+1][j+1], midpoint, gethue(datagrid[i+1][j].phase, datagrid[i+1][j+1].phase)))
 536         triangles.append(triangle(datagrid[i+1][j+1], datagrid[i][j+1], midpoint, gethue(datagrid[i+1][j+1].phase, datagrid[i][j+1].phase)))
 537         triangles.append(triangle(datagrid[i][j+1], datagrid[i][j], midpoint, gethue(datagrid[i][j+1].phase, datagrid[i][j].phase)))
 538
 539
 540 triangles.sort()
 541
 542 c = canvas.canvas()
 543
 544 colors = [t[1] for t in triangles if t[1] > 0]
 545 mincol = min(colors)
 546 maxcol = max(colors)
 547 for triangle in triangles:
 548     if triangle[1] < 0:
 549         c.fill(triangle[2], [color.gray.black])
 550     else:
 551         usecolor = color.hsb(triangle[3], 1, 0.25 + 0.65 * ((triangle[1] - mincol)/(maxcol - mincol)) ** 0.5 )
 552         c.fill(triangle[2], [usecolor])
 553
 554 c.writetofile("example")
 555