xgboost_algorithm.py

import xgboost as xgb
import numpy as np
from .base import BaseAlgorithm

class XGBoostAlgorithm(BaseAlgorithm):
    """Implementação do XGBoost para classificação."""
    
    def __init__(
        self,
        n_estimators=100,
        max_depth=3,
        learning_rate=0.1,
        gamma=0,
        min_child_weight=1,
        subsample=1.0,
        colsample_bytree=1.0,
        reg_alpha=0,
        reg_lambda=1,
        random_state=None,
        early_stopping_rounds=None
    ):
        self.n_estimators = n_estimators
        self.max_depth = max_depth
        self.learning_rate = learning_rate
        self.gamma = gamma
        self.min_child_weight = min_child_weight
        self.subsample = subsample
        self.colsample_bytree = colsample_bytree
        self.reg_alpha = reg_alpha
        self.reg_lambda = reg_lambda
        self.random_state = random_state
        self.early_stopping_rounds = early_stopping_rounds
        
        self.model = xgb.XGBClassifier(
            n_estimators=n_estimators,
            max_depth=max_depth,
            learning_rate=learning_rate,
            gamma=gamma,
            min_child_weight=min_child_weight,
            subsample=subsample,
            colsample_bytree=colsample_bytree,
            reg_alpha=reg_alpha,
            reg_lambda=reg_lambda,
            random_state=random_state
        )
    
    def fit(self, X, y):
        """Treina o modelo com early stopping opcional."""
        if self.early_stopping_rounds is not None:
            # Divide dados em treino e validação
            n_samples = len(X)
            n_val = int(0.2 * n_samples)
            indices = np.random.permutation(n_samples)
            val_idx = indices[:n_val]
            train_idx = indices[n_val:]
            
            X_train, X_val = X[train_idx], X[val_idx]
            y_train, y_val = y[train_idx], y[val_idx]
            
            # Treina com early stopping
            eval_set = [(X_val, y_val)]
            self.model.fit(
                X_train, y_train,
                eval_set=eval_set,
                early_stopping_rounds=self.early_stopping_rounds,
                eval_metric=['error', 'auc'],
                verbose=True
            )
        else:
            self.model.fit(X, y)
        
        return self
    
    def predict(self, X):
        """Faz previsões usando o modelo."""
        return self.model.predict(X)
    
    def predict_proba(self, X):
        """Retorna probabilidades das previsões."""
        return self.model.predict_proba(X)
    
    def get_params(self):
        """Retorna os parâmetros do modelo."""
        return {
            'n_estimators': self.n_estimators,
            'max_depth': self.max_depth,
            'learning_rate': self.learning_rate,
            'gamma': self.gamma,
            'min_child_weight': self.min_child_weight,
            'subsample': self.subsample,
            'colsample_bytree': self.colsample_bytree,
            'reg_alpha': self.reg_alpha,
            'reg_lambda': self.reg_lambda,
            'random_state': self.random_state,
            'early_stopping_rounds': self.early_stopping_rounds
        }
    
    def get_feature_importance(self, importance_type='weight'):
        """Retorna a importância das features."""
        return self.model.get_booster().get_score(importance_type=importance_type)
    
    def plot_feature_importance(self, feature_names=None, importance_type='weight'):
        """Plota a importância das features."""
        import matplotlib.pyplot as plt
        
        importance = self.get_feature_importance(importance_type)
        if not importance:
            print("Não foi possível obter importância das features.")
            return
        
        if feature_names is None:
            feature_names = list(importance.keys())
        
        values = [importance.get(f, 0) for f in feature_names]
        
        plt.figure(figsize=(10, 6))
        plt.bar(range(len(values)), values)
        plt.xticks(range(len(values)), feature_names, rotation=45)
        plt.title(f'Feature Importance ({importance_type})')
        plt.tight_layout()
        plt.show()
    
    def plot_learning_curves(self):
        """Plota as curvas de aprendizado do modelo."""
        results = self.model.evals_result()
        if not results:
            print("Não há resultados de validação disponíveis.")
            return
        
        import matplotlib.pyplot as plt
        
        plt.figure(figsize=(12, 5))
        
        # Plot error
        plt.subplot(1, 2, 1)
        for eval_set in results.keys():
            plt.plot(results[eval_set]['error'], label=eval_set)
        plt.xlabel('Iteration')
        plt.ylabel('Classification Error')
        plt.title('XGBoost Classification Error')
        plt.legend()
        
        # Plot AUC
        plt.subplot(1, 2, 2)
        for eval_set in results.keys():
            plt.plot(results[eval_set]['auc'], label=eval_set)
        plt.xlabel('Iteration')
        plt.ylabel('AUC')
        plt.title('XGBoost AUC')
        plt.legend()
        
        plt.tight_layout()
        plt.show()
    
    def save_model(self, filepath):
        """Salva o modelo em formato binário."""
        self.model.save_model(filepath)
    
    def load_model(self, filepath):
        """Carrega o modelo de um arquivo binário."""
        self.model.load_model(filepath)