Merge pull request westpa#339 from westpa/develop

Pushing recent `develop` changes to `westpa2`

.flake8

@@ -1,10 +1,17 @@
 [flake8]
 max-line-length = 132
 ignore =
-    E203,  # space before : (needed for how black formats slicing)
-    E266,  # too many leading '#' for block comment 
-    E501,  # line lengths too long
-    W503,  # line break before binary operator
-    W504,  # line break after binary operator
-    E402,  # module level import not at top of file
-    E731,  # do not assign a lambda expression, use a def
+    # space before : (needed for how black formats slicing)
+    E203,
+    # too many leading '#' for block comment 
+    E266,
+    # line lengths too long
+    E501,
+    # line break before binary operator
+    W503,
+    # line break after binary operator
+    W504,
+    # module level import not at top of file
+    E402,
+    # do not assign a lambda expression, use a def
+    E731,

.pre-commit-config.yaml

@@ -1,14 +1,14 @@
 repos:
   - repo: https://github.com/pre-commit/pre-commit-hooks
-    rev: v4.1.0
+    rev: v4.4.0
     hooks:
     -   id: end-of-file-fixer
         types: [python]
     -   id: trailing-whitespace
         types: [python]
 
   - repo: https://github.com/psf/black
-    rev: 22.3.0
+    rev: 23.7.0
     hooks:
       - id: black
         language_version: python3
@@ -19,7 +19,7 @@ repos:
               | doc
           )
   - repo: https://github.com/pycqa/flake8
-    rev: 4.0.1
+    rev: 6.1.0
     hooks:
     -   id: flake8
         exclude: > 
@@ -28,3 +28,6 @@ repos:
               | _version.py
               | doc
           )
+
+ci:
+  autoupdate_branch: 'develop'

examples/ex_graphtest.py

@@ -56,9 +56,9 @@ def get_all_leaves(n_iter, iter_group):
 # print(leaves_by_branchpoint)
 print('there are {:d} branch points leading to approximately independent trajectory bundles'.format(len(leaves_by_branchpoint)))
 # for ((branch,node),leaves) in leaves_by_branchpoint.iteritems():
-for (branchnode, subtrees) in leaves_by_branchpoint.items():
+for branchnode, subtrees in leaves_by_branchpoint.items():
     print('  branch point {}:{} has {} independent subtree(s)'.format(branchnode.n_iter, branchnode.seg_id, len(subtrees)))
-    for (subtree, leaves) in subtrees.items():
+    for subtree, leaves in subtrees.items():
         leaves = list(leaves)
         weights = np.fromiter((leaf.weight for leaf in leaves), dtype=np.float64)
         max_weight_node = leaves[np.argmax(weights)]
@@ -71,7 +71,7 @@ def get_all_leaves(n_iter, iter_group):
         # print('subtree rooted at point ({},{}) contains {:d} leaves'.format(node.n_iter, node.seg_id,len(leaves)))
         # print('  of which ({},{}) has the highest weight ({!r})'.format(max_weight_node.n_iter, max_weight_node.seg_id,
         #                                                                max_weight_node.weight))
-for (root, leaves) in leftovers_by_root.items():
+for root, leaves in leftovers_by_root.items():
     print('{} trajectories from root {}:{} pruned due to shared history'.format(len(leaves), root.n_iter, root.seg_id))
 
 

src/westpa/cli/core/w_fork.py

@@ -101,7 +101,7 @@ def entry_point():
     state_map = np.empty((n_segments,), dtype=state_map_dtype)
     state_map['old_n_iter'] = n_iter
 
-    for (iseg, (index_row, pcoord)) in enumerate(zip(old_index, old_final_pcoords)):
+    for iseg, (index_row, pcoord) in enumerate(zip(old_index, old_final_pcoords)):
         istate = istates[iseg]
         istate.iter_created = 0
         istate.iter_used = 1

src/westpa/cli/core/w_init.py

@@ -18,7 +18,6 @@
 
 
 def entry_point():
-
     parser = argparse.ArgumentParser(
         'w_init',
         description='''\

src/westpa/cli/core/w_run.py

@@ -30,7 +30,6 @@ def entry_point():
 
 
 def run_simulation():
-
     work_manager = westpa.rc.work_manager = make_work_manager()
 
     # Load the sim manager and other drivers

src/westpa/cli/core/w_states.py

@@ -102,7 +102,6 @@ def entry_point():
 
 # TODO: This would benefit from a refactor to set default args to None, and replace some of those "if <argument>" clauses
 def initialize(mode, bstates, _bstate_file, tstates, _tstate_file):
-
     work_manager = make_work_manager()
 
     system = westpa.rc.get_system_driver()
@@ -116,7 +115,6 @@ def initialize(mode, bstates, _bstate_file, tstates, _tstate_file):
 
             assert mode in ('show', 'replace', 'append')
             if mode == 'show':
-
                 basis_states = data_manager.get_basis_states(n_iter)
                 if basis_states:
                     bstate_file = sys.stdout if not _bstate_file else open(_bstate_file, 'wt')

src/westpa/cli/core/w_succ.py

@@ -57,7 +57,7 @@ def find_successful_trajs(self):
             # The above HDF5 selection always returns a vector; we want a 2-d array
             final_pcoords.shape = (len(recycled_seg_ids), pcoord_ndim)
 
-            for (ipc, seg_id) in enumerate(recycled_seg_ids):
+            for ipc, seg_id in enumerate(recycled_seg_ids):
                 self.output_file.write('%8d    %8d    %20.14g' % (n_iter, seg_id, seg_index[seg_id]['weight']))
                 fields = ['']
                 for field in final_pcoords[ipc]:

src/westpa/cli/tools/w_assign.py

@@ -33,7 +33,6 @@ def parse_pcoord_value(pc_str):
 def _assign_label_pop(
     n_iter, lb, ub, mapper, nstates, state_map, last_labels, parent_id_dsspec, weight_dsspec, pcoord_dsspec, subsample
 ):
-
     nbins = len(state_map) - 1
     parent_ids = parent_id_dsspec.get_iter_data(n_iter, index_exp[lb:ub])
     weights = weight_dsspec.get_iter_data(n_iter, index_exp[lb:ub])

src/westpa/cli/tools/w_bins.py

@@ -68,7 +68,6 @@ def add_args(self, parser):
         rebin_parser.set_defaults(func=self.cmd_rebin)
 
     def process_args(self, args):
-
         self.data_reader.process_args(args)
         self.data_reader.open(mode='r+')
         self.n_iter = getattr(args, 'n_iter', None) or self.data_reader.current_iteration

src/westpa/cli/tools/w_dumpsegs.py

@@ -64,7 +64,7 @@ def go(self):
             + '\n'
         )
         pcoord_lines = '  pcoord[0]  = {init_pcoord}\n  pcoord[-1] = {final_pcoord}' + '\n'
-        for (_seg_id, segment) in enumerate(segments):
+        for _seg_id, segment in enumerate(segments):
             parents_str = '[' + ', '.join(map(str, sorted(segment.wtg_parent_ids))) + ']'
             init_pcoord_str = '[' + ', '.join('{pcval:<12.6g}'.format(pcval=float(pce)) for pce in segment.pcoord[0]) + ']'
             final_pcoord_str = '[' + ', '.join('{pcval:<12.6g}'.format(pcval=float(pce)) for pce in segment.pcoord[-1]) + ']'

src/westpa/cli/tools/w_eddist.py

@@ -68,7 +68,6 @@ def _remote_min_max(ndim, dset_dtype, n_iter, dsspec):
 
 
 def _remote_bin_iter(iiter, n_iter, dsspec, wt_dsspec, initpoint, binbounds, ignore_out_of_range):
-
     iter_hist_shape = tuple(len(bounds) - 1 for bounds in binbounds)
     iter_hist = np.zeros(iter_hist_shape, dtype=np.float64)
 
@@ -220,7 +219,6 @@ def __init__(self):
         self.compress_output = False
 
     def add_args(self, parser):
-
         parser.add_argument(
             '-b',
             '--bins',
@@ -300,7 +298,6 @@ def process_args(self, args):
         self.compress_output = args.compress or False
 
     def go(self):
-
         pi = self.progress.indicator
         pi.operation = 'Initializing'
         with pi:
@@ -334,7 +331,6 @@ def go(self):
 
     @staticmethod
     def parse_binspec(binspec):
-
         namespace = {'numpy': np, 'np': np, 'inf': float('inf')}
 
         try:

src/westpa/cli/tools/w_ipa.py

@@ -641,7 +641,6 @@ def future(self, value=None):
         return self._future
 
     class Future(WIPIDataset):
-
         # This isn't a real fancy one.
         def __getitem__(self, value):
             if isinstance(value, str):

src/westpa/cli/tools/w_multi_west.py

@@ -156,7 +156,7 @@ def total_number_of_walkers(self):
             # Sometimes, we're smaller or larger by one.  Hm.
             try:
                 self.total_walkers[:] += west['summary'][:-1]['n_particles']
-            except (ValueError):
+            except ValueError:
                 self.total_walkers[:] += west['summary'][:-1]['n_particles'][: len(self.total_walkers)]
 
     class Segment:

src/westpa/cli/tools/w_pdist.py

@@ -51,7 +51,6 @@ def _remote_min_max(ndim, dset_dtype, n_iter, dsspec):
 
 
 def _remote_bin_iter(iiter, n_iter, dsspec, wt_dsspec, initpoint, binbounds, ignore_out_of_range):
-
     iter_hist_shape = tuple(len(bounds) - 1 for bounds in binbounds)
     iter_hist = np.zeros(iter_hist_shape, dtype=np.float64)
 

src/westpa/cli/tools/w_red.py

@@ -333,7 +333,6 @@ def _get_corrector(self):
         return self._dc
 
     def calc_rate(self, i_iter=None, red=False, **kwargs):
-
         if i_iter is None:
             i_iter = self.n_iters
 

src/westpa/core/_rc.py

@@ -644,7 +644,6 @@ def get_work_manager(self):
         return self.work_manager
 
     def clear_state(self):
-
         self._sim_manager = None
         self._system = None
         self._data_manager = None

src/westpa/core/binning/assign.py

@@ -429,7 +429,7 @@ def assign(self, coords, mask=None, output=None):
         # Which coordinates do we need to reassign, because they landed in
         # bins with embedded mappers?
         rmasks = {}
-        for (rindex, mapper) in self._recursion_targets.items():
+        for rindex, mapper in self._recursion_targets.items():
             omask = output == rindex
             mmask |= omask
             rmasks[rindex] = omask
@@ -441,7 +441,7 @@ def assign(self, coords, mask=None, output=None):
             output_map(output, omap, mask & ~mmask)
 
         # do any recursive assignments necessary
-        for (rindex, mapper) in self._recursion_targets.items():
+        for rindex, mapper in self._recursion_targets.items():
             mapper.assign(coords, mask & rmasks[rindex], output)
 
         return output

src/westpa/core/binning/binless_driver.py

@@ -33,7 +33,7 @@ def assign(self, segments, initializing=False):
         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):
+        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

src/westpa/core/binning/binless_manager.py

@@ -75,8 +75,31 @@ def propagate(self):
                 log.error('unknown future {!r} received from work manager'.format(future))
                 raise AssertionError('untracked future {!r}'.format(future))
 
+        # Collectively assign all segments to their bins...
         self.we_driver.assign(self.segments.values())
-        self.get_istate_futures()
+
+        # For cases where we need even more istates for recycled trajectories
+        # futures should be empty at this point.
+        istate_gen_futures = self.get_istate_futures()
+        futures.update(istate_gen_futures)
+
+        # Wait for istate_gen_futures and catch untracked futures.
+        while futures:
+            future = self.work_manager.wait_any(futures)
+            futures.remove(future)
+
+            if future in istate_gen_futures:
+                istate_gen_futures.remove(future)
+                _basis_state, initial_state = future.get_result()
+                log.debug('received newly-prepared initial state {!r}'.format(initial_state))
+                initial_state.istate_status = InitialState.ISTATE_STATUS_PREPARED
+                with self.data_manager.expiring_flushing_lock():
+                    self.data_manager.update_initial_states([initial_state], n_iter=self.n_iter + 1)
+                self.we_driver.avail_initial_states[initial_state.state_id] = initial_state
+            else:
+                log.error('unknown future {!r} received from work manager'.format(future))
+                raise AssertionError('untracked future {!r}'.format(future))
+
         log.debug('done with propagation')
         self.save_bin_data()
         self.data_manager.flush_backing()
@@ -132,7 +155,7 @@ def prepare_iteration(self):
         for iseg, segment in enumerate(segments.values()):
             initial_pcoords[iseg] = segment.pcoord[0]
         initial_assignments = self.system.bin_mapper.assign(initial_pcoords)
-        for (segment, assignment) in zip(iter(segments.values()), initial_assignments):
+        for segment, assignment in zip(iter(segments.values()), initial_assignments):
             initial_binning[assignment].add(segment)
         self.report_bin_statistics(initial_binning, [], save_summary=True)
         del initial_pcoords, initial_binning

src/westpa/core/binning/mab.py

@@ -212,7 +212,6 @@ def map_mab(coords, mask, output, *args, **kwargs):
         # the following are for the "linear" portion
         if not special:
             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
@@ -262,13 +261,13 @@ def __init__(self, nbins, direction=None, skip=None, bottleneck=True, pca=False,
             direction = [0] * ndim
         elif len(direction) != ndim:
             direction = [0] * ndim
-            log.warn("Direction list is not the correct dimensions, setting to defaults.")
+            log.warning("Direction list is not the correct dimensions, setting to defaults.")
 
         if skip is None:
             skip = [0] * ndim
         elif len(skip) != ndim:
             skip = [0] * ndim
-            log.warn("Skip list is not the correct dimensions, setting to defaults.")
+            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)
         # the following is neccessary because functional bin mappers need to "reserve"

src/westpa/core/binning/mab_driver.py

@@ -33,7 +33,7 @@ def assign(self, segments, initializing=False):
         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):
+        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

src/westpa/core/binning/mab_manager.py

@@ -75,8 +75,31 @@ def propagate(self):
                 log.error('unknown future {!r} received from work manager'.format(future))
                 raise AssertionError('untracked future {!r}'.format(future))
 
+        # Collectively assign all segments to their bins...
         self.we_driver.assign(self.segments.values())
-        self.get_istate_futures()
+
+        # For cases where we need even more istates for recycled trajectories
+        # futures should be empty at this point.
+        istate_gen_futures = self.get_istate_futures()
+        futures.update(istate_gen_futures)
+
+        # Wait for istate_gen_futures and catch untracked futures.
+        while futures:
+            future = self.work_manager.wait_any(futures)
+            futures.remove(future)
+
+            if future in istate_gen_futures:
+                istate_gen_futures.remove(future)
+                _basis_state, initial_state = future.get_result()
+                log.debug('received newly-prepared initial state {!r}'.format(initial_state))
+                initial_state.istate_status = InitialState.ISTATE_STATUS_PREPARED
+                with self.data_manager.expiring_flushing_lock():
+                    self.data_manager.update_initial_states([initial_state], n_iter=self.n_iter + 1)
+                self.we_driver.avail_initial_states[initial_state.state_id] = initial_state
+            else:
+                log.error('unknown future {!r} received from work manager'.format(future))
+                raise AssertionError('untracked future {!r}'.format(future))
+
         log.debug('done with propagation')
         self.save_bin_data()
         self.data_manager.flush_backing()
@@ -135,7 +158,7 @@ def prepare_iteration(self):
         # Assign this iteration's segments' initial points to bins and report on bin population
         initial_binning = self.system.bin_mapper.construct_bins()
         initial_assignments = self.system.bin_mapper.assign(pcoords_with_weights)
-        for (segment, assignment) in zip(iter(segments.values()), initial_assignments):
+        for segment, assignment in zip(iter(segments.values()), initial_assignments):
             initial_binning[assignment].add(segment)
         self.report_bin_statistics(initial_binning, [], save_summary=True)
         del pcoords_with_weights, initial_binning