MAINT: Save data and rebuild model (#49)

Author: eigenfoo

bayesalpha/author_model.py

@@ -9,6 +9,7 @@
 import pandas as pd
 from sklearn.preprocessing import LabelEncoder
 import pymc3 as pm
+import xarray as xr
 from .serialize import to_xarray
 from ._version import get_versions
 from .base import BayesAlphaResult
@@ -56,27 +57,27 @@ def __init__(self, data):
         self.model = self._build_model(data)
 
         self.coords = {
-            'author':    data.meta_user_id.drop_duplicates().values,
-            'algo':      data.meta_algorithm_id.drop_duplicates().values,
-            'backtest':  data.meta_code_id.values
+            'meta_user_id':      data.meta_user_id.drop_duplicates().values,
+            'meta_algorithm_id': data.meta_algorithm_id.drop_duplicates().values,
+            'meta_code_id':      data.meta_code_id.values
         }
 
         self.dims = {
             'mu_global':       (),
-            'mu_author':       ('author', ),
-            'mu_author_raw':   ('author', ),
+            'mu_author':       ('meta_user_id', ),
+            'mu_author_raw':   ('meta_user_id', ),
             'mu_author_sd':    (),
-            'sigma_author':    ('author', ),
+            'sigma_author':    ('meta_user_id', ),
             'sigma_author_sd': (),
-            'mu_algo':         ('algo', ),
-            'mu_algo_raw':     ('algo', ),
+            'mu_algo':         ('meta_algorithm_id', ),
+            'mu_algo_raw':     ('meta_algorithm_id', ),
             'mu_algo_sd':      (),
-            'sigma_algo':      ('algo', ),
+            'sigma_algo':      ('meta_algorithm_id', ),
             'sigma_algo_sd':   (),
-            'mu_backtest':     ('backtest', ),
-            'sigma_backtest':  ('backtest', ),
-            'alpha_author':    ('author', ),
-            'alpha_algo':      ('algo', )
+            'mu_backtest':     ('meta_code_id', ),
+            'sigma_backtest':  ('meta_code_id', ),
+            'alpha_author':    ('meta_user_id', ),
+            'alpha_algo':      ('meta_algorithm_id', )
         }
 
     def _build_model(self, data):
@@ -155,7 +156,12 @@ class AuthorModelResult(BayesAlphaResult):
     def rebuild_model(self, data=None):
         """ Return an AuthorModelBuilder that recreates the original model. """
         if data is None:
-            data = self.trace._data.to_pandas().copy()
+            data = (self.trace
+                        ._data
+                        .to_pandas()
+                        .rename('perf_sharpe_ratio_is')
+                        .reset_index()
+                        .copy())
 
         return AuthorModelBuilder(data)
 
@@ -233,6 +239,12 @@ def fit_authors(data,
     trace.attrs['model-version'] = get_versions()['version']
     trace.attrs['model-type'] = AUTHOR_MODEL_TYPE
 
+    if save_data:
+        d = data.set_index(['meta_user_id',
+                            'meta_algorithm_id',
+                            'meta_code_id']).squeeze()
+        trace['_data'] = xr.DataArray(d)
+
     return AuthorModelResult(trace)