test speedup ideas for simulated annealing

gcrsef-gang · Feb 24, 2023 · f57faec · f57faec
1 parent bb3bcf1
commit f57faec
Show file tree

Hide file tree

Showing 4 changed files with 117 additions and 188 deletions.
diff --git a/examples/ensemble_analysis.ipynb b/examples/ensemble_analysis.ipynb
diff --git a/examples/optimization.ipynb b/examples/optimization.ipynb
diff --git a/rba/data/2010/vra_md.json b/rba/data/2010/vra_md.json
@@ -0,0 +1,9 @@
+{
+    "black": 0,
+    "hispanic": 0,
+    "asian": 0,
+    "native": 0,
+    "islander": 0,
+    "combined": 4,
+    "opportunity_threshold": 0.51
+}
diff --git a/rba/optimization.py b/rba/optimization.py
@@ -15,7 +15,7 @@
 from gerrychain import Partition, Graph, MarkovChain, updaters, constraints, accept
 from gerrychain.constraints import Validator
 from gerrychain.proposals import recom
-from gerrychain.tree import recursive_tree_part, bipartition_tree
+from gerrychain.tree import recursive_tree_part, bipartition_tree, uniform_spanning_tree
 import gerrychain.random
 import networkx as nx
 import pandas as pd
@@ -24,6 +24,7 @@
 from .district_quantification import quantify_gerrymandering
 from .util import (get_num_vra_districts, load_districts, get_county_spanning_forest,
                    save_assignment, choose_cut)
+from .visualization import visualize_partition_geopandas
 
 
 class SimulatedAnnealingChain(MarkovChain):
@@ -179,25 +180,31 @@ def generate_districts_simulated_annealing(graph, edge_lifetimes, num_vra_distri
 
     initial_partition = Partition(graph, initial_assignment, sa_updaters)
 
-    county_recom_proposal = partial(
-        recom,
+    visualize_partition_geopandas(initial_partition)
+
+    # county_recom_proposal = partial(
+    #     recom,
+    #     pop_col="total_pop",
+    #     pop_target=ideal_population,
+    #     epsilon=constants.POP_EQUALITY_THRESHOLD,
+    #     node_repeats=2,
+    #     method=partial(
+    #         bipartition_tree,
+    #         spanning_tree_fn=get_county_spanning_forest,
+    #         choice=partial(choose_cut, graph=graph))
+    # )
+
+    recom_proposal = partial(recom,
         pop_col="total_pop",
         pop_target=ideal_population,
         epsilon=constants.POP_EQUALITY_THRESHOLD,
         node_repeats=2,
         method=partial(
             bipartition_tree,
-            spanning_tree_fn=get_county_spanning_forest,
-            choice=partial(choose_cut, graph=graph))
+            spanning_tree_fn=lambda G: Graph(uniform_spanning_tree(G, lambda x: random.choice(tuple(x))))
+        )
     )
 
-    # recom_proposal = partial(recom,
-    #     pop_col="total_pop",
-    #     pop_target=ideal_population,
-    #     epsilon=constants.POP_EQUALITY_THRESHOLD,
-    #     node_repeats=2
-    # )
-
     pop_constraint = constraints.within_percent_of_ideal_population(initial_partition,
                                                                     pop_equality_threshold)
 
@@ -221,8 +228,8 @@ def generate_districts_simulated_annealing(graph, edge_lifetimes, num_vra_distri
         get_temperature=partial(
             SimulatedAnnealingChain.COOLING_SCHEDULES[cooling_schedule],
             num_steps=num_steps),
-        proposal=county_recom_proposal,
-        # proposal=recom_proposal,
+        # proposal=county_recom_proposal,
+        proposal=recom_proposal,
         constraints=[
             pop_constraint,
             # compactness_bound