import logging import os import pickle import sqlite3 from datetime import datetime, timedelta from typing import Dict, List, Optional, Tuple import numpy as np import pandas as pd import torch from tqdm import tqdm from config import config from src.data.news_processor import NewsProcessor from src.data.social_processor import SocialMediaProcessor from src.utils.helpers import generate_intraday_timestamps from src.utils.memory_manager import MemoryManager logger = logging.getLogger(__name__) class StockDataPipeline: def __init__(self): self.db_path = os.path.join(config.PROCESSED_DATA_DIR, "stock_data.db") self._init_db() # Initialize memory manager self.memory_manager = MemoryManager() # Initialize components self.price_data = {} self.corporate_actions = {} self.sector_data = {} self.index_composition = {} self.delisted_tickers = set() self.intraday_data = {} # Initialize alternative data processors self.news_processor = NewsProcessor() self.social_processor = SocialMediaProcessor() # Load existing data if available self._load_data() # Load delisted tickers self._load_delisted_tickers() def _init_db(self): """Initialize the SQLite database with AMD-optimized settings""" with sqlite3.connect(self.db_path) as conn: cursor = conn.cursor() # Enable WAL mode for better concurrency (important for AMD GPUs) cursor.execute("PRAGMA journal_mode=WAL") # Increase cache size for better performance cursor.execute("PRAGMA cache_size=-10000") # 10MB cache # Create tables for different data types cursor.execute( """ CREATE TABLE IF NOT EXISTS price_data ( id INTEGER PRIMARY KEY AUTOINCREMENT, ticker TEXT, date TEXT, open REAL, high REAL, low REAL, close REAL, adj_close REAL, volume INTEGER, UNIQUE(ticker, date) ) """ ) cursor.execute( """ CREATE TABLE IF NOT EXISTS corporate_actions ( id INTEGER PRIMARY KEY AUTOINCREMENT, ticker TEXT, date TEXT, action_type TEXT, value REAL, details TEXT, UNIQUE(ticker, date, action_type) ) """ ) cursor.execute( """ CREATE TABLE IF NOT EXISTS sector_data ( id INTEGER PRIMARY KEY AUTOINCREMENT, ticker TEXT, date TEXT, sector TEXT, industry TEXT, UNIQUE(ticker, date) ) """ ) cursor.execute( """ CREATE TABLE IF NOT EXISTS index_composition ( id INTEGER PRIMARY KEY AUTOINCREMENT, index_ticker TEXT, date TEXT, member_ticker TEXT, UNIQUE(index_ticker, date, member_ticker) ) """ ) cursor.execute( """ CREATE TABLE IF NOT EXISTS features ( id INTEGER PRIMARY KEY AUTOINCREMENT, ticker TEXT, date TEXT, feature_name TEXT, feature_value REAL, UNIQUE(ticker, date, feature_name) ) """ ) # Create indexes for faster queries cursor.execute( "CREATE INDEX IF NOT EXISTS idx_price_data_ticker_date ON price_data(ticker, date)" ) cursor.execute( "CREATE INDEX IF NOT EXISTS idx_corporate_actions_ticker_date ON corporate_actions(ticker, date)" ) cursor.execute( "CREATE INDEX IF NOT EXISTS idx_sector_data_ticker_date ON sector_data(ticker, date)" ) cursor.execute( "CREATE INDEX IF NOT EXISTS idx_index_composition_date ON index_composition(date)" ) cursor.execute( "CREATE INDEX IF NOT EXISTS idx_features_ticker_date ON features(ticker, date)" ) conn.commit() def _load_data(self): """Load existing data from disk with memory management""" if not self.memory_manager.ensure_memory(2 * 1024**3): # 2GB logger.warning("Skipping data load due to memory constraints") return 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() self.memory_manager.log_memory_usage("[After Data Load]") def _save_data(self): """Save data to disk with memory management""" try: # Check memory before saving if not self.memory_manager.ensure_memory(1 * 1024**3): # 1GB logger.warning("Skipping data save due to memory constraints") return self._save_pickle(self.price_data, "price_data.pkl") self._save_pickle(self.corporate_actions, "corporate_actions.pkl") self._save_pickle(self.sector_data, "sector_data.pkl") self._save_pickle(self.index_composition, "index_composition.pkl") logger.info("Saved data to disk") self.memory_manager.log_memory_usage("[After Data Save]") except Exception as e: logger.error(f"Error saving data: {str(e)}", exc_info=True) self.memory_manager.empty_cache() def _load_pickle(self, filename: str): """Load data from pickle file with memory management""" filepath = os.path.join(config.PROCESSED_DATA_DIR, filename) if not os.path.exists(filepath): raise FileNotFoundError(f"File {filepath} not found") with open(filepath, "rb") as f: return pickle.load(f) def _save_pickle(self, data, filename: str): """Save data to pickle file with memory management""" filepath = os.path.join(config.PROCESSED_DATA_DIR, filename) with open(filepath, "wb") as f: pickle.dump(data, f) def _load_delisted_tickers(self): """Load delisted tickers from external file""" if os.path.exists(config.DELISTED_TICKERS_FILE): try: df = pd.read_csv(config.DELISTED_TICKERS_FILE) self.delisted_tickers = set(df["Ticker"].tolist()) logger.info(f"Loaded {len(self.delisted_tickers)} delisted tickers") except Exception as e: logger.error(f"Error loading delisted tickers: {str(e)}") else: logger.warning( "Delisted tickers file not found. Only using active tickers." ) def update_all_data(self): """Update all data sources with memory management""" # Get all tickers (active + delisted) all_tickers = config.INITIAL_TICKERS + list(self.delisted_tickers) # Update price data self.update_price_data(all_tickers) # Update corporate actions self.update_corporate_actions(all_tickers) # Update sector data self.update_sector_data(all_tickers) # Update index composition self.update_index_composition() # Update alternative data self.update_alternative_data(all_tickers) # Save updated data self._save_data() def update_price_data(self, tickers: List[str]): """Update price data for given tickers with memory management""" logger.info(f"Updating price data for {len(tickers)} tickers") # Use yfinance for historical data import yfinance as yf for ticker in tqdm(tickers, desc="Updating price data"): try: # Check memory before processing if not self.memory_manager.ensure_memory(100 * 1024**2): # 100MB logger.warning(f"Skipping {ticker} due to memory constraints") continue # Determine start date start_date = config.START_DATE if ticker in self.price_data and not self.price_data[ticker].empty: # If we already have data, start from the day after our last data point start_date = ( self.price_data[ticker].index[-1] + timedelta(days=1) ).strftime("%Y-%m-%d") # Download new data new_data = yf.download( ticker, start=start_date, end=config.END_DATE, progress=False, auto_adjust=True, # Use adjusted prices ) if not new_data.empty: # If we already have data for this ticker, append the new data if ticker in self.price_data and not self.price_data[ticker].empty: # Combine existing and new data combined = pd.concat([self.price_data[ticker], new_data]) # Remove duplicates (keep the new data) combined = combined[~combined.index.duplicated(keep="last")] self.price_data[ticker] = combined.sort_index() else: self.price_data[ticker] = new_data.sort_index() # Store in database self._store_price_data(ticker, new_data) # Memory management self.memory_manager.auto_manage_memory(threshold=0.7) except Exception as e: logger.warning(f"Error updating price data for {ticker}: {str(e)}") self.memory_manager.empty_cache() continue def _store_price_data(self, ticker: str, data: pd.DataFrame): """Store price data in the database""" if data.empty: return # Convert to list of dictionaries for bulk insert price_data = [] for date, row in data.iterrows(): price_data.append( { "ticker": ticker, "date": date.strftime("%Y-%m-%d"), "open": row["Open"], "high": row["High"], "low": row["Low"], "close": row["Close"], "adj_close": row[ "Close" ], # auto_adjust=True, 'Close' is already adjusted "volume": row["Volume"], } ) # Store in database with sqlite3.connect(self.db_path) as conn: cursor = conn.cursor() cursor.executemany( """ INSERT OR REPLACE INTO price_data (ticker, date, open, high, low, close, adj_close, volume) VALUES (:ticker, :date, :open, :high, :low, :close, :adj_close, :volume) """, price_data, ) conn.commit() def update_corporate_actions(self, tickers: List[str]): """Update corporate actions for given tickers with memory management""" logger.info(f"Updating corporate actions for {len(tickers)} tickers") import yfinance as yf for ticker in tqdm(tickers, desc="Updating corporate actions"): try: # Check memory before processing if not self.memory_manager.ensure_memory(50 * 1024**2): # 50MB logger.warning(f"Skipping {ticker} due to memory constraints") continue stock = yf.Ticker(ticker) # Initialize corporate actions dictionary if needed if ticker not in self.corporate_actions: self.corporate_actions[ticker] = { "splits": {}, "dividends": {}, "mergers": [], "spin-offs": [], } # Get splits splits = stock.splits if not splits.empty: for date, ratio in splits.items(): date_str = date.strftime("%Y-%m-%d") self.corporate_actions[ticker]["splits"][date_str] = float( ratio ) # Get dividends dividends = stock.dividends if not dividends.empty: for date, amount in dividends.items(): date_str = date.strftime("%Y-%m-%d") self.corporate_actions[ticker]["dividends"][date_str] = float( amount ) # Store in database self._store_corporate_actions(ticker) # Memory management self.memory_manager.auto_manage_memory(threshold=0.7) except Exception as e: logger.warning( f"Error updating corporate actions for {ticker}: {str(e)}" ) self.memory_manager.empty_cache() continue def _store_corporate_actions(self, ticker: str): """Store corporate actions in the database using a single connection.""" if ticker not in self.corporate_actions: return actions = self.corporate_actions[ticker] rows = [] for date, ratio in actions["splits"].items(): rows.append((ticker, date, "split", ratio, f"Split ratio: {ratio}")) for date, amount in actions["dividends"].items(): rows.append( (ticker, date, "dividend", amount, f"Dividend amount: {amount}") ) if not rows: return with sqlite3.connect(self.db_path) as conn: conn.executemany( """ INSERT OR REPLACE INTO corporate_actions (ticker, date, action_type, value, details) VALUES (?, ?, ?, ?, ?) """, rows, ) def update_sector_data(self, tickers: List[str]): """Update sector data for given tickers with memory management""" logger.info(f"Updating sector data for {len(tickers)} tickers") import yfinance as yf for ticker in tqdm(tickers, desc="Updating sector data"): try: # Check memory before processing if not self.memory_manager.ensure_memory(50 * 1024**2): # 50MB logger.warning(f"Skipping {ticker} due to memory constraints") continue stock = yf.Ticker(ticker) info = stock.info if "sector" in info: self.sector_data[ticker] = { "sector": info["sector"], "industry": info.get("industry", "Unknown"), } # Store in database with sqlite3.connect(self.db_path) as conn: cursor = conn.cursor() cursor.execute( """ INSERT OR REPLACE INTO sector_data (ticker, date, sector, industry) VALUES (?, ?, ?, ?) """, ( ticker, config.END_DATE, info["sector"], info.get("industry", "Unknown"), ), ) # Memory management self.memory_manager.auto_manage_memory(threshold=0.7) except Exception as e: logger.warning(f"Error updating sector data for {ticker}: {str(e)}") self.memory_manager.empty_cache() continue def update_index_composition(self): """Update index composition with memory management""" logger.info(f"Updating index composition for {config.INDEX_TICKER}") import yfinance as yf try: # Check memory before processing if not self.memory_manager.ensure_memory(100 * 1024**2): # 100MB logger.warning( "Skipping index composition update due to memory constraints" ) return # Get current constituents index = yf.Ticker(config.INDEX_TICKER) constituents = index.get_holders() if constituents is not None and not constituents.empty: current_date = datetime.now().strftime("%Y-%m-%d") if config.INDEX_TICKER not in self.index_composition: self.index_composition[config.INDEX_TICKER] = {} self.index_composition[config.INDEX_TICKER][current_date] = ( constituents["Symbol"].tolist() ) # Store in database self._store_index_composition() # Memory management self.memory_manager.auto_manage_memory(threshold=0.7) except Exception as e: logger.warning(f"Error updating index composition: {str(e)}") self.memory_manager.empty_cache() def _store_index_composition(self): """Store index composition in the database""" 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]): """ Update alternative data sources for given tickers with memory management Parameters: tickers: List of tickers to update """ # Calculate date range end_date = datetime.now().strftime("%Y-%m-%d") start_date = ( datetime.now() - timedelta( days=max(config.NEWS_LOOKBACK_DAYS, config.SOCIAL_MEDIA_LOOKBACK_DAYS) ) ).strftime("%Y-%m-%d") logger.info(f"Updating alternative data from {start_date} to {end_date}") # Update news data self.news_processor.fetch_news(tickers, start_date, end_date) # Update social media data self.social_processor.fetch_twitter_data(tickers, start_date, end_date) self.social_processor.fetch_reddit_data(tickers, start_date, end_date) def _compute_stock_sequence( self, ticker: str, date, pit_data: Dict ) -> Optional[np.ndarray]: """ Build a (SEQUENCE_LENGTH, NUM_FEATURES) float32 array for one stock at one date. NUM_FEATURES = 5 price + len(NEWS_FEATURES) + len(SOCIAL_FEATURES) + 1 corp action """ price_df = self.price_data.get(ticker) if price_df is None or price_df.empty: return None # Fetch enough history to compute SEQUENCE_LENGTH trading-day feature vectors buf_start = date - timedelta(days=config.SEQUENCE_LENGTH * 2 + 60) window = price_df.loc[buf_start:date] if len(window) < 2: return None closes = window["Close"].values.astype(np.float64) volumes = window["Volume"].values.astype(np.float64) raw_returns = np.diff(closes) / (closes[:-1] + 1e-10) n = len(raw_returns) seq_rows: List[List[float]] = [] for i in range(n): ret = float(raw_returns[i]) recent = raw_returns[max(0, i - 19) : i + 1] vol = float(np.std(recent)) if len(recent) > 1 else 0.0 mom = float(np.mean(recent)) log_vol = float(np.log(volumes[i + 1] + 1)) norm_price = float(closes[i + 1] / (closes[0] + 1e-10) - 1) seq_rows.append([ret, vol, mom, log_vol, norm_price]) seq_arr = np.array(seq_rows, dtype=np.float32) # Trim to SEQUENCE_LENGTH (pad with zeros if history is too short) if len(seq_arr) > config.SEQUENCE_LENGTH: seq_arr = seq_arr[-config.SEQUENCE_LENGTH :] elif len(seq_arr) < config.SEQUENCE_LENGTH: pad = np.zeros((config.SEQUENCE_LENGTH - len(seq_arr), 5), dtype=np.float32) seq_arr = np.vstack([pad, seq_arr]) # News features (current date, broadcast across all timesteps) date_str = date.strftime("%Y-%m-%d") nf = self.news_processor.get_news_features(ticker, date_str) news_vec = np.array( [ nf.get("news_sentiment", 0.0), nf.get("news_volume", 0.0), nf.get("news_recency", 0.0), nf.get("news_source_reliability", 0.0), nf.get("news_topic_relevance", 0.0), ], dtype=np.float32, ) # Social features (current date, broadcast) sf = self.social_processor.get_social_features(ticker, date_str) social_vec = np.array( [ sf.get("twitter_sentiment", 0.0), sf.get("twitter_volume", 0.0), sf.get("reddit_sentiment", 0.0), sf.get("reddit_volume", 0.0), sf.get("social_momentum", 0.0), ], dtype=np.float32, ) # Corporate action flag corp_flag = 0.0 if pit_data["upcoming_actions"]: soonest = min(pit_data["upcoming_actions"], key=lambda a: a["days_until"]) corp_flag = 1.0 if soonest["type"] == "split" else 2.0 alt_vec = np.concatenate([news_vec, social_vec, [corp_flag]]) # (11,) alt_broadcast = np.tile(alt_vec, (config.SEQUENCE_LENGTH, 1)) # (seq_len, 11) return np.concatenate([seq_arr, alt_broadcast], axis=1) # (seq_len, 16) def _build_edges( self, tickers: List[str], date, pit_cache: Dict ) -> Tuple[torch.Tensor, torch.Tensor]: """Build same-sector edges with return-correlation weights.""" lookback_start = date - timedelta(days=config.LOOKBACK_WINDOW) edge_index: List[List[int]] = [] edge_weight: List[float] = [] for i, t1 in enumerate(tickers): for j, t2 in enumerate(tickers): if j <= i: continue s1 = pit_cache[t1].get("sector") s2 = pit_cache[t2].get("sector") if not s1 or not s2 or s1 != s2: continue try: r1 = ( self.price_data[t1] .loc[lookback_start:date]["Close"] .pct_change() .dropna() ) r2 = ( self.price_data[t2] .loc[lookback_start:date]["Close"] .pct_change() .dropna() ) if len(r1) > 5 and len(r2) > 5: corr = float(r1.corr(r2)) if not np.isnan(corr): edge_index.append([i, j]) edge_weight.append(corr) except Exception as e: logger.warning(f"Edge {t1}-{t2}: {e}") if edge_index: ei = torch.tensor(edge_index, dtype=torch.long).t().contiguous() ew = torch.tensor(edge_weight, dtype=torch.float32).unsqueeze(1) else: ei = torch.empty((2, 0), dtype=torch.long) ew = torch.empty((0, 1), dtype=torch.float32) return ei, ew def _next_trading_date(self, date) -> Optional[object]: """Return the next date that has price data for at least one ticker.""" candidate = date + timedelta(days=1) for _ in range(7): for price_df in self.price_data.values(): if not price_df.empty and candidate in price_df.index: return candidate candidate += timedelta(days=1) return None def _create_dataset(self, start_date: str, end_date: str) -> List: """ Core dataset builder producing PyG Data objects with 3-D node features. x shape per object: (num_stocks, SEQUENCE_LENGTH, NUM_FEATURES) where NUM_FEATURES = 5 price + 5 news + 5 social + 1 corp_action = 16 """ from torch_geometric.data import Data all_tickers = config.INITIAL_TICKERS + list(self.delisted_tickers) dates = pd.bdate_range(start_date, end_date) # business days only dataset = [] for date in tqdm(dates, desc=f"Dataset {start_date}→{end_date}"): if not self.memory_manager.ensure_memory(500 * 1024**2): logger.warning(f"Skipping {date.date()} — low memory") self.memory_manager.empty_cache() continue # Stocks with price data spanning this date valid_tickers = [ t for t in all_tickers if t in self.price_data and not self.price_data[t].empty and self.price_data[t].index[0] <= date <= self.price_data[t].index[-1] ] if not valid_tickers: continue # Pre-fetch PIT data once per ticker (avoids O(n²) repeated calls) pit_cache = {t: self.get_point_in_time_data(t, date) for t in valid_tickers} # Build 3-D node feature matrix node_sequences: List[np.ndarray] = [] node_tickers: List[str] = [] for ticker in valid_tickers: if not self.memory_manager.ensure_memory(10 * 1024**2): continue try: seq = self._compute_stock_sequence(ticker, date, pit_cache[ticker]) if seq is not None: node_sequences.append(seq) node_tickers.append(ticker) except Exception as e: logger.warning(f"Sequence error {ticker} {date.date()}: {e}") if not node_sequences: continue # x: (num_stocks, seq_len, num_features) x = torch.tensor(np.array(node_sequences), dtype=torch.float32) # Cross-sectional z-score per feature per time step (normalise across stocks) mean = x.mean(dim=0, keepdim=True) std = x.std(dim=0, keepdim=True).clamp(min=1e-8) x = (x - mean) / std # Build graph edges edge_index, edge_weight = self._build_edges(node_tickers, date, pit_cache) # Build targets: next trading day's return for each stock. # Skip dates where no next trading date exists (e.g. tail end of dataset) # to avoid polluting training labels with fabricated zero returns. next_date = self._next_trading_date(date) if next_date is None: continue y_vals: List[float] = [] valid_seqs_for_y: List[np.ndarray] = [] valid_tickers_for_y: List[str] = [] for i, ticker in enumerate(node_tickers): if ( ticker in self.price_data and next_date in self.price_data[ticker].index and date in self.price_data[ticker].index ): ret = float( self.price_data[ticker].loc[next_date]["Close"] / self.price_data[ticker].loc[date]["Close"] - 1 ) y_vals.append(ret) valid_seqs_for_y.append(node_sequences[i]) valid_tickers_for_y.append(ticker) if not y_vals: continue if len(valid_tickers_for_y) < len(node_tickers): # Rebuild x and edges using only stocks with valid targets x = torch.tensor(np.array(valid_seqs_for_y), dtype=torch.float32) mean = x.mean(dim=0, keepdim=True) std = x.std(dim=0, keepdim=True).clamp(min=1e-8) x = (x - mean) / std edge_index, edge_weight = self._build_edges( valid_tickers_for_y, date, pit_cache ) node_tickers = valid_tickers_for_y y_arr = np.array(y_vals, dtype=np.float32) if len(y_arr) > 1: mu, sigma = y_arr.mean(), y_arr.std() if sigma > 1e-8: y_arr = (y_arr - mu) / sigma y = torch.tensor(y_arr, dtype=torch.float32).unsqueeze(1) data = Data(x=x, edge_index=edge_index, edge_attr=edge_weight, y=y) data.date = date data.tickers = node_tickers dataset.append(data) self.memory_manager.auto_manage_memory(threshold=0.7) return dataset def create_training_dataset(self) -> List: """Create the training dataset for config.START_DATE → config.TRAIN_END_DATE.""" logger.info("Creating training dataset") return self._create_dataset(config.START_DATE, config.TRAIN_END_DATE) def create_intraday_dataset( self, tickers: List[str], start_date: str, end_date: str ) -> List: """ Create a dataset for intraday trading with AMD optimizations Parameters: tickers: List of tickers start_date: Start date (YYYY-MM-DD) end_date: End date (YYYY-MM-DD) Returns: List of PyG Data objects for intraday trading """ from torch_geometric.data import Data logger.info( f"Creating intraday dataset from {start_date} to {end_date} with AMD optimizations" ) dataset = [] # Generate all trading days in the date range date_range = pd.date_range(start_date, end_date) trading_days = [ date.strftime("%Y-%m-%d") for date in date_range if date.weekday() < 5 ] # Weekdays only for date in trading_days: # Check memory before processing date if not self.memory_manager.ensure_memory(1 * 1024**3): # 1GB logger.warning(f"Skipping {date} due to memory constraints") 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 } # Determine next trading date once per day next_date = self._next_trading_date(pd_date) if next_date is None: continue # Pre-filter tickers that have target data on the next trading date target_valid_tickers = [ t for t in current_tickers if next_date in self.price_data[t].index ] if not target_valid_tickers: continue # Get all timestamps for this trading day timestamps = generate_intraday_timestamps(date) for i in range(config.SEQUENCE_LENGTH, len(timestamps)): current_timestamp = timestamps[i] sequence_start = timestamps[i - config.SEQUENCE_LENGTH] # Create node features for each stock in the sequence sequence_features = [] valid_tickers = [] for ticker in target_valid_tickers: # Check memory before processing ticker if not self.memory_manager.ensure_memory(50 * 1024**2): # 50MB logger.warning(f"Skipping {ticker} due to memory constraints") continue try: # Get price data for lookback period lookback_start = datetime.strptime( sequence_start, "%Y-%m-%d %H:%M:%S" ) lookback_end = datetime.strptime( current_timestamp, "%Y-%m-%d %H:%M:%S" ) # For intraday data, we need to get it from the live data service # In a real implementation, this would come from the live data feed # For this example, we'll simulate it with daily data # Get daily data for the date daily_data = self.price_data[ticker].loc[date] if isinstance(daily_data, pd.Series): # If only one day of data, create a sequence with the same values features = [ daily_data["Close"] / daily_data["Open"] - 1, # Return 0.2, # Volatility (placeholder) 0.0, # Momentum (placeholder) np.log(daily_data["Volume"] + 1), # Log volume daily_data["Close"], # Price (auto_adjust=True) ] # Repeat for the sequence sequence_data = np.tile( features, (config.SEQUENCE_LENGTH, 1) ) else: # This shouldn't happen with daily data continue sequence_features.append(sequence_data) valid_tickers.append(ticker) except Exception as e: logger.warning( f"Error processing {ticker} for {current_timestamp}: {str(e)}" ) continue # Skip if no features were created if not sequence_features: continue # Convert to tensor (num_stocks, sequence_length, num_features) x = torch.tensor(np.array(sequence_features), dtype=torch.float32) # Create edges based on sector relationships edge_index = [] edge_weight = [] for j, ticker1 in enumerate(valid_tickers): for k, ticker2 in enumerate(valid_tickers): if j < k: # Check memory before processing edge if not self.memory_manager.ensure_memory( 1 * 1024**2 ): # 1MB logger.warning( "Skipping edge creation due to memory constraints" ) continue try: # Use pre-fetched PIT data — avoids O(n²) repeated calls pit1 = pit_cache[ticker1] pit2 = pit_cache[ticker2] if ( pit1["sector"] and pit2["sector"] and pit1["sector"] == pit2["sector"] ): # Calculate correlation of recent returns # For intraday, we would use intraday returns # For this example, we'll use daily returns lookback_start = datetime.strptime( date, "%Y-%m-%d" ) - timedelta(days=config.LOOKBACK_WINDOW) returns1 = ( self.price_data[ticker1] .loc[lookback_start:date]["Close"] .pct_change() ) returns2 = ( self.price_data[ticker2] .loc[lookback_start:date]["Close"] .pct_change() ) if len(returns1) > 5 and len(returns2) > 5: corr = returns1.corr(returns2) if not np.isnan(corr): edge_index.append([j, k]) edge_weight.append(corr) except Exception as e: logger.warning( f"Error creating edge between {ticker1} and {ticker2}: {str(e)}" ) continue # Convert to tensors edge_index = ( torch.tensor(edge_index, dtype=torch.long).t() if edge_index else torch.empty((2, 0), dtype=torch.long) ) edge_weight = ( torch.tensor(edge_weight, dtype=torch.float32).unsqueeze(1) if edge_weight else torch.empty((0, 1), dtype=torch.float32) ) # Create target (returns over prediction horizon) y_vals = [] for ticker in valid_tickers: current_price = self.price_data[ticker].loc[date]["Close"] future_price = self.price_data[ticker].loc[next_date]["Close"] ret = future_price / current_price - 1 y_vals.append(ret) y_arr = np.array(y_vals, dtype=np.float32) if len(y_arr) > 1: mu, sigma = y_arr.mean(), y_arr.std() if sigma > 1e-8: y_arr = (y_arr - mu) / sigma y = torch.tensor(y_arr, dtype=torch.float32).unsqueeze(1) # Create Data object data = Data(x=x, edge_index=edge_index, edge_attr=edge_weight, y=y) data.timestamp = current_timestamp data.date = date data.tickers = valid_tickers data.sequence_start = sequence_start data.sequence_end = current_timestamp dataset.append(data) # Memory management self.memory_manager.auto_manage_memory(threshold=0.7) return dataset def get_point_in_time_data(self, ticker: str, date: datetime) -> Dict: """ Get point-in-time data for a stock at a specific date Parameters: ticker: Stock ticker date: Date as datetime object Returns: Dictionary with point-in-time data """ date_str = date.strftime("%Y-%m-%d") result = { "ticker": ticker, "date": date_str, "price": None, "sector": None, "industry": None, "in_index": False, "index": None, "upcoming_actions": [], } # Get price data if ticker in self.price_data and not self.price_data[ticker].empty: # Find the most recent price before or on the date price_data = self.price_data[ticker] idx = price_data.index.get_indexer([date], method="ffill")[0] if idx >= 0: result["price"] = price_data.iloc[idx]["Close"] # Get sector data if ticker in self.sector_data: result["sector"] = self.sector_data[ticker]["sector"] result["industry"] = self.sector_data[ticker]["industry"] # Check if in index (simplified - would need historical index composition) for index_ticker, composition in self.index_composition.items(): # Find the most recent composition before the date comp_dates = sorted(composition.keys()) for comp_date in reversed(comp_dates): if datetime.strptime(comp_date, "%Y-%m-%d") <= date: if ticker in composition[comp_date]: result["in_index"] = True result["index"] = index_ticker break # Get upcoming corporate actions if ticker in self.corporate_actions: actions = self.corporate_actions[ticker] # Check for upcoming splits (within 30 days) for action_date, ratio in actions["splits"].items(): action_date_dt = datetime.strptime(action_date, "%Y-%m-%d") if date < action_date_dt <= date + timedelta(days=30): result["upcoming_actions"].append( { "type": "split", "date": action_date, "ratio": ratio, "days_until": (action_date_dt - date).days, } ) # Check for upcoming dividends (within 7 days) for action_date, amount in actions["dividends"].items(): action_date_dt = datetime.strptime(action_date, "%Y-%m-%d") if date < action_date_dt <= date + timedelta(days=7): result["upcoming_actions"].append( { "type": "dividend", "date": action_date, "amount": amount, "days_until": (action_date_dt - date).days, } ) return result def get_latest_features( self, tickers: List[str], timestamp: str ) -> Dict[str, Dict]: """ Get the latest features for given tickers at a specific timestamp Parameters: tickers: List of tickers timestamp: Timestamp (YYYY-MM-DD HH:MM:SS) Returns: Dictionary of {ticker: features} for the latest available data """ features = {} for ticker in tickers: # Get price features price_features = self._get_latest_price_features(ticker, timestamp) # Get alternative data features news_features = self.news_processor.get_news_features( ticker, timestamp.split(" ")[0] ) social_features = self.social_processor.get_all_social_features( [ticker], timestamp.split(" ")[0] ) # Combine features if price_features: combined_features = { "ticker": ticker, "timestamp": timestamp, **price_features, **news_features, **( social_features.iloc[0].to_dict() if not social_features.empty else {} ), } features[ticker] = combined_features return features def _get_latest_price_features(self, ticker: str, timestamp: str) -> Dict: """ Get the latest price features for a ticker at a specific timestamp Parameters: ticker: Stock ticker timestamp: Timestamp (YYYY-MM-DD HH:MM:SS) Returns: Dictionary of price features """ if ticker not in self.price_data or self.price_data[ticker].empty: return {} # Get the date from the timestamp date_str = timestamp.split(" ")[0] date = datetime.strptime(date_str, "%Y-%m-%d") # Get price data for the date if date_str not in self.price_data[ticker].index: return {} price_data = self.price_data[ticker].loc[date_str] # Calculate features features = { "return": 0.0, # For intraday, this would be the return since open "volatility": 0.2, # Placeholder - would calculate from intraday data "momentum": 0.0, # Placeholder "volume": np.log(price_data["Volume"] + 1), "price": price_data["Close"], } return features def create_validation_dataset(self, start_date: str, end_date: str) -> List: """Create a validation dataset for the given date range.""" logger.info(f"Creating validation dataset {start_date} → {end_date}") return self._create_dataset(start_date, end_date) def create_walk_forward_splits(self) -> List[Tuple]: """ Build walk-forward CV splits as expanding-window train / fixed OOS val pairs. Each fold's training window starts at config.START_DATE and ends at a boundary that advances by config.WALK_FORWARD_TEST_YEARS each fold. The val window is the following config.WALK_FORWARD_TEST_YEARS of data. Folds that would push the val end beyond config.TRAIN_END_DATE (or available data) are dropped. Returns a list of (train_dataset, val_dataset, meta_dict) tuples. """ from datetime import datetime, timedelta train_start = config.START_DATE full_end = datetime.strptime(config.TRAIN_END_DATE, "%Y-%m-%d") test_years = config.WALK_FORWARD_TEST_YEARS n_folds = config.WALK_FORWARD_FOLDS splits: List[Tuple] = [] for fold in range(n_folds): # Each fold's val window is offset by fold * test_years from the full_end # minus (n_folds - fold) * test_years so folds are evenly spaced offset_years = (n_folds - fold) * test_years val_end_dt = full_end - timedelta(days=int((offset_years - test_years) * 365)) train_end_dt = val_end_dt - timedelta(days=int(test_years * 365)) val_start_dt = train_end_dt + timedelta(days=1) if train_end_dt <= datetime.strptime(train_start, "%Y-%m-%d"): logger.warning(f"Skipping fold {fold}: train window too short") continue train_end_str = train_end_dt.strftime("%Y-%m-%d") val_start_str = val_start_dt.strftime("%Y-%m-%d") val_end_str = val_end_dt.strftime("%Y-%m-%d") logger.info( f"Walk-forward fold {fold + 1}/{n_folds}: " f"train {train_start}→{train_end_str}, val {val_start_str}→{val_end_str}" ) train_ds = self._create_dataset(train_start, train_end_str) val_ds = self._create_dataset(val_start_str, val_end_str) meta = { "fold": fold, "train_start": train_start, "train_end": train_end_str, "val_start": val_start_str, "val_end": val_end_str, } splits.append((train_ds, val_ds, meta)) return splits