Update app.py

app.py

@@ -36,183 +36,80 @@ COMPANIES = [
 ]
 
 class StockPredictor:
-    def __init__(self, data, model_type='Prophet'):
+    def __init__(self, data):
         self.data = data
-        self.model_type = model_type
         self.model = None
-        self.scaler = None
-        self.lstm_scaler = None
 
     def preprocess_data(self):
-        self.data['Date'] = pd.to_datetime(self.data.index)
-        self.data = self.data.reset_index(drop=True)
-
-        # Enhanced Feature Engineering
-        self.data['DayOfWeek'] = self.data['Date'].dt.dayofweek
-        self.data['Month'] = self.data['Date'].dt.month
-        self.data['Year'] = self.data['Date'].dt.year
-        self.data['IsMonthEnd'] = self.data['Date'].dt.is_month_end.astype(int)
-
-        # Technical Indicators
-        self.data['SMA_20'] = SMAIndicator(close=self.data['Close'], window=20).sma_indicator()
-        self.data['EMA_20'] = EMAIndicator(close=self.data['Close'], window=20).ema_indicator()
-        self.data['RSI'] = RSIIndicator(close=self.data['Close']).rsi()
-        bb = BollingerBands(close=self.data['Close'], window=20, window_dev=2)
-        self.data['BB_High'] = bb.bollinger_hband()
-        self.data['BB_Low'] = bb.bollinger_lband()
+        # Prophet requires columns named 'ds' and 'y'
+        self.data = self.data.reset_index()
+        self.data = self.data.rename(columns={'Date': 'ds', 'Close': 'y'})
         
-        # Log returns
-        self.data['LogReturn'] = np.log(self.data['Close'] / self.data['Close'].shift(1))
-
-        # Handle NaN values 
-        self.data.dropna(inplace=True)
-
-        # Define features for the model
-        self.features = ['Open', 'High', 'Low', 'Close', 'Volume', 'SMA_20', 'EMA_20', 'RSI', 'BB_High', 'BB_Low', 'LogReturn', 'DayOfWeek', 'Month', 'Year', 'IsMonthEnd']
-
-        # Apply scaling for XGBoost and RandomForest
-        if self.model_type in ['XGBoost', 'RandomForest']:
-            self.scaler = StandardScaler()
-            self.data[self.features] = self.scaler.fit_transform(self.data[self.features])
-
-        # Additional preprocessing for LSTM
-        if self.model_type == 'LSTM':
-            self.lstm_scaler = MinMaxScaler(feature_range=(0, 1))
-            self.data['Scaled_Close'] = self.lstm_scaler.fit_transform(self.data[['Close']])
+        # Add any additional features you want to use
+        self.data['SMA_20'] = self.data['y'].rolling(window=20).mean()
+        self.data['EMA_20'] = self.data['y'].ewm(span=20, adjust=False).mean()
+        self.data['RSI'] = self.calculate_rsi(self.data['y'], periods=14)
+        
+        # Handle NaN values
+        self.data = self.data.dropna()
 
-    def create_lstm_dataset(self, look_back=60):
-        scaled_data = self.data['Scaled_Close'].values
-        x, y = [], []
-        for i in range(look_back, len(scaled_data)):
-            x.append(scaled_data[i-look_back:i])
-            y.append(scaled_data[i])
-        return np.array(x), np.array(y)
+    def calculate_rsi(self, prices, periods=14):
+        delta = prices.diff()
+        gain = (delta.where(delta > 0, 0)).rolling(window=periods).mean()
+        loss = (-delta.where(delta < 0, 0)).rolling(window=periods).mean()
+        rs = gain / loss
+        return 100 - (100 / (1 + rs))
 
     def train_model(self):
         try:
-            if self.model_type == 'LSTM':
-                x, y = self.create_lstm_dataset()
-                x_train, x_test, y_train, y_test = train_test_split(x, y, test_size=0.2, shuffle=False)
-
-                model = Sequential([
-                    LSTM(50, return_sequences=True, input_shape=(x_train.shape[1], 1)),
-                    LSTM(50, return_sequences=False),
-                    Dense(25),
-                    Dense(1)
-                ])
-
-                model.compile(optimizer=Adam(learning_rate=0.001), loss='mean_squared_error')
-                model.fit(x_train, y_train, epochs=50, batch_size=32, validation_data=(x_test, y_test), verbose=0)
-
-                self.model = model
-
-            elif self.model_type == 'SARIMA':
-                train_data = self.data['Close']
-                # Use auto_arima to find optimal parameters
-                from pmdarima import auto_arima
-                auto_model = auto_arima(train_data, start_p=1, start_q=1, max_p=3, max_q=3, m=12,
-                                        start_P=0, seasonal=True, d=1, D=1, trace=True,
-                                        error_action='ignore', suppress_warnings=True, stepwise=True)
-                
-                self.model = SARIMAX(train_data, order=auto_model.order, seasonal_order=auto_model.seasonal_order)
-                self.model = self.model.fit(disp=False)
-
-            elif self.model_type == 'Prophet':
-                df = self.data[['Date', 'Close']].rename(columns={'Date': 'ds', 'Close': 'y'})
-                self.model = Prophet(
-                    changepoint_prior_scale=0.05,
-                    seasonality_prior_scale=10,
-                    holidays_prior_scale=10,
-                    daily_seasonality=True,
-                    weekly_seasonality=True,
-                    yearly_seasonality=True
-                )
-                for feature in ['SMA_20', 'EMA_20', 'RSI', 'BB_High', 'BB_Low']:
-                    self.model.add_regressor(feature)
-                    df[feature] = self.data[feature]
-                self.model.fit(df)
-
-            elif self.model_type == 'XGBoost':
-                X = self.data[self.features]
-                y = self.data['Close']
-                X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, shuffle=False)
-                
-                param_grid = {
-                    'max_depth': [3, 5],
-                    'learning_rate': [0.01, 0.1],
-                    'n_estimators': [100, 200]
-                }
-                model = xgb.XGBRegressor(objective='reg:squarederror')
-                grid_search = GridSearchCV(estimator=model, param_grid=param_grid, cv=3, n_jobs=-1, verbose=0)
-                grid_search.fit(X_train, y_train)
-                
-                self.model = grid_search.best_estimator_
-
-            elif self.model_type == 'RandomForest':
-                X = self.data[self.features]
-                y = self.data['Close']
-                X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, shuffle=False)
-                
-                param_grid = {
-                    'n_estimators': [100, 200],
-                    'max_depth': [10, 20]
-                }
-                model = RandomForestRegressor(random_state=42)
-                grid_search = GridSearchCV(estimator=model, param_grid=param_grid, cv=3, n_jobs=-1, verbose=0)
-                grid_search.fit(X_train, y_train)
-                
-                self.model = grid_search.best_estimator_
-
+            self.model = Prophet(
+                changepoint_prior_scale=0.05,
+                seasonality_prior_scale=10,
+                holidays_prior_scale=10,
+                daily_seasonality=True,
+                weekly_seasonality=True,
+                yearly_seasonality=True
+            )
+            
+            # Add additional regressors
+            self.model.add_regressor('SMA_20')
+            self.model.add_regressor('EMA_20')
+            self.model.add_regressor('RSI')
+            
+            self.model.fit(self.data)
             return True
-
         except Exception as e:
-            print(f"Error training {self.model_type} model: {str(e)}")
+            print(f"Error training Prophet model: {str(e)}")
             return False
 
     def predict(self, days=30):
         try:
-            if self.model_type == 'LSTM':
-                last_sequence = self.data['Scaled_Close'].values[-60:].reshape(1, 60, 1)
-                predictions = []
-                for _ in range(days):
-                    pred = self.model.predict(last_sequence)
-                    predictions.append(pred[0, 0])
-                    last_sequence = np.roll(last_sequence, -1, axis=1)
-                    last_sequence[0, -1, 0] = pred[0, 0]
-                return self.lstm_scaler.inverse_transform(np.array(predictions).reshape(-1, 1)).flatten()
-
-            elif self.model_type == 'SARIMA':
-                forecast = self.model.get_forecast(steps=days)
-                return forecast.predicted_mean.values
-    
-            elif self.model_type == 'Prophet':
-                future = self.model.make_future_dataframe(periods=days)
-                for feature in ['SMA_20', 'EMA_20', 'RSI', 'BB_High', 'BB_Low']:
-                    future[feature] = self.data[feature].iloc[-1]  # Use last known value
-                forecast = self.model.predict(future)
-                return forecast[['ds', 'yhat', 'yhat_lower', 'yhat_upper']][-days:]
-
-            elif self.model_type in ['XGBoost', 'RandomForest']:
-                last_data = self.data[self.features].iloc[-1:].values
-                predictions = []
-                for _ in range(days):
-                    pred = self.model.predict(last_data)
-                    predictions.append(pred[0])
-                    # Update last_data for next prediction
-                    last_data = np.roll(last_data, -1, axis=1)
-                    last_data[0, -5] = pred[0]  # Assuming 'Close' is the 5th from last feature
-                return np.array(predictions)
-
+            future = self.model.make_future_dataframe(periods=days)
+            
+            # Add regressor values for future dates
+            for feature in ['SMA_20', 'EMA_20', 'RSI']:
+                future[feature] = self.data[feature].iloc[-1]  # Use last known value
+            
+            forecast = self.model.predict(future)
+            return forecast
         except Exception as e:
-            print(f"Error predicting with {self.model_type} model: {str(e)}")
+            print(f"Error predicting with Prophet model: {str(e)}")
             return None
 
-        def evaluate_model(self, test_data):
-            predictions = self.predict(len(test_data))
-            mse = mean_squared_error(test_data['Close'], predictions)
-            mape = mean_absolute_percentage_error(test_data['Close'], predictions)
-            rmse = np.sqrt(mse)
-            return mse, mape, rmse
+    def evaluate_model(self, test_data):
+        predictions = self.predict(days=len(test_data))
+        
+        if predictions is None:
+            return None, None, None
+
+        actual = test_data['Close'].values
+        predicted = predictions['yhat'].values[-len(test_data):]
+
+        mse = mean_squared_error(actual, predicted)
+        mape = mean_absolute_percentage_error(actual, predicted)
+        rmse = np.sqrt(mse)
+
+        return mse, mape, rmse
 
 def fetch_stock_data(ticker):
     try:
@@ -349,7 +246,7 @@ def main():
         predict_stock_prices()
 
 def test_model():
-    st.header("Test Stock Prediction Model")
+    st.header("Test Prophet Model")
 
     col1, col2 = st.columns(2)
 
@@ -357,9 +254,6 @@ def test_model():
         company = st.selectbox("Select Company", [company for company, _ in COMPANIES])
         test_split = st.slider("Test Data Split", 0.1, 0.5, 0.2, 0.05)
 
-    with col2:
-        model_type = st.selectbox("Select Model Type", ['Prophet', 'LSTM', 'SARIMA', 'XGBoost', 'RandomForest'])
-
     if st.button("Train and Test Model"):
         with st.spinner("Fetching data and training model..."):
             company_name, ticker = next((name, symbol) for name, symbol in COMPANIES if name == company)
@@ -378,24 +272,30 @@ def test_model():
                 train_data = data.iloc[:split_index]
                 test_data = data.iloc[split_index:]
 
-                predictor = StockPredictor(train_data, model_type)  # Updated: added model_type argument
+                predictor = StockPredictor(train_data)
                 predictor.preprocess_data()
                 if predictor.train_model():
                     test_pred = predictor.predict(days=len(test_data))
 
                     if test_pred is not None:
                         mse, mape, rmse = predictor.evaluate_model(test_data)
-                        accuracy = 100 - mape * 100
-
-                        st.subheader("Model Performance")
-                        st.metric("Prediction Accuracy", f"{accuracy:.2f}%")
-                        st.metric("Mean Squared Error", f"{mse:.4f}")
-                        st.metric("Root Mean Squared Error", f"{rmse:.4f}")
-
-                        plot = create_test_plot(predictor.data, test_data, test_pred, company_name)
-                        st.plotly_chart(plot, use_container_width=True)
+                        
+                        if mse is not None and mape is not None and rmse is not None:
+                            accuracy = 100 - mape * 100
+
+                            st.subheader("Model Performance")
+                            st.metric("Prediction Accuracy", f"{accuracy:.2f}%")
+                            st.metric("Mean Squared Error", f"{mse:.4f}")
+                            st.metric("Root Mean Squared Error", f"{rmse:.4f}")
+
+                            plot = create_test_plot(predictor.data, test_data, test_pred, company_name)
+                            st.plotly_chart(plot, use_container_width=True)
+                        else:
+                            st.error("Failed to evaluate the model. The evaluation metrics are None.")
+                    else:
+                        st.error("Failed to generate predictions. The predicted data is None.")
                 else:
-                    st.error(f"Failed to train the {model_type} model. Please try a different dataset or model type.")
+                    st.error("Failed to train the Prophet model. Please try a different dataset.")
 
 def predict_stock_prices():
     st.header("Predict Stock Prices")
@@ -406,9 +306,6 @@ def predict_stock_prices():
         company = st.selectbox("Select Company", [company for company, _ in COMPANIES])
         days_to_predict = st.slider("Days to Predict", 1, 365, 30)
 
-    with col2:
-        model_type = st.selectbox("Select Model Type", ['Prophet', 'LSTM', 'SARIMA', 'XGBoost', 'RandomForest'])
-
     if st.button("Predict Stock Prices"):
         with st.spinner("Fetching data and making predictions..."):
             company_name, ticker = next((name, symbol) for name, symbol in COMPANIES if name == company)
@@ -423,7 +320,7 @@ def predict_stock_prices():
 
                 st.markdown(get_table_download_link(data), unsafe_allow_html=True)
 
-                predictor = StockPredictor(data, model_type)  # Updated: added model_type argument
+                predictor = StockPredictor(data)
                 predictor.preprocess_data()
                 if predictor.train_model():
                     predictions = predictor.predict(days=days_to_predict)
@@ -444,8 +341,10 @@ def predict_stock_prices():
                         st.subheader("Latest News")
                         for item in news:
                             st.markdown(f"[{item['title']}]({item['link']}) ({item['pubDate']})")
+                    else:
+                        st.error("Failed to generate predictions. The predicted data is None.")
                 else:
-                    st.error(f"Failed to train the {model_type} model. Please try a different dataset or model type.")
+                    st.error("Failed to train the Prophet model. Please try a different dataset.")
 
 def explore_data():
     st.header("Explore Stock Data")