imdclient 0.1.1__py3-none-any.whl → 0.1.3__py3-none-any.whl

imdclient/tests/base.py

@@ -1,4 +1,5 @@
 from imdclient.IMDClient import IMDClient
+from imdclient.IMD import IMDReader
 import pytest
 from pathlib import Path
 import os
@@ -9,8 +10,11 @@ from numpy.testing import (
     assert_allclose,
 )
 import numpy as np
-
+import docker
 import logging
+import shutil
+import MDAnalysis as mda
+from .utils import get_free_port
 
 logger = logging.getLogger("imdclient.IMDClient")
 
@@ -61,62 +65,146 @@ def assert_allclose_with_logging(a, b, rtol=1e-07, atol=0, equal_nan=False):
 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,
+    def setup_command(self):
+        return None
+
+    @pytest.fixture()
+    def port(self):
+        yield get_free_port()
+
+    @pytest.fixture()
+    def docker_client(
+        self,
+        tmp_path,
+        input_files,
+        setup_command,
+        simulation_command,
+        port,
+    ):
+        # In CI, container process needs access to tmp_path
+        tmp_path.chmod(0o777)
+        docker_client = docker.from_env()
+        img = docker_client.images.pull(
+            "ghcr.io/becksteinlab/streaming-md-docker:main-Common-CPU"
+        )
+        # Copy input files into tmp_path
+        for inp in input_files:
+            shutil.copy(inp, tmp_path)
+
+        cmdstring = "cd '/tmp'"
+        # Run the setup command, if any
+        if setup_command is not None:
+            # This should be blocking
+            cmdstring += " && " + setup_command
+
+        cmdstring += " && " + simulation_command
+
+        # Start the container, mount tmp_path, run simulation
+        container = docker_client.containers.run(
+            img,
+            f"/bin/sh -c '{cmdstring}'",
+            detach=True,
+            volumes={tmp_path.as_posix(): {"bind": "/tmp", "mode": "rw"}},
+            ports={"8888/tcp": port},
+            name="sim",
+            remove=True,
         )
-        t = time.time()
 
-        for stdout_line in iter(p.stdout.readline, ""):
-            logger.debug(f"stdout: {stdout_line}")
-            if match_string in stdout_line:
+        # For now, just wait 30 seconds
+        # life is too short to figure out how to redirect all stdout from inside
+        # a container
+        time.sleep(30)
+
+        yield
+        try:
+            container.stop()
+        except docker.errors.NotFound:
+            pass
+
+    @pytest.fixture()
+    def imd_u(self, docker_client, topol, tmp_path, port):
+        u = mda.Universe((tmp_path / topol), f"imd://localhost:{port}")
+        with mda.Writer(
+            (tmp_path / "imd.trr").as_posix(), u.trajectory.n_atoms
+        ) as w:
+            for ts in u.trajectory:
+                w.write(u.atoms)
+        yield mda.Universe((tmp_path / topol), (tmp_path / "imd.trr"))
+
+    @pytest.fixture()
+    def true_u(self, topol, traj, imd_u, tmp_path):
+        u = mda.Universe(
+            (tmp_path / topol),
+            (tmp_path / traj),
+        )
+        yield u
 
-                break
-            if time.time() - t > 10:
-                raise TimeoutError("Timeout waiting for match string")
+    @pytest.fixture()
+    def comp_time(self):
+        return True
 
-        logger.debug("Match string found")
-        yield p
-        os.chdir(old_cwd)
-        os.killpg(os.getpgid(p.pid), signal.SIGTERM)
+    @pytest.fixture()
+    def comp_dt(self):
+        return True
 
     @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:
+    def comp_step(self):
+        return True
+
+    def test_compare_imd_to_true_traj(
+        self, imd_u, true_u, first_frame, comp_time, comp_dt, comp_step
+    ):
+        for i in range(first_frame, len(true_u.trajectory)):
+            if comp_time:
+                assert_allclose(
+                    true_u.trajectory[i].time,
+                    imd_u.trajectory[i - first_frame].time,
+                    atol=1e-03,
+                )
+            if comp_dt:
+                assert_allclose(
+                    true_u.trajectory[i].dt,
+                    imd_u.trajectory[i - first_frame].dt,
+                    atol=1e-03,
+                )
+            if comp_step:
+                assert_allclose(
+                    true_u.trajectory[i].data["step"],
+                    imd_u.trajectory[i - first_frame].data["step"],
+                )
+            if (
+                true_u.trajectory[i].dimensions is not None
+                and imd_u.trajectory[i - first_frame].dimensions is not None
+            ):
                 assert_allclose_with_logging(
-                    imdf.box,
-                    ts.triclinic_dimensions,
+                    true_u.trajectory[i].dimensions,
+                    imd_u.trajectory[i - first_frame].dimensions,
                     atol=1e-03,
                 )
-            if imdsinfo.positions:
+            if (
+                true_u.trajectory[i].has_positions
+                and imd_u.trajectory[i - first_frame].has_positions
+            ):
                 assert_allclose_with_logging(
-                    imdf.positions, ts.positions, atol=1e-03
+                    true_u.trajectory[i].positions,
+                    imd_u.trajectory[i - first_frame].positions,
+                    atol=1e-03,
                 )
-            if imdsinfo.velocities:
+            if (
+                true_u.trajectory[i].has_velocities
+                and imd_u.trajectory[i - first_frame].has_velocities
+            ):
                 assert_allclose_with_logging(
-                    imdf.velocities, ts.velocities, atol=1e-03
+                    true_u.trajectory[i].velocities,
+                    imd_u.trajectory[i - first_frame].velocities,
+                    atol=1e-03,
                 )
-            if imdsinfo.forces:
+            if (
+                true_u.trajectory[i].has_forces
+                and imd_u.trajectory[i - first_frame].has_forces
+            ):
                 assert_allclose_with_logging(
-                    imdf.forces, ts.forces, atol=1e-03
+                    true_u.trajectory[i].forces,
+                    imd_u.trajectory[i - first_frame].forces,
+                    atol=1e-03,
                 )

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/datafiles.py

@@ -16,18 +16,26 @@ 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"
+LAMMPS_IN_NST_1 = (
+    _data_ref / "lammps" / "md" / "lammps_v3_nst_1.in"
 ).as_posix()
-GROMACS_TOPOL = (
-    _data_ref / "gromacs" / "md" / "gromacs_struct.gro"
+LAMMPS_IN_NST_8 = (
+    _data_ref / "lammps" / "md" / "lammps_v3_nst_8.in"
 ).as_posix()
-GROMACS_TPR = (_data_ref / "gromacs" / "md" / "gromacs_v3_nst1.tpr").as_posix()
+
+
+GROMACS_GRO = (_data_ref / "gromacs" / "md" / "gromacs_struct.gro").as_posix()
+GROMACS_MDP_NST_1 = (
+    _data_ref / "gromacs" / "md" / "gromacs_v3_nst1.mdp"
+).as_posix()
+GROMACS_MDP_NST_8 = (
+    _data_ref / "gromacs" / "md" / "gromacs_v3_nst8.mdp"
+).as_posix()
+GROMACS_TOP = (_data_ref / "gromacs" / "md" / "gromacs_v3.top").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_CONF_NST_1 = (_data_ref / "namd" / "md" / "namd_v3_nst_1.namd").as_posix()
+NAMD_CONF_NST_8 = (_data_ref / "namd" / "md" / "namd_v3_nst_8.namd").as_posix()
 NAMD_PARAMS = (_data_ref / "namd" / "md" / "alanin.params").as_posix()
 NAMD_PSF = (_data_ref / "namd" / "md" / "alanin.psf").as_posix()
 

imdclient/tests/docker_testing/docker.md

@@ -0,0 +1,25 @@
+# Running simulation engines compiled for GPU acceleration from a container
+
+Ensure [docker](https://www.docker.com/) and the [NVIDIA container toolkit](https://docs.nvidia.com/datacenter/cloud-native/container-toolkit/latest/install-guide.html) are installed.
+
+
+To run the container:
+```bash
+docker pull ghcr.io/becksteinlab/streaming-md-docker:main-Common-GPU
+
+docker run -v $PWD/imdclient/data:/home/conda:rw -it --runtime=nvidia --gpus=all \
+    ghcr.io/becksteinlab/streaming-md-docker:main-Common-GPU
+```
+
+To run each simulation engine with repository simulation configurations:
+```bash
+cd /home/conda/namd/md
+namd3 +devices 0 namd_v3.namd
+
+cd /home/conda/gromacs/md
+gmx grompp -f gromacs_v3_nst1.mdp -c gromacs_struct.gro -p gromacs_v3.top
+gmx mdrun -s topol.tpr -nb gpu
+
+cd /home/conda/lammps/md
+lmp -sf gpu < lammps_v3_nst_1.in
+```

imdclient/tests/test_gromacs.py

@@ -2,10 +2,16 @@ import MDAnalysis as mda
 import pytest
 import logging
 from .base import IMDv3IntegrationTest
-from .datafiles import GROMACS_TOPOL, GROMACS_TRAJ, GROMACS_TPR
+from .datafiles import (
+    GROMACS_GRO,
+    GROMACS_TOP,
+    GROMACS_MDP_NST_1,
+    GROMACS_MDP_NST_8,
+)
+from pathlib import Path
 
 logger = logging.getLogger("imdclient.IMDClient")
-file_handler = logging.FileHandler("test.log")
+file_handler = logging.FileHandler("gromacs_test.log")
 formatter = logging.Formatter(
     "%(asctime)s - %(name)s - %(levelname)s - %(message)s"
 )
@@ -16,18 +22,34 @@ logger.setLevel(logging.DEBUG)
 
 class TestIMDv3Gromacs(IMDv3IntegrationTest):
 
+    @pytest.fixture(params=[GROMACS_MDP_NST_1, GROMACS_MDP_NST_8])
+    def mdp(self, request):
+        return request.param
+
     @pytest.fixture()
-    def command(self):
-        return f". /usr/local/gromacs/bin/GMXRC && gmx mdrun -s {GROMACS_TPR} -imdport 8888 -imdwait"
+    def setup_command(self, mdp):
+        return f"gmx grompp -f {Path(mdp).name} -c {Path(GROMACS_GRO).name} -p {Path(GROMACS_TOP).name} -o topol.tpr"
 
     @pytest.fixture()
-    def match_string(self):
-        return "IMD: Will wait until I have a connection and IMD_GO orders."
+    def simulation_command(self):
+        return f"gmx mdrun -s topol.tpr -o ci.trr -imdport 8888 -imdwait"
 
     @pytest.fixture()
-    def first_frame(self):
-        return 0
+    def input_files(self, mdp):
+        return [GROMACS_TOP, mdp, GROMACS_GRO]
+
+    @pytest.fixture()
+    def traj(self):
+        return "ci.trr"
 
     @pytest.fixture()
-    def universe(self):
-        return mda.Universe(GROMACS_TOPOL, GROMACS_TRAJ)
+    def topol(self):
+        return Path(GROMACS_GRO).name
+
+    # @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

imdclient/tests/test_imdclient.py

@@ -148,3 +148,54 @@ class TestIMDClientV3:
             client.get_imdframe()
         # server should receive disconnect from client (though it doesn't have to do anything)
         server.expect_packet(IMDHeaderType.IMD_DISCONNECT)
+
+
+class TestIMDClientV3ContextManager:
+    @pytest.fixture
+    def port(self):
+        return get_free_port()
+
+    @pytest.fixture
+    def universe(self):
+        return mda.Universe(COORDINATES_TOPOLOGY, COORDINATES_H5MD)
+
+    @pytest.fixture
+    def imdsinfo(self):
+        return create_default_imdsinfo_v3()
+
+    @pytest.fixture
+    def server(self, universe, imdsinfo, port):
+        server = InThreadIMDServer(universe.trajectory)
+        server.set_imdsessioninfo(imdsinfo)
+        yield server
+        server.cleanup()
+
+    def test_context_manager_traj_unchanged(self, server, port, universe):
+        server.handshake_sequence("localhost", port, first_frame=False)
+
+        i = 0
+        with IMDClient(
+            "localhost",
+            port,
+            universe.trajectory.n_atoms,
+        ) as client:
+            server.send_frames(0, 5)
+            while i < 5:
+
+                imdf = client.get_imdframe()
+                assert_allclose(universe.trajectory[i].time, imdf.time)
+                assert_allclose(universe.trajectory[i].dt, imdf.dt)
+                assert_allclose(universe.trajectory[i].data["step"], imdf.step)
+                assert_allclose(
+                    universe.trajectory[i].positions, imdf.positions
+                )
+                assert_allclose(
+                    universe.trajectory[i].velocities, imdf.velocities
+                )
+                assert_allclose(universe.trajectory[i].forces, imdf.forces)
+                assert_allclose(
+                    universe.trajectory[i].triclinic_dimensions, imdf.box
+                )
+                i += 1
+        server.expect_packet(IMDHeaderType.IMD_DISCONNECT)
+        assert i == 5

imdclient/tests/test_imdreader.py

@@ -31,7 +31,7 @@ import pytest
 from MDAnalysis.transformations import translate
 import pickle
 
-from imdclient.IMDREADER import IMDReader
+from imdclient.IMD import IMDReader
 
 logger = logging.getLogger("imdclient.IMDClient")
 file_handler = logging.FileHandler("test.log")
@@ -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
@@ -171,9 +171,7 @@ class TestIMDReaderBaseAPI(MultiframeReaderTest):
             decimal=ref.prec,
         )
 
-    @pytest.mark.skip(
-        reason="Stream-based reader can only be read iteratively"
-    )
+    @pytest.mark.skip(reason="Stream-based reader can only be read iteratively")
     def test_changing_dimensions(self, ref, reader):
         if ref.changing_dimensions:
             reader.rewind()
@@ -585,7 +583,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)
@@ -642,3 +640,19 @@ class TestStreamIteration:
         with pytest.raises(RuntimeError):
             for ts in sub_sliced_reader:
                 pass
+
+
+def test_n_atoms_mismatch():
+    universe = mda.Universe(COORDINATES_TOPOLOGY, COORDINATES_H5MD)
+    port = get_free_port()
+    server = InThreadIMDServer(universe.trajectory)
+    server.set_imdsessioninfo(create_default_imdsinfo_v3())
+    server.handshake_sequence("localhost", port, first_frame=True)
+    with pytest.raises(
+        EOFError,
+        match="IMDProducer: Expected n_atoms value 6, got 5. Ensure you are using the correct topology file.",
+    ):
+        IMDReader(
+            f"imd://localhost:{port}",
+            n_atoms=universe.trajectory.n_atoms + 1,
+        )

imdclient/tests/test_lammps.py

@@ -1,11 +1,12 @@
 import MDAnalysis as mda
 import pytest
+from pathlib import Path
 import logging
 from .base import IMDv3IntegrationTest
-from .datafiles import LAMMPS_IN, LAMMPS_TOPOL, LAMMPS_TRAJ
+from .datafiles import LAMMPS_TOPOL, LAMMPS_IN_NST_1, LAMMPS_IN_NST_8
 
 logger = logging.getLogger("imdclient.IMDClient")
-file_handler = logging.FileHandler("test.log")
+file_handler = logging.FileHandler("lammps_test.log")
 formatter = logging.Formatter(
     "%(asctime)s - %(name)s - %(levelname)s - %(message)s"
 )
@@ -16,23 +17,67 @@ logger.setLevel(logging.DEBUG)
 
 class TestIMDv3Lammps(IMDv3IntegrationTest):
 
+    @pytest.fixture(params=[LAMMPS_IN_NST_1, LAMMPS_IN_NST_8])
+    def inp(self, request):
+        return request.param
+
+    @pytest.fixture()
+    def simulation_command(self, inp):
+        return f"lmp < {Path(inp).name}"
+
+    @pytest.fixture()
+    def topol(self):
+        return Path(LAMMPS_TOPOL).name
+
+    @pytest.fixture()
+    def traj(self):
+        return "lammps_trj.h5md"
+
     @pytest.fixture()
-    def command(self):
-        return f"lmp < {LAMMPS_IN}"
+    def input_files(self, inp):
+        return [inp, LAMMPS_TOPOL]
+
+    # @pytest.fixture()
+    # def match_string(self):
+    #     return "Waiting for IMD connection on port 8888"
 
     @pytest.fixture()
-    def match_string(self):
-        return "Waiting for IMD connection on port 8888"
+    def first_frame(self, inp):
+        if inp == LAMMPS_IN_NST_1:
+            return 1
+        else:
+            return 0
 
+    # Not present in lammps-produced H5MD
     @pytest.fixture()
-    def first_frame(self):
-        return 1
+    def comp_dt(self):
+        return False
 
+    # This must wait until after imd stream has ended
     @pytest.fixture()
-    def universe(self):
-        return mda.Universe(
-            LAMMPS_TOPOL,
-            LAMMPS_TRAJ,
+    def true_u(self, topol, traj, imd_u, tmp_path):
+        u = mda.Universe(
+            (tmp_path / topol),
+            (tmp_path / traj),
             atom_style="id type x y z",
             convert_units=False,
         )
+        yield u
+
+    @pytest.fixture()
+    def imd_u(self, docker_client, topol, tmp_path, port):
+        u = mda.Universe(
+            (tmp_path / topol),
+            f"imd://localhost:{port}",
+            atom_style="id type x y z",
+        )
+        with mda.Writer(
+            (tmp_path / "imd.trr").as_posix(), u.trajectory.n_atoms
+        ) as w:
+            for ts in u.trajectory:
+                w.write(u.atoms)
+        yield mda.Universe(
+            (tmp_path / topol),
+            (tmp_path / "imd.trr"),
+            atom_style="id type x y z",
+        )