From e1037d1f95dac7d064b82f09ad9e100ea49ff0c3 Mon Sep 17 00:00:00 2001
From: "Josef M. Gallmetzer" <64498081+galjos@users.noreply.github.com>
Date: Mon, 24 Aug 2026 09:11:31 +0200
Subject: [PATCH] fix: write vel or forces alone in PQ restart files

The PQ atom-line guard required both velocities and forces to be
present before writing them, so a system carrying only one of the
two was written as a bare 6-field position line. The supplied data
was silently discarded and the resulting file could not be read
back by RestartFileReader, which accepts only 12 or 21 fields.
Writing whenever either quantity is present keeps the supplied
values and pads the missing block with zeros.
---
 PQAnalysis/io/restart_file/restart_writer.py |  2 +-
 tests/io/test_restartWriter.py               | 30 ++++++++++++++++++--
 2 files changed, 28 insertions(+), 4 deletions(-)

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