OpenBB/openbb_terminal/forecast/theta_model.py

"""Theta Model"""
__docformat__ = "numpy"

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

import warnings
from statsmodels.tools.sm_exceptions import ConvergenceWarning
import numpy as np
import pandas as pd
from darts import TimeSeries
from darts.models import Theta
from darts.utils.utils import SeasonalityMode
from darts.metrics import mape

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

warnings.simplefilter("ignore", ConvergenceWarning)

SEASONS = ["N", "A", "M"]
PERIODS = [4, 5, 7]

logger = logging.getLogger(__name__)


@log_start_end(log=logger)
def get_theta_data(
    data: Union[pd.Series, pd.DataFrame],
    target_column: str = "close",
    seasonal: str = "M",
    seasonal_periods: int = 7,
    n_predict: int = 5,
    start_window: float = 0.85,
    forecast_horizon: int = 5,
) -> Tuple[
    Optional[List[TimeSeries]],
    Optional[List[TimeSeries]],
    Optional[List[TimeSeries]],
    Optional[float],
    Optional[float],
    Optional[type[Theta]],
]:

    """Performs Theta forecasting
    An implementation of the 4Theta method with configurable theta parameter.

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

    Parameters
    ----------
    data : Union[pd.Series, np.ndarray]
        Input data.
    target_column: Optional[str]:
        Target column to forecast. Defaults to "close".
    seasonal: str
        Seasonal component.  One of [N, A, M]
        Defaults to MULTIPLICATIVE.
    seasonal_periods: int
        Number of seasonal periods in a year
        If not set, inferred from frequency of the series.
    n_predict: int
        Number of days to forecast
    start_window: float
        Size of sliding window from start of timeseries and onwards
    forecast_horizon: int
        Number of days to forecast when backtesting and retraining historical

    Returns
    -------
    Tuple[List[TimeSeries], List[TimeSeries], List[TimeSeries], float, float, type[Theta]]
        Adjusted Data series,
        Historical forecast by best theta,
        list of Predictions,
        Mean average precision error,
        Best Theta,
        Theta Model.
    """

    use_scalers = False
    _, ticker_series = helpers.get_series(data, target_column, is_scaler=use_scalers)
    train, val = ticker_series.split_before(start_window)

    if seasonal == "A":
        seasonal = SeasonalityMode.ADDITIVE
    elif seasonal == "N":
        seasonal = SeasonalityMode.NONE
    else:  # Default
        seasonal = SeasonalityMode.MULTIPLICATIVE

    thetas = np.linspace(-10, 10, 50)
    best_mape = float("inf")
    best_theta = 0
    error = ""
    for theta in thetas:
        model = Theta(
            theta=theta,
            season_mode=seasonal,
            seasonality_period=seasonal_periods,
        )
        try:
            model.fit(train)
            pred_theta = model.predict(len(val))
            res = mape(val, pred_theta)
            if res < best_mape:
                best_mape = res
                best_theta = theta
        except Exception as e:  # noqa
            error = str(e)

    if best_theta == 0:
        console.print(f"[red]{error}[/red]")
        return [], [], [], 0, 0, None

    best_theta_model = Theta(
        best_theta,
        season_mode=seasonal,
        seasonality_period=seasonal_periods,
    )

    # Training model based on historical backtesting
    historical_fcast_theta = best_theta_model.historical_forecasts(
        ticker_series,
        start=float(start_window),
        forecast_horizon=int(forecast_horizon),
        verbose=True,
    )

    best_theta_model_final = Theta(
        best_theta,
        season_mode=seasonal,
        seasonality_period=seasonal_periods,
    )

    # fit model on entire series for final prediction
    best_theta_model_final.fit(ticker_series)
    prediction = best_theta_model_final.predict(int(n_predict))
    precision = mape(
        actual_series=ticker_series, pred_series=historical_fcast_theta
    )  # mape = mean average precision error
    console.print(f"Theta Model obtains MAPE: {precision:.2f}% \n")

    return (
        ticker_series,
        historical_fcast_theta,
        prediction,
        precision,
        best_theta,
        best_theta_model,
    )