extratree - always split on missing if missing data is present

pityka · Feb 9, 2024 · 34e506e · 34e506e
1 parent bc57d94
commit 34e506e
Show file tree

Hide file tree

Showing 3 changed files with 143 additions and 156 deletions.
diff --git a/extratrees/src/main/scala/lamp/forest/extratrees.scala b/extratrees/src/main/scala/lamp/forest/extratrees.scala
@@ -11,15 +11,16 @@ object ClassificationLeaf {
   import upickle.default.{ReadWriter => RW, macroRW}
   implicit val rw: RW[ClassificationLeaf] = macroRW
 }
+/** Left bin contains elements < than cutpoint XOR missing elements  if cutpoint is empty
+ */
 case class ClassificationNonLeaf(
-    left: ClassificationTree,
+    left: ClassificationTree, 
     right: ClassificationTree,
     splitFeature: Int,
-    cutpoint: Double,
-    splitMissingIsLess: Boolean
+    cutpoint: Option[Double], 
 ) extends ClassificationTree {
   override def toString =
-    s"ClassificationTree(left=$left,right=$right,splitFeatures=$splitFeature,cutpoint=$cutpoint,splitMissingIsLess=$splitMissingIsLess)"
+    s"ClassificationTree(left=$left,right=$right,splitFeatures=$splitFeature,cutpoint=$cutpoint)"
 }
 object ClassificationNonLeaf {
   import upickle.default.{ReadWriter => RW, macroRW}
@@ -37,15 +38,17 @@ case class RegressionLeaf(targetMean: Double) extends RegressionTree {
   override def toString = s"RegressionLeaf(targetMean=$targetMean)"
 }
 
+/** Left bin contains elements < than cutpoint XOR missing elements  if cutpoint is empty
+ */
 case class RegressionNonLeaf(
     left: RegressionTree,
     right: RegressionTree,
     splitFeature: Int,
-    cutpoint: Double,
-    splitMissingIsLess: Boolean
+    cutpoint: Option[Double],
+
 ) extends RegressionTree {
   override def toString =
-    s"RegressionNonLeaf(left=$left,right=$right,splitFeatures=$splitFeature,cutpoint=$cutpoint,splitMissingIsLess=$splitMissingIsLess)"
+    s"RegressionNonLeaf(left=$left,right=$right,splitFeatures=$splitFeature,cutpoint=$cutpoint)"
 }
 object RegressionLeaf {
   import upickle.default.{ReadWriter => RW, macroRW}