PageRenderTime 83ms CodeModel.GetById 47ms app.highlight 30ms RepoModel.GetById 1ms app.codeStats 0ms

/scripts/tools/annotation_profiler/build_profile_indexes.py

https://bitbucket.org/cistrome/cistrome-harvard/
Python | 338 lines | 317 code | 13 blank | 8 comment | 31 complexity | 82ec483a55dc046aa41c662f0a94d06d MD5 | raw file
  1#!/usr/bin/env python
  2#Dan Blankenberg
  3
  4VERSION = '1.0.0' # version of this script
  5
  6from optparse import OptionParser
  7import os, gzip, struct, time
  8from 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
 11try:
 12    from hashlib import md5
 13except 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
 17try:
 18    from galaxy import eggs
 19    import pkg_resources
 20    pkg_resources.require( "bx-python" )
 21except: pass #Maybe there is a local installation available
 22from bx.bitset import BitSet
 23
 24#Define constants
 25STRUCT_FMT = '<I'
 26STRUCT_SIZE = struct.calcsize( STRUCT_FMT )
 27DEFAULT_BITSET_SIZE = 300000000
 28CHUNK_SIZE = 1024
 29
 30#Headers used to parse .sql files to determine column indexes for chromosome name, start and end
 31alias_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
 39trackDb_headers = ["tableName", "shortLabel", "type", "longLabel", "visibility", "priority", "colorR", "colorG", "colorB", "altColorR", "altColorG", "altColorB", "useScore", "private", "restrictCount", "restrictList", "url", "html", "grp", "canPack", "settings"]
 40
 41def 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
 78def 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...
100        fields = cur_buf.split( '\t' )
101        cur_buf = '' #reset buffer
102        assert len( fields ) == len( trackDb_headers ), 'Failed Parsing trackDb.txt.gz; fields: %s' % fields
103        table_name = fields[ 0 ]
104        tables[ table_name ] = {}
105        for field_name, field_value in zip( trackDb_headers, fields ):
106            tables[ table_name ][ field_name ] = field_value
107        #split settings fields into dict
108        fields = fields[-1].split( '\n' )
109        tables[ table_name ][ 'settings' ] = {}
110        for field in fields:
111            setting_fields = field.split( ' ', 1 )
112            setting_name = setting_value = setting_fields[ 0 ]
113            if len( setting_fields ) > 1:
114                setting_value = setting_fields[ 1 ]
115            if setting_name or setting_value:
116                tables[ table_name ][ 'settings' ][ setting_name ] = setting_value
117    #Load Groups
118    groups = load_groups()
119    in_groups = {}
120    for table_name, values in tables.iteritems():
121        if os.path.exists( os.path.join( output_dir, table_name ) ):
122            group = values['grp']
123            if group not in in_groups:
124                in_groups[group]={}
125            #***NAME CHANGE***, 'subTrack' no longer exists as a setting...use 'parent' instead
126            #subTrack = values.get('settings', {} ).get( 'subTrack', table_name )
127            subTrack = values.get('settings', {} ).get( 'parent', table_name ).split( ' ' )[0] #need to split, because could be e.g. 'trackgroup on'
128            if subTrack not in in_groups[group]:
129                in_groups[group][subTrack]=[]
130            in_groups[group][subTrack].append( table_name )
131    
132    assigned_tables = []
133    out.write( """<filter type="data_meta" data_ref="input1" meta_key="dbkey" value="%s">\n""" % ( dbkey ) )
134    out.write( "  <options>\n" )
135    for group, subTracks in sorted( in_groups.iteritems() ):
136        out.write( """    <option name="%s" value="group-%s">\n""" % ( groups[group]['desc'], group ) )
137        for sub_name, sub_tracks in subTracks.iteritems():
138            if len( sub_tracks ) > 1:
139                out.write( """      <option name="%s" value="subtracks-%s">\n""" % ( sub_name, sub_name ) )
140                sub_tracks.sort()
141                for track in sub_tracks:
142                    track_label = track
143                    if "$" not in tables[track]['shortLabel']:
144                        track_label = tables[track]['shortLabel']
145                    out.write( """        <option name="%s" value="%s"/>\n""" % ( track_label, track ) )
146                    assigned_tables.append( track )
147                out.write( "      </option>\n" )
148            else:
149                track = sub_tracks[0]
150                track_label = track
151                if "$" not in tables[track]['shortLabel']:
152                    track_label = tables[track]['shortLabel']
153                out.write( """        <option name="%s" value="%s"/>\n""" % ( track_label, track ) )
154                assigned_tables.append( track )
155        out.write( "    </option>\n" )
156    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 ) ) ] ) )
157    if unassigned_tables:
158        out.write( """    <option name="Uncategorized Tables" value="group-trackDbUnassigned">\n""" )
159        for table_name in unassigned_tables:
160            out.write( """        <option name="%s" value="%s"/>\n""" % ( table_name, table_name ) )
161        out.write( "    </option>\n" )
162    out.write( "  </options>\n" )
163    out.write( """</filter>\n""" )
164    out.close()
165
166def write_database_dump_info( input_dir, output_dir, dbkey, chrom_lengths, default_bitset_size ):
167    #generate hash for profiled table directories
168    #sort directories off output root (files in output root not hashed, including the profiler_info.txt file)
169    #sort files in each directory and hash file contents
170    profiled_hash = md5()
171    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 ) ) ] ):
172        for filename in sorted( os.listdir( os.path.join( output_dir, table_dir ) ) ):
173            f  = open( os.path.join( output_dir, table_dir, filename ), 'rb' )
174            while True:
175                hash_chunk = f.read( CHUNK_SIZE )
176                if not hash_chunk:
177                    break
178                profiled_hash.update( hash_chunk )
179    profiled_hash = profiled_hash.hexdigest()
180    
181    #generate hash for input dir
182    #sort directories off input root
183    #sort files in each directory and hash file contents
184    database_hash = md5()
185    for dirpath, dirnames, filenames in sorted( os.walk( input_dir ) ):
186        for filename in sorted( filenames ):
187            f  = open( os.path.join( input_dir, dirpath, filename ), 'rb' )
188            while True:
189                hash_chunk = f.read( CHUNK_SIZE )
190                if not hash_chunk:
191                    break
192                database_hash.update( hash_chunk )
193    database_hash = database_hash.hexdigest()
194    
195    #write out info file
196    out = open( os.path.join( output_dir, 'profiler_info.txt' ), 'wb' )
197    out.write( 'dbkey\t%s\n' % ( dbkey ) )
198    out.write( 'chromosomes\t%s\n' % ( ','.join( [ '%s=%s' % ( chrom_name, chrom_len ) for chrom_name, chrom_len in chrom_lengths.iteritems() ] ) ) )
199    out.write( 'bitset_size\t%s\n' % ( default_bitset_size ) )
200    for line in open( os.path.join( input_dir, 'trackDb.sql' ) ):
201        line = line.strip()
202        if line.startswith( '-- Dump completed on ' ):
203            line = line[ len( '-- Dump completed on ' ): ]
204            out.write( 'dump_time\t%s\n' % ( line ) )
205            break
206    out.write( 'dump_hash\t%s\n' % ( database_hash ) )
207    out.write( 'profiler_time\t%s\n' % ( time.time() ) )
208    out.write( 'profiler_hash\t%s\n' % ( profiled_hash ) )
209    out.write( 'profiler_version\t%s\n' % ( VERSION ) )
210    out.write( 'profiler_struct_format\t%s\n' % ( STRUCT_FMT ) )
211    out.write( 'profiler_struct_size\t%s\n' % ( STRUCT_SIZE ) )
212    out.close()
213    
214def __main__():
215    usage = "usage: %prog options"
216    parser = OptionParser( usage=usage )
217    parser.add_option( '-d', '--dbkey', dest='dbkey', default='hg18', help='dbkey to process' )
218    parser.add_option( '-i', '--input_dir', dest='input_dir', default=os.path.join( 'golden_path','%s', 'database' ), help='Input Directory' )
219    parser.add_option( '-o', '--output_dir', dest='output_dir', default=os.path.join( 'profiled_annotations','%s' ), help='Output Directory' )
220    parser.add_option( '-c', '--chromosomes', dest='chromosomes', default='', help='Comma separated list of: ChromName1[=length],ChromName2[=length],...' )
221    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' )
222    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' )
223    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' )
224    
225    ( options, args ) = parser.parse_args()
226    
227    input_dir = options.input_dir
228    if '%' in input_dir:
229        input_dir = input_dir % options.dbkey
230    assert os.path.exists( input_dir ), 'Input directory does not exist'
231    output_dir = options.output_dir
232    if '%' in output_dir:
233        output_dir = output_dir % options.dbkey
234    assert not os.path.exists( output_dir ), 'Output directory already exists'
235    os.makedirs( output_dir )
236    ftp_path = options.ftp_path
237    if '%' in ftp_path:
238        ftp_path = ftp_path % options.dbkey
239    
240    #Get chromosome names and lengths
241    chrom_lengths = {}
242    if options.chromosomes:
243        for chrom in options.chromosomes.split( ',' ):
244            fields = chrom.split( '=' )
245            chrom = fields[0]
246            if len( fields ) > 1:
247                chrom_len = int( fields[1] )
248            else:
249                chrom_len = options.bitset_size
250            chrom_lengths[ chrom ] = chrom_len
251        chroms = chrom_lengths.keys()
252        print 'Chrom info taken from command line option.'
253    else:
254        try:
255            for line in gzip.open( os.path.join( input_dir, 'chromInfo.txt.gz' ) ):
256                fields = line.strip().split( '\t' )
257                chrom_lengths[ fields[0] ] = int( fields[ 1 ] )
258            chroms = chrom_lengths.keys()
259            print 'Chrom info taken from chromInfo.txt.gz.'
260        except Exception, e:
261            print 'Error loading chrom info from chromInfo.txt.gz, trying FTP method.'
262            chrom_lengths = {} #zero out chrom_lengths
263            chroms = []
264            ftp = FTP( options.ftp_site )
265            ftp.login()
266            for name in ftp.nlst( ftp_path ):
267                if name.endswith( '.fa.gz' ):
268                    chroms.append( name.split( '/' )[-1][ :-len( '.fa.gz' ) ] )
269            ftp.close()
270            for chrom in chroms:
271                chrom_lengths[ chrom ] = options.bitset_size
272    #sort chroms by length of name, decending; necessary for when table names start with chrom name
273    chroms = list( reversed( [ chrom for chrom_len, chrom in sorted( [ ( len( chrom ), chrom ) for chrom in chroms ] ) ] ) )
274    
275    #parse tables from local files
276    #loop through directory contents, if file ends in '.sql', process table
277    for filename in os.listdir( input_dir ):
278        if filename.endswith ( '.sql' ):
279            base_filename = filename[ 0:-len( '.sql' ) ]
280            table_out_dir = os.path.join( output_dir, base_filename )
281            #some tables are chromosome specific, lets strip off the chrom name
282            for chrom in chroms:
283                if base_filename.startswith( "%s_" % chrom ):
284                    #found chromosome
285                    table_out_dir = os.path.join( output_dir, base_filename[len( "%s_" % chrom ):] )
286                    break
287            #create table dir
288            if not os.path.exists( table_out_dir ):
289                os.mkdir( table_out_dir ) #table dir may already exist in the case of single chrom tables
290                print "Created table dir (%s)." % table_out_dir
291            else:
292                print "Table dir (%s) already exists." % table_out_dir
293            #find column assignments
294            table_name, chrom_col, start_col, end_col = get_columns( "%s.sql" % os.path.join( input_dir, base_filename ) )
295            if chrom_col is None or start_col is None or end_col is None:
296                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 ) )
297                if not os.listdir( table_out_dir ):
298                    print "Removing empty table (%s) directory (%s)." % ( table_name, table_out_dir )
299                    os.rmdir( table_out_dir )
300                continue
301            #build bitsets from table
302            bitset_dict = {}
303            for line in gzip.open( '%s.txt.gz' % os.path.join( input_dir, base_filename )  ):
304                fields = line.strip().split( '\t' )
305                chrom = fields[ chrom_col ]
306                start = int( fields[ start_col ] )
307                end = int( fields[ end_col ] )
308                if chrom not in bitset_dict:
309                    bitset_dict[ chrom ] = BitSet( chrom_lengths.get( chrom, options.bitset_size ) )
310                bitset_dict[ chrom ].set_range(  start, end - start  )
311            #write bitsets as profiled annotations
312            for chrom_name, chrom_bits in bitset_dict.iteritems():
313                out = open( os.path.join( table_out_dir, '%s.covered' % chrom_name  ), 'wb' )
314                end = 0
315                total_regions = 0
316                total_coverage = 0
317                max_size = chrom_lengths.get( chrom_name, options.bitset_size )
318                while True:
319                    start = chrom_bits.next_set( end )
320                    if start >= max_size:
321                        break
322                    end = chrom_bits.next_clear( start )
323                    out.write( struct.pack( STRUCT_FMT, start ) )
324                    out.write( struct.pack( STRUCT_FMT, end ) )
325                    total_regions += 1
326                    total_coverage += end - start
327                    if end >= max_size:
328                        break
329                out.close()
330                open( os.path.join( table_out_dir, '%s.total_regions' % chrom_name  ), 'wb' ).write( str( total_regions ) )
331                open( os.path.join( table_out_dir, '%s.total_coverage' % chrom_name  ), 'wb' ).write( str( total_coverage ) )
332    
333    #create xml
334    create_grouping_xml( input_dir, output_dir, options.dbkey )
335    #create database dump info file, for database version control
336    write_database_dump_info( input_dir, output_dir, options.dbkey, chrom_lengths, options.bitset_size )
337    
338if __name__ == "__main__": __main__()