# Source: content/notes/ml/trees-and-ensembles.md
# Independent CPU example; use the curriculum environment.
# See /notes/ml/#example-environment or /notes/deep-learning/#example-environment.

import numpy as np
from sklearn.datasets import make_friedman1
from sklearn.dummy import DummyRegressor
from sklearn.ensemble import GradientBoostingRegressor, HistGradientBoostingRegressor
from sklearn.metrics import mean_squared_error
from sklearn.model_selection import train_test_split

X, y = make_friedman1(n_samples=600, n_features=7, noise=1.0, random_state=6)
X_train, X_test, y_train, y_test = train_test_split(X, y, random_state=6)
models = {
    "gradient": GradientBoostingRegressor(n_estimators=160, learning_rate=0.05,
        max_depth=2, validation_fraction=0.2, n_iter_no_change=10, random_state=6),
    "histogram": HistGradientBoostingRegressor(max_iter=160, learning_rate=0.08,
        max_leaf_nodes=15, early_stopping=True, random_state=6),
}
baseline = DummyRegressor().fit(X_train, y_train).predict(X_test)
for name, model in models.items():
    model.fit(X_train, y_train)
    pred = model.predict(X_test)
    assert np.isfinite(pred).all()
    assert mean_squared_error(y_test, pred) < mean_squared_error(y_test, baseline)
    print(name, "RMSE", np.sqrt(mean_squared_error(y_test, pred)))