imdclient 0.1.0__py3-none-any.whl

imdclient/IMDProtocol.py

@@ -0,0 +1,164 @@
+import struct
+import logging
+from enum import Enum, auto
+from typing import Union
+from dataclasses import dataclass
+import numpy as np
+
+
+IMDHEADERSIZE = 8
+IMDENERGYPACKETLENGTH = 40
+IMDBOXPACKETLENGTH = 36
+IMDTIMEPACKETLENGTH = 24
+IMDVERSIONS = {2, 3}
+IMDAWAITGOTIME = 1
+
+logger = logging.getLogger("imdclient.IMDClient")
+
+
+class IMDHeaderType(Enum):
+    IMD_DISCONNECT = 0
+    IMD_ENERGIES = 1
+    IMD_FCOORDS = 2
+    IMD_GO = 3
+    IMD_HANDSHAKE = 4
+    IMD_KILL = 5
+    IMD_MDCOMM = 6
+    IMD_PAUSE = 7
+    IMD_TRATE = 8
+    IMD_IOERROR = 9
+    # New in IMD v3
+    IMD_SESSIONINFO = 10
+    IMD_RESUME = 11
+    IMD_TIME = 12
+    IMD_BOX = 13
+    IMD_VELOCITIES = 14
+    IMD_FORCES = 15
+
+
+def parse_energy_bytes(data, endianness):
+    keys = [
+        "step",
+        "temperature",
+        "total_energy",
+        "potential_energy",
+        "van_der_walls_energy",
+        "coulomb_energy",
+        "bonds_energy",
+        "angles_energy",
+        "dihedrals_energy",
+        "improper_dihedrals_energy",
+    ]
+    values = struct.unpack(f"{endianness}ifffffffff", data)
+    return dict(zip(keys, values))
+
+
+def create_energy_bytes(
+    step,
+    temperature,
+    total_energy,
+    potential_energy,
+    van_der_walls_energy,
+    coulomb_energy,
+    bonds_energy,
+    angles_energy,
+    dihedrals_energy,
+    improper_dihedrals_energy,
+    endianness,
+):
+    return struct.pack(
+        f"{endianness}ifffffffff",
+        step,
+        temperature,
+        total_energy,
+        potential_energy,
+        van_der_walls_energy,
+        coulomb_energy,
+        bonds_energy,
+        angles_energy,
+        dihedrals_energy,
+        improper_dihedrals_energy,
+    )
+
+
+def parse_box_bytes(data, endianness):
+    """Box is a 3x3 matrix of floats"""
+
+    vals = struct.unpack(f"{endianness}fffffffff", data)
+    return np.array(vals).reshape(3, 3)
+
+
+class IMDHeader:
+    """Convenience class to represent the header of an IMD packet"""
+
+    def __init__(self, data):
+        msg_type, length = struct.unpack("!ii", data)
+        h_type = IMDHeaderType(msg_type)
+
+        self.type = h_type
+        self.length = length
+
+
+class IMDTime:
+    """Convenience class to represent the body of time packet"""
+
+    def __init__(self, data, endianness):
+        self.dt, self.time, self.step = struct.unpack(f"{endianness}ddq", data)
+
+
+@dataclass
+class IMDSessionInfo:
+    """Convenience class to represent the session information of an IMD connection
+
+    '<' represents little endian and '>' represents big endian
+    """
+
+    version: int
+    endianness: str
+    # In IMDv2, we don't know if the server sends wrapped coordinates until
+    # we receive the first packet.
+    wrapped_coords: bool
+    # In IMDv2, we don't know if the server sends energies until
+    # we receive the first packet.
+    energies: bool
+    time: bool
+    box: bool
+    positions: bool
+    velocities: bool
+    forces: bool
+
+
+def parse_imdv3_session_info(data, end):
+    """Parses the session information packet of an IMD v3 connection"""
+    logger.debug(f"parse_imdv3_session_info: {data}")
+    time, energies, box, positions, wrapped_coords, velocties, forces = (
+        struct.unpack(f"{end}BBBBBBB", data)
+    )
+    logger.debug(f"parse_imdv3_session_info2 : {data}")
+    imdsinfo = IMDSessionInfo(
+        version=3,
+        endianness=end,
+        time=(time != 0),
+        box=(box != 0),
+        positions=(positions != 0),
+        wrapped_coords=(wrapped_coords != 0),
+        velocities=(velocties != 0),
+        forces=(forces != 0),
+        energies=(energies != 0),
+    )
+    logger.debug(f"parse_imdv3_session_info3: {data}")
+    return imdsinfo
+
+
+def create_header_bytes(msg_type: IMDHeaderType, length: int):
+    # NOTE: add error checking for invalid packet msg_type here
+
+    type = msg_type.value
+    return struct.pack("!ii", type, length)
+
+
+def parse_header_bytes(data):
+    msg_type, length = struct.unpack("!ii", data)
+    type = IMDHeaderType(msg_type)
+    # NOTE: add error checking for invalid packet msg_type here
+    return IMDHeader(type, length)

imdclient/IMDREADER.py

@@ -0,0 +1,129 @@
+"""
+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")
+
+
+class IMDReader(StreamReaderBase):
+    """
+    Reader for IMD protocol packets.
+    """
+
+    format = "IMD"
+    one_pass = True
+
+    @store_init_arguments
+    def __init__(
+        self,
+        filename,
+        convert_units=True,
+        n_atoms=None,
+        **kwargs,
+    ):
+        """
+        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.
+        """
+
+        super(IMDReader, self).__init__(filename, **kwargs)
+
+        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:
+            # Not strictly necessary, but for clarity
+            raise StopIteration
+
+        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(imdf.energies)
+        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
+
+    def close(self):
+        """Gracefully shut down the reader. Stops the producer thread."""
+        logger.debug("IMDReader close() called")
+        self._imdclient.stop()
+        # NOTE: removeme after testing
+        logger.debug("IMDReader shut down gracefully.")

imdclient/__init__.py

@@ -0,0 +1,14 @@
+"""
+IMDClient
+"""
+
+# Don't import IMDReader here, eventually it may be moved to a separate package
+from .IMDClient import IMDClient
+from importlib.metadata import version
+
+from .streamanalysis import AnalysisBase, StackableAnalysis
+from MDAnalysis.analysis import base
+
+base.AnalysisBase = AnalysisBase
+
+__version__ = version("imdclient")

imdclient/backends.py

@@ -0,0 +1,352 @@
+# Copy of backends from MDA 2.8.0
+"""Analysis backends --- :mod:`MDAnalysis.analysis.backends`
+============================================================
+
+.. versionadded:: 2.8.0
+
+
+The :mod:`backends` module provides :class:`BackendBase` base class to
+implement custom execution backends for
+:meth:`MDAnalysis.analysis.base.AnalysisBase.run` and its
+subclasses.
+
+.. SeeAlso:: :ref:`parallel-analysis`
+
+.. _backends:
+
+Backends
+--------
+
+Three built-in backend classes are provided:
+
+* *serial*: :class:`BackendSerial`, that is equivalent to using no
+  parallelization and is the default
+
+* *multiprocessing*: :class:`BackendMultiprocessing` that supports
+  parallelization via standard Python :mod:`multiprocessing` module
+  and uses default :mod:`pickle` serialization
+
+* *dask*: :class:`BackendDask`, that uses the same process-based
+  parallelization as :class:`BackendMultiprocessing`, but different
+  serialization algorithm via `dask <https://dask.org/>`_ (see `dask
+  serialization algorithms
+  <https://distributed.dask.org/en/latest/serialization.html>`_ for details)
+
+Classes
+-------
+
+"""
+import warnings
+from typing import Callable
+import importlib.util
+
+
+def is_installed(modulename: str):
+    """Checks if module is installed
+
+    Parameters
+    ----------
+    modulename : str
+        name of the module to be tested
+
+
+    .. versionadded:: 2.8.0
+    """
+    return importlib.util.find_spec(modulename) is not None
+
+
+class BackendBase:
+    """Base class for backend implementation.
+
+    Initializes an instance and performs checks for its validity, such as
+    ``n_workers`` and possibly other ones.
+
+    Parameters
+    ----------
+    n_workers : int
+        number of workers (usually, processes) over which the work is split
+
+    Examples
+    --------
+    .. code-block:: python
+
+        from MDAnalysis.analysis.backends import BackendBase
+
+        class ThreadsBackend(BackendBase):
+            def apply(self, func, computations):
+                from multiprocessing.dummy import Pool
+
+                with Pool(processes=self.n_workers) as pool:
+                    results = pool.map(func, computations)
+                return results
+
+        import MDAnalysis as mda
+        from MDAnalysis.tests.datafiles import PSF, DCD
+        from MDAnalysis.analysis.rms import RMSD
+
+        u = mda.Universe(PSF, DCD)
+        ref = mda.Universe(PSF, DCD)
+
+        R = RMSD(u, ref)
+
+        n_workers = 2
+        backend = ThreadsBackend(n_workers=n_workers)
+        R.run(backend=backend, unsupported_backend=True)
+
+    .. warning::
+        Using `ThreadsBackend` above will lead to erroneous results, since it
+        is an educational example. Do not use it for real analysis.
+
+
+    .. versionadded:: 2.8.0
+
+    """
+
+    def __init__(self, n_workers: int):
+        self.n_workers = n_workers
+        self._validate()
+
+    def _get_checks(self):
+        """Get dictionary with ``condition: error_message`` pairs that ensure the
+        validity of the backend instance
+
+        Returns
+        -------
+        dict
+            dictionary with ``condition: error_message`` pairs that will get
+            checked during ``_validate()`` run
+        """
+        return {
+            isinstance(self.n_workers, int)
+            and self.n_workers
+            > 0: f"n_workers should be positive integer, got {self.n_workers=}",
+        }
+
+    def _get_warnings(self):
+        """Get dictionary with ``condition: warning_message`` pairs that ensure
+        the good usage of the backend instance
+
+        Returns
+        -------
+        dict
+            dictionary with ``condition: warning_message`` pairs that will get
+            checked during ``_validate()`` run
+        """
+        return dict()
+
+    def _validate(self):
+        """Check correctness (e.g. ``dask`` is installed if using ``backend='dask'``)
+        and good usage (e.g. ``n_workers=1`` if backend is serial) of the backend
+
+        Raises
+        ------
+        ValueError
+            if one of the conditions in :meth:`_get_checks` is ``True``
+        """
+        for check, msg in self._get_checks().items():
+            if not check:
+                raise ValueError(msg)
+        for check, msg in self._get_warnings().items():
+            if not check:
+                warnings.warn(msg)
+
+    def apply(self, func: Callable, computations: list) -> list:
+        """map function `func` to all tasks in the `computations` list
+
+        Main method that will get called when using an instance of
+        ``BackendBase``.  It is equivalent to running ``[func(item) for item in
+        computations]`` while using the parallel backend capabilities.
+
+        Parameters
+        ----------
+        func : Callable
+            function to be called on each of the tasks in computations list
+        computations : list
+            computation tasks to apply function to
+
+        Returns
+        -------
+        list
+            list of results of the function
+
+        """
+        raise NotImplementedError
+
+
+class BackendSerial(BackendBase):
+    """A built-in backend that does serial execution of the function, without any
+    parallelization.
+
+    Parameters
+    ----------
+    n_workers : int
+        Is ignored in this class, and if ``n_workers`` > 1, a warning will be
+        given.
+
+
+    .. versionadded:: 2.8.0
+    """
+
+    def _get_warnings(self):
+        """Get dictionary with ``condition: warning_message`` pairs that ensure
+        the good usage of the backend instance. Here, it checks if the number
+        of workers is not 1, otherwise gives warning.
+
+        Returns
+        -------
+        dict
+            dictionary with ``condition: warning_message`` pairs that will get
+            checked during ``_validate()`` run
+        """
+        return {
+            self.n_workers
+            == 1: "n_workers is ignored when executing with backend='serial'"
+        }
+
+    def apply(self, func: Callable, computations: list) -> list:
+        """
+        Serially applies `func` to each task object in ``computations``.
+
+        Parameters
+        ----------
+        func : Callable
+            function to be called on each of the tasks in computations list
+        computations : list
+            computation tasks to apply function to
+
+        Returns
+        -------
+        list
+            list of results of the function
+        """
+        return [func(task) for task in computations]
+
+
+class BackendMultiprocessing(BackendBase):
+    """A built-in backend that executes a given function using the
+    :meth:`multiprocessing.Pool.map <multiprocessing.pool.Pool.map>` method.
+
+    Parameters
+    ----------
+    n_workers : int
+        number of processes in :class:`multiprocessing.Pool
+        <multiprocessing.pool.Pool>` to distribute the workload
+        between. Must be a positive integer.
+
+    Examples
+    --------
+
+    .. code-block:: python
+
+        from MDAnalysis.analysis.backends import BackendMultiprocessing
+        import multiprocessing as mp
+
+        backend_obj = BackendMultiprocessing(n_workers=mp.cpu_count())
+
+
+    .. versionadded:: 2.8.0
+
+    """
+
+    def apply(self, func: Callable, computations: list) -> list:
+        """Applies `func` to each object in ``computations`` using `multiprocessing`'s `Pool.map`.
+
+        Parameters
+        ----------
+        func : Callable
+            function to be called on each of the tasks in computations list
+        computations : list
+            computation tasks to apply function to
+
+        Returns
+        -------
+        list
+            list of results of the function
+        """
+        from multiprocessing import Pool
+
+        with Pool(processes=self.n_workers) as pool:
+            results = pool.map(func, computations)
+        return results
+
+
+class BackendDask(BackendBase):
+    """A built-in backend that executes a given function with *dask*.
+
+    Execution is performed with the :func:`dask.compute` function of
+    :class:`dask.delayed.Delayed` object (created with
+    :func:`dask.delayed.delayed`) with ``scheduler='processes'`` and
+    ``chunksize=1`` (this ensures uniform distribution of tasks among
+    processes). Requires the `dask package <https://docs.dask.org/en/stable/>`_
+    to be `installed <https://docs.dask.org/en/stable/install.html>`_.
+
+    Parameters
+    ----------
+    n_workers : int
+        number of processes in to distribute the workload
+        between. Must be a positive integer. Workers are actually
+        :class:`multiprocessing.pool.Pool` processes, but they use a different and
+        more flexible `serialization protocol
+        <https://docs.dask.org/en/stable/phases-of-computation.html#graph-serialization>`_.
+
+    Examples
+    --------
+
+    .. code-block:: python
+
+        from MDAnalysis.analysis.backends import BackendDask
+        import multiprocessing as mp
+
+        backend_obj = BackendDask(n_workers=mp.cpu_count())
+
+
+    .. versionadded:: 2.8.0
+
+    """
+
+    def apply(self, func: Callable, computations: list) -> list:
+        """Applies `func` to each object in ``computations``.
+
+        Parameters
+        ----------
+        func : Callable
+            function to be called on each of the tasks in computations list
+        computations : list
+            computation tasks to apply function to
+
+        Returns
+        -------
+        list
+            list of results of the function
+        """
+        from dask.delayed import delayed
+        import dask
+
+        computations = [delayed(func)(task) for task in computations]
+        results = dask.compute(
+            computations,
+            scheduler="processes",
+            chunksize=1,
+            num_workers=self.n_workers,
+        )[0]
+        return results
+
+    def _get_checks(self):
+        """Get dictionary with ``condition: error_message`` pairs that ensure the
+        validity of the backend instance. Here checks if ``dask`` module is
+        installed in the environment.
+
+        Returns
+        -------
+        dict
+            dictionary with ``condition: error_message`` pairs that will get
+            checked during ``_validate()`` run
+        """
+        base_checks = super()._get_checks()
+        checks = {
+            is_installed("dask"): (
+                "module 'dask' is missing. Please install 'dask': "
+                "https://docs.dask.org/en/stable/install.html"
+            )
+        }
+        return base_checks | checks