get rid of numpy warning for kaplan-meier

better · Apr 14, 2018 · 14df4e6 · 14df4e6
1 parent 20849bd
commit 14df4e6
Showing 1 changed file with 7 additions and 4 deletions.
diff --git a/convoys/single.py b/convoys/single.py
@@ -23,11 +23,14 @@ def fit(self, B, T):
             self._ts.append(t)
             prod_s_terms *= 1 - d/n
             self._ss.append(prod_s_terms)
-            try:
-                sum_var_terms += d / (n*(n-d))
-            except ZeroDivisionError:
+            if d == n == 1:
                 sum_var_terms = float('inf')
-            self._vs.append(1 / numpy.log(prod_s_terms)**2 * sum_var_terms)
+            else:
+                sum_var_terms += d / (n*(n-d))
+            if sum_var_terms > 0:
+                self._vs.append(1 / numpy.log(prod_s_terms)**2 * sum_var_terms)
+            else:
+                self._vs.append(0)
             n -= 1
         self.get_j = lambda t: numpy.searchsorted(self._ts, t)