fixed script

Author: jermp

script/plot_space.py

@@ -11,7 +11,7 @@
 def format_ticks(x, pos):
     return f'{x:.2f}'
 
-def main(json_file, pdf_filename, alpha):
+def main(json_file, pdf_filename, alpha, bucketer):
     # Load the data from the provided JSON file
     with open(json_file, 'r') as f:
         records = [json.loads(line) for line in f]
@@ -27,13 +27,13 @@ def main(json_file, pdf_filename, alpha):
     # Define configurations for filtering
     configurations = [
 
-        ((df['avg_partition_size'] == "0") & (df['num_partitions'] == "0") & (df['dense_partitioning'] == "false")
+        ((df['avg_partition_size'] == "0") & (df['num_partitions'] == "0") & (df['dense_partitioning'] == "false") & (df['bucketer_type'] == bucketer)
                 , "SINGLE"),
 
-        ((df['avg_partition_size'] != "0") & (df['num_partitions'] != "0") & (df['dense_partitioning'] == "false")
+        ((df['avg_partition_size'] != "0") & (df['num_partitions'] != "0") & (df['dense_partitioning'] == "false")  & (df['bucketer_type'] == bucketer)
                 , "PARTITIONED"),
 
-        ((df['avg_partition_size'] != "0") & (df['num_partitions'] != "0") & (df['dense_partitioning'] == "true")
+        ((df['avg_partition_size'] != "0") & (df['num_partitions'] != "0") & (df['dense_partitioning'] == "true")  & (df['bucketer_type'] == bucketer)
                 , "DENSE-PARTITIONED")
     ]
 
@@ -50,7 +50,8 @@ def main(json_file, pdf_filename, alpha):
         # Group by the specified fields
         grouped_avg = filtered_df.groupby([
             'n', 'lambda', 'alpha', 'minimal',
-            'bucketer_type', 'avg_partition_size',
+            'bucketer_type',
+            'avg_partition_size',
             'num_partitions', 'dense_partitioning', 'seed', 'num_threads',
             'external_memory', 'encoder_type'
         ])['bits_per_key'].mean().reset_index()
@@ -68,23 +69,28 @@ def main(json_file, pdf_filename, alpha):
     # Define different marker symbols for each encoder_type
     # marker_symbols = ['o', 'v', '^', '<', '>', '8', 's', 'p', '*', 'h', 'H', 'D', 'd', 'P', 'X']
     colors = plt.get_cmap('tab20', 12)  # Use 'tab20'
-
     # Create a new PDF file to save plots
     with PdfPages(pdf_filename) as pdf:
         fig = plt.figure(figsize=(20, 8))
         gs = gridspec.GridSpec(1, 4, width_ratios=[3, 3, 3, 1])
         axs = [fig.add_subplot(gs[i]) for i in range(3)]
 
-        inter_mono_handles = []
-        other_encoder_handles = []
+        encoder_handles = []
+        encoder_color_map = {}
+        i = 0
+        for (grouped_avg, _) in grouped_data:
+            for e in sorted(grouped_avg['encoder_type'].unique()):
+                if e not in encoder_color_map.keys():
+                    encoder_color_map[e] = i
+                    i += 1
 
         for ax, (grouped_avg, title) in zip(axs, grouped_data):
 
             encoder_types = sorted(grouped_avg['encoder_type'].unique())
 
-            for i, encoder_type in enumerate(encoder_types):
+            for encoder_type in encoder_types:
 
-                encoder_color = colors(i)
+                encoder_color = colors(encoder_color_map[encoder_type])
                 for alpha_value in sorted(grouped_avg['alpha'].unique(), reverse=True):
 
                     if alpha_value == alpha: # filter on specific alpha
@@ -100,11 +106,8 @@ def main(json_file, pdf_filename, alpha):
                                  color=encoder_color,
                                  linestyle='none')[0]
 
-                if 'inter' in encoder_type or 'mono' in encoder_type:
-                    inter_mono_handles.append(handle)
-                else:
-                    if not any(encoder_type == l.get_label() for l in other_encoder_handles):
-                        other_encoder_handles.append(handle)
+                if not encoder_type in [l.get_label() for l in encoder_handles]:
+                    encoder_handles.append(handle)
 
             # Set plot labels and title with LaTeX formatting
             ax.set_xlabel(r'$\lambda$', fontsize=14)
@@ -119,23 +122,19 @@ def main(json_file, pdf_filename, alpha):
             ax.tick_params(axis='both', which='major', labelsize=12)
 
         # Create three columns in the legend
-        other_encoder_labels = [h.get_label() for h in other_encoder_handles]
-        inter_mono_labels = [h.get_label() for h in inter_mono_handles]
+        encoder_labels = [h.get_label() for h in encoder_handles]
 
         # Create a new axis for the legend at the bottom of the main figure
         legend_ax = fig.add_subplot(gs[-1, :])  # Use GridSpec to create the legend axis
         legend_ax.axis('off')
 
         # Adjusting the `bbox_to_anchor` to move legends further to the right
-        other_legend = legend_ax.legend(other_encoder_handles, other_encoder_labels, loc='upper right',
-            title='SINGLE and PARTITIONED')  # Move further right
-        inter_mono_legend = legend_ax.legend(inter_mono_handles, inter_mono_labels, loc='center right',
-            bbox_to_anchor=(1, 0.45),
-            title='DENSE-PARTITIONED')  # Adjusted position
+        other_legend = legend_ax.legend(encoder_handles, encoder_labels, loc='upper right',
+            title='Encoders')  # Move further right
 
         # Add the legends to the axis
         legend_ax.add_artist(other_legend)
-        legend_ax.add_artist(inter_mono_legend)
+        # legend_ax.add_artist(inter_mono_legend)
 
         # Adjust the layout of the main figure
         plt.tight_layout()
@@ -154,9 +153,10 @@ def main(json_file, pdf_filename, alpha):
     parser.add_argument('-i', '--input_json_filename', required=True, type=str, help='Path to the input JSON file.')
     parser.add_argument('-o', '--output_pdf_filename', required=True, type=str, help='Path for the output PDF file.')
     parser.add_argument('-a', '--alpha', required=True, type=float, help='Value of alpha (a float).')
+    parser.add_argument('-b', '--bucketer', required=True, type=str, help='Bucketer type: values are "skew" or "opt".')
 
     # Parse the arguments
     args = parser.parse_args()
 
     # Call the main function with parsed arguments
-    main(args.input_json_filename, args.output_pdf_filename, args.alpha)
+    main(args.input_json_filename, args.output_pdf_filename, args.alpha, args.bucketer)