001/*
002 * Copyright (c) 2015-2020, Oracle and/or its affiliates. All rights reserved.
003 *
004 * Licensed under the Apache License, Version 2.0 (the "License");
005 * you may not use this file except in compliance with the License.
006 * You may obtain a copy of the License at
007 *
008 *     http://www.apache.org/licenses/LICENSE-2.0
009 *
010 * Unless required by applicable law or agreed to in writing, software
011 * distributed under the License is distributed on an "AS IS" BASIS,
012 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express implied.
013 * See the License for the specific language governing permissions and
014 * limitations under the License.
015 */
016
017package org.tribuo.classification.dtree;
018
019import com.oracle.labs.mlrg.olcut.config.Option;
020import org.tribuo.Trainer;
021import org.tribuo.classification.ClassificationOptions;
022import org.tribuo.classification.dtree.impurity.Entropy;
023import org.tribuo.classification.dtree.impurity.GiniIndex;
024import org.tribuo.classification.dtree.impurity.LabelImpurity;
025
026/**
027 * Options for building a classification tree trainer.
028 */
029public class CARTClassificationOptions implements ClassificationOptions<CARTClassificationTrainer> {
030
031    @Override
032    public String getOptionsDescription() {
033        return "Options for decision/classification trees.";
034    }
035
036    /**
037     * Type of decision tree algorithm.
038     */
039    public enum TreeType {
040        /**
041         * Builds a CART model.
042         */
043        CART
044    }
045
046    /**
047     * The impurity algorithm.
048     */
049    public enum ImpurityType {
050        /**
051         * Uses {@link GiniIndex}.
052         */
053        GINI,
054        /**
055         * Uses {@link Entropy}.
056         */
057        ENTROPY
058    }
059
060    /**
061     * Maximum depth in the decision tree.
062     */
063    @Option(longName = "cart-max-depth", usage = "Maximum depth in the decision tree.")
064    public int cartMaxDepth = 6;
065    /**
066     * Minimum child weight.
067     */
068    @Option(longName = "cart-min-child-weight", usage = "Minimum child weight.")
069    public float cartMinChildWeight = 5.0f;
070    /**
071     * Minimum impurity decrease.
072     */
073    @Option(longName = "cart-min-impurity-decrease", usage = "Minimum impurity decrease.")
074    public float cartMinImpurityDecrease = 0.0f;
075    /**
076     * Fraction of features in split.
077     */
078    @Option(longName = "cart-split-fraction", usage = "Fraction of features in split.")
079    public float cartSplitFraction = 1.0f;
080    /**
081     * Whether to choose split points for features at random.
082     */
083    @Option(longName = "cart-random-split", usage = "Whether to choose split points for features at random.")
084    public boolean cartRandomSplit = false;
085    /**
086     * Impurity measure to use. Defaults to GINI.
087     */
088    @Option(longName = "cart-impurity", usage = "Impurity measure to use. Defaults to GINI.")
089    public ImpurityType cartImpurity = ImpurityType.GINI;
090    /**
091     * Prints the decision tree.
092     */
093    @Option(longName = "cart-print-tree", usage = "Prints the decision tree.")
094    public boolean cartPrintTree;
095    /**
096     * Tree algorithm to use (options are CART).
097     */
098    @Option(longName = "cart-tree-algorithm", usage = "Tree algorithm to use (options are CART).")
099    public TreeType cartTreeAlgorithm = TreeType.CART;
100    /**
101     * RNG seed.
102     */
103    @Option(longName = "cart-seed", usage = "RNG seed.")
104    public long cartSeed = Trainer.DEFAULT_SEED;
105
106    @Override
107    public CARTClassificationTrainer getTrainer() {
108        LabelImpurity impurity;
109        switch (cartImpurity) {
110            case GINI:
111                impurity = new GiniIndex();
112                break;
113            case ENTROPY:
114                impurity = new Entropy();
115                break;
116            default:
117                throw new IllegalArgumentException("unknown impurity type " + cartImpurity);
118        }
119
120        CARTClassificationTrainer trainer;
121        switch (cartTreeAlgorithm) {
122            case CART:
123                trainer = new CARTClassificationTrainer(cartMaxDepth, cartMinChildWeight, cartMinImpurityDecrease,
124                        cartSplitFraction, cartRandomSplit, impurity, cartSeed);
125                break;
126            default:
127                throw new IllegalArgumentException("Unknown tree type " + cartTreeAlgorithm);
128        }
129
130        return trainer;
131    }
132
133}