Implement ion velocities

src/py4vasp/_util/parser.py

@@ -128,6 +128,11 @@ def ion_positions_and_selective_dynamics(self):
             reciprocal_lattice_vectors = self.get_reciprocal_lattice_vectors(self.cell)
             direct_positions = cartesian_positions @ reciprocal_lattice_vectors.T
             positions = np.remainder(direct_positions, 1)
+        else:
+            raise ParserError(
+                "The type of positions is not specified in the right format. Choose\
+                either 'Direct' or 'Cartesian'."
+            )
         if self.has_selective_dynamics:
             selective_dynamics = [x.split()[3:6] for x in positions_and_selective_dyn]
             selective_dynamics = [
@@ -139,10 +144,28 @@ def ion_positions_and_selective_dynamics(self):
 
         return positions, selective_dynamics
 
+    @property
+    def has_lattice_velocities(self):
+        num_species = self.topology.number_ion_types.data.sum()
+        idx_start = 7 + num_species
+        if self.has_selective_dynamics:
+            idx_start += 1
+        if self.species_name is None:
+            idx_start += 1
+        if len(self.split_poscar) <= idx_start:
+            raise ParserError("No lattice velocities found in POSCAR.")
+        lattice_velocities_header = self.split_poscar[idx_start]
+        if lattice_velocities_header == "Lattice velocities and vectors":
+            return True
+        else:
+            return False
+
     @property
     def lattice_velocities(self):
         num_species = self.topology.number_ion_types.data.sum()
         idx_start = 7 + num_species
+        if not self.has_lattice_velocities:
+            raise ParserError("No lattice velocities found in POSCAR.")
         if self.has_selective_dynamics:
             idx_start += 1
         if self.species_name is None:
@@ -152,6 +175,39 @@ def lattice_velocities(self):
         lattice_velocities = VaspData(np.array(lattice_velocities, dtype=float))
         return lattice_velocities
 
+    @classmethod
+    def _convert_direct_to_cartesian(cls, cell, x, scale=True):
+        if scale:
+            lattice_vectors = cell.lattice_vectors.data * cell.scale
+        else:
+            lattice_vectors = np.array(cell.lattice_vectors.data)
+
+        cartesian_positions = x @ lattice_vectors.T
+        return cartesian_positions
+
+    @property
+    def ion_velocities(self):
+        num_species = self.topology.number_ion_types.data.sum()
+        idx_start = 7 + num_species
+        if self.has_selective_dynamics:
+            idx_start += 1
+        if self.species_name is None:
+            idx_start += 1
+        if self.has_lattice_velocities:
+            idx_start += 8
+        if len(self.split_poscar) <= idx_start:
+            raise ParserError("No ion velocities found in POSCAR.")
+        coordinate_system = self.split_poscar[idx_start]
+        ion_velocities = self.split_poscar[idx_start + 1 : idx_start + 1 + num_species]
+        ion_velocities = [x.split() for x in ion_velocities]
+        if coordinate_system == "Direct":
+            ion_velocities = self._convert_direct_to_cartesian(
+                self.cell, np.array(ion_velocities, dtype=float), scale=False
+            )
+            ion_velocities = ion_velocities.tolist()
+        ion_velocities = VaspData(np.array(ion_velocities, dtype=float))
+        return ion_velocities
+
     def to_contcar(self):
         ion_positions, selective_dynamics = self.ion_positions_and_selective_dynamics
         structure = Structure(

tests/util/test_parser.py

@@ -85,7 +85,8 @@ def _cubic_BN(
         has_ion_positions: bool = True,
         has_lattice_velocities: bool = False,
         has_ion_velocities: bool = False,
-        coordinate_system: str = "Direct",
+        ions_coordinate_system: str = "Direct",
+        velocity_coordinate_system: str = "Cartesian",
     ):
         comment_line = "Cubic BN" if has_comment_line else None
         scaling_factor = (
@@ -100,10 +101,10 @@ def _cubic_BN(
         ions_per_species = [1, 1] if has_ion_per_species else None
         selective_dynamics = "Selective dynamics" if has_selective_dynamics else None
         positions = [[0.0, 0.0, 0.0], [0.25, 0.0, 0.25]]
-        if coordinate_system == "Direct":
+        if ions_coordinate_system == "Direct":
             ion_positions = [["Direct"]] + positions if has_ion_positions else None
             ion_positions_direct = copy.deepcopy(ion_positions)
-        elif coordinate_system == "Cartesian":
+        elif ions_coordinate_system == "Cartesian":
             if positions is None:
                 ion_positions_direct = None
                 ion_positions = None
@@ -132,9 +133,22 @@ def _cubic_BN(
             if has_lattice_velocities
             else None
         )
-        ion_velocities = (
-            [[0.0, 0.0, 0.0], [0.0, 0.0, 0.0]] if has_ion_velocities else None
+        direct_velocities = [[0.2, 0.4, -0.2], [0.4, 0.6, -0.3]]
+        cartesian_velocities = np.array(direct_velocities) @ np.array(lattice).T
+        cartesian_velocities = np.round(cartesian_velocities, 5)
+        cartesian_velocities = cartesian_velocities.tolist()
+        if velocity_coordinate_system == "Direct":
+            ion_velocities = (
+                [["Direct"]] + direct_velocities if has_ion_velocities else None
+            )
+        elif velocity_coordinate_system == "Cartesian":
+            ion_velocities = (
+                [["Cartesian"]] + cartesian_velocities if has_ion_velocities else None
+            )
+        output_ion_velocities = (
+            [["Cartesian"]] + cartesian_velocities if has_ion_velocities else None
         )
+
         componentwise_input = [
             comment_line,
             scaling_factor,
@@ -156,7 +170,7 @@ def _cubic_BN(
             selective_dynamics,
             ion_positions_direct,
             lattice_velocities,
-            ion_velocities,
+            output_ion_velocities,
         ]
         arguments = {}
         if not has_species_name:
@@ -227,7 +241,7 @@ def test_cubic_BN_fixture_selective_dynamics(cubic_BN):
 
 
 def test_cubic_BN_fixture_cartesian(cubic_BN):
-    output_poscar_string, *_ = cubic_BN(coordinate_system="Cartesian")
+    output_poscar_string, *_ = cubic_BN(ions_coordinate_system="Cartesian")
     expected_poscar_string = """Cubic BN
 2.0
 0.0 0.5 0.5
@@ -264,6 +278,36 @@ def test_cubic_BN_fixture_lattice_velocities(cubic_BN):
     assert output_poscar_string == expected_poscar_string
 
 
+def test_cubic_BN_fixture_ion_velocities(cubic_BN):
+    output_poscar_string, *_ = cubic_BN(
+        has_lattice_velocities=True,
+        has_ion_velocities=True,
+        velocity_coordinate_system="Cartesian",
+    )
+    expected_poscar_string = """Cubic BN
+2.0
+0.0 0.5 0.5
+0.5 0.0 0.5
+0.5 0.5 0.0
+B N
+1 1
+Direct
+0.0 0.0 0.0
+0.25 0.0 0.25
+Lattice velocities and vectors
+1
+0.0 -0.6 0.2
+0.1 0.3 -0.2
+0.2 -0.4 0.4
+0.0 1.0 1.0
+1.0 0.0 1.0
+1.0 1.0 0.0
+Cartesian
+0.1 0.0 0.3
+0.15 0.05 0.5"""
+    assert output_poscar_string == expected_poscar_string
+
+
 def test_comment_line(cubic_BN):
     poscar_string, componentwise_inputs, _ = cubic_BN()
     comment_line = componentwise_inputs[0]
@@ -394,7 +438,7 @@ def test_positions_cartesian(
     poscar_string, componentwise_inputs, arguments = cubic_BN(
         has_species_name=has_species_name,
         has_selective_dynamics=has_selective_dynamics,
-        coordinate_system="Cartesian",
+        ions_coordinate_system="Cartesian",
         num_scaling_factors=num_scaling_factors,
     )
     ion_positions = componentwise_inputs[6]
@@ -434,10 +478,35 @@ def test_lattice_velocities(cubic_BN, Assert):
     Assert.allclose(expected_lattice_velocities, output_lattice_velocities)
 
 
+def test_no_lattice_velocities(cubic_BN):
+    poscar_string, *_ = cubic_BN(has_lattice_velocities=False)
+    with pytest.raises(ParserError):
+        ParsePoscar(poscar_string).lattice_velocities
+
+
+@pytest.mark.parametrize("velocity_coordinate_system", ["Cartesian", "Direct"])
+def test_ion_velocities(cubic_BN, velocity_coordinate_system):
+    poscar_string, componentwise_inputs, arguments = cubic_BN(
+        has_lattice_velocities=True,
+        has_ion_velocities=True,
+        velocity_coordinate_system=velocity_coordinate_system,
+    )
+    ion_velocities = componentwise_inputs[8]
+    expected_ion_velocities = [x[0:3] for x in ion_velocities[1:]]
+    expected_ion_velocities = VaspData(expected_ion_velocities)
+    output_ion_velocities = ParsePoscar(poscar_string, **arguments).ion_velocities
+    print(expected_ion_velocities)
+    print(output_ion_velocities)
+    assert np.allclose(expected_ion_velocities, output_ion_velocities)
+
+
 @pytest.mark.parametrize("has_species_name", [True, False])
 def test_to_contcar(cubic_BN, has_species_name, Assert):
     poscar_string, componentwise_inputs, arguments = cubic_BN(
-        has_species_name=has_species_name
+        has_species_name=has_species_name,
+        has_selective_dynamics=True,
+        has_lattice_velocities=True,
+        has_ion_velocities=True,
     )
     output_contcar = ParsePoscar(poscar_string, **arguments).to_contcar()
     assert isinstance(output_contcar, CONTCAR)