biopipen 0.21.0__py3-none-any.whl → 0.34.26__py3-none-any.whl

biopipen/scripts/tcr/TCRClusterStats.R

@@ -1,5 +1,5 @@
-source("{{biopipen_dir}}/utils/misc.R")
-source("{{biopipen_dir}}/utils/plot.R")
+{{ biopipen_dir | joinpaths: "utils", "misc.R" | source_r }}
+{{ biopipen_dir | joinpaths: "utils", "plot.R" | source_r }}
 library(tibble)
 library(tidyr)
 library(dplyr)
@@ -7,12 +7,14 @@ library(rlang)
 library(immunarch)
 library(ggprism)
 
-immfile = {{in.immfile | quote}}
-outdir = {{out.outdir | quote}}
+immfile = {{in.immfile | r}}
+outdir = {{out.outdir | r}}
 cluster_size_envs = {{envs.cluster_size | r}}
 shared_clusters_envs = {{envs.shared_clusters | r}}
 sample_diversity_envs = {{envs.sample_diversity | r}}
+joboutdir = {{job.outdir | r}}
 
+log_info("Expanding analysis cases ...")
 expand_cases = function(envs) {
     cases = envs$cases
     envs$cases = NULL
@@ -51,8 +53,9 @@ shared_clusters_cases = expand_cases(shared_clusters_envs)
 sample_diversity_cases = expand_cases(sample_diversity_envs)
 
 cluster_size_distribution = function(name) {
-    print(paste0("- Working on cluster size distribution: ", name))
-    odir = file.path(outdir, "ClusterSizeDistribution", name)
+    log_info("- Working on cluster size distribution: {name}")
+
+    odir = file.path(outdir, "ClusterSizeDistribution", slugify(name))
     dir.create(odir, showWarnings = FALSE, recursive = TRUE)
     case = cluster_size_cases[[name]]
 
@@ -66,6 +69,7 @@ cluster_size_distribution = function(name) {
 
     outfile = file.path(odir, "cluster_size_distribution.txt")
     outplot = file.path(odir, "cluster_size_distribution.png")
+    outplot_pdf = file.path(odir, "cluster_size_distribution.pdf")
     write.table(clsizes, outfile, quote = FALSE, sep = "\t", row.names = FALSE, col.names = TRUE)
 
     plotGG(
@@ -75,16 +79,29 @@ cluster_size_distribution = function(name) {
         ggs = c(
             "theme_prism()",
             "scale_y_continuous(trans='log10')",
-            "labs(x='TCR cluster size', y='Count')"
+            "labs(x='TCR cluster size', y='Count')",
+            "scale_fill_biopipen()"
         ),
         devpars = case$devpars,
-        outfile = outplot
+        outfile = c(outplot, outplot_pdf)
+    )
+
+    add_report(
+        list(
+            src = outplot,
+            name = ifelse(name == "DEFAULT", FALSE, name),
+            descr = paste0("Cluster size distribution for each ", case$by),
+            download = outplot_pdf
+        ),
+        ui = "table_of_images",
+        h1 = "Cluster Size Distribution"
     )
 }
 
 shared_clusters = function(name) {
-    print(paste0("- Working on shared clusters: ", name))
-    odir = file.path(outdir, "SharedClusters", name)
+    log_info("- Working on shared clusters: {name}")
+
+    odir = file.path(outdir, "SharedClusters", slugify(name))
     dir.create(odir, showWarnings = FALSE, recursive = TRUE)
     case = shared_clusters_cases[[name]]
     if (!is.null(case$grouping)) {
@@ -115,13 +132,40 @@ shared_clusters = function(name) {
         row.names=TRUE, col.names=TRUE, quote=FALSE, sep="\t"
     )
 
+    if (!is.null(case$sample_order) && length(case$sample_order) > 0) {
+        if (length(case$sample_order) == 1) {
+            case$sample_order = trimws(strsplit(case$sample_order, ",")[[1]])
+        }
+        nonexisting = setdiff(case$sample_order, samples)
+        if (length(nonexisting) > 0) {
+            stop(paste("  The following samples do not exist in `sample_order`:", paste(nonexisting, collapse=", ")))
+        }
+        plotdata = plotdata[, case$sample_order, drop=FALSE]
+    }
+
     if (is.null(case$heatmap_meta) || length(case$heatmap_meta) == 0) {
         anno = NULL
     } else {
-        anno = as.list(immdata$meta[, case$heatmap_meta, drop=FALSE])
+        anno = as.list(
+            immdata$meta[
+                match(colnames(plotdata), immdata$meta$Sample),
+                case$heatmap_meta,
+                drop=FALSE
+            ])
         anno = do_call(ComplexHeatmap::HeatmapAnnotation, anno)
     }
 
+    cluster_rows = case$cluster_rows && nrow(plotdata) > 2
+    col_samples = colnames(plotdata)
+    if (!cluster_rows) {
+        plotdata = plotdata[col_samples, ]
+        row_samples = col_samples
+    } else {
+        row_samples = samples
+    }
+
+    hmplot = file.path(odir, "shared_clusters.png")
+    hmplot_pdf = file.path(odir, "shared_clusters.pdf")
     # Plot heatmap
     plotHeatmap(
         plotdata,
@@ -129,22 +173,32 @@ shared_clusters = function(name) {
             name = "Shared TCR Clusters",
             col = c("#ffe1e1", "red3"),
             cluster_columns = FALSE,
-            cluster_rows = nrow(plotdata) > 2,
+            cluster_rows = cluster_rows,
             top_annotation = anno,
             cell_fun = if (
                 is.null(case$numbers_on_heatmap) || !case$numbers_on_heatmap
             ) NULL else function(j, i, x, y, width, height, fill) {
-                grid.text(plotdata[samples[i], samples[j]], x, y, gp = gpar(fontsize = 10))
+                grid.text(row_samples[i], col_samples[j], x, y, gp = gpar(fontsize = 10))
             }
         ),
         devpars = case$devpars,
-        outfile = file.path(odir, "shared_clusters.png")
+        outfile = c(hmplot, hmplot_pdf)
+    )
 
+    add_report(
+        list(
+            src = hmplot,
+            download = hmplot_pdf,
+            name = ifelse(name == "DEFAULT", FALSE, name),
+            descr = paste0("Shared TCR clusters across samples")
+        ),
+        ui = "table_of_images",
+        h1 = "Shared TCR Clusters"
     )
 }
 
 shared_clusters_by_grouping = function(name) {
-    odir = file.path(outdir, "SharedClusters", name)
+    odir = file.path(outdir, "SharedClusters", slugify(name))
     case = shared_clusters_cases[[name]]
 
     data = list()
@@ -170,18 +224,55 @@ shared_clusters_by_grouping = function(name) {
     }
 
     outfile = file.path(odir, "shared_clusters.png")
+    outfile_pdf = file.path(odir, "shared_clusters.pdf")
     plotVenn(
         data,
         ggs = 'ggtitle("Shared TCR Clusters")',
         devpars = case$devpars,
-        outfile = outfile
+        outfile = c(outfile, outfile_pdf)
+    )
+
+    add_report(
+        list(
+            src = outfile,
+            download = outfile_pdf,
+            name = ifelse(name == "DEFAULT", FALSE, name),
+            descr = paste0("Shared TCR clusters across ", grouping)
+        ),
+        ui = "table_of_images",
+        h1 = "Shared TCR Clusters"
     )
 }
 
 
+div_methods = list(
+    gini = list(
+        name = "The Gini coefficient",
+        descr = "The Gini coefficient is a measure of statistical dispersion intended to represent the income or wealth distribution of a nation's residents, and is the most commonly used measurement of inequality."
+    ),
+    gini.simp = list(
+        name = "The Gini-Simpson index",
+        descr = "The Gini-Simpson index is a measure of diversity. It is one of the most commonly used in ecology. It is also known as the Simpson index, the Simpson concentration index, the Simpson dominance index, or the Simpson diversity index."
+    ),
+    inv.simp = list(
+        name = "The inverse Simpson index",
+        descr = "It is the effective number of types that is obtained when
+                 the weighted arithmetic mean is used to quantify average
+                 proportional abundance of types in the dataset of interest."
+    ),
+    div = list(
+        name = "The true diversity",
+        descr = "It refers to the number of equally abundant types needed
+                 for the average proportional abundance of the types to
+                 equal that observed in the dataset of interest where all
+                 types may not be equally abundant."
+    )
+)
+
 sample_diversity = function(name) {
-    print(paste0("- Working on sample diversity: ", name))
-    odir = file.path(outdir, "SampleDiversity", name)
+    log_info("- Working on sample diversity: {name}")
+
+    odir = file.path(outdir, "SampleDiversity", slugify(name))
     dir.create(odir, showWarnings = FALSE, recursive = TRUE)
     case = sample_diversity_cases[[name]]
 
@@ -191,8 +282,21 @@ sample_diversity = function(name) {
     }
     outfile = file.path(odir, "diversity.txt")
     outplot = file.path(odir, "diversity.png")
+    outplot_pdf = file.path(odir, "diversity.pdf")
     div = repDiversity(data, .method = case$method)
-    write.table(div, outfile, row.names=TRUE, col.names=TRUE, quote=FALSE, sep="\t")
+    write.table(
+        if (ncol(div) == 1) {
+            as.data.frame(div) %>% rownames_to_column("Sample")
+        } else {
+            div
+        },
+        outfile,
+        row.names=TRUE,
+        col.names=TRUE,
+        quote=FALSE,
+        sep="\t"
+    )
+
     if (case$method == "gini") {
         div = as.data.frame(div) %>% rownames_to_column("Sample")
         colnames(div)[2] = "gini"
@@ -201,7 +305,8 @@ sample_diversity = function(name) {
         mapping = aes(x = Sample, y = gini, fill = Sample)
         ggs = c(
             "theme_prism(axis_text_angle = 90)",
-            "labs(title='Gini coefficient', subtitle='Sample diversity estimation using the Gini coefficient')"
+            "labs(title='Gini coefficient', subtitle='Sample diversity estimation using the Gini coefficient')",
+            "scale_fill_biopipen()"
         )
         if (is.null(case$by) || length(case$by) == 0) {
 
@@ -223,9 +328,8 @@ sample_diversity = function(name) {
             args = list(mapping = mapping),
             ggs = ggs,
             devpars = case$devpars,
-            outfile = outplot
+            outfile = c(outplot, outplot_pdf)
         )
-
     } else {
         if (is.null(case$by) || length(case$by) == 0) {
             p = vis(div)
@@ -242,7 +346,51 @@ sample_diversity = function(name) {
         )
         print(p)
         dev.off()
+
+        pdf(
+            outplot_pdf,
+            width=case$devpars$width / case$devpars$res,
+            height=case$devpars$height / case$devpars$res
+        )
+        print(p)
+        dev.off()
     }
+
+    add_report(
+        list(
+            ui = "flat",
+            label = "Diversity Plot",
+            contents = list(
+                list(
+                    kind = "descr",
+                    content = paste(
+                        div_methods[[case$method]]$name,
+                        ifelse(
+                            is.null(case$by) || length(case$by) == 0,
+                            "",
+                            paste0(" grouped by ", paste(case$by, collapse = ", "))
+                        ),
+                        div_methods[[case$method]]$descr
+                    )
+                ),
+                list(
+                    kind = "image",
+                    src = outplot,
+                    download = outplot_pdf
+                )
+            )
+        ),
+        list(
+            ui = "flat",
+            label = "Diversity Table",
+            contents = list(
+                list(kind = "table", src = outfile, data = list(index_col = 0))
+            )
+        ),
+        ui = "tabs",
+        h2 = ifelse(name == "DEFAULT", "#", name),
+        h1 = "Sample Diversity using TCR clusters"
+    )
 }
 
 
@@ -250,14 +398,20 @@ sample_diversity = function(name) {
     # main
     # --------------------------------------------------
     # Load immunarch data
+    log_info("Loading immunarch data ...")
     immdata = readRDS(immfile)
 
     # Cluster size distribution
+    log_info("Performing cluster size distribution analysis ...")
     sapply(names(cluster_size_cases), cluster_size_distribution)
 
     # Shared clusters
+    log_info("Performing shared clusters analysis ...")
     sapply(names(shared_clusters_cases), shared_clusters)
 
     # Diversity
+    log_info("Performing sample diversity analysis ...")
     sapply(names(sample_diversity_cases), sample_diversity)
+
+    save_report(joboutdir)
 }

biopipen/scripts/tcr/TCRDock.py

@@ -0,0 +1,110 @@
+from __future__ import annotations
+
+import os
+import sys
+from pathlib import Path
+
+import rtoml
+import pandas as pd
+from tempfile import gettempdir
+from biopipen.utils.misc import logger, run_command
+
+configfile: str = {{in.configfile | quote}}  # pyright: ignore  # noqa
+outdir = Path({{out.outdir | quote}})  # pyright: ignore
+envs: dict = {{envs | dict | repr}}  # pyright: ignore
+python: str | list[str] = sys.executable
+
+args = envs.copy()
+config = rtoml.load(Path(configfile))
+args.update(config)
+model_name = args.pop("model_name")
+model_file = Path(args.pop("model_file"))
+data_dir = args.pop("data_dir", None)
+tcrdock: Path | str | None = args.pop("tcrdock", None)
+tmpdir: str = args.pop("tmpdir", gettempdir())
+python = args.pop("python", python)
+
+if not isinstance(python, (list, tuple)):
+    python = [python]
+
+if not data_dir:
+    raise ValueError("`envs.data_dir` is required")
+
+if not tcrdock:
+    logger.info("- `envs.tcrdock` is not provided, cloning the repository ... ")
+    repo_url = "https://github.com/phbradley/TCRdock"
+    commit_id = "c5a7af42eeb0c2a4492a4d4fe803f1f9aafb6193"
+    branch = "main"
+
+    from git import Repo
+    repo = Repo.clone_from(repo_url, tmpdir, branch=branch, no_checkout=True)
+    repo.git.checkout(commit_id)
+    tcrdock = Path(tmpdir) / "TCRdock"
+
+    logger.info("- Running download_blast.py ...")
+    cmd = [
+        *python,
+        tcrdock / "download_blast.py",
+    ]
+    run_command(cmd, fg=True, cwd=str(tcrdock))
+
+tcrdock = str(tcrdock)
+
+if not model_file.is_absolute():
+    model_file = Path(data_dir) / "params" / model_file
+
+os.environ['TF_FORCE_UNIFIED_MEMORY'] = '1'
+os.environ['XLA_PYTHON_CLIENT_MEM_FRACTION'] = '4.0'
+
+logger.info("- Composing targets file ... ")
+targets_file = outdir / "user_targets.tsv"
+targets = pd.DataFrame(
+    [
+        dict(
+            organism=args['organism'],
+            mhc_class=args['mhc_class'],
+            mhc=args['mhc'],
+            peptide=args['peptide'],
+            va=args['va'],
+            ja=args['ja'],
+            cdr3a=args['cdr3a'],
+            vb=args['vb'],
+            jb=args['jb'],
+            cdr3b=args['cdr3b'],
+        )
+    ]
+)
+targets.to_csv(targets_file, sep="\t", index=False)
+
+logger.info("- Generating inputs for AlphaFold modeling ... ")
+cmd = [
+    *python,
+    tcrdock + "/setup_for_alphafold.py",
+    "--targets_tsvfile", targets_file,
+    "--output_dir", outdir / "user_output",
+    "--new_docking",
+]
+run_command(cmd, fg=True)
+
+logger.info("- Running AlphaFold modeling ... ")
+cmd = [
+    *python,
+    tcrdock + "/run_prediction.py",
+    "--verbose",
+    "--targets", outdir / "user_output/targets.tsv",
+    "--outfile_prefix", f"{outdir}/{args['peptide']}",
+    "--model_names", model_name,
+    "--data_dir", data_dir,
+    "--model_params_files", model_file,
+]
+run_command(cmd, fg=True, env={"XLA_FLAGS": "--xla_gpu_force_compilation_parallelism=1"})
+
+logger.info("- Calculating the PAE ... ")
+cmd = [
+    *python,
+    tcrdock + "/add_pmhc_tcr_pae_to_tsvfile.py",
+    "--infile", f"{outdir}/{args['peptide']}_final.tsv",
+    "--outfile", f"{outdir}/{args['peptide']}_w_pae.tsv",
+]
+
+run_command(cmd, fg=True)