prettify scores

folding_studio_demo/correlate.py

@@ -7,29 +7,31 @@ from scipy.stats import spearmanr, pearsonr, linregress
 
 logger = logging.getLogger(__name__)
 
-SCORE_COLUMNS = [
-        "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"
-    ]
+SCORE_COLUMN_NAMES = {
+    "confidence_score_boltz": "Boltz Confidence Score",
+    "ptm_boltz": "Boltz pTM Score",
+    "iptm_boltz": "Boltz ipTM Score",
+    "complex_plddt_boltz": "Boltz Complex pLDDT",
+    "complex_iplddt_boltz": "Boltz Complex ipLDDT",
+    "complex_pde_boltz": "Boltz Complex pDE",
+    "complex_ipde_boltz": "Boltz Complex ipDE",
+    "interchain_pae_monomer": "Monomer Interchain PAE",
+    "interface_pae_monomer": "Monomer Interface PAE", 
+    "overall_pae_monomer": "Monomer Overall PAE",
+    "interface_plddt_monomer": "Monomer Interface pLDDT",
+    "average_plddt_monomer": "Monomer Average pLDDT",
+    "ptm_monomer": "Monomer pTM Score",
+    "interface_ptm_monomer": "Monomer Interface pTM",
+    "interchain_pae_multimer": "Multimer Interchain PAE",
+    "interface_pae_multimer": "Multimer Interface PAE",
+    "overall_pae_multimer": "Multimer Overall PAE", 
+    "interface_plddt_multimer": "Multimer Interface pLDDT",
+    "average_plddt_multimer": "Multimer Average pLDDT",
+    "ptm_multimer": "Multimer pTM Score",
+    "interface_ptm_multimer": "Multimer Interface pTM"
+}
+
+SCORE_COLUMNS = list(SCORE_COLUMN_NAMES.values())
 
 def compute_correlation_data(spr_data_with_scores: pd.DataFrame, score_cols: list[str]) -> pd.DataFrame:
     corr_data_file = Path("corr_data.csv")
@@ -93,6 +95,9 @@ def plot_correlation_ranking(corr_data: pd.DataFrame, correlation_type: str) ->
 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."""
     
+    # Rename score columns using the mapping in SCORE_COLUMN_NAMES
+    spr_data_with_scores = spr_data_with_scores.rename(columns=SCORE_COLUMN_NAMES)
+    
     corr_data = compute_correlation_data(spr_data_with_scores, score_cols)
     corr_ranking_plot = plot_correlation_ranking(corr_data, "Spearman")