1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
|
from typing import Any, Optional
import numpy as np
from sklearn.tree import _tree # type: ignore[import-untyped]
class DecisionTreeNode:
def __init__(
self,
feature: Optional[str] = None,
threshold: Optional[float] = None,
left: Optional["DecisionTreeNode"] = None,
right: Optional["DecisionTreeNode"] = None,
class_probs: Any = None,
num_samples: int = 0,
node_id: int = 0,
) -> None:
self.feature = feature
self.threshold = threshold
self.left = left
self.right = right
self.class_probs = class_probs
self.num_samples = num_samples
self.id = node_id
def is_leaf(self) -> bool:
return self.left is None or self.right is None
class DecisionTree:
"""
Custom decision tree implementation that mimics some of the sklearn API.
The purpose of this class it to be able to perform transformations, such as custom pruning, which
does not seem to be easy with sklearn.
"""
def __init__(self, sklearn_tree: Any, feature_names: list[str]) -> None:
self.feature_names = feature_names
self.root = self._convert_sklearn_tree(sklearn_tree.tree_)
self.classes_: list[str] = sklearn_tree.classes_
def _convert_sklearn_tree(
self, sklearn_tree: Any, node_id: int = 0
) -> DecisionTreeNode:
class_probs = sklearn_tree.value[node_id][0]
num_samples = sklearn_tree.n_node_samples[node_id]
if sklearn_tree.feature[node_id] != _tree.TREE_UNDEFINED:
feature_index = sklearn_tree.feature[node_id]
feature = self.feature_names[feature_index]
left = self._convert_sklearn_tree(
sklearn_tree, sklearn_tree.children_left[node_id]
)
right = self._convert_sklearn_tree(
sklearn_tree, sklearn_tree.children_right[node_id]
)
return DecisionTreeNode(
feature=feature,
threshold=sklearn_tree.threshold[node_id],
left=left,
right=right,
class_probs=class_probs,
num_samples=num_samples,
node_id=node_id,
)
else:
return DecisionTreeNode(
class_probs=class_probs, num_samples=num_samples, node_id=node_id
)
def prune(self, df: Any, target_col: str, k: int) -> None:
self.root = self._prune_tree(self.root, df, target_col, k)
def _prune_tree(
self, node: DecisionTreeNode, df: Any, target_col: str, k: int
) -> DecisionTreeNode:
if node.is_leaf():
return node
left_df = df[df[node.feature] <= node.threshold]
right_df = df[df[node.feature] > node.threshold]
# number of unique classes in the left and right subtrees
left_counts = left_df[target_col].nunique()
right_counts = right_df[target_col].nunique()
# for ranking, we want to ensure that we return at least k classes, so if we have less than k classes in the
# left or right subtree, we remove the split and make this node a leaf node
if left_counts < k or right_counts < k:
return DecisionTreeNode(class_probs=node.class_probs)
assert node.left is not None, "expected left child to exist"
node.left = self._prune_tree(node.left, left_df, target_col, k)
assert node.right is not None, "expected right child to exist"
node.right = self._prune_tree(node.right, right_df, target_col, k)
return node
def to_dot(self) -> str:
dot = "digraph DecisionTree {\n"
dot += ' node [fontname="helvetica"];\n'
dot += ' edge [fontname="helvetica"];\n'
dot += self._node_to_dot(self.root)
dot += "}"
return dot
def _node_to_dot(
self, node: DecisionTreeNode, parent_id: int = 0, edge_label: str = ""
) -> str:
if node is None:
return ""
node_id = id(node)
# Format class_probs array with line breaks
class_probs_str = self._format_class_probs_array(
node.class_probs, node.num_samples
)
if node.is_leaf():
label = class_probs_str
shape = "box"
else:
feature_name = f"{node.feature}"
label = f"{feature_name} <= {node.threshold:.2f}\\n{class_probs_str}"
shape = "oval"
dot = f' {node_id} [label="{label}", shape={shape}];\n'
if parent_id != 0:
dot += f' {parent_id} -> {node_id} [label="{edge_label}"];\n'
if not node.is_leaf():
assert node.left is not None, "expected left child to exist"
dot += self._node_to_dot(node.left, node_id, "<=")
assert node.right is not None, "expected right child to exist"
dot += self._node_to_dot(node.right, node_id, ">")
return dot
def _format_class_prob(self, num: float) -> str:
if num == 0:
return "0"
return f"{num:.2f}"
def _format_class_probs_array(
self, class_probs: Any, num_samples: int, max_per_line: int = 5
) -> str:
# add line breaks to avoid very long lines
flat_class_probs = class_probs.flatten()
formatted = [self._format_class_prob(v) for v in flat_class_probs]
lines = [
formatted[i : i + max_per_line]
for i in range(0, len(formatted), max_per_line)
]
return f"num_samples={num_samples}\\n" + "\\n".join(
[", ".join(line) for line in lines]
)
def predict(self, X: Any) -> Any:
predictions = [self._predict_single(x) for _, x in X.iterrows()]
return np.array(predictions)
def predict_proba(self, X: Any) -> Any:
return np.array([self._predict_proba_single(x) for _, x in X.iterrows()])
def _get_leaf(self, X: Any) -> DecisionTreeNode:
node = self.root
while not node.is_leaf():
if X[node.feature] <= node.threshold:
assert node.left is not None, "expected left child to exist"
node = node.left
else:
assert node.right is not None, "expected right child to exist"
node = node.right
return node
def _predict_single(self, x: Any) -> str:
node = self._get_leaf(x)
# map index to class name
return self.classes_[np.argmax(node.class_probs)]
def _predict_proba_single(self, x: Any) -> Any:
node = self._get_leaf(x)
return node.class_probs
def apply(self, X: Any) -> Any:
ids = [self._apply_single(x) for _, x in X.iterrows()]
return np.array(ids)
def _apply_single(self, x: Any) -> int:
node = self._get_leaf(x)
return node.id
def codegen(
self,
dummy_col_2_col_val: dict[str, tuple[str, Any]],
lines: list[str],
unsafe_leaves: list[int],
) -> None:
# generates python code for the decision tree
def codegen_node(node: DecisionTreeNode, depth: int) -> None:
indent = " " * (depth + 1)
if node.is_leaf():
lines.append(handle_leaf(node, indent, unsafe_leaves))
else:
name = node.feature
threshold = node.threshold
if name in dummy_col_2_col_val:
(orig_name, value) = dummy_col_2_col_val[name]
predicate = f"{indent}if str(context.get_value('{orig_name}')) != '{value}':"
assert (
threshold == 0.5
), f"expected threshold to be 0.5 but is {threshold}"
else:
predicate = (
f"{indent}if context.get_value('{name}') <= {threshold}:"
)
lines.append(predicate)
assert node.left is not None, "expected left child to exist"
codegen_node(node.left, depth + 1)
lines.append(f"{indent}else:")
assert node.right is not None, "expected right child to exist"
codegen_node(node.right, depth + 1)
def handle_leaf(
node: DecisionTreeNode, indent: str, unsafe_leaves: list[int]
) -> str:
"""
This generates the code for a leaf node in the decision tree. If the leaf is unsafe, the learned heuristic
will return "unsure" (i.e. None).
"""
if node.id in unsafe_leaves:
return f"{indent}return None"
class_probas = node.class_probs
return f"{indent}return {best_probas_and_indices(class_probas)}"
def best_probas_and_indices(class_probas: Any) -> str:
"""
Given a list of tuples (proba, idx), this function returns a string in which the tuples are
sorted by proba in descending order. E.g.:
Given class_probas=[(0.3, 0), (0.5, 1), (0.2, 2)]
this function returns
"[(0.5, 1), (0.3, 0), (0.2, 2)]"
"""
# we generate a list of tuples (proba, idx) sorted by proba in descending order
# idx is the index of a choice
# we only generate a tuple if proba > 0
probas_indices_sorted = sorted(
[
(proba, index)
for index, proba in enumerate(class_probas)
if proba > 0
],
key=lambda x: x[0],
reverse=True,
)
probas_indices_sorted_str = ", ".join(
f"({value:.3f}, {index})" for value, index in probas_indices_sorted
)
return f"[{probas_indices_sorted_str}]"
codegen_node(self.root, 1)
|
