<repo_name>ctl43/sarlacc_prototype<file_sep>/sam2ranges/split_sam.R
sam <- sam2ranges("/Users/law02/Dropbox (Cambridge University)/Univeristy of Cambridge/Academic/Internship/NanoporeProject/test_20180612/analysis/2h-RNA-unbl/data/genome_pass-both-chopped_sarlacc_2h-RNA-unbl_full-oligo-dT_TSO-with-TruSeq.sam.gz", minq = 0)
undebug(sam2ranges)
library(GenomicAlignments)
cigar <- "6S39M2D9M5D7M1D12M2D5M1D30M8710N3I50M2D18M1I4M2D9M1D1M1D7M2I27M1D20M1D14M1D9M2I3M2D2M1D3M7D3M2I7M1752N2M3D10M1D3M1D22M1D47M1D17M1I14M4D85945N17M1I17M1D22M1D22M1D4M2D87M2D4M1I8M6I10M1D6M123692N1M2D8M2D15M1I5M1D11M83869N19M1I22M1I17M1D10M59493N64M1D15M1D5M1D1M3D16M5462N8M2D32M2I7M2D7M1I20M2D1M1D23M4D4M4D5M3D6M2D16M594N1M1D37M1D49M2D14M1I34M3D22M1I9M2D22M2D8M1D2M1D3M2D13M4D48M2I14M1D3M1D37M21S"

load("internship/NanoporeProject/sarlacc_prototype/sam2ranges/sam2ranges.RData")

mapping <- mapping[1:1000,]

mapping$cigar.groups <- cumsum(!duplicated(mapping$QNAME))
cigar.rle <- cigarToRleList(mapping$CIGAR)
run.val <- runValue(cigar.rle)
run.len <- runLength(cigar.rle)
is.skip <- run.val=="N"
is.skip <- lapply(is.skip, function(x){x[which(x)+1] <- 1; x})
group <- lapply(is.skip, cumsum)
split.len <- mapply(function(x,y)sum(IntegerList(split(x,f=y))), x=run.len, y=group, SIMPLIFY = FALSE)
start <- mapply(function(x,y)cumsum(c(x,y)),x = mapping$POS,y = split.len)
end <- lapply(start, function(x)(x[-1]-1)[seq(1,length(x)-1,2)])
start <- lapply(start, function(x)x[-length(x)][seq(1,length(x)-1,2)])
split.granges <- mapply(function(x,y,z,s)GRanges(x, IRanges(y, end=z), strand=ifelse(bitwAnd(s, 0x10), "-", "+")),
                        x=mapping$RNAME,y=start,z=end,s=mapping$FLAG)

GRangesList(split.granges)<file_sep>/homopolymerFixer/homopolymerFixer.R
homopolymerTrainer <- function(grouped.reads, ref, each.set=10, rep=30, min.cov,BPPARAM=SerialParam()){
    grouped.reads <- grouped.reads[lengths(grouped.reads)>(each.set)]
    max.times <- pmin(rep,floor(lengths(grouped.reads)/each.set))
    sampled.reads <- mapply(function(x,y)replicate(y,sample(x,each.set)), 
                            x=grouped.reads, y=max.times, SIMPLIFY = FALSE)
    sampled.reads <- unlist(sampled.reads, recursive = FALSE, use.names = FALSE)
    ercc.label <- rep(names(grouped.reads),max.times)
    msa <- bplapply(sampled.reads, multiReadAlign, BPPARAM = BPPARAM)
    train <- bplapply(seq_along(ercc.label), function(x){
        homopolymerCounter(msa[[x]], ref[ercc.label[x]], min.cov=min.cov)},BPPARAM=BPPARAM)
    train <- do.call(rbind, unlist(train))
    base.len <- tapply(train$observed, list(train$base,train$len), unlist, use.names = FALSE)
    debase.len <- apply(base.len, 2, function(x)unlist(x))
    debase.learned.distr <- sapply(debase.len, .proportion_calculator)
    debase.learned.distr <- rep(debase.learned.distr, nrow(base.len))
    debase.learned.distr <- matrix(debase.learned.distr, nrow(base.len), byrow = TRUE)
    learned.distr <- sapply(base.len, .proportion_calculator)
    learned.distr <- matrix(learned.distr, nrow(base.len), byrow = FALSE)
    dimnames(debase.learned.distr) <- dimnames(learned.distr) <- dimnames(base.len)
    learned.len.distr <- mapply(function(x,y)learned.distr[x,y], 
                                  x=train$base, y=as.character(train$len))
    len.distr <- lapply(train$observed, function(x)table(x)/each.set)
    msd <- .msd_calculator(len.distr, learned.len.distr)
    msd.distr <- tapply(msd, list(train$base,train$len),unlist, use.names = FALSE)
    debase.msd.distr <- apply(msd.distr, 2, unlist)
    debase.msd.distr <- matrix(rep(NumericList(debase.msd.distr), nrow(base.len)),nrow(base.len), byrow=TRUE)
    is.empty <- sapply(learned.distr, is.null)
    is.empty[c(lengths(base.len)<each.set*rep)] <- TRUE
    learned.distr[is.empty] <- debase.learned.distr[is.empty]
    msd.distr[is.empty] <- debase.msd.distr[is.empty]
    dimnames(debase.msd.distr) <- dimnames(msd.distr) <- dimnames(base.len)

    list(base=list(len.distr=learned.distr, msd.distr=msd.distr),
         debase=list(len.distr=debase.learned.distr, msd.distr=debase.msd.distr))
}

homopolymerCounter <- function(msa, ref.seq, end.dist=10, min.cov){
    # Get the alignment matrix
    aln.matrix <- unname(do.call(cbind,strsplit(unlist(msa$alignments), split="")))
    con.seq <- QUALFUN(msa$alignments[[1]], msa$qualities[[1]], min.cov, gapped = TRUE)
    cons.aln <- .get_translation_table(con.seq$seq, ref.seq)
    aln.ref <- cons.aln[[2]]
    cons.aln <- cons.aln[[1]]
    # Get homopolymer matches between consensus and the reference
    # Get the corresponding location in alignment matrix
    matched <- MATCHCHECK(aln.ref)
    pass <- start(matched$ref.pos)>end.dist&end(matched$ref.pos)<((nchar(ref.seq)-end.dist)+1)
    
    matched$aligned.pos <- matched$aligned.pos[pass]
    matched$ref.pos <- matched$ref.pos[pass]
    
    aliged.start <- start(matched$aligned.pos)
    aligned.end <- end(matched$aligned.pos)
    aligned.loc <- mapply(":",(aliged.start-1), (aligned.end+1))
    trigger <- as.numeric(elementMetadata(matched$ref.pos)$observed)
    
    # Get the position in alignment matrix
    # aln.mat.pos <- lapply(aligned.loc, function(x)range(cons.aln$aln.pos[x], na.rm = TRUE))
    aln.mat.pos <- mapply(function(x,y)if(x==0) c(-1,1)else range(cons.aln$aln.pos[y], na.rm = TRUE),
                          x=trigger,y=aligned.loc, SIMPLIFY = FALSE)
    # aln.mat.pos <- mapply(function(x,y)if(x==0)y<-c(-1,1)else y, x=trigger, y=aln.mat.pos, SIMPLIFY=FALSE)
    stopifnot(all(unlist(lapply(aln.mat.pos, function(x)all(is.finite(x))))))
    # Get the corresponding seq from alignment matrix
    aln.seq <- lapply(aln.mat.pos, function(x)DNAStringSet(substr(msa$alignments[[1]], start=x[1]+1,stop=x[2]-1)))
    # Find homopolymer in aligned seqeunces
    seq.num <- length(msa$alignments[[1]])
    aln.seq <- do.call(c, aln.seq)
    specific.bases <- rep(elementMetadata(matched$ref.pos)$base,each=seq.num)
    obs.homo.len <- .max_homo_length(aln.seq, specific.bases)
    obs.homo.len <- split(obs.homo.len, f=rep(seq_len(length(obs.homo.len)/seq.num),each=seq.num))
    
    # Count the number of reads with corresponding base and the observed length of the homopolymer
    out <- DataFrame(len=width(matched$ref.pos), base=elementMetadata(matched$ref.pos)$base)
    out$observed <- obs.homo.len
    out
}

homopolymerFixer <- function(learned, msa, min.cov=0.6, threshold.p=0.2){
    learned.len.distr <- learned$len.distr
    learned.msd.distr <- learned$msd.distr
    aln.matrix <- unname(do.call(cbind,strsplit(unlist(msa$alignments), split="")))
    con.seq <- QUALFUN(msa$alignments[[1]], msa$qualities[[1]], min.cov)
    con.seq$aln.pos <- .get_aln_pos(con.seq$seq)
    cons.rle <- homopolymerFinder(DNAStringSet(con.seq$seq), min.homo = 1)[[1]]
    cons.homo <- cons.rle[width(cons.rle)>1]
    L <- length(cons.homo)

    cons.start <- start(cons.homo)-1
    cons.end <- end(cons.homo)+1
    aln.loc <- mapply(function(x,y,z)con.seq$aln.pos[x:y],x=cons.start, y=cons.end)
    search.range <- lapply(aln.loc, range, na.rm=TRUE)
    check.seq <- lapply(search.range, function(x)DNAStringSet(substr(msa$alignments[[1]], start=x[1]+1,stop=x[2]-1)))
    check.seq <- do.call(c, check.seq)
    seq.num <- length(msa$alignments[[1]])
    specific.bases <- rep(elementMetadata(cons.homo)$base,each=seq.num)
    check.count <- .max_homo_length(check.seq, specific.bases)
    check.count <- split(check.count, f=rep(seq_len(length(check.count)/seq.num),each=seq.num))
    len.distr <- lapply(check.count, .proportion_calculator)
    
    msd <- mapply(function(x,y).msd_calculator(x, learned.len.distr[y,]), 
                    x= len.distr, y=elementMetadata(cons.homo)$base, 
                    SIMPLIFY = FALSE)
    p <- mapply(function(x,y)sum(x<NumericList(learned.msd.distr[y,]))/lengths(learned.msd.distr[y,]), 
                    x= msd, y=elementMetadata(cons.homo)$base, 
                  SIMPLIFY = FALSE)

    corrected.homo <- as.numeric(sapply(p, function(x)names(which.max(x))))
    
    # Safeguard1: if all p<threshold, then it will not be corrected
    dunno <- sapply(p,function(x)all(x<threshold.p))
    corrected.homo[dunno] <- width(cons.homo)[dunno]
    
    # Safeguard2: Only lengthen the homopolymer regions
    shortern <- corrected.homo<width(cons.homo)
    corrected.homo[shortern] <- width(cons.homo)[shortern]
    
    # Safeguard3: If the search region is exactly the same as the intitial homopolymer length, do nothing
    search.length <- sapply(search.range, function(x)diff(x)+1)
    no.change <- search.length == width(cons.homo)
    corrected.homo[no.change] <- width(cons.homo)[no.change]
    
    width(cons.rle)[width(cons.rle)>1] <- as.integer(corrected.homo)
    fixed.cons <- rep(elementMetadata(cons.rle)$base, times = width(cons.rle))
    fixed.cons <- paste(fixed.cons, collapse = "")
    result <- list(degapped = gsub("-","",con.seq$seq) ,fixed = fixed.cons)
    result
}

.get_aln_pos <- function(gapped){
    vec <- strsplit(gapped, split="")[[1]]
    aln.pos <- which(!grepl("-",vec))
    aln.pos
}

.max_homo_length <- function(seq, base){
    obs.homo.len <- homopolymerFinder(seq, base, report0 = TRUE, min.homo = 1)
    max.len <- lapply(obs.homo.len, function(x)max(width(x)))
    unlist(max.len)
}

.get_translation_table <- function(gapped.con, ref.seq){
    no.gap.seq <- gsub("-","",gapped.con)
    
    # Align the consensus sequence to the reference sequence
    aln.ref <- pairwiseAlignment(DNAStringSet(no.gap.seq), ref.seq, gapOpening=5, gapExtension=1)
    cons.aln <- strsplit(as.character(c(alignedSubject(aln.ref), alignedPattern(aln.ref))), split="")
    cons.aln <- do.call(cbind.data.frame, cons.aln)
    names(cons.aln) <- c("ref", "cons")
    
    # Create a table for indexing the reference positino and the position in alignment matrix
    cons.aln <- cbind.data.frame(cons.aln, ref.pos = NA, aln.pos=NA)
    cons.aln$ref.pos[cons.aln$ref!="-"] <- 1:nchar(ref.seq)
    cons.aln$aln.pos[cons.aln$cons!="-"] <- .get_aln_pos(gapped.con)
    list(cons.aln,aln.ref)
}

.msd_calculator <- function(check, learned){
    if(!is.list(check)){
        check <- list(check)
    }
    df1 <- lapply(learned, data.frame)
    df2 <- lapply(check, data.frame)
    merged <- mapply(merge, x=df1, y=df2, all=TRUE, by.x=names(df1[[1]])[1], 
                     by.y=names(df2[[1]])[1], SIMPLIFY = FALSE)
    merged <- lapply(merged, function(x){x[is.na(x)]<-0; x})
    msd <- sapply(merged, function(x)mean((x[,2]-x[,3])^2))
    msd
}

.proportion_calculator <- function(vec){
    table(vec)/length(vec)
}
<file_sep>/demultiplexer/ad_hoc_demultiplexer.R
library(stringr)
library(sarlacc)

demultiplex <- function(adap.stat, barcode, position=NULL){
    search_region <- umiExtract(adap.stat, position=position)
    aln.pid <- mclapply(barcode, internal_align, search_region=search_region, mc.cores=6L)
    aln.pid <- do.call(cbind, aln.pid)
    # replace all gapped alignment
    aln.pid[is.nan(aln.pid)] <- 0
    
    # Ambigurous alignment
    sorted.pid <- t(apply(aln.pid, 1, sort, decreasing = TRUE, na.last=TRUE))
    vague <- (sorted.pid[,1]-sorted.pid[,2])<0.05
    
    # sufficient match
    good.match <- sorted.pid[,1]>=0.7
    
    # Choose final barcode
    chosen.bc <- unlist(apply(aln.pid, 1, which.max))
    chosen.bc[c(vague|!good.match)] <- NA
    return(chosen.bc)
}

# Store parameter for alignment
internal_align <- function(search_region, x){
    aln <- pairwiseAlignment(search_region, x, type="local-global", gapOpening=5, gapExtension=1)
    nmatch(aln)/(str_count(subject(aln), "-")+nchar(x))
}
<file_sep>/homopolymerFixer/temp_consensus.R
nucleotides <- c("A","C","G","T")
QUALFUN <- function(alignments, qualities, min.coverage, gapped=TRUE) {
    aln.mat <- do.call(rbind, strsplit(as.character(alignments), ""))
    
    # Forming a quality matrix.
    qual.mat <- matrix(NA_integer_, nrow(aln.mat), ncol(aln.mat))
    for (i in seq_len(nrow(aln.mat))) { 
        current <- aln.mat[i,]
        chosen <- which(current!="-")
        qual.mat[i,chosen] <- qualities[[i]]
    }
    
    # Subsetting based on the minimum coverage.
    keep1 <- colSums(aln.mat!="-")/nrow(aln.mat) >= min.coverage
    aln.pos <- which(keep1)
    aln.mat <- aln.mat[,keep1,drop=FALSE]
    qual.mat <- qual.mat[,keep1,drop=FALSE]
    

    # Running through the alignment matrix.
    consensus <- character(ncol(aln.mat))
    out.err <- numeric(ncol(aln.mat))
    for (i in seq_len(ncol(aln.mat))) { 
        curcol <- aln.mat[,i]
        keep2 <- curcol %in% nucleotides
        curcol <- curcol[keep2]
        
        curqual <- qual.mat[keep2,i]
        correct <- log(1 - curqual)
        incorrect <- log(curqual/3)
        
        # Computing the probability of the current position being a particular base. 
        all.probs <- c(A=sum(correct[curcol=="A"]) + sum(incorrect[curcol!="A"]),
                       C=sum(correct[curcol=="C"]) + sum(incorrect[curcol!="C"]),
                       G=sum(correct[curcol=="G"]) + sum(incorrect[curcol!="G"]),
                       T=sum(correct[curcol=="T"]) + sum(incorrect[curcol!="T"]))
        all.probs <- exp(all.probs)
        all.probs <- all.probs/sum(all.probs)
        
        chosen <- which.max(all.probs)
        consensus[i] <- names(all.probs)[chosen]
        out.err[i] <- log(sum(all.probs[-chosen]))  # not using 1-chosen for numerical precision purposes.
    }
    if(gapped){
        gapped.consensus <- rep("-",length(keep1))
        gapped.consensus[keep1] <- consensus
        consensus <- gapped.consensus
    }
    
    return(list(seq=paste(consensus, collapse=""), qual=out.err))
}
<file_sep>/homopolymerFixer/homopolymerFixer_test.R
library(ShortRead)
library(sarlacc)
#ERCC sequence
ERCC_fasta <- readFasta("ERCC92.fa")
names(ERCC_fasta@sread) <- ERCC_fasta@id
ERCC_fasta <- ERCC_fasta@sread
ref <- ERCC_fasta
# Remove polyA tail
vec <- strsplit(as.character(ref),split="")
polyA.len <- sapply(vec, function(x)tail(rle(x)$lengths,1))
ref <- subseq(ref, start=1, end=(lengths(ref)-polyA.len))
reads <- readRDS("no.polyA.ercc.rds")

# Separate reads into train and test set
set.seed(100)
train.set <- sample(seq_along(reads), floor(length(reads)/2))
test.set <- seq_along(reads)[-train.set]
train.reads <- reads[train.set]
test.reads <- reads[test.set]
train.reads <- lapply(names(ref),function(x)train.reads[grep(x, names(train.reads))])
test.reads <- lapply(names(ref),function(x)test.reads[grep(x, names(test.reads))])
names(test.reads) <- names(train.reads) <- names(ref)

homo.correction.check <- function(test){
    test.match <- lapply(test, function(x)lapply(seq_along(x), function(y){
        homopolymerMatcher(x[y])
    }))
    original.homo <- unlist(IRangesList(lapply(test.match,function(x)x[[1]])))
    fixed.homo <- unlist(IRangesList(lapply(test.match,function(x)x[[2]])))
    fixed.obs.len <- as.numeric(unlist(elementMetadata(fixed.homo)$observed))
    original.obs.len <- as.numeric(unlist(elementMetadata(original.homo)$observed))
    true.len <- width(original.homo)
    fixed.tb <- aggregate(fixed.obs.len, list(true.len), table)
    original.tb <- aggregate(original.obs.len, list(true.len), table)
    fixed.mean <- aggregate(fixed.obs.len, list(true.len), mean)
    original.mean <- aggregate(original.obs.len, list(true.len), mean)
    fixed.diff <- aggregate((true.len-fixed.obs.len), list(true.len), table)
    original.diff <- aggregate((true.len-original.obs.len), list(true.len), table)
    fixed.mean.diff <- aggregate(abs(true.len-fixed.obs.len), list(true.len), mean)
    original.mean.diff <- aggregate(abs(true.len-original.obs.len), list(true.len), mean)
    mean.length <- cbind(original=original.mean, fixed=fixed.mean[,2, drop=FALSE])
    diff.length <- cbind(original=original.diff, fixed=fixed.diff[,2, drop=FALSE], 
                         original.mean=original.mean.diff[,2, drop=FALSE], fixed.mean=fixed.mean.diff[,2, drop=FALSE])
    colnames(mean.length) <- c("homo", "original.length", "fixed.length")
    colnames(diff.length) <- c("homo","original","fixed","original.mean.diff","fixed.mean.diff")
    list(mean=mean.length,diff=diff.length)
}

# Validate in the test set
test.reads <- test.reads[lengths(test.reads)>50]

validate.test <- function(learned.distr , sample.size){
    sample.size <- as.integer(sample.size)
    rep.time <- 20
    max.rep <- pmin(rep.time, floor(lengths(test.reads)/sample.size))
    
    sampled.reads <- mapply(function(x,y,z)replicate(z,sample(x,y)),
                            x=test.reads, y=sample.size, z=max.rep, SIMPLIFY = FALSE)
    ref.names <- rep(names(sampled.reads),lengths(sampled.reads))
    sampled.reads <- unlist(sampled.reads, recursive = FALSE)
    
    test.msa <- bplapply(sampled.reads, multiReadAlign, BPPARAM = MulticoreParam(30L))
    
    base.test <- bpmapply(function(x,y){
        fixed <- homopolymerFixer(learned.distr$base, x, threshold.p=0.5)
        aln <- pairwiseAlignment(c(fixed$degapped, fixed$fixed), 
                                 ref[y],gapOpening=5, gapExtension=1)
    }, x=test.msa, y=ref.names, BPPARAM=MulticoreParam(30L))
    
    debase.test <- bpmapply(function(x,y){
        fixed <- homopolymerFixer(learned.distr$debase, x, threshold.p=0.5)
        aln <- pairwiseAlignment(c(fixed$degapped, fixed$fixed), 
                                 ref[y],gapOpening=5, gapExtension=1)
    }, x=test.msa, y=ref.names, BPPARAM=MulticoreParam(30L))
    list(base=base.test, debase=debase.test)
}

all.learned.distr <- readRDS("combined.learned.distr.rds")
all.size <- names(all.learned.distr)
system.time(all.validate <- lapply(all.size, function(x)validate.test(all.learned.distr[[x]], x)))<file_sep>/demultiplexer/demultiplex.R
# If the barcodes are on the adaptor2, the input barcode need to be reverse complmeneted.
reads <- readQualityScaledDNAStringSet("fastq/testset.fastq")
source("ad_hoc_demultiplexer.R")
set.seed(119)
reads <- sample(reads,50000)
adaptor1 <- "GTCTCGTGGGCTCGGAGATGTGTATAAGAGACAGNNNNNNNNNNNNNNNNNNNNNNNNNNNNNNNNNNNNNNNNNNNNNNACTGGCCGTCGTTTTACATGGCGTAGCGGGTTCGAGCGCACCGCAGGGTATCCGGCTATTTTTTTTTTTTTTT"
adaptor2 <- "TCGTCGGCAGCGTCAGATGTGTATAAGAGACAGGG"
barcode <- read.csv("barcode.csv",header = TRUE,stringsAsFactors = FALSE)
barcode <- DNAStringSet(barcode$seq)
adap.aln <- adaptorAlign(reads, adaptor1 = adaptor1, adaptor2 = adaptor2, tolerance=300,BPPARAM = MulticoreParam(4L))
demulti <- demultiplex(adap.aln$adaptor1, barcode)