BUG: Rename coordinates for data in author model (#51)

quantopian · Jul 18, 2018 · 4556ad0 · 4556ad0
1 parent 48bf273
commit 4556ad0
Showing 1 changed file with 13 additions and 4 deletions.
diff --git a/bayesalpha/author_model.py b/bayesalpha/author_model.py
@@ -159,8 +159,12 @@ def rebuild_model(self, data=None):
             data = (self.trace
                         ._data
                         .to_pandas()
-                        .rename('perf_sharpe_ratio_is')
                         .reset_index()
+                        # Rename back to the expected column names.
+                        .rename_axis(['meta_user_id',
+                                      'meta_algorithm_id',
+                                      'meta_code_id',
+                                      'perf_sharpe_ratio_is'], axis=1)
                         .copy())
 
         return AuthorModelBuilder(data)
@@ -240,9 +244,14 @@ def fit_authors(data,
     trace.attrs['model-type'] = AUTHOR_MODEL_TYPE
 
     if save_data:
-        d = data.set_index(['meta_user_id',
-                            'meta_algorithm_id',
-                            'meta_code_id']).squeeze()
+        d = (data.set_index(['meta_user_id',
+                             'meta_algorithm_id',
+                             'meta_code_id'])
+                 # Rename index names to avoid name collision.
+                 .rename_axis(['data_meta_user_id',
+                               'data_meta_algorithm_id',
+                               'data_meta_code_id'], axis=0)
+                 .squeeze())
         trace['_data'] = xr.DataArray(d)
 
     return AuthorModelResult(trace)