Change flows_from_changes to expect partitions as arguments to enable caching (#384)

Author: InnovativeInventor

gerrychain/partition/partition.py

@@ -78,7 +78,7 @@ def _from_parent(self, parent, flips):
         self.graph = parent.graph
         self.updaters = parent.updaters
 
-        self.flows = flows_from_changes(parent.assignment, flips)
+        self.flows = flows_from_changes(parent, self)  # careful
 
         self.assignment = parent.assignment.copy()
         self.assignment.update_flows(self.flows)

gerrychain/updaters/flows.py

@@ -16,10 +16,11 @@ def create_flow():
     return {'in': set(), 'out': set()}
 
 
-def flows_from_changes(old_assignment, flips):
+@functools.lru_cache(maxsize=2)
+def flows_from_changes(old_partition, new_partition):
     flows = collections.defaultdict(create_flow)
-    for node, target in flips.items():
-        source = old_assignment.mapping[node]
+    for node, target in new_partition.flips.items():
+        source = old_partition.assignment.mapping[node]
         if source != target:
             flows[target]['in'].add(node)
             flows[source]['out'].add(node)

gerrychain/updaters/tally.py

@@ -102,14 +102,13 @@ def _update_tally(self, partition):
         :param partition: :class:`Partition` class.
         """
         parent = partition.parent
-        flips = partition.flips
 
         old_tally = parent[self.alias]
         new_tally = dict(old_tally)
 
         graph = partition.graph
 
-        for part, flow in flows_from_changes(parent.assignment, flips).items():
+        for part, flow in flows_from_changes(parent, partition).items():
             out_flow = compute_out_flow(graph, self.fields, flow)
             in_flow = compute_in_flow(graph, self.fields, flow)
             new_tally[part] = old_tally[part] - out_flow + in_flow