Hack ALSSuite to support NNLS testing.

Author: tmyklebu

mllib/src/test/scala/org/apache/spark/mllib/recommendation/ALSSuite.scala

@@ -48,12 +48,18 @@ object ALSSuite {
       features: Int,
       samplingRate: Double,
       implicitPrefs: Boolean = false,
-      negativeWeights: Boolean = false): (Seq[Rating], DoubleMatrix, DoubleMatrix) = {
+      negativeWeights: Boolean = false,
+      negativeFactors: Boolean = true): (Seq[Rating], DoubleMatrix, DoubleMatrix) = {
     val rand = new Random(42)
 
     // Create a random matrix with uniform values from -1 to 1
-    def randomMatrix(m: Int, n: Int) =
-      new DoubleMatrix(m, n, Array.fill(m * n)(rand.nextDouble() * 2 - 1): _*)
+    def randomMatrix(m: Int, n: Int) = {
+      if (negativeFactors) {
+        new DoubleMatrix(m, n, Array.fill(m * n)(rand.nextDouble() * 2 - 1): _*)
+      } else {
+        new DoubleMatrix(m, n, Array.fill(m * n)(rand.nextDouble()): _*)
+      }
+    }
 
     val userMatrix = randomMatrix(users, features)
     val productMatrix = randomMatrix(features, products)
@@ -128,6 +134,27 @@ class ALSSuite extends FunSuite with LocalSparkContext {
     assert(u11 != u2)
   }
 
+  test("negative ids") {
+    val data = ALSSuite.generateRatings(50, 50, 2, 0.7, false, false)
+    val ratings = sc.parallelize(data._1.map { case Rating(u,p,r) => Rating(u-25,p-25,r) })
+    val correct = data._2
+    val model = ALS.train(ratings, 5, 15)
+
+    val pairs = Array.tabulate(50, 50)((u,p) => (u-25,p-25)).flatten
+    val ans = model.predict(sc.parallelize(pairs)).collect
+    ans.foreach { r =>
+      val u = r.user + 25
+      val p = r.product + 25
+      val v = r.rating
+      val error = v - correct.get(u,p)
+      assert(math.abs(error) < 0.4)
+    }
+  }
+
+  test("NNALS, rank 2") {
+    testALS(100, 200, 2, 15, 0.7, 0.4, false, false, false, -1, false)
+  }
+
   /**
    * Test if we can correctly factorize R = U * P where U and P are of known rank.
    *
@@ -140,16 +167,21 @@ class ALSSuite extends FunSuite with LocalSparkContext {
    * @param implicitPrefs  flag to test implicit feedback
    * @param bulkPredict    flag to test bulk prediciton
    * @param negativeWeights whether the generated data can contain negative values
+   * @param numBlocks      number of blocks to partition users and products into
+   * @param negativeFactors whether the generated user/product factors can have negative entries
    */
   def testALS(users: Int, products: Int, features: Int, iterations: Int,
     samplingRate: Double, matchThreshold: Double, implicitPrefs: Boolean = false,
-    bulkPredict: Boolean = false, negativeWeights: Boolean = false)
+    bulkPredict: Boolean = false, negativeWeights: Boolean = false, numBlocks: Int = -1,
+    negativeFactors: Boolean = true)
   {
     val (sampledRatings, trueRatings, truePrefs) = ALSSuite.generateRatings(users, products,
-      features, samplingRate, implicitPrefs, negativeWeights)
+      features, samplingRate, implicitPrefs, negativeWeights, negativeFactors)
     val model = implicitPrefs match {
-      case false => ALS.train(sc.parallelize(sampledRatings), features, iterations)
-      case true => ALS.trainImplicit(sc.parallelize(sampledRatings), features, iterations)
+      case false => ALS.train(sc.parallelize(sampledRatings), features, iterations, 0.01,
+          numBlocks, 0L, !negativeFactors)
+      case true => ALS.trainImplicit(sc.parallelize(sampledRatings), features, iterations, 0.01,
+          numBlocks, 1.0, 0L, !negativeFactors)
     }
 
     val predictedU = new DoubleMatrix(users, features)