Skip to content

Commit cc7be93

Browse files
Switch to using more performant/preferred lookup call
1 parent 15d4cff commit cc7be93

File tree

1 file changed

+5
-5
lines changed

1 file changed

+5
-5
lines changed

gerrychain/tree.py

+5-5
Original file line numberDiff line numberDiff line change
@@ -202,7 +202,7 @@ def bipartition_tree(
202202
tree is not provided
203203
:param choice: :func:`random.choice`. Can be substituted for testing.
204204
"""
205-
populations = {node: graph.nodes[node][pop_col] for node in graph.node_indicies}
205+
populations = {node: graph.lookup(node, pop_col) for node in graph.node_indicies}
206206

207207
possible_cuts = []
208208
if spanning_tree is None:
@@ -232,7 +232,7 @@ def _bipartition_tree_random_all(
232232
choice=random.choice,
233233
):
234234
"""Randomly bipartitions a graph and returns all cuts."""
235-
populations = {node: graph.nodes[node][pop_col] for node in graph.node_indicies}
235+
populations = {node: graph.lookup(node, pop_col) for node in graph.node_indicies}
236236

237237
possible_cuts = []
238238
if spanning_tree is None:
@@ -343,7 +343,7 @@ def recursive_tree_part(
343343
part_pop = 0
344344
for node in nodes:
345345
flips[node] = part
346-
part_pop += graph.nodes[node][pop_col]
346+
part_pop += graph.lookup(node, pop_col)
347347
debt += part_pop - pop_target
348348
remaining_nodes -= nodes
349349

@@ -386,7 +386,7 @@ def get_seed_chunks(
386386

387387
chunk_pop = 0
388388
for node in graph.node_indicies:
389-
chunk_pop += graph.nodes[node][pop_col]
389+
chunk_pop += graph.lookup(node, pop_col)
390390

391391
while True:
392392
epsilon = abs(epsilon)
@@ -426,7 +426,7 @@ def get_seed_chunks(
426426

427427
part_pop = 0
428428
for node in remaining_nodes:
429-
part_pop += graph.nodes[node][pop_col]
429+
part_pop += graph.lookup(node, pop_col)
430430
part_pop_as_dist = part_pop / num_chunks_left
431431
fake_epsilon = epsilon
432432
if num_chunks_left != 1:

0 commit comments

Comments
 (0)