diff --git a/config.py b/config.py index 2def8db..85b0014 100644 --- a/config.py +++ b/config.py @@ -7,6 +7,17 @@ from dotenv import load_dotenv load_dotenv() +def _check_amd_gpu() -> bool: + """Detect AMD GPU via ROCm device name.""" + if not torch.cuda.is_available(): + return False + try: + name = torch.cuda.get_device_name(0) + return any(kw in name for kw in ("AMD", "Radeon", "gfx")) + except Exception: + return False + + class Config: # Project settings PROJECT_NAME = "StockGNN_R9700" @@ -20,12 +31,6 @@ class Config: EXTERNAL_DATA_DIR = os.path.join(DATA_DIR, "external") MODEL_DIR = os.path.join(BASE_DIR, "models") - # Ensure directories exist - os.makedirs(RAW_DATA_DIR, exist_ok=True) - os.makedirs(PROCESSED_DATA_DIR, exist_ok=True) - os.makedirs(EXTERNAL_DATA_DIR, exist_ok=True) - os.makedirs(MODEL_DIR, exist_ok=True) - # Stock universe settings INITIAL_TICKERS = [ "AAPL", @@ -71,7 +76,7 @@ class Config: # AMD GPU settings (Radeon R9700 AI Pro - 32GB) DEVICE = "cuda" if torch.cuda.is_available() else "cpu" - AMD_GPU = True + AMD_GPU = _check_amd_gpu() GPU_MEMORY_LIMIT = 0.9 # Use 90% of 32GB = 28.8GB ROCM_OPT_LEVEL = "O2" # Optimization level for ROCm ('O0', 'O1', 'O2') PIN_MEMORY = True # Enable pinned memory for faster data transfer @@ -195,5 +200,11 @@ class Config: EXECUTION_TIME_HORIZON = "5min" # Time to complete execution MARKET_IMPACT_MODEL = "kyle" # 'kyle', 'almgren_chriss', or 'none' + def __init__(self): + os.makedirs(self.RAW_DATA_DIR, exist_ok=True) + os.makedirs(self.PROCESSED_DATA_DIR, exist_ok=True) + os.makedirs(self.EXTERNAL_DATA_DIR, exist_ok=True) + os.makedirs(self.MODEL_DIR, exist_ok=True) + config = Config() diff --git a/live_trading.py b/live_trading.py index d68376b..a917e65 100644 --- a/live_trading.py +++ b/live_trading.py @@ -376,60 +376,50 @@ class LiveTradingSystem: if not current_price: continue - # Skip if we already have a position in this stock - if ticker in self.trader.current_positions: - continue + held_qty = self.trader.current_positions.get(ticker, 0) - # Check holding period constraints - if self.trader._check_holding_period(ticker, timestamp): - continue - - # Generate signals based on prediction - if prediction > 0.002: # Buy signal - # Calculate position size - position_size = self.trader._calculate_position_size( - ticker, current_price - ) - - if position_size > 0: - # Create buy order + if held_qty > 0: + # Existing position: sell on negative signal (holding-period check still applies) + if prediction < -0.002 and not self.trader._check_holding_period( + ticker, timestamp + ): order = { "ticker": ticker, - "action": "buy", - "quantity": position_size, + "action": "sell", + "quantity": held_qty, "price": current_price, "timestamp": timestamp, "type": "market", } - - # Submit order order_id = self.broker.submit_order(order) if order_id: self.trader.pending_orders[order_id] = order logger.info( - f"Submitted buy order for {position_size} shares of {ticker} at {current_price}" + f"Submitted sell order for {held_qty} shares of {ticker} at {current_price}" ) - - elif ( - prediction < -0.002 and ticker in self.trader.current_positions - ): # Sell signal - # Create sell order - order = { - "ticker": ticker, - "action": "sell", - "quantity": self.trader.current_positions[ticker], - "price": current_price, - "timestamp": timestamp, - "type": "market", - } - - # Submit order - order_id = self.broker.submit_order(order) - if order_id: - self.trader.pending_orders[order_id] = order - logger.info( - f"Submitted sell order for {self.trader.current_positions[ticker]} shares of {ticker} at {current_price}" + else: + # No position: buy on positive signal + if prediction > 0.002 and not self.trader._check_holding_period( + ticker, timestamp + ): + position_size = self.trader._calculate_position_size( + ticker, current_price ) + if position_size > 0: + order = { + "ticker": ticker, + "action": "buy", + "quantity": position_size, + "price": current_price, + "timestamp": timestamp, + "type": "market", + } + order_id = self.broker.submit_order(order) + if order_id: + self.trader.pending_orders[order_id] = order + logger.info( + f"Submitted buy order for {position_size} shares of {ticker} at {current_price}" + ) except Exception as e: logger.error(f"Error generating trading signals: {str(e)}", exc_info=True) diff --git a/src/amd/optimizations.py b/src/amd/optimizations.py index 79dcfd4..28026dd 100644 --- a/src/amd/optimizations.py +++ b/src/amd/optimizations.py @@ -290,8 +290,8 @@ class AMDGATConv(MessagePassing): alpha_dst = (x_dst * self.att_dst).sum(dim=-1) alpha = (alpha_src, alpha_dst) - # Propagate - out = self.propagate(edge_index, x=(x_src, x_dst), alpha=alpha, size=size) + # Propagate — pass edge_attr so correlation weights reach message() + out = self.propagate(edge_index, x=(x_src, x_dst), alpha=alpha, edge_attr=edge_attr, size=size) # Concatenate or average heads if self.concat: @@ -312,7 +312,8 @@ class AMDGATConv(MessagePassing): loop_index = loop_index.unsqueeze(0).repeat(2, 1) if edge_attr is not None: - loop_attr = edge_attr.new_zeros((num_nodes,) + edge_attr.size()[1:]) + # Self-loop weight = 1.0 so self-messages are not zeroed out + loop_attr = edge_attr.new_ones((num_nodes,) + edge_attr.size()[1:]) edge_attr = torch.cat([edge_attr, loop_attr], dim=0) edge_index = torch.cat([edge_index, loop_index], dim=1) @@ -321,15 +322,21 @@ class AMDGATConv(MessagePassing): def propagate(self, edge_index, size=None, **kwargs): return super().propagate(edge_index, size=size, **kwargs) - def message(self, x_j, alpha_j, alpha_i, index, ptr, size_i): + def message(self, x_j, alpha_j, alpha_i, edge_attr, index, ptr, size_i): # Compute attention weights alpha = alpha_j + alpha_i alpha = F.leaky_relu(alpha, negative_slope=0.2) alpha = self._softmax(alpha, index, ptr, size_i) alpha = F.dropout(alpha, p=self.dropout, training=self.training) - # Weighted sum of values - return x_j * alpha.unsqueeze(-1) + # (E, heads, out_channels) scaled by attention + msg = x_j * alpha.unsqueeze(-1) + + # Scale by edge correlation weight: (E, 1) → broadcasts to (E, heads, out_channels) + if edge_attr is not None: + msg = msg * edge_attr.unsqueeze(1) + + return msg def _softmax(self, src, index, ptr, num_nodes): # Memory-efficient softmax diff --git a/src/data/pipeline.py b/src/data/pipeline.py index 7d8234e..5fba448 100644 --- a/src/data/pipeline.py +++ b/src/data/pipeline.py @@ -147,25 +147,26 @@ class StockDataPipeline: def _load_data(self): """Load existing data from disk with memory management""" - try: - # Check memory before loading - if not self.memory_manager.ensure_memory(2 * 1024**3): # 2GB - logger.warning("Skipping data load due to memory constraints") - return + if not self.memory_manager.ensure_memory(2 * 1024**3): # 2GB + logger.warning("Skipping data load due to memory constraints") + return - self.price_data = self._load_pickle("price_data.pkl") - self.corporate_actions = self._load_pickle("corporate_actions.pkl") - self.sector_data = self._load_pickle("sector_data.pkl") - self.index_composition = self._load_pickle("index_composition.pkl") + for attr, filename in [ + ("price_data", "price_data.pkl"), + ("corporate_actions", "corporate_actions.pkl"), + ("sector_data", "sector_data.pkl"), + ("index_composition", "index_composition.pkl"), + ]: + try: + setattr(self, attr, self._load_pickle(filename)) + logger.info(f"Loaded {filename}") + except FileNotFoundError: + logger.info(f"{filename} not found, using empty store") + except Exception as e: + logger.error(f"Error loading {filename}: {e}", exc_info=True) + self.memory_manager.empty_cache() - logger.info("Loaded existing data from disk") - self.memory_manager.log_memory_usage("[After Data Load]") - - except FileNotFoundError: - logger.info("No existing data found. Starting with empty databases.") - except Exception as e: - logger.error(f"Error loading data: {str(e)}", exc_info=True) - self.memory_manager.empty_cache() + self.memory_manager.log_memory_usage("[After Data Load]") def _save_data(self): """Save data to disk with memory management""" @@ -497,19 +498,23 @@ class StockDataPipeline: def _store_index_composition(self): """Store index composition in the database""" - for index_ticker, composition in self.index_composition.items(): - for date, members in composition.items(): - for member in members: - with sqlite3.connect(self.db_path) as conn: - cursor = conn.cursor() - cursor.execute( - """ - INSERT OR REPLACE INTO index_composition - (index_ticker, date, member_ticker) - VALUES (?, ?, ?) - """, - (index_ticker, date, member), - ) + rows = [ + (index_ticker, date, member) + for index_ticker, composition in self.index_composition.items() + for date, members in composition.items() + for member in members + ] + if not rows: + return + with sqlite3.connect(self.db_path) as conn: + conn.executemany( + """ + INSERT OR REPLACE INTO index_composition + (index_ticker, date, member_ticker) + VALUES (?, ?, ?) + """, + rows, + ) def update_alternative_data(self, tickers: List[str]): """ @@ -806,6 +811,20 @@ class StockDataPipeline: self.memory_manager.empty_cache() continue + # Compute once per day — stock universe and PIT data don't change intraday + pd_date = pd.Timestamp(date) + current_tickers = [ + t for t in tickers + if t in self.price_data + and not self.price_data[t].empty + and self.price_data[t].index[0] <= pd_date <= self.price_data[t].index[-1] + ] + if not current_tickers: + continue + + date_dt = datetime.strptime(date, "%Y-%m-%d") + pit_cache = {t: self.get_point_in_time_data(t, date_dt) for t in current_tickers} + # Get all timestamps for this trading day timestamps = generate_intraday_timestamps(date) @@ -813,20 +832,6 @@ class StockDataPipeline: current_timestamp = timestamps[i] sequence_start = timestamps[i - config.SEQUENCE_LENGTH] - # Get current universe of stocks - current_tickers = [] - for ticker in tickers: - if ticker in self.price_data and not self.price_data[ticker].empty: - if ( - date >= self.price_data[ticker].index[0] - and date <= self.price_data[ticker].index[-1] - ): - current_tickers.append(ticker) - - # Skip if no stocks available - if not current_tickers: - continue - # Create node features for each stock in the sequence sequence_features = [] valid_tickers = [] @@ -905,13 +910,9 @@ class StockDataPipeline: continue try: - # Get sector relationship - pit1 = self.get_point_in_time_data( - ticker1, datetime.strptime(date, "%Y-%m-%d") - ) - pit2 = self.get_point_in_time_data( - ticker2, datetime.strptime(date, "%Y-%m-%d") - ) + # Use pre-fetched PIT data — avoids O(n²) repeated calls + pit1 = pit_cache[ticker1] + pit2 = pit_cache[ticker2] if ( pit1["sector"] diff --git a/src/evaluation/backtester.py b/src/evaluation/backtester.py index bbec96a..7f92c14 100644 --- a/src/evaluation/backtester.py +++ b/src/evaluation/backtester.py @@ -7,6 +7,7 @@ from datetime import timedelta from typing import Dict, List, Optional, Tuple import pandas as pd +import torch from config import config from src.models.trainer import GNNTrainer @@ -61,9 +62,11 @@ class GNNBacktester: # Move data to the same device as the model data_device = data.to(config.DEVICE) - with __import__("torch").no_grad(): + with torch.no_grad(): predictions = self.model(data_device) + portfolio_val = self._portfolio_value(date) + for i, ticker in enumerate(data.tickers): pred = predictions[i].item() price = self._get_current_price(ticker, date) @@ -71,18 +74,40 @@ class GNNBacktester: if price is None or price <= 0: continue - if pred > 0.002: - order = { - "ticker": ticker, - "action": "buy", - "quantity": 100, - "price": price, - "timestamp": str(date), - "type": "market", - } - order_id = self.broker.submit_order(order) - if order_id: - trade_log.append({**order, "order_id": order_id}) + held_qty = self.broker.positions.get(ticker, 0) + + if held_qty > 0: + # Existing position: sell on negative signal + if pred < -0.002: + order = { + "ticker": ticker, + "action": "sell", + "quantity": held_qty, + "price": price, + "timestamp": str(date), + "type": "market", + } + order_id = self.broker.submit_order(order) + if order_id: + trade_log.append({**order, "order_id": order_id}) + else: + # No position: buy on positive signal + if pred > 0.002: + max_spend = portfolio_val * config.MAX_POSITION_SIZE + quantity = int(max_spend / price) + if quantity <= 0: + continue + order = { + "ticker": ticker, + "action": "buy", + "quantity": quantity, + "price": price, + "timestamp": str(date), + "type": "market", + } + order_id = self.broker.submit_order(order) + if order_id: + trade_log.append({**order, "order_id": order_id}) portfolio_values.append(self._portfolio_value(date)) dates.append(date) diff --git a/src/models/gnn_model.py b/src/models/gnn_model.py index c3a8660..9fa532c 100644 --- a/src/models/gnn_model.py +++ b/src/models/gnn_model.py @@ -208,124 +208,82 @@ class CorporateActionAwareGNN(nn.Module): nn.Softmax(dim=1), ) - # LSTM for temporal dependencies with AMD optimizations - self.lstm = nn.LSTM( - input_size=config.HIDDEN_CHANNELS, - hidden_size=config.HIDDEN_CHANNELS, - num_layers=2, - batch_first=True, - dropout=config.DROPOUT if config.NUM_HEADS > 1 else 0, - bidirectional=False, + # MLP applied after GNN (replaces single-timestep LSTM which was equivalent + # to a linear layer but 4× more parameters) + self.post_gnn_mlp = nn.Sequential( + nn.Linear(config.HIDDEN_CHANNELS, config.HIDDEN_CHANNELS), + nn.SiLU(), + nn.LayerNorm(config.HIDDEN_CHANNELS), ) - # Initialize LSTM weights - self._init_lstm_weights() - - # Final prediction layer - self.linear = nn.Linear( - config.HIDDEN_CHANNELS * 2, 1 - ) # *2 for concatenating GNN output and attention features + # Final prediction: GNN output (HIDDEN) concat with weighted modal sum (HIDDEN) + self.linear = nn.Linear(config.HIDDEN_CHANNELS * 2, 1) # Initialize weights self._init_weights() def _init_weights(self): """Initialize weights with AMD-friendly initialization""" - # Initialize linear layer weights nn.init.xavier_uniform_(self.linear.weight) nn.init.zeros_(self.linear.bias) - # Initialize attention layers - for layer in self.alternative_data_attention: - if isinstance(layer, nn.Linear): - nn.init.xavier_uniform_(layer.weight) - nn.init.zeros_(layer.bias) + for module in (self.alternative_data_attention, self.post_gnn_mlp): + for layer in module: + if isinstance(layer, nn.Linear): + nn.init.xavier_uniform_(layer.weight) + nn.init.zeros_(layer.bias) - def _init_lstm_weights(self): - """Initialize LSTM weights with orthogonal initialization for better convergence""" - for name, param in self.lstm.named_parameters(): - if "weight_ih" in name: - nn.init.orthogonal_(param) - elif "weight_hh" in name: - nn.init.orthogonal_(param) - elif "bias" in name: - nn.init.zeros_(param) - # Set forget gate bias to 1 for better gradient flow - n = param.size(0) - param.data[n // 4 : n // 2].fill_(1) + def _split_features(self, x): + """Split node feature tensor into (price, news, social, corporate_action).""" + num_price = 5 + num_news = len(config.NEWS_FEATURES) + num_social = len(config.SOCIAL_FEATURES) + return ( + x[:, :, :num_price], + x[:, :, num_price : num_price + num_news], + x[:, :, num_price + num_news : num_price + num_news + num_social], + x[:, :, -1:], + ) def forward(self, data): x, edge_index, edge_attr = data.x, data.edge_index, data.edge_attr - # Split features into different types - num_price_features = 5 - num_news_features = len(config.NEWS_FEATURES) - num_social_features = len(config.SOCIAL_FEATURES) + price_features, news_features, social_features, corporate_action_flags = ( + self._split_features(x) + ) - # Price features (first 5) - price_features = x[:, :, :num_price_features] price_processed = self.price_processor(price_features) - - # News features (next num_news_features) - news_start = num_price_features - news_end = news_start + num_news_features - news_features = x[:, :, news_start:news_end] news_processed = self.news_processor(news_features) - - # Social features (next num_social_features) - social_start = news_end - social_end = social_start + num_social_features - social_features = x[:, :, social_start:social_end] social_processed = self.social_processor(social_features) - - # Corporate action flag (last feature) - corporate_action_flags = x[:, :, -1:] corporate_action_features = self.corporate_action_mlp(corporate_action_flags) - # Apply temporal attention to each feature type price_attended = self.temporal_attention(price_processed) news_attended = self.temporal_attention(news_processed) social_attended = self.temporal_attention(social_processed) corporate_attended = self.temporal_attention(corporate_action_features) - # Calculate attention weights for alternative data alternative_features = torch.cat( [price_attended, news_attended, social_attended], dim=1 ) - attention_weights = self.alternative_data_attention(alternative_features) - # Apply attention weights weighted_price = price_attended * attention_weights[:, 0].unsqueeze(1) weighted_news = news_attended * attention_weights[:, 1].unsqueeze(1) weighted_social = social_attended * attention_weights[:, 2].unsqueeze(1) - # Combine features combined_features = torch.cat( [weighted_price, weighted_news, weighted_social, corporate_attended], dim=1 ) - # Process through GNN with gradient checkpointing for memory efficiency x = self._gnn_forward(combined_features, edge_index, edge_attr) + x = self.post_gnn_mlp(x) # (num_stocks, HIDDEN_CHANNELS) - # Process through LSTM for temporal dependencies - lstm_input = x.unsqueeze(1) # (num_stocks, 1, hidden_size) - lstm_out, _ = self.lstm(lstm_input) - lstm_out = lstm_out.squeeze(1) + # Weighted combination of attended modalities: (N, HIDDEN_CHANNELS) + modal_combined = weighted_price + weighted_news + weighted_social - # Combine GNN output with attention features - attention_features = torch.cat( - [ - attention_weights[:, 0].unsqueeze(1), - attention_weights[:, 1].unsqueeze(1), - attention_weights[:, 2].unsqueeze(1), - ], - dim=1, - ) + # Concat GNN output with modal summary → (N, HIDDEN_CHANNELS * 2) + x = torch.cat([x, modal_combined], dim=1) - x = torch.cat([lstm_out, attention_features], dim=1) - - # Final prediction return self.linear(x) def _gnn_forward(self, x, edge_index, edge_attr): @@ -345,39 +303,14 @@ class CorporateActionAwareGNN(nn.Module): def get_attention_weights(self, data): """Get attention weights for interpretability""" - x, edge_index, edge_attr = data.x, data.edge_index, data.edge_attr + x = data.x + price_features, news_features, social_features, _ = self._split_features(x) - # Process through the model up to attention - num_price_features = 5 - num_news_features = len(config.NEWS_FEATURES) - num_social_features = len(config.SOCIAL_FEATURES) - - price_features = x[:, :, :num_price_features] - news_features = x[ - :, :, num_price_features : num_price_features + num_news_features - ] - social_features = x[ - :, - :, - num_price_features + num_news_features : num_price_features - + num_news_features - + num_social_features, - ] - corporate_action_flags = x[:, :, -1:] - - price_processed = self.price_processor(price_features) - news_processed = self.news_processor(news_features) - social_processed = self.social_processor(social_features) - corporate_action_features = self.corporate_action_mlp(corporate_action_flags) - - price_attended = self.temporal_attention(price_processed) - news_attended = self.temporal_attention(news_processed) - social_attended = self.temporal_attention(social_processed) + price_attended = self.temporal_attention(self.price_processor(price_features)) + news_attended = self.temporal_attention(self.news_processor(news_features)) + social_attended = self.temporal_attention(self.social_processor(social_features)) alternative_features = torch.cat( [price_attended, news_attended, social_attended], dim=1 ) - - attention_weights = self.alternative_data_attention(alternative_features) - - return attention_weights + return self.alternative_data_attention(alternative_features) diff --git a/src/models/trainer.py b/src/models/trainer.py index c72544b..f35ff4e 100644 --- a/src/models/trainer.py +++ b/src/models/trainer.py @@ -6,7 +6,8 @@ from typing import Dict, List, Tuple import numpy as np import torch import torch.nn as nn -from torch.cuda.amp import GradScaler, autocast +from torch.amp import autocast +from torch.cuda.amp import GradScaler from torch_geometric.loader import DataLoader from tqdm import tqdm @@ -92,6 +93,7 @@ class GNNTrainer: # Training self.model.train() epoch_train_loss = 0.0 + num_train_batches = 0 start_time = time.time() for batch in tqdm( @@ -110,6 +112,7 @@ class GNNTrainer: # Mixed precision training with autocast( + config.DEVICE, enabled=config.MIXED_PRECISION, dtype=self.amd_optimizer.get_precision_dtype(), ): @@ -130,6 +133,7 @@ class GNNTrainer: self.scaler.update() epoch_train_loss += loss.item() + num_train_batches += 1 # Memory management self.memory_manager.auto_manage_memory(threshold=0.8) @@ -139,12 +143,13 @@ class GNNTrainer: self.memory_manager.empty_cache() continue - epoch_train_loss /= len(train_loader) + epoch_train_loss = epoch_train_loss / num_train_batches if num_train_batches else 0.0 self.train_losses.append(epoch_train_loss) # Validation self.model.eval() epoch_val_loss = 0.0 + num_val_batches = 0 with torch.no_grad(): for batch in val_loader: @@ -159,6 +164,7 @@ class GNNTrainer: batch = batch.to(self.device, non_blocking=config.PIN_MEMORY) with autocast( + config.DEVICE, enabled=config.MIXED_PRECISION, dtype=self.amd_optimizer.get_precision_dtype(), ): @@ -166,6 +172,7 @@ class GNNTrainer: loss = self.criterion(out, batch.y) epoch_val_loss += loss.item() + num_val_batches += 1 # Memory management self.memory_manager.auto_manage_memory(threshold=0.8) @@ -178,7 +185,7 @@ class GNNTrainer: self.memory_manager.empty_cache() continue - epoch_val_loss /= len(val_loader) + epoch_val_loss = epoch_val_loss / num_val_batches if num_val_batches else 0.0 self.val_losses.append(epoch_val_loss) # Update learning rate scheduler