#!/usr/bin/env python

'''
an executable script that operates on Dacapo generated
netcdf files and fits volumetric equations of state to the data.
in that case a typical usage could be:

eos.py -m -s myfit.png vol_*.nc

this would perform a Murnaghan EOS fit to the volumes and energies from
all netcdf files in the current directory that match the pattern
vol_*.nc and save a plot of the data in the file myfit.png.

See the docstrings in this file for more details of supported EOS and
graphics files.


If no ncfiles are specified it operates on all the ncfiles in the
current directory.

eos.py --all -q
will fit all the equations of state available, but not make any graphs.
Typical output looks like:
V0(A^3)       B(Gpa)   E0 (eV)    pressure(GPa) Murnaghan
 63.10        143.72     -4412.4899       -0.00
chisq = 0.0000

kitchin@thlc000 AgPd3/murn_500pw_12k > eos.py --all
V0(A^3)       B(Gpa)   E0 (eV)    pressure(GPa) Murnaghan
 63.10        143.72     -4412.4899       -0.00
chisq = 0.0000

V0(A^3)       B(Gpa)   E0 (eV)    pressure(GPa) Birch
 63.10        144.39     -4412.4900       -0.00
chisq = 0.0000

V0(A^3)       B(Gpa)   E0 (eV)    pressure(GPa) BirchMurnaghan
 63.10        143.32     -4412.4907       -0.00
chisq = 0.0003

V0(A^3)       B(Gpa)   E0 (eV)    pressure(GPa) Taylor
 63.10        151.37     -4412.4929       -0.00
chisq = 0.0000

V0(A^3)       B(Gpa)   E0 (eV)    pressure(GPa) PourierTarantola
 63.10        145.56     -4412.4902       -0.00
chisq = 0.0000

V0(A^3)       B(Gpa)   E0 (eV)    pressure(GPa) Vinet
 63.10        144.88     -4412.4901       -0.00
chisq = 0.0000

V0(A^3)       B(Gpa)   E0 (eV)    pressure(GPa) AntonSchmidt
 63.16        142.50     -4412.4851       -0.00
chisq = 0.0002


To get the Vinet and Birch-Murnaghan fits, no plots:

eos.py -v --bm -q

I would encourage you to cite the papers for the equations of state used
in ASE.Utilities.EquationOfState if you publish results from this module.
        
John Kitchin <jkitchin@andrew.cmu.edu>
05/20/05
'''

import glob,sys
from ASE.Utilities.EquationOfState import EquationOfState
from Dacapo import Dacapo
from optparse import OptionParser

parser = OptionParser(usage='eos.py [-m -b --bm -a -t -v -p --all -s [filename]] [ncfiles]',
                      version='0.1')

parser.add_option('-m',
                  nargs=0,
                  help = 'Murnaghan EOS')

parser.add_option('--bm',
                  nargs=0,
                  help = 'Birch-Murnaghan EOS')

parser.add_option('-b',
                  nargs=0,
                  help = 'Birch EOS')

parser.add_option('-a',
                  nargs=0,
                  help = 'AntonSchmidt EOS')

parser.add_option('-t',
                  nargs=0,
                  help = 'Taylor series expansion EOS')

parser.add_option('-v',
                  nargs=0,
                  help = 'Vinet EOS')

parser.add_option('-p',
                  nargs=0,
                  help = 'PourierTarantola EOS')

parser.add_option('--all',
                  nargs=0,
                  help = 'run all EOS')

parser.add_option('-s',
                  nargs=1,
                  help = 'save figure as filename')

parser.add_option('-q',
                  nargs=0,
                  help = 'quiet, no plot')

options,args = parser.parse_args()

if args == []:
    files = glob.glob('*.nc')
else:
    files = args

volumes = []
energies = []

for file in files:
    atoms = Dacapo.ReadAtoms(file)

    energy = atoms.GetPotentialEnergy()
    if energy is not None:
        energies.append(energy)
        volume = atoms.GetUnitCellVolume()
        volumes.append(volume)

if options.all is not None:
    for e in ['Murnaghan',
              'Birch',
              'BirchMurnaghan',
              'Taylor',
              'PourierTarantola',
              'Vinet',
              'AntonSchmidt']:
        x = EquationOfState(e,volumes,energies)
        print x
        if options.q is None:
            x.GetPlot()
        if options.s is not None:
            x.SavePlot(options.s)
        print 
    sys.exit()

eos = []
if options.m is not None:
    eos.append('Murnaghan')

if options.b is not None:
    oos.append('Birch')

if options.bm is not None:
    eos.append('BirchMurnaghan')

if options.t is not None:
    eos.append('Taylor')

if options.p is not None:
    eos.append('PourierTarantola')

if options.v is not None:
    eos.append('Vinet')

if options.a is not None:
    eos.append('AntonSchmidt')

#do murnaghan by default
if eos == []:
    eos = ['Murnaghan']
for e in eos:
    x = EquationOfState(e,volumes,energies)
    print x
    if options.q is None:
        x.GetPlot()
    if options.s is not None:
        x.SavePlot(options.s)
    print #an extra space
