add initial model to kmeans

Author: yinxusen

mllib/src/main/scala/org/apache/spark/ml/clustering/KMeans.scala

@@ -34,7 +34,7 @@ import org.apache.spark.sql.types.{IntegerType, StructType}
  * Common params for KMeans and KMeansModel
  */
 private[clustering] trait KMeansParams extends Params with HasMaxIter with HasFeaturesCol
-  with HasSeed with HasPredictionCol with HasTol {
+  with HasSeed with HasPredictionCol with HasTol with HasInitialModel[KMeansModel] {
 
   /**
    * Set the number of clusters to create (k). Must be > 1. Default: 2.
@@ -96,7 +96,7 @@ private[clustering] trait KMeansParams extends Params with HasMaxIter with HasFe
 @Experimental
 class KMeansModel private[ml] (
     @Since("1.5.0") override val uid: String,
-    private val parentModel: MLlibKMeansModel)
+    private[ml] val parentModel: MLlibKMeansModel)
   extends Model[KMeansModel] with KMeansParams with MLWritable {
 
   @Since("1.5.0")
@@ -237,6 +237,10 @@ class KMeans @Since("1.5.0") (
   @Since("1.5.0")
   def setSeed(value: Long): this.type = set(seed, value)
 
+  /** @group setParam */
+  @Since("2.0.0")
+  def setInitialModel(value: KMeansModel): this.type = set(initialModel, value)
+
   @Since("1.5.0")
   override def fit(dataset: DataFrame): KMeansModel = {
     val rdd = dataset.select(col($(featuresCol))).map { case Row(point: Vector) => point }
@@ -248,6 +252,11 @@ class KMeans @Since("1.5.0") (
       .setMaxIterations($(maxIter))
       .setSeed($(seed))
       .setEpsilon($(tol))
+
+    if (isSet(initialModel)) {
+      algo.setInitialModel($(initialModel).parentModel)
+    }
+
     val parentModel = algo.run(rdd)
     val model = new KMeansModel(uid, parentModel)
     copyValues(model)

mllib/src/main/scala/org/apache/spark/ml/param/shared/SharedParamsCodeGen.scala

@@ -78,7 +78,24 @@ private[shared] object SharedParamsCodeGen {
       ParamDesc[String]("solver", "the solver algorithm for optimization. If this is not set or " +
         "empty, default value is 'auto'.", Some("\"auto\"")))
 
-    val code = genSharedParams(params)
+    // scalastyle:off
+    val extras: Seq[String] = Seq(
+      """
+        |private[ml] trait HasInitialModel[T <: Model[T]] extends Params {
+        |
+        |  /**
+        |   * Param for initial model of warm start.
+        |   * @group param
+        |   */
+        |  final val initialModel: Param[T] = new Param[T](this, "initial model", "initial model of warm-start")
+        |
+        |  /** @group getParam */
+        |  final def getInitialModel: T = $(initialModel)
+        |}
+        |""".stripMargin)
+    // scalastyle:on
+
+    val code = genSharedParams(params, extras)
     val file = "src/main/scala/org/apache/spark/ml/param/shared/sharedParams.scala"
     val writer = new PrintWriter(file)
     writer.write(code)
@@ -174,7 +191,7 @@ private[shared] object SharedParamsCodeGen {
   }
 
   /** Generates Scala source code for the input params with header. */
-  private def genSharedParams(params: Seq[ParamDesc[_]]): String = {
+  private def genSharedParams(params: Seq[ParamDesc[_]], extras: Seq[String] = Nil): String = {
     val header =
       """/*
         | * Licensed to the Apache Software Foundation (ASF) under one or more
@@ -195,6 +212,7 @@ private[shared] object SharedParamsCodeGen {
         |
         |package org.apache.spark.ml.param.shared
         |
+        |import org.apache.spark.ml.Model
         |import org.apache.spark.ml.param._
         |
         |// DO NOT MODIFY THIS FILE! It was generated by SharedParamsCodeGen.
@@ -205,7 +223,8 @@ private[shared] object SharedParamsCodeGen {
     val footer = "// scalastyle:on\n"
 
     val traits = params.map(genHasParamTrait).mkString
+    val extraTraits = extras.mkString
 
-    header + traits + footer
+    header + traits + extraTraits + footer
   }
 }

mllib/src/main/scala/org/apache/spark/ml/param/shared/sharedParams.scala

@@ -17,6 +17,7 @@
 
 package org.apache.spark.ml.param.shared
 
+import org.apache.spark.ml.Model
 import org.apache.spark.ml.param._
 
 // DO NOT MODIFY THIS FILE! It was generated by SharedParamsCodeGen.
@@ -389,4 +390,16 @@ private[ml] trait HasSolver extends Params {
   /** @group getParam */
   final def getSolver: String = $(solver)
 }
+
+private[ml] trait HasInitialModel[T <: Model[T]] extends Params {
+
+  /**
+   * Param for initial model of warm start.
+   * @group param
+   */
+  final val initialModel: Param[T] = new Param[T](this, "initial model", "initial model of warm-start")
+
+  /** @group getParam */
+  final def getInitialModel: T = $(initialModel)
+}
 // scalastyle:on

mllib/src/test/scala/org/apache/spark/ml/clustering/KMeansSuite.scala

@@ -19,7 +19,7 @@ package org.apache.spark.ml.clustering
 
 import org.apache.spark.SparkFunSuite
 import org.apache.spark.ml.util.DefaultReadWriteTest
-import org.apache.spark.mllib.clustering.{KMeans => MLlibKMeans}
+import org.apache.spark.mllib.clustering.{KMeans => MLlibKMeans, KMeansModel => MLlibKMeansModel}
 import org.apache.spark.mllib.linalg.{Vector, Vectors}
 import org.apache.spark.mllib.util.MLlibTestSparkContext
 import org.apache.spark.sql.{DataFrame, SQLContext}
@@ -106,6 +106,38 @@ class KMeansSuite extends SparkFunSuite with MLlibTestSparkContext with DefaultR
     val kmeans = new KMeans()
     testEstimatorAndModelReadWrite(kmeans, dataset, KMeansSuite.allParamSettings, checkModelData)
   }
+
+  test("Initialize using given cluster centers") {
+    val points = Array(
+      Vectors.dense(0.0, 0.0, 0.0),
+      Vectors.dense(1.0, 1.0, 1.0),
+      Vectors.dense(2.0, 2.0, 2.0),
+      Vectors.dense(3.0, 3.0, 3.0),
+      Vectors.dense(4.0, 4.0, 4.0)
+    )
+
+    // creating an initial model
+    val initialModel = new KMeansModel("test model", new MLlibKMeansModel(points))
+
+    val predictionColName = "kmeans_prediction"
+    val kmeans = new KMeans()
+      .setK(k)
+      .setPredictionCol(predictionColName)
+      .setSeed(1)
+      .setInitialModel(initialModel)
+    val model = kmeans.fit(dataset)
+    assert(model.clusterCenters.length === k)
+
+    val transformed = model.transform(dataset)
+    val expectedColumns = Array("features", predictionColName)
+    expectedColumns.foreach { column =>
+      assert(transformed.columns.contains(column))
+    }
+    val clusters =
+      transformed.select(predictionColName).map(_.getInt(0)).distinct().collect().toSet
+    assert(clusters.size === k)
+    assert(clusters === Set(0, 1, 2, 3, 4))
+  }
 }
 
 object KMeansSuite {