phonopy/scripts/propplot

#!/usr/bin/env python

# Copyright (C) 2011 Atsushi Togo
# All rights reserved.
#
# This file is part of phonopy.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions
# are met:
#
# * Redistributions of source code must retain the above copyright
#   notice, this list of conditions and the following disclaimer.
#
# * Redistributions in binary form must reproduce the above copyright
#   notice, this list of conditions and the following disclaimer in
#   the documentation and/or other materials provided with the
#   distribution.
#
# * Neither the name of the phonopy project nor the names of its
#   contributors may be used to endorse or promote products derived
#   from this software without specific prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
# "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
# LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS
# FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE
# COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT,
# INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING,
# BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
# LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT
# LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN
# ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
# POSSIBILITY OF SUCH DAMAGE.

# Thermal properties plot (propplot)
#
# Usage:
#   propplot

import numpy as np

try:
    import yaml
except ImportError:
    print "You need to install python-yaml."
    exit(1)
    
try:
    from yaml import CLoader as Loader
    from yaml import CDumper as Dumper
except ImportError:
    from yaml import Loader, Dumper

# Parse options
from optparse import OptionParser
parser = OptionParser()
parser.set_defaults( output_filename = None,
                     factor = 1.0,
                     is_heat_capacity = False,
                     is_entropy = False,
                     is_free_energy = False,
                     is_diff = False,
                     is_gnuplot = False,
                     ymax = None,
                     ymin = None,
                     tmax = None,
                     tmin = None )
parser.add_option("--factor", dest="factor", type="float",
                  help="Conversion factor of energy unit to internal electronic energy")
parser.add_option("--ymax", dest="ymax", type="float",
                  help="Maximum value of y axis")
parser.add_option("--ymin", dest="ymin", type="float",
                  help="Minimum value of y axis")
parser.add_option("--tmax", dest="tmax", type="float",
                  help="Maximum value of temperature")
parser.add_option("--tmin", dest="tmin", type="float",
                  help="Minimum value of temperature")
parser.add_option("--cv", "--heat_capacity", dest="is_heat_capacity",
                  action="store_true",
                  help="Plot heat capacity")
parser.add_option("-s", "--entropy", dest="is_entropy",
                  action="store_true",
                  help="Plot entropy")
parser.add_option("--fe", "--free_energy", dest="is_free_energy",
                  action="store_true",
                  help="Plot free energy")
parser.add_option("-o", "--output", dest="output_filename",
                  action="store", type="string",
                  help="Output filename")
parser.add_option("-d", "--diff", dest="is_diff",
                  action="store_true",
                  help="Calculate difference")
parser.add_option("--gnuplot", dest="is_gnuplot",
                  action="store_true",
                  help="Damp thermal properties in gnuplot data style")
(options, args) = parser.parse_args()

if options.is_gnuplot:
    print "# temperature[K] Free energy [kJ/mol] Entropy [J/K/mol] Heat capacity [J/K/mol] Energy [kJ/mol]"
else:
    import matplotlib.pyplot as plt
    if ( ( not options.is_heat_capacity ) and 
         ( not options.is_entropy ) and 
         ( not options.is_free_energy ) ):
        print "Set --cv, --entropy or --fe"
        exit(0)

if len( args ) == 0:
    filenames = ['thermal_properties.yaml']
else:
    filenames = args

if options.is_heat_capacity:
    prop_target = 'heat_capacity'
elif options.is_entropy:
    prop_target = 'entropy'
elif options.is_free_energy:
    prop_target = 'free_energy'

thermal_properties_0 = yaml.load( open( filenames[0] ).read(),
                                  Loader=Loader )['thermal_properties']
temperatures = [ v['temperature'] for v in thermal_properties_0 ]

tmin_index = 0
tmax_index = len( temperatures )
if not options.tmin == None:
    for i, t in enumerate( temperatures ):
        if t > options.tmin - ( temperatures[1] - temperatures[0] ) * 0.1:
            tmin_index = i
            break

if not options.tmax == None:
    for i, t in enumerate( temperatures ):
        if t > options.tmax + ( temperatures[1] - temperatures[0] ) * 0.1:
            tmax_index = i
            break

if options.is_diff:
    props_0 = [ v[prop_target] for v in thermal_properties_0 ]
    props_0 = np.array( props_0 ) * options.factor

for filename in filenames:
    thermal_properties = yaml.load( open( filename ).read(),
                                    Loader=Loader )['thermal_properties']

    if options.is_gnuplot:
        props = []
        for name in ('temperature', 'free_energy', 'entropy', 'heat_capacity'):
            props.append([ v[name] for v in thermal_properties ])
            
        for t, f, e, h in zip( props[0][tmin_index:tmax_index],
                               props[1][tmin_index:tmax_index],
                               props[2][tmin_index:tmax_index],
                               props[3][tmin_index:tmax_index] ):
            print t, f, e, h, f+e*t/1000
        print
        print
    else:
        temperatures = [ v['temperature'] for v in thermal_properties ]
        props = [ v[prop_target] for v in thermal_properties ]
        props = np.array( props ) * options.factor
        if options.is_diff:
            props -= props_0
        plt.plot( temperatures[tmin_index:tmax_index],
                  props[tmin_index:tmax_index] )

        

if not options.is_gnuplot:
    if ( not options.ymin == None ) and ( not options.ymax == None ):
        plt.ylim( options.ymin, options.ymax )
    elif options.ymin == None and ( not options.ymax == None ):
        plt.ylim( ymax=options.ymax )
    elif ( not options.ymin == None ) and options.ymax == None:
        plt.ylim( ymin=options.ymin )
    
    if not options.output_filename == None:
        plt.rcParams['pdf.fonttype'] = 42
        plt.rcParams['font.family'] = 'serif'
        plt.savefig(options.output_filename)
    else:
        plt.show()