Fix velocity when resetting structure (#356)

* Fix velocity when resetting structure

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Add test

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Move vector_to_lammps to a separate function (#357)

* Implement _vector_to_lammps

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* run black

* Remove flatten because it's always flattened

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Sam Dareska <[email protected]>

Author: 3 people

pylammpsmpi/wrapper/ase.py

@@ -96,9 +96,7 @@ def interactive_positions_setter(self, positions: list[list[float]]) -> None:
         Args:
             positions (List[List[float]]): The positions of atoms.
         """
-        if not _check_ortho_prism(prism=self._prism):
-            positions = self._prism.vector_to_lammps(positions)
-        positions = np.array(positions).flatten()
+        positions = _vector_to_lammps(vector=positions, prism=self._prism)
         if self._cores == 1:
             self._interactive_library.scatter_atoms(
                 "x", 1, 3, (len(positions) * c_double)(*positions)
@@ -279,10 +277,11 @@ def interactive_structure_setter(
                 self.interactive_lib_command(
                     command=f"mass {id_eam + 1:3d} {1.00:f}",
                 )
-        if not _check_ortho_prism(prism=self._prism):
-            positions = self._prism.vector_to_lammps(structure.positions).flatten()
-        else:
-            positions = structure.positions.flatten()
+        positions = _vector_to_lammps(vector=structure.positions, prism=self._prism)
+        velocities = _vector_to_lammps(
+            vector=structure.get_velocities(),
+            prism=self._prism,
+        )
         try:
             elem_all = get_lammps_indicies_from_ase_structure(
                 structure=structure, el_eam_lst=el_eam_lst
@@ -299,7 +298,7 @@ def interactive_structure_setter(
                 id=None,
                 type=elem_all,
                 x=positions,
-                v=None,
+                v=velocities,
                 image=None,
                 shrinkexceed=False,
             )
@@ -309,7 +308,7 @@ def interactive_structure_setter(
                 id=range(1, len(structure) + 1),
                 type=elem_all,
                 x=positions,
-                v=None,
+                v=velocities,
                 image=None,
                 shrinkexceed=False,
             )
@@ -476,6 +475,15 @@ def cell_is_skewed(cell, tolerance=1.0e-8):
     return not (volume > 0 and abs(volume - prod) / volume < tolerance)
 
 
+def _vector_to_lammps(vector, prism):
+    if vector is not None and np.any(vector):
+        if not _check_ortho_prism(prism=prism):
+            vector = prism.vector_to_lammps(vector)
+        return vector.flatten()
+    else:
+        return None
+
+
 def _check_ortho_prism(prism, rtol=0.0, atol=1e-08):
     """
     Check if the rotation matrix of the UnfoldingPrism object is sufficiently close to a unit matrix

tests/test_ase_interface.py

@@ -121,6 +121,57 @@ def test_small_displacement(self):
             np.all(np.isclose(lmp.interactive_positions_getter(), positions))
         )
 
+    def test_velocities(self):
+        lmp = LammpsASELibrary(
+            working_directory=None,
+            cores=1,
+            comm=None,
+            logger=logging.getLogger("TestStaticLogger"),
+            log_file=None,
+            library=LammpsLibrary(cores=2),
+            disable_log_file=True,
+        )
+        structure = bulk("Al", cubic=True)
+        velocities_initial = np.array(
+            [[1.0, 1.0, 1.0], [1.0, 1.0, 1.0], [1.0, 1.0, 1.0], [1.0, 1.0, 1.0]]
+        )
+        structure.set_velocities(velocities_initial)
+        lmp.interactive_structure_setter(
+            structure=structure,
+            units="lj",
+            dimension=3,
+            boundary=" ".join(["p" if coord else "f" for coord in structure.pbc]),
+            atom_style="atomic",
+            el_eam_lst=["Al"],
+            calc_md=False,
+        )
+        lmp.interactive_lib_command("pair_style lj/cut 6.0")
+        lmp.interactive_lib_command("pair_coeff 1 1 1.0 1.0 4.04")
+        lmp.interactive_lib_command(
+            command="thermo_style custom step temp pe etotal pxx pxy pxz pyy pyz pzz vol"
+        )
+        lmp.interactive_lib_command(command="thermo_modify format float %20.15g")
+        lmp.interactive_lib_command("run 0")
+        self.assertTrue(
+            np.all(np.isclose(lmp.interactive_positions_getter(), structure.positions))
+        )
+        self.assertTrue(
+            np.isclose(lmp.interactive_energy_pot_getter(), -0.04342932384411341)
+        )
+        self.assertTrue(
+            np.all(np.equal(lmp.interactive_velocities_getter(), velocities_initial))
+        )
+        positions = structure.positions.copy()
+        positions[0] += np.array([0.1, 0.1, 0.1])
+        lmp.interactive_positions_setter(positions=positions)
+        lmp.interactive_lib_command("run 0")
+        self.assertTrue(
+            np.isclose(lmp.interactive_energy_pot_getter(), -0.043829529274767506)
+        )
+        self.assertTrue(
+            np.all(np.isclose(lmp.interactive_positions_getter(), positions))
+        )
+
     def test_small_displacement_skewed(self):
         lmp = LammpsASELibrary(
             working_directory=None,