Skip to content

Instantly share code, notes, and snippets.

@eherrerosj
Created March 15, 2018 23:21
Show Gist options
  • Save eherrerosj/ed3400b06f3c5c7668c62653e2a695c2 to your computer and use it in GitHub Desktop.
Save eherrerosj/ed3400b06f3c5c7668c62653e2a695c2 to your computer and use it in GitHub Desktop.
Generate if-else statements out of a DecisionTree classifier
from sklearn.tree import _tree
def tree_to_code(tree, feature_names):
tree_ = tree.tree_
feature_name = [
feature_names[i] if i != _tree.TREE_UNDEFINED else "undefined!"
for i in tree_.feature
]
print ("def tree({}):".format(", ".join(feature_names)))
def recurse(node, depth):
indent = " " * depth
if tree_.feature[node] != _tree.TREE_UNDEFINED:
name = feature_name[node]
threshold = tree_.threshold[node]
print ("{}if {} <= {}:".format(indent, name, threshold))
recurse(tree_.children_left[node], depth + 1)
print ("{}else: # if {} > {}".format(indent, name, threshold))
recurse(tree_.children_right[node], depth + 1)
else:
print ("{}return {}".format(indent, np.argmax(tree_.value[node])))
recurse(0, 1)
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment