Fix creating Structure.from_POSCAR (#95)

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):

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):