modified: estimators.py

autonlab · Mar 25, 2022 · fec3d6f · fec3d6f
1 parent 529c6cc
commit fec3d6f
Showing 1 changed file with 10 additions and 4 deletions.
diff --git a/auton_survival/estimators.py b/auton_survival/estimators.py
@@ -725,10 +725,16 @@ def __init__(self, treated_model, control_model):
     self.treated_model = treated_model
     self.control_model = control_model
 
-  def predict_counterfactuals(self, features, times):
+  def predict_counterfactual_survival(self, features, times):
 
-    control_outcomes = self.control_model.predict(features, times)
-    treated_outcomes = self.treated_model.predict(features, times)
+    control_outcomes = self.control_model.predict_survival(features, times)
+    treated_outcomes = self.treated_model.predict_survival(features, times)
 
     return treated_outcomes, control_outcomes
-
+
+  def predict_counterfactual_risk(self, features, times):
+
+    control_outcomes = self.control_model.predict_risk(features, times)
+    treated_outcomes = self.treated_model.predict_risk(features, times)
+
+    return treated_outcomes, control_outcomes