catboost_algorithm.py

from catboost import CatBoostClassifier
import numpy as np
from .base import BaseAlgorithm

class CatBoostAlgorithm(BaseAlgorithm):
    """Implementação do CatBoost para classificação."""
    
    def __init__(
        self,
        iterations=1000,
        learning_rate=0.1,
        depth=6,
        l2_leaf_reg=3,
        random_strength=1,
        bagging_temperature=1,
        border_count=254,
        random_state=None,
        early_stopping_rounds=None,
        task_type='CPU',
        thread_count=-1,
        verbose=100
    ):
        self.iterations = iterations
        self.learning_rate = learning_rate
        self.depth = depth
        self.l2_leaf_reg = l2_leaf_reg
        self.random_strength = random_strength
        self.bagging_temperature = bagging_temperature
        self.border_count = border_count
        self.random_state = random_state
        self.early_stopping_rounds = early_stopping_rounds
        self.task_type = task_type
        self.thread_count = thread_count
        self.verbose = verbose
        
        self.model = CatBoostClassifier(
            iterations=iterations,
            learning_rate=learning_rate,
            depth=depth,
            l2_leaf_reg=l2_leaf_reg,
            random_strength=random_strength,
            bagging_temperature=bagging_temperature,
            border_count=border_count,
            random_seed=random_state,
            task_type=task_type,
            thread_count=thread_count,
            verbose=verbose
        )
    
    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
            self.model.fit(
                X_train, y_train,
                eval_set=(X_val, y_val),
                early_stopping_rounds=self.early_stopping_rounds,
                verbose=self.verbose
            )
        else:
            self.model.fit(X, y, verbose=self.verbose)
        
        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 {
            'iterations': self.iterations,
            'learning_rate': self.learning_rate,
            'depth': self.depth,
            'l2_leaf_reg': self.l2_leaf_reg,
            'random_strength': self.random_strength,
            'bagging_temperature': self.bagging_temperature,
            'border_count': self.border_count,
            'random_state': self.random_state,
            'early_stopping_rounds': self.early_stopping_rounds,
            'task_type': self.task_type,
            'thread_count': self.thread_count,
            'verbose': self.verbose
        }
    
    def get_feature_importance(self, type='FeatureImportance'):
        """Retorna a importância das features."""
        return self.model.get_feature_importance(type=type)
    
    def plot_feature_importance(self, feature_names=None, top_k=None):
        """Plota a importância das features."""
        self.model.plot_feature_importance(feature_names, top_k)
    
    def plot_learning_curves(self):
        """Plota as curvas de aprendizado do modelo."""
        self.model.plot_metrics()
    
    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)
    
    def get_best_iteration(self):
        """Retorna a melhor iteração se early stopping foi usado."""
        return self.model.get_best_iteration()
    
    def get_best_score(self):
        """Retorna o melhor score se early stopping foi usado."""
        return self.model.get_best_score()
    
    def get_evals_result(self):
        """Retorna os resultados da avaliação durante o treinamento."""
        return self.model.get_evals_result()
    
    def get_model_params(self):
        """Retorna os parâmetros internos do modelo."""
        return self.model.get_params()
    
    def shrink(self, ntree_start=0, ntree_end=0):
        """Reduz o modelo para usar apenas as árvores especificadas."""
        self.model.shrink(ntree_start, ntree_end)
        return self