# Import the basic libraries

# ASE system
import ase
from ase import Atom, Atoms
from ase import io
from ase.lattice.spacegroup import crystal

# Spacegroup/symmetry library
from pyspglib import spglib

# The qe-util package
from qeutil import QuantumEspresso

# iPython utility function
from IPython.core.display import Image

# Configure qe-util for local execution of the Quantum Espresso on four processors
QuantumEspresso.pw_cmd='mpiexec -n 4 pw.x < pw.in > pw.out'

a=4.3596                                         # Lattice constant in Angstrom
cryst = crystal(['Si', 'C'],                     # Atoms in the crystal
                [(0, 0, 0), (0.25, 0.25, 0.25)], # Atomic positions (fractional coordinates)
                spacegroup=216,                  # International number of the spacegroup of the crystal
                cellpar=[a, a, a, 90, 90, 90])   # Unit cell (a, b, c, alpha, beta, gamma) in Angstrom, Degrees

# Write the image to disk file
ase.io.write('crystal.png',       # The file where the picture get stored
             cryst,               # The object holding the crystal definition
             format='png',        # Format of the file
             show_unit_cell=2,    # Draw the unit cell boundaries
             rotation='115y,15x', # Rotate the scene by 115deg around Y axis and 15deg around X axis
             scale=30)            # Scale of the picture

# Display the image
Image(filename='crystal.png')

print 'Space group:', spglib.get_spacegroup(cryst)

# Create a Quantum Espresso calculator for our work. 
# This object encapsulates all parameters of the calculation, 
# not the system we are investigating.
qe=QuantumEspresso(label='SiC',                    # Label for calculations
                       wdir='calc',                # Working directory
                       pseudo_dir='../../pspot',   # Directory with pseudopotentials
                       kpts=[8,8,8],               # K-space sampling for the SCF calculation
                       xc='pz',                    # Exchange functional type in the name of the pseudopotentials
                       pp_type='vbc',              # Variant of the pseudopotential
                       pp_format='UPF',            # Format of the pseudopotential files
                       ecutwfc=70,                 # Energy cut-off (in Rydberg)
                       use_symmetry=True)          # Use symmetry in the calculation ?

# Check where the calculation files will reside.
print qe.directory

# Assign the calculator to our system
cryst.set_calculator(qe)

# Run the calculation to get stress tensor (Voigt notation, in eV/A^3) and pressure (in kBar)
print "Stress tensor   (Voigt notation eV/A^3):", cryst.get_stress()
print "External pressure                (kBar):", cryst.get_isotropic_pressure(cryst.get_stress())*1e-3

# Do the same but get the results in GPa
# Note that this time we get the results immediately. 
# We did not change the system, so it is not necessary to repeat the calculation.
print "Stress tensor  (Voigt notation, GPa):", cryst.get_stress()/ase.units.GPa
print 
print "Stress tensor (Tensor notation, GPa):"
print cryst.get_stress(voigt=False)/ase.units.GPa
print
print "External pressure              (GPa):", cryst.get_isotropic_pressure(cryst.get_stress())*1e-4

# A sequential run for a series of lattice constants

# We will store the results in this list
result=[]

# Our prototype crystal is just a copy of the structure defined above
cr=Atoms(cryst)

# It needs a calculator as well. This may be the same calculator or we can define a separate one.
cr.set_calculator(qe)

print " Scale     A(A)     Energy(eV)   Pressure(GPa) "
print "==============================================="

# Iterate over scales between 98% and 102% of the starting unit cell size. 
# We use 11 points in the interval
for x in linspace(0.98,1.02,11):
    # Modify the crystal by scaling the lattice vectors
    cr.set_cell(cryst.get_cell()*x,scale_atoms=True)
    
    # Calculate energy and stress and store the results in the result list
    result.append([x, x*cryst.get_cell()[0,0], cr.get_potential_energy(), 1e-4*cr.get_isotropic_pressure(cr.get_stress())])
    
    # Print it as well
    print "% 5.03f   % 6.04f    %+6.4f      % +8.3f " % tuple(result[-1])

# Prepare the collected data for plotting
# This will make an array (matrix) out of a list and transpose it for easier access later
# Transposing the matrix means that we can specify just a column to get the whole column as a vector.
result=array(result).T

# Let us save our calculated data to a file. 
# To have data in columns we need to transpose the array again.
# This is a consequence of the row-column conventions and has no deeper meaning.
savetxt('e+p-vs-a.dat',result.T)

# Let us plot the results and save the figure
figsize(12,7)

# To make the plot nicer we define a shift to energy. 
# Rounding to second decimal digit in eV
E0=round(min(result[2])-(max(result[2])-min(result[2]))/20,2)

# Plot the result
plot(result[1],     # Arguments (x-axis)
     result[2]-E0,  # Values (y-axis)
     'o-',          # Symbol and line style
     label='Total internal energy')

legend()            # Add a legend

# Set the axes labels
xlabel('Lattice vector length ($\AA$)')  
ylabel('Energy %+8.2f (eV)' % (-E0))

# Store the figure
savefig('e-vs-a.pdf')

# Lets do the same with pressure. 
# But this time let us fit a Birch-Murnaghan equation of state to the data

# We need a fitting package from scipy
from scipy import optimize

# Define a B-M eos function
def BMEOS(v,v0,b0,b0p):
        return (b0/b0p)*(pow(v0/v,b0p) - 1)

# Define functions for fitting
# The B-M EOS is defined as a function of volume.
# Our data is a function of lattice parameter A^3=V
# We need to convert them on-the-fly
fitfunc = lambda p, x: [BMEOS(xv**3,p[0]**3,p[1],p[2]) for xv in x]
errfunc = lambda p, x, y: fitfunc(p, x) - y


figsize(12,7)
# Plot the data
plot(result[1],result[3],'+',markersize=10,markeredgewidth=2,label='Pressure')

# Fit the EOS

# Create a data array: lattice constant vs. isotropic pressure
ap=array([result[1],result[3]])

# Estimate the initial guess assuming b0p=1
# Limiting arguments
a1=min(ap[0])
a2=max(ap[0])
# The pressure is falling with the growing volume
p2=min(ap[1])
p1=max(ap[1])

# Estimate the slope
b0=(p1*a1-p2*a2)/(a2-a1)
a0=(a1)*(p1+b0)/b0

# Set the initial guess
p0=[a0,b0,1]

# Fitting
# fit will receive the fitted parameters, 
# and value of succ indicates if fitting was successful
fit, succ = optimize.leastsq(errfunc, p0[:], args=(ap[0],ap[1]))

# Ranges - the ordering in ap is not guaranteed at all!
# In fact it may be purely random.
x=numpy.array([min(ap[0]),max(ap[0])])
y=numpy.array([min(ap[1]),max(ap[1])])

# Plot the P(V) curves and points for the crystal

# Mark the center at P=0, A=A0 with dashed lines
axvline(fit[0],ls='--')
axhline(0,ls='--')

# Plot the fitted B-M EOS through the points,
# and put the fitting results on the figure.
xa=numpy.linspace(x[0],x[-1],20)
plot(xa,fitfunc(fit,xa),'-', 
     label="\nB-M fit:\n$A_0$=%6.4f $\AA$,\n$B_0$=%6.1f GPa,\n$B'_0$=%5.3f  " % (fit[0], fit[1], fit[2]) )

legend()
xlabel('Lattice vector length ($\AA$)')
ylabel('Pressure (GPa)')

# Save our figure
savefig('p-vs-a.pdf')