OpenBB/openbb_terminal/forecast/tft_model.py

# pylint: disable=too-many-arguments
"""Temporal Fusion Transformer Model"""
__docformat__ = "numpy"

import logging
from typing import Tuple, Union, List, Optional

import warnings
from statsmodels.tools.sm_exceptions import ConvergenceWarning
import pandas as pd
from darts import TimeSeries
from darts.models import TFTModel
from darts.utils.likelihood_models import QuantileRegression

from openbb_terminal.decorators import log_start_end
from openbb_terminal.forecast import helpers


warnings.simplefilter("ignore", ConvergenceWarning)


logger = logging.getLogger(__name__)


@log_start_end(log=logger)
def get_tft_data(
    data: Union[pd.Series, pd.DataFrame],
    target_column: str = "close",
    n_predict: int = 5,
    past_covariates: str = None,
    train_split: float = 0.85,
    forecast_horizon: int = 5,
    input_chunk_length: int = 14,
    output_chunk_length: int = 5,
    hidden_size: int = 16,
    lstm_layers: int = 1,
    num_attention_heads: int = 4,
    full_attention: bool = False,
    dropout: float = 0.1,
    hidden_continuous_size: int = 8,
    n_epochs: int = 200,
    batch_size: int = 32,
    model_save_name: str = "tft_model",
    force_reset: bool = True,
    save_checkpoints: bool = True,
) -> Tuple[
    Optional[List[TimeSeries]],
    Optional[List[TimeSeries]],
    Optional[List[TimeSeries]],
    Optional[float],
    Optional[type[TFTModel]],
]:

    """Performs Temporal Fusion Transformer forecasting
    The TFT applies multi-head attention queries on future inputs from mandatory future_covariates.
    Specifying future encoders with add_encoders (read below) can automatically generate future
    covariates and allows to use the model without having to pass any future_covariates to fit()
    and predict().

    https://unit8co.github.io/darts/generated_api/darts.models.forecasting.tft_model.html

    Parameters
    ----------
    data (Union[pd.Series, pd.DataFrame]):
        Input Data
    target_column: Optional[str]
        Target column to forecast. Defaults to "close".
    n_predict: (int, optional)
        Days to predict. Defaults to 5.
    train_split: (float, optional)
        Train/val split. Defaults to 0.85.
    past_covariates: (str, optional)
        Multiple secondary columns to factor in when forecasting. Defaults to None.
    forecast_horizon: (int, optional)
        Forecast horizon when performing historical forecasting. Defaults to 5.
    input_chunk_length: (int, optional)
        Number of past time steps that are fed to the forecasting module at prediction time.
        Defaults to 14.
    output_chunk_length: (int, optional)
        The length of the forecast of the model. Defaults to 5.
    hidden_size: (int, optional)
        Hidden state size of the TFT. Defaults to 16.
    lstm_layers: (int, optional)
        Number of layers for the Long Short Term Memory Encoder and Decoder. Defaults to 16.
    num_attention_headers: (int, optional)
        Number of attention heads. Defaults to 4.
    full_attention: (bool, optional)
        Whether to apply a multi-head attention query. Defaults to False>
    dropout: (float, optional)
        Fraction of neurons affected by dropout. Defaults to 0.1.
    hidden_continuous_size: (int, optional)
        Default hidden size for processing continuous variables. Defaults to 8.
    n_epochs: (int, optional)
        Number of epochs to run during training. Defaults to 200.
    batch_size: (int, optional)
        Number of samples to pass through network during a single epoch. Defaults to 32.
    model_save_name: (str, optional)
        The name for the model. Defaults to tft_model
    force_reset: (bool, optional)
        If set to True, any previously-existing model with the same name will be reset
        (all checkpoints will be discarded). Defaults to True.
    save_checkpoints: (bool, optional)
        Whether or not to automatically save the untrained model and checkpoints from training.
        Defaults to True.

    Returns
    -------
    Tuple[List[TimeSeries], List[TimeSeries], List[TimeSeries], Optional[float], type[TFTModel]]:
        Adjusted Data series,
        List of historical fcast values,
        List of predicted fcast values,
        Optional[float] - precision,
        Fit Prob. TFT model object.
    """

    use_scalers = True
    probabilistic = True

    scaler, ticker_series = helpers.get_series(
        data, target_column, is_scaler=use_scalers
    )
    train, val = ticker_series.split_before(train_split)
    valid = helpers.check_data_length(
        train, val, input_chunk_length, output_chunk_length
    )
    if not valid:
        return [], [], [], None, None

    (
        past_covariate_whole,
        past_covariate_train,
        past_covariate_val,
    ) = helpers.past_covs(past_covariates, data, train_split, use_scalers)

    quantiles = [
        0.01,
        0.05,
        0.1,
        0.15,
        0.2,
        0.25,
        0.3,
        0.4,
        0.5,
        0.6,
        0.7,
        0.75,
        0.8,
        0.85,
        0.9,
        0.95,
        0.99,
    ]

    tft_model = TFTModel(
        input_chunk_length=input_chunk_length,
        output_chunk_length=output_chunk_length,
        hidden_size=hidden_size,
        lstm_layers=lstm_layers,
        num_attention_heads=num_attention_heads,
        full_attention=full_attention,
        dropout=dropout,
        hidden_continuous_size=hidden_continuous_size,
        model_name=model_save_name,
        force_reset=force_reset,
        save_checkpoints=save_checkpoints,
        random_state=42,
        n_epochs=n_epochs,
        batch_size=batch_size,
        pl_trainer_kwargs=helpers.get_pl_kwargs(accelerator="cpu"),
        likelihood=QuantileRegression(
            quantiles=quantiles
        ),  # QuantileRegression is set per default
        add_relative_index=True,  # TODO There is a bug with this. Must fix. Should be false
    )

    # fit model on train series for historical forecasting
    with warnings.catch_warnings():
        warnings.simplefilter("ignore")
        helpers.fit_model(
            tft_model,
            train,
            val,
            past_covariate_train,
            past_covariate_val,
        )
    best_model = TFTModel.load_from_checkpoint(model_name=model_save_name, best=True)

    # Showing historical backtesting without retraining model (too slow)
    return helpers.get_prediction(
        "TFT",
        probabilistic,
        use_scalers,
        scaler,
        past_covariates,
        best_model,
        ticker_series,
        past_covariate_whole,
        train_split,
        forecast_horizon,
        n_predict,
    )