Use node_repeats in _bipartitions_tree_random_all (#406)

mggg · Apr 29, 2023 · 04ff4f9 · 04ff4f9
1 parent ee00d5d
commit 04ff4f9
Showing 1 changed file with 7 additions and 2 deletions.
diff --git a/gerrychain/tree.py b/gerrychain/tree.py
@@ -253,18 +253,23 @@ def _bipartition_tree_random_all(
         spanning_tree = spanning_tree_fn(graph)
 
     repeat = True
+    restarts = 0
     attempts = 0
     while max_attempts is None or attempts < max_attempts:
-        spanning_tree = spanning_tree_fn(graph)
+        if restarts == node_repeats:
+            spanning_tree = spanning_tree_fn(graph)
+            restarts = 0
         h = PopulatedGraph(spanning_tree, populations, pop_target, epsilon)
         possible_cuts = balance_edge_fn(h, choice=choice)
 
         repeat = repeat_until_valid
-        attempts += 1
 
         if not (repeat and len(possible_cuts) == 0):
             return possible_cuts
 
+        restarts += 1
+        attempts += 1
+
     raise RuntimeError(f"Could not find a possible cut after {max_attempts} attempts.")