305 lines
10 KiB
Python
305 lines
10 KiB
Python
|
#!/util/bin/python
|
||
|
|
#
|
||
|
|
#
|
||
|
|
# This file really needs to be cleaned up and functions need to be improved
|
||
|
|
#
|
||
|
|
# (1) utility operations should be moved to a utility file
|
||
|
|
# (2) Need to add support for windows paired reads, information about where
|
||
|
|
# reads come from (which file), and n-base pair loci
|
||
|
|
# (3) Move to C++ based SAM/BAM i/o system
|
||
|
|
# (4) Figure out a way to get the reference to stream, and use a automatically pickled version
|
||
|
|
# (5) Add capability to walk over all reference bases, regardless of coverage. Currently
|
||
|
|
# we only walk over covered bases.
|
||
|
|
#
|
||
|
|
#
|
||
|
|
from optparse import OptionParser
|
||
|
|
import os
|
||
|
|
import sys
|
||
|
|
import random
|
||
|
|
import string
|
||
|
|
import os.path
|
||
|
|
import heapq
|
||
|
|
from itertools import *
|
||
|
|
|
||
|
|
from SAM import *
|
||
|
|
import SimpleSAM
|
||
|
|
#import SimulateReads # utility functions should be imported
|
||
|
|
|
||
|
|
# Global variable containing command line arguments
|
||
|
|
OPTIONS = None
|
||
|
|
|
||
|
|
#
|
||
|
|
# Trivial class that added pushback support to a stream
|
||
|
|
#
|
||
|
|
class peakStream(object):
|
||
|
|
"""A simple wrapper around a stream that adds unget operation"""
|
||
|
|
def __init__( self, stream ):
|
||
|
|
self.ungot = []
|
||
|
|
self.s = stream
|
||
|
|
|
||
|
|
def unget( self, elt ):
|
||
|
|
"""Put elt at the front of the underlying stream, making the next get call return it"""
|
||
|
|
|
||
|
|
#print 'unget', self.ungot
|
||
|
|
self.ungot.append(elt)
|
||
|
|
#print ' -> unget', self.ungot
|
||
|
|
|
||
|
|
# we hold the next() operation
|
||
|
|
def __iter__( self ): return self
|
||
|
|
|
||
|
|
def next(self):
|
||
|
|
"""For python iterator support"""
|
||
|
|
#print 'ungot', self.ungot
|
||
|
|
if self.ungot <> []:
|
||
|
|
elt = self.ungot.pop(0)
|
||
|
|
else:
|
||
|
|
elt = self.s.next()
|
||
|
|
|
||
|
|
#print ' -> ungot', self.ungot
|
||
|
|
#print ' elt', elt
|
||
|
|
return elt
|
||
|
|
|
||
|
|
#
|
||
|
|
# Returns a single, combined stream of SAM records in chromosome order from a list of
|
||
|
|
# input streams, each in reference order
|
||
|
|
#
|
||
|
|
def rawReadStream( ref, inputs ):
|
||
|
|
"""Returns a single iterable that returns records in reference order
|
||
|
|
from each of the input streams"""
|
||
|
|
def includePriorities( stream ):
|
||
|
|
return imap( lambda x: (x.getPos(), x), stream )
|
||
|
|
def prunePriorities( stream ):
|
||
|
|
return imap( lambda p: SimpleSAM.MappedReadFromSamRecord(ref, p[1]), stream )
|
||
|
|
|
||
|
|
with_priorities = map( includePriorities, inputs )
|
||
|
|
return prunePriorities( heapq.merge( *with_priorities ) )
|
||
|
|
|
||
|
|
#
|
||
|
|
# Just wraps the raw read stream objects in a list as they come out
|
||
|
|
#
|
||
|
|
def readStream( ref, inputs ):
|
||
|
|
"""Returns a single iterable that returns SAM records in reference order
|
||
|
|
from each of the input streams"""
|
||
|
|
rs = rawReadStream( ref, inputs )
|
||
|
|
return imap( lambda x: [x], rs )
|
||
|
|
|
||
|
|
#
|
||
|
|
# More complex stream object. Takes a list of input streams and creates a stream
|
||
|
|
# returning successive loci covered by the reads in the combined input stream.
|
||
|
|
#
|
||
|
|
class lociStream():
|
||
|
|
"""A streaming data structure that returns reads spanning each covered loci in
|
||
|
|
the input reads, offsets into them
|
||
|
|
where the bases are equivalent, and the position of the locus.
|
||
|
|
"""
|
||
|
|
def __init__( self, ref, inputs ):
|
||
|
|
self.ref = ref
|
||
|
|
self.rs = peakStream(rawReadStream( ref, inputs ))
|
||
|
|
self.window = []
|
||
|
|
|
||
|
|
self.chr = None
|
||
|
|
self.pos = None
|
||
|
|
|
||
|
|
def __iter__(self):
|
||
|
|
return self
|
||
|
|
|
||
|
|
def pushRead( self, read ):
|
||
|
|
self.window.append(read)
|
||
|
|
|
||
|
|
def cleanWindow( self, chr, pos ):
|
||
|
|
"""Walk over the window of reads, deleting those who no longer cover the current pos"""
|
||
|
|
if OPTIONS.debug:
|
||
|
|
print 'cleanWindow start:', chr, pos, self.window
|
||
|
|
|
||
|
|
def keepReadP( read ):
|
||
|
|
return read.chr == chr and pos >= read.start and pos <= read.end
|
||
|
|
self.window = filter( keepReadP, self.window )
|
||
|
|
|
||
|
|
if OPTIONS.debug:
|
||
|
|
print 'cleanWindow stop:', chr, pos, self.window
|
||
|
|
|
||
|
|
def expandWindow( self, chr, pos ):
|
||
|
|
"""Keep reading from the read stream until we've added all reads from the stream covering pos"""
|
||
|
|
if OPTIONS.debug:
|
||
|
|
print 'expandWindow start:', chr, pos, self.window
|
||
|
|
for read in self.rs:
|
||
|
|
#print 'read', read, pos
|
||
|
|
if read.chr == chr and read.start <= pos and read.end >= pos:
|
||
|
|
self.pushRead(read)
|
||
|
|
else:
|
||
|
|
self.rs.unget( read )
|
||
|
|
#self.rs = chain( [read], self.rs )
|
||
|
|
break
|
||
|
|
if OPTIONS.debug:
|
||
|
|
print 'expandWindow stop:', chr, pos, self.window
|
||
|
|
|
||
|
|
|
||
|
|
#
|
||
|
|
# This is the workhorse routine. It constructs a window of reads covering the
|
||
|
|
# next locus in the genome, and returns the reference, the contig, its position
|
||
|
|
# in the genome, along with a vector of offsets into the list of reads that denote
|
||
|
|
# the equivalent base among all the reads and reference.
|
||
|
|
#
|
||
|
|
def next(self):
|
||
|
|
if self.pos <> None:
|
||
|
|
self.cleanWindow(self.chr, self.pos) # consume reads not covering pos
|
||
|
|
self.expandWindow(self.chr, self.pos) # add reads to window covering pos
|
||
|
|
|
||
|
|
if self.window == []:
|
||
|
|
# the window is empty, we need to jump to the first pos of the first read in the stream:
|
||
|
|
nextRead = self.rs.next()
|
||
|
|
self.pushRead( nextRead )
|
||
|
|
self.chr = nextRead.chr
|
||
|
|
self.pos = nextRead.start
|
||
|
|
return self.next()
|
||
|
|
else:
|
||
|
|
# at this point, window contains all reads covering the pos, we need to return them
|
||
|
|
# and the offsets into each read for this loci
|
||
|
|
def calcOffset( read ):
|
||
|
|
offset = self.pos - read.start
|
||
|
|
return offset
|
||
|
|
|
||
|
|
offsets = map(calcOffset, self.window)
|
||
|
|
currPos = self.pos
|
||
|
|
self.pos += 1 # we are going to try to get the next position
|
||
|
|
return self.ref, self.chr, currPos, offsets, self.window
|
||
|
|
|
||
|
|
#
|
||
|
|
# Reference reader
|
||
|
|
#
|
||
|
|
def readRef(referenceFasta):
|
||
|
|
ref = {}
|
||
|
|
|
||
|
|
header = None
|
||
|
|
seq = ''
|
||
|
|
for line in open(referenceFasta):
|
||
|
|
if line[0] == '>':
|
||
|
|
if header <> None:
|
||
|
|
ref[header] = seq.lower()
|
||
|
|
seq = ''
|
||
|
|
header = line[1:].strip()
|
||
|
|
else:
|
||
|
|
seq += line.strip()
|
||
|
|
|
||
|
|
ref[header] = seq.lower()
|
||
|
|
#print ref
|
||
|
|
return ref
|
||
|
|
|
||
|
|
#
|
||
|
|
# Main() procedure
|
||
|
|
#
|
||
|
|
def main():
|
||
|
|
global OPTIONS, ROOT
|
||
|
|
|
||
|
|
# ------------------------------------------------------------
|
||
|
|
#
|
||
|
|
# Setup command line options
|
||
|
|
#
|
||
|
|
# ------------------------------------------------------------
|
||
|
|
usage = "usage: %prog [options] Walker [sam or bam file or file list]+"
|
||
|
|
parser = OptionParser(usage=usage)
|
||
|
|
parser.add_option("-r", "--ref", dest="reference",
|
||
|
|
type="string", default=None,
|
||
|
|
help="Reference DNA seqeunce in FASTA format")
|
||
|
|
parser.add_option("-m", "--maxRecords", dest="maxRecords",
|
||
|
|
type=int, default=None,
|
||
|
|
help="Max number of SAM records to process")
|
||
|
|
parser.add_option("-q", "--quiet", dest="quiet",
|
||
|
|
action='store_true', default=False,
|
||
|
|
help="Be quiet when generating output")
|
||
|
|
parser.add_option("-d", "--debug", dest="debug",
|
||
|
|
action='store_true', default=False,
|
||
|
|
help="Verbose debugging output?")
|
||
|
|
|
||
|
|
(OPTIONS, args) = parser.parse_args()
|
||
|
|
print 'args', args
|
||
|
|
if len(args) < 2:
|
||
|
|
parser.error("Incorrect number of arguments")
|
||
|
|
|
||
|
|
|
||
|
|
# ------------------------------------------------------------
|
||
|
|
#
|
||
|
|
# Dynamically load the walker class (the first cmd line arg)
|
||
|
|
# and initialize a walker class object
|
||
|
|
#
|
||
|
|
#
|
||
|
|
# load walkers from standard location
|
||
|
|
#
|
||
|
|
# ------------------------------------------------------------
|
||
|
|
walkerName = args[0]
|
||
|
|
walkerModule = __import__(walkerName, globals(), locals(), [], -1)
|
||
|
|
walkerClass = walkerModule.__dict__[walkerName]
|
||
|
|
walker = walkerClass()
|
||
|
|
|
||
|
|
print walkerName
|
||
|
|
print walkerModule
|
||
|
|
print walkerClass
|
||
|
|
print walker
|
||
|
|
|
||
|
|
print '------------------------------------------------------------'
|
||
|
|
print 'program:', sys.argv[0]
|
||
|
|
print ''
|
||
|
|
print ' ref: ', OPTIONS.reference
|
||
|
|
print ' walker: ', walker
|
||
|
|
if walker.hasOption('name'):
|
||
|
|
print ' -> ', walker.getOption('name')
|
||
|
|
if walker.hasOption('desc'):
|
||
|
|
print ' -> ', walker.getOption('desc')
|
||
|
|
print '------------------------------------------------------------'
|
||
|
|
|
||
|
|
|
||
|
|
# ------------------------------------------------------------
|
||
|
|
#
|
||
|
|
# Initialize the engine
|
||
|
|
#
|
||
|
|
# ------------------------------------------------------------
|
||
|
|
|
||
|
|
# read the reference
|
||
|
|
refs = readRef(OPTIONS.reference)
|
||
|
|
|
||
|
|
# create the low-level SAMIO streams
|
||
|
|
files = args[1:]
|
||
|
|
inputs = map( lambda file: SAMIO( file, debugging=OPTIONS.debug ), files )
|
||
|
|
|
||
|
|
# build the higher level stream object, either by record or by loci
|
||
|
|
if walker.isRecordWalker():
|
||
|
|
stream = readStream( refs, inputs )
|
||
|
|
if walker.isLociWalker():
|
||
|
|
stream = lociStream( refs, inputs )
|
||
|
|
|
||
|
|
# ------------------------------------------------------------
|
||
|
|
#
|
||
|
|
# Move the walker object over the stream
|
||
|
|
#
|
||
|
|
# For each element in the record or loci stream, invoke the walker
|
||
|
|
# object on it. Use filter, map, and reduce to construct the sum
|
||
|
|
#
|
||
|
|
# ------------------------------------------------------------
|
||
|
|
sum = walker.reduceDefault
|
||
|
|
counter = 0
|
||
|
|
for elt in stream:
|
||
|
|
counter += 1
|
||
|
|
#print 'processing elt', counter, elt, OPTIONS.maxRecords
|
||
|
|
if OPTIONS.maxRecords <> None:
|
||
|
|
if (counter * 10) % OPTIONS.maxRecords == 0:
|
||
|
|
print ' => ', (counter * 100.0) / OPTIONS.maxRecords, '% done'
|
||
|
|
|
||
|
|
if counter > OPTIONS.maxRecords:
|
||
|
|
break
|
||
|
|
if walker.filterfunc( *elt ):
|
||
|
|
x = walker.mapfunc( *elt )
|
||
|
|
sum = walker.reducefunc( x, sum )
|
||
|
|
print 'Result is', sum
|
||
|
|
|
||
|
|
if __name__ == "__main__":
|
||
|
|
if True:
|
||
|
|
import cProfile, pstats
|
||
|
|
cProfile.run('main()', 'prof')
|
||
|
|
stats = pstats.Stats('prof')
|
||
|
|
stats.sort_stats('time')
|
||
|
|
stats.print_stats(20)
|
||
|
|
else:
|
||
|
|
main()
|
||
|
|
|
||
|
|
|