initial commit

This commit is contained in:
2026-05-26 12:43:47 +02:00
commit b43fbe2170
20 changed files with 1619 additions and 0 deletions
+30
View File
@@ -0,0 +1,30 @@
stock_gnn_project/
├── data/
│ ├── raw/
│ │ ├── news/ # Raw news data
│ │ ├── social_media/ # Raw social media data
│ │ └── ...
│ ├── processed/
│ │ ├── news_features/ # Processed news features
│ │ ├── social_features/ # Processed social media features
│ │ └── ...
│ └── external/
├── src/
│ ├── data/
│ │ ├── news_processor.py # News data processing
│ │ ├── social_processor.py # Social media processing
│ │ ├── sentiment.py # Sentiment analysis
│ │ └── ...
│ │
│ ├── models/
│ │ ├── gnn_model.py # Updated GNN model
│ │ └── ...
│ │
│ └── utils/
│ ├── text_processing.py # Text processing utilities
│ └── ...
├── config.py # Updated configuration
└── ...
Binary file not shown.
Binary file not shown.
+103
View File
@@ -0,0 +1,103 @@
import os
from datetime import datetime
class Config:
# Data settings
DATA_DIR = os.path.join(os.path.dirname(__file__), "data")
RAW_DATA_DIR = os.path.join(DATA_DIR, "raw")
PROCESSED_DATA_DIR = os.path.join(DATA_DIR, "processed")
EXTERNAL_DATA_DIR = os.path.join(DATA_DIR, "external")
# 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)
# Stock universe settings
INITIAL_TICKERS = [
"AAPL",
"MSFT",
"GOOGL",
"AMZN",
"META",
"TSLA",
"NVDA",
"JPM",
"V",
"WMT",
"PG",
"DIS",
"NFLX",
"ADBE",
"PYPL",
"INTC",
"CSCO",
"PEP",
"KO",
"XOM",
]
INDEX_TICKER = "^GSPC" # S&P 500
DELISTED_TICKERS_FILE = os.path.join(EXTERNAL_DATA_DIR, "delisted_stocks.csv")
# Date settings
START_DATE = "2010-01-01"
END_DATE = datetime.now().strftime("%Y-%m-%d")
TRAIN_END_DATE = "2020-12-31"
VAL_END_DATE = "2021-12-31"
# Model settings
MODEL_DIR = os.path.join(os.path.dirname(__file__), "models")
os.makedirs(MODEL_DIR, exist_ok=True)
MODEL_NAME = "stock_gnn"
HIDDEN_CHANNELS = 64
NUM_HEADS = 8
DROPOUT = 0.6
LEARNING_RATE = 0.001
EPOCHS = 100
BATCH_SIZE = 32
LOOKBACK_WINDOW = 30 # Days for feature calculation
# Backtesting settings
INITIAL_CAPITAL = 100000
TRANSACTION_COST = 0.001 # 0.1% per trade
# Evaluation settings
BENCHMARK_TICKER = "^GSPC"
# News data settings
NEWS_API_KEY = "your_news_api_key" # For NewsAPI or similar
NEWS_SOURCES = ["reuters", "bloomberg", "financial-times", "wsj"]
NEWS_CATEGORIES = ["business", "financial", "economy"]
NEWS_LOOKBACK_DAYS = 7 # Number of days to look back for news
# Social media settings
TWITTER_BEARER_TOKEN = "your_twitter_bearer_token"
REDDIT_CLIENT_ID = "your_reddit_client_id"
REDDIT_CLIENT_SECRET = "your_reddit_client_secret"
SOCIAL_MEDIA_LOOKBACK_DAYS = 3 # Number of days to look back for social media
# Sentiment analysis settings
SENTIMENT_MODEL = "vader" # 'vader', 'finbert', or 'custom'
FINBERT_MODEL_PATH = "yiyanghkust/finbert-tone" # HuggingFace model path
# Alternative data features
NEWS_FEATURES = [
"sentiment_score",
"mention_count",
"positive_score",
"negative_score",
]
SOCIAL_FEATURES = [
"twitter_sentiment",
"reddit_sentiment",
"twitter_volume",
"reddit_volume",
]
ALTERNATIVE_DATA_WEIGHT = 0.3 # Weight for alternative data in final prediction
# Database settings for alternative data
ALTERNATIVE_DATA_DB = os.path.join(DATA_DIR, "alternative_data.db")
config = Config()
+117
View File
@@ -0,0 +1,117 @@
import logging
import matplotlib.pyplot as plt
import pandas as pd
from config import config
from src.data.pipeline import StockDataPipeline
from src.evaluation.metrics import calculate_performance_metrics, compare_to_benchmark
from src.models.backtester import GNNBacktester
from src.models.gnn_model import CorporateActionAwareGNN
from src.models.trainer import GNNTrainer
from src.utils.visualization import plot_performance, plot_trade_log
# Configure logging
logging.basicConfig(
level=logging.INFO,
format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
handlers=[logging.FileHandler("stock_gnn.log"), logging.StreamHandler()],
)
logger = logging.getLogger(__name__)
def main():
# Initialize data pipeline
logger.info("Initializing data pipeline")
pipeline = StockDataPipeline()
# Update all data
logger.info("Updating all data sources")
pipeline.update_all_data()
# Create datasets
logger.info("Creating training and validation datasets")
train_dataset = pipeline.create_training_dataset()
val_dataset = pipeline.create_validation_dataset()
logger.info(f"Training dataset size: {len(train_dataset)}")
logger.info(f"Validation dataset size: {len(val_dataset)}")
# Initialize model
logger.info("Initializing GNN model")
# Get number of features from first data point
num_features = train_dataset[0].x.shape[1]
model = CorporateActionAwareGNN(num_features)
# Train model
logger.info("Training model")
trainer = GNNTrainer(model)
train_losses, val_losses = trainer.train(train_dataset, val_dataset)
# Plot training curves
plt.figure(figsize=(10, 5))
plt.plot(train_losses, label="Training Loss")
plt.plot(val_losses, label="Validation Loss")
plt.title("Training and Validation Loss")
plt.xlabel("Epoch")
plt.ylabel("Loss")
plt.legend()
plt.savefig("training_curves.png")
plt.close()
# Load best model
trainer.load_model()
# Run backtest on validation set
logger.info("Running backtest on validation set")
backtester = GNNBacktester(model, pipeline.price_data)
portfolio_values, trade_log = backtester.run_backtest(val_dataset)
# Get benchmark data
benchmark_data = pipeline.price_data[config.BENCHMARK_TICKER]
benchmark_values = benchmark_data.loc[portfolio_values.index]["Adj Close"]
# Calculate performance metrics
logger.info("Calculating performance metrics")
portfolio_returns = portfolio_values.pct_change().dropna()
benchmark_returns = benchmark_values.pct_change().dropna()
metrics = calculate_performance_metrics(portfolio_returns, benchmark_returns)
comparison = compare_to_benchmark(portfolio_values, benchmark_values)
# Print metrics
logger.info("\nPerformance Metrics:")
for metric, value in metrics.items():
if isinstance(value, float):
logger.info(f"{metric.replace('_', ' ').title()}: {value:.4f}")
else:
logger.info(f"{metric.replace('_', ' ').title()}: {value}")
logger.info("\nComparison to Benchmark:")
for metric, value in comparison.items():
if isinstance(value, float):
logger.info(f"{metric.replace('_', ' ').title()}: {value:.4f}")
else:
logger.info(f"{metric.replace('_', ' ').title()}: {value}")
# Plot performance
plot_performance(portfolio_values, benchmark_values, "portfolio_performance.png")
# Plot trade log
plot_trade_log(trade_log, "trade_log.png")
# Save results
results_df = pd.DataFrame(
{
"date": portfolio_values.index,
"portfolio_value": portfolio_values.values,
"benchmark_value": benchmark_values.values,
}
)
results_df.to_csv("backtest_results.csv", index=False)
logger.info("Backtest completed. Results saved to backtest_results.csv")
if __name__ == "__main__":
main()
+8
View File
@@ -0,0 +1,8 @@
numpy>=1.21.0
pandas>=1.3.0
torch>=1.9.0
torch-geometric>=2.0.0
yfinance>=0.1.63
tqdm>=4.62.0
matplotlib>=3.4.0
scikit-learn>=0.24.0
Binary file not shown.
+611
View File
@@ -0,0 +1,611 @@
import logging
import os
import pickle
from datetime import datetime, timedelta
from typing import Dict, List
import numpy as np
import pandas as pd
import yfinance as yf
from tqdm import tqdm
from config import config
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)
class StockDataPipeline:
def __init__(self):
self.price_data = {}
self.corporate_actions = {}
self.sector_data = {}
self.index_composition = {}
self.delisted_tickers = set()
# Load existing data if available
self._load_data()
# Load delisted tickers
self._load_delisted_tickers()
def _load_data(self):
"""Load existing data from disk"""
try:
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")
logger.info("Loaded existing data from disk")
except FileNotFoundError:
logger.info("No existing data found. Starting with empty databases.")
def _save_data(self):
"""Save data to disk"""
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")
def _load_pickle(self, filename: str):
"""Load data from pickle file"""
filepath = os.path.join(config.PROCESSED_DATA_DIR, filename)
with open(filepath, "rb") as f:
return pickle.load(f)
def _save_pickle(self, data, filename: str):
"""Save data to pickle file"""
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):
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")
else:
logger.warning(
"Delisted tickers file not found. Only using active tickers."
)
def update_all_data(self):
"""Update all data sources"""
# 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()
# Save updated data
self._save_data()
def update_price_data(self, tickers: List[str]):
"""Update price data for given tickers"""
logger.info(f"Updating price data for {len(tickers)} tickers")
for ticker in tqdm(tickers, desc="Updating price data"):
try:
# Determine start date
start_date = config.START_DATE
if (
ticker in self.price_data
and self.price_data[ticker] is not None
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 new_data is not None and not new_data.empty:
# If we already have data for this ticker, append the new data
if (
ticker in self.price_data
and self.price_data[ticker] is not None
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()
except Exception as e:
logger.warning(f"Error updating price data for {ticker}: {str(e)}")
def update_corporate_actions(self, tickers: List[str]):
"""Update corporate actions for given tickers"""
logger.info(f"Updating corporate actions for {len(tickers)} tickers")
for ticker in tqdm(tickers, desc="Updating corporate actions"):
try:
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 = pd.Timestamp(date).strftime("%Y-%m-%d") # type: ignore
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 = pd.Timestamp(date).strftime("%Y-%m-%d") # type: ignore
self.corporate_actions[ticker]["dividends"][date_str] = float(
amount
)
# Note: Yahoo Finance doesn't provide merger/spin-off data directly
# You would need to supplement with other data sources
except Exception as e:
logger.warning(
f"Error updating corporate actions for {ticker}: {str(e)}"
)
def update_sector_data(self, tickers: List[str]):
"""Update sector data for given tickers"""
logger.info(f"Updating sector data for {len(tickers)} tickers")
for ticker in tqdm(tickers, desc="Updating sector data"):
try:
stock = yf.Ticker(ticker)
info = stock.info
if "sector" in info:
self.sector_data[ticker] = info["sector"]
except Exception as e:
logger.warning(f"Error updating sector data for {ticker}: {str(e)}")
def update_index_composition(self):
"""Update index composition for the main index"""
logger.info(f"Updating index composition for {config.INDEX_TICKER}")
try:
# Get current constituents
index = yf.Ticker(config.INDEX_TICKER)
constituents = getattr(
index, "get_index_major_holders", lambda: pd.DataFrame()
)()
if constituents is not None and not constituents.empty:
current_date = datetime.now().strftime("%Y-%m-%d")
self.index_composition[config.INDEX_TICKER] = {
current_date: constituents["Symbol"].tolist()
}
# Note: Yahoo Finance doesn't provide historical index composition
# You would need to supplement with other data sources for historical data
except Exception as e:
logger.warning(f"Error updating index composition: {str(e)}")
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,
"in_index": False,
"index": None,
"upcoming_actions": [],
}
# Get price data
if (
ticker in self.price_data
and self.price_data[ticker] is not None
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]["Adj Close"]
# Get sector data
if ticker in self.sector_data:
result["sector"] = self.sector_data[ticker]
# 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 create_training_dataset(self) -> List:
"""
Create a training dataset with proper handling of survivorship bias and corporate actions
Returns:
List of PyG Data objects
"""
import torch
from torch_geometric.data import Data
logger.info("Creating training dataset")
dates = pd.date_range(config.START_DATE, config.TRAIN_END_DATE)
dataset = []
for date in tqdm(dates, desc="Creating dataset"):
# Get current universe of stocks (those that existed at this date)
current_tickers = []
for ticker in config.INITIAL_TICKERS + list(self.delisted_tickers):
if (
ticker in self.price_data
and self.price_data[ticker] is not None
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
node_features = []
corporate_action_flags = []
for ticker in current_tickers:
# Get price data for lookback period
lookback_start = date - timedelta(days=config.LOOKBACK_WINDOW)
price_data = self.price_data[ticker].loc[lookback_start:date]
if len(price_data) < 5: # Need at least 5 days for meaningful features
# Use default values
features = [
0,
0,
0,
0,
0,
] # [return, volatility, momentum, volume, price]
else:
returns = price_data["Adj Close"].pct_change().dropna()
features = [
returns.iloc[-1], # Last return
returns.std(), # Volatility
returns.mean(), # Momentum
np.log(price_data["Volume"].iloc[-1] + 1), # Log volume
price_data["Adj Close"].iloc[-1], # Last price
]
node_features.append(features)
# Get corporate action flag
pit_data = self.get_point_in_time_data(ticker, date)
flag = 0
if pit_data["upcoming_actions"]:
# Use the type of the soonest upcoming action
soonest = min(
pit_data["upcoming_actions"], key=lambda x: x["days_until"]
)
if soonest["type"] == "split":
flag = 1
elif soonest["type"] == "dividend":
flag = 2
corporate_action_flags.append(flag)
# Convert to tensors
x = torch.tensor(node_features, dtype=torch.float)
# Add corporate action flags as additional features
corporate_action_tensor = torch.tensor(
corporate_action_flags, dtype=torch.float
).unsqueeze(1)
x = torch.cat([x, corporate_action_tensor], dim=1)
# Create edges based on sector relationships
edge_index = []
edge_weight = []
for i, ticker1 in enumerate(current_tickers):
for j, ticker2 in enumerate(current_tickers):
if i < j:
# Get sector data
pit1 = self.get_point_in_time_data(ticker1, date)
pit2 = self.get_point_in_time_data(ticker2, date)
# Create edge if same sector
if (
pit1["sector"]
and pit2["sector"]
and pit1["sector"] == pit2["sector"]
):
# Calculate correlation as edge weight
lookback_start = date - timedelta(
days=config.LOOKBACK_WINDOW
)
returns1 = (
self.price_data[ticker1]
.loc[lookback_start:date]["Adj Close"]
.pct_change()
)
returns2 = (
self.price_data[ticker2]
.loc[lookback_start:date]["Adj Close"]
.pct_change()
)
if len(returns1) > 5 and len(returns2) > 5:
corr = returns1.corr(returns2)
if not np.isnan(corr):
edge_index.append([i, j])
edge_weight.append(corr)
# 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.float).unsqueeze(1)
if edge_weight
else torch.empty((0, 1), dtype=torch.float)
)
# Create target (next day's return)
y = []
for ticker in current_tickers:
next_date = date + timedelta(days=1)
if (
ticker in self.price_data
and self.price_data[ticker] is not None
and next_date in self.price_data[ticker].index
):
ret = (
self.price_data[ticker].loc[next_date]["Adj Close"]
/ self.price_data[ticker].loc[date]["Adj Close"]
- 1
)
y.append(ret)
else:
y.append(0) # Default value
y = torch.tensor(y, dtype=torch.float).unsqueeze(1)
# Create Data object
data = Data(x=x, edge_index=edge_index, edge_attr=edge_weight, y=y)
data.date = date
data.tickers = current_tickers
dataset.append(data)
return dataset
def create_validation_dataset(self) -> List:
"""Create validation dataset (similar to training dataset but for validation period)"""
import torch
from torch_geometric.data import Data
logger.info("Creating validation dataset")
dates = pd.date_range(config.TRAIN_END_DATE, config.VAL_END_DATE)
dataset = []
for date in tqdm(dates, desc="Creating validation dataset"):
# Get current universe of stocks
current_tickers = []
for ticker in config.INITIAL_TICKERS + list(self.delisted_tickers):
if (
ticker in self.price_data
and self.price_data[ticker] is not None
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
node_features = []
corporate_action_flags = []
for ticker in current_tickers:
# Get price data for lookback period
lookback_start = date - timedelta(days=config.LOOKBACK_WINDOW)
price_data = self.price_data[ticker].loc[lookback_start:date]
if len(price_data) < 5:
features = [0, 0, 0, 0, 0]
else:
returns = price_data["Adj Close"].pct_change().dropna()
features = [
returns.iloc[-1],
returns.std(),
returns.mean(),
np.log(price_data["Volume"].iloc[-1] + 1),
price_data["Adj Close"].iloc[-1],
]
node_features.append(features)
# Get corporate action flag
pit_data = self.get_point_in_time_data(ticker, date)
flag = 0
if pit_data["upcoming_actions"]:
soonest = min(
pit_data["upcoming_actions"], key=lambda x: x["days_until"]
)
if soonest["type"] == "split":
flag = 1
elif soonest["type"] == "dividend":
flag = 2
corporate_action_flags.append(flag)
# Convert to tensors
x = torch.tensor(node_features, dtype=torch.float)
corporate_action_tensor = torch.tensor(
corporate_action_flags, dtype=torch.float
).unsqueeze(1)
x = torch.cat([x, corporate_action_tensor], dim=1)
# Create edges based on sector relationships
edge_index = []
edge_weight = []
for i, ticker1 in enumerate(current_tickers):
for j, ticker2 in enumerate(current_tickers):
if i < j:
pit1 = self.get_point_in_time_data(ticker1, date)
pit2 = self.get_point_in_time_data(ticker2, date)
if (
pit1["sector"]
and pit2["sector"]
and pit1["sector"] == pit2["sector"]
):
lookback_start = date - timedelta(
days=config.LOOKBACK_WINDOW
)
returns1 = (
self.price_data[ticker1]
.loc[lookback_start:date]["Adj Close"]
.pct_change()
)
returns2 = (
self.price_data[ticker2]
.loc[lookback_start:date]["Adj Close"]
.pct_change()
)
if len(returns1) > 5 and len(returns2) > 5:
corr = returns1.corr(returns2)
if not np.isnan(corr):
edge_index.append([i, j])
edge_weight.append(corr)
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.float).unsqueeze(1)
if edge_weight
else torch.empty((0, 1), dtype=torch.float)
)
# Create target
y = []
for ticker in current_tickers:
next_date = date + timedelta(days=1)
if (
ticker in self.price_data
and self.price_data[ticker] is not None
and next_date in self.price_data[ticker].index
):
ret = (
self.price_data[ticker].loc[next_date]["Adj Close"]
/ self.price_data[ticker].loc[date]["Adj Close"]
- 1
)
y.append(ret)
else:
y.append(0)
y = torch.tensor(y, dtype=torch.float).unsqueeze(1)
# Create Data object
data = Data(x=x, edge_index=edge_index, edge_attr=edge_weight, y=y)
data.date = date
data.tickers = current_tickers
dataset.append(data)
return dataset
+89
View File
@@ -0,0 +1,89 @@
import logging
from datetime import datetime
from typing import Dict, List, Set
logger = logging.getLogger(__name__)
class SurvivorshipBiasHandler:
def __init__(self, delisted_tickers: Set[str], index_composition: Dict):
self.delisted_tickers = delisted_tickers
self.index_composition = index_composition
def get_investable_universe(self, date: datetime) -> List[str]:
"""
Get the investable universe at a specific date
Parameters:
date: Date for which to get the universe
Returns:
List of tickers in the investable universe
"""
# Get current index members
current_members = set()
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:
current_members.update(composition[comp_date])
break
# Add delisted stocks that were in the index but haven't been delisted yet
investable_universe = list(current_members)
# Add any delisted stocks that were in the index but are still trading
for ticker in self.delisted_tickers:
if ticker in current_members:
investable_universe.append(ticker)
return investable_universe
def filter_available_stocks(
self, tickers: List[str], price_data: Dict, date: datetime
) -> List[str]:
"""
Filter stocks to only those available at a specific date
Parameters:
tickers: List of tickers to filter
price_data: Dictionary of price data
date: Date to check availability
Returns:
List of available tickers
"""
available_tickers = []
for ticker in tickers:
if ticker in price_data and not price_data[ticker].empty:
if (
date >= price_data[ticker].index[0]
and date <= price_data[ticker].index[-1]
):
available_tickers.append(ticker)
return available_tickers
def get_point_in_time_index_membership(self, ticker: str, date: datetime) -> bool:
"""
Check if a stock was in the index at a specific date
Parameters:
ticker: Stock ticker
date: Date to check
Returns:
Boolean indicating if the stock was in the index
"""
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]:
return True
break
return False
Binary file not shown.
+124
View File
@@ -0,0 +1,124 @@
import logging
from typing import Dict
import numpy as np
import pandas as pd
logger = logging.getLogger(__name__)
def calculate_performance_metrics(
portfolio_returns: pd.Series, benchmark_returns: pd.Series
) -> Dict:
"""
Calculate performance metrics for a trading strategy
Parameters:
portfolio_returns: Series of portfolio returns
benchmark_returns: Series of benchmark returns
Returns:
Dictionary of performance metrics
"""
metrics = {}
# Total return
metrics["total_return"] = (portfolio_returns + 1).prod() - 1
# Annualized return
years = len(portfolio_returns) / 252 # Trading days in a year
metrics["annualized_return"] = (1 + metrics["total_return"]) ** (1 / years) - 1
# Volatility
metrics["volatility"] = portfolio_returns.std() * np.sqrt(252)
# Sharpe ratio (assuming risk-free rate = 0)
metrics["sharpe_ratio"] = metrics["annualized_return"] / metrics["volatility"]
# Sortino ratio
downside_returns = portfolio_returns[portfolio_returns < 0]
downside_volatility = downside_returns.std() * np.sqrt(252)
metrics["sortino_ratio"] = (
metrics["annualized_return"] / downside_volatility
if downside_volatility > 0
else np.inf
)
# Maximum drawdown
cumulative_returns = (1 + portfolio_returns).cumprod()
running_max = cumulative_returns.cummax()
drawdown = (cumulative_returns - running_max) / running_max
metrics["max_drawdown"] = drawdown.min()
# Calmar ratio
metrics["calmar_ratio"] = (
metrics["annualized_return"] / abs(metrics["max_drawdown"])
if metrics["max_drawdown"] < 0
else np.inf
)
# Alpha and Beta
if len(benchmark_returns) > 1:
cov = portfolio_returns.cov(benchmark_returns)
var = benchmark_returns.var()
metrics["beta"] = cov / var
# Alpha = portfolio return - (risk-free rate + beta * (benchmark return - risk-free rate))
# Assuming risk-free rate = 0
metrics["alpha"] = metrics["annualized_return"] - metrics["beta"] * (
(benchmark_returns + 1).prod() ** (252 / len(benchmark_returns)) - 1
)
# Win rate
metrics["win_rate"] = (portfolio_returns > 0).mean()
# Profit factor
gains = portfolio_returns[portfolio_returns > 0].sum()
losses = -portfolio_returns[portfolio_returns < 0].sum()
metrics["profit_factor"] = gains / losses if losses > 0 else np.inf
# Return over maximum drawdown
metrics["return_over_max_drawdown"] = (
metrics["annualized_return"] / abs(metrics["max_drawdown"])
if metrics["max_drawdown"] < 0
else np.inf
)
return metrics
def compare_to_benchmark(
portfolio_values: pd.Series, benchmark_values: pd.Series
) -> Dict:
"""
Compare portfolio performance to benchmark
Parameters:
portfolio_values: Series of portfolio values
benchmark_values: Series of benchmark values
Returns:
Dictionary of comparison metrics
"""
# Calculate returns
portfolio_returns = portfolio_values.pct_change().dropna()
benchmark_returns = benchmark_values.pct_change().dropna()
# Align the returns
common_index = portfolio_returns.index.intersection(benchmark_returns.index)
portfolio_returns = portfolio_returns.loc[common_index]
benchmark_returns = benchmark_returns.loc[common_index]
# Calculate metrics
metrics = calculate_performance_metrics(portfolio_returns, benchmark_returns)
# Additional comparison metrics
metrics["benchmark_total_return"] = (benchmark_returns + 1).prod() - 1
metrics["benchmark_annualized_return"] = (
1 + metrics["benchmark_total_return"]
) ** (252 / len(benchmark_returns)) - 1
metrics["excess_return"] = (
metrics["annualized_return"] - metrics["benchmark_annualized_return"]
)
return metrics
Binary file not shown.
Binary file not shown.
Binary file not shown.
+177
View File
@@ -0,0 +1,177 @@
import logging
from datetime import datetime
from typing import Dict, List, Tuple
import pandas as pd
import torch
from config import config
from src.models.gnn_model import CorporateActionAwareGNN
logger = logging.getLogger(__name__)
class GNNBacktester:
def __init__(
self,
model: CorporateActionAwareGNN,
price_data: Dict,
initial_capital: float = config.INITIAL_CAPITAL,
):
self.model = model
self.price_data = price_data
self.initial_capital = initial_capital
self.portfolio_value = initial_capital
self.portfolio = {} # {ticker: shares}
self.trade_log = []
self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
def run_backtest(self, dataset: List) -> Tuple[pd.Series, pd.DataFrame]:
"""
Run backtest on the given dataset
Parameters:
dataset: List of PyG Data objects
Returns:
Tuple of (portfolio_values, trade_log)
"""
portfolio_values = []
dates = []
for data in dataset:
date = data.date
current_tickers = data.tickers
# Get current prices
current_prices = {}
for ticker in current_tickers:
if ticker in self.price_data and date in self.price_data[ticker].index:
current_prices[ticker] = self.price_data[ticker].loc[date][
"Adj Close"
]
# Calculate current portfolio value
current_value = sum(
shares * current_prices[ticker]
for ticker, shares in self.portfolio.items()
if ticker in current_prices
)
cash = self.portfolio_value - current_value
current_portfolio_value = current_value + cash
# Store portfolio value
portfolio_values.append(current_portfolio_value)
dates.append(date)
# Get predictions from GNN
with torch.no_grad():
data = data.to(self.device)
predictions = self.model(data).squeeze().cpu().numpy()
# Create trading signals
signals = {}
for i, ticker in enumerate(current_tickers):
if ticker in current_prices:
# Get corporate action flag (last feature)
corporate_action_flag = data.x[i, -1].item()
# Buy if predicted return > threshold and no upcoming corporate action
if predictions[i] > 0.005 and corporate_action_flag == 0:
signals[ticker] = "buy"
# Sell if predicted return < threshold or upcoming corporate action
elif predictions[i] < -0.005 or corporate_action_flag > 0:
signals[ticker] = "sell"
# Execute trades
for ticker, signal in signals.items():
if signal == "buy" and cash > 0:
# Buy with 10% of available cash
price = current_prices[ticker]
shares_to_buy = int(
(cash * 0.1) / (price * (1 + config.TRANSACTION_COST))
)
if shares_to_buy > 0:
cost = shares_to_buy * price * (1 + config.TRANSACTION_COST)
self.portfolio[ticker] = (
self.portfolio.get(ticker, 0) + shares_to_buy
)
cash -= cost
self.trade_log.append(
(date, ticker, "buy", shares_to_buy, price)
)
logger.debug(
f"Bought {shares_to_buy} shares of {ticker} at {price:.2f}"
)
elif signal == "sell" and ticker in self.portfolio:
# Sell all shares
shares = self.portfolio.pop(ticker)
proceeds = (
shares * current_prices[ticker] * (1 - config.TRANSACTION_COST)
)
cash += proceeds
self.trade_log.append(
(date, ticker, "sell", shares, current_prices[ticker])
)
logger.debug(
f"Sold {shares} shares of {ticker} at {current_prices[ticker]:.2f}"
)
# Update portfolio value
new_value = sum(
shares * current_prices[ticker]
for ticker, shares in self.portfolio.items()
if ticker in current_prices
)
self.portfolio_value = new_value + cash
# Create portfolio value series
portfolio_series = pd.Series(portfolio_values, index=dates)
# Create trade log DataFrame
if self.trade_log:
trade_log_df = pd.DataFrame(self.trade_log)
trade_log_df.columns = ["date", "ticker", "action", "shares", "price"]
else:
trade_log_df = pd.DataFrame()
trade_log_df.columns = ["date", "ticker", "action", "shares", "price"]
return portfolio_series, trade_log_df
def get_portfolio_composition(self, date: datetime) -> Dict[str, float]:
"""
Get portfolio composition at a specific date
Parameters:
date: Date to get composition for
Returns:
Dictionary of {ticker: weight} where weight is the percentage of portfolio
"""
# Get current prices
current_prices = {}
for ticker in self.portfolio:
if ticker in self.price_data and date in self.price_data[ticker].index:
current_prices[ticker] = self.price_data[ticker].loc[date]["Adj Close"]
# Calculate current value
current_value = sum(
shares * current_prices[ticker]
for ticker, shares in self.portfolio.items()
if ticker in current_prices
)
cash = self.portfolio_value - current_value
# Calculate weights
composition = {}
for ticker, shares in self.portfolio.items():
if ticker in current_prices:
composition[ticker] = (
shares * current_prices[ticker]
) / self.portfolio_value
# Add cash
composition["cash"] = cash / self.portfolio_value
return composition
+45
View File
@@ -0,0 +1,45 @@
import torch.nn as nn
import torch.nn.functional as F
from torch_geometric.nn import GATConv, LayerNorm
class CorporateActionAwareGNN(nn.Module):
def __init__(
self,
num_features: int,
hidden_channels: int = 64,
num_heads: int = 8,
dropout: float = 0.6,
):
super().__init__()
self.conv1 = GATConv(
num_features,
hidden_channels,
heads=num_heads,
dropout=dropout,
concat=True,
)
self.norm1 = LayerNorm(hidden_channels * num_heads)
self.conv2 = GATConv(
hidden_channels * num_heads,
hidden_channels,
heads=num_heads,
dropout=dropout,
concat=True,
)
self.norm2 = LayerNorm(hidden_channels * num_heads)
self.fc = nn.Linear(hidden_channels * num_heads, 1)
self.dropout = nn.Dropout(dropout)
def forward(self, data):
x, edge_index = data.x, data.edge_index
x = self.conv1(x, edge_index)
x = self.norm1(x)
x = F.elu(x)
x = self.dropout(x)
x = self.conv2(x, edge_index)
x = self.norm2(x)
x = F.elu(x)
x = self.dropout(x)
x = self.fc(x)
return x
+178
View File
@@ -0,0 +1,178 @@
import logging
import os
from datetime import datetime
from typing import Dict, List, Set, Tuple
import torch
import torch.nn as nn
from config import config
from src.models.gnn_model import CorporateActionAwareGNN
logger = logging.getLogger(__name__)
class GNNTrainer:
def __init__(self, model: CorporateActionAwareGNN):
self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
self.model = model.to(self.device)
self.optimizer = torch.optim.Adam(
self.model.parameters(), lr=config.LEARNING_RATE
)
self.criterion = nn.MSELoss()
self.model_dir = config.MODEL_DIR
self.model_name = config.MODEL_NAME
def train(
self, train_dataset: List, val_dataset: List
) -> Tuple[List[float], List[float]]:
train_losses = []
val_losses = []
for epoch in range(config.EPOCHS):
self.model.train()
total_loss = 0.0
for data in train_dataset:
data = data.to(self.device)
self.optimizer.zero_grad()
out = self.model(data).squeeze()
if hasattr(data, "y") and data.y is not None:
target = data.y.to(self.device)
if out.dim() == 0:
out = out.unsqueeze(0)
if target.dim() == 0:
target = target.unsqueeze(0)
loss = self.criterion(out, target)
loss.backward()
self.optimizer.step()
total_loss += loss.item()
avg_train_loss = total_loss / len(train_dataset) if train_dataset else 0.0
train_losses.append(avg_train_loss)
self.model.eval()
total_val_loss = 0.0
with torch.no_grad():
for data in val_dataset:
data = data.to(self.device)
out = self.model(data).squeeze()
if hasattr(data, "y") and data.y is not None:
target = data.y.to(self.device)
if out.dim() == 0:
out = out.unsqueeze(0)
if target.dim() == 0:
target = target.unsqueeze(0)
loss = self.criterion(out, target)
total_val_loss += loss.item()
avg_val_loss = total_val_loss / len(val_dataset) if val_dataset else 0.0
val_losses.append(avg_val_loss)
logger.info(
f"Epoch {epoch + 1}/{config.EPOCHS}, Train Loss: {avg_train_loss:.4f}, Val Loss: {avg_val_loss:.4f}"
)
self.save_model()
return train_losses, val_losses
def save_model(self):
os.makedirs(self.model_dir, exist_ok=True)
path = os.path.join(self.model_dir, f"{self.model_name}.pt")
torch.save(self.model.state_dict(), path)
logger.info(f"Model saved to {path}")
def load_model(self):
path = os.path.join(self.model_dir, f"{self.model_name}.pt")
if os.path.exists(path):
self.model.load_state_dict(torch.load(path, map_location=self.device))
logger.info(f"Model loaded from {path}")
else:
logger.warning(f"No model found at {path}")
class SurvivorshipBiasHandler:
def __init__(self, delisted_tickers: Set[str], index_composition: Dict):
self.delisted_tickers = delisted_tickers
self.index_composition = index_composition
def get_investable_universe(self, date: datetime) -> List[str]:
"""
Get the investable universe at a specific date
Parameters:
date: Date for which to get the universe
Returns:
List of tickers in the investable universe
"""
# Get current index members
current_members = set()
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:
current_members.update(composition[comp_date])
break
# Add delisted stocks that were in the index but haven't been delisted yet
investable_universe = list(current_members)
# Add any delisted stocks that were in the index but are still trading
for ticker in self.delisted_tickers:
if ticker in current_members:
investable_universe.append(ticker)
return investable_universe
def filter_available_stocks(
self, tickers: List[str], price_data: Dict, date: datetime
) -> List[str]:
"""
Filter stocks to only those available at a specific date
Parameters:
tickers: List of tickers to filter
price_data: Dictionary of price data
date: Date to check availability
Returns:
List of available tickers
"""
available_tickers = []
for ticker in tickers:
if (
ticker in price_data
and price_data[ticker] is not None
and not price_data[ticker].empty
):
if (
date >= price_data[ticker].index[0]
and date <= price_data[ticker].index[-1]
):
available_tickers.append(ticker)
return available_tickers
def get_point_in_time_index_membership(self, ticker: str, date: datetime) -> bool:
"""
Check if a stock was in the index at a specific date
Parameters:
ticker: Stock ticker
date: Date to check
Returns:
Boolean indicating if the stock was in the index
"""
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]:
return True
break
return False
Binary file not shown.
+137
View File
@@ -0,0 +1,137 @@
from typing import Dict, List, Optional
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
def plot_performance(
portfolio_values: pd.Series,
benchmark_values: pd.Series,
filename: Optional[str] = None,
):
"""
Plot portfolio performance vs benchmark
Parameters:
portfolio_values: Series of portfolio values
benchmark_values: Series of benchmark values
filename: Filename to save plot (optional)
"""
plt.figure(figsize=(12, 6))
# Normalize to start at 1
portfolio_normalized = portfolio_values / portfolio_values.iloc[0]
benchmark_normalized = benchmark_values / benchmark_values.iloc[0]
plt.plot(portfolio_normalized, label="Portfolio")
plt.plot(benchmark_normalized, label="Benchmark")
plt.title("Portfolio Performance vs Benchmark")
plt.xlabel("Date")
plt.ylabel("Normalized Value")
plt.legend()
plt.grid(True)
if filename:
plt.savefig(filename)
plt.close()
else:
plt.show()
def plot_trade_log(trade_log: pd.DataFrame, filename: Optional[str] = None):
"""
Plot trade log
Parameters:
trade_log: DataFrame of trades
filename: Filename to save plot (optional)
"""
if trade_log.empty:
return
plt.figure(figsize=(12, 6))
# Plot buy and sell points
buys = trade_log[trade_log["action"] == "buy"]
sells = trade_log[trade_log["action"] == "sell"]
plt.scatter(
buys["date"], buys["price"], color="g", label="Buy", marker="^", alpha=0.7
)
plt.scatter(
sells["date"], sells["price"], color="r", label="Sell", marker="v", alpha=0.7
)
plt.title("Trade Log")
plt.xlabel("Date")
plt.ylabel("Price")
plt.legend()
plt.grid(True)
if filename:
plt.savefig(filename)
plt.close()
else:
plt.show()
def plot_portfolio_composition(
composition: Dict[str, float], filename: Optional[str] = None
):
"""
Plot portfolio composition
Parameters:
composition: Dictionary of {ticker: weight}
filename: Filename to save plot (optional)
"""
if not composition:
return
plt.figure(figsize=(10, 6))
# Sort by weight
sorted_composition = sorted(composition.items(), key=lambda x: x[1], reverse=True)
# Extract tickers and weights
tickers = [item[0] for item in sorted_composition]
weights = [item[1] for item in sorted_composition]
# Create pie chart
plt.pie(weights, labels=tickers, autopct="%1.1f%%", startangle=140)
plt.title("Portfolio Composition")
if filename:
plt.savefig(filename)
plt.close()
else:
plt.show()
def plot_feature_importance(
importance: np.ndarray, feature_names: List[str], filename: Optional[str] = None
):
"""
Plot feature importance
Parameters:
importance: Array of feature importance scores
feature_names: List of feature names
filename: Filename to save plot (optional)
"""
plt.figure(figsize=(10, 6))
# Sort features by importance
sorted_idx = importance.argsort()
plt.barh(range(len(sorted_idx)), importance[sorted_idx], align="center")
plt.yticks(range(len(sorted_idx)), [feature_names[i] for i in sorted_idx])
plt.title("Feature Importance")
plt.xlabel("Importance Score")
if filename:
plt.savefig(filename)
plt.close()
else:
plt.show()
View File