Update utils.py

utils.py

@@ -4,7 +4,6 @@ import numpy as np
 import torch
 from datetime import datetime, timedelta
 import plotly.graph_objects as go
-import plotly.express as px
 from plotly.subplots import make_subplots
 import spaces
 
@@ -41,10 +40,7 @@ def calculate_technical_indicators(data):
         rs = gain / loss
         rsi = 100 - (100 / (1 + rs))
         return rsi
-    indicators['rsi'] = {
-        'current': calculate_rsi(data['Close']).iloc[-1],
-        'values': calculate_rsi(data['Close'])
-    }
+    indicators['rsi'] = {'current': calculate_rsi(data['Close']).iloc[-1], 'values': calculate_rsi(data['Close'])}
     def calculate_macd(prices, fast=12, slow=26, signal=9):
         exp1 = prices.ewm(span=fast).mean()
         exp2 = prices.ewm(span=slow).mean()
@@ -53,14 +49,7 @@ def calculate_technical_indicators(data):
         histogram = macd - signal_line
         return macd, signal_line, histogram
     macd, signal_line, histogram = calculate_macd(data['Close'])
-    indicators['macd'] = {
-        'macd': macd.iloc[-1],
-        'signal': signal_line.iloc[-1],
-        'histogram': histogram.iloc[-1],
-        'signal_text': 'BUY' if histogram.iloc[-1] > 0 else 'SELL',
-        'macd_values': macd,
-        'signal_values': signal_line
-    }
+    indicators['macd'] = {'macd': macd.iloc[-1], 'signal': signal_line.iloc[-1], 'histogram': histogram.iloc[-1], 'signal_text': 'BUY' if histogram.iloc[-1] > 0 else 'SELL', 'macd_values': macd, 'signal_values': signal_line}
     def calculate_bollinger_bands(prices, period=20, std_dev=2):
         sma = prices.rolling(window=period).mean()
         std = prices.rolling(window=period).std()
@@ -70,28 +59,11 @@ def calculate_technical_indicators(data):
     upper, middle, lower = calculate_bollinger_bands(data['Close'])
     current_price = data['Close'].iloc[-1]
     bb_position = (current_price - lower.iloc[-1]) / (upper.iloc[-1] - lower.iloc[-1])
-    indicators['bollinger'] = {
-        'upper': upper.iloc[-1],
-        'middle': middle.iloc[-1],
-        'lower': lower.iloc[-1],
-        'position': 'UPPER' if bb_position > 0.8 else 'LOWER' if bb_position < 0.2 else 'MIDDLE'
-    }
+    indicators['bollinger'] = {'upper': upper.iloc[-1], 'middle': middle.iloc[-1], 'lower': lower.iloc[-1], 'position': 'UPPER' if bb_position > 0.8 else 'LOWER' if bb_position < 0.2 else 'MIDDLE'}
     sma_20_series = data['Close'].rolling(20).mean()
     sma_50_series = data['Close'].rolling(50).mean()
-    indicators['moving_averages'] = {
-        'sma_20': sma_20_series.iloc[-1],
-        'sma_50': sma_50_series.iloc[-1],
-        'sma_200': data['Close'].rolling(200).mean().iloc[-1],
-        'ema_12': data['Close'].ewm(span=12).mean().iloc[-1],
-        'ema_26': data['Close'].ewm(span=26).mean().iloc[-1],
-        'sma_20_values': sma_20_series,
-        'sma_50_values': sma_50_series
-    }
-    indicators['volume'] = {
-        'current': data['Volume'].iloc[-1],
-        'avg_20': data['Volume'].rolling(20).mean().iloc[-1],
-        'ratio': data['Volume'].iloc[-1] / data['Volume'].rolling(20).mean().iloc[-1]
-    }
+    indicators['moving_averages'] = {'sma_20': sma_20_series.iloc[-1], 'sma_50': sma_50_series.iloc[-1], 'sma_200': data['Close'].rolling(200).mean().iloc[-1], 'ema_12': data['Close'].ewm(span=12).mean().iloc[-1], 'ema_26': data['Close'].ewm(span=26).mean().iloc[-1], 'sma_20_values': sma_20_series, 'sma_50_values': sma_50_series}
+    indicators['volume'] = {'current': data['Volume'].iloc[-1], 'avg_20': data['Volume'].rolling(20).mean().iloc[-1], 'ratio': data['Volume'].iloc[-1] / data['Volume'].rolling(20).mean().iloc[-1]}
     return indicators
 
 def generate_trading_signals(data, indicators):
@@ -146,59 +118,21 @@ def generate_trading_signals(data, indicators):
         signal_details.append(f"⚪ Normal volume ({volume_ratio:.1f}x avg)")
     total_signals = buy_signals + sell_signals
     signal_strength = (buy_signals / max(total_signals, 1)) * 100
-    if buy_signals > sell_signals:
-        overall_signal = "BUY"
-    elif sell_signals > buy_signals:
-        overall_signal = "SELL"
-    else:
-        overall_signal = "HOLD"
+    overall_signal = "BUY" if buy_signals > sell_signals else "SELL" if sell_signals > buy_signals else "HOLD"
     recent_high = data['High'].tail(20).max()
     recent_low = data['Low'].tail(20).min()
-    signals = {
-        'overall': overall_signal,
-        'strength': signal_strength,
-        'details': '\n'.join(signal_details),
-        'support': recent_low,
-        'resistance': recent_high,
-        'stop_loss': recent_low * 0.95 if overall_signal == "BUY" else recent_high * 1.05
-    }
+    signals = {'overall': overall_signal, 'strength': signal_strength, 'details': '\n'.join(signal_details), 'support': recent_low, 'resistance': recent_high, 'stop_loss': recent_low * 0.95 if overall_signal == "BUY" else recent_high * 1.05}
     return signals
 
 def get_fundamental_data(stock):
     try:
         info = stock.info
         history = stock.history(period="1d")
-        fundamental_info = {
-            'name': info.get('longName', 'N/A'),
-            'current_price': history['Close'].iloc[-1] if not history.empty else 0,
-            'market_cap': info.get('marketCap', 0),
-            'pe_ratio': info.get('forwardPE', 0),
-            'dividend_yield': info.get('dividendYield', 0) * 100 if info.get('dividendYield') else 0,
-            'volume': history['Volume'].iloc[-1] if not history.empty else 0,
-            'info': f"""
-Sector: {info.get('sector', 'N/A')}
-Industry: {info.get('industry', 'N/A')}
-Market Cap: {format_large_number(info.get('marketCap', 0))}
-52 Week High: {info.get('fiftyTwoWeekHigh', 'N/A')}
-52 Week Low: {info.get('fiftyTwoWeekLow', 'N/A')}
-Beta: {info.get('beta', 'N/A')}
-EPS: {info.get('forwardEps', 'N/A')}
-Book Value: {info.get('bookValue', 'N/A')}
-Price to Book: {info.get('priceToBook', 'N/A')}
-            """.strip()
-        }
+        fundamental_info = {'name': info.get('longName', 'N/A'), 'current_price': history['Close'].iloc[-1] if not history.empty else 0, 'market_cap': info.get('marketCap', 0), 'pe_ratio': info.get('forwardPE', 0), 'dividend_yield': info.get('dividendYield', 0) * 100 if info.get('dividendYield') else 0, 'volume': history['Volume'].iloc[-1] if not history.empty else 0, 'info': f"Sector: {info.get('sector', 'N/A')}\nIndustry: {info.get('industry', 'N/A')}\nMarket Cap: {format_large_number(info.get('marketCap', 0))}\n52 Week High: {info.get('fiftyTwoWeekHigh', 'N/A')}\n52 Week Low: {info.get('fiftyTwoWeekLow', 'N/A')}\nBeta: {info.get('beta', 'N/A')}\nEPS: {info.get('forwardEps', 'N/A')}\nBook Value: {info.get('bookValue', 'N/A')}\nPrice to Book: {info.get('priceToBook', 'N/A')}"}
         return fundamental_info
     except Exception as e:
         print(f"Error getting fundamental data: {e}")
-        return {
-            'name': 'N/A',
-            'current_price': 0,
-            'market_cap': 0,
-            'pe_ratio': 0,
-            'dividend_yield': 0,
-            'volume': 0,
-            'info': 'Unable to fetch fundamental data'
-        }
+        return {'name': 'N/A', 'current_price': 0, 'market_cap': 0, 'pe_ratio': 0, 'dividend_yield': 0, 'volume': 0, 'info': 'Unable to fetch fundamental data'}
 
 def format_large_number(num):
     if num >= 1e12:
@@ -218,49 +152,37 @@ def predict_prices(data, model=None, tokenizer=None, prediction_days=30):
         prices = data['Close'].values.astype(np.float32)
         try:
             from chronos import BaseChronosPipeline
-        except Exception as ie:
-            return {
-                'values': [],
-                'dates': [],
-                'high_30d': 0,
-                'low_30d': 0,
-                'mean_30d': 0,
-                'change_pct': 0,
-                'summary': 'chronos package not installed. install with: pip install chronos-forecasting'
-            }
+        except Exception:
+            return {'values': [], 'dates': [], 'high_30d': 0, 'low_30d': 0, 'mean_30d': 0, 'change_pct': 0, 'summary': 'chronos package not installed. install with: pip install chronos-forecasting'}
         pipeline = BaseChronosPipeline.from_pretrained("amazon/chronos-bolt-base", device_map="auto")
         with torch.no_grad():
             forecast = pipeline.predict(context=torch.tensor(prices), prediction_length=prediction_days)
-        if hasattr(forecast, 'mean'):
-            mean_forecast = forecast.mean(dim=1).squeeze().cpu().numpy()
+        if isinstance(forecast, torch.Tensor):
+            forecast_np = forecast.squeeze().cpu().numpy()
+        elif hasattr(forecast, 'numpy'):
+            forecast_np = forecast.numpy()
+        else:
+            forecast_np = np.array(forecast)
+        if forecast_np.ndim == 2:
+            mean_forecast = forecast_np.mean(axis=0)
+        elif forecast_np.ndim == 3:
+            mean_forecast = forecast_np.mean(axis=(0, 1))
+        elif forecast_np.ndim == 1:
+            mean_forecast = forecast_np
         else:
-            mean_forecast = forecast.mean(dim=1).squeeze().cpu().numpy()
+            mean_forecast = np.array([])
         pred_len = len(mean_forecast)
+        if pred_len == 0:
+            return {'values': [], 'dates': [], 'high_30d': 0, 'low_30d': 0, 'mean_30d': 0, 'change_pct': 0, 'summary': 'Model did not return valid prediction output.'}
         last_price = prices[-1]
-        predicted_high = np.max(mean_forecast) if pred_len > 0 else 0
-        predicted_low = np.min(mean_forecast) if pred_len > 0 else 0
-        predicted_mean = np.mean(mean_forecast) if pred_len > 0 else 0
-        change_pct = ((predicted_mean - last_price) / last_price) * 100 if last_price != 0 and pred_len > 0 else 0
-        return {
-            'values': mean_forecast,
-            'dates': pd.date_range(start=data.index[-1] + timedelta(days=1), periods=pred_len, freq='D') if pred_len > 0 else [],
-            'high_30d': predicted_high,
-            'low_30d': predicted_low,
-            'mean_30d': predicted_mean,
-            'change_pct': change_pct,
-            'summary': f"AI Model: Amazon Chronos-Bolt (Base)\nPrediction Period: {pred_len} days\nExpected Change: {change_pct:.2f}%\nConfidence: Medium\nNote: AI predictions are for reference only and not financial advice"
-        }
+        predicted_high = float(np.max(mean_forecast))
+        predicted_low = float(np.min(mean_forecast))
+        predicted_mean = float(np.mean(mean_forecast))
+        change_pct = ((predicted_mean - last_price) / last_price) * 100 if last_price != 0 else 0
+        return {'values': mean_forecast, 'dates': pd.date_range(start=data.index[-1] + timedelta(days=1), periods=pred_len, freq='D'), 'high_30d': predicted_high, 'low_30d': predicted_low, 'mean_30d': predicted_mean, 'change_pct': change_pct, 'summary': f"AI Model: Amazon Chronos-Bolt (Base)\nPrediction Period: {pred_len} days\nPredicted High: {predicted_high:.2f}\nPredicted Low: {predicted_low:.2f}\nExpected Change: {change_pct:.2f}%\nConfidence: Medium\nNote: AI predictions are for reference only and not financial advice"}
     except Exception as e:
         print(f"Error in prediction: {e}")
-        return {
-            'values': [],
-            'dates': [],
-            'high_30d': 0,
-            'low_30d': 0,
-            'mean_30d': 0,
-            'change_pct': 0,
-            'summary': f'Prediction unavailable due to model error: {e}'
-        }
+        return {'values': [], 'dates': [], 'high_30d': 0, 'low_30d': 0, 'mean_30d': 0, 'change_pct': 0, 'summary': f'Prediction unavailable due to model error: {e}'}
 
 def create_price_chart(data, indicators):
     fig = make_subplots(rows=3, cols=1, shared_xaxes=True, vertical_spacing=0.05, subplot_titles=('Price & Moving Averages', 'RSI', 'MACD'), row_width=[0.2, 0.2, 0.7])
@@ -275,15 +197,6 @@ def create_price_chart(data, indicators):
     fig.update_layout(title='Technical Analysis Dashboard', height=900, showlegend=True, xaxis_rangeslider_visible=False)
     return fig
 
-def create_technical_chart(data, indicators):
-    fig = make_subplots(rows=2, cols=2, subplot_titles=('Bollinger Bands', 'Volume', 'Price vs MA', 'RSI Analysis'), specs=[[{"secondary_y": False}, {"secondary_y": False}], [{"secondary_y": False}, {"secondary_y": False}]])
-    fig.add_trace(go.Scatter(x=data.index, y=data['Close'], name='Price', line=dict(color='black')), row=1, col=1)
-    fig.add_trace(go.Bar(x=data.index, y=data['Volume'], name='Volume', marker_color='lightblue'), row=1, col=2)
-    fig.add_trace(go.Scatter(x=data.index, y=data['Close'], name='Price', line=dict(color='black')), row=2, col=1)
-    fig.add_trace(go.Scatter(x=data.index, y=indicators['moving_averages']['sma_20_values'], name='SMA 20', line=dict(color='orange', dash='dash')), row=2, col=1)
-    fig.update_layout(title='Technical Indicators Overview', height=600, showlegend=False)
-    return fig
-
 def create_prediction_chart(data, predictions):
     if not len(predictions['values']):
         return go.Figure()