Added support for the 'nbins_cats' hyperparameter. HT: Eric Blood. Fixes

#4, fixes #16
jpmml · Aug 22, 2024 · 4f50e83 · 4f50e83
1 parent 698ffb2
commit 4f50e83
Show file tree

Hide file tree

Showing 8 changed files with 1,218 additions and 1,196 deletions.
diff --git a/pmml-h2o/src/main/java/org/jpmml/h2o/SharedTreeMojoModelConverter.java b/pmml-h2o/src/main/java/org/jpmml/h2o/SharedTreeMojoModelConverter.java
@@ -201,23 +201,9 @@ public Node encodeNode(SharedTree sharedTree, ByteBufferWrapper byteBuffer, Inte
 		} else
 
 		{
-			if(equal != 0){
+			if(feature instanceof CategoricalFeature){
 				CategoricalFeature categoricalFeature = (CategoricalFeature)feature;
 
-				GenmodelBitSet bitSet = new GenmodelBitSet(0);
-
-				if(equal == 8){
-					bitSet.fill2(compressedTree, byteBuffer);
-				} else
-
-				if(equal == 12){
-					bitSet.fill3(compressedTree, byteBuffer);
-				} else
-
-				{
-					throw new IllegalArgumentException("Node type " + equal + " is not supported");
-				}
-
 				String name = categoricalFeature.getName();
 				List<?> values = categoricalFeature.getValues();
 
@@ -226,19 +212,55 @@ public Node encodeNode(SharedTree sharedTree, ByteBufferWrapper byteBuffer, Inte
 				List<Object> leftValues = new ArrayList<>();
 				List<Object> rightValues = new ArrayList<>();
 
-				for(int i = 0; i < values.size(); i++){
-					Object value = values.get(i);
+				if(equal != 0){
+					GenmodelBitSet bitSet = new GenmodelBitSet(0);
 
-					if(!valueFilter.test(value)){
-						continue;
-					} // End if
+					if(equal == 8){
+						bitSet.fill2(compressedTree, byteBuffer);
+					} else
 
-					if(!bitSet.contains(i)){
-						leftValues.add(value);
+					if(equal == 12){
+						bitSet.fill3(compressedTree, byteBuffer);
 					} else
 
 					{
-						rightValues.add(value);
+						throw new IllegalArgumentException("Node type " + equal + " is not supported");
+					}
+
+					for(int i = 0; i < values.size(); i++){
+						Object value = values.get(i);
+
+						if(!valueFilter.test(value)){
+							continue;
+						} // End if
+
+						if(!bitSet.contains(i)){
+							leftValues.add(value);
+						} else
+
+						{
+							rightValues.add(value);
+						}
+					}
+				} else
+
+				{
+					Double splitVal = (double)byteBuffer.get4f();
+
+					for(int i = 0; i < values.size(); i++){
+						Object value = values.get(i);
+
+						if(!valueFilter.test(value)){
+							continue;
+						} // End if
+
+						if(i < splitVal){
+							leftValues.add(value);
+						} else
+
+						{
+							rightValues.add(value);
+						}
 					}
 				}