westpa 2022.10__cp312-cp312-macosx_10_9_x86_64.whl

westpa/cli/tools/w_select.py

@@ -0,0 +1,226 @@
+from westpa.tools import WESTParallelTool, WESTDataReader, IterRangeSelection, ProgressIndicatorComponent
+
+import numpy as np
+
+from westpa.core import h5io
+from westpa.core.data_manager import seg_id_dtype, n_iter_dtype, weight_dtype
+from westpa.core.extloader import get_object
+
+
+def _find_matching_segments(west_datafile_name, n_iter, predicate, invert=False):
+    '''Find all segments in iteration ``n_iter`` that match (or do not match, if
+    ``invert`` is true) the given ``predicate``. Returns a sequence of matching
+    seg_ids.'''
+
+    with h5io.WESTPAH5File(west_datafile_name, 'r') as west_datafile:
+        iter_group = west_datafile.get_iter_group(n_iter)
+        nsegs = iter_group['seg_index'].shape[0]
+        matching_ids = set(map(int, predicate(n_iter, iter_group)))
+
+        if invert:
+            matching_ids = set(range(nsegs)) - matching_ids
+
+        matchvec = sorted(np.fromiter(matching_ids, dtype=seg_id_dtype, count=len(matching_ids)))
+        return n_iter, matchvec
+
+
+class WSelectTool(WESTParallelTool):
+    prog = 'w_select'
+    description = '''\
+Select dynamics segments matching various criteria. This requires a
+user-provided prediate function. By default, only matching segments are
+stored. If the -a/--include-ancestors option is given, then matching segments
+and their ancestors will be stored.
+
+
+-----------------------------------------------------------------------------
+Predicate function
+-----------------------------------------------------------------------------
+
+Segments are selected based on a predicate function, which must be callable
+as ``predicate(n_iter, iter_group)`` and return a collection of segment IDs
+matching the predicate in that iteration.
+
+The predicate may be inverted by specifying the -v/--invert command-line
+argument.
+
+
+-----------------------------------------------------------------------------
+Output format
+-----------------------------------------------------------------------------
+
+The output file (-o/--output, by default "select.h5") contains the following
+datasets:
+
+  ``/n_iter`` [iteration]
+    *(Integer)* Iteration numbers for each entry in other datasets.
+
+  ``/n_segs`` [iteration]
+    *(Integer)* Number of segment IDs matching the predicate (or inverted
+    predicate, if -v/--invert is specified) in the given iteration.
+
+  ``/seg_ids`` [iteration][segment]
+    *(Integer)* Matching segments in each iteration. For an iteration
+    ``n_iter``, only the first ``n_iter`` entries are valid. For example,
+    the full list of matching seg_ids in the first stored iteration is
+    ``seg_ids[0][:n_segs[0]]``.
+
+  ``/weights`` [iteration][segment]
+    *(Floating-point)* Weights for each matching segment in ``/seg_ids``.
+
+
+-----------------------------------------------------------------------------
+Command-line arguments
+-----------------------------------------------------------------------------
+'''
+
+    def __init__(self):
+        super().__init__()
+
+        self.data_reader = WESTDataReader()
+        self.iter_range = IterRangeSelection()
+        self.progress = ProgressIndicatorComponent()
+        self.output_file = None
+        self.output_filename = None
+        self.predicate = None
+        self.invert = False
+        self.include_ancestors = False
+
+    def add_args(self, parser):
+        self.data_reader.add_args(parser)
+        self.iter_range.add_args(parser)
+
+        sgroup = parser.add_argument_group('selection options')
+        sgroup.add_argument(
+            '-p',
+            '--predicate-function',
+            metavar='MODULE.FUNCTION',
+            help='''Use the given predicate function to match segments. This function
+                             should take an iteration number and the HDF5 group corresponding to that
+                             iteration and return a sequence of seg_ids matching the predicate, as in
+                             ``match_predicate(n_iter, iter_group)``.''',
+        )
+        sgroup.add_argument('-v', '--invert', dest='invert', action='store_true', help='''Invert the match predicate.''')
+        sgroup.add_argument(
+            '-a', '--include-ancestors', action='store_true', help='''Include ancestors of matched segments in output.'''
+        )
+
+        ogroup = parser.add_argument_group('output options')
+        ogroup.add_argument('-o', '--output', default='select.h5', help='''Write output to OUTPUT (default: %(default)s).''')
+        self.progress.add_args(parser)
+
+    def process_args(self, args):
+        self.progress.process_args(args)
+        self.data_reader.process_args(args)
+        with self.data_reader:
+            self.iter_range.process_args(args)
+
+        predicate = get_object(args.predicate_function, path=['.'])
+        if not callable(predicate):
+            raise TypeError('predicate object {!r} is not callable'.format(predicate))
+        self.predicate = predicate
+        self.invert = bool(args.invert)
+        self.include_ancestors = bool(args.include_ancestors)
+        self.output_filename = args.output
+
+    def go(self):
+        self.data_reader.open('r')
+        output_file = h5io.WESTPAH5File(self.output_filename, mode='w')
+        pi = self.progress.indicator
+
+        iter_start, iter_stop = self.iter_range.iter_start, self.iter_range.iter_stop
+        iter_count = iter_stop - iter_start
+
+        output_file.create_dataset('n_iter', dtype=n_iter_dtype, data=list(range(iter_start, iter_stop)))
+        current_seg_count = 0
+        seg_count_ds = output_file.create_dataset('n_segs', dtype=np.uint, shape=(iter_count,))
+        matching_segs_ds = output_file.create_dataset(
+            'seg_ids',
+            shape=(iter_count, 0),
+            maxshape=(iter_count, None),
+            dtype=seg_id_dtype,
+            chunks=h5io.calc_chunksize((iter_count, 1000000), seg_id_dtype),
+            shuffle=True,
+            compression=9,
+        )
+        weights_ds = output_file.create_dataset(
+            'weights',
+            shape=(iter_count, 0),
+            maxshape=(iter_count, None),
+            dtype=weight_dtype,
+            chunks=h5io.calc_chunksize((iter_count, 1000000), weight_dtype),
+            shuffle=True,
+            compression=9,
+        )
+
+        with pi:
+            pi.new_operation('Finding matching segments', extent=iter_count)
+            #             futures = set()
+            #             for n_iter in xrange(iter_start,iter_stop):
+            #                 futures.add(self.work_manager.submit(_find_matching_segments,
+            #                                                      args=(self.data_reader.we_h5filename,n_iter,self.predicate,self.invert)))
+
+            #             for future in self.work_manager.as_completed(futures):
+            for future in self.work_manager.submit_as_completed(
+                (
+                    (_find_matching_segments, (self.data_reader.we_h5filename, n_iter, self.predicate, self.invert), {})
+                    for n_iter in range(iter_start, iter_stop)
+                ),
+                self.max_queue_len,
+            ):
+                n_iter, matching_ids = future.get_result()
+                n_matches = len(matching_ids)
+
+                if n_matches:
+                    if n_matches > current_seg_count:
+                        current_seg_count = len(matching_ids)
+                        matching_segs_ds.resize((iter_count, n_matches))
+                        weights_ds.resize((iter_count, n_matches))
+                        current_seg_count = n_matches
+
+                    seg_count_ds[n_iter - iter_start] = n_matches
+                    matching_segs_ds[n_iter - iter_start, :n_matches] = matching_ids
+                    weights_ds[n_iter - iter_start, :n_matches] = self.data_reader.get_iter_group(n_iter)['seg_index']['weight'][
+                        sorted(matching_ids)
+                    ]
+                del matching_ids
+                pi.progress += 1
+
+            if self.include_ancestors:
+                pi.new_operation('Tracing ancestors of matching segments', extent=iter_count)
+                from_previous = set()
+                current_seg_count = matching_segs_ds.shape[1]
+                for n_iter in range(iter_stop - 1, iter_start - 1, -1):
+                    iiter = n_iter - iter_start
+                    n_matches = seg_count_ds[iiter]
+                    matching_ids = set(from_previous)
+                    if n_matches:
+                        matching_ids.update(matching_segs_ds[iiter, : seg_count_ds[iiter]])
+                    from_previous.clear()
+
+                    n_matches = len(matching_ids)
+                    if n_matches > current_seg_count:
+                        matching_segs_ds.resize((iter_count, n_matches))
+                        weights_ds.resize((iter_count, n_matches))
+                        current_seg_count = n_matches
+
+                    if n_matches > 0:
+                        seg_count_ds[iiter] = n_matches
+                        matching_ids = sorted(matching_ids)
+                        matching_segs_ds[iiter, :n_matches] = matching_ids
+                        weights_ds[iiter, :n_matches] = self.data_reader.get_iter_group(n_iter)['seg_index']['weight'][
+                            sorted(matching_ids)
+                        ]
+                        parent_ids = self.data_reader.get_iter_group(n_iter)['seg_index']['parent_id'][sorted(matching_ids)]
+                        from_previous.update(parent_id for parent_id in parent_ids if parent_id >= 0)  # filter initial states
+                        del parent_ids
+                    del matching_ids
+                    pi.progress += 1
+
+
+def entry_point():
+    WSelectTool().main()
+
+
+if __name__ == '__main__':
+    entry_point()

westpa/cli/tools/w_stateprobs.py

@@ -0,0 +1,111 @@
+from westpa.tools import WESTMasterCommand, WESTParallelTool
+from warnings import warn
+
+from westpa.cli.tools.w_direct import DStateProbs
+
+# Just a shim to make sure everything works and is backwards compatible.
+# We're making sure it has the appropriate functions so that it can be called
+# as a regular tool, and not a subcommand.
+
+
+class WStateProbs(DStateProbs):
+    subcommand = 'trace'
+    help_text = 'averages and CIs for path-tracing kinetics analysis'
+    default_output_file = 'stateprobs.h5'
+    # This isn't strictly necessary, but for the moment, here it is.
+    # We really need to modify the underlying class so that we don't pull this sort of stuff if it isn't necessary.
+    # That'll take some case handling, which is fine.
+    default_kinetics_file = 'assign.h5'
+
+
+class WDirect(WESTMasterCommand, WESTParallelTool):
+    prog = 'w_stateprobs'
+    subcommands = [WStateProbs]
+    subparsers_title = 'calculate state-to-state kinetics by tracing trajectories'
+    description = '''\
+Calculate average populations and associated errors in state populations from
+weighted ensemble data. Bin assignments, including macrostate definitions,
+are required. (See "w_assign --help" for more information).
+
+-----------------------------------------------------------------------------
+Output format
+-----------------------------------------------------------------------------
+
+The output file (-o/--output, usually "stateprobs.h5") contains the following
+dataset:
+
+  /avg_state_pops [state]
+    (Structured -- see below) Population of each state across entire
+    range specified.
+
+If --evolution-mode is specified, then the following additional dataset is
+available:
+
+  /state_pop_evolution [window][state]
+    (Structured -- see below). State populations based on windows of
+    iterations of varying width.  If --evolution-mode=cumulative, then
+    these windows all begin at the iteration specified with
+    --start-iter and grow in length by --step-iter for each successive
+    element. If --evolution-mode=blocked, then these windows are all of
+    width --step-iter (excluding the last, which may be shorter), the first
+    of which begins at iteration --start-iter.
+
+The structure of these datasets is as follows:
+
+  iter_start
+    (Integer) Iteration at which the averaging window begins (inclusive).
+
+  iter_stop
+    (Integer) Iteration at which the averaging window ends (exclusive).
+
+  expected
+    (Floating-point) Expected (mean) value of the rate as evaluated within
+    this window, in units of inverse tau.
+
+  ci_lbound
+    (Floating-point) Lower bound of the confidence interval on the rate
+    within this window, in units of inverse tau.
+
+  ci_ubound
+    (Floating-point) Upper bound of the confidence interval on the rate
+    within this window, in units of inverse tau.
+
+  corr_len
+    (Integer) Correlation length of the rate within this window, in units
+    of tau.
+
+Each of these datasets is also stamped with a number of attributes:
+
+  mcbs_alpha
+    (Floating-point) Alpha value of confidence intervals. (For example,
+    *alpha=0.05* corresponds to a 95% confidence interval.)
+
+  mcbs_nsets
+    (Integer) Number of bootstrap data sets used in generating confidence
+    intervals.
+
+  mcbs_acalpha
+    (Floating-point) Alpha value for determining correlation lengths.
+
+
+-----------------------------------------------------------------------------
+Command-line options
+-----------------------------------------------------------------------------
+'''
+
+
+def entry_point():
+    warn('{} is being deprecated.  Please use w_direct instead.'.format(WDirect.prog))
+    # If we're not really supporting subcommands...
+    import sys
+
+    try:
+        if sys.argv[1] != 'trace':
+            sys.argv.insert(1, 'trace')
+    except Exception:
+        sys.argv.insert(1, 'trace')
+    WDirect().main()
+
+
+if __name__ == '__main__':
+    entry_point()