westpa 2022.12__cp313-cp313-manylinux_2_17_x86_64.manylinux2014_x86_64.whl

westpa/westext/weed/BinCluster.py

@@ -0,0 +1,180 @@
+import numpy as np
+
+
+from . import UncertMath
+
+
+class ClusterList:
+    def __init__(self, ratios, nbins):
+        super().__init__()
+        self.nbins = nbins
+        self.ratios = ratios
+
+        # Create an array to hold bin assignments and initial set all to -1 == not clustered
+        self.bin_assign = np.empty((nbins,), dtype=int)
+        self.bin_assign.fill(-1)
+
+        # Initialize an ucert container to hold per bin information; initially mask all elements
+        # note: masking is implicit since rates set to 0
+        dum_data = np.zeros((nbins,))
+        self.bin_data = UncertMath.UncertContainer(dum_data.copy(), dum_data.copy(), dum_data.copy())
+
+        self.cluster_id = 0  # ID of newest cluster
+        self.cluster_contents = {}  # Dictionary containing sets of bin ids with keys = cluster ids
+
+    def join(self, pairs):
+        """Join clusters given a tuple (i,j) of bin pairs"""
+
+        for i, j in zip(*pairs):
+            # Both bins not joined
+            if self.bin_assign[i] == -1 and self.bin_assign[j] == -1:
+                # Create new cluster
+                self.bin_assign[i] = self.cluster_id
+                self.bin_assign[j] = self.cluster_id
+
+                self.cluster_contents[self.cluster_id] = {i, j}
+                self.cluster_id += 1
+
+                rij = self.ratios[i, j]
+                denom = rij + 1.0
+                self.bin_data[i] = rij / denom  # relative probability for bin i
+                self.bin_data[j] = denom.recip()  # relative probability for bin j
+
+            # Only one bin previously assigned to a cluster
+            elif self.bin_assign[i] == -1 or self.bin_assign[j] == -1:
+                if self.bin_assign[i] == -1:
+                    idum, jdum = i, j
+                else:
+                    idum, jdum = j, i
+
+                jclust = self.bin_assign[jdum]
+                jclust_mid = np.where(self.bin_assign == jclust)[0]  # index of bins in jclust
+
+                rik = self.ratios[idum, jclust_mid]
+                pk = self.bin_data[jclust_mid]
+                piTmp = rik * pk  # estimate for p_idum / P_cluster based on 'path' through bin k
+                # Note that here P_cluster is value before addition of bin idum
+                piTmp_avg = piTmp.weighted_average(axis=0)
+
+                # now, compute relative prob of each bin in *new* cluster (including bin idum)
+                denom = piTmp_avg + 1.0
+                self.bin_data[idum] = piTmp_avg / denom
+
+                # Update bins already in cluster
+                self.bin_data[jclust_mid] = self.bin_data[jclust_mid] / denom
+
+                # Move bin idum into cluster jclust
+                self.bin_assign[idum] = jclust
+                self.cluster_contents[jclust].update({idum})
+
+            # Both bins previously assigned to different cluster; Join clusters
+            elif not self.bin_assign[i] == self.bin_assign[j]:
+                iclust = self.bin_assign[i]
+                jclust = self.bin_assign[j]
+
+                iclust_mid = np.where(self.bin_assign == iclust)[0]  # indx of bins in cluster i
+                jclust_mid = np.where(self.bin_assign == jclust)[0]  # indx of bins in cluster j
+
+                niclust = iclust_mid.size
+                njclust = jclust_mid.size
+
+                dum_data = np.zeros((niclust * njclust,))
+                ij_cluster_ratio = UncertMath.UncertContainer(dum_data.copy(), dum_data.copy(), dum_data.copy())
+
+                for count, im in enumerate(iclust_mid):
+                    rij = self.ratios[im, jclust_mid]
+                    pi = self.bin_data[im]
+                    pj = self.bin_data[jclust_mid]
+
+                    ij_cluster_ratio[count * njclust : (count + 1) * njclust] = rij * pj / pi
+
+                ij_cluster_ratio = ij_cluster_ratio.weighted_average(axis=0)
+
+                idenom = ij_cluster_ratio.recip() + 1.0
+                jdenom = ij_cluster_ratio + 1.0
+
+                self.bin_data[iclust_mid] = self.bin_data[iclust_mid] / idenom
+                self.bin_data[jclust_mid] = self.bin_data[jclust_mid] / jdenom
+
+                # Join all bins in cluster j into cluster iclust
+                self.bin_assign[jclust_mid] = iclust
+
+                # Move contents of jclust into iclust
+                self.cluster_contents[iclust].update(self.cluster_contents[jclust])
+
+                # Clear contents of jclust
+                self.cluster_contents[jclust].clear()
+
+                if len(self.cluster_contents[iclust]) == self.nbins:
+                    break
+
+    def join_simple(self, pairs):
+        """Join clusters using direct ratios given a tuple (i,j) of bin pairs"""
+
+        for i, j in zip(*pairs):
+            # Both bins not joined
+            if self.bin_assign[i] == -1 and self.bin_assign[j] == -1:
+                # Create new cluster
+                self.bin_assign[i] = self.cluster_id
+                self.bin_assign[j] = self.cluster_id
+
+                self.cluster_contents[self.cluster_id] = {i, j}
+                self.cluster_id += 1
+
+                rij = self.ratios[i, j]
+                denom = rij + 1.0
+                self.bin_data[i] = rij / denom  # relative probability for bin i
+                self.bin_data[j] = denom.recip()  # relative probability for bin j
+
+            # Only one bin previously assigned to a cluster
+            elif self.bin_assign[i] == -1 or self.bin_assign[j] == -1:
+                if self.bin_assign[i] == -1:
+                    idum, jdum = i, j
+                else:
+                    idum, jdum = j, i
+
+                jclust = self.bin_assign[jdum]
+                rik = self.ratios[idum, jdum]
+                pk = self.bin_data[jdum]
+                piTmp = rik * pk  # estimate for p_idum / P_cluster based on 'path' through bin k
+                # Note that here P_cluster is value before addition of bin idum
+                # now, compute relative prob of each bin in *new* cluster (including bin idum)
+                denom = piTmp + 1.0
+                self.bin_data[idum] = piTmp / denom
+
+                # Update bins already in cluster
+                jclust_mid = np.where(self.bin_assign == jclust)  # index of bins in jclust
+                self.bin_data[jclust_mid] = self.bin_data[jclust_mid] / denom
+
+                # Move bin idum into cluster jclust
+                self.bin_assign[idum] = jclust
+                self.cluster_contents[jclust].update({idum})
+
+            # Both bins previously assigned to different cluster; Join clusters
+            elif not self.bin_assign[i] == self.bin_assign[j]:
+                iclust = self.bin_assign[i]
+                jclust = self.bin_assign[j]
+                rij = self.ratios[i, j]
+                pi = self.bin_data[i]
+                pj = self.bin_data[j]
+                ij_cluster_ratio = rij * pj / pi
+                idenom = ij_cluster_ratio.recip() + 1.0
+                jdenom = ij_cluster_ratio + 1.0
+
+                iclust_mid = np.where(self.bin_assign == iclust)
+                self.bin_data[iclust_mid] = self.bin_data[iclust_mid] / idenom
+
+                jclust_mid = np.where(self.bin_assign == jclust)
+                self.bin_data[jclust_mid] = self.bin_data[jclust_mid] / jdenom
+
+                # Join all bins in cluster j into cluster iclust
+                self.bin_assign[jclust_mid] = iclust
+
+                # Move contents of jclust into iclust
+                self.cluster_contents[iclust].update(self.cluster_contents[jclust])
+
+                # Clear contents of jclust
+                self.cluster_contents[jclust].clear()
+
+                if len(self.cluster_contents[iclust]) == self.nbins:
+                    break

westpa/westext/weed/ProbAdjustEquil.py

@@ -0,0 +1,100 @@
+import numpy as np
+import numpy.ma as ma
+
+from . import UncertMath
+from . import BinCluster
+
+
+def probAdjustEquil(binProb, rates, uncert, threshold=0.0, fullCalcClust=False, fullCalcBins=False):
+    """This function adjusts bin pops in binProb using rates and uncert matrices
+    fullCalcBins --> True for weighted avg, False for simple calc
+    fullCalcClust --> True for weighted avg, False for simple calc
+    threshold --> minimum weight (relative to max) for another value to be averaged
+            only matters if fullCalcBins == True (or later perhaps if fullCalcClust == True)
+    """
+
+    # Check that rate matrix is square
+    Ni, Nj = rates.shape
+    if Ni != Nj:
+        print('\nWARNING: Not a square matrix!\n')
+
+    zi = np.where(binProb == 0.0)[0]  # indices of bins with zero probability
+
+    rates_uncert = UncertMath.UncertContainer(rates, rates - uncert, rates + uncert)
+
+    # STEP 1a: Create matrix of ratios of probabilities based on DIRECT estimates
+    # that is, ij element is p_i / p_j = k_ji / k_ij
+
+    ratios_direct = rates_uncert.transpose() / rates_uncert
+
+    # STEP 1b: Create averaged matrix of ratios of probabilities based on both direct and indirect estimates
+    # Indirect means '3rd bin' estimates: p_i / p_j = ( k_ki / k_ik ) ( k_jk / k_kj )
+    # Turns out this is not helpful, so generally set fullCalcBins = 0
+    if fullCalcBins:
+        # Calculate indirect ratios using Einstein Summation convention where
+        # ratios_indirect_kij  = ( k_ki / k_ik ) ( k_jk / k_kj ) = ratios_direct_ik * ratios_direct_kj
+        ri_vals = np.einsum('ik,kj->kij', ratios_direct.vals, ratios_direct.vals)
+        ri_min = np.einsum('ik,kj->kij', ratios_direct.dmin, ratios_direct.dmin)
+        ri_max = np.einsum('ik,kj->kij', ratios_direct.dmax, ratios_direct.dmax)
+        ratios_indirect = UncertMath.UncertContainer(ri_vals, ri_min, ri_max, mask=ratios_direct.vals.mask)
+
+        # Threshold indirect ratios
+        ti = ratios_indirect.wt < ratios_direct * threshold
+        ratios_indirect.mask = ti
+        ratios_indirect.update_mask()
+
+        ratios_indirect.concatenate(ratios_direct, axis=0)
+        ratios_average = ratios_indirect.weighted_average(axis=0)
+
+    else:
+        ratios_average = ratios_direct.weighted_average(axis=0, expaxis=0)
+
+    # STEP 2: Form clusters
+
+    # STEP 2a: Sort probability ratios based on uncertainty
+    # Sort uncertainties of ratios_average subject to the convention that p_i < p_j
+
+    i, j = np.triu_indices(Ni, 1)  # indices of ij pairs where i != j
+
+    # Remove pairs that include a bin that has zero probability
+    nzi = (binProb[i] != 0.0) & (binProb[j] != 0.0)
+    i = i[nzi]
+    j = j[nzi]
+
+    vals = ma.vstack((ratios_average.vals[i, j], ratios_average.vals[j, i]))
+    ias = ma.argsort(vals, axis=0, kind='stable', fill_value=np.inf)
+
+    ordered_ind = np.vstack((i, j))
+    flip_ind = np.nonzero(ias[0, :])  # Find pairs in which to select ji rather than ij
+    ordered_ind[:, flip_ind[0]] = ordered_ind[:, flip_ind[0]][::-1]
+
+    iind = ordered_ind[0, :]
+    jind = ordered_ind[1, :]
+    uncertij = ratios_average.uncert[iind, jind]  # Get the uncert for ij pairs
+
+    count = uncertij.count()  # Count of the unmasked uncertainties
+    ias = ma.argsort(uncertij, kind='stable', fill_value=np.inf)  # Get the indices that would sort uncertij
+    iind = iind[ias[:count]]  # Sort the indices excluding masked/undefined values
+    jind = jind[ias[:count]]
+
+    # STEP 2b: Create ClusterList object and cluster bins
+    clusters = BinCluster.ClusterList(ratios_average, Ni)
+
+    if fullCalcClust:
+        clusters.join((iind, jind))
+    else:
+        clusters.join_simple((iind, jind))
+
+    total_prob = 0.0  # total probability in all clusters
+    for cid in clusters.cluster_contents:
+        binlist = list(clusters.cluster_contents[cid])
+        if len(binlist):
+            prob_cluster = binProb[binlist].sum()
+            total_prob += prob_cluster
+
+            binProb[binlist] = prob_cluster * clusters.bin_data[binlist].vals
+
+    binProb[zi] = 0.0  # re-zero bins that previously had zero prob
+    # for bi,p in enumerate(binProb):
+    #    print('bin: {} -- {}'.format(bi,p))
+    print('.........Total Probability: {}'.format(binProb.sum()))

westpa/westext/weed/UncertMath.py

@@ -0,0 +1,247 @@
+import numpy as np
+import numpy.ma as ma
+import itertools
+
+
+TOL = 1.0e-6
+
+
+class UncertContainer:
+    """Container to hold uncertainty measurements. Data is convert to np masked arrays
+    to avoid possible numerical problems
+    """
+
+    def __init__(self, vals, vals_dmin, vals_dmax, mask=ma.nomask):
+        super().__init__()
+
+        # If input data already masked arrays extract unmasked data
+        if ma.isMaskedArray(vals):
+            vals = vals.data
+        if ma.isMaskedArray(vals_dmin):
+            vals_dmin = vals_dmin.data
+        if ma.isMaskedArray(vals_dmax):
+            vals_dmax = vals_dmax.data
+
+        # Adjust negative values
+        ineg = np.where(vals_dmin <= 0.0)
+        vals_dmin[ineg] = TOL * vals[ineg]
+
+        # Calculate weight based on fractional uncertainty
+        diff = vals_dmax - vals_dmin
+        diff_m = ma.masked_where(vals_dmax == vals_dmin, diff)
+
+        self.vals = ma.masked_where(vals == 0.0, vals)
+
+        self.wt = (self.vals / diff_m) ** 2
+        self.uncert = diff_m / self.vals
+
+        self.wt.fill_value = np.inf
+        self.uncert.fill_vaule = np.inf
+
+        assert np.all(self.wt.mask == self.uncert.mask)
+
+        # Mask data if uncertainty is not finite or if any of the inputs were
+        # already masked
+
+        mm = ma.mask_or(self.wt.mask, mask)
+
+        self.vals.mask = mm
+        self.wt.mask = mm
+        self.uncert.mask = mm
+        self.dmin = ma.array(vals_dmin, mask=mm, fill_value=np.inf)
+        self.dmax = ma.array(vals_dmax, mask=mm, fill_value=np.inf)
+
+        self.mask = ma.getmaskarray(self.vals)
+
+    def __getitem__(self, indx):
+        vals = self.vals[indx]
+        dmin = self.dmin[indx]
+        dmax = self.dmax[indx]
+
+        if isinstance(vals, ma.core.MaskedConstant):
+            dum = np.zeros((1,))
+            return UncertContainer(dum.copy(), dum.copy(), dum.copy())
+        elif isinstance(vals, (float, int, np.floating, np.integer)):
+            return UncertContainer(np.array([vals]), np.array([dmin]), np.array([dmax]))
+        elif isinstance(vals, np.ndarray):
+            return UncertContainer(vals, dmin, dmax, mask=vals.mask)
+        else:
+            raise TypeError
+
+    def __setitem__(self, indx, value):
+        if isinstance(value, UncertContainer):
+            self.vals[indx] = value.vals
+            self.dmin[indx] = value.dmin
+            self.dmax[indx] = value.dmax
+            self.wt[indx] = value.wt
+            self.uncert[indx] = value.uncert
+        else:
+            raise TypeError('Can only set values with an UncertContainer object')
+
+    def __repr__(self):
+        return 'shape={} vals={} dmin={} dmax={} vals.mask={}'.format(
+            self.vals.shape, self.vals, self.dmin, self.dmax, self.vals.mask
+        )
+
+    def __add__(self, value):
+        if isinstance(value, UncertContainer):
+            vals = self.vals + value.vals
+            dmin = self.dmin + value.dmin
+            dmax = self.dmax + value.dmax
+
+            return UncertContainer(vals, dmin, dmax, mask=vals.mask)
+        elif isinstance(value, (float, int, np.floating, np.integer)):
+            vals = self.vals + value
+            dmin = self.dmin + value
+            dmax = self.dmax + value
+
+            return UncertContainer(vals, dmin, dmax, mask=vals.mask)
+        else:
+            raise TypeError('Attempt to add value of unsupported type')
+
+    def __sub__(self, value):
+        if isinstance(value, UncertContainer):
+            vals = self.vals - value.vals
+            dmin = self.dmin - value.dmin
+            dmax = self.dmax - value.dmax
+
+            return UncertContainer(vals, dmin, dmax, mask=vals.mask)
+        else:
+            raise TypeError('Attempted to subtract by value of unsupported type')
+
+    def __mul__(self, value):
+        if isinstance(value, UncertContainer):
+            vals = self.vals * value.vals
+            dmin = self.dmin * value.dmin
+            dmax = self.dmax * value.dmax
+
+            return UncertContainer(vals, dmin, dmax, mask=vals.mask)
+
+        elif isinstance(value, (float, int, np.floating, np.integer)):
+            vals = self.vals * value
+            dmin = self.dmin * value
+            dmax = self.dmax * value
+
+            return UncertContainer(vals, dmin, dmax, mask=vals.mask)
+        else:
+            raise TypeError('Attempted to multiply by value of unsupported type')
+
+    def __truediv__(self, value):
+        if isinstance(value, UncertContainer):
+            vals = self.vals / value.vals
+            dmin = self.dmin / value.dmax
+            dmax = self.dmax / value.dmin
+
+            return UncertContainer(vals, dmin, dmax, mask=vals.mask)
+        else:
+            raise TypeError('Attempted to divide by unsupported type')
+
+    def transpose(self):
+        vals = self.vals.T
+        dmin = self.dmin.T
+        dmax = self.dmax.T
+
+        return UncertContainer(vals, dmin, dmax, mask=vals.mask)
+
+    def recip(self):
+        vals = 1.0 / self.vals
+        dmin = 1.0 / self.dmax
+        dmax = 1.0 / self.dmin
+
+        return UncertContainer(vals, dmin, dmax, mask=vals.mask)
+
+    def update_mask(self):
+        self.vals.mask = self.mask
+        self.dmin.mask = self.mask
+        self.dmax.mask = self.mask
+        self.wt.mask = self.mask
+        self.uncert.mask = self.mask
+
+    def concatenate(self, value, axis=0):
+        """Concatentate UncertContainer value to self.
+        Assumes that if dimensions of self and value do not match, to
+        add a np.newaxis along axis of value
+        """
+
+        if isinstance(value, UncertContainer):
+            if value.vals.ndim == self.vals.ndim:
+                vals = value.vals
+                dmin = value.dmin
+                dmax = value.dmax
+                wt = value.wt
+                uncert = value.uncert
+                mask = value.mask
+            elif (value.vals.ndim + 1) == self.vals.ndim:
+                vals = ma.expand_dims(value.vals, axis)
+                dmin = ma.expand_dims(value.dmin, axis)
+                dmax = ma.expand_dims(value.dmax, axis)
+                wt = ma.expand_dims(value.wt, axis)
+                uncert = ma.expand_dims(value.uncert, axis)
+                mask = np.expand_dims(value.mask, axis)
+            else:
+                raise ValueError('Could not propery match dimensionality')
+
+            self.vals = ma.concatenate((self.vals, vals), axis=axis)
+            self.dmin = ma.concatenate((self.dmin, dmin), axis=axis)
+            self.dmax = ma.concatenate((self.dmax, dmax), axis=axis)
+            self.wt = ma.concatenate((self.wt, wt), axis=axis)
+            self.uncert = ma.concatenate((self.uncert, uncert), axis=axis)
+
+            self.mask = np.concatenate((self.mask, mask), axis=axis)
+        else:
+            raise ValueError('Can only concatenate with an UncertContainer object')
+
+    def weighted_average(self, axis=0, expaxis=None):
+        """Calculate weighted average of data along axis
+        after optionally inserting a new dimension into the
+        shape array at position expaxis
+        """
+
+        if expaxis is not None:
+            vals = ma.expand_dims(self.vals, expaxis)
+            dmin = ma.expand_dims(self.dmin, expaxis)
+            dmax = ma.expand_dims(self.dmax, expaxis)
+            wt = ma.expand_dims(self.wt, expaxis)
+        else:
+            vals = self.vals
+            wt = self.wt
+            dmin = self.dmin
+            dmax = self.dmax
+
+        # Get average value
+        avg, norm = ma.average(vals, axis=axis, weights=wt, returned=True)
+        avg_ex = ma.expand_dims(avg, 0)
+
+        # Calculate weighted uncertainty
+        wtmax = ma.max(wt, axis=axis)
+        neff = norm / wtmax  # Effective number of samples based on uncertainties
+
+        # Seeking max deviation from the average; if above avg use max, if below use min
+        term = np.empty_like(vals)
+
+        indices = np.where(vals > avg_ex)
+        i0 = indices[0]
+        irest = indices[1:]
+        ii = tuple(x for x in itertools.chain([i0], irest))
+        jj = tuple(x for x in itertools.chain([np.zeros_like(i0)], irest))
+        term[ii] = (dmax[ii] - avg_ex[jj]) ** 2
+
+        indices = np.where(vals <= avg_ex)
+        i0 = indices[0]
+        irest = indices[1:]
+        ii = tuple(x for x in itertools.chain([i0], irest))
+        jj = tuple(x for x in itertools.chain([np.zeros_like(i0)], irest))
+        term[ii] = (avg_ex[jj] - dmin[ii]) ** 2
+
+        dsum = ma.sum(term * wt, axis=0)  # Sum for weighted average of deviations
+
+        dev = 0.5 * np.sqrt(dsum / (norm * neff))
+
+        if isinstance(avg, (float, np.floating)):
+            avg = avg_ex
+
+        tmp_min = avg - dev
+        ii = np.where(tmp_min < 0)
+        tmp_min[ii] = TOL * avg[ii]
+
+        return UncertContainer(avg, tmp_min, avg + dev)

westpa/westext/weed/__init__.py

@@ -0,0 +1,10 @@
+'''westext.weed -- Support for weighted ensemble equilibrium dynamics
+
+Initial code by Dan Zuckerman (May 2011), integration by Matt Zwier,
+and testing by Carsen Stringer. Re-factoring and optimization of probability
+adjustment routines by Joshua L. Adelman (January 2012).
+'''
+
+from . import BinCluster, ProbAdjustEquil, UncertMath, weed_driver  # noqa
+from .ProbAdjustEquil import probAdjustEquil  # noqa
+from .weed_driver import WEEDDriver  # noqa