diff --git a/pyforecaster/forecasting_models/neural_forecasters.py b/pyforecaster/forecasting_models/neural_forecasters.py
index d899069..294769f 100644
--- a/pyforecaster/forecasting_models/neural_forecasters.py
+++ b/pyforecaster/forecasting_models/neural_forecasters.py
@@ -12,6 +12,8 @@
 import matplotlib.pyplot as plt
 from sklearn.preprocessing import StandardScaler
 from typing import Union
+from jax import jit
+from inspect import getmro
 
 def positive_lecun(key, shape, dtype=jnp.float32, init_type='normal'):
     # Start with standard lecun_normal initialization
@@ -45,6 +47,7 @@ def __call__(self, x):
         if isinstance(self.n_layers, int):
             layers = np.ones(self.n_layers) * self.n_neurons
             layers[-1] = self.n_out
+            layers = layers.astype(int)
         else:
             layers = self.n_layers
         for i, n in enumerate(layers):
@@ -53,6 +56,23 @@ def __call__(self, x):
                 x = nn.relu(x)
         return x
 
+def jitting_wrapper(fun, model):
+    return jit(partial(fun, model=model))
+
+def loss_fn(params, inputs, targets, model=None):
+    predictions = model(params, inputs)
+    return jnp.mean((predictions - targets) ** 2)
+
+def train_step(params, optimizer_state, inputs_batch, targets_batch, model=None, loss_fn=None):
+    values, grads = value_and_grad(loss_fn)(params, inputs_batch, targets_batch)
+    updates, opt_state = model.update(grads, optimizer_state, params)
+    return optax.apply_updates(params, updates), opt_state, values
+
+def predict_batch(pars, inputs, model=None):
+    return model.apply(pars, inputs)
+
+def predict_batch_picnn(pars, inputs, model=None):
+    return model.apply(pars, *inputs)
 
 class NN(ScenarioGenerator):
     scaler: StandardScaler = None
@@ -60,98 +80,143 @@ class NN(ScenarioGenerator):
     batch_size: int = None
     load_path: str = None
     n_out: int = None
-    n_epochs: int = 10
-    savepath_tr_plots: str = None
+    n_layers: int = None
+    n_hidden_x: int = None
+    pars: dict = None
+    n_epochs:int = 10
+    savepath_tr_plots:str = None
     stats_step: int = 50
     rel_tol: float = 1e-4
+    unnormalized_inputs: list = None
+    to_be_normalized:list = None
+    target_columns: list = None
+    iterate = None
+    def __init__(self, learning_rate: float = 0.01, batch_size: int = None, load_path: str = None,
+                 n_hidden_x: int = 100, n_out: int = None, n_layers: int = 3, pars: dict = None, q_vect=None,
+                 val_ratio=None, nodes_at_step=None, n_epochs: int = 10, savepath_tr_plots: str = None,
+                 stats_step: int = 50, rel_tol: float = 1e-4, unnormalized_inputs=None, **scengen_kwgs):
 
+        super().__init__(q_vect, val_ratio=val_ratio, nodes_at_step=nodes_at_step, **scengen_kwgs)
+        self.set_attr({"learning_rate": learning_rate,
+                       "batch_size": batch_size,
+                       "load_path": load_path,
+                       "n_hidden_x": n_hidden_x,
+                       "n_out": n_out,
+                       "n_layers": n_layers,
+                       "pars": pars,
+                       "n_epochs": n_epochs,
+                       "savepath_tr_plots": savepath_tr_plots,
+                       "stats_step": stats_step,
+                       "rel_tol": rel_tol,
+                       "unnormalized_inputs": unnormalized_inputs
+                       })
 
-    def __init__(self, n_out=1, q_vect=None, n_epochs=10, val_ratio=None, nodes_at_step=None, learning_rate=1e-3,
-                 nn_module=None, scengen_dict={}, batch_size=None,  **model_kwargs):
-        super().__init__(q_vect, val_ratio=val_ratio, nodes_at_step=nodes_at_step, **scengen_dict)
-        model = nn_module(n_out=n_out, **model_kwargs)
-        self.batch_size = batch_size
-        self.learning_rate = learning_rate
-        self.model = model
-        self.optimizer = optax.adam(learning_rate=self.learning_rate)
-        self.n_epochs = n_epochs
+        if self.load_path is not None:
+            self.load(self.load_path)
+        self.model = self.set_arch()
+        self.optimizer = optax.adamw(learning_rate=self.learning_rate)
 
-        @jit
-        def loss_fn(params, x, y):
-            predictions = model.apply(params, x)
-            return jnp.mean((predictions - y) ** 2)
-        @jit
-        def train_step(params, optimizer_state, x_batch, y_batch):
-            values, grads = value_and_grad(loss_fn)(params, x_batch, y_batch)
-            updates, opt_state = self.optimizer.update(grads, optimizer_state, params)
-            return optax.apply_updates(params, updates), opt_state, values
 
-        @jit
-        @partial(vmap, in_axes=(None, 0))
-        def predict_batch(pars, x):
-            return model.apply(pars, x)
+        self.predict_batch = vmap(jitting_wrapper(predict_batch, self.model), in_axes=(None, 0))
+        self.loss_fn = jitting_wrapper(loss_fn, self.predict_batch)
+        self.train_step = jitting_wrapper(partial(train_step, loss_fn=self.loss_fn), self.optimizer)
 
-        self.train_step = train_step
-        self.loss_fn = loss_fn
-        self.predict_batch = predict_batch
-        self.iterate = None
+    def get_class_properties_names(cls):
+        attributes = []
+        cls = cls if isinstance(cls, type) else cls.__class__
+        # Loop through the MRO (Method Resolution Order) to include parent classes
+        for base_class in getmro(cls):
+            attributes.extend(
+                key for key, value in base_class.__dict__.items()
+                if not callable(value)
+                and not key.startswith('__')
+                and not isinstance(value, (classmethod, staticmethod))
+            )
+        return list(set(attributes))  # Remove duplicates
+
+    def get_class_properties(self):
+        return {k: getattr(self, k) for k in self.get_class_properties_names()}
+
+    def save(self, save_path):
+        attrdict = self.get_class_properties()
+        with open(save_path, 'wb') as f:
+            pk.dump(attrdict, f, protocol=pk.HIGHEST_PROTOCOL)
+
+    def set_attr(self, attrdict):
+        [self.__setattr__(k, v) for k, v in attrdict.items()]
+
+    def load(self, save_path):
+        with open(save_path, 'rb') as f:
+            attrdict = pk.load(f)
+        self.set_attr(attrdict)
 
     @staticmethod
     def init_arch(nn_init, n_inputs_x=1):
         "divides data into training and test sets "
         key1, key2 = random.split(random.key(0))
-        x = random.normal(key1, (n_inputs_x, ))  # Dummy input data (for the first input)
+        x = random.normal(key1, (n_inputs_x,))  # Dummy input data (for the first input)
         init_params = nn_init.init(key2, x)  # Initialization call
         return init_params
 
-    def fit(self, x, y, n_epochs=None, savepath_tr_plots=None, stats_step=None, rel_tol=None):
+    def set_arch(self):
+        model = FeedForwardModule(n_layers=self.n_layers, n_neurons=self.n_hidden_x,
+                              n_out=self.n_out)
+        return model
+
+    def fit(self, inputs, targets, n_epochs=None, savepath_tr_plots=None, stats_step=None, rel_tol=None):
+        self.to_be_normalized = [c for c in inputs.columns if
+                                 c not in self.unnormalized_inputs] if self.unnormalized_inputs is not None else inputs.columns
         rel_tol = rel_tol if rel_tol is not None else self.rel_tol
         n_epochs = n_epochs if n_epochs is not None else self.n_epochs
         stats_step = stats_step if stats_step is not None else self.stats_step
-        self.scaler = StandardScaler().set_output(transform='pandas').fit(x)
+        self.scaler = StandardScaler().set_output(transform='pandas').fit(inputs[self.to_be_normalized])
+
+        inputs, targets, inputs_val_0, targets_val = self.train_val_split(inputs, targets)
+        training_len = inputs.shape[0]
+        validation_len = inputs_val_0.shape[0]
 
-        x, y, x_val_0, y_val = self.train_val_split(x, y)
-        self.target_columns = y.columns
+        self.target_columns = targets.columns
+        batch_size = self.batch_size if self.batch_size is not None else inputs.shape[0] // 10
+        num_batches = inputs.shape[0] // batch_size
 
-        batch_size = self.batch_size if self.batch_size is not None else x.shape[0] // 10
-        num_batches = y.shape[0] // batch_size
-        y = y.values
-        x = self.get_normalized_inputs(x)
-        x_val = self.get_normalized_inputs(x_val_0)
+        inputs = self.get_normalized_inputs(inputs)
+        inputs_val = self.get_normalized_inputs(inputs_val_0)
+        inputs_len = [i.shape[1] for i in inputs] if isinstance(inputs, tuple) else inputs.shape[1]
 
-        pars = self.init_arch(self.model, x.shape[1])
+        pars = self.init_arch(self.model, *np.atleast_1d(inputs_len))
         opt_state = self.optimizer.init(pars)
 
         tr_loss, val_loss = [], []
         k = 0
         finished = False
         for epoch in range(n_epochs):
-            rand_idx_all = np.random.choice(x.shape[0], x.shape[0], replace=False)
+            rand_idx_all = np.random.choice(training_len, training_len, replace=False)
             for i in tqdm(range(num_batches), desc='epoch {}/{}'.format(epoch, n_epochs)):
-                rand_idx = rand_idx_all[i*batch_size:(i+1)*batch_size]
-                x_batch = x[rand_idx, :]
-                y_batch = y[rand_idx, :]
+                rand_idx = rand_idx_all[i * batch_size:(i + 1) * batch_size]
+                inputs_batch = [i[rand_idx, :] for i in inputs] if isinstance(inputs, tuple) else inputs[rand_idx, :]
+                targets_batch = targets.values[rand_idx, :]
 
-                pars, opt_state, values = self.train_step(pars, opt_state, x_batch, y_batch)
+                pars, opt_state, values = self.train_step(pars, opt_state, inputs_batch, targets_batch)
 
                 if k % stats_step == 0 and k > 0:
                     self.pars = pars
 
-                    te_loss_i = self.loss_fn(pars, x_val, y_val.values)
-                    tr_loss_i = self.loss_fn(pars, x, y)
+                    te_loss_i = self.loss_fn(pars, inputs_val, targets_val.values)
+                    tr_loss_i = self.loss_fn(pars, inputs, targets.values)
                     val_loss.append(np.array(jnp.mean(te_loss_i)))
                     tr_loss.append(np.array(jnp.mean(tr_loss_i)))
 
-                    self.logger.warning('tr loss: {:0.2e}, te loss: {:0.2e}'.format(tr_loss[-1], val_loss[-1]))
+                    self.logger.info('tr loss: {:0.2e}, te loss: {:0.2e}'.format(tr_loss[-1], val_loss[-1]))
                     if len(tr_loss) > 1:
                         if savepath_tr_plots is not None or self.savepath_tr_plots is not None:
                             savepath_tr_plots = savepath_tr_plots if savepath_tr_plots is not None else self.savepath_tr_plots
-                            rand_idx_plt = np.random.choice(x_val.shape[0], 9)
-                            self.training_plots(x_val[rand_idx_plt, :],
-                                                y_val.values[rand_idx_plt, :], tr_loss, val_loss, savepath_tr_plots, k)
+
+                            rand_idx_plt = np.random.choice(validation_len, 9)
+                            self.training_plots([i[rand_idx_plt, :] for i in inputs] if isinstance(inputs, tuple) else inputs_val[rand_idx_plt, :],
+                                                targets_val.values[rand_idx_plt, :], tr_loss, val_loss, savepath_tr_plots, k)
 
                         rel_te_err = (val_loss[-2] - val_loss[-1]) / np.abs(val_loss[-2] + 1e-6)
-                        if rel_te_err<rel_tol:
+                        if rel_te_err < rel_tol:
                             finished = True
                             break
                 k += 1
@@ -159,19 +224,16 @@ def fit(self, x, y, n_epochs=None, savepath_tr_plots=None, stats_step=None, rel_
                 break
 
             self.pars = pars
-        super().fit(x_val_0, y_val)
+        super().fit(inputs_val_0, targets_val)
         return self
-    def get_normalized_inputs(self, inputs):
-        inputs = self.scaler.transform(inputs)
-        return inputs.values
 
-    def training_plots(self, x, y, target, tr_loss, te_loss, savepath, k):
-        n_instances = x.shape[0]
-        y_hat = self.predict_batch(self.pars, x, y)
+    def training_plots(self, inputs, target, tr_loss, te_loss, savepath, k):
+        n_instances = target.shape[0]
+        y_hat = self.predict_batch(self.pars, inputs)
 
         # make the appropriate numbers of subplots disposed as a square
         fig, ax = plt.subplots(int(np.ceil(np.sqrt(n_instances))), int(np.ceil(np.sqrt(n_instances))))
-        for i, a  in enumerate(ax.ravel()):
+        for i, a in enumerate(ax.ravel()):
             a.plot(y_hat[i, :])
             a.plot(target[i, :])
             a.set_title('instance {}, iter {}'.format(i, k))
@@ -188,6 +250,18 @@ def predict(self, inputs, **kwargs):
         y_hat = self.predict_batch(self.pars, x)
         return pd.DataFrame(y_hat, index=inputs.index, columns=self.target_columns)
 
+    def get_normalized_inputs(self, inputs):
+        inputs = inputs.copy()
+        self.to_be_normalized = [c for c in inputs.columns if
+                                 c not in self.unnormalized_inputs] if self.unnormalized_inputs is not None else inputs.columns
+        normalized_inputs = self.scaler.transform(inputs[self.to_be_normalized])
+        inputs.loc[:, self.to_be_normalized] = normalized_inputs.copy().values
+
+        return inputs.values
+
+    def get_denormalized_output(self, output):
+        self.scaler.inverse_transform(output.values)
+
 
 class FastLinReg(NN):
     def __init__(self, n_out=1, q_vect=None, n_epochs=10, val_ratio=None, nodes_at_step=None, learning_rate=1e-3,
@@ -275,9 +349,126 @@ def reproject_weights(params, rec_stable=False):
 
     return params
 
+class PICNN(NN):
+    inverter_learning_rate: float = 0.1
+    n_hidden_y: int = 100
+    optimization_vars: list = ()
+    rec_stable: bool = False
+    init_type: str = 'normal'
+    augment_ctrl_inputs: bool = False
+    layer_normalization: bool = False
+    def __init__(self, learning_rate: float = 0.01, batch_size: int = None, load_path: str = None,
+                 n_hidden_x: int = 100, n_out: int = None, n_layers: int = 3, pars: dict = None, q_vect=None,
+                 val_ratio=None, nodes_at_step=None, n_epochs: int = 10, savepath_tr_plots: str = None,
+                 stats_step: int = 50, rel_tol: float = 1e-4, unnormalized_inputs=None,
+                 inverter_learning_rate: float = 0.1, optimization_vars: list = (), target_columns: list = None,
+                 init_type='normal', augment_ctrl_inputs=False, layer_normalization=False,**scengen_kwgs):
+
+        super().__init__(learning_rate, batch_size, load_path, n_hidden_x, n_out, n_layers, pars, q_vect, val_ratio,
+                         nodes_at_step, n_epochs, savepath_tr_plots, stats_step, rel_tol, unnormalized_inputs,
+                         **scengen_kwgs)
+        self.set_attr({"inverter_learning_rate":inverter_learning_rate,
+                       "optimization_vars":optimization_vars,
+                       "target_columns":target_columns,
+                       "init_type":init_type,
+                       "augment_ctrl_inputs":augment_ctrl_inputs,
+                       "layer_normalization":layer_normalization
+                       })
+
+        if load_path is not None:
+            self.load(load_path)
+        self.n_hidden_y = 2 * len(self.optimization_vars) if augment_ctrl_inputs else len(self.optimization_vars)
+        self.model = self.set_arch()
+        self.optimizer = optax.adamw(learning_rate=self.learning_rate)
+        self.inverter_optimizer = optax.adabelief(learning_rate=self.inverter_learning_rate)
+
+        self.predict_batch = vmap(jitting_wrapper(predict_batch_picnn, self.model), in_axes=(None, 0))
+        self.loss_fn = jitting_wrapper(loss_fn, self.predict_batch)
+        self.train_step = jitting_wrapper(partial(train_step, loss_fn=self.loss_fn), self.optimizer)
+
+    def set_arch(self):
+        model = PartiallyICNN(num_layers=self.n_layers, features_x=self.n_hidden_x, features_y=self.n_hidden_y,
+                              features_out=self.n_out, init_type=self.init_type, augment_ctrl_inputs=self.augment_ctrl_inputs)
+        return model
+
+    @staticmethod
+    def init_arch(nn_init, n_inputs_x=1, n_inputs_opt=1):
+        "divides data into training and test sets "
+        key1, key2 = random.split(random.key(0))
+        x = random.normal(key1, (n_inputs_x, ))  # Dummy input data (for the first input)
+        y = random.normal(key1, (n_inputs_opt, ))  # Dummy input data (for the second input)
+        init_params = nn_init.init(key2, x, y)  # Initialization call
+        return init_params
+    def get_normalized_inputs(self, inputs):
+        inputs = inputs.copy()
+        self.to_be_normalized = [c for c in inputs.columns if c not in self.unnormalized_inputs] if self.unnormalized_inputs is not None else inputs.columns
+        normalized_inputs = self.scaler.transform(inputs[self.to_be_normalized])
+        inputs.loc[:, self.to_be_normalized] = normalized_inputs.copy().values
+
+        x = inputs[[c for c in inputs.columns if not c in self.optimization_vars]].values
+        y = inputs[self.optimization_vars].values
+        return x, y
+    def optimize(self, inputs, objective, n_iter=200, rel_tol=1e-4, recompile_obj=True, vanilla_gd=False, **objective_kwargs):
+        rel_tol = rel_tol if rel_tol is not None else self.rel_tol
+        inputs = inputs.copy()
+        x, y = self.get_normalized_inputs(inputs)
+
+        def _objective(y, x, **objective_kwargs):
+            return objective(self.model.apply(self.pars, x, y), y, **objective_kwargs)
+
+        # if the objective changes from one call to another, you need to recompile it. Slower but necessary
+        if recompile_obj:
+            @jit
+            def iterate(x, y, opt_state, **objective_kwargs):
+                for i in range(10):
+                    values, grads = value_and_grad(partial(_objective, **objective_kwargs))(y, x)
+                    if vanilla_gd:
+                        y -= grads * 1e-1
+                    else:
+                        updates, opt_state = self.inverter_optimizer.update(grads, opt_state, y)
+                        y = optax.apply_updates(y, updates)
+                return y, values
+            self.iterate = iterate
+        else:
+            if self.iterate is None:
+                @jit
+                def iterate(x, y, opt_state, **objective_kwargs):
+                    for i in range(10):
+                        values, grads = value_and_grad(partial(_objective, **objective_kwargs))(y, x)
+                        if vanilla_gd:
+                            y -= grads * 1e-1
+                        else:
+                            updates, opt_state = self.inverter_optimizer.update(grads, opt_state, y)
+                            y = optax.apply_updates(y, updates)
+                    return y, values
+
+                self.iterate = iterate
 
+        opt_state = self.inverter_optimizer.init(y)
+        y, values_old = self.iterate(x, y, opt_state, **objective_kwargs)
+        values_init = np.copy(values_old)
 
-class PICNN(ScenarioGenerator):
+        # do 10 iterations at a time to speed up, check for convergence
+        for i in range(n_iter//10):
+            y, values = self.iterate(x, y, opt_state, **objective_kwargs)
+            rel_improvement = (values_old - values) / (np.abs(values_old)+ 1e-12)
+            values_old = values
+            if rel_improvement < rel_tol:
+                break
+
+        print('optimization terminated at iter {}, final objective value: {:0.2e} '
+                         'rel improvement: {:0.2e}'.format((i+1)*10, values,
+                                                          (values_init-values)/(np.abs(values_init)+1e-12)))
+
+        inputs.loc[:, self.optimization_vars] = y.ravel()
+        inputs.loc[:, [c for c in inputs.columns if c not in  self.optimization_vars]] = x.ravel()
+        inputs.loc[:, self.to_be_normalized] = self.scaler.inverse_transform(inputs[self.to_be_normalized].values)
+        target_opt = self.predict(inputs)
+
+        y_opt = inputs.loc[:, self.optimization_vars].values.ravel()
+        return y_opt, inputs, target_opt, values
+
+class PICNN_old(ScenarioGenerator):
     "Partially input-convex neural network"
     learning_rate: float = 0.01
     inverter_learning_rate: float = 0.1
@@ -559,14 +750,17 @@ def get_denormalized_output(self, output):
 
 class RecStablePICNN(PICNN):
     rec_stable = True
-    def __init__(self, learning_rate: float = 0.01, inverter_learning_rate: float = 0.1, batch_size: int = None,
-                 load_path: str = None, n_hidden_x: int = 100, n_out: int = None,
-                 n_layers: int = 3, optimization_vars: list = (), pars: dict = None, target_columns: list = None,
-                 q_vect=None, val_ratio=None, nodes_at_step=None, n_epochs:int=10, savepath_tr_plots:str = None,
-                 stats_step:int=50, rel_tol:float=1e-4, init_type='normal', **scengen_kwgs):
-        super().__init__(learning_rate, inverter_learning_rate, batch_size, load_path, n_hidden_x, n_out, n_layers,
-                         optimization_vars, pars, target_columns, q_vect, val_ratio, nodes_at_step, n_epochs,
-                         savepath_tr_plots, stats_step, rel_tol, init_type, **scengen_kwgs)
+    def __init__(self, learning_rate: float = 0.01, batch_size: int = None, load_path: str = None,
+                 n_hidden_x: int = 100, n_out: int = None, n_layers: int = 3, pars: dict = None, q_vect=None,
+                 val_ratio=None, nodes_at_step=None, n_epochs: int = 10, savepath_tr_plots: str = None,
+                 stats_step: int = 50, rel_tol: float = 1e-4, unnormalized_inputs=None,
+                 inverter_learning_rate: float = 0.1, optimization_vars: list = (), target_columns: list = None,
+                 init_type='normal', augment_ctrl_inputs=False, layer_normalization=False,**scengen_kwgs):
+
+        super().__init__(learning_rate, batch_size, load_path, n_hidden_x, n_out, n_layers, pars, q_vect, val_ratio,
+                         nodes_at_step, n_epochs, savepath_tr_plots, stats_step, rel_tol, unnormalized_inputs,
+                         inverter_learning_rate, optimization_vars, target_columns, init_type, augment_ctrl_inputs,
+                         layer_normalization, **scengen_kwgs)
 
     def set_arch(self):
         model = PartiallyICNN(num_layers=self.n_layers, features_x=self.n_hidden_x, features_y=self.n_hidden_y,
diff --git a/pyforecaster/formatter.py b/pyforecaster/formatter.py
index 40ba62d..65ac883 100644
--- a/pyforecaster/formatter.py
+++ b/pyforecaster/formatter.py
@@ -273,7 +273,8 @@ def _simulate_transform(self, x=None):
         for tr in self.target_transformers:
             _ = tr.transform(x, simulate=True)
 
-    def plot_transformed_feature(self, x, feature, frames=100):
+    def plot_transformed_feature(self, x, feature, frames=100, ax_labels=None, legend_kwargs={},
+                   remove_spines=True, **kwargs):
         x_tr, target = self.transform(x)
         all_feats = pd.concat([x_tr, target], axis=1)
         fs = []
@@ -289,7 +290,8 @@ def plot_transformed_feature(self, x, feature, frames=100):
                     end_times.append(
                         t.metadata.loc[derived_features.index[derived_features['function'] == function], 'end_time'])
 
-        ani = ts_animation_bars(fs, start_times, end_times, names, frames=frames)
+        ani = ts_animation_bars(fs, start_times, end_times, frames=frames, ax_labels=ax_labels,
+                                legend_kwargs=legend_kwargs, remove_spines=remove_spines, **kwargs)
         return ani
 
     def cat_encoding(self, df_train, df_test, target, encoder=None, cat_features=None):
diff --git a/pyforecaster/plot_utils.py b/pyforecaster/plot_utils.py
index 8417bde..ead0bef 100644
--- a/pyforecaster/plot_utils.py
+++ b/pyforecaster/plot_utils.py
@@ -8,7 +8,7 @@
 import matplotlib.colors as colors
 from matplotlib.gridspec import GridSpec
 import networkx as nx
-
+from pyforecaster.scenred import plot_from_graph
 
 def basic_setup(subplots_tuple, width, height, b=0.15, l=0.15, w=0.22, r=None, style ='seaborn', **kwargs):
     plt.style.use(style)
@@ -190,14 +190,25 @@ def animate(i):
     return ani
 
 
-"""
-def tree_animation(ys:list, y_gt=None, times=None, frames=150):
+
+def plot_trees(ys:list, y_gt=None, times=None, frames=150, ax_labels=None, legend_kwargs={},
+                   remove_spines=True, savepath=None, **kwargs):
     "plot the first n_rows of the two y_te and y_hat matrices"
-    fig, ax = plt.subplots(1)
+    fig, ax = plt.subplots(1, **kwargs)
     f_min = np.min([np.min(np.array(list(nx.get_node_attributes(y, 'v').values()))) for y in ys])
     f_max = np.max([np.max(np.array(list(nx.get_node_attributes(y, 'v').values()))) for y in ys])
     lines = None
-    lines = scrd.plot_from_graph(ys[0], alpha=0.2, linewidth=1)
+    lines = plot_from_graph(ys[0], alpha=0.2, linewidth=1, ax=ax, color='r')
+    if ax_labels is not None:
+        for k, v in ax_labels.items():
+            if k == 'x':
+                ax.set_xlabel(v)
+            elif k == 'y':
+                ax.set_ylabel(v)
+    if remove_spines:
+        ax.spines['right'].set_visible(False)
+        ax.spines['top'].set_visible(False)
+
     if y_gt is not None:
         t = np.arange(len(y_gt[0][0])) if times is None else times
         lines_ground_truth = []
@@ -207,25 +218,37 @@ def tree_animation(ys:list, y_gt=None, times=None, frames=150):
     ax.set_ylim(f_min - np.abs(f_min) * 0.1, f_max + np.abs(f_max) * 0.1)
 
     def animate(i):
-        _ = scrd.plot_from_graph(ys[i], lines, alpha=0.2)
+        _ = plot_from_graph(ys[i], lines, alpha=0.2, ax=ax, color='r')
         if y_gt is not None:
             for y, l in zip(y_gt, lines_ground_truth):
                 l.set_data(t, y[i, :])
         return
+    if savepath is not None:
+        writervideo = animation.FFMpegWriter(fps=60)
+        animation.FuncAnimation(fig, animate, blit=False, frames=np.minimum(len(ys) - 1, frames), interval=100,
+                                repeat=False).save(savepath, writer=writervideo)
+    else:
+        return animation.FuncAnimation(fig, animate,  blit=False, frames=np.minimum(len(ys)-1, frames), interval=100,
+                                       repeat=False)
 
-    plt.pause(1e-5)
-    ani = animation.FuncAnimation(fig, animate,  blit=False, frames=np.minimum(len(ys)-1, frames), interval=100, repeat=False)
-
-    return ani
-"""
 
-def ts_animation_bars(ys:list, start_t:list, end_t:list, names:list, frames=150):
+def ts_animation_bars(ys:list, start_t:list, end_t:list, frames=150, ax_labels=None, legend_kwargs={},
+                   remove_spines=True, savepath=None, **kwargs):
     "plot the first n_rows of the two y_te and y_hat matrices"
-    fig, ax = plt.subplots(1)
+    fig, ax = plt.subplots(1, **kwargs)
     lines = []
     f_min = np.min([np.min(y) for y in ys])
     f_max = np.max([np.max(y) for y in ys])
     cm = plt.get_cmap('Set1')
+    if ax_labels is not None:
+        for k, v in  ax_labels.items():
+            if k == 'x':
+                ax.set_xlabel(v)
+            elif k == 'y':
+                ax.set_ylabel(v)
+    if remove_spines:
+        ax.spines['right'].set_visible(False)
+        ax.spines['top'].set_visible(False)
     def init():
         for y, s_t, e_t, idx in zip(ys, start_t, end_t, range(len(ys))):
             l = [ax.stairs(y_j.ravel(), np.array([s_t.iloc[j]/np.timedelta64(3600*24,'s'), e_t.iloc[j]/np.timedelta64(3600*24,'s')]), color=cm(idx)) for j, y_j in enumerate(y[0, :])]
@@ -238,18 +261,23 @@ def animate(i):
         for y, l, s_t, e_t in zip(ys, lines, start_t, end_t):
             for j, l_j in enumerate(l):
                 l_j.set_data(y[i, j].ravel(), np.array([s_t.iloc[j]/np.timedelta64(3600*24,'s'), e_t.iloc[j]/np.timedelta64(3600*24,'s')]))
+        ax.legend(**legend_kwargs)
         return [item for sublist in lines for item in sublist]
 
-    plt.pause(1e-5)
-    ani = animation.FuncAnimation(fig, animate, init_func=init,  blit=False, frames=np.minimum(ys[0].shape[0]-1, frames), interval=100, repeat=False)
-
-    return ani
+    if savepath is not None:
+        writervideo = animation.FFMpegWriter(fps=60)
+        animation.FuncAnimation(fig, animate, init_func=init, blit=False, frames=np.minimum(ys[0].shape[0] - 1, frames),
+                                interval=100, repeat=False).save(savepath, writer=writervideo)
+    else:
+        return animation.FuncAnimation(fig, animate, init_func=init, blit=False,
+                                       frames=np.minimum(ys[0].shape[0] - 1, frames), interval=100, repeat=False)
 
-def plot_quantiles(signals, qs, labels, n_rows=50, interval=1, step=1, repeat=False):
+def plot_quantiles(signals, qs, labels, n_rows=50, interval=1, step=1, repeat=False, ax_labels=None, legend_kwargs={},
+                   remove_spines=True, savepath=None, **kwargs):
     n_max = np.minimum(signals[0].shape[0], int(n_rows*step))
     n_rows = np.minimum(int(np.floor(signals[0].shape[0]/step)), n_rows)
     qs = qs.values if isinstance(qs, pd.DataFrame) else qs
-    fig, ax = plt.subplots(1)
+    fig, ax = plt.subplots(1, **kwargs)
     signals = signals if isinstance(signals, list) else [signals]
     t = np.arange(signals[0].shape[1])
     lines= []
@@ -260,11 +288,23 @@ def plot_quantiles(signals, qs, labels, n_rows=50, interval=1, step=1, repeat=Fa
     lineq = ax.plot(np.squeeze(qs[0, :, :]), 'r', lw=2, alpha=0.3)
     ax.set_ylim(np.min(qs[:n_max]), np.max(qs[:n_max])*1.05)
 
+    if ax_labels is not None:
+        for k, v in  ax_labels.items():
+            if k == 'x':
+                ax.set_xlabel(v)
+            elif k == 'y':
+                ax.set_ylabel(v)
+    if remove_spines:
+        ax.spines['right'].set_visible(False)
+        ax.spines['top'].set_visible(False)
+
+
     def animate(i):
         i = i * step
         for l, s in zip(lines, signals):
             l.set_data(t, s[i, :])
         [lineq[j].set_data(t, qsi) for j, qsi in enumerate(qs[i, :, :].T)]
+        ax.legend(**legend_kwargs)
         return (*lines, *lineq, )
 
     def init():
@@ -272,9 +312,14 @@ def init():
         plt.legend()
         return (lines[0],)
 
-    ani = animation.FuncAnimation(fig, animate, init_func=init, frames=n_rows, interval=interval, blit=True,
+
+    if savepath is not None:
+        writervideo = animation.FFMpegWriter(fps=60)
+        animation.FuncAnimation(fig, animate, init_func=init, frames=n_rows, interval=interval, blit=True,
+                                repeat=repeat).save(savepath, writer=writervideo)
+    else:
+        return animation.FuncAnimation(fig, animate, init_func=init, frames=n_rows, interval=interval, blit=True,
                                   repeat=repeat)
-    return ani
 
 
 def plot_scenarios_from_multilevel(scens, i=0, ax=None):
diff --git a/tests/test_boosters.py b/tests/test_boosters.py
index c88e6b9..b8e1b63 100644
--- a/tests/test_boosters.py
+++ b/tests/test_boosters.py
@@ -6,7 +6,7 @@
 import logging
 from pyforecaster.forecasting_models.gradientboosters import LGBMHybrid
 from pyforecaster.forecaster import LGBForecaster, LinearForecaster
-from pyforecaster.plot_utils import plot_quantiles
+from pyforecaster.plot_utils import plot_quantiles, plot_trees
 from pyforecaster.formatter import Formatter
 
 
@@ -44,5 +44,36 @@ def test_linear_val_split(self):
         plot_quantiles([y_te.iloc[:10, :], y_hat_lin.iloc[:10, :], y_hat_lgbh.iloc[:10, :], y_hat_lgb.iloc[:10, :]], q[:10, :, :], ['y_te', 'y_lin', 'y_lgbhybrid_1', 'y_hat_lgb'])
         plt.close('all')
 
+    def do_not_test_linear_val_split(self):
+
+        formatter = Formatter(logger=self.logger).add_transform(['all'], lags=np.arange(144),
+                                                                    relative_lags=True)
+        formatter.add_transform(['all'], lags=np.arange(144)+144*6)
+        formatter.add_transform(['all'], ['min', 'max'], agg_bins=[1, 2, 15, 24, 48, 144])
+        formatter.add_target_transform(['all'], lags=-np.arange(144)-1)
+
+        x, y = formatter.transform(self.data)
+        n_tr = int(len(x) * 0.8)
+        x_tr, x_te, y_tr, y_te = [x.iloc[:n_tr, :].copy(), x.iloc[n_tr:, :].copy(), y.iloc[:n_tr].copy(),
+                                  y.iloc[n_tr:].copy()]
+
+        m_lin = LinearForecaster(q_vect=np.linspace(0.01, 0.99, 11), cov_est_method='vanilla', val_ratio=0.6, online_tree_reduction=False,
+                                 nodes_at_step=np.logspace(0.5, 2, 144,base=5).astype(int)).fit(x_tr, y_tr)
+        y_hat_lgbh = m_lin.predict(x_te)
+        q = m_lin.predict_quantiles(x_te)
+
+        trees = m_lin.predict_trees(x_te.iloc[:500, :])
+        plot_trees(trees, [y_te.iloc[:500, :].values], frames=500,
+                              ax_labels={'x':"step [10 min]", "y": 'P [kW]'}, layout='tight', figsize=(8, 4),
+                              legend_kwargs={'loc':'upper right'}, savepath='linear_preds_tree.mp4')
+        plt.close('all')
+        plot_quantiles([y_te, y_hat_lgbh], q, ['y_te', 'y_hat'], n_rows=500,
+                              ax_labels={'x':"step [10 min]", "y": 'P [kW]'}, layout='tight', figsize=(8, 4),
+                              legend_kwargs={'loc':'upper right'}, savepath='linear_preds.mp4')
+
+
+        formatter.plot_transformed_feature(self.data, 'all',
+                             ax_labels={'x':"step [1 day]", "y": 'P [kW]'}, layout='tight', figsize=(8, 4),
+                             legend_kwargs={'loc':'upper right'}, frames=500, savepath='transformed_features.mp4')
 if __name__ == '__main__':
     unittest.main()
diff --git a/tests/test_formatter.py b/tests/test_formatter.py
index 5a185f3..6c71668 100644
--- a/tests/test_formatter.py
+++ b/tests/test_formatter.py
@@ -136,7 +136,7 @@ def test_aggregated_transform(self):
         formatter.add_transform([0], ['min', 'max'], agg_bins=[-10, -15, -20], nested=False)
         formatter.add_target_transform([0], lags=-np.arange(30)-1)
         x_transformed, y_transformed = formatter.transform(self.x2)
-        formatter.plot_transformed_feature(self.x2, 0, frames=100)
+        #formatter.plot_transformed_feature(self.x2, 0, frames=100)
 
         # test nested transform (much faster)
         formatter = pyf.Formatter().add_transform([0], lags=np.arange(10), agg_freq='20min',
@@ -144,7 +144,7 @@ def test_aggregated_transform(self):
         formatter.add_transform([0], ['min', 'max'], agg_bins=[-10, -15, -20], nested=True)
         formatter.add_target_transform([0], lags=-np.arange(30) - 1)
         x_transformed_nested, y_transformed = formatter.transform(self.x2)
-        formatter.plot_transformed_feature(self.x2, 0, frames=100)
+        #formatter.plot_transformed_feature(self.x2, 0, frames=100)
 
 
         formatter = pyf.Formatter().add_transform([0], ['min', 'max'], agg_bins=[0, 1, 4, 10])
@@ -152,10 +152,10 @@ def test_aggregated_transform(self):
         formatter.add_transform([0], ['mean'], lags=np.arange(10), agg_freq='20min', relative_lags=True)
         formatter.add_target_transform([0], lags=-np.arange(30)-1)
         x_transformed, y_transformed = formatter.transform(self.x3)
-        formatter.plot_transformed_feature(self.x3, 0, frames=30)
+        #formatter.plot_transformed_feature(self.x3, 0, frames=30)
 
         x_transformed, y_transformed = formatter.transform(self.x2)
-        formatter.plot_transformed_feature(self.x2, 0, frames=100)
+        #formatter.plot_transformed_feature(self.x2, 0, frames=100)
         plt.close('all')
 
     def test_speed(self):
@@ -288,11 +288,11 @@ def test_normalizers_complex(self):
 
 
     def test_normalizers_impossible(self):
-        x_private = pd.DataFrame(np.random.randn(500, 15),
-                                 index=pd.date_range('01-01-2020', '01-05-2020', 500, tz='Europe/Zurich'),
+        x_private = pd.DataFrame(np.random.randn(100, 15),
+                                 index=pd.date_range('01-01-2020', '01-05-2020', 100, tz='Europe/Zurich'),
                                  columns=pd.MultiIndex.from_product([['b1', 'b2', 'b3'], ['a', 'b', 'c', 'd', 'e']]))
-        x_shared = pd.DataFrame(np.random.randn(500, 5),
-                                index=pd.date_range('01-01-2020', '01-05-2020', 500, tz='Europe/Zurich'),
+        x_shared = pd.DataFrame(np.random.randn(100, 5),
+                                index=pd.date_range('01-01-2020', '01-05-2020', 100, tz='Europe/Zurich'),
                                 columns=pd.MultiIndex.from_product([['shared'], [0, 1, 2, 3, 4]]))
 
         df_mi = pd.concat([x_private, x_shared], axis=1)
diff --git a/tests/test_nns.py b/tests/test_nns.py
index cea8451..a1a7d5b 100644
--- a/tests/test_nns.py
+++ b/tests/test_nns.py
@@ -3,7 +3,7 @@
 import pandas as pd
 import numpy as np
 import logging
-from pyforecaster.forecasting_models.neural_forecasters import PICNN, RecStablePICNN
+from pyforecaster.forecasting_models.neural_forecasters import PICNN, RecStablePICNN, NN
 from pyforecaster.trainer import hyperpar_optimizer
 from pyforecaster.formatter import Formatter
 from pyforecaster.metrics import nmae
@@ -37,6 +37,27 @@ def test_ffnn(self):
 
         savepath_tr_plots = 'tests/results/ffnn_tr_plots'
 
+        # if not there, create directory savepath_tr_plots
+        if not exists(savepath_tr_plots):
+            makedirs(savepath_tr_plots)
+
+        m = NN(learning_rate=1e-3,  batch_size=1000, load_path=None, n_hidden_x=200,
+               n_out=y_tr.shape[1], n_layers=3).fit(x_tr,y_tr, n_epochs=1, savepath_tr_plots=savepath_tr_plots, stats_step=40)
+
+        y_hat_1 = m.predict(x_te.iloc[:100, :])
+
+
+    def test_picnn(self):
+        # normalize inputs
+        x = (self.x - self.x.mean(axis=0)) / (self.x.std(axis=0)+0.01)
+        y = (self.y - self.y.mean(axis=0)) / (self.y.std(axis=0)+0.01)
+
+        n_tr = int(len(x) * 0.8)
+        x_tr, x_te, y_tr, y_te = [x.iloc[:n_tr, :].copy(), x.iloc[n_tr:, :].copy(), y.iloc[:n_tr].copy(),
+                                  y.iloc[n_tr:].copy()]
+
+        savepath_tr_plots = 'tests/results/ffnn_tr_plots'
+
         # if not there, create directory savepath_tr_plots
         if not exists(savepath_tr_plots):
             makedirs(savepath_tr_plots)
@@ -45,18 +66,20 @@ def test_ffnn(self):
         optimization_vars = x_tr.columns[:100]
 
 
-        m = PICNN(learning_rate=1e-3,  batch_size=1000, load_path=None, n_hidden_x=200, n_hidden_y=200,
-               n_out=y_tr.shape[1], n_layers=3, optimization_vars=optimization_vars).fit(x_tr,
-                                                                                           y_tr,
-                                                                                           n_epochs=1,
-                                                                                           savepath_tr_plots=savepath_tr_plots,
-                                                                                           stats_step=40)
+        m_1 = PICNN(learning_rate=1e-3, batch_size=1000, load_path=None, n_hidden_x=200, n_hidden_y=200,
+                  n_out=y_tr.shape[1], n_layers=3, optimization_vars=optimization_vars).fit(x_tr,
+                                                                                            y_tr,
+                                                                                            n_epochs=1,
+                                                                                            savepath_tr_plots=savepath_tr_plots,
+                                                                                            stats_step=40)
+
+        y_hat_1 = m_1.predict(x_te)
+        m_1.save('tests/results/ffnn_model.pk')
 
-        y_hat_1 = m.predict(x_te.iloc[:100, :])
-        m.save('tests/results/ffnn_model.pk')
         n = PICNN(load_path='tests/results/ffnn_model.pk')
         y_hat_2 = n.predict(x_te.iloc[:100, :])
 
+
         m = RecStablePICNN(learning_rate=1e-3, batch_size=1000, load_path=None, n_hidden_x=200, n_hidden_y=200,
                   n_out=1, n_layers=3, optimization_vars=optimization_vars).fit(x_tr,y_tr.iloc[:, [0]],
                                                                                             n_epochs=1,