Cache node attribute fields

Author: InnovativeInventor

gerrychain/updaters/tally.py

@@ -9,6 +9,11 @@ class DataTally:
     """An updater for tallying numerical data that is not necessarily stored as
     node attributes
     """
+    __slots__ = [
+        "data",
+        "alias",
+        "_call"
+    ]
 
     def __init__(self, data, alias):
         """
@@ -54,6 +59,11 @@ def __call__(self, partition, previous=None):
 class Tally:
     """An updater for keeping a tally of one or more node attributes.
     """
+    __slots__ = [
+        "fields",
+        "alias",
+        "dtype"
+    ]
 
     def __init__(self, fields, alias=None, dtype=int):
         """
@@ -117,12 +127,12 @@ def _update_tally(self, partition):
         return new_tally
 
     def _get_tally_from_node(self, partition, node):
-        return sum(partition.graph.nodes[node][field] for field in self.fields)
+        return sum(partition.graph.cached_node_data_lookup[node][field] for field in self.fields)
 
 
 def compute_out_flow(graph, fields, flow):
-    return sum(graph.nodes[node][field] for node in flow["out"] for field in fields)
+    return sum(graph.cached_node_data_lookup[node][field] for node in flow["out"] for field in fields)
 
 
 def compute_in_flow(graph, fields, flow):
-    return sum(graph.nodes[node][field] for node in flow["in"] for field in fields)
+    return sum(graph.cached_node_data_lookup[node][field] for node in flow["in"] for field in fields)

tests/test_tree.py

@@ -5,6 +5,7 @@
 
 from gerrychain import MarkovChain
 from gerrychain.constraints import contiguous, within_percent_of_ideal_population
+from gerrychain.graph import Graph
 from gerrychain.partition import Partition
 from gerrychain.proposals import recom
 from gerrychain.tree import (
@@ -22,7 +23,7 @@
 def graph_with_pop(three_by_three_grid):
     for node in three_by_three_grid:
         three_by_three_grid.nodes[node]["pop"] = 1
-    return three_by_three_grid
+    return Graph.from_networkx(three_by_three_grid)
 
 
 @pytest.fixture
@@ -41,7 +42,7 @@ def twelve_by_twelve_with_pop():
     grid = networkx.relabel_nodes(xy_grid, nodes)
     for node in grid:
         grid.nodes[node]["pop"] = 1
-    return grid
+    return Graph.from_networkx(grid)
 
 
 def test_bipartition_tree_returns_a_subset_of_nodes(graph_with_pop):