#!/usr/bin/python
# -*- coding: utf8 -*-
import os
import sys
import matplotlib.pyplot as plt
from matplotlib import ticker
from matplotlib import animation, colors, colorbar
import numpy as np
import colorsys
from scipy.interpolate import interp1d
from math import *
plt.rc('path', snap=False)
plt.rc('text', usetex=False)
plt.rc('mathtext', default='regular')
plt.rc('font', size=11)
plt.rcParams['font.sans-serif'] = ['DejaVu Sans']
# image settings
fname = 'Quantum-rigid-rotator_1+2-animation-color'
width, height = 300, 200
ml, mr, mt, mb, mh, mc = 35, 20, 22, 45, 12, 6
x0, x1 = -pi, pi
y0, y1 = 0.0, 0.4
nframes = 80
fps = 20
a = {1:1, 2:1}
def color(phase):
phase1 = ((phase / (2*pi)) % 1 + 1) % 1
hue = (interp1d([0, 1./3, 1.2/3, 0.5, 1], # spread yellow a bit
[0, 1./3, 1.3/3, 0.5, 1])(phase1) + 2./3.) % 1
light = interp1d([0, 1, 2, 3, 4, 5, 6], # adjust lightness
[0.64, 0.5, 0.56, 0.48, 0.75, 0.57, 0.64])(6 * hue)
hls = (hue, light, 1.0) # maximum saturation
rgb = colorsys.hls_to_rgb(*hls)
return rgb
def qrr_wavefunc(a, phi, t):
# Wavefunction of a quantum rigid rotator
# https://dx.doi.org/10.1119/1.17340
ai = [i for i in a.iterkeys()]
s = sqrt(np.sum(np.abs([v for v in a.itervalues()])**2))
av = [a[i] / s for i in ai]
psi = 0j * phi
for n, aa in zip(ai, av):
omega = 2*pi * n**2
psi += aa * np.exp(1j*(n*phi - omega*t)) / sqrt(2*pi)
return psi
def animate(nframe):
print str(nframe) + ' ', ; sys.stdout.flush()
t = float(nframe) / nframes
ax.cla()
ax.grid(True)
ax.axis((x0, x1, y0, y1))
x = np.linspace(-pi, pi, int(ceil(1+w_px)))
x2 = x - px_w/2.
psi = qrr_wavefunc(a, x, t)
psi2 = qrr_wavefunc(a, x2, t)
y = np.abs(psi)**2
phase = np.angle(psi2)
# plot color filling
for x_, phase_, y_ in zip(x, phase, y):
ax.plot([x_, x_], [0, y_], color=color(phase_), lw=2*0.72)
ax.plot(x, y, lw=2, color='black')
ax.yaxis.set_ticks(np.arange(0, 0.4, 0.1))
# create figure and axes
plt.close('all')
fig, ax = plt.subplots(1, figsize=(width/100., height/100.))
bounds = [float(ml)/width, float(mb)/height,
1.0 - float(mr+mc+mh)/width, 1.0 - float(mt)/height]
fig.subplots_adjust(left=bounds[0], bottom=bounds[1],
right=bounds[2], top=bounds[3], hspace=0)
w_px = width - (ml+mr+mh+mc) # plot width in pixels
px_w = float(x1 - x0) / w_px # width of one pixel in plot units
# axes labels
fig.text(0.5 + 0.5 * float(ml-mh-mc-mr)/width, 5./height,
r'$\phi$', ha='center')
fig.text(5./width, 1.0, '$|\psi|^2$', va='top')
# colorbar for phase
cax = fig.add_axes([1.0 - float(mr+mc)/width, float(mb)/height,
float(mc)/width, 1.0 - float(mb+mt)/height])
cax.yaxis.set_tick_params(length=2)
cmap = colors.ListedColormap([color(phase) for phase in
np.linspace(0, 2*pi, 384, endpoint=False)])
norm = colors.Normalize(0, 2*pi)
cbar = colorbar.ColorbarBase(cax, cmap=cmap, norm=norm,
orientation='vertical', ticks=np.linspace(0, 2*pi, 3))
cax.set_yticklabels(['$0$', r'$\pi$', r'$2\pi$'], rotation=90)
fig.text(1.0 - 12./width, 1.0, '$arg(\psi)$', ha='right', va='top')
plt.sca(ax)
# start animation
if 0 != os.system('convert -version > ' + os.devnull):
print 'imagemagick not installed!'
# warning: imagemagick produces somewhat jagged and therefore large gifs
anim = animation.FuncAnimation(fig, animate, frames=nframes)
anim.save(fname + '.gif', writer='imagemagick', fps=fps)
else:
# unfortunately the matplotlib imagemagick backend does not support
# options which are necessary to generate high quality output without
# framewise color palettes. Therefore save all frames and convert then.
if not os.path.isdir(fname):
os.mkdir(fname)
fnames = []
for frame in range(nframes):
animate(frame)
imgname = os.path.join(fname, fname + '{:03d}'.format(frame) + '.png')
fig.savefig(imgname)
fnames.append(imgname)
# compile optimized animation with ImageMagick
cmd = 'convert -loop 0 -delay ' + str(100 / fps) + ' '
cmd += ' '.join(fnames) # now create optimized palette from all frames
cmd += r' \( -clone 0--1 \( -clone 0--1 -fill black -colorize 100% \) '
cmd += '-append +dither -colors 255 -unique-colors '
cmd += '-write mpr:colormap +delete \) +dither -map mpr:colormap '
cmd += '-alpha activate -layers OptimizeTransparency '
cmd += fname + '.gif'
os.system(cmd)
for fnamei in fnames:
os.remove(fnamei)
os.rmdir(fname)