#! / usr / bin / python # - * - кодировка: utf8 - * -from math import *import matplotlib.pyplot as pltfrom matplotlib import animation, colors, colorbarimport numpy as npimport colorsysfrom scipy.interpolate import interp1dimport os, sysplt.rc('path', snap=False)plt.rc('mathtext', default='regular')# image settingsfname = 'QHO-catstate-even3-animation-color'width, height = 300, 200ml, mr, mt, mb, mh, mc = 35, 19, 22, 45, 12, 6x0, x1 = -6.5, 6.5y0, y1 = 0.0, 1.2nframes = 150fps = 20# physics settingsalpha0 = 3.0omega = 2*pidef color(phase): hue = (phase / (2*pi) + 2./3.) % 1 light = interp1d([0, 1, 2, 3, 4, 5, 6], # adjust lightness [0.64, 0.5, 0.55, 0.48, 0.70, 0.57, 0.64])(6 * hue) hls = (hue, light, 1.0) # maximum saturation rgb = colorsys.hls_to_rgb(*hls) return rgbdef coherent(alpha, x, omega, t, l=1.0): # Definition of coherent states # https://en.wikipedia.org/wiki/Coherent_states psi = (pi*l**2)**-0.25 * np.exp( -0.5/l**2 * (x - sqrt(2)*l * alpha.real)**2 + 1j*sqrt(2)/l * alpha.imag * x + 0.5j * (alpha0**2*sin(2*omega*t) - omega*t)) return psidef animate(nframe): print str(nframe) + ' ',; sys.stdout.flush() t = float(nframe) / nframes * 0.5 # animation repeats after t=0.5 alpha = e ** (-1j * omega * t) * alpha0 ax.cla() ax.grid(True) ax.axis((x0, x1, y0, y1)) x = np.linspace(x0, x1, int(ceil(1+w_px))) x2 = x - px_w/2. # Definition of cat states in terms of coherent states: # https://en.wikipedia.org/wiki/Cat_state psi = coherent(alpha, x, omega, t) + coherent(-alpha, x, omega, t) psi /= sqrt(2 * (1 + exp(-2*alpha0**2))) # Let's cheat a bit: discard the constant phase from the zero-point energy! # This will reduce the period from T=2*(2pi/omega) to T=0.5*(2pi/omega) # and allow fewer frames and less file size for repetition. # For big alpha the change is hardly visible psi *= np.exp(0.5j * omega * t) y = np.abs(psi)**2 psi2 = coherent(alpha, x2, omega, t) + coherent(-alpha, x2, omega, t) psi2 *= np.exp(0.5j * omega * t) phi = np.angle(psi2) # plot color filling for x_, phi_, y_ in zip(x, phi, y): ax.plot([x_, x_], [0, y_], color=color(phi_), lw=2*0.72) ax.plot(x, y, lw=2, color='black') ax.set_yticks(ax.get_yticks()[:-1]) # create figure and axesplt.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 pixelspx_w = float(x1 - x0) / w_px # width of one pixel in plot units# axes labelsfig.text(0.5 + 0.5 * float(ml-mh-mc-mr)/width, 4./height, r'$x\ \ [(\hbar/(m\omega))^{1/2}]$', ha='center')fig.text(5./width, 1.0, '$|\psi|^2$', va='top')# colorbar for phasecax = 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 - 10./width, 1.0, '$arg(\psi)$', ha='right', va='top')plt.sca(ax)# start animationif 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)