Add survival proba along with events in SurvTRACE (#45)

soda-inria · Jan 26, 2024 · 76d7c81 · 76d7c81
1 parent bc57614
commit 76d7c81
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/hazardous/survtrace/_model.py b/hazardous/survtrace/_model.py
@@ -299,7 +299,9 @@ def predict_survival_function(self, X):
         return surv
 
     def predict_cumulative_incidence(self, X):
-        return 1 - self.predict_survival_function(X)
+        risks = 1 - self.predict_survival_function(X)
+        surv = (1 - risks.sum(axis=0))[None, :, :]
+        return np.concatenate([surv, risks], axis=0)
 
 
 class _SurvTRACEModule(nn.Module):