Vasprun type hints

pymatgen/io/vasp/outputs.py

@@ -106,15 +106,13 @@ def _parse_v_parameters(val_type, val, filename, param_name):
     return val
 
 
-def _parse_varray(elem):
+def _parse_vasp_array(elem) -> list[list[float]]:
     if elem.get("type") == "logical":
-        m = [[i == "T" for i in v.text.split()] for v in elem]
-    else:
-        m = [[_vasprun_float(i) for i in v.text.split()] for v in elem]
-    return m
+        return [[i == "T" for i in v.text.split()] for v in elem]
+    return [[_vasprun_float(i) for i in v.text.split()] for v in elem]
 
 
-def _parse_from_incar(filename, key):
+def _parse_from_incar(filename: str, key: str) -> str | None:
     """Helper function to parse a parameter from the INCAR."""
     dirname = os.path.dirname(filename)
     for f in os.listdir(dirname):
@@ -127,19 +125,19 @@ def _parse_from_incar(filename, key):
     return None
 
 
-def _vasprun_float(f):
+def _vasprun_float(flt: float | str) -> float:
     """
     Large numbers are often represented as ********* in the vasprun.
     This function parses these values as np.nan.
     """
     try:
-        return float(f)
-    except ValueError as e:
-        f = f.strip()
-        if f == "*" * len(f):
+        return float(flt)
+    except ValueError as exc:
+        flt = flt.strip()  # type: ignore[union-attr]
+        if flt == "*" * len(flt):
             warnings.warn("Float overflow (*******) encountered in vasprun")
             return np.nan
-        raise e
+        raise exc
 
 
 class Vasprun(MSONable):
@@ -200,17 +198,17 @@ class Vasprun(MSONable):
 
     def __init__(
         self,
-        filename,
-        ionic_step_skip=None,
-        ionic_step_offset=0,
-        parse_dos=True,
-        parse_eigen=True,
-        parse_projected_eigen=False,
-        parse_potcar_file=True,
-        occu_tol=1e-8,
-        separate_spins=False,
-        exception_on_bad_xml=True,
-    ):
+        filename: str | Path,
+        ionic_step_skip: int | None = None,
+        ionic_step_offset: int = 0,
+        parse_dos: bool = True,
+        parse_eigen: bool = True,
+        parse_projected_eigen: bool = False,
+        parse_potcar_file: bool = True,
+        occu_tol: float = 1e-8,
+        separate_spins: bool = False,
+        exception_on_bad_xml: bool = True,
+    ) -> None:
         """
         Args:
             filename (str): Filename to parse
@@ -272,7 +270,7 @@ def __init__(
                 # The text before the first <calculation> is the preamble!
                 preamble = steps.pop(0)
                 self.nionic_steps = len(steps)
-                new_steps = steps[ionic_step_offset :: int(ionic_step_skip)]
+                new_steps = steps[ionic_step_offset :: int(ionic_step_skip or 1)]
                 # add the tailing information in the last step from the run
                 to_parse = "<calculation>".join(new_steps)
                 if steps[-1] != new_steps[-1]:
@@ -377,7 +375,7 @@ def _parse(self, stream, parse_dos, parse_eigen, parse_projected_eigen):
                             self.other_dielectric[comment] = self._parse_diel(elem)
 
                 elif tag == "varray" and elem.attrib.get("name") == "opticaltransitions":
-                    self.optical_transition = np.array(_parse_varray(elem))
+                    self.optical_transition = np.array(_parse_vasp_array(elem))
                 elif tag == "structure" and elem.attrib.get("name") == "finalpos":
                     self.final_structure = self._parse_structure(elem)
                 elif tag == "dynmat":
@@ -400,7 +398,7 @@ def _parse(self, stream, parse_dos, parse_eigen, parse_projected_eigen):
                         md_data.append({})
                         md_data[-1]["structure"] = self._parse_structure(elem)
                     elif tag == "varray" and elem.attrib.get("name") == "forces":
-                        md_data[-1]["forces"] = _parse_varray(elem)
+                        md_data[-1]["forces"] = _parse_vasp_array(elem)
                     elif tag == "energy":
                         d = {i.attrib["name"]: float(i.text) for i in elem.findall("i")}
                         if "kinetic" in d:
@@ -417,15 +415,15 @@ def _parse(self, stream, parse_dos, parse_eigen, parse_projected_eigen):
         self.vasp_version = self.generator["version"]
 
     @property
-    def structures(self):
+    def structures(self) -> list[Structure]:
         """
         Returns:
             List of Structure objects for the structure at each ionic step.
         """
         return [step["structure"] for step in self.ionic_steps]
 
     @property
-    def epsilon_static(self):
+    def epsilon_static(self) -> list[float]:
         """
         Property only available for DFPT calculations.
 
@@ -436,7 +434,7 @@ def epsilon_static(self):
         return self.ionic_steps[-1].get("epsilon", [])
 
     @property
-    def epsilon_static_wolfe(self):
+    def epsilon_static_wolfe(self) -> list[float]:
         """
         Property only available for DFPT calculations.
 
@@ -447,7 +445,7 @@ def epsilon_static_wolfe(self):
         return self.ionic_steps[-1].get("epsilon_rpa", [])
 
     @property
-    def epsilon_ionic(self):
+    def epsilon_ionic(self) -> list[float]:
         """
         Property only available for DFPT calculations and when IBRION=5, 6, 7 or 8.
 
@@ -473,14 +471,14 @@ def dielectric(self):
         return self.dielectric_data["density"]
 
     @property
-    def optical_absorption_coeff(self):
+    def optical_absorption_coeff(self) -> list[float]:
         """
         Calculate the optical absorption coefficient
         from the dielectric constants. Note that this method is only
         implemented for optical properties calculated with GGA and BSE.
 
         Returns:
-            optical absorption coefficient in list
+            list[float]: optical absorption coefficient
         """
         if self.dielectric_data["density"]:
             real_avg = [
@@ -506,7 +504,7 @@ def f(freq, real, imag):
         return absorption_coeff
 
     @property
-    def converged_electronic(self):
+    def converged_electronic(self) -> bool:
         """
         Returns:
             bool: True if electronic step convergence has been reached in the final ionic step.
@@ -524,7 +522,7 @@ def converged_electronic(self):
         return len(final_elec_steps) < self.parameters["NELM"]
 
     @property
-    def converged_ionic(self):
+    def converged_ionic(self) -> bool:
         """
         Returns:
             bool: True if ionic step convergence has been reached, i.e. that vasp
@@ -534,7 +532,7 @@ def converged_ionic(self):
         return nsw <= 1 or len(self.ionic_steps) < nsw
 
     @property
-    def converged(self):
+    def converged(self) -> bool:
         """
         Returns:
             bool: True if a relaxation run is both ionically and electronically converged.
@@ -578,17 +576,17 @@ def complete_dos(self):
     @property
     def complete_dos_normalized(self) -> CompleteDos:
         """
-        A CompleteDos object which incorporates the total DOS and all
-        projected DOS. Normalized by the volume of the unit cell with
-        units of states/eV/unit cell volume.
+        A CompleteDos object which incorporates the total DOS and all projected DOS.
+        Normalized by the volume of the unit cell with units of states/eV/unit cell
+        volume.
         """
         final_struct = self.final_structure
         pdoss = {final_struct[i]: pdos for i, pdos in enumerate(self.pdos)}
         return CompleteDos(self.final_structure, self.tdos, pdoss, normalize=True)
 
     @property
-    def hubbards(self):
-        """Hubbard U values used if a vasprun is a GGA+U run. {} otherwise."""
+    def hubbards(self) -> dict[str, float]:
+        """Hubbard U values used if a vasprun is a GGA+U run. Otherwise an empty dict."""
         symbols = [s.split()[1] for s in self.potcar_symbols]
         symbols = [re.split(r"_", s)[0] for s in symbols]
         if not self.incar.get("LDAU", False):
@@ -1224,9 +1222,9 @@ def _parse_kpoints(elem):
         for va in elem.findall("varray"):
             name = va.attrib["name"]
             if name == "kpointlist":
-                actual_kpoints = _parse_varray(va)
+                actual_kpoints = _parse_vasp_array(va)
             elif name == "weights":
-                weights = [i[0] for i in _parse_varray(va)]
+                weights = [i[0] for i in _parse_vasp_array(va)]
         elem.clear()
         if k.style == Kpoints.supported_modes.Reciprocal:
             k = Kpoints(
@@ -1239,12 +1237,12 @@ def _parse_kpoints(elem):
         return k, actual_kpoints, weights
 
     def _parse_structure(self, elem):
-        latt = _parse_varray(elem.find("crystal").find("varray"))
-        pos = _parse_varray(elem.find("varray"))
+        latt = _parse_vasp_array(elem.find("crystal").find("varray"))
+        pos = _parse_vasp_array(elem.find("varray"))
         struct = Structure(latt, self.atomic_symbols, pos)
         sdyn = elem.find("varray/[@name='selective']")
         if sdyn:
-            struct.add_site_property("selective_dynamics", _parse_varray(sdyn))
+            struct.add_site_property("selective_dynamics", _parse_vasp_array(sdyn))
         return struct
 
     @staticmethod
@@ -1265,8 +1263,8 @@ def _parse_optical_transition(elem):
         for va in elem.findall("varray"):
             if va.attrib.get("name") == "opticaltransitions":
                 # opticaltransitions array contains oscillator strength and probability of transition
-                oscillator_strength = np.array(_parse_varray(va))[0:]
-                probability_transition = np.array(_parse_varray(va))[0:, 1]
+                oscillator_strength = np.array(_parse_vasp_array(va))[0:]
+                probability_transition = np.array(_parse_vasp_array(va))[0:, 1]
         return oscillator_strength, probability_transition
 
     def _parse_chemical_shielding_calculation(self, elem):
@@ -1277,7 +1275,7 @@ def _parse_chemical_shielding_calculation(self, elem):
         except AttributeError:  # not all calculations have a structure
             struct = None
         for va in elem.findall("varray"):
-            istep[va.attrib["name"]] = _parse_varray(va)
+            istep[va.attrib["name"]] = _parse_vasp_array(va)
         istep["structure"] = struct
         istep["electronic_steps"] = []
         calculation.append(istep)
@@ -1316,7 +1314,7 @@ def _parse_calculation(self, elem):
         except AttributeError:  # not all calculations have a structure
             struct = None
         for va in elem.findall("varray"):
-            istep[va.attrib["name"]] = _parse_varray(va)
+            istep[va.attrib["name"]] = _parse_vasp_array(va)
         istep["electronic_steps"] = esteps
         istep["structure"] = struct
         elem.clear()
@@ -1330,7 +1328,7 @@ def _parse_dos(elem):
         idensities = {}
 
         for s in elem.find("total").find("array").find("set").findall("set"):
-            data = np.array(_parse_varray(s))
+            data = np.array(_parse_vasp_array(s))
             energies = data[:, 0]
             spin = Spin.up if s.attrib["comment"] == "spin 1" else Spin.down
             tdensities[spin] = data[:, 1]
@@ -1347,7 +1345,7 @@ def _parse_dos(elem):
 
                 for ss in s.findall("set"):
                     spin = Spin.up if ss.attrib["comment"] == "spin 1" else Spin.down
-                    data = np.array(_parse_varray(ss))
+                    data = np.array(_parse_vasp_array(ss))
                     nrow, ncol = data.shape
                     for j in range(1, ncol):
                         orb = Orbital(j - 1) if lm else OrbitalType(j - 1)
@@ -1366,7 +1364,7 @@ def _parse_eigen(elem):
         for s in elem.find("array").find("set").findall("set"):
             spin = Spin.up if s.attrib["comment"] == "spin 1" else Spin.down
             for ss in s.findall("set"):
-                eigenvalues[spin].append(_parse_varray(ss))
+                eigenvalues[spin].append(_parse_vasp_array(ss))
         eigenvalues = {spin: np.array(v) for spin, v in eigenvalues.items()}
         elem.clear()
         return eigenvalues
@@ -1382,7 +1380,7 @@ def _parse_projected_eigen(elem):
             for ss in s.findall("set"):
                 dk = []
                 for sss in ss.findall("set"):
-                    db = _parse_varray(sss)
+                    db = _parse_vasp_array(sss)
                     dk.append(db)
                 proj_eigen[spin].append(dk)
         proj_eigen = {spin: np.array(v) for spin, v in proj_eigen.items()}

tests/io/vasp/test_outputs.py

@@ -170,9 +170,9 @@ def test_standard(self):
         filepath = f"{TEST_FILES_DIR}/vasprun.xml"
         vasp_run = Vasprun(filepath, parse_potcar_file=False)
 
-        # Test NELM parsing.
+        # Test NELM parsing
         assert vasp_run.parameters["NELM"] == 60
-        # test pdos parsing
+        # test pDOS parsing
 
         assert vasp_run.complete_dos.spin_polarization == 1.0
         assert Vasprun(f"{TEST_FILES_DIR}/vasprun.xml.etest1.gz").complete_dos.spin_polarization is None
@@ -210,7 +210,7 @@ def test_standard(self):
 
         filepath2 = f"{TEST_FILES_DIR}/lifepo4.xml"
         vasprun_ggau = Vasprun(filepath2, parse_projected_eigen=True, parse_potcar_file=False)
-        totalscsteps = sum(len(i["electronic_steps"]) for i in vasp_run.ionic_steps)
+        total_sc_steps = sum(len(i["electronic_steps"]) for i in vasp_run.ionic_steps)
         assert len(vasp_run.ionic_steps) == 29
         assert len(vasp_run.structures) == len(vasp_run.ionic_steps)
 
@@ -225,7 +225,7 @@ def test_standard(self):
             vasp_run.structures[i] == vasp_run.ionic_steps[i]["structure"] for i in range(len(vasp_run.ionic_steps))
         )
 
-        assert totalscsteps == 308, "Incorrect number of energies read from vasprun.xml"
+        assert total_sc_steps == 308, "Incorrect number of energies read from vasprun.xml"
 
         assert ["Li"] + 4 * ["Fe"] + 4 * ["P"] + 16 * ["O"] == vasp_run.atomic_symbols
         assert vasp_run.final_structure.composition.reduced_formula == "LiFe4(PO4)4"
@@ -234,12 +234,8 @@ def test_standard(self):
         assert vasp_run.eigenvalues is not None, "Eigenvalues cannot be read"
         assert vasp_run.final_energy == approx(-269.38319884, abs=1e-7)
         assert vasp_run.tdos.get_gap() == approx(2.0589, abs=1e-4)
-        expectedans = (2.539, 4.0906, 1.5516, False)
-        (gap, cbm, vbm, direct) = vasp_run.eigenvalue_band_properties
-        assert gap == approx(expectedans[0])
-        assert cbm == approx(expectedans[1])
-        assert vbm == approx(expectedans[2])
-        assert direct == expectedans[3]
+        expected = (2.539, 4.0906, 1.5516, False)
+        assert vasp_run.eigenvalue_band_properties == approx(expected)
         assert not vasp_run.is_hubbard
         assert vasp_run.potcar_symbols == [
             "PAW_PBE Li 17Jan2003",
@@ -251,9 +247,9 @@ def test_standard(self):
         assert vasp_run.kpoints is not None, "Kpoints cannot be read"
         assert vasp_run.actual_kpoints is not None, "Actual kpoints cannot be read"
         assert vasp_run.actual_kpoints_weights is not None, "Actual kpoints weights cannot be read"
-        for atomdoses in vasp_run.pdos:
-            for orbitaldos in atomdoses:
-                assert orbitaldos is not None, "Partial Dos cannot be read"
+        for atom_doses in vasp_run.pdos:
+            for orbital_dos in atom_doses:
+                assert orbital_dos is not None, "Partial Dos cannot be read"
 
         # test skipping ionic steps.
         vasprun_skip = Vasprun(filepath, 3, parse_potcar_file=False)