Add / fix tests for new features

gerrychain/constraints/contiguity.py

@@ -155,6 +155,9 @@ def affected_parts(partition: Partition) -> Set[int]:
     if flips is None:
         return partition.parts
 
+    if parent is None:
+        return set(flips.values())
+
     affected = set()
     for node, part in flips.items():
         affected.add(part)

tests/conftest.py

@@ -3,7 +3,8 @@
 from gerrychain import Graph, Partition
 import random
 from gerrychain.updaters import cut_edges
-import networkx
+import networkx as nx
+
 random.seed(2018)
 
 
@@ -34,6 +35,83 @@ def three_by_three_grid():
     return graph
 
 
+@pytest.fixture
+def four_by_five_grid_for_opt():
+    #  1  2  2  2  2
+    #  1  2  1  1  2
+    #  1  2  2  1  2
+    #  1  1  1  1  2
+
+    # 15 16 17 18 19
+    # 10 11 12 13 14
+    #  5  6  7  8  9
+    #  0  1  2  3  4
+
+    graph = Graph()
+    graph.add_nodes_from(
+        [
+            (0, {"population": 10, "opt_value": 1, "MVAP": 2}),
+            (1, {"population": 10, "opt_value": 1, "MVAP": 2}),
+            (2, {"population": 10, "opt_value": 1, "MVAP": 2}),
+            (3, {"population": 10, "opt_value": 1, "MVAP": 2}),
+            (4, {"population": 10, "opt_value": 2, "MVAP": 2}),
+            (5, {"population": 10, "opt_value": 1, "MVAP": 2}),
+            (6, {"population": 10, "opt_value": 2, "MVAP": 2}),
+            (7, {"population": 10, "opt_value": 2, "MVAP": 2}),
+            (8, {"population": 10, "opt_value": 1, "MVAP": 2}),
+            (9, {"population": 10, "opt_value": 2, "MVAP": 2}),
+            (10, {"population": 10, "opt_value": 1, "MVAP": 6}),
+            (11, {"population": 10, "opt_value": 2, "MVAP": 6}),
+            (12, {"population": 10, "opt_value": 1, "MVAP": 6}),
+            (13, {"population": 10, "opt_value": 1, "MVAP": 4}),
+            (14, {"population": 10, "opt_value": 2, "MVAP": 4}),
+            (15, {"population": 10, "opt_value": 1, "MVAP": 6}),
+            (16, {"population": 10, "opt_value": 2, "MVAP": 6}),
+            (17, {"population": 10, "opt_value": 2, "MVAP": 6}),
+            (18, {"population": 10, "opt_value": 2, "MVAP": 4}),
+            (19, {"population": 10, "opt_value": 2, "MVAP": 4}),
+        ]
+    )
+
+    graph.add_edges_from(
+        [
+            (0, 1),
+            (0, 5),
+            (1, 2),
+            (1, 6),
+            (2, 3),
+            (2, 7),
+            (3, 4),
+            (3, 8),
+            (4, 9),
+            (5, 6),
+            (5, 10),
+            (6, 7),
+            (6, 11),
+            (7, 8),
+            (7, 12),
+            (8, 9),
+            (8, 13),
+            (9, 14),
+            (10, 11),
+            (10, 15),
+            (11, 12),
+            (11, 16),
+            (12, 13),
+            (12, 17),
+            (13, 14),
+            (13, 18),
+            (14, 19),
+            (15, 16),
+            (16, 17),
+            (17, 18),
+            (18, 19),
+        ]
+    )
+
+    return graph
+
+
 @pytest.fixture
 def graph_with_random_data_factory(three_by_three_grid):
     def factory(columns):
@@ -57,11 +135,12 @@ def graph(three_by_three_grid):
 
 @pytest.fixture
 def example_partition():
-    graph = Graph.from_networkx(networkx.complete_graph(3))
+    graph = Graph.from_networkx(nx.complete_graph(3))
     assignment = {0: 1, 1: 1, 2: 2}
     partition = Partition(graph, assignment, {"cut_edges": cut_edges})
     return partition
 
+
 # From the docs: https://docs.pytest.org/en/latest/example/simple.html#control-skipping-of-tests-according-to-command-line-option
 def pytest_addoption(parser):
     parser.addoption(

tests/optimization/test_gingleator.py

@@ -0,0 +1,229 @@
+from gerrychain import Partition
+from gerrychain.optimization import Gingleator
+from gerrychain.constraints import contiguous
+from gerrychain.proposals import recom
+from gerrychain.updaters import Tally
+from functools import partial
+import pytest
+import numpy as np
+import random
+
+random.seed(2024)
+
+
+def simple_cut_edge_count(partition):
+    return len(partition["cut_edges"])
+
+
+def gingleator_test_partition(four_by_five_grid_for_opt):
+    return Partition(
+        graph=four_by_five_grid_for_opt,
+        assignment={
+            0: 0,
+            1: 0,
+            2: 0,
+            3: 0,
+            4: 0,
+            5: 1,
+            6: 1,
+            7: 1,
+            8: 1,
+            9: 1,
+            10: 2,
+            11: 2,
+            12: 2,
+            13: 2,
+            14: 2,
+            15: 3,
+            16: 3,
+            17: 3,
+            18: 3,
+            19: 3,
+        },
+        updaters={
+            "population": Tally("population", alias="population"),
+            "MVAP": Tally("MVAP", alias="MVAP"),
+            "m_perc": lambda p_dict: {
+                key: p_dict["MVAP"][key] / p_dict["population"][key]
+                for key in p_dict["MVAP"]
+            },
+            "my_cut_edges": simple_cut_edge_count,
+        },
+    )
+
+
+def test_ginglator_needs_min_perc_or_min_pop_col(four_by_five_grid_for_opt):
+    initial_partition = Partition.from_random_assignment(
+        graph=four_by_five_grid_for_opt,
+        n_parts=4,
+        epsilon=0.0,
+        pop_col="population",
+        updaters={
+            "population": Tally("population", alias="population"),
+            "MVAP": Tally("MVAP", alias="MVAP"),
+            "my_cut_edges": simple_cut_edge_count,
+        },
+    )
+
+    ideal_pop = sum(initial_partition["population"].values()) / 4
+
+    proposal = partial(
+        recom,
+        pop_col="population",
+        pop_target=ideal_pop,
+        epsilon=0.0,
+        node_repeats=1,
+    )
+
+    with pytest.raises(ValueError) as gingle_err:
+        gingles = Gingleator(
+            proposal=proposal,
+            constraints=[contiguous],
+            initial_state=initial_partition,
+            total_pop_col="population",
+            score_function=Gingleator.num_opportunity_dists,
+        )
+
+    assert "`minority_perc_col` and `minority_pop_col` cannot both be `None`" in str(
+        gingle_err.value
+    )
+
+
+def test_ginglator_warns_if_min_perc_and_min_pop_col_set(four_by_five_grid_for_opt):
+    initial_partition = Partition.from_random_assignment(
+        graph=four_by_five_grid_for_opt,
+        n_parts=4,
+        epsilon=0.0,
+        pop_col="population",
+        updaters={
+            "population": Tally("population", alias="population"),
+            "MVAP": Tally("MVAP", alias="MVAP"),
+            "m_perc": lambda p_dict: {
+                key: p_dict["MVAP"][key] / p_dict["population"][key]
+                for key in p_dict["MVAP"]
+            },
+            "my_cut_edges": simple_cut_edge_count,
+        },
+    )
+
+    ideal_pop = sum(initial_partition["population"].values()) / 4
+
+    proposal = partial(
+        recom,
+        pop_col="population",
+        pop_target=ideal_pop,
+        epsilon=0.0,
+        node_repeats=1,
+    )
+
+    with pytest.warns() as record:
+        gingles = Gingleator(
+            proposal=proposal,
+            constraints=[contiguous],
+            initial_state=initial_partition,
+            total_pop_col="population",
+            minority_pop_col="MVAP",
+            minority_perc_col="m_perc",
+            score_function=Gingleator.num_opportunity_dists,
+        )
+
+    assert "`minority_perc_col` and `minority_pop_col` are both specified" in str(
+        record.list[0].message
+    )
+
+
+def test_ginglator_finds_best_partition(four_by_five_grid_for_opt):
+    initial_partition = Partition.from_random_assignment(
+        graph=four_by_five_grid_for_opt,
+        n_parts=4,
+        epsilon=0.0,
+        pop_col="population",
+        updaters={
+            "population": Tally("population", alias="population"),
+            "MVAP": Tally("MVAP", alias="MVAP"),
+            "my_cut_edges": simple_cut_edge_count,
+        },
+    )
+
+    ideal_pop = sum(initial_partition["population"].values()) / 4
+
+    proposal = partial(
+        recom,
+        pop_col="population",
+        pop_target=ideal_pop,
+        epsilon=0.0,
+        node_repeats=1,
+    )
+
+    gingles = Gingleator(
+        proposal=proposal,
+        constraints=[contiguous],
+        initial_state=initial_partition,
+        minority_pop_col="MVAP",
+        total_pop_col="population",
+        score_function=Gingleator.num_opportunity_dists,
+    )
+
+    total_steps = 1000
+    burst_length = 10
+
+    max_scores_sb = np.zeros(total_steps)
+    for i, part in enumerate(
+        gingles.short_bursts(
+            burst_length=burst_length,
+            num_bursts=total_steps // burst_length,
+            with_progress_bar=True,
+        )
+    ):
+        max_scores_sb[i] = gingles.best_score
+
+    assert max_scores_sb[-1] == 2
+
+
+def test_count_num_opportunity_dists(four_by_five_grid_for_opt):
+    initial_partition = gingleator_test_partition(four_by_five_grid_for_opt)
+
+    assert Gingleator.num_opportunity_dists(initial_partition, "m_perc", 0.5) == 2
+    assert Gingleator.num_opportunity_dists(initial_partition, "m_perc", 0.6) == 0
+
+
+def test_reward_partial_dist(four_by_five_grid_for_opt):
+    initial_partition = gingleator_test_partition(four_by_five_grid_for_opt)
+
+    assert Gingleator.reward_partial_dist(initial_partition, "m_perc", 0.5) == 2 + 0.2
+    assert Gingleator.reward_partial_dist(initial_partition, "m_perc", 0.6) == 0.52
+
+
+def test_reward_next_highest_close(four_by_five_grid_for_opt):
+    initial_partition = gingleator_test_partition(four_by_five_grid_for_opt)
+
+    assert Gingleator.reward_next_highest_close(initial_partition, "m_perc", 0.5) == 2
+    # Rounding needed here because of floating point arithmetic
+    assert (
+        round(
+            Gingleator.reward_next_highest_close(initial_partition, "m_perc", 0.29), 5
+        )
+        == 2 + 0.1
+    )
+
+
+def test_penalize_maximum_over(four_by_five_grid_for_opt):
+    initial_partition = gingleator_test_partition(four_by_five_grid_for_opt)
+
+    assert (
+        Gingleator.penalize_maximum_over(initial_partition, "m_perc", 0.5)
+        == 2.0 + 0.48 / 0.50
+    )
+
+    assert Gingleator.penalize_maximum_over(initial_partition, "m_perc", 0.6) == 0
+
+
+def test_penalize_avg_over(four_by_five_grid_for_opt):
+    initial_partition = gingleator_test_partition(four_by_five_grid_for_opt)
+
+    assert (
+        Gingleator.penalize_avg_over(initial_partition, "m_perc", 0.5)
+        == 2.0 + 0.48 / 0.50
+    )
+
+    assert Gingleator.penalize_avg_over(initial_partition, "m_perc", 0.6) == 0