imdclient 0.1.3__py3-none-any.whl → 0.2.0b0__py3-none-any.whl

imdclient/tests/test_manual.py

@@ -1,14 +1,10 @@
-from imdclient.IMD import IMDReader
-import pytest
+from .minimalreader import MinimalReader
 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,74 +16,236 @@ 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,
         )
-        yield imd_u
+    return mda.Universe(topol_path, traj_path)
 
-    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 load_imd_universe(topol_path, tmp_path):
+    # Pass atom_style (ignored if not using LAMMPS topol)
+    n_atoms = mda.Universe(
+        topol_path,
+        atom_style="id type x y z",
+        convert_units=False,
+    ).atoms.n_atoms
+    tmp_u = MinimalReader(
+        f"imd://localhost:8888", n_atoms=n_atoms, process_stream=True
+    )
+    return tmp_u
+
+
+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,
+        )
+
+
+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, step
+):
+    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,
+        )
+        # Issue #63
+        # if dt:
+        #     assert_allclose(
+        #         true_u.trajectory[i].dt,
+        #         imd_u.trajectory[i - first_frame].dt,
+        #         atol=1e-03,
+        #     )
+        if step:
             assert_allclose(
-                true_u.trajectory[i].time,
-                imd_u.trajectory[i - first_frame_arg].time,
+                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,
             )
-            assert_allclose(
-                true_u.trajectory[i].data["step"],
-                imd_u.trajectory[i - first_frame_arg].data["step"],
+        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")
+        # True when not using DCDReader
+        step_in_trr = not args.traj_path.endswith(".coor")
+
+        test_compare_imd_to_true_traj(
+            imd_u,
+            true_u,
+            args.first_frame,
+            vel_in_trr,
+            force_in_trr,
+            dt_in_trr,
+            step_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 - first_frame_arg].dimensions,
-                    atol=1e-03,
-                )
-            if true_u.trajectory[i].has_positions:
-                assert_allclose_with_logging(
-                    true_u.trajectory[i].positions,
-                    imd_u.trajectory[i - first_frame_arg].positions,
-                    atol=1e-03,
-                )
-            if true_u.trajectory[i].has_velocities:
-                assert_allclose_with_logging(
-                    true_u.trajectory[i].velocities,
-                    imd_u.trajectory[i - first_frame_arg].velocities,
-                    atol=1e-03,
-                )
-            if true_u.trajectory[i].has_forces:
-                assert_allclose_with_logging(
-                    true_u.trajectory[i].forces,
-                    imd_u.trajectory[i - first_frame_arg].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,6 +1,13 @@
-import MDAnalysis as mda
-import pytest
 import logging
+from pathlib import Path
+import re
+
+import pytest
+from numpy.testing import (
+    assert_allclose,
+)
+import MDAnalysis as mda
+
 from .base import IMDv3IntegrationTest, assert_allclose_with_logging
 from .datafiles import (
     NAMD_TOPOL,
@@ -9,10 +16,6 @@ from .datafiles import (
     NAMD_PARAMS,
     NAMD_PSF,
 )
-from pathlib import Path
-from numpy.testing import (
-    assert_allclose,
-)
 
 logger = logging.getLogger("imdclient.IMDClient")
 file_handler = logging.FileHandler("namd_test.log")
@@ -26,6 +29,10 @@ logger.setLevel(logging.DEBUG)
 
 class TestIMDv3NAMD(IMDv3IntegrationTest):
 
+    @pytest.fixture()
+    def container_name(self):
+        return "ghcr.io/becksteinlab/streaming-namd-docker:main-common-cpu"
+
     @pytest.fixture(params=[NAMD_CONF_NST_1, NAMD_CONF_NST_8])
     def inp(self, request):
         return request.param
@@ -42,6 +49,18 @@ class TestIMDv3NAMD(IMDv3IntegrationTest):
     def topol(self):
         return Path(NAMD_TOPOL).name
 
+    @pytest.fixture()
+    def dt(self, inp):
+        pattern = re.compile(r"^\s*timestep\s*(\S+)")
+        with open(inp, "r") as file:
+            for line in file:
+                match = pattern.match(line)
+                if match:
+                    # NAMD timestep is in femtoseconds, convert to picoseconds
+                    # as IMDv3 expects dt in ps, 1fs = 0.001ps
+                    return float(match.group(1)) * 0.001
+        raise ValueError(f"No dt found in {inp}")
+
     @pytest.fixture()
     def true_u(self, topol, imd_u, tmp_path):
         u = mda.Universe(
@@ -75,40 +94,43 @@ class TestIMDv3NAMD(IMDv3IntegrationTest):
         return 0
 
     # Compare coords, box, time, dt, step
-    def test_compare_imd_to_true_traj(self, imd_u, true_u, first_frame):
+    def test_compare_imd_to_true_traj(self, imd_u, true_u, first_frame, 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,
             )
+
             assert_allclose(
-                true_u.trajectory[i].dt,
+                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"],
-            )
+            # step in DCDReader is frame index, not integration step
+            # don't compare 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,
             )
 
+    # Since NAMD does not write velocities, forces to the DCD file, we need to do so seperately by extracting that info from their respective DCD files
     # Compare velocities
-    def test_compare_imd_to_true_traj_vel(
-        self, imd_u, true_u_vel, first_frame
-    ):
+    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,
+                # Unit conversion
+                true_u_vel.trajectory[i].positions * 20.45482706,
                 imd_u.trajectory[i - first_frame].velocities,
                 atol=1e-03,
             )
@@ -118,6 +140,7 @@ class TestIMDv3NAMD(IMDv3IntegrationTest):
         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,

imdclient/tests/test_utils.py

@@ -0,0 +1,31 @@
+# tests/test_utils.py
+import pytest
+from imdclient.utils import parse_host_port
+
+
+@pytest.mark.parametrize(
+    "server_address,host_port",
+    [
+        ("imd://localhost:8888", ("localhost", 8888)),
+        ("imd://example.com:12345", ("example.com", 12345)),
+    ],
+)
+def test_parse_host_port_valid(server_address, host_port):
+    assert parse_host_port(server_address) == host_port
+
+
+@pytest.mark.parametrize(
+    "server_address",
+    [
+        "",  # empty
+        "http://localhost:80",  # wrong protocol prefix
+        "imd://",  # missing host and port
+        "imd://localhost:",  # missing port
+        "imd://:8080",  # missing host
+        "imd://host:notaport",  # port not integer
+        "imd://host:80:90",  # too many segments
+    ],
+)
+def test_parse_host_port_invalid(server_address):
+    with pytest.raises(ValueError):
+        parse_host_port(server_address)

imdclient/utils.py

@@ -43,23 +43,6 @@ 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("imd://")[1].split(":")
-    if len(parts) == 2:
-        host = parts[0] 
-        try:
-            port = int(parts[1])
-            return (host, port)
-        except ValueError:
-            raise ValueError("IMDReader: Port must be an integer")
-    else:
-        raise ValueError("IMDReader: URL must be in the format 'imd://host:port'")
-
 
 def approximate_timestep_memsize(
     n_atoms, energies, dimensions, positions, velocities, forces
@@ -116,3 +99,53 @@ def sock_contains_data(sock, timeout) -> bool:
         [sock], [], [], timeout
     )
     return sock in ready_to_read
+
+
+def parse_host_port(filename):
+    """
+    Parses a URL in the format 'imd://host:port' and returns the host and port.
+    Parameters
+    ----------
+    filename : str
+        The URL to parse, must be in the format 'imd://host:port'.
+
+    Returns
+    -------
+    tuple[str, int]
+        A 2-tuple ``(host, port)`` where `host` is the host server name
+        and `port` is the TCP port number.
+    Raises
+    ------
+    ValueError
+        If the URL is not in the correct format or if the host or port is invalid.
+
+    Examples
+    --------
+    >>> parse_host_port("imd://localhost:8888")
+    ('localhost', 8888)
+    >>> parse_host_port("invalid://localhost:12345")
+    Traceback (most recent call last):
+        ... ValueError: IMDClient: URL must be in the format 'imd://host:port'
+    """
+    if not filename.startswith("imd://"):
+        raise ValueError(
+            "IMDClient: URL must be in the format 'imd://host:port'"
+        )
+
+    # Check if the format is correct
+    parts = filename.split("imd://")[1].split(":")
+    if len(parts) == 2:
+        host = parts[0]
+        if not host:
+            raise ValueError("IMDClient: Host cannot be empty")
+        if not parts[1]:
+            raise ValueError("IMDClient: Port cannot be empty")
+        try:
+            port = int(parts[1])
+            return (host, port)
+        except ValueError:
+            raise ValueError("IMDClient: Port must be an integer")
+    else:
+        raise ValueError(
+            "IMDClient: URL must be in the format 'imd://host:port'"
+        )