Repository: aehrc/VariantSpark (language: python)

Issue:
Predict Labels with VariantSpark
VariantSpark offers RandomForest training with the "importance" command. I was wondering if you could implement prediction command too. I expect the prediction command to
accept feature values for the new set of samples similarly that the importance command accepts it ("-if" input file, "-it" input type, ....)
accept the model previously trained by "importance"  command through "-im" (input model parameter)
output a two-column file with sample id and predicted label. But that would be great to see more columns to show how many trees vote for each label.
An extension would be to implement probability prediction that for each sample reports the probability of each class. A further extension is to report it per tree (what is the probability of each class for each sample based on each tree)

I would like to see this in CLI first and then Hail API. If possible it would be nice to appear in Python API and then other API too.

Given that not all features from the input file are used in the RandomForest model, I strongly recommend not to load the entire input file in the memory but only features which are used in the model. Note that the input file could be really big.

Also if there was a feature in the model that is not reported in the input file the program should exit with a proper message and list all the features missing in the input file.



Discussion / hints:
I've been tossing up between serialising the file to disk as a json object, or possibly even converting to ONNX format. For the sake of time I'd recommend saving as json and adding ONNX later on.

To compare with this example:

```
./variant-spark --local -- importance
  -if data/chr22_1000.vcf  
  -ff data/chr22-labels.csv  
  -fc 22_16051249  
 -om modelFile.model  //new parameter to add to this command  
  -v -rn 500 -rbs 20 -ro
```

how about something like:

```
./variant-spark --local -- predict
  -if modelFile.model
  -im 22_16051249
  -of 22_16051249_predictions.file
```

Is this close to what you're after?
I think it should be something like: 

```
./variant-spark --local -- predict
  -if  to_predict.vcf  # or other supported input format/specification
  -im modelFile.model  # the serialized model from training 
  -of to_predict_predictions.csv
  # other options e.g.  if should output class probabilities
```

Where the output csv has the following format:
```sample, predicted_class, p_0, p_1, p_2, .... ```
where:
-```sample``` : sample id
-```prediced_class```:  the predicted class (max probability class)
_```p_0, p_1, p2```:  probability of each class
 
@ArashBayatDev what do you think? 
@BMJHayward wrote: 
> Looking at ModelConverter.scala and ModelOutputArgs.scala, I'd like to be able to load a model previously serialised by the saveModel function. I.e. I want to be able to perform the inverse of the functionality here.
> @piotrszul what is the recommended approach here? Or even better, do you know of some code to do this already?
> AnalyzeRFCmd.run() might be good too

Yes, the code for serialization and (partialy) deserialization is in the classes you mentioned.
But ideally it should be refactored to a lower level API trait and classed that can handle serialisation and deserialisation in different formats.



I think we need to splits this into a few issues (and do them in this order):

- Implementing command line  prediction of class  form java serialised model 
- Add  command line option for predicing class probabilities
- Implementing command line  predictions form JSON serialised model
- Adding option for checking if the features in the prediction data match the model
- Implementing the same functionality in Hail API
- Implementing the same functionality in Python API
- Add support for ONYX 

@ArashBayatDev ???
Yes I think so as well. To clarify something in point 4:
> Adding option for checking if the features in the prediction data match the model    

Do you mean validating features of input data before running prediction?
> Yes I think so as well. To clarify something in point 4:
> 
> > Adding option for checking if the features in the prediction data match the model
> > Do you mean validating features of input data before running prediction?

Yes, at minimum or more more general, the ability to deal with input data that have to some degree  different format, shape etc, than the training data.  I think this needs a bit of further discussion a for how much robustness needs to be built, the weakest option being none - garbage in garbage out. 
In the first stage I think we can make the requirement that the data for prediction exactly match the format, shape etc. of the data used for training. And then optionally add functionality to relax parts of this assumption.

Ok I've a got a table here to represent the target output from @piotrszul above. I realised there could be a number of different ways to do this, and it could get complicated with things like multi-indices etc. Are my changes to the headers ok? @ArashBayatDev 

[predictionOutput.txt](https://github.com/aehrc/VariantSpark/files/4615172/predictionOutput.txt)

Notice the repetition of the sample ID, once for each sample location in the source vcf. Under this schema there would be (number of samples x number of locations) number of rows in the output.

Is this the _right_ implementation of what you mentioned above Peter?

Also note that I've got homozygous and heterozygous column names in place of 0|0, 0|1 and 1|1 for the alleles in each sample.

Fix patch:
diff --git a/data/ch22-model.ser b/data/ch22-model.ser
index 562bfccf..7c6f467d 100644
Binary files a/data/ch22-model.ser and b/data/ch22-model.ser differ
diff --git a/examples/local_run-predict-ch22.sh b/examples/local_run-predict-ch22.sh
new file mode 100755
index 00000000..0ce95a4d
--- /dev/null
+++ b/examples/local_run-predict-ch22.sh
@@ -0,0 +1,11 @@
+#!/bin/bash
+
+PWD=$(
+  cd $(dirname "$0")/..
+  pwd
+)
+PATH=${PATH}:${PWD}/bin
+export VS_ECHO_CMDLINE=YES
+
+variant-spark --spark --master 'local[*]' -- \
+  predict -if ${PWD}/data/chr22_1000.vcf -im ${PWD}/data/ch22-model.ser "$@"
diff --git a/src/main/scala/au/csiro/variantspark/algo/RandomForest.scala b/src/main/scala/au/csiro/variantspark/algo/RandomForest.scala
index 6714f5de..ebcc6cd4 100644
--- a/src/main/scala/au/csiro/variantspark/algo/RandomForest.scala
+++ b/src/main/scala/au/csiro/variantspark/algo/RandomForest.scala
@@ -3,7 +3,7 @@ package au.csiro.variantspark.algo
 import au.csiro.pbdava.ssparkle.common.utils.FastUtilConversions._
 import au.csiro.pbdava.ssparkle.common.utils.Logging
 import au.csiro.pbdava.ssparkle.common.utils.Timed._
-import au.csiro.variantspark.data.VariableType
+import au.csiro.variantspark.data.Feature
 import au.csiro.variantspark.metrics.Metrics
 import au.csiro.variantspark.utils.IndexedRDDFunction._
 import au.csiro.variantspark.utils.{Sample, defRng}
@@ -12,10 +12,6 @@ import it.unimi.dsi.util.XorShift1024StarRandomGenerator
 import org.apache.commons.lang3.builder.ToStringBuilder
 import org.apache.spark.rdd.RDD
 
-import scala.reflect.ClassTag
-import au.csiro.variantspark.data.Feature
-import au.csiro.variantspark.data.FeatureBuilder
-
 /** Allows for normalization(scaling)of the input map values
   */
 trait VarImportanceNormalizer {
@@ -70,7 +66,20 @@ case class VotingAggregator(nLabels: Int, nSamples: Int) {
   /** Maps votes to predictions
     *
     */
-  def predictions: Array[Int] = votes.map(_.zipWithIndex.maxBy(_._1)._2)
+  def predictions: Array[Int] = votes.map(v => v.indices.maxBy(v))
+
+  /**
+    * Computes class probabilities.
+    * The result is an array with one item per sample, where
+    * each item is a vector with class probabilities for this sample.
+    * @return predicted class probabilities for each sample.
+    */
+  def classProbabilities: Array[Array[Double]] = {
+    votes.map { row =>
+      val sampleTotal = row.sum.toDouble
+      row.map(classCount => classCount / sampleTotal)
+    }
+  }
 }
 
 /** Implements random forest members conditionally
@@ -90,9 +99,6 @@ case class RandomForestMember(predictor: PredictiveModelWithImportance,
 case class RandomForestModel(members: List[RandomForestMember], labelCount: Int,
     oobErrors: List[Double] = List.empty, params: RandomForestParams = null) {
 
-  def size: Int = members.size
-  def trees: List[PredictiveModelWithImportance] = members.map(_.predictor)
-
   def oobError: Double = oobErrors.last
 
   def printout() {
@@ -103,6 +109,8 @@ case class RandomForestModel(members: List[RandomForestMember], labelCount: Int,
     }
   }
 
+  def trees: List[PredictiveModelWithImportance] = members.map(_.predictor)
+
   def normalizedVariableImportance(
       norm: VarImportanceNormalizer = To100ImportanceNormalizer): Map[Long, Double] =
     norm.normalize(variableImportance)
@@ -119,6 +127,8 @@ case class RandomForestModel(members: List[RandomForestMember], labelCount: Int,
       .mapValues(_ / size)
   }
 
+  def size: Int = members.size
+
   def predict(indexedData: RDD[(Feature, Long)]): Array[Int] =
     predict(indexedData, indexedData.size)
 
@@ -128,6 +138,16 @@ case class RandomForestModel(members: List[RandomForestMember], labelCount: Int,
       .foldLeft(VotingAggregator(labelCount, nSamples))(_.addVote(_))
       .predictions
   }
+
+  def predictProb(indexedData: RDD[(Feature, Long)]): Array[Array[Double]] =
+    predictProb(indexedData, indexedData.size)
+
+  def predictProb(indexedData: RDD[(Feature, Long)], nSamples: Int): Array[Array[Double]] = {
+    val treeVotes = trees
+      .map(_.predict(indexedData))
+      .foldLeft(VotingAggregator(labelCount, nSamples))(_.addVote(_))
+    treeVotes.classProbabilities
+  }
 }
 
 /** Implements random forest params conditionally
@@ -186,6 +206,7 @@ trait BatchTreeModel {
 
 object RandomForest {
   type ModelBuilderFactory = DecisionTreeParams => BatchTreeModel
+  val defaultBatchSize: Int = 10
 
   def wideDecisionTreeBuilder(params: DecisionTreeParams): BatchTreeModel = {
     val decisionTree = new DecisionTree(params)
@@ -199,8 +220,6 @@ object RandomForest {
           models.asInstanceOf[Seq[DecisionTreeModel]], indexes)
     }
   }
-
-  val defaultBatchSize: Int = 10
 }
 
 /** Implements random forest
diff --git a/src/main/scala/au/csiro/variantspark/cli/ImportanceCmd.scala b/src/main/scala/au/csiro/variantspark/cli/ImportanceCmd.scala
index 2f858457..bba16714 100644
--- a/src/main/scala/au/csiro/variantspark/cli/ImportanceCmd.scala
+++ b/src/main/scala/au/csiro/variantspark/cli/ImportanceCmd.scala
@@ -1,53 +1,24 @@
 package au.csiro.variantspark.cli
 
+import au.csiro.pbdava.ssparkle.common.arg4j.{AppRunner, TestArgs}
+import au.csiro.pbdava.ssparkle.common.utils.{CSVUtils, Logging, ReusablePrintStream, Timer}
+import au.csiro.pbdava.ssparkle.spark.{SparkApp, SparkUtils}
 import au.csiro.sparkle.common.args4j.ArgsApp
-import au.csiro.sparkle.cmd.CmdApp
-import org.kohsuke.args4j.Option
-import au.csiro.pbdava.ssparkle.common.arg4j.AppRunner
-import au.csiro.pbdava.ssparkle.spark.SparkApp
-import collection.JavaConverters._
-import au.csiro.variantspark.input.VCFSource
-import au.csiro.variantspark.input.VCFFeatureSource
-import au.csiro.variantspark.input.HashingLabelSource
-import org.apache.spark.mllib.linalg.{Vector, Vectors}
-import au.csiro.variantspark.input.CsvLabelSource
+import au.csiro.variantspark.algo.{RandomForest, RandomForestCallback, RandomForestParams, _}
+import au.csiro.variantspark.cli.args.{
+  FeatureSourceArgs,
+  ImportanceArgs,
+  ModelOutputArgs,
+  RandomForestArgs
+}
+import au.csiro.variantspark.cmd.EchoUtils._
 import au.csiro.variantspark.cmd.Echoable
-import au.csiro.pbdava.ssparkle.common.utils.Logging
+import au.csiro.variantspark.input.CsvLabelSource
+import au.csiro.variantspark.utils.{HdfsPath, defRng}
 import org.apache.commons.lang3.builder.ToStringBuilder
-import au.csiro.variantspark.cmd.EchoUtils._
-import au.csiro.pbdava.ssparkle.common.utils.LoanUtils
-import com.github.tototoshi.csv.CSVWriter
-import au.csiro.pbdava.ssparkle.common.arg4j.TestArgs
-import org.apache.hadoop.fs.FileSystem
-import org.apache.commons.math3.stat.correlation.PearsonsCorrelation
-import au.csiro.pbdava.ssparkle.spark.SparkUtils
-import au.csiro.pbdava.ssparkle.common.utils.ReusablePrintStream
-import au.csiro.variantspark.algo.RandomForestCallback
-import au.csiro.variantspark.utils.VectorRDDFunction._
-import au.csiro.variantspark.input.CsvFeatureSource
-import au.csiro.variantspark.algo.RandomForestParams
-import au.csiro.variantspark.data.BoundedOrdinalVariable
-import au.csiro.pbdava.ssparkle.common.utils.Timer
-import au.csiro.variantspark.utils.defRng
-import au.csiro.variantspark.input.ParquetFeatureSource
-import au.csiro.variantspark.utils.IndexedRDDFunction._
-import java.io.ObjectOutputStream
-import java.io.FileOutputStream
 import org.apache.hadoop.conf.Configuration
-import au.csiro.variantspark.utils.HdfsPath
-import au.csiro.pbdava.ssparkle.common.utils.CSVUtils
-import au.csiro.variantspark.cli.args.ImportanceArgs
-import au.csiro.variantspark.cli.args.RandomForestArgs
-import au.csiro.variantspark.cli.args.FeatureSourceArgs
-import au.csiro.variantspark.data.ContinuousVariable
-import au.csiro.variantspark.algo.RandomForest
-import au.csiro.variantspark.data.FeatureBuilder
-import scala.reflect.ClassTag
-import au.csiro.variantspark.input._
-import au.csiro.variantspark.algo._
-import au.csiro.variantspark.data.VariableType
-import org.apache.spark.rdd.RDD
-import au.csiro.variantspark.cli.args.ModelOutputArgs
+import org.apache.hadoop.fs.FileSystem
+import org.kohsuke.args4j.Option
 
 class ImportanceCmd
     extends ArgsApp with SparkApp with FeatureSourceArgs with ImportanceArgs with RandomForestArgs
@@ -82,9 +53,10 @@ class ImportanceCmd
   val randomSeed: Long = defRng.nextLong
 
   override def testArgs: Array[String] =
-    Array("-if", "data/chr22_1000.vcf", "-ff", "data/chr22-labels.csv", "-fc", "22_16051249",
-      "-ovn", "raw", "-on", "1988", "-rn", "1000", "-rbs", "100", "-ic", "-om",
-      "target/ch22-model.json", "-omf", "json", "-sr", "13", "-v", "-io", """{"separator":":"}""")
+    Array("-if", "data/chr22_1000.vcf", "-ff", "data/chr22-labels.csv", "-fc", "22_16051347",
+      "-ovn", "raw", "-on", "1988", "-rn", "1000", "-rbs", "250", "-ic", "-om",
+      "target/ch22-model.ser", "-omf", "java", "-sr", "13", "-v", "-io", """{"separator":":"}""",
+      "-ro", "-rmtf", "0.1")
 
   override def run(): Unit = {
     implicit val fs: FileSystem = FileSystem.get(sc.hadoopConfiguration)
@@ -202,6 +174,7 @@ class ImportanceCmd
                                             } else { Nil })
         }))
     }
+
     saveModel(result, index.toMap)
   }
 }
diff --git a/src/main/scala/au/csiro/variantspark/cli/PredictCmd.scala b/src/main/scala/au/csiro/variantspark/cli/PredictCmd.scala
new file mode 100644
index 00000000..a8424015
--- /dev/null
+++ b/src/main/scala/au/csiro/variantspark/cli/PredictCmd.scala
@@ -0,0 +1,93 @@
+package au.csiro.variantspark.cli
+
+import java.io.FileInputStream
+
+import au.csiro.pbdava.ssparkle.common.arg4j.{AppRunner, TestArgs}
+import au.csiro.pbdava.ssparkle.common.utils._
+import au.csiro.pbdava.ssparkle.spark.SparkApp
+import au.csiro.sparkle.common.args4j.ArgsApp
+import au.csiro.variantspark.algo.RandomForestModel
+import au.csiro.variantspark.cli.args.FeatureSourceArgs
+import au.csiro.variantspark.cmd.EchoUtils._
+import au.csiro.variantspark.cmd.Echoable
+import au.csiro.variantspark.utils.HdfsPath
+import org.apache.commons.lang3.builder.ToStringBuilder
+import org.apache.hadoop.conf.Configuration
+import org.apache.hadoop.fs.FileSystem
+import org.apache.spark.serializer.JavaSerializer
+import org.kohsuke.args4j.{Option => ArgsOption}
+
+class PredictCmd
+    extends ArgsApp with FeatureSourceArgs with Echoable with SparkApp with Logging
+    with TestArgs {
+
+  /* e.g run with
+   * ./bin/variant-spark predict -if data/chr22_1000.vcf -ff data/chr22-labels.csv
+   * -fc 22_16051249 -om target/ch22-model.json -omf json -sr 13 -ii data/hipster_labels.txt
+   *  -im data/ch22-model.java
+   */
+  @ArgsOption(name = "-im", required = true, usage = "Path to input model",
+    aliases = Array("--input-model"))
+  val inputModel: String = null
+
+  @ArgsOption(name = "-of", required = false, usage = "Path to output file (def = stdout)",
+    aliases = Array("--output-file"))
+  val outputFile: String = null
+
+  override def testArgs: Array[String] =
+    Array("-if", "data/chr22_1000.vcf", "-im", "data/ch22-model.ser", "-v")
+
+  override def run(): Unit = {
+    implicit val fs: FileSystem = FileSystem.get(sc.hadoopConfiguration)
+
+    implicit val hadoopConf: Configuration = sc.hadoopConfiguration
+    logDebug(s"Running with filesystem: ${fs}, home: ${fs.getHomeDirectory}")
+    logInfo("Running with params: " + ToStringBuilder.reflectionToString(this))
+
+    echo(s"Running random forest prediction")
+
+    val javaSerializer = new JavaSerializer(conf)
+    val si = javaSerializer.newInstance()
+
+    val rfModel =
+      LoanUtils.withCloseable(new FileInputStream(inputModel)) { in =>
+        si.deserializeStream(in).readObject().asInstanceOf[RandomForestModel]
+      }
+
+    echo(s"Loaded model of size: ${rfModel.size}")
+
+    val dataLoadingTimer = Timer()
+    echo(s"Loaded rows: ${dumpList(featureSource.sampleNames)}")
+
+    val inputData = featureSource.features.zipWithIndex().cache()
+    val totalVariables = inputData.count()
+
+    val variablePreview = inputData.map(_._1.label).take(defaultPreviewSize).toList
+    echo(s"Loaded variables: ${dumpListHead(variablePreview, totalVariables)},"
+        + s" took: ${dataLoadingTimer.durationInSec}")
+    echoDataPreview()
+
+    val classProbabilities: Array[Array[Double]] =
+      rfModel.predictProb(inputData)
+    val predictionRows = (featureSource.sampleNames zip classProbabilities)
+      .map {
+        case (sampleName, classProb) =>
+          sampleName :: classProb.indices.maxBy(classProb) :: Nil ::: classProb.toList
+      }
+
+    CSVUtils.withStream(if (outputFile != null) {
+      HdfsPath(outputFile).create()
+    } else ReusablePrintStream.stdout) { writer =>
+      val header =
+        List("sample", "class") ::: Range(0, rfModel.labelCount).map(ci => s"p_${ci}").toList
+      writer.writeRow(header)
+      writer.writeAll(predictionRows)
+    }
+  }
+}
+
+object PredictCmd {
+  def main(args: Array[String]) {
+    AppRunner.mains[PredictCmd](args)
+  }
+}
diff --git a/src/main/scala/au/csiro/variantspark/cli/VariantSparkApp.scala b/src/main/scala/au/csiro/variantspark/cli/VariantSparkApp.scala
index 82ee6c19..517d34d5 100644
--- a/src/main/scala/au/csiro/variantspark/cli/VariantSparkApp.scala
+++ b/src/main/scala/au/csiro/variantspark/cli/VariantSparkApp.scala
@@ -16,7 +16,8 @@ class VariantSparkApp extends MultiCmdApp {
   registerClass("analyze-rf", classOf[AnalyzeRFCmd])
   registerClass("build-index", classOf[BuildVarIndexCmd])
   registerClass("pdist", classOf[PairWiseDistanceCmd])
-  registerClass("version", classOf[VersionCmd])
+  registerClass("predict", classOf[PredictCmd])
+  registerClass("vcf2labels", classOf[VcfToLabels])
 }
 
 object VariantSparkApp {
diff --git a/src/main/scala/au/csiro/variantspark/cli/VcfToLabels.scala b/src/main/scala/au/csiro/variantspark/cli/VcfToLabels.scala
index 230d699d..20266407 100644
--- a/src/main/scala/au/csiro/variantspark/cli/VcfToLabels.scala
+++ b/src/main/scala/au/csiro/variantspark/cli/VcfToLabels.scala
@@ -1,28 +1,22 @@
 package au.csiro.variantspark.cli
 
-import au.csiro.sparkle.common.args4j.ArgsApp
-import au.csiro.sparkle.cmd.CmdApp
-import org.kohsuke.args4j.Option
-import au.csiro.pbdava.ssparkle.common.arg4j.AppRunner
-import au.csiro.pbdava.ssparkle.spark.SparkApp
-import collection.JavaConverters._
-import au.csiro.variantspark.input.VCFSource
-import au.csiro.variantspark.input.VCFFeatureSource
-import au.csiro.variantspark.input.HashingLabelSource
-import org.apache.spark.mllib.linalg.Vectors
 import java.io.File
-import au.csiro.pbdava.ssparkle.common.utils.LoanUtils
+
+import au.csiro.pbdava.ssparkle.common.arg4j.AppRunner
 import au.csiro.pbdava.ssparkle.common.utils.CSVUtils
-import au.csiro.variantspark.input._
+import au.csiro.pbdava.ssparkle.spark.SparkApp
+import au.csiro.sparkle.common.args4j.ArgsApp
+import au.csiro.variantspark.input.{VCFFeatureSource, VCFSource}
+import org.kohsuke.args4j.Option
 
 class VcfToLabels extends ArgsApp with SparkApp {
 
   @Option(name = "-if", required = false, usage = "This is input files",
     aliases = Array("--input-file"))
-  val inputFile: String = "data/small.vcf"
+  val inputFile: String = null
 
   @Option(name = "-of", required = false, usage = "Output file", aliases = Array("--output-file"))
-  val outputFile: String = "data/small-labels.csv"
+  val outputFile: String = null
 
   @Option(name = "-l", required = false)
   val limit: Int = 10


Test patch:
diff --git a/src/test/scala/au/csiro/variantspark/algo/WideRandomForrestModelTest.scala b/src/test/scala/au/csiro/variantspark/algo/WideRandomForrestModelTest.scala
index c2673ad2..ebb2ce7f 100644
--- a/src/test/scala/au/csiro/variantspark/algo/WideRandomForrestModelTest.scala
+++ b/src/test/scala/au/csiro/variantspark/algo/WideRandomForrestModelTest.scala
@@ -1,17 +1,18 @@
 package au.csiro.variantspark.algo
 
+import au.csiro.variantspark.data._
 import au.csiro.variantspark.test.SparkTest
 import it.unimi.dsi.fastutil.longs.Long2DoubleOpenHashMap
 import org.apache.spark.mllib.linalg.Vectors
+import org.apache.spark.rdd.RDD
 import org.junit.Assert._
 import org.junit.Test
-import au.csiro.variantspark.data._
-import au.csiro.variantspark.input._
 
 class WideRandomForrestModelTest extends SparkTest {
+  val doubleComparisonDelta = 1e-6
   val nLabels = 4
   val nSamples = 2
-  val testData =
+  val testData: RDD[(Feature, Long)] =
     sc.parallelize(List(Vectors.zeros(nSamples))).asFeature(BoundedOrdinalVariable(3))
 
   @Test
@@ -47,10 +48,23 @@ class WideRandomForrestModelTest extends SparkTest {
   def whenManyPredictorsThenPredictsByVoting() {
     val assumedPredictions = List(Array(1, 0), Array(1, 2), Array(1, 0))
     val model =
-      new RandomForestModel(assumedPredictions.map(TestPredictorWithImportance(_, null).toMember).toList,
-        nLabels)
+      new RandomForestModel(
+          assumedPredictions
+            .map(TestPredictorWithImportance(_, null).toMember), nLabels)
     val prediction = model.predict(testData)
     assertArrayEquals(Array(1, 0), prediction)
   }
 
+  @Test
+  def predictProbabilities() {
+    val assumedPredictions = List(Array(1, 1), Array(1, 2), Array(1, 1), Array(1, 1))
+    val model =
+      new RandomForestModel(
+          assumedPredictions
+            .map(TestPredictorWithImportance(_, null).toMember), nLabels)
+    val classProbabilities = model.predictProb(testData)
+    assertArrayEquals(Array(0.0, 1.0, 0.0, 0.0), classProbabilities(0), doubleComparisonDelta)
+    assertArrayEquals(Array(0.0, 0.75, 0.25, 0.0), classProbabilities(1), doubleComparisonDelta)
+  }
+
 }