refactor

folding_studio_demo/app.py

@@ -158,7 +158,6 @@ def fake_predict_and_correlate(spr_data_with_scores: pd.DataFrame, score_cols: l
     corr_data = pd.DataFrame(corr_data)
     # Find the lines in corr_data with NaN values and remove them
     corr_data = corr_data[corr_data["correlation"].notna()]
-    logger.info("Correlation data: %s", corr_data)
     # Sort correlation data by correlation value
     corr_data = corr_data.sort_values('correlation', ascending=True)
     
@@ -199,6 +198,59 @@ def fake_predict_and_correlate(spr_data_with_scores: pd.DataFrame, score_cols: l
 
     return spr_data_with_scores[cols_to_show], corr_ranking_plot, corr_plot
 
+def create_correlation_tab():
+    gr.Markdown("# Upload binding affinity data")
+    spr_data_with_scores = pd.read_csv("spr_af_scores_mapped.csv")
+    with gr.Row():
+        csv_file = gr.File(label="Upload CSV file", file_types=[".csv"])
+    with gr.Row():
+        dataframe = gr.Dataframe(label="Binding Affinity Data")
+
+    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")
+
+    cols = [
+        "confidence_score_boltz",
+        "ptm_boltz",
+        "iptm_boltz", 
+        "complex_plddt_boltz",
+        "complex_iplddt_boltz",
+        "complex_pde_boltz",
+        "complex_ipde_boltz",
+        "interchain_pae_monomer",
+        "interface_pae_monomer",
+        "overall_pae_monomer",
+        "interface_plddt_monomer",
+        "average_plddt_monomer",
+        "ptm_monomer",
+        "interface_ptm_monomer", 
+        "interchain_pae_multimer",
+        "interface_pae_multimer",
+        "overall_pae_multimer",
+        "interface_plddt_multimer",
+        "average_plddt_multimer",
+        "ptm_multimer",
+        "interface_ptm_multimer"
+    ]
+    csv_file.change(
+        fn=lambda file: spr_data_with_scores.drop(columns=cols) if file else None,
+        inputs=csv_file,
+        outputs=dataframe
+    )
+    fake_predict_btn.click(
+        fn=lambda x: fake_predict_and_correlate(spr_data_with_scores, cols),
+        inputs=None,
+        outputs=[prediction_dataframe, correlation_ranking_plot, correlation_plot]
+    )
+
 def __main__():
     with gr.Blocks(title="Folding Studio Demo") as demo:
         gr.Markdown(
@@ -222,58 +274,6 @@ def __main__():
         with gr.Tab("📊 Model Comparison"):
             model_comparison(api_key)
         with gr.Tab("🔍 Correlations"):
-            gr.Markdown("# Upload binding affinity data")
-            spr_data_with_scores = pd.read_csv("spr_af_scores_mapped.csv")
-            with gr.Row():
-                csv_file = gr.File(label="Upload CSV file", file_types=[".csv"])
-            with gr.Row():
-                dataframe = gr.Dataframe(label="Binding Affinity Data")
-
-            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")
-
-            
-
-            cols = [
-                "confidence_score_boltz",
-                "ptm_boltz",
-                "iptm_boltz", 
-                "complex_plddt_boltz",
-                "complex_iplddt_boltz",
-                "complex_pde_boltz",
-                "complex_ipde_boltz",
-                "interchain_pae_monomer",
-                "interface_pae_monomer",
-                "overall_pae_monomer",
-                "interface_plddt_monomer",
-                "average_plddt_monomer",
-                "ptm_monomer",
-                "interface_ptm_monomer", 
-                "interchain_pae_multimer",
-                "interface_pae_multimer",
-                "overall_pae_multimer",
-                "interface_plddt_multimer",
-                "average_plddt_multimer",
-                "ptm_multimer",
-                "interface_ptm_multimer"
-            ]
-            csv_file.change(
-                fn=lambda file: spr_data_with_scores.drop(columns=cols) if file else None,
-                inputs=csv_file,
-                outputs=dataframe
-            )
-            fake_predict_btn.click(
-                fn=lambda x: fake_predict_and_correlate(spr_data_with_scores, cols),
-                inputs=None,
-                outputs=[prediction_dataframe, correlation_ranking_plot, correlation_plot]
-            )
+            create_correlation_tab()
 
     demo.launch()