Update src/vis_utils.py

src/vis_utils.py

@@ -183,6 +183,8 @@ def plot_family_results(methods_selected, dataset, metric, family_path="/tmp/fam
     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
 
+    df = df.fillna(0)
+
     # Set up the plot
     sns.set(rc={'figure.figsize': (13.7, 18.27)})
     sns.set_theme(style="whitegrid", color_codes=True)
@@ -214,23 +216,25 @@ def plot_family_results(methods_selected, dataset, metric, family_path="/tmp/fam
 
     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)
+def plot_affinity_results(method_names, metric, affinity_path="/tmp/affinity_results.csv"):
+    if not os.path.exists(affinity_path):
+        benchmark_types = ["similarity", "function", "family", "affinity"] #download all files for faster results later
+        download_from_hub(benchmark_types)
+
+    df = pd.read_csv(affinity_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
 
+    df = df.fillna(0)
+
     # Set up the plot
     sns.set(rc={'figure.figsize': (13.7, 8.27)})
     sns.set_theme(style="whitegrid", color_codes=True)
@@ -246,12 +250,15 @@ def plot_affinity_results(file_path, method_names, metric, save_path="./plot_ima
     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)
+    for label in ax.get_yticklabels():
+        method = label.get_text()
+        label.set_color(get_method_color(method))
 
     # Ensure save path exists
     os.makedirs(save_path, exist_ok=True)
 
     # Save the plot
+    save_path = "/tmp"
     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