Update src/vis_utils.py

src/vis_utils.py

@@ -17,9 +17,67 @@ from about import *
 
 global data_component, filter_component
 
+def get_baseline_df(selected_methods, selected_metrics):
+    df = pd.read_csv(CSV_RESULT_PATH)
+    present_columns = ["method_name"] + selected_metrics
+    df = df[df['method_name'].isin(selected_methods)][present_columns]
+    return df
+
 def get_method_color(method):
     return color_dict.get(method, 'black')  # If method is not in color_dict, use black
 
+def set_colors_and_marks_for_representation_groups(ax):
+    for label in ax.get_xticklabels():
+        text = label.get_text()
+        color = group_color_dict.get(text, 'black')  # Default to black if label not in dict
+        label.set_color(color)
+        label.set_fontweight('bold')
+        
+        # Add a caret symbol to specific labels
+        if text in {'MUT2VEC', 'PFAM', 'GENE2VEC', 'BERT-PFAM'}:
+            label.set_text(f"^ {text}")
+
+def benchmark_plot(benchmark_type, methods_selected, x_metric, y_metric):
+    if benchmark_type == 'flexible':
+        # Use general visualizer logic
+        return general_visualizer_plot(methods_selected, x_metric=x_metric, y_metric=y_metric)
+    elif benchmark_type == 'similarity':
+        title = f"{x_metric} vs {y_metric}"
+        return draw_scatter_plot_similarity(methods_selected, x_metric, y_metric, title)
+    elif benchmark_type == 'Benchmark 3':
+        return benchmark_3_plot(x_metric, y_metric)
+    elif benchmark_type == 'Benchmark 4':
+        return benchmark_4_plot(x_metric, y_metric)
+    else:
+        return "Invalid benchmark type selected."
+
+def general_visualizer(methods_selected, x_metric, y_metric):
+    df = pd.read_csv(CSV_RESULT_PATH)
+    filtered_df = df[df['method_name'].isin(methods_selected)]
+
+    # Create a Seaborn lineplot with method as hue
+    plt.figure(figsize=(10, 8))  # Increase figure size
+    sns.lineplot(
+        data=filtered_df, 
+        x=x_metric, 
+        y=y_metric, 
+        hue="method_name",  # Different colors for different methods
+        marker="o",  # Add markers to the line plot
+    )
+    
+    # Add labels and title
+    plt.xlabel(x_metric)
+    plt.ylabel(y_metric)
+    plt.title(f'{y_metric} vs {x_metric} for selected methods')
+    plt.grid(True)
+    
+    # Save the plot to display it in Gradio
+    plot_path = "plot.png"
+    plt.savefig(plot_path)
+    plt.close()
+    
+    return plot_path
+
 def draw_scatter_plot_similarity(methods_selected, x_metric, y_metric, title):
     df = pd.read_csv(CSV_RESULT_PATH)
     # Filter the dataframe based on selected methods
@@ -64,50 +122,33 @@ def draw_scatter_plot_similarity(methods_selected, x_metric, y_metric, title):
     
     return filename
 
-def benchmark_plot(benchmark_type, methods_selected, x_metric, y_metric):
-    if benchmark_type == 'flexible':
-        # Use general visualizer logic
-        return general_visualizer_plot(methods_selected, x_metric=x_metric, y_metric=y_metric)
-    elif benchmark_type == 'similarity':
-        title = f"{x_metric} vs {y_metric}"
-        return draw_scatter_plot_similarity(methods_selected, x_metric, y_metric, title)
-    elif benchmark_type == 'Benchmark 3':
-        return benchmark_3_plot(x_metric, y_metric)
-    elif benchmark_type == 'Benchmark 4':
-        return benchmark_4_plot(x_metric, y_metric)
-    else:
-        return "Invalid benchmark type selected."
-
-
-def get_baseline_df(selected_methods, selected_metrics):
-    df = pd.read_csv(CSV_RESULT_PATH)
-    present_columns = ["method_name"] + selected_metrics
-    df = df[df['method_name'].isin(selected_methods)][present_columns]
-    return df
-
-def general_visualizer(methods_selected, x_metric, y_metric):
-    df = pd.read_csv(CSV_RESULT_PATH)
-    filtered_df = df[df['method_name'].isin(methods_selected)]
-
-    # Create a Seaborn lineplot with method as hue
-    plt.figure(figsize=(10, 8))  # Increase figure size
-    sns.lineplot(
-        data=filtered_df, 
-        x=x_metric, 
-        y=y_metric, 
-        hue="method_name",  # Different colors for different methods
-        marker="o",  # Add markers to the line plot
-    )
+def visualize_aspect_metric_clustermap(file_path, aspect, metric, method_names):
+    # Load data
+    df = pd.read_csv(file_path)
     
-    # Add labels and title
-    plt.xlabel(x_metric)
-    plt.ylabel(y_metric)
-    plt.title(f'{y_metric} vs {x_metric} for selected methods')
-    plt.grid(True)
+    # Filter for selected methods
+    df = df[df['Method'].isin(method_names)]
     
-    # Save the plot to display it in Gradio
-    plot_path = "plot.png"
-    plt.savefig(plot_path)
-    plt.close()
+    # Filter columns for specified aspect and metric
+    columns_to_plot = [col for col in df.columns if col.startswith(f"{aspect}_") and col.endswith(f"_{metric}")]
+    df = df[['Method'] + columns_to_plot]
+    df.set_index('Method', inplace=True)
+    
+    # Create clustermap
+    g = sns.clustermap(df, annot=True, cmap="YlGnBu", row_cluster=False, col_cluster=False, figsize=(15, 15))
+    
+    # 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
+    os.makedirs(save_path, exist_ok=True)  # Create directory if it doesn't exist
+    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 plot_path
+    return filename