Update src/vis_utils.py

src/vis_utils.py

@@ -230,32 +230,43 @@ 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}_")]
-    
-    # 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)
 
+    df = df.T
+
+    print(df)
     # Set up the plot
-    sns.set(rc={'figure.figsize': (13.7, 8.27)})
+    sns.set(rc={'figure.figsize': (11.7, 8.27)})
     sns.set_theme(style="whitegrid", color_codes=True)
 
-    # Create a boxplot for the metric
-    ax = sns.boxplot(data=df_long, x='Value', y='Method', hue='Fold', whis=np.inf, orient="h")
-    
-    # Customize x-axis and y-axis tickers and grid
+    # Create the boxplot
+    ax = sns.boxplot(data=df_corr, whis=np.inf, orient="h")
+
+    # Add a swarmplot on top of the boxplot
+    sns.swarmplot(data=df_corr, orient="h", color=".1", ax=ax)
+
+    # Set labels and x-axis formatting
+    ax.set_xlabel("Percent Pearson Correlation")
     ax.xaxis.set_major_locator(ticker.MultipleLocator(5))
-    ax.get_xaxis().set_minor_locator(ticker.AutoMinorLocator())
-    ax.get_yaxis().set_minor_locator(ticker.AutoMinorLocator())
+    ax.xaxis.set_minor_locator(ticker.AutoMinorLocator())
+    ax.yaxis.set_minor_locator(ticker.AutoMinorLocator())
     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_corr.columns) + 1), color='black', label='Best PIPR')
+    ax.plot((baseline_ac, baseline_ac), (-1, len(df_corr.columns) + 1), color='purple', linestyle='dashed', label='Baseline AC')
+    ax.plot((baseline_ctd, baseline_ctd), (-1, len(df_corr.columns) + 1), color='brown', linestyle='dashed', label='Baseline CTD')
+
     # Apply custom color settings to y-axis labels
     for label in ax.get_yticklabels():
         method = label.get_text()
         label.set_color(get_method_color(method))
 
+    # Add legend
+    ax.legend(loc='best', frameon=True)
+
     # Save the plot
     save_path = "/tmp"
     filename = os.path.join(save_path, f"{metric}_affinity_results.png")