diff --git a/pymatgen/io/vasp/outputs.py b/pymatgen/io/vasp/outputs.py index 9af84743be7..e1275e25c6b 100644 --- a/pymatgen/io/vasp/outputs.py +++ b/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 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 = "".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,7 +415,7 @@ 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. @@ -425,7 +423,7 @@ def structures(self): 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()} diff --git a/tests/io/vasp/test_outputs.py b/tests/io/vasp/test_outputs.py index 6b722f17928..f40adfc9778 100644 --- a/tests/io/vasp/test_outputs.py +++ b/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)