select correlation plots

folding_studio_demo/app.py

@@ -8,7 +8,7 @@ from gradio_molecule3d import Molecule3D
 import pandas as pd
 
 from folding_studio_demo.predict import predict
-from folding_studio_demo.correlate import fake_predict_and_correlate, SCORE_COLUMNS
+from folding_studio_demo.correlate import fake_predict_and_correlate, SCORE_COLUMNS, select_correlation_plot
 
 logger = logging.getLogger(__name__)
 
@@ -167,21 +167,29 @@ def create_correlation_tab():
 
     gr.Markdown("# Prediction and correlation")
     with gr.Row():
-        with gr.Column():
-            with gr.Row():
-                fake_predict_btn = gr.Button("Predict structures of all complexes")
-            with gr.Row():
-                prediction_dataframe = gr.Dataframe(label="Predicted Structures Data")
-        with gr.Column():
-            correlation_ranking_plot = gr.Plot(label="Correlation ranking")
-            correlation_plot = gr.Plot(label="Correlation with binding affinity")
+        fake_predict_btn = gr.Button("Predict structures of all complexes")
+    with gr.Row():
+        prediction_dataframe = gr.Dataframe(label="Predicted Structures Data")
+    with gr.Row():
+        correlation_ranking_plot = gr.Plot(label="Correlation ranking")
+    with gr.Row():
+        # User can select the columns to display in the correlation plot
+        correlation_column = gr.Dropdown(label="Score data to display", choices=SCORE_COLUMNS, multiselect=False)
+        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),
         inputs=None,
-        outputs=[prediction_dataframe, correlation_ranking_plot, correlation_plot]
+        outputs=[prediction_dataframe, correlation_ranking_plot]
     )
 
+    # Call function to update the correlation plot when the user selects the columns
+    correlation_column.change(
+        fn=lambda score: select_correlation_plot(spr_data_with_scores, score),
+        inputs=correlation_column,
+        outputs=correlation_plot
+    )
+    
 def __main__():
     with gr.Blocks(title="Folding Studio Demo") as demo:
         gr.Markdown(

folding_studio_demo/correlate.py

@@ -65,26 +65,37 @@ def fake_predict_and_correlate(spr_data_with_scores: pd.DataFrame, score_cols: l
         showlegend=False
     )
 
+    cols_to_show = [kd_col]
+    cols_to_show.extend(score_cols)
+
+    return spr_data_with_scores[cols_to_show].round(2), corr_ranking_plot
+
+def select_correlation_plot(spr_data_with_scores: pd.DataFrame, score: str) -> go.Figure:
+    """Select the correlation plot to display."""
     # corr_plot is a scatter plot of the correlation between the binding affinity and each of the scores
-    scatters = []
-    for score_col in score_cols:
-        scatters.append(
-            go.Scatter(
-                x=spr_data_with_scores[kd_col],
-                y=spr_data_with_scores[score_col],
-                name=f"{kd_col} vs {score_col}",
-                mode='markers',  # Only show markers/dots, no lines
-                hovertemplate="<i>Score:</i> %{y}<br><i>KD:</i> %{x:.2f}<br>"
-            )   
-        )
-    corr_plot = go.Figure(data=scatters)
+    scatter =  go.Scatter(
+            x=spr_data_with_scores["KD (nM)"],
+            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>"
+        )   
+    corr_plot = go.Figure(data=scatter)
     corr_plot.update_layout(
         xaxis_title="KD (nM)",
-        yaxis_title="Score",
+        yaxis_title=score,
         template="simple_white",
-        xaxis_type="log"  # Set x-axis to logarithmic scale
+        # xaxis_type="log"  # Set x-axis to logarithmic scale
     )
-    cols_to_show = [kd_col]
-    cols_to_show.extend(score_cols)
-
-    return spr_data_with_scores[cols_to_show].round(2), corr_ranking_plot, corr_plot
+    # compute the correlation line
+    corr_line = np.polyfit(spr_data_with_scores["KD (nM)"], spr_data_with_scores[score], 1)
+    corr_line_x = np.linspace(min(spr_data_with_scores["KD (nM)"]), max(spr_data_with_scores["KD (nM)"]), 100)
+    corr_line_y = corr_line[0] * corr_line_x + corr_line[1]
+    # add the correlation line to the plot
+    corr_plot.add_trace(go.Scatter(
+        x=corr_line_x,
+        y=corr_line_y,
+        mode='lines',
+        name=f"Correlation line for {score}"
+    ))
+    return corr_plot