westpa 2022.12__cp313-cp313-macosx_10_13_x86_64.whl

westpa/core/binning/__init__.py

@@ -0,0 +1,55 @@
+from . import _assign
+from . import assign, bins
+
+from .assign import (
+    NopMapper,
+    FuncBinMapper,
+    PiecewiseBinMapper,
+    RectilinearBinMapper,
+    RecursiveBinMapper,
+    VectorizingFuncBinMapper,
+    VoronoiBinMapper,
+)
+
+from .mab import map_mab, MABBinMapper
+from .binless import map_binless, BinlessMapper
+
+from .mab_driver import MABDriver
+from .mab_manager import MABSimManager
+from .binless_manager import BinlessSimManager
+from .binless_driver import BinlessDriver
+
+from ._assign import accumulate_labeled_populations, assign_and_label, accumulate_state_populations_from_labeled
+from ._assign import assignments_list_to_table
+
+from .assign import coord_dtype, index_dtype
+from .bins import Bin
+
+
+__all__ = [
+    '_assign',
+    'assign',
+    'bins',
+    'NopMapper',
+    'FuncBinMapper',
+    'PiecewiseBinMapper',
+    'RectilinearBinMapper',
+    'RecursiveBinMapper',
+    'VectorizingFuncBinMapper',
+    'VoronoiBinMapper',
+    'map_mab',
+    'map_binless',
+    'MABBinMapper',
+    'BinlessMapper',
+    'MABDriver',
+    'MABSimManager',
+    'BinlessDriver',
+    'BinlessSimManager',
+    'accumulate_labeled_populations',
+    'assign_and_label',
+    'accumulate_state_populations_from_labeled',
+    'assignments_list_to_table',
+    'coord_dtype',
+    'index_dtype',
+    'Bin',
+]

westpa/core/binning/assign.py

@@ -0,0 +1,455 @@
+'''
+Bin assignment for WEST simulations. This module defines "bin mappers" which take
+vectors of coordinates (or rather, coordinate tuples), and assign each a definite
+integer value identifying a bin. Critical portions are implemented in a Cython
+extension module.
+
+A number of pre-defined bin mappers are available here:
+
+  * :class:`RectilinearBinMapper`, for bins divided by N-dimensional grids
+  * :class:`FuncBinMapper`, for functions which directly calculate bin assignments
+    for a number of coordinate values. This is best used with C/Cython/Numba
+    functions, or intellegently-tuned numpy-based Python functions.
+  * :class:`VectorizingFuncBinMapper`, for functions which calculate a bin
+    assignment for a single coordinate value. This is best used for arbitrary
+    Python functions.
+  * :class:`PiecewiseBinMapper`, for using a set of boolean-valued functions, one
+    per bin, to determine assignments. This is likely to be much slower than a
+    `FuncBinMapper` or `VectorizingFuncBinMapper` equipped with an appropriate
+    function, and its use is discouraged.
+
+One "super-mapper" is available, for assembling more complex bin spaces from
+simpler components:
+
+  * :class:`RecursiveBinMapper`, for nesting one set of bins within another.
+
+Users are also free to implement their own mappers. A bin mapper must implement, at
+least, an ``assign(coords, mask=None, output=None)`` method, which is responsible
+for mapping each of the vector of coordinate tuples ``coords`` to an integer
+(np.uint16) indicating a what bin that coordinate tuple falls into. The optional
+``mask`` (a numpy bool array) specifies that some coordinates are to be skipped; this is used,
+for instance, by the recursive (nested) bin mapper to minimize the number of calculations
+required to definitively assign a coordinate tuple to a bin. Similarly, the optional
+``output`` must be an integer (uint16) array of the same length as ``coords``, into which
+assignments are written. The ``assign()`` function must return a reference to ``output``.
+(This is used to avoid allocating many temporary output arrays in complex binning
+scenarios.)
+
+A user-defined bin mapper must also make an ``nbins`` property available, containing
+the total number of bins within the mapper.
+
+'''
+
+import hashlib
+import logging
+import pickle
+
+import numpy as np
+
+from .bins import Bin
+from ._assign import output_map, apply_down, apply_down_argmin_across, rectilinear_assign
+
+# All bin numbers are 16-bit unsigned ints, with one element (65525) reserved to
+# indicate unknown or unassigned points. This allows up to 65,536 bins, making
+# rate and flux matrices up to 32 GB (2**32 elements * 8 bytes). If you need more
+# bins, change index_dtype here and index_dtype and index_t in _assign.pyx.
+index_dtype = np.uint16
+UNKNOWN_INDEX = 65535
+
+# All coordinates are currently 32-bit floats. If you need 64-bit, change
+# coord_dtype here and coord_t in _assign.pyx.
+coord_dtype = np.float32
+
+
+log = logging.getLogger(__name__)
+
+
+class BinMapper:
+    hashfunc = hashlib.sha256
+
+    def __init__(self):
+        self.labels = None
+        self.nbins = 0
+
+    def construct_bins(self, type_=Bin):
+        '''Construct and return an array of bins of type ``type``'''
+        return np.array([type_() for _i in range(self.nbins)], dtype=np.object_)
+
+    def pickle_and_hash(self):
+        '''Pickle this mapper and calculate a hash of the result (thus identifying the
+        contents of the pickled data), returning a tuple ``(pickled_data, hash)``.
+        This will raise PickleError if this mapper cannot be pickled, in which case
+        code that would otherwise rely on detecting a topology change must assume
+        a topology change happened, even if one did not.
+        '''
+
+        pkldat = pickle.dumps(self, pickle.HIGHEST_PROTOCOL)
+        hash = self.hashfunc(pkldat)
+        return (pkldat, hash.hexdigest())
+
+    def __repr__(self):
+        return '<{} at 0x{:x} with {:d} bins>'.format(self.__class__.__name__, id(self), self.nbins or 0)
+
+
+class NopMapper(BinMapper):
+    '''Put everything into one bin.'''
+
+    def __init__(self):
+        super().__init__()
+        self.nbins = 1
+        self.labels = ['nop']
+
+    def assign(self, coords, mask=None, output=None):
+        if output is None:
+            output = np.zeros((len(coords),), dtype=index_dtype)
+
+        if mask is None:
+            mask = np.ones((len(coords),), dtype=np.bool_)
+        else:
+            mask = np.require(mask, dtype=np.bool_)
+
+        output[mask] = 0
+
+
+class RectilinearBinMapper(BinMapper):
+    '''Bin into a rectangular grid based on tuples of float values'''
+
+    def __init__(self, boundaries):
+        super().__init__()
+        self._boundaries = None
+        self._boundlens = None
+        self.ndim = 0
+        self.nbins = 0
+
+        # the setter function below handles all of the required wrangling
+        self.boundaries = boundaries
+
+    @property
+    def boundaries(self):
+        return self._boundaries
+
+    @boundaries.setter
+    def boundaries(self, boundaries):
+        del self._boundaries, self.labels
+        self._boundaries = []
+        self.labels = []
+        for boundset in boundaries:
+            boundarray = np.ascontiguousarray(boundset, dtype=coord_dtype)
+            db = np.diff(boundarray)
+            if (db <= 0).any():
+                raise ValueError('boundary set must be strictly monotonically increasing')
+            self._boundaries.append(boundarray)
+        self._boundlens = np.array([len(boundset) for boundset in self._boundaries], dtype=index_dtype)
+        self.ndim = len(self._boundaries)
+        self.nbins = np.multiply.accumulate([1] + [len(bounds) - 1 for bounds in self._boundaries])[-1]
+
+        _boundaries = self._boundaries
+        binspace_shape = tuple(self._boundlens[:] - 1)
+        for index in np.ndindex(binspace_shape):
+            label = (
+                '['
+                + ', '.join(
+                    f'({boundarray[index[idim]]!s}, {boundarray[index[idim] + 1]!s})' for idim, boundarray in enumerate(_boundaries)
+                )
+                + ']'
+            )
+            self.labels.append(label)
+
+    def assign(self, coords, mask=None, output=None):
+        try:
+            passed_coord_dtype = coords.dtype
+        except AttributeError:
+            coords = np.require(coords, dtype=coord_dtype)
+        else:
+            if passed_coord_dtype != coord_dtype:
+                coords = np.require(coords, dtype=coord_dtype)
+
+        if coords.ndim != 2:
+            raise TypeError('coords must be 2-dimensional')
+
+        if mask is None:
+            mask = np.ones((len(coords),), dtype=np.bool_)
+        elif len(mask) != len(coords):
+            raise TypeError('mask [shape {}] has different length than coords [shape {}]'.format(mask.shape, coords.shape))
+
+        if output is None:
+            output = np.empty((len(coords),), dtype=index_dtype)
+        elif len(output) != len(coords):
+            raise TypeError('output has different length than coords')
+
+        rectilinear_assign(coords, mask, output, self.boundaries, self._boundlens)
+
+        return output
+
+
+class PiecewiseBinMapper(BinMapper):
+    '''Binning using a set of functions returing boolean values; if the Nth function
+    returns True for a coordinate tuple, then that coordinate is in the Nth bin.'''
+
+    def __init__(self, functions):
+        self.functions = functions
+        self.nbins = len(functions)
+        self.index_dtype = np.min_scalar_type(self.nbins)
+        self.labels = [str(func) for func in functions]
+
+    def assign(self, coords, mask=None, output=None):
+        if output is None:
+            output = np.zeros((len(coords),), dtype=index_dtype)
+
+        if mask is None:
+            mask = np.ones((len(coords),), dtype=np.bool_)
+        else:
+            mask = np.require(mask, dtype=np.bool_)
+
+        coord_subset = coords[mask]
+        fnvals = np.empty((len(coord_subset), len(self.functions)), dtype=index_dtype)
+        for ifn, fn in enumerate(self.functions):
+            rsl = np.apply_along_axis(fn, 0, coord_subset)
+            if rsl.ndim > 1:
+                # this should work like a squeeze, unless the function returned something truly
+                # stupid (e.g., a 3d array with at least two dimensions greater than 1), in which
+                # case a broadcast error will occur
+                fnvals[:, ifn] = rsl.flat
+            else:
+                fnvals[:, ifn] = rsl
+        amask = np.require(fnvals.argmax(axis=1), dtype=index_dtype)
+        output[mask] = amask
+        return output
+
+
+class FuncBinMapper(BinMapper):
+    '''Binning using a custom function which must iterate over input coordinate
+    sets itself.'''
+
+    def __init__(self, func, nbins, args=None, kwargs=None):
+        self.func = func
+        self.nbins = nbins
+        self.args = args or ()
+        self.kwargs = kwargs or {}
+        self.labels = ['{!r} bin {:d}'.format(func, ibin) for ibin in range(nbins)]
+
+    def assign(self, coords, mask=None, output=None):
+        try:
+            passed_coord_dtype = coords.dtype
+        except AttributeError:
+            coords = np.require(coords, dtype=coord_dtype)
+        else:
+            if passed_coord_dtype != coord_dtype:
+                coords = np.require(coords, dtype=coord_dtype)
+
+        if coords.ndim != 2:
+            raise TypeError('coords must be 2-dimensional')
+        if mask is None:
+            mask = np.ones((len(coords),), dtype=np.bool_)
+        elif len(mask) != len(coords):
+            raise TypeError('mask [shape {}] has different length than coords [shape {}]'.format(mask.shape, coords.shape))
+
+        if output is None:
+            output = np.empty((len(coords),), dtype=index_dtype)
+        elif len(output) != len(coords):
+            raise TypeError('output has different length than coords')
+
+        self.func(coords, mask, output, *self.args, **self.kwargs)
+
+        return output
+
+
+class VectorizingFuncBinMapper(BinMapper):
+    '''Binning using a custom function which is evaluated once for each (unmasked)
+    coordinate tuple provided.'''
+
+    def __init__(self, func, nbins, args=None, kwargs=None):
+        self.func = func
+        self.args = args or ()
+        self.kwargs = kwargs or {}
+        self.nbins = nbins
+        self.index_dtype = np.min_scalar_type(self.nbins)
+        self.labels = ['{!r} bin {:d}'.format(func, ibin) for ibin in range(nbins)]
+
+    def assign(self, coords, mask=None, output=None):
+        try:
+            passed_coord_dtype = coords.dtype
+        except AttributeError:
+            coords = np.require(coords, dtype=coord_dtype)
+        else:
+            if passed_coord_dtype != coord_dtype:
+                coords = np.require(coords, dtype=coord_dtype)
+
+        if coords.ndim != 2:
+            raise TypeError('coords must be 2-dimensional')
+        if mask is None:
+            mask = np.ones((len(coords),), dtype=np.bool_)
+        elif len(mask) != len(coords):
+            raise TypeError('mask [shape {}] has different length than coords [shape {}]'.format(mask.shape, coords.shape))
+
+        if output is None:
+            output = np.empty((len(coords),), dtype=index_dtype)
+        elif len(output) != len(coords):
+            raise TypeError('output has different length than coords')
+
+        apply_down(self.func, self.args, self.kwargs, coords, mask, output)
+
+        return output
+
+
+class VoronoiBinMapper(BinMapper):
+    '''A one-dimensional mapper which assigns a multidimensional pcoord to the
+    closest center based on a distance metric. Both the list of centers and the
+    distance function must be supplied.'''
+
+    def __init__(self, dfunc, centers, dfargs=None, dfkwargs=None):
+        self.dfunc = dfunc
+        self.dfargs = dfargs or ()
+        self.dfkwargs = dfkwargs or {}
+        self.centers = np.asarray(centers)
+        self.nbins = self.centers.shape[0]
+        self.ndim = self.centers.shape[1]
+        self.labels = ['center={!r}'.format(center) for center in self.centers]
+
+        # Sanity check: does the distance map the centers to themselves?
+        check = self.assign(self.centers)
+        if (check != np.arange(len(self.centers))).any():
+            raise TypeError('dfunc does not map centers to themselves')
+
+    def assign(self, coords, mask=None, output=None):
+        try:
+            passed_coord_dtype = coords.dtype
+        except AttributeError:
+            coords = np.require(coords, dtype=coord_dtype)
+        else:
+            if passed_coord_dtype != coord_dtype:
+                coords = np.require(coords, dtype=coord_dtype)
+
+        if coords.ndim != 2:
+            raise TypeError('coords must be 2-dimensional')
+        if mask is None:
+            mask = np.ones((len(coords),), dtype=np.bool_)
+        elif len(mask) != len(coords):
+            raise TypeError('mask [shape {}] has different length than coords [shape {}]'.format(mask.shape, coords.shape))
+
+        if output is None:
+            output = np.empty((len(coords),), dtype=index_dtype)
+        elif len(output) != len(coords):
+            raise TypeError('output has different length than coords')
+
+        apply_down_argmin_across(self.dfunc, (self.centers,) + self.dfargs, self.dfkwargs, self.nbins, coords, mask, output)
+
+        return output
+
+
+class RecursiveBinMapper(BinMapper):
+    '''Nest mappers one within another.'''
+
+    def __init__(self, base_mapper, start_index=0):
+        self.base_mapper = base_mapper
+        self.nbins = base_mapper.nbins
+
+        # Targets for recursion
+        self._recursion_targets = {}
+
+        # Which bins must we recurse into?
+        self._recursion_map = np.zeros((self.base_mapper.nbins,), dtype=np.bool_)
+
+        self.start_index = start_index
+
+    @property
+    def labels(self):
+        for ilabel in range(self.base_mapper.nbins):
+            if self._recursion_map[ilabel]:
+                for label in self._recursion_targets[ilabel].labels:
+                    yield label
+            else:
+                yield self.base_mapper.labels[ilabel]
+
+    @property
+    def start_index(self):
+        return self._start_index
+
+    @start_index.setter
+    def start_index(self, new_index):
+        self._start_index = new_index
+        not_recursed = ~self._recursion_map
+        n_not_recursed = not_recursed.sum()
+        if n_not_recursed == self.nbins:
+            self._output_map = np.arange(self._start_index, self._start_index + self.nbins, dtype=index_dtype)
+        elif n_not_recursed > 0:
+            # This looks like uninitialized access, but self._output_map is always set during __init__
+            # (by self.start_index = 0, or whatever value was passed in), so this modifies the existing
+            # set chosen above
+            self._output_map[not_recursed] = np.arange(self._start_index, self._start_index + n_not_recursed, dtype=index_dtype)
+        else:
+            # No un-replaced bins
+            self._output_map = None
+
+        n_own_bins = self.base_mapper.nbins - self._recursion_map.sum()
+        startindex = self.start_index + n_own_bins
+        for mapper in self._recursion_targets.values():
+            mapper.start_index = startindex
+            startindex += mapper.nbins
+
+    def add_mapper(self, mapper, replaces_bin_at):
+        '''Replace the bin containing the coordinate tuple ``replaces_bin_at`` with the
+        specified ``mapper``.'''
+
+        replaces_bin_at = np.require(replaces_bin_at, dtype=coord_dtype)
+        if replaces_bin_at.ndim < 1:
+            replaces_bin_at.shape = (1, 1)
+        elif replaces_bin_at.ndim < 2:
+            replaces_bin_at.shape = (1, replaces_bin_at.shape[0])
+        elif replaces_bin_at.ndim > 2 or replaces_bin_at.shape[1] > 1:
+            raise TypeError('a single coordinate vector is required')
+
+        self.nbins += mapper.nbins - 1
+
+        ibin = self.base_mapper.assign(replaces_bin_at)[0]
+        log.debug('replacing bin {!r} containing {!r} with {!r}'.format(ibin, replaces_bin_at, mapper))
+        if self._recursion_map[ibin]:
+            # recursively add; this doesn't change anything for us except our
+            # total bin count, which has been accounted for above
+            self._recursion_targets[ibin].add_mapper(mapper, replaces_bin_at[0])
+        else:
+            # replace a bin on our mapper
+            self._recursion_map[ibin] = True
+            mapper = RecursiveBinMapper(mapper)
+            self._recursion_targets[ibin] = mapper
+
+        # we have updated our list of recursed bins, so set our own start index to trigger a recursive
+        # reassignment of mappers' output values
+        # Note that we're reordering the recursion targets based on outer bin numbers (dict keys) first,
+        # so the order the mappers were added no longer matters...
+        self._recursion_targets = {k: self._recursion_targets[k] for k in sorted(self._recursion_targets)}
+        self.start_index = self.start_index
+
+    def assign(self, coords, mask=None, output=None):
+        if mask is None:
+            mask = np.ones((len(coords),), dtype=np.bool_)
+
+        if output is None:
+            output = np.empty((len(coords),), dtype=index_dtype)
+
+        # mapping mask -- which output values come from our base
+        # region set and therefore must be remapped
+        mmask = np.zeros((len(coords),), dtype=np.bool_)
+
+        # Assign based on this mapper
+        self.base_mapper.assign(coords, mask, output)
+
+        # Which coordinates do we need to reassign, because they landed in
+        # bins with embedded mappers?
+        rmasks = {}
+        for rindex, mapper in self._recursion_targets.items():
+            omask = output == rindex
+            mmask |= omask
+            rmasks[rindex] = omask
+
+        # remap output from our (base) mapper
+        # omap may be None if every bin has a recursive mapper in it
+        omap = self._output_map
+        if omap is not None:
+            output_map(output, omap, mask & ~mmask)
+
+        # do any recursive assignments necessary
+        for rindex, mapper in self._recursion_targets.items():
+            mapper.assign(coords, mask & rmasks[rindex], output)
+
+        return output

westpa/core/binning/binless.py

@@ -0,0 +1,96 @@
+import numpy as np
+from westpa.core.binning import FuncBinMapper
+from westpa.core.extloader import get_object
+import logging
+
+log = logging.getLogger(__name__)
+
+
+def map_binless(coords, mask, output, *args, **kwargs):
+    '''Adaptively groups walkers according to a user-defined grouping function
+    that is defined externally. Very general implementation but limited to
+    only a two dimensional progress coordinate (for now).'''
+
+    n_groups = kwargs.get("n_groups")
+    n_dims = kwargs.get("n_dims")
+    group_function = get_object(kwargs.get("group_function"))
+    log.debug(f'binless arguments: {kwargs}')
+    try:
+        group_function_kwargs = kwargs.get('group_function_kwargs')['group_arguments']
+    except KeyError:
+        group_function_kwargs = {}
+    ndim = n_dims
+
+    if not np.any(mask):
+        return output
+
+    allcoords = np.copy(coords)
+    allmask = np.copy(mask)
+
+    isfinal = None
+    splitting = 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.shape[1] > ndim + 1:
+            isfinal = allcoords[:, ndim + 1].astype(np.bool_)
+        else:
+            isfinal = np.ones(coords.shape[0], dtype=np.bool_)
+        coords = coords[isfinal, :ndim]
+        mask = mask[isfinal]
+        splitting = True
+
+    # in case where there is no final segments but initial ones in range
+    if not np.any(mask):
+        coords = allcoords[:, :ndim]
+        mask = allmask
+        splitting = False
+
+    # filter the list of coordinates (which contains coordinates outside of the binless region)
+    # to obtain only the ones we want to cluster
+    # this is done with all dimensions at once
+    binless_coords = coords[mask]
+    nsegs_binless = len(binless_coords)
+
+    # we need to make sure that the number of segments in the binless region is greater than
+    # the number of clusters we request
+    # if only one segment in the binless region, assign it to a single cluster
+    if nsegs_binless == 1:
+        clusters = [0]
+
+    # if there are more than one segment in the binless region but still the total is less than
+    # our target number, adjust our target number to be the number of segments in the binless
+    # region minus one
+    elif nsegs_binless < n_groups:
+        clusters = group_function(binless_coords, nsegs_binless, splitting, **group_function_kwargs)
+    # if there are enough segments in the binless region, proceed as planned
+    elif nsegs_binless >= n_groups:
+        clusters = group_function(binless_coords, n_groups, splitting, **group_function_kwargs)
+
+    # this is a good place to say this... output is a list which matches the length of allcoords
+    # allcoords is a collection of all initial and final segment coords for that iteration
+    # we first filtered those to only contain the final data points, since those are the ones we care
+    # about clustering
+    # we then filtered to only have the coords in the binless region, since, again, those are what we care about
+    # we then assigned each to a cluster which is essentially a slitting index
+    # all that's left is to find where each binless segment is in the output and insert the cluster index there
+    for idx, val in enumerate(binless_coords):
+        if ndim > 1:
+            mask2 = np.logical_and(allcoords[:, 0] == val[0], allcoords[:, 1] == val[1])
+        else:
+            mask2 = allcoords[:, 0] == val[0]
+        output[mask2] = clusters[idx]
+
+    return output
+
+
+class BinlessMapper(FuncBinMapper):
+    '''Adaptively group walkers according to a user-defined grouping
+    function that is defined externally.'''
+
+    def __init__(self, ngroups, ndims, group_function, **group_function_kwargs):
+        kwargs = dict(n_groups=ngroups, n_dims=ndims, group_function=group_function, group_function_kwargs=group_function_kwargs)
+        n_total_groups = ngroups
+        super().__init__(map_binless, n_total_groups, kwargs=kwargs)

westpa/core/binning/binless_driver.py

@@ -0,0 +1,54 @@
+import logging
+
+import numpy as np
+from westpa.core.we_driver import WEDriver
+
+log = logging.getLogger(__name__)
+
+
+class BinlessDriver(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("BinlessDriver 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