对有序特征进行离散化(继承Spark的机器学习Estimator类)
采用信息增益或基尼指数寻找最优离散化点
package org.apache.spark.ml.feature
import org.apache.spark.sql.SparkSession
import org.apache.spark.sql.functions._
import scala.collection.mutable.ArrayBuffer
import org.apache.spark.sql.Dataset
import org.apache.spark.sql.types._
import org.apache.spark.ml.param._
import org.apache.spark.ml.util.Identifiable
import org.apache.spark.ml.param.shared._
import org.apache.spark.ml.util._
import org.apache.spark.ml.Estimator
import org.apache.spark.ml.attribute._
/**
* 连续有序特征离散类。
*
* 相关参数参照决策树同名参数
*
*
* 采用二分,每次分割都对一个或多个分段进行二分。寻找信息增益(或基尼指数)最大的分割点。
*
* 停止条件:达到指定分段数(numBuckets),或分割的后信息增益小于指定值(minInfoGain)
*
*/
private object DiscretizerTest {
def main(args: Array[String]): Unit = {
val time1 = System.currentTimeMillis()
val spark = spark = SparkSession.builder().getOrCreate()
import spark.implicits._
val inputCol1 = "f1"
val inputCol2 = "f2"
val labelCol = "label"
val outputCol1 = "discretizer1"
val outputCol2 = "discretizer2"
val train = spark.createDataFrame(
List(
(1, 2.3, 0),
(2, 8.1, 0),
(3, 1.1, 1),
(4, 2.2, 1),
(5, 3.3, 0),
(6, 7.0, 1))).toDF(inputCol1, inputCol2, labelCol)
val test = spark.createDataFrame(
List(
(1, 7),
(7, 8))).toDF(inputCol1, inputCol2)
val discretizer = new Discretizer().
setInputCols(Array(inputCol1, inputCol2)).
setOutputCols(Array(outputCol1, outputCol2)).
setNumBucketsArray(Array(2, 3)).
setLabelCol(labelCol).
setMinInstancesPerBucket(1)
val model = discretizer.fit(train)
model.transform(test).show()
model.getSplitsArray.foreach {
arr => println(arr.mkString(","))
}
val time2 = System.currentTimeMillis()
println(time2 - time1)
}
}
private[feature] trait DiscretizerBase extends Params
with HasHandleInvalid with HasInputCol with HasOutputCol
with HasInputCols with HasOutputCols with HasLabelCol {
final val impurity: Param[String] = new Param[String](this, "impurity", "Criterion used for information gain calculation (case-insensitive). " +
" Supported: \"entropy\" and \"gini\". (default = gini)",
ParamValidators.inArray(Array("gini", "entropy")))
final val minInfoGain: DoubleParam = new DoubleParam(this, "minInfoGain", "分组的最小信息增益(不包含),需非负数,默认0.0", ParamValidators.gtEq(0.0))
final val numBuckets: IntParam = new IntParam(this, "numBuckets", "离散分桶数量,正整数", ParamValidators.gtEq(2))
val numBucketsArray = new IntArrayParam(this, "numBucketsArray", "Array of number of buckets " +
"(quantiles, or categories) into which data points are grouped. This is for multiple " +
"columns input. If transforming multiple columns and numBucketsArray is not set, but " +
"numBuckets is set, then numBuckets will be applied across all columns.",
(arrayOfNumBuckets: Array[Int]) => arrayOfNumBuckets.forall(ParamValidators.gtEq(2)))
final val minInstancesPerBucket: IntParam = new IntParam(this, "minInstancesPerBucket", "每个桶最少记录数量(包含),默认1")
def getImpurity() = $(minInfoGain)
def getMinInfoGain() = $(minInfoGain)
def getNumBuckets() = $(numBuckets)
def getNumBucketsArray: Array[Int] = $(numBucketsArray)
def getMinInstancesPerBucket() = $(minInstancesPerBucket)
override val handleInvalid: Param[String] = new Param[String](
this,
"handleInvalid", "how to handle invalid entries. Options are skip (which will filter out rows with bad values), " +
"or error (which will throw an error) or keep (keep invalid values in a special additional bucket).",
ParamValidators.inArray(Array("skip", "error", "keep")))
def setImpurity(value: String): this.type = set(impurity, value)
def setHandleInvalid(value: String): this.type = set(handleInvalid, value)
def setLabelCol(value: String): this.type = set(labelCol, value)
def setInputCol(value: String): this.type = set(inputCol, value)
def setOutputCol(value: String): this.type = set(outputCol, value)
def setInputCols(value: Array[String]): this.type = set(inputCols, value)
def setOutputCols(value: Array[String]): this.type = set(outputCols, value)
def setMinInfoGain(value: Double): this.type = set(minInfoGain, value)
def setNumBuckets(value: Int): this.type = set(numBuckets, value)
def setNumBucketsArray(value: Array[Int]): this.type = set(numBucketsArray, value)
def setMinInstancesPerBucket(value: Int): this.type = set(minInstancesPerBucket, value)
setDefault(minInfoGain -> 0.0, labelCol -> "label", minInstancesPerBucket -> 1, handleInvalid -> "error", impurity -> "gini")
protected def getInOutCols: (Array[String], Array[String]) = {
require(
(isSet(inputCol) && isSet(outputCol) && !isSet(inputCols) && !isSet(outputCols)) ||
(!isSet(inputCol) && !isSet(outputCol) && isSet(inputCols) && isSet(outputCols)),
"Discretizer only supports setting either inputCol/outputCol or" +
"inputCols/outputCols.")
if (isSet(inputCol)) {
(Array($(inputCol)), Array($(outputCol)))
} else {
require(
$(inputCols).length == $(outputCols).length,
"inputCols number do not match outputCols")
($(inputCols), $(outputCols))
}
}
}
class Discretizer(override val uid: String) extends Estimator[Bucketizer]
with DiscretizerBase with DefaultParamsWritable {
def this() = this(Identifiable.randomUID("Discretizer"))
override def copy(extra: ParamMap): this.type = defaultCopy(extra)
override def fit(dataset: Dataset[_]): Bucketizer = {
transformSchema(dataset.schema, true)
val bucketizer = new Bucketizer(uid).setHandleInvalid($(handleInvalid))
val (inputColNames, outputColNames) = getInOutCols
val numBucketsArray_t = if (isSet(numBucketsArray) && isSet(inputCols)) {
$(numBucketsArray)
} else {
Array.fill[Int](inputColNames.size)($(numBuckets))
}
val splitsArrayBuffer = new ArrayBuffer[Array[Double]]()
inputColNames.zip(numBucketsArray_t).foreach {
case (inputColName, numBuckets_t) =>
val splits = deiscretizeCol(dataset, inputColName, numBuckets_t)
splitsArrayBuffer += splits.sorted
}
if (splitsArrayBuffer.size == 1) {
val splits = splitsArrayBuffer.head
bucketizer.setSplits(splits)
} else {
var splitsArray = splitsArrayBuffer.toArray
splitsArray.foreach(f => f.foreach(println))
bucketizer.setSplitsArray(splitsArray)
}
copyValues(bucketizer.setParent(this))
}
override def transformSchema(schema: StructType): StructType = {
val (inputColNames, outputColNames) = getInOutCols
val existingFields = schema.fields
var outputFields = existingFields
inputColNames.zip(outputColNames).foreach {
case (inputColName, outputColName) =>
require(
existingFields.exists(_.name == inputColName),
s"Iutput column ${inputColName} not exists.")
require(
existingFields.forall(_.name != outputColName),
s"Output column ${outputColName} already exists.")
val inputColType = schema(inputColName).dataType
require(
inputColType.isInstanceOf[NumericType],
s"The input column $inputColName must be numeric type, " +
s"but got $inputColType.")
val attr = NominalAttribute.defaultAttr.withName(outputColName)
outputFields :+= attr.toStructField()
}
StructType(outputFields)
}
def deiscretizeCol(dataset: Dataset[_], inputColName: String, numBuckets_t: Int) = {
val input_arr = dataset.select(col(inputColName).cast(DoubleType)).distinct().orderBy(inputColName).rdd.map(_.getDouble(0)).collect()
val splits = new ArrayBuffer[Double]()
splits.append(Double.MinValue)
splits.append(Double.MaxValue)
var split_map_arr = new ArrayBuffer[scala.collection.mutable.Map[String, Any]]()
split_map_arr.append(scala.collection.mutable.Map(
"arr" -> input_arr,
"closure" -> true,
"node" -> null))
var flag = true
while (flag) {
for (split_map <- split_map_arr) {
if (split_map("node") == null) {
getBestPoint(split_map, dataset, inputColName)
}
}
split_map_arr = split_map_arr.filter {
split_map => split_map("node").asInstanceOf[Map[String, Double]]("value") > $(minInfoGain)
}
if (split_map_arr.length > 0) {
val entropy_idxs = (Map[Double, Array[Int]]() /: split_map_arr.zipWithIndex) { (r, split_map_idx) =>
val (split_map, idx) = split_map_idx
val value = split_map("node").asInstanceOf[Map[String, Double]]("value")
r + (value -> (r.get(value) match {
case Some(arr: Array[Int]) => arr :+ idx
case None => Array[Int](idx)
}))
}
val split_map_arr_break = new ArrayBuffer[scala.collection.mutable.Map[String, Any]]()
entropy_idxs(entropy_idxs.keys.max).zipWithIndex.foreach {
case (idx, i) => {
split_map_arr_break.append(split_map_arr.remove(idx - i))
}
}
split_map_arr_break.foreach(
split_map => {
val point = split_map("node").asInstanceOf[Map[String, Double]]("point")
splits.append(point)
val arr = split_map("arr").asInstanceOf[Array[Double]]
val closure = split_map("closure").asInstanceOf[Boolean]
var left_arr = Array[Double]()
var right_arr = Array[Double]()
for (e <- arr) {
if (e < point) {
left_arr :+= e
} else {
right_arr :+= e
}
}
left_arr :+= point
val left_split_map = scala.collection.mutable.Map(
"arr" -> left_arr,
"closure" -> false,
"node" -> null)
val right_split_map = scala.collection.mutable.Map(
"arr" -> right_arr,
"closure" -> closure,
"node" -> null)
split_map_arr.append(left_split_map)
split_map_arr.append(right_split_map)
})
if (splits.length - 1 >= numBuckets_t) {
flag = false
}
} else {
flag = false
}
}
splits.toArray
}
/*
* 获取最大信息增益分割点
*/
def getBestPoint(split_map: scala.collection.mutable.Map[String, Any], dataset: Dataset[_], inputColName: String) {
val arr: Array[Double] = split_map("arr").asInstanceOf[Array[Double]]
if (arr.length <= 1) {
split_map("node") = Map("point" -> arr(0), "value" -> 0.0)
return
}
val start = arr(0)
val end = arr(arr.length - 1)
val closure = split_map("closure").asInstanceOf[Boolean]
var ds = dataset.filter(col(inputColName) >= start and col(inputColName) < end)
if (closure) {
ds = dataset.filter(col(inputColName) >= start and col(inputColName) <= end)
}
val point_set = ds.select(col(inputColName).cast(DoubleType), col($(labelCol))).
groupBy(col(inputColName)).
pivot($(labelCol)).count.
orderBy(col(inputColName)).collect()
val classNums = point_set(0).size - 1
val all_seq = (new Array[Long](classNums) /: point_set) { (arr, p) =>
val idx = (1 to (p.size - 1)).foreach {
i =>
arr(i - 1) += p(i).asInstanceOf[Long]
}
arr
}
val info_entropy = entropy(all_seq)
val all_count = all_seq.sum
var left_count = 0L
var left_seq = new Array[Long](classNums)
val gains = arr.zipWithIndex.map {
case (point, idx) =>
if (idx == 0) {
(point, 0.0)
} else {
left_seq.zipWithIndex.foreach {
case (e, i) =>
left_seq(i) = e + point_set(idx - 1)(i + 1).asInstanceOf[Long]
}
val right_seq = all_seq.zip(left_seq).map {
case (all_e, left_e) =>
all_e - left_e
}
val left_count = left_seq.sum
val right_count = right_seq.sum
if (left_count < $(minInstancesPerBucket) || right_count < $(minInstancesPerBucket)) {
(point, -1.0)
} else {
val conditional_entropy = left_count * 1.0 / all_count * entropy(left_seq) + right_count * 1.0 / all_count * entropy(right_seq)
val gain = info_entropy - conditional_entropy
(point, gain)
}
}
}
val gain_max = gains.map(_._2).max
val point = gains.filter(_._2 >= gain_max).map(_._1).apply(0)
val node = Map("point" -> point, "value" -> gain_max)
split_map("node") = node
}
/*
* 计算信息熵, 单位nat
*/
def entropy(groupCounts: Seq[Long]) = {
val count = groupCounts.sum
if ($(impurity) == "gini") {
(0.0 /: groupCounts) { (sum, groupCount) =>
if (groupCount == 0) {
sum
} else {
val p = groupCount * 1.0 / count
sum + (p * (1 - p))
}
}
} else if ($(impurity) == "entropy") {
(0.0 /: groupCounts) { (sum, groupCount) =>
if (groupCount == 0) {
sum
} else {
val p = groupCount * 1.0 / count
sum + (-p * math.log(p))
}
}
} else { //增加算法
Double.MaxValue
}
}
}