imdclient 0.1.0__py3-none-any.whl

imdclient/streambase.py

@@ -0,0 +1,199 @@
+from MDAnalysis.coordinates.base import (
+    ReaderBase,
+    FrameIteratorBase,
+    FrameIteratorAll,
+)
+import numbers
+import warnings
+
+
+class StreamReaderBase(ReaderBase):
+
+    def __init__(self, filename, convert_units=True, **kwargs):
+        super(StreamReaderBase, self).__init__(
+            filename, convert_units=convert_units, **kwargs
+        )
+        self._init_scope = True
+        self._reopen_called = False
+        self._first_ts = None
+
+    def _read_next_timestep(self):
+        # No rewinding- to both load the first frame after  __init__
+        # and access it again during iteration, we need to store first ts in mem
+        if not self._init_scope and self._frame == -1:
+            self._frame += 1
+            # can't simply return the same ts again- transformations would be applied twice
+            # instead, return the pre-transformed copy
+            return self._first_ts
+
+        ts = self._read_frame(self._frame + 1)
+
+        if self._init_scope:
+            self._first_ts = self.ts.copy()
+            self._init_scope = False
+
+        return ts
+
+    @property
+    def n_frames(self):
+        """Changes as stream is processed unlike other readers"""
+        raise RuntimeError(
+            "{}: n_frames is unknown".format(self.__class__.__name__)
+        )
+
+    def __len__(self):
+        raise RuntimeError(
+            "{} has unknown length".format(self.__class__.__name__)
+        )
+
+    def next(self):
+        """Don't rewind after iteration. When _reopen() is called,
+        an error will be raised
+        """
+        try:
+            ts = self._read_next_timestep()
+        except (EOFError, IOError):
+            # Don't rewind here like we normally would
+            raise StopIteration from None
+        else:
+            for auxname, reader in self._auxs.items():
+                ts = self._auxs[auxname].update_ts(ts)
+
+            ts = self._apply_transformations(ts)
+
+        return ts
+
+    def rewind(self):
+        """Raise error on rewind"""
+        raise RuntimeError(
+            "{}: Stream-based readers can't be rewound".format(
+                self.__class__.__name__
+            )
+        )
+
+    # Incompatible methods
+    def copy(self):
+        raise NotImplementedError(
+            "{} does not support copying".format(self.__class__.__name__)
+        )
+
+    def _reopen(self):
+        if self._reopen_called:
+            raise RuntimeError(
+                "{}: Cannot reopen stream".format(self.__class__.__name__)
+            )
+        self._frame = -1
+        self._reopen_called = True
+
+    def __getitem__(self, frame):
+        """Return the Timestep corresponding to *frame*.
+
+        If `frame` is a integer then the corresponding frame is
+        returned. Negative numbers are counted from the end.
+
+        If frame is a :class:`slice` then an iterator is returned that
+        allows iteration over that part of the trajectory.
+
+        Note
+        ----
+        *frame* is a 0-based frame index.
+        """
+        if isinstance(frame, slice):
+            _, _, step = self.check_slice_indices(
+                frame.start, frame.stop, frame.step
+            )
+            if step is None:
+                return FrameIteratorAll(self)
+            else:
+                return StreamFrameIteratorSliced(self, step)
+        else:
+            raise TypeError(
+                "Streamed trajectories must be an indexed using a slice"
+            )
+
+    def check_slice_indices(self, start, stop, step):
+        if start is not None:
+            raise ValueError(
+                "{}: Cannot expect a start index from a stream, 'start' must be None".format(
+                    self.__class__.__name__
+                )
+            )
+        if stop is not None:
+            raise ValueError(
+                "{}: Cannot expect a stop index from a stream, 'stop' must be None".format(
+                    self.__class__.__name__
+                )
+            )
+        if step is not None:
+            if isinstance(step, numbers.Integral):
+                if step < 1:
+                    raise ValueError(
+                        "{}: Cannot go backwards in a stream, 'step' must be > 0".format(
+                            self.__class__.__name__
+                        )
+                    )
+            else:
+                raise ValueError(
+                    "{}: 'step' must be an integer".format(
+                        self.__class__.__name__
+                    )
+                )
+
+        return start, stop, step
+
+    def __getstate__(self):
+        raise NotImplementedError(
+            "{} does not support pickling".format(self.__class__.__name__)
+        )
+
+    def __setstate__(self, state: object):
+        raise NotImplementedError(
+            "{} does not support pickling".format(self.__class__.__name__)
+        )
+
+    def __repr__(self):
+        return (
+            "<{cls} {fname} with continuous stream of {natoms} atoms>"
+            "".format(
+                cls=self.__class__.__name__,
+                fname=self.filename,
+                natoms=self.n_atoms,
+            )
+        )
+
+
+class StreamFrameIteratorSliced(FrameIteratorBase):
+
+    def __init__(self, trajectory, step):
+        super().__init__(trajectory)
+        self._step = step
+
+    def __iter__(self):
+        # Calling reopen tells reader
+        # it can't be reopened again
+        self.trajectory._reopen()
+        return self
+
+    def __next__(self):
+        try:
+            # Burn the timesteps until we reach the desired step
+            # Don't use next() to avoid unnecessary transformations
+            while self.trajectory._frame + 1 % self.step != 0:
+                self.trajectory._read_next_timestep()
+        except (EOFError, IOError):
+            # Don't rewind here like we normally would
+            raise StopIteration from None
+
+        return self.trajectory.next()
+
+    def __len__(self):
+        raise RuntimeError(
+            "{} has unknown length".format(self.__class__.__name__)
+        )
+
+    def __getitem__(self, frame):
+        raise RuntimeError("Sliced iterator does not support indexing")
+
+    @property
+    def step(self):
+        return self._step

imdclient/tests/base.py

@@ -0,0 +1,122 @@
+from imdclient.IMDClient import IMDClient
+import pytest
+from pathlib import Path
+import os
+import signal
+import subprocess
+import time
+from numpy.testing import (
+    assert_allclose,
+)
+import numpy as np
+
+import logging
+
+logger = logging.getLogger("imdclient.IMDClient")
+
+
+def assert_allclose_with_logging(a, b, rtol=1e-07, atol=0, equal_nan=False):
+    """
+    Custom function to compare two arrays element-wise, similar to np.testing.assert_allclose,
+    but logs all non-matching values.
+
+    Parameters:
+    a, b : array_like
+        Input arrays to compare.
+    rtol : float
+        Relative tolerance.
+    atol : float
+        Absolute tolerance.
+    equal_nan : bool
+        Whether to compare NaNs as equal.
+    """
+    # Convert inputs to numpy arrays
+    a = np.asarray(a)
+    b = np.asarray(b)
+
+    # Compute the absolute difference
+    diff = np.abs(a - b)
+
+    # Check if values are within tolerance
+    not_close = diff > (atol + rtol * np.abs(b))
+
+    # Check if there are any NaNs and handle them if necessary
+    if equal_nan:
+        nan_mask = np.isnan(a) & np.isnan(b)
+        not_close &= ~nan_mask
+
+    # Log all the values that are not close
+    if np.any(not_close):
+        print("The following values do not match within tolerance:")
+        for idx in np.argwhere(not_close):
+            logger.debug(
+                f"a[{tuple(idx)}]: {a[tuple(idx)]}, b[{tuple(idx)}]: {b[tuple(idx)]}, diff: {diff[tuple(idx)]}"
+            )
+        # Optionally raise an error after logging if you want it to behave like assert
+        raise AssertionError("Arrays are not almost equal.")
+    else:
+        print("All values are within tolerance.")
+
+
+class IMDv3IntegrationTest:
+
+    @pytest.fixture()
+    def run_sim_and_wait(self, tmp_path, command, match_string):
+        old_cwd = Path.cwd()
+        os.chdir(tmp_path)
+        p = subprocess.Popen(
+            command,
+            stdout=subprocess.PIPE,
+            stderr=subprocess.STDOUT,
+            shell=True,
+            text=True,
+            bufsize=0,
+            preexec_fn=os.setsid,
+        )
+        t = time.time()
+
+        for stdout_line in iter(p.stdout.readline, ""):
+            logger.debug(f"stdout: {stdout_line}")
+            if match_string in stdout_line:
+
+                break
+            if time.time() - t > 10:
+                raise TimeoutError("Timeout waiting for match string")
+
+        logger.debug("Match string found")
+        yield p
+        os.chdir(old_cwd)
+        os.killpg(os.getpgid(p.pid), signal.SIGTERM)
+
+    @pytest.fixture()
+    def client(self, run_sim_and_wait, universe):
+        client = IMDClient("localhost", 8888, universe.trajectory.n_atoms)
+        yield client
+        client.stop()
+
+    def test_compare_imd_to_true_traj(self, universe, client, first_frame):
+        imdsinfo = client.get_imdsessioninfo()
+
+        for ts in universe.trajectory[first_frame:]:
+            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:
+                assert_allclose_with_logging(
+                    imdf.box,
+                    ts.triclinic_dimensions,
+                    atol=1e-03,
+                )
+            if imdsinfo.positions:
+                assert_allclose_with_logging(
+                    imdf.positions, ts.positions, atol=1e-03
+                )
+            if imdsinfo.velocities:
+                assert_allclose_with_logging(
+                    imdf.velocities, ts.velocities, atol=1e-03
+                )
+            if imdsinfo.forces:
+                assert_allclose_with_logging(
+                    imdf.forces, ts.forces, atol=1e-03
+                )

imdclient/tests/conftest.py

@@ -0,0 +1,38 @@
+"""
+Global pytest fixtures
+"""
+
+
+# Command line arguments for 'test_manual.py'
+def pytest_addoption(parser):
+    parser.addoption(
+        "--topol_arg",
+        action="store",
+    )
+    parser.addoption(
+        "--traj_arg",
+        action="store",
+    )
+    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
+    first_frame = metafunc.config.option.first_frame_arg
+
+    if all(
+        arg in metafunc.fixturenames
+        for arg in ["topol_arg", "traj_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>', "
+                + "'--first_frame_arg <first traj frame to compare to IMD>"
+            )
+        metafunc.parametrize("topol_arg", [topol])
+        metafunc.parametrize("traj_arg", [traj])
+        metafunc.parametrize("first_frame_arg", [first_frame])

imdclient/tests/datafiles.py

@@ -0,0 +1,34 @@
+"""
+Location of data files
+======================
+
+Use as ::
+
+    from imdclient.tests.datafiles import *
+
+"""
+
+__all__ = ["LAMMPS_IN", "LAMMPS_TOPOL", "GROMACS_TRAJ", "GROMACS_MDP"]
+
+from importlib import resources
+from pathlib import Path
+
+_data_ref = resources.files("imdclient.data")
+
+LAMMPS_TOPOL = (_data_ref / "lammps" / "md" / "lammps_topol.data").as_posix()
+LAMMPS_IN = (_data_ref / "lammps" / "md" / "lammps_v3.in").as_posix()
+LAMMPS_TRAJ = (_data_ref / "lammps" / "md" / "lammps_trj.h5md").as_posix()
+GROMACS_TRAJ = (
+    _data_ref / "gromacs" / "md" / "gromacs_v3_nst1.trr"
+).as_posix()
+GROMACS_TOPOL = (
+    _data_ref / "gromacs" / "md" / "gromacs_struct.gro"
+).as_posix()
+GROMACS_TPR = (_data_ref / "gromacs" / "md" / "gromacs_v3_nst1.tpr").as_posix()
+NAMD_TOPOL = (_data_ref / "namd" / "md" / "alanin.pdb").as_posix()
+NAMD_CONF = (_data_ref / "namd" / "md" / "namd_v3.namd").as_posix()
+NAMD_TRAJ = (_data_ref / "namd" / "md" / "alanin.dcd").as_posix()
+NAMD_PARAMS = (_data_ref / "namd" / "md" / "alanin.params").as_posix()
+NAMD_PSF = (_data_ref / "namd" / "md" / "alanin.psf").as_posix()
+
+del resources

imdclient/tests/server.py

@@ -0,0 +1,212 @@
+import socket
+import struct
+import numpy as np
+import threading
+import time
+from imdclient.IMDProtocol import (
+    IMDHeaderType,
+    create_header_bytes,
+    create_energy_bytes,
+    IMDHEADERSIZE,
+)
+import logging
+from imdclient.IMDClient import sock_contains_data, read_into_buf, IMDHeader
+
+logger = logging.getLogger("imdclient.IMDClient")
+
+
+class InThreadIMDServer:
+    """Server operates in the same thread as the client (except for the method
+    'handshake_sequence', which spawns a new thread to get past the
+    reader connecting to the server).
+
+    Note that the trajectory units of the provided traj must be in MDAnalysis units
+    """
+
+    def __init__(self, traj):
+        self.traj = traj
+        self.listen_socket = None
+        self.conn = None
+        self.accept_thread = None
+
+    def set_imdsessioninfo(self, imdsinfo):
+        self.imdsinfo = imdsinfo
+
+    def handshake_sequence(self, host, port, first_frame=True):
+        s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
+        s.bind((host, port))
+        logger.debug(f"InThreadIMDServer: Listening on {host}:{port}")
+        s.listen(60)
+        self.listen_socket = s
+
+        self.accept_thread = threading.Thread(
+            target=self._accept_handshake, args=(first_frame,)
+        )
+        self.accept_thread.start()
+
+    def _accept_handshake(self, first_frame):
+        """Accepts the connection and sends the handshake, imdsessioninfo (in the case of v3),
+        the first frame (if first_frame is `True`), and expects IMD_GO"""
+        waited = 0
+
+        if sock_contains_data(self.listen_socket, 5):
+            logger.debug(f"InThreadIMDServer: Accepting connection")
+            self.conn, _ = self.listen_socket.accept()
+            # NOTE: may need to reorganize this
+            self.conn.settimeout(5)
+            if self.imdsinfo.version == 2:
+                self._send_handshakeV2()
+            elif self.imdsinfo.version == 3:
+                self._send_handshakeV3()
+            self.expect_packet(IMDHeaderType.IMD_GO)
+            if first_frame:
+                self.send_frame(0)
+            return
+        else:
+            # IMDReader will fail out if it fails to connect
+            return
+
+    def _send_handshakeV2(self):
+        header = struct.pack("!i", IMDHeaderType.IMD_HANDSHAKE.value)
+        header += struct.pack(f"{self.imdsinfo.endianness}i", 2)
+        self.conn.sendall(header)
+
+    def _send_handshakeV3(self):
+        logger.debug(f"InThreadIMDServer: Sending handshake V3")
+        packet = struct.pack("!i", IMDHeaderType.IMD_HANDSHAKE.value)
+        packet += struct.pack(f"{self.imdsinfo.endianness}i", 3)
+        self.conn.sendall(packet)
+
+        sinfo = struct.pack("!ii", IMDHeaderType.IMD_SESSIONINFO.value, 7)
+        time = 1 if self.imdsinfo.time else 0
+        energies = 1 if self.imdsinfo.energies else 0
+        box = 1 if self.imdsinfo.box else 0
+        positions = 1 if self.imdsinfo.positions else 0
+        velocities = 1 if self.imdsinfo.velocities else 0
+        forces = 1 if self.imdsinfo.forces else 0
+        wrapped_coords = 0
+        sinfo += struct.pack(
+            f"{self.imdsinfo.endianness}BBBBBBB",
+            time,
+            energies,
+            box,
+            positions,
+            wrapped_coords,
+            velocities,
+            forces
+        )
+        logger.debug(f"InThreadIMDServer: Sending session info")
+        self.conn.sendall(sinfo)
+
+    def join_accept_thread(self):
+        self.accept_thread.join()
+
+    def _expect_go(self):
+        logger.debug(f"InThreadIMDServer: Waiting for go")
+        head_buf = bytearray(IMDHEADERSIZE)
+        read_into_buf(self.conn, head_buf)
+        header = IMDHeader(head_buf)
+        if header.type != IMDHeaderType.IMD_GO:
+            raise ValueError("Expected IMD_GO packet, got something else")
+
+    def send_frames(self, start, end):
+        for i in range(start, end):
+            self.send_frame(i)
+
+    def send_frame(self, i):
+        endianness = self.imdsinfo.endianness
+
+        if self.imdsinfo.time:
+            time_header = create_header_bytes(IMDHeaderType.IMD_TIME, 1)
+            time = struct.pack(
+                f"{endianness}ddQ",
+                self.traj[i].dt,
+                self.traj[i].time,
+                self.traj[i].frame,
+            )
+
+            self.conn.sendall(time_header + time)
+        
+        if self.imdsinfo.energies:
+            energy_header = create_header_bytes(IMDHeaderType.IMD_ENERGIES, 1)
+            energies = create_energy_bytes(
+                i,
+                i + 1,
+                i + 2,
+                i + 3,
+                i + 4,
+                i + 5,
+                i + 6,
+                i + 7,
+                i + 8,
+                i + 9,
+                endianness,
+            )
+            self.conn.sendall(energy_header + energies)
+
+        if self.imdsinfo.box:
+            box_header = create_header_bytes(IMDHeaderType.IMD_BOX, 1)
+            box = np.ascontiguousarray(
+                self.traj[i].triclinic_dimensions, dtype=f"{endianness}f"
+            ).tobytes()
+
+            self.conn.sendall(box_header + box)
+
+        if self.imdsinfo.positions:
+            pos_header = create_header_bytes(
+                IMDHeaderType.IMD_FCOORDS, self.traj.n_atoms
+            )
+
+            pos = np.ascontiguousarray(
+                self.traj[i].positions, dtype=f"{endianness}f"
+            ).tobytes()
+
+            self.conn.sendall(pos_header + pos)
+
+        if self.imdsinfo.velocities:
+            vel_header = create_header_bytes(
+                IMDHeaderType.IMD_VELOCITIES, self.traj.n_atoms
+            )
+            vel = np.ascontiguousarray(
+                self.traj[i].velocities, dtype=f"{endianness}f"
+            ).tobytes()
+
+            self.conn.sendall(vel_header + vel)
+
+        if self.imdsinfo.forces:
+            force_header = create_header_bytes(
+                IMDHeaderType.IMD_FORCES, self.traj.n_atoms
+            )
+            force = np.ascontiguousarray(
+                self.traj[i].forces, dtype=f"{endianness}f"
+            ).tobytes()
+
+            self.conn.sendall(force_header + force)
+
+        
+    def expect_packet(self, packet_type, expected_length=None):
+        head_buf = bytearray(IMDHEADERSIZE)
+        read_into_buf(self.conn, head_buf)
+        header = IMDHeader(head_buf)
+        if header.type != packet_type:
+            raise ValueError(
+                f"Expected {packet_type} packet, got {header.type}"
+            )
+        if expected_length is not None and header.length != expected_length:
+            raise ValueError(
+                f"Expected packet length {expected_length}, got {header.length}"
+            )
+
+    def disconnect(self):
+        self.conn.shutdown(socket.SHUT_RD)
+        self.conn.close()
+
+    def cleanup(self):
+        try:
+            self.conn.close()
+        except:
+            pass
+        try:
+            self.listen_socket.close()
+        except:
+            pass

imdclient/tests/test_gromacs.py

@@ -0,0 +1,33 @@
+import MDAnalysis as mda
+import pytest
+import logging
+from .base import IMDv3IntegrationTest
+from .datafiles import GROMACS_TOPOL, GROMACS_TRAJ, GROMACS_TPR
+
+logger = logging.getLogger("imdclient.IMDClient")
+file_handler = logging.FileHandler("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 TestIMDv3Gromacs(IMDv3IntegrationTest):
+
+    @pytest.fixture()
+    def command(self):
+        return f". /usr/local/gromacs/bin/GMXRC && gmx mdrun -s {GROMACS_TPR} -imdport 8888 -imdwait"
+
+    @pytest.fixture()
+    def match_string(self):
+        return "IMD: Will wait until I have a connection and IMD_GO orders."
+
+    @pytest.fixture()
+    def first_frame(self):
+        return 0
+
+    @pytest.fixture()
+    def universe(self):
+        return mda.Universe(GROMACS_TOPOL, GROMACS_TRAJ)