#34 Add "observed" variable option to Survival Analysis

fractalis/analytics/tasks/survival/main.py

@@ -58,7 +58,7 @@ class SurvivalTask(AnalyticTask):
                 if event_observed:
                     # find observation boolean value for every duration
                     E = event_observed[0].merge(sub_df, how='right', on='id')
-                    E = [bool(x) and not np.isnan(x) for x in E['value']]
+                    E = [not x for x in pd.isnull(E['value_x'])]
                     assert len(E) == len(T)
                 if estimator == 'NelsonAalen':
                     fitter = NelsonAalenFitter()

tests/unit/analytics/survival/test_main.py

@@ -37,8 +37,11 @@ class TestSurvivalTask:
         duration.columns.values[2] = 'value'
         event_observed = df[['id', 'E']].copy()
         event_observed.insert(1, 'feature', 'was_observed')
+        event_observed.columns.values[2] = 'value'
+        event_observed = event_observed[event_observed['value'] == 1]
         categories = df[['id', 'group']].copy()
         categories.insert(1, 'feature', 'group')
+        categories.columns.values[2] = 'value'
         results = self.task.main(durations=[duration],
                                  categories=[categories],
                                  event_observed=[event_observed],