""" Demo Data for Portfolio Optimization This module provides sample stock data for the portfolio optimization quickstart. The data includes 20 stocks across 4 sectors with ML-predicted returns. In a real application, these predictions would come from an ML model trained on historical stock data. For this quickstart, we use hardcoded realistic values. FINANCE CONCEPTS: - predicted_return: Expected percentage gain (0.12 = 12% expected return) - sector: Industry classification for diversification - Equal weight: Each selected stock gets 100%/20 = 5% of the portfolio """ from enum import Enum from dataclasses import dataclass from .domain import StockSelection, PortfolioOptimizationPlan, PortfolioConfig class DemoData(Enum): """Available demo datasets.""" SMALL = 'SMALL' # 20 stocks - good for learning LARGE = 'LARGE' # 50 stocks - more realistic @dataclass class DemoDataConfig: """Configuration for demo data generation.""" target_position_count: int max_sector_percentage: float demo_data_configs = { DemoData.SMALL: DemoDataConfig( target_position_count=20, max_sector_percentage=0.25, ), DemoData.LARGE: DemoDataConfig( target_position_count=20, max_sector_percentage=0.25, ), } # Stock data with realistic ML predictions # Format: (ticker, name, sector, predicted_return) # # SMALL dataset: 25 stocks, need to select 20 # This is FEASIBLE because we have 5+ stocks in each of 4 sectors (5*4=20 max from limits) # Plus we have extra stocks to choose from in each sector SMALL_DATASET_STOCKS = [ # TECHNOLOGY (7 stocks) - typically higher predicted returns # Solver can pick max 5, so must choose best 5 from 7 ("AAPL", "Apple Inc.", "Technology", 0.12), ("GOOGL", "Alphabet (Google)", "Technology", 0.15), ("MSFT", "Microsoft Corp.", "Technology", 0.10), ("NVDA", "NVIDIA Corp.", "Technology", 0.18), ("META", "Meta Platforms", "Technology", 0.08), ("TSLA", "Tesla Inc.", "Technology", 0.20), ("AMD", "AMD Inc.", "Technology", 0.14), # HEALTHCARE (6 stocks) - moderate returns # Solver can pick max 5, so must choose best 5 from 6 ("JNJ", "Johnson & Johnson", "Healthcare", 0.09), ("UNH", "UnitedHealth Group", "Healthcare", 0.11), ("PFE", "Pfizer Inc.", "Healthcare", 0.07), ("ABBV", "AbbVie Inc.", "Healthcare", 0.10), ("TMO", "Thermo Fisher", "Healthcare", 0.13), ("DHR", "Danaher Corp.", "Healthcare", 0.12), # FINANCE (6 stocks) - stable returns # Solver can pick max 5, so must choose best 5 from 6 ("JPM", "JPMorgan Chase", "Finance", 0.08), ("BAC", "Bank of America", "Finance", 0.06), ("WFC", "Wells Fargo", "Finance", 0.07), ("GS", "Goldman Sachs", "Finance", 0.09), ("MS", "Morgan Stanley", "Finance", 0.08), ("C", "Citigroup", "Finance", 0.05), # ENERGY (6 stocks) - variable returns # Solver can pick max 5, so must choose best 5 from 6 ("XOM", "Exxon Mobil", "Energy", 0.04), ("CVX", "Chevron Corp.", "Energy", 0.05), ("COP", "ConocoPhillips", "Energy", 0.06), ("SLB", "Schlumberger", "Energy", 0.03), ("EOG", "EOG Resources", "Energy", 0.07), ("PXD", "Pioneer Natural", "Energy", 0.08), ] LARGE_DATASET_STOCKS = SMALL_DATASET_STOCKS + [ # Additional TECHNOLOGY (6 more -> 13 total) ("CRM", "Salesforce", "Technology", 0.11), ("ADBE", "Adobe Inc.", "Technology", 0.09), ("ORCL", "Oracle Corp.", "Technology", 0.07), ("CSCO", "Cisco Systems", "Technology", 0.06), ("IBM", "IBM Corp.", "Technology", 0.04), ("QCOM", "Qualcomm", "Technology", 0.13), # Additional HEALTHCARE (6 more -> 12 total) ("MRK", "Merck & Co.", "Healthcare", 0.08), ("LLY", "Eli Lilly", "Healthcare", 0.16), ("BMY", "Bristol-Myers", "Healthcare", 0.06), ("AMGN", "Amgen Inc.", "Healthcare", 0.09), ("GILD", "Gilead Sciences", "Healthcare", 0.05), ("ISRG", "Intuitive Surgical", "Healthcare", 0.14), # Additional FINANCE (4 more -> 10 total, no duplicates) ("AXP", "American Express", "Finance", 0.10), ("BLK", "BlackRock", "Finance", 0.11), ("SCHW", "Charles Schwab", "Finance", 0.07), ("USB", "U.S. Bancorp", "Finance", 0.04), # Additional ENERGY (2 more -> 8 total, no duplicates) ("OXY", "Occidental Petroleum", "Energy", 0.06), ("HAL", "Halliburton", "Energy", 0.05), # CONSUMER (new sector - 8 stocks) ("AMZN", "Amazon.com", "Consumer", 0.14), ("WMT", "Walmart", "Consumer", 0.06), ("HD", "Home Depot", "Consumer", 0.08), ("MCD", "McDonald's", "Consumer", 0.07), ("NKE", "Nike Inc.", "Consumer", 0.09), ("SBUX", "Starbucks", "Consumer", 0.05), ("PG", "Procter & Gamble", "Consumer", 0.04), ("KO", "Coca-Cola", "Consumer", 0.05), ] # LARGE total: 25 + 6 + 6 + 4 + 2 + 8 = 51 stocks def generate_demo_data(demo_data: DemoData) -> PortfolioOptimizationPlan: """ Generate demo data for portfolio optimization. Args: demo_data: Which demo dataset to generate (SMALL or LARGE) Returns: PortfolioOptimizationPlan with candidate stocks (all unselected initially) Example: >>> plan = generate_demo_data(DemoData.SMALL) >>> len(plan.stocks) 20 >>> plan.stocks[0].stock_id 'AAPL' """ config = demo_data_configs[demo_data] stock_data = SMALL_DATASET_STOCKS if demo_data == DemoData.SMALL else LARGE_DATASET_STOCKS stocks = [ StockSelection( stock_id=ticker, stock_name=name, sector=sector, predicted_return=predicted_return, selection=None, # To be decided by solver ) for ticker, name, sector, predicted_return in stock_data ] # Calculate max_per_sector from percentage target_count = config.target_position_count max_per_sector = max(1, int(config.max_sector_percentage * target_count)) # Create PortfolioConfig for constraints to access portfolio_config = PortfolioConfig( target_count=target_count, max_per_sector=max_per_sector, unselected_penalty=10000, ) return PortfolioOptimizationPlan( stocks=stocks, target_position_count=config.target_position_count, max_sector_percentage=config.max_sector_percentage, portfolio_config=portfolio_config, ) def get_stock_summary(plan: PortfolioOptimizationPlan) -> str: """ Generate a human-readable summary of the portfolio. Useful for debugging and understanding the solution. """ lines = [ "=" * 60, "PORTFOLIO SUMMARY", "=" * 60, ] selected = plan.get_selected_stocks() if not selected: lines.append("No stocks selected yet.") return "\n".join(lines) weight = plan.get_weight_per_stock() expected_return = plan.get_expected_return() lines.append(f"Selected: {len(selected)} stocks @ {weight*100:.1f}% each") lines.append(f"Expected Return: {expected_return*100:.2f}%") lines.append("") # Group by sector sector_stocks: dict[str, list[StockSelection]] = {} for stock in selected: if stock.sector not in sector_stocks: sector_stocks[stock.sector] = [] sector_stocks[stock.sector].append(stock) lines.append("BY SECTOR:") for sector, stocks in sorted(sector_stocks.items()): sector_weight = len(stocks) * weight * 100 lines.append(f" {sector}: {len(stocks)} stocks = {sector_weight:.1f}%") for stock in sorted(stocks, key=lambda s: -s.predicted_return): lines.append(f" - {stock.stock_id}: {stock.stock_name} ({stock.predicted_return*100:.1f}% pred)") lines.append("") lines.append(f"Score: {plan.score}") lines.append("=" * 60) return "\n".join(lines)