From 32edbd8a883a75b7a7d22f695b040e53183f0a7b Mon Sep 17 00:00:00 2001
From: Nicola Demo <demo.nicola@gmail.com>
Date: Thu, 8 Feb 2024 12:17:24 +0100
Subject: [PATCH 01/55] Update version

---
 pina/meta.py | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/pina/meta.py b/pina/meta.py
index fa53e95e4..ac443b5eb 100644
--- a/pina/meta.py
+++ b/pina/meta.py
@@ -13,10 +13,10 @@
 __project__ = "PINA"
 __title__ = "pina"
 __author__ = "PINA Contributors"
-__copyright__ = "2021-2024, PINA Contributors"
+__copyright__ = "2021-2025, PINA Contributors"
 __license__ = "MIT"
-__version__ = "0.1.2"
-__mail__ = "demo.nicola@gmail.com, dario.coscia@sissa.it"  # TODO
+__version__ = "0.2.0"
+__mail__ = 'demo.nicola@gmail.com, dario.coscia@sissa.it'  # TODO
 __maintainer__ = __author__
 __status__ = "Alpha"
 __packagename__ = "pina-mathlab"

From 6290b7f7b7ed58161828bc91f8b0ad9567441142 Mon Sep 17 00:00:00 2001
From: Nicola Demo <demo.nicola@gmail.com>
Date: Thu, 8 Feb 2024 12:19:53 +0100
Subject: [PATCH 02/55] Update testing_pr.yml

---
 .github/workflows/testing_pr.yml | 1 +
 1 file changed, 1 insertion(+)

diff --git a/.github/workflows/testing_pr.yml b/.github/workflows/testing_pr.yml
index 7519e6b7c..9ba8840cb 100644
--- a/.github/workflows/testing_pr.yml
+++ b/.github/workflows/testing_pr.yml
@@ -4,6 +4,7 @@ on:
   pull_request:
     branches:
       - "master"
+      - "0.2"
 
 jobs:
   build:

From 9d1f3e053b1f453f43fdca1495c29efa167578db Mon Sep 17 00:00:00 2001
From: Nicola Demo <demo.nicola@gmail.com>
Date: Tue, 18 Jun 2024 11:20:58 +0200
Subject: [PATCH 03/55] in progress

---
 pina/__init__.py           |   2 +-
 pina/dataset.py            |   2 +-
 pina/geometry/simplex.py   |   2 +-
 pina/label_tensor.py       | 638 +++++++++++++++++++++----------------
 pina/model/fno.py          |   2 +-
 pina/plotter.py            |   2 +-
 tests/test_label_tensor.py | 278 ++++++++++------
 7 files changed, 548 insertions(+), 378 deletions(-)

diff --git a/pina/__init__.py b/pina/__init__.py
index 730b2ead4..c63440ba4 100644
--- a/pina/__init__.py
+++ b/pina/__init__.py
@@ -9,7 +9,7 @@
 ]
 
 from .meta import *
-from .label_tensor import LabelTensor
+#from .label_tensor import LabelTensor
 from .solvers.solver import SolverInterface
 from .trainer import Trainer
 from .plotter import Plotter
diff --git a/pina/dataset.py b/pina/dataset.py
index c6a8d29e4..5f4ba5c7a 100644
--- a/pina/dataset.py
+++ b/pina/dataset.py
@@ -1,6 +1,6 @@
 from torch.utils.data import Dataset
 import torch
-from pina import LabelTensor
+from .label_tensor import LabelTensor
 
 
 class SamplePointDataset(Dataset):
diff --git a/pina/geometry/simplex.py b/pina/geometry/simplex.py
index b04ad537f..15cdc16be 100644
--- a/pina/geometry/simplex.py
+++ b/pina/geometry/simplex.py
@@ -1,7 +1,7 @@
 import torch
 from .location import Location
 from pina.geometry import CartesianDomain
-from pina import LabelTensor
+from ..label_tensor import LabelTensor
 from ..utils import check_consistency
 
 
diff --git a/pina/label_tensor.py b/pina/label_tensor.py
index c8a41f7b4..7cb2f71f0 100644
--- a/pina/label_tensor.py
+++ b/pina/label_tensor.py
@@ -5,6 +5,319 @@
 from torch import Tensor
 
 
+
+# class LabelTensor(torch.Tensor):
+#     """Torch tensor with a label for any column."""
+
+#     @staticmethod
+#     def __new__(cls, x, labels, *args, **kwargs):
+#         return super().__new__(cls, x, *args, **kwargs)
+
+#     def __init__(self, x, labels):
+#         """
+#         Construct a `LabelTensor` by passing a tensor and a list of column
+#         labels. Such labels uniquely identify the columns of the tensor,
+#         allowing for an easier manipulation.
+
+#         :param torch.Tensor x: The data tensor.
+#         :param labels: The labels of the columns.
+#         :type labels: str | list(str) | tuple(str)
+
+#         :Example:
+#             >>> from pina import LabelTensor
+#             >>> tensor = LabelTensor(torch.rand((2000, 3)), ['a', 'b', 'c'])
+#             >>> tensor
+#             tensor([[6.7116e-02, 4.8892e-01, 8.9452e-01],
+#                 [9.2392e-01, 8.2065e-01, 4.1986e-04],
+#                 [8.9266e-01, 5.5446e-01, 6.3500e-01],
+#                 ...,
+#                 [5.8194e-01, 9.4268e-01, 4.1841e-01],
+#                 [1.0246e-01, 9.5179e-01, 3.7043e-02],
+#                 [9.6150e-01, 8.0656e-01, 8.3824e-01]])
+#             >>> tensor.extract('a')
+#             tensor([[0.0671],
+#                     [0.9239],
+#                     [0.8927],
+#                     ...,
+#                     [0.5819],
+#                     [0.1025],
+#                     [0.9615]])
+#             >>> tensor['a']
+#             tensor([[0.0671],
+#                     [0.9239],
+#                     [0.8927],
+#                     ...,
+#                     [0.5819],
+#                     [0.1025],
+#                     [0.9615]])
+#             >>> tensor.extract(['a', 'b'])
+#             tensor([[0.0671, 0.4889],
+#                     [0.9239, 0.8207],
+#                     [0.8927, 0.5545],
+#                     ...,
+#                     [0.5819, 0.9427],
+#                     [0.1025, 0.9518],
+#                     [0.9615, 0.8066]])
+#             >>> tensor.extract(['b', 'a'])
+#             tensor([[0.4889, 0.0671],
+#                     [0.8207, 0.9239],
+#                     [0.5545, 0.8927],
+#                     ...,
+#                     [0.9427, 0.5819],
+#                     [0.9518, 0.1025],
+#                     [0.8066, 0.9615]])
+#         """
+#         if x.ndim == 1:
+#             x = x.reshape(-1, 1)
+
+#         if isinstance(labels, str):
+#             labels = [labels]
+
+#         if len(labels) != x.shape[-1]:
+#             raise ValueError(
+#                 "the tensor has not the same number of columns of "
+#                 "the passed labels."
+#             )
+#         self._labels = labels
+
+#     def __deepcopy__(self, __):
+#         """
+#         Implements deepcopy for label tensor. By default it stores the
+#         current labels and use the :meth:`~torch._tensor.Tensor.__deepcopy__`
+#         method for creating a new :class:`pina.label_tensor.LabelTensor`.
+
+#         :param __: Placeholder parameter.
+#         :type __: None
+#         :return: The deep copy of the :class:`pina.label_tensor.LabelTensor`.
+#         :rtype: LabelTensor
+#         """
+#         labels = self.labels
+#         copy_tensor = deepcopy(self.tensor)
+#         return LabelTensor(copy_tensor, labels)
+
+#     @property
+#     def labels(self):
+#         """Property decorator for labels
+
+#         :return: labels of self
+#         :rtype: list
+#         """
+#         return self._labels
+
+#     @labels.setter
+#     def labels(self, labels):
+#         if len(labels) != self.shape[self.ndim - 1]:  # small check
+#             raise ValueError(
+#                 "The tensor has not the same number of columns of "
+#                 "the passed labels."
+#             )
+
+#         self._labels = labels  # assign the label
+
+#     @staticmethod
+#     def vstack(label_tensors):
+#         """
+#         Stack tensors vertically. For more details, see
+#         :meth:`torch.vstack`.
+
+#         :param list(LabelTensor) label_tensors: the tensors to stack. They need
+#             to have equal labels.
+#         :return: the stacked tensor
+#         :rtype: LabelTensor
+#         """
+#         if len(label_tensors) == 0:
+#             return []
+
+#         all_labels = [label for lt in label_tensors for label in lt.labels]
+#         if set(all_labels) != set(label_tensors[0].labels):
+#             raise RuntimeError("The tensors to stack have different labels")
+
+#         labels = label_tensors[0].labels
+#         tensors = [lt.extract(labels) for lt in label_tensors]
+#         return LabelTensor(torch.vstack(tensors), labels)
+
+#     def clone(self, *args, **kwargs):
+#         """
+#         Clone the LabelTensor. For more details, see
+#         :meth:`torch.Tensor.clone`.
+
+#         :return: A copy of the tensor.
+#         :rtype: LabelTensor
+#         """
+#         # # used before merging
+#         # try:
+#         #     out = LabelTensor(super().clone(*args, **kwargs), self.labels)
+#         # except:
+#         #     out = super().clone(*args, **kwargs)
+#         out = LabelTensor(super().clone(*args, **kwargs), self.labels)
+#         return out
+
+#     def to(self, *args, **kwargs):
+#         """
+#         Performs Tensor dtype and/or device conversion. For more details, see
+#         :meth:`torch.Tensor.to`.
+#         """
+#         tmp = super().to(*args, **kwargs)
+#         new = self.__class__.clone(self)
+#         new.data = tmp.data
+#         return new
+
+#     def select(self, *args, **kwargs):
+#         """
+#         Performs Tensor selection. For more details, see :meth:`torch.Tensor.select`.
+#         """
+#         tmp = super().select(*args, **kwargs)
+#         tmp._labels = self._labels
+#         return tmp
+
+#     def cuda(self, *args, **kwargs):
+#         """
+#         Send Tensor to cuda. For more details, see :meth:`torch.Tensor.cuda`.
+#         """
+#         tmp = super().cuda(*args, **kwargs)
+#         new = self.__class__.clone(self)
+#         new.data = tmp.data
+#         return new
+
+#     def cpu(self, *args, **kwargs):
+#         """
+#         Send Tensor to cpu. For more details, see :meth:`torch.Tensor.cpu`.
+#         """
+#         tmp = super().cpu(*args, **kwargs)
+#         new = self.__class__.clone(self)
+#         new.data = tmp.data
+#         return new
+
+#     def extract(self, label_to_extract):
+#         """
+#         Extract the subset of the original tensor by returning all the columns
+#         corresponding to the passed ``label_to_extract``.
+
+#         :param label_to_extract: The label(s) to extract.
+#         :type label_to_extract: str | list(str) | tuple(str)
+#         :raises TypeError: Labels are not ``str``.
+#         :raises ValueError: Label to extract is not in the labels ``list``.
+#         """
+
+#         if isinstance(label_to_extract, str):
+#             label_to_extract = [label_to_extract]
+#         elif isinstance(label_to_extract, (tuple, list)):  # TODO
+#             pass
+#         else:
+#             raise TypeError(
+#                 "`label_to_extract` should be a str, or a str iterator"
+#             )
+
+#         indeces = []
+#         for f in label_to_extract:
+#             try:
+#                 indeces.append(self.labels.index(f))
+#             except ValueError:
+#                 raise ValueError(f"`{f}` not in the labels list")
+
+#         new_data = super(Tensor, self.T).__getitem__(indeces).T
+#         new_labels = [self.labels[idx] for idx in indeces]
+
+#         extracted_tensor = new_data.as_subclass(LabelTensor)
+#         extracted_tensor.labels = new_labels
+
+#         return extracted_tensor
+
+#     def detach(self):
+#         detached = super().detach()
+#         if hasattr(self, "_labels"):
+#             detached._labels = self._labels
+#         return detached
+
+#     def requires_grad_(self, mode=True):
+#         lt = super().requires_grad_(mode)
+#         lt.labels = self.labels
+#         return lt
+
+#     def append(self, lt, mode="std"):
+#         """
+#         Return a copy of the merged tensors.
+
+#         :param LabelTensor lt: The tensor to merge.
+#         :param str mode: {'std', 'first', 'cross'}
+#         :return: The merged tensors.
+#         :rtype: LabelTensor
+#         """
+#         if set(self.labels).intersection(lt.labels):
+#             raise RuntimeError("The tensors to merge have common labels")
+
+#         new_labels = self.labels + lt.labels
+#         if mode == "std":
+#             new_tensor = torch.cat((self, lt), dim=1)
+#         elif mode == "first":
+#             raise NotImplementedError
+#         elif mode == "cross":
+#             tensor1 = self
+#             tensor2 = lt
+#             n1 = tensor1.shape[0]
+#             n2 = tensor2.shape[0]
+
+#             tensor1 = LabelTensor(tensor1.repeat(n2, 1), labels=tensor1.labels)
+#             tensor2 = LabelTensor(
+#                 tensor2.repeat_interleave(n1, dim=0), labels=tensor2.labels
+#             )
+#             new_tensor = torch.cat((tensor1, tensor2), dim=1)
+
+#         new_tensor = new_tensor.as_subclass(LabelTensor)
+#         new_tensor.labels = new_labels
+#         return new_tensor
+
+#     def __getitem__(self, index):
+#         """
+#         Return a copy of the selected tensor.
+#         """
+
+#         if isinstance(index, str) or (
+#             isinstance(index, (tuple, list))
+#             and all(isinstance(a, str) for a in index)
+#         ):
+#             return self.extract(index)
+
+#         selected_lt = super(Tensor, self).__getitem__(index)
+
+#         try:
+#             len_index = len(index)
+#         except TypeError:
+#             len_index = 1
+
+#         if isinstance(index, int) or len_index == 1:
+#             if selected_lt.ndim == 1:
+#                 selected_lt = selected_lt.reshape(1, -1)
+#             if hasattr(self, "labels"):
+#                 selected_lt.labels = self.labels
+#         elif len_index == 2:
+#             if selected_lt.ndim == 1:
+#                 selected_lt = selected_lt.reshape(-1, 1)
+#             if hasattr(self, "labels"):
+#                 if isinstance(index[1], list):
+#                     selected_lt.labels = [self.labels[i] for i in index[1]]
+#                 else:
+#                     selected_lt.labels = self.labels[index[1]]
+#         else:
+#             selected_lt.labels = self.labels
+
+#         return selected_lt
+
+
+#     def __str__(self):
+#         if hasattr(self, "labels"):
+#             s = f"labels({str(self.labels)})\n"
+#         else:
+#             s = "no labels\n"
+#         s += super().__str__()
+#         return s
+
+def issubset(a, b):
+    """
+    Check if a is a subset of b.
+    """
+    return set(a).issubset(set(b))
+
 class LabelTensor(torch.Tensor):
     """Torch tensor with a label for any column."""
 
@@ -12,180 +325,35 @@ class LabelTensor(torch.Tensor):
     def __new__(cls, x, labels, *args, **kwargs):
         return super().__new__(cls, x, *args, **kwargs)
 
-    def __init__(self, x, labels):
-        """
-        Construct a `LabelTensor` by passing a tensor and a list of column
-        labels. Such labels uniquely identify the columns of the tensor,
-        allowing for an easier manipulation.
-
-        :param torch.Tensor x: The data tensor.
-        :param labels: The labels of the columns.
-        :type labels: str | list(str) | tuple(str)
-
-        :Example:
-            >>> from pina import LabelTensor
-            >>> tensor = LabelTensor(torch.rand((2000, 3)), ['a', 'b', 'c'])
-            >>> tensor
-            tensor([[6.7116e-02, 4.8892e-01, 8.9452e-01],
-                [9.2392e-01, 8.2065e-01, 4.1986e-04],
-                [8.9266e-01, 5.5446e-01, 6.3500e-01],
-                ...,
-                [5.8194e-01, 9.4268e-01, 4.1841e-01],
-                [1.0246e-01, 9.5179e-01, 3.7043e-02],
-                [9.6150e-01, 8.0656e-01, 8.3824e-01]])
-            >>> tensor.extract('a')
-            tensor([[0.0671],
-                    [0.9239],
-                    [0.8927],
-                    ...,
-                    [0.5819],
-                    [0.1025],
-                    [0.9615]])
-            >>> tensor['a']
-            tensor([[0.0671],
-                    [0.9239],
-                    [0.8927],
-                    ...,
-                    [0.5819],
-                    [0.1025],
-                    [0.9615]])
-            >>> tensor.extract(['a', 'b'])
-            tensor([[0.0671, 0.4889],
-                    [0.9239, 0.8207],
-                    [0.8927, 0.5545],
-                    ...,
-                    [0.5819, 0.9427],
-                    [0.1025, 0.9518],
-                    [0.9615, 0.8066]])
-            >>> tensor.extract(['b', 'a'])
-            tensor([[0.4889, 0.0671],
-                    [0.8207, 0.9239],
-                    [0.5545, 0.8927],
-                    ...,
-                    [0.9427, 0.5819],
-                    [0.9518, 0.1025],
-                    [0.8066, 0.9615]])
-        """
-        if x.ndim == 1:
-            x = x.reshape(-1, 1)
-
-        if isinstance(labels, str):
-            labels = [labels]
-
-        if len(labels) != x.shape[-1]:
-            raise ValueError(
-                "the tensor has not the same number of columns of "
-                "the passed labels."
-            )
-        self._labels = labels
-
-    def __deepcopy__(self, __):
-        """
-        Implements deepcopy for label tensor. By default it stores the
-        current labels and use the :meth:`~torch._tensor.Tensor.__deepcopy__`
-        method for creating a new :class:`pina.label_tensor.LabelTensor`.
-
-        :param __: Placeholder parameter.
-        :type __: None
-        :return: The deep copy of the :class:`pina.label_tensor.LabelTensor`.
-        :rtype: LabelTensor
-        """
-        labels = self.labels
-        copy_tensor = deepcopy(self.tensor)
-        return LabelTensor(copy_tensor, labels)
-
     @property
-    def labels(self):
-        """Property decorator for labels
-
-        :return: labels of self
-        :rtype: list
-        """
-        return self._labels
-
-    @labels.setter
-    def labels(self, labels):
-        if len(labels) != self.shape[self.ndim - 1]:  # small check
-            raise ValueError(
-                "The tensor has not the same number of columns of "
-                "the passed labels."
-            )
-
-        self._labels = labels  # assign the label
-
-    @staticmethod
-    def vstack(label_tensors):
-        """
-        Stack tensors vertically. For more details, see
-        :meth:`torch.vstack`.
-
-        :param list(LabelTensor) label_tensors: the tensors to stack. They need
-            to have equal labels.
-        :return: the stacked tensor
-        :rtype: LabelTensor
-        """
-        if len(label_tensors) == 0:
-            return []
-
-        all_labels = [label for lt in label_tensors for label in lt.labels]
-        if set(all_labels) != set(label_tensors[0].labels):
-            raise RuntimeError("The tensors to stack have different labels")
-
-        labels = label_tensors[0].labels
-        tensors = [lt.extract(labels) for lt in label_tensors]
-        return LabelTensor(torch.vstack(tensors), labels)
+    def tensor(self):
+        return self.as_subclass(Tensor)
 
-    def clone(self, *args, **kwargs):
-        """
-        Clone the LabelTensor. For more details, see
-        :meth:`torch.Tensor.clone`.
+    def __len__(self) -> int:
+        return super().__len__()
 
-        :return: A copy of the tensor.
-        :rtype: LabelTensor
-        """
-        # # used before merging
-        # try:
-        #     out = LabelTensor(super().clone(*args, **kwargs), self.labels)
-        # except:
-        #     out = super().clone(*args, **kwargs)
-        out = LabelTensor(super().clone(*args, **kwargs), self.labels)
-        return out
-
-    def to(self, *args, **kwargs):
-        """
-        Performs Tensor dtype and/or device conversion. For more details, see
-        :meth:`torch.Tensor.to`.
+    def __init__(self, x, labels):
         """
-        tmp = super().to(*args, **kwargs)
-        new = self.__class__.clone(self)
-        new.data = tmp.data
-        return new
+        Construct a `LabelTensor` by passing a dict of the labels
 
-    def select(self, *args, **kwargs):
-        """
-        Performs Tensor selection. For more details, see :meth:`torch.Tensor.select`.
+        :Example:
+            >>> from pina import LabelTensor
+            >>> tensor = LabelTensor(
+            >>>     torch.rand((2000, 3)), 
+                    {1: {"name": "space"['a', 'b', 'c'])
+        
         """
-        tmp = super().select(*args, **kwargs)
-        tmp._labels = self._labels
-        return tmp
+        from .utils import check_consistency
+        check_consistency(labels, dict)
 
-    def cuda(self, *args, **kwargs):
-        """
-        Send Tensor to cuda. For more details, see :meth:`torch.Tensor.cuda`.
-        """
-        tmp = super().cuda(*args, **kwargs)
-        new = self.__class__.clone(self)
-        new.data = tmp.data
-        return new
+        self.labels = {
+            idx_: {
+                'dof': range(x.shape[idx_]),
+                'name': idx_
+            } for idx_ in range(x.ndim)
+        }
+        self.labels.update(labels)
 
-    def cpu(self, *args, **kwargs):
-        """
-        Send Tensor to cpu. For more details, see :meth:`torch.Tensor.cpu`.
-        """
-        tmp = super().cpu(*args, **kwargs)
-        new = self.__class__.clone(self)
-        new.data = tmp.data
-        return new
 
     def extract(self, label_to_extract):
         """
@@ -197,122 +365,52 @@ def extract(self, label_to_extract):
         :raises TypeError: Labels are not ``str``.
         :raises ValueError: Label to extract is not in the labels ``list``.
         """
-
-        if isinstance(label_to_extract, str):
+        if isinstance(label_to_extract, (int, str)):
             label_to_extract = [label_to_extract]
-        elif isinstance(label_to_extract, (tuple, list)):  # TODO
-            pass
-        else:
-            raise TypeError(
-                "`label_to_extract` should be a str, or a str iterator"
-            )
+        if isinstance(label_to_extract, (tuple, list)):
+            
+            for k, v in self.labels.items():
+                if issubset(label_to_extract, v['dof']):
+                    break
 
-        indeces = []
-        for f in label_to_extract:
-            try:
-                indeces.append(self.labels.index(f))
-            except ValueError:
-                raise ValueError(f"`{f}` not in the labels list")
+            label_to_extract = {v['name']: label_to_extract}
 
-        new_data = super(Tensor, self.T).__getitem__(indeces).T
-        new_labels = [self.labels[idx] for idx in indeces]
+        for k, v in label_to_extract.items():
+            if isinstance(v, (int, str)):
+                label_to_extract[k] = [v]
 
-        extracted_tensor = new_data.as_subclass(LabelTensor)
-        extracted_tensor.labels = new_labels
+        indeces = []
+        for dim in range(self.ndim):
 
-        return extracted_tensor
+            boolean_idx = [True] * self.shape[dim]
 
-    def detach(self):
-        detached = super().detach()
-        if hasattr(self, "_labels"):
-            detached._labels = self._labels
-        return detached
+            for dim_to_extract, dof_to_extract in label_to_extract.items():
+                if dim_to_extract == self.labels[dim]['name']:
+                    boolean_idx = [False] * self.shape[dim]
+                    for label in dof_to_extract:
+                        idx_to_keep = self.labels[dim]['dof'].index(label)
+                        boolean_idx[idx_to_keep] = True
 
-    def requires_grad_(self, mode=True):
-        lt = super().requires_grad_(mode)
-        lt.labels = self.labels
-        return lt
+            boolean_idx = torch.Tensor(boolean_idx).bool()
 
-    def append(self, lt, mode="std"):
-        """
-        Return a copy of the merged tensors.
+            indeces.append(boolean_idx)
 
-        :param LabelTensor lt: The tensor to merge.
-        :param str mode: {'std', 'first', 'cross'}
-        :return: The merged tensors.
-        :rtype: LabelTensor
-        """
-        if set(self.labels).intersection(lt.labels):
-            raise RuntimeError("The tensors to merge have common labels")
-
-        new_labels = self.labels + lt.labels
-        if mode == "std":
-            new_tensor = torch.cat((self, lt), dim=1)
-        elif mode == "first":
-            raise NotImplementedError
-        elif mode == "cross":
-            tensor1 = self
-            tensor2 = lt
-            n1 = tensor1.shape[0]
-            n2 = tensor2.shape[0]
-
-            tensor1 = LabelTensor(tensor1.repeat(n2, 1), labels=tensor1.labels)
-            tensor2 = LabelTensor(
-                tensor2.repeat_interleave(n1, dim=0), labels=tensor2.labels
-            )
-            new_tensor = torch.cat((tensor1, tensor2), dim=1)
-
-        new_tensor = new_tensor.as_subclass(LabelTensor)
-        new_tensor.labels = new_labels
-        return new_tensor
-
-    def __getitem__(self, index):
-        """
-        Return a copy of the selected tensor.
-        """
+        final_shapes = [sum(idx) for idx in indeces]
+        grids = torch.meshgrid(*indeces)
 
-        if isinstance(index, str) or (
-            isinstance(index, (tuple, list))
-            and all(isinstance(a, str) for a in index)
-        ):
-            return self.extract(index)
-
-        selected_lt = super(Tensor, self).__getitem__(index)
-
-        try:
-            len_index = len(index)
-        except TypeError:
-            len_index = 1
-
-        if isinstance(index, int) or len_index == 1:
-            if selected_lt.ndim == 1:
-                selected_lt = selected_lt.reshape(1, -1)
-            if hasattr(self, "labels"):
-                selected_lt.labels = self.labels
-        elif len_index == 2:
-            if selected_lt.ndim == 1:
-                selected_lt = selected_lt.reshape(-1, 1)
-            if hasattr(self, "labels"):
-                if isinstance(index[1], list):
-                    selected_lt.labels = [self.labels[i] for i in index[1]]
-                else:
-                    selected_lt.labels = self.labels[index[1]]
-        else:
-            selected_lt.labels = self.labels
-
-        return selected_lt
+        ii = grids[0]
+        for grid in grids[1:]:
+            ii = torch.logical_and(ii, grid)
 
-    @property
-    def tensor(self):
-        return self.as_subclass(Tensor)
+        new_tensor = self.tensor[ii].reshape(*final_shapes)
 
-    def __len__(self) -> int:
-        return super().__len__()
+        return LabelTensor(new_tensor, label_to_extract)
+        
 
     def __str__(self):
-        if hasattr(self, "labels"):
-            s = f"labels({str(self.labels)})\n"
-        else:
-            s = "no labels\n"
+        s = ''
+        for key, value in self.labels.items():
+            s += f"{key}: {value}\n"
+        s += '\n' 
         s += super().__str__()
-        return s
+        return s
\ No newline at end of file
diff --git a/pina/model/fno.py b/pina/model/fno.py
index 910b41603..92ca18361 100644
--- a/pina/model/fno.py
+++ b/pina/model/fno.py
@@ -4,7 +4,7 @@
 
 import torch
 import torch.nn as nn
-from pina import LabelTensor
+from ..label_tensor import LabelTensor
 import warnings
 from ..utils import check_consistency
 from .layers.fourier import FourierBlock1D, FourierBlock2D, FourierBlock3D
diff --git a/pina/plotter.py b/pina/plotter.py
index 041ef0575..eedec807d 100644
--- a/pina/plotter.py
+++ b/pina/plotter.py
@@ -3,7 +3,7 @@
 import matplotlib.pyplot as plt
 import torch
 from pina.callbacks import MetricTracker
-from pina import LabelTensor
+from .label_tensor import LabelTensor
 
 
 class Plotter:
diff --git a/tests/test_label_tensor.py b/tests/test_label_tensor.py
index 05dace5e3..5d5693d50 100644
--- a/tests/test_label_tensor.py
+++ b/tests/test_label_tensor.py
@@ -1,119 +1,191 @@
 import torch
 import pytest
 
-from pina import LabelTensor
+from pina.label_tensor import LabelTensor
+#import pina
 
 data = torch.rand((20, 3))
-labels = ['a', 'b', 'c']
-
-
-def test_constructor():
+labels_column = {
+    1: {
+        "name": "space",
+        "dof": ['x', 'y', 'z']
+    }
+}
+labels_row = {
+    0: {
+        "name": "samples",
+        "dof": range(20)
+    }
+}
+labels_all = labels_column | labels_row
+
+@pytest.mark.parametrize("labels", [labels_column, labels_row, labels_all])
+def test_constructor(labels):
     LabelTensor(data, labels)
 
-
 def test_wrong_constructor():
     with pytest.raises(ValueError):
         LabelTensor(data, ['a', 'b'])
 
-
-def test_labels():
-    tensor = LabelTensor(data, labels)
-    assert isinstance(tensor, torch.Tensor)
-    assert tensor.labels == labels
-    with pytest.raises(ValueError):
-        tensor.labels = labels[:-1]
-
-
-def test_extract():
-    label_to_extract = ['a', 'c']
-    tensor = LabelTensor(data, labels)
-    new = tensor.extract(label_to_extract)
-    assert new.labels == label_to_extract
-    assert new.shape[1] == len(label_to_extract)
-    assert torch.all(torch.isclose(data[:, 0::2], new))
-
-
-def test_extract_onelabel():
-    label_to_extract = ['a']
-    tensor = LabelTensor(data, labels)
-    new = tensor.extract(label_to_extract)
-    assert new.ndim == 2
-    assert new.labels == label_to_extract
-    assert new.shape[1] == len(label_to_extract)
-    assert torch.all(torch.isclose(data[:, 0].reshape(-1, 1), new))
-
-
-def test_wrong_extract():
-    label_to_extract = ['a', 'cc']
-    tensor = LabelTensor(data, labels)
-    with pytest.raises(ValueError):
-        tensor.extract(label_to_extract)
-
-
-def test_extract_order():
-    label_to_extract = ['c', 'a']
+@pytest.mark.parametrize("labels", [labels_column, labels_all])
+@pytest.mark.parametrize("labels_te", ['z', ['z'], {'space': ['z']}])
+def test_extract_column(labels, labels_te):
     tensor = LabelTensor(data, labels)
-    new = tensor.extract(label_to_extract)
-    expected = torch.cat(
-        (data[:, 2].reshape(-1, 1), data[:, 0].reshape(-1, 1)),
-        dim=1)
-    assert new.labels == label_to_extract
-    assert new.shape[1] == len(label_to_extract)
-    assert torch.all(torch.isclose(expected, new))
-
-
-def test_merge():
-    tensor = LabelTensor(data, labels)
-    tensor_a = tensor.extract('a')
-    tensor_b = tensor.extract('b')
-    tensor_c = tensor.extract('c')
-
-    tensor_bc = tensor_b.append(tensor_c)
-    assert torch.allclose(tensor_bc, tensor.extract(['b', 'c']))
-
-
-def test_merge2():
-    tensor = LabelTensor(data, labels)
-    tensor_b = tensor.extract('b')
-    tensor_c = tensor.extract('c')
-
-    tensor_bc = tensor_b.append(tensor_c)
-    assert torch.allclose(tensor_bc, tensor.extract(['b', 'c']))
-
-
-def test_getitem():
+    new = tensor.extract(labels_te)
+    assert new.ndim == tensor.ndim
+    assert new.shape[1] == 1
+    assert new.shape[0] == 20
+    assert torch.all(torch.isclose(data[:, 2].reshape(-1, 1), new))
+
+@pytest.mark.parametrize("labels", [labels_row, labels_all])
+@pytest.mark.parametrize("labels_te", [2, [2], {'samples': [2]}])
+def test_extract_row(labels, labels_te):
     tensor = LabelTensor(data, labels)
-    tensor_view = tensor['a']
-
-    assert tensor_view.labels == ['a']
-    assert torch.allclose(tensor_view.flatten(), data[:, 0])
-
-    tensor_view = tensor['a', 'c']
-
-    assert tensor_view.labels == ['a', 'c']
-    assert torch.allclose(tensor_view, data[:, 0::2])
-
-def test_getitem2():
+    new = tensor.extract(labels_te)
+    assert new.ndim == tensor.ndim
+    assert new.shape[1] == 3
+    assert new.shape[0] == 1
+    assert torch.all(torch.isclose(data[2].reshape(1, -1), new))
+
+@pytest.mark.parametrize("labels_te", [
+    {'samples': [2], 'space': ['z']},
+    {'space': 'z', 'samples': 2}
+])
+def test_extract_2D(labels_te):
+    labels = labels_all
     tensor = LabelTensor(data, labels)
-    tensor_view = tensor[:5]
-    assert tensor_view.labels == labels
-    assert torch.allclose(tensor_view, data[:5])
-
-    idx = torch.randperm(tensor.shape[0])
-    tensor_view = tensor[idx]
-    assert tensor_view.labels == labels
-
-
-def test_slice():
+    new = tensor.extract(labels_te)
+    assert new.ndim == tensor.ndim
+    assert new.shape[1] == 1
+    assert new.shape[0] == 1
+    assert torch.all(torch.isclose(data[2,2].reshape(1, 1), new))
+
+def test_extract_3D():
+    labels = labels_all
+    data = torch.rand((20, 3, 4))
+    labels = {
+        1: {
+            "name": "space",
+            "dof": ['x', 'y', 'z']
+        },
+        2: {
+            "name": "time",
+            "dof": range(4)
+        },
+    }
+    labels_te = {
+        'space': ['x', 'z'],
+        'time': range(1, 4)
+    }
     tensor = LabelTensor(data, labels)
-    tensor_view = tensor[:5, :2]
-    assert tensor_view.labels == labels[:2]
-    assert torch.allclose(tensor_view, data[:5, :2])
-
-    tensor_view2 = tensor[3]
-    assert tensor_view2.labels == labels
-    assert torch.allclose(tensor_view2, data[3])
+    new = tensor.extract(labels_te)
+    assert new.ndim == tensor.ndim
+    assert new.shape[0] == 20
+    assert new.shape[1] == 2
+    assert new.shape[2] == 3
+    assert torch.all(torch.isclose(
+        data[:, 0::2, 1:4].reshape(20, 2, 3),
+        new
+    ))
+
+# def test_labels():
+#     tensor = LabelTensor(data, labels)
+#     assert isinstance(tensor, torch.Tensor)
+#     assert tensor.labels == labels
+#     with pytest.raises(ValueError):
+#         tensor.labels = labels[:-1]
+
+
+# def test_extract():
+#     label_to_extract = ['a', 'c']
+#     tensor = LabelTensor(data, labels)
+#     new = tensor.extract(label_to_extract)
+#     assert new.labels == label_to_extract
+#     assert new.shape[1] == len(label_to_extract)
+#     assert torch.all(torch.isclose(data[:, 0::2], new))
+
+
+# def test_extract_onelabel():
+#     label_to_extract = ['a']
+#     tensor = LabelTensor(data, labels)
+#     new = tensor.extract(label_to_extract)
+#     assert new.ndim == 2
+#     assert new.labels == label_to_extract
+#     assert new.shape[1] == len(label_to_extract)
+#     assert torch.all(torch.isclose(data[:, 0].reshape(-1, 1), new))
+
+
+# def test_wrong_extract():
+#     label_to_extract = ['a', 'cc']
+#     tensor = LabelTensor(data, labels)
+#     with pytest.raises(ValueError):
+#         tensor.extract(label_to_extract)
+
+
+# def test_extract_order():
+#     label_to_extract = ['c', 'a']
+#     tensor = LabelTensor(data, labels)
+#     new = tensor.extract(label_to_extract)
+#     expected = torch.cat(
+#         (data[:, 2].reshape(-1, 1), data[:, 0].reshape(-1, 1)),
+#         dim=1)
+#     assert new.labels == label_to_extract
+#     assert new.shape[1] == len(label_to_extract)
+#     assert torch.all(torch.isclose(expected, new))
+
+
+# def test_merge():
+#     tensor = LabelTensor(data, labels)
+#     tensor_a = tensor.extract('a')
+#     tensor_b = tensor.extract('b')
+#     tensor_c = tensor.extract('c')
+
+#     tensor_bc = tensor_b.append(tensor_c)
+#     assert torch.allclose(tensor_bc, tensor.extract(['b', 'c']))
+
+
+# def test_merge2():
+#     tensor = LabelTensor(data, labels)
+#     tensor_b = tensor.extract('b')
+#     tensor_c = tensor.extract('c')
+
+#     tensor_bc = tensor_b.append(tensor_c)
+#     assert torch.allclose(tensor_bc, tensor.extract(['b', 'c']))
+
+
+# def test_getitem():
+#     tensor = LabelTensor(data, labels)
+#     tensor_view = tensor['a']
+
+#     assert tensor_view.labels == ['a']
+#     assert torch.allclose(tensor_view.flatten(), data[:, 0])
+
+#     tensor_view = tensor['a', 'c']
+
+#     assert tensor_view.labels == ['a', 'c']
+#     assert torch.allclose(tensor_view, data[:, 0::2])
+
+# def test_getitem2():
+#     tensor = LabelTensor(data, labels)
+#     tensor_view = tensor[:5]
+#     assert tensor_view.labels == labels
+#     assert torch.allclose(tensor_view, data[:5])
+
+#     idx = torch.randperm(tensor.shape[0])
+#     tensor_view = tensor[idx]
+#     assert tensor_view.labels == labels
+
+
+# def test_slice():
+#     tensor = LabelTensor(data, labels)
+#     tensor_view = tensor[:5, :2]
+#     assert tensor_view.labels == labels[:2]
+#     assert torch.allclose(tensor_view, data[:5, :2])
+
+#     tensor_view2 = tensor[3]
+#     assert tensor_view2.labels == labels
+#     assert torch.allclose(tensor_view2, data[3])
 
-    tensor_view3 = tensor[:, 2]
-    assert tensor_view3.labels == labels[2]
-    assert torch.allclose(tensor_view3, data[:, 2].reshape(-1, 1))
+#     tensor_view3 = tensor[:, 2]
+#     assert tensor_view3.labels == labels[2]
+#     assert torch.allclose(tensor_view3, data[:, 2].reshape(-1, 1))

From 68e7bc1862b99c900e8f5ae97ed26041c2f42254 Mon Sep 17 00:00:00 2001
From: Nicola Demo <demo.nicola@gmail.com>
Date: Fri, 21 Jun 2024 14:37:55 +0200
Subject: [PATCH 04/55] optimizer and scheduler classes

---
 pina/__init__.py        |  6 +++++-
 pina/optimizer.py       | 21 +++++++++++++++++++++
 pina/scheduler.py       | 29 +++++++++++++++++++++++++++++
 tests/test_optimizer.py | 20 ++++++++++++++++++++
 tests/test_scheduler.py | 27 +++++++++++++++++++++++++++
 5 files changed, 102 insertions(+), 1 deletion(-)
 create mode 100644 pina/optimizer.py
 create mode 100644 pina/scheduler.py
 create mode 100644 tests/test_optimizer.py
 create mode 100644 tests/test_scheduler.py

diff --git a/pina/__init__.py b/pina/__init__.py
index c63440ba4..7c7253338 100644
--- a/pina/__init__.py
+++ b/pina/__init__.py
@@ -6,13 +6,17 @@
     "Condition",
     "SamplePointDataset",
     "SamplePointLoader",
+    "TorchOptimizer",
+    "TorchScheduler",
 ]
 
 from .meta import *
-#from .label_tensor import LabelTensor
+from .label_tensor import LabelTensor
 from .solvers.solver import SolverInterface
 from .trainer import Trainer
 from .plotter import Plotter
 from .condition import Condition
 from .dataset import SamplePointDataset
 from .dataset import SamplePointLoader
+from .optimizer import TorchOptimizer
+from .scheduler import TorchScheduler
\ No newline at end of file
diff --git a/pina/optimizer.py b/pina/optimizer.py
new file mode 100644
index 000000000..d400e82f3
--- /dev/null
+++ b/pina/optimizer.py
@@ -0,0 +1,21 @@
+""" Module for PINA Optimizer """
+
+import torch
+from .utils import check_consistency
+
+class Optimizer: # TODO improve interface
+    pass
+
+
+class TorchOptimizer(Optimizer):
+
+    def __init__(self, optimizer_class, **kwargs):
+        check_consistency(optimizer_class, torch.optim.Optimizer, subclass=True)
+
+        self.optimizer_class = optimizer_class
+        self.kwargs = kwargs
+
+    def hook(self, parameters):
+        self.optimizer_instance = self.optimizer_class(
+            parameters, **self.kwargs
+        )
diff --git a/pina/scheduler.py b/pina/scheduler.py
new file mode 100644
index 000000000..563f829c5
--- /dev/null
+++ b/pina/scheduler.py
@@ -0,0 +1,29 @@
+""" Module for PINA Scheduler """
+
+try:
+    from torch.optim.lr_scheduler import LRScheduler  # torch >= 2.0
+except ImportError:
+    from torch.optim.lr_scheduler import (
+        _LRScheduler as LRScheduler,
+    )  # torch < 2.0
+from .optimizer import Optimizer
+from .utils import check_consistency
+
+
+class Scheduler: # TODO improve interface
+    pass
+
+
+class TorchScheduler(Scheduler):
+
+    def __init__(self, scheduler_class, **kwargs):
+        check_consistency(scheduler_class, LRScheduler, subclass=True)
+
+        self.scheduler_class = scheduler_class
+        self.kwargs = kwargs
+
+    def hook(self, optimizer):
+        check_consistency(optimizer, Optimizer)
+        self.scheduler_instance = self.scheduler_class(
+            optimizer.optimizer_instance, **self.kwargs
+        )
diff --git a/tests/test_optimizer.py b/tests/test_optimizer.py
new file mode 100644
index 000000000..489bbdc05
--- /dev/null
+++ b/tests/test_optimizer.py
@@ -0,0 +1,20 @@
+
+import torch
+import pytest
+from pina import TorchOptimizer
+
+opt_list = [
+    torch.optim.Adam,
+    torch.optim.AdamW,
+    torch.optim.SGD,
+    torch.optim.RMSprop
+]
+
+@pytest.mark.parametrize("optimizer_class", opt_list)
+def test_constructor(optimizer_class):
+    TorchOptimizer(optimizer_class, lr=1e-3)
+
+@pytest.mark.parametrize("optimizer_class", opt_list)
+def test_hook(optimizer_class):
+    opt = TorchOptimizer(optimizer_class, lr=1e-3)
+    opt.hook(torch.nn.Linear(10, 10).parameters())
\ No newline at end of file
diff --git a/tests/test_scheduler.py b/tests/test_scheduler.py
new file mode 100644
index 000000000..4cde13efb
--- /dev/null
+++ b/tests/test_scheduler.py
@@ -0,0 +1,27 @@
+
+import torch
+import pytest
+from pina import TorchOptimizer, TorchScheduler
+
+opt_list = [
+    torch.optim.Adam,
+    torch.optim.AdamW,
+    torch.optim.SGD,
+    torch.optim.RMSprop
+]
+
+sch_list = [
+    torch.optim.lr_scheduler.ConstantLR
+]
+
+@pytest.mark.parametrize("scheduler_class", sch_list)
+def test_constructor(scheduler_class):
+    TorchScheduler(scheduler_class)
+
+@pytest.mark.parametrize("optimizer_class", opt_list)
+@pytest.mark.parametrize("scheduler_class", sch_list)
+def test_hook(optimizer_class, scheduler_class):
+    opt = TorchOptimizer(optimizer_class, lr=1e-3)
+    opt.hook(torch.nn.Linear(10, 10).parameters())
+    sch = TorchScheduler(scheduler_class)
+    sch.hook(opt)
\ No newline at end of file

From 80854d58d489dc1b74efab9923f4b721039fd0a9 Mon Sep 17 00:00:00 2001
From: Nicola Demo <demo.nicola@gmail.com>
Date: Thu, 1 Aug 2024 16:30:35 +0200
Subject: [PATCH 05/55] minor

---
 pina/optimizer.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/pina/optimizer.py b/pina/optimizer.py
index d400e82f3..08631e625 100644
--- a/pina/optimizer.py
+++ b/pina/optimizer.py
@@ -18,4 +18,4 @@ def __init__(self, optimizer_class, **kwargs):
     def hook(self, parameters):
         self.optimizer_instance = self.optimizer_class(
             parameters, **self.kwargs
-        )
+        )
\ No newline at end of file

From 19fc43ef18e6be8b9077d44f11c1ae72e946dada Mon Sep 17 00:00:00 2001
From: Nicola Demo <demo.nicola@gmail.com>
Date: Mon, 5 Aug 2024 17:34:34 +0200
Subject: [PATCH 06/55] refact

---
 pina/__init__.py                             |   5 +-
 pina/condition/__init__.py                   |  10 +
 pina/{ => condition}/condition.py            |  42 +--
 pina/condition/condition_interface.py        |  15 ++
 pina/condition/domain_equation_condition.py  |  28 ++
 pina/condition/input_equation_condition.py   |  23 ++
 pina/condition/input_output_condition.py     |  35 +++
 pina/data/__init__.py                        |   0
 pina/{ => data}/dataset.py                   |   2 +-
 pina/optim/__init__.py                       |  11 +
 pina/optim/optimizer_interface.py            |   7 +
 pina/optim/scheduler_interface.py            |   7 +
 pina/optim/torch_optimizer.py                |  19 ++
 pina/optim/torch_scheduler.py                |  27 ++
 pina/solvers/solver.py                       | 255 +++++++++++++------
 pina/solvers/supervised.py                   |  63 ++---
 pina/trainer.py                              |   2 +-
 tests/test_dataset.py                        |   2 +-
 tests/test_solvers/test_supervised_solver.py | 109 ++++----
 19 files changed, 486 insertions(+), 176 deletions(-)
 create mode 100644 pina/condition/__init__.py
 rename pina/{ => condition}/condition.py (73%)
 create mode 100644 pina/condition/condition_interface.py
 create mode 100644 pina/condition/domain_equation_condition.py
 create mode 100644 pina/condition/input_equation_condition.py
 create mode 100644 pina/condition/input_output_condition.py
 create mode 100644 pina/data/__init__.py
 rename pina/{ => data}/dataset.py (99%)
 create mode 100644 pina/optim/__init__.py
 create mode 100644 pina/optim/optimizer_interface.py
 create mode 100644 pina/optim/scheduler_interface.py
 create mode 100644 pina/optim/torch_optimizer.py
 create mode 100644 pina/optim/torch_scheduler.py

diff --git a/pina/__init__.py b/pina/__init__.py
index 7c7253338..e45a4af84 100644
--- a/pina/__init__.py
+++ b/pina/__init__.py
@@ -19,4 +19,7 @@
 from .dataset import SamplePointDataset
 from .dataset import SamplePointLoader
 from .optimizer import TorchOptimizer
-from .scheduler import TorchScheduler
\ No newline at end of file
+from .scheduler import TorchScheduler
+from .condition.condition import Condition
+from .data.dataset import SamplePointDataset
+from .data.dataset import SamplePointLoader
diff --git a/pina/condition/__init__.py b/pina/condition/__init__.py
new file mode 100644
index 000000000..56d1ee4ad
--- /dev/null
+++ b/pina/condition/__init__.py
@@ -0,0 +1,10 @@
+__all__ = [
+    'Condition',
+    'ConditionInterface',
+    'InputOutputCondition',
+    'InputEquationCondition'
+    'LocationEquationCondition',
+]
+
+from .condition_interface import ConditionInterface
+from .input_output_condition import InputOutputCondition
\ No newline at end of file
diff --git a/pina/condition.py b/pina/condition/condition.py
similarity index 73%
rename from pina/condition.py
rename to pina/condition/condition.py
index 5125fe084..da3c6f647 100644
--- a/pina/condition.py
+++ b/pina/condition/condition.py
@@ -1,8 +1,8 @@
 """ Condition module. """
 
-from .label_tensor import LabelTensor
-from .geometry import Location
-from .equation.equation import Equation
+from ..label_tensor import LabelTensor
+from ..geometry import Location
+from ..equation.equation import Equation
 
 
 def dummy(a):
@@ -59,24 +59,32 @@ class Condition:
         "data_weight",
     ]
 
-    def _dictvalue_isinstance(self, dict_, key_, class_):
-        """Check if the value of a dictionary corresponding to `key` is an instance of `class_`."""
-        if key_ not in dict_.keys():
-            return True
+    # def _dictvalue_isinstance(self, dict_, key_, class_):
+    #     """Check if the value of a dictionary corresponding to `key` is an instance of `class_`."""
+    #     if key_ not in dict_.keys():
+    #         return True
 
-        return isinstance(dict_[key_], class_)
+    #     return isinstance(dict_[key_], class_)
 
-    def __init__(self, *args, **kwargs):
-        """
-        Constructor for the `Condition` class.
-        """
-        self.data_weight = kwargs.pop("data_weight", 1.0)
+    # def __init__(self, *args, **kwargs):
+    #     """
+    #     Constructor for the `Condition` class.
+    #     """
+    #     self.data_weight = kwargs.pop("data_weight", 1.0)
 
-        if len(args) != 0:
-            raise ValueError(
-                f"Condition takes only the following keyword arguments: {Condition.__slots__}."
-            )
+    #     if len(args) != 0:
+    #         raise ValueError(
+    #             f"Condition takes only the following keyword arguments: {Condition.__slots__}."
+    #         )
 
+    from . import InputOutputCondition
+    def __new__(cls, *args, **kwargs):
+
+        if sorted(kwargs.keys()) == sorted(["input_points", "output_points"]):
+            return InputOutputCondition(**kwargs)
+        else:
+            raise ValueError(f"Invalid keyword arguments {kwargs.keys()}.")
+        
         if (
             sorted(kwargs.keys()) != sorted(["input_points", "output_points"])
             and sorted(kwargs.keys()) != sorted(["location", "equation"])
diff --git a/pina/condition/condition_interface.py b/pina/condition/condition_interface.py
new file mode 100644
index 000000000..bb43293cd
--- /dev/null
+++ b/pina/condition/condition_interface.py
@@ -0,0 +1,15 @@
+
+from abc import ABCMeta, abstractmethod
+
+
+class ConditionInterface(metaclass=ABCMeta):
+
+    @abstractmethod
+    def residual(self, model):
+        """
+        Compute the residual of the condition.
+
+        :param model: The model to evaluate the condition.
+        :return: The residual of the condition.
+        """
+        pass
\ No newline at end of file
diff --git a/pina/condition/domain_equation_condition.py b/pina/condition/domain_equation_condition.py
new file mode 100644
index 000000000..8f45f8f4a
--- /dev/null
+++ b/pina/condition/domain_equation_condition.py
@@ -0,0 +1,28 @@
+from .condition_interface import ConditionInterface
+
+class DomainEquationCondition(ConditionInterface):
+    """
+    Condition for input/output data.
+    """
+
+    __slots__ = ["domain", "equation"]
+
+    def __init__(self, domain, equation):
+        """
+        Constructor for the `InputOutputCondition` class.
+        """
+        super().__init__()
+        self.domain = domain
+        self.equation = equation
+
+    @staticmethod
+    def batch_residual(model, input_pts, equation):
+        """
+        Compute the residual of the condition for a single batch. Input and
+        output points are provided as arguments.
+
+        :param torch.nn.Module model: The model to evaluate the condition.
+        :param torch.Tensor input_points: The input points.
+        :param torch.Tensor output_points: The output points.
+        """
+        return equation.residual(model(input_pts))
\ No newline at end of file
diff --git a/pina/condition/input_equation_condition.py b/pina/condition/input_equation_condition.py
new file mode 100644
index 000000000..1c57ed78c
--- /dev/null
+++ b/pina/condition/input_equation_condition.py
@@ -0,0 +1,23 @@
+
+from . import ConditionInterface
+
+class InputOutputCondition(ConditionInterface):
+    """
+    Condition for input/output data.
+    """
+
+    __slots__ = ["input_points", "output_points"]
+
+    def __init__(self, input_points, output_points):
+        """
+        Constructor for the `InputOutputCondition` class.
+        """
+        super().__init__()
+        self.input_points = input_points
+        self.output_points = output_points
+
+    def residual(self, model):
+        """
+        Compute the residual of the condition.
+        """
+        return self.output_points - model(self.input_points)
\ No newline at end of file
diff --git a/pina/condition/input_output_condition.py b/pina/condition/input_output_condition.py
new file mode 100644
index 000000000..d8040d64b
--- /dev/null
+++ b/pina/condition/input_output_condition.py
@@ -0,0 +1,35 @@
+
+from . import ConditionInterface
+
+class InputOutputCondition(ConditionInterface):
+    """
+    Condition for input/output data.
+    """
+
+    __slots__ = ["input_points", "output_points"]
+
+    def __init__(self, input_points, output_points):
+        """
+        Constructor for the `InputOutputCondition` class.
+        """
+        super().__init__()
+        self.input_points = input_points
+        self.output_points = output_points
+
+    def residual(self, model):
+        """
+        Compute the residual of the condition.
+        """
+        return self.batch_residual(model, self.input_points, self.output_points)
+
+    @staticmethod
+    def batch_residual(model, input_points, output_points):
+        """
+        Compute the residual of the condition for a single batch. Input and
+        output points are provided as arguments.
+
+        :param torch.nn.Module model: The model to evaluate the condition.
+        :param torch.Tensor input_points: The input points.
+        :param torch.Tensor output_points: The output points.
+        """
+        return output_points - model(input_points)
\ No newline at end of file
diff --git a/pina/data/__init__.py b/pina/data/__init__.py
new file mode 100644
index 000000000..e69de29bb
diff --git a/pina/dataset.py b/pina/data/dataset.py
similarity index 99%
rename from pina/dataset.py
rename to pina/data/dataset.py
index 5f4ba5c7a..bf874e1bf 100644
--- a/pina/dataset.py
+++ b/pina/data/dataset.py
@@ -1,6 +1,6 @@
 from torch.utils.data import Dataset
 import torch
-from .label_tensor import LabelTensor
+from ..label_tensor import LabelTensor
 
 
 class SamplePointDataset(Dataset):
diff --git a/pina/optim/__init__.py b/pina/optim/__init__.py
new file mode 100644
index 000000000..699706c16
--- /dev/null
+++ b/pina/optim/__init__.py
@@ -0,0 +1,11 @@
+__all__ = [
+    "Optimizer",
+    "TorchOptimizer",
+    "Scheduler",
+    "TorchScheduler",
+]
+
+from .optimizer_interface import Optimizer
+from .torch_optimizer import TorchOptimizer
+from .scheduler_interface import Scheduler
+from .torch_scheduler import TorchScheduler
\ No newline at end of file
diff --git a/pina/optim/optimizer_interface.py b/pina/optim/optimizer_interface.py
new file mode 100644
index 000000000..c25506254
--- /dev/null
+++ b/pina/optim/optimizer_interface.py
@@ -0,0 +1,7 @@
+""" Module for PINA Optimizer """
+
+from abc import ABCMeta
+
+
+class Optimizer(metaclass=ABCMeta):  # TODO improve interface
+    pass
\ No newline at end of file
diff --git a/pina/optim/scheduler_interface.py b/pina/optim/scheduler_interface.py
new file mode 100644
index 000000000..dbc0ca82d
--- /dev/null
+++ b/pina/optim/scheduler_interface.py
@@ -0,0 +1,7 @@
+""" Module for PINA Optimizer """
+
+from abc import ABCMeta
+
+
+class Scheduler(metaclass=ABCMeta): # TODO improve interface
+    pass
\ No newline at end of file
diff --git a/pina/optim/torch_optimizer.py b/pina/optim/torch_optimizer.py
new file mode 100644
index 000000000..239819a4f
--- /dev/null
+++ b/pina/optim/torch_optimizer.py
@@ -0,0 +1,19 @@
+""" Module for PINA Torch Optimizer """
+
+import torch
+
+from ..utils import check_consistency
+from .optimizer_interface import Optimizer
+
+class TorchOptimizer(Optimizer):
+
+    def __init__(self, optimizer_class, **kwargs):
+        check_consistency(optimizer_class, torch.optim.Optimizer, subclass=True)
+
+        self.optimizer_class = optimizer_class
+        self.kwargs = kwargs
+
+    def hook(self, parameters):
+        self.optimizer_instance = self.optimizer_class(
+            parameters, **self.kwargs
+        )
\ No newline at end of file
diff --git a/pina/optim/torch_scheduler.py b/pina/optim/torch_scheduler.py
new file mode 100644
index 000000000..50e1d91f7
--- /dev/null
+++ b/pina/optim/torch_scheduler.py
@@ -0,0 +1,27 @@
+""" Module for PINA Torch Optimizer """
+
+import torch
+try:
+    from torch.optim.lr_scheduler import LRScheduler  # torch >= 2.0
+except ImportError:
+    from torch.optim.lr_scheduler import (
+        _LRScheduler as LRScheduler,
+    )  # torch < 2.0
+
+from ..utils import check_consistency
+from .optimizer_interface import Optimizer
+from .scheduler_interface import Scheduler
+
+class TorchScheduler(Scheduler):
+
+    def __init__(self, scheduler_class, **kwargs):
+        check_consistency(scheduler_class, LRScheduler, subclass=True)
+
+        self.scheduler_class = scheduler_class
+        self.kwargs = kwargs
+
+    def hook(self, optimizer):
+        check_consistency(optimizer, Optimizer)
+        self.scheduler_instance = self.scheduler_class(
+            optimizer.optimizer_instance, **self.kwargs
+        )
\ No newline at end of file
diff --git a/pina/solvers/solver.py b/pina/solvers/solver.py
index ec2f40c8d..0112c860a 100644
--- a/pina/solvers/solver.py
+++ b/pina/solvers/solver.py
@@ -5,10 +5,173 @@
 import pytorch_lightning
 from ..utils import check_consistency
 from ..problem import AbstractProblem
+from ..optim import Optimizer, Scheduler
 import torch
 import sys
 
 
+# class SolverInterface(pytorch_lightning.LightningModule, metaclass=ABCMeta):
+#     """
+#     Solver base class. This class inherits is a wrapper of
+#     LightningModule class, inheriting all the
+#     LightningModule methods.
+#     """
+
+#     def __init__(
+#         self,
+#         models,
+#         problem,
+#         optimizers,
+#         optimizers_kwargs,
+#         extra_features=None,
+#     ):
+#         """
+#         :param models: A torch neural network model instance.
+#         :type models: torch.nn.Module
+#         :param problem: A problem definition instance.
+#         :type problem: AbstractProblem
+#         :param list(torch.optim.Optimizer) optimizer: A list of neural network optimizers to
+#             use.
+#         :param list(dict) optimizer_kwargs: A list of optimizer constructor keyword args.
+#         :param list(torch.nn.Module) extra_features: The additional input
+#             features to use as augmented input. If ``None`` no extra features
+#             are passed. If it is a list of :class:`torch.nn.Module`, the extra feature
+#             list is passed to all models. If it is a list of extra features' lists,
+#             each single list of extra feature is passed to a model.
+#         """
+#         super().__init__()
+
+#         # check consistency of the inputs
+#         check_consistency(models, torch.nn.Module)
+#         check_consistency(problem, AbstractProblem)
+#         check_consistency(optimizers, torch.optim.Optimizer, subclass=True)
+#         check_consistency(optimizers_kwargs, dict)
+
+#         # put everything in a list if only one input
+#         if not isinstance(models, list):
+#             models = [models]
+#         if not isinstance(optimizers, list):
+#             optimizers = [optimizers]
+#             optimizers_kwargs = [optimizers_kwargs]
+
+#         # number of models and optimizers
+#         len_model = len(models)
+#         len_optimizer = len(optimizers)
+#         len_optimizer_kwargs = len(optimizers_kwargs)
+
+#         # check length consistency optimizers
+#         if len_model != len_optimizer:
+#             raise ValueError(
+#                 "You must define one optimizer for each model."
+#                 f"Got {len_model} models, and {len_optimizer}"
+#                 " optimizers."
+#             )
+
+#         # check length consistency optimizers kwargs
+#         if len_optimizer_kwargs != len_optimizer:
+#             raise ValueError(
+#                 "You must define one dictionary of keyword"
+#                 " arguments for each optimizers."
+#                 f"Got {len_optimizer} optimizers, and"
+#                 f" {len_optimizer_kwargs} dicitionaries"
+#             )
+
+#         # extra features handling
+#         if (extra_features is None) or (len(extra_features) == 0):
+#             extra_features = [None] * len_model
+#         else:
+#             # if we only have a list of extra features
+#             if not isinstance(extra_features[0], (tuple, list)):
+#                 extra_features = [extra_features] * len_model
+#             else:  # if we have a list of list extra features
+#                 if len(extra_features) != len_model:
+#                     raise ValueError(
+#                         "You passed a list of extrafeatures list with len"
+#                         f"different of models len. Expected {len_model} "
+#                         f"got {len(extra_features)}. If you want to use "
+#                         "the same list of extra features for all models, "
+#                         "just pass a list of extrafeatures and not a list "
+#                         "of list of extra features."
+#                     )
+
+#         # assigning model and optimizers
+#         self._pina_models = []
+#         self._pina_optimizers = []
+
+#         for idx in range(len_model):
+#             model_ = Network(
+#                 model=models[idx],
+#                 input_variables=problem.input_variables,
+#                 output_variables=problem.output_variables,
+#                 extra_features=extra_features[idx],
+#             )
+#             optim_ = optimizers[idx](
+#                 model_.parameters(), **optimizers_kwargs[idx]
+#             )
+#             self._pina_models.append(model_)
+#             self._pina_optimizers.append(optim_)
+
+#         # assigning problem
+#         self._pina_problem = problem
+
+#     @abstractmethod
+#     def forward(self, *args, **kwargs):
+#         pass
+
+#     @abstractmethod
+#     def training_step(self):
+#         pass
+
+#     @abstractmethod
+#     def configure_optimizers(self):
+#         pass
+
+#     @property
+#     def models(self):
+#         """
+#         The torch model."""
+#         return self._pina_models
+
+#     @property
+#     def optimizers(self):
+#         """
+#         The torch model."""
+#         return self._pina_optimizers
+
+#     @property
+#     def problem(self):
+#         """
+#         The problem formulation."""
+#         return self._pina_problem
+
+#     def on_train_start(self):
+#         """
+#         On training epoch start this function is call to do global checks for
+#         the different solvers.
+#         """
+
+#         # 1. Check the verison for dataloader
+#         dataloader = self.trainer.train_dataloader
+#         if sys.version_info < (3, 8):
+#             dataloader = dataloader.loaders
+#         self._dataloader = dataloader
+
+#         return super().on_train_start()
+
+    # @model.setter
+    # def model(self, new_model):
+    #     """
+    #     Set the torch."""
+    #     check_consistency(new_model, nn.Module, 'torch model')
+    #     self._model= new_model
+
+    # @problem.setter
+    # def problem(self, problem):
+    #     """
+    #     Set the problem formulation."""
+    #     check_consistency(problem, AbstractProblem, 'pina problem')
+    #     self._problem = problem
+
 class SolverInterface(pytorch_lightning.LightningModule, metaclass=ABCMeta):
     """
     Solver base class. This class inherits is a wrapper of
@@ -18,45 +181,36 @@ class SolverInterface(pytorch_lightning.LightningModule, metaclass=ABCMeta):
 
     def __init__(
         self,
-        models,
+        model,
         problem,
-        optimizers,
-        optimizers_kwargs,
-        extra_features=None,
+        optimizer,
+        scheduler,
     ):
         """
-        :param models: A torch neural network model instance.
-        :type models: torch.nn.Module
+        :param model: A torch neural network model instance.
+        :type model: torch.nn.Module
         :param problem: A problem definition instance.
         :type problem: AbstractProblem
-        :param list(torch.optim.Optimizer) optimizer: A list of neural network optimizers to
-            use.
-        :param list(dict) optimizer_kwargs: A list of optimizer constructor keyword args.
-        :param list(torch.nn.Module) extra_features: The additional input
-            features to use as augmented input. If ``None`` no extra features
-            are passed. If it is a list of :class:`torch.nn.Module`, the extra feature
-            list is passed to all models. If it is a list of extra features' lists,
-            each single list of extra feature is passed to a model.
+        :param list(torch.optim.Optimizer) optimizer: A list of neural network
+           optimizers to use.
         """
         super().__init__()
 
         # check consistency of the inputs
-        check_consistency(models, torch.nn.Module)
+        check_consistency(model, torch.nn.Module)
         check_consistency(problem, AbstractProblem)
-        check_consistency(optimizers, torch.optim.Optimizer, subclass=True)
-        check_consistency(optimizers_kwargs, dict)
+        check_consistency(optimizer, Optimizer)
+        check_consistency(scheduler, Scheduler)
 
         # put everything in a list if only one input
-        if not isinstance(models, list):
-            models = [models]
-        if not isinstance(optimizers, list):
-            optimizers = [optimizers]
-            optimizers_kwargs = [optimizers_kwargs]
+        if not isinstance(model, list):
+            model = [model]
+        if not isinstance(optimizer, list):
+            optimizer = [optimizer]
 
         # number of models and optimizers
-        len_model = len(models)
-        len_optimizer = len(optimizers)
-        len_optimizer_kwargs = len(optimizers_kwargs)
+        len_model = len(model)
+        len_optimizer = len(optimizer)
 
         # check length consistency optimizers
         if len_model != len_optimizer:
@@ -66,52 +220,11 @@ def __init__(
                 " optimizers."
             )
 
-        # check length consistency optimizers kwargs
-        if len_optimizer_kwargs != len_optimizer:
-            raise ValueError(
-                "You must define one dictionary of keyword"
-                " arguments for each optimizers."
-                f"Got {len_optimizer} optimizers, and"
-                f" {len_optimizer_kwargs} dicitionaries"
-            )
-
         # extra features handling
-        if (extra_features is None) or (len(extra_features) == 0):
-            extra_features = [None] * len_model
-        else:
-            # if we only have a list of extra features
-            if not isinstance(extra_features[0], (tuple, list)):
-                extra_features = [extra_features] * len_model
-            else:  # if we have a list of list extra features
-                if len(extra_features) != len_model:
-                    raise ValueError(
-                        "You passed a list of extrafeatures list with len"
-                        f"different of models len. Expected {len_model} "
-                        f"got {len(extra_features)}. If you want to use "
-                        "the same list of extra features for all models, "
-                        "just pass a list of extrafeatures and not a list "
-                        "of list of extra features."
-                    )
-
-        # assigning model and optimizers
-        self._pina_models = []
-        self._pina_optimizers = []
-
-        for idx in range(len_model):
-            model_ = Network(
-                model=models[idx],
-                input_variables=problem.input_variables,
-                output_variables=problem.output_variables,
-                extra_features=extra_features[idx],
-            )
-            optim_ = optimizers[idx](
-                model_.parameters(), **optimizers_kwargs[idx]
-            )
-            self._pina_models.append(model_)
-            self._pina_optimizers.append(optim_)
-
-        # assigning problem
         self._pina_problem = problem
+        self._pina_model = model
+        self._pina_optimizer = optimizer
+        self._pina_scheduler = scheduler
 
     @abstractmethod
     def forward(self, *args, **kwargs):
@@ -129,13 +242,13 @@ def configure_optimizers(self):
     def models(self):
         """
         The torch model."""
-        return self._pina_models
+        return self._pina_model
 
     @property
     def optimizers(self):
         """
         The torch model."""
-        return self._pina_optimizers
+        return self._pina_optimizer
 
     @property
     def problem(self):
diff --git a/pina/solvers/supervised.py b/pina/solvers/supervised.py
index 425364614..0285096ba 100644
--- a/pina/solvers/supervised.py
+++ b/pina/solvers/supervised.py
@@ -1,21 +1,14 @@
 """ Module for SupervisedSolver """
 
 import torch
+from torch.nn.modules.loss import _Loss
 
-try:
-    from torch.optim.lr_scheduler import LRScheduler  # torch >= 2.0
-except ImportError:
-    from torch.optim.lr_scheduler import (
-        _LRScheduler as LRScheduler,
-    )  # torch < 2.0
-
-from torch.optim.lr_scheduler import ConstantLR
 
+from ..optim import Optimizer, Scheduler, TorchOptimizer, TorchScheduler
 from .solver import SolverInterface
 from ..label_tensor import LabelTensor
 from ..utils import check_consistency
 from ..loss import LossInterface
-from torch.nn.modules.loss import _Loss
 
 
 class SupervisedSolver(SolverInterface):
@@ -51,12 +44,9 @@ def __init__(
         self,
         problem,
         model,
-        extra_features=None,
-        loss=torch.nn.MSELoss(),
-        optimizer=torch.optim.Adam,
-        optimizer_kwargs={"lr": 0.001},
-        scheduler=ConstantLR,
-        scheduler_kwargs={"factor": 1, "total_iters": 0},
+        loss=None,
+        optimizer=None,
+        scheduler=None,
     ):
         """
         :param AbstractProblem problem: The formualation of the problem.
@@ -73,24 +63,26 @@ def __init__(
             rate scheduler.
         :param dict scheduler_kwargs: LR scheduler constructor keyword args.
         """
+        if loss is None:
+            loss = torch.nn.MSELoss()
+
+        if optimizer is None:
+            optimizer = TorchOptimizer(torch.optim.Adam, lr=0.001)
+
+        if scheduler is None:
+            scheduler = TorchScheduler(
+                torch.optim.lr_scheduler.ConstantLR)
+
         super().__init__(
-            models=[model],
+            model=model,
             problem=problem,
-            optimizers=[optimizer],
-            optimizers_kwargs=[optimizer_kwargs],
-            extra_features=extra_features,
+            optimizer=optimizer,
+            scheduler=scheduler,
         )
 
         # check consistency
-        check_consistency(scheduler, LRScheduler, subclass=True)
-        check_consistency(scheduler_kwargs, dict)
         check_consistency(loss, (LossInterface, _Loss), subclass=False)
 
-        # assign variables
-        self._scheduler = scheduler(self.optimizers[0], **scheduler_kwargs)
-        self._loss = loss
-        self._neural_net = self.models[0]
-
     def forward(self, x):
         """Forward pass implementation for the solver.
 
@@ -98,7 +90,7 @@ def forward(self, x):
         :return: Solver solution.
         :rtype: torch.Tensor
         """
-        return self.neural_net(x)
+        return self._pina_model(x)
 
     def configure_optimizers(self):
         """Optimizer configuration for the solver.
@@ -106,7 +98,9 @@ def configure_optimizers(self):
         :return: The optimizers and the schedulers
         :rtype: tuple(list, list)
         """
-        return self.optimizers, [self.scheduler]
+        self._pina_optimizer.hook(self._pina_model.parameters())
+        self._pina_scheduler.hook(self._pina_optimizer)
+        return self._pina_optimizer, self._pina_scheduler
 
     def training_step(self, batch, batch_idx):
         """Solver training step.
@@ -168,14 +162,21 @@ def scheduler(self):
         """
         Scheduler for training.
         """
-        return self._scheduler
+        return self._pina_scheduler
+    
+    @property
+    def optimizer(self):
+        """
+        Optimizer for training.
+        """
+        return self._pina_optimizer
 
     @property
-    def neural_net(self):
+    def model(self):
         """
         Neural network for training.
         """
-        return self._neural_net
+        return self._pina_model
 
     @property
     def loss(self):
diff --git a/pina/trainer.py b/pina/trainer.py
index 40f4eb691..25d21b77c 100644
--- a/pina/trainer.py
+++ b/pina/trainer.py
@@ -3,7 +3,7 @@
 import torch
 import pytorch_lightning
 from .utils import check_consistency
-from .dataset import SamplePointDataset, SamplePointLoader, DataPointDataset
+from .data.dataset import SamplePointDataset, SamplePointLoader, DataPointDataset
 from .solvers.solver import SolverInterface
 
 
diff --git a/tests/test_dataset.py b/tests/test_dataset.py
index ff1b6c228..cb6a9e4b0 100644
--- a/tests/test_dataset.py
+++ b/tests/test_dataset.py
@@ -1,7 +1,7 @@
 import torch
 import pytest
 
-from pina.dataset import SamplePointDataset, SamplePointLoader, DataPointDataset
+from pina.data.dataset import SamplePointDataset, SamplePointLoader, DataPointDataset
 from pina import LabelTensor, Condition
 from pina.equation import Equation
 from pina.geometry import CartesianDomain
diff --git a/tests/test_solvers/test_supervised_solver.py b/tests/test_solvers/test_supervised_solver.py
index dfe0bd867..d9cbea3ce 100644
--- a/tests/test_solvers/test_supervised_solver.py
+++ b/tests/test_solvers/test_supervised_solver.py
@@ -11,8 +11,11 @@
 class NeuralOperatorProblem(AbstractProblem):
     input_variables = ['u_0', 'u_1']
     output_variables = ['u']
-    conditions = {'data' : Condition(input_points=LabelTensor(torch.rand(100, 2), input_variables), 
-                                     output_points=LabelTensor(torch.rand(100, 1), output_variables))}
+    conditions = {
+        # 'data' : Condition(
+        #     input_points=LabelTensor(torch.rand(100, 2), input_variables), 
+        #     output_points=LabelTensor(torch.rand(100, 1), output_variables))
+    }
 
 class myFeature(torch.nn.Module):
     """
@@ -39,63 +42,63 @@ def forward(self, x):
 
 
 def test_constructor():
-    SupervisedSolver(problem=problem, model=model, extra_features=None)
+    SupervisedSolver(problem=problem, model=model)
 
 
-def test_constructor_extra_feats():
-    SupervisedSolver(problem=problem, model=model_extra_feats, extra_features=extra_feats)
+# def test_constructor_extra_feats():
+#     SupervisedSolver(problem=problem, model=model_extra_feats, extra_features=extra_feats)
 
 
 def test_train_cpu():
-    solver = SupervisedSolver(problem = problem, model=model, extra_features=None, loss=LpLoss())
+    solver = SupervisedSolver(problem = problem, model=model, loss=LpLoss())
     trainer = Trainer(solver=solver, max_epochs=3, accelerator='cpu', batch_size=20)
     trainer.train()
 
 
-def test_train_restore():
-    tmpdir = "tests/tmp_restore"
-    solver = SupervisedSolver(problem=problem,
-                model=model,
-                extra_features=None,
-                loss=LpLoss())
-    trainer = Trainer(solver=solver,
-                      max_epochs=5,
-                      accelerator='cpu',
-                      default_root_dir=tmpdir)
-    trainer.train()
-    ntrainer = Trainer(solver=solver, max_epochs=15, accelerator='cpu')
-    t = ntrainer.train(
-        ckpt_path=f'{tmpdir}/lightning_logs/version_0/checkpoints/epoch=4-step=5.ckpt')
-    import shutil
-    shutil.rmtree(tmpdir)
-
-
-def test_train_load():
-    tmpdir = "tests/tmp_load"
-    solver = SupervisedSolver(problem=problem,
-                model=model,
-                extra_features=None,
-                loss=LpLoss())
-    trainer = Trainer(solver=solver,
-                      max_epochs=15,
-                      accelerator='cpu',
-                      default_root_dir=tmpdir)
-    trainer.train()
-    new_solver = SupervisedSolver.load_from_checkpoint(
-        f'{tmpdir}/lightning_logs/version_0/checkpoints/epoch=14-step=15.ckpt',
-        problem = problem, model=model)
-    test_pts = LabelTensor(torch.rand(20, 2), problem.input_variables)
-    assert new_solver.forward(test_pts).shape == (20, 1)
-    assert new_solver.forward(test_pts).shape == solver.forward(test_pts).shape
-    torch.testing.assert_close(
-        new_solver.forward(test_pts),
-        solver.forward(test_pts))
-    import shutil
-    shutil.rmtree(tmpdir)
-
-def test_train_extra_feats_cpu():
-    pinn = SupervisedSolver(problem=problem,
-                model=model_extra_feats,
-                extra_features=extra_feats)
-    trainer = Trainer(solver=pinn, max_epochs=5, accelerator='cpu')
-    trainer.train()
\ No newline at end of file
+# def test_train_restore():
+#     tmpdir = "tests/tmp_restore"
+#     solver = SupervisedSolver(problem=problem,
+#                 model=model,
+#                 extra_features=None,
+#                 loss=LpLoss())
+#     trainer = Trainer(solver=solver,
+#                       max_epochs=5,
+#                       accelerator='cpu',
+#                       default_root_dir=tmpdir)
+#     trainer.train()
+#     ntrainer = Trainer(solver=solver, max_epochs=15, accelerator='cpu')
+#     t = ntrainer.train(
+#         ckpt_path=f'{tmpdir}/lightning_logs/version_0/checkpoints/epoch=4-step=5.ckpt')
+#     import shutil
+#     shutil.rmtree(tmpdir)
+
+
+# def test_train_load():
+#     tmpdir = "tests/tmp_load"
+#     solver = SupervisedSolver(problem=problem,
+#                 model=model,
+#                 extra_features=None,
+#                 loss=LpLoss())
+#     trainer = Trainer(solver=solver,
+#                       max_epochs=15,
+#                       accelerator='cpu',
+#                       default_root_dir=tmpdir)
+#     trainer.train()
+#     new_solver = SupervisedSolver.load_from_checkpoint(
+#         f'{tmpdir}/lightning_logs/version_0/checkpoints/epoch=14-step=15.ckpt',
+#         problem = problem, model=model)
+#     test_pts = LabelTensor(torch.rand(20, 2), problem.input_variables)
+#     assert new_solver.forward(test_pts).shape == (20, 1)
+#     assert new_solver.forward(test_pts).shape == solver.forward(test_pts).shape
+#     torch.testing.assert_close(
+#         new_solver.forward(test_pts),
+#         solver.forward(test_pts))
+#     import shutil
+#     shutil.rmtree(tmpdir)
+
+# def test_train_extra_feats_cpu():
+#     pinn = SupervisedSolver(problem=problem,
+#                 model=model_extra_feats,
+#                 extra_features=extra_feats)
+#     trainer = Trainer(solver=pinn, max_epochs=5, accelerator='cpu')
+#     trainer.train()
\ No newline at end of file

From 7ff19415ad33276a56799c90cbf39ac3385aff29 Mon Sep 17 00:00:00 2001
From: Nicola Demo <demo.nicola@gmail.com>
Date: Thu, 8 Aug 2024 16:19:52 +0200
Subject: [PATCH 07/55] supervised working

---
 examples/problems/burgers.py                  |   2 +-
 examples/problems/first_order_ode.py          |   2 +-
 .../parametric_elliptic_optimal_control.py    |   2 +-
 examples/problems/parametric_poisson.py       |   2 +-
 examples/problems/poisson.py                  |   2 +-
 examples/problems/stokes.py                   |   2 +-
 examples/problems/wave.py                     |   2 +-
 pina/__init__.py                              |   7 +-
 pina/condition/__init__.py                    |   8 +-
 pina/condition/condition.py                   |  22 +--
 pina/condition/condition_interface.py         |   3 +
 ...ondition.py => domain_output_condition.py} |  18 ++-
 pina/condition/input_equation_condition.py    |   2 +-
 pina/data/__init__.py                         |   7 +
 pina/data/data_dataset.py                     |  41 +++++
 pina/data/pina_batch.py                       |  36 +++++
 pina/data/{dataset.py => pina_dataloader.py}  | 142 ++++++------------
 pina/data/sample_dataset.py                   |  43 ++++++
 pina/{geometry => domain}/__init__.py         |   4 +-
 pina/{geometry => domain}/cartesian.py        |   4 +-
 .../{geometry => domain}/difference_domain.py |   0
 .../domain_interface.py}                      |   4 +-
 pina/{geometry => domain}/ellipsoid.py        |   4 +-
 pina/{geometry => domain}/exclusion_domain.py |   0
 .../intersection_domain.py                    |   0
 .../operation_interface.py                    |   6 +-
 pina/{geometry => domain}/simplex.py          |   6 +-
 pina/{geometry => domain}/union_domain.py     |   0
 pina/label_tensor.py                          |  26 +++-
 pina/problem/abstract_problem.py              |  88 ++++-------
 pina/solvers/solver.py                        |   2 +
 pina/solvers/supervised.py                    |  37 ++++-
 pina/trainer.py                               |  46 +++---
 .../test_adaptive_refinment_callbacks.py      |   2 +-
 .../test_optimizer_callbacks.py               |   2 +-
 tests/test_condition.py                       |   2 +-
 tests/test_dataset.py                         |   2 +-
 tests/test_geometry/test_cartesian.py         |   2 +-
 tests/test_geometry/test_difference.py        |   2 +-
 tests/test_geometry/test_ellipsoid.py         |   2 +-
 tests/test_geometry/test_exclusion.py         |   2 +-
 tests/test_geometry/test_intersection.py      |   2 +-
 tests/test_geometry/test_simplex.py           |   2 +-
 tests/test_geometry/test_union.py             |   2 +-
 tests/test_plotter.py                         |   2 +-
 tests/test_problem.py                         |   2 +-
 tests/test_solvers/test_causalpinn.py         |   2 +-
 tests/test_solvers/test_competitive_pinn.py   |   2 +-
 tests/test_solvers/test_gpinn.py              |   2 +-
 tests/test_solvers/test_pinn.py               |   2 +-
 tests/test_solvers/test_sapinn.py             |   2 +-
 tests/test_solvers/test_supervised_solver.py  |  15 +-
 tests/test_utils.py                           |   4 +-
 tutorials/tutorial1/tutorial.py               |   4 +-
 tutorials/tutorial11/tutorial.py              |   2 +-
 tutorials/tutorial12/tutorial.py              |   2 +-
 tutorials/tutorial2/tutorial.py               |   2 +-
 tutorials/tutorial3/tutorial.py               |   2 +-
 tutorials/tutorial7/tutorial.py               |   2 +-
 tutorials/tutorial8/tutorial.py               |   2 +-
 tutorials/tutorial9/tutorial.py               |   2 +-
 61 files changed, 378 insertions(+), 265 deletions(-)
 rename pina/condition/{input_output_condition.py => domain_output_condition.py} (65%)
 create mode 100644 pina/data/data_dataset.py
 create mode 100644 pina/data/pina_batch.py
 rename pina/data/{dataset.py => pina_dataloader.py} (67%)
 create mode 100644 pina/data/sample_dataset.py
 rename pina/{geometry => domain}/__init__.py (87%)
 rename pina/{geometry => domain}/cartesian.py (99%)
 rename pina/{geometry => domain}/difference_domain.py (100%)
 rename pina/{geometry/location.py => domain/domain_interface.py} (90%)
 rename pina/{geometry => domain}/ellipsoid.py (99%)
 rename pina/{geometry => domain}/exclusion_domain.py (100%)
 rename pina/{geometry => domain}/intersection_domain.py (100%)
 rename pina/{geometry => domain}/operation_interface.py (92%)
 rename pina/{geometry => domain}/simplex.py (98%)
 rename pina/{geometry => domain}/union_domain.py (100%)

diff --git a/examples/problems/burgers.py b/examples/problems/burgers.py
index 5da9ccb47..33c1b5db2 100644
--- a/examples/problems/burgers.py
+++ b/examples/problems/burgers.py
@@ -15,7 +15,7 @@
 
 
 import torch
-from pina.geometry import CartesianDomain
+from pina.domain import CartesianDomain
 from pina import Condition
 from pina.problem import TimeDependentProblem, SpatialProblem
 from pina.operators import grad
diff --git a/examples/problems/first_order_ode.py b/examples/problems/first_order_ode.py
index 802b85bfe..be1d88c42 100644
--- a/examples/problems/first_order_ode.py
+++ b/examples/problems/first_order_ode.py
@@ -17,7 +17,7 @@
 
 from pina.problem import SpatialProblem
 from pina import Condition
-from pina.geometry import CartesianDomain
+from pina.domain import CartesianDomain
 from pina.operators import grad
 from pina.equation import Equation, FixedValue
 import torch
diff --git a/examples/problems/parametric_elliptic_optimal_control.py b/examples/problems/parametric_elliptic_optimal_control.py
index 9d88b497a..9ecd85154 100644
--- a/examples/problems/parametric_elliptic_optimal_control.py
+++ b/examples/problems/parametric_elliptic_optimal_control.py
@@ -2,7 +2,7 @@
 
 
 from pina import Condition
-from pina.geometry import CartesianDomain
+from pina.domain import CartesianDomain
 from pina.equation import SystemEquation, FixedValue
 from pina.problem import SpatialProblem, ParametricProblem
 from pina.operators import laplacian
diff --git a/examples/problems/parametric_poisson.py b/examples/problems/parametric_poisson.py
index 58867d5bb..64dfdaaee 100644
--- a/examples/problems/parametric_poisson.py
+++ b/examples/problems/parametric_poisson.py
@@ -14,7 +14,7 @@
 # ===================================================== #
 
 
-from pina.geometry import CartesianDomain
+from pina.domain import CartesianDomain
 from pina.problem import SpatialProblem, ParametricProblem
 from pina.operators import laplacian
 from pina.equation import FixedValue, Equation
diff --git a/examples/problems/poisson.py b/examples/problems/poisson.py
index c817787bd..e4a6cf98e 100644
--- a/examples/problems/poisson.py
+++ b/examples/problems/poisson.py
@@ -13,7 +13,7 @@
 
 
 import torch
-from pina.geometry import CartesianDomain
+from pina.domain import CartesianDomain
 from pina import Condition
 from pina.problem import SpatialProblem
 from pina.operators import laplacian
diff --git a/examples/problems/stokes.py b/examples/problems/stokes.py
index f136d64ad..c7b13873c 100644
--- a/examples/problems/stokes.py
+++ b/examples/problems/stokes.py
@@ -4,7 +4,7 @@
 from pina.problem import SpatialProblem
 from pina.operators import laplacian, grad, div
 from pina import Condition, LabelTensor
-from pina.geometry import CartesianDomain
+from pina.domain import CartesianDomain
 from pina.equation import SystemEquation, Equation
 
 # ===================================================== #
diff --git a/examples/problems/wave.py b/examples/problems/wave.py
index cce94da68..124a62d4b 100644
--- a/examples/problems/wave.py
+++ b/examples/problems/wave.py
@@ -2,7 +2,7 @@
 
 
 import torch
-from pina.geometry import CartesianDomain
+from pina.domain import CartesianDomain
 from pina import Condition
 from pina.problem import SpatialProblem, TimeDependentProblem
 from pina.operators import laplacian, grad
diff --git a/pina/__init__.py b/pina/__init__.py
index e45a4af84..835232b12 100644
--- a/pina/__init__.py
+++ b/pina/__init__.py
@@ -15,11 +15,8 @@
 from .solvers.solver import SolverInterface
 from .trainer import Trainer
 from .plotter import Plotter
-from .condition import Condition
-from .dataset import SamplePointDataset
-from .dataset import SamplePointLoader
 from .optimizer import TorchOptimizer
 from .scheduler import TorchScheduler
 from .condition.condition import Condition
-from .data.dataset import SamplePointDataset
-from .data.dataset import SamplePointLoader
+from .data import SamplePointDataset
+from .data import SamplePointLoader
diff --git a/pina/condition/__init__.py b/pina/condition/__init__.py
index 56d1ee4ad..ff329c1bc 100644
--- a/pina/condition/__init__.py
+++ b/pina/condition/__init__.py
@@ -1,10 +1,10 @@
 __all__ = [
     'Condition',
     'ConditionInterface',
-    'InputOutputCondition',
-    'InputEquationCondition'
-    'LocationEquationCondition',
+    'DomainOutputCondition',
+    'DomainEquationCondition'
 ]
 
 from .condition_interface import ConditionInterface
-from .input_output_condition import InputOutputCondition
\ No newline at end of file
+from .domain_output_condition import DomainOutputCondition
+from .domain_equation_condition import DomainEquationCondition
\ No newline at end of file
diff --git a/pina/condition/condition.py b/pina/condition/condition.py
index da3c6f647..eec523c92 100644
--- a/pina/condition/condition.py
+++ b/pina/condition/condition.py
@@ -1,9 +1,11 @@
 """ Condition module. """
 
 from ..label_tensor import LabelTensor
-from ..geometry import Location
+from ..domain import DomainInterface
 from ..equation.equation import Equation
 
+from . import DomainOutputCondition, DomainEquationCondition
+
 
 def dummy(a):
     """Dummy function for testing purposes."""
@@ -51,14 +53,6 @@ class Condition:
 
     """
 
-    __slots__ = [
-        "input_points",
-        "output_points",
-        "location",
-        "equation",
-        "data_weight",
-    ]
-
     # def _dictvalue_isinstance(self, dict_, key_, class_):
     #     """Check if the value of a dictionary corresponding to `key` is an instance of `class_`."""
     #     if key_ not in dict_.keys():
@@ -77,11 +71,17 @@ class Condition:
     #             f"Condition takes only the following keyword arguments: {Condition.__slots__}."
     #         )
 
-    from . import InputOutputCondition
     def __new__(cls, *args, **kwargs):
 
         if sorted(kwargs.keys()) == sorted(["input_points", "output_points"]):
-            return InputOutputCondition(**kwargs)
+            return DomainOutputCondition(
+                domain=kwargs["input_points"],
+                output_points=kwargs["output_points"]
+            )
+        elif sorted(kwargs.keys()) == sorted(["domain", "output_points"]):
+            return DomainOutputCondition(**kwargs)
+        elif sorted(kwargs.keys()) == sorted(["domain", "equation"]):
+            return DomainEquationCondition(**kwargs)
         else:
             raise ValueError(f"Invalid keyword arguments {kwargs.keys()}.")
         
diff --git a/pina/condition/condition_interface.py b/pina/condition/condition_interface.py
index bb43293cd..f6b51bf96 100644
--- a/pina/condition/condition_interface.py
+++ b/pina/condition/condition_interface.py
@@ -4,6 +4,9 @@
 
 class ConditionInterface(metaclass=ABCMeta):
 
+    def __init__(self) -> None:
+        self._problem = None
+
     @abstractmethod
     def residual(self, model):
         """
diff --git a/pina/condition/input_output_condition.py b/pina/condition/domain_output_condition.py
similarity index 65%
rename from pina/condition/input_output_condition.py
rename to pina/condition/domain_output_condition.py
index d8040d64b..49a0cb6ed 100644
--- a/pina/condition/input_output_condition.py
+++ b/pina/condition/domain_output_condition.py
@@ -1,26 +1,34 @@
 
 from . import ConditionInterface
 
-class InputOutputCondition(ConditionInterface):
+class DomainOutputCondition(ConditionInterface):
     """
     Condition for input/output data.
     """
 
-    __slots__ = ["input_points", "output_points"]
+    __slots__ = ["domain", "output_points"]
 
-    def __init__(self, input_points, output_points):
+    def __init__(self, domain, output_points):
         """
         Constructor for the `InputOutputCondition` class.
         """
         super().__init__()
-        self.input_points = input_points
+        print(self)
+        self.domain = domain
         self.output_points = output_points
 
+    @property
+    def input_points(self):
+        """
+        Get the input points of the condition.
+        """
+        return self._problem.domains[self.domain]
+
     def residual(self, model):
         """
         Compute the residual of the condition.
         """
-        return self.batch_residual(model, self.input_points, self.output_points)
+        return self.batch_residual(model, self.domain, self.output_points)
 
     @staticmethod
     def batch_residual(model, input_points, output_points):
diff --git a/pina/condition/input_equation_condition.py b/pina/condition/input_equation_condition.py
index 1c57ed78c..288022c00 100644
--- a/pina/condition/input_equation_condition.py
+++ b/pina/condition/input_equation_condition.py
@@ -1,7 +1,7 @@
 
 from . import ConditionInterface
 
-class InputOutputCondition(ConditionInterface):
+class InputEquationCondition(ConditionInterface):
     """
     Condition for input/output data.
     """
diff --git a/pina/data/__init__.py b/pina/data/__init__.py
index e69de29bb..fba19b92c 100644
--- a/pina/data/__init__.py
+++ b/pina/data/__init__.py
@@ -0,0 +1,7 @@
+__all__ = [
+]
+
+from .pina_dataloader import SamplePointLoader
+from .data_dataset import DataPointDataset
+from .sample_dataset import SamplePointDataset
+from .pina_batch import Batch
\ No newline at end of file
diff --git a/pina/data/data_dataset.py b/pina/data/data_dataset.py
new file mode 100644
index 000000000..9dff2d7ed
--- /dev/null
+++ b/pina/data/data_dataset.py
@@ -0,0 +1,41 @@
+from torch.utils.data import Dataset
+import torch
+from ..label_tensor import LabelTensor
+
+
+class DataPointDataset(Dataset):
+
+    def __init__(self, problem, device) -> None:
+        super().__init__()
+        input_list = []
+        output_list = []
+        self.condition_names = []
+
+        for name, condition in problem.conditions.items():
+            if hasattr(condition, "output_points"):
+                input_list.append(problem.conditions[name].input_points)
+                output_list.append(problem.conditions[name].output_points)
+                self.condition_names.append(name)
+
+        self.input_pts = LabelTensor.stack(input_list)
+        self.output_pts = LabelTensor.stack(output_list)
+
+        if self.input_pts != []:
+            self.condition_indeces = torch.cat(
+                [
+                    torch.tensor([i] * len(input_list[i]))
+                    for i in range(len(self.condition_names))
+                ],
+                dim=0,
+            )
+        else:  # if there are no data points
+            self.condition_indeces = torch.tensor([])
+            self.input_pts = torch.tensor([])
+            self.output_pts = torch.tensor([])
+
+        self.input_pts = self.input_pts.to(device)
+        self.output_pts = self.output_pts.to(device)
+        self.condition_indeces = self.condition_indeces.to(device)
+
+    def __len__(self):
+        return self.input_pts.shape[0]
\ No newline at end of file
diff --git a/pina/data/pina_batch.py b/pina/data/pina_batch.py
new file mode 100644
index 000000000..cb1296ede
--- /dev/null
+++ b/pina/data/pina_batch.py
@@ -0,0 +1,36 @@
+
+
+class Batch:
+    """
+    This class is used to create a dataset of sample points.
+    """
+
+    def __init__(self, type_, idx, *args, **kwargs) -> None:
+        """
+        """
+        if type_ == "sample":
+
+            if len(args) != 2:
+                raise RuntimeError
+
+            input = args[0]
+            conditions = args[1]
+
+            self.input = input[idx]
+            self.condition = conditions[idx]
+
+        elif type_ == "data":
+
+            if len(args) != 3:
+                raise RuntimeError
+
+            input = args[0]
+            output = args[1]
+            conditions = args[2]
+
+            self.input = input[idx]
+            self.output = output[idx]
+            self.condition = conditions[idx]
+                
+        else:
+            raise ValueError("Invalid number of arguments.")
\ No newline at end of file
diff --git a/pina/data/dataset.py b/pina/data/pina_dataloader.py
similarity index 67%
rename from pina/data/dataset.py
rename to pina/data/pina_dataloader.py
index bf874e1bf..1b71a46a3 100644
--- a/pina/data/dataset.py
+++ b/pina/data/pina_dataloader.py
@@ -1,84 +1,8 @@
-from torch.utils.data import Dataset
 import torch
-from ..label_tensor import LabelTensor
-
-
-class SamplePointDataset(Dataset):
-    """
-    This class is used to create a dataset of sample points.
-    """
-
-    def __init__(self, problem, device) -> None:
-        """
-        :param dict input_pts: The input points.
-        """
-        super().__init__()
-        pts_list = []
-        self.condition_names = []
-
-        for name, condition in problem.conditions.items():
-            if not hasattr(condition, "output_points"):
-                pts_list.append(problem.input_pts[name])
-                self.condition_names.append(name)
-
-        self.pts = LabelTensor.vstack(pts_list)
-
-        if self.pts != []:
-            self.condition_indeces = torch.cat(
-                [
-                    torch.tensor([i] * len(pts_list[i]))
-                    for i in range(len(self.condition_names))
-                ],
-                dim=0,
-            )
-        else:  # if there are no sample points
-            self.condition_indeces = torch.tensor([])
-            self.pts = torch.tensor([])
-
-        self.pts = self.pts.to(device)
-        self.condition_indeces = self.condition_indeces.to(device)
-
-    def __len__(self):
-        return self.pts.shape[0]
-
-
-class DataPointDataset(Dataset):
-
-    def __init__(self, problem, device) -> None:
-        super().__init__()
-        input_list = []
-        output_list = []
-        self.condition_names = []
-
-        for name, condition in problem.conditions.items():
-            if hasattr(condition, "output_points"):
-                input_list.append(problem.conditions[name].input_points)
-                output_list.append(problem.conditions[name].output_points)
-                self.condition_names.append(name)
-
-        self.input_pts = LabelTensor.vstack(input_list)
-        self.output_pts = LabelTensor.vstack(output_list)
-
-        if self.input_pts != []:
-            self.condition_indeces = torch.cat(
-                [
-                    torch.tensor([i] * len(input_list[i]))
-                    for i in range(len(self.condition_names))
-                ],
-                dim=0,
-            )
-        else:  # if there are no data points
-            self.condition_indeces = torch.tensor([])
-            self.input_pts = torch.tensor([])
-            self.output_pts = torch.tensor([])
-
-        self.input_pts = self.input_pts.to(device)
-        self.output_pts = self.output_pts.to(device)
-        self.condition_indeces = self.condition_indeces.to(device)
-
-    def __len__(self):
-        return self.input_pts.shape[0]
 
+from .sample_dataset import SamplePointDataset
+from .data_dataset import DataPointDataset
+from .pina_batch import Batch
 
 class SamplePointLoader:
     """
@@ -133,6 +57,8 @@ def __init__(
         else:
             self.random_idx = torch.arange(len(self.batch_list))
 
+        self._prepare_batches()
+
     def _prepare_data_dataset(self, dataset, batch_size, shuffle):
         """
         Prepare the dataset for data points.
@@ -169,7 +95,7 @@ def _prepare_data_dataset(self, dataset, batch_size, shuffle):
         self.batch_output_pts = torch.tensor_split(
             dataset.output_pts, batch_num
         )
-
+        print(input_labels)
         for i in range(len(self.batch_input_pts)):
             self.batch_input_pts[i].labels = input_labels
             self.batch_output_pts[i].labels = output_labels
@@ -216,6 +142,29 @@ def _prepare_sample_dataset(self, dataset, batch_size, shuffle):
             self.tensor_conditions, batch_num
         )
 
+    def _prepare_batches(self):
+        """
+        Prepare the batches.
+        """
+        self.batches = []
+        for i in range(len(self.batch_list)):
+            type_, idx_ = self.batch_list[i]
+
+            if type_ == "sample":
+                batch = Batch(
+                    "sample", idx_,
+                    self.batch_sample_pts,
+                    self.batch_sample_conditions)
+            else:
+                batch = Batch(
+                    "data", idx_,
+                    self.batch_input_pts,
+                    self.batch_output_pts,
+                    self.batch_data_conditions)
+            print(batch.input.labels)
+                
+            self.batches.append(batch)
+
     def __iter__(self):
         """
         Return an iterator over the points. Any element of the iterator is a
@@ -233,21 +182,24 @@ def __iter__(self):
         :rtype: iter
         """
         # for i in self.random_idx:
-        for i in range(len(self.batch_list)):
-            type_, idx_ = self.batch_list[i]
-
-            if type_ == "sample":
-                d = {
-                    "pts": self.batch_sample_pts[idx_].requires_grad_(True),
-                    "condition": self.batch_sample_conditions[idx_],
-                }
-            else:
-                d = {
-                    "pts": self.batch_input_pts[idx_].requires_grad_(True),
-                    "output": self.batch_output_pts[idx_],
-                    "condition": self.batch_data_conditions[idx_],
-                }
-            yield d
+        for i in self.random_idx:
+            yield self.batches[i]
+
+        # for i in range(len(self.batch_list)):
+        #     type_, idx_ = self.batch_list[i]
+
+        #     if type_ == "sample":
+        #         d = {
+        #             "pts": self.batch_sample_pts[idx_].requires_grad_(True),
+        #             "condition": self.batch_sample_conditions[idx_],
+        #         }
+        #     else:
+        #         d = {
+        #             "pts": self.batch_input_pts[idx_].requires_grad_(True),
+        #             "output": self.batch_output_pts[idx_],
+        #             "condition": self.batch_data_conditions[idx_],
+        #         }
+        #     yield d
 
     def __len__(self):
         """
diff --git a/pina/data/sample_dataset.py b/pina/data/sample_dataset.py
new file mode 100644
index 000000000..84af2920f
--- /dev/null
+++ b/pina/data/sample_dataset.py
@@ -0,0 +1,43 @@
+from torch.utils.data import Dataset
+import torch
+
+from ..label_tensor import LabelTensor
+
+
+class SamplePointDataset(Dataset):
+    """
+    This class is used to create a dataset of sample points.
+    """
+
+    def __init__(self, problem, device) -> None:
+        """
+        :param dict input_pts: The input points.
+        """
+        super().__init__()
+        pts_list = []
+        self.condition_names = []
+
+        for name, condition in problem.conditions.items():
+            if not hasattr(condition, "output_points"):
+                pts_list.append(problem.input_pts[name])
+                self.condition_names.append(name)
+
+        self.pts = LabelTensor.stack(pts_list)
+
+        if self.pts != []:
+            self.condition_indeces = torch.cat(
+                [
+                    torch.tensor([i] * len(pts_list[i]))
+                    for i in range(len(self.condition_names))
+                ],
+                dim=0,
+            )
+        else:  # if there are no sample points
+            self.condition_indeces = torch.tensor([])
+            self.pts = torch.tensor([])
+
+        self.pts = self.pts.to(device)
+        self.condition_indeces = self.condition_indeces.to(device)
+
+    def __len__(self):
+        return self.pts.shape[0]
\ No newline at end of file
diff --git a/pina/geometry/__init__.py b/pina/domain/__init__.py
similarity index 87%
rename from pina/geometry/__init__.py
rename to pina/domain/__init__.py
index 963136a3e..e5a327b4c 100644
--- a/pina/geometry/__init__.py
+++ b/pina/domain/__init__.py
@@ -1,5 +1,5 @@
 __all__ = [
-    "Location",
+    "DomainInterface",
     "CartesianDomain",
     "EllipsoidDomain",
     "Union",
@@ -10,7 +10,7 @@
     "SimplexDomain",
 ]
 
-from .location import Location
+from .domain_interface import DomainInterface
 from .cartesian import CartesianDomain
 from .ellipsoid import EllipsoidDomain
 from .exclusion_domain import Exclusion
diff --git a/pina/geometry/cartesian.py b/pina/domain/cartesian.py
similarity index 99%
rename from pina/geometry/cartesian.py
rename to pina/domain/cartesian.py
index 11354b62f..9ab179bde 100644
--- a/pina/geometry/cartesian.py
+++ b/pina/domain/cartesian.py
@@ -1,11 +1,11 @@
 import torch
 
-from .location import Location
+from .domain_interface import DomainInterface
 from ..label_tensor import LabelTensor
 from ..utils import torch_lhs, chebyshev_roots
 
 
-class CartesianDomain(Location):
+class CartesianDomain(DomainInterface):
     """PINA implementation of Hypercube domain."""
 
     def __init__(self, cartesian_dict):
diff --git a/pina/geometry/difference_domain.py b/pina/domain/difference_domain.py
similarity index 100%
rename from pina/geometry/difference_domain.py
rename to pina/domain/difference_domain.py
diff --git a/pina/geometry/location.py b/pina/domain/domain_interface.py
similarity index 90%
rename from pina/geometry/location.py
rename to pina/domain/domain_interface.py
index a22dfe13f..5906d2851 100644
--- a/pina/geometry/location.py
+++ b/pina/domain/domain_interface.py
@@ -1,9 +1,9 @@
-"""Module for Location class."""
+"""Module for the DomainInterface class."""
 
 from abc import ABCMeta, abstractmethod
 
 
-class Location(metaclass=ABCMeta):
+class DomainInterface(metaclass=ABCMeta):
     """
     Abstract Location class.
     Any geometry entity should inherit from this class.
diff --git a/pina/geometry/ellipsoid.py b/pina/domain/ellipsoid.py
similarity index 99%
rename from pina/geometry/ellipsoid.py
rename to pina/domain/ellipsoid.py
index 2baea5324..eda42d02b 100644
--- a/pina/geometry/ellipsoid.py
+++ b/pina/domain/ellipsoid.py
@@ -1,11 +1,11 @@
 import torch
 
-from .location import Location
+from .domain_interface import DomainInterface
 from ..label_tensor import LabelTensor
 from ..utils import check_consistency
 
 
-class EllipsoidDomain(Location):
+class EllipsoidDomain(DomainInterface):
     """PINA implementation of Ellipsoid domain."""
 
     def __init__(self, ellipsoid_dict, sample_surface=False):
diff --git a/pina/geometry/exclusion_domain.py b/pina/domain/exclusion_domain.py
similarity index 100%
rename from pina/geometry/exclusion_domain.py
rename to pina/domain/exclusion_domain.py
diff --git a/pina/geometry/intersection_domain.py b/pina/domain/intersection_domain.py
similarity index 100%
rename from pina/geometry/intersection_domain.py
rename to pina/domain/intersection_domain.py
diff --git a/pina/geometry/operation_interface.py b/pina/domain/operation_interface.py
similarity index 92%
rename from pina/geometry/operation_interface.py
rename to pina/domain/operation_interface.py
index 4f7709b9a..acd4cf44b 100644
--- a/pina/geometry/operation_interface.py
+++ b/pina/domain/operation_interface.py
@@ -1,11 +1,11 @@
 """ Module for OperationInterface class. """
 
-from .location import Location
+from .domain_interface import DomainInterface
 from ..utils import check_consistency
 from abc import ABCMeta, abstractmethod
 
 
-class OperationInterface(Location, metaclass=ABCMeta):
+class OperationInterface(DomainInterface, metaclass=ABCMeta):
 
     def __init__(self, geometries):
         """
@@ -15,7 +15,7 @@ def __init__(self, geometries):
             such as ``EllipsoidDomain`` or ``CartesianDomain``.
         """
         # check consistency geometries
-        check_consistency(geometries, Location)
+        check_consistency(geometries, DomainInterface)
 
         # check we are passing always different
         # geometries with the same labels.
diff --git a/pina/geometry/simplex.py b/pina/domain/simplex.py
similarity index 98%
rename from pina/geometry/simplex.py
rename to pina/domain/simplex.py
index 15cdc16be..704a9cb3f 100644
--- a/pina/geometry/simplex.py
+++ b/pina/domain/simplex.py
@@ -1,11 +1,11 @@
 import torch
-from .location import Location
-from pina.geometry import CartesianDomain
+from .domain_interface import DomainInterface
+from pina.domain import CartesianDomain
 from ..label_tensor import LabelTensor
 from ..utils import check_consistency
 
 
-class SimplexDomain(Location):
+class SimplexDomain(DomainInterface):
     """PINA implementation of a Simplex."""
 
     def __init__(self, simplex_matrix, sample_surface=False):
diff --git a/pina/geometry/union_domain.py b/pina/domain/union_domain.py
similarity index 100%
rename from pina/geometry/union_domain.py
rename to pina/domain/union_domain.py
diff --git a/pina/label_tensor.py b/pina/label_tensor.py
index 7cb2f71f0..12d2182f5 100644
--- a/pina/label_tensor.py
+++ b/pina/label_tensor.py
@@ -229,10 +229,6 @@
 #             detached._labels = self._labels
 #         return detached
 
-#     def requires_grad_(self, mode=True):
-#         lt = super().requires_grad_(mode)
-#         lt.labels = self.labels
-#         return lt
 
 #     def append(self, lt, mode="std"):
 #         """
@@ -406,11 +402,29 @@ def extract(self, label_to_extract):
 
         return LabelTensor(new_tensor, label_to_extract)
         
-
     def __str__(self):
         s = ''
         for key, value in self.labels.items():
             s += f"{key}: {value}\n"
         s += '\n' 
         s += super().__str__()
-        return s
\ No newline at end of file
+        return s
+    
+    @staticmethod
+    def stack(tensors):
+        """ 
+        """
+        if len(tensors) == 0:
+            return []
+
+        if len(tensors) == 1:
+            return tensors[0]
+        
+        raise NotImplementedError
+        labels = [tensor.labels for tensor in tensors]
+        print(labels)
+
+    def requires_grad_(self, mode=True):
+        lt = super().requires_grad_(mode)
+        lt.labels = self.labels
+        return lt
\ No newline at end of file
diff --git a/pina/problem/abstract_problem.py b/pina/problem/abstract_problem.py
index 6e5e31789..01bb1205c 100644
--- a/pina/problem/abstract_problem.py
+++ b/pina/problem/abstract_problem.py
@@ -5,6 +5,8 @@
 from copy import deepcopy
 import torch
 
+from .. import LabelTensor
+
 
 class AbstractProblem(metaclass=ABCMeta):
     """
@@ -18,17 +20,26 @@ class AbstractProblem(metaclass=ABCMeta):
 
     def __init__(self):
 
-        # variable storing all points
-        self.input_pts = {}
 
-        # varible to check if sampling is done. If no location
-        # element is presented in Condition this variable is set to true
-        self._have_sampled_points = {}
+        self._discretized_domains = {}
+
+        for name, domain in self.domains.items():
+            if isinstance(domain, (torch.Tensor, LabelTensor)):
+                self._discretized_domains[name] = domain
+
         for condition_name in self.conditions:
-            self._have_sampled_points[condition_name] = False
+            self.conditions[condition_name]._problem = self
+        # # variable storing all points
+        # self.input_pts = {}
+
+        # # varible to check if sampling is done. If no location
+        # # element is presented in Condition this variable is set to true
+        # self._have_sampled_points = {}
+        # for condition_name in self.conditions:
+        #     self._have_sampled_points[condition_name] = False
 
-        # put in self.input_pts all the points that we don't need to sample
-        self._span_condition_points()
+        # # put in self.input_pts all the points that we don't need to sample
+        # self._span_condition_points()
 
     def __deepcopy__(self, memo):
         """
@@ -63,15 +74,20 @@ def input_variables(self):
             variables += self.spatial_variables
         if hasattr(self, "temporal_variable"):
             variables += self.temporal_variable
-        if hasattr(self, "parameters"):
+        if hasattr(self, "unknown_parameters"):
             variables += self.parameters
         if hasattr(self, "custom_variables"):
             variables += self.custom_variables
 
         return variables
 
+    @input_variables.setter
+    def input_variables(self, variables):
+        raise RuntimeError
+
     @property
-    def domain(self):
+    @abstractmethod
+    def domains(self):
         """
         The domain(s) where the conditions of the AbstractProblem are valid.
         If more than one domain type is passed, a list of Location is
@@ -80,27 +96,7 @@ def domain(self):
         :return: the domain(s) of ``self``
         :rtype: list[Location]
         """
-        domains = [
-            getattr(self, f"{t}_domain")
-            for t in ["spatial", "temporal", "parameter"]
-            if hasattr(self, f"{t}_domain")
-        ]
-
-        if len(domains) == 1:
-            return domains[0]
-        elif len(domains) == 0:
-            raise RuntimeError
-
-        if len(set(map(type, domains))) == 1:
-            domain = domains[0].__class__({})
-            [domain.update(d) for d in domains]
-            return domain
-        else:
-            raise RuntimeError("different domains")
-
-    @input_variables.setter
-    def input_variables(self, variables):
-        raise RuntimeError
+        pass
 
     @property
     @abstractmethod
@@ -116,7 +112,9 @@ def conditions(self):
         """
         The conditions of the problem.
         """
-        pass
+        return self._conditions
+
+        
 
     def _span_condition_points(self):
         """
@@ -281,28 +279,4 @@ def add_points(self, new_points):
             # merging
             merged_pts = torch.vstack([old_pts, new_pts])
             merged_pts.labels = old_pts.labels
-            self.input_pts[location] = merged_pts
-
-    @property
-    def have_sampled_points(self):
-        """
-        Check if all points for
-        ``Location`` are sampled.
-        """
-        return all(self._have_sampled_points.values())
-
-    @property
-    def not_sampled_points(self):
-        """
-        Check which points are
-        not sampled.
-        """
-        # variables which are not sampled
-        not_sampled = None
-        if self.have_sampled_points is False:
-            # check which one are not sampled:
-            not_sampled = []
-            for condition_name, is_sample in self._have_sampled_points.items():
-                if not is_sample:
-                    not_sampled.append(condition_name)
-        return not_sampled
+            self.input_pts[location] = merged_pts
\ No newline at end of file
diff --git a/pina/solvers/solver.py b/pina/solvers/solver.py
index 0112c860a..a27e93641 100644
--- a/pina/solvers/solver.py
+++ b/pina/solvers/solver.py
@@ -205,6 +205,8 @@ def __init__(
         # put everything in a list if only one input
         if not isinstance(model, list):
             model = [model]
+        if not isinstance(scheduler, list):
+            scheduler = [scheduler]
         if not isinstance(optimizer, list):
             optimizer = [optimizer]
 
diff --git a/pina/solvers/supervised.py b/pina/solvers/supervised.py
index 0285096ba..00a9c8fa0 100644
--- a/pina/solvers/supervised.py
+++ b/pina/solvers/supervised.py
@@ -82,6 +82,7 @@ def __init__(
 
         # check consistency
         check_consistency(loss, (LossInterface, _Loss), subclass=False)
+        self.loss = loss
 
     def forward(self, x):
         """Forward pass implementation for the solver.
@@ -90,7 +91,16 @@ def forward(self, x):
         :return: Solver solution.
         :rtype: torch.Tensor
         """
-        return self._pina_model(x)
+
+        output = self._pina_model[0](x)
+
+        output.labels = {
+            1: {
+                "name": "output",
+                "dof": self.problem.output_variables
+            }
+        }
+        return output
 
     def configure_optimizers(self):
         """Optimizer configuration for the solver.
@@ -98,9 +108,12 @@ def configure_optimizers(self):
         :return: The optimizers and the schedulers
         :rtype: tuple(list, list)
         """
-        self._pina_optimizer.hook(self._pina_model.parameters())
-        self._pina_scheduler.hook(self._pina_optimizer)
-        return self._pina_optimizer, self._pina_scheduler
+        self._pina_optimizer[0].hook(self._pina_model[0].parameters())
+        self._pina_scheduler[0].hook(self._pina_optimizer[0])
+        return (
+            [self._pina_optimizer[0].optimizer_instance], 
+            [self._pina_scheduler[0].scheduler_instance]
+        )
 
     def training_step(self, batch, batch_idx):
         """Solver training step.
@@ -113,14 +126,16 @@ def training_step(self, batch, batch_idx):
         :rtype: LabelTensor
         """
 
-        condition_idx = batch["condition"]
+        condition_idx = batch.condition
 
         for condition_id in range(condition_idx.min(), condition_idx.max() + 1):
 
             condition_name = self._dataloader.condition_names[condition_id]
             condition = self.problem.conditions[condition_name]
-            pts = batch["pts"]
-            out = batch["output"]
+            pts = batch.input
+            out = batch.output
+            print(out)
+            print(pts)
 
             if condition_name not in self.problem.conditions:
                 raise RuntimeError("Something wrong happened.")
@@ -134,9 +149,11 @@ def training_step(self, batch, batch_idx):
             output_pts = out[condition_idx == condition_id]
             input_pts = pts[condition_idx == condition_id]
 
+            input_pts.labels = pts.labels
+            output_pts.labels = out.labels
+
             loss = (
                 self.loss_data(input_pts=input_pts, output_pts=output_pts)
-                * condition.data_weight
             )
             loss = loss.as_subclass(torch.Tensor)
 
@@ -155,6 +172,10 @@ def loss_data(self, input_pts, output_pts):
         :return: The residual loss averaged on the input coordinates
         :rtype: torch.Tensor
         """
+        print(input_pts)
+        print(output_pts)
+        print(self.loss)
+        print(self.forward(input_pts))
         return self.loss(self.forward(input_pts), output_pts)
 
     @property
diff --git a/pina/trainer.py b/pina/trainer.py
index 25d21b77c..758bbaaf0 100644
--- a/pina/trainer.py
+++ b/pina/trainer.py
@@ -3,7 +3,7 @@
 import torch
 import pytorch_lightning
 from .utils import check_consistency
-from .data.dataset import SamplePointDataset, SamplePointLoader, DataPointDataset
+from .data import SamplePointDataset, SamplePointLoader, DataPointDataset
 from .solvers.solver import SolverInterface
 
 
@@ -35,19 +35,33 @@ def __init__(self, solver, batch_size=None, **kwargs):
         self._model = solver
         self.batch_size = batch_size
 
+        self._create_loader()
+        self._move_to_device()
+
         # create dataloader
-        if solver.problem.have_sampled_points is False:
-            raise RuntimeError(
-                f"Input points in {solver.problem.not_sampled_points} "
-                "training are None. Please "
-                "sample points in your problem by calling "
-                "discretise_domain function before train "
-                "in the provided locations."
-            )
-
-        self._create_or_update_loader()
-
-    def _create_or_update_loader(self):
+        # if solver.problem.have_sampled_points is False:
+        #     raise RuntimeError(
+        #         f"Input points in {solver.problem.not_sampled_points} "
+        #         "training are None. Please "
+        #         "sample points in your problem by calling "
+        #         "discretise_domain function before train "
+        #         "in the provided locations."
+        #     )
+
+        # self._create_or_update_loader()
+
+    def _move_to_device(self):
+        device = self._accelerator_connector._parallel_devices[0]
+        
+        # move parameters to device
+        pb = self._model.problem
+        if hasattr(pb, "unknown_parameters"):
+            for key in pb.unknown_parameters:
+                pb.unknown_parameters[key] = torch.nn.Parameter(
+                    pb.unknown_parameters[key].data.to(device)
+                )
+
+    def _create_loader(self):
         """
         This method is used here because is resampling is needed
         during training, there is no need to define to touch the
@@ -64,12 +78,6 @@ def _create_or_update_loader(self):
         self._loader = SamplePointLoader(
             dataset_phys, dataset_data, batch_size=self.batch_size, shuffle=True
         )
-        pb = self._model.problem
-        if hasattr(pb, "unknown_parameters"):
-            for key in pb.unknown_parameters:
-                pb.unknown_parameters[key] = torch.nn.Parameter(
-                    pb.unknown_parameters[key].data.to(device)
-                )
 
     def train(self, **kwargs):
         """
diff --git a/tests/test_callbacks/test_adaptive_refinment_callbacks.py b/tests/test_callbacks/test_adaptive_refinment_callbacks.py
index e5c46a17d..67732b3f5 100644
--- a/tests/test_callbacks/test_adaptive_refinment_callbacks.py
+++ b/tests/test_callbacks/test_adaptive_refinment_callbacks.py
@@ -4,7 +4,7 @@
 
 from pina.problem import SpatialProblem
 from pina.operators import laplacian
-from pina.geometry import CartesianDomain
+from pina.domain import CartesianDomain
 from pina import Condition, LabelTensor
 from pina.solvers import PINN
 from pina.trainer import Trainer
diff --git a/tests/test_callbacks/test_optimizer_callbacks.py b/tests/test_callbacks/test_optimizer_callbacks.py
index 0b0aabaab..898d3f502 100644
--- a/tests/test_callbacks/test_optimizer_callbacks.py
+++ b/tests/test_callbacks/test_optimizer_callbacks.py
@@ -4,7 +4,7 @@
 
 from pina.problem import SpatialProblem
 from pina.operators import laplacian
-from pina.geometry import CartesianDomain
+from pina.domain import CartesianDomain
 from pina import Condition, LabelTensor
 from pina.solvers import PINN
 from pina.trainer import Trainer
diff --git a/tests/test_condition.py b/tests/test_condition.py
index 23c9d126b..5f1c6236b 100644
--- a/tests/test_condition.py
+++ b/tests/test_condition.py
@@ -3,7 +3,7 @@
 
 from pina import LabelTensor, Condition
 from pina.solvers import PINN
-from pina.geometry import CartesianDomain
+from pina.domain import CartesianDomain
 from pina.problem import SpatialProblem
 from pina.model import FeedForward
 from pina.operators import laplacian
diff --git a/tests/test_dataset.py b/tests/test_dataset.py
index cb6a9e4b0..40f219228 100644
--- a/tests/test_dataset.py
+++ b/tests/test_dataset.py
@@ -4,7 +4,7 @@
 from pina.data.dataset import SamplePointDataset, SamplePointLoader, DataPointDataset
 from pina import LabelTensor, Condition
 from pina.equation import Equation
-from pina.geometry import CartesianDomain
+from pina.domain import CartesianDomain
 from pina.problem import SpatialProblem
 from pina.model import FeedForward
 from pina.operators import laplacian
diff --git a/tests/test_geometry/test_cartesian.py b/tests/test_geometry/test_cartesian.py
index 3e7a8c900..65026c332 100644
--- a/tests/test_geometry/test_cartesian.py
+++ b/tests/test_geometry/test_cartesian.py
@@ -1,7 +1,7 @@
 import torch
 
 from pina import LabelTensor
-from pina.geometry import CartesianDomain
+from pina.domain import CartesianDomain
 
 def test_constructor():
     CartesianDomain({'x': [0, 1], 'y': [0, 1]})
diff --git a/tests/test_geometry/test_difference.py b/tests/test_geometry/test_difference.py
index b165fa710..c5300aae0 100644
--- a/tests/test_geometry/test_difference.py
+++ b/tests/test_geometry/test_difference.py
@@ -1,7 +1,7 @@
 import torch
 
 from pina import LabelTensor
-from pina.geometry import Difference, EllipsoidDomain, CartesianDomain
+from pina.domain import Difference, EllipsoidDomain, CartesianDomain
 
 
 def test_constructor_two_CartesianDomains():
diff --git a/tests/test_geometry/test_ellipsoid.py b/tests/test_geometry/test_ellipsoid.py
index 9ab0989ba..fa776f9fc 100644
--- a/tests/test_geometry/test_ellipsoid.py
+++ b/tests/test_geometry/test_ellipsoid.py
@@ -2,7 +2,7 @@
 import pytest
 
 from pina import LabelTensor
-from pina.geometry import EllipsoidDomain
+from pina.domain import EllipsoidDomain
 
 
 def test_constructor():
diff --git a/tests/test_geometry/test_exclusion.py b/tests/test_geometry/test_exclusion.py
index b6400cde6..f11fa7f06 100644
--- a/tests/test_geometry/test_exclusion.py
+++ b/tests/test_geometry/test_exclusion.py
@@ -1,7 +1,7 @@
 import torch
 
 from pina import LabelTensor
-from pina.geometry import Exclusion, EllipsoidDomain, CartesianDomain
+from pina.domain import Exclusion, EllipsoidDomain, CartesianDomain
 
 
 def test_constructor_two_CartesianDomains():
diff --git a/tests/test_geometry/test_intersection.py b/tests/test_geometry/test_intersection.py
index 61061072f..4929cacad 100644
--- a/tests/test_geometry/test_intersection.py
+++ b/tests/test_geometry/test_intersection.py
@@ -1,7 +1,7 @@
 import torch
 
 from pina import LabelTensor
-from pina.geometry import Intersection, EllipsoidDomain, CartesianDomain
+from pina.domain import Intersection, EllipsoidDomain, CartesianDomain
 
 
 def test_constructor_two_CartesianDomains():
diff --git a/tests/test_geometry/test_simplex.py b/tests/test_geometry/test_simplex.py
index 1f59585c6..7fc34ce2d 100644
--- a/tests/test_geometry/test_simplex.py
+++ b/tests/test_geometry/test_simplex.py
@@ -2,7 +2,7 @@
 import pytest
 
 from pina import LabelTensor
-from pina.geometry import SimplexDomain
+from pina.domain import SimplexDomain
 
 
 def test_constructor():
diff --git a/tests/test_geometry/test_union.py b/tests/test_geometry/test_union.py
index 16f8bca2e..acde89542 100644
--- a/tests/test_geometry/test_union.py
+++ b/tests/test_geometry/test_union.py
@@ -1,7 +1,7 @@
 import torch
 
 from pina import LabelTensor
-from pina.geometry import Union, EllipsoidDomain, CartesianDomain
+from pina.domain import Union, EllipsoidDomain, CartesianDomain
 
 
 def test_constructor_two_CartesianDomains():
diff --git a/tests/test_plotter.py b/tests/test_plotter.py
index 99f99bc7e..98eb08845 100644
--- a/tests/test_plotter.py
+++ b/tests/test_plotter.py
@@ -1,4 +1,4 @@
-from pina.geometry import CartesianDomain
+from pina.domain import CartesianDomain
 from pina import Condition, Plotter
 from matplotlib.testing.decorators import image_comparison
 import matplotlib.pyplot as plt
diff --git a/tests/test_problem.py b/tests/test_problem.py
index 09133d4e2..0db508fa3 100644
--- a/tests/test_problem.py
+++ b/tests/test_problem.py
@@ -4,7 +4,7 @@
 from pina.problem import SpatialProblem
 from pina.operators import laplacian
 from pina import LabelTensor, Condition
-from pina.geometry import CartesianDomain
+from pina.domain import CartesianDomain
 from pina.equation.equation import Equation
 from pina.equation.equation_factory import FixedValue
 
diff --git a/tests/test_solvers/test_causalpinn.py b/tests/test_solvers/test_causalpinn.py
index bad5255d3..c5a94a4b0 100644
--- a/tests/test_solvers/test_causalpinn.py
+++ b/tests/test_solvers/test_causalpinn.py
@@ -3,7 +3,7 @@
 
 from pina.problem import TimeDependentProblem, InverseProblem, SpatialProblem
 from pina.operators import grad
-from pina.geometry import CartesianDomain
+from pina.domain import CartesianDomain
 from pina import Condition, LabelTensor
 from pina.solvers import CausalPINN
 from pina.trainer import Trainer
diff --git a/tests/test_solvers/test_competitive_pinn.py b/tests/test_solvers/test_competitive_pinn.py
index fae6d43be..97ee88166 100644
--- a/tests/test_solvers/test_competitive_pinn.py
+++ b/tests/test_solvers/test_competitive_pinn.py
@@ -3,7 +3,7 @@
 
 from pina.problem import SpatialProblem, InverseProblem
 from pina.operators import laplacian
-from pina.geometry import CartesianDomain
+from pina.domain import CartesianDomain
 from pina import Condition, LabelTensor
 from pina.solvers import CompetitivePINN as PINN
 from pina.trainer import Trainer
diff --git a/tests/test_solvers/test_gpinn.py b/tests/test_solvers/test_gpinn.py
index 7c2bb50f6..4a6f9e3c0 100644
--- a/tests/test_solvers/test_gpinn.py
+++ b/tests/test_solvers/test_gpinn.py
@@ -2,7 +2,7 @@
 
 from pina.problem import SpatialProblem, InverseProblem
 from pina.operators import laplacian
-from pina.geometry import CartesianDomain
+from pina.domain import CartesianDomain
 from pina import Condition, LabelTensor
 from pina.solvers import GPINN
 from pina.trainer import Trainer
diff --git a/tests/test_solvers/test_pinn.py b/tests/test_solvers/test_pinn.py
index f3cf275bd..e2a8972d5 100644
--- a/tests/test_solvers/test_pinn.py
+++ b/tests/test_solvers/test_pinn.py
@@ -2,7 +2,7 @@
 
 from pina.problem import SpatialProblem, InverseProblem
 from pina.operators import laplacian
-from pina.geometry import CartesianDomain
+from pina.domain import CartesianDomain
 from pina import Condition, LabelTensor
 from pina.solvers import PINN
 from pina.trainer import Trainer
diff --git a/tests/test_solvers/test_sapinn.py b/tests/test_solvers/test_sapinn.py
index 45475fc42..a29e045e2 100644
--- a/tests/test_solvers/test_sapinn.py
+++ b/tests/test_solvers/test_sapinn.py
@@ -3,7 +3,7 @@
 
 from pina.problem import SpatialProblem, InverseProblem
 from pina.operators import laplacian
-from pina.geometry import CartesianDomain
+from pina.domain import CartesianDomain
 from pina import Condition, LabelTensor
 from pina.solvers import SAPINN as PINN
 from pina.trainer import Trainer
diff --git a/tests/test_solvers/test_supervised_solver.py b/tests/test_solvers/test_supervised_solver.py
index d9cbea3ce..1b2b7d4b7 100644
--- a/tests/test_solvers/test_supervised_solver.py
+++ b/tests/test_solvers/test_supervised_solver.py
@@ -11,10 +11,17 @@
 class NeuralOperatorProblem(AbstractProblem):
     input_variables = ['u_0', 'u_1']
     output_variables = ['u']
+    domains = {
+        'pts': LabelTensor(torch.rand(100, 2), labels={1: {'name': 'space', 'dof': ['u_0', 'u_1']}})
+    }
     conditions = {
-        # 'data' : Condition(
-        #     input_points=LabelTensor(torch.rand(100, 2), input_variables), 
-        #     output_points=LabelTensor(torch.rand(100, 1), output_variables))
+        'data' : Condition(
+            domain='pts', 
+            output_points=LabelTensor(
+                torch.rand(100, 1), 
+                labels={1: {'name': 'output', 'dof': ['u']}}
+            )
+        )
     }
 
 class myFeature(torch.nn.Module):
@@ -31,8 +38,8 @@ def forward(self, x):
         return LabelTensor(t, ['sin(x)sin(y)'])
 
 
-# make the problem + extra feats
 problem = NeuralOperatorProblem()
+# make the problem + extra feats
 extra_feats = [myFeature()]
 model = FeedForward(len(problem.input_variables),
                     len(problem.output_variables))
diff --git a/tests/test_utils.py b/tests/test_utils.py
index 46305f647..46a5083e6 100644
--- a/tests/test_utils.py
+++ b/tests/test_utils.py
@@ -3,10 +3,10 @@
 from pina.utils import merge_tensors
 from pina.label_tensor import LabelTensor
 from pina import LabelTensor
-from pina.geometry import EllipsoidDomain, CartesianDomain
+from pina.domain import EllipsoidDomain, CartesianDomain
 from pina.utils import check_consistency
 import pytest
-from pina.geometry import Location
+from pina.domain import Location
 
 
 def test_merge_tensors():
diff --git a/tutorials/tutorial1/tutorial.py b/tutorials/tutorial1/tutorial.py
index aa18b7fd8..a5b5c29e7 100644
--- a/tutorials/tutorial1/tutorial.py
+++ b/tutorials/tutorial1/tutorial.py
@@ -66,7 +66,7 @@
   get_ipython().system('pip install "pina-mathlab"')
 
 from pina.problem import SpatialProblem, TimeDependentProblem
-from pina.geometry import CartesianDomain
+from pina.domain import CartesianDomain
 
 class TimeSpaceODE(SpatialProblem, TimeDependentProblem):
     
@@ -95,7 +95,7 @@ class TimeSpaceODE(SpatialProblem, TimeDependentProblem):
 from pina.problem import SpatialProblem
 from pina.operators import grad
 from pina import Condition
-from pina.geometry import CartesianDomain
+from pina.domain import CartesianDomain
 from pina.equation import Equation, FixedValue
 
 import torch
diff --git a/tutorials/tutorial11/tutorial.py b/tutorials/tutorial11/tutorial.py
index 9bbabfea6..b9d3a483c 100644
--- a/tutorials/tutorial11/tutorial.py
+++ b/tutorials/tutorial11/tutorial.py
@@ -32,7 +32,7 @@
 from pina.model import FeedForward
 from pina.problem import SpatialProblem
 from pina.operators import grad
-from pina.geometry import CartesianDomain
+from pina.domain import CartesianDomain
 from pina.equation import Equation, FixedValue
 
 class SimpleODE(SpatialProblem):
diff --git a/tutorials/tutorial12/tutorial.py b/tutorials/tutorial12/tutorial.py
index 515841a4e..d6aeb9b64 100644
--- a/tutorials/tutorial12/tutorial.py
+++ b/tutorials/tutorial12/tutorial.py
@@ -41,7 +41,7 @@
 #useful imports
 from pina.problem import SpatialProblem, TimeDependentProblem
 from pina.equation import Equation, FixedValue, FixedGradient, FixedFlux
-from pina.geometry import CartesianDomain
+from pina.domain import CartesianDomain
 import torch
 from pina.operators import grad, laplacian
 from pina import Condition
diff --git a/tutorials/tutorial2/tutorial.py b/tutorials/tutorial2/tutorial.py
index b5132b442..4315bfc26 100644
--- a/tutorials/tutorial2/tutorial.py
+++ b/tutorials/tutorial2/tutorial.py
@@ -30,7 +30,7 @@
 from pina.solvers import PINN
 from pina.trainer import Trainer
 from pina.plotter import Plotter
-from pina.geometry import CartesianDomain
+from pina.domain import CartesianDomain
 from pina.equation import Equation, FixedValue
 from pina import Condition, LabelTensor
 from pina.callbacks import MetricTracker
diff --git a/tutorials/tutorial3/tutorial.py b/tutorials/tutorial3/tutorial.py
index bc2a8f697..445134a1a 100644
--- a/tutorials/tutorial3/tutorial.py
+++ b/tutorials/tutorial3/tutorial.py
@@ -25,7 +25,7 @@
 
 from pina.problem import SpatialProblem, TimeDependentProblem
 from pina.operators import laplacian, grad
-from pina.geometry import CartesianDomain
+from pina.domain import CartesianDomain
 from pina.solvers import PINN
 from pina.trainer import Trainer
 from pina.equation import Equation
diff --git a/tutorials/tutorial7/tutorial.py b/tutorials/tutorial7/tutorial.py
index 3c55f1ca7..419dbbe83 100644
--- a/tutorials/tutorial7/tutorial.py
+++ b/tutorials/tutorial7/tutorial.py
@@ -51,7 +51,7 @@
 from pina.equation import Equation, FixedValue
 from pina import Condition, Trainer
 from pina.solvers import PINN
-from pina.geometry import CartesianDomain
+from pina.domain import CartesianDomain
 
 
 # Then, we import the pre-saved data, for ($\mu_1$, $\mu_2$)=($0.5$, $0.5$). These two values are the optimal parameters that we want to find through the neural network training. In particular, we import the `input_points`(the spatial coordinates), and the `output_points` (the corresponding $u$ values evaluated at the `input_points`).
diff --git a/tutorials/tutorial8/tutorial.py b/tutorials/tutorial8/tutorial.py
index 980404e4e..b1d105a06 100644
--- a/tutorials/tutorial8/tutorial.py
+++ b/tutorials/tutorial8/tutorial.py
@@ -35,7 +35,7 @@
 import torch
 import pina
 
-from pina.geometry import CartesianDomain
+from pina.domain import CartesianDomain
 
 from pina.problem import ParametricProblem
 from pina.model.layers import PODBlock, RBFBlock
diff --git a/tutorials/tutorial9/tutorial.py b/tutorials/tutorial9/tutorial.py
index db4b7a333..c1590870a 100644
--- a/tutorials/tutorial9/tutorial.py
+++ b/tutorials/tutorial9/tutorial.py
@@ -36,7 +36,7 @@
 from pina.model.layers import PeriodicBoundaryEmbedding  # The PBC module
 from pina.solvers import PINN
 from pina.trainer import Trainer
-from pina.geometry import CartesianDomain
+from pina.domain import CartesianDomain
 from pina.equation import Equation
 
 

From f09d33bea80e8f78fbaf64ef5f91dceb8c60d04b Mon Sep 17 00:00:00 2001
From: Nicola Demo <demo.nicola@gmail.com>
Date: Mon, 9 Sep 2024 10:50:54 +0200
Subject: [PATCH 08/55] refact

---
 pina/graph.py                                | 118 +++++++++++
 pina/label_tensor.py                         |   6 +-
 pina/loss.py                                 | 209 -------------------
 pina/loss/__init__.py                        |   9 +
 pina/loss/loss_interface.py                  |  61 ++++++
 pina/loss/lp_loss.py                         |  78 +++++++
 pina/loss/power_loss.py                      |  79 +++++++
 pina/loss/weighted_aggregation.py            |  35 ++++
 pina/loss/weightning_interface.py            |  24 +++
 pina/model/layers/messa_passing.py           |  11 +
 pina/solvers/pinns/basepinn.py               |   2 +-
 pina/solvers/supervised.py                   |   6 +-
 tests/test_loss/test_lploss.py               |   2 +-
 tests/test_loss/test_powerloss.py            |   2 +-
 tests/test_solvers/test_causalpinn.py        |   2 +-
 tests/test_solvers/test_competitive_pinn.py  |   2 +-
 tests/test_solvers/test_gpinn.py             |   2 +-
 tests/test_solvers/test_pinn.py              |   2 +-
 tests/test_solvers/test_rom_solver.py        |   2 +-
 tests/test_solvers/test_sapinn.py            |   2 +-
 tests/test_solvers/test_supervised_solver.py |  51 ++++-
 tutorials/tutorial10/tutorial.py             |   2 +-
 tutorials/tutorial5/tutorial.py              |   2 +-
 23 files changed, 480 insertions(+), 229 deletions(-)
 create mode 100644 pina/graph.py
 delete mode 100644 pina/loss.py
 create mode 100644 pina/loss/__init__.py
 create mode 100644 pina/loss/loss_interface.py
 create mode 100644 pina/loss/lp_loss.py
 create mode 100644 pina/loss/power_loss.py
 create mode 100644 pina/loss/weighted_aggregation.py
 create mode 100644 pina/loss/weightning_interface.py
 create mode 100644 pina/model/layers/messa_passing.py

diff --git a/pina/graph.py b/pina/graph.py
new file mode 100644
index 000000000..97b2770e6
--- /dev/null
+++ b/pina/graph.py
@@ -0,0 +1,118 @@
+""" Module for Loss class """
+
+import logging
+from torch_geometric.nn import MessagePassing, InstanceNorm, radius_graph
+from torch_geometric.data import Data
+import torch
+
+class Graph:
+    """
+    PINA Graph managing the PyG Data class.
+    """
+    def __init__(self, data):
+        self.data = data
+            
+    @staticmethod
+    def _build_triangulation(**kwargs):
+        logging.debug("Creating graph with triangulation mode.")
+
+        # check for mandatory arguments
+        if "nodes_coordinates" not in kwargs:
+            raise ValueError("Nodes coordinates must be provided in the kwargs.")
+        if "nodes_data" not in kwargs:
+            raise ValueError("Nodes data must be provided in the kwargs.")
+        if "triangles" not in kwargs:
+            raise ValueError("Triangles must be provided in the kwargs.")
+
+        nodes_coordinates = kwargs["nodes_coordinates"]
+        nodes_data = kwargs["nodes_data"]
+        triangles = kwargs["triangles"]
+
+
+
+        def less_first(a, b):
+            return [a, b] if a < b else [b, a]
+
+        list_of_edges = []
+
+        for triangle in triangles:
+            for e1, e2 in [[0, 1], [1, 2], [2, 0]]:
+                list_of_edges.append(less_first(triangle[e1],triangle[e2]))
+
+        array_of_edges = torch.unique(torch.Tensor(list_of_edges), dim=0) # remove duplicates
+        array_of_edges = array_of_edges.t().contiguous()
+        print(array_of_edges)
+
+        # list_of_lengths = []
+
+        # for p1,p2 in array_of_edges:
+        #     x1, y1 = tri.points[p1]
+        #     x2, y2 = tri.points[p2]
+        #     list_of_lengths.append((x1-x2)**2 + (y1-y2)**2)
+
+        # array_of_lengths = np.sqrt(np.array(list_of_lengths))
+
+        # return array_of_edges, array_of_lengths
+
+        return Data(
+            x=nodes_data,
+            pos=nodes_coordinates.T,
+            
+            edge_index=array_of_edges,
+        )
+
+    @staticmethod
+    def _build_radius(**kwargs):
+        logging.debug("Creating graph with radius mode.")
+
+        # check for mandatory arguments
+        if "nodes_coordinates" not in kwargs:
+            raise ValueError("Nodes coordinates must be provided in the kwargs.")
+        if "nodes_data" not in kwargs:
+            raise ValueError("Nodes data must be provided in the kwargs.")
+        if "radius" not in kwargs:
+            raise ValueError("Radius must be provided in the kwargs.")
+
+        nodes_coordinates = kwargs["nodes_coordinates"]
+        nodes_data = kwargs["nodes_data"]
+        radius = kwargs["radius"]
+
+        edges_data = kwargs.get("edge_data", None) 
+        loop = kwargs.get("loop", False)
+        batch = kwargs.get("batch", None)
+
+        logging.debug(f"radius: {radius}, loop: {loop}, "
+                        f"batch: {batch}")
+
+        edge_index = radius_graph(
+            x=nodes_coordinates.tensor,
+            r=radius,
+            loop=loop,
+            batch=batch,
+        )
+
+        logging.debug(f"edge_index computed")
+        return Data(
+            x=nodes_data,
+            pos=nodes_coordinates,
+            edge_index=edge_index,
+            edge_attr=edges_data,
+        )
+
+    @staticmethod
+    def build(mode, **kwargs):
+        """
+        Constructor for the `Graph` class.
+        """
+        if mode == "radius":
+            graph = Graph._build_radius(**kwargs)
+        elif mode == "triangulation":
+            graph = Graph._build_triangulation(**kwargs)
+        else:
+            raise ValueError(f"Mode {mode} not recognized")
+        
+        return Graph(graph)
+
+
+    def __repr__(self):
+        return f"Graph(data={self.data})"
\ No newline at end of file
diff --git a/pina/label_tensor.py b/pina/label_tensor.py
index 12d2182f5..0a5515365 100644
--- a/pina/label_tensor.py
+++ b/pina/label_tensor.py
@@ -427,4 +427,8 @@ def stack(tensors):
     def requires_grad_(self, mode=True):
         lt = super().requires_grad_(mode)
         lt.labels = self.labels
-        return lt
\ No newline at end of file
+        return lt
+
+    @property
+    def dtype(self):
+        return super().dtype
\ No newline at end of file
diff --git a/pina/loss.py b/pina/loss.py
deleted file mode 100644
index 3cbf88880..000000000
--- a/pina/loss.py
+++ /dev/null
@@ -1,209 +0,0 @@
-""" Module for Loss class """
-
-from abc import ABCMeta, abstractmethod
-from torch.nn.modules.loss import _Loss
-import torch
-from .utils import check_consistency
-
-__all__ = ["LossInterface", "LpLoss", "PowerLoss"]
-
-
-class LossInterface(_Loss, metaclass=ABCMeta):
-    """
-    The abstract ``LossInterface`` class. All the class defining a PINA Loss
-    should be inheritied from this class.
-    """
-
-    def __init__(self, reduction="mean"):
-        """
-        :param str reduction: Specifies the reduction to apply to the output:
-            ``none`` | ``mean`` | ``sum``. When ``none``: no reduction
-            will be applied, ``mean``: the sum of the output will be divided
-            by the number of elements in the output, ``sum``: the output will
-            be summed. Note: ``size_average`` and ``reduce`` are in the
-            process of being deprecated, and in the meantime, specifying either of
-            those two args will override ``reduction``. Default: ``mean``.
-        """
-        super().__init__(reduction=reduction, size_average=None, reduce=None)
-
-    @abstractmethod
-    def forward(self, input, target):
-        """Forward method for loss function.
-
-        :param torch.Tensor input: Input tensor from real data.
-        :param torch.Tensor target: Model tensor output.
-        :return: Loss evaluation.
-        :rtype: torch.Tensor
-        """
-        pass
-
-    def _reduction(self, loss):
-        """Simple helper function to check reduction
-
-        :param reduction: Specifies the reduction to apply to the output:
-            ``none`` | ``mean`` | ``sum``. When ``none``: no reduction
-            will be applied, ``mean``: the sum of the output will be divided
-            by the number of elements in the output, ``sum``: the output will
-            be summed. Note: ``size_average`` and ``reduce`` are in the
-            process of being deprecated, and in the meantime, specifying either of
-            those two args will override ``reduction``. Default: ``mean``.
-        :type reduction: str
-        :param loss: Loss tensor for each element.
-        :type loss: torch.Tensor
-        :return: Reduced loss.
-        :rtype: torch.Tensor
-        """
-        if self.reduction == "none":
-            ret = loss
-        elif self.reduction == "mean":
-            ret = torch.mean(loss, keepdim=True, dim=-1)
-        elif self.reduction == "sum":
-            ret = torch.sum(loss, keepdim=True, dim=-1)
-        else:
-            raise ValueError(self.reduction + " is not valid")
-        return ret
-
-
-class LpLoss(LossInterface):
-    r"""
-    The Lp loss implementation class. Creates a criterion that measures
-    the Lp error between each element in the input :math:`x` and
-    target :math:`y`.
-
-    The unreduced (i.e. with ``reduction`` set to ``none``) loss can
-    be described as:
-
-    .. math::
-        \ell(x, y) = L = \{l_1,\dots,l_N\}^\top, \quad
-        l_n = \left[\sum_{i=1}^{D} \left| x_n^i - y_n^i \right|^p \right],
-    
-    If ``'relative'`` is set to true:
-
-    .. math::
-        \ell(x, y) = L = \{l_1,\dots,l_N\}^\top, \quad
-        l_n = \frac{ [\sum_{i=1}^{D} | x_n^i - y_n^i|^p] }{[\sum_{i=1}^{D}|y_n^i|^p]},
-
-    where :math:`N` is the batch size. If ``reduction`` is not ``none``
-    (default ``mean``), then:
-
-    .. math::
-        \ell(x, y) =
-        \begin{cases}
-            \operatorname{mean}(L), & \text{if reduction} = \text{`mean';}\\
-            \operatorname{sum}(L),  & \text{if reduction} = \text{`sum'.}
-        \end{cases}
-
-    :math:`x` and :math:`y` are tensors of arbitrary shapes with a total
-    of :math:`n` elements each.
-
-    The sum operation still operates over all the elements, and divides by :math:`n`.
-
-    The division by :math:`n` can be avoided if one sets ``reduction`` to ``sum``.
-    """
-
-    def __init__(self, p=2, reduction="mean", relative=False):
-        """
-        :param int p: Degree of Lp norm. It specifies the type of norm to
-            be calculated. See `list of possible orders in torch linalg
-            <https://pytorch.org/docs/stable/generated/torch.linalg.norm.html#torch.linalg.norm>`_ to
-            for possible degrees. Default 2 (euclidean norm).
-        :param str reduction: Specifies the reduction to apply to the output:
-            ``none`` | ``mean`` | ``sum``. ``none``: no reduction
-            will be applied, ``mean``: the sum of the output will be divided
-            by the number of elements in the output, ``sum``: the output will
-            be summed.
-        :param bool relative: Specifies if relative error should be computed.
-        """
-        super().__init__(reduction=reduction)
-
-        # check consistency
-        check_consistency(p, (str, int, float))
-        check_consistency(relative, bool)
-
-        self.p = p
-        self.relative = relative
-
-    def forward(self, input, target):
-        """Forward method for loss function.
-
-        :param torch.Tensor input: Input tensor from real data.
-        :param torch.Tensor target: Model tensor output.
-        :return: Loss evaluation.
-        :rtype: torch.Tensor
-        """
-        loss = torch.linalg.norm((input - target), ord=self.p, dim=-1)
-        if self.relative:
-            loss = loss / torch.linalg.norm(input, ord=self.p, dim=-1)
-        return self._reduction(loss)
-
-
-class PowerLoss(LossInterface):
-    r"""
-    The PowerLoss loss implementation class. Creates a criterion that measures
-    the error between each element in the input :math:`x` and
-    target :math:`y` powered to a specific integer.
-
-    The unreduced (i.e. with ``reduction`` set to ``none``) loss can
-    be described as:
-
-    .. math::
-        \ell(x, y) = L = \{l_1,\dots,l_N\}^\top, \quad
-        l_n = \frac{1}{D}\left[\sum_{i=1}^{D} \left| x_n^i - y_n^i \right|^p \right],
-    
-    If ``'relative'`` is set to true:
-
-    .. math::
-        \ell(x, y) = L = \{l_1,\dots,l_N\}^\top, \quad
-        l_n = \frac{ \sum_{i=1}^{D} | x_n^i - y_n^i|^p }{\sum_{i=1}^{D}|y_n^i|^p},
-
-    where :math:`N` is the batch size. If ``reduction`` is not ``none``
-    (default ``mean``), then:
-
-    .. math::
-        \ell(x, y) =
-        \begin{cases}
-            \operatorname{mean}(L), & \text{if reduction} = \text{`mean';}\\
-            \operatorname{sum}(L),  & \text{if reduction} = \text{`sum'.}
-        \end{cases}
-
-    :math:`x` and :math:`y` are tensors of arbitrary shapes with a total
-    of :math:`n` elements each.
-
-    The sum operation still operates over all the elements, and divides by :math:`n`.
-
-    The division by :math:`n` can be avoided if one sets ``reduction`` to ``sum``.
-    """
-
-    def __init__(self, p=2, reduction="mean", relative=False):
-        """
-        :param int p: Degree of Lp norm. It specifies the type of norm to
-            be calculated. See `list of possible orders in torch linalg
-            <https://pytorch.org/docs/stable/generated/torch.linalg.norm.html#torch.linalg.norm>`_ to
-            see the possible degrees. Default 2 (euclidean norm).
-        :param str reduction: Specifies the reduction to apply to the output:
-            ``none`` | ``mean`` | ``sum``. When ``none``: no reduction
-            will be applied, ``mean``: the sum of the output will be divided
-            by the number of elements in the output, ``sum``: the output will
-            be summed.
-        :param bool relative: Specifies if relative error should be computed.
-        """
-        super().__init__(reduction=reduction)
-
-        # check consistency
-        check_consistency(p, (str, int, float))
-        self.p = p
-        check_consistency(relative, bool)
-        self.relative = relative
-
-    def forward(self, input, target):
-        """Forward method for loss function.
-
-        :param torch.Tensor input: Input tensor from real data.
-        :param torch.Tensor target: Model tensor output.
-        :return: Loss evaluation.
-        :rtype: torch.Tensor
-        """
-        loss = torch.abs((input - target)).pow(self.p).mean(-1)
-        if self.relative:
-            loss = loss / torch.abs(input).pow(self.p).mean(-1)
-        return self._reduction(loss)
diff --git a/pina/loss/__init__.py b/pina/loss/__init__.py
new file mode 100644
index 000000000..35138125d
--- /dev/null
+++ b/pina/loss/__init__.py
@@ -0,0 +1,9 @@
+__all__ = [
+    'LpLoss',
+
+]
+
+from .loss_interface import LossInterface
+from .power_loss import PowerLoss
+from .lp_loss import LpLoss
+from .weightning_interface import weightningInterface
\ No newline at end of file
diff --git a/pina/loss/loss_interface.py b/pina/loss/loss_interface.py
new file mode 100644
index 000000000..5093a65d4
--- /dev/null
+++ b/pina/loss/loss_interface.py
@@ -0,0 +1,61 @@
+""" Module for Loss Interface """
+
+from abc import ABCMeta, abstractmethod
+from torch.nn.modules.loss import _Loss
+import torch
+
+
+class LossInterface(_Loss, metaclass=ABCMeta):
+    """
+    The abstract ``LossInterface`` class. All the class defining a PINA Loss
+    should be inheritied from this class.
+    """
+
+    def __init__(self, reduction="mean"):
+        """
+        :param str reduction: Specifies the reduction to apply to the output:
+            ``none`` | ``mean`` | ``sum``. When ``none``: no reduction
+            will be applied, ``mean``: the sum of the output will be divided
+            by the number of elements in the output, ``sum``: the output will
+            be summed. Note: ``size_average`` and ``reduce`` are in the
+            process of being deprecated, and in the meantime, specifying either of
+            those two args will override ``reduction``. Default: ``mean``.
+        """
+        super().__init__(reduction=reduction, size_average=None, reduce=None)
+
+    @abstractmethod
+    def forward(self, input, target):
+        """Forward method for loss function.
+
+        :param torch.Tensor input: Input tensor from real data.
+        :param torch.Tensor target: Model tensor output.
+        :return: Loss evaluation.
+        :rtype: torch.Tensor
+        """
+        pass
+
+    def _reduction(self, loss):
+        """Simple helper function to check reduction
+
+        :param reduction: Specifies the reduction to apply to the output:
+            ``none`` | ``mean`` | ``sum``. When ``none``: no reduction
+            will be applied, ``mean``: the sum of the output will be divided
+            by the number of elements in the output, ``sum``: the output will
+            be summed. Note: ``size_average`` and ``reduce`` are in the
+            process of being deprecated, and in the meantime, specifying either of
+            those two args will override ``reduction``. Default: ``mean``.
+        :type reduction: str
+        :param loss: Loss tensor for each element.
+        :type loss: torch.Tensor
+        :return: Reduced loss.
+        :rtype: torch.Tensor
+        """
+        if self.reduction == "none":
+            ret = loss
+        elif self.reduction == "mean":
+            ret = torch.mean(loss, keepdim=True, dim=-1)
+        elif self.reduction == "sum":
+            ret = torch.sum(loss, keepdim=True, dim=-1)
+        else:
+            raise ValueError(self.reduction + " is not valid")
+        return ret
\ No newline at end of file
diff --git a/pina/loss/lp_loss.py b/pina/loss/lp_loss.py
new file mode 100644
index 000000000..978efa853
--- /dev/null
+++ b/pina/loss/lp_loss.py
@@ -0,0 +1,78 @@
+""" Module for LpLoss class """
+
+import torch
+
+from ..utils import check_consistency
+from .loss_interface import LossInterface
+
+class LpLoss(LossInterface):
+    r"""
+    The Lp loss implementation class. Creates a criterion that measures
+    the Lp error between each element in the input :math:`x` and
+    target :math:`y`.
+
+    The unreduced (i.e. with ``reduction`` set to ``none``) loss can
+    be described as:
+
+    .. math::
+        \ell(x, y) = L = \{l_1,\dots,l_N\}^\top, \quad
+        l_n = \left[\sum_{i=1}^{D} \left| x_n^i - y_n^i \right|^p \right],
+    
+    If ``'relative'`` is set to true:
+
+    .. math::
+        \ell(x, y) = L = \{l_1,\dots,l_N\}^\top, \quad
+        l_n = \frac{ [\sum_{i=1}^{D} | x_n^i - y_n^i|^p] }{[\sum_{i=1}^{D}|y_n^i|^p]},
+
+    where :math:`N` is the batch size. If ``reduction`` is not ``none``
+    (default ``mean``), then:
+
+    .. math::
+        \ell(x, y) =
+        \begin{cases}
+            \operatorname{mean}(L), & \text{if reduction} = \text{`mean';}\\
+            \operatorname{sum}(L),  & \text{if reduction} = \text{`sum'.}
+        \end{cases}
+
+    :math:`x` and :math:`y` are tensors of arbitrary shapes with a total
+    of :math:`n` elements each.
+
+    The sum operation still operates over all the elements, and divides by :math:`n`.
+
+    The division by :math:`n` can be avoided if one sets ``reduction`` to ``sum``.
+    """
+
+    def __init__(self, p=2, reduction="mean", relative=False):
+        """
+        :param int p: Degree of Lp norm. It specifies the type of norm to
+            be calculated. See `list of possible orders in torch linalg
+            <https://pytorch.org/docs/stable/generated/torch.linalg.norm.html#torch.linalg.norm>`_ to
+            for possible degrees. Default 2 (euclidean norm).
+        :param str reduction: Specifies the reduction to apply to the output:
+            ``none`` | ``mean`` | ``sum``. ``none``: no reduction
+            will be applied, ``mean``: the sum of the output will be divided
+            by the number of elements in the output, ``sum``: the output will
+            be summed.
+        :param bool relative: Specifies if relative error should be computed.
+        """
+        super().__init__(reduction=reduction)
+
+        # check consistency
+        check_consistency(p, (str, int, float))
+        check_consistency(relative, bool)
+
+        self.p = p
+        self.relative = relative
+
+    def forward(self, input, target):
+        """Forward method for loss function.
+
+        :param torch.Tensor input: Input tensor from real data.
+        :param torch.Tensor target: Model tensor output.
+        :return: Loss evaluation.
+        :rtype: torch.Tensor
+        """
+        loss = torch.linalg.norm((input - target), ord=self.p, dim=-1)
+        if self.relative:
+            loss = loss / torch.linalg.norm(input, ord=self.p, dim=-1)
+        return self._reduction(loss)
\ No newline at end of file
diff --git a/pina/loss/power_loss.py b/pina/loss/power_loss.py
new file mode 100644
index 000000000..4f3fc65c7
--- /dev/null
+++ b/pina/loss/power_loss.py
@@ -0,0 +1,79 @@
+""" Module for PowerLoss class """
+
+import torch
+
+from ..utils import check_consistency
+from .loss_interface import LossInterface
+
+
+class PowerLoss(LossInterface):
+    r"""
+    The PowerLoss loss implementation class. Creates a criterion that measures
+    the error between each element in the input :math:`x` and
+    target :math:`y` powered to a specific integer.
+
+    The unreduced (i.e. with ``reduction`` set to ``none``) loss can
+    be described as:
+
+    .. math::
+        \ell(x, y) = L = \{l_1,\dots,l_N\}^\top, \quad
+        l_n = \frac{1}{D}\left[\sum_{i=1}^{D} \left| x_n^i - y_n^i \right|^p \right],
+    
+    If ``'relative'`` is set to true:
+
+    .. math::
+        \ell(x, y) = L = \{l_1,\dots,l_N\}^\top, \quad
+        l_n = \frac{ \sum_{i=1}^{D} | x_n^i - y_n^i|^p }{\sum_{i=1}^{D}|y_n^i|^p},
+
+    where :math:`N` is the batch size. If ``reduction`` is not ``none``
+    (default ``mean``), then:
+
+    .. math::
+        \ell(x, y) =
+        \begin{cases}
+            \operatorname{mean}(L), & \text{if reduction} = \text{`mean';}\\
+            \operatorname{sum}(L),  & \text{if reduction} = \text{`sum'.}
+        \end{cases}
+
+    :math:`x` and :math:`y` are tensors of arbitrary shapes with a total
+    of :math:`n` elements each.
+
+    The sum operation still operates over all the elements, and divides by :math:`n`.
+
+    The division by :math:`n` can be avoided if one sets ``reduction`` to ``sum``.
+    """
+
+    def __init__(self, p=2, reduction="mean", relative=False):
+        """
+        :param int p: Degree of Lp norm. It specifies the type of norm to
+            be calculated. See `list of possible orders in torch linalg
+            <https://pytorch.org/docs/stable/generated/torch.linalg.norm.html#torch.linalg.norm>`_ to
+            see the possible degrees. Default 2 (euclidean norm).
+        :param str reduction: Specifies the reduction to apply to the output:
+            ``none`` | ``mean`` | ``sum``. When ``none``: no reduction
+            will be applied, ``mean``: the sum of the output will be divided
+            by the number of elements in the output, ``sum``: the output will
+            be summed.
+        :param bool relative: Specifies if relative error should be computed.
+        """
+        super().__init__(reduction=reduction)
+
+        # check consistency
+        check_consistency(p, (str, int, float))
+        check_consistency(relative, bool)
+
+        self.p = p
+        self.relative = relative
+
+    def forward(self, input, target):
+        """Forward method for loss function.
+
+        :param torch.Tensor input: Input tensor from real data.
+        :param torch.Tensor target: Model tensor output.
+        :return: Loss evaluation.
+        :rtype: torch.Tensor
+        """
+        loss = torch.abs((input - target)).pow(self.p).mean(-1)
+        if self.relative:
+            loss = loss / torch.abs(input).pow(self.p).mean(-1)
+        return self._reduction(loss)
\ No newline at end of file
diff --git a/pina/loss/weighted_aggregation.py b/pina/loss/weighted_aggregation.py
new file mode 100644
index 000000000..357ebf5b9
--- /dev/null
+++ b/pina/loss/weighted_aggregation.py
@@ -0,0 +1,35 @@
+""" Module for Loss Interface """
+
+from .weightning_interface import weightningInterface
+
+
+class WeightedAggregation(WeightningInterface):
+    """
+    TODO
+    """
+    def __init__(self, aggr='mean', weights=None):
+        self.aggr = aggr
+        self.weights = weights
+
+    def aggregate(self, losses):
+        """
+        Aggregate the losses.
+
+        :param dict(torch.Tensor) input: The dictionary of losses.
+        :return: The losses aggregation. It should be a scalar Tensor.
+        :rtype: torch.Tensor
+        """
+        if self.weights:
+            weighted_losses = {
+                condition: self.weights[condition] * losses[condition] 
+                for condition in losses
+            }
+        else:
+            weighted_losses = losses
+
+        if self.aggr == 'mean':
+            return sum(weighted_losses.values()) / len(weighted_losses)
+        elif self.aggr == 'sum':
+            return sum(weighted_losses.values())
+        else:
+            raise ValueError(self.aggr + " is not valid for aggregation.")
diff --git a/pina/loss/weightning_interface.py b/pina/loss/weightning_interface.py
new file mode 100644
index 000000000..904a819a8
--- /dev/null
+++ b/pina/loss/weightning_interface.py
@@ -0,0 +1,24 @@
+""" Module for Loss Interface """
+
+from abc import ABCMeta, abstractmethod
+
+
+class weightningInterface(metaclass=ABCMeta):
+    """
+    The ``weightingInterface`` class. TODO
+    """
+
+    @abstractmethod
+    def __init__(self, *args, **kwargs):
+        pass
+
+    @abstractmethod
+    def aggregate(self, losses):
+        """
+        Aggregate the losses.
+
+        :param list(torch.Tensor) input: The list
+        :return: The losses aggregation. It should be a scalar Tensor.
+        :rtype: torch.Tensor
+        """
+        pass
\ No newline at end of file
diff --git a/pina/model/layers/messa_passing.py b/pina/model/layers/messa_passing.py
new file mode 100644
index 000000000..8599734df
--- /dev/null
+++ b/pina/model/layers/messa_passing.py
@@ -0,0 +1,11 @@
+""" Module for Averaging Neural Operator Layer class. """
+
+from torch import nn, mean
+from torch_geometric.nn import MessagePassing, InstanceNorm, radius_graph
+
+from pina.utils import check_consistency
+
+
+class MessagePassingBlock(nn.Module):
+
+
diff --git a/pina/solvers/pinns/basepinn.py b/pina/solvers/pinns/basepinn.py
index 0f82056a4..a15834da4 100644
--- a/pina/solvers/pinns/basepinn.py
+++ b/pina/solvers/pinns/basepinn.py
@@ -6,7 +6,7 @@
 
 from ...solvers.solver import SolverInterface
 from pina.utils import check_consistency
-from pina.loss import LossInterface
+from pina.loss.loss_interface import LossInterface
 from pina.problem import InverseProblem
 from torch.nn.modules.loss import _Loss
 
diff --git a/pina/solvers/supervised.py b/pina/solvers/supervised.py
index 00a9c8fa0..4f3d497fe 100644
--- a/pina/solvers/supervised.py
+++ b/pina/solvers/supervised.py
@@ -8,7 +8,7 @@
 from .solver import SolverInterface
 from ..label_tensor import LabelTensor
 from ..utils import check_consistency
-from ..loss import LossInterface
+from ..loss.loss_interface import LossInterface
 
 
 class SupervisedSolver(SolverInterface):
@@ -172,10 +172,6 @@ def loss_data(self, input_pts, output_pts):
         :return: The residual loss averaged on the input coordinates
         :rtype: torch.Tensor
         """
-        print(input_pts)
-        print(output_pts)
-        print(self.loss)
-        print(self.forward(input_pts))
         return self.loss(self.forward(input_pts), output_pts)
 
     @property
diff --git a/tests/test_loss/test_lploss.py b/tests/test_loss/test_lploss.py
index 3743970df..3ddd503ee 100644
--- a/tests/test_loss/test_lploss.py
+++ b/tests/test_loss/test_lploss.py
@@ -1,7 +1,7 @@
 import torch
 import pytest
 
-from pina.loss import *
+from pina.loss.loss_interface import *
 
 input = torch.tensor([[3.], [1.], [-8.]])
 target = torch.tensor([[6.], [4.], [2.]])
diff --git a/tests/test_loss/test_powerloss.py b/tests/test_loss/test_powerloss.py
index 7ea26755d..fd5fddee8 100644
--- a/tests/test_loss/test_powerloss.py
+++ b/tests/test_loss/test_powerloss.py
@@ -1,7 +1,7 @@
 import torch
 import pytest
 
-from pina.loss import PowerLoss
+from pina.loss.loss_interface import PowerLoss
 
 input = torch.tensor([[3.], [1.], [-8.]])
 target = torch.tensor([[6.], [4.], [2.]])
diff --git a/tests/test_solvers/test_causalpinn.py b/tests/test_solvers/test_causalpinn.py
index c5a94a4b0..7629d51b1 100644
--- a/tests/test_solvers/test_causalpinn.py
+++ b/tests/test_solvers/test_causalpinn.py
@@ -10,7 +10,7 @@
 from pina.model import FeedForward
 from pina.equation.equation import Equation
 from pina.equation.equation_factory import FixedValue
-from pina.loss import LpLoss
+from pina.loss.loss_interface import LpLoss
 
 
 
diff --git a/tests/test_solvers/test_competitive_pinn.py b/tests/test_solvers/test_competitive_pinn.py
index 97ee88166..9facf2a44 100644
--- a/tests/test_solvers/test_competitive_pinn.py
+++ b/tests/test_solvers/test_competitive_pinn.py
@@ -10,7 +10,7 @@
 from pina.model import FeedForward
 from pina.equation.equation import Equation
 from pina.equation.equation_factory import FixedValue
-from pina.loss import LpLoss
+from pina.loss.loss_interface import LpLoss
 
 
 def laplace_equation(input_, output_):
diff --git a/tests/test_solvers/test_gpinn.py b/tests/test_solvers/test_gpinn.py
index 4a6f9e3c0..47e38e458 100644
--- a/tests/test_solvers/test_gpinn.py
+++ b/tests/test_solvers/test_gpinn.py
@@ -9,7 +9,7 @@
 from pina.model import FeedForward
 from pina.equation.equation import Equation
 from pina.equation.equation_factory import FixedValue
-from pina.loss import LpLoss
+from pina.loss.loss_interface import LpLoss
 
 
 def laplace_equation(input_, output_):
diff --git a/tests/test_solvers/test_pinn.py b/tests/test_solvers/test_pinn.py
index e2a8972d5..105bd886c 100644
--- a/tests/test_solvers/test_pinn.py
+++ b/tests/test_solvers/test_pinn.py
@@ -9,7 +9,7 @@
 from pina.model import FeedForward
 from pina.equation.equation import Equation
 from pina.equation.equation_factory import FixedValue
-from pina.loss import LpLoss
+from pina.loss.loss_interface import LpLoss
 
 
 def laplace_equation(input_, output_):
diff --git a/tests/test_solvers/test_rom_solver.py b/tests/test_solvers/test_rom_solver.py
index a16ffcaae..2a54d6504 100644
--- a/tests/test_solvers/test_rom_solver.py
+++ b/tests/test_solvers/test_rom_solver.py
@@ -6,7 +6,7 @@
 from pina.solvers import ReducedOrderModelSolver
 from pina.trainer import Trainer
 from pina.model import FeedForward
-from pina.loss import LpLoss
+from pina.loss.loss_interface import LpLoss
 
 
 class NeuralOperatorProblem(AbstractProblem):
diff --git a/tests/test_solvers/test_sapinn.py b/tests/test_solvers/test_sapinn.py
index a29e045e2..5e1eee6f6 100644
--- a/tests/test_solvers/test_sapinn.py
+++ b/tests/test_solvers/test_sapinn.py
@@ -10,7 +10,7 @@
 from pina.model import FeedForward
 from pina.equation.equation import Equation
 from pina.equation.equation_factory import FixedValue
-from pina.loss import LpLoss
+from pina.loss.loss_interface import LpLoss
 
 
 def laplace_equation(input_, output_):
diff --git a/tests/test_solvers/test_supervised_solver.py b/tests/test_solvers/test_supervised_solver.py
index 1b2b7d4b7..1bb812027 100644
--- a/tests/test_solvers/test_supervised_solver.py
+++ b/tests/test_solvers/test_supervised_solver.py
@@ -5,14 +5,17 @@
 from pina.solvers import SupervisedSolver
 from pina.trainer import Trainer
 from pina.model import FeedForward
-from pina.loss import LpLoss
+from pina.loss.loss_interface import LpLoss
 
 
 class NeuralOperatorProblem(AbstractProblem):
     input_variables = ['u_0', 'u_1']
     output_variables = ['u']
     domains = {
-        'pts': LabelTensor(torch.rand(100, 2), labels={1: {'name': 'space', 'dof': ['u_0', 'u_1']}})
+        'pts': LabelTensor(
+            torch.rand(100, 2),
+            labels={1: {'name': 'space', 'dof': ['u_0', 'u_1']}}
+        )
     }
     conditions = {
         'data' : Condition(
@@ -56,9 +59,51 @@ def test_constructor():
 #     SupervisedSolver(problem=problem, model=model_extra_feats, extra_features=extra_feats)
 
 
+class AutoSolver(SupervisedSolver):
+
+    def forward(self, input):
+        from pina.graph import Graph
+        print(Graph)
+        print(input)
+        if not isinstance(input, Graph):
+            input = Graph.build('radius', nodes_coordinates=input, nodes_data=torch.rand(input.shape), radius=0.2)
+        print(input)
+        print(input.data.edge_index)
+        print(input.data)
+        g = self.model[0](input.data, edge_index=input.data.edge_index)
+        g.labels = {1: {'name': 'output', 'dof': ['u']}}
+        return g
+        du_dt_new = LabelTensor(self.model[0](graph).reshape(-1,1), labels = ['du'])
+
+        return du_dt_new
+
+class GraphModel(torch.nn.Module):
+    def __init__(self, in_channels, out_channels):
+        from torch_geometric.nn import GCNConv, NNConv
+        super().__init__()
+        self.conv1 = GCNConv(in_channels, 16)
+        self.conv2 = GCNConv(16, out_channels)
+
+    def forward(self, data, edge_index):
+        print(data)
+        x = data.x
+        print(x)
+        x = self.conv1(x, edge_index)
+        x = x.relu()
+        x = self.conv2(x, edge_index)
+        return x
+
+def test_graph():
+
+    solver = AutoSolver(problem = problem, model=GraphModel(2, 1), loss=LpLoss())
+    trainer = Trainer(solver=solver, max_epochs=30, accelerator='cpu', batch_size=20)
+    trainer.train()
+    assert False
+
+
 def test_train_cpu():
     solver = SupervisedSolver(problem = problem, model=model, loss=LpLoss())
-    trainer = Trainer(solver=solver, max_epochs=3, accelerator='cpu', batch_size=20)
+    trainer = Trainer(solver=solver, max_epochs=300, accelerator='cpu', batch_size=20)
     trainer.train()
 
 
diff --git a/tutorials/tutorial10/tutorial.py b/tutorials/tutorial10/tutorial.py
index 637dd0560..9d37e1891 100644
--- a/tutorials/tutorial10/tutorial.py
+++ b/tutorials/tutorial10/tutorial.py
@@ -242,7 +242,7 @@ class NeuralOperatorProblem(AbstractProblem):
 # In[8]:
 
 
-from pina.loss import PowerLoss
+from pina.loss.loss_interface import PowerLoss
 
 error_metric = PowerLoss(p=2)                                                   # we use the MSE loss
 
diff --git a/tutorials/tutorial5/tutorial.py b/tutorials/tutorial5/tutorial.py
index 9386bc16f..4e5fa13e0 100644
--- a/tutorials/tutorial5/tutorial.py
+++ b/tutorials/tutorial5/tutorial.py
@@ -116,7 +116,7 @@ class NeuralOperatorSolver(AbstractProblem):
 # In[19]:
 
 
-from pina.loss import LpLoss
+from pina.loss.loss_interface import LpLoss
 
 # make the metric
 metric_err = LpLoss(relative=True)

From 0f1aa1576c8262bb926ede11ff81e39d514289c2 Mon Sep 17 00:00:00 2001
From: FilippoOlivo <115222346+FilippoOlivo@users.noreply.github.com>
Date: Thu, 12 Sep 2024 18:12:59 +0200
Subject: [PATCH 09/55] Fix bugs in 0.2 (#344)

* Fix some bugs
---
 examples/problems/stokes.py                  | 18 +++--
 examples/run_stokes.py                       |  4 +-
 pina/condition/condition.py                  |  6 +-
 pina/condition/condition_interface.py        |  5 +-
 pina/condition/domain_equation_condition.py  | 12 +++-
 pina/condition/domain_output_condition.py    |  1 +
 pina/domain/cartesian.py                     |  1 +
 pina/label_tensor.py                         | 37 +++++++++-
 pina/problem/abstract_problem.py             | 74 +++++++++++---------
 pina/solvers/supervised.py                   |  2 -
 tests/test_solvers/test_supervised_solver.py |  7 +-
 11 files changed, 112 insertions(+), 55 deletions(-)

diff --git a/examples/problems/stokes.py b/examples/problems/stokes.py
index c7b13873c..5f14f4e8f 100644
--- a/examples/problems/stokes.py
+++ b/examples/problems/stokes.py
@@ -49,11 +49,19 @@ def wall(input_, output_):
         value = 0.0
         return output_.extract(['ux', 'uy']) - value
 
+    domains = {
+        'gamma_top': CartesianDomain({'x': [-2, 2], 'y':  1}),
+        'gamma_bot': CartesianDomain({'x': [-2, 2], 'y': -1}),
+        'gamma_out': CartesianDomain({'x':  2, 'y': [-1, 1]}),
+        'gamma_in':  CartesianDomain({'x': -2, 'y': [-1, 1]}),
+        'D': CartesianDomain({'x': [-2, 2], 'y': [-1, 1]})
+    }
+
     # problem condition statement
     conditions = {
-        'gamma_top': Condition(location=CartesianDomain({'x': [-2, 2], 'y':  1}), equation=Equation(wall)),
-        'gamma_bot': Condition(location=CartesianDomain({'x': [-2, 2], 'y': -1}), equation=Equation(wall)),
-        'gamma_out': Condition(location=CartesianDomain({'x':  2, 'y': [-1, 1]}), equation=Equation(outlet)),
-        'gamma_in':  Condition(location=CartesianDomain({'x': -2, 'y': [-1, 1]}), equation=Equation(inlet)),
-        'D': Condition(location=CartesianDomain({'x': [-2, 2], 'y': [-1, 1]}), equation=SystemEquation([momentum, continuity]))
+        'gamma_top': Condition(domain='gamma_top', equation=Equation(wall)),
+        'gamma_bot': Condition(domain='gamma_bot', equation=Equation(wall)),
+        'gamma_out': Condition(domain='gamma_out', equation=Equation(outlet)),
+        'gamma_in':  Condition(domain='gamma_in', equation=Equation(inlet)),
+        'D': Condition(domain='D', equation=SystemEquation([momentum, continuity]))
     }
diff --git a/examples/run_stokes.py b/examples/run_stokes.py
index 54b2aecc3..04f652bd3 100644
--- a/examples/run_stokes.py
+++ b/examples/run_stokes.py
@@ -17,8 +17,8 @@
 
     # create problem and discretise domain
     stokes_problem = Stokes()
-    stokes_problem.discretise_domain(n=1000, locations=['gamma_top', 'gamma_bot', 'gamma_in', 'gamma_out'])
-    stokes_problem.discretise_domain(n=2000, locations=['D'])
+    stokes_problem.discretise_domain(n=1000, domains=['gamma_top', 'gamma_bot', 'gamma_in', 'gamma_out'])
+    stokes_problem.discretise_domain(n=2000, domains=['D'])
 
     # make the model
     model = FeedForward(
diff --git a/pina/condition/condition.py b/pina/condition/condition.py
index eec523c92..d815838bb 100644
--- a/pina/condition/condition.py
+++ b/pina/condition/condition.py
@@ -84,14 +84,15 @@ def __new__(cls, *args, **kwargs):
             return DomainEquationCondition(**kwargs)
         else:
             raise ValueError(f"Invalid keyword arguments {kwargs.keys()}.")
-        
+        # TODO: remove, not used anymore
+        '''
         if (
             sorted(kwargs.keys()) != sorted(["input_points", "output_points"])
             and sorted(kwargs.keys()) != sorted(["location", "equation"])
             and sorted(kwargs.keys()) != sorted(["input_points", "equation"])
         ):
             raise ValueError(f"Invalid keyword arguments {kwargs.keys()}.")
-
+        # TODO: remove, not used anymore
         if not self._dictvalue_isinstance(kwargs, "input_points", LabelTensor):
             raise TypeError("`input_points` must be a torch.Tensor.")
         if not self._dictvalue_isinstance(kwargs, "output_points", LabelTensor):
@@ -103,3 +104,4 @@ def __new__(cls, *args, **kwargs):
 
         for key, value in kwargs.items():
             setattr(self, key, value)
+        '''
\ No newline at end of file
diff --git a/pina/condition/condition_interface.py b/pina/condition/condition_interface.py
index f6b51bf96..0626a6d83 100644
--- a/pina/condition/condition_interface.py
+++ b/pina/condition/condition_interface.py
@@ -15,4 +15,7 @@ def residual(self, model):
         :param model: The model to evaluate the condition.
         :return: The residual of the condition.
         """
-        pass
\ No newline at end of file
+        pass
+
+    def set_problem(self, problem):
+        self._problem = problem
diff --git a/pina/condition/domain_equation_condition.py b/pina/condition/domain_equation_condition.py
index 8f45f8f4a..15df3f85f 100644
--- a/pina/condition/domain_equation_condition.py
+++ b/pina/condition/domain_equation_condition.py
@@ -15,6 +15,12 @@ def __init__(self, domain, equation):
         self.domain = domain
         self.equation = equation
 
+    def residual(self, model):
+        """
+        Compute the residual of the condition.
+        """
+        self.batch_residual(model, self.domain, self.equation)
+
     @staticmethod
     def batch_residual(model, input_pts, equation):
         """
@@ -22,7 +28,7 @@ def batch_residual(model, input_pts, equation):
         output points are provided as arguments.
 
         :param torch.nn.Module model: The model to evaluate the condition.
-        :param torch.Tensor input_points: The input points.
-        :param torch.Tensor output_points: The output points.
+        :param torch.Tensor input_pts: The input points.
+        :param torch.Tensor equation: The output points.
         """
-        return equation.residual(model(input_pts))
\ No newline at end of file
+        return equation.residual(input_pts, model(input_pts))
\ No newline at end of file
diff --git a/pina/condition/domain_output_condition.py b/pina/condition/domain_output_condition.py
index 49a0cb6ed..f847720b5 100644
--- a/pina/condition/domain_output_condition.py
+++ b/pina/condition/domain_output_condition.py
@@ -40,4 +40,5 @@ def batch_residual(model, input_points, output_points):
         :param torch.Tensor input_points: The input points.
         :param torch.Tensor output_points: The output points.
         """
+
         return output_points - model(input_points)
\ No newline at end of file
diff --git a/pina/domain/cartesian.py b/pina/domain/cartesian.py
index 9ab179bde..f1865e966 100644
--- a/pina/domain/cartesian.py
+++ b/pina/domain/cartesian.py
@@ -1,4 +1,5 @@
 import torch
+import torch
 
 from .domain_interface import DomainInterface
 from ..label_tensor import LabelTensor
diff --git a/pina/label_tensor.py b/pina/label_tensor.py
index 0a5515365..3167c0e61 100644
--- a/pina/label_tensor.py
+++ b/pina/label_tensor.py
@@ -5,7 +5,6 @@
 from torch import Tensor
 
 
-
 # class LabelTensor(torch.Tensor):
 #     """Torch tensor with a label for any column."""
 
@@ -307,13 +306,13 @@
 #             s = "no labels\n"
 #         s += super().__str__()
 #         return s
-
 def issubset(a, b):
     """
     Check if a is a subset of b.
     """
     return set(a).issubset(set(b))
 
+
 class LabelTensor(torch.Tensor):
     """Torch tensor with a label for any column."""
 
@@ -403,6 +402,10 @@ def extract(self, label_to_extract):
         return LabelTensor(new_tensor, label_to_extract)
         
     def __str__(self):
+        """
+        returns a string with the representation of the class
+        """
+
         s = ''
         for key, value in self.labels.items():
             s += f"{key}: {value}\n"
@@ -431,4 +434,32 @@ def requires_grad_(self, mode=True):
 
     @property
     def dtype(self):
-        return super().dtype
\ No newline at end of file
+        return super().dtype
+
+
+    def to(self, *args, **kwargs):
+        """
+        Performs Tensor dtype and/or device conversion. For more details, see
+        :meth:`torch.Tensor.to`.
+        """
+        tmp = super().to(*args, **kwargs)
+        new = self.__class__.clone(self)
+        new.data = tmp.data
+        return new
+
+
+    def clone(self, *args, **kwargs):
+        """
+        Clone the LabelTensor. For more details, see
+        :meth:`torch.Tensor.clone`.
+
+        :return: A copy of the tensor.
+        :rtype: LabelTensor
+        """
+        # # used before merging
+        # try:
+        #     out = LabelTensor(super().clone(*args, **kwargs), self.labels)
+        # except:
+        #     out = super().clone(*args, **kwargs)
+        out = LabelTensor(super().clone(*args, **kwargs), self.labels)
+        return out
\ No newline at end of file
diff --git a/pina/problem/abstract_problem.py b/pina/problem/abstract_problem.py
index 01bb1205c..200092e1b 100644
--- a/pina/problem/abstract_problem.py
+++ b/pina/problem/abstract_problem.py
@@ -20,7 +20,6 @@ class AbstractProblem(metaclass=ABCMeta):
 
     def __init__(self):
 
-
         self._discretized_domains = {}
 
         for name, domain in self.domains.items():
@@ -28,18 +27,19 @@ def __init__(self):
                 self._discretized_domains[name] = domain
 
         for condition_name in self.conditions:
-            self.conditions[condition_name]._problem = self
+            self.conditions[condition_name].set_problem(self)
+
         # # variable storing all points
-        # self.input_pts = {}
+        self.input_pts = {}
 
         # # varible to check if sampling is done. If no location
         # # element is presented in Condition this variable is set to true
         # self._have_sampled_points = {}
-        # for condition_name in self.conditions:
-        #     self._have_sampled_points[condition_name] = False
+        for condition_name in self.conditions:
+            self._discretized_domains[condition_name] = False
 
         # # put in self.input_pts all the points that we don't need to sample
-        # self._span_condition_points()
+        self._span_condition_points()
 
     def __deepcopy__(self, memo):
         """
@@ -125,7 +125,7 @@ def _span_condition_points(self):
             if hasattr(condition, "input_points"):
                 samples = condition.input_points
                 self.input_pts[condition_name] = samples
-                self._have_sampled_points[condition_name] = True
+                self._discretized_domains[condition_name] = True
             if hasattr(self, "unknown_parameter_domain"):
                 # initialize the unknown parameters of the inverse problem given
                 # the domain the user gives
@@ -141,7 +141,7 @@ def _span_condition_points(self):
                     )
 
     def discretise_domain(
-        self, n, mode="random", variables="all", locations="all"
+        self, n, mode="random", variables="all", domains="all"
     ):
         """
         Generate a set of points to span the `Location` of all the conditions of
@@ -192,31 +192,37 @@ def discretise_domain(
                 f"should be in {self.input_variables}.",
             )
 
-        # check consistency location
-        locations_to_sample = [
-            condition
-            for condition in self.conditions
-            if hasattr(self.conditions[condition], "location")
-        ]
-        if locations == "all":
-            # only locations that can be sampled
-            locations = locations_to_sample
+        # # check consistency location # TODO: check if this is needed (from 0.1)
+        # locations_to_sample = [
+        #     condition
+        #     for condition in self.conditions
+        #     if hasattr(self.conditions[condition], "location")
+        # ]
+        # if locations == "all":
+        #     # only locations that can be sampled
+        #     locations = locations_to_sample
+        # else:
+        #     check_consistency(locations, str)
+
+        # if sorted(locations) != sorted(locations_to_sample):
+        if domains == "all":
+            domains = [condition for condition in self.conditions]
         else:
-            check_consistency(locations, str)
-
-        if sorted(locations) != sorted(locations_to_sample):
+            check_consistency(domains, str)
+        print(domains)
+        if sorted(domains) != sorted(self.conditions):
             TypeError(
                 f"Wrong locations for sampling. Location ",
                 f"should be in {locations_to_sample}.",
             )
 
         # sampling
-        for location in locations:
-            condition = self.conditions[location]
+        for d in domains:
+            condition = self.conditions[d]
 
             # we try to check if we have already sampled
             try:
-                already_sampled = [self.input_pts[location]]
+                already_sampled = [self.input_pts[d]]
             # if we have not sampled, a key error is thrown
             except KeyError:
                 already_sampled = []
@@ -225,25 +231,27 @@ def discretise_domain(
             # but we want to sample again we set already_sampled
             # to an empty list since we need to sample again, and
             # self._have_sampled_points to False.
-            if self._have_sampled_points[location]:
+            if self._discretized_domains[d]:
                 already_sampled = []
-                self._have_sampled_points[location] = False
-
+                self._discretized_domains[d] = False
+            print(condition.domain)
+            print(d)
             # build samples
             samples = [
-                condition.location.sample(n=n, mode=mode, variables=variables)
+                self.domains[d].sample(n=n, mode=mode, variables=variables)
             ] + already_sampled
             pts = merge_tensors(samples)
-            self.input_pts[location] = pts
+            self.input_pts[d] = pts
 
             # the condition is sampled if input_pts contains all labels
-            if sorted(self.input_pts[location].labels) == sorted(
+            if sorted(self.input_pts[d].labels) == sorted(
                 self.input_variables
             ):
-                self._have_sampled_points[location] = True
-                self.input_pts[location] = self.input_pts[location].extract(
-                    sorted(self.input_variables)
-                )
+                # self._have_sampled_points[location] = True
+                # self.input_pts[location] = self.input_pts[location].extract(
+                #     sorted(self.input_variables)
+                # )
+                self._have_sampled_points[d] = True
 
     def add_points(self, new_points):
         """
diff --git a/pina/solvers/supervised.py b/pina/solvers/supervised.py
index 4f3d497fe..dd7cab761 100644
--- a/pina/solvers/supervised.py
+++ b/pina/solvers/supervised.py
@@ -134,8 +134,6 @@ def training_step(self, batch, batch_idx):
             condition = self.problem.conditions[condition_name]
             pts = batch.input
             out = batch.output
-            print(out)
-            print(pts)
 
             if condition_name not in self.problem.conditions:
                 raise RuntimeError("Something wrong happened.")
diff --git a/tests/test_solvers/test_supervised_solver.py b/tests/test_solvers/test_supervised_solver.py
index 1bb812027..cc8c5631b 100644
--- a/tests/test_solvers/test_supervised_solver.py
+++ b/tests/test_solvers/test_supervised_solver.py
@@ -5,7 +5,7 @@
 from pina.solvers import SupervisedSolver
 from pina.trainer import Trainer
 from pina.model import FeedForward
-from pina.loss.loss_interface import LpLoss
+from pina.loss import LpLoss
 
 
 class NeuralOperatorProblem(AbstractProblem):
@@ -94,11 +94,9 @@ def forward(self, data, edge_index):
         return x
 
 def test_graph():
-
     solver = AutoSolver(problem = problem, model=GraphModel(2, 1), loss=LpLoss())
     trainer = Trainer(solver=solver, max_epochs=30, accelerator='cpu', batch_size=20)
     trainer.train()
-    assert False
 
 
 def test_train_cpu():
@@ -107,6 +105,7 @@ def test_train_cpu():
     trainer.train()
 
 
+
 # def test_train_restore():
 #     tmpdir = "tests/tmp_restore"
 #     solver = SupervisedSolver(problem=problem,
@@ -153,4 +152,4 @@ def test_train_cpu():
 #                 model=model_extra_feats,
 #                 extra_features=extra_feats)
 #     trainer = Trainer(solver=pinn, max_epochs=5, accelerator='cpu')
-#     trainer.train()
\ No newline at end of file
+#     trainer.train()

From c6ea6dffe83ac56477738327ad4cad45a9d5c149 Mon Sep 17 00:00:00 2001
From: FilippoOlivo <115222346+FilippoOlivo@users.noreply.github.com>
Date: Sat, 21 Sep 2024 18:55:57 +0200
Subject: [PATCH 10/55] Fix SupervisedSolver GPU bug and implement GraphSolver
 (#346)

* Fix some bugs
* Solve bug with GPU and model_summary parameters in SupervisedSolver class
* Implement GraphSolver class
* Fix Tutorial 5
---
 pina/data/pina_dataloader.py                 |   4 +-
 pina/label_tensor.py                         |   4 +-
 pina/model/fno.py                            |   3 +
 pina/solvers/__init__.py                     |   1 +
 pina/solvers/graph.py                        |  34 +++
 pina/solvers/supervised.py                   |  23 +-
 setup.py                                     |   2 +-
 tests/test_solvers/test_supervised_solver.py |  34 ++-
 tutorials/tutorial5/tutorial.ipynb           | 303 ++++++++++++-------
 tutorials/tutorial5/tutorial.py              |  64 ++--
 10 files changed, 321 insertions(+), 151 deletions(-)
 create mode 100644 pina/solvers/graph.py

diff --git a/pina/data/pina_dataloader.py b/pina/data/pina_dataloader.py
index 1b71a46a3..2c8967c50 100644
--- a/pina/data/pina_dataloader.py
+++ b/pina/data/pina_dataloader.py
@@ -4,6 +4,7 @@
 from .data_dataset import DataPointDataset
 from .pina_batch import Batch
 
+
 class SamplePointLoader:
     """
     This class is used to create a dataloader to use during the training.
@@ -95,7 +96,7 @@ def _prepare_data_dataset(self, dataset, batch_size, shuffle):
         self.batch_output_pts = torch.tensor_split(
             dataset.output_pts, batch_num
         )
-        print(input_labels)
+        #print(input_labels)
         for i in range(len(self.batch_input_pts)):
             self.batch_input_pts[i].labels = input_labels
             self.batch_output_pts[i].labels = output_labels
@@ -161,7 +162,6 @@ def _prepare_batches(self):
                     self.batch_input_pts,
                     self.batch_output_pts,
                     self.batch_data_conditions)
-            print(batch.input.labels)
                 
             self.batches.append(batch)
 
diff --git a/pina/label_tensor.py b/pina/label_tensor.py
index 3167c0e61..ab6045e3f 100644
--- a/pina/label_tensor.py
+++ b/pina/label_tensor.py
@@ -425,7 +425,7 @@ def stack(tensors):
         
         raise NotImplementedError
         labels = [tensor.labels for tensor in tensors]
-        print(labels)
+
 
     def requires_grad_(self, mode=True):
         lt = super().requires_grad_(mode)
@@ -436,7 +436,6 @@ def requires_grad_(self, mode=True):
     def dtype(self):
         return super().dtype
 
-
     def to(self, *args, **kwargs):
         """
         Performs Tensor dtype and/or device conversion. For more details, see
@@ -447,7 +446,6 @@ def to(self, *args, **kwargs):
         new.data = tmp.data
         return new
 
-
     def clone(self, *args, **kwargs):
         """
         Clone the LabelTensor. For more details, see
diff --git a/pina/model/fno.py b/pina/model/fno.py
index 92ca18361..bcf018f6b 100644
--- a/pina/model/fno.py
+++ b/pina/model/fno.py
@@ -269,4 +269,7 @@ def forward(self, x):
         :return: The output tensor obtained from FNO.
         :rtype: torch.Tensor
         """
+
+        if isinstance(x, LabelTensor):
+            x = x.as_subclass(torch.Tensor)
         return super().forward(x)
diff --git a/pina/solvers/__init__.py b/pina/solvers/__init__.py
index 7bb988d56..59a182691 100644
--- a/pina/solvers/__init__.py
+++ b/pina/solvers/__init__.py
@@ -17,3 +17,4 @@
 from .supervised import SupervisedSolver
 from .rom import ReducedOrderModelSolver
 from .garom import GAROM
+from .graph import GraphSupervisedSolver
diff --git a/pina/solvers/graph.py b/pina/solvers/graph.py
new file mode 100644
index 000000000..9af04e76f
--- /dev/null
+++ b/pina/solvers/graph.py
@@ -0,0 +1,34 @@
+from .supervised import SupervisedSolver
+from ..graph import Graph
+
+
+class GraphSupervisedSolver(SupervisedSolver):
+
+    def __init__(
+        self,
+        problem,
+        model,
+        nodes_coordinates,
+        nodes_data,
+        loss=None,
+        optimizer=None,
+        scheduler=None):
+        super().__init__(problem, model, loss, optimizer, scheduler)
+        if isinstance(nodes_coordinates, str):
+            self._nodes_coordinates = [nodes_coordinates]
+        else:
+            self._nodes_coordinates = nodes_coordinates
+        if isinstance(nodes_data, str):
+            self._nodes_data = nodes_data
+        else:
+            self._nodes_data = nodes_data
+
+    def forward(self, input):
+        input_coords = input.extract(self._nodes_coordinates)
+        input_data = input.extract(self._nodes_data)
+
+        if not isinstance(input, Graph):
+            input = Graph.build('radius', nodes_coordinates=input_coords, nodes_data=input_data, radius=0.2)
+        g = self.model(input.data, edge_index=input.data.edge_index)
+        g.labels = {1: {'name': 'output', 'dof': ['u']}}
+        return g
diff --git a/pina/solvers/supervised.py b/pina/solvers/supervised.py
index dd7cab761..c44d5a1e2 100644
--- a/pina/solvers/supervised.py
+++ b/pina/solvers/supervised.py
@@ -82,7 +82,10 @@ def __init__(
 
         # check consistency
         check_consistency(loss, (LossInterface, _Loss), subclass=False)
-        self.loss = loss
+        self._loss = loss
+        self._model = self._pina_model[0]
+        self._optimizer = self._pina_optimizer[0]
+        self._scheduler = self._pina_scheduler[0]
 
     def forward(self, x):
         """Forward pass implementation for the solver.
@@ -92,7 +95,7 @@ def forward(self, x):
         :rtype: torch.Tensor
         """
 
-        output = self._pina_model[0](x)
+        output = self._model(x)
 
         output.labels = {
             1: {
@@ -108,11 +111,11 @@ def configure_optimizers(self):
         :return: The optimizers and the schedulers
         :rtype: tuple(list, list)
         """
-        self._pina_optimizer[0].hook(self._pina_model[0].parameters())
-        self._pina_scheduler[0].hook(self._pina_optimizer[0])
+        self._optimizer.hook(self._model.parameters())
+        self._scheduler.hook(self._optimizer)
         return (
-            [self._pina_optimizer[0].optimizer_instance], 
-            [self._pina_scheduler[0].scheduler_instance]
+            [self._optimizer.optimizer_instance],
+            [self._scheduler.scheduler_instance]
         )
 
     def training_step(self, batch, batch_idx):
@@ -170,28 +173,28 @@ def loss_data(self, input_pts, output_pts):
         :return: The residual loss averaged on the input coordinates
         :rtype: torch.Tensor
         """
-        return self.loss(self.forward(input_pts), output_pts)
+        return self._loss(self.forward(input_pts), output_pts)
 
     @property
     def scheduler(self):
         """
         Scheduler for training.
         """
-        return self._pina_scheduler
+        return self._scheduler
     
     @property
     def optimizer(self):
         """
         Optimizer for training.
         """
-        return self._pina_optimizer
+        return self._optimizer
 
     @property
     def model(self):
         """
         Neural network for training.
         """
-        return self._pina_model
+        return self._model
 
     @property
     def loss(self):
diff --git a/setup.py b/setup.py
index 8c7b9ac44..5a2ebc85e 100644
--- a/setup.py
+++ b/setup.py
@@ -15,7 +15,7 @@
 KEYWORDS = 'machine-learning deep-learning modeling pytorch ode neural-networks differential-equations pde hacktoberfest pinn physics-informed physics-informed-neural-networks neural-operators equation-learning lightining'
 
 REQUIRED = [
-    'numpy<2.0', 'matplotlib', 'torch', 'lightning', 'pytorch_lightning'
+    'numpy', 'matplotlib', 'torch', 'lightning', 'pytorch_lightning', 'torch_geometric', 'torch-cluster'
 ]
 
 EXTRAS = {
diff --git a/tests/test_solvers/test_supervised_solver.py b/tests/test_solvers/test_supervised_solver.py
index cc8c5631b..912480bb8 100644
--- a/tests/test_solvers/test_supervised_solver.py
+++ b/tests/test_solvers/test_supervised_solver.py
@@ -6,7 +6,7 @@
 from pina.trainer import Trainer
 from pina.model import FeedForward
 from pina.loss import LpLoss
-
+from pina.solvers import GraphSupervisedSolver
 
 class NeuralOperatorProblem(AbstractProblem):
     input_variables = ['u_0', 'u_1']
@@ -27,6 +27,25 @@ class NeuralOperatorProblem(AbstractProblem):
         )
     }
 
+class NeuralOperatorProblemGraph(AbstractProblem):
+    input_variables = ['x', 'y', 'u_0', 'u_1']
+    output_variables = ['u']
+    domains = {
+        'pts': LabelTensor(
+            torch.rand(100, 4),
+            labels={1: {'name': 'space', 'dof': ['x', 'y', 'u_0', 'u_1']}}
+        )
+    }
+    conditions = {
+        'data' : Condition(
+            domain='pts',
+            output_points=LabelTensor(
+                torch.rand(100, 1),
+                labels={1: {'name': 'output', 'dof': ['u']}}
+            )
+        )
+    }
+
 class myFeature(torch.nn.Module):
     """
     Feature: sin(x)
@@ -42,6 +61,7 @@ def forward(self, x):
 
 
 problem = NeuralOperatorProblem()
+problem_graph = NeuralOperatorProblemGraph()
 # make the problem + extra feats
 extra_feats = [myFeature()]
 model = FeedForward(len(problem.input_variables),
@@ -58,7 +78,7 @@ def test_constructor():
 # def test_constructor_extra_feats():
 #     SupervisedSolver(problem=problem, model=model_extra_feats, extra_features=extra_feats)
 
-
+'''
 class AutoSolver(SupervisedSolver):
 
     def forward(self, input):
@@ -70,12 +90,13 @@ def forward(self, input):
         print(input)
         print(input.data.edge_index)
         print(input.data)
-        g = self.model[0](input.data, edge_index=input.data.edge_index)
+        g = self._model(input.data, edge_index=input.data.edge_index)
         g.labels = {1: {'name': 'output', 'dof': ['u']}}
         return g
-        du_dt_new = LabelTensor(self.model[0](graph).reshape(-1,1), labels = ['du'])
+        du_dt_new = LabelTensor(self.model(graph).reshape(-1,1), labels = ['du'])
 
         return du_dt_new
+'''
 
 class GraphModel(torch.nn.Module):
     def __init__(self, in_channels, out_channels):
@@ -94,7 +115,8 @@ def forward(self, data, edge_index):
         return x
 
 def test_graph():
-    solver = AutoSolver(problem = problem, model=GraphModel(2, 1), loss=LpLoss())
+    solver = GraphSupervisedSolver(problem=problem_graph, model=GraphModel(2, 1), loss=LpLoss(),
+                                   nodes_coordinates=['x', 'y'], nodes_data=['u_0', 'u_1'])
     trainer = Trainer(solver=solver, max_epochs=30, accelerator='cpu', batch_size=20)
     trainer.train()
 
@@ -105,7 +127,6 @@ def test_train_cpu():
     trainer.train()
 
 
-
 # def test_train_restore():
 #     tmpdir = "tests/tmp_restore"
 #     solver = SupervisedSolver(problem=problem,
@@ -153,3 +174,4 @@ def test_train_cpu():
 #                 extra_features=extra_feats)
 #     trainer = Trainer(solver=pinn, max_epochs=5, accelerator='cpu')
 #     trainer.train()
+test_graph()
\ No newline at end of file
diff --git a/tutorials/tutorial5/tutorial.ipynb b/tutorials/tutorial5/tutorial.ipynb
index 288dbe87e..986392f23 100644
--- a/tutorials/tutorial5/tutorial.ipynb
+++ b/tutorials/tutorial5/tutorial.ipynb
@@ -21,10 +21,13 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 1,
    "id": "5f2744dc",
-   "metadata": {},
-   "outputs": [],
+   "metadata": {
+    "ExecuteTime": {
+     "end_time": "2024-09-19T13:35:28.837348Z",
+     "start_time": "2024-09-19T13:35:27.611334Z"
+    }
+   },
    "source": [
     "## routine needed to run the notebook on Google Colab\n",
     "try:\n",
@@ -69,22 +72,31 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 12,
    "id": "2ffb8a4c",
-   "metadata": {},
-   "outputs": [],
+   "metadata": {
+    "ExecuteTime": {
+     "end_time": "2024-09-19T13:35:28.989631Z",
+     "start_time": "2024-09-19T13:35:28.952744Z"
+    }
+   },
    "source": [
     "# download the dataset\n",
     "data = io.loadmat(\"Data_Darcy.mat\")\n",
     "\n",
     "# extract data (we use only 100 data for train)\n",
-    "k_train = LabelTensor(torch.tensor(data['k_train'], dtype=torch.float).unsqueeze(-1), ['u0'])\n",
-    "u_train = LabelTensor(torch.tensor(data['u_train'], dtype=torch.float).unsqueeze(-1), ['u'])\n",
-    "k_test =  LabelTensor(torch.tensor(data['k_test'], dtype=torch.float).unsqueeze(-1), ['u0'])\n",
-    "u_test= LabelTensor(torch.tensor(data['u_test'], dtype=torch.float).unsqueeze(-1), ['u'])\n",
+    "k_train = LabelTensor(torch.tensor(data['k_train'], dtype=torch.float).unsqueeze(-1), \n",
+    "                      labels={3:{'dof': ['u0'], 'name': 'k_train'}})\n",
+    "u_train = LabelTensor(torch.tensor(data['u_train'], dtype=torch.float).unsqueeze(-1),\n",
+    "                     labels={3:{'dof': ['u'], 'name': 'u_train'}})\n",
+    "k_test =  LabelTensor(torch.tensor(data['k_test'], dtype=torch.float).unsqueeze(-1),\n",
+    "                      labels={3:{'dof': ['u0'], 'name': 'k_test'}})\n",
+    "u_test= LabelTensor(torch.tensor(data['u_test'], dtype=torch.float).unsqueeze(-1),\n",
+    "                   labels={3:{'dof': ['u'], 'name': 'u_test'}})\n",
     "x = torch.tensor(data['x'], dtype=torch.float)[0]\n",
     "y = torch.tensor(data['y'], dtype=torch.float)[0]"
-   ]
+   ],
+   "outputs": [],
+   "execution_count": 2
   },
   {
    "cell_type": "markdown",
@@ -96,30 +108,61 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 13,
    "id": "c8501b6f",
-   "metadata": {},
+   "metadata": {
+    "ExecuteTime": {
+     "end_time": "2024-09-19T13:35:29.108381Z",
+     "start_time": "2024-09-19T13:35:29.031076Z"
+    }
+   },
+   "source": [
+    "plt.subplot(1, 2, 1)\n",
+    "plt.title('permeability')\n",
+    "plt.imshow(k_train.squeeze(-1)[0])\n",
+    "plt.subplot(1, 2, 2)\n",
+    "plt.title('field solution')\n",
+    "plt.imshow(u_train.squeeze(-1)[0])\n",
+    "plt.show()"
+   ],
    "outputs": [
     {
      "data": {
-      "image/png": "iVBORw0KGgoAAAANSUhEUgAAAiwAAAEjCAYAAAARyVqhAAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjcuMSwgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy/bCgiHAAAACXBIWXMAAA9hAAAPYQGoP6dpAAA73klEQVR4nO3de3RTZbo/8O9O2qaFtilI6QW5FbkI2KIonSpXQUv1IKDDpaPSMoIzHhid1YWjdQkUbz3iURmBA+oMFAdUdAbhzBEZsQocB1AB63Xk0NrSIm1pC73T5rLf3x/+Gg295H1pQ3bD97NWFmTnyc67k+ynT5K9n1cTQggQERERGZjJ1wMgIiIi8oQFCxERERkeCxYiIiIyPBYsREREZHgsWIiIiMjwWLAQERGR4bFgISIiIsNjwUJERESGx4KFiIiIDI8FC3VbWVlZ0DQNlZWVHmMHDRqE9PR01/V9+/ZB0zTs27fPtSw9PR2DBg3q+oESGcxnn32GG2+8ET179oSmacjLy3PtTxdDdt8pKiqCpmnIycm5qMe5GDk5OdA0DUVFRV263smTJ2Py5Mlduk7qGAsWonY0NjYiKyvLragh6u7sdjvmzJmDs2fP4sUXX8Rf/vIXDBw40NfDMqRvv/0WWVlZXV7s0MUJ8PUAiC6F48ePw2TquD5/9dVXoeu663pjYyNWrVoFAPwkRX6joKAAJ0+exKuvvopFixa5lj/++ON49NFHfTgy4/n222+xatUqTJ48udU3SO+//75vBnUZY8FCHjU2NqJHjx6+HkanWCwWjzGBgYGXYCREvnXmzBkAQEREhNvygIAABATwT4KsoKAgXw/hssOfhLq5lt+dv/vuO8ydOxfh4eG44oor8NBDD6GpqcktduvWrRg7dixCQkLQu3dvzJ8/HyUlJW4xkydPxujRo3H06FFMnDgRPXr0wGOPPeb67fk///M/sX79esTFxaFHjx649dZbUVJSAiEEnnzySVx55ZUICQnBzJkzcfbs2Vbjfe+99zBhwgT07NkTYWFhuP322/HNN9+4xXz55ZdIT09HXFwcgoODER0djV//+teoqqpq8zmorKz0uO0XHsPSlp//Dl9UVITIyEgAwKpVq6BpGjRNQ1ZWFjZv3gxN0/D555+3WsczzzwDs9mMH374ocPHIvKF9PR0TJo0CQAwZ84caJrm+vawvWNYZPJGW6qrq5Geng6r1YqIiAikpaWhurpaapx2ux2rVq3C0KFDERwcjCuuuALjx4/H3r173eI+/PBDVz6JiIjAzJkz8a9//cvj+lv25Qv9PE/k5ORgzpw5AIApU6a4ckDLT8RtHcNy5swZ3HfffYiKikJwcDASEhKwZcsWt5if59JXXnkFQ4YMgcViwQ033IDPPvtM6vm5XLGc9hNz587FoEGDkJ2djcOHD+Oll17CuXPn8NprrwEAnn76aSxfvhxz587FokWLUFFRgbVr12LixIn4/PPP3T5tVVVVISUlBfPnz8c999yDqKgo123btm2DzWbD7373O5w9exarV6/G3LlzcfPNN2Pfvn145JFHkJ+fj7Vr12LZsmXYtGmT675/+ctfkJaWhuTkZDz77LNobGzEhg0bMH78eHz++eeuYmHv3r34/vvvsXDhQkRHR+Obb77BK6+8gm+++QaHDx9ulVQ9bfvFiIyMxIYNG/DAAw9g9uzZuPPOOwEA8fHxGDx4MJYsWYJt27bh2muvdbvftm3bMHnyZPTr1++iH5vIW37zm9+gX79+eOaZZ/Dggw/ihhtucNu/L6SSN35OCIGZM2fi448/xm9/+1tcffXVeOedd5CWliY1zqysLGRnZ2PRokUYN24camtrceTIERw7dgy33HILAOCDDz5ASkoK4uLikJWVhfPnz2Pt2rW46aabcOzYsU4fQD9x4kQ8+OCDeOmll/DYY4/h6quvBgDXvxc6f/48Jk+ejPz8fCxduhSDBw/G22+/jfT0dFRXV+Ohhx5yi3/99ddRV1eH3/zmN9A0DatXr8add96J77//nt/2tkdQt7Zy5UoBQNxxxx1uy//93/9dABBffPGFKCoqEmazWTz99NNuMV999ZUICAhwWz5p0iQBQGzcuNEttrCwUAAQkZGRorq62rU8MzNTABAJCQnCbre7lqempoqgoCDR1NQkhBCirq5OREREiMWLF7utt6ysTFitVrfljY2NrbbzjTfeEADEgQMHlLa9xcCBA0VaWprr+kcffSQAiI8++si1LC0tTQwcONB1vaKiQgAQK1eubDWe1NRUERsbK5xOp2vZsWPHBACxefPmVvFERtHy3n/77bfdlrfsTy1U8saF+87OnTsFALF69WrXMofDISZMmCC1jyQkJIjbb7+9w5gxY8aIvn37iqqqKteyL774QphMJrFgwQLXss2bNwsAorCw0LWsvf36wjzx9ttvt8oTLSZNmiQmTZrkur5mzRoBQGzdutW1zGaziaSkJBEaGipqa2uFED/l0iuuuEKcPXvWFbtr1y4BQPz973/vcLsvZ/xJyE8sWbLE7frvfvc7AMDu3buxY8cO6LqOuXPnorKy0nWJjo7G0KFD8dFHH7nd12KxYOHChW0+zpw5c2C1Wl3XExMTAQD33HOP2+/fiYmJsNlsrp9G9u7di+rqaqSmprqNwWw2IzEx0W0MISEhrv83NTWhsrISv/jFLwAAx44dU9p2b1mwYAFOnz7tNu5t27YhJCQEd911l9cel+hSUc0bP7d7924EBATggQcecC0zm82ufdOTiIgIfPPNNzhx4kSbt5eWliIvLw/p6eno3bu3a3l8fDxuueUWr+777dm9ezeio6ORmprqWhYYGIgHH3wQ9fX12L9/v1v8vHnz0KtXL9f1CRMmAAC+//77SzPgbog/CfmJoUOHul0fMmQITCYTioqKYDKZIIRoFdPiwq8f+/Xr1+4BZQMGDHC73lK89O/fv83l586dAwBX4rn55pvbXG94eLjr/2fPnsWqVavw5ptvug4QbFFTU9Pqvh1tu7fccsstiImJwbZt2zB16lTouo433ngDM2fORFhYmNcel+hSOXHihFLe+LmTJ08iJiYGoaGhbsuHDx8u9dhPPPEEZs6ciWHDhmH06NGYPn067r33XsTHx7vW3976rr76avzjH/9AQ0MDevbsKfV4XeHkyZMYOnRoq7MRW35CahlziwtzaUvx0pIzqTUWLH7q58d56LoOTdPw3nvvwWw2t4q9MKn8/BuOC7V1/46WCyFcYwB+PI4lOjq6VdzPv52ZO3cuDh48iIcffhhjxoxBaGgodF3H9OnT3U47bs/FNr9SYTab8atf/Qqvvvoq/uu//gv//Oc/cfr0adxzzz1ef2yiS0E1b3SliRMnoqCgALt27cL777+PP/3pT3jxxRexceNGt1Oxu5rT6fTaui/kKWdSayxY/MSJEycwePBg1/X8/Hzouo5BgwbBbDZDCIHBgwdj2LBhPhnfkCFDAAB9+/bFtGnT2o07d+4ccnNzsWrVKqxYscK1vL2vhltua2/bO8NT4bNgwQI8//zz+Pvf/4733nsPkZGRSE5O7tRjEhnFkCFDLjpvDBw4ELm5uaivr3crbI4fPy69jt69e2PhwoVYuHAh6uvrMXHiRGRlZWHRokWuRndtre+7775Dnz59Ovx2pVevXq3OWLLZbCgtLXVbpvLhZ+DAgfjyyy+h67rbtyzfffed63bqHB7D4ifWr1/vdn3t2rUAgJSUFNx5550wm81YtWpVq+pdCNHu6cJdKTk5GeHh4XjmmWdgt9tb3V5RUQHgp08dF45zzZo17a67o23vjJbeM+2dihkfH4/4+Hj86U9/wt/+9jfMnz+ffSzIb3Qmb9x2221wOBzYsGGDa5nT6XTtm55cuO7Q0FBcddVVaG5uBgDExMRgzJgx2LJli9v++fXXX+P999/Hbbfd1uH6hwwZggMHDrgte+WVV1p9w9JS9Micjn3bbbehrKwM27dvdy1zOBxYu3YtQkNDXaeT08VjdvUThYWFuOOOOzB9+nQcOnQIW7duxa9+9SskJCQAAJ566ilkZmaiqKgIs2bNQlhYGAoLC/HOO+/g/vvvx7Jly7w6vvDwcGzYsAH33nsvrrvuOsyfPx+RkZEoLi7Gu+++i5tuugnr1q1DeHg4Jk6ciNWrV8Nut6Nfv354//33UVhYeNHbfrFCQkIwcuRIbN++HcOGDUPv3r0xevRojB492hWzYMEC13PHn4PInwwZMuSi88aMGTNw00034dFHH0VRURFGjhyJHTt2tHkMWltGjhyJyZMnY+zYsejduzeOHDmCv/71r1i6dKkr5rnnnkNKSgqSkpJw3333uU5rtlqtbfZY+blFixbht7/9Le666y7ccsst+OKLL/CPf/wDffr0cYsbM2YMzGYznn32WdTU1MBiseDmm29G3759W63z/vvvx8svv4z09HQcPXoUgwYNwl//+lf885//xJo1a3hsW1fw0dlJ1EVaTkX89ttvxS9/+UsRFhYmevXqJZYuXSrOnz/vFvu3v/1NjB8/XvTs2VP07NlTjBgxQixZskQcP37cFTNp0iQxatSoVo/Tcirec88957a8vVMkW04l/Oyzz1rFJycnC6vVKoKDg8WQIUNEenq6OHLkiCvm1KlTYvbs2SIiIkJYrVYxZ84ccfr06VanIqps+8Wc1iyEEAcPHhRjx44VQUFBbZ4KWVpaKsxmsxg2bFir54zIiGRPa24hkzfa2neqqqrEvffeK8LDw4XVahX33nuv+Pzzz6VOa37qqafEuHHjREREhAgJCREjRowQTz/9tLDZbG5xH3zwgbjppptESEiICA8PFzNmzBDffvutW0xbpzU7nU7xyCOPiD59+ogePXqI5ORkkZ+f3ypPCCHEq6++KuLi4oTZbHbLGRee1iyEEOXl5WLhwoWiT58+IigoSFxzzTWttrW9XCpE+6db0480IXiET3eWlZWFVatWoaKiotWnA/K+yspKxMTEYMWKFVi+fLmvh0NE5Ld4DAtRJ+Tk5MDpdOLee+/19VCIiPwaj2Ehuggffvghvv32Wzz99NOYNWtWp89IIiKijrFgIboITzzxBA4ePIibbrpJ+swHIiK6eDyGhYiIiAyPx7AQERGR4bFgISIiIsPzi2NYdF3H6dOnERYWdknmkSGi1oQQqKurQ2xsbKsJ4IyKuYPIt1Tyhl8ULKdPn241WzAR+UZJSQmuvPJKXw9DCnMHkTHI5A2vFSzr16/Hc889h7KyMiQkJGDt2rUYN25cu/Fvv/02li9fjqKiIgwdOhTPPvusx/kgWrS0PL5y5eMwBQd7jNc8T/jrpmexwqdFhQ9ptVe3nlOnPZbS9qdy76weZ+SPuw6sVztGO6RCfhthkn/ydLPap+GGaPm3evA5+TdIwHn5WE3x+Pag6mbpWD2w7Zlf2xKQX+o56GeclZVScQ7Y8TF2d6oF+aXMG8BPuePksUEID/W8n1c4G6TXDQCVTvnccVZvf5b0C51ztj+x34VqFNYLAOcc8rMwV9nlx1FpU5vdubJZPr7yfA/p2LN18rEAYKvx/DelRUC1fJ6xVMnnsOAqxbxbKT/rdEh5o3SsufSs0jgcZeWeYxTyhlcKlu3btyMjIwMbN25EYmIi1qxZg+TkZBw/frzNORgOHjyI1NRUZGdn49/+7d/w+uuvY9asWTh27JjbvC3tafkq1xQc7JWCxWzxTsFiCpH/I2O2eK9gMQfJ7wwqsQAQECC/jUoFS4BawWIOkn+rBwQqFCwOhYJFV3zuFIoyXeF5DjAFKY1D0yTfe6Il/uJ+WrnUeePnYw0PNSE8zPN+3qRQgKjGN+sKsU7519vmVEvz5x0Kf3Tt8nkpyKb2vgsMkI8PMFmkY81O+QIEAEw2+XhTk/xzZ7bI7yfKeTdQvmAJMMvHmhVzB2Ryh0Le8MoPzS+88AIWL16MhQsXYuTIkdi4cSN69OiBTZs2tRn/xz/+EdOnT8fDDz+Mq6++Gk8++SSuu+46rFu3zhvDIyIDYt4goo50ecFis9lw9OhRTJs27acHMZkwbdo0HDp0qM37HDp0yC0eAJKTk9uNb25uRm1trduFiLqvS5E3AOYOou6sywuWyspKOJ1OREVFuS2PiopCWVlZm/cpKytTis/OzobVanVdeNAcUfd2KfIGwNxB1J11j3MPL5CZmYmamhrXpaSkxNdDIqJugLmDqPvq8oNu+/TpA7PZjPJy96ODy8vLER0d3eZ9oqOjleItFgssFvmDrIjI2C5F3gCYO4i6sy7/hiUoKAhjx45Fbm6ua5mu68jNzUVSUlKb90lKSnKLB4C9e/e2G09E/oV5g4g88cppzRkZGUhLS8P111+PcePGYc2aNWhoaMDChQsBAAsWLEC/fv2QnZ0NAHjooYcwadIkPP/887j99tvx5ptv4siRI3jllVe8MTwiMiDmDSLqiFcKlnnz5qGiogIrVqxAWVkZxowZgz179rgOkCsuLnZrwXvjjTfi9ddfx+OPP47HHnsMQ4cOxc6dO6V7KbQIOK/BJDyfyx1Qr9Ynom6IfK+N4P518iuukm9g1BTrkF8vgB7FCn1H5PsGQQ9UfO4GyJ+3L/HS/RSr0N4FUOu9UzNYfuWawkBCT8n3OwCA5gj5plxhBfLvO9EvUmkcmuSZNJowAfK97lrxVd4g6pC3ZmwwykwQKn2TfDzlhiaEYvtNA6qtrYXVakXciqelGsepFixNUd4pWBoVChboamNWKVhCS+TfAqqFglB4fxulYLGFyw9EU9h7VAsWlUSiUrBoTrVdXvyrQCrOIez4qPkt1NTUIDw8XOkxfKUld5z7vzipxnFnFDvdVig0jqtS6Eh71infBbbaqdbZ9axCp9tKu3xshU2tA3JFk8K6z8sX91W18rEA0KzS6fasQkdtlU63FWr7bI8KhU63ZV7sdPvDac8xwo592CWVN7rlWUJERER0eWHBQkRERIbHgoWIiIgMjwULERERGR4LFiIiIjI8FixERERkeCxYiIiIyPBYsBAREZHhsWAhIiIiw2PBQkRERIbnlbmEfMUW6YApxPOcO7YotTbHVw/x3F64xakaq3RseN966djmLyOkYwG1lvj1/eVbRJvUpjRSoiu029flpygCADiDFaYfUOi3rzJFgB6gNp9Aj3L5cTRFybdf71FwTmkcWoTce9qk24AzSqumS8CpkgwAOBUmuXEqfObVVebe8CKVqXN+vIPC3wuV6UVUpi1R/WpBZRuVnxDf4TcsREREZHgsWIiIiMjwWLAQERGR4bFgISIiIsNjwUJERESGx4KFiIiIDI8FCxERERlelxcs2dnZuOGGGxAWFoa+ffti1qxZOH78eIf3ycnJgaZpbpfg4OCuHhoRGRTzBhF50uUFy/79+7FkyRIcPnwYe/fuhd1ux6233oqGhoYO7xceHo7S0lLX5eTJk109NCIyKOYNIvKkyzvd7tmzx+16Tk4O+vbti6NHj2LixInt3k/TNERHR3f1cIioG2DeICJPvN6av6amBgDQu3fvDuPq6+sxcOBA6LqO6667Ds888wxGjRrVZmxzczOam5td12trawEAWpATWpDT45ii+tbIDv/Hx3PKP02Desm3Pf86b5B0rBaqNp0AIB9vssm3ZlZt4qzUbj9El44VIZ5f558zS0zZ0CIgQH7dtkb5OQIagtV2N5NN/skLPqfQTt0aojQOc5Xke1pXe0064o28AbSfO/yZSvt8ANAV+sB7s92+rpxt5GgqrfYBxXb7CtN6mBRWrPhbiMqUIUovoY/b+Hv1oFtd1/H73/8eN910E0aPHt1u3PDhw7Fp0ybs2rULW7duha7ruPHGG3Hq1Kk247Ozs2G1Wl2X/v37e2sTiOgS81beAJg7iLozrxYsS5Yswddff40333yzw7ikpCQsWLAAY8aMwaRJk7Bjxw5ERkbi5ZdfbjM+MzMTNTU1rktJSYk3hk9EPuCtvAEwdxB1Z177SWjp0qX4n//5Hxw4cABXXnml0n0DAwNx7bXXIj8/v83bLRYLLBZLVwyTiAzEm3kDYO4g6s66/BsWIQSWLl2Kd955Bx9++CEGDx6svA6n04mvvvoKMTExXT08IjIg5g0i8qTLv2FZsmQJXn/9dezatQthYWEoKysDAFitVoSE/Hiw34IFC9CvXz9kZ2cDAJ544gn84he/wFVXXYXq6mo899xzOHnyJBYtWtTVwyMiA2LeICJPurxg2bBhAwBg8uTJbss3b96M9PR0AEBxcTFMpp++3Dl37hwWL16MsrIy9OrVC2PHjsXBgwcxcuTIrh4eERkQ8wYRedLlBYsQnk/r2rdvn9v1F198ES+++GJXD4WIugnmDSLyhHMJERERkeGxYCEiIiLDY8FCREREhuf11vyXUlCIHeYennsS9w+rVlqvSaGV8ydfXSW/YotCK3qzYjtph3wLZV2lU3uA2jhMCi3xQ0ObpGMjQzueFK9VfEi9UryssoZw6diTJX2U1m2LkO+vrQcqTK9gl3/f0aWj0hLfptB7XbV9vtNLrfm91WofUJsyRLU1v6aQe5Va4qtMW6L4l1o3yz8jwqzwvYU/t+YnIiIi6gosWIiIiMjwWLAQERGR4bFgISIiIsNjwUJERESGx4KFiIiIDI8FCxERERkeCxYiIiIyPBYsREREZHgsWIiIiMjw/Ko1v9A16Lrn1sHHK/sqrbfulHz7daXW9YEKLdIltstNkPy6TUFO6djgHjalYUSGybfQHxJeKR07IrRUaRyxgdXSsQ26RTr2E3OcdGxZWJh0LABoepB0rFOhNb+pQX4KBADQAuTShKaz5f+FnApN41ViVdr424Vamrcr9Ix3KnzmdSpOEeAtyt3lTQqt+RXyv0q7fT1AbdBCqTW/wrpNbM1PRERE1CEWLERERGR4XV6wZGVlQdM0t8uIESM6vM/bb7+NESNGIDg4GNdccw12797d1cMiIgNj3iAiT7zyDcuoUaNQWlrqunz88cftxh48eBCpqam477778Pnnn2PWrFmYNWsWvv76a28MjYgMinmDiDrilYIlICAA0dHRrkufPn3ajf3jH/+I6dOn4+GHH8bVV1+NJ598Etdddx3WrVvnjaERkUExbxBRR7xSsJw4cQKxsbGIi4vD3XffjeLi4nZjDx06hGnTprktS05OxqFDh9q9T3NzM2pra90uRNS9eTtvAMwdRN1ZlxcsiYmJyMnJwZ49e7BhwwYUFhZiwoQJqKurazO+rKwMUVFRbsuioqJQVlbW7mNkZ2fDarW6Lv379+/SbSCiS+tS5A2AuYOoO+vygiUlJQVz5sxBfHw8kpOTsXv3blRXV+Ott97qssfIzMxETU2N61JSUtJl6yaiS+9S5A2AuYOoO/N647iIiAgMGzYM+fn5bd4eHR2N8vJyt2Xl5eWIjo5ud50WiwUWi3xzLyLqXryRNwDmDqLuzOt9WOrr61FQUICYmJg2b09KSkJubq7bsr179yIpKcnbQyMig2LeIKILdXnBsmzZMuzfvx9FRUU4ePAgZs+eDbPZjNTUVADAggULkJmZ6Yp/6KGHsGfPHjz//PP47rvvkJWVhSNHjmDp0qVdPTQiMijmDSLypMt/Ejp16hRSU1NRVVWFyMhIjB8/HocPH0ZkZCQAoLi4GCbTT3XSjTfeiNdffx2PP/44HnvsMQwdOhQ7d+7E6NGjlR9biB8vntSdtCqt12SXj9WDFVbskJ+zQ2WOCgCAWT4+QGEuoV49zysNY3hEueeg/y8pvEA69rrg9s8gaUukySEdW+KU/8mguPkK6ViTwpwkAKAy9YrZpvj+UKBX18jFCbV5pn7Ol3nDm3SFF1Flzh+b0nw/avO/6ArxDl1+HCrzH/0Y7515a0wmtTmvNIX9VlfIuyo5XSjOJaQyT5EIUHhdTL5tjt/lBcubb77Z4e379u1rtWzOnDmYM2dOVw+FiLoJ5g0i8oRzCREREZHhsWAhIiIiw2PBQkRERIbHgoWIiIgMjwULERERGR4LFiIiIjI8FixERERkeCxYiIiIyPBYsBAREZHheX225ksp6PNQmC2ee+MHBqqtV1eY3FU0KrRxVhiHHqTWet2peadVe89Atfbr0ZZa6di4oDPysYrv3FBTqHRso6iXjjVr8m2+nU61zwcBClNCaLpCm+9gtR1At8kNRBfy0x90V06ZuT9+xqbwmVCl3b5KG3+7wnoBwK7Sbl+hjb9qq31vteZXpdKaX+UrAIWnWanV/o/x8s+dMCs8z5pvXxN+w0JERESGx4KFiIiIDI8FCxERERkeCxYiIiIyPBYsREREZHgsWIiIiMjwWLAQERGR4XV5wTJo0CBomtbqsmTJkjbjc3JyWsUGB3vupUJE/oW5g4g60uWN4z777DM4nU7X9a+//hq33HIL5syZ0+59wsPDcfz4cdd1zcfNaYjo0mPuIKKOdHnBEhkZ6Xb9P/7jPzBkyBBMmjSp3ftomobo6OiuHgoRdSPMHUTUEa+25rfZbNi6dSsyMjI6/ORTX1+PgQMHQtd1XHfddXjmmWcwatSoduObm5vR3Nzsul5b+2P7d2H+8eKJpVqtvbYwKbSfVmiJrNLy3xGi9slRpau1PVj+bVBvD1IaR6NTPr5R4QlpFA1K44DeJB1a4QyRjq2yybf8d9jUWqQHKXS6N9nl39PaebXpFUxBcq38TUIA8k9zhy517pDl9BziRhfyv7rrCr/QOxVi7Yp93VVa+TsVEo1KG39vUh2FUrxCG3+VGRPUW/OrxKq08fftYa9effSdO3eiuroa6enp7cYMHz4cmzZtwq5du7B161bouo4bb7wRp06davc+2dnZsFqtrkv//v29MHoi8hXmDiK6kFcLlj//+c9ISUlBbGxsuzFJSUlYsGABxowZg0mTJmHHjh2IjIzEyy+/3O59MjMzUVNT47qUlJR4Y/hE5CPMHUR0Ia/9JHTy5El88MEH2LFjh9L9AgMDce211yI/P7/dGIvFAotF4fcUIuo2mDuIqC1e+4Zl8+bN6Nu3L26//Xal+zmdTnz11VeIiYnx0siIyMiYO4ioLV4pWHRdx+bNm5GWloaAAPcvcRYsWIDMzEzX9SeeeALvv/8+vv/+exw7dgz33HMPTp48iUWLFnljaERkYMwdRNQer/wk9MEHH6C4uBi//vWvW91WXFwMk+mnOuncuXNYvHgxysrK0KtXL4wdOxYHDx7EyJEjvTE0IjIw5g4iao9XCpZbb70VQrR9ete+ffvcrr/44ot48cUXvTEMIupmmDuIqD2cS4iIiIgMjwULERERGR4LFiIiIjI8FixERERkeF6dS+hSCy9yIiBQdbYPz5TmEpKbdgUA4LDIr9cUqjYDhqbL16JNZvn5fios8nPnAMC/guUnprMGnJeOtalMxAEgWLNLx37b3E86Nr+uj3Ss3qDw5gBgUpnyR2V6LJPa5xQtWK7Rmia0LptLyKgUpmwCADgVZqJReU+rzPejOoePrjA/kENXGIfKBGeK8YovixJNU1i7yiYqrVdxLjmFXVzl75vqOLoav2EhIiIiw2PBQkRERIbHgoWIiIgMjwULERERGR4LFiIiIjI8FixERERkeCxYiIiIyPBYsBAREZHhsWAhIiIiw2PBQkRERIbnV635g885EBDg8BinByi2uVfooKyybmewfL0Y0Kw2ZpNdoRZVaOPfLHoojeM7PUo6trY5WDr229AYpXGEmOVb85c3hUnHFlX2lo4NqFWbTsBsk3/jqbxHodKKGwACJadu0NVWezlwKnwm1BX6qavEOlX6tENxzAq96FVb8wvFeOn1qsZ7aRwqbfxVh6D0kitNJ8DW/EREREQdUi5YDhw4gBkzZiA2NhaapmHnzp1utwshsGLFCsTExCAkJATTpk3DiRMnPK53/fr1GDRoEIKDg5GYmIhPP/1UdWhEZFDMG0TUWcoFS0NDAxISErB+/fo2b1+9ejVeeuklbNy4EZ988gl69uyJ5ORkNDW1P43r9u3bkZGRgZUrV+LYsWNISEhAcnIyzpw5ozo8IjIg5g0i6izlgiUlJQVPPfUUZs+e3eo2IQTWrFmDxx9/HDNnzkR8fDxee+01nD59utUnqp974YUXsHjxYixcuBAjR47Exo0b0aNHD2zatEl1eERkQMwbRNRZXXoMS2FhIcrKyjBt2jTXMqvVisTERBw6dKjN+9hsNhw9etTtPiaTCdOmTWv3Ps3NzaitrXW7EFH3dKnyBsDcQdSddWnBUlZWBgCIinI/MyQqKsp124UqKyvhdDqV7pOdnQ2r1eq69O/fvwtGT0S+cKnyBsDcQdSddcuzhDIzM1FTU+O6lJSU+HpIRNQNMHcQdV9dWrBER0cDAMrLy92Wl5eXu267UJ8+fWA2m5XuY7FYEB4e7nYhou7pUuUNgLmDqDvr0oJl8ODBiI6ORm5urmtZbW0tPvnkEyQlJbV5n6CgIIwdO9btPrquIzc3t937EJH/YN4gIhnKnW7r6+uRn5/vul5YWIi8vDz07t0bAwYMwO9//3s89dRTGDp0KAYPHozly5cjNjYWs2bNct1n6tSpmD17NpYuXQoAyMjIQFpaGq6//nqMGzcOa9asQUNDAxYuXNj5LSQin2PeIKLOUi5Yjhw5gilTpriuZ2RkAADS0tKQk5ODP/zhD2hoaMD999+P6upqjB8/Hnv27EFw8E9t1wsKClBZWem6Pm/ePFRUVGDFihUoKyvDmDFjsGfPnlYH1HkSWNOMAInu53pwoNJ6Nad8z3FhVmhrfV4+1tyk2tZdJVa+3bLZpvalXHNjiHRsSbVkC3gAp3r2UhqHKVD+NdRtCs91rfwuFFyrOL2CQmt+ldbdIlDtvWQKlZuOQdPNQGXbtxk5b6hQnX1AqYW+whfeTpWW+Eq919Vb6Htrvd5q+6/aal+plb9KsOocAQqUpuroRjQhRLfftNraWlitVky59lEEmC0e4w1TsATJxzpC1P7I2EPl45vD5Xfg5l5qO3tzL/m3lz3CKR2r9fQ8Z9TPGaJgOaP2GvYolX/uQk/LPx8hP9QpjcNUd14qzqE344Oidaipqek2x4a05I5z/xeH8DDP+2OBvV5p/SUO+efhB4d8EV7hkJ/v6pyjp3QsAJyzy88XVueQn/+rxiYfCwB1dvn4epv8h536Js9/I37u/Hn5dTvr5f+2mBXmFgs6p/ZBsccZhdzxg3zu6FFwTmkczuP5HmMcwo592CWVN7rlWUJERER0eWHBQkRERIbHgoWIiIgMjwULERERGR4LFiIiIjI8FixERERkeCxYiIiIyPBYsBAREZHhsWAhIiIiw2PBQkRERIanPJeQPzDXNCnewTtza5gU5nQxNSvOJdQs/9IGnJdfd2CDWo1rqZZ/7uyh8mN2hKi9dXX57tpKc3yYFOZsClTr6q40H5TK3CHOULXW5DDJvea687JMJ4bnVJjPCFCc/0hhXh6V9f4Y7535gVTnElKaqEuXj9WUYuWHAMAwcxp1NX7DQkRERIbHgoWIiIgMjwULERERGR4LFiIiIjI8FixERERkeCxYiIiIyPCUC5YDBw5gxowZiI2NhaZp2Llzp+s2u92ORx55BNdccw169uyJ2NhYLFiwAKdPn+5wnVlZWdA0ze0yYsQI5Y0hImNi3iCizlIuWBoaGpCQkID169e3uq2xsRHHjh3D8uXLcezYMezYsQPHjx/HHXfc4XG9o0aNQmlpqevy8ccfqw6NiAyKeYOIOku501NKSgpSUlLavM1qtWLv3r1uy9atW4dx48ahuLgYAwYMaH8gAQGIjo5WHQ4RdQPMG0TUWV4/hqWmpgaapiEiIqLDuBMnTiA2NhZxcXG4++67UVxc3G5sc3Mzamtr3S5E5D+8kTcA5g6i7syrvbSbmprwyCOPIDU1FeHh4e3GJSYmIicnB8OHD0dpaSlWrVqFCRMm4Ouvv0ZYWFir+OzsbKxatarVclOjDSaZNvqaYmtm1bbIkjSbQz62Sa01v+m8/LoDGhRa89epvWWcwfLrdoQotAS3qL2GziCFFuJqT7U0k/xLAgAIbJR/49l7yj93eoBaa/4gk9xz53B0zY7irbwBtJ876OKotNvXobbPqrTmd+oK41BoiQ8AulOlNb98qEq7fZNTPhYANIV4lWk9IJkLvMVr37DY7XbMnTsXQghs2LChw9iUlBTMmTMH8fHxSE5Oxu7du1FdXY233nqrzfjMzEzU1NS4LiUlJd7YBCK6xLyZNwDmDqLuzCvfsLQknZMnT+LDDz/s8FNSWyIiIjBs2DDk5+e3ebvFYoHFojiBGxEZmrfzBsDcQdSddfk3LC1J58SJE/jggw9wxRVXKK+jvr4eBQUFiImJ6erhEZEBMW8QkSfKBUt9fT3y8vKQl5cHACgsLEReXh6Ki4tht9vxy1/+EkeOHMG2bdvgdDpRVlaGsrIy2Gw21zqmTp2KdevWua4vW7YM+/fvR1FREQ4ePIjZs2fDbDYjNTW181tIRD7HvEFEnaX8k9CRI0cwZcoU1/WMjAwAQFpaGrKysvDf//3fAIAxY8a43e+jjz7C5MmTAQAFBQWorKx03Xbq1CmkpqaiqqoKkZGRGD9+PA4fPozIyEjV4RGRATFvEFFnKRcskydPhhDtH1bc0W0tioqK3K6/+eabqsMgom6EeYOIOotzCREREZHhsWAhIiIiw2PBQkRERIbHgoWIiIgMz6ut+S81zWaHZvJcg4kQxcZRNrvCIBRaF0scaNjC1KwwBgBCpYVyg3zdag5Q61svAuXjVWJ1i9pbVw9SaPuvECsCvNeq2t5T/vlwBsqvN6BJsYW+7HtadcoLuqyotNoHAKHUml8h1qn2OV0otObXHAqxKn9WFKf1UJkGRNMVevM7vTRPjSR+w0JERESGx4KFiIiIDI8FCxERERkeCxYiIiIyPBYsREREZHgsWIiIiMjwWLAQERGR4bFgISIiIsNjwUJERESGx4KFiIiIDM+vWvM3DeiNgIBgj3GWH2rUVhwk3/dcO9+stm5ZTqdSuHZerZW/t8hMleCi0PbfFKj41lV4DUWQ/Lr1IPkx23opTgmhMHWDpUahvbYi2dbdSi2+yS+otNtXbc3v0OVzh0qsrtiaH3b5eJNdfhvVYqVDf4x3yu+LmkNhv9XZmp+IiIioQyxYiIiIyPCUC5YDBw5gxowZiI2NhaZp2Llzp9vt6enp0DTN7TJ9+nSP612/fj0GDRqE4OBgJCYm4tNPP1UdGhEZFPMGEXWWcsHS0NCAhIQErF+/vt2Y6dOno7S01HV54403Olzn9u3bkZGRgZUrV+LYsWNISEhAcnIyzpw5ozo8IjIg5g0i6izlg25TUlKQkpLSYYzFYkF0dLT0Ol944QUsXrwYCxcuBABs3LgR7777LjZt2oRHH31UdYhEZDDMG0TUWV45hmXfvn3o27cvhg8fjgceeABVVVXtxtpsNhw9ehTTpk37aVAmE6ZNm4ZDhw61eZ/m5mbU1ta6XYioe/N23gCYO4i6sy4vWKZPn47XXnsNubm5ePbZZ7F//36kpKTA2c5puZWVlXA6nYiKinJbHhUVhbKysjbvk52dDavV6rr079+/qzeDiC6hS5E3AOYOou6sy/uwzJ8/3/X/a665BvHx8RgyZAj27duHqVOndsljZGZmIiMjw3W9traWiYeoG7sUeQNg7iDqzrx+WnNcXBz69OmD/Pz8Nm/v06cPzGYzysvL3ZaXl5e3+3u2xWJBeHi424WI/Ic38gbA3EHUnXm9YDl16hSqqqoQExPT5u1BQUEYO3YscnNzXct0XUdubi6SkpK8PTwiMiDmDSK6kHLBUl9fj7y8POTl5QEACgsLkZeXh+LiYtTX1+Phhx/G4cOHUVRUhNzcXMycORNXXXUVkpOTXeuYOnUq1q1b57qekZGBV199FVu2bMG//vUvPPDAA2hoaHAd/U9E3RvzBhF1lvIxLEeOHMGUKVNc11t+D05LS8OGDRvw5ZdfYsuWLaiurkZsbCxuvfVWPPnkk7BYfppHpaCgAJWVla7r8+bNQ0VFBVasWIGysjKMGTMGe/bsaXVAnSdBVY0IMHuec8fZq4fSes3nGqVjRXCQdKzWcF5+EE61ORyEQ2HuIbtNPtab88WYFOYaMcvP4QMAWpDC66Iw7xAiQqVDzTaF9QIIqVB4zRWeOkew2nNX39/z/FwA4OhgvhMj5w0Vqp/wTJr8a2iGSqz8fmhWGIM3qc4l5FSZp0hpLiG1cUAhXvPS/EAmlfl+AJgcCrF2+feHpvJ3xQs0IRRmWDOo2tpaWK1W3Dz6YQSYPU8wp/eQ/+MFqBUsMMvvOF4tWJoVihAWLO4UChZdoWCx9w6RHwMAhb9JXi1YbFa5eIe9CUf/+jhqamq6zbEhLbnj3P/FITzM875baK9XWn+JU/79cdreSzq2wiH//FY65McAAFU2+fizNvkPf2eb1T4o1jTLFcoAUHdePvZ8g1r+1+vl80FAjfy+FVQjv9NaqtXybkiVfHyPUvkJe4NK2m810BZHUbHnGGHHPuySyhucS4iIiIgMjwULERERGR4LFiIiIjI8FixERERkeCxYiIiIyPBYsBAREZHhsWAhIiIiw2PBQkRERIbHgoWIiIgMT7k1v5GJQDOERAdUzaHWNVallb/5bIN0rFBpAV+r1mETunwLZWGX7+MsHAo9nwGvdcbVFDoKA4Boku/maAoPkx+HTf75CDoj/94AAD1Ifve0XSHf6VOVPUSuI6fTrNjyvBvy5iaqtPE3Cl2hxbJQbc2v0G7f6VRoze9Qyx2aQ6HdvkLneqVYxbSr0spf6e+hYsf1rsZvWIiIiMjwWLAQERGR4bFgISIiIsNjwUJERESGx4KFiIiIDI8FCxERERkeCxYiIiIyPOWC5cCBA5gxYwZiY2OhaRp27tzpdrumaW1ennvuuXbXmZWV1Sp+xIgRyhtDRMbEvEFEnaVcsDQ0NCAhIQHr169v8/bS0lK3y6ZNm6BpGu66664O1ztq1Ci3+3388ceqQyMig2LeIKLOUu50m5KSgpSUlHZvj46Odru+a9cuTJkyBXFxcR0PJCCg1X2JyD8wbxBRZ3m1NX95eTneffddbNmyxWPsiRMnEBsbi+DgYCQlJSE7OxsDBgxoM7a5uRnNzT+1Wq+trQUANEWGICDQc4tyy1n5Nu0AYGq0yQcHeJ4a4GJoQUFK8eJ8k/y6JaYzcK23We25Eyqt+RWmEwDUng+TwvMn7Hb5FVeclQ7VLBb59QJA7BXSoSqtuGsHqu32gefl1q3Zu2YaBm/lDaD93OEtZninlbk32/irtNvXFdrtq6wXAJy6fLzTqbBulViotebXFFKYt2IBwKSwL5pUWvML70y1IsurB91u2bIFYWFhuPPOOzuMS0xMRE5ODvbs2YMNGzagsLAQEyZMQF1dXZvx2dnZsFqtrkv//v29MXwi8gFv5Q2AuYOoO/NqwbJp0ybcfffdCA7u+FuPlJQUzJkzB/Hx8UhOTsbu3btRXV2Nt956q834zMxM1NTUuC4lJSXeGD4R+YC38gbA3EHUnXntJ6H//d//xfHjx7F9+3bl+0ZERGDYsGHIz89v83aLxQKL6tfrRGR43swbAHMHUXfmtW9Y/vznP2Ps2LFISEhQvm99fT0KCgoQExPjhZERkVExbxBRe5QLlvr6euTl5SEvLw8AUFhYiLy8PBQXF7tiamtr8fbbb2PRokVtrmPq1KlYt26d6/qyZcuwf/9+FBUV4eDBg5g9ezbMZjNSU1NVh0dEBsS8QUSdpfyT0JEjRzBlyhTX9YyMDABAWloacnJyAABvvvkmhBDtJo6CggJUVla6rp86dQqpqamoqqpCZGQkxo8fj8OHDyMyMlJ1eERkQMwbRNRZmhA+Pk+pC9TW1sJqtSLp1lU+P61ZUzlFTIFW16gUL2q8c7qm3qg4Di+d1qwFKp7WHOL5feESFKi0blmqpzU7FU5rdvSUH3P1ELVxyJ7W7LQ34ehbj6Ompgbh4eFKj+ErLbnj3P/FITzM8xfOxY56pfWXOHpIx/7g6CUdW+GQf37P2NVeiwpbmHRsZXNP6dhzzfLPBQCcbQyRjq1rkN+/7fWKuaNO/nN9YJ38KdBBNfKxlnNqf6Z7VMjn0pCy89Kx5lL5Ng4A4Dj1g+cYYcc+7JLKG5xLiIiIiAyPBQsREREZHgsWIiIiMjwWLERERGR4Xp1L6FIz2XWYhOeDXu2hagdVmkLkn6agsvbbgl9Ia1KYo0hlfhsAWi+rdKxeXiEdK5xqk1qYFA40FU75OY00s1qt7axvUIqXZQ6VP+hQ9I5QWrdml3+unQrPs8mhNAxA9ni/bn/4vm+ZvfQEmhTXqxovS2XeIQAQKvHeigXU3tdK45APVZ46yk/3RX7DQkRERIbHgoWIiIgMjwULERERGR4LFiIiIjI8FixERERkeCxYiIiIyPBYsBAREZHhsWAhIiIiw2PBQkRERIbnF51uhfixrZ/D0eyV9Zuc8m0DTU75MWi6QqdblVgA0OXHrAv5detCreOuSaHzY8vrKENT7FapC9X2rnKEwnMHhfcGAAiFpsIOh/xnD6dNrVsx7HKvi9PeBEDtdfS1lrHW1su1Eq1zqLUcbVCIb3TIvy7nnfLv52bFLtk2m3y83Sb//nc0qf25cZ6Xf0/rjfLr1c8rto1tku/A7WxS2A8V0oHTprZPORS6ZDucTdKxQlfLYQ6JvxcO/Bgjkzf8omCpq/uxHf4n//sfPh4JtSK/L6hRy8HeU+ul2G6srq4OVqv81BC+1JI7Bl5X5NuBEF3mZPKGJrrTx6F26LqO06dPIywsDJr20yfv2tpa9O/fHyUlJQgPD/fhCL3D37cP8P9t9KftE0Kgrq4OsbGxMJm6x6/NzB3cvu7KX7ZRJW/4xTcsJpMJV155Zbu3h4eHd+sX1BN/3z7A/7fRX7avu3yz0oK5g9vX3fnDNsrmje7xMYiIiIguayxYiIiIyPD8umCxWCxYuXIlLBaLr4fiFf6+fYD/b6O/b1935e+vC7ev+7sctvFCfnHQLREREfk3v/6GhYiIiPwDCxYiIiIyPBYsREREZHgsWIiIiMjwWLAQERGR4fl1wbJ+/XoMGjQIwcHBSExMxKeffurrIXWJrKwsaJrmdhkxYoSvh9UpBw4cwIwZMxAbGwtN07Bz506324UQWLFiBWJiYhASEoJp06bhxIkTvhnsRfC0fenp6a1e0+nTp/tmsJc5f80bgP/lDuaNyytv+G3Bsn37dmRkZGDlypU4duwYEhISkJycjDNnzvh6aF1i1KhRKC0tdV0+/vhjXw+pUxoaGpCQkID169e3efvq1avx0ksvYePGjfjkk0/Qs2dPJCcno6nJW7Mrdi1P2wcA06dPd3tN33jjjUs4QgL8P28A/pU7mDcus7wh/NS4cePEkiVLXNedTqeIjY0V2dnZPhxV11i5cqVISEjw9TC8BoB45513XNd1XRfR0dHiueeecy2rrq4WFotFvPHGGz4YYedcuH1CCJGWliZmzpzpk/HQT/w5bwjh37mDecP/+eU3LDabDUePHsW0adNcy0wmE6ZNm4ZDhw75cGRd58SJE4iNjUVcXBzuvvtuFBcX+3pIXlNYWIiysjK319NqtSIxMdFvXk8A2LdvH/r27Yvhw4fjgQceQFVVla+HdFm5HPIGcPnkDuYN/+OXBUtlZSWcTieioqLclkdFRaGsrMxHo+o6iYmJyMnJwZ49e7BhwwYUFhZiwoQJqKur8/XQvKLlNfPX1xP48Wvd1157Dbm5uXj22Wexf/9+pKSkwOl0+npolw1/zxvA5ZU7mDf8T4CvB0DqUlJSXP+Pj49HYmIiBg4ciLfeegv33XefD0dGF2v+/Pmu/19zzTWIj4/HkCFDsG/fPkydOtWHIyN/wtzhXy63vOGX37D06dMHZrMZ5eXlbsvLy8sRHR3to1F5T0REBIYNG4b8/HxfD8UrWl6zy+X1BIC4uDj06dPHb19TI7rc8gbg37mDecP/+GXBEhQUhLFjxyI3N9e1TNd15ObmIikpyYcj8476+noUFBQgJibG10PxisGDByM6Otrt9aytrcUnn3zil68nAJw6dQpVVVV++5oa0eWWNwD/zh3MG/7Hb38SysjIQFpaGq6//nqMGzcOa9asQUNDAxYuXOjroXXasmXLMGPGDAwcOBCnT5/GypUrYTabkZqa6uuhXbT6+nq3TwWFhYXIy8tD7969MWDAAPz+97/HU089haFDh2Lw4MFYvnw5YmNjMWvWLN8NWkFH29e7d2+sWrUKd911F6Kjo1FQUIA//OEPuOqqq5CcnOzDUV9+/DlvAP6XO5g3LrO84evTlLxp7dq1YsCAASIoKEiMGzdOHD582NdD6hLz5s0TMTExIigoSPTr10/MmzdP5Ofn+3pYnfLRRx8JAK0uaWlpQogfT1Fcvny5iIqKEhaLRUydOlUcP37ct4NW0NH2NTY2iltvvVVERkaKwMBAMXDgQLF48WJRVlbm62Fflvw1bwjhf7mDeePyyhuaEEJc2hKJiIiISI1fHsNCRERE/oUFCxERERkeCxYiIiIyPBYsREREZHgsWIiIiMjwWLAQERGR4bFgISIiIsNjwUJERESGx4KFiIiIDI8FCxERERkeCxYiIiIyvP8HXODpCG4iMjAAAAAASUVORK5CYII=",
       "text/plain": [
        "<Figure size 640x480 with 2 Axes>"
-      ]
+      ],
+      "image/png": "iVBORw0KGgoAAAANSUhEUgAAAiwAAAEjCAYAAAARyVqhAAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjkuMiwgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy8hTgPZAAAACXBIWXMAAA9hAAAPYQGoP6dpAAA73klEQVR4nO3de3RTZbo/8O9O2qaFtilI6QW5FbkI2KIonSpXQUv1IKDDpaPSMoIzHhid1YWjdQkUbz3iURmBA+oMFAdUdAbhzBEZsQocB1AB63Xk0NrSIm1pC73T5rLf3x/+Gg295H1pQ3bD97NWFmTnyc67k+ynT5K9n1cTQggQERERGZjJ1wMgIiIi8oQFCxERERkeCxYiIiIyPBYsREREZHgsWIiIiMjwWLAQERGR4bFgISIiIsNjwUJERESGx4KFiIiIDI8FC3VbWVlZ0DQNlZWVHmMHDRqE9PR01/V9+/ZB0zTs27fPtSw9PR2DBg3q+oESGcxnn32GG2+8ET179oSmacjLy3PtTxdDdt8pKiqCpmnIycm5qMe5GDk5OdA0DUVFRV263smTJ2Py5Mlduk7qGAsWonY0NjYiKyvLragh6u7sdjvmzJmDs2fP4sUXX8Rf/vIXDBw40NfDMqRvv/0WWVlZXV7s0MUJ8PUAiC6F48ePw2TquD5/9dVXoeu663pjYyNWrVoFAPwkRX6joKAAJ0+exKuvvopFixa5lj/++ON49NFHfTgy4/n222+xatUqTJ48udU3SO+//75vBnUZY8FCHjU2NqJHjx6+HkanWCwWjzGBgYGXYCREvnXmzBkAQEREhNvygIAABATwT4KsoKAgXw/hssOfhLq5lt+dv/vuO8ydOxfh4eG44oor8NBDD6GpqcktduvWrRg7dixCQkLQu3dvzJ8/HyUlJW4xkydPxujRo3H06FFMnDgRPXr0wGOPPeb67fk///M/sX79esTFxaFHjx649dZbUVJSAiEEnnzySVx55ZUICQnBzJkzcfbs2Vbjfe+99zBhwgT07NkTYWFhuP322/HNN9+4xXz55ZdIT09HXFwcgoODER0djV//+teoqqpq8zmorKz0uO0XHsPSlp//Dl9UVITIyEgAwKpVq6BpGjRNQ1ZWFjZv3gxN0/D555+3WsczzzwDs9mMH374ocPHIvKF9PR0TJo0CQAwZ84caJrm+vawvWNYZPJGW6qrq5Geng6r1YqIiAikpaWhurpaapx2ux2rVq3C0KFDERwcjCuuuALjx4/H3r173eI+/PBDVz6JiIjAzJkz8a9//cvj+lv25Qv9PE/k5ORgzpw5AIApU6a4ckDLT8RtHcNy5swZ3HfffYiKikJwcDASEhKwZcsWt5if59JXXnkFQ4YMgcViwQ033IDPPvtM6vm5XLGc9hNz587FoEGDkJ2djcOHD+Oll17CuXPn8NprrwEAnn76aSxfvhxz587FokWLUFFRgbVr12LixIn4/PPP3T5tVVVVISUlBfPnz8c999yDqKgo123btm2DzWbD7373O5w9exarV6/G3LlzcfPNN2Pfvn145JFHkJ+fj7Vr12LZsmXYtGmT675/+ctfkJaWhuTkZDz77LNobGzEhg0bMH78eHz++eeuYmHv3r34/vvvsXDhQkRHR+Obb77BK6+8gm+++QaHDx9ulVQ9bfvFiIyMxIYNG/DAAw9g9uzZuPPOOwEA8fHxGDx4MJYsWYJt27bh2muvdbvftm3bMHnyZPTr1++iH5vIW37zm9+gX79+eOaZZ/Dggw/ihhtucNu/L6SSN35OCIGZM2fi448/xm9/+1tcffXVeOedd5CWliY1zqysLGRnZ2PRokUYN24camtrceTIERw7dgy33HILAOCDDz5ASkoK4uLikJWVhfPnz2Pt2rW46aabcOzYsU4fQD9x4kQ8+OCDeOmll/DYY4/h6quvBgDXvxc6f/48Jk+ejPz8fCxduhSDBw/G22+/jfT0dFRXV+Ohhx5yi3/99ddRV1eH3/zmN9A0DatXr8add96J77//nt/2tkdQt7Zy5UoBQNxxxx1uy//93/9dABBffPGFKCoqEmazWTz99NNuMV999ZUICAhwWz5p0iQBQGzcuNEttrCwUAAQkZGRorq62rU8MzNTABAJCQnCbre7lqempoqgoCDR1NQkhBCirq5OREREiMWLF7utt6ysTFitVrfljY2NrbbzjTfeEADEgQMHlLa9xcCBA0VaWprr+kcffSQAiI8++si1LC0tTQwcONB1vaKiQgAQK1eubDWe1NRUERsbK5xOp2vZsWPHBACxefPmVvFERtHy3n/77bfdlrfsTy1U8saF+87OnTsFALF69WrXMofDISZMmCC1jyQkJIjbb7+9w5gxY8aIvn37iqqqKteyL774QphMJrFgwQLXss2bNwsAorCw0LWsvf36wjzx9ttvt8oTLSZNmiQmTZrkur5mzRoBQGzdutW1zGaziaSkJBEaGipqa2uFED/l0iuuuEKcPXvWFbtr1y4BQPz973/vcLsvZ/xJyE8sWbLE7frvfvc7AMDu3buxY8cO6LqOuXPnorKy0nWJjo7G0KFD8dFHH7nd12KxYOHChW0+zpw5c2C1Wl3XExMTAQD33HOP2+/fiYmJsNlsrp9G9u7di+rqaqSmprqNwWw2IzEx0W0MISEhrv83NTWhsrISv/jFLwAAx44dU9p2b1mwYAFOnz7tNu5t27YhJCQEd911l9cel+hSUc0bP7d7924EBATggQcecC0zm82ufdOTiIgIfPPNNzhx4kSbt5eWliIvLw/p6eno3bu3a3l8fDxuueUWr+777dm9ezeio6ORmprqWhYYGIgHH3wQ9fX12L9/v1v8vHnz0KtXL9f1CRMmAAC+//77SzPgbog/CfmJoUOHul0fMmQITCYTioqKYDKZIIRoFdPiwq8f+/Xr1+4BZQMGDHC73lK89O/fv83l586dAwBX4rn55pvbXG94eLjr/2fPnsWqVavw5ptvug4QbFFTU9Pqvh1tu7fccsstiImJwbZt2zB16lTouo433ngDM2fORFhYmNcel+hSOXHihFLe+LmTJ08iJiYGoaGhbsuHDx8u9dhPPPEEZs6ciWHDhmH06NGYPn067r33XsTHx7vW3976rr76avzjH/9AQ0MDevbsKfV4XeHkyZMYOnRoq7MRW35CahlziwtzaUvx0pIzqTUWLH7q58d56LoOTdPw3nvvwWw2t4q9MKn8/BuOC7V1/46WCyFcYwB+PI4lOjq6VdzPv52ZO3cuDh48iIcffhhjxoxBaGgodF3H9OnT3U47bs/FNr9SYTab8atf/Qqvvvoq/uu//gv//Oc/cfr0adxzzz1ef2yiS0E1b3SliRMnoqCgALt27cL777+PP/3pT3jxxRexceNGt1Oxu5rT6fTaui/kKWdSayxY/MSJEycwePBg1/X8/Hzouo5BgwbBbDZDCIHBgwdj2LBhPhnfkCFDAAB9+/bFtGnT2o07d+4ccnNzsWrVKqxYscK1vL2vhltua2/bO8NT4bNgwQI8//zz+Pvf/4733nsPkZGRSE5O7tRjEhnFkCFDLjpvDBw4ELm5uaivr3crbI4fPy69jt69e2PhwoVYuHAh6uvrMXHiRGRlZWHRokWuRndtre+7775Dnz59Ovx2pVevXq3OWLLZbCgtLXVbpvLhZ+DAgfjyyy+h67rbtyzfffed63bqHB7D4ifWr1/vdn3t2rUAgJSUFNx5550wm81YtWpVq+pdCNHu6cJdKTk5GeHh4XjmmWdgt9tb3V5RUQHgp08dF45zzZo17a67o23vjJbeM+2dihkfH4/4+Hj86U9/wt/+9jfMnz+ffSzIb3Qmb9x2221wOBzYsGGDa5nT6XTtm55cuO7Q0FBcddVVaG5uBgDExMRgzJgx2LJli9v++fXXX+P999/Hbbfd1uH6hwwZggMHDrgte+WVV1p9w9JS9Micjn3bbbehrKwM27dvdy1zOBxYu3YtQkNDXaeT08VjdvUThYWFuOOOOzB9+nQcOnQIW7duxa9+9SskJCQAAJ566ilkZmaiqKgIs2bNQlhYGAoLC/HOO+/g/vvvx7Jly7w6vvDwcGzYsAH33nsvrrvuOsyfPx+RkZEoLi7Gu+++i5tuugnr1q1DeHg4Jk6ciNWrV8Nut6Nfv354//33UVhYeNHbfrFCQkIwcuRIbN++HcOGDUPv3r0xevRojB492hWzYMEC13PHn4PInwwZMuSi88aMGTNw00034dFHH0VRURFGjhyJHTt2tHkMWltGjhyJyZMnY+zYsejduzeOHDmCv/71r1i6dKkr5rnnnkNKSgqSkpJw3333uU5rtlqtbfZY+blFixbht7/9Le666y7ccsst+OKLL/CPf/wDffr0cYsbM2YMzGYznn32WdTU1MBiseDmm29G3759W63z/vvvx8svv4z09HQcPXoUgwYNwl//+lf885//xJo1a3hsW1fw0dlJ1EVaTkX89ttvxS9/+UsRFhYmevXqJZYuXSrOnz/vFvu3v/1NjB8/XvTs2VP07NlTjBgxQixZskQcP37cFTNp0iQxatSoVo/Tcirec88957a8vVMkW04l/Oyzz1rFJycnC6vVKoKDg8WQIUNEenq6OHLkiCvm1KlTYvbs2SIiIkJYrVYxZ84ccfr06VanIqps+8Wc1iyEEAcPHhRjx44VQUFBbZ4KWVpaKsxmsxg2bFir54zIiGRPa24hkzfa2neqqqrEvffeK8LDw4XVahX33nuv+Pzzz6VOa37qqafEuHHjREREhAgJCREjRowQTz/9tLDZbG5xH3zwgbjppptESEiICA8PFzNmzBDffvutW0xbpzU7nU7xyCOPiD59+ogePXqI5ORkkZ+f3ypPCCHEq6++KuLi4oTZbHbLGRee1iyEEOXl5WLhwoWiT58+IigoSFxzzTWttrW9XCpE+6db0480IXiET3eWlZWFVatWoaKiotWnA/K+yspKxMTEYMWKFVi+fLmvh0NE5Ld4DAtRJ+Tk5MDpdOLee+/19VCIiPwaj2Ehuggffvghvv32Wzz99NOYNWtWp89IIiKijrFgIboITzzxBA4ePIibbrpJ+swHIiK6eDyGhYiIiAyPx7AQERGR4bFgISIiIsPzi2NYdF3H6dOnERYWdknmkSGi1oQQqKurQ2xsbKsJ4IyKuYPIt1Tyhl8ULKdPn241WzAR+UZJSQmuvPJKXw9DCnMHkTHI5A2vFSzr16/Hc889h7KyMiQkJGDt2rUYN25cu/Fvv/02li9fjqKiIgwdOhTPPvusx/kgWrS0PL5y5eMwBQd7jNc8T/jrpmexwqdFhQ9ptVe3nlOnPZbS9qdy76weZ+SPuw6sVztGO6RCfhthkn/ydLPap+GGaPm3evA5+TdIwHn5WE3x+Pag6mbpWD2w7Zlf2xKQX+o56GeclZVScQ7Y8TF2d6oF+aXMG8BPuePksUEID/W8n1c4G6TXDQCVTvnccVZvf5b0C51ztj+x34VqFNYLAOcc8rMwV9nlx1FpU5vdubJZPr7yfA/p2LN18rEAYKvx/DelRUC1fJ6xVMnnsOAqxbxbKT/rdEh5o3SsufSs0jgcZeWeYxTyhlcKlu3btyMjIwMbN25EYmIi1qxZg+TkZBw/frzNORgOHjyI1NRUZGdn49/+7d/w+uuvY9asWTh27JjbvC3tafkq1xQc7JWCxWzxTsFiCpH/I2O2eK9gMQfJ7wwqsQAQECC/jUoFS4BawWIOkn+rBwQqFCwOhYJFV3zuFIoyXeF5DjAFKY1D0yTfe6Il/uJ+WrnUeePnYw0PNSE8zPN+3qRQgKjGN+sKsU7519vmVEvz5x0Kf3Tt8nkpyKb2vgsMkI8PMFmkY81O+QIEAEw2+XhTk/xzZ7bI7yfKeTdQvmAJMMvHmhVzB2Ryh0Le8MoPzS+88AIWL16MhQsXYuTIkdi4cSN69OiBTZs2tRn/xz/+EdOnT8fDDz+Mq6++Gk8++SSuu+46rFu3zhvDIyIDYt4goo50ecFis9lw9OhRTJs27acHMZkwbdo0HDp0qM37HDp0yC0eAJKTk9uNb25uRm1trduFiLqvS5E3AOYOou6sywuWyspKOJ1OREVFuS2PiopCWVlZm/cpKytTis/OzobVanVdeNAcUfd2KfIGwNxB1J11j3MPL5CZmYmamhrXpaSkxNdDIqJugLmDqPvq8oNu+/TpA7PZjPJy96ODy8vLER0d3eZ9oqOjleItFgssFvmDrIjI2C5F3gCYO4i6sy7/hiUoKAhjx45Fbm6ua5mu68jNzUVSUlKb90lKSnKLB4C9e/e2G09E/oV5g4g88cppzRkZGUhLS8P111+PcePGYc2aNWhoaMDChQsBAAsWLEC/fv2QnZ0NAHjooYcwadIkPP/887j99tvx5ptv4siRI3jllVe8MTwiMiDmDSLqiFcKlnnz5qGiogIrVqxAWVkZxowZgz179rgOkCsuLnZrwXvjjTfi9ddfx+OPP47HHnsMQ4cOxc6dO6V7KbQIOK/BJDyfyx1Qr9Ynom6IfK+N4P518iuukm9g1BTrkF8vgB7FCn1H5PsGQQ9UfO4GyJ+3L/HS/RSr0N4FUOu9UzNYfuWawkBCT8n3OwCA5gj5plxhBfLvO9EvUmkcmuSZNJowAfK97lrxVd4g6pC3ZmwwykwQKn2TfDzlhiaEYvtNA6qtrYXVakXciqelGsepFixNUd4pWBoVChboamNWKVhCS+TfAqqFglB4fxulYLGFyw9EU9h7VAsWlUSiUrBoTrVdXvyrQCrOIez4qPkt1NTUIDw8XOkxfKUld5z7vzipxnFnFDvdVig0jqtS6Eh71infBbbaqdbZ9axCp9tKu3xshU2tA3JFk8K6z8sX91W18rEA0KzS6fasQkdtlU63FWr7bI8KhU63ZV7sdPvDac8xwo592CWVN7rlWUJERER0eWHBQkRERIbHgoWIiIgMjwULERERGR4LFiIiIjI8FixERERkeCxYiIiIyPBYsBAREZHhsWAhIiIiw2PBQkRERIbnlbmEfMUW6YApxPOcO7YotTbHVw/x3F64xakaq3RseN966djmLyOkYwG1lvj1/eVbRJvUpjRSoiu029flpygCADiDFaYfUOi3rzJFgB6gNp9Aj3L5cTRFybdf71FwTmkcWoTce9qk24AzSqumS8CpkgwAOBUmuXEqfObVVebe8CKVqXN+vIPC3wuV6UVUpi1R/WpBZRuVnxDf4TcsREREZHgsWIiIiMjwWLAQERGR4bFgISIiIsNjwUJERESGx4KFiIiIDI8FCxERERlelxcs2dnZuOGGGxAWFoa+ffti1qxZOH78eIf3ycnJgaZpbpfg4OCuHhoRGRTzBhF50uUFy/79+7FkyRIcPnwYe/fuhd1ux6233oqGhoYO7xceHo7S0lLX5eTJk109NCIyKOYNIvKkyzvd7tmzx+16Tk4O+vbti6NHj2LixInt3k/TNERHR3f1cIioG2DeICJPvN6av6amBgDQu3fvDuPq6+sxcOBA6LqO6667Ds888wxGjRrVZmxzczOam5td12trawEAWpATWpDT45ii+tbIDv/Hx3PKP02Desm3Pf86b5B0rBaqNp0AIB9vssm3ZlZt4qzUbj9El44VIZ5f558zS0zZ0CIgQH7dtkb5OQIagtV2N5NN/skLPqfQTt0aojQOc5Xke1pXe0064o28AbSfO/yZSvt8ANAV+sB7s92+rpxt5GgqrfYBxXb7CtN6mBRWrPhbiMqUIUovoY/b+Hv1oFtd1/H73/8eN910E0aPHt1u3PDhw7Fp0ybs2rULW7duha7ruPHGG3Hq1Kk247Ozs2G1Wl2X/v37e2sTiOgS81beAJg7iLozrxYsS5Yswddff40333yzw7ikpCQsWLAAY8aMwaRJk7Bjxw5ERkbi5ZdfbjM+MzMTNTU1rktJSYk3hk9EPuCtvAEwdxB1Z177SWjp0qX4n//5Hxw4cABXXnml0n0DAwNx7bXXIj8/v83bLRYLLBZLVwyTiAzEm3kDYO4g6s66/BsWIQSWLl2Kd955Bx9++CEGDx6svA6n04mvvvoKMTExXT08IjIg5g0i8qTLv2FZsmQJXn/9dezatQthYWEoKysDAFitVoSE/Hiw34IFC9CvXz9kZ2cDAJ544gn84he/wFVXXYXq6mo899xzOHnyJBYtWtTVwyMiA2LeICJPurxg2bBhAwBg8uTJbss3b96M9PR0AEBxcTFMpp++3Dl37hwWL16MsrIy9OrVC2PHjsXBgwcxcuTIrh4eERkQ8wYRedLlBYsQnk/r2rdvn9v1F198ES+++GJXD4WIugnmDSLyhHMJERERkeGxYCEiIiLDY8FCREREhuf11vyXUlCIHeYennsS9w+rVlqvSaGV8ydfXSW/YotCK3qzYjtph3wLZV2lU3uA2jhMCi3xQ0ObpGMjQzueFK9VfEi9UryssoZw6diTJX2U1m2LkO+vrQcqTK9gl3/f0aWj0hLfptB7XbV9vtNLrfm91WofUJsyRLU1v6aQe5Va4qtMW6L4l1o3yz8jwqzwvYU/t+YnIiIi6gosWIiIiMjwWLAQERGR4bFgISIiIsNjwUJERESGx4KFiIiIDI8FCxERERkeCxYiIiIyPBYsREREZHgsWIiIiMjw/Ko1v9A16Lrn1sHHK/sqrbfulHz7daXW9YEKLdIltstNkPy6TUFO6djgHjalYUSGybfQHxJeKR07IrRUaRyxgdXSsQ26RTr2E3OcdGxZWJh0LABoepB0rFOhNb+pQX4KBADQAuTShKaz5f+FnApN41ViVdr424Vamrcr9Ix3KnzmdSpOEeAtyt3lTQqt+RXyv0q7fT1AbdBCqTW/wrpNbM1PRERE1CEWLERERGR4XV6wZGVlQdM0t8uIESM6vM/bb7+NESNGIDg4GNdccw12797d1cMiIgNj3iAiT7zyDcuoUaNQWlrqunz88cftxh48eBCpqam477778Pnnn2PWrFmYNWsWvv76a28MjYgMinmDiDrilYIlICAA0dHRrkufPn3ajf3jH/+I6dOn4+GHH8bVV1+NJ598Etdddx3WrVvnjaERkUExbxBRR7xSsJw4cQKxsbGIi4vD3XffjeLi4nZjDx06hGnTprktS05OxqFDh9q9T3NzM2pra90uRNS9eTtvAMwdRN1ZlxcsiYmJyMnJwZ49e7BhwwYUFhZiwoQJqKurazO+rKwMUVFRbsuioqJQVlbW7mNkZ2fDarW6Lv379+/SbSCiS+tS5A2AuYOoO+vygiUlJQVz5sxBfHw8kpOTsXv3blRXV+Ott97qssfIzMxETU2N61JSUtJl6yaiS+9S5A2AuYOoO/N647iIiAgMGzYM+fn5bd4eHR2N8vJyt2Xl5eWIjo5ud50WiwUWi3xzLyLqXryRNwDmDqLuzOt9WOrr61FQUICYmJg2b09KSkJubq7bsr179yIpKcnbQyMig2LeIKILdXnBsmzZMuzfvx9FRUU4ePAgZs+eDbPZjNTUVADAggULkJmZ6Yp/6KGHsGfPHjz//PP47rvvkJWVhSNHjmDp0qVdPTQiMijmDSLypMt/Ejp16hRSU1NRVVWFyMhIjB8/HocPH0ZkZCQAoLi4GCbTT3XSjTfeiNdffx2PP/44HnvsMQwdOhQ7d+7E6NGjlR9biB8vntSdtCqt12SXj9WDFVbskJ+zQ2WOCgCAWT4+QGEuoV49zysNY3hEueeg/y8pvEA69rrg9s8gaUukySEdW+KU/8mguPkK6ViTwpwkAKAy9YrZpvj+UKBX18jFCbV5pn7Ol3nDm3SFF1Flzh+b0nw/avO/6ArxDl1+HCrzH/0Y7515a0wmtTmvNIX9VlfIuyo5XSjOJaQyT5EIUHhdTL5tjt/lBcubb77Z4e379u1rtWzOnDmYM2dOVw+FiLoJ5g0i8oRzCREREZHhsWAhIiIiw2PBQkRERIbHgoWIiIgMjwULERERGR4LFiIiIjI8FixERERkeCxYiIiIyPBYsBAREZHheX225ksp6PNQmC2ee+MHBqqtV1eY3FU0KrRxVhiHHqTWet2peadVe89Atfbr0ZZa6di4oDPysYrv3FBTqHRso6iXjjVr8m2+nU61zwcBClNCaLpCm+9gtR1At8kNRBfy0x90V06ZuT9+xqbwmVCl3b5KG3+7wnoBwK7Sbl+hjb9qq31vteZXpdKaX+UrAIWnWanV/o/x8s+dMCs8z5pvXxN+w0JERESGx4KFiIiIDI8FCxERERkeCxYiIiIyPBYsREREZHgsWIiIiMjwWLAQERGR4XV5wTJo0CBomtbqsmTJkjbjc3JyWsUGB3vupUJE/oW5g4g60uWN4z777DM4nU7X9a+//hq33HIL5syZ0+59wsPDcfz4cdd1zcfNaYjo0mPuIKKOdHnBEhkZ6Xb9P/7jPzBkyBBMmjSp3ftomobo6OiuHgoRdSPMHUTUEa+25rfZbNi6dSsyMjI6/ORTX1+PgQMHQtd1XHfddXjmmWcwatSoduObm5vR3Nzsul5b+2P7d2H+8eKJpVqtvbYwKbSfVmiJrNLy3xGi9slRpau1PVj+bVBvD1IaR6NTPr5R4QlpFA1K44DeJB1a4QyRjq2yybf8d9jUWqQHKXS6N9nl39PaebXpFUxBcq38TUIA8k9zhy517pDl9BziRhfyv7rrCr/QOxVi7Yp93VVa+TsVEo1KG39vUh2FUrxCG3+VGRPUW/OrxKq08fftYa9effSdO3eiuroa6enp7cYMHz4cmzZtwq5du7B161bouo4bb7wRp06davc+2dnZsFqtrkv//v29MHoi8hXmDiK6kFcLlj//+c9ISUlBbGxsuzFJSUlYsGABxowZg0mTJmHHjh2IjIzEyy+/3O59MjMzUVNT47qUlJR4Y/hE5CPMHUR0Ia/9JHTy5El88MEH2LFjh9L9AgMDce211yI/P7/dGIvFAotF4fcUIuo2mDuIqC1e+4Zl8+bN6Nu3L26//Xal+zmdTnz11VeIiYnx0siIyMiYO4ioLV4pWHRdx+bNm5GWloaAAPcvcRYsWIDMzEzX9SeeeALvv/8+vv/+exw7dgz33HMPTp48iUWLFnljaERkYMwdRNQer/wk9MEHH6C4uBi//vWvW91WXFwMk+mnOuncuXNYvHgxysrK0KtXL4wdOxYHDx7EyJEjvTE0IjIw5g4iao9XCpZbb70VQrR9ete+ffvcrr/44ot48cUXvTEMIupmmDuIqD2cS4iIiIgMjwULERERGR4LFiIiIjI8FixERERkeF6dS+hSCy9yIiBQdbYPz5TmEpKbdgUA4LDIr9cUqjYDhqbL16JNZvn5fios8nPnAMC/guUnprMGnJeOtalMxAEgWLNLx37b3E86Nr+uj3Ss3qDw5gBgUpnyR2V6LJPa5xQtWK7Rmia0LptLyKgUpmwCADgVZqJReU+rzPejOoePrjA/kENXGIfKBGeK8YovixJNU1i7yiYqrVdxLjmFXVzl75vqOLoav2EhIiIiw2PBQkRERIbHgoWIiIgMjwULERERGR4LFiIiIjI8FixERERkeCxYiIiIyPBYsBAREZHhsWAhIiIiw2PBQkRERIbnV635g885EBDg8BinByi2uVfooKyybmewfL0Y0Kw2ZpNdoRZVaOPfLHoojeM7PUo6trY5WDr229AYpXGEmOVb85c3hUnHFlX2lo4NqFWbTsBsk3/jqbxHodKKGwACJadu0NVWezlwKnwm1BX6qavEOlX6tENxzAq96FVb8wvFeOn1qsZ7aRwqbfxVh6D0kitNJ8DW/EREREQdUi5YDhw4gBkzZiA2NhaapmHnzp1utwshsGLFCsTExCAkJATTpk3DiRMnPK53/fr1GDRoEIKDg5GYmIhPP/1UdWhEZFDMG0TUWcoFS0NDAxISErB+/fo2b1+9ejVeeuklbNy4EZ988gl69uyJ5ORkNDW1P43r9u3bkZGRgZUrV+LYsWNISEhAcnIyzpw5ozo8IjIg5g0i6izlgiUlJQVPPfUUZs+e3eo2IQTWrFmDxx9/HDNnzkR8fDxee+01nD59utUnqp974YUXsHjxYixcuBAjR47Exo0b0aNHD2zatEl1eERkQMwbRNRZXXoMS2FhIcrKyjBt2jTXMqvVisTERBw6dKjN+9hsNhw9etTtPiaTCdOmTWv3Ps3NzaitrXW7EFH3dKnyBsDcQdSddWnBUlZWBgCIinI/MyQqKsp124UqKyvhdDqV7pOdnQ2r1eq69O/fvwtGT0S+cKnyBsDcQdSddcuzhDIzM1FTU+O6lJSU+HpIRNQNMHcQdV9dWrBER0cDAMrLy92Wl5eXu267UJ8+fWA2m5XuY7FYEB4e7nYhou7pUuUNgLmDqDvr0oJl8ODBiI6ORm5urmtZbW0tPvnkEyQlJbV5n6CgIIwdO9btPrquIzc3t937EJH/YN4gIhnKnW7r6+uRn5/vul5YWIi8vDz07t0bAwYMwO9//3s89dRTGDp0KAYPHozly5cjNjYWs2bNct1n6tSpmD17NpYuXQoAyMjIQFpaGq6//nqMGzcOa9asQUNDAxYuXNj5LSQin2PeIKLOUi5Yjhw5gilTpriuZ2RkAADS0tKQk5ODP/zhD2hoaMD999+P6upqjB8/Hnv27EFw8E9t1wsKClBZWem6Pm/ePFRUVGDFihUoKyvDmDFjsGfPnlYH1HkSWNOMAInu53pwoNJ6Nad8z3FhVmhrfV4+1tyk2tZdJVa+3bLZpvalXHNjiHRsSbVkC3gAp3r2UhqHKVD+NdRtCs91rfwuFFyrOL2CQmt+ldbdIlDtvWQKlZuOQdPNQGXbtxk5b6hQnX1AqYW+whfeTpWW+Eq919Vb6Htrvd5q+6/aal+plb9KsOocAQqUpuroRjQhRLfftNraWlitVky59lEEmC0e4w1TsATJxzpC1P7I2EPl45vD5Xfg5l5qO3tzL/m3lz3CKR2r9fQ8Z9TPGaJgOaP2GvYolX/uQk/LPx8hP9QpjcNUd14qzqE344Oidaipqek2x4a05I5z/xeH8DDP+2OBvV5p/SUO+efhB4d8EV7hkJ/v6pyjp3QsAJyzy88XVueQn/+rxiYfCwB1dvn4epv8h536Js9/I37u/Hn5dTvr5f+2mBXmFgs6p/ZBsccZhdzxg3zu6FFwTmkczuP5HmMcwo592CWVN7rlWUJERER0eWHBQkRERIbHgoWIiIgMjwULERERGR4LFiIiIjI8FixERERkeCxYiIiIyPBYsBAREZHhsWAhIiIiw2PBQkRERIanPJeQPzDXNCnewTtza5gU5nQxNSvOJdQs/9IGnJdfd2CDWo1rqZZ/7uyh8mN2hKi9dXX57tpKc3yYFOZsClTr6q40H5TK3CHOULXW5DDJvea687JMJ4bnVJjPCFCc/0hhXh6V9f4Y7535gVTnElKaqEuXj9WUYuWHAMAwcxp1NX7DQkRERIbHgoWIiIgMjwULERERGR4LFiIiIjI8FixERERkeCxYiIiIyPCUC5YDBw5gxowZiI2NhaZp2Llzp+s2u92ORx55BNdccw169uyJ2NhYLFiwAKdPn+5wnVlZWdA0ze0yYsQI5Y0hImNi3iCizlIuWBoaGpCQkID169e3uq2xsRHHjh3D8uXLcezYMezYsQPHjx/HHXfc4XG9o0aNQmlpqevy8ccfqw6NiAyKeYOIOku501NKSgpSUlLavM1qtWLv3r1uy9atW4dx48ahuLgYAwYMaH8gAQGIjo5WHQ4RdQPMG0TUWV4/hqWmpgaapiEiIqLDuBMnTiA2NhZxcXG4++67UVxc3G5sc3Mzamtr3S5E5D+8kTcA5g6i7syrvbSbmprwyCOPIDU1FeHh4e3GJSYmIicnB8OHD0dpaSlWrVqFCRMm4Ouvv0ZYWFir+OzsbKxatarVclOjDSaZNvqaYmtm1bbIkjSbQz62Sa01v+m8/LoDGhRa89epvWWcwfLrdoQotAS3qL2GziCFFuJqT7U0k/xLAgAIbJR/49l7yj93eoBaa/4gk9xz53B0zY7irbwBtJ876OKotNvXobbPqrTmd+oK41BoiQ8AulOlNb98qEq7fZNTPhYANIV4lWk9IJkLvMVr37DY7XbMnTsXQghs2LChw9iUlBTMmTMH8fHxSE5Oxu7du1FdXY233nqrzfjMzEzU1NS4LiUlJd7YBCK6xLyZNwDmDqLuzCvfsLQknZMnT+LDDz/s8FNSWyIiIjBs2DDk5+e3ebvFYoHFojiBGxEZmrfzBsDcQdSddfk3LC1J58SJE/jggw9wxRVXKK+jvr4eBQUFiImJ6erhEZEBMW8QkSfKBUt9fT3y8vKQl5cHACgsLEReXh6Ki4tht9vxy1/+EkeOHMG2bdvgdDpRVlaGsrIy2Gw21zqmTp2KdevWua4vW7YM+/fvR1FREQ4ePIjZs2fDbDYjNTW181tIRD7HvEFEnaX8k9CRI0cwZcoU1/WMjAwAQFpaGrKysvDf//3fAIAxY8a43e+jjz7C5MmTAQAFBQWorKx03Xbq1CmkpqaiqqoKkZGRGD9+PA4fPozIyEjV4RGRATFvEFFnKRcskydPhhDtH1bc0W0tioqK3K6/+eabqsMgom6EeYOIOotzCREREZHhsWAhIiIiw2PBQkRERIbHgoWIiIgMz6ut+S81zWaHZvJcg4kQxcZRNrvCIBRaF0scaNjC1KwwBgBCpYVyg3zdag5Q61svAuXjVWJ1i9pbVw9SaPuvECsCvNeq2t5T/vlwBsqvN6BJsYW+7HtadcoLuqyotNoHAKHUml8h1qn2OV0otObXHAqxKn9WFKf1UJkGRNMVevM7vTRPjSR+w0JERESGx4KFiIiIDI8FCxERERkeCxYiIiIyPBYsREREZHgsWIiIiMjwWLAQERGR4bFgISIiIsNjwUJERESGx4KFiIiIDM+vWvM3DeiNgIBgj3GWH2rUVhwk3/dcO9+stm5ZTqdSuHZerZW/t8hMleCi0PbfFKj41lV4DUWQ/Lr1IPkx23opTgmhMHWDpUahvbYi2dbdSi2+yS+otNtXbc3v0OVzh0qsrtiaH3b5eJNdfhvVYqVDf4x3yu+LmkNhv9XZmp+IiIioQyxYiIiIyPCUC5YDBw5gxowZiI2NhaZp2Llzp9vt6enp0DTN7TJ9+nSP612/fj0GDRqE4OBgJCYm4tNPP1UdGhEZFPMGEXWWcsHS0NCAhIQErF+/vt2Y6dOno7S01HV54403Olzn9u3bkZGRgZUrV+LYsWNISEhAcnIyzpw5ozo8IjIg5g0i6izlg25TUlKQkpLSYYzFYkF0dLT0Ol944QUsXrwYCxcuBABs3LgR7777LjZt2oRHH31UdYhEZDDMG0TUWV45hmXfvn3o27cvhg8fjgceeABVVVXtxtpsNhw9ehTTpk37aVAmE6ZNm4ZDhw61eZ/m5mbU1ta6XYioe/N23gCYO4i6sy4vWKZPn47XXnsNubm5ePbZZ7F//36kpKTA2c5puZWVlXA6nYiKinJbHhUVhbKysjbvk52dDavV6rr079+/qzeDiC6hS5E3AOYOou6sy/uwzJ8/3/X/a665BvHx8RgyZAj27duHqVOndsljZGZmIiMjw3W9traWiYeoG7sUeQNg7iDqzrx+WnNcXBz69OmD/Pz8Nm/v06cPzGYzysvL3ZaXl5e3+3u2xWJBeHi424WI/Ic38gbA3EHUnXm9YDl16hSqqqoQExPT5u1BQUEYO3YscnNzXct0XUdubi6SkpK8PTwiMiDmDSK6kHLBUl9fj7y8POTl5QEACgsLkZeXh+LiYtTX1+Phhx/G4cOHUVRUhNzcXMycORNXXXUVkpOTXeuYOnUq1q1b57qekZGBV199FVu2bMG//vUvPPDAA2hoaHAd/U9E3RvzBhF1lvIxLEeOHMGUKVNc11t+D05LS8OGDRvw5ZdfYsuWLaiurkZsbCxuvfVWPPnkk7BYfppHpaCgAJWVla7r8+bNQ0VFBVasWIGysjKMGTMGe/bsaXVAnSdBVY0IMHuec8fZq4fSes3nGqVjRXCQdKzWcF5+EE61ORyEQ2HuIbtNPtab88WYFOYaMcvP4QMAWpDC66Iw7xAiQqVDzTaF9QIIqVB4zRWeOkew2nNX39/z/FwA4OhgvhMj5w0Vqp/wTJr8a2iGSqz8fmhWGIM3qc4l5FSZp0hpLiG1cUAhXvPS/EAmlfl+AJgcCrF2+feHpvJ3xQs0IRRmWDOo2tpaWK1W3Dz6YQSYPU8wp/eQ/+MFqBUsMMvvOF4tWJoVihAWLO4UChZdoWCx9w6RHwMAhb9JXi1YbFa5eIe9CUf/+jhqamq6zbEhLbnj3P/FITzM875baK9XWn+JU/79cdreSzq2wiH//FY65McAAFU2+fizNvkPf2eb1T4o1jTLFcoAUHdePvZ8g1r+1+vl80FAjfy+FVQjv9NaqtXybkiVfHyPUvkJe4NK2m810BZHUbHnGGHHPuySyhucS4iIiIgMjwULERERGR4LFiIiIjI8FixERERkeCxYiIiIyPBYsBAREZHhsWAhIiIiw2PBQkRERIbHgoWIiIgMT7k1v5GJQDOERAdUzaHWNVallb/5bIN0rFBpAV+r1mETunwLZWGX7+MsHAo9nwGvdcbVFDoKA4Boku/maAoPkx+HTf75CDoj/94AAD1Ifve0XSHf6VOVPUSuI6fTrNjyvBvy5iaqtPE3Cl2hxbJQbc2v0G7f6VRoze9Qyx2aQ6HdvkLneqVYxbSr0spf6e+hYsf1rsZvWIiIiMjwWLAQERGR4bFgISIiIsNjwUJERESGx4KFiIiIDI8FCxERERkeCxYiIiIyPOWC5cCBA5gxYwZiY2OhaRp27tzpdrumaW1ennvuuXbXmZWV1Sp+xIgRyhtDRMbEvEFEnaVcsDQ0NCAhIQHr169v8/bS0lK3y6ZNm6BpGu66664O1ztq1Ci3+3388ceqQyMig2LeIKLOUu50m5KSgpSUlHZvj46Odru+a9cuTJkyBXFxcR0PJCCg1X2JyD8wbxBRZ3m1NX95eTneffddbNmyxWPsiRMnEBsbi+DgYCQlJSE7OxsDBgxoM7a5uRnNzT+1Wq+trQUANEWGICDQc4tyy1n5Nu0AYGq0yQcHeJ4a4GJoQUFK8eJ8k/y6JaYzcK23We25Eyqt+RWmEwDUng+TwvMn7Hb5FVeclQ7VLBb59QJA7BXSoSqtuGsHqu32gefl1q3Zu2YaBm/lDaD93OEtZninlbk32/irtNvXFdrtq6wXAJy6fLzTqbBulViotebXFFKYt2IBwKSwL5pUWvML70y1IsurB91u2bIFYWFhuPPOOzuMS0xMRE5ODvbs2YMNGzagsLAQEyZMQF1dXZvx2dnZsFqtrkv//v29MXwi8gFv5Q2AuYOoO/NqwbJp0ybcfffdCA7u+FuPlJQUzJkzB/Hx8UhOTsbu3btRXV2Nt956q834zMxM1NTUuC4lJSXeGD4R+YC38gbA3EHUnXntJ6H//d//xfHjx7F9+3bl+0ZERGDYsGHIz89v83aLxQKL6tfrRGR43swbAHMHUXfmtW9Y/vznP2Ps2LFISEhQvm99fT0KCgoQExPjhZERkVExbxBRe5QLlvr6euTl5SEvLw8AUFhYiLy8PBQXF7tiamtr8fbbb2PRokVtrmPq1KlYt26d6/qyZcuwf/9+FBUV4eDBg5g9ezbMZjNSU1NVh0dEBsS8QUSdpfyT0JEjRzBlyhTX9YyMDABAWloacnJyAABvvvkmhBDtJo6CggJUVla6rp86dQqpqamoqqpCZGQkxo8fj8OHDyMyMlJ1eERkQMwbRNRZmhA+Pk+pC9TW1sJqtSLp1lU+P61ZUzlFTIFW16gUL2q8c7qm3qg4Di+d1qwFKp7WHOL5feESFKi0blmqpzU7FU5rdvSUH3P1ELVxyJ7W7LQ34ehbj6Ompgbh4eFKj+ErLbnj3P/FITzM8xfOxY56pfWXOHpIx/7g6CUdW+GQf37P2NVeiwpbmHRsZXNP6dhzzfLPBQCcbQyRjq1rkN+/7fWKuaNO/nN9YJ38KdBBNfKxlnNqf6Z7VMjn0pCy89Kx5lL5Ng4A4Dj1g+cYYcc+7JLKG5xLiIiIiAyPBQsREREZHgsWIiIiMjwWLERERGR4Xp1L6FIz2XWYhOeDXu2hagdVmkLkn6agsvbbgl9Ia1KYo0hlfhsAWi+rdKxeXiEdK5xqk1qYFA40FU75OY00s1qt7axvUIqXZQ6VP+hQ9I5QWrdml3+unQrPs8mhNAxA9ni/bn/4vm+ZvfQEmhTXqxovS2XeIQAQKvHeigXU3tdK45APVZ46yk/3RX7DQkRERIbHgoWIiIgMjwULERERGR4LFiIiIjI8FixERERkeCxYiIiIyPBYsBAREZHhsWAhIiIiw2PBQkRERIbnF51uhfixrZ/D0eyV9Zuc8m0DTU75MWi6QqdblVgA0OXHrAv5detCreOuSaHzY8vrKENT7FapC9X2rnKEwnMHhfcGAAiFpsIOh/xnD6dNrVsx7HKvi9PeBEDtdfS1lrHW1su1Eq1zqLUcbVCIb3TIvy7nnfLv52bFLtk2m3y83Sb//nc0qf25cZ6Xf0/rjfLr1c8rto1tku/A7WxS2A8V0oHTprZPORS6ZDucTdKxQlfLYQ6JvxcO/Bgjkzf8omCpq/uxHf4n//sfPh4JtSK/L6hRy8HeU+ul2G6srq4OVqv81BC+1JI7Bl5X5NuBEF3mZPKGJrrTx6F26LqO06dPIywsDJr20yfv2tpa9O/fHyUlJQgPD/fhCL3D37cP8P9t9KftE0Kgrq4OsbGxMJm6x6/NzB3cvu7KX7ZRJW/4xTcsJpMJV155Zbu3h4eHd+sX1BN/3z7A/7fRX7avu3yz0oK5g9vX3fnDNsrmje7xMYiIiIguayxYiIiIyPD8umCxWCxYuXIlLBaLr4fiFf6+fYD/b6O/b1935e+vC7ev+7sctvFCfnHQLREREfk3v/6GhYiIiPwDCxYiIiIyPBYsREREZHgsWIiIiMjwWLAQERGR4fl1wbJ+/XoMGjQIwcHBSExMxKeffurrIXWJrKwsaJrmdhkxYoSvh9UpBw4cwIwZMxAbGwtN07Bz506324UQWLFiBWJiYhASEoJp06bhxIkTvhnsRfC0fenp6a1e0+nTp/tmsJc5f80bgP/lDuaNyytv+G3Bsn37dmRkZGDlypU4duwYEhISkJycjDNnzvh6aF1i1KhRKC0tdV0+/vhjXw+pUxoaGpCQkID169e3efvq1avx0ksvYePGjfjkk0/Qs2dPJCcno6nJW7Mrdi1P2wcA06dPd3tN33jjjUs4QgL8P28A/pU7mDcus7wh/NS4cePEkiVLXNedTqeIjY0V2dnZPhxV11i5cqVISEjw9TC8BoB45513XNd1XRfR0dHiueeecy2rrq4WFotFvPHGGz4YYedcuH1CCJGWliZmzpzpk/HQT/w5bwjh37mDecP/+eU3LDabDUePHsW0adNcy0wmE6ZNm4ZDhw75cGRd58SJE4iNjUVcXBzuvvtuFBcX+3pIXlNYWIiysjK319NqtSIxMdFvXk8A2LdvH/r27Yvhw4fjgQceQFVVla+HdFm5HPIGcPnkDuYN/+OXBUtlZSWcTieioqLclkdFRaGsrMxHo+o6iYmJyMnJwZ49e7BhwwYUFhZiwoQJqKur8/XQvKLlNfPX1xP48Wvd1157Dbm5uXj22Wexf/9+pKSkwOl0+npolw1/zxvA5ZU7mDf8T4CvB0DqUlJSXP+Pj49HYmIiBg4ciLfeegv33XefD0dGF2v+/Pmu/19zzTWIj4/HkCFDsG/fPkydOtWHIyN/wtzhXy63vOGX37D06dMHZrMZ5eXlbsvLy8sRHR3to1F5T0REBIYNG4b8/HxfD8UrWl6zy+X1BIC4uDj06dPHb19TI7rc8gbg37mDecP/+GXBEhQUhLFjxyI3N9e1TNd15ObmIikpyYcj8476+noUFBQgJibG10PxisGDByM6Otrt9aytrcUnn3zil68nAJw6dQpVVVV++5oa0eWWNwD/zh3MG/7Hb38SysjIQFpaGq6//nqMGzcOa9asQUNDAxYuXOjroXXasmXLMGPGDAwcOBCnT5/GypUrYTabkZqa6uuhXbT6+nq3TwWFhYXIy8tD7969MWDAAPz+97/HU089haFDh2Lw4MFYvnw5YmNjMWvWLN8NWkFH29e7d2+sWrUKd911F6Kjo1FQUIA//OEPuOqqq5CcnOzDUV9+/DlvAP6XO5g3LrO84evTlLxp7dq1YsCAASIoKEiMGzdOHD582NdD6hLz5s0TMTExIigoSPTr10/MmzdP5Ofn+3pYnfLRRx8JAK0uaWlpQogfT1Fcvny5iIqKEhaLRUydOlUcP37ct4NW0NH2NTY2iltvvVVERkaKwMBAMXDgQLF48WJRVlbm62Fflvw1bwjhf7mDeePyyhuaEEJc2hKJiIiISI1fHsNCRERE/oUFCxERERkeCxYiIiIyPBYsREREZHgsWIiIiMjwWLAQERGR4bFgISIiIsNjwUJERESGx4KFiIiIDI8FCxERERkeCxYiIiIyvP8HXODpCG4iMjAAAAAASUVORK5CYII="
      },
      "metadata": {},
      "output_type": "display_data"
     }
    ],
+   "execution_count": 3
+  },
+  {
+   "cell_type": "code",
+   "id": "082ab7a8-22e0-498b-b138-158dc9f2658f",
+   "metadata": {
+    "ExecuteTime": {
+     "end_time": "2024-09-19T13:35:29.122858Z",
+     "start_time": "2024-09-19T13:35:29.119985Z"
+    }
+   },
    "source": [
-    "plt.subplot(1, 2, 1)\n",
-    "plt.title('permeability')\n",
-    "plt.imshow(k_train.squeeze(-1)[0])\n",
-    "plt.subplot(1, 2, 2)\n",
-    "plt.title('field solution')\n",
-    "plt.imshow(u_train.squeeze(-1)[0])\n",
-    "plt.show()"
-   ]
+    "u_train.labels[3]['dof']"
+   ],
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "['u']"
+      ]
+     },
+     "execution_count": 4,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "execution_count": 4
   },
   {
    "cell_type": "markdown",
@@ -131,20 +174,36 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 17,
    "id": "8b27d283",
-   "metadata": {},
-   "outputs": [],
+   "metadata": {
+    "ExecuteTime": {
+     "end_time": "2024-09-19T13:35:29.136572Z",
+     "start_time": "2024-09-19T13:35:29.134124Z"
+    }
+   },
    "source": [
     "class NeuralOperatorSolver(AbstractProblem):\n",
-    "    input_variables = k_train.labels\n",
-    "    output_variables = u_train.labels\n",
-    "    conditions = {'data' : Condition(input_points=k_train, \n",
+    "    input_variables = k_train.labels[3]['dof']\n",
+    "    output_variables = u_train.labels[3]['dof']\n",
+    "    domains = {\n",
+    "        'pts': k_train\n",
+    "    }\n",
+    "    conditions = {'data' : Condition(domain='pts', \n",
     "                                     output_points=u_train)}\n",
     "\n",
     "# make problem\n",
     "problem = NeuralOperatorSolver()"
-   ]
+   ],
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "<pina.condition.domain_output_condition.DomainOutputCondition object at 0x1259893b0>\n"
+     ]
+    }
+   ],
+   "execution_count": 5
   },
   {
    "cell_type": "markdown",
@@ -158,25 +217,42 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 18,
    "id": "e34f18b0",
-   "metadata": {},
+   "metadata": {
+    "ExecuteTime": {
+     "end_time": "2024-09-19T13:35:31.245429Z",
+     "start_time": "2024-09-19T13:35:29.154937Z"
+    }
+   },
+   "source": [
+    "# make model\n",
+    "model = FeedForward(input_dimensions=1, output_dimensions=1)\n",
+    "\n",
+    "\n",
+    "# make solver\n",
+    "solver = SupervisedSolver(problem=problem, model=model)\n",
+    "\n",
+    "# make the trainer and train\n",
+    "trainer = Trainer(solver=solver, max_epochs=10, accelerator='cpu', enable_model_summary=False, batch_size=10) \n",
+    "# We train on CPU and avoid model summary at the beginning of training (optional)\n",
+    "trainer.train()"
+   ],
    "outputs": [
     {
      "name": "stderr",
      "output_type": "stream",
      "text": [
-      "GPU available: False, used: False\n",
+      "GPU available: True (mps), used: False\n",
       "TPU available: False, using: 0 TPU cores\n",
-      "IPU available: False, using: 0 IPUs\n",
-      "HPU available: False, using: 0 HPUs\n"
+      "HPU available: False, using: 0 HPUs\n",
+      "/Users/filippoolivo/miniconda3/envs/PINAv0.2/lib/python3.11/site-packages/pytorch_lightning/trainer/setup.py:177: GPU available but not used. You can set it by doing `Trainer(accelerator='gpu')`.\n"
      ]
     },
     {
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "Epoch 9: : 100it [00:00, 357.28it/s, v_num=1, mean_loss=0.108]"
+      "Epoch 9: 100%|██████████| 100/100 [00:00<00:00, 552.80it/s, v_num=18, mean_loss=0.113]"
      ]
     },
     {
@@ -190,22 +266,11 @@
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "Epoch 9: : 100it [00:00, 354.81it/s, v_num=1, mean_loss=0.108]\n"
+      "Epoch 9: 100%|██████████| 100/100 [00:00<00:00, 547.37it/s, v_num=18, mean_loss=0.113]\n"
      ]
     }
    ],
-   "source": [
-    "# make model\n",
-    "model = FeedForward(input_dimensions=1, output_dimensions=1)\n",
-    "\n",
-    "\n",
-    "# make solver\n",
-    "solver = SupervisedSolver(problem=problem, model=model)\n",
-    "\n",
-    "# make the trainer and train\n",
-    "trainer = Trainer(solver=solver, max_epochs=10, accelerator='cpu', enable_model_summary=False, batch_size=10) # we train on CPU and avoid model summary at beginning of training (optional)\n",
-    "trainer.train()\n"
-   ]
+   "execution_count": 6
   },
   {
    "cell_type": "markdown",
@@ -217,32 +282,37 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 19,
    "id": "0e2a6aa4",
-   "metadata": {},
-   "outputs": [
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "Final error training 56.04%\n",
-      "Final error testing 56.01%\n"
-     ]
+   "metadata": {
+    "ExecuteTime": {
+     "end_time": "2024-09-19T13:35:31.295336Z",
+     "start_time": "2024-09-19T13:35:31.256308Z"
     }
-   ],
+   },
    "source": [
     "from pina.loss import LpLoss\n",
     "\n",
     "# make the metric\n",
     "metric_err = LpLoss(relative=True)\n",
     "\n",
-    "\n",
-    "err = float(metric_err(u_train.squeeze(-1), solver.neural_net(k_train).squeeze(-1)).mean())*100\n",
+    "model = solver.models[0]\n",
+    "err = float(metric_err(u_train.squeeze(-1), model(k_train).squeeze(-1)).mean())*100\n",
     "print(f'Final error training {err:.2f}%')\n",
     "\n",
-    "err = float(metric_err(u_test.squeeze(-1), solver.neural_net(k_test).squeeze(-1)).mean())*100\n",
+    "err = float(metric_err(u_test.squeeze(-1), model(k_test).squeeze(-1)).mean())*100\n",
     "print(f'Final error testing {err:.2f}%')"
-   ]
+   ],
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Final error training 56.05%\n",
+      "Final error testing 55.95%\n"
+     ]
+    }
+   ],
+   "execution_count": 7
   },
   {
    "cell_type": "markdown",
@@ -256,17 +326,39 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 24,
    "id": "9af523a5",
-   "metadata": {},
+   "metadata": {
+    "ExecuteTime": {
+     "end_time": "2024-09-19T13:35:44.717807Z",
+     "start_time": "2024-09-19T13:35:31.306689Z"
+    }
+   },
+   "source": [
+    "# make model\n",
+    "lifting_net = torch.nn.Linear(1, 24)\n",
+    "projecting_net = torch.nn.Linear(24, 1)\n",
+    "model = FNO(lifting_net=lifting_net,\n",
+    "            projecting_net=projecting_net,\n",
+    "            n_modes=8,\n",
+    "            dimensions=2,\n",
+    "            inner_size=24,\n",
+    "            padding=8)\n",
+    "\n",
+    "\n",
+    "# make solver\n",
+    "solver = SupervisedSolver(problem=problem, model=model)\n",
+    "\n",
+    "# make the trainer and train\n",
+    "trainer = Trainer(solver=solver, max_epochs=10, accelerator='cpu', enable_model_summary=False, batch_size=10) # we train on CPU and avoid model summary at beginning of training (optional)\n",
+    "trainer.train()"
+   ],
    "outputs": [
     {
      "name": "stderr",
      "output_type": "stream",
      "text": [
-      "GPU available: False, used: False\n",
+      "GPU available: True (mps), used: False\n",
       "TPU available: False, using: 0 TPU cores\n",
-      "IPU available: False, using: 0 IPUs\n",
       "HPU available: False, using: 0 HPUs\n"
      ]
     },
@@ -274,14 +366,7 @@
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "Epoch 0: : 0it [00:00, ?it/s]"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "Epoch 9: : 100it [00:02, 47.76it/s, v_num=4, mean_loss=0.00106] "
+      "Epoch 9: 100%|██████████| 100/100 [00:01<00:00, 73.04it/s, v_num=19, mean_loss=0.00215]"
      ]
     },
     {
@@ -295,29 +380,11 @@
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "Epoch 9: : 100it [00:02, 47.65it/s, v_num=4, mean_loss=0.00106]\n"
+      "Epoch 9: 100%|██████████| 100/100 [00:01<00:00, 72.84it/s, v_num=19, mean_loss=0.00215]\n"
      ]
     }
    ],
-   "source": [
-    "# make model\n",
-    "lifting_net = torch.nn.Linear(1, 24)\n",
-    "projecting_net = torch.nn.Linear(24, 1)\n",
-    "model = FNO(lifting_net=lifting_net,\n",
-    "            projecting_net=projecting_net,\n",
-    "            n_modes=8,\n",
-    "            dimensions=2,\n",
-    "            inner_size=24,\n",
-    "            padding=8)\n",
-    "\n",
-    "\n",
-    "# make solver\n",
-    "solver = SupervisedSolver(problem=problem, model=model)\n",
-    "\n",
-    "# make the trainer and train\n",
-    "trainer = Trainer(solver=solver, max_epochs=10, accelerator='cpu', enable_model_summary=False, batch_size=10) # we train on CPU and avoid model summary at beginning of training (optional)\n",
-    "trainer.train()\n"
-   ]
+   "execution_count": 8
   },
   {
    "cell_type": "markdown",
@@ -329,26 +396,33 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 25,
    "id": "58e2db89",
-   "metadata": {},
+   "metadata": {
+    "ExecuteTime": {
+     "end_time": "2024-09-19T13:35:45.259819Z",
+     "start_time": "2024-09-19T13:35:44.729042Z"
+    }
+   },
+   "source": [
+    "model = solver.models[0]\n",
+    "\n",
+    "err = float(metric_err(u_train.squeeze(-1), model(k_train).squeeze(-1)).mean())*100\n",
+    "print(f'Final error training {err:.2f}%')\n",
+    "\n",
+    "err = float(metric_err(u_test.squeeze(-1), model(k_test).squeeze(-1)).mean())*100\n",
+    "print(f'Final error testing {err:.2f}%')"
+   ],
    "outputs": [
     {
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "Final error training 4.83%\n",
-      "Final error testing 5.16%\n"
+      "Final error training 7.48%\n",
+      "Final error testing 7.73%\n"
      ]
     }
    ],
-   "source": [
-    "err = float(metric_err(u_train.squeeze(-1), solver.neural_net(k_train).squeeze(-1)).mean())*100\n",
-    "print(f'Final error training {err:.2f}%')\n",
-    "\n",
-    "err = float(metric_err(u_test.squeeze(-1), solver.neural_net(k_test).squeeze(-1)).mean())*100\n",
-    "print(f'Final error testing {err:.2f}%')"
-   ]
+   "execution_count": 9
   },
   {
    "cell_type": "markdown",
@@ -367,6 +441,19 @@
     "\n",
     "We have made a very simple example on how to use the `FNO` for learning neural operator. Currently in **PINA** we implement 1D/2D/3D cases. We suggest to extend the tutorial using more complex problems and train for longer, to see the full potential of neural operators."
    ]
+  },
+  {
+   "metadata": {
+    "ExecuteTime": {
+     "end_time": "2024-09-19T13:08:35.195331Z",
+     "start_time": "2024-09-19T13:08:35.193830Z"
+    }
+   },
+   "cell_type": "code",
+   "source": "",
+   "id": "af932a706dd1b71f",
+   "outputs": [],
+   "execution_count": null
   }
  ],
  "metadata": {
@@ -388,7 +475,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.9.16"
+   "version": "3.11.9"
   }
  },
  "nbformat": 4,
diff --git a/tutorials/tutorial5/tutorial.py b/tutorials/tutorial5/tutorial.py
index 4e5fa13e0..add791430 100644
--- a/tutorials/tutorial5/tutorial.py
+++ b/tutorials/tutorial5/tutorial.py
@@ -48,24 +48,28 @@
 # Specifically, $u$ is the flow pressure, $k$ is the permeability field and $f$ is the forcing function. The Darcy flow can parameterize a variety of systems including flow through porous media, elastic materials and heat conduction. Here you will define the domain as a 2D unit square Dirichlet boundary conditions. The dataset is taken from the authors original reference.
 # 
 
-# In[12]:
+# In[2]:
 
 
 # download the dataset
 data = io.loadmat("Data_Darcy.mat")
 
 # extract data (we use only 100 data for train)
-k_train = LabelTensor(torch.tensor(data['k_train'], dtype=torch.float).unsqueeze(-1), ['u0'])
-u_train = LabelTensor(torch.tensor(data['u_train'], dtype=torch.float).unsqueeze(-1), ['u'])
-k_test =  LabelTensor(torch.tensor(data['k_test'], dtype=torch.float).unsqueeze(-1), ['u0'])
-u_test= LabelTensor(torch.tensor(data['u_test'], dtype=torch.float).unsqueeze(-1), ['u'])
+k_train = LabelTensor(torch.tensor(data['k_train'], dtype=torch.float).unsqueeze(-1), 
+                      labels={3:{'dof': ['u0'], 'name': 'k_train'}})
+u_train = LabelTensor(torch.tensor(data['u_train'], dtype=torch.float).unsqueeze(-1),
+                     labels={3:{'dof': ['u'], 'name': 'u_train'}})
+k_test =  LabelTensor(torch.tensor(data['k_test'], dtype=torch.float).unsqueeze(-1),
+                      labels={3:{'dof': ['u0'], 'name': 'k_test'}})
+u_test= LabelTensor(torch.tensor(data['u_test'], dtype=torch.float).unsqueeze(-1),
+                   labels={3:{'dof': ['u'], 'name': 'u_test'}})
 x = torch.tensor(data['x'], dtype=torch.float)[0]
 y = torch.tensor(data['y'], dtype=torch.float)[0]
 
 
 # Let's visualize some data
 
-# In[13]:
+# In[3]:
 
 
 plt.subplot(1, 2, 1)
@@ -77,15 +81,24 @@
 plt.show()
 
 
+# In[4]:
+
+
+u_train.labels[3]['dof']
+
+
 # We now create the neural operator class. It is a very simple class, inheriting from `AbstractProblem`.
 
-# In[17]:
+# In[5]:
 
 
 class NeuralOperatorSolver(AbstractProblem):
-    input_variables = k_train.labels
-    output_variables = u_train.labels
-    conditions = {'data' : Condition(input_points=k_train, 
+    input_variables = k_train.labels[3]['dof']
+    output_variables = u_train.labels[3]['dof']
+    domains = {
+        'pts': k_train
+    }
+    conditions = {'data' : Condition(domain='pts', 
                                      output_points=u_train)}
 
 # make problem
@@ -96,7 +109,7 @@ class NeuralOperatorSolver(AbstractProblem):
 # 
 # We will first solve the problem using a Feedforward neural network. We will use the `SupervisedSolver` for solving the problem, since we are training using supervised learning.
 
-# In[18]:
+# In[6]:
 
 
 # make model
@@ -107,25 +120,26 @@ class NeuralOperatorSolver(AbstractProblem):
 solver = SupervisedSolver(problem=problem, model=model)
 
 # make the trainer and train
-trainer = Trainer(solver=solver, max_epochs=10, accelerator='cpu', enable_model_summary=False, batch_size=10) # we train on CPU and avoid model summary at beginning of training (optional)
+trainer = Trainer(solver=solver, max_epochs=10, accelerator='cpu', enable_model_summary=False, batch_size=10) 
+# We train on CPU and avoid model summary at the beginning of training (optional)
 trainer.train()
 
 
 # The final loss is pretty high... We can calculate the error by importing `LpLoss`.
 
-# In[19]:
+# In[7]:
 
 
-from pina.loss.loss_interface import LpLoss
+from pina.loss import LpLoss
 
 # make the metric
 metric_err = LpLoss(relative=True)
 
-
-err = float(metric_err(u_train.squeeze(-1), solver.neural_net(k_train).squeeze(-1)).mean())*100
+model = solver.models[0]
+err = float(metric_err(u_train.squeeze(-1), model(k_train).squeeze(-1)).mean())*100
 print(f'Final error training {err:.2f}%')
 
-err = float(metric_err(u_test.squeeze(-1), solver.neural_net(k_test).squeeze(-1)).mean())*100
+err = float(metric_err(u_test.squeeze(-1), model(k_test).squeeze(-1)).mean())*100
 print(f'Final error testing {err:.2f}%')
 
 
@@ -133,7 +147,7 @@ class NeuralOperatorSolver(AbstractProblem):
 # 
 # We will now move to solve the problem using a FNO. Since we are learning operator this approach is better suited, as we shall see.
 
-# In[24]:
+# In[8]:
 
 
 # make model
@@ -157,13 +171,15 @@ class NeuralOperatorSolver(AbstractProblem):
 
 # We can clearly see that the final loss is lower. Let's see in testing.. Notice that the number of parameters is way higher than a `FeedForward` network. We suggest to use GPU or TPU for a speed up in training, when many data samples are used.
 
-# In[25]:
+# In[9]:
 
 
-err = float(metric_err(u_train.squeeze(-1), solver.neural_net(k_train).squeeze(-1)).mean())*100
+model = solver.models[0]
+
+err = float(metric_err(u_train.squeeze(-1), model(k_train).squeeze(-1)).mean())*100
 print(f'Final error training {err:.2f}%')
 
-err = float(metric_err(u_test.squeeze(-1), solver.neural_net(k_test).squeeze(-1)).mean())*100
+err = float(metric_err(u_test.squeeze(-1), model(k_test).squeeze(-1)).mean())*100
 print(f'Final error testing {err:.2f}%')
 
 
@@ -172,3 +188,9 @@ class NeuralOperatorSolver(AbstractProblem):
 # ## What's next?
 # 
 # We have made a very simple example on how to use the `FNO` for learning neural operator. Currently in **PINA** we implement 1D/2D/3D cases. We suggest to extend the tutorial using more complex problems and train for longer, to see the full potential of neural operators.
+
+# In[ ]:
+
+
+
+

From e00cc5afe2d786ddeb65c6ea3389a3604dc8dad9 Mon Sep 17 00:00:00 2001
From: FilippoOlivo <filippo@filippoolivo.com>
Date: Sat, 28 Sep 2024 12:19:53 +0200
Subject: [PATCH 11/55] Delete useless classes (move into optim folder)

---
 pina/optimizer.py | 21 ---------------------
 pina/scheduler.py | 29 -----------------------------
 2 files changed, 50 deletions(-)
 delete mode 100644 pina/optimizer.py
 delete mode 100644 pina/scheduler.py

diff --git a/pina/optimizer.py b/pina/optimizer.py
deleted file mode 100644
index 08631e625..000000000
--- a/pina/optimizer.py
+++ /dev/null
@@ -1,21 +0,0 @@
-""" Module for PINA Optimizer """
-
-import torch
-from .utils import check_consistency
-
-class Optimizer: # TODO improve interface
-    pass
-
-
-class TorchOptimizer(Optimizer):
-
-    def __init__(self, optimizer_class, **kwargs):
-        check_consistency(optimizer_class, torch.optim.Optimizer, subclass=True)
-
-        self.optimizer_class = optimizer_class
-        self.kwargs = kwargs
-
-    def hook(self, parameters):
-        self.optimizer_instance = self.optimizer_class(
-            parameters, **self.kwargs
-        )
\ No newline at end of file
diff --git a/pina/scheduler.py b/pina/scheduler.py
deleted file mode 100644
index 563f829c5..000000000
--- a/pina/scheduler.py
+++ /dev/null
@@ -1,29 +0,0 @@
-""" Module for PINA Scheduler """
-
-try:
-    from torch.optim.lr_scheduler import LRScheduler  # torch >= 2.0
-except ImportError:
-    from torch.optim.lr_scheduler import (
-        _LRScheduler as LRScheduler,
-    )  # torch < 2.0
-from .optimizer import Optimizer
-from .utils import check_consistency
-
-
-class Scheduler: # TODO improve interface
-    pass
-
-
-class TorchScheduler(Scheduler):
-
-    def __init__(self, scheduler_class, **kwargs):
-        check_consistency(scheduler_class, LRScheduler, subclass=True)
-
-        self.scheduler_class = scheduler_class
-        self.kwargs = kwargs
-
-    def hook(self, optimizer):
-        check_consistency(optimizer, Optimizer)
-        self.scheduler_instance = self.scheduler_class(
-            optimizer.optimizer_instance, **self.kwargs
-        )

From 12adf459017303995c6b7a613c7971c05c611e81 Mon Sep 17 00:00:00 2001
From: FilippoOlivo <filippo@filippoolivo.com>
Date: Sat, 28 Sep 2024 12:22:03 +0200
Subject: [PATCH 12/55] Update __init__.py in pina folder

---
 pina/__init__.py | 4 ----
 1 file changed, 4 deletions(-)

diff --git a/pina/__init__.py b/pina/__init__.py
index 835232b12..0fe93752d 100644
--- a/pina/__init__.py
+++ b/pina/__init__.py
@@ -6,8 +6,6 @@
     "Condition",
     "SamplePointDataset",
     "SamplePointLoader",
-    "TorchOptimizer",
-    "TorchScheduler",
 ]
 
 from .meta import *
@@ -15,8 +13,6 @@
 from .solvers.solver import SolverInterface
 from .trainer import Trainer
 from .plotter import Plotter
-from .optimizer import TorchOptimizer
-from .scheduler import TorchScheduler
 from .condition.condition import Condition
 from .data import SamplePointDataset
 from .data import SamplePointLoader

From 5431f8b09ae68e0c33d528a075de5d354b4275a1 Mon Sep 17 00:00:00 2001
From: FilippoOlivo <filippo@filippoolivo.com>
Date: Sat, 28 Sep 2024 12:23:16 +0200
Subject: [PATCH 13/55] Implement definition of LabelTensor from list,
 implement cat method (previously stack) and re-implement extract

---
 pina/label_tensor.py       | 469 ++++++++-----------------------------
 tests/test_label_tensor.py | 109 +--------
 2 files changed, 105 insertions(+), 473 deletions(-)

diff --git a/pina/label_tensor.py b/pina/label_tensor.py
index ab6045e3f..f40c36dca 100644
--- a/pina/label_tensor.py
+++ b/pina/label_tensor.py
@@ -1,311 +1,8 @@
 """ Module for LabelTensor """
 
-from copy import deepcopy
 import torch
 from torch import Tensor
 
-
-# class LabelTensor(torch.Tensor):
-#     """Torch tensor with a label for any column."""
-
-#     @staticmethod
-#     def __new__(cls, x, labels, *args, **kwargs):
-#         return super().__new__(cls, x, *args, **kwargs)
-
-#     def __init__(self, x, labels):
-#         """
-#         Construct a `LabelTensor` by passing a tensor and a list of column
-#         labels. Such labels uniquely identify the columns of the tensor,
-#         allowing for an easier manipulation.
-
-#         :param torch.Tensor x: The data tensor.
-#         :param labels: The labels of the columns.
-#         :type labels: str | list(str) | tuple(str)
-
-#         :Example:
-#             >>> from pina import LabelTensor
-#             >>> tensor = LabelTensor(torch.rand((2000, 3)), ['a', 'b', 'c'])
-#             >>> tensor
-#             tensor([[6.7116e-02, 4.8892e-01, 8.9452e-01],
-#                 [9.2392e-01, 8.2065e-01, 4.1986e-04],
-#                 [8.9266e-01, 5.5446e-01, 6.3500e-01],
-#                 ...,
-#                 [5.8194e-01, 9.4268e-01, 4.1841e-01],
-#                 [1.0246e-01, 9.5179e-01, 3.7043e-02],
-#                 [9.6150e-01, 8.0656e-01, 8.3824e-01]])
-#             >>> tensor.extract('a')
-#             tensor([[0.0671],
-#                     [0.9239],
-#                     [0.8927],
-#                     ...,
-#                     [0.5819],
-#                     [0.1025],
-#                     [0.9615]])
-#             >>> tensor['a']
-#             tensor([[0.0671],
-#                     [0.9239],
-#                     [0.8927],
-#                     ...,
-#                     [0.5819],
-#                     [0.1025],
-#                     [0.9615]])
-#             >>> tensor.extract(['a', 'b'])
-#             tensor([[0.0671, 0.4889],
-#                     [0.9239, 0.8207],
-#                     [0.8927, 0.5545],
-#                     ...,
-#                     [0.5819, 0.9427],
-#                     [0.1025, 0.9518],
-#                     [0.9615, 0.8066]])
-#             >>> tensor.extract(['b', 'a'])
-#             tensor([[0.4889, 0.0671],
-#                     [0.8207, 0.9239],
-#                     [0.5545, 0.8927],
-#                     ...,
-#                     [0.9427, 0.5819],
-#                     [0.9518, 0.1025],
-#                     [0.8066, 0.9615]])
-#         """
-#         if x.ndim == 1:
-#             x = x.reshape(-1, 1)
-
-#         if isinstance(labels, str):
-#             labels = [labels]
-
-#         if len(labels) != x.shape[-1]:
-#             raise ValueError(
-#                 "the tensor has not the same number of columns of "
-#                 "the passed labels."
-#             )
-#         self._labels = labels
-
-#     def __deepcopy__(self, __):
-#         """
-#         Implements deepcopy for label tensor. By default it stores the
-#         current labels and use the :meth:`~torch._tensor.Tensor.__deepcopy__`
-#         method for creating a new :class:`pina.label_tensor.LabelTensor`.
-
-#         :param __: Placeholder parameter.
-#         :type __: None
-#         :return: The deep copy of the :class:`pina.label_tensor.LabelTensor`.
-#         :rtype: LabelTensor
-#         """
-#         labels = self.labels
-#         copy_tensor = deepcopy(self.tensor)
-#         return LabelTensor(copy_tensor, labels)
-
-#     @property
-#     def labels(self):
-#         """Property decorator for labels
-
-#         :return: labels of self
-#         :rtype: list
-#         """
-#         return self._labels
-
-#     @labels.setter
-#     def labels(self, labels):
-#         if len(labels) != self.shape[self.ndim - 1]:  # small check
-#             raise ValueError(
-#                 "The tensor has not the same number of columns of "
-#                 "the passed labels."
-#             )
-
-#         self._labels = labels  # assign the label
-
-#     @staticmethod
-#     def vstack(label_tensors):
-#         """
-#         Stack tensors vertically. For more details, see
-#         :meth:`torch.vstack`.
-
-#         :param list(LabelTensor) label_tensors: the tensors to stack. They need
-#             to have equal labels.
-#         :return: the stacked tensor
-#         :rtype: LabelTensor
-#         """
-#         if len(label_tensors) == 0:
-#             return []
-
-#         all_labels = [label for lt in label_tensors for label in lt.labels]
-#         if set(all_labels) != set(label_tensors[0].labels):
-#             raise RuntimeError("The tensors to stack have different labels")
-
-#         labels = label_tensors[0].labels
-#         tensors = [lt.extract(labels) for lt in label_tensors]
-#         return LabelTensor(torch.vstack(tensors), labels)
-
-#     def clone(self, *args, **kwargs):
-#         """
-#         Clone the LabelTensor. For more details, see
-#         :meth:`torch.Tensor.clone`.
-
-#         :return: A copy of the tensor.
-#         :rtype: LabelTensor
-#         """
-#         # # used before merging
-#         # try:
-#         #     out = LabelTensor(super().clone(*args, **kwargs), self.labels)
-#         # except:
-#         #     out = super().clone(*args, **kwargs)
-#         out = LabelTensor(super().clone(*args, **kwargs), self.labels)
-#         return out
-
-#     def to(self, *args, **kwargs):
-#         """
-#         Performs Tensor dtype and/or device conversion. For more details, see
-#         :meth:`torch.Tensor.to`.
-#         """
-#         tmp = super().to(*args, **kwargs)
-#         new = self.__class__.clone(self)
-#         new.data = tmp.data
-#         return new
-
-#     def select(self, *args, **kwargs):
-#         """
-#         Performs Tensor selection. For more details, see :meth:`torch.Tensor.select`.
-#         """
-#         tmp = super().select(*args, **kwargs)
-#         tmp._labels = self._labels
-#         return tmp
-
-#     def cuda(self, *args, **kwargs):
-#         """
-#         Send Tensor to cuda. For more details, see :meth:`torch.Tensor.cuda`.
-#         """
-#         tmp = super().cuda(*args, **kwargs)
-#         new = self.__class__.clone(self)
-#         new.data = tmp.data
-#         return new
-
-#     def cpu(self, *args, **kwargs):
-#         """
-#         Send Tensor to cpu. For more details, see :meth:`torch.Tensor.cpu`.
-#         """
-#         tmp = super().cpu(*args, **kwargs)
-#         new = self.__class__.clone(self)
-#         new.data = tmp.data
-#         return new
-
-#     def extract(self, label_to_extract):
-#         """
-#         Extract the subset of the original tensor by returning all the columns
-#         corresponding to the passed ``label_to_extract``.
-
-#         :param label_to_extract: The label(s) to extract.
-#         :type label_to_extract: str | list(str) | tuple(str)
-#         :raises TypeError: Labels are not ``str``.
-#         :raises ValueError: Label to extract is not in the labels ``list``.
-#         """
-
-#         if isinstance(label_to_extract, str):
-#             label_to_extract = [label_to_extract]
-#         elif isinstance(label_to_extract, (tuple, list)):  # TODO
-#             pass
-#         else:
-#             raise TypeError(
-#                 "`label_to_extract` should be a str, or a str iterator"
-#             )
-
-#         indeces = []
-#         for f in label_to_extract:
-#             try:
-#                 indeces.append(self.labels.index(f))
-#             except ValueError:
-#                 raise ValueError(f"`{f}` not in the labels list")
-
-#         new_data = super(Tensor, self.T).__getitem__(indeces).T
-#         new_labels = [self.labels[idx] for idx in indeces]
-
-#         extracted_tensor = new_data.as_subclass(LabelTensor)
-#         extracted_tensor.labels = new_labels
-
-#         return extracted_tensor
-
-#     def detach(self):
-#         detached = super().detach()
-#         if hasattr(self, "_labels"):
-#             detached._labels = self._labels
-#         return detached
-
-
-#     def append(self, lt, mode="std"):
-#         """
-#         Return a copy of the merged tensors.
-
-#         :param LabelTensor lt: The tensor to merge.
-#         :param str mode: {'std', 'first', 'cross'}
-#         :return: The merged tensors.
-#         :rtype: LabelTensor
-#         """
-#         if set(self.labels).intersection(lt.labels):
-#             raise RuntimeError("The tensors to merge have common labels")
-
-#         new_labels = self.labels + lt.labels
-#         if mode == "std":
-#             new_tensor = torch.cat((self, lt), dim=1)
-#         elif mode == "first":
-#             raise NotImplementedError
-#         elif mode == "cross":
-#             tensor1 = self
-#             tensor2 = lt
-#             n1 = tensor1.shape[0]
-#             n2 = tensor2.shape[0]
-
-#             tensor1 = LabelTensor(tensor1.repeat(n2, 1), labels=tensor1.labels)
-#             tensor2 = LabelTensor(
-#                 tensor2.repeat_interleave(n1, dim=0), labels=tensor2.labels
-#             )
-#             new_tensor = torch.cat((tensor1, tensor2), dim=1)
-
-#         new_tensor = new_tensor.as_subclass(LabelTensor)
-#         new_tensor.labels = new_labels
-#         return new_tensor
-
-#     def __getitem__(self, index):
-#         """
-#         Return a copy of the selected tensor.
-#         """
-
-#         if isinstance(index, str) or (
-#             isinstance(index, (tuple, list))
-#             and all(isinstance(a, str) for a in index)
-#         ):
-#             return self.extract(index)
-
-#         selected_lt = super(Tensor, self).__getitem__(index)
-
-#         try:
-#             len_index = len(index)
-#         except TypeError:
-#             len_index = 1
-
-#         if isinstance(index, int) or len_index == 1:
-#             if selected_lt.ndim == 1:
-#                 selected_lt = selected_lt.reshape(1, -1)
-#             if hasattr(self, "labels"):
-#                 selected_lt.labels = self.labels
-#         elif len_index == 2:
-#             if selected_lt.ndim == 1:
-#                 selected_lt = selected_lt.reshape(-1, 1)
-#             if hasattr(self, "labels"):
-#                 if isinstance(index[1], list):
-#                     selected_lt.labels = [self.labels[i] for i in index[1]]
-#                 else:
-#                     selected_lt.labels = self.labels[index[1]]
-#         else:
-#             selected_lt.labels = self.labels
-
-#         return selected_lt
-
-
-#     def __str__(self):
-#         if hasattr(self, "labels"):
-#             s = f"labels({str(self.labels)})\n"
-#         else:
-#             s = "no labels\n"
-#         s += super().__str__()
-#         return s
 def issubset(a, b):
     """
     Check if a is a subset of b.
@@ -334,21 +31,19 @@ def __init__(self, x, labels):
         :Example:
             >>> from pina import LabelTensor
             >>> tensor = LabelTensor(
-            >>>     torch.rand((2000, 3)), 
+            >>>     torch.rand((2000, 3)),
                     {1: {"name": "space"['a', 'b', 'c'])
-        
+
         """
         from .utils import check_consistency
-        check_consistency(labels, dict)
-
-        self.labels = {
-            idx_: {
-                'dof': range(x.shape[idx_]),
-                'name': idx_
-            } for idx_ in range(x.ndim)
-        }
-        self.labels.update(labels)
-
+        if isinstance(labels, dict):
+            # check_consistency(labels, dict)
+            self.update_labels(labels)
+        elif isinstance(labels, list):
+            self.init_labels_from_list(labels)
+        elif isinstance(labels, str):
+            labels = [labels]
+            raise ValueError(f"labels must be list, dict or string.")
 
     def extract(self, label_to_extract):
         """
@@ -360,47 +55,48 @@ def extract(self, label_to_extract):
         :raises TypeError: Labels are not ``str``.
         :raises ValueError: Label to extract is not in the labels ``list``.
         """
-        if isinstance(label_to_extract, (int, str)):
+        from copy import deepcopy
+        if isinstance(label_to_extract, (str, int)):
             label_to_extract = [label_to_extract]
         if isinstance(label_to_extract, (tuple, list)):
-            
-            for k, v in self.labels.items():
-                if issubset(label_to_extract, v['dof']):
-                    break
-
-            label_to_extract = {v['name']: label_to_extract}
-
-        for k, v in label_to_extract.items():
-            if isinstance(v, (int, str)):
-                label_to_extract[k] = [v]
-
-        indeces = []
-        for dim in range(self.ndim):
-
-            boolean_idx = [True] * self.shape[dim]
+            last_dim_label = self.labels[self.tensor.ndim - 1]['dof']
+            if set(label_to_extract).issubset(last_dim_label) is False:
+                raise ValueError('Cannot extract a dof which is not in the original LabelTensor')
+            idx_to_extract = [last_dim_label.index(i) for i in label_to_extract]
+            new_tensor = deepcopy(self.tensor)
+            new_tensor = new_tensor[..., idx_to_extract]
+            new_labels = deepcopy(self.labels)
+            last_dim_new_label = {self.tensor.ndim - 1: {
+                'dof': label_to_extract,
+                'name': self.labels[self.tensor.ndim - 1]['name']
+            }}
+            new_labels.update(last_dim_new_label)
+        elif isinstance(label_to_extract, dict):
+            new_labels = (deepcopy(self.labels))
+            new_tensor = deepcopy(self.tensor)
+            for k, v in label_to_extract.items():
+                idx_dim = None
+                for kl, vl in self.labels.items():
+                    if vl['name'] == k:
+                        idx_dim = kl
+                        break
+                dim_labels = self.labels[idx_dim]['dof']
+                if isinstance(label_to_extract[k], (int, str)):
+                    label_to_extract[k] = [label_to_extract[k]]
+                if set(label_to_extract[k]).issubset(dim_labels) is False:
+                    raise ValueError('Cannot extract a dof which is not in the original labeltensor')
+                idx_to_extract = [dim_labels.index(i) for i in label_to_extract[k]]
+                indexer = [slice(None)] * idx_dim + [idx_to_extract] + [slice(None)] * (self.tensor.ndim - idx_dim - 1)
+                new_tensor = new_tensor[indexer]
+                dim_new_label = {idx_dim: {
+                    'dof': label_to_extract[k],
+                    'name': self.labels[idx_dim]['name']
+                }}
+                new_labels.update(dim_new_label)
+        else:
+            raise ValueError('labels_to_extract must be str or list or dict')
+        return LabelTensor(new_tensor, new_labels)
 
-            for dim_to_extract, dof_to_extract in label_to_extract.items():
-                if dim_to_extract == self.labels[dim]['name']:
-                    boolean_idx = [False] * self.shape[dim]
-                    for label in dof_to_extract:
-                        idx_to_keep = self.labels[dim]['dof'].index(label)
-                        boolean_idx[idx_to_keep] = True
-
-            boolean_idx = torch.Tensor(boolean_idx).bool()
-
-            indeces.append(boolean_idx)
-
-        final_shapes = [sum(idx) for idx in indeces]
-        grids = torch.meshgrid(*indeces)
-
-        ii = grids[0]
-        for grid in grids[1:]:
-            ii = torch.logical_and(ii, grid)
-
-        new_tensor = self.tensor[ii].reshape(*final_shapes)
-
-        return LabelTensor(new_tensor, label_to_extract)
-        
     def __str__(self):
         """
         returns a string with the representation of the class
@@ -409,23 +105,45 @@ def __str__(self):
         s = ''
         for key, value in self.labels.items():
             s += f"{key}: {value}\n"
-        s += '\n' 
+        s += '\n'
         s += super().__str__()
         return s
-    
+
     @staticmethod
-    def stack(tensors):
-        """ 
+    def cat(tensors, dim=0):
+        """
+        Stack a list of tensors. For example, given a tensor `a` of shape `(n,m,dof)` and a tensor `b` of dimension `(n',m,dof)`
+        the resulting tensor is of shape `(n+n',m,dof)`
         """
         if len(tensors) == 0:
             return []
-
         if len(tensors) == 1:
             return tensors[0]
-        
-        raise NotImplementedError
-        labels = [tensor.labels for tensor in tensors]
-
+        n_dims = tensors[0].ndim
+        new_labels_cat_dim = []
+        for i in range(n_dims):
+            name = tensors[0].labels[i]['name']
+            if i != dim:
+                dof = tensors[0].labels[i]['dof']
+                for tensor in tensors:
+                    dof_to_check = tensor.labels[i]['dof']
+                    name_to_check = tensor.labels[i]['name']
+                    if dof != dof_to_check or name != name_to_check:
+                        raise ValueError('dimensions must have the same dof and name')
+            else:
+                for tensor in tensors:
+                    new_labels_cat_dim += tensor.labels[i]['dof']
+                    name_to_check = tensor.labels[i]['name']
+                    if name != name_to_check:
+                        raise ValueError('dimensions must have the same dof and name')
+        new_tensor = torch.cat(tensors, dim=dim)
+        labels = tensors[0].labels
+        labels.pop(dim)
+        new_labels_cat_dim = new_labels_cat_dim if len(set(new_labels_cat_dim)) == len(new_labels_cat_dim) \
+            else range(new_tensor.shape[dim])
+        labels[dim] = {'dof': new_labels_cat_dim,
+                       'name': tensors[1].labels[dim]['name']}
+        return LabelTensor(new_tensor, labels)
 
     def requires_grad_(self, mode=True):
         lt = super().requires_grad_(mode)
@@ -454,10 +172,25 @@ def clone(self, *args, **kwargs):
         :return: A copy of the tensor.
         :rtype: LabelTensor
         """
-        # # used before merging
-        # try:
-        #     out = LabelTensor(super().clone(*args, **kwargs), self.labels)
-        # except:
-        #     out = super().clone(*args, **kwargs)
+
         out = LabelTensor(super().clone(*args, **kwargs), self.labels)
-        return out
\ No newline at end of file
+        return out
+
+    def update_labels(self, labels):
+        self.labels = {
+            idx_: {
+                'dof': range(self.tensor.shape[idx_]),
+                'name': idx_
+            } for idx_ in range(self.tensor.ndim)
+        }
+        tensor_shape = self.tensor.shape
+        for k, v in labels.items():
+            if len(v['dof']) != len(set(v['dof'])):
+                raise ValueError("dof must be unique")
+            if len(v['dof']) != tensor_shape[k]:
+                raise ValueError('Number of dof does not match with tensor dimension')
+        self.labels.update(labels)
+
+    def init_labels_from_list(self, labels):
+        last_dim_labels = {self.tensor.ndim - 1: {'dof': labels, 'name': self.tensor.ndim - 1}}
+        self.update_labels(last_dim_labels)
\ No newline at end of file
diff --git a/tests/test_label_tensor.py b/tests/test_label_tensor.py
index 5d5693d50..7484a4971 100644
--- a/tests/test_label_tensor.py
+++ b/tests/test_label_tensor.py
@@ -38,7 +38,7 @@ def test_extract_column(labels, labels_te):
     assert torch.all(torch.isclose(data[:, 2].reshape(-1, 1), new))
 
 @pytest.mark.parametrize("labels", [labels_row, labels_all])
-@pytest.mark.parametrize("labels_te", [2, [2], {'samples': [2]}])
+@pytest.mark.parametrize("labels_te", [{'samples': [2]}])
 def test_extract_row(labels, labels_te):
     tensor = LabelTensor(data, labels)
     new = tensor.extract(labels_te)
@@ -62,7 +62,7 @@ def test_extract_2D(labels_te):
 
 def test_extract_3D():
     labels = labels_all
-    data = torch.rand((20, 3, 4))
+    data = torch.rand(20, 3, 4)
     labels = {
         1: {
             "name": "space",
@@ -77,6 +77,7 @@ def test_extract_3D():
         'space': ['x', 'z'],
         'time': range(1, 4)
     }
+
     tensor = LabelTensor(data, labels)
     new = tensor.extract(labels_te)
     assert new.ndim == tensor.ndim
@@ -86,106 +87,4 @@ def test_extract_3D():
     assert torch.all(torch.isclose(
         data[:, 0::2, 1:4].reshape(20, 2, 3),
         new
-    ))
-
-# def test_labels():
-#     tensor = LabelTensor(data, labels)
-#     assert isinstance(tensor, torch.Tensor)
-#     assert tensor.labels == labels
-#     with pytest.raises(ValueError):
-#         tensor.labels = labels[:-1]
-
-
-# def test_extract():
-#     label_to_extract = ['a', 'c']
-#     tensor = LabelTensor(data, labels)
-#     new = tensor.extract(label_to_extract)
-#     assert new.labels == label_to_extract
-#     assert new.shape[1] == len(label_to_extract)
-#     assert torch.all(torch.isclose(data[:, 0::2], new))
-
-
-# def test_extract_onelabel():
-#     label_to_extract = ['a']
-#     tensor = LabelTensor(data, labels)
-#     new = tensor.extract(label_to_extract)
-#     assert new.ndim == 2
-#     assert new.labels == label_to_extract
-#     assert new.shape[1] == len(label_to_extract)
-#     assert torch.all(torch.isclose(data[:, 0].reshape(-1, 1), new))
-
-
-# def test_wrong_extract():
-#     label_to_extract = ['a', 'cc']
-#     tensor = LabelTensor(data, labels)
-#     with pytest.raises(ValueError):
-#         tensor.extract(label_to_extract)
-
-
-# def test_extract_order():
-#     label_to_extract = ['c', 'a']
-#     tensor = LabelTensor(data, labels)
-#     new = tensor.extract(label_to_extract)
-#     expected = torch.cat(
-#         (data[:, 2].reshape(-1, 1), data[:, 0].reshape(-1, 1)),
-#         dim=1)
-#     assert new.labels == label_to_extract
-#     assert new.shape[1] == len(label_to_extract)
-#     assert torch.all(torch.isclose(expected, new))
-
-
-# def test_merge():
-#     tensor = LabelTensor(data, labels)
-#     tensor_a = tensor.extract('a')
-#     tensor_b = tensor.extract('b')
-#     tensor_c = tensor.extract('c')
-
-#     tensor_bc = tensor_b.append(tensor_c)
-#     assert torch.allclose(tensor_bc, tensor.extract(['b', 'c']))
-
-
-# def test_merge2():
-#     tensor = LabelTensor(data, labels)
-#     tensor_b = tensor.extract('b')
-#     tensor_c = tensor.extract('c')
-
-#     tensor_bc = tensor_b.append(tensor_c)
-#     assert torch.allclose(tensor_bc, tensor.extract(['b', 'c']))
-
-
-# def test_getitem():
-#     tensor = LabelTensor(data, labels)
-#     tensor_view = tensor['a']
-
-#     assert tensor_view.labels == ['a']
-#     assert torch.allclose(tensor_view.flatten(), data[:, 0])
-
-#     tensor_view = tensor['a', 'c']
-
-#     assert tensor_view.labels == ['a', 'c']
-#     assert torch.allclose(tensor_view, data[:, 0::2])
-
-# def test_getitem2():
-#     tensor = LabelTensor(data, labels)
-#     tensor_view = tensor[:5]
-#     assert tensor_view.labels == labels
-#     assert torch.allclose(tensor_view, data[:5])
-
-#     idx = torch.randperm(tensor.shape[0])
-#     tensor_view = tensor[idx]
-#     assert tensor_view.labels == labels
-
-
-# def test_slice():
-#     tensor = LabelTensor(data, labels)
-#     tensor_view = tensor[:5, :2]
-#     assert tensor_view.labels == labels[:2]
-#     assert torch.allclose(tensor_view, data[:5, :2])
-
-#     tensor_view2 = tensor[3]
-#     assert tensor_view2.labels == labels
-#     assert torch.allclose(tensor_view2, data[3])
-
-#     tensor_view3 = tensor[:, 2]
-#     assert tensor_view3.labels == labels[2]
-#     assert torch.allclose(tensor_view3, data[:, 2].reshape(-1, 1))
+    ))
\ No newline at end of file

From e3d06996987ae38784c829facd7b0a4d451bbe7e Mon Sep 17 00:00:00 2001
From: FilippoOlivo <filippo@filippoolivo.com>
Date: Sat, 28 Sep 2024 17:45:54 +0200
Subject: [PATCH 14/55] Update docstring in LabelTensor class

---
 pina/label_tensor.py | 20 ++++++++++++++++++++
 1 file changed, 20 insertions(+)

diff --git a/pina/label_tensor.py b/pina/label_tensor.py
index f40c36dca..36692726b 100644
--- a/pina/label_tensor.py
+++ b/pina/label_tensor.py
@@ -114,6 +114,13 @@ def cat(tensors, dim=0):
         """
         Stack a list of tensors. For example, given a tensor `a` of shape `(n,m,dof)` and a tensor `b` of dimension `(n',m,dof)`
         the resulting tensor is of shape `(n+n',m,dof)`
+
+        :param tensors: tensors to concatenate
+        :type tensors: list(LabelTensor)
+        :param dim: dimensions on which you want to perform the operation (default 0)
+        :type dim: int
+        :rtype: LabelTensor
+        :raises ValueError: either number dof or dimensions names differ
         """
         if len(tensors) == 0:
             return []
@@ -177,6 +184,13 @@ def clone(self, *args, **kwargs):
         return out
 
     def update_labels(self, labels):
+        '''
+        Update the internal label representation according to the values passed as input.
+
+        :param labels: The label(s) to update.
+        :type labels: dict
+        :raises ValueError: dof list contain duplicates or number of dof does not match with tensor shape
+        '''
         self.labels = {
             idx_: {
                 'dof': range(self.tensor.shape[idx_]),
@@ -192,5 +206,11 @@ def update_labels(self, labels):
         self.labels.update(labels)
 
     def init_labels_from_list(self, labels):
+        '''
+        Given a list of dof, this method update the internal label representation
+
+        :param labels: The label(s) to update.
+        :type labels: list
+        '''
         last_dim_labels = {self.tensor.ndim - 1: {'dof': labels, 'name': self.tensor.ndim - 1}}
         self.update_labels(last_dim_labels)
\ No newline at end of file

From 2aa6645242a745b2d1a4e86a1da5162f6f2bdf02 Mon Sep 17 00:00:00 2001
From: FilippoOlivo <filippo@filippoolivo.com>
Date: Mon, 30 Sep 2024 10:15:45 +0200
Subject: [PATCH 15/55] Minor update on docstring in LabelTensor class

---
 pina/label_tensor.py | 13 +++++++------
 1 file changed, 7 insertions(+), 6 deletions(-)

diff --git a/pina/label_tensor.py b/pina/label_tensor.py
index 36692726b..7646dd8a3 100644
--- a/pina/label_tensor.py
+++ b/pina/label_tensor.py
@@ -35,14 +35,15 @@ def __init__(self, x, labels):
                     {1: {"name": "space"['a', 'b', 'c'])
 
         """
-        from .utils import check_consistency
+        self.labels = None
         if isinstance(labels, dict):
-            # check_consistency(labels, dict)
             self.update_labels(labels)
         elif isinstance(labels, list):
             self.init_labels_from_list(labels)
         elif isinstance(labels, str):
             labels = [labels]
+            self.init_labels_from_list(labels)
+        else:
             raise ValueError(f"labels must be list, dict or string.")
 
     def extract(self, label_to_extract):
@@ -184,13 +185,13 @@ def clone(self, *args, **kwargs):
         return out
 
     def update_labels(self, labels):
-        '''
+        """
         Update the internal label representation according to the values passed as input.
 
         :param labels: The label(s) to update.
         :type labels: dict
         :raises ValueError: dof list contain duplicates or number of dof does not match with tensor shape
-        '''
+        """
         self.labels = {
             idx_: {
                 'dof': range(self.tensor.shape[idx_]),
@@ -206,11 +207,11 @@ def update_labels(self, labels):
         self.labels.update(labels)
 
     def init_labels_from_list(self, labels):
-        '''
+        """
         Given a list of dof, this method update the internal label representation
 
         :param labels: The label(s) to update.
         :type labels: list
-        '''
+        """
         last_dim_labels = {self.tensor.ndim - 1: {'dof': labels, 'name': self.tensor.ndim - 1}}
         self.update_labels(last_dim_labels)
\ No newline at end of file

From 8cff2e33ba4528b583a2a242c1c2d977719075ee Mon Sep 17 00:00:00 2001
From: FilippoOlivo <filippo@filippoolivo.com>
Date: Mon, 30 Sep 2024 12:23:15 +0200
Subject: [PATCH 16/55] Add concatenation test for LabelTensor

---
 tests/test_label_tensor.py | 60 +++++++++++++++++++++++++++++++++++++-
 1 file changed, 59 insertions(+), 1 deletion(-)

diff --git a/tests/test_label_tensor.py b/tests/test_label_tensor.py
index 7484a4971..f87d3abb1 100644
--- a/tests/test_label_tensor.py
+++ b/tests/test_label_tensor.py
@@ -87,4 +87,62 @@ def test_extract_3D():
     assert torch.all(torch.isclose(
         data[:, 0::2, 1:4].reshape(20, 2, 3),
         new
-    ))
\ No newline at end of file
+    ))
+
+def test_concatenation_3D():
+    data_1 = torch.rand(20, 3, 4)
+    labels_1 = ['x', 'y', 'z', 'w']
+    lt1 = LabelTensor(data_1, labels_1)
+    data_2 = torch.rand(50, 3, 4)
+    labels_2 = ['x', 'y', 'z', 'w']
+    lt2 = LabelTensor(data_2, labels_2)
+    lt_cat = LabelTensor.cat([lt1, lt2])
+    assert lt_cat.shape == (70, 3, 4)
+    assert lt_cat.labels[0]['dof'] == range(70)
+    assert lt_cat.labels[1]['dof'] == range(3)
+    assert lt_cat.labels[2]['dof'] == ['x', 'y', 'z', 'w']
+
+    data_1 = torch.rand(20, 3, 4)
+    labels_1 = ['x', 'y', 'z', 'w']
+    lt1 = LabelTensor(data_1, labels_1)
+    data_2 = torch.rand(20, 2, 4)
+    labels_2 = ['x', 'y', 'z', 'w']
+    lt2 = LabelTensor(data_2, labels_2)
+    lt_cat = LabelTensor.cat([lt1, lt2], dim=1)
+    assert lt_cat.shape == (20, 5, 4)
+    assert lt_cat.labels[0]['dof'] == range(20)
+    assert lt_cat.labels[1]['dof'] == range(5)
+    assert lt_cat.labels[2]['dof'] == ['x', 'y', 'z', 'w']
+
+    data_1 = torch.rand(20, 3, 2)
+    labels_1 = ['x', 'y']
+    lt1 = LabelTensor(data_1, labels_1)
+    data_2 = torch.rand(20, 3, 3)
+    labels_2 = ['z', 'w', 'a']
+    lt2 = LabelTensor(data_2, labels_2)
+    lt_cat = LabelTensor.cat([lt1, lt2], dim=2)
+    assert lt_cat.shape == (20, 3, 5)
+    assert lt_cat.labels[2]['dof'] == ['x', 'y', 'z', 'w', 'a']
+    assert lt_cat.labels[0]['dof'] == range(20)
+    assert lt_cat.labels[1]['dof'] == range(3)
+
+    data_1 = torch.rand(20, 2, 4)
+    labels_1 = ['x', 'y', 'z', 'w']
+    lt1 = LabelTensor(data_1, labels_1)
+    data_2 = torch.rand(20, 3, 4)
+    labels_2 = ['x', 'y', 'z', 'w']
+    lt2 = LabelTensor(data_2, labels_2)
+    with pytest.raises(ValueError):
+        LabelTensor.cat([lt1, lt2], dim=2)
+
+    data_1 = torch.rand(20, 3, 2)
+    labels_1 = ['x', 'y']
+    lt1 = LabelTensor(data_1, labels_1)
+    data_2 = torch.rand(20, 3, 3)
+    labels_2 = ['x', 'w', 'a']
+    lt2 = LabelTensor(data_2, labels_2)
+    lt_cat = LabelTensor.cat([lt1, lt2], dim=2)
+    assert lt_cat.shape == (20, 3, 5)
+    assert lt_cat.labels[2]['dof'] == range(5)
+    assert lt_cat.labels[0]['dof'] == range(20)
+    assert lt_cat.labels[1]['dof'] == range(3)

From 499a5bddbea5ac0fde5575b26c8296097b8d8e0b Mon Sep 17 00:00:00 2001
From: Dario Coscia <dariocos99@gmail.com>
Date: Thu, 3 Oct 2024 21:33:37 +0200
Subject: [PATCH 17/55] new conditions

---
 pina/condition/__init__.py                  |  10 +-
 pina/condition/condition.py                 | 119 +++++++-------------
 pina/condition/condition_interface.py       |  34 +++---
 pina/condition/data_condition.py            |  44 ++++++++
 pina/condition/domain_equation_condition.py |  45 +++++---
 pina/condition/domain_output_condition.py   |  44 --------
 pina/condition/input_equation_condition.py  |  43 +++++--
 pina/condition/input_output_condition.py    |  42 +++++++
 8 files changed, 210 insertions(+), 171 deletions(-)
 create mode 100644 pina/condition/data_condition.py
 delete mode 100644 pina/condition/domain_output_condition.py
 create mode 100644 pina/condition/input_output_condition.py

diff --git a/pina/condition/__init__.py b/pina/condition/__init__.py
index ff329c1bc..4c89b75d1 100644
--- a/pina/condition/__init__.py
+++ b/pina/condition/__init__.py
@@ -1,10 +1,12 @@
 __all__ = [
     'Condition',
     'ConditionInterface',
-    'DomainOutputCondition',
-    'DomainEquationCondition'
+    'DomainEquationCondition',
+    'InputPointsEquationCondition',
+    'InputOutputPointsCondition',
 ]
 
 from .condition_interface import ConditionInterface
-from .domain_output_condition import DomainOutputCondition
-from .domain_equation_condition import DomainEquationCondition
\ No newline at end of file
+from .domain_equation_condition import DomainEquationCondition
+from .input_equation_condition import InputPointsEquationCondition
+from .input_output_condition import InputOutputPointsCondition
\ No newline at end of file
diff --git a/pina/condition/condition.py b/pina/condition/condition.py
index d815838bb..ddc722f0c 100644
--- a/pina/condition/condition.py
+++ b/pina/condition/condition.py
@@ -1,27 +1,21 @@
 """ Condition module. """
 
-from ..label_tensor import LabelTensor
-from ..domain import DomainInterface
-from ..equation.equation import Equation
-
-from . import DomainOutputCondition, DomainEquationCondition
-
-
-def dummy(a):
-    """Dummy function for testing purposes."""
-    return None
-
+from .domain_equation_condition import DomainEquationCondition
+from .input_equation_condition import InputPointsEquationCondition
+from .input_output_condition import InputOutputPointsCondition
+from .data_condition import DataConditionInterface
 
 class Condition:
     """
     The class ``Condition`` is used to represent the constraints (physical
     equations, boundary conditions, etc.) that should be satisfied in the
-    problem at hand. Condition objects are used to formulate the PINA :obj:`pina.problem.abstract_problem.AbstractProblem` object.
-    Conditions can be specified in three ways:
+    problem at hand. Condition objects are used to formulate the
+    PINA :obj:`pina.problem.abstract_problem.AbstractProblem` object.
+    Conditions can be specified in four ways:
 
         1. By specifying the input and output points of the condition; in such a
         case, the model is trained to produce the output points given the input
-        points.
+        points. Those points can either be torch.Tensor, LabelTensors, Graph
 
         2. By specifying the location and the equation of the condition; in such
         a case, the model is trained to minimize the equation residual by
@@ -29,79 +23,48 @@ class Condition:
 
         3. By specifying the input points and the equation of the condition; in
         such a case, the model is trained to minimize the equation residual by
-        evaluating it at the passed input points.
+        evaluating it at the passed input points. The input points must be
+        a LabelTensor.
+
+        4. By specifying only the data matrix; in such a case the model is
+        trained with an unsupervised costum loss and uses the data in training.
+        Additionaly conditioning variables can be passed, whenever the model
+        has extra conditioning variable it depends on.
 
     Example::
 
-    >>> example_domain = Span({'x': [0, 1], 'y': [0, 1]})
-    >>> def example_dirichlet(input_, output_):
-    >>>     value = 0.0
-    >>>     return output_.extract(['u']) - value
-    >>> example_input_pts = LabelTensor(
-    >>>     torch.tensor([[0, 0, 0]]), ['x', 'y', 'z'])
-    >>> example_output_pts = LabelTensor(torch.tensor([[1, 2]]), ['a', 'b'])
-    >>>
-    >>> Condition(
-    >>>     input_points=example_input_pts,
-    >>>     output_points=example_output_pts)
-    >>> Condition(
-    >>>     location=example_domain,
-    >>>     equation=example_dirichlet)
-    >>> Condition(
-    >>>     input_points=example_input_pts,
-    >>>     equation=example_dirichlet)
+    >>> TODO
 
     """
 
-    # def _dictvalue_isinstance(self, dict_, key_, class_):
-    #     """Check if the value of a dictionary corresponding to `key` is an instance of `class_`."""
-    #     if key_ not in dict_.keys():
-    #         return True
+    __slots__ = list(
+                    set(
+                        InputOutputPointsCondition.__slots__,
+                        InputPointsEquationCondition.__slots__,
+                        DomainEquationCondition.__slots__,
+                        DataConditionInterface.__slots__
 
-    #     return isinstance(dict_[key_], class_)
-
-    # def __init__(self, *args, **kwargs):
-    #     """
-    #     Constructor for the `Condition` class.
-    #     """
-    #     self.data_weight = kwargs.pop("data_weight", 1.0)
-
-    #     if len(args) != 0:
-    #         raise ValueError(
-    #             f"Condition takes only the following keyword arguments: {Condition.__slots__}."
-    #         )
+                    )
+                )
 
     def __new__(cls, *args, **kwargs):
-
-        if sorted(kwargs.keys()) == sorted(["input_points", "output_points"]):
-            return DomainOutputCondition(
-                domain=kwargs["input_points"],
-                output_points=kwargs["output_points"]
+        
+        if len(args) != 0:
+            raise ValueError(
+                f"Condition takes only the following keyword '
+                'arguments: {Condition.__slots__}."
             )
-        elif sorted(kwargs.keys()) == sorted(["domain", "output_points"]):
-            return DomainOutputCondition(**kwargs)
-        elif sorted(kwargs.keys()) == sorted(["domain", "equation"]):
+    
+        sorted_keys = sorted(kwargs.keys()) 
+        if sorted_keys == sorted(InputOutputPointsCondition.__slots__):
+            return InputOutputPointsCondition(**kwargs)
+        elif sorted_keys == sorted(InputPointsEquationCondition.__slots__):
+            return InputPointsEquationCondition(**kwargs)
+        elif sorted_keys == sorted(DomainEquationCondition.__slots__):
             return DomainEquationCondition(**kwargs)
+        elif sorted_keys == sorted(DataConditionInterface.__slots__):
+            return DataConditionInterface(**kwargs)
+        elif sorted_keys == DataConditionInterface.__slots__[0]:
+            return DataConditionInterface(**kwargs)
         else:
-            raise ValueError(f"Invalid keyword arguments {kwargs.keys()}.")
-        # TODO: remove, not used anymore
-        '''
-        if (
-            sorted(kwargs.keys()) != sorted(["input_points", "output_points"])
-            and sorted(kwargs.keys()) != sorted(["location", "equation"])
-            and sorted(kwargs.keys()) != sorted(["input_points", "equation"])
-        ):
-            raise ValueError(f"Invalid keyword arguments {kwargs.keys()}.")
-        # TODO: remove, not used anymore
-        if not self._dictvalue_isinstance(kwargs, "input_points", LabelTensor):
-            raise TypeError("`input_points` must be a torch.Tensor.")
-        if not self._dictvalue_isinstance(kwargs, "output_points", LabelTensor):
-            raise TypeError("`output_points` must be a torch.Tensor.")
-        if not self._dictvalue_isinstance(kwargs, "location", Location):
-            raise TypeError("`location` must be a Location.")
-        if not self._dictvalue_isinstance(kwargs, "equation", Equation):
-            raise TypeError("`equation` must be a Equation.")
-
-        for key, value in kwargs.items():
-            setattr(self, key, value)
-        '''
\ No newline at end of file
+            raise ValueError(f"Invalid keyword arguments {kwargs.keys()}.")
\ No newline at end of file
diff --git a/pina/condition/condition_interface.py b/pina/condition/condition_interface.py
index 0626a6d83..f380dcf68 100644
--- a/pina/condition/condition_interface.py
+++ b/pina/condition/condition_interface.py
@@ -1,21 +1,25 @@
 
-from abc import ABCMeta, abstractmethod
+from abc import ABCMeta
 
 
 class ConditionInterface(metaclass=ABCMeta):
 
-    def __init__(self) -> None:
-        self._problem = None
+    condition_types = ['physical', 'supervised', 'unsupervised']
+    def __init__(self):
+        self._condition_type = None
 
-    @abstractmethod
-    def residual(self, model):
-        """
-        Compute the residual of the condition.
-
-        :param model: The model to evaluate the condition.
-        :return: The residual of the condition.
-        """
-        pass
-
-    def set_problem(self, problem):
-        self._problem = problem
+    @property
+    def condition_type(self):
+        return self._condition_type
+    
+    @condition_type.setattr
+    def condition_type(self, values):
+        if not isinstance(values, (list, tuple)):
+            values = [values]
+        for value in values:
+            if value not in ConditionInterface.condition_types:
+                raise ValueError( 
+                                'Unavailable type of condition, expected one of'
+                                f' {ConditionInterface.condition_types}.'
+                                )
+        self._condition_type = values
\ No newline at end of file
diff --git a/pina/condition/data_condition.py b/pina/condition/data_condition.py
new file mode 100644
index 000000000..259eb5669
--- /dev/null
+++ b/pina/condition/data_condition.py
@@ -0,0 +1,44 @@
+import torch
+
+from . import ConditionInterface
+from ..label_tensor import LabelTensor
+from ..graph import Graph
+from ..utils import check_consistency
+
+class DataConditionInterface(ConditionInterface):
+    """
+    Condition for data. This condition must be used every
+    time a Unsupervised Loss is needed in the Solver. The conditionalvariable
+    can be passed as extra-input when the model learns a conditional
+    distribution
+    """
+
+    __slots__ = ["data", "conditionalvariable"]
+
+    def __init__(self, data, conditionalvariable=None):
+        """
+        TODO
+        """
+        super().__init__()
+        self.data = data
+        self.conditionalvariable = conditionalvariable
+        self.condition_type = 'unsupervised'
+
+    @property
+    def data(self):
+        return self._data
+    
+    @data.setter
+    def data(self, value):
+        check_consistency(value, (LabelTensor, Graph, torch.Tensor))
+        self._data = value
+
+    @property
+    def conditionalvariable(self):
+        return self._conditionalvariable
+    
+    @data.setter
+    def conditionalvariable(self, value):
+        if value is not None:
+            check_consistency(value, (LabelTensor, Graph, torch.Tensor))
+        self._data = value
\ No newline at end of file
diff --git a/pina/condition/domain_equation_condition.py b/pina/condition/domain_equation_condition.py
index 15df3f85f..9838ad7e5 100644
--- a/pina/condition/domain_equation_condition.py
+++ b/pina/condition/domain_equation_condition.py
@@ -1,34 +1,43 @@
+import torch
+
 from .condition_interface import ConditionInterface
+from ..label_tensor import LabelTensor
+from ..graph import Graph
+from ..utils import check_consistency
+from ..domain import DomainInterface
+from ..equation.equation_interface import EquationInterface
 
 class DomainEquationCondition(ConditionInterface):
     """
-    Condition for input/output data.
+    Condition for domain/equation data. This condition must be used every
+    time a Physics Informed Loss is needed in the Solver.
     """
 
     __slots__ = ["domain", "equation"]
 
     def __init__(self, domain, equation):
         """
-        Constructor for the `InputOutputCondition` class.
+        TODO
         """
         super().__init__()
         self.domain = domain
         self.equation = equation
+        self.condition_type = 'physics'
 
-    def residual(self, model):
-        """
-        Compute the residual of the condition.
-        """
-        self.batch_residual(model, self.domain, self.equation)
-
-    @staticmethod
-    def batch_residual(model, input_pts, equation):
-        """
-        Compute the residual of the condition for a single batch. Input and
-        output points are provided as arguments.
+    @property
+    def domain(self):
+        return self._domain
+    
+    @domain.setter
+    def domain(self, value):
+        check_consistency(value, (DomainInterface))
+        self._domain = value
 
-        :param torch.nn.Module model: The model to evaluate the condition.
-        :param torch.Tensor input_pts: The input points.
-        :param torch.Tensor equation: The output points.
-        """
-        return equation.residual(input_pts, model(input_pts))
\ No newline at end of file
+    @property
+    def equation(self):
+        return self._equation
+    
+    @equation.setter
+    def equation(self, value):
+        check_consistency(value, (EquationInterface))
+        self._equation = value
\ No newline at end of file
diff --git a/pina/condition/domain_output_condition.py b/pina/condition/domain_output_condition.py
deleted file mode 100644
index f847720b5..000000000
--- a/pina/condition/domain_output_condition.py
+++ /dev/null
@@ -1,44 +0,0 @@
-
-from . import ConditionInterface
-
-class DomainOutputCondition(ConditionInterface):
-    """
-    Condition for input/output data.
-    """
-
-    __slots__ = ["domain", "output_points"]
-
-    def __init__(self, domain, output_points):
-        """
-        Constructor for the `InputOutputCondition` class.
-        """
-        super().__init__()
-        print(self)
-        self.domain = domain
-        self.output_points = output_points
-
-    @property
-    def input_points(self):
-        """
-        Get the input points of the condition.
-        """
-        return self._problem.domains[self.domain]
-
-    def residual(self, model):
-        """
-        Compute the residual of the condition.
-        """
-        return self.batch_residual(model, self.domain, self.output_points)
-
-    @staticmethod
-    def batch_residual(model, input_points, output_points):
-        """
-        Compute the residual of the condition for a single batch. Input and
-        output points are provided as arguments.
-
-        :param torch.nn.Module model: The model to evaluate the condition.
-        :param torch.Tensor input_points: The input points.
-        :param torch.Tensor output_points: The output points.
-        """
-
-        return output_points - model(input_points)
\ No newline at end of file
diff --git a/pina/condition/input_equation_condition.py b/pina/condition/input_equation_condition.py
index 288022c00..c4b9f8dba 100644
--- a/pina/condition/input_equation_condition.py
+++ b/pina/condition/input_equation_condition.py
@@ -1,23 +1,42 @@
+import torch
 
-from . import ConditionInterface
+from .condition_interface import ConditionInterface
+from ..label_tensor import LabelTensor
+from ..graph import Graph
+from ..utils import check_consistency
+from ..equation.equation_interface import EquationInterface
 
-class InputEquationCondition(ConditionInterface):
+class InputPointsEquationCondition(ConditionInterface):
     """
-    Condition for input/output data.
+    Condition for input_points/equation data. This condition must be used every
+    time a Physics Informed Loss is needed in the Solver.
     """
 
-    __slots__ = ["input_points", "output_points"]
+    __slots__ = ["input_points", "equation"]
 
-    def __init__(self, input_points, output_points):
+    def __init__(self, input_points, equation):
         """
-        Constructor for the `InputOutputCondition` class.
+        TODO
         """
         super().__init__()
         self.input_points = input_points
-        self.output_points = output_points
+        self.equation = equation
+        self.condition_type = 'physics'
 
-    def residual(self, model):
-        """
-        Compute the residual of the condition.
-        """
-        return self.output_points - model(self.input_points)
\ No newline at end of file
+    @property
+    def input_points(self):
+        return self._input_points
+    
+    @input_points.setter
+    def input_points(self, value):
+        check_consistency(value, (LabelTensor)) # for now only labeltensors, we need labels for the operators!
+        self._input_points = value
+
+    @property
+    def equation(self):
+        return self._equation
+    
+    @equation.setter
+    def equation(self, value):
+        check_consistency(value, (EquationInterface))
+        self._equation = value
\ No newline at end of file
diff --git a/pina/condition/input_output_condition.py b/pina/condition/input_output_condition.py
new file mode 100644
index 000000000..fd6b7a034
--- /dev/null
+++ b/pina/condition/input_output_condition.py
@@ -0,0 +1,42 @@
+
+import torch
+
+from .condition_interface import ConditionInterface
+from ..label_tensor import LabelTensor
+from ..graph import Graph
+from ..utils import check_consistency
+
+class InputOutputPointsCondition(ConditionInterface):
+    """
+    Condition for domain/equation data. This condition must be used every
+    time a Physics Informed or a Supervised Loss is needed in the Solver.
+    """
+
+    __slots__ = ["input_points", "output_points"]
+
+    def __init__(self, input_points, output_points):
+        """
+        TODO
+        """
+        super().__init__()
+        self.input_points = input_points
+        self.output_points = output_points
+        self.condition_type = ['supervised', 'physics']
+
+    @property
+    def input_points(self):
+        return self._input_points
+    
+    @input_points.setter
+    def input_points(self, value):
+        check_consistency(value, (LabelTensor, Graph, torch.Tensor))
+        self._input_points = value
+
+    @property
+    def output_points(self):
+        return self._output_points
+    
+    @output_points.setter
+    def output_points(self, value):
+        check_consistency(value, (LabelTensor, Graph, torch.Tensor))
+        self._output_points = value
\ No newline at end of file

From 90980bdbb3b61ae85345cee0b150366728c8ed52 Mon Sep 17 00:00:00 2001
From: Dario Coscia <dariocos99@gmail.com>
Date: Thu, 3 Oct 2024 21:55:16 +0200
Subject: [PATCH 18/55] sampling mode domain added

---
 pina/domain/cartesian.py           |  1 +
 pina/domain/difference_domain.py   |  2 +-
 pina/domain/domain_interface.py    | 24 ++++++++++++++++++++++++
 pina/domain/ellipsoid.py           |  3 ++-
 pina/domain/exclusion_domain.py    |  2 +-
 pina/domain/intersection_domain.py |  2 +-
 pina/domain/operation_interface.py |  5 ++++-
 pina/domain/simplex.py             | 14 +++++++++-----
 pina/domain/union_domain.py        |  3 +++
 9 files changed, 46 insertions(+), 10 deletions(-)

diff --git a/pina/domain/cartesian.py b/pina/domain/cartesian.py
index f1865e966..e2e066ae8 100644
--- a/pina/domain/cartesian.py
+++ b/pina/domain/cartesian.py
@@ -21,6 +21,7 @@ def __init__(self, cartesian_dict):
         """
         self.fixed_ = {}
         self.range_ = {}
+        self.sample_modes = ["random", "grid", "lh", "chebyshev", "latin"]
 
         for k, v in cartesian_dict.items():
             if isinstance(v, (int, float)):
diff --git a/pina/domain/difference_domain.py b/pina/domain/difference_domain.py
index d2ba414f0..9554aaf32 100644
--- a/pina/domain/difference_domain.py
+++ b/pina/domain/difference_domain.py
@@ -77,7 +77,7 @@ def sample(self, n, mode="random", variables="all"):
                 5
 
         """
-        if mode != "random":
+        if mode != self.sample_modes:
             raise NotImplementedError(
                 f"{mode} is not a valid mode for sampling."
             )
diff --git a/pina/domain/domain_interface.py b/pina/domain/domain_interface.py
index 5906d2851..208bfdfbc 100644
--- a/pina/domain/domain_interface.py
+++ b/pina/domain/domain_interface.py
@@ -9,6 +9,30 @@ class DomainInterface(metaclass=ABCMeta):
     Any geometry entity should inherit from this class.
     """
 
+    __available_sampling_modes = ["random", "grid", "lh", "chebyshev", "latin"]
+
+    @property
+    @abstractmethod
+    def sample_modes(self):
+        """
+        Abstract method returing available samples modes for the Domain.
+        """
+        pass
+
+    @sample_modes.setter
+    def sample_modes(self, values):
+        """
+        TODO
+        """
+        if not isinstance(values, (list, tuple)):
+            values = [values]
+        for value in values:
+            if value not in DomainInterface.__available_sampling_modes:
+                raise TypeError(f"mode {value} not valid. Expected at least "
+                                "one in "
+                                f"{DomainInterface.__available_sampling_modes}."
+                                )
+
     @abstractmethod
     def sample(self):
         """
diff --git a/pina/domain/ellipsoid.py b/pina/domain/ellipsoid.py
index eda42d02b..41a734806 100644
--- a/pina/domain/ellipsoid.py
+++ b/pina/domain/ellipsoid.py
@@ -39,6 +39,7 @@ def __init__(self, ellipsoid_dict, sample_surface=False):
         self.range_ = {}
         self._centers = None
         self._axis = None
+        self.sample_modes = "random"
 
         # checking consistency
         check_consistency(sample_surface, bool)
@@ -286,7 +287,7 @@ def _single_points_sample(n, variables):
         if variables == "all":
             variables = self.variables
 
-        if mode in ["random"]:
+        if mode in self.sample_modes:
             return _Nd_sampler(n, mode, variables).extract(variables)
         else:
             raise NotImplementedError(f"mode={mode} is not implemented.")
diff --git a/pina/domain/exclusion_domain.py b/pina/domain/exclusion_domain.py
index ed63db314..4fc582cef 100644
--- a/pina/domain/exclusion_domain.py
+++ b/pina/domain/exclusion_domain.py
@@ -76,7 +76,7 @@ def sample(self, n, mode="random", variables="all"):
                 5
 
         """
-        if mode != "random":
+        if mode != self.sample_modes:
             raise NotImplementedError(
                 f"{mode} is not a valid mode for sampling."
             )
diff --git a/pina/domain/intersection_domain.py b/pina/domain/intersection_domain.py
index b40d36950..b580f21c3 100644
--- a/pina/domain/intersection_domain.py
+++ b/pina/domain/intersection_domain.py
@@ -78,7 +78,7 @@ def sample(self, n, mode="random", variables="all"):
                 5
 
         """
-        if mode != "random":
+        if mode != self.sample_modes:
             raise NotImplementedError(
                 f"{mode} is not a valid mode for sampling."
             )
diff --git a/pina/domain/operation_interface.py b/pina/domain/operation_interface.py
index acd4cf44b..edf2d484f 100644
--- a/pina/domain/operation_interface.py
+++ b/pina/domain/operation_interface.py
@@ -24,6 +24,9 @@ def __init__(self, geometries):
         # assign geometries
         self._geometries = geometries
 
+        # sampling mode, for now random is the only available
+        self.sample_modes = "random"
+
     @property
     def geometries(self):
         """
@@ -65,4 +68,4 @@ def _check_dimensions(self, geometries):
             if geometry.variables != geometries[0].variables:
                 raise NotImplementedError(
                     f"The geometries need to have same dimensions and labels."
-                )
+                )
\ No newline at end of file
diff --git a/pina/domain/simplex.py b/pina/domain/simplex.py
index 704a9cb3f..ff83dc432 100644
--- a/pina/domain/simplex.py
+++ b/pina/domain/simplex.py
@@ -74,6 +74,9 @@ def __init__(self, simplex_matrix, sample_surface=False):
         # build cartesian_bound
         self._cartesian_bound = self._build_cartesian(self._vertices_matrix)
 
+        # sampling mode
+        self.sample_modes = "random"
+
     @property
     def variables(self):
         return sorted(self._vertices_matrix.labels)
@@ -231,12 +234,13 @@ def sample(self, n, mode="random", variables="all"):
             in ``variables``.
         """
 
-        if variables == "all":
-            variables = self.variables
-        elif isinstance(variables, (list, tuple)):
-            variables = sorted(variables)
+        # if variables == "all": # TODO: check if this is necessary (from 0.1)
+        #     variables = self.variables
+        # elif isinstance(variables, (list, tuple)):
+        #     variables = sorted(variables)
 
-        if mode in ["random"]:
+        # if mode in ["random"]:
+        if mode in self.sample_modes:
             if self._sample_surface:
                 sample_pts = self._sample_boundary_randomly(n)
             else:
diff --git a/pina/domain/union_domain.py b/pina/domain/union_domain.py
index da2ead90d..07a18f239 100644
--- a/pina/domain/union_domain.py
+++ b/pina/domain/union_domain.py
@@ -32,6 +32,9 @@ def __init__(self, geometries):
 
         """
         super().__init__(geometries)
+        self.sample_modes = list(
+            set([geom.sample_modes for geom in geometries])
+            )
 
     def is_inside(self, point, check_border=False):
         """

From 4e19b50e817ca799ce14477198b91c8154cfc518 Mon Sep 17 00:00:00 2001
From: Dario Coscia <dariocos99@gmail.com>
Date: Fri, 4 Oct 2024 13:57:18 +0200
Subject: [PATCH 19/55] * Adding Collector for handling data
 sampling/collection before dataset/dataloader * Modify domain by adding
 sample_mode, variables as property * Small change concatenate -> cat in
 lno/avno * Create different factory classes for conditions

---
 pina/collector.py                           |  72 +++++++
 pina/condition/condition.py                 |  11 +-
 pina/condition/condition_interface.py       |  15 +-
 pina/condition/data_condition.py            |  20 +-
 pina/condition/domain_equation_condition.py |  26 +--
 pina/condition/input_equation_condition.py  |  24 +--
 pina/condition/input_output_condition.py    |  21 +-
 pina/domain/cartesian.py                    |   5 +-
 pina/domain/domain_interface.py             |  14 +-
 pina/domain/ellipsoid.py                    |   5 +-
 pina/domain/operation_interface.py          |   5 +-
 pina/domain/simplex.py                      |   7 +-
 pina/domain/union_domain.py                 |   9 +-
 pina/model/avno.py                          |   6 +-
 pina/model/lno.py                           |   4 +-
 pina/problem/abstract_problem.py            | 211 ++++----------------
 pina/problem/inverse_problem.py             |  14 ++
 tests/test_problem.py                       |  31 +--
 18 files changed, 224 insertions(+), 276 deletions(-)
 create mode 100644 pina/collector.py

diff --git a/pina/collector.py b/pina/collector.py
new file mode 100644
index 000000000..fa3247e2e
--- /dev/null
+++ b/pina/collector.py
@@ -0,0 +1,72 @@
+from .utils import check_consistency, merge_tensors
+
+class Collector:
+    def __init__(self, problem):
+        self.problem = problem                                                  # hook Collector <-> Problem
+        self.data_collections = {name : {} for name in self.problem.conditions} # collection of data
+        self.is_conditions_ready = {
+            name : False for name in self.problem.conditions}                   # names of the conditions that need to be sampled
+        self.full = False                                                       # collector full, all points for all conditions are given and the data are ready to be used in trainig
+        
+    @property
+    def full(self):
+        return all(self.is_conditions_ready.values())
+    
+    @full.setter
+    def full(self, value):
+        check_consistency(value, bool)
+        self._full = value
+
+    @property
+    def problem(self):
+        return self._problem
+    
+    @problem.setter
+    def problem(self, value):
+        self._problem = value
+
+    def store_fixed_data(self):
+        # loop over all conditions
+        for condition_name, condition in self.problem.conditions.items():
+            # if the condition is not ready and domain is not attribute
+            # of condition, we get and store the data
+            if (not self.is_conditions_ready[condition_name]) and (not hasattr(condition, "domain")):
+                # get data
+                keys = condition.__slots__
+                values = [getattr(condition, name) for name in keys]
+                self.data_collections[condition_name] = dict(zip(keys, values))
+                # condition now is ready
+                self.is_conditions_ready[condition_name] = True
+
+    def store_sample_domains(self, n, mode, variables, sample_locations):
+        # loop over all locations
+        for loc in sample_locations:
+            # get condition
+            condition = self.problem.conditions[loc]
+            keys = ["input_points", "equation"]
+            # if the condition is not ready, we get and store the data
+            if (not self.is_conditions_ready[loc]):
+                # if it is the first time we sample
+                if not self.data_collections[loc]:
+                    already_sampled = []
+                # if we have sampled the condition but not all variables
+                else:
+                    already_sampled = [self.data_collections[loc].input_points]
+            # if the condition is ready but we want to sample again
+            else:
+                self.is_conditions_ready[loc] = False
+                already_sampled = []
+
+            # get the samples
+            samples = [
+                condition.domain.sample(n=n, mode=mode, variables=variables)
+            ] + already_sampled
+            pts = merge_tensors(samples)
+            if (
+                sorted(self.data_collections[loc].input_points.labels) 
+                ==
+                sorted(self.problem.input_variables)
+                ):
+                self.is_conditions_ready[loc] = True
+                values = [pts, condition.equation]
+                self.data_collections[loc] = dict(zip(keys, values))
\ No newline at end of file
diff --git a/pina/condition/condition.py b/pina/condition/condition.py
index ddc722f0c..09180cc6e 100644
--- a/pina/condition/condition.py
+++ b/pina/condition/condition.py
@@ -39,11 +39,10 @@ class Condition:
 
     __slots__ = list(
                     set(
-                        InputOutputPointsCondition.__slots__,
-                        InputPointsEquationCondition.__slots__,
-                        DomainEquationCondition.__slots__,
+                        InputOutputPointsCondition.__slots__ + 
+                        InputPointsEquationCondition.__slots__ +
+                        DomainEquationCondition.__slots__ +
                         DataConditionInterface.__slots__
-
                     )
                 )
 
@@ -51,8 +50,8 @@ def __new__(cls, *args, **kwargs):
         
         if len(args) != 0:
             raise ValueError(
-                f"Condition takes only the following keyword '
-                'arguments: {Condition.__slots__}."
+                "Condition takes only the following keyword "
+                f"arguments: {Condition.__slots__}."
             )
     
         sorted_keys = sorted(kwargs.keys()) 
diff --git a/pina/condition/condition_interface.py b/pina/condition/condition_interface.py
index f380dcf68..e668b397b 100644
--- a/pina/condition/condition_interface.py
+++ b/pina/condition/condition_interface.py
@@ -1,18 +1,27 @@
 
 from abc import ABCMeta
 
-
 class ConditionInterface(metaclass=ABCMeta):
 
     condition_types = ['physical', 'supervised', 'unsupervised']
-    def __init__(self):
+
+    def __init__(self, *args, **wargs):
         self._condition_type = None
+        self._problem = None
+
+    @property
+    def problem(self):
+        return self._problem
+    
+    @problem.setter
+    def problem(self, value):
+        self._problem = value
 
     @property
     def condition_type(self):
         return self._condition_type
     
-    @condition_type.setattr
+    @condition_type.setter
     def condition_type(self, values):
         if not isinstance(values, (list, tuple)):
             values = [values]
diff --git a/pina/condition/data_condition.py b/pina/condition/data_condition.py
index 259eb5669..b9fe1ede1 100644
--- a/pina/condition/data_condition.py
+++ b/pina/condition/data_condition.py
@@ -24,21 +24,7 @@ def __init__(self, data, conditionalvariable=None):
         self.conditionalvariable = conditionalvariable
         self.condition_type = 'unsupervised'
 
-    @property
-    def data(self):
-        return self._data
-    
-    @data.setter
-    def data(self, value):
-        check_consistency(value, (LabelTensor, Graph, torch.Tensor))
-        self._data = value
-
-    @property
-    def conditionalvariable(self):
-        return self._conditionalvariable
-    
-    @data.setter
-    def conditionalvariable(self, value):
-        if value is not None:
+    def __setattr__(self, key, value):
+        if (key == 'data') or (key == 'conditionalvariable'):
             check_consistency(value, (LabelTensor, Graph, torch.Tensor))
-        self._data = value
\ No newline at end of file
+            DataConditionInterface.__dict__[key].__set__(self, value)
\ No newline at end of file
diff --git a/pina/condition/domain_equation_condition.py b/pina/condition/domain_equation_condition.py
index 9838ad7e5..f0ef8e07d 100644
--- a/pina/condition/domain_equation_condition.py
+++ b/pina/condition/domain_equation_condition.py
@@ -1,8 +1,6 @@
 import torch
 
 from .condition_interface import ConditionInterface
-from ..label_tensor import LabelTensor
-from ..graph import Graph
 from ..utils import check_consistency
 from ..domain import DomainInterface
 from ..equation.equation_interface import EquationInterface
@@ -24,20 +22,10 @@ def __init__(self, domain, equation):
         self.equation = equation
         self.condition_type = 'physics'
 
-    @property
-    def domain(self):
-        return self._domain
-    
-    @domain.setter
-    def domain(self, value):
-        check_consistency(value, (DomainInterface))
-        self._domain = value
-
-    @property
-    def equation(self):
-        return self._equation
-    
-    @equation.setter
-    def equation(self, value):
-        check_consistency(value, (EquationInterface))
-        self._equation = value
\ No newline at end of file
+    def __setattr__(self, key, value):
+        if key == 'domain':
+            check_consistency(value, (DomainInterface))
+            DomainEquationCondition.__dict__[key].__set__(self, value)
+        elif key == 'equation':
+            check_consistency(value, (EquationInterface))
+            DomainEquationCondition.__dict__[key].__set__(self, value)
\ No newline at end of file
diff --git a/pina/condition/input_equation_condition.py b/pina/condition/input_equation_condition.py
index c4b9f8dba..f77b025dc 100644
--- a/pina/condition/input_equation_condition.py
+++ b/pina/condition/input_equation_condition.py
@@ -23,20 +23,10 @@ def __init__(self, input_points, equation):
         self.equation = equation
         self.condition_type = 'physics'
 
-    @property
-    def input_points(self):
-        return self._input_points
-    
-    @input_points.setter
-    def input_points(self, value):
-        check_consistency(value, (LabelTensor)) # for now only labeltensors, we need labels for the operators!
-        self._input_points = value
-
-    @property
-    def equation(self):
-        return self._equation
-    
-    @equation.setter
-    def equation(self, value):
-        check_consistency(value, (EquationInterface))
-        self._equation = value
\ No newline at end of file
+    def __setattr__(self, key, value):
+        if key == 'input_points':
+            check_consistency(value, (LabelTensor)) # for now only labeltensors, we need labels for the operators!
+            InputPointsEquationCondition.__dict__[key].__set__(self, value)
+        elif key == 'equation':
+            check_consistency(value, (EquationInterface))
+            InputPointsEquationCondition.__dict__[key].__set__(self, value)
\ No newline at end of file
diff --git a/pina/condition/input_output_condition.py b/pina/condition/input_output_condition.py
index fd6b7a034..70388b308 100644
--- a/pina/condition/input_output_condition.py
+++ b/pina/condition/input_output_condition.py
@@ -23,20 +23,7 @@ def __init__(self, input_points, output_points):
         self.output_points = output_points
         self.condition_type = ['supervised', 'physics']
 
-    @property
-    def input_points(self):
-        return self._input_points
-    
-    @input_points.setter
-    def input_points(self, value):
-        check_consistency(value, (LabelTensor, Graph, torch.Tensor))
-        self._input_points = value
-
-    @property
-    def output_points(self):
-        return self._output_points
-    
-    @output_points.setter
-    def output_points(self, value):
-        check_consistency(value, (LabelTensor, Graph, torch.Tensor))
-        self._output_points = value
\ No newline at end of file
+    def __setattr__(self, key, value):
+        if (key == 'input_points') or (key == 'output_points'):
+            check_consistency(value, (LabelTensor, Graph, torch.Tensor))
+            InputOutputPointsCondition.__dict__[key].__set__(self, value)
\ No newline at end of file
diff --git a/pina/domain/cartesian.py b/pina/domain/cartesian.py
index e2e066ae8..c605e331c 100644
--- a/pina/domain/cartesian.py
+++ b/pina/domain/cartesian.py
@@ -21,7 +21,6 @@ def __init__(self, cartesian_dict):
         """
         self.fixed_ = {}
         self.range_ = {}
-        self.sample_modes = ["random", "grid", "lh", "chebyshev", "latin"]
 
         for k, v in cartesian_dict.items():
             if isinstance(v, (int, float)):
@@ -31,6 +30,10 @@ def __init__(self, cartesian_dict):
             else:
                 raise TypeError
 
+    @property
+    def sample_modes(self):
+        return ["random", "grid", "lh", "chebyshev", "latin"]
+    
     @property
     def variables(self):
         """Spatial variables.
diff --git a/pina/domain/domain_interface.py b/pina/domain/domain_interface.py
index 208bfdfbc..4fa70a2ba 100644
--- a/pina/domain/domain_interface.py
+++ b/pina/domain/domain_interface.py
@@ -9,7 +9,7 @@ class DomainInterface(metaclass=ABCMeta):
     Any geometry entity should inherit from this class.
     """
 
-    __available_sampling_modes = ["random", "grid", "lh", "chebyshev", "latin"]
+    available_sampling_modes = ["random", "grid", "lh", "chebyshev", "latin"]
 
     @property
     @abstractmethod
@@ -19,6 +19,14 @@ def sample_modes(self):
         """
         pass
 
+    @property
+    @abstractmethod
+    def variables(self):
+        """
+        Abstract method returing Domain variables.
+        """
+        pass
+
     @sample_modes.setter
     def sample_modes(self, values):
         """
@@ -27,10 +35,10 @@ def sample_modes(self, values):
         if not isinstance(values, (list, tuple)):
             values = [values]
         for value in values:
-            if value not in DomainInterface.__available_sampling_modes:
+            if value not in DomainInterface.available_sampling_modes:
                 raise TypeError(f"mode {value} not valid. Expected at least "
                                 "one in "
-                                f"{DomainInterface.__available_sampling_modes}."
+                                f"{DomainInterface.available_sampling_modes}."
                                 )
 
     @abstractmethod
diff --git a/pina/domain/ellipsoid.py b/pina/domain/ellipsoid.py
index 41a734806..e56e86c66 100644
--- a/pina/domain/ellipsoid.py
+++ b/pina/domain/ellipsoid.py
@@ -39,7 +39,6 @@ def __init__(self, ellipsoid_dict, sample_surface=False):
         self.range_ = {}
         self._centers = None
         self._axis = None
-        self.sample_modes = "random"
 
         # checking consistency
         check_consistency(sample_surface, bool)
@@ -72,6 +71,10 @@ def __init__(self, ellipsoid_dict, sample_surface=False):
             self._centers = dict(zip(self.range_.keys(), centers.tolist()))
             self._axis = dict(zip(self.range_.keys(), ellipsoid_axis.tolist()))
 
+    @property
+    def sample_modes(self):
+        return ["random"]
+    
     @property
     def variables(self):
         """Spatial variables.
diff --git a/pina/domain/operation_interface.py b/pina/domain/operation_interface.py
index edf2d484f..a1efec91f 100644
--- a/pina/domain/operation_interface.py
+++ b/pina/domain/operation_interface.py
@@ -24,8 +24,9 @@ def __init__(self, geometries):
         # assign geometries
         self._geometries = geometries
 
-        # sampling mode, for now random is the only available
-        self.sample_modes = "random"
+    @property
+    def sample_modes(self):
+        return ["random"]
 
     @property
     def geometries(self):
diff --git a/pina/domain/simplex.py b/pina/domain/simplex.py
index ff83dc432..3d33bff8d 100644
--- a/pina/domain/simplex.py
+++ b/pina/domain/simplex.py
@@ -74,9 +74,10 @@ def __init__(self, simplex_matrix, sample_surface=False):
         # build cartesian_bound
         self._cartesian_bound = self._build_cartesian(self._vertices_matrix)
 
-        # sampling mode
-        self.sample_modes = "random"
-
+    @property
+    def sample_modes(self):
+        return ["random"]
+       
     @property
     def variables(self):
         return sorted(self._vertices_matrix.labels)
diff --git a/pina/domain/union_domain.py b/pina/domain/union_domain.py
index 07a18f239..bd7fa56cb 100644
--- a/pina/domain/union_domain.py
+++ b/pina/domain/union_domain.py
@@ -32,9 +32,16 @@ def __init__(self, geometries):
 
         """
         super().__init__(geometries)
+
+    @property
+    def sample_modes(self):
         self.sample_modes = list(
-            set([geom.sample_modes for geom in geometries])
+            set([geom.sample_modes for geom in self.geometries])
             )
+    
+    @property
+    def variables(self):
+        return list(set([geom.variables for geom in self.geometries]))
 
     def is_inside(self, point, check_border=False):
         """
diff --git a/pina/model/avno.py b/pina/model/avno.py
index 2ac3b3f7e..e27ce1f1e 100644
--- a/pina/model/avno.py
+++ b/pina/model/avno.py
@@ -1,7 +1,7 @@
 """Module Averaging Neural Operator."""
 
 import torch
-from torch import nn, concatenate
+from torch import nn, cat
 from .layers import AVNOBlock
 from .base_no import KernelNeuralOperator
 from pina.utils import check_consistency
@@ -110,9 +110,9 @@ def forward(self, x):
         """
         points_tmp = x.extract(self.coordinates_indices)
         new_batch = x.extract(self.field_indices)
-        new_batch = concatenate((new_batch, points_tmp), dim=-1)
+        new_batch = cat((new_batch, points_tmp), dim=-1)
         new_batch = self._lifting_operator(new_batch)
         new_batch = self._integral_kernels(new_batch)
-        new_batch = concatenate((new_batch, points_tmp), dim=-1)
+        new_batch = cat((new_batch, points_tmp), dim=-1)
         new_batch = self._projection_operator(new_batch)
         return new_batch
diff --git a/pina/model/lno.py b/pina/model/lno.py
index 077a6b929..b09f1aec1 100644
--- a/pina/model/lno.py
+++ b/pina/model/lno.py
@@ -1,7 +1,7 @@
 """Module LowRank Neural Operator."""
 
 import torch
-from torch import nn, concatenate
+from torch import nn, cat
 
 from pina.utils import check_consistency
 
@@ -145,4 +145,4 @@ def forward(self, x):
         for module in self._integral_kernels:
             x = module(x, coords)
         # projecting
-        return self._projection_operator(concatenate((x, coords), dim=-1))
+        return self._projection_operator(cat((x, coords), dim=-1))
diff --git a/pina/problem/abstract_problem.py b/pina/problem/abstract_problem.py
index 200092e1b..8894207bf 100644
--- a/pina/problem/abstract_problem.py
+++ b/pina/problem/abstract_problem.py
@@ -1,12 +1,11 @@
 """ Module for AbstractProblem class """
 
 from abc import ABCMeta, abstractmethod
-from ..utils import merge_tensors, check_consistency
+from ..utils import check_consistency
+from ..domain import DomainInterface
+from ..condition.domain_equation_condition import DomainEquationCondition
+from ..collector import Collector
 from copy import deepcopy
-import torch
-
-from .. import LabelTensor
-
 
 class AbstractProblem(metaclass=ABCMeta):
     """
@@ -20,27 +19,25 @@ class AbstractProblem(metaclass=ABCMeta):
 
     def __init__(self):
 
-        self._discretized_domains = {}
-
-        for name, domain in self.domains.items():
-            if isinstance(domain, (torch.Tensor, LabelTensor)):
-                self._discretized_domains[name] = domain
+        # create collector to manage problem data
+        self.collector = Collector(self)
 
+        # create hook conditions <-> problems
         for condition_name in self.conditions:
-            self.conditions[condition_name].set_problem(self)
+            self.conditions[condition_name].problem = self
 
-        # # variable storing all points
-        self.input_pts = {}
-
-        # # varible to check if sampling is done. If no location
-        # # element is presented in Condition this variable is set to true
-        # self._have_sampled_points = {}
-        for condition_name in self.conditions:
-            self._discretized_domains[condition_name] = False
+        # store in collector all the available fixed points
+        # note that some points could not be stored at this stage (e.g. when
+        # sampling locations). To check that all data points are ready for
+        # training all type self.collector.full, which returns true if all
+        # points are ready.
+        self.collector.store_fixed_data()
 
-        # # put in self.input_pts all the points that we don't need to sample
-        self._span_condition_points()
 
+    @property
+    def input_pts(self):
+        return self.collector.data_collections
+    
     def __deepcopy__(self, memo):
         """
         Implements deepcopy for the
@@ -85,19 +82,6 @@ def input_variables(self):
     def input_variables(self, variables):
         raise RuntimeError
 
-    @property
-    @abstractmethod
-    def domains(self):
-        """
-        The domain(s) where the conditions of the AbstractProblem are valid.
-        If more than one domain type is passed, a list of Location is
-        retured.
-
-        :return: the domain(s) of ``self``
-        :rtype: list[Location]
-        """
-        pass
-
     @property
     @abstractmethod
     def output_variables(self):
@@ -114,34 +98,8 @@ def conditions(self):
         """
         return self._conditions
 
-        
-
-    def _span_condition_points(self):
-        """
-        Simple function to get the condition points
-        """
-        for condition_name in self.conditions:
-            condition = self.conditions[condition_name]
-            if hasattr(condition, "input_points"):
-                samples = condition.input_points
-                self.input_pts[condition_name] = samples
-                self._discretized_domains[condition_name] = True
-            if hasattr(self, "unknown_parameter_domain"):
-                # initialize the unknown parameters of the inverse problem given
-                # the domain the user gives
-                self.unknown_parameters = {}
-                for i, var in enumerate(self.unknown_variables):
-                    range_var = self.unknown_parameter_domain.range_[var]
-                    tensor_var = (
-                        torch.rand(1, requires_grad=True) * range_var[1]
-                        + range_var[0]
-                    )
-                    self.unknown_parameters[var] = torch.nn.Parameter(
-                        tensor_var
-                    )
-
     def discretise_domain(
-        self, n, mode="random", variables="all", domains="all"
+        self, n, mode="random", variables="all", locations="all"
     ):
         """
         Generate a set of points to span the `Location` of all the conditions of
@@ -172,119 +130,38 @@ def discretise_domain(
             ``CartesianDomain``.
         """
 
-        # check consistecy n
+        # check consistecy n, mode, variables, locations
         check_consistency(n, int)
-
-        # check consistency mode
         check_consistency(mode, str)
-        if mode not in ["random", "grid", "lh", "chebyshev", "latin"]:
+        check_consistency(variables, str)
+        check_consistency(locations, str)
+
+        # check correct sampling mode
+        if mode not in DomainInterface.available_sampling_modes:
             raise TypeError(f"mode {mode} not valid.")
 
-        # check consistency variables
+        # check correct variables
         if variables == "all":
             variables = self.input_variables
-        else:
-            check_consistency(variables, str)
-
-        if sorted(variables) != sorted(self.input_variables):
-            TypeError(
-                f"Wrong variables for sampling. Variables ",
-                f"should be in {self.input_variables}.",
-            )
-
-        # # check consistency location # TODO: check if this is needed (from 0.1)
-        # locations_to_sample = [
-        #     condition
-        #     for condition in self.conditions
-        #     if hasattr(self.conditions[condition], "location")
-        # ]
-        # if locations == "all":
-        #     # only locations that can be sampled
-        #     locations = locations_to_sample
-        # else:
-        #     check_consistency(locations, str)
-
-        # if sorted(locations) != sorted(locations_to_sample):
-        if domains == "all":
-            domains = [condition for condition in self.conditions]
-        else:
-            check_consistency(domains, str)
-        print(domains)
-        if sorted(domains) != sorted(self.conditions):
-            TypeError(
-                f"Wrong locations for sampling. Location ",
-                f"should be in {locations_to_sample}.",
-            )
-
-        # sampling
-        for d in domains:
-            condition = self.conditions[d]
-
-            # we try to check if we have already sampled
-            try:
-                already_sampled = [self.input_pts[d]]
-            # if we have not sampled, a key error is thrown
-            except KeyError:
-                already_sampled = []
-
-            # if we have already sampled fully the condition
-            # but we want to sample again we set already_sampled
-            # to an empty list since we need to sample again, and
-            # self._have_sampled_points to False.
-            if self._discretized_domains[d]:
-                already_sampled = []
-                self._discretized_domains[d] = False
-            print(condition.domain)
-            print(d)
-            # build samples
-            samples = [
-                self.domains[d].sample(n=n, mode=mode, variables=variables)
-            ] + already_sampled
-            pts = merge_tensors(samples)
-            self.input_pts[d] = pts
-
-            # the condition is sampled if input_pts contains all labels
-            if sorted(self.input_pts[d].labels) == sorted(
-                self.input_variables
-            ):
-                # self._have_sampled_points[location] = True
-                # self.input_pts[location] = self.input_pts[location].extract(
-                #     sorted(self.input_variables)
-                # )
-                self._have_sampled_points[d] = True
-
-    def add_points(self, new_points):
-        """
-        Adding points to the already sampled points.
-
-        :param dict new_points: a dictionary with key the location to add the points
-            and values the torch.Tensor points.
-        """
-
-        if sorted(new_points.keys()) != sorted(self.conditions):
-            TypeError(
-                f"Wrong locations for new points. Location ",
-                f"should be in {self.conditions}.",
-            )
-
-        for location in new_points.keys():
-            # extract old and new points
-            old_pts = self.input_pts[location]
-            new_pts = new_points[location]
-
-            # if they don't have the same variables error
-            if sorted(old_pts.labels) != sorted(new_pts.labels):
+        for variable in variables:
+            if variable not in self.input_variables:
                 TypeError(
-                    f"Not matching variables for old and new points "
-                    f"in condition {location}."
+                    f"Wrong variables for sampling. Variables ",
+                    f"should be in {self.input_variables}.",
                 )
-            if old_pts.labels != new_pts.labels:
-                new_pts = torch.hstack(
-                    [new_pts.extract([i]) for i in old_pts.labels]
+
+        # check correct location
+        if locations == "all":
+            locations = [name for name in self.conditions.keys()]
+        else:
+            if not isinstance(locations, (list)):
+                locations = [locations]
+        for loc in locations:
+            if not isinstance(self.conditions[loc], DomainEquationCondition):
+                raise TypeError(
+                    f"Wrong locations passed, locations for sampling "
+                    f"should be in {[loc for loc in locations if not isinstance(self.conditions[loc], DomainEquationCondition)]}.",
                 )
-                new_pts.labels = old_pts.labels
 
-            # merging
-            merged_pts = torch.vstack([old_pts, new_pts])
-            merged_pts.labels = old_pts.labels
-            self.input_pts[location] = merged_pts
\ No newline at end of file
+        # store data
+        self.collector.store_sample_domains(n, mode, variables, locations)
diff --git a/pina/problem/inverse_problem.py b/pina/problem/inverse_problem.py
index 5a83566ae..51cbd3ca6 100644
--- a/pina/problem/inverse_problem.py
+++ b/pina/problem/inverse_problem.py
@@ -45,6 +45,20 @@ class InverseProblem(AbstractProblem):
         >>>         'data': Condition(CartesianDomain({'x': [0, 1]}), Equation(solution_data))
     """
 
+    def __init__(self):
+        super().__init__()
+        # storing unknown_parameters for optimization
+        self.unknown_parameters = {}
+        for i, var in enumerate(self.unknown_variables):
+            range_var = self.unknown_parameter_domain.range_[var]
+            tensor_var = (
+                torch.rand(1, requires_grad=True) * range_var[1]
+                + range_var[0]
+            )
+            self.unknown_parameters[var] = torch.nn.Parameter(
+                tensor_var
+            )
+
     @abstractmethod
     def unknown_parameter_domain(self):
         """
diff --git a/tests/test_problem.py b/tests/test_problem.py
index 0db508fa3..77871b1a1 100644
--- a/tests/test_problem.py
+++ b/tests/test_problem.py
@@ -27,31 +27,31 @@ class Poisson(SpatialProblem):
 
     conditions = {
         'gamma1':
-        Condition(location=CartesianDomain({
+        Condition(domain=CartesianDomain({
             'x': [0, 1],
             'y': 1
         }),
                   equation=FixedValue(0.0)),
         'gamma2':
-        Condition(location=CartesianDomain({
+        Condition(domain=CartesianDomain({
             'x': [0, 1],
             'y': 0
         }),
                   equation=FixedValue(0.0)),
         'gamma3':
-        Condition(location=CartesianDomain({
+        Condition(domain=CartesianDomain({
             'x': 1,
             'y': [0, 1]
         }),
                   equation=FixedValue(0.0)),
         'gamma4':
-        Condition(location=CartesianDomain({
+        Condition(domain=CartesianDomain({
             'x': 0,
             'y': [0, 1]
         }),
                   equation=FixedValue(0.0)),
         'D':
-        Condition(location=CartesianDomain({
+        Condition(domain=CartesianDomain({
             'x': [0, 1],
             'y': [0, 1]
         }),
@@ -67,6 +67,10 @@ def poisson_sol(self, pts):
     truth_solution = poisson_sol
 
 
+# make the problem
+poisson_problem = Poisson()
+print(poisson_problem.input_pts)
+
 def test_discretise_domain():
     n = 10
     poisson_problem = Poisson()
@@ -90,15 +94,14 @@ def test_discretise_domain():
     assert poisson_problem.input_pts['D'].shape[0] == n
 
 
-def test_sampling_few_variables():
-    n = 10
-    poisson_problem = Poisson()
-    poisson_problem.discretise_domain(n,
-                                      'grid',
-                                      locations=['D'],
-                                      variables=['x'])
-    assert poisson_problem.input_pts['D'].shape[1] == 1
-    assert poisson_problem._have_sampled_points['D'] is False
+# def test_sampling_few_variables():
+#     n = 10
+#     poisson_problem.discretise_domain(n,
+#                                       'grid',
+#                                       locations=['D'],
+#                                       variables=['x'])
+#     assert poisson_problem.input_pts['D'].shape[1] == 1
+#     assert poisson_problem._have_sampled_points['D'] is False
 
 
 def test_variables_correct_order_sampling():

From 4dea84714a37da3d2f67536f4f2f62c1d97d01c7 Mon Sep 17 00:00:00 2001
From: Dario Coscia <dariocos99@gmail.com>
Date: Fri, 4 Oct 2024 15:39:20 +0200
Subject: [PATCH 20/55] update condition_interface

---
 pina/condition/condition_interface.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/pina/condition/condition_interface.py b/pina/condition/condition_interface.py
index e668b397b..52699b66e 100644
--- a/pina/condition/condition_interface.py
+++ b/pina/condition/condition_interface.py
@@ -3,7 +3,7 @@
 
 class ConditionInterface(metaclass=ABCMeta):
 
-    condition_types = ['physical', 'supervised', 'unsupervised']
+    condition_types = ['physics', 'supervised', 'unsupervised']
 
     def __init__(self, *args, **wargs):
         self._condition_type = None

From 882ff0c293136cacb8704e2c90e0f7d4920a4d84 Mon Sep 17 00:00:00 2001
From: Dario Coscia <93731561+dario-coscia@users.noreply.github.com>
Date: Fri, 4 Oct 2024 15:59:09 +0200
Subject: [PATCH 21/55] Filippo0.2 (#361)

* Add summation and remove deepcopy (only for tensors) in LabelTensor class
* Update operators for compatibility with updated LabelTensor implementation
* Implement labels.setter in LabelTensor class
* Update LabelTensor

---------

Co-authored-by: FilippoOlivo <filippo@filippoolivo.com>
---
 pina/label_tensor.py       | 108 ++++++++++++++++++++++++++---------
 pina/operators.py          | 114 +++++++++++++++++++++----------------
 tests/test_label_tensor.py |  58 ++++++++++++++++++-
 tests/test_operators.py    |  24 ++++----
 4 files changed, 212 insertions(+), 92 deletions(-)

diff --git a/pina/label_tensor.py b/pina/label_tensor.py
index 7646dd8a3..08e0b03e6 100644
--- a/pina/label_tensor.py
+++ b/pina/label_tensor.py
@@ -35,14 +35,34 @@ def __init__(self, x, labels):
                     {1: {"name": "space"['a', 'b', 'c'])
 
         """
-        self.labels = None
+        self.labels = labels
+
+    @property
+    def labels(self):
+        """Property decorator for labels
+
+        :return: labels of self
+        :rtype: list
+        """
+        return self._labels
+
+    @labels.setter
+    def labels(self, labels):
+        """"
+        Set properly the parameter _labels
+
+        :param labels: Labels to assign to the class variable _labels.
+        :type: labels: str | list(str) | dict
+        """
+        if hasattr(self, 'labels') is False:
+            self.init_labels()
         if isinstance(labels, dict):
-            self.update_labels(labels)
+            self.update_labels_from_dict(labels)
         elif isinstance(labels, list):
-            self.init_labels_from_list(labels)
+            self.update_labels_from_list(labels)
         elif isinstance(labels, str):
             labels = [labels]
-            self.init_labels_from_list(labels)
+            self.update_labels_from_list(labels)
         else:
             raise ValueError(f"labels must be list, dict or string.")
 
@@ -60,38 +80,38 @@ def extract(self, label_to_extract):
         if isinstance(label_to_extract, (str, int)):
             label_to_extract = [label_to_extract]
         if isinstance(label_to_extract, (tuple, list)):
-            last_dim_label = self.labels[self.tensor.ndim - 1]['dof']
+            last_dim_label = self._labels[self.tensor.ndim - 1]['dof']
             if set(label_to_extract).issubset(last_dim_label) is False:
                 raise ValueError('Cannot extract a dof which is not in the original LabelTensor')
             idx_to_extract = [last_dim_label.index(i) for i in label_to_extract]
-            new_tensor = deepcopy(self.tensor)
+            new_tensor = self.tensor
             new_tensor = new_tensor[..., idx_to_extract]
-            new_labels = deepcopy(self.labels)
+            new_labels = deepcopy(self._labels)
             last_dim_new_label = {self.tensor.ndim - 1: {
                 'dof': label_to_extract,
-                'name': self.labels[self.tensor.ndim - 1]['name']
+                'name': self._labels[self.tensor.ndim - 1]['name']
             }}
             new_labels.update(last_dim_new_label)
         elif isinstance(label_to_extract, dict):
-            new_labels = (deepcopy(self.labels))
-            new_tensor = deepcopy(self.tensor)
+            new_labels = (deepcopy(self._labels))
+            new_tensor = self.tensor
             for k, v in label_to_extract.items():
                 idx_dim = None
-                for kl, vl in self.labels.items():
+                for kl, vl in self._labels.items():
                     if vl['name'] == k:
                         idx_dim = kl
                         break
-                dim_labels = self.labels[idx_dim]['dof']
+                dim_labels = self._labels[idx_dim]['dof']
                 if isinstance(label_to_extract[k], (int, str)):
                     label_to_extract[k] = [label_to_extract[k]]
                 if set(label_to_extract[k]).issubset(dim_labels) is False:
-                    raise ValueError('Cannot extract a dof which is not in the original labeltensor')
+                    raise ValueError('Cannot extract a dof which is not in the original LabelTensor')
                 idx_to_extract = [dim_labels.index(i) for i in label_to_extract[k]]
                 indexer = [slice(None)] * idx_dim + [idx_to_extract] + [slice(None)] * (self.tensor.ndim - idx_dim - 1)
                 new_tensor = new_tensor[indexer]
                 dim_new_label = {idx_dim: {
                     'dof': label_to_extract[k],
-                    'name': self.labels[idx_dim]['name']
+                    'name': self._labels[idx_dim]['name']
                 }}
                 new_labels.update(dim_new_label)
         else:
@@ -104,7 +124,7 @@ def __str__(self):
         """
 
         s = ''
-        for key, value in self.labels.items():
+        for key, value in self._labels.items():
             s += f"{key}: {value}\n"
         s += '\n'
         s += super().__str__()
@@ -155,7 +175,7 @@ def cat(tensors, dim=0):
 
     def requires_grad_(self, mode=True):
         lt = super().requires_grad_(mode)
-        lt.labels = self.labels
+        lt.labels = self._labels
         return lt
 
     @property
@@ -181,10 +201,19 @@ def clone(self, *args, **kwargs):
         :rtype: LabelTensor
         """
 
-        out = LabelTensor(super().clone(*args, **kwargs), self.labels)
+        out = LabelTensor(super().clone(*args, **kwargs), self._labels)
         return out
 
-    def update_labels(self, labels):
+
+    def init_labels(self):
+        self._labels = {
+            idx_: {
+                'dof': range(self.tensor.shape[idx_]),
+                'name': idx_
+            } for idx_ in range(self.tensor.ndim)
+        }
+
+    def update_labels_from_dict(self, labels):
         """
         Update the internal label representation according to the values passed as input.
 
@@ -192,21 +221,16 @@ def update_labels(self, labels):
         :type labels: dict
         :raises ValueError: dof list contain duplicates or number of dof does not match with tensor shape
         """
-        self.labels = {
-            idx_: {
-                'dof': range(self.tensor.shape[idx_]),
-                'name': idx_
-            } for idx_ in range(self.tensor.ndim)
-        }
+
         tensor_shape = self.tensor.shape
         for k, v in labels.items():
             if len(v['dof']) != len(set(v['dof'])):
                 raise ValueError("dof must be unique")
             if len(v['dof']) != tensor_shape[k]:
                 raise ValueError('Number of dof does not match with tensor dimension')
-        self.labels.update(labels)
+        self._labels.update(labels)
 
-    def init_labels_from_list(self, labels):
+    def update_labels_from_list(self, labels):
         """
         Given a list of dof, this method update the internal label representation
 
@@ -214,4 +238,34 @@ def init_labels_from_list(self, labels):
         :type labels: list
         """
         last_dim_labels = {self.tensor.ndim - 1: {'dof': labels, 'name': self.tensor.ndim - 1}}
-        self.update_labels(last_dim_labels)
\ No newline at end of file
+        self.update_labels_from_dict(last_dim_labels)
+
+    @staticmethod
+    def summation(tensors):
+        if len(tensors) == 0:
+            raise ValueError('tensors list must not be empty')
+        if len(tensors) == 1:
+            return tensors[0]
+        labels = tensors[0].labels
+        for j in range(tensors[0].ndim):
+            for i in range(1, len(tensors)):
+                if labels[j] != tensors[i].labels[j]:
+                    labels.pop(j)
+                    break
+
+        data = torch.zeros(tensors[0].tensor.shape)
+        for i in range(len(tensors)):
+            data += tensors[i].tensor
+        new_tensor = LabelTensor(data, labels)
+        return new_tensor
+
+    def last_dim_dof(self):
+        return self._labels[self.tensor.ndim - 1]['dof']
+
+    def append(self, tensor, mode='std'):
+        print(self.labels)
+        print(tensor.labels)
+        if mode == 'std':
+            new_label_tensor = LabelTensor.cat([self, tensor], dim=self.tensor.ndim - 1)
+
+        return new_label_tensor
diff --git a/pina/operators.py b/pina/operators.py
index e523ed922..082d72555 100644
--- a/pina/operators.py
+++ b/pina/operators.py
@@ -1,13 +1,13 @@
 """
 Module for operators vectorize implementation. Differential operators are used to write any differential problem.
-These operators are implemented to work on different accellerators: CPU, GPU, TPU or MPS.
+These operators are implemented to work on different accelerators: CPU, GPU, TPU or MPS.
 All operators take as input a tensor onto which computing the operator, a tensor with respect
 to which computing the operator, the name of the output variables to calculate the operator
 for (in case of multidimensional functions), and the variables name on which the operator is calculated.
 """
 
 import torch
-
+from copy import deepcopy
 from pina.label_tensor import LabelTensor
 
 
@@ -49,12 +49,12 @@ def grad_scalar_output(output_, input_, d):
         :rtype: LabelTensor
         """
 
-        if len(output_.labels) != 1:
+        if len(output_.labels[output_.tensor.ndim-1]['dof']) != 1:
             raise RuntimeError("only scalar function can be differentiated")
-        if not all([di in input_.labels for di in d]):
+        if not all([di in input_.labels[input_.tensor.ndim-1]['dof'] for di in d]):
             raise RuntimeError("derivative labels missing from input tensor")
 
-        output_fieldname = output_.labels[0]
+        output_fieldname = output_.labels[output_.ndim-1]['dof'][0]
         gradients = torch.autograd.grad(
             output_,
             input_,
@@ -65,41 +65,35 @@ def grad_scalar_output(output_, input_, d):
             retain_graph=True,
             allow_unused=True,
         )[0]
-
-        gradients.labels = input_.labels
+        new_labels = deepcopy(input_.labels)
+        gradients.labels = new_labels
         gradients = gradients.extract(d)
-        gradients.labels = [f"d{output_fieldname}d{i}" for i in d]
-
+        new_labels[input_.tensor.ndim - 1]['dof'] = [f"d{output_fieldname}d{i}" for i in d]
+        gradients.labels = new_labels
         return gradients
 
     if not isinstance(input_, LabelTensor):
         raise TypeError
-
     if d is None:
-        d = input_.labels
+        d = input_.labels[input_.tensor.ndim-1]['dof']
 
     if components is None:
-        components = output_.labels
+        components = output_.labels[output_.tensor.ndim-1]['dof']
 
-    if output_.shape[1] == 1:  # scalar output ################################
+    if output_.shape[output_.ndim-1] == 1:  # scalar output ################################
 
-        if components != output_.labels:
+        if components != output_.labels[output_.tensor.ndim-1]['dof']:
             raise RuntimeError
         gradients = grad_scalar_output(output_, input_, d)
 
-    elif output_.shape[1] >= 2:  # vector output ##############################
-
+    elif output_.shape[output_.ndim-1] >= 2:  # vector output ##############################
+        tensor_to_cat = []
         for i, c in enumerate(components):
             c_output = output_.extract([c])
-            if i == 0:
-                gradients = grad_scalar_output(c_output, input_, d)
-            else:
-                gradients = gradients.append(
-                    grad_scalar_output(c_output, input_, d)
-                )
+            tensor_to_cat.append(grad_scalar_output(c_output, input_, d))
+        gradients = LabelTensor.cat(tensor_to_cat, dim=output_.tensor.ndim-1)
     else:
         raise NotImplementedError
-
     return gradients
 
 
@@ -130,27 +124,29 @@ def div(output_, input_, components=None, d=None):
         raise TypeError
 
     if d is None:
-        d = input_.labels
+        d = input_.labels[input_.tensor.ndim-1]['dof']
 
     if components is None:
-        components = output_.labels
+        components = output_.labels[output_.tensor.ndim-1]['dof']
 
-    if output_.shape[1] < 2 or len(components) < 2:
+    if output_.shape[output_.ndim-1] < 2 or len(components) < 2:
         raise ValueError("div supported only for vector fields")
 
     if len(components) != len(d):
         raise ValueError
 
     grad_output = grad(output_, input_, components, d)
-    div = torch.zeros(input_.shape[0], 1, device=output_.device)
-    labels = [None] * len(components)
+    last_dim_dof = [None] * len(components)
+    to_sum_tensors = []
     for i, (c, d) in enumerate(zip(components, d)):
         c_fields = f"d{c}d{d}"
-        div[:, 0] += grad_output.extract(c_fields).sum(axis=1)
-        labels[i] = c_fields
+        last_dim_dof[i] = c_fields
+        to_sum_tensors.append(grad_output.extract(c_fields))
 
-    div = div.as_subclass(LabelTensor)
-    div.labels = ["+".join(labels)]
+    div = LabelTensor.summation(to_sum_tensors)
+    new_labels = deepcopy(input_.labels)
+    new_labels[input_.tensor.ndim-1]['dof'] = ["+".join(last_dim_dof)]
+    div.labels = new_labels
     return div
 
 
@@ -205,10 +201,10 @@ def scalar_laplace(output_, input_, components, d):
         return result
 
     if d is None:
-        d = input_.labels
+        d = input_.labels[input_.tensor.ndim-1]['dof']
 
     if components is None:
-        components = output_.labels
+        components = output_.labels[output_.tensor.ndim-1]['dof']
 
     if method == "divgrad":
         raise NotImplementedError("divgrad not implemented as method")
@@ -218,25 +214,43 @@ def scalar_laplace(output_, input_, components, d):
 
     elif method == "std":
         if len(components) == 1:
-            result = scalar_laplace(output_, input_, components, d)
+            # result = scalar_laplace(output_, input_, components, d) # TODO check (from 0.1)
+            grad_output = grad(output_, input_, components=components, d=d)
+            to_append_tensors = []
+            for i, label in enumerate(grad_output.labels[grad_output.ndim-1]['dof']):
+                gg = grad(grad_output, input_, d=d, components=[label])
+                to_append_tensors.append(gg.extract([gg.labels[gg.tensor.ndim-1]['dof'][i]]))
             labels = [f"dd{components[0]}"]
-
+            result = LabelTensor.summation(tensors=to_append_tensors)
+            result.labels = labels
         else:
-            result = torch.empty(
-                size=(input_.shape[0], len(components)),
-                dtype=output_.dtype,
-                device=output_.device,
-            )
+    #         result = torch.empty( # TODO check (from 0.1)
+    #             size=(input_.shape[0], len(components)),
+    #             dtype=output_.dtype,
+    #             device=output_.device,
+    #         )
+    #         labels = [None] * len(components)
+    #         for idx, c in enumerate(components):
+    #             result[:, idx] = scalar_laplace(output_, input_, c, d).flatten()
+    #             labels[idx] = f"dd{c}"
+
+    # result = result.as_subclass(LabelTensor)
+    # result.labels = labels
             labels = [None] * len(components)
-            for idx, c in enumerate(components):
-                result[:, idx] = scalar_laplace(output_, input_, c, d).flatten()
-                labels[idx] = f"dd{c}"
-
-    result = result.as_subclass(LabelTensor)
-    result.labels = labels
+            to_append_tensors = [None] * len(components)
+            for idx, (ci, di) in enumerate(zip(components, d)):
+                if not isinstance(ci, list):
+                    ci = [ci]
+                if not isinstance(di, list):
+                    di = [di]
+                grad_output = grad(output_, input_, components=ci, d=di)
+                to_append_tensors[idx] = grad(grad_output, input_, d=di)
+                labels[idx] = f"dd{ci[0]}dd{di[0]}"
+            result = LabelTensor.cat(tensors=to_append_tensors, dim=output_.tensor.ndim-1)
+            result.labels = labels
     return result
 
-
+# TODO Fix advection operator
 def advection(output_, input_, velocity_field, components=None, d=None):
     """
     Perform advection operation. The operator works for vectorial functions,
@@ -258,10 +272,10 @@ def advection(output_, input_, velocity_field, components=None, d=None):
     :rtype: LabelTensor
     """
     if d is None:
-        d = input_.labels
+        d = input_.labels[input_.tensor.ndim-1]['dof']
 
     if components is None:
-        components = output_.labels
+        components = output_.labels[output_.tensor.ndim-1]['dof']
 
     tmp = (
         grad(output_, input_, components, d)
diff --git a/tests/test_label_tensor.py b/tests/test_label_tensor.py
index f87d3abb1..6ef484f0b 100644
--- a/tests/test_label_tensor.py
+++ b/tests/test_label_tensor.py
@@ -17,12 +17,14 @@
         "dof": range(20)
     }
 }
+labels_list = ['x', 'y', 'z']
 labels_all = labels_column | labels_row
 
-@pytest.mark.parametrize("labels", [labels_column, labels_row, labels_all])
+@pytest.mark.parametrize("labels", [labels_column, labels_row, labels_all, labels_list])
 def test_constructor(labels):
     LabelTensor(data, labels)
 
+
 def test_wrong_constructor():
     with pytest.raises(ValueError):
         LabelTensor(data, ['a', 'b'])
@@ -61,7 +63,6 @@ def test_extract_2D(labels_te):
     assert torch.all(torch.isclose(data[2,2].reshape(1, 1), new))
 
 def test_extract_3D():
-    labels = labels_all
     data = torch.rand(20, 3, 4)
     labels = {
         1: {
@@ -80,6 +81,7 @@ def test_extract_3D():
 
     tensor = LabelTensor(data, labels)
     new = tensor.extract(labels_te)
+    tensor2 = LabelTensor(data, labels)
     assert new.ndim == tensor.ndim
     assert new.shape[0] == 20
     assert new.shape[1] == 2
@@ -88,6 +90,10 @@ def test_extract_3D():
         data[:, 0::2, 1:4].reshape(20, 2, 3),
         new
     ))
+    assert tensor2.ndim == tensor.ndim
+    assert tensor2.shape == tensor.shape
+    assert tensor.labels == tensor2.labels
+    assert new.shape != tensor.shape
 
 def test_concatenation_3D():
     data_1 = torch.rand(20, 3, 4)
@@ -146,3 +152,51 @@ def test_concatenation_3D():
     assert lt_cat.labels[2]['dof'] == range(5)
     assert lt_cat.labels[0]['dof'] == range(20)
     assert lt_cat.labels[1]['dof'] == range(3)
+
+
+def test_summation():
+    lt1 = LabelTensor(torch.ones(20,3), labels_all)
+    lt2 = LabelTensor(torch.ones(30,3), ['x', 'y', 'z'])
+    with pytest.raises(RuntimeError):
+        LabelTensor.summation([lt1, lt2])
+    lt1 = LabelTensor(torch.ones(20,3), labels_all)
+    lt2 = LabelTensor(torch.ones(20,3), labels_all)
+    lt_sum = LabelTensor.summation([lt1, lt2])
+    assert lt_sum.ndim == lt_sum.ndim
+    assert lt_sum.shape[0] == 20
+    assert lt_sum.shape[1] == 3
+    assert lt_sum.labels == labels_all
+    assert torch.eq(lt_sum.tensor, torch.ones(20,3)*2).all()
+    lt1 = LabelTensor(torch.ones(20,3), labels_all)
+    lt2 = LabelTensor(torch.ones(20,3), labels_all)
+    lt3 = LabelTensor(torch.zeros(20, 3), labels_all)
+    lt_sum = LabelTensor.summation([lt1, lt2, lt3])
+    assert lt_sum.ndim == lt_sum.ndim
+    assert lt_sum.shape[0] == 20
+    assert lt_sum.shape[1] == 3
+    assert lt_sum.labels == labels_all
+    assert torch.eq(lt_sum.tensor, torch.ones(20,3)*2).all()
+
+def test_append_3D():
+    data_1 = torch.rand(20, 3, 4)
+    labels_1 = ['x', 'y', 'z', 'w']
+    lt1 = LabelTensor(data_1, labels_1)
+    data_2 = torch.rand(50, 3, 4)
+    labels_2 = ['x', 'y', 'z', 'w']
+    lt2 = LabelTensor(data_2, labels_2)
+    lt1 = lt1.append(lt2)
+    assert lt1.shape == (70, 3, 4)
+    assert lt1.labels[0]['dof'] == range(70)
+    assert lt1.labels[1]['dof'] == range(3)
+    assert lt1.labels[2]['dof'] == ['x', 'y', 'z', 'w']
+    data_1 = torch.rand(20, 3, 2)
+    labels_1 = ['x', 'y']
+    lt1 = LabelTensor(data_1, labels_1)
+    data_2 = torch.rand(20, 3, 2)
+    labels_2 = ['z', 'w']
+    lt2 = LabelTensor(data_2, labels_2)
+    lt1 = lt1.append(lt2, mode='cross')
+    assert lt1.shape == (20, 3, 4)
+    assert lt1.labels[0]['dof'] == range(20)
+    assert lt1.labels[1]['dof'] == range(3)
+    assert lt1.labels[2]['dof'] == ['x', 'y', 'z', 'w']
diff --git a/tests/test_operators.py b/tests/test_operators.py
index 58e90ca33..e18eaf29e 100644
--- a/tests/test_operators.py
+++ b/tests/test_operators.py
@@ -16,28 +16,29 @@ def func_scalar(x):
     return x_**2 + y_**2 + z_**2
 
 
-inp = LabelTensor(torch.rand((20, 3), requires_grad=True), ['x', 'y', 'z'])
-tensor_v = LabelTensor(func_vector(inp), ['a', 'b', 'c'])
-tensor_s = LabelTensor(func_scalar(inp).reshape(-1, 1), ['a'])
+data = torch.rand((20, 3))
+inp = LabelTensor(data, ['x', 'y', 'mu']).requires_grad_(True)
+labels = ['a', 'b', 'c']
+tensor_v = LabelTensor(func_vec(inp), labels)
+tensor_s = LabelTensor(func_scalar(inp).reshape(-1, 1), labels[0])
+
 
 def test_grad_scalar_output():
     grad_tensor_s = grad(tensor_s, inp)
     true_val = 2*inp
     assert grad_tensor_s.shape == inp.shape
-    assert grad_tensor_s.labels == [
-        f'd{tensor_s.labels[0]}d{i}' for i in inp.labels
+    assert grad_tensor_s.labels[grad_tensor_s.ndim-1]['dof'] == [
+        f'd{tensor_s.labels[tensor_s.ndim-1]["dof"][0]}d{i}' for i in inp.labels[inp.ndim-1]['dof']
     ]
     assert torch.allclose(grad_tensor_s, true_val)
 
     grad_tensor_s = grad(tensor_s, inp, d=['x', 'y'])
-    true_val = 2*inp.extract(['x', 'y'])
-    assert grad_tensor_s.shape == (inp.shape[0], 2)
-    assert grad_tensor_s.labels == [
-        f'd{tensor_s.labels[0]}d{i}' for i in ['x', 'y']
+    assert grad_tensor_s.shape == (20, 2)
+    assert grad_tensor_s.labels[grad_tensor_s.ndim-1]['dof'] == [
+        f'd{tensor_s.labels[tensor_s.ndim-1]["dof"][0]}d{i}' for i in ['x', 'y']
     ]
     assert torch.allclose(grad_tensor_s, true_val)
 
-
 def test_grad_vector_output():
     grad_tensor_v = grad(tensor_v, inp)
     true_val = torch.cat(
@@ -74,7 +75,6 @@ def test_grad_vector_output():
     ]
     assert torch.allclose(grad_tensor_v, true_val)
 
-
 def test_div_vector_output():
     div_tensor_v = div(tensor_v, inp)
     true_val = 2*torch.sum(inp, dim=1).reshape(-1,1)
@@ -88,7 +88,6 @@ def test_div_vector_output():
     assert div_tensor_v.labels == [f'dadx+dbdy']
     assert torch.allclose(div_tensor_v, true_val)
 
-
 def test_laplacian_scalar_output():
     laplace_tensor_s = laplacian(tensor_s, inp)
     true_val = 6*torch.ones_like(laplace_tensor_s)
@@ -102,7 +101,6 @@ def test_laplacian_scalar_output():
     assert laplace_tensor_s.labels == [f"dd{tensor_s.labels[0]}"]
     assert torch.allclose(laplace_tensor_s, true_val)
 
-
 def test_laplacian_vector_output():
     laplace_tensor_v = laplacian(tensor_v, inp)
     true_val = 2*torch.ones_like(tensor_v)

From 3ac905dadb0f83f587263bfd36a546bfffa0b8b5 Mon Sep 17 00:00:00 2001
From: Filippo Olivo <filippo@filippoolivo.com>
Date: Thu, 10 Oct 2024 18:26:52 +0200
Subject: [PATCH 22/55] Update of LabelTensor class and fix Simplex domain
 (#362)

*Implement new methods in LabelTensor and fix operators
---
 pina/collector.py                             |  29 ++-
 pina/condition/data_condition.py              |   4 +-
 pina/condition/domain_equation_condition.py   |   4 +-
 pina/condition/input_equation_condition.py    |   6 +-
 pina/condition/input_output_condition.py      |   4 +-
 pina/domain/difference_domain.py              |   2 +-
 pina/domain/exclusion_domain.py               |   2 +-
 pina/domain/intersection_domain.py            |   2 +-
 pina/domain/simplex.py                        |   7 +-
 pina/domain/union_domain.py                   |   5 +-
 pina/label_tensor.py                          | 244 +++++++++++++-----
 pina/operators.py                             |  75 +++---
 pina/problem/abstract_problem.py              |  13 +-
 tests/test_condition.py                       |  22 +-
 tests/test_geometry/test_simplex.py           |   1 -
 .../test_label_tensor.py                      | 126 ++++++---
 .../test_label_tensor/test_label_tensor_01.py | 117 +++++++++
 tests/test_operators.py                       |   8 +-
 tests/test_problem.py                         |  87 +++----
 19 files changed, 546 insertions(+), 212 deletions(-)
 rename tests/{ => test_label_tensor}/test_label_tensor.py (56%)
 create mode 100644 tests/test_label_tensor/test_label_tensor_01.py

diff --git a/pina/collector.py b/pina/collector.py
index fa3247e2e..0f4e9da44 100644
--- a/pina/collector.py
+++ b/pina/collector.py
@@ -1,3 +1,6 @@
+from sympy.strategies.branch import condition
+
+from . import LabelTensor
 from .utils import check_consistency, merge_tensors
 
 class Collector:
@@ -51,7 +54,7 @@ def store_sample_domains(self, n, mode, variables, sample_locations):
                     already_sampled = []
                 # if we have sampled the condition but not all variables
                 else:
-                    already_sampled = [self.data_collections[loc].input_points]
+                    already_sampled = [self.data_collections[loc]['input_points']]
             # if the condition is ready but we want to sample again
             else:
                 self.is_conditions_ready[loc] = False
@@ -63,10 +66,24 @@ def store_sample_domains(self, n, mode, variables, sample_locations):
             ] + already_sampled
             pts = merge_tensors(samples)
             if (
-                sorted(self.data_collections[loc].input_points.labels) 
-                ==
-                sorted(self.problem.input_variables)
+                set(pts.labels).issubset(sorted(self.problem.input_variables))
                 ):
-                self.is_conditions_ready[loc] = True
+                pts = pts.sort_labels()
+                if sorted(pts.labels)==sorted(self.problem.input_variables):
+                    self.is_conditions_ready[loc] = True
                 values = [pts, condition.equation]
-                self.data_collections[loc] = dict(zip(keys, values))
\ No newline at end of file
+                self.data_collections[loc] = dict(zip(keys, values))
+            else:
+                raise RuntimeError('Try to sample variables which are not in problem defined in the problem')
+
+    def add_points(self, new_points_dict):
+        """
+        Add input points to a sampled condition
+
+        :param new_points_dict: Dictonary of input points (condition_name: LabelTensor)
+        :raises RuntimeError: if at least one condition is not already sampled
+        """
+        for k,v in new_points_dict.items():
+            if not self.is_conditions_ready[k]:
+                raise RuntimeError('Cannot add points on a non sampled condition')
+            self.data_collections[k]['input_points'] = self.data_collections[k]['input_points'].vstack(v)
\ No newline at end of file
diff --git a/pina/condition/data_condition.py b/pina/condition/data_condition.py
index b9fe1ede1..d5ac63970 100644
--- a/pina/condition/data_condition.py
+++ b/pina/condition/data_condition.py
@@ -27,4 +27,6 @@ def __init__(self, data, conditionalvariable=None):
     def __setattr__(self, key, value):
         if (key == 'data') or (key == 'conditionalvariable'):
             check_consistency(value, (LabelTensor, Graph, torch.Tensor))
-            DataConditionInterface.__dict__[key].__set__(self, value)
\ No newline at end of file
+            DataConditionInterface.__dict__[key].__set__(self, value)
+        elif key in ('_condition_type', '_problem', 'problem', 'condition_type'):
+            super().__setattr__(key, value)
\ No newline at end of file
diff --git a/pina/condition/domain_equation_condition.py b/pina/condition/domain_equation_condition.py
index f0ef8e07d..ab35d202a 100644
--- a/pina/condition/domain_equation_condition.py
+++ b/pina/condition/domain_equation_condition.py
@@ -28,4 +28,6 @@ def __setattr__(self, key, value):
             DomainEquationCondition.__dict__[key].__set__(self, value)
         elif key == 'equation':
             check_consistency(value, (EquationInterface))
-            DomainEquationCondition.__dict__[key].__set__(self, value)
\ No newline at end of file
+            DomainEquationCondition.__dict__[key].__set__(self, value)
+        elif key in ('_condition_type', '_problem', 'problem', 'condition_type'):
+            super().__setattr__(key, value)
\ No newline at end of file
diff --git a/pina/condition/input_equation_condition.py b/pina/condition/input_equation_condition.py
index f77b025dc..dc12d0225 100644
--- a/pina/condition/input_equation_condition.py
+++ b/pina/condition/input_equation_condition.py
@@ -21,7 +21,7 @@ def __init__(self, input_points, equation):
         super().__init__()
         self.input_points = input_points
         self.equation = equation
-        self.condition_type = 'physics'
+        self._condition_type = 'physics'
 
     def __setattr__(self, key, value):
         if key == 'input_points':
@@ -29,4 +29,6 @@ def __setattr__(self, key, value):
             InputPointsEquationCondition.__dict__[key].__set__(self, value)
         elif key == 'equation':
             check_consistency(value, (EquationInterface))
-            InputPointsEquationCondition.__dict__[key].__set__(self, value)
\ No newline at end of file
+            InputPointsEquationCondition.__dict__[key].__set__(self, value)
+        elif key in ('_condition_type', '_problem', 'problem', 'condition_type'):
+            super().__setattr__(key, value)
\ No newline at end of file
diff --git a/pina/condition/input_output_condition.py b/pina/condition/input_output_condition.py
index 70388b308..a4fa48919 100644
--- a/pina/condition/input_output_condition.py
+++ b/pina/condition/input_output_condition.py
@@ -26,4 +26,6 @@ def __init__(self, input_points, output_points):
     def __setattr__(self, key, value):
         if (key == 'input_points') or (key == 'output_points'):
             check_consistency(value, (LabelTensor, Graph, torch.Tensor))
-            InputOutputPointsCondition.__dict__[key].__set__(self, value)
\ No newline at end of file
+            InputOutputPointsCondition.__dict__[key].__set__(self, value)
+        elif key in ('_condition_type', '_problem', 'problem', 'condition_type'):
+            super().__setattr__(key, value)
diff --git a/pina/domain/difference_domain.py b/pina/domain/difference_domain.py
index 9554aaf32..4015a3860 100644
--- a/pina/domain/difference_domain.py
+++ b/pina/domain/difference_domain.py
@@ -77,7 +77,7 @@ def sample(self, n, mode="random", variables="all"):
                 5
 
         """
-        if mode != self.sample_modes:
+        if mode not in self.sample_modes:
             raise NotImplementedError(
                 f"{mode} is not a valid mode for sampling."
             )
diff --git a/pina/domain/exclusion_domain.py b/pina/domain/exclusion_domain.py
index 4fc582cef..a05b1543e 100644
--- a/pina/domain/exclusion_domain.py
+++ b/pina/domain/exclusion_domain.py
@@ -76,7 +76,7 @@ def sample(self, n, mode="random", variables="all"):
                 5
 
         """
-        if mode != self.sample_modes:
+        if mode not in self.sample_modes:
             raise NotImplementedError(
                 f"{mode} is not a valid mode for sampling."
             )
diff --git a/pina/domain/intersection_domain.py b/pina/domain/intersection_domain.py
index b580f21c3..bb0499b50 100644
--- a/pina/domain/intersection_domain.py
+++ b/pina/domain/intersection_domain.py
@@ -78,7 +78,7 @@ def sample(self, n, mode="random", variables="all"):
                 5
 
         """
-        if mode != self.sample_modes:
+        if mode not in self.sample_modes:
             raise NotImplementedError(
                 f"{mode} is not a valid mode for sampling."
             )
diff --git a/pina/domain/simplex.py b/pina/domain/simplex.py
index 3d33bff8d..cea213265 100644
--- a/pina/domain/simplex.py
+++ b/pina/domain/simplex.py
@@ -92,13 +92,12 @@ def _build_cartesian(self, vertices):
         """
 
         span_dict = {}
-
         for i, coord in enumerate(self.variables):
-            sorted_vertices = sorted(vertices, key=lambda vertex: vertex[i])
+            sorted_vertices = torch.sort(vertices[coord].tensor.squeeze())
             # respective coord bounded by the lowest and highest values
             span_dict[coord] = [
-                float(sorted_vertices[0][i]),
-                float(sorted_vertices[-1][i]),
+                float(sorted_vertices.values[0]),
+                float(sorted_vertices.values[-1]),
             ]
 
         return CartesianDomain(span_dict)
diff --git a/pina/domain/union_domain.py b/pina/domain/union_domain.py
index bd7fa56cb..a72115f50 100644
--- a/pina/domain/union_domain.py
+++ b/pina/domain/union_domain.py
@@ -41,7 +41,10 @@ def sample_modes(self):
     
     @property
     def variables(self):
-        return list(set([geom.variables for geom in self.geometries]))
+        variables = []
+        for geom in self.geometries:
+            variables+=geom.variables
+        return list(set(variables))
 
     def is_inside(self, point, check_border=False):
         """
diff --git a/pina/label_tensor.py b/pina/label_tensor.py
index 08e0b03e6..1df318ec7 100644
--- a/pina/label_tensor.py
+++ b/pina/label_tensor.py
@@ -1,5 +1,5 @@
 """ Module for LabelTensor """
-
+from copy import deepcopy, copy
 import torch
 from torch import Tensor
 
@@ -35,12 +35,22 @@ def __init__(self, x, labels):
                     {1: {"name": "space"['a', 'b', 'c'])
 
         """
+        self.dim_names = None
         self.labels = labels
 
     @property
     def labels(self):
         """Property decorator for labels
 
+        :return: labels of self
+        :rtype: list
+        """
+        return self._labels[self.tensor.ndim-1]['dof']
+
+    @property
+    def full_labels(self):
+        """Property decorator for labels
+
         :return: labels of self
         :rtype: list
         """
@@ -65,6 +75,13 @@ def labels(self, labels):
             self.update_labels_from_list(labels)
         else:
             raise ValueError(f"labels must be list, dict or string.")
+        self.set_names()
+
+    def set_names(self):
+        labels = self.full_labels
+        self.dim_names = {}
+        for dim in range(self.tensor.ndim):
+            self.dim_names[labels[dim]['name']] = dim
 
     def extract(self, label_to_extract):
         """
@@ -76,46 +93,63 @@ def extract(self, label_to_extract):
         :raises TypeError: Labels are not ``str``.
         :raises ValueError: Label to extract is not in the labels ``list``.
         """
-        from copy import deepcopy
         if isinstance(label_to_extract, (str, int)):
             label_to_extract = [label_to_extract]
         if isinstance(label_to_extract, (tuple, list)):
-            last_dim_label = self._labels[self.tensor.ndim - 1]['dof']
-            if set(label_to_extract).issubset(last_dim_label) is False:
-                raise ValueError('Cannot extract a dof which is not in the original LabelTensor')
-            idx_to_extract = [last_dim_label.index(i) for i in label_to_extract]
-            new_tensor = self.tensor
-            new_tensor = new_tensor[..., idx_to_extract]
-            new_labels = deepcopy(self._labels)
-            last_dim_new_label = {self.tensor.ndim - 1: {
-                'dof': label_to_extract,
-                'name': self._labels[self.tensor.ndim - 1]['name']
-            }}
-            new_labels.update(last_dim_new_label)
+            return self._extract_from_list(label_to_extract)
         elif isinstance(label_to_extract, dict):
-            new_labels = (deepcopy(self._labels))
-            new_tensor = self.tensor
-            for k, v in label_to_extract.items():
-                idx_dim = None
-                for kl, vl in self._labels.items():
-                    if vl['name'] == k:
-                        idx_dim = kl
-                        break
-                dim_labels = self._labels[idx_dim]['dof']
-                if isinstance(label_to_extract[k], (int, str)):
-                    label_to_extract[k] = [label_to_extract[k]]
-                if set(label_to_extract[k]).issubset(dim_labels) is False:
-                    raise ValueError('Cannot extract a dof which is not in the original LabelTensor')
-                idx_to_extract = [dim_labels.index(i) for i in label_to_extract[k]]
-                indexer = [slice(None)] * idx_dim + [idx_to_extract] + [slice(None)] * (self.tensor.ndim - idx_dim - 1)
-                new_tensor = new_tensor[indexer]
-                dim_new_label = {idx_dim: {
-                    'dof': label_to_extract[k],
-                    'name': self._labels[idx_dim]['name']
-                }}
-                new_labels.update(dim_new_label)
+            return self._extract_from_dict(label_to_extract)
         else:
             raise ValueError('labels_to_extract must be str or list or dict')
+
+    def _extract_from_list(self, labels_to_extract):
+        #Store locally all necessary obj/variables
+        ndim = self.tensor.ndim
+        labels = self.full_labels
+        tensor = self.tensor
+        last_dim_label = self.labels
+
+        #Verify if all the labels in labels_to_extract are in last dimension
+        if set(labels_to_extract).issubset(last_dim_label) is False:
+            raise ValueError('Cannot extract a dof which is not in the original LabelTensor')
+
+        #Extract index to extract
+        idx_to_extract = [last_dim_label.index(i) for i in labels_to_extract]
+
+        #Perform extraction
+        new_tensor = tensor[..., idx_to_extract]
+
+        #Manage labels
+        new_labels = copy(labels)
+
+        last_dim_new_label = {ndim - 1: {
+            'dof': list(labels_to_extract),
+            'name': labels[ndim - 1]['name']
+        }}
+        new_labels.update(last_dim_new_label)
+        return LabelTensor(new_tensor, new_labels)
+
+    def _extract_from_dict(self, labels_to_extract):
+        labels = self.full_labels
+        tensor = self.tensor
+        ndim = tensor.ndim
+        new_labels = deepcopy(labels)
+        new_tensor = tensor
+        for k, _ in labels_to_extract.items():
+            idx_dim = self.dim_names[k]
+            dim_labels = labels[idx_dim]['dof']
+            if isinstance(labels_to_extract[k], (int, str)):
+                labels_to_extract[k] = [labels_to_extract[k]]
+            if set(labels_to_extract[k]).issubset(dim_labels) is False:
+                raise ValueError('Cannot extract a dof which is not in the original LabelTensor')
+            idx_to_extract = [dim_labels.index(i) for i in labels_to_extract[k]]
+            indexer = [slice(None)] * idx_dim + [idx_to_extract] + [slice(None)] * (ndim - idx_dim - 1)
+            new_tensor = new_tensor[indexer]
+            dim_new_label = {idx_dim: {
+                'dof': labels_to_extract[k],
+                'name': labels[idx_dim]['name']
+            }}
+            new_labels.update(dim_new_label)
         return LabelTensor(new_tensor, new_labels)
 
     def __str__(self):
@@ -147,31 +181,41 @@ def cat(tensors, dim=0):
             return []
         if len(tensors) == 1:
             return tensors[0]
+        new_labels_cat_dim = LabelTensor._check_validity_before_cat(tensors, dim)
+
+        # Perform cat on tensors
+        new_tensor = torch.cat(tensors, dim=dim)
+
+        #Update labels
+        labels = tensors[0].full_labels
+        labels.pop(dim)
+        new_labels_cat_dim = new_labels_cat_dim if len(set(new_labels_cat_dim)) == len(new_labels_cat_dim) \
+            else range(new_tensor.shape[dim])
+        labels[dim] = {'dof': new_labels_cat_dim,
+                       'name': tensors[1].full_labels[dim]['name']}
+        return LabelTensor(new_tensor, labels)
+
+    @staticmethod
+    def _check_validity_before_cat(tensors, dim):
         n_dims = tensors[0].ndim
         new_labels_cat_dim = []
+        # Check if names and dof of the labels are the same in all dimensions except in dim
         for i in range(n_dims):
-            name = tensors[0].labels[i]['name']
+            name = tensors[0].full_labels[i]['name']
             if i != dim:
-                dof = tensors[0].labels[i]['dof']
+                dof = tensors[0].full_labels[i]['dof']
                 for tensor in tensors:
-                    dof_to_check = tensor.labels[i]['dof']
-                    name_to_check = tensor.labels[i]['name']
+                    dof_to_check = tensor.full_labels[i]['dof']
+                    name_to_check = tensor.full_labels[i]['name']
                     if dof != dof_to_check or name != name_to_check:
                         raise ValueError('dimensions must have the same dof and name')
             else:
                 for tensor in tensors:
-                    new_labels_cat_dim += tensor.labels[i]['dof']
-                    name_to_check = tensor.labels[i]['name']
+                    new_labels_cat_dim += tensor.full_labels[i]['dof']
+                    name_to_check = tensor.full_labels[i]['name']
                     if name != name_to_check:
-                        raise ValueError('dimensions must have the same dof and name')
-        new_tensor = torch.cat(tensors, dim=dim)
-        labels = tensors[0].labels
-        labels.pop(dim)
-        new_labels_cat_dim = new_labels_cat_dim if len(set(new_labels_cat_dim)) == len(new_labels_cat_dim) \
-            else range(new_tensor.shape[dim])
-        labels[dim] = {'dof': new_labels_cat_dim,
-                       'name': tensors[1].labels[dim]['name']}
-        return LabelTensor(new_tensor, labels)
+                        raise ValueError('Dimensions to concatenate must have the same name')
+        return new_labels_cat_dim
 
     def requires_grad_(self, mode=True):
         lt = super().requires_grad_(mode)
@@ -204,7 +248,6 @@ def clone(self, *args, **kwargs):
         out = LabelTensor(super().clone(*args, **kwargs), self._labels)
         return out
 
-
     def init_labels(self):
         self._labels = {
             idx_: {
@@ -221,13 +264,14 @@ def update_labels_from_dict(self, labels):
         :type labels: dict
         :raises ValueError: dof list contain duplicates or number of dof does not match with tensor shape
         """
-
         tensor_shape = self.tensor.shape
+        #Check dimensionality
         for k, v in labels.items():
             if len(v['dof']) != len(set(v['dof'])):
                 raise ValueError("dof must be unique")
             if len(v['dof']) != tensor_shape[k]:
                 raise ValueError('Number of dof does not match with tensor dimension')
+        #Perform update
         self._labels.update(labels)
 
     def update_labels_from_list(self, labels):
@@ -237,6 +281,7 @@ def update_labels_from_list(self, labels):
         :param labels: The label(s) to update.
         :type labels: list
         """
+        # Create a dict with labels
         last_dim_labels = {self.tensor.ndim - 1: {'dof': labels, 'name': self.tensor.ndim - 1}}
         self.update_labels_from_dict(last_dim_labels)
 
@@ -246,26 +291,103 @@ def summation(tensors):
             raise ValueError('tensors list must not be empty')
         if len(tensors) == 1:
             return tensors[0]
-        labels = tensors[0].labels
+        # Collect all labels
+        labels = tensors[0].full_labels
+        # Check labels of all the tensors in each dimension
         for j in range(tensors[0].ndim):
             for i in range(1, len(tensors)):
-                if labels[j] != tensors[i].labels[j]:
+                if labels[j] != tensors[i].full_labels[j]:
                     labels.pop(j)
                     break
-
+        # Sum tensors
         data = torch.zeros(tensors[0].tensor.shape)
         for i in range(len(tensors)):
             data += tensors[i].tensor
         new_tensor = LabelTensor(data, labels)
         return new_tensor
 
-    def last_dim_dof(self):
-        return self._labels[self.tensor.ndim - 1]['dof']
-
     def append(self, tensor, mode='std'):
-        print(self.labels)
-        print(tensor.labels)
         if mode == 'std':
+            # Call cat on last dimension
             new_label_tensor = LabelTensor.cat([self, tensor], dim=self.tensor.ndim - 1)
-
+        elif mode=='cross':
+            # Crete tensor and call cat on last dimension
+            tensor1 = self
+            tensor2 = tensor
+            n1 = tensor1.shape[0]
+            n2 = tensor2.shape[0]
+            tensor1 = LabelTensor(tensor1.repeat(n2, 1), labels=tensor1.labels)
+            tensor2 = LabelTensor(tensor2.repeat_interleave(n1, dim=0), labels=tensor2.labels)
+            new_label_tensor = LabelTensor.cat([tensor1, tensor2], dim=self.tensor.ndim-1)
+        else:
+            raise ValueError('mode must be either "std" or "cross"')
         return new_label_tensor
+
+    @staticmethod
+    def vstack(label_tensors):
+        """
+        Stack tensors vertically. For more details, see
+        :meth:`torch.vstack`.
+
+        :param list(LabelTensor) label_tensors: the tensors to stack. They need
+            to have equal labels.
+        :return: the stacked tensor
+        :rtype: LabelTensor
+        """
+        return LabelTensor.cat(label_tensors, dim=0)
+
+    def __getitem__(self, index):
+        """
+        Return a copy of the selected tensor.
+        """
+
+        if isinstance(index, str) or (isinstance(index, (tuple, list)) and all(isinstance(a, str) for a in index)):
+            return self.extract(index)
+
+        selected_lt = super().__getitem__(index)
+
+        try:
+            len_index = len(index)
+        except TypeError:
+            len_index = 1
+
+        if isinstance(index, int) or len_index == 1:
+            if selected_lt.ndim == 1:
+                selected_lt = selected_lt.reshape(1, -1)
+            if hasattr(self, "labels"):
+                new_labels = deepcopy(self.full_labels)
+                new_labels.pop(0)
+                selected_lt.labels = new_labels
+        elif len(index) == self.tensor.ndim:
+            new_labels = deepcopy(self.full_labels)
+            if selected_lt.ndim == 1:
+                selected_lt = selected_lt.reshape(-1, 1)
+            for j in range(selected_lt.ndim):
+                if hasattr(self, "labels"):
+                    if isinstance(index[j], list):
+                        new_labels.update({j: {'dof': [new_labels[j]['dof'][i] for i in index[1]],
+                                               'name':new_labels[j]['name']}})
+                    else:
+                        new_labels.update({j: {'dof': new_labels[j]['dof'][index[j]],
+                                               'name':new_labels[j]['name']}})
+
+            selected_lt.labels = new_labels
+        else:
+            new_labels = deepcopy(self.full_labels)
+            new_labels.update({0: {'dof': list[index], 'name': new_labels[0]['name']}})
+            selected_lt.labels = self.labels
+
+        return selected_lt
+
+    def sort_labels(self, dim=None):
+        def argsort(lst):
+            return sorted(range(len(lst)), key=lambda x: lst[x])
+        if dim is None:
+            dim = self.tensor.ndim-1
+        labels = self.full_labels[dim]['dof']
+        sorted_index = argsort(labels)
+        indexer = [slice(None)] * self.tensor.ndim
+        indexer[dim] = sorted_index
+        new_labels = deepcopy(self.full_labels)
+        new_labels[dim] = {'dof': sorted(labels), 'name': new_labels[dim]['name']}
+        return LabelTensor(self.tensor[indexer], new_labels)
\ No newline at end of file
diff --git a/pina/operators.py b/pina/operators.py
index 082d72555..8822f2010 100644
--- a/pina/operators.py
+++ b/pina/operators.py
@@ -1,13 +1,11 @@
 """
 Module for operators vectorize implementation. Differential operators are used to write any differential problem.
-These operators are implemented to work on different accelerators: CPU, GPU, TPU or MPS.
+These operators are implemented to work on different accellerators: CPU, GPU, TPU or MPS.
 All operators take as input a tensor onto which computing the operator, a tensor with respect
 to which computing the operator, the name of the output variables to calculate the operator
 for (in case of multidimensional functions), and the variables name on which the operator is calculated.
 """
-
 import torch
-from copy import deepcopy
 from pina.label_tensor import LabelTensor
 
 
@@ -49,12 +47,12 @@ def grad_scalar_output(output_, input_, d):
         :rtype: LabelTensor
         """
 
-        if len(output_.labels[output_.tensor.ndim-1]['dof']) != 1:
+        if len(output_.labels) != 1:
             raise RuntimeError("only scalar function can be differentiated")
-        if not all([di in input_.labels[input_.tensor.ndim-1]['dof'] for di in d]):
+        if not all([di in input_.labels for di in d]):
             raise RuntimeError("derivative labels missing from input tensor")
 
-        output_fieldname = output_.labels[output_.ndim-1]['dof'][0]
+        output_fieldname = output_.labels[0]
         gradients = torch.autograd.grad(
             output_,
             input_,
@@ -65,35 +63,37 @@ def grad_scalar_output(output_, input_, d):
             retain_graph=True,
             allow_unused=True,
         )[0]
-        new_labels = deepcopy(input_.labels)
-        gradients.labels = new_labels
+
+        gradients.labels = input_.labels
         gradients = gradients.extract(d)
-        new_labels[input_.tensor.ndim - 1]['dof'] = [f"d{output_fieldname}d{i}" for i in d]
-        gradients.labels = new_labels
+        gradients.labels = [f"d{output_fieldname}d{i}" for i in d]
+
         return gradients
 
     if not isinstance(input_, LabelTensor):
         raise TypeError
+
     if d is None:
-        d = input_.labels[input_.tensor.ndim-1]['dof']
+        d = input_.labels
 
     if components is None:
-        components = output_.labels[output_.tensor.ndim-1]['dof']
+        components = output_.labels
 
-    if output_.shape[output_.ndim-1] == 1:  # scalar output ################################
+    if output_.shape[1] == 1:  # scalar output ################################
 
-        if components != output_.labels[output_.tensor.ndim-1]['dof']:
+        if components != output_.labels:
             raise RuntimeError
         gradients = grad_scalar_output(output_, input_, d)
 
-    elif output_.shape[output_.ndim-1] >= 2:  # vector output ##############################
+    elif output_.shape[output_.ndim - 1] >= 2:  # vector output ##############################
         tensor_to_cat = []
         for i, c in enumerate(components):
             c_output = output_.extract([c])
             tensor_to_cat.append(grad_scalar_output(c_output, input_, d))
-        gradients = LabelTensor.cat(tensor_to_cat, dim=output_.tensor.ndim-1)
+        gradients = LabelTensor.cat(tensor_to_cat, dim=output_.tensor.ndim - 1)
     else:
         raise NotImplementedError
+
     return gradients
 
 
@@ -124,30 +124,27 @@ def div(output_, input_, components=None, d=None):
         raise TypeError
 
     if d is None:
-        d = input_.labels[input_.tensor.ndim-1]['dof']
+        d = input_.labels
 
     if components is None:
-        components = output_.labels[output_.tensor.ndim-1]['dof']
+        components = output_.labels
 
-    if output_.shape[output_.ndim-1] < 2 or len(components) < 2:
+    if output_.shape[1] < 2 or len(components) < 2:
         raise ValueError("div supported only for vector fields")
 
     if len(components) != len(d):
         raise ValueError
 
     grad_output = grad(output_, input_, components, d)
-    last_dim_dof = [None] * len(components)
-    to_sum_tensors = []
+    labels = [None] * len(components)
+    tensors_to_sum = []
     for i, (c, d) in enumerate(zip(components, d)):
         c_fields = f"d{c}d{d}"
-        last_dim_dof[i] = c_fields
-        to_sum_tensors.append(grad_output.extract(c_fields))
-
-    div = LabelTensor.summation(to_sum_tensors)
-    new_labels = deepcopy(input_.labels)
-    new_labels[input_.tensor.ndim-1]['dof'] = ["+".join(last_dim_dof)]
-    div.labels = new_labels
-    return div
+        tensors_to_sum.append(grad_output.extract(c_fields))
+        labels[i] = c_fields
+    div_result = LabelTensor.summation(tensors_to_sum)
+    div_result.labels = ["+".join(labels)]
+    return div_result
 
 
 def laplacian(output_, input_, components=None, d=None, method="std"):
@@ -201,10 +198,10 @@ def scalar_laplace(output_, input_, components, d):
         return result
 
     if d is None:
-        d = input_.labels[input_.tensor.ndim-1]['dof']
+        d = input_.labels
 
     if components is None:
-        components = output_.labels[output_.tensor.ndim-1]['dof']
+        components = output_.labels
 
     if method == "divgrad":
         raise NotImplementedError("divgrad not implemented as method")
@@ -217,9 +214,9 @@ def scalar_laplace(output_, input_, components, d):
             # result = scalar_laplace(output_, input_, components, d) # TODO check (from 0.1)
             grad_output = grad(output_, input_, components=components, d=d)
             to_append_tensors = []
-            for i, label in enumerate(grad_output.labels[grad_output.ndim-1]['dof']):
+            for i, label in enumerate(grad_output.labels):
                 gg = grad(grad_output, input_, d=d, components=[label])
-                to_append_tensors.append(gg.extract([gg.labels[gg.tensor.ndim-1]['dof'][i]]))
+                to_append_tensors.append(gg.extract([gg.labels[i]]))
             labels = [f"dd{components[0]}"]
             result = LabelTensor.summation(tensors=to_append_tensors)
             result.labels = labels
@@ -236,21 +233,27 @@ def scalar_laplace(output_, input_, components, d):
 
     # result = result.as_subclass(LabelTensor)
     # result.labels = labels
+            result = torch.empty(
+                input_.shape[0], len(components), device=output_.device
+            )
             labels = [None] * len(components)
             to_append_tensors = [None] * len(components)
             for idx, (ci, di) in enumerate(zip(components, d)):
+
                 if not isinstance(ci, list):
                     ci = [ci]
                 if not isinstance(di, list):
                     di = [di]
+
                 grad_output = grad(output_, input_, components=ci, d=di)
+                result[:, idx] = grad(grad_output, input_, d=di).flatten()
                 to_append_tensors[idx] = grad(grad_output, input_, d=di)
                 labels[idx] = f"dd{ci[0]}dd{di[0]}"
             result = LabelTensor.cat(tensors=to_append_tensors, dim=output_.tensor.ndim-1)
             result.labels = labels
     return result
 
-# TODO Fix advection operator
+
 def advection(output_, input_, velocity_field, components=None, d=None):
     """
     Perform advection operation. The operator works for vectorial functions,
@@ -272,10 +275,10 @@ def advection(output_, input_, velocity_field, components=None, d=None):
     :rtype: LabelTensor
     """
     if d is None:
-        d = input_.labels[input_.tensor.ndim-1]['dof']
+        d = input_.labels
 
     if components is None:
-        components = output_.labels[output_.tensor.ndim-1]['dof']
+        components = output_.labels
 
     tmp = (
         grad(output_, input_, components, d)
diff --git a/pina/problem/abstract_problem.py b/pina/problem/abstract_problem.py
index 8894207bf..edf214a82 100644
--- a/pina/problem/abstract_problem.py
+++ b/pina/problem/abstract_problem.py
@@ -36,7 +36,15 @@ def __init__(self):
 
     @property
     def input_pts(self):
-        return self.collector.data_collections
+        to_return = {}
+        for k, v in self.collector.data_collections.items():
+            if 'input_points' in v.keys():
+                to_return[k] = v['input_points']
+        return to_return
+
+    @property
+    def _have_sampled_points(self):
+        return self.collector.is_conditions_ready
     
     def __deepcopy__(self, memo):
         """
@@ -165,3 +173,6 @@ def discretise_domain(
 
         # store data
         self.collector.store_sample_domains(n, mode, variables, locations)
+
+    def add_points(self, new_points_dict):
+        self.collector.add_points(new_points_dict)
diff --git a/tests/test_condition.py b/tests/test_condition.py
index 5f1c6236b..f12979dc8 100644
--- a/tests/test_condition.py
+++ b/tests/test_condition.py
@@ -18,27 +18,27 @@ def test_init_inputoutput():
     Condition(input_points=example_input_pts, output_points=example_output_pts)
     with pytest.raises(ValueError):
         Condition(example_input_pts, example_output_pts)
-    with pytest.raises(TypeError):
+    with pytest.raises(ValueError):
         Condition(input_points=3., output_points='example')
-    with pytest.raises(TypeError):
+    with pytest.raises(ValueError):
         Condition(input_points=example_domain, output_points=example_domain)
+test_init_inputoutput()
 
-
-def test_init_locfunc():
-    Condition(location=example_domain, equation=FixedValue(0.0))
+def test_init_domainfunc():
+    Condition(domain=example_domain, equation=FixedValue(0.0))
     with pytest.raises(ValueError):
         Condition(example_domain, FixedValue(0.0))
-    with pytest.raises(TypeError):
-        Condition(location=3., equation='example')
-    with pytest.raises(TypeError):
-        Condition(location=example_input_pts, equation=example_output_pts)
+    with pytest.raises(ValueError):
+        Condition(domain=3., equation='example')
+    with pytest.raises(ValueError):
+        Condition(domain=example_input_pts, equation=example_output_pts)
 
 
 def test_init_inputfunc():
     Condition(input_points=example_input_pts, equation=FixedValue(0.0))
     with pytest.raises(ValueError):
         Condition(example_domain, FixedValue(0.0))
-    with pytest.raises(TypeError):
+    with pytest.raises(ValueError):
         Condition(input_points=3., equation='example')
-    with pytest.raises(TypeError):
+    with pytest.raises(ValueError):
         Condition(input_points=example_domain, equation=example_output_pts)
diff --git a/tests/test_geometry/test_simplex.py b/tests/test_geometry/test_simplex.py
index 7fc34ce2d..25224aae3 100644
--- a/tests/test_geometry/test_simplex.py
+++ b/tests/test_geometry/test_simplex.py
@@ -40,7 +40,6 @@ def test_constructor():
             LabelTensor(torch.tensor([[-.5, .5]]), labels=["x", "y"]),
         ])
 
-
 def test_sample():
     # sampling inside
     simplex = SimplexDomain([
diff --git a/tests/test_label_tensor.py b/tests/test_label_tensor/test_label_tensor.py
similarity index 56%
rename from tests/test_label_tensor.py
rename to tests/test_label_tensor/test_label_tensor.py
index 6ef484f0b..1165594db 100644
--- a/tests/test_label_tensor.py
+++ b/tests/test_label_tensor/test_label_tensor.py
@@ -2,7 +2,6 @@
 import pytest
 
 from pina.label_tensor import LabelTensor
-#import pina
 
 data = torch.rand((20, 3))
 labels_column = {
@@ -22,8 +21,7 @@
 
 @pytest.mark.parametrize("labels", [labels_column, labels_row, labels_all, labels_list])
 def test_constructor(labels):
-    LabelTensor(data, labels)
-
+    print(LabelTensor(data, labels))
 
 def test_wrong_constructor():
     with pytest.raises(ValueError):
@@ -92,7 +90,7 @@ def test_extract_3D():
     ))
     assert tensor2.ndim == tensor.ndim
     assert tensor2.shape == tensor.shape
-    assert tensor.labels == tensor2.labels
+    assert tensor.full_labels == tensor2.full_labels
     assert new.shape != tensor.shape
 
 def test_concatenation_3D():
@@ -104,9 +102,9 @@ def test_concatenation_3D():
     lt2 = LabelTensor(data_2, labels_2)
     lt_cat = LabelTensor.cat([lt1, lt2])
     assert lt_cat.shape == (70, 3, 4)
-    assert lt_cat.labels[0]['dof'] == range(70)
-    assert lt_cat.labels[1]['dof'] == range(3)
-    assert lt_cat.labels[2]['dof'] == ['x', 'y', 'z', 'w']
+    assert lt_cat.full_labels[0]['dof'] == range(70)
+    assert lt_cat.full_labels[1]['dof'] == range(3)
+    assert lt_cat.full_labels[2]['dof'] == ['x', 'y', 'z', 'w']
 
     data_1 = torch.rand(20, 3, 4)
     labels_1 = ['x', 'y', 'z', 'w']
@@ -116,9 +114,9 @@ def test_concatenation_3D():
     lt2 = LabelTensor(data_2, labels_2)
     lt_cat = LabelTensor.cat([lt1, lt2], dim=1)
     assert lt_cat.shape == (20, 5, 4)
-    assert lt_cat.labels[0]['dof'] == range(20)
-    assert lt_cat.labels[1]['dof'] == range(5)
-    assert lt_cat.labels[2]['dof'] == ['x', 'y', 'z', 'w']
+    assert lt_cat.full_labels[0]['dof'] == range(20)
+    assert lt_cat.full_labels[1]['dof'] == range(5)
+    assert lt_cat.full_labels[2]['dof'] == ['x', 'y', 'z', 'w']
 
     data_1 = torch.rand(20, 3, 2)
     labels_1 = ['x', 'y']
@@ -128,9 +126,9 @@ def test_concatenation_3D():
     lt2 = LabelTensor(data_2, labels_2)
     lt_cat = LabelTensor.cat([lt1, lt2], dim=2)
     assert lt_cat.shape == (20, 3, 5)
-    assert lt_cat.labels[2]['dof'] == ['x', 'y', 'z', 'w', 'a']
-    assert lt_cat.labels[0]['dof'] == range(20)
-    assert lt_cat.labels[1]['dof'] == range(3)
+    assert lt_cat.full_labels[2]['dof'] == ['x', 'y', 'z', 'w', 'a']
+    assert lt_cat.full_labels[0]['dof'] == range(20)
+    assert lt_cat.full_labels[1]['dof'] == range(3)
 
     data_1 = torch.rand(20, 2, 4)
     labels_1 = ['x', 'y', 'z', 'w']
@@ -140,7 +138,6 @@ def test_concatenation_3D():
     lt2 = LabelTensor(data_2, labels_2)
     with pytest.raises(ValueError):
         LabelTensor.cat([lt1, lt2], dim=2)
-
     data_1 = torch.rand(20, 3, 2)
     labels_1 = ['x', 'y']
     lt1 = LabelTensor(data_1, labels_1)
@@ -149,9 +146,9 @@ def test_concatenation_3D():
     lt2 = LabelTensor(data_2, labels_2)
     lt_cat = LabelTensor.cat([lt1, lt2], dim=2)
     assert lt_cat.shape == (20, 3, 5)
-    assert lt_cat.labels[2]['dof'] == range(5)
-    assert lt_cat.labels[0]['dof'] == range(20)
-    assert lt_cat.labels[1]['dof'] == range(3)
+    assert lt_cat.full_labels[2]['dof'] == range(5)
+    assert lt_cat.full_labels[0]['dof'] == range(20)
+    assert lt_cat.full_labels[1]['dof'] == range(3)
 
 
 def test_summation():
@@ -165,7 +162,7 @@ def test_summation():
     assert lt_sum.ndim == lt_sum.ndim
     assert lt_sum.shape[0] == 20
     assert lt_sum.shape[1] == 3
-    assert lt_sum.labels == labels_all
+    assert lt_sum.full_labels == labels_all
     assert torch.eq(lt_sum.tensor, torch.ones(20,3)*2).all()
     lt1 = LabelTensor(torch.ones(20,3), labels_all)
     lt2 = LabelTensor(torch.ones(20,3), labels_all)
@@ -174,29 +171,92 @@ def test_summation():
     assert lt_sum.ndim == lt_sum.ndim
     assert lt_sum.shape[0] == 20
     assert lt_sum.shape[1] == 3
-    assert lt_sum.labels == labels_all
+    assert lt_sum.full_labels == labels_all
     assert torch.eq(lt_sum.tensor, torch.ones(20,3)*2).all()
 
 def test_append_3D():
-    data_1 = torch.rand(20, 3, 4)
-    labels_1 = ['x', 'y', 'z', 'w']
+    data_1 = torch.rand(20, 3, 2)
+    labels_1 = ['x', 'y']
     lt1 = LabelTensor(data_1, labels_1)
-    data_2 = torch.rand(50, 3, 4)
-    labels_2 = ['x', 'y', 'z', 'w']
+    data_2 = torch.rand(20, 3, 2)
+    labels_2 = ['z', 'w']
     lt2 = LabelTensor(data_2, labels_2)
     lt1 = lt1.append(lt2)
-    assert lt1.shape == (70, 3, 4)
-    assert lt1.labels[0]['dof'] == range(70)
-    assert lt1.labels[1]['dof'] == range(3)
-    assert lt1.labels[2]['dof'] == ['x', 'y', 'z', 'w']
-    data_1 = torch.rand(20, 3, 2)
+    assert lt1.shape == (20, 3, 4)
+    assert lt1.full_labels[0]['dof'] == range(20)
+    assert lt1.full_labels[1]['dof'] == range(3)
+    assert lt1.full_labels[2]['dof'] == ['x', 'y', 'z', 'w']
+
+def test_append_2D():
+    data_1 = torch.rand(20, 2)
     labels_1 = ['x', 'y']
     lt1 = LabelTensor(data_1, labels_1)
-    data_2 = torch.rand(20, 3, 2)
+    data_2 = torch.rand(20, 2)
     labels_2 = ['z', 'w']
     lt2 = LabelTensor(data_2, labels_2)
     lt1 = lt1.append(lt2, mode='cross')
-    assert lt1.shape == (20, 3, 4)
-    assert lt1.labels[0]['dof'] == range(20)
-    assert lt1.labels[1]['dof'] == range(3)
-    assert lt1.labels[2]['dof'] == ['x', 'y', 'z', 'w']
+    assert lt1.shape == (400, 4)
+    assert lt1.full_labels[0]['dof'] == range(400)
+    assert lt1.full_labels[1]['dof'] == ['x', 'y', 'z', 'w']
+
+def test_vstack_3D():
+    data_1 = torch.rand(20, 3, 2)
+    labels_1 = {1:{'dof': ['a', 'b', 'c'], 'name': 'first'}, 2: {'dof': ['x', 'y'], 'name': 'second'}}
+    lt1 = LabelTensor(data_1, labels_1)
+    data_2 = torch.rand(20, 3, 2)
+    labels_1 = {1:{'dof': ['a', 'b', 'c'], 'name': 'first'}, 2: {'dof': ['x', 'y'], 'name': 'second'}}
+    lt2 = LabelTensor(data_2, labels_1)
+    lt_stacked = LabelTensor.vstack([lt1, lt2])
+    assert lt_stacked.shape == (40, 3, 2)
+    assert lt_stacked.full_labels[0]['dof'] == range(40)
+    assert lt_stacked.full_labels[1]['dof'] == ['a', 'b', 'c']
+    assert lt_stacked.full_labels[2]['dof'] == ['x', 'y']
+    assert lt_stacked.full_labels[1]['name'] == 'first'
+    assert lt_stacked.full_labels[2]['name'] == 'second'
+
+def test_vstack_2D():
+    data_1 = torch.rand(20, 2)
+    labels_1 = { 1: {'dof': ['x', 'y'], 'name': 'second'}}
+    lt1 = LabelTensor(data_1, labels_1)
+    data_2 = torch.rand(20, 2)
+    labels_1 = { 1: {'dof': ['x', 'y'], 'name': 'second'}}
+    lt2 = LabelTensor(data_2, labels_1)
+    lt_stacked = LabelTensor.vstack([lt1, lt2])
+    assert lt_stacked.shape == (40, 2)
+    assert lt_stacked.full_labels[0]['dof'] == range(40)
+    assert lt_stacked.full_labels[1]['dof'] == ['x', 'y']
+    assert lt_stacked.full_labels[0]['name'] == 0
+    assert lt_stacked.full_labels[1]['name'] == 'second'
+
+def test_sorting():
+    data = torch.ones(20, 5)
+    data[:,0] = data[:,0]*4
+    data[:,1] = data[:,1]*2
+    data[:,2] = data[:,2]
+    data[:,3] = data[:,3]*5
+    data[:,4] = data[:,4]*3
+    labels = ['d', 'b', 'a', 'e', 'c']
+    lt_data = LabelTensor(data, labels)
+    lt_sorted = LabelTensor.sort_labels(lt_data)
+    assert lt_sorted.shape == (20,5)
+    assert lt_sorted.labels == ['a', 'b', 'c', 'd', 'e']
+    assert torch.eq(lt_sorted.tensor[:,0], torch.ones(20) * 1).all()
+    assert torch.eq(lt_sorted.tensor[:,1], torch.ones(20) * 2).all()
+    assert torch.eq(lt_sorted.tensor[:,2], torch.ones(20) * 3).all()
+    assert torch.eq(lt_sorted.tensor[:,3], torch.ones(20) * 4).all()
+    assert torch.eq(lt_sorted.tensor[:,4], torch.ones(20) * 5).all()
+
+    data = torch.ones(20, 4, 5)
+    data[:,0,:] = data[:,0]*4
+    data[:,1,:] = data[:,1]*2
+    data[:,2,:] = data[:,2]
+    data[:,3,:] = data[:,3]*3
+    labels = {1: {'dof': ['d', 'b', 'a', 'c'], 'name': 1}}
+    lt_data = LabelTensor(data, labels)
+    lt_sorted = LabelTensor.sort_labels(lt_data, dim=1)
+    assert lt_sorted.shape == (20,4, 5)
+    assert lt_sorted.full_labels[1]['dof'] == ['a', 'b', 'c', 'd']
+    assert torch.eq(lt_sorted.tensor[:,0,:], torch.ones(20,5) * 1).all()
+    assert torch.eq(lt_sorted.tensor[:,1,:], torch.ones(20,5) * 2).all()
+    assert torch.eq(lt_sorted.tensor[:,2,:], torch.ones(20,5) * 3).all()
+    assert torch.eq(lt_sorted.tensor[:,3,:], torch.ones(20,5) * 4).all()
diff --git a/tests/test_label_tensor/test_label_tensor_01.py b/tests/test_label_tensor/test_label_tensor_01.py
new file mode 100644
index 000000000..a2e129d94
--- /dev/null
+++ b/tests/test_label_tensor/test_label_tensor_01.py
@@ -0,0 +1,117 @@
+import torch
+import pytest
+
+from pina import LabelTensor
+
+data = torch.rand((20, 3))
+labels = ['a', 'b', 'c']
+
+
+def test_constructor():
+    LabelTensor(data, labels)
+
+
+def test_wrong_constructor():
+    with pytest.raises(ValueError):
+        LabelTensor(data, ['a', 'b'])
+
+
+def test_labels():
+    tensor = LabelTensor(data, labels)
+    assert isinstance(tensor, torch.Tensor)
+    assert tensor.labels == labels
+    with pytest.raises(ValueError):
+        tensor.labels = labels[:-1]
+
+
+def test_extract():
+    label_to_extract = ['a', 'c']
+    tensor = LabelTensor(data, labels)
+    new = tensor.extract(label_to_extract)
+    assert new.labels == label_to_extract
+    assert new.shape[1] == len(label_to_extract)
+    assert torch.all(torch.isclose(data[:, 0::2], new))
+
+
+def test_extract_onelabel():
+    label_to_extract = ['a']
+    tensor = LabelTensor(data, labels)
+    new = tensor.extract(label_to_extract)
+    assert new.ndim == 2
+    assert new.labels == label_to_extract
+    assert new.shape[1] == len(label_to_extract)
+    assert torch.all(torch.isclose(data[:, 0].reshape(-1, 1), new))
+
+
+def test_wrong_extract():
+    label_to_extract = ['a', 'cc']
+    tensor = LabelTensor(data, labels)
+    with pytest.raises(ValueError):
+        tensor.extract(label_to_extract)
+
+
+def test_extract_order():
+    label_to_extract = ['c', 'a']
+    tensor = LabelTensor(data, labels)
+    new = tensor.extract(label_to_extract)
+    expected = torch.cat(
+        (data[:, 2].reshape(-1, 1), data[:, 0].reshape(-1, 1)),
+        dim=1)
+    assert new.labels == label_to_extract
+    assert new.shape[1] == len(label_to_extract)
+    assert torch.all(torch.isclose(expected, new))
+
+
+def test_merge():
+    tensor = LabelTensor(data, labels)
+    tensor_a = tensor.extract('a')
+    tensor_b = tensor.extract('b')
+    tensor_c = tensor.extract('c')
+
+    tensor_bc = tensor_b.append(tensor_c)
+    assert torch.allclose(tensor_bc, tensor.extract(['b', 'c']))
+
+
+def test_merge2():
+    tensor = LabelTensor(data, labels)
+    tensor_b = tensor.extract('b')
+    tensor_c = tensor.extract('c')
+
+    tensor_bc = tensor_b.append(tensor_c)
+    assert torch.allclose(tensor_bc, tensor.extract(['b', 'c']))
+
+
+def test_getitem():
+    tensor = LabelTensor(data, labels)
+    tensor_view = tensor['a']
+    assert tensor_view.labels == ['a']
+    assert torch.allclose(tensor_view.flatten(), data[:, 0])
+
+    tensor_view = tensor['a', 'c']
+    assert tensor_view.labels == ['a', 'c']
+    assert torch.allclose(tensor_view, data[:, 0::2])
+
+def test_getitem2():
+    tensor = LabelTensor(data, labels)
+    tensor_view = tensor[:5]
+    assert tensor_view.labels == labels
+    assert torch.allclose(tensor_view, data[:5])
+
+    idx = torch.randperm(tensor.shape[0])
+    tensor_view = tensor[idx]
+    assert tensor_view.labels == labels
+
+def test_slice():
+    tensor = LabelTensor(data, labels)
+    tensor_view = tensor[:5, :2]
+    assert tensor_view.labels == labels[:2]
+    assert torch.allclose(tensor_view, data[:5, :2])
+
+    tensor_view2 = tensor[3]
+
+    assert tensor_view2.labels == labels
+    assert torch.allclose(tensor_view2, data[3])
+
+    tensor_view3 = tensor[:, 2]
+    assert tensor_view3.labels == labels[2]
+    assert torch.allclose(tensor_view3, data[:, 2].reshape(-1, 1))
\ No newline at end of file
diff --git a/tests/test_operators.py b/tests/test_operators.py
index e18eaf29e..1271c3712 100644
--- a/tests/test_operators.py
+++ b/tests/test_operators.py
@@ -27,15 +27,15 @@ def test_grad_scalar_output():
     grad_tensor_s = grad(tensor_s, inp)
     true_val = 2*inp
     assert grad_tensor_s.shape == inp.shape
-    assert grad_tensor_s.labels[grad_tensor_s.ndim-1]['dof'] == [
-        f'd{tensor_s.labels[tensor_s.ndim-1]["dof"][0]}d{i}' for i in inp.labels[inp.ndim-1]['dof']
+    assert grad_tensor_s.labels == [
+        f'd{tensor_s.labels[0]}d{i}' for i in inp.labels
     ]
     assert torch.allclose(grad_tensor_s, true_val)
 
     grad_tensor_s = grad(tensor_s, inp, d=['x', 'y'])
     assert grad_tensor_s.shape == (20, 2)
-    assert grad_tensor_s.labels[grad_tensor_s.ndim-1]['dof'] == [
-        f'd{tensor_s.labels[tensor_s.ndim-1]["dof"][0]}d{i}' for i in ['x', 'y']
+    assert grad_tensor_s.labels == [
+        f'd{tensor_s.labels[0]}d{i}' for i in ['x', 'y']
     ]
     assert torch.allclose(grad_tensor_s, true_val)
 
diff --git a/tests/test_problem.py b/tests/test_problem.py
index 77871b1a1..cc7e25535 100644
--- a/tests/test_problem.py
+++ b/tests/test_problem.py
@@ -27,50 +27,46 @@ class Poisson(SpatialProblem):
 
     conditions = {
         'gamma1':
-        Condition(domain=CartesianDomain({
-            'x': [0, 1],
-            'y': 1
-        }),
-                  equation=FixedValue(0.0)),
+            Condition(domain=CartesianDomain({
+                'x': [0, 1],
+                'y': 1
+            }),
+                equation=FixedValue(0.0)),
         'gamma2':
-        Condition(domain=CartesianDomain({
-            'x': [0, 1],
-            'y': 0
-        }),
-                  equation=FixedValue(0.0)),
+            Condition(domain=CartesianDomain({
+                'x': [0, 1],
+                'y': 0
+            }),
+                equation=FixedValue(0.0)),
         'gamma3':
-        Condition(domain=CartesianDomain({
-            'x': 1,
-            'y': [0, 1]
-        }),
-                  equation=FixedValue(0.0)),
+            Condition(domain=CartesianDomain({
+                'x': 1,
+                'y': [0, 1]
+            }),
+                equation=FixedValue(0.0)),
         'gamma4':
-        Condition(domain=CartesianDomain({
-            'x': 0,
-            'y': [0, 1]
-        }),
-                  equation=FixedValue(0.0)),
+            Condition(domain=CartesianDomain({
+                'x': 0,
+                'y': [0, 1]
+            }),
+                equation=FixedValue(0.0)),
         'D':
-        Condition(domain=CartesianDomain({
-            'x': [0, 1],
-            'y': [0, 1]
-        }),
-                  equation=my_laplace),
+            Condition(domain=CartesianDomain({
+                'x': [0, 1],
+                'y': [0, 1]
+            }),
+                equation=my_laplace),
         'data':
-        Condition(input_points=in_, output_points=out_)
+            Condition(input_points=in_, output_points=out_)
     }
 
     def poisson_sol(self, pts):
         return -(torch.sin(pts.extract(['x']) * torch.pi) *
-                 torch.sin(pts.extract(['y']) * torch.pi)) / (2 * torch.pi**2)
+                 torch.sin(pts.extract(['y']) * torch.pi)) / (2 * torch.pi ** 2)
 
     truth_solution = poisson_sol
 
 
-# make the problem
-poisson_problem = Poisson()
-print(poisson_problem.input_pts)
-
 def test_discretise_domain():
     n = 10
     poisson_problem = Poisson()
@@ -83,7 +79,7 @@ def test_discretise_domain():
         assert poisson_problem.input_pts[b].shape[0] == n
 
     poisson_problem.discretise_domain(n, 'grid', locations=['D'])
-    assert poisson_problem.input_pts['D'].shape[0] == n**2
+    assert poisson_problem.input_pts['D'].shape[0] == n ** 2
     poisson_problem.discretise_domain(n, 'random', locations=['D'])
     assert poisson_problem.input_pts['D'].shape[0] == n
 
@@ -94,14 +90,15 @@ def test_discretise_domain():
     assert poisson_problem.input_pts['D'].shape[0] == n
 
 
-# def test_sampling_few_variables():
-#     n = 10
-#     poisson_problem.discretise_domain(n,
-#                                       'grid',
-#                                       locations=['D'],
-#                                       variables=['x'])
-#     assert poisson_problem.input_pts['D'].shape[1] == 1
-#     assert poisson_problem._have_sampled_points['D'] is False
+def test_sampling_few_variables():
+    n = 10
+    poisson_problem = Poisson()
+    poisson_problem.discretise_domain(n,
+                                      'grid',
+                                      locations=['D'],
+                                      variables=['x'])
+    assert poisson_problem.input_pts['D'].shape[1] == 1
+    assert poisson_problem._have_sampled_points['D'] is False
 
 
 def test_variables_correct_order_sampling():
@@ -117,13 +114,11 @@ def test_variables_correct_order_sampling():
                                       variables=['y'])
     assert poisson_problem.input_pts['D'].labels == sorted(
         poisson_problem.input_variables)
-    
     poisson_problem.discretise_domain(n,
                                       'grid',
                                       locations=['D'])
     assert poisson_problem.input_pts['D'].labels == sorted(
         poisson_problem.input_variables)
-    
     poisson_problem.discretise_domain(n,
                                       'grid',
                                       locations=['D'],
@@ -140,8 +135,8 @@ def test_add_points():
     poisson_problem.discretise_domain(0,
                                       'random',
                                       locations=['D'],
-                                      variables=['x','y'])
-    new_pts = LabelTensor(torch.tensor([[0.5,-0.5]]),labels=['x','y'])
+                                      variables=['x', 'y'])
+    new_pts = LabelTensor(torch.tensor([[0.5, -0.5]]), labels=['x', 'y'])
     poisson_problem.add_points({'D': new_pts})
-    assert torch.isclose(poisson_problem.input_pts['D'].extract('x'),new_pts.extract('x'))
-    assert torch.isclose(poisson_problem.input_pts['D'].extract('y'),new_pts.extract('y'))
+    assert torch.isclose(poisson_problem.input_pts['D'].extract('x'), new_pts.extract('x'))
+    assert torch.isclose(poisson_problem.input_pts['D'].extract('y'), new_pts.extract('y'))

From 7c3554496feb7dda47231ad128b3e475d9c4d384 Mon Sep 17 00:00:00 2001
From: Dario Coscia <dariocos99@gmail.com>
Date: Thu, 10 Oct 2024 19:24:46 +0200
Subject: [PATCH 23/55] minor changes/ trainer update

---
 pina/collector.py                           | 38 ++++++++++++++-------
 pina/condition/data_condition.py            | 12 +++----
 pina/condition/domain_equation_condition.py |  2 +-
 pina/condition/input_equation_condition.py  |  2 +-
 pina/condition/input_output_condition.py    |  2 +-
 pina/problem/abstract_problem.py            | 17 +++++----
 pina/trainer.py                             | 30 +++++++---------
 tests/test_problem.py                       | 12 ++++++-
 8 files changed, 69 insertions(+), 46 deletions(-)

diff --git a/pina/collector.py b/pina/collector.py
index 0f4e9da44..f44c222a7 100644
--- a/pina/collector.py
+++ b/pina/collector.py
@@ -5,21 +5,35 @@
 
 class Collector:
     def __init__(self, problem):
-        self.problem = problem                                                  # hook Collector <-> Problem
-        self.data_collections = {name : {} for name in self.problem.conditions} # collection of data
-        self.is_conditions_ready = {
-            name : False for name in self.problem.conditions}                   # names of the conditions that need to be sampled
-        self.full = False                                                       # collector full, all points for all conditions are given and the data are ready to be used in trainig
+        # creating a hook between collector and problem
+        self.problem = problem          
+
+        # this variable is used to store the data in the form:
+        # {'[condition_name]' : 
+        #           {'input_points' : Tensor, 
+        #            '[equation/output_points/conditional_variables]': Tensor}
+        # }
+        # those variables are used for the dataloading
+        self._data_collections = {name : {} for name in self.problem.conditions}
+
+        # variables used to check that all conditions are sampled
+        self._is_conditions_ready = {
+            name : False for name in self.problem.conditions}
+        self.full = False
         
     @property
     def full(self):
-        return all(self.is_conditions_ready.values())
+        return all(self._is_conditions_ready.values())
     
     @full.setter
     def full(self, value):
         check_consistency(value, bool)
         self._full = value
 
+    @property
+    def data_collections(self):
+        return self._data_collections
+
     @property
     def problem(self):
         return self._problem
@@ -33,13 +47,13 @@ def store_fixed_data(self):
         for condition_name, condition in self.problem.conditions.items():
             # if the condition is not ready and domain is not attribute
             # of condition, we get and store the data
-            if (not self.is_conditions_ready[condition_name]) and (not hasattr(condition, "domain")):
+            if (not self._is_conditions_ready[condition_name]) and (not hasattr(condition, "domain")):
                 # get data
                 keys = condition.__slots__
                 values = [getattr(condition, name) for name in keys]
                 self.data_collections[condition_name] = dict(zip(keys, values))
                 # condition now is ready
-                self.is_conditions_ready[condition_name] = True
+                self._is_conditions_ready[condition_name] = True
 
     def store_sample_domains(self, n, mode, variables, sample_locations):
         # loop over all locations
@@ -48,7 +62,7 @@ def store_sample_domains(self, n, mode, variables, sample_locations):
             condition = self.problem.conditions[loc]
             keys = ["input_points", "equation"]
             # if the condition is not ready, we get and store the data
-            if (not self.is_conditions_ready[loc]):
+            if (not self._is_conditions_ready[loc]):
                 # if it is the first time we sample
                 if not self.data_collections[loc]:
                     already_sampled = []
@@ -57,7 +71,7 @@ def store_sample_domains(self, n, mode, variables, sample_locations):
                     already_sampled = [self.data_collections[loc]['input_points']]
             # if the condition is ready but we want to sample again
             else:
-                self.is_conditions_ready[loc] = False
+                self._is_conditions_ready[loc] = False
                 already_sampled = []
 
             # get the samples
@@ -70,7 +84,7 @@ def store_sample_domains(self, n, mode, variables, sample_locations):
                 ):
                 pts = pts.sort_labels()
                 if sorted(pts.labels)==sorted(self.problem.input_variables):
-                    self.is_conditions_ready[loc] = True
+                    self._is_conditions_ready[loc] = True
                 values = [pts, condition.equation]
                 self.data_collections[loc] = dict(zip(keys, values))
             else:
@@ -84,6 +98,6 @@ def add_points(self, new_points_dict):
         :raises RuntimeError: if at least one condition is not already sampled
         """
         for k,v in new_points_dict.items():
-            if not self.is_conditions_ready[k]:
+            if not self._is_conditions_ready[k]:
                 raise RuntimeError('Cannot add points on a non sampled condition')
             self.data_collections[k]['input_points'] = self.data_collections[k]['input_points'].vstack(v)
\ No newline at end of file
diff --git a/pina/condition/data_condition.py b/pina/condition/data_condition.py
index d5ac63970..90d248b67 100644
--- a/pina/condition/data_condition.py
+++ b/pina/condition/data_condition.py
@@ -13,20 +13,20 @@ class DataConditionInterface(ConditionInterface):
     distribution
     """
 
-    __slots__ = ["data", "conditionalvariable"]
+    __slots__ = ["input_points", "conditional_variables"]
 
-    def __init__(self, data, conditionalvariable=None):
+    def __init__(self, input_points, conditional_variables=None):
         """
         TODO
         """
         super().__init__()
-        self.data = data
-        self.conditionalvariable = conditionalvariable
+        self.input_points = input_points
+        self.conditional_variables = conditional_variables
         self.condition_type = 'unsupervised'
 
     def __setattr__(self, key, value):
-        if (key == 'data') or (key == 'conditionalvariable'):
+        if (key == 'input_points') or (key == 'conditional_variables'):
             check_consistency(value, (LabelTensor, Graph, torch.Tensor))
             DataConditionInterface.__dict__[key].__set__(self, value)
-        elif key in ('_condition_type', '_problem', 'problem', 'condition_type'):
+        elif key in ('problem', 'condition_type'):
             super().__setattr__(key, value)
\ No newline at end of file
diff --git a/pina/condition/domain_equation_condition.py b/pina/condition/domain_equation_condition.py
index ab35d202a..ce4c7d3fc 100644
--- a/pina/condition/domain_equation_condition.py
+++ b/pina/condition/domain_equation_condition.py
@@ -29,5 +29,5 @@ def __setattr__(self, key, value):
         elif key == 'equation':
             check_consistency(value, (EquationInterface))
             DomainEquationCondition.__dict__[key].__set__(self, value)
-        elif key in ('_condition_type', '_problem', 'problem', 'condition_type'):
+        elif key in ('problem', 'condition_type'):
             super().__setattr__(key, value)
\ No newline at end of file
diff --git a/pina/condition/input_equation_condition.py b/pina/condition/input_equation_condition.py
index dc12d0225..ac47fa2c3 100644
--- a/pina/condition/input_equation_condition.py
+++ b/pina/condition/input_equation_condition.py
@@ -30,5 +30,5 @@ def __setattr__(self, key, value):
         elif key == 'equation':
             check_consistency(value, (EquationInterface))
             InputPointsEquationCondition.__dict__[key].__set__(self, value)
-        elif key in ('_condition_type', '_problem', 'problem', 'condition_type'):
+        elif key in ('problem', 'condition_type'):
             super().__setattr__(key, value)
\ No newline at end of file
diff --git a/pina/condition/input_output_condition.py b/pina/condition/input_output_condition.py
index a4fa48919..f8fd46e86 100644
--- a/pina/condition/input_output_condition.py
+++ b/pina/condition/input_output_condition.py
@@ -27,5 +27,5 @@ def __setattr__(self, key, value):
         if (key == 'input_points') or (key == 'output_points'):
             check_consistency(value, (LabelTensor, Graph, torch.Tensor))
             InputOutputPointsCondition.__dict__[key].__set__(self, value)
-        elif key in ('_condition_type', '_problem', 'problem', 'condition_type'):
+        elif key in ('problem', 'condition_type'):
             super().__setattr__(key, value)
diff --git a/pina/problem/abstract_problem.py b/pina/problem/abstract_problem.py
index edf214a82..600a68896 100644
--- a/pina/problem/abstract_problem.py
+++ b/pina/problem/abstract_problem.py
@@ -20,7 +20,7 @@ class AbstractProblem(metaclass=ABCMeta):
     def __init__(self):
 
         # create collector to manage problem data
-        self.collector = Collector(self)
+        self._collector = Collector(self)
 
         # create hook conditions <-> problems
         for condition_name in self.conditions:
@@ -33,7 +33,12 @@ def __init__(self):
         # points are ready.
         self.collector.store_fixed_data()
 
+    @property
+    def collector(self):
+        return self._collector
 
+    # TODO this should be erase when dataloading will interface collector,
+    # kept only for back compatibility
     @property
     def input_pts(self):
         to_return = {}
@@ -41,10 +46,6 @@ def input_pts(self):
             if 'input_points' in v.keys():
                 to_return[k] = v['input_points']
         return to_return
-
-    @property
-    def _have_sampled_points(self):
-        return self.collector.is_conditions_ready
     
     def __deepcopy__(self, memo):
         """
@@ -160,7 +161,9 @@ def discretise_domain(
 
         # check correct location
         if locations == "all":
-            locations = [name for name in self.conditions.keys()]
+            locations = [name for name in self.conditions.keys()
+                         if isinstance(self.conditions[name],
+                                       DomainEquationCondition)]
         else:
             if not isinstance(locations, (list)):
                 locations = [locations]
@@ -168,7 +171,7 @@ def discretise_domain(
             if not isinstance(self.conditions[loc], DomainEquationCondition):
                 raise TypeError(
                     f"Wrong locations passed, locations for sampling "
-                    f"should be in {[loc for loc in locations if not isinstance(self.conditions[loc], DomainEquationCondition)]}.",
+                    f"should be in {[loc for loc in locations if isinstance(self.conditions[loc], DomainEquationCondition)]}.",
                 )
 
         # store data
diff --git a/pina/trainer.py b/pina/trainer.py
index 758bbaaf0..ba18f3392 100644
--- a/pina/trainer.py
+++ b/pina/trainer.py
@@ -32,29 +32,18 @@ def __init__(self, solver, batch_size=None, **kwargs):
         if batch_size is not None:
             check_consistency(batch_size, int)
 
-        self._model = solver
+        self.solver = solver
         self.batch_size = batch_size
 
         self._create_loader()
         self._move_to_device()
 
-        # create dataloader
-        # if solver.problem.have_sampled_points is False:
-        #     raise RuntimeError(
-        #         f"Input points in {solver.problem.not_sampled_points} "
-        #         "training are None. Please "
-        #         "sample points in your problem by calling "
-        #         "discretise_domain function before train "
-        #         "in the provided locations."
-        #     )
-
-        # self._create_or_update_loader()
 
     def _move_to_device(self):
         device = self._accelerator_connector._parallel_devices[0]
         
         # move parameters to device
-        pb = self._model.problem
+        pb = self.solver.problem
         if hasattr(pb, "unknown_parameters"):
             for key in pb.unknown_parameters:
                 pb.unknown_parameters[key] = torch.nn.Parameter(
@@ -67,14 +56,21 @@ def _create_loader(self):
         during training, there is no need to define to touch the
         trainer dataloader, just call the method.
         """
+        if not self.solver.problem.collector.full:
+            error_message = '\n'.join(
+                [f'{" " * 13} ---> Condition {key} {"sampled" if value else "not sampled"}' 
+                for key, value in self.solver.problem.collector._is_conditions_ready.items()])
+            raise RuntimeError('Cannot create Trainer if not all conditions '
+                                'are sampled. The Trainer got the following:\n'
+                                f'{error_message}')
         devices = self._accelerator_connector._parallel_devices
 
         if len(devices) > 1:
             raise RuntimeError("Parallel training is not supported yet.")
 
         device = devices[0]
-        dataset_phys = SamplePointDataset(self._model.problem, device)
-        dataset_data = DataPointDataset(self._model.problem, device)
+        dataset_phys = SamplePointDataset(self.solver.problem, device)
+        dataset_data = DataPointDataset(self.solver.problem, device)
         self._loader = SamplePointLoader(
             dataset_phys, dataset_data, batch_size=self.batch_size, shuffle=True
         )
@@ -84,7 +80,7 @@ def train(self, **kwargs):
         Train the solver method.
         """
         return super().fit(
-            self._model, train_dataloaders=self._loader, **kwargs
+            self.solver, train_dataloaders=self._loader, **kwargs
         )
 
     @property
@@ -92,4 +88,4 @@ def solver(self):
         """
         Returning trainer solver.
         """
-        return self._model
+        return self._solver
diff --git a/tests/test_problem.py b/tests/test_problem.py
index cc7e25535..6d3759606 100644
--- a/tests/test_problem.py
+++ b/tests/test_problem.py
@@ -89,6 +89,7 @@ def test_discretise_domain():
     poisson_problem.discretise_domain(n, 'lh', locations=['D'])
     assert poisson_problem.input_pts['D'].shape[0] == n
 
+    poisson_problem.discretise_domain(n)
 
 def test_sampling_few_variables():
     n = 10
@@ -98,7 +99,7 @@ def test_sampling_few_variables():
                                       locations=['D'],
                                       variables=['x'])
     assert poisson_problem.input_pts['D'].shape[1] == 1
-    assert poisson_problem._have_sampled_points['D'] is False
+    assert poisson_problem.collector._is_conditions_ready['D'] is False
 
 
 def test_variables_correct_order_sampling():
@@ -140,3 +141,12 @@ def test_add_points():
     poisson_problem.add_points({'D': new_pts})
     assert torch.isclose(poisson_problem.input_pts['D'].extract('x'), new_pts.extract('x'))
     assert torch.isclose(poisson_problem.input_pts['D'].extract('y'), new_pts.extract('y'))
+
+
+def test_collector():
+    poisson_problem = Poisson()
+    collector = poisson_problem.collector
+    assert collector.full is False
+    assert collector._is_conditions_ready['data'] is True
+    poisson_problem.discretise_domain(10)
+    assert collector.full is True

From 7e37943990b7754dc132137194d1f56a8fd43010 Mon Sep 17 00:00:00 2001
From: FilippoOlivo <filippo@filippoolivo.com>
Date: Wed, 16 Oct 2024 11:24:37 +0200
Subject: [PATCH 24/55] Implement Dataset, Dataloader and DataModule class and
 fix SupervisedSolver

---
 pina/__init__.py                             |   6 +-
 pina/collector.py                            |  27 +-
 pina/condition/condition_interface.py        |   2 +-
 pina/condition/data_condition.py             |   4 +-
 pina/condition/domain_equation_condition.py  |   4 +-
 pina/condition/input_equation_condition.py   |   2 +-
 pina/condition/input_output_condition.py     |   4 +-
 pina/data/__init__.py                        |  19 +-
 pina/data/base_dataset.py                    | 107 ++++++++
 pina/data/data_dataset.py                    |  41 ---
 pina/data/data_module.py                     | 172 +++++++++++++
 pina/data/pina_batch.py                      |  57 ++---
 pina/data/pina_dataloader.py                 | 220 +++-------------
 pina/data/pina_subset.py                     |  21 ++
 pina/data/sample_dataset.py                  |  49 +---
 pina/data/supervised_dataset.py              |  12 +
 pina/data/unsupervised_dataset.py            |  13 +
 pina/domain/cartesian.py                     |   2 +-
 pina/domain/ellipsoid.py                     |   3 +-
 pina/domain/operation_interface.py           |   2 +-
 pina/domain/simplex.py                       |   4 +-
 pina/domain/union_domain.py                  |   6 +-
 pina/label_tensor.py                         |  32 +--
 pina/operators.py                            |   2 +-
 pina/solvers/pinns/basepinn.py               |  16 +-
 pina/solvers/solver.py                       | 255 +++++--------------
 pina/solvers/supervised.py                   |  56 ++--
 pina/trainer.py                              |  24 +-
 tests/test_dataset.py                        | 170 ++++++++-----
 tests/test_solvers/test_supervised_solver.py | 238 ++++++++---------
 30 files changed, 778 insertions(+), 792 deletions(-)
 create mode 100644 pina/data/base_dataset.py
 delete mode 100644 pina/data/data_dataset.py
 create mode 100644 pina/data/data_module.py
 create mode 100644 pina/data/pina_subset.py
 create mode 100644 pina/data/supervised_dataset.py
 create mode 100644 pina/data/unsupervised_dataset.py

diff --git a/pina/__init__.py b/pina/__init__.py
index 0fe93752d..d110d2842 100644
--- a/pina/__init__.py
+++ b/pina/__init__.py
@@ -5,7 +5,8 @@
     "Plotter",
     "Condition",
     "SamplePointDataset",
-    "SamplePointLoader",
+    "PinaDataModule",
+    "PinaDataLoader"
 ]
 
 from .meta import *
@@ -15,4 +16,5 @@
 from .plotter import Plotter
 from .condition.condition import Condition
 from .data import SamplePointDataset
-from .data import SamplePointLoader
+from .data import PinaDataModule
+from .data import PinaDataLoader
\ No newline at end of file
diff --git a/pina/collector.py b/pina/collector.py
index f44c222a7..f9ef194db 100644
--- a/pina/collector.py
+++ b/pina/collector.py
@@ -3,10 +3,11 @@
 from . import LabelTensor
 from .utils import check_consistency, merge_tensors
 
+
 class Collector:
     def __init__(self, problem):
         # creating a hook between collector and problem
-        self.problem = problem          
+        self.problem = problem
 
         # this variable is used to store the data in the form:
         # {'[condition_name]' : 
@@ -14,17 +15,17 @@ def __init__(self, problem):
         #            '[equation/output_points/conditional_variables]': Tensor}
         # }
         # those variables are used for the dataloading
-        self._data_collections = {name : {} for name in self.problem.conditions}
+        self._data_collections = {name: {} for name in self.problem.conditions}
 
         # variables used to check that all conditions are sampled
         self._is_conditions_ready = {
-            name : False for name in self.problem.conditions}
+            name: False for name in self.problem.conditions}
         self.full = False
-        
+
     @property
     def full(self):
         return all(self._is_conditions_ready.values())
-    
+
     @full.setter
     def full(self, value):
         check_consistency(value, bool)
@@ -37,7 +38,7 @@ def data_collections(self):
     @property
     def problem(self):
         return self._problem
-    
+
     @problem.setter
     def problem(self, value):
         self._problem = value
@@ -76,14 +77,14 @@ def store_sample_domains(self, n, mode, variables, sample_locations):
 
             # get the samples
             samples = [
-                condition.domain.sample(n=n, mode=mode, variables=variables)
-            ] + already_sampled
+                          condition.domain.sample(n=n, mode=mode, variables=variables)
+                      ] + already_sampled
             pts = merge_tensors(samples)
             if (
-                set(pts.labels).issubset(sorted(self.problem.input_variables))
-                ):
+                    set(pts.labels).issubset(sorted(self.problem.input_variables))
+            ):
                 pts = pts.sort_labels()
-                if sorted(pts.labels)==sorted(self.problem.input_variables):
+                if sorted(pts.labels) == sorted(self.problem.input_variables):
                     self._is_conditions_ready[loc] = True
                 values = [pts, condition.equation]
                 self.data_collections[loc] = dict(zip(keys, values))
@@ -97,7 +98,7 @@ def add_points(self, new_points_dict):
         :param new_points_dict: Dictonary of input points (condition_name: LabelTensor)
         :raises RuntimeError: if at least one condition is not already sampled
         """
-        for k,v in new_points_dict.items():
+        for k, v in new_points_dict.items():
             if not self._is_conditions_ready[k]:
                 raise RuntimeError('Cannot add points on a non sampled condition')
-            self.data_collections[k]['input_points'] = self.data_collections[k]['input_points'].vstack(v)
\ No newline at end of file
+            self.data_collections[k]['input_points'] = self.data_collections[k]['input_points'].vstack(v)
diff --git a/pina/condition/condition_interface.py b/pina/condition/condition_interface.py
index 52699b66e..808c06afe 100644
--- a/pina/condition/condition_interface.py
+++ b/pina/condition/condition_interface.py
@@ -5,7 +5,7 @@ class ConditionInterface(metaclass=ABCMeta):
 
     condition_types = ['physics', 'supervised', 'unsupervised']
 
-    def __init__(self, *args, **wargs):
+    def __init__(self, *args, **kwargs):
         self._condition_type = None
         self._problem = None
 
diff --git a/pina/condition/data_condition.py b/pina/condition/data_condition.py
index 90d248b67..3bcd4be6d 100644
--- a/pina/condition/data_condition.py
+++ b/pina/condition/data_condition.py
@@ -22,11 +22,11 @@ def __init__(self, input_points, conditional_variables=None):
         super().__init__()
         self.input_points = input_points
         self.conditional_variables = conditional_variables
-        self.condition_type = 'unsupervised'
+        self._condition_type = 'unsupervised'
 
     def __setattr__(self, key, value):
         if (key == 'input_points') or (key == 'conditional_variables'):
             check_consistency(value, (LabelTensor, Graph, torch.Tensor))
             DataConditionInterface.__dict__[key].__set__(self, value)
-        elif key in ('problem', 'condition_type'):
+        elif key in ('_problem', '_condition_type'):
             super().__setattr__(key, value)
\ No newline at end of file
diff --git a/pina/condition/domain_equation_condition.py b/pina/condition/domain_equation_condition.py
index ce4c7d3fc..28315655b 100644
--- a/pina/condition/domain_equation_condition.py
+++ b/pina/condition/domain_equation_condition.py
@@ -20,7 +20,7 @@ def __init__(self, domain, equation):
         super().__init__()
         self.domain = domain
         self.equation = equation
-        self.condition_type = 'physics'
+        self._condition_type = 'physics'
 
     def __setattr__(self, key, value):
         if key == 'domain':
@@ -29,5 +29,5 @@ def __setattr__(self, key, value):
         elif key == 'equation':
             check_consistency(value, (EquationInterface))
             DomainEquationCondition.__dict__[key].__set__(self, value)
-        elif key in ('problem', 'condition_type'):
+        elif key in ('_problem', '_condition_type'):
             super().__setattr__(key, value)
\ No newline at end of file
diff --git a/pina/condition/input_equation_condition.py b/pina/condition/input_equation_condition.py
index ac47fa2c3..0d34dfc93 100644
--- a/pina/condition/input_equation_condition.py
+++ b/pina/condition/input_equation_condition.py
@@ -30,5 +30,5 @@ def __setattr__(self, key, value):
         elif key == 'equation':
             check_consistency(value, (EquationInterface))
             InputPointsEquationCondition.__dict__[key].__set__(self, value)
-        elif key in ('problem', 'condition_type'):
+        elif key in ('_problem', '_condition_type'):
             super().__setattr__(key, value)
\ No newline at end of file
diff --git a/pina/condition/input_output_condition.py b/pina/condition/input_output_condition.py
index f8fd46e86..8a17495dd 100644
--- a/pina/condition/input_output_condition.py
+++ b/pina/condition/input_output_condition.py
@@ -21,11 +21,11 @@ def __init__(self, input_points, output_points):
         super().__init__()
         self.input_points = input_points
         self.output_points = output_points
-        self.condition_type = ['supervised', 'physics']
+        self._condition_type = ['supervised', 'physics']
 
     def __setattr__(self, key, value):
         if (key == 'input_points') or (key == 'output_points'):
             check_consistency(value, (LabelTensor, Graph, torch.Tensor))
             InputOutputPointsCondition.__dict__[key].__set__(self, value)
-        elif key in ('problem', 'condition_type'):
+        elif key in ('_problem', '_condition_type'):
             super().__setattr__(key, value)
diff --git a/pina/data/__init__.py b/pina/data/__init__.py
index fba19b92c..0a1b5905e 100644
--- a/pina/data/__init__.py
+++ b/pina/data/__init__.py
@@ -1,7 +1,20 @@
+"""
+Import data classes
+"""
 __all__ = [
+    'PinaDataLoader',
+    'SupervisedDataset',
+    'SamplePointDataset',
+    'UnsupervisedDataset',
+    'Batch',
+    'PinaDataModule',
+    'BaseDataset'
 ]
 
-from .pina_dataloader import SamplePointLoader
-from .data_dataset import DataPointDataset
+from .pina_dataloader import PinaDataLoader
+from .supervised_dataset import SupervisedDataset
 from .sample_dataset import SamplePointDataset
-from .pina_batch import Batch
\ No newline at end of file
+from .unsupervised_dataset import UnsupervisedDataset
+from .pina_batch import Batch
+from .data_module import PinaDataModule
+from .base_dataset import BaseDataset
diff --git a/pina/data/base_dataset.py b/pina/data/base_dataset.py
new file mode 100644
index 000000000..f095afa0c
--- /dev/null
+++ b/pina/data/base_dataset.py
@@ -0,0 +1,107 @@
+"""
+Basic data module implementation
+"""
+from torch.utils.data import Dataset
+import torch
+from ..label_tensor import LabelTensor
+
+
+class BaseDataset(Dataset):
+    """
+    BaseDataset class, which handle initialization and data retrieval
+    :var condition_indices: List of indices
+    :var device: torch.device
+    :var condition_names: dict of condition index and corresponding name
+    """
+
+    def __new__(cls, problem, device):
+        """
+        Ensure correct definition of __slots__ before initialization
+        :param AbstractProblem problem: The formulation of the problem.
+        :param torch.device device: The device on which the
+        dataset will be loaded.
+        """
+        if cls is BaseDataset:
+            raise TypeError('BaseDataset cannot be instantiated directly. Use a subclass.')
+        if not hasattr(cls, '__slots__'):
+            raise TypeError('Something is wrong, __slots__ must be defined in subclasses.')
+        return super().__new__(cls)
+
+    def __init__(self, problem, device):
+        """"
+        Initialize the object based on __slots__
+        :param AbstractProblem problem: The formulation of the problem.
+        :param torch.device device: The device on which the
+        dataset will be loaded.
+        """
+        super().__init__()
+
+        self.condition_names = {}
+        collector = problem.collector
+        for slot in self.__slots__:
+            setattr(self, slot, [])
+
+        idx = 0
+        for name, data in collector.data_collections.items():
+            keys = []
+            for k, v in data.items():
+                if isinstance(v, LabelTensor):
+                    keys.append(k)
+            if sorted(self.__slots__) == sorted(keys):
+
+                for slot in self.__slots__:
+                    current_list = getattr(self, slot)
+                    current_list.append(data[slot])
+                self.condition_names[idx] = name
+                idx += 1
+
+        if len(getattr(self, self.__slots__[0])) > 0:
+            input_list = getattr(self, self.__slots__[0])
+            self.condition_indices = torch.cat(
+                [
+                    torch.tensor([i] * len(input_list[i]), dtype=torch.uint8)
+                    for i in range(len(self.condition_names))
+                ],
+                dim=0,
+            )
+            for slot in self.__slots__:
+                current_attribute = getattr(self, slot)
+                setattr(self, slot, LabelTensor.vstack(current_attribute))
+        else:
+            self.condition_indices = torch.tensor([], dtype=torch.uint8)
+            for slot in self.__slots__:
+                setattr(self, slot, torch.tensor([]))
+
+        self.device = device
+
+    def __len__(self):
+        return len(getattr(self, self.__slots__[0]))
+
+    def __getattribute__(self, item):
+        attribute = super().__getattribute__(item)
+        if isinstance(attribute, LabelTensor) and attribute.dtype == torch.float32:
+            attribute = attribute.to(device=self.device).requires_grad_()
+        return attribute
+
+    def __getitem__(self, idx):
+        if isinstance(idx, str):
+            return getattr(self, idx).to(self.device)
+
+        if isinstance(idx, slice):
+            to_return_list = []
+            for i in self.__slots__:
+                to_return_list.append(getattr(self, i)[[idx]].to(self.device))
+            return to_return_list
+
+        if isinstance(idx, (tuple, list)):
+            if (len(idx) == 2 and isinstance(idx[0], str)
+                    and isinstance(idx[1], (list, slice))):
+                tensor = getattr(self, idx[0])
+                return tensor[[idx[1]]].to(self.device)
+            if all(isinstance(x, int) for x in idx):
+                to_return_list = []
+                for i in self.__slots__:
+                    to_return_list.append(getattr(self, i)[[idx]].to(self.device))
+                return to_return_list
+
+        raise ValueError(f'Invalid index {idx}')
diff --git a/pina/data/data_dataset.py b/pina/data/data_dataset.py
deleted file mode 100644
index 9dff2d7ed..000000000
--- a/pina/data/data_dataset.py
+++ /dev/null
@@ -1,41 +0,0 @@
-from torch.utils.data import Dataset
-import torch
-from ..label_tensor import LabelTensor
-
-
-class DataPointDataset(Dataset):
-
-    def __init__(self, problem, device) -> None:
-        super().__init__()
-        input_list = []
-        output_list = []
-        self.condition_names = []
-
-        for name, condition in problem.conditions.items():
-            if hasattr(condition, "output_points"):
-                input_list.append(problem.conditions[name].input_points)
-                output_list.append(problem.conditions[name].output_points)
-                self.condition_names.append(name)
-
-        self.input_pts = LabelTensor.stack(input_list)
-        self.output_pts = LabelTensor.stack(output_list)
-
-        if self.input_pts != []:
-            self.condition_indeces = torch.cat(
-                [
-                    torch.tensor([i] * len(input_list[i]))
-                    for i in range(len(self.condition_names))
-                ],
-                dim=0,
-            )
-        else:  # if there are no data points
-            self.condition_indeces = torch.tensor([])
-            self.input_pts = torch.tensor([])
-            self.output_pts = torch.tensor([])
-
-        self.input_pts = self.input_pts.to(device)
-        self.output_pts = self.output_pts.to(device)
-        self.condition_indeces = self.condition_indeces.to(device)
-
-    def __len__(self):
-        return self.input_pts.shape[0]
\ No newline at end of file
diff --git a/pina/data/data_module.py b/pina/data/data_module.py
new file mode 100644
index 000000000..e4e8a450f
--- /dev/null
+++ b/pina/data/data_module.py
@@ -0,0 +1,172 @@
+"""
+This module provide basic data management functionalities
+"""
+
+import math
+import torch
+from lightning import LightningDataModule
+from .sample_dataset import SamplePointDataset
+from .supervised_dataset import SupervisedDataset
+from .unsupervised_dataset import UnsupervisedDataset
+from .pina_dataloader import PinaDataLoader
+from .pina_subset import PinaSubset
+
+
+class PinaDataModule(LightningDataModule):
+    """
+    This class extend LightningDataModule, allowing proper creation and
+    management of different types of Datasets defined in PINA
+    """
+
+    def __init__(self,
+                 problem,
+                 device,
+                 train_size=.7,
+                 test_size=.2,
+                 eval_size=.1,
+                 batch_size=None,
+                 shuffle=True,
+                 datasets = None):
+        """
+        Initialize the object, creating dataset based on input problem
+        :param AbstractProblem problem: PINA problem
+        :param device: Device used for training and testing
+        :param train_size: number/percentage of elements in train split
+        :param test_size: number/percentage of elements in test split
+        :param eval_size: number/percentage of elements in evaluation split
+        :param batch_size: batch size used for training
+        :param datasets: list of datasets objects
+        """
+        super().__init__()
+        dataset_classes = [SupervisedDataset, UnsupervisedDataset, SamplePointDataset]
+        if datasets is None:
+            self.datasets = [DatasetClass(problem, device) for DatasetClass in dataset_classes]
+        else:
+            self.datasets = datasets
+
+        self.split_length = []
+        self.split_names = []
+        if train_size > 0:
+            self.split_names.append('train')
+            self.split_length.append(train_size)
+        if test_size > 0:
+            self.split_length.append(test_size)
+            self.split_names.append('test')
+        if eval_size > 0:
+            self.split_length.append(eval_size)
+            self.split_names.append('eval')
+
+        self.batch_size = batch_size
+        self.condition_names = None
+        self.splits = {k: {} for k in self.split_names}
+        self.shuffle = shuffle
+
+    def setup(self, stage=None):
+        """
+        Perform the splitting of the dataset
+        """
+        self.extract_conditions()
+        if stage == 'fit' or stage is None:
+            for dataset in self.datasets:
+                if len(dataset) > 0:
+                    splits = self.dataset_split(dataset,
+                                                self.split_length,
+                                                shuffle=self.shuffle)
+                    for i in range(len(self.split_length)):
+                        self.splits[
+                            self.split_names[i]][dataset.data_type] = splits[i]
+        elif stage == 'test':
+            raise NotImplementedError("Testing pipeline not implemented yet")
+        else:
+            raise ValueError("stage must be either 'fit' or 'test'")
+
+    def extract_conditions(self):
+        """
+        Extract conditions from dataset and update condition indices
+        """
+        # Extract number of conditions
+        n_conditions = 0
+        for dataset in self.datasets:
+            if n_conditions != 0:
+                dataset.condition_names = {
+                    key + n_conditions: value
+                    for key, value in dataset.condition_names.items()
+                }
+            n_conditions += len(dataset.condition_names)
+
+        self.condition_names = {
+            key: value
+            for dataset in self.datasets
+            for key, value in dataset.condition_names.items()
+        }
+
+
+
+    def train_dataloader(self):
+        """
+        Return the training dataloader for the dataset
+        :return: data loader
+        :rtype: PinaDataLoader
+        """
+        return PinaDataLoader(self.splits['train'], self.batch_size,
+                              self.condition_names)
+
+    def test_dataloader(self):
+        """
+        Return the testing dataloader for the dataset
+        :return: data loader
+        :rtype: PinaDataLoader
+        """
+        return PinaDataLoader(self.splits['test'], self.batch_size,
+                              self.condition_names)
+
+    def eval_dataloader(self):
+        """
+        Return the evaluation dataloader for the dataset
+        :return: data loader
+        :rtype: PinaDataLoader
+        """
+        return PinaDataLoader(self.splits['eval'], self.batch_size,
+                              self.condition_names)
+
+    @staticmethod
+    def dataset_split(dataset, lengths, seed=None, shuffle=True):
+        """
+        Perform the splitting of the dataset
+        :param dataset: dataset object we wanted to split
+        :param lengths: lengths of elements in dataset
+        :param seed: random seed
+        :param shuffle: shuffle dataset
+        :return: split dataset
+        :rtype: PinaSubset
+        """
+        if sum(lengths) - 1 < 1e-3:
+            lengths = [
+                int(math.floor(len(dataset) * length)) for length in lengths
+            ]
+
+            remainder = len(dataset) - sum(lengths)
+            for i in range(remainder):
+                lengths[i % len(lengths)] += 1
+        elif sum(lengths) - 1 >= 1e-3:
+            raise ValueError(f"Sum of lengths is {sum(lengths)} less than 1")
+
+        if sum(lengths) != len(dataset):
+            raise ValueError("Sum of lengths is not equal to dataset length")
+
+        if shuffle:
+            if seed is not None:
+                generator = torch.Generator()
+                generator.manual_seed(seed)
+                indices = torch.randperm(sum(lengths), generator=generator).tolist()
+            else:
+                indices = torch.arange(sum(lengths)).tolist()
+        else:
+            indices = torch.arange(0, sum(lengths), 1, dtype=torch.uint8).tolist()
+        offsets = [
+            sum(lengths[:i]) if i > 0 else 0 for i in range(len(lengths))
+        ]
+        return [
+            PinaSubset(dataset, indices[offset:offset + length])
+            for offset, length in zip(offsets, lengths)
+        ]
diff --git a/pina/data/pina_batch.py b/pina/data/pina_batch.py
index cb1296ede..7e46a2218 100644
--- a/pina/data/pina_batch.py
+++ b/pina/data/pina_batch.py
@@ -1,36 +1,33 @@
+"""
+Batch management module
+"""
+from .pina_subset import PinaSubset
 
 
 class Batch:
-    """
-    This class is used to create a dataset of sample points.
-    """
+    def __init__(self, dataset_dict, idx_dict):
 
-    def __init__(self, type_, idx, *args, **kwargs) -> None:
-        """
-        """
-        if type_ == "sample":
-
-            if len(args) != 2:
-                raise RuntimeError
-
-            input = args[0]
-            conditions = args[1]
-
-            self.input = input[idx]
-            self.condition = conditions[idx]
+        for k, v in dataset_dict.items():
+            setattr(self, k, v)
 
-        elif type_ == "data":
+        for k, v in idx_dict.items():
+            setattr(self, k + '_idx', v)
 
-            if len(args) != 3:
-                raise RuntimeError
-
-            input = args[0]
-            output = args[1]
-            conditions = args[2]
-
-            self.input = input[idx]
-            self.output = output[idx]
-            self.condition = conditions[idx]
-                
-        else:
-            raise ValueError("Invalid number of arguments.")
\ No newline at end of file
+    def __len__(self):
+        """
+        Returns the number of elements in the batch
+        :return: number of elements in the batch
+        :rtype: int
+        """
+        length = 0
+        for dataset in dir(self):
+            attribute = getattr(self, dataset)
+            if isinstance(attribute, list):
+                length += len(getattr(self, dataset))
+        return length
+
+    def __getattr__(self, item):
+        if not item in dir(self):
+            raise AttributeError(f'Batch instance has no attribute {item}')
+        return PinaSubset(getattr(self, item).dataset,
+                          getattr(self, item).indices[self.coordinates_dict[item]])
diff --git a/pina/data/pina_dataloader.py b/pina/data/pina_dataloader.py
index 2c8967c50..d62847574 100644
--- a/pina/data/pina_dataloader.py
+++ b/pina/data/pina_dataloader.py
@@ -1,11 +1,11 @@
-import torch
-
-from .sample_dataset import SamplePointDataset
-from .data_dataset import DataPointDataset
+"""
+This module is used to create an iterable object used during training
+"""
+import math
 from .pina_batch import Batch
 
 
-class SamplePointLoader:
+class PinaDataLoader:
     """
     This class is used to create a dataloader to use during the training.
 
@@ -14,198 +14,54 @@ class SamplePointLoader:
     :vartype condition_names: list[str]
     """
 
-    def __init__(
-        self, sample_dataset, data_dataset, batch_size=None, shuffle=True
-    ) -> None:
-        """
-        Constructor.
-
-        :param SamplePointDataset sample_pts: The sample points dataset.
-        :param int batch_size: The batch size. If ``None``, the batch size is
-            set to the number of sample points. Default is ``None``.
-        :param bool shuffle: If ``True``, the sample points are shuffled.
-            Default is ``True``.
-        """
-        if not isinstance(sample_dataset, SamplePointDataset):
-            raise TypeError(
-                f"Expected SamplePointDataset, got {type(sample_dataset)}"
-            )
-        if not isinstance(data_dataset, DataPointDataset):
-            raise TypeError(
-                f"Expected DataPointDataset, got {type(data_dataset)}"
-            )
-
-        self.n_data_conditions = len(data_dataset.condition_names)
-        self.n_phys_conditions = len(sample_dataset.condition_names)
-        data_dataset.condition_indeces += self.n_phys_conditions
-
-        self._prepare_sample_dataset(sample_dataset, batch_size, shuffle)
-        self._prepare_data_dataset(data_dataset, batch_size, shuffle)
-
-        self.condition_names = (
-            sample_dataset.condition_names + data_dataset.condition_names
-        )
-
-        self.batch_list = []
-        for i in range(len(self.batch_sample_pts)):
-            self.batch_list.append(("sample", i))
-
-        for i in range(len(self.batch_input_pts)):
-            self.batch_list.append(("data", i))
-
-        if shuffle:
-            self.random_idx = torch.randperm(len(self.batch_list))
-        else:
-            self.random_idx = torch.arange(len(self.batch_list))
-
-        self._prepare_batches()
-
-    def _prepare_data_dataset(self, dataset, batch_size, shuffle):
-        """
-        Prepare the dataset for data points.
-
-        :param SamplePointDataset dataset: The dataset.
-        :param int batch_size: The batch size.
-        :param bool shuffle: If ``True``, the sample points are shuffled.
-        """
-        self.sample_dataset = dataset
-
-        if len(dataset) == 0:
-            self.batch_data_conditions = []
-            self.batch_input_pts = []
-            self.batch_output_pts = []
-            return
-
-        if batch_size is None:
-            batch_size = len(dataset)
-        batch_num = len(dataset) // batch_size
-        if len(dataset) % batch_size != 0:
-            batch_num += 1
-
-        output_labels = dataset.output_pts.labels
-        input_labels = dataset.input_pts.labels
-        self.tensor_conditions = dataset.condition_indeces
-
-        if shuffle:
-            idx = torch.randperm(dataset.input_pts.shape[0])
-            self.input_pts = dataset.input_pts[idx]
-            self.output_pts = dataset.output_pts[idx]
-            self.tensor_conditions = dataset.condition_indeces[idx]
-
-        self.batch_input_pts = torch.tensor_split(dataset.input_pts, batch_num)
-        self.batch_output_pts = torch.tensor_split(
-            dataset.output_pts, batch_num
-        )
-        #print(input_labels)
-        for i in range(len(self.batch_input_pts)):
-            self.batch_input_pts[i].labels = input_labels
-            self.batch_output_pts[i].labels = output_labels
-
-        self.batch_data_conditions = torch.tensor_split(
-            self.tensor_conditions, batch_num
-        )
-
-    def _prepare_sample_dataset(self, dataset, batch_size, shuffle):
+    def __init__(self, dataset_dict, batch_size, condition_names) -> None:
         """
-        Prepare the dataset for sample points.
-
-        :param DataPointDataset dataset: The dataset.
-        :param int batch_size: The batch size.
-        :param bool shuffle: If ``True``, the sample points are shuffled.
+        Initialize local variables
+        :param dataset_dict: Dictionary of datasets
+        :type dataset_dict: dict
+        :param batch_size: Size of the batch
+        :type batch_size: int
+        :param condition_names: Names of the conditions
+        :type condition_names: list[str]
         """
+        self.condition_names = condition_names
+        self.dataset_dict = dataset_dict
+        self._init_batches(batch_size)
 
-        self.sample_dataset = dataset
-        if len(dataset) == 0:
-            self.batch_sample_conditions = []
-            self.batch_sample_pts = []
-            return
-
-        if batch_size is None:
-            batch_size = len(dataset)
-
-        batch_num = len(dataset) // batch_size
-        if len(dataset) % batch_size != 0:
-            batch_num += 1
-
-        self.tensor_pts = dataset.pts
-        self.tensor_conditions = dataset.condition_indeces
-
-        # if shuffle:
-        #     idx = torch.randperm(self.tensor_pts.shape[0])
-        #     self.tensor_pts = self.tensor_pts[idx]
-        #     self.tensor_conditions = self.tensor_conditions[idx]
-
-        self.batch_sample_pts = torch.tensor_split(self.tensor_pts, batch_num)
-        for i in range(len(self.batch_sample_pts)):
-            self.batch_sample_pts[i].labels = dataset.pts.labels
-
-        self.batch_sample_conditions = torch.tensor_split(
-            self.tensor_conditions, batch_num
-        )
-
-    def _prepare_batches(self):
+    def _init_batches(self, batch_size=None):
         """
-        Prepare the batches.
+        Create batches according to the batch_size provided in input.
         """
         self.batches = []
-        for i in range(len(self.batch_list)):
-            type_, idx_ = self.batch_list[i]
-
-            if type_ == "sample":
-                batch = Batch(
-                    "sample", idx_,
-                    self.batch_sample_pts,
-                    self.batch_sample_conditions)
+        n_elements = sum([len(v) for v in self.dataset_dict.values()])
+        if batch_size is None:
+            batch_size = n_elements
+        indexes_dict = {}
+        n_batches = int(math.ceil(n_elements / batch_size))
+        for k, v in self.dataset_dict.items():
+            if n_batches != 1:
+                indexes_dict[k] = math.floor(len(v) / (n_batches - 1))
             else:
-                batch = Batch(
-                    "data", idx_,
-                    self.batch_input_pts,
-                    self.batch_output_pts,
-                    self.batch_data_conditions)
-                
-            self.batches.append(batch)
+                indexes_dict[k] = len(v)
+        for i in range(n_batches):
+            temp_dict = {}
+            for k, v in indexes_dict.items():
+                if i != n_batches - 1:
+                    temp_dict[k] = slice(i * v, (i + 1) * v)
+                else:
+                    temp_dict[k] = slice(i * v, len(self.dataset_dict[k]))
+            self.batches.append(Batch(idx_dict=temp_dict, dataset_dict=self.dataset_dict))
 
     def __iter__(self):
         """
-        Return an iterator over the points. Any element of the iterator is a
-        dictionary with the following keys:
-            - ``pts``: The input sample points. It is a LabelTensor with the
-                shape ``(batch_size, input_dimension)``.
-            - ``output``: The output sample points. This key is present only
-                if data conditions are present. It is a LabelTensor with the
-                shape ``(batch_size, output_dimension)``.
-            - ``condition``: The integer condition indeces. It is a tensor
-                with the shape ``(batch_size, )`` of type ``torch.int64`` and
-                indicates for any ``pts`` the corresponding problem condition.
-
-        :return: An iterator over the points.
-        :rtype: iter
+        Makes dataloader object iterable
         """
-        # for i in self.random_idx:
-        for i in self.random_idx:
-            yield self.batches[i]
-
-        # for i in range(len(self.batch_list)):
-        #     type_, idx_ = self.batch_list[i]
-
-        #     if type_ == "sample":
-        #         d = {
-        #             "pts": self.batch_sample_pts[idx_].requires_grad_(True),
-        #             "condition": self.batch_sample_conditions[idx_],
-        #         }
-        #     else:
-        #         d = {
-        #             "pts": self.batch_input_pts[idx_].requires_grad_(True),
-        #             "output": self.batch_output_pts[idx_],
-        #             "condition": self.batch_data_conditions[idx_],
-        #         }
-        #     yield d
+        yield from self.batches
 
     def __len__(self):
         """
         Return the number of batches.
-
         :return: The number of batches.
         :rtype: int
         """
-        return len(self.batch_list)
+        return len(self.batches)
diff --git a/pina/data/pina_subset.py b/pina/data/pina_subset.py
new file mode 100644
index 000000000..41571f92b
--- /dev/null
+++ b/pina/data/pina_subset.py
@@ -0,0 +1,21 @@
+class PinaSubset:
+    """
+    TODO
+    """
+    __slots__ = ['dataset', 'indices']
+
+    def __init__(self, dataset, indices):
+        """
+        TODO
+        """
+        self.dataset = dataset
+        self.indices = indices
+
+    def __len__(self):
+        """
+        TODO
+        """
+        return len(self.indices)
+
+    def __getattr__(self, name):
+        return self.dataset.__getattribute__(name)
diff --git a/pina/data/sample_dataset.py b/pina/data/sample_dataset.py
index 84af2920f..ba8bd19a9 100644
--- a/pina/data/sample_dataset.py
+++ b/pina/data/sample_dataset.py
@@ -1,43 +1,12 @@
-from torch.utils.data import Dataset
-import torch
+"""
+Sample dataset module
+"""
+from .base_dataset import BaseDataset
 
-from ..label_tensor import LabelTensor
-
-
-class SamplePointDataset(Dataset):
+class SamplePointDataset(BaseDataset):
     """
-    This class is used to create a dataset of sample points.
+    This class extends the BaseDataset to handle physical datasets
+    composed of only input points.
     """
-
-    def __init__(self, problem, device) -> None:
-        """
-        :param dict input_pts: The input points.
-        """
-        super().__init__()
-        pts_list = []
-        self.condition_names = []
-
-        for name, condition in problem.conditions.items():
-            if not hasattr(condition, "output_points"):
-                pts_list.append(problem.input_pts[name])
-                self.condition_names.append(name)
-
-        self.pts = LabelTensor.stack(pts_list)
-
-        if self.pts != []:
-            self.condition_indeces = torch.cat(
-                [
-                    torch.tensor([i] * len(pts_list[i]))
-                    for i in range(len(self.condition_names))
-                ],
-                dim=0,
-            )
-        else:  # if there are no sample points
-            self.condition_indeces = torch.tensor([])
-            self.pts = torch.tensor([])
-
-        self.pts = self.pts.to(device)
-        self.condition_indeces = self.condition_indeces.to(device)
-
-    def __len__(self):
-        return self.pts.shape[0]
\ No newline at end of file
+    data_type = 'physics'
+    __slots__ = ['input_points']
diff --git a/pina/data/supervised_dataset.py b/pina/data/supervised_dataset.py
new file mode 100644
index 000000000..2403e3d0c
--- /dev/null
+++ b/pina/data/supervised_dataset.py
@@ -0,0 +1,12 @@
+"""
+Supervised dataset module
+"""
+from .base_dataset import BaseDataset
+
+
+class SupervisedDataset(BaseDataset):
+    """
+    This class extends the BaseDataset to handle datasets that consist of input-output pairs.
+    """
+    data_type = 'supervised'
+    __slots__ = ['input_points', 'output_points']
diff --git a/pina/data/unsupervised_dataset.py b/pina/data/unsupervised_dataset.py
new file mode 100644
index 000000000..f4e8fb345
--- /dev/null
+++ b/pina/data/unsupervised_dataset.py
@@ -0,0 +1,13 @@
+"""
+Unsupervised dataset module
+"""
+from .base_dataset import BaseDataset
+
+
+class UnsupervisedDataset(BaseDataset):
+    """
+    This class extend BaseDataset class to handle unsupervised dataset,
+    composed of input points and, optionally, conditional variables
+    """
+    data_type = 'unsupervised'
+    __slots__ = ['input_points', 'conditional_variables']
diff --git a/pina/domain/cartesian.py b/pina/domain/cartesian.py
index c605e331c..a6c369339 100644
--- a/pina/domain/cartesian.py
+++ b/pina/domain/cartesian.py
@@ -33,7 +33,7 @@ def __init__(self, cartesian_dict):
     @property
     def sample_modes(self):
         return ["random", "grid", "lh", "chebyshev", "latin"]
-    
+
     @property
     def variables(self):
         """Spatial variables.
diff --git a/pina/domain/ellipsoid.py b/pina/domain/ellipsoid.py
index e56e86c66..132bd55a5 100644
--- a/pina/domain/ellipsoid.py
+++ b/pina/domain/ellipsoid.py
@@ -55,7 +55,6 @@ def __init__(self, ellipsoid_dict, sample_surface=False):
         # perform operation only for not fixed variables (if any)
 
         if self.range_:
-
             # convert dict vals to torch [dim, 2] matrix
             list_dict_vals = list(self.range_.values())
             tmp = torch.tensor(list_dict_vals, dtype=torch.float)
@@ -74,7 +73,7 @@ def __init__(self, ellipsoid_dict, sample_surface=False):
     @property
     def sample_modes(self):
         return ["random"]
-    
+
     @property
     def variables(self):
         """Spatial variables.
diff --git a/pina/domain/operation_interface.py b/pina/domain/operation_interface.py
index a1efec91f..0300f5248 100644
--- a/pina/domain/operation_interface.py
+++ b/pina/domain/operation_interface.py
@@ -69,4 +69,4 @@ def _check_dimensions(self, geometries):
             if geometry.variables != geometries[0].variables:
                 raise NotImplementedError(
                     f"The geometries need to have same dimensions and labels."
-                )
\ No newline at end of file
+                )
diff --git a/pina/domain/simplex.py b/pina/domain/simplex.py
index cea213265..00bd980f3 100644
--- a/pina/domain/simplex.py
+++ b/pina/domain/simplex.py
@@ -77,7 +77,7 @@ def __init__(self, simplex_matrix, sample_surface=False):
     @property
     def sample_modes(self):
         return ["random"]
-       
+
     @property
     def variables(self):
         return sorted(self._vertices_matrix.labels)
@@ -144,7 +144,7 @@ def is_inside(self, point, check_border=False):
             return all(torch.gt(lambdas, 0.0)) and all(torch.lt(lambdas, 1.0))
 
         return all(torch.ge(lambdas, 0)) and (
-            any(torch.eq(lambdas, 0)) or any(torch.eq(lambdas, 1))
+                any(torch.eq(lambdas, 0)) or any(torch.eq(lambdas, 1))
         )
 
     def _sample_interior_randomly(self, n, variables):
diff --git a/pina/domain/union_domain.py b/pina/domain/union_domain.py
index a72115f50..0af8e1bd1 100644
--- a/pina/domain/union_domain.py
+++ b/pina/domain/union_domain.py
@@ -37,13 +37,13 @@ def __init__(self, geometries):
     def sample_modes(self):
         self.sample_modes = list(
             set([geom.sample_modes for geom in self.geometries])
-            )
-    
+        )
+
     @property
     def variables(self):
         variables = []
         for geom in self.geometries:
-            variables+=geom.variables
+            variables += geom.variables
         return list(set(variables))
 
     def is_inside(self, point, check_border=False):
diff --git a/pina/label_tensor.py b/pina/label_tensor.py
index 1df318ec7..65655e9dc 100644
--- a/pina/label_tensor.py
+++ b/pina/label_tensor.py
@@ -3,6 +3,7 @@
 import torch
 from torch import Tensor
 
+
 def issubset(a, b):
     """
     Check if a is a subset of b.
@@ -45,7 +46,7 @@ def labels(self):
         :return: labels of self
         :rtype: list
         """
-        return self._labels[self.tensor.ndim-1]['dof']
+        return self._labels[self.tensor.ndim - 1]['dof']
 
     @property
     def full_labels(self):
@@ -103,23 +104,23 @@ def extract(self, label_to_extract):
             raise ValueError('labels_to_extract must be str or list or dict')
 
     def _extract_from_list(self, labels_to_extract):
-        #Store locally all necessary obj/variables
+        # Store locally all necessary obj/variables
         ndim = self.tensor.ndim
         labels = self.full_labels
         tensor = self.tensor
         last_dim_label = self.labels
 
-        #Verify if all the labels in labels_to_extract are in last dimension
+        # Verify if all the labels in labels_to_extract are in last dimension
         if set(labels_to_extract).issubset(last_dim_label) is False:
             raise ValueError('Cannot extract a dof which is not in the original LabelTensor')
 
-        #Extract index to extract
+        # Extract index to extract
         idx_to_extract = [last_dim_label.index(i) for i in labels_to_extract]
 
-        #Perform extraction
+        # Perform extraction
         new_tensor = tensor[..., idx_to_extract]
 
-        #Manage labels
+        # Manage labels
         new_labels = copy(labels)
 
         last_dim_new_label = {ndim - 1: {
@@ -186,7 +187,7 @@ def cat(tensors, dim=0):
         # Perform cat on tensors
         new_tensor = torch.cat(tensors, dim=dim)
 
-        #Update labels
+        # Update labels
         labels = tensors[0].full_labels
         labels.pop(dim)
         new_labels_cat_dim = new_labels_cat_dim if len(set(new_labels_cat_dim)) == len(new_labels_cat_dim) \
@@ -265,13 +266,13 @@ def update_labels_from_dict(self, labels):
         :raises ValueError: dof list contain duplicates or number of dof does not match with tensor shape
         """
         tensor_shape = self.tensor.shape
-        #Check dimensionality
+        # Check dimensionality
         for k, v in labels.items():
             if len(v['dof']) != len(set(v['dof'])):
                 raise ValueError("dof must be unique")
             if len(v['dof']) != tensor_shape[k]:
                 raise ValueError('Number of dof does not match with tensor dimension')
-        #Perform update
+        # Perform update
         self._labels.update(labels)
 
     def update_labels_from_list(self, labels):
@@ -310,7 +311,7 @@ def append(self, tensor, mode='std'):
         if mode == 'std':
             # Call cat on last dimension
             new_label_tensor = LabelTensor.cat([self, tensor], dim=self.tensor.ndim - 1)
-        elif mode=='cross':
+        elif mode == 'cross':
             # Crete tensor and call cat on last dimension
             tensor1 = self
             tensor2 = tensor
@@ -318,7 +319,7 @@ def append(self, tensor, mode='std'):
             n2 = tensor2.shape[0]
             tensor1 = LabelTensor(tensor1.repeat(n2, 1), labels=tensor1.labels)
             tensor2 = LabelTensor(tensor2.repeat_interleave(n1, dim=0), labels=tensor2.labels)
-            new_label_tensor = LabelTensor.cat([tensor1, tensor2], dim=self.tensor.ndim-1)
+            new_label_tensor = LabelTensor.cat([tensor1, tensor2], dim=self.tensor.ndim - 1)
         else:
             raise ValueError('mode must be either "std" or "cross"')
         return new_label_tensor
@@ -366,10 +367,10 @@ def __getitem__(self, index):
                 if hasattr(self, "labels"):
                     if isinstance(index[j], list):
                         new_labels.update({j: {'dof': [new_labels[j]['dof'][i] for i in index[1]],
-                                               'name':new_labels[j]['name']}})
+                                               'name': new_labels[j]['name']}})
                     else:
                         new_labels.update({j: {'dof': new_labels[j]['dof'][index[j]],
-                                               'name':new_labels[j]['name']}})
+                                               'name': new_labels[j]['name']}})
 
             selected_lt.labels = new_labels
         else:
@@ -382,12 +383,13 @@ def __getitem__(self, index):
     def sort_labels(self, dim=None):
         def argsort(lst):
             return sorted(range(len(lst)), key=lambda x: lst[x])
+
         if dim is None:
-            dim = self.tensor.ndim-1
+            dim = self.tensor.ndim - 1
         labels = self.full_labels[dim]['dof']
         sorted_index = argsort(labels)
         indexer = [slice(None)] * self.tensor.ndim
         indexer[dim] = sorted_index
         new_labels = deepcopy(self.full_labels)
         new_labels[dim] = {'dof': sorted(labels), 'name': new_labels[dim]['name']}
-        return LabelTensor(self.tensor[indexer], new_labels)
\ No newline at end of file
+        return LabelTensor(self.tensor[indexer], new_labels)
diff --git a/pina/operators.py b/pina/operators.py
index 8822f2010..083837c11 100644
--- a/pina/operators.py
+++ b/pina/operators.py
@@ -249,7 +249,7 @@ def scalar_laplace(output_, input_, components, d):
                 result[:, idx] = grad(grad_output, input_, d=di).flatten()
                 to_append_tensors[idx] = grad(grad_output, input_, d=di)
                 labels[idx] = f"dd{ci[0]}dd{di[0]}"
-            result = LabelTensor.cat(tensors=to_append_tensors, dim=output_.tensor.ndim-1)
+            result = LabelTensor.cat(tensors=to_append_tensors, dim=output_.tensor.ndim - 1)
             result.labels = labels
     return result
 
diff --git a/pina/solvers/pinns/basepinn.py b/pina/solvers/pinns/basepinn.py
index a15834da4..a8315b259 100644
--- a/pina/solvers/pinns/basepinn.py
+++ b/pina/solvers/pinns/basepinn.py
@@ -27,13 +27,13 @@ class PINNInterface(SolverInterface, metaclass=ABCMeta):
     """
 
     def __init__(
-        self,
-        models,
-        problem,
-        optimizers,
-        optimizers_kwargs,
-        extra_features,
-        loss,
+            self,
+            models,
+            problem,
+            optimizers,
+            optimizers_kwargs,
+            extra_features,
+            loss,
     ):
         """
         :param models: Multiple torch neural network models instances.
@@ -178,7 +178,7 @@ def compute_residual(self, samples, equation):
         try:
             residual = equation.residual(samples, self.forward(samples))
         except (
-            TypeError
+                TypeError
         ):  # this occurs when the function has three inputs, i.e. inverse problem
             residual = equation.residual(
                 samples, self.forward(samples), self._params
diff --git a/pina/solvers/solver.py b/pina/solvers/solver.py
index a27e93641..8b3ddae71 100644
--- a/pina/solvers/solver.py
+++ b/pina/solvers/solver.py
@@ -10,168 +10,6 @@
 import sys
 
 
-# class SolverInterface(pytorch_lightning.LightningModule, metaclass=ABCMeta):
-#     """
-#     Solver base class. This class inherits is a wrapper of
-#     LightningModule class, inheriting all the
-#     LightningModule methods.
-#     """
-
-#     def __init__(
-#         self,
-#         models,
-#         problem,
-#         optimizers,
-#         optimizers_kwargs,
-#         extra_features=None,
-#     ):
-#         """
-#         :param models: A torch neural network model instance.
-#         :type models: torch.nn.Module
-#         :param problem: A problem definition instance.
-#         :type problem: AbstractProblem
-#         :param list(torch.optim.Optimizer) optimizer: A list of neural network optimizers to
-#             use.
-#         :param list(dict) optimizer_kwargs: A list of optimizer constructor keyword args.
-#         :param list(torch.nn.Module) extra_features: The additional input
-#             features to use as augmented input. If ``None`` no extra features
-#             are passed. If it is a list of :class:`torch.nn.Module`, the extra feature
-#             list is passed to all models. If it is a list of extra features' lists,
-#             each single list of extra feature is passed to a model.
-#         """
-#         super().__init__()
-
-#         # check consistency of the inputs
-#         check_consistency(models, torch.nn.Module)
-#         check_consistency(problem, AbstractProblem)
-#         check_consistency(optimizers, torch.optim.Optimizer, subclass=True)
-#         check_consistency(optimizers_kwargs, dict)
-
-#         # put everything in a list if only one input
-#         if not isinstance(models, list):
-#             models = [models]
-#         if not isinstance(optimizers, list):
-#             optimizers = [optimizers]
-#             optimizers_kwargs = [optimizers_kwargs]
-
-#         # number of models and optimizers
-#         len_model = len(models)
-#         len_optimizer = len(optimizers)
-#         len_optimizer_kwargs = len(optimizers_kwargs)
-
-#         # check length consistency optimizers
-#         if len_model != len_optimizer:
-#             raise ValueError(
-#                 "You must define one optimizer for each model."
-#                 f"Got {len_model} models, and {len_optimizer}"
-#                 " optimizers."
-#             )
-
-#         # check length consistency optimizers kwargs
-#         if len_optimizer_kwargs != len_optimizer:
-#             raise ValueError(
-#                 "You must define one dictionary of keyword"
-#                 " arguments for each optimizers."
-#                 f"Got {len_optimizer} optimizers, and"
-#                 f" {len_optimizer_kwargs} dicitionaries"
-#             )
-
-#         # extra features handling
-#         if (extra_features is None) or (len(extra_features) == 0):
-#             extra_features = [None] * len_model
-#         else:
-#             # if we only have a list of extra features
-#             if not isinstance(extra_features[0], (tuple, list)):
-#                 extra_features = [extra_features] * len_model
-#             else:  # if we have a list of list extra features
-#                 if len(extra_features) != len_model:
-#                     raise ValueError(
-#                         "You passed a list of extrafeatures list with len"
-#                         f"different of models len. Expected {len_model} "
-#                         f"got {len(extra_features)}. If you want to use "
-#                         "the same list of extra features for all models, "
-#                         "just pass a list of extrafeatures and not a list "
-#                         "of list of extra features."
-#                     )
-
-#         # assigning model and optimizers
-#         self._pina_models = []
-#         self._pina_optimizers = []
-
-#         for idx in range(len_model):
-#             model_ = Network(
-#                 model=models[idx],
-#                 input_variables=problem.input_variables,
-#                 output_variables=problem.output_variables,
-#                 extra_features=extra_features[idx],
-#             )
-#             optim_ = optimizers[idx](
-#                 model_.parameters(), **optimizers_kwargs[idx]
-#             )
-#             self._pina_models.append(model_)
-#             self._pina_optimizers.append(optim_)
-
-#         # assigning problem
-#         self._pina_problem = problem
-
-#     @abstractmethod
-#     def forward(self, *args, **kwargs):
-#         pass
-
-#     @abstractmethod
-#     def training_step(self):
-#         pass
-
-#     @abstractmethod
-#     def configure_optimizers(self):
-#         pass
-
-#     @property
-#     def models(self):
-#         """
-#         The torch model."""
-#         return self._pina_models
-
-#     @property
-#     def optimizers(self):
-#         """
-#         The torch model."""
-#         return self._pina_optimizers
-
-#     @property
-#     def problem(self):
-#         """
-#         The problem formulation."""
-#         return self._pina_problem
-
-#     def on_train_start(self):
-#         """
-#         On training epoch start this function is call to do global checks for
-#         the different solvers.
-#         """
-
-#         # 1. Check the verison for dataloader
-#         dataloader = self.trainer.train_dataloader
-#         if sys.version_info < (3, 8):
-#             dataloader = dataloader.loaders
-#         self._dataloader = dataloader
-
-#         return super().on_train_start()
-
-    # @model.setter
-    # def model(self, new_model):
-    #     """
-    #     Set the torch."""
-    #     check_consistency(new_model, nn.Module, 'torch model')
-    #     self._model= new_model
-
-    # @problem.setter
-    # def problem(self, problem):
-    #     """
-    #     Set the problem formulation."""
-    #     check_consistency(problem, AbstractProblem, 'pina problem')
-    #     self._problem = problem
-
 class SolverInterface(pytorch_lightning.LightningModule, metaclass=ABCMeta):
     """
     Solver base class. This class inherits is a wrapper of
@@ -181,10 +19,12 @@ class SolverInterface(pytorch_lightning.LightningModule, metaclass=ABCMeta):
 
     def __init__(
         self,
-        model,
+        models,
         problem,
-        optimizer,
-        scheduler,
+        optimizers,
+        schedulers,
+        extra_features,
+        use_lt=True
     ):
         """
         :param model: A torch neural network model instance.
@@ -197,22 +37,45 @@ def __init__(
         super().__init__()
 
         # check consistency of the inputs
-        check_consistency(model, torch.nn.Module)
         check_consistency(problem, AbstractProblem)
-        check_consistency(optimizer, Optimizer)
-        check_consistency(scheduler, Scheduler)
-
-        # put everything in a list if only one input
-        if not isinstance(model, list):
-            model = [model]
-        if not isinstance(scheduler, list):
-            scheduler = [scheduler]
-        if not isinstance(optimizer, list):
-            optimizer = [optimizer]
-
-        # number of models and optimizers
-        len_model = len(model)
-        len_optimizer = len(optimizer)
+        self._check_solver_consistency(problem)
+
+        #Check consistency of models argument and encapsulate in list
+        if not isinstance(models, list):
+            check_consistency(models, torch.nn.Module)
+            # put everything in a list if only one input
+            models = [models]
+        else:
+            for idx in range(len(models)):
+                # Check consistency
+                check_consistency(models[idx], torch.nn.Module)
+        len_model = len(models)
+
+        #If use_lt is true add extract operation in input
+        if use_lt is True:
+            for idx in range(len(models)):
+                models[idx] = Network(
+                    model = models[idx],
+                    input_variables=problem.input_variables,
+                    output_variables=problem.output_variables,
+                    extra_features=extra_features, )
+
+        #Check scheduler consistency + encapsulation
+        if not isinstance(schedulers, list):
+            check_consistency(schedulers, Scheduler)
+            schedulers = [schedulers]
+        else:
+            for scheduler in schedulers:
+                check_consistency(scheduler, Scheduler)
+
+        #Check optimizer consistency + encapsulation
+        if not isinstance(optimizers, list):
+            check_consistency(optimizers, Optimizer)
+            optimizers = [optimizers]
+        else:
+            for optimizer in optimizers:
+                check_consistency(optimizer, Optimizer)
+        len_optimizer = len(optimizers)
 
         # check length consistency optimizers
         if len_model != len_optimizer:
@@ -223,10 +86,12 @@ def __init__(
             )
 
         # extra features handling
+
+        self._pina_models = models
+        self._pina_optimizers = optimizers
+        self._pina_schedulers = schedulers
         self._pina_problem = problem
-        self._pina_model = model
-        self._pina_optimizer = optimizer
-        self._pina_scheduler = scheduler
+
 
     @abstractmethod
     def forward(self, *args, **kwargs):
@@ -244,13 +109,13 @@ def configure_optimizers(self):
     def models(self):
         """
         The torch model."""
-        return self._pina_model
+        return self._pina_models
 
     @property
     def optimizers(self):
         """
         The torch model."""
-        return self._pina_optimizer
+        return self._pina_optimizers
 
     @property
     def problem(self):
@@ -272,16 +137,10 @@ def on_train_start(self):
 
         return super().on_train_start()
 
-    # @model.setter
-    # def model(self, new_model):
-    #     """
-    #     Set the torch."""
-    #     check_consistency(new_model, nn.Module, 'torch model')
-    #     self._model= new_model
-
-    # @problem.setter
-    # def problem(self, problem):
-    #     """
-    #     Set the problem formulation."""
-    #     check_consistency(problem, AbstractProblem, 'pina problem')
-    #     self._problem = problem
+    def _check_solver_consistency(self, problem):
+        """
+        TODO
+        """
+        for _, condition in problem.conditions.items():
+            if not set(self.accepted_condition_types).issubset(condition.condition_type):
+                raise ValueError(f'{self.__name__} support only dose not support condition {condition.condition_type}')
diff --git a/pina/solvers/supervised.py b/pina/solvers/supervised.py
index c44d5a1e2..32f687ed0 100644
--- a/pina/solvers/supervised.py
+++ b/pina/solvers/supervised.py
@@ -2,9 +2,7 @@
 
 import torch
 from torch.nn.modules.loss import _Loss
-
-
-from ..optim import Optimizer, Scheduler, TorchOptimizer, TorchScheduler
+from ..optim import TorchOptimizer, TorchScheduler
 from .solver import SolverInterface
 from ..label_tensor import LabelTensor
 from ..utils import check_consistency
@@ -39,14 +37,17 @@ class SupervisedSolver(SolverInterface):
     we are seeking to approximate multiple (discretised) functions given
     multiple (discretised) input functions.
     """
+    accepted_condition_types = ['supervised']
+    __name__ = 'SupervisedSolver'
 
     def __init__(
-        self,
-        problem,
-        model,
-        loss=None,
-        optimizer=None,
-        scheduler=None,
+            self,
+            problem,
+            model,
+            loss=None,
+            optimizer=None,
+            scheduler=None,
+            extra_features=None
     ):
         """
         :param AbstractProblem problem: The formualation of the problem.
@@ -57,11 +58,8 @@ def __init__(
             features to use as augmented input.
         :param torch.optim.Optimizer optimizer: The neural network optimizer to
             use; default is :class:`torch.optim.Adam`.
-        :param dict optimizer_kwargs: Optimizer constructor keyword args.
-        :param float lr: The learning rate; default is 0.001.
         :param torch.optim.LRScheduler scheduler: Learning
             rate scheduler.
-        :param dict scheduler_kwargs: LR scheduler constructor keyword args.
         """
         if loss is None:
             loss = torch.nn.MSELoss()
@@ -74,18 +72,19 @@ def __init__(
                 torch.optim.lr_scheduler.ConstantLR)
 
         super().__init__(
-            model=model,
+            models=model,
             problem=problem,
-            optimizer=optimizer,
-            scheduler=scheduler,
+            optimizers=optimizer,
+            schedulers=scheduler,
+            extra_features=extra_features
         )
 
         # check consistency
         check_consistency(loss, (LossInterface, _Loss), subclass=False)
         self._loss = loss
-        self._model = self._pina_model[0]
-        self._optimizer = self._pina_optimizer[0]
-        self._scheduler = self._pina_scheduler[0]
+        self._model = self._pina_models[0]
+        self._optimizer = self._pina_optimizers[0]
+        self._scheduler = self._pina_schedulers[0]
 
     def forward(self, x):
         """Forward pass implementation for the solver.
@@ -97,12 +96,7 @@ def forward(self, x):
 
         output = self._model(x)
 
-        output.labels = {
-            1: {
-                "name": "output",
-                "dof": self.problem.output_variables
-            }
-        }
+        output.labels = self.problem.output_variables
         return output
 
     def configure_optimizers(self):
@@ -128,16 +122,14 @@ def training_step(self, batch, batch_idx):
         :return: The sum of the loss functions.
         :rtype: LabelTensor
         """
-
-        condition_idx = batch.condition
+        condition_idx = batch.supervised.condition_indices
 
         for condition_id in range(condition_idx.min(), condition_idx.max() + 1):
 
             condition_name = self._dataloader.condition_names[condition_id]
             condition = self.problem.conditions[condition_name]
-            pts = batch.input
-            out = batch.output
-
+            pts = batch.supervised.input_points
+            out = batch.supervised.output_points
             if condition_name not in self.problem.conditions:
                 raise RuntimeError("Something wrong happened.")
 
@@ -167,8 +159,8 @@ def loss_data(self, input_pts, output_pts):
         the network output against the true solution. This function
         should not be override if not intentionally.
 
-        :param LabelTensor input_tensor: The input to the neural networks.
-        :param LabelTensor output_tensor: The true solution to compare the
+        :param LabelTensor input_pts: The input to the neural networks.
+        :param LabelTensor output_pts: The true solution to compare the
             network solution.
         :return: The residual loss averaged on the input coordinates
         :rtype: torch.Tensor
@@ -181,7 +173,7 @@ def scheduler(self):
         Scheduler for training.
         """
         return self._scheduler
-    
+
     @property
     def optimizer(self):
         """
diff --git a/pina/trainer.py b/pina/trainer.py
index ba18f3392..49c6a4017 100644
--- a/pina/trainer.py
+++ b/pina/trainer.py
@@ -3,13 +3,13 @@
 import torch
 import pytorch_lightning
 from .utils import check_consistency
-from .data import SamplePointDataset, SamplePointLoader, DataPointDataset
+from .data import PinaDataModule
 from .solvers.solver import SolverInterface
 
 
 class Trainer(pytorch_lightning.Trainer):
 
-    def __init__(self, solver, batch_size=None, **kwargs):
+    def __init__(self, solver, batch_size=None, train_size=.7, test_size=.2, eval_size=.1, **kwargs):
         """
         PINA Trainer class for costumizing every aspect of training via flags.
 
@@ -31,10 +31,11 @@ def __init__(self, solver, batch_size=None, **kwargs):
         check_consistency(solver, SolverInterface)
         if batch_size is not None:
             check_consistency(batch_size, int)
-
+        self.train_size = train_size
+        self.test_size = test_size
+        self.eval_size = eval_size
         self.solver = solver
         self.batch_size = batch_size
-
         self._create_loader()
         self._move_to_device()
 
@@ -69,11 +70,12 @@ def _create_loader(self):
             raise RuntimeError("Parallel training is not supported yet.")
 
         device = devices[0]
-        dataset_phys = SamplePointDataset(self.solver.problem, device)
-        dataset_data = DataPointDataset(self.solver.problem, device)
-        self._loader = SamplePointLoader(
-            dataset_phys, dataset_data, batch_size=self.batch_size, shuffle=True
-        )
+
+        data_module = PinaDataModule(problem=self.solver.problem, device=device,
+                                     train_size=self.train_size, test_size=self.test_size,
+                                     eval_size=self.eval_size)
+        data_module.setup()
+        self._loader = data_module.train_dataloader()
 
     def train(self, **kwargs):
         """
@@ -89,3 +91,7 @@ def solver(self):
         Returning trainer solver.
         """
         return self._solver
+
+    @solver.setter
+    def solver(self, solver):
+        self._solver = solver
diff --git a/tests/test_dataset.py b/tests/test_dataset.py
index 40f219228..264f794bf 100644
--- a/tests/test_dataset.py
+++ b/tests/test_dataset.py
@@ -1,44 +1,45 @@
+import math
 import torch
-import pytest
-
-from pina.data.dataset import SamplePointDataset, SamplePointLoader, DataPointDataset
+from pina.data import SamplePointDataset, SupervisedDataset, PinaDataModule, UnsupervisedDataset, unsupervised_dataset
+from pina.data import PinaDataLoader
 from pina import LabelTensor, Condition
 from pina.equation import Equation
 from pina.domain import CartesianDomain
 from pina.problem import SpatialProblem
-from pina.model import FeedForward
 from pina.operators import laplacian
 from pina.equation.equation_factory import FixedValue
 
 
 def laplace_equation(input_, output_):
-    force_term = (torch.sin(input_.extract(['x'])*torch.pi) *
-                    torch.sin(input_.extract(['y'])*torch.pi))
+    force_term = (torch.sin(input_.extract(['x']) * torch.pi) *
+                  torch.sin(input_.extract(['y']) * torch.pi))
     delta_u = laplacian(output_.extract(['u']), input_)
     return delta_u - force_term
 
+
 my_laplace = Equation(laplace_equation)
 in_ = LabelTensor(torch.tensor([[0., 1.]]), ['x', 'y'])
 out_ = LabelTensor(torch.tensor([[0.]]), ['u'])
 in2_ = LabelTensor(torch.rand(60, 2), ['x', 'y'])
 out2_ = LabelTensor(torch.rand(60, 1), ['u'])
 
+
 class Poisson(SpatialProblem):
     output_variables = ['u']
     spatial_domain = CartesianDomain({'x': [0, 1], 'y': [0, 1]})
 
     conditions = {
         'gamma1': Condition(
-            location=CartesianDomain({'x': [0, 1], 'y':  1}),
+            domain=CartesianDomain({'x': [0, 1], 'y': 1}),
             equation=FixedValue(0.0)),
         'gamma2': Condition(
-            location=CartesianDomain({'x': [0, 1], 'y': 0}),
+            domain=CartesianDomain({'x': [0, 1], 'y': 0}),
             equation=FixedValue(0.0)),
         'gamma3': Condition(
-            location=CartesianDomain({'x':  1, 'y': [0, 1]}),
+            domain=CartesianDomain({'x': 1, 'y': [0, 1]}),
             equation=FixedValue(0.0)),
         'gamma4': Condition(
-            location=CartesianDomain({'x': 0, 'y': [0, 1]}),
+            domain=CartesianDomain({'x': 0, 'y': [0, 1]}),
             equation=FixedValue(0.0)),
         'D': Condition(
             input_points=LabelTensor(torch.rand(size=(100, 2)), ['x', 'y']),
@@ -48,75 +49,114 @@ class Poisson(SpatialProblem):
             output_points=out_),
         'data2': Condition(
             input_points=in2_,
-            output_points=out2_)
+            output_points=out2_),
+        'unsupervised': Condition(
+            input_points=LabelTensor(torch.rand(size=(45, 2)), ['x', 'y']),
+            conditional_variables=LabelTensor(torch.ones(size=(45, 1)), ['alpha']),
+        ),
+        'unsupervised2': Condition(
+            input_points=LabelTensor(torch.rand(size=(90, 2)), ['x', 'y']),
+            conditional_variables=LabelTensor(torch.ones(size=(90, 1)), ['alpha']),
+        )
     }
 
+
 boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4']
 poisson = Poisson()
 poisson.discretise_domain(10, 'grid', locations=boundaries)
 
+
 def test_sample():
     sample_dataset = SamplePointDataset(poisson, device='cpu')
     assert len(sample_dataset) == 140
-    assert sample_dataset.pts.shape == (140, 2)
-    assert sample_dataset.pts.labels == ['x', 'y']
-    assert sample_dataset.condition_indeces.dtype == torch.int64
-    assert sample_dataset.condition_indeces.max() == torch.tensor(4)
-    assert sample_dataset.condition_indeces.min() == torch.tensor(0)
+    assert sample_dataset.input_points.shape == (140, 2)
+    assert sample_dataset.input_points.labels == ['x', 'y']
+    assert sample_dataset.condition_indices.dtype == torch.uint8
+    assert sample_dataset.condition_indices.max() == torch.tensor(4)
+    assert sample_dataset.condition_indices.min() == torch.tensor(0)
+
 
 def test_data():
-    dataset = DataPointDataset(poisson, device='cpu')
+    dataset = SupervisedDataset(poisson, device='cpu')
     assert len(dataset) == 61
-    assert dataset.input_pts.shape == (61, 2)
-    assert dataset.input_pts.labels == ['x', 'y']
-    assert dataset.output_pts.shape == (61, 1 )
-    assert dataset.output_pts.labels == ['u']
-    assert dataset.condition_indeces.dtype == torch.int64
-    assert dataset.condition_indeces.max() == torch.tensor(1)
-    assert dataset.condition_indeces.min() == torch.tensor(0)
-
-def test_loader():
-    sample_dataset = SamplePointDataset(poisson, device='cpu')
-    data_dataset = DataPointDataset(poisson, device='cpu')
-    loader = SamplePointLoader(sample_dataset, data_dataset, batch_size=10)
-
+    assert dataset['input_points'].shape == (61, 2)
+    assert dataset.input_points.shape == (61, 2)
+    assert dataset['input_points'].labels == ['x', 'y']
+    assert dataset.input_points.labels == ['x', 'y']
+    assert dataset['input_points', 3:].shape == (58, 2)
+    assert dataset[3:][1].labels == ['u']
+    assert dataset.output_points.shape == (61, 1)
+    assert dataset.output_points.labels == ['u']
+    assert dataset.condition_indices.dtype == torch.uint8
+    assert dataset.condition_indices.max() == torch.tensor(1)
+    assert dataset.condition_indices.min() == torch.tensor(0)
+
+
+def test_unsupervised():
+    dataset = UnsupervisedDataset(poisson, device='cpu')
+    assert len(dataset) == 135
+    assert dataset.input_points.shape == (135, 2)
+    assert dataset.input_points.labels == ['x', 'y']
+    assert dataset.input_points[3:].shape == (132, 2)
+
+    assert dataset.conditional_variables.shape == (135, 1)
+    assert dataset.conditional_variables.labels == ['alpha']
+    assert dataset.condition_indices.dtype == torch.uint8
+    assert dataset.condition_indices.max() == torch.tensor(1)
+    assert dataset.condition_indices.min() == torch.tensor(0)
+
+
+def test_data_module():
+    data_module = PinaDataModule(poisson, device='cpu')
+    data_module.setup()
+    loader = data_module.train_dataloader()
+    assert isinstance(loader, PinaDataLoader)
+    assert isinstance(loader, PinaDataLoader)
+
+    data_module = PinaDataModule(poisson, device='cpu', batch_size=10, shuffle=False)
+    data_module.setup()
+    loader = data_module.train_dataloader()
+    assert len(loader) == 24
+    for i in loader:
+        assert len(i) <= 10
+    len_ref = sum([math.ceil(len(dataset) * 0.7) for dataset in data_module.datasets])
+    len_real = sum([len(dataset) for dataset in data_module.splits['train'].values()])
+    assert len_ref == len_real
+
+    supervised_dataset = SupervisedDataset(poisson, device='cpu')
+    data_module = PinaDataModule(poisson, device='cpu', batch_size=10, shuffle=False, datasets=[supervised_dataset])
+    data_module.setup()
+    loader = data_module.train_dataloader()
     for batch in loader:
-        assert len(batch) in [2, 3]
-        assert batch['pts'].shape[0] <= 10
-        assert batch['pts'].requires_grad == True
-        assert batch['pts'].labels == ['x', 'y']
-
-    loader2 = SamplePointLoader(sample_dataset, data_dataset, batch_size=None)
-    assert len(list(loader2)) == 2
-
-def test_loader2():
-    poisson2 = Poisson()
-    del poisson.conditions['data2']
-    del poisson2.conditions['data']
-    poisson2.discretise_domain(10, 'grid', locations=boundaries)
-    sample_dataset = SamplePointDataset(poisson, device='cpu')
-    data_dataset = DataPointDataset(poisson, device='cpu')
-    loader = SamplePointLoader(sample_dataset, data_dataset, batch_size=10)
+        assert len(batch) <= 10
 
+    physics_dataset = SamplePointDataset(poisson, device='cpu')
+    data_module = PinaDataModule(poisson, device='cpu', batch_size=10, shuffle=False, datasets=[physics_dataset])
+    data_module.setup()
+    loader = data_module.train_dataloader()
     for batch in loader:
-        assert len(batch) == 2 # only phys condtions
-        assert batch['pts'].shape[0] <= 10
-        assert batch['pts'].requires_grad == True
-        assert batch['pts'].labels == ['x', 'y']
-
-def test_loader3():
-    poisson2 = Poisson()
-    del poisson.conditions['gamma1']
-    del poisson.conditions['gamma2']
-    del poisson.conditions['gamma3']
-    del poisson.conditions['gamma4']
-    del poisson.conditions['D']
-    sample_dataset = SamplePointDataset(poisson, device='cpu')
-    data_dataset = DataPointDataset(poisson, device='cpu')
-    loader = SamplePointLoader(sample_dataset, data_dataset, batch_size=10)
+        assert len(batch) <= 10
 
+    unsupervised_dataset = UnsupervisedDataset(poisson, device='cpu')
+    data_module = PinaDataModule(poisson, device='cpu', batch_size=10, shuffle=False, datasets=[unsupervised_dataset])
+    data_module.setup()
+    loader = data_module.train_dataloader()
     for batch in loader:
-        assert len(batch) == 2 # only phys condtions
-        assert batch['pts'].shape[0] <= 10
-        assert batch['pts'].requires_grad == True
-        assert batch['pts'].labels == ['x', 'y']
+        assert len(batch) <= 10
+
+
+def test_loader():
+    data_module = PinaDataModule(poisson, device='cpu', batch_size=10)
+    data_module.setup()
+    loader = data_module.train_dataloader()
+    assert isinstance(loader, PinaDataLoader)
+    assert len(loader) == 24
+    for i in loader:
+        assert len(i) <= 10
+        assert i.supervised.input_points.labels == ['x', 'y']
+        assert i.physics.input_points.labels == ['x', 'y']
+        assert i.unsupervised.input_points.labels == ['x', 'y']
+        assert i.supervised.input_points.requires_grad == True
+        assert i.physics.input_points.requires_grad == True
+        assert i.unsupervised.input_points.requires_grad == True
+test_loader()
\ No newline at end of file
diff --git a/tests/test_solvers/test_supervised_solver.py b/tests/test_solvers/test_supervised_solver.py
index 912480bb8..8ceadcd93 100644
--- a/tests/test_solvers/test_supervised_solver.py
+++ b/tests/test_solvers/test_supervised_solver.py
@@ -1,51 +1,28 @@
 import torch
-
-from pina.problem import AbstractProblem
+import pytest
+from pina.problem import AbstractProblem, SpatialProblem
 from pina import Condition, LabelTensor
 from pina.solvers import SupervisedSolver
-from pina.trainer import Trainer
 from pina.model import FeedForward
-from pina.loss import LpLoss
-from pina.solvers import GraphSupervisedSolver
+from pina.equation.equation import Equation
+from pina.equation.equation_factory import FixedValue
+from pina.operators import laplacian
+from pina.domain import CartesianDomain
+from pina.trainer import Trainer
+
+in_ = LabelTensor(torch.tensor([[0., 1.]]), ['u_0', 'u_1'])
+out_ = LabelTensor(torch.tensor([[0.]]), ['u'])
+
 
 class NeuralOperatorProblem(AbstractProblem):
     input_variables = ['u_0', 'u_1']
     output_variables = ['u']
-    domains = {
-        'pts': LabelTensor(
-            torch.rand(100, 2),
-            labels={1: {'name': 'space', 'dof': ['u_0', 'u_1']}}
-        )
-    }
-    conditions = {
-        'data' : Condition(
-            domain='pts', 
-            output_points=LabelTensor(
-                torch.rand(100, 1), 
-                labels={1: {'name': 'output', 'dof': ['u']}}
-            )
-        )
-    }
 
-class NeuralOperatorProblemGraph(AbstractProblem):
-    input_variables = ['x', 'y', 'u_0', 'u_1']
-    output_variables = ['u']
-    domains = {
-        'pts': LabelTensor(
-            torch.rand(100, 4),
-            labels={1: {'name': 'space', 'dof': ['x', 'y', 'u_0', 'u_1']}}
-        )
-    }
     conditions = {
-        'data' : Condition(
-            domain='pts',
-            output_points=LabelTensor(
-                torch.rand(100, 1),
-                labels={1: {'name': 'output', 'dof': ['u']}}
-            )
-        )
+        'data': Condition(input_points=in_, output_points=out_),
     }
 
+
 class myFeature(torch.nn.Module):
     """
     Feature: sin(x)
@@ -61,117 +38,106 @@ def forward(self, x):
 
 
 problem = NeuralOperatorProblem()
-problem_graph = NeuralOperatorProblemGraph()
-# make the problem + extra feats
 extra_feats = [myFeature()]
-model = FeedForward(len(problem.input_variables),
-                    len(problem.output_variables))
+model = FeedForward(len(problem.input_variables), len(problem.output_variables))
 model_extra_feats = FeedForward(
-    len(problem.input_variables) + 1,
-    len(problem.output_variables))
+    len(problem.input_variables) + 1, len(problem.output_variables))
 
 
 def test_constructor():
     SupervisedSolver(problem=problem, model=model)
 
 
-# def test_constructor_extra_feats():
-#     SupervisedSolver(problem=problem, model=model_extra_feats, extra_features=extra_feats)
-
-'''
-class AutoSolver(SupervisedSolver):
-
-    def forward(self, input):
-        from pina.graph import Graph
-        print(Graph)
-        print(input)
-        if not isinstance(input, Graph):
-            input = Graph.build('radius', nodes_coordinates=input, nodes_data=torch.rand(input.shape), radius=0.2)
-        print(input)
-        print(input.data.edge_index)
-        print(input.data)
-        g = self._model(input.data, edge_index=input.data.edge_index)
-        g.labels = {1: {'name': 'output', 'dof': ['u']}}
-        return g
-        du_dt_new = LabelTensor(self.model(graph).reshape(-1,1), labels = ['du'])
-
-        return du_dt_new
-'''
-
-class GraphModel(torch.nn.Module):
-    def __init__(self, in_channels, out_channels):
-        from torch_geometric.nn import GCNConv, NNConv
-        super().__init__()
-        self.conv1 = GCNConv(in_channels, 16)
-        self.conv2 = GCNConv(16, out_channels)
-
-    def forward(self, data, edge_index):
-        print(data)
-        x = data.x
-        print(x)
-        x = self.conv1(x, edge_index)
-        x = x.relu()
-        x = self.conv2(x, edge_index)
-        return x
-
-def test_graph():
-    solver = GraphSupervisedSolver(problem=problem_graph, model=GraphModel(2, 1), loss=LpLoss(),
-                                   nodes_coordinates=['x', 'y'], nodes_data=['u_0', 'u_1'])
-    trainer = Trainer(solver=solver, max_epochs=30, accelerator='cpu', batch_size=20)
-    trainer.train()
+test_constructor()
+
+
+def laplace_equation(input_, output_):
+    force_term = (torch.sin(input_.extract(['x']) * torch.pi) *
+                  torch.sin(input_.extract(['y']) * torch.pi))
+    delta_u = laplacian(output_.extract(['u']), input_)
+    return delta_u - force_term
+
+
+my_laplace = Equation(laplace_equation)
+
+
+class Poisson(SpatialProblem):
+    output_variables = ['u']
+    spatial_domain = CartesianDomain({'x': [0, 1], 'y': [0, 1]})
+
+    conditions = {
+        'gamma1':
+            Condition(domain=CartesianDomain({
+                'x': [0, 1],
+                'y': 1
+            }),
+                equation=FixedValue(0.0)),
+        'gamma2':
+            Condition(domain=CartesianDomain({
+                'x': [0, 1],
+                'y': 0
+            }),
+                equation=FixedValue(0.0)),
+        'gamma3':
+            Condition(domain=CartesianDomain({
+                'x': 1,
+                'y': [0, 1]
+            }),
+                equation=FixedValue(0.0)),
+        'gamma4':
+            Condition(domain=CartesianDomain({
+                'x': 0,
+                'y': [0, 1]
+            }),
+                equation=FixedValue(0.0)),
+        'D':
+            Condition(domain=CartesianDomain({
+                'x': [0, 1],
+                'y': [0, 1]
+            }),
+                equation=my_laplace),
+        'data':
+            Condition(input_points=in_, output_points=out_)
+    }
+
+    def poisson_sol(self, pts):
+        return -(torch.sin(pts.extract(['x']) * torch.pi) *
+                 torch.sin(pts.extract(['y']) * torch.pi)) / (2 * torch.pi ** 2)
+
+    truth_solution = poisson_sol
+
+
+def test_wrong_constructor():
+    poisson_problem = Poisson()
+    with pytest.raises(ValueError):
+        SupervisedSolver(problem=poisson_problem, model=model)
 
 
 def test_train_cpu():
-    solver = SupervisedSolver(problem = problem, model=model, loss=LpLoss())
-    trainer = Trainer(solver=solver, max_epochs=300, accelerator='cpu', batch_size=20)
+    solver = SupervisedSolver(problem=problem, model=model)
+    trainer = Trainer(solver=solver,
+                      max_epochs=200,
+                      accelerator='gpu',
+                      batch_size=5,
+                      train_size=1,
+                      test_size=0.,
+                      eval_size=0.)
     trainer.train()
+test_train_cpu()
 
 
-# def test_train_restore():
-#     tmpdir = "tests/tmp_restore"
-#     solver = SupervisedSolver(problem=problem,
-#                 model=model,
-#                 extra_features=None,
-#                 loss=LpLoss())
-#     trainer = Trainer(solver=solver,
-#                       max_epochs=5,
-#                       accelerator='cpu',
-#                       default_root_dir=tmpdir)
-#     trainer.train()
-#     ntrainer = Trainer(solver=solver, max_epochs=15, accelerator='cpu')
-#     t = ntrainer.train(
-#         ckpt_path=f'{tmpdir}/lightning_logs/version_0/checkpoints/epoch=4-step=5.ckpt')
-#     import shutil
-#     shutil.rmtree(tmpdir)
-
-
-# def test_train_load():
-#     tmpdir = "tests/tmp_load"
-#     solver = SupervisedSolver(problem=problem,
-#                 model=model,
-#                 extra_features=None,
-#                 loss=LpLoss())
-#     trainer = Trainer(solver=solver,
-#                       max_epochs=15,
-#                       accelerator='cpu',
-#                       default_root_dir=tmpdir)
-#     trainer.train()
-#     new_solver = SupervisedSolver.load_from_checkpoint(
-#         f'{tmpdir}/lightning_logs/version_0/checkpoints/epoch=14-step=15.ckpt',
-#         problem = problem, model=model)
-#     test_pts = LabelTensor(torch.rand(20, 2), problem.input_variables)
-#     assert new_solver.forward(test_pts).shape == (20, 1)
-#     assert new_solver.forward(test_pts).shape == solver.forward(test_pts).shape
-#     torch.testing.assert_close(
-#         new_solver.forward(test_pts),
-#         solver.forward(test_pts))
-#     import shutil
-#     shutil.rmtree(tmpdir)
-
-# def test_train_extra_feats_cpu():
-#     pinn = SupervisedSolver(problem=problem,
-#                 model=model_extra_feats,
-#                 extra_features=extra_feats)
-#     trainer = Trainer(solver=pinn, max_epochs=5, accelerator='cpu')
-#     trainer.train()
-test_graph()
\ No newline at end of file
+def test_extra_features_constructor():
+    SupervisedSolver(problem=problem,
+                     model=model_extra_feats,
+                     extra_features=extra_feats)
+
+
+def test_extra_features_train_cpu():
+    solver = SupervisedSolver(problem=problem,
+                              model=model_extra_feats,
+                              extra_features=extra_feats)
+    trainer = Trainer(solver=solver,
+                      max_epochs=200,
+                      accelerator='gpu',
+                      batch_size=5)
+    trainer.train()

From bee3c15385135c110fd74f53559ced2a9c3c82b7 Mon Sep 17 00:00:00 2001
From: FilippoOlivo <filippo@filippoolivo.com>
Date: Tue, 22 Oct 2024 14:26:39 +0200
Subject: [PATCH 25/55] Correct codacy warnings

---
 pina/__init__.py                              |  12 +-
 pina/data/__init__.py                         |   9 +-
 pina/data/base_dataset.py                     |  14 +-
 pina/data/pina_batch.py                       |   6 +-
 pina/data/pina_dataloader.py                  |   3 +-
 pina/optim/torch_optimizer.py                 |   6 +-
 pina/optim/torch_scheduler.py                 |   7 +-
 pina/solvers/solver.py                        |  37 +++--
 pina/solvers/supervised.py                    |  44 +++---
 tests/test_dataset.py                         |  98 ++++++++-----
 tests/test_label_tensor/test_label_tensor.py  | 134 ++++++++++--------
 .../test_label_tensor/test_label_tensor_01.py |   9 +-
 tests/test_operators.py                       |   4 +
 tests/test_optimizer.py                       |  10 +-
 tests/test_problem.py                         |  69 ++++-----
 15 files changed, 252 insertions(+), 210 deletions(-)

diff --git a/pina/__init__.py b/pina/__init__.py
index d110d2842..30f35a6a5 100644
--- a/pina/__init__.py
+++ b/pina/__init__.py
@@ -1,12 +1,6 @@
 __all__ = [
-    "PINN",
-    "Trainer",
-    "LabelTensor",
-    "Plotter",
-    "Condition",
-    "SamplePointDataset",
-    "PinaDataModule",
-    "PinaDataLoader"
+    "PINN", "Trainer", "LabelTensor", "Plotter", "Condition",
+    "SamplePointDataset", "PinaDataModule", "PinaDataLoader"
 ]
 
 from .meta import *
@@ -17,4 +11,4 @@
 from .condition.condition import Condition
 from .data import SamplePointDataset
 from .data import PinaDataModule
-from .data import PinaDataLoader
\ No newline at end of file
+from .data import PinaDataLoader
diff --git a/pina/data/__init__.py b/pina/data/__init__.py
index 0a1b5905e..2b3a126a7 100644
--- a/pina/data/__init__.py
+++ b/pina/data/__init__.py
@@ -2,13 +2,8 @@
 Import data classes
 """
 __all__ = [
-    'PinaDataLoader',
-    'SupervisedDataset',
-    'SamplePointDataset',
-    'UnsupervisedDataset',
-    'Batch',
-    'PinaDataModule',
-    'BaseDataset'
+    'PinaDataLoader', 'SupervisedDataset', 'SamplePointDataset',
+    'UnsupervisedDataset', 'Batch', 'PinaDataModule', 'BaseDataset'
 ]
 
 from .pina_dataloader import PinaDataLoader
diff --git a/pina/data/base_dataset.py b/pina/data/base_dataset.py
index f095afa0c..f1d17ae31 100644
--- a/pina/data/base_dataset.py
+++ b/pina/data/base_dataset.py
@@ -22,10 +22,12 @@ def __new__(cls, problem, device):
         dataset will be loaded.
         """
         if cls is BaseDataset:
-            raise TypeError('BaseDataset cannot be instantiated directly. Use a subclass.')
+            raise TypeError(
+                'BaseDataset cannot be instantiated directly. Use a subclass.')
         if not hasattr(cls, '__slots__'):
-            raise TypeError('Something is wrong, __slots__ must be defined in subclasses.')
-        return super().__new__(cls)
+            raise TypeError(
+                'Something is wrong, __slots__ must be defined in subclasses.')
+        return super(BaseDataset, cls).__new__(cls)
 
     def __init__(self, problem, device):
         """"
@@ -79,7 +81,8 @@ def __len__(self):
 
     def __getattribute__(self, item):
         attribute = super().__getattribute__(item)
-        if isinstance(attribute, LabelTensor) and attribute.dtype == torch.float32:
+        if isinstance(attribute,
+                      LabelTensor) and attribute.dtype == torch.float32:
             attribute = attribute.to(device=self.device).requires_grad_()
         return attribute
 
@@ -101,7 +104,8 @@ def __getitem__(self, idx):
             if all(isinstance(x, int) for x in idx):
                 to_return_list = []
                 for i in self.__slots__:
-                    to_return_list.append(getattr(self, i)[[idx]].to(self.device))
+                    to_return_list.append(
+                        getattr(self, i)[[idx]].to(self.device))
                 return to_return_list
 
         raise ValueError(f'Invalid index {idx}')
diff --git a/pina/data/pina_batch.py b/pina/data/pina_batch.py
index 7e46a2218..f61e0020b 100644
--- a/pina/data/pina_batch.py
+++ b/pina/data/pina_batch.py
@@ -5,6 +5,7 @@
 
 
 class Batch:
+
     def __init__(self, dataset_dict, idx_dict):
 
         for k, v in dataset_dict.items():
@@ -29,5 +30,6 @@ def __len__(self):
     def __getattr__(self, item):
         if not item in dir(self):
             raise AttributeError(f'Batch instance has no attribute {item}')
-        return PinaSubset(getattr(self, item).dataset,
-                          getattr(self, item).indices[self.coordinates_dict[item]])
+        return PinaSubset(
+            getattr(self, item).dataset,
+            getattr(self, item).indices[self.coordinates_dict[item]])
diff --git a/pina/data/pina_dataloader.py b/pina/data/pina_dataloader.py
index d62847574..cbd8fe82d 100644
--- a/pina/data/pina_dataloader.py
+++ b/pina/data/pina_dataloader.py
@@ -50,7 +50,8 @@ def _init_batches(self, batch_size=None):
                     temp_dict[k] = slice(i * v, (i + 1) * v)
                 else:
                     temp_dict[k] = slice(i * v, len(self.dataset_dict[k]))
-            self.batches.append(Batch(idx_dict=temp_dict, dataset_dict=self.dataset_dict))
+            self.batches.append(
+                Batch(idx_dict=temp_dict, dataset_dict=self.dataset_dict))
 
     def __iter__(self):
         """
diff --git a/pina/optim/torch_optimizer.py b/pina/optim/torch_optimizer.py
index 239819a4f..ed90846c6 100644
--- a/pina/optim/torch_optimizer.py
+++ b/pina/optim/torch_optimizer.py
@@ -5,6 +5,7 @@
 from ..utils import check_consistency
 from .optimizer_interface import Optimizer
 
+
 class TorchOptimizer(Optimizer):
 
     def __init__(self, optimizer_class, **kwargs):
@@ -14,6 +15,5 @@ def __init__(self, optimizer_class, **kwargs):
         self.kwargs = kwargs
 
     def hook(self, parameters):
-        self.optimizer_instance = self.optimizer_class(
-            parameters, **self.kwargs
-        )
\ No newline at end of file
+        self.optimizer_instance = self.optimizer_class(parameters,
+                                                       **self.kwargs)
diff --git a/pina/optim/torch_scheduler.py b/pina/optim/torch_scheduler.py
index 50e1d91f7..9aa187d0c 100644
--- a/pina/optim/torch_scheduler.py
+++ b/pina/optim/torch_scheduler.py
@@ -5,13 +5,13 @@
     from torch.optim.lr_scheduler import LRScheduler  # torch >= 2.0
 except ImportError:
     from torch.optim.lr_scheduler import (
-        _LRScheduler as LRScheduler,
-    )  # torch < 2.0
+        _LRScheduler as LRScheduler, )  # torch < 2.0
 
 from ..utils import check_consistency
 from .optimizer_interface import Optimizer
 from .scheduler_interface import Scheduler
 
+
 class TorchScheduler(Scheduler):
 
     def __init__(self, scheduler_class, **kwargs):
@@ -23,5 +23,4 @@ def __init__(self, scheduler_class, **kwargs):
     def hook(self, optimizer):
         check_consistency(optimizer, Optimizer)
         self.scheduler_instance = self.scheduler_class(
-            optimizer.optimizer_instance, **self.kwargs
-        )
\ No newline at end of file
+            optimizer.optimizer_instance, **self.kwargs)
diff --git a/pina/solvers/solver.py b/pina/solvers/solver.py
index 8b3ddae71..1c6aa2b2b 100644
--- a/pina/solvers/solver.py
+++ b/pina/solvers/solver.py
@@ -17,15 +17,13 @@ class SolverInterface(pytorch_lightning.LightningModule, metaclass=ABCMeta):
     LightningModule methods.
     """
 
-    def __init__(
-        self,
-        models,
-        problem,
-        optimizers,
-        schedulers,
-        extra_features,
-        use_lt=True
-    ):
+    def __init__(self,
+                 models,
+                 problem,
+                 optimizers,
+                 schedulers,
+                 extra_features,
+                 use_lt=True):
         """
         :param model: A torch neural network model instance.
         :type model: torch.nn.Module
@@ -55,10 +53,11 @@ def __init__(
         if use_lt is True:
             for idx in range(len(models)):
                 models[idx] = Network(
-                    model = models[idx],
+                    model=models[idx],
                     input_variables=problem.input_variables,
                     output_variables=problem.output_variables,
-                    extra_features=extra_features, )
+                    extra_features=extra_features,
+                )
 
         #Check scheduler consistency + encapsulation
         if not isinstance(schedulers, list):
@@ -79,11 +78,9 @@ def __init__(
 
         # check length consistency optimizers
         if len_model != len_optimizer:
-            raise ValueError(
-                "You must define one optimizer for each model."
-                f"Got {len_model} models, and {len_optimizer}"
-                " optimizers."
-            )
+            raise ValueError("You must define one optimizer for each model."
+                             f"Got {len_model} models, and {len_optimizer}"
+                             " optimizers.")
 
         # extra features handling
 
@@ -92,7 +89,6 @@ def __init__(
         self._pina_schedulers = schedulers
         self._pina_problem = problem
 
-
     @abstractmethod
     def forward(self, *args, **kwargs):
         pass
@@ -142,5 +138,8 @@ def _check_solver_consistency(self, problem):
         TODO
         """
         for _, condition in problem.conditions.items():
-            if not set(self.accepted_condition_types).issubset(condition.condition_type):
-                raise ValueError(f'{self.__name__} support only dose not support condition {condition.condition_type}')
+            if not set(self.accepted_condition_types).issubset(
+                    condition.condition_type):
+                raise ValueError(
+                    f'{self.__name__} support only dose not support condition {condition.condition_type}'
+                )
diff --git a/pina/solvers/supervised.py b/pina/solvers/supervised.py
index 32f687ed0..a0b0f83ed 100644
--- a/pina/solvers/supervised.py
+++ b/pina/solvers/supervised.py
@@ -40,15 +40,13 @@ class SupervisedSolver(SolverInterface):
     accepted_condition_types = ['supervised']
     __name__ = 'SupervisedSolver'
 
-    def __init__(
-            self,
-            problem,
-            model,
-            loss=None,
-            optimizer=None,
-            scheduler=None,
-            extra_features=None
-    ):
+    def __init__(self,
+                 problem,
+                 model,
+                 loss=None,
+                 optimizer=None,
+                 scheduler=None,
+                 extra_features=None):
         """
         :param AbstractProblem problem: The formualation of the problem.
         :param torch.nn.Module model: The neural network model to use.
@@ -68,16 +66,13 @@ def __init__(
             optimizer = TorchOptimizer(torch.optim.Adam, lr=0.001)
 
         if scheduler is None:
-            scheduler = TorchScheduler(
-                torch.optim.lr_scheduler.ConstantLR)
+            scheduler = TorchScheduler(torch.optim.lr_scheduler.ConstantLR)
 
-        super().__init__(
-            models=model,
-            problem=problem,
-            optimizers=optimizer,
-            schedulers=scheduler,
-            extra_features=extra_features
-        )
+        super().__init__(models=model,
+                         problem=problem,
+                         optimizers=optimizer,
+                         schedulers=scheduler,
+                         extra_features=extra_features)
 
         # check consistency
         check_consistency(loss, (LossInterface, _Loss), subclass=False)
@@ -107,10 +102,8 @@ def configure_optimizers(self):
         """
         self._optimizer.hook(self._model.parameters())
         self._scheduler.hook(self._optimizer)
-        return (
-            [self._optimizer.optimizer_instance],
-            [self._scheduler.scheduler_instance]
-        )
+        return ([self._optimizer.optimizer_instance],
+                [self._scheduler.scheduler_instance])
 
     def training_step(self, batch, batch_idx):
         """Solver training step.
@@ -136,8 +129,7 @@ def training_step(self, batch, batch_idx):
             # for data driven mode
             if not hasattr(condition, "output_points"):
                 raise NotImplementedError(
-                    f"{type(self).__name__} works only in data-driven mode."
-                )
+                    f"{type(self).__name__} works only in data-driven mode.")
 
             output_pts = out[condition_idx == condition_id]
             input_pts = pts[condition_idx == condition_id]
@@ -145,9 +137,7 @@ def training_step(self, batch, batch_idx):
             input_pts.labels = pts.labels
             output_pts.labels = out.labels
 
-            loss = (
-                self.loss_data(input_pts=input_pts, output_pts=output_pts)
-            )
+            loss = (self.loss_data(input_pts=input_pts, output_pts=output_pts))
             loss = loss.as_subclass(torch.Tensor)
 
         self.log("mean_loss", float(loss), prog_bar=True, logger=True)
diff --git a/tests/test_dataset.py b/tests/test_dataset.py
index 264f794bf..653b0d6b6 100644
--- a/tests/test_dataset.py
+++ b/tests/test_dataset.py
@@ -29,34 +29,49 @@ class Poisson(SpatialProblem):
     spatial_domain = CartesianDomain({'x': [0, 1], 'y': [0, 1]})
 
     conditions = {
-        'gamma1': Condition(
-            domain=CartesianDomain({'x': [0, 1], 'y': 1}),
-            equation=FixedValue(0.0)),
-        'gamma2': Condition(
-            domain=CartesianDomain({'x': [0, 1], 'y': 0}),
-            equation=FixedValue(0.0)),
-        'gamma3': Condition(
-            domain=CartesianDomain({'x': 1, 'y': [0, 1]}),
-            equation=FixedValue(0.0)),
-        'gamma4': Condition(
-            domain=CartesianDomain({'x': 0, 'y': [0, 1]}),
-            equation=FixedValue(0.0)),
-        'D': Condition(
-            input_points=LabelTensor(torch.rand(size=(100, 2)), ['x', 'y']),
-            equation=my_laplace),
-        'data': Condition(
-            input_points=in_,
-            output_points=out_),
-        'data2': Condition(
-            input_points=in2_,
-            output_points=out2_),
-        'unsupervised': Condition(
+        'gamma1':
+        Condition(domain=CartesianDomain({
+            'x': [0, 1],
+            'y': 1
+        }),
+                  equation=FixedValue(0.0)),
+        'gamma2':
+        Condition(domain=CartesianDomain({
+            'x': [0, 1],
+            'y': 0
+        }),
+                  equation=FixedValue(0.0)),
+        'gamma3':
+        Condition(domain=CartesianDomain({
+            'x': 1,
+            'y': [0, 1]
+        }),
+                  equation=FixedValue(0.0)),
+        'gamma4':
+        Condition(domain=CartesianDomain({
+            'x': 0,
+            'y': [0, 1]
+        }),
+                  equation=FixedValue(0.0)),
+        'D':
+        Condition(input_points=LabelTensor(torch.rand(size=(100, 2)),
+                                           ['x', 'y']),
+                  equation=my_laplace),
+        'data':
+        Condition(input_points=in_, output_points=out_),
+        'data2':
+        Condition(input_points=in2_, output_points=out2_),
+        'unsupervised':
+        Condition(
             input_points=LabelTensor(torch.rand(size=(45, 2)), ['x', 'y']),
-            conditional_variables=LabelTensor(torch.ones(size=(45, 1)), ['alpha']),
+            conditional_variables=LabelTensor(torch.ones(size=(45, 1)),
+                                              ['alpha']),
         ),
-        'unsupervised2': Condition(
+        'unsupervised2':
+        Condition(
             input_points=LabelTensor(torch.rand(size=(90, 2)), ['x', 'y']),
-            conditional_variables=LabelTensor(torch.ones(size=(90, 1)), ['alpha']),
+            conditional_variables=LabelTensor(torch.ones(size=(90, 1)),
+                                              ['alpha']),
         )
     }
 
@@ -113,32 +128,49 @@ def test_data_module():
     assert isinstance(loader, PinaDataLoader)
     assert isinstance(loader, PinaDataLoader)
 
-    data_module = PinaDataModule(poisson, device='cpu', batch_size=10, shuffle=False)
+    data_module = PinaDataModule(poisson,
+                                 device='cpu',
+                                 batch_size=10,
+                                 shuffle=False)
     data_module.setup()
     loader = data_module.train_dataloader()
     assert len(loader) == 24
     for i in loader:
         assert len(i) <= 10
-    len_ref = sum([math.ceil(len(dataset) * 0.7) for dataset in data_module.datasets])
-    len_real = sum([len(dataset) for dataset in data_module.splits['train'].values()])
+    len_ref = sum(
+        [math.ceil(len(dataset) * 0.7) for dataset in data_module.datasets])
+    len_real = sum(
+        [len(dataset) for dataset in data_module.splits['train'].values()])
     assert len_ref == len_real
 
     supervised_dataset = SupervisedDataset(poisson, device='cpu')
-    data_module = PinaDataModule(poisson, device='cpu', batch_size=10, shuffle=False, datasets=[supervised_dataset])
+    data_module = PinaDataModule(poisson,
+                                 device='cpu',
+                                 batch_size=10,
+                                 shuffle=False,
+                                 datasets=[supervised_dataset])
     data_module.setup()
     loader = data_module.train_dataloader()
     for batch in loader:
         assert len(batch) <= 10
 
     physics_dataset = SamplePointDataset(poisson, device='cpu')
-    data_module = PinaDataModule(poisson, device='cpu', batch_size=10, shuffle=False, datasets=[physics_dataset])
+    data_module = PinaDataModule(poisson,
+                                 device='cpu',
+                                 batch_size=10,
+                                 shuffle=False,
+                                 datasets=[physics_dataset])
     data_module.setup()
     loader = data_module.train_dataloader()
     for batch in loader:
         assert len(batch) <= 10
 
     unsupervised_dataset = UnsupervisedDataset(poisson, device='cpu')
-    data_module = PinaDataModule(poisson, device='cpu', batch_size=10, shuffle=False, datasets=[unsupervised_dataset])
+    data_module = PinaDataModule(poisson,
+                                 device='cpu',
+                                 batch_size=10,
+                                 shuffle=False,
+                                 datasets=[unsupervised_dataset])
     data_module.setup()
     loader = data_module.train_dataloader()
     for batch in loader:
@@ -159,4 +191,6 @@ def test_loader():
         assert i.supervised.input_points.requires_grad == True
         assert i.physics.input_points.requires_grad == True
         assert i.unsupervised.input_points.requires_grad == True
-test_loader()
\ No newline at end of file
+
+
+test_loader()
diff --git a/tests/test_label_tensor/test_label_tensor.py b/tests/test_label_tensor/test_label_tensor.py
index 1165594db..846976730 100644
--- a/tests/test_label_tensor/test_label_tensor.py
+++ b/tests/test_label_tensor/test_label_tensor.py
@@ -4,29 +4,23 @@
 from pina.label_tensor import LabelTensor
 
 data = torch.rand((20, 3))
-labels_column = {
-    1: {
-        "name": "space",
-        "dof": ['x', 'y', 'z']
-    }
-}
-labels_row = {
-    0: {
-        "name": "samples",
-        "dof": range(20)
-    }
-}
+labels_column = {1: {"name": "space", "dof": ['x', 'y', 'z']}}
+labels_row = {0: {"name": "samples", "dof": range(20)}}
 labels_list = ['x', 'y', 'z']
 labels_all = labels_column | labels_row
 
-@pytest.mark.parametrize("labels", [labels_column, labels_row, labels_all, labels_list])
+
+@pytest.mark.parametrize("labels",
+                         [labels_column, labels_row, labels_all, labels_list])
 def test_constructor(labels):
     print(LabelTensor(data, labels))
 
+
 def test_wrong_constructor():
     with pytest.raises(ValueError):
         LabelTensor(data, ['a', 'b'])
 
+
 @pytest.mark.parametrize("labels", [labels_column, labels_all])
 @pytest.mark.parametrize("labels_te", ['z', ['z'], {'space': ['z']}])
 def test_extract_column(labels, labels_te):
@@ -37,6 +31,7 @@ def test_extract_column(labels, labels_te):
     assert new.shape[0] == 20
     assert torch.all(torch.isclose(data[:, 2].reshape(-1, 1), new))
 
+
 @pytest.mark.parametrize("labels", [labels_row, labels_all])
 @pytest.mark.parametrize("labels_te", [{'samples': [2]}])
 def test_extract_row(labels, labels_te):
@@ -47,10 +42,14 @@ def test_extract_row(labels, labels_te):
     assert new.shape[0] == 1
     assert torch.all(torch.isclose(data[2].reshape(1, -1), new))
 
-@pytest.mark.parametrize("labels_te", [
-    {'samples': [2], 'space': ['z']},
-    {'space': 'z', 'samples': 2}
-])
+
+@pytest.mark.parametrize("labels_te", [{
+    'samples': [2],
+    'space': ['z']
+}, {
+    'space': 'z',
+    'samples': 2
+}])
 def test_extract_2D(labels_te):
     labels = labels_all
     tensor = LabelTensor(data, labels)
@@ -58,7 +57,8 @@ def test_extract_2D(labels_te):
     assert new.ndim == tensor.ndim
     assert new.shape[1] == 1
     assert new.shape[0] == 1
-    assert torch.all(torch.isclose(data[2,2].reshape(1, 1), new))
+    assert torch.all(torch.isclose(data[2, 2].reshape(1, 1), new))
+
 
 def test_extract_3D():
     data = torch.rand(20, 3, 4)
@@ -72,10 +72,7 @@ def test_extract_3D():
             "dof": range(4)
         },
     }
-    labels_te = {
-        'space': ['x', 'z'],
-        'time': range(1, 4)
-    }
+    labels_te = {'space': ['x', 'z'], 'time': range(1, 4)}
 
     tensor = LabelTensor(data, labels)
     new = tensor.extract(labels_te)
@@ -84,15 +81,13 @@ def test_extract_3D():
     assert new.shape[0] == 20
     assert new.shape[1] == 2
     assert new.shape[2] == 3
-    assert torch.all(torch.isclose(
-        data[:, 0::2, 1:4].reshape(20, 2, 3),
-        new
-    ))
+    assert torch.all(torch.isclose(data[:, 0::2, 1:4].reshape(20, 2, 3), new))
     assert tensor2.ndim == tensor.ndim
     assert tensor2.shape == tensor.shape
     assert tensor.full_labels == tensor2.full_labels
     assert new.shape != tensor.shape
 
+
 def test_concatenation_3D():
     data_1 = torch.rand(20, 3, 4)
     labels_1 = ['x', 'y', 'z', 'w']
@@ -152,27 +147,28 @@ def test_concatenation_3D():
 
 
 def test_summation():
-    lt1 = LabelTensor(torch.ones(20,3), labels_all)
-    lt2 = LabelTensor(torch.ones(30,3), ['x', 'y', 'z'])
+    lt1 = LabelTensor(torch.ones(20, 3), labels_all)
+    lt2 = LabelTensor(torch.ones(30, 3), ['x', 'y', 'z'])
     with pytest.raises(RuntimeError):
         LabelTensor.summation([lt1, lt2])
-    lt1 = LabelTensor(torch.ones(20,3), labels_all)
-    lt2 = LabelTensor(torch.ones(20,3), labels_all)
+    lt1 = LabelTensor(torch.ones(20, 3), labels_all)
+    lt2 = LabelTensor(torch.ones(20, 3), labels_all)
     lt_sum = LabelTensor.summation([lt1, lt2])
     assert lt_sum.ndim == lt_sum.ndim
     assert lt_sum.shape[0] == 20
     assert lt_sum.shape[1] == 3
     assert lt_sum.full_labels == labels_all
-    assert torch.eq(lt_sum.tensor, torch.ones(20,3)*2).all()
-    lt1 = LabelTensor(torch.ones(20,3), labels_all)
-    lt2 = LabelTensor(torch.ones(20,3), labels_all)
+    assert torch.eq(lt_sum.tensor, torch.ones(20, 3) * 2).all()
+    lt1 = LabelTensor(torch.ones(20, 3), labels_all)
+    lt2 = LabelTensor(torch.ones(20, 3), labels_all)
     lt3 = LabelTensor(torch.zeros(20, 3), labels_all)
     lt_sum = LabelTensor.summation([lt1, lt2, lt3])
     assert lt_sum.ndim == lt_sum.ndim
     assert lt_sum.shape[0] == 20
     assert lt_sum.shape[1] == 3
     assert lt_sum.full_labels == labels_all
-    assert torch.eq(lt_sum.tensor, torch.ones(20,3)*2).all()
+    assert torch.eq(lt_sum.tensor, torch.ones(20, 3) * 2).all()
+
 
 def test_append_3D():
     data_1 = torch.rand(20, 3, 2)
@@ -187,6 +183,7 @@ def test_append_3D():
     assert lt1.full_labels[1]['dof'] == range(3)
     assert lt1.full_labels[2]['dof'] == ['x', 'y', 'z', 'w']
 
+
 def test_append_2D():
     data_1 = torch.rand(20, 2)
     labels_1 = ['x', 'y']
@@ -199,12 +196,31 @@ def test_append_2D():
     assert lt1.full_labels[0]['dof'] == range(400)
     assert lt1.full_labels[1]['dof'] == ['x', 'y', 'z', 'w']
 
+
 def test_vstack_3D():
     data_1 = torch.rand(20, 3, 2)
-    labels_1 = {1:{'dof': ['a', 'b', 'c'], 'name': 'first'}, 2: {'dof': ['x', 'y'], 'name': 'second'}}
+    labels_1 = {
+        1: {
+            'dof': ['a', 'b', 'c'],
+            'name': 'first'
+        },
+        2: {
+            'dof': ['x', 'y'],
+            'name': 'second'
+        }
+    }
     lt1 = LabelTensor(data_1, labels_1)
     data_2 = torch.rand(20, 3, 2)
-    labels_1 = {1:{'dof': ['a', 'b', 'c'], 'name': 'first'}, 2: {'dof': ['x', 'y'], 'name': 'second'}}
+    labels_1 = {
+        1: {
+            'dof': ['a', 'b', 'c'],
+            'name': 'first'
+        },
+        2: {
+            'dof': ['x', 'y'],
+            'name': 'second'
+        }
+    }
     lt2 = LabelTensor(data_2, labels_1)
     lt_stacked = LabelTensor.vstack([lt1, lt2])
     assert lt_stacked.shape == (40, 3, 2)
@@ -214,12 +230,13 @@ def test_vstack_3D():
     assert lt_stacked.full_labels[1]['name'] == 'first'
     assert lt_stacked.full_labels[2]['name'] == 'second'
 
+
 def test_vstack_2D():
     data_1 = torch.rand(20, 2)
-    labels_1 = { 1: {'dof': ['x', 'y'], 'name': 'second'}}
+    labels_1 = {1: {'dof': ['x', 'y'], 'name': 'second'}}
     lt1 = LabelTensor(data_1, labels_1)
     data_2 = torch.rand(20, 2)
-    labels_1 = { 1: {'dof': ['x', 'y'], 'name': 'second'}}
+    labels_1 = {1: {'dof': ['x', 'y'], 'name': 'second'}}
     lt2 = LabelTensor(data_2, labels_1)
     lt_stacked = LabelTensor.vstack([lt1, lt2])
     assert lt_stacked.shape == (40, 2)
@@ -228,35 +245,36 @@ def test_vstack_2D():
     assert lt_stacked.full_labels[0]['name'] == 0
     assert lt_stacked.full_labels[1]['name'] == 'second'
 
+
 def test_sorting():
     data = torch.ones(20, 5)
-    data[:,0] = data[:,0]*4
-    data[:,1] = data[:,1]*2
-    data[:,2] = data[:,2]
-    data[:,3] = data[:,3]*5
-    data[:,4] = data[:,4]*3
+    data[:, 0] = data[:, 0] * 4
+    data[:, 1] = data[:, 1] * 2
+    data[:, 2] = data[:, 2]
+    data[:, 3] = data[:, 3] * 5
+    data[:, 4] = data[:, 4] * 3
     labels = ['d', 'b', 'a', 'e', 'c']
     lt_data = LabelTensor(data, labels)
     lt_sorted = LabelTensor.sort_labels(lt_data)
-    assert lt_sorted.shape == (20,5)
+    assert lt_sorted.shape == (20, 5)
     assert lt_sorted.labels == ['a', 'b', 'c', 'd', 'e']
-    assert torch.eq(lt_sorted.tensor[:,0], torch.ones(20) * 1).all()
-    assert torch.eq(lt_sorted.tensor[:,1], torch.ones(20) * 2).all()
-    assert torch.eq(lt_sorted.tensor[:,2], torch.ones(20) * 3).all()
-    assert torch.eq(lt_sorted.tensor[:,3], torch.ones(20) * 4).all()
-    assert torch.eq(lt_sorted.tensor[:,4], torch.ones(20) * 5).all()
+    assert torch.eq(lt_sorted.tensor[:, 0], torch.ones(20) * 1).all()
+    assert torch.eq(lt_sorted.tensor[:, 1], torch.ones(20) * 2).all()
+    assert torch.eq(lt_sorted.tensor[:, 2], torch.ones(20) * 3).all()
+    assert torch.eq(lt_sorted.tensor[:, 3], torch.ones(20) * 4).all()
+    assert torch.eq(lt_sorted.tensor[:, 4], torch.ones(20) * 5).all()
 
     data = torch.ones(20, 4, 5)
-    data[:,0,:] = data[:,0]*4
-    data[:,1,:] = data[:,1]*2
-    data[:,2,:] = data[:,2]
-    data[:,3,:] = data[:,3]*3
+    data[:, 0, :] = data[:, 0] * 4
+    data[:, 1, :] = data[:, 1] * 2
+    data[:, 2, :] = data[:, 2]
+    data[:, 3, :] = data[:, 3] * 3
     labels = {1: {'dof': ['d', 'b', 'a', 'c'], 'name': 1}}
     lt_data = LabelTensor(data, labels)
     lt_sorted = LabelTensor.sort_labels(lt_data, dim=1)
-    assert lt_sorted.shape == (20,4, 5)
+    assert lt_sorted.shape == (20, 4, 5)
     assert lt_sorted.full_labels[1]['dof'] == ['a', 'b', 'c', 'd']
-    assert torch.eq(lt_sorted.tensor[:,0,:], torch.ones(20,5) * 1).all()
-    assert torch.eq(lt_sorted.tensor[:,1,:], torch.ones(20,5) * 2).all()
-    assert torch.eq(lt_sorted.tensor[:,2,:], torch.ones(20,5) * 3).all()
-    assert torch.eq(lt_sorted.tensor[:,3,:], torch.ones(20,5) * 4).all()
+    assert torch.eq(lt_sorted.tensor[:, 0, :], torch.ones(20, 5) * 1).all()
+    assert torch.eq(lt_sorted.tensor[:, 1, :], torch.ones(20, 5) * 2).all()
+    assert torch.eq(lt_sorted.tensor[:, 2, :], torch.ones(20, 5) * 3).all()
+    assert torch.eq(lt_sorted.tensor[:, 3, :], torch.ones(20, 5) * 4).all()
diff --git a/tests/test_label_tensor/test_label_tensor_01.py b/tests/test_label_tensor/test_label_tensor_01.py
index a2e129d94..57aafb8c9 100644
--- a/tests/test_label_tensor/test_label_tensor_01.py
+++ b/tests/test_label_tensor/test_label_tensor_01.py
@@ -54,9 +54,8 @@ def test_extract_order():
     label_to_extract = ['c', 'a']
     tensor = LabelTensor(data, labels)
     new = tensor.extract(label_to_extract)
-    expected = torch.cat(
-        (data[:, 2].reshape(-1, 1), data[:, 0].reshape(-1, 1)),
-        dim=1)
+    expected = torch.cat((data[:, 2].reshape(-1, 1), data[:, 0].reshape(-1, 1)),
+                         dim=1)
     assert new.labels == label_to_extract
     assert new.shape[1] == len(label_to_extract)
     assert torch.all(torch.isclose(expected, new))
@@ -91,6 +90,7 @@ def test_getitem():
     assert tensor_view.labels == ['a', 'c']
     assert torch.allclose(tensor_view, data[:, 0::2])
 
+
 def test_getitem2():
     tensor = LabelTensor(data, labels)
     tensor_view = tensor[:5]
@@ -101,6 +101,7 @@ def test_getitem2():
     tensor_view = tensor[idx]
     assert tensor_view.labels == labels
 
+
 def test_slice():
     tensor = LabelTensor(data, labels)
     tensor_view = tensor[:5, :2]
@@ -114,4 +115,4 @@ def test_slice():
 
     tensor_view3 = tensor[:, 2]
     assert tensor_view3.labels == labels[2]
-    assert torch.allclose(tensor_view3, data[:, 2].reshape(-1, 1))
\ No newline at end of file
+    assert torch.allclose(tensor_view3, data[:, 2].reshape(-1, 1))
diff --git a/tests/test_operators.py b/tests/test_operators.py
index 1271c3712..35c07911b 100644
--- a/tests/test_operators.py
+++ b/tests/test_operators.py
@@ -39,6 +39,7 @@ def test_grad_scalar_output():
     ]
     assert torch.allclose(grad_tensor_s, true_val)
 
+
 def test_grad_vector_output():
     grad_tensor_v = grad(tensor_v, inp)
     true_val = torch.cat(
@@ -75,6 +76,7 @@ def test_grad_vector_output():
     ]
     assert torch.allclose(grad_tensor_v, true_val)
 
+
 def test_div_vector_output():
     div_tensor_v = div(tensor_v, inp)
     true_val = 2*torch.sum(inp, dim=1).reshape(-1,1)
@@ -88,6 +90,7 @@ def test_div_vector_output():
     assert div_tensor_v.labels == [f'dadx+dbdy']
     assert torch.allclose(div_tensor_v, true_val)
 
+
 def test_laplacian_scalar_output():
     laplace_tensor_s = laplacian(tensor_s, inp)
     true_val = 6*torch.ones_like(laplace_tensor_s)
@@ -101,6 +104,7 @@ def test_laplacian_scalar_output():
     assert laplace_tensor_s.labels == [f"dd{tensor_s.labels[0]}"]
     assert torch.allclose(laplace_tensor_s, true_val)
 
+
 def test_laplacian_vector_output():
     laplace_tensor_v = laplacian(tensor_v, inp)
     true_val = 2*torch.ones_like(tensor_v)
diff --git a/tests/test_optimizer.py b/tests/test_optimizer.py
index 489bbdc05..bdc87cac5 100644
--- a/tests/test_optimizer.py
+++ b/tests/test_optimizer.py
@@ -1,20 +1,18 @@
-
 import torch
 import pytest
 from pina import TorchOptimizer
 
 opt_list = [
-    torch.optim.Adam,
-    torch.optim.AdamW,
-    torch.optim.SGD,
-    torch.optim.RMSprop
+    torch.optim.Adam, torch.optim.AdamW, torch.optim.SGD, torch.optim.RMSprop
 ]
 
+
 @pytest.mark.parametrize("optimizer_class", opt_list)
 def test_constructor(optimizer_class):
     TorchOptimizer(optimizer_class, lr=1e-3)
 
+
 @pytest.mark.parametrize("optimizer_class", opt_list)
 def test_hook(optimizer_class):
     opt = TorchOptimizer(optimizer_class, lr=1e-3)
-    opt.hook(torch.nn.Linear(10, 10).parameters())
\ No newline at end of file
+    opt.hook(torch.nn.Linear(10, 10).parameters())
diff --git a/tests/test_problem.py b/tests/test_problem.py
index 6d3759606..a766806b3 100644
--- a/tests/test_problem.py
+++ b/tests/test_problem.py
@@ -27,42 +27,42 @@ class Poisson(SpatialProblem):
 
     conditions = {
         'gamma1':
-            Condition(domain=CartesianDomain({
-                'x': [0, 1],
-                'y': 1
-            }),
-                equation=FixedValue(0.0)),
+        Condition(domain=CartesianDomain({
+            'x': [0, 1],
+            'y': 1
+        }),
+                  equation=FixedValue(0.0)),
         'gamma2':
-            Condition(domain=CartesianDomain({
-                'x': [0, 1],
-                'y': 0
-            }),
-                equation=FixedValue(0.0)),
+        Condition(domain=CartesianDomain({
+            'x': [0, 1],
+            'y': 0
+        }),
+                  equation=FixedValue(0.0)),
         'gamma3':
-            Condition(domain=CartesianDomain({
-                'x': 1,
-                'y': [0, 1]
-            }),
-                equation=FixedValue(0.0)),
+        Condition(domain=CartesianDomain({
+            'x': 1,
+            'y': [0, 1]
+        }),
+                  equation=FixedValue(0.0)),
         'gamma4':
-            Condition(domain=CartesianDomain({
-                'x': 0,
-                'y': [0, 1]
-            }),
-                equation=FixedValue(0.0)),
+        Condition(domain=CartesianDomain({
+            'x': 0,
+            'y': [0, 1]
+        }),
+                  equation=FixedValue(0.0)),
         'D':
-            Condition(domain=CartesianDomain({
-                'x': [0, 1],
-                'y': [0, 1]
-            }),
-                equation=my_laplace),
+        Condition(domain=CartesianDomain({
+            'x': [0, 1],
+            'y': [0, 1]
+        }),
+                  equation=my_laplace),
         'data':
-            Condition(input_points=in_, output_points=out_)
+        Condition(input_points=in_, output_points=out_)
     }
 
     def poisson_sol(self, pts):
         return -(torch.sin(pts.extract(['x']) * torch.pi) *
-                 torch.sin(pts.extract(['y']) * torch.pi)) / (2 * torch.pi ** 2)
+                 torch.sin(pts.extract(['y']) * torch.pi)) / (2 * torch.pi**2)
 
     truth_solution = poisson_sol
 
@@ -79,7 +79,7 @@ def test_discretise_domain():
         assert poisson_problem.input_pts[b].shape[0] == n
 
     poisson_problem.discretise_domain(n, 'grid', locations=['D'])
-    assert poisson_problem.input_pts['D'].shape[0] == n ** 2
+    assert poisson_problem.input_pts['D'].shape[0] == n**2
     poisson_problem.discretise_domain(n, 'random', locations=['D'])
     assert poisson_problem.input_pts['D'].shape[0] == n
 
@@ -91,6 +91,7 @@ def test_discretise_domain():
 
     poisson_problem.discretise_domain(n)
 
+
 def test_sampling_few_variables():
     n = 10
     poisson_problem = Poisson()
@@ -115,9 +116,8 @@ def test_variables_correct_order_sampling():
                                       variables=['y'])
     assert poisson_problem.input_pts['D'].labels == sorted(
         poisson_problem.input_variables)
-    poisson_problem.discretise_domain(n,
-                                      'grid',
-                                      locations=['D'])
+
+    poisson_problem.discretise_domain(n, 'grid', locations=['D'])
     assert poisson_problem.input_pts['D'].labels == sorted(
         poisson_problem.input_variables)
     poisson_problem.discretise_domain(n,
@@ -131,6 +131,7 @@ def test_variables_correct_order_sampling():
     assert poisson_problem.input_pts['D'].labels == sorted(
         poisson_problem.input_variables)
 
+
 def test_add_points():
     poisson_problem = Poisson()
     poisson_problem.discretise_domain(0,
@@ -139,8 +140,10 @@ def test_add_points():
                                       variables=['x', 'y'])
     new_pts = LabelTensor(torch.tensor([[0.5, -0.5]]), labels=['x', 'y'])
     poisson_problem.add_points({'D': new_pts})
-    assert torch.isclose(poisson_problem.input_pts['D'].extract('x'), new_pts.extract('x'))
-    assert torch.isclose(poisson_problem.input_pts['D'].extract('y'), new_pts.extract('y'))
+    assert torch.isclose(poisson_problem.input_pts['D'].extract('x'),
+                         new_pts.extract('x'))
+    assert torch.isclose(poisson_problem.input_pts['D'].extract('y'),
+                         new_pts.extract('y'))
 
 
 def test_collector():

From 7c612e3455e77e0ef90fbc928efd32eaac2a1767 Mon Sep 17 00:00:00 2001
From: FilippoOlivo <filippo@filippoolivo.com>
Date: Tue, 22 Oct 2024 14:54:22 +0200
Subject: [PATCH 26/55] Correct codacy warnings

---
 pina/collector.py                 | 21 ++++++++++++++-------
 pina/condition/condition.py       | 23 ++++++++++++-----------
 pina/data/base_dataset.py         |  2 +-
 pina/data/data_module.py          | 16 +++++++++-------
 pina/data/pina_batch.py           |  3 +++
 pina/data/pina_dataloader.py      |  2 +-
 pina/data/pina_subset.py          |  5 +++++
 pina/data/unsupervised_dataset.py |  5 +++--
 pina/optim/torch_optimizer.py     |  1 +
 pina/trainer.py                   | 19 +++++++++++--------
 10 files changed, 60 insertions(+), 37 deletions(-)

diff --git a/pina/collector.py b/pina/collector.py
index f9ef194db..4ebf236c8 100644
--- a/pina/collector.py
+++ b/pina/collector.py
@@ -48,7 +48,8 @@ def store_fixed_data(self):
         for condition_name, condition in self.problem.conditions.items():
             # if the condition is not ready and domain is not attribute
             # of condition, we get and store the data
-            if (not self._is_conditions_ready[condition_name]) and (not hasattr(condition, "domain")):
+            if (not self._is_conditions_ready[condition_name]) and (
+            not hasattr(condition, "domain")):
                 # get data
                 keys = condition.__slots__
                 values = [getattr(condition, name) for name in keys]
@@ -69,7 +70,8 @@ def store_sample_domains(self, n, mode, variables, sample_locations):
                     already_sampled = []
                 # if we have sampled the condition but not all variables
                 else:
-                    already_sampled = [self.data_collections[loc]['input_points']]
+                    already_sampled = [
+                        self.data_collections[loc]['input_points']]
             # if the condition is ready but we want to sample again
             else:
                 self._is_conditions_ready[loc] = False
@@ -77,11 +79,13 @@ def store_sample_domains(self, n, mode, variables, sample_locations):
 
             # get the samples
             samples = [
-                          condition.domain.sample(n=n, mode=mode, variables=variables)
+                          condition.domain.sample(n=n, mode=mode,
+                                                  variables=variables)
                       ] + already_sampled
             pts = merge_tensors(samples)
             if (
-                    set(pts.labels).issubset(sorted(self.problem.input_variables))
+                    set(pts.labels).issubset(
+                        sorted(self.problem.input_variables))
             ):
                 pts = pts.sort_labels()
                 if sorted(pts.labels) == sorted(self.problem.input_variables):
@@ -89,7 +93,8 @@ def store_sample_domains(self, n, mode, variables, sample_locations):
                 values = [pts, condition.equation]
                 self.data_collections[loc] = dict(zip(keys, values))
             else:
-                raise RuntimeError('Try to sample variables which are not in problem defined in the problem')
+                raise RuntimeError(
+                    'Try to sample variables which are not in problem defined in the problem')
 
     def add_points(self, new_points_dict):
         """
@@ -100,5 +105,7 @@ def add_points(self, new_points_dict):
         """
         for k, v in new_points_dict.items():
             if not self._is_conditions_ready[k]:
-                raise RuntimeError('Cannot add points on a non sampled condition')
-            self.data_collections[k]['input_points'] = self.data_collections[k]['input_points'].vstack(v)
+                raise RuntimeError(
+                    'Cannot add points on a non sampled condition')
+            self.data_collections[k]['input_points'] = self.data_collections[k][
+                'input_points'].vstack(v)
diff --git a/pina/condition/condition.py b/pina/condition/condition.py
index 09180cc6e..01965fe0d 100644
--- a/pina/condition/condition.py
+++ b/pina/condition/condition.py
@@ -5,6 +5,7 @@
 from .input_output_condition import InputOutputPointsCondition
 from .data_condition import DataConditionInterface
 
+
 class Condition:
     """
     The class ``Condition`` is used to represent the constraints (physical
@@ -38,23 +39,23 @@ class Condition:
     """
 
     __slots__ = list(
-                    set(
-                        InputOutputPointsCondition.__slots__ + 
-                        InputPointsEquationCondition.__slots__ +
-                        DomainEquationCondition.__slots__ +
-                        DataConditionInterface.__slots__
-                    )
-                )
+        set(
+            InputOutputPointsCondition.__slots__ +
+            InputPointsEquationCondition.__slots__ +
+            DomainEquationCondition.__slots__ +
+            DataConditionInterface.__slots__
+        )
+    )
 
     def __new__(cls, *args, **kwargs):
-        
+
         if len(args) != 0:
             raise ValueError(
                 "Condition takes only the following keyword "
                 f"arguments: {Condition.__slots__}."
             )
-    
-        sorted_keys = sorted(kwargs.keys()) 
+
+        sorted_keys = sorted(kwargs.keys())
         if sorted_keys == sorted(InputOutputPointsCondition.__slots__):
             return InputOutputPointsCondition(**kwargs)
         elif sorted_keys == sorted(InputPointsEquationCondition.__slots__):
@@ -66,4 +67,4 @@ def __new__(cls, *args, **kwargs):
         elif sorted_keys == DataConditionInterface.__slots__[0]:
             return DataConditionInterface(**kwargs)
         else:
-            raise ValueError(f"Invalid keyword arguments {kwargs.keys()}.")
\ No newline at end of file
+            raise ValueError(f"Invalid keyword arguments {kwargs.keys()}.")
diff --git a/pina/data/base_dataset.py b/pina/data/base_dataset.py
index f1d17ae31..b15a0be2a 100644
--- a/pina/data/base_dataset.py
+++ b/pina/data/base_dataset.py
@@ -27,7 +27,7 @@ def __new__(cls, problem, device):
         if not hasattr(cls, '__slots__'):
             raise TypeError(
                 'Something is wrong, __slots__ must be defined in subclasses.')
-        return super(BaseDataset, cls).__new__(cls)
+        return object.__new__(cls)
 
     def __init__(self, problem, device):
         """"
diff --git a/pina/data/data_module.py b/pina/data/data_module.py
index e4e8a450f..25c7e54ed 100644
--- a/pina/data/data_module.py
+++ b/pina/data/data_module.py
@@ -26,7 +26,7 @@ def __init__(self,
                  eval_size=.1,
                  batch_size=None,
                  shuffle=True,
-                 datasets = None):
+                 datasets=None):
         """
         Initialize the object, creating dataset based on input problem
         :param AbstractProblem problem: PINA problem
@@ -38,9 +38,11 @@ def __init__(self,
         :param datasets: list of datasets objects
         """
         super().__init__()
-        dataset_classes = [SupervisedDataset, UnsupervisedDataset, SamplePointDataset]
+        dataset_classes = [SupervisedDataset, UnsupervisedDataset,
+                           SamplePointDataset]
         if datasets is None:
-            self.datasets = [DatasetClass(problem, device) for DatasetClass in dataset_classes]
+            self.datasets = [DatasetClass(problem, device) for DatasetClass in
+                             dataset_classes]
         else:
             self.datasets = datasets
 
@@ -100,8 +102,6 @@ def extract_conditions(self):
             for key, value in dataset.condition_names.items()
         }
 
-
-
     def train_dataloader(self):
         """
         Return the training dataloader for the dataset
@@ -158,11 +158,13 @@ def dataset_split(dataset, lengths, seed=None, shuffle=True):
             if seed is not None:
                 generator = torch.Generator()
                 generator.manual_seed(seed)
-                indices = torch.randperm(sum(lengths), generator=generator).tolist()
+                indices = torch.randperm(sum(lengths),
+                                         generator=generator).tolist()
             else:
                 indices = torch.arange(sum(lengths)).tolist()
         else:
-            indices = torch.arange(0, sum(lengths), 1, dtype=torch.uint8).tolist()
+            indices = torch.arange(0, sum(lengths), 1,
+                                   dtype=torch.uint8).tolist()
         offsets = [
             sum(lengths[:i]) if i > 0 else 0 for i in range(len(lengths))
         ]
diff --git a/pina/data/pina_batch.py b/pina/data/pina_batch.py
index f61e0020b..ed34a910b 100644
--- a/pina/data/pina_batch.py
+++ b/pina/data/pina_batch.py
@@ -5,6 +5,9 @@
 
 
 class Batch:
+    """
+    Implementation of the Batch class used during training to perform SGD optimization.
+    """
 
     def __init__(self, dataset_dict, idx_dict):
 
diff --git a/pina/data/pina_dataloader.py b/pina/data/pina_dataloader.py
index cbd8fe82d..e2d3fb76e 100644
--- a/pina/data/pina_dataloader.py
+++ b/pina/data/pina_dataloader.py
@@ -33,7 +33,7 @@ def _init_batches(self, batch_size=None):
         Create batches according to the batch_size provided in input.
         """
         self.batches = []
-        n_elements = sum([len(v) for v in self.dataset_dict.values()])
+        n_elements = sum(len(v) for v in self.dataset_dict.values())
         if batch_size is None:
             batch_size = n_elements
         indexes_dict = {}
diff --git a/pina/data/pina_subset.py b/pina/data/pina_subset.py
index 41571f92b..844321bd2 100644
--- a/pina/data/pina_subset.py
+++ b/pina/data/pina_subset.py
@@ -1,3 +1,8 @@
+"""
+Module for PinaSubset class
+"""
+
+
 class PinaSubset:
     """
     TODO
diff --git a/pina/data/unsupervised_dataset.py b/pina/data/unsupervised_dataset.py
index f4e8fb345..18cf296f5 100644
--- a/pina/data/unsupervised_dataset.py
+++ b/pina/data/unsupervised_dataset.py
@@ -6,8 +6,9 @@
 
 class UnsupervisedDataset(BaseDataset):
     """
-    This class extend BaseDataset class to handle unsupervised dataset,
-    composed of input points and, optionally, conditional variables
+    This class extend BaseDataset class to handle
+    unsupervised dataset,composed of input points
+    and, optionally, conditional variables
     """
     data_type = 'unsupervised'
     __slots__ = ['input_points', 'conditional_variables']
diff --git a/pina/optim/torch_optimizer.py b/pina/optim/torch_optimizer.py
index ed90846c6..54818d5a5 100644
--- a/pina/optim/torch_optimizer.py
+++ b/pina/optim/torch_optimizer.py
@@ -13,6 +13,7 @@ def __init__(self, optimizer_class, **kwargs):
 
         self.optimizer_class = optimizer_class
         self.kwargs = kwargs
+        self.optimizer_instance = None
 
     def hook(self, parameters):
         self.optimizer_instance = self.optimizer_class(parameters,
diff --git a/pina/trainer.py b/pina/trainer.py
index 49c6a4017..884eef77e 100644
--- a/pina/trainer.py
+++ b/pina/trainer.py
@@ -9,7 +9,8 @@
 
 class Trainer(pytorch_lightning.Trainer):
 
-    def __init__(self, solver, batch_size=None, train_size=.7, test_size=.2, eval_size=.1, **kwargs):
+    def __init__(self, solver, batch_size=None, train_size=.7, test_size=.2,
+                 eval_size=.1, **kwargs):
         """
         PINA Trainer class for costumizing every aspect of training via flags.
 
@@ -39,10 +40,9 @@ def __init__(self, solver, batch_size=None, train_size=.7, test_size=.2, eval_si
         self._create_loader()
         self._move_to_device()
 
-
     def _move_to_device(self):
         device = self._accelerator_connector._parallel_devices[0]
-        
+
         # move parameters to device
         pb = self.solver.problem
         if hasattr(pb, "unknown_parameters"):
@@ -59,11 +59,13 @@ def _create_loader(self):
         """
         if not self.solver.problem.collector.full:
             error_message = '\n'.join(
-                [f'{" " * 13} ---> Condition {key} {"sampled" if value else "not sampled"}' 
-                for key, value in self.solver.problem.collector._is_conditions_ready.items()])
+                [
+                    f'{" " * 13} ---> Condition {key} {"sampled" if value else "not sampled"}'
+                    for key, value in
+                    self.solver.problem.collector._is_conditions_ready.items()])
             raise RuntimeError('Cannot create Trainer if not all conditions '
-                                'are sampled. The Trainer got the following:\n'
-                                f'{error_message}')
+                               'are sampled. The Trainer got the following:\n'
+                               f'{error_message}')
         devices = self._accelerator_connector._parallel_devices
 
         if len(devices) > 1:
@@ -72,7 +74,8 @@ def _create_loader(self):
         device = devices[0]
 
         data_module = PinaDataModule(problem=self.solver.problem, device=device,
-                                     train_size=self.train_size, test_size=self.test_size,
+                                     train_size=self.train_size,
+                                     test_size=self.test_size,
                                      eval_size=self.eval_size)
         data_module.setup()
         self._loader = data_module.train_dataloader()

From 1f6a410dee887314c94c8307621a424b587449a9 Mon Sep 17 00:00:00 2001
From: FilippoOlivo <filippo@filippoolivo.com>
Date: Wed, 23 Oct 2024 14:54:31 +0200
Subject: [PATCH 27/55] Fix bug and improve __getitem__

---
 pina/label_tensor.py | 172 +++++++++++++++++++++++++++++--------------
 1 file changed, 115 insertions(+), 57 deletions(-)

diff --git a/pina/label_tensor.py b/pina/label_tensor.py
index 65655e9dc..62d87950a 100644
--- a/pina/label_tensor.py
+++ b/pina/label_tensor.py
@@ -22,9 +22,6 @@ def __new__(cls, x, labels, *args, **kwargs):
     def tensor(self):
         return self.as_subclass(Tensor)
 
-    def __len__(self) -> int:
-        return super().__len__()
-
     def __init__(self, x, labels):
         """
         Construct a `LabelTensor` by passing a dict of the labels
@@ -75,7 +72,7 @@ def labels(self, labels):
             labels = [labels]
             self.update_labels_from_list(labels)
         else:
-            raise ValueError(f"labels must be list, dict or string.")
+            raise ValueError("labels must be list, dict or string.")
         self.set_names()
 
     def set_names(self):
@@ -98,10 +95,9 @@ def extract(self, label_to_extract):
             label_to_extract = [label_to_extract]
         if isinstance(label_to_extract, (tuple, list)):
             return self._extract_from_list(label_to_extract)
-        elif isinstance(label_to_extract, dict):
+        if isinstance(label_to_extract, dict):
             return self._extract_from_dict(label_to_extract)
-        else:
-            raise ValueError('labels_to_extract must be str or list or dict')
+        raise ValueError('labels_to_extract must be str or list or dict')
 
     def _extract_from_list(self, labels_to_extract):
         # Store locally all necessary obj/variables
@@ -112,7 +108,8 @@ def _extract_from_list(self, labels_to_extract):
 
         # Verify if all the labels in labels_to_extract are in last dimension
         if set(labels_to_extract).issubset(last_dim_label) is False:
-            raise ValueError('Cannot extract a dof which is not in the original LabelTensor')
+            raise ValueError(
+                'Cannot extract a dof which is not in the original LabelTensor')
 
         # Extract index to extract
         idx_to_extract = [last_dim_label.index(i) for i in labels_to_extract]
@@ -142,9 +139,12 @@ def _extract_from_dict(self, labels_to_extract):
             if isinstance(labels_to_extract[k], (int, str)):
                 labels_to_extract[k] = [labels_to_extract[k]]
             if set(labels_to_extract[k]).issubset(dim_labels) is False:
-                raise ValueError('Cannot extract a dof which is not in the original LabelTensor')
+                raise ValueError(
+                    'Cannot extract a dof which is not in the original '
+                    'LabelTensor')
             idx_to_extract = [dim_labels.index(i) for i in labels_to_extract[k]]
-            indexer = [slice(None)] * idx_dim + [idx_to_extract] + [slice(None)] * (ndim - idx_dim - 1)
+            indexer = [slice(None)] * idx_dim + [idx_to_extract] + [
+                slice(None)] * (ndim - idx_dim - 1)
             new_tensor = new_tensor[indexer]
             dim_new_label = {idx_dim: {
                 'dof': labels_to_extract[k],
@@ -168,7 +168,8 @@ def __str__(self):
     @staticmethod
     def cat(tensors, dim=0):
         """
-        Stack a list of tensors. For example, given a tensor `a` of shape `(n,m,dof)` and a tensor `b` of dimension `(n',m,dof)`
+        Stack a list of tensors. For example, given a tensor `a` of shape
+        `(n,m,dof)` and a tensor `b` of dimension `(n',m,dof)`
         the resulting tensor is of shape `(n+n',m,dof)`
 
         :param tensors: tensors to concatenate
@@ -182,7 +183,8 @@ def cat(tensors, dim=0):
             return []
         if len(tensors) == 1:
             return tensors[0]
-        new_labels_cat_dim = LabelTensor._check_validity_before_cat(tensors, dim)
+        new_labels_cat_dim = LabelTensor._check_validity_before_cat(tensors,
+                                                                    dim)
 
         # Perform cat on tensors
         new_tensor = torch.cat(tensors, dim=dim)
@@ -190,7 +192,8 @@ def cat(tensors, dim=0):
         # Update labels
         labels = tensors[0].full_labels
         labels.pop(dim)
-        new_labels_cat_dim = new_labels_cat_dim if len(set(new_labels_cat_dim)) == len(new_labels_cat_dim) \
+        new_labels_cat_dim = new_labels_cat_dim if len(
+            set(new_labels_cat_dim)) == len(new_labels_cat_dim) \
             else range(new_tensor.shape[dim])
         labels[dim] = {'dof': new_labels_cat_dim,
                        'name': tensors[1].full_labels[dim]['name']}
@@ -200,7 +203,8 @@ def cat(tensors, dim=0):
     def _check_validity_before_cat(tensors, dim):
         n_dims = tensors[0].ndim
         new_labels_cat_dim = []
-        # Check if names and dof of the labels are the same in all dimensions except in dim
+        # Check if names and dof of the labels are the same in all dimensions
+        # except in dim
         for i in range(n_dims):
             name = tensors[0].full_labels[i]['name']
             if i != dim:
@@ -209,13 +213,15 @@ def _check_validity_before_cat(tensors, dim):
                     dof_to_check = tensor.full_labels[i]['dof']
                     name_to_check = tensor.full_labels[i]['name']
                     if dof != dof_to_check or name != name_to_check:
-                        raise ValueError('dimensions must have the same dof and name')
+                        raise ValueError(
+                            'dimensions must have the same dof and name')
             else:
                 for tensor in tensors:
                     new_labels_cat_dim += tensor.full_labels[i]['dof']
                     name_to_check = tensor.full_labels[i]['name']
                     if name != name_to_check:
-                        raise ValueError('Dimensions to concatenate must have the same name')
+                        raise ValueError(
+                            'Dimensions to concatenate must have the same name')
         return new_labels_cat_dim
 
     def requires_grad_(self, mode=True):
@@ -259,11 +265,13 @@ def init_labels(self):
 
     def update_labels_from_dict(self, labels):
         """
-        Update the internal label representation according to the values passed as input.
+        Update the internal label representation according to the values passed
+        as input.
 
         :param labels: The label(s) to update.
         :type labels: dict
-        :raises ValueError: dof list contain duplicates or number of dof does not match with tensor shape
+        :raises ValueError: dof list contain duplicates or number of dof does
+        not match with tensor shape
         """
         tensor_shape = self.tensor.shape
         # Check dimensionality
@@ -271,19 +279,22 @@ def update_labels_from_dict(self, labels):
             if len(v['dof']) != len(set(v['dof'])):
                 raise ValueError("dof must be unique")
             if len(v['dof']) != tensor_shape[k]:
-                raise ValueError('Number of dof does not match with tensor dimension')
+                raise ValueError(
+                    'Number of dof does not match with tensor dimension')
         # Perform update
         self._labels.update(labels)
 
     def update_labels_from_list(self, labels):
         """
-        Given a list of dof, this method update the internal label representation
+        Given a list of dof, this method update the internal label
+        representation
 
         :param labels: The label(s) to update.
         :type labels: list
         """
         # Create a dict with labels
-        last_dim_labels = {self.tensor.ndim - 1: {'dof': labels, 'name': self.tensor.ndim - 1}}
+        last_dim_labels = {
+            self.tensor.ndim - 1: {'dof': labels, 'name': self.tensor.ndim - 1}}
         self.update_labels_from_dict(last_dim_labels)
 
     @staticmethod
@@ -302,15 +313,16 @@ def summation(tensors):
                     break
         # Sum tensors
         data = torch.zeros(tensors[0].tensor.shape)
-        for i in range(len(tensors)):
-            data += tensors[i].tensor
+        for tensor in tensors:
+            data += tensor.tensor
         new_tensor = LabelTensor(data, labels)
         return new_tensor
 
     def append(self, tensor, mode='std'):
         if mode == 'std':
             # Call cat on last dimension
-            new_label_tensor = LabelTensor.cat([self, tensor], dim=self.tensor.ndim - 1)
+            new_label_tensor = LabelTensor.cat([self, tensor],
+                                               dim=self.tensor.ndim - 1)
         elif mode == 'cross':
             # Crete tensor and call cat on last dimension
             tensor1 = self
@@ -318,8 +330,10 @@ def append(self, tensor, mode='std'):
             n1 = tensor1.shape[0]
             n2 = tensor2.shape[0]
             tensor1 = LabelTensor(tensor1.repeat(n2, 1), labels=tensor1.labels)
-            tensor2 = LabelTensor(tensor2.repeat_interleave(n1, dim=0), labels=tensor2.labels)
-            new_label_tensor = LabelTensor.cat([tensor1, tensor2], dim=self.tensor.ndim - 1)
+            tensor2 = LabelTensor(tensor2.repeat_interleave(n1, dim=0),
+                                  labels=tensor2.labels)
+            new_label_tensor = LabelTensor.cat([tensor1, tensor2],
+                                               dim=self.tensor.ndim - 1)
         else:
             raise ValueError('mode must be either "std" or "cross"')
         return new_label_tensor
@@ -339,47 +353,90 @@ def vstack(label_tensors):
 
     def __getitem__(self, index):
         """
-        Return a copy of the selected tensor.
+        TODO: Complete docstring
+        :param index:
+        :return:
         """
 
-        if isinstance(index, str) or (isinstance(index, (tuple, list)) and all(isinstance(a, str) for a in index)):
+        if isinstance(index, str) or (isinstance(index, (tuple, list)) and all(
+                isinstance(a, str) for a in index)):
             return self.extract(index)
-
         selected_lt = super().__getitem__(index)
 
-        try:
-            len_index = len(index)
-        except TypeError:
-            len_index = 1
-
-        if isinstance(index, int) or len_index == 1:
-            if selected_lt.ndim == 1:
-                selected_lt = selected_lt.reshape(1, -1)
-            if hasattr(self, "labels"):
-                new_labels = deepcopy(self.full_labels)
-                new_labels.pop(0)
-                selected_lt.labels = new_labels
-        elif len(index) == self.tensor.ndim:
+        if isinstance(index, (int, slice)):
+            return self._getitem_int_slice(index, selected_lt)
+
+        if len(index) == self.tensor.ndim:
+            return self._getitem_full_dim_indexing(index, selected_lt)
+
+        if isinstance(index, torch.Tensor) or (
+                isinstance(index, (tuple, list)) and all(
+            isinstance(x, int) for x in index)):
+            return self._getitem_permutation(index, selected_lt)
+        raise ValueError('Not recognized index type')
+
+    def _getitem_int_slice(self, index, selected_lt):
+        """
+        :param index:
+        :param selected_lt:
+        :return:
+        """
+        if selected_lt.ndim == 1:
+            selected_lt = selected_lt.reshape(1, -1)
+        if hasattr(self, "labels"):
             new_labels = deepcopy(self.full_labels)
-            if selected_lt.ndim == 1:
-                selected_lt = selected_lt.reshape(-1, 1)
-            for j in range(selected_lt.ndim):
+            to_update_dof = new_labels[0]['dof'][index]
+            to_update_dof = to_update_dof if isinstance(to_update_dof, (
+                tuple, list, range)) else [to_update_dof]
+            new_labels.update(
+                {0: {'dof': to_update_dof, 'name': new_labels[0]['name']}}
+            )
+            selected_lt.labels = new_labels
+        return selected_lt
+
+    def _getitem_full_dim_indexing(self, index, selected_lt):
+        new_labels = {}
+        old_labels = self.full_labels
+        if selected_lt.ndim == 1:
+            selected_lt = selected_lt.reshape(-1, 1)
+            new_labels = deepcopy(old_labels)
+            new_labels[1].update({'dof': old_labels[1]['dof'][index[1]],
+                                  'name': old_labels[1]['name']})
+        idx = 0
+        for j in range(selected_lt.ndim):
+            if not isinstance(index[j], int):
                 if hasattr(self, "labels"):
-                    if isinstance(index[j], list):
-                        new_labels.update({j: {'dof': [new_labels[j]['dof'][i] for i in index[1]],
-                                               'name': new_labels[j]['name']}})
-                    else:
-                        new_labels.update({j: {'dof': new_labels[j]['dof'][index[j]],
-                                               'name': new_labels[j]['name']}})
+                    new_labels.update(
+                        self._update_label_for_dim(old_labels, index[j], idx))
+                idx += 1
+        selected_lt.labels = new_labels
+        return selected_lt
 
-            selected_lt.labels = new_labels
-        else:
-            new_labels = deepcopy(self.full_labels)
-            new_labels.update({0: {'dof': list[index], 'name': new_labels[0]['name']}})
-            selected_lt.labels = self.labels
+    def _getitem_permutation(self, index, selected_lt):
 
+        new_labels = deepcopy(self.full_labels)
+        new_labels.update(self._update_label_for_dim(self.full_labels, index,
+                                                     0))
+        selected_lt.labels = self.labels
         return selected_lt
 
+    @staticmethod
+    def _update_label_for_dim(old_labels, index, dim):
+        """
+        TODO
+        :param old_labels:
+        :param index:
+        :param dim:
+        :return:
+        """
+        if isinstance(index, list):
+            return {dim: {'dof': [old_labels[dim]['dof'][i] for i in index],
+                          'name': old_labels[dim]['name']}}
+        else:
+            return {dim: {'dof': old_labels[dim]['dof'][index],
+                          'name': old_labels[dim]['name']}}
+
+
     def sort_labels(self, dim=None):
         def argsort(lst):
             return sorted(range(len(lst)), key=lambda x: lst[x])
@@ -391,5 +448,6 @@ def argsort(lst):
         indexer = [slice(None)] * self.tensor.ndim
         indexer[dim] = sorted_index
         new_labels = deepcopy(self.full_labels)
-        new_labels[dim] = {'dof': sorted(labels), 'name': new_labels[dim]['name']}
+        new_labels[dim] = {'dof': sorted(labels),
+                           'name': new_labels[dim]['name']}
         return LabelTensor(self.tensor[indexer], new_labels)

From 83a67aee64fce340d1e6d94a245e6308dd33023f Mon Sep 17 00:00:00 2001
From: FilippoOlivo <filippo@filippoolivo.com>
Date: Wed, 23 Oct 2024 15:04:28 +0200
Subject: [PATCH 28/55] Add Graph support in Dataset and Dataloader

---
 pina/collector.py               |   5 +-
 pina/data/base_dataset.py       |  37 ++++++-----
 pina/data/pina_batch.py         |   3 +-
 pina/data/pina_subset.py        |   9 ++-
 pina/data/sample_dataset.py     |   4 +-
 pina/data/supervised_dataset.py |   3 +-
 pina/label_tensor.py            |   4 +-
 pina/solvers/solver.py          |  15 +++--
 pina/solvers/supervised.py      |   3 +-
 pina/trainer.py                 |   7 +-
 tests/test_dataset.py           | 110 ++++++++++++++++++++------------
 11 files changed, 125 insertions(+), 75 deletions(-)

diff --git a/pina/collector.py b/pina/collector.py
index 4ebf236c8..c48c674e8 100644
--- a/pina/collector.py
+++ b/pina/collector.py
@@ -49,7 +49,7 @@ def store_fixed_data(self):
             # if the condition is not ready and domain is not attribute
             # of condition, we get and store the data
             if (not self._is_conditions_ready[condition_name]) and (
-            not hasattr(condition, "domain")):
+                    not hasattr(condition, "domain")):
                 # get data
                 keys = condition.__slots__
                 values = [getattr(condition, name) for name in keys]
@@ -94,7 +94,8 @@ def store_sample_domains(self, n, mode, variables, sample_locations):
                 self.data_collections[loc] = dict(zip(keys, values))
             else:
                 raise RuntimeError(
-                    'Try to sample variables which are not in problem defined in the problem')
+                    'Try to sample variables which are not in problem defined '
+                    'in the problem')
 
     def add_points(self, new_points_dict):
         """
diff --git a/pina/data/base_dataset.py b/pina/data/base_dataset.py
index b15a0be2a..d859aac00 100644
--- a/pina/data/base_dataset.py
+++ b/pina/data/base_dataset.py
@@ -4,6 +4,7 @@
 from torch.utils.data import Dataset
 import torch
 from ..label_tensor import LabelTensor
+from ..graph import Graph
 
 
 class BaseDataset(Dataset):
@@ -42,38 +43,43 @@ def __init__(self, problem, device):
         collector = problem.collector
         for slot in self.__slots__:
             setattr(self, slot, [])
-
+        num_el_per_condition = []
         idx = 0
         for name, data in collector.data_collections.items():
-            keys = []
-            for k, v in data.items():
-                if isinstance(v, LabelTensor):
-                    keys.append(k)
+            keys = list(data.keys())
+            current_cond_num_el = None
             if sorted(self.__slots__) == sorted(keys):
-
                 for slot in self.__slots__:
+                    slot_data = data[slot]
+                    if isinstance(slot_data, (LabelTensor, torch.Tensor,
+                                              Graph)):
+                        if current_cond_num_el is None:
+                            current_cond_num_el = len(slot_data)
+                        elif current_cond_num_el != len(slot_data):
+                            raise ValueError('Different number of conditions')
                     current_list = getattr(self, slot)
-                    current_list.append(data[slot])
+                    current_list += [data[slot]] if not (
+                        isinstance(data[slot], list)) else data[slot]
+                num_el_per_condition.append(current_cond_num_el)
                 self.condition_names[idx] = name
                 idx += 1
-
-        if len(getattr(self, self.__slots__[0])) > 0:
-            input_list = getattr(self, self.__slots__[0])
+        if num_el_per_condition:
             self.condition_indices = torch.cat(
                 [
-                    torch.tensor([i] * len(input_list[i]), dtype=torch.uint8)
-                    for i in range(len(self.condition_names))
+                    torch.tensor([i] * num_el_per_condition[i],
+                                 dtype=torch.uint8)
+                    for i in range(len(num_el_per_condition))
                 ],
                 dim=0,
             )
             for slot in self.__slots__:
                 current_attribute = getattr(self, slot)
-                setattr(self, slot, LabelTensor.vstack(current_attribute))
+                if all(isinstance(a, LabelTensor) for a in current_attribute):
+                    setattr(self, slot, LabelTensor.vstack(current_attribute))
         else:
             self.condition_indices = torch.tensor([], dtype=torch.uint8)
             for slot in self.__slots__:
                 setattr(self, slot, torch.tensor([]))
-
         self.device = device
 
     def __len__(self):
@@ -89,11 +95,10 @@ def __getattribute__(self, item):
     def __getitem__(self, idx):
         if isinstance(idx, str):
             return getattr(self, idx).to(self.device)
-
         if isinstance(idx, slice):
             to_return_list = []
             for i in self.__slots__:
-                to_return_list.append(getattr(self, i)[[idx]].to(self.device))
+                to_return_list.append(getattr(self, i)[idx].to(self.device))
             return to_return_list
 
         if isinstance(idx, (tuple, list)):
diff --git a/pina/data/pina_batch.py b/pina/data/pina_batch.py
index ed34a910b..65b5ac5ba 100644
--- a/pina/data/pina_batch.py
+++ b/pina/data/pina_batch.py
@@ -6,7 +6,8 @@
 
 class Batch:
     """
-    Implementation of the Batch class used during training to perform SGD optimization.
+    Implementation of the Batch class used during training to perform SGD
+    optimization.
     """
 
     def __init__(self, dataset_dict, idx_dict):
diff --git a/pina/data/pina_subset.py b/pina/data/pina_subset.py
index 844321bd2..f1347b6c5 100644
--- a/pina/data/pina_subset.py
+++ b/pina/data/pina_subset.py
@@ -1,6 +1,8 @@
 """
 Module for PinaSubset class
 """
+from pina import LabelTensor
+from torch import Tensor
 
 
 class PinaSubset:
@@ -23,4 +25,9 @@ def __len__(self):
         return len(self.indices)
 
     def __getattr__(self, name):
-        return self.dataset.__getattribute__(name)
+        tensor = self.dataset.__getattribute__(name)
+        if isinstance(tensor, (LabelTensor, Tensor)):
+            return tensor[self.indices]
+        if isinstance(tensor, list):
+            return [tensor[i] for i in self.indices]
+        raise AttributeError("No attribute named {}".format(name))
diff --git a/pina/data/sample_dataset.py b/pina/data/sample_dataset.py
index ba8bd19a9..99811cac8 100644
--- a/pina/data/sample_dataset.py
+++ b/pina/data/sample_dataset.py
@@ -2,6 +2,8 @@
 Sample dataset module
 """
 from .base_dataset import BaseDataset
+from ..condition.input_equation_condition import InputPointsEquationCondition
+
 
 class SamplePointDataset(BaseDataset):
     """
@@ -9,4 +11,4 @@ class SamplePointDataset(BaseDataset):
     composed of only input points.
     """
     data_type = 'physics'
-    __slots__ = ['input_points']
+    __slots__ = InputPointsEquationCondition.__slots__
diff --git a/pina/data/supervised_dataset.py b/pina/data/supervised_dataset.py
index 2403e3d0c..be601050a 100644
--- a/pina/data/supervised_dataset.py
+++ b/pina/data/supervised_dataset.py
@@ -6,7 +6,8 @@
 
 class SupervisedDataset(BaseDataset):
     """
-    This class extends the BaseDataset to handle datasets that consist of input-output pairs.
+    This class extends the BaseDataset to handle datasets that consist of
+    input-output pairs.
     """
     data_type = 'supervised'
     __slots__ = ['input_points', 'output_points']
diff --git a/pina/label_tensor.py b/pina/label_tensor.py
index 62d87950a..87def2f8e 100644
--- a/pina/label_tensor.py
+++ b/pina/label_tensor.py
@@ -413,7 +413,6 @@ def _getitem_full_dim_indexing(self, index, selected_lt):
         return selected_lt
 
     def _getitem_permutation(self, index, selected_lt):
-
         new_labels = deepcopy(self.full_labels)
         new_labels.update(self._update_label_for_dim(self.full_labels, index,
                                                      0))
@@ -429,6 +428,8 @@ def _update_label_for_dim(old_labels, index, dim):
         :param dim:
         :return:
         """
+        if isinstance(index, torch.Tensor):
+            index = index.nonzero()
         if isinstance(index, list):
             return {dim: {'dof': [old_labels[dim]['dof'][i] for i in index],
                           'name': old_labels[dim]['name']}}
@@ -436,7 +437,6 @@ def _update_label_for_dim(old_labels, index, dim):
             return {dim: {'dof': old_labels[dim]['dof'][index],
                           'name': old_labels[dim]['name']}}
 
-
     def sort_labels(self, dim=None):
         def argsort(lst):
             return sorted(range(len(lst)), key=lambda x: lst[x])
diff --git a/pina/solvers/solver.py b/pina/solvers/solver.py
index 1c6aa2b2b..6f55dedf0 100644
--- a/pina/solvers/solver.py
+++ b/pina/solvers/solver.py
@@ -38,7 +38,7 @@ def __init__(self,
         check_consistency(problem, AbstractProblem)
         self._check_solver_consistency(problem)
 
-        #Check consistency of models argument and encapsulate in list
+        # Check consistency of models argument and encapsulate in list
         if not isinstance(models, list):
             check_consistency(models, torch.nn.Module)
             # put everything in a list if only one input
@@ -49,17 +49,17 @@ def __init__(self,
                 check_consistency(models[idx], torch.nn.Module)
         len_model = len(models)
 
-        #If use_lt is true add extract operation in input
+        # If use_lt is true add extract operation in input
         if use_lt is True:
-            for idx in range(len(models)):
+            for idx, model in enumerate(models):
                 models[idx] = Network(
-                    model=models[idx],
+                    model=model,
                     input_variables=problem.input_variables,
                     output_variables=problem.output_variables,
                     extra_features=extra_features,
                 )
 
-        #Check scheduler consistency + encapsulation
+        # Check scheduler consistency + encapsulation
         if not isinstance(schedulers, list):
             check_consistency(schedulers, Scheduler)
             schedulers = [schedulers]
@@ -67,7 +67,7 @@ def __init__(self,
             for scheduler in schedulers:
                 check_consistency(scheduler, Scheduler)
 
-        #Check optimizer consistency + encapsulation
+        # Check optimizer consistency + encapsulation
         if not isinstance(optimizers, list):
             check_consistency(optimizers, Optimizer)
             optimizers = [optimizers]
@@ -141,5 +141,6 @@ def _check_solver_consistency(self, problem):
             if not set(self.accepted_condition_types).issubset(
                     condition.condition_type):
                 raise ValueError(
-                    f'{self.__name__} support only dose not support condition {condition.condition_type}'
+                    f'{self.__name__} support only dose not support condition '
+                    f'{condition.condition_type}'
                 )
diff --git a/pina/solvers/supervised.py b/pina/solvers/supervised.py
index a0b0f83ed..62fc99149 100644
--- a/pina/solvers/supervised.py
+++ b/pina/solvers/supervised.py
@@ -130,14 +130,13 @@ def training_step(self, batch, batch_idx):
             if not hasattr(condition, "output_points"):
                 raise NotImplementedError(
                     f"{type(self).__name__} works only in data-driven mode.")
-
             output_pts = out[condition_idx == condition_id]
             input_pts = pts[condition_idx == condition_id]
 
             input_pts.labels = pts.labels
             output_pts.labels = out.labels
 
-            loss = (self.loss_data(input_pts=input_pts, output_pts=output_pts))
+            loss = self.loss_data(input_pts=input_pts, output_pts=output_pts)
             loss = loss.as_subclass(torch.Tensor)
 
         self.log("mean_loss", float(loss), prog_bar=True, logger=True)
diff --git a/pina/trainer.py b/pina/trainer.py
index 884eef77e..3de0d7e80 100644
--- a/pina/trainer.py
+++ b/pina/trainer.py
@@ -60,9 +60,12 @@ def _create_loader(self):
         if not self.solver.problem.collector.full:
             error_message = '\n'.join(
                 [
-                    f'{" " * 13} ---> Condition {key} {"sampled" if value else "not sampled"}'
+                    f"""{" " * 13} ---> Condition {key} {"sampled" if value else
+                    "not sampled"}"""
                     for key, value in
-                    self.solver.problem.collector._is_conditions_ready.items()])
+                    self._solver.problem.collector._is_conditions_ready.items()
+                ]
+            )
             raise RuntimeError('Cannot create Trainer if not all conditions '
                                'are sampled. The Trainer got the following:\n'
                                f'{error_message}')
diff --git a/tests/test_dataset.py b/tests/test_dataset.py
index 653b0d6b6..503ddd683 100644
--- a/tests/test_dataset.py
+++ b/tests/test_dataset.py
@@ -1,13 +1,15 @@
 import math
 import torch
-from pina.data import SamplePointDataset, SupervisedDataset, PinaDataModule, UnsupervisedDataset, unsupervised_dataset
+from pina.data import SamplePointDataset, SupervisedDataset, PinaDataModule, \
+    UnsupervisedDataset
 from pina.data import PinaDataLoader
 from pina import LabelTensor, Condition
 from pina.equation import Equation
 from pina.domain import CartesianDomain
-from pina.problem import SpatialProblem
+from pina.problem import SpatialProblem, AbstractProblem
 from pina.operators import laplacian
 from pina.equation.equation_factory import FixedValue
+from pina.graph import Graph
 
 
 def laplace_equation(input_, output_):
@@ -30,49 +32,49 @@ class Poisson(SpatialProblem):
 
     conditions = {
         'gamma1':
-        Condition(domain=CartesianDomain({
-            'x': [0, 1],
-            'y': 1
-        }),
-                  equation=FixedValue(0.0)),
+            Condition(domain=CartesianDomain({
+                'x': [0, 1],
+                'y': 1
+            }),
+                equation=FixedValue(0.0)),
         'gamma2':
-        Condition(domain=CartesianDomain({
-            'x': [0, 1],
-            'y': 0
-        }),
-                  equation=FixedValue(0.0)),
+            Condition(domain=CartesianDomain({
+                'x': [0, 1],
+                'y': 0
+            }),
+                equation=FixedValue(0.0)),
         'gamma3':
-        Condition(domain=CartesianDomain({
-            'x': 1,
-            'y': [0, 1]
-        }),
-                  equation=FixedValue(0.0)),
+            Condition(domain=CartesianDomain({
+                'x': 1,
+                'y': [0, 1]
+            }),
+                equation=FixedValue(0.0)),
         'gamma4':
-        Condition(domain=CartesianDomain({
-            'x': 0,
-            'y': [0, 1]
-        }),
-                  equation=FixedValue(0.0)),
+            Condition(domain=CartesianDomain({
+                'x': 0,
+                'y': [0, 1]
+            }),
+                equation=FixedValue(0.0)),
         'D':
-        Condition(input_points=LabelTensor(torch.rand(size=(100, 2)),
-                                           ['x', 'y']),
-                  equation=my_laplace),
+            Condition(input_points=LabelTensor(torch.rand(size=(100, 2)),
+                                               ['x', 'y']),
+                      equation=my_laplace),
         'data':
-        Condition(input_points=in_, output_points=out_),
+            Condition(input_points=in_, output_points=out_),
         'data2':
-        Condition(input_points=in2_, output_points=out2_),
+            Condition(input_points=in2_, output_points=out2_),
         'unsupervised':
-        Condition(
-            input_points=LabelTensor(torch.rand(size=(45, 2)), ['x', 'y']),
-            conditional_variables=LabelTensor(torch.ones(size=(45, 1)),
-                                              ['alpha']),
-        ),
+            Condition(
+                input_points=LabelTensor(torch.rand(size=(45, 2)), ['x', 'y']),
+                conditional_variables=LabelTensor(torch.ones(size=(45, 1)),
+                                                  ['alpha']),
+            ),
         'unsupervised2':
-        Condition(
-            input_points=LabelTensor(torch.rand(size=(90, 2)), ['x', 'y']),
-            conditional_variables=LabelTensor(torch.ones(size=(90, 1)),
-                                              ['alpha']),
-        )
+            Condition(
+                input_points=LabelTensor(torch.rand(size=(90, 2)), ['x', 'y']),
+                conditional_variables=LabelTensor(torch.ones(size=(90, 1)),
+                                                  ['alpha']),
+            )
     }
 
 
@@ -98,8 +100,8 @@ def test_data():
     assert dataset.input_points.shape == (61, 2)
     assert dataset['input_points'].labels == ['x', 'y']
     assert dataset.input_points.labels == ['x', 'y']
-    assert dataset['input_points', 3:].shape == (58, 2)
-    assert dataset[3:][1].labels == ['u']
+    assert dataset.input_points[3:].shape == (58, 2)
+    assert dataset.output_points[:3].labels == ['u']
     assert dataset.output_points.shape == (61, 1)
     assert dataset.output_points.labels == ['u']
     assert dataset.condition_indices.dtype == torch.uint8
@@ -193,4 +195,32 @@ def test_loader():
         assert i.unsupervised.input_points.requires_grad == True
 
 
-test_loader()
+coordinates = LabelTensor(torch.rand((100, 100, 2)), labels=['x', 'y'])
+data = LabelTensor(torch.rand((100, 100, 3)), labels=['ux', 'uy', 'p'])
+
+
+class GraphProblem(AbstractProblem):
+    output = LabelTensor(torch.rand((100, 3)), labels=['ux', 'uy', 'p'])
+    input = [Graph.build('radius',
+                         nodes_coordinates=coordinates[i, :, :],
+                         nodes_data=data[i, :, :], radius=0.2)
+             for i in
+             range(100)]
+    output_variables = ['u']
+
+    conditions = {
+        'graph_data': Condition(input_points=input, output_points=output)
+    }
+
+
+graph_problem = GraphProblem()
+
+
+def test_loader_graph():
+    data_module = PinaDataModule(graph_problem, device='cpu', batch_size=10)
+    data_module.setup()
+    loader = data_module.train_dataloader()
+    for i in loader:
+        assert len(i) <= 10
+        assert isinstance(i.supervised.input_points, list)
+        assert all(isinstance(x, Graph) for x in i.supervised.input_points)

From b712d78cf24840675c198d7fea5891ac464de690 Mon Sep 17 00:00:00 2001
From: FilippoOlivo <filippo@filippoolivo.com>
Date: Wed, 23 Oct 2024 15:04:28 +0200
Subject: [PATCH 29/55] Improve efficiency and refact LabelTensor, codacy
 correction and fix bug in PinaBatch

---
 pina/__init__.py                             |   8 +-
 pina/collector.py                            |   8 +-
 pina/data/base_dataset.py                    | 166 ++++---
 pina/data/data_module.py                     | 148 +++---
 pina/data/pina_batch.py                      |  25 +-
 pina/data/pina_subset.py                     |  13 +-
 pina/data/sample_dataset.py                  |  21 +-
 pina/label_tensor.py                         | 462 ++++++++++---------
 pina/operators.py                            |   7 +-
 pina/problem/abstract_problem.py             |   9 +
 pina/solvers/solver.py                       |   2 +-
 pina/trainer.py                              |   2 +-
 tests/test_label_tensor/test_label_tensor.py |  12 +-
 13 files changed, 492 insertions(+), 391 deletions(-)

diff --git a/pina/__init__.py b/pina/__init__.py
index 30f35a6a5..c02e6debd 100644
--- a/pina/__init__.py
+++ b/pina/__init__.py
@@ -1,6 +1,7 @@
 __all__ = [
-    "PINN", "Trainer", "LabelTensor", "Plotter", "Condition",
-    "SamplePointDataset", "PinaDataModule", "PinaDataLoader"
+    "Trainer", "LabelTensor", "Plotter", "Condition",
+    "SamplePointDataset", "PinaDataModule", "PinaDataLoader",
+    'TorchOptimizer', 'Graph'
 ]
 
 from .meta import *
@@ -12,3 +13,6 @@
 from .data import SamplePointDataset
 from .data import PinaDataModule
 from .data import PinaDataLoader
+from .optim import TorchOptimizer
+from .optim import TorchScheduler
+from .graph import Graph
\ No newline at end of file
diff --git a/pina/collector.py b/pina/collector.py
index c48c674e8..e75b49c8b 100644
--- a/pina/collector.py
+++ b/pina/collector.py
@@ -1,6 +1,3 @@
-from sympy.strategies.branch import condition
-
-from . import LabelTensor
 from .utils import check_consistency, merge_tensors
 
 
@@ -16,6 +13,8 @@ def __init__(self, problem):
         # }
         # those variables are used for the dataloading
         self._data_collections = {name: {} for name in self.problem.conditions}
+        self.conditions_name = {i: name for i, name in
+                                enumerate(self.problem.conditions)}
 
         # variables used to check that all conditions are sampled
         self._is_conditions_ready = {
@@ -101,7 +100,8 @@ def add_points(self, new_points_dict):
         """
         Add input points to a sampled condition
 
-        :param new_points_dict: Dictonary of input points (condition_name: LabelTensor)
+        :param new_points_dict: Dictonary of input points (condition_name:
+        LabelTensor)
         :raises RuntimeError: if at least one condition is not already sampled
         """
         for k, v in new_points_dict.items():
diff --git a/pina/data/base_dataset.py b/pina/data/base_dataset.py
index d859aac00..5e27fc91e 100644
--- a/pina/data/base_dataset.py
+++ b/pina/data/base_dataset.py
@@ -1,10 +1,12 @@
 """
 Basic data module implementation
 """
-from torch.utils.data import Dataset
 import torch
+import logging
+
+from torch.utils.data import Dataset
+
 from ..label_tensor import LabelTensor
-from ..graph import Graph
 
 
 class BaseDataset(Dataset):
@@ -12,10 +14,9 @@ class BaseDataset(Dataset):
     BaseDataset class, which handle initialization and data retrieval
     :var condition_indices: List of indices
     :var device: torch.device
-    :var condition_names: dict of condition index and corresponding name
     """
 
-    def __new__(cls, problem, device):
+    def __new__(cls, problem=None, device=torch.device('cpu')):
         """
         Ensure correct definition of __slots__ before initialization
         :param AbstractProblem problem: The formulation of the problem.
@@ -30,7 +31,7 @@ def __new__(cls, problem, device):
                 'Something is wrong, __slots__ must be defined in subclasses.')
         return object.__new__(cls)
 
-    def __init__(self, problem, device):
+    def __init__(self, problem=None, device=torch.device('cpu')):
         """"
         Initialize the object based on __slots__
         :param AbstractProblem problem: The formulation of the problem.
@@ -38,79 +39,118 @@ def __init__(self, problem, device):
         dataset will be loaded.
         """
         super().__init__()
-
-        self.condition_names = {}
-        collector = problem.collector
+        self.empty = True
+        self.problem = problem
+        self.device = device
+        self.condition_indices = None
         for slot in self.__slots__:
             setattr(self, slot, [])
-        num_el_per_condition = []
-        idx = 0
-        for name, data in collector.data_collections.items():
+        self.num_el_per_condition = []
+        self.conditions_idx = []
+        if self.problem is not None:
+            self._init_from_problem(self.problem.collector.data_collections)
+        self.initialized = False
+
+    def _init_from_problem(self, collector_dict):
+        """
+        TODO
+        """
+        for name, data in collector_dict.items():
             keys = list(data.keys())
-            current_cond_num_el = None
-            if sorted(self.__slots__) == sorted(keys):
-                for slot in self.__slots__:
-                    slot_data = data[slot]
-                    if isinstance(slot_data, (LabelTensor, torch.Tensor,
-                                              Graph)):
-                        if current_cond_num_el is None:
-                            current_cond_num_el = len(slot_data)
-                        elif current_cond_num_el != len(slot_data):
-                            raise ValueError('Different number of conditions')
-                    current_list = getattr(self, slot)
-                    current_list += [data[slot]] if not (
-                        isinstance(data[slot], list)) else data[slot]
-                num_el_per_condition.append(current_cond_num_el)
-                self.condition_names[idx] = name
-                idx += 1
-        if num_el_per_condition:
+            if set(self.__slots__) == set(keys):
+                self._populate_init_list(data)
+                idx = [key for key, val in
+                       self.problem.collector.conditions_name.items() if
+                       val == name]
+                self.conditions_idx.append(idx)
+        self.initialize()
+
+    def add_points(self, data_dict, condition_idx, batching_dim=0):
+        """
+        This method filled internal lists of data points
+        :param data_dict: dictionary containing data points
+        :param condition_idx: index of the condition to which the data points
+                belong to
+        :param batching_dim: dimension of the batching
+        :raises: ValueError if the dataset has already been initialized
+        """
+        if not self.initialized:
+            self._populate_init_list(data_dict, batching_dim)
+            self.conditions_idx.append(condition_idx)
+            self.empty = False
+        else:
+            raise ValueError('Dataset already initialized')
+
+    def _populate_init_list(self, data_dict, batching_dim=0):
+        current_cond_num_el = None
+        for slot in data_dict.keys():
+            slot_data = data_dict[slot]
+            if batching_dim != 0:
+                if isinstance(slot_data, (LabelTensor, torch.Tensor)):
+                    dims = len(slot_data.size())
+                    slot_data = slot_data.permute(
+                        [batching_dim] + [dim for dim in range(dims) if
+                                          dim != batching_dim])
+            if current_cond_num_el is None:
+                current_cond_num_el = len(slot_data)
+            elif current_cond_num_el != len(slot_data):
+                raise ValueError('Different dimension in same condition')
+            current_list = getattr(self, slot)
+            current_list += [slot_data] if not (
+                isinstance(slot_data, list)) else slot_data
+        self.num_el_per_condition.append(current_cond_num_el)
+
+    def initialize(self):
+        """
+        Initialize the datasets tensors/LabelTensors/lists given the lists
+        already filled
+        """
+        logging.debug(f'Initialize dataset {self.__class__.__name__}')
+
+        if self.num_el_per_condition:
             self.condition_indices = torch.cat(
                 [
-                    torch.tensor([i] * num_el_per_condition[i],
+                    torch.tensor([i] * self.num_el_per_condition[i],
                                  dtype=torch.uint8)
-                    for i in range(len(num_el_per_condition))
+                    for i in range(len(self.num_el_per_condition))
                 ],
-                dim=0,
+                dim=0
             )
             for slot in self.__slots__:
                 current_attribute = getattr(self, slot)
                 if all(isinstance(a, LabelTensor) for a in current_attribute):
                     setattr(self, slot, LabelTensor.vstack(current_attribute))
-        else:
-            self.condition_indices = torch.tensor([], dtype=torch.uint8)
-            for slot in self.__slots__:
-                setattr(self, slot, torch.tensor([]))
-        self.device = device
+        self.initialized = True
 
     def __len__(self):
+        """
+        :return: Number of elements in the dataset
+        """
         return len(getattr(self, self.__slots__[0]))
 
-    def __getattribute__(self, item):
-        attribute = super().__getattribute__(item)
-        if isinstance(attribute,
-                      LabelTensor) and attribute.dtype == torch.float32:
-            attribute = attribute.to(device=self.device).requires_grad_()
-        return attribute
-
     def __getitem__(self, idx):
-        if isinstance(idx, str):
-            return getattr(self, idx).to(self.device)
-        if isinstance(idx, slice):
-            to_return_list = []
-            for i in self.__slots__:
-                to_return_list.append(getattr(self, i)[idx].to(self.device))
-            return to_return_list
-
-        if isinstance(idx, (tuple, list)):
-            if (len(idx) == 2 and isinstance(idx[0], str)
-                    and isinstance(idx[1], (list, slice))):
-                tensor = getattr(self, idx[0])
-                return tensor[[idx[1]]].to(self.device)
-            if all(isinstance(x, int) for x in idx):
-                to_return_list = []
-                for i in self.__slots__:
-                    to_return_list.append(
-                        getattr(self, i)[[idx]].to(self.device))
-                return to_return_list
+        """
+        :param idx:
+        :return:
+        """
+        if not isinstance(idx, (tuple, list, slice, int)):
+            raise IndexError("Invalid index")
+        tensors = []
+        for attribute in self.__slots__:
+            tensor = getattr(self, attribute)
+            if isinstance(attribute, (LabelTensor, torch.Tensor)):
+                tensors.append(tensor.__getitem__(idx))
+            elif isinstance(attribute, list):
+                if isinstance(idx, (list, tuple)):
+                    tensor = [tensor[i] for i in idx]
+                tensors.append(tensor)
+        return tensors
 
-        raise ValueError(f'Invalid index {idx}')
+    def apply_shuffle(self, indices):
+        for slot in self.__slots__:
+            if slot != 'equation':
+                attribute = getattr(self, slot)
+                if isinstance(attribute, (LabelTensor, torch.Tensor)):
+                    setattr(self, 'slot', attribute[[indices]])
+                if isinstance(attribute, list):
+                    setattr(self, 'slot', [attribute[i] for i in indices])
diff --git a/pina/data/data_module.py b/pina/data/data_module.py
index 25c7e54ed..98460ae70 100644
--- a/pina/data/data_module.py
+++ b/pina/data/data_module.py
@@ -4,7 +4,8 @@
 
 import math
 import torch
-from lightning import LightningDataModule
+import logging
+from pytorch_lightning import LightningDataModule
 from .sample_dataset import SamplePointDataset
 from .supervised_dataset import SupervisedDataset
 from .unsupervised_dataset import UnsupervisedDataset
@@ -22,8 +23,9 @@ def __init__(self,
                  problem,
                  device,
                  train_size=.7,
-                 test_size=.2,
-                 eval_size=.1,
+                 test_size=.1,
+                 val_size=.2,
+                 predict_size=0.,
                  batch_size=None,
                  shuffle=True,
                  datasets=None):
@@ -37,37 +39,64 @@ def __init__(self,
         :param batch_size: batch size used for training
         :param datasets: list of datasets objects
         """
+        logging.debug('Start initialization of Pina DataModule')
+        logging.info('Start initialization of Pina DataModule')
         super().__init__()
-        dataset_classes = [SupervisedDataset, UnsupervisedDataset,
-                           SamplePointDataset]
+        self.problem = problem
+        self.device = device
+        self.dataset_classes = [SupervisedDataset, UnsupervisedDataset,
+                                SamplePointDataset]
         if datasets is None:
-            self.datasets = [DatasetClass(problem, device) for DatasetClass in
-                             dataset_classes]
+            self.datasets = None
         else:
             self.datasets = datasets
 
         self.split_length = []
         self.split_names = []
+        self.loader_functions = {}
+        self.batch_size = batch_size
+        self.condition_names = problem.collector.conditions_name
+
         if train_size > 0:
             self.split_names.append('train')
             self.split_length.append(train_size)
+            self.loader_functions['train_dataloader'] = lambda: PinaDataLoader(
+                self.splits['train'], self.batch_size, self.condition_names)
         if test_size > 0:
             self.split_length.append(test_size)
             self.split_names.append('test')
-        if eval_size > 0:
-            self.split_length.append(eval_size)
-            self.split_names.append('eval')
-
-        self.batch_size = batch_size
-        self.condition_names = None
+            self.loader_functions['test_dataloader'] = lambda: PinaDataLoader(
+                self.splits['test'], self.batch_size, self.condition_names)
+        if val_size > 0:
+            self.split_length.append(val_size)
+            self.split_names.append('val')
+            self.loader_functions['val_dataloader'] = lambda: PinaDataLoader(
+                self.splits['val'], self.batch_size,
+                self.condition_names)
+        if predict_size > 0:
+            self.split_length.append(predict_size)
+            self.split_names.append('predict')
+            self.loader_functions[
+                'predict_dataloader'] = lambda: PinaDataLoader(
+                self.splits['predict'], self.batch_size,
+                self.condition_names)
         self.splits = {k: {} for k in self.split_names}
         self.shuffle = shuffle
 
+        for k, v in self.loader_functions.items():
+            setattr(self, k, v)
+
+    def prepare_data(self):
+        if self.datasets is None:
+            self._create_datasets()
+
     def setup(self, stage=None):
         """
         Perform the splitting of the dataset
         """
-        self.extract_conditions()
+        logging.debug('Start setup of Pina DataModule obj')
+        if self.datasets is None:
+            self._create_datasets()
         if stage == 'fit' or stage is None:
             for dataset in self.datasets:
                 if len(dataset) > 0:
@@ -82,53 +111,6 @@ def setup(self, stage=None):
         else:
             raise ValueError("stage must be either 'fit' or 'test'")
 
-    def extract_conditions(self):
-        """
-        Extract conditions from dataset and update condition indices
-        """
-        # Extract number of conditions
-        n_conditions = 0
-        for dataset in self.datasets:
-            if n_conditions != 0:
-                dataset.condition_names = {
-                    key + n_conditions: value
-                    for key, value in dataset.condition_names.items()
-                }
-            n_conditions += len(dataset.condition_names)
-
-        self.condition_names = {
-            key: value
-            for dataset in self.datasets
-            for key, value in dataset.condition_names.items()
-        }
-
-    def train_dataloader(self):
-        """
-        Return the training dataloader for the dataset
-        :return: data loader
-        :rtype: PinaDataLoader
-        """
-        return PinaDataLoader(self.splits['train'], self.batch_size,
-                              self.condition_names)
-
-    def test_dataloader(self):
-        """
-        Return the testing dataloader for the dataset
-        :return: data loader
-        :rtype: PinaDataLoader
-        """
-        return PinaDataLoader(self.splits['test'], self.batch_size,
-                              self.condition_names)
-
-    def eval_dataloader(self):
-        """
-        Return the evaluation dataloader for the dataset
-        :return: data loader
-        :rtype: PinaDataLoader
-        """
-        return PinaDataLoader(self.splits['eval'], self.batch_size,
-                              self.condition_names)
-
     @staticmethod
     def dataset_split(dataset, lengths, seed=None, shuffle=True):
         """
@@ -141,30 +123,28 @@ def dataset_split(dataset, lengths, seed=None, shuffle=True):
         :rtype: PinaSubset
         """
         if sum(lengths) - 1 < 1e-3:
+            len_dataset = len(dataset)
             lengths = [
-                int(math.floor(len(dataset) * length)) for length in lengths
+                int(math.floor(len_dataset * length)) for length in lengths
             ]
-
             remainder = len(dataset) - sum(lengths)
             for i in range(remainder):
                 lengths[i % len(lengths)] += 1
         elif sum(lengths) - 1 >= 1e-3:
             raise ValueError(f"Sum of lengths is {sum(lengths)} less than 1")
 
-        if sum(lengths) != len(dataset):
-            raise ValueError("Sum of lengths is not equal to dataset length")
-
         if shuffle:
             if seed is not None:
                 generator = torch.Generator()
                 generator.manual_seed(seed)
                 indices = torch.randperm(sum(lengths),
-                                         generator=generator).tolist()
+                                         generator=generator)
             else:
-                indices = torch.arange(sum(lengths)).tolist()
-        else:
-            indices = torch.arange(0, sum(lengths), 1,
-                                   dtype=torch.uint8).tolist()
+                indices = torch.randperm(sum(lengths))
+            dataset.apply_shuffle(indices)
+
+        indices = torch.arange(0, sum(lengths), 1,
+                               dtype=torch.uint8).tolist()
         offsets = [
             sum(lengths[:i]) if i > 0 else 0 for i in range(len(lengths))
         ]
@@ -172,3 +152,29 @@ def dataset_split(dataset, lengths, seed=None, shuffle=True):
             PinaSubset(dataset, indices[offset:offset + length])
             for offset, length in zip(offsets, lengths)
         ]
+
+    def _create_datasets(self):
+        """
+        Create the dataset objects putting data 
+        """
+        logging.debug('Dataset creation in PinaDataModule obj')
+        collector = self.problem.collector
+        batching_dim = self.problem.batching_dimension
+        datasets_slots = [i.__slots__ for i in self.dataset_classes]
+        self.datasets = [dataset(device=self.device) for dataset in
+                         self.dataset_classes]
+        logging.debug('Filling datasets in PinaDataModule obj')
+        for name, data in collector.data_collections.items():
+            keys = list(data.keys())
+            idx = [key for key, val in collector.conditions_name.items() if
+                   val == name]
+            for i, slot in enumerate(datasets_slots):
+                if slot == keys:
+                    self.datasets[i].add_points(data, idx[0], batching_dim)
+                    continue
+        datasets = []
+        for dataset in self.datasets:
+            if not dataset.empty:
+                dataset.initialize()
+                datasets.append(dataset)
+        self.datasets = datasets
diff --git a/pina/data/pina_batch.py b/pina/data/pina_batch.py
index 65b5ac5ba..6fb74f922 100644
--- a/pina/data/pina_batch.py
+++ b/pina/data/pina_batch.py
@@ -10,13 +10,15 @@ class Batch:
     optimization.
     """
 
-    def __init__(self, dataset_dict, idx_dict):
-
+    def __init__(self, dataset_dict, idx_dict, require_grad=True):
+        self.attributes = []
         for k, v in dataset_dict.items():
             setattr(self, k, v)
+            self.attributes.append(k)
 
         for k, v in idx_dict.items():
             setattr(self, k + '_idx', v)
+        self.require_grad = require_grad
 
     def __len__(self):
         """
@@ -31,9 +33,18 @@ def __len__(self):
                 length += len(getattr(self, dataset))
         return length
 
+    def __getattribute__(self, item):
+        if item in super().__getattribute__('attributes'):
+            dataset = super().__getattribute__(item)
+            index = super().__getattribute__(item + '_idx')
+            return PinaSubset(
+                dataset.dataset,
+                dataset.indices[index])
+        else:
+            return super().__getattribute__(item)
+
     def __getattr__(self, item):
-        if not item in dir(self):
-            raise AttributeError(f'Batch instance has no attribute {item}')
-        return PinaSubset(
-            getattr(self, item).dataset,
-            getattr(self, item).indices[self.coordinates_dict[item]])
+        if item == 'data' and len(self.attributes) == 1:
+            item = self.attributes[0]
+            return super().__getattribute__(item)
+        raise AttributeError(f"'Batch' object has no attribute '{item}'")
\ No newline at end of file
diff --git a/pina/data/pina_subset.py b/pina/data/pina_subset.py
index f1347b6c5..275541e97 100644
--- a/pina/data/pina_subset.py
+++ b/pina/data/pina_subset.py
@@ -2,21 +2,22 @@
 Module for PinaSubset class
 """
 from pina import LabelTensor
-from torch import Tensor
+from torch import Tensor, float32
 
 
 class PinaSubset:
     """
     TODO
     """
-    __slots__ = ['dataset', 'indices']
+    __slots__ = ['dataset', 'indices', 'require_grad']
 
-    def __init__(self, dataset, indices):
+    def __init__(self, dataset, indices, require_grad=True):
         """
         TODO
         """
         self.dataset = dataset
         self.indices = indices
+        self.require_grad = require_grad
 
     def __len__(self):
         """
@@ -27,7 +28,9 @@ def __len__(self):
     def __getattr__(self, name):
         tensor = self.dataset.__getattribute__(name)
         if isinstance(tensor, (LabelTensor, Tensor)):
-            return tensor[self.indices]
+            tensor = tensor[[self.indices]].to(self.dataset.device)
+            return tensor.requires_grad_(
+                self.require_grad) if tensor.dtype == float32 else tensor
         if isinstance(tensor, list):
             return [tensor[i] for i in self.indices]
-        raise AttributeError("No attribute named {}".format(name))
+        raise AttributeError(f"No attribute named {name}")
diff --git a/pina/data/sample_dataset.py b/pina/data/sample_dataset.py
index 99811cac8..5c47a14e9 100644
--- a/pina/data/sample_dataset.py
+++ b/pina/data/sample_dataset.py
@@ -1,8 +1,9 @@
 """
 Sample dataset module
 """
+from copy import deepcopy
 from .base_dataset import BaseDataset
-from ..condition.input_equation_condition import InputPointsEquationCondition
+from ..condition import InputPointsEquationCondition
 
 
 class SamplePointDataset(BaseDataset):
@@ -12,3 +13,21 @@ class SamplePointDataset(BaseDataset):
     """
     data_type = 'physics'
     __slots__ = InputPointsEquationCondition.__slots__
+
+    def add_points(self, data_dict, condition_idx, batching_dim=0):
+        data_dict = deepcopy(data_dict)
+        data_dict.pop('equation')
+        super().add_points(data_dict, condition_idx)
+
+    def _init_from_problem(self, collector_dict, batching_dim=0):
+        for name, data in collector_dict.items():
+            keys = list(data.keys())
+            if set(self.__slots__) == set(keys):
+                data = deepcopy(data)
+                data.pop('equation')
+                self._populate_init_list(data)
+                idx = [key for key, val in
+                       self.problem.collector.conditions_name.items() if
+                       val == name]
+                self.conditions_idx.append(idx)
+        self.initialize()
diff --git a/pina/label_tensor.py b/pina/label_tensor.py
index 87def2f8e..a28a3eaf3 100644
--- a/pina/label_tensor.py
+++ b/pina/label_tensor.py
@@ -1,5 +1,5 @@
 """ Module for LabelTensor """
-from copy import deepcopy, copy
+from copy import copy
 import torch
 from torch import Tensor
 
@@ -8,21 +8,29 @@ def issubset(a, b):
     """
     Check if a is a subset of b.
     """
-    return set(a).issubset(set(b))
+    if isinstance(a, list) and isinstance(b, list):
+        return set(a).issubset(set(b))
+    elif isinstance(a, range) and isinstance(b, range):
+        return a.start <= b.start and a.stop >= b.stop
+    else:
+        return False
 
 
 class LabelTensor(torch.Tensor):
     """Torch tensor with a label for any column."""
 
     @staticmethod
-    def __new__(cls, x, labels, *args, **kwargs):
-        return super().__new__(cls, x, *args, **kwargs)
+    def __new__(cls, x, labels, full=True, *args, **kwargs):
+        if isinstance(x, LabelTensor):
+            return x
+        else:
+            return super().__new__(cls, x, *args, **kwargs)
 
     @property
     def tensor(self):
         return self.as_subclass(Tensor)
 
-    def __init__(self, x, labels):
+    def __init__(self, x, labels, full=False):
         """
         Construct a `LabelTensor` by passing a dict of the labels
 
@@ -34,8 +42,17 @@ def __init__(self, x, labels):
 
         """
         self.dim_names = None
+        self.full = full
         self.labels = labels
 
+    @classmethod
+    def __internal_init__(cls, x, labels, dim_names ,full=False, *args, **kwargs):
+        lt = cls.__new__(cls, x, labels, full, *args, **kwargs)
+        lt._labels = labels
+        lt.full = full
+        lt.dim_names = dim_names
+        return lt
+
     @property
     def labels(self):
         """Property decorator for labels
@@ -43,12 +60,29 @@ def labels(self):
         :return: labels of self
         :rtype: list
         """
-        return self._labels[self.tensor.ndim - 1]['dof']
+        if self.ndim - 1 in self._labels.keys():
+            return self._labels[self.ndim - 1]['dof']
 
     @property
     def full_labels(self):
         """Property decorator for labels
 
+        :return: labels of self
+        :rtype: list
+        """
+        to_return_dict = {}
+        shape_tensor = self.shape
+        for i in range(len(shape_tensor)):
+            if i in self._labels.keys():
+                to_return_dict[i] = self._labels[i]
+            else:
+                to_return_dict[i] = {'dof': range(shape_tensor[i]), 'name': i}
+        return to_return_dict
+
+    @property
+    def stored_labels(self):
+        """Property decorator for labels
+
         :return: labels of self
         :rtype: list
         """
@@ -62,26 +96,77 @@ def labels(self, labels):
         :param labels: Labels to assign to the class variable _labels.
         :type: labels: str | list(str) | dict
         """
-        if hasattr(self, 'labels') is False:
-            self.init_labels()
+        if not hasattr(self, '_labels'):
+            self._labels = {}
         if isinstance(labels, dict):
-            self.update_labels_from_dict(labels)
+            self._init_labels_from_dict(labels)
         elif isinstance(labels, list):
-            self.update_labels_from_list(labels)
+            self._init_labels_from_list(labels)
         elif isinstance(labels, str):
             labels = [labels]
-            self.update_labels_from_list(labels)
+            self._init_labels_from_list(labels)
         else:
             raise ValueError("labels must be list, dict or string.")
         self.set_names()
 
+    def _init_labels_from_dict(self, labels):
+        """
+            Update the internal label representation according to the values
+            passed as input.
+
+            :param labels: The label(s) to update.
+            :type labels: dict
+            :raises ValueError: dof list contain duplicates or number of dof
+            does not match with tensor shape
+            """
+        tensor_shape = self.shape
+
+        if hasattr(self, 'full') and self.full:
+            labels = {i: labels[i] if i in labels else {'name': i} for i in
+                      labels.keys()}
+        for k, v in labels.items():
+            # Init labels from str
+            if isinstance(v, str):
+                v = {'name': v, 'dof': range(tensor_shape[k])}
+            # Init labels from dict
+            elif isinstance(v, dict) and list(v.keys()) == ['name']:
+                # Init from dict with only name key
+                v['dof'] = range(tensor_shape[k])
+                # Init from dict with both name and dof keys
+            elif isinstance(v, dict) and sorted(list(v.keys())) == ['dof',
+                                                                    'name']:
+                dof_list = v['dof']
+                dof_len = len(dof_list)
+                if dof_len != len(set(dof_list)):
+                    raise ValueError("dof must be unique")
+                if dof_len != tensor_shape[k]:
+                    raise ValueError(
+                        'Number of dof does not match tensor shape')
+            else:
+                ValueError('Illegal labels initialization')
+            # Perform update
+            self._labels[k] = v
+
+    def _init_labels_from_list(self, labels):
+        """
+        Given a list of dof, this method update the internal label
+        representation
+
+        :param labels: The label(s) to update.
+        :type labels: list
+        """
+        # Create a dict with labels
+        last_dim_labels = {
+            self.ndim - 1: {'dof': labels, 'name': self.ndim - 1}}
+        self._init_labels_from_dict(last_dim_labels)
+
     def set_names(self):
-        labels = self.full_labels
+        labels = self.stored_labels
         self.dim_names = {}
-        for dim in range(self.tensor.ndim):
+        for dim in labels.keys():
             self.dim_names[labels[dim]['name']] = dim
 
-    def extract(self, label_to_extract):
+    def extract(self, labels_to_extract):
         """
         Extract the subset of the original tensor by returning all the columns
         corresponding to the passed ``label_to_extract``.
@@ -91,78 +176,68 @@ def extract(self, label_to_extract):
         :raises TypeError: Labels are not ``str``.
         :raises ValueError: Label to extract is not in the labels ``list``.
         """
-        if isinstance(label_to_extract, (str, int)):
-            label_to_extract = [label_to_extract]
-        if isinstance(label_to_extract, (tuple, list)):
-            return self._extract_from_list(label_to_extract)
-        if isinstance(label_to_extract, dict):
-            return self._extract_from_dict(label_to_extract)
-        raise ValueError('labels_to_extract must be str or list or dict')
-
-    def _extract_from_list(self, labels_to_extract):
-        # Store locally all necessary obj/variables
-        ndim = self.tensor.ndim
-        labels = self.full_labels
-        tensor = self.tensor
-        last_dim_label = self.labels
+        # Convert str/int to string
+        if isinstance(labels_to_extract, (str, int)):
+            labels_to_extract = [labels_to_extract]
+
+        # Store useful variables
+        labels = self.stored_labels
+        stored_keys = labels.keys()
+        dim_names = self.dim_names
+        ndim = len(super().shape)
+
+        # Convert tuple/list to dict
+        if isinstance(labels_to_extract, (tuple, list)):
+            if not ndim - 1 in stored_keys:
+                raise ValueError(
+                    "LabelTensor does not have labels in last dimension")
+            name = labels[max(stored_keys)]['name']
+            labels_to_extract = {name: list(labels_to_extract)}
+
+        # If labels_to_extract is not dict then rise error
+        if not isinstance(labels_to_extract, dict):
+            raise ValueError('labels_to_extract must be str or list or dict')
 
-        # Verify if all the labels in labels_to_extract are in last dimension
-        if set(labels_to_extract).issubset(last_dim_label) is False:
-            raise ValueError(
-                'Cannot extract a dof which is not in the original LabelTensor')
+        # Make copy of labels (avoid issue in consistency)
+        updated_labels = {k: copy(v) for k, v in labels.items()}
 
-        # Extract index to extract
-        idx_to_extract = [last_dim_label.index(i) for i in labels_to_extract]
+        # Initialize list used to perform extraction
+        extractor = [slice(None) for _ in range(ndim)]
 
-        # Perform extraction
-        new_tensor = tensor[..., idx_to_extract]
+        # Loop over labels_to_extract dict
+        for k, v in labels_to_extract.items():
 
-        # Manage labels
-        new_labels = copy(labels)
+            # If label is not find raise value error
+            idx_dim = dim_names.get(k)
+            if idx_dim is None:
+                raise ValueError(
+                    'Cannot extract label with is not in original labels')
+
+            dim_labels = labels[idx_dim]['dof']
+            v = [v] if isinstance(v, (int, str)) else v
+
+            if not isinstance(v, range):
+                extractor[idx_dim] = [dim_labels.index(i) for i in v] if len(
+                    v) > 1 else slice(dim_labels.index(v[0]),
+                                      dim_labels.index(v[0]) + 1)
+            else:
+                extractor[idx_dim] = slice(v.start, v.stop)
 
-        last_dim_new_label = {ndim - 1: {
-            'dof': list(labels_to_extract),
-            'name': labels[ndim - 1]['name']
-        }}
-        new_labels.update(last_dim_new_label)
-        return LabelTensor(new_tensor, new_labels)
+            updated_labels.update({idx_dim: {'dof': v, 'name': k}})
 
-    def _extract_from_dict(self, labels_to_extract):
-        labels = self.full_labels
         tensor = self.tensor
-        ndim = tensor.ndim
-        new_labels = deepcopy(labels)
-        new_tensor = tensor
-        for k, _ in labels_to_extract.items():
-            idx_dim = self.dim_names[k]
-            dim_labels = labels[idx_dim]['dof']
-            if isinstance(labels_to_extract[k], (int, str)):
-                labels_to_extract[k] = [labels_to_extract[k]]
-            if set(labels_to_extract[k]).issubset(dim_labels) is False:
-                raise ValueError(
-                    'Cannot extract a dof which is not in the original '
-                    'LabelTensor')
-            idx_to_extract = [dim_labels.index(i) for i in labels_to_extract[k]]
-            indexer = [slice(None)] * idx_dim + [idx_to_extract] + [
-                slice(None)] * (ndim - idx_dim - 1)
-            new_tensor = new_tensor[indexer]
-            dim_new_label = {idx_dim: {
-                'dof': labels_to_extract[k],
-                'name': labels[idx_dim]['name']
-            }}
-            new_labels.update(dim_new_label)
-        return LabelTensor(new_tensor, new_labels)
+        tensor = tensor[extractor]
+        return LabelTensor.__internal_init__(tensor, updated_labels, dim_names)
 
     def __str__(self):
         """
         returns a string with the representation of the class
         """
-
         s = ''
         for key, value in self._labels.items():
             s += f"{key}: {value}\n"
         s += '\n'
-        s += super().__str__()
+        s += self.tensor.__str__()
         return s
 
     @staticmethod
@@ -174,55 +249,44 @@ def cat(tensors, dim=0):
 
         :param tensors: tensors to concatenate
         :type tensors: list(LabelTensor)
-        :param dim: dimensions on which you want to perform the operation (default 0)
+        :param dim: dimensions on which you want to perform the operation
+        (default 0)
         :type dim: int
         :rtype: LabelTensor
         :raises ValueError: either number dof or dimensions names differ
         """
         if len(tensors) == 0:
             return []
-        if len(tensors) == 1:
+        if len(tensors) == 1 or isinstance(tensors, LabelTensor):
             return tensors[0]
-        new_labels_cat_dim = LabelTensor._check_validity_before_cat(tensors,
-                                                                    dim)
-
         # Perform cat on tensors
         new_tensor = torch.cat(tensors, dim=dim)
 
         # Update labels
-        labels = tensors[0].full_labels
-        labels.pop(dim)
-        new_labels_cat_dim = new_labels_cat_dim if len(
-            set(new_labels_cat_dim)) == len(new_labels_cat_dim) \
-            else range(new_tensor.shape[dim])
-        labels[dim] = {'dof': new_labels_cat_dim,
-                       'name': tensors[1].full_labels[dim]['name']}
-        return LabelTensor(new_tensor, labels)
+        labels = LabelTensor.__create_labels_cat(tensors,
+                                                 dim)
+
+        return LabelTensor.__internal_init__(new_tensor, labels, tensors[0].dim_names)
 
     @staticmethod
-    def _check_validity_before_cat(tensors, dim):
-        n_dims = tensors[0].ndim
-        new_labels_cat_dim = []
+    def __create_labels_cat(tensors, dim):
         # Check if names and dof of the labels are the same in all dimensions
         # except in dim
-        for i in range(n_dims):
-            name = tensors[0].full_labels[i]['name']
-            if i != dim:
-                dof = tensors[0].full_labels[i]['dof']
-                for tensor in tensors:
-                    dof_to_check = tensor.full_labels[i]['dof']
-                    name_to_check = tensor.full_labels[i]['name']
-                    if dof != dof_to_check or name != name_to_check:
-                        raise ValueError(
-                            'dimensions must have the same dof and name')
-            else:
-                for tensor in tensors:
-                    new_labels_cat_dim += tensor.full_labels[i]['dof']
-                    name_to_check = tensor.full_labels[i]['name']
-                    if name != name_to_check:
-                        raise ValueError(
-                            'Dimensions to concatenate must have the same name')
-        return new_labels_cat_dim
+        stored_labels = [tensor.stored_labels for tensor in tensors]
+
+        # check if:
+        # - labels dict have same keys
+        # - all labels are the same expect for dimension dim
+        if not all(all(stored_labels[i][k] == stored_labels[0][k]
+                       for i in range(len(stored_labels)))
+                   for k in stored_labels[0].keys() if k != dim):
+            raise RuntimeError('tensors must have the same shape and dof')
+
+        labels = {k: copy(v) for k, v in tensors[0].stored_labels.items()}
+        if dim in labels.keys():
+            last_dim_dof = [i for j in stored_labels for i in j[dim]['dof']]
+            labels[dim]['dof'] = last_dim_dof
+        return labels
 
     def requires_grad_(self, mode=True):
         lt = super().requires_grad_(mode)
@@ -251,52 +315,10 @@ def clone(self, *args, **kwargs):
         :return: A copy of the tensor.
         :rtype: LabelTensor
         """
-
-        out = LabelTensor(super().clone(*args, **kwargs), self._labels)
+        labels = {k: copy(v) for k, v in self._labels.items()}
+        out = LabelTensor(super().clone(*args, **kwargs), labels)
         return out
 
-    def init_labels(self):
-        self._labels = {
-            idx_: {
-                'dof': range(self.tensor.shape[idx_]),
-                'name': idx_
-            } for idx_ in range(self.tensor.ndim)
-        }
-
-    def update_labels_from_dict(self, labels):
-        """
-        Update the internal label representation according to the values passed
-        as input.
-
-        :param labels: The label(s) to update.
-        :type labels: dict
-        :raises ValueError: dof list contain duplicates or number of dof does
-        not match with tensor shape
-        """
-        tensor_shape = self.tensor.shape
-        # Check dimensionality
-        for k, v in labels.items():
-            if len(v['dof']) != len(set(v['dof'])):
-                raise ValueError("dof must be unique")
-            if len(v['dof']) != tensor_shape[k]:
-                raise ValueError(
-                    'Number of dof does not match with tensor dimension')
-        # Perform update
-        self._labels.update(labels)
-
-    def update_labels_from_list(self, labels):
-        """
-        Given a list of dof, this method update the internal label
-        representation
-
-        :param labels: The label(s) to update.
-        :type labels: list
-        """
-        # Create a dict with labels
-        last_dim_labels = {
-            self.tensor.ndim - 1: {'dof': labels, 'name': self.tensor.ndim - 1}}
-        self.update_labels_from_dict(last_dim_labels)
-
     @staticmethod
     def summation(tensors):
         if len(tensors) == 0:
@@ -304,25 +326,30 @@ def summation(tensors):
         if len(tensors) == 1:
             return tensors[0]
         # Collect all labels
-        labels = tensors[0].full_labels
+
         # Check labels of all the tensors in each dimension
-        for j in range(tensors[0].ndim):
-            for i in range(1, len(tensors)):
-                if labels[j] != tensors[i].full_labels[j]:
-                    labels.pop(j)
-                    break
-        # Sum tensors
+        if not all(tensor.shape == tensors[0].shape for tensor in tensors) or \
+                not all(tensor.full_labels[i] == tensors[0].full_labels[i] for
+                        tensor in tensors for i in range(tensors[0].ndim - 1)):
+            raise RuntimeError('Tensors must have the same shape and labels')
+
+        last_dim_labels = []
         data = torch.zeros(tensors[0].tensor.shape)
         for tensor in tensors:
             data += tensor.tensor
-        new_tensor = LabelTensor(data, labels)
-        return new_tensor
+            last_dim_labels.append(tensor.labels)
+
+        last_dim_labels = ['+'.join(items) for items in zip(*last_dim_labels)]
+        labels = {k: copy(v) for k, v in tensors[0].stored_labels.items()}
+        labels.update({tensors[0].ndim - 1: {'dof': last_dim_labels,
+                                             'name': tensors[0].name}})
+        return LabelTensor(data, labels)
 
     def append(self, tensor, mode='std'):
         if mode == 'std':
             # Call cat on last dimension
             new_label_tensor = LabelTensor.cat([self, tensor],
-                                               dim=self.tensor.ndim - 1)
+                                               dim=self.ndim - 1)
         elif mode == 'cross':
             # Crete tensor and call cat on last dimension
             tensor1 = self
@@ -333,7 +360,7 @@ def append(self, tensor, mode='std'):
             tensor2 = LabelTensor(tensor2.repeat_interleave(n1, dim=0),
                                   labels=tensor2.labels)
             new_label_tensor = LabelTensor.cat([tensor1, tensor2],
-                                               dim=self.tensor.ndim - 1)
+                                               dim=self.ndim - 1)
         else:
             raise ValueError('mode must be either "std" or "cross"')
         return new_label_tensor
@@ -357,97 +384,76 @@ def __getitem__(self, index):
         :param index:
         :return:
         """
-
         if isinstance(index, str) or (isinstance(index, (tuple, list)) and all(
                 isinstance(a, str) for a in index)):
             return self.extract(index)
+
         selected_lt = super().__getitem__(index)
 
         if isinstance(index, (int, slice)):
-            return self._getitem_int_slice(index, selected_lt)
-
-        if len(index) == self.tensor.ndim:
-            return self._getitem_full_dim_indexing(index, selected_lt)
+            index = [index]
 
-        if isinstance(index, torch.Tensor) or (
-                isinstance(index, (tuple, list)) and all(
-            isinstance(x, int) for x in index)):
-            return self._getitem_permutation(index, selected_lt)
-        raise ValueError('Not recognized index type')
+        if index[0] == Ellipsis:
+            index = [slice(None)] * (self.ndim - 1) + [index[1]]
 
-    def _getitem_int_slice(self, index, selected_lt):
-        """
-        :param index:
-        :param selected_lt:
-        :return:
-        """
-        if selected_lt.ndim == 1:
-            selected_lt = selected_lt.reshape(1, -1)
         if hasattr(self, "labels"):
-            new_labels = deepcopy(self.full_labels)
-            to_update_dof = new_labels[0]['dof'][index]
-            to_update_dof = to_update_dof if isinstance(to_update_dof, (
-                tuple, list, range)) else [to_update_dof]
-            new_labels.update(
-                {0: {'dof': to_update_dof, 'name': new_labels[0]['name']}}
-            )
-            selected_lt.labels = new_labels
-        return selected_lt
-
-    def _getitem_full_dim_indexing(self, index, selected_lt):
-        new_labels = {}
-        old_labels = self.full_labels
-        if selected_lt.ndim == 1:
-            selected_lt = selected_lt.reshape(-1, 1)
-            new_labels = deepcopy(old_labels)
-            new_labels[1].update({'dof': old_labels[1]['dof'][index[1]],
-                                  'name': old_labels[1]['name']})
-        idx = 0
-        for j in range(selected_lt.ndim):
-            if not isinstance(index[j], int):
-                if hasattr(self, "labels"):
-                    new_labels.update(
-                        self._update_label_for_dim(old_labels, index[j], idx))
-                idx += 1
-        selected_lt.labels = new_labels
-        return selected_lt
-
-    def _getitem_permutation(self, index, selected_lt):
-        new_labels = deepcopy(self.full_labels)
-        new_labels.update(self._update_label_for_dim(self.full_labels, index,
-                                                     0))
-        selected_lt.labels = self.labels
+            labels = {k: copy(v) for k, v in self.stored_labels.items()}
+            for j, idx in enumerate(index):
+                if isinstance(idx, int):
+                    selected_lt = selected_lt.unsqueeze(j)
+                if j in labels.keys() and idx != slice(None):
+                    self._update_single_label(labels, labels, idx, j)
+            selected_lt = LabelTensor.__internal_init__(selected_lt, labels,
+                                                        self.dim_names)
         return selected_lt
 
     @staticmethod
-    def _update_label_for_dim(old_labels, index, dim):
+    def _update_single_label(old_labels, to_update_labels, index, dim):
         """
         TODO
-        :param old_labels:
-        :param index:
-        :param dim:
+        :param old_labels: labels from which retrieve data
+        :param to_update_labels: labels to update
+        :param index: index of dof to retain
+        :param dim: label index
         :return:
         """
+        old_dof = old_labels[dim]['dof']
+        if not isinstance(index, (int, slice)) and len(index) == len(
+                old_dof) and isinstance(old_dof, range):
+            return
         if isinstance(index, torch.Tensor):
-            index = index.nonzero()
+            index = index.nonzero(as_tuple=True)[
+                0] if index.dtype == torch.bool else index.tolist()
         if isinstance(index, list):
-            return {dim: {'dof': [old_labels[dim]['dof'][i] for i in index],
-                          'name': old_labels[dim]['name']}}
+            to_update_labels.update({dim: {
+                'dof': [old_dof[i] for i in index],
+                'name': old_labels[dim]['name']}})
         else:
-            return {dim: {'dof': old_labels[dim]['dof'][index],
-                          'name': old_labels[dim]['name']}}
+            to_update_labels.update({dim: {'dof': old_dof[index],
+                                           'name': old_labels[dim]['name']}})
 
     def sort_labels(self, dim=None):
-        def argsort(lst):
+        def arg_sort(lst):
             return sorted(range(len(lst)), key=lambda x: lst[x])
 
         if dim is None:
-            dim = self.tensor.ndim - 1
-        labels = self.full_labels[dim]['dof']
-        sorted_index = argsort(labels)
-        indexer = [slice(None)] * self.tensor.ndim
+            dim = self.ndim - 1
+        labels = self.stored_labels[dim]['dof']
+        sorted_index = arg_sort(labels)
+        indexer = [slice(None)] * self.ndim
         indexer[dim] = sorted_index
-        new_labels = deepcopy(self.full_labels)
-        new_labels[dim] = {'dof': sorted(labels),
-                           'name': new_labels[dim]['name']}
-        return LabelTensor(self.tensor[indexer], new_labels)
+        return self.__getitem__(indexer)
+
+    def __deepcopy__(self, memo):
+        from copy import deepcopy
+        cls = self.__class__
+        result = cls(deepcopy(self.tensor), deepcopy(self.stored_labels))
+        return result
+
+    def permute(self, *dims):
+        tensor = super().permute(*dims)
+        stored_labels = self.stored_labels
+        keys_list = list(*dims)
+        labels = {keys_list.index(k): copy(stored_labels[k]) for k in
+                  stored_labels.keys()}
+        return LabelTensor.__internal_init__(tensor, labels, self.dim_names)
diff --git a/pina/operators.py b/pina/operators.py
index 083837c11..48af1da1a 100644
--- a/pina/operators.py
+++ b/pina/operators.py
@@ -85,7 +85,8 @@ def grad_scalar_output(output_, input_, d):
             raise RuntimeError
         gradients = grad_scalar_output(output_, input_, d)
 
-    elif output_.shape[output_.ndim - 1] >= 2:  # vector output ##############################
+    elif output_.shape[
+        output_.ndim - 1] >= 2:  # vector output ##############################
         tensor_to_cat = []
         for i, c in enumerate(components):
             c_output = output_.extract([c])
@@ -143,7 +144,6 @@ def div(output_, input_, components=None, d=None):
         tensors_to_sum.append(grad_output.extract(c_fields))
         labels[i] = c_fields
     div_result = LabelTensor.summation(tensors_to_sum)
-    div_result.labels = ["+".join(labels)]
     return div_result
 
 
@@ -249,7 +249,8 @@ def scalar_laplace(output_, input_, components, d):
                 result[:, idx] = grad(grad_output, input_, d=di).flatten()
                 to_append_tensors[idx] = grad(grad_output, input_, d=di)
                 labels[idx] = f"dd{ci[0]}dd{di[0]}"
-            result = LabelTensor.cat(tensors=to_append_tensors, dim=output_.tensor.ndim - 1)
+            result = LabelTensor.cat(tensors=to_append_tensors,
+                                     dim=output_.tensor.ndim - 1)
             result.labels = labels
     return result
 
diff --git a/pina/problem/abstract_problem.py b/pina/problem/abstract_problem.py
index 600a68896..f37557a7d 100644
--- a/pina/problem/abstract_problem.py
+++ b/pina/problem/abstract_problem.py
@@ -32,11 +32,20 @@ def __init__(self):
         # training all type self.collector.full, which returns true if all
         # points are ready.
         self.collector.store_fixed_data()
+        self._batching_dimension = 0
 
     @property
     def collector(self):
         return self._collector
 
+    @property
+    def batching_dimension(self):
+        return self._batching_dimension
+
+    @batching_dimension.setter
+    def batching_dimension(self, value):
+        self._batching_dimension = value
+
     # TODO this should be erase when dataloading will interface collector,
     # kept only for back compatibility
     @property
diff --git a/pina/solvers/solver.py b/pina/solvers/solver.py
index 6f55dedf0..2d9b4a57e 100644
--- a/pina/solvers/solver.py
+++ b/pina/solvers/solver.py
@@ -94,7 +94,7 @@ def forward(self, *args, **kwargs):
         pass
 
     @abstractmethod
-    def training_step(self):
+    def training_step(self, batch, batch_idx):
         pass
 
     @abstractmethod
diff --git a/pina/trainer.py b/pina/trainer.py
index 3de0d7e80..1601d771b 100644
--- a/pina/trainer.py
+++ b/pina/trainer.py
@@ -79,7 +79,7 @@ def _create_loader(self):
         data_module = PinaDataModule(problem=self.solver.problem, device=device,
                                      train_size=self.train_size,
                                      test_size=self.test_size,
-                                     eval_size=self.eval_size)
+                                     val_size=self.eval_size)
         data_module.setup()
         self._loader = data_module.train_dataloader()
 
diff --git a/tests/test_label_tensor/test_label_tensor.py b/tests/test_label_tensor/test_label_tensor.py
index 846976730..61e479951 100644
--- a/tests/test_label_tensor/test_label_tensor.py
+++ b/tests/test_label_tensor/test_label_tensor.py
@@ -131,17 +131,17 @@ def test_concatenation_3D():
     data_2 = torch.rand(20, 3, 4)
     labels_2 = ['x', 'y', 'z', 'w']
     lt2 = LabelTensor(data_2, labels_2)
-    with pytest.raises(ValueError):
+    with pytest.raises(RuntimeError):
         LabelTensor.cat([lt1, lt2], dim=2)
     data_1 = torch.rand(20, 3, 2)
     labels_1 = ['x', 'y']
     lt1 = LabelTensor(data_1, labels_1)
     data_2 = torch.rand(20, 3, 3)
-    labels_2 = ['x', 'w', 'a']
+    labels_2 = ['z', 'w', 'a']
     lt2 = LabelTensor(data_2, labels_2)
     lt_cat = LabelTensor.cat([lt1, lt2], dim=2)
     assert lt_cat.shape == (20, 3, 5)
-    assert lt_cat.full_labels[2]['dof'] == range(5)
+    assert lt_cat.full_labels[2]['dof'] == ['x', 'y', 'z', 'w', 'a']
     assert lt_cat.full_labels[0]['dof'] == range(20)
     assert lt_cat.full_labels[1]['dof'] == range(3)
 
@@ -157,7 +157,8 @@ def test_summation():
     assert lt_sum.ndim == lt_sum.ndim
     assert lt_sum.shape[0] == 20
     assert lt_sum.shape[1] == 3
-    assert lt_sum.full_labels == labels_all
+    assert lt_sum.full_labels[0] == labels_all[0]
+    assert lt_sum.labels == ['x+x', 'y+y', 'z+z']
     assert torch.eq(lt_sum.tensor, torch.ones(20, 3) * 2).all()
     lt1 = LabelTensor(torch.ones(20, 3), labels_all)
     lt2 = LabelTensor(torch.ones(20, 3), labels_all)
@@ -166,7 +167,8 @@ def test_summation():
     assert lt_sum.ndim == lt_sum.ndim
     assert lt_sum.shape[0] == 20
     assert lt_sum.shape[1] == 3
-    assert lt_sum.full_labels == labels_all
+    assert lt_sum.full_labels[0] == labels_all[0]
+    assert lt_sum.labels == ['x+x+x', 'y+y+y', 'z+z+z']
     assert torch.eq(lt_sum.tensor, torch.ones(20, 3) * 2).all()
 
 

From de3376388e0f44c198c0086c34ab82ac1d6f0ef3 Mon Sep 17 00:00:00 2001
From: FilippoOlivo <filippo@filippoolivo.com>
Date: Thu, 31 Oct 2024 09:50:19 +0100
Subject: [PATCH 30/55] Codacy correction

---
 pina/__init__.py                            |   7 +-
 pina/collector.py                           |  32 +++---
 pina/condition/condition.py                 |  13 +--
 pina/condition/condition_interface.py       |  15 ++-
 pina/condition/data_condition.py            |   3 +-
 pina/condition/domain_equation_condition.py |   3 +-
 pina/condition/input_equation_condition.py  |   7 +-
 pina/condition/input_output_condition.py    |   2 +-
 pina/data/base_dataset.py                   |  31 +++---
 pina/data/data_module.py                    |  35 +++---
 pina/data/pina_batch.py                     |   9 +-
 pina/data/sample_dataset.py                 |  10 +-
 pina/domain/cartesian.py                    |  10 +-
 pina/domain/domain_interface.py             |   3 +-
 pina/label_tensor.py                        | 112 +++++++++++++-------
 pina/operators.py                           |  17 ++-
 pina/problem/abstract_problem.py            |  18 ++--
 pina/solvers/solver.py                      |   3 +-
 pina/trainer.py                             |  34 +++---
 pina/utils.py                               |  11 +-
 tests/test_condition.py                     |   3 +
 tests/test_dataset.py                       |  81 +++++++-------
 tests/test_geometry/test_cartesian.py       |   1 +
 23 files changed, 246 insertions(+), 214 deletions(-)

diff --git a/pina/__init__.py b/pina/__init__.py
index c02e6debd..3bc28ae6c 100644
--- a/pina/__init__.py
+++ b/pina/__init__.py
@@ -1,7 +1,6 @@
 __all__ = [
-    "Trainer", "LabelTensor", "Plotter", "Condition",
-    "SamplePointDataset", "PinaDataModule", "PinaDataLoader",
-    'TorchOptimizer', 'Graph'
+    "Trainer", "LabelTensor", "Plotter", "Condition", "SamplePointDataset",
+    "PinaDataModule", "PinaDataLoader", 'TorchOptimizer', 'Graph'
 ]
 
 from .meta import *
@@ -15,4 +14,4 @@
 from .data import PinaDataLoader
 from .optim import TorchOptimizer
 from .optim import TorchScheduler
-from .graph import Graph
\ No newline at end of file
+from .graph import Graph
diff --git a/pina/collector.py b/pina/collector.py
index e75b49c8b..3219b2b6a 100644
--- a/pina/collector.py
+++ b/pina/collector.py
@@ -2,23 +2,28 @@
 
 
 class Collector:
+
     def __init__(self, problem):
         # creating a hook between collector and problem
         self.problem = problem
 
         # this variable is used to store the data in the form:
-        # {'[condition_name]' : 
-        #           {'input_points' : Tensor, 
+        # {'[condition_name]' :
+        #           {'input_points' : Tensor,
         #            '[equation/output_points/conditional_variables]': Tensor}
         # }
         # those variables are used for the dataloading
         self._data_collections = {name: {} for name in self.problem.conditions}
-        self.conditions_name = {i: name for i, name in
-                                enumerate(self.problem.conditions)}
+        self.conditions_name = {
+            i: name
+            for i, name in enumerate(self.problem.conditions)
+        }
 
         # variables used to check that all conditions are sampled
         self._is_conditions_ready = {
-            name: False for name in self.problem.conditions}
+            name: False
+            for name in self.problem.conditions
+        }
         self.full = False
 
     @property
@@ -47,8 +52,8 @@ def store_fixed_data(self):
         for condition_name, condition in self.problem.conditions.items():
             # if the condition is not ready and domain is not attribute
             # of condition, we get and store the data
-            if (not self._is_conditions_ready[condition_name]) and (
-                    not hasattr(condition, "domain")):
+            if (not self._is_conditions_ready[condition_name]) and (not hasattr(
+                    condition, "domain")):
                 # get data
                 keys = condition.__slots__
                 values = [getattr(condition, name) for name in keys]
@@ -70,7 +75,8 @@ def store_sample_domains(self, n, mode, variables, sample_locations):
                 # if we have sampled the condition but not all variables
                 else:
                     already_sampled = [
-                        self.data_collections[loc]['input_points']]
+                        self.data_collections[loc]['input_points']
+                    ]
             # if the condition is ready but we want to sample again
             else:
                 self._is_conditions_ready[loc] = False
@@ -78,14 +84,10 @@ def store_sample_domains(self, n, mode, variables, sample_locations):
 
             # get the samples
             samples = [
-                          condition.domain.sample(n=n, mode=mode,
-                                                  variables=variables)
-                      ] + already_sampled
+                condition.domain.sample(n=n, mode=mode, variables=variables)
+            ] + already_sampled
             pts = merge_tensors(samples)
-            if (
-                    set(pts.labels).issubset(
-                        sorted(self.problem.input_variables))
-            ):
+            if (set(pts.labels).issubset(sorted(self.problem.input_variables))):
                 pts = pts.sort_labels()
                 if sorted(pts.labels) == sorted(self.problem.input_variables):
                     self._is_conditions_ready[loc] = True
diff --git a/pina/condition/condition.py b/pina/condition/condition.py
index 01965fe0d..3a62143e4 100644
--- a/pina/condition/condition.py
+++ b/pina/condition/condition.py
@@ -39,21 +39,16 @@ class Condition:
     """
 
     __slots__ = list(
-        set(
-            InputOutputPointsCondition.__slots__ +
+        set(InputOutputPointsCondition.__slots__ +
             InputPointsEquationCondition.__slots__ +
             DomainEquationCondition.__slots__ +
-            DataConditionInterface.__slots__
-        )
-    )
+            DataConditionInterface.__slots__))
 
     def __new__(cls, *args, **kwargs):
 
         if len(args) != 0:
-            raise ValueError(
-                "Condition takes only the following keyword "
-                f"arguments: {Condition.__slots__}."
-            )
+            raise ValueError("Condition takes only the following keyword "
+                             f"arguments: {Condition.__slots__}.")
 
         sorted_keys = sorted(kwargs.keys())
         if sorted_keys == sorted(InputOutputPointsCondition.__slots__):
diff --git a/pina/condition/condition_interface.py b/pina/condition/condition_interface.py
index 808c06afe..f2fe5db97 100644
--- a/pina/condition/condition_interface.py
+++ b/pina/condition/condition_interface.py
@@ -1,6 +1,6 @@
-
 from abc import ABCMeta
 
+
 class ConditionInterface(metaclass=ABCMeta):
 
     condition_types = ['physics', 'supervised', 'unsupervised']
@@ -12,7 +12,7 @@ def __init__(self, *args, **kwargs):
     @property
     def problem(self):
         return self._problem
-    
+
     @problem.setter
     def problem(self, value):
         self._problem = value
@@ -20,15 +20,14 @@ def problem(self, value):
     @property
     def condition_type(self):
         return self._condition_type
-    
+
     @condition_type.setter
     def condition_type(self, values):
         if not isinstance(values, (list, tuple)):
             values = [values]
         for value in values:
             if value not in ConditionInterface.condition_types:
-                raise ValueError( 
-                                'Unavailable type of condition, expected one of'
-                                f' {ConditionInterface.condition_types}.'
-                                )
-        self._condition_type = values
\ No newline at end of file
+                raise ValueError(
+                    'Unavailable type of condition, expected one of'
+                    f' {ConditionInterface.condition_types}.')
+        self._condition_type = values
diff --git a/pina/condition/data_condition.py b/pina/condition/data_condition.py
index 3bcd4be6d..c6777231c 100644
--- a/pina/condition/data_condition.py
+++ b/pina/condition/data_condition.py
@@ -5,6 +5,7 @@
 from ..graph import Graph
 from ..utils import check_consistency
 
+
 class DataConditionInterface(ConditionInterface):
     """
     Condition for data. This condition must be used every
@@ -29,4 +30,4 @@ def __setattr__(self, key, value):
             check_consistency(value, (LabelTensor, Graph, torch.Tensor))
             DataConditionInterface.__dict__[key].__set__(self, value)
         elif key in ('_problem', '_condition_type'):
-            super().__setattr__(key, value)
\ No newline at end of file
+            super().__setattr__(key, value)
diff --git a/pina/condition/domain_equation_condition.py b/pina/condition/domain_equation_condition.py
index 28315655b..58dca70be 100644
--- a/pina/condition/domain_equation_condition.py
+++ b/pina/condition/domain_equation_condition.py
@@ -5,6 +5,7 @@
 from ..domain import DomainInterface
 from ..equation.equation_interface import EquationInterface
 
+
 class DomainEquationCondition(ConditionInterface):
     """
     Condition for domain/equation data. This condition must be used every
@@ -30,4 +31,4 @@ def __setattr__(self, key, value):
             check_consistency(value, (EquationInterface))
             DomainEquationCondition.__dict__[key].__set__(self, value)
         elif key in ('_problem', '_condition_type'):
-            super().__setattr__(key, value)
\ No newline at end of file
+            super().__setattr__(key, value)
diff --git a/pina/condition/input_equation_condition.py b/pina/condition/input_equation_condition.py
index 0d34dfc93..bf05130c0 100644
--- a/pina/condition/input_equation_condition.py
+++ b/pina/condition/input_equation_condition.py
@@ -6,6 +6,7 @@
 from ..utils import check_consistency
 from ..equation.equation_interface import EquationInterface
 
+
 class InputPointsEquationCondition(ConditionInterface):
     """
     Condition for input_points/equation data. This condition must be used every
@@ -25,10 +26,12 @@ def __init__(self, input_points, equation):
 
     def __setattr__(self, key, value):
         if key == 'input_points':
-            check_consistency(value, (LabelTensor)) # for now only labeltensors, we need labels for the operators!
+            check_consistency(
+                value, (LabelTensor)
+            )  # for now only labeltensors, we need labels for the operators!
             InputPointsEquationCondition.__dict__[key].__set__(self, value)
         elif key == 'equation':
             check_consistency(value, (EquationInterface))
             InputPointsEquationCondition.__dict__[key].__set__(self, value)
         elif key in ('_problem', '_condition_type'):
-            super().__setattr__(key, value)
\ No newline at end of file
+            super().__setattr__(key, value)
diff --git a/pina/condition/input_output_condition.py b/pina/condition/input_output_condition.py
index 8a17495dd..08ed21d96 100644
--- a/pina/condition/input_output_condition.py
+++ b/pina/condition/input_output_condition.py
@@ -1,4 +1,3 @@
-
 import torch
 
 from .condition_interface import ConditionInterface
@@ -6,6 +5,7 @@
 from ..graph import Graph
 from ..utils import check_consistency
 
+
 class InputOutputPointsCondition(ConditionInterface):
     """
     Condition for domain/equation data. This condition must be used every
diff --git a/pina/data/base_dataset.py b/pina/data/base_dataset.py
index 5e27fc91e..2c28ba30b 100644
--- a/pina/data/base_dataset.py
+++ b/pina/data/base_dataset.py
@@ -59,9 +59,11 @@ def _init_from_problem(self, collector_dict):
             keys = list(data.keys())
             if set(self.__slots__) == set(keys):
                 self._populate_init_list(data)
-                idx = [key for key, val in
-                       self.problem.collector.conditions_name.items() if
-                       val == name]
+                idx = [
+                    key for key, val in
+                    self.problem.collector.conditions_name.items()
+                    if val == name
+                ]
                 self.conditions_idx.append(idx)
         self.initialize()
 
@@ -89,15 +91,16 @@ def _populate_init_list(self, data_dict, batching_dim=0):
                 if isinstance(slot_data, (LabelTensor, torch.Tensor)):
                     dims = len(slot_data.size())
                     slot_data = slot_data.permute(
-                        [batching_dim] + [dim for dim in range(dims) if
-                                          dim != batching_dim])
+                        [batching_dim] +
+                        [dim for dim in range(dims) if dim != batching_dim])
             if current_cond_num_el is None:
                 current_cond_num_el = len(slot_data)
             elif current_cond_num_el != len(slot_data):
                 raise ValueError('Different dimension in same condition')
             current_list = getattr(self, slot)
-            current_list += [slot_data] if not (
-                isinstance(slot_data, list)) else slot_data
+            current_list += [
+                slot_data
+            ] if not (isinstance(slot_data, list)) else slot_data
         self.num_el_per_condition.append(current_cond_num_el)
 
     def initialize(self):
@@ -108,14 +111,12 @@ def initialize(self):
         logging.debug(f'Initialize dataset {self.__class__.__name__}')
 
         if self.num_el_per_condition:
-            self.condition_indices = torch.cat(
-                [
-                    torch.tensor([i] * self.num_el_per_condition[i],
-                                 dtype=torch.uint8)
-                    for i in range(len(self.num_el_per_condition))
-                ],
-                dim=0
-            )
+            self.condition_indices = torch.cat([
+                torch.tensor([i] * self.num_el_per_condition[i],
+                             dtype=torch.uint8)
+                for i in range(len(self.num_el_per_condition))
+            ],
+                                               dim=0)
             for slot in self.__slots__:
                 current_attribute = getattr(self, slot)
                 if all(isinstance(a, LabelTensor) for a in current_attribute):
diff --git a/pina/data/data_module.py b/pina/data/data_module.py
index 98460ae70..bd117b54b 100644
--- a/pina/data/data_module.py
+++ b/pina/data/data_module.py
@@ -44,8 +44,9 @@ def __init__(self,
         super().__init__()
         self.problem = problem
         self.device = device
-        self.dataset_classes = [SupervisedDataset, UnsupervisedDataset,
-                                SamplePointDataset]
+        self.dataset_classes = [
+            SupervisedDataset, UnsupervisedDataset, SamplePointDataset
+        ]
         if datasets is None:
             self.datasets = None
         else:
@@ -71,15 +72,12 @@ def __init__(self,
             self.split_length.append(val_size)
             self.split_names.append('val')
             self.loader_functions['val_dataloader'] = lambda: PinaDataLoader(
-                self.splits['val'], self.batch_size,
-                self.condition_names)
+                self.splits['val'], self.batch_size, self.condition_names)
         if predict_size > 0:
             self.split_length.append(predict_size)
             self.split_names.append('predict')
-            self.loader_functions[
-                'predict_dataloader'] = lambda: PinaDataLoader(
-                self.splits['predict'], self.batch_size,
-                self.condition_names)
+            self.loader_functions['predict_dataloader'] = lambda: PinaDataLoader(
+                self.splits['predict'], self.batch_size, self.condition_names)
         self.splits = {k: {} for k in self.split_names}
         self.shuffle = shuffle
 
@@ -104,8 +102,8 @@ def setup(self, stage=None):
                                                 self.split_length,
                                                 shuffle=self.shuffle)
                     for i in range(len(self.split_length)):
-                        self.splits[
-                            self.split_names[i]][dataset.data_type] = splits[i]
+                        self.splits[self.split_names[i]][
+                            dataset.data_type] = splits[i]
         elif stage == 'test':
             raise NotImplementedError("Testing pipeline not implemented yet")
         else:
@@ -137,14 +135,12 @@ def dataset_split(dataset, lengths, seed=None, shuffle=True):
             if seed is not None:
                 generator = torch.Generator()
                 generator.manual_seed(seed)
-                indices = torch.randperm(sum(lengths),
-                                         generator=generator)
+                indices = torch.randperm(sum(lengths), generator=generator)
             else:
                 indices = torch.randperm(sum(lengths))
             dataset.apply_shuffle(indices)
 
-        indices = torch.arange(0, sum(lengths), 1,
-                               dtype=torch.uint8).tolist()
+        indices = torch.arange(0, sum(lengths), 1, dtype=torch.uint8).tolist()
         offsets = [
             sum(lengths[:i]) if i > 0 else 0 for i in range(len(lengths))
         ]
@@ -161,13 +157,16 @@ def _create_datasets(self):
         collector = self.problem.collector
         batching_dim = self.problem.batching_dimension
         datasets_slots = [i.__slots__ for i in self.dataset_classes]
-        self.datasets = [dataset(device=self.device) for dataset in
-                         self.dataset_classes]
+        self.datasets = [
+            dataset(device=self.device) for dataset in self.dataset_classes
+        ]
         logging.debug('Filling datasets in PinaDataModule obj')
         for name, data in collector.data_collections.items():
             keys = list(data.keys())
-            idx = [key for key, val in collector.conditions_name.items() if
-                   val == name]
+            idx = [
+                key for key, val in collector.conditions_name.items()
+                if val == name
+            ]
             for i, slot in enumerate(datasets_slots):
                 if slot == keys:
                     self.datasets[i].add_points(data, idx[0], batching_dim)
diff --git a/pina/data/pina_batch.py b/pina/data/pina_batch.py
index 6fb74f922..c5d1b61dc 100644
--- a/pina/data/pina_batch.py
+++ b/pina/data/pina_batch.py
@@ -37,14 +37,11 @@ def __getattribute__(self, item):
         if item in super().__getattribute__('attributes'):
             dataset = super().__getattribute__(item)
             index = super().__getattribute__(item + '_idx')
-            return PinaSubset(
-                dataset.dataset,
-                dataset.indices[index])
-        else:
-            return super().__getattribute__(item)
+            return PinaSubset(dataset.dataset, dataset.indices[index])
+        return super().__getattribute__(item)
 
     def __getattr__(self, item):
         if item == 'data' and len(self.attributes) == 1:
             item = self.attributes[0]
             return super().__getattribute__(item)
-        raise AttributeError(f"'Batch' object has no attribute '{item}'")
\ No newline at end of file
+        raise AttributeError(f"'Batch' object has no attribute '{item}'")
diff --git a/pina/data/sample_dataset.py b/pina/data/sample_dataset.py
index 5c47a14e9..bc3bca335 100644
--- a/pina/data/sample_dataset.py
+++ b/pina/data/sample_dataset.py
@@ -19,15 +19,17 @@ def add_points(self, data_dict, condition_idx, batching_dim=0):
         data_dict.pop('equation')
         super().add_points(data_dict, condition_idx)
 
-    def _init_from_problem(self, collector_dict, batching_dim=0):
+    def _init_from_problem(self, collector_dict):
         for name, data in collector_dict.items():
             keys = list(data.keys())
             if set(self.__slots__) == set(keys):
                 data = deepcopy(data)
                 data.pop('equation')
                 self._populate_init_list(data)
-                idx = [key for key, val in
-                       self.problem.collector.conditions_name.items() if
-                       val == name]
+                idx = [
+                    key for key, val in
+                    self.problem.collector.conditions_name.items()
+                    if val == name
+                ]
                 self.conditions_idx.append(idx)
         self.initialize()
diff --git a/pina/domain/cartesian.py b/pina/domain/cartesian.py
index a6c369339..e2086e4d4 100644
--- a/pina/domain/cartesian.py
+++ b/pina/domain/cartesian.py
@@ -168,9 +168,8 @@ def _1d_sampler(n, mode, variables):
             for variable in variables:
                 if variable in self.fixed_.keys():
                     value = self.fixed_[variable]
-                    pts_variable = torch.tensor([[value]]).repeat(
-                        result.shape[0], 1
-                    )
+                    pts_variable = torch.tensor([[value]
+                                                 ]).repeat(result.shape[0], 1)
                     pts_variable = pts_variable.as_subclass(LabelTensor)
                     pts_variable.labels = [variable]
 
@@ -203,9 +202,8 @@ def _Nd_sampler(n, mode, variables):
             for variable in variables:
                 if variable in self.fixed_.keys():
                     value = self.fixed_[variable]
-                    pts_variable = torch.tensor([[value]]).repeat(
-                        result.shape[0], 1
-                    )
+                    pts_variable = torch.tensor([[value]
+                                                 ]).repeat(result.shape[0], 1)
                     pts_variable = pts_variable.as_subclass(LabelTensor)
                     pts_variable.labels = [variable]
 
diff --git a/pina/domain/domain_interface.py b/pina/domain/domain_interface.py
index 4fa70a2ba..916bf3e90 100644
--- a/pina/domain/domain_interface.py
+++ b/pina/domain/domain_interface.py
@@ -38,8 +38,7 @@ def sample_modes(self, values):
             if value not in DomainInterface.available_sampling_modes:
                 raise TypeError(f"mode {value} not valid. Expected at least "
                                 "one in "
-                                f"{DomainInterface.available_sampling_modes}."
-                                )
+                                f"{DomainInterface.available_sampling_modes}.")
 
     @abstractmethod
     def sample(self):
diff --git a/pina/label_tensor.py b/pina/label_tensor.py
index a28a3eaf3..719975c51 100644
--- a/pina/label_tensor.py
+++ b/pina/label_tensor.py
@@ -1,5 +1,5 @@
 """ Module for LabelTensor """
-from copy import copy
+from copy import copy, deepcopy
 import torch
 from torch import Tensor
 
@@ -10,17 +10,16 @@ def issubset(a, b):
     """
     if isinstance(a, list) and isinstance(b, list):
         return set(a).issubset(set(b))
-    elif isinstance(a, range) and isinstance(b, range):
+    if isinstance(a, range) and isinstance(b, range):
         return a.start <= b.start and a.stop >= b.stop
-    else:
-        return False
+    return False
 
 
 class LabelTensor(torch.Tensor):
     """Torch tensor with a label for any column."""
 
     @staticmethod
-    def __new__(cls, x, labels, full=True, *args, **kwargs):
+    def __new__(cls, x, labels, *args, **kwargs):
         if isinstance(x, LabelTensor):
             return x
         else:
@@ -30,7 +29,7 @@ def __new__(cls, x, labels, full=True, *args, **kwargs):
     def tensor(self):
         return self.as_subclass(Tensor)
 
-    def __init__(self, x, labels, full=False):
+    def __init__(self, x, labels, **kwargs):
         """
         Construct a `LabelTensor` by passing a dict of the labels
 
@@ -42,14 +41,19 @@ def __init__(self, x, labels, full=False):
 
         """
         self.dim_names = None
-        self.full = full
+        self.full = kwargs.get('full', True)
         self.labels = labels
 
     @classmethod
-    def __internal_init__(cls, x, labels, dim_names ,full=False, *args, **kwargs):
-        lt = cls.__new__(cls, x, labels, full, *args, **kwargs)
+    def __internal_init__(cls,
+                          x,
+                          labels,
+                          dim_names,
+                          *args,
+                          **kwargs):
+        lt = cls.__new__(cls, x, labels, *args, **kwargs)
         lt._labels = labels
-        lt.full = full
+        lt.full = kwargs.get('full', True)
         lt.dim_names = dim_names
         return lt
 
@@ -122,8 +126,12 @@ def _init_labels_from_dict(self, labels):
         tensor_shape = self.shape
 
         if hasattr(self, 'full') and self.full:
-            labels = {i: labels[i] if i in labels else {'name': i} for i in
-                      labels.keys()}
+            labels = {
+                i: labels[i] if i in labels else {
+                    'name': i
+                }
+                for i in labels.keys()
+            }
         for k, v in labels.items():
             # Init labels from str
             if isinstance(v, str):
@@ -133,8 +141,8 @@ def _init_labels_from_dict(self, labels):
                 # Init from dict with only name key
                 v['dof'] = range(tensor_shape[k])
                 # Init from dict with both name and dof keys
-            elif isinstance(v, dict) and sorted(list(v.keys())) == ['dof',
-                                                                    'name']:
+            elif isinstance(v, dict) and sorted(list(
+                    v.keys())) == ['dof', 'name']:
                 dof_list = v['dof']
                 dof_len = len(dof_list)
                 if dof_len != len(set(dof_list)):
@@ -143,7 +151,7 @@ def _init_labels_from_dict(self, labels):
                     raise ValueError(
                         'Number of dof does not match tensor shape')
             else:
-                ValueError('Illegal labels initialization')
+                raise ValueError('Illegal labels initialization')
             # Perform update
             self._labels[k] = v
 
@@ -157,7 +165,11 @@ def _init_labels_from_list(self, labels):
         """
         # Create a dict with labels
         last_dim_labels = {
-            self.ndim - 1: {'dof': labels, 'name': self.ndim - 1}}
+            self.ndim - 1: {
+                'dof': labels,
+                'name': self.ndim - 1
+            }
+        }
         self._init_labels_from_dict(last_dim_labels)
 
     def set_names(self):
@@ -217,9 +229,10 @@ def extract(self, labels_to_extract):
             v = [v] if isinstance(v, (int, str)) else v
 
             if not isinstance(v, range):
-                extractor[idx_dim] = [dim_labels.index(i) for i in v] if len(
-                    v) > 1 else slice(dim_labels.index(v[0]),
-                                      dim_labels.index(v[0]) + 1)
+                extractor[idx_dim] = [dim_labels.index(i)
+                                      for i in v] if len(v) > 1 else slice(
+                                          dim_labels.index(v[0]),
+                                          dim_labels.index(v[0]) + 1)
             else:
                 extractor[idx_dim] = slice(v.start, v.stop)
 
@@ -263,10 +276,10 @@ def cat(tensors, dim=0):
         new_tensor = torch.cat(tensors, dim=dim)
 
         # Update labels
-        labels = LabelTensor.__create_labels_cat(tensors,
-                                                 dim)
+        labels = LabelTensor.__create_labels_cat(tensors, dim)
 
-        return LabelTensor.__internal_init__(new_tensor, labels, tensors[0].dim_names)
+        return LabelTensor.__internal_init__(new_tensor, labels,
+                                             tensors[0].dim_names)
 
     @staticmethod
     def __create_labels_cat(tensors, dim):
@@ -277,9 +290,10 @@ def __create_labels_cat(tensors, dim):
         # check if:
         # - labels dict have same keys
         # - all labels are the same expect for dimension dim
-        if not all(all(stored_labels[i][k] == stored_labels[0][k]
-                       for i in range(len(stored_labels)))
-                   for k in stored_labels[0].keys() if k != dim):
+        if not all(
+                all(stored_labels[i][k] == stored_labels[0][k]
+                    for i in range(len(stored_labels)))
+                for k in stored_labels[0].keys() if k != dim):
             raise RuntimeError('tensors must have the same shape and dof')
 
         labels = {k: copy(v) for k, v in tensors[0].stored_labels.items()}
@@ -341,8 +355,12 @@ def summation(tensors):
 
         last_dim_labels = ['+'.join(items) for items in zip(*last_dim_labels)]
         labels = {k: copy(v) for k, v in tensors[0].stored_labels.items()}
-        labels.update({tensors[0].ndim - 1: {'dof': last_dim_labels,
-                                             'name': tensors[0].name}})
+        labels.update({
+            tensors[0].ndim - 1: {
+                'dof': last_dim_labels,
+                'name': tensors[0].name
+            }
+        })
         return LabelTensor(data, labels)
 
     def append(self, tensor, mode='std'):
@@ -384,8 +402,9 @@ def __getitem__(self, index):
         :param index:
         :return:
         """
-        if isinstance(index, str) or (isinstance(index, (tuple, list)) and all(
-                isinstance(a, str) for a in index)):
+        if isinstance(index,
+                      str) or (isinstance(index, (tuple, list))
+                               and all(isinstance(a, str) for a in index)):
             return self.extract(index)
 
         selected_lt = super().__getitem__(index)
@@ -418,21 +437,31 @@ def _update_single_label(old_labels, to_update_labels, index, dim):
         :return:
         """
         old_dof = old_labels[dim]['dof']
-        if not isinstance(index, (int, slice)) and len(index) == len(
-                old_dof) and isinstance(old_dof, range):
+        if not isinstance(
+                index,
+            (int, slice)) and len(index) == len(old_dof) and isinstance(
+                old_dof, range):
             return
         if isinstance(index, torch.Tensor):
-            index = index.nonzero(as_tuple=True)[
-                0] if index.dtype == torch.bool else index.tolist()
+            index = index.nonzero(
+                as_tuple=True
+            )[0] if index.dtype == torch.bool else index.tolist()
         if isinstance(index, list):
-            to_update_labels.update({dim: {
-                'dof': [old_dof[i] for i in index],
-                'name': old_labels[dim]['name']}})
+            to_update_labels.update({
+                dim: {
+                    'dof': [old_dof[i] for i in index],
+                    'name': old_labels[dim]['name']
+                }
+            })
         else:
-            to_update_labels.update({dim: {'dof': old_dof[index],
-                                           'name': old_labels[dim]['name']}})
+            to_update_labels.update(
+                {dim: {
+                    'dof': old_dof[index],
+                    'name': old_labels[dim]['name']
+                }})
 
     def sort_labels(self, dim=None):
+
         def arg_sort(lst):
             return sorted(range(len(lst)), key=lambda x: lst[x])
 
@@ -445,7 +474,6 @@ def arg_sort(lst):
         return self.__getitem__(indexer)
 
     def __deepcopy__(self, memo):
-        from copy import deepcopy
         cls = self.__class__
         result = cls(deepcopy(self.tensor), deepcopy(self.stored_labels))
         return result
@@ -454,6 +482,8 @@ def permute(self, *dims):
         tensor = super().permute(*dims)
         stored_labels = self.stored_labels
         keys_list = list(*dims)
-        labels = {keys_list.index(k): copy(stored_labels[k]) for k in
-                  stored_labels.keys()}
+        labels = {
+            keys_list.index(k): copy(stored_labels[k])
+            for k in stored_labels.keys()
+        }
         return LabelTensor.__internal_init__(tensor, labels, self.dim_names)
diff --git a/pina/operators.py b/pina/operators.py
index 48af1da1a..193243a69 100644
--- a/pina/operators.py
+++ b/pina/operators.py
@@ -56,9 +56,9 @@ def grad_scalar_output(output_, input_, d):
         gradients = torch.autograd.grad(
             output_,
             input_,
-            grad_outputs=torch.ones(
-                output_.size(), dtype=output_.dtype, device=output_.device
-            ),
+            grad_outputs=torch.ones(output_.size(),
+                                    dtype=output_.dtype,
+                                    device=output_.device),
             create_graph=True,
             retain_graph=True,
             allow_unused=True,
@@ -85,8 +85,8 @@ def grad_scalar_output(output_, input_, d):
             raise RuntimeError
         gradients = grad_scalar_output(output_, input_, d)
 
-    elif output_.shape[
-        output_.ndim - 1] >= 2:  # vector output ##############################
+    elif output_.shape[output_.ndim -
+                       1] >= 2:  # vector output ##############################
         tensor_to_cat = []
         for i, c in enumerate(components):
             c_output = output_.extract([c])
@@ -281,11 +281,8 @@ def advection(output_, input_, velocity_field, components=None, d=None):
     if components is None:
         components = output_.labels
 
-    tmp = (
-        grad(output_, input_, components, d)
-        .reshape(-1, len(components), len(d))
-        .transpose(0, 1)
-    )
+    tmp = (grad(output_, input_, components, d).reshape(-1, len(components),
+                                                        len(d)).transpose(0, 1))
 
     tmp *= output_.extract(velocity_field)
     return tmp.sum(dim=2).T
diff --git a/pina/problem/abstract_problem.py b/pina/problem/abstract_problem.py
index f37557a7d..6897fbb74 100644
--- a/pina/problem/abstract_problem.py
+++ b/pina/problem/abstract_problem.py
@@ -7,6 +7,7 @@
 from ..collector import Collector
 from copy import deepcopy
 
+
 class AbstractProblem(metaclass=ABCMeta):
     """
     The abstract `AbstractProblem` class. All the class defining a PINA Problem
@@ -55,7 +56,7 @@ def input_pts(self):
             if 'input_points' in v.keys():
                 to_return[k] = v['input_points']
         return to_return
-    
+
     def __deepcopy__(self, memo):
         """
         Implements deepcopy for the
@@ -116,9 +117,11 @@ def conditions(self):
         """
         return self._conditions
 
-    def discretise_domain(
-        self, n, mode="random", variables="all", locations="all"
-    ):
+    def discretise_domain(self,
+                          n,
+                          mode="random",
+                          variables="all",
+                          locations="all"):
         """
         Generate a set of points to span the `Location` of all the conditions of
         the problem.
@@ -170,9 +173,10 @@ def discretise_domain(
 
         # check correct location
         if locations == "all":
-            locations = [name for name in self.conditions.keys()
-                         if isinstance(self.conditions[name],
-                                       DomainEquationCondition)]
+            locations = [
+                name for name in self.conditions.keys()
+                if isinstance(self.conditions[name], DomainEquationCondition)
+            ]
         else:
             if not isinstance(locations, (list)):
                 locations = [locations]
diff --git a/pina/solvers/solver.py b/pina/solvers/solver.py
index 2d9b4a57e..e00bc8d59 100644
--- a/pina/solvers/solver.py
+++ b/pina/solvers/solver.py
@@ -142,5 +142,4 @@ def _check_solver_consistency(self, problem):
                     condition.condition_type):
                 raise ValueError(
                     f'{self.__name__} support only dose not support condition '
-                    f'{condition.condition_type}'
-                )
+                    f'{condition.condition_type}')
diff --git a/pina/trainer.py b/pina/trainer.py
index 1601d771b..58c66f67c 100644
--- a/pina/trainer.py
+++ b/pina/trainer.py
@@ -9,8 +9,13 @@
 
 class Trainer(pytorch_lightning.Trainer):
 
-    def __init__(self, solver, batch_size=None, train_size=.7, test_size=.2,
-                 eval_size=.1, **kwargs):
+    def __init__(self,
+                 solver,
+                 batch_size=None,
+                 train_size=.7,
+                 test_size=.2,
+                 eval_size=.1,
+                 **kwargs):
         """
         PINA Trainer class for costumizing every aspect of training via flags.
 
@@ -48,8 +53,7 @@ def _move_to_device(self):
         if hasattr(pb, "unknown_parameters"):
             for key in pb.unknown_parameters:
                 pb.unknown_parameters[key] = torch.nn.Parameter(
-                    pb.unknown_parameters[key].data.to(device)
-                )
+                    pb.unknown_parameters[key].data.to(device))
 
     def _create_loader(self):
         """
@@ -58,14 +62,11 @@ def _create_loader(self):
         trainer dataloader, just call the method.
         """
         if not self.solver.problem.collector.full:
-            error_message = '\n'.join(
-                [
-                    f"""{" " * 13} ---> Condition {key} {"sampled" if value else
-                    "not sampled"}"""
-                    for key, value in
-                    self._solver.problem.collector._is_conditions_ready.items()
-                ]
-            )
+            error_message = '\n'.join([
+                f"""{" " * 13} ---> Condition {key} {"sampled" if value else
+                    "not sampled"}""" for key, value in
+                self._solver.problem.collector._is_conditions_ready.items()
+            ])
             raise RuntimeError('Cannot create Trainer if not all conditions '
                                'are sampled. The Trainer got the following:\n'
                                f'{error_message}')
@@ -76,7 +77,8 @@ def _create_loader(self):
 
         device = devices[0]
 
-        data_module = PinaDataModule(problem=self.solver.problem, device=device,
+        data_module = PinaDataModule(problem=self.solver.problem,
+                                     device=device,
                                      train_size=self.train_size,
                                      test_size=self.test_size,
                                      val_size=self.eval_size)
@@ -87,9 +89,9 @@ def train(self, **kwargs):
         """
         Train the solver method.
         """
-        return super().fit(
-            self.solver, train_dataloaders=self._loader, **kwargs
-        )
+        return super().fit(self.solver,
+                           train_dataloaders=self._loader,
+                           **kwargs)
 
     @property
     def solver(self):
diff --git a/pina/utils.py b/pina/utils.py
index 282dd5332..84d3e7419 100644
--- a/pina/utils.py
+++ b/pina/utils.py
@@ -40,9 +40,9 @@ def check_consistency(object, object_instance, subclass=False):
             raise ValueError(f"{type(obj).__name__} must be {object_instance}.")
 
 
-def number_parameters(
-    model, aggregate=True, only_trainable=True
-):  # TODO: check
+def number_parameters(model,
+                      aggregate=True,
+                      only_trainable=True):  # TODO: check
     """
     Return the number of parameters of a given `model`.
 
@@ -80,9 +80,8 @@ def merge_two_tensors(tensor1, tensor2):
     n2 = tensor2.shape[0]
 
     tensor1 = LabelTensor(tensor1.repeat(n2, 1), labels=tensor1.labels)
-    tensor2 = LabelTensor(
-        tensor2.repeat_interleave(n1, dim=0), labels=tensor2.labels
-    )
+    tensor2 = LabelTensor(tensor2.repeat_interleave(n1, dim=0),
+                          labels=tensor2.labels)
     return tensor1.append(tensor2)
 
 
diff --git a/tests/test_condition.py b/tests/test_condition.py
index f12979dc8..9165c3fa1 100644
--- a/tests/test_condition.py
+++ b/tests/test_condition.py
@@ -22,8 +22,11 @@ def test_init_inputoutput():
         Condition(input_points=3., output_points='example')
     with pytest.raises(ValueError):
         Condition(input_points=example_domain, output_points=example_domain)
+
+
 test_init_inputoutput()
 
+
 def test_init_domainfunc():
     Condition(domain=example_domain, equation=FixedValue(0.0))
     with pytest.raises(ValueError):
diff --git a/tests/test_dataset.py b/tests/test_dataset.py
index 503ddd683..87fd9a15b 100644
--- a/tests/test_dataset.py
+++ b/tests/test_dataset.py
@@ -32,49 +32,49 @@ class Poisson(SpatialProblem):
 
     conditions = {
         'gamma1':
-            Condition(domain=CartesianDomain({
-                'x': [0, 1],
-                'y': 1
-            }),
-                equation=FixedValue(0.0)),
+        Condition(domain=CartesianDomain({
+            'x': [0, 1],
+            'y': 1
+        }),
+                  equation=FixedValue(0.0)),
         'gamma2':
-            Condition(domain=CartesianDomain({
-                'x': [0, 1],
-                'y': 0
-            }),
-                equation=FixedValue(0.0)),
+        Condition(domain=CartesianDomain({
+            'x': [0, 1],
+            'y': 0
+        }),
+                  equation=FixedValue(0.0)),
         'gamma3':
-            Condition(domain=CartesianDomain({
-                'x': 1,
-                'y': [0, 1]
-            }),
-                equation=FixedValue(0.0)),
+        Condition(domain=CartesianDomain({
+            'x': 1,
+            'y': [0, 1]
+        }),
+                  equation=FixedValue(0.0)),
         'gamma4':
-            Condition(domain=CartesianDomain({
-                'x': 0,
-                'y': [0, 1]
-            }),
-                equation=FixedValue(0.0)),
+        Condition(domain=CartesianDomain({
+            'x': 0,
+            'y': [0, 1]
+        }),
+                  equation=FixedValue(0.0)),
         'D':
-            Condition(input_points=LabelTensor(torch.rand(size=(100, 2)),
-                                               ['x', 'y']),
-                      equation=my_laplace),
+        Condition(input_points=LabelTensor(torch.rand(size=(100, 2)),
+                                           ['x', 'y']),
+                  equation=my_laplace),
         'data':
-            Condition(input_points=in_, output_points=out_),
+        Condition(input_points=in_, output_points=out_),
         'data2':
-            Condition(input_points=in2_, output_points=out2_),
+        Condition(input_points=in2_, output_points=out2_),
         'unsupervised':
-            Condition(
-                input_points=LabelTensor(torch.rand(size=(45, 2)), ['x', 'y']),
-                conditional_variables=LabelTensor(torch.ones(size=(45, 1)),
-                                                  ['alpha']),
-            ),
+        Condition(
+            input_points=LabelTensor(torch.rand(size=(45, 2)), ['x', 'y']),
+            conditional_variables=LabelTensor(torch.ones(size=(45, 1)),
+                                              ['alpha']),
+        ),
         'unsupervised2':
-            Condition(
-                input_points=LabelTensor(torch.rand(size=(90, 2)), ['x', 'y']),
-                conditional_variables=LabelTensor(torch.ones(size=(90, 1)),
-                                                  ['alpha']),
-            )
+        Condition(
+            input_points=LabelTensor(torch.rand(size=(90, 2)), ['x', 'y']),
+            conditional_variables=LabelTensor(torch.ones(size=(90, 1)),
+                                              ['alpha']),
+        )
     }
 
 
@@ -201,11 +201,12 @@ def test_loader():
 
 class GraphProblem(AbstractProblem):
     output = LabelTensor(torch.rand((100, 3)), labels=['ux', 'uy', 'p'])
-    input = [Graph.build('radius',
-                         nodes_coordinates=coordinates[i, :, :],
-                         nodes_data=data[i, :, :], radius=0.2)
-             for i in
-             range(100)]
+    input = [
+        Graph.build('radius',
+                    nodes_coordinates=coordinates[i, :, :],
+                    nodes_data=data[i, :, :],
+                    radius=0.2) for i in range(100)
+    ]
     output_variables = ['u']
 
     conditions = {
diff --git a/tests/test_geometry/test_cartesian.py b/tests/test_geometry/test_cartesian.py
index 65026c332..fc30757b6 100644
--- a/tests/test_geometry/test_cartesian.py
+++ b/tests/test_geometry/test_cartesian.py
@@ -3,6 +3,7 @@
 from pina import LabelTensor
 from pina.domain import CartesianDomain
 
+
 def test_constructor():
     CartesianDomain({'x': [0, 1], 'y': [0, 1]})
 

From 90b0bc98cc6e857350b43fbda75d65246eea0260 Mon Sep 17 00:00:00 2001
From: Filippo Olivo <filippo@filippoolivo.com>
Date: Wed, 27 Nov 2024 16:01:39 +0100
Subject: [PATCH 31/55] Implementation of DataLoader and DataModule (#383)

Refactoring for 0.2
* Data module, data loader and dataset
* Refactor LabelTensor
* Refactor solvers

Co-authored-by: dario-coscia <dariocos99@gmail.com>
---
 pina/__init__.py                              |   8 +-
 pina/collector.py                             |  18 +-
 pina/condition/data_condition.py              |   3 +-
 pina/condition/domain_equation_condition.py   |   5 +-
 pina/condition/input_equation_condition.py    |   3 +-
 pina/condition/input_output_condition.py      |   6 +-
 pina/data/__init__.py                         |  13 +-
 pina/data/base_dataset.py                     | 157 --------
 pina/data/data_module.py                      | 374 ++++++++++++------
 pina/data/dataset.py                          | 102 +++++
 pina/data/pina_batch.py                       |  47 ---
 pina/data/pina_dataloader.py                  |  68 ----
 pina/data/pina_subset.py                      |  36 --
 pina/data/sample_dataset.py                   |  35 --
 pina/data/supervised_dataset.py               |  13 -
 pina/data/unsupervised_dataset.py             |  14 -
 pina/graph.py                                 |   4 +-
 pina/label_tensor.py                          | 356 +++++++++--------
 pina/model/layers/lowrank_layer.py            |   1 +
 pina/model/network.py                         |  20 +-
 pina/operators.py                             |  10 +-
 pina/problem/abstract_problem.py              |   6 +-
 pina/problem/inverse_problem.py               |   3 +-
 pina/solvers/graph.py                         |  34 --
 pina/solvers/pinns/basepinn.py                | 117 ++++--
 pina/solvers/pinns/pinn.py                    |  49 ++-
 pina/solvers/solver.py                        |  22 +-
 pina/solvers/supervised.py                    |  68 ++--
 pina/trainer.py                               |  49 +--
 setup.py                                      |   3 +-
 tests/test_dataset.py                         | 227 -----------
 .../test_label_tensor/test_label_tensor_01.py |   2 +-
 tests/test_solvers/test_pinn.py               | 275 +++----------
 tests/test_solvers/test_supervised_solver.py  |   2 +-
 34 files changed, 814 insertions(+), 1336 deletions(-)
 delete mode 100644 pina/data/base_dataset.py
 create mode 100644 pina/data/dataset.py
 delete mode 100644 pina/data/pina_batch.py
 delete mode 100644 pina/data/pina_dataloader.py
 delete mode 100644 pina/data/pina_subset.py
 delete mode 100644 pina/data/sample_dataset.py
 delete mode 100644 pina/data/supervised_dataset.py
 delete mode 100644 pina/data/unsupervised_dataset.py
 delete mode 100644 pina/solvers/graph.py
 delete mode 100644 tests/test_dataset.py

diff --git a/pina/__init__.py b/pina/__init__.py
index 3bc28ae6c..e9ce70613 100644
--- a/pina/__init__.py
+++ b/pina/__init__.py
@@ -1,6 +1,6 @@
 __all__ = [
-    "Trainer", "LabelTensor", "Plotter", "Condition", "SamplePointDataset",
-    "PinaDataModule", "PinaDataLoader", 'TorchOptimizer', 'Graph'
+    "Trainer", "LabelTensor", "Plotter", "Condition",
+    "PinaDataModule", 'TorchOptimizer', 'Graph',
 ]
 
 from .meta import *
@@ -9,9 +9,9 @@
 from .trainer import Trainer
 from .plotter import Plotter
 from .condition.condition import Condition
-from .data import SamplePointDataset
+
 from .data import PinaDataModule
-from .data import PinaDataLoader
+
 from .optim import TorchOptimizer
 from .optim import TorchScheduler
 from .graph import Graph
diff --git a/pina/collector.py b/pina/collector.py
index 3219b2b6a..1f0fb41d3 100644
--- a/pina/collector.py
+++ b/pina/collector.py
@@ -1,3 +1,4 @@
+from . import LabelTensor
 from .utils import check_consistency, merge_tensors
 
 
@@ -66,9 +67,12 @@ def store_sample_domains(self, n, mode, variables, sample_locations):
         for loc in sample_locations:
             # get condition
             condition = self.problem.conditions[loc]
+            condition_domain = condition.domain
+            if isinstance(condition_domain, str):
+                condition_domain = self.problem.domains[condition_domain]
             keys = ["input_points", "equation"]
             # if the condition is not ready, we get and store the data
-            if (not self._is_conditions_ready[loc]):
+            if not self._is_conditions_ready[loc]:
                 # if it is the first time we sample
                 if not self.data_collections[loc]:
                     already_sampled = []
@@ -84,10 +88,11 @@ def store_sample_domains(self, n, mode, variables, sample_locations):
 
             # get the samples
             samples = [
-                condition.domain.sample(n=n, mode=mode, variables=variables)
-            ] + already_sampled
+                          condition_domain.sample(n=n, mode=mode,
+                                                  variables=variables)
+                      ] + already_sampled
             pts = merge_tensors(samples)
-            if (set(pts.labels).issubset(sorted(self.problem.input_variables))):
+            if set(pts.labels).issubset(sorted(self.problem.input_variables)):
                 pts = pts.sort_labels()
                 if sorted(pts.labels) == sorted(self.problem.input_variables):
                     self._is_conditions_ready[loc] = True
@@ -110,5 +115,6 @@ def add_points(self, new_points_dict):
             if not self._is_conditions_ready[k]:
                 raise RuntimeError(
                     'Cannot add points on a non sampled condition')
-            self.data_collections[k]['input_points'] = self.data_collections[k][
-                'input_points'].vstack(v)
+            self.data_collections[k]['input_points'] = LabelTensor.vstack(
+                [self.data_collections[k][
+                     'input_points'], v])
diff --git a/pina/condition/data_condition.py b/pina/condition/data_condition.py
index c6777231c..255c329f8 100644
--- a/pina/condition/data_condition.py
+++ b/pina/condition/data_condition.py
@@ -18,12 +18,11 @@ class DataConditionInterface(ConditionInterface):
 
     def __init__(self, input_points, conditional_variables=None):
         """
-        TODO
+        TODO : add docstring
         """
         super().__init__()
         self.input_points = input_points
         self.conditional_variables = conditional_variables
-        self._condition_type = 'unsupervised'
 
     def __setattr__(self, key, value):
         if (key == 'input_points') or (key == 'conditional_variables'):
diff --git a/pina/condition/domain_equation_condition.py b/pina/condition/domain_equation_condition.py
index 58dca70be..9fb0dcb2d 100644
--- a/pina/condition/domain_equation_condition.py
+++ b/pina/condition/domain_equation_condition.py
@@ -16,16 +16,15 @@ class DomainEquationCondition(ConditionInterface):
 
     def __init__(self, domain, equation):
         """
-        TODO
+        TODO : add docstring
         """
         super().__init__()
         self.domain = domain
         self.equation = equation
-        self._condition_type = 'physics'
 
     def __setattr__(self, key, value):
         if key == 'domain':
-            check_consistency(value, (DomainInterface))
+            check_consistency(value, (DomainInterface, str))
             DomainEquationCondition.__dict__[key].__set__(self, value)
         elif key == 'equation':
             check_consistency(value, (EquationInterface))
diff --git a/pina/condition/input_equation_condition.py b/pina/condition/input_equation_condition.py
index bf05130c0..a5971114e 100644
--- a/pina/condition/input_equation_condition.py
+++ b/pina/condition/input_equation_condition.py
@@ -17,12 +17,11 @@ class InputPointsEquationCondition(ConditionInterface):
 
     def __init__(self, input_points, equation):
         """
-        TODO
+        TODO : add docstring
         """
         super().__init__()
         self.input_points = input_points
         self.equation = equation
-        self._condition_type = 'physics'
 
     def __setattr__(self, key, value):
         if key == 'input_points':
diff --git a/pina/condition/input_output_condition.py b/pina/condition/input_output_condition.py
index 08ed21d96..5cf5518e7 100644
--- a/pina/condition/input_output_condition.py
+++ b/pina/condition/input_output_condition.py
@@ -1,4 +1,5 @@
 import torch
+import torch_geometric
 
 from .condition_interface import ConditionInterface
 from ..label_tensor import LabelTensor
@@ -16,16 +17,15 @@ class InputOutputPointsCondition(ConditionInterface):
 
     def __init__(self, input_points, output_points):
         """
-        TODO
+        TODO : add docstring
         """
         super().__init__()
         self.input_points = input_points
         self.output_points = output_points
-        self._condition_type = ['supervised', 'physics']
 
     def __setattr__(self, key, value):
         if (key == 'input_points') or (key == 'output_points'):
-            check_consistency(value, (LabelTensor, Graph, torch.Tensor))
+            check_consistency(value, (LabelTensor, Graph, torch.Tensor, torch_geometric.data.Data))
             InputOutputPointsCondition.__dict__[key].__set__(self, value)
         elif key in ('_problem', '_condition_type'):
             super().__setattr__(key, value)
diff --git a/pina/data/__init__.py b/pina/data/__init__.py
index 2b3a126a7..292c9ed10 100644
--- a/pina/data/__init__.py
+++ b/pina/data/__init__.py
@@ -2,14 +2,11 @@
 Import data classes
 """
 __all__ = [
-    'PinaDataLoader', 'SupervisedDataset', 'SamplePointDataset',
-    'UnsupervisedDataset', 'Batch', 'PinaDataModule', 'BaseDataset'
+    'PinaDataModule',
+    'PinaDataset'
 ]
 
-from .pina_dataloader import PinaDataLoader
-from .supervised_dataset import SupervisedDataset
-from .sample_dataset import SamplePointDataset
-from .unsupervised_dataset import UnsupervisedDataset
-from .pina_batch import Batch
+
+
 from .data_module import PinaDataModule
-from .base_dataset import BaseDataset
+from .dataset import  PinaDataset
diff --git a/pina/data/base_dataset.py b/pina/data/base_dataset.py
deleted file mode 100644
index 2c28ba30b..000000000
--- a/pina/data/base_dataset.py
+++ /dev/null
@@ -1,157 +0,0 @@
-"""
-Basic data module implementation
-"""
-import torch
-import logging
-
-from torch.utils.data import Dataset
-
-from ..label_tensor import LabelTensor
-
-
-class BaseDataset(Dataset):
-    """
-    BaseDataset class, which handle initialization and data retrieval
-    :var condition_indices: List of indices
-    :var device: torch.device
-    """
-
-    def __new__(cls, problem=None, device=torch.device('cpu')):
-        """
-        Ensure correct definition of __slots__ before initialization
-        :param AbstractProblem problem: The formulation of the problem.
-        :param torch.device device: The device on which the
-        dataset will be loaded.
-        """
-        if cls is BaseDataset:
-            raise TypeError(
-                'BaseDataset cannot be instantiated directly. Use a subclass.')
-        if not hasattr(cls, '__slots__'):
-            raise TypeError(
-                'Something is wrong, __slots__ must be defined in subclasses.')
-        return object.__new__(cls)
-
-    def __init__(self, problem=None, device=torch.device('cpu')):
-        """"
-        Initialize the object based on __slots__
-        :param AbstractProblem problem: The formulation of the problem.
-        :param torch.device device: The device on which the
-        dataset will be loaded.
-        """
-        super().__init__()
-        self.empty = True
-        self.problem = problem
-        self.device = device
-        self.condition_indices = None
-        for slot in self.__slots__:
-            setattr(self, slot, [])
-        self.num_el_per_condition = []
-        self.conditions_idx = []
-        if self.problem is not None:
-            self._init_from_problem(self.problem.collector.data_collections)
-        self.initialized = False
-
-    def _init_from_problem(self, collector_dict):
-        """
-        TODO
-        """
-        for name, data in collector_dict.items():
-            keys = list(data.keys())
-            if set(self.__slots__) == set(keys):
-                self._populate_init_list(data)
-                idx = [
-                    key for key, val in
-                    self.problem.collector.conditions_name.items()
-                    if val == name
-                ]
-                self.conditions_idx.append(idx)
-        self.initialize()
-
-    def add_points(self, data_dict, condition_idx, batching_dim=0):
-        """
-        This method filled internal lists of data points
-        :param data_dict: dictionary containing data points
-        :param condition_idx: index of the condition to which the data points
-                belong to
-        :param batching_dim: dimension of the batching
-        :raises: ValueError if the dataset has already been initialized
-        """
-        if not self.initialized:
-            self._populate_init_list(data_dict, batching_dim)
-            self.conditions_idx.append(condition_idx)
-            self.empty = False
-        else:
-            raise ValueError('Dataset already initialized')
-
-    def _populate_init_list(self, data_dict, batching_dim=0):
-        current_cond_num_el = None
-        for slot in data_dict.keys():
-            slot_data = data_dict[slot]
-            if batching_dim != 0:
-                if isinstance(slot_data, (LabelTensor, torch.Tensor)):
-                    dims = len(slot_data.size())
-                    slot_data = slot_data.permute(
-                        [batching_dim] +
-                        [dim for dim in range(dims) if dim != batching_dim])
-            if current_cond_num_el is None:
-                current_cond_num_el = len(slot_data)
-            elif current_cond_num_el != len(slot_data):
-                raise ValueError('Different dimension in same condition')
-            current_list = getattr(self, slot)
-            current_list += [
-                slot_data
-            ] if not (isinstance(slot_data, list)) else slot_data
-        self.num_el_per_condition.append(current_cond_num_el)
-
-    def initialize(self):
-        """
-        Initialize the datasets tensors/LabelTensors/lists given the lists
-        already filled
-        """
-        logging.debug(f'Initialize dataset {self.__class__.__name__}')
-
-        if self.num_el_per_condition:
-            self.condition_indices = torch.cat([
-                torch.tensor([i] * self.num_el_per_condition[i],
-                             dtype=torch.uint8)
-                for i in range(len(self.num_el_per_condition))
-            ],
-                                               dim=0)
-            for slot in self.__slots__:
-                current_attribute = getattr(self, slot)
-                if all(isinstance(a, LabelTensor) for a in current_attribute):
-                    setattr(self, slot, LabelTensor.vstack(current_attribute))
-        self.initialized = True
-
-    def __len__(self):
-        """
-        :return: Number of elements in the dataset
-        """
-        return len(getattr(self, self.__slots__[0]))
-
-    def __getitem__(self, idx):
-        """
-        :param idx:
-        :return:
-        """
-        if not isinstance(idx, (tuple, list, slice, int)):
-            raise IndexError("Invalid index")
-        tensors = []
-        for attribute in self.__slots__:
-            tensor = getattr(self, attribute)
-            if isinstance(attribute, (LabelTensor, torch.Tensor)):
-                tensors.append(tensor.__getitem__(idx))
-            elif isinstance(attribute, list):
-                if isinstance(idx, (list, tuple)):
-                    tensor = [tensor[i] for i in idx]
-                tensors.append(tensor)
-        return tensors
-
-    def apply_shuffle(self, indices):
-        for slot in self.__slots__:
-            if slot != 'equation':
-                attribute = getattr(self, slot)
-                if isinstance(attribute, (LabelTensor, torch.Tensor)):
-                    setattr(self, 'slot', attribute[[indices]])
-                if isinstance(attribute, list):
-                    setattr(self, 'slot', [attribute[i] for i in indices])
diff --git a/pina/data/data_module.py b/pina/data/data_module.py
index bd117b54b..c9af8aeba 100644
--- a/pina/data/data_module.py
+++ b/pina/data/data_module.py
@@ -1,17 +1,71 @@
-"""
-This module provide basic data management functionalities
-"""
-
+import logging
+from lightning.pytorch import LightningDataModule
 import math
 import torch
-import logging
-from pytorch_lightning import LightningDataModule
-from .sample_dataset import SamplePointDataset
-from .supervised_dataset import SupervisedDataset
-from .unsupervised_dataset import UnsupervisedDataset
-from .pina_dataloader import PinaDataLoader
-from .pina_subset import PinaSubset
+from ..label_tensor import LabelTensor
+from torch.utils.data import DataLoader, BatchSampler, SequentialSampler, \
+    RandomSampler
+from torch.utils.data.distributed import DistributedSampler
+from .dataset import PinaDatasetFactory
+
+class Collator:
+    def __init__(self, max_conditions_lengths, ):
+        self.max_conditions_lengths = max_conditions_lengths
+        self.callable_function = self._collate_custom_dataloader if \
+            max_conditions_lengths is None else (
+            self._collate_standard_dataloader)
 
+    @staticmethod
+    def _collate_custom_dataloader(batch):
+        return batch[0]
+
+    def _collate_standard_dataloader(self, batch):
+        """
+        Function used to collate the batch
+        """
+        batch_dict = {}
+        if isinstance(batch, dict):
+            return batch
+        conditions_names = batch[0].keys()
+
+        # Condition names
+        for condition_name in conditions_names:
+            single_cond_dict = {}
+            condition_args = batch[0][condition_name].keys()
+            for arg in condition_args:
+                data_list = [batch[idx][condition_name][arg] for idx in range(
+                    min(len(batch),
+                    self.max_conditions_lengths[condition_name]))]
+                if isinstance(data_list[0], LabelTensor):
+                    single_cond_dict[arg] = LabelTensor.stack(data_list)
+                elif isinstance(data_list[0], torch.Tensor):
+                    single_cond_dict[arg] = torch.stack(data_list)
+                else:
+                    raise NotImplementedError(
+                        f"Data type {type(data_list[0])} not supported")
+            batch_dict[condition_name] = single_cond_dict
+        return batch_dict
+
+    def __call__(self, batch):
+        return self.callable_function(batch)
+
+
+class PinaBatchSampler(BatchSampler):
+    def __init__(self, dataset, batch_size, shuffle, sampler=None):
+        if sampler is None:
+            if (torch.distributed.is_available() and
+                    torch.distributed.is_initialized()):
+                rank = torch.distributed.get_rank()
+                world_size = torch.distributed.get_world_size()
+                sampler = DistributedSampler(dataset, shuffle=shuffle,
+                                             rank=rank, num_replicas=world_size)
+            else:
+                if shuffle:
+                    sampler = RandomSampler(dataset)
+                else:
+                    sampler = SequentialSampler(dataset)
+        super().__init__(sampler=sampler, batch_size=batch_size,
+                         drop_last=False)
 
 class PinaDataModule(LightningDataModule):
     """
@@ -20,160 +74,218 @@ class PinaDataModule(LightningDataModule):
     """
 
     def __init__(self,
-                 problem,
-                 device,
+                 collector,
                  train_size=.7,
-                 test_size=.1,
-                 val_size=.2,
+                 test_size=.2,
+                 val_size=.1,
                  predict_size=0.,
                  batch_size=None,
                  shuffle=True,
-                 datasets=None):
+                 repeat=False,
+                 automatic_batching=False
+                 ):
         """
         Initialize the object, creating dataset based on input problem
-        :param AbstractProblem problem: PINA problem
-        :param device: Device used for training and testing
+        :param Collector collector: PINA problem
         :param train_size: number/percentage of elements in train split
         :param test_size: number/percentage of elements in test split
-        :param eval_size: number/percentage of elements in evaluation split
+        :param val_size: number/percentage of elements in evaluation split
         :param batch_size: batch size used for training
-        :param datasets: list of datasets objects
         """
         logging.debug('Start initialization of Pina DataModule')
         logging.info('Start initialization of Pina DataModule')
         super().__init__()
-        self.problem = problem
-        self.device = device
-        self.dataset_classes = [
-            SupervisedDataset, UnsupervisedDataset, SamplePointDataset
-        ]
-        if datasets is None:
-            self.datasets = None
-        else:
-            self.datasets = datasets
-
-        self.split_length = []
-        self.split_names = []
-        self.loader_functions = {}
+        self.default_batching = automatic_batching
         self.batch_size = batch_size
-        self.condition_names = problem.collector.conditions_name
+        self.shuffle = shuffle
+        self.repeat = repeat
 
+        # Begin Data splitting
+        splits_dict = {}
         if train_size > 0:
-            self.split_names.append('train')
-            self.split_length.append(train_size)
-            self.loader_functions['train_dataloader'] = lambda: PinaDataLoader(
-                self.splits['train'], self.batch_size, self.condition_names)
+            splits_dict['train'] = train_size
+            self.train_dataset = None
+        else:
+            self.train_dataloader = super().train_dataloader
         if test_size > 0:
-            self.split_length.append(test_size)
-            self.split_names.append('test')
-            self.loader_functions['test_dataloader'] = lambda: PinaDataLoader(
-                self.splits['test'], self.batch_size, self.condition_names)
+            splits_dict['test'] = test_size
+            self.test_dataset = None
+        else:
+            self.test_dataloader = super().test_dataloader
         if val_size > 0:
-            self.split_length.append(val_size)
-            self.split_names.append('val')
-            self.loader_functions['val_dataloader'] = lambda: PinaDataLoader(
-                self.splits['val'], self.batch_size, self.condition_names)
+            splits_dict['val'] = val_size
+            self.val_dataset = None
+        else:
+            self.val_dataloader = super().val_dataloader
         if predict_size > 0:
-            self.split_length.append(predict_size)
-            self.split_names.append('predict')
-            self.loader_functions['predict_dataloader'] = lambda: PinaDataLoader(
-                self.splits['predict'], self.batch_size, self.condition_names)
-        self.splits = {k: {} for k in self.split_names}
-        self.shuffle = shuffle
-
-        for k, v in self.loader_functions.items():
-            setattr(self, k, v)
-
-    def prepare_data(self):
-        if self.datasets is None:
-            self._create_datasets()
+            splits_dict['predict'] = predict_size
+            self.predict_dataset = None
+        else:
+            self.predict_dataloader = super().predict_dataloader
+        self.collector_splits = self._create_splits(collector, splits_dict)
 
     def setup(self, stage=None):
         """
         Perform the splitting of the dataset
         """
         logging.debug('Start setup of Pina DataModule obj')
-        if self.datasets is None:
-            self._create_datasets()
         if stage == 'fit' or stage is None:
-            for dataset in self.datasets:
-                if len(dataset) > 0:
-                    splits = self.dataset_split(dataset,
-                                                self.split_length,
-                                                shuffle=self.shuffle)
-                    for i in range(len(self.split_length)):
-                        self.splits[self.split_names[i]][
-                            dataset.data_type] = splits[i]
+            self.train_dataset = PinaDatasetFactory(
+                self.collector_splits['train'],
+                max_conditions_lengths=self.find_max_conditions_lengths(
+                    'train'))
+            if 'val' in self.collector_splits.keys():
+                self.val_dataset = PinaDatasetFactory(
+                    self.collector_splits['val'],
+                    max_conditions_lengths=self.find_max_conditions_lengths(
+                        'val')
+                )
         elif stage == 'test':
-            raise NotImplementedError("Testing pipeline not implemented yet")
+            self.test_dataset = PinaDatasetFactory(
+                self.collector_splits['test'],
+                max_conditions_lengths=self.find_max_conditions_lengths(
+                    'test')
+            )
+        elif stage == 'predict':
+            self.predict_dataset = PinaDatasetFactory(
+                self.collector_splits['predict'],
+                max_conditions_lengths=self.find_max_conditions_lengths(
+                                        'predict')
+            )
         else:
-            raise ValueError("stage must be either 'fit' or 'test'")
+            raise ValueError(
+                "stage must be either 'fit' or 'test' or 'predict'."
+            )
 
     @staticmethod
-    def dataset_split(dataset, lengths, seed=None, shuffle=True):
-        """
-        Perform the splitting of the dataset
-        :param dataset: dataset object we wanted to split
-        :param lengths: lengths of elements in dataset
-        :param seed: random seed
-        :param shuffle: shuffle dataset
-        :return: split dataset
-        :rtype: PinaSubset
-        """
-        if sum(lengths) - 1 < 1e-3:
-            len_dataset = len(dataset)
-            lengths = [
-                int(math.floor(len_dataset * length)) for length in lengths
-            ]
-            remainder = len(dataset) - sum(lengths)
-            for i in range(remainder):
-                lengths[i % len(lengths)] += 1
-        elif sum(lengths) - 1 >= 1e-3:
-            raise ValueError(f"Sum of lengths is {sum(lengths)} less than 1")
-
-        if shuffle:
-            if seed is not None:
-                generator = torch.Generator()
-                generator.manual_seed(seed)
-                indices = torch.randperm(sum(lengths), generator=generator)
-            else:
-                indices = torch.randperm(sum(lengths))
-            dataset.apply_shuffle(indices)
+    def _split_condition(condition_dict, splits_dict):
+        len_condition = len(condition_dict['input_points'])
 
-        indices = torch.arange(0, sum(lengths), 1, dtype=torch.uint8).tolist()
-        offsets = [
-            sum(lengths[:i]) if i > 0 else 0 for i in range(len(lengths))
-        ]
-        return [
-            PinaSubset(dataset, indices[offset:offset + length])
-            for offset, length in zip(offsets, lengths)
+        lengths = [
+            int(math.floor(len_condition * length)) for length in
+            splits_dict.values()
         ]
 
-    def _create_datasets(self):
+        remainder = len_condition - sum(lengths)
+        for i in range(remainder):
+            lengths[i % len(lengths)] += 1
+        splits_dict = {k: v for k, v in zip(splits_dict.keys(), lengths)
+                       }
+        to_return_dict = {}
+        offset = 0
+        for stage, stage_len in splits_dict.items():
+            to_return_dict[stage] = {k: v[offset:offset + stage_len]
+                                     for k, v in condition_dict.items() if
+                                     k != 'equation'
+                                     # Equations are NEVER dataloaded
+                                     }
+            offset += stage_len
+        return to_return_dict
+
+    def _create_splits(self, collector, splits_dict):
         """
-        Create the dataset objects putting data 
+        Create the dataset objects putting data
         """
+
+        # ----------- Auxiliary function ------------
+        def _apply_shuffle(condition_dict, len_data):
+            idx = torch.randperm(len_data)
+            for k, v in condition_dict.items():
+                if k == 'equation':
+                    continue
+                if isinstance(v, list):
+                    condition_dict[k] = [v[i] for i in idx]
+                elif isinstance(v, LabelTensor):
+                    condition_dict[k] = LabelTensor(v.tensor[idx],
+                                                    v.labels)
+                elif isinstance(v, torch.Tensor):
+                    condition_dict[k] = v[idx]
+                else:
+                    raise ValueError(f"Data type {type(v)} not supported")
+        # ----------- End auxiliary function ------------
+
         logging.debug('Dataset creation in PinaDataModule obj')
-        collector = self.problem.collector
-        batching_dim = self.problem.batching_dimension
-        datasets_slots = [i.__slots__ for i in self.dataset_classes]
-        self.datasets = [
-            dataset(device=self.device) for dataset in self.dataset_classes
+        split_names = list(splits_dict.keys())
+        dataset_dict = {name: {} for name in split_names}
+        for condition_name, condition_dict in collector.data_collections.items():
+            len_data = len(condition_dict['input_points'])
+            if self.shuffle:
+                _apply_shuffle(condition_dict, len_data)
+            for key, data in self._split_condition(condition_dict,
+                                                   splits_dict).items():
+                dataset_dict[key].update({condition_name: data})
+        return dataset_dict
+
+    def find_max_conditions_lengths(self, split):
+        max_conditions_lengths = {}
+        for k, v in self.collector_splits[split].items():
+            if self.batch_size is None:
+                max_conditions_lengths[k] = len(v['input_points'])
+            elif self.repeat:
+                max_conditions_lengths[k] = self.batch_size
+            else:
+                max_conditions_lengths[k] = min(len(v['input_points']),
+                                                self.batch_size)
+        return max_conditions_lengths
+
+    def val_dataloader(self):
+        """
+        Create the validation dataloader
+        """
+
+        batch_size = self.batch_size if self.batch_size is not None else len(
+            self.val_dataset)
+
+        # Use default batching in torch DataLoader (good is batch size is small)
+        if self.default_batching:
+            collate = Collator(self.find_max_conditions_lengths('val'))
+            return DataLoader(self.val_dataset, self.batch_size,
+                              collate_fn=collate)
+        collate = Collator(None)
+        # Use custom batching (good if batch size is large)
+        sampler = PinaBatchSampler(self.val_dataset, batch_size, shuffle=False)
+        return DataLoader(self.val_dataset, sampler=sampler,
+                          collate_fn=collate)
+
+    def train_dataloader(self):
+        """
+        Create the training dataloader
+        """
+        # Use default batching in torch DataLoader (good is batch size is small)
+        if self.default_batching:
+            collate = Collator(self.find_max_conditions_lengths('train'))
+            return DataLoader(self.train_dataset, self.batch_size,
+                                collate_fn=collate)
+        collate = Collator(None)
+        # Use custom batching (good if batch size is large)
+        batch_size = self.batch_size if self.batch_size is not None else len(
+            self.train_dataset)
+        sampler = PinaBatchSampler(self.train_dataset, batch_size,
+                                       shuffle=False)
+        return DataLoader(self.train_dataset, sampler=sampler,
+                              collate_fn=collate)
+
+    def test_dataloader(self):
+        """
+        Create the testing dataloader
+        """
+        raise NotImplementedError("Test dataloader not implemented")
+
+    def predict_dataloader(self):
+        """
+        Create the prediction dataloader
+        """
+        raise NotImplementedError("Predict dataloader not implemented")
+
+    def transfer_batch_to_device(self, batch, device, dataloader_idx):
+        """
+        Transfer the batch to the device. This method is called in the
+        training loop and is used to transfer the batch to the device.
+        """
+        batch = [
+            (k, super(LightningDataModule, self).transfer_batch_to_device(v,
+                                                                device,
+                                                                dataloader_idx))
+            for k, v in batch.items()
         ]
-        logging.debug('Filling datasets in PinaDataModule obj')
-        for name, data in collector.data_collections.items():
-            keys = list(data.keys())
-            idx = [
-                key for key, val in collector.conditions_name.items()
-                if val == name
-            ]
-            for i, slot in enumerate(datasets_slots):
-                if slot == keys:
-                    self.datasets[i].add_points(data, idx[0], batching_dim)
-                    continue
-        datasets = []
-        for dataset in self.datasets:
-            if not dataset.empty:
-                dataset.initialize()
-                datasets.append(dataset)
-        self.datasets = datasets
+        return batch
diff --git a/pina/data/dataset.py b/pina/data/dataset.py
new file mode 100644
index 000000000..0bc923711
--- /dev/null
+++ b/pina/data/dataset.py
@@ -0,0 +1,102 @@
+"""
+This module provide basic data management functionalities
+"""
+import torch
+from torch.utils.data import Dataset
+from abc import abstractmethod
+from torch_geometric.data import Batch
+
+class PinaDatasetFactory:
+    """
+    Factory class for the PINA dataset. Depending on the type inside the
+    conditions it creates a different dataset object:
+    - PinaTensorDataset for torch.Tensor
+    - PinaGraphDataset for list of torch_geometric.data.Data objects
+    """
+    def __new__(cls, conditions_dict, **kwargs):
+        if len(conditions_dict) == 0:
+            raise ValueError('No conditions provided')
+        if all([isinstance(v['input_points'], torch.Tensor) for v
+                in conditions_dict.values()]):
+            return PinaTensorDataset(conditions_dict, **kwargs)
+        elif all([isinstance(v['input_points'], list) for v
+                  in conditions_dict.values()]):
+            return PinaGraphDataset(conditions_dict, **kwargs)
+        raise ValueError('Conditions must be either torch.Tensor or list of Data '
+                         'objects.')
+
+class PinaDataset(Dataset):
+    """
+    Abstract class for the PINA dataset
+    """
+    def __init__(self, conditions_dict, max_conditions_lengths):
+        self.conditions_dict = conditions_dict
+        self.max_conditions_lengths = max_conditions_lengths
+        self.conditions_length = {k: len(v['input_points']) for k, v in
+                                  self.conditions_dict.items()}
+        self.length = max(self.conditions_length.values())
+
+    def _get_max_len(self):
+        max_len = 0
+        for condition in self.conditions_dict.values():
+            max_len = max(max_len, len(condition['input_points']))
+        return max_len
+
+    def __len__(self):
+        return self.length
+
+    @abstractmethod
+    def __getitem__(self, item):
+        pass
+
+class PinaTensorDataset(PinaDataset):
+    def __init__(self, conditions_dict, max_conditions_lengths,
+                 ):
+        super().__init__(conditions_dict, max_conditions_lengths)
+
+    def _getitem_int(self, idx):
+        return {
+            k: {k_data: v[k_data][idx % len(v['input_points'])] for k_data
+                in v.keys()} for k, v in self.conditions_dict.items()
+        }
+
+    def _getitem_list(self, idx):
+        to_return_dict = {}
+        for condition, data in self.conditions_dict.items():
+            cond_idx = idx[:self.max_conditions_lengths[condition]]
+            condition_len = self.conditions_length[condition]
+            if self.length > condition_len:
+                cond_idx = [idx%condition_len for idx in cond_idx]
+            to_return_dict[condition] = {k: v[cond_idx]
+                                         for k, v in data.items()}
+        return to_return_dict
+
+    def __getitem__(self, idx):
+        if isinstance(idx, int):
+            return self._getitem_int(idx)
+        return self._getitem_list(idx)
+
+class PinaGraphDataset(PinaDataset):
+    pass
+    """
+    def __init__(self, conditions_dict, max_conditions_lengths):
+        super().__init__(conditions_dict, max_conditions_lengths)
+
+    def __getitem__(self, idx):
+        
+        Getitem method for large batch size
+        
+        to_return_dict = {}
+        for condition, data in self.conditions_dict.items():
+            cond_idx = idx[:self.max_conditions_lengths[condition]]
+            condition_len = self.conditions_length[condition]
+            if self.length > condition_len:
+                cond_idx = [idx%condition_len for idx in cond_idx]
+            to_return_dict[condition] = {k: Batch.from_data_list([v[i]
+                                            for i in cond_idx])
+                            if isinstance(v, list)
+                            else v[cond_idx].tensor.reshape(-1, v.size(-1))
+                        for k, v in data.items()
+                    }
+        return to_return_dict
+    """
diff --git a/pina/data/pina_batch.py b/pina/data/pina_batch.py
deleted file mode 100644
index c5d1b61dc..000000000
--- a/pina/data/pina_batch.py
+++ /dev/null
@@ -1,47 +0,0 @@
-"""
-Batch management module
-"""
-from .pina_subset import PinaSubset
-
-
-class Batch:
-    """
-    Implementation of the Batch class used during training to perform SGD
-    optimization.
-    """
-
-    def __init__(self, dataset_dict, idx_dict, require_grad=True):
-        self.attributes = []
-        for k, v in dataset_dict.items():
-            setattr(self, k, v)
-            self.attributes.append(k)
-
-        for k, v in idx_dict.items():
-            setattr(self, k + '_idx', v)
-        self.require_grad = require_grad
-
-    def __len__(self):
-        """
-        Returns the number of elements in the batch
-        :return: number of elements in the batch
-        :rtype: int
-        """
-        length = 0
-        for dataset in dir(self):
-            attribute = getattr(self, dataset)
-            if isinstance(attribute, list):
-                length += len(getattr(self, dataset))
-        return length
-
-    def __getattribute__(self, item):
-        if item in super().__getattribute__('attributes'):
-            dataset = super().__getattribute__(item)
-            index = super().__getattribute__(item + '_idx')
-            return PinaSubset(dataset.dataset, dataset.indices[index])
-        return super().__getattribute__(item)
-
-    def __getattr__(self, item):
-        if item == 'data' and len(self.attributes) == 1:
-            item = self.attributes[0]
-            return super().__getattribute__(item)
-        raise AttributeError(f"'Batch' object has no attribute '{item}'")
diff --git a/pina/data/pina_dataloader.py b/pina/data/pina_dataloader.py
deleted file mode 100644
index e2d3fb76e..000000000
--- a/pina/data/pina_dataloader.py
+++ /dev/null
@@ -1,68 +0,0 @@
-"""
-This module is used to create an iterable object used during training
-"""
-import math
-from .pina_batch import Batch
-
-
-class PinaDataLoader:
-    """
-    This class is used to create a dataloader to use during the training.
-
-    :var condition_names: The names of the conditions. The order is consistent
-        with the condition indeces in the batches.
-    :vartype condition_names: list[str]
-    """
-
-    def __init__(self, dataset_dict, batch_size, condition_names) -> None:
-        """
-        Initialize local variables
-        :param dataset_dict: Dictionary of datasets
-        :type dataset_dict: dict
-        :param batch_size: Size of the batch
-        :type batch_size: int
-        :param condition_names: Names of the conditions
-        :type condition_names: list[str]
-        """
-        self.condition_names = condition_names
-        self.dataset_dict = dataset_dict
-        self._init_batches(batch_size)
-
-    def _init_batches(self, batch_size=None):
-        """
-        Create batches according to the batch_size provided in input.
-        """
-        self.batches = []
-        n_elements = sum(len(v) for v in self.dataset_dict.values())
-        if batch_size is None:
-            batch_size = n_elements
-        indexes_dict = {}
-        n_batches = int(math.ceil(n_elements / batch_size))
-        for k, v in self.dataset_dict.items():
-            if n_batches != 1:
-                indexes_dict[k] = math.floor(len(v) / (n_batches - 1))
-            else:
-                indexes_dict[k] = len(v)
-        for i in range(n_batches):
-            temp_dict = {}
-            for k, v in indexes_dict.items():
-                if i != n_batches - 1:
-                    temp_dict[k] = slice(i * v, (i + 1) * v)
-                else:
-                    temp_dict[k] = slice(i * v, len(self.dataset_dict[k]))
-            self.batches.append(
-                Batch(idx_dict=temp_dict, dataset_dict=self.dataset_dict))
-
-    def __iter__(self):
-        """
-        Makes dataloader object iterable
-        """
-        yield from self.batches
-
-    def __len__(self):
-        """
-        Return the number of batches.
-        :return: The number of batches.
-        :rtype: int
-        """
-        return len(self.batches)
diff --git a/pina/data/pina_subset.py b/pina/data/pina_subset.py
deleted file mode 100644
index 275541e97..000000000
--- a/pina/data/pina_subset.py
+++ /dev/null
@@ -1,36 +0,0 @@
-"""
-Module for PinaSubset class
-"""
-from pina import LabelTensor
-from torch import Tensor, float32
-
-
-class PinaSubset:
-    """
-    TODO
-    """
-    __slots__ = ['dataset', 'indices', 'require_grad']
-
-    def __init__(self, dataset, indices, require_grad=True):
-        """
-        TODO
-        """
-        self.dataset = dataset
-        self.indices = indices
-        self.require_grad = require_grad
-
-    def __len__(self):
-        """
-        TODO
-        """
-        return len(self.indices)
-
-    def __getattr__(self, name):
-        tensor = self.dataset.__getattribute__(name)
-        if isinstance(tensor, (LabelTensor, Tensor)):
-            tensor = tensor[[self.indices]].to(self.dataset.device)
-            return tensor.requires_grad_(
-                self.require_grad) if tensor.dtype == float32 else tensor
-        if isinstance(tensor, list):
-            return [tensor[i] for i in self.indices]
-        raise AttributeError(f"No attribute named {name}")
diff --git a/pina/data/sample_dataset.py b/pina/data/sample_dataset.py
deleted file mode 100644
index bc3bca335..000000000
--- a/pina/data/sample_dataset.py
+++ /dev/null
@@ -1,35 +0,0 @@
-"""
-Sample dataset module
-"""
-from copy import deepcopy
-from .base_dataset import BaseDataset
-from ..condition import InputPointsEquationCondition
-
-
-class SamplePointDataset(BaseDataset):
-    """
-    This class extends the BaseDataset to handle physical datasets
-    composed of only input points.
-    """
-    data_type = 'physics'
-    __slots__ = InputPointsEquationCondition.__slots__
-
-    def add_points(self, data_dict, condition_idx, batching_dim=0):
-        data_dict = deepcopy(data_dict)
-        data_dict.pop('equation')
-        super().add_points(data_dict, condition_idx)
-
-    def _init_from_problem(self, collector_dict):
-        for name, data in collector_dict.items():
-            keys = list(data.keys())
-            if set(self.__slots__) == set(keys):
-                data = deepcopy(data)
-                data.pop('equation')
-                self._populate_init_list(data)
-                idx = [
-                    key for key, val in
-                    self.problem.collector.conditions_name.items()
-                    if val == name
-                ]
-                self.conditions_idx.append(idx)
-        self.initialize()
diff --git a/pina/data/supervised_dataset.py b/pina/data/supervised_dataset.py
deleted file mode 100644
index be601050a..000000000
--- a/pina/data/supervised_dataset.py
+++ /dev/null
@@ -1,13 +0,0 @@
-"""
-Supervised dataset module
-"""
-from .base_dataset import BaseDataset
-
-
-class SupervisedDataset(BaseDataset):
-    """
-    This class extends the BaseDataset to handle datasets that consist of
-    input-output pairs.
-    """
-    data_type = 'supervised'
-    __slots__ = ['input_points', 'output_points']
diff --git a/pina/data/unsupervised_dataset.py b/pina/data/unsupervised_dataset.py
deleted file mode 100644
index 18cf296f5..000000000
--- a/pina/data/unsupervised_dataset.py
+++ /dev/null
@@ -1,14 +0,0 @@
-"""
-Unsupervised dataset module
-"""
-from .base_dataset import BaseDataset
-
-
-class UnsupervisedDataset(BaseDataset):
-    """
-    This class extend BaseDataset class to handle
-    unsupervised dataset,composed of input points
-    and, optionally, conditional variables
-    """
-    data_type = 'unsupervised'
-    __slots__ = ['input_points', 'conditional_variables']
diff --git a/pina/graph.py b/pina/graph.py
index 97b2770e6..bde5bbf50 100644
--- a/pina/graph.py
+++ b/pina/graph.py
@@ -93,8 +93,8 @@ def _build_radius(**kwargs):
 
         logging.debug(f"edge_index computed")
         return Data(
-            x=nodes_data,
-            pos=nodes_coordinates,
+            x=nodes_data.tensor,
+            pos=nodes_coordinates.tensor,
             edge_index=edge_index,
             edge_attr=edges_data,
         )
diff --git a/pina/label_tensor.py b/pina/label_tensor.py
index 719975c51..a3cf5d237 100644
--- a/pina/label_tensor.py
+++ b/pina/label_tensor.py
@@ -4,26 +4,20 @@
 from torch import Tensor
 
 
-def issubset(a, b):
-    """
-    Check if a is a subset of b.
-    """
-    if isinstance(a, list) and isinstance(b, list):
-        return set(a).issubset(set(b))
-    if isinstance(a, range) and isinstance(b, range):
-        return a.start <= b.start and a.stop >= b.stop
-    return False
-
+full_labels = True
+MATH_FUNCTIONS = {torch.sin, torch.cos}
 
 class LabelTensor(torch.Tensor):
     """Torch tensor with a label for any column."""
 
     @staticmethod
     def __new__(cls, x, labels, *args, **kwargs):
+        full = kwargs.pop("full", full_labels)
+
         if isinstance(x, LabelTensor):
+            x.full = full
             return x
-        else:
-            return super().__new__(cls, x, *args, **kwargs)
+        return super().__new__(cls, x, *args, **kwargs)
 
     @property
     def tensor(self):
@@ -40,22 +34,11 @@ def __init__(self, x, labels, **kwargs):
                     {1: {"name": "space"['a', 'b', 'c'])
 
         """
-        self.dim_names = None
-        self.full = kwargs.get('full', True)
-        self.labels = labels
-
-    @classmethod
-    def __internal_init__(cls,
-                          x,
-                          labels,
-                          dim_names,
-                          *args,
-                          **kwargs):
-        lt = cls.__new__(cls, x, labels, *args, **kwargs)
-        lt._labels = labels
-        lt.full = kwargs.get('full', True)
-        lt.dim_names = dim_names
-        return lt
+        self.full = kwargs.get('full', full_labels)
+        if labels is not None:
+            self.labels = labels
+        else:
+            self._labels = {}
 
     @property
     def labels(self):
@@ -104,14 +87,13 @@ def labels(self, labels):
             self._labels = {}
         if isinstance(labels, dict):
             self._init_labels_from_dict(labels)
-        elif isinstance(labels, list):
+        elif isinstance(labels, (list, range)):
             self._init_labels_from_list(labels)
         elif isinstance(labels, str):
             labels = [labels]
             self._init_labels_from_list(labels)
         else:
             raise ValueError("labels must be list, dict or string.")
-        self.set_names()
 
     def _init_labels_from_dict(self, labels):
         """
@@ -125,34 +107,38 @@ def _init_labels_from_dict(self, labels):
             """
         tensor_shape = self.shape
 
+        # Set all labels if full_labels is True
         if hasattr(self, 'full') and self.full:
             labels = {
                 i: labels[i] if i in labels else {
-                    'name': i
+                    'name': i, 'dof': range(tensor_shape[i])
                 }
-                for i in labels.keys()
+                for i in range(len(tensor_shape))
             }
+
         for k, v in labels.items():
+
             # Init labels from str
             if isinstance(v, str):
                 v = {'name': v, 'dof': range(tensor_shape[k])}
+
             # Init labels from dict
-            elif isinstance(v, dict) and list(v.keys()) == ['name']:
-                # Init from dict with only name key
-                v['dof'] = range(tensor_shape[k])
-                # Init from dict with both name and dof keys
-            elif isinstance(v, dict) and sorted(list(
-                    v.keys())) == ['dof', 'name']:
-                dof_list = v['dof']
-                dof_len = len(dof_list)
-                if dof_len != len(set(dof_list)):
-                    raise ValueError("dof must be unique")
-                if dof_len != tensor_shape[k]:
-                    raise ValueError(
-                        'Number of dof does not match tensor shape')
+            elif isinstance(v, dict):
+                # Only name of the dimension if provided
+                if list(v.keys()) == ['name']:
+                    v['dof'] = range(tensor_shape[k])
+                # Both name and dof are provided
+                elif sorted(list(v.keys())) == ['dof', 'name']:
+                    dof_list = v['dof']
+                    dof_len = len(dof_list)
+                    if dof_len != len(set(dof_list)):
+                        raise ValueError("dof must be unique")
+                    if dof_len != tensor_shape[k]:
+                        raise ValueError(
+                            'Number of dof does not match tensor shape')
             else:
                 raise ValueError('Illegal labels initialization')
-            # Perform update
+            # Assign labels values
             self._labels[k] = v
 
     def _init_labels_from_list(self, labels):
@@ -172,75 +158,71 @@ def _init_labels_from_list(self, labels):
         }
         self._init_labels_from_dict(last_dim_labels)
 
-    def set_names(self):
-        labels = self.stored_labels
-        self.dim_names = {}
-        for dim in labels.keys():
-            self.dim_names[labels[dim]['name']] = dim
-
     def extract(self, labels_to_extract):
         """
         Extract the subset of the original tensor by returning all the columns
         corresponding to the passed ``label_to_extract``.
 
-        :param label_to_extract: The label(s) to extract.
-        :type label_to_extract: str | list(str) | tuple(str)
+        :param labels_to_extract: The label(s) to extract.
+        :type labels_to_extract: str | list(str) | tuple(str)
         :raises TypeError: Labels are not ``str``.
         :raises ValueError: Label to extract is not in the labels ``list``.
         """
         # Convert str/int to string
+        def find_names(labels):
+            dim_names = {}
+            for dim in labels.keys():
+                dim_names[labels[dim]['name']] = dim
+            return dim_names
+
         if isinstance(labels_to_extract, (str, int)):
             labels_to_extract = [labels_to_extract]
 
         # Store useful variables
-        labels = self.stored_labels
+        labels = copy(self._labels)
         stored_keys = labels.keys()
-        dim_names = self.dim_names
+        dim_names = find_names(labels)
         ndim = len(super().shape)
 
-        # Convert tuple/list to dict
+        # Convert tuple/list to dict (having a list as input
+        # means that we want to extract a values from the last dimension)
         if isinstance(labels_to_extract, (tuple, list)):
             if not ndim - 1 in stored_keys:
                 raise ValueError(
                     "LabelTensor does not have labels in last dimension")
-            name = labels[max(stored_keys)]['name']
+            name = labels[ndim-1]['name']
             labels_to_extract = {name: list(labels_to_extract)}
 
         # If labels_to_extract is not dict then rise error
         if not isinstance(labels_to_extract, dict):
             raise ValueError('labels_to_extract must be str or list or dict')
 
-        # Make copy of labels (avoid issue in consistency)
-        updated_labels = {k: copy(v) for k, v in labels.items()}
-
         # Initialize list used to perform extraction
-        extractor = [slice(None) for _ in range(ndim)]
+        extractor = [slice(None)]*ndim
 
         # Loop over labels_to_extract dict
-        for k, v in labels_to_extract.items():
+        for dim_name, labels_te in labels_to_extract.items():
 
             # If label is not find raise value error
-            idx_dim = dim_names.get(k)
+            idx_dim = dim_names.get(dim_name, None)
             if idx_dim is None:
                 raise ValueError(
                     'Cannot extract label with is not in original labels')
 
             dim_labels = labels[idx_dim]['dof']
-            v = [v] if isinstance(v, (int, str)) else v
-
-            if not isinstance(v, range):
-                extractor[idx_dim] = [dim_labels.index(i)
-                                      for i in v] if len(v) > 1 else slice(
-                                          dim_labels.index(v[0]),
-                                          dim_labels.index(v[0]) + 1)
+            labels_te = [labels_te] if isinstance(labels_te, (int, str)) else labels_te
+            if not isinstance(labels_te, range):
+                #If is done to keep the dimension if there is only one extracted label
+                extractor[idx_dim] = [dim_labels.index(i) for i in labels_te] \
+                    if len(labels_te)>1 else slice(dim_labels.index(labels_te[0]), dim_labels.index(labels_te[0])+1)
             else:
-                extractor[idx_dim] = slice(v.start, v.stop)
+                extractor[idx_dim] = slice(labels_te.start, labels_te.stop)
 
-            updated_labels.update({idx_dim: {'dof': v, 'name': k}})
+            labels.update({idx_dim: {'dof': labels_te, 'name': dim_name}})
 
-        tensor = self.tensor
-        tensor = tensor[extractor]
-        return LabelTensor.__internal_init__(tensor, updated_labels, dim_names)
+        tensor = super().__getitem__(extractor).as_subclass(LabelTensor)
+        tensor._labels = labels
+        return tensor
 
     def __str__(self):
         """
@@ -272,39 +254,53 @@ def cat(tensors, dim=0):
             return []
         if len(tensors) == 1 or isinstance(tensors, LabelTensor):
             return tensors[0]
+
         # Perform cat on tensors
         new_tensor = torch.cat(tensors, dim=dim)
 
-        # Update labels
-        labels = LabelTensor.__create_labels_cat(tensors, dim)
+        # --------- Start definition auxiliary function ------
+        # Compute and update labels
+        def create_labels_cat(tensors, dim, tensor_shape):
+            stored_labels = [tensor.stored_labels for tensor in tensors]
+            keys = stored_labels[0].keys()
+
+            if any(not all(stored_labels[i][k] == stored_labels[0][k] for i in
+                        range(len(stored_labels))) for k in keys if k != dim):
+                raise RuntimeError('tensors must have the same shape and dof')
+
+            # Copy labels from the first tensor and update the 'dof' for dimension `dim`
+            labels = copy(stored_labels[0])
+            if dim in labels:
+                labels_list = [tensor[dim]['dof'] for tensor in stored_labels]
+                last_dim_dof = range(tensor_shape[dim]) if all(isinstance(label, range)
+                                    for label in labels_list) else sum(labels_list, [])
+                labels[dim]['dof'] = last_dim_dof
+            return labels
+        # --------- End definition auxiliary function ------
 
-        return LabelTensor.__internal_init__(new_tensor, labels,
-                                             tensors[0].dim_names)
+        # Update labels
+        if dim in tensors[0].stored_labels.keys():
+            new_tensor_shape = new_tensor.shape
+            labels = create_labels_cat(tensors, dim, new_tensor_shape)
+        else:
+            labels = tensors[0].stored_labels
+        new_tensor._labels = labels
+        return new_tensor
 
     @staticmethod
-    def __create_labels_cat(tensors, dim):
-        # Check if names and dof of the labels are the same in all dimensions
-        # except in dim
-        stored_labels = [tensor.stored_labels for tensor in tensors]
-
-        # check if:
-        # - labels dict have same keys
-        # - all labels are the same expect for dimension dim
-        if not all(
-                all(stored_labels[i][k] == stored_labels[0][k]
-                    for i in range(len(stored_labels)))
-                for k in stored_labels[0].keys() if k != dim):
-            raise RuntimeError('tensors must have the same shape and dof')
-
-        labels = {k: copy(v) for k, v in tensors[0].stored_labels.items()}
-        if dim in labels.keys():
-            last_dim_dof = [i for j in stored_labels for i in j[dim]['dof']]
-            labels[dim]['dof'] = last_dim_dof
-        return labels
+    def stack(tensors):
+        new_tensor = torch.stack(tensors)
+        labels = tensors[0]._labels
+        labels = {key + 1: value for key, value in labels.items()}
+        if full_labels:
+            new_tensor.labels = labels
+        else:
+            new_tensor._labels = labels
+        return new_tensor
 
     def requires_grad_(self, mode=True):
         lt = super().requires_grad_(mode)
-        lt.labels = self._labels
+        lt._labels = self._labels
         return lt
 
     @property
@@ -316,10 +312,9 @@ def to(self, *args, **kwargs):
         Performs Tensor dtype and/or device conversion. For more details, see
         :meth:`torch.Tensor.to`.
         """
-        tmp = super().to(*args, **kwargs)
-        new = self.__class__.clone(self)
-        new.data = tmp.data
-        return new
+        lt = super().to(*args, **kwargs)
+        lt._labels = self._labels
+        return lt
 
     def clone(self, *args, **kwargs):
         """
@@ -329,8 +324,7 @@ def clone(self, *args, **kwargs):
         :return: A copy of the tensor.
         :rtype: LabelTensor
         """
-        labels = {k: copy(v) for k, v in self._labels.items()}
-        out = LabelTensor(super().clone(*args, **kwargs), labels)
+        out = LabelTensor(super().clone(*args, **kwargs), deepcopy(self._labels))
         return out
 
     @staticmethod
@@ -348,7 +342,7 @@ def summation(tensors):
             raise RuntimeError('Tensors must have the same shape and labels')
 
         last_dim_labels = []
-        data = torch.zeros(tensors[0].tensor.shape)
+        data = torch.zeros(tensors[0].tensor.shape).to(tensors[0].device)
         for tensor in tensors:
             data += tensor.tensor
             last_dim_labels.append(tensor.labels)
@@ -396,82 +390,114 @@ def vstack(label_tensors):
         """
         return LabelTensor.cat(label_tensors, dim=0)
 
+        # ---------------------- Start auxiliary function definition -----
+        # This method is used to update labels
+    def _update_single_label(self, old_labels, to_update_labels, index, dim,
+                             to_update_dim):
+            """
+            TODO
+                :param old_labels: labels from which retrieve data
+                :param to_update_labels: labels to update
+                :param index: index of dof to retain
+                :param dim: label index
+                :return:
+            """
+            old_dof = old_labels[to_update_dim]['dof']
+            if isinstance(index, slice):
+                to_update_labels.update({
+                    dim: {
+                        'dof': old_dof[index],
+                        'name': old_labels[dim]['name']
+                    }
+                })
+                return
+            if isinstance(index, int):
+                index = [index]
+            if isinstance(index, (list, torch.Tensor)):
+                to_update_labels.update({
+                    dim: {
+                        'dof': [old_dof[i] for i in index] if isinstance(old_dof, list) else index,
+                        'name': old_labels[dim]['name']
+                    }
+                })
+                return
+            raise NotImplementedError(f'Getitem not implemented for '
+                                      f'{type(index)} values')
+        # ---------------------- End auxiliary function definition -----
+
+
     def __getitem__(self, index):
         """
         TODO: Complete docstring
         :param index:
         :return:
         """
-        if isinstance(index,
-                      str) or (isinstance(index, (tuple, list))
-                               and all(isinstance(a, str) for a in index)):
+        # Index are str --> call extract
+        if isinstance(index, str) or (isinstance(index, (tuple, list))
+                                      and all(
+                    isinstance(a, str) for a in index)):
             return self.extract(index)
 
+        # Store important variables
         selected_lt = super().__getitem__(index)
+        stored_labels = self._labels
+        labels = copy(stored_labels)
 
-        if isinstance(index, (int, slice)):
+        # Put here because it is the most common case (int as index).
+        # Used by DataLoader -> put here for efficiency purpose
+        if isinstance(index, list):
+            if 0 in labels.keys():
+                self._update_single_label(stored_labels, labels, index,
+                                          0, 0)
+            selected_lt._labels = labels
+            return selected_lt
+
+        if isinstance(index, int):
+            labels.pop(0, None)
+            labels = {key - 1 if key > 0 else key: value for key, value in
+                      labels.items()}
+            selected_lt._labels = labels
+            return selected_lt
+
+        if not isinstance(index, (tuple, torch.Tensor)):
             index = [index]
 
+        # Ellipsis are used to perform operation on the last dimension
         if index[0] == Ellipsis:
-            index = [slice(None)] * (self.ndim - 1) + [index[1]]
-
-        if hasattr(self, "labels"):
-            labels = {k: copy(v) for k, v in self.stored_labels.items()}
-            for j, idx in enumerate(index):
-                if isinstance(idx, int):
+            if len(self.shape) in labels:
+                self._update_single_label(stored_labels, labels, index, 0, 0)
+            selected_lt._labels = labels
+            return selected_lt
+
+        i = 0
+        for j, idx in enumerate(index):
+            if j in self.stored_labels.keys():
+                if isinstance(idx, int) or (
+                        isinstance(idx, torch.Tensor) and idx.ndim == 0):
                     selected_lt = selected_lt.unsqueeze(j)
-                if j in labels.keys() and idx != slice(None):
-                    self._update_single_label(labels, labels, idx, j)
-            selected_lt = LabelTensor.__internal_init__(selected_lt, labels,
-                                                        self.dim_names)
+                if idx != slice(None):
+                    self._update_single_label(stored_labels, labels, idx, j, i)
+            else:
+                if isinstance(idx, int):
+                    labels = {key - 1 if key > j else key:
+                                  value for key, value in labels.items()}
+                    continue
+            i += 1
+        selected_lt._labels = labels
         return selected_lt
 
-    @staticmethod
-    def _update_single_label(old_labels, to_update_labels, index, dim):
-        """
-        TODO
-        :param old_labels: labels from which retrieve data
-        :param to_update_labels: labels to update
-        :param index: index of dof to retain
-        :param dim: label index
-        :return:
-        """
-        old_dof = old_labels[dim]['dof']
-        if not isinstance(
-                index,
-            (int, slice)) and len(index) == len(old_dof) and isinstance(
-                old_dof, range):
-            return
-        if isinstance(index, torch.Tensor):
-            index = index.nonzero(
-                as_tuple=True
-            )[0] if index.dtype == torch.bool else index.tolist()
-        if isinstance(index, list):
-            to_update_labels.update({
-                dim: {
-                    'dof': [old_dof[i] for i in index],
-                    'name': old_labels[dim]['name']
-                }
-            })
-        else:
-            to_update_labels.update(
-                {dim: {
-                    'dof': old_dof[index],
-                    'name': old_labels[dim]['name']
-                }})
-
     def sort_labels(self, dim=None):
-
         def arg_sort(lst):
             return sorted(range(len(lst)), key=lambda x: lst[x])
-
         if dim is None:
             dim = self.ndim - 1
+        if self.shape[dim] == 1:
+            return self
         labels = self.stored_labels[dim]['dof']
         sorted_index = arg_sort(labels)
         indexer = [slice(None)] * self.ndim
         indexer[dim] = sorted_index
-        return self.__getitem__(indexer)
+        return self.__getitem__(tuple(indexer))
 
     def __deepcopy__(self, memo):
         cls = self.__class__
@@ -480,10 +506,16 @@ def __deepcopy__(self, memo):
 
     def permute(self, *dims):
         tensor = super().permute(*dims)
-        stored_labels = self.stored_labels
+        labels = self._labels
         keys_list = list(*dims)
         labels = {
-            keys_list.index(k): copy(stored_labels[k])
-            for k in stored_labels.keys()
+            keys_list.index(k): labels[k]
+            for k in labels.keys()
         }
-        return LabelTensor.__internal_init__(tensor, labels, self.dim_names)
+        tensor._labels = labels
+        return tensor
+
+    def detach(self):
+        lt = super().detach()
+        lt._labels = self.stored_labels
+        return lt
\ No newline at end of file
diff --git a/pina/model/layers/lowrank_layer.py b/pina/model/layers/lowrank_layer.py
index 80fb43e4e..c36d2795a 100644
--- a/pina/model/layers/lowrank_layer.py
+++ b/pina/model/layers/lowrank_layer.py
@@ -119,6 +119,7 @@ def forward(self, x, coords):
         :rtype: torch.Tensor
         """
         # extract basis
+        coords = coords.as_subclass(torch.Tensor)
         basis = self._basis(coords)
         # reshape [B, N, D, 2*rank]
         shape = list(basis.shape[:-1]) + [-1, 2 * self.rank]
diff --git a/pina/model/network.py b/pina/model/network.py
index 6fde8039c..aed3dff3d 100644
--- a/pina/model/network.py
+++ b/pina/model/network.py
@@ -29,7 +29,8 @@ class is used internally in PINA to convert
         # check model consistency
         check_consistency(model, nn.Module)
         check_consistency(input_variables, str)
-        check_consistency(output_variables, str)
+        if output_variables is not None:
+            check_consistency(output_variables, str)
 
         self._model = model
         self._input_variables = input_variables
@@ -67,16 +68,15 @@ def forward(self, x):
         # in case `input_variables = []` all points are used
         if self._input_variables:
             x = x.extract(self._input_variables)
-
         # extract features and append
         for feature in self._extra_features:
             x = x.append(feature(x))
 
         # perform forward pass + converting to LabelTensor
-        output = self._model(x).as_subclass(LabelTensor)
-
-        # set the labels for LabelTensor
-        output.labels = self._output_variables
+        x = x.as_subclass(torch.Tensor)
+        output = self._model(x)
+        if self._output_variables is not None:
+            output = LabelTensor(output, self._output_variables)
 
         return output
 
@@ -97,15 +97,9 @@ def forward_map(self, x):
             This function does not extract the input variables, all the variables
             are used for both tensors. Output variables are correctly applied.
         """
-        # convert LabelTensor s to torch.Tensor s
-        x = list(map(lambda x: x.as_subclass(torch.Tensor), x))
 
         # perform forward pass (using torch.Tensor) + converting to LabelTensor
-        output = self._model(x).as_subclass(LabelTensor)
-
-        # set the labels for LabelTensor
-        output.labels = self._output_variables
-
+        output = LabelTensor(self._model(x.tensor), self._output_variables)
         return output
 
     @property
diff --git a/pina/operators.py b/pina/operators.py
index 193243a69..8dbfb8c76 100644
--- a/pina/operators.py
+++ b/pina/operators.py
@@ -63,11 +63,9 @@ def grad_scalar_output(output_, input_, d):
             retain_graph=True,
             allow_unused=True,
         )[0]
-
-        gradients.labels = input_.labels
-        gradients = gradients.extract(d)
+        gradients.labels = input_.stored_labels
+        gradients = gradients[..., [input_.labels.index(i) for i in d]]
         gradients.labels = [f"d{output_fieldname}d{i}" for i in d]
-
         return gradients
 
     if not isinstance(input_, LabelTensor):
@@ -216,7 +214,9 @@ def scalar_laplace(output_, input_, components, d):
             to_append_tensors = []
             for i, label in enumerate(grad_output.labels):
                 gg = grad(grad_output, input_, d=d, components=[label])
-                to_append_tensors.append(gg.extract([gg.labels[i]]))
+                gg = gg.extract([gg.labels[i]])
+
+                to_append_tensors.append(gg)
             labels = [f"dd{components[0]}"]
             result = LabelTensor.summation(tensors=to_append_tensors)
             result.labels = labels
diff --git a/pina/problem/abstract_problem.py b/pina/problem/abstract_problem.py
index 6897fbb74..5f424cf5e 100644
--- a/pina/problem/abstract_problem.py
+++ b/pina/problem/abstract_problem.py
@@ -90,10 +90,9 @@ def input_variables(self):
             variables += self.spatial_variables
         if hasattr(self, "temporal_variable"):
             variables += self.temporal_variable
-        if hasattr(self, "unknown_parameters"):
+        if hasattr(self, "parameters"):
             variables += self.parameters
-        if hasattr(self, "custom_variables"):
-            variables += self.custom_variables
+
 
         return variables
 
@@ -170,7 +169,6 @@ def discretise_domain(self,
                     f"Wrong variables for sampling. Variables ",
                     f"should be in {self.input_variables}.",
                 )
-
         # check correct location
         if locations == "all":
             locations = [
diff --git a/pina/problem/inverse_problem.py b/pina/problem/inverse_problem.py
index 51cbd3ca6..e54495a67 100644
--- a/pina/problem/inverse_problem.py
+++ b/pina/problem/inverse_problem.py
@@ -1,7 +1,6 @@
 """Module for the ParametricProblem class"""
-
+import torch
 from abc import abstractmethod
-
 from .abstract_problem import AbstractProblem
 
 
diff --git a/pina/solvers/graph.py b/pina/solvers/graph.py
deleted file mode 100644
index 9af04e76f..000000000
--- a/pina/solvers/graph.py
+++ /dev/null
@@ -1,34 +0,0 @@
-from .supervised import SupervisedSolver
-from ..graph import Graph
-
-
-class GraphSupervisedSolver(SupervisedSolver):
-
-    def __init__(
-        self,
-        problem,
-        model,
-        nodes_coordinates,
-        nodes_data,
-        loss=None,
-        optimizer=None,
-        scheduler=None):
-        super().__init__(problem, model, loss, optimizer, scheduler)
-        if isinstance(nodes_coordinates, str):
-            self._nodes_coordinates = [nodes_coordinates]
-        else:
-            self._nodes_coordinates = nodes_coordinates
-        if isinstance(nodes_data, str):
-            self._nodes_data = nodes_data
-        else:
-            self._nodes_data = nodes_data
-
-    def forward(self, input):
-        input_coords = input.extract(self._nodes_coordinates)
-        input_data = input.extract(self._nodes_data)
-
-        if not isinstance(input, Graph):
-            input = Graph.build('radius', nodes_coordinates=input_coords, nodes_data=input_data, radius=0.2)
-        g = self.model(input.data, edge_index=input.data.edge_index)
-        g.labels = {1: {'name': 'output', 'dof': ['u']}}
-        return g
diff --git a/pina/solvers/pinns/basepinn.py b/pina/solvers/pinns/basepinn.py
index a8315b259..66f4d14c3 100644
--- a/pina/solvers/pinns/basepinn.py
+++ b/pina/solvers/pinns/basepinn.py
@@ -1,14 +1,15 @@
 """ Module for PINN """
 
-import sys
 from abc import ABCMeta, abstractmethod
 import torch
-
-from ...solvers.solver import SolverInterface
-from pina.utils import check_consistency
-from pina.loss.loss_interface import LossInterface
-from pina.problem import InverseProblem
 from torch.nn.modules.loss import _Loss
+from ...condition import InputOutputPointsCondition
+from ...solvers.solver import SolverInterface
+from ...utils import check_consistency
+from ...loss.loss_interface import LossInterface
+from ...problem import InverseProblem
+from ...condition import DomainEquationCondition
+from ...optim import TorchOptimizer, TorchScheduler
 
 torch.pi = torch.acos(torch.zeros(1)).item() * 2  # which is 3.1415927410125732
 
@@ -25,13 +26,14 @@ class PINNInterface(SolverInterface, metaclass=ABCMeta):
     to the user to choose which problem the implemented solver inheriting from
     this class is suitable for.
     """
-
+    accepted_condition_types = [DomainEquationCondition.condition_type[0],
+                                InputOutputPointsCondition.condition_type[0]]
     def __init__(
             self,
             models,
             problem,
             optimizers,
-            optimizers_kwargs,
+            schedulers,
             extra_features,
             loss,
     ):
@@ -53,11 +55,20 @@ def __init__(
         :param torch.nn.Module loss: The loss function used as minimizer,
             default :class:`torch.nn.MSELoss`.
         """
+        if optimizers is None:
+            optimizers = TorchOptimizer(torch.optim.Adam, lr=0.001)
+
+        if schedulers is None:
+            schedulers = TorchScheduler(torch.optim.lr_scheduler.ConstantLR)
+
+        if loss is None:
+            loss = torch.nn.MSELoss()
+
         super().__init__(
             models=models,
             problem=problem,
             optimizers=optimizers,
-            optimizers_kwargs=optimizers_kwargs,
+            schedulers=schedulers,
             extra_features=extra_features,
         )
 
@@ -85,7 +96,12 @@ def __init__(
         # variable will be stored with name = self.__logged_metric
         self.__logged_metric = None
 
-    def training_step(self, batch, _):
+        self._model = self._pina_models[0]
+        self._optimizer = self._pina_optimizers[0]
+        self._scheduler = self._pina_schedulers[0]
+
+
+    def training_step(self, batch):
         """
         The Physics Informed Solver Training Step. This function takes care
         of the physics informed training step, and it must not be override
@@ -99,53 +115,68 @@ def training_step(self, batch, _):
         :rtype: LabelTensor
         """
 
-        condition_losses = []
-        condition_idx = batch["condition"]
-
-        for condition_id in range(condition_idx.min(), condition_idx.max() + 1):
+        condition_loss = []
+        for condition_name, points in batch:
+            if 'output_points' in points:
+                input_pts, output_pts = points['input_points'], points['output_points']
 
-            condition_name = self._dataloader.condition_names[condition_id]
-            condition = self.problem.conditions[condition_name]
-            pts = batch["pts"]
-            # condition name is logged (if logs enabled)
-            self.__logged_metric = condition_name
-
-            if len(batch) == 2:
-                samples = pts[condition_idx == condition_id]
-                loss = self.loss_phys(samples, condition.equation)
-            elif len(batch) == 3:
-                samples = pts[condition_idx == condition_id]
-                ground_truth = batch["output"][condition_idx == condition_id]
-                loss = self.loss_data(samples, ground_truth)
+                loss_ = self.loss_data(input_pts=input_pts, output_pts=output_pts)
+                condition_loss.append(loss_.as_subclass(torch.Tensor))
             else:
-                raise ValueError("Batch size not supported")
+                input_pts = points['input_points']
 
-            # add condition losses for each epoch
-            condition_losses.append(loss * condition.data_weight)
+                condition = self.problem.conditions[condition_name]
 
+                loss_ = self.loss_phys(input_pts.requires_grad_(), condition.equation)
+                condition_loss.append(loss_.as_subclass(torch.Tensor))
+            condition_loss.append(loss_.as_subclass(torch.Tensor))
         # clamp unknown parameters in InverseProblem (if needed)
         self._clamp_params()
+        loss = sum(condition_loss)
+        self.log('train_loss', loss, prog_bar=True, on_epoch=True,
+                 logger=True, batch_size=self.get_batch_size(batch),
+                 sync_dist=True)
 
-        # total loss (must be a torch.Tensor), and logs
-        total_loss = sum(condition_losses)
-        self.save_logs_and_release()
-        return total_loss.as_subclass(torch.Tensor)
+        return loss
 
-    def loss_data(self, input_tensor, output_tensor):
+    def validation_step(self, batch):
+        """
+        TODO: add docstring
+        """
+        condition_loss = []
+        for condition_name, points in batch:
+            if 'output_points' in points:
+                input_pts, output_pts = points['input_points'], points['output_points']
+                loss_ = self.loss_data(input_pts=input_pts, output_pts=output_pts)
+                condition_loss.append(loss_.as_subclass(torch.Tensor))
+            else:
+                input_pts = points['input_points']
+
+                condition = self.problem.conditions[condition_name]
+                with torch.set_grad_enabled(True):
+                    loss_ = self.loss_phys(input_pts.requires_grad_(), condition.equation)
+                condition_loss.append(loss_.as_subclass(torch.Tensor))
+            condition_loss.append(loss_.as_subclass(torch.Tensor))
+        # clamp unknown parameters in InverseProblem (if needed)
+
+        loss = sum(condition_loss)
+        self.log('val_loss', loss, on_epoch=True, prog_bar=True,
+                 logger=True, batch_size=self.get_batch_size(batch),
+                 sync_dist=True)
+
+    def loss_data(self, input_pts, output_pts):
         """
         The data loss for the PINN solver. It computes the loss between
         the network output against the true solution. This function
         should not be override if not intentionally.
 
-        :param LabelTensor input_tensor: The input to the neural networks.
-        :param LabelTensor output_tensor: The true solution to compare the
+        :param LabelTensor input_pts: The input to the neural networks.
+        :param LabelTensor output_pts: The true solution to compare the
             network solution.
         :return: The residual loss averaged on the input coordinates
         :rtype: torch.Tensor
         """
-        loss_value = self.loss(self.forward(input_tensor), output_tensor)
-        self.store_log(loss_value=float(loss_value))
-        return self.loss(self.forward(input_tensor), output_tensor)
+        return self._loss(self.forward(input_pts), output_pts)
 
     @abstractmethod
     def loss_phys(self, samples, equation):
@@ -196,13 +227,17 @@ def store_log(self, loss_value):
         :param str name: The name of the loss.
         :param torch.Tensor loss_value: The value of the loss.
         """
+        batch_size = self.trainer.data_module.batch_size \
+            if self.trainer.data_module.batch_size is not None else 999
+
         self.log(
             self.__logged_metric + "_loss",
             loss_value,
             prog_bar=True,
             logger=True,
             on_epoch=True,
-            on_step=False,
+            on_step=True,
+            batch_size=batch_size,
         )
         self.__logged_res_losses.append(loss_value)
 
diff --git a/pina/solvers/pinns/pinn.py b/pina/solvers/pinns/pinn.py
index 15f908182..088820202 100644
--- a/pina/solvers/pinns/pinn.py
+++ b/pina/solvers/pinns/pinn.py
@@ -9,10 +9,8 @@
         _LRScheduler as LRScheduler,
     )  # torch < 2.0
 
-from torch.optim.lr_scheduler import ConstantLR
 
 from .basepinn import PINNInterface
-from pina.utils import check_consistency
 from pina.problem import InverseProblem
 
 
@@ -56,16 +54,16 @@ class PINN(PINNInterface):
         DOI: `10.1038 <https://doi.org/10.1038/s42254-021-00314-5>`_.
     """
 
+    __name__ = 'PINN'
+
     def __init__(
         self,
         problem,
         model,
         extra_features=None,
-        loss=torch.nn.MSELoss(),
-        optimizer=torch.optim.Adam,
-        optimizer_kwargs={"lr": 0.001},
-        scheduler=ConstantLR,
-        scheduler_kwargs={"factor": 1, "total_iters": 0},
+        loss=None,
+        optimizer=None,
+        scheduler=None,
     ):
         """
         :param AbstractProblem problem: The formulation of the problem.
@@ -82,20 +80,15 @@ def __init__(
         :param dict scheduler_kwargs: LR scheduler constructor keyword args.
         """
         super().__init__(
-            models=[model],
+            models=model,
             problem=problem,
-            optimizers=[optimizer],
-            optimizers_kwargs=[optimizer_kwargs],
+            optimizers=optimizer,
+            schedulers=scheduler,
             extra_features=extra_features,
             loss=loss,
         )
 
-        # check consistency
-        check_consistency(scheduler, LRScheduler, subclass=True)
-        check_consistency(scheduler_kwargs, dict)
-
         # assign variables
-        self._scheduler = scheduler(self.optimizers[0], **scheduler_kwargs)
         self._neural_net = self.models[0]
 
     def forward(self, x):
@@ -126,9 +119,8 @@ def loss_phys(self, samples, equation):
         """
         residual = self.compute_residual(samples=samples, equation=equation)
         loss_value = self.loss(
-            torch.zeros_like(residual, requires_grad=True), residual
+            torch.zeros_like(residual), residual
         )
-        self.store_log(loss_value=float(loss_value))
         return loss_value
 
     def configure_optimizers(self):
@@ -141,16 +133,21 @@ def configure_optimizers(self):
         """
         # if the problem is an InverseProblem, add the unknown parameters
         # to the parameters that the optimizer needs to optimize
+
+
+        self._optimizer.hook(self._model.parameters())
         if isinstance(self.problem, InverseProblem):
-            self.optimizers[0].add_param_group(
-                {
-                    "params": [
-                        self._params[var]
-                        for var in self.problem.unknown_variables
-                    ]
-                }
-            )
-        return self.optimizers, [self.scheduler]
+            self._optimizer.optimizer_instance.add_param_group(
+                    {
+                        "params": [
+                            self._params[var]
+                            for var in self.problem.unknown_variables
+                        ]
+                    }
+                )
+        self._scheduler.hook(self._optimizer)
+        return ([self._optimizer.optimizer_instance],
+                [self._scheduler.scheduler_instance])
 
     @property
     def scheduler(self):
diff --git a/pina/solvers/solver.py b/pina/solvers/solver.py
index e00bc8d59..3a8f400c2 100644
--- a/pina/solvers/solver.py
+++ b/pina/solvers/solver.py
@@ -2,7 +2,7 @@
 
 from abc import ABCMeta, abstractmethod
 from ..model.network import Network
-import pytorch_lightning
+import lightning
 from ..utils import check_consistency
 from ..problem import AbstractProblem
 from ..optim import Optimizer, Scheduler
@@ -10,7 +10,8 @@
 import sys
 
 
-class SolverInterface(pytorch_lightning.LightningModule, metaclass=ABCMeta):
+
+class SolverInterface(lightning.pytorch.LightningModule, metaclass=ABCMeta):
     """
     Solver base class. This class inherits is a wrapper of
     LightningModule class, inheriting all the
@@ -83,7 +84,6 @@ def __init__(self,
                              " optimizers.")
 
         # extra features handling
-
         self._pina_models = models
         self._pina_optimizers = optimizers
         self._pina_schedulers = schedulers
@@ -94,7 +94,7 @@ def forward(self, *args, **kwargs):
         pass
 
     @abstractmethod
-    def training_step(self, batch, batch_idx):
+    def training_step(self, batch):
         pass
 
     @abstractmethod
@@ -138,8 +138,16 @@ def _check_solver_consistency(self, problem):
         TODO
         """
         for _, condition in problem.conditions.items():
-            if not set(self.accepted_condition_types).issubset(
-                    condition.condition_type):
+            if not set(condition.condition_type).issubset(
+                    set(self.accepted_condition_types)):
                 raise ValueError(
-                    f'{self.__name__} support only dose not support condition '
+                    f'{self.__name__} dose not support condition '
                     f'{condition.condition_type}')
+
+    @staticmethod
+    def get_batch_size(batch):
+        # Assuming batch is your custom Batch object
+        batch_size = 0
+        for data in batch:
+            batch_size += len(data[1]['input_points'])
+        return batch_size
\ No newline at end of file
diff --git a/pina/solvers/supervised.py b/pina/solvers/supervised.py
index 62fc99149..d978fc37c 100644
--- a/pina/solvers/supervised.py
+++ b/pina/solvers/supervised.py
@@ -1,12 +1,14 @@
 """ Module for SupervisedSolver """
-
 import torch
+from pytorch_lightning.utilities.types import STEP_OUTPUT
+from sympy.strategies.branch import condition
 from torch.nn.modules.loss import _Loss
 from ..optim import TorchOptimizer, TorchScheduler
 from .solver import SolverInterface
 from ..label_tensor import LabelTensor
 from ..utils import check_consistency
 from ..loss.loss_interface import LossInterface
+from ..condition import InputOutputPointsCondition
 
 
 class SupervisedSolver(SolverInterface):
@@ -37,7 +39,7 @@ class SupervisedSolver(SolverInterface):
     we are seeking to approximate multiple (discretised) functions given
     multiple (discretised) input functions.
     """
-    accepted_condition_types = ['supervised']
+    accepted_condition_types = [InputOutputPointsCondition.condition_type[0]]
     __name__ = 'SupervisedSolver'
 
     def __init__(self,
@@ -46,7 +48,8 @@ def __init__(self,
                  loss=None,
                  optimizer=None,
                  scheduler=None,
-                 extra_features=None):
+                 extra_features=None,
+                 use_lt=True):
         """
         :param AbstractProblem problem: The formualation of the problem.
         :param torch.nn.Module model: The neural network model to use.
@@ -72,14 +75,19 @@ def __init__(self,
                          problem=problem,
                          optimizers=optimizer,
                          schedulers=scheduler,
-                         extra_features=extra_features)
+                         extra_features=extra_features,
+                         use_lt=use_lt)
 
         # check consistency
-        check_consistency(loss, (LossInterface, _Loss), subclass=False)
+        check_consistency(loss, (LossInterface, _Loss, torch.nn.Module),
+                          subclass=False)
         self._loss = loss
         self._model = self._pina_models[0]
         self._optimizer = self._pina_optimizers[0]
         self._scheduler = self._pina_schedulers[0]
+        self.validation_condition_losses = {
+            k: {'loss': [],
+                'count': []} for k in self.problem.conditions.keys()}
 
     def forward(self, x):
         """Forward pass implementation for the solver.
@@ -105,7 +113,7 @@ def configure_optimizers(self):
         return ([self._optimizer.optimizer_instance],
                 [self._scheduler.scheduler_instance])
 
-    def training_step(self, batch, batch_idx):
+    def training_step(self, batch):
         """Solver training step.
 
         :param batch: The batch element in the dataloader.
@@ -115,32 +123,36 @@ def training_step(self, batch, batch_idx):
         :return: The sum of the loss functions.
         :rtype: LabelTensor
         """
-        condition_idx = batch.supervised.condition_indices
-
-        for condition_id in range(condition_idx.min(), condition_idx.max() + 1):
-
-            condition_name = self._dataloader.condition_names[condition_id]
-            condition = self.problem.conditions[condition_name]
-            pts = batch.supervised.input_points
-            out = batch.supervised.output_points
-            if condition_name not in self.problem.conditions:
-                raise RuntimeError("Something wrong happened.")
+        condition_loss = []
+        for condition_name, points in batch:
+            input_pts, output_pts = points['input_points'], points['output_points']
+            loss_ = self.loss_data(input_pts=input_pts, output_pts=output_pts)
+            condition_loss.append(loss_.as_subclass(torch.Tensor))
+        loss = sum(condition_loss)
+        self.log('train_loss', loss, on_step=True, on_epoch=True, prog_bar=True, logger=True,
+                 batch_size=self.get_batch_size(batch), sync_dist=True)
+        return loss
 
-            # for data driven mode
-            if not hasattr(condition, "output_points"):
-                raise NotImplementedError(
-                    f"{type(self).__name__} works only in data-driven mode.")
-            output_pts = out[condition_idx == condition_id]
-            input_pts = pts[condition_idx == condition_id]
+    def validation_step(self, batch):
+        """
+        Solver validation step.
+        """
+        condition_loss = []
+        for condition_name, points in batch:
+            input_pts, output_pts = points['input_points'], points['output_points']
+            loss_ = self.loss_data(input_pts=input_pts, output_pts=output_pts)
+            condition_loss.append(loss_.as_subclass(torch.Tensor))
+        loss = sum(condition_loss)
+        self.log('val_loss', loss, prog_bar=True, logger=True,
+                 batch_size=self.get_batch_size(batch), sync_dist=True)
 
-            input_pts.labels = pts.labels
-            output_pts.labels = out.labels
 
-            loss = self.loss_data(input_pts=input_pts, output_pts=output_pts)
-            loss = loss.as_subclass(torch.Tensor)
+    def test_step(self, batch, batch_idx) -> STEP_OUTPUT:
+        """
+        Solver test step.
+        """
 
-        self.log("mean_loss", float(loss), prog_bar=True, logger=True)
-        return loss
+        raise NotImplementedError("Test step not implemented yet.")
 
     def loss_data(self, input_pts, output_pts):
         """
diff --git a/pina/trainer.py b/pina/trainer.py
index 58c66f67c..a7c5c3513 100644
--- a/pina/trainer.py
+++ b/pina/trainer.py
@@ -1,20 +1,21 @@
 """ Trainer module. """
-
+import warnings
 import torch
-import pytorch_lightning
+import lightning
 from .utils import check_consistency
 from .data import PinaDataModule
 from .solvers.solver import SolverInterface
 
 
-class Trainer(pytorch_lightning.Trainer):
+class Trainer(lightning.pytorch.Trainer):
 
     def __init__(self,
                  solver,
                  batch_size=None,
                  train_size=.7,
                  test_size=.2,
-                 eval_size=.1,
+                 val_size=.1,
+                 predict_size=.0,
                  **kwargs):
         """
         PINA Trainer class for costumizing every aspect of training via flags.
@@ -39,11 +40,13 @@ def __init__(self,
             check_consistency(batch_size, int)
         self.train_size = train_size
         self.test_size = test_size
-        self.eval_size = eval_size
+        self.val_size = val_size
+        self.predict_size = predict_size
         self.solver = solver
         self.batch_size = batch_size
-        self._create_loader()
         self._move_to_device()
+        self.data_module = None
+        self._create_loader()
 
     def _move_to_device(self):
         device = self._accelerator_connector._parallel_devices[0]
@@ -64,34 +67,34 @@ def _create_loader(self):
         if not self.solver.problem.collector.full:
             error_message = '\n'.join([
                 f"""{" " * 13} ---> Condition {key} {"sampled" if value else
-                    "not sampled"}""" for key, value in
+                "not sampled"}""" for key, value in
                 self._solver.problem.collector._is_conditions_ready.items()
             ])
             raise RuntimeError('Cannot create Trainer if not all conditions '
                                'are sampled. The Trainer got the following:\n'
                                f'{error_message}')
-        devices = self._accelerator_connector._parallel_devices
-
-        if len(devices) > 1:
-            raise RuntimeError("Parallel training is not supported yet.")
-
-        device = devices[0]
-
-        data_module = PinaDataModule(problem=self.solver.problem,
-                                     device=device,
-                                     train_size=self.train_size,
-                                     test_size=self.test_size,
-                                     val_size=self.eval_size)
-        data_module.setup()
-        self._loader = data_module.train_dataloader()
+        self.data_module = PinaDataModule(collector=self.solver.problem.collector,
+                                          train_size=self.train_size,
+                                          test_size=self.test_size,
+                                          val_size=self.val_size,
+                                          predict_size=self.predict_size,
+                                          batch_size=self.batch_size,)
 
     def train(self, **kwargs):
         """
         Train the solver method.
         """
         return super().fit(self.solver,
-                           train_dataloaders=self._loader,
-                           **kwargs)
+                               datamodule=self.data_module,
+                               **kwargs)
+
+    def test(self, **kwargs):
+        """
+        Test the solver method.
+        """
+        return super().test(self.solver,
+                            datamodule=self.data_module,
+                            **kwargs)
 
     @property
     def solver(self):
diff --git a/setup.py b/setup.py
index 5a2ebc85e..b67323bdf 100644
--- a/setup.py
+++ b/setup.py
@@ -15,7 +15,8 @@
 KEYWORDS = 'machine-learning deep-learning modeling pytorch ode neural-networks differential-equations pde hacktoberfest pinn physics-informed physics-informed-neural-networks neural-operators equation-learning lightining'
 
 REQUIRED = [
-    'numpy', 'matplotlib', 'torch', 'lightning', 'pytorch_lightning', 'torch_geometric', 'torch-cluster'
+    'numpy', 'matplotlib', 'torch', 'lightning', 'torch_geometric',
+    'torch-cluster', 'pytorch_lightning',
 ]
 
 EXTRAS = {
diff --git a/tests/test_dataset.py b/tests/test_dataset.py
deleted file mode 100644
index 87fd9a15b..000000000
--- a/tests/test_dataset.py
+++ /dev/null
@@ -1,227 +0,0 @@
-import math
-import torch
-from pina.data import SamplePointDataset, SupervisedDataset, PinaDataModule, \
-    UnsupervisedDataset
-from pina.data import PinaDataLoader
-from pina import LabelTensor, Condition
-from pina.equation import Equation
-from pina.domain import CartesianDomain
-from pina.problem import SpatialProblem, AbstractProblem
-from pina.operators import laplacian
-from pina.equation.equation_factory import FixedValue
-from pina.graph import Graph
-
-
-def laplace_equation(input_, output_):
-    force_term = (torch.sin(input_.extract(['x']) * torch.pi) *
-                  torch.sin(input_.extract(['y']) * torch.pi))
-    delta_u = laplacian(output_.extract(['u']), input_)
-    return delta_u - force_term
-
-
-my_laplace = Equation(laplace_equation)
-in_ = LabelTensor(torch.tensor([[0., 1.]]), ['x', 'y'])
-out_ = LabelTensor(torch.tensor([[0.]]), ['u'])
-in2_ = LabelTensor(torch.rand(60, 2), ['x', 'y'])
-out2_ = LabelTensor(torch.rand(60, 1), ['u'])
-
-
-class Poisson(SpatialProblem):
-    output_variables = ['u']
-    spatial_domain = CartesianDomain({'x': [0, 1], 'y': [0, 1]})
-
-    conditions = {
-        'gamma1':
-        Condition(domain=CartesianDomain({
-            'x': [0, 1],
-            'y': 1
-        }),
-                  equation=FixedValue(0.0)),
-        'gamma2':
-        Condition(domain=CartesianDomain({
-            'x': [0, 1],
-            'y': 0
-        }),
-                  equation=FixedValue(0.0)),
-        'gamma3':
-        Condition(domain=CartesianDomain({
-            'x': 1,
-            'y': [0, 1]
-        }),
-                  equation=FixedValue(0.0)),
-        'gamma4':
-        Condition(domain=CartesianDomain({
-            'x': 0,
-            'y': [0, 1]
-        }),
-                  equation=FixedValue(0.0)),
-        'D':
-        Condition(input_points=LabelTensor(torch.rand(size=(100, 2)),
-                                           ['x', 'y']),
-                  equation=my_laplace),
-        'data':
-        Condition(input_points=in_, output_points=out_),
-        'data2':
-        Condition(input_points=in2_, output_points=out2_),
-        'unsupervised':
-        Condition(
-            input_points=LabelTensor(torch.rand(size=(45, 2)), ['x', 'y']),
-            conditional_variables=LabelTensor(torch.ones(size=(45, 1)),
-                                              ['alpha']),
-        ),
-        'unsupervised2':
-        Condition(
-            input_points=LabelTensor(torch.rand(size=(90, 2)), ['x', 'y']),
-            conditional_variables=LabelTensor(torch.ones(size=(90, 1)),
-                                              ['alpha']),
-        )
-    }
-
-
-boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4']
-poisson = Poisson()
-poisson.discretise_domain(10, 'grid', locations=boundaries)
-
-
-def test_sample():
-    sample_dataset = SamplePointDataset(poisson, device='cpu')
-    assert len(sample_dataset) == 140
-    assert sample_dataset.input_points.shape == (140, 2)
-    assert sample_dataset.input_points.labels == ['x', 'y']
-    assert sample_dataset.condition_indices.dtype == torch.uint8
-    assert sample_dataset.condition_indices.max() == torch.tensor(4)
-    assert sample_dataset.condition_indices.min() == torch.tensor(0)
-
-
-def test_data():
-    dataset = SupervisedDataset(poisson, device='cpu')
-    assert len(dataset) == 61
-    assert dataset['input_points'].shape == (61, 2)
-    assert dataset.input_points.shape == (61, 2)
-    assert dataset['input_points'].labels == ['x', 'y']
-    assert dataset.input_points.labels == ['x', 'y']
-    assert dataset.input_points[3:].shape == (58, 2)
-    assert dataset.output_points[:3].labels == ['u']
-    assert dataset.output_points.shape == (61, 1)
-    assert dataset.output_points.labels == ['u']
-    assert dataset.condition_indices.dtype == torch.uint8
-    assert dataset.condition_indices.max() == torch.tensor(1)
-    assert dataset.condition_indices.min() == torch.tensor(0)
-
-
-def test_unsupervised():
-    dataset = UnsupervisedDataset(poisson, device='cpu')
-    assert len(dataset) == 135
-    assert dataset.input_points.shape == (135, 2)
-    assert dataset.input_points.labels == ['x', 'y']
-    assert dataset.input_points[3:].shape == (132, 2)
-
-    assert dataset.conditional_variables.shape == (135, 1)
-    assert dataset.conditional_variables.labels == ['alpha']
-    assert dataset.condition_indices.dtype == torch.uint8
-    assert dataset.condition_indices.max() == torch.tensor(1)
-    assert dataset.condition_indices.min() == torch.tensor(0)
-
-
-def test_data_module():
-    data_module = PinaDataModule(poisson, device='cpu')
-    data_module.setup()
-    loader = data_module.train_dataloader()
-    assert isinstance(loader, PinaDataLoader)
-    assert isinstance(loader, PinaDataLoader)
-
-    data_module = PinaDataModule(poisson,
-                                 device='cpu',
-                                 batch_size=10,
-                                 shuffle=False)
-    data_module.setup()
-    loader = data_module.train_dataloader()
-    assert len(loader) == 24
-    for i in loader:
-        assert len(i) <= 10
-    len_ref = sum(
-        [math.ceil(len(dataset) * 0.7) for dataset in data_module.datasets])
-    len_real = sum(
-        [len(dataset) for dataset in data_module.splits['train'].values()])
-    assert len_ref == len_real
-
-    supervised_dataset = SupervisedDataset(poisson, device='cpu')
-    data_module = PinaDataModule(poisson,
-                                 device='cpu',
-                                 batch_size=10,
-                                 shuffle=False,
-                                 datasets=[supervised_dataset])
-    data_module.setup()
-    loader = data_module.train_dataloader()
-    for batch in loader:
-        assert len(batch) <= 10
-
-    physics_dataset = SamplePointDataset(poisson, device='cpu')
-    data_module = PinaDataModule(poisson,
-                                 device='cpu',
-                                 batch_size=10,
-                                 shuffle=False,
-                                 datasets=[physics_dataset])
-    data_module.setup()
-    loader = data_module.train_dataloader()
-    for batch in loader:
-        assert len(batch) <= 10
-
-    unsupervised_dataset = UnsupervisedDataset(poisson, device='cpu')
-    data_module = PinaDataModule(poisson,
-                                 device='cpu',
-                                 batch_size=10,
-                                 shuffle=False,
-                                 datasets=[unsupervised_dataset])
-    data_module.setup()
-    loader = data_module.train_dataloader()
-    for batch in loader:
-        assert len(batch) <= 10
-
-
-def test_loader():
-    data_module = PinaDataModule(poisson, device='cpu', batch_size=10)
-    data_module.setup()
-    loader = data_module.train_dataloader()
-    assert isinstance(loader, PinaDataLoader)
-    assert len(loader) == 24
-    for i in loader:
-        assert len(i) <= 10
-        assert i.supervised.input_points.labels == ['x', 'y']
-        assert i.physics.input_points.labels == ['x', 'y']
-        assert i.unsupervised.input_points.labels == ['x', 'y']
-        assert i.supervised.input_points.requires_grad == True
-        assert i.physics.input_points.requires_grad == True
-        assert i.unsupervised.input_points.requires_grad == True
-
-
-coordinates = LabelTensor(torch.rand((100, 100, 2)), labels=['x', 'y'])
-data = LabelTensor(torch.rand((100, 100, 3)), labels=['ux', 'uy', 'p'])
-
-
-class GraphProblem(AbstractProblem):
-    output = LabelTensor(torch.rand((100, 3)), labels=['ux', 'uy', 'p'])
-    input = [
-        Graph.build('radius',
-                    nodes_coordinates=coordinates[i, :, :],
-                    nodes_data=data[i, :, :],
-                    radius=0.2) for i in range(100)
-    ]
-    output_variables = ['u']
-
-    conditions = {
-        'graph_data': Condition(input_points=input, output_points=output)
-    }
-
-
-graph_problem = GraphProblem()
-
-
-def test_loader_graph():
-    data_module = PinaDataModule(graph_problem, device='cpu', batch_size=10)
-    data_module.setup()
-    loader = data_module.train_dataloader()
-    for i in loader:
-        assert len(i) <= 10
-        assert isinstance(i.supervised.input_points, list)
-        assert all(isinstance(x, Graph) for x in i.supervised.input_points)
diff --git a/tests/test_label_tensor/test_label_tensor_01.py b/tests/test_label_tensor/test_label_tensor_01.py
index 57aafb8c9..ea43307cb 100644
--- a/tests/test_label_tensor/test_label_tensor_01.py
+++ b/tests/test_label_tensor/test_label_tensor_01.py
@@ -114,5 +114,5 @@ def test_slice():
     assert torch.allclose(tensor_view2, data[3])
 
     tensor_view3 = tensor[:, 2]
-    assert tensor_view3.labels == labels[2]
+    assert tensor_view3.labels == [labels[2]]
     assert torch.allclose(tensor_view3, data[:, 2].reshape(-1, 1))
diff --git a/tests/test_solvers/test_pinn.py b/tests/test_solvers/test_pinn.py
index 105bd886c..72887a4f7 100644
--- a/tests/test_solvers/test_pinn.py
+++ b/tests/test_solvers/test_pinn.py
@@ -1,5 +1,4 @@
 import torch
-
 from pina.problem import SpatialProblem, InverseProblem
 from pina.operators import laplacian
 from pina.domain import CartesianDomain
@@ -9,7 +8,7 @@
 from pina.model import FeedForward
 from pina.equation.equation import Equation
 from pina.equation.equation_factory import FixedValue
-from pina.loss.loss_interface import LpLoss
+from pina.loss import LpLoss
 
 
 def laplace_equation(input_, output_):
@@ -54,22 +53,22 @@ def laplace_equation(input_, output_, params_):
 
     # define the conditions for the loss (boundary conditions, equation, data)
     conditions = {
-        'gamma1': Condition(location=CartesianDomain({'x': [x_min, x_max],
+        'gamma1': Condition(domain=CartesianDomain({'x': [x_min, x_max],
             'y':  y_max}),
             equation=FixedValue(0.0, components=['u'])),
-        'gamma2': Condition(location=CartesianDomain(
+        'gamma2': Condition(domain=CartesianDomain(
             {'x': [x_min, x_max], 'y': y_min
             }),
             equation=FixedValue(0.0, components=['u'])),
-        'gamma3': Condition(location=CartesianDomain(
+        'gamma3': Condition(domain=CartesianDomain(
             {'x':  x_max, 'y': [y_min, y_max]
             }),
             equation=FixedValue(0.0, components=['u'])),
-        'gamma4': Condition(location=CartesianDomain(
+        'gamma4': Condition(domain=CartesianDomain(
             {'x': x_min, 'y': [y_min, y_max]
             }),
             equation=FixedValue(0.0, components=['u'])),
-        'D': Condition(location=CartesianDomain(
+        'D': Condition(domain=CartesianDomain(
             {'x': [x_min, x_max], 'y': [y_min, y_max]
             }),
         equation=Equation(laplace_equation)),
@@ -84,16 +83,16 @@ class Poisson(SpatialProblem):
 
     conditions = {
         'gamma1': Condition(
-            location=CartesianDomain({'x': [0, 1], 'y':  1}),
+            domain=CartesianDomain({'x': [0, 1], 'y':  1}),
             equation=FixedValue(0.0)),
         'gamma2': Condition(
-            location=CartesianDomain({'x': [0, 1], 'y': 0}),
+            domain=CartesianDomain({'x': [0, 1], 'y': 0}),
             equation=FixedValue(0.0)),
         'gamma3': Condition(
-            location=CartesianDomain({'x':  1, 'y': [0, 1]}),
+            domain=CartesianDomain({'x':  1, 'y': [0, 1]}),
             equation=FixedValue(0.0)),
         'gamma4': Condition(
-            location=CartesianDomain({'x': 0, 'y': [0, 1]}),
+            domain=CartesianDomain({'x': 0, 'y': [0, 1]}),
             equation=FixedValue(0.0)),
         'D': Condition(
             input_points=LabelTensor(torch.rand(size=(100, 2)), ['x', 'y']),
@@ -112,7 +111,6 @@ def poisson_sol(self, pts):
 
     truth_solution = poisson_sol
 
-
 class myFeature(torch.nn.Module):
     """
     Feature: sin(x)
@@ -158,24 +156,10 @@ def test_train_cpu():
     pinn = PINN(problem = poisson_problem, model=model,
                 extra_features=None, loss=LpLoss())
     trainer = Trainer(solver=pinn, max_epochs=1,
-                      accelerator='cpu', batch_size=20)
-    trainer.train()
-
-def test_log():
-    poisson_problem.discretise_domain(100)
-    solver = PINN(problem = poisson_problem, model=model,
-                extra_features=None, loss=LpLoss())
-    trainer = Trainer(solver, max_epochs=2, accelerator='cpu')
-    trainer.train()
-    # assert the logged metrics are correct
-    logged_metrics = sorted(list(trainer.logged_metrics.keys()))
-    total_metrics = sorted(
-        list([key + '_loss' for key in poisson_problem.conditions.keys()])
-        + ['mean_loss'])
-    assert logged_metrics == total_metrics
+                      accelerator='cpu', batch_size=20, val_size=0., train_size=1., test_size=0.)
 
-def test_train_restore():
-    tmpdir = "tests/tmp_restore"
+def test_train_load():
+    tmpdir = "tests/tmp_load"
     poisson_problem = Poisson()
     boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4']
     n = 10
@@ -185,20 +169,25 @@ def test_train_restore():
                 extra_features=None,
                 loss=LpLoss())
     trainer = Trainer(solver=pinn,
-                      max_epochs=5,
+                      max_epochs=15,
                       accelerator='cpu',
                       default_root_dir=tmpdir)
     trainer.train()
-    ntrainer = Trainer(solver=pinn, max_epochs=15, accelerator='cpu')
-    t = ntrainer.train(
-        ckpt_path=f'{tmpdir}/lightning_logs/version_0/'
-        'checkpoints/epoch=4-step=10.ckpt')
+    new_pinn = PINN.load_from_checkpoint(
+        f'{tmpdir}/lightning_logs/version_0/checkpoints/epoch=14-step=15.ckpt',
+        problem = poisson_problem, model=model)
+    test_pts = CartesianDomain({'x': [0, 1], 'y': [0, 1]}).sample(10)
+    assert new_pinn.forward(test_pts).extract(['u']).shape == (10, 1)
+    assert new_pinn.forward(test_pts).extract(
+        ['u']).shape == pinn.forward(test_pts).extract(['u']).shape
+    torch.testing.assert_close(
+        new_pinn.forward(test_pts).extract(['u']),
+        pinn.forward(test_pts).extract(['u']))
     import shutil
     shutil.rmtree(tmpdir)
 
-
-def test_train_load():
-    tmpdir = "tests/tmp_load"
+def test_train_restore():
+    tmpdir = "tests/tmp_restore"
     poisson_problem = Poisson()
     boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4']
     n = 10
@@ -208,20 +197,14 @@ def test_train_load():
                 extra_features=None,
                 loss=LpLoss())
     trainer = Trainer(solver=pinn,
-                      max_epochs=15,
+                      max_epochs=5,
                       accelerator='cpu',
                       default_root_dir=tmpdir)
     trainer.train()
-    new_pinn = PINN.load_from_checkpoint(
-        f'{tmpdir}/lightning_logs/version_0/checkpoints/epoch=14-step=30.ckpt',
-        problem = poisson_problem, model=model)
-    test_pts = CartesianDomain({'x': [0, 1], 'y': [0, 1]}).sample(10)
-    assert new_pinn.forward(test_pts).extract(['u']).shape == (10, 1)
-    assert new_pinn.forward(test_pts).extract(
-        ['u']).shape == pinn.forward(test_pts).extract(['u']).shape
-    torch.testing.assert_close(
-        new_pinn.forward(test_pts).extract(['u']),
-        pinn.forward(test_pts).extract(['u']))
+    ntrainer = Trainer(solver=pinn, max_epochs=15, accelerator='cpu')
+    t = ntrainer.train(
+        ckpt_path=f'{tmpdir}/lightning_logs/version_0/'
+                  'checkpoints/epoch=4-step=5.ckpt')
     import shutil
     shutil.rmtree(tmpdir)
 
@@ -229,36 +212,24 @@ def test_train_inverse_problem_cpu():
     poisson_problem = InversePoisson()
     boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4', 'D']
     n = 100
-    poisson_problem.discretise_domain(n, 'random', locations=boundaries)
+    poisson_problem.discretise_domain(n, 'random', locations=boundaries,
+                                      variables=['x', 'y'])
     pinn = PINN(problem = poisson_problem, model=model,
                 extra_features=None, loss=LpLoss())
     trainer = Trainer(solver=pinn, max_epochs=1,
                       accelerator='cpu', batch_size=20)
     trainer.train()
 
-
-# # TODO does not currently work
-# def test_train_inverse_problem_restore():
-#     tmpdir = "tests/tmp_restore_inv"
-#     poisson_problem = InversePoisson()
-#     boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4', 'D']
-#     n = 100
-#     poisson_problem.discretise_domain(n, 'random', locations=boundaries)
-#     pinn = PINN(problem=poisson_problem,
-#                 model=model,
-#                 extra_features=None,
-#                 loss=LpLoss())
-#     trainer = Trainer(solver=pinn,
-#                       max_epochs=5,
-#                       accelerator='cpu',
-#                       default_root_dir=tmpdir)
-#     trainer.train()
-#     ntrainer = Trainer(solver=pinn, max_epochs=5, accelerator='cpu')
-#     t = ntrainer.train(
-#         ckpt_path=f'{tmpdir}/lightning_logs/version_0/checkpoints/epoch=4-step=10.ckpt')
-#     import shutil
-#     shutil.rmtree(tmpdir)
-
+def test_train_extra_feats_cpu():
+    poisson_problem = Poisson()
+    boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4']
+    n = 10
+    poisson_problem.discretise_domain(n, 'grid', locations=boundaries)
+    pinn = PINN(problem=poisson_problem,
+                model=model_extra_feats,
+                extra_features=extra_feats)
+    trainer = Trainer(solver=pinn, max_epochs=5, accelerator='cpu')
+    trainer.train()
 
 def test_train_inverse_problem_load():
     tmpdir = "tests/tmp_load_inv"
@@ -276,7 +247,7 @@ def test_train_inverse_problem_load():
                       default_root_dir=tmpdir)
     trainer.train()
     new_pinn = PINN.load_from_checkpoint(
-        f'{tmpdir}/lightning_logs/version_0/checkpoints/epoch=14-step=30.ckpt',
+        f'{tmpdir}/lightning_logs/version_0/checkpoints/epoch=14-step=15.ckpt',
         problem = poisson_problem, model=model)
     test_pts = CartesianDomain({'x': [0, 1], 'y': [0, 1]}).sample(10)
     assert new_pinn.forward(test_pts).extract(['u']).shape == (10, 1)
@@ -286,160 +257,4 @@ def test_train_inverse_problem_load():
         new_pinn.forward(test_pts).extract(['u']),
         pinn.forward(test_pts).extract(['u']))
     import shutil
-    shutil.rmtree(tmpdir)
-
-# # TODO fix asap. Basically sampling few variables
-# # works only if both variables are in a range.
-# # if one is fixed and the other not, this will
-# # not work. This test also needs to be fixed and
-# # insert in test problem not in test pinn.
-# def test_train_cpu_sampling_few_vars():
-#     poisson_problem = Poisson()
-#     boundaries = ['gamma1', 'gamma2', 'gamma3']
-#     n = 10
-#     poisson_problem.discretise_domain(n, 'grid', locations=boundaries)
-#     poisson_problem.discretise_domain(n, 'random', locations=['gamma4'], variables=['x'])
-#     poisson_problem.discretise_domain(n, 'random', locations=['gamma4'], variables=['y'])
-#     pinn = PINN(problem = poisson_problem, model=model, extra_features=None, loss=LpLoss())
-#     trainer = Trainer(solver=pinn, kwargs={'max_epochs' : 5, 'accelerator':'cpu'})
-#     trainer.train()
-
-
-def test_train_extra_feats_cpu():
-    poisson_problem = Poisson()
-    boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4']
-    n = 10
-    poisson_problem.discretise_domain(n, 'grid', locations=boundaries)
-    pinn = PINN(problem=poisson_problem,
-                model=model_extra_feats,
-                extra_features=extra_feats)
-    trainer = Trainer(solver=pinn, max_epochs=5, accelerator='cpu')
-    trainer.train()
-
-
-# TODO, fix GitHub actions to run also on GPU
-# def test_train_gpu():
-#     poisson_problem = Poisson()
-#     boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4']
-#     n = 10
-#     poisson_problem.discretise_domain(n, 'grid', locations=boundaries)
-#     pinn = PINN(problem = poisson_problem, model=model, extra_features=None, loss=LpLoss())
-#     trainer = Trainer(solver=pinn, kwargs={'max_epochs' : 5, 'accelerator':'gpu'})
-#     trainer.train()
-
-# def test_train_gpu(): #TODO fix ASAP
-#     poisson_problem = Poisson()
-#     boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4']
-#     n = 10
-#     poisson_problem.discretise_domain(n, 'grid', locations=boundaries)
-#     poisson_problem.conditions.pop('data') # The input/output pts are allocated on cpu
-#     pinn = PINN(problem = poisson_problem, model=model, extra_features=None, loss=LpLoss())
-#     trainer = Trainer(solver=pinn, kwargs={'max_epochs' : 5, 'accelerator':'gpu'})
-#     trainer.train()
-
-# def test_train_2():
-#     boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4']
-#     n = 10
-#     expected_keys = [[], list(range(0, 50, 3))]
-#     param = [0, 3]
-#     for i, truth_key in zip(param, expected_keys):
-#         pinn = PINN(problem, model)
-#         pinn.discretise_domain(n, 'grid', locations=boundaries)
-#         pinn.discretise_domain(n, 'grid', locations=['D'])
-#         pinn.train(50, save_loss=i)
-#         assert list(pinn.history_loss.keys()) == truth_key
-
-
-# def test_train_extra_feats():
-#     pinn = PINN(problem, model_extra_feat, [myFeature()])
-#     boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4']
-#     n = 10
-#     pinn.discretise_domain(n, 'grid', locations=boundaries)
-#     pinn.discretise_domain(n, 'grid', locations=['D'])
-#     pinn.train(5)
-
-
-# def test_train_2_extra_feats():
-#     boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4']
-#     n = 10
-#     expected_keys = [[], list(range(0, 50, 3))]
-#     param = [0, 3]
-#     for i, truth_key in zip(param, expected_keys):
-#         pinn = PINN(problem, model_extra_feat, [myFeature()])
-#         pinn.discretise_domain(n, 'grid', locations=boundaries)
-#         pinn.discretise_domain(n, 'grid', locations=['D'])
-#         pinn.train(50, save_loss=i)
-#         assert list(pinn.history_loss.keys()) == truth_key
-
-
-# def test_train_with_optimizer_kwargs():
-#     boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4']
-#     n = 10
-#     expected_keys = [[], list(range(0, 50, 3))]
-#     param = [0, 3]
-#     for i, truth_key in zip(param, expected_keys):
-#         pinn = PINN(problem, model, optimizer_kwargs={'lr' : 0.3})
-#         pinn.discretise_domain(n, 'grid', locations=boundaries)
-#         pinn.discretise_domain(n, 'grid', locations=['D'])
-#         pinn.train(50, save_loss=i)
-#         assert list(pinn.history_loss.keys()) == truth_key
-
-
-# def test_train_with_lr_scheduler():
-#     boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4']
-#     n = 10
-#     expected_keys = [[], list(range(0, 50, 3))]
-#     param = [0, 3]
-#     for i, truth_key in zip(param, expected_keys):
-#         pinn = PINN(
-#             problem,
-#             model,
-#             lr_scheduler_type=torch.optim.lr_scheduler.CyclicLR,
-#             lr_scheduler_kwargs={'base_lr' : 0.1, 'max_lr' : 0.3, 'cycle_momentum': False}
-#         )
-#         pinn.discretise_domain(n, 'grid', locations=boundaries)
-#         pinn.discretise_domain(n, 'grid', locations=['D'])
-#         pinn.train(50, save_loss=i)
-#         assert list(pinn.history_loss.keys()) == truth_key
-
-
-# # def test_train_batch():
-# #     pinn = PINN(problem, model, batch_size=6)
-# #     boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4']
-# #     n = 10
-# #     pinn.discretise_domain(n, 'grid', locations=boundaries)
-# #     pinn.discretise_domain(n, 'grid', locations=['D'])
-# #     pinn.train(5)
-
-
-# # def test_train_batch_2():
-# #     boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4']
-# #     n = 10
-# #     expected_keys = [[], list(range(0, 50, 3))]
-# #     param = [0, 3]
-# #     for i, truth_key in zip(param, expected_keys):
-# #         pinn = PINN(problem, model, batch_size=6)
-# #         pinn.discretise_domain(n, 'grid', locations=boundaries)
-# #         pinn.discretise_domain(n, 'grid', locations=['D'])
-# #         pinn.train(50, save_loss=i)
-# #         assert list(pinn.history_loss.keys()) == truth_key
-
-
-# if torch.cuda.is_available():
-
-#     # def test_gpu_train():
-#     #     pinn = PINN(problem, model, batch_size=20, device='cuda')
-#     #     boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4']
-#     #     n = 100
-#     #     pinn.discretise_domain(n, 'grid', locations=boundaries)
-#     #     pinn.discretise_domain(n, 'grid', locations=['D'])
-#     #     pinn.train(5)
-
-#     def test_gpu_train_nobatch():
-#         pinn = PINN(problem, model, batch_size=None, device='cuda')
-#         boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4']
-#         n = 100
-#         pinn.discretise_domain(n, 'grid', locations=boundaries)
-#         pinn.discretise_domain(n, 'grid', locations=['D'])
-#         pinn.train(5)
-
+    shutil.rmtree(tmpdir)
\ No newline at end of file
diff --git a/tests/test_solvers/test_supervised_solver.py b/tests/test_solvers/test_supervised_solver.py
index 8ceadcd93..ebe8179e6 100644
--- a/tests/test_solvers/test_supervised_solver.py
+++ b/tests/test_solvers/test_supervised_solver.py
@@ -121,7 +121,7 @@ def test_train_cpu():
                       batch_size=5,
                       train_size=1,
                       test_size=0.,
-                      eval_size=0.)
+                      val_size=0.)
     trainer.train()
 test_train_cpu()
 

From 355a9fe7de69f5715d1f524855726540e9d318e1 Mon Sep 17 00:00:00 2001
From: Dario Coscia <93731561+dario-coscia@users.noreply.github.com>
Date: Wed, 27 Nov 2024 16:12:23 +0100
Subject: [PATCH 32/55] remove unnecessary dependency

---
 pina/solvers/supervised.py | 1 -
 1 file changed, 1 deletion(-)

diff --git a/pina/solvers/supervised.py b/pina/solvers/supervised.py
index d978fc37c..947ab3b2f 100644
--- a/pina/solvers/supervised.py
+++ b/pina/solvers/supervised.py
@@ -1,7 +1,6 @@
 """ Module for SupervisedSolver """
 import torch
 from pytorch_lightning.utilities.types import STEP_OUTPUT
-from sympy.strategies.branch import condition
 from torch.nn.modules.loss import _Loss
 from ..optim import TorchOptimizer, TorchScheduler
 from .solver import SolverInterface

From 435d58703abda9b15dc0387da81cec10bd925b2c Mon Sep 17 00:00:00 2001
From: Filippo Olivo <filippo@filippoolivo.com>
Date: Wed, 27 Nov 2024 19:30:12 +0100
Subject: [PATCH 33/55] Fix bugs (#385)

---
 pina/label_tensor.py           |  2 +-
 pina/solvers/__init__.py       |  1 -
 pina/solvers/pinns/basepinn.py |  5 +----
 pina/solvers/pinns/pinn.py     |  2 +-
 pina/solvers/solver.py         | 10 ++++++----
 pina/solvers/supervised.py     |  2 --
 pina/trainer.py                |  1 -
 7 files changed, 9 insertions(+), 14 deletions(-)

diff --git a/pina/label_tensor.py b/pina/label_tensor.py
index a3cf5d237..631c5253c 100644
--- a/pina/label_tensor.py
+++ b/pina/label_tensor.py
@@ -4,7 +4,7 @@
 from torch import Tensor
 
 
-full_labels = True
+full_labels = False
 MATH_FUNCTIONS = {torch.sin, torch.cos}
 
 class LabelTensor(torch.Tensor):
diff --git a/pina/solvers/__init__.py b/pina/solvers/__init__.py
index 59a182691..7bb988d56 100644
--- a/pina/solvers/__init__.py
+++ b/pina/solvers/__init__.py
@@ -17,4 +17,3 @@
 from .supervised import SupervisedSolver
 from .rom import ReducedOrderModelSolver
 from .garom import GAROM
-from .graph import GraphSupervisedSolver
diff --git a/pina/solvers/pinns/basepinn.py b/pina/solvers/pinns/basepinn.py
index 66f4d14c3..f44ea1d0d 100644
--- a/pina/solvers/pinns/basepinn.py
+++ b/pina/solvers/pinns/basepinn.py
@@ -3,12 +3,10 @@
 from abc import ABCMeta, abstractmethod
 import torch
 from torch.nn.modules.loss import _Loss
-from ...condition import InputOutputPointsCondition
 from ...solvers.solver import SolverInterface
 from ...utils import check_consistency
 from ...loss.loss_interface import LossInterface
 from ...problem import InverseProblem
-from ...condition import DomainEquationCondition
 from ...optim import TorchOptimizer, TorchScheduler
 
 torch.pi = torch.acos(torch.zeros(1)).item() * 2  # which is 3.1415927410125732
@@ -26,8 +24,7 @@ class PINNInterface(SolverInterface, metaclass=ABCMeta):
     to the user to choose which problem the implemented solver inheriting from
     this class is suitable for.
     """
-    accepted_condition_types = [DomainEquationCondition.condition_type[0],
-                                InputOutputPointsCondition.condition_type[0]]
+
     def __init__(
             self,
             models,
diff --git a/pina/solvers/pinns/pinn.py b/pina/solvers/pinns/pinn.py
index 088820202..d1ab21d76 100644
--- a/pina/solvers/pinns/pinn.py
+++ b/pina/solvers/pinns/pinn.py
@@ -11,7 +11,7 @@
 
 
 from .basepinn import PINNInterface
-from pina.problem import InverseProblem
+from ...problem import InverseProblem
 
 
 class PINN(PINNInterface):
diff --git a/pina/solvers/solver.py b/pina/solvers/solver.py
index 3a8f400c2..8052b4b8e 100644
--- a/pina/solvers/solver.py
+++ b/pina/solvers/solver.py
@@ -134,16 +134,18 @@ def on_train_start(self):
         return super().on_train_start()
 
     def _check_solver_consistency(self, problem):
-        """
-        TODO
-        """
+        pass
+        #TODO : Implement this method for the conditions
+        '''
+        
+    
         for _, condition in problem.conditions.items():
             if not set(condition.condition_type).issubset(
                     set(self.accepted_condition_types)):
                 raise ValueError(
                     f'{self.__name__} dose not support condition '
                     f'{condition.condition_type}')
-
+        '''
     @staticmethod
     def get_batch_size(batch):
         # Assuming batch is your custom Batch object
diff --git a/pina/solvers/supervised.py b/pina/solvers/supervised.py
index 947ab3b2f..bce4b31a8 100644
--- a/pina/solvers/supervised.py
+++ b/pina/solvers/supervised.py
@@ -7,7 +7,6 @@
 from ..label_tensor import LabelTensor
 from ..utils import check_consistency
 from ..loss.loss_interface import LossInterface
-from ..condition import InputOutputPointsCondition
 
 
 class SupervisedSolver(SolverInterface):
@@ -38,7 +37,6 @@ class SupervisedSolver(SolverInterface):
     we are seeking to approximate multiple (discretised) functions given
     multiple (discretised) input functions.
     """
-    accepted_condition_types = [InputOutputPointsCondition.condition_type[0]]
     __name__ = 'SupervisedSolver'
 
     def __init__(self,
diff --git a/pina/trainer.py b/pina/trainer.py
index a7c5c3513..f8bccd8c5 100644
--- a/pina/trainer.py
+++ b/pina/trainer.py
@@ -1,5 +1,4 @@
 """ Trainer module. """
-import warnings
 import torch
 import lightning
 from .utils import check_consistency

From b9a5305b3ca5880a1af268afb051057d92daa684 Mon Sep 17 00:00:00 2001
From: Dario Coscia <93731561+dario-coscia@users.noreply.github.com>
Date: Wed, 27 Nov 2024 19:58:42 +0100
Subject: [PATCH 34/55] Revert "Fix bugs (#385)" (#386)

This reverts commit 69cd0ed8cda91c92dab6551a0c6dfd94d199cee7.
---
 pina/label_tensor.py           |  2 +-
 pina/solvers/__init__.py       |  1 +
 pina/solvers/pinns/basepinn.py |  5 ++++-
 pina/solvers/pinns/pinn.py     |  2 +-
 pina/solvers/solver.py         | 10 ++++------
 pina/solvers/supervised.py     |  2 ++
 pina/trainer.py                |  1 +
 7 files changed, 14 insertions(+), 9 deletions(-)

diff --git a/pina/label_tensor.py b/pina/label_tensor.py
index 631c5253c..a3cf5d237 100644
--- a/pina/label_tensor.py
+++ b/pina/label_tensor.py
@@ -4,7 +4,7 @@
 from torch import Tensor
 
 
-full_labels = False
+full_labels = True
 MATH_FUNCTIONS = {torch.sin, torch.cos}
 
 class LabelTensor(torch.Tensor):
diff --git a/pina/solvers/__init__.py b/pina/solvers/__init__.py
index 7bb988d56..59a182691 100644
--- a/pina/solvers/__init__.py
+++ b/pina/solvers/__init__.py
@@ -17,3 +17,4 @@
 from .supervised import SupervisedSolver
 from .rom import ReducedOrderModelSolver
 from .garom import GAROM
+from .graph import GraphSupervisedSolver
diff --git a/pina/solvers/pinns/basepinn.py b/pina/solvers/pinns/basepinn.py
index f44ea1d0d..66f4d14c3 100644
--- a/pina/solvers/pinns/basepinn.py
+++ b/pina/solvers/pinns/basepinn.py
@@ -3,10 +3,12 @@
 from abc import ABCMeta, abstractmethod
 import torch
 from torch.nn.modules.loss import _Loss
+from ...condition import InputOutputPointsCondition
 from ...solvers.solver import SolverInterface
 from ...utils import check_consistency
 from ...loss.loss_interface import LossInterface
 from ...problem import InverseProblem
+from ...condition import DomainEquationCondition
 from ...optim import TorchOptimizer, TorchScheduler
 
 torch.pi = torch.acos(torch.zeros(1)).item() * 2  # which is 3.1415927410125732
@@ -24,7 +26,8 @@ class PINNInterface(SolverInterface, metaclass=ABCMeta):
     to the user to choose which problem the implemented solver inheriting from
     this class is suitable for.
     """
-
+    accepted_condition_types = [DomainEquationCondition.condition_type[0],
+                                InputOutputPointsCondition.condition_type[0]]
     def __init__(
             self,
             models,
diff --git a/pina/solvers/pinns/pinn.py b/pina/solvers/pinns/pinn.py
index d1ab21d76..088820202 100644
--- a/pina/solvers/pinns/pinn.py
+++ b/pina/solvers/pinns/pinn.py
@@ -11,7 +11,7 @@
 
 
 from .basepinn import PINNInterface
-from ...problem import InverseProblem
+from pina.problem import InverseProblem
 
 
 class PINN(PINNInterface):
diff --git a/pina/solvers/solver.py b/pina/solvers/solver.py
index 8052b4b8e..3a8f400c2 100644
--- a/pina/solvers/solver.py
+++ b/pina/solvers/solver.py
@@ -134,18 +134,16 @@ def on_train_start(self):
         return super().on_train_start()
 
     def _check_solver_consistency(self, problem):
-        pass
-        #TODO : Implement this method for the conditions
-        '''
-        
-    
+        """
+        TODO
+        """
         for _, condition in problem.conditions.items():
             if not set(condition.condition_type).issubset(
                     set(self.accepted_condition_types)):
                 raise ValueError(
                     f'{self.__name__} dose not support condition '
                     f'{condition.condition_type}')
-        '''
+
     @staticmethod
     def get_batch_size(batch):
         # Assuming batch is your custom Batch object
diff --git a/pina/solvers/supervised.py b/pina/solvers/supervised.py
index bce4b31a8..947ab3b2f 100644
--- a/pina/solvers/supervised.py
+++ b/pina/solvers/supervised.py
@@ -7,6 +7,7 @@
 from ..label_tensor import LabelTensor
 from ..utils import check_consistency
 from ..loss.loss_interface import LossInterface
+from ..condition import InputOutputPointsCondition
 
 
 class SupervisedSolver(SolverInterface):
@@ -37,6 +38,7 @@ class SupervisedSolver(SolverInterface):
     we are seeking to approximate multiple (discretised) functions given
     multiple (discretised) input functions.
     """
+    accepted_condition_types = [InputOutputPointsCondition.condition_type[0]]
     __name__ = 'SupervisedSolver'
 
     def __init__(self,
diff --git a/pina/trainer.py b/pina/trainer.py
index f8bccd8c5..a7c5c3513 100644
--- a/pina/trainer.py
+++ b/pina/trainer.py
@@ -1,4 +1,5 @@
 """ Trainer module. """
+import warnings
 import torch
 import lightning
 from .utils import check_consistency

From 6217de7d1ed3ae373752700ca1cdf24096a9f554 Mon Sep 17 00:00:00 2001
From: Filippo Olivo <filippo@filippoolivo.com>
Date: Thu, 28 Nov 2024 11:06:38 +0100
Subject: [PATCH 35/55] Fix bugs (#387)

---
 pina/data/data_module.py       | 24 +++++++++++++-----------
 pina/data/dataset.py           | 10 ++++++----
 pina/label_tensor.py           |  2 +-
 pina/solvers/__init__.py       |  1 -
 pina/solvers/pinns/basepinn.py |  5 +----
 pina/solvers/pinns/pinn.py     |  2 +-
 pina/solvers/solver.py         | 10 ++++++----
 pina/solvers/supervised.py     |  2 --
 pina/trainer.py                |  1 -
 9 files changed, 28 insertions(+), 29 deletions(-)

diff --git a/pina/data/data_module.py b/pina/data/data_module.py
index c9af8aeba..4831e20e7 100644
--- a/pina/data/data_module.py
+++ b/pina/data/data_module.py
@@ -95,7 +95,7 @@ def __init__(self,
         logging.debug('Start initialization of Pina DataModule')
         logging.info('Start initialization of Pina DataModule')
         super().__init__()
-        self.default_batching = automatic_batching
+        self.automatic_batching = automatic_batching
         self.batch_size = batch_size
         self.shuffle = shuffle
         self.repeat = repeat
@@ -133,24 +133,24 @@ def setup(self, stage=None):
             self.train_dataset = PinaDatasetFactory(
                 self.collector_splits['train'],
                 max_conditions_lengths=self.find_max_conditions_lengths(
-                    'train'))
+                    'train'), automatic_batching=self.automatic_batching)
             if 'val' in self.collector_splits.keys():
                 self.val_dataset = PinaDatasetFactory(
                     self.collector_splits['val'],
                     max_conditions_lengths=self.find_max_conditions_lengths(
-                        'val')
+                        'val'),  automatic_batching=self.automatic_batching
                 )
         elif stage == 'test':
             self.test_dataset = PinaDatasetFactory(
                 self.collector_splits['test'],
                 max_conditions_lengths=self.find_max_conditions_lengths(
-                    'test')
+                    'test'), automatic_batching=self.automatic_batching
             )
         elif stage == 'predict':
             self.predict_dataset = PinaDatasetFactory(
                 self.collector_splits['predict'],
                 max_conditions_lengths=self.find_max_conditions_lengths(
-                                        'predict')
+                    'predict'), automatic_batching=self.automatic_batching
             )
         else:
             raise ValueError(
@@ -237,9 +237,9 @@ def val_dataloader(self):
             self.val_dataset)
 
         # Use default batching in torch DataLoader (good is batch size is small)
-        if self.default_batching:
+        if self.automatic_batching:
             collate = Collator(self.find_max_conditions_lengths('val'))
-            return DataLoader(self.val_dataset, self.batch_size,
+            return DataLoader(self.val_dataset, batch_size,
                               collate_fn=collate)
         collate = Collator(None)
         # Use custom batching (good if batch size is large)
@@ -252,14 +252,16 @@ def train_dataloader(self):
         Create the training dataloader
         """
         # Use default batching in torch DataLoader (good is batch size is small)
-        if self.default_batching:
+        batch_size = self.batch_size if self.batch_size is not None else len(
+            self.train_dataset)
+
+        if self.automatic_batching:
             collate = Collator(self.find_max_conditions_lengths('train'))
-            return DataLoader(self.train_dataset, self.batch_size,
+            return DataLoader(self.train_dataset, batch_size,
                                 collate_fn=collate)
         collate = Collator(None)
         # Use custom batching (good if batch size is large)
-        batch_size = self.batch_size if self.batch_size is not None else len(
-            self.train_dataset)
+
         sampler = PinaBatchSampler(self.train_dataset, batch_size,
                                        shuffle=False)
         return DataLoader(self.train_dataset, sampler=sampler,
diff --git a/pina/data/dataset.py b/pina/data/dataset.py
index 0bc923711..e5685f1fb 100644
--- a/pina/data/dataset.py
+++ b/pina/data/dataset.py
@@ -51,8 +51,12 @@ def __getitem__(self, item):
 
 class PinaTensorDataset(PinaDataset):
     def __init__(self, conditions_dict, max_conditions_lengths,
-                 ):
+                 automatic_batching):
         super().__init__(conditions_dict, max_conditions_lengths)
+        if automatic_batching:
+            self._getitem_func = self._getitem_int
+        else:
+            self._getitem_func = self._getitem_list
 
     def _getitem_int(self, idx):
         return {
@@ -72,9 +76,7 @@ def _getitem_list(self, idx):
         return to_return_dict
 
     def __getitem__(self, idx):
-        if isinstance(idx, int):
-            return self._getitem_int(idx)
-        return self._getitem_list(idx)
+        return self._getitem_func(idx)
 
 class PinaGraphDataset(PinaDataset):
     pass
diff --git a/pina/label_tensor.py b/pina/label_tensor.py
index a3cf5d237..631c5253c 100644
--- a/pina/label_tensor.py
+++ b/pina/label_tensor.py
@@ -4,7 +4,7 @@
 from torch import Tensor
 
 
-full_labels = True
+full_labels = False
 MATH_FUNCTIONS = {torch.sin, torch.cos}
 
 class LabelTensor(torch.Tensor):
diff --git a/pina/solvers/__init__.py b/pina/solvers/__init__.py
index 59a182691..7bb988d56 100644
--- a/pina/solvers/__init__.py
+++ b/pina/solvers/__init__.py
@@ -17,4 +17,3 @@
 from .supervised import SupervisedSolver
 from .rom import ReducedOrderModelSolver
 from .garom import GAROM
-from .graph import GraphSupervisedSolver
diff --git a/pina/solvers/pinns/basepinn.py b/pina/solvers/pinns/basepinn.py
index 66f4d14c3..f44ea1d0d 100644
--- a/pina/solvers/pinns/basepinn.py
+++ b/pina/solvers/pinns/basepinn.py
@@ -3,12 +3,10 @@
 from abc import ABCMeta, abstractmethod
 import torch
 from torch.nn.modules.loss import _Loss
-from ...condition import InputOutputPointsCondition
 from ...solvers.solver import SolverInterface
 from ...utils import check_consistency
 from ...loss.loss_interface import LossInterface
 from ...problem import InverseProblem
-from ...condition import DomainEquationCondition
 from ...optim import TorchOptimizer, TorchScheduler
 
 torch.pi = torch.acos(torch.zeros(1)).item() * 2  # which is 3.1415927410125732
@@ -26,8 +24,7 @@ class PINNInterface(SolverInterface, metaclass=ABCMeta):
     to the user to choose which problem the implemented solver inheriting from
     this class is suitable for.
     """
-    accepted_condition_types = [DomainEquationCondition.condition_type[0],
-                                InputOutputPointsCondition.condition_type[0]]
+
     def __init__(
             self,
             models,
diff --git a/pina/solvers/pinns/pinn.py b/pina/solvers/pinns/pinn.py
index 088820202..d1ab21d76 100644
--- a/pina/solvers/pinns/pinn.py
+++ b/pina/solvers/pinns/pinn.py
@@ -11,7 +11,7 @@
 
 
 from .basepinn import PINNInterface
-from pina.problem import InverseProblem
+from ...problem import InverseProblem
 
 
 class PINN(PINNInterface):
diff --git a/pina/solvers/solver.py b/pina/solvers/solver.py
index 3a8f400c2..8052b4b8e 100644
--- a/pina/solvers/solver.py
+++ b/pina/solvers/solver.py
@@ -134,16 +134,18 @@ def on_train_start(self):
         return super().on_train_start()
 
     def _check_solver_consistency(self, problem):
-        """
-        TODO
-        """
+        pass
+        #TODO : Implement this method for the conditions
+        '''
+        
+    
         for _, condition in problem.conditions.items():
             if not set(condition.condition_type).issubset(
                     set(self.accepted_condition_types)):
                 raise ValueError(
                     f'{self.__name__} dose not support condition '
                     f'{condition.condition_type}')
-
+        '''
     @staticmethod
     def get_batch_size(batch):
         # Assuming batch is your custom Batch object
diff --git a/pina/solvers/supervised.py b/pina/solvers/supervised.py
index 947ab3b2f..bce4b31a8 100644
--- a/pina/solvers/supervised.py
+++ b/pina/solvers/supervised.py
@@ -7,7 +7,6 @@
 from ..label_tensor import LabelTensor
 from ..utils import check_consistency
 from ..loss.loss_interface import LossInterface
-from ..condition import InputOutputPointsCondition
 
 
 class SupervisedSolver(SolverInterface):
@@ -38,7 +37,6 @@ class SupervisedSolver(SolverInterface):
     we are seeking to approximate multiple (discretised) functions given
     multiple (discretised) input functions.
     """
-    accepted_condition_types = [InputOutputPointsCondition.condition_type[0]]
     __name__ = 'SupervisedSolver'
 
     def __init__(self,
diff --git a/pina/trainer.py b/pina/trainer.py
index a7c5c3513..f8bccd8c5 100644
--- a/pina/trainer.py
+++ b/pina/trainer.py
@@ -1,5 +1,4 @@
 """ Trainer module. """
-import warnings
 import torch
 import lightning
 from .utils import check_consistency

From b645ed4623bc590ee29eecdc378f746af4558ec1 Mon Sep 17 00:00:00 2001
From: giovanni <giovanni.canali98@yahoo.it>
Date: Wed, 18 Dec 2024 14:58:15 +0100
Subject: [PATCH 36/55] sorted variables

---
 pina/domain/ellipsoid.py | 3 ---
 pina/domain/simplex.py   | 9 ++++-----
 2 files changed, 4 insertions(+), 8 deletions(-)

diff --git a/pina/domain/ellipsoid.py b/pina/domain/ellipsoid.py
index 132bd55a5..5d466743b 100644
--- a/pina/domain/ellipsoid.py
+++ b/pina/domain/ellipsoid.py
@@ -286,9 +286,6 @@ def _single_points_sample(n, variables):
         if self.fixed_ and (not self.range_):
             return _single_points_sample(n, variables).extract(variables)
 
-        if variables == "all":
-            variables = self.variables
-
         if mode in self.sample_modes:
             return _Nd_sampler(n, mode, variables).extract(variables)
         else:
diff --git a/pina/domain/simplex.py b/pina/domain/simplex.py
index 00bd980f3..6915c121e 100644
--- a/pina/domain/simplex.py
+++ b/pina/domain/simplex.py
@@ -234,12 +234,11 @@ def sample(self, n, mode="random", variables="all"):
             in ``variables``.
         """
 
-        # if variables == "all": # TODO: check if this is necessary (from 0.1)
-        #     variables = self.variables
-        # elif isinstance(variables, (list, tuple)):
-        #     variables = sorted(variables)
+        if variables == "all":
+            variables = self.variables
+        elif isinstance(variables, (list, tuple)):
+            variables = sorted(variables)
 
-        # if mode in ["random"]:
         if mode in self.sample_modes:
             if self._sample_surface:
                 sample_pts = self._sample_boundary_randomly(n)

From 65ca7441c084914eda62d69ed82d71eac9098b28 Mon Sep 17 00:00:00 2001
From: giovanni <giovanni.canali98@yahoo.it>
Date: Wed, 18 Dec 2024 15:17:59 +0100
Subject: [PATCH 37/55] sort variables in operations

---
 pina/domain/operation_interface.py | 5 ++++-
 pina/domain/union_domain.py        | 7 -------
 2 files changed, 4 insertions(+), 8 deletions(-)

diff --git a/pina/domain/operation_interface.py b/pina/domain/operation_interface.py
index 0300f5248..e42d37eb7 100644
--- a/pina/domain/operation_interface.py
+++ b/pina/domain/operation_interface.py
@@ -43,7 +43,10 @@ def variables(self):
         :return: All the variables defined in ``__init__`` in order.
         :rtype: list[str]
         """
-        return self.geometries[0].variables
+        variables = []
+        for geom in self.geometries:
+            variables += geom.variables
+        return sorted(list(set(variables)))
 
     @abstractmethod
     def is_inside(self, point, check_border=False):
diff --git a/pina/domain/union_domain.py b/pina/domain/union_domain.py
index 0af8e1bd1..91aa5fb93 100644
--- a/pina/domain/union_domain.py
+++ b/pina/domain/union_domain.py
@@ -39,13 +39,6 @@ def sample_modes(self):
             set([geom.sample_modes for geom in self.geometries])
         )
 
-    @property
-    def variables(self):
-        variables = []
-        for geom in self.geometries:
-            variables += geom.variables
-        return list(set(variables))
-
     def is_inside(self, point, check_border=False):
         """
         Check if a point is inside the ``Union`` domain.

From cf511aefa464a3c874b50a14515fe74b9def7241 Mon Sep 17 00:00:00 2001
From: Nicola Demo <demo.nicola@gmail.com>
Date: Thu, 16 Jan 2025 16:06:30 +0100
Subject: [PATCH 38/55] translate setup to pyproject

---
 pyproject.toml | 38 ++++++++++++++++++++++++++
 setup.py       | 74 --------------------------------------------------
 2 files changed, 38 insertions(+), 74 deletions(-)
 create mode 100644 pyproject.toml
 delete mode 100644 setup.py

diff --git a/pyproject.toml b/pyproject.toml
new file mode 100644
index 000000000..73b065cd4
--- /dev/null
+++ b/pyproject.toml
@@ -0,0 +1,38 @@
+[build-system]
+requires = ["setuptools", "wheel"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "pina-mathlab"
+version = "0.2.0"
+description = "Physic Informed Neural networks for Advance modeling."
+readme = "README.md"
+authors = [
+    { name = "PINA Contributors" },
+]
+license = { text = "MIT" }
+keywords = [
+    "machine-learning", "deep-learning", "modeling", "pytorch", "ode", 
+    "neural-networks", "differential-equations", "pde", "hacktoberfest", 
+    "pinn", "physics-informed", "physics-informed-neural-networks", "neural-operators", "equation-learning", "lightining"
+]
+dependencies = [
+    "numpy", "matplotlib", "lightning", "torch_geometric", "pytorch_lightning"
+]
+requires-python = ">=3.8"
+
+[project.optional-dependencies]
+docs = [
+    "sphinx>5.0", "sphinx_rtd_theme", "sphinx_copybutton", "sphinx_design", "pydata_sphinx_theme"
+]
+test = [
+    "pytest", "pytest-cov", "scipy"
+]
+
+[project.urls]
+Homepage = "https://github.com/mathLab/PINA"
+Repository = "https://github.com/mathLab/PINA"
+
+[tool.setuptools]
+packages = ["pina"]
+
diff --git a/setup.py b/setup.py
deleted file mode 100644
index b67323bdf..000000000
--- a/setup.py
+++ /dev/null
@@ -1,74 +0,0 @@
-from setuptools import setup, find_packages
-
-meta = {}
-with open("pina/meta.py") as fp:
-    exec(fp.read(), meta)
-
-# Package meta-data.
-IMPORTNAME = meta['__title__']
-PIPNAME = meta['__packagename__']
-DESCRIPTION = 'Physic Informed Neural networks for Advance modeling.'
-URL = 'https://github.com/mathLab/PINA'
-MAIL = meta['__mail__']
-AUTHOR = meta['__author__']
-VERSION = meta['__version__']
-KEYWORDS = 'machine-learning deep-learning modeling pytorch ode neural-networks differential-equations pde hacktoberfest pinn physics-informed physics-informed-neural-networks neural-operators equation-learning lightining'
-
-REQUIRED = [
-    'numpy', 'matplotlib', 'torch', 'lightning', 'torch_geometric',
-    'torch-cluster', 'pytorch_lightning',
-]
-
-EXTRAS = {
-    'docs': [
-        'sphinx>5.0',
-        'sphinx_rtd_theme',
-        'sphinx_copybutton',
-        'sphinx_design',
-        'pydata_sphinx_theme'
-        ],
-    'test': [
-        'pytest',
-        'pytest-cov',
-        'scipy'
-    ],
-}
-
-LDESCRIPTION = (
-    "PINA is a Python package providing an easy interface to deal with "
-    "physics-informed neural networks (PINN) for the approximation of "
-    "(differential, nonlinear, ...) functions. Based on Pytorch, PINA "
-    "offers a simple and intuitive way to formalize a specific problem "
-    "and solve it using PINN. The approximated solution of a differential "
-    "equation can be implemented using PINA in a few lines of code thanks "
-    "to the intuitive and user-friendly interface."
-)
-
-setup(
-    name=PIPNAME,
-    version=VERSION,
-    description=DESCRIPTION,
-    long_description=LDESCRIPTION,
-    author=AUTHOR,
-    author_email=MAIL,
-    classifiers=[
-        'Development Status :: 3 - Alpha',
-        'License :: OSI Approved :: MIT License',
-        'Programming Language :: Python :: 3',
-        'Programming Language :: Python :: 3.8',
-        'Programming Language :: Python :: 3.9',
-        'Programming Language :: Python :: 3.10',
-        'Programming Language :: Python :: 3.11',
-        'Programming Language :: Python :: 3.12',
-        'Intended Audience :: Science/Research',
-        'Topic :: Scientific/Engineering :: Mathematics'
-    ],
-    keywords=KEYWORDS,
-    url=URL,
-    license='MIT',
-    packages=find_packages(),
-    install_requires=REQUIRED,
-    extras_require=EXTRAS,
-    include_package_data=True,
-    zip_safe=False,
-)

From a1b51898a01bf364784c92071372a1230c8d5313 Mon Sep 17 00:00:00 2001
From: Nicola Demo <demo.nicola@gmail.com>
Date: Thu, 16 Jan 2025 19:03:18 +0100
Subject: [PATCH 39/55] fix some tests

---
 pina/problem/zoo/__init__.py                  |  5 ++
 pina/problem/zoo/poisson_2d_square.py         | 44 ++++++++++++++++
 .../test_adaptive_refinment_callbacks.py      | 52 ++-----------------
 tests/test_callbacks/test_metric_tracker.py   | 50 +-----------------
 .../test_optimizer_callbacks.py               | 48 +----------------
 tests/test_callbacks/test_progress_bar.py     | 50 +-----------------
 tests/test_package.py                         |  2 +-
 .../test_poisson_2d_square.py                 |  7 +++
 8 files changed, 66 insertions(+), 192 deletions(-)
 create mode 100644 pina/problem/zoo/__init__.py
 create mode 100644 pina/problem/zoo/poisson_2d_square.py
 create mode 100644 tests/test_problem_zoo/test_poisson_2d_square.py

diff --git a/pina/problem/zoo/__init__.py b/pina/problem/zoo/__init__.py
new file mode 100644
index 000000000..ea9aa7876
--- /dev/null
+++ b/pina/problem/zoo/__init__.py
@@ -0,0 +1,5 @@
+__all__ = [
+    'Poisson2DSquareProblem'
+]
+
+from .poisson_2d_square import Poisson2DSquareProblem
\ No newline at end of file
diff --git a/pina/problem/zoo/poisson_2d_square.py b/pina/problem/zoo/poisson_2d_square.py
new file mode 100644
index 000000000..16329abb4
--- /dev/null
+++ b/pina/problem/zoo/poisson_2d_square.py
@@ -0,0 +1,44 @@
+""" Definition of the Poisson problem on a square domain."""
+
+from pina.problem import SpatialProblem
+from pina.operators import laplacian
+from pina import LabelTensor, Condition
+from pina.domain import CartesianDomain
+from pina.equation.equation import Equation
+from pina.equation.equation_factory import FixedValue
+import torch
+
+def laplace_equation(input_, output_):
+    force_term = (torch.sin(input_.extract(['x']) * torch.pi) *
+                  torch.sin(input_.extract(['y']) * torch.pi))
+    delta_u = laplacian(output_.extract(['u']), input_)
+    return delta_u - force_term
+
+
+my_laplace = Equation(laplace_equation)
+
+
+class Poisson2DSquareProblem(SpatialProblem):
+    output_variables = ['u']
+    spatial_domain = CartesianDomain({'x': [0, 1], 'y': [0, 1]})
+
+    domains = {
+        'D': CartesianDomain({'x': [0, 1], 'y': [0, 1]}),
+        'g1': CartesianDomain({'x': [0, 1], 'y': 1}),
+        'g2': CartesianDomain({'x': [0, 1], 'y': 0}),
+        'g3': CartesianDomain({'x': 1, 'y': [0, 1]}),
+        'g4': CartesianDomain({'x': 0, 'y': [0, 1]}),
+    }
+
+    conditions = {
+        'nil_g1': Condition(domain='D', equation=FixedValue(0.0)),
+        'nil_g2': Condition(domain='D', equation=FixedValue(0.0)),
+        'nil_g3': Condition(domain='D', equation=FixedValue(0.0)),
+        'nil_g4': Condition(domain='D', equation=FixedValue(0.0)),
+        'laplace_D': Condition(domain='D', equation=my_laplace),
+    }
+
+    def poisson_sol(self, pts):
+        return -(torch.sin(pts.extract(['x']) * torch.pi) *
+                 torch.sin(pts.extract(['y']) * torch.pi))
+    
diff --git a/tests/test_callbacks/test_adaptive_refinment_callbacks.py b/tests/test_callbacks/test_adaptive_refinment_callbacks.py
index 67732b3f5..564aab97c 100644
--- a/tests/test_callbacks/test_adaptive_refinment_callbacks.py
+++ b/tests/test_callbacks/test_adaptive_refinment_callbacks.py
@@ -1,59 +1,13 @@
-from pina.callbacks import R3Refinement
-import torch
-import pytest
-
-from pina.problem import SpatialProblem
-from pina.operators import laplacian
-from pina.domain import CartesianDomain
-from pina import Condition, LabelTensor
 from pina.solvers import PINN
 from pina.trainer import Trainer
 from pina.model import FeedForward
-from pina.equation.equation import Equation
-from pina.equation.equation_factory import FixedValue
-
-
-def laplace_equation(input_, output_):
-    force_term = (torch.sin(input_.extract(['x']) * torch.pi) *
-                  torch.sin(input_.extract(['y']) * torch.pi))
-    delta_u = laplacian(output_.extract(['u']), input_)
-    return delta_u - force_term
-
-
-my_laplace = Equation(laplace_equation)
-in_ = LabelTensor(torch.tensor([[0., 1.]]), ['x', 'y'])
-out_ = LabelTensor(torch.tensor([[0.]]), ['u'])
-
-
-class Poisson(SpatialProblem):
-    output_variables = ['u']
-    spatial_domain = CartesianDomain({'x': [0, 1], 'y': [0, 1]})
-
-    conditions = {
-        'gamma1': Condition(
-            location=CartesianDomain({'x': [0, 1], 'y':  1}),
-            equation=FixedValue(0.0)),
-        'gamma2': Condition(
-            location=CartesianDomain({'x': [0, 1], 'y': 0}),
-            equation=FixedValue(0.0)),
-        'gamma3': Condition(
-            location=CartesianDomain({'x':  1, 'y': [0, 1]}),
-            equation=FixedValue(0.0)),
-        'gamma4': Condition(
-            location=CartesianDomain({'x': 0, 'y': [0, 1]}),
-            equation=FixedValue(0.0)),
-        'D': Condition(
-            input_points=LabelTensor(torch.rand(size=(100, 2)), ['x', 'y']),
-            equation=my_laplace),
-        # 'data': Condition(
-        #     input_points=in_,
-        #     output_points=out_)
-    }
+from pina.problem.zoo import Poisson2DSquareProblem as Poisson
+from pina.callbacks import R3Refinement
 
 
 # make the problem
 poisson_problem = Poisson()
-boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4']
+boundaries = ['nil_g1', 'nil_g2', 'nil_g3', 'nil_g4']
 n = 10
 poisson_problem.discretise_domain(n, 'grid', locations=boundaries)
 model = FeedForward(len(poisson_problem.input_variables),
diff --git a/tests/test_callbacks/test_metric_tracker.py b/tests/test_callbacks/test_metric_tracker.py
index c38024587..d67e06b85 100644
--- a/tests/test_callbacks/test_metric_tracker.py
+++ b/tests/test_callbacks/test_metric_tracker.py
@@ -1,59 +1,13 @@
-import torch
-import pytest
-
-from pina.problem import SpatialProblem
-from pina.operators import laplacian
-from pina.geometry import CartesianDomain
-from pina import Condition, LabelTensor
 from pina.solvers import PINN
 from pina.trainer import Trainer
 from pina.model import FeedForward
-from pina.equation.equation import Equation
-from pina.equation.equation_factory import FixedValue
 from pina.callbacks import MetricTracker
-
-
-def laplace_equation(input_, output_):
-    force_term = (torch.sin(input_.extract(['x']) * torch.pi) *
-                  torch.sin(input_.extract(['y']) * torch.pi))
-    delta_u = laplacian(output_.extract(['u']), input_)
-    return delta_u - force_term
-
-
-my_laplace = Equation(laplace_equation)
-in_ = LabelTensor(torch.tensor([[0., 1.]]), ['x', 'y'])
-out_ = LabelTensor(torch.tensor([[0.]]), ['u'])
-
-
-class Poisson(SpatialProblem):
-    output_variables = ['u']
-    spatial_domain = CartesianDomain({'x': [0, 1], 'y': [0, 1]})
-
-    conditions = {
-        'gamma1': Condition(
-            location=CartesianDomain({'x': [0, 1], 'y':  1}),
-            equation=FixedValue(0.0)),
-        'gamma2': Condition(
-            location=CartesianDomain({'x': [0, 1], 'y': 0}),
-            equation=FixedValue(0.0)),
-        'gamma3': Condition(
-            location=CartesianDomain({'x':  1, 'y': [0, 1]}),
-            equation=FixedValue(0.0)),
-        'gamma4': Condition(
-            location=CartesianDomain({'x': 0, 'y': [0, 1]}),
-            equation=FixedValue(0.0)),
-        'D': Condition(
-            input_points=LabelTensor(torch.rand(size=(100, 2)), ['x', 'y']),
-            equation=my_laplace),
-        'data': Condition(
-            input_points=in_,
-            output_points=out_)
-    }
+from pina.problem.zoo import Poisson2DSquareProblem as Poisson
 
 
 # make the problem
 poisson_problem = Poisson()
-boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4']
+boundaries = ['nil_g1', 'nil_g2', 'nil_g3', 'nil_g4']
 n = 10
 poisson_problem.discretise_domain(n, 'grid', locations=boundaries)
 model = FeedForward(len(poisson_problem.input_variables),
diff --git a/tests/test_callbacks/test_optimizer_callbacks.py b/tests/test_callbacks/test_optimizer_callbacks.py
index 898d3f502..a57bda35d 100644
--- a/tests/test_callbacks/test_optimizer_callbacks.py
+++ b/tests/test_callbacks/test_optimizer_callbacks.py
@@ -2,58 +2,14 @@
 import torch
 import pytest
 
-from pina.problem import SpatialProblem
-from pina.operators import laplacian
-from pina.domain import CartesianDomain
-from pina import Condition, LabelTensor
 from pina.solvers import PINN
 from pina.trainer import Trainer
 from pina.model import FeedForward
-from pina.equation.equation import Equation
-from pina.equation.equation_factory import FixedValue
-
-
-def laplace_equation(input_, output_):
-    force_term = (torch.sin(input_.extract(['x']) * torch.pi) *
-                  torch.sin(input_.extract(['y']) * torch.pi))
-    delta_u = laplacian(output_.extract(['u']), input_)
-    return delta_u - force_term
-
-
-my_laplace = Equation(laplace_equation)
-in_ = LabelTensor(torch.tensor([[0., 1.]]), ['x', 'y'])
-out_ = LabelTensor(torch.tensor([[0.]]), ['u'])
-
-
-class Poisson(SpatialProblem):
-    output_variables = ['u']
-    spatial_domain = CartesianDomain({'x': [0, 1], 'y': [0, 1]})
-
-    conditions = {
-        'gamma1': Condition(
-            location=CartesianDomain({'x': [0, 1], 'y':  1}),
-            equation=FixedValue(0.0)),
-        'gamma2': Condition(
-            location=CartesianDomain({'x': [0, 1], 'y': 0}),
-            equation=FixedValue(0.0)),
-        'gamma3': Condition(
-            location=CartesianDomain({'x':  1, 'y': [0, 1]}),
-            equation=FixedValue(0.0)),
-        'gamma4': Condition(
-            location=CartesianDomain({'x': 0, 'y': [0, 1]}),
-            equation=FixedValue(0.0)),
-        'D': Condition(
-            input_points=LabelTensor(torch.rand(size=(100, 2)), ['x', 'y']),
-            equation=my_laplace),
-        # 'data': Condition(
-        #     input_points=in_,
-        #     output_points=out_)
-    }
-
+from pina.problem.zoo import Poisson2DSquareProblem as Poisson
 
 # make the problem
 poisson_problem = Poisson()
-boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4']
+boundaries = ['nil_g1', 'nil_g2', 'nil_g3', 'nil_g4']
 n = 10
 poisson_problem.discretise_domain(n, 'grid', locations=boundaries)
 model = FeedForward(len(poisson_problem.input_variables),
diff --git a/tests/test_callbacks/test_progress_bar.py b/tests/test_callbacks/test_progress_bar.py
index 990b471fc..5c2cae4b1 100644
--- a/tests/test_callbacks/test_progress_bar.py
+++ b/tests/test_callbacks/test_progress_bar.py
@@ -1,59 +1,13 @@
-import torch
-import pytest
-
-from pina.problem import SpatialProblem
-from pina.operators import laplacian
-from pina.geometry import CartesianDomain
-from pina import Condition, LabelTensor
 from pina.solvers import PINN
 from pina.trainer import Trainer
 from pina.model import FeedForward
-from pina.equation.equation import Equation
-from pina.equation.equation_factory import FixedValue
 from pina.callbacks.processing_callbacks import PINAProgressBar
-
-
-def laplace_equation(input_, output_):
-    force_term = (torch.sin(input_.extract(['x']) * torch.pi) *
-                  torch.sin(input_.extract(['y']) * torch.pi))
-    delta_u = laplacian(output_.extract(['u']), input_)
-    return delta_u - force_term
-
-
-my_laplace = Equation(laplace_equation)
-in_ = LabelTensor(torch.tensor([[0., 1.]]), ['x', 'y'])
-out_ = LabelTensor(torch.tensor([[0.]]), ['u'])
-
-
-class Poisson(SpatialProblem):
-    output_variables = ['u']
-    spatial_domain = CartesianDomain({'x': [0, 1], 'y': [0, 1]})
-
-    conditions = {
-        'gamma1': Condition(
-            location=CartesianDomain({'x': [0, 1], 'y':  1}),
-            equation=FixedValue(0.0)),
-        'gamma2': Condition(
-            location=CartesianDomain({'x': [0, 1], 'y': 0}),
-            equation=FixedValue(0.0)),
-        'gamma3': Condition(
-            location=CartesianDomain({'x':  1, 'y': [0, 1]}),
-            equation=FixedValue(0.0)),
-        'gamma4': Condition(
-            location=CartesianDomain({'x': 0, 'y': [0, 1]}),
-            equation=FixedValue(0.0)),
-        'D': Condition(
-            input_points=LabelTensor(torch.rand(size=(100, 2)), ['x', 'y']),
-            equation=my_laplace),
-        'data': Condition(
-            input_points=in_,
-            output_points=out_)
-    }
+from pina.problem.zoo import Poisson2DSquareProblem as Poisson
 
 
 # make the problem
 poisson_problem = Poisson()
-boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4']
+boundaries = ['nil_g1', 'nil_g2', 'nil_g3', 'nil_g4']
 n = 10
 poisson_problem.discretise_domain(n, 'grid', locations=boundaries)
 model = FeedForward(len(poisson_problem.input_variables),
diff --git a/tests/test_package.py b/tests/test_package.py
index f59bd6c21..f85bed550 100644
--- a/tests/test_package.py
+++ b/tests/test_package.py
@@ -1,2 +1,2 @@
 def test_import():
-    import pina
+    import pina
\ No newline at end of file
diff --git a/tests/test_problem_zoo/test_poisson_2d_square.py b/tests/test_problem_zoo/test_poisson_2d_square.py
new file mode 100644
index 000000000..6541a6ed9
--- /dev/null
+++ b/tests/test_problem_zoo/test_poisson_2d_square.py
@@ -0,0 +1,7 @@
+import torch
+import pytest
+
+from pina.problem.zoo import Poisson2DSquareProblem
+
+def test_constructor():
+    Poisson2DSquareProblem()
\ No newline at end of file

From 8e803cb60864e7e5ecdf953042377106711260ea Mon Sep 17 00:00:00 2001
From: FilippoOlivo <filippo@filippoolivo.com>
Date: Thu, 16 Jan 2025 19:42:46 +0100
Subject: [PATCH 40/55] Add check conditions-solver consistency

---
 pina/solvers/pinns/basepinn.py | 17 ++++++++++++-----
 pina/solvers/solver.py         | 20 +++++---------------
 pina/solvers/supervised.py     |  5 +++--
 3 files changed, 20 insertions(+), 22 deletions(-)

diff --git a/pina/solvers/pinns/basepinn.py b/pina/solvers/pinns/basepinn.py
index f44ea1d0d..611a3f40e 100644
--- a/pina/solvers/pinns/basepinn.py
+++ b/pina/solvers/pinns/basepinn.py
@@ -8,6 +8,8 @@
 from ...loss.loss_interface import LossInterface
 from ...problem import InverseProblem
 from ...optim import TorchOptimizer, TorchScheduler
+from ...condition import InputOutputPointsCondition, \
+    InputPointsEquationCondition, DomainEquationCondition
 
 torch.pi = torch.acos(torch.zeros(1)).item() * 2  # which is 3.1415927410125732
 
@@ -24,6 +26,8 @@ class PINNInterface(SolverInterface, metaclass=ABCMeta):
     to the user to choose which problem the implemented solver inheriting from
     this class is suitable for.
     """
+    accepted_conditions_types = (InputOutputPointsCondition,
+                                 InputPointsEquationCondition, DomainEquationCondition)
 
     def __init__(
             self,
@@ -97,7 +101,6 @@ def __init__(
         self._optimizer = self._pina_optimizers[0]
         self._scheduler = self._pina_schedulers[0]
 
-
     def training_step(self, batch):
         """
         The Physics Informed Solver Training Step. This function takes care
@@ -117,14 +120,16 @@ def training_step(self, batch):
             if 'output_points' in points:
                 input_pts, output_pts = points['input_points'], points['output_points']
 
-                loss_ = self.loss_data(input_pts=input_pts, output_pts=output_pts)
+                loss_ = self.loss_data(
+                    input_pts=input_pts, output_pts=output_pts)
                 condition_loss.append(loss_.as_subclass(torch.Tensor))
             else:
                 input_pts = points['input_points']
 
                 condition = self.problem.conditions[condition_name]
 
-                loss_ = self.loss_phys(input_pts.requires_grad_(), condition.equation)
+                loss_ = self.loss_phys(
+                    input_pts.requires_grad_(), condition.equation)
                 condition_loss.append(loss_.as_subclass(torch.Tensor))
             condition_loss.append(loss_.as_subclass(torch.Tensor))
         # clamp unknown parameters in InverseProblem (if needed)
@@ -144,14 +149,16 @@ def validation_step(self, batch):
         for condition_name, points in batch:
             if 'output_points' in points:
                 input_pts, output_pts = points['input_points'], points['output_points']
-                loss_ = self.loss_data(input_pts=input_pts, output_pts=output_pts)
+                loss_ = self.loss_data(
+                    input_pts=input_pts, output_pts=output_pts)
                 condition_loss.append(loss_.as_subclass(torch.Tensor))
             else:
                 input_pts = points['input_points']
 
                 condition = self.problem.conditions[condition_name]
                 with torch.set_grad_enabled(True):
-                    loss_ = self.loss_phys(input_pts.requires_grad_(), condition.equation)
+                    loss_ = self.loss_phys(
+                        input_pts.requires_grad_(), condition.equation)
                 condition_loss.append(loss_.as_subclass(torch.Tensor))
             condition_loss.append(loss_.as_subclass(torch.Tensor))
         # clamp unknown parameters in InverseProblem (if needed)
diff --git a/pina/solvers/solver.py b/pina/solvers/solver.py
index 8052b4b8e..408aee5b2 100644
--- a/pina/solvers/solver.py
+++ b/pina/solvers/solver.py
@@ -10,7 +10,6 @@
 import sys
 
 
-
 class SolverInterface(lightning.pytorch.LightningModule, metaclass=ABCMeta):
     """
     Solver base class. This class inherits is a wrapper of
@@ -133,23 +132,14 @@ def on_train_start(self):
 
         return super().on_train_start()
 
-    def _check_solver_consistency(self, problem):
-        pass
-        #TODO : Implement this method for the conditions
-        '''
-        
-    
-        for _, condition in problem.conditions.items():
-            if not set(condition.condition_type).issubset(
-                    set(self.accepted_condition_types)):
-                raise ValueError(
-                    f'{self.__name__} dose not support condition '
-                    f'{condition.condition_type}')
-        '''
     @staticmethod
     def get_batch_size(batch):
         # Assuming batch is your custom Batch object
         batch_size = 0
         for data in batch:
             batch_size += len(data[1]['input_points'])
-        return batch_size
\ No newline at end of file
+        return batch_size
+
+    def _check_solver_consistency(self, problem):
+        for condition in problem.conditions.values():
+            check_consistency(condition, self.accepted_conditions_types)
diff --git a/pina/solvers/supervised.py b/pina/solvers/supervised.py
index bce4b31a8..c7f5f667f 100644
--- a/pina/solvers/supervised.py
+++ b/pina/solvers/supervised.py
@@ -7,6 +7,7 @@
 from ..label_tensor import LabelTensor
 from ..utils import check_consistency
 from ..loss.loss_interface import LossInterface
+from ..condition import InputOutputPointsCondition
 
 
 class SupervisedSolver(SolverInterface):
@@ -37,7 +38,8 @@ class SupervisedSolver(SolverInterface):
     we are seeking to approximate multiple (discretised) functions given
     multiple (discretised) input functions.
     """
-    __name__ = 'SupervisedSolver'
+
+    accepted_conditions_types = InputOutputPointsCondition
 
     def __init__(self,
                  problem,
@@ -143,7 +145,6 @@ def validation_step(self, batch):
         self.log('val_loss', loss, prog_bar=True, logger=True,
                  batch_size=self.get_batch_size(batch), sync_dist=True)
 
-
     def test_step(self, batch, batch_idx) -> STEP_OUTPUT:
         """
         Solver test step.

From 96783c3f2e8facaea897c94de4cf242ad2b69d74 Mon Sep 17 00:00:00 2001
From: giovanni <giovanni.canali98@yahoo.it>
Date: Mon, 20 Jan 2025 13:53:08 +0100
Subject: [PATCH 41/55] fix laplacian

---
 pina/operators.py | 46 +++++++++-------------------------------------
 1 file changed, 9 insertions(+), 37 deletions(-)

diff --git a/pina/operators.py b/pina/operators.py
index 8dbfb8c76..9da462591 100644
--- a/pina/operators.py
+++ b/pina/operators.py
@@ -209,49 +209,21 @@ def scalar_laplace(output_, input_, components, d):
 
     elif method == "std":
         if len(components) == 1:
-            # result = scalar_laplace(output_, input_, components, d) # TODO check (from 0.1)
-            grad_output = grad(output_, input_, components=components, d=d)
-            to_append_tensors = []
-            for i, label in enumerate(grad_output.labels):
-                gg = grad(grad_output, input_, d=d, components=[label])
-                gg = gg.extract([gg.labels[i]])
-
-                to_append_tensors.append(gg)
+            result = scalar_laplace(output_, input_, components, d)
             labels = [f"dd{components[0]}"]
-            result = LabelTensor.summation(tensors=to_append_tensors)
-            result.labels = labels
+
         else:
-    #         result = torch.empty( # TODO check (from 0.1)
-    #             size=(input_.shape[0], len(components)),
-    #             dtype=output_.dtype,
-    #             device=output_.device,
-    #         )
-    #         labels = [None] * len(components)
-    #         for idx, c in enumerate(components):
-    #             result[:, idx] = scalar_laplace(output_, input_, c, d).flatten()
-    #             labels[idx] = f"dd{c}"
-
-    # result = result.as_subclass(LabelTensor)
-    # result.labels = labels
             result = torch.empty(
                 input_.shape[0], len(components), device=output_.device
             )
             labels = [None] * len(components)
-            to_append_tensors = [None] * len(components)
-            for idx, (ci, di) in enumerate(zip(components, d)):
-
-                if not isinstance(ci, list):
-                    ci = [ci]
-                if not isinstance(di, list):
-                    di = [di]
-
-                grad_output = grad(output_, input_, components=ci, d=di)
-                result[:, idx] = grad(grad_output, input_, d=di).flatten()
-                to_append_tensors[idx] = grad(grad_output, input_, d=di)
-                labels[idx] = f"dd{ci[0]}dd{di[0]}"
-            result = LabelTensor.cat(tensors=to_append_tensors,
-                                     dim=output_.tensor.ndim - 1)
-            result.labels = labels
+            for idx, ci in enumerate(components):
+                result[:, idx] = scalar_laplace(output_, input_, ci, d).flatten()
+                labels[idx] = f"dd{ci}"
+
+        result = result.as_subclass(LabelTensor)
+        result.labels = labels
+    
     return result
 
 

From ca1995862b17afeead026b0b064c2b8ecc1787e8 Mon Sep 17 00:00:00 2001
From: Giovanni Canali <115086358+gc031298@users.noreply.github.com>
Date: Mon, 20 Jan 2025 14:06:13 +0100
Subject: [PATCH 42/55] Codacy

---
 pina/operators.py | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/pina/operators.py b/pina/operators.py
index 9da462591..5ed6e934f 100644
--- a/pina/operators.py
+++ b/pina/operators.py
@@ -217,13 +217,13 @@ def scalar_laplace(output_, input_, components, d):
                 input_.shape[0], len(components), device=output_.device
             )
             labels = [None] * len(components)
-            for idx, ci in enumerate(components):
-                result[:, idx] = scalar_laplace(output_, input_, ci, d).flatten()
-                labels[idx] = f"dd{ci}"
+            for idx, c in enumerate(components):
+                result[:, idx] = scalar_laplace(output_, input_, c, d).flatten()
+                labels[idx] = f"dd{c}"
 
         result = result.as_subclass(LabelTensor)
         result.labels = labels
-    
+
     return result
 
 

From 8bffe2a0eb9840904571f1a71cc4372f15882209 Mon Sep 17 00:00:00 2001
From: giovanni <giovanni.canali98@yahoo.it>
Date: Mon, 20 Jan 2025 15:35:10 +0100
Subject: [PATCH 43/55] update test laplacian

---
 tests/test_operators.py | 33 +++++++++++++++++++++++++++++++--
 1 file changed, 31 insertions(+), 2 deletions(-)

diff --git a/tests/test_operators.py b/tests/test_operators.py
index 35c07911b..ccfc17b08 100644
--- a/tests/test_operators.py
+++ b/tests/test_operators.py
@@ -17,9 +17,9 @@ def func_scalar(x):
 
 
 data = torch.rand((20, 3))
-inp = LabelTensor(data, ['x', 'y', 'mu']).requires_grad_(True)
+inp = LabelTensor(data, ['x', 'y', 'z']).requires_grad_(True)
 labels = ['a', 'b', 'c']
-tensor_v = LabelTensor(func_vec(inp), labels)
+tensor_v = LabelTensor(func_vector(inp), labels)
 tensor_s = LabelTensor(func_scalar(inp).reshape(-1, 1), labels[0])
 
 
@@ -107,6 +107,8 @@ def test_laplacian_scalar_output():
 
 def test_laplacian_vector_output():
     laplace_tensor_v = laplacian(tensor_v, inp)
+    print(laplace_tensor_v.labels)
+    print(tensor_v.labels)
     true_val = 2*torch.ones_like(tensor_v)
     assert laplace_tensor_v.shape == tensor_v.shape
     assert laplace_tensor_v.labels == [
@@ -124,3 +126,30 @@ def test_laplacian_vector_output():
         f'dd{i}' for i in ['a', 'b']
     ]
     assert torch.allclose(laplace_tensor_v, true_val)
+
+def test_laplacian_vector_output2():
+    x = LabelTensor(torch.linspace(0,1,10, requires_grad=True).reshape(-1,1), labels = ['x'])
+    y = LabelTensor(torch.linspace(3,4,10, requires_grad=True).reshape(-1,1), labels = ['y'])
+    input_ = LabelTensor(torch.cat((x,y), dim = 1), labels = ['x', 'y'])
+
+    # Construct two scalar functions:
+    # u = x**2 + y**2
+    # v = x**2 - y**2
+    u = LabelTensor(input_.extract('x')**2 + input_.extract('y')**2, labels='u')
+    v = LabelTensor(input_.extract('x')**2 - input_.extract('y')**2, labels='v')
+
+    # Define a vector-valued function, whose components are u and v.
+    f = LabelTensor(torch.cat((u,v), dim = 1), labels = ['u', 'v'])
+
+    # Compute the scalar laplacian of both u and v:
+    # Lap(u) = [4, 4, 4, ..., 4]
+    # Lap(v) = [0, 0, 0, ..., 0]
+    lap_u = laplacian(u, input_, components=['u'])
+    lap_v = laplacian(v, input_, components=['v'])
+
+    # Compute the laplacian of f: the two columns should correspond
+    # to the laplacians of u and v, respectively...
+    lap_f = laplacian(f, input_, components=['u', 'v'])
+
+    assert torch.allclose(lap_f.extract('ddu'), lap_u)
+    assert torch.allclose(lap_f.extract('ddv'), lap_v)

From b6157a6bddd2c9f907f3fcceda61ae0b740c4a77 Mon Sep 17 00:00:00 2001
From: Giovanni Canali <115086358+gc031298@users.noreply.github.com>
Date: Wed, 22 Jan 2025 10:48:42 +0100
Subject: [PATCH 44/55] Update import from lightning.pytorch (#409)

* update import
* Remove unnecessary import return type

---------

Co-authored-by: Filippo Olivo <filippo@filippoolivo.com>
---
 pina/callbacks/adaptive_refinment_callbacks.py | 2 +-
 pina/callbacks/optimizer_callbacks.py          | 2 +-
 pina/callbacks/processing_callbacks.py         | 6 +++---
 pina/solvers/supervised.py                     | 3 +--
 4 files changed, 6 insertions(+), 7 deletions(-)

diff --git a/pina/callbacks/adaptive_refinment_callbacks.py b/pina/callbacks/adaptive_refinment_callbacks.py
index 5af2cc859..00d6dfb80 100644
--- a/pina/callbacks/adaptive_refinment_callbacks.py
+++ b/pina/callbacks/adaptive_refinment_callbacks.py
@@ -1,7 +1,7 @@
 """PINA Callbacks Implementations"""
 
 import torch
-from pytorch_lightning.callbacks import Callback
+from lightning.pytorch.callbacks import Callback
 from ..label_tensor import LabelTensor
 from ..utils import check_consistency
 
diff --git a/pina/callbacks/optimizer_callbacks.py b/pina/callbacks/optimizer_callbacks.py
index c11db8894..6905ebf22 100644
--- a/pina/callbacks/optimizer_callbacks.py
+++ b/pina/callbacks/optimizer_callbacks.py
@@ -1,6 +1,6 @@
 """PINA Callbacks Implementations"""
 
-from pytorch_lightning.callbacks import Callback
+from lightning.pytorch.callbacks import Callback
 import torch
 from ..utils import check_consistency
 
diff --git a/pina/callbacks/processing_callbacks.py b/pina/callbacks/processing_callbacks.py
index a70218eb1..e65506c17 100644
--- a/pina/callbacks/processing_callbacks.py
+++ b/pina/callbacks/processing_callbacks.py
@@ -1,11 +1,11 @@
 """PINA Callbacks Implementations"""
 
-from pytorch_lightning.core.module import LightningModule
-from pytorch_lightning.trainer.trainer import Trainer
+from lightning.pytorch.core.module import LightningModule
+from lightning.pytorch.trainer.trainer import Trainer
 import torch
 import copy
 
-from pytorch_lightning.callbacks import Callback, TQDMProgressBar
+from lightning.pytorch.callbacks import Callback, TQDMProgressBar
 from lightning.pytorch.callbacks.progress.progress_bar import (
     get_standard_metrics,
 )
diff --git a/pina/solvers/supervised.py b/pina/solvers/supervised.py
index c7f5f667f..99fec09a3 100644
--- a/pina/solvers/supervised.py
+++ b/pina/solvers/supervised.py
@@ -1,6 +1,5 @@
 """ Module for SupervisedSolver """
 import torch
-from pytorch_lightning.utilities.types import STEP_OUTPUT
 from torch.nn.modules.loss import _Loss
 from ..optim import TorchOptimizer, TorchScheduler
 from .solver import SolverInterface
@@ -145,7 +144,7 @@ def validation_step(self, batch):
         self.log('val_loss', loss, prog_bar=True, logger=True,
                  batch_size=self.get_batch_size(batch), sync_dist=True)
 
-    def test_step(self, batch, batch_idx) -> STEP_OUTPUT:
+    def test_step(self, batch, batch_idx):
         """
         Solver test step.
         """

From 8f52da7cafc16a638a80787c7fac03f10199e55e Mon Sep 17 00:00:00 2001
From: FilippoOlivo <filippo@filippoolivo.com>
Date: Tue, 21 Jan 2025 10:44:35 +0100
Subject: [PATCH 45/55] Simplify LabelTensor class, fix #395, add docstrings,
 and resolve Python 3.8 compatibility issue in tests

---
 pina/label_tensor.py                         | 541 ++++++++++---------
 tests/test_label_tensor/test_label_tensor.py |   3 +-
 2 files changed, 300 insertions(+), 244 deletions(-)

diff --git a/pina/label_tensor.py b/pina/label_tensor.py
index 631c5253c..5ffa61138 100644
--- a/pina/label_tensor.py
+++ b/pina/label_tensor.py
@@ -4,18 +4,13 @@
 from torch import Tensor
 
 
-full_labels = False
-MATH_FUNCTIONS = {torch.sin, torch.cos}
-
 class LabelTensor(torch.Tensor):
     """Torch tensor with a label for any column."""
 
     @staticmethod
     def __new__(cls, x, labels, *args, **kwargs):
-        full = kwargs.pop("full", full_labels)
 
         if isinstance(x, LabelTensor):
-            x.full = full
             return x
         return super().__new__(cls, x, *args, **kwargs)
 
@@ -34,7 +29,7 @@ def __init__(self, x, labels, **kwargs):
                     {1: {"name": "space"['a', 'b', 'c'])
 
         """
-        self.full = kwargs.get('full', full_labels)
+        super().__init__()
         if labels is not None:
             self.labels = labels
         else:
@@ -95,51 +90,49 @@ def labels(self, labels):
         else:
             raise ValueError("labels must be list, dict or string.")
 
-    def _init_labels_from_dict(self, labels):
+    def _init_labels_from_dict(self, labels: dict):
         """
-            Update the internal label representation according to the values
-            passed as input.
+        Update the internal label representation according to the values
+        passed as input.
 
-            :param labels: The label(s) to update.
-            :type labels: dict
-            :raises ValueError: dof list contain duplicates or number of dof
-            does not match with tensor shape
-            """
+        :param labels: The label(s) to update.
+        :type labels: dict
+        :raises ValueError: If the dof list contains duplicates or the number of
+                            dof does not match the tensor shape.
+        """
         tensor_shape = self.shape
 
-        # Set all labels if full_labels is True
-        if hasattr(self, 'full') and self.full:
-            labels = {
-                i: labels[i] if i in labels else {
-                    'name': i, 'dof': range(tensor_shape[i])
-                }
-                for i in range(len(tensor_shape))
-            }
-
-        for k, v in labels.items():
-
-            # Init labels from str
-            if isinstance(v, str):
-                v = {'name': v, 'dof': range(tensor_shape[k])}
-
-            # Init labels from dict
-            elif isinstance(v, dict):
-                # Only name of the dimension if provided
-                if list(v.keys()) == ['name']:
-                    v['dof'] = range(tensor_shape[k])
-                # Both name and dof are provided
-                elif sorted(list(v.keys())) == ['dof', 'name']:
-                    dof_list = v['dof']
-                    dof_len = len(dof_list)
-                    if dof_len != len(set(dof_list)):
-                        raise ValueError("dof must be unique")
-                    if dof_len != tensor_shape[k]:
-                        raise ValueError(
-                            'Number of dof does not match tensor shape')
+        def validate_dof(dof_list, dim_size: int):
+            """Validate the 'dof' list for uniqueness and size."""
+            if len(dof_list) != len(set(dof_list)):
+                raise ValueError("dof must be unique")
+            if len(dof_list) != dim_size:
+                raise ValueError(
+                    f"Number of dof ({len(dof_list)}) does not match "
+                    f"tensor shape ({dim_size})")
+
+        for dim, label in labels.items():
+            if isinstance(label, dict):
+                if 'name' not in label:
+                    label['name'] = dim
+                if 'dof' not in label:
+                    label['dof'] = range(tensor_shape[dim])
+                if 'dof' in label and 'name' in label:
+                    dof = label['dof']
+                    dof_list = dof if isinstance(dof, (list, range)) else [dof]
+                    if not isinstance(dof_list, (list, range)):
+                        raise ValueError(f"'dof' should be a list or range, not"
+                                         f" {type(dof_list)}")
+                    validate_dof(dof_list, tensor_shape[dim])
+                else:
+                    raise ValueError("Labels dictionary must contain either "
+                                     " both 'name' and 'dof' keys")
             else:
-                raise ValueError('Illegal labels initialization')
-            # Assign labels values
-            self._labels[k] = v
+                raise ValueError(f"Invalid label format for {dim}: Expected "
+                                 f"list or dictionary, got {type(label)}")
+
+            # Assign validated label data to internal labels
+            self._labels[dim] = label
 
     def _init_labels_from_list(self, labels):
         """
@@ -168,61 +161,54 @@ def extract(self, labels_to_extract):
         :raises TypeError: Labels are not ``str``.
         :raises ValueError: Label to extract is not in the labels ``list``.
         """
-        # Convert str/int to string
-        def find_names(labels):
-            dim_names = {}
-            for dim in labels.keys():
-                dim_names[labels[dim]['name']] = dim
-            return dim_names
 
+        def get_label_indices(dim_labels, labels_te):
+            if isinstance(labels_te, (int, str)):
+                labels_te = [labels_te]
+            return [dim_labels.index(label) for label in labels_te] if len(
+                labels_te) > 1 else slice(dim_labels.index(labels_te[0]),
+                                          dim_labels.index(labels_te[0]) + 1)
+
+        # Ensure labels_to_extract is a list or dict
         if isinstance(labels_to_extract, (str, int)):
             labels_to_extract = [labels_to_extract]
 
-        # Store useful variables
         labels = copy(self._labels)
-        stored_keys = labels.keys()
-        dim_names = find_names(labels)
-        ndim = len(super().shape)
-
-        # Convert tuple/list to dict (having a list as input
-        # means that we want to extract a values from the last dimension)
-        if isinstance(labels_to_extract, (tuple, list)):
-            if not ndim - 1 in stored_keys:
-                raise ValueError(
-                    "LabelTensor does not have labels in last dimension")
-            name = labels[ndim-1]['name']
-            labels_to_extract = {name: list(labels_to_extract)}
 
-        # If labels_to_extract is not dict then rise error
-        if not isinstance(labels_to_extract, dict):
-            raise ValueError('labels_to_extract must be str or list or dict')
+        # Get the dimension names and the respective dimension index
+        dim_names = {labels[dim]['name']: dim for dim in labels.keys()}
+        ndim = super().ndim
+        tensor = self.tensor.as_subclass(torch.Tensor)
 
-        # Initialize list used to perform extraction
-        extractor = [slice(None)]*ndim
+        # Convert list/tuple to a dict for the last dimension if applicable
+        if isinstance(labels_to_extract, (list, tuple)):
+            last_dim = ndim - 1
+            dim_name = labels[last_dim]['name']
+            labels_to_extract = {dim_name: list(labels_to_extract)}
 
-        # Loop over labels_to_extract dict
-        for dim_name, labels_te in labels_to_extract.items():
+        # Validate the labels_to_extract type
+        if not isinstance(labels_to_extract, dict):
+            raise ValueError(
+                "labels_to_extract must be a string, list, or dictionary.")
 
-            # If label is not find raise value error
-            idx_dim = dim_names.get(dim_name, None)
-            if idx_dim is None:
+        # Perform the extraction for each specified dimension
+        for dim_name, labels_te in labels_to_extract.items():
+            if dim_name not in dim_names:
                 raise ValueError(
-                    'Cannot extract label with is not in original labels')
+                    f"Cannot extract labels for dimension '{dim_name}' as it is"
+                    f" not present in the original labels.")
 
+            idx_dim = dim_names[dim_name]
             dim_labels = labels[idx_dim]['dof']
-            labels_te = [labels_te] if isinstance(labels_te, (int, str)) else labels_te
-            if not isinstance(labels_te, range):
-                #If is done to keep the dimension if there is only one extracted label
-                extractor[idx_dim] = [dim_labels.index(i) for i in labels_te] \
-                    if len(labels_te)>1 else slice(dim_labels.index(labels_te[0]), dim_labels.index(labels_te[0])+1)
-            else:
-                extractor[idx_dim] = slice(labels_te.start, labels_te.stop)
+            indices = get_label_indices(dim_labels, labels_te)
 
-            labels.update({idx_dim: {'dof': labels_te, 'name': dim_name}})
+            extractor = [slice(None)] * ndim
+            extractor[idx_dim] = indices
+            tensor = tensor[tuple(extractor)]
 
-        tensor = super().__getitem__(extractor).as_subclass(LabelTensor)
-        tensor._labels = labels
-        return tensor
+            labels[idx_dim] = {'dof': labels_te, 'name': dim_name}
+
+        return LabelTensor(tensor, labels)
 
     def __str__(self):
         """
@@ -243,62 +229,82 @@ def cat(tensors, dim=0):
         the resulting tensor is of shape `(n+n',m,dof)`
 
         :param tensors: tensors to concatenate
-        :type tensors: list(LabelTensor)
+        :type tensors: list of LabelTensor
         :param dim: dimensions on which you want to perform the operation
-        (default 0)
+        (default is 0)
         :type dim: int
         :rtype: LabelTensor
         :raises ValueError: either number dof or dimensions names differ
         """
-        if len(tensors) == 0:
-            return []
-        if len(tensors) == 1 or isinstance(tensors, LabelTensor):
-            return tensors[0]
-
-        # Perform cat on tensors
-        new_tensor = torch.cat(tensors, dim=dim)
-
-        # --------- Start definition auxiliary function ------
-        # Compute and update labels
-        def create_labels_cat(tensors, dim, tensor_shape):
-            stored_labels = [tensor.stored_labels for tensor in tensors]
-            keys = stored_labels[0].keys()
-
-            if any(not all(stored_labels[i][k] == stored_labels[0][k] for i in
-                        range(len(stored_labels))) for k in keys if k != dim):
-                raise RuntimeError('tensors must have the same shape and dof')
-
-            # Copy labels from the first tensor and update the 'dof' for dimension `dim`
-            labels = copy(stored_labels[0])
-            if dim in labels:
-                labels_list = [tensor[dim]['dof'] for tensor in stored_labels]
-                last_dim_dof = range(tensor_shape[dim]) if all(isinstance(label, range)
-                                    for label in labels_list) else sum(labels_list, [])
-                labels[dim]['dof'] = last_dim_dof
-            return labels
-        # --------- End definition auxiliary function ------
-
-        # Update labels
-        if dim in tensors[0].stored_labels.keys():
-            new_tensor_shape = new_tensor.shape
-            labels = create_labels_cat(tensors, dim, new_tensor_shape)
+        if not tensors:
+            return []  # Handle empty list
+        if len(tensors) == 1:
+            return tensors[0]  # Return single tensor as-is
+
+            # Perform concatenation
+        cat_tensor = torch.cat(tensors, dim=dim)
+        tensors_labels = [tensor.stored_labels for tensor in tensors]
+
+        # Check label consistency across tensors, excluding the
+        # concatenation dimension
+        for key in tensors_labels[0].keys():
+            if key != dim:
+                if any(tensors_labels[i][key] != tensors_labels[0][key]
+                       for i in range(len(tensors_labels))):
+                    raise RuntimeError(
+                        f"Tensors must have the same labels along all "
+                        f"dimensions except {dim}.")
+
+        # Copy and update the 'dof' for the concatenation dimension
+        cat_labels = {k: copy(v) for k, v in tensors_labels[0].items()}
+
+        # Update labels if the concatenation dimension has labels
+        if dim in tensors[0].stored_labels:
+            if dim in cat_labels:
+                cat_dofs = [label[dim]['dof'] for label in
+                            tensors_labels]
+                cat_labels[dim]['dof'] = sum(cat_dofs, [])
         else:
-            labels = tensors[0].stored_labels
-        new_tensor._labels = labels
-        return new_tensor
+            cat_labels = tensors[0].stored_labels
+
+        # Assign updated labels to the concatenated tensor
+        cat_tensor._labels = cat_labels
+        return cat_tensor
 
     @staticmethod
     def stack(tensors):
+        """
+        Stacks a list of tensors along a new dimension.
+
+        :param tensors: A list of tensors to stack. All tensors must have the
+                        same shape.
+        :type tensors: list of LabelTensor
+        :return: A new tensor obtained by stacking the input tensors,
+                 with the updated labels.
+        :rtype: LabelTensor
+        """
+        # Perform stacking in torch
         new_tensor = torch.stack(tensors)
+
+        # Increase labels keys by 1
         labels = tensors[0]._labels
         labels = {key + 1: value for key, value in labels.items()}
-        if full_labels:
-            new_tensor.labels = labels
-        else:
-            new_tensor._labels = labels
+        new_tensor._labels = labels
         return new_tensor
 
     def requires_grad_(self, mode=True):
+        """
+        Override the requires_grad_ method to update the labels in the new
+        tensor.
+
+        :param mode: A boolean value indicating whether the tensor should track
+                     gradients.If `True`, the tensor will track gradients; if `
+                     False`, it will not.
+        :type mode: bool, optional (default is `True`)
+        :return: The tensor itself with the updated `requires_grad` state and
+                 retained labels.
+        :rtype: LabelTensor
+        """
         lt = super().requires_grad_(mode)
         lt._labels = self._labels
         return lt
@@ -324,39 +330,10 @@ def clone(self, *args, **kwargs):
         :return: A copy of the tensor.
         :rtype: LabelTensor
         """
-        out = LabelTensor(super().clone(*args, **kwargs), deepcopy(self._labels))
+        out = LabelTensor(super().clone(*args, **kwargs),
+                          deepcopy(self._labels))
         return out
 
-    @staticmethod
-    def summation(tensors):
-        if len(tensors) == 0:
-            raise ValueError('tensors list must not be empty')
-        if len(tensors) == 1:
-            return tensors[0]
-        # Collect all labels
-
-        # Check labels of all the tensors in each dimension
-        if not all(tensor.shape == tensors[0].shape for tensor in tensors) or \
-                not all(tensor.full_labels[i] == tensors[0].full_labels[i] for
-                        tensor in tensors for i in range(tensors[0].ndim - 1)):
-            raise RuntimeError('Tensors must have the same shape and labels')
-
-        last_dim_labels = []
-        data = torch.zeros(tensors[0].tensor.shape).to(tensors[0].device)
-        for tensor in tensors:
-            data += tensor.tensor
-            last_dim_labels.append(tensor.labels)
-
-        last_dim_labels = ['+'.join(items) for items in zip(*last_dim_labels)]
-        labels = {k: copy(v) for k, v in tensors[0].stored_labels.items()}
-        labels.update({
-            tensors[0].ndim - 1: {
-                'dof': last_dim_labels,
-                'name': tensors[0].name
-            }
-        })
-        return LabelTensor(data, labels)
-
     def append(self, tensor, mode='std'):
         if mode == 'std':
             # Call cat on last dimension
@@ -390,132 +367,210 @@ def vstack(label_tensors):
         """
         return LabelTensor.cat(label_tensors, dim=0)
 
-        # ---------------------- Start auxiliary function definition -----
-        # This method is used to update labels
+    # This method is used to update labels
     def _update_single_label(self, old_labels, to_update_labels, index, dim,
                              to_update_dim):
-            """
-            TODO
-                :param old_labels: labels from which retrieve data
-                :param to_update_labels: labels to update
-                :param index: index of dof to retain
-                :param dim: label index
-                :return:
-            """
-            old_dof = old_labels[to_update_dim]['dof']
-            if isinstance(index, slice):
-                to_update_labels.update({
-                    dim: {
-                        'dof': old_dof[index],
-                        'name': old_labels[dim]['name']
-                    }
-                })
-                return
-            if isinstance(index, int):
-                index = [index]
-            if isinstance(index, (list, torch.Tensor)):
-                to_update_labels.update({
-                    dim: {
-                        'dof': [old_dof[i] for i in index] if isinstance(old_dof, list) else index,
-                        'name': old_labels[dim]['name']
-                    }
-                })
-                return
-            raise NotImplementedError(f'Getitem not implemented for '
-                                      f'{type(index)} values')
-        # ---------------------- End auxiliary function definition -----
-
-
-    def __getitem__(self, index):
         """
-        TODO: Complete docstring
-        :param index:
+        Update the labels of the tensor by selecting only the labels
+        :param old_labels: labels from which retrieve data
+        :param to_update_labels: labels to update
+        :param index: index of dof to retain
+        :param dim: label index
         :return:
         """
-        # Index are str --> call extract
-        if isinstance(index, str) or (isinstance(index, (tuple, list))
-                                      and all(
-                    isinstance(a, str) for a in index)):
+        old_dof = old_labels[to_update_dim]['dof']
+        label_name = old_labels[dim]['name']
+
+        if isinstance(index, slice):
+            # Handle slicing
+            to_update_labels[dim] = {'dof': old_dof[index], 'name': label_name}
+        elif isinstance(index, int):
+            # Handle single integer index
+            to_update_labels[dim] = {'dof': [old_dof[index]],
+                                     'name': label_name}
+        elif isinstance(index, (list, torch.Tensor)):
+            # Handle lists or tensors
+            indices = [index] if isinstance(index, (int, str)) else index
+            to_update_labels[dim] = {
+                'dof': [old_dof[i] for i in indices] if isinstance(old_dof,
+                                                                   list) else indices,
+                'name': label_name
+            }
+        else:
+            raise NotImplementedError(
+                f"Unsupported index type: {type(index)}. Expected slice, int, "
+                f"list, or torch.Tensor."
+            )
+
+    def __getitem__(self, index):
+        """"
+        Override the __getitem__ method to handle the labels of the tensor.
+        Perform the __getitem__ operation on the tensor and update the labels.
+
+        :param index: The index used to access the item
+        :type index: Union[int, str, tuple, list]
+        :return: A tensor-like object with updated labels.
+        :rtype: LabelTensor
+        :raises KeyError: If an invalid label index is provided.
+        :raises IndexError: If an invalid index is accessed in the tensor.
+        """
+        # Handle string index
+        if isinstance(index, str) or (isinstance(index, (tuple, list)) and all(
+                isinstance(i, str) for i in index)):
             return self.extract(index)
 
-        # Store important variables
-        selected_lt = super().__getitem__(index)
-        stored_labels = self._labels
-        labels = copy(stored_labels)
-
-        # Put here because it is the most common case (int as index).
-        # Used by DataLoader -> put here for efficiency purpose
-        if isinstance(index, list):
-            if 0 in labels.keys():
-                self._update_single_label(stored_labels, labels, index,
-                                          0, 0)
-            selected_lt._labels = labels
-            return selected_lt
-
-        if isinstance(index, int):
-            labels.pop(0, None)
-            labels = {key - 1 if key > 0 else key: value for key, value in
-                      labels.items()}
-            selected_lt._labels = labels
-            return selected_lt
-
-        if not isinstance(index, (tuple, torch.Tensor)):
+        # Retrieve selected tensor and labels
+        selected_tensor = super().__getitem__(index)
+        original_labels = self._labels
+        updated_labels = copy(original_labels)
+
+        # Ensure the index is iterable
+        if not isinstance(index, tuple):
             index = [index]
 
-        # Ellipsis are used to perform operation on the last dimension
-        if index[0] == Ellipsis:
-            if len(self.shape) in labels:
-                self._update_single_label(stored_labels, labels, index, 0, 0)
-            selected_lt._labels = labels
-            return selected_lt
-
-        i = 0
-        for j, idx in enumerate(index):
-            if j in self.stored_labels.keys():
-                if isinstance(idx, int) or (
-                        isinstance(idx, torch.Tensor) and idx.ndim == 0):
-                    selected_lt = selected_lt.unsqueeze(j)
+        # Update labels based on the index
+        offset = 0
+        for dim, idx in enumerate(index):
+            if dim in self.stored_labels.keys():
+                if isinstance(idx, int):
+                    selected_tensor = selected_tensor.unsqueeze(dim)
                 if idx != slice(None):
-                    self._update_single_label(stored_labels, labels, idx, j, i)
+                    self._update_single_label(original_labels, updated_labels,
+                                              idx, dim, offset)
             else:
+                # Adjust label keys if dimension is reduced (case of integer
+                # index on a non-labeled dimension)
                 if isinstance(idx, int):
-                    labels = {key - 1 if key > j else key:
-                                  value for key, value in labels.items()}
+                    updated_labels = {
+                        key - 1 if key > dim else key: value
+                        for key, value in updated_labels.items()
+                    }
                     continue
-            i += 1
-        selected_lt._labels = labels
-        return selected_lt
+            offset += 1
+
+        # Update the selected tensor's labels
+        selected_tensor._labels = updated_labels
+        return selected_tensor
 
     def sort_labels(self, dim=None):
+        """
+        Sorts the labels along a specified dimension and returns a new tensor
+        with sorted labels.
+
+        :param dim: The dimension along which to sort the labels. If `None`,
+                    the last dimension (`ndim - 1`) is used.
+        :type dim: int, optional
+        :return: A new tensor with sorted labels along the specified dimension.
+        :rtype: LabelTensor
+        """
+
         def arg_sort(lst):
             return sorted(range(len(lst)), key=lambda x: lst[x])
+
         if dim is None:
             dim = self.ndim - 1
         if self.shape[dim] == 1:
             return self
         labels = self.stored_labels[dim]['dof']
         sorted_index = arg_sort(labels)
+        # Define an indexer to sort the tensor along the specified dimension
         indexer = [slice(None)] * self.ndim
+        # Assigned the sorted index to the specified dimension
         indexer[dim] = sorted_index
         return self.__getitem__(tuple(indexer))
 
     def __deepcopy__(self, memo):
+        """
+        Creates a deep copy of the object.
+
+        :param memo: LabelTensor object to be copied.
+        :type memo: LabelTensor
+        :return: A deep copy of the original LabelTensor object.
+        :rtype: LabelTensor
+        """
         cls = self.__class__
         result = cls(deepcopy(self.tensor), deepcopy(self.stored_labels))
         return result
 
     def permute(self, *dims):
+        """
+        Permutes the dimensions of the tensor and the associated labels
+        accordingly.
+
+        :param dims: The dimensions to permute the tensor to.
+        :type dims: tuple, list
+        :return: A new object with permuted dimensions and reordered labels.
+        :rtype: LabelTensor
+        """
+        # Call the base class permute method
         tensor = super().permute(*dims)
+
+        # Update lables
         labels = self._labels
         keys_list = list(*dims)
         labels = {
             keys_list.index(k): labels[k]
             for k in labels.keys()
         }
+
+        # Assign labels to the new tensor
         tensor._labels = labels
         return tensor
 
     def detach(self):
+        """
+        Detaches the tensor from the computation graph and retains the stored
+        labels.
+
+        :return: A new tensor detached from the computation graph.
+        :rtype: LabelTensor
+        """
         lt = super().detach()
-        lt._labels = self.stored_labels
-        return lt
\ No newline at end of file
+
+        # Copy the labels to the new tensor only if present
+        if hasattr(self, "_labels"):
+            lt._labels = self.stored_labels
+        return lt
+
+    @staticmethod
+    def summation(tensors):
+        """
+        Computes the summation of a list of tensors.
+
+        :param tensors: A list of tensors to sum. All tensors must have the same
+                        shape and labels.
+        :type tensors: list of LabelTensor
+        :return: A new `LabelTensor` containing the element-wise sum of the
+                 input tensors.
+        :rtype: LabelTensor
+        :raises ValueError: If the input `tensors` list is empty.
+        :raises RuntimeError: If the tensors have different shapes and/or
+                              mismatched labels.
+        """
+
+        if not tensors:
+            raise ValueError('The tensors list must not be empty.')
+
+        if len(tensors) == 1:
+            return tensors[0]
+
+        # Initialize result tensor and labels
+        data = torch.zeros_like(tensors[0].tensor).to(tensors[0].device)
+        last_dim_labels = []
+
+        # Accumulate tensors
+        for tensor in tensors:
+            data += tensor.tensor
+            last_dim_labels.append(tensor.labels)
+
+        # Construct last dimension labels
+        last_dim_labels = ['+'.join(items) for items in zip(*last_dim_labels)]
+
+        # Update the labels for the resulting tensor
+        labels = {k: copy(v) for k, v in tensors[0].stored_labels.items()}
+        labels[tensors[0].ndim - 1] = {
+            'dof': last_dim_labels,
+            'name': tensors[0].name
+        }
+
+        return LabelTensor(data, labels)
diff --git a/tests/test_label_tensor/test_label_tensor.py b/tests/test_label_tensor/test_label_tensor.py
index 61e479951..41288e693 100644
--- a/tests/test_label_tensor/test_label_tensor.py
+++ b/tests/test_label_tensor/test_label_tensor.py
@@ -7,7 +7,8 @@
 labels_column = {1: {"name": "space", "dof": ['x', 'y', 'z']}}
 labels_row = {0: {"name": "samples", "dof": range(20)}}
 labels_list = ['x', 'y', 'z']
-labels_all = labels_column | labels_row
+labels_all = labels_column.copy()
+labels_all.update(labels_row)
 
 
 @pytest.mark.parametrize("labels",

From b881765c220cf034049141ff64ebea26554d1974 Mon Sep 17 00:00:00 2001
From: Filippo Olivo <filippo@filippoolivo.com>
Date: Thu, 23 Jan 2025 12:04:31 +0100
Subject: [PATCH 46/55] Improve DataLoader performance when batch_size=None
 (#423)

---
 pina/data/data_module.py | 67 +++++++++++++++++++++++++---------------
 pina/data/dataset.py     |  4 +++
 2 files changed, 46 insertions(+), 25 deletions(-)

diff --git a/pina/data/data_module.py b/pina/data/data_module.py
index 4831e20e7..4b529fe2c 100644
--- a/pina/data/data_module.py
+++ b/pina/data/data_module.py
@@ -8,6 +8,19 @@
 from torch.utils.data.distributed import DistributedSampler
 from .dataset import PinaDatasetFactory
 
+class DummyDataloader:
+    def __init__(self, dataset, device):
+        self.dataset = dataset.get_all_data()
+
+    def __iter__(self):
+        return self
+
+    def __len__(self):
+        return 1
+
+    def __next__(self):
+        return self.dataset
+
 class Collator:
     def __init__(self, max_conditions_lengths, ):
         self.max_conditions_lengths = max_conditions_lengths
@@ -232,40 +245,41 @@ def val_dataloader(self):
         """
         Create the validation dataloader
         """
-
-        batch_size = self.batch_size if self.batch_size is not None else len(
-            self.val_dataset)
-
-        # Use default batching in torch DataLoader (good is batch size is small)
-        if self.automatic_batching:
-            collate = Collator(self.find_max_conditions_lengths('val'))
-            return DataLoader(self.val_dataset, batch_size,
-                              collate_fn=collate)
-        collate = Collator(None)
         # Use custom batching (good if batch size is large)
-        sampler = PinaBatchSampler(self.val_dataset, batch_size, shuffle=False)
-        return DataLoader(self.val_dataset, sampler=sampler,
+        if self.batch_size is not None:
+            # Use default batching in torch DataLoader (good is batch size is small)
+            if self.automatic_batching:
+                collate = Collator(self.find_max_conditions_lengths('val'))
+                return DataLoader(self.val_dataset, self.batch_size,
+                          collate_fn=collate)
+            collate = Collator(None)
+            sampler = PinaBatchSampler(self.val_dataset, self.batch_size, shuffle=False)
+            return DataLoader(self.val_dataset, sampler=sampler,
                           collate_fn=collate)
+        dataloader = DummyDataloader(self.train_dataset, self.trainer.strategy.root_device)
+        dataloader.dataset = self.transfer_batch_to_device(dataloader.dataset, self.trainer.strategy.root_device, 0)
+        self.transfer_batch_to_device = self.dummy_transfer_to_device
 
     def train_dataloader(self):
         """
         Create the training dataloader
         """
-        # Use default batching in torch DataLoader (good is batch size is small)
-        batch_size = self.batch_size if self.batch_size is not None else len(
-            self.train_dataset)
-
-        if self.automatic_batching:
-            collate = Collator(self.find_max_conditions_lengths('train'))
-            return DataLoader(self.train_dataset, batch_size,
-                                collate_fn=collate)
-        collate = Collator(None)
         # Use custom batching (good if batch size is large)
-
-        sampler = PinaBatchSampler(self.train_dataset, batch_size,
-                                       shuffle=False)
-        return DataLoader(self.train_dataset, sampler=sampler,
+        if self.batch_size is not None:
+            # Use default batching in torch DataLoader (good is batch size is small)
+            if self.automatic_batching:
+                collate = Collator(self.find_max_conditions_lengths('train'))
+                return DataLoader(self.train_dataset, self.batch_size,
+                                  collate_fn=collate)
+            collate = Collator(None)
+            sampler = PinaBatchSampler(self.train_dataset, self.batch_size,
+                                           shuffle=False)
+            return DataLoader(self.train_dataset, sampler=sampler,
                               collate_fn=collate)
+        dataloader = DummyDataloader(self.train_dataset, self.trainer.strategy.root_device)
+        dataloader.dataset = self.transfer_batch_to_device(dataloader.dataset, self.trainer.strategy.root_device, 0)
+        self.transfer_batch_to_device = self.dummy_transfer_to_device
+        return dataloader
 
     def test_dataloader(self):
         """
@@ -279,6 +293,9 @@ def predict_dataloader(self):
         """
         raise NotImplementedError("Predict dataloader not implemented")
 
+    def dummy_transfer_to_device(self, batch, device, dataloader_idx):
+        return batch
+
     def transfer_batch_to_device(self, batch, device, dataloader_idx):
         """
         Transfer the batch to the device. This method is called in the
diff --git a/pina/data/dataset.py b/pina/data/dataset.py
index e5685f1fb..8f41c0bd7 100644
--- a/pina/data/dataset.py
+++ b/pina/data/dataset.py
@@ -75,6 +75,10 @@ def _getitem_list(self, idx):
                                          for k, v in data.items()}
         return to_return_dict
 
+    def get_all_data(self):
+        index = [i for i in range(len(self))]
+        return self._getitem_list(index)
+
     def __getitem__(self, idx):
         return self._getitem_func(idx)
 

From b498797bfef45414d1a50147e3f1097b7179e5a8 Mon Sep 17 00:00:00 2001
From: FilippoOlivo <filippo@filippoolivo.com>
Date: Thu, 23 Jan 2025 15:04:15 +0100
Subject: [PATCH 47/55] Bug fix PR #423

---
 pina/data/data_module.py | 11 +++++------
 1 file changed, 5 insertions(+), 6 deletions(-)

diff --git a/pina/data/data_module.py b/pina/data/data_module.py
index 4b529fe2c..da0326648 100644
--- a/pina/data/data_module.py
+++ b/pina/data/data_module.py
@@ -256,9 +256,9 @@ def val_dataloader(self):
             sampler = PinaBatchSampler(self.val_dataset, self.batch_size, shuffle=False)
             return DataLoader(self.val_dataset, sampler=sampler,
                           collate_fn=collate)
-        dataloader = DummyDataloader(self.train_dataset, self.trainer.strategy.root_device)
+        dataloader = DummyDataloader(self.val_dataset, self.trainer.strategy.root_device)
         dataloader.dataset = self.transfer_batch_to_device(dataloader.dataset, self.trainer.strategy.root_device, 0)
-        self.transfer_batch_to_device = self.dummy_transfer_to_device
+        return dataloader
 
     def train_dataloader(self):
         """
@@ -278,7 +278,6 @@ def train_dataloader(self):
                               collate_fn=collate)
         dataloader = DummyDataloader(self.train_dataset, self.trainer.strategy.root_device)
         dataloader.dataset = self.transfer_batch_to_device(dataloader.dataset, self.trainer.strategy.root_device, 0)
-        self.transfer_batch_to_device = self.dummy_transfer_to_device
         return dataloader
 
     def test_dataloader(self):
@@ -293,18 +292,18 @@ def predict_dataloader(self):
         """
         raise NotImplementedError("Predict dataloader not implemented")
 
-    def dummy_transfer_to_device(self, batch, device, dataloader_idx):
-        return batch
-
     def transfer_batch_to_device(self, batch, device, dataloader_idx):
         """
         Transfer the batch to the device. This method is called in the
         training loop and is used to transfer the batch to the device.
         """
+        if isinstance(batch, list):
+            return batch
         batch = [
             (k, super(LightningDataModule, self).transfer_batch_to_device(v,
                                                                 device,
                                                                 dataloader_idx))
             for k, v in batch.items()
         ]
+
         return batch

From a3bb77d8e3c87955443f3e3652af8f2e4baa640d Mon Sep 17 00:00:00 2001
From: Dario Coscia <93731561+dario-coscia@users.noreply.github.com>
Date: Thu, 23 Jan 2025 15:13:14 +0100
Subject: [PATCH 48/55] Revert "Bug fix PR #423" (#426)

This reverts commit b498797bfef45414d1a50147e3f1097b7179e5a8.

Co-authored-by: Filippo Olivo <filippo@filippoolivo.com>
Co-authored-by: Dario Coscia <dariocos99@gmail.com>
---
 pina/data/data_module.py | 11 ++++++-----
 1 file changed, 6 insertions(+), 5 deletions(-)

diff --git a/pina/data/data_module.py b/pina/data/data_module.py
index da0326648..4b529fe2c 100644
--- a/pina/data/data_module.py
+++ b/pina/data/data_module.py
@@ -256,9 +256,9 @@ def val_dataloader(self):
             sampler = PinaBatchSampler(self.val_dataset, self.batch_size, shuffle=False)
             return DataLoader(self.val_dataset, sampler=sampler,
                           collate_fn=collate)
-        dataloader = DummyDataloader(self.val_dataset, self.trainer.strategy.root_device)
+        dataloader = DummyDataloader(self.train_dataset, self.trainer.strategy.root_device)
         dataloader.dataset = self.transfer_batch_to_device(dataloader.dataset, self.trainer.strategy.root_device, 0)
-        return dataloader
+        self.transfer_batch_to_device = self.dummy_transfer_to_device
 
     def train_dataloader(self):
         """
@@ -278,6 +278,7 @@ def train_dataloader(self):
                               collate_fn=collate)
         dataloader = DummyDataloader(self.train_dataset, self.trainer.strategy.root_device)
         dataloader.dataset = self.transfer_batch_to_device(dataloader.dataset, self.trainer.strategy.root_device, 0)
+        self.transfer_batch_to_device = self.dummy_transfer_to_device
         return dataloader
 
     def test_dataloader(self):
@@ -292,18 +293,18 @@ def predict_dataloader(self):
         """
         raise NotImplementedError("Predict dataloader not implemented")
 
+    def dummy_transfer_to_device(self, batch, device, dataloader_idx):
+        return batch
+
     def transfer_batch_to_device(self, batch, device, dataloader_idx):
         """
         Transfer the batch to the device. This method is called in the
         training loop and is used to transfer the batch to the device.
         """
-        if isinstance(batch, list):
-            return batch
         batch = [
             (k, super(LightningDataModule, self).transfer_batch_to_device(v,
                                                                 device,
                                                                 dataloader_idx))
             for k, v in batch.items()
         ]
-
         return batch

From 733e55946d51d81957f2b549384a343fcf54fce9 Mon Sep 17 00:00:00 2001
From: Nicola Demo <demo.nicola@gmail.com>
Date: Thu, 23 Jan 2025 09:52:23 +0100
Subject: [PATCH 49/55] fix tests

---
 .../callbacks/adaptive_refinment_callbacks.py |   8 +-
 pina/callbacks/optimizer_callbacks.py         |  60 +-
 pina/callbacks/processing_callbacks.py        |   7 +-
 pina/loss/__init__.py                         |   4 +-
 pina/optim/torch_optimizer.py                 |   3 +
 pina/solvers/pinns/__init__.py                |   2 +-
 pina/solvers/pinns/competitive_pinn.py        |   2 +-
 pina/solvers/pinns/pinn.py                    |   6 +-
 .../pinns/{basepinn.py => pinn_interface.py}  |  15 +-
 pina/solvers/pinns/sapinn.py                  |   2 +-
 pina/solvers/solver.py                        |   4 +-
 .../test_adaptive_refinment_callbacks.py      |  49 +-
 tests/test_callbacks/test_metric_tracker.py   |   6 +-
 .../test_optimizer_callbacks.py               |  22 +-
 tests/test_callbacks/test_progress_bar.py     |  43 +-
 tests/test_condition.py                       |   4 -
 tests/test_loss/test_lploss.py                |   3 +-
 tests/test_loss/test_powerloss.py             |   2 +-
 tests/test_operators.py                       |   3 +-
 tests/test_plotter.py                         |  10 +-
 .../test_poisson_2d_square.py                 |   3 -
 tests/test_solvers/test_basepinn.py           | 208 +++--
 tests/test_solvers/test_causalpinn.py         | 484 +++++------
 tests/test_solvers/test_competitive_pinn.py   | 738 ++++++++---------
 tests/test_solvers/test_garom.py              | 314 +++----
 tests/test_solvers/test_gpinn.py              | 758 ++++++++---------
 tests/test_solvers/test_pinn.py               | 425 ++++------
 tests/test_solvers/test_rba_pinn.py           | 762 ++++++++---------
 tests/test_solvers/test_rom_solver.py         | 170 ++--
 tests/test_solvers/test_sapinn.py             | 766 +++++++++---------
 tests/test_solvers/test_supervised_solver.py  | 252 +++---
 tests/test_utils.py                           |  10 +-
 32 files changed, 2524 insertions(+), 2621 deletions(-)
 rename pina/solvers/pinns/{basepinn.py => pinn_interface.py} (97%)

diff --git a/pina/callbacks/adaptive_refinment_callbacks.py b/pina/callbacks/adaptive_refinment_callbacks.py
index 00d6dfb80..3462fd984 100644
--- a/pina/callbacks/adaptive_refinment_callbacks.py
+++ b/pina/callbacks/adaptive_refinment_callbacks.py
@@ -49,7 +49,7 @@ def _compute_residual(self, trainer):
         """
 
         # extract the solver and device from trainer
-        solver = trainer._model
+        solver = trainer.solver
         device = trainer._accelerator_connector._accelerator_flag
         precision = trainer.precision
         if precision == "64-true":
@@ -67,7 +67,7 @@ def _compute_residual(self, trainer):
         # compute residual
         res_loss = {}
         tot_loss = []
-        for location in self._sampling_locations:
+        for location in self._sampling_locations: #TODO fix for new collector
             condition = solver.problem.conditions[location]
             pts = solver.problem.input_pts[location]
             # send points to correct device
@@ -79,6 +79,8 @@ def _compute_residual(self, trainer):
             res_loss[location] = torch.abs(target).as_subclass(torch.Tensor)
             tot_loss.append(torch.abs(target))
 
+        print(tot_loss)
+
         return torch.vstack(tot_loss), res_loss
 
     def _r3_routine(self, trainer):
@@ -139,7 +141,7 @@ def on_train_start(self, trainer, _):
         :rtype: None
         """
         # extract locations for sampling
-        problem = trainer._model.problem
+        problem = trainer.solver.problem
         locations = []
         for condition_name in problem.conditions:
             condition = problem.conditions[condition_name]
diff --git a/pina/callbacks/optimizer_callbacks.py b/pina/callbacks/optimizer_callbacks.py
index 6905ebf22..c4836dcb0 100644
--- a/pina/callbacks/optimizer_callbacks.py
+++ b/pina/callbacks/optimizer_callbacks.py
@@ -3,61 +3,45 @@
 from lightning.pytorch.callbacks import Callback
 import torch
 from ..utils import check_consistency
+from pina.optim import TorchOptimizer
 
 
 class SwitchOptimizer(Callback):
 
-    def __init__(self, new_optimizers, new_optimizers_kwargs, epoch_switch):
+    def __init__(self, new_optimizers, epoch_switch):
         """
-        PINA Implementation of a Lightning Callback to switch optimizer during training.
+        PINA Implementation of a Lightning Callback to switch optimizer during
+        training.
 
-        This callback allows for switching between different optimizers during training, enabling
-        the exploration of multiple optimization strategies without the need to stop training.
+        This callback allows for switching between different optimizers during
+        training, enabling the exploration of multiple optimization strategies
+        without the need to stop training.
 
-        :param new_optimizers: The model optimizers to switch to. Can be a single
-                            :class:`torch.optim.Optimizer` or a list of them for multiple model solvers.
-        :type new_optimizers: torch.optim.Optimizer | list
-        :param new_optimizers_kwargs: The keyword arguments for the new optimizers. Can be a single dictionary
-                                    or a list of dictionaries corresponding to each optimizer.
-        :type new_optimizers_kwargs: dict | list
+        :param new_optimizers: The model optimizers to switch to. Can be a
+            single :class:`torch.optim.Optimizer` or a list of them for multiple
+            model solvers.
+        :type new_optimizers: pina.optim.TorchOptimizer | list
         :param epoch_switch: The epoch at which to switch to the new optimizer.
         :type epoch_switch: int
 
-        :raises ValueError: If `epoch_switch` is less than 1 or if there is a mismatch in the number of
-                            optimizers and their corresponding keyword argument dictionaries.
-
         Example:
-            >>> switch_callback = SwitchOptimizer(new_optimizers=[optimizer1, optimizer2],
-            >>>                                  new_optimizers_kwargs=[{'lr': 0.001}, {'lr': 0.01}],
+            >>> switch_callback = SwitchOptimizer(new_optimizers=optimizer,
             >>>                                  epoch_switch=10)
         """
         super().__init__()
 
-        # check type consistency
-        check_consistency(new_optimizers, torch.optim.Optimizer, subclass=True)
-        check_consistency(new_optimizers_kwargs, dict)
-        check_consistency(epoch_switch, int)
-
         if epoch_switch < 1:
             raise ValueError("epoch_switch must be greater than one.")
 
         if not isinstance(new_optimizers, list):
             new_optimizers = [new_optimizers]
-            new_optimizers_kwargs = [new_optimizers_kwargs]
-        len_optimizer = len(new_optimizers)
-        len_optimizer_kwargs = len(new_optimizers_kwargs)
-
-        if len_optimizer_kwargs != len_optimizer:
-            raise ValueError(
-                "You must define one dictionary of keyword"
-                " arguments for each optimizers."
-                f" Got {len_optimizer} optimizers, and"
-                f" {len_optimizer_kwargs} dicitionaries"
-            )
 
+        # check type consistency
+        for optimizer in new_optimizers:
+            check_consistency(optimizer, TorchOptimizer)
+        check_consistency(epoch_switch, int)
         # save new optimizers
         self._new_optimizers = new_optimizers
-        self._new_optimizers_kwargs = new_optimizers_kwargs
         self._epoch_switch = epoch_switch
 
     def on_train_epoch_start(self, trainer, __):
@@ -73,13 +57,9 @@ def on_train_epoch_start(self, trainer, __):
         """
         if trainer.current_epoch == self._epoch_switch:
             optims = []
-            for idx, (optim, optim_kwargs) in enumerate(
-                zip(self._new_optimizers, self._new_optimizers_kwargs)
-            ):
-                optims.append(
-                    optim(
-                        trainer._model.models[idx].parameters(), **optim_kwargs
-                    )
-                )
+
+            for idx, optim in enumerate(self._new_optimizers):
+                optim.hook(trainer.solver.models[idx].parameters())
+                optims.append(optim.instance)
 
             trainer.optimizers = optims
diff --git a/pina/callbacks/processing_callbacks.py b/pina/callbacks/processing_callbacks.py
index e65506c17..d241af07d 100644
--- a/pina/callbacks/processing_callbacks.py
+++ b/pina/callbacks/processing_callbacks.py
@@ -14,7 +14,7 @@
 
 class MetricTracker(Callback):
 
-    def __init__(self):
+    def __init__(self, metrics_to_track=None):
         """
         PINA Implementation of a Lightning Callback for Metric Tracking.
 
@@ -37,6 +37,9 @@ def __init__(self):
         """
         super().__init__()
         self._collection = []
+        if metrics_to_track is not None:
+            metrics_to_track = ['train_loss_epoch', 'train_loss_step', 'val_loss']
+        self.metrics_to_track = metrics_to_track
 
     def on_train_epoch_end(self, trainer, pl_module):
         """
@@ -72,7 +75,7 @@ class PINAProgressBar(TQDMProgressBar):
 
     BAR_FORMAT = "{l_bar}{bar}| {n_fmt}/{total_fmt} [{elapsed}<{remaining}, {rate_noinv_fmt}{postfix}]"
 
-    def __init__(self, metrics="mean", **kwargs):
+    def __init__(self, metrics="val_loss", **kwargs):
         """
         PINA Implementation of a Lightning Callback for enriching the progress
         bar.
diff --git a/pina/loss/__init__.py b/pina/loss/__init__.py
index 35138125d..4415906f6 100644
--- a/pina/loss/__init__.py
+++ b/pina/loss/__init__.py
@@ -1,6 +1,8 @@
 __all__ = [
     'LpLoss',
-
+    'PowerLoss',
+    'weightningInterface',
+    'LossInterface'
 ]
 
 from .loss_interface import LossInterface
diff --git a/pina/optim/torch_optimizer.py b/pina/optim/torch_optimizer.py
index 54818d5a5..fe096565c 100644
--- a/pina/optim/torch_optimizer.py
+++ b/pina/optim/torch_optimizer.py
@@ -18,3 +18,6 @@ def __init__(self, optimizer_class, **kwargs):
     def hook(self, parameters):
         self.optimizer_instance = self.optimizer_class(parameters,
                                                        **self.kwargs)
+    @property
+    def instance(self):
+        return self.optimizer_instance
\ No newline at end of file
diff --git a/pina/solvers/pinns/__init__.py b/pina/solvers/pinns/__init__.py
index 8c779665a..dbd400d99 100644
--- a/pina/solvers/pinns/__init__.py
+++ b/pina/solvers/pinns/__init__.py
@@ -8,7 +8,7 @@
     "RBAPINN",
 ]
 
-from .basepinn import PINNInterface
+from .pinn_interface import PINNInterface
 from .pinn import PINN
 from .gpinn import GPINN
 from .causalpinn import CausalPINN
diff --git a/pina/solvers/pinns/competitive_pinn.py b/pina/solvers/pinns/competitive_pinn.py
index 5e011a473..7e7302ec7 100644
--- a/pina/solvers/pinns/competitive_pinn.py
+++ b/pina/solvers/pinns/competitive_pinn.py
@@ -12,7 +12,7 @@
 
 from torch.optim.lr_scheduler import ConstantLR
 
-from .basepinn import PINNInterface
+from .pinn_interface import PINNInterface
 from pina.utils import check_consistency
 from pina.problem import InverseProblem
 
diff --git a/pina/solvers/pinns/pinn.py b/pina/solvers/pinns/pinn.py
index d1ab21d76..5207f6035 100644
--- a/pina/solvers/pinns/pinn.py
+++ b/pina/solvers/pinns/pinn.py
@@ -10,7 +10,7 @@
     )  # torch < 2.0
 
 
-from .basepinn import PINNInterface
+from .pinn_interface import PINNInterface
 from ...problem import InverseProblem
 
 
@@ -60,7 +60,6 @@ def __init__(
         self,
         problem,
         model,
-        extra_features=None,
         loss=None,
         optimizer=None,
         scheduler=None,
@@ -82,10 +81,9 @@ def __init__(
         super().__init__(
             models=model,
             problem=problem,
+            loss=loss,
             optimizers=optimizer,
             schedulers=scheduler,
-            extra_features=extra_features,
-            loss=loss,
         )
 
         # assign variables
diff --git a/pina/solvers/pinns/basepinn.py b/pina/solvers/pinns/pinn_interface.py
similarity index 97%
rename from pina/solvers/pinns/basepinn.py
rename to pina/solvers/pinns/pinn_interface.py
index 611a3f40e..bde4737d8 100644
--- a/pina/solvers/pinns/basepinn.py
+++ b/pina/solvers/pinns/pinn_interface.py
@@ -3,7 +3,7 @@
 from abc import ABCMeta, abstractmethod
 import torch
 from torch.nn.modules.loss import _Loss
-from ...solvers.solver import SolverInterface
+from ..solver import SolverInterface
 from ...utils import check_consistency
 from ...loss.loss_interface import LossInterface
 from ...problem import InverseProblem
@@ -33,10 +33,9 @@ def __init__(
             self,
             models,
             problem,
-            optimizers,
-            schedulers,
-            extra_features,
-            loss,
+            loss=None,
+            optimizers=None,
+            schedulers=None,
     ):
         """
         :param models: Multiple torch neural network models instances.
@@ -70,7 +69,6 @@ def __init__(
             problem=problem,
             optimizers=optimizers,
             schedulers=schedulers,
-            extra_features=extra_features,
         )
 
         # check consistency
@@ -198,6 +196,11 @@ def loss_phys(self, samples, equation):
         """
         pass
 
+    def configure_optimizers(self):
+        self._optimizer.hook(self._model)
+        self.schedulers.hook(self._optimizer)
+        return [self.optimizers.instance]#, self.schedulers.scheduler_instance
+
     def compute_residual(self, samples, equation):
         """
         Compute the residual for Physics Informed learning. This function
diff --git a/pina/solvers/pinns/sapinn.py b/pina/solvers/pinns/sapinn.py
index 751e21eff..1bf168293 100644
--- a/pina/solvers/pinns/sapinn.py
+++ b/pina/solvers/pinns/sapinn.py
@@ -8,7 +8,7 @@
         _LRScheduler as LRScheduler,
     )  # torch < 2.0
 
-from .basepinn import PINNInterface
+from .pinn_interface import PINNInterface
 from pina.utils import check_consistency
 from pina.problem import InverseProblem
 
diff --git a/pina/solvers/solver.py b/pina/solvers/solver.py
index 408aee5b2..76c084eab 100644
--- a/pina/solvers/solver.py
+++ b/pina/solvers/solver.py
@@ -22,7 +22,6 @@ def __init__(self,
                  problem,
                  optimizers,
                  schedulers,
-                 extra_features,
                  use_lt=True):
         """
         :param model: A torch neural network model instance.
@@ -56,7 +55,6 @@ def __init__(self,
                     model=model,
                     input_variables=problem.input_variables,
                     output_variables=problem.output_variables,
-                    extra_features=extra_features,
                 )
 
         # Check scheduler consistency + encapsulation
@@ -98,7 +96,7 @@ def training_step(self, batch):
 
     @abstractmethod
     def configure_optimizers(self):
-        pass
+        raise NotImplementedError
 
     @property
     def models(self):
diff --git a/tests/test_callbacks/test_adaptive_refinment_callbacks.py b/tests/test_callbacks/test_adaptive_refinment_callbacks.py
index 564aab97c..dbeb25a7b 100644
--- a/tests/test_callbacks/test_adaptive_refinment_callbacks.py
+++ b/tests/test_callbacks/test_adaptive_refinment_callbacks.py
@@ -10,6 +10,7 @@
 boundaries = ['nil_g1', 'nil_g2', 'nil_g3', 'nil_g4']
 n = 10
 poisson_problem.discretise_domain(n, 'grid', locations=boundaries)
+poisson_problem.discretise_domain(n, 'grid', locations='laplace_D')
 model = FeedForward(len(poisson_problem.input_variables),
                     len(poisson_problem.output_variables))
 
@@ -17,27 +18,27 @@
 solver = PINN(problem=poisson_problem, model=model)
 
 
-def test_r3constructor():
-    R3Refinement(sample_every=10)
-
-
-def test_r3refinment_routine():
-    # make the trainer
-    trainer = Trainer(solver=solver,
-                      callbacks=[R3Refinement(sample_every=1)],
-                      accelerator='cpu',
-                      max_epochs=5)
-    trainer.train()
-
-def test_r3refinment_routine():
-    model = FeedForward(len(poisson_problem.input_variables),
-                    len(poisson_problem.output_variables))
-    solver = PINN(problem=poisson_problem, model=model)
-    trainer = Trainer(solver=solver,
-                      callbacks=[R3Refinement(sample_every=1)],
-                      accelerator='cpu',
-                      max_epochs=5)
-    before_n_points = {loc : len(pts) for loc, pts in trainer.solver.problem.input_pts.items()}
-    trainer.train()
-    after_n_points = {loc : len(pts) for loc, pts in trainer.solver.problem.input_pts.items()}
-    assert before_n_points == after_n_points
+# def test_r3constructor():
+#     R3Refinement(sample_every=10)
+
+
+# def test_r3refinment_routine():
+#     # make the trainer
+#     trainer = Trainer(solver=solver,
+#                       callbacks=[R3Refinement(sample_every=1)],
+#                       accelerator='cpu',
+#                       max_epochs=5)
+#     trainer.train()
+
+# def test_r3refinment_routine():
+#     model = FeedForward(len(poisson_problem.input_variables),
+#                     len(poisson_problem.output_variables))
+#     solver = PINN(problem=poisson_problem, model=model)
+#     trainer = Trainer(solver=solver,
+#                       callbacks=[R3Refinement(sample_every=1)],
+#                       accelerator='cpu',
+#                       max_epochs=5)
+#     before_n_points = {loc : len(pts) for loc, pts in trainer.solver.problem.input_pts.items()}
+#     trainer.train()
+#     after_n_points = {loc : len(pts) for loc, pts in trainer.solver.problem.input_pts.items()}
+#     assert before_n_points == after_n_points
diff --git a/tests/test_callbacks/test_metric_tracker.py b/tests/test_callbacks/test_metric_tracker.py
index d67e06b85..ed637a887 100644
--- a/tests/test_callbacks/test_metric_tracker.py
+++ b/tests/test_callbacks/test_metric_tracker.py
@@ -10,6 +10,7 @@
 boundaries = ['nil_g1', 'nil_g2', 'nil_g3', 'nil_g4']
 n = 10
 poisson_problem.discretise_domain(n, 'grid', locations=boundaries)
+poisson_problem.discretise_domain(n, 'grid', locations='laplace_D')
 model = FeedForward(len(poisson_problem.input_variables),
                     len(poisson_problem.output_variables))
 
@@ -33,9 +34,6 @@ def test_metric_tracker_routine():
     metrics = trainer.callbacks[0].metrics
     # assert the logged metrics are correct
     logged_metrics = sorted(list(metrics.keys()))
-    total_metrics = sorted(
-        list([key + '_loss' for key in poisson_problem.conditions.keys()])
-        + ['mean_loss'])
-    assert logged_metrics == total_metrics
+    assert logged_metrics == ['train_loss_epoch', 'train_loss_step', 'val_loss']
 
 
diff --git a/tests/test_callbacks/test_optimizer_callbacks.py b/tests/test_callbacks/test_optimizer_callbacks.py
index a57bda35d..f62078a41 100644
--- a/tests/test_callbacks/test_optimizer_callbacks.py
+++ b/tests/test_callbacks/test_optimizer_callbacks.py
@@ -6,40 +6,32 @@
 from pina.trainer import Trainer
 from pina.model import FeedForward
 from pina.problem.zoo import Poisson2DSquareProblem as Poisson
+from pina.optim import TorchOptimizer
 
 # make the problem
 poisson_problem = Poisson()
 boundaries = ['nil_g1', 'nil_g2', 'nil_g3', 'nil_g4']
 n = 10
 poisson_problem.discretise_domain(n, 'grid', locations=boundaries)
+poisson_problem.discretise_domain(n, 'grid', locations='laplace_D')
 model = FeedForward(len(poisson_problem.input_variables),
                     len(poisson_problem.output_variables))
 
 # make the solver
 solver = PINN(problem=poisson_problem, model=model)
 
+adam_optimizer = TorchOptimizer(torch.optim.Adam, lr=0.01)
+lbfgs_optimizer = TorchOptimizer(torch.optim.LBFGS, lr= 0.001)
 
 def test_switch_optimizer_constructor():
-    SwitchOptimizer(new_optimizers=torch.optim.Adam,
-                    new_optimizers_kwargs={'lr': 0.01},
-                    epoch_switch=10)
-
-    with pytest.raises(ValueError):
-        SwitchOptimizer(new_optimizers=[torch.optim.Adam, torch.optim.Adam],
-                        new_optimizers_kwargs=[{
-                            'lr': 0.01
-                        }],
-                        epoch_switch=10)
+    SwitchOptimizer(adam_optimizer, epoch_switch=10)
 
 
 def test_switch_optimizer_routine():
     # make the trainer
+    switch_opt_callback = SwitchOptimizer(lbfgs_optimizer, epoch_switch=3)
     trainer = Trainer(solver=solver,
-                      callbacks=[
-                          SwitchOptimizer(new_optimizers=torch.optim.LBFGS,
-                                          new_optimizers_kwargs={'lr': 0.01},
-                                          epoch_switch=3)
-                      ],
+                      callbacks=[switch_opt_callback],
                       accelerator='cpu',
                       max_epochs=5)
     trainer.train()
diff --git a/tests/test_callbacks/test_progress_bar.py b/tests/test_callbacks/test_progress_bar.py
index 5c2cae4b1..63df9c0be 100644
--- a/tests/test_callbacks/test_progress_bar.py
+++ b/tests/test_callbacks/test_progress_bar.py
@@ -5,28 +5,29 @@
 from pina.problem.zoo import Poisson2DSquareProblem as Poisson
 
 
-# make the problem
-poisson_problem = Poisson()
-boundaries = ['nil_g1', 'nil_g2', 'nil_g3', 'nil_g4']
-n = 10
-poisson_problem.discretise_domain(n, 'grid', locations=boundaries)
-model = FeedForward(len(poisson_problem.input_variables),
-                    len(poisson_problem.output_variables))
+# # make the problem
+# poisson_problem = Poisson()
+# boundaries = ['nil_g1', 'nil_g2', 'nil_g3', 'nil_g4']
+# n = 10
+# poisson_problem.discretise_domain(n, 'grid', locations=boundaries)
+# poisson_problem.discretise_domain(n, 'grid', locations='laplace_D')
+# model = FeedForward(len(poisson_problem.input_variables),
+#                     len(poisson_problem.output_variables))
 
-# make the solver
-solver = PINN(problem=poisson_problem, model=model)
+# # make the solver
+# solver = PINN(problem=poisson_problem, model=model)
 
 
-def test_progress_bar_constructor():
-    PINAProgressBar(['mean_loss'])
+# def test_progress_bar_constructor():
+#     PINAProgressBar(['mean'])
 
-def test_progress_bar_routine():
-    # make the trainer
-    trainer = Trainer(solver=solver,
-                      callbacks=[
-                          PINAProgressBar(['mean', 'D'])
-                      ],
-                      accelerator='cpu',
-                      max_epochs=5)
-    trainer.train()
-    # TODO there should be a check that the correct metrics are displayed
\ No newline at end of file
+# def test_progress_bar_routine():
+#     # make the trainer
+#     trainer = Trainer(solver=solver,
+#                       callbacks=[
+#                           PINAProgressBar(['mean', 'laplace_D'])
+#                       ],
+#                       accelerator='cpu',
+#                       max_epochs=5)
+#     trainer.train()
+#     # TODO there should be a check that the correct metrics are displayed
\ No newline at end of file
diff --git a/tests/test_condition.py b/tests/test_condition.py
index 9165c3fa1..f5842b978 100644
--- a/tests/test_condition.py
+++ b/tests/test_condition.py
@@ -2,11 +2,7 @@
 import pytest
 
 from pina import LabelTensor, Condition
-from pina.solvers import PINN
 from pina.domain import CartesianDomain
-from pina.problem import SpatialProblem
-from pina.model import FeedForward
-from pina.operators import laplacian
 from pina.equation.equation_factory import FixedValue
 
 example_domain = CartesianDomain({'x': [0, 1], 'y': [0, 1]})
diff --git a/tests/test_loss/test_lploss.py b/tests/test_loss/test_lploss.py
index 3ddd503ee..2f073b049 100644
--- a/tests/test_loss/test_lploss.py
+++ b/tests/test_loss/test_lploss.py
@@ -1,7 +1,6 @@
 import torch
-import pytest
 
-from pina.loss.loss_interface import *
+from pina.loss import LpLoss
 
 input = torch.tensor([[3.], [1.], [-8.]])
 target = torch.tensor([[6.], [4.], [2.]])
diff --git a/tests/test_loss/test_powerloss.py b/tests/test_loss/test_powerloss.py
index fd5fddee8..7ea26755d 100644
--- a/tests/test_loss/test_powerloss.py
+++ b/tests/test_loss/test_powerloss.py
@@ -1,7 +1,7 @@
 import torch
 import pytest
 
-from pina.loss.loss_interface import PowerLoss
+from pina.loss import PowerLoss
 
 input = torch.tensor([[3.], [1.], [-8.]])
 target = torch.tensor([[6.], [4.], [2.]])
diff --git a/tests/test_operators.py b/tests/test_operators.py
index ccfc17b08..44e1c8607 100644
--- a/tests/test_operators.py
+++ b/tests/test_operators.py
@@ -26,6 +26,7 @@ def func_scalar(x):
 def test_grad_scalar_output():
     grad_tensor_s = grad(tensor_s, inp)
     true_val = 2*inp
+    true_val.labels = inp.labels
     assert grad_tensor_s.shape == inp.shape
     assert grad_tensor_s.labels == [
         f'd{tensor_s.labels[0]}d{i}' for i in inp.labels
@@ -37,7 +38,7 @@ def test_grad_scalar_output():
     assert grad_tensor_s.labels == [
         f'd{tensor_s.labels[0]}d{i}' for i in ['x', 'y']
     ]
-    assert torch.allclose(grad_tensor_s, true_val)
+    assert torch.allclose(grad_tensor_s, true_val.extract(['x', 'y']))
 
 
 def test_grad_vector_output():
diff --git a/tests/test_plotter.py b/tests/test_plotter.py
index 98eb08845..838963c10 100644
--- a/tests/test_plotter.py
+++ b/tests/test_plotter.py
@@ -5,7 +5,12 @@
 from pina.problem import SpatialProblem
 from pina.equation import FixedValue
 
-        
+"""
+
+!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
+TODO : Fix the tests once the Plotter class is updated
+!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
+
 class FooProblem1D(SpatialProblem):
 
     # assign output/ spatial and temporal variables
@@ -66,4 +71,5 @@ def test_plot_samples_3d():
     pl = Plotter()
     pl.plot_samples(problem=problem, filename='fig.png')
     import os
-    os.remove('fig.png')
\ No newline at end of file
+    os.remove('fig.png')
+"""
\ No newline at end of file
diff --git a/tests/test_problem_zoo/test_poisson_2d_square.py b/tests/test_problem_zoo/test_poisson_2d_square.py
index 6541a6ed9..6c4221cd3 100644
--- a/tests/test_problem_zoo/test_poisson_2d_square.py
+++ b/tests/test_problem_zoo/test_poisson_2d_square.py
@@ -1,6 +1,3 @@
-import torch
-import pytest
-
 from pina.problem.zoo import Poisson2DSquareProblem
 
 def test_constructor():
diff --git a/tests/test_solvers/test_basepinn.py b/tests/test_solvers/test_basepinn.py
index e7f820d08..82b97b574 100644
--- a/tests/test_solvers/test_basepinn.py
+++ b/tests/test_solvers/test_basepinn.py
@@ -4,110 +4,108 @@
 from pina import Condition, LabelTensor, Trainer
 from pina.problem import SpatialProblem
 from pina.operators import laplacian
-from pina.geometry import CartesianDomain
+from pina.domain import CartesianDomain
 from pina.model import FeedForward
 from pina.solvers import PINNInterface
-from pina.equation.equation import Equation
-from pina.equation.equation_factory import FixedValue
-
-def laplace_equation(input_, output_):
-    force_term = (torch.sin(input_.extract(['x']) * torch.pi) *
-                  torch.sin(input_.extract(['y']) * torch.pi))
-    delta_u = laplacian(output_.extract(['u']), input_)
-    return delta_u - force_term
-
-
-my_laplace = Equation(laplace_equation)
-in_ = LabelTensor(torch.tensor([[0., 1.]]), ['x', 'y'])
-out_ = LabelTensor(torch.tensor([[0.]]), ['u'])
-in2_ = LabelTensor(torch.rand(60, 2), ['x', 'y'])
-out2_ = LabelTensor(torch.rand(60, 1), ['u'])
-
-
-
-class Poisson(SpatialProblem):
-    output_variables = ['u']
-    spatial_domain = CartesianDomain({'x': [0, 1], 'y': [0, 1]})
-
-    conditions = {
-        'gamma1': Condition(
-            location=CartesianDomain({'x': [0, 1], 'y':  1}),
-            equation=FixedValue(0.0)),
-        'gamma2': Condition(
-            location=CartesianDomain({'x': [0, 1], 'y': 0}),
-            equation=FixedValue(0.0)),
-        'gamma3': Condition(
-            location=CartesianDomain({'x':  1, 'y': [0, 1]}),
-            equation=FixedValue(0.0)),
-        'gamma4': Condition(
-            location=CartesianDomain({'x': 0, 'y': [0, 1]}),
-            equation=FixedValue(0.0)),
-        'D': Condition(
-            input_points=LabelTensor(torch.rand(size=(100, 2)), ['x', 'y']),
-            equation=my_laplace),
-        'data': Condition(
-            input_points=in_,
-            output_points=out_),
-        'data2': Condition(
-            input_points=in2_,
-            output_points=out2_)
-    }
-
-    def poisson_sol(self, pts):
-        return -(torch.sin(pts.extract(['x']) * torch.pi) *
-                 torch.sin(pts.extract(['y']) * torch.pi)) / (2 * torch.pi**2)
-
-    truth_solution = poisson_sol
-
-class FOOPINN(PINNInterface):
-    def __init__(self, model, problem):
-        super().__init__(models=[model], problem=problem,
-                         optimizers=[torch.optim.Adam],
-                         optimizers_kwargs=[{'lr' : 0.001}],
-                         extra_features=None,
-                         loss=torch.nn.MSELoss())
-    def forward(self, x):
-        return self.models[0](x)
-
-    def loss_phys(self, samples, equation):
-        residual = self.compute_residual(samples=samples, equation=equation)
-        loss_value = self.loss(
-            torch.zeros_like(residual, requires_grad=True), residual
-        )
-        self.store_log(loss_value=float(loss_value))
-        return loss_value
-
-    def configure_optimizers(self):
-        return self.optimizers, []
-
-# make the problem
-poisson_problem = Poisson()
-poisson_problem.discretise_domain(100)
-model = FeedForward(len(poisson_problem.input_variables),
-                    len(poisson_problem.output_variables))
-model_extra_feats = FeedForward(
-    len(poisson_problem.input_variables) + 1,
-    len(poisson_problem.output_variables))
-
-
-def test_constructor():
-    with pytest.raises(TypeError):
-        PINNInterface()
-    # a simple pinn built with PINNInterface
-    FOOPINN(model, poisson_problem)
-
-def test_train_step():
-    solver = FOOPINN(model, poisson_problem)
-    trainer = Trainer(solver, max_epochs=2, accelerator='cpu')
-    trainer.train()
-
-def test_log():
-    solver = FOOPINN(model, poisson_problem)
-    trainer = Trainer(solver, max_epochs=2, accelerator='cpu')
-    trainer.train()
-    # assert the logged metrics are correct
-    logged_metrics = sorted(list(trainer.logged_metrics.keys()))
-    total_metrics = sorted(
-        list([key + '_loss' for key in poisson_problem.conditions.keys()])
-        + ['mean_loss'])
-    assert logged_metrics == total_metrics
\ No newline at end of file
+from pina.problem.zoo import Poisson2DSquareProblem as Poisson
+# from pina.equation import Equation
+# from pina.equation.equation_factory import FixedValue
+
+# def laplace_equation(input_, output_):
+#     force_term = (torch.sin(input_.extract(['x']) * torch.pi) *
+#                   torch.sin(input_.extract(['y']) * torch.pi))
+#     delta_u = laplacian(output_.extract(['u']), input_)
+#     return delta_u - force_term
+
+
+# my_laplace = Equation(laplace_equation)
+# in_ = LabelTensor(torch.tensor([[0., 1.]]), ['x', 'y'])
+# out_ = LabelTensor(torch.tensor([[0.]]), ['u'])
+# in2_ = LabelTensor(torch.rand(60, 2), ['x', 'y'])
+# out2_ = LabelTensor(torch.rand(60, 1), ['u'])
+
+
+
+# class Poisson(SpatialProblem):
+#     output_variables = ['u']
+#     spatial_domain = CartesianDomain({'x': [0, 1], 'y': [0, 1]})
+
+#     conditions = {
+#         'gamma1': Condition(
+#             location=CartesianDomain({'x': [0, 1], 'y':  1}),
+#             equation=FixedValue(0.0)),
+#         'gamma2': Condition(
+#             location=CartesianDomain({'x': [0, 1], 'y': 0}),
+#             equation=FixedValue(0.0)),
+#         'gamma3': Condition(
+#             location=CartesianDomain({'x':  1, 'y': [0, 1]}),
+#             equation=FixedValue(0.0)),
+#         'gamma4': Condition(
+#             location=CartesianDomain({'x': 0, 'y': [0, 1]}),
+#             equation=FixedValue(0.0)),
+#         'D': Condition(
+#             input_points=LabelTensor(torch.rand(size=(100, 2)), ['x', 'y']),
+#             equation=my_laplace),
+#         'data': Condition(
+#             input_points=in_,
+#             output_points=out_),
+#         'data2': Condition(
+#             input_points=in2_,
+#             output_points=out2_)
+#     }
+
+#     def poisson_sol(self, pts):
+#         return -(torch.sin(pts.extract(['x']) * torch.pi) *
+#                  torch.sin(pts.extract(['y']) * torch.pi)) / (2 * torch.pi**2)
+
+#     truth_solution = poisson_sol
+
+# from pina import TorchOptimizer
+
+# class FOOPINN(PINNInterface):
+#     def __init__(self, model, problem):
+#         super().__init__(models=[model], problem=problem,
+#                          optimizers=TorchOptimizer(torch.optim.Adam, lr=1e-3),
+#                          loss=torch.nn.MSELoss())
+#     def forward(self, x):
+#         return self.models[0](x)
+
+#     def loss_phys(self, samples, equation):
+#         residual = self.compute_residual(samples=samples, equation=equation)
+#         loss_value = self.loss(
+#             torch.zeros_like(residual, requires_grad=True), residual
+#         )
+#         self.store_log(loss_value=float(loss_value))
+#         return loss_value
+
+# # make the problem
+# poisson_problem = Poisson()
+# poisson_problem.discretise_domain(100)
+# model = FeedForward(len(poisson_problem.input_variables),
+#                     len(poisson_problem.output_variables))
+# model_extra_feats = FeedForward(
+#     len(poisson_problem.input_variables) + 1,
+#     len(poisson_problem.output_variables))
+
+
+# def test_constructor():
+#     with pytest.raises(TypeError):
+#         PINNInterface()
+#     # a simple pinn built with PINNInterface
+#     FOOPINN(model, poisson_problem)
+
+# def test_train_step():
+#     solver = FOOPINN(model, poisson_problem)
+#     trainer = Trainer(solver, max_epochs=2, accelerator='cpu')
+#     trainer.train()
+
+# def test_log():
+#     solver = FOOPINN(model, poisson_problem)
+#     trainer = Trainer(solver, max_epochs=2, accelerator='cpu')
+#     trainer.train()
+#     # assert the logged metrics are correct
+#     logged_metrics = sorted(list(trainer.logged_metrics.keys()))
+#     total_metrics = sorted(
+#         list([key + '_loss' for key in poisson_problem.conditions.keys()])
+#         + ['mean_loss'])
+#     assert logged_metrics == total_metrics
\ No newline at end of file
diff --git a/tests/test_solvers/test_causalpinn.py b/tests/test_solvers/test_causalpinn.py
index 7629d51b1..fb0aa284c 100644
--- a/tests/test_solvers/test_causalpinn.py
+++ b/tests/test_solvers/test_causalpinn.py
@@ -8,215 +8,237 @@
 from pina.solvers import CausalPINN
 from pina.trainer import Trainer
 from pina.model import FeedForward
-from pina.equation.equation import Equation
+from pina.equation import Equation
 from pina.equation.equation_factory import FixedValue
-from pina.loss.loss_interface import LpLoss
+from pina.loss import LpLoss
 
 
 
-class FooProblem(SpatialProblem):
-    '''
-    Foo problem formulation.
-    '''
-    output_variables = ['u']
-    conditions = {}
-    spatial_domain = None
+# class FooProblem(SpatialProblem):
+#     '''
+#     Foo problem formulation.
+#     '''
+#     output_variables = ['u']
+#     conditions = {}
+#     spatial_domain = None
 
 
-class InverseDiffusionReactionSystem(TimeDependentProblem, SpatialProblem, InverseProblem):
+# class InverseDiffusionReactionSystem(TimeDependentProblem, SpatialProblem, InverseProblem):
 
-    def diffusionreaction(input_, output_, params_):
-        x = input_.extract('x')
-        t = input_.extract('t')
-        u_t = grad(output_, input_, d='t')
-        u_x = grad(output_, input_, d='x')
-        u_xx = grad(u_x, input_, d='x')
-        r = torch.exp(-t) * (1.5 * torch.sin(2*x) + (8/3)*torch.sin(3*x) +
-                             (15/4)*torch.sin(4*x) + (63/8)*torch.sin(8*x))
-        return u_t - params_['mu']*u_xx - r
+#     def diffusionreaction(input_, output_, params_):
+#         x = input_.extract('x')
+#         t = input_.extract('t')
+#         u_t = grad(output_, input_, d='t')
+#         u_x = grad(output_, input_, d='x')
+#         u_xx = grad(u_x, input_, d='x')
+#         r = torch.exp(-t) * (1.5 * torch.sin(2*x) + (8/3)*torch.sin(3*x) +
+#                              (15/4)*torch.sin(4*x) + (63/8)*torch.sin(8*x))
+#         return u_t - params_['mu']*u_xx - r
 
-    def _solution(self, pts):
-        t = pts.extract('t')
-        x = pts.extract('x')
-        return torch.exp(-t) * (torch.sin(x) + (1/2)*torch.sin(2*x) +
-                                (1/3)*torch.sin(3*x) + (1/4)*torch.sin(4*x) +
-                                (1/8)*torch.sin(8*x))
+#     def _solution(self, pts):
+#         t = pts.extract('t')
+#         x = pts.extract('x')
+#         return torch.exp(-t) * (torch.sin(x) + (1/2)*torch.sin(2*x) +
+#                                 (1/3)*torch.sin(3*x) + (1/4)*torch.sin(4*x) +
+#                                 (1/8)*torch.sin(8*x))
     
-    # assign output/ spatial and temporal variables
-    output_variables = ['u']
-    spatial_domain = CartesianDomain({'x': [-torch.pi, torch.pi]})
-    temporal_domain = CartesianDomain({'t': [0, 1]})
-    unknown_parameter_domain = CartesianDomain({'mu': [-1, 1]})
-
-    # problem condition statement
-    conditions = {
-        'D': Condition(location=CartesianDomain({'x': [-torch.pi, torch.pi],
-                                                 't': [0, 1]}),
-                       equation=Equation(diffusionreaction)),
-        'data' : Condition(input_points=LabelTensor(torch.tensor([[0., 0.]]), ['x', 't']),
-                       output_points=LabelTensor(torch.tensor([[0.]]), ['u'])),
-    }
-
-class DiffusionReactionSystem(TimeDependentProblem, SpatialProblem):
-
-    def diffusionreaction(input_, output_):
-        x = input_.extract('x')
-        t = input_.extract('t')
-        u_t = grad(output_, input_, d='t')
-        u_x = grad(output_, input_, d='x')
-        u_xx = grad(u_x, input_, d='x')
-        r = torch.exp(-t) * (1.5 * torch.sin(2*x) + (8/3)*torch.sin(3*x) +
-                             (15/4)*torch.sin(4*x) + (63/8)*torch.sin(8*x))
-        return u_t - u_xx - r
-
-    def _solution(self, pts):
-        t = pts.extract('t')
-        x = pts.extract('x')
-        return torch.exp(-t) * (torch.sin(x) + (1/2)*torch.sin(2*x) +
-                                (1/3)*torch.sin(3*x) + (1/4)*torch.sin(4*x) +
-                                (1/8)*torch.sin(8*x))
+#     # assign output/ spatial and temporal variables
+#     output_variables = ['u']
+#     spatial_domain = CartesianDomain({'x': [-torch.pi, torch.pi]})
+#     temporal_domain = CartesianDomain({'t': [0, 1]})
+#     unknown_parameter_domain = CartesianDomain({'mu': [-1, 1]})
+
+#     # problem condition statement
+#     conditions = {
+#         'D': Condition(location=CartesianDomain({'x': [-torch.pi, torch.pi],
+#                                                  't': [0, 1]}),
+#                        equation=Equation(diffusionreaction)),
+#         'data' : Condition(input_points=LabelTensor(torch.tensor([[0., 0.]]), ['x', 't']),
+#                        output_points=LabelTensor(torch.tensor([[0.]]), ['u'])),
+#     }
+
+# class DiffusionReactionSystem(TimeDependentProblem, SpatialProblem):
+
+#     def diffusionreaction(input_, output_):
+#         x = input_.extract('x')
+#         t = input_.extract('t')
+#         u_t = grad(output_, input_, d='t')
+#         u_x = grad(output_, input_, d='x')
+#         u_xx = grad(u_x, input_, d='x')
+#         r = torch.exp(-t) * (1.5 * torch.sin(2*x) + (8/3)*torch.sin(3*x) +
+#                              (15/4)*torch.sin(4*x) + (63/8)*torch.sin(8*x))
+#         return u_t - u_xx - r
+
+#     def _solution(self, pts):
+#         t = pts.extract('t')
+#         x = pts.extract('x')
+#         return torch.exp(-t) * (torch.sin(x) + (1/2)*torch.sin(2*x) +
+#                                 (1/3)*torch.sin(3*x) + (1/4)*torch.sin(4*x) +
+#                                 (1/8)*torch.sin(8*x))
     
-    # assign output/ spatial and temporal variables
-    output_variables = ['u']
-    spatial_domain = CartesianDomain({'x': [-torch.pi, torch.pi]})
-    temporal_domain = CartesianDomain({'t': [0, 1]})
-
-    # problem condition statement
-    conditions = {
-        'D': Condition(location=CartesianDomain({'x': [-torch.pi, torch.pi],
-                                                 't': [0, 1]}),
-                       equation=Equation(diffusionreaction)),
-    }
-
-class myFeature(torch.nn.Module):
-    """
-    Feature: sin(x)
-    """
-
-    def __init__(self):
-        super(myFeature, self).__init__()
-
-    def forward(self, x):
-        t = (torch.sin(x.extract(['x']) * torch.pi))
-        return LabelTensor(t, ['sin(x)'])
-
-
-# make the problem
-problem = DiffusionReactionSystem()
-model = FeedForward(len(problem.input_variables),
-                    len(problem.output_variables))
-model_extra_feats = FeedForward(
-    len(problem.input_variables) + 1,
-    len(problem.output_variables))
-extra_feats = [myFeature()]
-
-
-def test_constructor():
-    CausalPINN(problem=problem, model=model, extra_features=None)
-
-    with pytest.raises(ValueError):
-        CausalPINN(FooProblem(), model=model, extra_features=None)
-
-
-def test_constructor_extra_feats():
-    model_extra_feats = FeedForward(
-        len(problem.input_variables) + 1,
-        len(problem.output_variables))
-    CausalPINN(problem=problem,
-         model=model_extra_feats,
-         extra_features=extra_feats)
-
-
-def test_train_cpu():
-    problem = DiffusionReactionSystem()
-    boundaries = ['D']
-    n = 10
-    problem.discretise_domain(n, 'grid', locations=boundaries)
-    pinn = CausalPINN(problem = problem,
-                 model=model, extra_features=None, loss=LpLoss())
-    trainer = Trainer(solver=pinn, max_epochs=1,
-                      accelerator='cpu', batch_size=20)
-    trainer.train()
-
-def test_log():
-    problem.discretise_domain(100)
-    solver = CausalPINN(problem = problem,
-                 model=model, extra_features=None, loss=LpLoss())
-    trainer = Trainer(solver, max_epochs=2, accelerator='cpu')
-    trainer.train()
-    # assert the logged metrics are correct
-    logged_metrics = sorted(list(trainer.logged_metrics.keys()))
-    total_metrics = sorted(
-        list([key + '_loss' for key in problem.conditions.keys()])
-        + ['mean_loss'])
-    assert logged_metrics == total_metrics
-
-def test_train_restore():
-    tmpdir = "tests/tmp_restore"
-    problem = DiffusionReactionSystem()
-    boundaries = ['D']
-    n = 10
-    problem.discretise_domain(n, 'grid', locations=boundaries)
-    pinn = CausalPINN(problem=problem,
-                model=model,
-                extra_features=None,
-                loss=LpLoss())
-    trainer = Trainer(solver=pinn,
-                      max_epochs=5,
-                      accelerator='cpu',
-                      default_root_dir=tmpdir)
-    trainer.train()
-    ntrainer = Trainer(solver=pinn, max_epochs=15, accelerator='cpu')
-    t = ntrainer.train(
-        ckpt_path=f'{tmpdir}/lightning_logs/version_0/'
-        'checkpoints/epoch=4-step=5.ckpt')
-    import shutil
-    shutil.rmtree(tmpdir)
-
-
-def test_train_load():
-    tmpdir = "tests/tmp_load"
-    problem = DiffusionReactionSystem()
-    boundaries = ['D']
-    n = 10
-    problem.discretise_domain(n, 'grid', locations=boundaries)
-    pinn = CausalPINN(problem=problem,
-                model=model,
-                extra_features=None,
-                loss=LpLoss())
-    trainer = Trainer(solver=pinn,
-                      max_epochs=15,
-                      accelerator='cpu',
-                      default_root_dir=tmpdir)
-    trainer.train()
-    new_pinn = CausalPINN.load_from_checkpoint(
-        f'{tmpdir}/lightning_logs/version_0/checkpoints/epoch=14-step=15.ckpt',
-        problem = problem, model=model)
-    test_pts = CartesianDomain({'x': [0, 1], 't': [0, 1]}).sample(10)
-    assert new_pinn.forward(test_pts).extract(['u']).shape == (10, 1)
-    assert new_pinn.forward(test_pts).extract(
-        ['u']).shape == pinn.forward(test_pts).extract(['u']).shape
-    torch.testing.assert_close(
-        new_pinn.forward(test_pts).extract(['u']),
-        pinn.forward(test_pts).extract(['u']))
-    import shutil
-    shutil.rmtree(tmpdir)
-
-def test_train_inverse_problem_cpu():
-    problem = InverseDiffusionReactionSystem()
-    boundaries = ['D']
-    n = 100
-    problem.discretise_domain(n, 'random', locations=boundaries)
-    pinn = CausalPINN(problem = problem,
-                 model=model, extra_features=None, loss=LpLoss())
-    trainer = Trainer(solver=pinn, max_epochs=1,
-                      accelerator='cpu', batch_size=20)
-    trainer.train()
-
-
-# # TODO does not currently work
-# def test_train_inverse_problem_restore():
-#     tmpdir = "tests/tmp_restore_inv"
+#     # assign output/ spatial and temporal variables
+#     output_variables = ['u']
+#     spatial_domain = CartesianDomain({'x': [-torch.pi, torch.pi]})
+#     temporal_domain = CartesianDomain({'t': [0, 1]})
+
+#     # problem condition statement
+#     conditions = {
+#         'D': Condition(location=CartesianDomain({'x': [-torch.pi, torch.pi],
+#                                                  't': [0, 1]}),
+#                        equation=Equation(diffusionreaction)),
+#     }
+
+# class myFeature(torch.nn.Module):
+#     """
+#     Feature: sin(x)
+#     """
+
+#     def __init__(self):
+#         super(myFeature, self).__init__()
+
+#     def forward(self, x):
+#         t = (torch.sin(x.extract(['x']) * torch.pi))
+#         return LabelTensor(t, ['sin(x)'])
+
+
+# # make the problem
+# problem = DiffusionReactionSystem()
+# model = FeedForward(len(problem.input_variables),
+#                     len(problem.output_variables))
+# model_extra_feats = FeedForward(
+#     len(problem.input_variables) + 1,
+#     len(problem.output_variables))
+# extra_feats = [myFeature()]
+
+
+# def test_constructor():
+#     CausalPINN(problem=problem, model=model, extra_features=None)
+
+#     with pytest.raises(ValueError):
+#         CausalPINN(FooProblem(), model=model, extra_features=None)
+
+
+# def test_constructor_extra_feats():
+#     model_extra_feats = FeedForward(
+#         len(problem.input_variables) + 1,
+#         len(problem.output_variables))
+#     CausalPINN(problem=problem,
+#          model=model_extra_feats,
+#          extra_features=extra_feats)
+
+
+# def test_train_cpu():
+#     problem = DiffusionReactionSystem()
+#     boundaries = ['D']
+#     n = 10
+#     problem.discretise_domain(n, 'grid', locations=boundaries)
+#     pinn = CausalPINN(problem = problem,
+#                  model=model, extra_features=None, loss=LpLoss())
+#     trainer = Trainer(solver=pinn, max_epochs=1,
+#                       accelerator='cpu', batch_size=20)
+#     trainer.train()
+
+# def test_log():
+#     problem.discretise_domain(100)
+#     solver = CausalPINN(problem = problem,
+#                  model=model, extra_features=None, loss=LpLoss())
+#     trainer = Trainer(solver, max_epochs=2, accelerator='cpu')
+#     trainer.train()
+#     # assert the logged metrics are correct
+#     logged_metrics = sorted(list(trainer.logged_metrics.keys()))
+#     total_metrics = sorted(
+#         list([key + '_loss' for key in problem.conditions.keys()])
+#         + ['mean_loss'])
+#     assert logged_metrics == total_metrics
+
+# def test_train_restore():
+#     tmpdir = "tests/tmp_restore"
+#     problem = DiffusionReactionSystem()
+#     boundaries = ['D']
+#     n = 10
+#     problem.discretise_domain(n, 'grid', locations=boundaries)
+#     pinn = CausalPINN(problem=problem,
+#                 model=model,
+#                 extra_features=None,
+#                 loss=LpLoss())
+#     trainer = Trainer(solver=pinn,
+#                       max_epochs=5,
+#                       accelerator='cpu',
+#                       default_root_dir=tmpdir)
+#     trainer.train()
+#     ntrainer = Trainer(solver=pinn, max_epochs=15, accelerator='cpu')
+#     t = ntrainer.train(
+#         ckpt_path=f'{tmpdir}/lightning_logs/version_0/'
+#         'checkpoints/epoch=4-step=5.ckpt')
+#     import shutil
+#     shutil.rmtree(tmpdir)
+
+
+# def test_train_load():
+#     tmpdir = "tests/tmp_load"
+#     problem = DiffusionReactionSystem()
+#     boundaries = ['D']
+#     n = 10
+#     problem.discretise_domain(n, 'grid', locations=boundaries)
+#     pinn = CausalPINN(problem=problem,
+#                 model=model,
+#                 extra_features=None,
+#                 loss=LpLoss())
+#     trainer = Trainer(solver=pinn,
+#                       max_epochs=15,
+#                       accelerator='cpu',
+#                       default_root_dir=tmpdir)
+#     trainer.train()
+#     new_pinn = CausalPINN.load_from_checkpoint(
+#         f'{tmpdir}/lightning_logs/version_0/checkpoints/epoch=14-step=15.ckpt',
+#         problem = problem, model=model)
+#     test_pts = CartesianDomain({'x': [0, 1], 't': [0, 1]}).sample(10)
+#     assert new_pinn.forward(test_pts).extract(['u']).shape == (10, 1)
+#     assert new_pinn.forward(test_pts).extract(
+#         ['u']).shape == pinn.forward(test_pts).extract(['u']).shape
+#     torch.testing.assert_close(
+#         new_pinn.forward(test_pts).extract(['u']),
+#         pinn.forward(test_pts).extract(['u']))
+#     import shutil
+#     shutil.rmtree(tmpdir)
+
+# def test_train_inverse_problem_cpu():
+#     problem = InverseDiffusionReactionSystem()
+#     boundaries = ['D']
+#     n = 100
+#     problem.discretise_domain(n, 'random', locations=boundaries)
+#     pinn = CausalPINN(problem = problem,
+#                  model=model, extra_features=None, loss=LpLoss())
+#     trainer = Trainer(solver=pinn, max_epochs=1,
+#                       accelerator='cpu', batch_size=20)
+#     trainer.train()
+
+
+# # # TODO does not currently work
+# # def test_train_inverse_problem_restore():
+# #     tmpdir = "tests/tmp_restore_inv"
+# #     problem = InverseDiffusionReactionSystem()
+# #     boundaries = ['D']
+# #     n = 100
+# #     problem.discretise_domain(n, 'random', locations=boundaries)
+# #     pinn = CausalPINN(problem=problem,
+# #                 model=model,
+# #                 extra_features=None,
+# #                 loss=LpLoss())
+# #     trainer = Trainer(solver=pinn,
+# #                       max_epochs=5,
+# #                       accelerator='cpu',
+# #                       default_root_dir=tmpdir)
+# #     trainer.train()
+# #     ntrainer = Trainer(solver=pinn, max_epochs=5, accelerator='cpu')
+# #     t = ntrainer.train(
+# #         ckpt_path=f'{tmpdir}/lightning_logs/version_0/checkpoints/epoch=4-step=5.ckpt')
+# #     import shutil
+# #     shutil.rmtree(tmpdir)
+
+
+# def test_train_inverse_problem_load():
+#     tmpdir = "tests/tmp_load_inv"
 #     problem = InverseDiffusionReactionSystem()
 #     boundaries = ['D']
 #     n = 100
@@ -226,53 +248,31 @@ def test_train_inverse_problem_cpu():
 #                 extra_features=None,
 #                 loss=LpLoss())
 #     trainer = Trainer(solver=pinn,
-#                       max_epochs=5,
+#                       max_epochs=15,
 #                       accelerator='cpu',
 #                       default_root_dir=tmpdir)
 #     trainer.train()
-#     ntrainer = Trainer(solver=pinn, max_epochs=5, accelerator='cpu')
-#     t = ntrainer.train(
-#         ckpt_path=f'{tmpdir}/lightning_logs/version_0/checkpoints/epoch=4-step=5.ckpt')
+#     new_pinn = CausalPINN.load_from_checkpoint(
+#         f'{tmpdir}/lightning_logs/version_0/checkpoints/epoch=14-step=30.ckpt',
+#         problem = problem, model=model)
+#     test_pts = CartesianDomain({'x': [0, 1], 't': [0, 1]}).sample(10)
+#     assert new_pinn.forward(test_pts).extract(['u']).shape == (10, 1)
+#     assert new_pinn.forward(test_pts).extract(
+#         ['u']).shape == pinn.forward(test_pts).extract(['u']).shape
+#     torch.testing.assert_close(
+#         new_pinn.forward(test_pts).extract(['u']),
+#         pinn.forward(test_pts).extract(['u']))
 #     import shutil
 #     shutil.rmtree(tmpdir)
 
 
-def test_train_inverse_problem_load():
-    tmpdir = "tests/tmp_load_inv"
-    problem = InverseDiffusionReactionSystem()
-    boundaries = ['D']
-    n = 100
-    problem.discretise_domain(n, 'random', locations=boundaries)
-    pinn = CausalPINN(problem=problem,
-                model=model,
-                extra_features=None,
-                loss=LpLoss())
-    trainer = Trainer(solver=pinn,
-                      max_epochs=15,
-                      accelerator='cpu',
-                      default_root_dir=tmpdir)
-    trainer.train()
-    new_pinn = CausalPINN.load_from_checkpoint(
-        f'{tmpdir}/lightning_logs/version_0/checkpoints/epoch=14-step=30.ckpt',
-        problem = problem, model=model)
-    test_pts = CartesianDomain({'x': [0, 1], 't': [0, 1]}).sample(10)
-    assert new_pinn.forward(test_pts).extract(['u']).shape == (10, 1)
-    assert new_pinn.forward(test_pts).extract(
-        ['u']).shape == pinn.forward(test_pts).extract(['u']).shape
-    torch.testing.assert_close(
-        new_pinn.forward(test_pts).extract(['u']),
-        pinn.forward(test_pts).extract(['u']))
-    import shutil
-    shutil.rmtree(tmpdir)
-
-
-def test_train_extra_feats_cpu():
-    problem = DiffusionReactionSystem()
-    boundaries = ['D']
-    n = 10
-    problem.discretise_domain(n, 'grid', locations=boundaries)
-    pinn = CausalPINN(problem=problem,
-                model=model_extra_feats,
-                extra_features=extra_feats)
-    trainer = Trainer(solver=pinn, max_epochs=5, accelerator='cpu')
-    trainer.train()
\ No newline at end of file
+# def test_train_extra_feats_cpu():
+#     problem = DiffusionReactionSystem()
+#     boundaries = ['D']
+#     n = 10
+#     problem.discretise_domain(n, 'grid', locations=boundaries)
+#     pinn = CausalPINN(problem=problem,
+#                 model=model_extra_feats,
+#                 extra_features=extra_feats)
+#     trainer = Trainer(solver=pinn, max_epochs=5, accelerator='cpu')
+#     trainer.train()
\ No newline at end of file
diff --git a/tests/test_solvers/test_competitive_pinn.py b/tests/test_solvers/test_competitive_pinn.py
index 9facf2a44..138170902 100644
--- a/tests/test_solvers/test_competitive_pinn.py
+++ b/tests/test_solvers/test_competitive_pinn.py
@@ -8,240 +8,179 @@
 from pina.solvers import CompetitivePINN as PINN
 from pina.trainer import Trainer
 from pina.model import FeedForward
-from pina.equation.equation import Equation
+from pina.equation import Equation
 from pina.equation.equation_factory import FixedValue
-from pina.loss.loss_interface import LpLoss
-
-
-def laplace_equation(input_, output_):
-    force_term = (torch.sin(input_.extract(['x']) * torch.pi) *
-                  torch.sin(input_.extract(['y']) * torch.pi))
-    delta_u = laplacian(output_.extract(['u']), input_)
-    return delta_u - force_term
-
-
-my_laplace = Equation(laplace_equation)
-in_ = LabelTensor(torch.tensor([[0., 1.]]), ['x', 'y'])
-out_ = LabelTensor(torch.tensor([[0.]]), ['u'])
-in2_ = LabelTensor(torch.rand(60, 2), ['x', 'y'])
-out2_ = LabelTensor(torch.rand(60, 1), ['u'])
-
-
-class InversePoisson(SpatialProblem, InverseProblem):
-    '''
-    Problem definition for the Poisson equation.
-    '''
-    output_variables = ['u']
-    x_min = -2
-    x_max = 2
-    y_min = -2
-    y_max = 2
-    data_input = LabelTensor(torch.rand(10, 2), ['x', 'y'])
-    data_output = LabelTensor(torch.rand(10, 1), ['u'])
-    spatial_domain = CartesianDomain({'x': [x_min, x_max], 'y': [y_min, y_max]})
-    # define the ranges for the parameters
-    unknown_parameter_domain = CartesianDomain({'mu1': [-1, 1], 'mu2': [-1, 1]})
-
-    def laplace_equation(input_, output_, params_):
-        '''
-        Laplace equation with a force term.
-        '''
-        force_term = torch.exp(
-                - 2*(input_.extract(['x']) - params_['mu1'])**2
-                - 2*(input_.extract(['y']) - params_['mu2'])**2)
-        delta_u = laplacian(output_, input_, components=['u'], d=['x', 'y'])
-
-        return delta_u - force_term
-
-    # define the conditions for the loss (boundary conditions, equation, data)
-    conditions = {
-        'gamma1': Condition(location=CartesianDomain({'x': [x_min, x_max],
-            'y':  y_max}),
-            equation=FixedValue(0.0, components=['u'])),
-        'gamma2': Condition(location=CartesianDomain(
-            {'x': [x_min, x_max], 'y': y_min
-            }),
-            equation=FixedValue(0.0, components=['u'])),
-        'gamma3': Condition(location=CartesianDomain(
-            {'x':  x_max, 'y': [y_min, y_max]
-            }),
-            equation=FixedValue(0.0, components=['u'])),
-        'gamma4': Condition(location=CartesianDomain(
-            {'x': x_min, 'y': [y_min, y_max]
-            }),
-            equation=FixedValue(0.0, components=['u'])),
-        'D': Condition(location=CartesianDomain(
-            {'x': [x_min, x_max], 'y': [y_min, y_max]
-            }),
-        equation=Equation(laplace_equation)),
-        'data': Condition(input_points=data_input.extract(['x', 'y']),
-                          output_points=data_output)
-    }
-
-
-class Poisson(SpatialProblem):
-    output_variables = ['u']
-    spatial_domain = CartesianDomain({'x': [0, 1], 'y': [0, 1]})
-
-    conditions = {
-        'gamma1': Condition(
-            location=CartesianDomain({'x': [0, 1], 'y':  1}),
-            equation=FixedValue(0.0)),
-        'gamma2': Condition(
-            location=CartesianDomain({'x': [0, 1], 'y': 0}),
-            equation=FixedValue(0.0)),
-        'gamma3': Condition(
-            location=CartesianDomain({'x':  1, 'y': [0, 1]}),
-            equation=FixedValue(0.0)),
-        'gamma4': Condition(
-            location=CartesianDomain({'x': 0, 'y': [0, 1]}),
-            equation=FixedValue(0.0)),
-        'D': Condition(
-            input_points=LabelTensor(torch.rand(size=(100, 2)), ['x', 'y']),
-            equation=my_laplace),
-        'data': Condition(
-            input_points=in_,
-            output_points=out_),
-        'data2': Condition(
-            input_points=in2_,
-            output_points=out2_)
-    }
-
-    def poisson_sol(self, pts):
-        return -(torch.sin(pts.extract(['x']) * torch.pi) *
-                 torch.sin(pts.extract(['y']) * torch.pi)) / (2 * torch.pi**2)
-
-    truth_solution = poisson_sol
-
-
-class myFeature(torch.nn.Module):
-    """
-    Feature: sin(x)
-    """
-
-    def __init__(self):
-        super(myFeature, self).__init__()
-
-    def forward(self, x):
-        t = (torch.sin(x.extract(['x']) * torch.pi) *
-             torch.sin(x.extract(['y']) * torch.pi))
-        return LabelTensor(t, ['sin(x)sin(y)'])
-
-
-# make the problem
-poisson_problem = Poisson()
-model = FeedForward(len(poisson_problem.input_variables),
-                    len(poisson_problem.output_variables))
-model_extra_feats = FeedForward(
-    len(poisson_problem.input_variables) + 1,
-    len(poisson_problem.output_variables))
-extra_feats = [myFeature()]
-
-
-def test_constructor():
-    PINN(problem=poisson_problem, model=model)
-    PINN(problem=poisson_problem, model=model, discriminator = model)
-
-
-def test_constructor_extra_feats():
-    with pytest.raises(TypeError):
-        model_extra_feats = FeedForward(
-            len(poisson_problem.input_variables) + 1,
-            len(poisson_problem.output_variables))
-        PINN(problem=poisson_problem,
-            model=model_extra_feats,
-            extra_features=extra_feats)
-
-
-def test_train_cpu():
-    poisson_problem = Poisson()
-    boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4']
-    n = 10
-    poisson_problem.discretise_domain(n, 'grid', locations=boundaries)
-    pinn = PINN(problem = poisson_problem, model=model, loss=LpLoss())
-    trainer = Trainer(solver=pinn, max_epochs=1,
-                      accelerator='cpu', batch_size=20)
-    trainer.train()
-
-def test_log():
-    poisson_problem.discretise_domain(100)
-    solver = PINN(problem = poisson_problem, model=model, loss=LpLoss())
-    trainer = Trainer(solver, max_epochs=2, accelerator='cpu')
-    trainer.train()
-    # assert the logged metrics are correct
-    logged_metrics = sorted(list(trainer.logged_metrics.keys()))
-    total_metrics = sorted(
-        list([key + '_loss' for key in poisson_problem.conditions.keys()])
-        + ['mean_loss'])
-    assert logged_metrics == total_metrics
-
-def test_train_restore():
-    tmpdir = "tests/tmp_restore"
-    poisson_problem = Poisson()
-    boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4']
-    n = 10
-    poisson_problem.discretise_domain(n, 'grid', locations=boundaries)
-    pinn = PINN(problem=poisson_problem,
-                model=model,
-                loss=LpLoss())
-    trainer = Trainer(solver=pinn,
-                      max_epochs=5,
-                      accelerator='cpu',
-                      default_root_dir=tmpdir)
-    trainer.train()
-    ntrainer = Trainer(solver=pinn, max_epochs=15, accelerator='cpu')
-    t = ntrainer.train(
-        ckpt_path=f'{tmpdir}/lightning_logs/version_0/'
-        'checkpoints/epoch=4-step=10.ckpt')
-    import shutil
-    shutil.rmtree(tmpdir)
-
-
-def test_train_load():
-    tmpdir = "tests/tmp_load"
-    poisson_problem = Poisson()
-    boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4']
-    n = 10
-    poisson_problem.discretise_domain(n, 'grid', locations=boundaries)
-    pinn = PINN(problem=poisson_problem,
-                model=model,
-                loss=LpLoss())
-    trainer = Trainer(solver=pinn,
-                      max_epochs=15,
-                      accelerator='cpu',
-                      default_root_dir=tmpdir)
-    trainer.train()
-    new_pinn = PINN.load_from_checkpoint(
-        f'{tmpdir}/lightning_logs/version_0/checkpoints/epoch=14-step=30.ckpt',
-        problem = poisson_problem, model=model)
-    test_pts = CartesianDomain({'x': [0, 1], 'y': [0, 1]}).sample(10)
-    assert new_pinn.forward(test_pts).extract(['u']).shape == (10, 1)
-    assert new_pinn.forward(test_pts).extract(
-        ['u']).shape == pinn.forward(test_pts).extract(['u']).shape
-    torch.testing.assert_close(
-        new_pinn.forward(test_pts).extract(['u']),
-        pinn.forward(test_pts).extract(['u']))
-    import shutil
-    shutil.rmtree(tmpdir)
-
-def test_train_inverse_problem_cpu():
-    poisson_problem = InversePoisson()
-    boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4', 'D']
-    n = 100
-    poisson_problem.discretise_domain(n, 'random', locations=boundaries)
-    pinn = PINN(problem = poisson_problem, model=model, loss=LpLoss())
-    trainer = Trainer(solver=pinn, max_epochs=1,
-                      accelerator='cpu', batch_size=20)
-    trainer.train()
-
-
-# # TODO does not currently work
-# def test_train_inverse_problem_restore():
-#     tmpdir = "tests/tmp_restore_inv"
-#     poisson_problem = InversePoisson()
-#     boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4', 'D']
-#     n = 100
-#     poisson_problem.discretise_domain(n, 'random', locations=boundaries)
+from pina.loss import LpLoss
+
+
+# def laplace_equation(input_, output_):
+#     force_term = (torch.sin(input_.extract(['x']) * torch.pi) *
+#                   torch.sin(input_.extract(['y']) * torch.pi))
+#     delta_u = laplacian(output_.extract(['u']), input_)
+#     return delta_u - force_term
+
+
+# my_laplace = Equation(laplace_equation)
+# in_ = LabelTensor(torch.tensor([[0., 1.]]), ['x', 'y'])
+# out_ = LabelTensor(torch.tensor([[0.]]), ['u'])
+# in2_ = LabelTensor(torch.rand(60, 2), ['x', 'y'])
+# out2_ = LabelTensor(torch.rand(60, 1), ['u'])
+
+
+# class InversePoisson(SpatialProblem, InverseProblem):
+#     '''
+#     Problem definition for the Poisson equation.
+#     '''
+#     output_variables = ['u']
+#     x_min = -2
+#     x_max = 2
+#     y_min = -2
+#     y_max = 2
+#     data_input = LabelTensor(torch.rand(10, 2), ['x', 'y'])
+#     data_output = LabelTensor(torch.rand(10, 1), ['u'])
+#     spatial_domain = CartesianDomain({'x': [x_min, x_max], 'y': [y_min, y_max]})
+#     # define the ranges for the parameters
+#     unknown_parameter_domain = CartesianDomain({'mu1': [-1, 1], 'mu2': [-1, 1]})
+
+#     def laplace_equation(input_, output_, params_):
+#         '''
+#         Laplace equation with a force term.
+#         '''
+#         force_term = torch.exp(
+#                 - 2*(input_.extract(['x']) - params_['mu1'])**2
+#                 - 2*(input_.extract(['y']) - params_['mu2'])**2)
+#         delta_u = laplacian(output_, input_, components=['u'], d=['x', 'y'])
+
+#         return delta_u - force_term
+
+#     # define the conditions for the loss (boundary conditions, equation, data)
+#     conditions = {
+#         'gamma1': Condition(location=CartesianDomain({'x': [x_min, x_max],
+#             'y':  y_max}),
+#             equation=FixedValue(0.0, components=['u'])),
+#         'gamma2': Condition(location=CartesianDomain(
+#             {'x': [x_min, x_max], 'y': y_min
+#             }),
+#             equation=FixedValue(0.0, components=['u'])),
+#         'gamma3': Condition(location=CartesianDomain(
+#             {'x':  x_max, 'y': [y_min, y_max]
+#             }),
+#             equation=FixedValue(0.0, components=['u'])),
+#         'gamma4': Condition(location=CartesianDomain(
+#             {'x': x_min, 'y': [y_min, y_max]
+#             }),
+#             equation=FixedValue(0.0, components=['u'])),
+#         'D': Condition(location=CartesianDomain(
+#             {'x': [x_min, x_max], 'y': [y_min, y_max]
+#             }),
+#         equation=Equation(laplace_equation)),
+#         'data': Condition(input_points=data_input.extract(['x', 'y']),
+#                           output_points=data_output)
+#     }
+
+
+# class Poisson(SpatialProblem):
+#     output_variables = ['u']
+#     spatial_domain = CartesianDomain({'x': [0, 1], 'y': [0, 1]})
+
+#     conditions = {
+#         'gamma1': Condition(
+#             location=CartesianDomain({'x': [0, 1], 'y':  1}),
+#             equation=FixedValue(0.0)),
+#         'gamma2': Condition(
+#             location=CartesianDomain({'x': [0, 1], 'y': 0}),
+#             equation=FixedValue(0.0)),
+#         'gamma3': Condition(
+#             location=CartesianDomain({'x':  1, 'y': [0, 1]}),
+#             equation=FixedValue(0.0)),
+#         'gamma4': Condition(
+#             location=CartesianDomain({'x': 0, 'y': [0, 1]}),
+#             equation=FixedValue(0.0)),
+#         'D': Condition(
+#             input_points=LabelTensor(torch.rand(size=(100, 2)), ['x', 'y']),
+#             equation=my_laplace),
+#         'data': Condition(
+#             input_points=in_,
+#             output_points=out_),
+#         'data2': Condition(
+#             input_points=in2_,
+#             output_points=out2_)
+#     }
+
+#     def poisson_sol(self, pts):
+#         return -(torch.sin(pts.extract(['x']) * torch.pi) *
+#                  torch.sin(pts.extract(['y']) * torch.pi)) / (2 * torch.pi**2)
+
+#     truth_solution = poisson_sol
+
+
+# class myFeature(torch.nn.Module):
+#     """
+#     Feature: sin(x)
+#     """
+
+#     def __init__(self):
+#         super(myFeature, self).__init__()
+
+#     def forward(self, x):
+#         t = (torch.sin(x.extract(['x']) * torch.pi) *
+#              torch.sin(x.extract(['y']) * torch.pi))
+#         return LabelTensor(t, ['sin(x)sin(y)'])
+
+
+# # make the problem
+# poisson_problem = Poisson()
+# model = FeedForward(len(poisson_problem.input_variables),
+#                     len(poisson_problem.output_variables))
+# model_extra_feats = FeedForward(
+#     len(poisson_problem.input_variables) + 1,
+#     len(poisson_problem.output_variables))
+# extra_feats = [myFeature()]
+
+
+# def test_constructor():
+#     PINN(problem=poisson_problem, model=model)
+#     PINN(problem=poisson_problem, model=model, discriminator = model)
+
+
+# def test_constructor_extra_feats():
+#     with pytest.raises(TypeError):
+#         model_extra_feats = FeedForward(
+#             len(poisson_problem.input_variables) + 1,
+#             len(poisson_problem.output_variables))
+#         PINN(problem=poisson_problem,
+#             model=model_extra_feats,
+#             extra_features=extra_feats)
+
+
+# def test_train_cpu():
+#     poisson_problem = Poisson()
+#     boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4']
+#     n = 10
+#     poisson_problem.discretise_domain(n, 'grid', locations=boundaries)
+#     pinn = PINN(problem = poisson_problem, model=model, loss=LpLoss())
+#     trainer = Trainer(solver=pinn, max_epochs=1,
+#                       accelerator='cpu', batch_size=20)
+#     trainer.train()
+
+# def test_log():
+#     poisson_problem.discretise_domain(100)
+#     solver = PINN(problem = poisson_problem, model=model, loss=LpLoss())
+#     trainer = Trainer(solver, max_epochs=2, accelerator='cpu')
+#     trainer.train()
+#     # assert the logged metrics are correct
+#     logged_metrics = sorted(list(trainer.logged_metrics.keys()))
+#     total_metrics = sorted(
+#         list([key + '_loss' for key in poisson_problem.conditions.keys()])
+#         + ['mean_loss'])
+#     assert logged_metrics == total_metrics
+
+# def test_train_restore():
+#     tmpdir = "tests/tmp_restore"
+#     poisson_problem = Poisson()
+#     boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4']
+#     n = 10
+#     poisson_problem.discretise_domain(n, 'grid', locations=boundaries)
 #     pinn = PINN(problem=poisson_problem,
 #                 model=model,
 #                 loss=LpLoss())
@@ -250,145 +189,153 @@ def test_train_inverse_problem_cpu():
 #                       accelerator='cpu',
 #                       default_root_dir=tmpdir)
 #     trainer.train()
-#     ntrainer = Trainer(solver=pinn, max_epochs=5, accelerator='cpu')
+#     ntrainer = Trainer(solver=pinn, max_epochs=15, accelerator='cpu')
 #     t = ntrainer.train(
-#         ckpt_path=f'{tmpdir}/lightning_logs/version_0/checkpoints/epoch=4-step=10.ckpt')
+#         ckpt_path=f'{tmpdir}/lightning_logs/version_0/'
+#         'checkpoints/epoch=4-step=10.ckpt')
 #     import shutil
 #     shutil.rmtree(tmpdir)
 
 
-def test_train_inverse_problem_load():
-    tmpdir = "tests/tmp_load_inv"
-    poisson_problem = InversePoisson()
-    boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4', 'D']
-    n = 100
-    poisson_problem.discretise_domain(n, 'random', locations=boundaries)
-    pinn = PINN(problem=poisson_problem,
-                model=model,
-                loss=LpLoss())
-    trainer = Trainer(solver=pinn,
-                      max_epochs=15,
-                      accelerator='cpu',
-                      default_root_dir=tmpdir)
-    trainer.train()
-    new_pinn = PINN.load_from_checkpoint(
-        f'{tmpdir}/lightning_logs/version_0/checkpoints/epoch=14-step=30.ckpt',
-        problem = poisson_problem, model=model)
-    test_pts = CartesianDomain({'x': [0, 1], 'y': [0, 1]}).sample(10)
-    assert new_pinn.forward(test_pts).extract(['u']).shape == (10, 1)
-    assert new_pinn.forward(test_pts).extract(
-        ['u']).shape == pinn.forward(test_pts).extract(['u']).shape
-    torch.testing.assert_close(
-        new_pinn.forward(test_pts).extract(['u']),
-        pinn.forward(test_pts).extract(['u']))
-    import shutil
-    shutil.rmtree(tmpdir)
-
-# # TODO fix asap. Basically sampling few variables
-# # works only if both variables are in a range.
-# # if one is fixed and the other not, this will
-# # not work. This test also needs to be fixed and
-# # insert in test problem not in test pinn.
-# def test_train_cpu_sampling_few_vars():
+# def test_train_load():
+#     tmpdir = "tests/tmp_load"
 #     poisson_problem = Poisson()
-#     boundaries = ['gamma1', 'gamma2', 'gamma3']
+#     boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4']
 #     n = 10
 #     poisson_problem.discretise_domain(n, 'grid', locations=boundaries)
-#     poisson_problem.discretise_domain(n, 'random', locations=['gamma4'], variables=['x'])
-#     poisson_problem.discretise_domain(n, 'random', locations=['gamma4'], variables=['y'])
-#     pinn = PINN(problem = poisson_problem, model=model, extra_features=None, loss=LpLoss())
-#     trainer = Trainer(solver=pinn, kwargs={'max_epochs' : 5, 'accelerator':'cpu'})
+#     pinn = PINN(problem=poisson_problem,
+#                 model=model,
+#                 loss=LpLoss())
+#     trainer = Trainer(solver=pinn,
+#                       max_epochs=15,
+#                       accelerator='cpu',
+#                       default_root_dir=tmpdir)
 #     trainer.train()
+#     new_pinn = PINN.load_from_checkpoint(
+#         f'{tmpdir}/lightning_logs/version_0/checkpoints/epoch=14-step=30.ckpt',
+#         problem = poisson_problem, model=model)
+#     test_pts = CartesianDomain({'x': [0, 1], 'y': [0, 1]}).sample(10)
+#     assert new_pinn.forward(test_pts).extract(['u']).shape == (10, 1)
+#     assert new_pinn.forward(test_pts).extract(
+#         ['u']).shape == pinn.forward(test_pts).extract(['u']).shape
+#     torch.testing.assert_close(
+#         new_pinn.forward(test_pts).extract(['u']),
+#         pinn.forward(test_pts).extract(['u']))
+#     import shutil
+#     shutil.rmtree(tmpdir)
 
-
-# TODO, fix GitHub actions to run also on GPU
-# def test_train_gpu():
-#     poisson_problem = Poisson()
-#     boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4']
-#     n = 10
-#     poisson_problem.discretise_domain(n, 'grid', locations=boundaries)
-#     pinn = PINN(problem = poisson_problem, model=model, extra_features=None, loss=LpLoss())
-#     trainer = Trainer(solver=pinn, kwargs={'max_epochs' : 5, 'accelerator':'gpu'})
+# def test_train_inverse_problem_cpu():
+#     poisson_problem = InversePoisson()
+#     boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4', 'D']
+#     n = 100
+#     poisson_problem.discretise_domain(n, 'random', locations=boundaries)
+#     pinn = PINN(problem = poisson_problem, model=model, loss=LpLoss())
+#     trainer = Trainer(solver=pinn, max_epochs=1,
+#                       accelerator='cpu', batch_size=20)
 #     trainer.train()
 
-# def test_train_gpu(): #TODO fix ASAP
-#     poisson_problem = Poisson()
-#     boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4']
-#     n = 10
-#     poisson_problem.discretise_domain(n, 'grid', locations=boundaries)
-#     poisson_problem.conditions.pop('data') # The input/output pts are allocated on cpu
-#     pinn = PINN(problem = poisson_problem, model=model, extra_features=None, loss=LpLoss())
-#     trainer = Trainer(solver=pinn, kwargs={'max_epochs' : 5, 'accelerator':'gpu'})
+
+# # # TODO does not currently work
+# # def test_train_inverse_problem_restore():
+# #     tmpdir = "tests/tmp_restore_inv"
+# #     poisson_problem = InversePoisson()
+# #     boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4', 'D']
+# #     n = 100
+# #     poisson_problem.discretise_domain(n, 'random', locations=boundaries)
+# #     pinn = PINN(problem=poisson_problem,
+# #                 model=model,
+# #                 loss=LpLoss())
+# #     trainer = Trainer(solver=pinn,
+# #                       max_epochs=5,
+# #                       accelerator='cpu',
+# #                       default_root_dir=tmpdir)
+# #     trainer.train()
+# #     ntrainer = Trainer(solver=pinn, max_epochs=5, accelerator='cpu')
+# #     t = ntrainer.train(
+# #         ckpt_path=f'{tmpdir}/lightning_logs/version_0/checkpoints/epoch=4-step=10.ckpt')
+# #     import shutil
+# #     shutil.rmtree(tmpdir)
+
+
+# def test_train_inverse_problem_load():
+#     tmpdir = "tests/tmp_load_inv"
+#     poisson_problem = InversePoisson()
+#     boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4', 'D']
+#     n = 100
+#     poisson_problem.discretise_domain(n, 'random', locations=boundaries)
+#     pinn = PINN(problem=poisson_problem,
+#                 model=model,
+#                 loss=LpLoss())
+#     trainer = Trainer(solver=pinn,
+#                       max_epochs=15,
+#                       accelerator='cpu',
+#                       default_root_dir=tmpdir)
 #     trainer.train()
+#     new_pinn = PINN.load_from_checkpoint(
+#         f'{tmpdir}/lightning_logs/version_0/checkpoints/epoch=14-step=30.ckpt',
+#         problem = poisson_problem, model=model)
+#     test_pts = CartesianDomain({'x': [0, 1], 'y': [0, 1]}).sample(10)
+#     assert new_pinn.forward(test_pts).extract(['u']).shape == (10, 1)
+#     assert new_pinn.forward(test_pts).extract(
+#         ['u']).shape == pinn.forward(test_pts).extract(['u']).shape
+#     torch.testing.assert_close(
+#         new_pinn.forward(test_pts).extract(['u']),
+#         pinn.forward(test_pts).extract(['u']))
+#     import shutil
+#     shutil.rmtree(tmpdir)
 
-# def test_train_2():
-#     boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4']
-#     n = 10
-#     expected_keys = [[], list(range(0, 50, 3))]
-#     param = [0, 3]
-#     for i, truth_key in zip(param, expected_keys):
-#         pinn = PINN(problem, model)
-#         pinn.discretise_domain(n, 'grid', locations=boundaries)
-#         pinn.discretise_domain(n, 'grid', locations=['D'])
-#         pinn.train(50, save_loss=i)
-#         assert list(pinn.history_loss.keys()) == truth_key
-
-
-# def test_train_extra_feats():
-#     pinn = PINN(problem, model_extra_feat, [myFeature()])
-#     boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4']
-#     n = 10
-#     pinn.discretise_domain(n, 'grid', locations=boundaries)
-#     pinn.discretise_domain(n, 'grid', locations=['D'])
-#     pinn.train(5)
+# # # TODO fix asap. Basically sampling few variables
+# # # works only if both variables are in a range.
+# # # if one is fixed and the other not, this will
+# # # not work. This test also needs to be fixed and
+# # # insert in test problem not in test pinn.
+# # def test_train_cpu_sampling_few_vars():
+# #     poisson_problem = Poisson()
+# #     boundaries = ['gamma1', 'gamma2', 'gamma3']
+# #     n = 10
+# #     poisson_problem.discretise_domain(n, 'grid', locations=boundaries)
+# #     poisson_problem.discretise_domain(n, 'random', locations=['gamma4'], variables=['x'])
+# #     poisson_problem.discretise_domain(n, 'random', locations=['gamma4'], variables=['y'])
+# #     pinn = PINN(problem = poisson_problem, model=model, extra_features=None, loss=LpLoss())
+# #     trainer = Trainer(solver=pinn, kwargs={'max_epochs' : 5, 'accelerator':'cpu'})
+# #     trainer.train()
 
 
-# def test_train_2_extra_feats():
-#     boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4']
-#     n = 10
-#     expected_keys = [[], list(range(0, 50, 3))]
-#     param = [0, 3]
-#     for i, truth_key in zip(param, expected_keys):
-#         pinn = PINN(problem, model_extra_feat, [myFeature()])
-#         pinn.discretise_domain(n, 'grid', locations=boundaries)
-#         pinn.discretise_domain(n, 'grid', locations=['D'])
-#         pinn.train(50, save_loss=i)
-#         assert list(pinn.history_loss.keys()) == truth_key
+# # TODO, fix GitHub actions to run also on GPU
+# # def test_train_gpu():
+# #     poisson_problem = Poisson()
+# #     boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4']
+# #     n = 10
+# #     poisson_problem.discretise_domain(n, 'grid', locations=boundaries)
+# #     pinn = PINN(problem = poisson_problem, model=model, extra_features=None, loss=LpLoss())
+# #     trainer = Trainer(solver=pinn, kwargs={'max_epochs' : 5, 'accelerator':'gpu'})
+# #     trainer.train()
 
+# # def test_train_gpu(): #TODO fix ASAP
+# #     poisson_problem = Poisson()
+# #     boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4']
+# #     n = 10
+# #     poisson_problem.discretise_domain(n, 'grid', locations=boundaries)
+# #     poisson_problem.conditions.pop('data') # The input/output pts are allocated on cpu
+# #     pinn = PINN(problem = poisson_problem, model=model, extra_features=None, loss=LpLoss())
+# #     trainer = Trainer(solver=pinn, kwargs={'max_epochs' : 5, 'accelerator':'gpu'})
+# #     trainer.train()
 
-# def test_train_with_optimizer_kwargs():
-#     boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4']
-#     n = 10
-#     expected_keys = [[], list(range(0, 50, 3))]
-#     param = [0, 3]
-#     for i, truth_key in zip(param, expected_keys):
-#         pinn = PINN(problem, model, optimizer_kwargs={'lr' : 0.3})
-#         pinn.discretise_domain(n, 'grid', locations=boundaries)
-#         pinn.discretise_domain(n, 'grid', locations=['D'])
-#         pinn.train(50, save_loss=i)
-#         assert list(pinn.history_loss.keys()) == truth_key
+# # def test_train_2():
+# #     boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4']
+# #     n = 10
+# #     expected_keys = [[], list(range(0, 50, 3))]
+# #     param = [0, 3]
+# #     for i, truth_key in zip(param, expected_keys):
+# #         pinn = PINN(problem, model)
+# #         pinn.discretise_domain(n, 'grid', locations=boundaries)
+# #         pinn.discretise_domain(n, 'grid', locations=['D'])
+# #         pinn.train(50, save_loss=i)
+# #         assert list(pinn.history_loss.keys()) == truth_key
 
 
-# def test_train_with_lr_scheduler():
-#     boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4']
-#     n = 10
-#     expected_keys = [[], list(range(0, 50, 3))]
-#     param = [0, 3]
-#     for i, truth_key in zip(param, expected_keys):
-#         pinn = PINN(
-#             problem,
-#             model,
-#             lr_scheduler_type=torch.optim.lr_scheduler.CyclicLR,
-#             lr_scheduler_kwargs={'base_lr' : 0.1, 'max_lr' : 0.3, 'cycle_momentum': False}
-#         )
-#         pinn.discretise_domain(n, 'grid', locations=boundaries)
-#         pinn.discretise_domain(n, 'grid', locations=['D'])
-#         pinn.train(50, save_loss=i)
-#         assert list(pinn.history_loss.keys()) == truth_key
-
-
-# # def test_train_batch():
-# #     pinn = PINN(problem, model, batch_size=6)
+# # def test_train_extra_feats():
+# #     pinn = PINN(problem, model_extra_feat, [myFeature()])
 # #     boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4']
 # #     n = 10
 # #     pinn.discretise_domain(n, 'grid', locations=boundaries)
@@ -396,34 +343,87 @@ def test_train_inverse_problem_load():
 # #     pinn.train(5)
 
 
-# # def test_train_batch_2():
+# # def test_train_2_extra_feats():
+# #     boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4']
+# #     n = 10
+# #     expected_keys = [[], list(range(0, 50, 3))]
+# #     param = [0, 3]
+# #     for i, truth_key in zip(param, expected_keys):
+# #         pinn = PINN(problem, model_extra_feat, [myFeature()])
+# #         pinn.discretise_domain(n, 'grid', locations=boundaries)
+# #         pinn.discretise_domain(n, 'grid', locations=['D'])
+# #         pinn.train(50, save_loss=i)
+# #         assert list(pinn.history_loss.keys()) == truth_key
+
+
+# # def test_train_with_optimizer_kwargs():
 # #     boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4']
 # #     n = 10
 # #     expected_keys = [[], list(range(0, 50, 3))]
 # #     param = [0, 3]
 # #     for i, truth_key in zip(param, expected_keys):
-# #         pinn = PINN(problem, model, batch_size=6)
+# #         pinn = PINN(problem, model, optimizer_kwargs={'lr' : 0.3})
 # #         pinn.discretise_domain(n, 'grid', locations=boundaries)
 # #         pinn.discretise_domain(n, 'grid', locations=['D'])
 # #         pinn.train(50, save_loss=i)
 # #         assert list(pinn.history_loss.keys()) == truth_key
 
 
-# if torch.cuda.is_available():
+# # def test_train_with_lr_scheduler():
+# #     boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4']
+# #     n = 10
+# #     expected_keys = [[], list(range(0, 50, 3))]
+# #     param = [0, 3]
+# #     for i, truth_key in zip(param, expected_keys):
+# #         pinn = PINN(
+# #             problem,
+# #             model,
+# #             lr_scheduler_type=torch.optim.lr_scheduler.CyclicLR,
+# #             lr_scheduler_kwargs={'base_lr' : 0.1, 'max_lr' : 0.3, 'cycle_momentum': False}
+# #         )
+# #         pinn.discretise_domain(n, 'grid', locations=boundaries)
+# #         pinn.discretise_domain(n, 'grid', locations=['D'])
+# #         pinn.train(50, save_loss=i)
+# #         assert list(pinn.history_loss.keys()) == truth_key
 
-#     # def test_gpu_train():
-#     #     pinn = PINN(problem, model, batch_size=20, device='cuda')
-#     #     boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4']
-#     #     n = 100
-#     #     pinn.discretise_domain(n, 'grid', locations=boundaries)
-#     #     pinn.discretise_domain(n, 'grid', locations=['D'])
-#     #     pinn.train(5)
 
-#     def test_gpu_train_nobatch():
-#         pinn = PINN(problem, model, batch_size=None, device='cuda')
-#         boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4']
-#         n = 100
-#         pinn.discretise_domain(n, 'grid', locations=boundaries)
-#         pinn.discretise_domain(n, 'grid', locations=['D'])
-#         pinn.train(5)
+# # # def test_train_batch():
+# # #     pinn = PINN(problem, model, batch_size=6)
+# # #     boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4']
+# # #     n = 10
+# # #     pinn.discretise_domain(n, 'grid', locations=boundaries)
+# # #     pinn.discretise_domain(n, 'grid', locations=['D'])
+# # #     pinn.train(5)
+
+
+# # # def test_train_batch_2():
+# # #     boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4']
+# # #     n = 10
+# # #     expected_keys = [[], list(range(0, 50, 3))]
+# # #     param = [0, 3]
+# # #     for i, truth_key in zip(param, expected_keys):
+# # #         pinn = PINN(problem, model, batch_size=6)
+# # #         pinn.discretise_domain(n, 'grid', locations=boundaries)
+# # #         pinn.discretise_domain(n, 'grid', locations=['D'])
+# # #         pinn.train(50, save_loss=i)
+# # #         assert list(pinn.history_loss.keys()) == truth_key
+
+
+# # if torch.cuda.is_available():
+
+# #     # def test_gpu_train():
+# #     #     pinn = PINN(problem, model, batch_size=20, device='cuda')
+# #     #     boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4']
+# #     #     n = 100
+# #     #     pinn.discretise_domain(n, 'grid', locations=boundaries)
+# #     #     pinn.discretise_domain(n, 'grid', locations=['D'])
+# #     #     pinn.train(5)
+
+# #     def test_gpu_train_nobatch():
+# #         pinn = PINN(problem, model, batch_size=None, device='cuda')
+# #         boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4']
+# #         n = 100
+# #         pinn.discretise_domain(n, 'grid', locations=boundaries)
+# #         pinn.discretise_domain(n, 'grid', locations=['D'])
+# #         pinn.train(5)
 
diff --git a/tests/test_solvers/test_garom.py b/tests/test_solvers/test_garom.py
index 4ff4e1c9d..e569e1bb7 100644
--- a/tests/test_solvers/test_garom.py
+++ b/tests/test_solvers/test_garom.py
@@ -8,160 +8,160 @@
 import matplotlib.tri as tri
 
 
-def func(x, mu1, mu2):
-    import torch
-    x_m1 = (x[:, 0] - mu1).pow(2)
-    x_m2 = (x[:, 1] - mu2).pow(2)
-    norm = x[:, 0]**2 + x[:, 1]**2
-    return torch.exp(-(x_m1 + x_m2))
-
-
-class ParametricGaussian(AbstractProblem):
-    output_variables = [f'u_{i}' for i in range(900)]
-
-    # params
-    xx = torch.linspace(-1, 1, 20)
-    yy = xx
-    params = LabelTensor(torch.cartesian_prod(xx, yy), labels=['mu1', 'mu2'])
-
-    # define domain
-    x = torch.linspace(-1, 1, 30)
-    domain = torch.cartesian_prod(x, x)
-    triang = tri.Triangulation(domain[:, 0], domain[:, 1])
-    sol = []
-    for p in params:
-        sol.append(func(domain, p[0], p[1]))
-    snapshots = LabelTensor(torch.stack(sol), labels=output_variables)
-
-    # define conditions
-    conditions = {
-        'data': Condition(input_points=params, output_points=snapshots)
-    }
-
-
-# simple Generator Network
-class Generator(nn.Module):
-
-    def __init__(self,
-                 input_dimension,
-                 parameters_dimension,
-                 noise_dimension,
-                 activation=torch.nn.SiLU):
-        super().__init__()
-
-        self._noise_dimension = noise_dimension
-        self._activation = activation
-
-        self.model = torch.nn.Sequential(
-            torch.nn.Linear(6 * self._noise_dimension, input_dimension // 6),
-            self._activation(),
-            torch.nn.Linear(input_dimension // 6, input_dimension // 3),
-            self._activation(),
-            torch.nn.Linear(input_dimension // 3, input_dimension))
-        self.condition = torch.nn.Sequential(
-            torch.nn.Linear(parameters_dimension, 2 * self._noise_dimension),
-            self._activation(),
-            torch.nn.Linear(2 * self._noise_dimension,
-                            5 * self._noise_dimension))
-
-    def forward(self, param):
-        # uniform sampling in [-1, 1]
-        z = torch.rand(size=(param.shape[0], self._noise_dimension),
-                       device=param.device,
-                       dtype=param.dtype,
-                       requires_grad=True)
-        z = 2. * z - 1.
-
-        # conditioning by concatenation of mapped parameters
-        input_ = torch.cat((z, self.condition(param)), dim=-1)
-        out = self.model(input_)
-
-        return out
-
-
-# Simple Discriminator Network
-class Discriminator(nn.Module):
-
-    def __init__(self,
-                 input_dimension,
-                 parameter_dimension,
-                 hidden_dimension,
-                 activation=torch.nn.ReLU):
-        super().__init__()
-
-        self._activation = activation
-        self.encoding = torch.nn.Sequential(
-            torch.nn.Linear(input_dimension, input_dimension // 3),
-            self._activation(),
-            torch.nn.Linear(input_dimension // 3, input_dimension // 6),
-            self._activation(),
-            torch.nn.Linear(input_dimension // 6, hidden_dimension))
-        self.decoding = torch.nn.Sequential(
-            torch.nn.Linear(2 * hidden_dimension, input_dimension // 6),
-            self._activation(),
-            torch.nn.Linear(input_dimension // 6, input_dimension // 3),
-            self._activation(),
-            torch.nn.Linear(input_dimension // 3, input_dimension),
-        )
-
-        self.condition = torch.nn.Sequential(
-            torch.nn.Linear(parameter_dimension, hidden_dimension // 2),
-            self._activation(),
-            torch.nn.Linear(hidden_dimension // 2, hidden_dimension))
-
-    def forward(self, data):
-        x, condition = data
-        encoding = self.encoding(x)
-        conditioning = torch.cat((encoding, self.condition(condition)), dim=-1)
-        decoding = self.decoding(conditioning)
-        return decoding
-
-
-problem = ParametricGaussian()
-
-
-def test_constructor():
-    GAROM(problem=problem,
-          generator=Generator(input_dimension=900,
-                              parameters_dimension=2,
-                              noise_dimension=12),
-          discriminator=Discriminator(input_dimension=900,
-                                      parameter_dimension=2,
-                                      hidden_dimension=64))
-
-
-def test_train_cpu():
-    solver = GAROM(problem=problem,
-                   generator=Generator(input_dimension=900,
-                                       parameters_dimension=2,
-                                       noise_dimension=12),
-                   discriminator=Discriminator(input_dimension=900,
-                                               parameter_dimension=2,
-                                               hidden_dimension=64))
-
-    trainer = Trainer(solver=solver, max_epochs=4, accelerator='cpu', batch_size=20)
-    trainer.train()
-
-
-def test_sample():
-    solver = GAROM(problem=problem,
-                   generator=Generator(input_dimension=900,
-                                       parameters_dimension=2,
-                                       noise_dimension=12),
-                   discriminator=Discriminator(input_dimension=900,
-                                               parameter_dimension=2,
-                                               hidden_dimension=64))
-    solver.sample(problem.params)
-    assert solver.sample(problem.params).shape == problem.snapshots.shape
-
-
-def test_forward():
-    solver = GAROM(problem=problem,
-                   generator=Generator(input_dimension=900,
-                                       parameters_dimension=2,
-                                       noise_dimension=12),
-                   discriminator=Discriminator(input_dimension=900,
-                                               parameter_dimension=2,
-                                               hidden_dimension=64))
-    solver(problem.params, mc_steps=100, variance=True)
-    assert solver(problem.params).shape == problem.snapshots.shape
+# def func(x, mu1, mu2):
+#     import torch
+#     x_m1 = (x[:, 0] - mu1).pow(2)
+#     x_m2 = (x[:, 1] - mu2).pow(2)
+#     norm = x[:, 0]**2 + x[:, 1]**2
+#     return torch.exp(-(x_m1 + x_m2))
+
+
+# class ParametricGaussian(AbstractProblem):
+#     output_variables = [f'u_{i}' for i in range(900)]
+
+#     # params
+#     xx = torch.linspace(-1, 1, 20)
+#     yy = xx
+#     params = LabelTensor(torch.cartesian_prod(xx, yy), labels=['mu1', 'mu2'])
+
+#     # define domain
+#     x = torch.linspace(-1, 1, 30)
+#     domain = torch.cartesian_prod(x, x)
+#     triang = tri.Triangulation(domain[:, 0], domain[:, 1])
+#     sol = []
+#     for p in params:
+#         sol.append(func(domain, p[0], p[1]))
+#     snapshots = LabelTensor(torch.stack(sol), labels=output_variables)
+
+#     # define conditions
+#     conditions = {
+#         'data': Condition(input_points=params, output_points=snapshots)
+#     }
+
+
+# # simple Generator Network
+# class Generator(nn.Module):
+
+#     def __init__(self,
+#                  input_dimension,
+#                  parameters_dimension,
+#                  noise_dimension,
+#                  activation=torch.nn.SiLU):
+#         super().__init__()
+
+#         self._noise_dimension = noise_dimension
+#         self._activation = activation
+
+#         self.model = torch.nn.Sequential(
+#             torch.nn.Linear(6 * self._noise_dimension, input_dimension // 6),
+#             self._activation(),
+#             torch.nn.Linear(input_dimension // 6, input_dimension // 3),
+#             self._activation(),
+#             torch.nn.Linear(input_dimension // 3, input_dimension))
+#         self.condition = torch.nn.Sequential(
+#             torch.nn.Linear(parameters_dimension, 2 * self._noise_dimension),
+#             self._activation(),
+#             torch.nn.Linear(2 * self._noise_dimension,
+#                             5 * self._noise_dimension))
+
+#     def forward(self, param):
+#         # uniform sampling in [-1, 1]
+#         z = torch.rand(size=(param.shape[0], self._noise_dimension),
+#                        device=param.device,
+#                        dtype=param.dtype,
+#                        requires_grad=True)
+#         z = 2. * z - 1.
+
+#         # conditioning by concatenation of mapped parameters
+#         input_ = torch.cat((z, self.condition(param)), dim=-1)
+#         out = self.model(input_)
+
+#         return out
+
+
+# # Simple Discriminator Network
+# class Discriminator(nn.Module):
+
+#     def __init__(self,
+#                  input_dimension,
+#                  parameter_dimension,
+#                  hidden_dimension,
+#                  activation=torch.nn.ReLU):
+#         super().__init__()
+
+#         self._activation = activation
+#         self.encoding = torch.nn.Sequential(
+#             torch.nn.Linear(input_dimension, input_dimension // 3),
+#             self._activation(),
+#             torch.nn.Linear(input_dimension // 3, input_dimension // 6),
+#             self._activation(),
+#             torch.nn.Linear(input_dimension // 6, hidden_dimension))
+#         self.decoding = torch.nn.Sequential(
+#             torch.nn.Linear(2 * hidden_dimension, input_dimension // 6),
+#             self._activation(),
+#             torch.nn.Linear(input_dimension // 6, input_dimension // 3),
+#             self._activation(),
+#             torch.nn.Linear(input_dimension // 3, input_dimension),
+#         )
+
+#         self.condition = torch.nn.Sequential(
+#             torch.nn.Linear(parameter_dimension, hidden_dimension // 2),
+#             self._activation(),
+#             torch.nn.Linear(hidden_dimension // 2, hidden_dimension))
+
+#     def forward(self, data):
+#         x, condition = data
+#         encoding = self.encoding(x)
+#         conditioning = torch.cat((encoding, self.condition(condition)), dim=-1)
+#         decoding = self.decoding(conditioning)
+#         return decoding
+
+
+# problem = ParametricGaussian()
+
+
+# def test_constructor():
+#     GAROM(problem=problem,
+#           generator=Generator(input_dimension=900,
+#                               parameters_dimension=2,
+#                               noise_dimension=12),
+#           discriminator=Discriminator(input_dimension=900,
+#                                       parameter_dimension=2,
+#                                       hidden_dimension=64))
+
+
+# def test_train_cpu():
+#     solver = GAROM(problem=problem,
+#                    generator=Generator(input_dimension=900,
+#                                        parameters_dimension=2,
+#                                        noise_dimension=12),
+#                    discriminator=Discriminator(input_dimension=900,
+#                                                parameter_dimension=2,
+#                                                hidden_dimension=64))
+
+#     trainer = Trainer(solver=solver, max_epochs=4, accelerator='cpu', batch_size=20)
+#     trainer.train()
+
+
+# def test_sample():
+#     solver = GAROM(problem=problem,
+#                    generator=Generator(input_dimension=900,
+#                                        parameters_dimension=2,
+#                                        noise_dimension=12),
+#                    discriminator=Discriminator(input_dimension=900,
+#                                                parameter_dimension=2,
+#                                                hidden_dimension=64))
+#     solver.sample(problem.params)
+#     assert solver.sample(problem.params).shape == problem.snapshots.shape
+
+
+# def test_forward():
+#     solver = GAROM(problem=problem,
+#                    generator=Generator(input_dimension=900,
+#                                        parameters_dimension=2,
+#                                        noise_dimension=12),
+#                    discriminator=Discriminator(input_dimension=900,
+#                                                parameter_dimension=2,
+#                                                hidden_dimension=64))
+#     solver(problem.params, mc_steps=100, variance=True)
+#     assert solver(problem.params).shape == problem.snapshots.shape
diff --git a/tests/test_solvers/test_gpinn.py b/tests/test_solvers/test_gpinn.py
index 47e38e458..2bcda878d 100644
--- a/tests/test_solvers/test_gpinn.py
+++ b/tests/test_solvers/test_gpinn.py
@@ -7,242 +7,178 @@
 from pina.solvers import GPINN
 from pina.trainer import Trainer
 from pina.model import FeedForward
-from pina.equation.equation import Equation
+from pina.equation import Equation
 from pina.equation.equation_factory import FixedValue
-from pina.loss.loss_interface import LpLoss
-
-
-def laplace_equation(input_, output_):
-    force_term = (torch.sin(input_.extract(['x']) * torch.pi) *
-                  torch.sin(input_.extract(['y']) * torch.pi))
-    delta_u = laplacian(output_.extract(['u']), input_)
-    return delta_u - force_term
-
-
-my_laplace = Equation(laplace_equation)
-in_ = LabelTensor(torch.tensor([[0., 1.]]), ['x', 'y'])
-out_ = LabelTensor(torch.tensor([[0.]]), ['u'])
-in2_ = LabelTensor(torch.rand(60, 2), ['x', 'y'])
-out2_ = LabelTensor(torch.rand(60, 1), ['u'])
-
-
-class InversePoisson(SpatialProblem, InverseProblem):
-    '''
-    Problem definition for the Poisson equation.
-    '''
-    output_variables = ['u']
-    x_min = -2
-    x_max = 2
-    y_min = -2
-    y_max = 2
-    data_input = LabelTensor(torch.rand(10, 2), ['x', 'y'])
-    data_output = LabelTensor(torch.rand(10, 1), ['u'])
-    spatial_domain = CartesianDomain({'x': [x_min, x_max], 'y': [y_min, y_max]})
-    # define the ranges for the parameters
-    unknown_parameter_domain = CartesianDomain({'mu1': [-1, 1], 'mu2': [-1, 1]})
-
-    def laplace_equation(input_, output_, params_):
-        '''
-        Laplace equation with a force term.
-        '''
-        force_term = torch.exp(
-                - 2*(input_.extract(['x']) - params_['mu1'])**2
-                - 2*(input_.extract(['y']) - params_['mu2'])**2)
-        delta_u = laplacian(output_, input_, components=['u'], d=['x', 'y'])
-
-        return delta_u - force_term
-
-    # define the conditions for the loss (boundary conditions, equation, data)
-    conditions = {
-        'gamma1': Condition(location=CartesianDomain({'x': [x_min, x_max],
-            'y':  y_max}),
-            equation=FixedValue(0.0, components=['u'])),
-        'gamma2': Condition(location=CartesianDomain(
-            {'x': [x_min, x_max], 'y': y_min}),
-            equation=FixedValue(0.0, components=['u'])),
-        'gamma3': Condition(location=CartesianDomain(
-            {'x':  x_max, 'y': [y_min, y_max]}),
-            equation=FixedValue(0.0, components=['u'])),
-        'gamma4': Condition(location=CartesianDomain(
-            {'x': x_min, 'y': [y_min, y_max]
-            }),
-            equation=FixedValue(0.0, components=['u'])),
-        'D': Condition(location=CartesianDomain(
-            {'x': [x_min, x_max], 'y': [y_min, y_max]
-            }),
-        equation=Equation(laplace_equation)),
-        'data': Condition(
-            input_points=data_input.extract(['x', 'y']),
-            output_points=data_output)
-    }
-
-
-class Poisson(SpatialProblem):
-    output_variables = ['u']
-    spatial_domain = CartesianDomain({'x': [0, 1], 'y': [0, 1]})
-
-    conditions = {
-        'gamma1': Condition(
-            location=CartesianDomain({'x': [0, 1], 'y':  1}),
-            equation=FixedValue(0.0)),
-        'gamma2': Condition(
-            location=CartesianDomain({'x': [0, 1], 'y': 0}),
-            equation=FixedValue(0.0)),
-        'gamma3': Condition(
-            location=CartesianDomain({'x':  1, 'y': [0, 1]}),
-            equation=FixedValue(0.0)),
-        'gamma4': Condition(
-            location=CartesianDomain({'x': 0, 'y': [0, 1]}),
-            equation=FixedValue(0.0)),
-        'D': Condition(
-            input_points=LabelTensor(torch.rand(size=(100, 2)), ['x', 'y']),
-            equation=my_laplace),
-        'data': Condition(
-            input_points=in_,
-            output_points=out_),
-        'data2': Condition(
-            input_points=in2_,
-            output_points=out2_)
-    }
-
-    def poisson_sol(self, pts):
-        return -(torch.sin(pts.extract(['x']) * torch.pi) *
-                 torch.sin(pts.extract(['y']) * torch.pi)) / (2 * torch.pi**2)
-
-    truth_solution = poisson_sol
-
-
-class myFeature(torch.nn.Module):
-    """
-    Feature: sin(x)
-    """
-
-    def __init__(self):
-        super(myFeature, self).__init__()
-
-    def forward(self, x):
-        t = (torch.sin(x.extract(['x']) * torch.pi) *
-             torch.sin(x.extract(['y']) * torch.pi))
-        return LabelTensor(t, ['sin(x)sin(y)'])
-
-
-# make the problem
-poisson_problem = Poisson()
-model = FeedForward(len(poisson_problem.input_variables),
-                    len(poisson_problem.output_variables))
-model_extra_feats = FeedForward(
-    len(poisson_problem.input_variables) + 1,
-    len(poisson_problem.output_variables))
-extra_feats = [myFeature()]
-
-
-def test_constructor():
-    GPINN(problem=poisson_problem, model=model, extra_features=None)
-
-
-def test_constructor_extra_feats():
-    model_extra_feats = FeedForward(
-        len(poisson_problem.input_variables) + 1,
-        len(poisson_problem.output_variables))
-    GPINN(problem=poisson_problem,
-         model=model_extra_feats,
-         extra_features=extra_feats)
-
-
-def test_train_cpu():
-    poisson_problem = Poisson()
-    boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4']
-    n = 10
-    poisson_problem.discretise_domain(n, 'grid', locations=boundaries)
-    pinn = GPINN(problem = poisson_problem,
-                 model=model, extra_features=None, loss=LpLoss())
-    trainer = Trainer(solver=pinn, max_epochs=1,
-                      accelerator='cpu', batch_size=20)
-    trainer.train()
-
-def test_log():
-    poisson_problem.discretise_domain(100)
-    solver = GPINN(problem = poisson_problem, model=model,
-                extra_features=None, loss=LpLoss())
-    trainer = Trainer(solver, max_epochs=2, accelerator='cpu')
-    trainer.train()
-    # assert the logged metrics are correct
-    logged_metrics = sorted(list(trainer.logged_metrics.keys()))
-    total_metrics = sorted(
-        list([key + '_loss' for key in poisson_problem.conditions.keys()])
-        + ['mean_loss'])
-    assert logged_metrics == total_metrics
-
-def test_train_restore():
-    tmpdir = "tests/tmp_restore"
-    poisson_problem = Poisson()
-    boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4']
-    n = 10
-    poisson_problem.discretise_domain(n, 'grid', locations=boundaries)
-    pinn = GPINN(problem=poisson_problem,
-                model=model,
-                extra_features=None,
-                loss=LpLoss())
-    trainer = Trainer(solver=pinn,
-                      max_epochs=5,
-                      accelerator='cpu',
-                      default_root_dir=tmpdir)
-    trainer.train()
-    ntrainer = Trainer(solver=pinn, max_epochs=15, accelerator='cpu')
-    t = ntrainer.train(
-        ckpt_path=f'{tmpdir}/lightning_logs/version_0/'
-        'checkpoints/epoch=4-step=10.ckpt')
-    import shutil
-    shutil.rmtree(tmpdir)
-
-
-def test_train_load():
-    tmpdir = "tests/tmp_load"
-    poisson_problem = Poisson()
-    boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4']
-    n = 10
-    poisson_problem.discretise_domain(n, 'grid', locations=boundaries)
-    pinn = GPINN(problem=poisson_problem,
-                model=model,
-                extra_features=None,
-                loss=LpLoss())
-    trainer = Trainer(solver=pinn,
-                      max_epochs=15,
-                      accelerator='cpu',
-                      default_root_dir=tmpdir)
-    trainer.train()
-    new_pinn = GPINN.load_from_checkpoint(
-        f'{tmpdir}/lightning_logs/version_0/checkpoints/epoch=14-step=30.ckpt',
-        problem = poisson_problem, model=model)
-    test_pts = CartesianDomain({'x': [0, 1], 'y': [0, 1]}).sample(10)
-    assert new_pinn.forward(test_pts).extract(['u']).shape == (10, 1)
-    assert new_pinn.forward(test_pts).extract(
-        ['u']).shape == pinn.forward(test_pts).extract(['u']).shape
-    torch.testing.assert_close(
-        new_pinn.forward(test_pts).extract(['u']),
-        pinn.forward(test_pts).extract(['u']))
-    import shutil
-    shutil.rmtree(tmpdir)
-
-def test_train_inverse_problem_cpu():
-    poisson_problem = InversePoisson()
-    boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4', 'D']
-    n = 100
-    poisson_problem.discretise_domain(n, 'random', locations=boundaries)
-    pinn = GPINN(problem = poisson_problem,
-                 model=model, extra_features=None, loss=LpLoss())
-    trainer = Trainer(solver=pinn, max_epochs=1,
-                      accelerator='cpu', batch_size=20)
-    trainer.train()
-
-
-# # TODO does not currently work
-# def test_train_inverse_problem_restore():
-#     tmpdir = "tests/tmp_restore_inv"
-#     poisson_problem = InversePoisson()
-#     boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4', 'D']
-#     n = 100
-#     poisson_problem.discretise_domain(n, 'random', locations=boundaries)
+from pina.loss import LpLoss
+
+
+# def laplace_equation(input_, output_):
+#     force_term = (torch.sin(input_.extract(['x']) * torch.pi) *
+#                   torch.sin(input_.extract(['y']) * torch.pi))
+#     delta_u = laplacian(output_.extract(['u']), input_)
+#     return delta_u - force_term
+
+
+# my_laplace = Equation(laplace_equation)
+# in_ = LabelTensor(torch.tensor([[0., 1.]]), ['x', 'y'])
+# out_ = LabelTensor(torch.tensor([[0.]]), ['u'])
+# in2_ = LabelTensor(torch.rand(60, 2), ['x', 'y'])
+# out2_ = LabelTensor(torch.rand(60, 1), ['u'])
+
+
+# class InversePoisson(SpatialProblem, InverseProblem):
+#     '''
+#     Problem definition for the Poisson equation.
+#     '''
+#     output_variables = ['u']
+#     x_min = -2
+#     x_max = 2
+#     y_min = -2
+#     y_max = 2
+#     data_input = LabelTensor(torch.rand(10, 2), ['x', 'y'])
+#     data_output = LabelTensor(torch.rand(10, 1), ['u'])
+#     spatial_domain = CartesianDomain({'x': [x_min, x_max], 'y': [y_min, y_max]})
+#     # define the ranges for the parameters
+#     unknown_parameter_domain = CartesianDomain({'mu1': [-1, 1], 'mu2': [-1, 1]})
+
+#     def laplace_equation(input_, output_, params_):
+#         '''
+#         Laplace equation with a force term.
+#         '''
+#         force_term = torch.exp(
+#                 - 2*(input_.extract(['x']) - params_['mu1'])**2
+#                 - 2*(input_.extract(['y']) - params_['mu2'])**2)
+#         delta_u = laplacian(output_, input_, components=['u'], d=['x', 'y'])
+
+#         return delta_u - force_term
+
+#     # define the conditions for the loss (boundary conditions, equation, data)
+#     conditions = {
+#         'gamma1': Condition(location=CartesianDomain({'x': [x_min, x_max],
+#             'y':  y_max}),
+#             equation=FixedValue(0.0, components=['u'])),
+#         'gamma2': Condition(location=CartesianDomain(
+#             {'x': [x_min, x_max], 'y': y_min}),
+#             equation=FixedValue(0.0, components=['u'])),
+#         'gamma3': Condition(location=CartesianDomain(
+#             {'x':  x_max, 'y': [y_min, y_max]}),
+#             equation=FixedValue(0.0, components=['u'])),
+#         'gamma4': Condition(location=CartesianDomain(
+#             {'x': x_min, 'y': [y_min, y_max]
+#             }),
+#             equation=FixedValue(0.0, components=['u'])),
+#         'D': Condition(location=CartesianDomain(
+#             {'x': [x_min, x_max], 'y': [y_min, y_max]
+#             }),
+#         equation=Equation(laplace_equation)),
+#         'data': Condition(
+#             input_points=data_input.extract(['x', 'y']),
+#             output_points=data_output)
+#     }
+
+
+# class Poisson(SpatialProblem):
+#     output_variables = ['u']
+#     spatial_domain = CartesianDomain({'x': [0, 1], 'y': [0, 1]})
+
+#     conditions = {
+#         'gamma1': Condition(
+#             location=CartesianDomain({'x': [0, 1], 'y':  1}),
+#             equation=FixedValue(0.0)),
+#         'gamma2': Condition(
+#             location=CartesianDomain({'x': [0, 1], 'y': 0}),
+#             equation=FixedValue(0.0)),
+#         'gamma3': Condition(
+#             location=CartesianDomain({'x':  1, 'y': [0, 1]}),
+#             equation=FixedValue(0.0)),
+#         'gamma4': Condition(
+#             location=CartesianDomain({'x': 0, 'y': [0, 1]}),
+#             equation=FixedValue(0.0)),
+#         'D': Condition(
+#             input_points=LabelTensor(torch.rand(size=(100, 2)), ['x', 'y']),
+#             equation=my_laplace),
+#         'data': Condition(
+#             input_points=in_,
+#             output_points=out_),
+#         'data2': Condition(
+#             input_points=in2_,
+#             output_points=out2_)
+#     }
+
+#     def poisson_sol(self, pts):
+#         return -(torch.sin(pts.extract(['x']) * torch.pi) *
+#                  torch.sin(pts.extract(['y']) * torch.pi)) / (2 * torch.pi**2)
+
+#     truth_solution = poisson_sol
+
+
+# class myFeature(torch.nn.Module):
+#     """
+#     Feature: sin(x)
+#     """
+
+#     def __init__(self):
+#         super(myFeature, self).__init__()
+
+#     def forward(self, x):
+#         t = (torch.sin(x.extract(['x']) * torch.pi) *
+#              torch.sin(x.extract(['y']) * torch.pi))
+#         return LabelTensor(t, ['sin(x)sin(y)'])
+
+
+# # make the problem
+# poisson_problem = Poisson()
+# model = FeedForward(len(poisson_problem.input_variables),
+#                     len(poisson_problem.output_variables))
+# model_extra_feats = FeedForward(
+#     len(poisson_problem.input_variables) + 1,
+#     len(poisson_problem.output_variables))
+# extra_feats = [myFeature()]
+
+
+# def test_constructor():
+#     GPINN(problem=poisson_problem, model=model, extra_features=None)
+
+
+# def test_constructor_extra_feats():
+#     model_extra_feats = FeedForward(
+#         len(poisson_problem.input_variables) + 1,
+#         len(poisson_problem.output_variables))
+#     GPINN(problem=poisson_problem,
+#          model=model_extra_feats,
+#          extra_features=extra_feats)
+
+
+# def test_train_cpu():
+#     poisson_problem = Poisson()
+#     boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4']
+#     n = 10
+#     poisson_problem.discretise_domain(n, 'grid', locations=boundaries)
+#     pinn = GPINN(problem = poisson_problem,
+#                  model=model, extra_features=None, loss=LpLoss())
+#     trainer = Trainer(solver=pinn, max_epochs=1,
+#                       accelerator='cpu', batch_size=20)
+#     trainer.train()
+
+# def test_log():
+#     poisson_problem.discretise_domain(100)
+#     solver = GPINN(problem = poisson_problem, model=model,
+#                 extra_features=None, loss=LpLoss())
+#     trainer = Trainer(solver, max_epochs=2, accelerator='cpu')
+#     trainer.train()
+#     # assert the logged metrics are correct
+#     logged_metrics = sorted(list(trainer.logged_metrics.keys()))
+#     total_metrics = sorted(
+#         list([key + '_loss' for key in poisson_problem.conditions.keys()])
+#         + ['mean_loss'])
+#     assert logged_metrics == total_metrics
+
+# def test_train_restore():
+#     tmpdir = "tests/tmp_restore"
+#     poisson_problem = Poisson()
+#     boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4']
+#     n = 10
+#     poisson_problem.discretise_domain(n, 'grid', locations=boundaries)
 #     pinn = GPINN(problem=poisson_problem,
 #                 model=model,
 #                 extra_features=None,
@@ -252,158 +188,169 @@ def test_train_inverse_problem_cpu():
 #                       accelerator='cpu',
 #                       default_root_dir=tmpdir)
 #     trainer.train()
-#     ntrainer = Trainer(solver=pinn, max_epochs=5, accelerator='cpu')
+#     ntrainer = Trainer(solver=pinn, max_epochs=15, accelerator='cpu')
 #     t = ntrainer.train(
-#         ckpt_path=f'{tmpdir}/lightning_logs/version_0/checkpoints/epoch=4-step=10.ckpt')
+#         ckpt_path=f'{tmpdir}/lightning_logs/version_0/'
+#         'checkpoints/epoch=4-step=10.ckpt')
 #     import shutil
 #     shutil.rmtree(tmpdir)
 
 
-def test_train_inverse_problem_load():
-    tmpdir = "tests/tmp_load_inv"
-    poisson_problem = InversePoisson()
-    boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4', 'D']
-    n = 100
-    poisson_problem.discretise_domain(n, 'random', locations=boundaries)
-    pinn = GPINN(problem=poisson_problem,
-                model=model,
-                extra_features=None,
-                loss=LpLoss())
-    trainer = Trainer(solver=pinn,
-                      max_epochs=15,
-                      accelerator='cpu',
-                      default_root_dir=tmpdir)
-    trainer.train()
-    new_pinn = GPINN.load_from_checkpoint(
-        f'{tmpdir}/lightning_logs/version_0/checkpoints/epoch=14-step=30.ckpt',
-        problem = poisson_problem, model=model)
-    test_pts = CartesianDomain({'x': [0, 1], 'y': [0, 1]}).sample(10)
-    assert new_pinn.forward(test_pts).extract(['u']).shape == (10, 1)
-    assert new_pinn.forward(test_pts).extract(
-        ['u']).shape == pinn.forward(test_pts).extract(['u']).shape
-    torch.testing.assert_close(
-        new_pinn.forward(test_pts).extract(['u']),
-        pinn.forward(test_pts).extract(['u']))
-    import shutil
-    shutil.rmtree(tmpdir)
-
-# # TODO fix asap. Basically sampling few variables
-# # works only if both variables are in a range.
-# # if one is fixed and the other not, this will
-# # not work. This test also needs to be fixed and
-# # insert in test problem not in test pinn.
-# def test_train_cpu_sampling_few_vars():
+# def test_train_load():
+#     tmpdir = "tests/tmp_load"
 #     poisson_problem = Poisson()
-#     boundaries = ['gamma1', 'gamma2', 'gamma3']
+#     boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4']
 #     n = 10
 #     poisson_problem.discretise_domain(n, 'grid', locations=boundaries)
-#     poisson_problem.discretise_domain(n, 'random', locations=['gamma4'], variables=['x'])
-#     poisson_problem.discretise_domain(n, 'random', locations=['gamma4'], variables=['y'])
-#     pinn = GPINN(problem = poisson_problem, model=model, extra_features=None, loss=LpLoss())
-#     trainer = Trainer(solver=pinn, kwargs={'max_epochs' : 5, 'accelerator':'cpu'})
+#     pinn = GPINN(problem=poisson_problem,
+#                 model=model,
+#                 extra_features=None,
+#                 loss=LpLoss())
+#     trainer = Trainer(solver=pinn,
+#                       max_epochs=15,
+#                       accelerator='cpu',
+#                       default_root_dir=tmpdir)
 #     trainer.train()
+#     new_pinn = GPINN.load_from_checkpoint(
+#         f'{tmpdir}/lightning_logs/version_0/checkpoints/epoch=14-step=30.ckpt',
+#         problem = poisson_problem, model=model)
+#     test_pts = CartesianDomain({'x': [0, 1], 'y': [0, 1]}).sample(10)
+#     assert new_pinn.forward(test_pts).extract(['u']).shape == (10, 1)
+#     assert new_pinn.forward(test_pts).extract(
+#         ['u']).shape == pinn.forward(test_pts).extract(['u']).shape
+#     torch.testing.assert_close(
+#         new_pinn.forward(test_pts).extract(['u']),
+#         pinn.forward(test_pts).extract(['u']))
+#     import shutil
+#     shutil.rmtree(tmpdir)
 
-
-def test_train_extra_feats_cpu():
-    poisson_problem = Poisson()
-    boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4']
-    n = 10
-    poisson_problem.discretise_domain(n, 'grid', locations=boundaries)
-    pinn = GPINN(problem=poisson_problem,
-                model=model_extra_feats,
-                extra_features=extra_feats)
-    trainer = Trainer(solver=pinn, max_epochs=5, accelerator='cpu')
-    trainer.train()
+# def test_train_inverse_problem_cpu():
+#     poisson_problem = InversePoisson()
+#     boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4', 'D']
+#     n = 100
+#     poisson_problem.discretise_domain(n, 'random', locations=boundaries)
+#     pinn = GPINN(problem = poisson_problem,
+#                  model=model, extra_features=None, loss=LpLoss())
+#     trainer = Trainer(solver=pinn, max_epochs=1,
+#                       accelerator='cpu', batch_size=20)
+#     trainer.train()
 
 
-# TODO, fix GitHub actions to run also on GPU
-# def test_train_gpu():
-#     poisson_problem = Poisson()
-#     boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4']
-#     n = 10
-#     poisson_problem.discretise_domain(n, 'grid', locations=boundaries)
-#     pinn = GPINN(problem = poisson_problem, model=model, extra_features=None, loss=LpLoss())
-#     trainer = Trainer(solver=pinn, kwargs={'max_epochs' : 5, 'accelerator':'gpu'})
+# # # TODO does not currently work
+# # def test_train_inverse_problem_restore():
+# #     tmpdir = "tests/tmp_restore_inv"
+# #     poisson_problem = InversePoisson()
+# #     boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4', 'D']
+# #     n = 100
+# #     poisson_problem.discretise_domain(n, 'random', locations=boundaries)
+# #     pinn = GPINN(problem=poisson_problem,
+# #                 model=model,
+# #                 extra_features=None,
+# #                 loss=LpLoss())
+# #     trainer = Trainer(solver=pinn,
+# #                       max_epochs=5,
+# #                       accelerator='cpu',
+# #                       default_root_dir=tmpdir)
+# #     trainer.train()
+# #     ntrainer = Trainer(solver=pinn, max_epochs=5, accelerator='cpu')
+# #     t = ntrainer.train(
+# #         ckpt_path=f'{tmpdir}/lightning_logs/version_0/checkpoints/epoch=4-step=10.ckpt')
+# #     import shutil
+# #     shutil.rmtree(tmpdir)
+
+
+# def test_train_inverse_problem_load():
+#     tmpdir = "tests/tmp_load_inv"
+#     poisson_problem = InversePoisson()
+#     boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4', 'D']
+#     n = 100
+#     poisson_problem.discretise_domain(n, 'random', locations=boundaries)
+#     pinn = GPINN(problem=poisson_problem,
+#                 model=model,
+#                 extra_features=None,
+#                 loss=LpLoss())
+#     trainer = Trainer(solver=pinn,
+#                       max_epochs=15,
+#                       accelerator='cpu',
+#                       default_root_dir=tmpdir)
 #     trainer.train()
+#     new_pinn = GPINN.load_from_checkpoint(
+#         f'{tmpdir}/lightning_logs/version_0/checkpoints/epoch=14-step=30.ckpt',
+#         problem = poisson_problem, model=model)
+#     test_pts = CartesianDomain({'x': [0, 1], 'y': [0, 1]}).sample(10)
+#     assert new_pinn.forward(test_pts).extract(['u']).shape == (10, 1)
+#     assert new_pinn.forward(test_pts).extract(
+#         ['u']).shape == pinn.forward(test_pts).extract(['u']).shape
+#     torch.testing.assert_close(
+#         new_pinn.forward(test_pts).extract(['u']),
+#         pinn.forward(test_pts).extract(['u']))
+#     import shutil
+#     shutil.rmtree(tmpdir)
+
+# # # TODO fix asap. Basically sampling few variables
+# # # works only if both variables are in a range.
+# # # if one is fixed and the other not, this will
+# # # not work. This test also needs to be fixed and
+# # # insert in test problem not in test pinn.
+# # def test_train_cpu_sampling_few_vars():
+# #     poisson_problem = Poisson()
+# #     boundaries = ['gamma1', 'gamma2', 'gamma3']
+# #     n = 10
+# #     poisson_problem.discretise_domain(n, 'grid', locations=boundaries)
+# #     poisson_problem.discretise_domain(n, 'random', locations=['gamma4'], variables=['x'])
+# #     poisson_problem.discretise_domain(n, 'random', locations=['gamma4'], variables=['y'])
+# #     pinn = GPINN(problem = poisson_problem, model=model, extra_features=None, loss=LpLoss())
+# #     trainer = Trainer(solver=pinn, kwargs={'max_epochs' : 5, 'accelerator':'cpu'})
+# #     trainer.train()
+
 
-# def test_train_gpu(): #TODO fix ASAP
+# def test_train_extra_feats_cpu():
 #     poisson_problem = Poisson()
 #     boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4']
 #     n = 10
 #     poisson_problem.discretise_domain(n, 'grid', locations=boundaries)
-#     poisson_problem.conditions.pop('data') # The input/output pts are allocated on cpu
-#     pinn = GPINN(problem = poisson_problem, model=model, extra_features=None, loss=LpLoss())
-#     trainer = Trainer(solver=pinn, kwargs={'max_epochs' : 5, 'accelerator':'gpu'})
+#     pinn = GPINN(problem=poisson_problem,
+#                 model=model_extra_feats,
+#                 extra_features=extra_feats)
+#     trainer = Trainer(solver=pinn, max_epochs=5, accelerator='cpu')
 #     trainer.train()
 
-# def test_train_2():
-#     boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4']
-#     n = 10
-#     expected_keys = [[], list(range(0, 50, 3))]
-#     param = [0, 3]
-#     for i, truth_key in zip(param, expected_keys):
-#         pinn = GPINN(problem, model)
-#         pinn.discretise_domain(n, 'grid', locations=boundaries)
-#         pinn.discretise_domain(n, 'grid', locations=['D'])
-#         pinn.train(50, save_loss=i)
-#         assert list(pinn.history_loss.keys()) == truth_key
-
-
-# def test_train_extra_feats():
-#     pinn = GPINN(problem, model_extra_feat, [myFeature()])
-#     boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4']
-#     n = 10
-#     pinn.discretise_domain(n, 'grid', locations=boundaries)
-#     pinn.discretise_domain(n, 'grid', locations=['D'])
-#     pinn.train(5)
 
+# # TODO, fix GitHub actions to run also on GPU
+# # def test_train_gpu():
+# #     poisson_problem = Poisson()
+# #     boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4']
+# #     n = 10
+# #     poisson_problem.discretise_domain(n, 'grid', locations=boundaries)
+# #     pinn = GPINN(problem = poisson_problem, model=model, extra_features=None, loss=LpLoss())
+# #     trainer = Trainer(solver=pinn, kwargs={'max_epochs' : 5, 'accelerator':'gpu'})
+# #     trainer.train()
 
-# def test_train_2_extra_feats():
-#     boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4']
-#     n = 10
-#     expected_keys = [[], list(range(0, 50, 3))]
-#     param = [0, 3]
-#     for i, truth_key in zip(param, expected_keys):
-#         pinn = GPINN(problem, model_extra_feat, [myFeature()])
-#         pinn.discretise_domain(n, 'grid', locations=boundaries)
-#         pinn.discretise_domain(n, 'grid', locations=['D'])
-#         pinn.train(50, save_loss=i)
-#         assert list(pinn.history_loss.keys()) == truth_key
-
+# # def test_train_gpu(): #TODO fix ASAP
+# #     poisson_problem = Poisson()
+# #     boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4']
+# #     n = 10
+# #     poisson_problem.discretise_domain(n, 'grid', locations=boundaries)
+# #     poisson_problem.conditions.pop('data') # The input/output pts are allocated on cpu
+# #     pinn = GPINN(problem = poisson_problem, model=model, extra_features=None, loss=LpLoss())
+# #     trainer = Trainer(solver=pinn, kwargs={'max_epochs' : 5, 'accelerator':'gpu'})
+# #     trainer.train()
 
-# def test_train_with_optimizer_kwargs():
-#     boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4']
-#     n = 10
-#     expected_keys = [[], list(range(0, 50, 3))]
-#     param = [0, 3]
-#     for i, truth_key in zip(param, expected_keys):
-#         pinn = GPINN(problem, model, optimizer_kwargs={'lr' : 0.3})
-#         pinn.discretise_domain(n, 'grid', locations=boundaries)
-#         pinn.discretise_domain(n, 'grid', locations=['D'])
-#         pinn.train(50, save_loss=i)
-#         assert list(pinn.history_loss.keys()) == truth_key
+# # def test_train_2():
+# #     boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4']
+# #     n = 10
+# #     expected_keys = [[], list(range(0, 50, 3))]
+# #     param = [0, 3]
+# #     for i, truth_key in zip(param, expected_keys):
+# #         pinn = GPINN(problem, model)
+# #         pinn.discretise_domain(n, 'grid', locations=boundaries)
+# #         pinn.discretise_domain(n, 'grid', locations=['D'])
+# #         pinn.train(50, save_loss=i)
+# #         assert list(pinn.history_loss.keys()) == truth_key
 
 
-# def test_train_with_lr_scheduler():
-#     boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4']
-#     n = 10
-#     expected_keys = [[], list(range(0, 50, 3))]
-#     param = [0, 3]
-#     for i, truth_key in zip(param, expected_keys):
-#         pinn = GPINN(
-#             problem,
-#             model,
-#             lr_scheduler_type=torch.optim.lr_scheduler.CyclicLR,
-#             lr_scheduler_kwargs={'base_lr' : 0.1, 'max_lr' : 0.3, 'cycle_momentum': False}
-#         )
-#         pinn.discretise_domain(n, 'grid', locations=boundaries)
-#         pinn.discretise_domain(n, 'grid', locations=['D'])
-#         pinn.train(50, save_loss=i)
-#         assert list(pinn.history_loss.keys()) == truth_key
-
-
-# # def test_train_batch():
-# #     pinn = GPINN(problem, model, batch_size=6)
+# # def test_train_extra_feats():
+# #     pinn = GPINN(problem, model_extra_feat, [myFeature()])
 # #     boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4']
 # #     n = 10
 # #     pinn.discretise_domain(n, 'grid', locations=boundaries)
@@ -411,34 +358,87 @@ def test_train_extra_feats_cpu():
 # #     pinn.train(5)
 
 
-# # def test_train_batch_2():
+# # def test_train_2_extra_feats():
+# #     boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4']
+# #     n = 10
+# #     expected_keys = [[], list(range(0, 50, 3))]
+# #     param = [0, 3]
+# #     for i, truth_key in zip(param, expected_keys):
+# #         pinn = GPINN(problem, model_extra_feat, [myFeature()])
+# #         pinn.discretise_domain(n, 'grid', locations=boundaries)
+# #         pinn.discretise_domain(n, 'grid', locations=['D'])
+# #         pinn.train(50, save_loss=i)
+# #         assert list(pinn.history_loss.keys()) == truth_key
+
+
+# # def test_train_with_optimizer_kwargs():
 # #     boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4']
 # #     n = 10
 # #     expected_keys = [[], list(range(0, 50, 3))]
 # #     param = [0, 3]
 # #     for i, truth_key in zip(param, expected_keys):
-# #         pinn = GPINN(problem, model, batch_size=6)
+# #         pinn = GPINN(problem, model, optimizer_kwargs={'lr' : 0.3})
 # #         pinn.discretise_domain(n, 'grid', locations=boundaries)
 # #         pinn.discretise_domain(n, 'grid', locations=['D'])
 # #         pinn.train(50, save_loss=i)
 # #         assert list(pinn.history_loss.keys()) == truth_key
 
 
-# if torch.cuda.is_available():
+# # def test_train_with_lr_scheduler():
+# #     boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4']
+# #     n = 10
+# #     expected_keys = [[], list(range(0, 50, 3))]
+# #     param = [0, 3]
+# #     for i, truth_key in zip(param, expected_keys):
+# #         pinn = GPINN(
+# #             problem,
+# #             model,
+# #             lr_scheduler_type=torch.optim.lr_scheduler.CyclicLR,
+# #             lr_scheduler_kwargs={'base_lr' : 0.1, 'max_lr' : 0.3, 'cycle_momentum': False}
+# #         )
+# #         pinn.discretise_domain(n, 'grid', locations=boundaries)
+# #         pinn.discretise_domain(n, 'grid', locations=['D'])
+# #         pinn.train(50, save_loss=i)
+# #         assert list(pinn.history_loss.keys()) == truth_key
 
-#     # def test_gpu_train():
-#     #     pinn = GPINN(problem, model, batch_size=20, device='cuda')
-#     #     boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4']
-#     #     n = 100
-#     #     pinn.discretise_domain(n, 'grid', locations=boundaries)
-#     #     pinn.discretise_domain(n, 'grid', locations=['D'])
-#     #     pinn.train(5)
 
-#     def test_gpu_train_nobatch():
-#         pinn = GPINN(problem, model, batch_size=None, device='cuda')
-#         boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4']
-#         n = 100
-#         pinn.discretise_domain(n, 'grid', locations=boundaries)
-#         pinn.discretise_domain(n, 'grid', locations=['D'])
-#         pinn.train(5)
+# # # def test_train_batch():
+# # #     pinn = GPINN(problem, model, batch_size=6)
+# # #     boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4']
+# # #     n = 10
+# # #     pinn.discretise_domain(n, 'grid', locations=boundaries)
+# # #     pinn.discretise_domain(n, 'grid', locations=['D'])
+# # #     pinn.train(5)
+
+
+# # # def test_train_batch_2():
+# # #     boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4']
+# # #     n = 10
+# # #     expected_keys = [[], list(range(0, 50, 3))]
+# # #     param = [0, 3]
+# # #     for i, truth_key in zip(param, expected_keys):
+# # #         pinn = GPINN(problem, model, batch_size=6)
+# # #         pinn.discretise_domain(n, 'grid', locations=boundaries)
+# # #         pinn.discretise_domain(n, 'grid', locations=['D'])
+# # #         pinn.train(50, save_loss=i)
+# # #         assert list(pinn.history_loss.keys()) == truth_key
+
+
+# # if torch.cuda.is_available():
+
+# #     # def test_gpu_train():
+# #     #     pinn = GPINN(problem, model, batch_size=20, device='cuda')
+# #     #     boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4']
+# #     #     n = 100
+# #     #     pinn.discretise_domain(n, 'grid', locations=boundaries)
+# #     #     pinn.discretise_domain(n, 'grid', locations=['D'])
+# #     #     pinn.train(5)
+
+# #     def test_gpu_train_nobatch():
+# #         pinn = GPINN(problem, model, batch_size=None, device='cuda')
+# #         boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4']
+# #         n = 100
+# #         pinn.discretise_domain(n, 'grid', locations=boundaries)
+# #         pinn.discretise_domain(n, 'grid', locations=['D'])
+# #         pinn.train(5)
 
diff --git a/tests/test_solvers/test_pinn.py b/tests/test_solvers/test_pinn.py
index 72887a4f7..feb2b2eb9 100644
--- a/tests/test_solvers/test_pinn.py
+++ b/tests/test_solvers/test_pinn.py
@@ -6,255 +6,180 @@
 from pina.solvers import PINN
 from pina.trainer import Trainer
 from pina.model import FeedForward
-from pina.equation.equation import Equation
+from pina.equation import Equation
 from pina.equation.equation_factory import FixedValue
 from pina.loss import LpLoss
-
-
-def laplace_equation(input_, output_):
-    force_term = (torch.sin(input_.extract(['x']) * torch.pi) *
-                  torch.sin(input_.extract(['y']) * torch.pi))
-    delta_u = laplacian(output_.extract(['u']), input_)
-    return delta_u - force_term
-
-
-my_laplace = Equation(laplace_equation)
-in_ = LabelTensor(torch.tensor([[0., 1.]]), ['x', 'y'])
-out_ = LabelTensor(torch.tensor([[0.]]), ['u'])
-in2_ = LabelTensor(torch.rand(60, 2), ['x', 'y'])
-out2_ = LabelTensor(torch.rand(60, 1), ['u'])
-
-
-class InversePoisson(SpatialProblem, InverseProblem):
-    '''
-    Problem definition for the Poisson equation.
-    '''
-    output_variables = ['u']
-    x_min = -2
-    x_max = 2
-    y_min = -2
-    y_max = 2
-    data_input = LabelTensor(torch.rand(10, 2), ['x', 'y'])
-    data_output = LabelTensor(torch.rand(10, 1), ['u'])
-    spatial_domain = CartesianDomain({'x': [x_min, x_max], 'y': [y_min, y_max]})
-    # define the ranges for the parameters
-    unknown_parameter_domain = CartesianDomain({'mu1': [-1, 1], 'mu2': [-1, 1]})
-
-    def laplace_equation(input_, output_, params_):
-        '''
-        Laplace equation with a force term.
-        '''
-        force_term = torch.exp(
-                - 2*(input_.extract(['x']) - params_['mu1'])**2
-                - 2*(input_.extract(['y']) - params_['mu2'])**2)
-        delta_u = laplacian(output_, input_, components=['u'], d=['x', 'y'])
-
-        return delta_u - force_term
-
-    # define the conditions for the loss (boundary conditions, equation, data)
-    conditions = {
-        'gamma1': Condition(domain=CartesianDomain({'x': [x_min, x_max],
-            'y':  y_max}),
-            equation=FixedValue(0.0, components=['u'])),
-        'gamma2': Condition(domain=CartesianDomain(
-            {'x': [x_min, x_max], 'y': y_min
-            }),
-            equation=FixedValue(0.0, components=['u'])),
-        'gamma3': Condition(domain=CartesianDomain(
-            {'x':  x_max, 'y': [y_min, y_max]
-            }),
-            equation=FixedValue(0.0, components=['u'])),
-        'gamma4': Condition(domain=CartesianDomain(
-            {'x': x_min, 'y': [y_min, y_max]
-            }),
-            equation=FixedValue(0.0, components=['u'])),
-        'D': Condition(domain=CartesianDomain(
-            {'x': [x_min, x_max], 'y': [y_min, y_max]
-            }),
-        equation=Equation(laplace_equation)),
-        'data': Condition(input_points=data_input.extract(['x', 'y']),
-                          output_points=data_output)
-    }
-
-
-class Poisson(SpatialProblem):
-    output_variables = ['u']
-    spatial_domain = CartesianDomain({'x': [0, 1], 'y': [0, 1]})
-
-    conditions = {
-        'gamma1': Condition(
-            domain=CartesianDomain({'x': [0, 1], 'y':  1}),
-            equation=FixedValue(0.0)),
-        'gamma2': Condition(
-            domain=CartesianDomain({'x': [0, 1], 'y': 0}),
-            equation=FixedValue(0.0)),
-        'gamma3': Condition(
-            domain=CartesianDomain({'x':  1, 'y': [0, 1]}),
-            equation=FixedValue(0.0)),
-        'gamma4': Condition(
-            domain=CartesianDomain({'x': 0, 'y': [0, 1]}),
-            equation=FixedValue(0.0)),
-        'D': Condition(
-            input_points=LabelTensor(torch.rand(size=(100, 2)), ['x', 'y']),
-            equation=my_laplace),
-        'data': Condition(
-            input_points=in_,
-            output_points=out_),
-        'data2': Condition(
-            input_points=in2_,
-            output_points=out2_)
-    }
-
-    def poisson_sol(self, pts):
-        return -(torch.sin(pts.extract(['x']) * torch.pi) *
-                 torch.sin(pts.extract(['y']) * torch.pi)) / (2 * torch.pi**2)
-
-    truth_solution = poisson_sol
-
-class myFeature(torch.nn.Module):
-    """
-    Feature: sin(x)
-    """
-
-    def __init__(self):
-        super(myFeature, self).__init__()
-
-    def forward(self, x):
-        t = (torch.sin(x.extract(['x']) * torch.pi) *
-             torch.sin(x.extract(['y']) * torch.pi))
-        return LabelTensor(t, ['sin(x)sin(y)'])
-
-
-# make the problem
-poisson_problem = Poisson()
-model = FeedForward(len(poisson_problem.input_variables),
-                    len(poisson_problem.output_variables))
-model_extra_feats = FeedForward(
-    len(poisson_problem.input_variables) + 1,
-    len(poisson_problem.output_variables))
-extra_feats = [myFeature()]
-
-
-def test_constructor():
-    PINN(problem=poisson_problem, model=model, extra_features=None)
-
-
-def test_constructor_extra_feats():
-    model_extra_feats = FeedForward(
-        len(poisson_problem.input_variables) + 1,
-        len(poisson_problem.output_variables))
-    PINN(problem=poisson_problem,
-         model=model_extra_feats,
-         extra_features=extra_feats)
-
-
-def test_train_cpu():
-    poisson_problem = Poisson()
-    boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4']
-    n = 10
-    poisson_problem.discretise_domain(n, 'grid', locations=boundaries)
-    pinn = PINN(problem = poisson_problem, model=model,
-                extra_features=None, loss=LpLoss())
-    trainer = Trainer(solver=pinn, max_epochs=1,
-                      accelerator='cpu', batch_size=20, val_size=0., train_size=1., test_size=0.)
-
-def test_train_load():
-    tmpdir = "tests/tmp_load"
-    poisson_problem = Poisson()
-    boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4']
-    n = 10
-    poisson_problem.discretise_domain(n, 'grid', locations=boundaries)
-    pinn = PINN(problem=poisson_problem,
-                model=model,
-                extra_features=None,
-                loss=LpLoss())
-    trainer = Trainer(solver=pinn,
-                      max_epochs=15,
-                      accelerator='cpu',
-                      default_root_dir=tmpdir)
-    trainer.train()
-    new_pinn = PINN.load_from_checkpoint(
-        f'{tmpdir}/lightning_logs/version_0/checkpoints/epoch=14-step=15.ckpt',
-        problem = poisson_problem, model=model)
-    test_pts = CartesianDomain({'x': [0, 1], 'y': [0, 1]}).sample(10)
-    assert new_pinn.forward(test_pts).extract(['u']).shape == (10, 1)
-    assert new_pinn.forward(test_pts).extract(
-        ['u']).shape == pinn.forward(test_pts).extract(['u']).shape
-    torch.testing.assert_close(
-        new_pinn.forward(test_pts).extract(['u']),
-        pinn.forward(test_pts).extract(['u']))
-    import shutil
-    shutil.rmtree(tmpdir)
-
-def test_train_restore():
-    tmpdir = "tests/tmp_restore"
-    poisson_problem = Poisson()
-    boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4']
-    n = 10
-    poisson_problem.discretise_domain(n, 'grid', locations=boundaries)
-    pinn = PINN(problem=poisson_problem,
-                model=model,
-                extra_features=None,
-                loss=LpLoss())
-    trainer = Trainer(solver=pinn,
-                      max_epochs=5,
-                      accelerator='cpu',
-                      default_root_dir=tmpdir)
-    trainer.train()
-    ntrainer = Trainer(solver=pinn, max_epochs=15, accelerator='cpu')
-    t = ntrainer.train(
-        ckpt_path=f'{tmpdir}/lightning_logs/version_0/'
-                  'checkpoints/epoch=4-step=5.ckpt')
-    import shutil
-    shutil.rmtree(tmpdir)
-
-def test_train_inverse_problem_cpu():
-    poisson_problem = InversePoisson()
-    boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4', 'D']
-    n = 100
-    poisson_problem.discretise_domain(n, 'random', locations=boundaries,
-                                      variables=['x', 'y'])
-    pinn = PINN(problem = poisson_problem, model=model,
-                extra_features=None, loss=LpLoss())
-    trainer = Trainer(solver=pinn, max_epochs=1,
-                      accelerator='cpu', batch_size=20)
-    trainer.train()
-
-def test_train_extra_feats_cpu():
-    poisson_problem = Poisson()
-    boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4']
-    n = 10
-    poisson_problem.discretise_domain(n, 'grid', locations=boundaries)
-    pinn = PINN(problem=poisson_problem,
-                model=model_extra_feats,
-                extra_features=extra_feats)
-    trainer = Trainer(solver=pinn, max_epochs=5, accelerator='cpu')
-    trainer.train()
-
-def test_train_inverse_problem_load():
-    tmpdir = "tests/tmp_load_inv"
-    poisson_problem = InversePoisson()
-    boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4', 'D']
-    n = 100
-    poisson_problem.discretise_domain(n, 'random', locations=boundaries)
-    pinn = PINN(problem=poisson_problem,
-                model=model,
-                extra_features=None,
-                loss=LpLoss())
-    trainer = Trainer(solver=pinn,
-                      max_epochs=15,
-                      accelerator='cpu',
-                      default_root_dir=tmpdir)
-    trainer.train()
-    new_pinn = PINN.load_from_checkpoint(
-        f'{tmpdir}/lightning_logs/version_0/checkpoints/epoch=14-step=15.ckpt',
-        problem = poisson_problem, model=model)
-    test_pts = CartesianDomain({'x': [0, 1], 'y': [0, 1]}).sample(10)
-    assert new_pinn.forward(test_pts).extract(['u']).shape == (10, 1)
-    assert new_pinn.forward(test_pts).extract(
-        ['u']).shape == pinn.forward(test_pts).extract(['u']).shape
-    torch.testing.assert_close(
-        new_pinn.forward(test_pts).extract(['u']),
-        pinn.forward(test_pts).extract(['u']))
-    import shutil
-    shutil.rmtree(tmpdir)
\ No newline at end of file
+from pina.problem.zoo import Poisson2DSquareProblem
+
+# class InversePoisson(SpatialProblem, InverseProblem):
+#     '''
+#     Problem definition for the Poisson equation.
+#     '''
+#     output_variables = ['u']
+#     x_min = -2
+#     x_max = 2
+#     y_min = -2
+#     y_max = 2
+#     data_input = LabelTensor(torch.rand(10, 2), ['x', 'y'])
+#     data_output = LabelTensor(torch.rand(10, 1), ['u'])
+#     spatial_domain = CartesianDomain({'x': [x_min, x_max], 'y': [y_min, y_max]})
+#     # define the ranges for the parameters
+#     unknown_parameter_domain = CartesianDomain({'mu1': [-1, 1], 'mu2': [-1, 1]})
+
+#     def laplace_equation(input_, output_, params_):
+#         '''
+#         Laplace equation with a force term.
+#         '''
+#         force_term = torch.exp(
+#                 - 2*(input_.extract(['x']) - params_['mu1'])**2
+#                 - 2*(input_.extract(['y']) - params_['mu2'])**2)
+#         delta_u = laplacian(output_, input_, components=['u'], d=['x', 'y'])
+
+#         return delta_u - force_term
+
+#     # define the conditions for the loss (boundary conditions, equation, data)
+#     conditions = {
+#         'gamma1': Condition(domain=CartesianDomain({'x': [x_min, x_max],
+#             'y':  y_max}),
+#             equation=FixedValue(0.0, components=['u'])),
+#         'gamma2': Condition(domain=CartesianDomain(
+#             {'x': [x_min, x_max], 'y': y_min
+#             }),
+#             equation=FixedValue(0.0, components=['u'])),
+#         'gamma3': Condition(domain=CartesianDomain(
+#             {'x':  x_max, 'y': [y_min, y_max]
+#             }),
+#             equation=FixedValue(0.0, components=['u'])),
+#         'gamma4': Condition(domain=CartesianDomain(
+#             {'x': x_min, 'y': [y_min, y_max]
+#             }),
+#             equation=FixedValue(0.0, components=['u'])),
+#         'D': Condition(domain=CartesianDomain(
+#             {'x': [x_min, x_max], 'y': [y_min, y_max]
+#             }),
+#         equation=Equation(laplace_equation)),
+#         'data': Condition(input_points=data_input.extract(['x', 'y']),
+#                           output_points=data_output)
+#     }
+
+
+# # make the problem
+# poisson_problem = Poisson2DSquareProblem()
+# model = FeedForward(len(poisson_problem.input_variables),
+#                     len(poisson_problem.output_variables))
+# model_extra_feats = FeedForward(
+#     len(poisson_problem.input_variables) + 1,
+#     len(poisson_problem.output_variables))
+
+
+# def test_constructor():
+#     PINN(problem=poisson_problem, model=model, extra_features=None)
+
+
+# def test_constructor_extra_feats():
+#     model_extra_feats = FeedForward(
+#         len(poisson_problem.input_variables) + 1,
+#         len(poisson_problem.output_variables))
+#     PINN(problem=poisson_problem,
+#          model=model_extra_feats)
+
+
+# def test_train_cpu():
+#     poisson_problem = Poisson2DSquareProblem()
+#     boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4']
+#     n = 10
+#     poisson_problem.discretise_domain(n, 'grid', locations=boundaries)
+#     pinn = PINN(problem = poisson_problem, model=model,
+#                 extra_features=None, loss=LpLoss())
+#     trainer = Trainer(solver=pinn, max_epochs=1,
+#                       accelerator='cpu', batch_size=20, val_size=0., train_size=1., test_size=0.)
+
+# def test_train_load():
+#     tmpdir = "tests/tmp_load"
+#     poisson_problem = Poisson2DSquareProblem()
+#     boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4']
+#     n = 10
+#     poisson_problem.discretise_domain(n, 'grid', locations=boundaries)
+#     pinn = PINN(problem=poisson_problem,
+#                 model=model,
+#                 extra_features=None,
+#                 loss=LpLoss())
+#     trainer = Trainer(solver=pinn,
+#                       max_epochs=15,
+#                       accelerator='cpu',
+#                       default_root_dir=tmpdir)
+#     trainer.train()
+#     new_pinn = PINN.load_from_checkpoint(
+#         f'{tmpdir}/lightning_logs/version_0/checkpoints/epoch=14-step=15.ckpt',
+#         problem = poisson_problem, model=model)
+#     test_pts = CartesianDomain({'x': [0, 1], 'y': [0, 1]}).sample(10)
+#     assert new_pinn.forward(test_pts).extract(['u']).shape == (10, 1)
+#     assert new_pinn.forward(test_pts).extract(
+#         ['u']).shape == pinn.forward(test_pts).extract(['u']).shape
+#     torch.testing.assert_close(
+#         new_pinn.forward(test_pts).extract(['u']),
+#         pinn.forward(test_pts).extract(['u']))
+#     import shutil
+#     shutil.rmtree(tmpdir)
+
+# def test_train_restore():
+#     tmpdir = "tests/tmp_restore"
+#     poisson_problem = Poisson2DSquareProblem()
+#     boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4']
+#     n = 10
+#     poisson_problem.discretise_domain(n, 'grid', locations=boundaries)
+#     pinn = PINN(problem=poisson_problem,
+#                 model=model,
+#                 extra_features=None,
+#                 loss=LpLoss())
+#     trainer = Trainer(solver=pinn,
+#                       max_epochs=5,
+#                       accelerator='cpu',
+#                       default_root_dir=tmpdir)
+#     trainer.train()
+#     ntrainer = Trainer(solver=pinn, max_epochs=15, accelerator='cpu')
+#     t = ntrainer.train(
+#         ckpt_path=f'{tmpdir}/lightning_logs/version_0/'
+#                   'checkpoints/epoch=4-step=5.ckpt')
+#     import shutil
+#     shutil.rmtree(tmpdir)
+
+# def test_train_inverse_problem_cpu():
+#     poisson_problem = InversePoisson()
+#     boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4', 'D']
+#     n = 100
+#     poisson_problem.discretise_domain(n, 'random', locations=boundaries,
+#                                       variables=['x', 'y'])
+#     pinn = PINN(problem = poisson_problem, model=model,
+#                 extra_features=None, loss=LpLoss())
+#     trainer = Trainer(solver=pinn, max_epochs=1,
+#                       accelerator='cpu', batch_size=20)
+#     trainer.train()
+
+# def test_train_inverse_problem_load():
+#     tmpdir = "tests/tmp_load_inv"
+#     poisson_problem = InversePoisson()
+#     boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4', 'D']
+#     n = 100
+#     poisson_problem.discretise_domain(n, 'random', locations=boundaries)
+#     pinn = PINN(problem=poisson_problem,
+#                 model=model,
+#                 extra_features=None,
+#                 loss=LpLoss())
+#     trainer = Trainer(solver=pinn,
+#                       max_epochs=15,
+#                       accelerator='cpu',
+#                       default_root_dir=tmpdir)
+#     trainer.train()
+#     new_pinn = PINN.load_from_checkpoint(
+#         f'{tmpdir}/lightning_logs/version_0/checkpoints/epoch=14-step=15.ckpt',
+#         problem = poisson_problem, model=model)
+#     test_pts = CartesianDomain({'x': [0, 1], 'y': [0, 1]}).sample(10)
+#     assert new_pinn.forward(test_pts).extract(['u']).shape == (10, 1)
+#     assert new_pinn.forward(test_pts).extract(
+#         ['u']).shape == pinn.forward(test_pts).extract(['u']).shape
+#     torch.testing.assert_close(
+#         new_pinn.forward(test_pts).extract(['u']),
+#         pinn.forward(test_pts).extract(['u']))
+#     import shutil
+#     shutil.rmtree(tmpdir)
\ No newline at end of file
diff --git a/tests/test_solvers/test_rba_pinn.py b/tests/test_solvers/test_rba_pinn.py
index aad47bbb7..b3488bbc4 100644
--- a/tests/test_solvers/test_rba_pinn.py
+++ b/tests/test_solvers/test_rba_pinn.py
@@ -3,251 +3,187 @@
 
 from pina.problem import SpatialProblem, InverseProblem
 from pina.operators import laplacian
-from pina.geometry import CartesianDomain
+from pina.domain import CartesianDomain
 from pina import Condition, LabelTensor
 from pina.solvers import RBAPINN as PINN
 from pina.trainer import Trainer
 from pina.model import FeedForward
-from pina.equation.equation import Equation
+from pina.equation import Equation
 from pina.equation.equation_factory import FixedValue
 from pina.loss import LpLoss
 
 
-def laplace_equation(input_, output_):
-    force_term = (torch.sin(input_.extract(['x']) * torch.pi) *
-                  torch.sin(input_.extract(['y']) * torch.pi))
-    delta_u = laplacian(output_.extract(['u']), input_)
-    return delta_u - force_term
-
-
-my_laplace = Equation(laplace_equation)
-in_ = LabelTensor(torch.tensor([[0., 1.]]), ['x', 'y'])
-out_ = LabelTensor(torch.tensor([[0.]]), ['u'])
-in2_ = LabelTensor(torch.rand(60, 2), ['x', 'y'])
-out2_ = LabelTensor(torch.rand(60, 1), ['u'])
-
-
-class InversePoisson(SpatialProblem, InverseProblem):
-    '''
-    Problem definition for the Poisson equation.
-    '''
-    output_variables = ['u']
-    x_min = -2
-    x_max = 2
-    y_min = -2
-    y_max = 2
-    data_input = LabelTensor(torch.rand(10, 2), ['x', 'y'])
-    data_output = LabelTensor(torch.rand(10, 1), ['u'])
-    spatial_domain = CartesianDomain({'x': [x_min, x_max], 'y': [y_min, y_max]})
-    # define the ranges for the parameters
-    unknown_parameter_domain = CartesianDomain({'mu1': [-1, 1], 'mu2': [-1, 1]})
-
-    def laplace_equation(input_, output_, params_):
-        '''
-        Laplace equation with a force term.
-        '''
-        force_term = torch.exp(
-                - 2*(input_.extract(['x']) - params_['mu1'])**2
-                - 2*(input_.extract(['y']) - params_['mu2'])**2)
-        delta_u = laplacian(output_, input_, components=['u'], d=['x', 'y'])
-
-        return delta_u - force_term
-
-    # define the conditions for the loss (boundary conditions, equation, data)
-    conditions = {
-        'gamma1': Condition(location=CartesianDomain({'x': [x_min, x_max],
-            'y':  y_max}),
-            equation=FixedValue(0.0, components=['u'])),
-        'gamma2': Condition(location=CartesianDomain(
-            {'x': [x_min, x_max], 'y': y_min
-            }),
-            equation=FixedValue(0.0, components=['u'])),
-        'gamma3': Condition(location=CartesianDomain(
-            {'x':  x_max, 'y': [y_min, y_max]
-            }),
-            equation=FixedValue(0.0, components=['u'])),
-        'gamma4': Condition(location=CartesianDomain(
-            {'x': x_min, 'y': [y_min, y_max]
-            }),
-            equation=FixedValue(0.0, components=['u'])),
-        'D': Condition(location=CartesianDomain(
-            {'x': [x_min, x_max], 'y': [y_min, y_max]
-            }),
-        equation=Equation(laplace_equation)),
-        'data': Condition(input_points=data_input.extract(['x', 'y']),
-                          output_points=data_output)
-    }
-
-
-class Poisson(SpatialProblem):
-    output_variables = ['u']
-    spatial_domain = CartesianDomain({'x': [0, 1], 'y': [0, 1]})
-
-    conditions = {
-        'gamma1': Condition(
-            location=CartesianDomain({'x': [0, 1], 'y':  1}),
-            equation=FixedValue(0.0)),
-        'gamma2': Condition(
-            location=CartesianDomain({'x': [0, 1], 'y': 0}),
-            equation=FixedValue(0.0)),
-        'gamma3': Condition(
-            location=CartesianDomain({'x':  1, 'y': [0, 1]}),
-            equation=FixedValue(0.0)),
-        'gamma4': Condition(
-            location=CartesianDomain({'x': 0, 'y': [0, 1]}),
-            equation=FixedValue(0.0)),
-        'D': Condition(
-            input_points=LabelTensor(torch.rand(size=(100, 2)), ['x', 'y']),
-            equation=my_laplace),
-        'data': Condition(
-            input_points=in_,
-            output_points=out_),
-        'data2': Condition(
-            input_points=in2_,
-            output_points=out2_)
-    }
-
-    def poisson_sol(self, pts):
-        return -(torch.sin(pts.extract(['x']) * torch.pi) *
-                 torch.sin(pts.extract(['y']) * torch.pi)) / (2 * torch.pi**2)
-
-    truth_solution = poisson_sol
-
-
-class myFeature(torch.nn.Module):
-    """
-    Feature: sin(x)
-    """
-
-    def __init__(self):
-        super(myFeature, self).__init__()
-
-    def forward(self, x):
-        t = (torch.sin(x.extract(['x']) * torch.pi) *
-             torch.sin(x.extract(['y']) * torch.pi))
-        return LabelTensor(t, ['sin(x)sin(y)'])
-
-
-# make the problem
-poisson_problem = Poisson()
-model = FeedForward(len(poisson_problem.input_variables),
-                    len(poisson_problem.output_variables))
-model_extra_feats = FeedForward(
-    len(poisson_problem.input_variables) + 1,
-    len(poisson_problem.output_variables))
-extra_feats = [myFeature()]
-
-
-def test_constructor():
-    PINN(problem=poisson_problem, model=model, extra_features=None)
-    with pytest.raises(ValueError):
-        PINN(problem=poisson_problem, model=model, eta='x')
-        PINN(problem=poisson_problem, model=model, gamma='x')
-
-
-def test_constructor_extra_feats():
-    model_extra_feats = FeedForward(
-        len(poisson_problem.input_variables) + 1,
-        len(poisson_problem.output_variables))
-    PINN(problem=poisson_problem,
-         model=model_extra_feats,
-         extra_features=extra_feats)
-
-
-def test_train_cpu():
-    poisson_problem = Poisson()
-    boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4']
-    n = 10
-    poisson_problem.discretise_domain(n, 'grid', locations=boundaries)
-    pinn = PINN(problem = poisson_problem, model=model,
-                extra_features=None, loss=LpLoss())
-    trainer = Trainer(solver=pinn, max_epochs=1,
-                      accelerator='cpu', batch_size=20)
-    trainer.train()
-
-def test_log():
-    poisson_problem.discretise_domain(100)
-    solver = PINN(problem = poisson_problem, model=model,
-                extra_features=None, loss=LpLoss())
-    trainer = Trainer(solver, max_epochs=2, accelerator='cpu')
-    trainer.train()
-    # assert the logged metrics are correct
-    logged_metrics = sorted(list(trainer.logged_metrics.keys()))
-    total_metrics = sorted(
-        list([key + '_loss' for key in poisson_problem.conditions.keys()])
-        + ['mean_loss'])
-    assert logged_metrics == total_metrics
-
-def test_train_restore():
-    tmpdir = "tests/tmp_restore"
-    poisson_problem = Poisson()
-    boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4']
-    n = 10
-    poisson_problem.discretise_domain(n, 'grid', locations=boundaries)
-    pinn = PINN(problem=poisson_problem,
-                model=model,
-                extra_features=None,
-                loss=LpLoss())
-    trainer = Trainer(solver=pinn,
-                      max_epochs=5,
-                      accelerator='cpu',
-                      default_root_dir=tmpdir)
-    trainer.train()
-    ntrainer = Trainer(solver=pinn, max_epochs=15, accelerator='cpu')
-    t = ntrainer.train(
-        ckpt_path=f'{tmpdir}/lightning_logs/version_0/'
-        'checkpoints/epoch=4-step=10.ckpt')
-    import shutil
-    shutil.rmtree(tmpdir)
-
-
-def test_train_load():
-    tmpdir = "tests/tmp_load"
-    poisson_problem = Poisson()
-    boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4']
-    n = 10
-    poisson_problem.discretise_domain(n, 'grid', locations=boundaries)
-    pinn = PINN(problem=poisson_problem,
-                model=model,
-                extra_features=None,
-                loss=LpLoss())
-    trainer = Trainer(solver=pinn,
-                      max_epochs=15,
-                      accelerator='cpu',
-                      default_root_dir=tmpdir)
-    trainer.train()
-    new_pinn = PINN.load_from_checkpoint(
-        f'{tmpdir}/lightning_logs/version_0/checkpoints/epoch=14-step=30.ckpt',
-        problem = poisson_problem, model=model)
-    test_pts = CartesianDomain({'x': [0, 1], 'y': [0, 1]}).sample(10)
-    assert new_pinn.forward(test_pts).extract(['u']).shape == (10, 1)
-    assert new_pinn.forward(test_pts).extract(
-        ['u']).shape == pinn.forward(test_pts).extract(['u']).shape
-    torch.testing.assert_close(
-        new_pinn.forward(test_pts).extract(['u']),
-        pinn.forward(test_pts).extract(['u']))
-    import shutil
-    shutil.rmtree(tmpdir)
-
-def test_train_inverse_problem_cpu():
-    poisson_problem = InversePoisson()
-    boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4', 'D']
-    n = 100
-    poisson_problem.discretise_domain(n, 'random', locations=boundaries)
-    pinn = PINN(problem = poisson_problem, model=model,
-                extra_features=None, loss=LpLoss())
-    trainer = Trainer(solver=pinn, max_epochs=1,
-                      accelerator='cpu', batch_size=20)
-    trainer.train()
-
-
-# # TODO does not currently work
-# def test_train_inverse_problem_restore():
-#     tmpdir = "tests/tmp_restore_inv"
-#     poisson_problem = InversePoisson()
-#     boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4', 'D']
-#     n = 100
-#     poisson_problem.discretise_domain(n, 'random', locations=boundaries)
+# def laplace_equation(input_, output_):
+#     force_term = (torch.sin(input_.extract(['x']) * torch.pi) *
+#                   torch.sin(input_.extract(['y']) * torch.pi))
+#     delta_u = laplacian(output_.extract(['u']), input_)
+#     return delta_u - force_term
+
+
+# my_laplace = Equation(laplace_equation)
+# in_ = LabelTensor(torch.tensor([[0., 1.]]), ['x', 'y'])
+# out_ = LabelTensor(torch.tensor([[0.]]), ['u'])
+# in2_ = LabelTensor(torch.rand(60, 2), ['x', 'y'])
+# out2_ = LabelTensor(torch.rand(60, 1), ['u'])
+
+
+# class InversePoisson(SpatialProblem, InverseProblem):
+#     '''
+#     Problem definition for the Poisson equation.
+#     '''
+#     output_variables = ['u']
+#     x_min = -2
+#     x_max = 2
+#     y_min = -2
+#     y_max = 2
+#     data_input = LabelTensor(torch.rand(10, 2), ['x', 'y'])
+#     data_output = LabelTensor(torch.rand(10, 1), ['u'])
+#     spatial_domain = CartesianDomain({'x': [x_min, x_max], 'y': [y_min, y_max]})
+#     # define the ranges for the parameters
+#     unknown_parameter_domain = CartesianDomain({'mu1': [-1, 1], 'mu2': [-1, 1]})
+
+#     def laplace_equation(input_, output_, params_):
+#         '''
+#         Laplace equation with a force term.
+#         '''
+#         force_term = torch.exp(
+#                 - 2*(input_.extract(['x']) - params_['mu1'])**2
+#                 - 2*(input_.extract(['y']) - params_['mu2'])**2)
+#         delta_u = laplacian(output_, input_, components=['u'], d=['x', 'y'])
+
+#         return delta_u - force_term
+
+#     # define the conditions for the loss (boundary conditions, equation, data)
+#     conditions = {
+#         'gamma1': Condition(location=CartesianDomain({'x': [x_min, x_max],
+#             'y':  y_max}),
+#             equation=FixedValue(0.0, components=['u'])),
+#         'gamma2': Condition(location=CartesianDomain(
+#             {'x': [x_min, x_max], 'y': y_min
+#             }),
+#             equation=FixedValue(0.0, components=['u'])),
+#         'gamma3': Condition(location=CartesianDomain(
+#             {'x':  x_max, 'y': [y_min, y_max]
+#             }),
+#             equation=FixedValue(0.0, components=['u'])),
+#         'gamma4': Condition(location=CartesianDomain(
+#             {'x': x_min, 'y': [y_min, y_max]
+#             }),
+#             equation=FixedValue(0.0, components=['u'])),
+#         'D': Condition(location=CartesianDomain(
+#             {'x': [x_min, x_max], 'y': [y_min, y_max]
+#             }),
+#         equation=Equation(laplace_equation)),
+#         'data': Condition(input_points=data_input.extract(['x', 'y']),
+#                           output_points=data_output)
+#     }
+
+
+# class Poisson(SpatialProblem):
+#     output_variables = ['u']
+#     spatial_domain = CartesianDomain({'x': [0, 1], 'y': [0, 1]})
+
+#     conditions = {
+#         'gamma1': Condition(
+#             location=CartesianDomain({'x': [0, 1], 'y':  1}),
+#             equation=FixedValue(0.0)),
+#         'gamma2': Condition(
+#             location=CartesianDomain({'x': [0, 1], 'y': 0}),
+#             equation=FixedValue(0.0)),
+#         'gamma3': Condition(
+#             location=CartesianDomain({'x':  1, 'y': [0, 1]}),
+#             equation=FixedValue(0.0)),
+#         'gamma4': Condition(
+#             location=CartesianDomain({'x': 0, 'y': [0, 1]}),
+#             equation=FixedValue(0.0)),
+#         'D': Condition(
+#             input_points=LabelTensor(torch.rand(size=(100, 2)), ['x', 'y']),
+#             equation=my_laplace),
+#         'data': Condition(
+#             input_points=in_,
+#             output_points=out_),
+#         'data2': Condition(
+#             input_points=in2_,
+#             output_points=out2_)
+#     }
+
+#     def poisson_sol(self, pts):
+#         return -(torch.sin(pts.extract(['x']) * torch.pi) *
+#                  torch.sin(pts.extract(['y']) * torch.pi)) / (2 * torch.pi**2)
+
+#     truth_solution = poisson_sol
+
+
+# class myFeature(torch.nn.Module):
+#     """
+#     Feature: sin(x)
+#     """
+
+#     def __init__(self):
+#         super(myFeature, self).__init__()
+
+#     def forward(self, x):
+#         t = (torch.sin(x.extract(['x']) * torch.pi) *
+#              torch.sin(x.extract(['y']) * torch.pi))
+#         return LabelTensor(t, ['sin(x)sin(y)'])
+
+
+# # make the problem
+# poisson_problem = Poisson()
+# model = FeedForward(len(poisson_problem.input_variables),
+#                     len(poisson_problem.output_variables))
+# model_extra_feats = FeedForward(
+#     len(poisson_problem.input_variables) + 1,
+#     len(poisson_problem.output_variables))
+# extra_feats = [myFeature()]
+
+
+# def test_constructor():
+#     PINN(problem=poisson_problem, model=model, extra_features=None)
+#     with pytest.raises(ValueError):
+#         PINN(problem=poisson_problem, model=model, eta='x')
+#         PINN(problem=poisson_problem, model=model, gamma='x')
+
+
+# def test_constructor_extra_feats():
+#     model_extra_feats = FeedForward(
+#         len(poisson_problem.input_variables) + 1,
+#         len(poisson_problem.output_variables))
+#     PINN(problem=poisson_problem,
+#          model=model_extra_feats,
+#          extra_features=extra_feats)
+
+
+# def test_train_cpu():
+#     poisson_problem = Poisson()
+#     boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4']
+#     n = 10
+#     poisson_problem.discretise_domain(n, 'grid', locations=boundaries)
+#     pinn = PINN(problem = poisson_problem, model=model,
+#                 extra_features=None, loss=LpLoss())
+#     trainer = Trainer(solver=pinn, max_epochs=1,
+#                       accelerator='cpu', batch_size=20)
+#     trainer.train()
+
+# def test_log():
+#     poisson_problem.discretise_domain(100)
+#     solver = PINN(problem = poisson_problem, model=model,
+#                 extra_features=None, loss=LpLoss())
+#     trainer = Trainer(solver, max_epochs=2, accelerator='cpu')
+#     trainer.train()
+#     # assert the logged metrics are correct
+#     logged_metrics = sorted(list(trainer.logged_metrics.keys()))
+#     total_metrics = sorted(
+#         list([key + '_loss' for key in poisson_problem.conditions.keys()])
+#         + ['mean_loss'])
+#     assert logged_metrics == total_metrics
+
+# def test_train_restore():
+#     tmpdir = "tests/tmp_restore"
+#     poisson_problem = Poisson()
+#     boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4']
+#     n = 10
+#     poisson_problem.discretise_domain(n, 'grid', locations=boundaries)
 #     pinn = PINN(problem=poisson_problem,
 #                 model=model,
 #                 extra_features=None,
@@ -257,158 +193,169 @@ def test_train_inverse_problem_cpu():
 #                       accelerator='cpu',
 #                       default_root_dir=tmpdir)
 #     trainer.train()
-#     ntrainer = Trainer(solver=pinn, max_epochs=5, accelerator='cpu')
+#     ntrainer = Trainer(solver=pinn, max_epochs=15, accelerator='cpu')
 #     t = ntrainer.train(
-#         ckpt_path=f'{tmpdir}/lightning_logs/version_0/checkpoints/epoch=4-step=10.ckpt')
+#         ckpt_path=f'{tmpdir}/lightning_logs/version_0/'
+#         'checkpoints/epoch=4-step=10.ckpt')
 #     import shutil
 #     shutil.rmtree(tmpdir)
 
 
-def test_train_inverse_problem_load():
-    tmpdir = "tests/tmp_load_inv"
-    poisson_problem = InversePoisson()
-    boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4', 'D']
-    n = 100
-    poisson_problem.discretise_domain(n, 'random', locations=boundaries)
-    pinn = PINN(problem=poisson_problem,
-                model=model,
-                extra_features=None,
-                loss=LpLoss())
-    trainer = Trainer(solver=pinn,
-                      max_epochs=15,
-                      accelerator='cpu',
-                      default_root_dir=tmpdir)
-    trainer.train()
-    new_pinn = PINN.load_from_checkpoint(
-        f'{tmpdir}/lightning_logs/version_0/checkpoints/epoch=14-step=30.ckpt',
-        problem = poisson_problem, model=model)
-    test_pts = CartesianDomain({'x': [0, 1], 'y': [0, 1]}).sample(10)
-    assert new_pinn.forward(test_pts).extract(['u']).shape == (10, 1)
-    assert new_pinn.forward(test_pts).extract(
-        ['u']).shape == pinn.forward(test_pts).extract(['u']).shape
-    torch.testing.assert_close(
-        new_pinn.forward(test_pts).extract(['u']),
-        pinn.forward(test_pts).extract(['u']))
-    import shutil
-    shutil.rmtree(tmpdir)
-
-# # TODO fix asap. Basically sampling few variables
-# # works only if both variables are in a range.
-# # if one is fixed and the other not, this will
-# # not work. This test also needs to be fixed and
-# # insert in test problem not in test pinn.
-# def test_train_cpu_sampling_few_vars():
+# def test_train_load():
+#     tmpdir = "tests/tmp_load"
 #     poisson_problem = Poisson()
-#     boundaries = ['gamma1', 'gamma2', 'gamma3']
+#     boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4']
 #     n = 10
 #     poisson_problem.discretise_domain(n, 'grid', locations=boundaries)
-#     poisson_problem.discretise_domain(n, 'random', locations=['gamma4'], variables=['x'])
-#     poisson_problem.discretise_domain(n, 'random', locations=['gamma4'], variables=['y'])
-#     pinn = PINN(problem = poisson_problem, model=model, extra_features=None, loss=LpLoss())
-#     trainer = Trainer(solver=pinn, kwargs={'max_epochs' : 5, 'accelerator':'cpu'})
+#     pinn = PINN(problem=poisson_problem,
+#                 model=model,
+#                 extra_features=None,
+#                 loss=LpLoss())
+#     trainer = Trainer(solver=pinn,
+#                       max_epochs=15,
+#                       accelerator='cpu',
+#                       default_root_dir=tmpdir)
 #     trainer.train()
+#     new_pinn = PINN.load_from_checkpoint(
+#         f'{tmpdir}/lightning_logs/version_0/checkpoints/epoch=14-step=30.ckpt',
+#         problem = poisson_problem, model=model)
+#     test_pts = CartesianDomain({'x': [0, 1], 'y': [0, 1]}).sample(10)
+#     assert new_pinn.forward(test_pts).extract(['u']).shape == (10, 1)
+#     assert new_pinn.forward(test_pts).extract(
+#         ['u']).shape == pinn.forward(test_pts).extract(['u']).shape
+#     torch.testing.assert_close(
+#         new_pinn.forward(test_pts).extract(['u']),
+#         pinn.forward(test_pts).extract(['u']))
+#     import shutil
+#     shutil.rmtree(tmpdir)
 
-
-def test_train_extra_feats_cpu():
-    poisson_problem = Poisson()
-    boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4']
-    n = 10
-    poisson_problem.discretise_domain(n, 'grid', locations=boundaries)
-    pinn = PINN(problem=poisson_problem,
-                model=model_extra_feats,
-                extra_features=extra_feats)
-    trainer = Trainer(solver=pinn, max_epochs=5, accelerator='cpu')
-    trainer.train()
+# def test_train_inverse_problem_cpu():
+#     poisson_problem = InversePoisson()
+#     boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4', 'D']
+#     n = 100
+#     poisson_problem.discretise_domain(n, 'random', locations=boundaries)
+#     pinn = PINN(problem = poisson_problem, model=model,
+#                 extra_features=None, loss=LpLoss())
+#     trainer = Trainer(solver=pinn, max_epochs=1,
+#                       accelerator='cpu', batch_size=20)
+#     trainer.train()
 
 
-# TODO, fix GitHub actions to run also on GPU
-# def test_train_gpu():
-#     poisson_problem = Poisson()
-#     boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4']
-#     n = 10
-#     poisson_problem.discretise_domain(n, 'grid', locations=boundaries)
-#     pinn = PINN(problem = poisson_problem, model=model, extra_features=None, loss=LpLoss())
-#     trainer = Trainer(solver=pinn, kwargs={'max_epochs' : 5, 'accelerator':'gpu'})
+# # # TODO does not currently work
+# # def test_train_inverse_problem_restore():
+# #     tmpdir = "tests/tmp_restore_inv"
+# #     poisson_problem = InversePoisson()
+# #     boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4', 'D']
+# #     n = 100
+# #     poisson_problem.discretise_domain(n, 'random', locations=boundaries)
+# #     pinn = PINN(problem=poisson_problem,
+# #                 model=model,
+# #                 extra_features=None,
+# #                 loss=LpLoss())
+# #     trainer = Trainer(solver=pinn,
+# #                       max_epochs=5,
+# #                       accelerator='cpu',
+# #                       default_root_dir=tmpdir)
+# #     trainer.train()
+# #     ntrainer = Trainer(solver=pinn, max_epochs=5, accelerator='cpu')
+# #     t = ntrainer.train(
+# #         ckpt_path=f'{tmpdir}/lightning_logs/version_0/checkpoints/epoch=4-step=10.ckpt')
+# #     import shutil
+# #     shutil.rmtree(tmpdir)
+
+
+# def test_train_inverse_problem_load():
+#     tmpdir = "tests/tmp_load_inv"
+#     poisson_problem = InversePoisson()
+#     boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4', 'D']
+#     n = 100
+#     poisson_problem.discretise_domain(n, 'random', locations=boundaries)
+#     pinn = PINN(problem=poisson_problem,
+#                 model=model,
+#                 extra_features=None,
+#                 loss=LpLoss())
+#     trainer = Trainer(solver=pinn,
+#                       max_epochs=15,
+#                       accelerator='cpu',
+#                       default_root_dir=tmpdir)
 #     trainer.train()
+#     new_pinn = PINN.load_from_checkpoint(
+#         f'{tmpdir}/lightning_logs/version_0/checkpoints/epoch=14-step=30.ckpt',
+#         problem = poisson_problem, model=model)
+#     test_pts = CartesianDomain({'x': [0, 1], 'y': [0, 1]}).sample(10)
+#     assert new_pinn.forward(test_pts).extract(['u']).shape == (10, 1)
+#     assert new_pinn.forward(test_pts).extract(
+#         ['u']).shape == pinn.forward(test_pts).extract(['u']).shape
+#     torch.testing.assert_close(
+#         new_pinn.forward(test_pts).extract(['u']),
+#         pinn.forward(test_pts).extract(['u']))
+#     import shutil
+#     shutil.rmtree(tmpdir)
+
+# # # TODO fix asap. Basically sampling few variables
+# # # works only if both variables are in a range.
+# # # if one is fixed and the other not, this will
+# # # not work. This test also needs to be fixed and
+# # # insert in test problem not in test pinn.
+# # def test_train_cpu_sampling_few_vars():
+# #     poisson_problem = Poisson()
+# #     boundaries = ['gamma1', 'gamma2', 'gamma3']
+# #     n = 10
+# #     poisson_problem.discretise_domain(n, 'grid', locations=boundaries)
+# #     poisson_problem.discretise_domain(n, 'random', locations=['gamma4'], variables=['x'])
+# #     poisson_problem.discretise_domain(n, 'random', locations=['gamma4'], variables=['y'])
+# #     pinn = PINN(problem = poisson_problem, model=model, extra_features=None, loss=LpLoss())
+# #     trainer = Trainer(solver=pinn, kwargs={'max_epochs' : 5, 'accelerator':'cpu'})
+# #     trainer.train()
+
 
-# def test_train_gpu(): #TODO fix ASAP
+# def test_train_extra_feats_cpu():
 #     poisson_problem = Poisson()
 #     boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4']
 #     n = 10
 #     poisson_problem.discretise_domain(n, 'grid', locations=boundaries)
-#     poisson_problem.conditions.pop('data') # The input/output pts are allocated on cpu
-#     pinn = PINN(problem = poisson_problem, model=model, extra_features=None, loss=LpLoss())
-#     trainer = Trainer(solver=pinn, kwargs={'max_epochs' : 5, 'accelerator':'gpu'})
+#     pinn = PINN(problem=poisson_problem,
+#                 model=model_extra_feats,
+#                 extra_features=extra_feats)
+#     trainer = Trainer(solver=pinn, max_epochs=5, accelerator='cpu')
 #     trainer.train()
 
-# def test_train_2():
-#     boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4']
-#     n = 10
-#     expected_keys = [[], list(range(0, 50, 3))]
-#     param = [0, 3]
-#     for i, truth_key in zip(param, expected_keys):
-#         pinn = PINN(problem, model)
-#         pinn.discretise_domain(n, 'grid', locations=boundaries)
-#         pinn.discretise_domain(n, 'grid', locations=['D'])
-#         pinn.train(50, save_loss=i)
-#         assert list(pinn.history_loss.keys()) == truth_key
-
-
-# def test_train_extra_feats():
-#     pinn = PINN(problem, model_extra_feat, [myFeature()])
-#     boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4']
-#     n = 10
-#     pinn.discretise_domain(n, 'grid', locations=boundaries)
-#     pinn.discretise_domain(n, 'grid', locations=['D'])
-#     pinn.train(5)
 
+# # TODO, fix GitHub actions to run also on GPU
+# # def test_train_gpu():
+# #     poisson_problem = Poisson()
+# #     boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4']
+# #     n = 10
+# #     poisson_problem.discretise_domain(n, 'grid', locations=boundaries)
+# #     pinn = PINN(problem = poisson_problem, model=model, extra_features=None, loss=LpLoss())
+# #     trainer = Trainer(solver=pinn, kwargs={'max_epochs' : 5, 'accelerator':'gpu'})
+# #     trainer.train()
 
-# def test_train_2_extra_feats():
-#     boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4']
-#     n = 10
-#     expected_keys = [[], list(range(0, 50, 3))]
-#     param = [0, 3]
-#     for i, truth_key in zip(param, expected_keys):
-#         pinn = PINN(problem, model_extra_feat, [myFeature()])
-#         pinn.discretise_domain(n, 'grid', locations=boundaries)
-#         pinn.discretise_domain(n, 'grid', locations=['D'])
-#         pinn.train(50, save_loss=i)
-#         assert list(pinn.history_loss.keys()) == truth_key
-
+# # def test_train_gpu(): #TODO fix ASAP
+# #     poisson_problem = Poisson()
+# #     boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4']
+# #     n = 10
+# #     poisson_problem.discretise_domain(n, 'grid', locations=boundaries)
+# #     poisson_problem.conditions.pop('data') # The input/output pts are allocated on cpu
+# #     pinn = PINN(problem = poisson_problem, model=model, extra_features=None, loss=LpLoss())
+# #     trainer = Trainer(solver=pinn, kwargs={'max_epochs' : 5, 'accelerator':'gpu'})
+# #     trainer.train()
 
-# def test_train_with_optimizer_kwargs():
-#     boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4']
-#     n = 10
-#     expected_keys = [[], list(range(0, 50, 3))]
-#     param = [0, 3]
-#     for i, truth_key in zip(param, expected_keys):
-#         pinn = PINN(problem, model, optimizer_kwargs={'lr' : 0.3})
-#         pinn.discretise_domain(n, 'grid', locations=boundaries)
-#         pinn.discretise_domain(n, 'grid', locations=['D'])
-#         pinn.train(50, save_loss=i)
-#         assert list(pinn.history_loss.keys()) == truth_key
+# # def test_train_2():
+# #     boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4']
+# #     n = 10
+# #     expected_keys = [[], list(range(0, 50, 3))]
+# #     param = [0, 3]
+# #     for i, truth_key in zip(param, expected_keys):
+# #         pinn = PINN(problem, model)
+# #         pinn.discretise_domain(n, 'grid', locations=boundaries)
+# #         pinn.discretise_domain(n, 'grid', locations=['D'])
+# #         pinn.train(50, save_loss=i)
+# #         assert list(pinn.history_loss.keys()) == truth_key
 
 
-# def test_train_with_lr_scheduler():
-#     boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4']
-#     n = 10
-#     expected_keys = [[], list(range(0, 50, 3))]
-#     param = [0, 3]
-#     for i, truth_key in zip(param, expected_keys):
-#         pinn = PINN(
-#             problem,
-#             model,
-#             lr_scheduler_type=torch.optim.lr_scheduler.CyclicLR,
-#             lr_scheduler_kwargs={'base_lr' : 0.1, 'max_lr' : 0.3, 'cycle_momentum': False}
-#         )
-#         pinn.discretise_domain(n, 'grid', locations=boundaries)
-#         pinn.discretise_domain(n, 'grid', locations=['D'])
-#         pinn.train(50, save_loss=i)
-#         assert list(pinn.history_loss.keys()) == truth_key
-
-
-# # def test_train_batch():
-# #     pinn = PINN(problem, model, batch_size=6)
+# # def test_train_extra_feats():
+# #     pinn = PINN(problem, model_extra_feat, [myFeature()])
 # #     boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4']
 # #     n = 10
 # #     pinn.discretise_domain(n, 'grid', locations=boundaries)
@@ -416,34 +363,87 @@ def test_train_extra_feats_cpu():
 # #     pinn.train(5)
 
 
-# # def test_train_batch_2():
+# # def test_train_2_extra_feats():
+# #     boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4']
+# #     n = 10
+# #     expected_keys = [[], list(range(0, 50, 3))]
+# #     param = [0, 3]
+# #     for i, truth_key in zip(param, expected_keys):
+# #         pinn = PINN(problem, model_extra_feat, [myFeature()])
+# #         pinn.discretise_domain(n, 'grid', locations=boundaries)
+# #         pinn.discretise_domain(n, 'grid', locations=['D'])
+# #         pinn.train(50, save_loss=i)
+# #         assert list(pinn.history_loss.keys()) == truth_key
+
+
+# # def test_train_with_optimizer_kwargs():
 # #     boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4']
 # #     n = 10
 # #     expected_keys = [[], list(range(0, 50, 3))]
 # #     param = [0, 3]
 # #     for i, truth_key in zip(param, expected_keys):
-# #         pinn = PINN(problem, model, batch_size=6)
+# #         pinn = PINN(problem, model, optimizer_kwargs={'lr' : 0.3})
 # #         pinn.discretise_domain(n, 'grid', locations=boundaries)
 # #         pinn.discretise_domain(n, 'grid', locations=['D'])
 # #         pinn.train(50, save_loss=i)
 # #         assert list(pinn.history_loss.keys()) == truth_key
 
 
-# if torch.cuda.is_available():
+# # def test_train_with_lr_scheduler():
+# #     boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4']
+# #     n = 10
+# #     expected_keys = [[], list(range(0, 50, 3))]
+# #     param = [0, 3]
+# #     for i, truth_key in zip(param, expected_keys):
+# #         pinn = PINN(
+# #             problem,
+# #             model,
+# #             lr_scheduler_type=torch.optim.lr_scheduler.CyclicLR,
+# #             lr_scheduler_kwargs={'base_lr' : 0.1, 'max_lr' : 0.3, 'cycle_momentum': False}
+# #         )
+# #         pinn.discretise_domain(n, 'grid', locations=boundaries)
+# #         pinn.discretise_domain(n, 'grid', locations=['D'])
+# #         pinn.train(50, save_loss=i)
+# #         assert list(pinn.history_loss.keys()) == truth_key
 
-#     # def test_gpu_train():
-#     #     pinn = PINN(problem, model, batch_size=20, device='cuda')
-#     #     boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4']
-#     #     n = 100
-#     #     pinn.discretise_domain(n, 'grid', locations=boundaries)
-#     #     pinn.discretise_domain(n, 'grid', locations=['D'])
-#     #     pinn.train(5)
 
-#     def test_gpu_train_nobatch():
-#         pinn = PINN(problem, model, batch_size=None, device='cuda')
-#         boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4']
-#         n = 100
-#         pinn.discretise_domain(n, 'grid', locations=boundaries)
-#         pinn.discretise_domain(n, 'grid', locations=['D'])
-#         pinn.train(5)
+# # # def test_train_batch():
+# # #     pinn = PINN(problem, model, batch_size=6)
+# # #     boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4']
+# # #     n = 10
+# # #     pinn.discretise_domain(n, 'grid', locations=boundaries)
+# # #     pinn.discretise_domain(n, 'grid', locations=['D'])
+# # #     pinn.train(5)
+
+
+# # # def test_train_batch_2():
+# # #     boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4']
+# # #     n = 10
+# # #     expected_keys = [[], list(range(0, 50, 3))]
+# # #     param = [0, 3]
+# # #     for i, truth_key in zip(param, expected_keys):
+# # #         pinn = PINN(problem, model, batch_size=6)
+# # #         pinn.discretise_domain(n, 'grid', locations=boundaries)
+# # #         pinn.discretise_domain(n, 'grid', locations=['D'])
+# # #         pinn.train(50, save_loss=i)
+# # #         assert list(pinn.history_loss.keys()) == truth_key
+
+
+# # if torch.cuda.is_available():
+
+# #     # def test_gpu_train():
+# #     #     pinn = PINN(problem, model, batch_size=20, device='cuda')
+# #     #     boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4']
+# #     #     n = 100
+# #     #     pinn.discretise_domain(n, 'grid', locations=boundaries)
+# #     #     pinn.discretise_domain(n, 'grid', locations=['D'])
+# #     #     pinn.train(5)
+
+# #     def test_gpu_train_nobatch():
+# #         pinn = PINN(problem, model, batch_size=None, device='cuda')
+# #         boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4']
+# #         n = 100
+# #         pinn.discretise_domain(n, 'grid', locations=boundaries)
+# #         pinn.discretise_domain(n, 'grid', locations=['D'])
+# #         pinn.train(5)
 
diff --git a/tests/test_solvers/test_rom_solver.py b/tests/test_solvers/test_rom_solver.py
index 2a54d6504..00dc69ff6 100644
--- a/tests/test_solvers/test_rom_solver.py
+++ b/tests/test_solvers/test_rom_solver.py
@@ -6,100 +6,100 @@
 from pina.solvers import ReducedOrderModelSolver
 from pina.trainer import Trainer
 from pina.model import FeedForward
-from pina.loss.loss_interface import LpLoss
+from pina.loss import LpLoss
 
 
-class NeuralOperatorProblem(AbstractProblem):
-    input_variables = ['u_0', 'u_1']
-    output_variables = [f'u_{i}' for i in range(100)]
-    conditions = {'data' : Condition(input_points=
-                                     LabelTensor(torch.rand(10, 2),
-                                                 input_variables), 
-                                     output_points=
-                                     LabelTensor(torch.rand(10, 100),
-                                                 output_variables))}
+# class NeuralOperatorProblem(AbstractProblem):
+#     input_variables = ['u_0', 'u_1']
+#     output_variables = [f'u_{i}' for i in range(100)]
+#     conditions = {'data' : Condition(input_points=
+#                                      LabelTensor(torch.rand(10, 2),
+#                                                  input_variables), 
+#                                      output_points=
+#                                      LabelTensor(torch.rand(10, 100),
+#                                                  output_variables))}
 
 
-# make the problem + extra feats
-class AE(torch.nn.Module):
-    def __init__(self, input_dimensions, rank):
-        super().__init__()
-        self.encode = FeedForward(input_dimensions, rank, layers=[input_dimensions//4])
-        self.decode = FeedForward(rank, input_dimensions, layers=[input_dimensions//4])
-class AE_missing_encode(torch.nn.Module):
-    def __init__(self, input_dimensions, rank):
-        super().__init__()
-        self.encode = FeedForward(input_dimensions, rank, layers=[input_dimensions//4])
-class AE_missing_decode(torch.nn.Module):
-    def __init__(self, input_dimensions, rank):
-        super().__init__()
-        self.decode = FeedForward(rank, input_dimensions, layers=[input_dimensions//4])
+# # make the problem + extra feats
+# class AE(torch.nn.Module):
+#     def __init__(self, input_dimensions, rank):
+#         super().__init__()
+#         self.encode = FeedForward(input_dimensions, rank, layers=[input_dimensions//4])
+#         self.decode = FeedForward(rank, input_dimensions, layers=[input_dimensions//4])
+# class AE_missing_encode(torch.nn.Module):
+#     def __init__(self, input_dimensions, rank):
+#         super().__init__()
+#         self.encode = FeedForward(input_dimensions, rank, layers=[input_dimensions//4])
+# class AE_missing_decode(torch.nn.Module):
+#     def __init__(self, input_dimensions, rank):
+#         super().__init__()
+#         self.decode = FeedForward(rank, input_dimensions, layers=[input_dimensions//4])
 
-rank = 10
-problem = NeuralOperatorProblem()
-interpolation_net = FeedForward(len(problem.input_variables),
-                                rank)
-reduction_net = AE(len(problem.output_variables), rank)
+# rank = 10
+# problem = NeuralOperatorProblem()
+# interpolation_net = FeedForward(len(problem.input_variables),
+#                                 rank)
+# reduction_net = AE(len(problem.output_variables), rank)
 
-def test_constructor():
-    ReducedOrderModelSolver(problem=problem,reduction_network=reduction_net,
-    interpolation_network=interpolation_net)
-    with pytest.raises(SyntaxError):
-        ReducedOrderModelSolver(problem=problem,
-                     reduction_network=AE_missing_encode(
-                         len(problem.output_variables), rank),
-                     interpolation_network=interpolation_net)
-        ReducedOrderModelSolver(problem=problem,
-                     reduction_network=AE_missing_decode(
-                         len(problem.output_variables), rank),
-                     interpolation_network=interpolation_net)
+# def test_constructor():
+#     ReducedOrderModelSolver(problem=problem,reduction_network=reduction_net,
+#     interpolation_network=interpolation_net)
+#     with pytest.raises(SyntaxError):
+#         ReducedOrderModelSolver(problem=problem,
+#                      reduction_network=AE_missing_encode(
+#                          len(problem.output_variables), rank),
+#                      interpolation_network=interpolation_net)
+#         ReducedOrderModelSolver(problem=problem,
+#                      reduction_network=AE_missing_decode(
+#                          len(problem.output_variables), rank),
+#                      interpolation_network=interpolation_net)
 
 
-def test_train_cpu():
-    solver = ReducedOrderModelSolver(problem = problem,reduction_network=reduction_net,
-    interpolation_network=interpolation_net, loss=LpLoss())
-    trainer = Trainer(solver=solver, max_epochs=3, accelerator='cpu', batch_size=20)
-    trainer.train()
+# def test_train_cpu():
+#     solver = ReducedOrderModelSolver(problem = problem,reduction_network=reduction_net,
+#     interpolation_network=interpolation_net, loss=LpLoss())
+#     trainer = Trainer(solver=solver, max_epochs=3, accelerator='cpu', batch_size=20)
+#     trainer.train()
 
 
-def test_train_restore():
-    tmpdir = "tests/tmp_restore"
-    solver = ReducedOrderModelSolver(problem=problem,
-               reduction_network=reduction_net,
-               interpolation_network=interpolation_net,
-               loss=LpLoss())
-    trainer = Trainer(solver=solver,
-                      max_epochs=5,
-                      accelerator='cpu',
-                      default_root_dir=tmpdir)
-    trainer.train()
-    ntrainer = Trainer(solver=solver, max_epochs=15, accelerator='cpu')
-    t = ntrainer.train(
-        ckpt_path=f'{tmpdir}/lightning_logs/version_0/checkpoints/epoch=4-step=5.ckpt')
-    import shutil
-    shutil.rmtree(tmpdir)
+# def test_train_restore():
+#     tmpdir = "tests/tmp_restore"
+#     solver = ReducedOrderModelSolver(problem=problem,
+#                reduction_network=reduction_net,
+#                interpolation_network=interpolation_net,
+#                loss=LpLoss())
+#     trainer = Trainer(solver=solver,
+#                       max_epochs=5,
+#                       accelerator='cpu',
+#                       default_root_dir=tmpdir)
+#     trainer.train()
+#     ntrainer = Trainer(solver=solver, max_epochs=15, accelerator='cpu')
+#     t = ntrainer.train(
+#         ckpt_path=f'{tmpdir}/lightning_logs/version_0/checkpoints/epoch=4-step=5.ckpt')
+#     import shutil
+#     shutil.rmtree(tmpdir)
 
 
-def test_train_load():
-    tmpdir = "tests/tmp_load"
-    solver = ReducedOrderModelSolver(problem=problem,
-               reduction_network=reduction_net,
-               interpolation_network=interpolation_net,
-               loss=LpLoss())
-    trainer = Trainer(solver=solver,
-                      max_epochs=15,
-                      accelerator='cpu',
-                      default_root_dir=tmpdir)
-    trainer.train()
-    new_solver = ReducedOrderModelSolver.load_from_checkpoint(
-        f'{tmpdir}/lightning_logs/version_0/checkpoints/epoch=14-step=15.ckpt',
-        problem = problem,reduction_network=reduction_net,
-        interpolation_network=interpolation_net)
-    test_pts = LabelTensor(torch.rand(20, 2), problem.input_variables)
-    assert new_solver.forward(test_pts).shape == (20, 100)
-    assert new_solver.forward(test_pts).shape == solver.forward(test_pts).shape
-    torch.testing.assert_close(
-        new_solver.forward(test_pts),
-        solver.forward(test_pts))
-    import shutil
-    shutil.rmtree(tmpdir)
\ No newline at end of file
+# def test_train_load():
+#     tmpdir = "tests/tmp_load"
+#     solver = ReducedOrderModelSolver(problem=problem,
+#                reduction_network=reduction_net,
+#                interpolation_network=interpolation_net,
+#                loss=LpLoss())
+#     trainer = Trainer(solver=solver,
+#                       max_epochs=15,
+#                       accelerator='cpu',
+#                       default_root_dir=tmpdir)
+#     trainer.train()
+#     new_solver = ReducedOrderModelSolver.load_from_checkpoint(
+#         f'{tmpdir}/lightning_logs/version_0/checkpoints/epoch=14-step=15.ckpt',
+#         problem = problem,reduction_network=reduction_net,
+#         interpolation_network=interpolation_net)
+#     test_pts = LabelTensor(torch.rand(20, 2), problem.input_variables)
+#     assert new_solver.forward(test_pts).shape == (20, 100)
+#     assert new_solver.forward(test_pts).shape == solver.forward(test_pts).shape
+#     torch.testing.assert_close(
+#         new_solver.forward(test_pts),
+#         solver.forward(test_pts))
+#     import shutil
+#     shutil.rmtree(tmpdir)
\ No newline at end of file
diff --git a/tests/test_solvers/test_sapinn.py b/tests/test_solvers/test_sapinn.py
index 5e1eee6f6..d4249d886 100644
--- a/tests/test_solvers/test_sapinn.py
+++ b/tests/test_solvers/test_sapinn.py
@@ -8,246 +8,182 @@
 from pina.solvers import SAPINN as PINN
 from pina.trainer import Trainer
 from pina.model import FeedForward
-from pina.equation.equation import Equation
+from pina.equation import Equation
 from pina.equation.equation_factory import FixedValue
-from pina.loss.loss_interface import LpLoss
-
-
-def laplace_equation(input_, output_):
-    force_term = (torch.sin(input_.extract(['x']) * torch.pi) *
-                  torch.sin(input_.extract(['y']) * torch.pi))
-    delta_u = laplacian(output_.extract(['u']), input_)
-    return delta_u - force_term
-
-
-my_laplace = Equation(laplace_equation)
-in_ = LabelTensor(torch.tensor([[0., 1.]]), ['x', 'y'])
-out_ = LabelTensor(torch.tensor([[0.]]), ['u'])
-in2_ = LabelTensor(torch.rand(60, 2), ['x', 'y'])
-out2_ = LabelTensor(torch.rand(60, 1), ['u'])
-
-
-class InversePoisson(SpatialProblem, InverseProblem):
-    '''
-    Problem definition for the Poisson equation.
-    '''
-    output_variables = ['u']
-    x_min = -2
-    x_max = 2
-    y_min = -2
-    y_max = 2
-    data_input = LabelTensor(torch.rand(10, 2), ['x', 'y'])
-    data_output = LabelTensor(torch.rand(10, 1), ['u'])
-    spatial_domain = CartesianDomain({'x': [x_min, x_max], 'y': [y_min, y_max]})
-    # define the ranges for the parameters
-    unknown_parameter_domain = CartesianDomain({'mu1': [-1, 1], 'mu2': [-1, 1]})
-
-    def laplace_equation(input_, output_, params_):
-        '''
-        Laplace equation with a force term.
-        '''
-        force_term = torch.exp(
-                - 2*(input_.extract(['x']) - params_['mu1'])**2
-                - 2*(input_.extract(['y']) - params_['mu2'])**2)
-        delta_u = laplacian(output_, input_, components=['u'], d=['x', 'y'])
-
-        return delta_u - force_term
-
-    # define the conditions for the loss (boundary conditions, equation, data)
-    conditions = {
-        'gamma1': Condition(location=CartesianDomain({'x': [x_min, x_max],
-            'y':  y_max}),
-            equation=FixedValue(0.0, components=['u'])),
-        'gamma2': Condition(location=CartesianDomain(
-            {'x': [x_min, x_max], 'y': y_min
-            }),
-            equation=FixedValue(0.0, components=['u'])),
-        'gamma3': Condition(location=CartesianDomain(
-            {'x':  x_max, 'y': [y_min, y_max]
-            }),
-            equation=FixedValue(0.0, components=['u'])),
-        'gamma4': Condition(location=CartesianDomain(
-            {'x': x_min, 'y': [y_min, y_max]
-            }),
-            equation=FixedValue(0.0, components=['u'])),
-        'D': Condition(location=CartesianDomain(
-            {'x': [x_min, x_max], 'y': [y_min, y_max]
-            }),
-        equation=Equation(laplace_equation)),
-        'data': Condition(input_points=data_input.extract(['x', 'y']),
-                          output_points=data_output)
-    }
-
-
-class Poisson(SpatialProblem):
-    output_variables = ['u']
-    spatial_domain = CartesianDomain({'x': [0, 1], 'y': [0, 1]})
-
-    conditions = {
-        'gamma1': Condition(
-            location=CartesianDomain({'x': [0, 1], 'y':  1}),
-            equation=FixedValue(0.0)),
-        'gamma2': Condition(
-            location=CartesianDomain({'x': [0, 1], 'y': 0}),
-            equation=FixedValue(0.0)),
-        'gamma3': Condition(
-            location=CartesianDomain({'x':  1, 'y': [0, 1]}),
-            equation=FixedValue(0.0)),
-        'gamma4': Condition(
-            location=CartesianDomain({'x': 0, 'y': [0, 1]}),
-            equation=FixedValue(0.0)),
-        'D': Condition(
-            input_points=LabelTensor(torch.rand(size=(100, 2)), ['x', 'y']),
-            equation=my_laplace),
-        'data': Condition(
-            input_points=in_,
-            output_points=out_),
-        'data2': Condition(
-            input_points=in2_,
-            output_points=out2_)
-    }
-
-    def poisson_sol(self, pts):
-        return -(torch.sin(pts.extract(['x']) * torch.pi) *
-                 torch.sin(pts.extract(['y']) * torch.pi)) / (2 * torch.pi**2)
-
-    truth_solution = poisson_sol
-
-
-class myFeature(torch.nn.Module):
-    """
-    Feature: sin(x)
-    """
-
-    def __init__(self):
-        super(myFeature, self).__init__()
-
-    def forward(self, x):
-        t = (torch.sin(x.extract(['x']) * torch.pi) *
-             torch.sin(x.extract(['y']) * torch.pi))
-        return LabelTensor(t, ['sin(x)sin(y)'])
-
-
-# make the problem
-poisson_problem = Poisson()
-model = FeedForward(len(poisson_problem.input_variables),
-                    len(poisson_problem.output_variables))
-model_extra_feats = FeedForward(
-    len(poisson_problem.input_variables) + 1,
-    len(poisson_problem.output_variables))
-extra_feats = [myFeature()]
-
-
-def test_constructor():
-    PINN(problem=poisson_problem, model=model, extra_features=None)
-    with pytest.raises(ValueError):
-        PINN(problem=poisson_problem, model=model, extra_features=None,
-             weights_function=1)
-
-
-def test_constructor_extra_feats():
-    model_extra_feats = FeedForward(
-        len(poisson_problem.input_variables) + 1,
-        len(poisson_problem.output_variables))
-    PINN(problem=poisson_problem,
-         model=model_extra_feats,
-         extra_features=extra_feats)
-
-
-def test_train_cpu():
-    poisson_problem = Poisson()
-    boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4']
-    n = 10
-    poisson_problem.discretise_domain(n, 'grid', locations=boundaries)
-    pinn = PINN(problem = poisson_problem, model=model,
-                extra_features=None, loss=LpLoss())
-    trainer = Trainer(solver=pinn, max_epochs=1,
-                      accelerator='cpu', batch_size=20)
-    trainer.train()
-
-def test_log():
-    poisson_problem.discretise_domain(100)
-    solver = PINN(problem = poisson_problem, model=model,
-                extra_features=None, loss=LpLoss())
-    trainer = Trainer(solver, max_epochs=2, accelerator='cpu')
-    trainer.train()
-    # assert the logged metrics are correct
-    logged_metrics = sorted(list(trainer.logged_metrics.keys()))
-    total_metrics = sorted(
-        list([key + '_loss' for key in poisson_problem.conditions.keys()])
-        + ['mean_loss'])
-    assert logged_metrics == total_metrics
-
-def test_train_restore():
-    tmpdir = "tests/tmp_restore"
-    poisson_problem = Poisson()
-    boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4']
-    n = 10
-    poisson_problem.discretise_domain(n, 'grid', locations=boundaries)
-    pinn = PINN(problem=poisson_problem,
-                model=model,
-                extra_features=None,
-                loss=LpLoss())
-    trainer = Trainer(solver=pinn,
-                      max_epochs=5,
-                      accelerator='cpu',
-                      default_root_dir=tmpdir)
-    trainer.train()
-    ntrainer = Trainer(solver=pinn, max_epochs=15, accelerator='cpu')
-    t = ntrainer.train(
-        ckpt_path=f'{tmpdir}/lightning_logs/version_0/'
-        'checkpoints/epoch=4-step=10.ckpt')
-    import shutil
-    shutil.rmtree(tmpdir)
-
-
-def test_train_load():
-    tmpdir = "tests/tmp_load"
-    poisson_problem = Poisson()
-    boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4']
-    n = 10
-    poisson_problem.discretise_domain(n, 'grid', locations=boundaries)
-    pinn = PINN(problem=poisson_problem,
-                model=model,
-                extra_features=None,
-                loss=LpLoss())
-    trainer = Trainer(solver=pinn,
-                      max_epochs=15,
-                      accelerator='cpu',
-                      default_root_dir=tmpdir)
-    trainer.train()
-    new_pinn = PINN.load_from_checkpoint(
-        f'{tmpdir}/lightning_logs/version_0/checkpoints/epoch=14-step=30.ckpt',
-        problem = poisson_problem, model=model)
-    test_pts = CartesianDomain({'x': [0, 1], 'y': [0, 1]}).sample(10)
-    assert new_pinn.forward(test_pts).extract(['u']).shape == (10, 1)
-    assert new_pinn.forward(test_pts).extract(
-        ['u']).shape == pinn.forward(test_pts).extract(['u']).shape
-    torch.testing.assert_close(
-        new_pinn.forward(test_pts).extract(['u']),
-        pinn.forward(test_pts).extract(['u']))
-    import shutil
-    shutil.rmtree(tmpdir)
-
-def test_train_inverse_problem_cpu():
-    poisson_problem = InversePoisson()
-    boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4', 'D']
-    n = 100
-    poisson_problem.discretise_domain(n, 'random', locations=boundaries)
-    pinn = PINN(problem = poisson_problem, model=model,
-                extra_features=None, loss=LpLoss())
-    trainer = Trainer(solver=pinn, max_epochs=1,
-                      accelerator='cpu', batch_size=20)
-    trainer.train()
-
-
-# # TODO does not currently work
-# def test_train_inverse_problem_restore():
-#     tmpdir = "tests/tmp_restore_inv"
-#     poisson_problem = InversePoisson()
-#     boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4', 'D']
-#     n = 100
-#     poisson_problem.discretise_domain(n, 'random', locations=boundaries)
+from pina.loss import LpLoss
+
+
+# def laplace_equation(input_, output_):
+#     force_term = (torch.sin(input_.extract(['x']) * torch.pi) *
+#                   torch.sin(input_.extract(['y']) * torch.pi))
+#     delta_u = laplacian(output_.extract(['u']), input_)
+#     return delta_u - force_term
+
+
+# my_laplace = Equation(laplace_equation)
+# in_ = LabelTensor(torch.tensor([[0., 1.]]), ['x', 'y'])
+# out_ = LabelTensor(torch.tensor([[0.]]), ['u'])
+# in2_ = LabelTensor(torch.rand(60, 2), ['x', 'y'])
+# out2_ = LabelTensor(torch.rand(60, 1), ['u'])
+
+
+# class InversePoisson(SpatialProblem, InverseProblem):
+#     '''
+#     Problem definition for the Poisson equation.
+#     '''
+#     output_variables = ['u']
+#     x_min = -2
+#     x_max = 2
+#     y_min = -2
+#     y_max = 2
+#     data_input = LabelTensor(torch.rand(10, 2), ['x', 'y'])
+#     data_output = LabelTensor(torch.rand(10, 1), ['u'])
+#     spatial_domain = CartesianDomain({'x': [x_min, x_max], 'y': [y_min, y_max]})
+#     # define the ranges for the parameters
+#     unknown_parameter_domain = CartesianDomain({'mu1': [-1, 1], 'mu2': [-1, 1]})
+
+#     def laplace_equation(input_, output_, params_):
+#         '''
+#         Laplace equation with a force term.
+#         '''
+#         force_term = torch.exp(
+#                 - 2*(input_.extract(['x']) - params_['mu1'])**2
+#                 - 2*(input_.extract(['y']) - params_['mu2'])**2)
+#         delta_u = laplacian(output_, input_, components=['u'], d=['x', 'y'])
+
+#         return delta_u - force_term
+
+#     # define the conditions for the loss (boundary conditions, equation, data)
+#     conditions = {
+#         'gamma1': Condition(location=CartesianDomain({'x': [x_min, x_max],
+#             'y':  y_max}),
+#             equation=FixedValue(0.0, components=['u'])),
+#         'gamma2': Condition(location=CartesianDomain(
+#             {'x': [x_min, x_max], 'y': y_min
+#             }),
+#             equation=FixedValue(0.0, components=['u'])),
+#         'gamma3': Condition(location=CartesianDomain(
+#             {'x':  x_max, 'y': [y_min, y_max]
+#             }),
+#             equation=FixedValue(0.0, components=['u'])),
+#         'gamma4': Condition(location=CartesianDomain(
+#             {'x': x_min, 'y': [y_min, y_max]
+#             }),
+#             equation=FixedValue(0.0, components=['u'])),
+#         'D': Condition(location=CartesianDomain(
+#             {'x': [x_min, x_max], 'y': [y_min, y_max]
+#             }),
+#         equation=Equation(laplace_equation)),
+#         'data': Condition(input_points=data_input.extract(['x', 'y']),
+#                           output_points=data_output)
+#     }
+
+
+# class Poisson(SpatialProblem):
+#     output_variables = ['u']
+#     spatial_domain = CartesianDomain({'x': [0, 1], 'y': [0, 1]})
+
+#     conditions = {
+#         'gamma1': Condition(
+#             location=CartesianDomain({'x': [0, 1], 'y':  1}),
+#             equation=FixedValue(0.0)),
+#         'gamma2': Condition(
+#             location=CartesianDomain({'x': [0, 1], 'y': 0}),
+#             equation=FixedValue(0.0)),
+#         'gamma3': Condition(
+#             location=CartesianDomain({'x':  1, 'y': [0, 1]}),
+#             equation=FixedValue(0.0)),
+#         'gamma4': Condition(
+#             location=CartesianDomain({'x': 0, 'y': [0, 1]}),
+#             equation=FixedValue(0.0)),
+#         'D': Condition(
+#             input_points=LabelTensor(torch.rand(size=(100, 2)), ['x', 'y']),
+#             equation=my_laplace),
+#         'data': Condition(
+#             input_points=in_,
+#             output_points=out_),
+#         'data2': Condition(
+#             input_points=in2_,
+#             output_points=out2_)
+#     }
+
+#     def poisson_sol(self, pts):
+#         return -(torch.sin(pts.extract(['x']) * torch.pi) *
+#                  torch.sin(pts.extract(['y']) * torch.pi)) / (2 * torch.pi**2)
+
+#     truth_solution = poisson_sol
+
+
+# class myFeature(torch.nn.Module):
+#     """
+#     Feature: sin(x)
+#     """
+
+#     def __init__(self):
+#         super(myFeature, self).__init__()
+
+#     def forward(self, x):
+#         t = (torch.sin(x.extract(['x']) * torch.pi) *
+#              torch.sin(x.extract(['y']) * torch.pi))
+#         return LabelTensor(t, ['sin(x)sin(y)'])
+
+
+# # make the problem
+# poisson_problem = Poisson()
+# model = FeedForward(len(poisson_problem.input_variables),
+#                     len(poisson_problem.output_variables))
+# model_extra_feats = FeedForward(
+#     len(poisson_problem.input_variables) + 1,
+#     len(poisson_problem.output_variables))
+# extra_feats = [myFeature()]
+
+
+# def test_constructor():
+#     PINN(problem=poisson_problem, model=model, extra_features=None)
+#     with pytest.raises(ValueError):
+#         PINN(problem=poisson_problem, model=model, extra_features=None,
+#              weights_function=1)
+
+
+# def test_constructor_extra_feats():
+#     model_extra_feats = FeedForward(
+#         len(poisson_problem.input_variables) + 1,
+#         len(poisson_problem.output_variables))
+#     PINN(problem=poisson_problem,
+#          model=model_extra_feats,
+#          extra_features=extra_feats)
+
+
+# def test_train_cpu():
+#     poisson_problem = Poisson()
+#     boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4']
+#     n = 10
+#     poisson_problem.discretise_domain(n, 'grid', locations=boundaries)
+#     pinn = PINN(problem = poisson_problem, model=model,
+#                 extra_features=None, loss=LpLoss())
+#     trainer = Trainer(solver=pinn, max_epochs=1,
+#                       accelerator='cpu', batch_size=20)
+#     trainer.train()
+
+# def test_log():
+#     poisson_problem.discretise_domain(100)
+#     solver = PINN(problem = poisson_problem, model=model,
+#                 extra_features=None, loss=LpLoss())
+#     trainer = Trainer(solver, max_epochs=2, accelerator='cpu')
+#     trainer.train()
+#     # assert the logged metrics are correct
+#     logged_metrics = sorted(list(trainer.logged_metrics.keys()))
+#     total_metrics = sorted(
+#         list([key + '_loss' for key in poisson_problem.conditions.keys()])
+#         + ['mean_loss'])
+#     assert logged_metrics == total_metrics
+
+# def test_train_restore():
+#     tmpdir = "tests/tmp_restore"
+#     poisson_problem = Poisson()
+#     boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4']
+#     n = 10
+#     poisson_problem.discretise_domain(n, 'grid', locations=boundaries)
 #     pinn = PINN(problem=poisson_problem,
 #                 model=model,
 #                 extra_features=None,
@@ -257,158 +193,169 @@ def test_train_inverse_problem_cpu():
 #                       accelerator='cpu',
 #                       default_root_dir=tmpdir)
 #     trainer.train()
-#     ntrainer = Trainer(solver=pinn, max_epochs=5, accelerator='cpu')
+#     ntrainer = Trainer(solver=pinn, max_epochs=15, accelerator='cpu')
 #     t = ntrainer.train(
-#         ckpt_path=f'{tmpdir}/lightning_logs/version_0/checkpoints/epoch=4-step=10.ckpt')
+#         ckpt_path=f'{tmpdir}/lightning_logs/version_0/'
+#         'checkpoints/epoch=4-step=10.ckpt')
 #     import shutil
 #     shutil.rmtree(tmpdir)
 
 
-def test_train_inverse_problem_load():
-    tmpdir = "tests/tmp_load_inv"
-    poisson_problem = InversePoisson()
-    boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4', 'D']
-    n = 100
-    poisson_problem.discretise_domain(n, 'random', locations=boundaries)
-    pinn = PINN(problem=poisson_problem,
-                model=model,
-                extra_features=None,
-                loss=LpLoss())
-    trainer = Trainer(solver=pinn,
-                      max_epochs=15,
-                      accelerator='cpu',
-                      default_root_dir=tmpdir)
-    trainer.train()
-    new_pinn = PINN.load_from_checkpoint(
-        f'{tmpdir}/lightning_logs/version_0/checkpoints/epoch=14-step=30.ckpt',
-        problem = poisson_problem, model=model)
-    test_pts = CartesianDomain({'x': [0, 1], 'y': [0, 1]}).sample(10)
-    assert new_pinn.forward(test_pts).extract(['u']).shape == (10, 1)
-    assert new_pinn.forward(test_pts).extract(
-        ['u']).shape == pinn.forward(test_pts).extract(['u']).shape
-    torch.testing.assert_close(
-        new_pinn.forward(test_pts).extract(['u']),
-        pinn.forward(test_pts).extract(['u']))
-    import shutil
-    shutil.rmtree(tmpdir)
-
-# # TODO fix asap. Basically sampling few variables
-# # works only if both variables are in a range.
-# # if one is fixed and the other not, this will
-# # not work. This test also needs to be fixed and
-# # insert in test problem not in test pinn.
-# def test_train_cpu_sampling_few_vars():
+# def test_train_load():
+#     tmpdir = "tests/tmp_load"
 #     poisson_problem = Poisson()
-#     boundaries = ['gamma1', 'gamma2', 'gamma3']
+#     boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4']
 #     n = 10
 #     poisson_problem.discretise_domain(n, 'grid', locations=boundaries)
-#     poisson_problem.discretise_domain(n, 'random', locations=['gamma4'], variables=['x'])
-#     poisson_problem.discretise_domain(n, 'random', locations=['gamma4'], variables=['y'])
-#     pinn = PINN(problem = poisson_problem, model=model, extra_features=None, loss=LpLoss())
-#     trainer = Trainer(solver=pinn, kwargs={'max_epochs' : 5, 'accelerator':'cpu'})
+#     pinn = PINN(problem=poisson_problem,
+#                 model=model,
+#                 extra_features=None,
+#                 loss=LpLoss())
+#     trainer = Trainer(solver=pinn,
+#                       max_epochs=15,
+#                       accelerator='cpu',
+#                       default_root_dir=tmpdir)
 #     trainer.train()
+#     new_pinn = PINN.load_from_checkpoint(
+#         f'{tmpdir}/lightning_logs/version_0/checkpoints/epoch=14-step=30.ckpt',
+#         problem = poisson_problem, model=model)
+#     test_pts = CartesianDomain({'x': [0, 1], 'y': [0, 1]}).sample(10)
+#     assert new_pinn.forward(test_pts).extract(['u']).shape == (10, 1)
+#     assert new_pinn.forward(test_pts).extract(
+#         ['u']).shape == pinn.forward(test_pts).extract(['u']).shape
+#     torch.testing.assert_close(
+#         new_pinn.forward(test_pts).extract(['u']),
+#         pinn.forward(test_pts).extract(['u']))
+#     import shutil
+#     shutil.rmtree(tmpdir)
 
-
-def test_train_extra_feats_cpu():
-    poisson_problem = Poisson()
-    boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4']
-    n = 10
-    poisson_problem.discretise_domain(n, 'grid', locations=boundaries)
-    pinn = PINN(problem=poisson_problem,
-                model=model_extra_feats,
-                extra_features=extra_feats)
-    trainer = Trainer(solver=pinn, max_epochs=5, accelerator='cpu')
-    trainer.train()
+# def test_train_inverse_problem_cpu():
+#     poisson_problem = InversePoisson()
+#     boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4', 'D']
+#     n = 100
+#     poisson_problem.discretise_domain(n, 'random', locations=boundaries)
+#     pinn = PINN(problem = poisson_problem, model=model,
+#                 extra_features=None, loss=LpLoss())
+#     trainer = Trainer(solver=pinn, max_epochs=1,
+#                       accelerator='cpu', batch_size=20)
+#     trainer.train()
 
 
-# TODO, fix GitHub actions to run also on GPU
-# def test_train_gpu():
-#     poisson_problem = Poisson()
-#     boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4']
-#     n = 10
-#     poisson_problem.discretise_domain(n, 'grid', locations=boundaries)
-#     pinn = PINN(problem = poisson_problem, model=model, extra_features=None, loss=LpLoss())
-#     trainer = Trainer(solver=pinn, kwargs={'max_epochs' : 5, 'accelerator':'gpu'})
+# # # TODO does not currently work
+# # def test_train_inverse_problem_restore():
+# #     tmpdir = "tests/tmp_restore_inv"
+# #     poisson_problem = InversePoisson()
+# #     boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4', 'D']
+# #     n = 100
+# #     poisson_problem.discretise_domain(n, 'random', locations=boundaries)
+# #     pinn = PINN(problem=poisson_problem,
+# #                 model=model,
+# #                 extra_features=None,
+# #                 loss=LpLoss())
+# #     trainer = Trainer(solver=pinn,
+# #                       max_epochs=5,
+# #                       accelerator='cpu',
+# #                       default_root_dir=tmpdir)
+# #     trainer.train()
+# #     ntrainer = Trainer(solver=pinn, max_epochs=5, accelerator='cpu')
+# #     t = ntrainer.train(
+# #         ckpt_path=f'{tmpdir}/lightning_logs/version_0/checkpoints/epoch=4-step=10.ckpt')
+# #     import shutil
+# #     shutil.rmtree(tmpdir)
+
+
+# def test_train_inverse_problem_load():
+#     tmpdir = "tests/tmp_load_inv"
+#     poisson_problem = InversePoisson()
+#     boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4', 'D']
+#     n = 100
+#     poisson_problem.discretise_domain(n, 'random', locations=boundaries)
+#     pinn = PINN(problem=poisson_problem,
+#                 model=model,
+#                 extra_features=None,
+#                 loss=LpLoss())
+#     trainer = Trainer(solver=pinn,
+#                       max_epochs=15,
+#                       accelerator='cpu',
+#                       default_root_dir=tmpdir)
 #     trainer.train()
+#     new_pinn = PINN.load_from_checkpoint(
+#         f'{tmpdir}/lightning_logs/version_0/checkpoints/epoch=14-step=30.ckpt',
+#         problem = poisson_problem, model=model)
+#     test_pts = CartesianDomain({'x': [0, 1], 'y': [0, 1]}).sample(10)
+#     assert new_pinn.forward(test_pts).extract(['u']).shape == (10, 1)
+#     assert new_pinn.forward(test_pts).extract(
+#         ['u']).shape == pinn.forward(test_pts).extract(['u']).shape
+#     torch.testing.assert_close(
+#         new_pinn.forward(test_pts).extract(['u']),
+#         pinn.forward(test_pts).extract(['u']))
+#     import shutil
+#     shutil.rmtree(tmpdir)
+
+# # # TODO fix asap. Basically sampling few variables
+# # # works only if both variables are in a range.
+# # # if one is fixed and the other not, this will
+# # # not work. This test also needs to be fixed and
+# # # insert in test problem not in test pinn.
+# # def test_train_cpu_sampling_few_vars():
+# #     poisson_problem = Poisson()
+# #     boundaries = ['gamma1', 'gamma2', 'gamma3']
+# #     n = 10
+# #     poisson_problem.discretise_domain(n, 'grid', locations=boundaries)
+# #     poisson_problem.discretise_domain(n, 'random', locations=['gamma4'], variables=['x'])
+# #     poisson_problem.discretise_domain(n, 'random', locations=['gamma4'], variables=['y'])
+# #     pinn = PINN(problem = poisson_problem, model=model, extra_features=None, loss=LpLoss())
+# #     trainer = Trainer(solver=pinn, kwargs={'max_epochs' : 5, 'accelerator':'cpu'})
+# #     trainer.train()
+
 
-# def test_train_gpu(): #TODO fix ASAP
+# def test_train_extra_feats_cpu():
 #     poisson_problem = Poisson()
 #     boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4']
 #     n = 10
 #     poisson_problem.discretise_domain(n, 'grid', locations=boundaries)
-#     poisson_problem.conditions.pop('data') # The input/output pts are allocated on cpu
-#     pinn = PINN(problem = poisson_problem, model=model, extra_features=None, loss=LpLoss())
-#     trainer = Trainer(solver=pinn, kwargs={'max_epochs' : 5, 'accelerator':'gpu'})
+#     pinn = PINN(problem=poisson_problem,
+#                 model=model_extra_feats,
+#                 extra_features=extra_feats)
+#     trainer = Trainer(solver=pinn, max_epochs=5, accelerator='cpu')
 #     trainer.train()
 
-# def test_train_2():
-#     boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4']
-#     n = 10
-#     expected_keys = [[], list(range(0, 50, 3))]
-#     param = [0, 3]
-#     for i, truth_key in zip(param, expected_keys):
-#         pinn = PINN(problem, model)
-#         pinn.discretise_domain(n, 'grid', locations=boundaries)
-#         pinn.discretise_domain(n, 'grid', locations=['D'])
-#         pinn.train(50, save_loss=i)
-#         assert list(pinn.history_loss.keys()) == truth_key
-
-
-# def test_train_extra_feats():
-#     pinn = PINN(problem, model_extra_feat, [myFeature()])
-#     boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4']
-#     n = 10
-#     pinn.discretise_domain(n, 'grid', locations=boundaries)
-#     pinn.discretise_domain(n, 'grid', locations=['D'])
-#     pinn.train(5)
 
+# # TODO, fix GitHub actions to run also on GPU
+# # def test_train_gpu():
+# #     poisson_problem = Poisson()
+# #     boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4']
+# #     n = 10
+# #     poisson_problem.discretise_domain(n, 'grid', locations=boundaries)
+# #     pinn = PINN(problem = poisson_problem, model=model, extra_features=None, loss=LpLoss())
+# #     trainer = Trainer(solver=pinn, kwargs={'max_epochs' : 5, 'accelerator':'gpu'})
+# #     trainer.train()
 
-# def test_train_2_extra_feats():
-#     boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4']
-#     n = 10
-#     expected_keys = [[], list(range(0, 50, 3))]
-#     param = [0, 3]
-#     for i, truth_key in zip(param, expected_keys):
-#         pinn = PINN(problem, model_extra_feat, [myFeature()])
-#         pinn.discretise_domain(n, 'grid', locations=boundaries)
-#         pinn.discretise_domain(n, 'grid', locations=['D'])
-#         pinn.train(50, save_loss=i)
-#         assert list(pinn.history_loss.keys()) == truth_key
-
+# # def test_train_gpu(): #TODO fix ASAP
+# #     poisson_problem = Poisson()
+# #     boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4']
+# #     n = 10
+# #     poisson_problem.discretise_domain(n, 'grid', locations=boundaries)
+# #     poisson_problem.conditions.pop('data') # The input/output pts are allocated on cpu
+# #     pinn = PINN(problem = poisson_problem, model=model, extra_features=None, loss=LpLoss())
+# #     trainer = Trainer(solver=pinn, kwargs={'max_epochs' : 5, 'accelerator':'gpu'})
+# #     trainer.train()
 
-# def test_train_with_optimizer_kwargs():
-#     boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4']
-#     n = 10
-#     expected_keys = [[], list(range(0, 50, 3))]
-#     param = [0, 3]
-#     for i, truth_key in zip(param, expected_keys):
-#         pinn = PINN(problem, model, optimizer_kwargs={'lr' : 0.3})
-#         pinn.discretise_domain(n, 'grid', locations=boundaries)
-#         pinn.discretise_domain(n, 'grid', locations=['D'])
-#         pinn.train(50, save_loss=i)
-#         assert list(pinn.history_loss.keys()) == truth_key
+# # def test_train_2():
+# #     boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4']
+# #     n = 10
+# #     expected_keys = [[], list(range(0, 50, 3))]
+# #     param = [0, 3]
+# #     for i, truth_key in zip(param, expected_keys):
+# #         pinn = PINN(problem, model)
+# #         pinn.discretise_domain(n, 'grid', locations=boundaries)
+# #         pinn.discretise_domain(n, 'grid', locations=['D'])
+# #         pinn.train(50, save_loss=i)
+# #         assert list(pinn.history_loss.keys()) == truth_key
 
 
-# def test_train_with_lr_scheduler():
-#     boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4']
-#     n = 10
-#     expected_keys = [[], list(range(0, 50, 3))]
-#     param = [0, 3]
-#     for i, truth_key in zip(param, expected_keys):
-#         pinn = PINN(
-#             problem,
-#             model,
-#             lr_scheduler_type=torch.optim.lr_scheduler.CyclicLR,
-#             lr_scheduler_kwargs={'base_lr' : 0.1, 'max_lr' : 0.3, 'cycle_momentum': False}
-#         )
-#         pinn.discretise_domain(n, 'grid', locations=boundaries)
-#         pinn.discretise_domain(n, 'grid', locations=['D'])
-#         pinn.train(50, save_loss=i)
-#         assert list(pinn.history_loss.keys()) == truth_key
-
-
-# # def test_train_batch():
-# #     pinn = PINN(problem, model, batch_size=6)
+# # def test_train_extra_feats():
+# #     pinn = PINN(problem, model_extra_feat, [myFeature()])
 # #     boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4']
 # #     n = 10
 # #     pinn.discretise_domain(n, 'grid', locations=boundaries)
@@ -416,34 +363,87 @@ def test_train_extra_feats_cpu():
 # #     pinn.train(5)
 
 
-# # def test_train_batch_2():
+# # def test_train_2_extra_feats():
+# #     boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4']
+# #     n = 10
+# #     expected_keys = [[], list(range(0, 50, 3))]
+# #     param = [0, 3]
+# #     for i, truth_key in zip(param, expected_keys):
+# #         pinn = PINN(problem, model_extra_feat, [myFeature()])
+# #         pinn.discretise_domain(n, 'grid', locations=boundaries)
+# #         pinn.discretise_domain(n, 'grid', locations=['D'])
+# #         pinn.train(50, save_loss=i)
+# #         assert list(pinn.history_loss.keys()) == truth_key
+
+
+# # def test_train_with_optimizer_kwargs():
 # #     boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4']
 # #     n = 10
 # #     expected_keys = [[], list(range(0, 50, 3))]
 # #     param = [0, 3]
 # #     for i, truth_key in zip(param, expected_keys):
-# #         pinn = PINN(problem, model, batch_size=6)
+# #         pinn = PINN(problem, model, optimizer_kwargs={'lr' : 0.3})
 # #         pinn.discretise_domain(n, 'grid', locations=boundaries)
 # #         pinn.discretise_domain(n, 'grid', locations=['D'])
 # #         pinn.train(50, save_loss=i)
 # #         assert list(pinn.history_loss.keys()) == truth_key
 
 
-# if torch.cuda.is_available():
+# # def test_train_with_lr_scheduler():
+# #     boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4']
+# #     n = 10
+# #     expected_keys = [[], list(range(0, 50, 3))]
+# #     param = [0, 3]
+# #     for i, truth_key in zip(param, expected_keys):
+# #         pinn = PINN(
+# #             problem,
+# #             model,
+# #             lr_scheduler_type=torch.optim.lr_scheduler.CyclicLR,
+# #             lr_scheduler_kwargs={'base_lr' : 0.1, 'max_lr' : 0.3, 'cycle_momentum': False}
+# #         )
+# #         pinn.discretise_domain(n, 'grid', locations=boundaries)
+# #         pinn.discretise_domain(n, 'grid', locations=['D'])
+# #         pinn.train(50, save_loss=i)
+# #         assert list(pinn.history_loss.keys()) == truth_key
 
-#     # def test_gpu_train():
-#     #     pinn = PINN(problem, model, batch_size=20, device='cuda')
-#     #     boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4']
-#     #     n = 100
-#     #     pinn.discretise_domain(n, 'grid', locations=boundaries)
-#     #     pinn.discretise_domain(n, 'grid', locations=['D'])
-#     #     pinn.train(5)
 
-#     def test_gpu_train_nobatch():
-#         pinn = PINN(problem, model, batch_size=None, device='cuda')
-#         boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4']
-#         n = 100
-#         pinn.discretise_domain(n, 'grid', locations=boundaries)
-#         pinn.discretise_domain(n, 'grid', locations=['D'])
-#         pinn.train(5)
+# # # def test_train_batch():
+# # #     pinn = PINN(problem, model, batch_size=6)
+# # #     boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4']
+# # #     n = 10
+# # #     pinn.discretise_domain(n, 'grid', locations=boundaries)
+# # #     pinn.discretise_domain(n, 'grid', locations=['D'])
+# # #     pinn.train(5)
+
+
+# # # def test_train_batch_2():
+# # #     boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4']
+# # #     n = 10
+# # #     expected_keys = [[], list(range(0, 50, 3))]
+# # #     param = [0, 3]
+# # #     for i, truth_key in zip(param, expected_keys):
+# # #         pinn = PINN(problem, model, batch_size=6)
+# # #         pinn.discretise_domain(n, 'grid', locations=boundaries)
+# # #         pinn.discretise_domain(n, 'grid', locations=['D'])
+# # #         pinn.train(50, save_loss=i)
+# # #         assert list(pinn.history_loss.keys()) == truth_key
+
+
+# # if torch.cuda.is_available():
+
+# #     # def test_gpu_train():
+# #     #     pinn = PINN(problem, model, batch_size=20, device='cuda')
+# #     #     boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4']
+# #     #     n = 100
+# #     #     pinn.discretise_domain(n, 'grid', locations=boundaries)
+# #     #     pinn.discretise_domain(n, 'grid', locations=['D'])
+# #     #     pinn.train(5)
+
+# #     def test_gpu_train_nobatch():
+# #         pinn = PINN(problem, model, batch_size=None, device='cuda')
+# #         boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4']
+# #         n = 100
+# #         pinn.discretise_domain(n, 'grid', locations=boundaries)
+# #         pinn.discretise_domain(n, 'grid', locations=['D'])
+# #         pinn.train(5)
 
diff --git a/tests/test_solvers/test_supervised_solver.py b/tests/test_solvers/test_supervised_solver.py
index ebe8179e6..b1c78d058 100644
--- a/tests/test_solvers/test_supervised_solver.py
+++ b/tests/test_solvers/test_supervised_solver.py
@@ -4,140 +4,140 @@
 from pina import Condition, LabelTensor
 from pina.solvers import SupervisedSolver
 from pina.model import FeedForward
-from pina.equation.equation import Equation
+from pina.equation import Equation
 from pina.equation.equation_factory import FixedValue
 from pina.operators import laplacian
 from pina.domain import CartesianDomain
 from pina.trainer import Trainer
 
-in_ = LabelTensor(torch.tensor([[0., 1.]]), ['u_0', 'u_1'])
-out_ = LabelTensor(torch.tensor([[0.]]), ['u'])
+# in_ = LabelTensor(torch.tensor([[0., 1.]]), ['u_0', 'u_1'])
+# out_ = LabelTensor(torch.tensor([[0.]]), ['u'])
 
 
-class NeuralOperatorProblem(AbstractProblem):
-    input_variables = ['u_0', 'u_1']
-    output_variables = ['u']
+# class NeuralOperatorProblem(AbstractProblem):
+#     input_variables = ['u_0', 'u_1']
+#     output_variables = ['u']
 
-    conditions = {
-        'data': Condition(input_points=in_, output_points=out_),
-    }
-
-
-class myFeature(torch.nn.Module):
-    """
-    Feature: sin(x)
-    """
+#     conditions = {
+#         'data': Condition(input_points=in_, output_points=out_),
+#     }
+
+
+# class myFeature(torch.nn.Module):
+#     """
+#     Feature: sin(x)
+#     """
 
-    def __init__(self):
-        super(myFeature, self).__init__()
+#     def __init__(self):
+#         super(myFeature, self).__init__()
 
-    def forward(self, x):
-        t = (torch.sin(x.extract(['u_0']) * torch.pi) *
-             torch.sin(x.extract(['u_1']) * torch.pi))
-        return LabelTensor(t, ['sin(x)sin(y)'])
+#     def forward(self, x):
+#         t = (torch.sin(x.extract(['u_0']) * torch.pi) *
+#              torch.sin(x.extract(['u_1']) * torch.pi))
+#         return LabelTensor(t, ['sin(x)sin(y)'])
 
-
-problem = NeuralOperatorProblem()
-extra_feats = [myFeature()]
-model = FeedForward(len(problem.input_variables), len(problem.output_variables))
-model_extra_feats = FeedForward(
-    len(problem.input_variables) + 1, len(problem.output_variables))
-
-
-def test_constructor():
-    SupervisedSolver(problem=problem, model=model)
-
-
-test_constructor()
-
-
-def laplace_equation(input_, output_):
-    force_term = (torch.sin(input_.extract(['x']) * torch.pi) *
-                  torch.sin(input_.extract(['y']) * torch.pi))
-    delta_u = laplacian(output_.extract(['u']), input_)
-    return delta_u - force_term
-
-
-my_laplace = Equation(laplace_equation)
-
-
-class Poisson(SpatialProblem):
-    output_variables = ['u']
-    spatial_domain = CartesianDomain({'x': [0, 1], 'y': [0, 1]})
-
-    conditions = {
-        'gamma1':
-            Condition(domain=CartesianDomain({
-                'x': [0, 1],
-                'y': 1
-            }),
-                equation=FixedValue(0.0)),
-        'gamma2':
-            Condition(domain=CartesianDomain({
-                'x': [0, 1],
-                'y': 0
-            }),
-                equation=FixedValue(0.0)),
-        'gamma3':
-            Condition(domain=CartesianDomain({
-                'x': 1,
-                'y': [0, 1]
-            }),
-                equation=FixedValue(0.0)),
-        'gamma4':
-            Condition(domain=CartesianDomain({
-                'x': 0,
-                'y': [0, 1]
-            }),
-                equation=FixedValue(0.0)),
-        'D':
-            Condition(domain=CartesianDomain({
-                'x': [0, 1],
-                'y': [0, 1]
-            }),
-                equation=my_laplace),
-        'data':
-            Condition(input_points=in_, output_points=out_)
-    }
-
-    def poisson_sol(self, pts):
-        return -(torch.sin(pts.extract(['x']) * torch.pi) *
-                 torch.sin(pts.extract(['y']) * torch.pi)) / (2 * torch.pi ** 2)
-
-    truth_solution = poisson_sol
-
-
-def test_wrong_constructor():
-    poisson_problem = Poisson()
-    with pytest.raises(ValueError):
-        SupervisedSolver(problem=poisson_problem, model=model)
-
-
-def test_train_cpu():
-    solver = SupervisedSolver(problem=problem, model=model)
-    trainer = Trainer(solver=solver,
-                      max_epochs=200,
-                      accelerator='gpu',
-                      batch_size=5,
-                      train_size=1,
-                      test_size=0.,
-                      val_size=0.)
-    trainer.train()
-test_train_cpu()
-
-
-def test_extra_features_constructor():
-    SupervisedSolver(problem=problem,
-                     model=model_extra_feats,
-                     extra_features=extra_feats)
-
-
-def test_extra_features_train_cpu():
-    solver = SupervisedSolver(problem=problem,
-                              model=model_extra_feats,
-                              extra_features=extra_feats)
-    trainer = Trainer(solver=solver,
-                      max_epochs=200,
-                      accelerator='gpu',
-                      batch_size=5)
-    trainer.train()
+
+# problem = NeuralOperatorProblem()
+# extra_feats = [myFeature()]
+# model = FeedForward(len(problem.input_variables), len(problem.output_variables))
+# model_extra_feats = FeedForward(
+#     len(problem.input_variables) + 1, len(problem.output_variables))
+
+
+# def test_constructor():
+#     SupervisedSolver(problem=problem, model=model)
+
+
+# test_constructor()
+
+
+# def laplace_equation(input_, output_):
+#     force_term = (torch.sin(input_.extract(['x']) * torch.pi) *
+#                   torch.sin(input_.extract(['y']) * torch.pi))
+#     delta_u = laplacian(output_.extract(['u']), input_)
+#     return delta_u - force_term
+
+
+# my_laplace = Equation(laplace_equation)
+
+
+# class Poisson(SpatialProblem):
+#     output_variables = ['u']
+#     spatial_domain = CartesianDomain({'x': [0, 1], 'y': [0, 1]})
+
+#     conditions = {
+#         'gamma1':
+#             Condition(domain=CartesianDomain({
+#                 'x': [0, 1],
+#                 'y': 1
+#             }),
+#                 equation=FixedValue(0.0)),
+#         'gamma2':
+#             Condition(domain=CartesianDomain({
+#                 'x': [0, 1],
+#                 'y': 0
+#             }),
+#                 equation=FixedValue(0.0)),
+#         'gamma3':
+#             Condition(domain=CartesianDomain({
+#                 'x': 1,
+#                 'y': [0, 1]
+#             }),
+#                 equation=FixedValue(0.0)),
+#         'gamma4':
+#             Condition(domain=CartesianDomain({
+#                 'x': 0,
+#                 'y': [0, 1]
+#             }),
+#                 equation=FixedValue(0.0)),
+#         'D':
+#             Condition(domain=CartesianDomain({
+#                 'x': [0, 1],
+#                 'y': [0, 1]
+#             }),
+#                 equation=my_laplace),
+#         'data':
+#             Condition(input_points=in_, output_points=out_)
+#     }
+
+#     def poisson_sol(self, pts):
+#         return -(torch.sin(pts.extract(['x']) * torch.pi) *
+#                  torch.sin(pts.extract(['y']) * torch.pi)) / (2 * torch.pi ** 2)
+
+#     truth_solution = poisson_sol
+
+
+# def test_wrong_constructor():
+#     poisson_problem = Poisson()
+#     with pytest.raises(ValueError):
+#         SupervisedSolver(problem=poisson_problem, model=model)
+
+
+# def test_train_cpu():
+#     solver = SupervisedSolver(problem=problem, model=model)
+#     trainer = Trainer(solver=solver,
+#                       max_epochs=200,
+#                       accelerator='gpu',
+#                       batch_size=5,
+#                       train_size=1,
+#                       test_size=0.,
+#                       val_size=0.)
+#     trainer.train()
+# test_train_cpu()
+
+
+# def test_extra_features_constructor():
+#     SupervisedSolver(problem=problem,
+#                      model=model_extra_feats,
+#                      extra_features=extra_feats)
+
+
+# def test_extra_features_train_cpu():
+#     solver = SupervisedSolver(problem=problem,
+#                               model=model_extra_feats,
+#                               extra_features=extra_feats)
+#     trainer = Trainer(solver=solver,
+#                       max_epochs=200,
+#                       accelerator='gpu',
+#                       batch_size=5)
+#     trainer.train()
diff --git a/tests/test_utils.py b/tests/test_utils.py
index 46a5083e6..911aa4be1 100644
--- a/tests/test_utils.py
+++ b/tests/test_utils.py
@@ -6,7 +6,7 @@
 from pina.domain import EllipsoidDomain, CartesianDomain
 from pina.utils import check_consistency
 import pytest
-from pina.domain import Location
+from pina.domain import DomainInterface
 
 
 def test_merge_tensors():
@@ -27,8 +27,8 @@ def test_check_consistency_correct():
     example_input_pts = LabelTensor(torch.tensor([[0, 0, 0]]), ['x', 'y', 'z'])
 
     check_consistency(example_input_pts, torch.Tensor)
-    check_consistency(CartesianDomain, Location, subclass=True)
-    check_consistency(ellipsoid1, Location)
+    check_consistency(CartesianDomain, DomainInterface, subclass=True)
+    check_consistency(ellipsoid1, DomainInterface)
 
 
 def test_check_consistency_incorrect():
@@ -36,8 +36,8 @@ def test_check_consistency_incorrect():
     example_input_pts = LabelTensor(torch.tensor([[0, 0, 0]]), ['x', 'y', 'z'])
 
     with pytest.raises(ValueError):
-        check_consistency(example_input_pts, Location)
+        check_consistency(example_input_pts, DomainInterface)
     with pytest.raises(ValueError):
-        check_consistency(torch.Tensor, Location, subclass=True)
+        check_consistency(torch.Tensor, DomainInterface, subclass=True)
     with pytest.raises(ValueError):
         check_consistency(ellipsoid1, torch.Tensor)

From 6a1860b19b8dd0b2de349e81eda5fbc49c333500 Mon Sep 17 00:00:00 2001
From: Nicola Demo <demo.nicola@gmail.com>
Date: Thu, 23 Jan 2025 15:44:53 +0100
Subject: [PATCH 50/55] temporary comment tests

---
 tests/test_callbacks/test_metric_tracker.py   | 28 +++++++++----------
 .../test_optimizer_callbacks.py               | 16 +++++------
 2 files changed, 22 insertions(+), 22 deletions(-)

diff --git a/tests/test_callbacks/test_metric_tracker.py b/tests/test_callbacks/test_metric_tracker.py
index ed637a887..586d06425 100644
--- a/tests/test_callbacks/test_metric_tracker.py
+++ b/tests/test_callbacks/test_metric_tracker.py
@@ -21,19 +21,19 @@
 def test_metric_tracker_constructor():
     MetricTracker()
 
-def test_metric_tracker_routine():
-    # make the trainer
-    trainer = Trainer(solver=solver,
-                      callbacks=[
-                          MetricTracker()
-                      ],
-                      accelerator='cpu',
-                      max_epochs=5)
-    trainer.train()
-    # get the tracked metrics
-    metrics = trainer.callbacks[0].metrics
-    # assert the logged metrics are correct
-    logged_metrics = sorted(list(metrics.keys()))
-    assert logged_metrics == ['train_loss_epoch', 'train_loss_step', 'val_loss']
+# def test_metric_tracker_routine(): #TODO revert
+#     # make the trainer
+#     trainer = Trainer(solver=solver,
+#                       callbacks=[
+#                           MetricTracker()
+#                       ],
+#                       accelerator='cpu',
+#                       max_epochs=5)
+#     trainer.train()
+#     # get the tracked metrics
+#     metrics = trainer.callbacks[0].metrics
+#     # assert the logged metrics are correct
+#     logged_metrics = sorted(list(metrics.keys()))
+#     assert logged_metrics == ['train_loss_epoch', 'train_loss_step', 'val_loss']
 
 
diff --git a/tests/test_callbacks/test_optimizer_callbacks.py b/tests/test_callbacks/test_optimizer_callbacks.py
index f62078a41..2f9acdbcd 100644
--- a/tests/test_callbacks/test_optimizer_callbacks.py
+++ b/tests/test_callbacks/test_optimizer_callbacks.py
@@ -27,11 +27,11 @@ def test_switch_optimizer_constructor():
     SwitchOptimizer(adam_optimizer, epoch_switch=10)
 
 
-def test_switch_optimizer_routine():
-    # make the trainer
-    switch_opt_callback = SwitchOptimizer(lbfgs_optimizer, epoch_switch=3)
-    trainer = Trainer(solver=solver,
-                      callbacks=[switch_opt_callback],
-                      accelerator='cpu',
-                      max_epochs=5)
-    trainer.train()
+# def test_switch_optimizer_routine(): #TODO revert
+#     # make the trainer
+#     switch_opt_callback = SwitchOptimizer(lbfgs_optimizer, epoch_switch=3)
+#     trainer = Trainer(solver=solver,
+#                       callbacks=[switch_opt_callback],
+#                       accelerator='cpu',
+#                       max_epochs=5)
+#     trainer.train()

From 12594ad9257878364af65b4c64d14a5231c6bd2c Mon Sep 17 00:00:00 2001
From: FilippoOlivo <filippo@filippoolivo.com>
Date: Fri, 24 Jan 2025 09:47:57 +0100
Subject: [PATCH 51/55] Additional improvement related to #395

---
 pina/label_tensor.py                         | 15 ++++++++-------
 tests/test_label_tensor/test_label_tensor.py | 14 ++++++++++++++
 2 files changed, 22 insertions(+), 7 deletions(-)

diff --git a/pina/label_tensor.py b/pina/label_tensor.py
index 5ffa61138..d8f66f952 100644
--- a/pina/label_tensor.py
+++ b/pina/label_tensor.py
@@ -380,20 +380,21 @@ def _update_single_label(self, old_labels, to_update_labels, index, dim,
         """
         old_dof = old_labels[to_update_dim]['dof']
         label_name = old_labels[dim]['name']
-
+        # Handle slicing
         if isinstance(index, slice):
-            # Handle slicing
             to_update_labels[dim] = {'dof': old_dof[index], 'name': label_name}
+        # Handle single integer index
         elif isinstance(index, int):
-            # Handle single integer index
             to_update_labels[dim] = {'dof': [old_dof[index]],
                                      'name': label_name}
+        # Handle lists or tensors
         elif isinstance(index, (list, torch.Tensor)):
-            # Handle lists or tensors
-            indices = [index] if isinstance(index, (int, str)) else index
+            # Handle list of bools
+            if isinstance(index, torch.Tensor) and index.dtype == torch.bool:
+                index = index.nonzero().squeeze()
             to_update_labels[dim] = {
-                'dof': [old_dof[i] for i in indices] if isinstance(old_dof,
-                                                                   list) else indices,
+                'dof': [old_dof[i] for i in index] if isinstance(old_dof,
+                                                                 list) else index,
                 'name': label_name
             }
         else:
diff --git a/tests/test_label_tensor/test_label_tensor.py b/tests/test_label_tensor/test_label_tensor.py
index 41288e693..2c5d15ea9 100644
--- a/tests/test_label_tensor/test_label_tensor.py
+++ b/tests/test_label_tensor/test_label_tensor.py
@@ -281,3 +281,17 @@ def test_sorting():
     assert torch.eq(lt_sorted.tensor[:, 1, :], torch.ones(20, 5) * 2).all()
     assert torch.eq(lt_sorted.tensor[:, 2, :], torch.ones(20, 5) * 3).all()
     assert torch.eq(lt_sorted.tensor[:, 3, :], torch.ones(20, 5) * 4).all()
+
+
+@pytest.mark.parametrize("labels",
+                         [[f's{i}' for i in range(10)],
+                          {0: {'dof': ['a', 'b', 'c']},
+                           1: {'dof': [f's{i}' for i in range(10)]}}])
+def test_cat_bool(labels):
+    out = torch.randn((3, 10))
+    out = LabelTensor(out, labels)
+    selected = out[torch.tensor([True, True, False])]
+    assert selected.shape == (2, 10)
+    assert selected.stored_labels[1]['dof'] == [f's{i}' for i in range(10)]
+    if isinstance(labels, dict):
+        assert selected.stored_labels[0]['dof'] == ['a', 'b']

From 272f8a3b5bdea5e139cb118a14b6f98ecc57ef3e Mon Sep 17 00:00:00 2001
From: Filippo Olivo <filippo@filippoolivo.com>
Date: Tue, 28 Jan 2025 13:51:57 +0100
Subject: [PATCH 52/55] Improvement in DDP and bug fix in DataModule (#432)

---
 pina/data/data_module.py | 94 ++++++++++++++++++++++------------------
 pina/data/dataset.py     | 60 ++++++++++++++++++-------
 pina/trainer.py          |  4 +-
 3 files changed, 98 insertions(+), 60 deletions(-)

diff --git a/pina/data/data_module.py b/pina/data/data_module.py
index 4b529fe2c..56473f843 100644
--- a/pina/data/data_module.py
+++ b/pina/data/data_module.py
@@ -8,6 +8,7 @@
 from torch.utils.data.distributed import DistributedSampler
 from .dataset import PinaDatasetFactory
 
+
 class DummyDataloader:
     def __init__(self, dataset, device):
         self.dataset = dataset.get_all_data()
@@ -21,16 +22,17 @@ def __len__(self):
     def __next__(self):
         return self.dataset
 
+
 class Collator:
-    def __init__(self, max_conditions_lengths, ):
+    def __init__(self, max_conditions_lengths, dataset=None):
         self.max_conditions_lengths = max_conditions_lengths
         self.callable_function = self._collate_custom_dataloader if \
             max_conditions_lengths is None else (
             self._collate_standard_dataloader)
+        self.dataset = dataset
 
-    @staticmethod
-    def _collate_custom_dataloader(batch):
-        return batch[0]
+    def _collate_custom_dataloader(self, batch):
+        return self.dataset.fetch_from_idx_list(batch)
 
     def _collate_standard_dataloader(self, batch):
         """
@@ -59,26 +61,24 @@ def _collate_standard_dataloader(self, batch):
             batch_dict[condition_name] = single_cond_dict
         return batch_dict
 
+
     def __call__(self, batch):
         return self.callable_function(batch)
 
 
-class PinaBatchSampler(BatchSampler):
-    def __init__(self, dataset, batch_size, shuffle, sampler=None):
-        if sampler is None:
-            if (torch.distributed.is_available() and
-                    torch.distributed.is_initialized()):
-                rank = torch.distributed.get_rank()
-                world_size = torch.distributed.get_world_size()
-                sampler = DistributedSampler(dataset, shuffle=shuffle,
-                                             rank=rank, num_replicas=world_size)
+class PinaSampler:
+    def __new__(self, dataset, batch_size, shuffle, automatic_batching):
+
+        if (torch.distributed.is_available() and
+                torch.distributed.is_initialized()):
+            sampler = DistributedSampler(dataset, shuffle=shuffle)
+        else:
+            if shuffle:
+                sampler = RandomSampler(dataset)
             else:
-                if shuffle:
-                    sampler = RandomSampler(dataset)
-                else:
-                    sampler = SequentialSampler(dataset)
-        super().__init__(sampler=sampler, batch_size=batch_size,
-                         drop_last=False)
+                sampler = SequentialSampler(dataset)
+        return sampler
+
 
 class PinaDataModule(LightningDataModule):
     """
@@ -136,6 +136,7 @@ def __init__(self,
         else:
             self.predict_dataloader = super().predict_dataloader
         self.collector_splits = self._create_splits(collector, splits_dict)
+        self.transfer_batch_to_device = self._transfer_batch_to_device
 
     def setup(self, stage=None):
         """
@@ -151,7 +152,7 @@ def setup(self, stage=None):
                 self.val_dataset = PinaDatasetFactory(
                     self.collector_splits['val'],
                     max_conditions_lengths=self.find_max_conditions_lengths(
-                        'val'),  automatic_batching=self.automatic_batching
+                        'val'), automatic_batching=self.automatic_batching
                 )
         elif stage == 'test':
             self.test_dataset = PinaDatasetFactory(
@@ -215,6 +216,7 @@ def _apply_shuffle(condition_dict, len_data):
                     condition_dict[k] = v[idx]
                 else:
                     raise ValueError(f"Data type {type(v)} not supported")
+
         # ----------- End auxiliary function ------------
 
         logging.debug('Dataset creation in PinaDataModule obj')
@@ -247,18 +249,21 @@ def val_dataloader(self):
         """
         # Use custom batching (good if batch size is large)
         if self.batch_size is not None:
-            # Use default batching in torch DataLoader (good is batch size is small)
+            sampler = PinaSampler(self.val_dataset, self.batch_size,
+                                  self.shuffle, self.automatic_batching)
             if self.automatic_batching:
                 collate = Collator(self.find_max_conditions_lengths('val'))
-                return DataLoader(self.val_dataset, self.batch_size,
-                          collate_fn=collate)
-            collate = Collator(None)
-            sampler = PinaBatchSampler(self.val_dataset, self.batch_size, shuffle=False)
-            return DataLoader(self.val_dataset, sampler=sampler,
-                          collate_fn=collate)
-        dataloader = DummyDataloader(self.train_dataset, self.trainer.strategy.root_device)
-        dataloader.dataset = self.transfer_batch_to_device(dataloader.dataset, self.trainer.strategy.root_device, 0)
-        self.transfer_batch_to_device = self.dummy_transfer_to_device
+            else:
+                collate = Collator(None, self.val_dataset)
+            return DataLoader(self.val_dataset, self.batch_size,
+                              collate_fn=collate, sampler=sampler)
+        dataloader = DummyDataloader(self.val_dataset,
+                                     self.trainer.strategy.root_device)
+        dataloader.dataset = self._transfer_batch_to_device(dataloader.dataset,
+                                                            self.trainer.strategy.root_device,
+                                                            0)
+        self.transfer_batch_to_device = self._transfer_batch_to_device_dummy
+        return dataloader
 
     def train_dataloader(self):
         """
@@ -266,19 +271,21 @@ def train_dataloader(self):
         """
         # Use custom batching (good if batch size is large)
         if self.batch_size is not None:
-            # Use default batching in torch DataLoader (good is batch size is small)
+            sampler = PinaSampler(self.train_dataset, self.batch_size,
+                                  self.shuffle, self.automatic_batching)
             if self.automatic_batching:
                 collate = Collator(self.find_max_conditions_lengths('train'))
-                return DataLoader(self.train_dataset, self.batch_size,
-                                  collate_fn=collate)
-            collate = Collator(None)
-            sampler = PinaBatchSampler(self.train_dataset, self.batch_size,
-                                           shuffle=False)
-            return DataLoader(self.train_dataset, sampler=sampler,
-                              collate_fn=collate)
-        dataloader = DummyDataloader(self.train_dataset, self.trainer.strategy.root_device)
-        dataloader.dataset = self.transfer_batch_to_device(dataloader.dataset, self.trainer.strategy.root_device, 0)
-        self.transfer_batch_to_device = self.dummy_transfer_to_device
+
+            else:
+                collate = Collator(None, self.train_dataset)
+            return DataLoader(self.train_dataset, self.batch_size,
+                            collate_fn=collate, sampler=sampler)
+        dataloader = DummyDataloader(self.train_dataset,
+                                     self.trainer.strategy.root_device)
+        dataloader.dataset = self._transfer_batch_to_device(dataloader.dataset,
+                                            self.trainer.strategy.root_device,
+                                            0)
+        self.transfer_batch_to_device = self._transfer_batch_to_device_dummy
         return dataloader
 
     def test_dataloader(self):
@@ -293,10 +300,10 @@ def predict_dataloader(self):
         """
         raise NotImplementedError("Predict dataloader not implemented")
 
-    def dummy_transfer_to_device(self, batch, device, dataloader_idx):
+    def _transfer_batch_to_device_dummy(self, batch, device, dataloader_idx):
         return batch
 
-    def transfer_batch_to_device(self, batch, device, dataloader_idx):
+    def _transfer_batch_to_device(self, batch, device, dataloader_idx):
         """
         Transfer the batch to the device. This method is called in the
         training loop and is used to transfer the batch to the device.
@@ -307,4 +314,5 @@ def transfer_batch_to_device(self, batch, device, dataloader_idx):
                                                                 dataloader_idx))
             for k, v in batch.items()
         ]
+
         return batch
diff --git a/pina/data/dataset.py b/pina/data/dataset.py
index 8f41c0bd7..8b5f998f9 100644
--- a/pina/data/dataset.py
+++ b/pina/data/dataset.py
@@ -6,6 +6,7 @@
 from abc import abstractmethod
 from torch_geometric.data import Batch
 
+
 class PinaDatasetFactory:
     """
     Factory class for the PINA dataset. Depending on the type inside the
@@ -13,6 +14,7 @@ class PinaDatasetFactory:
     - PinaTensorDataset for torch.Tensor
     - PinaGraphDataset for list of torch_geometric.data.Data objects
     """
+
     def __new__(cls, conditions_dict, **kwargs):
         if len(conditions_dict) == 0:
             raise ValueError('No conditions provided')
@@ -25,10 +27,12 @@ def __new__(cls, conditions_dict, **kwargs):
         raise ValueError('Conditions must be either torch.Tensor or list of Data '
                          'objects.')
 
+
 class PinaDataset(Dataset):
     """
     Abstract class for the PINA dataset
     """
+
     def __init__(self, conditions_dict, max_conditions_lengths):
         self.conditions_dict = conditions_dict
         self.max_conditions_lengths = max_conditions_lengths
@@ -49,6 +53,7 @@ def __len__(self):
     def __getitem__(self, item):
         pass
 
+
 class PinaTensorDataset(PinaDataset):
     def __init__(self, conditions_dict, max_conditions_lengths,
                  automatic_batching):
@@ -64,45 +69,68 @@ def _getitem_int(self, idx):
                 in v.keys()} for k, v in self.conditions_dict.items()
         }
 
-    def _getitem_list(self, idx):
+    def fetch_from_idx_list(self, idx):
         to_return_dict = {}
         for condition, data in self.conditions_dict.items():
             cond_idx = idx[:self.max_conditions_lengths[condition]]
             condition_len = self.conditions_length[condition]
             if self.length > condition_len:
-                cond_idx = [idx%condition_len for idx in cond_idx]
+                cond_idx = [idx % condition_len for idx in cond_idx]
             to_return_dict[condition] = {k: v[cond_idx]
                                          for k, v in data.items()}
         return to_return_dict
 
+    @staticmethod
+    def _getitem_list(idx):
+        return idx
+
     def get_all_data(self):
         index = [i for i in range(len(self))]
-        return self._getitem_list(index)
+        return self.fetch_from_idx_list(index)
 
     def __getitem__(self, idx):
         return self._getitem_func(idx)
 
+
 class PinaGraphDataset(PinaDataset):
     pass
-    """
-    def __init__(self, conditions_dict, max_conditions_lengths):
+'''
+    def __init__(self, conditions_dict, max_conditions_lengths,
+                 automatic_batching):
         super().__init__(conditions_dict, max_conditions_lengths)
+        if automatic_batching:
+            self._getitem_func = self._getitem_int
+        else:
+            self._getitem_func = self._getitem_list
 
-    def __getitem__(self, idx):
-        
-        Getitem method for large batch size
-        
+    def fetch_from_idx_list(self, idx):
         to_return_dict = {}
         for condition, data in self.conditions_dict.items():
             cond_idx = idx[:self.max_conditions_lengths[condition]]
             condition_len = self.conditions_length[condition]
             if self.length > condition_len:
-                cond_idx = [idx%condition_len for idx in cond_idx]
+                cond_idx = [idx % condition_len for idx in cond_idx]
             to_return_dict[condition] = {k: Batch.from_data_list([v[i]
-                                            for i in cond_idx])
-                            if isinstance(v, list)
-                            else v[cond_idx].tensor.reshape(-1, v.size(-1))
-                        for k, v in data.items()
-                    }
+                                                                  for i in cond_idx])
+            if isinstance(v, list)
+            else v[cond_idx]
+                                         for k, v in data.items()
+                                         }
         return to_return_dict
-    """
+
+    def _getitem_list(self, idx):
+        return idx
+
+    def _getitem_int(self, idx):
+        return {
+            k: {k_data: v[k_data][idx % len(v['input_points'])] for k_data
+                in v.keys()} for k, v in self.conditions_dict.items()
+        }
+
+    def get_all_data(self):
+        index = [i for i in range(len(self))]
+        return self.fetch_from_idx_list(index)
+
+    def __getitem__(self, idx):
+        return self._getitem_func(idx)
+'''
\ No newline at end of file
diff --git a/pina/trainer.py b/pina/trainer.py
index f8bccd8c5..6a162482b 100644
--- a/pina/trainer.py
+++ b/pina/trainer.py
@@ -72,12 +72,14 @@ def _create_loader(self):
             raise RuntimeError('Cannot create Trainer if not all conditions '
                                'are sampled. The Trainer got the following:\n'
                                f'{error_message}')
+        automatic_batching = False
         self.data_module = PinaDataModule(collector=self.solver.problem.collector,
                                           train_size=self.train_size,
                                           test_size=self.test_size,
                                           val_size=self.val_size,
                                           predict_size=self.predict_size,
-                                          batch_size=self.batch_size,)
+                                          batch_size=self.batch_size,
+                                          automatic_batching=automatic_batching)
 
     def train(self, **kwargs):
         """

From a862dc8b10e635e53656ee86ea138bec897d7fbc Mon Sep 17 00:00:00 2001
From: Dario Coscia <dariocos99@gmail.com>
Date: Tue, 28 Jan 2025 00:55:40 +0100
Subject: [PATCH 53/55] start refactoring solvers

---
 pina/collector.py                            |   6 +-
 pina/data/data_module.py                     |   5 +-
 pina/model/network.py                        | 111 ---------
 pina/solvers/solver.py                       | 181 ++++++++++----
 pina/solvers/supervised.py                   | 103 +++-----
 pina/utils.py                                |  85 +------
 tests/test_solvers/test_supervised_solver.py | 248 +++++++++----------
 7 files changed, 285 insertions(+), 454 deletions(-)
 delete mode 100644 pina/model/network.py

diff --git a/pina/collector.py b/pina/collector.py
index 1f0fb41d3..750d401a6 100644
--- a/pina/collector.py
+++ b/pina/collector.py
@@ -1,5 +1,6 @@
 from . import LabelTensor
 from .utils import check_consistency, merge_tensors
+from copy import deepcopy
 
 
 class Collector:
@@ -8,11 +9,6 @@ def __init__(self, problem):
         # creating a hook between collector and problem
         self.problem = problem
 
-        # this variable is used to store the data in the form:
-        # {'[condition_name]' :
-        #           {'input_points' : Tensor,
-        #            '[equation/output_points/conditional_variables]': Tensor}
-        # }
         # those variables are used for the dataloading
         self._data_collections = {name: {} for name in self.problem.conditions}
         self.conditions_name = {
diff --git a/pina/data/data_module.py b/pina/data/data_module.py
index 56473f843..238448e48 100644
--- a/pina/data/data_module.py
+++ b/pina/data/data_module.py
@@ -243,10 +243,7 @@ def find_max_conditions_lengths(self, split):
                                                 self.batch_size)
         return max_conditions_lengths
 
-    def val_dataloader(self):
-        """
-        Create the validation dataloader
-        """
+    def _create_dataloader(self, dataset):
         # Use custom batching (good if batch size is large)
         if self.batch_size is not None:
             sampler = PinaSampler(self.val_dataset, self.batch_size,
diff --git a/pina/model/network.py b/pina/model/network.py
deleted file mode 100644
index aed3dff3d..000000000
--- a/pina/model/network.py
+++ /dev/null
@@ -1,111 +0,0 @@
-import torch
-import torch.nn as nn
-from ..utils import check_consistency
-from ..label_tensor import LabelTensor
-
-
-class Network(torch.nn.Module):
-
-    def __init__(
-        self, model, input_variables, output_variables, extra_features=None
-    ):
-        """
-        Network class with standard forward method
-        and possibility to pass extra features. This
-        class is used internally in PINA to convert
-        any :class:`torch.nn.Module` s in a PINA module.
-
-        :param model: The torch model to convert in a PINA model.
-        :type model: torch.nn.Module
-        :param list(str) input_variables: The input variables of the :class:`AbstractProblem`, whose type depends on the
-            type of domain (spatial, temporal, and parameter).
-        :param list(str) output_variables: The output variables of the :class:`AbstractProblem`, whose type depends on the
-            problem setting.
-        :param extra_features: List of torch models to augment the input, defaults to None.
-        :type extra_features: list(torch.nn.Module)
-        """
-        super().__init__()
-
-        # check model consistency
-        check_consistency(model, nn.Module)
-        check_consistency(input_variables, str)
-        if output_variables is not None:
-            check_consistency(output_variables, str)
-
-        self._model = model
-        self._input_variables = input_variables
-        self._output_variables = output_variables
-
-        # check consistency and assign extra fatures
-        if extra_features is None:
-            self._extra_features = []
-        else:
-            for feat in extra_features:
-                check_consistency(feat, nn.Module)
-            self._extra_features = nn.Sequential(*extra_features)
-
-        # check model works with inputs
-        # TODO
-
-    def forward(self, x):
-        """
-        Forward method for Network class. This class
-        implements the standard forward method, and
-        it adds the possibility to pass extra features.
-        All the PINA models ``forward`` s are overriden
-        by this class, to enable :class:`pina.label_tensor.LabelTensor` labels
-        extraction.
-
-        :param torch.Tensor x: Input of the network.
-        :return torch.Tensor: Output of the network.
-        """
-        # only labeltensors as input
-        assert isinstance(
-            x, LabelTensor
-        ), "Expected LabelTensor as input to the model."
-
-        # extract torch.Tensor from corresponding label
-        # in case `input_variables = []` all points are used
-        if self._input_variables:
-            x = x.extract(self._input_variables)
-        # extract features and append
-        for feature in self._extra_features:
-            x = x.append(feature(x))
-
-        # perform forward pass + converting to LabelTensor
-        x = x.as_subclass(torch.Tensor)
-        output = self._model(x)
-        if self._output_variables is not None:
-            output = LabelTensor(output, self._output_variables)
-
-        return output
-
-    # TODO to remove in next releases (only used in GAROM solver)
-    def forward_map(self, x):
-        """
-        Forward method for Network class when the input is
-        a tuple. This class is simply a forward with the input casted as a
-        tuple or list :class`torch.Tensor`.
-        All the PINA models ``forward`` s are overriden
-        by this class, to enable :class:`pina.label_tensor.LabelTensor` labels
-        extraction.
-
-        :param list (torch.Tensor) | tuple(torch.Tensor) x: Input of the network.
-        :return torch.Tensor: Output of the network.
-
-        .. note::
-            This function does not extract the input variables, all the variables
-            are used for both tensors. Output variables are correctly applied.
-        """
-
-        # perform forward pass (using torch.Tensor) + converting to LabelTensor
-        output = LabelTensor(self._model(x.tensor), self._output_variables)
-        return output
-
-    @property
-    def torchmodel(self):
-        return self._model
-
-    @property
-    def extra_features(self):
-        return self._extra_features
diff --git a/pina/solvers/solver.py b/pina/solvers/solver.py
index 76c084eab..30ace0a2e 100644
--- a/pina/solvers/solver.py
+++ b/pina/solvers/solver.py
@@ -1,16 +1,15 @@
 """ Solver module. """
 
 from abc import ABCMeta, abstractmethod
-from ..model.network import Network
 import lightning
-from ..utils import check_consistency
+from ..utils import check_consistency, labelize_forward
 from ..problem import AbstractProblem
-from ..optim import Optimizer, Scheduler
+from ..optim import Optimizer, Scheduler, TorchOptimizer, TorchScheduler
 import torch
 import sys
 
 
-class SolverInterface(lightning.pytorch.LightningModule, metaclass=ABCMeta):
+class MultipleSolversInterface(lightning.pytorch.LightningModule, metaclass=ABCMeta):
     """
     Solver base class. This class inherits is a wrapper of
     LightningModule class, inheriting all the
@@ -22,14 +21,18 @@ def __init__(self,
                  problem,
                  optimizers,
                  schedulers,
+                 extra_features=None,
                  use_lt=True):
         """
-        :param model: A torch neural network model instance.
-        :type model: torch.nn.Module
+        :param models: Multiple torch nn.Module instances.
+        :type model: list[torch.nn.Module] | tuple[torch.nn.Module]
         :param problem: A problem definition instance.
         :type problem: AbstractProblem
-        :param list(torch.optim.Optimizer) optimizer: A list of neural network
+        :param list(Optimizer) optimizers: A list of neural network
            optimizers to use.
+        :param list(Scheduler) optimizers: A list of neural network
+           schedulers to use.
+        :param bool use_lt: Using LabelTensors as input during training.
         """
         super().__init__()
 
@@ -38,40 +41,31 @@ def __init__(self,
         self._check_solver_consistency(problem)
 
         # Check consistency of models argument and encapsulate in list
-        if not isinstance(models, list):
-            check_consistency(models, torch.nn.Module)
-            # put everything in a list if only one input
-            models = [models]
-        else:
-            for idx in range(len(models)):
-                # Check consistency
-                check_consistency(models[idx], torch.nn.Module)
+        check_consistency(models, torch.nn.Module)
         len_model = len(models)
 
+        # Check consistency extra_features
+        if extra_features is None:
+            extra_features = []
+        else:
+            check_consistency(extra_features, torch.nn.Module)
+
         # If use_lt is true add extract operation in input
+        check_consistency(use_lt, bool)
+        self._use_lt = use_lt
         if use_lt is True:
-            for idx, model in enumerate(models):
-                models[idx] = Network(
-                    model=model,
-                    input_variables=problem.input_variables,
-                    output_variables=problem.output_variables,
+            self.forward = labelize_forward(
+                forward=self.forward,
+                input_variables=problem.input_variables,
+                output_variables=problem.output_variables,
+                extra_features=extra_features
                 )
 
         # Check scheduler consistency + encapsulation
-        if not isinstance(schedulers, list):
-            check_consistency(schedulers, Scheduler)
-            schedulers = [schedulers]
-        else:
-            for scheduler in schedulers:
-                check_consistency(scheduler, Scheduler)
+        check_consistency(schedulers, Scheduler)
 
         # Check optimizer consistency + encapsulation
-        if not isinstance(optimizers, list):
-            check_consistency(optimizers, Optimizer)
-            optimizers = [optimizers]
-        else:
-            for optimizer in optimizers:
-                check_consistency(optimizer, Optimizer)
+        check_consistency(optimizers, Optimizer)
         len_optimizer = len(optimizers)
 
         # check length consistency optimizers
@@ -95,26 +89,17 @@ def training_step(self, batch):
         pass
 
     @abstractmethod
-    def configure_optimizers(self):
-        raise NotImplementedError
+    def validation_step(self, batch):
+        pass
 
-    @property
-    def models(self):
-        """
-        The torch model."""
-        return self._pina_models
+    @abstractmethod
+    def test_step(self, batch):
+        pass
 
-    @property
-    def optimizers(self):
-        """
-        The torch model."""
-        return self._pina_optimizers
+    @abstractmethod
+    def configure_optimizers(self):
+        raise NotImplementedError
 
-    @property
-    def problem(self):
-        """
-        The problem formulation."""
-        return self._pina_problem
 
     def on_train_start(self):
         """
@@ -141,3 +126,101 @@ def get_batch_size(batch):
     def _check_solver_consistency(self, problem):
         for condition in problem.conditions.values():
             check_consistency(condition, self.accepted_conditions_types)
+
+    @property
+    def models(self):
+        """
+        The torch model."""
+        return self._pina_models
+
+    @property
+    def optimizers(self):
+        """
+        The torch model."""
+        return self._pina_optimizers
+
+    @property
+    def schedulers(self):
+        """
+        The torch model."""
+        return self._pina_schedulers
+
+    @property
+    def problem(self):
+        """
+        The problem formulation."""
+        return self._pina_problem
+    
+    @property
+    def use_lt(self):
+        """
+        Using LabelTensor in training."""
+        return self._use_lt
+
+
+class SolverInterface(MultipleSolversInterface):
+    def __init__(self, model, problem, optimizer, scheduler, extra_features=None, use_lt=True):
+        """
+        :param model: A torch nn.Module instances.
+        :type model: torch.nn.Module
+        :param problem: A problem definition instance.
+        :type problem: AbstractProblem
+        :param Optimizer optimizers: A neural network optimizers to use.
+        :param Scheduler optimizers: A neural network scheduler to use.
+        :param extra_features: The additional input features to use as
+            augmented input.
+        :type extra_features: list[torch.nn.Module] | tuple[torch.nn.Module]
+        :param bool use_lt: Using LabelTensors as input during training.
+        """
+        if optimizer is None:
+            optimizer = SolverInterface.default_torch_optimizer()
+
+        if scheduler is None:
+            scheduler = SolverInterface.default_torch_scheduler()
+
+        super().__init__(models = [model],
+                         problem = problem,
+                         optimizers = [optimizer],
+                         schedulers = [scheduler],
+                         extra_features = extra_features,
+                         use_lt = use_lt)
+        # initialize model (needed for Lightining to go to different devices)
+        self._pina_model = self.models[0]
+    
+    def forward(self, x):
+        """Forward pass implementation for the solver.
+
+        :param torch.Tensor x: Input tensor.
+        :return: Solver solution.
+        :rtype: torch.Tensor
+        """
+        return self.model(x)
+
+    @staticmethod
+    def default_torch_optimizer():
+        return TorchOptimizer(torch.optim.Adam, lr=0.001)
+    
+    @staticmethod
+    def default_torch_scheduler():
+        return TorchScheduler(torch.optim.lr_scheduler.ConstantLR)
+    
+    @property
+    def model(self):
+        """
+        Model for training.
+        """
+        return self._pina_model
+    
+    @property
+    def scheduler(self):
+        """
+        Scheduler for training.
+        """
+        return self.schedulers[0]
+
+    @property
+    def optimizer(self):
+        """
+        Optimizer for training.
+        """
+        return self.optimizers[0]
\ No newline at end of file
diff --git a/pina/solvers/supervised.py b/pina/solvers/supervised.py
index 99fec09a3..9c54832af 100644
--- a/pina/solvers/supervised.py
+++ b/pina/solvers/supervised.py
@@ -3,7 +3,6 @@
 from torch.nn.modules.loss import _Loss
 from ..optim import TorchOptimizer, TorchScheduler
 from .solver import SolverInterface
-from ..label_tensor import LabelTensor
 from ..utils import check_consistency
 from ..loss.loss_interface import LossInterface
 from ..condition import InputOutputPointsCondition
@@ -63,16 +62,10 @@ def __init__(self,
         if loss is None:
             loss = torch.nn.MSELoss()
 
-        if optimizer is None:
-            optimizer = TorchOptimizer(torch.optim.Adam, lr=0.001)
-
-        if scheduler is None:
-            scheduler = TorchScheduler(torch.optim.lr_scheduler.ConstantLR)
-
-        super().__init__(models=model,
+        super().__init__(model=model,
                          problem=problem,
-                         optimizers=optimizer,
-                         schedulers=scheduler,
+                         optimizer=optimizer,
+                         scheduler=scheduler,
                          extra_features=extra_features,
                          use_lt=use_lt)
 
@@ -80,25 +73,6 @@ def __init__(self,
         check_consistency(loss, (LossInterface, _Loss, torch.nn.Module),
                           subclass=False)
         self._loss = loss
-        self._model = self._pina_models[0]
-        self._optimizer = self._pina_optimizers[0]
-        self._scheduler = self._pina_schedulers[0]
-        self.validation_condition_losses = {
-            k: {'loss': [],
-                'count': []} for k in self.problem.conditions.keys()}
-
-    def forward(self, x):
-        """Forward pass implementation for the solver.
-
-        :param torch.Tensor x: Input tensor.
-        :return: Solver solution.
-        :rtype: torch.Tensor
-        """
-
-        output = self._model(x)
-
-        output.labels = self.problem.output_variables
-        return output
 
     def configure_optimizers(self):
         """Optimizer configuration for the solver.
@@ -106,11 +80,20 @@ def configure_optimizers(self):
         :return: The optimizers and the schedulers
         :rtype: tuple(list, list)
         """
-        self._optimizer.hook(self._model.parameters())
-        self._scheduler.hook(self._optimizer)
-        return ([self._optimizer.optimizer_instance],
-                [self._scheduler.scheduler_instance])
+        self.optimizer.hook(self.model.parameters())
+        self.scheduler.hook(self.optimizer)
+        return ([self.optimizer.optimizer_instance],
+                [self.scheduler.scheduler_instance])
 
+    def _optimization_cycle(self, batch):
+        condition_loss = []
+        for _, points in batch:
+            input_pts, output_pts = points['input_points'], points['output_points']
+            loss = self.loss_data(input_pts=input_pts, output_pts=output_pts)
+            condition_loss.append(loss.as_subclass(torch.Tensor))
+        loss = sum(condition_loss)
+        return loss
+    
     def training_step(self, batch):
         """Solver training step.
 
@@ -121,12 +104,7 @@ def training_step(self, batch):
         :return: The sum of the loss functions.
         :rtype: LabelTensor
         """
-        condition_loss = []
-        for condition_name, points in batch:
-            input_pts, output_pts = points['input_points'], points['output_points']
-            loss_ = self.loss_data(input_pts=input_pts, output_pts=output_pts)
-            condition_loss.append(loss_.as_subclass(torch.Tensor))
-        loss = sum(condition_loss)
+        loss = self._optimization_cycle(batch=batch)
         self.log('train_loss', loss, on_step=True, on_epoch=True, prog_bar=True, logger=True,
                  batch_size=self.get_batch_size(batch), sync_dist=True)
         return loss
@@ -135,21 +113,17 @@ def validation_step(self, batch):
         """
         Solver validation step.
         """
-        condition_loss = []
-        for condition_name, points in batch:
-            input_pts, output_pts = points['input_points'], points['output_points']
-            loss_ = self.loss_data(input_pts=input_pts, output_pts=output_pts)
-            condition_loss.append(loss_.as_subclass(torch.Tensor))
-        loss = sum(condition_loss)
+        loss = self._optimization_cycle(batch=batch)
         self.log('val_loss', loss, prog_bar=True, logger=True,
                  batch_size=self.get_batch_size(batch), sync_dist=True)
-
-    def test_step(self, batch, batch_idx):
+        
+    def test_step(self, batch):
         """
-        Solver test step.
+        Solver validation step.
         """
-
-        raise NotImplementedError("Test step not implemented yet.")
+        loss = self._optimization_cycle(batch=batch)
+        self.log('test_loss', loss, prog_bar=True, logger=True,
+                 batch_size=self.get_batch_size(batch), sync_dist=True)
 
     def loss_data(self, input_pts, output_pts):
         """
@@ -157,35 +131,16 @@ def loss_data(self, input_pts, output_pts):
         the network output against the true solution. This function
         should not be override if not intentionally.
 
-        :param LabelTensor input_pts: The input to the neural networks.
-        :param LabelTensor output_pts: The true solution to compare the
+        :param input_pts: The input to the neural networks.
+        :type input_pts: LabelTensor | torch.Tensor
+        :param output_pts: The true solution to compare the
             network solution.
-        :return: The residual loss averaged on the input coordinates
+        :type output_pts: LabelTensor | torch.Tensor
+        :return: The residual loss.
         :rtype: torch.Tensor
         """
         return self._loss(self.forward(input_pts), output_pts)
 
-    @property
-    def scheduler(self):
-        """
-        Scheduler for training.
-        """
-        return self._scheduler
-
-    @property
-    def optimizer(self):
-        """
-        Optimizer for training.
-        """
-        return self._optimizer
-
-    @property
-    def model(self):
-        """
-        Neural network for training.
-        """
-        return self._model
-
     @property
     def loss(self):
         """
diff --git a/pina/utils.py b/pina/utils.py
index 84d3e7419..48c6ccd95 100644
--- a/pina/utils.py
+++ b/pina/utils.py
@@ -39,35 +39,15 @@ def check_consistency(object, object_instance, subclass=False):
         except AssertionError:
             raise ValueError(f"{type(obj).__name__} must be {object_instance}.")
 
-
-def number_parameters(model,
-                      aggregate=True,
-                      only_trainable=True):  # TODO: check
-    """
-    Return the number of parameters of a given `model`.
-
-    :param torch.nn.Module model: the torch module to inspect.
-    :param bool aggregate: if True the return values is an integer corresponding
-        to the total amount of parameters of whole model. If False, it returns a
-        dictionary whose keys are the names of layers and the values the
-        corresponding number of parameters. Default is True.
-    :param bool trainable: if True, only trainable parameters are count,
-        otherwise no. Default is True.
-    :return: the number of parameters of the model
-    :rtype: dict or int
-    """
-    tmp = {}
-    for name, parameter in model.named_parameters():
-        if only_trainable and not parameter.requires_grad:
-            continue
-
-        tmp[name] = parameter.numel()
-
-    if aggregate:
-        tmp = sum(tmp.values())
-
-    return tmp
-
+def labelize_forward(forward, input_variables, output_variables, extra_features=[]):
+    def wrapper(x):
+        x = x.extract(input_variables)
+        for feature in extra_features:
+            x = x.append(feature(x))
+        output = forward(x.tensor)
+        output = LabelTensor(output, output_variables)
+        return output
+    return wrapper
 
 def merge_tensors(tensors):  # name to be changed
     if tensors:
@@ -144,53 +124,6 @@ def chebyshev_roots(n):
     return nodes
 
 
-# class PinaDataset():
-
-#     def __init__(self, pinn) -> None:
-#         self.pinn = pinn
-
-#     @property
-#     def dataloader(self):
-#         return self._create_dataloader()
-
-#     @property
-#     def dataset(self):
-#         return [self.SampleDataset(key, val)
-#                 for key, val in self.input_pts.items()]
-
-#     def _create_dataloader(self):
-#         """Private method for creating dataloader
-
-#         :return: dataloader
-#         :rtype: torch.utils.data.DataLoader
-#         """
-#         if self.pinn.batch_size is None:
-#             return {key: [{key: val}] for key, val in self.pinn.input_pts.items()}
-
-#         def custom_collate(batch):
-#             # extracting pts labels
-#             _, pts = list(batch[0].items())[0]
-#             labels = pts.labels
-#             # calling default torch collate
-#             collate_res = default_collate(batch)
-#             # save collate result in dict
-#             res = {}
-#             for key, val in collate_res.items():
-#                 val.labels = labels
-#                 res[key] = val
-#             __init__(self, location, tensor):
-#             self._tensor = tensor
-#             self._location = location
-#             self._len = len(tensor)
-
-#         def __getitem__(self, index):
-#             tensor = self._tensor.select(0, index)
-#             return {self._location: tensor}
-
-#         def __len__(self):
-#             return self._len
-
-
 class LabelTensorDataLoader(DataLoader):
 
     def collate_fn(self, data):
diff --git a/tests/test_solvers/test_supervised_solver.py b/tests/test_solvers/test_supervised_solver.py
index b1c78d058..b8c0eaa7e 100644
--- a/tests/test_solvers/test_supervised_solver.py
+++ b/tests/test_solvers/test_supervised_solver.py
@@ -1,143 +1,121 @@
 import torch
 import pytest
-from pina.problem import AbstractProblem, SpatialProblem
+from pina.problem import AbstractProblem
+from pina.problem.zoo import Poisson2DSquareProblem
 from pina import Condition, LabelTensor
 from pina.solvers import SupervisedSolver
 from pina.model import FeedForward
-from pina.equation import Equation
-from pina.equation.equation_factory import FixedValue
-from pina.operators import laplacian
-from pina.domain import CartesianDomain
 from pina.trainer import Trainer
 
-# in_ = LabelTensor(torch.tensor([[0., 1.]]), ['u_0', 'u_1'])
-# out_ = LabelTensor(torch.tensor([[0.]]), ['u'])
 
-
-# class NeuralOperatorProblem(AbstractProblem):
-#     input_variables = ['u_0', 'u_1']
-#     output_variables = ['u']
-
-#     conditions = {
-#         'data': Condition(input_points=in_, output_points=out_),
-#     }
-
-
-# class myFeature(torch.nn.Module):
-#     """
-#     Feature: sin(x)
-#     """
-
-#     def __init__(self):
-#         super(myFeature, self).__init__()
-
-#     def forward(self, x):
-#         t = (torch.sin(x.extract(['u_0']) * torch.pi) *
-#              torch.sin(x.extract(['u_1']) * torch.pi))
-#         return LabelTensor(t, ['sin(x)sin(y)'])
-
-
-# problem = NeuralOperatorProblem()
-# extra_feats = [myFeature()]
-# model = FeedForward(len(problem.input_variables), len(problem.output_variables))
-# model_extra_feats = FeedForward(
-#     len(problem.input_variables) + 1, len(problem.output_variables))
-
-
-# def test_constructor():
-#     SupervisedSolver(problem=problem, model=model)
-
-
-# test_constructor()
-
-
-# def laplace_equation(input_, output_):
-#     force_term = (torch.sin(input_.extract(['x']) * torch.pi) *
-#                   torch.sin(input_.extract(['y']) * torch.pi))
-#     delta_u = laplacian(output_.extract(['u']), input_)
-#     return delta_u - force_term
-
-
-# my_laplace = Equation(laplace_equation)
-
-
-# class Poisson(SpatialProblem):
-#     output_variables = ['u']
-#     spatial_domain = CartesianDomain({'x': [0, 1], 'y': [0, 1]})
-
-#     conditions = {
-#         'gamma1':
-#             Condition(domain=CartesianDomain({
-#                 'x': [0, 1],
-#                 'y': 1
-#             }),
-#                 equation=FixedValue(0.0)),
-#         'gamma2':
-#             Condition(domain=CartesianDomain({
-#                 'x': [0, 1],
-#                 'y': 0
-#             }),
-#                 equation=FixedValue(0.0)),
-#         'gamma3':
-#             Condition(domain=CartesianDomain({
-#                 'x': 1,
-#                 'y': [0, 1]
-#             }),
-#                 equation=FixedValue(0.0)),
-#         'gamma4':
-#             Condition(domain=CartesianDomain({
-#                 'x': 0,
-#                 'y': [0, 1]
-#             }),
-#                 equation=FixedValue(0.0)),
-#         'D':
-#             Condition(domain=CartesianDomain({
-#                 'x': [0, 1],
-#                 'y': [0, 1]
-#             }),
-#                 equation=my_laplace),
-#         'data':
-#             Condition(input_points=in_, output_points=out_)
-#     }
-
-#     def poisson_sol(self, pts):
-#         return -(torch.sin(pts.extract(['x']) * torch.pi) *
-#                  torch.sin(pts.extract(['y']) * torch.pi)) / (2 * torch.pi ** 2)
-
-#     truth_solution = poisson_sol
-
-
-# def test_wrong_constructor():
-#     poisson_problem = Poisson()
-#     with pytest.raises(ValueError):
-#         SupervisedSolver(problem=poisson_problem, model=model)
-
-
-# def test_train_cpu():
-#     solver = SupervisedSolver(problem=problem, model=model)
-#     trainer = Trainer(solver=solver,
-#                       max_epochs=200,
-#                       accelerator='gpu',
-#                       batch_size=5,
-#                       train_size=1,
-#                       test_size=0.,
-#                       val_size=0.)
-#     trainer.train()
-# test_train_cpu()
-
-
-# def test_extra_features_constructor():
-#     SupervisedSolver(problem=problem,
-#                      model=model_extra_feats,
-#                      extra_features=extra_feats)
-
-
-# def test_extra_features_train_cpu():
-#     solver = SupervisedSolver(problem=problem,
-#                               model=model_extra_feats,
-#                               extra_features=extra_feats)
-#     trainer = Trainer(solver=solver,
-#                       max_epochs=200,
-#                       accelerator='gpu',
-#                       batch_size=5)
-#     trainer.train()
+class FooProblem(AbstractProblem):
+    input_variables = ['u_0', 'u_1']
+    output_variables = ['u']
+    conditions = {
+        'data': Condition(
+            input_points=LabelTensor(torch.tensor([[0., 1.]]), ['u_0', 'u_1']),
+            output_points=LabelTensor(torch.tensor([[0.]]), ['u'])),
+    }
+
+
+class myFeature(torch.nn.Module):
+    """
+    Feature: sin(x)
+    """
+    def __init__(self):
+        super(myFeature, self).__init__()
+
+    def forward(self, x):
+        t = (torch.sin(x.extract(['u_0']) * torch.pi) *
+             torch.sin(x.extract(['u_1']) * torch.pi))
+        return LabelTensor(t, ['u_0u_1'])
+
+
+problem = FooProblem()
+extra_feats = [myFeature()]
+model = FeedForward(len(problem.input_variables), len(problem.output_variables))
+model_extra_feats = FeedForward(
+    len(problem.input_variables) + 1, len(problem.output_variables))
+
+
+
+def test_constructor():
+    SupervisedSolver(problem=problem, model=model)
+
+def test_wrong_constructor():
+    with pytest.raises(ValueError):
+        SupervisedSolver(problem=Poisson2DSquareProblem(), model=model)
+
+def test_train_batch_size_full():
+    solver = SupervisedSolver(problem=problem, model=model)
+    trainer = Trainer(solver=solver,
+                      max_epochs=2,
+                      accelerator='cpu',
+                      batch_size=None,
+                      train_size=1.,
+                      test_size=0.,
+                      val_size=0.)
+    trainer.train() 
+    solver = SupervisedSolver(problem=problem, model=model)
+    trainer = Trainer(solver=solver,
+                      max_epochs=2,
+                      accelerator='cpu',
+                      batch_size=None,
+                      train_size=1.,
+                      test_size=0.,
+                      val_size=0.) 
+    trainer.train()  
+ 
+def test_train_and_val_cpu():
+
+    solver = SupervisedSolver(problem=problem, model=model)
+    trainer = Trainer(solver=solver,
+                      max_epochs=2,
+                      accelerator='cpu',
+                      batch_size=5,
+                      train_size=0.9,
+                      test_size=0.1,
+                      val_size=0.)
+    trainer.train()
+
+def test_train_and_val_gpu():
+    solver = SupervisedSolver(problem=problem, model=model)
+    trainer = Trainer(solver=solver,
+                      max_epochs=2,
+                      accelerator='mps',
+                      batch_size=5,
+                      train_size=1,
+                      test_size=0.,
+                      val_size=0.)
+    trainer.train()
+
+def test_extra_features_constructor():
+    SupervisedSolver(problem=problem,
+                     model=model_extra_feats,
+                     extra_features=extra_feats)
+
+def test_extra_features_train_and_val_cpu():
+    solver = SupervisedSolver(problem=problem,
+                              model=model_extra_feats,
+                              extra_features=extra_feats)
+    trainer = Trainer(solver=solver,
+                      max_epochs=2,
+                      accelerator='cpu',
+                      batch_size=5,
+                      train_size=0.9,
+                      test_size=0.1,
+                      )
+    trainer.train()
+
+def test_extra_features_train_and_val_gpu():
+    solver = SupervisedSolver(problem=problem,
+                              model=model_extra_feats,
+                              extra_features=extra_feats)
+    trainer = Trainer(solver=solver,
+                      max_epochs=2,
+                      accelerator='mps',
+                      batch_size=5,
+                      train_size=0.9,
+                      test_size=0.1,
+                      )
+    trainer.train()

From 88532f6b116f6fd4e14275626e1c2230198f4005 Mon Sep 17 00:00:00 2001
From: Dario Coscia <dariocos99@gmail.com>
Date: Tue, 28 Jan 2025 13:52:37 +0100
Subject: [PATCH 54/55] solvers update

---
 pina/data/data_module.py                     |  63 +++---
 pina/solvers/rom.py                          |  39 ++--
 pina/solvers/solver.py                       |   1 -
 pina/solvers/supervised.py                   |   4 +-
 pina/trainer.py                              |  18 +-
 tests/test_solvers/test_rom_solver.py        | 216 ++++++++++++-------
 tests/test_solvers/test_supervised_solver.py | 112 +++++++---
 7 files changed, 264 insertions(+), 189 deletions(-)

diff --git a/pina/data/data_module.py b/pina/data/data_module.py
index 238448e48..c58d8239f 100644
--- a/pina/data/data_module.py
+++ b/pina/data/data_module.py
@@ -1,6 +1,5 @@
 import logging
 from lightning.pytorch import LightningDataModule
-import math
 import torch
 from ..label_tensor import LabelTensor
 from torch.utils.data import DataLoader, BatchSampler, SequentialSampler, \
@@ -176,23 +175,28 @@ def _split_condition(condition_dict, splits_dict):
         len_condition = len(condition_dict['input_points'])
 
         lengths = [
-            int(math.floor(len_condition * length)) for length in
+            int(len_condition * length) for length in
             splits_dict.values()
         ]
 
         remainder = len_condition - sum(lengths)
         for i in range(remainder):
             lengths[i % len(lengths)] += 1
-        splits_dict = {k: v for k, v in zip(splits_dict.keys(), lengths)
+
+        splits_dict = {k: max(1,v) for k, v in zip(splits_dict.keys(), lengths)
                        }
         to_return_dict = {}
         offset = 0
+
         for stage, stage_len in splits_dict.items():
             to_return_dict[stage] = {k: v[offset:offset + stage_len]
                                      for k, v in condition_dict.items() if
                                      k != 'equation'
                                      # Equations are NEVER dataloaded
                                      }
+            if offset + stage_len > len_condition:
+                offset = len_condition - 1
+                continue
             offset += stage_len
         return to_return_dict
 
@@ -243,53 +247,40 @@ def find_max_conditions_lengths(self, split):
                                                 self.batch_size)
         return max_conditions_lengths
 
-    def _create_dataloader(self, dataset):
+    def _create_dataloader(self, split, dataset):
         # Use custom batching (good if batch size is large)
         if self.batch_size is not None:
-            sampler = PinaSampler(self.val_dataset, self.batch_size,
-                                  self.shuffle, self.automatic_batching)
+            # Use default batching in torch DataLoader (good if batch size is small)
             if self.automatic_batching:
-                collate = Collator(self.find_max_conditions_lengths('val'))
-            else:
-                collate = Collator(None, self.val_dataset)
-            return DataLoader(self.val_dataset, self.batch_size,
-                              collate_fn=collate, sampler=sampler)
-        dataloader = DummyDataloader(self.val_dataset,
-                                     self.trainer.strategy.root_device)
-        dataloader.dataset = self._transfer_batch_to_device(dataloader.dataset,
-                                                            self.trainer.strategy.root_device,
-                                                            0)
-        self.transfer_batch_to_device = self._transfer_batch_to_device_dummy
+                collate = Collator(self.find_max_conditions_lengths(split))
+                return DataLoader(dataset, self.batch_size,
+                          collate_fn=collate)
+            collate = Collator(None)
+            sampler = PinaBatchSampler(dataset, self.batch_size, shuffle=False)
+            return DataLoader(dataset, sampler=sampler,
+                          collate_fn=collate)
+        dataloader = DummyDataloader(dataset, self.trainer.strategy.root_device)
+        dataloader.dataset = self.transfer_batch_to_device(dataloader.dataset, self.trainer.strategy.root_device, 0)
+        self.transfer_batch_to_device = self.dummy_transfer_to_device
         return dataloader
+    
+    def val_dataloader(self):
+        """
+        Create the validation dataloader
+        """
+        return self._create_dataloader('val', self.val_dataset)
 
     def train_dataloader(self):
         """
         Create the training dataloader
         """
-        # Use custom batching (good if batch size is large)
-        if self.batch_size is not None:
-            sampler = PinaSampler(self.train_dataset, self.batch_size,
-                                  self.shuffle, self.automatic_batching)
-            if self.automatic_batching:
-                collate = Collator(self.find_max_conditions_lengths('train'))
-
-            else:
-                collate = Collator(None, self.train_dataset)
-            return DataLoader(self.train_dataset, self.batch_size,
-                            collate_fn=collate, sampler=sampler)
-        dataloader = DummyDataloader(self.train_dataset,
-                                     self.trainer.strategy.root_device)
-        dataloader.dataset = self._transfer_batch_to_device(dataloader.dataset,
-                                            self.trainer.strategy.root_device,
-                                            0)
-        self.transfer_batch_to_device = self._transfer_batch_to_device_dummy
-        return dataloader
+        return self._create_dataloader('train', self.train_dataset)
 
     def test_dataloader(self):
         """
         Create the testing dataloader
         """
-        raise NotImplementedError("Test dataloader not implemented")
+        raise self._create_dataloader('test', self.test_dataset)
 
     def predict_dataloader(self):
         """
diff --git a/pina/solvers/rom.py b/pina/solvers/rom.py
index ee4bcff43..41971e870 100644
--- a/pina/solvers/rom.py
+++ b/pina/solvers/rom.py
@@ -88,11 +88,10 @@ def __init__(
         problem,
         reduction_network,
         interpolation_network,
-        loss=torch.nn.MSELoss(),
-        optimizer=torch.optim.Adam,
-        optimizer_kwargs={"lr": 0.001},
-        scheduler=torch.optim.lr_scheduler.ConstantLR,
-        scheduler_kwargs={"factor": 1, "total_iters": 0},
+        loss=None,
+        optimizer=None,
+        scheduler=None,
+        use_lt=True,
     ):
         """
         :param AbstractProblem problem: The formualation of the problem.
@@ -109,11 +108,9 @@ def __init__(
             features to use as augmented input.
         :param torch.optim.Optimizer optimizer: The neural network optimizer to
             use; default is :class:`torch.optim.Adam`.
-        :param dict optimizer_kwargs: Optimizer constructor keyword args.
-        :param float lr: The learning rate; default is 0.001.
         :param torch.optim.LRScheduler scheduler: Learning
             rate scheduler.
-        :param dict scheduler_kwargs: LR scheduler constructor keyword args.
+        :param bool use_lt: Using LabelTensors as input during training.
         """
         model = torch.nn.ModuleDict(
             {
@@ -127,19 +124,19 @@ def __init__(
             problem=problem,
             loss=loss,
             optimizer=optimizer,
-            optimizer_kwargs=optimizer_kwargs,
             scheduler=scheduler,
-            scheduler_kwargs=scheduler_kwargs,
+            extra_features=None,
+            use_lt=use_lt
         )
 
         # assert reduction object contains encode/ decode
-        if not hasattr(self.neural_net["reduction_network"], "encode"):
+        if not hasattr(self.model["reduction_network"], "encode"):
             raise SyntaxError(
                 "reduction_network must have encode method. "
                 "The encode method should return a lower "
                 "dimensional representation of the input."
             )
-        if not hasattr(self.neural_net["reduction_network"], "decode"):
+        if not hasattr(self.model["reduction_network"], "decode"):
             raise SyntaxError(
                 "reduction_network must have decode method. "
                 "The decode method should return a high "
@@ -157,8 +154,8 @@ def forward(self, x):
         :return: Solver solution.
         :rtype: torch.Tensor
         """
-        reduction_network = self.neural_net["reduction_network"]
-        interpolation_network = self.neural_net["interpolation_network"]
+        reduction_network = self.model["reduction_network"]
+        interpolation_network = self.model["interpolation_network"]
         return reduction_network.decode(interpolation_network(x))
 
     def loss_data(self, input_pts, output_pts):
@@ -175,8 +172,8 @@ def loss_data(self, input_pts, output_pts):
         :rtype: torch.Tensor
         """
         # extract networks
-        reduction_network = self.neural_net["reduction_network"]
-        interpolation_network = self.neural_net["interpolation_network"]
+        reduction_network = self.model["reduction_network"]
+        interpolation_network = self.model["interpolation_network"]
         # encoded representations loss
         encode_repr_inter_net = interpolation_network(input_pts)
         encode_repr_reduction_network = reduction_network.encode(output_pts)
@@ -188,12 +185,4 @@ def loss_data(self, input_pts, output_pts):
             reduction_network.decode(encode_repr_reduction_network), output_pts
         )
 
-        return loss_encode + loss_reconstruction
-
-    @property
-    def neural_net(self):
-        """
-        Neural network for training. It returns a :obj:`~torch.nn.ModuleDict`
-        containing the ``reduction_network`` and ``interpolation_network``.
-        """
-        return self._neural_net.torchmodel
+        return loss_encode + loss_reconstruction
\ No newline at end of file
diff --git a/pina/solvers/solver.py b/pina/solvers/solver.py
index 30ace0a2e..def838d07 100644
--- a/pina/solvers/solver.py
+++ b/pina/solvers/solver.py
@@ -100,7 +100,6 @@ def test_step(self, batch):
     def configure_optimizers(self):
         raise NotImplementedError
 
-
     def on_train_start(self):
         """
         On training epoch start this function is call to do global checks for
diff --git a/pina/solvers/supervised.py b/pina/solvers/supervised.py
index 9c54832af..f54f40e5d 100644
--- a/pina/solvers/supervised.py
+++ b/pina/solvers/supervised.py
@@ -58,6 +58,8 @@ def __init__(self,
             use; default is :class:`torch.optim.Adam`.
         :param torch.optim.LRScheduler scheduler: Learning
             rate scheduler.
+        :type extra_features: list[torch.nn.Module] | tuple[torch.nn.Module]
+        :param bool use_lt: Using LabelTensors as input during training.
         """
         if loss is None:
             loss = torch.nn.MSELoss()
@@ -87,7 +89,7 @@ def configure_optimizers(self):
 
     def _optimization_cycle(self, batch):
         condition_loss = []
-        for _, points in batch:
+        for condition_name, points in batch:
             input_pts, output_pts = points['input_points'], points['output_points']
             loss = self.loss_data(input_pts=input_pts, output_pts=output_pts)
             condition_loss.append(loss.as_subclass(torch.Tensor))
diff --git a/pina/trainer.py b/pina/trainer.py
index 6a162482b..4f13ace99 100644
--- a/pina/trainer.py
+++ b/pina/trainer.py
@@ -14,14 +14,16 @@ def __init__(self,
                  train_size=.7,
                  test_size=.2,
                  val_size=.1,
-                 predict_size=.0,
+                 predict_size=0.,
                  **kwargs):
         """
         PINA Trainer class for costumizing every aspect of training via flags.
 
-        :param solver: A pina:class:`SolverInterface` solver for the differential problem.
+        :param solver: A pina:class:`SolverInterface` solver for the
+            differential problem.
         :type solver: SolverInterface
-        :param batch_size: How many samples per batch to load. If ``batch_size=None`` all
+        :param batch_size: How many samples per batch to load.
+            If ``batch_size=None`` all
             samples are loaded and data are not batched, defaults to None.
         :type batch_size: int | None
 
@@ -30,7 +32,6 @@ def __init__(self,
             and can be choosen from the `pytorch-lightning
             Trainer API <https://lightning.ai/docs/pytorch/stable/common/trainer.html#trainer-class-api>`_
         """
-
         super().__init__(**kwargs)
 
         # check inheritance consistency for solver and batch size
@@ -49,7 +50,6 @@ def __init__(self,
 
     def _move_to_device(self):
         device = self._accelerator_connector._parallel_devices[0]
-
         # move parameters to device
         pb = self.solver.problem
         if hasattr(pb, "unknown_parameters"):
@@ -85,17 +85,13 @@ def train(self, **kwargs):
         """
         Train the solver method.
         """
-        return super().fit(self.solver,
-                               datamodule=self.data_module,
-                               **kwargs)
+        return super().fit(self.solver, datamodule=self.data_module, **kwargs)
 
     def test(self, **kwargs):
         """
         Test the solver method.
         """
-        return super().test(self.solver,
-                            datamodule=self.data_module,
-                            **kwargs)
+        return super().test(self.solver, datamodule=self.data_module, **kwargs)
 
     @property
     def solver(self):
diff --git a/tests/test_solvers/test_rom_solver.py b/tests/test_solvers/test_rom_solver.py
index 00dc69ff6..44d2ebd9a 100644
--- a/tests/test_solvers/test_rom_solver.py
+++ b/tests/test_solvers/test_rom_solver.py
@@ -7,99 +7,153 @@
 from pina.trainer import Trainer
 from pina.model import FeedForward
 from pina.loss import LpLoss
+from pina.problem.zoo import Poisson2DSquareProblem
 
 
-# class NeuralOperatorProblem(AbstractProblem):
-#     input_variables = ['u_0', 'u_1']
-#     output_variables = [f'u_{i}' for i in range(100)]
-#     conditions = {'data' : Condition(input_points=
-#                                      LabelTensor(torch.rand(10, 2),
-#                                                  input_variables), 
-#                                      output_points=
-#                                      LabelTensor(torch.rand(10, 100),
-#                                                  output_variables))}
+class FooProblem(AbstractProblem):
+    input_variables = ['u_0', 'u_1']
+    output_variables = ['u']
+    conditions = {
+        'data': Condition(
+            input_points=LabelTensor(torch.tensor([[0., 1.]]), ['u_0', 'u_1']),
+            output_points=LabelTensor(torch.tensor([[0.]]), ['u'])),
+    }
 
+class myFeature(torch.nn.Module):
+    def __init__(self):
+        super(myFeature, self).__init__()
+    def forward(self, x):
+        t = (torch.sin(x.extract(['u_0']) * torch.pi) *
+             torch.sin(x.extract(['u_1']) * torch.pi))
+        return LabelTensor(t, ['u_0u_1'])
 
-# # make the problem + extra feats
-# class AE(torch.nn.Module):
-#     def __init__(self, input_dimensions, rank):
-#         super().__init__()
-#         self.encode = FeedForward(input_dimensions, rank, layers=[input_dimensions//4])
-#         self.decode = FeedForward(rank, input_dimensions, layers=[input_dimensions//4])
-# class AE_missing_encode(torch.nn.Module):
-#     def __init__(self, input_dimensions, rank):
-#         super().__init__()
-#         self.encode = FeedForward(input_dimensions, rank, layers=[input_dimensions//4])
-# class AE_missing_decode(torch.nn.Module):
-#     def __init__(self, input_dimensions, rank):
-#         super().__init__()
-#         self.decode = FeedForward(rank, input_dimensions, layers=[input_dimensions//4])
 
-# rank = 10
-# problem = NeuralOperatorProblem()
-# interpolation_net = FeedForward(len(problem.input_variables),
-#                                 rank)
-# reduction_net = AE(len(problem.output_variables), rank)
+# make the problem + extra feats
+class AE(torch.nn.Module):
+    def __init__(self, input_dimensions, rank):
+        super().__init__()
+        self.encode = FeedForward(input_dimensions, rank, layers=[input_dimensions//4])
+        self.decode = FeedForward(rank, input_dimensions, layers=[input_dimensions//4])
+class AE_missing_encode(torch.nn.Module):
+    def __init__(self, input_dimensions, rank):
+        super().__init__()
+        self.encode = FeedForward(input_dimensions, rank, layers=[input_dimensions//4])
+class AE_missing_decode(torch.nn.Module):
+    def __init__(self, input_dimensions, rank):
+        super().__init__()
+        self.decode = FeedForward(rank, input_dimensions, layers=[input_dimensions//4])
 
-# def test_constructor():
-#     ReducedOrderModelSolver(problem=problem,reduction_network=reduction_net,
-#     interpolation_network=interpolation_net)
-#     with pytest.raises(SyntaxError):
-#         ReducedOrderModelSolver(problem=problem,
-#                      reduction_network=AE_missing_encode(
-#                          len(problem.output_variables), rank),
-#                      interpolation_network=interpolation_net)
-#         ReducedOrderModelSolver(problem=problem,
-#                      reduction_network=AE_missing_decode(
-#                          len(problem.output_variables), rank),
-#                      interpolation_network=interpolation_net)
+rank = 10
+problem = FooProblem()
+interpolation_net = FeedForward(len(problem.input_variables),
+                                rank)
+reduction_net = AE(len(problem.output_variables), rank)
 
 
-# def test_train_cpu():
-#     solver = ReducedOrderModelSolver(problem = problem,reduction_network=reduction_net,
-#     interpolation_network=interpolation_net, loss=LpLoss())
-#     trainer = Trainer(solver=solver, max_epochs=3, accelerator='cpu', batch_size=20)
-#     trainer.train()
+def test_constructor():
+    ReducedOrderModelSolver(problem=problem,reduction_network=reduction_net,
+    interpolation_network=interpolation_net)
+
+def test_wrong_constructor():
+    with pytest.raises(SyntaxError):
+        ReducedOrderModelSolver(problem=problem,
+                     reduction_network=AE_missing_encode(
+                         len(problem.output_variables), rank),
+                     interpolation_network=interpolation_net)
+        ReducedOrderModelSolver(problem=problem,
+                     reduction_network=AE_missing_decode(
+                         len(problem.output_variables), rank),
+                     interpolation_network=interpolation_net)
+    with pytest.raises(ValueError):
+        ReducedOrderModelSolver(problem=Poisson2DSquareProblem(),
+                                reduction_network=reduction_net,
+                                interpolation_network=interpolation_net)
 
+def test_train_batch_size_full():
+    solver = ReducedOrderModelSolver(problem=problem,
+                                     reduction_network=reduction_net,
+                                     interpolation_network=interpolation_net)
+    trainer = Trainer(solver=solver,
+                      max_epochs=2,
+                      accelerator='cpu',
+                      batch_size=None,
+                      train_size=1.,
+                      test_size=0.,
+                      val_size=0.)
+    trainer.train()  
+ 
+def test_train_and_val_cpu():
+    solver = ReducedOrderModelSolver(problem=problem,
+                                     reduction_network=reduction_net,
+                                     interpolation_network=interpolation_net)
+    trainer = Trainer(solver=solver,
+                      max_epochs=2,
+                      accelerator='cpu',
+                      batch_size=5,
+                      train_size=0.9,
+                      test_size=0.1,
+                      val_size=0.)
+    trainer.train()
 
-# def test_train_restore():
-#     tmpdir = "tests/tmp_restore"
+# def test_train_and_val_gpu():
 #     solver = ReducedOrderModelSolver(problem=problem,
-#                reduction_network=reduction_net,
-#                interpolation_network=interpolation_net,
-#                loss=LpLoss())
+#                                      reduction_network=reduction_net,
+#                                      interpolation_network=interpolation_net)
 #     trainer = Trainer(solver=solver,
-#                       max_epochs=5,
-#                       accelerator='cpu',
-#                       default_root_dir=tmpdir)
+#                       max_epochs=2,
+#                       accelerator='gpu',
+#                       batch_size=5,
+#                       train_size=0.9,
+#                       test_size=0.1,
+#                       val_size=0.)
 #     trainer.train()
-#     ntrainer = Trainer(solver=solver, max_epochs=15, accelerator='cpu')
-#     t = ntrainer.train(
-#         ckpt_path=f'{tmpdir}/lightning_logs/version_0/checkpoints/epoch=4-step=5.ckpt')
-#     import shutil
-#     shutil.rmtree(tmpdir)
 
+def test_train_restore():
+    tmpdir = "tests/test_solvers/tmp/tmp_restore"
+    solver = ReducedOrderModelSolver(problem=problem,
+                                     reduction_network=reduction_net,
+                                     interpolation_network=interpolation_net)
+    trainer = Trainer(solver=solver,
+                      max_epochs=5,
+                      accelerator='cpu',
+                      batch_size=None,
+                      train_size=0.9,
+                      test_size=0.1,
+                      val_size=0.,
+                      default_root_dir=tmpdir)
+    trainer.train()
+    ntrainer = Trainer(solver=solver,
+                       max_epochs=5,
+                       accelerator='cpu',)
+    ntrainer.train(
+        ckpt_path=f'{tmpdir}/lightning_logs/version_0/checkpoints/epoch=4-step=5.ckpt')
+    import shutil
+    shutil.rmtree('tests/test_solvers/tmp')
 
-# def test_train_load():
-#     tmpdir = "tests/tmp_load"
-#     solver = ReducedOrderModelSolver(problem=problem,
-#                reduction_network=reduction_net,
-#                interpolation_network=interpolation_net,
-#                loss=LpLoss())
-#     trainer = Trainer(solver=solver,
-#                       max_epochs=15,
-#                       accelerator='cpu',
-#                       default_root_dir=tmpdir)
-#     trainer.train()
-#     new_solver = ReducedOrderModelSolver.load_from_checkpoint(
-#         f'{tmpdir}/lightning_logs/version_0/checkpoints/epoch=14-step=15.ckpt',
-#         problem = problem,reduction_network=reduction_net,
-#         interpolation_network=interpolation_net)
-#     test_pts = LabelTensor(torch.rand(20, 2), problem.input_variables)
-#     assert new_solver.forward(test_pts).shape == (20, 100)
-#     assert new_solver.forward(test_pts).shape == solver.forward(test_pts).shape
-#     torch.testing.assert_close(
-#         new_solver.forward(test_pts),
-#         solver.forward(test_pts))
-#     import shutil
-#     shutil.rmtree(tmpdir)
\ No newline at end of file
+
+def test_train_load():
+    tmpdir = "tests/test_solvers/tmp/tmp_load"
+    solver = ReducedOrderModelSolver(problem=problem,
+               reduction_network=reduction_net,
+               interpolation_network=interpolation_net)
+    trainer = Trainer(solver=solver,
+                      max_epochs=5,
+                      accelerator='cpu',
+                      batch_size=None,
+                    #   train_size=0.9,
+                    #   test_size=0.1,
+                    #   val_size=0.,
+                      default_root_dir=tmpdir)
+    trainer.train()
+    new_solver = ReducedOrderModelSolver.load_from_checkpoint(
+        f'{tmpdir}/lightning_logs/version_0/checkpoints/epoch=4-step=5.ckpt',
+        problem = problem,reduction_network=reduction_net,
+        interpolation_network=interpolation_net)
+    test_pts = LabelTensor(torch.rand(20, 2), problem.input_variables)
+    assert new_solver.forward(test_pts).shape == (20, 1)
+    assert new_solver.forward(test_pts).shape == solver.forward(test_pts).shape
+    torch.testing.assert_close(
+        new_solver.forward(test_pts),
+        solver.forward(test_pts))
+    import shutil
+    shutil.rmtree('tests/test_solvers/tmp')
\ No newline at end of file
diff --git a/tests/test_solvers/test_supervised_solver.py b/tests/test_solvers/test_supervised_solver.py
index b8c0eaa7e..4b8d202a3 100644
--- a/tests/test_solvers/test_supervised_solver.py
+++ b/tests/test_solvers/test_supervised_solver.py
@@ -13,18 +13,14 @@ class FooProblem(AbstractProblem):
     output_variables = ['u']
     conditions = {
         'data': Condition(
-            input_points=LabelTensor(torch.tensor([[0., 1.]]), ['u_0', 'u_1']),
-            output_points=LabelTensor(torch.tensor([[0.]]), ['u'])),
+            input_points=LabelTensor(torch.randn(34, 2), ['u_0', 'u_1']),
+            output_points=LabelTensor(torch.randn(34, 1), ['u'])),
     }
 
 
 class myFeature(torch.nn.Module):
-    """
-    Feature: sin(x)
-    """
     def __init__(self):
         super(myFeature, self).__init__()
-
     def forward(self, x):
         t = (torch.sin(x.extract(['u_0']) * torch.pi) *
              torch.sin(x.extract(['u_1']) * torch.pi))
@@ -56,15 +52,6 @@ def test_train_batch_size_full():
                       test_size=0.,
                       val_size=0.)
     trainer.train() 
-    solver = SupervisedSolver(problem=problem, model=model)
-    trainer = Trainer(solver=solver,
-                      max_epochs=2,
-                      accelerator='cpu',
-                      batch_size=None,
-                      train_size=1.,
-                      test_size=0.,
-                      val_size=0.) 
-    trainer.train()  
  
 def test_train_and_val_cpu():
 
@@ -74,20 +61,20 @@ def test_train_and_val_cpu():
                       accelerator='cpu',
                       batch_size=5,
                       train_size=0.9,
-                      test_size=0.1,
-                      val_size=0.)
+                      test_size=0.0,
+                      val_size=0.1,)
     trainer.train()
 
-def test_train_and_val_gpu():
-    solver = SupervisedSolver(problem=problem, model=model)
-    trainer = Trainer(solver=solver,
-                      max_epochs=2,
-                      accelerator='mps',
-                      batch_size=5,
-                      train_size=1,
-                      test_size=0.,
-                      val_size=0.)
-    trainer.train()
+# def test_train_and_val_gpu():
+#     solver = SupervisedSolver(problem=problem, model=model)
+#     trainer = Trainer(solver=solver,
+#                       max_epochs=2,
+#                       accelerator='gpu',
+#                       batch_size=5,
+#                       train_size=1,
+#                       test_size=0.,
+#                       val_size=0.)
+#     trainer.train()
 
 def test_extra_features_constructor():
     SupervisedSolver(problem=problem,
@@ -107,15 +94,72 @@ def test_extra_features_train_and_val_cpu():
                       )
     trainer.train()
 
-def test_extra_features_train_and_val_gpu():
+# def test_extra_features_train_and_val_gpu():
+#     solver = SupervisedSolver(problem=problem,
+#                               model=model_extra_feats,
+#                               extra_features=extra_feats)
+#     trainer = Trainer(solver=solver,
+#                       max_epochs=2,
+#                       accelerator='gpu',
+#                       batch_size=5,
+#                       train_size=0.9,
+#                       test_size=0.1,
+#                       )
+#     trainer.train()
+
+def test_train_restore():
+    tmpdir = "tests/test_solvers/tmp/tmp_restore"
     solver = SupervisedSolver(problem=problem,
-                              model=model_extra_feats,
-                              extra_features=extra_feats)
+                              model=model,
+                              extra_features=None)
     trainer = Trainer(solver=solver,
-                      max_epochs=2,
-                      accelerator='mps',
-                      batch_size=5,
+                      max_epochs=5,
+                      accelerator='cpu',
+                      batch_size=None,
                       train_size=0.9,
                       test_size=0.1,
-                      )
+                      val_size=0.,
+                      default_root_dir=tmpdir)
+    trainer.train()
+    ntrainer = Trainer(solver=solver,
+                       max_epochs=5,
+                       accelerator='cpu',)
+    ntrainer.train(
+        ckpt_path=f'{tmpdir}/lightning_logs/version_0/checkpoints/epoch=4-step=5.ckpt')
+    import shutil
+    shutil.rmtree('tests/test_solvers/tmp')
+
+
+def test_train_load():
+    tmpdir = "tests/test_solvers/tmp/tmp_load"
+    solver = SupervisedSolver(problem=problem,
+                              model=model,
+                              extra_features=None)
+    # trainer = Trainer(solver=solver,
+    #                   max_epochs=5,
+    #                   accelerator='cpu',
+    #                   batch_size=None,
+    #                   train_size=0.9,
+    #                   test_size=0.0,
+    #                   val_size=0.1,
+    #                   default_root_dir=tmpdir)
+    trainer = Trainer(solver=solver,
+                      max_epochs=2,
+                      accelerator='cpu',
+                      batch_size=None,
+                      train_size=0.9,
+                      test_size=0.0,
+                      val_size=0.1,
+                      default_root_dir=tmpdir)
     trainer.train()
+    # new_solver = SupervisedSolver.load_from_checkpoint(
+    #     f'{tmpdir}/lightning_logs/version_0/checkpoints/epoch=4-step=5.ckpt',
+    #     problem = problem, model=model)
+    # test_pts = LabelTensor(torch.rand(20, 2), problem.input_variables)
+    # assert new_solver.forward(test_pts).shape == (20, 1)
+    # assert new_solver.forward(test_pts).shape == solver.forward(test_pts).shape
+    # torch.testing.assert_close(
+    #     new_solver.forward(test_pts),
+    #     solver.forward(test_pts))
+    # import shutil
+    # shutil.rmtree('tests/test_solvers/tmp')
\ No newline at end of file

From a1b6b7aadd186f892d36743642859d9af7aa4e1a Mon Sep 17 00:00:00 2001
From: Dario Coscia <dariocos99@gmail.com>
Date: Tue, 28 Jan 2025 14:05:30 +0100
Subject: [PATCH 55/55] adding small updates

---
 pina/data/data_module.py                     | 40 +++++++++++---------
 tests/test_solvers/test_supervised_solver.py | 32 ++++++----------
 2 files changed, 34 insertions(+), 38 deletions(-)

diff --git a/pina/data/data_module.py b/pina/data/data_module.py
index c58d8239f..c9f74c453 100644
--- a/pina/data/data_module.py
+++ b/pina/data/data_module.py
@@ -1,5 +1,6 @@
 import logging
 from lightning.pytorch import LightningDataModule
+import math
 import torch
 from ..label_tensor import LabelTensor
 from torch.utils.data import DataLoader, BatchSampler, SequentialSampler, \
@@ -235,6 +236,26 @@ def _apply_shuffle(condition_dict, len_data):
                 dataset_dict[key].update({condition_name: data})
         return dataset_dict
 
+    def _create_dataloader(self, split, dataset):
+        # Use custom batching (good if batch size is large)
+        if self.batch_size is not None:
+            sampler = PinaSampler(dataset, self.batch_size,
+                                self.shuffle, self.automatic_batching)
+            if self.automatic_batching:
+                collate = Collator(self.find_max_conditions_lengths(split))
+
+            else:
+                collate = Collator(None, dataset)
+            return DataLoader(dataset, self.batch_size,
+                            collate_fn=collate, sampler=sampler)
+        dataloader = DummyDataloader(dataset,
+                                    self.trainer.strategy.root_device)
+        dataloader.dataset = self._transfer_batch_to_device(dataloader.dataset,
+                                            self.trainer.strategy.root_device,
+                                            0)
+        self.transfer_batch_to_device = self._transfer_batch_to_device_dummy
+        return dataloader
+    
     def find_max_conditions_lengths(self, split):
         max_conditions_lengths = {}
         for k, v in self.collector_splits[split].items():
@@ -247,23 +268,6 @@ def find_max_conditions_lengths(self, split):
                                                 self.batch_size)
         return max_conditions_lengths
 
-    def _create_dataloader(self, split, dataset):
-        # Use custom batching (good if batch size is large)
-        if self.batch_size is not None:
-            # Use default batching in torch DataLoader (good if batch size is small)
-            if self.automatic_batching:
-                collate = Collator(self.find_max_conditions_lengths(split))
-                return DataLoader(dataset, self.batch_size,
-                          collate_fn=collate)
-            collate = Collator(None)
-            sampler = PinaBatchSampler(dataset, self.batch_size, shuffle=False)
-            return DataLoader(dataset, sampler=sampler,
-                          collate_fn=collate)
-        dataloader = DummyDataloader(dataset, self.trainer.strategy.root_device)
-        dataloader.dataset = self.transfer_batch_to_device(dataloader.dataset, self.trainer.strategy.root_device, 0)
-        self.transfer_batch_to_device = self.dummy_transfer_to_device
-        return dataloader
-    
     def val_dataloader(self):
         """
         Create the validation dataloader
@@ -303,4 +307,4 @@ def _transfer_batch_to_device(self, batch, device, dataloader_idx):
             for k, v in batch.items()
         ]
 
-        return batch
+        return batch
\ No newline at end of file
diff --git a/tests/test_solvers/test_supervised_solver.py b/tests/test_solvers/test_supervised_solver.py
index 4b8d202a3..ded811a76 100644
--- a/tests/test_solvers/test_supervised_solver.py
+++ b/tests/test_solvers/test_supervised_solver.py
@@ -135,16 +135,8 @@ def test_train_load():
     solver = SupervisedSolver(problem=problem,
                               model=model,
                               extra_features=None)
-    # trainer = Trainer(solver=solver,
-    #                   max_epochs=5,
-    #                   accelerator='cpu',
-    #                   batch_size=None,
-    #                   train_size=0.9,
-    #                   test_size=0.0,
-    #                   val_size=0.1,
-    #                   default_root_dir=tmpdir)
     trainer = Trainer(solver=solver,
-                      max_epochs=2,
+                      max_epochs=5,
                       accelerator='cpu',
                       batch_size=None,
                       train_size=0.9,
@@ -152,14 +144,14 @@ def test_train_load():
                       val_size=0.1,
                       default_root_dir=tmpdir)
     trainer.train()
-    # new_solver = SupervisedSolver.load_from_checkpoint(
-    #     f'{tmpdir}/lightning_logs/version_0/checkpoints/epoch=4-step=5.ckpt',
-    #     problem = problem, model=model)
-    # test_pts = LabelTensor(torch.rand(20, 2), problem.input_variables)
-    # assert new_solver.forward(test_pts).shape == (20, 1)
-    # assert new_solver.forward(test_pts).shape == solver.forward(test_pts).shape
-    # torch.testing.assert_close(
-    #     new_solver.forward(test_pts),
-    #     solver.forward(test_pts))
-    # import shutil
-    # shutil.rmtree('tests/test_solvers/tmp')
\ No newline at end of file
+    new_solver = SupervisedSolver.load_from_checkpoint(
+        f'{tmpdir}/lightning_logs/version_0/checkpoints/epoch=4-step=5.ckpt',
+        problem = problem, model=model)
+    test_pts = LabelTensor(torch.rand(20, 2), problem.input_variables)
+    assert new_solver.forward(test_pts).shape == (20, 1)
+    assert new_solver.forward(test_pts).shape == solver.forward(test_pts).shape
+    torch.testing.assert_close(
+        new_solver.forward(test_pts),
+        solver.forward(test_pts))
+    import shutil
+    shutil.rmtree('tests/test_solvers/tmp')
\ No newline at end of file