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
[docs]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,
ProblemTypes.TIME_SERIES_REGRESSION,
]
[docs] def __init__(
self,
criterion="mse",
max_features="auto",
max_depth=6,
min_samples_split=2,
min_weight_fraction_leaf=0.0,
random_seed=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_seed, **parameters)
super().__init__(
parameters=parameters, component_obj=dt_regressor, random_seed=random_seed
)