Update src/vis_utils.py

src/vis_utils.py

@@ -26,13 +26,13 @@ def benchmark_plot(benchmark_type, methods_selected, x_metric, y_metric, aspect,
     if benchmark_type == 'similarity':
         return plot_similarity_results(methods_selected, x_metric, y_metric)
     elif benchmark_type == 'function':
-        return plot_function_results(aspect, single_metric, methods_selected)
+        return plot_function_results(methods_selected, aspect, single_metric)
     elif benchmark_type == 'family':
-        return plot_family_results("./data/family_results.csv", methods_selected, x_metric, save_path="./plot_images")
+        return plot_family_results(methods_selected, dataset, single_metric)
     elif benchmark_type == "affinity":
-        return plot_affinity_results("./data/affinity_results.csv", methods_selected, x_metric, save_path="./plot_images")
-
-    return 0
+        return plot_affinity_results(methods_selected, single_metric)
+    else:
+        return -1 
 
 def get_method_color(method):
     return color_dict.get(method, 'black')  # If method is not in color_dict, use black
@@ -119,7 +119,7 @@ def plot_similarity_results(methods_selected, x_metric, y_metric, similarity_pat
     
     return filename
 
-def plot_function_results(aspect, metric, method_names, function_path="/tmp/function_results.csv"):
+def plot_function_results(method_names, aspect, metric, function_path="/tmp/function_results.csv"):
     if not os.path.exists(function_path):
         benchmark_types = ["similarity", "function", "family", "affinity"] #download all files for faster results later
         download_from_hub(benchmark_types)
@@ -168,18 +168,16 @@ def plot_function_results(aspect, metric, method_names, function_path="/tmp/func
 
     return filename
 
-def plot_family_results(file_path, method_names, metric, save_path="./plot_images"):
-    # Load data
-    df = pd.read_csv(file_path)
-    
-    # Filter by method names and selected metric columns
+def plot_family_results(methods_selected, dataset, metric, family_path="/tmp/family_results.csv"):
+    if not os.path.exists(function_path):
+        benchmark_types = ["similarity", "function", "family", "affinity"] #download all files for faster results later
+        download_from_hub(benchmark_types)
+
+    df = pd.read_csv(family_path)
+
+    # Filter by method names and selected dataset columns
     df = df[df['Method'].isin(method_names)]
-    metric_columns = [col for col in df.columns if col.startswith(f"{metric}_")]
-    
-    # Check if there are columns matching the selected metric
-    if not metric_columns:
-        print(f"No columns found for metric '{metric}'.")
-        return None
+    metric_columns = [col for col in df.columns if col.startswith(f"{dataset}_{metric}_")]
     
     # Reshape data for plotting
     df_long = pd.melt(df[['Method'] + metric_columns], id_vars=['Method'], var_name='Fold', value_name='Value')
@@ -204,12 +202,12 @@ def plot_family_results(file_path, method_names, metric, save_path="./plot_image
         ax.hlines(ytick + 0.5, -0.1, 1, linestyles='dashed')
 
     # Apply color settings to y-axis labels
-    set_colors_and_marks_for_representation_groups(ax)
-    
-    # Ensure save directory exists
-    os.makedirs(save_path, exist_ok=True)
-    
+    for label in ax.get_yticklabels():
+        method = label.get_text()
+        label.set_color(get_method_color(method))
+
     # Save the plot
+    save_path = "/tmp"
     filename = os.path.join(save_path, f"{metric}_family_results.png")
     ax.get_figure().savefig(filename, dpi=400, bbox_inches='tight')
     plt.close()  # Close the plot to free memory