-
Notifications
You must be signed in to change notification settings - Fork 0
/
main.py
73 lines (58 loc) · 2.27 KB
/
main.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
import streamlit as sl
from datetime import date
import yfinance as yahoo
from prophet import Prophet
from prophet.plot import plot_plotly
from plotly import graph_objs as go
START = "2015-01-01"
TODAY = date.today().strftime("%Y-%m-%d")
sl.title("Stock Prediction Site")
# Use a text input instead of a selectbox for ticker symbols
selected_stock = sl.text_input("Enter the stock ticker for prediction (e.g., AAPL for Apple Inc.)", "AAPL")
num_years = sl.slider("Years of prediction", 1, 10)
period = num_years * 365
# helpful cache decorator for streamlit
@sl.cache_data
def get_data(ticker):
# this data is a pandas dataframe
data = yahoo.download(ticker, START, TODAY)
# puts date in first column
data.reset_index(inplace=True)
return data
def plot_data(data):
fig = go.Figure()
fig.add_trace(go.Scatter(x=data['Date'], y=data['Open'], name='stock_open'))
fig.add_trace(go.Scatter(x=data['Date'], y=data['Close'], name='stock_close'))
fig.layout.update(title_text="Time Series Data", xaxis_rangeslider_visible=True)
sl.plotly_chart(fig)
# Ensure a ticker is entered before fetching data
if selected_stock:
data_state = sl.text("Loading data...")
try:
data = get_data(selected_stock)
data_state.text("Loaded data!")
sl.subheader("Raw stock data")
sl.write(data.tail())
plot_data(data)
# prediction using prophet
df_train = data[['Date', 'Close']]
# how prophet takes the data, look at documentation
df_train = df_train.rename(columns={"Date": "ds", "Close": "y"})
# init prophet model and start training
model = Prophet()
model.fit(df_train)
future = model.make_future_dataframe(periods=period)
prediction = model.predict(future)
sl.subheader("Prediction data")
sl.write(prediction.tail())
sl.write('Prediction data')
fig1 = plot_plotly(model, prediction)
sl.plotly_chart(fig1)
sl.write('Prediction components')
fig2 = model.plot_components(prediction)
# not a plot so dont need to plotly plot
sl.write(fig2)
except Exception as e:
data_state.text(f"Error loading data for {selected_stock}: {str(e)}")
else:
sl.text("Please enter a stock ticker to get started.")