using new transform function in plotting

plotting.py

@@ -255,12 +255,16 @@ def plot_hist_2D_events(plotname, valsx, valsy, nbinsx, xmin, xmax, nbinsy, ymin
 def plot_cond_avg_actmax_2D(plotname, model, layer, neuron, ranges,
                             varx_index, vary_index,
                             nbinsx, xmin, xmax, nbinsy, ymin, ymax,
-                            scaler=None,
+                            transform=None, inverse_transform=None,
                             ntries=20,
                             step=1,
                             maxit=1,
                             **kwargs):
 
+    transform_given = [fn is not None for fn in [transform, inverse_transform]]
+    if any(transform_given) and not all(transform_given):
+        raise ValueError("Need to pass both transform and inverse_transform if data should be transformed")
+
     xedges = np.linspace(xmin, xmax, nbinsx)
     yedges = np.linspace(ymin, ymax, nbinsy)
 
@@ -271,12 +275,12 @@ def plot_cond_avg_actmax_2D(plotname, model, layer, neuron, ranges,
     for ix, x in enumerate(xedges):
         for iy, y in enumerate(yedges):
             random_event = create_random_event(ranges)
-            if scaler is not None:
-                random_event = scaler.inverse_transform(random_event)
+            if inverse_transform is not None:
+                random_event = inverse_transform(random_event)
             for index, val in [(varx_index, x), (vary_index, y)]:
                 random_event[0][index] = val
-            if scaler is not None:
-                random_event = scaler.transform(random_event)
+            if transform is not None:
+                random_event = transform(random_event)
             act = np.mean([max_activation_wrt_input(gradient_function, random_event, maxit=maxit, step=step, const_indices=[varx_index, vary_index])[0][0] for i in range(ntries)])
             hist[ix][iy] = act
 
@@ -374,15 +378,19 @@ if __name__ == "__main__":
 
 
         plot_NN_vs_var_2D_all("mt_vs_met_all.pdf", means=mean_signal,
-                              model=c.model, transform_function=c.scaler.transform,
+                              model=c.model, transform_function=c.transform,
                               varx_index=c.fields.index("met"),
                               vary_index=c.fields.index("mt"),
                               nbinsx=100, xmin=0, xmax=1000,
                               nbinsy=100, ymin=0, ymax=500,
                               varx_label="met [GeV]", vary_label="mt [GeV]")
 
+        input_transform = c.transform
+        if hasattr(c, "get_input_list"):
+            input_transform = lambda x : c.get_input_list(c.transform(x))
+
         plot_NN_vs_var_2D("mt_vs_met_crosscheck.pdf", means=mean_signal,
-                          scorefun=get_single_neuron_function(c.model, layer=3, neuron=0, scaler=c.scaler),
+                          scorefun=get_single_neuron_function(c.model, layer=3, neuron=0, input_transform=input_transform),
                           varx_index=c.fields.index("met"),
                           vary_index=c.fields.index("mt"),
                           nbinsx=100, xmin=0, xmax=1000,
@@ -398,7 +406,7 @@ if __name__ == "__main__":
 
         losses, events = get_max_activation_events(c.model, ranges, ntries=100000, layer=3, neuron=0, threshold=0.2)
 
-        events = c.scaler.inverse_transform(events)
+        events = c.inverse_transform(events)
 
         plot_hist_2D_events(
             "mt_vs_met_actmaxhist.pdf",
@@ -432,7 +440,7 @@ if __name__ == "__main__":
             c.fields.index("mt"),
             30, 0, 1000,
             30, 0, 500,
-            scaler=c.scaler,
+            transform=c.transform, inverse_transform=c.inverse_transform,
             varx_label="met [GeV]", vary_label="mt [GeV]",
         )
 
@@ -441,7 +449,7 @@ if __name__ == "__main__":
 
         c.load(reload=True)
 
-        utrf_x_test = c.scaler.inverse_transform(c.x_test)
+        utrf_x_test = c.inverse_transform(c.x_test)
 
         plot_hist_2D_events(
             "mt_vs_output_signal_test.pdf",
@@ -489,7 +497,7 @@ if __name__ == "__main__":
     def test_profile():
 
         c.load(reload=True)
-        utrf_x_test = c.scaler.inverse_transform(c.x_test)
+        utrf_x_test = c.inverse_transform(c.x_test)
 
         plot_profile_2D(
             "mt_vs_met_profilemean_sig.pdf",

utils.py

@@ -13,16 +13,18 @@ from meme import cache
 logger = logging.getLogger(__name__)
 logger.addHandler(logging.NullHandler())
 
-def get_single_neuron_function(model, layer, neuron, scaler=None, input_transform=None):
+def get_single_neuron_function(model, layer, neuron, input_transform=None):
 
     f = K.function([model.input]+[K.learning_phase()], [model.layers[layer].output[:,neuron]])
 
     def eval_single_neuron(x):
-        if scaler is not None:
-            x_eval = scaler.transform(x)
+        if input_transform is not None:
+            x_eval = input_transform(x)
         else:
             x_eval = x
-        return f([x_eval])[0]
+        if not isinstance(x_eval, list):
+            x_eval = [x_eval]
+        return f(x_eval)[0]
 
     return eval_single_neuron