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   
10 -def CollectLabelLevels(tree, levels, level=0, maxDepth=1e8): 
11    if level < maxDepth: 
12      if not tree.GetTerminal(): 
13        l = tree.GetLabel() 
14        currLevel = levels.get(l, 1e8) 
15        if level < currLevel: 
16          levels[l] = level 
17        for child in tree.GetChildren(): 
18          CollectLabelLevels(child, levels, level + 1, maxDepth) 
19    return levels 
20   
21   
22 -def CollectDescriptorNames(tree, names, level=0, maxDepth=1e8): 
23    if level < maxDepth: 
24      if not tree.GetTerminal(): 
25        names[tree.GetLabel()] = tree.GetName() 
26        for child in tree.GetChildren(): 
27          CollectDescriptorNames(child, names, level + 1, maxDepth) 
28    return names 
29