Skip to content
Snippets Groups Projects
Commit c0e48dfa authored by Yu ISHIKAWA's avatar Yu ISHIKAWA Committed by Davies Liu
Browse files

[SPARK-11566] [MLLIB] [PYTHON] Refactoring GaussianMixtureModel.gaussians in Python

cc jkbradley

Author: Yu ISHIKAWA <yuu.ishikawa@gmail.com>

Closes #9534 from yu-iskw/SPARK-11566.
parent a3989058
No related branches found
No related tags found
No related merge requests found
......@@ -17,14 +17,11 @@
package org.apache.spark.mllib.api.python
import java.util.{List => JList}
import scala.collection.JavaConverters._
import scala.collection.mutable.ArrayBuffer
import scala.collection.JavaConverters
import org.apache.spark.SparkContext
import org.apache.spark.mllib.linalg.{Vector, Vectors, Matrix}
import org.apache.spark.mllib.clustering.GaussianMixtureModel
import org.apache.spark.mllib.linalg.{Vector, Vectors}
/**
* Wrapper around GaussianMixtureModel to provide helper methods in Python
......@@ -36,17 +33,11 @@ private[python] class GaussianMixtureModelWrapper(model: GaussianMixtureModel) {
/**
* Returns gaussians as a List of Vectors and Matrices corresponding each MultivariateGaussian
*/
val gaussians: JList[Object] = {
val modelGaussians = model.gaussians
var i = 0
var mu = ArrayBuffer.empty[Vector]
var sigma = ArrayBuffer.empty[Matrix]
while (i < k) {
mu += modelGaussians(i).mu
sigma += modelGaussians(i).sigma
i += 1
val gaussians: Array[Byte] = {
val modelGaussians = model.gaussians.map { gaussian =>
Array[Any](gaussian.mu, gaussian.sigma)
}
List(mu.toArray, sigma.toArray).map(_.asInstanceOf[Object]).asJava
SerDe.dumps(JavaConverters.seqAsJavaListConverter(modelGaussians).asJava)
}
def save(sc: SparkContext, path: String): Unit = model.save(sc, path)
......
......@@ -266,7 +266,7 @@ class GaussianMixtureModel(JavaModelWrapper, JavaSaveable, JavaLoader):
"""
return [
MultivariateGaussian(gaussian[0], gaussian[1])
for gaussian in zip(*self.call("gaussians"))]
for gaussian in self.call("gaussians")]
@property
@since('1.4.0')
......
0% Loading or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment