Repository: esa/neuralg (language: python)

Issue:
Try out if GPU is used 



Discussion / hints:

Fix patch:
diff --git a/README.md b/README.md
index 30cae30..8b13789 100644
--- a/README.md
+++ b/README.md
@@ -1,2 +1,1 @@
-# NLAAP
-NeuralLinearAlgebraAPproximators
+
diff --git a/neuralg/__init__.py b/neuralg/__init__.py
index 5e1df9a..4844fbb 100644
--- a/neuralg/__init__.py
+++ b/neuralg/__init__.py
@@ -7,6 +7,7 @@
 from .utils.set_log_level import set_log_level
 from .utils.set_precision import set_precision
 from .utils.clear_loaded_models import clear_loaded_models
+from .utils.set_up_torch import set_up_torch
 
 set_log_level("INFO")
 
@@ -16,11 +17,17 @@
 
 logger.info(f"Initialized neuralg for {os.environ['TORCH_DEVICE']}")
 
-# Set default precision
+# Set default precision to float32 and uses CUDA if initialized
 set_precision()
 
 # Initialize global model handler
 neuralg_ModelHandler = ModelHandler()
 
 
-__all__ = ["eig", "set_precision", "set_log_level", "clear_loaded_models"]
+__all__ = [
+    "eig",
+    "clear_loaded_models",
+    "set_precision",
+    "set_log_level",
+    "set_up_torch",
+]
diff --git a/neuralg/evaluation/evaluate.py b/neuralg/evaluation/evaluate.py
index bf555ab..5da7e41 100644
--- a/neuralg/evaluation/evaluate.py
+++ b/neuralg/evaluation/evaluate.py
@@ -20,11 +20,11 @@ def evaluate(run_cfg):
 
     for j in range(no_models):
         d = ms[j]  # Matrix dimension
-        model = run_cfg[str(d)].model
-        run_cfg[str(d)] = deepcopy(_load_test_settings(run_cfg[str(d)]))
+        model = run_cfg[d].model
+        run_cfg[d] = deepcopy(_load_test_settings(run_cfg[d]))
         # Initialize a DotMap for every different test set setting
-        run_cfg[str(d)].test_results = DotMap()
-        temp_test_settings = deepcopy(run_cfg[str(d)].test_cfg.test_settings)
+        run_cfg[d].test_results = DotMap()
+        temp_test_settings = deepcopy(run_cfg[d].test_cfg.test_settings)
 
         for i, k in enumerate(temp_test_settings):
             test_parameters = temp_test_settings[k]
@@ -32,10 +32,10 @@ def evaluate(run_cfg):
 
             results = evaluate_model(model, test_parameters)
             acc = []
-            for tol in run_cfg[str(d)].test_cfg.tolerances:
+            for tol in run_cfg[d].test_cfg.tolerances:
                 acc.append(compute_accuracy(tol, results).item())
 
-            run_cfg[str(d)].test_results[k] = acc
+            run_cfg[d].test_results[k] = acc
 
     return run_cfg
 
diff --git a/neuralg/io/load_model.py b/neuralg/io/load_model.py
index c4839d3..9a15067 100644
--- a/neuralg/io/load_model.py
+++ b/neuralg/io/load_model.py
@@ -18,8 +18,11 @@ def load_model(model_name):
     available_models = DotMap()
     # Added some trained nerf models for first minimal module
     for d in range(NEURALG_MIN_MATRIX_SIZE, NEURALG_MAX_MATRIX_SIZE + 1):
-        state_dict_path = os.path.join(
-            os.path.dirname(__file__), "../models/saved_models/eigval{}.pt".format(d),
+        state_dict_path = os.path.realpath(
+            os.path.join(
+                os.path.dirname(__file__),
+                "../models/saved_models/eigval{}.pt".format(d),
+            )
         )
         available_models["eigval{}".format(d)] = [state_dict_path, "nerf", d]
 
diff --git a/neuralg/models/nerf.py b/neuralg/models/nerf.py
index 639c9b2..204d72f 100644
--- a/neuralg/models/nerf.py
+++ b/neuralg/models/nerf.py
@@ -1,7 +1,6 @@
 # From https://github.com/darioizzo/geodesyNets/blob/master/gravann/networks/_nerf.py,
 # Implementation of architecture from "NeRF: Representing Scenes as
 # Neural Radiance Fields for View Synthesis" , https://arxiv.org/pdf/2003.08934.pdf
-
 import torch
 import torch.nn as nn
 
@@ -43,7 +42,10 @@ def __init__(
             self.out_features = matrix_dimension  # For e.g. full eigval problem
 
         self.net = nn.ModuleList()
-        self.flatten = nn.Flatten(start_dim=1)
+        # Added this for more robust against different batch dimensions
+        self.flatten_batch = nn.Flatten(start_dim=0, end_dim=-3)
+        self.flatten = nn.Flatten(start_dim=-2)
+
         self.net.append(NERFLayer(in_features, n_neurons))
 
         for i in range(hidden_layers):
@@ -52,21 +54,28 @@ def __init__(
             else:
                 self.net.append(NERFLayer(n_neurons, n_neurons))
 
-        self.net.append(nn.Linear(n_neurons, self.out_features))
+        self.net.append(nn.Linear(n_neurons, matrix_dimension))
 
     def forward(self, x):
-        x = self.flatten(x)
+        batch_dim = x.shape[0:-2]
+
+        # Is there a smarter way to deal with different batch dimensions?
+        if len(batch_dim) > 1:
+            x = self.flatten_batch(x)
+
+        x_flat = self.flatten(x)
+
         # save for skip connection
-        identity = x
+        identity = x_flat
 
         # compute first layer
-        out = self.net[0].forward(x)
+        out = self.net[0].forward(x_flat)
 
         # compute all other layers and apply skip where requested
         for layer_idx in range(1, len(self.net)):
             out = self.net[layer_idx].forward(out)
             if layer_idx in self.skip:
-                out = torch.cat([out, identity], dim=1)
-
-        return out[:, None, :]  # Dummy index to fit framework
-
+                out = torch.cat([out, identity], dim=-1)
+        if len(batch_dim) > 1:
+            out = nn.Unflatten(0, batch_dim)(out)
+        return out
diff --git a/neuralg/models/saved_models/__init__.py b/neuralg/models/saved_models/__init__.py
new file mode 100644
index 0000000..e69de29
diff --git a/neuralg/models/saved_models/eigval11.pt b/neuralg/models/saved_models/eigval11.pt
new file mode 100644
index 0000000..6210e9b
Binary files /dev/null and b/neuralg/models/saved_models/eigval11.pt differ
diff --git a/neuralg/models/saved_models/eigval12.pt b/neuralg/models/saved_models/eigval12.pt
new file mode 100644
index 0000000..49b9402
Binary files /dev/null and b/neuralg/models/saved_models/eigval12.pt differ
diff --git a/neuralg/models/saved_models/eigval13.pt b/neuralg/models/saved_models/eigval13.pt
new file mode 100644
index 0000000..0123870
Binary files /dev/null and b/neuralg/models/saved_models/eigval13.pt differ
diff --git a/neuralg/models/saved_models/eigval14.pt b/neuralg/models/saved_models/eigval14.pt
new file mode 100644
index 0000000..e227464
Binary files /dev/null and b/neuralg/models/saved_models/eigval14.pt differ
diff --git a/neuralg/models/saved_models/eigval15.pt b/neuralg/models/saved_models/eigval15.pt
new file mode 100644
index 0000000..4909954
Binary files /dev/null and b/neuralg/models/saved_models/eigval15.pt differ
diff --git a/neuralg/ops/eig.py b/neuralg/ops/eig.py
index b1a7d28..710c64e 100644
--- a/neuralg/ops/eig.py
+++ b/neuralg/ops/eig.py
@@ -1,7 +1,6 @@
-import numpy as np
 import torch
-from loguru import logger
 import neuralg
+from ..utils.constants import NEURALG_MIN_MATRIX_SIZE, NEURALG_MAX_MATRIX_SIZE
 
 # Potentially, we should perhaps call this eigvals, since it only computes eigenvalues
 # Also note that is is only trained on symmetric matrices and will only output real eigenvalues.
@@ -10,21 +9,12 @@ def eig(A):
     """Approximates (real) eigenvalues of a square matrix. 
     Supports batches of matrices, and if A is a batch of matrices then the output has the same batch dimensions. 
     Supports input of float and double dtypes.
-
     Args:
         A (tensor): Tensor of shape [*,d,d] where * can be zero or more batch dimensions.
-
     Returns:
         tensor: Containing the real-valued eigenvalue approximations to A. If A is a n-dimensional, resulting output is n-1 dimensional with the same batch dimension.
     """
-    try:
-        _validate_input(A)
-    except ValueError:
-        return None
-    # if type(A) == np.ndarray:
-    #     logger.info("Creating a tensor from input numpy ndarray ")
-    #     A = torch.from_numpy(A)
-    #     print(A.dtype)
+    _validate_input(A)
     # Load the right model via model handler
     matrix_size = A.shape[-1]
     model = neuralg.neuralg_ModelHandler.get_model("eigval", matrix_size)
@@ -35,43 +25,41 @@ def eig(A):
 
 def _predict(model, A):
     """Model prediction that allows keeping input dimension
-
     Args:
         model (torch.nn): Model to perform eigenvalue prediction
         A (tensor): Tensor of shape [*,d,d] where * can be zero or more batch dimension
-
     Returns:
         tensor: Containing eigenvalue approximations, of shape [*,d]
     """
-    input_shape = A.shape
-    if len(input_shape) == 2:
-        # Add dummy dimension
-        A = A[None, :]
-        assert len(A.shape) == 3
-        # Returns a one dimensional tensor with
-        out = model(A).squeeze()
-    elif len(input_shape) == 3:
-        out = model(A).squeeze(1)
-    elif len(input_shape) == 4 and input_shape[1] == 1:
-        out = model(A)
-    else:
-        raise NotImplementedError(
-            "Prediction on this input dimension is not implemented yet"
-        )
+    out = model(A)
     return out
 
 
 def _validate_input(input):
-    """ Checks that the input has correct shape. If not
+    """ Checks that the input has correct shape.
     Args:
         input (tensor): Batch to be validated for eigenvalue approximation
-
     Raises:
         ValueError: If matrices are not quadratic
         ValueError: If input is not at least two dimensional
     """
-    if len(input.shape) < 2:
-        raise ValueError("Input must be at least two dimensional, but had shape" +  str(input.shape))
+    if not torch.is_tensor(input):
+        raise ValueError("Only torch tensors are supported as input")
 
+    if len(input.shape) < 2:
+        raise ValueError(
+            "Input must be at least two dimensional, but had shape" + str(input.shape)
+        )
     if input.shape[-2] != input.shape[-1]:
-        raise ValueError("Matrices must be quadratic", but had shape" +  str(input.shape))
+        raise ValueError("Matrices must be quadratic but had shape" + str(input.shape))
+
+    if (
+        input.shape[-1] < NEURALG_MIN_MATRIX_SIZE
+        or input.shape[-1] > NEURALG_MAX_MATRIX_SIZE
+    ):
+        raise ValueError(
+            "Matrix dimension must be between {} and {}, but had dimension"
+            + str(input.shape[-1]).format(
+                NEURALG_MIN_MATRIX_SIZE, NEURALG_MAX_MATRIX_SIZE
+            )
+        )
diff --git a/neuralg/plots/plot_eigval_results.py b/neuralg/plots/plot_eigval_results.py
index b693776..978b96d 100644
--- a/neuralg/plots/plot_eigval_results.py
+++ b/neuralg/plots/plot_eigval_results.py
@@ -11,7 +11,7 @@ def plot_eigval_results(run_cfg, save_path=None):
     ms = run_cfg.matrix_sizes
     no_plots = len(ms)
     for j in range(no_plots):
-        d = str(ms[j])
+        d = ms[j]
         fig = plt.figure(figsize=(14, 6), dpi=150)
         fig.patch.set_facecolor("white")
         ax = fig.add_subplot()
diff --git a/neuralg/plots/plot_loss_logs.py b/neuralg/plots/plot_loss_logs.py
index 0537369..7ced203 100644
--- a/neuralg/plots/plot_loss_logs.py
+++ b/neuralg/plots/plot_loss_logs.py
@@ -11,7 +11,7 @@ def plot_loss_logs(run_cfg, save_path=None):
     ms = run_cfg.matrix_sizes
     no_plots = len(ms)
     for j in range(no_plots):
-        d = str(ms[j])
+        d = ms[j]
         fig = plt.figure(figsize=(14, 6), dpi=150)
         fig.patch.set_facecolor("white")
         k = run_cfg[d].run_params.iterations
diff --git a/neuralg/training/RandomMatrixDataSet.py b/neuralg/training/RandomMatrixDataSet.py
index ef66e72..bd2ea72 100644
--- a/neuralg/training/RandomMatrixDataSet.py
+++ b/neuralg/training/RandomMatrixDataSet.py
@@ -1,3 +1,4 @@
+from cmath import sqrt
 import torch
 from neuralg.training.RandomSymmetricMatrix import RandomSymmetricMatrix
 
@@ -15,8 +16,8 @@ def __init__(self, N, d, operation):
     def from_rand(self, r1=-10, r2=10):
         self.X = (r1 - r2) * torch.rand(self.N, 1, self.d, self.d) + r2
 
-    def from_randn(self):
-        self.X = torch.randn(self.N, 1, self.d, self.d)
+    def from_randn(self, sigma=10 / sqrt(3)):
+        self.X = sigma * torch.randn(self.N, 1, self.d, self.d)
 
     def from_dist(self, dist):
         self.X = RandomSymmetricMatrix(N=self.N, d=self.d, dist=dist).X
diff --git a/neuralg/training/get_sample.py b/neuralg/training/get_sample.py
index 3560043..081945e 100644
--- a/neuralg/training/get_sample.py
+++ b/neuralg/training/get_sample.py
@@ -12,18 +12,23 @@ def get_sample(matrix_parameters):
         RandomMatrixDataSet: Batch of random matrices with given matrix parameters
     """
     # Instantiate batch
-    N, d = matrix_parameters["N"], matrix_parameters["d"]
+    N, dim = matrix_parameters["N"], matrix_parameters["d"]
     if "operation" in matrix_parameters:
         op = matrix_parameters["operation"]
-        M = RandomMatrixDataSet(N, d, op)
+        M = RandomMatrixDataSet(N, dim, op)
     else:
-        M = RandomMatrixDataSet(N, d)
+        M = RandomMatrixDataSet(N, dim)
 
     if "dist" in matrix_parameters:
         M.from_dist(matrix_parameters["dist"])
     else:
-        # Otherwise just sample a matrix with uniform iid entries
-        M.from_rand()  # M.from_randn() #fix so one can choose between these
+        if "normal" in matrix_parameters:
+            # Create matrices with iid gaussian entries
+            if matrix_parameters["normal"]:
+                M.from_randn()
+        else:
+            # Otherwise just sample a matrix with uniform iid entries
+            M.from_rand()  # M.from_randn() #fix so one can choose between these
         if "wigner" in matrix_parameters:
             # Create Wigner matrix
             if matrix_parameters["wigner"]:
diff --git a/neuralg/training/train_models.py b/neuralg/training/train_models.py
index 3f764a3..7ca9286 100644
--- a/neuralg/training/train_models.py
+++ b/neuralg/training/train_models.py
@@ -35,7 +35,7 @@ def train_models(cfg, save_training_run=False):
             raise NotImplementedError("Model type is not available for training")
         temp_config.model = model
         logger.trace("Calling training for matrix size {}".format(d))
-        run_results[str(d)] = run_training(temp_config)
+        run_results[d] = run_training(temp_config)
 
     logger.trace("Finalized training all requested models...")
     # For plotting
diff --git a/neuralg/training/utils/__init__.py b/neuralg/training/utils/__init__.py
new file mode 100644
index 0000000..e69de29
diff --git a/neuralg/training/utils/sorting.py b/neuralg/training/utils/sorting.py
new file mode 100644
index 0000000..30a3bac
--- /dev/null
+++ b/neuralg/training/utils/sorting.py
@@ -0,0 +1,33 @@
+import torch
+
+
+def real_sort(tensor, dim=-1):
+    """Sorts the elements of a potentially complex valued tensor along a given dimension in ascending order by real part.
+        If dim is not given, the last dimension of the input is chosen.
+        If descending is True then the elements are sorted in descending order.
+   Args:
+        tensor (tensor): input tensor to sort.
+        dim (int, optional): Dimension to sort along. Defaults to -1.
+
+    Returns:
+        tensor: Containing the sorted values of the elements in the original input tensor
+    """
+    indices = torch.sort(torch.real(tensor))[1]
+    t = torch.gather(tensor, dim=dim, index=indices)
+    return t
+
+
+def abs_sort(tensor, dim=-1, descending=False):
+    """Sorts the elements of a potentially complex valued tensor along a given dimension in ascending order by absolute value.
+        If dim is not given, the last dimension of the input is chosen.
+        If  descending is True then the elements are sorted in descending order.
+    Args:
+        tensor (tensor): input tensor to sort.
+        dim (int, optional): Dimension to sort along. Defaults to -1.
+
+    Returns:
+        tensor: Containing the sorted values of the elements in the original input tensor
+    """
+    indices = torch.sort(tensor.abs(), descending=descending)[1]
+    t = torch.gather(tensor, dim=dim, index=indices)
+    return t
diff --git a/neuralg/utils/ModelHandler.py b/neuralg/utils/ModelHandler.py
index 64a4467..cda96f3 100644
--- a/neuralg/utils/ModelHandler.py
+++ b/neuralg/utils/ModelHandler.py
@@ -59,9 +59,8 @@ def clear_loaded_models(self):
         """ Frees allocated memory from loaded models
         """
         logger.info("Clearing loaded models")
-        loaded_models_copy = deepcopy(self.loaded_models)
+        loaded_models_copy = deepcopy(list(self.loaded_models.keys()))
         for model in loaded_models_copy:
             del self.loaded_models[model]
         del loaded_models_copy  # Not sure if needed
-        None
 
diff --git a/neuralg/utils/constants.py b/neuralg/utils/constants.py
index 659a448..dcbc194 100644
--- a/neuralg/utils/constants.py
+++ b/neuralg/utils/constants.py
@@ -1,4 +1,4 @@
 NEURALG_MIN_MATRIX_SIZE = 3
-NEURALG_MAX_MATRIX_SIZE = 10
+NEURALG_MAX_MATRIX_SIZE = 15
 
 NEURALG_SUPPORTED_OPERATIONS = ["eig"]
diff --git a/neuralg/utils/enable_cuda.py b/neuralg/utils/enable_cuda.py
new file mode 100644
index 0000000..a79ec99
--- /dev/null
+++ b/neuralg/utils/enable_cuda.py
@@ -0,0 +1,24 @@
+from loguru import logger
+
+from .set_precision import set_precision
+
+
+def enable_cuda(data_type="float32"):
+    """This function sets torch's default device to CUDA if possible. Call before declaring any variables!
+    The default precision can be set here initially, or using set_precision later.
+    Args:
+        data_type ("float32", "float64" or None, optional): Data type to use. If None, skip the call to set_precision. Defaults to "float32".
+    """
+    import torch
+
+    if torch.cuda.is_available():
+        logger.info("PyTorch VERSION: " + str(torch.__version__))
+        logger.info("CUDNN VERSION: " + str(torch.backends.cudnn.version()))
+        logger.info("Number of CUDA Devices: " + str(torch.cuda.device_count()))
+        logger.info("Active CUDA Device: GPU" + str(torch.cuda.current_device()))
+        if data_type is not None:
+            set_precision(data_type)
+    else:
+        logger.warning(
+            "Error enabling CUDA. cuda.is_available() returned False. CPU will be used."
+        )
diff --git a/neuralg/utils/set_up_torch.py b/neuralg/utils/set_up_torch.py
new file mode 100644
index 0000000..08f17bf
--- /dev/null
+++ b/neuralg/utils/set_up_torch.py
@@ -0,0 +1,20 @@
+from .set_precision import set_precision
+from .enable_cuda import enable_cuda
+
+
+def set_up_torch(data_type=None, torch_enable_cuda=True):
+    """ Configure device and potentially precision for neuralg backend.
+    Calls neuralg.enable_cuda unless torch_enable_cuda is False.
+    If a data type is passed, set the default floating point precision with neuralg.set_precision.
+    Args:
+        data_type ("float32", "float64" or None, optional): Data type which is passed to set_precision. If None, do not call set_precision except if CUDA is enabled for torch. Defaults to None.
+        torch_enable_cuda (Bool, optional): If True and backend is "torch", call enable_cuda. Defaults to True.
+    """
+    if torch_enable_cuda:
+        if data_type is None:
+            enable_cuda()
+        else:
+            # Do not call set_precision twice.
+            enable_cuda(data_type=None)
+    if data_type is not None:
+        set_precision(data_type)
diff --git a/notebooks/example_problem.ipynb b/notebooks/example_problem.ipynb
new file mode 100644
index 0000000..63eeaf3
--- /dev/null
+++ b/notebooks/example_problem.ipynb
@@ -0,0 +1,376 @@
+{
+    "cells": [
+        {
+            "cell_type": "markdown",
+            "metadata": {},
+            "source": [
+                "# Demo of ```neuralg``` module\n",
+                "\n",
+                "## Spectral graph theory application\n",
+                "\n",
+                "In spectral graph theory, graph properties are studied in light of the eigendecompositions of their associated matrices. For instance, graph characteristics can be inferred from the eigenvalue of the adjacency matrix $A$ of a graph, i.e \n",
+                "$$ A_{ij} = \\begin{cases} 1 \\; \\text{if there is an edge from $i$ to $j$} \\\\ 0 \\; \\text{otherwise} \\end{cases}$$\n",
+                "A more detailed read can be found e.g. [here](https://en.wikipedia.org/wiki/Spectral_graph_theory).\n",
+                "## A machine learning task using the neuralg module\n",
+                "To carry out automatic differentation to allow, for example, gradient-based optimization, the eigenvalue approximations must not break the gradient flow. Training a neural network with PyTorch will serve as a demonstration that the `neuralg` module indeed preserves gradients. Lets say we want a neural network that approximates the total number of length $k$ cycles starting from any node in a graph. In exact arithmetic, this quantity is given by \n",
+                "\n",
+                "$$trace(A^k),$$\n",
+                "\n",
+                "where $A$ is the adjacency matrix of the graph, but we do not want to perform the $k$ matrix multiplications. From the properties of the trace operator, we can relate this quantity to  the spectrum of $A$ such that  \n",
+                "\n",
+                "\\begin{equation} \\# \\text{ cycles of length } k = \\sum_{i} \\lambda_i(A)^k \n",
+                "\\end{equation}\n",
+                "\n",
+                "\n",
+                "To this end, we can use `neuralg.eig()` to approximate the ground truths in the supervised learning of predicting this quantity.  "
+            ]
+        },
+        {
+            "cell_type": "code",
+            "execution_count": null,
+            "metadata": {},
+            "outputs": [],
+            "source": [
+                "import sys \n",
+                "sys.path.append("..") \n",
+                "import neuralg \n",
+                "from neuralg.evaluation.compute_accuracy import compute_accuracy\n",
+                "from neuralg.training.losses import relative_L1_evaluation_error\n",
+                "import torch\n",
+                "import networkx # additional requirement for processing graphs \n",
+                "from copy import deepcopy\n",
+                "import time\n",
+                "import matplotlib.pyplot as plt"
+            ]
+        },
+        {
+            "cell_type": "markdown",
+            "metadata": {},
+            "source": [
+                "### `neuralg` allows for setting float precision and can potentially be employed on a GPU, if available."
+            ]
+        },
+        {
+            "cell_type": "code",
+            "execution_count": null,
+            "metadata": {},
+            "outputs": [],
+            "source": [
+                "neuralg.set_up_torch(torch_enable_cuda=True)\n",
+                "neuralg.set_precision(\"float64\")"
+            ]
+        },
+        {
+            "cell_type": "markdown",
+            "metadata": {},
+            "source": [
+                "### Define a simple convolutional net for the regression task. A convolutional block is followed by two dense layers."
+            ]
+        },
+        {
+            "cell_type": "code",
+            "execution_count": null,
+            "metadata": {},
+            "outputs": [],
+            "source": [
+                "import torch.nn as nn \n",
+                "import torch.nn.functional as F\n",
+                "\n",
+                "class CycleCNN(nn.Module):\n",
+                "    def __init__(self, n_graph_nodes, conv_layers, filters, kernel_size):\n",
+                "        super(CycleCNN, self).__init__()\n",
+                "        self.net = []\n",
+                "        self.n_graph_nodes = n_graph_nodes\n",
+                "        self.net.append(nn.Conv2d(1,filters,kernel_size, padding = \"same\"))\n",
+                "        self.net.append(nn.BatchNorm2d(filters))\n",
+                "        self.net.append(nn.ReLU())\n",
+                "        for i in range(conv_layers-1):\n",
+                "            self.net.append(nn.Conv2d(filters,filters,kernel_size, padding = \"same\"))\n",
+                "            self.net.append(nn.BatchNorm2d(filters))\n",
+                "            self.net.append(nn.ReLU())\n",
+                "        \n",
+                "        self.net.append(nn.Conv2d(filters,1,kernel_size, padding = \"same\"))\n",
+                "        self.net.append(nn.Flatten())\n",
+                "        self.net.append(DenseLayer(n_graph_nodes**2,n_graph_nodes))\n",
+                "        self.net.append(DenseLayer(n_graph_nodes,1, is_last = True))\n",
+                "        self.net = nn.Sequential(*self.net)\n",
+                "    \n",
+                "    def forward(self, x):\n",
+                "        out = self.net(x)\n",
+                "        return out\n",
+                "\n",
+                "class DenseLayer(nn.Module):\n",
+                "\n",
+                "    def __init__(self, in_features, out_features, bias=True, is_last = False):\n",
+                "        super().__init__()\n",
+                "        self.is_last = is_last\n",
+                "        self.in_features = in_features\n",
+                "        self.linear = nn.Linear(in_features, out_features, bias=bias)\n",
+                "\n",
+                "    def forward(self, input):\n",
+                "        if self.is_last: \n",
+                "            return self.linear(input)\n",
+                "        else:\n",
+                "            return F.relu(self.linear(input))"
+            ]
+        },
+        {
+            "cell_type": "markdown",
+            "metadata": {},
+            "source": [
+                "#### Define graph generation properties and target cycle length and necessary training parameters"
+            ]
+        },
+        {
+            "cell_type": "code",
+            "execution_count": null,
+            "metadata": {},
+            "outputs": [],
+            "source": [
+                "#Data parameters \n",
+                "n_graph_nodes = 5 # Number of nodes in the graph, defining the size of the adjacancy matrices\n",
+                "k = 5 # Cycle length of interest\n",
+                "p = 0.5 # Probability of two nodes having a connecting edge \n",
+                "\n",
+                "\n",
+                "#Training parameters\n",
+                "iterations = 2000 # Iterations per epoch\n",
+                "batch_size = 32 # Number of matrices per forward pass\n",
+                "epochs = 5\n",
+                "criterion = nn.MSELoss() # We use a mean square error loss "
+            ]
+        },
+        {
+            "cell_type": "code",
+            "execution_count": null,
+            "metadata": {},
+            "outputs": [],
+            "source": [
+                "def get_adjacency_batch(batch_size,n_graph_nodes, fixed_seed = False):\n",
+                "    \"\"\" Simple generator of random graph adjacency matrices.       \n",
+                "    Args:\n",
+                "        batch_size (int): Size of batch to be generated\n",
+                "        n_graph_nodes (int): Corresponds to size of the generaed adjacency matrices. \n",
+                "        fixed_seed (bool, optional): If true, the seed is fixed in the sampling. Defaults to False.\n",
+                "\n",
+                "    Returns:\n",
+                "        tensor: batch of adjacency matrices, of shape [batchsize,1,n_graph_nodes,n_graph_nodes] \n",
+                "    \"\"\"\n",
+                "    if fixed_seed:\n",
+                "        torch.manual_seed(1)\n",
+                "    for i in range(batch_size):\n",
+                "        g = networkx.erdos_renyi_graph(n=n_graph_nodes, p= 0.5)\n",
+                "        adj_matrix = torch.tensor(networkx.to_numpy_array(g))[None,:]\n",
+                "        if i == 0: \n",
+                "            A = adj_matrix\n",
+                "        else:   \n",
+                "            A = torch.cat((A,adj_matrix))\n",
+                "        \n",
+                "    return A[:,None,:] # broadcasts to an extra channel dimension "
+            ]
+        },
+        {
+            "cell_type": "code",
+            "execution_count": null,
+            "metadata": {},
+            "outputs": [],
+            "source": [
+                "def train(model: nn.Module, use_neuralg=True) -> None:\n",
+                "    model.train()  # turn on train mode\n",
+                "    total_loss = 0.\n",
+                "    log_interval = 1000\n",
+                "    start_time = time.time()\n",
+                "\n",
+                "    for i in range(iterations):\n",
+                "\n",
+                "        A = get_adjacency_batch(batch_size, n_graph_nodes=n_graph_nodes) #Sample a batch\n",
+                "\n",
+                "        output = model(A) #Predict # of k-cycles with model\n",
+                "\n",
+                "    \n",
+                "        if use_neuralg:\n",
+                "            target_eigvals = neuralg.eig(A) #Use neuralg module to compute ground truth\n",
+                "        else:\n",
+                "            target_eigvals = torch.linalg.eigvalsh(A) #Or, use torch built-in \n",
+                "        \n",
+                "        target = torch.pow(target_eigvals,k).sum(-1) #The target is the quantiti in Eq.(1)\n",
+                "        \n",
+                "        loss = criterion(output,target) # Compute loss\n",
+                "    \n",
+                "        optimizer.zero_grad()\n",
+                "        \n",
+                "        loss.backward()\n",
+                "        \n",
+                "        optimizer.step()\n",
+                "\n",
+                "        total_loss += loss.item()\n",
+                "        \n",
+                "        if i % log_interval == 0 and i > 0:\n",
+                "            lr = scheduler.get_last_lr()[0]\n",
+                "            ms_per_batch = (time.time() - start_time) * 1000 / log_interval\n",
+                "            cur_loss = total_loss / log_interval\n",
+                "\n",
+                "            print(f'| epoch {epoch:3d} | {i:5d} batches | '\n",
+                "                  f'lr {lr:02.4f} | ms/batch {ms_per_batch:5.2f} | '\n",
+                "                  f'loss {cur_loss:5.5f}')\n",
+                "\n",
+                "            total_loss = 0\n",
+                "            start_time = time.time()\n",
+                "    return ms_per_batch\n",
+                "    "
+            ]
+        },
+        {
+            "cell_type": "markdown",
+            "metadata": {},
+            "source": [
+                "### Training a network \n",
+                "##### For reference, two identically initialized models are trained using either `neuralg.eig()` or `torch.linalg.eigvalsh()` to compute the eigenvalue ground truths. The loss is computed between the model outputs and the target function of the eigenvalues from Eq.(1), i.e. the analytical solution to number of total paths of length $k$. "
+            ]
+        },
+        {
+            "cell_type": "code",
+            "execution_count": null,
+            "metadata": {},
+            "outputs": [],
+            "source": [
+                "trained_models = []\n",
+                "for use_neuralg in [True,False]:\n",
+                "    model = CycleCNN(n_graph_nodes, conv_layers = 3,filters = 32,kernel_size=3) #Instantiate model\n",
+                "    lr = 3e-4 # learning rate\n",
+                "    optimizer = torch.optim.Adam(model.parameters(), lr=lr)\n",
+                "    scheduler = torch.optim.lr_scheduler.StepLR(optimizer, 1.0, gamma=0.95)\n",
+                "    model.train(); # turn on train modeepochs = 1\n",
+                "    \n",
+                "    for epoch in range(1, epochs + 1):\n",
+                "        epoch_start_time = time.time()\n",
+                "        ms_per_batch = train(model,use_neuralg)\n",
+                "        scheduler.step()\n",
+                "    trained_models.append([deepcopy(model),ms_per_batch])"
+            ]
+        },
+        {
+            "cell_type": "markdown",
+            "metadata": {},
+            "source": [
+                "#### Evaluation from training\n",
+                "When evaluating on the test set, we consider a predicted output $\\hat{y}$ to be a correct solution to the problem if it approximates the ground truth $y$ to a given tolerance $\\tau$. Specifically, we check that $|y-\\hat{y}| \\leq \\tau|y|$ is satisfied. For completeness, the test set ground truths are calculated both using `neuralg.eig()` and `torch.eigvalsh()`."
+            ]
+        },
+        {
+            "cell_type": "code",
+            "execution_count": null,
+            "metadata": {},
+            "outputs": [],
+            "source": [
+                "eval_size = 10000 # Size of evaluation set \n",
+                "eval_set = get_adjacency_batch(eval_size, n_graph_nodes, fixed_seed=True) #Sample a test set\n",
+                "\n",
+                "#Since path length should be an integer, we round the final predictions with the two trained models\n",
+                "neuralg_pred = torch.round(trained_models[0][0](eval_set)) # Model trained with neuralg ground truths\n",
+                "torch_pred = torch.round(trained_models[1][0](eval_set)) # Model trained with torch ground truths\n",
+                "\n",
+                "torch_targets = torch.round(torch.pow(torch.linalg.eigvalsh(eval_set),k).sum(-1)) # We evaluation eigenvalues with torch built-int\n",
+                "neuralg_targets = torch.round(torch.pow(neuralg.eig(eval_set),k).sum(-1))\n",
+                "\n",
+                "#Model trained with neuralg \n",
+                "neuralg_to_neuralg_eval_errors= relative_L1_evaluation_error(neuralg_pred,neuralg_targets)\n",
+                "neuralg_to_torch_eval_errors= relative_L1_evaluation_error(neuralg_pred,torch_targets)\n",
+                "\n",
+                "#Model trained with torch\n",
+                "torch_to_neuralg_eval_errors = relative_L1_evaluation_error(torch_pred,neuralg_targets) \n",
+                "torch_to_torch_eval_errors = relative_L1_evaluation_error(torch_pred,torch_targets) \n",
+                "\n",
+                "tolerances = [0.1,0.08,0.06,0.05,0.025,0.02,0.01,0.005,0.0001]\n",
+                "n_neuralg_acc = []\n",
+                "n_torch_acc = []\n",
+                "t_neuralg_acc = []\n",
+                "t_torch_acc = []\n",
+                "for tol in tolerances:\n",
+                "    n_neuralg_acc.append(compute_accuracy(tol,neuralg_to_neuralg_eval_errors).item())\n",
+                "    n_torch_acc.append(compute_accuracy(tol,neuralg_to_torch_eval_errors).item())\n",
+                "    t_neuralg_acc.append(compute_accuracy(tol,torch_to_neuralg_eval_errors).item())\n",
+                "    t_torch_acc.append(compute_accuracy(tol,torch_to_torch_eval_errors).item())"
+            ]
+        },
+        {
+            "cell_type": "code",
+            "execution_count": null,
+            "metadata": {},
+            "outputs": [],
+            "source": [
+                "fig = plt.figure(figsize=(14, 6), dpi=150)\n",
+                "fig.patch.set_facecolor(\"white\")\n",
+                "ax = fig.add_subplot(1,2,1)\n",
+                "ax.set_title(\"neuralg trained, {:.3} ms per batch. \\n Test accuracy versus tolerance, {}x{}\".format(trained_models[0][1],n_graph_nodes, n_graph_nodes), fontsize=18)\n",
+                "ax.plot(tolerances,n_neuralg_acc,label=\"neuralg evaluation ground truths\")\n",
+                "ax.plot(tolerances,n_torch_acc,label=\"torch evaluation ground truths\")\n",
+                "\n",
+                "ax.legend(fontsize=14)\n",
+                "ax.set_xlabel(\"$\\\\tau $\", fontsize=18)\n",
+                "ax.set_ylabel(\"Accuracy\", fontsize=16);\n",
+                "\n",
+                "ax = fig.add_subplot(1,2,2)\n",
+                "ax.set_title(\"torch trained, {:.3} ms per batch.\\n Test accuracy versus tolerance, {}x{}\".format(trained_models[1][1],n_graph_nodes, n_graph_nodes), fontsize=18)\n",
+                "ax.plot(tolerances,t_neuralg_acc,label=\"neuralg evaluation ground truths\")\n",
+                "ax.plot(tolerances,t_torch_acc,label=\"torch evaluation ground truths\")\n",
+                "\n",
+                "ax.legend(fontsize=14)\n",
+                "ax.set_xlabel(\"$\\\\tau $\", fontsize=18)\n",
+                "ax.set_ylabel(\"Accuracy\", fontsize=16);"
+            ]
+        },
+        {
+            "cell_type": "markdown",
+            "metadata": {},
+            "source": [
+                "### Conclusion \n",
+                "- In evaluation, the model trained with `neuralg` eigenvalues performs well also on the more accurate torch evaluatiom ground truths. \n",
+                "- Training with `neuralg` or `torch` modules seem to yield similarly performing models. \n",
+                "- When run On CPU there is no synchronization needed and for these small matrices, the computational cost of the numerical methods is very small.  For larger matrices and on the GPU, we expect `neuralg`to have an edge. "
+            ]
+        },
+        {
+            "cell_type": "markdown",
+            "metadata": {},
+            "source": [
+                "#### After using the module, we can clear the loaded models to free the allocated memory"
+            ]
+        },
+        {
+            "cell_type": "code",
+            "execution_count": null,
+            "metadata": {},
+            "outputs": [],
+            "source": [
+                "neuralg.clear_loaded_models()\n",
+                "assert neuralg.neuralg_ModelHandler.loaded_models == {}"
+            ]
+        }
+    ],
+    "metadata": {
+        "interpreter": {
+            "hash": "33b13c0f5861550def7c5fd31dfce4700e1b96089b8cc03c6412b21d90f85adf"
+        },
+        "kernelspec": {
+            "display_name": "PyCharm (pythonProject1)",
+            "language": "python",
+            "name": "python3"
+        },
+        "language_info": {
+            "codemirror_mode": {
+                "name": "ipython",
+                "version": 3
+            },
+            "file_extension": ".py",
+            "mimetype": "text/x-python",
+            "name": "python",
+            "nbconvert_exporter": "python",
+            "pygments_lexer": "ipython3",
+            "version": "3.8.12"
+        },
+        "orig_nbformat": 4
+    },
+    "nbformat": 4,
+    "nbformat_minor": 2
+}
diff --git a/notebooks/training_results.ipynb b/notebooks/training_results.ipynb
index 26b6215..bc91ab8 100644
--- a/notebooks/training_results.ipynb
+++ b/notebooks/training_results.ipynb
@@ -26,7 +26,17 @@
     "from neuralg.training.train_models import train_models\n",
     "from neuralg.plots.plot_eigval_results import plot_eigval_results\n",
     "from neuralg.plots.plot_loss_logs import plot_loss_logs\n",
-    "from neuralg.evaluation.evaluate import evaluate\n"
+    "from neuralg.evaluation.evaluate import evaluate\n",
+    "from neuralg.utils.set_up_torch import set_up_torch"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "set_up_torch(\"float64\", torch_enable_cuda= True)"
    ]
   },
   {
@@ -42,7 +52,8 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "cfg = load_default_training_cfg()"
+    "cfg = load_default_training_cfg()\n",
+    "#del cfg.batch_parameters[\"dist\"]"
    ]
   },
   {
@@ -181,6 +192,13 @@
    "source": [
     "save_run(run_cfg)"
    ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": []
   }
  ],
  "metadata": {
@@ -202,7 +220,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.9.7"
+   "version": "3.8.12"
   },
   "orig_nbformat": 4
  },
diff --git a/prototyping/EigDecomp_Baseline.ipynb b/prototyping/EigDecomp_Baseline.ipynb
index 5f5f8c3..5f13bdb 100644
--- a/prototyping/EigDecomp_Baseline.ipynb
+++ b/prototyping/EigDecomp_Baseline.ipynb
@@ -230,21 +230,16 @@
     "ax.set_title(\"Evaluation results on test set, N = {} \".format(eval_size), fontsize = 20);\n",
     "#plt.savefig(\"eigval_10by10_Nerf.png\")"
    ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "id": "7456fe67",
-   "metadata": {},
-   "outputs": [],
-   "source": []
   }
  ],
  "metadata": {
+  "interpreter": {
+   "hash": "4375ee4ebde01def14bb595c7a3c603f6ae87c7ce25d1cfd1652406174cf8b32"
+  },
   "kernelspec": {
    "display_name": "PyCharm (pythonProject1)",
    "language": "python",
-   "name": "pycharm-8a7ee7e4"
+   "name": "python3"
   },
   "language_info": {
    "codemirror_mode": {
@@ -256,7 +251,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.7.7"
+   "version": "3.9.7"
   }
  },
  "nbformat": 4,
diff --git a/prototyping/nerf.py b/prototyping/nerf.py
index b1d76b7..d18afda 100644
--- a/prototyping/nerf.py
+++ b/prototyping/nerf.py
@@ -1,5 +1,7 @@
 import torch
 import torch.nn as nn
+
+
 class NERFLayer(nn.Module):
     def __init__(self, in_features, out_features, bias=True, activation=nn.ReLU()):
         super().__init__()
@@ -16,20 +18,27 @@ def forward(self, input):
 
 
 class EigNERF(nn.Module):
-    def __init__(self,matrix_dimension,in_features, skip=[], n_neurons=100, activation=nn.Sigmoid(), hidden_layers=8):
+    def __init__(
+        self,
+        matrix_dimension,
+        in_features,
+        skip=[],
+        n_neurons=100,
+        activation=nn.Sigmoid(),
+        hidden_layers=8,
+    ):
         super().__init__()
         self.matrix_dimension = matrix_dimension
         self.in_features = in_features
         self.skip = skip
 
         self.net = nn.ModuleList()
-        self.flatten = nn.Flatten(start_dim = 1)
-        self.net.append(
-            NERFLayer(in_features, n_neurons))
+        self.flatten = nn.Flatten(start_dim=1)
+        self.net.append(NERFLayer(in_features, n_neurons))
 
         for i in range(hidden_layers):
             if i in self.skip:
-                self.net.append(NERFLayer(n_neurons+in_features, n_neurons))
+                self.net.append(NERFLayer(n_neurons + in_features, n_neurons))
             else:
                 self.net.append(NERFLayer(n_neurons, n_neurons))
 
@@ -49,4 +58,5 @@ def forward(self, x):
             if layer_idx in self.skip:
                 out = torch.cat([out, identity], dim=1)
 
-        return out[:,None,:] #Dummy index to fit framework
\ No newline at end of file
+        return out[:, None, :]  # Dummy index to fit framework
+
diff --git a/setup.py b/setup.py
index 8db46e3..32bb000 100644
--- a/setup.py
+++ b/setup.py
@@ -9,6 +9,14 @@
 
 setup(
     name="neuralg",
+    version="0.0.1",
+    install_requires=[
+        "matplotlib>=3.3.3",
+        "torch==1.9.0",
+        "loguru>=0.5.3",
+        "dotmap>=1.3.24",
+        "toml>=0.10.2",
+    ],
     packages=[
         "neuralg",
         "neuralg.io",
@@ -18,7 +26,7 @@
         "neuralg.training",
         "neuralg.plots",
         "neuralg.utils",
-        "neuralg.scripts",
     ],
+    python_requires=">=3.8, <4",
 )
 


Test patch:
diff --git a/.github/workflows/run_tests.yml b/.github/workflows/run_tests.yml
new file mode 100644
index 0000000..819df02
--- /dev/null
+++ b/.github/workflows/run_tests.yml
@@ -0,0 +1,30 @@
+# This workflow will install Python dependencies, run tests and lint with a single version of Python
+# For more information see: https://help.github.com/actions/language-and-framework-guides/using-python-with-github-actions
+
+name: Running tests
+
+on:
+  push:
+    branches:
+    - main
+  pull_request:
+    branches:
+    - main
+jobs:
+  build:
+
+    runs-on: ubuntu-latest
+
+    steps:
+    - uses: actions/checkout@v2
+    - name: Set up Python 3.8
+      uses: actions/setup-python@v2
+      with:
+        python-version: 3.8
+    - name: Install dependencies
+      run: |
+        $CONDA/bin/conda env update --file environment.yml --name base
+    - name: Test with pytest
+      run: |
+        conda install pytest
+        $CONDA/bin/pytest
\ No newline at end of file
diff --git a/neuralg/tests/eig_test.py b/neuralg/tests/eig_test.py
index 9012223..a58a45b 100644
--- a/neuralg/tests/eig_test.py
+++ b/neuralg/tests/eig_test.py
@@ -21,7 +21,7 @@ def test_eig():
     test_parameters = {"N": 10000, "operation": "eig", "wigner": True}
 
     # These are set pretty high right now
-    tol = 0.15  # Preliminary tolerace
+    tol = 0.16  # Preliminary tolerace
     error_bound = 0.95  # Percentage of matrices we require to lie within the tolerance
     results = DotMap()
     # Track which models fails the error bound
@@ -29,11 +29,14 @@ def test_eig():
     results.accuracy = []
     for matrix_size in supported_sizes:
         test_parameters["d"] = matrix_size
-        # Make sure that it supports different types of input e.g. in and out of batch mode
+
+        # Make sure that it supports different types of input e.g. in and out of batch mode with different batch dimensions
         test_matrix = torch.rand(matrix_size, matrix_size)
         test_batch = torch.rand(2, matrix_size, matrix_size)
-        out1, out2 = eig(test_matrix), eig(test_batch)
-
+        test_batch2 = torch.rand(2, 3, 4, matrix_size, matrix_size)
+        out1 = eig(test_matrix)
+        out2 = eig(test_batch)
+        out3 = eig(test_batch2)
         assert out1 is not None
         assert out1.__class__.__name__ == "Tensor"
         assert out1.shape == torch.Size([matrix_size])
@@ -42,6 +45,10 @@ def test_eig():
         assert out2.__class__.__name__ == "Tensor"
         assert out2.shape == torch.Size([2, matrix_size])
 
+        assert out3 is not None
+        assert out3.__class__.__name__ == "Tensor"
+        assert out3.shape == torch.Size([2, 3, 4, matrix_size])
+
         # Should also handle invalid input by throwing value errors
         try:
             out = eig(torch.rand(3, 4))
@@ -52,22 +59,6 @@ def test_eig():
         except ValueError:
             None
 
-        # This is if we want it to be numpy compatible, it actually seemed to cause some trouble
-        # when I tried to fix it. But did not want to spend time on it at this stage
-
-        # test_numpy = np.random.rand(matrix_size, matrix_size)
-        # out3 = eig(test_numpy)
-        # assert out3 is not None
-        # assert out3.__class__.__name__ == "Tensor"
-        # assert out3.shape == torch.Size([matrix_size])
-
-        # # Should also work in batch mode but with numpy
-        # test_numpy = np.random.rand(2, 1, matrix_size, matrix_size)
-        # out4 = eig(test_numpy)
-        # assert out4 is not None
-        # assert out4.__class__.__name__ == "Tensor"
-        #  # assert out4.shape == torch.Size([2, 1, matrix_size])
-
         # Test that errors are within the bound
         error = evaluate_model(eig, test_parameters)
         accuracy = compute_accuracy(tol, error)