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

Issue:
Adding cosmetic changes to make plotting terminal output
Added tqdm to plotting terminal output
Overall code cleaning done
Overall code cleaning done: mostly removing redundancies and minimal changes of variable names for logical consistency


Discussion / hints:

Fix patch:
diff --git a/.github/workflows/docker.yaml b/.github/workflows/docker.yaml
index 2cccb0b4..5dd115af 100644
--- a/.github/workflows/docker.yaml
+++ b/.github/workflows/docker.yaml
@@ -1,13 +1,16 @@
 name: Docker-CI
 on: [pull_request, workflow_dispatch]
+branches:
+  - main
+  - development
 
 jobs:
   build-docker-image:
     runs-on: ubuntu-latest
 
     steps:
-      - uses: actions/checkout@v3
-      - name: Build the Docker image
-        run: |
-          ./bin/build_docker.sh
-          ./bin/test_docker.sh
+    - uses: actions/checkout@v3
+    - name: Build the Docker image
+      run: |
+        ./bin/build_docker.sh
+        ./bin/test_docker.sh
diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
new file mode 100644
index 00000000..8c682da5
--- /dev/null
+++ b/.pre-commit-config.yaml
@@ -0,0 +1,14 @@
+# See https://pre-commit.com for more information
+# See https://pre-commit.com/hooks.html for more hooks
+repos:
+-   repo: https://github.com/pre-commit/pre-commit-hooks
+    rev: v3.2.0
+    hooks:
+    -   id: trailing-whitespace
+    -   id: end-of-file-fixer
+    -   id: check-yaml
+    -   id: check-added-large-files
+-   repo: https://github.com/psf/black
+    rev: 22.10.0
+    hooks:
+    -   id: black
diff --git a/Dockerfile b/Dockerfile
index b021a3ec..de1b4602 100644
--- a/Dockerfile
+++ b/Dockerfile
@@ -71,7 +71,7 @@ COPY --from=python-base /baler-root/baler/dist/*.whl ./
 # Install wheel
 RUN pip install *.whl
 
-# Copy source 
+# Copy source
 COPY --from=python-base /baler-root/baler/modules/ ./modules
 COPY --from=python-base /baler-root/baler/*.py /baler-root/baler/README.md ./
 
diff --git a/Dockerfile.arm64 b/Dockerfile.arm64
index d7e863a2..2c347e5e 100644
--- a/Dockerfile.arm64
+++ b/Dockerfile.arm64
@@ -72,7 +72,7 @@ COPY --from=python-base /baler-root/baler/dist/*.whl ./
 # Install wheel
 RUN pip install *.whl
 
-# Copy source 
+# Copy source
 COPY --from=python-base /baler-root/baler/modules/ ./modules
 COPY --from=python-base /baler-root/baler/*.py /baler-root/baler/README.md ./
 
diff --git a/baler/baler.py b/baler/baler.py
index 47e67859..d2a04faf 100644
--- a/baler/baler.py
+++ b/baler/baler.py
@@ -7,49 +7,77 @@
 
 
 def main():
+    """Calls different functions depending on argument parsed in command line.
+
+        - if --mode=newProject: call `helper.create_new_project` and create a new project sub directory with config file
+        - if --mode=train: call `perform_training` and train the network on given data and based on the config file
+        - if --mode=compress: call `perform_compression` and compress the given data using the model trained in `--mode=train`
+        - if --mode=decompress: call `perform_decompression` and decompress the compressed file outputted from `--mode=compress`
+        - if --mode=plot: call `perform_plotting` and plot the comparison between the original data and the decompressed data from `--mode=decompress`. Also plots the loss plot from the trained network.
+
+
+    Raises:
+        NameError: Raises error if the chosen mode does not exist.
+    """
     config, mode, project_name = helper.get_arguments()
     project_path = f"projects/{project_name}/"
     if mode == "newProject":
         helper.create_new_project(project_name)
     elif mode == "train":
-        perform_training(config, project_path)
-    elif mode == "plot":
-        perform_plotting(project_path, config)
+        perform_training(project_path, config)
     elif mode == "compress":
-        perform_compression(config, project_path)
+        perform_compression(project_path, config)
     elif mode == "decompress":
-        perform_decompression(
-            config.save_as_root, config.model_name, project_path, config
-        )
+        perform_decompression(config.model_name, project_path, config)
+    elif mode == "plot":
+        perform_plotting(project_path, config)
     elif mode == "info":
-        print_info(project_path)
+        print_info(project_path, config)
+    else:
+        raise NameError(
+            "Baler mode "
+            + mode
+            + " not recognised. Use baler --help to see available modes."
+        )
+
 
+def perform_training(project_path, config):
+    """Main function calling the training functions, ran when --mode=train is selected.
+        The three main functions this calls are: `helper.process`, `helper.mode_init` and `helper.training`.
 
-def perform_training(config, project_path):
-    (train_set_norm, test_set_norm, normalization_features,) = helper.process(
+        Depending on `config.data_dimensions`, the calculated latent space size will differ.
+
+    Args:
+        project_path (string): Selects base path for determining output path
+        config (dataClass): Base class selecting user inputs
+
+    Raises:
+        NameError: Baler currently only supports 1D (e.g. HEP) or 2D (e.g. CFD) data as inputs.
+    """
+    train_set_norm, test_set_norm, normalization_features = helper.process(
         config.input_path,
         config.custom_norm,
         config.test_size,
-        config.energy_conversion,
         config.apply_normalization,
     )
 
     try:
         if config.data_dimension == 1:
-            number_of_columns = len(train_set_norm[0])
+            number_of_columns = train_set_norm.shape[1]
             config.latent_space_size = int(
                 number_of_columns // config.compression_ratio
             )
             config.number_of_columns = number_of_columns
         elif config.data_dimension == 2:
-            number_of_columns = len(train_set_norm[0])
+            number_of_rows = train_set_norm.shape[1]
+            number_of_columns = train_set_norm.shape[2]
             config.latent_space_size = int(
-                (number_of_columns * number_of_columns) // config.compression_ratio
+                (number_of_rows * number_of_columns) // config.compression_ratio
             )
             config.number_of_columns = number_of_columns
         else:
             raise NameError(
-                "Data dimension can only be 1 or 2. Introduced value = "
+                "Data dimension can only be 1 or 2. Got config.data_dimension value = "
                 + str(config.data_dimension)
             )
     except AttributeError:
@@ -59,9 +87,8 @@ def perform_training(config, project_path):
     device = helper.get_device()
 
     model_object = helper.model_init(config.model_name)
-    model = model_object(
-        device=device, n_features=number_of_columns, z_dim=config.latent_space_size
-    )
+    model = model_object(n_features=number_of_columns, z_dim=config.latent_space_size)
+    model.to(device)
 
     output_path = project_path + "training/"
     trained_model = helper.train(
@@ -77,28 +104,50 @@ def perform_training(config, project_path):
 
 
 def perform_plotting(project_path, config):
+    """Main function calling the two plotting functions, ran when --mode=plot is selected.
+       The two main functions this calls are: `helper.plotter` and `helper.loss_plotter`
+
+    Args:
+        project_path (string): Selects base path for determining output path
+        config (dataClass): Base class selecting user inputs
+    """
     output_path = project_path + "plotting/"
-    helper.plot(project_path, config)
     helper.loss_plotter(project_path + "training/loss_data.npy", output_path, config)
+    helper.plotter(project_path, config)
+
 
+def perform_compression(project_path, config):
+    """Main function calling the compression functions, ran when --mode=compress is selected.
+       The main function being called here is: `helper.compress`
 
-def perform_compression(config, project_path):
+        If `config.extra_compression` is selected, the compressed file is further compressed via zip
+        Else, the function returns a compressed file of `.npz`, only compressed by Baler.
+
+    Args:
+        project_path (string): Selects base path for determining output path
+        config (dataClass): Base class selecting user inputs
+
+    Outputs:
+        An `.npz` file which includes:
+        - The compressed data
+        - The data headers
+        - Normalization features if `config.apply_normalization=True`
+    """
     print("Compressing...")
     start = time.time()
+    normalization_features = []
 
     if config.apply_normalization:
         normalization_features = np.load(
             project_path + "training/normalization_features.npy"
         )
-    else:
-        normalization_features = []
 
     compressed = helper.compress(
         model_path=project_path + "compressed_output/model.pt",
         config=config,
     )
     # Converting back to numpyarray
-    compressed = helper.detach(compressed)
+    compressed = helper.detacher(compressed)
     end = time.time()
 
     print("Compression took:", f"{(end - start) / 60:.3} minutes")
@@ -121,7 +170,17 @@ def perform_compression(config, project_path):
         )
 
 
-def perform_decompression(save_as_root, model_name, project_path, config):
+def perform_decompression(model_name, project_path, config):
+    """Main function calling the decompression functions, ran when --mode=decompress is selected.
+       The main function being called here is: `helper.decompress`
+
+        If `config.apply_normalization=True` the output is un-normalized with the same normalization features saved from `perform_training()`.
+
+    Args:
+        model_name (string): Name of the model you want to use for decompression
+        project_path (string): Selects base path for determining output path
+        config (dataClass): Base class selecting user inputs
+    """
     print("Decompressing...")
 
     start = time.time()
@@ -132,7 +191,7 @@ def perform_decompression(save_as_root, model_name, project_path, config):
     )
 
     # Converting back to numpyarray
-    decompressed = helper.detach(decompressed)
+    decompressed = helper.detacher(decompressed)
 
     if config.apply_normalization:
         print("Un-normalizing...")
@@ -161,14 +220,15 @@ def perform_decompression(save_as_root, model_name, project_path, config):
         )
 
 
-def print_info(project_path):
+def print_info(project_path, config):
     print(
         "================================== \n Information about your compression \n================================== "
     )
 
-    pre_compression = project_path + "compressed_output/cleandata_pre_comp.pickle"
-    compressed = project_path + "compressed_output/compressed.pickle"
-    decompressed = project_path + "decompressed_output/decompressed.pickle"
+    pre_compression = config.input_path
+    model = project_path + "compressed_output/model.pt"
+    compressed = project_path + "compressed_output/compressed.npz"
+    decompressed = project_path + "decompressed_output/decompressed.npz"
 
     files = [pre_compression, compressed, decompressed]
     q = []
@@ -180,7 +240,7 @@ def print_info(project_path):
     )
     print(f"File size before compression: {round(q[0], 2)} MB")
     print(f"Compressed file size: {round(q[1], 2)} MB")
-    print(f"De-compressed file size: {round(q[2], 2),} MB")
+    print(f"De-compressed file size: {round(q[2], 2)} MB")
     print(f"Compression ratio: {round(q[0] / q[1], 2)}")
 
 
diff --git a/baler/modules/data_processing.py b/baler/modules/data_processing.py
index 58312879..2123235f 100644
--- a/baler/modules/data_processing.py
+++ b/baler/modules/data_processing.py
@@ -7,17 +7,46 @@
 
 
 def save_model(model, model_path: str) -> None:
+    """Saves the models state dictionary as a `.pt` file to the given path.
+
+    Args:
+        model (nn.Module): The PyTorch model to save.
+        model_path (str): String defining the models save path.
+
+    Returns:
+        None: Saved model state dictionary as `.pt` file.
+    """
     return torch.save(model.state_dict(), model_path)
 
 
 def initialise_model(model_name):
+    """Initializing the models attributes to a model_object variable.
+
+    Args:
+        model_name (function): The name of the model you wish to initialize. This should correspond to what your Model name.
+
+    Returns:
+        class: Object with the models class attributes
+    """
     model_object = getattr(models, model_name)
     return model_object
 
 
 def load_model(model_object, model_path, n_features, z_dim):
+    """Loads the state dictionary of the trained model into a model variable. This variable is then used for passing data through the encoding and decoding functions.
+
+    Args:
+        model_object (object): Object with the models attributels
+        model_path (str): Path to model
+        n_features (int): Input dimension size
+        z_dim (int): Latent space size
+
+    Returns:
+        nn.Module: Returns a model object with the attributes of the model class, with the a selected state dictionary loaded into it.
+    """
     device = helper.get_device()
-    model = model_object(device, n_features, z_dim)
+    model = model_object(n_features, z_dim)
+    model.to(device)
 
     # Loading the state_dict into the model
     model.load_state_dict(torch.load(str(model_path)), strict=False)
@@ -25,10 +54,19 @@ def load_model(model_object, model_path, n_features, z_dim):
 
 
 def find_minmax(data):
+    """Obains the minimum and maximum values for each column.
+
+    Args:
+        data (ndarray): Any dataset as a `ndarray` which one eventually wants to normalize using the Min-Max method.
+
+    Returns:
+        ndarray: An array of two lists. One of the lists contains the minimum of each column, while the other list contains `feature_range = max - min` for each column.
+    """
     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)
 
+    # Computes the range
     feature_range_list = true_max_list - true_min_list
 
     normalization_features = np.array([true_min_list, feature_range_list])
@@ -36,6 +74,16 @@ def find_minmax(data):
 
 
 def normalize(data, custom_norm):
+    """This function scales the data to be in the range [0,1], based on the Min Max normalization method. It finds the minimum and maximum values of each column
+        and computes the values according to: x_norm = (x - x_min) / (x_max - x_min).
+
+    Args:
+        data (ndarray): A dataset of type `ndarray`.
+        custom_norm (boolean): If you want to do Min Max normalization or any custom normalization. Custom normalization is not supported at the moment.
+
+    Returns:
+        ndarray: If not custom_norm: Input data where every column is scaled to be in the range [0,1]. Otherwise, the input data is returned
+    """
     data = np.array(data)
     if custom_norm:
         pass
@@ -48,11 +96,33 @@ def normalize(data, custom_norm):
     return data
 
 
-def split(df, test_size, random_state):
-    return train_test_split(df, test_size=test_size, random_state=random_state)
+def split(data, test_size, random_state):
+    """Splits the given data according to a test size into two sets, train and test. This is sklearns function, opted from:
+    https://scikit-learn.org/stable/modules/generated/sklearn.model_selection.train_test_split.html
+
+    Args:
+        data (ndarray): Dataset which you want to split
+        test_size (float): Value between 0.0 and 1.0 determining how large the test size will be.
+                           Example: test_size = 0.15 randomly selectes 15% of the original dataset and puts it into a new array.
+        random_state (int): Shuffling control. Passing an int here shuffels the data the same way every time. Used for reproducibility.
+
+    Returns:
+        ndarray, ndarray: Two datasets splitted in size according to the `test_size` value.
+    """
+    return train_test_split(data, test_size=test_size, random_state=random_state)
 
 
 def renormalize_std(data, true_min, feature_range):
+    """Computes the un-normalization of normalized arrays. This function is used in combination with `renormalize_func` to perform this operation on an entire dataset.
+
+    Args:
+        data (ndarray): Normalized data you want to un-normalize from the range [0,1]
+        true_min (float): A columns minimum
+        feature_range (float): A columns feature range, computed as `max - min`
+
+    Returns:
+        ndarray: Un-normalized data array
+    """
     data = list(data)
     data = [((i * feature_range) + true_min) for i in data]
     data = np.array(data)
@@ -60,6 +130,17 @@ def renormalize_std(data, true_min, feature_range):
 
 
 def renormalize_func(norm_data, min_list, range_list):
+    """Un-normalizes an entire dataset. Applies the `renormalize_std`function across an entire dataset.
+        `min_list` and `range_list` are obtained from the `normalization_features.npy` file.
+
+    Args:
+        norm_data (ndarray): Normalized data which one wants to un-normalized.
+        min_list (list): List of minimum values. Should have the same length as the amount of columns.
+        range_list (list): List of feature range values. Should have the same length as the amount of columns.
+
+    Returns:
+        ndarray: Array with the un-normalized values.
+    """
     norm_data = np.array(norm_data)
     renormalized = [
         renormalize_std(norm_data, min_list[i], range_list[i])
diff --git a/baler/modules/helper.py b/baler/modules/helper.py
index cfa3e278..b17c5d0d 100644
--- a/baler/modules/helper.py
+++ b/baler/modules/helper.py
@@ -1,9 +1,8 @@
 import argparse
 import importlib
 import os
-import pickle
-from dataclasses import dataclass
 import sys
+from dataclasses import dataclass
 
 sys.path.append(os.getcwd())
 import numpy as np
@@ -14,6 +13,11 @@
 
 
 def get_arguments():
+    """Determines the arguments one is able to apply in the command line when running Baler. Use `--help` to see what options are avaliable.
+
+    Returns:
+        .py, string, folder: `.py` file containing the config options, string determining what mode to run, projects directory where outputs go.
+    """
     parser = argparse.ArgumentParser(
         prog="baler.py",
         description=(
@@ -53,6 +57,12 @@ def get_arguments():
 
 
 def create_new_project(project_name: str, base_path: str = "projects") -> None:
+    """Creates a new project directory with all necessary sub-directories and config files.
+
+    Args:
+        project_name (str): Determines what you want to call your new project as.
+        base_path (str, optional): Defaults to "projects"
+    """
     project_path = os.path.join(base_path, project_name)
     if os.path.exists(project_path):
         print(f"The project {project_path} already exists.")
@@ -79,8 +89,9 @@ class Config:
     compression_ratio: float
     epochs: int
     early_stopping: bool
+    early_stoppin_patience: int
     lr_scheduler: bool
-    patience: int
+    lr_scheduler_patience: int
     min_delta: int
     model_name: str
     custom_norm: bool
@@ -89,75 +100,101 @@ class Config:
     RHO: float
     lr: float
     batch_size: int
-    save_as_root: bool
     test_size: float
-    energy_conversion: bool
     data_dimension: int
+    intermittent_model_saving: bool
+    intermittent_saving_patience: int
 
 
 def create_default_config(project_name: str) -> str:
     return f"""
 def set_config(c):
-    c.input_path          = "data/{project_name}/{project_name}_data.npz"
-    c.compression_ratio   = 2.0
-    c.epochs              = 5
-    c.energy_conversion = False
-    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.energy_conversion   = False
-    c.data_dimension      = 1
-    c.apply_normalization = True
+    c.input_path                   = "data/{project_name}/{project_name}_data.npz"
+    c.compression_ratio            = 2.0
+    # c.number_of_columns            = 24
+    # c.latent_space_size            = 12
+    c.epochs                       = 5
+    c.early_stopping               = True
+    c.early_stopping_patience      = 100
+    c.min_delta                    = 0
+    c.lr_scheduler                 = False
+    c.lr_scheduler_patience        = 100
+    c.model_name                   = "AE"
+    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.test_size                    = 0
+    c.data_dimension               = 1
+    c.apply_normalization          = True
+    c.extra_compression            = False
+    c.intermittent_model_saving    = False
+    c.intermittent_saving_patience = 100
 
 """
 
 
-def to_pickle(data, path):
-    with open(path, "wb") as handle:
-        pickle.dump(data, handle)
-
-
-def from_pickle(path):
-    with open(path, "rb") as handle:
-        return pickle.load(handle)
+def model_init(model_name):
+    """Calls `data_processing.initialise_model`.
 
+    Args:
+        model_name (string): The name of the model you wish to initialize
 
-def model_init(model_name):
+    Returns:
+       nn.Module : The initialized model
+    """
     # This is used when we don't have saved model parameters.
     model_object = data_processing.initialise_model(model_name)
     return model_object
 
 
 def numpy_to_tensor(data):
+    """Converts ndarrays to torch.Tensors.
+
+    Args:
+        data (ndarray): The data you wish to convert from ndarray to torch.Tensor.
+
+    Returns:
+        torch.Tensor: Your data as a tensor
+    """
     return torch.from_numpy(data)
 
 
 def normalize(data, custom_norm):
+    """Applies `data_processing.normalize()` along every axis of given data
+
+    Args:
+        data (ndarray): Data you wish to normalize
+        custom_norm (boolean): Wether or not you wish to use MinMax normalization
+
+    Returns:
+        ndarray: Normalized data
+    """
     data = np.apply_along_axis(
         data_processing.normalize, axis=0, arr=data, custom_norm=custom_norm
     )
     return data
 
 
-def process(input_path, custom_norm, test_size, energy_conversion, apply_normalization):
+def process(input_path, custom_norm, test_size, apply_normalization):
+    """Loads the input data into an ndarray, splits it into train/test splits and normalizes if chosen.
+
+    Args:
+        input_path (string): Path to input data
+        custom_norm (boolean): `False` if you want to use MinMax normalization. `True` otherwise.
+        test_size (float): How much of your data to use as validation.
+        apply_normalization (boolean): `True` if you want to normalize. `False` if you don't want to normalize.
+
+    Returns:
+        ndarray, ndarray, ndarray: Array with the train set, array with the test set and array with the normalization features.
+    """
     loaded = np.load(input_path)
     data = loaded["data"]
     names = loaded["names"]
+    normalization_features = []
 
-    # TODO Fix this
-    # if energy_conversion:
-    #     print("Converting mass to energy with eta, pt & mass")
-    #     df = convert_mass_to_energy(df, cleared_col_names)
     normalization_features = data_processing.find_minmax(data)
     if apply_normalization:
         print("Normalizing the data...")
@@ -178,33 +215,119 @@ def process(input_path, custom_norm, test_size, energy_conversion, apply_normali
 
 
 def renormalize(data, true_min_list, feature_range_list):
+    """Calls `data_processing.renormalize_func()`.
+
+    Args:
+        data (ndarray): Data you wish to un-normalize
+        true_min_list (ndarray): List of column minimums
+        feature_range_list (ndarray): List of column feature ranges
+
+    Returns:
+        ndarray: Un-normalized array
+    """
     return data_processing.renormalize_func(data, true_min_list, feature_range_list)
 
 
 def train(model, number_of_columns, train_set, test_set, project_path, config):
+    """Calls `training.train()`
+
+    Args:
+        model (modelObject): The model you wish to train
+        number_of_columns (int): Amount of columns in the initial dataset
+        train_set (ndarray): Array consisting of the train set
+        test_set (ndarray): Array consisting of the test set
+        project_path (string): Path to the project directory
+        config (dataClass): Base class selecting user inputs
+
+    Returns:
+        _type_: _description_
+    """
     return training.train(
         model, number_of_columns, train_set, test_set, project_path, config
     )
 
 
-def plot(project_path, config):
+def plotter(project_path, config):
+    """Calls `plotting.plot()`
+
+    Args:
+        project_path (string): Path to the project directory
+        config (dataClass): Base class selecting user inputs
+
+    """
+
     plotting.plot(project_path, config)
+    print("### Done ###")
+    print("Your plots are available in:", project_path + "plotting/")
 
 
 def loss_plotter(path_to_loss_data, output_path, config):
+    """Calls `plotting.loss_plot()`
+
+    Args:
+        path_to_loss_data (string): Path to the values for the loss plot
+        output_path (string): Path to output the data
+        config (dataClass): Base class selecting user inputs
+
+    Returns:
+        .pdf file: Plot containing the loss curves
+    """
     return plotting.loss_plot(path_to_loss_data, output_path, config)
 
 
 def model_saver(model, model_path):
+    """Calls `data_processing.save_model()`
+
+    Args:
+        model (nn.Module): The PyTorch model to save.
+        model_path (str): String defining the models save path.
+
+    Returns:
+        .pt file: `.pt` File containing the model state dictionary
+    """
     return data_processing.save_model(model, model_path)
 
 
-def detach(tensor):
+def detacher(tensor):
+    """Detaches a given tensor to ndarray
+
+    Args:
+        tensor (torch.Tensor): The PyTorch tensor one wants to convert to a ndarray
+
+    Returns:
+        ndarray: Converted torch.Tensor to ndarray
+    """
     return tensor.cpu().detach().numpy()
 
 
+def get_device():
+    device = None
+    if torch.cuda.is_available():
+        dev = "cuda:0"
+        device = torch.device(dev)
+    else:
+        dev = "cpu"
+        device = torch.device(dev)
+    return device
+
+
 def compress(model_path, config):
-    # Give the encoding function the correct input as tensor
+    """Function which performs the compression of the input file. In order to compress, you must have a dataset whose path is
+        determined by `input_path` in the `config`. You also need a trained model from path `model_path`. The model path is then used to initialize the model
+        used for compression. The data is then converted into a `torch.tensor` and then passed through `model.encode`.
+
+    Args:
+        model_path (string): Path to where the model is located
+        config (dataClass): Base class selecting user inputs
+
+    Raises:
+        NameError: Baler currently only supports 1D (e.g. HEP) or 2D (e.g. CFD) data as inputs.
+
+    Returns:
+        torch.Tensor: Compressed data as PyTorch tensor
+    """
+
+    # Loads the data and applies normalization if config.apply_normalization = True
     loaded = np.load(config.input_path)
     data_before = loaded["data"]
     if config.apply_normalization:
@@ -218,44 +341,67 @@ def compress(model_path, config):
         if config.data_dimension == 1:
             column_names = np.load(config.input_path)["names"]
             number_of_columns = len(column_names)
-            latent_space_size = int(number_of_columns // config.compression_ratio)
+            config.latent_space_size = int(
+                number_of_columns // config.compression_ratio
+            )
             config.number_of_columns = number_of_columns
         elif config.data_dimension == 2:
-            data = np.load(config.input_path)["data"]
-            number_of_columns = len(data)
-            latent_space_size = int(
-                (number_of_columns * number_of_columns) // config.compression_ratio
+            data = data_before
+            number_of_rows = data.shape[1]
+            config.number_of_columns = data.shape[2]
+            config.latent_space_size = int(
+                (number_of_rows * config.number_of_columns) // config.compression_ratio
             )
         else:
             raise NameError(
-                "Data dimension can only be 1 or 2. Introduced value = "
+                "Data dimension can only be 1 or 2. Got config.data_dimension = "
                 + str(config.data_dimension)
             )
     except AttributeError:
         number_of_columns = config.number_of_columns
         latent_space_size = config.latent_space_size
-        print(number_of_columns, latent_space_size)
+        print(f"{number_of_columns} -> {latent_space_size} dimensions")
 
     # Initialise and load the model correctly.
+    latent_space_size = config.latent_space_size
+    device = get_device()
     model_object = data_processing.initialise_model(config.model_name)
     model = data_processing.load_model(
         model_object,
         model_path=model_path,
-        n_features=number_of_columns,
-        z_dim=latent_space_size,
+        n_features=config.number_of_columns,
+        z_dim=config.latent_space_size,
     )
 
+    # Give the encoding function the correct input as tensor
     if config.data_dimension == 2:
-        data_tensor = torch.from_numpy(data).to(model.device).view(1, 1, 50, 50)
+        data_tensor = (
+            torch.from_numpy(data.astype("float32", casting="same_kind"))
+            .to(device)
+            .view(data.shape[0], 1, data.shape[1], data.shape[2])
+        )
     elif config.data_dimension == 1:
-        data_tensor = torch.from_numpy(data).to(model.device)
+        data_tensor = torch.from_numpy(data).to(device)
 
     compressed = model.encode(data_tensor)
     return compressed
 
 
 def decompress(model_path, input_path, model_name):
-    # Load the data & convert to tensor
+    """Function which performs the decompression of the compressed file. In order to decompress, you must have a compressed file, whose path is
+        determined by `input_path`, a model from path `model_path` and a model_name. The model path and model names are used to initialize the model
+        used for decompression. The data is then converted into a `torch.tensor` and then passed through `model.decode`.
+
+    Args:
+        model_path (string): Path to where the model is located
+        input_path (string): Path to the data you want to decompress
+        model_name (string): Name of trained model from which you want to use decode
+
+    Returns:
+        torch.tensor, ndarray, ndarray: decompressed data as tensor, ndarray of column names, ndarray of normalization features
+    """
+
+    # Load the data & define necessary variables
     loaded = np.load(input_path)
     data = loaded["data"]
     names = loaded["names"]
@@ -265,6 +411,7 @@ def decompress(model_path, input_path, model_name):
     number_of_columns = len(model_dict[list(model_dict.keys())[-1]])
 
     # Initialise and load the model correctly.
+    device = get_device()
     model_object = data_processing.initialise_model(model_name)
     model = data_processing.load_model(
         model_object,
@@ -274,61 +421,8 @@ def decompress(model_path, input_path, model_name):
     )
 
     # Load the data & convert to tensor
-    data_tensor = torch.from_numpy(data).to(model.device)
+    data_tensor = torch.from_numpy(data).to(device)
 
+    # Decompress the data using the trained models decode function
     decompressed = model.decode(data_tensor)
     return decompressed, names, normalization_features
-
-
-def get_device():
-    device = None
-    if torch.cuda.is_available():
-        dev = "cuda:0"
-        device = torch.device(dev)
-    else:
-        dev = "cpu"
-        device = torch.device(dev)
-    return device
-
-
-def compute_e(mass, eta, pt):
-    masspt = pt**2 + mass**2
-    cosh = (np.cosh(eta)) ** 2
-    total = np.sqrt(masspt * cosh)
-    return total
-
-
-def convert_mass_to_energy(df, col_names):
-    # TODO Not used right now
-
-    for i in range(len(col_names)):
-        if col_names[i].split(".")[-1] == "pt":
-            pt = df.iloc[:, i]
-
-        if col_names[i].split(".")[-1] == "mass_":
-            mass = df.iloc[:, i]
-
-            # Store name to rename & replace mass in df:
-            mass_name: str = str(col_names[i])
-
-        if col_names[i].split(".")[-1] == "99":
-            eta = df.iloc[:, i]
-
-        else:
-            print(
-                "Can't convert to energy. Please turn off `energy_conversion` in the config to continue"
-            )
-            exit(1)
-
-    # Compute mass
-    energy = compute_e(mass=mass, eta=eta, pt=pt)
-
-    # Get correct new column name
-    energy_name = mass_name.replace("mass_", "energy_")
-
-    # Replace mass with energy
-    df[mass_name] = energy
-
-    # Replace column name
-    df.columns = df.columns.str.replace(mass_name, energy_name, regex=True)
-    return df
diff --git a/baler/modules/models.py b/baler/modules/models.py
index 62122660..7bf5c42d 100644
--- a/baler/modules/models.py
+++ b/baler/modules/models.py
@@ -3,27 +3,26 @@
 from torch.nn import functional as F
 
 
-class george_SAE(nn.Module):
+class AE(nn.Module):
+
     # This class is a modified version of the original class by George Dialektakis found at
     # https://github.com/Autoencoders-compression-anomaly/Deep-Autoencoders-Data-Compression-GSoC-2021
     # Released under the Apache License 2.0 found at https://www.apache.org/licenses/LICENSE-2.0.txt
     # Copyright 2021 George Dialektakis
 
-    def __init__(self, device, n_features, z_dim, *args, **kwargs):
-        super().__init__(*args, **kwargs)
-
-        self.device = device
+    def __init__(self, n_features, z_dim, *args, **kwargs):
+        super(AE, self).__init__(*args, **kwargs)
 
         # encoder
-        self.en1 = nn.Linear(n_features, 200, dtype=torch.float64, device=device)
-        self.en2 = nn.Linear(200, 100, dtype=torch.float64, device=device)
-        self.en3 = nn.Linear(100, 50, dtype=torch.float64, device=device)
-        self.en4 = nn.Linear(50, z_dim, dtype=torch.float64, device=device)
+        self.en1 = nn.Linear(n_features, 200, dtype=torch.float64)
+        self.en2 = nn.Linear(200, 100, dtype=torch.float64)
+        self.en3 = nn.Linear(100, 50, dtype=torch.float64)
+        self.en4 = nn.Linear(50, z_dim, dtype=torch.float64)
         # decoder
-        self.de1 = nn.Linear(z_dim, 50, dtype=torch.float64, device=device)
-        self.de2 = nn.Linear(50, 100, dtype=torch.float64, device=device)
-        self.de3 = nn.Linear(100, 200, dtype=torch.float64, device=device)
-        self.de4 = nn.Linear(200, n_features, dtype=torch.float64, device=device)
+        self.de1 = nn.Linear(z_dim, 50, dtype=torch.float64)
+        self.de2 = nn.Linear(50, 100, dtype=torch.float64)
+        self.de3 = nn.Linear(100, 200, dtype=torch.float64)
+        self.de4 = nn.Linear(200, n_features, dtype=torch.float64)
 
         self.n_features = n_features
         self.z_dim = z_dim
@@ -46,71 +45,9 @@ def forward(self, x):
         return self.decode(z)
 
 
-class george_SAE_BN(nn.Module):
-    def __init__(self, device, n_features, z_dim):
-        super(george_SAE_BN, self).__init__()
-
-        self.device = device
-
-        # encoder
-        self.enc_nn = nn.Sequential(
-            nn.Linear(n_features, 200, dtype=torch.float64, device=device),
-            # nn.Dropout(p=0.5),
-            nn.LeakyReLU(),
-            nn.BatchNorm1d(200, dtype=torch.float64, device=device),
-            nn.Linear(200, 100, dtype=torch.float64, device=device),
-            # nn.Dropout(p=0.4),
-            nn.LeakyReLU(),
-            nn.BatchNorm1d(100, dtype=torch.float64, device=device),
-            nn.Linear(100, 50, dtype=torch.float64, device=device),
-            # nn.Dropout(p=0.3),
-            nn.LeakyReLU(),
-            nn.BatchNorm1d(50, dtype=torch.float64, device=device),
-            nn.Linear(50, z_dim, dtype=torch.float64, device=device),
-            # nn.Dropout(p=0.2),
-            nn.LeakyReLU(),
-            nn.BatchNorm1d(z_dim, dtype=torch.float64, device=device),
-        )
-
-        # decoder
-        self.dec_nn = nn.Sequential(
-            nn.Linear(z_dim, 50, dtype=torch.float64, device=device),
-            # nn.Dropout(p=0.2),
-            nn.LeakyReLU(),
-            # nn.BatchNorm1d(50,dtype=torch.float64),
-            nn.Linear(50, 100, dtype=torch.float64, device=device),
-            # nn.Dropout(p=0.3),
-            nn.LeakyReLU(),
-            # nn.BatchNorm1d(100,dtype=torch.float64),
-            nn.Linear(100, 200, dtype=torch.float64, device=device),
-            # nn.Dropout(p=0.4),
-            nn.LeakyReLU(),
-            # nn.BatchNorm1d(200,dtype=torch.float64),
-            nn.Linear(200, n_features, dtype=torch.float64, device=device),
-            # nn.Dropout(p=0.5),
-            # nn.BatchNorm1d(n_features,dtype=torch.float64),
-            nn.ReLU(),
-        )
-
-        self.n_features = n_features
-        self.z_dim = z_dim
-
-    def encode(self, x):
-        out = self.enc_nn(x)
-        return out
-
-    def decode(self, z):
-        out = self.dec_nn(z)
-        return out
-
-    def forward(self, x):
-        z = self.encode(x)
-        return self.decode(z)
-
-
-class george_SAE_Dropout_BN(nn.Module):
-    def __init__(self, n_features, z_dim):
-        super(george_SAE_Dropout_BN, self).__init__()
+class AE_Dropout_BN(nn.Module):
+    def __init__(self, n_features, z_dim, *args, **kwargs):
+        super(AE_Dropout_BN, self).__init__(*args, **kwargs)
 
         # encoder
         self.enc_nn = nn.Sequential(
@@ -168,78 +105,15 @@ def forward(self, x):
         return self.decode(z)
 
 
-class george_SAE_Dropout(nn.Module):
-    def __init__(self, n_features, z_dim):
-        super(george_SAE_Dropout, self).__init__()
-
-        # encoder
-        self.enc_nn = nn.Sequential(
-            nn.Linear(n_features, 200, dtype=torch.float64),
-            nn.Dropout(p=0.5),
-            nn.LeakyReLU(),
-            nn.Linear(200, 100, dtype=torch.float64),
-            nn.Dropout(p=0.4),
-            nn.LeakyReLU(),
-            nn.Linear(100, 50, dtype=torch.float64),
-            nn.Dropout(p=0.3),
-            nn.LeakyReLU(),
-            nn.Linear(50, z_dim, dtype=torch.float64),
-            nn.Dropout(p=0.2),
-            nn.LeakyReLU(),
-        )
-
-        # decoder
-        self.dec_nn = nn.Sequential(
-            nn.Linear(z_dim, 50, dtype=torch.float64),
-            nn.Dropout(p=0.2),
-            nn.LeakyReLU(),
-            nn.Linear(50, 100, dtype=torch.float64),
-            nn.Dropout(p=0.3),
-            nn.LeakyReLU(),
-            nn.Linear(100, 200, dtype=torch.float64),
-            nn.Dropout(p=0.4),
-            nn.LeakyReLU(),
-            nn.Linear(200, n_features, dtype=torch.float64),
-            nn.Dropout(p=0.5),
-            nn.ReLU(),
-        )
-
-        self.n_features = n_features
-        self.z_dim = z_dim
-
-    def encode(self, x):
-        out = self.enc_nn(x)
-        return out
-
-    def decode(self, z):
-        out = self.dec_nn(z)
-        return out
-
-    def forward(self, x):
-        # z = x.view(batch_size,a,b,c) ? What is this
-        return self.decode(x)
-
-    def loss(self, model_children, true_data, reconstructed_data, reg_param):
-        mse = nn.MSELoss()
-        mse_loss = mse(reconstructed_data, true_data)
-        l1_loss = 0
-        values = true_data
-        for i in range(len(model_children)):
-            values = F.relu((model_children[i](values)))
-            l1_loss += torch.mean(torch.abs(values))
-        loss = mse_loss + reg_param * l1_loss
-        return loss
-
-
 class Conv_AE(nn.Module):
-    def __init__(self, device, n_features, z_dim, *args, **kwargs):
+    def __init__(self, n_features, z_dim, *args, **kwargs):
         super(Conv_AE, self).__init__(*args, **kwargs)
 
-        self.device = device
-
         self.q_z_mid_dim = 2000
         self.q_z_output_dim = 72128
+
         # Encoder
+
         # Conv Layers
         self.q_z_conv = nn.Sequential(
             nn.Conv2d(1, 8, kernel_size=(2, 5), stride=(1), padding=(1)),
@@ -264,6 +138,8 @@ def __init__(self, device, n_features, z_dim, *args, **kwargs):
         )
 
         # Decoder
+
+        # Linear layers
         self.p_x_lin = nn.Sequential(
             nn.Linear(z_dim, self.q_z_mid_dim),
             nn.ReLU(),
@@ -286,19 +162,18 @@ def __init__(self, device, n_features, z_dim, *args, **kwargs):
     def encode(self, x):
         # Conv
         out = self.q_z_conv(x)
-        # flatten
+        # Flatten
         out = self.flatten(out)
-        # out = out.view(out.size(0), -1)
-        # dense
+        # Dense
         out = self.q_z_lin(out)
         return out
 
     def decode(self, z):
-        # dense
+        # Dense
         out = self.p_x_lin(z)
-        # reshape
+        # Unflatten
         out = out.view(out.size(0), 32, 49, 46)
-        # DeConv/UnConv?
+        # Conv transpose
         out = self.p_x_conv(out)
         return out
 
diff --git a/baler/modules/plotting.py b/baler/modules/plotting.py
index 66d09915..eb05d400 100644
--- a/baler/modules/plotting.py
+++ b/baler/modules/plotting.py
@@ -1,19 +1,18 @@
-from mpl_toolkits.axes_grid1 import make_axes_locatable
-import numpy as np
 import matplotlib.pyplot as plt
-import pickle
+import numpy as np
+from tqdm import tqdm
+from tqdm import trange
 from matplotlib.backends.backend_pdf import PdfPages
 
 
-# pickle_file = "./data/cfd2/cfd.pickle"
-
-# directory = "ConvAEDes_75"
-# directory = "ConvAEDes_500"
-
-# decompressed_cfd = f"./projects/{directory}/decompressed_output/decompressed.pickle"
-
-
 def loss_plot(path_to_loss_data, output_path, config):
+    """This function Plots the loss from the training and saves it
+
+    Args:
+        path_to_loss_data (string): Path to file containing loss plot data generated during training
+        output_path (string): Directory path to which the loss plot is saved
+        config (dataclass): The config class containing attributes set in the config file
+    """
     loss_data = np.load(path_to_loss_data)
     str_list = ["Epochs:", "Model Name:", "Reg. Param:", "lr:", "BS:"]
 
@@ -41,20 +40,32 @@ def loss_plot(path_to_loss_data, output_path, config):
     # plt.show()
 
 
-def pickle_to_df(file):
-    # From pickle to df:
-    with open(file, "rb") as handle:
-        data = pickle.load(handle)
-
-        return data
+def get_index_to_cut(column_index, cut, array):
+    """Given an array column index and a threshold, this function returns the index of the
+        entries not passing the threshold.
 
+    Args:
+        column_index (int): The index for the column where cuts should be applied
+        cut (float): Threshold for which values below will have the whole entry removed
+        array (np.array): The full array to be edited
 
-def get_index_to_cut(column_index, cut, array):
+    Returns:
+        _type_: returns the index of the rows to be removed
+    """
     indices_to_cut = np.argwhere(array[column_index] < cut).flatten()
     return indices_to_cut
 
 
 def plot_1D(project_path, config):
+    """General plotting for 1D data, for example data from a '.csv' file. This function generates a pdf
+        document where each page contains the before/after performance
+        of each column of the 1D data
+
+    Args:
+        project_path (string): The path to the project directory
+        config (dataclass): The config class containing attributes set in the config file
+    """
+
     output_path = project_path + "training/"
     before_path = config.input_path
     after_path = project_path + "decompressed_output/decompressed.npz"
@@ -82,9 +93,11 @@ def plot_1D(project_path, config):
         ax4 = axsRight[1]
 
         number_of_columns = len(names)
-        for index, column in enumerate(names):
+
+        print("=== Plotting ===")
+
+        for index, column in enumerate(tqdm(names)):
             column_name = column.split(".")[-1]
-            print(f"Plotting: {column_name} ({index+1} of {number_of_columns})")
             rms = np.sqrt(np.mean(np.square(response[index])))
             residual_RMS = np.sqrt(np.mean(np.square(residual[index])))
 
@@ -189,51 +202,88 @@ def plot_1D(project_path, config):
             ax3.clear()
             ax4.clear()
 
-            # if index == 3:
-            #    break
-
 
 def plot_2D(project_path, config):
-    data = np.load(config.input_path)["data"]
+    """General plotting for 2D data, for example 2D arraysfrom computational fluid
+        dynamics or other image like data. This function generates a pdf
+        document where each page contains the before/after performance
+        of each column of the 1D data
 
+    Args:
+        project_path (string): The path to the project directory
+        config (dataclass): The config class containing attributes set in the config file
+    """
+
+    data = np.load(config.input_path)["data"]
     data_decompressed = np.load(project_path + "/decompressed_output/decompressed.npz")[
         "data"
-    ].reshape(50, 50)
-
-    diff = ((data_decompressed - data) / data) * 100
-
-    fig, axs = plt.subplots(
-        1, 3, figsize=(29.7 * (1 / 2.54), 21 * (1 / 2.54)), sharey=True
-    )
-    axs[0].set_title("Original", fontsize=11, y=-0.2)
-    im1 = axs[0].imshow(
-        data, vmin=-0.01, vmax=0.07, cmap="CMRmap", interpolation="nearest"
-    )
-    cb2 = plt.colorbar(im1, ax=[axs[0]], location="top")
-
-    axs[1].set_title("Decompressed", fontsize=11, y=-0.2)
-    im2 = axs[1].imshow(
-        data_decompressed, vmin=-0.01, vmax=0.07, cmap="CMRmap", interpolation="nearest"
-    )
-    cb2 = plt.colorbar(im2, ax=[axs[1]], location="top")
-
-    axs[2].set_title("Relative Diff. [%]", fontsize=11, y=-0.2)
-    im3 = axs[2].imshow(diff, vmin=-10, vmax=10, cmap="cool_r", interpolation="nearest")
-    cb2 = plt.colorbar(im3, ax=[axs[2]], location="top")
-
-    plt.ylim(0, 50)
-    plt.xlim(0, 50)
-    fig.suptitle(
-        "Compressed file is 10% the size of original,\n75 epochs (3.5 min)",
-        y=0.9,
-        fontsize=16,
-    )
-    # fig.suptitle('Compressed file is 10% the size of original,\n500 epochs (20 min)',y=0.9, fontsize=16)
+    ]
 
-    fig.savefig(project_path + "/plotting/" + "CFD.jpg", bbox_inches="tight")
+    if data.shape[0] > 1:
+        num_tiles = data.shape[0]
+    else:
+        num_tiles = 1
+
+    print("=== Plotting ===")
+
+    for ind in trange(num_tiles):
+        tile_data_decompressed = data_decompressed[ind].reshape(
+            data_decompressed.shape[2], data_decompressed.shape[3]
+        )
+        tile_data = data[ind].reshape(data.shape[1], data.shape[2])
+
+        diff = ((tile_data_decompressed - tile_data) / tile_data) * 100
+
+        fig, axs = plt.subplots(
+            1, 3, figsize=(29.7 * (1 / 2.54), 21 * (1 / 2.54)), sharey=True
+        )
+        axs[0].set_title("Original", fontsize=11, y=-0.2)
+        im1 = axs[0].imshow(
+            tile_data, vmin=-0.01, vmax=0.07, cmap="CMRmap", interpolation="nearest"
+        )
+        cb2 = plt.colorbar(im1, ax=[axs[0]], location="top")
+
+        axs[1].set_title("Decompressed", fontsize=11, y=-0.2)
+        im2 = axs[1].imshow(
+            tile_data_decompressed,
+            vmin=-0.01,
+            vmax=0.07,
+            cmap="CMRmap",
+            interpolation="nearest",
+        )
+        cb2 = plt.colorbar(im2, ax=[axs[1]], location="top")
+
+        axs[2].set_title("Relative Diff. [%]", fontsize=11, y=-0.2)
+        im3 = axs[2].imshow(
+            diff, vmin=-10, vmax=10, cmap="cool_r", interpolation="nearest"
+        )
+        cb2 = plt.colorbar(im3, ax=[axs[2]], location="top")
+
+        plt.ylim(0, 50)
+        plt.xlim(0, 50)
+        fig.suptitle(
+            "Compressed file is 10% the size of original,\n75 epochs (3.5 min)",
+            y=0.9,
+            fontsize=16,
+        )
+        fig.suptitle(
+            "Compressed file is 10% the size of original,\n500 epochs (20 min)",
+            y=0.9,
+            fontsize=16,
+        )
+
+        fig.savefig(
+            project_path + "/plotting/CFD" + str(ind) + ".jpg", bbox_inches="tight"
+        )
 
 
 def plot(project_path, config):
+    """Runs the appropriate plotting function based on the data dimension 1D or 2D
+
+    Args:
+        project_path (string): The path to the project directory
+        config (dataclass): The config class containing attributes set in the config file
+    """
     if config.data_dimension == 1:
         plot_1D(project_path, config)
     elif config.data_dimension == 2:
diff --git a/baler/modules/training.py b/baler/modules/training.py
index 79a3845b..19689082 100644
--- a/baler/modules/training.py
+++ b/baler/modules/training.py
@@ -1,3 +1,4 @@
+import os
 import random
 import time
 
@@ -6,33 +7,42 @@
 from torch.utils.data import DataLoader
 from tqdm import tqdm
 
-import modules.utils as utils
 import modules.helper as helper
-
-import os
+import modules.utils as utils
 
 
 def fit(
     model, train_dl, model_children, regular_param, optimizer, RHO, l1, n_dimensions
 ):
-    # This function is a modified version of the original function by George Dialektakis found at
-    # https://github.com/Autoencoders-compression-anomaly/Deep-Autoencoders-Data-Compression-GSoC-2021
-    # Released under the Apache License 2.0 found at https://www.apache.org/licenses/LICENSE-2.0.txt
-    # Copyright 2021 George Dialektakis
+
+    """This function trains the model on the train set. It computes the losses and does the backwards propagation, and updates the optimizer as well.
+
+    Args:
+        model (modelObject): The model you wish to train
+        train_dl (torch.DataLoader): Defines the batched data which the model is trained on
+        model_children (list): List of model parameters
+        regular_param (float): Determines proportionality constant for the gradient descent step.
+        optimizer (torch.optim): Chooses optimizer for gradient descent.
+        RHO (float): Float used for KL Divergence (Not currently a feature)
+        l1 (boolean): If `True`, use L1 regularization. Otherwise, don't.
+        n_dimensions (int): Number of dimensions.
+
+    Returns:
+        list, model object: Losses and trained model
+    """
 
     print("### Beginning Training")
 
     model.train()
 
     running_loss = 0.0
-    counter = 0
+    device = helper.get_device()
 
-    for inputs in tqdm(train_dl):
-        counter += 1
-        inputs = inputs.to(model.device)
+    for idx, inputs in enumerate(tqdm(train_dl)):
+        inputs = inputs.to(device)
         optimizer.zero_grad()
         reconstructions = model(inputs)
-        loss, mse_loss, l1_loss = utils.sparse_SumLoss_function_l1(
+        loss, mse_loss, l1_loss = utils.mse_loss_l1(
             model_children=model_children,
             true_data=inputs,
             reconstructed_data=reconstructions,
@@ -45,30 +55,36 @@ def fit(
 
         running_loss += loss.item()
 
-    epoch_loss = running_loss / counter
+    epoch_loss = running_loss / (idx + 1)
     print(f"# Finished. Training Loss: {loss:.6f}")
     return epoch_loss, mse_loss, l1_loss, model
 
 
 def validate(model, test_dl, model_children, reg_param):
-    print("### Beginning Validating")
+    """Function used to validate the training. Not necessary for doing compression, but gives a good indication of wether the model selected is a good fit or not.
 
-    # This function is a modified version of the original function by George Dialektakis found at
-    # https://github.com/Autoencoders-compression-anomaly/Deep-Autoencoders-Data-Compression-GSoC-2021
-    # Released under the Apache License 2.0 found at https://www.apache.org/licenses/LICENSE-2.0.txt
-    # Copyright 2021 George Dialektakis
+    Args:
+        model (modelObject): Defines the model one wants to validate. The model used here is passed directly from `fit()`.
+        test_dl (torch.DataLoader): Defines the batched data which the model is validated on
+        model_children (list): List of model parameters
+        regular_param (float): Determines proportionality constant for the gradient descent step.
+
+    Returns:
+        _type_: _description_
+    """
+    print("### Beginning Validating")
 
     model.eval()
-    counter = 0
+
     running_loss = 0.0
+    device = helper.get_device()
 
     with torch.no_grad():
-        for inputs in tqdm(test_dl):
-            counter += 1
-            inputs = inputs.to(model.device)
+        for idx, inputs in enumerate(tqdm(test_dl)):
+            inputs = inputs.to(device)
             reconstructions = model(inputs)
 
-            loss, _, _ = utils.sparse_loss_function_l1(
+            loss, _, _ = utils.mse_loss_l1(
                 model_children=model_children,
                 true_data=inputs,
                 reconstructed_data=reconstructions,
@@ -77,7 +93,7 @@ def validate(model, test_dl, model_children, reg_param):
             )
             running_loss += loss.item()
 
-    epoch_loss = running_loss / counter
+    epoch_loss = running_loss / (idx + 1)
     print(f"# Finished. Validation Loss: {loss:.6f}")
     return epoch_loss
 
@@ -88,12 +104,22 @@ def seed_worker(worker_id):
     random.seed(worker_seed)
 
 
-def train(model, variables, train_data, test_data, parent_path, config):
+def train(model, variables, train_data, test_data, project_path, config):
+
+    """Calls the `fit()` and `validate()` functions in a loop, which defines how many "times" the network should be trained.
 
-    # This function is a modified version of the original function by George Dialektakis found at
-    # https://github.com/Autoencoders-compression-anomaly/Deep-Autoencoders-Data-Compression-GSoC-2021
-    # Released under the Apache License 2.0 found at https://www.apache.org/licenses/LICENSE-2.0.txt
-    # Copyright 2021 George Dialektakis
+    Args:
+        model (modelObject): The model you wish to train
+        variables (_type_): _description_
+        train_set (ndarray): Array consisting of the train set
+        test_set (ndarray): Array consisting of the test set
+        project_path (string): Path to the project directory
+        config (dataClass): Base class selecting user inputs
+
+    Returns:
+        modelObject: fully trained model ready to perform compression and decompression
+    """
+    # Fix the random seed - TODO: add flag to make this optional
 
     random.seed(0)
     torch.manual_seed(0)
@@ -110,6 +136,8 @@ def train(model, variables, train_data, test_data, parent_path, config):
     l1 = config.l1
     epochs = config.epochs
     latent_space_size = config.latent_space_size
+    intermittent_model_saving = config.intermittent_model_saving
+    intermittent_saving_patience = config.intermittent_saving_patience
 
     model_children = list(model.children())
 
@@ -120,10 +148,10 @@ def train(model, variables, train_data, test_data, parent_path, config):
     # Converting data to tensors
     if config.data_dimension == 2:
         train_ds = torch.tensor(train_data, dtype=torch.float32, device=device).view(
-            1, 1, len(test_data), len(test_data)
+            train_data.shape[0], 1, train_data.shape[1], train_data.shape[2]
         )
         valid_ds = torch.tensor(test_data, dtype=torch.float32, device=device).view(
-            1, 1, len(test_data), len(test_data)
+            test_data.shape[0], 1, test_data.shape[1], test_data.shape[2]
         )
     elif config.data_dimension == 1:
         train_ds = torch.tensor(train_data, dtype=torch.float64, device=device)
@@ -145,7 +173,7 @@ def train(model, variables, train_data, test_data, parent_path, config):
         worker_init_fn=seed_worker,
         generator=g,
         drop_last=False,
-    )  # Used to be batch_size = bs * 2
+    )
 
     # Select Optimizer
     optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)
@@ -153,14 +181,16 @@ def train(model, variables, train_data, test_data, parent_path, config):
     # Activate early stopping
     if config.early_stopping:
         early_stopping = utils.EarlyStopping(
-            patience=config.patience, min_delta=config.min_delta
+            patience=config.early_stopping_patience, min_delta=config.min_delta
         )  # Changes to patience & min_delta can be made in configs
 
     # Activate LR Scheduler
     if config.lr_scheduler:
-        lr_scheduler = utils.LRScheduler(optimizer=optimizer, patience=config.patience)
+        lr_scheduler = utils.LRScheduler(
+            optimizer=optimizer, patience=config.lr_scheduler_patience
+        )
 
-    # train and validate the autoencoder neural network
+    # Training and Validation of the model
     train_loss = []
     val_loss = []
     start = time.time()
@@ -200,37 +230,15 @@ 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 = False
-        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")
-
+        ## Implementation to save models & values after every N epochs, where N is stored in 'intermittent_saving_patience':
+        if intermittent_model_saving:
+            if epoch % intermittent_saving_patience == 0:
+                path = os.path.join(project_path, f"model_{epoch}.pt")
                 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")
-    np.save(parent_path + "loss_data.npy", np.array([train_loss, val_loss]))
-
-    # print("1")
-    # data_as_tensor = torch.tensor(test_data, dtype=torch.float32)
-    # print("2")
-    # print(test_data.shape, data_as_tensor.shape)
-    # data_as_tensor = data_as_tensor.to(trained_model.device)
-    # print("3")
-    # print(test_data.shape, data_as_tensor.shape)
-    # pred_as_tensor = trained_model(data_as_tensor)
+    np.save(project_path + "loss_data.npy", np.array([train_loss, val_loss]))
 
     return trained_model
diff --git a/baler/modules/utils.py b/baler/modules/utils.py
index 6e794e19..f7efd0e5 100644
--- a/baler/modules/utils.py
+++ b/baler/modules/utils.py
@@ -1,16 +1,33 @@
 import torch
 import torch.nn as nn
 from scipy.stats import wasserstein_distance
-from torch.nn import functional as F
+from torch.nn import functional
 from tqdm import tqdm
 
 factor = 0.5
 min_lr = 1e-6
 
 
-def sparse_loss_function_emd_l1(
-    model_children, true_data, reconstructed_data, reg_param, validate
-):
+def mse_loss_emd_l1(model_children, true_data, reconstructed_data, reg_param, validate):
+    """
+    Computes a sparse loss function consisting of three terms: the Earth Mover's Distance (EMD) loss between the
+    true and reconstructed data, the mean squared error (MSE) loss between the reconstructed and true data, and a
+    L1 regularization term on the output of a list of model children.
+
+    Args: model_children (list): List of PyTorch modules representing the model architecture to be regularized.
+    true_data (torch.Tensor): The ground truth data, with shape (batch_size, num_features). reconstructed_data (
+    torch.Tensor): The reconstructed data, with shape (batch_size, num_features). reg_param (float): The weight of
+    the L1 regularization term in the loss function. validate (bool): If True, returns only the EMD loss. If False,
+    computes the full loss with the L1 regularization term.
+
+    Returns:
+        If validate is False, returns a tuple with three elements:
+        - loss (torch.Tensor): The full sparse loss function, with shape ().
+        - emd_loss (float): The EMD loss between the true and reconstructed data.
+        - l1_loss (float): The L1 regularization term on the output of the model children.
+
+        If validate is True, returns only the EMD loss as a float.
+    """
     mse = nn.MSELoss()
     mse_loss = mse(reconstructed_data, true_data)
     wasserstein_distance_list = [
@@ -34,14 +51,31 @@ def sparse_loss_function_emd_l1(
         return emd_loss
 
 
-def sparse_loss_function_l1(
-    model_children, true_data, reconstructed_data, reg_param, validate
-):
+def mse_loss_l1(model_children, true_data, reconstructed_data, reg_param, validate):
+
     # This function is a modified version of the original function by George Dialektakis found at
     # https://github.com/Autoencoders-compression-anomaly/Deep-Autoencoders-Data-Compression-GSoC-2021
     # Released under the Apache License 2.0 found at https://www.apache.org/licenses/LICENSE-2.0.txt
     # Copyright 2021 George Dialektakis
 
+    """
+    Computes a sparse loss function consisting of two terms: the mean squared error (MSE) loss between the
+    reconstructed and true data, and a L1 regularization term on the output of a list of model children.
+
+    Args: model_children (list): List of PyTorch modules representing the model architecture to be regularized.
+    true_data (torch.Tensor): The ground truth data, with shape (batch_size, num_features). reconstructed_data (
+    torch.Tensor): The reconstructed data, with shape (batch_size, num_features). reg_param (float): The weight of
+    the L1 regularization term in the loss function. validate (bool): If True, returns only the MSE loss. If False,
+    computes the full loss with the L1 regularization term.
+
+    Returns:
+        If validate is False, returns a tuple with three elements:
+        - loss (torch.Tensor): The full sparse loss function, with shape ().
+        - mse_loss (float): The MSE loss between the true and reconstructed data.
+        - l1_loss (float): The L1 regularization term on the output of the model children.
+
+        If validate is True, returns only the MSE loss as a float.
+    """
     mse = nn.MSELoss()
     mse_loss = mse(reconstructed_data, true_data)
 
@@ -49,7 +83,7 @@ def sparse_loss_function_l1(
     values = true_data
     if not validate:
         for i in range(len(model_children)):
-            values = F.relu(model_children[i](values))
+            values = functional.relu(model_children[i](values))
             l1_loss += torch.mean(torch.abs(values))
 
         loss = mse_loss + reg_param * l1_loss
@@ -58,9 +92,27 @@ def sparse_loss_function_l1(
         return mse_loss, 0, 0
 
 
-def sparse_SumLoss_function_l1(
-    model_children, true_data, reconstructed_data, reg_param, validate
-):
+def mse_sum_loss_l1(model_children, true_data, reconstructed_data, reg_param, validate):
+    """
+    Computes the sum of mean squared error (MSE) loss and L1 regularization loss.
+
+    Args:
+        model_children (list): List of PyTorch modules representing the encoder network.
+        true_data (tensor): Ground truth tensor of shape (batch_size, input_size).
+        reconstructed_data (tensor): Reconstructed tensor of shape (batch_size, input_size).
+        reg_param (float): Regularization parameter for L1 loss.
+        validate (bool): Whether to return only MSE loss or both MSE and L1 losses.
+
+    Returns:
+        If validate is False:
+            loss (tensor): Total loss consisting of MSE loss and L1 regularization loss.
+            mse_sum_loss (tensor): Mean squared error loss.
+            l1_loss (tensor): L1 regularization loss.
+        If validate is True:
+            mse_sum_loss (tensor): Mean squared error loss.
+            0 (int): Placeholder for MSE loss since it is not calculated during validation.
+            0 (int): Placeholder for L1 loss since it is not calculated during validation.
+    """
     mse_sum = nn.MSELoss(reduction="sum")
     mse_loss = mse_sum(reconstructed_data, true_data)
     number_of_columns = true_data.shape[1]
@@ -71,7 +123,7 @@ def sparse_SumLoss_function_l1(
     values = true_data
     if not validate:
         for i in range(len(model_children)):
-            values = F.relu(model_children[i](values))
+            values = functional.relu(model_children[i](values))
             l1_loss += torch.mean(torch.abs(values))
 
         loss = mse_sum_loss + reg_param * l1_loss
@@ -85,6 +137,16 @@ def sparse_SumLoss_function_l1(
 
 
 def accuracy(model, dataloader):
+    """
+    Computes the accuracy of a PyTorch model on a given dataset.
+
+    Args:
+        model (nn.Module): The PyTorch model to evaluate.
+        dataloader (DataLoader): DataLoader object containing the dataset to evaluate on.
+
+    Returns:
+        accuracy_frac (float): The fraction of correctly classified instances in the dataset.
+    """
     print("Accuracy")
     model.eval()
 
@@ -107,7 +169,20 @@ def accuracy(model, dataloader):
 
 
 class EarlyStopping:
-    def __init__(self, patience, min_delta):
+    """
+    Class to perform early stopping during model training.
+
+    Attributes:
+        patience (int): The number of epochs to wait before stopping the training process if the
+            validation loss doesn't improve.
+        min_delta (float): The minimum difference between the new loss and the previous best loss
+            for the new loss to be considered an improvement.
+        counter (int): Counts the number of times the validation loss hasn't improved.
+        best_loss (float): The best validation loss observed so far.
+        early_stop (bool): Flag that indicates whether early stopping criteria have been met.
+    """
+
+    def __init__(self, patience: int, min_delta: float):
         self.patience = patience  # Nr of times we allow val. loss to not improve before early stopping
         self.min_delta = min_delta  # min(new loss - best loss) for new loss to be considered improvement
         self.counter = 0  # counts nr of times val_loss dosent improve
@@ -132,6 +207,30 @@ def __call__(self, train_loss):
 
 
 class LRScheduler:
+    """
+    A learning rate scheduler that adjusts the learning rate of an optimizer based on the training loss.
+
+    Args:
+        optimizer (torch.optim.Optimizer): The optimizer whose learning rate will be adjusted.
+        patience (int): The number of epochs with no improvement in training loss after which the learning rate
+            will be reduced.
+        min_lr (float, optional): The minimum learning rate that can be reached (default: 1e-6).
+        factor (float, optional): The factor by which the learning rate will be reduced (default: 0.1).
+
+    Attributes:
+        lr_scheduler (torch.optim.lr_scheduler.ReduceLROnPlateau): The PyTorch learning rate scheduler that
+            actually performs the adjustments.
+
+    Example usage:
+        optimizer = torch.optim.Adam(model.parameters(), lr=0.01)
+        lr_scheduler = LRScheduler(optimizer, patience=3, min_lr=1e-5, factor=0.5)
+
+        for epoch in range(num_epochs):
+            train_loss = train(model, train_data_loader)
+            lr_scheduler(train_loss)
+            # ...
+    """
+
     def __init__(self, optimizer, patience, min_lr=min_lr, factor=factor):
         self.optimizer = optimizer
         self.patience = patience
diff --git a/data/example_CFD/example_CFD.npz b/data/example_CFD/example_CFD.npz
index 7852addc..c7710b9d 100644
Binary files a/data/example_CFD/example_CFD.npz and b/data/example_CFD/example_CFD.npz differ
diff --git a/poetry.lock b/poetry.lock
index 6181e029..8907154d 100644
--- a/poetry.lock
+++ b/poetry.lock
@@ -19,6 +19,18 @@ docs = ["furo", "myst-parser", "sphinx", "sphinx-notfound-page", "sphinxcontrib-
 tests = ["attrs[tests-no-zope]", "zope.interface"]
 tests-no-zope = ["cloudpickle", "cloudpickle", "hypothesis", "hypothesis", "mypy (>=0.971,<0.990)", "mypy (>=0.971,<0.990)", "pympler", "pympler", "pytest (>=4.3.0)", "pytest (>=4.3.0)", "pytest-mypy-plugins", "pytest-mypy-plugins", "pytest-xdist[psutil]", "pytest-xdist[psutil]"]
 
+[[package]]
+name = "cfgv"
+version = "3.3.1"
+description = "Validate configuration and produce human readable error messages."
+category = "main"
+optional = false
+python-versions = ">=3.6.1"
+files = [
+    {file = "cfgv-3.3.1-py2.py3-none-any.whl", hash = "sha256:c6a0883f3917a037485059700b9e75da2464e6c27051014ad85ba6aaa5884426"},
+    {file = "cfgv-3.3.1.tar.gz", hash = "sha256:f5a830efb9ce7a445376bb66ec94c638a9787422f96264c98edc6bdeed8ab736"},
+]
+
 [[package]]
 name = "colorama"
 version = "0.4.6"
@@ -118,6 +130,18 @@ files = [
     {file = "cycler-0.11.0.tar.gz", hash = "sha256:9c87405839a19696e837b3b818fed3f5f69f16f1eec1a1ad77e043dcea9c772f"},
 ]
 
+[[package]]
+name = "distlib"
+version = "0.3.6"
+description = "Distribution utilities"
+category = "main"
+optional = false
+python-versions = "*"
+files = [
+    {file = "distlib-0.3.6-py2.py3-none-any.whl", hash = "sha256:f35c4b692542ca110de7ef0bea44d73981caeb34ca0b9b6b2e6d7790dda8f80e"},
+    {file = "distlib-0.3.6.tar.gz", hash = "sha256:14bad2d9b04d3a36127ac97f30b12a19268f211063d8f8ee4f47108896e11b46"},
+]
+
 [[package]]
 name = "exceptiongroup"
 version = "1.1.1"
@@ -133,6 +157,22 @@ files = [
 [package.extras]
 test = ["pytest (>=6)"]
 
+[[package]]
+name = "filelock"
+version = "3.10.2"
+description = "A platform independent file lock."
+category = "main"
+optional = false
+python-versions = ">=3.7"
+files = [
+    {file = "filelock-3.10.2-py3-none-any.whl", hash = "sha256:eb8f0f2d37ed68223ea63e3bddf2fac99667e4362c88b3f762e434d160190d18"},
+    {file = "filelock-3.10.2.tar.gz", hash = "sha256:75997740323c5f12e18f10b494bc11c03e42843129f980f17c04352cc7b09d40"},
+]
+
+[package.extras]
+docs = ["furo (>=2022.12.7)", "sphinx (>=6.1.3)", "sphinx-autodoc-typehints (>=1.22,!=1.23.4)"]
+testing = ["covdefaults (>=2.3)", "coverage (>=7.2.2)", "pytest (>=7.2.2)", "pytest-cov (>=4)", "pytest-timeout (>=2.1)"]
+
 [[package]]
 name = "fonttools"
 version = "4.39.2"
@@ -159,6 +199,21 @@ ufo = ["fs (>=2.2.0,<3)"]
 unicode = ["unicodedata2 (>=15.0.0)"]
 woff = ["brotli (>=1.0.1)", "brotlicffi (>=0.8.0)", "zopfli (>=0.1.4)"]
 
+[[package]]
+name = "identify"
+version = "2.5.21"
+description = "File identification library for Python"
+category = "main"
+optional = false
+python-versions = ">=3.7"
+files = [
+    {file = "identify-2.5.21-py2.py3-none-any.whl", hash = "sha256:69edcaffa8e91ae0f77d397af60f148b6b45a8044b2cc6d99cafa5b04793ff00"},
+    {file = "identify-2.5.21.tar.gz", hash = "sha256:7671a05ef9cfaf8ff63b15d45a91a1147a03aaccb2976d4e9bd047cbbc508471"},
+]
+
+[package.extras]
+license = ["ukkonen"]
+
 [[package]]
 name = "importlib-resources"
 version = "5.12.0"
@@ -343,6 +398,21 @@ pillow = ">=6.2.0"
 pyparsing = ">=2.3.1"
 python-dateutil = ">=2.7"
 
+[[package]]
+name = "nodeenv"
+version = "1.7.0"
+description = "Node.js virtual environment builder"
+category = "main"
+optional = false
+python-versions = ">=2.7,!=3.0.*,!=3.1.*,!=3.2.*,!=3.3.*,!=3.4.*,!=3.5.*,!=3.6.*"
+files = [
+    {file = "nodeenv-1.7.0-py2.py3-none-any.whl", hash = "sha256:27083a7b96a25f2f5e1d8cb4b6317ee8aeda3bdd121394e5ac54e498028a042e"},
+    {file = "nodeenv-1.7.0.tar.gz", hash = "sha256:e0e7f7dfb85fc5394c6fe1e8fa98131a2473e04311a45afb6508f7cf1836fa2b"},
+]
+
+[package.dependencies]
+setuptools = "*"
+
 [[package]]
 name = "numpy"
 version = "1.24.2"
@@ -549,6 +619,22 @@ files = [
 docs = ["furo", "olefile", "sphinx (>=2.4)", "sphinx-copybutton", "sphinx-inline-tabs", "sphinx-issues (>=3.0.1)", "sphinx-removed-in", "sphinxext-opengraph"]
 tests = ["check-manifest", "coverage", "defusedxml", "markdown2", "olefile", "packaging", "pyroma", "pytest", "pytest-cov", "pytest-timeout"]
 
+[[package]]
+name = "platformdirs"
+version = "3.1.1"
+description = "A small Python package for determining appropriate platform-specific dirs, e.g. a \"user data dir\"."
+category = "main"
+optional = false
+python-versions = ">=3.7"
+files = [
+    {file = "platformdirs-3.1.1-py3-none-any.whl", hash = "sha256:e5986afb596e4bb5bde29a79ac9061aa955b94fca2399b7aaac4090860920dd8"},
+    {file = "platformdirs-3.1.1.tar.gz", hash = "sha256:024996549ee88ec1a9aa99ff7f8fc819bb59e2c3477b410d90a16d32d6e707aa"},
+]
+
+[package.extras]
+docs = ["furo (>=2022.12.7)", "proselint (>=0.13)", "sphinx (>=6.1.3)", "sphinx-autodoc-typehints (>=1.22,!=1.23.4)"]
+test = ["appdirs (==1.4.4)", "covdefaults (>=2.2.2)", "pytest (>=7.2.1)", "pytest-cov (>=4)", "pytest-mock (>=3.10)"]
+
 [[package]]
 name = "pluggy"
 version = "1.0.0"
@@ -565,6 +651,25 @@ files = [
 dev = ["pre-commit", "tox"]
 testing = ["pytest", "pytest-benchmark"]
 
+[[package]]
+name = "pre-commit"
+version = "3.2.0"
+description = "A framework for managing and maintaining multi-language pre-commit hooks."
+category = "main"
+optional = false
+python-versions = ">=3.8"
+files = [
+    {file = "pre_commit-3.2.0-py2.py3-none-any.whl", hash = "sha256:f712d3688102e13c8e66b7d7dbd8934a6dda157e58635d89f7d6fecdca39ce8a"},
+    {file = "pre_commit-3.2.0.tar.gz", hash = "sha256:818f0d998059934d0f81bb3667e3ccdc32da6ed7ccaac33e43dc231561ddaaa9"},
+]
+
+[package.dependencies]
+cfgv = ">=2.0.0"
+identify = ">=1.0.0"
+nodeenv = ">=0.11.1"
+pyyaml = ">=5.1"
+virtualenv = ">=20.10.0"
+
 [[package]]
 name = "pyparsing"
 version = "3.0.9"
@@ -619,6 +724,56 @@ files = [
 [package.dependencies]
 six = ">=1.5"
 
+[[package]]
+name = "pyyaml"
+version = "6.0"
+description = "YAML parser and emitter for Python"
+category = "main"
+optional = false
+python-versions = ">=3.6"
+files = [
+    {file = "PyYAML-6.0-cp310-cp310-macosx_10_9_x86_64.whl", hash = "sha256:d4db7c7aef085872ef65a8fd7d6d09a14ae91f691dec3e87ee5ee0539d516f53"},
+    {file = "PyYAML-6.0-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:9df7ed3b3d2e0ecfe09e14741b857df43adb5a3ddadc919a2d94fbdf78fea53c"},
+    {file = "PyYAML-6.0-cp310-cp310-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:77f396e6ef4c73fdc33a9157446466f1cff553d979bd00ecb64385760c6babdc"},
+    {file = "PyYAML-6.0-cp310-cp310-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:a80a78046a72361de73f8f395f1f1e49f956c6be882eed58505a15f3e430962b"},
+    {file = "PyYAML-6.0-cp310-cp310-manylinux_2_5_x86_64.manylinux1_x86_64.manylinux_2_12_x86_64.manylinux2010_x86_64.whl", hash = "sha256:f84fbc98b019fef2ee9a1cb3ce93e3187a6df0b2538a651bfb890254ba9f90b5"},
+    {file = "PyYAML-6.0-cp310-cp310-win32.whl", hash = "sha256:2cd5df3de48857ed0544b34e2d40e9fac445930039f3cfe4bcc592a1f836d513"},
+    {file = "PyYAML-6.0-cp310-cp310-win_amd64.whl", hash = "sha256:daf496c58a8c52083df09b80c860005194014c3698698d1a57cbcfa182142a3a"},
+    {file = "PyYAML-6.0-cp311-cp311-macosx_10_9_x86_64.whl", hash = "sha256:d4b0ba9512519522b118090257be113b9468d804b19d63c71dbcf4a48fa32358"},
+    {file = "PyYAML-6.0-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:81957921f441d50af23654aa6c5e5eaf9b06aba7f0a19c18a538dc7ef291c5a1"},
+    {file = "PyYAML-6.0-cp311-cp311-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:afa17f5bc4d1b10afd4466fd3a44dc0e245382deca5b3c353d8b757f9e3ecb8d"},
+    {file = "PyYAML-6.0-cp311-cp311-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:dbad0e9d368bb989f4515da330b88a057617d16b6a8245084f1b05400f24609f"},
+    {file = "PyYAML-6.0-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:432557aa2c09802be39460360ddffd48156e30721f5e8d917f01d31694216782"},
+    {file = "PyYAML-6.0-cp311-cp311-win32.whl", hash = "sha256:bfaef573a63ba8923503d27530362590ff4f576c626d86a9fed95822a8255fd7"},
+    {file = "PyYAML-6.0-cp311-cp311-win_amd64.whl", hash = "sha256:01b45c0191e6d66c470b6cf1b9531a771a83c1c4208272ead47a3ae4f2f603bf"},
+    {file = "PyYAML-6.0-cp36-cp36m-macosx_10_9_x86_64.whl", hash = "sha256:897b80890765f037df3403d22bab41627ca8811ae55e9a722fd0392850ec4d86"},
+    {file = "PyYAML-6.0-cp36-cp36m-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:50602afada6d6cbfad699b0c7bb50d5ccffa7e46a3d738092afddc1f9758427f"},
+    {file = "PyYAML-6.0-cp36-cp36m-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:48c346915c114f5fdb3ead70312bd042a953a8ce5c7106d5bfb1a5254e47da92"},
+    {file = "PyYAML-6.0-cp36-cp36m-manylinux_2_5_x86_64.manylinux1_x86_64.manylinux_2_12_x86_64.manylinux2010_x86_64.whl", hash = "sha256:98c4d36e99714e55cfbaaee6dd5badbc9a1ec339ebfc3b1f52e293aee6bb71a4"},
+    {file = "PyYAML-6.0-cp36-cp36m-win32.whl", hash = "sha256:0283c35a6a9fbf047493e3a0ce8d79ef5030852c51e9d911a27badfde0605293"},
+    {file = "PyYAML-6.0-cp36-cp36m-win_amd64.whl", hash = "sha256:07751360502caac1c067a8132d150cf3d61339af5691fe9e87803040dbc5db57"},
+    {file = "PyYAML-6.0-cp37-cp37m-macosx_10_9_x86_64.whl", hash = "sha256:819b3830a1543db06c4d4b865e70ded25be52a2e0631ccd2f6a47a2822f2fd7c"},
+    {file = "PyYAML-6.0-cp37-cp37m-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:473f9edb243cb1935ab5a084eb238d842fb8f404ed2193a915d1784b5a6b5fc0"},
+    {file = "PyYAML-6.0-cp37-cp37m-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:0ce82d761c532fe4ec3f87fc45688bdd3a4c1dc5e0b4a19814b9009a29baefd4"},
+    {file = "PyYAML-6.0-cp37-cp37m-manylinux_2_5_x86_64.manylinux1_x86_64.manylinux_2_12_x86_64.manylinux2010_x86_64.whl", hash = "sha256:231710d57adfd809ef5d34183b8ed1eeae3f76459c18fb4a0b373ad56bedcdd9"},
+    {file = "PyYAML-6.0-cp37-cp37m-win32.whl", hash = "sha256:c5687b8d43cf58545ade1fe3e055f70eac7a5a1a0bf42824308d868289a95737"},
+    {file = "PyYAML-6.0-cp37-cp37m-win_amd64.whl", hash = "sha256:d15a181d1ecd0d4270dc32edb46f7cb7733c7c508857278d3d378d14d606db2d"},
+    {file = "PyYAML-6.0-cp38-cp38-macosx_10_9_x86_64.whl", hash = "sha256:0b4624f379dab24d3725ffde76559cff63d9ec94e1736b556dacdfebe5ab6d4b"},
+    {file = "PyYAML-6.0-cp38-cp38-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:213c60cd50106436cc818accf5baa1aba61c0189ff610f64f4a3e8c6726218ba"},
+    {file = "PyYAML-6.0-cp38-cp38-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:9fa600030013c4de8165339db93d182b9431076eb98eb40ee068700c9c813e34"},
+    {file = "PyYAML-6.0-cp38-cp38-manylinux_2_5_x86_64.manylinux1_x86_64.manylinux_2_12_x86_64.manylinux2010_x86_64.whl", hash = "sha256:277a0ef2981ca40581a47093e9e2d13b3f1fbbeffae064c1d21bfceba2030287"},
+    {file = "PyYAML-6.0-cp38-cp38-win32.whl", hash = "sha256:d4eccecf9adf6fbcc6861a38015c2a64f38b9d94838ac1810a9023a0609e1b78"},
+    {file = "PyYAML-6.0-cp38-cp38-win_amd64.whl", hash = "sha256:1e4747bc279b4f613a09eb64bba2ba602d8a6664c6ce6396a4d0cd413a50ce07"},
+    {file = "PyYAML-6.0-cp39-cp39-macosx_10_9_x86_64.whl", hash = "sha256:055d937d65826939cb044fc8c9b08889e8c743fdc6a32b33e2390f66013e449b"},
+    {file = "PyYAML-6.0-cp39-cp39-macosx_11_0_arm64.whl", hash = "sha256:e61ceaab6f49fb8bdfaa0f92c4b57bcfbea54c09277b1b4f7ac376bfb7a7c174"},
+    {file = "PyYAML-6.0-cp39-cp39-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:d67d839ede4ed1b28a4e8909735fc992a923cdb84e618544973d7dfc71540803"},
+    {file = "PyYAML-6.0-cp39-cp39-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:cba8c411ef271aa037d7357a2bc8f9ee8b58b9965831d9e51baf703280dc73d3"},
+    {file = "PyYAML-6.0-cp39-cp39-manylinux_2_5_x86_64.manylinux1_x86_64.manylinux_2_12_x86_64.manylinux2010_x86_64.whl", hash = "sha256:40527857252b61eacd1d9af500c3337ba8deb8fc298940291486c465c8b46ec0"},
+    {file = "PyYAML-6.0-cp39-cp39-win32.whl", hash = "sha256:b5b9eccad747aabaaffbc6064800670f0c297e52c12754eb1d976c57e4f74dcb"},
+    {file = "PyYAML-6.0-cp39-cp39-win_amd64.whl", hash = "sha256:b3d267842bf12586ba6c734f89d1f5b871df0273157918b0ccefa29deb05c21c"},
+    {file = "PyYAML-6.0.tar.gz", hash = "sha256:68fb519c14306fec9720a2a5b45bc9f0c8d1b9c72adf45c37baedfcd949c35a2"},
+]
+
 [[package]]
 name = "scikit-learn"
 version = "1.2.2"
@@ -828,6 +983,27 @@ files = [
     {file = "typing_extensions-4.5.0.tar.gz", hash = "sha256:5cb5f4a79139d699607b3ef622a1dedafa84e115ab0024e0d9c044a9479ca7cb"},
 ]
 
+[[package]]
+name = "virtualenv"
+version = "20.21.0"
+description = "Virtual Python Environment builder"
+category = "main"
+optional = false
+python-versions = ">=3.7"
+files = [
+    {file = "virtualenv-20.21.0-py3-none-any.whl", hash = "sha256:31712f8f2a17bd06234fa97fdf19609e789dd4e3e4bf108c3da71d710651adbc"},
+    {file = "virtualenv-20.21.0.tar.gz", hash = "sha256:f50e3e60f990a0757c9b68333c9fdaa72d7188caa417f96af9e52407831a3b68"},
+]
+
+[package.dependencies]
+distlib = ">=0.3.6,<1"
+filelock = ">=3.4.1,<4"
+platformdirs = ">=2.4,<4"
+
+[package.extras]
+docs = ["furo (>=2022.12.7)", "proselint (>=0.13)", "sphinx (>=6.1.3)", "sphinx-argparse (>=0.4)", "sphinxcontrib-towncrier (>=0.2.1a0)", "towncrier (>=22.12)"]
+test = ["covdefaults (>=2.2.2)", "coverage (>=7.1)", "coverage-enable-subprocess (>=1)", "flaky (>=3.7)", "packaging (>=23)", "pytest (>=7.2.1)", "pytest-env (>=0.8.1)", "pytest-freezegun (>=0.4.2)", "pytest-mock (>=3.10)", "pytest-randomly (>=3.12)", "pytest-timeout (>=2.1)"]
+
 [[package]]
 name = "wheel"
 version = "0.40.0"
@@ -862,4 +1038,4 @@ testing = ["big-O", "flake8 (<5)", "jaraco.functools", "jaraco.itertools", "more
 [metadata]
 lock-version = "2.0"
 python-versions = "^3.8"
-content-hash = "7d2146077166ad2c9f96811ab7b557bda0bd71c3a051984b67a7bd4dc77fdfa2"
+content-hash = "2ee4c8621bff74145ffdd1231ebe4aa952993e791b8a3162be21074ca7cd5c1e"
diff --git a/projects/example_CFD/example_CFD_config.py b/projects/example_CFD/example_CFD_config.py
index 990ec57a..175bc242 100644
--- a/projects/example_CFD/example_CFD_config.py
+++ b/projects/example_CFD/example_CFD_config.py
@@ -1,12 +1,14 @@
 def set_config(c):
     c.input_path = "data/example_CFD/example_CFD.npz"
-    c.energy_conversion = False
     c.compression_ratio = 10.0
-    c.epochs = 2
+    # c.number_of_columns = 24
+    # c.latent_space_size = 15
+    c.epochs = 25
     c.early_stopping = True
-    c.lr_scheduler = True
-    c.patience = 100
+    c.early_stopping_patience = 100
     c.min_delta = 0
+    c.lr_scheduler = False
+    c.lr_scheduler_patience = 100
     c.model_name = "Conv_AE"
     c.custom_norm = True
     c.l1 = True
@@ -14,8 +16,9 @@ def set_config(c):
     c.RHO = 0.05
     c.lr = 0.001
     c.batch_size = 1
-    c.save_as_root = False
     c.test_size = 0
     c.data_dimension = 2
-    c.extra_compression = False
     c.apply_normalization = False
+    c.extra_compression = False
+    c.intermittent_model_saving = False
+    c.intermittent_saving_patience = 100
diff --git a/projects/example_CMS/example_CMS_config.py b/projects/example_CMS/example_CMS_config.py
index 4ee7ff79..d9bda795 100644
--- a/projects/example_CMS/example_CMS_config.py
+++ b/projects/example_CMS/example_CMS_config.py
@@ -1,22 +1,24 @@
 def set_config(c):
     c.input_path = "data/example_CMS/example_CMS.npz"
-    c.latent_space_size = 15
+    # c.compression_ratio            = 2.0
     c.number_of_columns = 24
-    c.epochs = 2
-    c.early_stopping = False
-    c.lr_scheduler = True
-    c.patience = 50
+    c.latent_space_size = 15
+    c.epochs = 25
+    c.early_stopping = True
+    c.early_stopping_patience = 100
     c.min_delta = 0
-    c.model_name = "george_SAE"
+    c.lr_scheduler = False
+    c.lr_scheduler_patience = 100
+    c.model_name = "AE"
     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.energy_conversion = False
-    c.extra_compression = False
     c.data_dimension = 1
     c.apply_normalization = True
+    c.extra_compression = False
+    c.intermittent_model_saving = False
+    c.intermittent_saving_patience = 100
diff --git a/pyproject.toml b/pyproject.toml
index 4b07eab9..6743ac90 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -12,6 +12,7 @@ tqdm = "^4.64.1"
 matplotlib = "^3.6.2"
 scikit-learn = "^1.2.0"
 pytest = "^7.2.1"
+pre-commit = "^3.2.0"
 
 [tool.pytest.ini_options]
 addopts = [


Test patch:
diff --git a/.github/workflows/test_and_lint.yaml b/.github/workflows/test_and_lint.yaml
index 7cebac2a..b764bd56 100644
--- a/.github/workflows/test_and_lint.yaml
+++ b/.github/workflows/test_and_lint.yaml
@@ -3,30 +3,28 @@ name: Code-Quality
 on: [pull_request, workflow_dispatch]
 
 jobs:
-  test_and_lint:    
+  test_and_lint:
     strategy:
         fail-fast: false
     runs-on: ubuntu-latest
-    
+
     steps:
     - name: "checkout repository"
       uses: actions/checkout@v3
-      
+
     - name: "install poetry"
       uses: abatilo/actions-poetry@v2
       with:
         poetry-version: 1.3.2
-        
+
     - name: "install dependencies"
       run: poetry install
-      
+
     - name: "set python path"
       run: echo "PYTHONPATH=.:baler" >> $GITHUB_ENV
     - name: "run tests"
       run: poetry run pytest
-      
+
     - name: "run black"
-      uses: psf/black@stable
-      with:
-        options: "--check"
-        src: "./baler"
+      run: |
+        black --check baler
diff --git a/tests/test_data_processing.py b/tests/test_data_processing.py
index 1c9e8a8b..7fe3e948 100644
--- a/tests/test_data_processing.py
+++ b/tests/test_data_processing.py
@@ -11,7 +11,7 @@
 
 def test_import_config_success():
     # Call the import_config function with the sample config file path
-    config = helper.configClass
+    config = helper.Config
     config.Foo = "Bar"
     config.Baz = 10