src/forecaster/models/xgbregressor.py

import logging

from xgboost import XGBRegressor as xgbreg
import numpy as np

from .base import ModelBase


class XGBRegressor(ModelBase):
    def __init__(
            self,
            **kwargs
    ) -> None:

        logging.debug('Init XGBRegressor')

        regressor = xgbreg(
            objective='reg:squarederror',
            random_state=42)

        param_grid = {
            'estimator__max_depth': [3, 5, 6, 10, 15, 20],
            'estimator__learning_rate': [0.01, 0.1, 0.2, 0.3],
            'estimator__subsample': np.arange(0.5, 1.0, 0.1),
            'estimator__colsample_bytree': np.arange(0.4, 1.0, 0.1),
            'estimator__colsample_bylevel': np.arange(0.4, 1.0, 0.1),
            'estimator__n_estimators': [100, 500, 1000]
        }

        super().__init__(
            regressor,
            param_grid,
            **kwargs)

    # def forecast(self):
    #     logging.debug('XGBRegressor forecast')

    #     return {'forecast': self.forecaster.predict(**self.predict_args),
    #             'evaluate': self.evaluate()}