Some tests and fixed a bug in kaplan meier plot

fractalis/analytics/tasks/kaplan_meier_survival/main.py

@@ -11,7 +11,6 @@ from fractalis.analytics.task import AnalyticTask
 from fractalis.analytics.tasks.shared import utils
 
 
-
 logger = logging.getLogger(__name__)
 
 
@@ -54,7 +53,7 @@ class KaplanMeierSurvivalTask(AnalyticTask):
                 E = None  # default is nothing is censored
                 if event_observed:
                     # find observation boolean value for every duration
-                    E = event_observed[0].merge(sub_df, how='left', on='id')
+                    E = event_observed[0].merge(sub_df, how='right', on='id')
                     E = [bool(x) and not np.isnan(x) for x in E['value']]
                     assert len(E) == len(T)
                 kmf.fit(durations=T, event_observed=E)

tests/unit/analytics/kaplan_meier_survival/test_main.py

+"""This module contains tests for the kaplan_meier_survival module."""
+
+from lifelines.datasets import load_waltons
+
+from fractalis.analytics.tasks.kaplan_meier_survival.main \
+    import KaplanMeierSurvivalTask
+
+
+class TestKaplanMeierSurvivalTask:
+
+    task = KaplanMeierSurvivalTask()
+
+    def test_correct_output_for_simple_input(self):
+        df = load_waltons()
+        df.insert(0, 'id', df.index)
+        duration = df[['id', 'T']].copy()
+        duration.insert(1, 'feature', 'duration')
+        duration.columns.values[2] = 'value'
+        results = self.task.main(durations=[duration],
+                                 categories=[],
+                                 event_observed=[],
+                                 id_filter=[],
+                                 subsets=[])
+        assert 'timeline' in results['stats'][''][0]
+        assert 'median' in results['stats'][''][0]
+        assert 'survival_function' in results['stats'][''][0]
+        assert 'confidence_interval' in results['stats'][''][0]
+
+    def test_correct_output_for_complex_input(self):
+        df = load_waltons()
+        df.insert(0, 'id', df.index)
+        duration = df[['id', 'T']].copy()
+        duration.insert(1, 'feature', 'duration')
+        duration.columns.values[2] = 'value'
+        event_observed = df[['id', 'E']].copy()
+        event_observed.insert(1, 'feature', 'was_observed')
+        categories = df[['id', 'group']].copy()
+        categories.insert(1, 'feature', 'group')
+        results = self.task.main(durations=[duration],
+                                 categories=[categories],
+                                 event_observed=[event_observed],
+                                 id_filter=[],
+                                 subsets=[])
+        assert results['stats']['control']
+        assert results['stats']['miR-137']