Merge remote-tracking branch 'origin/main'

Author: sschrod

bites/analyse/analyse_utils.py

@@ -129,15 +129,6 @@ def find_nearest_index(array, value):
         pred1_cf[i] = surv1_cf.axes[0][find_nearest_index(surv1_cf.iloc[:, i].values, death_probability)]
     ITE1 = pred1 - pred1_cf
 
-    correct_predicted_probability=None
-    if best_treatment is not None:
-        pred_best_choice0 = ((pred0 - pred0_cf) < 0) * 1
-        pred_best_choice1 = ((pred1 - pred1_cf) > 0) * 1
-        correct_predicted_probability = np.sum(np.append(best_treatment[treatment == 0] == pred_best_choice0,
-                                                         best_treatment[treatment == 1] == pred_best_choice1)) \
-                                        / (pred_best_choice0.size + pred_best_choice1.size)
-        print('Fraction best choice: ' + str(correct_predicted_probability))
-
     ITE = np.zeros(X.shape[0])
     k, j = 0, 0
     for i in range(X.shape[0]):
@@ -148,29 +139,30 @@ def find_nearest_index(array, value):
             ITE[i] = ITE1[j]
             j = j + 1
 
+    correct_predicted_probability=None
+    if best_treatment is not None:
+        correct_predicted_probability=np.sum(best_treatment==(ITE>0)*1)/best_treatment.shape[0]
+        print('Fraction best choice: ' + str(correct_predicted_probability))
+
     return ITE, correct_predicted_probability
 
 def get_ITE_CFRNet(model, X, treatment, best_treatment=None):
 
     pred,_ = model.predict_numpy(X, treatment)
     pred_cf,_ = model.predict_numpy(X, 1-treatment)
 
-    correct_predicted_probability=None
-    if best_treatment is not None:
-        pred_best_choice0=(pred_cf[treatment==0]-pred[treatment==0]>0) * 1
-        pred_best_choice1=(pred[treatment==1]-pred_cf[treatment==1]>0) * 1
-        correct_predicted_probability = np.sum(np.append(best_treatment[treatment == 0] == pred_best_choice0,
-                                                         best_treatment[treatment == 1] == pred_best_choice1)) \
-                                        / (pred_best_choice0.size + pred_best_choice1.size)
-        print('Fraction best choice: ' + str(correct_predicted_probability))
-
     ITE = np.zeros(X.shape[0])
     for i in range(X.shape[0]):
         if treatment[i] == 0:
             ITE[i] = pred_cf[i]-pred[i]
         else:
             ITE[i] = pred[i]-pred_cf[i]
 
+    correct_predicted_probability=None
+    if best_treatment is not None:
+        correct_predicted_probability=np.sum(best_treatment==(ITE>0)*1)/best_treatment.shape[0]
+        print('Fraction best choice: ' + str(correct_predicted_probability))
+
     return ITE, correct_predicted_probability
 
 
@@ -206,15 +198,6 @@ def find_nearest_index(array, value):
         pred1_cf[i] = surv1_cf.axes[0][find_nearest_index(surv1_cf.iloc[:, i].values, death_probability)]
     ITE1 = pred1 - pred1_cf
 
-    correct_predicted_probability=None
-    if best_treatment is not None:
-        pred_best_choice0 = ((pred0 - pred0_cf) < 0) * 1
-        pred_best_choice1 = ((pred1 - pred1_cf) > 0) * 1
-        correct_predicted_probability = np.sum(np.append(best_treatment[treatment == 0] == pred_best_choice0,
-                                                         best_treatment[treatment == 1] == pred_best_choice1)) \
-                                        / (pred_best_choice0.size + pred_best_choice1.size)
-        print('Fraction best choice: ' + str(correct_predicted_probability))
-
     ITE = np.zeros(X.shape[0])
     k, j = 0, 0
     for i in range(X.shape[0]):
@@ -225,6 +208,11 @@ def find_nearest_index(array, value):
             ITE[i] = ITE1[j]
             j = j + 1
 
+    correct_predicted_probability=None
+    if best_treatment is not None:
+        correct_predicted_probability=np.sum(best_treatment==(ITE>0)*1)/best_treatment.shape[0]
+        print('Fraction best choice: ' + str(correct_predicted_probability))
+
     return ITE, correct_predicted_probability
 
 def get_ITE_DeepSurv(model, X, treatment, best_treatment=None, death_probability=0.5):
@@ -257,15 +245,6 @@ def find_nearest_index(array, value):
         pred1_cf[i] = surv1_cf.axes[0][find_nearest_index(surv1_cf.iloc[:, i].values, death_probability)]
     ITE1 = pred1 - pred1_cf
 
-    correct_predicted_probability=None
-    if best_treatment is not None:
-        pred_best_choice0 = ((pred0 - pred0_cf) < 0) * 1
-        pred_best_choice1 = ((pred1 - pred1_cf) > 0) * 1
-        correct_predicted_probability = np.sum(np.append(best_treatment[treatment == 0] == pred_best_choice0,
-                                                         best_treatment[treatment == 1] == pred_best_choice1)) \
-                                        / (pred_best_choice0.size + pred_best_choice1.size)
-        print('Fraction best choice: ' + str(correct_predicted_probability))
-
     ITE = np.zeros(X.shape[0])
     k, j = 0, 0
     for i in range(X.shape[0]):
@@ -276,6 +255,11 @@ def find_nearest_index(array, value):
             ITE[i] = ITE1[j]
             j = j + 1
 
+    correct_predicted_probability=None
+    if best_treatment is not None:
+        correct_predicted_probability=np.sum(best_treatment==(ITE>0)*1)/best_treatment.shape[0]
+        print('Fraction best choice: ' + str(correct_predicted_probability))
+
     return ITE, correct_predicted_probability
 
 

bites/model/Fit.py

@@ -13,7 +13,7 @@
 from torch.utils.data import *
 
 
-def fit (config, X_train, Y_train, event_train, treatment_train=None,**kwargs):
+def fit (config, X_train, Y_train, event_train=None, treatment_train=None,**kwargs):
     """
     :param config: config file as given in the examples.
     :param X_train: np.array(num_samples, features)

examples/RGBSG_analyse.py

@@ -70,7 +70,7 @@
 
     elif method == 'CFRNet':
         model, config = get_best_model("ray_results/" + method + "_RGBSG")
-        pred_ite=get_ITE_CFRNet(model, X_test, treatment_test, best_treatment=None)
+        pred_ite,_=get_ITE_CFRNet(model, X_test, treatment_test, best_treatment=None)
 
 
     # The loaded model can be used for further analysis!