westpa 2022.10__cp312-cp312-macosx_10_9_x86_64.whl

westpa/westext/stringmethod/__init__.py

@@ -0,0 +1,11 @@
+"""westext.stringmethod - Plugin to drive the adaptive evolution of one or more
+strings of Voronoi bins
+
+Joshua L. Adelman 2011
+"""
+
+from .string_method import DefaultStringMethod, WESTStringMethod
+from .string_driver import StringDriver
+
+
+__all__ = ['DefaultStringMethod', 'WESTStringMethod', 'StringDriver']

westpa/westext/stringmethod/fourier_fitting.py

@@ -0,0 +1,69 @@
+import numpy as np
+import scipy
+import scipy.optimize
+
+
+class FourierFit:
+    def __init__(self, P=2, ndims=2, maxiters=100, tol=1.0e-6):
+        super().__init__()
+
+        self.P = P
+        self.maxiters = maxiters
+        self.ndims = ndims
+        self.tol = tol
+        self.pp = []
+        self.t0 = None
+        self.w0 = None
+
+    def calc_string(self, w, t, x_meas):
+        tlen = len(t)
+        t = np.linspace(0.0, 1.0, tlen)
+        x_est = x_meas[0, :] + (x_meas[-1, :] - x_meas[0, :]) * t[:, np.newaxis]
+        for i in range(self.ndims):
+            for j in range(self.P):
+                x_est[:, i] += w[i, j] * np.sin((j + 1) * np.pi * t)
+        return x_est
+
+    def _optimize_dist(self, tk, x_meas, w, k):
+        x_target = x_meas[k, :]
+        x_est = x_meas[0, :] + (x_meas[-1, :] - x_meas[0, :]) * tk
+        for i in range(self.ndims):
+            for j in range(self.P):
+                x_est[i] += w[i, j] * np.sin((j + 1) * np.pi * tk)
+
+        err = x_target - x_est
+        return err
+
+    def _optimize_w(self, w, x_meas, t, k, weight):
+        x_target = x_meas[:, k]
+        x_est = x_meas[0, k] + (x_meas[-1, k] - x_meas[0, k]) * t
+
+        for j in range(self.P):
+            x_est += w[j] * np.sin((j + 1) * np.pi * t)
+
+        err = weight * (x_target - x_est)
+        return err
+
+    def optimize(self, data, weight, w0, t0):
+        ncenters = data.shape[0]
+        self.w0 = w0
+        self.t0 = t0
+        if weight is None:
+            weight = np.ones_like(t0)
+
+        for iiter in range(self.maxiters):
+            self.pp.append(self.calc_string(self.w0, self.t0, data))
+            if iiter > 0:
+                err = np.sum((self.pp[-1] - self.pp[-2]) ** 2) / ncenters
+                print('{} -- {}'.format(iiter, err))
+                if err < self.tol:
+                    break
+            else:
+                print(iiter)
+            # Optimize tk
+            for ci in range(ncenters):
+                self.t0[ci] = scipy.optimize.leastsq(self._optimize_dist, self.t0[ci], args=(data, self.w0, ci))[0]
+
+            # Optimize wij
+            for k in range(self.ndims):
+                self.w0[k, :] = scipy.optimize.leastsq(self._optimize_w, self.w0[k, :], args=(data, self.t0, k, weight))[0]

westpa/westext/stringmethod/string_driver.py

@@ -0,0 +1,253 @@
+import logging
+import types
+
+import numpy as np
+
+import westpa
+from westpa.core import extloader
+from westpa.core.binning import VoronoiBinMapper
+from westpa.core.yamlcfg import check_bool, ConfigItemMissing
+
+from westpa.westext.stringmethod import WESTStringMethod, DefaultStringMethod
+
+
+log = logging.getLogger(__name__)
+
+
+class StringDriver:
+    def __init__(self, sim_manager, plugin_config):
+        super().__init__()
+
+        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
+
+        # Parameters from config file
+        self.windowsize = plugin_config.get('windowsize', 10)
+        self.update_interval = plugin_config.get('update_interval', 10)
+        self.initial_update = plugin_config.get('initial_update', 20)
+        self.priority = plugin_config.get('priority', 0)
+
+        self.write_avg_pos = check_bool(plugin_config.get('write_avgpos', True))
+        self.do_update = check_bool(plugin_config.get('do_update', True))
+        self.init_from_data = check_bool(plugin_config.get('init_from_data', True))
+
+        self.dfunc = self.get_dfunc_method(plugin_config)
+
+        # Load method to calculate average position in a bin
+        # If the method is defined in an external module, correctly bind it
+        ap = self.get_avgpos_method(plugin_config)
+        if hasattr(ap, 'im_class'):
+            self.get_avgpos = ap
+        else:
+            self.get_avgpos = types.MethodType(ap, self)
+
+        # Get initial set of string centers
+        centers = self.get_initial_centers()
+
+        try:
+            sm_params = self.system.sm_params
+        except AttributeError as e:
+            log.error(
+                'String Driver Error: system does not define sm_params. \
+                        This is required and should be added to the system definition; {}'.format(
+                    e
+                )
+            )
+            raise
+
+        # Initialize the string
+        str_method = self.get_string_method(plugin_config)
+
+        try:
+            self.strings = str_method(centers, **sm_params)
+        except (TypeError, AssertionError) as e:
+            log.error('String Driver Error: Failed during initialization of string method: {}'.format(e))
+            raise
+
+        # Update the BinMapper
+        self.update_bin_mapper()
+
+        # Register callback
+        sim_manager.register_callback(sim_manager.prepare_new_iteration, self.prepare_new_iteration, self.priority)
+
+        westpa.rc.pstatus('-westext.stringmethod -----------------\n')
+        westpa.rc.pstatus('windowsize: {}\n'.format(self.windowsize))
+        westpa.rc.pstatus('update interval: {}\n'.format(self.update_interval))
+        westpa.rc.pstatus('initial update: {}\n'.format(self.initial_update))
+        westpa.rc.pstatus('priority: {}\n'.format(self.priority))
+        westpa.rc.pstatus('write average positions: {}\n'.format(self.write_avg_pos))
+        westpa.rc.pstatus('do update: {}\n'.format(self.do_update))
+        westpa.rc.pstatus('initialize from WE data: {}\n'.format(self.init_from_data))
+        westpa.rc.pstatus('----------------------------------------\n')
+        westpa.rc.pflush()
+
+    def dfunc(self):
+        raise NotImplementedError
+
+    def get_avgpos(self, n_iter):
+        raise NotImplementedError
+
+    def get_dfunc_method(self, plugin_config):
+        try:
+            methodname = plugin_config['dfunc_method']
+        except KeyError:
+            raise ConfigItemMissing('dfunc_method')
+
+        dfunc_method = extloader.get_object(methodname)
+
+        log.info('loaded stringmethod dfunc method {!r}'.format(dfunc_method))
+
+        return dfunc_method
+
+    def get_avgpos_method(self, plugin_config):
+        try:
+            methodname = plugin_config['avgpos_method']
+        except KeyError:
+            raise ConfigItemMissing('avgpos_method')
+
+        if methodname.lower() == 'cartesian':
+            avgpos_method = self.avgpos_cartesian
+        else:
+            avgpos_method = extloader.get_object(methodname)
+
+        log.info('loaded stringmethod avgpos method {!r}'.format(avgpos_method))
+
+        return avgpos_method
+
+    def get_string_method(self, plugin_config):
+        try:
+            methodname = plugin_config['string_method']
+        except KeyError:
+            raise ConfigItemMissing('string_method')
+
+        if methodname.lower() == 'default':
+            str_method = DefaultStringMethod
+        else:
+            str_method = extloader.get_object(methodname)
+
+        assert issubclass(str_method, WESTStringMethod)
+        log.debug('loaded stringmethod string method {!r}'.format(str_method))
+
+        return str_method
+
+    def get_initial_centers(self):
+        self.data_manager.open_backing()
+
+        with self.data_manager.lock:
+            n_iter = max(self.data_manager.current_iteration - 1, 1)
+            iter_group = self.data_manager.get_iter_group(n_iter)
+
+            # First attempt to initialize string from data rather than system
+            centers = None
+            if self.init_from_data:
+                log.info('Attempting to initialize stringmethod from data')
+
+                try:
+                    binhash = iter_group.attrs['binhash'].encode()
+                    bin_mapper = self.data_manager.get_bin_mapper(binhash)
+
+                    centers = bin_mapper.centers
+
+                except Exception:
+                    log.warning('Initializing string centers from data failed; Using definition in system instead.')
+                    centers = self.system.bin_mapper.centers
+            else:
+                log.info('Initializing string centers from system definition')
+                centers = self.system.bin_mapper.centers
+
+        self.data_manager.close_backing()
+
+        return centers
+
+    def update_bin_mapper(self):
+        '''Update the bin_mapper using the current string'''
+
+        westpa.rc.pstatus('westext.stringmethod: Updating bin mapper\n')
+        westpa.rc.pflush()
+
+        try:
+            dfargs = getattr(self.system, 'dfargs', None)
+            dfkwargs = getattr(self.system, 'dfkwargs', None)
+            self.system.bin_mapper = VoronoiBinMapper(self.dfunc, self.strings.centers, dfargs=dfargs, dfkwargs=dfkwargs)
+        except (ValueError, TypeError) as e:
+            log.error('StringDriver Error: Failed updating the bin mapper: {}'.format(e))
+            raise
+
+    def avgpos_cartesian(self, n_iter):
+        '''Get average position of replicas in each bin as of n_iter for the
+        the user selected update interval'''
+
+        nbins = self.system.bin_mapper.nbins
+        ndim = self.system.pcoord_ndim
+
+        avg_pos = np.zeros((nbins, ndim), dtype=self.system.pcoord_dtype)
+        sum_bin_weight = np.zeros((nbins,), dtype=self.system.pcoord_dtype)
+
+        start_iter = max(n_iter - min(self.windowsize, n_iter), 2)
+        stop_iter = n_iter + 1
+
+        for n in range(start_iter, stop_iter):
+            with self.data_manager.lock:
+                iter_group = self.data_manager.get_iter_group(n)
+                seg_index = iter_group['seg_index'][...]
+
+                pcoords = iter_group['pcoord'][:, -1, :]  # Only read final point
+                bin_indices = self.system.bin_mapper.assign(pcoords)
+                weights = seg_index['weight']
+
+                pcoord_w = pcoords * weights[:, np.newaxis]
+                uniq_indices = np.unique(bin_indices)
+
+                for indx in uniq_indices:
+                    avg_pos[indx, :] += pcoord_w[bin_indices == indx].sum(axis=0)
+
+                sum_bin_weight += np.bincount(bin_indices.astype(int), weights=weights, minlength=nbins)
+
+        # Some bins might have zero samples so exclude to avoid divide by zero
+        occ_ind = np.nonzero(sum_bin_weight)
+        avg_pos[occ_ind] /= sum_bin_weight[occ_ind][:, np.newaxis]
+
+        return avg_pos, sum_bin_weight
+
+    def prepare_new_iteration(self):
+        n_iter = self.sim_manager.n_iter
+
+        with self.data_manager.lock:
+            iter_group = self.data_manager.get_iter_group(n_iter)
+
+            try:
+                del iter_group['stringmethod']
+            except KeyError:
+                pass
+
+            sm_global_group = self.data_manager.we_h5file.require_group('stringmethod')
+            last_update = int(sm_global_group.attrs.get('last_update', 0))
+
+        if n_iter - last_update < self.update_interval or n_iter < self.initial_update or not self.do_update:
+            log.debug('Not updating string this iteration')
+            return
+        else:
+            log.debug('Updating string - n_iter: {}'.format(n_iter))
+
+        westpa.rc.pstatus('-westext.stringmethod -----------------\n')
+        westpa.rc.pstatus('westext.stringmethod: Calculating average position in string images\n')
+        westpa.rc.pflush()
+
+        avg_pos, sum_bin_weight = self.get_avgpos(n_iter)
+
+        westpa.rc.pstatus('westext.stringmethod: Updating string\n')
+        westpa.rc.pflush()
+
+        self.strings.update_string_centers(avg_pos, sum_bin_weight)
+
+        westpa.rc.pstatus('westext.stringmethod: String lengths: {}\n'.format(self.strings.length))
+        westpa.rc.pflush()
+
+        # Update the bin definitions
+        self.update_bin_mapper()
+
+        sm_global_group.attrs['last_update'] = n_iter

westpa/westext/stringmethod/string_method.py

@@ -0,0 +1,306 @@
+from abc import ABCMeta, abstractmethod, abstractproperty
+import logging
+
+try:
+    from collections.abc import Iterable
+except ImportError:
+    from collections import Iterable
+
+import numpy as np
+
+
+try:
+    import scipy
+    import scipy.interpolate
+    import scipy.linalg
+
+    SCIPY_FLAG = True
+except Exception:
+    SCIPY_FLAG = False
+
+from .fourier_fitting import FourierFit
+
+
+log = logging.getLogger(__name__)
+
+
+class WESTStringMethod:
+    ___metaclass__ = ABCMeta
+
+    def __init__(self, centers, **kwargs):
+        pass
+
+    @abstractproperty
+    def centers(self):
+        """Return the centers of all of the strings"""
+        pass
+
+    @abstractproperty
+    def length(self):
+        """Return a list of the lengths of each string"""
+        pass
+
+    @abstractmethod
+    def update_string_centers(self, avgcoords, binprob):
+        """Given a set of average coordinates (avgcoords) in each bin
+        and the individual probabilities for each bin (binprob), update
+        the string centers
+        """
+        pass
+
+
+class DefaultStringMethod(WESTStringMethod):
+    """Implementation of a method to evolve one or more pseudo-1D strings in a high dimensional
+    progress coordinate space.
+
+    **Parameters**
+    centers:        A numpy array of size (number of total centers,pcoord dim) that stores
+                        the positions of all of the string images
+    slen:           An iterable containing the number of centers in each string
+    slabels:        An list containing the relative positions in each string of any state label
+                        progress coordinates if present. These progress coordinates will be ignored in the
+                        calculation. None if no labels
+    mpairs:         A list of lists containing the indices of pairs of centers that should move together.
+                        None if strings move independently
+    dtau:           Parameter controlling the rate at which centers move toward the average value in the bin
+    kappa:          Parameter controlling the smoothing of the string
+    fixed_ends:     Boolean flag specifying whether to fix ends of the strings
+    sciflag:        Boolean flag specifying whether to attempt to use scipy methods which are
+                        generally more efficient
+    fourierflag:    Boolean flag specifying whether to user fourier fitting method
+    fourier_P:      Integer value specifying how many fourier modes to use in fitting
+    fourier_maxiters: Maximum number of iterations of fourier fitting procedure
+    fourier_tol:    Tolerance for ending fourier fitting
+    """
+
+    def __init__(
+        self,
+        centers,
+        slen=None,
+        slabels=None,
+        mpairs=None,
+        dtau=0.1,
+        kappa=0.1,
+        sciflag=None,
+        fixed_ends=True,
+        fourierflag=False,
+        fourier_P=2,
+        fourier_maxiters=100,
+        fourier_tol=1.0e-6,
+        **kwargs
+    ):
+        super().__init__(centers, **kwargs)
+
+        self._SCIPY_FLAG = None
+        self._fixed_ends = fixed_ends
+
+        if sciflag is None:
+            self._SCIPY_FLAG = SCIPY_FLAG
+        else:
+            self._SCIPY_FLAG = SCIPY_FLAG & sciflag
+
+        self._dtau = dtau
+        self._kappa = kappa
+
+        self._centers = centers
+
+        self._nstrings = len(slen)  # Number of strings
+        self._N = centers.shape[0]
+        self._ndim = centers.shape[1]  # Number of progress coordinates
+
+        if slen and isinstance(slen, Iterable):
+            self._slen = np.array(slen)
+        elif slen:
+            self._slen = np.array([slen])
+        else:
+            self._slen = np.array([self._N])
+            log.warning('Input parameters do not define slen; assuming system is composed of a single string')
+
+        self._mpairs = mpairs
+
+        assert np.sum(self._slen) == self._N
+
+        self._skip_dim = np.array(slabels) if slabels is not None else np.array([])
+
+        # Get iterable of slicing objects to get per string center coordinates to perform calculation on
+        self._strindx = []
+
+        indx_all = np.arange(self._ndim)
+        self._indx_take = np.setdiff1d(indx_all, self._skip_dim)
+        self._ndim_take = len(self._indx_take)
+
+        start_count = 0
+        for sl in slen:
+            self._strindx.append(np.index_exp[start_count : (start_count + sl), self._indx_take])
+            start_count += sl
+
+        # Fourier fitting parameters
+        self._FFIT_FLAG = fourierflag
+        self._ffp = fourier_P
+        self._ffmaxiters = fourier_maxiters
+        self._fftol = fourier_tol
+
+        # Create dict to hold kappan and A objects for all unique lengths of strings
+        self._kappan = {}
+        self._A = {}
+
+        self.finalize_init()
+
+    @property
+    def centers(self):
+        return self._centers
+
+    @property
+    def length(self):
+        L = []
+        for sid, si in enumerate(self._strindx):
+            L.append(self.calculate_length(self._centers[si])[-1])
+
+        return L
+
+    def calculate_length(self, x):
+        dd = x - np.roll(x, 1, axis=0)
+        dd[0, :] = 0.0
+        return np.cumsum(np.sqrt((dd * dd).sum(axis=1)))
+
+    def finalize_init(self):
+        # Set up A and kappan for each string
+        uslen = np.unique(self._slen)
+
+        for ulen in uslen:
+            self._kappan[ulen] = self._kappa * self._dtau * ulen
+            self._A[ulen] = None
+
+            if self._SCIPY_FLAG:
+                ud = np.zeros((ulen,))
+                ld = np.zeros((ulen,))
+                d = np.ones((ulen,))
+
+                d[1:-1] = 2.0 * self._kappan[ulen] + 1.0
+                ud[2:] = -self._kappan[ulen]
+                ld[:-2] = -self._kappan[ulen]
+
+                self._A[ulen] = np.mat([ud, d, ld])
+
+            else:
+                self._A[ulen] = np.eye(ulen)
+                di = np.diag_indices(ulen, ndim=2)
+                ii = (di[0][1:-1], di[1][1:-1])
+
+                self._A[ulen][ii] = 2.0 * self._kappan[ulen] + 1.0
+
+                dd = np.zeros((ulen - 1,))
+                dd[1:] = -self._kappan[ulen]
+                self._A[ulen] += np.diag(dd, k=1)
+
+                dd = np.zeros((ulen - 1,))
+                dd[:-1] = -self._kappan[ulen]
+                self._A[ulen] += np.diag(dd, k=-1)
+
+    def update_string_centers(self, avgcoords, binprob):
+        """Update the position of all string centers
+        **Parameters**
+        avgcoords:      Average position of replicas in each voronoi cell
+        binprob:        The total weight in each voronoi cell
+
+        """
+        assert self.centers.shape == avgcoords.shape
+
+        # If centers are paired, calculate their average position
+        if self._mpairs is not None:
+            for pi in self._mpairs:
+                pprob = binprob[pi]
+                if np.sum(pprob) == 0.0:
+                    continue
+
+                idx = np.ix_(pi, self._indx_take)
+                pavg = np.ma.array(avgcoords[idx])
+                zind = np.where(pprob == 0)[0]
+                pavg[zind, :] = np.ma.masked
+
+                avgcoords[idx] = pavg.mean(axis=0)
+
+        for sid, si in enumerate(self._strindx):
+            x = avgcoords.copy()[si]
+            centers = self.centers[si]
+            occupied = np.nonzero(binprob[si[0]])
+
+            N = self._slen[sid]
+
+            # if avgcoords has missing values fill them by linearly interpolating
+            # present data
+            if occupied[0].shape != N:
+                notocc = np.ones((N,), dtype=bool)  # unoccupied
+                notocc[occupied] = False
+
+                # marked paired centers as occupied to avoid reseting averaged value if
+                # at least one is occupied
+                if self._mpairs is not None:
+                    for pi in self._mpairs:
+                        totprob = np.sum(binprob[pi])
+                        if totprob == 0.0:
+                            continue
+                        else:
+                            for m in pi:
+                                if (m >= si[0].start) and (m < si[0].stop):
+                                    notocc[m - si[0].start] = False
+
+                cfunc = lambda z: z.nonzero()[0]
+
+                # Handle ends first
+                if notocc[0]:
+                    x[0, :] = centers[0, :]
+                    notocc[0] = False
+                if notocc[-1]:
+                    x[-1, :] = centers[-1, :]
+                    notocc[-1] = False
+
+                # interpolate values for unoccupied bins
+                if self._SCIPY_FLAG:
+                    for k in range(self._ndim_take):
+                        f = scipy.interpolate.interp1d(cfunc(~notocc), x[~notocc, k], kind='linear')
+                        x[notocc, k] = f(cfunc(notocc))
+                else:
+                    for k in range(self._ndim_take):
+                        x[notocc, k] = np.interp(cfunc(notocc), cfunc(~notocc), x[~notocc, k])
+
+            if self._fixed_ends:
+                x[0, :] = centers[0, :]
+                x[-1, :] = centers[-1, :]
+
+            psi = centers
+            psi_new = np.zeros_like(psi)
+
+            b = psi - self._dtau * (psi - x)
+
+            # Update and smooth the string
+            if self._SCIPY_FLAG:
+                for k in range(self._ndim_take):
+                    psi_new[:, k] = scipy.linalg.solve_banded((1, 1), self._A[N], b[:, k])
+            else:
+                for k in range(self._ndim_take):
+                    psi_new[:, k] = np.linalg.solve(self._A[N], b[:, k])
+
+            # Optionally smooth using fourier method
+            if self._FFIT_FLAG:
+                w0 = np.zeros((self._ndim_take, self._ffp), np.float)
+                t0 = np.linspace(0, 1, psi_new.shape[0])
+
+                ff = FourierFit(P=self._ffp, maxiters=self._ffmaxiters)
+                ff.optimize(psi_new, None, w0, t0)
+                psi_new = ff.pp[-1][:]
+
+            # Enforce equal spacing between centers along the string
+            L = self.calculate_length(psi_new)
+            L /= L[-1]
+            g2 = np.linspace(0, 1, N)
+
+            if self._SCIPY_FLAG:
+                for k in range(self._ndim_take):
+                    f = scipy.interpolate.interp1d(L, psi_new[:, k], kind='linear')
+                    psi_new[:, k] = f(g2)
+            else:
+                for k in range(self._ndim_take):
+                    psi_new[:, k] = np.interp(g2, L, psi_new[:, k])
+
+            self.centers[si] = psi_new.copy()