westpa 2022.10__cp312-cp312-macosx_11_0_arm64.whl

westpa/analysis/statistics.py

@@ -0,0 +1,27 @@
+import numpy as np
+
+
+def time_average(observable, iterations):
+    """Compute the time average of an observable.
+
+    Parameters
+    ----------
+    observable : Callable[[Walker], ArrayLike]
+        Function that takes a walker as input and returns a number or
+        a fixed-size array of numbers.
+    iterations : Sequence[Iteration]
+        Sequence of iterations over which to compute the average.
+
+    Returns
+    -------
+    ArrayLike
+        The time average of `observable` over `iterations`.
+
+    """
+    for iteration in iterations:
+        values = [observable(walker) for walker in iteration]
+        try:
+            value += np.dot(iteration.weights, values)
+        except NameError:
+            value = np.dot(iteration.weights, values)
+    return value / len(iterations)

westpa/analysis/trajectories.py

@@ -0,0 +1,360 @@
+import concurrent.futures as cf
+import functools
+import inspect
+import operator
+import os
+
+import mdtraj
+import numpy as np
+
+from tqdm import tqdm
+from typing import Callable
+from westpa.analysis.core import Walker, Trace
+from westpa.core.states import InitialState
+from westpa.core.h5io import WESTIterationFile
+
+
+class Trajectory:
+    """A callable that returns the trajectory of a walker or trace.
+
+    Parameters
+    ----------
+    fget : callable
+        Function for retrieving a single trajectory segment. Must take a
+        :class:`Walker` instance as its first argument and accept a boolean
+        keyword argument `include_initpoint`. The function should return a
+        sequence (e.g., a list or ndarray) representing the trajectory of
+        the walker. If `include_initpoint` is True, the trajectory segment
+        should include its initial point. Otherwise, the trajectory segment
+        should exclude its initial point.
+    fconcat : callable, optional
+        Function for concatenating trajectory segments. Must take a sequence
+        of trajectory segments as input and return their concatenation. The
+        default concatenation function is :func:`concatenate`.
+
+    """
+
+    def __init__(self, fget=None, *, fconcat=None):
+        if fget is None:
+            return functools.partial(self.__init__, fconcat=fconcat)
+
+        if 'include_initpoint' not in inspect.signature(fget).parameters:
+            raise ValueError("'fget' must accept a parameter 'include_initpoint'")
+
+        self._fget = fget
+        self.fconcat = fconcat
+
+        self._segment_collector = SegmentCollector(self)
+
+    @property
+    def segment_collector(self):
+        """SegmentCollector: Segment retrieval manager."""
+        return self._segment_collector
+
+    @property
+    def fget(self):
+        """callable: Function for getting trajectory segments."""
+        return self._fget
+
+    @property
+    def fconcat(self):
+        """callable: Function for concatenating trajectory segments."""
+        return self._fconcat
+
+    @fconcat.setter
+    def fconcat(self, value):
+        if value is None:
+            value = concatenate
+        elif not isinstance(value, Callable):
+            raise TypeError("'fconcat' must be a callable object")
+        self._fconcat = value
+
+    def __call__(self, obj, include_initpoint=True, **kwargs):
+        if isinstance(obj, Walker):
+            value = self.fget(obj, include_initpoint=include_initpoint, **kwargs)
+            self._validate_segment(value)
+            return value
+        if isinstance(obj, Trace):
+            initpoint_mask = np.full(len(obj), False)
+            initpoint_mask[0] = include_initpoint
+            segments = self.segment_collector.get_segments(obj, initpoint_mask, **kwargs)
+            return self.fconcat(segments)
+        raise TypeError('argument must be a Walker or Trace instance')
+
+    def _validate_segment(self, value):
+        if not hasattr(value, '__getitem__'):
+            msg = f"{type(value).__name__!r} object can't be concatenated"
+            raise TypeError(msg)
+
+
+class SegmentCollector:
+    """An object that manages the retrieval of trajectory segments.
+
+    Parameters
+    ----------
+    trajectory : Trajectory
+        The trajectory to which the segment collector is attached.
+    use_threads : bool, default False
+        Whether to use a pool of threads to retrieve trajectory segments
+        asynchronously. Setting this parameter to True may be may be
+        useful when segment retrieval is an I/O bound task.
+    max_workers : int, optional
+        Maximum number of threads to use. The default value is specified in the
+        `ThreadPoolExecutor <https://docs.python.org/3/library/concurrent.futures.html#concurrent.futures.ThreadPoolExecutor>`_
+        documentation.
+    show_progress : bool, default False
+        Whether to show a progress bar when retrieving multiple segments.
+
+    """
+
+    def __init__(self, trajectory, use_threads=False, max_workers=None, show_progress=False):
+        self.trajectory = trajectory
+        self.use_threads = use_threads
+        self.max_workers = max_workers
+        self.show_progress = show_progress
+
+    @property
+    def trajectory(self):
+        return self._trajectory
+
+    @trajectory.setter
+    def trajectory(self, value):
+        if not isinstance(value, Trajectory):
+            msg = f'trajectory must be an instance of {Trajectory}'
+            raise TypeError(msg)
+        self._trajectory = value
+
+    @property
+    def use_threads(self):
+        return self._use_threads
+
+    @use_threads.setter
+    def use_threads(self, value):
+        if not isinstance(value, bool):
+            raise TypeError('use_threads must be True or False')
+        self._use_threads = value
+
+    @property
+    def max_workers(self):
+        return self._max_workers
+
+    @max_workers.setter
+    def max_workers(self, value):
+        if value is None:
+            self._max_workers = None
+            return
+        if value <= 0:
+            raise ValueError('max_workers must be greater than 0')
+        self._max_workers = value
+
+    @property
+    def show_progress(self):
+        return self._show_progress
+
+    @show_progress.setter
+    def show_progress(self, value):
+        if not isinstance(value, bool):
+            raise ValueError('show_progress must be True or False')
+        self._show_progress = value
+
+    def get_segments(self, walkers, initpoint_mask=None, **kwargs):
+        """Retrieve the trajectories of multiple walkers.
+
+        Parameters
+        ----------
+        walkers : sequence of Walker
+            The walkers for which to retrieve trajectories.
+        initpoint_mask : sequence of bool, optional
+            A Boolean mask indicating whether each trajectory segment should
+            include (True) or exclude (False) its initial point. Default is
+            all True.
+
+        Returns
+        -------
+        list of sequences
+            The trajectory of each walker.
+
+        """
+        if initpoint_mask is None:
+            initpoint_mask = np.full(len(walkers), True)
+        else:
+            initpoint_mask = np.asarray(initpoint_mask, dtype=bool)
+
+        get_segment = functools.partial(self.trajectory, **kwargs)
+
+        tqdm_kwargs = dict(
+            desc='Retrieving segments',
+            disable=(not self.show_progress),
+            position=0,
+            total=len(walkers),
+        )
+
+        if self.use_threads:
+            with cf.ThreadPoolExecutor(self.max_workers) as executor:
+                future_to_key = {
+                    executor.submit(get_segment, walker, include_initpoint=i): key
+                    for key, (walker, i) in enumerate(zip(walkers, initpoint_mask))
+                }
+                futures = list(tqdm(cf.as_completed(future_to_key), **tqdm_kwargs))
+                futures.sort(key=future_to_key.get)
+                segments = (future.result() for future in futures)
+        else:
+            it = (get_segment(walker, include_initpoint=i) for walker, i in zip(walkers, initpoint_mask))
+            segments = tqdm(it, **tqdm_kwargs)
+
+        return list(segments)
+
+
+class BasicMDTrajectory(Trajectory):
+    """Trajectory reader for MD trajectories stored as in the
+    `Basic Tutorial <https://github.com/westpa/westpa_tutorials/tree/main/basic_nacl>`_.
+
+    Parameters
+    ----------
+    top : str or mdtraj.Topology, default 'bstate.pdb'
+    traj_ext : str, default '.dcd'
+    state_ext : str, default '.xml'
+    sim_root : str, default '.'
+
+    """
+
+    def __init__(self, top='bstate.pdb', traj_ext='.dcd', state_ext='.xml', sim_root='.'):
+        self.top = top
+        self.traj_ext = traj_ext
+        self.state_ext = state_ext
+        self.sim_root = sim_root
+
+        def fget(walker, include_initpoint=True, atom_indices=None, sim_root=None):
+            sim_root = sim_root or self.sim_root
+
+            if isinstance(self.top, str):
+                top = os.path.join(sim_root, 'common_files', self.top)
+            else:
+                top = self.top
+
+            path = os.path.join(
+                sim_root,
+                'traj_segs',
+                format(walker.iteration.number, '06d'),
+                format(walker.index, '06d'),
+                'seg' + self.traj_ext,
+            )
+            if top is not None:
+                traj = mdtraj.load(path, top=top)
+            else:
+                traj = mdtraj.load(path)
+
+            if include_initpoint:
+                parent = walker.parent
+
+                if isinstance(parent, InitialState):
+                    if parent.istate_type == InitialState.ISTATE_TYPE_BASIS:
+                        path = os.path.join(
+                            sim_root,
+                            'bstates',
+                            parent.basis_state.auxref,
+                        )
+                    else:
+                        path = os.path.join(
+                            sim_root,
+                            'istates',
+                            str(parent.iter_created),
+                            str(parent.state_id) + self.state_ext,
+                        )
+                else:
+                    path = os.path.join(
+                        sim_root,
+                        'traj_segs',
+                        format(walker.iteration.number - 1, '06d'),
+                        format(parent.index, '06d'),
+                        'seg' + self.state_ext,
+                    )
+
+                frame = mdtraj.load(path, top=traj.top)
+                traj = frame.join(traj, check_topology=False)
+
+            if atom_indices is not None:
+                traj.atom_slice(atom_indices, inplace=True)
+
+            return traj
+
+        super().__init__(fget)
+
+        self.segment_collector.use_threads = True
+        self.segment_collector.show_progress = True
+
+
+class HDF5MDTrajectory(Trajectory):
+    """Trajectory reader for MD trajectories stored by the HDF5 framework."""
+
+    def __init__(self):
+        def fget(walker, include_initpoint=True, atom_indices=None):
+            iteration = walker.iteration
+
+            try:
+                link = iteration.h5group['trajectories']
+            except KeyError:
+                msg = 'the HDF5 framework does not appear to have been used to store trajectories for this run'
+                raise ValueError(msg)
+
+            with WESTIterationFile(link.file.filename) as traj_file:
+                traj = traj_file.read_as_traj(
+                    iteration=iteration.number,
+                    segment=walker.index,
+                    atom_indices=atom_indices,
+                )
+
+            if include_initpoint:
+                parent = walker.parent
+
+                if isinstance(parent, InitialState):
+                    if parent.istate_type == InitialState.ISTATE_TYPE_BASIS:
+                        link = walker.run.h5file.get_iter_group(0)['trajectories']
+                        with WESTIterationFile(link.file.filename) as traj_file:
+                            frame = traj_file.read_as_traj(
+                                iteration=0,
+                                segment=parent.basis_state_id,
+                                atom_indices=atom_indices,
+                            )
+                    elif parent.istate_type == InitialState.ISTATE_TYPE_GENERATED:
+                        link = walker.run.h5file.get_iter_group(0)['trajectories']
+                        istate_iter = -int(parent.iter_created)  # the conversion to int is because iter_created is uint
+                        with WESTIterationFile(link.file.filename) as traj_file:
+                            frame = traj_file.read_as_traj(
+                                iteration=istate_iter,
+                                segment=parent.state_id,
+                                atom_indices=atom_indices,
+                            )
+                    else:
+                        raise ValueError('unsupported initial state type: %d' % parent.istate_type)
+                else:
+                    frame = fget(parent, include_initpoint=False, atom_indices=atom_indices)[-1]
+                traj = frame.join(traj, check_topology=False)
+
+            return traj
+
+        super().__init__(fget)
+
+        self.segment_collector.use_threads = False
+        self.segment_collector.show_progress = True
+
+
+def concatenate(segments):
+    """Return the concatenation of a sequence of trajectory segments.
+
+    Parameters
+    ----------
+    segments : sequence of sequences
+        A sequence of trajectory segments.
+
+    Returns
+    -------
+    sequence
+        The concatenation of `segments`.
+
+    """
+    if isinstance(segments[0], np.ndarray):
+        return np.concatenate(segments)
+    if isinstance(segments[0], mdtraj.Trajectory):
+        return segments[0].join(segments[1:], check_topology=False)
+    return functools.reduce(operator.concat, segments)

westpa/cli/core/w_fork.py

@@ -0,0 +1,152 @@
+import argparse
+import logging
+
+import numpy as np
+
+import westpa
+from westpa.core.segment import Segment
+from westpa.core.states import InitialState
+from westpa.core.data_manager import n_iter_dtype, seg_id_dtype
+
+log = logging.getLogger('w_fork')
+
+
+def entry_point():
+    parser = argparse.ArgumentParser(
+        'w_fork',
+        description='''\
+    Prepare a new weighted ensemble simulation from an existing one at a particular
+    point. A new HDF5 file is generated. In the case of executable propagation,
+    it is the user's responsibility to prepare the new simulation directory
+    appropriately, particularly making the old simulation's restart data from the
+    appropriate iteration available as the new simulations initial state data; a
+    mapping of old simulation segment to new simulation initial states is
+    created, both in the new HDF5 file and as a flat text file, to aid in this.
+    Target states and basis states for the new simulation are taken from those
+    in the original simulation.
+    ''',
+    )
+
+    westpa.rc.add_args(parser)
+    parser.add_argument(
+        '-i',
+        '--input',
+        dest='input_h5file',
+        help='''Create simulation from the given INPUT_H5FILE (default: read from
+                                configuration file.''',
+    )
+    parser.add_argument(
+        '-I',
+        '--iteration',
+        dest='n_iter',
+        type=int,
+        help='''Take initial distribution for new simulation from iteration N_ITER
+                                (default: last complete iteration).''',
+    )
+    parser.add_argument(
+        '-o',
+        '--output',
+        dest='output_h5file',
+        default='forked.h5',
+        help='''Save new simulation HDF5 file as OUTPUT (default: %(default)s).''',
+    )
+    parser.add_argument(
+        '--istate-map',
+        default='istate_map.txt',
+        help='''Write text file describing mapping of existing segments to new initial
+                                states in ISTATE_MAP (default: %(default)s).''',
+    )
+    parser.add_argument('--no-headers', action='store_true', help='''Do not write header to ISTATE_MAP''')
+    args = parser.parse_args()
+    westpa.rc.process_args(args)
+
+    # Open old HDF5 file
+    dm_old = westpa.rc.new_data_manager()
+    if args.input_h5file:
+        dm_old.we_h5filename = args.input_h5file
+    dm_old.open_backing(mode='r')
+
+    # Get iteration if necessary
+    n_iter = args.n_iter or dm_old.current_iteration - 1
+
+    # Create and open new HDF5 file
+    dm_new = westpa.rc.new_data_manager()
+    dm_new.we_h5filename = args.output_h5file
+    dm_new.prepare_backing()
+    dm_new.open_backing()
+
+    # Copy target states
+    target_states = dm_old.get_target_states(n_iter)
+    dm_new.save_target_states(target_states, n_iter)
+
+    # Copy basis states
+    basis_states = dm_old.get_basis_states(n_iter)
+    dm_new.create_ibstate_group(basis_states, n_iter=1)
+
+    # Transform old segments into initial states and new segments
+    # We produce one initial state and one corresponding
+    # new segment for each old segment. Further adjustment
+    # can be accomplished by using w_binning.
+    old_iter_group = dm_old.get_iter_group(n_iter)
+    old_index = old_iter_group['seg_index'][...]
+    old_pcoord_ds = old_iter_group['pcoord']
+    n_segments = old_pcoord_ds.shape[0]
+    pcoord_len = old_pcoord_ds.shape[1]
+    pcoord_ndim = old_pcoord_ds.shape[2]
+    old_final_pcoords = old_pcoord_ds[:, pcoord_len - 1, :]
+
+    istates = dm_new.create_initial_states(n_segments, n_iter=1)
+    segments = []
+    state_map_dtype = np.dtype([('old_n_iter', n_iter_dtype), ('old_seg_id', seg_id_dtype), ('new_istate_id', seg_id_dtype)])
+    state_map = np.empty((n_segments,), dtype=state_map_dtype)
+    state_map['old_n_iter'] = n_iter
+
+    for iseg, (index_row, pcoord) in enumerate(zip(old_index, old_final_pcoords)):
+        istate = istates[iseg]
+        istate.iter_created = 0
+        istate.iter_used = 1
+        istate.istate_type = InitialState.ISTATE_TYPE_RESTART
+        istate.istate_status = InitialState.ISTATE_STATUS_PREPARED
+        istate.pcoord = pcoord
+
+        segment = Segment(
+            n_iter=1,
+            seg_id=iseg,
+            weight=index_row['weight'],
+            parent_id=-(istate.state_id + 1),
+            wtg_parent_ids=[-(istate.state_id + 1)],
+            status=Segment.SEG_STATUS_PREPARED,
+        )
+        segment.pcoord = np.zeros((pcoord_len, pcoord_ndim), dtype=pcoord.dtype)
+        segment.pcoord[0] = pcoord
+        segments.append(segment)
+        state_map[iseg]['old_seg_id'] = iseg
+        state_map[iseg]['new_istate_id'] = istate.state_id
+
+    dm_new.update_initial_states(istates, n_iter=0)
+    dm_new.prepare_iteration(n_iter=1, segments=segments)
+
+    # Update current iteration and close both files
+    dm_new.current_iteration = 1
+    dm_new.close_backing()
+    dm_old.close_backing()
+
+    # Write state map
+    istate_map_file = open(args.istate_map, 'wt')
+    if not args.no_headers:
+        istate_map_file.write('# mapping from previous segment IDs to new initial states\n')
+        istate_map_file.write('# generated by w_fork\n')
+        istate_map_file.write('# column 0: old simulation n_iter\n')
+        istate_map_file.write('# column 1: old simulation seg_id\n')
+        istate_map_file.write('# column 2: new simulation initial state ID\n')
+
+    for row in state_map:
+        istate_map_file.write(
+            '{old_n_iter:20d}    {old_seg_id:20d}    {new_istate_id:20d}\n'.format(
+                old_n_iter=int(row['old_n_iter']), old_seg_id=int(row['old_seg_id']), new_istate_id=int(row['new_istate_id'])
+            )
+        )
+
+
+if __name__ == '__main__':
+    entry_point()