# 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_classification
from sklearn.ensemble import ExtraTreesClassifier, RandomForestClassifier
from sklearn.inspection import permutation_importance
from sklearn.metrics import balanced_accuracy_score, log_loss
from sklearn.model_selection import train_test_split

X, y = make_classification(n_samples=700, n_features=10, n_informative=5,
                           n_redundant=2, random_state=13)
X_train, X_test, y_train, y_test = train_test_split(
    X, y, test_size=0.25, stratify=y, random_state=13)
forest = RandomForestClassifier(n_estimators=100, min_samples_leaf=3,
    max_features="sqrt", oob_score=True, random_state=13, n_jobs=1)
forest.fit(X_train, y_train)
pred = forest.predict(X_test)
prob = forest.predict_proba(X_test)
manual = np.mean([t.predict_proba(X_test) for t in forest.estimators_], axis=0)
assert np.allclose(prob, manual)
assert np.isfinite(forest.oob_score_)
print("OOB accuracy:", forest.oob_score_)
print("Test balanced accuracy:", balanced_accuracy_score(y_test, pred))
print("Test log loss:", log_loss(y_test, prob))
importance = permutation_importance(forest, X_test, y_test,
    scoring="balanced_accuracy", n_repeats=3, random_state=13, n_jobs=1)
print("Permutation importance:", importance.importances_mean.round(3))
extra = ExtraTreesClassifier(n_estimators=80, min_samples_leaf=3,
    random_state=13, n_jobs=1).fit(X_train, y_train)
extra_pred = extra.predict(X_test)
assert extra_pred.shape == y_test.shape
print("Extra Trees accuracy:", balanced_accuracy_score(y_test, extra_pred))