Merge branch 'main' of https://github.com/gcrsef-gang/the-rebalancing…

…-act into main

.gitignore

@@ -11,3 +11,5 @@ out.mp4
 
 district_assignment.csv
 maup_concated.csv
+
+examples/.ipynb_checkpoints

examples/vra_nh.json

@@ -0,0 +1,9 @@
+{
+    "black": 0,
+    "hispanic": 0,
+    "asian": 0,
+    "native": 0,
+    "islander": 0,
+    "combined": 0,
+    "opportunity_threshold": 0.51
+}

rba/__main__.py

@@ -7,7 +7,10 @@
     Generates `num_thresholds` community maps based on the precinct graph `graph`, and writes to a
     file storing a list of individual communities, containing data on constituent precincts, birth 
     and death times.
- - districtgen
+ - districtgen [--graph_file] [--edge_lifetimes_file] [--vra_config] [--output_dir]
+    Runs simulated annealing algorithm, saves the ten best maps, as well as a dataframe keeping
+    track of various statistics for each state of the chain. `vra_config` is a JSON file
+    containing information about minority-opportunity district constraints.
  - ensemblegen
  - quantify
  - draw
@@ -47,5 +50,11 @@
     draw_parser.add_argument("--num_frames", type=int, default=50)
     draw_parser.set_defaults(func=rba.visualization.visualize)
 
+    optimize_parser = subparsers.add_parser("optimize")
+    optimize_parser.add_argument("--graph_file", type=str, default=os.path.join(package_dir, "data/2010/new_hampshire_geodata_merged.json"))
+    optimize_parser.add_argument("--edge_lifetime_file", type=str)
+    optimize_parser.add_argument("--vra_config", type=str)
+    optimize_parser.add_argument("--output_dir", type=str)
+
     args = parser.parse_args()
     args.func(**{key: val for key, val in vars(args).items() if key != "func"})

rba/constants.py

@@ -0,0 +1,11 @@
+"""Arbitrarily decided parameters for all algorithms.
+"""
+
+
+# Edges between nodes in the same county should be weighted less than those that cross, because
+# the maximum spanning tree should be made more likely to choose an edge crossing county lines.
+SAME_COUNTY_PENALTY = 0.5
+
+POP_EQUALITY_THRESHOLD = 0.005
+
+MINORITY_NAMES = ["black", "hispanic", "asian", "native", "islander"]

rba/district_optimization.py

@@ -0,0 +1,64 @@
+"""
+Given supercommunity edge lifetimes, uses simulated annealing to generate a map that minimizes
+the average border edge lifetime while conforming to redistricting requirements.
+"""
+
+from functools import partial
+import random
+
+from gerrychain import (GeographicPartition, Partition, Graph, MarkovChain,
+                        proposals, updaters, constraints, accept, Election)
+from gerrychain.proposals import recom
+from gerrychain.tree import bipartition_tree
+from networkx import tree
+
+from .util import get_num_vra_districts, get_county_weighted_random_spanning_tree
+
+
+class SimulatedAnnealingChain(MarkovChain):
+    """Augments gerrychain.MarkovChain to take both the current state and proposal in the `accept`
+    function.
+    """
+    def __next__(self):
+        if self.counter == 0:
+            self.counter += 1
+            return self.state
+
+        while self.counter < self.total_steps:
+            proposed_next_state = self.proposal(self.state)
+            # Erase the parent of the parent, to avoid memory leak
+            if self.state is not None:
+                self.state.parent = None
+
+            if self.is_valid(proposed_next_state):
+                if self.accept(self.state, proposed_next_state):
+                    self.state = proposed_next_state
+                self.counter += 1
+                return self.state
+        raise StopIteration
+
+
+def accept_proposal(temperature, current_energy, proposed_energy):
+    """Simple simulated-annealing acceptance function.
+    """
+    if current_energy > proposed_energy or random.random() < temperature:
+        return True
+    return False
+
+
+def generate_districts_simulated_annealing(graph, edge_lifetimes, num_vra_districts, vra_threshold,
+                                           pop_equality_threshold):
+    """Returns the 10 best maps and a dataframe of statistics for the entire chain.
+    """
+
+    weighted_recom_proposal = partial(
+        recom,
+        method=partial(
+            bipartition_tree,
+            spanning_tree_fn=get_county_weighted_random_spanning_tree)
+    )
+
+
+def optimize():
+    """
+    """

rba/util.py

@@ -1,9 +1,12 @@
-from scipy.special import rel_entr
-import networkx as nx
+"""Miscellaneous utilities.
+"""
+
+import random
 
+import networkx as nx
 
-def jenson_shannon_divergence(distribution1, distribution2):
-    average = [(distribution1[i] + distribution2[i])/2 for i in range(distribution1)]
+from . import constants
+from .district_quantification import quantify_gerrymandering
 
 
 def copy_adjacency(graph):
@@ -14,4 +17,59 @@ def copy_adjacency(graph):
         copy_graph.add_node(node)
     for u, v in graph.edges:
         copy_graph.add_edge(u, v)
-    return copy_graph
+    return copy_graph
+
+
+def get_num_vra_districts(partition, label, threshold):
+    """Returns the number of minority-opportunity distrcts for a given minority and threshold.
+
+    Parameters
+    ----------
+    partition : gerrychain.Parition
+        Proposed district plan.
+    label : str
+        Node data key that returns the population of that minority.
+    threshold : float
+        Value between 0 and 1 indicating the percent population required for a district to be
+        considered minority opportunity.
+    """
+    num_vra_districts = 0
+    for part in partition.parts:
+        total_pop = 0
+        minority_pop = 0
+        for node in partition.parts[part]:
+            total_pop += partition.graph.nodes[node]["total_pop"]
+            if label == "total_combined":
+                for minority in constants.MINORITY_NAMES:
+                    minority_pop += partition.graph.nodes[node][f"total_{minority}"]
+            else:
+                minority_pop += partition.graph.nodes[node][label]
+        if minority_pop / total_pop >= threshold:
+            num_vra_districts += 1
+    return num_vra_districts
+
+
+def get_gerrymandering_score(partition, edge_lifetimes):
+    """Returns the gerrymandering score of a partition.
+    """
+    return quantify_gerrymandering(partition.graph, partition.subgraphs, edge_lifetimes)[1]
+
+
+def get_district_gerrymandering_scores(partition, edge_lifetimes):
+    """Returns the gerrymandering scores of the districts in a partition"""
+    return quantify_gerrymandering(partition.graph, partition.subgraphs, edge_lifetimes)[0]
+
+
+def get_county_weighted_random_spanning_tree(graph):
+    """Applies random edge weights to a graph, then multiplies those weights depending on whether or
+    not the edge crosses a county border. Then returns the maximum spanning tree for the graph."""
+    for u, v in graph.edges:
+        weight = random.random()
+        if graph[u]["COUNTYFP10"] == graph[v]["COUNTYFP10"]:
+            weight *= constants.SAME_COUNTY_PENALTY
+        graph[u][v]["random_weight"] = weight
+
+    spanning_tree = nx.tree.maximum_spanning_tree(
+        graph, algorithm="kruskal", weight="random_weight"
+    )
+    return spanning_tree

rba/visualization.py

@@ -10,6 +10,7 @@
 import math
 
 from PIL import Image, ImageDraw, ImageFont
+import geopandas
 import networkx as nx
 import shapely.geometry
 import shapely.ops
@@ -104,6 +105,19 @@ def modify_coords(coords, bounds):
     return new_coords
 
 
+def visualize_partition_geopandas(partition):
+    """Visualizes a gerrychain.Partition object using geopandas.
+    """
+    data = {"assignment": [], "geometry": []}
+    for node in partition.graph:
+        data["assignment"].append(partition.assignment[node])
+        data["geometry"].append(shapely.geometry.shape(partition.graph.nodes[node]['geometry']))
+
+    gdf = geopandas.GeoDataFrame(data)
+    del data
+    gdf.plot(column="assignment")
+
+
 def visualize_map(graph, output_fpath, node_coords, edge_coords, node_colors=None, edge_colors=None,
         edge_widths=None, node_list=None, additional_polygons=None, text=None, show=False):
     """Creates an image of a map and saves it to a file.