cleanup for next release

Author: dvilaverde

pom.xml

@@ -6,7 +6,7 @@
 
     <groupId>rocks.vilaverde</groupId>
     <artifactId>scikit-learn-2-java</artifactId>
-    <version>1.0.2-SNAPSHOT</version>
+    <version>1.0.2</version>
 
     <name>${project.groupId}:${project.artifactId}</name>
     <description>A sklearn exported_text models parser for executing in the Java runtime.</description>

src/main/java/rocks/vilaverde/classifier/dt/AbstractDecisionTreeVisitor.java

@@ -28,8 +28,7 @@ public void visit(ChoiceNode object) {
   public void visit(DecisionNode object) {
     visitBase(object);
 
-    for (TreeNode child : object.getChildren()) {
-      child.accept(this);
-    }
+    object.getLeft().accept(this);
+    object.getRight().accept(this);
   }
 }

src/main/java/rocks/vilaverde/classifier/dt/DecisionNode.java

@@ -5,26 +5,76 @@
 
 /**
  * Represents a decision in the DecisionTreeClassifier. The decision will have
- * a left and right hand choice to be evaluated. Choices may have nested DecisionNodes.
+ * a left and right hand {@link ChoiceNode} to be evaluated.
+ * A {@link ChoiceNode} may have nested {@link DecisionNode} or {@link EndNode}.
  */
 class DecisionNode extends TreeNode {
 
   private final String featureName;
 
-  private final List<TreeNode> children = new ArrayList<>(2);
-
-  public List<TreeNode> getChildren() {
-    return children;
-  }
+  private ChoiceNode left;
+  private ChoiceNode right;
 
+  /**
+   * Factory method to create a {@link DecisionNode}.
+   * @param feature the name of the feature.
+   * @return DecisionNode
+   */
   public static DecisionNode create(String feature) {
     return new DecisionNode(feature);
   }
 
+  /**
+   * Private Constructor.
+   * @param featureName
+   */
   private DecisionNode(String featureName) {
-    this.featureName = featureName;
+    this.featureName = featureName.intern();
+  }
+
+  /**
+   * Getter for the left.
+   * @return ChoiceNode
+   */
+  public ChoiceNode getLeft() {
+    return left;
+  }
+
+  /**
+   * Left hand side of the decision.
+   * @param left the left {@link ChoiceNode}
+   */
+  public void setLeft(ChoiceNode left) {
+    this.left = left;
+  }
+
+  /**
+   * Getter for the right.
+   * @return ChoiceNode
+   */
+  public ChoiceNode getRight() {
+    return right;
+  }
+
+  /**
+   * Right hand side of a decision.
+   * @param right the right {@link ChoiceNode}
+   */
+  public void setRight(ChoiceNode right) {
+    this.right = right;
   }
 
+  /**
+   * @return true when the left and right choice are set on this decision.
+   */
+  public boolean isComplete() {
+    return getLeft() != null && getRight() != null;
+  }
+
+  /**
+   * Getter for the feature used in this decision node.
+   * @return String
+   */
   public String getFeatureName() {
     return featureName;
   }
@@ -36,6 +86,15 @@ public void accept(AbstractDecisionTreeVisitor visitor) {
 
   @Override
   public String toString() {
-    return String.format("%s, %d operations", getFeatureName(), getChildren().size());
+    int count = 0;
+    if (getLeft() != null) {
+      count++;
+    }
+
+    if (getRight() != null) {
+      count++;
+    }
+
+    return String.format("%s, %d operations", getFeatureName(), count);
   }
 }

src/main/java/rocks/vilaverde/classifier/dt/DecisionTreeClassifier.java

@@ -3,6 +3,7 @@
 import rocks.vilaverde.classifier.Operator;
 import rocks.vilaverde.classifier.Prediction;
 
+import java.awt.*;
 import java.io.BufferedReader;
 import java.io.Reader;
 import java.util.Map;
@@ -36,7 +37,6 @@ public static <T> DecisionTreeClassifier<T> parse(Reader reader, PredictionFacto
 
   private final PredictionFactory<T> predictionFactory;
   private DecisionNode root;
-
   private Set<String> featureNames;
 
   /**
@@ -81,28 +81,19 @@ public Prediction<T> getClassification(Map<String, Double> features) {
       if (currentNode != null) {
         DecisionNode decisionNode = ((DecisionNode) currentNode);
 
-        ChoiceNode selection = null;
-        for (TreeNode child : decisionNode.getChildren()) {
-
-          Double featureValue = features.get(decisionNode.getFeatureName());
-          if (featureValue == null) {
-            featureValue = Double.NaN;
-          }
-
-          // find the path to traverse by evaluating all choices
-
-          if (((ChoiceNode) child).eval(featureValue)) {
-            selection = (ChoiceNode) child;
-            break;
-          }
+        Double featureValue = features.get(decisionNode.getFeatureName());
+        if (featureValue == null) {
+          featureValue = Double.NaN;
         }
 
-        if (selection != null) {
-          currentNode = selection.getChild();
+        if (decisionNode.getLeft().eval(featureValue)) {
+          currentNode = decisionNode.getLeft().getChild();
+        } else if (decisionNode.getRight().eval(featureValue)) {
+          currentNode = decisionNode.getRight().getChild();
         } else {
           // if I get here something is wrong since none of the branches evaluated to true
           throw new RuntimeException(String.format("no branches evaluated to true for feature '%s'",
-              decisionNode.getFeatureName()));
+                  decisionNode.getFeatureName()));
         }
       }
     }
@@ -174,7 +165,7 @@ private void processChildNode(Stack<TreeNode> stack, String line) throws Excepti
     // if the current choice has been popped check if the decision node
     // has 2 operations, and if so pop that one as well and any other
     // completed decision nodes.
-    while (stack.size() > 1 && ((DecisionNode)stack.peek()).getChildren().size() == 2 ) {
+    while (stack.size() > 1 && ((DecisionNode)stack.peek()).isComplete() ) {
       stack.pop();
     }
   }
@@ -215,7 +206,13 @@ private void processDecisionNode(Stack<TreeNode> stack, String line) {
     }
 
     ChoiceNode choice = ChoiceNode.create(op, value);
-    decisionNode.getChildren().add(choice);
+
+    if (decisionNode.getLeft() == null) {
+      decisionNode.setLeft(choice);
+    } else {
+      decisionNode.setRight(choice);
+    }
+
     stack.push(choice);
   }