Recursive implementation of decision tree.
| 10 | |
| 11 | |
| 12 | class Tree(object): |
| 13 | """Recursive implementation of decision tree.""" |
| 14 | |
| 15 | def __init__(self, regression=False, criterion=None, n_classes=None): |
| 16 | self.regression = regression |
| 17 | self.impurity = None |
| 18 | self.threshold = None |
| 19 | self.column_index = None |
| 20 | self.outcome = None |
| 21 | self.criterion = criterion |
| 22 | self.loss = None |
| 23 | self.n_classes = n_classes # Only for classification |
| 24 | |
| 25 | self.left_child = None |
| 26 | self.right_child = None |
| 27 | |
| 28 | @property |
| 29 | def is_terminal(self): |
| 30 | return not bool(self.left_child and self.right_child) |
| 31 | |
| 32 | def _find_splits(self, X): |
| 33 | """Find all possible split values.""" |
| 34 | split_values = set() |
| 35 | |
| 36 | # Get unique values in a sorted order |
| 37 | x_unique = list(np.unique(X)) |
| 38 | for i in range(1, len(x_unique)): |
| 39 | # Find a point between two values |
| 40 | average = (x_unique[i - 1] + x_unique[i]) / 2.0 |
| 41 | split_values.add(average) |
| 42 | |
| 43 | return list(split_values) |
| 44 | |
| 45 | def _find_best_split(self, X, target, n_features): |
| 46 | """Find best feature and value for a split. Greedy algorithm.""" |
| 47 | |
| 48 | # Sample random subset of features |
| 49 | subset = random.sample(list(range(0, X.shape[1])), n_features) |
| 50 | max_gain, max_col, max_val = None, None, None |
| 51 | |
| 52 | for column in subset: |
| 53 | split_values = self._find_splits(X[:, column]) |
| 54 | for value in split_values: |
| 55 | if self.loss is None: |
| 56 | # Random forest |
| 57 | splits = split(X[:, column], target["y"], value) |
| 58 | gain = self.criterion(target["y"], splits) |
| 59 | else: |
| 60 | # Gradient boosting |
| 61 | left, right = split_dataset( |
| 62 | X, target, column, value, return_X=False |
| 63 | ) |
| 64 | gain = xgb_criterion(target, left, right, self.loss) |
| 65 | |
| 66 | if (max_gain is None) or (gain > max_gain): |
| 67 | max_col, max_val, max_gain = column, value, gain |
| 68 | return max_col, max_val, max_gain |
| 69 |