From 84067990fcf75a7a8f43a36c40eb5288d7be1fb8 Mon Sep 17 00:00:00 2001
From: florianj77 <florian.joerg@aon.at>
Date: Tue, 5 Sep 2023 14:38:52 +0200
Subject: [PATCH] mda reader for openmm rst files

---
 protex/helpers.py            | 17 +++++++++++++++++
 protex/tests/test_helpers.py | 26 ++++++++++++++++++++++++++
 2 files changed, 43 insertions(+)
 create mode 100644 protex/tests/test_helpers.py

diff --git a/protex/helpers.py b/protex/helpers.py
index f20fbb84..cd3e8246 100644
--- a/protex/helpers.py
+++ b/protex/helpers.py
@@ -1,4 +1,7 @@
 """helpers.py contains additional functions for easier use of protex."""
+from MDAnalysis.coordinates.base import SingleFrameReaderBase
+from MDAnalysis.lib.util import openany
+from openmm import XmlSerializer
 from parmed.formats import PSFFile
 from parmed.topologyobjects import DrudeAtom, ExtraPoint, LocalCoordinatesFrame
 
@@ -8,6 +11,20 @@
 from parmed.utils.io import genopen
 
 
+class XMLSingleReader(SingleFrameReaderBase):
+    format = "rst"
+    units = {"time": "ps", "length": "Nanometer"}
+    _format_type = "xml"
+
+    def _read_first_frame(self):
+        with openany(self.filename) as xmlfile:
+            state = XmlSerializer.deserialize(xmlfile.read())
+            positions = state.getPositions(asNumpy=True)
+        self.n_atoms = positions.shape[0]
+        self.ts = self._Timestep.from_coordinates(positions, **self._ts_kwargs)
+        self.ts.frame = 0
+
+
 class CustomPSFFile(PSFFile):
     """
     Child of parmeds formats/psf.py PFSFile class
diff --git a/protex/tests/test_helpers.py b/protex/tests/test_helpers.py
new file mode 100644
index 00000000..2a0abfba
--- /dev/null
+++ b/protex/tests/test_helpers.py
@@ -0,0 +1,26 @@
+import os
+
+import pytest
+
+from protex.helpers import XMLSingleReader
+
+
+@pytest.mark.skipif(
+    os.getenv("CI") == "true",
+    reason="No MDAnalysis",
+)
+def test_XMLSingleReader():
+    a = XMLSingleReader("protex/forcefield/traj/im1h_oac_150_im1_hoac_350_npt_7.rst")
+
+@pytest.mark.skipif(
+    os.getenv("CI") == "true",
+    reason="No MDAnalysis",
+)
+def test_Universe():
+    import MDAnalysis
+    u = MDAnalysis.Universe("protex/forcefield/im1h_oac_150_im1_hoac_350.psf","protex/forcefield/traj/im1h_oac_150_im1_hoac_350_npt_7.rst")
+    #print(u.trajectory.n_frames)
+    #print(u.trajectory.n_atoms)
+    #print(dir(u.trajectory))
+    sel = u.select_atoms("all")
+    #print(sel.positions)
\ No newline at end of file