Update src/vis_utils.py

src/vis_utils.py

@@ -230,6 +230,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[['Method'] + metric_columns]
 
     df = df.fillna(0)
 
@@ -254,10 +255,6 @@ def plot_affinity_results(method_names, metric, affinity_path="/tmp/affinity_res
     ax.grid(visible=True, which='major', color='gainsboro', linewidth=1.0)
     ax.grid(visible=True, which='minor', color='whitesmoke', linewidth=0.5)
 
-    # Add reference lines
-    ax.plot((best_pipr, best_pipr), (-1, len(df.columns) + 1), color='black', label='Best PIPR')
-    ax.plot((baseline_ac, baseline_ac), (-1, len(df.columns) + 1), color='purple', linestyle='dashed', label='Baseline AC')
-    ax.plot((baseline_ctd, baseline_ctd), (-1, len(df.columns) + 1), color='brown', linestyle='dashed', label='Baseline CTD')
 
     # Apply custom color settings to y-axis labels
     for label in ax.get_yticklabels():