OpenBB/openbb_terminal/forecast/expo_model.py

# pylint: disable=too-many-arguments
"""Probabilistic Exponential Smoothing Model"""
__docformat__ = "numpy"

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

import warnings
from statsmodels.tools.sm_exceptions import ConvergenceWarning
from numpy import ndarray
import pandas as pd
from darts import TimeSeries
from darts.models import ExponentialSmoothing
from darts.utils.utils import ModelMode, 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)


TRENDS = ["N", "A", "M"]
SEASONS = ["N", "A", "M"]
PERIODS = [4, 5, 7]
DAMPEN = ["T", "F"]

logger = logging.getLogger(__name__)


@log_start_end(log=logger)
def get_expo_data(
    data: Union[pd.Series, pd.DataFrame],
    target_column: str = "close",
    trend: str = "A",
    seasonal: str = "A",
    seasonal_periods: int = 7,
    dampen: str = "F",
    n_predict: int = 5,
    start_window: float = 0.85,
    forecast_horizon: int = 5,
) -> Tuple[
    List[TimeSeries],
    List[TimeSeries],
    List[TimeSeries],
    Optional[Union[float, ndarray]],
    ExponentialSmoothing,
]:

    """Performs Probabilistic Exponential Smoothing forecasting
    This is a wrapper around statsmodels Holt-Winters' Exponential Smoothing;
    we refer to this link for the original and more complete documentation of the parameters.

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

    Parameters
    ----------
    data : Union[pd.Series, np.ndarray]
        Input data.
    target_column: Optional[str]:
        Target column to forecast. Defaults to "close".
    trend: str
        Trend component.  One of [N, A, M]
        Defaults to ADDITIVE.
    seasonal: str
        Seasonal component.  One of [N, A, M]
        Defaults to ADDITIVE.
    seasonal_periods: int
        Number of seasonal periods in a year (7 for daily data)
        If not set, inferred from frequency of the series.
    dampen: str
        Dampen the function
    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], Optional[Union[float, ndarray]], ExponentialSmoothing]
        Adjusted Data series,
        List of historical fcast values,
        List of predicted fcast values,
        Optional[float] - precision,
        Fit Prob. Expo model object.
    """

    use_scalers = False
    _, ticker_series = helpers.get_series(data, target_column, is_scaler=use_scalers)

    if trend == "M":
        trend_model = ModelMode.MULTIPLICATIVE
    elif trend == "N":
        trend_model = ModelMode.NONE
    else:  # Default
        trend_model = ModelMode.ADDITIVE

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

    damped = True
    if dampen == "F":
        damped = False

    # Model Init
    model_es = ExponentialSmoothing(
        trend=trend_model,
        seasonal=seasonal_model,
        seasonal_periods=int(seasonal_periods),
        damped=damped,
        random_state=42,
    )

    try:
        # Historical backtesting
        historical_fcast_es = model_es.historical_forecasts(
            ticker_series,  # backtest on entire ts
            start=float(start_window),
            forecast_horizon=int(forecast_horizon),
            verbose=True,
        )
    except Exception as e:  # noqa
        error = str(e)
        # lets translate this to something everyone understands
        if "with`overlap_end` set to `False`." in error:
            console.print(
                "[red]Dataset too small.[/red]"
                "[red] Please increase size to at least 100 data points.[/red]"
            )
        else:
            console.print(f"[red]{error}[/red]")
        return [], [], [], None, None

    # train new model on entire timeseries to provide best current forecast
    best_model = ExponentialSmoothing(
        trend=trend_model,
        seasonal=seasonal_model,
        seasonal_periods=int(seasonal_periods),
        damped=damped,
        random_state=42,
    )

    # we have the historical fcast, now lets train on entire set and predict.
    best_model.fit(ticker_series)
    probabilistic_forecast = best_model.predict(int(n_predict), num_samples=500)
    precision = mape(actual_series=ticker_series, pred_series=historical_fcast_es)
    console.print(f"Exponential smoothing obtains MAPE: {precision:.2f}% \n")

    return (
        ticker_series,
        historical_fcast_es,
        probabilistic_forecast,
        precision,
        best_model,
    )