Added code for importance testing

src/main.py

@@ -10,6 +10,7 @@ import tensorflow as tf
 import keras_tuner as kt
 
 import numpy as np
+import matplotlib.pyplot as plt
 
 from helpers import get_knots, get_params, generate_model
 from loaders import load_dataset, split_train_test_validation
@@ -96,6 +97,8 @@ def main():
     elif mode == "test":
         test(test_dataset, bs)
 
+    elif mode == "importance":
+        plot_importance(test_dataset)
 
 def train(model, train_dataset, val_dataset, bs):
     """Training function
@@ -136,6 +139,46 @@ def train(model, train_dataset, val_dataset, bs):
         w.writeheader()
         w.writerow(history.history)
 
+def plot_importance(test_dataset):
+    """Function to plot importance in results
+
+    Args:
+        test_dataset (tf.data.Dataset): Test dataset
+    """    
+
+    # Loading the model
+    model = tf.keras.models.load_model(checkpoint_filepath)
+
+    # Creating test labels and input datasets
+    test_dataset = test_dataset.map(lambda x, y: x).as_numpy()
+    
+    el_num = (np.random.rand(1) * len_db * 0.075 * len(knots)).astype(int)
+
+    for i, x_element in enumerate(el_num):
+        if i==el_num:
+            input_sequence = x_element
+
+    # ACTIVATION-BASED METHOD
+    input_tensor = model.input
+    output_tensor = model.output
+
+    for i, layer in enumerate(model.layers):
+        predicted_output, activations = tf.keras.backend.function([input_tensor], [output_tensor, *layer.output])(input_sequence)
+        activations_per_neuron = np.abs(activations).mean(axis=(0, 1))
+
+        plt.plot(np.arange(Nbeads), activations_per_neuron)
+        plt.xlabel("Bead Index")
+        plt.savefig(os.path.join(checkpoint_filepath,f"activation_layer_{i}_test_{el_num}.pdf"))
+        plt.cla()
+
+        # WEIGHT-BASED METHOD
+        weights = layer.get_weights()[0]
+        weights_per_neuron = np.abs(weights).sum(axis=0)
+
+        plt.plot(np.arange(Nbeads), weights_per_neuron)
+        plt.xlabel("Bead Index")
+        plt.savefig(os.path.join(checkpoint_filepath,f"weights_layer_{i}.pdf"))
+        plt.cla()
 
 def test(test_dataset, bs):
     """Testing function for the models created