Add stub

src/training/decision_tree/main.py

@@ -0,0 +1,71 @@
+import os
+
+from sklearn.ensemble import HistGradientBoostingRegressor, GradientBoostingRegressor, RandomForestRegressor
+from sklearn.model_selection import train_test_split, GroupShuffleSplit
+from sklearn.tree import DecisionTreeRegressor
+
+from src.training.decision_tree import preprocessor
+from src.training.decision_tree.model import Model
+from src.training.decision_tree.visualizer import Visualizer
+
+
+def create_split(x, y, game_ids):
+
+    splitter = GroupShuffleSplit(
+        n_splits=1,
+        test_size=0.10,
+        random_state=42
+    )
+
+    train_indices, test_indices = next(splitter.split(x, y, groups=game_ids))
+
+    x_train = x[train_indices]
+    x_test = x[test_indices]
+
+    y_train = y[train_indices]
+    y_test = y[test_indices]
+
+    train_games = set(game_ids[train_indices])
+    test_games = set(game_ids[test_indices])
+
+    print("Train positions:", len(x_train))
+    print("Test positions:", len(x_test))
+    print("Train games:", len(train_games))
+    print("Test games:", len(test_games))
+    print("Overlapping games:", len(train_games & test_games))
+
+    return x_train, x_test, y_train, y_test
+
+if __name__ == "__main__":
+
+    model_path = "C:/Users/Laurent/Desktop/chess-engine/src/training/decision_tree/cache/model.joblib"
+
+    if os.path.exists(model_path):
+        model = Model.load(model_path)
+
+    else:
+        x, y, game_ids = preprocessor.run("C:/Users/Laurent/Desktop/chess-engine/dataset/v1/output.csv")
+
+        split = create_split(x, y, game_ids)
+
+        depth = 35
+        leafs = int(depth * 3.5)
+
+        model = Model(
+            tree=RandomForestRegressor(
+                n_estimators=25,
+                max_depth=depth,
+                max_features=0.35,
+                min_samples_leaf=leafs,
+                random_state=42
+            ),
+            feature_names=preprocessor.get_feature_names(),
+            train_test_split=split
+        )
+
+        model.train()
+        model.evaluate()
+        #model.save_model(model_path)
+
+    visualizer = Visualizer(model=model)
+    visualizer.show_feature_importance(15)

src/training/decision_tree/model.py

@@ -0,0 +1,75 @@
+import os
+import joblib
+from sklearn.metrics import mean_absolute_error, mean_squared_error
+
+
+class Model:
+    def __init__(self, tree=None, feature_names=None, train_test_split=None):
+        self._model = tree
+        self._feature_names = feature_names
+        self._x_train = self._x_test = None
+        self._y_train = self._y_test = None
+
+        if train_test_split is not None:
+            self._x_train, self._x_test, \
+            self._y_train, self._y_test = train_test_split
+
+    @classmethod
+    def load(cls, path):
+        data = joblib.load(path)
+
+        model = cls()
+        model._model = data["model"]
+        model._feature_names = data["feature_names"]
+
+        print(f"Loaded model from {path}")
+
+        return model
+
+    def get_feature_names(self):
+        return self._feature_names
+
+    def get_feature_importance(self):
+        importances = self._model.feature_importances_
+
+        return sorted(
+            zip(self._feature_names, importances),
+            key=lambda x: x[1],
+            reverse=True
+        )
+
+    def train(self):
+        print("Start training...")
+        self._model.fit(self._x_train, self._y_train)
+        print("Training complete.")
+
+    def evaluate(self):
+        print("Evaluating model...")
+
+        train_predictions = self._model.predict(self._x_train)
+        test_predictions = self._model.predict(self._x_test)
+
+        train_mae = mean_absolute_error(self._y_train, train_predictions)
+        train_mse = mean_squared_error(self._y_train, train_predictions)
+
+        test_mae = mean_absolute_error(self._y_test, test_predictions)
+        test_mse = mean_squared_error(self._y_test, test_predictions)
+
+        print()
+        print("Training:")
+        print("MAE:", train_mae)
+        print("MSE:", train_mse)
+
+        print()
+        print("Test:")
+        print("MAE:", test_mae)
+        print("MSE:", test_mse)
+
+
+    def save_model(self, path):
+        joblib.dump({
+            "model": self._model,
+            "feature_names": self._feature_names
+        }, path)
+
+        print(f"Saved model to {path}")

src/training/decision_tree/visualizer.py

@@ -0,0 +1,50 @@
+import chess
+import numpy as np
+from matplotlib import pyplot as plt
+
+
+class Visualizer:
+    def __init__(self, model):
+        self._model = model
+
+    def show_feature_importance(self, top_n=15):
+        importances = self._model.get_feature_importance()
+        print("Feature importance")
+        print("-------------------------")
+
+        for feature_name, importance in importances[:top_n]:
+            print(f"{feature_name:<35} {importance:.6f}")
+
+    def show_piece_heatmaps(self):
+        importances = dict(self._model.get_feature_importance())
+
+        pieces = [
+            "P", "N", "B", "R", "Q", "K",
+            "p", "n", "b", "r", "q", "k"
+        ]
+
+        fig, axes = plt.subplots(3, 4, figsize=(16, 12))
+
+        for ax, piece in zip(axes.flat, pieces):
+            board = np.zeros((8, 8))
+
+            for row, rank in enumerate(range(7, -1, -1)):
+                for col, file in enumerate(range(8)):
+                    square = chess.square_name(chess.square(file, rank))
+                    feature_name = f"{piece}_{square}"
+                    board[row, col] = importances.get(feature_name, 0)
+
+            image = ax.imshow(board)
+
+            ax.set_xticks(range(8))
+            ax.set_xticklabels("abcdefgh")
+            ax.set_yticks(range(8))
+            ax.set_yticklabels("87654321")
+            ax.set_title(piece)
+
+            fig.colorbar(image, ax=ax, fraction=0.046, pad=0.04)
+
+        fig.suptitle("Piece-square feature importance", fontsize=16)
+        plt.tight_layout()
+        plt.show()
+        input()