processOUTCAR.py

#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
Process OUTCAR file generated by VASP

Usage: python ~/VASPtools/processOUTCAR.py -o OUTCAR-1 -b 10
optional arguments:
  -h, --help  show this help message and exit
  -o O        name of the VASP OUTCAR file
  -b B        band range for plotting above and below the Fermi energy

@author: Oleg Rubel
"""

import numpy as np
import os
import sys
import csv
import argparse # parse line arguments

def user_input():
    """Editable section where users define their input"""
    showplots = True

    return bandrange, showplots
    # end user_input

def read_OUTCAR(fnameOUTCAR):
    """Read OUTCAR file and determine
    - the Fermi energy (eV)"""
    WorkingDir = os.getcwd()
    print (f'Working directory = {WorkingDir}')
    os.chdir(WorkingDir)
    i = 0 # line counter
    readene = False # flag to enable reading of G vector
    ik = 0 # counter for k points read
    with open(fnameOUTCAR) as file:
        print (f'Reading {fnameOUTCAR} file ...')
        for line in file:
            i += 1
            if ('band energies' in line): # energy eigenvalues section
                readene = True
                iE = 0 # G vector lines counter
                G = np.zeros((3,3)) # allocate
            elif readene: # Read eigenvalues
                # line looks like this
                #      10     -22.7712      1.00000
                lsplit = line.rstrip().split()
                if not(len(lsplit) == 3):
                    raise ValueError(f'The line suppose to split into 3 values, but it does not. Here is the line: {line}')
                # Skip first 3 values. Those are real space lattice parameters
                # Units of G are (1/Ang)
                if ((ik >= nkpts) and lspin and not(lsorbit)): # filling up spin DN
                    eneEigvalsDN[ik-nkpts,iE] = float(lsplit[1])
                else: # spin UP or SOC
                    eneEigvals[ik,iE] = float(lsplit[1])
                if (iE == nbands-1): # all eigenvalues read?
                    readene = False # stop reading G matrix
                    ik += 1 # next k point
                else:
                    iE += 1 # next eiganvalue
            elif ('NBANDS=' in line): # get number of bands and k-points
                # working with line
                # k-points           NKPTS =      4   k-points in BZ     NKDIM =      4   number of bands    NBANDS=   3016
                lsplit = line.rstrip().split('NKPTS =')
                lsplit = lsplit[1].rstrip().split()
                nkpts = int(lsplit[0])
                lsplit = line.rstrip().split('NBANDS=')
                lsplit = lsplit[1].rstrip().split()
                nbands = int(lsplit[0])
            elif ('ISPIN  =' in line): # spin polarized?
                lsplit = line.rstrip().split('ISPIN  =')
                lsplit = lsplit[1].rstrip().split()
                if lsplit[0] == '1':
                    lspin = False
                else:
                    lspin = True
            elif ('LSORBIT =' in line): # SOC?
                lsplit = line.rstrip().split('LSORBIT =')
                lsplit = lsplit[1].rstrip().split()
                if lsplit[0] == 'T':
                    lsorbit = True
                else:
                    lsorbit = False
            elif ('E-fermi :' in line): # Fermi energy (eV) and allocate E(k,n)
                lsplit = line.rstrip().split()
                efermi = float(lsplit[2])
                eneEigvals = np.zeros((nkpts,nbands)) # E(k,n)
                eneEigvalsDN = np.zeros((nkpts,nbands))
                ik = 0 # reset counter for k points since new set of eiganvalues
                       # will follow E_Fermi in the OUTCAR file

        print(f'Number of bands from {fnameOUTCAR}: {nbands}')
        print(f'Number of k-points from {fnameOUTCAR}: {nkpts}')
        print(f'Spin-polarized calculation from {fnameOUTCAR}: {lspin}')
        print(f'SOC from {fnameOUTCAR}: {lsorbit}')
        print(f'Fermi energy from {fnameOUTCAR}: {efermi} (eV)')

    return nkpts, nbands, lspin, lsorbit, efermi, eneEigvals, eneEigvalsDN
    # end read_OUTCAR

def coordTransform(V,G):
    """Transform vector V(:,3) in G(3,3) coord. system -> W(:,3) in Cartesian 
    coordinates"""
    W = np.zeros(np.shape(V))
    for i in range(np.shape(V)[0]):
        W[i,:] = G[0,:]*V[i,0] + G[1,:]*V[i,1] + G[2,:]*V[i,2]

    return W
    # end coordTransform

# MAIN
if __name__=="__main__":
    # Set up parser for line arguments
    parser = argparse.ArgumentParser()
    parser.add_argument("-o",\
        help="name of the VASP OUTCAR file",\
        type=str,\
        required=False,\
        default='OUTCAR')
    parser.add_argument("-b",\
        help="band range for plotting above and below the Fermi energy",\
        type=int,\
        required=False,\
        default=5)
    args = parser.parse_args()

    # Set user parameters
    fnameOUTCAR = args.o
    print(f'OUTCAR file name: {fnameOUTCAR}')
    bandrange = args.b
    showplots = user_input()
    
    # Print input
    #print("User input:")
    #print(f'Energy range from {erange[0]} to {erange[1]} (eV) relative to the Fermi energy')
    
    # read OUTCAR
    nkpts, nbands, lspin, lsorbit, efermi, eneEigvals,eneEigvalsDN \
        = read_OUTCAR(fnameOUTCAR)

    # find band that crosses the Fermi energy
    for ik in range(nkpts):
        exitloop = False
        for ib in range(nbands):
            if (eneEigvals[ik,ib] >= efermi):
                ibandFermi = ib # keep in mind that this is a band index
                                # the actual band number is (i + 1)
                ikptFermi = ik
                exitloop = True
                break
        if exitloop:
            break
    print(f'Band crossing the Fermi energy: {ibandFermi+1} at k-point {ikptFermi+1}')

    # write csv file
    data = eneEigvals[:,ibandFermi-bandrange:ibandFermi+bandrange]
    bands = np.arange(ibandFermi-bandrange,ibandFermi+bandrange,1) + 1
    if (lspin and not(lsorbit)): # spin dn component only
        dataDN = eneEigvalsDN[:,ibandFermi-bandrange:ibandFermi+bandrange]
        dataDN = np.concatenate(([bands],dataDN),axis=0)
        dataDN = np.transpose(dataDN)

    data = np.concatenate(([bands],data),axis=0)
    data = np.transpose(data)
    with open("bands.csv", "wt") as fp:
        writer = csv.writer(fp, delimiter=",")
        # write header
        writer.writerow(["# band", "energy k1 (eV)", "energy k2 (eV)" ,"..."])
        if (lspin and not(lsorbit)): # spin dn component only
            writer.writerow(["# up"])
        # write data
        writer.writerows(data)
        if (lspin and not(lsorbit)): # spin dn component only
            writer.writerow(["# dn"])
            writer.writerows(dataDN)
        writer.writerow(["# Fermi energy (eV)", efermi])
        writer.writerow(["# Band crossing the Fermi energy", ibandFermi+1])

    # plot results
    import matplotlib.pyplot as plt
    if (not(lspin) or lsorbit): # one spin component component only or spinor
        # plot simple band structure
        f1 = plt.figure()
        kindex = np.arange(1,nkpts+1)
        for ie in range(ibandFermi-bandrange,ibandFermi+bandrange):
            plt.plot(kindex, eneEigvals[:,ie], label=str(ie+1), marker='.')

        plt.legend() # show legend
        ax = plt.gca()
        ax.set_xlim([1, nkpts])
        plt.grid(visible=True, which='major', axis='both')
        plt.axhline(y=efermi, color='black', linestyle='--') # line at E_Fermi
        plt.xlabel("k-point index")
        plt.ylabel("Energy (eV)")
        if showplots:
            plt.show()

        # save figure
        f1.savefig("band_struct.pdf", bbox_inches='tight')
    else:
        # plot both spins
        f2, (ax1, ax2) = plt.subplots(1, 2, sharey=True)
        kindex = np.arange(1,nkpts+1)
        for ie in range(ibandFermi-bandrange,ibandFermi+bandrange):
            ax1.plot(kindex, eneEigvals[:,ie], label=str(ie+1), marker='.')
            ax2.plot(kindex, eneEigvalsDN[:,ie], label=str(ie+1), marker='.')

        ax1.legend() # show legend
        ax1.set_xlim([1, nkpts])
        ax2.set_xlim([1, nkpts])
        ax1.grid(visible=True, which='major', axis='both')
        ax2.grid(visible=True, which='major', axis='both')
        ax1.axhline(y=efermi, color='black', linestyle='--') # line at E_Fermi
        ax2.axhline(y=efermi, color='black', linestyle='--') # line at E_Fermi
        ax1.set(xlabel="k-point index", ylabel="Energy (eV)", title="spin UP")
        ax2.set(xlabel="k-point index", title="spin DN")
        if showplots:
            plt.show()

        # save figure
        f2.savefig("band_struct.pdf", bbox_inches='tight')

    print('Figure stored in band_struct.pdf file')