[WIP] GH-6769: multinomial dt yuliia

h2o-algos/src/main/java/hex/tree/dt/DT.java

@@ -8,7 +8,6 @@
 import hex.tree.dt.binning.Histogram;
 import hex.tree.dt.mrtasks.GetClassCountsMRTask;
 import hex.tree.dt.mrtasks.ScoreDTTask;
-import org.apache.commons.math3.util.Precision;
 import org.apache.log4j.Logger;
 import water.DKV;
 import water.exceptions.H2OModelBuilderIllegalArgumentException;
@@ -19,7 +18,7 @@
 import java.util.stream.Collectors;
 import java.util.stream.IntStream;
 
-import static hex.tree.dt.binning.SplitStatistics.entropyBinarySplit;
+import static hex.tree.dt.binning.SplitStatistics.entropyBinarySplitMultinomial;
 
 /**
  * Decision Tree
@@ -108,8 +107,8 @@ private AbstractSplittingRule findBestSplit(Histogram histogram) {
 
     private AbstractSplittingRule findBestSplitForFeature(Histogram histogram, int featureIndex) {
         return (_train.vec(featureIndex).isNumeric()
-                ? histogram.calculateSplitStatisticsForNumericFeature(featureIndex)
-                : histogram.calculateSplitStatisticsForCategoricalFeature(featureIndex))
+                ? histogram.calculateSplitStatisticsForNumericFeature(featureIndex, _nclass)
+                : histogram.calculateSplitStatisticsForCategoricalFeature(featureIndex, _nclass))
                 .stream()
                 // todo - consider setting min count of samples in bin instead of filtering splits
                 .filter(binStatistics -> ((binStatistics._leftCount >= _min_rows)
@@ -128,6 +127,7 @@ private AbstractSplittingRule findBestSplitForFeature(Histogram histogram, int f
 
 
     private static double calculateCriterionOfSplit(SplitStatistics binStatistics) {
+        // if(binStatistics.() == 2) // todo - fix bin statistics first, they are binomial-only now
         return binStatistics.binaryEntropy();
     }
 
@@ -139,7 +139,7 @@ private static double calculateCriterionOfSplit(SplitStatistics binStatistics) {
      */
     private int selectDecisionValue(int[] countsByClass) {
         if (_nclass == 1) {
-            return countsByClass[0];
+            return 0;
         }
         int currentMaxClass = 0;
         int currentMax = countsByClass[currentMaxClass];
@@ -205,11 +205,7 @@ public void buildNextNode(Queue<DataFeaturesLimits> limitsQueue, int nodeIndex)
         // compute node depth
         int nodeDepth = (int) Math.floor(MathUtils.log2(nodeIndex + 1));
         // stop building from this node, the node will be a leaf
-        if ((nodeDepth >= _parms._max_depth)
-                || (countsByClass[0] <= _min_rows)
-                || (countsByClass[1] <= _min_rows)
-//                || zeroRatio > 0.999 || zeroRatio < 0.001
-        ) {
+        if ((nodeDepth >= _parms._max_depth) || Arrays.stream(countsByClass).anyMatch(c -> c <= _min_rows)) {
             // add imaginary left and right children to imitate valid tree structure
             // left child
             limitsQueue.add(null);
@@ -219,10 +215,10 @@ public void buildNextNode(Queue<DataFeaturesLimits> limitsQueue, int nodeIndex)
             return;
         }
 
-        Histogram histogram = new Histogram(_train, actualLimits, BinningStrategy.EQUAL_WIDTH/*, minNumSamplesInBin - todo consider*/);
+        Histogram histogram = new Histogram(_train, actualLimits, BinningStrategy.EQUAL_WIDTH, _nclass/*, minNumSamplesInBin - todo consider*/);
 
         AbstractSplittingRule bestSplittingRule = findBestSplit(histogram);
-        double criterionForTheParentNode = entropyBinarySplit(1.0 * countsByClass[0] / (countsByClass[0] + countsByClass[1]));
+        double criterionForTheParentNode = entropyBinarySplitMultinomial(countsByClass, Arrays.stream(countsByClass).sum());
         // if no split could be found, make a list from current node
         // if the information gain is low, make a leaf from current node
         if (bestSplittingRule == null
@@ -365,7 +361,7 @@ public BuilderVisibility builderVisibility() {
     public ModelCategory[] can_build() {
         return new ModelCategory[]{
                 ModelCategory.Binomial,
-//                ModelCategory.Multinomial,
+                ModelCategory.Multinomial,
 //                                            ModelCategory.Ordinal,
 //                ModelCategory.Regression
         };

h2o-algos/src/main/java/hex/tree/dt/binning/AbstractBin.java

@@ -5,11 +5,11 @@
  * Single bin holding limits (min excluded), count of samples and count of class 0.
  */
 public abstract class AbstractBin {
-    public int _count0;
+    public int[] _classesDistribution;
     public int _count;
 
-    public int getCount0() {
-        return _count0;
+    public int getClassCount(int i) {
+        return _classesDistribution[i];
     }
 
     public abstract AbstractBin clone();

h2o-algos/src/main/java/hex/tree/dt/binning/BinningStrategy.java

@@ -9,13 +9,12 @@
 import java.math.BigDecimal;
 import java.math.RoundingMode;
 import java.util.ArrayList;
+import java.util.Arrays;
 import java.util.List;
-
-import static hex.tree.dt.mrtasks.CountBinsSamplesCountsMRTask.COUNT;
-import static hex.tree.dt.mrtasks.CountBinsSamplesCountsMRTask.COUNT_0;
+import java.util.stream.DoubleStream;
 
 /**
- * Strategy for binning. Creates bins for single feature.
+ * Strategy for binning. Create bins for single feature.
  */
 public enum BinningStrategy {
 
@@ -39,13 +38,13 @@ public enum BinningStrategy {
             return roundToNDecimalPoints(number, DECIMALS_TO_CONSIDER);
         }
 
-        private List<AbstractBin> createEmptyBinsFromBinningValues(List<Double> binningValues, double realMin, double realMax) {
+        private List<AbstractBin> createEmptyBinsFromBinningValues(List<Double> binningValues, double realMin, double realMax, int nclass) {
             List<AbstractBin> emptyBins = new ArrayList<>();
             // create bins between nearest binning values, don't create bin starting with the last value (on index size - 1)
             for (int i = 0; i < binningValues.size() - 1; i++) {
                 emptyBins.add(
                         new NumericBin(roundToNDecimalPoints(binningValues.get(i)),
-                                roundToNDecimalPoints(binningValues.get(i + 1))));
+                                roundToNDecimalPoints(binningValues.get(i + 1)), nclass));
             }
             // set the firs min to some lower value (relative to step) so the actual value equal to min is not lost
             ((NumericBin) emptyBins.get(0)).setMin(realMin - MIN_REL_COEFF * (binningValues.get(1) - binningValues.get(0)));
@@ -55,7 +54,7 @@ private List<AbstractBin> createEmptyBinsFromBinningValues(List<Double> binningV
         }
 
         @Override
-        List<AbstractBin> createFeatureBins(Frame originData, DataFeaturesLimits featuresLimits, int feature) {
+        List<AbstractBin> createFeatureBins(Frame originData, DataFeaturesLimits featuresLimits, int feature, int nclass) {
             if (originData.vec(feature).isNumeric()) {
                 NumericFeatureLimits featureLimits = (NumericFeatureLimits) featuresLimits.getFeatureLimits(feature);
                 double step = (featureLimits._max - featureLimits._min) / NUM_BINS;
@@ -69,7 +68,7 @@ List<AbstractBin> createFeatureBins(Frame originData, DataFeaturesLimits feature
                     binningValues.add(value);
                 }
                 List<AbstractBin> emptyBins = createEmptyBinsFromBinningValues(
-                        binningValues, featureLimits._min, featureLimits._max);
+                        binningValues, featureLimits._min, featureLimits._max, nclass);
 
                 return calculateNumericBinSamplesCount(originData, emptyBins, featuresLimits.toDoubles(), feature);
             } else {
@@ -78,7 +77,7 @@ List<AbstractBin> createFeatureBins(Frame originData, DataFeaturesLimits feature
                 for (int category = 0; category < featureLimits._mask.length; category++) {
                     // if the category is present in feature values, add new bin for this category
                     if (featureLimits._mask[category]) {
-                        emptyBins.add(new CategoricalBin(category));
+                        emptyBins.add(new CategoricalBin(category, nclass));
                     }
                 }
 
@@ -95,7 +94,7 @@ List<AbstractBin> createFeatureBins(Frame originData, DataFeaturesLimits feature
      */
     EQUAL_HEIGHT {
         @Override
-        List<AbstractBin> createFeatureBins(Frame originData, DataFeaturesLimits featuresLimits, int feature) {
+        List<AbstractBin> createFeatureBins(Frame originData, DataFeaturesLimits featuresLimits, int feature, int nclass) {
             return null;
         }
 
@@ -106,7 +105,7 @@ List<AbstractBin> createFeatureBins(Frame originData, DataFeaturesLimits feature
      */
     CUSTOM_BINS {
         @Override
-        List<AbstractBin> createFeatureBins(Frame originData, DataFeaturesLimits featuresLimits, int feature) {
+        List<AbstractBin> createFeatureBins(Frame originData, DataFeaturesLimits featuresLimits, int feature, int nclass) {
             return null;
         }
     };
@@ -121,7 +120,7 @@ List<AbstractBin> createFeatureBins(Frame originData, DataFeaturesLimits feature
      * @param feature        selected feature index
      * @return list of created bins
      */
-    abstract List<AbstractBin> createFeatureBins(Frame originData, DataFeaturesLimits featuresLimits, int feature);
+    abstract List<AbstractBin> createFeatureBins(Frame originData, DataFeaturesLimits featuresLimits, int feature, int nclass);
 
     /**
      * Calculates samples count for given bins for categorical feature.
@@ -136,11 +135,14 @@ private static List<AbstractBin> calculateCategoricalBinSamplesCount(Frame data,
                                                                          double[][] featuresLimits, int feature) {
         // run MR task to compute accumulated statistic for bins - one task for one feature, calculates all bins at once
         double[][] binsArray = bins.stream().map(AbstractBin::toDoubles).toArray(double[][]::new);
-        CountBinsSamplesCountsMRTask task = new CountBinsSamplesCountsMRTask(feature, featuresLimits, binsArray);
+        int countsOffset = CountBinsSamplesCountsMRTask.CAT_COUNT_OFFSET;
+        CountBinsSamplesCountsMRTask task = new CountBinsSamplesCountsMRTask(feature, featuresLimits, binsArray, countsOffset);
         task.doAll(data);
-        for(int i = 0; i < binsArray.length; i ++) {
-            bins.get(i)._count = (int) task._bins[i][COUNT];
-            bins.get(i)._count0 = (int) task._bins[i][COUNT_0];
+        for (int i = 0; i < binsArray.length; i++) {
+            bins.get(i)._count = (int) task._bins[i][countsOffset];
+            bins.get(i)._classesDistribution =
+                    DoubleStream.of(Arrays.copyOfRange(task._bins[i], countsOffset + 1, task._bins[i].length))
+                            .mapToInt(c -> (int) c).toArray();
         }
         return bins;
     }
@@ -158,11 +160,15 @@ private static List<AbstractBin> calculateNumericBinSamplesCount(Frame data, Lis
                                                              double[][] featuresLimits, int feature) {
         // run MR task to compute accumulated statistic for bins - one task for one feature, calculates all bins at once
         double[][] binsArray = bins.stream().map(AbstractBin::toDoubles).toArray(double[][]::new);
-        CountBinsSamplesCountsMRTask task = new CountBinsSamplesCountsMRTask(feature, featuresLimits, binsArray);
+        int countsOffset = CountBinsSamplesCountsMRTask.NUM_COUNT_OFFSET;
+        CountBinsSamplesCountsMRTask task = new CountBinsSamplesCountsMRTask(feature, featuresLimits, binsArray, countsOffset);
         task.doAll(data);
-        for(int i = 0; i < binsArray.length; i ++) {
-            bins.get(i)._count = (int) task._bins[i][COUNT];
-            bins.get(i)._count0 = (int) task._bins[i][COUNT_0];
+
+        for (int i = 0; i < binsArray.length; i++) {
+            bins.get(i)._count = (int) task._bins[i][countsOffset];
+            bins.get(i)._classesDistribution =
+                    DoubleStream.of(Arrays.copyOfRange(task._bins[i], countsOffset + 1, task._bins[i].length))
+                            .mapToInt(c -> (int) c).toArray();
         }
         return bins;
     }

h2o-algos/src/main/java/hex/tree/dt/binning/CategoricalBin.java

@@ -1,33 +1,39 @@
 package hex.tree.dt.binning;
 
+import org.apache.commons.lang.ArrayUtils;
+
+import java.util.Arrays;
+
 /**
  * For categorical features values are already binned to categories - each bin corresponds to one value (category)
  */
 public class CategoricalBin extends AbstractBin {
     public int _category;
 
-    public CategoricalBin(int category, int count, int count0) {
+    public CategoricalBin(int category, int[] classesDistribution, int count) {
         _category = category;
+        _classesDistribution = classesDistribution;
         _count = count;
-        _count0 = count0;
     }
 
-    public CategoricalBin(int category) {
+    public CategoricalBin(int category, int nclass) {
         _category = category;
+        _classesDistribution = new int[nclass];
         _count = 0;
-        _count0 = 0;
     }
 
     public int getCategory() {
         return _category;
     }
 
     public CategoricalBin clone() {
-        return new CategoricalBin(_category, _count, _count0);
+        return new CategoricalBin(_category, _classesDistribution, _count);
     }
 
     public double[] toDoubles() {
-        return new double[]{_category, _count, _count0};
+        // category|count|class0|class1|...
+        return ArrayUtils.addAll(new double[]{_category, _count}, 
+                Arrays.stream(_classesDistribution).asDoubleStream().toArray());
     }
 
 }