Add categorical test

h2oai · Sep 14, 2023 · 665201f · 665201f
1 parent 814c847
commit 665201f
Showing 1 changed file with 34 additions and 0 deletions.
diff --git a/h2o-algos/src/test/java/hex/adaboost/AdaBoostTest.java b/h2o-algos/src/test/java/hex/adaboost/AdaBoostTest.java
@@ -1,5 +1,6 @@
 package hex.adaboost;
 
+import hex.Model;
 import hex.genmodel.algos.tree.SharedTreeSubgraph;
 import hex.glm.GLM;
 import hex.glm.GLMModel;
@@ -182,6 +183,39 @@ public void testBasicTrainAndScore() {
         }
     }
 
+    @Test
+    public void testBasicTrainAndScoreCategorical() {
+        try {
+            Scope.enter();
+            Frame train = parseTestFile("smalldata/prostate/prostate.csv");
+            Scope.track(train);
+            String response = "CAPSULE";
+            train.toCategoricalCol(response);
+            train.toCategoricalCol("RACE");
+            train.toCategoricalCol("DPROS");
+            train.toCategoricalCol("DCAPS");
+            train.toCategoricalCol("GLEASON");
+            AdaBoostModel.AdaBoostParameters p = new AdaBoostModel.AdaBoostParameters();
+            p._train = train._key;
+            p._seed = 0xDECAF;
+            p._n_estimators = 50;
+            p._response_column = response;
+            p._categorical_encoding = Model.Parameters.CategoricalEncodingScheme.OneHotExplicit;
+
+            AdaBoost adaBoost = new AdaBoost(p);
+            AdaBoostModel adaBoostModel = adaBoost.trainModel().get();
+            Scope.track_generic(adaBoostModel);
+            assertNotNull(adaBoostModel);
+
+            System.out.println("train.toTwoDimTable() = " + train.toTwoDimTable());
+
+            Frame score = adaBoostModel.score(train);
+            Scope.track(score);
+        } finally {
+            Scope.exit();
+        }
+    }
+
 //    @Test
 //    public void testBasicTrainAndScoreGLM() {
 //        try {