westpa 2022.12__cp313-cp313-manylinux_2_17_x86_64.manylinux2014_x86_64.whl

westpa/cli/tools/w_bins.py

@@ -0,0 +1,166 @@
+import logging
+import sys
+
+import numpy as np
+
+from westpa.tools import WESTTool, WESTDataReader, BinMappingComponent
+import westpa
+
+from westpa.tools.binning import write_bin_info
+
+
+log = logging.getLogger('w_bins')
+
+
+class WBinTool(WESTTool):
+    prog = 'w_bins'
+    description = '''\
+Display information and statistics about binning in a WEST simulation, or
+modify the binning for the current iteration of a WEST simulation.
+-------------------------------------------------------------------------------
+'''
+
+    def __init__(self):
+        super().__init__()
+        self.subcommand = None
+        self.data_reader = WESTDataReader()
+        self.binning = BinMappingComponent()
+        self.args = None
+        self.n_iter = None
+
+    # Interface for command-line tools
+    def add_args(self, parser):
+        self.data_reader.add_args(parser)
+
+        subparsers = parser.add_subparsers(help='available commands')
+
+        info_parser = subparsers.add_parser('info', help='Display information about binning.')
+        info_parser.add_argument(
+            '-n', '--n-iter', type=int, help='''Consider initial points of segment N_ITER (default: current iteration).'''
+        )
+        info_parser.add_argument(
+            '--detail',
+            action='store_true',
+            help='''Display detailed per-bin information in addition to summary
+                                 information.''',
+        )
+        self.binning.add_args(info_parser)
+        info_parser.set_defaults(func=self.cmd_info)
+
+        rebin_parser = subparsers.add_parser('rebin', help='Rebuild current iteration with new binning.')
+        rebin_parser.add_argument(
+            '--confirm',
+            action='store_true',
+            help='''Commit the revised iteration to HDF5; without this option, the effects of the
+                                  new binning are only calculated and printed.''',
+        )
+        rebin_parser.add_argument(
+            '--detail',
+            action='store_true',
+            help='''Display detailed per-bin information in addition to summary
+                                     information.''',
+        )
+        rebin_parser.add_argument(
+            '-n', '--n-iter', type=int, help='''Consider initial points of segment N_ITER (default: current iteration).'''
+        )
+        self.binning.add_args(rebin_parser, suppress=['--bins-from-file'])
+        self.binning.add_target_count_args(rebin_parser)
+        rebin_parser.set_defaults(func=self.cmd_rebin)
+
+    def process_args(self, args):
+        self.data_reader.process_args(args)
+        self.data_reader.open(mode='r+')
+        self.n_iter = getattr(args, 'n_iter', None) or self.data_reader.current_iteration
+
+        # we cannot read bin information during rebins
+        # interesting note: '==' is required here; 'is' fails
+        if args.func == self.cmd_rebin:
+            self.binning.target_counts_required = True
+        else:
+            self.binning.set_we_h5file_info(self.n_iter, self.data_reader)
+
+        self.binning.process_args(args)
+
+        self.args = args
+        self.subcommand = args.func
+
+    def go(self):
+        self.subcommand()
+
+    def cmd_info(self):
+        mapper = self.binning.mapper
+
+        # Get target states and their assignments
+        target_states = self.data_reader.get_target_states(self.n_iter)
+        n_target_states = len(target_states)
+
+        iter_group = self.data_reader.get_iter_group(self.n_iter)
+
+        # bin initial pcoords for iteration n_iter
+        initial_pcoords = iter_group['pcoord'][:, 0, :]
+        assignments = mapper.assign(initial_pcoords)
+        del initial_pcoords
+
+        print('Bin information for iteration {:d}'.format(self.n_iter))
+
+        # Get bin counts and weights
+        weights = iter_group['seg_index']['weight']
+
+        write_bin_info(mapper, assignments, weights, n_target_states, detailed=self.args.detail)
+
+    def cmd_rebin(self):
+        mapper = self.binning.mapper
+        assert mapper is not None
+        if self.n_iter == 1:
+            sys.stderr.write('rebin is not supported for the first iteration; reinitialize with w_init instead\n')
+            sys.exit(1)
+        n_target_states = len(self.data_reader.get_target_states(self.n_iter))
+        we_driver = westpa.rc.get_we_driver()
+        data_manager = self.data_reader.data_manager
+
+        segments = data_manager.get_segments(self.n_iter, load_pcoords=True)
+        last_iter_segments = data_manager.get_segments(self.n_iter - 1, load_pcoords=False)
+
+        # Bin on this iteration's initial points
+        # We don't have to worry about recycling because we are binning on
+        # initial points rather than final points, so recycling has already
+        # occurred for this iteration.
+        # We do need initial states, in case we merge a newly-created walker out of existence
+        # avail_initial_states = {state.state_id: state
+        #                        for state in data_manager.get_unused_initial_states(n_iter = self.n_iter)}
+        avail_initial_states = data_manager.get_unused_initial_states(n_iter=self.n_iter)
+        used_initial_states = data_manager.get_segment_initial_states(segments)
+        we_driver.new_iteration(
+            initial_states=avail_initial_states, bin_mapper=mapper, bin_target_counts=self.binning.bin_target_counts
+        )
+        we_driver.used_initial_states = {state.state_id: state for state in used_initial_states}
+        we_driver.assign(segments, initializing=True)
+        we_driver.rebin_current(parent_segments=last_iter_segments)
+
+        weights = np.array([segment.weight for segment in we_driver.next_iter_segments])
+        assignments = np.fromiter(we_driver.next_iter_assignments, dtype=int, count=len(weights))
+        write_bin_info(mapper, assignments, weights, n_target_states, detailed=self.args.detail)
+
+        if self.args.confirm:
+            data_manager.prepare_iteration(self.n_iter, list(we_driver.next_iter_segments))
+
+            # manually update endpoint statuses only
+            endpoint_types = sorted([(segment.seg_id, segment.endpoint_type) for segment in last_iter_segments])
+            last_iter_group = data_manager.get_iter_group(self.n_iter - 1)
+            last_iter_index = last_iter_group['seg_index'][...]
+            last_iter_index['endpoint_type'] = [pair[1] for pair in endpoint_types]
+            last_iter_group['seg_index'][...] = last_iter_index
+
+            data_manager.save_iter_binning(
+                self.n_iter, self.binning.mapper_hash, self.binning.mapper_pickle, we_driver.bin_target_counts
+            )
+            data_manager.update_initial_states(we_driver.all_initial_states)
+            data_manager.flush_backing()
+
+
+def entry_point():
+    WBinTool().main()
+
+
+if __name__ == '__main__':
+    entry_point()

westpa/cli/tools/w_crawl.py

@@ -0,0 +1,119 @@
+import logging
+
+from westpa.tools import WESTParallelTool, WESTDataReader, IterRangeSelection, ProgressIndicatorComponent
+import westpa
+from westpa.core.extloader import get_object
+
+log = logging.getLogger('w_crawl')
+
+
+class WESTPACrawler:
+    '''Base class for general crawling execution. This class
+    only exists on the master.'''
+
+    def initialize(self, iter_start, iter_stop):
+        '''Initialize this crawling process.'''
+        pass
+
+    def finalize(self):
+        '''Finalize this crawling process.'''
+        pass
+
+    def process_iter_result(self, n_iter, result):
+        '''Process the result of a per-iteration task.'''
+        pass
+
+
+def _remote_task(n_iter, taskfn):
+    data_manager = westpa.rc.get_data_manager()  # gaahhh...globals
+    data_manager.open_backing(mode='r')
+    return n_iter, taskfn(n_iter, data_manager.get_iter_group(n_iter))
+
+
+class WCrawl(WESTParallelTool):
+    prog = 'w_crawl'
+    description = '''\
+Crawl a weighted ensemble dataset, executing a function for each iteration.
+This can be used for postprocessing of trajectories, cleanup of datasets,
+or anything else that can be expressed as "do X for iteration N, then do
+something with the result". Tasks are parallelized by iteration, and
+no guarantees are made about evaluation order.
+
+
+-----------------------------------------------------------------------------
+Command-line options
+-----------------------------------------------------------------------------
+
+'''
+
+    def __init__(self):
+        super().__init__()
+
+        # These are used throughout
+        self.progress = ProgressIndicatorComponent()
+        self.data_reader = WESTDataReader()
+        self.iter_range = IterRangeSelection(self.data_reader)
+
+        self.crawler = None
+        self.task_callable = None
+
+    def add_args(self, parser):
+        self.data_reader.add_args(parser)
+        self.iter_range.add_args(parser)
+
+        tgroup = parser.add_argument_group('task options')
+        tgroup.add_argument(
+            '-c',
+            '--crawler-instance',
+            help='''Use CRAWLER_INSTANCE (specified as module.instance) as an instance of
+                            WESTPACrawler to coordinate the calculation. Required only if initialization,
+                            finalization, or task result processing is required.''',
+        )
+        tgroup.add_argument(
+            'task_callable',
+            help='''Run TASK_CALLABLE (specified as module.function) on each iteration.
+                            Required.''',
+        )
+        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)
+
+        self.task_callable = get_object(args.task_callable, path=['.'])
+        if args.crawler_instance is not None:
+            self.crawler = get_object(args.crawler_instance, path=['.'])
+        else:
+            self.crawler = WESTPACrawler()
+
+    def go(self):
+        iter_start = self.iter_range.iter_start
+        iter_stop = self.iter_range.iter_stop
+        iter_count = iter_stop - iter_start
+        self.data_reader.open('r')
+        pi = self.progress.indicator
+        with pi:
+            pi.operation = 'Initializing'
+            self.crawler.initialize(iter_start, iter_stop)
+
+            try:
+                pi.new_operation('Dispatching tasks & processing results', iter_count)
+                task_gen = ((_remote_task, (n_iter, self.task_callable), {}) for n_iter in range(iter_start, iter_stop))
+                for future in self.work_manager.submit_as_completed(task_gen, self.max_queue_len):
+                    n_iter, result = future.get_result(discard=True)
+                    if self.crawler is not None:
+                        self.crawler.process_iter_result(n_iter, result)
+                    pi.progress += 1
+            finally:
+                pi.new_operation('Finalizing')
+                self.crawler.finalize()
+
+
+def entry_point():
+    WCrawl().main()
+
+
+if __name__ == '__main__':
+    entry_point()