Generalize fatband plots from Lobster (#3688)

* 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 <[email protected]>
materialsproject · Mar 14, 2024 · 72977a9 · 72977a9
1 parent 668fa57
commit 72977a9
Show file tree

Hide file tree

Showing 3 changed files with 139 additions and 64 deletions.
diff --git a/pymatgen/io/lobster/outputs.py b/pymatgen/io/lobster/outputs.py
@@ -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:
@@ -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."""
@@ -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])
@@ -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
@@ -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]
                         )
 

diff --git a/tests/files/cohp/Fatband_SiO2/Test_p_x/POSCAR b/tests/files/cohp/Fatband_SiO2/Test_p_x/POSCAR
@@ -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
diff --git a/tests/io/lobster/test_inputs.py b/tests/io/lobster/test_inputs.py
@@ -40,7 +40,6 @@
 __email__ = "[email protected], [email protected]"
 __date__ = "Dec 10, 2017"
 
-
 module_dir = os.path.dirname(os.path.abspath(__file__))
 
 
@@ -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,
@@ -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
@@ -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"
         ):
@@ -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(
@@ -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):