Merge pull request #106 from pzzhang/main

Fixing rebalancing

Author: santoshlite

src/empyrial/main.py

@@ -65,6 +65,8 @@ def __init__(
         optimizer=None,
         max_vol=0.15,
         diversification=1,
+        expected_returns=None,
+        risk_model=None,
         # confidences=None,
         # view=None,
         min_weights=None,
@@ -84,6 +86,12 @@ def __init__(
         self.rebalance = rebalance
         self.max_vol = max_vol
         self.diversification = diversification
+        self.expected_returns = expected_returns
+        if expected_returns is not None:
+            assert expected_returns in ["mean_historical_return", "ema_historical_return", "capm_return"], f"Expected return method: {expected_returns} not supported yet!"
+        self.risk_model = risk_model
+        if risk_model is not None:
+            assert risk_model in ["sample_cov", "semicovariance", "exp_cov", "ledoit_wolf", "ledoit_wolf_constant_variance", "ledoit_wolf_single_factor", "ledoit_wolf_constant_correlation", "oracle_approximating"], f"Risk model: {risk_model} not supported yet!"
         self.max_weights = max_weights
         self.min_weights = min_weights
         self.risk_manager = risk_manager
@@ -115,26 +123,34 @@ def __init__(
                 self.diversification,
                 self.min_weights,
                 self.max_weights,
+                self.expected_returns,
+                self.risk_model
             )
 
 
 def get_returns(stocks, wts, start_date, end_date=TODAY):
     if len(stocks) > 1:
         assets = yf.download(stocks, start=start_date, end=end_date, progress=False)["Adj Close"]
         assets = assets.filter(stocks)
-        ret_data = assets.pct_change()[1:]
-        returns = (ret_data * wts).sum(axis=1)
+        initial_alloc = wts/assets.iloc[0]
+        if initial_alloc.isna().any():
+            raise ValueError("Some stock is not available at initial state!")
+        portfolio_value = (assets * initial_alloc).sum(axis=1)
+        returns = portfolio_value.pct_change()[1:]
         return returns
     else:
         df = yf.download(stocks, start=start_date, end=end_date, progress=False)["Adj Close"]
         df = pd.DataFrame(df)
-        returns = df.pct_change()
+        returns = df.pct_change()[1:]
         return returns
 
 
 def get_returns_from_data(data, wts):
-    ret_data = data.pct_change()[1:]
-    returns = (ret_data * wts).sum(axis=1)
+    initial_alloc = wts/data.iloc[0]
+    if initial_alloc.isna().any():
+        raise ValueError("Some stock is not available at initial state!")
+    portfolio_value = (data * initial_alloc).sum(axis=1)
+    returns = portfolio_value.pct_change()[1:]
     return returns
 
 
@@ -535,12 +551,12 @@ def efficient_frontier(my_portfolio, perf=True) -> list:
     # sometimes we will pick a date range where company isn't public we can't set price to 0 so it has to go to 1
     df = df.fillna(1)
 
-    mu = expected_returns.mean_historical_return(df)
-    S = risk_models.sample_cov(df)
+    mu = expected_returns.return_model(df, method=my_portfolio.expected_returns)
+    S = risk_models.risk_matrix(df, method=my_portfolio.risk_model)
 
     # optimize for max sharpe ratio
     ef = EfficientFrontier(mu, S)
-
+    ef.add_objective(objective_functions.L2_reg, gamma=my_portfolio.diversification)
     if my_portfolio.min_weights is not None:
         ef.add_constraint(lambda x: x >= my_portfolio.min_weights)
     if my_portfolio.max_weights is not None:
@@ -608,8 +624,8 @@ def mean_var(my_portfolio, vol_max=0.15, perf=True) -> list:
     # sometimes we will pick a date range where company isn't public we can't set price to 0 so it has to go to 1
     prices = prices.fillna(1)
 
-    mu = expected_returns.capm_return(prices)
-    S = risk_models.CovarianceShrinkage(prices).ledoit_wolf()
+    mu = expected_returns.return_model(prices, method=my_portfolio.expected_returns)
+    S = risk_models.risk_matrix(prices, method=my_portfolio.risk_model)
 
     ef = EfficientFrontier(mu, S)
     ef.add_objective(objective_functions.L2_reg, gamma=my_portfolio.diversification)
@@ -643,8 +659,8 @@ def min_var(my_portfolio, perf=True) -> list:
     prices = ohlc["Adj Close"].dropna(how="all")
     prices = prices.filter(my_portfolio.portfolio)
 
-    mu = expected_returns.capm_return(prices)
-    S = risk_models.CovarianceShrinkage(prices).ledoit_wolf()
+    mu = expected_returns.return_model(prices, method=my_portfolio.expected_returns)
+    S = risk_models.risk_matrix(prices, method=my_portfolio.risk_model)
 
     ef = EfficientFrontier(mu, S)
     ef.add_objective(objective_functions.L2_reg, gamma=my_portfolio.diversification)
@@ -794,6 +810,8 @@ def make_rebalance(
     div,
     min,
     max,
+    expected_returns,
+    risk_model,
 ) -> pd.DataFrame:
     sdate = str(start_date)[:10]
     if rebalance[0] != sdate:
@@ -836,6 +854,8 @@ def make_rebalance(
                 diversification=div,
                 min_weights=min,
                 max_weights=max,
+                expected_returns=expected_returns,
+                risk_model=risk_model,
             )
 
         except TypeError:
@@ -849,6 +869,8 @@ def make_rebalance(
                 diversification=div,
                 min_weights=min,
                 max_weights=max,
+                expected_returns=expected_returns,
+                risk_model=risk_model,
             )
 
         output_df["{}".format(dates[i + 1])] = portfolio.weights
@@ -864,6 +886,8 @@ def make_rebalance(
             diversification=div,
             min_weights=min,
             max_weights=max,
+            expected_returns=expected_returns,
+            risk_model=risk_model,
         )
 
     except TypeError:
@@ -876,11 +900,15 @@ def make_rebalance(
             diversification=div,
             min_weights=min,
             max_weights=max,
+            expected_returns=expected_returns,
+            risk_model=risk_model,
         )
 
     output_df["{}".format(TODAY)] = portfolio.weights
 
     make_rebalance.output = output_df
+    print("Rebalance schedule: ")
+    print(output_df)
     return output_df