blackopsrepl/portfolio-optimization-python
1
1"""
2Demo Data for Portfolio Optimization
3
4This module provides sample stock data for the portfolio optimization quickstart.
5The data includes 20 stocks across 4 sectors with ML-predicted returns.
6
7In a real application, these predictions would come from an ML model trained
8on historical stock data. For this quickstart, we use hardcoded realistic values.
9
10FINANCE CONCEPTS:
11- predicted_return: Expected percentage gain (0.12 = 12% expected return)
12- sector: Industry classification for diversification
13- Equal weight: Each selected stock gets 100%/20 = 5% of the portfolio
14"""
15from enum import Enum
16from dataclasses import dataclass
17
18from .domain import StockSelection, PortfolioOptimizationPlan, PortfolioConfig
19
20
21class DemoData(Enum):
22 """Available demo datasets."""
23 SMALL = 'SMALL' # 20 stocks - good for learning
24 LARGE = 'LARGE' # 50 stocks - more realistic
25
26
27@dataclass
28class DemoDataConfig:
29 """Configuration for demo data generation."""
30 target_position_count: int
31 max_sector_percentage: float
32
33
34demo_data_configs = {
35 DemoData.SMALL: DemoDataConfig(
36 target_position_count=20,
37 max_sector_percentage=0.25,
38 ),
39 DemoData.LARGE: DemoDataConfig(
40 target_position_count=20,
41 max_sector_percentage=0.25,
42 ),
43}
44
45
46# Stock data with realistic ML predictions
47# Format: (ticker, name, sector, predicted_return)
48#
49# SMALL dataset: 25 stocks, need to select 20
50# This is FEASIBLE because we have 5+ stocks in each of 4 sectors (5*4=20 max from limits)
51# Plus we have extra stocks to choose from in each sector
52SMALL_DATASET_STOCKS = [
53 # TECHNOLOGY (7 stocks) - typically higher predicted returns
54 # Solver can pick max 5, so must choose best 5 from 7
55 ("AAPL", "Apple Inc.", "Technology", 0.12),
56 ("GOOGL", "Alphabet (Google)", "Technology", 0.15),
57 ("MSFT", "Microsoft Corp.", "Technology", 0.10),
58 ("NVDA", "NVIDIA Corp.", "Technology", 0.18),
59 ("META", "Meta Platforms", "Technology", 0.08),
60 ("TSLA", "Tesla Inc.", "Technology", 0.20),
61 ("AMD", "AMD Inc.", "Technology", 0.14),
62
63 # HEALTHCARE (6 stocks) - moderate returns
64 # Solver can pick max 5, so must choose best 5 from 6
65 ("JNJ", "Johnson & Johnson", "Healthcare", 0.09),
66 ("UNH", "UnitedHealth Group", "Healthcare", 0.11),
67 ("PFE", "Pfizer Inc.", "Healthcare", 0.07),
68 ("ABBV", "AbbVie Inc.", "Healthcare", 0.10),
69 ("TMO", "Thermo Fisher", "Healthcare", 0.13),
70 ("DHR", "Danaher Corp.", "Healthcare", 0.12),
71
72 # FINANCE (6 stocks) - stable returns
73 # Solver can pick max 5, so must choose best 5 from 6
74 ("JPM", "JPMorgan Chase", "Finance", 0.08),
75 ("BAC", "Bank of America", "Finance", 0.06),
76 ("WFC", "Wells Fargo", "Finance", 0.07),
77 ("GS", "Goldman Sachs", "Finance", 0.09),
78 ("MS", "Morgan Stanley", "Finance", 0.08),
79 ("C", "Citigroup", "Finance", 0.05),
80
81 # ENERGY (6 stocks) - variable returns
82 # Solver can pick max 5, so must choose best 5 from 6
83 ("XOM", "Exxon Mobil", "Energy", 0.04),
84 ("CVX", "Chevron Corp.", "Energy", 0.05),
85 ("COP", "ConocoPhillips", "Energy", 0.06),
86 ("SLB", "Schlumberger", "Energy", 0.03),
87 ("EOG", "EOG Resources", "Energy", 0.07),
88 ("PXD", "Pioneer Natural", "Energy", 0.08),
89]
90
91LARGE_DATASET_STOCKS = SMALL_DATASET_STOCKS + [
92 # Additional TECHNOLOGY (6 more -> 13 total)
93 ("CRM", "Salesforce", "Technology", 0.11),
94 ("ADBE", "Adobe Inc.", "Technology", 0.09),
95 ("ORCL", "Oracle Corp.", "Technology", 0.07),
96 ("CSCO", "Cisco Systems", "Technology", 0.06),
97 ("IBM", "IBM Corp.", "Technology", 0.04),
98 ("QCOM", "Qualcomm", "Technology", 0.13),
99
100 # Additional HEALTHCARE (6 more -> 12 total)
101 ("MRK", "Merck & Co.", "Healthcare", 0.08),
102 ("LLY", "Eli Lilly", "Healthcare", 0.16),
103 ("BMY", "Bristol-Myers", "Healthcare", 0.06),
104 ("AMGN", "Amgen Inc.", "Healthcare", 0.09),
105 ("GILD", "Gilead Sciences", "Healthcare", 0.05),
106 ("ISRG", "Intuitive Surgical", "Healthcare", 0.14),
107
108 # Additional FINANCE (4 more -> 10 total, no duplicates)
109 ("AXP", "American Express", "Finance", 0.10),
110 ("BLK", "BlackRock", "Finance", 0.11),
111 ("SCHW", "Charles Schwab", "Finance", 0.07),
112 ("USB", "U.S. Bancorp", "Finance", 0.04),
113
114 # Additional ENERGY (2 more -> 8 total, no duplicates)
115 ("OXY", "Occidental Petroleum", "Energy", 0.06),
116 ("HAL", "Halliburton", "Energy", 0.05),
117
118 # CONSUMER (new sector - 8 stocks)
119 ("AMZN", "Amazon.com", "Consumer", 0.14),
120 ("WMT", "Walmart", "Consumer", 0.06),
121 ("HD", "Home Depot", "Consumer", 0.08),
122 ("MCD", "McDonald's", "Consumer", 0.07),
123 ("NKE", "Nike Inc.", "Consumer", 0.09),
124 ("SBUX", "Starbucks", "Consumer", 0.05),
125 ("PG", "Procter & Gamble", "Consumer", 0.04),
126 ("KO", "Coca-Cola", "Consumer", 0.05),
127]
128# LARGE total: 25 + 6 + 6 + 4 + 2 + 8 = 51 stocks
129
130
131def generate_demo_data(demo_data: DemoData) -> PortfolioOptimizationPlan:
132 """
133 Generate demo data for portfolio optimization.
134
135 Args:
136 demo_data: Which demo dataset to generate (SMALL or LARGE)
137
138 Returns:
139 PortfolioOptimizationPlan with candidate stocks (all unselected initially)
140
141 Example:
142 >>> plan = generate_demo_data(DemoData.SMALL)
143 >>> len(plan.stocks)
144 20
145 >>> plan.stocks[0].stock_id
146 'AAPL'
147 """
148 config = demo_data_configs[demo_data]
149 stock_data = SMALL_DATASET_STOCKS if demo_data == DemoData.SMALL else LARGE_DATASET_STOCKS
150
151 stocks = [
152 StockSelection(
153 stock_id=ticker,
154 stock_name=name,
155 sector=sector,
156 predicted_return=predicted_return,
157 selection=None, # To be decided by solver
158 )
159 for ticker, name, sector, predicted_return in stock_data
160 ]
161
162 # Calculate max_per_sector from percentage
163 target_count = config.target_position_count
164 max_per_sector = max(1, int(config.max_sector_percentage * target_count))
165
166 # Create PortfolioConfig for constraints to access
167 portfolio_config = PortfolioConfig(
168 target_count=target_count,
169 max_per_sector=max_per_sector,
170 unselected_penalty=10000,
171 )
172
173 return PortfolioOptimizationPlan(
174 stocks=stocks,
175 target_position_count=config.target_position_count,
176 max_sector_percentage=config.max_sector_percentage,
177 portfolio_config=portfolio_config,
178 )
179
180
181def get_stock_summary(plan: PortfolioOptimizationPlan) -> str:
182 """
183 Generate a human-readable summary of the portfolio.
184
185 Useful for debugging and understanding the solution.
186 """
187 lines = [
188 "=" * 60,
189 "PORTFOLIO SUMMARY",
190 "=" * 60,
191 ]
192
193 selected = plan.get_selected_stocks()
194 if not selected:
195 lines.append("No stocks selected yet.")
196 return "\n".join(lines)
197
198 weight = plan.get_weight_per_stock()
199 expected_return = plan.get_expected_return()
200
201 lines.append(f"Selected: {len(selected)} stocks @ {weight*100:.1f}% each")
202 lines.append(f"Expected Return: {expected_return*100:.2f}%")
203 lines.append("")
204
205 # Group by sector
206 sector_stocks: dict[str, list[StockSelection]] = {}
207 for stock in selected:
208 if stock.sector not in sector_stocks:
209 sector_stocks[stock.sector] = []
210 sector_stocks[stock.sector].append(stock)
211
212 lines.append("BY SECTOR:")
213 for sector, stocks in sorted(sector_stocks.items()):
214 sector_weight = len(stocks) * weight * 100
215 lines.append(f" {sector}: {len(stocks)} stocks = {sector_weight:.1f}%")
216 for stock in sorted(stocks, key=lambda s: -s.predicted_return):
217 lines.append(f" - {stock.stock_id}: {stock.stock_name} ({stock.predicted_return*100:.1f}% pred)")
218
219 lines.append("")
220 lines.append(f"Score: {plan.score}")
221 lines.append("=" * 60)
222
223 return "\n".join(lines)
224