Update src/vis_utils.py

src/vis_utils.py

@@ -119,9 +119,13 @@ def plot_similarity_results(methods_selected, x_metric, y_metric, similarity_pat
     
     return filename
 
-def plot_function_results(file_path, aspect, metric, method_names):
+def plot_function_results(file_path, aspect, metric, method_names, 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)
+
     # Load data
-    df = pd.read_csv(file_path)
+    df = pd.read_csv(function_path)
     
     # Filter for selected methods
     df = df[df['Method'].isin(method_names)]
@@ -131,24 +135,35 @@ def plot_function_results(file_path, aspect, metric, method_names):
     df = df[['Method'] + columns_to_plot]
     df.set_index('Method', inplace=True)
     
+    # Fill missing values with 0
+    df = df.fillna(0)
+
+    # Generate colors for methods
+    row_color_dict = {method: get_method_color(method) for method in df.index}
+
+    long_form_mapping = {
+        "MF": "Molecular Function",
+        "BP": "Biological Process",
+        "CC": "Cellular Component"
+    }
+
     # Create clustermap
     g = sns.clustermap(df, annot=True, cmap="YlGnBu", row_cluster=False, col_cluster=False, figsize=(15, 15))
-    
+
+    title = f"{long_form_mapping[aspect.capitalize()]} Results for {metric.capitalize()}"
+    g.fig.suptitle(title, x=0.5, y=1.02, fontsize=16, ha='center')  # Center the title above the plot
+
     # Get heatmap axis and customize labels
     ax = g.ax_heatmap
     ax.set_xlabel("")
     ax.set_ylabel("")
-    
-    # Apply color and caret adjustments to x-axis labels
-    set_colors_and_marks_for_representation_groups(ax)
 
     # Save the plot as an image
-    save_path = "./plot_images"  # Ensure this folder exists or adjust the path
-    os.makedirs(save_path, exist_ok=True)  # Create directory if it doesn't exist
+    save_path = "/tmp"
     filename = os.path.join(save_path, f"{aspect}_{metric}_heatmap.png")
     plt.savefig(filename, dpi=400, bbox_inches='tight')
     plt.close()  # Close the plot to free memory
-    
+
     return filename
 
 def plot_family_results(file_path, method_names, metric, save_path="./plot_images"):