Update stock_analysis.py

stock_analysis.py

@@ -10,32 +10,35 @@ def is_business_day(a_date):
     return a_date.weekday() < 5
 
 def forecast_series(series, model="ARIMA", forecast_period=FORECAST_PERIOD):
-    predictions = []
-    confidence_intervals = []
-    
     if series.shape[1] > 1:
-        series = series['Close'].values.tolist()
+        series = series['Close'].values
     
     if model == "ARIMA":
         model = ARIMA(series, order=(5, 1, 0))
         model_fit = model.fit()
-        forecast = model_fit.forecast(steps=forecast_period, alpha=(1 - CONFIDENCE_INTERVAL))
+        forecast = model_fit.forecast(steps=forecast_period)
+        
+        # Get confidence intervals
+        conf_int = model_fit.get_forecast(steps=forecast_period).conf_int()
+        lower_ci = conf_int.iloc[:, 0] if isinstance(conf_int, pd.DataFrame) else conf_int[:, 0]
+        upper_ci = conf_int.iloc[:, 1] if isinstance(conf_int, pd.DataFrame) else conf_int[:, 1]
+        
+        # Ensure all arrays have the same length
+        min_length = min(len(forecast), len(lower_ci), len(upper_ci))
+        predictions = forecast[:min_length]
+        lower_ci = lower_ci[:min_length]
+        upper_ci = upper_ci[:min_length]
         
-        # Check if forecast is a numpy array (newer statsmodels) or a ForecastResults object (older statsmodels)
-        if isinstance(forecast, np.ndarray):
-            predictions = forecast
-            confidence_intervals = model_fit.get_forecast(steps=forecast_period).conf_int()
-        else:
-            predictions = forecast.predicted_mean
-            confidence_intervals = forecast.conf_int()
     elif model == "Prophet":
         # Implement Prophet forecasting method
         pass
     elif model == "LSTM":
         # Implement LSTM forecasting method
         pass
+    else:
+        raise ValueError(f"Unsupported model: {model}")
 
-    return predictions, confidence_intervals
+    return predictions, pd.DataFrame({'Lower_CI': lower_ci, 'Upper_CI': upper_ci})
 
 def get_stock_graph_and_info(idx, stock, interval, graph_type, forecast_method, start_date, end_date):
     stock_name, ticker_name = stock.split(":")
@@ -49,14 +52,14 @@ def get_stock_graph_and_info(idx, stock, interval, graph_type, forecast_method,
     predictions, confidence_intervals = forecast_series(series, model=forecast_method)
 
     last_date = pd.to_datetime(series['Date'].values[-1])
-    forecast_dates = pd.date_range(start=last_date + timedelta(days=1), periods=FORECAST_PERIOD)
-    forecast_dates = [date for date in forecast_dates if is_business_day(date)]
+    forecast_dates = pd.date_range(start=last_date + timedelta(days=1), periods=len(predictions))
+    forecast_dates = [date for date in forecast_dates if is_business_day(date)][:len(predictions)]
 
     forecast = pd.DataFrame({
         "Date": forecast_dates,
         "Forecast": predictions,
-        "Lower_CI": confidence_intervals[:, 0],
-        "Upper_CI": confidence_intervals[:, 1]
+        "Lower_CI": confidence_intervals['Lower_CI'],
+        "Upper_CI": confidence_intervals['Upper_CI']
     })
 
     if graph_type == 'Line Graph':