from pylab import *

figformat = '.png'
saveopts = {'dpi':300} #, 'transparent':True, 'frameon':True, 'bbox_inches':'tight'}
seterr(divide='ignore')

temp_canonical = 4.1
energy_microcanonical = -2.0
range_microcanonical = 1.0
micro_e0 = energy_microcanonical - 0.5*range_microcanonical
micro_e1 = energy_microcanonical + 0.5*range_microcanonical
def potential(x):
    return x**6 + 4*x**3 - 5*x**2 - 4*x
x = linspace(-2.5,2.5,2001) ; dx = x[1] - x[0]
mass = 1.0
p = linspace(-15,15,2001) ; dp = p[1] - p[0]
psextent = (x[0]-0.5*dx, x[-1]+0.5*dx, p[0]-0.5*dp, p[-1]+0.5*dp)

# compute pixel edges, used for pcolormesh.
xcorners = zeros(len(x)+1)
xcorners[:len(x)] = x-0.5*dx
xcorners[-1] = x[-1] + 0.5*dx

X,P = meshgrid(x, p)
E = potential(X) + P**2/(2*mass) #Hamiltonian

# make an energy range, for plots vs energy.
Evals = arange(-8,10,0.1)
phaseV = array(list(sum(E <= Elim) for Elim in Evals))
Evals2 = (Evals + 0.5*(Evals[1]-Evals[0]))[:-1]
phaseDOS = diff(phaseV)

# also figure out the density of states function in position-energy.
xvals = list()
phasesump = array(list(sum(E <= Elim,axis=0) for Elim in Evals))
phasedosp = diff(phasesump,axis=0)

#define color map that is transparent for low values, and dark blue for high values.
# weighted to show low probabilities well
cdic = {'red':   [(0,0,0),(1,0,0)],
        'green': [(0,0,0),(1,0,0)],
        'blue':  [(0,0.7,0.7),(1,0.7,0.7)],
        'alpha': [(0,0,0),
                  (0.1,0.4,0.4),
                  (0.2,0.6,0.6),
                  (0.4,0.8,0.8),
                  (0.6,0.9,0.9),
                  (1,1,1)]}
cm_prob = matplotlib.colors.LinearSegmentedColormap('prob',cdic)

def energyplot(phaseDOS_E, phaseDOS, phasedosp, ensemble, doslighten=1.0, ensemblelighten=1.0):
    """
    Plot the potential with density of states on sidebar.
      Evals, phaseDOS: list of energies and DOS to plot on right panel
    """
    fig = figure()
    
    # energy-position plot
    ax = axes([0.08,0.06,0.73,0.43])
    plot(x,potential(x), linewidth=2, color='r', zorder=1)
    extent = (xcorners[0], xcorners[-1], Evals[0], Evals[-1])
    img = imshow(phasedosp, cmap=cm_prob, extent=extent, interpolation='none', aspect='auto', origin='lower', zorder=0)
    clim(0,amax(phasedosp)*doslighten)
    ax.xaxis.labelpad = 2
    ax.yaxis.labelpad = -3
    xlabel("position $x$")
    ylabel("energy")
    ax.xaxis.set_ticklabels([])
    ax.yaxis.set_ticklabels([])
    ylim(-9,9)
    xlim(-2.1,1.7)
    ax.xaxis.set_ticks([-2,-1,0,1])

    # density of states sidebar
    ax = axes([0.83,0.06,0.14,0.43]) #, axisbg=(0.95,0.95,0.95))
    xlabel("states")
    ax.xaxis.set_ticks([])
    ax.yaxis.set_ticklabels([])
    ax.yaxis.set_ticks_position('right')
    ylim(-9,9)
    fill_betweenx(phaseDOS_E, phaseDOS, linewidth=0, color=(0.5,0.5,0.85))
    xlim(-0.05*max(phaseDOS),max(phaseDOS)*1.1)

    # phase space plot
    ax = axes([0.08,0.50,0.73,0.455])
    img = imshow(ensemble, cmap=cm_prob, extent=psextent, interpolation='none', aspect='auto', origin='lower', zorder=0)
    clim(0,amax(ensemble)*ensemblelighten)
    ax.xaxis.labelpad = 4
    ax.xaxis.set_label_position('top')
    ax.xaxis.set_ticklabels([])
    ax.yaxis.set_ticks([])
    ax.xaxis.set_ticks_position('both')
    ax.yaxis.labelpad = 0
    xlabel("position $x$")
    ylabel("momentum $p$")
    ylim(-7.5,7.5)
    xlim(-2.1,1.7)
    ax.xaxis.set_ticks([-2,-1,0,1])

    fig.set_size_inches(3,4.5)
    fig.patch.set_alpha(0)

allensemble = (E > -999.0)
#viewensemble = (E < 9.0)
energyplot(Evals2, phaseDOS,phasedosp,allensemble, doslighten=0.8, ensemblelighten=16.0)
savefig("class_potential"+figformat, **saveopts)

#canonical phase space image
canonical = exp(-E/temp_canonical)
print "canonical (T =",temp_canonical,") avg energy",
canonical_avgE = sum(E*canonical)/sum(canonical)
print canonical_avgE

energyplot(Evals2, phaseDOS*exp(-Evals2/temp_canonical),
    phasedosp*(exp(-Evals2/temp_canonical))[:,newaxis],
    canonical, doslighten=0.3)
sca(gcf().axes[0])
annotate("$\\langle E\\rangle$", (-0.5,canonical_avgE),
    textcoords=None,verticalalignment='top',color=(0,0.4,0))
axhline(canonical_avgE, linestyle='dotted', linewidth=1,color=(0,0.4,0))
annotate('',(1.2,7.-temp_canonical),(1.2,7.),
    arrowprops = {'arrowstyle':'<->'})
text(1.15,7.-0.5*temp_canonical,'$kT$',
    horizontalalignment='right',verticalalignment='center')
sca(gcf().axes[1])
axhline(canonical_avgE, linestyle='dotted', linewidth=1,color=(0,0.4,0))
savefig("class_canonical_potential"+figformat, **saveopts)


micro = (E < micro_e1)*(E > micro_e0)
print "microcanonical (E0 =",energy_microcanonical,", Delta =",0.5*range_microcanonical,") avg energy",
print sum(E*micro)/sum(micro)

tmp = (Evals2 < micro_e1)*(Evals2 > micro_e0)
energyplot(Evals2, phaseDOS*tmp,phasedosp*tmp[:,newaxis], micro, doslighten=0.5, ensemblelighten=3.0)
sca(gcf().axes[0])
axhspan(micro_e0, micro_e1, color=(0.7,1,0.7),zorder=-2)
sca(gcf().axes[1])
axhspan(micro_e0, micro_e1, color=(0.7,1,0.7),zorder=-2)
savefig("class_microcanonical_potential"+figformat, **saveopts)


# Position expectation values 
fig = figure()
plot(x, sum(micro,axis=0)/float(sum(micro))/dx, label='microcanonical')
plot(x, sum(canonical,axis=0)/sum(canonical)/dx, label='canonical')
xlim(-2.1,1.7)
fig.get_axes()[0].xaxis.set_ticks([-2,-1,0,1])
xlabel("position $x$")
ylabel("PDF of position $P(x)$")
legend()
fig.set_size_inches(4,4)
fig.patch.set_alpha(0)
savefig("class_position_pdf"+figformat, **saveopts)

# Momentum expectation values 
fig = figure()
plot(p, sum(micro,axis=1)/float(sum(micro))/dp, label='microcanonical')
plot(p, sum(canonical,axis=1)/sum(canonical)/dp, label='canonical')
xlim(-7.5,7.5)
xlabel("momentum $p$")
ylabel("PDF of momentum $P(p)$")
legend()
fig.set_size_inches(4,4)
fig.patch.set_alpha(0)
savefig("class_momentum_pdf"+figformat, **saveopts)