westpa 2022.12__cp312-cp312-macosx_10_13_x86_64.whl

westpa/core/reweight/__init__.py

@@ -0,0 +1,14 @@
+'''
+Function(s) for the postanalysis toolkit
+'''
+
+import logging
+
+from . import _reweight
+from ._reweight import stats_process, reweight_for_c
+from .matrix import FluxMatrix
+
+__all__ = ['_reweight', 'stats_process', 'reweight_for_c', 'FluxMatrix']
+
+
+log = logging.getLogger(__name__)

westpa/core/reweight/matrix.py

@@ -0,0 +1,126 @@
+import warnings
+
+import numpy as np
+import scipy.sparse as sp
+
+from westpa.core.data_manager import weight_dtype
+from westpa.core import h5io
+
+# From postanalysis matrix
+from westpa.core.binning import index_dtype
+from westpa.core.reweight import stats_process
+
+
+warnings.filterwarnings('ignore', category=DeprecationWarning)
+warnings.filterwarnings('ignore', category=RuntimeWarning)
+warnings.filterwarnings('ignore', category=FutureWarning)
+
+
+def calc_stats(bin_assignments, weights, fluxes, populations, trans, mask, sampling_frequency):
+    fluxes.fill(0.0)
+    populations.fill(0.0)
+    trans.fill(0)
+
+    stats_process(bin_assignments, weights, fluxes, populations, trans, mask, interval=sampling_frequency)
+
+
+class FluxMatrix:
+    def w_postanalysis_matrix(self):
+        pi = self.progress.indicator
+        pi.new_operation('Initializing')
+
+        self.data_reader.open('r')
+        nbins = self.assignments_file.attrs['nbins']
+
+        state_labels = self.assignments_file['state_labels'][...]
+        # state_map = self.assignments_file['state_map'][...]
+        nstates = len(state_labels)
+
+        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
+
+        nfbins = nbins * nstates
+
+        flux_shape = (iter_count, nfbins, nfbins)
+        pop_shape = (iter_count, nfbins)
+
+        h5io.stamp_iter_range(self.output_file, start_iter, stop_iter)
+
+        bin_populations_ds = self.output_file.create_dataset('bin_populations', shape=pop_shape, dtype=weight_dtype)
+        h5io.stamp_iter_range(bin_populations_ds, start_iter, stop_iter)
+        h5io.label_axes(bin_populations_ds, ['iteration', 'bin'])
+
+        flux_grp = self.output_file.create_group('iterations')
+        self.output_file.attrs['nrows'] = nfbins
+        self.output_file.attrs['ncols'] = nfbins
+
+        fluxes = np.empty(flux_shape[1:], weight_dtype)
+        populations = np.empty(pop_shape[1:], weight_dtype)
+        trans = np.empty(flux_shape[1:], np.int64)
+
+        # Check to make sure this isn't a data set with target states
+        # tstates = self.data_reader.data_manager.get_target_states(0)
+        # if len(tstates) > 0:
+        #    raise ValueError('Postanalysis reweighting analysis does not support WE simulation run under recycling conditions')
+
+        pi.new_operation('Calculating flux matrices', iter_count)
+        # Calculate instantaneous statistics
+        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']
+
+            # 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
+            )
+
+            mask_unknown = np.zeros_like(bin_assignments, dtype=np.uint16)
+
+            macrostate_iiter = h5io.get_iteration_entry(self.assignments_file, n_iter)
+            macrostate_assignments = np.require(
+                self.assignments_file['trajlabels'][macrostate_iiter + np.s_[:nsegs, :npts]], dtype=index_dtype
+            )
+
+            # Transform bin_assignments to take macrostate membership into account
+            bin_assignments = nstates * bin_assignments + macrostate_assignments
+
+            mask_indx = np.where(macrostate_assignments == nstates)
+            mask_unknown[mask_indx] = 1
+
+            # Calculate bin-to-bin fluxes, bin populations and number of obs transitions
+            calc_stats(bin_assignments, weights, fluxes, populations, trans, mask_unknown, self.sampling_frequency)
+
+            # Store bin-based kinetics data
+            bin_populations_ds[iiter] = populations
+
+            # Setup sparse data structures for flux and obs
+            fluxes_sp = sp.coo_matrix(fluxes)
+            trans_sp = sp.coo_matrix(trans)
+
+            assert fluxes_sp.nnz == trans_sp.nnz
+
+            flux_iter_grp = flux_grp.create_group('iter_{:08d}'.format(n_iter))
+            flux_iter_grp.create_dataset('flux', data=fluxes_sp.data, dtype=weight_dtype)
+            flux_iter_grp.create_dataset('obs', data=trans_sp.data, dtype=np.int32)
+            flux_iter_grp.create_dataset('rows', data=fluxes_sp.row, dtype=np.int32)
+            flux_iter_grp.create_dataset('cols', data=fluxes_sp.col, dtype=np.int32)
+            flux_iter_grp.attrs['nrows'] = nfbins
+            flux_iter_grp.attrs['ncols'] = nfbins
+
+            # Do a little manual clean-up to prevent memory explosion
+            del iter_group, weights, bin_assignments
+            del macrostate_assignments
+
+            pi.progress += 1
+
+            # Check and save the number of intermediate time points; this will be used to normalize the
+            # flux and kinetics to tau in w_postanalysis_reweight.
+            if self.assignments_file.attrs['subsampled'] is True or self.sampling_frequency == 'iteration':
+                self.output_file.attrs['npts'] = 2
+            else:
+                # self.output_file.attrs['npts'] = npts if self.sampling_frequency == 'timepoint' else 2
+                self.output_file.attrs['npts'] = npts

westpa/core/segment.py

@@ -0,0 +1,119 @@
+import math
+
+import numpy as np
+
+
+class Segment:
+    '''A class wrapping segment data that must be passed through the work manager or data manager.
+    Most fields are self-explanatory.  One item worth noting is that a negative parent ID means that
+    the segment starts from the initial state with ID -(segment.parent_id+1)
+    '''
+
+    SEG_STATUS_UNSET = 0
+    SEG_STATUS_PREPARED = 1
+    SEG_STATUS_COMPLETE = 2
+    SEG_STATUS_FAILED = 3
+
+    SEG_INITPOINT_UNSET = 0
+    SEG_INITPOINT_CONTINUES = 1
+    SEG_INITPOINT_NEWTRAJ = 2
+
+    SEG_ENDPOINT_UNSET = 0
+    SEG_ENDPOINT_CONTINUES = 1
+    SEG_ENDPOINT_MERGED = 2
+    SEG_ENDPOINT_RECYCLED = 3
+
+    statuses = {}
+    initpoint_types = {}
+    endpoint_types = {}
+
+    status_names = {}
+    initpoint_type_names = {}
+    endpoint_type_names = {}
+
+    # convenience functions for binning
+    @staticmethod
+    def initial_pcoord(segment):
+        'Return the initial progress coordinate point of this segment.'
+        return segment.pcoord[0]
+
+    @staticmethod
+    def final_pcoord(segment):
+        'Return the final progress coordinate point of this segment.'
+        return segment.pcoord[-1]
+
+    def __init__(
+        self,
+        n_iter=None,
+        seg_id=None,
+        weight=None,
+        endpoint_type=None,
+        parent_id=None,
+        wtg_parent_ids=None,
+        pcoord=None,
+        status=None,
+        walltime=None,
+        cputime=None,
+        data=None,
+    ):
+        # NaNs appear sometimes if a WEST program is terminated unexpectedly; replace with zero
+        walltime = 0.0 if walltime is None or math.isnan(walltime) else walltime
+        cputime = 0.0 if cputime is None or math.isnan(cputime) else cputime
+
+        # the int() and float() calls are required so that new-style string formatting doesn't barf
+        # assuming that the respective fields are actually strings, probably after implicitly
+        # calling __str__() on them.  Not sure if this is a numpy, h5py, or python problem
+        self.n_iter = int(n_iter) if n_iter is not None else None
+        self.seg_id = int(seg_id) if seg_id is not None else None
+        self.status = int(status) if status is not None else None
+        self.parent_id = int(parent_id) if parent_id is not None else None
+        self.endpoint_type = int(endpoint_type) if endpoint_type else self.SEG_ENDPOINT_UNSET
+
+        self.weight = float(weight) if weight is not None else None
+        self.wtg_parent_ids = set(wtg_parent_ids or ())
+
+        self.pcoord = np.asarray(pcoord) if pcoord is not None else None
+        self.walltime = walltime
+        self.cputime = cputime
+        self.data = data if data else {}
+
+    def __repr__(self):
+        return '<%s(%s) n_iter=%r seg_id=%r weight=%r parent_id=%r wtg_parent_ids=%r pcoord[0]=%r pcoord[-1]=%r>' % (
+            self.__class__.__name__,
+            hex(id(self)),
+            self.n_iter,
+            self.seg_id,
+            self.weight,
+            self.parent_id,
+            tuple(self.wtg_parent_ids or ()),
+            self.pcoord[0] if self.pcoord is not None else None,
+            self.pcoord[-1] if self.pcoord is not None else None,
+        )
+
+    @property
+    def initpoint_type(self):
+        if self.parent_id < 0:
+            return Segment.SEG_INITPOINT_NEWTRAJ
+        else:
+            return Segment.SEG_INITPOINT_CONTINUES
+
+    @property
+    def initial_state_id(self):
+        if self.parent_id < 0:
+            return -(self.parent_id + 1)
+        else:
+            return None
+
+    status_text = property((lambda s: s.status_names[s.status]))
+    endpoint_type_text = property((lambda s: s.endpoint_type_names[s.endpoint_type]))
+
+
+Segment.statuses.update({_attr: getattr(Segment, _attr) for _attr in dir(Segment) if _attr.startswith('SEG_STATUS_')})
+Segment.initpoint_types.update({_attr: getattr(Segment, _attr) for _attr in dir(Segment) if _attr.startswith('SEG_INITPOINT_')})
+Segment.endpoint_types.update({_attr: getattr(Segment, _attr) for _attr in dir(Segment) if _attr.startswith('SEG_ENDPOINT_')})
+
+Segment.status_names.update({getattr(Segment, _attr): _attr for _attr in dir(Segment) if _attr.startswith('SEG_STATUS_')})
+Segment.initpoint_type_names.update(
+    {getattr(Segment, _attr): _attr for _attr in dir(Segment) if _attr.startswith('SEG_INITPOINT_')}
+)
+Segment.endpoint_type_names.update({getattr(Segment, _attr): _attr for _attr in dir(Segment) if _attr.startswith('SEG_ENDPOINT_')})