Update src/vis_utils.py

src/vis_utils.py

@@ -45,11 +45,11 @@ def benchmark_plot(benchmark_type, methods_selected, x_metric, y_metric):
         title = f"{x_metric} vs {y_metric}"
         return draw_scatter_plot_similarity(methods_selected, x_metric, y_metric, title)
     elif benchmark_type == 'function':
-        plot_function_results("./data/function_results.csv", x_metric, y_metric, methods_selected)
+        return plot_function_results("./data/function_results.csv", x_metric, y_metric, methods_selected)
     elif benchmark_type == 'family':
-        plot_family_results("./data/family_results.csv", methods_selected, x_metric, save_path="./plot_images")
-    else:
-        return "Invalid benchmark type selected."
+        return plot_family_results("./data/family_results.csv", methods_selected, x_metric, save_path="./plot_images")
+    elif benchmark_type == "affinity":
+        return plot_affinity_results("./data/affinity_results.csv", methods_selected, x_metric, save_path="./plot_images")
 
 def general_visualizer(methods_selected, x_metric, y_metric):
     df = pd.read_csv(CSV_RESULT_PATH)
@@ -200,4 +200,48 @@ def plot_family_results(file_path, method_names, metric, save_path="./plot_image
     ax.get_figure().savefig(filename, dpi=400, bbox_inches='tight')
     plt.close()  # Close the plot to free memory
 
+    return filename
+
+def plot_affinity_results(file_path, method_names, metric, save_path="./plot_images"):
+    # Load the CSV data
+    df = pd.read_csv(file_path)
+    
+    # Filter for selected methods
+    df = df[df['Method'].isin(method_names)]
+    
+    # Gather columns related to the specified metric and validate
+    metric_columns = [col for col in df.columns if col.startswith(f"{metric}_")]
+    if not metric_columns:
+        print(f"No columns found for metric '{metric}'.")
+        return None
+    
+    # Reshape data for plotting
+    df_long = pd.melt(df[['Method'] + metric_columns], id_vars=['Method'], var_name='Fold', value_name='Value')
+    df_long['Fold'] = df_long['Fold'].apply(lambda x: int(x.split('_')[-1]))  # Extract fold index for sorting
+
+    # Set up the plot
+    sns.set(rc={'figure.figsize': (13.7, 8.27)})
+    sns.set_theme(style="whitegrid", color_codes=True)
+
+    # Create a boxplot for the metric
+    ax = sns.boxplot(data=df_long, x='Value', y='Method', hue='Fold', whis=np.inf, orient="h")
+    
+    # Customize x-axis and y-axis tickers and grid
+    ax.xaxis.set_major_locator(ticker.MultipleLocator(5))
+    ax.get_xaxis().set_minor_locator(mpl.ticker.AutoMinorLocator())
+    ax.get_yaxis().set_minor_locator(mpl.ticker.AutoMinorLocator())
+    ax.grid(b=True, which='major', color='gainsboro', linewidth=1.0)
+    ax.grid(b=True, which='minor', color='whitesmoke', linewidth=0.5)
+
+    # Apply custom color settings to y-axis labels
+    set_colors_and_marks_for_representation_groups(ax)
+
+    # Ensure save path exists
+    os.makedirs(save_path, exist_ok=True)
+
+    # Save the plot
+    filename = os.path.join(save_path, f"{metric}_affinity_results.png")
+    ax.get_figure().savefig(filename, dpi=400, bbox_inches='tight')
+    plt.close()  # Close the plot to free memory
+
     return filename