Add residuals layer and diagnostic plots

bartab/app/app.py

@@ -20,10 +20,12 @@
 from bartab.io import load_anndata
 from bartab.models.anndata import AnnDataWLSModel, AnnDataHillModel
 from bartab.plotting import (
+    count_vs_resid,
     dose_response,
     expansion_vs_count,
     expansion_vs_ratio, 
     pred_vs_true,
+    pred_vs_resid,
     time_vs_count, 
     time_vs_ratio, 
     volcano
@@ -298,6 +300,38 @@ def _plot_pred_vs_true(
     )
 
 
+@_plot_wrapper(message="Plotting predicted vs residuals")
+def _plot_pred_vs_resid(
+    adata,
+    highlight=None,
+    control_prefix: str = "ctrl_",
+    mode: str = MODES["single"]
+):
+    do_dose_response = mode == MODES["dose response"]
+    return pred_vs_resid(
+        adata,
+        model_name="HillFitnessModel" if do_dose_response else "WLS",
+        highlight_barcodes=highlight,
+        control_prefix=control_prefix,
+    )
+
+
+@_plot_wrapper(message="Plotting counts vs residuals")
+def _plot_count_vs_resid(
+    adata,
+    highlight=None,
+    control_prefix: str = "ctrl_",
+    mode: str = MODES["single"]
+):
+    do_dose_response = mode == MODES["dose response"]
+    return count_vs_resid(
+        adata,
+        model_name="HillFitnessModel" if do_dose_response else "WLS",
+        highlight_barcodes=highlight,
+        control_prefix=control_prefix,
+    )
+
+
 @_plot_wrapper(message="Plotting volcano")
 def _plot_volcano(
     adata,
@@ -732,6 +766,17 @@ def _invisible_plot(**kwargs):
                 _plot_volcano,
             ),
         }
+    with gr.Row():
+        plots |= {
+            "pred_resid": (
+                _invisible_plot(label="Predicted vs residuals"),
+                _plot_pred_vs_resid,
+            ),
+            "count_resid": (
+                _invisible_plot(label="Counts vs residuals"),
+                _plot_count_vs_resid,
+            ),
+        }
     with gr.Row():
         download = gr.DownloadButton(
             label="Download parameters as CSV",

bartab/cli.py

@@ -98,6 +98,8 @@ def _plot(args: Namespace) -> None:
         expansion_vs_count,
         expansion_vs_ratio, 
         pred_vs_true,
+        pred_vs_resid,
+        count_vs_resid,
         time_vs_count, 
         time_vs_ratio, 
         volcano
@@ -156,6 +158,20 @@ def _plot(args: Namespace) -> None:
         model_name=model_type,
         filename=args.output + f"_pred-obs.{args.plot_format}",
     )
+    fig, axes = pred_vs_resid(
+        adata,
+        control_prefix=args.control,
+        highlight_barcodes=args.highlight,
+        model_name=model_type,
+        filename=args.output + f"_pred-resid.{args.plot_format}",
+    )
+    fig, axes = count_vs_resid(
+        adata,
+        control_prefix=args.control,
+        highlight_barcodes=args.highlight,
+        model_name=model_type,
+        filename=args.output + f"_pred-count.{args.plot_format}",
+    )
 
     fig, axes = volcano(
         adata,

bartab/models/anndata.py

@@ -68,6 +68,7 @@ def fit(
             **kwargs
         )
         adata.layers[f"{name}:predicted"] = preds
+        adata.layers[f"{name}:residual"] = y - preds
 
         results_with_index = []
         for i, (idx, res) in enumerate(zip(

bartab/plotting.py

@@ -272,12 +272,78 @@ def _layered_scatter_barcodes(
     return fig, ax
 
 
+def _rsq(obs, pred, log=False):
+    from scipy.stats import pearsonr
+    import numpy as np
+    if log:
+        obs, pred = np.log10(obs), np.log10(pred)
+    r = pearsonr(pred, obs).statistic
+    SS_res = np.sum(np.square(obs - pred))
+    SS_tot = np.sum(np.square(obs - np.mean(obs)))
+    R2 = 1. - SS_res / SS_tot
+    return r, R2
+
+
+def pred_vs_resid(
+    adata,
+    model_name,
+    filename: str = None,
+    **kwargs
+):
+    fig, ax = _layered_scatter_barcodes(
+        adata, 
+        x_layer=f"{model_name}:predicted", 
+        layer=f"{model_name}:residual", 
+        filename=filename,
+        exp_x=True, 
+        exp_y=True,
+        callback=lambda ax: ax.axhline(1., color="lightgrey", zorder=-1),
+        xlabel=f"Predicted: {model_name}",
+        ylabel="Residuals",
+        **kwargs,
+    )
+    return fig, ax
+
+
+def count_vs_resid(
+    adata,
+    model_name,
+    filename: str = None,
+    **kwargs
+):
+    fig, ax = _layered_scatter_barcodes(
+        adata, 
+        x_layer=None, 
+        layer=f"{model_name}:residual", 
+        filename=filename,
+        exp_x=False, 
+        exp_y=True,
+        callback=lambda ax: ax.axhline(1., color="lightgrey", zorder=-1),
+        xlabel="Counts",
+        ylabel="Residuals",
+        xscale="log",
+        **kwargs,
+    )
+    return fig, ax
+
+
 def pred_vs_true(
     adata,
     model_name,
     filename: str = None,
     **kwargs
 ):
+    from carabiner import print_err
+
+    pred = adata.layers[f"{model_name}:predicted"].ravel()
+    r, R2 = _rsq(
+        pred, 
+        adata.layers["__log_ratio__"].ravel(), 
+        log=False,
+    )
+
+    message = f"Pearson r: {r:.2f}; Rsq = {R2:.2f}, n={len(pred)}"
+    print_err(f"[INFO] {message}")
     fig, ax = _layered_scatter_barcodes(
         adata, 
         x_layer=f"{model_name}:predicted", 
@@ -288,13 +354,12 @@ def pred_vs_true(
         callback=lambda ax: ax.plot(ax.get_xlim(), ax.get_xlim(), color="lightgrey", zorder=-1),
         xlabel=f"Predicted: {model_name}",
         ylabel="Observed:\nbarcode expansion ratio",
-        # xscale="log",
+        title=message,
         **kwargs,
     )
     return fig, ax
 
 
-
 def expansion_vs_count(
     adata,
     filename: str = None,
@@ -457,7 +522,6 @@ def dose_response(
                 "label": "_none",
                 "zorder": 0,
             }
-        # print(bc_df)
         ax.plot(
             "_concentration",
             "fitness",

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "bartab"
-version = "0.0.7"
+version = "0.0.8"
 authors = [
   { name="Eachan Johnson", email="eachan.johnson@crick.ac.uk" },
 ]