westpa 2022.13__cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl

westpa/core/h5io.py

@@ -0,0 +1,1079 @@
+'''Miscellaneous routines to help with HDF5 input and output of WEST-related data.'''
+
+import collections
+import errno
+import getpass
+import os
+import posixpath
+import socket
+import sys
+import time
+import logging
+import warnings
+from pathlib import Path
+
+import h5py
+import numpy as np
+from numpy import index_exp
+from tables import NaturalNameWarning
+
+from mdtraj import Trajectory, join as join_traj
+from mdtraj.utils import in_units_of, import_, ensure_type
+from mdtraj.formats import HDF5TrajectoryFile
+from mdtraj.formats.hdf5 import _check_mode, Frames
+
+from .trajectory import WESTTrajectory
+
+try:
+    import psutil
+except ImportError:
+    psutil = None
+
+log = logging.getLogger(__name__)
+warnings.filterwarnings('ignore', category=NaturalNameWarning)
+
+#
+# Constants and globals
+#
+default_iter_prec = 8
+
+#
+# Helper functions
+#
+
+
+def resolve_filepath(path, constructor=h5py.File, cargs=None, ckwargs=None, **addtlkwargs):
+    '''Use a combined filesystem and HDF5 path to open an HDF5 file and return the
+    appropriate object. Returns (h5file, h5object). The file is opened using
+    ``constructor(filename, *cargs, **ckwargs)``.'''
+
+    cargs = cargs or ()
+    ckwargs = ckwargs or {}
+    ckwargs.update(addtlkwargs)
+    objpieces = collections.deque()
+    path = posixpath.normpath(path)
+
+    filepieces = path.split('/')
+    while filepieces:
+        testpath = '/'.join(filepieces)
+        if not testpath:
+            filepieces.pop()
+            continue
+        try:
+            h5file = constructor(testpath, *cargs, **ckwargs)
+        except IOError:
+            objpieces.appendleft(filepieces.pop())
+            continue
+        else:
+            return (h5file, h5file['/'.join([''] + list(objpieces)) if objpieces else '/'])
+    else:
+        # We don't provide a filename, because we're not sure where the filename stops
+        # and the HDF5 path begins.
+        raise IOError(errno.ENOENT, os.strerror(errno.ENOENT))
+
+
+def calc_chunksize(shape, dtype, max_chunksize=262144):
+    '''Calculate a chunk size for HDF5 data, anticipating that access will slice
+    along lower dimensions sooner than higher dimensions.'''
+
+    chunk_shape = list(shape)
+    dtype = np.dtype(dtype)
+    for idim in range(len(shape)):
+        chunk_nbytes = np.multiply.reduce(chunk_shape) * dtype.itemsize
+        while chunk_shape[idim] > 1 and chunk_nbytes > max_chunksize:
+            chunk_shape[idim] >>= 1  # divide by 2
+            chunk_nbytes = np.multiply.reduce(chunk_shape) * dtype.itemsize
+
+        if chunk_nbytes <= max_chunksize:
+            break
+
+    chunk_shape = tuple(chunk_shape)
+    return chunk_shape
+
+
+def tostr(b):
+    '''Convert a nonstandard string object ``b`` to str with the handling of the
+    case where ``b`` is bytes.'''
+
+    if b is None:
+        return None
+    elif isinstance(b, bytes):
+        return b.decode('utf-8')
+    else:
+        return str(b)
+
+
+def is_within_directory(directory, target):
+    abs_directory = os.path.abspath(directory)
+    abs_target = os.path.abspath(target)
+
+    prefix = os.path.commonprefix([abs_directory, abs_target])
+
+    return prefix == abs_directory
+
+
+def safe_extract(tar, path=".", members=None, *, numeric_owner=False):
+    for member in tar.getmembers():
+        member_path = os.path.join(path, member.name)
+        if not is_within_directory(path, member_path):
+            raise Exception("Attempted Path Traversal in Tar File")
+
+    tar.extractall(path, members, numeric_owner=numeric_owner, filter='fully_trusted')
+
+
+#
+# Group and dataset manipulation functions
+#
+
+
+def create_hdf5_group(parent_group, groupname, replace=False, creating_program=None):
+    '''Create (or delete and recreate) and HDF5 group named ``groupname`` within
+    the enclosing Group (object) ``parent_group``. If ``replace`` is True, then
+    the group is replaced if present; if False, then an error is raised if the
+    group is present. After the group is created, HDF5 attributes are set using
+    `stamp_creator_data`.
+    '''
+
+    if replace:
+        try:
+            del parent_group[groupname]
+        except KeyError:
+            pass
+
+    newgroup = parent_group.create_group(groupname)
+    stamp_creator_data(newgroup)
+    return newgroup
+
+
+#
+# Group and dataset labeling functions
+#
+
+
+def stamp_creator_data(h5group, creating_program=None):
+    '''Mark the following on the HDF5 group ``h5group``:
+
+      :creation_program:   The name of the program that created the group
+      :creation_user:      The username of the user who created the group
+      :creation_hostname:  The hostname of the machine on which the group was created
+      :creation_time:      The date and time at which the group was created, in the
+                           current locale.
+      :creation_unix_time: The Unix time (seconds from the epoch, UTC) at which the
+                           group was created.
+
+    This is meant to facilitate tracking the flow of data, but should not be considered
+    a secure paper trail (after all, anyone with write access to the HDF5 file can modify
+    these attributes).
+    '''
+    now = time.time()
+    attrs = h5group.attrs
+
+    attrs['creation_program'] = creating_program or sys.argv[0] or 'unknown program'
+    attrs['creation_user'] = getpass.getuser()
+    attrs['creation_hostname'] = socket.gethostname()
+    attrs['creation_unix_time'] = now
+    attrs['creation_time'] = time.strftime('%c', time.localtime(now))
+
+
+def get_creator_data(h5group):
+    '''Read back creator data as written by ``stamp_creator_data``, returning a dictionary with
+    keys as described for ``stamp_creator_data``. Missing fields are denoted with None.
+    The ``creation_time`` field is returned as a string.'''
+    attrs = h5group.attrs
+    d = dict()
+    for attr in ['creation_program', 'creation_user', 'creation_hostname', 'creation_unix_time', 'creation_time']:
+        d[attr] = attrs.get(attr)
+    return d
+
+
+def load_west(filename):
+    """Load WESTPA trajectory files from disk.
+
+    Parameters
+    ----------
+    filename : str
+        String filename of HDF Trajectory file.
+    """
+
+    with h5py.File(filename, 'r') as f:
+        iter_group_template = 'iter_{1:0{0}d}'
+        iter_prec = f.attrs['west_iter_prec']
+        trajectories = []
+        n = 0
+
+        iter_group_name = iter_group_template.format(iter_prec, n)
+        for iter_group_name in f['iterations']:
+            iter_group = f['iterations/' + iter_group_name]
+
+            if 'trajectories' in iter_group:
+                traj_link = iter_group['trajectories']
+                traj_filename = traj_link.file.filename
+
+                with WESTIterationFile(traj_filename) as traj_file:
+                    traj = traj_file.read_as_traj()
+            else:
+                # TODO: [HDF5] allow initializing trajectory without coordinates
+                raise ValueError("Missing trajectories for iteration %d" % n)
+
+            # pcoord is required
+            if 'pcoord' not in iter_group:
+                raise ValueError("Missing pcoords for iteration %d" % n)
+
+            raw_pcoord = iter_group['pcoord'][:]
+            if raw_pcoord.ndim != 3:
+                log.warning('pcoord is expected to be a 3-d ndarray instead of {}-d'.format(raw_pcoord.ndim))
+                continue
+            # ignore the first frame of each segment
+            if raw_pcoord.shape[1] == traj.n_frames + 1:
+                raw_pcoord = raw_pcoord[:, 1:, :]
+            elif raw_pcoord.shape[1] == traj.n_frames:
+                raw_pcoord = raw_pcoord[:, :, :]
+            else:
+                raise ValueError(
+                    "Inconsistent number of pcoords (%d) and frames (%d) for iteration %d" % (raw_pcoord.shape[1], traj.n_frames, n)
+                )
+
+            pcoords = np.concatenate(raw_pcoord, axis=0)
+            n_frames = raw_pcoord.shape[1]
+
+            if 'seg_index' in iter_group:
+                raw_pid = iter_group['seg_index']['parent_id'][:]
+
+                if np.any(raw_pid < 0):
+                    init_basis_ids = iter_group['ibstates']['istate_index']['basis_state_id'][:]
+                    init_ids = -(raw_pid[raw_pid < 0] + 1)
+                    raw_pid[raw_pid < 0] = [init_basis_ids[iid] for iid in init_ids]
+                parent_ids = raw_pid.repeat(n_frames, axis=0)
+            else:
+                parent_ids = None
+
+            traj.pcoords = pcoords
+            traj.parent_ids = parent_ids
+            trajectories.append(traj)
+
+            n += 1
+            iter_group_name = iter_group_template.format(iter_prec, n)
+
+    west_traj = join_traj(trajectories)
+
+    return west_traj
+
+
+###
+# Iteration range metadata
+###
+def stamp_iter_range(h5object, start_iter, stop_iter):
+    '''Mark that the HDF5 object ``h5object`` (dataset or group) contains data from iterations
+    start_iter <= n_iter < stop_iter.'''
+    h5object.attrs['iter_start'] = start_iter
+    h5object.attrs['iter_stop'] = stop_iter
+
+
+def get_iter_range(h5object):
+    '''Read back iteration range data written by ``stamp_iter_range``'''
+    return int(h5object.attrs['iter_start']), int(h5object.attrs['iter_stop'])
+
+
+def stamp_iter_step(h5group, iter_step):
+    '''Mark that the HDF5 object ``h5object`` (dataset or group) contains data with an
+    iteration step (stride) of iter_step).'''
+    h5group.attrs['iter_step'] = iter_step
+
+
+def get_iter_step(h5group):
+    '''Read back iteration step (stride) written by ``stamp_iter_step``'''
+    return int(h5group.attrs['iter_step'])
+
+
+def check_iter_range_least(h5object, iter_start, iter_stop):
+    '''Return True if the iteration range [iter_start, iter_stop) is
+    the same as or entirely contained within the iteration range stored
+    on ``h5object``.'''
+    obj_iter_start, obj_iter_stop = get_iter_range(h5object)
+    return obj_iter_start <= iter_start and obj_iter_stop >= iter_stop
+
+
+def check_iter_range_equal(h5object, iter_start, iter_stop):
+    '''Return True if the iteration range [iter_start, iter_stop) is
+    the same as the iteration range stored on ``h5object``.'''
+    obj_iter_start, obj_iter_stop = get_iter_range(h5object)
+    return obj_iter_start == iter_start and obj_iter_stop == iter_stop
+
+
+def get_iteration_entry(h5object, n_iter):
+    '''Create a slice for data corresponding to iteration ``n_iter`` in ``h5object``.'''
+    obj_iter_start, obj_iter_stop = get_iter_range(h5object)
+    if n_iter < obj_iter_start or n_iter >= obj_iter_stop:
+        raise IndexError('data for iteration {} not available in dataset {!r}'.format(n_iter, h5object))
+    return np.index_exp[n_iter - obj_iter_start]
+
+
+def get_iteration_slice(h5object, iter_start, iter_stop=None, iter_stride=None):
+    '''Create a slice for data corresponding to iterations [iter_start,iter_stop),
+    with stride iter_step, in the given ``h5object``.'''
+    obj_iter_start, obj_iter_stop = get_iter_range(h5object)
+
+    if iter_stop is None:
+        iter_stop = iter_start + 1
+    if iter_stride is None:
+        iter_stride = 1
+
+    if iter_start < obj_iter_start:
+        raise IndexError('data for iteration {} not available in dataset {!r}'.format(iter_start, h5object))
+    elif iter_start > obj_iter_stop:
+        raise IndexError('data for iteration {} not available in dataset {!r}'.format(iter_stop, h5object))
+
+    start_index = iter_start - obj_iter_start
+    stop_index = iter_stop - obj_iter_start
+    return np.index_exp[start_index:stop_index:iter_stride]
+
+
+###
+# Axis label metadata
+###
+def label_axes(h5object, labels, units=None):
+    '''Stamp the given HDF5 object with axis labels. This stores the axis labels
+    in an array of strings in an attribute called ``axis_labels`` on the given
+    object. ``units`` if provided is a corresponding list of units.'''
+
+    if len(labels) != len(h5object.shape):
+        raise ValueError('number of axes and number of labels do not match')
+
+    if units is None:
+        units = []
+
+    if len(units) and len(units) != len(labels):
+        raise ValueError('number of units labels does not match number of axes')
+
+    h5object.attrs['axis_labels'] = np.array([np.bytes_(i) for i in labels])
+
+    if len(units):
+        h5object.attrs['axis_units'] = np.array([np.bytes_(i) for i in units])
+
+
+NotGiven = object()
+
+
+def _get_one_attr(h5object, namelist, default=NotGiven):
+    attrs = dict(h5object.attrs)
+    for name in namelist:
+        try:
+            return attrs[name]
+        except KeyError:
+            pass
+    else:
+        if default is NotGiven:
+            raise KeyError('no such key')
+        else:
+            return default
+
+
+class WESTPAH5File(h5py.File):
+    '''Generalized input/output for WESTPA simulation (or analysis) data.'''
+
+    default_iter_prec = 8
+    _this_fileformat_version = 8
+
+    def __init__(self, *args, **kwargs):
+        # These values are used for creating files or reading files where this
+        # data is not stored. Otherwise, values stored as attributes on the root
+        # group are used instead.
+        arg_iter_prec = kwargs.pop('westpa_iter_prec', self.default_iter_prec)
+        arg_fileformat_version = kwargs.pop('westpa_fileformat_version', self._this_fileformat_version)
+        arg_creating_program = kwargs.pop('creating_program', None)
+
+        # Initialize h5py file
+        super().__init__(*args, **kwargs)
+
+        # Try to get iteration precision and I/O class version
+        h5file_iter_prec = _get_one_attr(self, ['westpa_iter_prec', 'west_iter_prec', 'wemd_iter_prec'], None)
+        h5file_fileformat_version = _get_one_attr(
+            self, ['westpa_fileformat_version', 'west_file_format_version', 'wemd_file_format_version'], None
+        )
+
+        self.iter_prec = h5file_iter_prec if h5file_iter_prec is not None else arg_iter_prec
+        self.fileformat_version = h5file_fileformat_version if h5file_fileformat_version is not None else arg_fileformat_version
+
+        # Ensure that file format attributes are stored, if the file is writable
+        if self.mode == 'r+':
+            self.attrs['westpa_iter_prec'] = self.iter_prec
+            self.attrs['westpa_fileformat_version'] = self.fileformat_version
+            if arg_creating_program:
+                stamp_creator_data(self, creating_program=arg_creating_program)
+
+    # Helper function to automatically replace a group, if it exists.
+    # Should really only be called when one is certain a dataset should be blown away.
+    def replace_dataset(self, *args, **kwargs):
+        try:
+            del self[args[0]]
+        except Exception:
+            pass
+        try:
+            del self[kwargs['name']]
+        except Exception:
+            pass
+
+        return self.create_dataset(*args, **kwargs)
+
+    # Iteration groups
+
+    def iter_object_name(self, n_iter, prefix='', suffix=''):
+        '''Return a properly-formatted per-iteration name for iteration
+        ``n_iter``. (This is used in create/require/get_iter_group, but may
+        also be useful for naming datasets on a per-iteration basis.)'''
+        return '{prefix}iter_{n_iter:0{prec}d}{suffix}'.format(n_iter=n_iter, prefix=prefix, suffix=suffix, prec=self.iter_prec)
+
+    def create_iter_group(self, n_iter, group=None):
+        '''Create a per-iteration data storage group for iteration number ``n_iter``
+        in the group ``group`` (which is '/iterations' by default).'''
+
+        if group is None:
+            group = self.require_group('/iterations')
+        return group.create_group(self.iter_object_name(n_iter))
+
+    def require_iter_group(self, n_iter, group=None):
+        '''Ensure that a per-iteration data storage group for iteration number ``n_iter``
+        is available in the group ``group`` (which is '/iterations' by default).'''
+        if group is None:
+            group = self.require_group('/iterations')
+        return group.require_group(self.iter_object_name(n_iter))
+
+    def get_iter_group(self, n_iter, group=None):
+        '''Get the per-iteration data group for iteration number ``n_iter`` from within
+        the group ``group`` ('/iterations' by default).'''
+        if group is None:
+            group = self['/iterations']
+        return group[self.iter_object_name(n_iter)]
+
+
+class WESTIterationFile(HDF5TrajectoryFile):
+    def __init__(self, file, mode='r', force_overwrite=True, compression='zlib', link=None):
+        if isinstance(file, (str, Path)):
+            super(WESTIterationFile, self).__init__(file, mode, force_overwrite, compression)
+        else:
+            try:
+                self._init_from_handle(file)  # If a WESTIterationFile object, just make sure it's open correctly
+            except AttributeError:
+                raise ValueError('unknown input type: %s' % str(type(file)))
+
+    def _init_from_handle(self, handle: HDF5TrajectoryFile):
+        self._handle = handle
+        self._open = handle.isopen != 0
+        self.mode = mode = handle.mode  # the mode in which the file was opened?
+
+        if mode not in ['r', 'w', 'a']:
+            raise ValueError("mode must be one of ['r', 'w', 'a']")
+
+        # import tables
+        self.tables = import_('tables')
+
+        if mode == 'w':
+            # what frame are we currently reading or writing at?
+            self._frame_index = 0
+            # do we need to write the header information?
+            self._needs_initialization = True
+
+        elif mode == 'a':
+            try:
+                self._frame_index = len(self._handle.root.coordinates)
+                self._needs_initialization = False
+            except self.tables.NoSuchNodeError:
+                self._frame_index = 0
+                self._needs_initialization = True
+        elif mode == 'r':
+            self._frame_index = 0
+            self._needs_initialization = False
+
+    def __contains__(self, path):
+        try:
+            self._get_node('/', path)
+        except self.tables.NoSuchNodeError:
+            return False
+
+        return True
+
+    def read(self, frame_indices=None, atom_indices=None):
+        _check_mode(self.mode, ('r',))
+
+        if frame_indices is None:
+            frame_slice = slice(None)
+            self._frame_index += frame_slice.stop - frame_slice.start
+        else:
+            frame_slice = ensure_type(frame_indices, dtype=int, ndim=1, name='frame_indices', warn_on_cast=False)
+            if not np.all(frame_slice < self._handle.root.coordinates.shape[0]):
+                raise ValueError(
+                    'As a zero-based index, the entries in '
+                    'frame_slice must all be less than the number of frames '
+                    'in the trajectory, %d' % self._handle.root.coordinates.shape[0]
+                )
+            if not np.all(frame_slice >= 0):
+                raise ValueError('The entries in frame_indices must be greater ' 'than or equal to zero')
+            self._frame_index += frame_slice[-1] - frame_slice[0]
+
+        if atom_indices is None:
+            # get all of the atoms
+            atom_slice = slice(None)
+        else:
+            atom_slice = ensure_type(atom_indices, dtype=int, ndim=1, name='atom_indices', warn_on_cast=False)
+            if not np.all(atom_slice < self._handle.root.coordinates.shape[1]):
+                raise ValueError(
+                    'As a zero-based index, the entries in '
+                    'atom_indices must all be less than the number of atoms '
+                    'in the trajectory, %d' % self._handle.root.coordinates.shape[1]
+                )
+            if not np.all(atom_slice >= 0):
+                raise ValueError('The entries in atom_indices must be greater ' 'than or equal to zero')
+
+        def get_item(node, key):
+            if not isinstance(key, tuple):
+                return node.__getitem__(key)
+
+            n_list_like = 0
+            new_keys = []
+            for item in key:
+                if not isinstance(item, slice):
+                    try:
+                        d = np.diff(item)
+                        if len(d) == 0:
+                            item = item[0]
+                        elif np.all(d == d[0]):
+                            item = slice(item[0], item[-1] + d[0], d[0])
+                        else:
+                            n_list_like += 1
+                    except Exception:
+                        n_list_like += 1
+                new_keys.append(item)
+            new_keys = tuple(new_keys)
+
+            if n_list_like <= 1:
+                return node.__getitem__(new_keys)
+
+            data = node
+            for i, item in enumerate(new_keys):
+                dkey = [slice(None)] * len(key)
+                dkey[i] = item
+                dkey = tuple(dkey)
+                data = data.__getitem__(dkey)
+
+            return data
+
+        def get_field(name, slice, out_units, can_be_none=True):
+            try:
+                node = self._get_node(where='/', name=name)
+                data = get_item(node, slice)
+                in_units = node.attrs.units
+                if not isinstance(in_units, str):
+                    in_units = in_units.decode()
+                data = in_units_of(data, in_units, out_units)
+                return data
+            except self.tables.NoSuchNodeError:
+                if can_be_none:
+                    return None
+                raise
+
+        frames = Frames(
+            coordinates=get_field('coordinates', (frame_slice, atom_slice, slice(None)), out_units='nanometers', can_be_none=False),
+            time=get_field('time', frame_slice, out_units='picoseconds'),
+            cell_lengths=get_field('cell_lengths', (frame_slice, slice(None)), out_units='nanometers'),
+            cell_angles=get_field('cell_angles', (frame_slice, slice(None)), out_units='degrees'),
+            velocities=get_field('velocities', (frame_slice, atom_slice, slice(None)), out_units='nanometers/picosecond'),
+            kineticEnergy=get_field('kineticEnergy', frame_slice, out_units='kilojoules_per_mole'),
+            potentialEnergy=get_field('potentialEnergy', frame_slice, out_units='kilojoules_per_mole'),
+            temperature=get_field('temperature', frame_slice, out_units='kelvin'),
+            alchemicalLambda=get_field('lambda', frame_slice, out_units='dimensionless'),
+        )
+
+        return frames
+
+    def _has_node(self, where, name):
+        try:
+            self._get_node(where, name=name)
+        except self.tables.NoSuchNodeError:
+            return False
+
+        return True
+
+    def has_topology(self):
+        return self._has_node('/', 'topology')
+
+    def has_pointer(self):
+        return self._has_node('/', 'pointer')
+
+    def has_restart(self, segment):
+        return self._has_node('/restart', '%d_%d' % (segment.n_iter, segment.seg_id))
+
+    def has_seglog(self, segment):
+        return self._has_node('/log', '%d_%d' % (segment.n_iter, segment.seg_id))
+
+    def write_data(self, where, name, data):
+        node = self._get_node(where=where, name=name)
+        node.append(data)
+
+    def read_data(self, where, name):
+        node = self._get_node(where=where, name=name)
+        return node.read()
+
+    def read_as_traj(self, iteration=None, segment=None, atom_indices=None):
+        _check_mode(self.mode, ('r',))
+
+        pnode = self._get_node(where='/', name='pointer')
+
+        iter_labels = pnode[:, 0]
+        seg_labels = pnode[:, 1]
+
+        if iteration is None and segment is None:
+            frame_indices = slice(None)
+        elif isinstance(iteration, (np.integer, int)) and isinstance(segment, (np.integer, int)):
+            frame_torf = np.logical_and(iter_labels == iteration, seg_labels == segment)
+            frame_indices = np.arange(len(iter_labels))[frame_torf]
+        else:
+            raise ValueError("iteration and segment must be integers and provided at the same time")
+
+        if len(frame_indices) == 0:
+            raise ValueError(f"no frame was selected: iteration={iteration}, segment={segment}, atom_indices={atom_indices}")
+
+        iter_labels = iter_labels[frame_indices]
+        seg_labels = seg_labels[frame_indices]
+
+        topology = self.topology
+        if atom_indices is not None:
+            topology = topology.subset(atom_indices)
+
+        data = self.read(frame_indices=frame_indices, atom_indices=atom_indices)
+        if len(data) == 0:
+            return Trajectory(xyz=np.zeros((0, topology.n_atoms, 3)), topology=topology)
+
+        in_units_of(data.coordinates, self.distance_unit, Trajectory._distance_unit, inplace=True)
+        in_units_of(data.cell_lengths, self.distance_unit, Trajectory._distance_unit, inplace=True)
+
+        return WESTTrajectory(
+            data.coordinates,
+            topology=topology,
+            time=data.time,
+            unitcell_lengths=data.cell_lengths,
+            unitcell_angles=data.cell_angles,
+            iter_labels=iter_labels,
+            seg_labels=seg_labels,
+            pcoords=None,
+        )
+
+    def read_restart(self, segment):
+        if self.has_restart(segment):
+            data = self.read_data('/restart/%d_%d' % (segment.n_iter, segment.seg_id), 'data')
+            segment.data['iterh5/restart'] = data
+        else:
+            raise ValueError('no restart data available for {}'.format(str(segment)))
+
+    def read_seglog(self, segment):
+        if self.has_seglog(segment):
+            data = self.read_data('/log/%d_%d' % (segment.niter, segment.seg_id), 'data')
+            segment.data['iterh5/log'] = data
+        else:
+            raise ValueError('no log data available for {}'.format(str(segment)))
+
+    def write_segment(self, segment, pop=False):
+        n_iter = segment.n_iter
+
+        self.root._v_attrs['n_iter'] = n_iter
+
+        if pop:
+            get_data = segment.data.pop
+        else:
+            get_data = segment.data.get
+
+        traj = get_data('iterh5/trajectory', None)
+        restart = get_data('iterh5/restart', None)
+        slog = get_data('iterh5/log', None)
+
+        if traj is not None:
+            # create trajectory object or if already is, skip.
+            if not isinstance(traj, WESTTrajectory):
+                traj = WESTTrajectory(traj, iter_labels=n_iter, seg_labels=segment.seg_id)
+            if traj.n_frames == 0:
+                # we may consider logging warnings instead throwing errors for later.
+                # right now this is good for debugging purposes
+                raise ValueError('no trajectory data present for %s' % repr(segment))
+
+            if n_iter == 0:
+                base_time = 0
+            else:
+                iter_duration = traj.time[-1] - traj.time[0]
+                base_time = iter_duration * (n_iter - 1)
+
+            traj.time -= traj.time[0]
+            traj.time += base_time
+
+            # pointers
+            if not self.has_pointer():
+                self._create_earray('/', name='pointer', atom=self.tables.Int64Atom(), shape=(0, 2))
+
+            try:
+                existing_labels = np.where(self.root['pointer'][:, 1] == traj.seg_labels[0])[0]
+            except ValueError:
+                existing_labels = []
+
+            iter_idx = traj.iter_labels
+            seg_idx = traj.seg_labels
+
+            pointers = np.stack((iter_idx, seg_idx)).T
+            # This is to deal with the case where the simulation ended mid-iteration and some segment data are saved but itself not marked as complete yet.
+            if [segment.seg_id] in self.root['pointer'][:, 1]:
+                needed_extra = len(pointers) - len(existing_labels)
+                # If previous run did not save all frames
+                if needed_extra > 0:
+                    # Write extra rows for extra frames
+                    self.write_data('/', 'pointer', pointers[-needed_extra:])
+
+                    # write trajectory for the extra rows
+                    output_dict = {
+                        'coordinates': in_units_of(traj.xyz[-needed_extra:], Trajectory._distance_unit, self.distance_unit),
+                        'time': traj.time[-needed_extra:],
+                    }
+                    if traj.unitcell_lengths:
+                        output_dict['cell_lengths'] = in_units_of(
+                            traj.unitcell_lengths[-needed_extra:], Trajectory._distance_unit, self.distance_unit
+                        )
+                        output_dict['cell_angles'] = (traj.unitcell_angles[-needed_extra:],)
+
+                    self.write(**output_dict)
+                    needed_extra = len(existing_labels)
+                elif needed_extra < 0:
+                    # Extra frames found, turning pointer for those rows to sentinel
+                    log.warning(
+                        f'Extra frames for segment {n_iter}_{segment.seg_id} found in WESTIterationFile. Overwriting extra frame pointers with sentinal [-n_iter, seg_id].'
+                    )
+                    for row_idx in range(needed_extra, 0):
+                        self.root['pointer'][existing_labels[row_idx]] = [-n_iter, segment.seg_id]
+                    needed_extra = len(pointers)
+                else:
+                    # Number of Frames match. None will return all frames and traj.
+                    needed_extra = None
+
+                # Replace existing rows with corresponding data, up until specified
+                self.replace_frames(existing_labels[:needed_extra], traj[:needed_extra])
+
+            else:
+                # Write pointers
+                self.write_data('/', 'pointer', pointers)
+
+                # trajectory
+                self.write(
+                    coordinates=in_units_of(traj.xyz, Trajectory._distance_unit, self.distance_unit),
+                    time=traj.time,
+                    cell_lengths=in_units_of(traj.unitcell_lengths, Trajectory._distance_unit, self.distance_unit),
+                    cell_angles=traj.unitcell_angles,
+                )
+
+            # topology
+            try:
+                if self.mode == 'a':
+                    if not self.has_topology():
+                        self.topology = traj.topology
+                elif self.mode == 'w':
+                    self.topology = traj.topology
+            except (ModuleNotFoundError, ImportError):
+                pass
+
+        # restart
+        if restart is not None:
+            if self.has_restart(segment):
+                self._remove_node('/restart', name='%d_%d' % (segment.n_iter, segment.seg_id), recursive=True)
+
+            self._create_array(
+                '/restart/%d_%d' % (segment.n_iter, segment.seg_id),
+                name='data',
+                atom=self.tables.StringAtom(itemsize=len(restart)),
+                obj=restart,
+                createparents=True,
+            )
+
+        if slog is not None:
+            if self._has_node('/log', name='%d_%d' % (segment.n_iter, segment.seg_id)):
+                self._remove_node('/log', name='%d_%d' % (segment.n_iter, segment.seg_id), recursive=True)
+
+            self._create_array(
+                '/log/%d_%d' % (segment.n_iter, segment.seg_id),
+                name='data',
+                atom=self.tables.StringAtom(itemsize=len(slog)),
+                obj=slog,
+                createparents=True,
+            )
+
+    def scrub_data(self):
+        '''Method to remove existing coordinates, pointers etc. while preserving topology'''
+        for node in ['log', 'restart', 'time', 'coordinates', 'pointer', 'cell_angles', 'cell_lengths']:
+            try:
+                self._remove_node('/', node, recursive=True)
+            except self.tables.exceptions.NoSuchNodeError:
+                pass
+        self._frame_index = 0
+        self.root._v_attrs.n_iter = 0
+        self.flush()
+
+    def replace_frames(self, rows, traj):
+        datasets = {'coordinates': 'xyz', 'time': 'time', 'cell_angles': 'unitcell_angles', 'cell_lengths': 'unitcell_lengths'}
+
+        for ptkey, mdkey in datasets.items():
+            if self._has_node('/', ptkey) and getattr(traj, mdkey) is not None:
+                for frame_idx, row in enumerate(rows):
+                    self.root[ptkey][row] = getattr(traj, mdkey)[frame_idx]
+
+    @property
+    def _create_group(self):
+        if self.tables.__version__ >= '3.0.0':
+            return self._handle.create_group
+        return self._handle.createGroup
+
+    @property
+    def _create_array(self):
+        if self.tables.__version__ >= '3.0.0':
+            return self._handle.create_array
+        return self._handle.createArray
+
+    @property
+    def _remove_node(self):
+        if self.tables.__version__ >= '3.0.0':
+            return self._handle.remove_node
+        return self._handle.removeNode
+
+
+### Generalized WE dataset access classes
+
+
+class DSSpec:
+    '''Generalized WE dataset access'''
+
+    def get_iter_data(self, n_iter, seg_slice=index_exp[:]):
+        raise NotImplementedError
+
+    def get_segment_data(self, n_iter, seg_id):
+        return self.get_iter_data(n_iter)[seg_id]
+
+    def __getstate__(self):
+        d = dict(self.__dict__)
+        if '_h5file' in d:
+            d['_h5file'] = None
+        return d
+
+    def __setstate__(self, state):
+        self.__dict__.update(state)
+
+
+class FileLinkedDSSpec(DSSpec):
+    '''Provide facilities for accessing WESTPA HDF5 files, including auto-opening and the ability
+    to pickle references to such files for transmission (through, e.g., the work manager), provided
+    that the HDF5 file can be accessed by the same path on both the sender and receiver.'''
+
+    def __init__(self, h5file_or_name):
+        self._h5file = None
+        self._h5filename = None
+
+        try:
+            self._h5filename = os.path.abspath(h5file_or_name.filename)
+        except AttributeError:
+            self._h5filename = h5file_or_name
+            self._h5file = None
+        else:
+            self._h5file = h5file_or_name
+
+    @property
+    def h5file(self):
+        '''Lazily open HDF5 file. This is required because allowing an open HDF5
+        file to cross a fork() boundary generally corrupts the internal state of
+        the HDF5 library.'''
+        if self._h5file is None:
+            self._h5file = WESTPAH5File(self._h5filename, 'r')
+        return self._h5file
+
+
+class SingleDSSpec(FileLinkedDSSpec):
+    @classmethod
+    def from_string(cls, dsspec_string, default_h5file):
+        alias = None
+
+        h5file = default_h5file
+        fields = dsspec_string.split(',')
+        dsname = fields[0]
+        slice = None
+
+        for field in (field.strip() for field in fields[1:]):
+            k, v = field.split('=')
+            k = k.lower()
+            if k == 'alias':
+                alias = v
+            elif k == 'slice':
+                try:
+                    slice = eval('np.index_exp' + v)
+                except SyntaxError:
+                    raise SyntaxError('invalid index expression {!r}'.format(v))
+            elif k == 'file':
+                h5file = v
+            else:
+                raise ValueError('invalid dataset option {!r}'.format(k))
+
+        return cls(h5file, dsname, alias, slice)
+
+    def __init__(self, h5file_or_name, dsname, alias=None, slice=None):
+        FileLinkedDSSpec.__init__(self, h5file_or_name)
+        self.dsname = dsname
+        self.alias = alias or dsname
+        self.slice = np.index_exp[slice] if slice else None
+
+
+class SingleIterDSSpec(SingleDSSpec):
+    def get_iter_data(self, n_iter, seg_slice=index_exp[:]):
+        if self.slice:
+            return self.h5file.get_iter_group(n_iter)[self.dsname][seg_slice + self.slice]
+        else:
+            return self.h5file.get_iter_group(n_iter)[self.dsname][seg_slice]
+
+
+class SingleSegmentDSSpec(SingleDSSpec):
+    def get_iter_data(self, n_iter, seg_slice=index_exp[:]):
+        if self.slice:
+            return self.h5file.get_iter_group(n_iter)[self.dsname][seg_slice + index_exp[:] + self.slice]
+        else:
+            return self.h5file.get_iter_group(n_iter)[self.dsname][seg_slice]
+
+    def get_segment_data(self, n_iter, seg_id):
+        if self.slice:
+            return self.h5file.get_iter_group(n_iter)[np.index_exp[seg_id, :] + self.slice]
+        else:
+            return self.h5file.get_iter_group(n_iter)[seg_id]
+
+
+class FnDSSpec(FileLinkedDSSpec):
+    def __init__(self, h5file_or_name, fn):
+        FileLinkedDSSpec.__init__(self, h5file_or_name)
+        self.fn = fn
+
+    def get_iter_data(self, n_iter, seg_slice=index_exp[:]):
+        return self.fn(n_iter, self.h5file.get_iter_group(n_iter))[seg_slice]
+
+
+class MultiDSSpec(DSSpec):
+    def __init__(self, dsspecs):
+        self.dsspecs = dsspecs
+
+    def get_iter_data(self, n_iter, seg_slice=index_exp[:]):
+        datasets = [dsspec.get_iter_data(n_iter) for dsspec in self.dsspecs]
+
+        ncols = 0
+        nsegs = None
+        npts = None
+        for iset, dset in enumerate(datasets):
+            if nsegs is None:
+                nsegs = dset.shape[0]
+            elif dset.shape[0] != nsegs:
+                raise TypeError('dataset {} has incorrect first dimension (number of segments)'.format(self.dsspecs[iset]))
+            if npts is None:
+                npts = dset.shape[1]
+            elif dset.shape[1] != npts:
+                raise TypeError('dataset {} has incorrect second dimension (number of time points)'.format(self.dsspecs[iset]))
+
+            if dset.ndim < 2:
+                # scalar per segment or scalar per iteration
+                raise TypeError('dataset {} has too few dimensions'.format(self.dsspecs[iset]))
+            elif dset.ndim > 3:
+                # array per timepoint
+                raise TypeError('dataset {} has too many dimensions'.format(self.dsspecs[iset]))
+            elif dset.ndim == 2:
+                # scalar per timepoint
+                ncols += 1
+            else:
+                # vector per timepoint
+                ncols += dset.shape[-1]
+
+        output_dtype = np.result_type(*[ds.dtype for ds in datasets])
+        output_array = np.empty((nsegs, npts, ncols), dtype=output_dtype)
+
+        ocol = 0
+        for iset, dset in enumerate(datasets):
+            if dset.ndim == 2:
+                output_array[:, :, ocol] = dset[...]
+                ocol += 1
+            elif dset.ndim == 3:
+                output_array[:, :, ocol : (ocol + dset.shape[-1])] = dset[...]
+                ocol += dset.shape[-1]
+
+        return output_array[seg_slice]
+
+
+class IterBlockedDataset:
+    @classmethod
+    def empty_like(cls, blocked_dataset):
+        source = blocked_dataset.data if blocked_dataset.data is not None else blocked_dataset.dataset
+
+        newbds = cls(
+            np.empty(source.shape, source.dtype),
+            attrs={'iter_start': blocked_dataset.iter_start, 'iter_stop': blocked_dataset.iter_stop},
+        )
+        return newbds
+
+    def __init__(self, dataset_or_array, attrs=None):
+        try:
+            dataset_or_array.attrs
+        except AttributeError:
+            self.dataset = None
+            self.data = dataset_or_array
+            if attrs is None:
+                raise ValueError('attribute dictionary containing iteration bounds must be provided')
+            self.iter_shape = self.data.shape[1:]
+            self.dtype = self.data.dtype
+        else:
+            self.dataset = dataset_or_array
+            attrs = self.dataset.attrs
+            self.data = None
+            self.iter_shape = self.dataset.shape[1:]
+            self.dtype = self.dataset.dtype
+
+        self.iter_start = attrs['iter_start']
+        self.iter_stop = attrs['iter_stop']
+
+    def cache_data(self, max_size=None):
+        '''Cache this dataset in RAM. If ``max_size`` is given, then only cache if the entire dataset
+        fits in ``max_size`` bytes. If ``max_size`` is the string 'available', then only cache if
+        the entire dataset fits in available RAM, as defined by the ``psutil`` module.'''
+
+        if max_size is not None:
+            dssize = self.dtype.itemsize * np.multiply.reduce(self.dataset.shape)
+            if max_size == 'available' and psutil is not None:
+                avail_bytes = psutil.virtual_memory().available
+                if dssize > avail_bytes:
+                    return
+            elif isinstance(max_size, str):
+                return
+            else:
+                if dssize > max_size:
+                    return
+        if self.dataset is not None:
+            if self.data is None:
+                self.data = self.dataset[...]
+
+    def drop_cache(self):
+        if self.dataset is not None:
+            del self.data
+            self.data = None
+
+    def iter_entry(self, n_iter):
+        if n_iter < self.iter_start:
+            raise IndexError('requested iteration {} less than first stored iteration {}'.format(n_iter, self.iter_start))
+
+        source = self.data if self.data is not None else self.dataset
+        return source[n_iter - self.iter_start]
+
+    def iter_slice(self, start=None, stop=None):
+        start = start or self.iter_start
+        stop = stop or self.iter_stop
+        step = 1  # strided retrieval not implemented yet
+
+        # if step % self.iter_step > 0:
+        #    raise TypeError('dataset {!r} stored with stride {} cannot be accessed with stride {}'
+        #                    .format(self.dataset, self.iter_step, step))
+        if start < self.iter_start:
+            raise IndexError('requested start {} less than stored start {}'.format(start, self.iter_start))
+        elif stop > self.iter_stop:
+            stop = self.iter_stop
+
+        source = self.data if self.data is not None else self.dataset
+        return source[start - self.iter_start : stop - self.iter_start : step]