westpa 2022.10__cp312-cp312-macosx_11_0_arm64.whl

westpa/tools/iter_range.py

@@ -0,0 +1,198 @@
+import logging
+
+import numpy as np
+
+import westpa
+from westpa.tools.core import WESTToolComponent
+from westpa.core import h5io
+
+log = logging.getLogger(__name__)
+
+
+class IterRangeSelection(WESTToolComponent):
+    '''Select and record limits on iterations used in analysis and/or reporting.
+    This class provides both the user-facing command-line options and parsing, and
+    the application-side API for recording limits in HDF5.
+
+    HDF5 datasets calculated based on a restricted set of iterations should be tagged
+    with the following attributes:
+
+        ``first_iter``
+            The first iteration included in the calculation.
+
+        ``last_iter``
+            One past the last iteration included in the calculation.
+
+        ``iter_step``
+            Blocking or sampling period for iterations included in the calculation.
+    '''
+
+    def __init__(self, data_manager=None):
+        super().__init__()
+
+        self.data_manager = data_manager
+
+        # First iteration on which to perform analysis/reporting
+        self.iter_start = None
+
+        # One past the last iteration on which to perform analysis/reporting
+        self.iter_stop = None
+
+        # Step
+        self.iter_step = None
+
+        self.iter_count = None
+
+        self.include_args.update({'iter_start': True, 'iter_stop': True, 'iter_step': False})
+
+    def add_args(self, parser):
+        group = parser.add_argument_group('iteration range')
+
+        if self.include_args['iter_start']:
+            group.add_argument(
+                '--first-iter',
+                dest='first_iter',
+                type=int,
+                metavar='N_ITER',
+                default=1,
+                help='''Begin analysis at iteration N_ITER (default: %(default)d).''',
+            )
+        if self.include_args['iter_stop']:
+            group.add_argument(
+                '--last-iter',
+                dest='last_iter',
+                type=int,
+                metavar='N_ITER',
+                help='''Conclude analysis with N_ITER, inclusive (default: last completed iteration).''',
+            )
+        if self.include_args['iter_step']:
+            group.add_argument(
+                '--step-iter', dest='step_iter', type=int, metavar='STEP', help='''Analyze/report in blocks of STEP iterations.'''
+            )
+
+    def process_args(self, args, override_iter_start=None, override_iter_stop=None, default_iter_step=1):
+        if override_iter_start is not None:
+            self.iter_start = override_iter_start
+        elif args.first_iter is not None:
+            self.iter_start = args.first_iter
+        else:
+            self.iter_start = 1
+
+        if override_iter_stop is not None:
+            self.iter_stop = override_iter_stop
+        elif args.last_iter is not None:
+            self.iter_stop = args.last_iter + 1
+        else:
+            self.iter_stop = (self.data_manager or westpa.rc.get_data_manager()).current_iteration
+
+        if self.include_args['iter_step']:
+            self.iter_step = args.step_iter or default_iter_step
+
+        try:
+            self.iter_count = self.iter_stop - self.iter_start
+        except TypeError:
+            # one or both are None
+            pass
+
+    def iter_block_iter(self):
+        '''Return an iterable of (block_start,block_end) over the blocks of iterations
+        selected by --first-iter/--last-iter/--step-iter.'''
+
+        for blkfirst in range(self.iter_start, self.iter_stop, self.iter_step):
+            yield (blkfirst, min(self.iter_stop, blkfirst + self.iter_step))
+
+    def n_iter_blocks(self):
+        '''Return the number of blocks of iterations (as returned by ``iter_block_iter``)
+        selected by --first-iter/--last-iter/--step-iter.'''
+        npoints = self.iter_stop - self.iter_start
+        if npoints % self.iter_step == 0:
+            return npoints // self.iter_step
+        else:
+            return npoints // self.iter_step + 1
+
+    def record_data_iter_range(self, h5object, iter_start=None, iter_stop=None):
+        '''Store attributes ``iter_start`` and ``iter_stop`` on the given HDF5 object (group/dataset)'''
+        iter_start = self.iter_start if iter_start is None else iter_start
+        iter_stop = self.iter_stop if iter_stop is None else iter_stop
+        h5object.attrs['iter_start'] = iter_start
+        h5object.attrs['iter_stop'] = iter_stop
+
+    def record_data_iter_step(self, h5object, iter_step=None):
+        '''Store attribute ``iter_step`` on the given HDF5 object (group/dataset).'''
+        iter_step = self.iter_step if iter_step is None else iter_step
+        h5object.attrs['iter_step'] = iter_step
+
+    def check_data_iter_range_least(self, h5object, iter_start=None, iter_stop=None):
+        '''Check that the given HDF5 object contains (as denoted by its ``iter_start``/``iter_stop`` attributes)
+        data at least for the iteration range specified.'''
+        iter_start = self.iter_start if iter_start is None else iter_start
+        iter_stop = self.iter_stop if iter_stop is None else iter_stop
+
+        return h5io.check_iter_range_least(h5object, iter_start, iter_stop)
+
+    def check_data_iter_range_equal(self, h5object, iter_start=None, iter_stop=None):
+        '''Check that the given HDF5 object contains (as denoted by its ``iter_start``/``iter_stop`` attributes)
+        data exactly for the iteration range specified.'''
+
+        iter_start = self.iter_start if iter_start is None else iter_start
+        iter_stop = self.iter_stop if iter_stop is None else iter_stop
+
+        return h5io.check_iter_range_equal(h5object, iter_start, iter_stop)
+
+    def check_data_iter_step_conformant(self, h5object, iter_step=None):
+        '''Check that the given HDF5 object contains per-iteration data at an iteration stride suitable for extracting data
+        with the given stride (in other words, the given ``iter_step`` is a multiple of the stride with
+        which data was recorded).'''
+
+        iter_step = iter_step or self.iter_step
+        obj_iter_step = h5object.attrs.get('iter_step')
+        return obj_iter_step % iter_step == 0
+
+    def check_data_iter_step_equal(self, h5object, iter_step=None):
+        '''Check that the given HDF5 object contains per-iteration data at an iteration stride the same as
+        that specified.'''
+        iter_step = iter_step or self.iter_step
+        obj_iter_step = h5object.attrs.get('iter_step')
+        return obj_iter_step == iter_step
+
+    def slice_per_iter_data(self, dataset, iter_start=None, iter_stop=None, iter_step=None, axis=0):
+        '''Return the subset of the given dataset corresponding to the given iteration range and stride. Unless
+        otherwise specified, the first dimension of the dataset is the one sliced.'''
+
+        iter_start = self.iter_start if iter_start is None else iter_start
+        iter_stop = self.iter_stop if iter_stop is None else iter_stop
+        iter_step = self.iter_step if iter_step is None else iter_step
+
+        ds_iter_start = dataset.attrs['iter_start']
+        ds_iter_stop = dataset.attrs['iter_stop']
+        ds_iter_step = dataset.attrs.get('iter_step', 1)
+
+        if iter_start < ds_iter_start or iter_stop > ds_iter_stop or ds_iter_step % iter_step > 0:
+            raise IndexError(
+                'Cannot slice requested iterations [{:d},{:d}) (stride={:d}) from dataset {!r} with range [{:d},{:d}) (stride={:d}).'.format(
+                    iter_start, iter_stop, iter_step, dataset, ds_iter_start, ds_iter_stop, ds_iter_step
+                )
+            )
+
+        dimslices = []
+        for idim in range(len(dataset.shape)):
+            if idim == axis:
+                dimslices.append(slice(iter_start - ds_iter_start, iter_stop - ds_iter_stop + iter_step, iter_step))
+            else:
+                dimslices.append(slice(None, None, None))
+
+        dimslices = tuple(dimslices)
+        log.debug('slicing {!r} with {!r}'.format(dataset, dimslices))
+        data = dataset[dimslices]
+        log.debug('resulting data is of shape {!r}'.format(data.shape))
+        return data
+
+    def iter_range(self, iter_start=None, iter_stop=None, iter_step=None, dtype=None):
+        '''Return a sequence for the given iteration numbers and stride, filling
+        in missing values from those stored on ``self``. The smallest data type capable of
+        holding ``iter_stop`` is returned unless otherwise specified using the ``dtype``
+        argument.'''
+        iter_start = self.iter_start if iter_start is None else iter_start
+        iter_stop = self.iter_stop if iter_stop is None else iter_stop
+        iter_step = self.iter_step if iter_step is None else iter_step
+        return np.arange(iter_start, iter_stop, iter_step, dtype=(dtype or np.min_scalar_type(iter_stop)))

westpa/tools/kinetics_tool.py

@@ -0,0 +1,340 @@
+import numpy as np
+
+from westpa.tools import WESTDataReader, IterRangeSelection, WESTSubcommand, ProgressIndicatorComponent
+
+from westpa import mclib
+from westpa.core import h5io
+
+from westpa.tools.dtypes import iter_block_ci_dtype as ci_dtype
+
+
+# A function to just help with creating future objects for the work manager.
+
+
+def generate_future(work_manager, name, eval_block, kwargs):
+    submit_kwargs = {'name': name}
+    submit_kwargs.update(kwargs)
+    future = work_manager.submit(eval_block, kwargs=submit_kwargs)
+    return future
+
+
+class WESTKineticsBase(WESTSubcommand):
+    '''
+    Common argument processing for w_direct/w_reweight subcommands.
+    Mostly limited to handling input and output from w_assign.
+    '''
+
+    def __init__(self, parent):
+        super().__init__(parent)
+
+        self.data_reader = WESTDataReader()
+        self.iter_range = IterRangeSelection()
+        self.progress = ProgressIndicatorComponent()
+
+        self.output_filename = None
+        # This is actually applicable to both.
+        self.assignment_filename = None
+
+        self.output_file = None
+        self.assignments_file = None
+
+        self.evolution_mode = None
+
+        self.mcbs_alpha = None
+        self.mcbs_acalpha = None
+        self.mcbs_nsets = None
+
+        # Now we're adding in things that come from the old w_kinetics
+        self.do_compression = True
+
+    def add_args(self, parser):
+        self.progress.add_args(parser)
+        self.data_reader.add_args(parser)
+        self.iter_range.include_args['iter_step'] = True
+        self.iter_range.add_args(parser)
+
+        iogroup = parser.add_argument_group('input/output options')
+        iogroup.add_argument(
+            '-a',
+            '--assignments',
+            default='assign.h5',
+            help='''Bin assignments and macrostate definitions are in ASSIGNMENTS
+                            (default: %(default)s).''',
+        )
+
+        iogroup.add_argument(
+            '-o',
+            '--output',
+            dest='output',
+            default=self.default_output_file,
+            help='''Store results in OUTPUT (default: %(default)s).''',
+        )
+
+    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, default_iter_step=None)
+        if self.iter_range.iter_step is None:
+            # use about 10 blocks by default
+            self.iter_range.iter_step = max(1, (self.iter_range.iter_stop - self.iter_range.iter_start) // 10)
+
+        self.output_filename = args.output
+        self.assignments_filename = args.assignments
+
+
+# This provides some convenience functions, modified from w_kinavg, to help with calculating evolution and averages for observables with the mclib library in a consistent manner.
+# It's used in both w_direct and w_reweight.
+class AverageCommands(WESTKineticsBase):
+    default_output_file = 'direct.h5'
+
+    def __init__(self, parent):
+        # Ideally, this is stuff general to all the calculations we want to perform.
+        super().__init__(parent)
+        self.kinetics_filename = None
+        self.kinetics_file = None
+
+    def add_args(self, parser):
+        iogroup = parser.add_argument_group('input/output options')
+        # self.default_kinetics_file will be picked up as a class attribute from the appropriate subclass
+        # We can do this with the output file, too...
+        # ... by default, however, we're going to use {direct/reweight}.h5 for everything.
+        # Modules which are called with different default values will, of course, still use those.
+        iogroup.add_argument(
+            '-k',
+            '--kinetics',
+            default=self.default_kinetics_file,
+            help='''Populations and transition rates are stored in KINETICS
+                            (default: %(default)s).''',
+        )
+
+        cgroup = parser.add_argument_group('confidence interval calculation options')
+        cgroup.add_argument(
+            '--disable-bootstrap',
+            '-db',
+            dest='bootstrap',
+            action='store_const',
+            const=False,
+            help='''Enable the use of Monte Carlo Block Bootstrapping.''',
+        )
+        cgroup.add_argument(
+            '--disable-correl',
+            '-dc',
+            dest='correl',
+            action='store_const',
+            const=False,
+            help='''Disable the correlation analysis.''',
+        )
+        cgroup.add_argument(
+            '--alpha',
+            type=float,
+            default=0.05,
+            help='''Calculate a (1-ALPHA) confidence interval'
+                             (default: %(default)s)''',
+        )
+        cgroup.add_argument(
+            '--autocorrel-alpha',
+            type=float,
+            dest='acalpha',
+            metavar='ACALPHA',
+            help='''Evaluate autocorrelation to (1-ACALPHA) significance.
+                             Note that too small an ACALPHA will result in failure to detect autocorrelation
+                             in a noisy flux signal. (Default: same as ALPHA.)''',
+        )
+        cgroup.add_argument('--nsets', type=int, help='''Use NSETS samples for bootstrapping (default: chosen based on ALPHA)''')
+
+        cogroup = parser.add_argument_group('calculation options')
+        cogroup.add_argument(
+            '-e',
+            '--evolution-mode',
+            choices=['cumulative', 'blocked', 'none'],
+            default='none',
+            help='''How to calculate time evolution of rate estimates.
+                             ``cumulative`` evaluates rates over windows starting with --start-iter and getting progressively
+                             wider to --stop-iter by steps of --step-iter.
+                             ``blocked`` evaluates rates over windows of width --step-iter, the first of which begins at
+                             --start-iter.
+                             ``none`` (the default) disables calculation of the time evolution of rate estimates.''',
+        )
+        cogroup.add_argument(
+            '--window-frac',
+            type=float,
+            default=1.0,
+            help='''Fraction of iterations to use in each window when running in ``cumulative`` mode.
+                             The (1 - frac) fraction of iterations will be discarded from the start of each window.''',
+        )
+
+        mgroup = parser.add_argument_group('misc options')
+        mgroup.add_argument(
+            '--disable-averages',
+            '-da',
+            dest='display_averages',
+            action='store_false',
+            help='''Whether or not the averages should be printed to the console (set to FALSE if flag is used).''',
+        )
+
+    def process_args(self, args):
+        self.kinetics_filename = args.kinetics
+
+        # Disable the bootstrap or the correlation analysis.
+        self.mcbs_enable = args.bootstrap if args.bootstrap is not None else True
+        self.do_correl = args.correl if args.correl is not None else True
+        self.mcbs_alpha = args.alpha
+        self.mcbs_acalpha = args.acalpha if args.acalpha else self.mcbs_alpha
+        self.mcbs_nsets = args.nsets if args.nsets else mclib.get_bssize(self.mcbs_alpha)
+
+        self.display_averages = args.display_averages
+
+        self.evolution_mode = args.evolution_mode
+        self.evol_window_frac = args.window_frac
+        if self.evol_window_frac <= 0 or self.evol_window_frac > 1:
+            raise ValueError('Parameter error -- fractional window defined by --window-frac must be in (0,1]')
+
+    def stamp_mcbs_info(self, dataset):
+        dataset.attrs['mcbs_alpha'] = self.mcbs_alpha
+        dataset.attrs['mcbs_acalpha'] = self.mcbs_acalpha
+        dataset.attrs['mcbs_nsets'] = self.mcbs_nsets
+
+    def open_files(self):
+        self.output_file = h5io.WESTPAH5File(self.output_filename, 'a', creating_program=True)
+        h5io.stamp_creator_data(self.output_file)
+        self.assignments_file = h5io.WESTPAH5File(self.assignments_filename, 'r')  # , driver='core', backing_store=False)
+        self.kinetics_file = h5io.WESTPAH5File(self.kinetics_filename, 'r')  # , driver='core', backing_store=False)
+        if not self.iter_range.check_data_iter_range_least(self.assignments_file):
+            raise ValueError('assignments data do not span the requested iterations')
+
+    def open_assignments(self):
+        # Actually, I should rename this, as we're not OPENING assignments.
+        # This seems to be stuff we're going to be using a lot, so.
+        self.nstates = self.assignments_file.attrs['nstates']
+        self.nbins = self.assignments_file.attrs['nbins']
+        self.state_labels = self.assignments_file['state_labels'][...]
+        assert self.nstates == len(self.state_labels)
+        self.start_iter, self.stop_iter, self.step_iter = (
+            self.iter_range.iter_start,
+            self.iter_range.iter_stop,
+            self.iter_range.iter_step,
+        )
+
+        # Import for the reweighting code.
+
+        # state_map = self.assignments_file['state_map'][...]
+
+        # We've moved this into a different step so that it's compatible with
+        # loading up from the all command.
+        # Otherwise, we try to load the kinetics (since we're just mixing subclasses)
+        # before it's actually run, and so we fail out.
+        if not self.iter_range.check_data_iter_range_least(self.kinetics_file):
+            raise ValueError('kinetics data do not span the requested iterations')
+
+    def print_averages(self, dataset, header, dim=1):
+        print(header)
+        maxlabellen = max(list(map(len, self.state_labels)))
+        for istate in range(self.nstates):
+            if dim == 1:
+                print(
+                    '{:{maxlabellen}s}: mean={:21.15e} CI=({:21.15e}, {:21.15e}) * tau^-1'.format(
+                        self.state_labels[istate],
+                        dataset['expected'][istate],
+                        dataset['ci_lbound'][istate],
+                        dataset['ci_ubound'][istate],
+                        maxlabellen=maxlabellen,
+                    )
+                )
+
+            else:
+                for jstate in range(self.nstates):
+                    if istate == jstate:
+                        continue
+                    print(
+                        '{:{maxlabellen}s} -> {:{maxlabellen}s}: mean={:21.15e} CI=({:21.15e}, {:21.15e}) * tau^-1'.format(
+                            self.state_labels[istate],
+                            self.state_labels[jstate],
+                            dataset['expected'][istate, jstate],
+                            dataset['ci_lbound'][istate, jstate],
+                            dataset['ci_ubound'][istate, jstate],
+                            maxlabellen=maxlabellen,
+                        )
+                    )
+
+    def run_calculation(
+        self, pi, nstates, start_iter, stop_iter, step_iter, dataset, eval_block, name, dim, do_averages=False, **extra
+    ):
+        # We want to use the same codepath to run a quick average as we do the longer evolution sets, so...
+        if do_averages:
+            start_pts = [start_iter, stop_iter]
+        else:
+            start_pts = list(range(start_iter, stop_iter, step_iter))
+        # Our evolution dataset!
+        if dim == 2:
+            evolution_dataset = np.zeros((len(start_pts), nstates, nstates), dtype=ci_dtype)
+        elif dim == 1:
+            evolution_dataset = np.zeros((len(start_pts), nstates), dtype=ci_dtype)
+        else:
+            # Temp.
+            print("What's wrong?")
+
+        # This is appropriate for bootstrapped quantities, I think.
+
+        if True:
+            futures = []
+            for iblock, start in enumerate(start_pts):
+                stop = min(start + step_iter, stop_iter)
+                if self.evolution_mode == 'cumulative' or do_averages is True:
+                    windowsize = int(self.evol_window_frac * (stop - start_iter))
+                    block_start = max(start_iter, stop - windowsize)
+                else:  # self.evolution_mode == 'blocked'
+                    block_start = start
+
+                # Create a basic set of kwargs for this iteration slice.
+                future_kwargs = dict(
+                    iblock=iblock,
+                    start=block_start,
+                    stop=stop,
+                    nstates=nstates,
+                    mcbs_alpha=self.mcbs_alpha,
+                    mcbs_nsets=self.mcbs_nsets,
+                    mcbs_acalpha=self.mcbs_acalpha,
+                    do_correl=self.do_correl,
+                    name=name,
+                    mcbs_enable=self.mcbs_enable,
+                    data_input={},
+                    **extra
+                )
+
+                # Slice up the datasets for this iteration slice.
+                # We're assuming they're all h5io iter blocked datasets; it's up to the calling routine
+                # to ensure this is true.
+
+                # Actually, I'm less sure how to handle this for pre-calculated datasets.  Need to consider this.  But for now...
+                for key, value in dataset.items():
+                    try:
+                        future_kwargs['data_input'][key] = (
+                            value.iter_slice(block_start, stop) if hasattr(value, 'iter_slice') else value[block_start:stop]
+                        )
+                    except Exception:
+                        future_kwargs['data_input'][key] = (
+                            value.iter_slice(block_start, stop) if hasattr(value, 'iter_slice') else value[block_start:stop, :]
+                        )
+                    # print(future_kwargs['data_input'][key])
+
+                # We create a future object with the appropriate name, and then append it to the work manager.
+                futures.append(generate_future(self.work_manager, name, eval_block, future_kwargs))
+
+            pi.new_operation('Calculating {}'.format(name), len(futures))
+
+            # Now, we wait to get the result back; we'll store it in the result, and return it.
+            for future in self.work_manager.as_completed(futures):
+                pi.progress = iblock / step_iter
+                future_result = future.get_result(discard=True)
+
+                if dim == 2:
+                    for result in future_result:
+                        name, iblock, istate, jstate, ci_result = result
+                        evolution_dataset[iblock, istate, jstate] = ci_result
+                elif dim == 1:
+                    for result in future_result:
+                        name, iblock, istate, ci_result = result
+                        evolution_dataset[iblock, istate] = ci_result
+
+        return evolution_dataset