mirror of
https://github.com/OpenBB-finance/OpenBB.git
synced 2026-10-30 22:11:55 +00:00
6426b76771
* decorated functions and removed assertion tests that are now caught by decorator logs * modified functions with @property * fixed bc of coingecko * fixed outedated tests
77 lines
1.8 KiB
Python
77 lines
1.8 KiB
Python
"""Arima Prediction Model"""
|
|
__docformat__ = "numpy"
|
|
|
|
import logging
|
|
from typing import Any, List, Tuple, Union
|
|
|
|
import pandas as pd
|
|
import pmdarima
|
|
from statsmodels.tsa.arima.model import ARIMA
|
|
|
|
from gamestonk_terminal.decorators import log_start_end
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
ICS = ["aic", "aicc", "bic", "hqic", "oob"]
|
|
|
|
|
|
@log_start_end(log=logger)
|
|
def get_arima_model(
|
|
values: Union[pd.Series, pd.DataFrame],
|
|
arima_order: str,
|
|
n_predict: int,
|
|
seasonal: bool,
|
|
ic: str,
|
|
) -> Tuple[List[float], Any]:
|
|
"""Get an ARIMA model for data
|
|
|
|
Parameters
|
|
----------
|
|
values : Union[pd.Series, pd.DataFrame]
|
|
Data to fit
|
|
arima_order : str
|
|
String of ARIMA params in form "p,q,d"
|
|
n_predict : int
|
|
Days to predict
|
|
seasonal : bool
|
|
Flag to use seasonal model
|
|
ic : str
|
|
Information Criteria for model evaluation
|
|
|
|
Returns
|
|
-------
|
|
List[float]
|
|
List of predicted values
|
|
Any
|
|
Fit ARIMA model object.
|
|
"""
|
|
if arima_order:
|
|
model = ARIMA(
|
|
values, order=tuple(int(ord) for ord in arima_order.split(","))
|
|
).fit()
|
|
l_predictions = list(
|
|
model.predict(
|
|
start=len(values.values) + 1,
|
|
end=len(values.values) + n_predict,
|
|
)
|
|
)
|
|
else:
|
|
if seasonal:
|
|
model = pmdarima.auto_arima(
|
|
values.values,
|
|
error_action="ignore",
|
|
seasonal=True,
|
|
m=5,
|
|
information_criteria=ic,
|
|
)
|
|
else:
|
|
model = pmdarima.auto_arima(
|
|
values.values,
|
|
error_action="ignore",
|
|
seasonal=False,
|
|
information_criteria=ic,
|
|
)
|
|
l_predictions = list(model.predict(n_predict))
|
|
|
|
return l_predictions, model
|