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