Skip to content

Commit

Permalink
Merge pull request #15 from Becksteinlab/new-order
Browse files Browse the repository at this point in the history 
StreamReaderBase, Gromacs fixes
  • Loading branch information
@ljwoods2
ljwoods2 authored Oct 7, 2024
2 parents 26ae789 + 8647503 commit 8eb88d7
Show file tree
Hide file tree
Showing 10 changed files with 12,110 additions and 175 deletions.
12 changes: 9 additions & 3 deletions imdclient/IMDClient.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -660,7 +660,9 @@ def __init__(
raise ValueError("pause_empty_proportion must be between 0 and 1")
self._pause_empty_proportion = pause_empty_proportion
if unpause_empty_proportion < 0 or unpause_empty_proportion > 1:
raise ValueError("unpause_empty_proportion must be between 0 and 1")
raise ValueError(
"unpause_empty_proportion must be between 0 and 1"
)
self._unpause_empty_proportion = unpause_empty_proportion

if buffer_size <= 0:
Expand Down Expand Up @@ -727,7 +729,9 @@ def wait_for_space(self):
logger.debug("IMDProducer: Noticing consumer finished")
raise EOFError
except Exception as e:
logger.debug(f"IMDProducer: Error waiting for space in buffer: {e}")
logger.debug(
f"IMDProducer: Error waiting for space in buffer: {e}"
)

def pop_empty_imdframe(self):
logger.debug("IMDProducer: Getting empty frame")
Expand Down Expand Up @@ -773,7 +777,9 @@ def pop_full_imdframe(self):
imdf = self._full_q.get()
else:
with self._full_imdf_avail:
while self._full_q.qsize() == 0 and not self._producer_finished:
while (
self._full_q.qsize() == 0 and not self._producer_finished
):
self._full_imdf_avail.wait()

if self._producer_finished and self._full_q.qsize() == 0:
Expand Down
164 changes: 19 additions & 145 deletions imdclient/IMDREADER.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -43,13 +43,6 @@
"""

import queue
from MDAnalysis.coordinates.base import (
ReaderBase,
FrameIteratorIndices,
FrameIteratorAll,
FrameIteratorSliced,
)
from MDAnalysis.coordinates import core
from MDAnalysis.lib.util import store_init_arguments

Expand All @@ -58,14 +51,15 @@
from .utils import *
import numpy as np
import logging
import warnings

from typing import Optional
import numbers

from .streambase import StreamReaderBase

logger = logging.getLogger("imdclient.IMDClient")


class IMDReader(ReaderBase):
class IMDReader(StreamReaderBase):
"""
Reader for IMD protocol packets.
"""
Expand All @@ -91,21 +85,23 @@ def __init__(
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.
"""
self._init_scope = True
self._reopen_called = False

super(IMDReader, self).__init__(filename, **kwargs)

logger.debug("IMDReader initializing")

if n_atoms is None:
raise ValueError("IMDReader: n_atoms must be specified")
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)
self._imdclient = IMDClient(
host, port, n_atoms, **kwargs
)

imdsinfo = self._imdclient.get_imdsessioninfo()
# NOTE: after testing phase, fail out on IMDv2
Expand All @@ -123,18 +119,9 @@ def __init__(
try:
self._read_next_timestep()
except StopIteration:
raise RuntimeError("IMDReader: No data found in stream")

def _read_next_timestep(self):
# No rewinding- to both load the first frame after __init__
# and access it again during iteration, we need to store first ts in mem
if not self._init_scope and self._frame == -1:
self._frame += 1
# can't simply return the same ts again- transformations would be applied twice
# instead, return the pre-transformed copy
return self._first_ts

return self._read_frame(self._frame + 1)
raise RuntimeError(
"IMDReader: No data found in stream"
)

def _read_frame(self, frame):

Expand All @@ -147,11 +134,9 @@ def _read_frame(self, frame):
self._frame = frame
self._load_imdframe_into_ts(imdf)

if self._init_scope:
self._first_ts = self.ts.copy()
self._init_scope = False

logger.debug(f"IMDReader: Loaded frame {self._frame}")
logger.debug(
f"IMDReader: Loaded frame {self._frame}"
)
return self.ts

def _load_imdframe_into_ts(self, imdf):
Expand All @@ -164,7 +149,9 @@ def _load_imdframe_into_ts(self, imdf):
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)
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
Expand All @@ -174,32 +161,6 @@ def _load_imdframe_into_ts(self, imdf):
if imdf.forces is not None:
self.ts.forces = imdf.forces

@property
def n_frames(self):
"""Changes as stream is processed unlike other readers"""
raise RuntimeError("IMDReader: n_frames is unknown")

def next(self):
"""Don't rewind after iteration. When _reopen() is called,
an error will be raised
"""
try:
ts = self._read_next_timestep()
except (EOFError, IOError):
# Don't rewind here like we normally would
raise StopIteration from None
else:
for auxname, reader in self._auxs.items():
ts = self._auxs[auxname].update_ts(ts)

ts = self._apply_transformations(ts)

return ts

def rewind(self):
"""Raise error on rewind"""
raise RuntimeError("IMDReader: Stream-based readers can't be rewound")

@staticmethod
def _format_hint(thing):
try:
Expand All @@ -214,90 +175,3 @@ def close(self):
self._imdclient.stop()
# NOTE: removeme after testing
logger.debug("IMDReader shut down gracefully.")

# Incompatible methods
def copy(self):
raise NotImplementedError("IMDReader does not support copying")

def _reopen(self):
if self._reopen_called:
raise RuntimeError("IMDReader: Cannot reopen IMD stream")
self._frame = -1
self._reopen_called = True

def __getitem__(self, frame):
"""This method from ProtoReader must be overridden
to prevent slicing that doesn't make sense in a stream.
"""
raise RuntimeError("IMDReader: Trajectory can only be read in for loop")

def check_slice_indices(self, start, stop, step):
"""Check frame indices are valid and clip to fit trajectory.
The usage follows standard Python conventions for :func:`range` but see
the warning below.
Parameters
----------
start : int or None
Starting frame index (inclusive). ``None`` corresponds to the default
of 0, i.e., the initial frame.
stop : int or None
Last frame index (exclusive). ``None`` corresponds to the default
of n_frames, i.e., it includes the last frame of the trajectory.
step : int or None
step size of the slice, ``None`` corresponds to the default of 1, i.e,
include every frame in the range `start`, `stop`.
Returns
-------
start, stop, step : tuple (int, int, int)
Integers representing the slice
Warning
-------
The returned values `start`, `stop` and `step` give the expected result
when passed in :func:`range` but gives unexpected behavior when passed
in a :class:`slice` when ``stop=None`` and ``step=-1``
This can be a problem for downstream processing of the output from this
method. For example, slicing of trajectories is implemented by passing
the values returned by :meth:`check_slice_indices` to :func:`range` ::
range(start, stop, step)
and using them as the indices to randomly seek to. On the other hand,
in :class:`MDAnalysis.analysis.base.AnalysisBase` the values returned
by :meth:`check_slice_indices` are used to splice the trajectory by
creating a :class:`slice` instance ::
slice(start, stop, step)
This creates a discrepancy because these two lines are not equivalent::
range(10, -1, -1) # [10, 9, 8, 7, 6, 5, 4, 3, 2, 1, 0]
range(10)[slice(10, -1, -1)] # []
"""
if start is not None:
raise ValueError(
"IMDReader: Cannot slice a stream, 'start' must be None"
)
if stop is not None:
raise ValueError(
"IMDReader: Cannot slice a stream, 'stop' must be None"
)
if step is not None:
if isinstance(step, numbers.Integral):
if step != 1:
raise ValueError(
"IMDReader: Cannot slice a stream, 'step' must be None or 1"
)

return start, stop, step

def __getstate__(self):
raise NotImplementedError("IMDReader does not support pickling")

def __setstate__(self, state: object):
raise NotImplementedError("IMDReader does not support pickling")
Loading

0 comments on commit 8eb88d7

Please sign in to comment.