Examples with GBMModel hex.tree.gbm.GBMModel used on opensource projects

Example 11 with GBMModel

use of hex.tree.gbm.GBMModel in project h2o-3 by h2oai.

the class PartialDependenceTest method prostateBinaryPickCols.

@Test
public void prostateBinaryPickCols() {
    Frame fr = null;
    GBMModel model = null;
    PartialDependence partialDependence = null;
    try {
        // Frame
        fr = parse_test_file("smalldata/prostate/prostate.csv");
        for (String s : new String[] { "RACE", "GLEASON", "DPROS", "DCAPS", "CAPSULE" }) {
            Vec v = fr.remove(s);
            fr.add(s, v.toCategoricalVec());
            v.remove();
        }
        DKV.put(fr);
        // Model
        GBMModel.GBMParameters parms = new GBMModel.GBMParameters();
        parms._train = fr._key;
        parms._ignored_columns = new String[] { "ID" };
        parms._response_column = "CAPSULE";
        model = new GBM(parms).trainModel().get();
        // PartialDependence
        partialDependence = new PartialDependence(Key.<PartialDependence>make());
        //pick columns manually
        partialDependence._cols = new String[] { "DPROS", "GLEASON" };
        partialDependence._nbins = 10;
        partialDependence._model_id = (Key) model._key;
        partialDependence._frame_id = fr._key;
        partialDependence.execImpl().get();
        for (TwoDimTable t : partialDependence._partial_dependence_data) Log.info(t);
        Assert.assertTrue(partialDependence._partial_dependence_data.length == 2);
    } finally {
        if (fr != null)
            fr.remove();
        if (model != null)
            model.remove();
        if (partialDependence != null)
            partialDependence.remove();
    }
}