initial commit

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()