feature: IMDReader Integration

.github/actions/setup-deps/action.yaml

@@ -82,6 +82,8 @@ inputs:
     default: 'seaborn>=0.7.0'
   tidynamics:
     default: 'tidynamics>=1.0.0'
+  imdclient:
+    default: 'imdclient'
   # pip-installed min dependencies
   coverage:
     default: 'coverage'
@@ -137,6 +139,7 @@ runs:
           ${{ inputs.distopia }}
           ${{ inputs.gsd }}
           ${{ inputs.h5py }}
+          ${{ inputs.imdclient }}
           ${{ inputs.hole2 }}
           ${{ inputs.joblib }}
           ${{ inputs.netcdf4 }}

azure-pipelines.yml

@@ -77,6 +77,7 @@ jobs:
     displayName: 'Install tools'
   - script: >-
       python -m pip install --only-binary=scipy,h5py
+      imdclient
       cython
       hypothesis
       h5py>=2.10

maintainer/conda/environment.yml

@@ -30,6 +30,7 @@ dependencies:
   - sphinxcontrib-bibtex
   - mdaencore
   - waterdynamics
+  - imdclient
   - pip:
       - mdahole2
       - pathsimanalysis

package/MDAnalysis/coordinates/IMD.py

@@ -0,0 +1,248 @@
+"""
+IMDReader --- :mod:`MDAnalysis.coordinates.IMD`
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+:class:`MDAnalysis.coordinates.IMD.IMDReader` is a class that implements the Interactive Molecular Dynamics (IMD) protocol for reading simulation
+data using the IMDClient (see `imdclient <https://github.com/Becksteinlab/imdclient>`_).
+The protocol allows two-way communicating molecular simulation data through a socket.
+Via IMD, a simulation engine sends data to a receiver (in this case, the IMDClient) and the receiver can send forces and specific control
+requests (such as pausing, resuming, or terminating the simulation) back to the simulation engine. It currently supports
+simulation running with GROMACS, LAMMPS, or NAMD.
+
+For example, when running a simulation with GROMACS that supports streaming, use the following commands:
+
+.. code-block:: bash
+
+    gmx grompp -f run-NPT_imd-v3.mdp -c conf.gro -p topol.top -o topol.tpr
+    gmx mdrun -v -nt 4 -imdwait -imdport 8889
+
+The :class:`MDAnalysis.coordinates.IMD.IMDReader` can then connect to the running simulation and stream data in real time:
+
+.. code-block:: python
+
+    import MDAnalysis as mda
+    u = mda.Universe("topol.tpr", "imd://localhost:8889", buffer_size=10*1024*1024)
+
+    print("    time [         position         ] [         velocity         ] [           force          ] [            box           ]")
+    sel = u.select_atoms("all")  # Select all atoms; adjust selection as needed
+    for ts in u.trajectory:
+        print(f'{ts.time:8.3f} {sel[0].position} {sel[0].velocity} {sel[0].force} {u.dimensions[0:3]}')
+
+Details about the IMD protocol and usage examples can be found in the
+`imdclient <https://github.com/Becksteinlab/imdclient>`_ repository.
+
+
+Units
+-----
+The units in IMDv3 are fixed.
+
+.. list-table::
+   :widths: 10 10
+   :header-rows: 1
+
+   * - Measurement
+     - Unit
+   * - Length
+     - angstrom
+   * - Velocity
+     - angstrom/picosecond
+   * - Force
+     - kilojoules/(mol*angstrom)
+   * - Time
+     - picosecond
+   * - Energy
+     - kilojoules/mol
+
+Classes
+-------
+
+.. autoclass:: IMDReader
+   :members:
+   :inherited-members:
+
+"""
+
+import numpy as np
+import logging
+import warnings
+
+from MDAnalysis.coordinates import core
+from MDAnalysis.lib.util import store_init_arguments
+from MDAnalysis.coordinates.base import StreamReaderBase
+
+
+from packaging.version import Version
+
+MIN_IMDCLIENT_VERSION = Version("0.1.4")
+
+try:
+    import imdclient
+    from imdclient.IMDClient import IMDClient
+except ImportError:
+    HAS_IMDCLIENT = False
+    imdclient_version = Version("0.0.0")
+
+    # Allow building documentation without imdclient
+    import types
+
+    class MockIMDClient:
+        pass
+
+    imdclient = types.ModuleType("imdclient")
+    imdclient.IMDClient = MockIMDClient
+    imdclient.__version__ = "0.0.0"
+
+else:
+    HAS_IMDCLIENT = True
+    imdclient_version = Version(imdclient.__version__)
+
+    # Check for compatibility: currently needs to be >=0.1.4
+    if imdclient_version < MIN_IMDCLIENT_VERSION:
+        warnings.warn(
+            f"imdclient version {imdclient_version} is too old; "
+            f"need at least {imdclient_version}, Your installed version of "
+            "imdclient will NOT be used.",
+            category=RuntimeWarning,
+        )
+        HAS_IMDCLIENT = False
+
+logger = logging.getLogger("MDAnalysis.coordinates.IMDReader")
+
+
+class IMDReader(StreamReaderBase):
+    """
+    Reader that supports the Interactive Molecular Dynamics (IMD) protocol for reading simulation
+    data using the IMDClient.
+
+    Parameters
+    ----------
+    filename : a string of the form "imd://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"
+    one_pass = True
+
+    @store_init_arguments
+    def __init__(
+        self,
+        filename,
+        convert_units=True,
+        n_atoms=None,
+        **kwargs,
+    ):
+        if not HAS_IMDCLIENT:
+            raise ImportError(
+                "IMDReader requires the imdclient package. "
+                "Please install it with 'pip install imdclient'."
+            )
+
+        super(IMDReader, self).__init__(filename, **kwargs)
+
+        self._imdclient = None
+        logger.debug("IMDReader initializing")
+
+        if n_atoms is None:
+            raise ValueError("IMDReader: n_atoms must be specified")
+        self.n_atoms = n_atoms
+
+        host, port = parse_host_port(filename)
+
+        # This starts the simulation
+        self._imdclient = IMDClient(host, port, n_atoms, **kwargs)
+
+        imdsinfo = self._imdclient.get_imdsessioninfo()
+        # NOTE: after testing phase, fail out on IMDv2
+
+        self.ts = self._Timestep(
+            self.n_atoms,
+            positions=imdsinfo.positions,
+            velocities=imdsinfo.velocities,
+            forces=imdsinfo.forces,
+            **self._ts_kwargs,
+        )
+
+        self._frame = -1
+
+        try:
+            self._read_next_timestep()
+        except StopIteration as e:
+            raise RuntimeError("IMDReader: No data found in stream") from e
+
+    def _read_frame(self, frame):
+
+        try:
+            imdf = self._imdclient.get_imdframe()
+        except EOFError as e:
+            raise e
+
+        self._frame = frame
+        self._load_imdframe_into_ts(imdf)
+
+        logger.debug("IMDReader: Loaded frame %d", self._frame)
+        return self.ts
+
+    def _load_imdframe_into_ts(self, imdf):
+        self.ts.frame = self._frame
+        if imdf.time is not None:
+            self.ts.time = imdf.time
+            # NOTE: timestep.pyx "dt" method is suspicious bc it uses "new" keyword for a float
+            self.ts.data["dt"] = imdf.dt
+            self.ts.data["step"] = imdf.step
+        if imdf.energies is not None:
+            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:
+            # must call copy because reference is expected to reset
+            # see 'test_frame_collect_all_same' in MDAnalysisTests.coordinates.base
+            np.copyto(self.ts.positions, imdf.positions)
+        if imdf.velocities is not None:
+            np.copyto(self.ts.velocities, imdf.velocities)
+        if imdf.forces is not None:
+            np.copyto(self.ts.forces, imdf.forces)
+
+    @staticmethod
+    def _format_hint(thing):
+        try:
+            parse_host_port(thing)
+        except ValueError:
+            return False
+        return HAS_IMDCLIENT and True
+
+    def close(self):
+        """Gracefully shut down the reader. Stops the producer thread."""
+        logger.debug("IMDReader close() called")
+        if self._imdclient is not None:
+            self._imdclient.stop()
+        # NOTE: removeme after testing
+        logger.debug("IMDReader shut down gracefully.")
+
+
+# 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 as e:
+            raise ValueError("IMDReader: Port must be an integer") from e
+    else:
+        raise ValueError(
+            "IMDReader: URL must be in the format 'imd://host:port'"
+        )

package/MDAnalysis/coordinates/__init__.py

@@ -274,6 +274,11 @@ class can choose an appropriate reader automatically.
    | library       |           |       | file formats`_ and                                   |
    |               |           |       | :mod:`MDAnalysis.coordinates.chemfiles`              |
    +---------------+-----------+-------+------------------------------------------------------+
+   | IMD           | IP address|  r/w  | Receive simulation trajectory data using interactive |
+   |               | and port  |       | molecular dynamics version 3 (IMDv3) by configuring  |
+   |               | number    |       | a socket address to a NAMD, GROMACS, or LAMMPS       |
+   |               |           |       | simulation.                                          |
+   +---------------+-----------+-------+------------------------------------------------------+
 
 .. [#a] This format can also be used to provide basic *topology*
    information (i.e. the list of atoms); it is possible to create a
@@ -770,6 +775,7 @@ class can choose an appropriate reader automatically.
 from . import DMS
 from . import GMS
 from . import GRO
+from . import IMD
 from . import INPCRD
 from . import LAMMPS
 from . import MOL2