v1.0 release

Author: pipidog

DFTtoolbox/__init__.py

@@ -0,0 +1,311 @@
+# This file defines the postprocess module
+
+import numpy as np
+import sys, os, re, shutil
+import matplotlib.pyplot as plt
+import matplotlib.colors as col
+import matplotlib.cm as cm
+
+class dftpp:
+    '''
+    **** ABSTRACT ****
+    This module defines the core class of postprocessing the calculated
+    results of DFT such as band plot, fatband plot and spectral plot, etc.
+    This module can be used either independent or inherited by other 
+    package specified classes. 
+    **** INITLALIZE ****
+    N/A
+    **** METHODS ****
+    bicolormap(self,colorcode):
+        * generate three default colormaps for fatband plot.
+        => [colorcode]: str, 'r'/'g'/'b'
+        three default bicolormaps can be generated by simply input a
+        single keyword: 'r': red, 'g': green, 'b': blue
+
+    band_plot(self,spin,Ek,Ef,kdiv,klabel='default',Ebound='default',lw=2,fontsize=18)
+        * this method plots the band structures
+        => [spin]: int, 1 / 2
+        if two spin are separate, spin=2, else spin=1
+        => [Ek]: numpy array, tot_k x tot_band / tot_k x 2*tot_band
+        if spin=1, tot_k x tot_band. if spin=2, tot_k x 2*tot_band
+        => [Ef]: float
+        Fermi energy
+        => [kdiv]: list, int, e.g. [0,25,42]
+        location of the k-point to label high symmetry point. the first one
+        must be 0 and the last one must be tot_k-1
+        => [klabel]: list, string, e.g. ['$\Gamma','X','M']
+        name of the high symmetry point. must the same length as kdiv
+        default=['k1','k2',...,'kn']
+        => [Ebound]: list, two float values, e.g. [-3.5, 3.5]
+        upper and lower energy bounds
+        default=min(Ek) ~ max(Ek)
+        => [lw]: int, default=2
+        the line width 
+        [fontsize]: int, default=18
+        the fontsize of the plot
+        
+    fatband_plot(self,Ek,Ek_weight,Ef,state_info,state_grp,kdiv,klabel='default'\
+        ,Ebound='default',marker_size=30,colorcode='b'):
+        * this method plots the fatband sturctures
+        => [Ek]: numpy array, tot_k x tot_band (all spin counts) 
+        the energy bands, tot_band are spin included
+        => [Ek_weight]: numpy array, tot_k x tot_band x tot_state
+        the weight of each state at each eigenvalues
+        => [Ef]: float
+        the Fermi level. Ek will be shift based on it, so the output
+        Ef is always zero. 
+        => [state_info]: list, string, size: tot_state
+        a list of strings that describes each state. each element corresponds
+        to each state.
+        => [state_grp]: list, e.g:[[1,2,3],[7,8,10],[12,17]]
+        a list which shows the states to be grouped into a fatband plot. 
+        therefore, the above example will output three fatband plots where
+        the first plot shows the wieght of sum over state 1 2 and 3. 
+        => [kdiv], [klabel], [Ebound]
+        the same as band_plot
+        => [ini_fig_num]: int, default=1
+        in case you plot many plots in your code, you don't want it conflict with 
+        others. so you can set the ini_fig_num to tell the code the figure number
+        of each fatband plot. e.g. if ini_fig_num=3 and you have 3 state_grp, then
+        your fatband plot will be fig.3 ~ fig.5
+        => [marker_size]: int, default=30
+        the size of your fatband markers
+        => [colorcode]: str, 'r' / 'g' / 'b' , default: 'b'
+        the default colormap for fatband plot. 'r': red, 'g': green, 'b': blue
+        
+    spectral_plot(self,E,PDOS,state_grp='default',xlabel='Energy',ylabel='PDOS'\
+        ,llabel='default',Ebound='default',lw=2,fontsize=18):
+        * this method plots all kinds of spectral functions such as DOS, PDOS,
+        DMFT spectral, etc.
+        =>[E]: Nx1 numpy array, 
+        the x-axis of your plot
+        =>[PDOS]: numpy array, len(E) X tot_state
+        the PDOS of different states. each column corresponds to each state
+        => [state_grp]: a list of list 
+        combine the data to plot a line in the spectral, e.g, [[1,2,3],[4,5]] 
+        means combine the data of state [1,2,3] to plot a line and [4,5] for
+        another line. 
+        
+    **** Version ****
+    02/01/2016: first built
+    **** Comment ****
+    1. Run /test/postprocess_test.py to test this module 
+    '''
+    def __str__(self):
+        return 'DFTtooolbox: Postprocess Object'
+        
+    def grep(self,txtlines,kws,reg=False):
+        # search for line numbers from a textlines where contains the keywords
+        if reg:
+            ln=[ n for n,txt in enumerate(txtlines) if (re.search(kws,txt)!=None)]
+        else:
+            ln=[ n for n,txt in enumerate(txtlines) if (txt.find(kws)!=-1)]
+            
+        return ln
+        
+    def bicolormap(self,colorcode):
+        if colorcode=='r':
+            colorlist=['#e6e6e6','#cc0000']
+        elif colorcode=='g':
+            colorlist=['#e6e6e6','#009900']
+        elif colorcode=='b':
+            colorlist=['#e6e6e6','#000099']
+            
+        # generate a colormap
+        my_cmap = col.LinearSegmentedColormap.from_list('cmap_tmp',colorlist,N=32)
+
+        # return the colormap
+        return my_cmap 
+        
+    def state_grp_trans(self,state_info,state_grp):
+        # this method turn the string format of states in DFTtoolbox into the
+        # correct state label. To use it, one must orgainze the state info 
+        # as ' 1 => xxx ( atn / l / ml / ms)\n'. In this format, each state
+        # is labeled by 4 numbers within the the bracket only. The first one 
+        # is atom number the other three are quantum numbers.
+        # if some of them are not available, use 'a' instead. 
+        # xxx is extra information about the state which will not be used for filter.
+        state_grp_new=[]
+        if type(state_grp[0][0])==str: 
+            state_info=np.array(\
+            [ state.replace(')',' ').replace('(',' ').replace('/',' ').split()[-4:] for state in state_info])
+            def state_translator(state_str):
+                state=state_str.replace(':',' ').replace('/',' ').split()
+                
+                # pick atoms
+                s=set(range(0,len(state_info)))
+                s=s & set(np.nonzero(\
+                (state_info[:,0].astype(np.int) >= int(state[0]))\
+                & (state_info[:,0].astype(np.int) <= int(state[1]))\
+                )[0].tolist())
+                
+                # pick index (other than atom label) 
+                for n in range(0,3):
+                    if (state[n+2] is not 'a') & (state_info[:,n+1] is not 'a'):
+                        s=s & set(np.nonzero(state_info[:,n+1].astype(np.float)==float(state[n+2]))[0].tolist())
+                        
+                s=list(s)
+                if s==[]:
+                    print('Error: assigned state {0} not exist!'.format(state_str))
+                    sys.exit()
+                return s
+                
+            for grp in state_grp:
+                for subgrp in grp:
+                    state_grp_new.append(state_translator(subgrp))
+
+        elif type(state_grp[0][0])==int:
+                state_grp_new=state_grp
+                
+        return state_grp_new
+            
+    def band_plot(self,spin,Ek,Ef,kdiv='default',klabel='default',Ebound='default',\
+    lw=2,fontsize=18,savefig_dir='default'):
+    
+        # give input parameter default values
+        if (kdiv is 'default'):
+            kdiv=[0,Ek.shape[0]-1]
+        if (klabel is 'default'):
+            klabel=['k'+str(n) for n, val in enumerate(kdiv)]
+        if (Ebound is 'default'):
+            Ebound=[np.min(Ek),np.max(Ek)]
+        
+        tot_k=Ek.shape[0]
+        tot_ban=Ek.shape[1]
+        Ek=Ek-Ef
+
+        # plot normal band structure  
+        if spin==1:          
+          plt.plot(range(0,tot_k),Ek,'b',lw=lw)
+        elif spin==2:
+          plt.plot(range(0,tot_k),Ek[:,0:int(tot_ban/2)],'b'\
+          ,range(0,tot_k),Ek[:,int(tot_ban/2):],'g',lw=lw)
+         
+        # plot k divider  
+        if len(kdiv)>=3:    
+            for val in kdiv[1:-1]:
+               plt.plot(val*np.ones(10),np.linspace(np.min(Ek),np.max(Ek),10),'k--')
+        
+        # plot Fermi level    
+        plt.plot(np.linspace(0,tot_k-1,10),np.zeros(10),'r--')
+           
+        # tweak figure
+        plt.xticks(kdiv,klabel,fontsize=fontsize)
+        plt.yticks(fontsize=fontsize)
+        plt.ylabel('Energy (eV)',fontsize=fontsize)
+        plt.title('bands',fontsize=fontsize)
+        plt.xlim(0,tot_k-1)    
+        plt.ylim(Ebound[0],Ebound[1])
+        
+        # save figure
+        if savefig_dir is not 'default':
+            plt.savefig(savefig_dir+'band.png',bbox_inches='tight')
+        
+        # display figure
+        plt.show()
+            
+    def fatband_plot(self,Ek,Ek_weight,Ef,state_info,state_grp,kdiv='default',klabel='default'\
+    ,Ebound='default',ini_fig_num=1,marker_size=30,colorcode='b',fontsize=18,savefig_dir='default'):
+
+        # give input parameter default values
+        if (kdiv is 'default'):
+            kdiv=[0,Ek.shape[0]-1]
+        if (klabel is 'default'):
+            klabel=['k'+str(n) for n, val in enumerate(kdiv)]      
+        if (Ebound is 'default'):
+            Ebound=[np.min(Ek),np.max(Ek)]
+            
+        tot_k=Ek.shape[0]
+        tot_ban=Ek.shape[1]
+        
+        #organize data to scatter form
+        x=np.tile(np.array(range(0,tot_k)),[tot_ban,1]).transpose().flatten()
+        y=Ek.flatten()-Ef
+        
+        
+        # screen out data out of E_bound
+        my_cmap=self.bicolormap(colorcode)
+        for n, state_list in enumerate(state_grp):
+            w=np.zeros(len(Ek_weight[:,:,0].flatten()))
+            
+            print('\n==== fatband-'+str(n)+' ====')
+            # sum over projected states
+            for state in state_list:
+                print(state_info[state])
+                w=w+Ek_weight[:,:,state].flatten()
+               
+                
+            # plot colorbar 
+            plt.figure(n+ini_fig_num)
+            cbar_obj = plt.contourf([[0,0],[0,0]],\
+            np.linspace(min(w),max(w),32), cmap=my_cmap)    
+            plt.clf()
+            plt.colorbar(cbar_obj)
+            
+            # lexsort data based on w, so low intensity will be plot first
+            w_sort=np.argsort(w)
+            plt.scatter(x=x[w_sort], y=y[w_sort],s=marker_size, c=w[w_sort], cmap=my_cmap,edgecolor='face')
+            
+            # plot k divider    
+            for val in kdiv[1:-1]:
+                plt.plot(val*np.ones(10),np.linspace(np.min(y),\
+                np.max(y),10),'k--')
+            
+            # plot Fermi level    
+            plt.plot(np.linspace(0,np.max(x),10),np.zeros(10),'r--')
+
+            # tweak figure
+            plt.xticks(kdiv,klabel,fontsize=fontsize)
+            plt.yticks(fontsize=fontsize)
+            plt.ylabel('Energy (eV)',fontsize=fontsize)
+            plt.title('fatband-'+str(n),fontsize=fontsize)
+            plt.xlim(0,np.max(x))    
+            plt.ylim( Ebound[0],Ebound[1])
+            
+            # save figure
+            if savefig_dir is not 'default':
+                plt.savefig(savefig_dir+'fatband-'+str(n)+'.png',bbox_inches='tight')
+            
+            # display figure
+            plt.show()
+
+    def spectral_plot(self,E,PDOS,state_info,state_grp='default',xlabel='Energy',ylabel='PDOS'\
+    ,llabel='default',Ebound='default',savefig_path='default',lw=2,fontsize=18):
+        if Ebound=='default':
+            Ebound=[np.min(E),np.max(E)]
+            
+        if state_grp is 'default':
+            state_grp=[n for n in range(0,PDOS.shape[1])]
+        
+        PDOS_grp=np.zeros((E.shape[0],len(state_grp)))
+        for n, state_list in enumerate(state_grp):
+            print('spectral data-{0}'.format(n))
+            for state in state_list:
+                print('  {0}'.format(state_info[state]))
+                PDOS_grp[:,n]+=PDOS[:,state]
+            
+            if llabel is 'default':
+                plt.plot(E,PDOS_grp[:,n],lw=lw,label='data-'+str(n))
+            else:
+                plt.plot(E,PDOS_grp[:,n],lw=lw,label=llabel[n])
+
+
+        # tweak figure
+        plt.plot(np.zeros(10),np.linspace(1.01*np.min(PDOS_grp),1.01*np.max(PDOS_grp),10),'r--')
+        plt.legend(loc='upper right', shadow=True, fontsize='large',framealpha=0.0)      
+        plt.xlabel(xlabel,fontsize=fontsize)
+        plt.ylabel(ylabel,fontsize=fontsize)
+        plt.title('spectral',fontsize=fontsize)
+        plt.xlim(Ebound[0],Ebound[1])  
+        plt.xticks(fontsize=fontsize)
+        plt.yticks(fontsize=fontsize)  
+        if savefig_path!='default':
+            plt.savefig(savefig_path,bbox_inches='tight')
+            print('spectral has been saved in: '+savefig_path)
+        plt.show()
+        
+
+        
+
+
+