Repository: mggg/GerryChain

Issue:
Four of six default proposal functions require a `cut_edges` updater
The `propose_chunk_flip`, `propose_random_flip`, `slow_reversible_propose_bi`, and `slow_reversible_propose` proposal functions all require that a `cut_edges` updater is present on the `partition` parameter passed to the function.

Returns a `KeyError`.

Comments:

Patch:
diff --git a/gerrychain/constraints/validity.py b/gerrychain/constraints/validity.py
index 4a1f663..cf6e352 100644
--- a/gerrychain/constraints/validity.py
+++ b/gerrychain/constraints/validity.py
@@ -1,5 +1,6 @@
 from ..updaters import CountySplit
 from .bounds import Bounds
+import numpy
 
 
 class Validator:
@@ -31,6 +32,10 @@ class Validator:
         # check each constraint function and fail when a constraint test fails
         for constraint in self.constraints:
             is_valid = constraint(partition)
+            # Coerce NumPy booleans
+            if isinstance(is_valid, numpy.bool_):
+                is_valid = bool(is_valid)
+
             if is_valid is False:
                 return False
             elif is_valid is True:
diff --git a/gerrychain/partition/partition.py b/gerrychain/partition/partition.py
index 83c4b9a..b893aee 100644
--- a/gerrychain/partition/partition.py
+++ b/gerrychain/partition/partition.py
@@ -6,6 +6,7 @@ from ..graph import Graph
 from ..updaters import compute_edge_flows, flows_from_changes
 from .assignment import get_assignment
 from .subgraphs import SubgraphView
+from ..updaters import cut_edges
 
 
 class Partition:
@@ -15,7 +16,9 @@ class Partition:
     aggregations and calculations that we want to optimize.
     """
 
-    default_updaters = {}
+    default_updaters = {
+        "cut_edges": cut_edges
+    }
 
     def __init__(
         self, graph=None, assignment=None, updaters=None, parent=None, flips=None
@@ -46,6 +49,7 @@ class Partition:
 
         if updaters is None:
             updaters = dict()
+
         self.updaters = self.default_updaters.copy()
         self.updaters.update(updaters)
 
diff --git a/gerrychain/tree.py b/gerrychain/tree.py
index 60e0f2f..c56761a 100644
--- a/gerrychain/tree.py
+++ b/gerrychain/tree.py
@@ -202,20 +202,31 @@ def recursive_tree_part(
     """
     flips = {}
     remaining_nodes = set(graph.nodes)
+    # We keep a running tally of deviation from ``epsilon`` at each partition
+    # and use it to tighten the population constraints on a per-partition
+    # basis such that every partition, including the last partition, has a
+    # population within +/-``epsilon`` of the target population.
+    # For instance, if district n's population exceeds the target by 2%
+    # with a +/-2% epsilon, then district n+1's population should be between
+    # 98% of the target population and the target population.
+    debt = 0
 
     for part in parts[:-1]:
+        min_pop = max(pop_target * (1 - epsilon), pop_target * (1 - epsilon) - debt)
+        max_pop = min(pop_target * (1 + epsilon), pop_target * (1 + epsilon) - debt)
         nodes = method(
             graph.subgraph(remaining_nodes),
             pop_col=pop_col,
-            pop_target=pop_target,
-            epsilon=epsilon,
+            pop_target=(min_pop + max_pop) / 2,
+            epsilon=(max_pop - min_pop) / (2 * pop_target),
             node_repeats=node_repeats,
         )
 
+        part_pop = 0
         for node in nodes:
             flips[node] = part
-        # update pop_target?
-
+            part_pop += graph.nodes[node][pop_col]
+        debt += part_pop - pop_target
         remaining_nodes -= nodes
 
     # All of the remaining nodes go in the last part


Test patch:
diff --git a/tests/constraints/test_validity.py b/tests/constraints/test_validity.py
index d4061a0..006b563 100644
--- a/tests/constraints/test_validity.py
+++ b/tests/constraints/test_validity.py
@@ -1,6 +1,7 @@
 from unittest.mock import MagicMock
 
 import networkx as nx
+import numpy
 import pytest
 
 from gerrychain.constraints import (SelfConfiguringLowerBound, Validator,
@@ -140,6 +141,17 @@ def test_validator_raises_TypeError_if_constraint_returns_non_boolean():
         validator(mock_partition)
 
 
+def test_validator_accepts_numpy_booleans():
+    mock_partition = MagicMock()
+
+    mock_constraint = MagicMock()
+    mock_constraint.return_value = numpy.bool_(True)
+    mock_constraint.__name__ = "mock_constraint"
+
+    is_valid = Validator([mock_constraint])
+    assert is_valid(mock_partition)
+
+
 def test_no_vanishing_districts_works():
     parent = MagicMock()
     parent.assignment = get_assignment({1: 1, 2: 2}, MagicMock())
diff --git a/tests/partition/test_partition.py b/tests/partition/test_partition.py
index 6b44a41..56b0b16 100644
--- a/tests/partition/test_partition.py
+++ b/tests/partition/test_partition.py
@@ -137,6 +137,16 @@ def districtr_plan_file():
             json.dump(districtr_plan, f)
         yield filename
 
-
 def test_repr(example_partition):
     assert repr(example_partition) == "<Partition [2 parts]>"
+
+def test_partition_has_default_updaters(example_partition):
+    partition = example_partition
+    default_updaters = partition.default_updaters
+    should_have_updaters = {
+        "cut_edges": cut_edges
+    }
+
+    for updater in should_have_updaters:
+        assert default_updaters.get(updater, None) is not None
+        assert should_have_updaters[updater](partition) == partition[updater]
diff --git a/tests/test_tree.py b/tests/test_tree.py
index 1cab047..1acefaf 100644
--- a/tests/test_tree.py
+++ b/tests/test_tree.py
@@ -11,6 +11,7 @@ from gerrychain.tree import (
     bipartition_tree,
     random_spanning_tree,
     find_balanced_edge_cuts,
+    recursive_tree_part,
     PopulatedGraph,
 )
 from gerrychain.updaters import Tally, cut_edges
@@ -32,6 +33,16 @@ def partition_with_pop(graph_with_pop):
     )
 
 
+@pytest.fixture
+def twelve_by_twelve_with_pop():
+    xy_grid = networkx.grid_graph([12, 12])
+    nodes = {node: node[1] + 12 * node[0] for node in xy_grid}
+    grid = networkx.relabel_nodes(xy_grid, nodes)
+    for node in grid:
+        grid.nodes[node]["pop"] = 1
+    return grid
+
+
 def test_bipartition_tree_returns_a_subset_of_nodes(graph_with_pop):
     ideal_pop = sum(graph_with_pop.nodes[node]["pop"] for node in graph_with_pop) / 2
     result = bipartition_tree(graph_with_pop, "pop", ideal_pop, 0.25, 10)
@@ -48,6 +59,19 @@ def test_bipartition_tree_returns_within_epsilon_of_target_pop(graph_with_pop):
     assert abs(part_pop - ideal_pop) / ideal_pop < epsilon
 
 
+def test_recursive_tree_part_returns_within_epsilon_of_target_pop(twelve_by_twelve_with_pop):
+    n_districts = 7  # 144/7 ≈ 20.5 nodes/subgraph (1 person/node)
+    ideal_pop = (sum(twelve_by_twelve_with_pop.nodes[node]["pop"]
+                     for node in twelve_by_twelve_with_pop)) / n_districts
+    epsilon = 0.05
+    result = recursive_tree_part(twelve_by_twelve_with_pop, range(n_districts),
+                                 ideal_pop, "pop", epsilon)
+    partition = Partition(twelve_by_twelve_with_pop, result,
+                          updaters={"pop": Tally("pop")})
+    return all(abs(part_pop - ideal_pop) / ideal_pop < epsilon
+               for part_pop in partition['pop'].values())
+
+
 def test_random_spanning_tree_returns_tree_with_pop_attribute(graph_with_pop):
     tree = random_spanning_tree(graph_with_pop)
     assert networkx.is_tree(tree)