Venn_plot_v2.r

#!/usr/bin/env Rscript

# =============================================================================
# 脚本功能：基于输入文件绘制Venn图
# 输入文件格式：两列，tab分隔，无表头
#   第一列：集合名称（label）
#   第二列：逗号分隔的集合元素（例如：gene1,gene2,gene3）
#   Control   geneA,geneB,geneC,geneD
#   KO        geneB,geneC,geneE,geneF
#   KD        geneA,geneC,geneG

# 输出：Venn图保存到输出目录（默认PDF格式）
# 依赖包：ggvenn, ggplot2, optparse
# 使用示例: Rscript venn_plot.R -i sets.txt -o ./results -f pdf -w 8 -H 8
# =============================================================================

# 加载或安装必要的包
source("/data/med-wangcq/01CondaEnv/02Git_repo/PlotScripts/01loadPackages.r")
# library(optparse)
# library(ggvenn)
# library(ggplot2)

# 解析命令行参数
option_list <- list(
  make_option(c("-i", "--infile"), type = "character", default = NULL,
              help = "输入文件路径（两列，tab分隔，第一列label，第二列逗号分隔的元素; 无列名）", metavar = "file"),
  make_option(c("-o", "--outdir"), type = "character", default = ".",
              help = "输出目录 [默认 %default]", metavar = "path"),
  make_option(c("-p", "--prefix"), type = "character", default = "no",
              help = "指定输出文件前缀 [默认 %default]"),
  make_option(c("-s", "--geneset"), type = "character", default = "all",
              help = "筛选需要的基因集列表，逗号分隔 [默认 %default]"),
  make_option(c("-w", "--width"), type = "numeric", default = 8,
              help = "图片宽度（英寸） [默认 %default]", metavar = "numeric"),
  make_option(c("-H", "--height"), type = "numeric", default = 8,
              help = "图片高度（英寸） [默认 %default]", metavar = "numeric")
)

opt_parser <- OptionParser(option_list = option_list, 
                           description = "绘制基于集合列表的Venn图")
opt <- parse_args(opt_parser)

prefix <- opt$prefix
infile <- opt$infile
outdir <- opt$outdir
geneset <- opt$geneset
width <- opt$width
height <- opt$height

cat("prefix: ",prefix,"\n")
cat("infile: ",infile,"\n")
cat("outdir: ",outdir,"\n")
cat("geneset: ",geneset,"\n")
cat("width: ",width,"\n")
cat("height: ",height,"\n")

# 检查必要参数
if (is.null(infile)) {
  print_help(opt_parser)
  stop("错误：必须提供输入文件（--infile）", call. = FALSE)
}

# 创建输出目录（如果不存在）
if (!dir.exists(outdir)) {
  dir.create(outdir, recursive = TRUE)
}

# 读取输入文件
message("正在读取输入文件：", infile)
input_data <- tryCatch({
  read.table(infile, sep = "\t", header = FALSE, stringsAsFactors = FALSE)
}, error = function(e) {
  stop("无法读取输入文件，请检查文件路径和格式：", e$message)
})

# 验证列数
if (ncol(input_data) < 2) {
  stop("输入文件必须至少包含两列（tab分隔）")
}

# 筛选需要基因集
if( geneset != "all"){
  gs <- unlist(strsplit(geneset,","))
  input_data <- input_data[which(input_data[,1] %in% gs),]
}
print(dim(input_data))
gs_use <- input_data[,1]
cat("使用 ",gs_use," 作图\n")

# 构建集合列表
sets <- list()
for (i in 1:nrow(input_data)) {
  label <- trimws(input_data[i, 1])
  elements_raw <- trimws(input_data[i, 2])
  # 处理空行或空元素
  if (elements_raw == "" || is.na(elements_raw)) {
    next
  } else {
    # 按逗号分割，并去除每个元素前后的空格
    elements <- trimws(unlist(strsplit(elements_raw, ",")))
    # 去除空字符串
    elements <- elements[elements != ""]
  }
  # 计算每个集合的大小
  message("集合 ", label, ": ", length(elements))
  label <- paste0(label,"_",length(elements))
  sets[[label]] <- elements
}

# 检查集合数量
n_sets <- length(sets)
message("检测到 ", n_sets, " 个集合：", paste(names(sets), collapse = ", "))
if (n_sets < 2) {
  stop("至少需要2个集合才能绘制Venn图，当前集合数：", n_sets)
}
# 建立交集结果文件
generate_intersection_tables <- function(sets, out_prefix) {
    all_elements <- unique(unlist(sets))
    set_names <- names(sets)
    
    # 构建归属矩阵 (0/1)
    membership <- sapply(sets, function(x) as.integer(all_elements %in% x))
    rownames(membership) <- all_elements
    
    # 生成交集组合标签（如：Control_4&KO_4）
    intersection_labels <- apply(membership, 1, function(x) {
      present <- set_names[as.logical(x)]
      if (length(present) == 0) return("None")
      paste(present, collapse = "&")
    })
    
    # 表格1：基因级别归属表
    gene_table <- as.data.frame(membership, stringsAsFactors = FALSE)
    gene_table$Gene <- all_elements
    gene_table$Intersection <- intersection_labels
    gene_table <- gene_table[, c("Gene", set_names, "Intersection")]
    
    # 表格2：交集统计表
    combo_counts <- table(intersection_labels)
    combo_genes <- tapply(gene_table$Gene, gene_table$Intersection, paste, collapse = ",")
    
    summary_df <- data.frame(
      Intersection = names(combo_genes),
      Count = as.integer(combo_counts[names(combo_genes)]),
      Genes = as.character(combo_genes),
      stringsAsFactors = FALSE,
      row.names = NULL
    )
    summary_df <- summary_df[order(summary_df$Count, decreasing = TRUE), ]
    rownames(summary_df) <- NULL
    
    # 保存为tab分隔文件
    gene_file <- paste0(out_prefix, "_genes_membership.txt")
    summary_file <- paste0(out_prefix, "_intersections_summary.txt")
    
    write.table(gene_table, gene_file, sep = "\t", quote = FALSE, row.names = FALSE)
    write.table(summary_df, summary_file, sep = "\t", quote = FALSE, row.names = FALSE)
    
    message("基因归属表已保存至：", gene_file)
    message("交集统计表已保存至：", summary_file)
    
    return(list(gene_table = gene_table, summary = summary_df))
}

# =============================================================================
# 绘图与输出
# =============================================================================
out_prefix <- file.path(outdir, prefix)

if (n_sets > 5) {
  message("集合数超过5，绘制 UpSet 图...")
  library(ComplexUpset)
  
  # 构建 UpSet 所需的宽格式二进制矩阵
  all_elements <- unique(unlist(sets))
  membership_df <- as.data.frame(
    sapply(sets, function(x) as.integer(all_elements %in% x)),
    stringsAsFactors = FALSE
  )
  rownames(membership_df) <- all_elements
  # 确保列名合法
  colnames(membership_df) <- make.names(names(sets), unique = TRUE)
  
  upset_plot <- upset(
    membership_df,
    intersect = colnames(membership_df),
    base_annotations = list(
      'Intersection size' = intersection_size(counts = TRUE)
    ),
    set_sizes = upset_set_size(),
    width_ratio = 0.1
  )
  
  outfile <- file.path(outdir, paste0(prefix, "_upset_plot.pdf"))
  ggsave(filename = outfile, plot = upset_plot, width = width, height = height, dpi = 300)
  message("UpSet图已保存至：", outfile)
  
} else {
  message("正在绘制 Venn 图...")
  venn_plot <- ggvenn(sets, 
                      fill_color = c("#E41A1C", "#377EB8", "#4DAF4A", "#984EA3", "#FF7F00")[1:n_sets],
                      stroke_size = 0.5,
                      set_name_size = 5,
                      text_size = 4)
  
  outfile <- file.path(outdir, paste0(prefix, "_venn_plot.pdf"))
  ggsave(filename = outfile, plot = venn_plot, width = width, height = height, dpi = 300)
  message("Venn图已保存至：", outfile)
}

# 统一调用函数输出交集表格（支持任意集合数）
message("正在生成交集统计表格...")
generate_intersection_tables(sets, out_prefix)

message("完成！")