westpa 2022.12__cp313-cp313-macosx_11_0_arm64.whl

westpa/core/binning/mab.py

@@ -0,0 +1,506 @@
+import logging
+from typing import List, Optional
+import numpy as np
+import westpa
+from westpa.core.binning import FuncBinMapper
+from os.path import expandvars
+
+log = logging.getLogger(__name__)
+
+
+class MABBinMapper(FuncBinMapper):
+    """
+    Adaptively place bins between minimum and maximum segments along
+    the progress coordinate. Extrema and bottleneck segments are assigned
+    to their own bins.
+    """
+
+    def __init__(
+        self,
+        nbins: List[int],
+        direction: Optional[List[int]] = None,
+        skip: Optional[List[int]] = None,
+        bottleneck: bool = True,
+        pca: bool = False,
+        mab_log: bool = False,
+        bin_log: bool = False,
+        bin_log_path: str = "$WEST_SIM_ROOT/binbounds.log",
+    ):
+        """
+        Parameters
+        ----------
+        nbins : list of int
+            List of number of bins in each dimension.
+        direction : Optional[list of int], default: None
+            List of directions in each dimension. Direction options:
+                0   : default split at leading and lagging boundaries
+                1   : split at leading boundary only
+                -1  : split at lagging boundary only
+                86  : no splitting at either leading or lagging boundary (both bottlenecks included)
+        skip : Optional[list of int], default: None
+            List of skip flags for each dimension. Default None (no skipping).
+        bottleneck : bool, default: True
+            Whether to enable bottleneck walker splitting.
+        pca : bool, default: False
+            Whether to perform PCA on progress coordinates before bin assignment.
+        mab_log : bool, default: False
+            Whether to output MAB info to west.log.
+        bin_log : bool, default: False
+            Whether to output MAB bin boundaries to a log file.
+        bin_log_path : str, default: "$WEST_SIM_ROOT/binbounds.log"
+            Path to output bin boundaries.
+        """
+        # Verifying parameters
+        if nbins is None:
+            raise ValueError("nbins is missing")
+        ndim = len(nbins)
+
+        direction = direction or [0] * ndim
+        if len(direction) != ndim:
+            direction = [0] * ndim
+            log.warning("Direction list is not the correct dimensions, setting to defaults.")
+
+        skip = skip or [0] * ndim
+        if len(skip) != ndim:
+            skip = [0] * ndim
+            log.warning("Skip list is not the correct dimensions, setting to defaults.")
+
+        kwargs = dict(
+            nbins_per_dim=nbins,
+            direction=direction,
+            skip=skip,
+            bottleneck=bottleneck,
+            pca=pca,
+            mab_log=mab_log,
+            bin_log=bin_log,
+            bin_log_path=bin_log_path,
+        )
+
+        n_total_bins = self.determine_total_bins(**kwargs)
+
+        super().__init__(map_mab, n_total_bins, kwargs=kwargs)
+
+    def determine_total_bins(
+        self, nbins_per_dim: List[int], direction: List[int], skip: List[int], bottleneck: bool, **kwargs
+    ) -> int:
+        """
+        Calculate the total number of bins needed, taking direction and skipping into account.
+        This function is necessary because functional bin mappers need to "reserve"
+        bins and tell the sim manager how many bins they will need to use, this is
+        determined by taking all direction/skipping info into account.
+
+        Parameters
+        ----------
+        nbins_per_dim : list of int
+            Number of total bins in each dimension within the linear portion.
+        direction : list of int
+            Direction in each dimension.
+        skip : list of int
+            List indicating whether to skip each dimension.
+        bottleneck : bool
+            Whether to include a separate bin for bottleneck walker(s).
+        **kwargs : dict
+            Additional MAB parameters (unused).
+
+        Returns
+        -------
+        n_total_bins : int
+            Number of total bins.
+        """
+        # Update nbins_per_dim with any skipped dimensions, setting number of bins along skipped dimensions to 1
+        skip = np.array([bool(s) for s in skip])
+        nbins_per_dim = np.array(nbins_per_dim)
+        nbins_per_dim[skip] = 1
+
+        # Total bins is product of all linear bins plus and special bins
+        n_total_bins = nbins_per_dim.prod()
+        for direct, skip_dim in zip(direction, skip):
+            if not skip_dim:
+                if direct in [-1, 1]:
+                    # 1 lead or lag bin + 1 bottleneck bin
+                    n_total_bins += 1 + 1 * bottleneck
+                elif direct == 0:
+                    # 2 lead/lag bins + 2 bottleneck bins
+                    n_total_bins += 2 + 2 * bottleneck
+                elif direct == 86:
+                    # 0 lead/lag + 2 bottleneck bins
+                    n_total_bins += 2 * bottleneck
+        return n_total_bins
+
+
+def map_mab(coords: np.ndarray, mask: np.ndarray, output: List[int], *args, **kwargs) -> List[int]:
+    """
+    Adaptively place bins based on extrema and bottleneck segments along the progress coordinate.
+
+    Bottleneck segments are where the difference in probability is the greatest
+    along the progress coordinate. Operates per dimension (unless skipped) and places a fixed number of
+    evenly spaced bins between the segments with the min and max pcoord values. Extrema and
+    bottleneck segments are assigned their own bins.
+
+    Parameters
+    ----------
+    coords : np.ndarray
+        An array with pcoord and weight info.
+    mask : np.ndarray
+        Boolean array to filter out unwanted segments.
+    output : list
+        The main list that, for each segment, holds the bin assignment.
+    *args : list
+        Additional arguments.
+    **kwargs : dict
+        Additional keyword arguments. Contains most of the MAB-needed parameters.
+
+    Returns
+    ------
+    output : list
+        List with bin assignments for each segment.
+    """
+
+    # Argument Processing
+    nbins_per_dim = kwargs.get("nbins_per_dim")
+    ndim = len(nbins_per_dim)
+    pca = kwargs.get("pca", False)
+    bottleneck = kwargs.get("bottleneck", True)
+    direction = kwargs.get("direction", [0] * ndim)
+    skip = kwargs.get("skip", [0] * ndim)
+    mab_log = kwargs.get("mab_log", False)
+    bin_log = kwargs.get("bin_log", False)
+    bin_log_path = kwargs.get("bin_log_path", "$WEST_SIM_ROOT/binbounds.log")
+
+    if not np.any(mask):
+        return output
+
+    if skip is None:
+        skip = [0] * ndim
+
+    allcoords = coords.copy()
+    allmask = mask.copy()
+
+    weights = None
+    isfinal = None
+    splitting = False
+    report = False
+
+    # the segments should be sent in by the driver as half initial segments and half final segments
+    # allcoords contains all segments
+    # coords should contain ONLY final segments
+    if coords.shape[1] > ndim:
+        if coords[0, -1] == 0:
+            report = True
+        if coords.shape[1] > ndim + 1:
+            isfinal = allcoords[:, ndim + 1].astype(bool)
+        else:
+            isfinal = np.ones(coords.shape[0], dtype=bool)
+        coords = coords[isfinal, :ndim]
+        weights = allcoords[isfinal, ndim]
+        mask = mask[isfinal]
+        splitting = True
+
+    if not np.any(mask):
+        coords = allcoords[:, :ndim]
+        mask = allmask
+        weights = None
+        splitting = False
+
+    originalcoords = np.copy(coords)
+    if pca and len(output) > 1:
+        coords = apply_pca(coords, weights)
+
+    # Computing special bins (bottleneck and boundary bins)
+    minlist, maxlist, bottlenecks_forward, bottlenecks_reverse = calculate_bin_boundaries(
+        originalcoords, weights, mask, skip, splitting, bottleneck
+    )
+
+    if mab_log and report:
+        log_mab_stats(minlist, maxlist, direction, skip)
+
+    # Assign segments to bins
+    n_bottleneck_filled = bin_assignment(
+        allcoords,
+        allmask,
+        minlist,
+        maxlist,
+        bottlenecks_forward,
+        bottlenecks_reverse,
+        nbins_per_dim,
+        direction,
+        skip,
+        splitting,
+        bottleneck,
+        output,
+    )
+
+    # Report MAB bin statistics
+    if bin_log and report and westpa.rc.sim_manager.n_iter:
+        log_bin_boundaries(
+            skip,
+            bottleneck,
+            direction,
+            bin_log_path,
+            minlist,
+            maxlist,
+            nbins_per_dim,
+            n_bottleneck_filled,
+            bottlenecks_forward,
+            bottlenecks_reverse,
+        )
+
+    return output
+
+
+def apply_pca(coords, weights):
+    colavg = np.mean(coords, axis=0)
+    varcoords = coords - colavg
+    covcoords = np.cov(varcoords.T, aweights=weights)
+    eigval, eigvec = np.linalg.eigh(covcoords)
+    eigvec = eigvec[:, np.argmax(np.abs(eigvec), axis=1)]
+    eigvec[:, np.diag(eigvec) < 0] *= -1
+    return np.dot(varcoords, eigvec)
+
+
+def calculate_bin_boundaries(coords, weights, mask, skip, splitting, bottleneck):
+    """
+    This function calculates minima, maxima, and bottleneck segments.
+    """
+    skip = np.array([bool(s) for s in skip])
+
+    # Initialize lists to hold minima and maxima along each dimension
+    minlist, maxlist = [], []
+    # Initialize lists to hold bottleneck segments along each dimension
+    bottlenecks_forward, bottlenecks_reverse = [None] * len(coords[0]), [None] * len(coords[0])
+    # number of unmasked coords
+    n_coords = mask.sum()
+    # Grabbing all unmasked coords and weights
+    unmasked_coords = coords[mask, :]
+    unmasked_weights = weights[mask] if weights is not None else None
+    # Replace any zero weights with non-zero values so that log(weight) is well-defined
+    if unmasked_weights is not None:
+        unmasked_weights[unmasked_weights == 0] = 10**-323
+    # Looping over each dimension of progress coordinate, even those being skipped
+    for n in range(len(coords[0])):
+        # We calculate the min and max pcoord along each dimension (boundary segments) even if skipping
+        maxlist.append(np.max(coords[mask, n]))
+        minlist.append(np.min(coords[mask, n]))
+        # Now we calculate the bottleneck segments
+        if splitting and bottleneck and not skip[n]:
+            bottlenecks_forward[n], bottlenecks_reverse[n] = detect_bottlenecks(unmasked_coords, unmasked_weights, n_coords, n)
+
+    return minlist, maxlist, bottlenecks_forward, bottlenecks_reverse
+
+
+def detect_bottlenecks(unmasked_coords, unmasked_weights, n_coords, n):
+    """
+    Detect the bottleneck segments along the given coordinate n, this uses the weights
+    """
+    # Grabbing all unmasked coords in current dimension, plus corresponding weights
+    # Sort by current dimension in coord, smallest to largest
+    sorted_indices = unmasked_coords[:, n].argsort(kind='stable')
+
+    # Grab sorted coords and weights
+    coords_srt = unmasked_coords[sorted_indices, :]
+    weights_srt = unmasked_weights[sorted_indices]
+
+    # Also sort in reverse order for opposite direction
+    coords_srt_flip = np.flipud(coords_srt)
+    weights_srt_flip = np.flipud(weights_srt)
+
+    # Initialize the max directional differences along current dimension as None (these may not be updated)
+    bottleneck_coords, bottleneck_coords_flip = None, None
+    maxdiff, maxdiff_flip = -np.inf, -np.inf
+
+    # Looping through all non-boundary coords
+    # Compute the cumulative weight on either side of each non-boundary walker
+    for i in range(1, n_coords - 1):
+        # Summing up weights of all walkers ahead of current walker along current dim in both directions
+        cumulative_prob = np.sum(weights_srt[i + 1 :])
+        cumulative_prob_flip = np.sum(weights_srt_flip[i + 1 :])
+        # Compute the difference of log cumulative weight of current walker and all walkers ahead of it (Z im the MAB paper)
+        # We use the log as weights vary over many orders of magnitude
+        # Note a negative Z indicates the cumulative weight ahead of the current walker is larger than the weight of the current walker,
+        # while a positive Z indicates the cumulative weight ahead of the current walker is smaller, indicating a barrier
+        Z = np.log(weights_srt[i]) - np.log(cumulative_prob)
+        Z_flip = np.log(weights_srt_flip[i]) - np.log(cumulative_prob_flip)
+        # Update ALL coords of the current walker into bottlenecks_forward if it is largest
+        # This way we uniquely identify a walker by its full set of coordinates
+        if Z > maxdiff:
+            bottleneck_coords = coords_srt[i, :]
+            maxdiff = Z
+        if Z_flip > maxdiff_flip:
+            bottleneck_coords_flip = coords_srt_flip[i, :]
+            maxdiff_flip = Z_flip
+    return bottleneck_coords, bottleneck_coords_flip
+
+
+def log_mab_stats(minlist, maxlist, direction, skip):
+    westpa.rc.pstatus("################ MAB stats ################")
+    westpa.rc.pstatus(f"minima in each dimension:      {minlist}")
+    westpa.rc.pstatus(f"maxima in each dimension:      {maxlist}")
+    westpa.rc.pstatus(f"direction in each dimension:   {direction}")
+    westpa.rc.pstatus(f"skip in each dimension:        {skip}")
+    westpa.rc.pstatus("###########################################")
+    westpa.rc.pflush()
+
+
+def bin_assignment(
+    coords,
+    mask,
+    minlist,
+    maxlist,
+    bottlenecks_forward,
+    bottlenecks_reverse,
+    nbins_per_dim,
+    direction,
+    skip,
+    splitting,
+    bottleneck,
+    output,
+):
+    """
+    Assign segments to bins based on the minima, maxima, and
+    bottleneck segments along the progress coordinate.
+    """
+    # Update nbins_per_dim with any skipped dimensions, setting number of bins along skipped dimensions to 1
+    skip = np.array([bool(s) for s in skip])
+    nbins_per_dim = np.array(nbins_per_dim)
+    nbins_per_dim[skip] = 1
+    direction = np.array(direction)
+
+    ndim = len(nbins_per_dim)
+    n_bottleneck_filled = 0
+
+    # Boolean arrays that track use of special bins along each dimension
+    skip_bneck_fwd = np.array([d == -1 if bottleneck else True for d in direction]) + skip
+    skip_bneck_rev = np.array([d == 1 if bottleneck else True for d in direction]) + skip
+    skip_lead = np.array([d in [86, -1] for d in direction]) + skip
+    skip_lag = np.array([d in [86, 1] for d in direction]) + skip
+
+    # List of dimensions that are not skipped
+    active_dims = np.array([n for n in range(ndim) if not skip[n]])
+
+    # Compute the boundary bin ID offsets
+    # In forward direction, this is all the linear bins
+    boundary_bin_id_offset_fwd = nbins_per_dim.prod()
+    # In reverse, we add the number of forward boundary bins to the offset
+    boundary_bin_id_offset_rev = boundary_bin_id_offset_fwd + (~skip_lead).sum()
+
+    # Compute the bottleneck bin ID offsets
+    # In forward direction, bin IDs are offset by all linear and boundary bins
+    bneck_bin_id_offset_fwd = boundary_bin_id_offset_rev + (~skip_lag).sum()
+    # In reverse, we add the number of forward bottleneck bins to the offset
+    bneck_bin_id_offset_rev = bneck_bin_id_offset_fwd + (~skip_bneck_fwd).sum()
+
+    # Bin assignment loop over all walkers
+    for i in range(len(output)):
+        # Skip masked walkers, these walkers bin IDs are unchanged
+        if not mask[i]:
+            continue
+        # Initialize bin ID and special tracker for current coord
+        # The special variable indicates a boundary or bottleneck walker (not assigned to the linear space)
+        bin_id, special = 0, False
+
+        # Searching for bottleneck bins first
+        if splitting and bottleneck:
+            for n in active_dims:
+                # Grab coord(s) of current walker
+                coord = coords[i][:ndim]
+                # Assign bottlenecks, taking directionality into account
+                # Check both directions when using 0 or 86
+                # Note: 86 implies no leading or lagging bins, but does add bottlenecks for *both* directions when bottleneck is enabled
+                # Note: All bottleneck bins will typically be filled unless a walker is simultaneously in bottleneck bins along multiple dimensions
+                # or there are too few walkers to compute free energy barriers
+                if (coord == bottlenecks_forward[n]).all() and not skip_bneck_fwd[n]:
+                    bin_id = bneck_bin_id_offset_fwd + n - skip_bneck_fwd[:n].sum()
+                    special = True
+                    n_bottleneck_filled += 1
+                    break
+                elif (coord == bottlenecks_reverse[n]).all() and not skip_bneck_rev[n]:
+                    bin_id = bneck_bin_id_offset_rev + n - skip_bneck_rev[:n].sum()
+                    special = True
+                    n_bottleneck_filled += 1
+                    break
+
+        # Now check for boundary walkers, taking directionality into account
+        # This should only be done after fully checking for bottleneck walkers
+        if splitting and not special:
+            for n in active_dims:
+                # Grab coord of current walker along current dimension
+                coord = coords[i, n]
+                if (coord == maxlist[n]) and not skip_lead[n]:
+                    bin_id = boundary_bin_id_offset_fwd + n - skip_lead[:n].sum()
+                    special = True
+                    break
+                elif (coord == minlist[n]) and not skip_lag[n]:
+                    bin_id = boundary_bin_id_offset_rev + n - skip_lag[:n].sum()
+                    special = True
+                    break
+
+        # Now check for linear bin walkers
+        if not special:
+            # Again we loop over the dimensions
+            # Note: no need to worry about skipping as we've already set all skipped dimensions to 1 bin
+            for n in range(ndim):
+                coord = coords[i][n]
+                nbins = nbins_per_dim[n]
+                minp = minlist[n]
+                maxp = maxlist[n]
+
+                # Generate the bins along this dimension
+                bins = np.linspace(minp, maxp, nbins + 1)
+
+                # Assign walker to a bin along this dimension
+                bin_number = np.digitize(coord, bins) - 1  # note np.digitize is 1-indexed
+
+                # Sometimes the walker is exactly at the max/min value,
+                # which would put it in the next bin
+                if bin_number == nbins:
+                    bin_number -= 1
+                elif bin_number == -1:
+                    bin_number = 0
+                elif bin_number > nbins or bin_number < -1:
+                    raise ValueError("Walker out of boundary.")
+
+                # Assign to bin within the full dimensional space
+                bin_id += bin_number * np.prod(nbins_per_dim[:n])
+
+        # Output is the main list that, for each segment, holds the bin assignment
+        output[i] = bin_id
+    return n_bottleneck_filled
+
+
+def log_bin_boundaries(
+    skip,
+    bottleneck,
+    direction,
+    bin_log_path,
+    minlist,
+    maxlist,
+    nbins_per_dim,
+    n_bottleneck_filled,
+    bottlenecks_forward,
+    bottlenecks_reverse,
+):
+    ndim = len(nbins_per_dim)
+    skip = np.array([bool(s) for s in skip])
+    active_dims = np.array([n for n in range(ndim) if not skip[n]])
+    max_bottleneck = np.sum([1 if direction[n] in [-1, 1] else 2 for n in active_dims]) if bottleneck else 0
+    with open(expandvars(bin_log_path), 'a') as bb_file:
+        # Iteration Number
+        bb_file.write(f'Iteration: {westpa.rc.sim_manager.n_iter}\n')
+        bb_file.write('MAB linear bin boundaries: ')
+        for n in range(ndim):
+            # Write binbounds per dim
+            bb_file.write(f'{np.linspace(minlist[n], maxlist[n], nbins_per_dim[n] + 1)}\t')
+        # Min/Max pcoord
+        bb_file.write(f'\nLagging pcoord in each dimension: {minlist}\n')
+        bb_file.write(f'Leading pcoord in each dimension: {maxlist}\n')
+        # Bottlenecks bins exist
+        if bottleneck:
+            bb_file.write(f'Number of bottleneck bins filled: {n_bottleneck_filled} / {max_bottleneck}\n')
+            for n in active_dims:
+                if direction[n] in [0, 1, 86]:
+                    bb_file.write(f'Dimension {n} forward bottleneck walker at: {list(bottlenecks_forward[n])}\n')
+                if direction[n] in [0, -1, 86]:
+                    bb_file.write(f'Dimension {n} backward bottleneck walker at: {list(bottlenecks_reverse[n])}\n')
+            bb_file.write('\n')
+        else:
+            bb_file.write('\n')

westpa/core/binning/mab_driver.py

@@ -0,0 +1,54 @@
+import logging
+
+import numpy as np
+from westpa.core.we_driver import WEDriver
+
+log = logging.getLogger(__name__)
+
+
+class MABDriver(WEDriver):
+    def assign(self, segments, initializing=False):
+        '''Assign segments to initial and final bins, and update the (internal) lists of used and available
+        initial states. This function is adapted to the MAB scheme, so that the inital and final segments are
+        sent to the bin mapper at the same time, otherwise the inital and final bin boundaries can be inconsistent.'''
+
+        log.debug("MABDriver in use.")
+        # collect initial and final coordinates into one place
+        n_segments = len(segments)
+        all_pcoords = np.empty((n_segments * 2, self.system.pcoord_ndim + 2), dtype=self.system.pcoord_dtype)
+
+        for iseg, segment in enumerate(segments):
+            all_pcoords[iseg] = np.append(segment.pcoord[0, :], [segment.weight, 0.0])
+            all_pcoords[n_segments + iseg] = np.append(segment.pcoord[-1, :], [segment.weight, 1.0])
+
+        # assign based on initial and final progress coordinates
+        assignments = self.bin_mapper.assign(all_pcoords)
+        initial_assignments = assignments[:n_segments]
+        if initializing:
+            final_assignments = initial_assignments
+        else:
+            final_assignments = assignments[n_segments:]
+
+        initial_binning = self.initial_binning
+        final_binning = self.final_binning
+        flux_matrix = self.flux_matrix
+        transition_matrix = self.transition_matrix
+        for segment, iidx, fidx in zip(segments, initial_assignments, final_assignments):
+            initial_binning[iidx].add(segment)
+            final_binning[fidx].add(segment)
+            flux_matrix[iidx, fidx] += segment.weight
+            transition_matrix[iidx, fidx] += 1
+
+        n_recycled_total = self.n_recycled_segs
+        n_new_states = n_recycled_total - len(self.avail_initial_states)
+
+        log.debug(
+            '{} walkers scheduled for recycling, {} initial states available'.format(
+                n_recycled_total, len(self.avail_initial_states)
+            )
+        )
+
+        if n_new_states > 0:
+            return n_new_states
+        else:
+            return 0