read_elai.Rmd

---
title: "Elai final results"
author: "Tram"
date: "`r format(Sys.time(), '%d %B, %Y')`"
output: html_document
---

```{r setup, include=FALSE}
knitr::opts_chunk$set(echo = FALSE)
knitr::opts_chunk$set(eval = TRUE)
mylib <- "~/R/x86_64-pc-linux-gnu-library/4.0"
.libPaths(mylib)
# dir.create(mylib, showWarnings = F)
if (!requireNamespace("DT", quietly = TRUE)) install.packages('DT', repos = "https://cloud.r-project.org")
if (!requireNamespace("ggplot2", quietly = TRUE)) install.packages('ggplot2', repos = "https://cloud.r-project.org")
if (!requireNamespace("dplyr", quietly = TRUE)) install.packages('dplyr', repos = "https://cloud.r-project.org")
if (!requireNamespace("data.table", quietly = TRUE)) install.packages('data.table', repos = "https://cloud.r-project.org")
library(DT)
library(dplyr)
library(ggplot2)
library(data.table)
```

```{r functions}
read_elai <- function(elai_outdir) {
  adm_file <- file.path(elai_outdir, "elai_r.admix.txt")
  dosage_file <- file.path(elai_outdir, "elai_r.ps21.txt")
  snpout_file <- file.path(elai_outdir, "elai_r.snpinfo.txt")
  log_file <- file.path(elai_outdir, "elai_r.log.txt")

  log_cmd <- readLines(log_file, n = 3)
  log_cmd <- log_cmd[grep("COMMAND", log_cmd)]
  # log_cmd <- stringr::str_trim(gsub("## COMMAND: ", "", log_cmd))
  # log_cmd <- unlist(strsplit(log_cmd, " "))
  test_file <- stringr::str_split(log_cmd, " -p 1 ", simplify = T)[1]
  test_file <- stringr::str_split(test_file, " -g ", simplify = T)
  test_file <- test_file[length(test_file)]
  stopifnot(file.exists(test_file))

  sources <- stringr::str_split(log_cmd, " -p 1\\d+ ", simplify = T)
  sources <- sources[!grepl("-p 1", sources)]
  sources <- gsub(".*-g ", "", sources) # **
  stopifnot(all(file.exists(sources)))

  ancestries <- basename(sources)

  admixture <- data.table::fread(adm_file, header = F, strip.white = T, sep = " ", stringsAsFactors = F, data.table = F)
  colnames(admixture) <- ancestries
  test_info <- read.csv(test_file, skip = 2, nrows = 5, header = T)
  ind_names <- colnames(test_info)[-1]
  rownames(admixture) <- ind_names


  ## read ancestry dosage
  ## snp_out_file <- file.path(elai_outdir, "output/vn_var.snpinfo.txt")
  snp_out <- data.table::fread(snpout_file, header = T, stringsAsFactors = F, data.table = F)
  n_snp <- nrow(snp_out) # number of snps in snpinfo.txt
  n_ind <- ncol(test_info)-1 # number of test individuals in the test genotype file
  n_src <- length(sources) # number of source populations

  dosage <- scan(dosage_file)
  dim(dosage) <- c(n_src, n_snp, n_ind)

  # get dosage info for each individual in a list
  dosage <- sapply(ind_names, function(ind) {
    # get dosage matrix of the individual
    i <- which(ind_names == ind)
    ind_dosage <- as.data.frame(t(dosage[, , i]))
    # ind_dosage has [n. sources] columns and [n. snps] rows
    colnames(ind_dosage) <- ancestries
    ind_dosage$rs <- snp_out$rs
    # merge dosage table with snp info table
    ind_dosage <- merge(snp_out, ind_dosage, by = "rs")
    ind_dosage <- reshape2::melt(ind_dosage,
                                 measure.vars = ancestries,
                                 variable.name = "ancestry",
                                 value.name = "dosage")
    return(ind_dosage)
  }, USE.NAMES = T, simplify = F)

  return(list("oveall_admixture" = admixture, "ancestral_dosage" = dosage))
}
```

```{r elai_input}
elai_dir <- snakemake@params$workdir
source_files <- snakemake@input$source
source_param <- do.call(paste, 
                        list(lapply(1:length(source_files), function(f) {
                          id <- f-1
                          par <- paste0("-g ", source_files[f], " -p 1", id)
                          return(par)}),
                          collapse = " "))
test_file <- snakemake@input$test_file
pos_dir <- dirname(snakemake@input$snp_file)
pos_files <- list.files(pos_dir, "*.pos", full.names = T)
upper_layer <- length(source_files)
opts <- snakemake@params$options
mem <- snakemake@resources$mem_gb

elai <- "/data3/projects/vietcaf/baotram/scripts/robusta_vn/elai/elai-lin"

all_params <- data.frame(source = source_param,
                         test = test_file,
                         pos = pos_files,
                         upper_layer = upper_layer,
                         other_params = opts,
                         elai_dir = elai_dir,
                         elai = elai)
```

```{r run_elai}
run_elai <- function(source, test, pos, upper_layer, other_params, elai_dir, elai) {
  ini_dir <- getwd()
  batch <- gsub("test_|.pos", "", basename(pos))
  workdir <- file.path(elai_dir, batch)
  setwd(workdir)
  cmd <- paste(elai, source_params, "-g", test_file, "-p 1", "-pos", pos_file, "-o elai_r", "-C", upper_layer, opts)
  system(cmd)
  outdir <- file.path(el)
  file.copy("./output", workdir, recursive = T, overwrite = T)
  file.remove("./output")
  setwd(ini_dir)
  return(workdir)
}

parse_jobname <- function(pos_file) {
  batch <- gsub("test_|.pos", "", basename(pos_file))
  snp_selection <- basename(dirname(pos_file))
  chromosome <- basename(dirname(dirname(pos_file)))
  jobname <- paste("elai", chromosome, snp_selection, batch, sep = "_")
  return(jobname)
}

logname <- snakemake@params$logname
logdir <- snakemake@params$logdir
logfile <- file.path(logdir, "slurm", paste0(logname, "_%j.log"))

submit_elai <- lapply(1:nrow(all_params), function(i) {
  slurm_apply(run_elai, all_params[i,], submit = T,
              jobname = parse_jobname(all_params$pos[i]),
              slurm_options = list("output" = logfile, "error" = logfile, "mem" = paste0(mem,"G")))
})
```

```{r read_elai}
elai_results <- lapply(submit_elai, function(j)
  elai_dir <- get_slurm_out(j, outtype = 'table', wait = T)
  elai_dir <- as.character(elai_dir)
  elai_res <- read_elai(elai_dir)
  adm <- data.frame(elai_res$oveall_admixture, "batch" = basename(elai_dir), stringsAsFactors = F)
  adm$id <- rownames(adm)
  dosage <- elai_res$ancestral_dosage
  return(list("oveall_admixture" = adm, "ancestral_dosage" = dosage))
  )

# elai_results <- lapply(elai_dir, function(d) {
#   elai_res <- read_elai(d)
#   adm <- data.frame(elai_res$oveall_admixture, "chromosome" = basename(dirname(d)), stringsAsFactors = F)
#   dosage <- elai_res$ancestral_dosage
#   return(list("oveall_admixture" = adm, "ancestral_dosage" = dosage))
# })

outdir <- dirname(snakemake@output)
all_adm <- do.call(rbind, lapply(elai_results, function(x) {
  a <- x$oveall_admixture
  return(a)
}))
write.table(all_adm, file.path(outdir, "overall_admixture.txt"), sep = "\t", row.names = F, col.names = T)

all_dosage <- do.call(rbind, lapply(elai_results, function(x) {
  a <- x$ancestral_dosage
  d <- do.call(rbind, lapply(names(a), function(n) data.frame(a[[n]], "individual" = n)))
  return(d)
}))
write.table(all_dosage, file.path(outdir, "local_dosage.txt"), sep = "\t", row.names = F, col.names = T)
```

## elai parameters
```{r elai_params}
snakemake@params$elai_options
```

## overall admixture
```{r overall_adm}
DT::datatable(all_adm)
```

## local ancestry inference
```{r plot_lai, fig.height=8, fig.width=15}
# all_dosage$ancestry <- factor(all_dosage$ancestry, levels = c("group_ER", "group_OB", "group_C", "group_D", "group_AG"))
for (i in unique(all_dosage$individual)) {
  gp <- ggplot(all_dosage %>% filter(individual == i)) +
    geom_point(aes(x = pos, y = dosage, color = ancestry), size = 1.5, shape = 16, alpha = .5) +
    geom_line(aes(x = pos, y = dosage, color = ancestry), alpha = .5) +
    scale_color_viridis_d(end = 0.9) +
    xlab(paste("position on chromosome", unique(all_dosage$chr))) +
    ylab("ancestral dosage") +
    ggtitle(i)
  print(gp)
}
```


```{r, eval=FALSE}
true_inference_file <- snakemake@params$true_dosage_file
print_cor <- file.exists(true_inference_file)
```


```{r, eval=print_cor, eval=FALSE}
knitr::asis_output("## correlation between estimated inference and true inference\\n")
for (i in unique(all_dosage$individual)) {
  true_dosage <- data.table::fread(true_inference_file, header = T, sep = "\t", stringsAsFactors = F)
  cat(i, "\n")
  test_dosage <- merge(all_dosage %>% filter(individual == i), true_dosage, by.x = c("rs", "ancestry"), by.y = c("id", "ancestry"), all.x = T)
  test_dosage <- na.omit(test_dosage)
  test_dosage <- test_dosage[test_dosage$ancestry == "group_ER",]
  r <- cor(test_dosage$dosage, test_dosage$true_dosage, method = "pearson")
  cat("r = ", r, "\n")
}
```


## snps distribution
```{r snps_distribution, fig.height=8, fig.width=15}
# hist(unique(all_dosage$pos), breaks = 1000)
ggplot(all_dosage %>% dplyr::distinct(pos)) +
  geom_histogram(aes(x = pos), binwidth = 1e5) +
  ylab("count per 100kb")
```

```{r save_data, include=FALSE}
save(all_adm,all_dosage, file = file.path(outdir, "elai_results.RData"))
```