westpa 2022.10__cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl

westpa/core/kinetics/__init__.py

@@ -0,0 +1,24 @@
+'''
+Kinetics analysis library
+'''
+
+import logging
+
+from .rate_averaging import RateAverager  # noqa
+
+from . import _kinetics  # noqa
+from ._kinetics import (  # noqa
+    calculate_labeled_fluxes,
+    labeled_flux_to_rate,
+    calculate_labeled_fluxes_alllags,
+    nested_to_flat_matrix,
+    nested_to_flat_vector,
+    flat_to_nested_matrix,
+    flat_to_nested_vector,
+    find_macrostate_transitions,
+    sequence_macro_flux_to_rate,
+)
+from .events import WKinetics  # noqa
+
+
+log = logging.getLogger(__name__)

westpa/core/kinetics/events.py

@@ -0,0 +1,147 @@
+import warnings
+
+import numpy as np
+
+from westpa.core.data_manager import weight_dtype
+from westpa.core import h5io
+
+# From w_kinetics.
+from westpa.tools.dtypes import ed_list_dtype
+from westpa.core.binning import index_dtype
+from westpa.core.kinetics._kinetics import _fast_transition_state_copy  # @UnresolvedImport
+from westpa.core.kinetics import find_macrostate_transitions
+
+
+warnings.filterwarnings('ignore', category=DeprecationWarning)
+warnings.filterwarnings('ignore', category=RuntimeWarning)
+warnings.filterwarnings('ignore', category=FutureWarning)
+
+
+# The old w_kinetics
+class WKinetics:
+    def w_kinetics(self):
+        pi = self.progress.indicator
+        pi.new_operation('Initializing')
+
+        self.data_reader.open('r')
+        self.open_files()
+        nstates = self.assignments_file.attrs['nstates']
+        start_iter, stop_iter = self.iter_range.iter_start, self.iter_range.iter_stop  # h5io.get_iter_range(self.assignments_file)
+        iter_count = stop_iter - start_iter
+        durations_ds = self.output_file.replace_dataset(
+            'durations',
+            shape=(iter_count, 0),
+            maxshape=(iter_count, None),
+            dtype=ed_list_dtype,
+            chunks=(1, 15360) if self.do_compression else None,
+            shuffle=self.do_compression,
+            compression=9 if self.do_compression else None,
+        )
+        durations_count_ds = self.output_file.replace_dataset(
+            'duration_count', shape=(iter_count,), dtype=np.int_, shuffle=True, compression=9
+        )
+        cond_fluxes_ds = self.output_file.replace_dataset(
+            'conditional_fluxes',
+            shape=(iter_count, nstates, nstates),
+            dtype=weight_dtype,
+            chunks=(h5io.calc_chunksize((iter_count, nstates, nstates), weight_dtype) if self.do_compression else None),
+            shuffle=self.do_compression,
+            compression=9 if self.do_compression else None,
+        )
+        total_fluxes_ds = self.output_file.replace_dataset(
+            'total_fluxes',
+            shape=(iter_count, nstates),
+            dtype=weight_dtype,
+            chunks=(h5io.calc_chunksize((iter_count, nstates), weight_dtype) if self.do_compression else None),
+            shuffle=self.do_compression,
+            compression=9 if self.do_compression else None,
+        )
+
+        cond_arrival_counts_ds = self.output_file.replace_dataset(
+            'conditional_arrivals',
+            shape=(iter_count, nstates, nstates),
+            dtype=np.uint,
+            chunks=(h5io.calc_chunksize((iter_count, nstates, nstates), np.uint) if self.do_compression else None),
+            shuffle=self.do_compression,
+            compression=9 if self.do_compression else None,
+        )
+        arrival_counts_ds = self.output_file.replace_dataset(
+            'arrivals',
+            shape=(iter_count, nstates),
+            dtype=np.uint,
+            chunks=(h5io.calc_chunksize((iter_count, nstates), np.uint) if self.do_compression else None),
+            shuffle=self.do_compression,
+            compression=9 if self.do_compression else None,
+        )
+
+        # copy state labels for convenience
+        self.output_file.replace_dataset('state_labels', data=self.assignments_file['state_labels'][...])
+
+        # Put nice labels on things
+        for ds in (self.output_file, durations_count_ds, cond_fluxes_ds, total_fluxes_ds):
+            h5io.stamp_iter_range(ds, start_iter, stop_iter)
+
+        # Calculate instantaneous rate matrices and trace trajectories
+        last_state = None
+        pi.new_operation('Tracing trajectories', iter_count)
+        for iiter, n_iter in enumerate(range(start_iter, stop_iter)):
+            # Get data from the main HDF5 file
+            iter_group = self.data_reader.get_iter_group(n_iter)
+            seg_index = iter_group['seg_index']
+            nsegs, npts = iter_group['pcoord'].shape[0:2]
+            weights = seg_index['weight']
+            # parent_ids = seg_index['parent_id']
+            parent_ids = self.data_reader.parent_id_dsspec.get_iter_data(n_iter)
+
+            # Get bin and traj. ensemble assignments from the previously-generated assignments file
+            assignment_iiter = h5io.get_iteration_entry(self.assignments_file, n_iter)
+            bin_assignments = np.require(
+                self.assignments_file['assignments'][assignment_iiter + np.s_[:nsegs, :npts]], dtype=index_dtype
+            )
+            label_assignments = np.require(
+                self.assignments_file['trajlabels'][assignment_iiter + np.s_[:nsegs, :npts]], dtype=index_dtype
+            )
+            state_assignments = np.require(
+                self.assignments_file['statelabels'][assignment_iiter + np.s_[:nsegs, :npts]], dtype=index_dtype
+            )
+
+            # Prepare to run analysis
+            cond_fluxes = np.zeros((nstates, nstates), weight_dtype)
+            total_fluxes = np.zeros((nstates,), weight_dtype)
+            cond_counts = np.zeros((nstates, nstates), np.uint)
+            total_counts = np.zeros((nstates,), np.uint)
+            durations = []
+
+            # Estimate macrostate fluxes and calculate event durations using trajectory tracing
+            # state is opaque to the find_macrostate_transitions function
+            dt = 1.0 if npts == 1 else 1.0 / (npts - 1)
+            state = _fast_transition_state_copy(iiter, nstates, parent_ids, last_state)
+            find_macrostate_transitions(
+                nstates,
+                weights,
+                label_assignments,
+                state_assignments,
+                dt,
+                state,
+                cond_fluxes,
+                cond_counts,
+                total_fluxes,
+                total_counts,
+                durations,
+            )
+            last_state = state
+
+            # Store trace-based kinetics data
+            cond_fluxes_ds[iiter] = cond_fluxes
+            total_fluxes_ds[iiter] = total_fluxes
+            arrival_counts_ds[iiter] = total_counts
+            cond_arrival_counts_ds[iiter] = cond_counts
+
+            durations_count_ds[iiter] = len(durations)
+            if len(durations) > 0:
+                durations_ds.resize((iter_count, max(len(durations), durations_ds.shape[1])))
+                durations_ds[iiter, : len(durations)] = durations
+
+            # Do a little manual clean-up to prevent memory explosion
+            del iter_group, weights, parent_ids, bin_assignments, label_assignments, state, cond_fluxes, total_fluxes
+            pi.progress += 1

westpa/core/kinetics/matrates.py

@@ -0,0 +1,156 @@
+"""
+Routines for implementing Letteri et al.'s macrostate-to-macrostate rate calculations
+using extrapolation to steady-state populations from average rate matrices
+
+Internally, "labeled" objects (bin populations labeled by history, rate matrix elements labeled
+by history) are stored as nested arrays -- e.g. rates[initial_label, final_label, initial_bin, final_bin].
+These are converted to the flat forms required for, say, eigenvalue calculations internally, and the
+results converted back. This is because these conversions are not expensive, and saves users of
+this code from having to know how the flattened indexing works (something I screwed up all too
+easily during development) -- mcz
+"""
+
+import logging
+import warnings
+
+import numpy as np
+import scipy.linalg
+
+from westpa.core.data_manager import weight_dtype
+
+from ._kinetics import (
+    calculate_labeled_fluxes,  # @UnresolvedImport
+    calculate_labeled_fluxes_alllags,  # @UnresolvedImport
+    labeled_flux_to_rate,  # @UnresolvedImport
+    nested_to_flat_matrix,
+    nested_to_flat_vector,  # @UnresolvedImport
+    flat_to_nested_vector,
+    _reduce_labeled_rate_matrix_to_macro,
+)  # @UnresolvedImport
+
+
+log = logging.getLogger(__name__)
+
+
+class ConsistencyWarning(UserWarning):
+    pass
+
+
+def get_steady_state(rates):
+    '''Get steady state solution for a rate matrix. As an optimization, returns the
+    flattened labeled population vector (of length nstates*nbins); to convert to the
+    nested vector used for storage, use nested_to_flat_vector().'''
+
+    rates = rates.copy()
+
+    # Convert to a transition probability matrix
+    for i in range(rates.shape[0]):
+        rowsum = rates[i, :].sum()
+        if rowsum > 0:
+            rates[i, :] = rates[i, :] / rowsum
+        else:
+            if rates[:, i].sum() != 0:
+                warnings.warn('sink microstate in rate matrix', ConsistencyWarning)
+            rates[i, :] = 0
+
+    try:
+        vals, vecs = scipy.linalg.eig(rates.T)
+    except Exception:
+        log.debug('exception obtaining eigenvectors', exc_info=True)
+        return None
+
+    vals = np.abs(vals)
+    log.debug('eigenvalues: {!r}'.format(list(reversed(sorted(vals)))))
+    asort = np.argsort(vals)
+    vec = vecs[:, asort[-1]]
+    ss = np.abs(vec)
+
+    ss /= ss.sum()
+    return ss
+
+
+def get_macrostate_rates(labeled_rates, labeled_pops, extrapolate=True):
+    '''Using a labeled rate matrix and labeled bin populations, calculate the steady state
+    probability distribution and consequent state-to-state rates.
+
+    Returns ``(ss, macro_rates)``, where ``ss`` is the steady-state probability distribution
+    and ``macro_rates`` is the state-to-state rate matrix.'''
+
+    nstates, nbins = labeled_pops.shape
+
+    rates = nested_to_flat_matrix(labeled_rates)
+
+    # Find steady-state solution
+    if extrapolate:
+        ss = get_steady_state(rates)
+        if ss is None:
+            warnings.warn('no well-defined steady state; using average populations', ConsistencyWarning)
+            ss = nested_to_flat_vector(labeled_pops)
+    else:
+        ss = nested_to_flat_vector(labeled_pops)
+
+    macro_rates = _reduce_labeled_rate_matrix_to_macro(nstates, nbins, rates, ss)
+
+    return flat_to_nested_vector(nstates, nbins, ss), macro_rates
+
+
+def estimate_rates(
+    nbins,
+    state_labels,
+    weights,
+    parent_ids,
+    bin_assignments,
+    label_assignments,
+    state_map,
+    labeled_pops,
+    all_lags=False,
+    labeled_fluxes=None,
+    labeled_rates=None,
+    unlabeled_rates=None,
+):
+    '''Estimate fluxes and rates over multiple iterations. The number of iterations is determined by how many
+    vectors of weights, parent IDs, bin assignments, and label assignments are passed.
+
+    If ``all_lags`` is true, then the average is over all possible lags within the length-N window given, otherwise
+    simply the length N lag.
+
+    Returns labeled flux matrix, labeled rate matrix, and unlabeled rate matrix.'''
+
+    assert len(weights) == len(parent_ids) == len(bin_assignments) == len(label_assignments)
+    nstates = len(state_labels)
+    nbins = labeled_pops.shape[1] - 1
+
+    # Prepare output arrays
+    if labeled_fluxes is None:
+        labeled_fluxes = np.zeros((nstates, nstates, nbins, nbins), weight_dtype)
+    else:
+        labeled_fluxes.fill(0.0)
+
+    if labeled_rates is None:
+        labeled_rates = np.zeros_like(labeled_fluxes)
+    else:
+        labeled_rates.fill(0.0)
+
+    if unlabeled_rates is None:
+        unlabeled_rates = np.zeros((nbins, nbins), weight_dtype)
+    else:
+        unlabeled_rates.fill(0.0)
+
+    # Loop over all possible windows to accumulate flux matrix
+    # flux matrix is [initial_label][final_label][initial_bin][final_bin]
+    if all_lags:
+        twindow = calculate_labeled_fluxes_alllags(nstates, weights, parent_ids, bin_assignments, label_assignments, labeled_fluxes)
+    else:
+        twindow = calculate_labeled_fluxes(nstates, weights, parent_ids, bin_assignments, label_assignments, labeled_fluxes)
+    labeled_fluxes /= twindow
+
+    # Calculate rate matrix for this window, using populations from the last iteration (which correspond
+    # to the weights that contribute to the flux matrix)
+    labeled_flux_to_rate(labeled_fluxes, labeled_pops, labeled_rates)
+
+    # Calculate an unlabeled rate matrix
+    unlabeled_fluxes = np.sum(labeled_fluxes, axis=(0, 1))
+    unlabeled_pops = labeled_pops.sum(axis=0)
+    unlabeled_rates[...] = labeled_flux_to_rate(unlabeled_fluxes[None, None, :, :], unlabeled_pops[None, :])[0, 0]
+
+    return labeled_fluxes, labeled_rates, unlabeled_rates

westpa/core/kinetics/rate_averaging.py

@@ -0,0 +1,266 @@
+from collections import namedtuple
+from itertools import zip_longest
+
+import numpy as np
+
+import westpa
+from westpa.core.kinetics._kinetics import (
+    flux_assign,
+    pop_assign,
+    calc_rates,
+    StreamingStats1D,
+    StreamingStats2D,
+)  # @UnresolvedImport
+
+
+# Named tuple proxy for StreamingStats class
+StreamingStatsTuple = namedtuple('StreamingStatsTuple', ['M1', 'M2', 'n'])
+
+
+def grouper(n, iterable, fillvalue=None):
+    "Collect data into fixed-length chunks or blocks"
+    # grouper(3, 'ABCDEFG', 'x') --> ABC DEF Gxx
+    args = [iter(iterable)] * n
+    return zip_longest(fillvalue=fillvalue, *args)
+
+
+def tuple2stats(stat_tuple):
+    ndims = stat_tuple.M1.ndim
+    assert ndims == 1 or ndims == 2
+
+    if ndims == 2:
+        stats = StreamingStats2D(stat_tuple.M1.shape)
+    elif ndims == 1:
+        stats = StreamingStats1D(stat_tuple.M1.shape[0])
+    else:
+        raise ValueError
+
+    stats.M1 = stat_tuple.M1
+    stats.M2 = stat_tuple.M2
+    stats.n = stat_tuple.n
+
+    return stats
+
+
+def process_iter_chunk(bin_mapper, iter_indices, iter_data=None):
+    '''Calculate the flux matrices and populations of a set of iterations specified
+    by iter_indices. Optionally provide the necessary arrays to perform the calculation
+    in iter_data. Otherwise get data from the data_manager directly.
+    '''
+
+    data_manager = westpa.rc.get_data_manager()
+    system = westpa.rc.get_system_driver()
+
+    # itercount = len(iter_indices)
+    nbins = bin_mapper.nbins
+
+    flux_stats = StreamingStats2D((nbins, nbins))
+    rate_stats = StreamingStats2D((nbins, nbins))
+    pop_stats = StreamingStats1D(nbins)
+
+    nomask1d = np.zeros((nbins,), np.uint8)
+    nomask2d = np.zeros((nbins, nbins), np.uint8)
+    rate_mask = np.zeros((nbins, nbins), np.uint8)
+
+    flux_matrix = np.zeros((nbins, nbins), np.float64)
+    rate_matrix = np.zeros((nbins, nbins), np.float64)
+    population_vector = np.zeros((nbins,), np.float64)
+
+    pcoord_len = system.pcoord_len
+    assign = bin_mapper.assign
+
+    for iiter, n_iter in enumerate(iter_indices):
+        flux_matrix.fill(0.0)
+        population_vector.fill(0.0)
+
+        if iter_data:
+            iter_group_name = 'iter_{:0{prec}d}'.format(int(n_iter), prec=data_manager.iter_prec)
+            iter_group = iter_data[iter_group_name]
+        else:
+            iter_group = data_manager.get_iter_group(n_iter)
+
+        # first, account for the flux due to recycling
+        # we access the hdf5 file directly to avoid nearly 50% overhead of creating a ton of
+        # tiny newweightentry objects
+        try:
+            nwgroup = iter_group['new_weights']
+        except KeyError:
+            # no new weight data
+            pass
+        else:
+            if iter_data:
+                index = None
+                weights = nwgroup['weight']
+                prev_init_pcoords = nwgroup['prev_init_pcoord']
+                new_init_pcoords = nwgroup['new_init_pcoord']
+            else:
+                index = nwgroup['index'][...]
+                weights = index['weight']
+                prev_init_pcoords = nwgroup['prev_init_pcoord'][...]
+                new_init_pcoords = nwgroup['new_init_pcoord'][...]
+
+            prev_init_assignments = assign(prev_init_pcoords)
+            new_init_assignments = assign(new_init_pcoords)
+
+            flux_assign(weights, prev_init_assignments, new_init_assignments, flux_matrix)
+            # for (weight,i,j) in izip (weights, prev_init_assignments, new_init_assignments):
+            #    flux_matrices[iiter,i,j] += weight
+            del index
+            del prev_init_pcoords, new_init_pcoords, prev_init_assignments, new_init_assignments, weights
+
+        # iter_group = data_manager.get_iter_group(n_iter)
+        if iter_data:
+            weights = iter_group['weight']
+            initial_pcoords = iter_group['initial_pcoords']
+            final_pcoords = iter_group['final_pcoords']
+        else:
+            weights = iter_group['seg_index']['weight']
+            initial_pcoords = iter_group['pcoord'][:, 0]
+            final_pcoords = iter_group['pcoord'][:, pcoord_len - 1]
+
+        initial_assignments = assign(initial_pcoords)
+        final_assignments = assign(final_pcoords)
+
+        flux_assign(weights, initial_assignments, final_assignments, flux_matrix)
+        pop_assign(weights, initial_assignments, population_vector)
+
+        flux_stats.update(flux_matrix, nomask2d)
+        pop_stats.update(population_vector, nomask1d)
+
+        calc_rates(flux_matrix, population_vector, rate_matrix, rate_mask)
+        rate_stats.update(rate_matrix, rate_mask)
+
+        del weights
+        del initial_assignments, final_assignments
+        del initial_pcoords, final_pcoords
+        del iter_group
+
+    # Create namedtuple proxies for the cython StreamingStats objects
+    # since the typed memoryviews class variables do not seem to return
+    # cleanly from the zmq workers
+    c_flux_stats = StreamingStatsTuple(flux_stats.M1, flux_stats.M2, flux_stats.n)
+    c_rate_stats = StreamingStatsTuple(rate_stats.M1, rate_stats.M2, rate_stats.n)
+    c_pop_stats = StreamingStatsTuple(pop_stats.M1, pop_stats.M2, pop_stats.n)
+
+    return c_flux_stats, c_rate_stats, c_pop_stats
+
+
+class RateAverager:
+    '''Calculate bin-to-bin kinetic properties (fluxes, rates, populations) at
+    1-tau resolution'''
+
+    def __init__(self, bin_mapper, system=None, data_manager=None, work_manager=None):
+        self.bin_mapper = bin_mapper
+        self.data_manager = data_manager or westpa.rc.get_data_manager()
+        self.system = system or westpa.rc.get_system_driver()
+        self.work_manager = work_manager or westpa.rc.get_work_manager()
+
+    def extract_data(self, iter_indices):
+        '''Extract data from the data_manger and place in dict mirroring the same
+        underlying layout.'''
+
+        data = {}
+        pcoord_len = self.system.pcoord_len
+
+        for n_iter in iter_indices:
+            iter_group_name = 'iter_{:0{prec}d}'.format(int(n_iter), prec=self.data_manager.iter_prec)
+            iter_group = self.data_manager.get_iter_group(n_iter)
+            di = data[iter_group_name] = {}
+
+            try:
+                nwgroup = iter_group['new_weights']
+            except KeyError:
+                # no new weight data
+                pass
+            else:
+                di_nw = di['new_weights'] = {}
+                di_nw['weight'] = nwgroup['index'][...]['weight']
+                di_nw['prev_init_pcoord'] = nwgroup['prev_init_pcoord'][...]
+                di_nw['new_init_pcoord'] = nwgroup['new_init_pcoord'][...]
+
+            di['weight'] = iter_group['seg_index']['weight']
+            di['initial_pcoords'] = iter_group['pcoord'][:, 0]
+            di['final_pcoords'] = iter_group['pcoord'][:, pcoord_len - 1]
+
+        return data
+
+    def task_generator(self, iter_start, iter_stop, block_size):
+        for iter_block in grouper(block_size, range(iter_start, iter_stop)):
+            iter_block = [x for x in iter_block if x is not None]
+            iter_data = self.extract_data(iter_block)
+            yield (process_iter_chunk, (self.bin_mapper, iter_block), {'iter_data': iter_data})
+            del iter_data
+
+    def calculate(self, iter_start=None, iter_stop=None, n_blocks=1, queue_size=1):
+        '''Read the HDF5 file and collect flux matrices and population vectors
+        for each bin for each iteration in the range [iter_start, iter_stop). Break
+        the calculation into n_blocks blocks. If the calculation is broken up into
+        more than one block, queue_size specifies the maxmimum number of tasks in
+        the work queue.'''
+
+        iter_start = iter_start or 1
+        iter_stop = iter_stop or self.data_manager.current_iteration
+
+        itercount = iter_stop - iter_start
+        block_size = max(1, itercount // n_blocks)
+        nbins = self.bin_mapper.nbins
+
+        if n_blocks == 1:
+            flux_stats_t, rate_stats_t, population_stats_t = process_iter_chunk(self.bin_mapper, list(range(iter_start, iter_stop)))
+
+            flux_stats = tuple2stats(flux_stats_t)
+            rate_stats = tuple2stats(rate_stats_t)
+            population_stats = tuple2stats(population_stats_t)
+        else:
+            flux_stats = StreamingStats2D((nbins, nbins))
+            rate_stats = StreamingStats2D((nbins, nbins))
+            population_stats = StreamingStats1D(nbins)
+
+            task_generator = self.task_generator(iter_start, iter_stop, block_size)
+
+            for future in self.work_manager.submit_as_completed(task_generator, queue_size):
+                chunk_flux_stats_t, chunk_rate_stats_t, chunk_pop_stats_t = future.get_result()
+
+                chunk_flux_stats = tuple2stats(chunk_flux_stats_t)
+                chunk_rate_stats = tuple2stats(chunk_rate_stats_t)
+                chunk_pop_stats = tuple2stats(chunk_pop_stats_t)
+
+                # Update statistics with chunked subsets
+                flux_stats += chunk_flux_stats
+                rate_stats += chunk_rate_stats
+                population_stats += chunk_pop_stats
+
+        self.average_flux = flux_stats.mean
+        self.stderr_flux = np.nan_to_num(np.sqrt(flux_stats.var) / flux_stats.n)
+
+        self.average_populations = population_stats.mean
+        self.stderr_populations = np.nan_to_num(np.sqrt(population_stats.var) / population_stats.n)
+
+        self.average_rate = rate_stats.mean
+        self.stderr_rate = np.nan_to_num(np.sqrt(rate_stats.var) / rate_stats.n)
+
+        assert ~np.any(np.isinf(self.stderr_flux))
+        assert ~np.any(np.isinf(self.stderr_rate))
+        assert ~np.any(np.isinf(self.stderr_populations))
+
+
+if __name__ == '__main__':
+    # Tests this file on the west.h5 data in the current (sim root) directory
+    westpa.rc.read_config()
+    system = westpa.rc.get_system_driver()
+    data_manager = westpa.rc.get_data_manager()
+    data_manager.open_backing('r')
+    averager = RateAverager(system.bin_mapper)
+    averager.calculate()
+
+    print('Population mean and standard error')
+    print(averager.average_populations)
+    print(averager.stderr_populations)
+
+    print('Flux matrix, mean and standard error')
+    print(averager.average_flux)
+    print(averager.stderr_flux)
+
+    print('Rate matrix, mean and standard error')
+    print(averager.average_rate)
+    print(averager.stderr_rate)