#!/usr/bin/env python """Classes for dealing with bitvectors: arrays of 1 and 0. bitvectors are often much faster than numpy arrays for bit operations, especially if several characters (e.g. of DNA) can be packed into a single byte. Usage: v = Bitvector('11000101') Provides the following major classes and factory functions: ShortBitvector and LongBitvector: subclass ImmutableBitvector and are produced by the factory function Bitvector(). Short uses an int, while Long uses a long. ShortBitvectors in particular are _very_ fast, but are limited to 31 bits. Bitvector() will return the appropriate kind of vector for your sequence. MutableBitvector: always a LongBitvector. Can be changed through __getitem__, e.g. vec[3] = 0. Use freeze() and thaw() methods to convert between mutable and immutable Bitvectors (both types define both methods). VectorFromCases: constructs a Bitvector from letters in a string. VectorFromMatches: constructs a Bitvector from a text and a pattern, with 1 at the positions where there's a match. Pattern can be a string or a regex. PackedBases: subclasses LongBitvector: provides a way of storing nucleic acid bases compactly and in a format convenient for assessing sequence similarity. Note: there isn't a general VectorFromX factory function because if you have a function, you can always do Bitvector(map(func, items)), making VectorFromX trivial. """ from cogent.util.misc import Delegator import re from string import maketrans from operator import and_, or_, xor from numpy import log2 from sys import maxint __author__ = "Jeremy Widmann" __copyright__ = "Copyright 2007-2009, The Cogent Project" __credits__ = ["Jeremy Widmann", "Rob Knight", "Daniel McDonald"] __license__ = "GPL" __version__ = "1.4.1" __maintainer__ = "Jeremy Widmann" __email__ = "jeremy.widmann@colorado.edu" __status__ = "Production" _bits_in_int = int(round(log2(maxint))) def is_nonzero_string_char(char): """Tests whether input character is not '0', returning '1' or '0'.""" if char == '0' or char == '': return '0' else: return '1' def is_nonzero_char(item): """Tests whether input item is nonzero, returning '1' or '0'.""" if isinstance(item, str): return is_nonzero_string_char(item) else: if item: return '1' else: return '0' def seq_to_bitstring(seq): """Converts sequence to string of 1 and 0, returning string of '1'/'0'.""" if not seq: return '' if isinstance(seq, str): return ''.join(map(is_nonzero_string_char, seq)) else: return ''.join(map(is_nonzero_char, seq)) def is_nonzero_string_int(char): """Tests whether input character is not '0', returning 1 or 0.""" if char == '0' or char == '': return 0 else: return 1 def is_nonzero_int(item): """Tests whether input item is nonzero, returning 1 or 0.""" if isinstance(item, str): return is_nonzero_string_int(item) else: if item: return 1 else: return 0 def seq_to_bitlist(seq): """Converts sequence to list of 1 and 0.""" if not seq: return [] if isinstance(seq, str): return map(is_nonzero_string_int, seq) else: return map(is_nonzero_int, seq) def num_to_bitstring(num, length): """Returns string of bits from number, truncated at length. Algorithm: While the number storing the bitvector is greater than zero, find the last digit as (num mod 2) and then bit-shift the string to the right to delete the last digit. Add the digits into a string from right to left (in reverse order). In other words, when length is too small for the number, takes the last n bits of the number in order. Warning: this is not terribly efficient, since the entire number must be rewritten at each shift step. In mutable bitvectors, use caching to reduce the number of conversions required. """ bits = [0] * length #start with all zeroes #successively find the last bit of curr, replacing the #appropriate element of bits as each bit is processed and #removed while num and length: bits[length-1] = 1 & num num >>= 1 length -= 1 return ''.join(map(str, bits)) def bitcount(num, length, value=1): """Counts the bits in num that are set to value (1 or 0, default 1).""" one_count = 0 curr_length = length while num and curr_length: one_count += 1 & num num >>= 1 curr_length -= 1 if value: return one_count else: return length - one_count class ImmutableBitvector(object): """Generic interface for immutable bitvectors.""" def __init__(self, items='', length=None): """Returns new Bitvector.""" if length is not None: self._length = length else: self._length = len(items) def __str__(self): """Returns string representation of bitvector.""" return num_to_bitstring(self, self._length) def __len__(self): """Returns length of the bitvector.""" return self._length def op(self, other, func): """Performs bitwise op on self and other, returning new Bitvector.""" self_length = self._length if not isinstance(other, ImmutableBitvector): #coerce to comparable type try: other = other.freeze() other_length = other._length except: other = Bitvector(other, self_length) other_length = self_length else: other_length = other._length if self_length == other_length: #if they're the same length, just combine the vectors return Bitvector(func(self, other), self_length) else: #need to find which is larger, right-shift, and combine diff = self_length - other_length if diff < 0: #re-bind self and other so self is now longest #tuple unpacking swaps self and other self_length, other_length = other_length, self_length self, other = other, self #shift the shorter vector and combine with bitwise function return Bitvector(func(self >> abs(diff), other), other_length) def __getitem__(self, item): """Returns self[item] as 1 or 0, i.e. as integers and not strings. key should be an index less than the length of the bitvector; negative indexes are handled in the usual manner. Algorithm: uses bit masks to figure out whether the specific position is occupied. Caches masks in Bitvector._masks, so should be very fast when only a few different lengths of bitvector are used in a program. Will be somewhat inefficient if many large bitvectors of different sizes are each used occasionally. """ #check that there actually are items in the vector length = self._length if not length: raise IndexError, "Can't find items in empty vector!" if isinstance(item, slice): return Bitvector(''.join(map(str,[self[i] \ for i in range(*item.indices(length))]))) else: #transform keys to allow negative index if item < 0: item += length #check that key is in bounds if (item < 0) or (item >= length): raise IndexError, "Index %s out of range." % (item,) move_to = length - item - 1 if move_to >= _bits_in_int: result = self & (1L << move_to) else: result = self & (1 << move_to) if result: return 1 else: return 0 def bitcount(self, value=1): """Counts the bits in self that match value.""" return bitcount(self, self._length, value) def __repr__(self): """Produces standard object representation instead of int.""" c = self.__class__ return "<%s.%s object at %s>" % (c.__module__,c.__name__,hex(id(self))) def freeze(self): """Returns ImmutableBitvector containing data from self.""" return self def thaw(self): """Returns MutableBitvector containing data from self.""" return MutableBitvector(self) def stateChanges(self): """Returns list of indices where state changes from 0->1 or 1->0.""" #bail out if no elements rather than raising IndexError length = len(self) if not length: return [] bits = list(self) changes = [i for i in range(1, length) if self[i] != self[i-1]] return [0] + changes + [length] #always implicitly add start and end def divideSequence(self, seq, state=None): """Divides sequence into a list of subsequences, cut using stateChanges. The list will contain the slices for the indices containing the given state. If no state given (state = None) the list will consist of all of the subsequences, cut at the indices. The list and self[0] are returned as a tuple. Truncates whichever sequence is shorter. """ #bail out rather than raising IndexError if vector or sequence is empty if not (len(self) and seq): return ([], 0) cut_list = self.stateChanges() cut_seq = [] first = 0 cut_index = len(cut_list)-1 #If user supplied a specific state, return only pieces in that state if state is not None: #test whether we want to include the first segment or not exclude_start = self[0] != state for i in range(cut_index): if i % 2 == exclude_start: cut_seq.append(seq[cut_list[i]:cut_list[i+1]]) else: #No state supplied: return pieces in all states for i in range(cut_index): cut_seq.append(seq[cut_list[i]:cut_list[i+1]]) first = self[0] return filter(None, cut_seq), first class ShortBitvector(ImmutableBitvector, int): """Short bitvector, stored as an int.""" __slots__ = ['_length'] def __new__(cls, items='', length='ignored'): """Creates new bitvector from sequence of items.""" if isinstance(items, str): if items: #guard against empty string return int.__new__(cls, items, 2) else: return int.__new__(cls, 0) else: return int.__new__(cls, items) def __or__(self, other): """Returns position-wise OR (true if either true).""" if isinstance(other, long): return LongBitvector(self).__or__(other) else: return self.op(other, int.__or__) def __and__(self, other): """Returns position-wise AND (true if both true).""" if isinstance(other, long): return LongBitvector(self).__and__(other) else: return self.op(other, int.__and__) def __xor__(self, other): """Returns position-wise XOR of self and other (true if states differ). """ if isinstance(other, long): return LongBitvector(self).__xor__(other) else: return self.op(other, int.__xor__) def __invert__(self): """Returns complement (replace 1 with 0 and vice versa).""" length = self._length if length >= _bits_in_int: mask = (1L << length) - 1 else: mask = (1 << length) - 1 return Bitvector(mask & ~int(self), length) class LongBitvector(ImmutableBitvector, long): """Long bitvector, stored as a long.""" def __new__(cls, items='', length='ignored'): """Creates new bitvector from sequence of items.""" if isinstance(items, str): if items: #guard against empty string return long.__new__(cls, items, 2) else: return long.__new__(cls, 0) else: return long.__new__(cls, items) def __or__(self, other): """Returns position-wise OR (true if either true).""" return self.op(other, long.__or__) def __and__(self, other): """Returns position-wise AND (true if both true).""" return self.op(other, long.__and__) def __xor__(self, other): """Returns position-wise XOR of self and other (true if states differ). """ return self.op(other, long.__xor__) def __invert__(self): """Returns complement (replace 1 with 0 and vice versa).""" length = self._length return Bitvector(((1L << length) - 1) & ~long(self), length) def Bitvector(items='', length=None, constructor=None): """Factory function returning short or long Bitvector depending on length. Note: uses explict test rather than try/except to fix memory leak. Now is the only way to convert an arbitrary sequence of true/false values into a bitvector. """ #convert whatever was passed as 'items' into a number if isinstance(items, ImmutableBitvector): num = items if length is None: length = len(items) elif isinstance(items, int) or isinstance(items, long): num = items if length is None: raise TypeError, "Must specify length if initializing with number." else: bitstring = seq_to_bitstring(items) if bitstring: num = long(bitstring, 2) if length is None: length = len(items) else: num = 0 length = 0 #if the constructor was not passed explicitly, guess it if not constructor: if 0 <= num <= maxint: constructor = ShortBitvector else: constructor = LongBitvector return constructor(num, length) # Original version with memory leak under Python 2.3 # try: # return ShortBitvector(items, length) # except (OverflowError, TypeError): # return LongBitvector(items, length) class MutableBitvector(Delegator): """Array of bits (0 or 1) supporting set operations. Supports __setitem__. """ def __init__(self, items='', length=None): """Initializes the Bitvector class, storing data in long _vector. Items can be any sequence with members that evaluate to True or False. Private data: _vector: long representing the bitvector _string: string representation of the vector _is_current: boolean indicating whether the string is known to match the long data """ Delegator.__init__(self, None) self.replace(items, length) def replace(self, items='', length=None): """Replaces the contents of self with data in items. Items should be a sequence with elements that evaluate to True or False. If items is another BitVector, uses an efficient method based on the internal data. Otherwise, will loop through items evaluating each to True or False. Primarily used internally, but exposed publically because it's often convenient to replace the contents of an existing bitvector. Usage: vec.update(foo) where foo is a sequence or bitvector. """ vec = Bitvector(items, length) self._handler = vec #make sure attributes go to the new vec self._is_current = False self._string = '' def __str__(self): """Prints the bits in the vector as a string of '1' and '0'.""" #check that the string is up to date; when it is, return it. if not self._is_current: self._string = str(self._handler) self._is_current = True return self._string def __len__(self): """Need to pass directly to handler.""" return len(self._handler) def __getitem__(self, *args): """Gets item from vector: needs to be passed through explicitly.""" return self._handler.__getitem__(*args) def __setitem__(self, item, value): """Sets self[item] to value: value must be 0 or 1, or '0' or '1'. Uses masks to alter the bit at a specific position. To set a bit to zero, make a string that is 1 everywhere except at that position and then use bitwise &. To set a bit to 1, make a string that is 0 everywhere except at that position and then use bitwise |. Warning: this is quite a slow operation, especially on long vectors. """ length = self._length #handle slice assignment. Warning: does not allow change in length! if isinstance(item, slice): for i, val in zip(range(*item.indices(length)), value): self[i] = is_nonzero_int(val) return #otherwise, check that the key is in range vec = self._handler curr_val = vec[item] val = is_nonzero_int(value) if val != curr_val: #need only do anything if value changed! #find index of negative values if item < 0: item += length #check that key is in bounds if (item < 0) or (item >= length): raise IndexError, "Index %s out of range." % (item,) #figure out offset move_to = length - item - 1 if value == 0: #make mask of '1', and combine with & result = vec &((1L << length) - 1) ^ (1L << move_to) elif value == 1: #make mask of '0', and combine with | result = vec | (1L << move_to) else: #complain if we got anything else raise ValueError, "Item not 0 or 1: " + str(value) self.replace(result) def __cmp__(self, other): """Comparison needs to be passed through explicitly.""" return cmp(self._handler, other) def op(self, other, func): """Returns position-wise op for self and other. self[i] == 1 if op(self[i], other[i]) == 1 Internal method: used to implement bitwise operations on pairs of bitvectors. Truncates the longer vector to fit the shorter one by right-shifting, i.e. the leftmost bits of the longer vector will be combined with all the bits of the shorter vector. Handles the case where both the operands are bitvectors specially for speed, to greatest advantage when they are equal length. Uses bool() to convert elements of sequences, in preference to checking whether they directly represent 0 or 1. This makes the operations more flexible, but beware of bad data! """ return self.__class__(self._handler.op(other, func)) def __or__(self, other): """Returns position-wise OR (true if either true). """ return self.op(other, or_) def __and__(self, other): """Returns position-wise AND (true if both true). """ return self.op(other, and_) def __xor__(self, other): """Returns position-wise XOR of self and other (true if states differ). """ return self.op(other, xor) def __invert__(self): """Returns complement (replace 1 with 0 and vice versa). """ return self.__class__(~self._handler) def thaw(self): """Returns mutable version of self.""" return self def freeze(self): """Returns immutable version of self.""" return self._handler def VectorFromCases(seq, constructor=LongBitvector): """Returns vector with 0(1) for lower(upper)case positions in sequence. Primarily used for identifying binding sites denoted by uppercase bases. """ return constructor(''.join(map(str, map(int, [i.isupper() for i in seq])))) def VectorFromMatches(seq, pattern, overlapping=1, constructor=LongBitvector): """Replaces self with 1(0) for each position in sequence (not) matching. Usage: vec.toMatches('agagac', 'aga', 1) == vec('111110') vec.toMatches('agagac', 'aga', 0) == vec('111000') The first argument must be a type that is searchable by a re object, typically a string. The second argument must be a pattern to search: this can be a string, a list of strings, or an re object. The third, optional argument specifies whether matches may overlap or not. If a list is passed, its elements will be converted into strings, and then into a regular expression using alternation. Lists of regular expressions are not supported directly. Zero-length patterns will typically never be counted as matching, rather than being counted as matching everywhere. """ #allocate space for the results seqlength = len(seq) result = [0] * seqlength if seqlength: #will return empty string if seq empty if isinstance(pattern, str): #handle strings by using the index string method patlength = len(pattern) #will skip zero-length patterns if patlength: #until we reach the end of the sequence, find the next #index that contains a match. When a match has been #processed by inserting '1' into the list at each position #that matches, either move to the next position (if matches #can overlap) or to the position after the end of the match #(if matches cannot overlap). start = 0 while start + patlength - 1 <= seqlength: try: match = seq.index(pattern, start) except ValueError: break #no more matches #insert '1' into the list in the appropriate slice result[match:match + patlength] = [1] * patlength #move to the first index where the next match could be if overlapping: start = match + 1 else: start = match + patlength else: #if not a string, assume regex #first, attempt to alternate the items in pattern as though it #were a sequence type try: pattern = re.compile('|'.join(map(str, pattern))) except: pass #use same strategy as above to locate matches of pattern, find #the lengths, insert '1' into the list at the appropriate #places, and move on to the next position where a match might #start. start = 0 while start < seqlength: match = pattern.search(seq, start) try: #find the length of the match, and the start and end #indexes. Use these to define the positions where '1' #will appear in the result. match_start, match_end = match.span() match_length = match_end - match_start result[match_start:match_end] = [1] * match_length #figure out how many positions to advance for the next #possible match start if overlapping or (match_length == 0): start = match_start + 1 else: start = match_end except AttributeError, TypeError: #match is None if no more in string break else: #sequence was zero-length result = [] return constructor(''.join(map(str, result))) def VectorFromRuns(seq, length, constructor=LongBitvector): """Returns vector with 1 at positions specified by sequence of (start,len). seq should be a sequence of 2-item sequences, where the first element is the index where the run of 1's starts and the second element is the number of 1's in the run. length should be an int giving the total length of the sequence. """ if not length: return constructor('') bits = ['0']*length for start, run in seq: if start + run > length: raise IndexError, "start %s + run %s exceeds seq length %s" % \ (start, run, length) bits[start:start+run] = ['1']*run return constructor(''.join(bits)) def VectorFromSpans(seq, length, constructor=LongBitvector): """Returns vector with 1 at positions specified by sequence of (start,end). seq should be a sequence of 2-item sequences, where the first element is the index where the run of 1's starts and the second element is the index where the run stops (standard Python slice notation, i.e. the start and stop you would normally use in a slice, where the slice does not include the element at i[stop]). length should be an int giving the total length of the sequence. """ if not length: return constructor('') bits = ['0']*length for start, stop in seq: if stop > length: raise IndexError, "stop %s exceeds seq length %s" % (stop, length) run = stop - start bits[start:stop] = ['1']*run return constructor(''.join(bits)) def VectorFromPositions(seq, length, constructor=LongBitvector): """Returns vector with 1 at positions specified by sequence of positions. seq should be a sequence of ints (position). length should be an int giving the total length of the sequence. """ if not length: return constructor('') bits = ['0']*length for i in seq: bits[i] = '1' return constructor(''.join(bits)) class PackedBases(LongBitvector): """Stores unambiguous nucleotide sequences as bitvectors, 4 per byte. Rna controls whether __str__ produces RNA or DNA. Each base is stored as 2 bits. The first bit encodes purine/pyrmidine, while the second encodes weak/strong. NOTE: len() returns the length in _bits_, not the length in bases. """ #translation table to turn bases into sequences. First bit encodes purine/ #pyrimindine; second bit encodes weak/strong. _sequence_to_bits = maketrans('AGUTCagutc', '0122301223') _rna_bases = {0:'A', 1:'G', 2:'U', 3:'C'} _dna_bases = {0:'A', 1:'G', 2:'T', 3:'C'} def __new__(cls, sequence='', Rna='ignored'): """Packs bases in sequence into a 2-bit per base bitvector. Usage: vec.toBases('AcgcaAc') == vec('00110111000011') Encoding is to use the first bit as purine vs pyrimidine (purine = 0), and to use the second bit as weak vs. strong (weak = 0). This means that A == 00, G == 01, U == 10, and C == 11. """ seq = sequence.translate(cls._sequence_to_bits) if not seq: return long.__new__(cls, 0L) else: return long.__new__(cls, long(seq, 4)) #note base 4 conversion def __init__(self, sequence='', Rna=True): """Returns new PackedBases object.""" self._length = len(sequence)*2 self.Rna = Rna def __str__(self): """Unpacks bases from sequence using encoding scheme in toBases.""" length = self._length/2 bits = [0] * length #allocate space using '00' pattern curr = self #successively find the last two bit of curr, replacing the #appropriate element of bits as each bit is processed and #removed while curr and length: bits[length-1] = curr & 3 curr >>= 2 length -= 1 #convert numbers back into string if self.Rna: num2base = self._rna_bases else: num2base = self._dna_bases return ''.join([num2base[i] for i in bits])