"""
MDAnalysis IMDReader
^^^^^^^^^^^^^^^^^^^^
.. autoclass:: IMDReader
:members:
:inherited-members:
"""
from MDAnalysis.coordinates import core
from MDAnalysis.lib.util import store_init_arguments
# NOTE: changeme
from .IMDClient import IMDClient
from .utils import *
import logging
from .streambase import StreamReaderBase
logger = logging.getLogger("imdclient.IMDClient")
[docs]
class IMDReader(StreamReaderBase):
"""
Reader for IMD protocol packets.
Parameters
----------
filename : a string of the form "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,
):
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:
raise RuntimeError("IMDReader: No data found in stream")
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(f"IMDReader: Loaded frame {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
self.ts.positions = imdf.positions
if imdf.velocities is not None:
self.ts.velocities = imdf.velocities
if imdf.forces is not None:
self.ts.forces = imdf.forces
@staticmethod
def _format_hint(thing):
try:
parse_host_port(thing)
except:
return False
return True
[docs]
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.")