# -*- coding:utf-8 -*-
"""
========================================================================
Provide electro-related file class which do operations on these files.
========================================================================
Written by PytLab <shaozhengjiang@gmail.com>, September 2015
Updated by PytLab <shaozhengjiang@gmail.com>, August 2016
========================================================================
"""
import copy
import logging
from string import whitespace
import numpy as np
from scipy.integrate import simps
from scipy.interpolate import interp2d
import mpl_toolkits.mplot3d
# whether pyplot installed
try:
import matplotlib.pyplot as plt
plt_installed = True
except ImportError:
print('Warning: Module matplotlib.pyplot is not installed')
plt_installed = False
#whether mayavi installed
try:
from mayavi import mlab
mayavi_installed = True
except ImportError:
mayavi_installed = False
from vaspy.plotter import DataPlotter
from vaspy.atomco import PosCar
from vaspy.functions import line2list
class DosX(DataPlotter):
def __init__(self, filename, field=' ', dtype=float):
"""
Create a DOS file class.
Example:
>>> a = DosX(filename='DOS1')
Class attributes descriptions
=======================================================
Attribute Description
============ =======================================
filename string, name of the SPLITED DOS file
field string, separator of a line
dtype type, convertion type of data
reset_data method, reset object data
plotsum method, 绘制多列加合的图像
============ =======================================
"""
DataPlotter.__init__(self, filename=filename, field=' ', dtype=float)
# Set logger.
self.__logger = logging.getLogger("vaspy.DosX")
def __deepcopy__(self, memo):
"""
Overload copy.deepcopy behavior, only deep copy data when call copy.deepcopy().
"""
cls = self.__class__
result = cls.__new__(cls)
memo[id(result)] = result
# Deepcopy object components, data ONLY.
result.data = copy.deepcopy(self.data, memo)
return result
def __add__(self, dosx_inst):
# Get a copy.
sum_dosx = copy.deepcopy(self)
# 相加之前判断能量分布是否相同
same = (self.data[:, 0] == dosx_inst.data[:, 0]).all()
if not same:
raise ValueError('Energy is different.')
sum_dosx.data[:, 1:] = self.data[:, 1:] + dosx_inst.data[:, 1:]
sum_dosx.filename = "DOS_SUM"
return sum_dosx
def reset_data(self):
"Reset data array to zeros."
self.data[:, 1:] = 0.0
return self
def plotsum(self, xcol, ycols, **kwargs):
'''
绘制多列加合的图像.
Parameter
---------
xcol: int
column number of data for x values
绘制图像的x轴数据的列号
ycols: tuple of int
column numbers of data for y values
(start, stop[, step])
绘制图像的Y轴数据的列号，可以是多个，并进行列向量自动合并
Optional kwargs:
----------------
fill: Fill the area below fermi level or not, bool.
The default value is True.
show_dbc: Show the label of dband-center or not, bool.
The default value is False.
show_fermi: Show the lable of fermi level or not, bool.
The default value is True.
Example:
--------
# Use the 0th column data as x, sum of 1st and 2nd column data as y.
>>> a.plotsum(0, (1, 3))
# Use the 0th column data as x, sum of #5, #7, #9 column data as y.
>>> a.plotsum(0, (5, 10, 2))
'''
# Get kwargs.
fill = kwargs.pop("fill", True)
show_fermi = kwargs.pop("show_fermi", True)
d_cols = kwargs.pop("d_cols", (0, 0))
show_dbc = kwargs.pop("show_dbc", False)
x = self.data[:, xcol]
if len(ycols) == 2:
start, stop = ycols
step = 1
else:
start, stop, step = ycols
ys = self.data[:, start:stop:step]
y = np.sum(ys, axis=1)
ymax = np.max(y)
ymin = np.min(y)
# Plot.
fig = plt.figure()
ax = fig.add_subplot(111)
ax.plot(x, y, linewidth=5, color='#104E8B')
# Plot fermi energy auxiliary line.
if show_fermi:
# Fermi verical line
xfermi = [0.0, 0.0]
yfermi = [int(ymin-1), int(ymax+1)]
ax.plot(xfermi, yfermi, linestyle='dashed',
color='#4A708B', linewidth=3)
# Fill area from minus infinit to 0.
if fill:
minus_x = np.array([i for i in x if i <= 0])
minus_y = y[: len(minus_x)]
ax.fill_between(minus_x, minus_y, facecolor='#B9D3EE',
interpolate=True)
# show d band center line
if show_dbc:
dbc = self.get_dband_center(d_cols)
x_dbc = [dbc]*2
y_dbc = [int(ymin-1), int(ymax+1)]
ax.plot(x_dbc, y_dbc, linestyle='dashed',
color='#C67171', linewidth=3)
ax.set_xlabel(r'$\bf{E - E_F(eV)}$', fontdict={'fontsize': 20})
ax.set_ylabel(r'$\bf{pDOS(arb. unit)}$', fontdict={'fontsize': 20})
margin = (ymax - ymin)*0.2
ax.set_ylim(ymin-margin, ymax+margin)
fig.show()
return
def tofile(self, filename=None):
"""
DosX object to DOSX file.
Parameters:
-----------
filename: The name of generated DOSX file, str.
"""
ndata = self.data.shape[1] # data number in a line
data = self.data.tolist()
content = ''
for datalist in data:
content += ('%12.8f'*ndata + '\n') % tuple(datalist)
if filename is None:
filename = self.filename
with open(filename, 'w') as f:
f.write(content)
return
def get_dband_center(self, d_cols):
"""
Get d-band center of the DosX object.
Parameters:
-----------
d_cols: The column number range for d orbitals, int or tuple of int.
Examples:
---------
# The 5 - 9 columns are state density for d orbitals.
>>> dos.get_dband_center(d_cols=(5, 10))
"""
d_cols = (d_cols, d_cols+1) if type(d_cols) is int else d_cols
# 合并d轨道DOS
start, end = d_cols
yd = np.sum(self.data[:, start:end], axis=1)
#获取feimi能级索引
for idx, E in enumerate(self.data[:, 0]):
if E >= 0:
nfermi = idx
break
E = self.data[: nfermi+1, 0] # negative inf to Fermi
dos = yd[: nfermi+1] # y values from negative inf to Fermi
# Use Simpson integration to get d-electron number
nelectro = simps(dos, E)
# Get total energy of dband
tot_E = simps(E*dos, E)
dband_center = tot_E/nelectro
self.dband_center = dband_center
return dband_center
class ElfCar(PosCar):
def __init__(self, filename='ELFCAR'):
"""
Create a ELFCAR file class.
Example:
>>> a = ElfCar()
Class attributes descriptions
==============================================================
Attribute Description
============== =============================================
filename string, name of the ELFCAR file
------------- ame as PosCar ------------
bases_const float, lattice bases constant
bases np.array, bases of POSCAR
atoms list of strings, atom types
ntot int, the number of total atom number
natoms list of int, same shape with atoms
atom number of atoms in atoms
tf list of list, T&F info of atoms
data np.array, coordinates of atoms, dtype=float64
------------- ame as PosCar ------------
elf_data 3d array
plot_contour method, use matplotlib to plot contours
plot_mcontours method, use Mayavi.mlab to plot beautiful contour
plot_contour3d method, use mayavi.mlab to plot 3d contour
plot_field method, plot scalar field for elf data
============== =============================================
"""
super(ElfCar, self).__init__(filename)
# Set logger.
self.__logger = logging.getLogger("vaspy.ElfCar")
def load(self):
"Rewrite load method"
PosCar.load(self)
with open(self.filename, 'r') as f:
for i in range(self.totline):
f.readline()
#get dimension of 3d array
grid = f.readline().strip(whitespace)
empty = not grid # empty row
while empty:
grid = f.readline().strip(whitespace)
empty = not grid
x, y, z = line2list(grid, dtype=int)
#read electron localization function data
elf_data = []
for line in f:
datalist = line2list(line)
elf_data.extend(datalist)
#########################################
# #
# !!! Notice !!! #
# NGX is the length of the **0th** axis #
# NGY is the length of the **1st** axis #
# NGZ is the length of the **2nd** axis #
# #
#########################################
#reshape to 3d array
elf_data = np.array(elf_data).reshape((x, y, z), order='F')
#set attrs
self.grid = x, y, z
self.elf_data = elf_data
return
@staticmethod
def expand_data(data, grid, widths):
'''
根据widths, 将三维矩阵沿着x, y, z轴方向进行扩展.
'''
# expand grid
widths = np.array(widths)
expanded_grid = np.array(grid)*widths # expanded grid
# expand eld_data matrix
expanded_data = copy.deepcopy(data)
nx, ny, nz = widths
# x axis
added_data = copy.deepcopy(expanded_data)
for i in range(nx - 1):
expanded_data = np.append(expanded_data, added_data, axis=0)
# y axis
added_data = copy.deepcopy(expanded_data)
for i in range(ny - 1):
expanded_data = np.append(expanded_data, added_data, axis=1)
# z axis
added_data = copy.deepcopy(expanded_data)
for i in range(nz - 1):
expanded_data = np.append(expanded_data, added_data, axis=2)
return expanded_data, expanded_grid
# 装饰器
def contour_decorator(func):
'''
等值线作图方法装饰器.
Decorator for contour plot methods.
Set ndim on x, y axis and z values.
'''
def contour_wrapper(self, axis_cut='z', distance=0.5,
show_mode='show', widths=(1, 1, 1)):
'''
绘制ELF等值线图
Parameter in kwargs
-------------------
axis_cut: str
['x', 'X', 'y', 'Y', 'z', 'Z'], axis which will be cut.
distance: float
(0.0 ~ 1.0), distance to origin
show_mode: str
'save' or 'show'
widths: tuple of int,
number of replication on x, y, z axis
'''
#expand elf_data and grid
elf_data, grid = self.expand_data(self.elf_data, self.grid,
widths=widths)
self.__logger.info('data shape = %s', str(elf_data.shape))
# now cut the cube
if abs(distance) > 1:
raise ValueError('Distance must be between 0 and 1.')
if axis_cut in ['X', 'x']: # cut vertical to x axis
nlayer = int(self.grid[0]*distance)
z = elf_data[nlayer, :, :]
ndim0, ndim1 = grid[2], grid[1] # y, z
elif axis_cut in ['Y', 'y']:
nlayer = int(self.grid[1]*distance)
z = elf_data[:, nlayer, :]
ndim0, ndim1 = grid[2], grid[0] # x, z
elif axis_cut in ['Z', 'z']:
nlayer = int(self.grid[2]*distance)
z = elf_data[:, :, nlayer]
ndim0, ndim1 = grid[1], grid[0] # x, y
return func(self, ndim0, ndim1, z, show_mode=show_mode)
return contour_wrapper
@contour_decorator
def plot_contour(self, ndim0, ndim1, z, show_mode):
'''
ndim0: int, point number on x-axis
ndim1: int, point number on y-axis
z : 2darray, values on plane perpendicular to z axis
'''
#do 2d interpolation
#get slice object
s = np.s_[0:ndim0:1, 0:ndim1:1]
x, y = np.ogrid[s]
self.__logger.info('z shape = %s, x shape = %s, y shape = %s',
str(z.shape), str(x.shape), str(y.shape))
mx, my = np.mgrid[s]
#use cubic 2d interpolation
interpfunc = interp2d(x, y, z, kind='cubic')
newx = np.linspace(0, ndim0, 600)
newy = np.linspace(0, ndim1, 600)
#-----------for plot3d---------------------
ms = np.s_[0:ndim0:600j, 0:ndim1:600j] # |
newmx, newmy = np.mgrid[ms] # |
#-----------for plot3d---------------------
newz = interpfunc(newx, newy)
#plot 2d contour map
fig2d_1, fig2d_2, fig2d_3 = plt.figure(), plt.figure(), plt.figure()
ax1 = fig2d_1.add_subplot(1, 1, 1)
extent = [np.min(newx), np.max(newx), np.min(newy), np.max(newy)]
img = ax1.imshow(newz, extent=extent, origin='lower')
#coutour plot
ax2 = fig2d_2.add_subplot(1, 1, 1)
cs = ax2.contour(newx.reshape(-1), newy.reshape(-1), newz, 20, extent=extent)
ax2.clabel(cs)
plt.colorbar(mappable=img)
# contourf plot
ax3 = fig2d_3.add_subplot(1, 1, 1)
ax3.contourf(newx.reshape(-1), newy.reshape(-1), newz, 20, extent=extent)
#3d plot
fig3d = plt.figure(figsize=(12, 8))
ax3d = fig3d.add_subplot(111, projection='3d')
ax3d.plot_surface(newmx, newmy, newz, cmap=plt.cm.RdBu_r)
#save or show
if show_mode == 'show':
plt.show()
elif show_mode == 'save':
fig2d_1.savefig('surface2d.png', dpi=500)
fig2d_2.savefig('contour2d.png', dpi=500)
fig2d_3.savefig('contourf2d.png', dpi=500)
fig3d.savefig('surface3d.png', dpi=500)
else:
raise ValueError('Unrecognized show mode parameter : ' +
show_mode)
return
@contour_decorator
def plot_mcontour(self, ndim0, ndim1, z, show_mode):
"use mayavi.mlab to plot contour."
if not mayavi_installed:
self.__logger.info("Mayavi is not installed on your device.")
return
#do 2d interpolation
#get slice object
s = np.s_[0:ndim0:1, 0:ndim1:1]
x, y = np.ogrid[s]
mx, my = np.mgrid[s]
#use cubic 2d interpolation
interpfunc = interp2d(x, y, z, kind='cubic')
newx = np.linspace(0, ndim0, 600)
newy = np.linspace(0, ndim1, 600)
newz = interpfunc(newx, newy)
#mlab
face = mlab.surf(newx, newy, newz, warp_scale=2)
mlab.axes(xlabel='x', ylabel='y', zlabel='z')
mlab.outline(face)
#save or show
if show_mode == 'show':
mlab.show()
elif show_mode == 'save':
mlab.savefig('mlab_contour3d.png')
else:
raise ValueError('Unrecognized show mode parameter : ' +
show_mode)
return
def plot_contour3d(self, **kwargs):
'''
use mayavi.mlab to plot 3d contour.
Parameter
---------
kwargs: {
'maxct' : float,max contour number,
'nct' : int, number of contours,
'opacity' : float, opacity of contour,
'widths' : tuple of int
number of replication on x, y, z axis,
}
'''
if not mayavi_installed:
self.__logger.warning("Mayavi is not installed on your device.")
return
# set parameters
widths = kwargs['widths'] if 'widths' in kwargs else (1, 1, 1)
elf_data, grid = self.expand_data(self.elf_data, self.grid, widths)
# import pdb; pdb.set_trace()
maxdata = np.max(elf_data)
maxct = kwargs['maxct'] if 'maxct' in kwargs else maxdata
# check maxct
if maxct > maxdata:
self.__logger.warning("maxct is larger than %f", maxdata)
opacity = kwargs['opacity'] if 'opacity' in kwargs else 0.6
nct = kwargs['nct'] if 'nct' in kwargs else 5
# plot surface
surface = mlab.contour3d(elf_data)
# set surface attrs
surface.actor.property.opacity = opacity
surface.contour.maximum_contour = maxct
surface.contour.number_of_contours = nct
# reverse axes labels
mlab.axes(xlabel='z', ylabel='y', zlabel='x') # 是mlab参数顺序问题?
mlab.outline()
mlab.show()
return
def plot_field(self, **kwargs):
"plot scalar field for elf data"
if not mayavi_installed:
self.__logger.warning("Mayavi is not installed on your device.")
return
# set parameters
vmin = kwargs['vmin'] if 'vmin' in kwargs else 0.0
vmax = kwargs['vmax'] if 'vmax' in kwargs else 1.0
axis_cut = kwargs['axis_cut'] if 'axis_cut' in kwargs else 'z'
nct = kwargs['nct'] if 'nct' in kwargs else 5
widths = kwargs['widths'] if 'widths' in kwargs else (1, 1, 1)
elf_data, grid = self.expand_data(self.elf_data, self.grid, widths)
#create pipeline
field = mlab.pipeline.scalar_field(elf_data) # data source
mlab.pipeline.volume(field, vmin=vmin, vmax=vmax) # put data into volumn to visualize
#cut plane
if axis_cut in ['Z', 'z']:
plane_orientation = 'z_axes'
elif axis_cut in ['Y', 'y']:
plane_orientation = 'y_axes'
elif axis_cut in ['X', 'x']:
plane_orientation = 'x_axes'
cut = mlab.pipeline.scalar_cut_plane(
field.children[0], plane_orientation=plane_orientation)
cut.enable_contours = True # 开启等值线显示
cut.contour.number_of_contours = nct
mlab.show()
#mlab.savefig('field.png', size=(2000, 2000))
return
class ChgCar(ElfCar):
def __init__(self, filename='CHGCAR'):
'''
Create a CHGCAR file class.
Example:
>>> a = ChgCar()
'''
ElfCar.__init__(self, filename)
