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>=0.2.2'
   # pip-installed min dependencies
   coverage:
     default: 'coverage'
@@ -147,6 +149,7 @@ runs:
           ${{ inputs.scikit-learn }}
           ${{ inputs.seaborn }}
           ${{ inputs.tidynamics }}
+          ${{ inputs.imdclient }}
 
       run: |
         # setup full variable

azure-pipelines.yml

@@ -113,6 +113,8 @@ jobs:
       pytng>=0.2.3
       rdkit>=2024.03.4
       tidynamics>=1.0.0
+      imdclient>=0.2.2
+
     # remove from azure to avoid test hanging #4707
     # "gsd>3.0.0"
     displayName: 'Install additional dependencies for 64-bit tests'

maintainer/conda/environment.yml

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

package/AUTHORS

@@ -260,6 +260,7 @@ Chronological list of authors
   - Gareth Elliott
   - Marc Schuh
   - Sirsha Ganguly
+  - Amruthesh Thirumalaiswamy
 
 External code
 -------------

package/CHANGELOG

@@ -16,7 +16,7 @@ The rules for this file:
 -------------------------------------------------------------------------------
 ??/??/?? IAlibay, orbeckst, BHM-Bob, TRY-ER, Abdulrahman-PROG, pbuslaev,
          yuxuanzhuang, yuyuan871111, tanishy7777, tulga-rdn, Gareth-elliott,
-         hmacdope, tylerjereddy, cbouy, talagayev, DrDomenicoMarson
+         hmacdope, tylerjereddy, cbouy, talagayev, DrDomenicoMarson, amruthesht
 
 
  * 2.10.0
@@ -71,6 +71,11 @@ Enhancements
    (PR #5038)
  * Moved distopia checking function to common import location in
    MDAnalysisTest.util (PR #5038)
+ * Added support for reading and processing streamed data in `coordinates.base`
+   with new `StreamFrameIteratorSliced` and `StreamReaderBase` (Issue #4827, PR #4923)
+ * New coordinate reader: Added `IMDReader` for reading real-time streamed
+   molecular dynamics simulation data using the IMDv3 protocol - requires
+   `imdclient` package (Issue #4827, PR #4923)  
  * Enables parallelization for `analysis.polymer.PersistenceLength` (Issue #4671, PR #5074)
 
 

package/MDAnalysis/coordinates/IMD.py

@@ -0,0 +1,290 @@
+"""
+IMDReader --- :mod:`MDAnalysis.coordinates.IMD`
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+This module provides support for reading molecular dynamics simulation data via the 
+`Interactive Molecular Dynamics (IMD) protocol v3 <https://imdclient.readthedocs.io/en/latest/protocol_v3.html>`_.
+The IMD 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. 
+
+.. note::
+   This reader only supports IMDv3, which is implemented in GROMACS, LAMMPS, and NAMD at varying
+   stages of development. See the `imdclient simulation engine docs`_ for more. 
+   While IMDv2 is widely available in simulation engines, it was designed primarily for visualization
+   and gaps are allowed in the stream (i.e., an inconsistent number of integrator time steps between transmitted coordinate arrays is allowed)
+
+The :class:`IMDReader` connects to a simulation via a socket and receives coordinate,
+velocity, force, and energy data as the simulation progresses. This allows for real-time
+monitoring and analysis of ongoing simulations. It uses the `imdclient package <https://github.com/Becksteinlab/imdclient>`_ 
+(dependency) to implement the IMDv3 protocol and manage the socket connection and data parsing.
+
+Usage Example
+-------------
+
+As an example of reading a stream, after configuring GROMACS to run a simulation with IMDv3 enabled
+(see the `imdclient simulation engine docs`_ for 
+up-to-date resources on configuring each simulation engine), use the following commands:
+
+.. code-block:: bash
+
+    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]}')
+
+Important Limitations
+---------------------
+
+Since IMD streams data in real-time from a running simulation, it has fundamental 
+constraints that differ from traditional trajectory readers:
+
+* **No random access**: Cannot jump to arbitrary frame numbers or seek backwards  
+* **Forward-only access**: You can only move forward through frames as they arrive
+* **No trajectory length**: The total number of frames is unknown until the simulation ends
+* **Single-use iteration**: Cannot restart or rewind once the stream has been consumed
+* **No independent copies**: Cannot create separate reader instances for the same stream
+* **No stream restart**: Cannot reconnect or reopen once the connection is closed
+* **No bulk operations**: Cannot extract all data at once using timeseries methods
+* **Limited multiprocessing**: Cannot split reader across processes for parallel analysis
+* **Single client connection**: Only one reader can connect to an IMD stream at a time
+* **No trajectory Writing**: Complimentary IMD Writer class is not available for streaming data
+
+.. warning::  
+   The IMDReader has some important limitations that are inherent in streaming data.  
+   See :class:`~MDAnalysis.coordinates.base.StreamReaderBase` for technical details.  
+
+.. seealso::
+   :class:`IMDReader`
+      Technical details and parameter options for the reader class
+   
+   `imdclient documentation <https://imdclient.readthedocs.io/>`_
+      Complete documentation for the IMDClient package
+      
+   `IMDClient GitHub repository <https://github.com/Becksteinlab/imdclient>`_
+      Source code and development resources
+
+.. _`imdclient simulation engine docs`: https://imdclient.readthedocs.io/en/latest/usage.html
+
+
+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.2.2")
+
+try:
+    import imdclient
+    from imdclient.IMDClient import IMDClient
+    from imdclient.utils import parse_host_port
+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.2.2
+    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):
+    """
+    Coordinate reader implementing the IMDv3 protocol for streaming simulation data.
+
+    This class handles the technical aspects of connecting to IMD-enabled simulation
+    engines and processing the incoming data stream. For usage examples and protocol
+    overview, see the module documentation above.
+
+    The reader manages socket connections, data buffering, and frame parsing according
+    to the IMDv3 specification. It automatically handles different data packet types
+    (coordinates, velocities, forces, energies, timing) and populates MDAnalysis
+    timestep objects accordingly.
+
+    Parameters
+    ----------
+    filename : a string of the form "imd://host:port" where host is the hostname
+        or IP address of the listening simulation engine's IMD 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.
+    buffer_size: int (optional) default=10*(1024**2)
+        number of bytes of memory to allocate to the :class:`~imdclient.IMDClient.IMDClient`'s
+        internal buffer. Defaults to 10 megabytes. Larger buffers can improve
+        performance for analyses with periodic heavy computation.
+    **kwargs : dict (optional)
+        keyword arguments passed to the constructed :class:`~imdclient.IMDClient.IMDClient`
+
+    Notes
+    -----
+    The IMDReader provides access to additional simulation data through the timestep's
+    `data` attribute (`ts.data`). The following keys may be available depending on
+    what the simulation engine transmits:
+
+    * `dt` : float
+        Time step size in picoseconds (from the `IMD_TIME`_ packet of the IMDv3 protocol)
+    * `step` : int
+        Current simulation step number (from the `IMD_TIME`_ packet of the IMDv3 protocol)
+    * Energy terms : float
+        Various energy components (e.g., 'potential', 'kinetic', 'total', etc.)
+        from the `IMD_ENERGIES`_ packet of the IMDv3 protocol.
+
+    .. _IMD_TIME: https://imdclient.readthedocs.io/en/latest/protocol_v3.html#time
+    .. _IMD_ENERGIES: https://imdclient.readthedocs.io/en/latest/protocol_v3.html#energies
+
+    .. note::
+       For important limitations inherent to streaming data, see the module documentation above
+       and :class:`~MDAnalysis.coordinates.base.StreamReaderBase` for more technical details.
+
+    .. versionadded:: 2.10.0
+    """
+
+    format = "IMD"
+
+    @store_init_arguments
+    def __init__(
+        self,
+        filename,
+        n_atoms=None,
+        buffer_size=10 * (1024**2),
+        **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
+
+        try:
+            host, port = parse_host_port(filename)
+        except ValueError as e:
+            raise ValueError(f"IMDReader: Invalid IMD URL '{filename}': {e}")
+
+        # This starts the simulation
+        self._imdclient = IMDClient(
+            host, port, n_atoms, buffer_size=buffer_size, **kwargs
+        )
+
+        imdsinfo = self._imdclient.get_imdsessioninfo()
+        if imdsinfo.version != 3:
+            raise ValueError(
+                f"IMDReader: Detected IMD version v{imdsinfo.version}, "
+                + "but IMDReader is only compatible with v3"
+            )
+
+        self.ts = self._Timestep(
+            self.n_atoms,
+            positions=imdsinfo.positions,
+            velocities=imdsinfo.velocities,
+            forces=imdsinfo.forces,
+            **self._ts_kwargs,
+        )
+
+        try:
+            self._read_next_timestep()
+        except EOFError as e:
+            raise RuntimeError(f"IMDReader: Read error: {e}") from e
+
+    def _read_frame(self, frame):
+
+        imdf = self._imdclient.get_imdframe()
+
+        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):
+        if not isinstance(thing, str):
+            return False
+        # a weaker check for type hint
+        if thing.startswith("imd://"):
+            return True
+        else:
+            return False
+
+    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()
+        logger.debug("IMDReader shut down gracefully.")

package/MDAnalysis/coordinates/__init__.py

@@ -53,7 +53,7 @@ class that defines a common :ref:`Trajectory API` and allows other code to
 :class:`~MDAnalysis.coordinates.base.ProtoReader` object; all Readers are accessible
 through this entry point in the same manner ("`duck typing`_").
 
-There are three types of base Reader which act as starting points for each
+There are four types of base Reader which act as starting points for each
 specific format. These are:
 
 :class:`~MDAnalysis.coordinates.base.ReaderBase`
@@ -66,6 +66,12 @@ class that defines a common :ref:`Trajectory API` and allows other code to
    frame of information.  This is used with formats such as GRO
    and CRD
 
+:class:`~MDAnalysis.coordinates.base.StreamReaderBase`
+   A specialized Reader for continuous data streams such as live 
+   simulation feeds. Unlike standard readers, streaming readers cannot
+   randomly access frames, rewind, or determine total length. This is
+   used for real-time trajectory data from simulations via IMD connections.
+
 :class:`~MDAnalysis.coordinates.chain.ChainReader`
    An advanced Reader designed to read a sequence of files, to
    provide iteration over all the frames in each file seamlessly.
@@ -277,6 +283,11 @@ class can choose an appropriate reader automatically.
    | library       |           |       | file formats`_ and                                   |
    |               |           |       | :mod:`MDAnalysis.coordinates.chemfiles`              |
    +---------------+-----------+-------+------------------------------------------------------+
+   | IMD           | imd://    |  r    | Receive simulation trajectory data using interactive |
+   |               | <host>:   |       | molecular dynamics version 3 (IMDv3) by configuring  |
+   |               | <port>    |       | a socket address to a NAMD, GROMACS, or LAMMPS       |
+   |               |           |       | simulation. :mod:`MDAnalysis.coordinates.IMD`        |
+   +---------------+-----------+-------+------------------------------------------------------+
 
 .. [#a] This format can also be used to provide basic *topology*
    information (i.e. the list of atoms); it is possible to create a
@@ -773,6 +784,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