#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
Tracés de lignes de niveau ou isolignes
Application : Potentiel Energy Surface de la réaction
H + H2 --> H2 + H

"""
# ref : http://bulldog2.redlands.edu/facultyfolder/deweerd/tutorials/Tutorial-ContourPlot.pdf

import matplotlib.pyplot as plt  # directive d'importation standard de Matplotlib
import numpy as np               # directive d'importation standard de numpy
from mpl_toolkits.mplot3d import Axes3D  # Axes3D

def Ebond(rAB):
    return D_e * (np.exp(-2.*beta*(rAB-r_e)) - 2.*np.exp(-beta*(rAB-r_e)))
def Eant(rAB):
    return 0.5 * D_e * (np.exp(-2.*beta*(rAB-r_e)) + 2.*np.exp(-beta*(rAB-r_e)))
def Q(rAB):
    return 0.5 * ((1.+k)*Ebond(rAB) + (1.-k)*Eant(rAB))
def a(rAB):
    return 0.5 * ((1.+k)*Ebond(rAB) - (1.-k)*Eant(rAB))

beta=19.3E-3 # pm-1
r_e=74.1 # pm
D_e = .76 # E-18 J
k=0.18
rmin=10.
rmax=400.
num=100
x_1d = np.linspace(rmin,rmax, num)
print x_1d.shape, x_1d.dtype, x_1d.ndim
y_1d = np.linspace(rmin,rmax, num)
print y_1d.shape, y_1d.dtype, y_1d.ndim
X, Y = np.meshgrid(x_1d, y_1d)
print X.shape, X.dtype, X.ndim, Y.shape, Y.dtype, Y.ndim
E=(Q(X)+Q(Y)+Q(X+Y)-np.sqrt(2.*( (a(X)-a(Y) )**2.+(a(Y)-a(X+Y) )**2.+(a(X+Y)-a(X) )**2.) ))/(1.+k)
print np.min(E)  #valeur minimale de E

fig = plt.figure(figsize=(12, 12), dpi=80)
ax = fig.add_subplot(111)
# cf. http://stackoverflow.com/questions/7965743/how-can-i-set-the-aspect-ratio-in-matplotlib
ax.set_aspect("equal")
levels = np.linspace(-1.7, 1.0, 53)
CS1 = plt.contour(X, Y, E, levels, colors='k')
plt.clabel(CS1, colors = 'k', fmt = '%2.2f', fontsize=14)
CS2 = plt.contourf(X, Y, E, levels)
#plt.colorbar(CS2)  # visualisation éventuelle de l'échelle de couleur

plt.title('Isolignes')
plt.xlabel('x (pm)')
plt.ylabel('y (pm)')

fig = plt.figure(2,figsize=(15, 15) )
ax = Axes3D(fig)
ax.plot_surface(X,Y,E, rstride=1,cstride=1 ,cmap=plt.cm.jet)
ax.set_xlabel('X')
ax.set_ylabel('Y')
ax.set_zlabel('E')
plt.show()