import numpy as np
import pickle
from ase.units import Hartree
from gpaw import GPAW
from gpaw.kpt_descriptor import to1bz
from gpaw.spinorbit import soc_eigenstates
from gpaw.pw.descriptor import PWDescriptor
import gpaw.mpi as mpi
[docs]class Unfold:
"""This Class is used to Unfold the Bands of a supercell (SC) calculations
into a the primitive cell (PC). As a convention (when possible) capital
letters variables are related to the SC while lowercase ones to the
PC """
def __init__(self,
name=None,
calc=None,
M=None,
spinorbit=None):
self.name = name
self.calc = GPAW(calc, txt=None, communicator=mpi.serial_comm)
self.M = np.array(M, dtype=float)
self.spinorbit = spinorbit
self.gd = self.calc.wfs.gd.new_descriptor()
self.kd = self.calc.wfs.kd
if self.calc.wfs.mode == 'pw':
self.pd = self.calc.wfs.pd
else:
self.pd = PWDescriptor(ecut=None, gd=self.gd, kd=self.kd,
dtype=complex)
self.acell_cv = self.gd.cell_cv
self.bcell_cv = 2 * np.pi * self.gd.icell_cv
self.nb = self.calc.get_number_of_bands()
self.v_Kmsn = None
if spinorbit:
if mpi.world.rank == 0:
print('Calculating spinorbit Corrections')
self.nb = 2 * self.calc.get_number_of_bands()
soc = soc_eigenstates(self.calc)
self.e_mK = soc.eigenvalues().T
self.v_Kmsn = soc.eigenvectors()
if mpi.world.rank == 0:
print('Done with the spinorbit Corrections')
def get_K_index(self, K):
"""Find the index of a given K."""
K = np.array([K])
bzKG = to1bz(K, self.acell_cv)[0]
iK = self.kd.where_is_q(bzKG, self.kd.bzk_kc)
return iK
def get_g(self, iK):
"""Not all the G vectors are relevant for the bands unfolding,
but only the ones that match the PC reciprocal vectors.
This function finds the relevant ones."""
G_Gv_temp = self.pd.get_reciprocal_vectors(q=iK, add_q=False)
G_Gc_temp = np.dot(G_Gv_temp, np.linalg.inv(self.bcell_cv))
iG_list = []
g_list = []
for iG, G in enumerate(G_Gc_temp):
a = np.dot(G, np.linalg.inv(self.M).T)
check = np.abs(a) % 1 < 1e-5
check2 = np.abs((np.abs(a[np.where(~check)]) % 1) - 1) < 1e-5
if all(check) or all(check2):
iG_list.append(iG)
g_list.append(G)
return np.array(iG_list), np.array(g_list)
def get_G_index(self, iK, G, G_list):
"""Find the index of a given G."""
G_list -= G
sumG = np.sum(abs(G_list), axis=1)
iG = np.where(sumG < 1e-5)[0]
return iG
def get_eigenvalues(self, iK):
"""Get the list of eigenvalues for a given iK."""
if not self.spinorbit:
e_m = self.calc.get_eigenvalues(kpt=iK, spin=0) / Hartree
else:
e_m = self.e_mK[:, iK] / Hartree
return np.array(e_m)
def get_pw_wavefunctions_k(self, iK):
"""Get the list of Fourier coefficients of the WaveFunction for a
given iK. For spinors the number of bands is doubled and a spin
dimension is added."""
psi_mgrid = get_rs_wavefunctions_k(self.calc, iK, self.spinorbit,
self.v_Kmsn)
if not self.spinorbit:
psi_list_mG = []
for i in range(len(psi_mgrid)):
psi_list_mG.append(self.pd.fft(psi_mgrid[i], iK))
psi_mG = np.array(psi_list_mG)
return psi_mG
else:
u0_list_mG = []
u1_list_mG = []
for i in range(psi_mgrid.shape[0]):
u0_list_mG.append(self.pd.fft(psi_mgrid[i, 0], iK))
u1_list_mG.append(self.pd.fft(psi_mgrid[i, 1], iK))
u0_mG = np.array(u0_list_mG)
u1_mG = np.array(u1_list_mG)
u_mG = np.zeros((len(u0_mG),
2,
u0_mG.shape[1]), complex)
u_mG[:, 0] = u0_mG
u_mG[:, 1] = u1_mG
return u_mG
def get_spectral_weights_k(self, k_t):
r"""Returns the spectral weights for a given k in the PC:
P_mK(k_t) = \sum_n |<Km|k_t n>|**2
which can be shown to be equivalent to:
P_mK(k_t) = \sum_g |C_Km(g+k_t-K)|**2
"""
K_c, G_t = find_K_from_k(k_t, self.M)
iK = self.get_K_index(K_c)
iG_list, g_list = self.get_g(iK)
gG_t_list = g_list + G_t
G_Gv = self.pd.get_reciprocal_vectors(q=iK, add_q=False)
G_Gc = np.dot(G_Gv, np.linalg.inv(self.bcell_cv))
igG_t_list = []
for g in gG_t_list:
igG_t_list.append(self.get_G_index(iK, g.copy(), G_Gc.copy()))
C_mG = self.get_pw_wavefunctions_k(iK)
P_m = []
if not self.spinorbit:
for m in range(self.nb):
P = 0.
norm = np.sum(np.linalg.norm(C_mG[m, :])**2)
for iG in igG_t_list:
P += np.linalg.norm(C_mG[m, iG])**2
P_m.append(P / norm)
else:
for m in range(self.nb):
P = 0.
norm = np.sum(np.linalg.norm(C_mG[m, 0, :])**2 +
np.linalg.norm(C_mG[m, 1, :])**2)
for iG in igG_t_list:
P += (np.linalg.norm(C_mG[m, 0, iG])**2 +
np.linalg.norm(C_mG[m, 1, iG])**2)
P_m.append(P / norm)
return np.array(P_m)
def get_spectral_weights(self, kpoints, filename=None):
"""Collect the spectral weights for the k points in the kpoints list.
This function is parallelized over k's."""
Nk = len(kpoints)
Nb = self.nb
world = mpi.world
if filename is None:
try:
e_mK, P_mK = pickle.load(open('weights_' + self.name +
'.pckl', 'rb'))
except OSError:
e_Km = []
P_Km = []
if world.rank == 0:
print('Getting EigenValues and Weights')
e_Km = np.zeros((Nk, Nb))
P_Km = np.zeros((Nk, Nb))
myk = range(0, Nk)[world.rank::world.size]
for ik in myk:
k = kpoints[ik]
print('kpoint: %s' % k)
K_c, G_c = find_K_from_k(k, self.M)
iK = self.get_K_index(K_c)
e_Km[ik] = self.get_eigenvalues(iK)
P_Km[ik] = self.get_spectral_weights_k(k)
world.barrier()
world.sum(e_Km)
world.sum(P_Km)
e_mK = np.array(e_Km).T
P_mK = np.array(P_Km).T
if world.rank == 0:
pickle.dump((e_mK, P_mK),
open('weights_' + self.name + '.pckl', 'wb'))
else:
e_mK, P_mK = pickle.load(open(filename, 'rb'))
return e_mK, P_mK
def spectral_function(self, kpts, x, X, points_name, width=0.002,
npts=10000, filename=None):
r"""Returns the spectral function for all the ks in kpoints:
eta / pi
A_k(e) = \sum_m P_mK(k) x ---------------------
(e - e_mk)**2 + eta**2
at each k-points defined on npts energy points in the range
[emin, emax]. The width keyword is FWHM = 2 * eta."""
Nk = len(kpts)
A_ke = np.zeros((Nk, npts), float)
world = mpi.world
e_mK, P_mK = self.get_spectral_weights(kpts, filename)
if world.rank == 0:
print('Calculating the Spectral Function')
emin = np.min(e_mK) - 5 * width
emax = np.max(e_mK) + 5 * width
e = np.linspace(emin, emax, npts)
for ik in range(Nk):
for ie in range(len(e_mK[:, ik])):
e0 = e_mK[ie, ik]
D = (width / 2 / np.pi) / ((e - e0)**2 + (width / 2)**2)
A_ke[ik] += P_mK[ie, ik] * D
if world.rank == 0:
pickle.dump((e * Hartree, A_ke, x, X, points_name),
open('sf_' + self.name + '.pckl', 'wb'))
print('Spectral Function calculation completed!')
return
def find_K_from_k(k, M):
"""Gets a k vector in scaled coordinates and returns a K vector and the
unfolding G in scaled Coordinates."""
KG = np.dot(M, k)
G = np.zeros(3, dtype=int)
for i in range(3):
if KG[i] > 0.5000001:
G[i] = int(np.round(KG[i]))
KG[i] -= np.round(KG[i])
elif KG[i] < -0.4999999:
G[i] = int(np.round(KG[i]))
KG[i] += abs(np.round(KG[i]))
return KG, G
def get_rs_wavefunctions_k(calc, iK, spinorbit=False, v_Kmsn=None):
"""Get the list of WaveFunction for a given iK. For spinors the number of
bands is doubled and a spin dimension is added."""
N_c = calc.wfs.gd.N_c
k_c = calc.wfs.kd.ibzk_kc[iK]
Nb = calc.wfs.bd.nbands
Ns = calc.wfs.nspins
eikr_R = np.exp(-2j * np.pi * np.dot(np.indices(N_c).T,
k_c / N_c).T)
if calc.wfs.mode == 'lcao' and not calc.wfs.positions_set:
calc.initialize_positions()
if not spinorbit:
psit_mgrid = np.array([calc.wfs.get_wave_function_array(m, iK, 0) *
eikr_R for m in range(Nb)])
return psit_mgrid
else:
v_msn = v_Kmsn[iK]
v0_mn = v_msn[:, 0]
v1_mn = v_msn[:, 1]
u0_ngrid = np.array(
[calc.wfs.get_wave_function_array(n, iK, 0) * eikr_R
for n in range(Nb)])
u1_ngrid = np.array(
[calc.wfs.get_wave_function_array(n, iK, (Ns - 1)) * eikr_R
for n in range(Nb)])
u0_mG = np.swapaxes(np.dot(v0_mn, np.swapaxes(u0_ngrid, 0, 2)), 1, 2)
u1_mG = np.swapaxes(np.dot(v1_mn, np.swapaxes(u1_ngrid, 0, 2)), 1, 2)
ut_mgrid = np.zeros((len(u0_mG),
2,
len(u0_mG[0]),
len(u0_mG[0, 0]),
len(u0_mG[0, 0, 0])), complex)
ut_mgrid[:, 0] = u0_mG
ut_mgrid[:, 1] = u1_mG
return ut_mgrid
[docs]def plot_spectral_function(filename, color='blue', eref=None,
emin=None, emax=None):
"""Function to plot spectral function corresponding to the bandstructure
along the kpoints path."""
try:
e, A_ke, x, X, points_name = pickle.load(open(filename + '.pckl',
'rb'))
except OSError:
print('You Need to Calculate the SF first!')
raise SystemExit()
import matplotlib.pyplot as plt
print('Plotting Spectral Function')
if eref is not None:
e -= eref
if emin is None:
emin = e.min()
if emax is None:
emax = e.max()
A_ke /= np.max(A_ke)
A_ek = A_ke.T
A_ekc = np.reshape(A_ek, (A_ek.shape[0], A_ek.shape[1]))
mycmap = make_colormap(color)
plt.figure()
plt.plot([0, x[-1]], 2 * [0.0], '--', c='0.5')
plt.imshow(A_ekc + 0.23,
cmap=mycmap,
aspect='auto',
origin='lower',
vmin=0.,
vmax=1,
extent=[0, x[-1], e.min(), e.max()])
for k in X[1:-1]:
plt.plot([k, k], [emin, emax], lw=0.5, c='0.5')
plt.xticks(X, points_name, size=20)
plt.yticks(size=20)
plt.ylabel('E(eV)', size=20)
plt.axis([0, x[-1], emin, emax])
plt.savefig(filename + '_spec.png')
plt.show()
def plot_band_structure(e_mK, P_mK, x, X, points_name,
weights_mK=None, color='red', fit=True, nfit=200):
"""Function to plot the bandstructure using the P_mK weights directly.
each point is represented with a filled circle, whose size and color
vary with as a function of P_mK."""
import matplotlib.pyplot as plt
print('Plotting Bands Structure')
emin = e_mK.min()
emax = e_mK.max()
new_cmap = make_colormap(color)
if weights_mK is None:
weights_mK = P_mK.copy()
else:
weights_mK *= P_mK.copy()
plt.figure()
plt.plot([0, x[-1]], 2 * [0.0], '--', c='0.5')
plt.scatter(np.tile(x, len(e_mK)), e_mK.reshape(-1),
c=P_mK.reshape(-1),
cmap=new_cmap,
vmin=0.,
vmax=1.,
s=20. * weights_mK.reshape(-1),
marker='o',
edgecolor='none')
for k in X[1:-1]:
plt.plot([k, k], [emin, emax], lw=0.5, c='0.5')
plt.xticks(X, points_name, size=20)
plt.yticks(size=20)
plt.ylabel('E(eV)', size=20)
plt.axis([0, x[-1], emin, emax])
plt.show()
def make_colormap(main_color):
"""Custom colormaps used in plot_spectral function and
plot_band_structure."""
from matplotlib.colors import LinearSegmentedColormap
if main_color == 'blue':
cdict = {'red': ((0.0, 1.0, 1.0),
(0.25, 1.0, 1.0),
(0.5, 0.0, 0.0),
(1.0, 0.0, 0.0)),
'green': ((0.0, 1.0, 1.0),
(0.25, 1.0, 1.0),
(0.5, 0.0, 0.0),
(1.0, 0.0, 0.0)),
'blue': ((0.0, 1.0, 1.0),
(0.25, 1.0, 1.0),
(0.5, 1.0, 1.0),
(1.0, 0.75, 0.75))}
elif main_color == 'red':
cdict = {'red': ((0.0, 1.0, 1.0),
(0.25, 1.0, 1.0),
(0.5, 1.0, 1.0),
(1.0, 0.75, 0.75)),
'green': ((0.0, 1.0, 1.0),
(0.25, 1.0, 1.0),
(0.5, 0.0, 0.0),
(1.0, 0.0, 0.0)),
'blue': ((0.0, 1.0, 1.0),
(0.25, 1.0, 1.0),
(0.5, 0.0, 0.0),
(1.0, 0.0, 0.0))}
elif main_color == 'green':
cdict = {'red': ((0.0, 1.0, 1.0),
(0.25, 1.0, 1.0),
(0.5, 0.0, 0.0),
(1.0, 0.0, 0.0)),
'green': ((0.0, 1.0, 1.0),
(0.25, 1.0, 1.0),
(0.5, 1.0, 1.0),
(1.0, 0.75, 0.75)),
'blue': ((0.0, 1.0, 1.0),
(0.25, 1.0, 1.0),
(0.5, 0.0, 0.0),
(1.0, 0.0, 0.0)),
'alpha': ((0.0, 0.0, 0.0),
(0.25, 0.1, 0.1),
(0.5, 1.0, 1.0),
(1.0, 1.0, 1.0))}
cmap = LinearSegmentedColormap('mymap', cdict)
return cmap
def get_vacuum_level(calc, plot_pot=False):
"""Get the vacuum energy level from a given calculator."""
calc.restore_state()
if calc.wfs.mode == 'pw':
vHt_g = calc.hamiltonian.pd3.ifft(calc.hamiltonian.vHt_q) * Hartree
else:
vHt_g = calc.hamiltonian.vHt_g * Hartree
vHt_z = np.mean(np.mean(vHt_g, axis=0), axis=0)
if plot_pot:
import matplotlib.pyplot as plt
plt.plot(vHt_z)
plt.show()
return vHt_z[0]