westpa 2022.13__cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl

westpa/westext/weed/weed_driver.py

@@ -0,0 +1,192 @@
+import logging
+import operator
+
+import numpy as np
+
+import westpa
+from westpa.core.yamlcfg import check_bool
+from westpa.core.kinetics import RateAverager
+from westpa.westext.weed.ProbAdjustEquil import probAdjustEquil
+from westpa.core._rc import bins_from_yaml_dict
+
+EPS = np.finfo(np.float64).eps
+
+log = logging.getLogger(__name__)
+
+
+class WEEDDriver:
+    def __init__(self, sim_manager, plugin_config):
+        if not sim_manager.work_manager.is_master:
+            return
+
+        self.sim_manager = sim_manager
+        self.data_manager = sim_manager.data_manager
+        self.system = sim_manager.system
+        self.work_manager = sim_manager.work_manager
+
+        self.do_reweight = check_bool(
+            plugin_config.get('do_equilibrium_reweighting', False) or plugin_config.get('do_reweighting', False)
+        )
+        self.windowsize = 0.5
+        self.windowtype = 'fraction'
+
+        windowsize = plugin_config.get('window_size')
+        if windowsize is not None:
+            if isinstance(windowsize, float):
+                self.windowsize = windowsize
+                self.windowtype = 'fraction'
+                if self.windowsize <= 0 or self.windowsize > 1:
+                    raise ValueError('WEED parameter error -- fractional window size must be in (0,1]')
+            elif isinstance(windowsize, int):
+                self.windowsize = int(windowsize)
+                self.windowtype = 'fixed'
+            else:
+                raise ValueError('WEED parameter error -- invalid window size {!r}'.format(windowsize))
+        log.info('using window size of {!r} ({})'.format(self.windowsize, self.windowtype))
+
+        self.max_windowsize = plugin_config.get('max_window_size')
+        if self.max_windowsize is not None:
+            log.info('Using max windowsize of {:d}'.format(self.max_windowsize))
+
+        self.reweight_period = plugin_config.get('reweight_period', 0)
+        self.priority = plugin_config.get('priority', 0)
+
+        self.rate_calc_queue_size = plugin_config.get('rate_calc_queue_size', 1)
+        self.rate_calc_n_blocks = plugin_config.get('rate_calc_n_blocks', 1)
+
+        bin_obj = plugin_config.get('bins', None)
+        if isinstance(bin_obj, dict):
+            bin_obj = bins_from_yaml_dict(bin_obj)
+        self.bin_mapper = bin_obj
+
+        if self.do_reweight:
+            sim_manager.register_callback(sim_manager.prepare_new_iteration, self.prepare_new_iteration, self.priority)
+
+    def get_rates(self, n_iter, mapper):
+        '''Get rates and associated uncertainties as of n_iter, according to the window size the user
+        has selected (self.windowsize)'''
+
+        if self.windowtype == 'fraction':
+            if self.max_windowsize is not None:
+                eff_windowsize = min(self.max_windowsize, int(n_iter * self.windowsize))
+            else:
+                eff_windowsize = int(n_iter * self.windowsize)
+        else:  # self.windowtype == 'fixed':
+            eff_windowsize = min(n_iter, self.windowsize or 0)
+
+        averager = RateAverager(mapper, self.system, self.data_manager, self.work_manager)
+        averager.calculate(max(1, n_iter - eff_windowsize), n_iter + 1, self.rate_calc_n_blocks, self.rate_calc_queue_size)
+        self.eff_windowsize = eff_windowsize
+
+        return averager
+
+    def prepare_new_iteration(self):
+        n_iter = self.sim_manager.n_iter
+        we_driver = self.sim_manager.we_driver
+
+        if we_driver.target_states and self.do_reweight:
+            log.warning('equilibrium reweighting requested but target states (sinks) present; reweighting disabled')
+            return
+
+        if not self.do_reweight:
+            # Reweighting not requested
+            log.debug('Equilibrium reweighting not enabled')
+            return
+
+        with self.data_manager.lock:
+            weed_global_group = self.data_manager.we_h5file.require_group('weed')
+            reweighting_history_dataset = weed_global_group.require_dataset(
+                'reweighting_history', (1,), maxshape=(None,), dtype=int
+            )
+            last_reweighting = int(weed_global_group.attrs.get('last_reweighting', 0))
+            if last_reweighting > n_iter:
+                last_reweighting = n_iter - 1
+                reweighting_history = reweighting_history_dataset[:]
+                reweighting_history = reweighting_history[reweighting_history < n_iter]
+                reweighting_history_dataset.resize((reweighting_history.size), axis=0)
+
+        if n_iter - last_reweighting < self.reweight_period:
+            # Not time to reweight yet
+            log.debug('not reweighting')
+            return
+        else:
+            log.debug('reweighting')
+
+        if self.bin_mapper is None:
+            mapper = we_driver.bin_mapper
+            bins = we_driver.next_iter_binning
+            westpa.rc.pstatus('\nReweighting using the simulation bin mapper:\n{}'.format(mapper))
+        else:
+            mapper = self.bin_mapper
+            bins = mapper.construct_bins()
+
+            segments = [s for s in we_driver.next_iter_segments]
+            pcoords = self.system.new_pcoord_array(len(segments))
+            for iseg, segment in enumerate(segments):
+                pcoords[iseg] = segment.pcoord[0]
+            assignments = mapper.assign(pcoords)
+            for segment, assignment in zip(segments, assignments):
+                bins[assignment].add(segment)
+
+            westpa.rc.pstatus('\nReweighting using a different bin mapper than simulation:\n{}'.format(mapper))
+
+        n_bins = len(bins)
+
+        # Create storage for ourselves
+        with self.data_manager.lock:
+            iter_group = self.data_manager.get_iter_group(n_iter)
+            try:
+                del iter_group['weed']
+            except KeyError:
+                pass
+
+            weed_iter_group = iter_group.create_group('weed')
+            avg_populations_ds = weed_iter_group.create_dataset('avg_populations', shape=(n_bins,), dtype=np.float64)
+            unc_populations_ds = weed_iter_group.create_dataset('unc_populations', shape=(n_bins,), dtype=np.float64)
+            avg_flux_ds = weed_iter_group.create_dataset('avg_fluxes', shape=(n_bins, n_bins), dtype=np.float64)
+            unc_flux_ds = weed_iter_group.create_dataset('unc_fluxes', shape=(n_bins, n_bins), dtype=np.float64)
+            avg_rates_ds = weed_iter_group.create_dataset('avg_rates', shape=(n_bins, n_bins), dtype=np.float64)
+            unc_rates_ds = weed_iter_group.create_dataset('unc_rates', shape=(n_bins, n_bins), dtype=np.float64)
+
+        averager = self.get_rates(n_iter, mapper)
+
+        with self.data_manager.flushing_lock():
+            avg_populations_ds[...] = averager.average_populations
+            unc_populations_ds[...] = averager.stderr_populations
+            avg_flux_ds[...] = averager.average_flux
+            unc_flux_ds[...] = averager.stderr_flux
+            avg_rates_ds[...] = averager.average_rate
+            unc_rates_ds[...] = averager.stderr_rate
+
+            binprobs = np.fromiter(map(operator.attrgetter('weight'), bins), dtype=np.float64, count=n_bins)
+            orig_binprobs = binprobs.copy()
+
+        westpa.rc.pstatus('Calculating equilibrium reweighting using window size of {:d}'.format(self.eff_windowsize))
+        westpa.rc.pstatus('\nBin probabilities prior to reweighting:\n{!s}'.format(binprobs))
+        westpa.rc.pflush()
+
+        probAdjustEquil(binprobs, averager.average_rate, averager.stderr_rate)
+
+        # Check to see if reweighting has set non-zero bins to zero probability (should never happen)
+        assert (~((orig_binprobs > 0) & (binprobs == 0))).all(), 'populated bin reweighted to zero probability'
+
+        # Check to see if reweighting has set zero bins to nonzero probability (may happen)
+        z2nz_mask = (orig_binprobs == 0) & (binprobs > 0)
+        if (z2nz_mask).any():
+            westpa.rc.pstatus('Reweighting would assign nonzero probability to an empty bin; not reweighting this iteration.')
+            westpa.rc.pstatus('Empty bins assigned nonzero probability: {!s}.'.format(np.array_str(np.arange(n_bins)[z2nz_mask])))
+        else:
+            westpa.rc.pstatus('\nBin populations after reweighting:\n{!s}'.format(binprobs))
+            for bin, newprob in zip(bins, binprobs):
+                bin.reweight(newprob)
+
+            reweighting_history_dataset.resize((reweighting_history_dataset.shape[0] + 1), axis=0)
+            reweighting_history_dataset[-1] = n_iter
+            weed_global_group.attrs['last_reweighting'] = n_iter
+
+        assert (
+            abs(1 - np.fromiter(map(operator.attrgetter('weight'), bins), dtype=np.float64, count=n_bins).sum())
+            < EPS * np.fromiter(map(len, bins), dtype=int, count=n_bins).sum()
+        )
+
+        westpa.rc.pflush()

westpa/westext/wess/ProbAdjust.py

@@ -0,0 +1,101 @@
+import numpy as np
+import scipy.optimize
+
+
+def solve_steady_state(T, U, target_bins_index):
+    ntarget = len(target_bins_index)  # Number of target states
+    nstates = T.shape[0]  # Number of total states
+
+    # Number of active states
+    nsolve = nstates - ntarget
+
+    # list of active states
+    sactive = sorted(list(set(range(nstates)) - set(target_bins_index)))
+
+    W = np.zeros((nsolve, nsolve))
+    W_unc = np.zeros((nsolve, nsolve))
+    S = np.zeros((nsolve,))
+
+    for iiw, iit in zip(range(nsolve), sactive):
+        for jjw, jjt in zip(range(nsolve), sactive):
+            W[iiw, jjw] = T[jjt, iit]
+            W_unc[iiw, jjw] = U[jjt, iit]
+
+    for ii in range(nsolve):
+        W[ii, ii] = 0.0
+        for jj in range(nstates):  # nstates
+            if jj != ii:
+                W[ii, ii] -= T[ii, jj]
+
+    # we have n equations, but only n-1 linearly independent equations
+    # set one equation to constraint sum(Pi)=1
+
+    # Pick row that contains the largest uncertainty to replace
+    ii = np.unravel_index(np.argmax(W_unc), W_unc.shape)[1]
+
+    S[ii] = 1.0
+    W[ii, :] = 1.0
+
+    # P = np.linalg.solve(W,S)
+    try:
+        P, err = scipy.optimize.nnls(W, S)
+    except RuntimeError:
+        print('Solve did not converge')
+        return None
+
+    # There are some instances where a single bin has its prob set
+    # to 1.0; In this case do not reweight
+    if np.allclose(np.max(P), 1.0):
+        return None
+
+    pp = np.matrix(W) * np.matrix(P.reshape(-1, 1))
+
+    print('WESS: M*p = {}'.format(pp))
+    print('WESS: max = {}'.format(np.max(pp)))
+
+    return P, err, sactive
+
+
+def prob_adjust(binprob, rates, uncert, oldindex, targets=[]):
+    nbins = binprob.size
+
+    result = solve_steady_state(rates, uncert, targets)
+
+    if result is None:
+        return binprob
+    else:
+        ss_estimate, err, active_bins = result
+
+    print('WESS NNLS norm: {}'.format(err))
+
+    # now remap new_weights onto the bin indices
+    MappedActiveBins = []
+    for ibin in active_bins:
+        MappedActiveBins.append(oldindex[ibin])
+
+    mapped_new_weights = np.zeros((nbins,))
+    mapped_new_weights[MappedActiveBins] = ss_estimate
+
+    # Check to make sure no bins with non-zero probability end up with zero prob after reweighting
+    orig_nzi = np.nonzero(binprob)[0]  # bins that originally had non-zero prob
+    new_zi = np.where(mapped_new_weights == 0.0)[0]  # bins that after reweight have zero prob
+    ri = np.intersect1d(orig_nzi, new_zi)  # bins to reset
+    mapped_new_weights[ri] = 1.0e-16
+
+    # Check to make sure no bins with zero probability end up with nonzero prob after reweighting
+    orig_zi = np.where(binprob == 0.0)[0]  # bins that originally had zero prob
+    new_nzi = np.nonzero(mapped_new_weights)[0]  # bins that after reweight have non-zero prob
+    ri = np.intersect1d(orig_zi, new_nzi)  # bins to reset
+    mapped_new_weights[ri] = 0.0
+
+    # Ensure that all target bins have their weight set to zero
+    MappedTargetBins = []
+    for ibin in targets:
+        MappedTargetBins.append(oldindex[ibin])
+    mapped_new_weights[MappedTargetBins] = 0.0
+
+    mapped_new_weights /= np.sum(mapped_new_weights)
+
+    assert not (mapped_new_weights < 0).any()
+
+    return mapped_new_weights

westpa/westext/wess/__init__.py

@@ -0,0 +1,6 @@
+from . import ProbAdjust
+from .ProbAdjust import prob_adjust
+from .wess_driver import WESSDriver
+
+
+__all__ = ['ProbAdjust', 'prob_adjust', 'WESSDriver']

westpa/westext/wess/wess_driver.py

@@ -0,0 +1,217 @@
+import logging
+import operator
+
+import numpy as np
+
+import westpa
+from westpa.core.yamlcfg import check_bool
+from westpa.core.kinetics import RateAverager
+from westpa.westext.wess.ProbAdjust import prob_adjust
+from westpa.core._rc import bins_from_yaml_dict
+
+EPS = np.finfo(np.float64).eps
+
+log = logging.getLogger(__name__)
+
+
+def reduce_array(Aij):
+    """Remove empty rows and columns from an array Aij and return the reduced
+    array Bij and the list of non-empty states"""
+
+    nonempty = list(range(0, Aij.shape[0]))
+    eps = np.finfo(Aij.dtype).eps
+
+    for i in range(0, Aij.shape[0]):
+        if (Aij[i, :] < eps).all() and (Aij[:, i] < eps).all():
+            nonempty.pop(nonempty.index(i))
+
+    nne = len(nonempty)
+    Bij = np.zeros((nne, nne))
+
+    for i in range(0, nne):
+        for j in range(0, nne):
+            Bij[i, j] = Aij[nonempty[i], nonempty[j]]
+
+    return Bij, nonempty
+
+
+class WESSDriver:
+    def __init__(self, sim_manager, plugin_config):
+        if not sim_manager.work_manager.is_master:
+            return
+
+        self.sim_manager = sim_manager
+        self.data_manager = sim_manager.data_manager
+        self.system = sim_manager.system
+        self.work_manager = sim_manager.work_manager
+
+        self.do_reweight = check_bool(plugin_config.get('do_reweighting', False))
+        self.windowsize = 0.5
+        self.windowtype = 'fraction'
+
+        windowsize = plugin_config.get('window_size')
+        if windowsize is not None:
+            if isinstance(windowsize, float):
+                self.windowsize = windowsize
+                self.windowtype = 'fraction'
+                if self.windowsize <= 0 or self.windowsize > 1:
+                    raise ValueError('WESS parameter error -- fractional window size must be in (0,1]')
+            elif isinstance(windowsize, int):
+                self.windowsize = int(windowsize)
+                self.windowtype = 'fixed'
+            else:
+                raise ValueError('WESS parameter error -- invalid window size {!r}'.format(windowsize))
+        log.info('using window size of {!r} ({})'.format(self.windowsize, self.windowtype))
+
+        self.max_windowsize = plugin_config.get('max_window_size')
+        if self.max_windowsize is not None:
+            log.info('Using max windowsize of {:d}'.format(self.max_windowsize))
+
+        self.reweight_period = plugin_config.get('reweight_period', 0)
+        self.priority = plugin_config.get('priority', 0)
+
+        self.rate_calc_queue_size = plugin_config.get('rate_calc_queue_size', 1)
+        self.rate_calc_n_blocks = plugin_config.get('rate_calc_n_blocks', 1)
+
+        bin_obj = plugin_config.get('bins', None)
+        if isinstance(bin_obj, dict):
+            bin_obj = bins_from_yaml_dict(bin_obj)
+        self.bin_mapper = bin_obj
+
+        if self.do_reweight:
+            sim_manager.register_callback(sim_manager.prepare_new_iteration, self.prepare_new_iteration, self.priority)
+
+        self.write_matrices = plugin_config.get('write_matrices', False)
+
+    def get_rates(self, n_iter, mapper):
+        '''Get rates and associated uncertainties as of n_iter, according to the window size the user
+        has selected (self.windowsize)'''
+
+        if self.windowtype == 'fraction':
+            if self.max_windowsize is not None:
+                eff_windowsize = min(self.max_windowsize, int(n_iter * self.windowsize))
+            else:
+                eff_windowsize = int(n_iter * self.windowsize)
+        else:  # self.windowtype == 'fixed':
+            eff_windowsize = min(n_iter, self.windowsize or 0)
+
+        averager = RateAverager(mapper, self.system, self.data_manager, self.work_manager)
+        averager.calculate(max(1, n_iter - eff_windowsize), n_iter + 1, self.rate_calc_n_blocks, self.rate_calc_queue_size)
+        self.eff_windowsize = eff_windowsize
+
+        return averager
+
+    def prepare_new_iteration(self):
+        n_iter = self.sim_manager.n_iter
+        we_driver = self.sim_manager.we_driver
+
+        if not self.do_reweight:
+            # Reweighting not requested (or not possible)
+            log.debug('Reweighting not enabled')
+            return
+
+        with self.data_manager.lock:
+            wess_global_group = self.data_manager.we_h5file.require_group('wess')
+            reweighting_history_dataset = wess_global_group.require_dataset(
+                'reweighting_history', (1,), maxshape=(None,), dtype=int
+            )
+            last_reweighting = int(wess_global_group.attrs.get('last_reweighting', 0))
+            if last_reweighting > n_iter:
+                last_reweighting = n_iter - 1
+                reweighting_history = reweighting_history_dataset[:]
+                reweighting_history = reweighting_history[reweighting_history < n_iter]
+                reweighting_history_dataset.resize((reweighting_history.size), axis=0)
+
+        if n_iter - last_reweighting < self.reweight_period:
+            # Not time to reweight yet
+            log.debug('not reweighting')
+            return
+        else:
+            log.debug('reweighting')
+
+        if self.bin_mapper is None:
+            mapper = we_driver.bin_mapper
+            bins = we_driver.next_iter_binning
+            target_regions = list(we_driver.target_states.keys())
+            westpa.rc.pstatus('\nReweighting using the simulation bin mapper:\n{}'.format(mapper))
+        else:
+            mapper = self.bin_mapper
+            bins = mapper.construct_bins()
+
+            segments = [s for s in we_driver.next_iter_segments]
+            pcoords = self.system.new_pcoord_array(len(segments))
+            for iseg, segment in enumerate(segments):
+                pcoords[iseg] = segment.pcoord[0]
+            assignments = mapper.assign(pcoords)
+            for segment, assignment in zip(segments, assignments):
+                bins[assignment].add(segment)
+
+            target_states = list(we_driver.target_states.values())
+            target_regions = []
+            for tstate in target_states:
+                tstate_assignment = mapper.assign([tstate.pcoord])[0]
+                target_regions.append(tstate_assignment)
+            westpa.rc.pstatus('\nReweighting using a different bin mapper than simulation:\n{}'.format(mapper))
+
+        n_bins = len(bins)
+        westpa.rc.pstatus('Averaging rates')
+        averager = self.get_rates(n_iter, mapper)
+        binprobs = np.fromiter(map(operator.attrgetter('weight'), bins), dtype=np.float64, count=n_bins)
+        orig_binprobs = binprobs.copy()
+
+        if self.write_matrices:
+            # Create storage for ourselves
+            with self.data_manager.lock:
+                iter_group = self.data_manager.get_iter_group(n_iter)
+                try:
+                    del iter_group['wess']
+                except KeyError:
+                    pass
+
+                wess_iter_group = iter_group.create_group('wess')
+                wess_iter_group.create_dataset('avg_populations', data=averager.average_populations, compression=4)
+                wess_iter_group.create_dataset('unc_populations', data=averager.stderr_populations, compression=4)
+                wess_iter_group.create_dataset('avg_fluxes', data=averager.average_flux, compression=4)
+                wess_iter_group.create_dataset('unc_fluxes', data=averager.stderr_flux, compression=4)
+                wess_iter_group.create_dataset('avg_rates', data=averager.average_rate, compression=4)
+                wess_iter_group.create_dataset('unc_rates', data=averager.stderr_rate, compression=4)
+
+        westpa.rc.pstatus('Calculating reweighting using window size of {:d}'.format(self.eff_windowsize))
+        westpa.rc.pstatus('\nBin probabilities prior to reweighting:\n{!s}'.format(binprobs))
+        westpa.rc.pflush()
+
+        rij, oldindex = reduce_array(averager.average_rate)
+        uij = averager.stderr_rate[np.ix_(oldindex, oldindex)]
+
+        flat_target_regions = []
+        for target_region in target_regions:
+            if target_region in oldindex:  # it is possible that the target region was removed (ie if no recycling has occurred)
+                new_ibin = oldindex.index(target_region)  # this is in terms of the Rij
+                flat_target_regions.append(new_ibin)
+
+        binprobs = prob_adjust(binprobs, rij, uij, oldindex, flat_target_regions)
+
+        # Check to see if reweighting has set non-zero bins to zero probability (should never happen)
+        assert (~((orig_binprobs > 0) & (binprobs == 0))).all(), 'populated bin reweighted to zero probability'
+
+        # Check to see if reweighting has set zero bins to nonzero probability (may happen)
+        z2nz_mask = (orig_binprobs == 0) & (binprobs > 0)
+        if (z2nz_mask).any():
+            westpa.rc.pstatus('Reweighting would assign nonzero probability to an empty bin; not reweighting this iteration.')
+            westpa.rc.pstatus('Empty bins assigned nonzero probability: {!s}.'.format(np.array_str(np.arange(n_bins)[z2nz_mask])))
+        else:
+            westpa.rc.pstatus('\nBin populations after reweighting:\n{!s}'.format(binprobs))
+            for bin, newprob in zip(bins, binprobs):
+                if len(bin):
+                    bin.reweight(newprob)
+
+            reweighting_history_dataset.resize((reweighting_history_dataset.shape[0] + 1), axis=0)
+            reweighting_history_dataset[-1] = n_iter
+            wess_global_group.attrs['last_reweighting'] = n_iter
+
+        assert (
+            abs(1 - np.fromiter(map(operator.attrgetter('weight'), bins), dtype=np.float64, count=n_bins).sum())
+            < EPS * np.fromiter(map(len, bins), dtype=int, count=n_bins).sum()
+        )
+
+        westpa.rc.pflush()

westpa/work_managers/__init__.py

@@ -0,0 +1,57 @@
+'''A system for parallel, remote execution of multiple arbitrary tasks.
+Much of this, both in concept and execution, was inspired by (and in some
+cases based heavily on) the ``concurrent.futures`` package from Python 3.2,
+with some simplifications and adaptations (thanks to Brian Quinlan and his
+futures implementation).
+'''
+
+import logging
+
+from .core import WorkManager, WMFuture, FutureWatcher  # noqa
+
+
+# Import core work managers, which should run most everywhere that
+# Python does
+from . import serial, threads, processes  # noqa
+from .serial import SerialWorkManager
+from .threads import ThreadsWorkManager
+from .processes import ProcessWorkManager
+
+log = logging.getLogger(__name__)
+
+_available_work_managers = {'serial': SerialWorkManager, 'threads': ThreadsWorkManager, 'processes': ProcessWorkManager}
+
+# Import ZeroMQ work manager if available
+try:
+    from . import zeromq  # noqa
+    from .zeromq import ZMQWorkManager
+except ImportError:
+    log.info('ZeroMQ work manager not available')
+    log.debug('traceback follows', exc_info=True)
+else:
+    _available_work_managers['zmq'] = ZMQWorkManager
+
+# Import MPI work manager if available
+try:
+    from . import mpi  # noqa
+    from .mpi import MPIWorkManager
+except ImportError:
+    log.info('MPI work manager not available')
+    log.debug('traceback follows', exc_info=True)
+else:
+    _available_work_managers['mpi'] = MPIWorkManager
+
+from . import environment  # noqa
+from .environment import make_work_manager  # noqa
+
+
+__all__ = [
+    'serial',
+    'threads',
+    'processes',
+    'SerialWorkManager',
+    'ThreadsWorkManager',
+    'ProcessWorkManager',
+    'environment',
+    'make_work_manager',
+]