Update src/optimization.py

src/optimization.py

@@ -1,21 +1,25 @@
-# src/optimization.py
 import os
 import pandas as pd
 import numpy as np
+import yfinance as yf
 from pypfopt.efficient_frontier import EfficientFrontier
 from pypfopt import risk_models, expected_returns, plotting
 import matplotlib.pyplot as plt
 
 TRADING_DAYS = 252
 
-def load_processed_prices(proc_dir: str, tickers: list[str]) -> dict[str, pd.DataFrame]:
-    """Load processed CSVs for each ticker."""
+def fetch_live_prices(tickers: list[str], start_date: str, end_date: str) -> dict[str, pd.DataFrame]:
+    """
+    Fetch live stock prices from yfinance for a list of tickers.
+    """
     data = {}
-    for t in tickers:
-        fn = os.path.join(proc_dir, f"{t}_processed.csv")
-        df = pd.read_csv(fn, parse_dates=["Date"], index_col="Date")
+    for ticker in tickers:
+        df = yf.download(ticker, start=start_date, end=end_date, progress=False)
         df = df.sort_index()
-        data[t] = df
+        if df.empty:
+            st.error(f"Could not fetch data for {ticker}. Please check the ticker symbol.")
+            st.stop()
+        data[ticker] = df
     return data
 
 
@@ -104,30 +108,22 @@ def optimize_portfolio(mu: pd.Series,
 
 def optimize_for_target(mu: pd.Series,
                         cov: pd.DataFrame,
-                        target_type: str = "return",
-                        target_value: float = 0.20) -> dict:
+                        target_type: str,
+                        target_value: float,
+                        bounds: tuple = (0, 1)) -> dict:
     """
-    Optimize portfolio for a specific target return or volatility.
-    
-    Args:
-        mu (pd.Series): Expected annual returns.
-        cov (pd.DataFrame): Annualized covariance matrix.
-        target_type (str): 'return' or 'volatility'.
-        target_value (float): The desired target value (e.g., 0.20 for 20%).
-        
-    Returns:
-        dict with portfolio weights and performance.
+    Optimize for a specific target return or volatility.
     """
-    ef = EfficientFrontier(mu, cov)
+    ef = EfficientFrontier(mu, cov, weight_bounds=bounds)
     
     try:
-        if target_type == "return":
+        if target_type == 'return':
             weights = ef.efficient_return(target_value)
-        elif target_type == "volatility":
+        elif target_type == 'volatility':
             weights = ef.efficient_risk(target_value)
         else:
-            raise ValueError("target_type must be 'return' or 'volatility'")
-            
+            return {"error": "Invalid target_type. Use 'return' or 'volatility'."}
+
         cleaned_weights = ef.clean_weights()
         perf = ef.portfolio_performance(verbose=False)
         
@@ -139,19 +135,13 @@ def optimize_for_target(mu: pd.Series,
                 "sharpe": perf[2],
             }
         }
-    except ValueError as e:
-        # Handle cases where the target is not achievable
+    except Exception as e:
         return {"weights": {}, "performance": {"error": str(e)}}
 
 
 def plot_efficient_frontier(mu, cov, results: dict):
     """
-    Plot the efficient frontier and mark key portfolios.
-    
-    Args:
-        mu (pd.Series): Expected annual returns.
-        cov (pd.DataFrame): Annualized covariance matrix.
-        results (dict): Dictionary containing the results from optimize_portfolio.
+    Plot the efficient frontier with key portfolios marked.
     """
     ef = EfficientFrontier(mu, cov)
     fig, ax = plt.subplots(figsize=(8, 6))
@@ -159,13 +149,17 @@ def plot_efficient_frontier(mu, cov, results: dict):
     # Plot the frontier
     plotting.plot_efficient_frontier(ef, ax=ax, show_assets=False)
     
-    # Plot Max Sharpe portfolio
-    sharpe_perf = results["perf_sharpe"]
-    ax.scatter(sharpe_perf["volatility"], sharpe_perf["return"], marker="*", s=150, c="r", label="Max Sharpe")
+    # Get weights for plotting
+    sharpe_weights = np.array(list(results["weights_sharpe"].values()))
+    minvol_weights = np.array(list(results["weights_minvol"].values()))
     
-    # Plot Min Volatility portfolio
+    # Find portfolio returns/volatility for plotting
+    sharpe_perf = results["perf_sharpe"]
     minvol_perf = results["perf_minvol"]
-    ax.scatter(minvol_perf["volatility"], minvol_perf["return"], marker="P", s=150, c="g", label="Min Volatility")
+    
+    # Plot markers
+    ax.scatter(sharpe_perf["volatility"], sharpe_perf["return"], marker="*", color="r", s=250, label="Max Sharpe")
+    ax.scatter(minvol_perf["volatility"], minvol_perf["return"], marker="X", color="g", s=200, label="Min Volatility")
     
     ax.set_title("Efficient Frontier")
     ax.legend()