""" Classes for index files that map genomic intervals to values. :Authors: James Taylor (james@bx.psu.edu), Bob Harris (rsharris@bx.psu.edu) An interval index file maps genomic intervals to values. This implementation writes version 1 file format, and reads versions 0 and 1. Index File Format ----------------- All fields are in big-endian format (most significant byte first). All intervals are origin-zero, inclusive start, exclusive end. The file begins with an index file header, then is immediately followed by an index table. The index table points to index headers, and index headers point to bins. Index headers and bins are referenced via pointers (file offsets), and can be placed more or less anywhere in the file. File header ~~~~~~~~~~~ ============ =========== ================================================= offset 0x00: 2C FF 80 0A magic number offset 0x04: 00 00 00 01 version (00 00 00 00 is also supported) offset 0x08: 00 00 00 2A (N) number of index sets offset 0x0C: ... index table ============ =========== ================================================= Index table ~~~~~~~~~~~ The index table is a list of N index headers, packed sequentially and sorted by name. The first begins at offset 0x0C. Each header describes one set of intervals. ============ =========== ================================================= offset: xx xx xx xx (L) length of index src name offset+4: ... index src name (e.g. canFam1.chr1) offset+4+L: xx xx xx xx offset (in this file) to index data offset+8+L: xx xx xx xx (B) number of bytes in each value; for version 0, this field is absent, and B is assumed to be 4 ============ =========== ================================================= Index data ~~~~~~~~~~ The index data for (for one index table) consists of the overall range of intervals followed by an array of pointers to bins. The length of the array is 1+binForRange(maxEnd-1,maxEnd), where maxEnd is the maximum interval end. ============ =========== ================================================= offset: xx xx xx xx minimum interval start offset+4: xx xx xx xx maximum interval end offset+8: xx xx xx xx offset (in this file) to bin 0 offset+12: xx xx xx xx number of intervals in bin 0 offset+16: xx xx xx xx offset (in this file) to bin 1 offset+20: xx xx xx xx number of intervals in bin 1 ... ... ... ============ =========== ================================================= Bin ~~~ A bin is an array of (start,end,val), sorted by increasing start (with end and val as tiebreakers). Note that bins may be empty (the number of intervals indicated in the index data is zero). Note that B is determined from the appropriate entry in the index table. ============ =========== ================================================= offset: xx xx xx xx start for interval 1 offset+4: xx xx xx xx end for interval 1 offset+8: ... (B bytes) value for interval 1 offset+8+B: xx xx xx xx start for interval 2 offset+12+B: xx xx xx xx end for interval 2 offset+16+B: ... (B bytes) value for interval 2 ... ... ... ============ =========== ================================================= """ import os.path import sys from bisect import ( insort, insort_right, ) from struct import ( calcsize, pack, unpack, ) from warnings import warn from bx.misc import filecache try: from bx.misc import seekbzip2 except ImportError: seekbzip2 = None try: from bx.misc import seeklzop except ImportError: seeklzop = None __all__ = ["Indexes", "Index"] MAGIC = 0x2CFF800A VERSION = 2 # These three constants determine the structure of the default binning strategy BIN_LEVELS = 6 # Number of levels of bins to build BIN_FIRST_SHIFT = 17 # Number of bits for the bottom level bin BIN_NEXT_SHIFT = 3 # Number of bits for each higher level bin # Build offset and max size arrays for each bin level BIN_OFFSETS = [1, 0] BIN_OFFSETS_MAX = [(1 << BIN_FIRST_SHIFT << BIN_NEXT_SHIFT), (1 << BIN_FIRST_SHIFT)] for i in range(BIN_LEVELS - 2): BIN_OFFSETS.insert(0, (2 ** (3 * (i + 1))) + BIN_OFFSETS[0]) BIN_OFFSETS_MAX.insert(0, (BIN_OFFSETS_MAX[0] << BIN_NEXT_SHIFT)) # The maximum size for the top bin is actually bigger than the signed integers # we use to store positions in the file, so we'll change it to prevent confusion BIN_OFFSETS_MAX[0] = 2**63 # Constants for the minimum and maximum size of the overall interval MIN = 0 OLD_MAX = 512 * 1024 * 1024 # Maximum size supported by versions < 2 DEFAULT_MAX = 512 * 1024 * 1024 # Default max size to use when none is passed MAX = 2**31 # Absolute max size (limited by file format) def offsets_for_max_size(max_size): """ Return the subset of offsets needed to contain intervals over (0,max_size) """ offset_index = None for i, off_max in enumerate(reversed(BIN_OFFSETS_MAX)): if max_size <= off_max: offset_index = i break if offset_index is None: raise Exception("%d is larger than the maximum possible size (%d)" % (max_size, BIN_OFFSETS_MAX[0])) return BIN_OFFSETS[(len(BIN_OFFSETS) - offset_index - 1):] def bin_for_range(start, end, offsets=None): """Find the smallest bin that can contain interval (start,end)""" if offsets is None: offsets = BIN_OFFSETS start_bin, end_bin = start, max(start, end - 1) start_bin >>= BIN_FIRST_SHIFT end_bin >>= BIN_FIRST_SHIFT for offset in offsets: if start_bin == end_bin: return offset + start_bin else: start_bin >>= BIN_NEXT_SHIFT end_bin >>= BIN_NEXT_SHIFT raise Exception("Interval (%d,%d) out of range" % (start, end)) class AbstractMultiIndexedAccess: """ Allows accessing multiple indexes / files as if they were one """ indexed_access_class = None def __init__(self, filenames, index_filenames=None, keep_open=False, use_cache=False, **kwargs): # TODO: Handle index_filenames argument self.indexes = [ self.new_indexed_access(fname, keep_open=keep_open, use_cache=use_cache, **kwargs) for fname in filenames ] def new_indexed_access(self, data_filename, index_filename=None, keep_open=False, **kwargs): return self.indexed_access_class(data_filename, index_filename, keep_open, **kwargs) def get(self, src, start, end): return list(self.get_as_iterator(src, start, end)) def get_as_iterator(self, src, start, end): for block, _index, _offset in self.get_as_iterator_with_index_and_offset(src, start, end): yield block def get_as_iterator_with_index_and_offset(self, src, start, end): for index in self.indexes: yield from index.get_as_iterator_with_index_and_offset(src, start, end) def close(self): for index in self.indexes: index.close() class AbstractIndexedAccess: """Indexed access to a data using overlap queries, requires an index file""" def __init__(self, data_filename, index_filename=None, keep_open=False, use_cache=False, **kwargs): self.data_kwargs = kwargs self.data_filename = data_filename if data_filename.endswith(".bz2"): if seekbzip2 is None: raise Exception("Trying to open .bz2 file but no seekbzip2 module found") table_filename = data_filename + "t" self.table_filename = table_filename if not os.path.exists(table_filename): raise Exception("Cannot find bz2t file for: " + data_filename) self.file_type = "bz2t" # Strip .bz2 from the filename before adding ".index" data_filename_root = data_filename[:-4] elif data_filename.endswith(".lzo"): if seeklzop is None: raise Exception("Trying to open .lzo file but no seeklzop module found") table_filename = data_filename + "t" self.table_filename = table_filename if not os.path.exists(table_filename): raise Exception("Cannot find lzot file for: " + data_filename) self.file_type = "lzot" # Strip .lzo from the filename before adding ".index" data_filename_root = data_filename[:-4] else: self.file_type = "plain" data_filename_root = data_filename # Open index if index_filename is None: index_filename = data_filename_root + ".index" self.indexes = Indexes(filename=index_filename) # Use a file cache? self.use_cache = use_cache # Open now? if keep_open: self.f = self.open_data() else: self.f = None def close(self): if self.f: self.f.close() self.f = None def open_data(self): if self.file_type == "plain": return open(self.data_filename, "rb") elif self.file_type == "bz2t": f = seekbzip2.SeekableBzip2File(self.data_filename, self.table_filename) if self.use_cache: return filecache.FileCache(f, f.size) else: return f elif self.file_type == "lzot": if self.use_cache: block_cache_size = 20 else: block_cache_size = 0 f = seeklzop.SeekableLzopFile(self.data_filename, self.table_filename, block_cache_size=block_cache_size) return f def get(self, src, start, end): return list(self.get_as_iterator(src, start, end)) def get_as_iterator(self, src, start, end): for val, _index, _offset in self.get_as_iterator_with_index_and_offset(src, start, end): yield val def get_as_iterator_with_index_and_offset(self, src, start, end): for _val_start, _val_end, val in self.indexes.find(src, start, end): yield self.get_at_offset(val), self, val def get_at_offset(self, offset): if self.f: self.f.seek(offset) return self.read_at_current_offset(self.f, **self.data_kwargs) else: f = self.open_data() try: f.seek(offset) return self.read_at_current_offset(f, **self.data_kwargs) finally: f.close() def read_at_current_offset(self, file, **kwargs): raise TypeError("Abstract Method") class Indexes: """A set of indexes, each identified by a unique name""" def __init__(self, filename=None): self.indexes = {} if filename is not None: self.open(filename) def add(self, name, start, end, val, max=DEFAULT_MAX): if name not in self.indexes: self.indexes[name] = Index(max=max) self.indexes[name].add(start, end, val) def get(self, name): if self.indexes[name] is None: offset, value_size = self.offsets[name] self.indexes[name] = Index( filename=self.filename, offset=offset, value_size=value_size, version=self.version ) return self.indexes[name] def find(self, name, start, end): if name in self.indexes: return self.get(name).find(start, end) else: return [] def open(self, filename): self.filename = filename self.offsets = {} # (will map key to (offset,value_size)) with open(filename, "rb") as f: magic, version, length = read_packed(f, ">3I") if magic != MAGIC: raise Exception("File does not have expected header") if version > VERSION: warn( "File claims version %d, I don't known anything about versions beyond %d. Attempting to continue", version, VERSION, ) self.version = version for _ in range(length): key_len = read_packed(f, ">I") key = f.read(key_len).decode() offset = read_packed(f, ">I") if version == 0: value_size = 4 else: value_size = read_packed(f, ">I") assert value_size % 4 == 0, f"unsupported value size: {value_size}" self.indexes[key] = None self.offsets[key] = (offset, value_size) def write(self, f): keys = sorted(self.indexes.keys()) # First determine the size of the header base = calcsize(">3I") for key in keys: key = str(key) base += calcsize(">I") base += len(key) base += calcsize(">2I") # Now actually write the header write_packed(f, ">3I", MAGIC, VERSION, len(self.indexes)) # And write the index table for key in keys: key = str(key) # Write the string prefixed by its length (pascal!) write_packed(f, ">I", len(key)) f.write(key.encode()) # Write offset write_packed(f, ">I", base) base += self.indexes[key].bytes_required() # Write value size write_packed(f, ">I", self.indexes[key].value_size) # And finally write each index in order for key in keys: self.indexes[key].write(f) class Index: def __init__(self, min=MIN, max=DEFAULT_MAX, filename=None, offset=0, value_size=None, version=None): self._value_size = value_size self.max_val = 1 # (1, rather than 0, to force value_size > 0) if filename is None: self.new(min, max) else: self.open(filename, offset, version) def get_value_size(self): if self._value_size is not None: return self._value_size else: return round_up_to_4(bytes_of(self.max_val)) value_size = property(fget=get_value_size) def new(self, min, max): """Create an empty index for intervals in the range min, max""" # Ensure the range will fit given the shifting strategy assert MIN <= min <= max <= MAX self.min = min self.max = max # Determine offsets to use self.offsets = offsets_for_max_size(max) # Determine the largest bin we will actually use self.bin_count = bin_for_range(max - 1, max, offsets=self.offsets) + 1 # Create empty bins self.bins = [[] for i in range(self.bin_count)] def open(self, filename, offset, version): self.filename = filename self.offset = offset # Open the file and seek to where we expect our header f = open(filename, "rb") f.seek(offset) # Read min/max min, max = read_packed(f, ">2I") self.new(min, max) # Decide how many levels of bins based on 'max' if version < 2: # Prior to version 2 all files used the bins for 512MB self.offsets = offsets_for_max_size(OLD_MAX - 1) else: self.offsets = offsets_for_max_size(max) # Read bin indexes self.bin_offsets = [] self.bin_sizes = [] for _ in range(self.bin_count): o, s = read_packed(f, ">2I") self.bin_offsets.append(o) self.bin_sizes.append(s) # Initialize bins to None, indicating that they need to be loaded self.bins = [None for _ in range(self.bin_count)] def add(self, start, end, val): """Add the interval (start,end) with associated value val to the index""" insort(self.bins[bin_for_range(start, end, offsets=self.offsets)], (start, end, val)) assert val >= 0 self.max_val = max(self.max_val, val) def find(self, start, end): rval = [] start_bin = (max(start, self.min)) >> BIN_FIRST_SHIFT end_bin = (min(end, self.max) - 1) >> BIN_FIRST_SHIFT for offset in self.offsets: for i in range(start_bin + offset, end_bin + offset + 1): if self.bins[i] is None: self.load_bin(i) # Iterate over bin and insert any overlapping elements into return value for el_start, el_end, val in self.bins[i]: if el_start < end and el_end > start: insort_right(rval, (el_start, el_end, val)) start_bin >>= BIN_NEXT_SHIFT end_bin >>= BIN_NEXT_SHIFT return rval def iterate(self): for i in range(self.bin_count): if self.bins[i] is None: self.load_bin(i) yield from self.bins[i] def load_bin(self, index): bin = [] if self.bin_sizes[index] == 0: self.bins[index] = bin return f = open(self.filename, "rb") f.seek(self.bin_offsets[index]) # One big read for happy NFS item_size = self.value_size + calcsize(">2I") buffer = f.read(self.bin_sizes[index] * item_size) for i in range(self.bin_sizes[index]): start, end = unpack(">2I", buffer[i * item_size : i * item_size + 8]) val = unpack_uints(buffer[i * item_size + 8 : (i + 1) * item_size]) bin.append((start, end, val)) self.bins[index] = bin f.close() def write(self, f): value_size = self.value_size item_size = value_size + calcsize(">2I") # Write min/max write_packed(f, ">2I", self.min, self.max) # Write table of bin sizes and offsets base = f.tell() + self.bin_count * calcsize(">2I") for bin in self.bins: write_packed(f, ">2I", base, len(bin)) base += len(bin) * item_size # Write contents of each bin for bin in self.bins: for start, end, val in bin: write_packed(f, ">2I", start, end) write_packed_uints(f, val, value_size) def bytes_required(self): item_size = self.value_size + calcsize(">2I") rval = calcsize(">2I") rval += self.bin_count * calcsize(">2I") for bin in self.bins: rval += len(bin) * item_size return rval def write_packed(f, pattern, *vals): f.write(pack(pattern, *vals)) def read_packed(f, pattern): rval = unpack(pattern, f.read(calcsize(pattern))) if len(rval) == 1: return rval[0] return rval def write_packed_uints(f, v, num_bytes): if num_bytes < 4: write_packed(f, ">I", v) else: parts = [] while num_bytes > 0: parts.append(v & 0xFFFFFFFF) v >>= 32 num_bytes -= 4 parts.reverse() # (write most-significant chunk first) write_packed(f, ">%dI" % len(parts), *parts) def unpack_uints(parts): chunks = len(parts) / 4 vals = unpack(">%dI" % chunks, parts) val = vals[0] for v in vals[1:]: val = (val << 32) + v return val def bytes_of(v): assert v > 0 b = 0 while v > 0: v >>= 8 b += 1 return b def round_up_to_4(v): if v % 4 == 0: return v else: return v + 4 - (v % 4)