commit b43fbe21709e034ac26875ec84aa890b3ef2522f Author: Florian Egger Date: Tue May 26 12:43:47 2026 +0200 initial commit diff --git a/README.md b/README.md new file mode 100644 index 0000000..fa2487b --- /dev/null +++ b/README.md @@ -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 +└── ... diff --git a/__pycache__/config.cpython-314.pyc b/__pycache__/config.cpython-314.pyc new file mode 100644 index 0000000..3993270 Binary files /dev/null and b/__pycache__/config.cpython-314.pyc differ diff --git a/__pycache__/main.cpython-314.pyc b/__pycache__/main.cpython-314.pyc new file mode 100644 index 0000000..ecc6ff4 Binary files /dev/null and b/__pycache__/main.cpython-314.pyc differ diff --git a/config.py b/config.py new file mode 100644 index 0000000..d3037d0 --- /dev/null +++ b/config.py @@ -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() diff --git a/main.py b/main.py new file mode 100644 index 0000000..df04aaf --- /dev/null +++ b/main.py @@ -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() diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000..cf6a764 --- /dev/null +++ b/requirements.txt @@ -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 diff --git a/src/data/__pycache__/pipeline.cpython-314.pyc b/src/data/__pycache__/pipeline.cpython-314.pyc new file mode 100644 index 0000000..74e6545 Binary files /dev/null and b/src/data/__pycache__/pipeline.cpython-314.pyc differ diff --git a/src/data/pipeline.py b/src/data/pipeline.py new file mode 100644 index 0000000..0153df5 --- /dev/null +++ b/src/data/pipeline.py @@ -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 diff --git a/src/data/survivorship_bias.py b/src/data/survivorship_bias.py new file mode 100644 index 0000000..c8dca4a --- /dev/null +++ b/src/data/survivorship_bias.py @@ -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 diff --git a/src/evaluation/__pycache__/metrics.cpython-314.pyc b/src/evaluation/__pycache__/metrics.cpython-314.pyc new file mode 100644 index 0000000..a24b106 Binary files /dev/null and b/src/evaluation/__pycache__/metrics.cpython-314.pyc differ diff --git a/src/evaluation/metrics.py b/src/evaluation/metrics.py new file mode 100644 index 0000000..6c38449 --- /dev/null +++ b/src/evaluation/metrics.py @@ -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 diff --git a/src/models/__pycache__/backtester.cpython-314.pyc b/src/models/__pycache__/backtester.cpython-314.pyc new file mode 100644 index 0000000..59ee519 Binary files /dev/null and b/src/models/__pycache__/backtester.cpython-314.pyc differ diff --git a/src/models/__pycache__/gnn_model.cpython-314.pyc b/src/models/__pycache__/gnn_model.cpython-314.pyc new file mode 100644 index 0000000..74e4ccb Binary files /dev/null and b/src/models/__pycache__/gnn_model.cpython-314.pyc differ diff --git a/src/models/__pycache__/trainer.cpython-314.pyc b/src/models/__pycache__/trainer.cpython-314.pyc new file mode 100644 index 0000000..fb29a9d Binary files /dev/null and b/src/models/__pycache__/trainer.cpython-314.pyc differ diff --git a/src/models/backtester.py b/src/models/backtester.py new file mode 100644 index 0000000..bbd3d06 --- /dev/null +++ b/src/models/backtester.py @@ -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 diff --git a/src/models/gnn_model.py b/src/models/gnn_model.py new file mode 100644 index 0000000..59f9f7d --- /dev/null +++ b/src/models/gnn_model.py @@ -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 diff --git a/src/models/trainer.py b/src/models/trainer.py new file mode 100644 index 0000000..0ecdb7a --- /dev/null +++ b/src/models/trainer.py @@ -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 diff --git a/src/utils/__pycache__/visualization.cpython-314.pyc b/src/utils/__pycache__/visualization.cpython-314.pyc new file mode 100644 index 0000000..e560fe3 Binary files /dev/null and b/src/utils/__pycache__/visualization.cpython-314.pyc differ diff --git a/src/utils/visualization.py b/src/utils/visualization.py new file mode 100644 index 0000000..e4a0f2f --- /dev/null +++ b/src/utils/visualization.py @@ -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() diff --git a/stock_gnn.log b/stock_gnn.log new file mode 100644 index 0000000..e69de29