#!/usr/bin/python
# -*- coding: utf8 -*-
from math import *
import matplotlib.pyplot as plt
import numpy as np
from numpy.polynomial.hermite import Hermite
plt.rc('path', snap=False)
plt.rc('mathtext', default='regular')
# image settings
fname = 'QHO-Fockstate1'
width, height = 300, 200
ml, mr, mt, mb = 35, 8, 22, 45
x0, x1 = -4, 4
y0, y1 = 0.0, 0.6
# physics settings
nfock = 1
def plot():
ax.cla()
ax.axis((x0, x1, y0, y1))
ax.grid(True)
psi_fock = np.eye(1, nfock+1, nfock).flatten()
# Definition of Fock-states in terms of Hermite functions:
# https://en.wikipedia.org/wiki/Quantum_harmonic_oscillator
a_hermite = [psi_fock[n] * pi**-0.25 / sqrt(2.**n*factorial(n))
for n in range(1+nfock)]
# doc: http://docs.scipy.org/doc/numpy/reference/generated/numpy.polynomial.hermite.Hermite.html
H = Hermite(a_hermite)
x = np.linspace(x0, x1, 2 * width)
psi_x = np.exp(-x**2 / 2.0) * H(x)
y = np.abs(psi_x)**2
plt.plot(x, y, lw=2, color='#0000cc')
ax.set_yticks(ax.get_yticks()[:-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)/width, 1.0 - float(mt)/height]
fig.subplots_adjust(left=bounds[0], bottom=bounds[1],
right=bounds[2], top=bounds[3], hspace=0)
# axes labels
fig.text(0.5 + 0.5 * float(ml-mr)/width, 4./height,
r'$x\ \ [(\hbar/(m\omega))^{1/2}]$', ha='center')
fig.text(5./width, 1.0, '$|\psi|^2$', va='top')
plot()
fig.savefig(fname + '.png')
# compress with optipng
from os import system
cmd = 'optipng -o9 "' + fname + '.png"'
if system(cmd) != 0:
print 'warning: optipng not found!'