From a3000e38588ce42153f22ab93eb5cca84645cdad Mon Sep 17 00:00:00 2001 From: Max Fan Date: Thu, 13 Jan 2022 21:23:17 -0500 Subject: [PATCH] Maybe faster record perf (todo: benchmark carefully) --- gerrychain/proposals/tree_proposals.py | 7 ++++--- gerrychain/tree.py | 5 +++-- 2 files changed, 7 insertions(+), 5 deletions(-) diff --git a/gerrychain/proposals/tree_proposals.py b/gerrychain/proposals/tree_proposals.py index fe4525fe..bba9e3ff 100644 --- a/gerrychain/proposals/tree_proposals.py +++ b/gerrychain/proposals/tree_proposals.py @@ -45,16 +45,17 @@ def recom( partition.parts[parts_to_merge[0]] | partition.parts[parts_to_merge[1]] ) - flips = recursive_tree_part( + flips_left, flips_right = bipartition_tree_retworkx( subgraph, - parts_to_merge, pop_col=pop_col, pop_target=pop_target, epsilon=epsilon, node_repeats=node_repeats, - method=method, ) + flips = {node: parts_to_merge[0] for node in flips_left} + flips |= {node: parts_to_merge[1] for node in flips_right} + return partition.flip(flips) diff --git a/gerrychain/tree.py b/gerrychain/tree.py index aefde56d..d442e593 100644 --- a/gerrychain/tree.py +++ b/gerrychain/tree.py @@ -182,8 +182,9 @@ def bipartition_tree_retworkx( ): pops = graph.pygraph_pop_lookup(pop_col) - balanced_nodes = retworkx.bipartition_tree(graph.pygraph, lambda x: random.random(), pops, pop_target, epsilon) - return {graph.retworkx_networkx_mapping[x] for x in choice(balanced_nodes)[1]} + balanced_node_choices = retworkx.bipartition_tree(graph.pygraph, lambda x: random.random(), pops, float(pop_target), float(epsilon)) + balanced_nodes = {graph.retworkx_networkx_mapping[x] for x in choice(balanced_node_choices)[1]} + return (balanced_nodes, graph.node_indicies - balanced_nodes) def bipartition_tree(