# paste this code at the end of "hydrogen.py" 1.1
basename = 'porbital_'
nframes = 60
n, l = 2, 1

for i in range(nframes):
    t = (i + 10) % nframes / float(nframes)
    fname = basename + '{:02}.png'.format(i)
    print fname
    
    # slow down animation around m=1 at the ring shape
    if t < 0.5:
        phi = sin(t * pi / 2)**2
    else:
        phi = sin((1-t) * pi / 2)**2
    
    # rotate population between two eigenstates
    a1 = cos(phi * pi / 2)
    a2 = sin(phi * pi / 2)
    b1 = (a1 + a2) / sqrt(2)
    b2 = (a1 - a2) / sqrt(2)
    
    draw_orbital([[n, l, 1, a1], [n, l, -1, a2]], w=220, zoom=1.5,
                 fname=fname, density=heuristic_density(n, l, 1))
    
    if '{:.2f}'.format(a1) == '1.00' and '{:.2f}'.format(a2) == '0.00':
        psi1 = r'\    |2p1⟩            '
    else:
        psi1 = '{:.2f}|2p1⟩ + {:.2f}|2p-1⟩'.format(a1, a2)
    if '{:.2f}'.format(b1) == '1.00' and '{:.2f}'.format(b2) == '0.00':
        psi2 = r'\    |2px⟩             '
    else:
        psi2 = '{:.2f}|2px⟩ + {:.2f}i|2py⟩'.format(b1, b2)
    
    add_text = 'convert {} +antialias -pointsize 15 \
-font "DejaVu-Sans-Mono-Bold" -annotate +12+41 "{}" -annotate +12+192 "{}" \
-crop 220x180+0+20 +repage {}'.format(fname, psi1, psi2, fname)
    os.system(add_text)

assemble_gif = r'convert -delay 5 -dispose Background -dither FloydSteinberg \
-colors 256 {}??.png \( -clone 20 -set delay 200 \) -swap 20,-1 +delete \( \
-clone 50 -set delay 200 \) -swap 50,-1 +delete Orbital_p1-px_animation.gif'.format(basename)
os.system(assemble_gif)