"""
VASP計算によるバンド構造をプロットするスクリプト。
概要:
VASPの計算結果（OUTCAR, EIGENVAL, KPOINTSなど）からバンド構造を読み込み、
matplotlibを使用してプロットし、関連するバンドデータをCSV/Excelファイルに保存します。
詳細説明:
本スクリプトは、指定されたVASP計算ディレクトリから必要な出力ファイルを解析し、
バンドエネルギー、K点情報、フェルミ準位、バンドギャップなどを抽出します。
抽出された情報に基づいて、バンド構造のグラフを作成し、
選択されたモード（'band', 'bandline', 'bandocc'）に応じて描画方法を調整します。
プロットされたグラフはPNGファイルとして保存でき、バンドデータはCSVまたはExcel形式で出力されます。
関連リンク:
:doc:`vasp_plot_band_usage`
"""
import os
import sys
import shutil
import glob
import csv
import re
import numpy as np
from numpy import exp, log, sin, cos, tan, arcsin, arccos, arctan, pi
from scipy.interpolate import interp1d
from pprint import pprint
from matplotlib import pyplot as plt
from tklib.tkfile import tkFile
import tklib.tkre as tkre
#from tklib.tkutils import save_csv
from tklib.tkvariousdata import tkVariousData
from tklib.tkutils import IsDir, IsFile, SplitFilePath
from tklib.tkutils import terminate, pint, pfloat, getarg, getintarg, getfloatarg
from tklib.tksci.tksci import Reduce01, Round
from tklib.tksci.tkmatrix import make_matrix1, make_matrix2, make_matrix3
from tklib.tkcrystal.tkcif import tkCIF, tkCIFData
from tklib.tkcrystal.tkcrystal import tkCrystal
from tklib.tkcrystal.tkatomtype import tkAtomType
from tklib.tkcrystal.tkvasp import tkVASP
from tklib.tkcrystal.tkbandstructure import plot_band
from tklib.tksci import tkequation
import tklib.tkcsv
#================================
# global parameters
#================================
debug = 0
# mode: 'band', 'bandline', 'bandocc'
validmodes = ['band', 'bandline', 'bandocc']
mode = 'band'
#mode = 'bandocc'
CAR_dir = '.'
# Energy to search edge energies
EF0 = 0.1
# Max number of electrons occupies one state
Nemax = 1.0
# Occupancy threshold to separate HOMO and LUMO
occ_th = 0.5
# Output CSV path
#band_path = 'band.csv'
band_path = 'band.xlsx'
save_figure = 1
plot_figure = 1
#===================================
# figure configuration
#===================================
band_marker_size = 4
band_marker_edge_width = 0.5
Emin = -10.0 # eV
Emax = 10.0 # eV
figsize = (6, 8)
#figsize = (6, 4)
#colors = ['#000000', '#ff0000', '#00aa00', '#0000ff', '#aaaa00', '#ff00ff', '#00ffff', '#aa0000', '#00aa00', '#0000aa']
#colors = ['k', 'r', 'g', 'b', 'y', 'm', 'c']
fontsize = 16
labelfontsize = 12
legend_fontsize = 8
#=============================
# Treat argments
#=============================
[ドキュメント]
def usage():
"""
スクリプトの使用方法を表示します。
詳細説明:
コマンドライン引数として指定可能なモードと引数の例を出力します。
無効なモードが選択されている場合は、デフォルトのモードに設定されます。
Parameters:
なし
Returns:
なし
"""
global mode
if mode not in validmodes:
mode = validmodes[0]
print("")
print("Usage:")
print(" (a) python {} mode CAR_dir EF0 Emin, Emax".format(sys.argv[0]))
print(" mode: {}".format(validmodes))
print(" ex: python {} {} {} {} {} {}".format(sys.argv[0], mode, CAR_dir, EF0, Emin, Emax))
[ドキュメント]
def updatevars():
"""
コマンドライン引数からグローバル変数を更新します。
詳細説明:
sys.argvから引数を読み込み、`mode`、`CAR_dir`、`EF0`、`Emin`、`Emax`、
`occ_th`、`save_figure`、`plot_figure`といったグローバル変数を更新します。
無効なモードが指定された場合は、エラーメッセージを表示してスクリプトを終了します。
Parameters:
なし (sys.argvから引数を読み込む)
Returns:
なし
"""
global mode
global CAR_dir
global EF0, occ_th
global Emin, Emax
global save_figure, plot_figure
mode = getarg (1, mode)
CAR_dir = getarg (2, CAR_dir)
EF0 = getfloatarg(3, EF0)
Emin = getfloatarg(4, Emin)
Emax = getfloatarg(5, Emax)
occ_th = getfloatarg(6, occ_th)
save_figure = getintarg(7, save_figure)
plot_figure = getintarg(8, plot_figure)
if mode == 'band' or mode == 'bandline' or mode == 'bandocc':
pass
else:
terminate("Error: Invalide mode [{}]".format(mode), usage = usage)
[ドキュメント]
def save_bandstructure(band_path, ISPIN, kvec, Eup, Edn):
"""
バンド構造データをCSVまたはExcelファイルに保存します。
詳細説明:
K点情報 (`kvec`) と各バンドのエネルギー情報 (`Eup`, `Edn`) を整理し、
指定されたパス (`band_path`) にCSVまたはExcel形式で出力します。
ISPINが2（スピン偏極計算）の場合、スピンダウンのエネルギーも保存します。
Parameters:
:param band_path: str - 出力ファイルのパス (例: 'band.xlsx', 'band.csv')。
:param ISPIN: int - スピン偏極計算の有無 (1: スピン非偏極, 2: スピン偏極)。
:param kvec: list - 各k点の座標と距離情報を含むリスト。各要素は `[kx, ky, kz, dk, ktot]`。
:param Eup: list - 各k点におけるスピンアップ電子のバンドエネルギーのリスト。
`Eup[ik][iL]` は `ik` 番目のk点の `iL` 番目のバンドエネルギー。
:param Edn: list - 各k点におけるスピンダウン電子のバンドエネルギーのリスト。
`Edn[ik][iL]` は `ik` 番目のk点の `iL` 番目のバンドエネルギー。
ISPIN=1の場合は空のリスト。
Returns:
:returns: bool - ファイル保存が成功した場合は `True`。
"""
nk = len(kvec)
nL = len(Eup[0])
print("")
print("Save band structure data to [{}]".format(band_path))
labels = ["kx", "ky", "kz", "dk", "ktot"]
labels.extend([f'Eup({iL})' for iL in range(nL)])
kv0 = [kvec[ik][0] for ik in range(nk)]
kv1 = [kvec[ik][1] for ik in range(nk)]
kv2 = [kvec[ik][2] for ik in range(nk)]
kv3 = [kvec[ik][3] for ik in range(nk)]
kv4 = [kvec[ik][4] for ik in range(nk)]
EupT = np.array(Eup).T
data_list = [kv0, kv1, kv2, kv3, kv4, *EupT]
if len(Edn) > 0:
labels.extend([f'Edn({iL})' for iL in range(nL)])
EdnT = np.array(Edn).T
# numpy配列としてextendすると、タプルが展開されるので注意。
# list.extend()はイテラブルを展開して要素を追加するため、EdnTが既に`numpy.ndarray`の場合は`*EdnT`で適切に展開される。
data_list.extend(EdnT)
tkVariousData().to_excel(band_path, labels = labels, data_list = data_list)
return True
# 以下はコメントアウトされた元のCSV保存ロジック
# out = tkFile(band_path, 'w')
# if not out:
# print("Erorr in save_bandstructure_csv: Can not read [{}]".format(band_path))
# return 0
# # print("Eup=", Eup)
# out.Write("kx,ky,kz,dk,ktot")
# for iL in range(nL):
# out.Write(",Eup({})".format(iL))
# if ISPIN == 2:
# for iL in range(nL):
# out.Write(",Edn({})".format(iL))
# out.Write("\n")
# for ik in range(nk):
# kv = kvec[ik]
# Eu = Eup[ik]
# out.Write("{},{},{},{},{}".format(kv[0], kv[1], kv[2], kv[3], kv[4]))
# for iL in range(nL):
# out.Write(",{}".format(Eu[iL]))
# if ISPIN == 2:
# Ed = Edn[ik]
# for iL in range(nL):
# out.Write(",{}".format(Ed[iL]))
# out.Write("\n")
# out.Close()
# return 1
[ドキュメント]
def plot_band_structure(mode, CAR_dir, Emin, Emax):
"""
VASPの計算結果からバンド構造をプロットし、関連データを保存します。
詳細説明:
指定されたVASP計算ディレクトリ (`CAR_dir`) から、INCAR, POSCAR, KPOINTS,
CONTCAR, OUTCAR, EIGENVAL, DOSCAR などのファイルを読み込み、
結晶構造情報、計算パラメータ、K点情報、バンドエネルギー、占有率、フェルミ準位、
バンドエッジ（HOMO/LUMO/価電子帯最大値/伝導帯最小値）を抽出します。
これらの情報を用いてmatplotlibでバンド構造のグラフを描画し、
指定されたモード ('band', 'bandline', 'bandocc') に応じてマーカーや線種を調整します。
プロット後、バンド構造データを`band_path`で指定されたファイルに保存し、
`save_figure`が`True`の場合はPNG画像としてグラフを保存します。
`plot_figure`が`True`の場合はグラフウィンドウを表示します。
Parameters:
:param mode: str - プロットモード ('band', 'bandline', 'bandocc')。
:param CAR_dir: str - VASPの計算結果が格納されているディレクトリのパス。
:param Emin: float - プロットするエネルギー範囲の下限 (eV)。
:param Emax: float - プロットするエネルギー範囲の上限 (eV)。
Returns:
なし
"""
vasp = tkVASP()
base_path = vasp.getdir(CAR_dir)
print("")
INCAR_path = vasp.get_INCAR(base_path)
POSCAR_path = vasp.get_POSCAR(base_path)
KPOINTS_path = vasp.get_VASPPath(base_path, 'KPOINTS')
CONTCAR_path = vasp.get_CONTCAR(base_path)
OUTCAR_path = vasp.get_OUTCAR(base_path)
EIGENVAL_path = vasp.get_VASPPath(base_path, 'EIGENVAL')
DOSCAR_path = vasp.get_VASPPath(base_path, 'DOSCAR')
print("CAR dir(ideal) : ", CAR_dir)
print(" INCAR : ", INCAR_path)
print(" POSCAR : ", POSCAR_path)
print(" KPOINTS : ", KPOINTS_path)
print(" CONTCAR : ", CONTCAR_path)
print(" OUTCAR : ", OUTCAR_path)
print(" EIGENVAL: ", EIGENVAL_path)
print(" DOSCAR : ", DOSCAR_path)
print("")
print(f"Occupancy threshold to seprate HOMO and LUMO: {occ_th}")
print(f"save_figure: {save_figure}")
print(f"plot_figure: {plot_figure}")
print("")
print("EF0: {} eV".format(EF0))
print("")
print("*** Read crystal structure from [{}]".format(POSCAR_path))
cry1 = vasp.read_poscar(POSCAR_path)
if cry1 is None:
terminate("Error: Can not read [{}]".format(POSCAR_path), usage = usage)
a1, b1, c1, alpha1, beta1, gamm1 = cry1.LatticeParameters()
cry1.PrintInf("cell")
print("")
incarinf = vasp.read_incar_inf(INCAR_path)
print("")
outcarinf = vasp.read_outcar_inf(OUTCAR_path)
EF = outcarinf.get("EF", None)
ISPIN = outcarinf.get("ISPIN", 1)
IsHF = outcarinf.get("IsHF", 0)
IsMETAGGA = outcarinf.get("IsMETAGGA", 0)
print("Information in OUTCAR:")
print(f" ISPIN : {ISPIN}")
print(f" IsHF : {IsHF}")
print(f" IsMETAGGA: {IsMETAGGA}")
print(f" EF = {EF} eV corrected to 0")
if IsMETAGGA:
IsHF = 1
kpoints_inf = vasp.read_kpoints(KPOINTS_path)
# print("kpoints: ", kpoints_inf)
# print("kpoints: ", kpoints_inf['kpoints'])
nk = kpoints_inf['nk']
kpoints = kpoints_inf['kpoints']
kptdic = kpoints_inf['kpointsdic']
print("")
print("Special k points in KPOINTS:")
if IsHF == 0: # IsHF == 0 の時だけkname0_convが使われる？条件が不明なのでそのまま残す。
kp = kptdic[0]
kx0, ky0, kz0, kname0 = kp["kx0"], kp["ky0"], kp["kz0"], kp["kname0_conv"]
kx1, ky1, kz1, kname1 = kp["kx1"], kp["ky1"], kp["kz1"], kp["kname1_conv"]
w, dk, ktot = kp["kw"], 0.0, 0.0
if kname0 == '':
kname0 = 'no_name'
print(" {:10} ({:8.4f} {:8.4f} {:8.4f}) w={:.3f} {:8.4f} {:12.4f}".format(kname0, kx0, ky0, kz0, w, dk, ktot))
for i in range(len(kptdic)):
kp = kptdic[i]
kx0, ky0, kz0, kname0 = kp["kx0"], kp["ky0"], kp["kz0"], kp["kname0_conv"]
kx1, ky1, kz1, kname1 = kp["kx1"], kp["ky1"], kp["kz1"], kp["kname1_conv"]
w, dk, ktot = kp["kw"], kp["dk"], kp["ktot"]
if kname1 == '':
kname1 = 'no_name'
print(" {:10} ({:8.4f} {:8.4f} {:8.4f}) w={:.3f} {:8.4f} {:12.4f}".format(kname1, kx1, ky1, kz1, w, dk, ktot))
print("")
print("k points in KPOINTS:")
if IsHF == 0 and IsMETAGGA == 0: # 同上、kname0_convの扱いが不明。そのまま残す。
kp = kpoints[0]
kx0, ky0, kz0, kname0 = kp["kx0"], kp["ky0"], kp["kz0"], kp["kname0_conv"]
kx1, ky1, kz1, kname1 = kp["kx1"], kp["ky1"], kp["kz1"], kp["kname1_conv"]
w, dk, ktot = kp["kw"], 0.0, 0.0
print(" {:10} ({:8.4f} {:8.4f} {:8.4f}) w={:.3f} {:8.4f} {:12.4f}".format(kname0, kx0, ky0, kz0, w, dk, ktot))
for i in range(nk):
kp = kpoints[i]
kx0, ky0, kz0, kname0 = kp["kx0"], kp["ky0"], kp["kz0"], kp["kname0_conv"]
kx1, ky1, kz1, kname1 = kp["kx1"], kp["ky1"], kp["kz1"], kp["kname1_conv"]
w, dk, ktot = kp["kw"], kp["dk"], kp["ktot"]
print(" {:10} ({:8.4f} {:8.4f} {:8.4f}) w={:.3f} {:8.4f} {:12.4f}".format(kname1, kx1, ky1, kz1, w, dk, ktot))
eigenvalinf = vasp.read_eigenval(EIGENVAL_path, KPOINTS_inf = kpoints_inf, ISPIN = ISPIN, IsHF = IsHF, EF = EF)
# print("eigenvalinf=", eigenvalinf)
nk = eigenvalinf["nk"]
nLevels = eigenvalinf["nLevels"]
EList = eigenvalinf['EList']
nk_band = len(EList)
print("")
print("k points in EIGENVAL:")
print("nk(band)=", nk_band)
print("nLevels =", nLevels)
for i in range(nk_band):
el = EList[i]
kx, ky, kz, wk, dk, ktot, Eups, occups, Edns, occdns = el
print(f" ({kx:8.4f} {ky:8.4f} {kz:8.4f}) w={wk:8.4g} dk={dk:8.4f} ktot={ktot:8.4f}")
bandedgeinf = vasp.find_band_edges_from_eigenval(EF0 = EF0, eigenvalinf = eigenvalinf, ISPIN = ISPIN, occ_th = occ_th)
print("bandedgeinf=", bandedgeinf)
bandedgeinf2 = vasp.gbandedges(CAR_dir, OUTCAR_path, EIGENVAL_path, EF)
EV1 = bandedgeinf["EV"]
EC1 = bandedgeinf["EC"]
Eg1 = bandedgeinf["Eg"]
EV = bandedgeinf2["EVBM"] - EF
EC = bandedgeinf2["ECBM"] - EF
Eg = bandedgeinf2["Eg"]
EHOMO = bandedgeinf["EHOMO"]
ELUMO = bandedgeinf["ELUMO"]
print("")
print(f"Band edge from EIGENVAL:")
print(f"EF from OUTCAR: {EF:10.6f} eV")
print(f"EF0={EF0:10.6f} eV")
print(f" find_band_edges: EV={EV1:10.6f} EC={EC1:10.6f} Eg={Eg:10.6f} eV")
print(f" HOMO:{EHOMO:10.6f} eV")
print(f" LUMO:{ELUMO:10.6f} eV")
print(f" gbandedges() : EV={EV:10.6f} EC={EC:10.6f} Eg={Eg:10.6f} eV")
# print(f" EHOMO={EHOMO+EF:10.6f} ELUMO={ELUMO+EF:10.6f}")
#=============================
# Plot graphs
#=============================
print("")
print("plot band structure")
fig = plt.figure(figsize = figsize)
axband = fig.add_subplot(1, 1, 1)
# バンド構造をプロット
# xk: プロットするk点の蓄積距離のリスト
# yE: E(xk[i])。入れ子になったリストで構わない
# ktotallist: k点境界における、最初のk点からの距離の和のリスト
# ktotal_namelist: k点境界における、k点の名称
#def plot_band(axis, xk, yE, Erange, ktotallist, ktotal_namelist):
if len(kptdic) == 0:
ktotallist = []
ktotal_namelist = []
else:
kp = kptdic[0]
kname0 = kp["kname0"]
ktotallist = [0.0]
ktotal_namelist = [vasp.convert_kname(kname0)]
for i in range(len(kptdic)):
kp = kptdic[i]
kx0, ky0, kz0, kname0 = kp["kx0"], kp["ky0"], kp["kz0"], kp["kname0"]
kx1, ky1, kz1, kname1 = kp["kx1"], kp["ky1"], kp["kz1"], kp["kname1"]
ktot = kp["ktot"]
ktotallist.append(ktot)
ktotal_namelist.append(vasp.convert_kname(kname1))
# exit()
xk = []
kvec = []
yEup = []
yNup = []
yEdn = []
yNdn = []
EList = eigenvalinf['EList']
for i in range(len(EList)):
el = EList[i]
kx0, ky0, kz0, wk, dk, ktot, Eups, occups, Edns, occdns = [*el]
xk.append(ktot)
kvec.append([kx0, ky0, kz0, dk, ktot])
yEup.append(Eups)
yNup.append(occups)
if ISPIN == 2:
yEdn.append(Edns)
yNdn.append(occdns)
# CSVファイルを保存
save_bandstructure(band_path, ISPIN, kvec, yEup, yEdn)
# HF/METAGGA計算の場合は、BZ端の線、k点の名称をプロットしない
# if IsHF or IsMETAGGA:
# ktotallist = None
# ktotal_namelist = None
if mode == 'bandocc':
plot_band(plt, axband, ISPIN, xk, yEup, yEdn, [Emin, Emax], occups = yNup, occdns = yNdn,
ktotallist = ktotallist, ktotal_namelist = ktotal_namelist,
EV = EV, EC = EC, EF = EF0, EHOMO = EHOMO, ELUMO = ELUMO,
EFlabel = '$E_{F0}$',
marker = 'o', linestyle = '',
markersize = band_marker_size, markeredgewidth = band_marker_edge_width,
legendloc = 'lower right')
elif mode == 'bandline':
plot_band(plt, axband, ISPIN, xk, yEup, yEdn, [Emin, Emax],
ktotallist = ktotallist, ktotal_namelist = ktotal_namelist,
EV = EV, EC = EC, EF = EF0, EHOMO = EHOMO, ELUMO = ELUMO,
EFlabel = '$E_{F0}$',
marker = None, linestyle = '-',
legendloc = 'lower right')
elif mode == 'band':
plot_band(plt, axband, ISPIN, xk, yEup, yEdn, [Emin, Emax],
ktotallist = ktotallist, ktotal_namelist = ktotal_namelist,
EV = EV, EC = EC, EF = EF0, EHOMO = EHOMO, ELUMO = ELUMO,
EFlabel = '$E_{F0}$',
marker = 'o', linestyle = '',
markersize = band_marker_size, markeredgewidth = band_marker_edge_width,
legendloc = 'lower right')
# Rearange the graph axes so that they are not overlapped
plt.tight_layout()
if save_figure:
outfig_path = 'band.png'
print("")
print(f"Save band structure to figure file [{outfig_path}]")
plt.savefig(outfig_path, dpi=300, transparent=True)
use_pause = True
if plot_figure:
if use_pause:
plt.pause(0.1)
print("")
input("Press ENTER to exit>>")
else:
plt.show()
print("")
print("Close graph window to terminate")
terminate("", usage = usage)
[ドキュメント]
def main():
"""
スクリプトのメイン処理を実行します。
詳細説明:
コマンドライン引数に基づいてグローバル変数を更新し、
指定されたプロットモード (`mode`) に応じてバンド構造のプロット関数を呼び出します。
無効なモードが指定された場合はエラーで終了します。
Parameters:
なし
Returns:
なし
"""
updatevars()
print("")
print("=============== Plot band structure calculated by VASP ============")
print("")
print("mode: ", mode)
if mode == 'band' or mode == 'bandline' or mode == 'bandocc':
plot_band_structure(mode, CAR_dir, Emin, Emax)
else:
terminate("Error: Invalide mode [{}]".format(mode), usage = usage)
if __name__ == "__main__":
main()