ppp_vcf_to_ima: vcf_reader_func.py comparison

comparison vcf_reader_func.py @ 0:1ad9c146ebb4 draft

Uploaded

author	jaredgk
date	Fri, 19 Oct 2018 14:36:02 -0400
parents
children

comparison

equal deleted inserted replaced

--1:000000000000
+:1ad9c146ebb4
+import sys
+import pysam
+import logging
+import struct
+from random import sample
+from collections import defaultdict
+import os
+import gzip
+def checkIfGzip(filename):
+try:
+gf = gzip.open(filename)
+gl = gf.readline()
+gf.close()
+vcf_check = b'##fileformat=VCF'
+if gl[0:3] == b'BCF':
+return 'bcf'
+elif gl[:len(vcf_check)] == vcf_check:
+return checkHeader(filename)
+else:
+return 'other'
+except:
+return 'nozip'
+def checkHeader(filename):
+f = open(filename,'rb')
+l = f.readline()
+f.close()
+BGZF_HEADER=b'\x1f\x8b\x08\x04\x00\x00\x00\x00\x00\xff\x06\x00\x42\x43\x02\x00'
+#BGZF_HEADER=b'\x1f\x8b\x08\x04\x00\x00\x00\x00\x00\xff'
+GZF_HEADER=b'\x1f\x8b'
+if l[:len(BGZF_HEADER)] == BGZF_HEADER:
+return 'bgzip'
+if l[:len(GZF_HEADER)] == GZF_HEADER:
+return 'gzip'
+return 'nozip'
+def checkFormat(vcfname):
+"""Checks header of given file for compression type
+Given a filename, opens file and reads first line to check if
+file has BGZF or GZIP header. May be extended to check for BCF format
+Parameters
+----------
+filename : str
+Name of file to be checked
+Returns
+-------
+extension : str {'bgzip','gzip','vcf','other'}
+File extension as indicated by header
+"""
+typ = checkIfGzip(vcfname)
+if typ != 'nozip':
+return typ
+f = open(vcfname)
+l = f.readline()
+f.close()
+VCF_TAG='##fileformat=VCF'
+if l[:len(VCF_TAG)] == VCF_TAG:
+return 'vcf'
+return 'other'
+def checkIfCpG(record,fasta_ref,offset=0,add_chr=False):
+dr = None
+pos = record.pos
+c = record.chrom
+if record.alts is None:
+return False
+if add_chr:
+c = 'chr'+record.chrom
+if record.ref == 'C' and 'T' in record.alts:
+seq = fasta_ref.fetch(c,pos-1,pos+1)
+if seq[0].upper() != 'C':
+logging.warning('%s %d has bad base %s' % (record.chrom,record.pos,seq[0]))
+#raise Exception("checkIfCpG function not lining up properly")
+if seq[1].upper() == 'G':
+return True
+return False
+elif record.ref == 'G' and 'A' in record.alts:
+seq = fasta_ref.fetch(c,pos-2,pos)
+if seq[1].upper() != 'G':
+logging.warning('%s %d has bad base %s' % (record.chrom,record.pos,seq[1]))
+#raise Exception("checkIfCpg function not lining up on negative strand")
+if seq[0].upper() == 'C':
+return True
+return False
+return False
+def checkForDuplicates(rec_list,pass_list):
+for i in range(len(rec_list)-1):
+if rec_list[i].pos == rec_list[i+1].pos:
+pass_list[i] = False
+pass_list[i+1] = False
+def checkForMultiallele(rec_list,pass_list):
+for i in range(len(rec_list)):
+if i != len(rec_list)-1 and rec_list[i].pos == rec_list[i+1].pos:
+pass_list[i] = False
+pass_list[i+1] = False
+if len(rec_list[i].alleles) > 2:
+pass_list[i] = False
+def flipChrom(chrom):
+if chrom[0:3] == 'chr':
+return chrom[0:3]
+return 'chr'+chrom
+def getAlleleCountDict(rec):
+alleles = defaultdict(int)
+total_sites = 0
+missing_inds = 0
+for j in range(len(rec.samples)):
+samp = rec.samples[j]
+if None in samp.alleles:
+missing_inds += 1
+for k in range(len(samp.alleles)):
+b = samp.alleles[k]
+if b is not None:
+alleles[b] += 1
+total_sites+=1
+return alleles, total_sites, missing_inds
+def isInvariant(rec):
+alleles, total_sites, missing_inds = getAlleleCountDict(rec)
+if len(alleles) == 1:
+return True
+return False
+def isInformative(rec, mincount=2, alleles=None):
+count = 0
+if alleles is None:
+alleles, total_sites, missing_inds = getAlleleCountDict(rec)
+if len(alleles) != 2:
+return False
+i1,i2 = alleles.keys()
+return (alleles[i1] >= mincount and alleles[i2] >= mincount)
+def getPassSites(record_list, remove_cpg=False, remove_indels=True,
+remove_multiallele=True, remove_missing=0,
+inform_level=2, fasta_ref=None):
+pass_list = [True for r in record_list]
+if remove_cpg == True and fasta_ref is None:
+raise Exception("CpG removal requires a reference")
+if inform_level > 2 or inform_level < 0:
+raise Exception("Inform level %d must be between 0 and 2" % inform_level)
+if remove_multiallele:
+checkForMultiallele(record_list,pass_list)
+for i in range(len(record_list)):
+rec = record_list[i]
+if remove_indels and not checkRecordIsSnp(rec):
+pass_list[i] = False
+if remove_cpg and checkIfCpG(rec,fasta_ref):
+pass_list[i] = False
+alleles,total_sites,missing_inds = getAlleleCountDict(rec)
+if remove_missing != -1 and missing_inds > int(remove_missing):
+pass_list[i] = False
+if inform_level != 0 and not isInformative(rec,mincount=inform_level,alleles=alleles):
+pass_list[i] = False
+#pp = zip([r.pos for r in record_list],pass_list)
+#for ppp in pp:
+#    logging.info(ppp)
+return pass_list
+def filterSites(record_list, remove_cpg=False, remove_indels=True,
+remove_multiallele=True, remove_missing=0, inform_level=2,
+fasta_ref=None):
+pass_list = getPassSites(record_list,remove_cpg,remove_indels,remove_multiallele,remove_missing,inform_level,fasta_ref)
+out_list = []
+for i in range(len(pass_list)):
+if pass_list[i]:
+out_list.append(record_list[i])
+return out_list
+class VcfReader():
+def __init__(self, vcfname, compress_flag=False, subsamp_num=None,
+subsamp_fn=None, subsamp_list=None, index=None, popmodel=None, use_allpop=False):
+ext = checkFormat(vcfname)
+if ext in ['gzip','other'] :
+raise Exception(('Input file %s is gzip-formatted, must be either '
+'uncompressed or zipped with bgzip' % vcfname))
+self.file_uncompressed = (ext == 'vcf')
+self.reader_uncompressed = (self.file_uncompressed and not compress_flag)
+self.popmodel = None
+self.popkeys = None
+if popmodel is not None and use_allpop:
+raise Exception("Popmodel and allpop cannot both be specified")
+if compress_flag and file_uncompressed:
+vcfname = compressVcf(vcfname)
+if subsamp_num is not None:
+if subsamp_list is not None:
+raise Exception('Multiple subsampling options called in getVcfReader')
+subsamp_list = getSubsampleList(vcfname, subsamp_num)
+elif subsamp_fn is not None:
+if subsamp_list is not None:
+raise Exception('Multiple subsampling options called in getVcfReader')
+subsamp_file = open(subsamp_fn,'r')
+subsamp_list = [l.strip() for l in subsamp_file.readlines()]
+subsamp_file.close()
+if index is None:
+self.reader = pysam.VariantFile(vcfname)
+else:
+self.reader = pysam.VariantFile(vcfname, index_filename=index)
+if popmodel is not None:
+self.popmodel = popmodel
+popsamp_list = popmodel.inds
+self.reader.subset_samples(popsamp_list)
+self.setPopIdx()
+if use_allpop:
+self.setAllPop()
+if subsamp_list is not None:
+logging.debug('Subsampling %d individuals from VCF file' %
+(len(subsamp_list)))
+self.reader.subset_samples(subsamp_list)
+self.prev_last_rec = next(self.reader)
+self.chr_in_chrom = (self.prev_last_rec.chrom[0:3] == 'chr')
+def fetch(self, chrom=None, start=None, end=None):
+return self.reader.fetch(chrom, start, end)
+def getRecordList(self, region=None, chrom=None, start=None,
+end=None):
+if self.reader_uncompressed:
+ret, self.prev_last_rec = getRecordListUnzipped(self.reader, self.prev_last_rec, region, add_chr=self.chr_in_chrom)
+return ret
+else:
+return getRecordList(self.reader, region, chrom, start, end, self.chr_in_chrom)
+def setPopIdx(self):
+self.popkeys = {}
+sample_names = [l for l in self.reader.header.samples]
+for p in self.popmodel.pop_list:
+self.popkeys[p] = []
+for ind in self.popmodel.ind_dict[p]:
+self.popkeys[p].append(sample_names.index(ind))
+def close(self):
+self.reader.close()
+def setAllPop(self):
+self.popkeys = {'ALL':[]}
+for i in range(len(self.reader.header.samples)):
+self.popkeys['ALL'].append(i)
+def modChrom(c,vcf_chr):
+if c is None:
+return None
+if vcf_chr and c[:3] != 'chr':
+return 'chr'+c
+if not vcf_chr and c[:3] == 'chr':
+return c[3:]
+return c
+def getRecordList(vcf_reader, region=None, chrom=None, start=None,
+end=None, add_chr=False):
+"""Returns list for use in subsampling from input file"""
+if region is not None:
+c = modChrom(region.chrom,add_chr)
+var_sites = vcf_reader.fetch(c, region.start, region.end)
+else:
+c = modChrom(chrom,add_chr)
+var_sites = vcf_reader.fetch(c, start, end)
+lst = []
+for rec in var_sites:
+lst.append(rec)
+return lst
+def getRecordListUnzipped(vcf_reader, prev_last_rec, region=None, chrom=None,
+start=None, end=None, add_chr=False):
+"""Method for getting record list from unzipped VCF file.
+This method will sequentially look through a VCF file until it finds
+the given `start` position on `chrom`, then add all records to a list
+until the `end` position has been reached. Note that `prev_last_rec`
+must be kept track of externally to ensure that if consecutive regions
+are called, the record of the first variant outside the first region
+is not lost.
+Parameters
+----------
+vcf_reader : pysam VariantFile object
+VCF reader initialized from other function
+region : region object
+Region object with start and end coordinates of region of interest.
+prev_last_rec : VariantRecord object
+Variable with last record read from VcfReader. Stored here so that
+if two adjacent regions are called, the overflow record from the
+first region is still included in the next region
+Returns
+-------
+lst : list
+List of records in given gene region
+prev_last_rec : VariantRecord object
+First record after target region, for use in next call
+"""
+lst = []
+if region is None:
+lst.append(prev_last_rec)
+for rec in vcf_reader:
+lst.append(rec)
+return lst, lst[-1]
+if (prev_last_rec is not None and
+region.containsRecord(prev_last_rec) == 'in'):
+lst.append(prev_last_rec)
+elif (prev_last_rec is not None and
+region.containsRecord(prev_last_rec) == 'after'):
+return []
+rec = next(vcf_reader,None)
+if rec is None:
+return lst,None
+place = region.containsRecord(rec)
+while rec is not None and place != 'after':
+if place == 'in':
+lst.append(rec)
+rec = next(vcf_reader,None)
+if rec is None:
+break
+place = region.containsRecord(rec)
+prev_last_rec = rec
+return lst, prev_last_rec
+def checkRecordIsSnp(rec):
+"""Checks if this record is a single nucleotide variant, returns bool."""
+logging.info(str(rec.pos))
+if len(rec.ref) != 1:
+return False
+if rec.alts is None:
+return False
+for allele in rec.alts:
+if len(allele) != 1:
+return False
+return True
+def getSubsampleList(vcfname, ss_count):
+"""Returns a list of the first `ss_count` individuals in `vcfname`
+"""
+vcf_o = pysam.VariantFile(vcfname)
+rec = next(vcf_o)
+vcf_o.close()
+lst = []
+for samp in rec.samples:
+lst.append(samp)
+return lst[:int(ss_count)]
+def compressVcf(vcfname,forceflag=False,remove=False):
+"""Runs bgzip and tabix on input VCF file.
+Using the pysam library, this function runs the bgzip and tabix utilities
+on the given input file. By default, this will not overwrite an existing
+zipped file, but will overwrite an existing index. `remove` can be set to
+delete the unzipped file.
+Parameters
+----------
+vcfname : str
+Name of uncompressed VCF file
+forceflag : bool (False)
+If true, will overwrite (vcfname).gz if it exists
+remove : bool (False)
+If true, will delete uncompressed source file
+Returns
+-------
+cvcfname : str
+Filepath to compressed VCF file
+"""
+cvcfname = vcfname+".gz"
+pysam.tabix_compress(vcfname,cvcfname,force=forceflag)
+pysam.tabix_index(cvcfname,preset="vcf",force=True)
+if remove:
+os.remove(vcfname)
+return cvcfname
+def vcfRegionName(prefix, region, ext, oneidx=False,
+halfopen=True, sep='-'):
+chrom = region.toStr(halfopen, oneidx, sep)
+return prefix+'_'+chrom+'.'+ext
+def getRecordsInRegion(region, record_list):
+sub_list = []
+for i in range(len(record_list)):
+loc = region.containsRecord(record_list[i])
+if loc == "in":
+sub_list.append(record_list[i])
+elif loc == "after":
+break
+return sub_list
+#def getVcfReader(args):
+def getVcfReader(vcfname, compress_flag=False, subsamp_num=None,
+subsamp_fn=None, subsamp_list=None, index=None):
+"""Returns a reader for a given input VCF file.
+Given a filename, filetype, compression option, and optional Subsampling
+options, will return a pysam.VariantFile object for iteration and
+a flag as to whether this file is compressed or uncompressed.
+Parameters
+----------
+vcfname : str
+Filename for VCF file. The extension of this file will be used to
+determine whether it is compressed or not unless `var_ext` is set.
+var_ext : str (None)
+Extension for VCF file if it is not included in the filename.
+compress_flag : bool (False)
+If filetype is uncompressed and this is set to true, will run
+compressVcf function.
+subsamp_num : int (None)
+If set, will randomly select `subsamp_num` individuals (not
+genotypes) from the input VCF file and return a reader with
+only those data.
+subsamp_fn : str (None)
+If set, will return a reader with only data from the samples listed
+in the file provided. Cannot be used with other subsampling options.
+subsamp_list : list (None)
+If set, will return reader with records containing only
+individuals named in the list. Cannot be used with other subsampling
+options.
+Returns
+-------
+vcf_reader : pysam.VariantFile
+A reader that can be iterated through for variant records. If
+compressed, it will be able to use the pysam fetch method, otherwise
+it must be read through sequentially
+reader_uncompressed : bool
+If True, VCF reader is uncompressed. This means the fetch method
+cannot be used and region access must be done using the
+"getRecordListUnzipped" method.
+"""
+ext = checkFormat(vcfname)
+if ext in ['gzip','other'] :
+raise Exception(('Input file %s is gzip-formatted, must be either '
+'uncompressed or zipped with bgzip' % vcfname))
+file_uncompressed = (ext == 'vcf')
+reader_uncompressed = (file_uncompressed and not compress_flag)
+if compress_flag and file_uncompressed:
+vcfname = compressVcf(vcfname)
+#subsamp_list = None
+if subsamp_num is not None:
+if subsamp_list is not None:
+raise Exception('Multiple subsampling options called in getVcfReader')
+subsamp_list = getSubsampleList(vcfname, subsamp_num)
+elif subsamp_fn is not None:
+if subsamp_list is not None:
+raise Exception('Multiple subsampling options called in getVcfReader')
+subsamp_file = open(subsamp_fn,'r')
+subsamp_list = [l.strip() for l in subsamp_file.readlines()]
+subsamp_file.close()
+if index is None:
+vcf_reader = pysam.VariantFile(vcfname)
+else:
+vcf_reader = pysam.VariantFile(vcfname, index_filename=index)
+if subsamp_list is not None:
+logging.debug('Subsampling %d individuals from VCF file' %
+(len(subsamp_list)))
+vcf_reader.subset_samples(subsamp_list)
+return vcf_reader, reader_uncompressed

Mercurial > repos > jaredgk > ppp_vcf_to_ima

comparison vcf_reader_func.py @ 0:1ad9c146ebb4 draft