westpa 2022.10__cp312-cp312-macosx_11_0_arm64.whl

westpa/oldtools/aframe/trajwalker.py

@@ -0,0 +1,126 @@
+import logging
+
+log = logging.getLogger(__name__)
+
+import numpy
+
+
+class TrajWalker:
+    """A class to perform analysis by walking the trajectory tree.  A stack is used rather than recursion, or else
+    the highest number of iterations capable of being considered would be the same as the Python recursion limit.
+    """
+
+    def __init__(self, data_reader, history_chunksize=100):
+        self.data_reader = data_reader
+        self.history_chunksize = history_chunksize
+        self.n_segs_visited = 0
+
+    # TrajTree.count_segs_in_range() is now DataReader.total_segs_in_range()
+
+    def trace_to_root(self, n_iter, seg_id):
+        '''Trace the given segment back to its starting point, returning a list of Segment
+        objects describing the entire trajectory.'''
+
+        segments = []
+        segment = self.data_reader.get_segments_by_id(n_iter, [seg_id])[0]
+        segments.append(segment)
+        while segment.p_parent_id >= 0:
+            segment = self.data_reader.get_segments_by_id(segment.n_iter - 1, [segment.p_parent_id])[0]
+            segments.append(segment)
+        return list(reversed(segments))
+
+    def get_trajectory_roots(self, first_iter, last_iter, include_pcoords=True):
+        '''Get segments which start new trajectories.  If min_iter or max_iter is specified, restrict the
+        set of iterations within which the search is conducted.'''
+
+        roots = []
+        for n_iter in range(first_iter, last_iter + 1):
+            seg_ids = self.data_reader.get_created_seg_ids(n_iter)
+            segments = self.data_reader.get_segments_by_id(n_iter, seg_ids, include_pcoords=include_pcoords)
+            roots.extend(segments)
+        return roots
+
+    def get_initial_nodes(self, first_iter, last_iter, include_pcoords=True):
+        '''Get segments with which to begin a tree walk -- those alive or created within [first_iter,last_iter].'''
+
+        root_ids = dict()
+
+        # All trajectories alive or newly created in first_iter are initial nodes
+        root_ids[first_iter] = set(self.data_reader.get_seg_ids(first_iter))
+
+        # Find trajectories created in [first_iter, last_iter]
+        for n_iter in range(first_iter, last_iter + 1):
+            seg_ids = self.data_reader.get_created_seg_ids(n_iter)
+            try:
+                root_ids[n_iter].update(seg_ids)
+            except KeyError:
+                root_ids[n_iter] = set(seg_ids)
+
+        # Convert to Segment objects
+        segments = []
+        for n_iter, id_set in root_ids.items():
+            segments.extend(self.data_reader.get_segments_by_id(n_iter, id_set, include_pcoords=include_pcoords))
+        return segments
+
+    def trace_trajectories(
+        self, first_iter, last_iter, callable, include_pcoords=True, cargs=None, ckwargs=None, get_state=None, set_state=None
+    ):
+        """
+        Walk the trajectory tree depth-first, calling
+          ``callable(segment, children, history, *cargs, **ckwargs)`` for each segment
+        visited. ``segment`` is the segment being visited, ``children`` is that
+        segment's children, ``history`` is the chain of segments leading
+        to ``segment`` (not including ``segment``). get_state and set_state are
+        used to record and reset, respectively, any state specific to
+        ``callable`` when a new branch is traversed.
+        """
+
+        cargs = cargs or tuple()
+        ckwargs = ckwargs or dict()
+
+        # Either both or neither of external state getter/setter required
+        if (get_state or set_state) and not (get_state and set_state):
+            raise ValueError('either both or neither of get_state/set_state must be specified')
+
+        # This will grow to contain the maximum trajectory length
+        history = numpy.empty((self.history_chunksize,), numpy.object_)
+        roots = self.get_initial_nodes(first_iter, last_iter, include_pcoords)
+
+        for root in roots:
+            children = self.data_reader.get_children(root, include_pcoords)
+
+            # Visit the root node of each tree unconditionally
+            callable(root, children, [], *cargs, **ckwargs)
+            self.n_segs_visited += 1
+
+            state_stack = [{'node': root, 'children': children, 'len_history': 0, 'ext': get_state() if get_state else None}]
+
+            # Walk the tree, depth-first
+            while state_stack:
+                state = state_stack.pop(-1)
+
+                node = state['node']
+                children = state['children']
+                len_history = state['len_history']
+                if set_state:
+                    set_state(state['ext'])
+
+                # Descend as far as we can
+                while node.n_iter < last_iter and len(children):
+                    # Save current state before descending
+                    state_stack.append(
+                        {'node': node, 'children': children, 'len_history': len_history, 'ext': get_state() if get_state else None}
+                    )
+
+                    # Add an item to the historical record
+                    if len_history >= history.shape[0]:
+                        history.resize((history.shape[0] + self.history_chunksize,))
+                    history[len_history] = node
+                    len_history += 1
+
+                    node = children.pop(-1)
+                    children = self.data_reader.get_children(node, include_pcoords)
+
+                    # Visit the new node as we descend
+                    callable(node, children, history[:len_history], *cargs, **ckwargs)
+                    self.n_segs_visited += 1

westpa/oldtools/aframe/transitions.py

@@ -0,0 +1,469 @@
+import logging
+
+import numpy as np
+
+import westpa
+from westpa.oldtools.aframe import AnalysisMixin
+from westpa.oldtools.aframe.trajwalker import TrajWalker
+
+log = logging.getLogger(__name__)
+
+
+class TransitionEventAccumulator:
+    index_dtype = np.uintp
+    count_dtype = np.uint64
+    weight_dtype = np.float64
+    output_tdat_chunksize = 4096  # HDF5 chunksize for transition data (~300 KiB)
+    tdat_buffersize = 524288  # Internal buffer length (~38 MiB)
+    max_acc = 32768
+
+    def __init__(self, n_bins, output_group, calc_fpts=True):
+        self.calc_fpts = calc_fpts
+        self.n_bins = n_bins
+        self.iibins = np.arange(n_bins)
+        self.iibdisc = np.empty((n_bins,), np.bool_)
+
+        self.bin_index_dtype = np.min_scalar_type(n_bins)
+
+        self.tdat_dtype = np.dtype(
+            [
+                ('traj', self.index_dtype),
+                ('n_iter', self.index_dtype),
+                ('timepoint', self.index_dtype),
+                ('initial_bin', self.bin_index_dtype),
+                ('final_bin', self.bin_index_dtype),
+                ('initial_weight', self.weight_dtype),
+                ('final_weight', self.weight_dtype),
+                ('initial_bin_pop', self.weight_dtype),
+                ('duration', self.index_dtype),
+                ('fpt', self.index_dtype),
+            ]
+        )
+
+        # HDF5 group in which to store results
+        self.output_group = output_group
+        self.tdat_buffer = np.empty((self.tdat_buffersize,), dtype=self.tdat_dtype)
+        self.tdat_buffer_offset = 0
+        self.output_tdat_offset = 0
+        self.output_tdat_ds = None
+
+        # Accumulators/counters
+        self.n_trans = None  # shape (n_bins,n_bins)
+
+        # Time points and per-timepoint data
+        self.last_exit = None  # (n_bins,)
+        self.last_entry = None  # (n_bins,)
+        self.last_completion = None  # (n_bins,n_bins)
+        self.weight_last_exit = None  # (n_bins)
+        self.bin_pops_last_exit = None  # (n_bins,)
+
+        # Analysis continuation information
+        self.timepoint = None  # current time index for separate calls on same trajectory
+        self.last_bin = None  # last region occupied, for separate calls on same trajectory
+        self.last_bin_pop = None  # total weight in self.last_region at end of last processing step
+
+        self.clear()
+
+    def clear(self):
+        self.clear_state()
+        self.n_trans = np.zeros((self.n_bins, self.n_bins), self.count_dtype)
+        self.tdat_buffer = np.empty((self.tdat_buffersize,), dtype=self.tdat_dtype)
+        self.tdat_buffer_offset = 0
+        self.output_tdat_offset = 0
+        self.output_tdat_ds = None
+
+    def clear_state(self):
+        self.last_exit = np.zeros((self.n_bins,), self.index_dtype)
+        self.last_entry = np.zeros((self.n_bins,), self.index_dtype)
+        self.last_completion = np.zeros((self.n_bins, self.n_bins), self.index_dtype)
+        self.weight_last_exit = np.zeros((self.n_bins,), self.weight_dtype)
+        self.bin_pops_last_exit = np.zeros((self.n_bins,), self.weight_dtype)
+        self.timepoint = 0
+        self.last_bin = None
+        self.last_bin_pop = None
+
+    def get_state(self):
+        return {
+            'last_entry': self.last_entry.copy(),
+            'last_exit': self.last_exit.copy(),
+            'last_completion': self.last_completion.copy(),
+            'weight_last_exit': self.weight_last_exit.copy(),
+            'bin_pops_last_exit': self.bin_pops_last_exit.copy(),
+            'timepoint': self.timepoint,
+            'last_bin': self.last_bin,
+            'last_bin_pop': self.last_bin_pop,
+        }
+
+    def set_state(self, state_dict):
+        self.last_entry = state_dict['last_entry']
+        self.last_exit = state_dict['last_exit']
+        self.last_completion = state_dict['last_completion']
+        self.weight_last_exit = state_dict['weight_last_exit']
+        self.bin_pops_last_exit = state_dict['bin_pops_last_exit']
+        self.timepoint = state_dict['timepoint']
+        self.last_bin = state_dict['last_bin']
+        self.last_bin_pop = state_dict['last_bin_pop']
+
+    def record_transition_data(self, tdat):
+        """Update running statistics and write transition data to HDF5 (with buffering)"""
+
+        # Write out accumulated transition data
+        if self.output_tdat_ds is None:
+            # Create dataset
+            try:
+                del self.output_group['transitions']
+            except KeyError:
+                pass
+
+            self.output_tdat_ds = self.output_group.create_dataset(
+                'transitions',
+                shape=(1,),
+                dtype=self.tdat_dtype,
+                maxshape=(None,),
+                chunks=(self.output_tdat_chunksize,),
+                compression='gzip',
+            )
+
+        # If the amount of data to write exceeds our remaining buffer space, flush the buffer, then
+        # write data directly to HDF5, otherwise just add to the buffer and wait for the last flush
+        if len(tdat) + self.tdat_buffer_offset > self.tdat_buffersize:
+            self.flush_transition_data()
+            ub = self.output_tdat_offset + len(tdat)
+            self.output_tdat_ds.resize((ub,))
+            self.output_tdat_ds[self.output_tdat_offset : ub] = tdat
+            self.output_tdat_offset += len(tdat)
+        else:
+            self.tdat_buffer[self.tdat_buffer_offset : (self.tdat_buffer_offset + len(tdat))] = tdat
+            self.tdat_buffer_offset += len(tdat)
+
+    def flush_transition_data(self):
+        """Flush any unwritten output that may be present"""
+        if self.output_tdat_ds is None:
+            return
+
+        # self.tdat_buffer_offset is the number of items in the buffer
+        nbuf = self.tdat_buffer_offset
+        if nbuf == 0:
+            return
+        ub = nbuf + self.output_tdat_offset
+        if ub > self.output_tdat_ds.len():
+            # Resize dataset to fit data
+            self.output_tdat_ds.resize((ub,))
+
+        self.output_tdat_ds[self.output_tdat_offset : ub] = self.tdat_buffer[:nbuf]
+        self.output_tdat_offset += nbuf
+        self.tdat_buffer_offset = 0
+
+    def start_accumulation(self, assignments, weights, bin_pops, traj=0, n_iter=0):
+        self.clear_state()
+        timepoints = np.arange(len(assignments))
+        self._accumulate_transitions(timepoints, assignments, weights, bin_pops, traj, n_iter)
+
+    def continue_accumulation(self, assignments, weights, bin_pops, traj=0, n_iter=0):
+        aug_assign = np.empty((len(assignments) + 1,), assignments.dtype)
+        aug_assign[0] = self.last_bin
+        aug_assign[1:] = assignments
+
+        aug_weights = np.empty((len(weights) + 1,), self.weight_dtype)
+        aug_weights[0] = 0
+        aug_weights[1:] = weights
+
+        aug_pops = np.empty((len(bin_pops) + 1, len(bin_pops[0])), self.weight_dtype)
+        aug_pops[0, :] = 0
+        aug_pops[0, self.last_bin] = self.last_bin_pop
+        aug_pops[1:] = bin_pops
+
+        timepoints = np.arange(self.timepoint, self.timepoint + len(aug_assign))
+        self._accumulate_transitions(timepoints, aug_assign, aug_weights, aug_pops, traj, n_iter)
+
+    def _accumulate_transitions(self, timepoints, assignments, weights, bin_pops, traj, n_iter):
+        tdat = []
+
+        assignments_from_1 = assignments[1:]
+        assignments_to_1 = assignments[:-1]
+
+        calc_fpts = self.calc_fpts
+
+        trans_occur = assignments_from_1 != assignments_to_1
+        trans_ibin = assignments_to_1[trans_occur]
+        trans_fbin = assignments_from_1[trans_occur]
+        trans_timepoints = timepoints[1:][trans_occur]
+        trans_weights = weights[1:][trans_occur]  # arrival weights
+        trans_ibinpops = bin_pops[:-1][trans_occur]
+
+        last_exit = self.last_exit
+        last_entry = self.last_entry
+        last_completion = self.last_completion
+        bin_pops_last_exit = self.bin_pops_last_exit
+        weight_last_exit = self.weight_last_exit
+        n_trans = self.n_trans
+        iibdisc = self.iibdisc
+        iibins = self.iibins
+        tdat_maxlen = self.max_acc
+        for trans_ti, weight, ibin, fbin, ibinpops in zip(trans_timepoints, trans_weights, trans_ibin, trans_fbin, trans_ibinpops):
+            # Record this crossing event's data
+            bin_pops_last_exit[ibin] = ibinpops[ibin]
+            last_exit[ibin] = trans_ti
+            last_entry[fbin] = trans_ti
+            weight_last_exit[ibin] = weight
+
+            # See what other transitions this crossing event completes
+            iibdisc[:] = last_exit > 0
+            iibdisc &= last_entry > last_completion[:, fbin]
+
+            # Calculate event durations, etc for each transition generated by this crossing event
+            durations = -last_exit + trans_ti + 1  # = time now - time of exit from initial bin
+            if calc_fpts:
+                fpts = -last_completion[fbin, :] + trans_ti  # = time now - time of last final->initial transition
+                fpts[last_completion[:, fbin] == 0] = 0
+
+                for iibin in iibins[iibdisc]:
+                    tdat.append(
+                        (
+                            traj,
+                            n_iter,
+                            trans_ti,
+                            iibin,
+                            fbin,
+                            weight_last_exit[iibin],
+                            weight,
+                            bin_pops_last_exit[iibin],
+                            durations[iibin],
+                            fpts[iibin],
+                        )
+                    )
+            else:
+                for iibin in iibins[iibdisc]:
+                    tdat.append(
+                        (
+                            traj,
+                            n_iter,
+                            trans_ti,
+                            iibin,
+                            fbin,
+                            weight_last_exit[iibin],
+                            weight,
+                            bin_pops_last_exit[iibin],
+                            durations[iibin],
+                            0,
+                        )
+                    )
+
+            # Update tracking and statistics
+            last_completion[iibdisc, fbin] = trans_ti
+            n_trans[iibdisc, fbin] += 1
+
+            if len(tdat) > tdat_maxlen:
+                self.record_transition_data(tdat)
+
+        self.record_transition_data(tdat)
+        self.timepoint = timepoints[-1]
+        self.last_bin = assignments[-1]
+        self.last_bin_pop = bin_pops[-1, assignments[-1]]
+
+
+class TransitionAnalysisMixin(AnalysisMixin):
+    def __init__(self):
+        super().__init__()
+        self.discard_transition_data = False
+        self.calc_fpts = False
+        self.trans_h5gname = 'transitions'
+        self.trans_h5group = None
+        self.__transitions_ds = None
+        self.n_trajs = 0
+
+    def require_transitions_group(self):
+        if self.trans_h5group is None:
+            self.trans_h5group = self.anal_h5file.require_group(self.trans_h5gname)
+        return self.trans_h5group
+
+    def delete_transitions_group(self):
+        self.trans_h5group = None
+        del self.anal_h5file[self.trans_h5gname]
+
+    def get_transitions_ds(self):
+        if self.__transitions_ds is not None:
+            return self.__transitions_ds
+        else:
+            self.__transitions_ds = self.trans_h5group['transitions']
+            return self.__transitions_ds
+
+    def add_args(self, parser, upcall=True):
+        if upcall:
+            try:
+                upfunc = super().add_args
+            except AttributeError:
+                pass
+            else:
+                upfunc(parser)
+
+        group = parser.add_argument_group('transition analysis options')
+        group.add_argument(
+            '--discard-transition-data',
+            dest='discard_transition_data',
+            action='store_true',
+            help='''Discard any existing transition data stored in the analysis HDF5 file.''',
+        )
+
+    def process_args(self, args, upcall=True):
+        self.discard_transition_data = args.discard_transition_data
+
+        if upcall:
+            try:
+                upfunc = super().process_args
+            except AttributeError:
+                pass
+            else:
+                upfunc(args)
+
+    def require_transitions(self):
+        self.require_bin_assignments()
+        self.require_transitions_group()
+        do_trans = False
+        if self.discard_transition_data:
+            westpa.rc.pstatus('Discarding existing transition data.')
+            do_trans = True
+        elif not self.check_data_binhash(self.trans_h5group):
+            westpa.rc.pstatus('Bin definitions have changed; deleting existing transition data.')
+            do_trans = True
+        elif 'transitions' in self.trans_h5group and not self.check_data_iter_range_least(self.trans_h5group):
+            westpa.rc.pstatus('Existing transition data is for different first/last iterations; deleting.')
+            do_trans = True
+
+        if do_trans:
+            self.delete_transitions_group()
+            self.find_transitions()
+
+    def find_transitions(self):
+        westpa.rc.pstatus('Finding transitions...')
+        output_group = self.require_transitions_group()
+
+        self.n_segs_visited = 0
+        self.n_total_segs = self.total_segs_in_range(self.first_iter, self.last_iter)
+        self.accumulator = TransitionEventAccumulator(self.n_bins, output_group, calc_fpts=self.calc_fpts)
+        self.bin_assignments = self.get_bin_assignments(self.first_iter, self.last_iter)
+        self.bin_populations = self.get_bin_populations(self.first_iter, self.last_iter)
+
+        walker = TrajWalker(data_reader=self)
+
+        self.__pcoord_len = self.get_pcoord_len(self.first_iter)
+        self.__quiet_mode = westpa.rc.quiet_mode
+
+        walker.trace_trajectories(
+            self.first_iter,
+            self.last_iter,
+            callable=self._segment_callback,
+            include_pcoords=False,
+            get_state=self.accumulator.get_state,
+            set_state=self.accumulator.set_state,
+        )
+        self.accumulator.flush_transition_data()
+        try:
+            del output_group['n_trans']
+        except KeyError:
+            pass
+        output_group['n_trans'] = self.accumulator.n_trans
+
+        for h5object in (output_group, output_group['n_trans'], output_group['transitions']):
+            self.record_data_iter_range(h5object)
+            self.record_data_binhash(h5object)
+            h5object.attrs['n_trajs'] = self.n_trajs
+            h5object.attrs['n_segs'] = self.n_segs_visited
+
+        self.accumulator.clear()
+        westpa.rc.pstatus()
+
+    def _segment_callback(self, segment, children, history):
+        iiter = segment.n_iter - self.first_iter
+        seg_id = segment.seg_id
+        weights = np.empty((self.__pcoord_len,), np.float64)
+        weights[:] = segment.weight
+        bin_pops = self.bin_populations[iiter, :, :]
+
+        if len(history) == 0:
+            # New trajectory
+            self.n_trajs += 1
+            self.accumulator.start_accumulation(
+                self.bin_assignments[iiter, seg_id, :], weights, bin_pops, traj=self.n_trajs, n_iter=segment.n_iter
+            )
+        else:
+            # Continuing trajectory
+            self.accumulator.continue_accumulation(
+                self.bin_assignments[iiter, seg_id, :], weights, bin_pops, traj=self.n_trajs, n_iter=segment.n_iter
+            )
+
+        self.n_segs_visited += 1
+
+        if not self.__quiet_mode and (self.n_segs_visited % 1000 == 0 or self.n_segs_visited == self.n_total_segs):
+            pct_visited = self.n_segs_visited / self.n_total_segs * 100
+            westpa.rc.pstatus(
+                '\r  {:d} of {:d} segments ({:.1f}%) analyzed ({:d} independent trajectories)'.format(
+                    int(self.n_segs_visited), int(self.n_total_segs), float(pct_visited), self.n_trajs
+                ),
+                end='',
+            )
+            westpa.rc.pflush()
+
+
+class BFTransitionAnalysisMixin(TransitionAnalysisMixin):
+    def require_transitions(self):
+        self.require_bin_assignments()
+        self.require_transitions_group()
+        do_trans = False
+        if self.discard_transition_data:
+            westpa.rc.pstatus('Discarding existing transition data.')
+            do_trans = True
+        elif not self.check_data_binhash(self.trans_h5group):
+            westpa.rc.pstatus('Bin definitions have changed; deleting existing transition data.')
+            do_trans = True
+
+        if do_trans:
+            self.delete_transitions_group()
+            self.find_transitions()
+
+    def find_transitions(self, chunksize=65536):
+        self.require_bf_h5file()
+        self.require_binning_group()
+        westpa.rc.pstatus('Finding transitions...')
+        output_group = self.require_analysis_group('transitions')
+
+        self.accumulator = TransitionEventAccumulator(self.n_bins, output_group, calc_fpts=True)
+        assignments_ds = self.binning_h5group['bin_assignments']
+
+        max_nrows = assignments_ds.len()
+        maxwidth_nrows = len(str(max_nrows))
+
+        for traj_id in range(self.get_n_trajs()):
+            nrows = self.get_traj_len(traj_id)
+            for istart in range(0, nrows, chunksize):
+                iend = min(istart + chunksize, nrows)
+                assignments = assignments_ds[traj_id, istart:iend]
+                weights = np.ones((len(assignments),))
+                binpops = np.ones((len(assignments), self.n_bins))
+
+                if istart == 0:
+                    self.accumulator.start_accumulation(assignments, weights, binpops, traj=traj_id)
+                else:
+                    self.accumulator.continue_accumulation(assignments, weights, binpops, traj=traj_id)
+
+                westpa.rc.pstatus(
+                    '\r  Trajectory {:d}: {:{mwnr}d}/{:<{mwnr}d} ({:.2f}%)'.format(
+                        int(traj_id), int(iend), int(nrows), iend / nrows * 100, mwnr=maxwidth_nrows
+                    ),
+                    end='',
+                )
+                westpa.rc.pflush()
+                self.accumulator.flush_transition_data()
+                del assignments, weights, binpops
+            westpa.rc.pstatus()
+
+        try:
+            del output_group['n_trans']
+        except KeyError:
+            pass
+        output_group['n_trans'] = self.accumulator.n_trans
+
+        for h5object in (output_group, output_group['n_trans'], output_group['transitions']):
+            self.record_data_binhash(h5object)
+
+        self.accumulator.clear()