Source Code for Module rdkit.DataStructs.BitVect

  1  # $Id: BitVect.py 997 2009-02-25 06:12:43Z glandrum $ 
  2  # 
  3  #  Copyright (C) 2001-2006  greg Landrum and Rational Discovery LLC 
  4  # 
  5  #   @@ All Rights Reserved  @@ 
  6  # 
  7  """ This has all be re-implemented in the C++ code 
  8  """ 
  9  import math 
 10   
 11 -class BitVect: 
 12 -  def __init__(self,nBits): 
 13      self.nBits = nBits 
 14      self.bits = [0]*nBits 
 15   
 16 -  def NumOnBits(self): 
 17      return len(self.GetOnBits()) 
 18   
 19 -  def GetOnBits(self,sort=1,reverse=0): 
 20      l = [idx for idx in xrange(self.nBits) if self.bits[idx] == 1] 
 21      if reverse: 
 22        l.reverse() 
 23      return l 
 24     
 25 -  def TanimotoSimilarity(self,other): 
 26      if not isinstance(other,BitVect): 
 27        raise TypeError,"Tanimoto similarities can only be calculated between two BitVects" 
 28      if len(self)!=len(other): 
 29        raise ValueError,"BitVects must be the same length" 
 30      bc = len(self & other) 
 31      b1 = self.NumOnBits() 
 32      b2 = other.NumOnBits() 
 33      return float(bc) / float(b1 + b2 - bc) 
 34    TanimotoSimilarity = TanimotoSimilarity 
 35   
 36 -  def EuclideanDistance(self,other): 
 37      if not isinstance(other,BitVect): 
 38        raise TypeError,"Tanimoto similarities can only be calculated between two BitVects" 
 39      bt = len(self) 
 40      bi = len(self ^ (~ other)) 
 41      return math.sqrt(bt-bi)/bt 
 42     
 43 -  def __getitem__(self,which): 
 44      if which >= self.nBits or which < 0: 
 45        raise ValueError,'bad index' 
 46      return self.bits[which] 
 47   
 48 -  def __setitem__(self,which,val): 
 49      if which >= self.nBits or which < 0: 
 50        raise ValueError,'bad index' 
 51      if val not in [0,1]: 
 52        raise ValueError,'val must be 0 or 1' 
 53   
 54      self.bits[which] = val 
 55         
 56 -  def __len__(self): 
 57      return self.nBits 
 58   
 59 -  def __and__(self,other): 
 60      if not isinstance(other,BitVect): 
 61        raise TypeError,"BitVects can only be &'ed with other BitVects" 
 62      if len(self) != len(other): 
 63        raise ValueError,"BitVects must be of the same length" 
 64   
 65      l1 = self.GetOnBits() 
 66      l2 = other.GetOnBits() 
 67      r = [bit for bit in l1 if bit in l2] 
 68      return r 
 69   
 70 -  def __or__(self,other): 
 71      if not isinstance(other,BitVect): 
 72        raise TypeError,"BitVects can only be |'ed with other BitVects" 
 73      if len(self) != len(other): 
 74        raise ValueError,"BitVects must be of the same length" 
 75      l1 = self.GetOnBits() 
 76      l2 = other.GetOnBits() 
 77      r = l1 + [bit for bit in l2 if bit not in l1] 
 78      r.sort() 
 79      return r 
 80   
 81 -  def __xor__(self,other): 
 82      if not isinstance(other,BitVect): 
 83        raise TypeError,"BitVects can only be ^'ed with other BitVects" 
 84      if len(self) != len(other): 
 85        raise ValueError,"BitVects must be of the same length" 
 86       
 87      l1 = self.GetOnBits() 
 88      l2 = other.GetOnBits() 
 89      r = [bit for bit in l1 if bit not in l2] + [bit for bit in l2 if bit not in l1] 
 90      r.sort() 
 91      return r 
 92   
 93 -  def __invert__(self): 
 94      res = BitVect(len(self)) 
 95      for i in xrange(len(self)): 
 96        res[i] = not self[i] 
 97      return res 
 98     
 99 -class SparseBitVect(BitVect): 
100 -  def __init__(self,nBits): 
101      self.nBits = nBits 
102      self.bits = [] 
103   
104 -  def NumOnBits(self): 
105      return len(self.bits) 
106   
107 -  def GetOnBits(self,sort=1,reverse=0): 
108      l = self.bits[:] 
109      if sort: 
110        l.sort() 
111      if reverse: 
112        l.reverse() 
113      return l 
114     
115 -  def __getitem__(self,which): 
116      if which >= self.nBits or which < 0: 
117        raise ValueError,'bad index' 
118      if which in self.bits: 
119        return 1 
120      else: 
121        return 0 
122   
123 -  def __setitem__(self,which,val): 
124      if which >= self.nBits or which < 0: 
125        raise ValueError,'bad index' 
126      if val == 0: 
127        if which in self.bits: 
128          self.bits.remove(which) 
129      else: 
130        self.bits.append(which) 
131         
132 -  def __len__(self): 
133      return self.nBits 
134   
135  if __name__ == '__main__': 
136    b1 = BitVect(10) 
137    b2 = SparseBitVect(10) 
138    b1[0] = 1 
139    b2[0] = 1 
140    b1[3] = 1 
141    b2[4] = 1 
142    b2[5] = 1 
143    b2[5] = 0 
144    print 'b1:',b1.GetOnBits() 
145    print 'b2:',b2.GetOnBits() 
146    print '&:', b1 & b2 
147    print '|:', b1 | b2 
148    print '^:', b1 ^ b2 
149    print 'b1.Tanimoto(b2):',b1.TanimotoSimilarity(b2) 
150    print 'b1.Tanimoto(b1):',b1.TanimotoSimilarity(b1) 
151    print 'b2.Tanimoto(b2):',b2.TanimotoSimilarity(b2) 
152    print 'b2.Tanimoto(b1):',b2.TanimotoSimilarity(b1) 
153