#!/usr/bin/env Rscript
suppressPackageStartupMessages({
  library(argparser, quietly = TRUE)
  library(pheatmap)
  library(tidyverse)
})
parser <- arg_parser("draw heatmap and gene expression figure for listed genes")
parser <- add_argument(parser, "expression", help = "input: gene expression file, TPM or FPKM", type = "character")
parser <- add_argument(parser, "genelist", help = "input: gene list for heatmap", type = "character")
parser <- add_argument(parser, "sampleinfo", help = "input: sample and replicate relationships", type = "character")
parser <- add_argument(parser, "output", help = "output prefix", type = "character")
argv <- parse_args(parser)
exp <- read.table(argv$expression, header = TRUE, row.names = 1, stringsAsFactors = FALSE, check.names = FALSE)
gene_id <- read.table(argv$genelist, header = FALSE, stringsAsFactors = FALSE)
sample <- read.table(argv$sampleinfo, header = FALSE, row.names = 1, stringsAsFactors = FALSE, check.names = FALSE)
keep_genes <- gene_id$V1[gene_id$V1 %in% rownames(exp)]
exp_sub <- exp[keep_genes, , drop = FALSE]
exp_sub <- exp_sub[apply(exp_sub, 1, sum, na.rm = TRUE) >= 1, , drop = FALSE]
pdf(paste(argv$output, "gene-exp.heatmap.pdf", sep = "."))
pheatmap(log2(exp_sub + 1), color = colorRampPalette(c("blue", "white", "red"))(100))
dev.off()
dat_cor <- cor(t(exp_sub), method = "pearson")
pdf(paste(argv$output, "gene-cor.heatmap.pdf", sep = "."), width = 10, height = 10)
pheatmap(dat_cor, color = colorRampPalette(c("blue", "white", "red"))(100), cluster_rows = FALSE, cluster_cols = FALSE, display_numbers = TRUE)
dev.off()
write.table(dat_cor, file = paste(argv$output, "gene-cor.txt", sep = "."), quote = FALSE, row.names = FALSE, sep = "	")
sample_tbl <- rownames_to_column(sample, var = "sample")
exp_tbl <- rownames_to_column(exp_sub, var = "geneID")
write.table(exp_tbl, file = paste(argv$output, "gene-exp.txt", sep = "."), quote = FALSE, row.names = FALSE, sep = "	")
exp_long <- pivot_longer(exp_tbl, -geneID, names_to = "sample", values_to = "exp")
exp_data <- left_join(exp_long, sample_tbl, by = "sample")
exp_stat <- exp_data %>% group_by(geneID, V2) %>% summarise(mean = mean(exp), sd = sd(exp, na.rm = TRUE), .groups = "drop") %>% mutate(up = mean + sd, down = if_else(mean < sd, 0, mean - sd))
pdf(paste(argv$output, "gene-exp.hist.pdf", sep = "."))
ggplot(exp_stat, aes(x = V2, y = mean, fill = V2)) + geom_bar(position = position_dodge(), stat = "identity", width = 0.5) + geom_errorbar(aes(ymin = down, ymax = up), width = 0.2, position = position_dodge(0.5)) + ylim(0, NA) + theme_classic() + facet_wrap(~ geneID, scales = "free", ncol = 4) + theme(strip.background = element_blank(), legend.position = "none", axis.title.x = element_blank())
dev.off()