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

westpa/cli/core/w_init.py

@@ -0,0 +1,230 @@
+import argparse
+import io
+import logging
+import sys
+
+import numpy as np
+
+import westpa
+
+from westpa.core.states import BasisState, TargetState
+
+import westpa.work_managers as work_managers
+from westpa.work_managers import make_work_manager
+
+
+log = logging.getLogger('w_init')
+EPS = np.finfo(np.float64).eps
+
+
+def entry_point():
+    parser = argparse.ArgumentParser(
+        'w_init',
+        description='''\
+    Initialize a new WEST simulation, creating the WEST HDF5 file and preparing the first
+    iteration's segments.
+
+    Initial states are generated from one or more "basis states" which are specified either in a file specified with --bstates-from, or
+    by one or more "--bstate" arguments. If neither --bstates-from nor at least one "--bstate" argument is provided, then a default
+    basis state of probability one identified by the state ID zero and label "basis" will be created (a warning will be printed in this
+    case, to remind you of this behavior, in case it is not what you wanted).
+
+    Target states for (non-equilibrium) steady-state simulations are specified either in a file specified with --tstates-from, or
+    by one or more --tstate arguments. If neither --tstates-from nor at least one --tstate argument is provided, then an equilibrium
+    simulation (without any sinks) will be performed.
+    ''',
+    )
+    westpa.rc.add_args(parser)
+    parser.add_argument('--force', dest='force', action='store_true', help='Overwrite any existing simulation data')
+    parser.add_argument(
+        '--bstate-file',
+        '--bstates-from',
+        metavar='BSTATE_FILE',
+        help='Read basis state names, probabilities, and (optionally) data references from BSTATE_FILE.',
+    )
+    parser.add_argument(
+        '--bstate',
+        action='append',
+        dest='bstates',
+        help='''Add the given basis state (specified as a string 'label,probability[,auxref]')
+                        to the list of basis states (after those specified in --bstates-from, if any). This argument
+                        may be specified more than once, in which case the given states are appended in the order
+                        they are given on the command line.''',
+    )
+    parser.add_argument(
+        '--tstate-file',
+        '--tstates-from',
+        metavar='TSTATE_FILE',
+        help='''Read target state names and representative progress coordinates from TSTATE_FILE. WESTPA uses the
+                        representative progress coordinate of a target state and converts the **entire** bin
+                        containing that progress coordinate into a recycling sink.''',
+    )
+    parser.add_argument(
+        '--tstate',
+        action='append',
+        dest='tstates',
+        help='''Add the given target state (specified as a string 'label,pcoord0[,pcoord1[,...]]') to the
+                        list of target states (after those specified in the file given by --tstates-from, if any).
+                        This argument may be specified more than once, in which case the given states are appended
+                        in the order they appear on the command line.''',
+    )
+
+    parser.add_argument(
+        '--sstate-file',
+        '--sstates-from',
+        metavar='SSTATE_FILE',
+        help='Read start state names, probabilities, and (optionally) data references from SSTATE_FILE.',
+    )
+    parser.add_argument(
+        '--sstate',
+        action='append',
+        dest='sstates',
+        help='''Add the given start state (specified as a string 'label,probability[,auxref]')
+                        to the list of start states (after those specified in --sstates-from, if any). This argument
+                        may be specified more than once, in which case the given states are appended in the order
+                        they are given on the command line.''',
+    )
+
+    parser.add_argument(
+        '--segs-per-state',
+        type=int,
+        metavar='N',
+        default=1,
+        help='''Initialize N segments from each basis state (default: %(default)s).''',
+    )
+    parser.add_argument(
+        '--no-we',
+        '--shotgun',
+        dest='shotgun',
+        action='store_true',
+        help='''Do not run the weighted ensemble bin/split/merge algorithm on newly-created segments.''',
+    )
+
+    # TODO: Does this belong here or not? I like that it's parsing arguments, which is the purpose of entry_point.
+    #   I don't necessarily like that it's setting state across different parts of the program.
+
+    work_managers.environment.add_wm_args(parser)
+    args = parser.parse_args()
+    westpa.rc.process_args(args)
+    work_managers.environment.process_wm_args(args)
+
+    initialize(
+        args.tstates,
+        args.tstate_file,
+        args.bstates,
+        args.bstate_file,
+        args.sstates,
+        args.sstate_file,
+        args.segs_per_state,
+        args.shotgun,
+    )
+
+
+def initialize(tstates, tstate_file, bstates, bstate_file, sstates=None, sstate_file=None, segs_per_state=1, shotgun=False):
+    """
+    Initialize a WESTPA simulation.
+
+    tstates : list of str
+
+    tstate_file : str
+
+    bstates : list of str
+
+    bstate_file : str
+
+    sstates : list of str
+
+    sstate_file : str
+
+    segs_per_state : int
+
+    shotgun : bool
+    """
+
+    westpa.rc.work_manager = work_manager = make_work_manager()
+
+    system = westpa.rc.get_system_driver()
+    sim_manager = westpa.rc.get_sim_manager()
+
+    data_manager = westpa.rc.get_data_manager()
+
+    data_manager.system = system
+
+    with work_manager:
+        if work_manager.is_master:
+            # Process target states
+            target_states = []
+            if tstate_file:
+                target_states.extend(TargetState.states_from_file(tstate_file, system.pcoord_dtype))
+            if tstates:
+                tstates_strio = io.StringIO('\n'.join(tstates).replace(',', ' '))
+                target_states.extend(TargetState.states_from_file(tstates_strio, system.pcoord_dtype))
+                del tstates_strio
+
+            # Process basis states
+            basis_states = []
+            if bstate_file:
+                basis_states.extend(BasisState.states_from_file(bstate_file))
+            if bstates:
+                for bstate_str in bstates:
+                    fields = bstate_str.split(',')
+                    label = fields[0]
+                    probability = float(fields[1])
+                    try:
+                        auxref = fields[2]
+                    except IndexError:
+                        auxref = None
+                    basis_states.append(BasisState(label=label, probability=probability, auxref=auxref))
+
+            # Process the list of start states, creating a BasisState from each
+            start_states = []
+            if sstate_file:
+                start_states.extend(BasisState.states_from_file(sstate_file))
+            if sstates:
+                for sstate_str in sstates:
+                    fields = sstate_str.split(',')
+                    label = fields[0]
+                    probability = float(fields[1])
+                    try:
+                        auxref = fields[2]
+                    except IndexError:
+                        auxref = None
+                    start_states.append(BasisState(label=label, probability=probability, auxref=auxref))
+
+            if not basis_states:
+                log.error('At least one basis state is required')
+                sys.exit(3)
+
+            # Check that the total probability of basis states adds to one
+            bstate_prob, sstate_prob = (
+                sum(bstate.probability for bstate in basis_states),
+                sum(sstate.probability for sstate in start_states),
+            )
+            # tprob = sum(bstate.probability for bstate in basis_states)
+            tprob = bstate_prob + sstate_prob
+            if abs(1.0 - tprob) > len(basis_states) * EPS:
+                pscale = 1 / tprob
+                log.warning(
+                    'Basis state probabilities do not add to unity (basis: {:.2f}, start states: {:.2f}); rescaling by {:g}. If using start states, some rescaling is normal.'.format(
+                        bstate_prob, sstate_prob, pscale
+                    )
+                )
+                for bstate in basis_states:
+                    bstate.probability *= pscale
+                for sstate in start_states:
+                    sstate.probability *= pscale
+
+            # Prepare simulation
+            sim_manager.initialize_simulation(
+                basis_states=basis_states,
+                target_states=target_states,
+                start_states=start_states,
+                segs_per_state=segs_per_state,
+                suppress_we=shotgun,
+            )
+        else:
+            work_manager.run()
+
+
+if __name__ == '__main__':
+    entry_point()

westpa/cli/core/w_run.py

@@ -0,0 +1,77 @@
+import argparse
+import logging
+import traceback
+
+import westpa
+import westpa.work_managers as work_managers
+from westpa.work_managers import make_work_manager
+
+
+log = logging.getLogger('w_run')
+
+
+def entry_point():
+    parser = argparse.ArgumentParser('w_run', 'start/continue a WEST simulation')
+    westpa.rc.add_args(parser)
+    parser.add_argument(
+        '--oneseg',
+        dest='only_one_segment',
+        action='store_true',
+        help='only propagate one segment (useful for debugging propagators)',
+    )
+
+    work_managers.environment.add_wm_args(parser)
+
+    args = parser.parse_args()
+    westpa.rc.process_args(args)
+    work_managers.environment.process_wm_args(args)
+
+    run_simulation()
+
+
+def run_simulation():
+    work_manager = westpa.rc.work_manager = make_work_manager()
+
+    # Load the sim manager and other drivers
+    sim_manager = westpa.rc.get_sim_manager()
+    system = westpa.rc.get_system_driver()
+    data_manager = westpa.rc.get_data_manager()
+    we_driver = westpa.rc.get_we_driver()
+    propagator = westpa.rc.get_propagator()
+
+    propagator.system = system
+    data_manager.system = system
+    we_driver.system = system
+
+    sim_manager.data_manager = data_manager
+    sim_manager.system = system
+    sim_manager.propagator = propagator
+    sim_manager.we_driver = we_driver
+
+    with work_manager:
+        if work_manager.is_master:
+            work_manager.install_sigint_handler()
+            sim_manager.load_plugins()
+
+            log.debug('preparing run')
+            sim_manager.prepare_run()
+
+            try:
+                log.debug('beginning run')
+                sim_manager.run()
+
+                log.debug('finalizing run')
+                sim_manager.finalize_run()
+            except KeyboardInterrupt:
+                westpa.rc.pstatus('interrupted; shutting down')
+            except Exception as e:
+                westpa.rc.pstatus('exception caught; shutting down')
+                if str(e) != '':
+                    log.error(f'error message: {e}')
+                log.error(traceback.format_exc())
+        else:
+            work_manager.run()
+
+
+if __name__ == '__main__':
+    entry_point()

westpa/cli/core/w_states.py

@@ -0,0 +1,212 @@
+import argparse
+import io
+import logging
+import sys
+
+import numpy as np
+
+import westpa.work_managers as work_managers
+from westpa.work_managers import make_work_manager
+
+import westpa
+from westpa.core.segment import Segment
+from westpa.core.states import BasisState, TargetState
+
+
+log = logging.getLogger('w_states')
+EPS = np.finfo(np.float64).eps
+
+
+def entry_point():
+    parser = argparse.ArgumentParser(
+        'w_states',
+        description='''\
+    Display or manipulate basis (initial) or target (recycling) states for a WEST simulation.  By default, states are
+    displayed (or dumped to files).  If ``--replace`` is specified, all basis/target states are replaced for the
+    next iteration.  If ``--append`` is specified, the given target state(s) are appended to the list for the
+    next iteration.
+
+    Appending basis states is not permitted, as this would require renormalizing basis state
+    probabilities in ways that may be error-prone. Instead, use ``w_states --show --bstate-file=bstates.txt``
+    and then edit the resulting ``bstates.txt`` file to include the new desired basis states, then use
+    ``w_states --replace --bstate-file=bstates.txt`` to update the WEST HDF5 file appropriately.
+    ''',
+    )
+    westpa.rc.add_args(parser)
+    smgroup = parser.add_argument_group('modes of operation')
+    mode_group = smgroup.add_mutually_exclusive_group()
+    mode_group.add_argument(
+        '--show', dest='mode', action='store_const', const='show', help='Display current basis/target states (or dump to files).'
+    )
+    mode_group.add_argument(
+        '--append',
+        dest='mode',
+        action='store_const',
+        const='append',
+        help='Append the given basis/target states to those currently in use.',
+    )
+    mode_group.add_argument(
+        '--replace',
+        dest='mode',
+        action='store_const',
+        const='replace',
+        help='Replace current basis/target states with those specified.',
+    )
+    parser.add_argument(
+        '--bstate-file',
+        metavar='BSTATE_FILE',
+        help='''Read (--append/--replace) or write (--show) basis state names, probabilities,
+                        and data references from/to BSTATE_FILE.''',
+    )
+    parser.add_argument(
+        '--bstate',
+        action='append',
+        dest='bstates',
+        help='''Add the given basis state (specified as a string 'label,probability[,auxref]')
+                        to the list of basis states (after those specified in --bstate-file, if any). This argument
+                        may be specified more than once, in which case the given states are appended in the order
+                        they are given on the command line.''',
+    )
+    parser.add_argument(
+        '--tstate-file',
+        metavar='TSTATE_FILE',
+        help='''Read (--append/--replace) or write (--show) target state names
+                        and representative progress coordinates from/to TSTATE_FILE''',
+    )
+    parser.add_argument(
+        '--tstate',
+        action='append',
+        dest='tstates',
+        help='''Add the given target state (specified as a string 'label,pcoord0[,pcoord1[,...]]') to the
+                        list of target states (after those specified in the file given by --tstates-from, if any).
+                        This argument may be specified more than once, in which case the given states are appended
+                        in the order they appear on the command line.''',
+    )
+    parser.set_defaults(mode='show')
+
+    work_managers.environment.add_wm_args(parser)
+    args = parser.parse_args()
+    westpa.rc.process_args(args)
+    work_managers.environment.process_wm_args(args)
+
+    # Need to have something to pass to initialize
+    if not hasattr(args, 'bstates'):
+        args.bstates = None
+    if not hasattr(args, 'tstates'):
+        args.tstates = None
+    if not hasattr(args, 'tstate_file'):
+        args.tstate_file = None
+
+    initialize(args.mode, args.bstates, args.bstate_file, args.tstates, args.tstate_file)
+
+
+# TODO: This would benefit from a refactor to set default args to None, and replace some of those "if <argument>" clauses
+def initialize(mode, bstates, _bstate_file, tstates, _tstate_file):
+    work_manager = make_work_manager()
+
+    system = westpa.rc.get_system_driver()
+
+    with work_manager:
+        if work_manager.is_master:
+            data_manager = westpa.rc.get_data_manager()
+            data_manager.open_backing(mode='a')
+            sim_manager = westpa.rc.get_sim_manager()
+            n_iter = data_manager.current_iteration
+
+            assert mode in ('show', 'replace', 'append')
+            if mode == 'show':
+                basis_states = data_manager.get_basis_states(n_iter)
+                if basis_states:
+                    bstate_file = sys.stdout if not _bstate_file else open(_bstate_file, 'wt')
+                    bstate_file.write('# Basis states for iteration {:d}\n'.format(n_iter))
+                    BasisState.states_to_file(basis_states, bstate_file)
+
+                target_states = data_manager.get_target_states(n_iter)
+                if target_states:
+                    tstate_file = sys.stdout if not _tstate_file else open(_tstate_file, 'wt')
+                    tstate_file.write('# Target states for iteration {:d}\n'.format(n_iter))
+                    TargetState.states_to_file(target_states, tstate_file)
+
+            elif mode == 'replace':
+                seg_index = data_manager.get_seg_index(n_iter)
+                if (seg_index['status'] == Segment.SEG_STATUS_COMPLETE).any():
+                    print('Iteration {:d} has completed segments; applying new states to iteration {:d}'.format(n_iter, n_iter + 1))
+                    n_iter += 1
+
+                basis_states = []
+                if _bstate_file:
+                    basis_states.extend(BasisState.states_from_file(_bstate_file))
+                if bstates:
+                    for bstate_str in bstates:
+                        fields = bstate_str.split(',')
+                        label = fields[0]
+                        probability = float(fields[1])
+                        try:
+                            auxref = fields[2]
+                        except IndexError:
+                            auxref = None
+                        basis_states.append(BasisState(label=label, probability=probability, auxref=auxref))
+
+                if basis_states:
+                    # Check that the total probability of basis states adds to one
+                    tprob = sum(bstate.probability for bstate in basis_states)
+                    if abs(1.0 - tprob) > len(basis_states) * EPS:
+                        pscale = 1 / tprob
+                        log.warning('Basis state probabilities do not add to unity; rescaling by {:g}'.format(pscale))
+                        for bstate in basis_states:
+                            bstate.probability *= pscale
+
+                    # Assign progress coordinates to basis states
+                    sim_manager.get_bstate_pcoords(basis_states)
+                    data_manager.create_ibstate_group(basis_states, n_iter)
+                    sim_manager.report_basis_states(basis_states)
+
+                # Now handle target states
+                target_states = []
+                if _tstate_file:
+                    target_states.extend(TargetState.states_from_file(_tstate_file, system.pcoord_dtype))
+                if tstates:
+                    tstates_strio = io.StringIO('\n'.join(tstates).replace(',', ' '))
+                    target_states.extend(TargetState.states_from_file(tstates_strio, system.pcoord_dtype))
+                    del tstates_strio
+
+                if not target_states:
+                    westpa.rc.pstatus('No target states specified.')
+                else:
+                    data_manager.save_target_states(target_states, n_iter)
+                    sim_manager.report_target_states(target_states)
+
+                data_manager.update_iter_group_links(n_iter)
+
+            else:  # args.mode == 'append'
+                if _bstate_file or bstates:
+                    sys.stderr.write('refusing to append basis states; use --show followed by --replace instead\n')
+                    sys.exit(2)
+
+                target_states = data_manager.get_target_states(n_iter)
+
+                seg_index = data_manager.get_seg_index(n_iter)
+                if (seg_index['status'] == Segment.SEG_STATUS_COMPLETE).any():
+                    print('Iteration {:d} has completed segments; applying new states to iteration {:d}'.format(n_iter, n_iter + 1))
+                    n_iter += 1
+
+                if _tstate_file:
+                    target_states.extend(TargetState.states_from_file(_tstate_file, system.pcoord_dtype))
+                if tstates:
+                    tstates_strio = io.StringIO('\n'.join(tstates).replace(',', ' '))
+                    target_states.extend(TargetState.states_from_file(tstates_strio, system.pcoord_dtype))
+                    del tstates_strio
+
+                if not target_states:
+                    westpa.rc.pstatus('No target states specified.')
+                else:
+                    data_manager.save_target_states(target_states, n_iter)
+                    sim_manager.report_target_states(target_states)
+
+                data_manager.update_iter_group_links(n_iter)
+        else:
+            work_manager.run()
+
+
+if __name__ == '__main__':
+    entry_point()

westpa/cli/core/w_succ.py

@@ -0,0 +1,99 @@
+import argparse
+import sys
+
+import numpy as np
+
+import westpa
+
+from westpa.core.segment import Segment
+from westpa.oldtools.aframe import WESTAnalysisTool, WESTDataReaderMixin, CommonOutputMixin
+
+import logging
+
+log = logging.getLogger('w_succ')
+
+
+class WSucc(CommonOutputMixin, WESTDataReaderMixin, WESTAnalysisTool):
+    def __init__(self):
+        super().__init__()
+        self.include_args['CommonOutputMixin']['print_bin_labels'] = False
+        self.output_file = sys.stdout
+
+    def find_successful_trajs(self):
+        pcoord_formats = {
+            'u8': '%20d',
+            'i8': '%20d',
+            'u4': '%10d',
+            'i4': '%11d',
+            'u2': '%5d',
+            'i2': '%6d',
+            'f4': '%14.7g',
+            'f8': '%23.15g',
+        }
+
+        if not self.output_suppress_headers:
+            self.output_file.write(
+                '''\
+# successful (recycled) segments
+# column 0:    iteration
+# column 1:    seg_id
+# column 2:    weight
+# column>2:    final progress coordinate value
+'''
+            )
+        for n_iter in range(1, self.data_manager.current_iteration):
+            seg_index = self.get_seg_index(n_iter)
+            all_seg_ids = np.arange(len(seg_index), dtype=np.int_)
+            recycled_seg_ids = all_seg_ids[seg_index[:]['endpoint_type'] == Segment.SEG_ENDPOINT_RECYCLED]
+
+            if len(recycled_seg_ids) == 0:
+                # Attemping to retrieve a 0-length selection from HDF5 (the pcoords below) fails
+                continue
+
+            pcoord_ds = self.get_pcoord_dataset(n_iter)
+            pcoord_len = pcoord_ds.shape[1]
+            pcoord_ndim = pcoord_ds.shape[2]
+            final_pcoords = self.get_pcoord_dataset(n_iter)[recycled_seg_ids, pcoord_len - 1, :]
+            # The above HDF5 selection always returns a vector; we want a 2-d array
+            final_pcoords.shape = (len(recycled_seg_ids), pcoord_ndim)
+
+            for ipc, seg_id in enumerate(recycled_seg_ids):
+                self.output_file.write('%8d    %8d    %20.14g' % (n_iter, seg_id, seg_index[seg_id]['weight']))
+                fields = ['']
+                for field in final_pcoords[ipc]:
+                    fields.append(pcoord_formats.get(field.dtype.str[1:], '%s') % field)
+                self.output_file.write('    '.join(fields))
+                self.output_file.write('\n')
+
+
+def entry_point():
+    wsucc = WSucc()
+
+    parser = argparse.ArgumentParser(
+        'w_succ',
+        description='''\
+    List segments which successfully reach a target state''',
+    )
+
+    westpa.rc.add_args(parser)
+    wsucc.add_args(parser)
+
+    parser.add_argument(
+        '-o',
+        '--output',
+        dest='output_file',
+        help='Store output in OUTPUT_FILE (default: write to standard output).',
+        type=argparse.FileType('wt'),
+        default=sys.stdout,
+    )
+
+    args = parser.parse_args()
+    westpa.rc.process_args(args, config_required=False)
+    wsucc.process_args(args)
+    wsucc.output_file = args.output_file
+
+    wsucc.find_successful_trajs()
+
+
+if __name__ == '__main__':
+    entry_point()

westpa/cli/core/w_truncate.py

@@ -0,0 +1,68 @@
+import logging
+from argparse import ArgumentParser
+
+import westpa
+
+
+log = logging.getLogger('w_truncate')
+
+warning_string = '''\
+NOTE: w_truncate only deletes iteration groups from the HDF5 data store.
+It is recommended that any iteration data saved to the file system (e.g. in the
+traj_segs directory) is deleted or moved for the corresponding iterations.
+'''
+
+
+def entry_point():
+    parser = ArgumentParser(
+        'w_truncate',
+        description='''\
+    Remove all iterations after a certain point in a WESTPA simulation.
+    ''',
+        epilog=warning_string,
+    )
+
+    westpa.rc.add_args(parser)
+    parser.add_argument(
+        '-W',
+        '--west-data',
+        dest='we_h5filename',
+        metavar='WEST_H5FILE',
+        help='''Take WEST data from WEST_H5FILE (default: read from the HDF5 file specified in west.cfg).''',
+    )
+    parser.add_argument('-n', '--iter', dest='n_iter', type=int, default=0, help='Truncate this iteration and those following.')
+
+    args = parser.parse_args()
+    westpa.rc.process_args(args, config_required=False)
+    dm = westpa.rc.get_data_manager()
+    if args.we_h5filename:
+        dm.we_h5filename = args.we_h5filename
+
+    dm.open_backing()
+    # max_iter = dm.current_iteration
+
+    if args.n_iter > dm.current_iteration:
+        parser.error(
+            'Provided iteration {} > current iteration {} of the {} HDF5 file. Exiting without doing anything.'.format(
+                args.n_iter, dm.current_iteration, dm.we_h5filename.split('/')[-1]
+            )
+        )
+    else:
+        n_iter = args.n_iter if args.n_iter > 0 else dm.current_iteration
+
+    for i in range(n_iter, dm.current_iteration + 1):
+        dm.del_iter_group(i)
+
+    dm.del_iter_summary(n_iter)
+    dm.current_iteration = n_iter - 1
+
+    westpa.rc.pstatus('simulation data truncated after iteration {}'.format(dm.current_iteration))
+    westpa.rc.pstatus('\n' + warning_string)
+    westpa.rc.pflush()
+
+    dm.flush_backing()
+    dm.close_backing()
+
+
+if __name__ == '__main__':
+    entry_point()