Repository: baler-collaboration/baler (language: python)

Issue:
149 move to numpy arrays

Permission issues when creating files using docker
Solve the issue of docker creating baler output files as root. Mainly, make output files readable and writeable by other users


Discussion / hints:

Fix patch:
diff --git a/Dockerfile b/Dockerfile
index 111181b5..d8a511c6 100644
--- a/Dockerfile
+++ b/Dockerfile
@@ -56,4 +56,27 @@ COPY --from=python-base /baler-root/baler/*.py /baler-root/baler/README.md ./
 # Configure run time
 ENV PYTHONUNBUFFERED=1
 WORKDIR /baler-root/
-ENTRYPOINT ["python", "baler"]
+
+# Configure fixuid env
+RUN addgroup --gid 1000 docker && \
+    adduser --uid 1000 \
+    --ingroup docker \
+    --home /home/docker \
+    --shell /bin/sh \
+    --disabled-password \
+    --gecos "" \
+    docker
+
+# Install fixuid
+RUN apt update && \
+    apt install --no-install-recommends -y curl && \
+    USER=docker && \
+    GROUP=docker && \
+    curl -SsL https://github.com/boxboat/fixuid/releases/download/v0.5.1/fixuid-0.5.1-linux-amd64.tar.gz | tar -C /usr/local/bin -xzf - && \
+    chown root:root /usr/local/bin/fixuid && \
+    chmod 4755 /usr/local/bin/fixuid && \
+    mkdir -p /etc/fixuid && \
+    printf "user: $USER\ngroup: $GROUP\n" > /etc/fixuid/config.yml
+
+COPY entrypoint.sh entrypoint.sh
+ENTRYPOINT ["./entrypoint.sh"]
diff --git a/README.md b/README.md
index a40cf518..aa620398 100644
--- a/README.md
+++ b/README.md
@@ -112,6 +112,7 @@ Running with Docker requires slight modifications to the above commands. The bas
 
 ```console
 docker run \
+-u ${UID}:${GID} \ 
 --mount type=bind,source=${PWD}/projects/,target=/baler-root/projects \
 --mount type=bind,source=${PWD}/data/,target=/baler-root/data \
 ghcr.io/uomresearchit/baler:latest 
@@ -120,6 +121,7 @@ ghcr.io/uomresearchit/baler:latest
 
 Where:
   * `docker run` invokes docker and specifies the running of a container
+  * `-u ${UID}:${GID}` tells the container to use your username to create files
   * `--mount type=bind,source=${PWD}/projects/,target=/baler-root/projects` mounts the local (host) directory `./projects` to the container at `/projects`
   * `ghcr.io/uomresearchit/baler:latest` specifies the container to run
   
@@ -129,6 +131,7 @@ Therefore the three commands detailed above become:
 
 ```console
 docker run \
+-u ${UID}:${GID} \ 
 --mount type=bind,source=${PWD}/projects/,target=/baler-root/projects \
 --mount type=bind,source=${PWD}/data/,target=/baler-root/data \
 ghcr.io/uomresearchit/baler:latest \
@@ -139,6 +142,7 @@ ghcr.io/uomresearchit/baler:latest \
 ### Compress: ### 
 ```console
 docker run \
+-u ${UID}:${GID} \ 
 --mount type=bind,source=${PWD}/projects/,target=/baler-root/projects \
 --mount type=bind,source=${PWD}/data/,target=/baler-root/data \
 ghcr.io/uomresearchit/baler:latest \
@@ -149,6 +153,7 @@ ghcr.io/uomresearchit/baler:latest \
 ### Decompress: ###
 ```console
 docker run \
+-u ${UID}:${GID} \ 
 --mount type=bind,source=${PWD}/projects/,target=/baler-root/projects \
 --mount type=bind,source=${PWD}/data/,target=/baler-root/data \
 ghcr.io/uomresearchit/baler:latest \
@@ -166,6 +171,7 @@ To allow the Docker image access to the system GPU a modification to the standar
 ```console
 docker run \
 --gpus all \
+-u ${UID}:${GID} \ 
 --mount type=bind,source=${PWD}/projects/,target=/baler-root/projects \
 --mount type=bind,source=${PWD}/data/,target=/baler-root/data \
 ghcr.io/uomresearchit/baler:latest \
@@ -188,6 +194,7 @@ This image may be run using (e.g.):
 
 ```console
 docker run \
+-u ${UID}:${GID} \ 
 --mount type=bind,source=${PWD}/projects/,target=/baler-root/projects \
 --mount type=bind,source=${PWD}/data/,target=/baler-root/data \
 baler:latest \
@@ -205,6 +212,7 @@ An example command is given here:
 
 ```console
 docker run \
+-u ${UID}:${GID} \ 
 --mount type=bind,source=${PWD}/projects/,target=/baler-root/projects \
 --mount type=bind,source=${PWD}/data/,target=/baler-root/data \
 --mount type=bind,source=${PWD}/baler/modules,target=/baler-root/baler/modules \
diff --git a/baler/baler.py b/baler/baler.py
index 20c1ece1..39b8d3e3 100644
--- a/baler/baler.py
+++ b/baler/baler.py
@@ -2,6 +2,7 @@
 import time
 
 import pandas as pd
+import numpy as np
 
 import modules.helper as helper
 
@@ -29,11 +30,9 @@ def perform_training(config, project_path):
         test_set_norm,
         number_of_columns,
         normalization_features,
-        full_norm,
-        full_pre_norm,
-        cleared_col_names,
     ) = helper.process(
-        config.input_path,
+        config.data_path,
+        config.names_path,
         config.custom_norm,
         config.test_size,
         config.energy_conversion,
@@ -45,6 +44,7 @@ def perform_training(config, project_path):
     except AttributeError:
         print(config.latent_space_size, config.number_of_columns)
         assert number_of_columns == config.number_of_columns
+
     device = helper.get_device()
 
     ModelObject = helper.model_init(config.model_name)
@@ -60,39 +60,45 @@ def perform_training(config, project_path):
     reconstructed_data = helper.detach(reconstructed_data_tensor)
 
     print("Un-normalzing...")
+    print(np.mean(test_data[0]))
     start = time.time()
     test_data_renorm = helper.renormalize(
         test_data,
-        normalization_features["True min"],
-        normalization_features["Feature Range"],
+        normalization_features[0],
+        normalization_features[1],
     )
     reconstructed_data_renorm = helper.renormalize(
         reconstructed_data,
-        normalization_features["True min"],
-        normalization_features["Feature Range"],
+        normalization_features[0],
+        normalization_features[1],
     )
     end = time.time()
     print("Un-normalization took:", f"{(end - start) / 60:.3} minutes")
 
-    helper.to_pickle(test_data_renorm, output_path + "before.pickle")
-    helper.to_pickle(reconstructed_data_renorm, output_path + "after.pickle")
-    helper.to_pickle(full_pre_norm, output_path + "fulldata_energy.pickle")
+    # helper.to_pickle(test_data_renorm, output_path + "before.pickle")
+    # helper.to_pickle(reconstructed_data_renorm, output_path + "after.pickle")
+    # helper.to_pickle(full_pre_norm, output_path + "fulldata_energy.pickle")
+    np.save(output_path + "before.npy", test_data_renorm)
+    np.save(output_path + "after.npy", reconstructed_data_renorm)
 
-    normalization_features.to_csv(project_path + "model/cms_normalization_features.csv")
-    helper.model_saver(trained_model, project_path + "model/model.pt")
+    np.save(
+        project_path + "compressed_output/normalization_features.npy",
+        normalization_features,
+    )
+    helper.model_saver(trained_model, project_path + "compressed_output/model.pt")
 
 
 def perform_plotting(project_path, config):
     output_path = project_path + "plotting/"
-    helper.plot(project_path)
-    helper.loss_plotter(project_path + "training/loss_data.csv", output_path, config)
+    helper.plot(project_path, config)
+    helper.loss_plotter(project_path + "training/loss_data.npy", output_path, config)
 
 
 def perform_compression(config, project_path):
     print("Compressing...")
     start = time.time()
     compressed, data_before, cleared_col_names = helper.compress(
-        model_path=project_path + "model/model.pt",
+        model_path=project_path + "compressed_output/model.pt",
         config=config,
     )
     # Converting back to numpyarray
@@ -101,55 +107,34 @@ def perform_compression(config, project_path):
 
     print("Compression took:", f"{(end - start) / 60:.3} minutes")
 
-    helper.to_pickle(compressed, project_path + "compressed_output/compressed.pickle")
-    helper.to_pickle(
-        data_before, project_path + "compressed_output/cleandata_pre_comp.pickle"
-    )
-    helper.to_pickle(
-        cleared_col_names, project_path + "compressed_output/column_names.pickle"
-    )
+    np.save(project_path + "compressed_output/compressed.npy", compressed)
+    np.save(project_path + "compressed_output/names.npy", cleared_col_names)
 
 
 def perform_decompression(save_as_root, model_name, project_path):
     print("Decompressing...")
-    cleared_col_names = helper.from_pickle(
-        project_path + "compressed_output/column_names.pickle"
-    )
+    cleared_col_names = np.load(project_path + "compressed_output/names.npy")
     start = time.time()
     decompressed = helper.decompress(
-        model_path=project_path + "model/model.pt",
-        input_path=project_path + "compressed_output/compressed.pickle",
+        model_path=project_path + "compressed_output/model.pt",
+        input_path=project_path + "compressed_output/compressed.npy",
         model_name=model_name,
     )
 
     # Converting back to numpyarray
     decompressed = helper.detach(decompressed)
-    normalization_features = pd.read_csv(
-        project_path + "model/cms_normalization_features.csv"
+    normalization_features = np.load(
+        project_path + "compressed_output/normalization_features.npy"
     )
 
     decompressed = helper.renormalize(
         decompressed,
-        normalization_features["True min"],
-        normalization_features["Feature Range"],
+        normalization_features[0],
+        normalization_features[1],
     )
     end = time.time()
     print("Decompression took:", f"{(end - start) / 60:.3} minutes")
-
-    # False by default
-    if save_as_root:
-        helper.to_root(
-            decompressed,
-            cleared_col_names,
-            project_path + "decompressed_output/decompressed.root",
-        )
-        helper.to_pickle(
-            decompressed, project_path + "decompressed_output/decompressed.pickle"
-        )
-    else:
-        helper.to_pickle(
-            decompressed, project_path + "decompressed_output/decompressed.pickle"
-        )
+    np.save(project_path + "decompressed_output/decompressed.npy", decompressed)
 
 
 def print_info(project_path):
diff --git a/baler/modules/data_processing.py b/baler/modules/data_processing.py
index 9a6bbcb5..827f0899 100644
--- a/baler/modules/data_processing.py
+++ b/baler/modules/data_processing.py
@@ -92,16 +92,13 @@ def clean_data(df, config):
 
 
 def find_minmax(data):
-    data = np.array(data)
     data = list(data)
     true_max_list = np.apply_along_axis(np.max, axis=0, arr=data)
     true_min_list = np.apply_along_axis(np.min, axis=0, arr=data)
 
     feature_range_list = true_max_list - true_min_list
 
-    normalization_features = pd.DataFrame(
-        {"True min": true_min_list, "Feature Range": feature_range_list}
-    )
+    normalization_features = np.array([true_min_list, feature_range_list])
     return normalization_features
 
 
diff --git a/baler/modules/helper.py b/baler/modules/helper.py
index 56bebe66..de15aea9 100644
--- a/baler/modules/helper.py
+++ b/baler/modules/helper.py
@@ -2,7 +2,7 @@
 import os
 import pickle
 import sys
-import numpy
+import numpy as np
 import pandas
 import torch
 
@@ -90,9 +90,11 @@ class configClass:
 def create_default_config(project_name) -> str:
     return f"""
 def set_config(c):
-    c.input_path          = "data/{project_name}/{project_name}.pickle"
+    c.data_path          = "data/{project_name}/{project_name}_data.npy"
+    c.names_path         = "data/{project_name}/{project_name}_names.npy"
     c.compression_ratio   = 2.0
     c.epochs              = 5
+    c.energy_conversion = False
     c.early_stopping      = True
     c.lr_scheduler        = False
     c.patience            = 100
@@ -138,44 +140,38 @@ def numpy_to_tensor(data):
     return torch.from_numpy(data)
 
 
-def normalize(data, custom_norm, cleared_col_names):
-    data = numpy.apply_along_axis(
+def normalize(data, custom_norm):
+    data = np.apply_along_axis(
         data_processing.normalize, axis=0, arr=data, custom_norm=custom_norm
     )
-    df = data_processing.numpy_to_df(data, cleared_col_names)
-    return df
+    return data
 
 
-def process(data_path, custom_norm, test_size, energy_conversion):
-    df = data_processing.load_data(data_path)
-    cleared_col_names = data_processing.get_columns(df)
+def process(data_path, names_path, custom_norm, test_size, energy_conversion):
+    data = np.load(data_path)
+    names = np.load(names_path)
 
     if energy_conversion:
         print("Converting mass to energy with eta, pt & mass")
         df = convert_mass_to_energy(df, cleared_col_names)
 
-    full_pre_norm = df
-    normalization_features = data_processing.find_minmax(df)
-    df = normalize(df, custom_norm, cleared_col_names)
-    full_norm = df
+    normalization_features = data_processing.find_minmax(data)
+    data = normalize(data, custom_norm)
     if not test_size:
-        train_set = df
+        train_set = data
         test_set = train_set
     else:
         train_set, test_set = data_processing.split(
-            df, test_size=test_size, random_state=1
+            data, test_size=test_size, random_state=1
         )
+        number_of_columns = len(names)
 
-    number_of_columns = len(data_processing.get_columns(df))
-
+    # train_set, test_set = data_processing.split(df, test_size=test_size, random_state=1)
     return (
         train_set,
         test_set,
         number_of_columns,
         normalization_features,
-        full_norm,
-        full_pre_norm,
-        cleared_col_names,
     )
 
 
@@ -189,8 +185,8 @@ def train(model, number_of_columns, train_set, test_set, project_path, config):
     )
 
 
-def plot(project_path):
-    plotting.plot(project_path)
+def plot(project_path, config):
+    plotting.plot(project_path, config)
 
 
 def loss_plotter(path_to_loss_data, output_path, config):
@@ -207,16 +203,15 @@ def detach(tensor):
 
 def compress(model_path, config):
     # Give the encoding function the correct input as tensor
-    data = data_loader(config.input_path)
-    cleared_col_names = data_processing.get_columns(data)
-    number_of_columns = len(data_processing.get_columns(data))
+    data_before = np.load(config.data_path)
+    data = normalize(data_before, config.custom_norm)
+    cleared_col_names = np.load(config.names_path)
+    number_of_columns = len(cleared_col_names)
     try:
         config.latent_space_size = int(number_of_columns // config.compression_ratio)
         config.number_of_columns = number_of_columns
     except AttributeError:
         assert number_of_columns == config.number_of_columns
-    data_before = numpy.array(data)
-    data = normalize(data, config.custom_norm, cleared_col_names)
 
     # Initialise and load the model correctly.
     ModelObject = data_processing.initialise_model(config.model_name)
@@ -227,12 +222,6 @@ def compress(model_path, config):
         z_dim=config.latent_space_size,
     )
 
-    # Give the encoding function the correct input as tensor
-    data = data_loader(config.input_path)
-    # data = data_processing.clean_data(data, config)
-    data_before = numpy.array(data)
-
-    data = normalize(data, config.custom_norm, cleared_col_names)
     data_tensor = numpy_to_tensor(data).to(model.device)
 
     compressed = model.encode(data_tensor)
@@ -241,7 +230,7 @@ def compress(model_path, config):
 
 def decompress(model_path, input_path, model_name):
     # Load the data & convert to tensor
-    data = data_loader(input_path)
+    data = np.load(input_path)
     latent_space_size = len(data[0])
     modelDict = torch.load(str(model_path))
     number_of_columns = len(modelDict[list(modelDict.keys())[-1]])
@@ -256,7 +245,6 @@ def decompress(model_path, input_path, model_name):
     )
 
     # Load the data & convert to tensor
-    data = data_loader(input_path)
     data_tensor = numpy_to_tensor(data).to(model.device)
 
     decompressed = model.decode(data_tensor)
@@ -271,7 +259,7 @@ def to_root(data_path, cleared_col_names, save_path):
         return data_processing.df_to_root(
             data_path, col_names=data_path.columns(), save_path=save_path
         )
-    elif isinstance(data_path, numpy.ndarray):
+    elif isinstance(data_path, np.ndarray):
         df = data_processing.numpy_to_df(data_path, cleared_col_names)
         df_names = df.columns
         return data_processing.df_to_root(df, col_names=df_names, save_path=save_path)
@@ -290,8 +278,8 @@ def get_device():
 
 def compute_E(mass, eta, pt):
     masspt = pt**2 + mass**2
-    cosh = (numpy.cosh(eta)) ** 2
-    total = numpy.sqrt(masspt * cosh)
+    cosh = (np.cosh(eta)) ** 2
+    total = np.sqrt(masspt * cosh)
     return total
 
 
diff --git a/baler/modules/plotting.py b/baler/modules/plotting.py
index 55e56c20..0468de41 100644
--- a/baler/modules/plotting.py
+++ b/baler/modules/plotting.py
@@ -9,6 +9,8 @@
 import modules.data_processing as data_processing
 import modules.helper as helper
 
+import sys
+
 
 def to_percent(y, position):
     # Ignore the passed in position. This has the effect of scaling the default
@@ -23,11 +25,11 @@ def to_percent(y, position):
 
 
 def loss_plot(path_to_loss_data, output_path, config):
-    loss_data = pd.read_csv(path_to_loss_data)
+    loss_data = np.load(path_to_loss_data)
     str_list = ["Epochs:", "Model Name:", "Reg. Param:", "lr:", "BS:"]
 
-    val_loss = loss_data["Val Loss"]
-    train_loss = loss_data["Train Loss"]
+    train_loss = loss_data[0]
+    val_loss = loss_data[1]
     conf_list = [
         len(train_loss),
         config.model_name,
@@ -43,31 +45,23 @@ def loss_plot(path_to_loss_data, output_path, config):
     for i in range(len(conf_list)):
         plt.plot([], [], " ", label=str_list[i] + " " + str(conf_list[i]))
     plt.xlabel("Epochs")
+    plt.yscale("log")
     plt.ylabel("Loss")
     plt.legend(loc="best")
     plt.savefig(output_path + "_Loss_plot.pdf")
     # plt.show()
 
 
-def plot(project_path):
-    output_path = project_path + "plotting/"
-    before_path = project_path + "training/before.pickle"
-    after_path = project_path + "training/after.pickle"
-
-    with open(before_path, "rb") as handle:
-        before = pickle.load(handle)
-    with open(after_path, "rb") as handle:
-        after = pickle.load(handle)
+def plot(project_path, config):
+    output_path = project_path + "training/"
+    names_path = config.names_path
+    before_path = output_path + "before.npy"
+    after_path = output_path + "after.npy"
 
-    before = np.array(before)
-    # Added because plotting is not supported for non-DataFrame objects yet.
+    column_names = np.load(names_path)
 
-    column_names = helper.from_pickle(
-        project_path + "compressed_output/column_names.pickle"
-    )
-
-    before = pd.DataFrame(before, columns=column_names)
-    after = pd.DataFrame(after, columns=column_names)
+    before = pd.DataFrame(np.load(before_path), columns=column_names)
+    after = pd.DataFrame(np.load(after_path), columns=column_names)
 
     response = (after - before) / before
 
@@ -158,3 +152,6 @@ def plot(project_path):
             ax2.clear()
             ax1.clear()
             ax3.clear()
+
+            # if index == 3:
+            #    break
diff --git a/baler/modules/training.py b/baler/modules/training.py
index 87bc16b7..17254f40 100644
--- a/baler/modules/training.py
+++ b/baler/modules/training.py
@@ -7,8 +7,10 @@
 from torch.utils.data import DataLoader
 from tqdm import tqdm
 
-import modules.helper as helper
 import modules.utils as utils
+import modules.helper as helper
+
+import os
 
 
 def fit(model, train_dl, train_ds, model_children, regular_param, optimizer, RHO, l1):
@@ -19,13 +21,13 @@ def fit(model, train_dl, train_ds, model_children, regular_param, optimizer, RHO
     running_loss = 0.0
     counter = 0
     n_data = int(len(train_ds) / train_dl.batch_size)
+
     for inputs in tqdm(train_dl):
         counter += 1
         inputs = inputs.to(model.device)
-
         optimizer.zero_grad()
         reconstructions = model(inputs)
-        loss, mse_loss, l1_loss = utils.sparse_loss_function_EMD_L1(
+        loss, mse_loss, l1_loss = utils.sparse_loss_function_L1(
             model_children=model_children,
             true_data=inputs,
             reconstructed_data=reconstructions,
@@ -54,9 +56,8 @@ def validate(model, test_dl, test_ds, model_children, reg_param):
         for inputs in tqdm(test_dl):
             counter += 1
             inputs = inputs.to(model.device)
-
             reconstructions = model(inputs)
-            loss = utils.sparse_loss_function_EMD_L1(
+            loss = utils.sparse_loss_function_L1(
                 model_children=model_children,
                 true_data=inputs,
                 reconstructed_data=reconstructions,
@@ -100,8 +101,8 @@ def train(model, variables, train_data, test_data, parent_path, config):
     model = model.to(device)
 
     # Converting data to tensors
-    train_ds = torch.tensor(train_data.values, dtype=torch.float64, device=device)
-    valid_ds = torch.tensor(train_data.values, dtype=torch.float64, device=device)
+    train_ds = torch.tensor(train_data, dtype=torch.float64, device=device)
+    valid_ds = torch.tensor(test_data, dtype=torch.float64, device=device)
 
     # Pushing input data into the torch-DataLoader object and combines into one DataLoaders object (a basic wrapper
     # around several DataLoader objects).
@@ -166,15 +167,32 @@ def train(model, variables, train_data, test_data, parent_path, config):
             if early_stopping.early_stop:
                 break
 
+        ## Make-shift implementation to save models & values after 100 epochs:
+        save_model_and_data = True
+        if save_model_and_data:
+            if epoch % 100 == 0:
+                path = os.path.join(parent_path, f"model_{epoch}.pt")
+                path_data = os.path.join(parent_path, f"after_{epoch}.pickle")
+                path_pred = os.path.join(parent_path, f"before_{epoch}.pickle")
+
+                helper.model_saver(model, path)
+                data_tensor = torch.tensor(test_data, dtype=torch.float64).to(
+                    model.device
+                )
+                pred_tensor = model(data_tensor)
+                data = helper.detach(data_tensor)
+                pred = helper.detach(pred_tensor)
+
+                helper.to_pickle(data, path_data)
+                helper.to_pickle(pred, path_pred)
+
     end = time.time()
 
     print(f"{(end - start) / 60:.3} minutes")
-    pd.DataFrame({"Train Loss": train_loss, "Val Loss": val_loss}).to_csv(
-        parent_path + "loss_data.csv"
-    )
+    np.save(parent_path + "loss_data.npy", np.array([train_loss, val_loss]))
 
-    data_as_tensor = torch.tensor(test_data.values, dtype=torch.float64)
-    data_as_tensor = data_as_tensor.to(model.device)
+    data_as_tensor = torch.tensor(test_data, dtype=torch.float64)
+    data_as_tensor = data_as_tensor.to(trained_model.device)
     pred_as_tensor = trained_model(data_as_tensor)
 
     return data_as_tensor, pred_as_tensor, trained_model
diff --git a/baler/modules/utils.py b/baler/modules/utils.py
index 82d5ce26..6f552e38 100644
--- a/baler/modules/utils.py
+++ b/baler/modules/utils.py
@@ -2,6 +2,7 @@
 import torch.nn as nn
 from tqdm import tqdm
 from scipy.stats import wasserstein_distance
+from torch.nn import functional as F
 
 ###############################################
 factor = 0.5
@@ -45,7 +46,7 @@ def sparse_loss_function_L1(
     values = true_data
     if not validate:
         for i in range(len(model_children)):
-            values = model_children[i](values)
+            values = F.relu(model_children[i](values))
             l1_loss += torch.mean(torch.abs(values))
 
         loss = mse_loss + reg_param * l1_loss
@@ -88,15 +89,15 @@ def __init__(self, patience, min_delta):
         self.best_loss = None
         self.early_stop = False
 
-    def __call__(self, val_loss):
+    def __call__(self, train_loss):
         if self.best_loss is None:
-            self.best_loss = val_loss
+            self.best_loss = train_loss
 
-        elif self.best_loss - val_loss > self.min_delta:
-            self.best_loss = val_loss
+        elif self.best_loss - train_loss > self.min_delta:
+            self.best_loss = train_loss
             self.counter = 0  ## Resets if val_loss improves
 
-        elif self.best_loss - val_loss < self.min_delta:
+        elif self.best_loss - train_loss < self.min_delta:
             self.counter += 1
 
             print(f"Early stopping counter {self.counter} of {self.patience}")
@@ -122,5 +123,5 @@ def __init__(self, optimizer, patience, min_lr=min_lr, factor=factor):
             verbose=True,
         )
 
-    def __call__(self, val_loss):
-        self.lr_scheduler.step(val_loss)
+    def __call__(self, train_loss):
+        self.lr_scheduler.step(train_loss)
diff --git a/entrypoint.sh b/entrypoint.sh
new file mode 100755
index 00000000..ab2f64d4
--- /dev/null
+++ b/entrypoint.sh
@@ -0,0 +1,5 @@
+#!/bin/bash
+set -e
+
+eval $( fixuid -q )
+exec python baler "$@"
diff --git a/projects/example_CMS/example_CMS_config.py b/projects/example_CMS/example_CMS_config.py
index 7df469c3..629eb3bb 100644
--- a/projects/example_CMS/example_CMS_config.py
+++ b/projects/example_CMS/example_CMS_config.py
@@ -1,5 +1,7 @@
 def set_config(c):
-    c.input_path = "data/example_CMS/example_CMS.pickle"
+    c.data_path = "data/example_CMS/example_CMS_data.npy"
+    c.names_path = "data/example_CMS/example_CMS_names.npy"
+    c.energy_conversion = False
     c.compression_ratio = 2.0
     c.epochs = 10
     c.early_stopping = False
diff --git a/projects/example_CMS/example_CMS_preprocessing.py b/projects/example_CMS/example_CMS_preprocessing.py
index 0407112f..e9dcb34c 100644
--- a/projects/example_CMS/example_CMS_preprocessing.py
+++ b/projects/example_CMS/example_CMS_preprocessing.py
@@ -1,7 +1,12 @@
 import uproot
+import numpy as np
+import pandas as pd
 
+PROJECT_NAME = "example_CMS"
+DATA_PATH = "./data/example_CMS/example_CMS.root"
 
-def pre_processing(input_path, output_path):
+
+def pre_processing():
 
     Branch = "Events"
     Collection = "recoGenJets_slimmedGenJets__PAT."
@@ -19,7 +24,7 @@ def pre_processing(input_path, output_path):
     ]
 
     # Load data
-    tree = uproot.open(input_path)[Branch][Collection][Objects]
+    tree = uproot.open(DATA_PATH)[Branch][Collection][Objects]
     # Type clearing
     names = type_clearing(tree)
     df = tree.arrays(names, library="pd")
@@ -28,8 +33,11 @@ def pre_processing(input_path, output_path):
     df = df.reset_index(drop=True)
     df = df.dropna()
     global cleared_column_names
-    cleared_column_names = list(df)
-    df.to_pickle(output_path)
+    cleared_column_names = np.array(list(df))
+    arr = df.to_numpy()
+
+    np.save(f"./data/{PROJECT_NAME}/{PROJECT_NAME}_names.npy", cleared_column_names)
+    np.save(f"./data/{PROJECT_NAME}/{PROJECT_NAME}_data.npy", arr)
 
 
 def type_clearing(tt_tree):
@@ -54,7 +62,4 @@ def type_clearing(tt_tree):
     return column_names
 
 
-data_path_beofre_pre_processing = "./data/example_CMS/example_CMS.root"
-output_path = "./data/example_CMS/example_CMS.pickle"
-
-pre_processing(data_path_beofre_pre_processing, output_path)
+pre_processing()
diff --git a/projects/example_LHCb/example_LHCb_analysis.py b/projects/example_LHCb/example_LHCb_analysis.py
index 3769b393..71667cbf 100644
--- a/projects/example_LHCb/example_LHCb_analysis.py
+++ b/projects/example_LHCb/example_LHCb_analysis.py
@@ -6,93 +6,180 @@
 from matplotlib.backends.backend_pdf import PdfPages
 import scipy.optimize as opt
 
+
 def create_branches(df):
-    K_mass=0.494
-    df["H1_E"]=np.sqrt(df["H1_PX"]**2+df["H1_PY"]**2+df["H1_PZ"]**2+K_mass**2)
-    df["H2_E"]=np.sqrt(df["H2_PX"]**2+df["H2_PY"]**2+df["H2_PZ"]**2+K_mass**2)
-    df["H3_E"]=np.sqrt(df["H3_PX"]**2+df["H3_PY"]**2+df["H3_PZ"]**2+K_mass**2)
-    df["B_mass"]=np.sqrt((df["H1_E"]+df["H2_E"]+df["H3_E"])**2-(df["H1_PX"]+df["H2_PX"]+df["H3_PX"])**2-(df["H1_PY"]+df["H2_PY"]+df["H3_PY"])**2-(df["H1_PZ"]+df["H2_PZ"]+df["H3_PZ"])**2)
+    K_mass = 0.494
+    df["H1_E"] = np.sqrt(
+        df["H1_PX"] ** 2 + df["H1_PY"] ** 2 + df["H1_PZ"] ** 2 + K_mass**2
+    )
+    df["H2_E"] = np.sqrt(
+        df["H2_PX"] ** 2 + df["H2_PY"] ** 2 + df["H2_PZ"] ** 2 + K_mass**2
+    )
+    df["H3_E"] = np.sqrt(
+        df["H3_PX"] ** 2 + df["H3_PY"] ** 2 + df["H3_PZ"] ** 2 + K_mass**2
+    )
+    df["B_mass"] = np.sqrt(
+        (df["H1_E"] + df["H2_E"] + df["H3_E"]) ** 2
+        - (df["H1_PX"] + df["H2_PX"] + df["H3_PX"]) ** 2
+        - (df["H1_PY"] + df["H2_PY"] + df["H3_PY"]) ** 2
+        - (df["H1_PZ"] + df["H2_PZ"] + df["H3_PZ"]) ** 2
+    )
+
 
 def fit(x, a, b, c, k):
-     return a*np.exp(-((x-b)**2)/(2*c**2)) + k
-
-with open("./projects/example_LHCb/compressed_output/column_names.pickle", "rb") as handle:
-    column_names = pickle.load(handle)
-with open("./data/example_LHCb/example_LHCb.pickle", "rb") as handle:
-    before = pickle.load(handle)
-with open("./projects/example_LHCb/decompressed_output/decompressed.pickle", "rb") as handle:
-    after = pickle.load(handle)
-    
-df_before = pd.DataFrame(before, columns=column_names)
+    return a * np.exp(-((x - b) ** 2) / (2 * c**2)) + k
+
+
+column_names = np.load("./data/example_LHCb/example_LHCb_names.npy")
+
+df_before = pd.DataFrame(
+    np.load("./data/example_LHCb/example_LHCb_data.npy"), columns=column_names
+)
 create_branches(df_before)
 before = df_before["B_mass"]
-df_after = pd.DataFrame(after, columns=column_names)
+df_after = pd.DataFrame(
+    np.load("./projects/example_LHCb/decompressed_output/decompressed.npy"),
+    columns=column_names,
+)
 create_branches(df_after)
 after = df_after["B_mass"]
 
 
-fig, (ax1,ax2) = plt.subplots(2,1,figsize=(8*2.5*(1/2.54), 6*2.5*(1/2.54)), sharex=True)
+fig, (ax1, ax2) = plt.subplots(
+    2, 1, figsize=(8 * 2.5 * (1 / 2.54), 6 * 2.5 * (1 / 2.54)), sharex=True
+)
 
-x_min = min(before+after)
-x_max = max(before+after)
+x_min = min(before + after)
+x_max = max(before + after)
 x_diff = abs(x_max - x_min)
 
-with PdfPages("./projects/example_LHCb/plotting/analysis.pdf") as pdf:    
+with PdfPages("./projects/example_LHCb/plotting/analysis.pdf") as pdf:
 
     # Before Histogram
-    counts_before, bins_before = np.histogram(
-        before, bins=np.linspace(0,10000,100)
-    )
+    counts_before, bins_before = np.histogram(before, bins=np.linspace(0, 10000, 100))
     hist1 = ax1.hist(
-        bins_before[:-1], bins_before, weights=counts_before, label="Before", histtype="step", color="black"
+        bins_before[:-1],
+        bins_before,
+        weights=counts_before,
+        label="Before",
+        histtype="step",
+        color="black",
     )
-    before_bin_centers = hist1[1][:-1]+(hist1[1][1:]-hist1[1][:-1])/2
-    before_bin_centers_error = (hist1[1][1:]-hist1[1][:-1])/2
+    before_bin_centers = hist1[1][:-1] + (hist1[1][1:] - hist1[1][:-1]) / 2
+    before_bin_centers_error = (hist1[1][1:] - hist1[1][:-1]) / 2
     before_bin_counts = hist1[0]
     before_count_error = np.sqrt(hist1[0])
-    ax1.errorbar(before_bin_centers,before_bin_counts, yerr=before_count_error, xerr=None, marker="",linewidth=0.75, markersize=1, linestyle='', color="black")
-    optimizedParameters1, pcov1 = opt.curve_fit(fit, before_bin_centers, before_bin_counts, p0=[1,4500,1000,1]);
+    ax1.errorbar(
+        before_bin_centers,
+        before_bin_counts,
+        yerr=before_count_error,
+        xerr=None,
+        marker="",
+        linewidth=0.75,
+        markersize=1,
+        linestyle="",
+        color="black",
+    )
+    optimizedParameters1, pcov1 = opt.curve_fit(
+        fit, before_bin_centers, before_bin_counts, p0=[1, 4500, 1000, 1]
+    )
     perr1 = np.sqrt(np.diag(pcov1))
-    ax1.plot(before_bin_centers, fit(before_bin_centers, *optimizedParameters1),linewidth=1, label="Fit",color="red")
-    leg1 = ax1.legend(borderpad=0.5, loc=1, ncol=2, frameon=True,facecolor="white",framealpha=1,fontsize='medium')
+    ax1.plot(
+        before_bin_centers,
+        fit(before_bin_centers, *optimizedParameters1),
+        linewidth=1,
+        label="Fit",
+        color="red",
+    )
+    leg1 = ax1.legend(
+        borderpad=0.5,
+        loc=1,
+        ncol=2,
+        frameon=True,
+        facecolor="white",
+        framealpha=1,
+        fontsize="medium",
+    )
     leg1._legend_box.align = "left"
-    leg1.set_title(f"Mass  : {round(optimizedParameters1[1],2)} +/- {round(perr1[1],2)}" + f"\nWidth : {round(optimizedParameters1[2],2)} +/- {round(perr1[2],2)}")
-    ax1.set_ylabel("Counts", fontsize=14, ha='right', y=1.0)
-    ax1.set_xlabel("Mass [GeV]", fontsize=14, ha='right',x=1.0)
-    ax1.set_title('Before Compression')
-    ax1.set_ylim(0,20000)
+    leg1.set_title(
+        f"Mass  : {round(optimizedParameters1[1],2)} +/- {round(perr1[1],2)}"
+        + f"\nWidth : {round(optimizedParameters1[2],2)} +/- {round(perr1[2],2)}"
+    )
+    ax1.set_ylabel("Counts", fontsize=14, ha="right", y=1.0)
+    ax1.set_xlabel("Mass [GeV]", fontsize=14, ha="right", x=1.0)
+    ax1.set_title("Before Compression")
+    ax1.set_ylim(0, 20000)
     print(f"Before compression:")
     print(f"Mass  : {round(optimizedParameters1[1],2)} +/- {round(perr1[1],2)}")
     print(f"Width : {round(optimizedParameters1[2],2)} +/- {round(perr1[2],2)}")
 
     # After Histogram
-    counts_after, bins_after = np.histogram(
-        after, bins=np.linspace(0,10000,100)
-    )
+    counts_after, bins_after = np.histogram(after, bins=np.linspace(0, 10000, 100))
     hist2 = ax2.hist(
-        bins_after[:-1], bins_after, weights=counts_after, label="After", histtype="step", color="black"
+        bins_after[:-1],
+        bins_after,
+        weights=counts_after,
+        label="After",
+        histtype="step",
+        color="black",
     )
-    after_bin_centers = hist2[1][:-1]+(hist1[1][1:]-hist1[1][:-1])/2
-    after_bin_centers_error = (hist2[1][1:]-hist1[1][:-1])/2
+    after_bin_centers = hist2[1][:-1] + (hist1[1][1:] - hist1[1][:-1]) / 2
+    after_bin_centers_error = (hist2[1][1:] - hist1[1][:-1]) / 2
     after_bin_counts = hist2[0]
     after_count_error = np.sqrt(hist2[0])
-    ax2.errorbar(after_bin_centers,after_bin_counts, yerr=after_count_error, xerr=None, marker="",linewidth=0.75, markersize=1, linestyle='', color="black")
-    optimizedParameters2, pcov2 = opt.curve_fit(fit, after_bin_centers, after_bin_counts, p0=[1,4500,1000,1]);
+    ax2.errorbar(
+        after_bin_centers,
+        after_bin_counts,
+        yerr=after_count_error,
+        xerr=None,
+        marker="",
+        linewidth=0.75,
+        markersize=1,
+        linestyle="",
+        color="black",
+    )
+    optimizedParameters2, pcov2 = opt.curve_fit(
+        fit, after_bin_centers, after_bin_counts, p0=[1, 4500, 1000, 1]
+    )
     perr2 = np.sqrt(np.diag(pcov2))
-    ax2.plot(after_bin_centers, fit(after_bin_centers, *optimizedParameters2),linewidth=1, label="Fit",color="red");
-    leg2 = ax2.legend(borderpad=0.5, loc=1, ncol=2, frameon=True,facecolor="white",framealpha=1,fontsize='medium')
+    ax2.plot(
+        after_bin_centers,
+        fit(after_bin_centers, *optimizedParameters2),
+        linewidth=1,
+        label="Fit",
+        color="red",
+    )
+    leg2 = ax2.legend(
+        borderpad=0.5,
+        loc=1,
+        ncol=2,
+        frameon=True,
+        facecolor="white",
+        framealpha=1,
+        fontsize="medium",
+    )
     leg2._legend_box.align = "left"
-    leg2.set_title(f"Mass  : {round(optimizedParameters2[1],2)} +/- {round(perr2[1],2)}" + f"\nWidth : {round(optimizedParameters2[2],2)} +/- {round(perr2[2],2)}")
-    ax2.set_ylabel("Counts", fontsize=14, ha='right', y=1.0)
-    ax2.set_xlabel("Mass [GeV]", fontsize=14, ha='right',x=1.0)
-    ax2.set_title('After Decompression')
-    ax2.set_ylim(0,20000)
+    leg2.set_title(
+        f"Mass  : {round(optimizedParameters2[1],2)} +/- {round(perr2[1],2)}"
+        + f"\nWidth : {round(optimizedParameters2[2],2)} +/- {round(perr2[2],2)}"
+    )
+    ax2.set_ylabel("Counts", fontsize=14, ha="right", y=1.0)
+    ax2.set_xlabel("Mass [GeV]", fontsize=14, ha="right", x=1.0)
+    ax2.set_title("After Decompression")
+    ax2.set_ylim(0, 20000)
 
     print(f"After compression:")
     print(f"Mass  : {round(optimizedParameters2[1],2)} +/- {round(perr2[1],2)}")
     print(f"Width : {round(optimizedParameters2[2],2)} +/- {round(perr2[2],2)}")
 
-    diff = round((abs(optimizedParameters1[1]-optimizedParameters2[1])/optimizedParameters2[1])*100,1)
+    diff = round(
+        (
+            abs(optimizedParameters1[1] - optimizedParameters2[1])
+            / optimizedParameters2[1]
+        )
+        * 100,
+        1,
+    )
     fig.suptitle(f"Relative Mass Difference = {diff} %", fontsize=16)
 
     pdf.savefig()
diff --git a/projects/example_LHCb/example_LHCb_config.py b/projects/example_LHCb/example_LHCb_config.py
index 137a9cde..11637a28 100644
--- a/projects/example_LHCb/example_LHCb_config.py
+++ b/projects/example_LHCb/example_LHCb_config.py
@@ -1,18 +1,19 @@
-
 def set_config(c):
-    c.input_path          = "data/example_LHCb/example_LHCb.pickle"
-    c.compression_ratio   = 2.0
-    c.epochs              = 200
-    c.early_stopping      = True
-    c.lr_scheduler        = False
-    c.patience            = 100
-    c.min_delta           = 0
-    c.model_name          = "george_SAE"
-    c.custom_norm         = False
-    c.l1                  = True
-    c.reg_param             = 0.001
-    c.RHO                 = 0.05
-    c.lr                  = 0.001
-    c.batch_size          = 512
-    c.save_as_root        = True
-    c.test_size           = 0.15
+    c.data_path = "data/example_LHCb/example_LHCb_data.npy"
+    c.names_path = "data/example_LHCb/example_LHCb_names.npy"
+    c.compression_ratio = 2.0
+    c.energy_conversion = False
+    c.epochs = 20
+    c.early_stopping = True
+    c.lr_scheduler = False
+    c.patience = 100
+    c.min_delta = 0
+    c.model_name = "george_SAE"
+    c.custom_norm = False
+    c.l1 = True
+    c.reg_param = 0.001
+    c.RHO = 0.05
+    c.lr = 0.001
+    c.batch_size = 512
+    c.save_as_root = True
+    c.test_size = 0.15
diff --git a/projects/example_LHCb/example_LHCb_preprocessing.py b/projects/example_LHCb/example_LHCb_preprocessing.py
index a6bc1a1d..05ab91b4 100644
--- a/projects/example_LHCb/example_LHCb_preprocessing.py
+++ b/projects/example_LHCb/example_LHCb_preprocessing.py
@@ -1,11 +1,16 @@
 import uproot
+import numpy as np
 
-#Open pre-Baler data as pandas df
-tuple_pre=uproot.open("./data/example_LHCb/example_LHCb.root")["DecayTree"]
-df=tuple_pre.arrays(tuple_pre.keys(), library='pd')
+# Open pre-Baler data as pandas df
+infile = uproot.open("./data/example_LHCb/example_LHCb.root")["DecayTree"]
+df = infile.arrays(infile.keys(), library="pd")
 
-#Perform PID cuts so we are looking at B->KKK
-df_pre=df[(df.H1_ProbK>0.5) & (df.H2_ProbK>0.5) & (df.H3_ProbK>0.5)].copy()
+# Perform PID cuts so we are looking at B->KKK
+df = df[(df.H1_ProbK > 0.5) & (df.H2_ProbK > 0.5) & (df.H3_ProbK > 0.5)].copy()
 
-#Export as pkl file for input to Baler
-df_pre.to_pickle("./data/example_LHCb/example_LHCb.pickle")
+# Export as pkl file for input to Baler
+names = np.array(list(df.columns))
+data = df.to_numpy()
+
+np.save("./data/example_LHCb/example_LHCb_names.npy", names)
+np.save("./data/example_LHCb/example_LHCb_data.npy", data)
diff --git a/projects/good_plots/compressed_output/.gitkeep b/projects/good_plots/compressed_output/.gitkeep
new file mode 100644
index 00000000..e69de29b
diff --git a/projects/good_plots/configClass.py b/projects/good_plots/configClass.py
new file mode 100644
index 00000000..f40704d0
--- /dev/null
+++ b/projects/good_plots/configClass.py
@@ -0,0 +1,44 @@
+class Configuration(object):
+    def __init__(self):
+        self.input_path = "./data/good_plots/george.pickle"
+
+        self.epochs = 10
+        self.early_stopping = True
+        self.lr_scheduler = False
+        self.patience = 100
+        self.min_delta = 0
+        self.model_name = "george_SAE"
+        self.custom_norm = False
+        self.l1 = True
+        self.reg_param = 0.001
+        self.RHO = 0.05
+        self.lr = 0.001
+        self.batch_size = 512
+        self.save_as_root = True
+        self.cleared_col_names = [
+            "pt",
+            "eta",
+            "phi",
+            "m",
+            "EmEnergy",
+            "HadEnergy",
+            "InvisEnergy",
+            "AuxilEnergy",
+        ]
+        self.test_size = 0.15
+        self.Branch = "Events"
+        self.Collection = "recoGenJets_slimmedGenJets__PAT."
+        self.Objects = "recoGenJets_slimmedGenJets__PAT.obj"
+        self.number_of_columns = 24
+        self.latent_space_size = 15
+        self.dropped_variables = [
+            "recoGenJets_slimmedGenJets__PAT.obj.m_state.vertex_.fCoordinates.fX",
+            "recoGenJets_slimmedGenJets__PAT.obj.m_state.vertex_.fCoordinates.fY",
+            "recoGenJets_slimmedGenJets__PAT.obj.m_state.vertex_.fCoordinates.fZ",
+            "recoGenJets_slimmedGenJets__PAT.obj.m_state.qx3_",
+            "recoGenJets_slimmedGenJets__PAT.obj.m_state.pdgId_",
+            "recoGenJets_slimmedGenJets__PAT.obj.m_state.status_",
+            "recoGenJets_slimmedGenJets__PAT.obj.mJetArea",
+            "recoGenJets_slimmedGenJets__PAT.obj.mPileupEnergy",
+            "recoGenJets_slimmedGenJets__PAT.obj.mPassNumber",
+        ]
diff --git a/projects/good_plots/decompressed_output/.gitkeep b/projects/good_plots/decompressed_output/.gitkeep
new file mode 100644
index 00000000..e69de29b
diff --git a/projects/good_plots/model/.gitkeep b/projects/good_plots/model/.gitkeep
new file mode 100644
index 00000000..e69de29b
diff --git a/projects/good_plots/plotting/.gitkeep b/projects/good_plots/plotting/.gitkeep
new file mode 100644
index 00000000..e69de29b
diff --git a/projects/good_plots/training/.gitkeep b/projects/good_plots/training/.gitkeep
new file mode 100644
index 00000000..e69de29b


Test patch:
diff --git a/tests/test_data_processing.py b/tests/test_data_processing.py
index ddf52802..2aa30fc5 100644
--- a/tests/test_data_processing.py
+++ b/tests/test_data_processing.py
@@ -41,25 +41,19 @@ def test_save_model():
 @pytest.fixture
 def minmax_test_data():
     return [
+        (np.array([[1, 2, 3], [4, 5, 6], [7, 8, 9]]), np.array([[1, 2, 3], [6, 6, 6]])),
         (
-            [[1, 2, 3], [4, 5, 6], [7, 8, 9]],
-            pd.DataFrame({"True min": [1, 2, 3], "Feature Range": [6, 6, 6]}),
-        ),
-        (
-            [[-1, -2, -3], [-4, -5, -6], [-7, -8, -9]],
-            pd.DataFrame({"True min": [-7, -8, -9], "Feature Range": [6, 6, 6]}),
-        ),
-        (
-            [[0, 0, 0], [1, 1, 1], [2, 2, 2]],
-            pd.DataFrame({"True min": [0, 0, 0], "Feature Range": [2, 2, 2]}),
+            np.array([[-1, -2, -3], [-4, -5, -6], [-7, -8, -9]]),
+            np.array([[-7, -8, -9], [6, 6, 6]]),
         ),
+        (np.array([[0, 0, 0], [1, 1, 1], [2, 2, 2]]), np.array([[0, 0, 0], [2, 2, 2]])),
     ]
 
 
 def test_find_minmax_success(minmax_test_data):
     for data, expected_result in minmax_test_data:
         result = data_processing.find_minmax(data)
-        assert result.equals(expected_result)
+        assert np.array_equal(result, expected_result)
 
 
 def test_normalize():