dont fail if no beta coefs estimated

hammerlab · Aug 5, 2016 · 5a022fb · 5a022fb
1 parent 871d7be
commit 5a022fb
Showing 1 changed file with 11 additions and 8 deletions.
diff --git a/survivalstan/survivalstan.py b/survivalstan/survivalstan.py
@@ -158,14 +158,17 @@ def fit_stan_survival_model(df, formula, event_col, model_code,
         **kwargs
     )
 
-    beta_coefs = pd.DataFrame(
-        survival_fit.extract()['beta'],
-        columns = x_df.columns
-    )
-    beta_coefs.reset_index(0, inplace = True)
-    beta_coefs = beta_coefs.rename(columns = {'index':'iter'})
-    beta_coefs = pd.melt(beta_coefs, id_vars = ['iter'])
-    beta_coefs['model_cohort'] = model_cohort
+    try:
+        beta_coefs = pd.DataFrame(
+            survival_fit.extract()['beta'],
+            columns = x_df.columns
+        )
+        beta_coefs.reset_index(0, inplace = True)
+        beta_coefs = beta_coefs.rename(columns = {'index':'iter'})
+        beta_coefs = pd.melt(beta_coefs, id_vars = ['iter'])
+        beta_coefs['model_cohort'] = model_cohort
+    except:
+        beta_coefs = None
 
     ## prep by-group coefs if group specified
     if group_id_col: