diff --git a/PQAnalysis/io/restart_file/restart_writer.py b/PQAnalysis/io/restart_file/restart_writer.py index 2325a125..7320fa89 100644 --- a/PQAnalysis/io/restart_file/restart_writer.py +++ b/PQAnalysis/io/restart_file/restart_writer.py @@ -220,7 +220,7 @@ def _get_atom_lines( line += f"{atom.name} {atom_counter[i]} {residue} " line += f"{pos[0]} {pos[1]} {pos[2]}" - if (frame.has_vel and frame.has_forces) or md_engine_format != MDEngineFormat.PQ: + if (frame.has_vel or frame.has_forces) or md_engine_format != MDEngineFormat.PQ: line += f" {vel[0]} {vel[1]} {vel[2]}" line += f" {force[0]} {force[1]} {force[2]}" diff --git a/tests/io/test_restartWriter.py b/tests/io/test_restartWriter.py index 78d80b18..5aab261b 100644 --- a/tests/io/test_restartWriter.py +++ b/tests/io/test_restartWriter.py @@ -5,6 +5,7 @@ from . import pytestmark from PQAnalysis.io import RestartFileWriter +from PQAnalysis.io.restart_file.api import read_restart_file, write_restart_file from PQAnalysis.traj import MDEngineFormat from PQAnalysis.core import Cell, Atom from PQAnalysis.atomic_system import AtomicSystem @@ -184,9 +185,9 @@ def test_write(self, capsys): captured = capsys.readouterr() assert captured.out == """ Box 10.0 10.0 10.0 90 90 90 -C 0 0 0.0 0.0 0.0 -H 1 0 1.0 1.0 1.0 -H 2 0 2.0 2.0 2.0 +C 0 0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 +H 1 0 1.0 1.0 1.0 1.0 1.0 1.0 0.0 0.0 0.0 +H 2 0 2.0 2.0 2.0 2.0 2.0 2.0 0.0 0.0 0.0 """ forces = np.array([[0.0, 0.0, 0.0], [1.0, 1.0, 1.0], [2.0, 2.0, 2.0]]) @@ -214,3 +215,26 @@ def test_write(self, capsys): H 1 0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 H 2 0 2.0 2.0 2.0 2.0 2.0 2.0 2.0 2.0 2.0 2.0 2.0 2.0 2.0 2.0 2.0 2.0 2.0 2.0 """ + + @pytest.mark.usefixtures("tmpdir") + def test_write_only_velocities_or_forces_round_trip(self): + atoms = [Atom("C"), Atom("H")] + positions = np.array([[0.0, 0.0, 0.0], [1.0, 0.0, 0.0]]) + values = np.array([[0.1, 0.2, 0.3], [0.4, 0.5, 0.6]]) + cell = Cell(10.0, 10.0, 10.0) + + frame = AtomicSystem(atoms=atoms, pos=positions, vel=values, cell=cell) + write_restart_file(frame, "vel_only.rst", mode="o") + system = read_restart_file("vel_only.rst") + + assert np.allclose(system.vel, values) + assert np.allclose(system.forces, np.zeros((2, 3))) + + frame = AtomicSystem( + atoms=atoms, pos=positions, forces=values, cell=cell + ) + write_restart_file(frame, "forces_only.rst", mode="o") + system = read_restart_file("forces_only.rst") + + assert np.allclose(system.forces, values) + assert np.allclose(system.vel, np.zeros((2, 3)))