imdclient 0.1.0__py3-none-any.whl → 0.1.2__py3-none-any.whl

imdclient/IMDClient.py

@@ -30,6 +30,23 @@ logger = logging.getLogger(__name__)
 
 
 class IMDClient:
+    """
+    Parameters
+    ----------
+    host : str
+        Hostname of the server
+    port : int
+        Port number of the server
+    n_atoms : int
+        Number of atoms in the simulation
+    socket_bufsize : int, (optional)
+        Size of the socket buffer in bytes. Default is to use the system default
+    buffer_size : int (optional)
+        IMDFramebuffer will be filled with as many :class:`IMDFrame` fit in `buffer_size` [``10MB``]
+    **kwargs : dict (optional)
+        Additional keyword arguments to pass to the :class:`BaseIMDProducer` and :class:`IMDFrameBuffer`
+    """
+
     def __init__(
         self,
         host,
@@ -39,22 +56,7 @@ class IMDClient:
         multithreaded=True,
         **kwargs,
     ):
-        """
-        Parameters
-        ----------
-        host : str
-            Hostname of the server
-        port : int
-            Port number of the server
-        n_atoms : int
-            Number of atoms in the simulation
-        socket_bufsize : int, optional
-            Size of the socket buffer in bytes. Default is to use the system default
-        buffer_size : int, optional
-            IMDFramebuffer will be filled with as many IMDFrames fit in `buffer_size` [``10MB``]
-        **kwargs : optional
-            Additional keyword arguments to pass to the IMDProducer and IMDFrameBuffer
-        """
+
         self._stopped = False
         self._conn = self._connect_to_server(host, port, socket_bufsize)
         self._imdsinfo = self._await_IMD_handshake()
@@ -254,6 +256,25 @@ class IMDClient:
 
 
 class BaseIMDProducer(threading.Thread):
+    """
+
+    Parameters
+    ----------
+    conn : socket.socket
+        Connection object to the server
+    buffer : IMDFrameBuffer
+        Buffer object to hold IMD frames. If `multithreaded` is False, this
+        argument is ignored
+    sinfo : IMDSessionInfo
+        Information about the IMD session
+    n_atoms : int
+        Number of atoms in the simulation
+    multithreaded : bool, optional
+        If True, socket interaction will occur in a separate thread &
+        frames will be buffered. Single-threaded, blocking IMDClient
+        should only be used in testing [[``True``]]
+
+    """
 
     def __init__(
         self,
@@ -265,24 +286,6 @@ class BaseIMDProducer(threading.Thread):
         timeout=5,
         **kwargs,
     ):
-        """
-        Parameters
-        ----------
-        conn : socket.socket
-            Connection object to the server
-        buffer : IMDFrameBuffer
-            Buffer object to hold IMD frames. If `multithreaded` is False, this
-            argument is ignored
-        sinfo : IMDSessionInfo
-            Information about the IMD session
-        n_atoms : int
-            Number of atoms in the simulation
-        multithreaded : bool, optional
-            If True, socket interaction will occur in a separate thread &
-            frames will be buffered. Single-threaded, blocking IMDClient
-            should only be used in testing [[``True``]]
-
-        """
         super(BaseIMDProducer, self).__init__(daemon=True)
         self._conn = conn
         self._imdsinfo = sinfo
@@ -638,6 +641,21 @@ class IMDFrameBuffer:
     """
     Acts as interface between producer (IMDProducer) and consumer (IMDClient) threads
     when IMDClient runs in multithreaded mode
+
+    Parameters
+    ----------
+    imdsinfo : IMDSessionInfo
+        Information about the IMD session
+    n_atoms : int
+        Number of atoms in the simulation
+    buffer_size : int, optional
+        Size of the buffer in bytes [``10MB``]
+    pause_empty_proportion : float, optional
+        Lower threshold proportion of the buffer's IMDFrames that are empty
+        before the simulation is paused [``0.25``]
+    unpause_empty_proportion : float, optional
+        Proportion of the buffer's IMDFrames that must be empty
+        before the simulation is unpaused [``0.5``]
     """
 
     def __init__(
@@ -649,23 +667,6 @@ class IMDFrameBuffer:
         unpause_empty_proportion=0.5,
         **kwargs,
     ):
-        """
-        Parameters
-        ----------
-        imdsinfo : IMDSessionInfo
-            Information about the IMD session
-        n_atoms : int
-            Number of atoms in the simulation
-        buffer_size : int, optional
-            Size of the buffer in bytes [``10MB``]
-        pause_empty_proportion : float, optional
-            Lower threshold proportion of the buffer's IMDFrames that are empty
-            before the simulation is paused [``0.25``]
-        unpause_empty_proportion : float, optional
-            Proportion of the buffer's IMDFrames that must be empty
-            before the simulation is unpaused [``0.5``]
-        """
-
         # Syncing reader and producer
         self._producer_finished = False
         self._consumer_finished = False

imdclient/IMDREADER.py

@@ -24,6 +24,17 @@ logger = logging.getLogger("imdclient.IMDClient")
 class IMDReader(StreamReaderBase):
     """
     Reader for IMD protocol packets.
+
+    Parameters
+    ----------
+    filename : a string of the form "host:port" where host is the hostname
+        or IP address of the listening GROMACS server and port
+        is the port number.
+    n_atoms : int (optional)
+        number of atoms in the system. defaults to number of atoms
+        in the topology. don't set this unless you know what you're doing.
+    kwargs : dict (optional)
+        keyword arguments passed to the constructed :class:`IMDClient`
     """
 
     format = "IMD"
@@ -37,17 +48,6 @@ class IMDReader(StreamReaderBase):
         n_atoms=None,
         **kwargs,
     ):
-        """
-        Parameters
-        ----------
-        filename : a string of the form "host:port" where host is the hostname
-            or IP address of the listening GROMACS server and port
-            is the port number.
-        n_atoms : int (optional)
-            number of atoms in the system. defaults to number of atoms
-            in the topology. don't set this unless you know what you're doing.
-        """
-
         super(IMDReader, self).__init__(filename, **kwargs)
 
         logger.debug("IMDReader initializing")
@@ -101,7 +101,9 @@ class IMDReader(StreamReaderBase):
             self.ts.data["dt"] = imdf.dt
             self.ts.data["step"] = imdf.step
         if imdf.energies is not None:
-            self.ts.data.update(imdf.energies)
+            self.ts.data.update(
+                {k: v for k, v in imdf.energies.items() if k != "step"}
+            )
         if imdf.box is not None:
             self.ts.dimensions = core.triclinic_box(*imdf.box)
         if imdf.positions is not None:

imdclient/tests/conftest.py

@@ -6,33 +6,37 @@ Global pytest fixtures
 # Command line arguments for 'test_manual.py'
 def pytest_addoption(parser):
     parser.addoption(
-        "--topol_arg",
+        "--topol_path_arg",
         action="store",
+        default=None,
     )
     parser.addoption(
-        "--traj_arg",
+        "--traj_path_arg",
         action="store",
+        default=None,
+    )
+    parser.addoption(
+        "--first_frame_arg", action="store", type=int, default=None
     )
-    parser.addoption("--first_frame_arg", action="store", type=int)
 
 
 def pytest_generate_tests(metafunc):
     # This is called for every test. Only get/set command line arguments
     # if the argument is specified in the list of test "fixturenames".
-    topol = metafunc.config.option.topol_arg
-    traj = metafunc.config.option.traj_arg
+    topol = metafunc.config.option.topol_path_arg
+    traj = metafunc.config.option.traj_path_arg
     first_frame = metafunc.config.option.first_frame_arg
 
     if all(
         arg in metafunc.fixturenames
-        for arg in ["topol_arg", "traj_arg", "first_frame_arg"]
+        for arg in ["topol_path_arg", "traj_path_arg", "first_frame_arg"]
     ):
         if topol is None or traj is None or first_frame is None:
             raise ValueError(
-                "Must pass all three of '--topol_arg <path/to/topology>', "
-                + "'--traj_arg <path/to/trajectory>', "
+                "Must pass all three of '--topol_path_arg <path/to/topology>', "
+                + "'--traj_path_arg <path/to/trajectory>', "
                 + "'--first_frame_arg <first traj frame to compare to IMD>"
             )
-        metafunc.parametrize("topol_arg", [topol])
-        metafunc.parametrize("traj_arg", [traj])
+        metafunc.parametrize("topol_path_arg", [topol])
+        metafunc.parametrize("traj_path_arg", [traj])
         metafunc.parametrize("first_frame_arg", [first_frame])

imdclient/tests/test_imdreader.py

@@ -55,7 +55,7 @@ class IMDReference(BaseReference):
         self.n_atoms = traj.n_atoms
         self.prec = 3
 
-        self.trajectory = f"localhost:{self.port}"
+        self.trajectory = f"imd://localhost:{self.port}"
         self.topology = COORDINATES_TOPOLOGY
         self.changing_dimensions = True
         self.reader = IMDReader
@@ -585,7 +585,7 @@ class TestStreamIteration:
         server.set_imdsessioninfo(imdsinfo)
         server.handshake_sequence("localhost", port, first_frame=True)
         reader = IMDReader(
-            f"localhost:{port}",
+            f"imd://localhost:{port}",
             n_atoms=universe.trajectory.n_atoms,
         )
         server.send_frames(1, 5)

imdclient/tests/test_manual.py

@@ -1,18 +1,24 @@
-from imdclient.IMDClient import IMDClient
+from imdclient.IMDREADER import IMDReader
 import pytest
 import MDAnalysis as mda
 from MDAnalysisTests.coordinates.base import assert_timestep_almost_equal
 from numpy.testing import (
-    assert_array_almost_equal,
-    assert_almost_equal,
     assert_allclose,
 )
 import numpy as np
 from .base import assert_allclose_with_logging
+from pathlib import Path
 
 import logging
 
 logger = logging.getLogger("imdclient.IMDClient")
+file_handler = logging.FileHandler("manual_test.log")
+formatter = logging.Formatter(
+    "%(asctime)s - %(name)s - %(levelname)s - %(message)s"
+)
+file_handler.setFormatter(formatter)
+logger.addHandler(file_handler)
+logger.setLevel(logging.DEBUG)
 
 
 class TestIMDv3Manual:
@@ -23,48 +29,63 @@ class TestIMDv3Manual:
     and then run this command relative to the root of the cloned respository:
 
     pytest -s imdclient/tests/test_manual.py \
-        --topol_arg <path/to/topology> \
-        --traj_arg <path/to/trajectory> \
+        --topol_path_arg <path/to/topology> \
+        --traj_path_arg <path/to/trajectory> \
         --first_frame_arg <first traj frame to compare to IMD>
 
-    Where the topology is the same topology as the IMD system, the trajectory is the trajectory
-    to compare to IMD data read from the socket, and the first frame is the first frame of the 
+    Where the topology is the same topology as the IMD system, the trajectory is the path where
+    the trajectory of the running simulation is being written, and the first frame is the first frame of the
     trajectory which should be compared to IMD data read from the socket (0 for GROMACS and NAMD, 1 for LAMMPS)
     """
 
     @pytest.fixture()
-    def universe(self, topol_arg, traj_arg):
-        return mda.Universe(topol_arg, traj_arg)
+    def true_u(self, imd_u, topol_path_arg, traj_path_arg):
+        return mda.Universe(topol_path_arg, traj_path_arg)
 
     @pytest.fixture()
-    def client(self, universe):
-        client = IMDClient("localhost", 8888, universe.trajectory.n_atoms)
-        yield client
-        client.stop()
+    def imd_u(self, topol_path_arg, tmp_path):
+        tmp_u = mda.Universe(topol_path_arg, "imd://localhost:8888")
+        with mda.Writer(
+            f"{tmp_path.as_posix()}/imd_test_traj.trr", tmp_u.atoms.n_atoms
+        ) as w:
+            for ts in tmp_u.trajectory:
+                w.write(tmp_u.atoms)
+        imd_u = mda.Universe(
+            topol_path_arg, f"{tmp_path.as_posix()}/imd_test_traj.trr"
+        )
+        yield imd_u
 
-    def test_compare_imd_to_true_traj(self, universe, client, first_frame_arg):
-        imdsinfo = client.get_imdsessioninfo()
+    def test_compare_imd_to_true_traj(self, true_u, imd_u, first_frame_arg):
 
-        for ts in universe.trajectory[first_frame_arg:]:
-            imdf = client.get_imdframe()
-            if imdsinfo.time:
-                assert_allclose(imdf.time, ts.time, atol=1e-03)
-                assert_allclose(imdf.step, ts.data["step"])
-            if imdsinfo.box:
+        for i in range(first_frame_arg, len(true_u.trajectory)):
+            assert_allclose(
+                true_u.trajectory[i].time, imd_u.trajectory[i].time, atol=1e-03
+            )
+            assert_allclose(
+                true_u.trajectory[i].data["step"],
+                imd_u.trajectory[i].data["step"],
+            )
+            if true_u.trajectory[i].dimensions is not None:
                 assert_allclose_with_logging(
-                    imdf.box,
-                    ts.triclinic_dimensions,
+                    true_u.trajectory[i].dimensions,
+                    imd_u.trajectory[i].dimensions,
                     atol=1e-03,
                 )
-            if imdsinfo.positions:
+            if true_u.trajectory[i].has_positions:
                 assert_allclose_with_logging(
-                    imdf.positions, ts.positions, atol=1e-03
+                    true_u.trajectory[i].positions,
+                    imd_u.trajectory[i].positions,
+                    atol=1e-03,
                 )
-            if imdsinfo.velocities:
+            if true_u.trajectory[i].has_velocities:
                 assert_allclose_with_logging(
-                    imdf.velocities, ts.velocities, atol=1e-03
+                    true_u.trajectory[i].velocities,
+                    imd_u.trajectory[i].velocities,
+                    atol=1e-03,
                 )
-            if imdsinfo.forces:
+            if true_u.trajectory[i].has_forces:
                 assert_allclose_with_logging(
-                    imdf.forces, ts.forces, atol=1e-03
+                    true_u.trajectory[i].forces,
+                    imd_u.trajectory[i].forces,
+                    atol=1e-03,
                 )

imdclient/tests/test_stream_analysis.py

@@ -38,7 +38,7 @@ class TestStreamAnalysis:
         server.set_imdsessioninfo(imdsinfo)
         server.handshake_sequence("localhost", port, first_frame=True)
 
-        imd_universe = mda.Universe(COORDINATES_TOPOLOGY, f"localhost:{port}")
+        imd_universe = mda.Universe(COORDINATES_TOPOLOGY, f"imd://localhost:{port}")
         server.send_frames(1, 5)
 
         yield imd_universe

imdclient/utils.py

@@ -43,22 +43,22 @@ class timeit(object):
         # always propagate exceptions forward
         return False
 
-
 # NOTE: think of other edge cases as well- should be robust
 def parse_host_port(filename):
+    if not filename.startswith("imd://"):
+        raise ValueError("IMDReader: URL must be in the format 'imd://host:port'")
+    
     # Check if the format is correct
-    parts = filename.split(":")
+    parts = filename.split("imd://")[1].split(":")
     if len(parts) == 2:
-        host = parts[0]  # Hostname part
+        host = parts[0] 
         try:
-            port = int(parts[1])  # Convert the port part to an integer
+            port = int(parts[1])
             return (host, port)
         except ValueError:
-            # Handle the case where the port is not a valid integer
-            raise ValueError("Port must be an integer")
+            raise ValueError("IMDReader: Port must be an integer")
     else:
-        # Handle the case where the format does not match "host:port"
-        raise ValueError("Filename must be in the format 'host:port'")
+        raise ValueError("IMDReader: URL must be in the format 'imd://host:port'")
 
 
 def approximate_timestep_memsize(