fegger
19ed77f4a0
Downsize GNN hidden channels/heads (128→64, 16→4) to match the 30-stock price-only universe before alternative data processors are ready. Add USE_ALTERNATIVE_DATA flag (default False) that skips news/social branches so the model trains on real data rather than zero-filled stubs. Standardize target returns per-batch in the data pipeline to stabilize training. Introduce walk-forward cross-validation with expanding windows: configurable fold count and out-of-sample years. Add IC loss weighting (0.7) to complement MSE, and wire it into the trainer alongside the new loss function.
1246 lines
48 KiB
Python
1246 lines
48 KiB
Python
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
|