Final edits to figures and figure scripts

Author: peterrrock2

.gitignore

@@ -9,4 +9,5 @@ dev_files
 
 # These might be added later when the full repo is complete
 mile_marker.txt
-checker.ipynb
+checker.ipynb
+*.pkl

figure_and_table_generation/figure_scripts/box_share_VA.py

@@ -168,9 +168,9 @@
     ax.legend(
         handles=handles,
         labels=[
-            "RRC CD 12 (5B proposed)",
-            "RRC CD 16 (5B proposed)",
-            "RRC Rand Plan (5B proposed)",
+            "RevReCom Seed 1 (5B proposed)",
+            "RevReCom Seed 2 (5B proposed)",
+            "RevReCom Seed 3 (5B proposed)",
             "Forest (10M proposed)",
         ],
         loc="upper left",

figure_and_table_generation/figure_scripts/helper_files/wasserstein_trace_tally.py

@@ -169,3 +169,120 @@ def wasserstein_trace_v_full(
             xticks.append(step)
             trace.append(distance)
     return xticks, trace
+
+
+def wasserstein_trace_shares(shares1_df, shares2_df, weights1, weights2, resolution):
+    """
+    Computes the Wasserstein trace between a full ensemble and an ongoing ensemble.
+    That is, given a full dataframe of counts and weights which generates some distribution,
+    and some ongoing array of counts and weights, the Wasserstein trace contains the Wasserstein
+    between the ongoing distribution and the totality of the full distribution at each step
+    equal to the resolution.
+
+    Parameters
+    ----------
+    shares_df : pandas.DataFrame
+        The dataframe of shares for the ongoing ensemble.
+    full_df : pandas.DataFrame
+        The dataframe of shares for the full ensemble.
+    weights : pandas.Series
+        The weights for the ongoing ensemble.
+    weights_full : pandas.Series
+        The weights for the full ensemble.
+    resolution : int
+        The resolution of the trace.
+
+    Returns
+    -------
+    (array-like, array-like):
+        The xticks for use in plotting and the trace of the Wasserstein distances.
+    """
+    assert all(shares1_df.columns == shares2_df.columns)
+
+    shares1 = shares1_df.sort_index(axis=1).to_numpy()
+    shares2 = shares2_df.sort_index(axis=1).to_numpy()
+
+    n_districts = len(shares1[0])
+
+    assert shares1_df.shape == shares2_df.shape
+
+    state1 = np.zeros(n_districts)
+    state2 = np.zeros(n_districts)
+    xticks = []
+    trace = []
+    hist1 = [Counter() for _ in range(n_districts)]
+    hist2 = [Counter() for _ in range(n_districts)]
+
+    for step, (s1, w1, s2, w2) in enumerate(
+        tqdm(zip(shares1, weights1, shares2, weights2), total=shares1.shape[0])
+    ):
+        # We assume 1-indexed districts.
+        for dist, v in enumerate(s1):
+            state1[dist] = v
+        for k, v in enumerate(sorted(state1)):
+            hist1[k][v] += w1
+        for dist, v in enumerate(s2):
+            state2[dist] = v
+        for k, v in enumerate(sorted(state2)):
+            hist2[k][v] += w2
+        if step > 0 and step % resolution == 0:
+            distance = 0
+            for dist1, dist2 in zip(hist1, hist2):
+                distance += wasserstein_distance(
+                    list(dist1.keys()),
+                    list(dist2.keys()),
+                    list(dist1.values()),
+                    list(dist2.values()),
+                )
+            xticks.append(step)
+            trace.append(distance)
+    return xticks, trace
+
+
+def wasserstein_trace_shares(shares1_df, shares2_df, weights1, weights2, resolution):
+    """
+    Computes the Wasserstein trace between a full ensemble and an ongoing ensemble.
+    """
+    # Ensure that the dataframes have the same columns
+    assert all(shares1_df.columns == shares2_df.columns)
+
+    # Convert dataframes to numpy arrays (columns sorted)
+    shares1 = shares1_df.sort_index(axis=1).to_numpy()
+    shares2 = shares2_df.sort_index(axis=1).to_numpy()
+
+    n_districts = shares1.shape[1]
+    assert shares1_df.shape == shares2_df.shape
+
+    xticks = []
+    trace = []
+    # Initialize a counter per district for each ensemble
+    hist1 = [Counter() for _ in range(n_districts)]
+    hist2 = [Counter() for _ in range(n_districts)]
+
+    for step, (s1, w1, s2, w2) in enumerate(
+        tqdm(zip(shares1, weights1, shares2, weights2), total=shares1.shape[0])
+    ):
+        # Directly sort the current row using NumPy
+        sorted_s1 = np.sort(s1)
+        for k, v in enumerate(sorted_s1):
+            hist1[k][v] += w1
+
+        sorted_s2 = np.sort(s2)
+        for k, v in enumerate(sorted_s2):
+            hist2[k][v] += w2
+
+        # Compute the Wasserstein trace at the specified resolution
+        if step > 0 and step % resolution == 0:
+            distance = sum(
+                wasserstein_distance(
+                    list(dist1.keys()),
+                    list(dist2.keys()),
+                    list(dist1.values()),
+                    list(dist2.values()),
+                )
+                for dist1, dist2 in zip(hist1, hist2)
+            )
+            xticks.append(step)
+            trace.append(distance)
+
+    return xticks, trace

figure_and_table_generation/figure_scripts/histogram_comparison_50x50.py

@@ -25,6 +25,9 @@
     "#80b1d3",
 ]
 
+script_dir = Path(__file__).resolve().parent
+top_dir = script_dir.parents[1]
+
 
 def make_method_plot(rrc_forest, lower, upper, n_dists, methods="forest_rrc"):
     """
@@ -50,7 +53,7 @@ def make_method_plot(rrc_forest, lower, upper, n_dists, methods="forest_rrc"):
     """
 
     methods = methods.replace(" ", "_")
-    out_path = Path("../figures")
+    out_path = Path(f"{script_dir}/../figures")
 
     _, ax = plt.subplots(figsize=(25, 10), dpi=400)
 
@@ -73,7 +76,7 @@ def make_method_plot(rrc_forest, lower, upper, n_dists, methods="forest_rrc"):
     ax.set_xticks(list(range(lower, upper, 50)))
     ax.set_xticklabels([str(i) for i in range(lower, upper, 50)], fontsize=16)
     ax.set_yticks([])
-    ax.legend(loc="right", bbox_to_anchor=(1.18, 0.5), prop={"size": 16})
+    ax.legend(loc="right", bbox_to_anchor=(1.21, 0.5), prop={"size": 16})
     plt.savefig(
         out_path.joinpath(f"50x50_{n_dists}_dist_{methods}_comparison.png"),
         bbox_inches="tight",
@@ -100,7 +103,7 @@ def make_recom_plot(lower, upper, n_dists, glob_expr):
     -------
     None
     """
-    out_path = Path("../figures")
+    out_path = Path(f"{script_dir}/../figures")
     data_path = Path(f"{top_dir}/hpc_files/hpc_processed_data/50x50")
     _, ax = plt.subplots(figsize=(25, 10), dpi=500)
 
@@ -114,8 +117,6 @@ def make_recom_plot(lower, upper, n_dists, glob_expr):
         lst = file.name.split("_")
         all_recom_files[f"{lst[1]} (1B proposed)"] = file
 
-    all_recom_files
-
     for i, (n, f) in enumerate(all_recom_files.items()):
         df = pd.read_parquet(f)
         prob_df = df.groupby("cut_edges").sum().reset_index()
@@ -128,7 +129,7 @@ def make_recom_plot(lower, upper, n_dists, glob_expr):
             edgecolor=None,
             color=colors[i + 3],
             alpha=0.8,
-            label=n.replace("Recom", "ReCom "),
+            label=n.replace("ReCom", "ReCom-"),
         )
 
     ax.set_xlim(lower - 20, upper + 20)
@@ -147,43 +148,43 @@ def make_recom_plot(lower, upper, n_dists, glob_expr):
     top_dir = script_dir.parents[1]
 
     rrc_forest_10 = {
-        "RRC (10B proposed)": f"{top_dir}/hpc_files/hpc_processed_data/50x50/50x50_RevReCom_steps_10000000000_plan_50x5_strip_20240618_174413_cut_edges.parquet",
+        "RevReCom (10B proposed)": f"{top_dir}/hpc_files/hpc_processed_data/50x50/50x50_RevReCom_steps_10000000000_plan_50x5_strip_20240618_174413_cut_edges.parquet",
         "Forest (10M proposed)": f"{top_dir}/hpc_files/hpc_processed_data/50x50/50x50_Forest_steps_10000000_rng_seed_278986_gamma_0.0_alpha_1.0_ndists_10_20240830_142334_cut_edges.parquet",
     }
 
     rrc_forest_smc_10 = rrc_forest_10.copy()
-    rrc_forest_smc_10["SMC (100k Samples)"] = (
+    rrc_forest_smc_10["SMC (100K)"] = (
         f"{top_dir}/hpc_files/hpc_processed_data/50x50/50x50_SMC_batch_size_100000_rng_seed_278986_dists_10_20250129_150813_cut_edges.parquet"
     )
 
-    make_method_plot(rrc_forest_10, 350, 601, 10, "forest_rrc")
-    make_method_plot(rrc_forest_smc_10, 350, 601, 10, "forest_rrc_smc")
+    # make_method_plot(rrc_forest_10, 350, 601, 10, "forest_rrc")
+    # make_method_plot(rrc_forest_smc_10, 350, 601, 10, "forest_rrc_smc")
     make_recom_plot(350, 601, 10, "*_ReCom*50x5_*")
 
     rrc_forest_25 = {
-        "RRC (10B proposed)": f"{top_dir}/hpc_files/hpc_processed_data/50x50/50x50_RevReCom_steps_10000000000_plan_10x10_square_20240618_174413_cut_edges.parquet",
+        "RevReCom (10B proposed)": f"{top_dir}/hpc_files/hpc_processed_data/50x50/50x50_RevReCom_steps_10000000000_plan_10x10_square_20240618_174413_cut_edges.parquet",
         "Forest (10M proposed)": f"{top_dir}/hpc_files/hpc_processed_data/50x50/50x50_Forest_steps_10000000_rng_seed_278986_gamma_0.0_alpha_1.0_ndists_25_20240830_142334_cut_edges.parquet",
     }
 
     rrc_forest_smc_25 = rrc_forest_25.copy()
-    rrc_forest_smc_25["SMC (100k Samples)"] = (
+    rrc_forest_smc_25["SMC (100K)"] = (
         f"{top_dir}/hpc_files/hpc_processed_data/50x50/50x50_SMC_batch_size_100000_rng_seed_278986_dists_25_20250129_150813_cut_edges.parquet"
     )
 
-    make_method_plot(rrc_forest_25, 650, 880, 25, "forest_rrc")
-    make_method_plot(rrc_forest_smc_25, 650, 880, 25, "forest_rrc_smc")
-    make_recom_plot(650, 880, 25, "*_ReCom*10x10_*")
+    # make_method_plot(rrc_forest_25, 650, 851, 25, "forest_rrc")
+    # make_method_plot(rrc_forest_smc_25, 650, 851, 25, "forest_rrc_smc")
+    make_recom_plot(650, 851, 25, "*_ReCom*10x10_*")
 
     rrc_forest_50 = {
-        "RRC (10B proposed)": f"{top_dir}/hpc_files/hpc_processed_data/50x50/50x50_RevReCom_steps_10000000000_plan_50x1_strip_20240618_174413_cut_edges.parquet",
+        "RevReCom (10B proposed)": f"{top_dir}/hpc_files/hpc_processed_data/50x50/50x50_RevReCom_steps_10000000000_plan_50x1_strip_20240618_174413_cut_edges.parquet",
         "Forest (10M proposed)": f"{top_dir}/hpc_files/hpc_processed_data/50x50/50x50_Forest_steps_10000000_rng_seed_278986_gamma_0.0_alpha_1.0_ndists_50_20240830_142334_cut_edges.parquet",
     }
 
     rrc_forest_smc_50 = rrc_forest_50.copy()
-    rrc_forest_smc_50["SMC (100k Samples)"] = (
+    rrc_forest_smc_50["SMC (100K)"] = (
         f"{top_dir}/hpc_files/hpc_processed_data/50x50/50x50_SMC_batch_size_100000_rng_seed_278986_dists_50_20250129_150813_cut_edges.parquet"
     )
 
-    make_method_plot(rrc_forest_50, 900, 1180, 50, "forest_rrc")
-    make_method_plot(rrc_forest_smc_50, 900, 1180, 50, "forest_rrc_smc")
+    # make_method_plot(rrc_forest_50, 900, 1151, 50, "forest_rrc")
+    # make_method_plot(rrc_forest_smc_50, 900, 1151, 50, "forest_rrc_smc")
     make_recom_plot(900, 1180, 50, "*_ReCom*50x1_*")