Files
VASP_calc/electrochemistry/echem/io_data/qe.py

130 lines
5.3 KiB
Python
Raw Normal View History

2026-01-08 19:47:32 +03:00
import numpy as np
from monty.re import regrep
import itertools
import typing
class QEOutput:
def __init__(self):
self.patterns = {'nkpts': r'number of k points=\s+([\d]+)',
'kpts_coord': r'k\s*=\s*(-?\d.[\d]+)\s*(-?\d.[\d]+)\s*(-?\d.[\d]+)\s*\([\d]+ PWs\)',
'occupations': 'occupation numbers',
'efermi': r'the Fermi energy is\s*(-?[\d]+.[\d]+) ev'}
self.eigenvalues = None
self.weights = None
self.occupations = None
self.efermi = None
self.nkpt = None
@staticmethod
def _GaussianSmearing(x, x0, sigma):
"""Simulate the Delta function by a Gaussian shape function"""
return 1 / (np.sqrt(2 * np.pi) * sigma) * np.exp(-(x - x0) ** 2 / (2 * sigma ** 2))
def from_file(self, filepath):
matches = regrep(filepath, self.patterns)
if len(matches['kpts_coord']) != 0:
with open(filepath, 'rt') as file:
file_data = file.readlines()
eigenvalues = []
for start, end in zip(matches['kpts_coord'], matches['occupations']):
data = file_data[start[1] + 2:end[1] - 1]
data = [float(i) for i in itertools.chain.from_iterable([line.split() for line in data])]
eigenvalues.append(data)
self.eigenvalues = np.array(eigenvalues)
occupations = []
n_strings_occups = matches['occupations'][0][1] - matches['kpts_coord'][0][1] - 1
for start in matches['occupations']:
data = file_data[start[1] + 1: start[1] + n_strings_occups]
data = [float(i) for i in itertools.chain.from_iterable([line.split() for line in data])]
occupations.append(data)
self.occupations = np.array(occupations)
self.efermi = float(matches['efermi'][0][0][0])
self.nkpt = int(matches['nkpts'][0][0][0])
weights = np.zeros(self.nkpt)
for i in range(self.nkpt):
weights[i] = file_data[matches['nkpts'][0][1]+2+i].split()[-1]
self.weights = weights
def get_band_eigs(self, bands):
if type(bands) is int:
return np.array([eig for eig in self.eigenvalues[:, bands]])
elif isinstance(bands, typing.Iterable):
return np.array([[eig for eig in self.eigenvalues[:, band]] for band in bands])
else:
raise ValueError('Variable bands should be int or iterable')
def get_band_occ(self, bands):
if type(bands) is int:
return [occ for occ in self.occupations[:, bands]]
elif isinstance(bands, typing.Iterable):
return np.array([[occ for occ in self.occupations[:, band]] for band in bands])
else:
raise ValueError('Variable bands should be int or iterable')
def get_DOS(self, **kwargs):
"""Calculate Density of States based on eigenvalues and its weights
Args:
dE (float): step of energy array in function output
zero_at_fermi (bool, optional): if True Fermi energy will be equal to zero
sm_param (dict, optional): parameters for smooth DOS.
E_min (float, str): minimum value in DOS calculation. If E_min == 'min' left border of energy
is equal to the minimum eigenvalue
E_max (float, str): maximum value in DOS calculation. If E_max == 'max' right border of energy
is equal to the maximum eigenvalue
bw_method (float): The method used to calculate the estimator bandwidth. This can be 'scott',
'silverman', a scalar constant or a callable. If a scalar, this will be used directly as `kde.factor`.
If a callable, it should take a `gaussian_kde` instance as only parameter and return a scalar.
If None (default), 'scott' is used.
nelec (int): Number of electrons in the system. DOS integral to efermi should be equal to the nelec
Returns:
E, DOS - Two 1D np.arrays that contain energy and according DOS values
"""
if 'zero_at_fermi' in kwargs:
zero_at_fermi = kwargs['zero_at_fermi']
else:
zero_at_fermi = False
if 'dE' in kwargs:
dE = kwargs['dE']
else:
dE = 0.01
if 'smearing' in kwargs:
smearing = kwargs['smearing']
else:
smearing = 'Gaussian'
if smearing == 'Gaussian':
if 'sigma' in kwargs:
sigma = kwargs['sigma']
else:
sigma = 0.02
if 'emin' in kwargs:
E_min = kwargs['emin']
else:
E_min = np.min(self.eigenvalues)
if 'emax' in kwargs:
E_max = kwargs['emax']
else:
E_max = np.max(self.eigenvalues)
E_arr = np.arange(E_min, E_max, dE)
DOS_arr = np.zeros_like(E_arr)
for energy_kpt, weight in zip(self.eigenvalues, self.weights):
for energy in energy_kpt:
DOS_arr += weight * self._GaussianSmearing(E_arr, energy, sigma)
# 2 is not used because sum(weights) = 2
if zero_at_fermi:
return E_arr - self.efermi, DOS_arr
else:
return E_arr, DOS_arr