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
|
/*
* Licensed to the Apache Software Foundation (ASF) under one or more
* contributor license agreements. See the NOTICE file distributed with
* this work for additional information regarding copyright ownership.
* The ASF licenses this file to You under the Apache License, Version 2.0
* (the "License"); you may not use this file except in compliance with
* the License. You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package org.apache.spark.ml.tree
import java.util.Objects
import org.apache.spark.annotation.{DeveloperApi, Since}
import org.apache.spark.mllib.linalg.Vector
import org.apache.spark.mllib.tree.configuration.{FeatureType => OldFeatureType}
import org.apache.spark.mllib.tree.model.{Split => OldSplit}
/**
* :: DeveloperApi ::
* Interface for a "Split," which specifies a test made at a decision tree node
* to choose the left or right path.
*/
@DeveloperApi
sealed trait Split extends Serializable {
/** Index of feature which this split tests */
def featureIndex: Int
/**
* Return true (split to left) or false (split to right).
* @param features Vector of features (original values, not binned).
*/
private[ml] def shouldGoLeft(features: Vector): Boolean
/**
* Return true (split to left) or false (split to right).
* @param binnedFeature Binned feature value.
* @param splits All splits for the given feature.
*/
private[tree] def shouldGoLeft(binnedFeature: Int, splits: Array[Split]): Boolean
/** Convert to old Split format */
private[tree] def toOld: OldSplit
}
private[tree] object Split {
def fromOld(oldSplit: OldSplit, categoricalFeatures: Map[Int, Int]): Split = {
oldSplit.featureType match {
case OldFeatureType.Categorical =>
new CategoricalSplit(featureIndex = oldSplit.feature,
_leftCategories = oldSplit.categories.toArray, categoricalFeatures(oldSplit.feature))
case OldFeatureType.Continuous =>
new ContinuousSplit(featureIndex = oldSplit.feature, threshold = oldSplit.threshold)
}
}
}
/**
* :: DeveloperApi ::
* Split which tests a categorical feature.
* @param featureIndex Index of the feature to test
* @param _leftCategories If the feature value is in this set of categories, then the split goes
* left. Otherwise, it goes right.
* @param numCategories Number of categories for this feature.
*/
@DeveloperApi
final class CategoricalSplit private[ml] (
override val featureIndex: Int,
_leftCategories: Array[Double],
@Since("2.0.0") val numCategories: Int)
extends Split {
require(_leftCategories.forall(cat => 0 <= cat && cat < numCategories), "Invalid leftCategories" +
s" (should be in range [0, $numCategories)): ${_leftCategories.mkString(",")}")
/**
* If true, then "categories" is the set of categories for splitting to the left, and vice versa.
*/
private val isLeft: Boolean = _leftCategories.length <= numCategories / 2
/** Set of categories determining the splitting rule, along with [[isLeft]]. */
private val categories: Set[Double] = {
if (isLeft) {
_leftCategories.toSet
} else {
setComplement(_leftCategories.toSet)
}
}
override private[ml] def shouldGoLeft(features: Vector): Boolean = {
if (isLeft) {
categories.contains(features(featureIndex))
} else {
!categories.contains(features(featureIndex))
}
}
override private[tree] def shouldGoLeft(binnedFeature: Int, splits: Array[Split]): Boolean = {
if (isLeft) {
categories.contains(binnedFeature.toDouble)
} else {
!categories.contains(binnedFeature.toDouble)
}
}
override def hashCode(): Int = {
val state = Seq(featureIndex, isLeft, categories)
state.map(Objects.hashCode).foldLeft(0)((a, b) => 31 * a + b)
}
override def equals(o: Any): Boolean = o match {
case other: CategoricalSplit => featureIndex == other.featureIndex &&
isLeft == other.isLeft && categories == other.categories
case _ => false
}
override private[tree] def toOld: OldSplit = {
val oldCats = if (isLeft) {
categories
} else {
setComplement(categories)
}
OldSplit(featureIndex, threshold = 0.0, OldFeatureType.Categorical, oldCats.toList)
}
/** Get sorted categories which split to the left */
def leftCategories: Array[Double] = {
val cats = if (isLeft) categories else setComplement(categories)
cats.toArray.sorted
}
/** Get sorted categories which split to the right */
def rightCategories: Array[Double] = {
val cats = if (isLeft) setComplement(categories) else categories
cats.toArray.sorted
}
/** [0, numCategories) \ cats */
private def setComplement(cats: Set[Double]): Set[Double] = {
Range(0, numCategories).map(_.toDouble).filter(cat => !cats.contains(cat)).toSet
}
}
/**
* :: DeveloperApi ::
* Split which tests a continuous feature.
* @param featureIndex Index of the feature to test
* @param threshold If the feature value is <= this threshold, then the split goes left.
* Otherwise, it goes right.
*/
@DeveloperApi
final class ContinuousSplit private[ml] (override val featureIndex: Int, val threshold: Double)
extends Split {
override private[ml] def shouldGoLeft(features: Vector): Boolean = {
features(featureIndex) <= threshold
}
override private[tree] def shouldGoLeft(binnedFeature: Int, splits: Array[Split]): Boolean = {
if (binnedFeature == splits.length) {
// > last split, so split right
false
} else {
val featureValueUpperBound = splits(binnedFeature).asInstanceOf[ContinuousSplit].threshold
featureValueUpperBound <= threshold
}
}
override def equals(o: Any): Boolean = {
o match {
case other: ContinuousSplit =>
featureIndex == other.featureIndex && threshold == other.threshold
case _ =>
false
}
}
override def hashCode(): Int = {
val state = Seq(featureIndex, threshold)
state.map(Objects.hashCode).foldLeft(0)((a, b) => 31 * a + b)
}
override private[tree] def toOld: OldSplit = {
OldSplit(featureIndex, threshold, OldFeatureType.Continuous, List.empty[Double])
}
}
|
