add predict for single doc

Author: hhbyyh

mllib/src/main/scala/org/apache/spark/mllib/clustering/LDAModel.scala

@@ -366,6 +366,34 @@ class LocalLDAModel private[clustering] (
     }
   }
 
+  /**
+   * Predicts the topic mixture distribution for a document (often called "theta" in the
+   * literature).  Returns a vector of zeros for an empty document.
+   *
+   * Note this means to allow quick query for single document. For batch documents, please refer
+   * to [[topicDistributions(documents: RDD[(Long, Vector)])]] to avoid overhead.
+   *
+   * @param document document to predict topic mixture distributions for
+   * @return (document ID, topic mixture distribution for document)
+   */
+  @Since("1.6.0")
+  def topicDistributions(document: (Long, Vector)): (Long, Vector) = {
+    val expElogbeta = exp(LDAUtils.dirichletExpectation(topicsMatrix.toBreeze.toDenseMatrix.t).t)
+    val id = document._1
+    val termCounts = document._2
+    if (termCounts.numNonzeros == 0) {
+      (id, Vectors.zeros(this.k))
+    } else {
+      val (gamma, _) = OnlineLDAOptimizer.variationalTopicInference(
+        termCounts,
+        expElogbeta,
+        this.docConcentration.toBreeze,
+        gammaShape,
+        this.k)
+      (id, Vectors.dense(normalize(gamma, 1.0).toArray))
+    }
+  }
+
   /**
    * Java-friendly version of [[topicDistributions]]
    */

mllib/src/test/scala/org/apache/spark/mllib/clustering/LDASuite.scala

@@ -374,9 +374,14 @@ class LDASuite extends SparkFunSuite with MLlibTestSparkContext {
       .values
       .collect()
 
-    expectedPredictions.zip(actualPredictions).forall { case (expected, actual) =>
-      expected._1 === actual._1 && (expected._2 ~== actual._2 relTol 1E-3D)
+    expectedPredictions.zip(actualPredictions).foreach { case (expected, actual) =>
+      assert(expected._1 === actual._1 && (expected._2 ~== actual._2 relTol 1E-3D))
     }
+
+    val topicsBz = ldaModel.topicDistributions(docs.first())._2.toBreeze.toDenseVector
+    val singlePrediction = (argmax(topicsBz), max(topicsBz))
+    assert(expectedPredictions(0)._1 === singlePrediction._1
+      && (expectedPredictions(0)._2 ~== singlePrediction._2 relTol 1E-3D))
   }
 
   test("OnlineLDAOptimizer with asymmetric prior") {