Update src/vis_utils.py

src/vis_utils.py

@@ -152,8 +152,17 @@ def plot_function_results(method_names, aspect, metric, function_path="/tmp/func
     }
 
     # Create clustermap
-    g = sns.clustermap(df, annot=True, cmap="YlGnBu", row_cluster=False, col_cluster=False, figsize=(15, 15))
+    g = sns.clustermap(df, annot=True, cmap="YlGnBu", row_cluster=True, col_cluster=True, figsize=(15, 15))
 
+    for label in g.ax_heatmap.get_yticklabels():
+        method = label.get_text()
+        label.set_color(get_method_color(method))
+
+    # Apply color to column labels
+    for label in g.ax_heatmap.get_xticklabels():
+        method = label.get_text()
+        label.set_color(get_method_color(method))
+        
     title = f"{long_form_mapping[aspect.upper()]} Results for {metric.capitalize()}"
     g.fig.suptitle(title, x=0.5, y=1.02, fontsize=16, ha='center')  # Center the title above the plot
 
@@ -180,6 +189,12 @@ def plot_family_results(method_names, dataset, family_path="/tmp/family_results.
     # Filter by method names and selected dataset columns
     df = df[df['Method'].isin(method_names)]
 
+    mcc_columns = [col for col in df.columns if col.startswith(f"{dataset}_mcc_")]
+    df['Mean_MCC'] = df[mcc_columns].mean(axis=1)
+
+    # Sort the DataFrame by the mean MCC
+    df = df.sort_values(by='Mean_MCC', ascending=False)
+
     # Filter columns based on the dataset and metrics
     value_vars = [col for col in df.columns if col.startswith(f"{dataset}_") and "_" in col]
 
@@ -244,6 +259,7 @@ def plot_affinity_results(method_names, metric, affinity_path="/tmp/affinity_res
     
     # Gather columns related to the specified metric and validate
     metric_columns = [col for col in df.columns if col.startswith(f"{metric}_")]
+    df = df.sort_values(by=metric_columns, ascending=False)
     df = df[['Method'] + metric_columns].set_index('Method')
 
     df = df.fillna(0)