Adds protobuf serialization to Tree models

Classification/DecisionTree/src/test/java/org/tribuo/classification/dtree/TestCART.java

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2015-2021, Oracle and/or its affiliates. All rights reserved.
+ * Copyright (c) 2015, 2022, Oracle and/or its affiliates. All rights reserved.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -44,7 +44,7 @@ public class TestCART {
     private static final CARTClassificationTrainer randomt = new CARTClassificationTrainer(5,2, 0.0f,1.0f, true,
             new GiniIndex(), Trainer.DEFAULT_SEED);
 
-    public static Model<Label> testCART(Pair<Dataset<Label>,Dataset<Label>> p, CARTClassificationTrainer trainer) {
+    public static TreeModel<Label> testCART(Pair<Dataset<Label>,Dataset<Label>> p, CARTClassificationTrainer trainer) {
         TreeModel<Label> m = trainer.train(p.getA());
         LabelEvaluator e = new LabelEvaluator();
         LabelEvaluation evaluation = e.evaluate(m,p.getB());
@@ -54,6 +54,7 @@ public static Model<Label> testCART(Pair<Dataset<Label>,Dataset<Label>> p, CARTC
         features = m.getTopFeatures(-1);
         Assertions.assertNotNull(features);
         Assertions.assertFalse(features.isEmpty());
+        Helpers.testModelProtoSerialization(m, Label.class, p.getB());
 
         return m;
     }
@@ -80,7 +81,7 @@ public void testRandomSingleClassTraining() {
         runSingleClassTraining(randomt);
     }
 
-    public Model<Label> runDenseData(CARTClassificationTrainer trainer) {
+    public static TreeModel<Label> runDenseData(CARTClassificationTrainer trainer) {
         Pair<Dataset<Label>,Dataset<Label>> p = LabelledDataGenerator.denseTrainTest();
         return testCART(p, trainer);
     }
@@ -91,6 +92,20 @@ public void testDenseData() {
         Helpers.testModelSerialization(model,Label.class);
     }
 
+    @Test
+    public void testDecisionStump() {
+        CARTClassificationTrainer stumpTrainer = new CARTClassificationTrainer(1);
+        TreeModel<Label> model = runDenseData(stumpTrainer);
+        assertEquals(3,model.countNodes(model.getRoot()));
+    }
+
+    @Test
+    public void testMajorityPrediction() {
+        CARTClassificationTrainer stumpTrainer = new CARTClassificationTrainer(0);
+        TreeModel<Label> model = runDenseData(stumpTrainer);
+        assertEquals(1,model.countNodes(model.getRoot()));
+    }
+
     @Test
     public void testRandomDenseData() {
         runDenseData(randomt);

Common/Trees/src/main/java/org/tribuo/common/tree/AbstractCARTTrainer.java

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2015-2021, Oracle and/or its affiliates. All rights reserved.
+ * Copyright (c) 2015, 2022, Oracle and/or its affiliates. All rights reserved.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -126,8 +126,8 @@ public synchronized void postConfig() {
             throw new IllegalArgumentException("minImpurityDecrease must be greater than or equal to 0");
         }
 
-        if (maxDepth < 1) {
-            throw new IllegalArgumentException("maxDepth must be greater than or equal to 1");
+        if (maxDepth < 0) {
+            throw new IllegalArgumentException("maxDepth must be non-negative");
         }
 
         if (minChildWeight <= 0.0f) {

Common/Trees/src/main/java/org/tribuo/common/tree/LeafNode.java

@@ -16,10 +16,14 @@
 
 package org.tribuo.common.tree;
 
+import com.google.protobuf.Any;
 import org.tribuo.Example;
 import org.tribuo.Output;
 import org.tribuo.Prediction;
+import org.tribuo.common.tree.protos.LeafNodeProto;
+import org.tribuo.common.tree.protos.TreeNodeProto;
 import org.tribuo.math.la.SparseVector;
+import org.tribuo.protos.core.OutputProto;
 
 import java.util.Collections;
 import java.util.HashMap;
@@ -35,6 +39,11 @@
 public class LeafNode<T extends Output<T>> implements Node<T> {
     private static final long serialVersionUID = 4L;
 
+    /**
+     * Protobuf serialization version.
+     */
+    public static final int CURRENT_VERSION = 0;
+
     private final double impurity;
 
     private final T output;
@@ -136,4 +145,102 @@ public String toString() {
         return "LeafNode(impurity="+impurity+",output="+output.toString()+",scores="+scores.toString()+",probability="+generatesProbabilities+")";
     }
 
+    TreeNodeProto serialize(int parentIdx, int curIdx) {
+        LeafNodeProto.Builder nodeBuilder = LeafNodeProto.newBuilder();
+        nodeBuilder.setParentIdx(parentIdx);
+        nodeBuilder.setCurIdx(curIdx);
+        nodeBuilder.setOutput(output.serialize());
+        for (Map.Entry<String, T> e : scores.entrySet()) {
+            nodeBuilder.putScore(e.getKey(), e.getValue().serialize());
+        }
+        nodeBuilder.setGeneratesProbabilities(generatesProbabilities);
+        nodeBuilder.setImpurity(impurity);
+
+
+        TreeNodeProto.Builder builder = TreeNodeProto.newBuilder();
+        builder.setVersion(CURRENT_VERSION);
+        builder.setClassName(LeafNode.class.getName());
+        builder.setSerializedData(Any.pack(nodeBuilder.build()));
+
+        return builder.build();
+    }
+
+    static final class LeafNodeBuilder<T extends Output<T>> extends TreeModel.NodeBuilder implements Node<T> {
+        private final int parentIdx;
+        private final int curIdx;
+        private final double impurity;
+        private final T output;
+        private final Map<String,T> scores;
+        private final boolean generatesProbabilities;
+
+        @SuppressWarnings("unchecked")
+        LeafNodeBuilder(LeafNodeProto proto) {
+            this.parentIdx = proto.getParentIdx();
+            this.curIdx = proto.getCurIdx();
+            this.impurity = proto.getImpurity();
+            this.output = (T) Output.deserialize(proto.getOutput());
+            this.scores = new HashMap<>();
+            for (Map.Entry<String, OutputProto> e : proto.getScoreMap().entrySet()) {
+                Output<?> curOutput = Output.deserialize(e.getValue());
+                if (!curOutput.getClass().equals(output.getClass())) {
+                    throw new IllegalStateException("Invalid protobuf, scores were not the same type as the most likely output, found " + curOutput.getClass() + ", expected " + output.getClass());
+                }
+                this.scores.put(e.getKey(), (T) curOutput);
+            }
+            this.generatesProbabilities = proto.getGeneratesProbabilities();
+        }
+
+        LeafNodeBuilder(int parentIdx, int curIdx, double impurity, T output, Map<String, T> scores, boolean generatesProbabilities) {
+            this.parentIdx = parentIdx;
+            this.curIdx = curIdx;
+            this.impurity = impurity;
+            this.output = output;
+            this.scores = scores;
+            this.generatesProbabilities = generatesProbabilities;
+        }
+
+        @Override
+        public boolean isLeaf() {
+            return true;
+        }
+
+        @Override
+        public Node<T> getNextNode(SparseVector example) {
+            return null;
+        }
+
+        @Override
+        public double getImpurity() {
+            return impurity;
+        }
+
+        @Override
+        public LeafNodeBuilder<T> copy() {
+            return new LeafNodeBuilder<>(parentIdx,curIdx,impurity,output.copy(),new HashMap<>(scores),generatesProbabilities);
+        }
+
+        /**
+         * Gets the index of the parent node.
+         * @return The parent index.
+         */
+        int getParentIdx() {
+            return parentIdx;
+        }
+
+        /**
+         * Gets the index of this node.
+         * @return The current node index.
+         */
+        int getCurIdx() {
+            return curIdx;
+        }
+
+        /**
+         * Builds this builder into a leaf node.
+         * @return The leaf node.
+         */
+        LeafNode<T> build() {
+            return new LeafNode<>(impurity,output.copy(),new HashMap<>(scores),generatesProbabilities);
+        }
+    }
 }

Common/Trees/src/main/java/org/tribuo/common/tree/Node.java

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2015-2020, Oracle and/or its affiliates. All rights reserved.
+ * Copyright (c) 2015, 2022, Oracle and/or its affiliates. All rights reserved.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -50,5 +50,4 @@ public interface Node<T extends Output<T>> extends Serializable {
      * @return A deep copy.
      */
     public Node<T> copy();
-
 }

Common/Trees/src/main/java/org/tribuo/common/tree/SplitNode.java

@@ -16,7 +16,10 @@
 
 package org.tribuo.common.tree;
 
+import com.google.protobuf.Any;
 import org.tribuo.Output;
+import org.tribuo.common.tree.protos.SplitNodeProto;
+import org.tribuo.common.tree.protos.TreeNodeProto;
 import org.tribuo.math.la.SparseVector;
 
 import java.util.Objects;
@@ -27,6 +30,11 @@
 public class SplitNode<T extends Output<T>> implements Node<T> {
     private static final long serialVersionUID = 3L;
 
+    /**
+     * Protobuf serialization version.
+     */
+    public static final int CURRENT_VERSION = 0;
+
     private final Node<T> greaterThan;
 
     private final Node<T> lessThanOrEqual;
@@ -138,5 +146,120 @@ public boolean equals(Object o) {
     public int hashCode() {
         return Objects.hash(greaterThan, lessThanOrEqual, splitFeature, splitValue, impurity);
     }
+
+    TreeNodeProto serialize(int parentIdx, int curIdx, int greaterThanIdx, int lessThanOrEqualIdx) {
+        SplitNodeProto.Builder nodeBuilder = SplitNodeProto.newBuilder();
+        nodeBuilder.setParentIdx(parentIdx);
+        nodeBuilder.setCurIdx(curIdx);
+        nodeBuilder.setGreaterThanIdx(greaterThanIdx);
+        nodeBuilder.setLessThanOrEqualIdx(lessThanOrEqualIdx);
+        nodeBuilder.setSplitFeatureIdx(splitFeature);
+        nodeBuilder.setSplitValue(splitValue);
+        nodeBuilder.setImpurity(impurity);
+
+        TreeNodeProto.Builder builder = TreeNodeProto.newBuilder();
+        builder.setVersion(CURRENT_VERSION);
+        builder.setClassName(LeafNode.class.getName());
+        builder.setSerializedData(Any.pack(nodeBuilder.build()));
+
+        return builder.build();
+    }
+
+    static final class SplitNodeBuilder<T extends Output<T>> extends TreeModel.NodeBuilder implements Node<T> {
+
+        private final int parentIdx;
+        private final int curIdx;
+        private final int greaterThanIdx;
+        private final int lessThanOrEqualIdx;
+        private final int splitFeature;
+        private final double splitValue;
+        private final double impurity;
+
+        private Node<T> greaterThan;
+        private Node<T> lessThanOrEqual;
+
+        SplitNodeBuilder(SplitNodeProto proto) {
+            this.parentIdx = proto.getParentIdx();
+            this.curIdx = proto.getCurIdx();
+            this.greaterThanIdx = proto.getGreaterThanIdx();
+            this.lessThanOrEqualIdx = proto.getLessThanOrEqualIdx();
+            this.splitFeature = proto.getSplitFeatureIdx();
+            this.splitValue = proto.getSplitValue();
+            this.impurity = proto.getImpurity();
+        }
+
+        SplitNodeBuilder(int parentIdx, int curIdx, int greaterThanIdx, int lessThanOrEqualIdx, int splitFeature, double splitValue, double impurity) {
+            this.parentIdx = parentIdx;
+            this.curIdx = curIdx;
+            this.greaterThanIdx = greaterThanIdx;
+            this.lessThanOrEqualIdx = lessThanOrEqualIdx;
+            this.splitFeature = splitFeature;
+            this.splitValue = splitValue;
+            this.impurity = impurity;
+        }
+
+        @Override
+        public boolean isLeaf() {
+            return false;
+        }
+
+        @Override
+        public Node<T> getNextNode(SparseVector example) {
+            return null;
+        }
+
+        @Override
+        public double getImpurity() {
+            return impurity;
+        }
+
+        @Override
+        public SplitNodeBuilder<T> copy() {
+            return new SplitNodeBuilder<>(parentIdx, curIdx, greaterThanIdx, lessThanOrEqualIdx, splitFeature, splitValue, impurity);
+        }
+
+        boolean canBuild() {
+            return greaterThan != null && lessThanOrEqual != null;
+        }
+
+        SplitNode<T> build() {
+            if (!canBuild()) {
+                throw new IllegalStateException("Not ready to build this split node, missing the children pointers");
+            }
+            return new SplitNode<>(splitValue,splitFeature,impurity,greaterThan,lessThanOrEqual);
+        }
+
+        void setGreaterThan(Node<T> greaterThan) {
+            if (this.greaterThan == null) {
+                this.greaterThan = greaterThan;
+            } else {
+                throw new IllegalStateException("Invalid protobuf, multiple nodes mapped to the greaterThanIdx");
+            }
+        }
+
+        void setLessThanOrEqual(Node<T> lessThanOrEqual) {
+            if (this.lessThanOrEqual == null) {
+                this.lessThanOrEqual = lessThanOrEqual;
+            } else {
+                throw new IllegalStateException("Invalid protobuf, multiple nodes mapped to the lessThanOrEqualIdx");
+            }
+        }
+
+        int getGreaterThanIdx() {
+            return greaterThanIdx;
+        }
+
+        int getLessThanOrEqualIdx() {
+            return lessThanOrEqualIdx;
+        }
+
+        int getParentIdx() {
+            return parentIdx;
+        }
+
+        int getCurIdx() {
+            return curIdx;
+        }
+    }
 }