1 files changed, 20 insertions, 18 deletions
diff --git a/python/exploration/beta_trade.py b/python/exploration/beta_trade.py
index 8563f11b..8f6825ef 100644
--- a/python/exploration/beta_trade.py
+++ b/python/exploration/beta_trade.py
@@ -11,29 +11,31 @@ from statsmodels.tsa.ar_model import AR
 
 import matplotlib.pyplot as plt
 
-def calc_returns(save_feather=False):
-    returns = (index_returns(index=['IG', 'HY'], tenor='5yr', years=None).
-               reset_index('tenor', drop=True))
-    returns_hy = (returns.
-                  xs('HY', level=1).
-                  dropna().
-                  groupby(level=['date'], as_index=False).
-                  nth(-1))
-    returns_ig = returns.xs('IG', level=1)
-    # hy starts trading later than ig, so we line it up based on hy series
-    df = pd.merge(returns_hy, returns_ig, left_index=True, right_index=True,
-                  suffixes=('_hy','_ig'))
-    returns = df[['price_return_hy', 'price_return_ig']]
-    returns.columns = ['hy', 'ig']
-    returns = returns.dropna()['2009-03-20':]
+def calc_returns(save_feather=False, index_pair=['HY','IG']):
+    returns = (index_returns(index=['EU','IG', 'HY'],tenor='5yr',years=None).
+               reset_index(['tenor','series'],drop=True))
+    returns1 = {}
+    returns1[index_pair[0]] = (returns.
+                               xs(index_pair[0], level=1).
+                               dropna().
+                               groupby(level=['date'], as_index=False).
+                               nth(-1))
+    returns1[index_pair[1]] = returns.xs(index_pair[1], level=1)
+    df = pd.merge(returns1[index_pair[0]],returns1[index_pair[1]],
+                  left_index=True,right_index=True,
+                  suffixes=('_'+index_pair[0],'_'+index_pair[1]))
+    returns = df[['price_return_'+index_pair[0], 'price_return_'+index_pair[1]]]
+    returns.columns = [index_pair[0], index_pair[1]]
+    min_date = max(pd.Timestamp('20090320'),returns.index.min()[0])
+    returns = returns.dropna()[min_date:]
     if save_feather:
         feather.write_dataframe(returns.reset_index(),
                                 os.path.join(os.environ["DATA_DIR"], "index_returns.fth"))
-    return returns.reset_index('series', drop=True)
+    return returns.reset_index('maturity', drop=True)
 
-def calc_betas(returns=None, spans=[5, 20]):
+def calc_betas(returns=None, spans=[5, 20], index_pair=['HY','IG']):
     if returns is None:
-        returns = calc_returns()
+        returns = calc_returns(index_pair=index_pair)
     return [(returns.
              ewm(span=span).
              cov().