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

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",
+        )

imdclient/tests/test_manual.py

@@ -1,14 +1,11 @@
-from imdclient.IMDREADER import IMDReader
-import pytest
+import imdclient
+from imdclient.IMD import IMDReader
 import MDAnalysis as mda
-from MDAnalysisTests.coordinates.base import assert_timestep_almost_equal
-from numpy.testing import (
-    assert_allclose,
-)
+from numpy.testing import assert_allclose
 import numpy as np
-from .base import assert_allclose_with_logging
 from pathlib import Path
-
+import time
+import argparse
 import logging
 
 logger = logging.getLogger("imdclient.IMDClient")
@@ -20,72 +17,233 @@ file_handler.setFormatter(formatter)
 logger.addHandler(file_handler)
 logger.setLevel(logging.DEBUG)
 
+"""
+Tool for running IMDv3 integration tests via the command line.
 
-class TestIMDv3Manual:
-    """
-    Tool for running IMDv3 integration tests via the command line.
+To use, start the simulation, wait for it to be ready for an IMD connection,
+and then run this command relative to the root of the cloned respository:
 
-    To use, start the simulation, wait for it to be ready for an IMD connection,
-    and then run this command relative to the root of the cloned respository:
+python imdclient/tests/test_manual.py \
+    --topol_path <path/to/topology> \
+    --traj_path <path/to/trajectory> \
+    --first_frame <first frame to compare> \
+    --tmp_path <temporary directory for IMD trajectory>
 
-    pytest -s imdclient/tests/test_manual.py \
-        --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 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)
+"""
+
+
+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.
 
-    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)
+    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.")
 
-    @pytest.fixture()
-    def true_u(self, imd_u, topol_path_arg, traj_path_arg):
-        return mda.Universe(topol_path_arg, traj_path_arg)
-
-    @pytest.fixture()
-    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"
+
+def load_true_universe(topol_path, traj_path):
+    if topol_path.endswith(".data"):
+        return mda.Universe(
+            topol_path,
+            traj_path,
+            atom_style="id type x y z",
+            convert_units=False,
+        )
+    return mda.Universe(topol_path, traj_path)
+
+
+def load_imd_universe(topol_path, tmp_path):
+    # Pass atom_style (ignored if not using LAMMPS topol)
+    tmp_u = mda.Universe(
+        topol_path,
+        "imd://localhost:8888",
+        atom_style="id type x y z",
+    )
+    tmp_traj_file = f"{tmp_path}/imd_test_traj.trr"
+    with mda.Writer(tmp_traj_file, tmp_u.atoms.n_atoms) as w:
+        for ts in tmp_u.trajectory:
+            w.write(tmp_u.atoms)
+    time.sleep(10)  # Give MPI ranks a chance to release FD
+    return mda.Universe(topol_path, tmp_traj_file, atom_style="id type x y z")
+
+
+def test_compare_imd_to_true_traj_vel(imd_u, true_u_vel, first_frame):
+    for i in range(first_frame, len(true_u_vel.trajectory)):
+        # Manually convert unit
+        assert_allclose_with_logging(
+            true_u_vel.trajectory[i].positions * 20.45482706,
+            imd_u.trajectory[i - first_frame].velocities,
+            atol=1e-03,
         )
-        yield imd_u
 
-    def test_compare_imd_to_true_traj(self, true_u, imd_u, first_frame_arg):
 
-        for i in range(first_frame_arg, len(true_u.trajectory)):
+def test_compare_imd_to_true_traj_forces(imd_u, true_u_force, first_frame):
+    for i in range(first_frame, len(true_u_force.trajectory)):
+        assert_allclose_with_logging(
+            true_u_force.trajectory[i].positions,
+            imd_u.trajectory[i - first_frame].forces,
+            atol=1e-03,
+        )
+
+
+def test_compare_imd_to_true_traj(imd_u, true_u, first_frame, vel, force, dt):
+    for i in range(first_frame, len(true_u.trajectory)):
+        assert_allclose(
+            true_u.trajectory[i].time,
+            imd_u.trajectory[i - first_frame].time,
+            atol=1e-03,
+        )
+        if dt:
             assert_allclose(
-                true_u.trajectory[i].time, imd_u.trajectory[i].time, atol=1e-03
+                true_u.trajectory[i].dt,
+                imd_u.trajectory[i - first_frame].dt,
+                atol=1e-03,
             )
-            assert_allclose(
-                true_u.trajectory[i].data["step"],
-                imd_u.trajectory[i].data["step"],
+        assert_allclose(
+            true_u.trajectory[i].data["step"],
+            imd_u.trajectory[i - first_frame].data["step"],
+        )
+        assert_allclose_with_logging(
+            true_u.trajectory[i].dimensions,
+            imd_u.trajectory[i - first_frame].dimensions,
+            atol=1e-03,
+        )
+        assert_allclose_with_logging(
+            true_u.trajectory[i].positions,
+            imd_u.trajectory[i - first_frame].positions,
+            atol=1e-03,
+        )
+        if vel:
+            assert_allclose_with_logging(
+                true_u.trajectory[i].velocities,
+                imd_u.trajectory[i - first_frame].velocities,
+                atol=1e-03,
+            )
+        if force:
+            assert_allclose_with_logging(
+                true_u.trajectory[i].forces,
+                imd_u.trajectory[i - first_frame].forces,
+                atol=1e-03,
+            )
+
+
+def main():
+    parser = argparse.ArgumentParser(description="IMDv3 Integration Test Tool")
+    parser.add_argument(
+        "--topol_path", required=True, help="Path to topology file"
+    )
+    parser.add_argument(
+        "--traj_path", required=True, help="Path to trajectory file"
+    )
+    parser.add_argument(
+        "--vel_path",
+        required=False,
+        help="Path to velocities trajectory file (NAMD only)",
+    )
+    parser.add_argument(
+        "--force_path",
+        required=False,
+        help="Path to forces trajectory file (NAMD only)",
+    )
+    parser.add_argument(
+        "--first_frame",
+        type=int,
+        required=True,
+        help="First frame to compare (0 for GROMACS and NAMD, 1 for LAMMPS)",
+    )
+    parser.add_argument(
+        "--tmp_path", default=".", help="Temporary directory for IMD traj"
+    )
+
+    args = parser.parse_args()
+
+    print(
+        "Writing IMD trajectory to temporary directory...\n===================="
+    )
+    imd_u = load_imd_universe(args.topol_path, args.tmp_path)
+
+    print("Loading source of truth trajectory...\n====================")
+    true_u = load_true_universe(args.topol_path, args.traj_path)
+
+    try:
+        print("Comparing trajectories...\n====================")
+        vel_in_trr = args.vel_path is None
+        force_in_trr = args.force_path is None
+        dt_in_trr = not args.topol_path.endswith(".data")
+
+        test_compare_imd_to_true_traj(
+            imd_u,
+            true_u,
+            args.first_frame,
+            vel_in_trr,
+            force_in_trr,
+            dt_in_trr,
+        )
+
+        if args.vel_path is not None:
+            print(
+                "Loading source of truth velocity trajectory...\n===================="
+            )
+            true_vel = load_true_universe(args.topol_path, args.vel_path)
+            print("Comparing velocities...")
+            test_compare_imd_to_true_traj_vel(
+                imd_u, true_vel, args.first_frame
+            )
+
+        if args.force_path is not None:
+            print(
+                "Loading source of truth force trajectory...\n===================="
             )
-            if true_u.trajectory[i].dimensions is not None:
-                assert_allclose_with_logging(
-                    true_u.trajectory[i].dimensions,
-                    imd_u.trajectory[i].dimensions,
-                    atol=1e-03,
-                )
-            if true_u.trajectory[i].has_positions:
-                assert_allclose_with_logging(
-                    true_u.trajectory[i].positions,
-                    imd_u.trajectory[i].positions,
-                    atol=1e-03,
-                )
-            if true_u.trajectory[i].has_velocities:
-                assert_allclose_with_logging(
-                    true_u.trajectory[i].velocities,
-                    imd_u.trajectory[i].velocities,
-                    atol=1e-03,
-                )
-            if true_u.trajectory[i].has_forces:
-                assert_allclose_with_logging(
-                    true_u.trajectory[i].forces,
-                    imd_u.trajectory[i].forces,
-                    atol=1e-03,
-                )
+            true_force = load_true_universe(args.topol_path, args.force_path)
+            print("Comparing forces...")
+            test_compare_imd_to_true_traj_forces(
+                imd_u, true_force, args.first_frame
+            )
+
+        print("All tests passed!")
+
+    except AssertionError as e:
+        logger.error("Comparison failed!")
+        print(f"Test failed: {e}")
+        raise e
+
+
+if __name__ == "__main__":
+    main()

imdclient/tests/test_namd.py

@@ -1,11 +1,21 @@
 import MDAnalysis as mda
 import pytest
 import logging
-from .base import IMDv3IntegrationTest
-from .datafiles import NAMD_TOPOL, NAMD_CONF, NAMD_TRAJ, NAMD_PARAMS, NAMD_PSF
+from .base import IMDv3IntegrationTest, assert_allclose_with_logging
+from .datafiles import (
+    NAMD_TOPOL,
+    NAMD_CONF_NST_1,
+    NAMD_CONF_NST_8,
+    NAMD_PARAMS,
+    NAMD_PSF,
+)
+from pathlib import Path
+from numpy.testing import (
+    assert_allclose,
+)
 
 logger = logging.getLogger("imdclient.IMDClient")
-file_handler = logging.FileHandler("test.log")
+file_handler = logging.FileHandler("namd_test.log")
 formatter = logging.Formatter(
     "%(asctime)s - %(name)s - %(levelname)s - %(message)s"
 )
@@ -16,23 +26,100 @@ logger.setLevel(logging.DEBUG)
 
 class TestIMDv3NAMD(IMDv3IntegrationTest):
 
+    @pytest.fixture(params=[NAMD_CONF_NST_1, NAMD_CONF_NST_8])
+    def inp(self, request):
+        return request.param
+
+    @pytest.fixture()
+    def simulation_command(self, inp):
+        return f"namd3 {Path(inp).name}"
+
+    @pytest.fixture()
+    def input_files(self, inp):
+        return [NAMD_TOPOL, inp, NAMD_PARAMS, NAMD_PSF]
+
+    @pytest.fixture()
+    def topol(self):
+        return Path(NAMD_TOPOL).name
+
     @pytest.fixture()
-    def command(self):
-        return (
-            f"cp {NAMD_PARAMS} {NAMD_PSF} {NAMD_TOPOL} . && namd3 {NAMD_CONF}"
+    def true_u(self, topol, imd_u, tmp_path):
+        u = mda.Universe(
+            (tmp_path / topol),
+            (tmp_path / "alanin.dcd"),
         )
+        yield u
 
     @pytest.fixture()
-    def match_string(self):
-        return "INTERACTIVE MD AWAITING CONNECTION"
+    def true_u_vel(self, topol, imd_u, tmp_path):
+        u = mda.Universe(
+            (tmp_path / topol),
+            (tmp_path / "alanin.vel.dcd"),
+        )
+        yield u
+
+    @pytest.fixture()
+    def true_u_force(self, topol, imd_u, tmp_path):
+        u = mda.Universe(
+            (tmp_path / topol),
+            (tmp_path / "alanin.force.dcd"),
+        )
+        yield u
+
+    # @pytest.fixture()
+    # def match_string(self):
+    #     return "INTERACTIVE MD AWAITING CONNECTION"
 
     @pytest.fixture()
     def first_frame(self):
         return 0
 
-    @pytest.fixture()
-    def universe(self):
-        return mda.Universe(
-            NAMD_TOPOL,
-            NAMD_TRAJ,
-        )
+    # Compare coords, box, time, dt, step
+    def test_compare_imd_to_true_traj(self, imd_u, true_u, first_frame):
+        for i in range(first_frame, len(true_u.trajectory)):
+            assert_allclose(
+                true_u.trajectory[i].time,
+                imd_u.trajectory[i - first_frame].time,
+                atol=1e-03,
+            )
+            assert_allclose(
+                true_u.trajectory[i].dt,
+                imd_u.trajectory[i - first_frame].dt,
+                atol=1e-03,
+            )
+            assert_allclose(
+                true_u.trajectory[i].data["step"],
+                imd_u.trajectory[i - first_frame].data["step"],
+            )
+            assert_allclose_with_logging(
+                true_u.trajectory[i].dimensions,
+                imd_u.trajectory[i - first_frame].dimensions,
+                atol=1e-03,
+            )
+            assert_allclose_with_logging(
+                true_u.trajectory[i].positions,
+                imd_u.trajectory[i - first_frame].positions,
+                atol=1e-03,
+            )
+
+    # Compare velocities
+    def test_compare_imd_to_true_traj_vel(
+        self, imd_u, true_u_vel, first_frame
+    ):
+        for i in range(first_frame, len(true_u_vel.trajectory)):
+            assert_allclose_with_logging(
+                true_u_vel.trajectory[i].positions,
+                imd_u.trajectory[i - first_frame].velocities,
+                atol=1e-03,
+            )
+
+    # Compare forces
+    def test_compare_imd_to_true_traj_forces(
+        self, imd_u, true_u_force, first_frame
+    ):
+        for i in range(first_frame, len(true_u_force.trajectory)):
+            assert_allclose_with_logging(
+                true_u_force.trajectory[i].positions,
+                imd_u.trajectory[i - first_frame].forces,
+                atol=1e-03,
+            )

imdclient/tests/test_stream_analysis.py

@@ -15,7 +15,7 @@ from numpy.testing import (
 )
 import numpy as np
 import pytest
-from imdclient.IMDREADER import IMDReader
+from imdclient.IMD import IMDReader
 
 
 class TestStreamAnalysis:

imdclient/tests/utils.py

@@ -1,6 +1,5 @@
 from imdclient.IMDProtocol import *
 import socket
-from imdclient.IMDProtocol import *
 import logging
 
 

imdclient-0.1.4.dist-info/LICENSE

@@ -0,0 +1,5 @@
+Copyright 2024 Lawson Woods
+
+Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the “Software”), to deal in the Software without restriction, including without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the Software is furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software.