compute_vdjdb_motifs.Rmd

---
title: "Inferring motifs from VDJdb and building motif PWMs"
author: "Mikhail Shugay"
date: "09-03-2022"
output:
  pdf_document: default
---

```{r setup, include=FALSE}
knitr::opts_chunk$set(echo = TRUE)
library(dplyr)
library(data.table)
library(stringr)
library(stringdist)
library(reshape2)
library(igraph)
library(ggplot2)
library(ggrepel)
library(ggseqlogo)
library(parallel)
library(cowplot)
select = dplyr::select

VDJTOOLS_PATH = "/software/bin/vdjtools"
VDJDB_SLIM_PATH = "vdjdb_dump/vdjdb.slim.txt"
POOLS_PATH = "pools/"
```

```{r}
if (!file.exists(VDJTOOLS_PATH)) {
  stop(paste("VDJtools binary not found at", VDJTOOLS_PATH))
}

if (!file.exists(VDJDB_SLIM_PATH)) {
  stop(paste("VDJdb database not found at", VDJDB_SLIM_PATH))
}

if (!dir.exists(POOLS_PATH)) {
  system("wget https://zenodo.org/record/6339774/files/pools.zip")
  system("unzip pools.zip")
}
```

## Export VDJdb to VDJtools format

Load latest DB for human and mouse

```{r}
df.vdjdb = fread(VDJDB_SLIM_PATH) %>%
  filter(species %in% c("HomoSapiens", "MusMusculus"),
         v.segm != "", j.segm != "", !is.na(j.segm)) %>%
  select(cdr3, v.segm, j.segm, v.end, j.start, antigen.epitope, gene, species) %>%
  unique
```

Select epitopes with at least $10$ specific TCR sequences

```{r}
dt.epi.count = df.vdjdb %>%
  group_by(species, gene, antigen.epitope) %>%
  summarise(total = n()) %>%
  arrange(-total)

dt.epi.count %>% head(n = 10)

df.vdjdb.good = df.vdjdb %>%
  merge(dt.epi.count %>% filter(total >= 10))
```
Back translate mock functions

```{r}
mock_codons <- c('GCT', 'TGT', 'GAT', 'GAA', 'TTT',
                 'GGT', 'ATT', 'CAT', 'AAA', 'TTA',
                 'ATG', 'AAT', 'CCT', 'CAA', 'CGT',
                 'TCT', 'ACT', 'GTT', 'TGG', 'TAT')

names(mock_codons) <- c('A', 'C', 'D', 'E', 'F',
                        'G', 'I', 'H', 'K', 'L',
                        'M', 'N', 'P', 'Q', 'R',
                        'S', 'T', 'V', 'W', 'Y')

mock_back_translate <- function(x) {
  paste0(mock_codons[x], collapse = "")
}

# "CASS" %>% strsplit('') %>% lapply(mock_back_translate)
```

Convert to VDJtools format

```{r}
df.vdjtools <- df.vdjdb.good %>%
  mutate(count = 1,
         cdr3aa = cdr3,
         cdr3nt = cdr3 %>% 
           strsplit('') %>% 
           lapply(mock_back_translate),
         v = v.segm, d = "", j = j.segm,
         vend = v.end, dstart = -1, dend = -1, jstart = j.start) %>% # 1-based, inclusive
  group_by(gene, antigen.epitope, species) %>%
  mutate(freq = count / sum(count)) %>%
  select(count, freq, cdr3nt, cdr3aa, v, d, j, vend, dstart, dend, jstart, 
         gene, antigen.epitope, species)
```

Write separate VDJtools-formatted file for each epitope

```{r}
system("mkdir vdjdb_export/")

df.vdjtools %>%
  group_by(gene, antigen.epitope, species) %>%
  do(data.frame(
    fwrite(., paste0("vdjdb_export/", .$species[1], ".",
                     .$gene[1], ".", .$antigen.epitope[1], ".txt"), sep = "\t", quote = F)
    ))
```

## Run TCRNET

List split and formatted files from VDJdb

```{r}
filenames <- list.files("vdjdb_export/")
tmp <- str_split_fixed(filenames, fixed("."), 4)[,1:3]
df.files <- data.frame(file.name = filenames, species = tmp[,1], gene = tmp[,2]) %>%
  mutate(species = ifelse(species == "HomoSapiens", "human", "mouse"),
         gene = tolower(gene)) %>%
  mutate(input = paste0("vdjdb_export/", file.name),
         background = paste0("pools/", species, ".", gene, ".aa.txt"))
```

Run TCRNET algorithm via VDJtools JAR
         
```{r}
df.files %>%
  group_by(background) %>%
  summarise(cmd = paste0(VDJTOOLS_PATH, " -Xmx100G CalcDegreeStats -o 1,0,1 -g2 dummy -b ",
         background[1], " ", paste(input, collapse = " "), " tcrnet/")) %>%
  group_by(background) %>%
  do(data.frame(system(.$cmd)))
```

Load TCRNET results

```{r}
df.tcrnet <- filenames %>%
  lapply(function(x) fread(paste0("tcrnet/", x)) %>%
                             mutate(file.name = x)) %>%
  rbindlist %>%
  merge(data.frame(file.name = filenames, 
                   species = tmp[,1], 
                   gene = tmp[,2],
                   antigen.epitope = tmp[,3]))
```

Adjust P-values and select enriched variants

```{r}
df.tcrnet <- df.tcrnet %>%
  group_by(species, gene, antigen.epitope) %>%
  mutate(p.adj = p.value.g) %>%
  ungroup %>%
  mutate(enriched = degree.s >= 2 & p.adj < 0.05)
```

## Analyze CDR3 graphs

Build hamming graph neighborhood around enriched CDR3s

```{r}
compute_edges <- function(cdr3.1, cdr3.2, combine) {
  if (combine) {
    cdr3.1 <- c(cdr3.1, cdr3.2) %>% unique
    cdr3.2 <- cdr3.1
  }
  
  stringdistmatrix(cdr3.1, cdr3.2, 
                   method = "hamming", 
                   nthread = 80,
                   useNames = "strings") %>%
    as.matrix %>%
    melt %>% 
    filter(value == 1) %>%
    mutate(cdr3.from = as.character(Var1), 
           cdr3.to = as.character(Var2)) %>%
    select(-Var1, -Var2, -value) %>%
    unique
}

# I - between enriched and neighbors
df.graph.e <- df.tcrnet %>%
  group_by(species, gene, antigen.epitope) %>%
  do(compute_edges(.$cdr3aa[which(.$enriched)], .$cdr3aa, F)) %>%
  ungroup

# II - between all from previous graph
df.graph <- df.graph.e %>%
  group_by(species, gene, antigen.epitope) %>%
  do(compute_edges(.$cdr3.from, .$cdr3.to, T)) %>%
  ungroup
```

Select connected componenets (or "clusters") that represent CDR3 motifs

```{r}
compute_cc <- function(from, to) {
  gg <- data.frame(from, to) %>%
    graph_from_data_frame %>%
    simplify
  
  cc <- gg %>%
    clusters
  
  coords <- gg %>%
      layout_with_graphopt(niter = 3000, charge = 0.005)
  
  data.frame(cdr3aa = names(V(gg)),
             x = coords[,1],
             y = coords[,2],
             stringsAsFactors = F) %>%
    merge(data.frame(cdr3aa = names(cc$membership),
                 cid = cc$membership,
                 stringsAsFactors = F))
}

df.cc <- df.graph %>%
  group_by(species, gene, antigen.epitope) %>%
  do(compute_cc(.$cdr3.from, .$cdr3.to)) %>%
  ungroup %>%
  mutate(cid = paste(substr(species, 1, 1),
                     substr(gene, 3, 3),
                     antigen.epitope,
                     cid, sep = ".")) %>%
  group_by(cid) %>%
  mutate(csz = n()) %>%
  ungroup

most_freq_segm <- function(segm) {
  segm <- str_split_fixed(segm, ",", 2)[,1]
  names(sort(table(segm), decreasing=T))[1]
}

df.cc <- df.cc %>%
  merge(df.vdjdb %>% 
          mutate(cdr3aa = cdr3) %>% 
          filter(cdr3aa %in% df.cc$cdr3aa) %>%
          select(-cdr3)) %>%
  group_by(cid) %>%
  # select representative V/J
  mutate(v.segm.repr = most_freq_segm(v.segm), 
         j.segm.repr = most_freq_segm(j.segm)) %>%
  ungroup
```

Select and write clusters with at least 5 members

```{r}
df.cc.sel <- df.cc %>%
  filter(csz >= 5)

# append HLA info
df.mhc.info <- fread(VDJDB_SLIM_PATH) %>%
  select(species, antigen.epitope, antigen.gene, antigen.species, 
         mhc.a, mhc.b, mhc.class) %>%
  group_by(species, antigen.epitope) %>%
  summarise(antigen.gene = most_freq_segm(antigen.gene),
            antigen.species = most_freq_segm(antigen.species),
            mhc.a = most_freq_segm(mhc.a),
            mhc.b = most_freq_segm(mhc.b),
            mhc.class = most_freq_segm(mhc.class)) %>%
  ungroup

df.mhc.info %>%
  merge(df.cc.sel) %>%
  fwrite("cluster_members.txt", sep = "\t")
```

Plot epitopes with at least 30 sequences combined in clusters for human

```{r fig.width=8, fig.height=8}
p1 <- df.cc.sel %>%
  group_by(species, gene, antigen.epitope) %>%
  mutate(total = n()) %>%
  filter(antigen.epitope != "KLGGALQAK", species == "HomoSapiens", gene == "TRA", total >= 30) %>%
  ggplot(aes(x = x, y = y, color = nchar(cdr3aa))) +
  geom_point(color = "black") +
  geom_point() +
  xlab("") + ylab("") +
  scale_color_distiller(palette = "Spectral", guide = F) +
  facet_wrap(~ antigen.epitope, scales = "free", ncol = 3) +
  ggtitle("TRA, epitopes w/ 30+ sequences, colored by CDR3len") +
  theme_minimal() +
  theme(aspect = 1, axis.text=element_blank())
p1

p2 <- df.cc.sel %>%
  group_by(species, gene, antigen.epitope) %>%
  mutate(total = n()) %>%
  filter(antigen.epitope != "KLGGALQAK", species == "HomoSapiens", gene == "TRB", total >= 30) %>%
  ggplot(aes(x = x, y = y, color = nchar(cdr3aa))) +
  geom_point(color = "black") +
  geom_point() +
  xlab("") + ylab("") +
  scale_color_distiller(palette = "Spectral", guide = F) +
  facet_wrap(~ antigen.epitope, scales = "free", ncol = 3) +
  ggtitle("TRB, epitopes w/ 30+ sequences, colored by CDR3len") +
  theme_minimal() +
  theme(aspect = 1, axis.text=element_blank())
p2
```

## Compute position-weight matrices for motifs

Flatten CDR3 sequences

```{r}
df.flataa <- df.cc.sel %>%
  .$cdr3aa %>% 
  unique %>%
  strsplit("") %>%
  mclapply(function(x) data.frame(cdr3aa = paste0(x, collapse = ""), 
                                  aa = x, stringsAsFactors = F) %>%
             mutate(pos = 0:(n()-1)),
           mc.cores = 80) %>%
  rbindlist
```

Append sequences and compute PWMs

```{r}
df.pwms <- df.cc.sel %>%
  merge(df.flataa) %>%
  mutate(len = nchar(cdr3aa)) %>%
  group_by(species, gene, antigen.epitope,
           v.segm.repr, j.segm.repr, cid, csz, aa, pos, len) %>%
  summarise(count = n()) %>%
  ungroup
```

Load background PWMs for normalization

```{r}
df.bg.pwms <- expand.grid(species = unique(df.pwms$species), gene = unique(df.pwms$gene)) %>%
  mutate(species.1 = ifelse(species == "HomoSapiens", "human", "mouse"),
         gene.1 = tolower(gene)) %>%
  mutate(file.name = paste0("pools/pwms/", species.1, ".", gene.1, ".aa.aa_cdr3_pwm.txt")) %>%
  group_by(species, gene) %>%
  do(fread(.$file.name)) %>%
  mutate(v.segm.repr = v, j.segm.repr = j, count.bg = value) %>%
  select(-sample.id, -v, -j, -value) %>%
  group_by(species, gene, v.segm.repr, j.segm.repr, len) %>%
  mutate(total.bg = sum(count.bg[which(pos == 0)])) %>%
  ungroup

# for imputation, to use where we don't have data
df.bg.pwms.i <- df.bg.pwms %>%
  group_by(species, gene, len, pos, aa) %>%
  summarise(count.bg.i = sum(count.bg)) %>%
  group_by(species, gene, len) %>%
  mutate(total.bg.i = sum(count.bg.i[which(pos == 0)])) %>%
  ungroup
```

Merge with PWMs for VDJdb motifs

```{r}
df.pwms.info <- df.pwms %>%
  merge(df.bg.pwms, all.x = T) %>%
  merge(df.bg.pwms.i, all.x = T) %>%
  filter(total.bg > 0) %>%
  group_by(cid) %>%
  mutate(need.impute = sum(count.bg[which(pos == 0)], na.rm = T) == 0,
         count.bg = ifelse(is.na(count.bg), 0, count.bg),
         count.bg.i = ifelse(is.na(count.bg.i), 0, count.bg.i),
         total.bg = max(total.bg, na.rm = T),
         total.bg.i = max(total.bg.i, na.rm = T)) %>%
  ungroup
```

Compute frequencies, information and normalized information

```{r fig.width=8, fig.height=6}
df.pwms.info.h <- df.pwms.info %>%
  group_by(species, antigen.epitope, gene, cid, csz,
           v.segm.repr, j.segm.repr, pos, len) %>%
  mutate(freq = count / csz, 
         freq.bg = ifelse(need.impute,
                          (count.bg.i + 1) / (total.bg.i + 1),
                          (count.bg + 1) / (total.bg + 1)),
         I = 1 + sum(freq * log(freq)) / log(20),
         I.norm = -sum(freq * log(freq.bg)) / log(20) / 2,
         height.I = freq * I,
         height.I.norm = freq * I.norm) %>%
  ungroup

# some plotting
p3 <- df.pwms.info.h %>%
  select(species, gene, cid, len, pos, I, I.norm) %>%
  unique %>%
  filter(species == "HomoSapiens", gene == "TRA", len > 10, len < 17) %>%
  melt(measure.vars = c("I", "I.norm")) %>%
  ggplot(aes(x = pos, y = value, group = paste(cid, variable),
             color = variable)) +
  geom_line(alpha = 0.5) + ylab("") +
  scale_x_continuous("position in CDR3", breaks = 0:16) +
  scale_color_brewer("", palette = "Set1") +
  facet_wrap(~len, scales = "free_x") +
  ggtitle("PWM information for TRA, CDR3 length in [11, 16]") +
  theme_classic() +
  theme(legend.position = "bottom")
p3

p4 <- df.pwms.info.h %>%
  select(species, gene, cid, len, pos, I, I.norm) %>%
  unique %>%
  filter(species == "HomoSapiens", gene == "TRB", len > 10, len < 17) %>%
  melt(measure.vars = c("I", "I.norm")) %>%
  ggplot(aes(x = pos, y = value, group = paste(cid, variable),
             color = variable)) +
  geom_line(alpha = 0.5) + ylab("") +
  scale_x_continuous("position in CDR3", breaks = 0:16) +
  scale_color_brewer("", palette = "Set1") +
  facet_wrap(~len, scales = "free_x") +
  ggtitle("PWM information for TRB, CDR3 length in [11, 16]") +
  theme_classic() +
  theme(legend.position = "bottom")
p4
```

Append MHC information

> NB present analysis doesn't consider cases when the same epitope is presented by different MHC

```{r}
df.pwms.info.h %>%
  merge(df.mhc.info) %>%
  fwrite("motif_pwms.txt", sep = "\t")
```

### An example of motif visualization

For A*02 YLQ. TCR similarity map

```{r fig.height=6, fig.width=8}
p5 <- df.cc.sel %>%
  filter(antigen.epitope == "YLQPRTFLL") %>%
  ggplot(aes(x = x, y = y)) +
  geom_density_2d_filled(color = "black", contour = F) +
  geom_point(aes(color = cid)) +
  #geom_text_repel(aes(label = cdr3aa), 
  #                min.segment.length = 0, segment.alpha = 0.5,
  #                color = "black", cex = 2.5) +
  xlab("") + ylab("") +
  scale_color_discrete("CID") +
  #scale_color_brewer(palette = "Paired", guide = F) +
  scale_fill_brewer(guide=F, palette = "Greys") +
  facet_wrap( ~ gene, scales = "free") +
  ggtitle("CDR3 amino acid sequence similarity map for A*02 YLQ") +
  theme_minimal() +
  theme(aspect = 1,
        axis.text = element_blank(),
        panel.grid.major = element_blank(),
        panel.grid.minor = element_blank()) 

p5
```
and motif PWMs

```{r fig.height=6, fig.width=8}
p6 <- plot_grid(plotlist = df.cc.sel %>%
                               filter(antigen.epitope == "YLQPRTFLL") %>%
                               mutate(cid = str_replace(cid, "H\\.", "TR")) %>%
                               group_by(cid) %>%
                               group_map(~ggseqlogo(.x$cdr3aa) +
                                           ggtitle(.y) + 
                                           theme(legend.position = "none",
                                                 axis.title.y = element_blank(),
                                                 axis.text.y = element_blank())),
                             ncol = 3)

p6
```

### For COVID-19 whitepaper

```{r fig.height=6, fig.width=8}
df.cc.sel.cov <- fread(VDJDB_SLIM_PATH) %>%
  filter(antigen.species == "SARS-CoV-2") %>%
  mutate(cdr3aa = cdr3) %>%
  select(cdr3aa, antigen.epitope, reference.id) %>%
  mutate(study = ifelse(grepl(",", reference.id, fixed = T), 
                        ifelse(grepl("215", reference.id), "multi-structure", "multi"), 
                        ifelse(grepl("215", reference.id), "structure", reference.id))) %>%
  select(-reference.id) %>%
  merge(df.cc.sel) %>%
  merge(
    data.frame(study = c(
      "PMID:33664060",
      "https://github.com/antigenomics/vdjdb-db/issues/326",
      "https://github.com/antigenomics/vdjdb-db/issues/315",
      "https://github.com/antigenomics/vdjdb-db/issues/323",
      "PMID:33951417",
      "PMID:34086357",
      "structure",
      "multi",
      "multi-structure"),
      study.name = c(
        "Agerer et al (2021)", 
        "Thomas lab (2021)",
        "Efimov lab (2020)",
        "Sewell lab (2021)",
        "Kedzierska lab (2021)",
        "Rowntree et al (2021)",
        "Structure",
        "Multi-lab",
        "Multi-lab-structure"),
      stringsAsFactors=F),
    by = "study"
  )


p7 <- df.cc.sel.cov %>%
  filter(antigen.epitope == "YLQPRTFLL", 
         !(study %in% c("PMID:33664060", "multi", "structure", "multi-structure"))) %>%
  ggplot(aes(x = x, y = y)) +
  #geom_density_2d_filled(color = "white") +
  geom_density_2d(color = "grey") +
  geom_point(color = "white", size = 2) +
  geom_point(aes(color = study.name), size = 1.5) +
  geom_point(data = df.cc.sel.cov %>%
               filter(antigen.epitope == "YLQPRTFLL", startsWith(study, "multi")),
             color = "white", size = 3) +
  geom_point(data = df.cc.sel.cov %>%
               filter(antigen.epitope == "YLQPRTFLL", startsWith(study, "multi")),
             color = "red", size = 2.5) +
  geom_label_repel(data = df.cc.sel.cov %>%
                    filter(antigen.epitope == "YLQPRTFLL", endsWith(study, "structure")),
                  aes(label = cdr3aa),
                  min.segment.length = 0, 
                  nudge_x = 200,
                  nudge_y = 200,
                  alpha = 0.9,
                  color = "black",
                  arrow = arrow(length = unit(0.05, "inches"), 
                                type = "closed")) +
  xlab("") + ylab("") +
  scale_color_brewer("", palette = "Set2") +
  scale_fill_brewer(guide=F, palette = "Greys") +
  facet_wrap( ~ gene, scales = "free") +
  theme_minimal() +
  theme(aspect = 1,
        legend.position = "bottom",
        axis.text = element_blank(),
        panel.grid.major = element_blank(),
        panel.grid.minor = element_blank()) 

p7

pdf("fig2a.pdf", width = 8, height = 6)
p7
dev.off()
```

```{r fig.height=3.5, fig.width=8}
p8 <- plot_grid(plotlist = df.cc.sel.cov %>%
                  filter(gene == "TRA", antigen.epitope == "YLQPRTFLL") %>%
                  group_by(study.name, cid) %>%
                  filter(n() >= 10) %>%
                  group_map(~ggseqlogo(.x$cdr3aa) +
                              ggtitle(.y) + 
                              theme(legend.position = "none",
                                    axis.title.y = element_blank(),
                                    axis.text.y = element_blank())),
                ncol = 4)
p8

pdf("fig2b.pdf", width = 8, height = 3.5)
p8
dev.off()

p9 <- plot_grid(plotlist = df.cc.sel.cov %>%
                  filter(gene == "TRB", antigen.epitope == "YLQPRTFLL") %>%
                  group_by(study.name, cid) %>%
                  filter(n() >= 10) %>%
                  group_map(~ggseqlogo(.x$cdr3aa) +
                              ggtitle(.y) + 
                              theme(legend.position = "none",
                                    axis.title.y = element_blank(),
                                    axis.text.y = element_blank())),
                ncol = 4)
p9

pdf("fig2c.pdf", width = 8, height = 3.5)
p9
dev.off()
```

###

```{r}
#
```