bandstruct.py

#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
Plot VASP band structure with orbital character.

Usage: python ~/VASPtools/bandstruct.py

@author: Oleg Rubel
"""

import numpy as np
import os
import sys
import matplotlib.pyplot as plt
from scipy.ndimage import gaussian_filter, gaussian_filter1d

def user_input():
    """Editable section where users define their input"""
    erange = [-2.5, 1]
    # [-22.5, -20] # energy range (eV) relative to the Fermi energy (for Sn-d states)
    krange = [1,73] # range of k points (1 is the 1st k point); None is for no boundary
    # [74, None] # points on a mesh with weights (for Sn-d states)
    # [1,73] # range of k points (1 is the 1st k point); None is for no boundary
    atoms =  [16, 17, 18, 44, 45, 46] # slab middle (FeSn)
    # [1, 11, 21, 31, 41, 51] # Co-termination side (CoSn)
    # [10, 20, 30, 40, 50, 60] # Sn-termination side (CoSn)
    # [6, 16, 26, 35, 45, 56] # middle of the slab (CoSn)
    # [16, 17, 18, 44, 45, 46] # slab middle (FeSn)
    # [28, 29, 30, 58, 59, 60] # Sn-termination side (FeSn)
    # [1, 2, 3, 31, 32, 33] # Fe-termination side (FeSn) # list of atoms for orbital contribution (1 is the 1st atom)
    # s  py  pz  px  dxy  dyz  dz2  dxz  dx2-y2  fy3x2  fxyz  fyz2  fz3  fxz2  fzx2  fx3  tot
    # 1  2   3   4    5    6    7    8      9      10    11    12   13    14    15   16   17
    # or
    # s  py  pz  px  dxy  dyz  dz2  dxz  dx2-y2  tot
    # 1  2   3   4    5    6    7    8      9     10
    orb = [10] # list of orbitals to include in weights
    xticks = {1:"M", 27:"G", 58:"K", 73:"M'"} # y axis ticks
    # {1:"G"} (for Sn-d states)
    # {1:"M", 27:"G", 58:"K", 73:"M'"} # y axis ticks
    readprocar = True # read PROCAR file
    showplots = True
    weighIPR = False # "fat bands" weighted by the inverse participation ratio
    spinpolar = False # spin polarization with vasp_std
    c = 10 # scaling coefficient for point plots

    return erange, krange, atoms, orb, xticks, readprocar, showplots, c, weighIPR, spinpolar
    # end user_input

def read_EIGENVAL(spinpolar):
    """Read eigenvalues and k points from EIGENVAL file"""
    WorkingDir = os.getcwd()
    print (f'Working directory = {WorkingDir}')
    os.chdir(WorkingDir)
    i = 0 # line counter
    ik = 0 # k point counter
    ie = 0 # eigenvalue counter
    nband = float("inf")
    with open('EIGENVAL') as file:
        print ('Reading EIGENVAL file ...')
        for line in file:
            i += 1
            if i == 6: # 6th line (heading)
                # get number of k points and bands
                lsplit = line.rstrip().split()
                if not(len(lsplit) == 3):
                    raise ValueError(f'lsplit list should have length = 3, '+
                            f'while you have {len(lsplit)}. '+
                            f'The line is: {line.rstrip()}')
                
                [dum, numk, nband] = [int(num) for num in lsplit]
                print (f'Number of k points = {numk}')
                print (f'Number of bands = {nband}')
                # allocate arrays
                kpt = np.zeros((numk,3)) # k-point coordinate
                kptw = np.zeros(numk) # k-point weight
                if not(spinpolar):
                    eig = np.zeros((numk,nband))
                else:
                    eig = np.zeros((numk,nband,2))
                    
            elif np.mod(i,nband+2) == 8: # k point line
                ik += 1
                lsplit = line.rstrip().split()
                # the line should contain 'k1 k2 k3 weight'
                if not(len(lsplit) == 4):
                    raise ValueError(f'lsplit list should have length = 4, '+
                            f'while you have {len(lsplit)}. The line is: {line.rstrip()}')
                
                kpti = [float(num) for num in lsplit]
                kpt[ik-1,:] = kpti[0:3] # store k point coordinates
                kptw[ik-1] = kpti[3] # store k point weight
            elif (np.mod(i,nband+2) > 8 and np.mod(i,nband+2) < 8+nband+1) or \
                    (i > 8 and np.mod(i,nband+2) < 7): # eigenvalue line
                ie += 1
                lsplit = line.rstrip().split()
                if not(len(lsplit) == 3) and not(spinpolar): # spin unpolarized or SOC or NCL (vasp_std, ISPIN = 1 or vasp_ncl)
                    raise ValueError(f'lsplit list should have length = 3, '+
                            f'while you have {len(lsplit)}. '+
                            f'The line is: {line.rstrip()}')
                elif not(len(lsplit) == 5) and spinpolar: # spin-polarized, no SOC (vasp_std, ISPIN = 2)
                    raise ValueError(f'lsplit list should have length = 3, '+
                            f'while you have {len(lsplit)}. '+
                            f'The line is: {line.rstrip()}')
                
                iband = np.mod(ie,nband)
                if not(spinpolar):
                    eig[ik-1,iband-1] = float(lsplit[1])
                else:
                    eig[ik-1,iband-1,0] = float(lsplit[1])
                    eig[ik-1,iband-1,1] = float(lsplit[2])

    print (f'Processed {ik} of {numk} k points and {ie} of {numk*nband} eigenvalues')
    if not(ik == numk):
        raise ValueError(f'The number of k points processed {ik} is not equal '+
                f'to {numk} read from heading of EIGENVAL file')
    
    if not(ie == numk*nband):
        raise ValueError(f'The number of bands processed {ie} is not equal to '+
                f'numk*nband={numk}*{nband}={numk*nband} read from heading of '+
                'EIGENVAL file')

    return numk, nband, eig, kpt, kptw
    # end read_EIGENVAL

def read_PROCAR(atoms, orb, weighIPR, spinpolar):
    """Read PROCAR and compute orbital contribution of individual atoms for 
    fat bands plot"""
    WorkingDir = os.getcwd()
    print (f'Working directory = {WorkingDir}')
    os.chdir(WorkingDir)
    i = 0 # line counter
    ik = 0 # k point counter
    ie = 0 # eigenvalue counter
    iion = 0 # ion counter
    ispin = 0 # spin counter
    readorbw = False # flag to enable reading of orbital projections
    nband = float("inf")
    with open('PROCAR') as file:
        print ('Reading PROCAR file ...')
        for line in file:
            i += 1
            if i == 2: # 2nd line (heading)
                # get number of k points, bands, ions
                lsplit = line.rstrip().split()
                if not(len(lsplit) == 12):
                    raise ValueError(f'lsplit list should have length = 12, '+
                            f'while you have {len(lsplit)}. '+
                            f'The line is: {line.rstrip()}')
                
                numk = int(lsplit[3])
                nband = int(lsplit[7])
                nion = int(lsplit[11])
                print (f'Number of k points = {numk}')
                print (f'Number of bands = {nband}')
                print (f'Number of ions = {nion}')
                
                if weighIPR: # collect projections for all atoms
                    print('Since IPR plot is requested, the "atoms" '+\
                            'variable is ignored')
                    atoms = list(range(1, nion+1)) # added 1 since range last element is n-1
                    print('Since IPR plot is requested, the "orb" '+\
                            'variable is ignored')

                if max(atoms) > nion:
                    raise ValueError(f'One of atoms in the list atoms={atoms} '+
                            f'exceeds the total number of atoms {nion}')
                
                # allocate array
                if not(spinpolar):
                    w = np.zeros((numk,nband))
                else:
                    w = np.zeros((numk,nband,2))
                    
            elif 'k-point ' in line: # k point line (space is important here not to confuse with "# of k-points: " line)
                ik += 1
                if ik > 1 and not(ie == nband): # check that all bands are read
                    raise ValueError(f'The number of bands processed {ie} is '+
                            f'not equal to {nband} read from heading of PROCAR '+
                            'file')
                elif ik > 1:
                    print(f'    processed {ie} bands of {nband}')
                
                ie = 0 # reset eigenvalue counter
                print(f'  reading k point {ik} of {numk}')
            elif 'band' in line: # band line
                ie += 1
            elif 'ion' in line: # weights by ion
                readorbw = True
                iion = 0 # reset ion counter
            elif readorbw and iion <= nion:
                iion += 1
                lsplit = line.rstrip().split()
                proj = lsplit[1:]
                if iion == 1: # first ion
                    # get ion and orbital resolved projections
                    nproj = len(proj)
                    if (nproj < max(orb)):
                        raise ValueError(f'The number of projection columns '+
                                f'in PROCAR is {nproj}, while you asked for '+
                                f'projection column {max(orb)}')

                    wi = np.zeros((nion,nproj)) # allocate
                    wtot = 0 # reset total weight
                    wtot2 = 0 # reset total weight^2
                
                wi[iion-1,:] = [float(num) for num in proj]
                wtot += wi[iion-1,-1] # add total weights of all ions
                if weighIPR:
                    wtot2 += wi[iion-1,-1]**2.0 # collect sum(w_i^2)

                if iion == nion: # last ion
                    if weighIPR:
                        # compute IPR sum(w^2)/[sum(w)]^2
                        if not(spinpolar):
                            w[ik-1,ie-1] = wtot2/(sum(wi[:,-1])**2.0)
                        else:
                            w[ik-1,ie-1,ispin] = wtot2/(sum(wi[:,-1])**2.0)
                            
                        # normalize by ideal IPR = 1/nion
                        # in this case the w=1 means no localization at all
                        if not(spinpolar):
                            w[ik-1,ie-1] = w[ik-1,ie-1]*nion
                        else:
                            w[ik-1,ie-1,ispin] = w[ik-1,ie-1,ispin]*nion
                    else:
                        for jion in atoms:
                            for jorb in orb:
                                # add all relevant projections
                                if not(spinpolar):
                                    w[ik-1,ie-1] += wi[jion-1,jorb-1]
                                else:
                                    w[ik-1,ie-1,ispin] += wi[jion-1,jorb-1]
                        
                        if (wtot > 1 + 0.0005*nion) or (wtot == 0):
                            # Here '0.0005*atoms.count' accounts for round-off errors
                            # associated with storing data in PROCAR using 3 digits after
                            # the decimal.
                            raise ValueError(f'The total weight of all ions is '+
                                    f'{wtot} in band {ie} kpoin {ik} not within '+
                                    f'bounds (0, 1]')

                        # w[ik-1,ie-1] = w[ik-1,ie-1]/wtot # normalize the total weight such that sum(w(all atoms, all orb))=1
                        # if (w[ik-1,ie-1] > 1.01):
                        #     raise ValueError(f'The weight is {w[ik-1,ie-1]} '+
                        #             f'in band {ie} kpoin {ik} exceedes 1.01')

                    readorbw = False
                    
            if i > 2 and spinpolar and ispin == 0 and ik == numk:
                # starting to read the second spin record
                # reset k-point counter and advance spin counter
                ik = 0
                ispin += 1
                

        if not(ik == numk):
            raise ValueError(f'The number of k points processed {ik} is not '+
                    f'equal to {numk} read from heading of PROCAR file')

    print("w.shape=",w.shape)
    return w
    # end read_PROCAR

def read_OUTCAR():
    """Read OUTCAR file and determine reciprocal lattice vectors G (1/Ang) 
    and the Fermi energy (eV)"""
    WorkingDir = os.getcwd()
    print (f'Working directory = {WorkingDir}')
    os.chdir(WorkingDir)
    i = 0 # line counter
    readG = False # flag to enable reading of G vector
    with open('OUTCAR') as file:
        print ('Reading OUTCAR file ...')
        for line in file:
            i += 1
            if ('direct lattice vectors' in line) and \
                    ('reciprocal lattice vectors' in line): # G vector lines
                readG = True
                iG = 0 # G vector lines counter
                G = np.zeros((3,3)) # allocate
            elif readG: # Read G vector
                iG += 1
                lsplit = line.rstrip().split()
                if not(len(lsplit) == 6):
                    raise ValueError(f'The line suppose to split into 6 '+
                            f'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)
                G[iG-1,:] = [float(num) for num in lsplit[3:]]
                if iG == 3:
                    readG = False # stop reading G matrix
            elif 'E-fermi :' in line: # Fermi energy (eV)
                lsplit = line.rstrip().split()
                efermi = float(lsplit[2])

        print(f'Fermi energy from OUTCAR: {efermi} (eV)')
        print(f'Reciprocal lattice vectors from OUTCAR are (1/Ang):')
        for i in range(3):
            print(f'  G({i+1})=[{G[i,0]}, {G[i,1]}, {G[i,2]}]')

    return G, efermi
    # 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 user parameters
    erange, krange, atoms, orb, xticks, readprocar, showplots, c, \
            weighIPR, spinpolar = user_input()

    # Check input
    if not(len(erange) == 2):
        raise ValueError(f'erange list should have length = 2, '+
                f'while you have {len(erange)}')
    
    # Print input
    print("User input:")
    print(f'Energy range from {erange[0]} to {erange[1]} (eV) relative to '+
            f'the Fermi energy')
    
    # read OUTCAR
    G, efermi = read_OUTCAR()
    # read eigenvalues and k points from EIGENVAL
    numk, nband, eig, kpt, kptw = read_EIGENVAL(spinpolar)
    eig = eig - efermi # subtract Fermi energy
    # crop k points
    if not(krange[0] == None) and krange[1] == None: # range [XX, None]
        ik1 = krange[0] - 1
        kpt = kpt[ik1:,:]
        kptw = kptw[ik1:]
        if not(spinpolar):
            eig = eig[ik1:]
        else:
            eig = eig[ik1:,:]
            
        if numk - ik1 <= 0: # senity check
            raise ValueError(f'Total number of k points = {numk}, '+\
                    f'while the input requested k points starting with {ik1}')
        else:
            numk = numk - ik1
            
    elif not(krange[0] == None) and krange[1] is not None: # range [XX, YY]
        ik1 = krange[0] - 1
        ik2 = krange[1] - 1
        kpt = kpt[ik1:ik2+1,:]
        kptw = kptw[ik1:ik2+1]
        if not(spinpolar):
            eig = eig[ik1:ik2+1]
        else:
            eig = eig[ik1:ik2+1,:]
            
        if numk - ik1 <= 0: # senity check
            raise ValueError(f'Total number of k points = {numk}, '+\
                    f'while the input requested k points starting with {ik1}')
        else:
            numk = ik2 - ik1 + 1
            
    else:
        raise ValueError(f'This choice of krange={krange} is not implemented')
    
    # convert k point fractional coordinates into Cartesian
    KPATH = coordTransform(kpt,G)
    print('KPATH.shape=', KPATH.shape)
    print('numk=', numk)
    # compute length along the k-path
    K = np.zeros((numk))
    for i in range(numk-1):
        B = KPATH[i+1,:] - KPATH[i,:]
        dk = np.sqrt(np.dot(B,B))
        K[i+1] = K[i] + dk
    
    # plot simple band structure
    if max(xticks) > numk:
        raise ValueError(f'The number of k points is {numk} but max(xticks)={max(xticks)}')

    print('Plotting simple band structure (PDF will be stored after closing the graphical window)...')
    f1 = plt.figure()
    for ie in range(nband):
        if not(spinpolar):
            if max(eig[:,ie]) > erange[0] and min(eig[:,ie]) < erange[1]:
                plt.plot(K, eig[:,ie], c='black')
        else:
            if max(eig[:,ie,0]) > erange[0] and min(eig[:,ie,0]) < erange[1]:
                # spin UP
                plt.plot(K, eig[:,ie,0], c='black')
                
            if max(eig[:,ie,1]) > erange[0] and min(eig[:,ie,1]) < erange[1]:
                # spin DN
                plt.plot(K, eig[:,ie,1], c='black')

    ax = plt.gca()
    ax.set_xlim([min(K), max(K)])
    xticksK = [K[ik-1] for ik in xticks.keys()]
    xticksLabel = list(xticks.values())
    plt.xticks(ticks=xticksK, labels=xticksLabel)
    plt.grid(visible=True, which='major', axis='x')
    ax.set_ylim([erange[0], erange[1]])
    plt.axhline(y=0, color='black', linestyle='--') # line at E = 0
    plt.xlabel("Wave vector")
    plt.ylabel("Energy (eV)")
    if showplots:
        plt.show()
    
    # save figure
    f1.savefig("band_struct.pdf", bbox_inches='tight')
    print('Figure stored in band_struct.pdf file')

    if not(readprocar):
        sys.exit('Finished without reading PROCAR')

    # plot FAT bands
    f2 = plt.figure()
    # read PROCAR
    print('Reading PROCAR...')
    weight = read_PROCAR(atoms, orb, weighIPR, spinpolar)
    print(f'max kptw={np.max(kptw)}')
    print(f'min kptw={np.min(kptw)}')
    weightk = np.zeros(np.shape(weight)) # this weights combine PROCAR weights with k-point weights
    for ik in range(weight.shape[0]):
        for je in range(weight.shape[1]):
            weightk[ik,je] = weight[ik,je] * kptw[i]

    # crop k points
    if not(krange[0] == None) and krange[1] == None:
        if not(spinpolar):
            weight = weight[ik1:,:]
            weightk = weightk[ik1:,:]
        else:
            weight = weight[ik1:,:,:]
            weightk = weightk[ik1:,:,:]
            
    elif not(krange[0] == None) and krange[1] is not None: # range [XX, YY]
        ik1 = krange[0] - 1
        ik2 = krange[1] - 1
        if not(spinpolar):
            weight = weight[ik1:ik2+1]
            weightk = weightk[ik1:ik2+1]
        else:
            weight = weight[ik1:ik2+1,:]
            weightk = weightk[ik1:ik2+1,:]
            
    else:
        raise ValueError(f'This choice of krange={krange} is not implemented')

    # plot bandstructure 2
    if not(spinpolar):
        maxw = 0.0
        minw = 999.0
        maxwk = 0.0
        minwk = 999.0
    else:
        maxwUP = 0.0
        minwUP = 999.0
        maxwDN = 0.0
        minwDN = 999.0
        maxwkUP = 0.0
        minwkUP = 999.0
        maxwkDN = 0.0
        minwkDN = 999.0
        
    eigs = np.empty( shape=(0, ) ) # to gather data for plotting
    weights = np.empty( shape=(0, ) )
    weightsk = np.empty( shape=(0, ) )
    Ks = np.empty( shape=(0, ) )
    for ie in range(nband):
        if not(spinpolar):
            if max(eig[:,ie]) > erange[0] and min(eig[:,ie]) < erange[1]:
                eigs = np.append(eigs, eig[:,ie])
                weights = np.append(weights, weight[:,ie])
                weightsk = np.append(weightsk, weightk[:,ie])
                Ks = np.append(Ks, K)
                minw = np.minimum(minw, np.min(weight[:,ie]))
                maxw = np.maximum(maxw, np.max(weight[:,ie]))
                minwk = np.minimum(minwk, np.min(weightk[:,ie]))
                maxwk = np.maximum(maxwk, np.max(weightk[:,ie]))
        
        else:
            if max(eig[:,ie,0]) > erange[0] and min(eig[:,ie,0]) < erange[1]:
                # spin UP
                eigs = np.append(eigs, eig[:,ie,0])
                weights = np.append(weights, weight[:,ie,0])
                weightsk = np.append(weightsk, weightk[:,ie,0])
                Ks = np.append(Ks, K)
                minwUP = np.minimum(minwUP, np.min(weight[:,ie,0]))
                maxwUP = np.maximum(maxwUP, np.max(weight[:,ie,0]))
                minwkUP = np.minimum(minwkUP, np.min(weightk[:,ie,0]))
                maxwkUP = np.maximum(maxwkUP, np.max(weightk[:,ie,0]))
            
            if max(eig[:,ie,1]) > erange[0] and min(eig[:,ie,1]) < erange[1]:
                # spin DN
                eigs = np.append(eigs, eig[:,ie,1])
                weights = np.append(weights, weight[:,ie,1])
                weightsk = np.append(weightsk, weightk[:,ie,1])
                Ks = np.append(Ks, K)
                minwDN = np.minimum(minwDN, np.min(weight[:,ie,1]))
                maxwDN = np.maximum(maxwDN, np.max(weight[:,ie,1]))
                minwkDN = np.minimum(minwkDN, np.min(weightk[:,ie,1]))
                maxwkDN = np.maximum(maxwkDN, np.max(weightk[:,ie,1]))
            
            minw = np.minimum(minwUP, minwDN)
            maxw = np.maximum(maxwUP, maxwDN)
            minwk = np.minimum(minwkUP, minwkDN)
            maxwk = np.maximum(maxwkUP, maxwkDN)
                

    print(f'Min-max weight in the plot: {minw} - {maxw}')
    print(f'Min-max weight in the plot incl. k-point weights: {minwk} - {maxwk}')
    print('Plotting fat band structure (PDF will be stored after closing the graphical window)...')
    scatter = plt.scatter(Ks, eigs, s=c*weights, c='black', alpha=0.5)
    ax = plt.gca()
    ax.set_xlim([min(K), max(K)])
    xticksK = [K[ik-1] for ik in xticks.keys()]
    xticksLabel = list(xticks.values())
    plt.xticks(ticks=xticksK, labels=xticksLabel)
    plt.grid(visible=True, which='major', axis='x')
    ax.set_ylim([erange[0], erange[1]])
    plt.axhline(y=0, color='black', linestyle='--') # line at E = 0
    plt.xlabel("Wave vector")
    plt.ylabel("Energy (eV)")
    # legend
    handles, labels = scatter.legend_elements(prop="sizes", num=6, fmt=None, func=lambda s: s/c)
    plt.legend(handles, labels, title='Weight:')

    if showplots:
        plt.show()
    
    # save figure
    f2.savefig("band_struct_fat.pdf", bbox_inches='tight')
    print('Figure stored in band_struct_fat.pdf file')

    # plot 3 blurred fat bands (ARPES-like)
    bins = 1000                       # grid resolution (increase for finer detail)
    hist, xedges, yedges = np.histogram2d(x=Ks, y=eigs, bins=bins, weights=c*weights, density=False)
    sigma = 20.0                      # controls how blurry the result is
    blurred = gaussian_filter(hist, sigma=sigma)
    f3 = plt.figure()
    plt.imshow(
        blurred.T,
        origin="lower",
        extent=[xedges[0], xedges[-1], yedges[0], yedges[-1]],
        aspect="auto",
    )
    ax = plt.gca()
    ax.set_xlim([min(K), max(K)])
    xticksK = [K[ik-1] for ik in xticks.keys()]
    xticksLabel = list(xticks.values())
    plt.xticks(ticks=xticksK, labels=xticksLabel)
    plt.grid(visible=True, which='major', axis='x')
    ax.set_ylim([erange[0], erange[1]])
    plt.axhline(y=0, color='black', linestyle='--') # line at E = 0
    plt.xlabel("Wave vector")
    plt.ylabel("Energy (eV)")
    
    if showplots:
        plt.show()
        
    # save figure
    f3.savefig("band_struct_fat_blurred.pdf", bbox_inches='tight')
    print('Figure stored in band_struct_fat_blurred.pdf file')
    
    # plot 4 blurred line plot
    bins = 1000                       # grid resolution (increase for finer detail)
    sigma = 30.0                      # controls how blurry the result is
    xmin = np.min(eigs)
    xmax = np.max(eigs)
    dx = xmax - xmin
    xmin = xmin - 3*sigma/bins*dx
    xmax = xmax + 3*sigma/bins*dx
    hist, xedges = np.histogram(a=eigs, bins=bins, range=(xmin, xmax), weights=weightsk, density=False)    
    blurred = gaussian_filter1d(hist, sigma=sigma)
    f4 = plt.figure()
    plt.plot(xedges[:-1], blurred)
    plt.xlabel("Energy (eV)")
    plt.ylabel("Density")
    
    if showplots:
        plt.show()
        
    # save figure
    f4.savefig("DOS.pdf", bbox_inches='tight')
    print('Figure stored in DOS.pdf file')