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