Plots and models of variant effect sizes

# Univariate analysis, mashr-adjusted, data-driven approach
univariate_lmm_results <- tbl(db, "univariate_lmm_results") %>% 
  select(-contains("canonical"), 
         -contains("raw")) %>% 
  inner_join(tbl(db, "variants") %>% 
               select(SNP, FBID, site.class, distance.to.gene, MAF), 
             by = "SNP") %>%
  left_join(
    tbl(db, "genes") %>% 
      select(FBID, gene_name), by = "FBID") %>%
  collect(n = Inf) %>%
  rename_all(~ gsub("beta_", "", .x)) %>%
  rename_all(~ gsub("_mashr_ED", "", .x))

univariate_lmm_results <- univariate_lmm_results %>%
  mutate(site.class = gsub("5_", "5-", site.class),
         site.class = gsub("3_", "3-", site.class),
         site.class = gsub("NON_", "NON-", site.class),
         site.class = gsub("_", " ", site.class),
         site.class = capitalize(tolower(site.class)),
         site.class = gsub("Utr", "UTR", site.class)
  )

# Results of one multivariate linear mixed model run in GEMMA:
# multivariate_lmm_results <- 
#   read_tsv("data/derived/output/all_four_traits.assoc.txt") %>%
#   filter(p_fdr < 0.05) %>% 
#   inner_join(tbl(db, "variants") %>% 
#                select(-chr, -position) %>%
#                collect(), by = "SNP") %>% 
#   arrange(p_wald) %>% 
#   select(-contains("Vbeta")) %>%
#   left_join(
#     tbl(db, "genes") %>% 
#       select(FBID, gene_name) %>% 
#       collect(n=Inf), by = "FBID") %>%
#   mutate(gene_name = replace(gene_name, is.na(FBID), NA)) 
# 
# multivariate_lmm_results <- multivariate_lmm_results[,c(1:8, 14, 9:13)]
# 
# multivariate_lmm_results <- multivariate_lmm_results %>%
#   mutate_at(vars(starts_with("p_")), ~ -log10(.x)) %>%
#   mutate_if(is.numeric, ~ format(round(.x, 2), nsmall = 2)) %>%
#   left_join(tbl(db, "univariate_lmm_results") %>% 
#               select(SNP, SNP_clump) %>% 
#               collect(),
#             by = "SNP") %>%
#   distinct(SNP_clump, .keep_all = TRUE) %>%
#   select(SNP, SNP_clump, everything())

p1 <- SNP_clumps %>%
  filter(LD_subset) %>% # The figure looks the same whether or not the data are thinned to the LD set of SNPs
  arrange(beta_female_early_raw) %>%
  mutate(bin = c(rep(1:floor(n()/1000), each = 1000), 
                 rep(floor(n()/1000) + 1, each = n() %% 1000)),
         age = "A. Early life") %>%
  group_by(bin, age) %>%
  summarise(females = mean(beta_female_early_raw), males = mean(beta_male_early_raw)) %>%
  ggplot(aes(females, males)) +
  geom_hline(yintercept = 0, linetype = 2) +
  geom_vline(xintercept = 0, linetype = 2) +
  geom_point() +
  stat_smooth(method = "lm", formula = y ~ x + I(x^2), size = 0.6) +
  facet_wrap(~ age) +
  scale_y_continuous(limits = c(-0.1, 0.2)) +
  xlab(NULL) + ylab(NULL) + 
  theme_bw() +
  theme(strip.background = element_blank(),
        strip.text = element_text(hjust=0))   + 
  theme(text = element_text(family = "Raleway", size = 12))

p2 <- SNP_clumps %>%
  filter(LD_subset) %>% 
  arrange(beta_female_late_raw) %>%
  mutate(bin = c(rep(1:floor(n()/1000), each = 1000), 
                 rep(floor(n()/1000) + 1, each = n() %% 1000)),
         age = "B. Late life") %>%
  group_by(bin, age) %>%
  summarise(females = mean(beta_female_late_raw), males = mean(beta_male_late_raw)) %>%
  ggplot(aes(females, males)) +
  geom_hline(yintercept = 0, linetype = 2) +
  geom_vline(xintercept = 0, linetype = 2) +
  geom_point() +
  stat_smooth(method = "lm", formula = y ~ x + I(x^2), size = 0.6) +
  facet_wrap(~ age) +
  scale_y_continuous(limits = c(-0.1, 0.2)) +
  xlab(NULL) + ylab(NULL) + 
  theme_bw() +
  theme(strip.background = element_blank(),
        strip.text = element_text(hjust=0))   + 
  theme(text = element_text(family = "Raleway", size = 12))

grid.arrange(p1, p2, nrow = 1, bottom = "Mean effect size on female fitness", 
  left = "Mean effect size on male fitness")

# Make plot for the GWAS
mashr_results_canonical <- read_rds("data/derived/mashr_results_canonical.rds")
mashr_2L <- readRDS("data/derived/mashr_results_canonical_chr2L.rds")
mashr_2R <- readRDS("data/derived/mashr_results_canonical_chr2R.rds")
mashr_3L <- readRDS("data/derived/mashr_results_canonical_chr3L.rds")
mashr_3R <- readRDS("data/derived/mashr_results_canonical_chr3R.rds")
mashr_X <- readRDS("data/derived/mashr_results_canonical_chrX.rds")

mix <- bind_rows(
  enframe(sort(get_estimated_pi(mashr_results_canonical))) %>%
    mutate(Chromosome = "All"),
  enframe(sort(get_estimated_pi(mashr_2L))) %>%
    mutate(Chromosome = "2L"),
  enframe(sort(get_estimated_pi(mashr_2R))) %>%
    mutate(Chromosome = "2R"),
  enframe(sort(get_estimated_pi(mashr_3L))) %>%
    mutate(Chromosome = "3L"),
  enframe(sort(get_estimated_pi(mashr_3R))) %>%
    mutate(Chromosome = "3R"),
  enframe(sort(get_estimated_pi(mashr_X))) %>%
    mutate(Chromosome = "X")) %>% 
  rename(Mixture_component = name) 

to_keep <- mix %>%
  group_by(Mixture_component) %>%
  summarise(value = max(value), .groups = "drop") %>%
  filter(value > 0.01) %>%
  pull(Mixture_component)

mix <- mix %>%
  filter(Mixture_component %in% to_keep) %>%
  spread(Mixture_component, value) %>%
  rename(`Sexually concordant effect` = equal_effects,
         `Female-specific effect` = Female_specific_1,
         `Male-specific effect` = Male_specific_1,
         `Sexually antagonistic effect` = Sex_antag_0.25,
         `No effect on fitness` = null) %>%
  gather(Mixture_component, value, -Chromosome) %>%
  arrange(-value) 

chr_levels <- mix %>% 
  filter(Mixture_component == "Sexually antagonistic effect") %>%
  arrange(value) %>% pull(Chromosome)

mix <- mix %>%
  mutate(Chromosome = factor(Chromosome, chr_levels),
         Mixture_component = factor(Mixture_component, 
                                    c("Sexually antagonistic effect",
                                      "Sexually concordant effect",
                                      "Female-specific effect",
                                      "Male-specific effect",
                                      "No effect on fitness"))) 


p1 <- mix %>%
  ggplot(aes(Chromosome, 100 * value)) + 
  geom_bar(stat = "identity",aes(fill = Chromosome), colour = "grey10",  size = 0.3) + 
  scale_fill_brewer(palette = "Spectral", direction = -1) + 
  coord_flip() + 
  theme_bw() + 
  scale_y_continuous(expand = c(0,0), limits = c(0, 70)) + 
  scale_x_discrete(expand = c(0.14, 0.14)) + 
  theme(axis.ticks.y = element_blank(),
        legend.position = "none",
        panel.border = element_rect(size = 0.8),
        text = element_text(family = "Raleway", size = 12),
        strip.background = element_blank(),
        panel.grid.major.y = element_blank()) +
  ylab("Estimated % of loci") +
  facet_wrap(~ Mixture_component, ncol = 1)


# Make plot for the TWAS
mashr_results_canonical <- readRDS("data/derived/TWAS/TWAS_canonical.rds")
mashr_2L <- readRDS("data/derived/TWAS/TWAS_canonical_2L.rds")
mashr_2R <- readRDS("data/derived/TWAS/TWAS_canonical_2R.rds")
mashr_3L <- readRDS("data/derived/TWAS/TWAS_canonical_3L.rds")
mashr_3R <- readRDS("data/derived/TWAS/TWAS_canonical_3R.rds")
mashr_X <- readRDS("data/derived/TWAS/TWAS_canonical_X.rds")

mix <- bind_rows(
  enframe(sort(get_estimated_pi(mashr_results_canonical))) %>%
    mutate(Chromosome = "All"),
  enframe(sort(get_estimated_pi(mashr_2L))) %>%
    mutate(Chromosome = "2L"),
  enframe(sort(get_estimated_pi(mashr_2R))) %>%
    mutate(Chromosome = "2R"),
  enframe(sort(get_estimated_pi(mashr_3L))) %>%
    mutate(Chromosome = "3L"),
  enframe(sort(get_estimated_pi(mashr_3R))) %>%
    mutate(Chromosome = "3R"),
  enframe(sort(get_estimated_pi(mashr_X))) %>%
    mutate(Chromosome = "X")) %>% 
  rename(Mixture_component = name) %>%
  mutate(Mixture_component = str_remove_all(Mixture_component, "_0.25"),
         Mixture_component = str_remove_all(Mixture_component, "_0.5"),
         Mixture_component = str_remove_all(Mixture_component, "_0.75"),
         Mixture_component = str_remove_all(Mixture_component, "_1.0")) %>%
  group_by(Mixture_component, Chromosome) %>%
  summarise(value = sum(value), .groups = "drop")

to_keep <- mix %>%
  group_by(Mixture_component) %>%
  summarise(value = max(value), .groups = "drop") %>%
  filter(value > 0.01) %>%
  pull(Mixture_component)

mix <- mix %>%
  filter(Mixture_component %in% to_keep) %>%
  spread(Mixture_component, value) %>%
  rename(`Sexually concordant effect` = equal_effects,
         `Female-specific effect` = Female_specific_1,
         `Male-specific effect` = Male_specific_1,
         `Sexually antagonistic effect` = Sex_antag,
         `No effect on fitness` = null) %>%
  gather(Mixture_component, value, -Chromosome) %>%
  arrange(-value) 

chr_levels <- rev(c("X", "3L", "3R", "All", "2R", "2L"))

mix <- mix %>%
  mutate(Chromosome = factor(Chromosome, chr_levels),
         Mixture_component = factor(Mixture_component, 
                                    c("Sexually antagonistic effect",
                                      "Sexually concordant effect",
                                      "Female-specific effect",
                                      "Male-specific effect",
                                      "No effect on fitness"))) 

p2 <- mix %>%
  ggplot(aes(Chromosome, 100 * value)) + 
  geom_bar(stat = "identity",aes(fill = Chromosome), colour = "grey10",  size = 0.3) + 
  scale_fill_brewer(palette = "Spectral", direction = -1) + 
  coord_flip() + 
  theme_bw() + 
  scale_y_continuous(expand = c(0,0), limits = c(0, 65)) + # 
  scale_x_discrete(expand = c(0.14, 0.14)) + 
  theme(axis.ticks.y = element_blank(),
        legend.position = "none",
        panel.border = element_rect(size = 0.8),
        text = element_text(family = "Raleway", size = 12),
        strip.background = element_blank(),
        panel.grid.major.y = element_blank()) +
  ylab("Estimated % of transcripts") +
  xlab(" ") + 
  facet_wrap(~ Mixture_component, ncol = 1)


# Save composite figure of the GWAS and TWAS mixture proportions
ggsave(filename = "figures/composite_mixture_figure.pdf", 
       cowplot::plot_grid(p1, p2, labels = c('A', 'B'), label_size = 12),
       width = 5, height = 8)

cowplot::plot_grid(p1, p2, labels = c('A', 'B'), label_size = 12)

histo_data <- tbl(db, "univariate_lmm_results") %>% 
  left_join(tbl(db, "variants") %>% 
               select(SNP, MAF, site.class), 
             by = "SNP") %>%
  filter(!is.na(LFSR_female_early_mashr_ED)) %>% 
  collect(n = Inf) %>% 
  mutate(class = 
           case_when(
             P_sex_antag > quantile(P_sex_antag, probs = 0.999)  ~ "Top 0.1%\nsexually antagonistic loci",
             P_equal_effects > quantile(P_equal_effects, probs = 0.999)  ~ "Top 0.1%\nsexually concordant loci",
             P_female_specific > quantile(P_female_specific, probs = 0.999)  ~ "Top 0.1%\nfemale-specific loci",
             P_male_specific > quantile(P_male_specific, probs = 0.999)  ~ "Top 0.1%\nmale-specific loci",
           )) %>%
  mutate(MAF = ifelse(MAF < 0.2, "0.05 < MAF < 0.2", "0.2 < MAF < 0.5")) %>%
  filter(!is.na(class)) %>%
  select(beta_female_early_mashr_ED, beta_male_early_mashr_ED, class, MAF)

eff_size_histos <- ggplot(histo_data, aes(beta_female_early_mashr_ED)) + 
  geom_vline(xintercept = 0, linetype = 2) +
  geom_histogram(bins=30, fill="tomato", colour=NA, alpha = 0.5) + 
  geom_histogram(aes(beta_male_early_mashr_ED), bins=30, fill="steelblue", colour=NA, alpha = 0.5) + 
  geom_histogram(bins=30, fill=NA, colour="black") + 
  geom_histogram(aes(beta_male_early_mashr_ED), bins=30,fill=NA,colour="black") + 
  facet_grid(MAF ~ class, scales = "free_y") +
  xlab("Effect of the minor allele on early-life fitness\n(females:red, males:blue)") + 
  theme_bw() +
  theme(legend.position = "none",
        strip.background = element_blank(), 
        text = element_text(family = "Raleway", size = 12)) +
  ylab("Number of loci")

ggsave("figures/eff_size_histos.pdf", width = 8, height = 5)
eff_size_histos

manhattan_data <- tbl(db, "univariate_lmm_results") %>% 
  select(SNP, starts_with("P_"),
         pvalue_female_early_raw, pvalue_male_early_raw) %>%
  distinct() %>%
  collect(n=Inf) %>%
  mutate(position = str_split(SNP, "_"),
         chr = map_chr(position, ~ .x[1]),
         position = as.numeric(map_chr(position, ~ .x[2]))) %>% 
         # P_pleiotropy = 1 - P_null - P_female_specific - P_male_specific,
         # SA_not_concord = P_sex_antag > P_equal_effects,
         # top_SA = P_sex_antag >= quantile(P_sex_antag, probs=0.9999)) %>% 
  filter(chr != "4") 

max_pos <- manhattan_data %>%
  group_by(chr) %>%
  summarise(max_pos = max(position), .groups = "drop") %>%
  as.data.frame()
max_pos$max_pos <- c(0, cumsum(max_pos$max_pos[1:4]))

manhattan_data <- manhattan_data %>%
    left_join(max_pos, by = "chr") %>%
    mutate(position = position + max_pos) 


p1 <- manhattan_data %>%
  ggplot(aes(position, -1 * log10(pvalue_female_early_raw), group = chr, colour = chr, stroke = 0.2)) + 
  geom_point(size = 0.5) +
  # geom_point(data = manhattan_data %>% filter(top_SA), 
  #            aes(y = 8), size = 3, pch = 6) +
  scale_colour_brewer(palette = "Paired", name = "Chromosome") + 
  scale_y_continuous(limits = c(0,8)) +
  ylab(expression(paste("Effect on female early-life fitness (-", Log[10], " p)"))) + 
  xlab("") +
  theme_bw() + 
  theme(axis.text.x = element_blank(),
        # text = element_text(family = "Raleway", size = 12),
        panel.border = element_blank(),
        axis.ticks.x = element_blank())

p2 <- manhattan_data %>%
  ggplot(aes(position, -1 * log10(pvalue_male_early_raw), group = chr, colour = chr, stroke = 0.2)) + 
  geom_point(size = 0.5) +
  scale_colour_brewer(palette = "Paired", name = "Chromosome") + 
  ylab(expression(paste("Effect on male early-life fitness (-", Log[10], " p)"))) + 
  xlab("Position") +
  scale_y_reverse(limits = c(8,0)) +
  theme_bw() + 
  theme(axis.text.x = element_blank(),
        # text = element_text(family = "Raleway", size = 12),
        panel.border = element_blank(),
        axis.ticks.x = element_blank())

grid_arrange_shared_legend <- function(..., ncol = length(list(...)), nrow = 1) {

  plots <- list(...)
  # position <- match.arg(position)
  g <- ggplotGrob(plots[[1]] + theme(legend.position = "right"))$grobs
  legend <- g[[which(sapply(g, function(x) x$name) == "guide-box")]]
  lheight <- sum(legend$height)
  lwidth <- sum(legend$width)
  gl <- lapply(plots, function(x) x + theme(legend.position="none"))
  gl <- c(gl, ncol = ncol, nrow = nrow)
  
  combined <- arrangeGrob(arrangeGrob(gl[[1]], gl[[2]], nrow = 2),
                          legend,
                          ncol = 2,
                          widths = unit.c(unit(1, "npc") - lwidth, lwidth))
  
  grid.newpage()
  grid.draw(combined)
  invisible(combined) # return gtable invisibly
}

grid_arrange_shared_legend(p1, p2)

make_evidence_ratio_plot <- function(dat, ymax, ylab){
  # Argument needs to be dataframe of loci or transcripts, which are identified by col called "identifier"
  # The FE, FL,ME, ML effect sizes need to be called pheno1-pheno4
  # The corresponding LFSR columns need to be called LFSR1-LFSR4
  
  antagonism_evidence_ratios <- dat %>%
    
    # Convert the LFSR to probability that beta is positive
    # Here, we use the ED results because they are agnostic to our expectations for the results
    # However, running this with the Canonocial results instead gives an essentially identical graph
    mutate(pp_female_early = ifelse(pheno1 > 0, LFSR1, 1 - LFSR1),
           pp_female_late  = ifelse(pheno2 > 0, LFSR2, 1 - LFSR2),
           pp_male_early   = ifelse(pheno3 > 0, LFSR3, 1 - LFSR3),
           pp_male_late    = ifelse(pheno4 > 0, LFSR4, 1 - LFSR4)) %>%
    
    # Calculate the probabilities that beta_i and beta_j have the same/opposite signs
    mutate(p_sex_concord_early = pp_female_early * pp_male_early + 
             (1 - pp_female_early) * (1 - pp_male_early),
           p_sex_antag_early = pp_female_early * (1 - pp_male_early) + 
             (1 - pp_female_early) * pp_male_early,
           p_sex_concord_late  = pp_female_late * pp_male_late + 
             (1 - pp_female_late) * (1 - pp_male_late),
           p_sex_antag_late = pp_female_late * (1 - pp_male_late) + 
             (1 - pp_female_late) * pp_male_late,
           p_age_concord_females = pp_female_early * pp_female_late + 
             (1 - pp_female_early) * (1 - pp_female_late),
           p_age_antag_females = pp_female_early * (1 - pp_female_late) + 
             (1 - pp_female_early) * pp_female_late,
           p_age_concord_males = pp_male_early * pp_male_late + (1 - pp_male_early) * (1 - pp_male_late),
           p_age_antag_males = pp_male_early * (1 - pp_male_late) + (1 - pp_male_early) * pp_male_late) %>%
    
    # Find the ratios of some of these probabilities (i.e. "evidence ratios")
    mutate(`Inter-sex (early life)` = p_sex_concord_early / p_sex_antag_early,
           `Inter-sex (late life)` = p_sex_concord_late / p_sex_antag_late,
           `Inter-age (females)` = p_age_concord_females / p_age_antag_females,
           `Inter-age (males)` = p_age_concord_males / p_age_antag_males) %>%
    select(identifier,  starts_with("Inter")) %>%
    gather(trait, evidence_ratio, -identifier) %>%
    mutate(trait = factor(trait, c("Inter-sex (early life)",
                                   "Inter-sex (late life)",
                                   "Inter-age (females)",
                                   "Inter-age (males)")))
  print(antagonism_evidence_ratios %>% arrange(evidence_ratio))
  antagonism_evidence_ratios %>%
    ggplot(aes(log2(evidence_ratio))) + 
    geom_histogram(data=subset(antagonism_evidence_ratios, evidence_ratio < 1), 
                   bins = 500, fill = "#FF635C") + 
    geom_histogram(data=subset(antagonism_evidence_ratios, evidence_ratio > 1), 
                   bins = 500, fill = "#5B8AFD") +
    coord_cartesian(xlim = c(-10, 10), ylim = c(0, ymax)) + 
    scale_x_continuous(breaks = c(-10,  -6,  -2,   2,   6,  10), 
                       labels = c(paste("1/",2 ^ c(10,  6,  2), sep = ""), 2 ^ c(2,6,10))) + 
    facet_wrap(~ trait) + 
    xlab("Evidence ratio (log2 scale)") + ylab(ylab) + 
    theme_bw() + 
    theme(panel.border = element_rect(size = 0.8),
         # text = element_text(family = "Raleway", size = 12),
          axis.text.x = element_text(angle = 45, vjust = 1, hjust=1),
          strip.background = element_blank())
}

GWAS_ratios_plot <- SNP_clumps %>%
  rename(pheno1 = beta_female_early_mashr_ED, 
         pheno2 = beta_female_late_mashr_ED, 
         pheno3 = beta_male_early_mashr_ED, 
         pheno4 = beta_male_late_mashr_ED,
         LFSR1 = LFSR_female_early_mashr_ED, 
         LFSR2 = LFSR_female_late_mashr_ED, 
         LFSR3 = LFSR_male_early_mashr_ED, 
         LFSR4 = LFSR_male_late_mashr_ED,
         identifier = SNP_clump) %>%
  make_evidence_ratio_plot(ymax = 10000, ylab = "Number of loci")

TWAS_ED <- readRDS("data/derived/TWAS/TWAS_ED.rds")
TWAS_canonical <- readRDS("data/derived/TWAS/TWAS_canonical.rds")


TWAS_mashr_results <- 
  data.frame(
    FBID = read_csv("data/derived/TWAS/TWAS_results.csv")$FBID,
    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())

TWAS_ratios_plot <- TWAS_mashr_results %>%
  rename(pheno1 = ED_beta_FE, 
         pheno2 = ED_beta_FL, 
         pheno3 = ED_beta_ME, 
         pheno4 = ED_beta_ML,
         LFSR1 = ED_LFSR_FE, 
         LFSR2 = ED_LFSR_FL, 
         LFSR3 = ED_LFSR_ME, 
         LFSR4 = ED_LFSR_ML,
         identifier = FBID) %>%
  make_evidence_ratio_plot(ymax = 500, ylab = "Number of transcripts")

pp <- plot_grid(GWAS_ratios_plot, TWAS_ratios_plot, 
                   labels = c('A', 'B'), label_size = 12)
# ggdraw(add_sub(pp, "Evidence ratio (log scale)", 
#                vpadding=grid::unit(0,"lines"), 
#                y=6.3, x=0.5, vjust=6))
ggsave("figures/antagonism_ratios.pdf", pp, width = 8.5, height = 4.9)

pp