prettify columns

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, COLUMNS
+from folding_studio_demo.correlate import fake_predict_and_correlate, SCORE_COLUMNS
 
 logger = logging.getLogger(__name__)
 
@@ -144,12 +144,26 @@ def model_comparison(api_key: str) -> None:
 
 
 def create_correlation_tab():
-    gr.Markdown("# Upload binding affinity data")
+    gr.Markdown("# Correlation with experimental binding affinity data")
     spr_data_with_scores = pd.read_csv("spr_af_scores_mapped.csv")
+    prettified_columns = {
+            "antibody_name": "Antibody Name",
+            "KD (nM)": "KD (nM)",
+            "antibody_vh_sequence": "Antibody VH Sequence",
+            "antibody_vl_sequence": "Antibody VL Sequence",
+            "antigen_sequence": "Antigen Sequence"
+        }
+    spr_data_with_scores = spr_data_with_scores.rename(columns=prettified_columns)
     with gr.Row():
-        csv_file = gr.File(label="Upload CSV file", file_types=[".csv"])
-    with gr.Row():
-        dataframe = gr.Dataframe(label="Binding Affinity Data")
+        columns = [
+            "Antibody Name",
+            "KD (nM)",
+            "Antibody VH Sequence",
+            "Antibody VL Sequence",
+            "Antigen Sequence"
+        ]
+        # Display dataframe with floating point values rounded to 2 decimal places
+        spr_data = gr.DataFrame(value=spr_data_with_scores[columns].round(2), label="Experimental Antibody-Antigen Binding Affinity Data")
 
     gr.Markdown("# Prediction and correlation")
     with gr.Row():
@@ -162,13 +176,8 @@ def create_correlation_tab():
             correlation_ranking_plot = gr.Plot(label="Correlation ranking")
             correlation_plot = gr.Plot(label="Correlation with binding affinity")
 
-    csv_file.change(
-        fn=lambda file: spr_data_with_scores.drop(columns=COLUMNS) if file else None,
-        inputs=csv_file,
-        outputs=dataframe
-    )
     fake_predict_btn.click(
-        fn=lambda x: fake_predict_and_correlate(spr_data_with_scores, COLUMNS),
+        fn=lambda x: fake_predict_and_correlate(spr_data_with_scores, SCORE_COLUMNS),
         inputs=None,
         outputs=[prediction_dataframe, correlation_ranking_plot, correlation_plot]
     )

folding_studio_demo/correlate.py

@@ -6,7 +6,7 @@ from scipy.stats import spearmanr
 
 logger = logging.getLogger(__name__)
 
-COLUMNS = [
+SCORE_COLUMNS = [
         "confidence_score_boltz",
         "ptm_boltz",
         "iptm_boltz", 
@@ -87,4 +87,4 @@ def fake_predict_and_correlate(spr_data_with_scores: pd.DataFrame, score_cols: l
     cols_to_show = [kd_col]
     cols_to_show.extend(score_cols)
 
-    return spr_data_with_scores[cols_to_show], corr_ranking_plot, corr_plot
+    return spr_data_with_scores[cols_to_show].round(2), corr_ranking_plot, corr_plot