improve corr plot

folding_studio_demo/app.py

@@ -178,7 +178,7 @@ def create_correlation_tab():
         correlation_plot = gr.Plot(label="Correlation with binding affinity")
 
     fake_predict_btn.click(
-        fn=lambda x: fake_predict_and_correlate(spr_data_with_scores, SCORE_COLUMNS),
+        fn=lambda x: fake_predict_and_correlate(spr_data_with_scores, SCORE_COLUMNS, ["Antibody Name", "KD (nM)"]),
         inputs=None,
         outputs=[prediction_dataframe, correlation_ranking_plot]
     )

folding_studio_demo/correlate.py

@@ -30,7 +30,7 @@ SCORE_COLUMNS = [
         "interface_ptm_multimer"
     ]
 
-def fake_predict_and_correlate(spr_data_with_scores: pd.DataFrame, score_cols: list[str]) -> tuple[pd.DataFrame, go.Figure]:
+def fake_predict_and_correlate(spr_data_with_scores: pd.DataFrame, score_cols: list[str], main_cols: list[str]) -> tuple[pd.DataFrame, go.Figure]:
     """Fake predict structures of all complexes and correlate the results."""
     corr_data = []
     spr_data_with_scores["log_kd"] = np.log10(spr_data_with_scores["KD (nM)"])
@@ -65,7 +65,7 @@ def fake_predict_and_correlate(spr_data_with_scores: pd.DataFrame, score_cols: l
         showlegend=False
     )
 
-    cols_to_show = [kd_col]
+    cols_to_show = main_cols[:]
     cols_to_show.extend(score_cols)
 
     return spr_data_with_scores[cols_to_show].round(2), corr_ranking_plot
@@ -78,13 +78,21 @@ def select_correlation_plot(spr_data_with_scores: pd.DataFrame, score: str) -> g
             y=spr_data_with_scores[score],
             name=f"KD (nM) vs {score}",
             mode='markers',  # Only show markers/dots, no lines
-            hovertemplate="<i>Score:</i> %{y}<br><i>KD:</i> %{x:.2f}<br>"
+            hovertemplate="<i>Score:</i> %{y}<br><i>KD:</i> %{x:.2f}<br>",
+            marker=dict(color='#1f77b4')  # Set color to match default first color
         )   
     corr_plot = go.Figure(data=scatter)
     corr_plot.update_layout(
         xaxis_title="KD (nM)",
         yaxis_title=score,
         template="simple_white",
+        legend=dict(
+            orientation="h",
+            yanchor="bottom",
+            y=1.02,
+            xanchor="right",
+            x=1
+        )
         # xaxis_type="log"  # Set x-axis to logarithmic scale
     )
     # compute the correlation line
@@ -96,6 +104,7 @@ def select_correlation_plot(spr_data_with_scores: pd.DataFrame, score: str) -> g
         x=corr_line_x,
         y=corr_line_y,
         mode='lines',
-        name=f"Correlation line for {score}"
+        name=f"Correlation",
+        line=dict(color='#1f77b4')  # Set same color as scatter points
     ))
     return corr_plot