Resolve "Adding new datasets"

Closes #17 (closed)

code/bolsonaro/models/kmeans_forest_regressor.py

+from bolsonaro.utils import tqdm_joblib
+
+from sklearn.ensemble import RandomForestRegressor
+from sklearn.metrics import mean_squared_error
+from sklearn.base import BaseEstimator
+from sklearn.cluster import KMeans
+from abc import abstractmethod, ABCMeta
+import numpy as np
+from scipy.stats import mode
+from joblib import Parallel, delayed
+from tqdm import tqdm
+
+
+class KMeansForestRegressor(BaseEstimator, metaclass=ABCMeta):
+    """
+    On extreme pruning of random forest ensembles for ral-time predictive applications', by Khaled Fawagreh, Mohamed Medhat Gaber and Eyad Elyan.
+    """
+
+    def __init__(self, models_parameters, score_metric=mean_squared_error):
+        self._models_parameters = models_parameters
+        self._estimator = RandomForestRegressor(**self._models_parameters.hyperparameters,
+            random_state=self._models_parameters.seed, n_jobs=-1)
+        self._extracted_forest_size = self._models_parameters.extracted_forest_size
+        self._score_metric = score_metric
+
+    @property
+    def models_parameters(self):
+        return self._models_parameters
+
+    def fit(self, X_train, y_train, X_val, y_val):
+        self._estimator.fit(X_train, y_train)
+
+        predictions = list()
+        for tree in self._estimator.estimators_:
+            predictions.append(tree.predict(X_train))
+        predictions = np.array(predictions)
+
+        kmeans = KMeans(n_clusters=self._extracted_forest_size, random_state=self._models_parameters.seed).fit(predictions)
+        labels = np.array(kmeans.labels_)
+
+        # For each cluster select the best tree on the validation set
+        extracted_forest_sizes = list(range(self._extracted_forest_size))
+        with tqdm_joblib(tqdm(total=self._extracted_forest_size, disable=True)) as prune_forest_job_pb:
+            pruned_forest = Parallel(n_jobs=-1)(delayed(self._prune_forest_job)(prune_forest_job_pb,
+                extracted_forest_sizes[i], labels, X_val, y_val, self._score_metric)
+                for i in range(self._extracted_forest_size))
+
+        self._estimator.estimators_ = pruned_forest
+
+    def _prune_forest_job(self, prune_forest_job_pb, c, labels, X_val, y_val, score_metric):
+        index = np.where(labels == c)[0]
+        with tqdm_joblib(tqdm(total=len(index), disable=True)) as cluster_job_pb:
+            cluster = Parallel(n_jobs=-1)(delayed(self._cluster_job)(cluster_job_pb, index[i], X_val, 
+                y_val, score_metric) for i in range(len(index)))
+        best_tree_index = np.argmax(cluster)
+        prune_forest_job_pb.update()
+        return self._estimator.estimators_[index[best_tree_index]]
+
+    def _cluster_job(self, cluster_job_pb, i, X_val, y_val, score_metric):
+        y_val_pred = self._estimator.estimators_[i].predict(X_val)
+        tree_pred = score_metric(y_val, y_val_pred)
+        cluster_job_pb.update()
+        return tree_pred
+
+    def predict(self, X):
+        return self._estimator.predict(X)
+
+    def score(self, X, y):
+        predictions = list()
+        for tree in self._estimator.estimators_:
+            predictions.append(tree.predict(X))
+        predictions = np.array(predictions)
+        mean_predictions = np.mean(predictions, axis=0)
+        score = self._score_metric(mean_predictions, y)
+        return score
+
+    def predict_base_estimator(self, X):
+        return self._estimator.predict(X)