Transcriptome-wide association study (TWAS)

library(edgeR)
library(tidyverse)
library(glue)
library(future)
library(future.apply)
library(parallel)
library(kableExtra)
library(DT)
library(ashr)
library(mashr)

options(stringsAsFactors = FALSE)

# Connect to the database of annotations
db <- DBI::dbConnect(RSQLite::SQLite(), "data/derived/annotations.sqlite3")

# Helper to run shell commands
run_command <- function(shell_command, wd = getwd(), path = ""){
  cat(system(glue("cd ", wd, path, "\n",shell_command), intern = TRUE), sep = '\n')
}

kable_table <- function(df) { # cool tables
  kable(df, "html") %>%
  kable_styling() %>%
  scroll_box(height = "300px")
}

my_data_table <- function(df){ # Make html tables:
  datatable(
    df, rownames=FALSE,
    autoHideNavigation = TRUE,
    extensions = c("Scroller",  "Buttons"),
    options = list(
      dom = 'Bfrtip',
      deferRender=TRUE,
      scrollX=TRUE, scrollY=400,
      scrollCollapse=TRUE,
      buttons = 
        list('pageLength', 'colvis', 'csv', list(
          extend = 'pdf',
          pageSize = 'A4',
          orientation = 'landscape',
          filename = 'TWAS_enrichment')),
       columnDefs = list(list(targets = c(8,10), visible = FALSE)),
      pageLength = 50
    )
  )
}

# Helper to load Huang et al.'s data
load_expression_data <- function(sex, both_sex_chromosomes = TRUE){
  
  # Note: Huang et al's data contains weird stuff like supposedly Y-linked genes that have
  # higher/equal expression in *females* in all lines, presumably microarray issues/errors.
  # To be conservative, we restrict our analyses to genes that are known to be on a
  # chromosomes that is present in both sexes
  if(both_sex_chromosomes){
    genes_allowed <- tbl(db, "genes") %>%
      filter(chromosome %in% c("2L", "2R", "3L", "3R", "4", "X")) %>%
      pull(FBID)
  } else {
    genes_allowed <- tbl(db, "genes") %>%
      pull(FBID)
  }
  
  if(sex != "both"){
    expression <- glue("data/input/huang_transcriptome/dgrp.array.exp.{sex}.txt") %>% 
      read_delim(delim = " ") %>%
      filter(gene %in% genes_allowed)
    sample_names <- names(expression)[names(expression) != "gene"] %>% str_remove(":[12]") 
    gene_names <- expression$gene
    expression <- expression %>% select(-gene) %>% as.matrix() %>% t() 
    rownames(expression) <- sample_names # rows are samples, columns are genes
    colnames(expression) <- gene_names 
    return(expression %>% as.data.frame() %>% 
             tibble::rownames_to_column("line") %>% 
             as_tibble() %>%
             mutate(line = str_remove_all(line, "[.]1")))
  }
  
  females <- read_delim("data/input/huang_transcriptome/dgrp.array.exp.female.txt", delim = " ") %>%
    filter(gene %in% genes_allowed)
  names(females)[-1] <- paste("F_", names(females)[-1],sep="") #%>% str_remove(":[12]") 
  females <- females %>% 
    left_join(read_delim("data/input/huang_transcriptome/dgrp.array.exp.male.txt", delim = " "), by = "gene")
  sample_names <- names(females)[names(females) != "gene"] %>% str_remove(":[12]") 
  gene_names <- females$gene
  sex <- ifelse(str_detect(sample_names, "F_"), "female", "male")
  line <- str_remove_all(sample_names, "F_")
  females <- females %>% select(-gene) %>% t()
  colnames(females) <- gene_names
  list(
    sampleIDs = tibble(sex, line),
    expression = females
  )
}

# Table S5: heritability of the expression level of each transcript (as measured in males, or females)
huang_heritability <- read_csv("data/input/huang_2015_tableS5_transcript_heritability.csv")

# Load the predicted line means from present study, as calculated in get_predicted_line_means.Rmd
predicted_line_means <- read_csv("data/derived/predicted_line_means.csv")

transcript_selection_analysis <- function(expression_data, phenotypes){
  
  if("block" %in% names(phenotypes)) phenotypes <- phenotypes %>% select(-block)
  
  expression_data <- expression_data %>% 
    filter(line %in% phenotypes$line)
  
  # Find line mean expression for each gene (average across the c. 2 replicate samples per line) 
  chunk_cols <- split(2:ncol(expression_data), ceiling(seq_along(2:ncol(expression_data)) / 500))
  
  mean_expression_data <- mclapply(1:2, function(i){  
    expression_data[, c(1, chunk_cols[[i]])] %>% 
      group_by(line) %>% 
      summarise_all(mean) %>% 
      ungroup()
  }) %>% bind_rows()

  # Scale each transcript's expression level so that mean is 0, and the variance is 1, across all the lines measured by Huang et al.
  for(i in 2:ncol(mean_expression_data)) mean_expression_data[,i] <- as.numeric(scale(mean_expression_data[,i]))
  
  # Join the microarray data with the phenotypes (i.e. our fitness data), and keep only the lines where we have both sets of measurements
  expression_data <- phenotypes %>% left_join(expression_data, by = "line")
  expression_data <- expression_data[complete.cases(expression_data), ] %>% select(-line)
  
  print("Data ready for analysis. Starting TWAS...")

  # Create chunks of transcript names, which will be used to facilitate parallel processing
  transcripts <- names(expression_data)[-c(1:4)]
  transcripts <- split(transcripts, ceiling(seq_along(transcripts) / 100))
  
  # Define a function to run 4 linear models, and get the beta and SE for regressions of expression level on the 4 fitness traits
  do_one_transcript <- function(transcript){
    expression_level <- expression_data %>% pull(!!transcript)
    FE <- summary(lm(female.fitness.early ~ expression_level, data = expression_data))$coefficients
    FL <- summary(lm(female.fitness.late ~ expression_level, data = expression_data))$coefficients
    ME <- summary(lm(male.fitness.early ~ expression_level, data = expression_data))$coefficients
    ML <- summary(lm(male.fitness.late ~ expression_level, data = expression_data))$coefficients
    
    c(FE[2,1], FL[2,1], ME[2,1], ML[2,1], # effect size
      FE[2,2], FL[2,2], ME[2,2], ML[2,2], # SE
      FE[2,4], FL[2,4], ME[2,4], ML[2,4]) # p-value
  }
  
  # Runs do_one_transcript() on all the transcripts listed in the vector 'transcripts'
  do_chunk_of_transcripts <- function(transcripts){
    output <- data.frame(transcripts, lapply(transcripts, do_one_transcript) %>% do.call("rbind", .))
    names(output) <- c("gene", "beta_FE", "beta_FL", "beta_ME", "beta_ML", 
                       "SE_FE", "SE_FL", "SE_ME", "SE_ML",
                       "pval_FE", "pval_FL", "pval_ME", "pval_ML")
    output
  }
  
  # Run it all, in parallel
  transcripts %>% 
    mclapply(do_chunk_of_transcripts) %>% 
    do.call("rbind", .) %>% as_tibble() %>% mutate(gene = as.character(gene))
}

if(!file.exists("data/derived/TWAS/TWAS_result_males.csv")){

  TWAS_result_females <- load_expression_data("female") %>% transcript_selection_analysis(predicted_line_means) 
  TWAS_result_females %>% write_csv("data/derived/TWAS/TWAS_result_females.csv")
  TWAS_result_males <- load_expression_data("male") %>% transcript_selection_analysis(predicted_line_means)
  TWAS_result_males %>% write_csv("data/derived/TWAS/TWAS_result_males.csv")
} else {
  TWAS_result_females <- read_csv("data/derived/TWAS/TWAS_result_females.csv")
  TWAS_result_males <- read_csv("data/derived/TWAS/TWAS_result_males.csv")
}

if(!file.exists("data/derived/TWAS/TWAS_ED.rds")){
  
  input_data <- data.frame(TWAS_result_females[,1:3], 
                           TWAS_result_males[,4:5], 
                           TWAS_result_females[,6:7], 
                           TWAS_result_males[,8:9])
  
  with_chrs <- input_data %>%
    left_join(tbl(db, "genes") %>% 
                select(FBID, chromosome) %>% 
                collect(), by = c("gene" = "FBID"))
  
  TWAS_canonical <- input_data %>% 
    run_mashr(mashr_mode = "canonical") 
  
  TWAS_canonical_2L <- with_chrs %>% 
    filter(chromosome == "2L") %>% select(-chromosome) %>%
    run_mashr(mashr_mode = "canonical") 
  
  TWAS_canonical_2R <- with_chrs %>% 
    filter(chromosome == "2R") %>% select(-chromosome) %>%
    run_mashr(mashr_mode = "canonical") 
  
  TWAS_canonical_3L <- with_chrs %>% 
    filter(chromosome == "3L") %>% select(-chromosome) %>%
    run_mashr(mashr_mode = "canonical") 
  
  TWAS_canonical_3R <- with_chrs %>% 
    filter(chromosome == "3R") %>% select(-chromosome) %>%
    run_mashr(mashr_mode = "canonical") 
  
  TWAS_canonical_X <- with_chrs %>% 
    filter(chromosome == "X") %>% select(-chromosome) %>%
    run_mashr(mashr_mode = "canonical") 
  
  saveRDS(TWAS_canonical, "data/derived/TWAS/TWAS_canonical.rds")
  saveRDS(TWAS_canonical_2L, "data/derived/TWAS/TWAS_canonical_2L.rds")
  saveRDS(TWAS_canonical_2R, "data/derived/TWAS/TWAS_canonical_2R.rds")
  saveRDS(TWAS_canonical_3L, "data/derived/TWAS/TWAS_canonical_3L.rds")
  saveRDS(TWAS_canonical_3R, "data/derived/TWAS/TWAS_canonical_3R.rds")
  saveRDS(TWAS_canonical_X, "data/derived/TWAS/TWAS_canonical_X.rds")
  
  TWAS_ED <- input_data %>% 
    run_mashr(mashr_mode = "ED", 
              ED_p_cutoff = 0.4) 
  saveRDS(TWAS_ED, "data/derived/TWAS/TWAS_ED.rds")
} else {
  TWAS_canonical <- readRDS("data/derived/TWAS/TWAS_canonical.rds")
  TWAS_ED <- readRDS("data/derived/TWAS/TWAS_ED.rds")
}

TWAS_mashr_results <- 
  data.frame(
    FBID = colnames(load_expression_data("female"))[-1],
    as.data.frame(get_pm(TWAS_canonical)) %>% rename_all(~str_c("canonical_", .)),
    as.data.frame(get_lfsr(TWAS_canonical)) %>% 
      rename_all(~str_replace_all(., "beta", "LFSR")) %>% rename_all(~str_c("canonical_", .))) %>% 
  as_tibble() %>% 
  bind_cols(
    data.frame(
      as.data.frame(get_pm(TWAS_ED)) %>% rename_all(~str_c("ED_", .)),
      as.data.frame(get_lfsr(TWAS_ED)) %>% 
        rename_all(~str_replace_all(., "beta", "LFSR")) %>% rename_all(~str_c("ED_", .))) %>% 
      as_tibble()) 

# Get the mixture weights, as advised by mashr authors here: https://github.com/stephenslab/mashr/issues/68
posterior_weights_cov <- TWAS_canonical$posterior_weights 
colnames(posterior_weights_cov) <- sapply(
  str_split(colnames(posterior_weights_cov), '\\.'), 
  function(x) {
    if(length(x) == 1) return(x)
    else if(length(x) == 2) return(x[1])
    else if(length(x) == 3) return(paste(x[1], x[2], sep = "."))
  })
posterior_weights_cov <- t(rowsum(t(posterior_weights_cov), 
                                  colnames(posterior_weights_cov)))

# Make a neat dataframe
TWAS_mixture_assignment_probabilities <- data.frame(
  FBID = TWAS_mashr_results$FBID,
  posterior_weights_cov,
  stringsAsFactors = FALSE
) %>% as_tibble() %>%
  rename(P_equal_effects = equal_effects,
         P_female_specific = Female_specific_1,
         P_male_specific = Male_specific_1,
         P_null = null,
         P_sex_antag = Sex_antag_0.25)

TWAS_mixture_assignment_probabilities <- TWAS_mixture_assignment_probabilities %>%
  left_join(huang_heritability, by = "FBID")

saveRDS(TWAS_mixture_assignment_probabilities, 
        "data/derived/TWAS/TWAS_mixture_assignment_probabilities.rds")