Skip to content

Commit

Permalink
Add type annotations for io.vasp.outputs (#3776)
Browse files Browse the repository at this point in the history
* remove a lot of ignore tags

* first go: quick look and comment/type tweaks

* remove ALL type: ignore[reportPossiblyUnboundVariable]

* a quick look

* tweak module docstring

* fix/supress PossiblyUnboundVariable

* pre-commit fix

* revert changes from other branch

* [Need Discussion] set IVDW default to 0

* suppress None error for ElementTree searches

* replace dict update with |=

* fix typo in update -> |=, and mypy fixes

* fix another typo in update -> |=

* add DEBUG tag and mypy fixes

* remove DEBUG tag

* more mypy fixes

* more mypy fixes

* mypy fixes

* finish mypy errors
  • Loading branch information
DanielYang59 authored May 31, 2024
1 parent cd8846b commit c70e5b8
Show file tree
Hide file tree
Showing 7 changed files with 1,095 additions and 873 deletions.
5 changes: 4 additions & 1 deletion pymatgen/core/trajectory.py
Original file line number Diff line number Diff line change
Expand Up @@ -16,7 +16,6 @@

from pymatgen.core.structure import Composition, DummySpecies, Element, Lattice, Molecule, Species, Structure
from pymatgen.io.ase import AseAtomsAdaptor
from pymatgen.io.vasp.outputs import Vasprun, Xdatcar

if TYPE_CHECKING:
from collections.abc import Iterator
Expand Down Expand Up @@ -547,9 +546,13 @@ def from_file(cls, filename: str | Path, constant_lattice: bool = True, **kwargs
structures = []

if fnmatch(filename, "*XDATCAR*"):
from pymatgen.io.vasp.outputs import Xdatcar

structures = Xdatcar(filename).structures

elif fnmatch(filename, "vasprun*.xml*"):
from pymatgen.io.vasp.outputs import Vasprun

structures = Vasprun(filename).structures

elif fnmatch(filename, "*.traj"):
Expand Down
11 changes: 7 additions & 4 deletions pymatgen/entries/computed_entries.py
Original file line number Diff line number Diff line change
Expand Up @@ -291,7 +291,7 @@ def __init__(
energy_adjustments: list | None = None,
parameters: dict | None = None,
data: dict | None = None,
entry_id: object | None = None,
entry_id: str | None = None,
):
"""Initialize a ComputedEntry.
Expand Down Expand Up @@ -557,7 +557,7 @@ def __init__(
energy_adjustments: list | None = None,
parameters: dict | None = None,
data: dict | None = None,
entry_id: object | None = None,
entry_id: str | None = None,
) -> None:
"""Initialize a ComputedStructureEntry.
Expand Down Expand Up @@ -647,7 +647,10 @@ def from_dict(cls, dct: dict) -> Self:
entry_id=dct.get("entry_id"),
)

def normalize(self, mode: Literal["formula_unit", "atom"] = "formula_unit") -> ComputedStructureEntry:
def normalize(
self,
mode: Literal["formula_unit", "atom"] = "formula_unit",
) -> ComputedStructureEntry:
"""Normalize the entry's composition and energy. The structure remains unchanged.
Args:
Expand Down Expand Up @@ -698,7 +701,7 @@ def __init__(
energy_adjustments: list | None = None,
parameters: dict | None = None,
data: dict | None = None,
entry_id: object | None = None,
entry_id: str | None = None,
):
"""
Args:
Expand Down
4 changes: 2 additions & 2 deletions pymatgen/io/lobster/outputs.py
Original file line number Diff line number Diff line change
Expand Up @@ -1380,14 +1380,14 @@ def __init__(
iband = 0
if line.split()[0] != "#":
if linenumber < self.nbands:
if ifilename == 0:
if ifilename == 0 and self.efermi is not None:
eigenvals[Spin.up][iband][idx_kpt] = float(line.split()[1]) + self.efermi

p_eigenvals[Spin.up][iband][idx_kpt][atom_names[ifilename]][orbital_names[ifilename]] = float(
line.split()[2]
)
if linenumber >= self.nbands and self.is_spinpolarized:
if ifilename == 0:
if ifilename == 0 and self.efermi is not None:
eigenvals[Spin.down][iband][idx_kpt] = float(line.split()[1]) + self.efermi
p_eigenvals[Spin.down][iband][idx_kpt][atom_names[ifilename]][orbital_names[ifilename]] = float(
line.split()[2]
Expand Down
8 changes: 7 additions & 1 deletion pymatgen/io/vasp/optics.py
Original file line number Diff line number Diff line change
Expand Up @@ -72,14 +72,17 @@ class DielectricFunctionCalculator(MSONable):

@classmethod
def from_vasp_objects(cls, vrun: Vasprun, waveder: Waveder) -> Self:
"""Construct a DielectricFunction from Vasprun, Kpoint, and Waveder objects.
"""Construct a DielectricFunction from Vasprun, Kpoint, and Waveder.
Args:
vrun: Vasprun object
kpoint: Kpoint object
waveder: Waveder object
"""
bands = vrun.eigenvalues
if bands is None:
raise RuntimeError("eigenvalues cannot be None.")

sspins = [Spin.up, Spin.down]
eigs = np.stack([bands[spin] for spin in sspins[: vrun.parameters["ISPIN"]]], axis=2)[..., 0]
eigs = np.swapaxes(eigs, 0, 1)
Expand All @@ -95,6 +98,9 @@ def from_vasp_objects(cls, vrun: Vasprun, waveder: Waveder) -> Self:
if vrun.parameters["ISYM"] != 0:
raise NotImplementedError("ISYM != 0 is not implemented yet")

if efermi is None:
raise ValueError("efermi cannot be None.")

return cls(
cder_real=waveder.cder_real,
cder_imag=waveder.cder_imag,
Expand Down
Loading

0 comments on commit c70e5b8

Please sign in to comment.