annotate scripts/build_profile_indexes.py @ 0:3b33da018e74 draft default tip

Imported from capsule None

1 #!/usr/bin/env python
2 #Dan Blankenberg
3 
4 VERSION = '1.0.0' # version of this script
5 
6 from optparse import OptionParser
7 import os, gzip, struct, time
8 from ftplib import FTP #do we want a diff method than using FTP to determine Chrom Names, eg use local copy
9 
10 #import md5 from hashlib; if python2.4 or less, use old md5
11 try:
12     from hashlib import md5
13 except ImportError:
14     from md5 import new as md5
15 
16 #import BitSet from bx-python, try using eggs and package resources, fall back to any local installation
17 try:
18     from galaxy import eggs
19     import pkg_resources
20     pkg_resources.require( "bx-python" )
21 except: pass #Maybe there is a local installation available
22 from bx.bitset import BitSet
23 
24 #Define constants
25 STRUCT_FMT = '<I'
26 STRUCT_SIZE = struct.calcsize( STRUCT_FMT )
27 DEFAULT_BITSET_SIZE = 300000000
28 CHUNK_SIZE = 1024
29 
30 #Headers used to parse .sql files to determine column indexes for chromosome name, start and end
31 alias_spec = { 
32     'chromCol'  : [ 'chrom' , 'CHROMOSOME' , 'CHROM', 'Chromosome Name', 'tName' ],  
33     'startCol'  : [ 'start' , 'START', 'chromStart', 'txStart', 'Start Position (bp)', 'tStart', 'genoStart' ],
34     'endCol'    : [ 'end'   , 'END'  , 'STOP', 'chromEnd', 'txEnd', 'End Position (bp)', 'tEnd', 'genoEnd' ], 
35 }
36 
37 #Headers used to parse trackDb.txt.gz
38 #TODO: these should be parsed directly from trackDb.sql
39 trackDb_headers = ["tableName", "shortLabel", "type", "longLabel", "visibility", "priority", "colorR", "colorG", "colorB", "altColorR", "altColorG", "altColorB", "useScore", "private", "restrictCount", "restrictList", "url", "html", "grp", "canPack", "settings"]
40 
41 def get_columns( filename ):
42     input_sql = open( filename ).read()
43     input_sql = input_sql.split( 'CREATE TABLE ' )[1].split( ';' )[0]
44     input_sql = input_sql.split( ' (', 1 )
45     table_name = input_sql[0].strip().strip( '`' )
46     input_sql = [ split.strip().split( ' ' )[0].strip().strip( '`' ) for split in input_sql[1].rsplit( ')', 1 )[0].strip().split( '\n' ) ]
47     print input_sql
48     chrom_col = None
49     start_col = None
50     end_col = None
51     for col_name in alias_spec['chromCol']:
52         for i, header_name in enumerate( input_sql ):
53             if col_name == header_name:
54                 chrom_col = i
55                 break
56         if chrom_col is not None:
57             break
58     
59     for col_name in alias_spec['startCol']:
60         for i, header_name in enumerate( input_sql ):
61             if col_name == header_name:
62                 start_col = i
63                 break
64         if start_col is not None:
65             break
66 
67     for col_name in alias_spec['endCol']:
68         for i, header_name in enumerate( input_sql ):
69             if col_name == header_name:
70                 end_col = i
71                 break
72         if end_col is not None:
73             break
74 
75     return table_name, chrom_col, start_col, end_col
76 
77 
78 def create_grouping_xml( input_dir, output_dir, dbkey ):
79     output_filename = os.path.join( output_dir, '%s_tables.xml' % dbkey )
80     def load_groups( file_name = 'grp.txt.gz' ):
81         groups = {}
82         for line in gzip.open( os.path.join( input_dir, file_name ) ):
83             fields = line.split( '\t' )
84             groups[fields[0]] = { 'desc': fields[1], 'priority':  fields[2] }
85         return groups
86     f = gzip.open( os.path.join( input_dir, 'trackDb.txt.gz' ) )
87     out = open( output_filename, 'wb' )
88     tables = {}
89     cur_buf = ''
90     while True:
91         line = f.readline()
92         if not line: break
93         #remove new lines
94         line = line.rstrip( '\n\r' )
95         line = line.replace( '\\\t', ' ' ) #replace escaped tabs with space
96         cur_buf += "%s\n" % line.rstrip( '\\' )
97         if line.endswith( '\\' ):
98             continue #line is wrapped, next line
99         #all fields should be loaded now...
00         fields = cur_buf.split( '\t' )
01         cur_buf = '' #reset buffer
02         assert len( fields ) == len( trackDb_headers ), 'Failed Parsing trackDb.txt.gz; fields: %s' % fields
03         table_name = fields[ 0 ]
04         tables[ table_name ] = {}
05         for field_name, field_value in zip( trackDb_headers, fields ):
06             tables[ table_name ][ field_name ] = field_value
07         #split settings fields into dict
08         fields = fields[-1].split( '\n' )
09         tables[ table_name ][ 'settings' ] = {}
10         for field in fields:
11             setting_fields = field.split( ' ', 1 )
12             setting_name = setting_value = setting_fields[ 0 ]
13             if len( setting_fields ) > 1:
14                 setting_value = setting_fields[ 1 ]
15             if setting_name or setting_value:
16                 tables[ table_name ][ 'settings' ][ setting_name ] = setting_value
17     #Load Groups
18     groups = load_groups()
19     in_groups = {}
20     for table_name, values in tables.iteritems():
21         if os.path.exists( os.path.join( output_dir, table_name ) ):
22             group = values['grp']
23             if group not in in_groups:
24                 in_groups[group]={}
25             #***NAME CHANGE***, 'subTrack' no longer exists as a setting...use 'parent' instead
26             #subTrack = values.get('settings', {} ).get( 'subTrack', table_name )
27             subTrack = values.get('settings', {} ).get( 'parent', table_name ).split( ' ' )[0] #need to split, because could be e.g. 'trackgroup on'
28             if subTrack not in in_groups[group]:
29                 in_groups[group][subTrack]=[]
30             in_groups[group][subTrack].append( table_name )
31     
32     assigned_tables = []
33     out.write( """<filter type="data_meta" data_ref="input1" meta_key="dbkey" value="%s">\n""" % ( dbkey ) )
34     out.write( "  <options>\n" )
35     for group, subTracks in sorted( in_groups.iteritems() ):
36         out.write( """    <option name="%s" value="group-%s">\n""" % ( groups[group]['desc'], group ) )
37         for sub_name, sub_tracks in subTracks.iteritems():
38             if len( sub_tracks ) > 1:
39                 out.write( """      <option name="%s" value="subtracks-%s">\n""" % ( sub_name, sub_name ) )
40                 sub_tracks.sort()
41                 for track in sub_tracks:
42                     track_label = track
43                     if "$" not in tables[track]['shortLabel']:
44                         track_label = tables[track]['shortLabel']
45                     out.write( """        <option name="%s" value="%s"/>\n""" % ( track_label, track ) )
46                     assigned_tables.append( track )
47                 out.write( "      </option>\n" )
48             else:
49                 track = sub_tracks[0]
50                 track_label = track
51                 if "$" not in tables[track]['shortLabel']:
52                     track_label = tables[track]['shortLabel']
53                 out.write( """        <option name="%s" value="%s"/>\n""" % ( track_label, track ) )
54                 assigned_tables.append( track )
55         out.write( "    </option>\n" )
56     unassigned_tables = list( sorted( [ table_dir for table_dir in os.listdir( output_dir ) if table_dir not in assigned_tables and os.path.isdir( os.path.join( output_dir, table_dir ) ) ] ) )
57     if unassigned_tables:
58         out.write( """    <option name="Uncategorized Tables" value="group-trackDbUnassigned">\n""" )
59         for table_name in unassigned_tables:
60             out.write( """        <option name="%s" value="%s"/>\n""" % ( table_name, table_name ) )
61         out.write( "    </option>\n" )
62     out.write( "  </options>\n" )
63     out.write( """</filter>\n""" )
64     out.close()
65 
66 def write_database_dump_info( input_dir, output_dir, dbkey, chrom_lengths, default_bitset_size ):
67     #generate hash for profiled table directories
68     #sort directories off output root (files in output root not hashed, including the profiler_info.txt file)
69     #sort files in each directory and hash file contents
70     profiled_hash = md5()
71     for table_dir in sorted( [ table_dir for table_dir in os.listdir( output_dir ) if os.path.isdir( os.path.join( output_dir, table_dir ) ) ] ):
72         for filename in sorted( os.listdir( os.path.join( output_dir, table_dir ) ) ):
73             f  = open( os.path.join( output_dir, table_dir, filename ), 'rb' )
74             while True:
75                 hash_chunk = f.read( CHUNK_SIZE )
76                 if not hash_chunk:
77                     break
78                 profiled_hash.update( hash_chunk )
79     profiled_hash = profiled_hash.hexdigest()
80     
81     #generate hash for input dir
82     #sort directories off input root
83     #sort files in each directory and hash file contents
84     database_hash = md5()
85     for dirpath, dirnames, filenames in sorted( os.walk( input_dir ) ):
86         for filename in sorted( filenames ):
87             f  = open( os.path.join( input_dir, dirpath, filename ), 'rb' )
88             while True:
89                 hash_chunk = f.read( CHUNK_SIZE )
90                 if not hash_chunk:
91                     break
92                 database_hash.update( hash_chunk )
93     database_hash = database_hash.hexdigest()
94     
95     #write out info file
96     out = open( os.path.join( output_dir, 'profiler_info.txt' ), 'wb' )
97     out.write( 'dbkey\t%s\n' % ( dbkey ) )
98     out.write( 'chromosomes\t%s\n' % ( ','.join( [ '%s=%s' % ( chrom_name, chrom_len ) for chrom_name, chrom_len in chrom_lengths.iteritems() ] ) ) )
99     out.write( 'bitset_size\t%s\n' % ( default_bitset_size ) )
00     for line in open( os.path.join( input_dir, 'trackDb.sql' ) ):
01         line = line.strip()
02         if line.startswith( '-- Dump completed on ' ):
03             line = line[ len( '-- Dump completed on ' ): ]
04             out.write( 'dump_time\t%s\n' % ( line ) )
05             break
06     out.write( 'dump_hash\t%s\n' % ( database_hash ) )
07     out.write( 'profiler_time\t%s\n' % ( time.time() ) )
08     out.write( 'profiler_hash\t%s\n' % ( profiled_hash ) )
09     out.write( 'profiler_version\t%s\n' % ( VERSION ) )
10     out.write( 'profiler_struct_format\t%s\n' % ( STRUCT_FMT ) )
11     out.write( 'profiler_struct_size\t%s\n' % ( STRUCT_SIZE ) )
12     out.close()
13     
14 def __main__():
15     usage = "usage: %prog options"
16     parser = OptionParser( usage=usage )
17     parser.add_option( '-d', '--dbkey', dest='dbkey', default='hg18', help='dbkey to process' )
18     parser.add_option( '-i', '--input_dir', dest='input_dir', default=os.path.join( 'golden_path','%s', 'database' ), help='Input Directory' )
19     parser.add_option( '-o', '--output_dir', dest='output_dir', default=os.path.join( 'profiled_annotations','%s' ), help='Output Directory' )
20     parser.add_option( '-c', '--chromosomes', dest='chromosomes', default='', help='Comma separated list of: ChromName1[=length],ChromName2[=length],...' )
21     parser.add_option( '-b', '--bitset_size', dest='bitset_size', default=DEFAULT_BITSET_SIZE, type='int', help='Default BitSet size; overridden by sizes specified in chromInfo.txt.gz or by --chromosomes' )
22     parser.add_option( '-f', '--ftp_site', dest='ftp_site', default='hgdownload.cse.ucsc.edu', help='FTP site; used for chromosome info when chromInfo.txt.gz method fails' )
23     parser.add_option( '-p', '--ftp_path', dest='ftp_path', default='/goldenPath/%s/chromosomes/', help='FTP Path; used for chromosome info when chromInfo.txt.gz method fails' )
24     
25     ( options, args ) = parser.parse_args()
26     
27     input_dir = options.input_dir
28     if '%' in input_dir:
29         input_dir = input_dir % options.dbkey
30     assert os.path.exists( input_dir ), 'Input directory does not exist'
31     output_dir = options.output_dir
32     if '%' in output_dir:
33         output_dir = output_dir % options.dbkey
34     assert not os.path.exists( output_dir ), 'Output directory already exists'
35     os.makedirs( output_dir )
36     ftp_path = options.ftp_path
37     if '%' in ftp_path:
38         ftp_path = ftp_path % options.dbkey
39     
40     #Get chromosome names and lengths
41     chrom_lengths = {}
42     if options.chromosomes:
43         for chrom in options.chromosomes.split( ',' ):
44             fields = chrom.split( '=' )
45             chrom = fields[0]
46             if len( fields ) > 1:
47                 chrom_len = int( fields[1] )
48             else:
49                 chrom_len = options.bitset_size
50             chrom_lengths[ chrom ] = chrom_len
51         chroms = chrom_lengths.keys()
52         print 'Chrom info taken from command line option.'
53     else:
54         try:
55             for line in gzip.open( os.path.join( input_dir, 'chromInfo.txt.gz' ) ):
56                 fields = line.strip().split( '\t' )
57                 chrom_lengths[ fields[0] ] = int( fields[ 1 ] )
58             chroms = chrom_lengths.keys()
59             print 'Chrom info taken from chromInfo.txt.gz.'
60         except Exception, e:
61             print 'Error loading chrom info from chromInfo.txt.gz, trying FTP method.'
62             chrom_lengths = {} #zero out chrom_lengths
63             chroms = []
64             ftp = FTP( options.ftp_site )
65             ftp.login()
66             for name in ftp.nlst( ftp_path ):
67                 if name.endswith( '.fa.gz' ):
68                     chroms.append( name.split( '/' )[-1][ :-len( '.fa.gz' ) ] )
69             ftp.close()
70             for chrom in chroms:
71                 chrom_lengths[ chrom ] = options.bitset_size
72     #sort chroms by length of name, decending; necessary for when table names start with chrom name
73     chroms = list( reversed( [ chrom for chrom_len, chrom in sorted( [ ( len( chrom ), chrom ) for chrom in chroms ] ) ] ) )
74     
75     #parse tables from local files
76     #loop through directory contents, if file ends in '.sql', process table
77     for filename in os.listdir( input_dir ):
78         if filename.endswith ( '.sql' ):
79             base_filename = filename[ 0:-len( '.sql' ) ]
80             table_out_dir = os.path.join( output_dir, base_filename )
81             #some tables are chromosome specific, lets strip off the chrom name
82             for chrom in chroms:
83                 if base_filename.startswith( "%s_" % chrom ):
84                     #found chromosome
85                     table_out_dir = os.path.join( output_dir, base_filename[len( "%s_" % chrom ):] )
86                     break
87             #create table dir
88             if not os.path.exists( table_out_dir ):
89                 os.mkdir( table_out_dir ) #table dir may already exist in the case of single chrom tables
90                 print "Created table dir (%s)." % table_out_dir
91             else:
92                 print "Table dir (%s) already exists." % table_out_dir
93             #find column assignments
94             table_name, chrom_col, start_col, end_col = get_columns( "%s.sql" % os.path.join( input_dir, base_filename ) )
95             if chrom_col is None or start_col is None or end_col is None:
96                 print "Table %s (%s) does not appear to have a chromosome, a start, or a stop." % ( table_name, "%s.sql" % os.path.join( input_dir, base_filename ) )
97                 if not os.listdir( table_out_dir ):
98                     print "Removing empty table (%s) directory (%s)." % ( table_name, table_out_dir )
99                     os.rmdir( table_out_dir )
00                 continue
01             #build bitsets from table
02             bitset_dict = {}
03             for line in gzip.open( '%s.txt.gz' % os.path.join( input_dir, base_filename )  ):
04                 fields = line.strip().split( '\t' )
05                 chrom = fields[ chrom_col ]
06                 start = int( fields[ start_col ] )
07                 end = int( fields[ end_col ] )
08                 if chrom not in bitset_dict:
09                     bitset_dict[ chrom ] = BitSet( chrom_lengths.get( chrom, options.bitset_size ) )
10                 bitset_dict[ chrom ].set_range(  start, end - start  )
11             #write bitsets as profiled annotations
12             for chrom_name, chrom_bits in bitset_dict.iteritems():
13                 out = open( os.path.join( table_out_dir, '%s.covered' % chrom_name  ), 'wb' )
14                 end = 0
15                 total_regions = 0
16                 total_coverage = 0
17                 max_size = chrom_lengths.get( chrom_name, options.bitset_size )
18                 while True:
19                     start = chrom_bits.next_set( end )
20                     if start >= max_size:
21                         break
22                     end = chrom_bits.next_clear( start )
23                     out.write( struct.pack( STRUCT_FMT, start ) )
24                     out.write( struct.pack( STRUCT_FMT, end ) )
25                     total_regions += 1
26                     total_coverage += end - start
27                     if end >= max_size:
28                         break
29                 out.close()
30                 open( os.path.join( table_out_dir, '%s.total_regions' % chrom_name  ), 'wb' ).write( str( total_regions ) )
31                 open( os.path.join( table_out_dir, '%s.total_coverage' % chrom_name  ), 'wb' ).write( str( total_coverage ) )
32     
33     #create xml
34     create_grouping_xml( input_dir, output_dir, options.dbkey )
35     #create database dump info file, for database version control
36     write_database_dump_info( input_dir, output_dir, options.dbkey, chrom_lengths, options.bitset_size )
37     
38 if __name__ == "__main__": __main__()