Search in sources :

Example 1 with DecisionTreeModel

use of com.airbnb.aerosolve.core.models.DecisionTreeModel in project aerosolve by airbnb.

the class DecisionTreeTransformTest method makeTree.

/*
   * XOR like decision regions
   * 
   *         x = 2
   *          |
   *    -0.5  |   1.0
   *          |-------- y = 1
   *   -------- y = 0
   *          |
   *     0.25 |   -0.75
   *          |
   */
public DecisionTreeModel makeTree() {
    ArrayList<ModelRecord> records = new ArrayList<>();
    DecisionTreeModel tree = new DecisionTreeModel();
    tree.setStumps(records);
    // 0 - an x split at 2
    ModelRecord record = new ModelRecord();
    record.setFeatureFamily("loc");
    record.setFeatureName("x");
    record.setThreshold(2.0);
    record.setLeftChild(1);
    record.setRightChild(2);
    records.add(record);
    // 1 - a y split at 0
    record = new ModelRecord();
    record.setFeatureFamily("loc");
    record.setFeatureName("y");
    record.setThreshold(0.0);
    record.setLeftChild(3);
    record.setRightChild(4);
    records.add(record);
    // 2 - a y split at 1
    record = new ModelRecord();
    record.setFeatureFamily("loc");
    record.setFeatureName("y");
    record.setThreshold(1.0);
    record.setLeftChild(5);
    record.setRightChild(6);
    records.add(record);
    // 3  a leaf
    record = new ModelRecord();
    record.setFeatureWeight(0.25);
    records.add(record);
    // 4  a leaf
    record = new ModelRecord();
    record.setFeatureWeight(-0.5);
    records.add(record);
    // 5  a leaf
    record = new ModelRecord();
    record.setFeatureWeight(-0.75);
    records.add(record);
    // 6  a leaf
    record = new ModelRecord();
    record.setFeatureWeight(1.0);
    records.add(record);
    return tree;
}
Also used : DecisionTreeModel(com.airbnb.aerosolve.core.models.DecisionTreeModel) ModelRecord(com.airbnb.aerosolve.core.ModelRecord)

Example 2 with DecisionTreeModel

use of com.airbnb.aerosolve.core.models.DecisionTreeModel in project aerosolve by airbnb.

the class DecisionTreeTransformTest method testToHumanReadableConfig.

@Test
public void testToHumanReadableConfig() {
    DecisionTreeModel tree = makeTree();
    String result = tree.toHumanReadableTransform();
    log.info(result);
    String[] tokens = result.split("\n");
    assertEquals(9, tokens.length);
    assertTrue(tokens[3].contains("P,2,loc,y,1.000000,5,6"));
    assertTrue(tokens[4].contains("L,3,0.250000,LEAF_3"));
}
Also used : DecisionTreeModel(com.airbnb.aerosolve.core.models.DecisionTreeModel) Test(org.junit.Test)

Aggregations

DecisionTreeModel (com.airbnb.aerosolve.core.models.DecisionTreeModel)2 ModelRecord (com.airbnb.aerosolve.core.ModelRecord)1 Test (org.junit.Test)1