Repository: alteryx/evalml (python)

Issue:
Add sklearn DecisionTreeRegressor and DecisionTreeClassifier estimators
https://scikit-learn.org/stable/modules/generated/sklearn.tree.DecisionTreeRegressor.html

Params we should consider exposing to automl: max_depth, criterion, max_features.

Pull request:
Add DecisionTreeClassifier and DecisionTreeRegressor classes
Closes #1196 by adding DecisionTreeClassifier and DecisionTreeRegressor

- [x] Add DecisionTree* classes (+API reference)


TODO in later PR: Add DecisionTree* to AutoML and do perf testing


Patch:
diff --git a/docs/source/release_notes.rst b/docs/source/release_notes.rst
index a622a4ad..82a3ddfc 100644
--- a/docs/source/release_notes.rst
+++ b/docs/source/release_notes.rst
@@ -8,7 +8,8 @@ Release Notes
         * Added `detect_problem_type` to `problem_type/utils.py` to automatically detect the problem type given targets :pr:`1194`
         * Added LightGBM to AutoMLSearch :pr:`1199`
         * Updates scikit-learn and scikit-optimize to use latest versions - 0.23.2 and 0.8.1 respectively :pr:`1141`
-        * Add `ProblemTypes.all_problem_types` helper to get list of supported problem types :pr:`1219`
+        * Added `DecisionTreeClassifier` and `DecisionTreeRegressor` classes :pr:`1223`
+        * Added `ProblemTypes.all_problem_types` helper to get list of supported problem types :pr:`1219`
         * `DataChecks` can now be parametrized by passing a list of `DataCheck` classes and a parameter dictionary :pr:`1167`
         * Updated `flake8` configuration to enable linting on `__init__.py` files :pr:`1234`
     * Fixes
diff --git a/evalml/model_family/model_family.py b/evalml/model_family/model_family.py
index 0534605e..0bd151df 100644
--- a/evalml/model_family/model_family.py
+++ b/evalml/model_family/model_family.py
@@ -9,6 +9,7 @@ class ModelFamily(Enum):
     LINEAR_MODEL = 'linear_model'
     CATBOOST = 'catboost'
     EXTRA_TREES = 'extra_trees'
+    DECISION_TREE = 'decision_tree'
     BASELINE = 'baseline'
     NONE = 'none'
 
@@ -19,6 +20,7 @@ class ModelFamily(Enum):
                              ModelFamily.LINEAR_MODEL.name: "Linear",
                              ModelFamily.CATBOOST.name: "CatBoost",
                              ModelFamily.EXTRA_TREES.name: "Extra Trees",
+                             ModelFamily.DECISION_TREE.name: "Decision Tree",
                              ModelFamily.BASELINE.name: "Baseline",
                              ModelFamily.NONE.name: "None"}
         return model_family_dict[self.name]
@@ -27,7 +29,7 @@ class ModelFamily(Enum):
         return "ModelFamily." + self.name
 
     def is_tree_estimator(self):
-        """Checks whether the estimator's model family uses tree ensembles."""
+        """Checks whether the estimator's model family uses trees."""
         tree_estimators = {self.CATBOOST, self.EXTRA_TREES, self.RANDOM_FOREST,
-                           self.XGBOOST, self.LIGHTGBM}
+                           self.DECISION_TREE, self.XGBOOST, self.LIGHTGBM}
         return self in tree_estimators
diff --git a/evalml/pipelines/__init__.py b/evalml/pipelines/__init__.py
index bbbce77e..f4383db1 100644
--- a/evalml/pipelines/__init__.py
+++ b/evalml/pipelines/__init__.py
@@ -20,7 +20,9 @@ from .components import (
     ElasticNetClassifier,
     ElasticNetRegressor,
     ExtraTreesClassifier,
-    ExtraTreesRegressor
+    ExtraTreesRegressor,
+    DecisionTreeClassifier,
+    DecisionTreeRegressor
 )
 
 from .pipeline_base import PipelineBase
diff --git a/evalml/pipelines/components/__init__.py b/evalml/pipelines/components/__init__.py
index 53b12401..b2b0a61f 100644
--- a/evalml/pipelines/components/__init__.py
+++ b/evalml/pipelines/components/__init__.py
@@ -15,7 +15,9 @@ from .estimators import (
     ElasticNetClassifier,
     ElasticNetRegressor,
     BaselineClassifier,
-    BaselineRegressor
+    BaselineRegressor,
+    DecisionTreeClassifier,
+    DecisionTreeRegressor
 )
 from .transformers import (
     Transformer,
diff --git a/evalml/pipelines/components/estimators/__init__.py b/evalml/pipelines/components/estimators/__init__.py
index 60b769f4..08a8f5fa 100644
--- a/evalml/pipelines/components/estimators/__init__.py
+++ b/evalml/pipelines/components/estimators/__init__.py
@@ -6,11 +6,13 @@ from .classifiers import (LogisticRegressionClassifier,
                           CatBoostClassifier,
                           ElasticNetClassifier,
                           ExtraTreesClassifier,
-                          BaselineClassifier)
+                          BaselineClassifier,
+                          DecisionTreeClassifier)
 from .regressors import (LinearRegressor,
                          RandomForestRegressor,
                          CatBoostRegressor,
                          XGBoostRegressor,
                          ElasticNetRegressor,
                          ExtraTreesRegressor,
-                         BaselineRegressor)
+                         BaselineRegressor,
+                         DecisionTreeRegressor)
diff --git a/evalml/pipelines/components/estimators/classifiers/__init__.py b/evalml/pipelines/components/estimators/classifiers/__init__.py
index cc21992e..b6d33a56 100644
--- a/evalml/pipelines/components/estimators/classifiers/__init__.py
+++ b/evalml/pipelines/components/estimators/classifiers/__init__.py
@@ -6,3 +6,4 @@ from .elasticnet_classifier import ElasticNetClassifier
 from .et_classifier import ExtraTreesClassifier
 from .baseline_classifier import BaselineClassifier
 from .lightgbm_classifier import LightGBMClassifier
+from .decision_tree_classifier import DecisionTreeClassifier
diff --git a/evalml/pipelines/components/estimators/classifiers/decision_tree_classifier.py b/evalml/pipelines/components/estimators/classifiers/decision_tree_classifier.py
new file mode 100644
index 00000000..ed573c30
--- /dev/null
+++ b/evalml/pipelines/components/estimators/classifiers/decision_tree_classifier.py
@@ -0,0 +1,39 @@
+from sklearn.tree import DecisionTreeClassifier as SKDecisionTreeClassifier
+from skopt.space import Integer
+
+from evalml.model_family import ModelFamily
+from evalml.pipelines.components.estimators import Estimator
+from evalml.problem_types import ProblemTypes
+
+
+class DecisionTreeClassifier(Estimator):
+    """Decision Tree Classifier."""
+    name = "Decision Tree Classifier."
+    hyperparameter_ranges = {
+        "criterion": ["gini", "entropy"],
+        "max_features": ["auto", "sqrt", "log2"],
+        "max_depth": Integer(4, 10)
+    }
+    model_family = ModelFamily.DECISION_TREE
+    supported_problem_types = [ProblemTypes.BINARY, ProblemTypes.MULTICLASS]
+
+    def __init__(self,
+                 criterion="gini",
+                 max_features="auto",
+                 max_depth=6,
+                 min_samples_split=2,
+                 min_weight_fraction_leaf=0.0,
+                 random_state=0,
+                 **kwargs):
+        parameters = {"criterion": criterion,
+                      "max_features": max_features,
+                      "max_depth": max_depth,
+                      "min_samples_split": min_samples_split,
+                      "min_weight_fraction_leaf": min_weight_fraction_leaf}
+        parameters.update(kwargs)
+
+        dt_classifier = SKDecisionTreeClassifier(random_state=random_state,
+                                                 **parameters)
+        super().__init__(parameters=parameters,
+                         component_obj=dt_classifier,
+                         random_state=random_state)
diff --git a/evalml/pipelines/components/estimators/classifiers/rf_classifier.py b/evalml/pipelines/components/estimators/classifiers/rf_classifier.py
index 28f2996d..6f4b3d53 100644
--- a/evalml/pipelines/components/estimators/classifiers/rf_classifier.py
+++ b/evalml/pipelines/components/estimators/classifiers/rf_classifier.py
@@ -27,7 +27,3 @@ class RandomForestClassifier(Estimator):
         super().__init__(parameters=parameters,
                          component_obj=rf_classifier,
                          random_state=random_state)
-
-    @property
-    def feature_importance(self):
-        return self._component_obj.feature_importances_
diff --git a/evalml/pipelines/components/estimators/regressors/__init__.py b/evalml/pipelines/components/estimators/regressors/__init__.py
index e906c9e7..584c93a8 100644
--- a/evalml/pipelines/components/estimators/regressors/__init__.py
+++ b/evalml/pipelines/components/estimators/regressors/__init__.py
@@ -5,3 +5,4 @@ from .catboost_regressor import CatBoostRegressor
 from .xgboost_regressor import XGBoostRegressor
 from .et_regressor import ExtraTreesRegressor
 from .baseline_regressor import BaselineRegressor
+from .decision_tree_regressor import DecisionTreeRegressor
diff --git a/evalml/pipelines/components/estimators/regressors/decision_tree_regressor.py b/evalml/pipelines/components/estimators/regressors/decision_tree_regressor.py
new file mode 100644
index 00000000..13b4ed35
--- /dev/null
+++ b/evalml/pipelines/components/estimators/regressors/decision_tree_regressor.py
@@ -0,0 +1,39 @@
+from sklearn.tree import DecisionTreeRegressor as SKDecisionTreeRegressor
+from skopt.space import Integer
+
+from evalml.model_family import ModelFamily
+from evalml.pipelines.components.estimators import Estimator
+from evalml.problem_types import ProblemTypes
+
+
+class DecisionTreeRegressor(Estimator):
+    """Decision Tree Regressor."""
+    name = "Decision Tree Regressor."
+    hyperparameter_ranges = {
+        "criterion": ["mse", "friedman_mse", "mae"],
+        "max_features": ["auto", "sqrt", "log2"],
+        "max_depth": Integer(4, 10)
+    }
+    model_family = ModelFamily.DECISION_TREE
+    supported_problem_types = [ProblemTypes.REGRESSION]
+
+    def __init__(self,
+                 criterion="mse",
+                 max_features="auto",
+                 max_depth=6,
+                 min_samples_split=2,
+                 min_weight_fraction_leaf=0.0,
+                 random_state=0,
+                 **kwargs):
+        parameters = {"criterion": criterion,
+                      "max_features": max_features,
+                      "max_depth": max_depth,
+                      "min_samples_split": min_samples_split,
+                      "min_weight_fraction_leaf": min_weight_fraction_leaf}
+        parameters.update(kwargs)
+
+        dt_regressor = SKDecisionTreeRegressor(random_state=random_state,
+                                               **parameters)
+        super().__init__(parameters=parameters,
+                         component_obj=dt_regressor,
+                         random_state=random_state)
diff --git a/evalml/pipelines/components/estimators/regressors/rf_regressor.py b/evalml/pipelines/components/estimators/regressors/rf_regressor.py
index 769b4d97..9dc1a7b2 100644
--- a/evalml/pipelines/components/estimators/regressors/rf_regressor.py
+++ b/evalml/pipelines/components/estimators/regressors/rf_regressor.py
@@ -27,7 +27,3 @@ class RandomForestRegressor(Estimator):
         super().__init__(parameters=parameters,
                          component_obj=rf_regressor,
                          random_state=random_state)
-
-    @property
-    def feature_importance(self):
-        return self._component_obj.feature_importances_
diff --git a/evalml/utils/gen_utils.py b/evalml/utils/gen_utils.py
index 537ada3c..db0144b2 100644
--- a/evalml/utils/gen_utils.py
+++ b/evalml/utils/gen_utils.py
@@ -165,7 +165,7 @@ def _get_subclasses(base_class):
     return subclasses
 
 
-_not_used_in_automl = {'BaselineClassifier', 'BaselineRegressor',
+_not_used_in_automl = {'BaselineClassifier', 'BaselineRegressor', 'DecisionTreeClassifier', 'DecisionTreeRegressor',
                        'ModeBaselineBinaryPipeline', 'BaselineBinaryPipeline', 'MeanBaselineRegressionPipeline',
                        'BaselineRegressionPipeline', 'ModeBaselineMulticlassPipeline', 'BaselineMulticlassPipeline'}
 


Test patch:
diff --git a/evalml/tests/component_tests/test_decision_tree_classifier.py b/evalml/tests/component_tests/test_decision_tree_classifier.py
new file mode 100644
index 00000000..ae8ef172
--- /dev/null
+++ b/evalml/tests/component_tests/test_decision_tree_classifier.py
@@ -0,0 +1,62 @@
+import numpy as np
+from sklearn.tree import DecisionTreeClassifier as SKDecisionTreeClassifier
+
+from evalml.model_family import ModelFamily
+from evalml.pipelines import DecisionTreeClassifier
+from evalml.problem_types import ProblemTypes
+
+
+def test_model_family():
+    assert DecisionTreeClassifier.model_family == ModelFamily.DECISION_TREE
+
+
+def test_problem_types():
+    assert set(DecisionTreeClassifier.supported_problem_types) == {ProblemTypes.BINARY, ProblemTypes.MULTICLASS}
+
+
+def test_fit_predict_binary(X_y_binary):
+    X, y = X_y_binary
+
+    sk_clf = SKDecisionTreeClassifier(max_depth=6, max_features='auto', random_state=0)
+    sk_clf.fit(X, y)
+    y_pred_sk = sk_clf.predict(X)
+    y_pred_proba_sk = sk_clf.predict_proba(X)
+
+    clf = DecisionTreeClassifier()
+    clf.fit(X, y)
+    y_pred = clf.predict(X)
+    y_pred_proba = clf.predict_proba(X)
+
+    np.testing.assert_almost_equal(y_pred, y_pred_sk, decimal=5)
+    np.testing.assert_almost_equal(y_pred_proba, y_pred_proba_sk, decimal=5)
+
+
+def test_fit_predict_multi(X_y_multi):
+    X, y = X_y_multi
+
+    sk_clf = SKDecisionTreeClassifier(max_depth=6, max_features='auto', random_state=0)
+    sk_clf.fit(X, y)
+    y_pred_sk = sk_clf.predict(X)
+    y_pred_proba_sk = sk_clf.predict_proba(X)
+
+    clf = DecisionTreeClassifier()
+    clf.fit(X, y)
+    y_pred = clf.predict(X)
+    y_pred_proba = clf.predict_proba(X)
+
+    np.testing.assert_almost_equal(y_pred, y_pred_sk, decimal=5)
+    np.testing.assert_almost_equal(y_pred_proba, y_pred_proba_sk, decimal=5)
+
+
+def test_feature_importance(X_y_binary):
+    X, y = X_y_binary
+
+    clf = DecisionTreeClassifier()
+    sk_clf = SKDecisionTreeClassifier(max_depth=6, max_features='auto', random_state=0)
+    sk_clf.fit(X, y)
+    sk_feature_importance = sk_clf.feature_importances_
+
+    clf.fit(X, y)
+    feature_importance = clf.feature_importance
+
+    np.testing.assert_almost_equal(sk_feature_importance, feature_importance, decimal=5)
diff --git a/evalml/tests/component_tests/test_decision_tree_regressor.py b/evalml/tests/component_tests/test_decision_tree_regressor.py
new file mode 100644
index 00000000..e6e3ee24
--- /dev/null
+++ b/evalml/tests/component_tests/test_decision_tree_regressor.py
@@ -0,0 +1,42 @@
+import numpy as np
+from sklearn.tree import DecisionTreeRegressor as SKDecisionTreeRegressor
+
+from evalml.model_family import ModelFamily
+from evalml.pipelines import DecisionTreeRegressor
+from evalml.problem_types import ProblemTypes
+
+
+def test_model_family():
+    assert DecisionTreeRegressor.model_family == ModelFamily.DECISION_TREE
+
+
+def test_problem_types():
+    assert set(DecisionTreeRegressor.supported_problem_types) == {ProblemTypes.REGRESSION}
+
+
+def test_fit_predict(X_y_regression):
+    X, y = X_y_regression
+
+    sk_clf = SKDecisionTreeRegressor(max_depth=6, max_features='auto', random_state=0)
+    sk_clf.fit(X, y)
+    y_pred_sk = sk_clf.predict(X)
+
+    clf = DecisionTreeRegressor()
+    clf.fit(X, y)
+    y_pred = clf.predict(X)
+
+    np.testing.assert_almost_equal(y_pred, y_pred_sk, decimal=5)
+
+
+def test_feature_importance(X_y_regression):
+    X, y = X_y_regression
+
+    clf = DecisionTreeRegressor()
+    sk_clf = SKDecisionTreeRegressor(max_depth=6, max_features='auto', random_state=0)
+    sk_clf.fit(X, y)
+    sk_feature_importance = sk_clf.feature_importances_
+
+    clf.fit(X, y)
+    feature_importance = clf.feature_importance
+
+    np.testing.assert_almost_equal(sk_feature_importance, feature_importance, decimal=5)
diff --git a/evalml/tests/component_tests/test_en_classifier.py b/evalml/tests/component_tests/test_en_classifier.py
index 50cb91f7..aaa6b986 100644
--- a/evalml/tests/component_tests/test_en_classifier.py
+++ b/evalml/tests/component_tests/test_en_classifier.py
@@ -27,9 +27,7 @@ def test_en_parameters():
 
 
 def test_problem_types():
-    assert ProblemTypes.BINARY in ElasticNetClassifier.supported_problem_types
-    assert ProblemTypes.MULTICLASS in ElasticNetClassifier.supported_problem_types
-    assert len(ElasticNetClassifier.supported_problem_types) == 2
+    assert set(ElasticNetClassifier.supported_problem_types) == {ProblemTypes.BINARY, ProblemTypes.MULTICLASS}
 
 
 def test_fit_predict_binary(X_y_binary):
diff --git a/evalml/tests/component_tests/test_en_regressor.py b/evalml/tests/component_tests/test_en_regressor.py
index 53678aa6..af782b0c 100644
--- a/evalml/tests/component_tests/test_en_regressor.py
+++ b/evalml/tests/component_tests/test_en_regressor.py
@@ -24,8 +24,7 @@ def test_en_parameters():
 
 
 def test_problem_types():
-    assert ProblemTypes.REGRESSION in ElasticNetRegressor.supported_problem_types
-    assert len(ElasticNetRegressor.supported_problem_types) == 1
+    assert set(ElasticNetRegressor.supported_problem_types) == {ProblemTypes.REGRESSION}
 
 
 def test_fit_predict(X_y_regression):
diff --git a/evalml/tests/component_tests/test_et_classifier.py b/evalml/tests/component_tests/test_et_classifier.py
index 8e126e67..a089b658 100644
--- a/evalml/tests/component_tests/test_et_classifier.py
+++ b/evalml/tests/component_tests/test_et_classifier.py
@@ -11,9 +11,7 @@ def test_model_family():
 
 
 def test_problem_types():
-    assert ProblemTypes.BINARY in ExtraTreesClassifier.supported_problem_types
-    assert ProblemTypes.MULTICLASS in ExtraTreesClassifier.supported_problem_types
-    assert len(ExtraTreesClassifier.supported_problem_types) == 2
+    assert set(ExtraTreesClassifier.supported_problem_types) == {ProblemTypes.BINARY, ProblemTypes.MULTICLASS}
 
 
 def test_et_parameters():
diff --git a/evalml/tests/component_tests/test_et_regressor.py b/evalml/tests/component_tests/test_et_regressor.py
index a0a7f563..be4efa35 100644
--- a/evalml/tests/component_tests/test_et_regressor.py
+++ b/evalml/tests/component_tests/test_et_regressor.py
@@ -11,8 +11,7 @@ def test_model_family():
 
 
 def test_problem_types():
-    assert ProblemTypes.REGRESSION in ExtraTreesRegressor.supported_problem_types
-    assert len(ExtraTreesRegressor.supported_problem_types) == 1
+    assert set(ExtraTreesRegressor.supported_problem_types) == {ProblemTypes.REGRESSION}
 
 
 def test_et_parameters():
diff --git a/evalml/tests/component_tests/test_utils.py b/evalml/tests/component_tests/test_utils.py
index ad4b0574..6ea5e8bf 100644
--- a/evalml/tests/component_tests/test_utils.py
+++ b/evalml/tests/component_tests/test_utils.py
@@ -12,9 +12,9 @@ from evalml.pipelines.components.utils import (
 
 def test_all_components(has_minimal_dependencies):
     if has_minimal_dependencies:
-        assert len(all_components()) == 23
+        assert len(all_components()) == 25
     else:
-        assert len(all_components()) == 28
+        assert len(all_components()) == 30
 
 
 def test_handle_component_class_names():
diff --git a/evalml/tests/model_family_tests/test_model_family.py b/evalml/tests/model_family_tests/test_model_family.py
index c59e2ee1..a34f5bb4 100644
--- a/evalml/tests/model_family_tests/test_model_family.py
+++ b/evalml/tests/model_family_tests/test_model_family.py
@@ -5,12 +5,16 @@ from evalml.model_family import ModelFamily, handle_model_family
 
 @pytest.fixture
 def correct_model_families():
-    correct_model_families = [ModelFamily.LINEAR_MODEL, ModelFamily.LIGHTGBM, ModelFamily.RANDOM_FOREST, ModelFamily.XGBOOST, ModelFamily.CATBOOST, ModelFamily.EXTRA_TREES, ModelFamily.BASELINE, ModelFamily.NONE]
+    correct_model_families = [ModelFamily.LINEAR_MODEL, ModelFamily.LIGHTGBM,
+                              ModelFamily.RANDOM_FOREST, ModelFamily.XGBOOST,
+                              ModelFamily.CATBOOST, ModelFamily.EXTRA_TREES,
+                              ModelFamily.DECISION_TREE, ModelFamily.BASELINE, ModelFamily.NONE]
     yield correct_model_families
 
 
 def test_handle_string(correct_model_families):
-    model_families = ['linear_model', 'lightgbm', 'random_forest', 'xgboost', 'catboost', 'extra_trees', 'baseline', 'none']
+    model_families = ['linear_model', 'lightgbm', 'random_forest', 'xgboost',
+                      'catboost', 'extra_trees', 'decision_tree', 'baseline', 'none']
     for model_family in zip(model_families, correct_model_families):
         assert handle_model_family(model_family[0]) == model_family[1]
 
diff --git a/evalml/tests/pipeline_tests/test_pipelines.py b/evalml/tests/pipeline_tests/test_pipelines.py
index 59c7870f..0e241ad8 100644
--- a/evalml/tests/pipeline_tests/test_pipelines.py
+++ b/evalml/tests/pipeline_tests/test_pipelines.py
@@ -59,11 +59,8 @@ def test_allowed_model_families(has_minimal_dependencies):
     expected_model_families_binary = set(families)
     expected_model_families_regression = set(families)
     if not has_minimal_dependencies:
-        expected_model_families_binary.add(ModelFamily.XGBOOST)
-        expected_model_families_binary.add(ModelFamily.CATBOOST)
-        expected_model_families_binary.add(ModelFamily.LIGHTGBM)
-        expected_model_families_regression.add(ModelFamily.CATBOOST)
-        expected_model_families_regression.add(ModelFamily.XGBOOST)
+        expected_model_families_binary.update([ModelFamily.XGBOOST, ModelFamily.CATBOOST, ModelFamily.LIGHTGBM])
+        expected_model_families_regression.update([ModelFamily.CATBOOST, ModelFamily.XGBOOST])
     assert set(allowed_model_families(ProblemTypes.BINARY)) == expected_model_families_binary
     assert set(allowed_model_families(ProblemTypes.REGRESSION)) == expected_model_families_regression