Repository: CitrineInformatics/lolo (language: python)

Issue:
Reproducible training
The intention of this PR is to allow complete reproducibility in Lolo. Given the same training data and the same configuration, we should be able to produce identical models when starting from the same random seed. The biggest change made to accomplish this goal is to **move random seeds from objects to methods**. Currently, a random seed can be specified when the user creates a `RandomForest` object. This is the same pattern as scikit-learn, but I have changed the random seed to be an argument of the `RandomForest`'s `train` method. There are several reasons.

1. Locality makes debugging easier. Putting the random seed on the function makes each function call reproducible, independent of global state. This allows us to drill down into the behavior of a specific piece of code when trying to understand an error.
2. It's more flexible. The same persisted configuration can be retrained with a specific random seed if we need to reproducible an earlier result. But it can also be retrained with an arbitrary seed.
3. It facilitates determinism in a multi-threaded environment. Logic can be fanned out to parallel processes and as long as each function call gets the same random seed, the result is deterministic. If the seed is on the objects then we need to construct a new object holding that seed in order to pass to each parallel process.

The clearest user-facing change is that learners, like `RandomForest` no longer have a random seed and the `train` method does have a random seed. But there are several internal methods that now have a random seed:
* getting the best split
* hyperparameter optimization
* running cross-validation and other model quality tests
* randomly rotating features

I'm submitting this PR as-is because it accomplishes the goal and is quite large already. But there are a few things to change in follow-on PRs.
* Currently, a multi-task RF is not able to pass the random seed all the way down to training the individual leaf learners. It's not an issue in practice because the leaf learners are deterministic, but it's still a potential bug. We'll need to re-architect the multi-task node creation structure to make this possible.
* Similarly, the regression and classification node structure leads to random seeds in some object constructors. These are internal objects that are created as part of training and expose no methods to the outside, so it's not an immediate issue. But it would be better to change how they are structured.

I also removed Scala 2.12 and upgraded the Travis build file from using jdk8 to jdk11


Discussion / hints:
> Also, do we have an explicit test here checking serializability? Would be good to have if not.

The serializability of a Random Forest is checked as part of [LoloPyDataLoaderTest](https://github.com/CitrineInformatics/lolo/blob/main/src/test/scala/io/citrine/lolo/util/LoloPyDataLoaderTest.scala#L121)

Fix patch:
diff --git a/src/main/scala/io/citrine/lolo/TrainingResult.scala b/src/main/scala/io/citrine/lolo/TrainingResult.scala
index 223da1b6..98138213 100644
--- a/src/main/scala/io/citrine/lolo/TrainingResult.scala
+++ b/src/main/scala/io/citrine/lolo/TrainingResult.scala
@@ -1,8 +1,6 @@
 package io.citrine.lolo
 
-/**
-  * Created by maxhutch on 12/4/16.
-  */
+// TODO (PLA-10433): Can we put a type parameter on TrainingResult?
 trait TrainingResult extends Serializable {
 
   /**
diff --git a/src/main/scala/io/citrine/lolo/learners/RandomForest.scala b/src/main/scala/io/citrine/lolo/learners/RandomForest.scala
index 4360a30c..14c20c8a 100644
--- a/src/main/scala/io/citrine/lolo/learners/RandomForest.scala
+++ b/src/main/scala/io/citrine/lolo/learners/RandomForest.scala
@@ -7,7 +7,7 @@ import io.citrine.lolo.transformers.{FeatureRotator, MultiTaskFeatureRotator, Mu
 import io.citrine.lolo.trees.classification.ClassificationTreeLearner
 import io.citrine.lolo.trees.multitask.MultiTaskTreeLearner
 import io.citrine.lolo.trees.regression.RegressionTreeLearner
-import io.citrine.lolo.trees.splits.{ClassificationSplitter, RegressionSplitter}
+import io.citrine.lolo.trees.splits.{ClassificationSplitter, MultiTaskSplitter, RegressionSplitter}
 import io.citrine.lolo.{Learner, TrainingResult}
 
 /**
@@ -90,7 +90,7 @@ case class RandomForest(
             numFeatures = numFeatures,
             maxDepth = maxDepth,
             minLeafInstances = minLeafInstances,
-            randomizePivotLocation = randomizePivotLocation
+            splitter = MultiTaskSplitter(randomizePivotLocation)
           )
         )
         val bagger = MultiTaskBagger(
diff --git a/src/main/scala/io/citrine/lolo/trees/Nodes.scala b/src/main/scala/io/citrine/lolo/trees/Nodes.scala
index 48b8b244..67b8ee92 100644
--- a/src/main/scala/io/citrine/lolo/trees/Nodes.scala
+++ b/src/main/scala/io/citrine/lolo/trees/Nodes.scala
@@ -2,37 +2,33 @@ package io.citrine.lolo.trees
 
 import breeze.linalg.DenseMatrix
 import io.citrine.lolo.trees.splits.Split
-import io.citrine.lolo.{Learner, Model, PredictionResult}
-import io.citrine.random.Random
+import io.citrine.lolo.{Model, PredictionResult, TrainingResult}
 
 import scala.collection.mutable
 
 /**
-  * Class to provide getNode interface for internal and leaf training nodes
+  * Trait to provide getNode interface for internal and leaf training nodes
   *
-  * @param trainingData   that this node sees
-  * @param remainingDepth to stop growing the node
-  * @tparam T type of the input vector
-  * @tparam S type of the model output
+  * @tparam T type of the model output
   */
-abstract class TrainingNode[T <: AnyVal, S](
-    trainingData: Seq[(Vector[T], S, Double)],
-    remainingDepth: Int = Int.MaxValue
-) extends Serializable {
+trait TrainingNode[+T] extends Serializable {
+
+  // TODO (PLA-10433): make this a structured type
+  def trainingData: Seq[(Vector[AnyVal], T, Double)]
 
   /**
     * Get the lightweight prediction node for the output tree
     *
     * @return lightweight prediction node
     */
-  def getNode(): ModelNode[PredictionResult[S]]
+  def modelNode: ModelNode[PredictionResult[T]]
 
   /**
     * Get the feature importance of the subtree below this node
     *
     * @return feature importance as a vector
     */
-  def getFeatureImportance(): mutable.ArraySeq[Double]
+  def featureImportance: mutable.ArraySeq[Double]
 }
 
 trait ModelNode[+T <: PredictionResult[Any]] extends Serializable {
@@ -88,7 +84,7 @@ trait ModelNode[+T <: PredictionResult[Any]] extends Serializable {
   * @param trainingWeight  weight of training data in subtree (i.e. size of unweighted training set)
   * @tparam T type of the output
   */
-class InternalModelNode[T <: PredictionResult[Any]](
+case class InternalModelNode[T <: PredictionResult[Any]](
     split: Split,
     left: ModelNode[T],
     right: ModelNode[T],
@@ -128,7 +124,7 @@ class InternalModelNode[T <: PredictionResult[Any]](
       omitFeatures: Set[Int],
       featureWeights: Map[Int, FeatureWeightFactor]
   ): DenseMatrix[Double] = {
-    val featureIndex = split.getIndex()
+    val featureIndex = split.index
 
     // The hot node is the one that is selected when the feature is present
     val (hot, cold) = if (this.split.turnLeft(input)) {
@@ -181,38 +177,20 @@ class InternalModelNode[T <: PredictionResult[Any]](
   override def getTrainingWeight(): Double = trainingWeight
 }
 
-/**
-  * Average the training data to make a leaf prediction
-  *
-  * @param trainingData to train on
-  */
-class TrainingLeaf[T](
-    trainingData: Seq[(Vector[AnyVal], T, Double)],
-    leafLearner: Learner,
-    depth: Int,
-    rng: Random
-) extends TrainingNode(
-      trainingData = trainingData,
-      remainingDepth = 0
-    ) {
+/** A leaf defined by a training result. */
+trait TrainingLeaf[T] extends TrainingNode[T] {
 
-  /**
-    * Average the training data
-    *
-    * @return lightweight prediction node
-    */
-  def getNode(): ModelNode[PredictionResult[T]] = {
-    new ModelLeaf(
-      leafLearner.train(trainingData, rng = rng).getModel().asInstanceOf[Model[PredictionResult[T]]],
-      depth,
-      trainingData.size.toDouble
-    )
-  }
+  def trainingResult: TrainingResult
+
+  def depth: Int
+
+  def modelNode: ModelNode[PredictionResult[T]] =
+    ModelLeaf(model, depth, trainingData.map(_._3).sum)
 
-  override def getFeatureImportance(): mutable.ArraySeq[Double] = mutable.ArraySeq.fill(trainingData.head._1.size)(0.0)
+  def model: Model[PredictionResult[T]] = trainingResult.getModel().asInstanceOf[Model[PredictionResult[T]]]
 }
 
-class ModelLeaf[T](model: Model[PredictionResult[T]], depth: Int, trainingWeight: Double)
+case class ModelLeaf[T](model: Model[PredictionResult[T]], depth: Int, trainingWeight: Double)
     extends ModelNode[PredictionResult[T]] {
   override def transform(input: Vector[AnyVal]): (PredictionResult[T], TreeMeta) = {
     (model.transform(Seq(input)), TreeMeta(depth))
diff --git a/src/main/scala/io/citrine/lolo/trees/classification/ClassificationTrainingLeaf.scala b/src/main/scala/io/citrine/lolo/trees/classification/ClassificationTrainingLeaf.scala
new file mode 100644
index 00000000..01a5b321
--- /dev/null
+++ b/src/main/scala/io/citrine/lolo/trees/classification/ClassificationTrainingLeaf.scala
@@ -0,0 +1,35 @@
+package io.citrine.lolo.trees.classification
+
+import io.citrine.lolo.{Learner, TrainingResult}
+import io.citrine.lolo.trees.TrainingLeaf
+import io.citrine.random.Random
+
+import scala.collection.mutable
+
+/** Training leaf for a classification tree. */
+case class ClassificationTrainingLeaf(
+    trainingData: Seq[(Vector[AnyVal], Char, Double)],
+    trainingResult: TrainingResult,
+    depth: Int
+) extends TrainingLeaf[Char] {
+
+  override def featureImportance: mutable.ArraySeq[Double] = mutable.ArraySeq.fill(trainingData.head._1.size)(0.0)
+}
+
+object ClassificationTrainingLeaf {
+
+  /**
+    * @param trainingData for this leaf node
+    * @param leafLearner on which to train the data
+    * @param depth depth in the tree
+    * @param rng random number generator, for reproducibility
+    * @return a trained classification leaf node
+    */
+  def build(
+      trainingData: Seq[(Vector[AnyVal], Char, Double)],
+      leafLearner: Learner,
+      depth: Int,
+      rng: Random
+  ): ClassificationTrainingLeaf =
+    ClassificationTrainingLeaf(trainingData, leafLearner.train(trainingData, rng = rng), depth)
+}
diff --git a/src/main/scala/io/citrine/lolo/trees/classification/ClassificationTrainingNode.scala b/src/main/scala/io/citrine/lolo/trees/classification/ClassificationTrainingNode.scala
index 6cbbd9a4..5f49e1b5 100644
--- a/src/main/scala/io/citrine/lolo/trees/classification/ClassificationTrainingNode.scala
+++ b/src/main/scala/io/citrine/lolo/trees/classification/ClassificationTrainingNode.scala
@@ -1,152 +1,99 @@
 package io.citrine.lolo.trees.classification
 
 import io.citrine.lolo.trees.splits.{NoSplit, Split, Splitter}
-import io.citrine.lolo.trees.{InternalModelNode, ModelNode, TrainingLeaf, TrainingNode}
+import io.citrine.lolo.trees.{InternalModelNode, ModelNode, TrainingNode}
 import io.citrine.lolo.{Learner, PredictionResult}
 import io.citrine.random.Random
 
-/**
-  * @param trainingData     that this node sees
-  * @param leafLearner      to train on leaves
-  * @param split            used to split this node's direct descendants
-  * @param deltaImpurity    change in training set impurity between this node vs direct descendants
-  * @param numFeatures      number of features used in selecting splits
-  * @param minLeafInstances minimum training instances per node
-  * @param remainingDepth   to stop growing the node
-  * @param maxDepth         maximum depth of tree
-  * @param numClasses       total number of classes, used as an output dimensionality
-  * @param splitter         used to select splits
-  * @param rng              random number generator for reproducibility
-  */
-class ClassificationTrainingNode(
+case class ClassificationTrainingNode(
     trainingData: Seq[(Vector[AnyVal], Char, Double)],
+    leftNode: TrainingNode[Char],
+    rightNode: TrainingNode[Char],
     leafLearner: Learner,
     split: Split,
     deltaImpurity: Double,
-    numFeatures: Int,
-    minLeafInstances: Int,
-    remainingDepth: Int,
-    maxDepth: Int,
-    numClasses: Int,
-    splitter: Splitter[Char],
-    rng: Random
-) extends TrainingNode(trainingData, remainingDepth) {
-
-  assert(trainingData.size > 1, "If we are going to split, we need at least 2 training rows")
-  assert(
-    !split.isInstanceOf[NoSplit],
-    s"Empty split split for training data: \n${trainingData.map(_.toString() + "\n")}"
-  )
-  private val leftRng = rng.split()
-  private val rightRng = rng.split()
-
-  lazy val (leftTrain, rightTrain) = trainingData.partition(r => split.turnLeft(r._1))
-  assert(
-    leftTrain.nonEmpty && rightTrain.nonEmpty,
-    s"Split ${split} resulted in zero size: ${trainingData.map(_._1(split.getIndex()))}"
-  )
-
-  lazy val leftChild = ClassificationTrainingNode.buildChild(
-    leftTrain,
-    leafLearner,
-    minLeafInstances,
-    remainingDepth,
-    maxDepth,
-    numFeatures,
-    numClasses,
-    splitter,
-    leftRng
-  )
-
-  lazy val rightChild = ClassificationTrainingNode.buildChild(
-    rightTrain,
-    leafLearner,
-    minLeafInstances,
-    remainingDepth,
-    maxDepth,
-    numFeatures,
-    numClasses,
-    splitter,
-    rightRng
-  )
+    numClasses: Int
+) extends TrainingNode[Char] {
 
   /**
     * Get the lightweight prediction node for the output tree
     *
     * @return lightweight prediction node
     */
-  override def getNode(): ModelNode[PredictionResult[Char]] =
-    new InternalModelNode(
+  override def modelNode: ModelNode[PredictionResult[Char]] =
+    InternalModelNode(
       split,
-      leftChild.getNode(),
-      rightChild.getNode(),
+      leftNode.modelNode,
+      rightNode.modelNode,
       numClasses,
-      trainingData.size.toDouble
+      trainingData.map(_._3).sum
     )
 
-  override def getFeatureImportance(): scala.collection.mutable.ArraySeq[Double] = {
+  override def featureImportance: scala.collection.mutable.ArraySeq[Double] = {
     val improvement = deltaImpurity
-    val ans = leftChild.getFeatureImportance().zip(rightChild.getFeatureImportance()).map(p => p._1 + p._2)
-    ans(split.getIndex()) = ans(split.getIndex()) + improvement
+    val ans = leftNode.featureImportance.zip(rightNode.featureImportance).map(p => p._1 + p._2)
+    ans(split.index) = ans(split.index) + improvement
     ans
   }
 }
 
 object ClassificationTrainingNode {
 
-  /**
-    * Build a child node by pre-computing a split
-    *
-    * If there isn't a split, the child is a leaf; otherwise, the child is
-    * another training node
-    *
-    * @param trainingData     for the child
-    * @param leafLearner      to pass through
-    * @param minLeafInstances minimum training instances per node
-    * @param remainingDepth   the number of splits left
-    * @param maxDepth         to compute depth via remainingDepth
-    * @param numFeatures      to consider in the split
-    * @param numClasses       total number of classes, used as an output dimensionality
-    * @param splitter         used to select splits
-    * @param rng              random number generator for reproducibility
-    * @return the child node, either a RegressionTrainingNode or TrainingLeaf
-    */
-  def buildChild(
+  def build(
       trainingData: Seq[(Vector[AnyVal], Char, Double)],
       leafLearner: Learner,
+      splitter: Splitter[Char],
+      numFeatures: Int,
       minLeafInstances: Int,
       remainingDepth: Int,
       maxDepth: Int,
-      numFeatures: Int,
       numClasses: Int,
-      splitter: Splitter[Char],
       rng: Random
-  ): TrainingNode[AnyVal, Char] = {
-    if (
-      trainingData.size >= 2 * minLeafInstances && remainingDepth > 0 && trainingData.exists(
-        _._2 != trainingData.head._2
-      )
-    ) {
-      val (leftSplit, leftDelta) = splitter.getBestSplit(trainingData, numFeatures, minLeafInstances, rng)
-      if (!leftSplit.isInstanceOf[NoSplit]) {
-        new ClassificationTrainingNode(
-          trainingData,
-          leafLearner,
-          leftSplit,
-          leftDelta,
-          numFeatures,
-          minLeafInstances,
-          remainingDepth - 1,
-          maxDepth,
-          numClasses,
-          splitter,
-          rng
-        )
-      } else {
-        new TrainingLeaf(trainingData, leafLearner, maxDepth - remainingDepth, rng)
-      }
+  ): TrainingNode[Char] = {
+    val sufficientData = trainingData.size >= 2 * minLeafInstances &&
+      remainingDepth > 0 &&
+      trainingData.exists(_._2 != trainingData.head._2)
+    val (split: Split, deltaImpurity: Double) = if (sufficientData) {
+      splitter.getBestSplit(trainingData, numFeatures, minLeafInstances, rng)
     } else {
-      new TrainingLeaf(trainingData, leafLearner, maxDepth - remainingDepth, rng)
+      (NoSplit(), 0.0)
+    }
+    split match {
+      case _: NoSplit =>
+        ClassificationTrainingLeaf.build(trainingData, leafLearner, maxDepth - remainingDepth, rng)
+      case split: Split =>
+        val (leftTrain, rightTrain) = trainingData.partition(r => split.turnLeft(r._1))
+        val leftNode = ClassificationTrainingNode.build(
+          trainingData = leftTrain,
+          leafLearner = leafLearner,
+          splitter = splitter,
+          numFeatures = numFeatures,
+          minLeafInstances = minLeafInstances,
+          remainingDepth = remainingDepth - 1,
+          maxDepth = maxDepth,
+          numClasses = numClasses,
+          rng = rng
+        )
+        val rightNode = ClassificationTrainingNode.build(
+          trainingData = rightTrain,
+          leafLearner = leafLearner,
+          splitter = splitter,
+          numFeatures = numFeatures,
+          minLeafInstances = minLeafInstances,
+          remainingDepth = remainingDepth - 1,
+          maxDepth = maxDepth,
+          numClasses = numClasses,
+          rng = rng
+        )
+        ClassificationTrainingNode(
+          trainingData = trainingData,
+          leftNode = leftNode,
+          rightNode = rightNode,
+          leafLearner = leafLearner,
+          split = split,
+          deltaImpurity = deltaImpurity,
+          numClasses = numClasses
+        )
     }
   }
 }
diff --git a/src/main/scala/io/citrine/lolo/trees/classification/ClassificationTree.scala b/src/main/scala/io/citrine/lolo/trees/classification/ClassificationTree.scala
index f9890109..6701ca89 100644
--- a/src/main/scala/io/citrine/lolo/trees/classification/ClassificationTree.scala
+++ b/src/main/scala/io/citrine/lolo/trees/classification/ClassificationTree.scala
@@ -73,41 +73,32 @@ case class ClassificationTreeLearner(
       finalTraining.head._1.size
     }
 
-    /* The tree is built of training nodes */
-    val (split, delta) = splitter.getBestSplit(finalTraining, numFeaturesActual, minLeafInstances, rng)
-    val rootTrainingNode = if (split.isInstanceOf[NoSplit] || maxDepth == 0) {
-      new TrainingLeaf(finalTraining, myLeafLearner, 0, rng)
-    } else {
-      new ClassificationTrainingNode(
-        finalTraining,
-        myLeafLearner,
-        split,
-        delta,
-        numFeaturesActual,
-        remainingDepth = maxDepth - 1,
-        maxDepth = maxDepth,
-        minLeafInstances = minLeafInstances,
-        numClasses = trainingData.map { _._2 }.distinct.length,
-        splitter = splitter,
-        rng = rng
-      )
-    }
-
-    /* Wrap them up in a regression tree */
+    // Recursively build the tree via its nodes and wrap the top node in a ClassificationTreeTrainingResult
+    val rootTrainingNode = ClassificationTrainingNode.build(
+      trainingData = finalTraining,
+      leafLearner = myLeafLearner,
+      splitter = splitter,
+      numFeatures = numFeaturesActual,
+      minLeafInstances = minLeafInstances,
+      remainingDepth = maxDepth,
+      maxDepth = maxDepth,
+      numClasses = trainingData.map(_._2).distinct.length,
+      rng = rng
+    )
     new ClassificationTrainingResult(rootTrainingNode, inputEncoders, outputEncoder)
   }
 }
 
 class ClassificationTrainingResult(
-    rootTrainingNode: TrainingNode[AnyVal, Char],
+    rootTrainingNode: TrainingNode[Char],
     inputEncoders: Seq[Option[CategoricalEncoder[Any]]],
     outputEncoder: CategoricalEncoder[Any]
 ) extends TrainingResult {
   /* Grab a prediction node.  The partitioning happens here */
-  lazy val model = new ClassificationTree(rootTrainingNode.getNode(), inputEncoders, outputEncoder)
+  lazy val model = new ClassificationTree(rootTrainingNode.modelNode, inputEncoders, outputEncoder)
 
   /* Grab the feature influences */
-  lazy val importance = rootTrainingNode.getFeatureImportance()
+  lazy val importance = rootTrainingNode.featureImportance
   private lazy val importanceNormalized = {
     if (Math.abs(importance.sum) > 0) {
       importance.map(_ / importance.sum)
diff --git a/src/main/scala/io/citrine/lolo/trees/impurity/MultiImpurityCalculator.scala b/src/main/scala/io/citrine/lolo/trees/impurity/MultiImpurityCalculator.scala
index a456ae0c..250926b3 100644
--- a/src/main/scala/io/citrine/lolo/trees/impurity/MultiImpurityCalculator.scala
+++ b/src/main/scala/io/citrine/lolo/trees/impurity/MultiImpurityCalculator.scala
@@ -7,7 +7,7 @@ package io.citrine.lolo.trees.impurity
   */
 class MultiImpurityCalculator(
     calculators: Seq[ImpurityCalculator[AnyVal]]
-) extends ImpurityCalculator[Array[AnyVal]] {
+) extends ImpurityCalculator[Seq[AnyVal]] {
 
   /**
     * Add the value to each calculator
@@ -16,7 +16,7 @@ class MultiImpurityCalculator(
     * @param weight of the value
     * @return the impurity after adding
     */
-  def add(value: Array[AnyVal], weight: Double): Double = {
+  def add(value: Seq[AnyVal], weight: Double): Double = {
     value.zip(calculators).map {
       case (v, calc) =>
         if (v.isInstanceOf[Double]) {
@@ -36,7 +36,7 @@ class MultiImpurityCalculator(
     * @param weight of the value
     * @return the impurity after removing
     */
-  def remove(value: Array[AnyVal], weight: Double): Double = {
+  def remove(value: Seq[AnyVal], weight: Double): Double = {
     value.zip(calculators).map {
       case (v, calc) =>
         if (v.isInstanceOf[Double]) {
@@ -78,7 +78,7 @@ object MultiImpurityCalculator {
     * @param weights which are assumed to be constant over the labels at each row
     * @return MultiImpurityCalculator that sum the impurity of each label index
     */
-  def build(labels: Seq[Array[AnyVal]], weights: Seq[Double]): MultiImpurityCalculator = {
+  def build(labels: Seq[Seq[AnyVal]], weights: Seq[Double]): MultiImpurityCalculator = {
     val calculators: Seq[ImpurityCalculator[AnyVal]] = labels.transpose.map { labelSeq =>
       if (labelSeq.head.isInstanceOf[Double]) {
         VarianceCalculator
diff --git a/src/main/scala/io/citrine/lolo/trees/multitask/MultiTaskTrainingNode.scala b/src/main/scala/io/citrine/lolo/trees/multitask/MultiTaskTrainingNode.scala
index c087a945..0c43e323 100644
--- a/src/main/scala/io/citrine/lolo/trees/multitask/MultiTaskTrainingNode.scala
+++ b/src/main/scala/io/citrine/lolo/trees/multitask/MultiTaskTrainingNode.scala
@@ -2,171 +2,195 @@ package io.citrine.lolo.trees.multitask
 
 import io.citrine.lolo.PredictionResult
 import io.citrine.lolo.linear.GuessTheMeanLearner
+import io.citrine.lolo.trees.classification.ClassificationTrainingLeaf
 import io.citrine.lolo.trees.regression.RegressionTrainingLeaf
-import io.citrine.lolo.trees.splits.{MultiTaskSplitter, NoSplit}
-import io.citrine.lolo.trees.{InternalModelNode, ModelNode, TrainingLeaf}
+import io.citrine.lolo.trees.splits.{MultiTaskSplitter, NoSplit, Split}
+import io.citrine.lolo.trees.{InternalModelNode, ModelNode, TrainingNode}
 import io.citrine.random.Random
 
 import scala.collection.mutable
 
-/**
-  * Node in a multi-task training tree, which can produce nodes for its model trees.
-  * Splits are chosen using a MultiTaskSplitter, which considers the sum impurity decrease across all tasks.
-  *
-  * @param inputs data on which to select splits and form models
-  * @param numFeatures to randomly select from at each split (negative values indicate that all features should be considered)
-  * @param maxDepth to grow the tree to
-  * @param minInstances minimum number of training instances per leaf
-  * @param randomizePivotLocation whether to generate splits randomly between the data points
-  * @param rng random number generator, for reproducibility
-  */
-class MultiTaskTrainingNode(
-    inputs: Seq[(Vector[AnyVal], Array[AnyVal], Double)],
-    numFeatures: Int,
-    maxDepth: Int,
-    minInstances: Int,
-    randomizePivotLocation: Boolean = false,
-    rng: Random
+case class MultiTaskTrainingNode(
+    trainingData: Seq[(Vector[AnyVal], Vector[AnyVal], Double)],
+    labelWiseInstructions: Seq[MultiTaskLabelInstruction],
+    deltaImpurity: Double
 ) {
 
-  // Compute a split
-  val (split, deltaImpurity) = if (maxDepth <= 0) {
-    (new NoSplit, 0.0)
-  } else {
-    MultiTaskSplitter(randomizePivotLocation = randomizePivotLocation).getBestSplit(
-      inputs,
-      numFeatures,
-      minInstances,
-      rng
-    )
-  }
-
-  // Try to construct left and right children
-  val (leftChild: Option[MultiTaskTrainingNode], rightChild: Option[MultiTaskTrainingNode]) = split match {
-    case _: NoSplit => (None, None)
-    case _: Any =>
-      val (leftData, rightData) = inputs.partition(row => split.turnLeft(row._1))
-      (
-        Some(
-          new MultiTaskTrainingNode(
-            leftData,
-            numFeatures,
-            maxDepth - 1,
-            minInstances,
-            randomizePivotLocation,
-            rng = rng.split()
-          )
-        ),
-        Some(
-          new MultiTaskTrainingNode(
-            rightData,
-            numFeatures,
-            maxDepth - 1,
-            minInstances,
-            randomizePivotLocation,
-            rng = rng.split()
-          )
-        )
-      )
-  }
-
-  // get feature importance for the i'th label
-  def getFeatureImportance(index: Int): mutable.ArraySeq[Double] = {
-    // Filter out "missing" values, which are NaN for regression and 0 for encoded categoricals
-    val label = inputs.head._2(index)
-    val reducedData = if (label.isInstanceOf[Double]) {
-      inputs.map(x => (x._1, x._2(index).asInstanceOf[Double], x._3)).filterNot(_._2.isNaN)
-    } else {
-      inputs.map(x => (x._1, x._2(index).asInstanceOf[Char], x._3)).filter(_._2 > 0)
-    }
-    // Compute the valid data for each child
-    val (left, right) = reducedData.partition(r => split.turnLeft(r._1))
-
-    // get feature importance from the children if they exist, or from this node if it is a leaf
-    (leftChild, rightChild) match {
-      case (Some(theLeftChild), Some(theRightChild)) if left.nonEmpty && right.nonEmpty =>
-        val ans =
-          theLeftChild.getFeatureImportance(index).zip(theRightChild.getFeatureImportance(index)).map(p => p._1 + p._2)
-        ans(split.getIndex()) = ans(split.getIndex()) + deltaImpurity
+  // get feature importance for the `index`th label
+  def featureImportanceByLabelIndex(index: Int): mutable.ArraySeq[Double] = {
+    labelWiseInstructions(index) match {
+      case Stop(leaf)             => leaf.featureImportance
+      case FollowChild(childNode) => childNode.featureImportanceByLabelIndex(index)
+      case DoSplit(split, leftNode, rightNode) =>
+        val ans = leftNode
+          .featureImportanceByLabelIndex(index)
+          .zip(rightNode.featureImportanceByLabelIndex(index))
+          .map(p => p._1 + p._2)
+        ans(split.index) = ans(split.index) + deltaImpurity
         ans
-      case (Some(theLeftChild), _) if left.nonEmpty =>
-        theLeftChild.getFeatureImportance(index)
-      case (_, Some(theRightChild)) if right.nonEmpty =>
-        theRightChild.getFeatureImportance(index)
-      case (_, _) =>
-        // TODO (PLA-10415): get the rng in here somehow (right now it's instantiating a random rng)
-        if (label.isInstanceOf[Double]) {
-          new RegressionTrainingLeaf(
-            reducedData.asInstanceOf[Seq[(Vector[AnyVal], Double, Double)]],
-            GuessTheMeanLearner(),
-            1,
-            rng = Random()
-          ).getFeatureImportance()
-        } else {
-          new TrainingLeaf[Char](
-            reducedData.asInstanceOf[Seq[(Vector[AnyVal], Char, Double)]],
-            GuessTheMeanLearner(),
-            1,
-            rng = Random()
-          ).getFeatureImportance()
-        }
+      case Inaccessible() => throw new RuntimeException(s"No valid training data present for label $index")
     }
   }
 
   // Construct the model node for the `index`th label
-  def getNode(index: Int): ModelNode[PredictionResult[Any]] = {
-    // Filter out "missing" values, which are NaN for regression and 0 for encoded categoricals
-    val label = inputs.head._2(index)
-    val reducedData = if (label.isInstanceOf[Double]) {
-      inputs.map(x => (x._1, x._2(index).asInstanceOf[Double], x._3)).filterNot(_._2.isNaN)
-    } else {
-      inputs.map(x => (x._1, x._2(index).asInstanceOf[Char], x._3)).filter(_._2 > 0)
-    }
-    // Compute the valid data for each child
-    val (left, right) = reducedData.partition(r => split.turnLeft(r._1))
-
-    // Construct an internal node if the children are defined and actually have valid data
-    (leftChild, rightChild) match {
-      case (Some(theLeftChild), Some(theRightChild)) if left.nonEmpty && right.nonEmpty =>
-        if (label.isInstanceOf[Double]) {
+  def modelNodeByLabelIndex(index: Int): ModelNode[PredictionResult[Any]] = {
+    labelWiseInstructions(index) match {
+      case Stop(leaf)             => leaf.modelNode
+      case FollowChild(childNode) => childNode.modelNodeByLabelIndex(index)
+      case DoSplit(split, leftNode, rightNode) =>
+        if (trainingData.head._2(index).isInstanceOf[Double]) {
           new InternalModelNode[PredictionResult[Double]](
-            split,
-            theLeftChild.getNode(index).asInstanceOf[ModelNode[PredictionResult[Double]]],
-            theRightChild.getNode(index).asInstanceOf[ModelNode[PredictionResult[Double]]],
-            outputDimension = 0, // Don't support multitask SHAP at this time.
-            trainingWeight = reducedData.length.toDouble
+            split = split,
+            left = leftNode.modelNodeByLabelIndex(index).asInstanceOf[ModelNode[PredictionResult[Double]]],
+            right = rightNode.modelNodeByLabelIndex(index).asInstanceOf[ModelNode[PredictionResult[Double]]],
+            outputDimension = 0, // Shapley is not implemented for multi-task nodes
+            trainingWeight = trainingData.map(_._3).sum
           )
         } else {
-          if (!label.isInstanceOf[Char]) throw new IllegalArgumentException("Training data wasn't double or char")
           new InternalModelNode[PredictionResult[Char]](
-            split,
-            theLeftChild.getNode(index).asInstanceOf[ModelNode[PredictionResult[Char]]],
-            theRightChild.getNode(index).asInstanceOf[ModelNode[PredictionResult[Char]]],
-            outputDimension = 0, // Don't support multitask SHAP at this time.
-            trainingWeight = reducedData.length.toDouble
+            split = split,
+            left = leftNode.modelNodeByLabelIndex(index).asInstanceOf[ModelNode[PredictionResult[Char]]],
+            right = rightNode.modelNodeByLabelIndex(index).asInstanceOf[ModelNode[PredictionResult[Char]]],
+            outputDimension = 0,
+            trainingWeight = trainingData.map(_._3).sum
           )
         }
-      case (Some(theLeftChild), _) if left.nonEmpty =>
-        theLeftChild.getNode(index)
-      case (_, Some(theRightChild)) if right.nonEmpty =>
-        theRightChild.getNode(index)
-      case (_, _) =>
-        // TODO (PLA-10415): get the rng in here somehow (right now it's instantiating a random rng)
-        if (label.isInstanceOf[Double]) {
-          new RegressionTrainingLeaf(
-            reducedData.asInstanceOf[Seq[(Vector[AnyVal], Double, Double)]],
-            GuessTheMeanLearner(),
-            1,
-            rng = Random()
-          ).getNode()
+      case Inaccessible() => throw new RuntimeException(s"No valid training data present for label $index")
+    }
+  }
+}
+
+object MultiTaskTrainingNode {
+
+  /**
+    * Build a node by computing splits and recursively building the child nodes.
+    *
+    * @param trainingData on which to build the split
+    * @param numFeatures to considering when determining the split
+    * @param remainingDepth the maximum number of splits left
+    * @param maxDepth maximum number of splits (used to compute depth of this node)
+    * @param minInstances minimum training instances per node
+    * @param splitter determines the best split
+    * @param rng random number generator, for reproducibility
+    * @return
+    */
+  def build(
+      trainingData: Seq[(Vector[AnyVal], Vector[AnyVal], Double)],
+      numFeatures: Int,
+      remainingDepth: Int,
+      maxDepth: Int,
+      minInstances: Int,
+      splitter: MultiTaskSplitter,
+      rng: Random = Random()
+  ): MultiTaskTrainingNode = {
+    // Determine the overall split.
+    val sufficientData = trainingData.size >= 2 * minInstances &&
+      remainingDepth > 0 &&
+      trainingData.exists(row => row._2 != trainingData.head._2)
+    val (split: Split, deltaImpurity: Double) = if (sufficientData) {
+      splitter.getBestSplit(
+        trainingData,
+        numFeatures,
+        minInstances,
+        rng
+      )
+    } else {
+      (NoSplit(), 0.0)
+    }
+    // Build the left and right nodes (assuming the split exists).
+    val nodesOpt: Option[(MultiTaskTrainingNode, MultiTaskTrainingNode)] = split match {
+      case _: NoSplit => None
+      case split: Split =>
+        val (leftTrain, rightTrain) = trainingData.partition(r => split.turnLeft(r._1))
+        val leftNode = MultiTaskTrainingNode.build(
+          trainingData = leftTrain,
+          numFeatures = numFeatures,
+          remainingDepth = remainingDepth - 1,
+          maxDepth = maxDepth,
+          minInstances = minInstances,
+          splitter = splitter,
+          rng = rng
+        )
+        val rightNode = MultiTaskTrainingNode.build(
+          trainingData = rightTrain,
+          numFeatures = numFeatures,
+          remainingDepth = remainingDepth - 1,
+          maxDepth = maxDepth,
+          minInstances = minInstances,
+          splitter = splitter,
+          rng = rng
+        )
+        Some((leftNode, rightNode))
+    }
+    // Determine what to do for each label (evaluate a model, evaluate the split, go left, or go right).
+    val exampleRow = trainingData.head
+    val labelWiseInstructions = exampleRow._2.indices.map { index =>
+      // Determine how much data *with this label* goes down each branch
+      val exampleLabel = exampleRow._2(index)
+      val reducedData = if (exampleLabel.isInstanceOf[Double]) {
+        trainingData.map(x => (x._1, x._2(index).asInstanceOf[Double], x._3)).filterNot(_._2.isNaN)
+      } else {
+        trainingData.map(x => (x._1, x._2(index).asInstanceOf[Char], x._3)).filter(_._2 > 0)
+      }
+      val (left, right) = reducedData.partition(r => split.turnLeft(r._1))
+
+      if (reducedData.isEmpty) {
+        Inaccessible()
+      } else if (nodesOpt.isEmpty || reducedData.length <= minInstances) {
+        // Either there's no split or there's not enough data remaining with this label. Build a leaf node.
+        val trainingLeaf: TrainingNode[AnyVal] = if (exampleLabel.isInstanceOf[Double]) {
+          RegressionTrainingLeaf
+            .build(
+              reducedData.asInstanceOf[Seq[(Vector[AnyVal], Double, Double)]],
+              GuessTheMeanLearner(),
+              maxDepth - remainingDepth,
+              rng
+            )
         } else {
-          new TrainingLeaf[Char](
-            reducedData.asInstanceOf[Seq[(Vector[AnyVal], Char, Double)]],
-            GuessTheMeanLearner(),
-            1,
-            rng = Random()
-          ).getNode()
+          ClassificationTrainingLeaf
+            .build(
+              reducedData.asInstanceOf[Seq[(Vector[AnyVal], Char, Double)]],
+              GuessTheMeanLearner(),
+              maxDepth - remainingDepth,
+              rng
+            )
         }
+        Stop(trainingLeaf)
+      } else if (left.nonEmpty && right.nonEmpty) {
+        // From this point onwards we know that nodesOpt is defined. Follow the split.
+        DoSplit(split, nodesOpt.get._1, nodesOpt.get._2)
+      } else if (left.nonEmpty) {
+        // This label only has data on the left-hand side, so go left.
+        FollowChild(nodesOpt.get._1)
+      } else {
+        // This label only has data on the right-hand side, so go right.
+        FollowChild(nodesOpt.get._2)
+      }
     }
+
+    MultiTaskTrainingNode(
+      trainingData = trainingData,
+      labelWiseInstructions = labelWiseInstructions,
+      deltaImpurity = deltaImpurity
+    )
   }
 }
+
+/** An enumeration specifying what behavior to perform for a specific label. */
+sealed trait MultiTaskLabelInstruction
+
+/** Follow the proscribed split. */
+case class DoSplit(
+    split: Split,
+    leftNode: MultiTaskTrainingNode,
+    rightNode: MultiTaskTrainingNode
+) extends MultiTaskLabelInstruction
+
+/** Always go to the specified child node. */
+case class FollowChild(childNode: MultiTaskTrainingNode) extends MultiTaskLabelInstruction
+
+/** Stop and evaluate the provided leaf node. */
+case class Stop(leaf: TrainingNode[AnyVal]) extends MultiTaskLabelInstruction
+
+/** A placeholder for inaccessible sections of the tree, such as underneath a Stop() instruction. */
+case class Inaccessible() extends MultiTaskLabelInstruction
diff --git a/src/main/scala/io/citrine/lolo/trees/multitask/MultiTaskTree.scala b/src/main/scala/io/citrine/lolo/trees/multitask/MultiTaskTree.scala
index fac307f4..5eacb3e2 100644
--- a/src/main/scala/io/citrine/lolo/trees/multitask/MultiTaskTree.scala
+++ b/src/main/scala/io/citrine/lolo/trees/multitask/MultiTaskTree.scala
@@ -5,6 +5,7 @@ import io.citrine.lolo.encoders.CategoricalEncoder
 import io.citrine.lolo.trees.ModelNode
 import io.citrine.lolo.trees.classification.ClassificationTree
 import io.citrine.lolo.trees.regression.RegressionTree
+import io.citrine.lolo.trees.splits.MultiTaskSplitter
 import io.citrine.lolo.{Model, MultiTaskLearner, MultiTaskTrainingResult, ParallelModels, PredictionResult}
 
 /**
@@ -13,13 +14,13 @@ import io.citrine.lolo.{Model, MultiTaskLearner, MultiTaskTrainingResult, Parall
   * @param numFeatures to random select from at each split (numbers less than 0 indicate that all features are used)
   * @param maxDepth to grow the tree to
   * @param minLeafInstances minimum number of training instances per leaf
-  * @param randomizePivotLocation whether to generate splits randomly between the data points
+  * @param splitter to determine the best split given data
   */
 case class MultiTaskTreeLearner(
     numFeatures: Int = -1,
     maxDepth: Int = 30,
     minLeafInstances: Int = 1,
-    randomizePivotLocation: Boolean = false
+    splitter: MultiTaskSplitter = MultiTaskSplitter(randomizePivotLocation = true)
 ) extends MultiTaskLearner {
 
   /**
@@ -65,7 +66,7 @@ case class MultiTaskTreeLearner(
     // Encode the inputs, outputs, and filter out zero weight rows
     val collectedData = inputs.indices
       .map { i =>
-        (encodedInputs(i), encodedLabels(i).toArray, weights.map(_(i)).getOrElse(1.0))
+        (encodedInputs(i), encodedLabels(i), weights.map(_(i)).getOrElse(1.0))
       }
       .filter(_._3 > 0.0)
 
@@ -77,17 +78,18 @@ case class MultiTaskTreeLearner(
     }
 
     // Construct the training tree
-    val root = new MultiTaskTrainingNode(
-      inputs = collectedData,
+    val root = MultiTaskTrainingNode.build(
+      trainingData = collectedData,
       numFeatures = numFeaturesActual,
+      remainingDepth = maxDepth,
       maxDepth = maxDepth,
       minInstances = minLeafInstances,
-      randomizePivotLocation = randomizePivotLocation,
+      splitter = splitter,
       rng = rng
     )
 
     // Construct the model trees
-    val nodes = labelIndices.map(root.getNode)
+    val nodes = labelIndices.map(root.modelNodeByLabelIndex)
 
     // Stick the model trees into RegressionTree and ClassificationTree objects
     val models = labelIndices.map { i =>
@@ -108,7 +110,7 @@ case class MultiTaskTreeLearner(
     val sumFeatureImportance: Vector[Double] = {
       val startingImportances = Vector.fill(repInput.length)(0.0)
       labelIndices.foldLeft(startingImportances) { (importance, i) =>
-        root.getFeatureImportance(i).toVector.zip(importance).map(p => p._1 + p._2)
+        root.featureImportanceByLabelIndex(i).toVector.zip(importance).map(p => p._1 + p._2)
       }
     }
 
diff --git a/src/main/scala/io/citrine/lolo/trees/regression/RegressionTrainingLeaf.scala b/src/main/scala/io/citrine/lolo/trees/regression/RegressionTrainingLeaf.scala
index a92f7824..c013c46a 100644
--- a/src/main/scala/io/citrine/lolo/trees/regression/RegressionTrainingLeaf.scala
+++ b/src/main/scala/io/citrine/lolo/trees/regression/RegressionTrainingLeaf.scala
@@ -1,37 +1,25 @@
 package io.citrine.lolo.trees.regression
 
-import io.citrine.lolo.trees.{ModelLeaf, ModelNode, TrainingNode}
-import io.citrine.lolo.{Learner, Model, PredictionResult}
+import io.citrine.lolo.trees.TrainingLeaf
+import io.citrine.lolo.{Learner, TrainingResult}
 import io.citrine.random.Random
 
 import scala.collection.mutable
 
-/**
-  * Training leaf node for regression trees
-  */
-class RegressionTrainingLeaf(
+/** Training leaf node for regression trees. */
+case class RegressionTrainingLeaf(
     trainingData: Seq[(Vector[AnyVal], Double, Double)],
-    leafLearner: Learner,
-    depth: Int,
-    rng: Random
-) extends TrainingNode(trainingData, depth) {
-
-  /**
-    * Wrap the leaf model (previously trained) in a lightweight leaf node
-    *
-    * @return lightweight prediction node
-    */
-  def getNode(): ModelNode[PredictionResult[Double]] = {
-    new ModelLeaf(model, depth, trainingData.size.toDouble)
-  }
+    trainingResult: TrainingResult,
+    depth: Int
+) extends TrainingLeaf[Double] {
 
   /**
     * Pull the leaf model's feature importance and rescale it by the remaining impurity
     *
     * @return feature importance as a vector
     */
-  def getFeatureImportance(): scala.collection.mutable.ArraySeq[Double] = {
-    importance match {
+  override def featureImportance: scala.collection.mutable.ArraySeq[Double] = {
+    trainingResult.getFeatureImportance() match {
       case Some(x) =>
         // Compute the weighted sum of the label, the square label, and the weights
         val expectations: (Double, Double, Double) = trainingData
@@ -47,13 +35,22 @@ class RegressionTrainingLeaf(
       case None => mutable.ArraySeq.fill(trainingData.head._1.size)(0.0)
     }
   }
+}
 
-  /** Train the leaf learner on the training data */
-  val leafTrainingResult = leafLearner.train(trainingData, rng = rng)
-
-  /** Pull out the model for future use */
-  val model = leafTrainingResult.getModel().asInstanceOf[Model[PredictionResult[Double]]]
+object RegressionTrainingLeaf {
 
-  /** Pull out the importance for future use */
-  val importance = leafTrainingResult.getFeatureImportance()
+  /**
+    * @param trainingData for this leaf node
+    * @param leafLearner on which to train the data
+    * @param depth depth in the tree
+    * @param rng random number generator, for reproducibility
+    * @return a trained regression leaf node
+    */
+  def build(
+      trainingData: Seq[(Vector[AnyVal], Double, Double)],
+      leafLearner: Learner,
+      depth: Int,
+      rng: Random
+  ): RegressionTrainingLeaf =
+    RegressionTrainingLeaf(trainingData, leafLearner.train(trainingData, rng = rng), depth)
 }
diff --git a/src/main/scala/io/citrine/lolo/trees/regression/RegressionTrainingNode.scala b/src/main/scala/io/citrine/lolo/trees/regression/RegressionTrainingNode.scala
index ea2cc62b..d8dd98ca 100644
--- a/src/main/scala/io/citrine/lolo/trees/regression/RegressionTrainingNode.scala
+++ b/src/main/scala/io/citrine/lolo/trees/regression/RegressionTrainingNode.scala
@@ -5,74 +5,27 @@ import io.citrine.lolo.trees.{InternalModelNode, ModelNode, TrainingNode}
 import io.citrine.lolo.{Learner, PredictionResult}
 import io.citrine.random.Random
 
-/**
-  * Created by maxhutch on 1/12/17.
-  */
-class RegressionTrainingNode(
+case class RegressionTrainingNode(
     trainingData: Seq[(Vector[AnyVal], Double, Double)],
+    leftNode: TrainingNode[Double],
+    rightNode: TrainingNode[Double],
     leafLearner: Learner,
-    splitter: Splitter[Double],
     split: Split,
-    deltaImpurity: Double,
-    numFeatures: Int,
-    minLeafInstances: Int,
-    remainingDepth: Int,
-    maxDepth: Int,
-    rng: Random
-) extends TrainingNode(
-      trainingData = trainingData,
-      remainingDepth = remainingDepth
-    ) {
-
-  // TODO (PLA-10415): see if there's a way to get the rng for this (and the other nodes) out of the constructor
-  // val (split: Split, deltaImpurity: Double) = RegressionSplitter.getBestSplit(trainingData, numFeatures)
-  assert(trainingData.size > 1, "If we are going to split, we need at least 2 training rows")
-  assert(
-    !split.isInstanceOf[NoSplit],
-    s"Empty split split for training data: \n${trainingData.map(_.toString() + "\n")}"
-  )
-  private val leftRng = rng.split()
-  private val rightRng = rng.split()
-
-  lazy val (leftTrain, rightTrain) = trainingData.partition(r => split.turnLeft(r._1))
-  assert(
-    leftTrain.nonEmpty && rightTrain.nonEmpty,
-    s"Split ${split} resulted in zero size: ${trainingData.map(_._1(split.getIndex()))}"
-  )
-
-  lazy val leftChild = RegressionTrainingNode.buildChild(
-    leftTrain,
-    leafLearner,
-    splitter,
-    minLeafInstances,
-    remainingDepth,
-    maxDepth,
-    numFeatures,
-    leftRng
-  )
-  lazy val rightChild = RegressionTrainingNode.buildChild(
-    rightTrain,
-    leafLearner,
-    splitter,
-    minLeafInstances,
-    remainingDepth,
-    maxDepth,
-    numFeatures,
-    rightRng
-  )
+    deltaImpurity: Double
+) extends TrainingNode[Double] {
 
   /**
     * Get the lightweight prediction node for the output tree
     *
     * @return lightweight prediction node
     */
-  override def getNode(): ModelNode[PredictionResult[Double]] = {
+  override def modelNode: ModelNode[PredictionResult[Double]] = {
     new InternalModelNode[PredictionResult[Double]](
       split,
-      leftChild.getNode(),
-      rightChild.getNode(),
+      leftNode.modelNode,
+      rightNode.modelNode,
       1,
-      trainingData.size.toDouble
+      trainingData.map(_._3).sum
     )
   }
 
@@ -84,68 +37,81 @@ class RegressionTrainingNode(
     *
     * @return feature importance as a vector
     */
-  override def getFeatureImportance(): scala.collection.mutable.ArraySeq[Double] = {
+  override def featureImportance: scala.collection.mutable.ArraySeq[Double] = {
     val improvement = deltaImpurity
-    val ans = leftChild.getFeatureImportance().zip(rightChild.getFeatureImportance()).map(p => p._1 + p._2)
-    ans(split.getIndex()) = ans(split.getIndex()) + improvement
+    val ans = leftNode.featureImportance.zip(rightNode.featureImportance).map(p => p._1 + p._2)
+    ans(split.index) = ans(split.index) + improvement
     ans
   }
 }
 
-/**
-  * Companion object to hold helper functions
-  */
 object RegressionTrainingNode {
 
   /**
-    * Build a child node by pre-computing a split
+    * Build a node by computing the split and recursively building the child nodes if they exist.
+    * If there isn't a split, this becomes a leaf node.
     *
-    * If there isn't a split, the child is a leaf; otherwise, the child is
-    * another training node
-    *
-    * @param trainingData     for the child
-    * @param leafLearner      to pass through
+    * @param trainingData     on which to split
+    * @param leafLearner      to pass through until we get to a leaf node
+    * @param splitter         determines the best split
     * @param minLeafInstances minimum training instances per node
-    * @param remainingDepth   the number of splits left
+    * @param remainingDepth   the maximum number of splits left
     * @param maxDepth         to compute depth via remainingDepth
     * @param numFeatures      to consider in the split
-    * @param rng              random number generator for reproducibility
+    * @param rng              random number generator, for reproducibility
     * @return the child node, either a RegressionTrainingNode or TrainingLeaf
     */
-  def buildChild(
+  def build(
       trainingData: Seq[(Vector[AnyVal], Double, Double)],
       leafLearner: Learner,
       splitter: Splitter[Double],
+      numFeatures: Int,
       minLeafInstances: Int,
       remainingDepth: Int,
       maxDepth: Int,
-      numFeatures: Int,
       rng: Random
-  ): TrainingNode[AnyVal, Double] = {
-    if (
-      trainingData.size >= 2 * minLeafInstances && remainingDepth > 0 && trainingData.exists(
-        _._2 != trainingData.head._2
-      )
-    ) {
-      val (leftSplit, leftDelta) = splitter.getBestSplit(trainingData, numFeatures, minLeafInstances, rng = rng)
-      if (!leftSplit.isInstanceOf[NoSplit]) {
-        new RegressionTrainingNode(
-          trainingData,
-          leafLearner,
-          splitter,
-          leftSplit,
-          leftDelta,
-          numFeatures,
-          minLeafInstances,
-          remainingDepth - 1,
-          maxDepth,
-          rng
-        )
-      } else {
-        new RegressionTrainingLeaf(trainingData, leafLearner, maxDepth - remainingDepth, rng)
-      }
+  ): TrainingNode[Double] = {
+    val sufficientData = trainingData.size >= 2 * minLeafInstances &&
+      remainingDepth > 0 &&
+      trainingData.exists(_._2 != trainingData.head._2)
+    val (split: Split, deltaImpurity: Double) = if (sufficientData) {
+      splitter.getBestSplit(trainingData, numFeatures, minLeafInstances, rng)
     } else {
-      new RegressionTrainingLeaf(trainingData, leafLearner, maxDepth - remainingDepth, rng)
+      (NoSplit(), 0.0)
+    }
+    split match {
+      case _: NoSplit =>
+        RegressionTrainingLeaf.build(trainingData, leafLearner, maxDepth - remainingDepth, rng)
+      case split: Split =>
+        val (leftTrain, rightTrain) = trainingData.partition(r => split.turnLeft(r._1))
+        val leftNode = RegressionTrainingNode.build(
+          trainingData = leftTrain,
+          leafLearner = leafLearner,
+          splitter = splitter,
+          numFeatures = numFeatures,
+          minLeafInstances = minLeafInstances,
+          remainingDepth = remainingDepth - 1,
+          maxDepth = maxDepth,
+          rng = rng
+        )
+        val rightNode = RegressionTrainingNode.build(
+          trainingData = rightTrain,
+          leafLearner = leafLearner,
+          splitter = splitter,
+          numFeatures = numFeatures,
+          minLeafInstances = minLeafInstances,
+          remainingDepth = remainingDepth - 1,
+          maxDepth = maxDepth,
+          rng = rng
+        )
+        RegressionTrainingNode(
+          trainingData = trainingData,
+          leftNode = leftNode,
+          rightNode = rightNode,
+          leafLearner = leafLearner,
+          split = split,
+          deltaImpurity = deltaImpurity
+        )
     }
   }
 }
diff --git a/src/main/scala/io/citrine/lolo/trees/regression/RegressionTree.scala b/src/main/scala/io/citrine/lolo/trees/regression/RegressionTree.scala
index 0b4fdef3..bc8b6dcc 100644
--- a/src/main/scala/io/citrine/lolo/trees/regression/RegressionTree.scala
+++ b/src/main/scala/io/citrine/lolo/trees/regression/RegressionTree.scala
@@ -4,7 +4,7 @@ import breeze.linalg.DenseMatrix
 import io.citrine.random.Random
 import io.citrine.lolo.encoders.CategoricalEncoder
 import io.citrine.lolo.linear.GuessTheMeanLearner
-import io.citrine.lolo.trees.splits.{NoSplit, RegressionSplitter, Splitter}
+import io.citrine.lolo.trees.splits.{RegressionSplitter, Splitter}
 import io.citrine.lolo.trees.{ModelNode, TrainingNode, TreeMeta}
 import io.citrine.lolo.{Learner, Model, PredictionResult, TrainingResult}
 
@@ -59,10 +59,10 @@ case class RegressionTreeLearner(
         }
     }
 
-    /* Encode the training data */
+    // Encode the training data.
     val encodedTraining = trainingData.map(p => (CategoricalEncoder.encodeInput(p._1, encoders), p._2))
 
-    /* Add the weights to the (features, label) tuples and remove any with zero weight */
+    // Add the weights to the (features, label) tuples and remove any with zero weight.
     val finalTraining = encodedTraining
       .zip(weights.getOrElse(Seq.fill(trainingData.size)(1.0)))
       .map {
@@ -77,44 +77,35 @@ case class RegressionTreeLearner(
       s"We need to have at least 4 rows with non-zero weights, only ${finalTraining.size} given"
     )
 
-    /* If the number of features isn't specified, use all of them */
+    // If the number of features isn't specified, use all of them.
     val numFeaturesActual = if (numFeatures > 0) {
       numFeatures
     } else {
       finalTraining.head._1.size
     }
 
-    /* The tree is built of training nodes */
-    val (split, delta) = splitter.getBestSplit(finalTraining, numFeaturesActual, minLeafInstances, rng)
-    val rootTrainingNode: TrainingNode[AnyVal, Double] = if (split.isInstanceOf[NoSplit] || maxDepth == 0) {
-      new RegressionTrainingLeaf(finalTraining, myLeafLearner, 0, rng)
-    } else {
-      new RegressionTrainingNode(
-        finalTraining,
-        myLeafLearner,
-        splitter,
-        split,
-        delta,
-        numFeaturesActual,
-        minLeafInstances = minLeafInstances,
-        remainingDepth = maxDepth - 1,
-        maxDepth,
-        rng
-      )
-    }
-
-    /* Wrap them up in a regression tree */
+    // Recursively build the tree via its nodes and wrap the top node in a RegressionTreeTrainingResult.
+    val rootTrainingNode = RegressionTrainingNode.build(
+      trainingData = finalTraining,
+      leafLearner = myLeafLearner,
+      splitter = splitter,
+      numFeatures = numFeaturesActual,
+      minLeafInstances = minLeafInstances,
+      remainingDepth = maxDepth,
+      maxDepth = maxDepth,
+      rng = rng
+    )
     new RegressionTreeTrainingResult(rootTrainingNode, encoders)
   }
 
 }
 
 class RegressionTreeTrainingResult(
-    rootTrainingNode: TrainingNode[AnyVal, Double],
+    rootTrainingNode: TrainingNode[Double],
     encoders: Seq[Option[CategoricalEncoder[Any]]]
 ) extends TrainingResult {
-  lazy val model = new RegressionTree(rootTrainingNode.getNode(), encoders)
-  lazy val importance = rootTrainingNode.getFeatureImportance()
+  lazy val model = new RegressionTree(rootTrainingNode.modelNode, encoders)
+  lazy val importance = rootTrainingNode.featureImportance
   private lazy val importanceNormalized = {
     if (Math.abs(importance.sum) > 0) {
       importance.map(_ / importance.sum)
diff --git a/src/main/scala/io/citrine/lolo/trees/splits/BoltzmannSplitter.scala b/src/main/scala/io/citrine/lolo/trees/splits/BoltzmannSplitter.scala
index b4d370bd..81f62db2 100644
--- a/src/main/scala/io/citrine/lolo/trees/splits/BoltzmannSplitter.scala
+++ b/src/main/scala/io/citrine/lolo/trees/splits/BoltzmannSplitter.scala
@@ -58,7 +58,7 @@ case class BoltzmannSplitter(temperature: Double) extends Splitter[Double] {
 
     // Don't split if there is no impurity to reduce
     if (initialVariance == 0) {
-      return (new NoSplit(), 0.0)
+      return (NoSplit(), 0.0)
     }
     val beta = 1.0 / (temperature * initialVariance)
 
@@ -78,7 +78,7 @@ case class BoltzmannSplitter(temperature: Double) extends Splitter[Double] {
 
     // If we couldn't find a split, then return NoSplit with no variance reduction
     if (possibleSplits.isEmpty) {
-      return (new NoSplit(), 0.0)
+      return (NoSplit(), 0.0)
     }
 
     // Re-based the probabilities, such that the largest probability is order-1.0
@@ -174,7 +174,7 @@ object BoltzmannSplitter {
       case (score, pivot, variance) =>
         cumSum = cumSum + Math.exp(score - base)
         if (draw < cumSum) {
-          return Some(SplitterResult(new RealSplit(index, pivot), variance, totalScore, base))
+          return Some(SplitterResult(RealSplit(index, pivot), variance, totalScore, base))
         }
     }
     // This should never be hit; it would mean there's a bug in the logic above ^^
@@ -250,7 +250,7 @@ object BoltzmannSplitter {
       case (score, includeSet, variance) =>
         cumSum = cumSum + Math.exp(score - base)
         if (draw < cumSum) {
-          return Some(SplitterResult(new CategoricalSplit(index, includeSet), variance, totalScore, base))
+          return Some(SplitterResult(CategoricalSplit(index, includeSet), variance, totalScore, base))
         }
     }
     // This should never be hit; it would mean there's a bug in the logic above ^^
diff --git a/src/main/scala/io/citrine/lolo/trees/splits/ClassificationSplitter.scala b/src/main/scala/io/citrine/lolo/trees/splits/ClassificationSplitter.scala
index 41b0e1c3..f8e1624e 100644
--- a/src/main/scala/io/citrine/lolo/trees/splits/ClassificationSplitter.scala
+++ b/src/main/scala/io/citrine/lolo/trees/splits/ClassificationSplitter.scala
@@ -25,7 +25,7 @@ case class ClassificationSplitter(randomizedPivotLocation: Boolean = false) exte
       minInstances: Int,
       rng: Random
   ): (Split, Double) = {
-    var bestSplit: Split = new NoSplit()
+    var bestSplit: Split = NoSplit()
     var bestImpurity = Double.MaxValue
 
     /* Pre-compute these for the variance calculation */
@@ -51,7 +51,7 @@ case class ClassificationSplitter(randomizedPivotLocation: Boolean = false) exte
       }
     }
     if (bestImpurity == Double.MaxValue) {
-      (new NoSplit(), 0.0)
+      (NoSplit(), 0.0)
     } else {
       val deltaImpurity = initialImpurity - bestImpurity
       (bestSplit, deltaImpurity)
@@ -101,7 +101,7 @@ case class ClassificationSplitter(randomizedPivotLocation: Boolean = false) exte
         bestSet = orderedNames.slice(0, j + 1).toSet
       }
     }
-    (new CategoricalSplit(index, new scala.collection.mutable.BitSet() ++ bestSet.map(_.toInt)), bestPurity)
+    (CategoricalSplit(index, new scala.collection.mutable.BitSet() ++ bestSet.map(_.toInt)), bestPurity)
   }
 
 }
diff --git a/src/main/scala/io/citrine/lolo/trees/splits/ExtraRandomClassificationSplitter.scala b/src/main/scala/io/citrine/lolo/trees/splits/ExtraRandomClassificationSplitter.scala
index 2a34c0f2..37b1a3fe 100644
--- a/src/main/scala/io/citrine/lolo/trees/splits/ExtraRandomClassificationSplitter.scala
+++ b/src/main/scala/io/citrine/lolo/trees/splits/ExtraRandomClassificationSplitter.scala
@@ -29,7 +29,7 @@ case class ExtraRandomClassificationSplitter() extends Splitter[Char] {
 
     val calculator = GiniCalculator.build(data.map { p => (p._2, p._3) })
     val initialVariance = calculator.getImpurity
-    var bestSplit: Split = new NoSplit()
+    var bestSplit: Split = NoSplit()
     var bestVariance = Double.MaxValue
 
     val rep = data.head
@@ -51,7 +51,7 @@ case class ExtraRandomClassificationSplitter() extends Splitter[Char] {
       }
     }
     if (bestVariance >= initialVariance) {
-      (new NoSplit(), 0.0)
+      (NoSplit(), 0.0)
     } else {
       val deltaImpurity = initialVariance - bestVariance
       (bestSplit, deltaImpurity)
@@ -80,7 +80,7 @@ case class ExtraRandomClassificationSplitter() extends Splitter[Char] {
     val upperBound = axis.max
     val pivot = lowerBound + (upperBound - lowerBound) * rng.nextDouble()
 
-    val split = new RealSplit(index, pivot)
+    val split = RealSplit(index, pivot)
 
     /* Move the data from the right to the left partition one value at a time */
     calculator.reset()
@@ -122,12 +122,12 @@ case class ExtraRandomClassificationSplitter() extends Splitter[Char] {
     val orderedNames: Seq[Char] = categoryAverages.toSeq.sortBy(_._2).map(_._1)
     if (orderedNames.size == 1) {
       calculator.reset()
-      return (new NoSplit(), calculator.getImpurity)
+      return (NoSplit(), calculator.getImpurity)
     }
 
     val pivot = rng.nextInt(orderedNames.size - 1)
     val bestSet = orderedNames.slice(0, pivot + 1).toSet
-    val split = new CategoricalSplit(index, new scala.collection.mutable.BitSet() ++ bestSet.map(_.toInt))
+    val split = CategoricalSplit(index, new scala.collection.mutable.BitSet() ++ bestSet.map(_.toInt))
 
     /* Move the data from the right to the left partition one value at a time */
     calculator.reset()
diff --git a/src/main/scala/io/citrine/lolo/trees/splits/ExtraRandomRegressionSplitter.scala b/src/main/scala/io/citrine/lolo/trees/splits/ExtraRandomRegressionSplitter.scala
index 29b08740..a813594a 100644
--- a/src/main/scala/io/citrine/lolo/trees/splits/ExtraRandomRegressionSplitter.scala
+++ b/src/main/scala/io/citrine/lolo/trees/splits/ExtraRandomRegressionSplitter.scala
@@ -29,7 +29,7 @@ case class ExtraRandomRegressionSplitter() extends Splitter[Double] {
 
     val calculator = VarianceCalculator.build(data.map(_._2), data.map(_._3))
     val initialVariance = calculator.getImpurity
-    var bestSplit: Split = new NoSplit()
+    var bestSplit: Split = NoSplit()
     var bestVariance = Double.MaxValue
 
     val rep = data.head
@@ -51,7 +51,7 @@ case class ExtraRandomRegressionSplitter() extends Splitter[Double] {
       }
     }
     if (bestVariance >= initialVariance) {
-      (new NoSplit(), 0.0)
+      (NoSplit(), 0.0)
     } else {
       val deltaImpurity = initialVariance - bestVariance
       (bestSplit, deltaImpurity)
@@ -80,7 +80,7 @@ case class ExtraRandomRegressionSplitter() extends Splitter[Double] {
     val upperBound = axis.max
     val pivot = lowerBound + (upperBound - lowerBound) * rng.nextDouble()
 
-    val split = new RealSplit(index, pivot)
+    val split = RealSplit(index, pivot)
 
     /* Move the data from the right to the left partition one value at a time */
     calculator.reset()
@@ -122,12 +122,12 @@ case class ExtraRandomRegressionSplitter() extends Splitter[Double] {
     val orderedNames: Seq[Char] = categoryAverages.toSeq.sortBy(_._2).map(_._1)
     if (orderedNames.size == 1) {
       calculator.reset()
-      return (new NoSplit(), calculator.getImpurity)
+      return (NoSplit(), calculator.getImpurity)
     }
 
     val pivot = rng.nextInt(orderedNames.size - 1)
     val bestSet = orderedNames.slice(0, pivot + 1).toSet
-    val split = new CategoricalSplit(index, new scala.collection.mutable.BitSet() ++ bestSet.map(_.toInt))
+    val split = CategoricalSplit(index, new scala.collection.mutable.BitSet() ++ bestSet.map(_.toInt))
 
     /* Move the data from the right to the left partition one value at a time */
     calculator.reset()
diff --git a/src/main/scala/io/citrine/lolo/trees/splits/MultiTaskSplitter.scala b/src/main/scala/io/citrine/lolo/trees/splits/MultiTaskSplitter.scala
index 5efee046..000c20af 100644
--- a/src/main/scala/io/citrine/lolo/trees/splits/MultiTaskSplitter.scala
+++ b/src/main/scala/io/citrine/lolo/trees/splits/MultiTaskSplitter.scala
@@ -3,7 +3,7 @@ package io.citrine.lolo.trees.splits
 import io.citrine.random.Random
 import io.citrine.lolo.trees.impurity.MultiImpurityCalculator
 
-case class MultiTaskSplitter(randomizePivotLocation: Boolean = false) extends Splitter[Array[AnyVal]] {
+case class MultiTaskSplitter(randomizePivotLocation: Boolean = false) extends Splitter[Seq[AnyVal]] {
 
   /**
     * Get the best split, considering numFeature random features (w/o replacement)
@@ -15,38 +15,38 @@ case class MultiTaskSplitter(randomizePivotLocation: Boolean = false) extends Sp
     * @return a split object that optimally divides data
     */
   def getBestSplit(
-      data: Seq[(Vector[AnyVal], Array[AnyVal], Double)],
+      data: Seq[(Vector[AnyVal], Seq[AnyVal], Double)],
       numFeatures: Int,
       minInstances: Int,
       rng: Random
   ): (Split, Double) = {
-    var bestSplit: Split = new NoSplit()
+    var bestSplit: Split = NoSplit()
     var bestImpurity = Double.MaxValue
     val calculator = MultiImpurityCalculator.build(data.map(_._2), data.map(_._3))
     val initialImpurity = calculator.getImpurity
 
-    /* Pre-compute these for the variance calculation */
+    // Pre-compute these for the variance calculation.
     val rep = data.head
 
-    /* Try every feature index */
+    // Try every feature index.
     val featureIndices: Seq[Int] = rep._1.indices
     rng.shuffle(featureIndices).take(numFeatures).foreach { index =>
-      /* Use different spliters for each type */
+      // Use different splitters for each type.
       val (possibleSplit, possibleImpurity) = rep._1(index) match {
         case _: Double =>
-          Splitter.getBestRealSplit[Array[AnyVal]](data, calculator, index, minInstances, randomizePivotLocation, rng)
+          Splitter.getBestRealSplit[Seq[AnyVal]](data, calculator, index, minInstances, randomizePivotLocation, rng)
         case _: Char => getBestCategoricalSplit(data, calculator, index, minInstances)
         case _: Any  => throw new IllegalArgumentException("Trying to split unknown feature type")
       }
 
-      /* Keep track of the best split */
+      // Keep track of the best split.
       if (possibleImpurity < bestImpurity) {
         bestImpurity = possibleImpurity
         bestSplit = possibleSplit
       }
     }
     if (bestImpurity == Double.MaxValue) {
-      (new NoSplit(), 0.0)
+      (NoSplit(), 0.0)
     } else {
       val deltaImpurity = initialImpurity - bestImpurity
       (bestSplit, deltaImpurity)
@@ -54,25 +54,26 @@ case class MultiTaskSplitter(randomizePivotLocation: Boolean = false) extends Sp
   }
 
   /**
-    * Get find the best categorical splitter.
+    * Get the best categorical split.
     *
-    * @param data        to split
-    * @param totalWeight Pre-computed data.map(d => d._3).sum
-    * @param index       of the feature to split on
+    * @param data         to split
+    * @param calculator   calculates the impurity of training data subsets
+    * @param index        of the feature to split on
+    * @param minInstances minimum instances permitted in a post-split partition
     * @return the best split of this feature
     */
   def getBestCategoricalSplit(
-      data: Seq[(Vector[AnyVal], Array[AnyVal], Double)],
+      data: Seq[(Vector[AnyVal], Seq[AnyVal], Double)],
       calculator: MultiImpurityCalculator,
       index: Int,
-      minCount: Int
+      minInstances: Int
   ): (Split, Double) = {
-    /* Extract the features at the index */
-    val thinData: Seq[(Char, Array[AnyVal], Double)] =
+    // Extract the features at the index.
+    val thinData: Seq[(Char, Seq[AnyVal], Double)] =
       data.map(dat => (dat._1(index).asInstanceOf[Char], dat._2, dat._3))
     val totalWeight = thinData.map(_._3).sum
 
-    /* Group the data by categorical feature and compute the weighted sum and sum of the weights for each */
+    // Group the data by categorical feature and compute the weighted sum and sum of the weights for each.
     val groupedData: Map[Char, (Double, Double, Double)] = thinData
       .groupBy(_._1)
       .view
@@ -82,7 +83,7 @@ case class MultiTaskSplitter(randomizePivotLocation: Boolean = false) extends Sp
     /* Make sure there is more than one member for most of the classes */
     val nonTrivial: Double = groupedData.filter(_._2._3 > 1).map(_._2._2).sum
     if (nonTrivial / totalWeight < 0.25) {
-      return (new NoSplit, Double.MaxValue)
+      return (NoSplit(), Double.MaxValue)
     }
 
     /* Create an orderd list of the categories by average label */
@@ -91,14 +92,14 @@ case class MultiTaskSplitter(randomizePivotLocation: Boolean = false) extends Sp
     /* Add the categories one at a time in order of their average label */
     var leftNum = 0
     calculator.reset()
-    val pivots = (0 until orderedNames.size).flatMap { j =>
-      thinData.filter(r => orderedNames(j) == r._1).map { r =>
+    val pivots = orderedNames.indices.flatMap { j =>
+      thinData.filter(r => orderedNames(j) == r._1).foreach { r =>
         calculator.add(r._2, r._3)
         leftNum = leftNum + 1
       }
       val totalImpurity = calculator.getImpurity
 
-      if (leftNum >= minCount && thinData.size - leftNum >= minCount) {
+      if (leftNum >= minInstances && thinData.size - leftNum >= minInstances) {
         val set = orderedNames.take(j + 1).toSet
         Some(set, totalImpurity)
       } else {
@@ -106,10 +107,10 @@ case class MultiTaskSplitter(randomizePivotLocation: Boolean = false) extends Sp
       }
     }
     if (pivots.isEmpty) {
-      (new NoSplit, Double.MaxValue)
+      (NoSplit(), Double.MaxValue)
     } else {
       val best = pivots.minBy(_._2)
-      (new CategoricalSplit(index, new scala.collection.mutable.BitSet() ++ best._1.map(_.toInt)), best._2)
+      (CategoricalSplit(index, new scala.collection.mutable.BitSet() ++ best._1.map(_.toInt)), best._2)
     }
   }
 
@@ -119,7 +120,7 @@ case class MultiTaskSplitter(randomizePivotLocation: Boolean = false) extends Sp
     * @param labels is a seq of (Array of multiple labels, single weight)
     * @return the impurity, which is in [0, number of labels * sum of weights]
     */
-  def computeImpurity(labels: Seq[(Array[AnyVal], Double)]): Double = {
+  def computeImpurity(labels: Seq[(Seq[AnyVal], Double)]): Double = {
     val calculator = MultiImpurityCalculator.build(labels.map(_._1), labels.map(_._2))
     calculator.getImpurity
   }
diff --git a/src/main/scala/io/citrine/lolo/trees/splits/RegressionSplitter.scala b/src/main/scala/io/citrine/lolo/trees/splits/RegressionSplitter.scala
index 19f4883f..da98baf4 100644
--- a/src/main/scala/io/citrine/lolo/trees/splits/RegressionSplitter.scala
+++ b/src/main/scala/io/citrine/lolo/trees/splits/RegressionSplitter.scala
@@ -40,7 +40,7 @@ case class RegressionSplitter(randomizePivotLocation: Boolean = false) extends S
 
     val calculator = VarianceCalculator.build(data.map(_._2), data.map(_._3))
     val initialVariance = calculator.getImpurity
-    var bestSplit: Split = new NoSplit()
+    var bestSplit: Split = NoSplit()
     var bestVariance = Double.MaxValue
 
     val rep = data.head
@@ -63,7 +63,7 @@ case class RegressionSplitter(randomizePivotLocation: Boolean = false) extends S
       }
     }
     if (bestVariance == Double.MaxValue) {
-      (new NoSplit(), 0.0)
+      (NoSplit(), 0.0)
     } else {
       val deltaImpurity = initialVariance - bestVariance
       (bestSplit, deltaImpurity)
@@ -93,7 +93,7 @@ case class RegressionSplitter(randomizePivotLocation: Boolean = false) extends S
     /* Make sure there is more than one member for most of the classes */
     val nonTrivial: Double = groupedData.filter(_._2._3 > 1).map(_._2._2).sum
     if (nonTrivial / totalWeight < 0.5) {
-      return (new CategoricalSplit(index, BitSet()), Double.MaxValue)
+      return (CategoricalSplit(index, BitSet()), Double.MaxValue)
     }
 
     /* Compute the average label for each categorical value */
@@ -121,7 +121,7 @@ case class RegressionSplitter(randomizePivotLocation: Boolean = false) extends S
         bestSet = orderedNames.slice(0, j + 1).toSet
       }
     }
-    (new CategoricalSplit(index, new scala.collection.mutable.BitSet() ++ bestSet.map(_.toInt)), bestVariance)
+    (CategoricalSplit(index, new scala.collection.mutable.BitSet() ++ bestSet.map(_.toInt)), bestVariance)
   }
 
 }
diff --git a/src/main/scala/io/citrine/lolo/trees/splits/Split.scala b/src/main/scala/io/citrine/lolo/trees/splits/Split.scala
index c9f74062..0c2a5b2b 100644
--- a/src/main/scala/io/citrine/lolo/trees/splits/Split.scala
+++ b/src/main/scala/io/citrine/lolo/trees/splits/Split.scala
@@ -5,7 +5,7 @@ import scala.collection.BitSet
 /**
   * Splits are used by decision trees to partition the input space
   */
-abstract trait Split extends Serializable {
+sealed trait Split extends Serializable {
 
   /**
     * Take the left branch in the binary split?
@@ -20,13 +20,13 @@ abstract trait Split extends Serializable {
     *
     * @return index of the input vector used by this split
     */
-  def getIndex(): Int
+  def index: Int
 }
 
 /**
   * If no split was found
   */
-class NoSplit extends Split {
+case class NoSplit() extends Split {
 
   /**
     * Take the left branch in the binary split?
@@ -41,7 +41,7 @@ class NoSplit extends Split {
     *
     * @return index of the input vector used by this split
     */
-  override def getIndex(): Int = -1
+  override def index: Int = -1
 }
 
 /**
@@ -50,7 +50,7 @@ class NoSplit extends Split {
   * @param index position of the real value to inspect
   * @param pivot value at or below which to take the left split
   */
-class RealSplit(index: Int, pivot: Double) extends Split {
+case class RealSplit(index: Int, pivot: Double) extends Split {
 
   /**
     * If the value is at or less than the pivot, turn left
@@ -68,19 +68,12 @@ class RealSplit(index: Int, pivot: Double) extends Split {
     }
   }
 
-  /**
-    * ${inherit_doc}
-    *
-    * @return index of the input vector used by this split
-    */
-  override def getIndex: Int = index
-
   /**
     * Pretty print
     *
     * @return debug string
     */
-  override def toString: String = s"Split index ${index} @ ${pivot}"
+  override def toString: String = s"Split index $index @ $pivot"
 }
 
 /**
@@ -89,7 +82,7 @@ class RealSplit(index: Int, pivot: Double) extends Split {
   * @param index      of the categorical feature
   * @param includeSet set of values that turn left
   */
-class CategoricalSplit(index: Int, includeSet: BitSet) extends Split {
+case class CategoricalSplit(index: Int, includeSet: BitSet) extends Split {
 
   /**
     * If the value at the index position is in the set, turn left
@@ -100,11 +93,4 @@ class CategoricalSplit(index: Int, includeSet: BitSet) extends Split {
   override def turnLeft(input: Vector[AnyVal]): Boolean = {
     includeSet.contains(input(index).asInstanceOf[Char].toInt)
   }
-
-  /**
-    * ${inherit_doc}
-    *
-    * @return index of the input vector used by this split
-    */
-  override def getIndex: Int = index
 }
diff --git a/src/main/scala/io/citrine/lolo/trees/splits/Splitter.scala b/src/main/scala/io/citrine/lolo/trees/splits/Splitter.scala
index 7aa0369e..3b07d693 100644
--- a/src/main/scala/io/citrine/lolo/trees/splits/Splitter.scala
+++ b/src/main/scala/io/citrine/lolo/trees/splits/Splitter.scala
@@ -67,6 +67,6 @@ object Splitter {
         }
       }
     }
-    (new RealSplit(index, bestPivot), bestImpurity)
+    (RealSplit(index, bestPivot), bestImpurity)
   }
 }


Test patch:
diff --git a/src/test/scala/io/citrine/lolo/learners/RandomForestTest.scala b/src/test/scala/io/citrine/lolo/learners/RandomForestTest.scala
index 85145cbf..1566a9d3 100644
--- a/src/test/scala/io/citrine/lolo/learners/RandomForestTest.scala
+++ b/src/test/scala/io/citrine/lolo/learners/RandomForestTest.scala
@@ -239,7 +239,7 @@ class RandomForestTest extends SeedRandomMixIn {
       trainingData: Seq[(Vector[Any], Double)],
       evalLocation: Vector[Any],
       expected: Vector[Double],
-      rtol: Double = 5e-2
+      rtol: Double = 0.1
   ): Unit = {
     val actual = RandomForest().train(trainingData, rng = rng).getModel().shapley(evalLocation) match {
       case None => fail("Unexpected None returned by shapley.")
@@ -310,14 +310,3 @@ class RandomForestTest extends SeedRandomMixIn {
     shapleyCompare((1 to 8).map { _ => trainingData4 }.flatten, Vector.fill[Any](5)(1.0), expected4)
   }
 }
-
-object RandomForestTest {
-  def main(argv: Array[String]): Unit = {
-    new RandomForestTest()
-      .testClassificationForest()
-    new RandomForestTest()
-      .testClassificationForestUnbiased()
-    new RandomForestTest()
-      .testShapley()
-  }
-}
diff --git a/src/test/scala/io/citrine/lolo/trees/multitask/MultiTaskTreeTest.scala b/src/test/scala/io/citrine/lolo/trees/multitask/MultiTaskTreeTest.scala
index d414071b..f71b1e34 100644
--- a/src/test/scala/io/citrine/lolo/trees/multitask/MultiTaskTreeTest.scala
+++ b/src/test/scala/io/citrine/lolo/trees/multitask/MultiTaskTreeTest.scala
@@ -12,7 +12,7 @@ import org.scalatest.Assertions._
 @Test
 class MultiTaskTreeTest extends SeedRandomMixIn {
 
-  /* Setup some data */
+  // Set up some data.
   val raw: Seq[(Vector[Double], Double)] =
     TestUtils.generateTrainingData(1024, 12, noise = 0.1, function = Friedman.friedmanSilverman, rng = rng)
   val (inputs: Seq[Vector[Double]], realLabel: Seq[Double]) = raw.unzip
@@ -25,7 +25,7 @@ class MultiTaskTreeTest extends SeedRandomMixIn {
   @Test
   def testTwoLabels(): Unit = {
     val learner = MultiTaskTreeLearner()
-    val models = learner.train(inputs.zip(labels)).getModels()
+    val models = learner.train(inputs.zip(labels), rng = rng).getModels()
     assert(models.size == 2)
     assert(models.head.isInstanceOf[RegressionTree])
     assert(models.last.isInstanceOf[ClassificationTree])
@@ -74,11 +74,11 @@ class MultiTaskTreeTest extends SeedRandomMixIn {
     val sparseLabels = Vector(realLabel, sparseCat).transpose
 
     val learner = MultiTaskTreeLearner()
-    val models = learner.train(inputs.zip(sparseLabels)).getModels()
+    val models = learner.train(inputs.zip(sparseLabels), rng = rng).getModels()
     val catResults = models.last.transform(inputs).getExpected().asInstanceOf[Seq[Boolean]]
 
     val reference = ClassificationTreeLearner()
-      .train(inputs.zip(sparseCat).filterNot(_._2 == null))
+      .train(inputs.zip(sparseCat).filterNot(_._2 == null), rng = rng)
       .getModel()
       .transform(inputs)
       .getExpected()
diff --git a/src/test/scala/io/citrine/lolo/trees/regression/RegressionTreeTest.scala b/src/test/scala/io/citrine/lolo/trees/regression/RegressionTreeTest.scala
index 8968b137..9dca5b4e 100644
--- a/src/test/scala/io/citrine/lolo/trees/regression/RegressionTreeTest.scala
+++ b/src/test/scala/io/citrine/lolo/trees/regression/RegressionTreeTest.scala
@@ -40,11 +40,11 @@ class RegressionTreeTest extends SeedRandomMixIn {
     /* We should be able to memorize the inputs */
     val output = DT.transform(trainingData.map(_._1))
     trainingData.zip(output.getExpected()).foreach {
-      case ((x, a), p) =>
+      case ((_, a), p) =>
         assert(Math.abs(a - p) < 1.0e-9)
     }
     assert(output.getGradient().isEmpty)
-    output.getDepth().foreach(d => assert(d > 3 && d < 9, s"Depth is ${d}"))
+    output.getDepth().foreach(d => assert(d > 3 && d < 9, s"Depth is $d"))
   }
 
   /**
@@ -65,7 +65,7 @@ class RegressionTreeTest extends SeedRandomMixIn {
         assert(Math.abs(a - p) < 1.0e-9)
     }
     assert(output.getGradient().isEmpty)
-    output.getDepth().foreach(d => assert(d > 3 && d < 9, s"Depth is ${d}"))
+    output.getDepth().foreach(d => assert(d > 3 && d < 9, s"Depth is $d"))
   }
 
   /**
@@ -162,7 +162,7 @@ class RegressionTreeTest extends SeedRandomMixIn {
         assert(Math.abs(a - p) < 1.0e-9)
     }
     assert(output.getGradient().isDefined)
-    output.getDepth().foreach(d => assert(d > 4 && d < 18, s"Depth is ${d}"))
+    output.getDepth().foreach(d => assert(d > 4 && d < 18, s"Depth is $d"))
 
     /* The first feature should be the most important */
     val importances = DTMeta.getFeatureImportance().get
@@ -346,16 +346,3 @@ class RegressionTreeTest extends SeedRandomMixIn {
     RegressionTreeLearner(numFeatures = 1).train(trainingData4).getModel().shapley(Vector.fill[Any](5)(0.0), Set())
   }
 }
-
-/** Companion driver */
-object RegressionTreeTest {
-
-  /**
-    * Test driver
-    *
-    * @param argv args
-    */
-  def main(argv: Array[String]): Unit = {
-    new RegressionTreeTest().testSimpleBoltzmannTree()
-  }
-}
diff --git a/src/test/scala/io/citrine/lolo/trees/splits/ExtraRandomRegressionSplitterTest.scala b/src/test/scala/io/citrine/lolo/trees/splits/ExtraRandomRegressionSplitterTest.scala
index f1c3dc66..e80ea270 100644
--- a/src/test/scala/io/citrine/lolo/trees/splits/ExtraRandomRegressionSplitterTest.scala
+++ b/src/test/scala/io/citrine/lolo/trees/splits/ExtraRandomRegressionSplitterTest.scala
@@ -127,7 +127,7 @@ class ExtraRandomRegressionSplitterTest extends SeedRandomMixIn {
 
             // Finally, we can ensure that the index on which we split is correct...
             val testCaveatMessage = "NOTE: this test may inaccurately fail due to changes in the sequence of rng calls."
-            assert(bestSplit._1.getIndex() == indexOfBest, s"Incorrect index of best split. $testCaveatMessage")
+            assert(bestSplit._1.index == indexOfBest, s"Incorrect index of best split. $testCaveatMessage")
             // Do a sanity check about the directionality of turnLeft to ensure this test is valid.
             assert(
               bestSplit._1.turnLeft(Vector.fill(numFeatures)(xTrain.flatten.min)),
diff --git a/src/test/scala/io/citrine/lolo/trees/splits/MultiTaskSplitterTest.scala b/src/test/scala/io/citrine/lolo/trees/splits/MultiTaskSplitterTest.scala
index 3b8c5409..04ceb67e 100644
--- a/src/test/scala/io/citrine/lolo/trees/splits/MultiTaskSplitterTest.scala
+++ b/src/test/scala/io/citrine/lolo/trees/splits/MultiTaskSplitterTest.scala
@@ -16,10 +16,10 @@ class MultiTaskSplitterTest extends SeedRandomMixIn {
   @Test
   def testBestRealSplit(): Unit = {
     val data = Seq(
-      (Vector(0.0), Array[AnyVal](0.0, 1.toChar), 1.0),
-      (Vector(1.0), Array[AnyVal](1.0, 2.toChar), 1.0),
-      (Vector(2.0), Array[AnyVal](2.0, 2.toChar), 1.0),
-      (Vector(100.0), Array[AnyVal](200.0, 3.toChar), 0.0)
+      (Vector(0.0), Vector[AnyVal](0.0, 1.toChar), 1.0),
+      (Vector(1.0), Vector[AnyVal](1.0, 2.toChar), 1.0),
+      (Vector(2.0), Vector[AnyVal](2.0, 2.toChar), 1.0),
+      (Vector(100.0), Vector[AnyVal](200.0, 3.toChar), 0.0)
     )
     val calculator = MultiImpurityCalculator.build(data.map(_._2), data.map(_._3))
 
@@ -36,10 +36,10 @@ class MultiTaskSplitterTest extends SeedRandomMixIn {
   @Test
   def testBestCategoricalSplit(): Unit = {
     val data = Seq(
-      (Vector(1.toChar), Array[AnyVal](0.0, 1.toChar), 1.0),
-      (Vector(2.toChar), Array[AnyVal](1.0, 2.toChar), 1.0),
-      (Vector(2.toChar), Array[AnyVal](2.0, 2.toChar), 1.0),
-      (Vector(3.toChar), Array[AnyVal](200.0, 3.toChar), 0.0)
+      (Vector(1.toChar), Vector[AnyVal](0.0, 1.toChar), 1.0),
+      (Vector(2.toChar), Vector[AnyVal](1.0, 2.toChar), 1.0),
+      (Vector(2.toChar), Vector[AnyVal](2.0, 2.toChar), 1.0),
+      (Vector(3.toChar), Vector[AnyVal](200.0, 3.toChar), 0.0)
     )
     val calculator = MultiImpurityCalculator.build(data.map(_._2), data.map(_._3))
 
@@ -60,10 +60,10 @@ class MultiTaskSplitterTest extends SeedRandomMixIn {
       Vector(0.4566490222691504, 0.21730793346810628, 0.08289198941227605),
       Vector(0.0025661310641983404, 0.21416453420489545, 0.15783628254609094)
     )
-    val labels = Vector(Array(1.238, null), Array(1.180, null))
+    val labels = Vector(Vector(1.238, null), Vector(1.180, null))
     val weights = Vector(1.0, 1.0)
     val data = inputs.indices.map { i =>
-      (inputs(i), labels(i).asInstanceOf[Array[AnyVal]], weights(i))
+      (inputs(i), labels(i).asInstanceOf[Vector[AnyVal]], weights(i))
     }
     val splitter = MultiTaskSplitter()
     val (pivot, _) = splitter.getBestSplit(data, data.head._1.size, 1, rng = rng)
diff --git a/src/test/scala/io/citrine/lolo/trees/splits/SplitTest.scala b/src/test/scala/io/citrine/lolo/trees/splits/SplitTest.scala
index d32be0cb..927b8ace 100644
--- a/src/test/scala/io/citrine/lolo/trees/splits/SplitTest.scala
+++ b/src/test/scala/io/citrine/lolo/trees/splits/SplitTest.scala
@@ -9,7 +9,7 @@ class SplitTest {
     */
   @Test
   def testSplitNaN(): Unit = {
-    val split = new RealSplit(0, Double.NaN)
+    val split = RealSplit(0, Double.NaN)
     assert(
       split.turnLeft(Vector(0.0)) != split.turnLeft(Vector(Double.NaN)),
       "NaN and non-Nan values are split in the same direction."