minor improvments

materialsproject · Oct 1, 2024 · 8786b1c · 8786b1c
1 parent 6e34a44
commit 8786b1c
Show file tree

Hide file tree

Showing 3 changed files with 25 additions and 28 deletions.
diff --git a/src/pymatgen/electronic_structure/dos.py b/src/pymatgen/electronic_structure/dos.py
@@ -1468,17 +1468,14 @@ def get_site_t2g_eg_resolved_dos(
         for s, atom_dos in self.pdos.items():
             if s == site:
                 for orb, pdos in atom_dos.items():
-                    try:
-                        orbital = _get_orb_lobster(str(orb))
-                        if orbital is None:
-                            raise ValueError("orbital is None")
-
-                        if orbital in (Orbital.dxy, Orbital.dxz, Orbital.dyz):
-                            t2g_dos.append(pdos)
-                        elif orbital in (Orbital.dx2, Orbital.dz2):
-                            eg_dos.append(pdos)
-                    except (AttributeError, ValueError):
-                        pass
+                    orbital = _get_orb_lobster(str(orb))
+                    if orbital is None:
+                        raise ValueError("orbital is None")
+
+                    if orbital in (Orbital.dxy, Orbital.dxz, Orbital.dyz):
+                        t2g_dos.append(pdos)
+                    elif orbital in (Orbital.dx2, Orbital.dz2):
+                        eg_dos.append(pdos)
         return {
             "t2g": Dos(self.efermi, self.energies, functools.reduce(add_densities, t2g_dos)),
             "e_g": Dos(self.efermi, self.energies, functools.reduce(add_densities, eg_dos)),
@@ -1497,14 +1494,11 @@ def get_spd_dos(self) -> dict[str, Dos]:
         orb = None
         for atom_dos in self.pdos.values():
             for orb, pdos in atom_dos.items():
-                try:
-                    orbital_type = _get_orb_type_lobster(str(orb))
-                    if orbital_type not in spd_dos:
-                        spd_dos[orbital_type] = pdos
-                    else:
-                        spd_dos[orbital_type] = add_densities(spd_dos[orbital_type], pdos)
-                except (AttributeError, ValueError):
-                    pass
+                orbital_type = _get_orb_type_lobster(str(orb))
+                if orbital_type not in spd_dos:
+                    spd_dos[orbital_type] = pdos
+                else:
+                    spd_dos[orbital_type] = add_densities(spd_dos[orbital_type], pdos)
 
         return {orb: Dos(self.efermi, self.energies, densities) for orb, densities in spd_dos.items()}  # type: ignore[misc]
 
@@ -1522,14 +1516,11 @@ def get_element_spd_dos(self, el: SpeciesLike) -> dict[str, Dos]:
         for site, atom_dos in self.pdos.items():
             if site.specie == el:
                 for orb, pdos in atom_dos.items():
-                    try:
-                        orbital_type = _get_orb_type_lobster(str(orb))
-                        if orbital_type not in el_dos:
-                            el_dos[orbital_type] = pdos
-                        else:
-                            el_dos[orbital_type] = add_densities(el_dos[orbital_type], pdos)
-                    except (AttributeError, ValueError):
-                        pass
+                    orbital_type = _get_orb_type_lobster(str(orb))
+                    if orbital_type not in el_dos:
+                        el_dos[orbital_type] = pdos
+                    else:
+                        el_dos[orbital_type] = add_densities(el_dos[orbital_type], pdos)
 
         return {orb: Dos(self.efermi, self.energies, densities) for orb, densities in el_dos.items()}  # type: ignore[misc]
 

diff --git a/src/pymatgen/io/lobster/__init__.py b/src/pymatgen/io/lobster/__init__.py
@@ -10,6 +10,7 @@
 from .inputs import Lobsterin
 from .outputs import (
     Bandoverlaps,
+    Bwdf,
     Charge,
     Cohpcar,
     Doscar,
@@ -20,6 +21,7 @@
     Lobsterout,
     MadelungEnergies,
     NciCobiList,
+    Polarization,
     SitePotential,
     Wavefunction,
 )
diff --git a/src/pymatgen/io/lobster/outputs.py b/src/pymatgen/io/lobster/outputs.py
@@ -2418,5 +2418,9 @@ def __init__(
                     self.bwdf[Spin.up] = np.append(self.bwdf[Spin.up], float(clean_line[1]))
                     self.bwdf[Spin.down] = np.append(self.bwdf[Spin.down], float(clean_line[2]))
                 else:
-                    self.bwdf[Spin.down] = np.append(self.bwdf[Spin.up], float(clean_line[1]))
+                    self.bwdf[Spin.up] = np.append(self.bwdf[Spin.up], float(clean_line[1]))
+
+            if len(self.bwdf[Spin.down]) == 0:  # remove down spin key if not spin polarized calculation
+                del self.bwdf[Spin.down]
+
             self.bin_width = np.diff(self.centers)[0]