From 6ad82e56eec4cace65e2fac284971de0b506186e Mon Sep 17 00:00:00 2001 From: Martin Schlipf Date: Wed, 17 May 2023 14:49:08 +0200 Subject: [PATCH] Fix creating Structure.from_POSCAR (#95) --- src/py4vasp/_data/structure.py | 31 +++++++++++++++++++++++++++++-- tests/data/test_structure.py | 28 ++++++++++++++++++++++++++++ 2 files changed, 57 insertions(+), 2 deletions(-) diff --git a/src/py4vasp/_data/structure.py b/src/py4vasp/_data/structure.py index 29dfc9e7..3b82c0c1 100644 --- a/src/py4vasp/_data/structure.py +++ b/src/py4vasp/_data/structure.py @@ -63,9 +63,11 @@ class Structure(slice_.Mixin, base.Refinery): "Converting Å to nm used for mdtraj trajectories." @classmethod - def from_POSCAR(cls, poscar): + def from_POSCAR(cls, poscar, *, elements=None): """Generate a structure from string in POSCAR format.""" - poscar = io.StringIO(str(poscar)) + poscar = _replace_or_set_elements(str(poscar), elements) + print(poscar) + poscar = io.StringIO(poscar) structure = ase.io.read(poscar, format="vasp") return cls.from_ase(structure) @@ -316,6 +318,31 @@ def _cell_from_ase(structure): return raw.Cell(lattice_vectors=np.array([structure.get_cell()])) +def _replace_or_set_elements(poscar, elements): + line_with_elements = 5 + elements = "" if not elements else " ".join(elements) + lines = poscar.split("\n") + if _elements_not_in_poscar(lines[line_with_elements]): + _raise_error_if_elements_not_set(elements) + lines.insert(line_with_elements, elements) + elif elements: + lines[line_with_elements] = elements + return "\n".join(lines) + + +def _elements_not_in_poscar(elements): + elements = elements.split() + return any(element.isdecimal() for element in elements) + + +def _raise_error_if_elements_not_set(elements): + if not elements: + message = """The POSCAR file does not specify the elements needed to create a + Structure. Please pass `elements=[...]` to the `from_POSCAR` routine where + ... are the elements in the same order as in the POSCAR.""" + raise exception.IncorrectUsage(message) + + class Mixin: @property def _structure(self): diff --git a/tests/data/test_structure.py b/tests/data/test_structure.py index a3f6681b..09cce967 100644 --- a/tests/data/test_structure.py +++ b/tests/data/test_structure.py @@ -104,6 +104,34 @@ def test_to_poscar(Sr2TiO4): def test_from_poscar(Sr2TiO4, Assert): structure = Structure.from_POSCAR(REF_POSCAR) check_Sr2TiO4_structure(structure.read(), Sr2TiO4.ref, -1, Assert) + structure = Structure.from_POSCAR(REF_POSCAR, elements=["Ba", "Zr", "S"]) + actual = structure.read() + Assert.allclose(actual["lattice_vectors"], Sr2TiO4.ref.lattice_vectors[-1]) + Assert.allclose(actual["positions"], Sr2TiO4.ref.positions[-1]) + assert actual["elements"] == ["Ba", "Ba", "Zr", "S", "S", "S", "S"] + assert actual["names"] == ["Ba_1", "Ba_2", "Zr_1", "S_1", "S_2", "S_3", "S_4"] + + +def test_from_poscar_without_elements(Sr2TiO4, Assert): + poscar = """\ +POSCAR without elements +1.0 + 6.9229000000000003 0.0000000000000000 0.0000000000000000 + 4.6945030167999979 5.0880434191000035 0.0000000000000000 + -5.8086962205000017 -2.5440193935999971 2.7773292841999986 +2 1 4 +Direct + 0.6452900000000000 0.6452900000000000 0.0000000000000000 + 0.3547100000000000 0.3547100000000000 0.0000000000000000 + 0.0000000000000000 0.0000000000000000 0.0000000000000000 + 0.8417800000000000 0.8417800000000000 0.0000000000000000 + 0.1582300000000000 0.1582300000000000 0.0000000000000000 + 0.5000000000000000 0.0000000000000000 0.5000000000000000 + 0.0000000000000000 0.5000000000000000 0.5000000000000000""" + with pytest.raises(exception.IncorrectUsage): + structure = Structure.from_POSCAR(poscar) + structure = Structure.from_POSCAR(poscar, elements=["Sr", "Ti", "O"]) + check_Sr2TiO4_structure(structure.read(), Sr2TiO4.ref, -1, Assert) def test_to_ase_Sr2TiO4(Sr2TiO4, Assert):