Skip to content

Commit

Permalink
LOBSTER IO improvements (#3649)
Browse files Browse the repository at this point in the history 
* add as_dict method to Bandoverlaps, Grosspop, Charge, Madelung, Sitepotential parser

* add tests for as_dict method

* Charges MSonable

* adapt charges test

* add deprecation warning to Charge.Loewdin and Charge.Mulliken

* fix msonable sitepotential, grosspop, bandoverlap, icohplist, madelung

* add robust test for msonable

* make lobsterout msonable

* add safety check for kwargs

* Lobsterout._ATTRIBUTES list->tuple

include value in ValueError msg

* cover ValueError in test_msonable

---------

Co-authored-by: Janosh Riebesell <janosh.riebesell@gmail.com>
  • Loading branch information
@naik-aakash @janosh
naik-aakash and janosh committed Feb 24, 2024
1 parent 40afffb commit 6bdd406
Show file tree
Hide file tree
Showing 2 changed files with 127 additions and 63 deletions.
166 changes: 108 additions & 58 deletions pymatgen/io/lobster/outputs.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -757,7 +757,7 @@ def Loewdin(self):
return self.loewdin


class Lobsterout:
class Lobsterout(MSONable):
"""
Class to read in the lobsterout and evaluate the spilling, save the basis, save warnings, save infos.
Expand Down Expand Up @@ -793,80 +793,122 @@ class Lobsterout:
warning_lines (str): String with all warnings.
"""

# valid Lobsterout instance attributes
_ATTRIBUTES = (
"filename",
"is_restart_from_projection",
"lobster_version",
"number_of_threads",
"dft_program",
"number_of_spins",
"charge_spilling",
"total_spilling",
"elements",
"basis_type",
"basis_functions",
"timing",
"warning_lines",
"info_orthonormalization",
"info_lines",
"has_doscar",
"has_doscar_lso",
"has_cohpcar",
"has_coopcar",
"has_cobicar",
"has_charge",
"has_madelung",
"has_projection",
"has_bandoverlaps",
"has_fatbands",
"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="lobsterout"):
def __init__(self, filename: str | None, **kwargs) -> None:
"""
Args:
filename: filename of lobsterout.
**kwargs:dict to initialize Lobsterout instance (see > INIT_ATTRIBUTES_DEFAULTS)
"""
# read in file
with zopen(filename, mode="rt") as file:
data = file.read().split("\n") # [3:-3]
if len(data) == 0:
raise OSError("lobsterout does not contain any data")
self.filename = filename
if kwargs:
for attr, val in kwargs.items():
if attr in self.ATTRIBUTE_DEFAULTS:
setattr(self, attr, val)
else:
raise ValueError(f"{attr}={val} is not a valid attribute for Lobsterout")
else:
with zopen(filename, mode="rt") as file: # read in file
data = file.read().split("\n")
if len(data) == 0:
raise OSError("lobsterout does not contain any data")

# check if Lobster starts from a projection
self.is_restart_from_projection = "loading projection from projectionData.lobster..." in data
# check if Lobster starts from a projection
self.is_restart_from_projection = "loading projection from projectionData.lobster..." in data

self.lobster_version = self._get_lobster_version(data=data)
self.lobster_version = self._get_lobster_version(data=data)

self.number_of_threads = int(self._get_threads(data=data))
self.dft_program = self._get_dft_program(data=data)
self.number_of_threads = int(self._get_threads(data=data))
self.dft_program = self._get_dft_program(data=data)

self.number_of_spins = self._get_number_of_spins(data=data)
chargespilling, totalspilling = self._get_spillings(data=data, number_of_spins=self.number_of_spins)
self.charge_spilling = chargespilling
self.total_spilling = totalspilling
self.number_of_spins = self._get_number_of_spins(data=data)
chargespilling, totalspilling = self._get_spillings(data=data, number_of_spins=self.number_of_spins)
self.charge_spilling = chargespilling
self.total_spilling = totalspilling

elements, basistype, basisfunctions = self._get_elements_basistype_basisfunctions(data=data)
self.elements = elements
self.basis_type = basistype
self.basis_functions = basisfunctions
elements, basistype, basisfunctions = self._get_elements_basistype_basisfunctions(data=data)
self.elements = elements
self.basis_type = basistype
self.basis_functions = basisfunctions

wall_time, user_time, sys_time = self._get_timing(data=data)
timing = {}
timing["wall_time"] = wall_time
timing["user_time"] = user_time
timing["sys_time"] = sys_time
self.timing = timing
wall_time, user_time, sys_time = self._get_timing(data=data)
timing = {}
timing["wall_time"] = wall_time
timing["user_time"] = user_time
timing["sys_time"] = sys_time
self.timing = timing

warninglines = self._get_all_warning_lines(data=data)
self.warning_lines = warninglines
warninglines = self._get_all_warning_lines(data=data)
self.warning_lines = warninglines

orthowarning = self._get_warning_orthonormalization(data=data)
self.info_orthonormalization = orthowarning
orthowarning = self._get_warning_orthonormalization(data=data)
self.info_orthonormalization = orthowarning

infos = self._get_all_info_lines(data=data)
self.info_lines = infos
infos = self._get_all_info_lines(data=data)
self.info_lines = infos

self.has_doscar = "writing DOSCAR.lobster..." in data and "SKIPPING writing DOSCAR.lobster..." not in data
self.has_doscar_lso = (
"writing DOSCAR.LSO.lobster..." in data and "SKIPPING writing DOSCAR.LSO.lobster..." not in data
)
self.has_cohpcar = (
"writing COOPCAR.lobster and ICOOPLIST.lobster..." in data
and "SKIPPING writing COOPCAR.lobster and ICOOPLIST.lobster..." not in data
)
self.has_coopcar = (
"writing COHPCAR.lobster and ICOHPLIST.lobster..." in data
and "SKIPPING writing COHPCAR.lobster and ICOHPLIST.lobster..." not in data
)
self.has_cobicar = (
"writing COBICAR.lobster and ICOBILIST.lobster..." in data
and "SKIPPING writing COBICAR.lobster and ICOBILIST.lobster..." not in data
)
self.has_doscar = "writing DOSCAR.lobster..." in data and "SKIPPING writing DOSCAR.lobster..." not in data
self.has_doscar_lso = (
"writing DOSCAR.LSO.lobster..." in data and "SKIPPING writing DOSCAR.LSO.lobster..." not in data
)
self.has_cohpcar = (
"writing COOPCAR.lobster and ICOOPLIST.lobster..." in data
and "SKIPPING writing COOPCAR.lobster and ICOOPLIST.lobster..." not in data
)
self.has_coopcar = (
"writing COHPCAR.lobster and ICOHPLIST.lobster..." in data
and "SKIPPING writing COHPCAR.lobster and ICOHPLIST.lobster..." not in data
)
self.has_cobicar = (
"writing COBICAR.lobster and ICOBILIST.lobster..." in data
and "SKIPPING writing COBICAR.lobster and ICOBILIST.lobster..." not in data
)

self.has_charge = "SKIPPING writing CHARGE.lobster..." not in data
self.has_projection = "saving projection to projectionData.lobster..." in data
self.has_bandoverlaps = "WARNING: I dumped the band overlap matrices to the file bandOverlaps.lobster." in data
self.has_fatbands = self._has_fatband(data=data)
self.has_grosspopulation = "writing CHARGE.lobster and GROSSPOP.lobster..." in data
self.has_density_of_energies = "writing DensityOfEnergy.lobster..." in data
self.has_madelung = (
"writing SitePotentials.lobster and MadelungEnergies.lobster..." in data
and "skipping writing SitePotentials.lobster and MadelungEnergies.lobster..." not in data
)
self.has_charge = "SKIPPING writing CHARGE.lobster..." not in data
self.has_projection = "saving projection to projectionData.lobster..." in data
self.has_bandoverlaps = (
"WARNING: I dumped the band overlap matrices to the file bandOverlaps.lobster." in data
)
self.has_fatbands = self._has_fatband(data=data)
self.has_grosspopulation = "writing CHARGE.lobster and GROSSPOP.lobster..." in data
self.has_density_of_energies = "writing DensityOfEnergy.lobster..." in data
self.has_madelung = (
"writing SitePotentials.lobster and MadelungEnergies.lobster..." in data
and "skipping writing SitePotentials.lobster and MadelungEnergies.lobster..." not in data
)

def get_doc(self):
"""Returns: LobsterDict with all the information stored in lobsterout."""
Expand Down Expand Up @@ -907,6 +949,14 @@ def get_doc(self):

return LobsterDict

def as_dict(self):
"""MSONable dict"""
dct = vars(self)
dct["@module"] = type(self).__module__
dct["@class"] = type(self).__name__

return dct

@staticmethod
def _get_lobster_version(data):
for row in data:
Expand Down
24 changes: 19 additions & 5 deletions tests/io/lobster/test_inputs.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -1192,7 +1192,7 @@ def test_attributes(self):
assert self.lobsterout_skipping_cobi_madelung.has_madelung is False

def test_get_doc(self):
comparedict = {
ref_data = {
"restart_from_projection": False,
"lobster_version": "v3.1.0",
"threads": 8,
Expand Down Expand Up @@ -1247,13 +1247,27 @@ def test_get_doc(self):
for key, item in self.lobsterout_normal.get_doc().items():
if key not in ["has_cobicar", "has_madelung"]:
if isinstance(item, str):
assert comparedict[key], item
assert ref_data[key], item
elif isinstance(item, int):
assert comparedict[key] == item
assert ref_data[key] == item
elif key in ("charge_spilling", "total_spilling"):
assert item[0] == approx(comparedict[key][0])
assert item[0] == approx(ref_data[key][0])
elif isinstance(item, (list, dict)):
assert item == comparedict[key]
assert item == ref_data[key]

def test_msonable(self):
dict_data = self.lobsterout_normal.as_dict()
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
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]

with pytest.raises(ValueError, match="invalid=val is not a valid attribute for Lobsterout"):
Lobsterout(filename=None, invalid="val")


class TestFatband(PymatgenTest):
Expand Down

0 comments on commit 6bdd406

Please sign in to comment.