Merge pull request #383 from mggg/perf-assignment-flows-opt

Refactor out Assignment.update (replace with Assignment.update_flows) to prevent double-calculating flows

gerrychain/partition/assignment.py

@@ -4,8 +4,6 @@
 
 import pandas
 
-from ..updaters.flows import flows_from_changes
-
 
 class Assignment(Mapping):
     """An assignment of nodes into parts.
@@ -58,10 +56,9 @@ def copy(self):
         """
         return Assignment(self.parts.copy(), self.mapping.copy(), validate=False)
 
-    def update(self, mapping):
-        """Update the assignment for some nodes using the given mapping.
+    def update_flows(self, flows):
+        """Update the assignment for some nodes using the given flows.
         """
-        flows = flows_from_changes(self, mapping)
         for part, flow in flows.items():
             # Union between frozenset and set returns an object whose type
             # matches the object on the left, which here is a frozenset

gerrychain/partition/partition.py

@@ -75,14 +75,14 @@ def _from_parent(self, parent, flips):
         self.parent = parent
         self.flips = flips
 
-        self.assignment = parent.assignment.copy()
-        self.assignment.update(flips)
-
         self.graph = parent.graph
         self.updaters = parent.updaters
 
         self.flows = flows_from_changes(parent.assignment, flips)
 
+        self.assignment = parent.assignment.copy()
+        self.assignment.update_flows(self.flows)
+
         if "cut_edges" in self.updaters:
             self.edge_flows = compute_edge_flows(self)
 

tests/constraints/test_validity.py

@@ -169,6 +169,6 @@ def test_no_vanishing_districts_works():
     partition = MagicMock()
     partition.parent = parent
     partition.assignment = parent.assignment.copy()
-    partition.assignment.update({2: 1})
+    partition.assignment.update_flows({1: {"out": set(), "in": {2}}, 2: {"out": {2}, "in": set()}})
 
     assert not no_vanishing_districts(partition)

tests/partition/test_assignment.py

@@ -12,7 +12,7 @@ def assignment():
 
 class TestAssignment:
     def test_assignment_can_be_updated(self, assignment):
-        assignment.update({2: 1})
+        assignment.update_flows({1: {"out": set(), "in": {2}}, 2: {"out": {2}, "in": set()}})
         assert assignment[2] == 1
 
     def test_assignment_copy_does_not_copy_the_node_sets(self, assignment):