westpa 2022.10__cp310-cp310-macosx_11_0_arm64.whl → 2022.12__cp310-cp310-macosx_11_0_arm64.whl

westpa/core/binning/mab.py

@@ -1,73 +1,67 @@
 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 in between minimum and maximum segments along
-    the progress coordinte. Extrema and bottleneck segments are assigned
+    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,
-        direction=None,
-        skip=None,
-        bottleneck=True,
-        pca=False,
-        mab_log=False,
-        bin_log=False,
-        bin_log_path="$WEST_SIM_ROOT/binbounds.log",
+        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 int for nbins in each dimension.
-        direction : Union(list of int, None), default: None
-            List of int for 'direction' in each dimension.
-            Direction options are as follows:
+            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
-        skip : Union(list of int, None), default: None
-            List of int for each dimension. Default None for skip=0.
-            Set to 1 to 'skip' running mab in a dimension.
+                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 turn on or off bottleneck walker splitting.
+            Whether to enable bottleneck walker splitting.
         pca : bool, default: False
-            Can be True or False (default) to run PCA on pcoords before bin assignment.
+            Whether to perform PCA on progress coordinates before bin assignment.
         mab_log : bool, default: False
-            Whether to output mab info to west.log.
+            Whether to output MAB info to west.log.
         bin_log : bool, default: False
-            Whether to output mab bin boundaries to bin_log_path file.
+            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_per_dim is missing")
+            raise ValueError("nbins is missing")
         ndim = len(nbins)
 
-        if direction is None:
-            direction = [0] * ndim
-        elif len(direction) != ndim:
+        direction = direction or [0] * ndim
+        if len(direction) != ndim:
             direction = [0] * ndim
             log.warning("Direction list is not the correct dimensions, setting to defaults.")
 
-        if skip is None:
-            skip = [0] * ndim
-        elif len(skip) != ndim:
+        skip = skip or [0] * ndim
+        if len(skip) != ndim:
             skip = [0] * ndim
             log.warning("Skip list is not the correct dimensions, setting to defaults.")
 
@@ -86,71 +80,80 @@ class MABBinMapper(FuncBinMapper):
 
         super().__init__(map_mab, n_total_bins, kwargs=kwargs)
 
-    def determine_total_bins(self, nbins_per_dim, direction, skip, bottleneck, **kwargs):
+    def determine_total_bins(
+        self, nbins_per_dim: List[int], direction: List[int], skip: List[int], bottleneck: bool, **kwargs
+    ) -> int:
         """
-        The following is neccessary because functional bin mappers need to "reserve"
+        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 : int
-            Number of total bins in each direction.
+        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. See __init__ for more information.
+            Direction in each dimension.
         skip : list of int
-            List of 0s and 1s indicating whether to skip each dimension.
+            List indicating whether to skip each dimension.
         bottleneck : bool
-            Whether to include separate bin for bottleneck walker(s).
+            Whether to include a separate bin for bottleneck walker(s).
         **kwargs : dict
-            Arbitary keyword arguments. Contains unneeded MAB parameters.
+            Additional MAB parameters (unused).
 
         Returns
         -------
         n_total_bins : int
             Number of total bins.
-
         """
-        n_total_bins = np.prod(nbins_per_dim)
-        ndim = len(nbins_per_dim)
-        for i in range(ndim):
-            if skip[i] == 0:
-                if direction[i] != 0:
+        # 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
-                else:
+                elif direct == 0:
+                    # 2 lead/lag bins + 2 bottleneck bins
                     n_total_bins += 2 + 2 * bottleneck
-            else:
-                n_total_bins -= nbins_per_dim[i] - 1
-                n_total_bins += 1 * ndim  # or else it will be one bin short
+                elif direct == 86:
+                    # 0 lead/lag + 2 bottleneck bins
+                    n_total_bins += 2 * bottleneck
         return n_total_bins
 
 
-def map_mab(coords, mask, output, *args, **kwargs):
+def map_mab(coords: np.ndarray, mask: np.ndarray, output: List[int], *args, **kwargs) -> List[int]:
     """
-    Binning which adaptively places bins based on the positions of extrema segments and
-    bottleneck segments, which are where the difference in probability is the greatest
-    along the progress coordinate. Operates per dimension and places a fixed number of
+    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 : ndarray
+    coords : np.ndarray
         An array with pcoord and weight info.
-    mask : ndarray
-        Array of 1 (True) and 0 (False), to filter out unwanted segment 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
-        Variable length arguments.
+        Additional arguments.
     **kwargs : dict
-        Arbitary keyword arguments. Contains most of the MAB-needed parameters.
+        Additional keyword arguments. Contains most of the MAB-needed parameters.
 
     Returns
     ------
     output : list
-        The main list that, for each segment, holds the bin assignment.
-
+        List with bin assignments for each segment.
     """
 
     # Argument Processing
@@ -158,8 +161,8 @@ def map_mab(coords, mask, output, *args, **kwargs):
     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))
+    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")
@@ -170,8 +173,8 @@ def map_mab(coords, mask, output, *args, **kwargs):
     if skip is None:
         skip = [0] * ndim
 
-    allcoords = np.copy(coords)
-    allmask = np.copy(mask)
+    allcoords = coords.copy()
+    allmask = mask.copy()
 
     weights = None
     isfinal = None
@@ -185,11 +188,11 @@ def map_mab(coords, mask, output, *args, **kwargs):
         if coords[0, -1] == 0:
             report = True
         if coords.shape[1] > ndim + 1:
-            isfinal = allcoords[:, ndim + 1].astype(np.bool_)
+            isfinal = allcoords[:, ndim + 1].astype(bool)
         else:
-            isfinal = np.ones(coords.shape[0], dtype=np.bool_)
+            isfinal = np.ones(coords.shape[0], dtype=bool)
         coords = coords[isfinal, :ndim]
-        weights = allcoords[isfinal, ndim + 0]
+        weights = allcoords[isfinal, ndim]
         mask = mask[isfinal]
         splitting = True
 
@@ -199,229 +202,305 @@ def map_mab(coords, mask, output, *args, **kwargs):
         weights = None
         splitting = False
 
-    varcoords = np.copy(coords)
     originalcoords = np.copy(coords)
     if pca and len(output) > 1:
-        colavg = np.mean(coords, axis=0)
-        for i in range(len(coords)):
-            for j in range(len(coords[i])):
-                varcoords[i][j] = coords[i][j] - colavg[j]
-        covcoords = np.cov(np.transpose(varcoords), aweights=weights)
-        eigval, eigvec = np.linalg.eigh(covcoords)
-        eigvec = eigvec[:, np.argmax(np.absolute(eigvec), axis=1)]
-        for i in range(len(eigvec)):
-            if eigvec[i, i] < 0:
-                eigvec[:, i] = -1 * eigvec[:, i]
-        for i in range(ndim):
-            for j in range(len(output)):
-                coords[j][i] = np.dot(varcoords[j], eigvec[:, i])
-
-    maxlist = []
-    minlist = []
-    difflist = []
-    flipdifflist = []
-    for n in range(ndim):
-        # identify the boundary segments
-        maxcoord = np.max(coords[mask, n])
-        mincoord = np.min(coords[mask, n])
-        maxlist.append(maxcoord)
-        minlist.append(mincoord)
-
-        # detect the bottleneck segments, this uses the weights
-        if splitting:
-            temp = np.column_stack((originalcoords[mask, n], weights[mask]))
-            sorted_indices = temp[:, 0].argsort()
-            temp = temp[sorted_indices]
-            for p in range(len(temp)):
-                if temp[p][1] == 0:
-                    temp[p][1] = 10**-323
-            fliptemp = np.flipud(temp)
-
-            difflist.append(None)
-            flipdifflist.append(None)
-            maxdiff = 0
-            flipmaxdiff = 0
-            for i in range(1, len(temp) - 1):
-                comprob = 0
-                flipcomprob = 0
-                j = i + 1
-                while j < len(temp):
-                    comprob = comprob + temp[j][1]
-                    flipcomprob = flipcomprob + fliptemp[j][1]
-                    j = j + 1
-                diff = -np.log(comprob) + np.log(temp[i][1])
-                if diff > maxdiff:
-                    difflist[n] = temp[i][0]
-                    maxdiff = diff
-                flipdiff = -np.log(flipcomprob) + np.log(fliptemp[i][1])
-                if flipdiff > flipmaxdiff:
-                    flipdifflist[n] = fliptemp[i][0]
-                    flipmaxdiff = flipdiff
+        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:
-        westpa.rc.pstatus("################ MAB stats ################")
-        westpa.rc.pstatus("minima in each dimension:      {}".format(minlist))
-        westpa.rc.pstatus("maxima in each dimension:      {}".format(maxlist))
-        westpa.rc.pstatus("direction in each dimension:   {}".format(direction))
-        westpa.rc.pstatus("skip in each dimension:        {}".format(skip))
-        westpa.rc.pstatus("###########################################")
-        westpa.rc.pflush()
-
-    # assign segments to bins
-    # the total number of linear bins is the boundary base
-    boundary_base = np.prod(nbins_per_dim)
-
-    # the bottleneck base is offset by the number of boundary walkers,
-    # which is two per dimension unless there is a direction specified
-    # in a particluar dimension, then it's just one
-    bottleneck_base = boundary_base
+        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
 
-    for i in range(0, ndim):
-        # for single direction, 1 boundary walker
-        if direction[i] == 1 or direction[i] == -1:
-            bottleneck_base += 1
-        # 2 boundary walkers with 0 direction
-        elif direction[i] == 0:
-            bottleneck_base += 2
-        # for 86 direction, no boundary walkers so offset of 0
-        elif direction[i] == 86:
-            bottleneck_base += 0
-
-    # if a dimension is being "skipped", leave only one bin total as
-    # the offset
-    for i in range(0, ndim):
-        if skip[i] != 0:
-            boundary_base -= nbins_per_dim[i] - 1
+    # 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
 
-    for i in range(len(output)):
-        if not allmask[i]:
-            continue
+    # List of dimensions that are not skipped
+    active_dims = np.array([n for n in range(ndim) if not skip[n]])
 
-        # special means either a boundary or bottleneck walker (not a walker in the linear space)
-        special = False
-        # this holder is the bin number, which only needs to be unique for different walker groups
-        holder = 0
-        if splitting:
-            for n in range(ndim):
-                coord = allcoords[i][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()
 
-                # if skipped, just assign the walkers to the same bin (offset of boundary base)
-                if skip[n] != 0:
-                    holder = boundary_base + n
+    # 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
 
-                # assign bottlenecks, taking directionality into account
-                if bottleneck:
-                    if direction[n] == -1:
-                        if coord == flipdifflist[n]:
-                            holder = bottleneck_base + n
-                            special = True
-                            n_bottleneck_filled += 1
-                            break
-
-                    if direction[n] == 1:
-                        if coord == difflist[n]:
-                            holder = bottleneck_base + n
-                            special = True
-                            n_bottleneck_filled += 1
-                            break
-
-                    # both directions when using 0 or with
-                    # special value of 86 for no lead/lag split
-                    if direction[n] == 0 or direction[n] == 86:
-                        if coord == difflist[n]:
-                            holder = bottleneck_base + n
-                            special = True
-                            n_bottleneck_filled += 1
-                            break
-                        elif coord == flipdifflist[n]:
-                            holder = bottleneck_base + n + 1
-                            special = True
-                            n_bottleneck_filled += 1
-                            break
-
-                # assign boundary walkers, taking directionality into account
-                if direction[n] == -1:
-                    if coord == minlist[n]:
-                        holder = boundary_base + n
-                        special = True
-                        break
-
-                elif direction[n] == 1:
-                    if coord == maxlist[n]:
-                        holder = boundary_base + n
-                        special = True
-                        break
-
-                elif direction[n] == 0:
-                    if coord == minlist[n]:
-                        holder = boundary_base + n
-                        special = True
-                        break
-                    elif coord == maxlist[n]:
-                        holder = boundary_base + n + 1
-                        special = True
-                        break
-
-                # special value for direction with no lead/lag split
-                elif direction[n] == 86:
-                    # westpa.rc.pstatus(f"No lead/lag split for dim {n}")
-                    # westpa.rc.pflush()
-                    # nornmally adds to special bin but here just leaving it forever empty
-                    # holder = boundary_base + n
+        # 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
 
-        # the following are for the "linear" portion
+        # 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):
-                # if skipped, it's added to the same bin as the special walkers above
-                if skip[n] != 0:
-                    holder = boundary_base + n
-                    break
-
-                coord = allcoords[i][n]
+                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)
-                bin_number = np.digitize(coord, bins) - 1
-
-                if isfinal is None or not isfinal[i]:
-                    if bin_number >= nbins:
-                        bin_number = nbins - 1
-                    elif bin_number < 0:
-                        bin_number = 0
-                elif bin_number >= nbins or bin_number < 0:
-                    if np.isclose(bins[-1], coord):
-                        bin_number = nbins - 1
-                    elif np.isclose(bins[0], coord):
-                        bin_number = 0
-                    else:
-                        raise ValueError("Walker out of boundary")
-
-                holder += bin_number * np.prod(nbins_per_dim[:n])
-
-        # output is the main list that, for each segment, holds the bin assignment
-        output[i] = holder
-
-    if bin_log and report:
-        if westpa.rc.sim_manager.n_iter:
-            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('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'\nmin/max pcoord: {minlist} {maxlist}\n')
-                bb_file.write(f'bottleneck bins: {n_bottleneck_filled}\n')
-                if n_bottleneck_filled > 0:
-                    # Bottlenecks bins exist (passes any of the if bottleneck: checks)
-                    bb_file.write(f'bottleneck pcoord: {flipdifflist} {difflist}\n\n')
-                else:
-                    bb_file.write('\n')
 
-    return output
+                # 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')