Files
OpenBB/prediction_techniques/arima.py
2021-02-13 23:46:54 +00:00

79 lines
4.3 KiB
Python

import argparse
from helper_funcs import *
import pandas as pd
import matplotlib.pyplot as plt
from pandas.plotting import register_matplotlib_converters
register_matplotlib_converters()
import pmdarima
from statsmodels.tsa.arima.model import ARIMA
# -------------------------------------------------- ARIMA --------------------------------------------------
def arima(l_args, s_ticker, s_interval, df_stock):
parser = argparse.ArgumentParser(prog='arima',
description="""In statistics and econometrics, and in particular in time
series analysis, an autoregressive integrated moving average (ARIMA) model
is a generalization of an autoregressive moving average (ARMA) model. Both
of these models are fitted to time series data either to better understand
the data or to predict future points in the series (forecasting). """)
parser.add_argument('-d', "--days", action="store", dest="n_days", type=check_positive, default=5, help='prediction days')
parser.add_argument('-i', "--ic", action="store", dest="s_ic", type=str, default='aic', choices=['aic', 'aicc', 'bic', 'hqic', 'oob'], help='information criteria')
parser.add_argument('-s', "--seasonal", action="store_true", default=False, dest="b_seasonal", help='Use the fact that the market data is weekly')
parser.add_argument('-o', "--order", action="store", dest="s_order", type=str, help='arima model order')
parser.add_argument('-r', "--results", action="store_true", dest="b_results", default=False, help='results about ARIMA summary')
try:
(ns_parser, l_unknown_args) = parser.parse_known_args(l_args)
if l_unknown_args:
print(f"The following args couldn't be interpreted: {l_unknown_args}\n")
return
# Machine Learning model
if ns_parser.s_order:
t_order = tuple([int(ord) for ord in list(ns_parser.s_order)])
model = ARIMA(df_stock['5. adjusted close'].values, order=t_order).fit()
l_predictions = model.predict(start=len(df_stock['5. adjusted close'])+1, end=len(df_stock['5. adjusted close'])+ns_parser.n_days)
else:
if ns_parser.b_seasonal:
model = pmdarima.auto_arima(df_stock['5. adjusted close'].values, error_action='ignore', seasonal=True, m=5, information_criteria=ns_parser.s_ic)
else:
model = pmdarima.auto_arima(df_stock['5. adjusted close'].values, error_action='ignore', seasonal=False, information_criteria=ns_parser.s_ic)
l_predictions = model.predict(n_periods=ns_parser.n_days)
# Prediction data
l_pred_days = get_next_stock_market_days(last_stock_day=df_stock['5. adjusted close'].index[-1], n_next_days=ns_parser.n_days)
df_pred = pd.Series(l_predictions, index=l_pred_days, name='Price')
if ns_parser.b_results:
print(model.summary())
print("")
# Plotting
plt.plot(df_stock.index, df_stock['5. adjusted close'], lw=2)
if ns_parser.s_order:
plt.title(f"ARIMA {str(t_order)} on {s_ticker} - {ns_parser.n_days} days prediction")
else:
plt.title(f"ARIMA {model.order} on {s_ticker} - {ns_parser.n_days} days prediction")
plt.xlim(df_stock.index[0], get_next_stock_market_days(df_pred.index[-1], 1)[-1] )
plt.xlabel('Time')
plt.ylabel('Share Price ($)')
plt.grid(b=True, which='major', color='#666666', linestyle='-')
plt.minorticks_on()
plt.grid(b=True, which='minor', color='#999999', linestyle='-', alpha=0.2)
plt.plot([df_stock.index[-1], df_pred.index[0]], [df_stock['5. adjusted close'].values[-1], df_pred.values[0]], lw=1, c='tab:green', linestyle='--')
plt.plot(df_pred.index, df_pred, lw=2, c='tab:green')
plt.axvspan(df_stock.index[-1], df_pred.index[-1], facecolor='tab:orange', alpha=0.2)
xmin, xmax, ymin, ymax = plt.axis()
plt.vlines(df_stock.index[-1], ymin, ymax, linewidth=1, linestyle='--', color='k')
plt.show()
# Print prediction data
print("Predicted share price:")
df_pred = df_pred.apply(lambda x: f"{x:.2f} $")
print(df_pred.to_string())
print("")
except:
print("")