initial commit
This commit is contained in:
@@ -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.
@@ -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()
|
||||
@@ -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()
|
||||
@@ -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.
@@ -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
|
||||
@@ -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.
@@ -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.
@@ -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
|
||||
@@ -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
|
||||
@@ -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.
@@ -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()
|
||||
Reference in New Issue
Block a user