Skip to content

Commit

Permalink
Generalize fatband plots from Lobster (#3688)
Browse files Browse the repository at this point in the history 
* fatband fix

* pre-commit auto-fixes

* update plot

* pre-commit auto-fixes

* Fix ruff

* Fix linting

* Fix linting

* pre-commit auto-fixes

* Fix linting

* Fix linting

* pre-commit auto-fixes

* Fix linting

* pre-commit auto-fixes

* Type

* Type

* pre-commit auto-fixes

* new test

* Review comments

* pre-commit auto-fixes

* linting

* remove redundant Lobsterout.ATTRIBUTE_DEFAULTS

* fix: handle kwargs and filename missing case

* Fix typo in Fatband init docs

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Janosh Riebesell <janosh.riebesell@gmail.com>
  • Loading branch information
@JaGeo @pre-commit-ci @janosh
3 people committed Mar 14, 2024
1 parent 668fa57 commit 72977a9
Show file tree
Hide file tree
Showing 3 changed files with 139 additions and 64 deletions.
109 changes: 64 additions & 45 deletions pymatgen/io/lobster/outputs.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -823,24 +823,23 @@ class Lobsterout(MSONable):
"has_grosspopulation",
"has_density_of_energies",
)
ATTRIBUTE_DEFAULTS = dict.fromkeys(_ATTRIBUTES, None)

# TODO: add tests for skipping COBI and madelung
# TODO: add tests for including COBI and madelung
def __init__(self, filename: str | None, **kwargs) -> None:
"""
Args:
filename: filename of lobsterout.
**kwargs:dict to initialize Lobsterout instance (see > ATTRIBUTE_DEFAULTS)
**kwargs: dict to initialize Lobsterout instance
"""
self.filename = filename
if kwargs:
for attr, val in kwargs.items():
if attr in self.ATTRIBUTE_DEFAULTS:
if attr in self._ATTRIBUTES:
setattr(self, attr, val)
else:
raise ValueError(f"{attr}={val} is not a valid attribute for Lobsterout")
else:
elif filename:
with zopen(filename, mode="rt") as file: # read in file
data = file.read().split("\n")
if len(data) == 0:
Expand Down Expand Up @@ -909,6 +908,8 @@ def __init__(self, filename: str | None, **kwargs) -> None:
"writing SitePotentials.lobster and MadelungEnergies.lobster..." in data
and "skipping writing SitePotentials.lobster and MadelungEnergies.lobster..." not in data
)
else:
raise ValueError("must provide either filename or kwargs to initialize Lobsterout")

def get_doc(self):
"""Returns: LobsterDict with all the information stored in lobsterout."""
Expand Down Expand Up @@ -1116,62 +1117,80 @@ class Fatband:
The indices of the array are [band_index, kpoint_index].
The dict is then built the following way: {"string of element": "string of orbital as read in
from FATBAND file"}. If the band structure is not spin polarized, we only store one data set under Spin.up.
structure (Structure): Structure read in from vasprun.xml.
structure (Structure): Structure read in from Structure object.
"""

def __init__(self, filenames=".", vasprun="vasprun.xml", Kpointsfile="KPOINTS"):
def __init__(
self,
filenames: str | list = ".",
kpoints_file: str = "KPOINTS",
vasprun_file: str | None = "vasprun.xml",
structure: Structure | IStructure | None = None,
efermi: float | None = None,
):
"""
Args:
filenames (list or string): can be a list of file names or a path to a folder from which all
"FATBAND_*" files will be read
vasprun: corresponding vasprun file
Kpointsfile: KPOINTS file for bandstructure calculation, typically "KPOINTS".
kpoints_file (str): KPOINTS file for bandstructure calculation, typically "KPOINTS".
vasprun_file (str): Corresponding vasprun file.
Instead, the Fermi energy from the DFT run can be provided. Then,
this value should be set to None.
structure (Structure): Structure object.
efermi (float): fermi energy in eV
"""
warnings.warn("Make sure all relevant FATBAND files were generated and read in!")
warnings.warn("Use Lobster 3.2.0 or newer for fatband calculations!")

vasp_run = Vasprun(
filename=vasprun,
ionic_step_skip=None,
ionic_step_offset=0,
parse_dos=True,
parse_eigen=False,
parse_projected_eigen=False,
parse_potcar_file=False,
occu_tol=1e-8,
exception_on_bad_xml=True,
)
self.structure = vasp_run.final_structure
if structure is None:
raise ValueError("A structure object has to be provided")
self.structure = structure
if vasprun_file is None and efermi is None:
raise ValueError("vasprun_file or efermi have to be provided")

self.lattice = self.structure.lattice.reciprocal_lattice
self.efermi = vasp_run.efermi
kpoints_object = Kpoints.from_file(Kpointsfile)
if vasprun_file is not None:
self.efermi = Vasprun(
filename=vasprun_file,
ionic_step_skip=None,
ionic_step_offset=0,
parse_dos=True,
parse_eigen=False,
parse_projected_eigen=False,
parse_potcar_file=False,
occu_tol=1e-8,
exception_on_bad_xml=True,
).efermi
else:
self.efermi = efermi
kpoints_object = Kpoints.from_file(kpoints_file)

atomtype = []
atomnames = []
atom_type = []
atom_names = []
orbital_names = []

if not isinstance(filenames, list) or filenames is None:
filenames_new = []
if filenames is None:
filenames = "."
for file in os.listdir(filenames):
if fnmatch.fnmatch(file, "FATBAND_*.lobster"):
filenames_new.append(os.path.join(filenames, file))
for name in os.listdir(filenames):
if fnmatch.fnmatch(name, "FATBAND_*.lobster"):
filenames_new.append(os.path.join(filenames, name))
filenames = filenames_new
if len(filenames) == 0:
raise ValueError("No FATBAND files in folder or given")
for filename in filenames:
with zopen(filename, mode="rt") as file:
for name in filenames:
with zopen(name, mode="rt") as file:
contents = file.read().split("\n")

atomnames.append(os.path.split(filename)[1].split("_")[1].capitalize())
atom_names.append(os.path.split(name)[1].split("_")[1].capitalize())
parameters = contents[0].split()
atomtype.append(re.split(r"[0-9]+", parameters[3])[0].capitalize())
atom_type.append(re.split(r"[0-9]+", parameters[3])[0].capitalize())
orbital_names.append(parameters[4])

# get atomtype orbital dict
atom_orbital_dict = {}
for iatom, atom in enumerate(atomnames):
atom_orbital_dict = {} # type: dict
for iatom, atom in enumerate(atom_names):
if atom not in atom_orbital_dict:
atom_orbital_dict[atom] = []
atom_orbital_dict[atom].append(orbital_names[iatom])
Expand Down Expand Up @@ -1213,37 +1232,37 @@ def __init__(self, filenames=".", vasprun="vasprun.xml", Kpointsfile="KPOINTS"):
self.is_spinpolarized = len(linenumbers) == 2

if ifilename == 0:
eigenvals = {}
eigenvals = {} # type: dict
eigenvals[Spin.up] = [
[collections.defaultdict(float) for i in range(self.number_kpts)] for j in range(self.nbands)
[collections.defaultdict(float) for _ in range(self.number_kpts)] for _ in range(self.nbands)
]
if self.is_spinpolarized:
eigenvals[Spin.down] = [
[collections.defaultdict(float) for i in range(self.number_kpts)] for j in range(self.nbands)
[collections.defaultdict(float) for _ in range(self.number_kpts)] for _ in range(self.nbands)
]

p_eigenvals = {}
p_eigenvals = {} # type: dict
p_eigenvals[Spin.up] = [
[
{
str(e): {str(orb): collections.defaultdict(float) for orb in atom_orbital_dict[e]}
for e in atomnames
for e in atom_names
}
for i in range(self.number_kpts)
for _ in range(self.number_kpts)
]
for j in range(self.nbands)
for _ in range(self.nbands)
]

if self.is_spinpolarized:
p_eigenvals[Spin.down] = [
[
{
str(e): {str(orb): collections.defaultdict(float) for orb in atom_orbital_dict[e]}
for e in atomnames
for e in atom_names
}
for i in range(self.number_kpts)
for _ in range(self.number_kpts)
]
for j in range(self.nbands)
for _ in range(self.nbands)
]

ikpoint = -1
Expand All @@ -1269,13 +1288,13 @@ def __init__(self, filenames=".", vasprun="vasprun.xml", Kpointsfile="KPOINTS"):
if ifilename == 0:
eigenvals[Spin.up][iband][ikpoint] = float(line.split()[1]) + self.efermi

p_eigenvals[Spin.up][iband][ikpoint][atomnames[ifilename]][orbital_names[ifilename]] = float(
p_eigenvals[Spin.up][iband][ikpoint][atom_names[ifilename]][orbital_names[ifilename]] = float(
line.split()[2]
)
if linenumber >= self.nbands and self.is_spinpolarized:
if ifilename == 0:
eigenvals[Spin.down][iband][ikpoint] = float(line.split()[1]) + self.efermi
p_eigenvals[Spin.down][iband][ikpoint][atomnames[ifilename]][orbital_names[ifilename]] = float(
p_eigenvals[Spin.down][iband][ikpoint][atom_names[ifilename]][orbital_names[ifilename]] = float(
line.split()[2]
)

Expand Down
17 changes: 17 additions & 0 deletions tests/files/cohp/Fatband_SiO2/Test_p_x/POSCAR
View file
Original file line number Diff line number Diff line change
@@ -0,0 +1,17 @@
Si3 O6
1.0
5.0218978900000000 0.0000000000000000 0.0000000000000000
-2.5109484399999999 4.3490909799999997 0.0000000000000000
0.0000000000000000 0.0000000000000000 5.5119294099999996
Si O
3 6
direct
0.0000000000000000 0.4763431500000000 0.6666670000000000 Si
0.5236568500000000 0.5236568500000000 0.0000000000000000 Si
0.4763431500000000 0.0000000000000000 0.3333330000000000 Si
0.1589037800000000 0.7440031600000000 0.4613477300000000 O
0.2559968400000000 0.4149006200000000 0.7946807299999999 O
0.5850993799999999 0.8410962200000000 0.1280147300000000 O
0.7440031600000000 0.1589037800000000 0.5386522700000000 O
0.4149006200000000 0.2559968400000000 0.2053192700000000 O
0.8410962200000000 0.5850993799999999 0.8719852700000000 O
77 changes: 58 additions & 19 deletions tests/io/lobster/test_inputs.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -40,7 +40,6 @@
__email__ = "janine.george@uclouvain.be, esters@uoregon.edu"
__date__ = "Dec 10, 2017"


module_dir = os.path.dirname(os.path.abspath(__file__))


Expand Down Expand Up @@ -1260,41 +1259,59 @@ def test_msonable(self):
lobsterout_from_dict = Lobsterout.from_dict(dict_data)
assert dict_data == lobsterout_from_dict.as_dict()
# test initialization with empty attributes (ensure file is not read again)
dict_data_empty = self.lobsterout_doscar_lso.ATTRIBUTE_DEFAULTS
dict_data_empty = dict.fromkeys(self.lobsterout_doscar_lso._ATTRIBUTES, None)
lobsterout_empty_init_dict = Lobsterout.from_dict(dict_data_empty).as_dict()
for attribute in lobsterout_empty_init_dict:
if "@" not in attribute:
assert dict_data_empty[attribute] == lobsterout_empty_init_dict[attribute]
assert lobsterout_empty_init_dict[attribute] is None

with pytest.raises(ValueError, match="invalid=val is not a valid attribute for Lobsterout"):
Lobsterout(filename=None, invalid="val")


class TestFatband(PymatgenTest):
def setUp(self):
self.structure = Vasprun(
filename=f"{TEST_FILES_DIR}/cohp/Fatband_SiO2/Test_p_x/vasprun.xml",
ionic_step_skip=None,
ionic_step_offset=0,
parse_dos=True,
parse_eigen=False,
parse_projected_eigen=False,
parse_potcar_file=False,
occu_tol=1e-8,
exception_on_bad_xml=True,
).final_structure
self.fatband_SiO2_p_x = Fatband(
filenames=f"{TEST_FILES_DIR}/cohp/Fatband_SiO2/Test_p_x",
Kpointsfile=f"{TEST_FILES_DIR}/cohp/Fatband_SiO2/Test_p_x/KPOINTS",
vasprun=f"{TEST_FILES_DIR}/cohp/Fatband_SiO2/Test_p_x/vasprun.xml",
kpoints_file=f"{TEST_FILES_DIR}/cohp/Fatband_SiO2/Test_p_x/KPOINTS",
structure=self.structure,
vasprun_file=f"{TEST_FILES_DIR}/cohp/Fatband_SiO2/Test_p_x/vasprun.xml",
)
self.vasprun_SiO2_p_x = Vasprun(filename=f"{TEST_FILES_DIR}/cohp/Fatband_SiO2/Test_p_x/vasprun.xml")
self.bs_symmline = self.vasprun_SiO2_p_x.get_band_structure(line_mode=True, force_hybrid_mode=True)
self.fatband_SiO2_p = Fatband(
filenames=f"{TEST_FILES_DIR}/cohp/Fatband_SiO2/Test_p",
Kpointsfile=f"{TEST_FILES_DIR}/cohp/Fatband_SiO2/Test_p/KPOINTS",
vasprun=f"{TEST_FILES_DIR}/cohp/Fatband_SiO2/Test_p/vasprun.xml",
kpoints_file=f"{TEST_FILES_DIR}/cohp/Fatband_SiO2/Test_p/KPOINTS",
vasprun_file=f"{TEST_FILES_DIR}/cohp/Fatband_SiO2/Test_p/vasprun.xml",
structure=self.structure,
)
self.fatband_SiO2_p2 = Fatband(
filenames=f"{TEST_FILES_DIR}/cohp/Fatband_SiO2/Test_p",
kpoints_file=f"{TEST_FILES_DIR}/cohp/Fatband_SiO2/Test_p/KPOINTS",
structure=self.structure,
vasprun_file=None,
efermi=1.0647039,
)
self.vasprun_SiO2_p = Vasprun(filename=f"{TEST_FILES_DIR}/cohp/Fatband_SiO2/Test_p/vasprun.xml")
self.bs_symmline2 = self.vasprun_SiO2_p.get_band_structure(line_mode=True, force_hybrid_mode=True)
self.fatband_SiO2_spin = Fatband(
filenames=f"{TEST_FILES_DIR}/cohp/Fatband_SiO2/Test_Spin",
Kpointsfile=f"{TEST_FILES_DIR}/cohp/Fatband_SiO2/Test_Spin/KPOINTS",
vasprun=os.path.join(
TEST_FILES_DIR,
"cohp",
"Fatband_SiO2/Test_Spin/vasprun.xml",
),
kpoints_file=f"{TEST_FILES_DIR}/cohp/Fatband_SiO2/Test_Spin/KPOINTS",
vasprun_file=os.path.join(TEST_FILES_DIR, "cohp", "Fatband_SiO2/Test_Spin/vasprun.xml"),
structure=self.structure,
)

self.vasprun_SiO2_spin = Vasprun(
filename=os.path.join(
TEST_FILES_DIR,
Expand Down Expand Up @@ -1334,6 +1351,7 @@ def test_attributes(self):
assert self.fatband_SiO2_p.structure[0].frac_coords == approx([0.0, 0.47634315, 0.666667])
assert self.fatband_SiO2_p.structure[0].species_string == "Si"
assert self.fatband_SiO2_p.structure[0].coords == approx([-1.19607309, 2.0716597, 3.67462144])
assert self.fatband_SiO2_p.efermi == approx(1.0647039)

assert list(self.fatband_SiO2_spin.label_dict["M"]) == approx([0.5, 0.0, 0.0])
assert self.fatband_SiO2_spin.efermi == self.vasprun_SiO2_spin.efermi
Expand All @@ -1353,6 +1371,13 @@ def test_attributes(self):
assert self.fatband_SiO2_spin.structure[0].coords == approx([-1.19607309, 2.0716597, 3.67462144])

def test_raises(self):
with pytest.raises(ValueError, match="vasprun_file or efermi have to be provided"):
Fatband(
filenames=f"{TEST_FILES_DIR}/cohp/Fatband_SiO2/Test_Spin",
kpoints_file=f"{TEST_FILES_DIR}/cohp/Fatband_SiO2/Test_Spin/KPOINTS",
vasprun_file=None,
structure=self.structure,
)
with pytest.raises(
ValueError, match="The are two FATBAND files for the same atom and orbital. The program will stop"
):
Expand All @@ -1361,8 +1386,20 @@ def test_raises(self):
f"{TEST_FILES_DIR}/cohp/Fatband_SiO2/Test_p_x/FATBAND_si1_3p_x.lobster",
f"{TEST_FILES_DIR}/cohp/Fatband_SiO2/Test_p_x/FATBAND_si1_3p_x.lobster",
],
Kpointsfile=f"{TEST_FILES_DIR}/cohp/Fatband_SiO2/Test_p_x/KPOINTS",
vasprun=f"{TEST_FILES_DIR}/cohp/Fatband_SiO2/Test_p_x/vasprun.xml",
kpoints_file=f"{TEST_FILES_DIR}/cohp/Fatband_SiO2/Test_p_x/KPOINTS",
vasprun_file=f"{TEST_FILES_DIR}/cohp/Fatband_SiO2/Test_p_x/vasprun.xml",
structure=self.structure,
)

with pytest.raises(ValueError, match="A structure object has to be provided"):
self.fatband_SiO2_p_x = Fatband(
filenames=[
f"{TEST_FILES_DIR}/cohp/Fatband_SiO2/Test_p_x/FATBAND_si1_3p_x.lobster",
f"{TEST_FILES_DIR}/cohp/Fatband_SiO2/Test_p_x/FATBAND_si1_3p_x.lobster",
],
kpoints_file=f"{TEST_FILES_DIR}/cohp/Fatband_SiO2/Test_p_x/KPOINTS",
vasprun_file=f"{TEST_FILES_DIR}/cohp/Fatband_SiO2/Test_p_x/vasprun.xml",
structure=None,
)

with pytest.raises(
Expand All @@ -1374,15 +1411,17 @@ def test_raises(self):
f"{TEST_FILES_DIR}/cohp/Fatband_SiO2/Test_p_x/FATBAND_si1_3p_x.lobster",
f"{TEST_FILES_DIR}/cohp/Fatband_SiO2/Test_p/FATBAND_si1_3p.lobster",
],
Kpointsfile=f"{TEST_FILES_DIR}/cohp/Fatband_SiO2/Test_p_x/KPOINTS",
vasprun=f"{TEST_FILES_DIR}/cohp/Fatband_SiO2/Test_p_x/vasprun.xml",
kpoints_file=f"{TEST_FILES_DIR}/cohp/Fatband_SiO2/Test_p_x/KPOINTS",
vasprun_file=f"{TEST_FILES_DIR}/cohp/Fatband_SiO2/Test_p_x/vasprun.xml",
structure=self.structure,
)

with pytest.raises(ValueError, match="No FATBAND files in folder or given"):
self.fatband_SiO2_p_x = Fatband(
filenames=".",
Kpointsfile=f"{TEST_FILES_DIR}/cohp/Fatband_SiO2/Test_p_x/KPOINTS",
vasprun=f"{TEST_FILES_DIR}/cohp/Fatband_SiO2/Test_p_x/vasprun.xml",
kpoints_file=f"{TEST_FILES_DIR}/cohp/Fatband_SiO2/Test_p_x/KPOINTS",
vasprun_file=f"{TEST_FILES_DIR}/cohp/Fatband_SiO2/Test_p_x/vasprun.xml",
structure=self.structure,
)

def test_get_bandstructure(self):
Expand Down

0 comments on commit 72977a9

Please sign in to comment.