chkSNPrmask.py

   1 try:
   2     import psyco
   3     psyco.full()
   4 except:
   5     pass
   6
   7 import sys
   8 import tempfile
   9 import shutil
  10 import os
  11 import optparse
  12 import sqlite3 as sqlite
  13 from commoncode import getConfigParser, getConfigOption, getConfigBoolOption
  14
  15 configParser = getConfigParser()
  16 cisTemp = getConfigOption(configParser, "general", "cistematic_temp", default="/tmp")
  17 tempfile.tempdir = cisTemp
  18
  19 print "chkSNPrmask: version 3.4"
  20
  21
  22 def main(argv=None):
  23     if not argv:
  24         argv = sys.argv
  25
  26     usage = "usage: python %s dbfile snpsfile nr_snps_outfile [--cache numPages] [--repeats]"
  27
  28     parser = makeParser(usage)
  29     (options, args) = parser.parse_args(argv[1:])
  30
  31     if len(args) < 3:
  32         print usage
  33         sys.exit(1)
  34
  35     dbfile = args[0]
  36     filename = args[1]
  37     outfile = args[2]
  38
  39     chkSNPrmask(dbfile, filename, outfile, options.repeats, options.cachePages)
  40
  41
  42 def makeParser(usage=""):
  43     parser = optparse.OptionParser(usage=usage)
  44     parser.add_option("--repeats", action="store_true", dest="repeats")
  45     parser.add_option("--cache", type="int", dest="cachePages")
  46
  47     configParser = getConfigParser()
  48     section = "checkSNPrmask"
  49     repeats = getConfigBoolOption(configParser, section, "repeats", False)
  50     cachePages = getConfigOption(configParser, section, "cachePages", None)
  51
  52     parser.set_defaults(repeats=repeats, cachePages=cachePages)
  53
  54     return parser
  55
  56
  57 def chkSNPrmask(dbfile, filename, outfile, repeats=False, cachePages=None):
  58     print dbfile
  59
  60     if cachePages is not None:
  61         if cachePages < 250000:
  62             cachePages = 250000
  63
  64         print "caching locally..."
  65         cachefile = tempfile.mktemp() + ".db"
  66         shutil.copyfile(dbfile, cachefile)
  67         db = sqlite.connect(cachefile)
  68         doCache = True
  69         print "cached..."
  70     else:
  71         cachePages = 500000
  72         doCache = False
  73         db = sqlite.connect(dbfile)
  74
  75     sql = db.cursor()
  76     sql.execute("PRAGMA CACHE_SIZE = %d" % cachePages)
  77     sql.execute("PRAGMA temp_store = MEMORY")
  78     sql.execute("ANALYZE")
  79
  80     infile = open(filename)
  81     featureList = []
  82     featureDict = {}
  83
  84     for line in infile:
  85         if doNotProcessLine(line):
  86             continue
  87
  88         fields = line.strip().split("\t")
  89         chrom = fields[2][3:]
  90         pos = int(fields[3])
  91         featureList.append((chrom,pos))
  92         featureDict[(chrom, pos)] = line.strip()
  93
  94     featureList.sort()
  95
  96     index = 0
  97     currentChrom=None
  98     for (chrom, pos) in featureList:
  99         index += 1
 100         if chrom != currentChrom:
 101             print "\n%s" % chrom
 102             currentChrom = chrom
 103
 104         results = []
 105         try:
 106             sql.execute("select family from repeats where chrom = '%s' and %d between start and stop" % (chrom, pos))
 107             results = sql.fetchall()
 108         except:
 109             pass
 110
 111         if repeats: # if user wants to keep track of the SNPs in repeats
 112             featureDict[(chrom,pos)] += "\tN\A"
 113             for x in results:
 114                 featureDict[(chrom,pos)] += "\t" + str(x)
 115         else:
 116             for x in results:
 117                 try:
 118                     del featureDict[(chrom,pos)]
 119                 except KeyError:
 120                     pass
 121
 122         if index % 100 == 0:
 123             print ".",
 124             sys.stdout.flush()
 125
 126     if doCache:
 127         print "removing cache"
 128         del db
 129         os.remove(cachefile)
 130
 131     outFile = open(outfile, "w")
 132     for key, value in featureDict.iteritems():
 133         outStr = str(value) + "\n"
 134         outFile.write(outStr)
 135
 136     outFile.close()
 137
 138
 139 def doNotProcessLine(line):
 140     return line[0] == "#"
 141
 142
 143 if __name__ == "__main__":
 144     main(sys.argv)