bigger training set for test and more verbose error msg

CeON · Sep 25, 2015 · 2b2595a · 2b2595a
1 parent 3f521f8
commit 2b2595a
Show file tree

Hide file tree

Showing 2 changed files with 15 additions and 0 deletions.
diff --git a/src/main/scala/pl/edu/icm/sparkling_ferns/Fern.scala b/src/main/scala/pl/edu/icm/sparkling_ferns/Fern.scala
@@ -207,6 +207,7 @@ object Fern {
     val numSamples = objectsPerLabel.sum
 
     val countOfZeros = objectsInLeafPerLabel.map(_.count(_ == 0)).sum
+    assert(countOfZeros < numLabels * numDistinctPoints, "Empty fern, please rerun the algorithm and/or supply a bigger training set")
     val countOfMin = objectsInLeafPerLabel.flatMap(_.filter(_ > 0)).groupBy(identity).minBy(_._1)._2.length
     val epsilon =
       if (countOfZeros > 0) {

diff --git a/src/test/scala/pl/edu/icm/sparkling_ferns/FernForestSuite.scala b/src/test/scala/pl/edu/icm/sparkling_ferns/FernForestSuite.scala
@@ -13,6 +13,20 @@ class FernForestSuite extends FunSuite with LocalSparkContext with FixedRandomNu
       LabeledPoint(1.0, Vectors.dense(1.0, 1.0, 1.0)),
       LabeledPoint(1.0, Vectors.dense(1.0, 1.0, 1.0)),
       LabeledPoint(1.0, Vectors.dense(1.0, 1.0, 1.0)),
+      LabeledPoint(1.0, Vectors.dense(1.0, 1.0, 1.0)),
+      LabeledPoint(1.0, Vectors.dense(1.0, 1.0, 1.0)),
+      LabeledPoint(1.0, Vectors.dense(1.0, 1.0, 1.0)),
+      LabeledPoint(1.0, Vectors.dense(1.0, 1.0, 1.0)),
+      LabeledPoint(1.0, Vectors.dense(1.0, 1.0, 1.0)),
+      LabeledPoint(1.0, Vectors.dense(1.0, 1.0, 1.0)),
+      LabeledPoint(1.0, Vectors.dense(1.0, 1.0, 1.0)),
+      LabeledPoint(-1.0, Vectors.dense(-1.0, -1.0, -1.0)),
+      LabeledPoint(-1.0, Vectors.dense(-1.0, -1.0, -1.0)),
+      LabeledPoint(-1.0, Vectors.dense(-1.0, -1.0, -1.0)),
+      LabeledPoint(-1.0, Vectors.dense(-1.0, -1.0, -1.0)),
+      LabeledPoint(-1.0, Vectors.dense(-1.0, -1.0, -1.0)),
+      LabeledPoint(-1.0, Vectors.dense(-1.0, -1.0, -1.0)),
+      LabeledPoint(-1.0, Vectors.dense(-1.0, -1.0, -1.0)),
       LabeledPoint(-1.0, Vectors.dense(-1.0, -1.0, -1.0)),
       LabeledPoint(-1.0, Vectors.dense(-1.0, -1.0, -1.0))
     )