Package rdkit :: Package ML :: Package DecTree :: Module TreeUtils

[hide private]
[frames] | no frames]

# Source Code for Module rdkit.ML.DecTree.TreeUtils

      1  #
      2  #  Copyright (C) 2001-2004  greg Landrum and Rational Discovery LLC
      3  #  All Rights Reserved
      4  #
      5  """ Utilities for working with trees
      6
      7  """
      8
      9 -def CollectLabelLevels(tree,levels,level=0,maxDepth=1e8):
     10    if level < maxDepth:
     11      if not tree.GetTerminal():
     12        l = tree.GetLabel()
     13        currLevel = levels.get(l,1e8)
     14        if level < currLevel:
     15          levels[l] = level
     16        for child in tree.GetChildren():
     17          CollectLabelLevels(child,levels,level+1,maxDepth)
     18    return levels
     19
     20 -def CollectDescriptorNames(tree,names,level=0,maxDepth=1e8):
     21    if level < maxDepth:
     22      if not tree.GetTerminal():
     23        names[tree.GetLabel()] = tree.GetName()
     24        for child in tree.GetChildren():
     25          CollectDescriptorNames(child,names,level+1,maxDepth)
     26    return names
     27
     28  #------------------------------------
     29  #
     30  #  doctest boilerplate
     31  #
     32  _test1="""
     33  >>> from DecTree import DecTreeNode as Node
     34  >>> t1 = Node(None,'d1',1)
     35  >>> t2 = Node(None,'d2',2)
     36  >>> t1.AddChildNode(t2)
     37  >>> t2 = Node(None,'d3',3)
     38  >>> t1.AddChildNode(t2)
     39  >>> t3 = Node(None,'d4',4)
     40  >>> t2.AddChildNode(t3)
     41  >>> t3 = Node(None,'d2',2)
     42  >>> t2.AddChildNode(t3)
     43  >>> r = CollectLabelLevels(t1,{})
     44  >>> r[2]
     45  1
     46  >>> r[1]
     47  0
     48  >>> r[3]
     49  1
     50  >>> r[4]
     51  2
     52  >>> r = CollectLabelLevels(t1,{},0,2)
     53  >>> r[2]
     54  1
     55  >>> r[1]
     56  0
     57  >>> r[3]
     58  1
     59  >>> 4 in r
     60  0
     61
     62  Check that we can handle subtrees:
     63  >>> r = CollectLabelLevels(t1,{},1,2)
     64  >>> r[1]
     65  1
     66  >>> 2 in r
     67  0
     68  >>> 3 in r
     69  0
     70  >>> 4 in r
     71  0
     72
     73  >>> names = CollectDescriptorNames(t1,{})
     74  >>> names[1]
     75  'd1'
     76  >>> names[2]
     77  'd2'
     78  >>> names[3]
     79  'd3'
     80  >>> names[4]
     81  'd4'
     82
     83  >>> names = CollectDescriptorNames(t1,{},0,2)
     84  >>> names[1]
     85  'd1'
     86  >>> names[2]
     87  'd2'
     88  >>> names[3]
     89  'd3'
     90  >>> 4 in names
     91  0
     92
     93  >>> names = CollectDescriptorNames(t1,{},1,2)
     94  >>> names[1]
     95  'd1'
     96  >>> 2 in names
     97  0
     98  >>> 3 in names
     99  0
    100  >>> 4 in names
    101  0
    102
    103
    104  """
    105
    106  __test__={'_test1':_test1}
    107 -def _test():
    108    import doctest,sys
    109    return doctest.testmod(sys.modules["__main__"])
    110
    111  if __name__ == '__main__':
    112    import sys
    113    failed,tried = _test()
    114    sys.exit(failed)
    115