chkSNPrmask.py

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