compare.py

@@ -44,7 +44,7 @@ def overlay_ROC(filename, *projects, **kwargs):
     colors = prop_cycle.by_key()['color']
 
     for p, color in zip(projects, colors):
-        fpr, tpr, threshold = roc_curve(p.y_test, p.scores_test, sample_weight = p.w_test)
+        fpr, tpr, threshold = roc_curve(p.l_test, p.scores_test, sample_weight = p.w_test)
         fpr = 1.0 - fpr
         try:
             roc_auc = auc(tpr, fpr)
@@ -61,8 +61,8 @@ def overlay_ROC(filename, *projects, **kwargs):
         if plot_thresholds:
             ax2.plot(tpr, threshold, "--", color=color)
         if lumifactor is not None:
-            sumw_b = p.w_test[p.y_test==0].sum()*lumifactor
-            sumw_s = p.w_test[p.y_test==1].sum()*lumifactor
+            sumw_b = p.w_test[p.l_test==0].sum()*lumifactor
+            sumw_s = p.w_test[p.l_test==1].sum()*lumifactor
             ax_abs_b.plot(tpr, (1.-fpr)*sumw_b, alpha=0)
             ax_abs_b.invert_yaxis()
             ax_abs_s.plot(tpr*sumw_s, fpr, alpha=0)

test/test_toolkit.py

@@ -7,6 +7,7 @@ from keras.layers import GRU
 
 from KerasROOTClassification import ClassificationProject, ClassificationProjectRNN
 
+
 def create_dataset(path):
 
     # create example dataset with (low-weighted) noise added
@@ -50,6 +51,7 @@ def test_ClassificationProject(tmp_path):
         layers=3,
         nodes=128,
     )
+
     c.train(epochs=200)
     c.plot_all_inputs()
     c.plot_loss()

utils.py

@@ -197,14 +197,26 @@ def weighted_quantile(values, quantiles, sample_weight=None, values_sorted=False
 
 class WeightedRobustScaler(RobustScaler):
 
-    def fit(self, X, y=None, weights=None):
-        if not np.isnan(X).any():
+    def fit(self, X, y=None, weights=None, mask_value=None):
+        if not np.isnan(X).any() and mask_value is not None and weights is None:
             # these checks don't work for nan values
-            super(WeightedRobustScaler, self).fit(X, y)
-        if weights is None:
-            return self
+            return super(WeightedRobustScaler, self).fit(X, y)
         else:
-            wqs = np.array([weighted_quantile(X[:,i][~np.isnan(X[:,i])], [0.25, 0.5, 0.75], sample_weight=weights) for i in range(X.shape[1])])
+            if weights is None:
+                weights = np.ones(len(self.X))
+            wqs = []
+            for i in range(X.shape[1]):
+                mask = ~np.isnan(X[:,i])
+                if mask_value is not None:
+                    mask &= (X[:,i] != mask_value)
+                wqs.append(
+                    weighted_quantile(
+                        X[:,i][mask],
+                        [0.25, 0.5, 0.75],
+                        sample_weight=weights[mask]
+                    )
+                )
+            wqs = np.array(wqs)
             self.center_ = wqs[:,1]
             self.scale_ = wqs[:,2]-wqs[:,0]
             self.scale_ = _handle_zeros_in_scale(self.scale_, copy=False)