utils.R

#!/usr/bin/Rscript

## IMPORT

##############################
# INPUT

#' Transform a molten data.frame into a squared data.frame
#' For lists of pairwise comparisons
#' @input df Molten data.frame w/ three columns (two w/ labels and one w/ values)
#' @input col1 Column name containing labels (1)
#' @input col2 Column name containing labels (2)
#' @return a data.frame with labels from two label columns as row and column names
dcast_sq <- function(df, col1, col2){
    # make sure the labels are identical
    testit::assert(all( sort(df[,col1]) == sort(df[,col2]) ))
    # reshape given data.frame using given formula
    df <- reshape2::dcast(df, as.formula(sprintf("%s ~ %s", col1, col2)))
    # use col. w/ tool names as rownames and remove from table
    rownames(df) <- df[,col1]
    df <- df[,setdiff(colnames(df), col1)]
    return(df)
} 

read_nanostats <- function(fname){
    print(sprintf("Reading: %s", fname))
    df <- read.csv(
        file=fname,
        sep="\t",
        header=TRUE,
        check.names=FALSE,
        stringsAsFactors=FALSE
    )
    df_cols <- c("stat"="Statistic", "value"="Value")
    colnames(df) <- df_cols[colnames(df)]
    return(df)
}

read_fastp <- function(fname){
    print(sprintf("Reading: %s", fname))
    df <- read.csv(
        file=fname,
        sep="\t",
        header=TRUE,
        row.names=1,
        check.names=FALSE,
        stringsAsFactors=FALSE
    )
    df <- df[,c("total_reads", "total_bases", "q20_rate", "q30_rate")]
    return(df)
}

read_mappability <- function(fname){
    print(sprintf("Reading: %s", fname))
    df <- read.csv(
        file=fname,
        sep="\t",
        header=TRUE,
        check.names=FALSE,
        stringsAsFactors=FALSE
    )
    testit::assert(all(df$tool %in% names(ASM_TOOL_NAMES)))
    df$tool <- ASM_TOOL_NAMES[df$tool]
    return(df)
}

read_prodigal <- function(fname){
    print(sprintf("Reading: %s", fname))
    df <- read.csv(
        file=fname,
        sep="\t",
        header=TRUE,
        check.names=FALSE,
        stringsAsFactors=FALSE
    )
    testit::assert(all(df$tool %in% names(ASM_TOOL_NAMES)))
    df$tool <- ASM_TOOL_NAMES[df$tool]
    return(df)
}

read_prodigal_gcounts <- function(fname){
    df <- read_prodigal(fname)
    df$partial_pct <- 100 * df$partial / df$total
    return(df)
}

read_prodigal_glength <- function(fname){
    df <- read_prodigal(fname)
    return(df)
}

read_quast <- function(fname){
    print(sprintf("Reading: %s", fname))
    df <- read.csv(
        file=fname,
        sep="\t",
        header=TRUE,
        row.names=1,
        check.names=FALSE,
        stringsAsFactors=FALSE
    )
    testit::assert(all(colnames(df) %in% names(ASM_TOOL_NAMES)))
    colnames(df) <- ASM_TOOL_NAMES[colnames(df)]
    df <- df[QUAST_VARS, ASM_TOOL_NAMES]
    return(df)
}

read_plasflow <- function(fname){
    print(sprintf("Reading: %s", fname))
    df <- read.csv(
        file=fname,
        sep="\t",
        header=TRUE,
        stringsAsFactors=FALSE,
        check.names=FALSE
    )
    testit::assert(all(df$tool %in% names(ASM_TOOL_NAMES)))
    df$tool <- ASM_TOOL_NAMES[df$tool]
    df <- df[df$label %in% names(PLASFLOW_NAMES$labels),]
    df$label <- PLASFLOW_NAMES$labels[df$label]
    return(df)
}

read_rgi <- function(fname){
    print(sprintf("Reading: %s", fname))
    df <- read.csv(
        file=fname,
        sep="\t",
        header=TRUE,
        stringsAsFactors=FALSE,
        check.names=FALSE
    )
    colnames(df) <- sapply(colnames(df), function(x){ ifelse(x %in% names(ASM_TOOL_NAMES), ASM_TOOL_NAMES[x], x) })
    testit::assert(all(df$col %in% names(RGI_NAMES$col)))
    df$col <- RGI_NAMES$col[df$col]
    return(df)
}

read_barrnap <- function(fname){
    print(sprintf("Reading: %s", fname))
    df <- read.csv(
        file=fname,
        sep="\t",
        header=TRUE,
        stringsAsFactors=FALSE,
        check.names=FALSE
    )
    testit::assert(all(df$tool %in% names(ASM_TOOL_NAMES)))
    df$tool <- ASM_TOOL_NAMES[df$tool]
    testit::assert(all(df$kingdom %in% names(BARRNAP_KINGDOM_NAMES)))
    df$kingdom <- BARRNAP_KINGDOM_NAMES[df$kingdom]
    return(df)
}

proc_barrnap <- function(df){
    # total counts per tool, kingdom and partial/complete categories
    df_total <- aggregate(
        df$count,
        by=list(tool=df$tool, kingdom=df$kingdom, partial=grepl("partial", df$gene)),
        FUN=sum
    )
    # partial: FALSE/TRUE -> label
    df_total$partial <- c("complete", "partial")[df_total$partial + 1]
    # add partial + complete = total
    df_total <- rbind(
        df_total,
        aggregate(
            df_total$x,
            by=list(tool=df_total$tool, kingdom=df_total$kingdom, partial=rep("total", nrow(df_total))),
            FUN=sum
        )
    )
    return(df_total)
}

read_crispr <- function(fname){
    print(sprintf("Reading: %s", fname))
    df <- read.csv(
        file=fname,
        sep="\t",
        header=TRUE,
        stringsAsFactors=FALSE,
        check.names=FALSE
    )
    # testit::assert(all(df$crispr_tool %in% names(CRISPR_TOOL_NAMES)))
    testit::assert(all(df$tool    %in% names(ASM_TOOL_NAMES)))
    # df$crispr_tool <- CRISPR_TOOL_NAMES[df$crispr_tool]
    df$tool <- ASM_TOOL_NAMES[df$tool]
    return(df)
}

aggr_crispr <- function(df){
    df_spacers <- aggregate(
        x=df$spacers,
        by=list(tool=df$tool),
        FUN=sum
    )
    rownames(df_spacers) <- df_spacers$tool
    df_arrays <- aggregate(
        x=df$seq_id,
        by=list(tool=df$tool),
        FUN=length
    )
    rownames(df_arrays) <- df_arrays$tool
    return(list(spacers=df_spacers, arrays=df_arrays))
}

read_diamondDB <- function(fname){
    print(sprintf("Reading: %s", fname))
    df <- read.csv(
        file=fname,
        sep="\t",
        header=TRUE,
        check.names=FALSE,
        stringsAsFactors=FALSE
    )
    testit::assert(all(df$tool %in% names(ASM_TOOL_NAMES)))
    df$tool <- ASM_TOOL_NAMES[df$tool]
    return(df)
}

read_ugenes <- function(fname){
    print(sprintf("Reading: %s", fname))
    df <- read.csv(
        file=fname,
        sep="\t",
        header=TRUE,
        check.names=FALSE,
        stringsAsFactors=FALSE
    )
    df$uniq_pct                 <- 100 * df$uniq / df$total
    df$highcovuniq_pct_total    <- 100 * df$highcovuniq / df$total
    df$highcovuniq_pct_uniq     <- 100 * df$highcovuniq / df$uniq
    testit::assert(all(df$tool1 %in% names(ASM_TOOL_NAMES)))
    testit::assert(all(df$tool2 %in% names(ASM_TOOL_NAMES)))
    df$tool1 <- ASM_TOOL_NAMES[df$tool1]
    df$tool2 <- ASM_TOOL_NAMES[df$tool2]
    return(df)
}

read_fastani_many2many <- function(fname){
    print(sprintf("Reading: %s", fname))
    proc_name <- function(x){
        x <- basename(dirname(x))
        testit::assert(x %in% names(ASM_TOOL_NAMES))
        return(ASM_TOOL_NAMES[x])
    }
    dm <- read.csv(file=fname, sep='\t', header=FALSE, check.names=FALSE, stringsAsFactors=FALSE, col.names=c("tool1", "tool2", "ani", "mappings", "queries"))
    dm <- dcast_sq(df=dm[,c("tool1", "tool2", "ani")], col1="tool1", col2="tool2")
    # proc. names
    colnames(dm) <- sapply(colnames(dm), proc_name)
    rownames(dm) <- sapply(rownames(dm), proc_name)
    return(dm)
}

read_mummer_dnadiff <- function(fname){
    print(sprintf("Reading: %s", fname))
    df <- read.csv(file=fname, sep='\t', header=TRUE, check.names=FALSE, stringsAsFactors=FALSE)
    df$seqs_pct  <- 100 * df$seqs_aligned  / df$seqs_total
    df$bases_pct <- 100 * df$bases_aligned / df$bases_total
    testit::assert(all(df$tool1 %in% names(ASM_TOOL_NAMES)))
    testit::assert(all(df$tool2 %in% names(ASM_TOOL_NAMES)))
    df$tool1 <- ASM_TOOL_NAMES[df$tool1]
    df$tool2 <- ASM_TOOL_NAMES[df$tool2]
    return(df)
}

read_mash_dist_reads <- function(fname){
    print(sprintf("Reading: %s", fname))
    proc_name <- function(x){
        rtype <- basename(dirname(x))
        mtype <- basename(dirname(dirname(x)))
        testit::assert(rtype %in% names(READ_TYPES))
        testit::assert(mtype %in% names(META_TYPES))
        return(sprintf("%s %s", META_TYPES[mtype], READ_TYPES[rtype]))
    }
    dm <- read.csv(file=fname, sep='\t', header=TRUE, row.names=1, check.names=FALSE)
    colnames(dm) <- sapply(colnames(dm), proc_name)
    rownames(dm) <- sapply(rownames(dm), proc_name)
    return(dm)
}

read_mash_dist_asm <- function(fname){
    print(sprintf("Reading: %s", fname))
    proc_name <- function(x){
        x <- basename(dirname(x))
        testit::assert(x %in% names(ASM_TOOL_NAMES))
        return(ASM_TOOL_NAMES[x])
    }
    dm <- read.csv(file=fname, sep='\t', header=TRUE, row.names=1, check.names=FALSE)
    colnames(dm) <- sapply(colnames(dm), proc_name)
    rownames(dm) <- sapply(rownames(dm), proc_name)
    return(dm)
}

##############################
# PLOTS

set_tool_order <- function(df, cols){
    for(cname in cols){
        df[,cname] <- factor(df[,cname], ordered=TRUE, levels=ASM_TOOL_NAMES)
    }
    return(df)
}

plot_mappability <- function(df){
    df <- set_tool_order(df, c("tool"))
    pp <- 
        ggplot(data=df, aes(x=tool, y=value, fill=tool)) +
        geom_col() +
        geom_text(aes(label=value, y=0.5*value), color="black", size=4, angle=90) +
        scale_fill_manual(values=ASM_TOOL_COLORS, guide=NULL) +
        facet_wrap(vars(minlength), ncol=2, scales="fixed") +
        labs(
            x="",
            y="Mapped reads [%]"
        ) +
        mappability_theme
    return(pp)
}

plot_prodigal <- function(df_gc, df_gl){
    df_gcm <- reshape2::melt(df_gc[,c("tool", "total", "partial")], id.vars=c("tool"))
    df_gc  <- set_tool_order(df_gc,  c("tool"))
    df_gl  <- set_tool_order(df_gl,  c("tool"))
    df_gcm <- set_tool_order(df_gcm, c("tool"))
    pp1 <-
        ggplot(data=df_gcm, aes(x=variable, y=value, fill=tool)) +
        geom_col(position="dodge") +
        scale_fill_manual(values=ASM_TOOL_COLORS) +
        labs(
            x="",
            y="Gene count"
        ) +
        prodigal_theme
    
    pp2 <-
        ggplot(data=df_gc, aes(x=tool, y=partial_pct, fill=tool)) +
        geom_col() +
        scale_fill_manual(values=ASM_TOOL_COLORS, guide=NULL) +
        labs(
            x="",
            y="Percentage of partial genes"
        ) +
        prodigal_theme
    
    pp3 <-
        ggplot(data=df_gl, aes(x=tool, y=gene_length, fill=tool)) +
        geom_violin(draw_quantiles = c(0.25, 0.5, 0.75)) +
        scale_fill_manual(values=ASM_TOOL_COLORS, guide=NULL) +
        labs(
            x="",
            y="Gene length [bp]"
        ) +
        diamondDB_theme
    
    pp4 <- pp3 + coord_cartesian(ylim=c(0, 2000))

    return(list(
        gcounts=pp1,
        gpct=pp2,
        glen=pp3,
        glen_zoom=pp4
    ))
}

plot_barrnap <- function(df){
    df <- set_tool_order(df, c("tool"))
    pp <-
        ggplot(data=df, aes(x=tool, y=x, fill=tool)) +
        geom_col() +
        scale_fill_manual(values=ASM_TOOL_COLORS, guide=NULL) +
        facet_grid(vars(partial), vars(kingdom), scales="fixed") +
        labs(
            x="",
            y=sprintf("rRNA gene hits")
        ) +
        barrnap_theme
    return(pp)
}

plot_barrnap_genes <- function(df, subtitle){
    df <- set_tool_order(df, c("tool"))
    pp <-
        ggplot(data=df, aes(x=tool, y=count, fill=tool)) +
        geom_col() +
        scale_fill_manual(values=ASM_TOOL_COLORS, guide=NULL) +
        facet_wrap(vars(gene), ncol=1, scales="free_y") +
        labs(
            subtitle=subtitle,
            x="",
            y="rRNA gene hits"
        ) +
        barrnap_theme
    return(pp)
}

plot_rgi_overlap <- function(df, ctype, col){
    df <- df[df$type == ctype & df$col == col,]
    df_list <- lapply(ASM_TOOL_NAMES, function(x){ df[df[,x] > 0,"label"] })
    names(df_list) <- ASM_TOOL_NAMES[names(df_list)]
    UpSetR::upset(
        data=UpSetR::fromList(df_list),
        # overlap order
        order.by="degree",
        decreasing=FALSE,
        # number of sets to plot
        nsets=length(ASM_TOOL_NAMES),
        # y-label title
        mainbar.y.label=sprintf("Intersection size (%s hits, %s)", ctype, col),
        # text size
        text.scale = c(1.2, 1.2, 1.2, 1.2, 1.2, 1.2)#,
        # colors
        # set.metadata=list(
        #     data=data.frame(
        #         sets=names(df_list),
        #         Tool=names(df_list),
        #         stringsAsFactors=FALSE
        #     ), 
        #     plots=list(list(type="matrix_rows", column="Tool", colors=ASM_TOOL_COLORS, alpha=0.7))
        # )
    )
}

plot_crispr <- function(df){
    df_m <- reshape2::melt(df, id.vars="tool")
    df_m <- set_tool_order(df_m, c("tool"))
    pp <-
        ggplot(data=df_m, aes(x=tool, y=value, fill=tool)) +
        geom_col(position="dodge") +
        scale_fill_manual(values=ASM_TOOL_COLORS, guide=NULL) +
        facet_wrap(vars(variable), ncol=1, scales="free_y") +
        labs(
            x="",
            y="Number of features"
        ) +
        crispr_theme
    return(pp)
}

plot_plasflow <- function(df, ylab=""){
    df <- set_tool_order(df, c("tool"))
    pp <-
        ggplot(data=df, aes_string(x="tool", y="value", fill="label")) +
        geom_col(position="dodge") +
        scale_fill_manual(values=PLASFLOW_COLORS$labels, guide="legend") +
        labs(
            x="",
            y=ylab
        ) +
        plasflow_theme
    return(pp)
}

plot_quast <- function(df){
    df_m <- reshape2::melt(
        cbind(stat_vars=rownames(df), df),
        id.vars="stat_vars"
    )
    df_m <- set_tool_order(df_m, c("variable"))
    pp <-
        ggplot(data=df_m, aes(x=variable, y=value)) +
        geom_col(aes(fill=variable)) +
        scale_fill_manual(values=ASM_TOOL_COLORS, guide=NULL) +
        facet_wrap(vars(stat_vars), ncol=2, scales="free_y") +
        labs(
            x="",
            y="QUAST statistic"
        ) +
        quast_theme
    return(pp)
}

plot_ugenes_barplots <- function(df, ycol, ylab="", subtitle=""){
    df <- set_tool_order(df, c("tool1", "tool2"))
    pp <-
        ggplot(data=df, aes_string(x="tool2", y=ycol, fill="tool2")) +
        geom_col() +
        facet_wrap(vars(tool1), ncol=1) +
        scale_fill_manual(values=ASM_TOOL_COLORS, guide=NULL) +
        labs(
            subtitle=subtitle,
            x="Assembly 2",
            y=ylab
        ) +
        ugenes_theme
    return(pp)
}

plot_ugenes_scatterplot <- function(df, subtitle=""){
    df <- set_tool_order(df, c("tool1", "tool2"))
    pp <-
        ggplot(data=df, aes(x=uniq_pct, y=highcovuniq_pct_uniq, fill=tool2, shape=tool2)) +
        geom_point(colour="white", size=6) +
        scale_fill_manual(values=ASM_TOOL_COLORS, name="") +
        scale_shape_manual(values=ASM_TOOL_SHAPES, name="") +
        labs(
            subtitle=subtitle,
            x="Unique proteins [%, total]",
            y="Unique proteins w/ high mean cov. [% of unique]"
        ) +
        default_theme +
        theme(
            legend.position="bottom",
            legend.direction="horizontal"
        )
    return(pp)
}

plot_diamondDB_density <- function(df, col, xlim=NULL){
    df <- set_tool_order(df, c("tool"))
    pp <-
        ggplot(data=df, aes_string(x=col, colour="tool", fill="tool")) +
        geom_density(alpha=0.2) +
        scale_colour_manual(values=ASM_TOOL_COLORS, guide=NULL) +
        scale_fill_manual(values=ASM_TOOL_COLORS, guide=NULL) +
        facet_wrap(vars(tool), nrow=4, scales="free_y") +
        labs(
            subtitle=ifelse(!is.null(xlim), "zoomed in", ""),
            x=DIAMOND_VAR_LABELLER(col),
            y="Density"
        ) +
        diamondDB_theme
    if(!is.null(xlim)){
        pp <- pp + coord_cartesian(xlim=xlim)
    }
    return(pp)
}

plot_diamondDB_density2d <- function(df){
    df <- set_tool_order(df, c("tool"))
    pp <-
        ggplot(data=df, aes(x=qcov, y=scov)) +
        geom_bin2d(bins=25) +
        scale_fill_continuous(type="viridis") +
        facet_wrap(vars(tool), ncol=3, nrow=3) +
        labs(
            x=DIAMOND_VAR_LABELLER("qcov"),
            y=DIAMOND_VAR_LABELLER("scov")
        ) +
        diamondDB_theme
    return(pp)
}

#' Contig coverage and segmentation plot
#' @input df_cov Coverage data.frame incl. contig ID, base, coverage and state
#' @input cid Contig ID
#' @return ggplot2 object
plot_contig_cov_seg <- function(df_cov, cid){
    # subset cov.
    df_cov <- df_cov[df_cov$contig==cid,]
    # segments/states medians
    df_cov$state_median <- NA
    for(s in unique(df_cov$state)){
        df_cov$state_median[df_cov$state == s] <- median(df_cov[df_cov$state==s, "cov"])
    }
    # plot
    df_plot <-
        ggplot(data=df_cov, aes(x=base, y=cov)) +
        geom_line(colour="#666666") +
        geom_line(aes(x=base, y=state_median), colour="#0066CC", size=2) +
        scale_x_continuous(breaks=seq(0, max(df_cov$base), by=5000)) +
        scale_y_log10(breaks=trans_breaks("log10", function(x) 10^x), labels=trans_format("log10", math_format(10^.x))) +
        labs(
            title=cid,
            x="base",
            y="coverage"
        ) +
        theme_bw() +
        theme(
            axis.text.x=element_text(angle=90, vjust=0.5, hjust=1)
        )
    return(df_plot)
}

#' Coverage and segmentation scatter plot
#' @input df_states Coverage segmentation summary data.frame
#' @input title Title
#' @input subtitle Sub-title
#' @return ggplot object
plot_contig_cov_seg_scatterplot <- function(df_states, title="", subtitle=""){
    pp <-
        ggplot(data=df_states[df_states$states > 1,], aes(x=length, y=mean)) +
        geom_point(data=df_states[df_states$states < 2,], aes(x=length, y=mean), shape=4, color="#CCCCCC", size=0.5) +
        geom_point(aes(size=states_median_sd, fill=as.factor(states), color=as.factor(states)), shape=21, alpha=0.75) +
        scale_x_log10(
            breaks=trans_breaks("log10", function(x) 10^x),
            labels=trans_format("log10", math_format(10^.x))
        ) +
        scale_y_log10(
            breaks=trans_breaks("log10", function(x) 10^x),
            labels=trans_format("log10", math_format(10^.x))
        ) +
        labs(
            title=title, #"Contig coverage segmentation",
            subtitle=subtitle, #"Flye + Racon(metaG SR)",
            fill="Number of states",
            color="Number of states",
            size="State median SD",
            x="Contig length [bp]",
            y="Mean contig coverage"
        ) +
        theme_bw()
    return(pp)
}

# THEMES
default_theme <-
    # theme_bw() +
    theme_minimal(
        base_size=12
    ) +
    theme(
        plot.title=element_text(size=14, face="bold"),
        plot.subtitle=element_text(size=12, face="italic"),
        # legend
        legend.title=element_blank(),
        legend.text=element_text(size=12),
        # grid
        panel.grid=element_blank(),
        # strip
        strip.background=element_rect(fill="white"),
        strip.text=element_text(size=12),
        # axes
        axis.title=element_text(size=12, color="black"),
        axis.text.y=element_text(size=12, color="black"),
        axis.text.x=element_text(size=12, color="black")
    )

default_theme_axis_text_x <-
    theme(
        axis.text.x=element_text(size=12, color="black", angle=90, vjust=0.5, hjust=1)
    )

mappability_theme <- default_theme + default_theme_axis_text_x

crispr_theme <- default_theme + default_theme_axis_text_x

plasflow_theme <- default_theme + default_theme_axis_text_x

prodigal_theme <- default_theme + default_theme_axis_text_x

diamondDB_theme <- default_theme + default_theme_axis_text_x

rgi_theme <- default_theme + default_theme_axis_text_x

barrnap_theme <- default_theme + default_theme_axis_text_x

quast_theme <- default_theme + default_theme_axis_text_x

ugenes_theme <- default_theme + default_theme_axis_text_x