class IpdModel:
"""
Predicts the IPD of an any context, possibly containing multiple modifications.
We use a 4^12 entry LUT to get the predictions for contexts without modifications,
then we use the GbmModel to get predictions in the presence of arbitrary mods.
Note on the coding scheme. For each contig we store a byte-array that has size = contig.length + 2*self.pad
The upper 4 bits contain a lookup into seqReverseMap, which can contains N's. This is used for giving
template snippets that may contains N's if the reference sequence does, or if the snippet
The lowe 4 bits contain a lookup into lutReverseMap, which
"""
def __init__(self, fastaRecords, modelFile=None, modelIterations=-1):
"""
Load the reference sequences and the ipd lut into shared arrays that can be
used as numpy arrays in worker processes.
fastaRecords is a list of FastaRecords, in the cmp.h5 file order
"""
self.pre = 10
self.post = 4
self.pad = 30
self.base4 = 4 ** np.array(range(self.pre + self.post + 1))
self.refDict = {}
self.refLengthDict = {}
for contig in fastaRecords:
if contig.id is None:
# This contig has no mapped reads -- skip it
continue
rawSeq = contig.sequence
refSeq = np.fromstring(rawSeq, dtype=byte)
# Store the reference length
self.refLengthDict[contig.id] = len(rawSeq)
# Make a shared array
sa = SharedArray(dtype='B', shape=len(rawSeq) + self.pad * 2)
saWrap = sa.getNumpyWrapper()
# Lut Codes convert Ns to As so that we don't put Ns into the Gbm Model
# Seq Codes leaves Ns as Ns for getting reference snippets out
innerLutCodes = lutCodeMap[refSeq]
innerSeqCodes = seqCodeMap[refSeq]
innerCodes = np.bitwise_or(innerLutCodes, np.left_shift(innerSeqCodes, 4))
saWrap[self.pad:(len(rawSeq) + self.pad)] = innerCodes
# Padding codes -- the lut array is padded with 0s the sequence array is padded with N's (4)
outerCodes = np.left_shift(np.ones(self.pad, dtype=uint8) * 4, 4)
saWrap[0:self.pad] = outerCodes
saWrap[(len(rawSeq) + self.pad):(len(rawSeq) + 2 * self.pad)] = outerCodes
self.refDict[contig.id] = sa
# No correction factor for IPDs everything is normalized to 1
self.meanIpd = 1
# Find and open the ipd model file
if modelFile:
self.lutPath = modelFile
else:
self.lutPath = os.path.dirname(os.path.abspath(__file__)) + os.path.sep + "kineticLut.h5"
if os.path.exists(self.lutPath):
h5File = h5py.File(self.lutPath, mode='r')
gbmModelGroup = h5File["/AllMods_GbmModel"]
self.gbmModel = GbmContextModel(gbmModelGroup, modelIterations)
# We always use the model -- no more LUTS
self.predictIpdFunc = self.predictIpdFuncModel
self.predictManyIpdFunc = self.predictManyIpdFuncModel
else:
logging.info("Couldn't find model file: %s" % self.lutPath)
def _loadIpdTable(self, nullModelGroup):
"""
Read the null kinetic model into a shared numpy array dataset
"""
nullModelDataset = nullModelGroup["KineticValues"]
# assert that the dataset is a uint8
assert(nullModelDataset.dtype == uint8)
# Construct a 'shared array' (a numpy wrapper around some shared memory
# Read the LUT into this table
self.sharedArray = SharedArray('B', nullModelDataset.shape[0])
lutArray = self.sharedArray.getNumpyWrapper()
nullModelDataset.read_direct(lutArray)
# Load the second-level LUT
self.floatLut = nullModelGroup["Lut"][:]
def refLength(self, refId):
return self.refLengthDict[refId]
def cognateBaseFunc(self, refId):
"""
Return a function that returns a snippet of the reference sequence around a given position
"""
# FIXME -- what is the correct strand to return?!
# FIXME -- what to do about padding when the snippet runs off the end of the reference
# how do we account for / indicate what is happening
refArray = self.refDict[refId].getNumpyWrapper()
def f(tplPos, tplStrand):
# skip over the padding
tplPos += self.pad
# Forward strand
if tplStrand == 0:
slc = refArray[tplPos]
slc = np.right_shift(slc, 4)
return seqMap[slc]
# Reverse strand
else:
slc = refArray[tplPos]
slc = np.right_shift(slc, 4)
return seqMapComplement[slc]
return f
def snippetFunc(self, refId, pre, post):
"""
Return a function that returns a snippet of the reference sequence around a given position
"""
refArray = self.refDict[refId].getNumpyWrapper()
def f(tplPos, tplStrand):
"""Closure for returning a reference snippet. The reference is padded with N's for bases falling outside the extents of the reference"""
# skip over the padding
tplPos += self.pad
# Forward strand
if tplStrand == 0:
slc = refArray[(tplPos - pre):(tplPos + 1 + post)]
slc = np.right_shift(slc, 4)
return seqMapNp[slc].tostring()
# Reverse strand
else:
slc = refArray[(tplPos + pre):(tplPos - post - 1):-1]
slc = np.right_shift(slc, 4)
return seqMapComplementNp[slc].tostring()
return f
def getReferenceWindow(self, refId, tplStrand, start, end):
"""
Return a snippet of the reference sequence
"""
refArray = self.refDict[refId].getNumpyWrapper()
# adjust position for reference padding
start += self.pad
end += self.pad
# Forward strand
if tplStrand == 0:
slc = refArray[start:end]
slc = np.right_shift(slc, 4)
return "".join(seqMap[x] for x in slc)
# Reverse strand
else:
slc = refArray[end:start:-1]
slc = np.right_shift(slc, 4)
return "".join(seqMapComplement[x] for x in slc)
def predictIpdFuncLut(self, refId):
"""
Each (pre+post+1) base context gets mapped to an integer
by converting each nucleotide to a base-4 number A=0, C=1, etc,
and treating the 'pre' end of the context of the least significant
digit. This code is used to lookup the expected IPD in a
pre-computed table. Contexts near the ends of the reference
are coded by padding the context with 0
"""
# Materialized the numpy wrapper around the shared data
refArray = self.refDict[refId].getNumpyWrapper()
lutArray = self.sharedArray.getNumpyWrapper()
floatLut = self.floatLut
def f(tplPos, tplStrand):
# skip over the padding
tplPos += self.pad
# Forward strand
if tplStrand == 0:
slc = np.bitwise_and(refArray[(tplPos + self.pre):(tplPos - self.post - 1):-1], 0xf)
# Reverse strand
else:
slc = 3 - np.bitwise_and(refArray[(tplPos - self.pre):(tplPos + 1 + self.post)], 0xf)
code = (self.base4 * slc).sum()
return floatLut[max(1, lutArray[code])]
return f
def predictIpdFuncModel(self, refId):
"""
Each (pre+post+1) base context gets mapped to an integer
by converting each nucleotide to a base-4 number A=0, C=1, etc,
and treating the 'pre' end of the context of the least significant
digit. This code is used to lookup the expected IPD in a
pre-computed table. Contexts near the ends of the reference
are coded by padding the context with 0
"""
# Materialized the numpy wrapper around the shared data
snipFunction = self.snippetFunc(refId, self.post, self.pre)
def f(tplPos, tplStrand):
# Get context string
context = snipFunction(tplPos, tplStrand)
# Get prediction
return self.gbmModel.getPredictions([context])[0]
return f
def predictManyIpdFuncModel(self, refId):
"""
Each (pre+post+1) base context gets mapped to an integer
by converting each nucleotide to a base-4 number A=0, C=1, etc,
and treating the 'pre' end of the context of the least significant
digit. This code is used to lookup the expected IPD in a
pre-computed table. Contexts near the ends of the reference
are coded by padding the context with 0
"""
# Materialized the numpy wrapper around the shared data
snipFunction = self.snippetFunc(refId, self.post, self.pre)
def fMany(sites):
contexts = [snipFunction(x[0], x[1]) for x in sites]
return self.gbmModel.getPredictions(contexts)
return fMany
def modPredictIpdFunc(self, refId, mod):
"""
Each (pre+post+1) base context gets mapped to an integer
by converting each nucleotide to a base-4 number A=0, C=1, etc,
and treating the 'pre' end of the context of the least significant
digit. This code is used to lookup the expected IPD in a
pre-computed table. Contexts near the ends of the reference
are coded by padding the context with 0
"""
refArray = self.refDict[refId].getNumpyWrapper()
def f(tplPos, relativeModPos, readStrand):
# skip over the padding
tplPos += self.pad
# Read sequence matches forward strand
if readStrand == 0:
slc = 3 - np.bitwise_and(refArray[(tplPos - self.pre):(tplPos + 1 + self.post)], 0xf)
# Reverse strand
else:
slc = np.bitwise_and(refArray[(tplPos + self.pre):(tplPos - self.post - 1):-1], 0xf)
# Modify the indicated position
slc[relativeModPos + self.pre] = baseToCode[mod]
slcString = "".join([codeToBase[x] for x in slc])
# Get the prediction for this context
# return self.gbmModel.getPredictions([slcString])[0]
return 0.0
return f
def __init__(self, fastaRecords, modelFile=None, modelIterations=-1):
"""
Load the reference sequences and the ipd lut into shared arrays that can be
used as numpy arrays in worker processes.
fastaRecords is a list of FastaRecords, in the cmp.h5 file order
"""
self.pre = 10
self.post = 4
self.pad = 30
self.base4 = 4 ** np.array(range(self.pre + self.post + 1))
self.refDict = {}
self.refLengthDict = {}
for contig in fastaRecords:
if contig.id is None:
# This contig has no mapped reads -- skip it
continue
rawSeq = contig.sequence
refSeq = np.fromstring(rawSeq, dtype=byte)
# Store the reference length
self.refLengthDict[contig.id] = len(rawSeq)
# Make a shared array
sa = SharedArray(dtype='B', shape=len(rawSeq) + self.pad * 2)
saWrap = sa.getNumpyWrapper()
# Lut Codes convert Ns to As so that we don't put Ns into the Gbm Model
# Seq Codes leaves Ns as Ns for getting reference snippets out
innerLutCodes = lutCodeMap[refSeq]
innerSeqCodes = seqCodeMap[refSeq]
innerCodes = np.bitwise_or(innerLutCodes, np.left_shift(innerSeqCodes, 4))
saWrap[self.pad:(len(rawSeq) + self.pad)] = innerCodes
# Padding codes -- the lut array is padded with 0s the sequence array is padded with N's (4)
outerCodes = np.left_shift(np.ones(self.pad, dtype=uint8) * 4, 4)
saWrap[0:self.pad] = outerCodes
saWrap[(len(rawSeq) + self.pad):(len(rawSeq) + 2 * self.pad)] = outerCodes
self.refDict[contig.id] = sa
# No correction factor for IPDs everything is normalized to 1
self.meanIpd = 1
# Find and open the ipd model file
if modelFile:
self.lutPath = modelFile
else:
self.lutPath = os.path.dirname(os.path.abspath(__file__)) + os.path.sep + "kineticLut.h5"
if os.path.exists(self.lutPath):
h5File = h5py.File(self.lutPath, mode='r')
gbmModelGroup = h5File["/AllMods_GbmModel"]
self.gbmModel = GbmContextModel(gbmModelGroup, modelIterations)
# We always use the model -- no more LUTS
self.predictIpdFunc = self.predictIpdFuncModel
self.predictManyIpdFunc = self.predictManyIpdFuncModel
else:
logging.info("Couldn't find model file: %s" % self.lutPath)
def __init__(self, fastaRecords, modelFile, modelIterations=-1):
"""
Load the reference sequences and the ipd lut into shared arrays that can be
used as numpy arrays in worker processes.
fastaRecords is a list of FastaRecords, in the alignments file order
"""
self.pre = 10
self.post = 4
self.pad = 30
self.base4 = 4**np.array(range(self.pre + self.post + 1))
self.refDict = {}
self.refLengthDict = {}
for contig in fastaRecords:
if contig.alignmentID is None:
# This contig has no mapped reads -- skip it
continue
rawSeq = contig.sequence[:]
refSeq = np.frombuffer(rawSeq.encode("utf-8"), dtype=byte)
# Store the reference length
self.refLengthDict[contig.alignmentID] = len(rawSeq)
# Make a shared array
sa = SharedArray(dtype='B', shape=len(rawSeq) + self.pad * 2)
saWrap = sa.getNumpyWrapper()
# Lut Codes convert Ns to As so that we don't put Ns into the Gbm Model
# Seq Codes leaves Ns as Ns for getting reference snippets out
innerLutCodes = lutCodeMap[refSeq]
innerSeqCodes = seqCodeMap[refSeq]
innerCodes = np.bitwise_or(innerLutCodes,
np.left_shift(innerSeqCodes, 4))
saWrap[self.pad:(len(rawSeq) + self.pad)] = innerCodes
# Padding codes -- the lut array is padded with 0s the sequence
# array is padded with N's (4)
outerCodes = np.left_shift(np.ones(self.pad, dtype=uint8) * 4, 4)
saWrap[0:self.pad] = outerCodes
saWrap[(len(rawSeq) + self.pad):(len(rawSeq) +
2 * self.pad)] = outerCodes
self.refDict[contig.alignmentID] = sa
# No correction factor for IPDs everything is normalized to 1
self.meanIpd = 1
# Find and open the ipd model file
self.lutPath = modelFile
if os.path.exists(self.lutPath):
with gzip.open(self.lutPath, "rb") as npz_in:
gbmModelData = np.load(npz_in, allow_pickle=True)
self.gbmModel = GbmContextModel(gbmModelData, modelIterations)
# We always use the model -- no more LUTS
self.predictIpdFunc = self.predictIpdFuncModel
self.predictManyIpdFunc = self.predictManyIpdFuncModel
else:
logging.info("Couldn't find model file: %s" % self.lutPath)