make log a global option

folding_studio_demo/app.py

@@ -321,13 +321,18 @@ def create_correlation_tab():
         outputs=[prediction_dataframe],
     )
     with gr.Row():
-        correlation_type = gr.Radio(
-            choices=["Spearman", "Pearson", "R²"],
-            value="Spearman",
-            label="Correlation Type",
-            interactive=True,
-            scale=0,
-        )
+        with gr.Row():
+            correlation_type = gr.Radio(
+                choices=["Spearman", "Pearson"],
+                value="Spearman",
+                label="Correlation Type",
+                interactive=True,
+                scale=0,
+            )
+        with gr.Row():
+            log_scale = gr.Checkbox(
+                label="Use log scale for KD", value=False
+            )
         correlation_ranking_plot = gr.Plot(label="Correlation ranking")
     with gr.Row():
         with gr.Column(scale=0):
@@ -338,10 +343,6 @@ def create_correlation_tab():
                 multiselect=False,
                 value=SCORE_COLUMNS[0],
             )
-            # Add checkbox for log scale and update plot when either input changes
-            log_scale = gr.Checkbox(
-                label="Display x-axis on logarithmic scale", value=False
-            )
             score_description = gr.Markdown(
                 get_score_description(correlation_column.value)
             )
@@ -361,31 +362,29 @@ def create_correlation_tab():
         outputs=[prediction_dataframe, correlation_ranking_plot, regression_plot],
     )
 
-    def update_regression_plot(score, use_log):
-        return make_regression_plot(spr_data_with_scores, score, use_log)
-
-    def update_correlation_plot(correlation_type):
+    def update_plots_with_log(correlation_type, score, use_log):
         logger.info(f"Updating correlation plot for {correlation_type}")
         corr_data = compute_correlation_data(spr_data_with_scores, SCORE_COLUMNS)
         logger.info(f"Correlation data: {corr_data}")
-        return plot_correlation_ranking(corr_data, correlation_type)
+        corr_ranking_plot = plot_correlation_ranking(corr_data, correlation_type, kd_col="KD (nM)" if not use_log else "log_kd")
+        regression_plot = make_regression_plot(spr_data_with_scores, score, use_log)
+        return regression_plot, corr_ranking_plot
 
     correlation_column.change(
-        fn=update_regression_plot,
-        inputs=[correlation_column, log_scale],
-        outputs=regression_plot,
+        fn=update_plots_with_log,
+        inputs=[correlation_type, correlation_column, log_scale],
+        outputs=[regression_plot, correlation_ranking_plot],
     )
 
     correlation_type.change(
-        fn=update_correlation_plot,
-        inputs=[correlation_type],
-        outputs=correlation_ranking_plot,
+        fn=update_plots_with_log,
+        inputs=[correlation_type, correlation_column, log_scale],
+        outputs=[regression_plot, correlation_ranking_plot],
     )
-
     log_scale.change(
-        fn=update_regression_plot,
-        inputs=[correlation_column, log_scale],
-        outputs=regression_plot,
+        fn=update_plots_with_log,
+        inputs=[correlation_type, correlation_column, log_scale],
+        outputs=[regression_plot, correlation_ranking_plot],
     )
 
 

folding_studio_demo/correlate.py

@@ -75,30 +75,26 @@ def compute_correlation_data(
     corr_funcs = {}
     corr_funcs["Spearman"] = spearmanr
     corr_funcs["Pearson"] = pearsonr
-    corr_funcs["R²"] = linregress
-    for correlation_type, corr_func in corr_funcs.items():
-        for score_col in score_cols:
-            logger.info(
-                f"Computing {correlation_type} correlation between {score_col} and KD (nM)"
-            )
-            res = corr_func(
-                spr_data_with_scores[kd_col], spr_data_with_scores[score_col]
-            )
-            logger.info(f"Correlation function: {corr_func}")
-            correlation_value = (
-                res.rvalue**2 if correlation_type == "R²" else res.statistic
-            )
-            corr_data.append(
-                {
-                    "correlation_type": correlation_type,
-                    "score": score_col,
-                    "correlation": correlation_value,
-                    "p-value": res.pvalue,
-                }
-            )
-            logger.info(
-                f"Correlation {correlation_type} between {score_col} and KD (nM): {correlation_value}"
-            )
+    for kd_col in ["KD (nM)", "log_kd"]:
+        for correlation_type, corr_func in corr_funcs.items():
+            for score_col in score_cols:
+                logger.info(
+                    f"Computing {correlation_type} correlation between {score_col} and {kd_col}"
+                )
+                res = corr_func(
+                    spr_data_with_scores[kd_col], spr_data_with_scores[score_col]
+                )
+                logger.info(f"Correlation function: {corr_func}")
+                correlation_value = res.statistic
+                corr_data.append(
+                    {
+                        "correlation_type": correlation_type,
+                        "kd_col": kd_col,
+                        "score": score_col,
+                        "correlation": correlation_value,
+                        "p-value": res.pvalue,
+                    }
+                )
 
     corr_data = pd.DataFrame(corr_data)
     # Find the lines in corr_data with NaN values and remove them
@@ -112,10 +108,13 @@ def compute_correlation_data(
 
 
 def plot_correlation_ranking(
-    corr_data: pd.DataFrame, correlation_type: str
+    corr_data: pd.DataFrame, correlation_type: str, kd_col: str
 ) -> go.Figure:
     # Create bar plot of correlations
-    data = corr_data[corr_data["correlation_type"] == correlation_type]
+    data = corr_data[
+        (corr_data["correlation_type"] == correlation_type)
+        & (corr_data["kd_col"] == kd_col)
+    ]
     corr_ranking_plot = go.Figure(
         data=[
             go.Bar(
@@ -144,7 +143,7 @@ def fake_predict_and_correlate(
     """Fake predict structures of all complexes and correlate the results."""
 
     corr_data = compute_correlation_data(spr_data_with_scores, score_cols)
-    corr_ranking_plot = plot_correlation_ranking(corr_data, "Spearman")
+    corr_ranking_plot = plot_correlation_ranking(corr_data, "Spearman", kd_col="KD (nM)")
 
     cols_to_show = main_cols[:]
     cols_to_show.extend(score_cols)