Package rdkit :: Package DataStructs :: Module BitVect
[hide private]
[frames] | no frames]

Source Code for Module rdkit.DataStructs.BitVect

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