Repository: mathLab/PINA (language: python)

Issue:
Conditions
I agree to divide conditions in multiple classes but I do not think that the `residual` method should be inside conditions. I see conditions more like an object that encapsulates data, if we constrain also to what the model should do we can't cover all test cases. I think that the loss to minimize should be defined only inside the solver (physics-informed, data-driven whatever), in this way is user responsibility to know how to use the data

Maybe we can think more of something like PhysicsConditions, DataConditions, GraphConditions (like the division in the dataloading)... These conditions are needed to call loss_data, loss_phys, and loss_graph (which are defined inside the solver). This helps us to do two things:
1. We know which functions need to be defined in the solver by looking at the conditions the problem has (we can raise errors before training)
2. We don't need to cover all classes of possible ways residuals are defined, we only need to pass the data that are needed to compute it.

_Originally posted by @dario-coscia in https://github.com/mathLab/PINA/pull/312#discussion_r1730291720_
            


Discussion / hints:

Fix patch:
diff --git a/pina/__init__.py b/pina/__init__.py
index 0fe93752..3bc28ae6 100644
--- a/pina/__init__.py
+++ b/pina/__init__.py
@@ -1,11 +1,6 @@
 __all__ = [
-    "PINN",
-    "Trainer",
-    "LabelTensor",
-    "Plotter",
-    "Condition",
-    "SamplePointDataset",
-    "SamplePointLoader",
+    "Trainer", "LabelTensor", "Plotter", "Condition", "SamplePointDataset",
+    "PinaDataModule", "PinaDataLoader", 'TorchOptimizer', 'Graph'
 ]
 
 from .meta import *
@@ -15,4 +10,8 @@
 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
+from .optim import TorchOptimizer
+from .optim import TorchScheduler
+from .graph import Graph
diff --git a/pina/collector.py b/pina/collector.py
new file mode 100644
index 00000000..3219b2b6
--- /dev/null
+++ b/pina/collector.py
@@ -0,0 +1,114 @@
+from .utils import check_consistency, merge_tensors
+
+
+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,
+        #            '[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)
+        }
+
+        # 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())
+
+    @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
+
+    @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 (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
+                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')
+
+    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)
diff --git a/pina/condition/__init__.py b/pina/condition/__init__.py
index ff329c1b..4c89b75d 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 d815838b..3a62143e 100644
--- a/pina/condition/condition.py
+++ b/pina/condition/condition.py
@@ -1,27 +1,22 @@
 """ 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 +24,42 @@ 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
-
-    #     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__}."
-    #         )
+    __slots__ = list(
+        set(InputOutputPointsCondition.__slots__ +
+            InputPointsEquationCondition.__slots__ +
+            DomainEquationCondition.__slots__ +
+            DataConditionInterface.__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"]
-            )
-        elif sorted(kwargs.keys()) == sorted(["domain", "output_points"]):
-            return DomainOutputCondition(**kwargs)
-        elif sorted(kwargs.keys()) == sorted(["domain", "equation"]):
+        if len(args) != 0:
+            raise ValueError("Condition takes only the following keyword "
+                             f"arguments: {Condition.__slots__}.")
+
+        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
diff --git a/pina/condition/condition_interface.py b/pina/condition/condition_interface.py
index 0626a6d8..f2fe5db9 100644
--- a/pina/condition/condition_interface.py
+++ b/pina/condition/condition_interface.py
@@ -1,21 +1,33 @@
-
-from abc import ABCMeta, abstractmethod
+from abc import ABCMeta
 
 
 class ConditionInterface(metaclass=ABCMeta):
 
-    def __init__(self) -> None:
+    condition_types = ['physics', 'supervised', 'unsupervised']
+
+    def __init__(self, *args, **kwargs):
+        self._condition_type = None
         self._problem = None
 
-    @abstractmethod
-    def residual(self, model):
-        """
-        Compute the residual of the condition.
+    @property
+    def problem(self):
+        return self._problem
+
+    @problem.setter
+    def problem(self, value):
+        self._problem = value
 
-        :param model: The model to evaluate the condition.
-        :return: The residual of the condition.
-        """
-        pass
+    @property
+    def condition_type(self):
+        return self._condition_type
 
-    def set_problem(self, problem):
-        self._problem = problem
+    @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
diff --git a/pina/condition/data_condition.py b/pina/condition/data_condition.py
new file mode 100644
index 00000000..c6777231
--- /dev/null
+++ b/pina/condition/data_condition.py
@@ -0,0 +1,33 @@
+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__ = ["input_points", "conditional_variables"]
+
+    def __init__(self, input_points, conditional_variables=None):
+        """
+        TODO
+        """
+        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'):
+            check_consistency(value, (LabelTensor, Graph, torch.Tensor))
+            DataConditionInterface.__dict__[key].__set__(self, value)
+        elif key in ('_problem', '_condition_type'):
+            super().__setattr__(key, value)
diff --git a/pina/condition/domain_equation_condition.py b/pina/condition/domain_equation_condition.py
index 15df3f85..58dca70b 100644
--- a/pina/condition/domain_equation_condition.py
+++ b/pina/condition/domain_equation_condition.py
@@ -1,34 +1,34 @@
+import torch
+
 from .condition_interface import ConditionInterface
+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.
-
-        :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
+    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)
+        elif key in ('_problem', '_condition_type'):
+            super().__setattr__(key, value)
diff --git a/pina/condition/domain_output_condition.py b/pina/condition/domain_output_condition.py
deleted file mode 100644
index f847720b..00000000
--- 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 288022c0..bf05130c 100644
--- a/pina/condition/input_equation_condition.py
+++ b/pina/condition/input_equation_condition.py
@@ -1,23 +1,37 @@
+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
+    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)
+        elif key in ('_problem', '_condition_type'):
+            super().__setattr__(key, value)
diff --git a/pina/condition/input_output_condition.py b/pina/condition/input_output_condition.py
new file mode 100644
index 00000000..08ed21d9
--- /dev/null
+++ b/pina/condition/input_output_condition.py
@@ -0,0 +1,31 @@
+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']
+
+    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'):
+            super().__setattr__(key, value)
diff --git a/pina/data/__init__.py b/pina/data/__init__.py
index fba19b92..2b3a126a 100644
--- a/pina/data/__init__.py
+++ b/pina/data/__init__.py
@@ -1,7 +1,15 @@
+"""
+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 00000000..2c28ba30
--- /dev/null
+++ b/pina/data/base_dataset.py
@@ -0,0 +1,157 @@
+"""
+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_dataset.py b/pina/data/data_dataset.py
deleted file mode 100644
index 9dff2d7e..00000000
--- 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 00000000..bd117b54
--- /dev/null
+++ b/pina/data/data_module.py
@@ -0,0 +1,179 @@
+"""
+This module provide basic data management functionalities
+"""
+
+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
+
+
+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=.1,
+                 val_size=.2,
+                 predict_size=0.,
+                 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
+        """
+        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.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')
+            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
+        """
+        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]
+        elif stage == 'test':
+            raise NotImplementedError("Testing pipeline not implemented yet")
+        else:
+            raise ValueError("stage must be either 'fit' or 'test'")
+
+    @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)
+
+        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)
+        ]
+
+    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 cb1296ed..c5d1b61d 100644
--- a/pina/data/pina_batch.py
+++ b/pina/data/pina_batch.py
@@ -1,36 +1,47 @@
+"""
+Batch management module
+"""
+from .pina_subset import PinaSubset
 
 
 class Batch:
     """
-    This class is used to create a dataset of sample points.
+    Implementation of the Batch class used during training to perform SGD
+    optimization.
     """
 
-    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]
+    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)
 
-        elif type_ == "data":
+        for k, v in idx_dict.items():
+            setattr(self, k + '_idx', v)
+        self.require_grad = require_grad
 
-            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 __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
index 2c8967c5..e2d3fb76 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,55 @@ 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 00000000..275541e9
--- /dev/null
+++ b/pina/data/pina_subset.py
@@ -0,0 +1,36 @@
+"""
+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
index 84af2920..bc3bca33 100644
--- a/pina/data/sample_dataset.py
+++ b/pina/data/sample_dataset.py
@@ -1,43 +1,35 @@
-from torch.utils.data import Dataset
-import torch
+"""
+Sample dataset module
+"""
+from copy import deepcopy
+from .base_dataset import BaseDataset
+from ..condition import InputPointsEquationCondition
 
-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__ = 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
new file mode 100644
index 00000000..be601050
--- /dev/null
+++ b/pina/data/supervised_dataset.py
@@ -0,0 +1,13 @@
+"""
+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 00000000..18cf296f
--- /dev/null
+++ b/pina/data/unsupervised_dataset.py
@@ -0,0 +1,14 @@
+"""
+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 51270549..5fe99d64 100644
--- a/pina/domain/cartesian.py
+++ b/pina/domain/cartesian.py
@@ -30,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.
@@ -164,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]
 
@@ -199,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/difference_domain.py b/pina/domain/difference_domain.py
index d2ba414f..4015a386 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 not in 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 5906d285..916bf3e9 100644
--- a/pina/domain/domain_interface.py
+++ b/pina/domain/domain_interface.py
@@ -9,6 +9,37 @@ 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
+
+    @property
+    @abstractmethod
+    def variables(self):
+        """
+        Abstract method returing Domain variables.
+        """
+        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 dcbc5554..18e28d4b 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)
@@ -71,6 +70,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.
@@ -281,7 +284,7 @@ def _single_points_sample(n, variables):
         if variables == "all":
             variables = list(self.range_.keys()) + list(self.fixed_.keys())
 
-        if mode in ["random"]:
+        if mode in self.sample_modes:
             return _Nd_sampler(n, mode, 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 ed63db31..a05b1543 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 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 b40d3695..bb0499b5 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 not in 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 acd4cf44..0300f524 100644
--- a/pina/domain/operation_interface.py
+++ b/pina/domain/operation_interface.py
@@ -24,6 +24,10 @@ def __init__(self, geometries):
         # assign geometries
         self._geometries = geometries
 
+    @property
+    def sample_modes(self):
+        return ["random"]
+
     @property
     def geometries(self):
         """
diff --git a/pina/domain/simplex.py b/pina/domain/simplex.py
index 8d26422a..931f861a 100644
--- a/pina/domain/simplex.py
+++ b/pina/domain/simplex.py
@@ -74,6 +74,10 @@ def __init__(self, simplex_matrix, sample_surface=False):
         # build cartesian_bound
         self._cartesian_bound = self._build_cartesian(self._vertices_matrix)
 
+    @property
+    def sample_modes(self):
+        return ["random"]
+
     @property
     def variables(self):
         return self._vertices_matrix.labels
@@ -88,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)
@@ -141,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):
@@ -231,7 +234,7 @@ def sample(self, n, mode="random", variables="all"):
             in ``variables``.
         """
 
-        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 da2ead90..0af8e1bd 100644
--- a/pina/domain/union_domain.py
+++ b/pina/domain/union_domain.py
@@ -33,6 +33,19 @@ def __init__(self, geometries):
         """
         super().__init__(geometries)
 
+    @property
+    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
+        return list(set(variables))
+
     def is_inside(self, point, check_border=False):
         """
         Check if a point is inside the ``Union`` domain.
diff --git a/pina/label_tensor.py b/pina/label_tensor.py
index 7646dd8a..719975c5 100644
--- a/pina/label_tensor.py
+++ b/pina/label_tensor.py
@@ -1,13 +1,18 @@
 """ Module for LabelTensor """
-
+from copy import copy, deepcopy
 import torch
 from torch import Tensor
 
+
 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))
+    if isinstance(a, range) and isinstance(b, range):
+        return a.start <= b.start and a.stop >= b.stop
+    return False
 
 
 class LabelTensor(torch.Tensor):
@@ -15,16 +20,16 @@ class LabelTensor(torch.Tensor):
 
     @staticmethod
     def __new__(cls, x, labels, *args, **kwargs):
-        return super().__new__(cls, x, *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 __len__(self) -> int:
-        return super().__len__()
-
-    def __init__(self, x, labels):
+    def __init__(self, x, labels, **kwargs):
         """
         Construct a `LabelTensor` by passing a dict of the labels
 
@@ -35,18 +40,145 @@ def __init__(self, x, labels):
                     {1: {"name": "space"['a', 'b', 'c'])
 
         """
-        self.labels = None
+        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
+
+    @property
+    def labels(self):
+        """Property decorator for labels
+
+        :return: labels of self
+        :rtype: list
+        """
+        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
+        """
+        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 not hasattr(self, '_labels'):
+            self._labels = {}
         if isinstance(labels, dict):
-            self.update_labels(labels)
+            self._init_labels_from_dict(labels)
         elif isinstance(labels, list):
-            self.init_labels_from_list(labels)
+            self._init_labels_from_list(labels)
         elif isinstance(labels, str):
             labels = [labels]
-            self.init_labels_from_list(labels)
+            self._init_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 _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:
+                raise 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
 
-    def extract(self, label_to_extract):
+        :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.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``.
@@ -56,106 +188,123 @@ 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 = 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:
+        # 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')
-        return LabelTensor(new_tensor, new_labels)
+
+        # 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)]
+
+        # Loop over labels_to_extract dict
+        for k, v in labels_to_extract.items():
+
+            # 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)
+
+            updated_labels.update({idx_dim: {'dof': v, 'name': k}})
+
+        tensor = self.tensor
+        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():
+        for key, value in self._labels.items():
             s += f"{key}: {value}\n"
         s += '\n'
-        s += super().__str__()
+        s += self.tensor.__str__()
         return s
 
     @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
         :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]
-        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')
+        # Perform cat on tensors
         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)
+
+        # Update labels
+        labels = LabelTensor.__create_labels_cat(tensors, dim)
+
+        return LabelTensor.__internal_init__(new_tensor, labels,
+                                             tensors[0].dim_names)
+
+    @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 requires_grad_(self, mode=True):
         lt = super().requires_grad_(mode)
-        lt.labels = self.labels
+        lt.labels = self._labels
         return lt
 
     @property
@@ -180,38 +329,161 @@ 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 update_labels(self, labels):
+    @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)
+        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
+            new_label_tensor = LabelTensor.cat([self, tensor],
+                                               dim=self.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.ndim - 1)
+        else:
+            raise ValueError('mode must be either "std" or "cross"')
+        return new_label_tensor
+
+    @staticmethod
+    def vstack(label_tensors):
         """
-        Update the internal label representation according to the values passed as input.
+        Stack tensors vertically. For more details, see
+        :meth:`torch.vstack`.
 
-        :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_]),
-                '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)
+        :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 init_labels_from_list(self, labels):
+    def __getitem__(self, index):
         """
-        Given a list of dof, this method update the internal label representation
+        TODO: Complete docstring
+        :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)
 
-        :param labels: The label(s) to update.
-        :type labels: list
+        selected_lt = super().__getitem__(index)
+
+        if isinstance(index, (int, slice)):
+            index = [index]
+
+        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):
+                    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_single_label(old_labels, to_update_labels, index, dim):
         """
-        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
+        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
+        labels = self.stored_labels[dim]['dof']
+        sorted_index = arg_sort(labels)
+        indexer = [slice(None)] * self.ndim
+        indexer[dim] = sorted_index
+        return self.__getitem__(indexer)
+
+    def __deepcopy__(self, memo):
+        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/model/avno.py b/pina/model/avno.py
index 2ac3b3f7..e27ce1f1 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 077a6b92..b09f1aec 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/operators.py b/pina/operators.py
index 17b45d81..0b306dfb 100644
--- a/pina/operators.py
+++ b/pina/operators.py
@@ -5,9 +5,7 @@
 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 pina.label_tensor import LabelTensor
 
 
@@ -58,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,
@@ -87,16 +85,13 @@ def grad_scalar_output(output_, input_, d):
             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
 
@@ -142,16 +137,14 @@ def div(output_, input_, components=None, d=None):
         raise ValueError
 
     grad_output = grad(output_, input_, components, d)
-    div = torch.zeros(input_.shape[0], 1, device=output_.device)
     labels = [None] * len(components)
+    tensors_to_sum = []
     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)
+        tensors_to_sum.append(grad_output.extract(c_fields))
         labels[i] = c_fields
-
-    div = div.as_subclass(LabelTensor)
-    div.labels = ["+".join(labels)]
-    return div
+    div_result = LabelTensor.summation(tensors_to_sum)
+    return div_result
 
 
 def laplacian(output_, input_, components=None, d=None, method="std"):
@@ -194,19 +187,19 @@ def laplacian(output_, input_, components=None, d=None, method="std"):
 
         if len(components) == 1:
             grad_output = grad(output_, input_, components=components, d=d)
-            result = torch.zeros(output_.shape[0], 1, device=output_.device)
+            to_append_tensors = []
             for i, label in enumerate(grad_output.labels):
                 gg = grad(grad_output, input_, d=d, components=[label])
-                result[:, 0] += super(torch.Tensor, gg.T).__getitem__(
-                    i
-                )  # TODO improve
+                to_append_tensors.append(gg.extract([gg.labels[i]]))
             labels = [f"dd{components[0]}"]
-
+            result = LabelTensor.summation(tensors=to_append_tensors)
+            result.labels = labels
         else:
-            result = torch.empty(
-                input_.shape[0], len(components), device=output_.device
-            )
+            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):
@@ -216,10 +209,11 @@ def laplacian(output_, input_, components=None, d=None, method="std"):
 
                 grad_output = grad(output_, input_, components=ci, d=di)
                 result[:, idx] = grad(grad_output, input_, d=di).flatten()
-                labels[idx] = f"dd{ci}dd{di}"
-
-    result = result.as_subclass(LabelTensor)
-    result.labels = labels
+                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
 
 
@@ -249,11 +243,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/optim/torch_optimizer.py b/pina/optim/torch_optimizer.py
index 239819a4..54818d5a 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):
@@ -12,8 +13,8 @@ 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, **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 50e1d91f..9aa187d0 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/problem/abstract_problem.py b/pina/problem/abstract_problem.py
index bf0a2dcc..6897fbb7 100644
--- a/pina/problem/abstract_problem.py
+++ b/pina/problem/abstract_problem.py
@@ -1,11 +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,26 +20,42 @@ 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 = {}
+        # 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()
+        self._batching_dimension = 0
 
-        # # 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
+    @property
+    def collector(self):
+        return self._collector
 
-        # # put in self.input_pts all the points that we don't need to sample
-        self._span_condition_points()
+    @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
+    def input_pts(self):
+        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
 
     def __deepcopy__(self, memo):
         """
@@ -85,19 +101,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,35 +117,11 @@ 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"
-    ):
+    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.
@@ -172,103 +151,44 @@ 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
-        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 {self.conditions}.",
-            )
-
-        # 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[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()
+                if isinstance(self.conditions[name], DomainEquationCondition)
+            ]
+        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 isinstance(self.conditions[loc], DomainEquationCondition)]}.",
                 )
-                new_pts.labels = old_pts.labels
 
-            # merging
-            merged_pts = torch.vstack([old_pts, new_points[location]])
-            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)
+
+    def add_points(self, new_points_dict):
+        self.collector.add_points(new_points_dict)
diff --git a/pina/problem/inverse_problem.py b/pina/problem/inverse_problem.py
index 5a83566a..51cbd3ca 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/pina/solvers/pinns/basepinn.py b/pina/solvers/pinns/basepinn.py
index 9e841d65..543f823f 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.
@@ -177,7 +177,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 a27e9364..e00bc8d5 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
@@ -179,13 +17,13 @@ class SolverInterface(pytorch_lightning.LightningModule, metaclass=ABCMeta):
     LightningModule methods.
     """
 
-    def __init__(
-        self,
-        model,
-        problem,
-        optimizer,
-        scheduler,
-    ):
+    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
@@ -197,43 +35,66 @@ 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, model in enumerate(models):
+                models[idx] = Network(
+                    model=model,
+                    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:
-            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
+
+        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):
         pass
 
     @abstractmethod
-    def training_step(self):
+    def training_step(self, batch, batch_idx):
         pass
 
     @abstractmethod
@@ -244,13 +105,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 +133,13 @@ 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 '
+                    f'{condition.condition_type}')
diff --git a/pina/solvers/supervised.py b/pina/solvers/supervised.py
index c44d5a1e..62fc9914 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,15 +37,16 @@ class SupervisedSolver(SolverInterface):
     we are seeking to approximate multiple (discretised) functions given
     multiple (discretised) input functions.
     """
-
-    def __init__(
-        self,
-        problem,
-        model,
-        loss=None,
-        optimizer=None,
-        scheduler=None,
-    ):
+    accepted_condition_types = ['supervised']
+    __name__ = 'SupervisedSolver'
+
+    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.
@@ -57,11 +56,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()
@@ -70,22 +66,20 @@ 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__(
-            model=model,
-            problem=problem,
-            optimizer=optimizer,
-            scheduler=scheduler,
-        )
+        super().__init__(models=model,
+                         problem=problem,
+                         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 +91,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):
@@ -113,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.
@@ -128,34 +115,28 @@ 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.")
 
             # 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]
 
             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)
@@ -167,8 +148,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 +162,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 758bbaaf..58c66f67 100644
--- a/pina/trainer.py
+++ b/pina/trainer.py
@@ -3,13 +3,19 @@
 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,35 +37,23 @@ def __init__(self, solver, batch_size=None, **kwargs):
         check_consistency(solver, SolverInterface)
         if batch_size is not None:
             check_consistency(batch_size, int)
-
-        self._model = solver
+        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()
 
-        # 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(
-                    pb.unknown_parameters[key].data.to(device)
-                )
+                    pb.unknown_parameters[key].data.to(device))
 
     def _create_loader(self):
         """
@@ -67,29 +61,45 @@ 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)
-        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,
+                                     val_size=self.eval_size)
+        data_module.setup()
+        self._loader = data_module.train_dataloader()
 
     def train(self, **kwargs):
         """
         Train the solver method.
         """
-        return super().fit(
-            self._model, train_dataloaders=self._loader, **kwargs
-        )
+        return super().fit(self.solver,
+                           train_dataloaders=self._loader,
+                           **kwargs)
 
     @property
     def solver(self):
         """
         Returning trainer solver.
         """
-        return self._model
+        return self._solver
+
+    @solver.setter
+    def solver(self, solver):
+        self._solver = solver
diff --git a/pina/utils.py b/pina/utils.py
index 282dd533..84d3e741 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)
 
 


Test patch:
diff --git a/tests/test_condition.py b/tests/test_condition.py
index 5f1c6236..9165c3fa 100644
--- a/tests/test_condition.py
+++ b/tests/test_condition.py
@@ -18,27 +18,30 @@ 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)
 
 
-def test_init_locfunc():
-    Condition(location=example_domain, equation=FixedValue(0.0))
+test_init_inputoutput()
+
+
+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_dataset.py b/tests/test_dataset.py
index 40f21922..87fd9a15 100644
--- a/tests/test_dataset.py
+++ b/tests/test_dataset.py
@@ -1,122 +1,227 @@
+import math
 import torch
-import pytest
-
-from pina.data.dataset import SamplePointDataset, SamplePointLoader, DataPointDataset
+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.model import FeedForward
+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))
+    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_)
+        '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.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)
+    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():
-    sample_dataset = SamplePointDataset(poisson, device='cpu')
-    data_dataset = DataPointDataset(poisson, device='cpu')
-    loader = SamplePointLoader(sample_dataset, data_dataset, batch_size=10)
+    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']
 
-    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)
+    conditions = {
+        'graph_data': Condition(input_points=input, output_points=output)
+    }
 
-    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)
 
-    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']
+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_geometry/test_cartesian.py b/tests/test_geometry/test_cartesian.py
index 65026c33..fc30757b 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]})
 
diff --git a/tests/test_geometry/test_simplex.py b/tests/test_geometry/test_simplex.py
index 7fc34ce2..25224aae 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.py
deleted file mode 100644
index f87d3abb..00000000
--- a/tests/test_label_tensor.py
+++ /dev/null
@@ -1,148 +0,0 @@
-import torch
-import pytest
-
-from pina.label_tensor import LabelTensor
-#import pina
-
-data = torch.rand((20, 3))
-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'])
-
-@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(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", [{'samples': [2]}])
-def test_extract_row(labels, labels_te):
-    tensor = LabelTensor(data, labels)
-    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)
-    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)
-    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_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)
diff --git a/tests/test_label_tensor/test_label_tensor.py b/tests/test_label_tensor/test_label_tensor.py
new file mode 100644
index 00000000..61e47995
--- /dev/null
+++ b/tests/test_label_tensor/test_label_tensor.py
@@ -0,0 +1,282 @@
+import torch
+import pytest
+
+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_list = ['x', 'y', 'z']
+labels_all = labels_column | labels_row
+
+
+@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):
+    tensor = LabelTensor(data, labels)
+    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", [{'samples': [2]}])
+def test_extract_row(labels, labels_te):
+    tensor = LabelTensor(data, labels)
+    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)
+    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():
+    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)
+    new = tensor.extract(labels_te)
+    tensor2 = LabelTensor(data, labels)
+    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))
+    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']
+    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.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']
+    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.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']
+    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.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']
+    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(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 = ['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'] == ['x', 'y', 'z', 'w', 'a']
+    assert lt_cat.full_labels[0]['dof'] == range(20)
+    assert lt_cat.full_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.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)
+    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[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()
+
+
+def test_append_3D():
+    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)
+    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, 2)
+    labels_2 = ['z', 'w']
+    lt2 = LabelTensor(data_2, labels_2)
+    lt1 = lt1.append(lt2, mode='cross')
+    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 00000000..57aafb8c
--- /dev/null
+++ b/tests/test_label_tensor/test_label_tensor_01.py
@@ -0,0 +1,118 @@
+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))
diff --git a/tests/test_operators.py b/tests/test_operators.py
index aa9ea9d8..e446c8f3 100644
--- a/tests/test_operators.py
+++ b/tests/test_operators.py
@@ -10,15 +10,14 @@ def func_vec(x):
 
 
 def func_scalar(x):
-    print('X')
     x_ = x.extract(['x'])
     y_ = x.extract(['y'])
     mu_ = x.extract(['mu'])
     return x_**2 + y_**2 + mu_**3
 
 
-data = torch.rand((20, 3), requires_grad=True)
-inp = LabelTensor(data, ['x', 'y', 'mu'])
+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])
@@ -31,7 +30,7 @@ def test_grad_scalar_output():
         f'd{tensor_s.labels[0]}d{i}' for i in inp.labels
     ]
     grad_tensor_s = grad(tensor_s, inp, d=['x', 'y'])
-    assert grad_tensor_s.shape == (inp.shape[0], 2)
+    assert grad_tensor_s.shape == (20, 2)
     assert grad_tensor_s.labels == [
         f'd{tensor_s.labels[0]}d{i}' for i in ['x', 'y']
     ]
diff --git a/tests/test_optimizer.py b/tests/test_optimizer.py
index 489bbdc0..bdc87cac 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 3e1e1ee2..f1aafcd9 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,12 +67,9 @@ def poisson_sol(self, pts):
     truth_solution = poisson_sol
 
 
-# make the problem
-poisson_problem = Poisson()
-
-
 def test_discretise_domain():
     n = 10
+    poisson_problem = Poisson()
     boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4']
     poisson_problem.discretise_domain(n, 'grid', locations=boundaries)
     for b in boundaries:
@@ -92,31 +89,68 @@ 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
+    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
+    assert poisson_problem.collector._is_conditions_ready['D'] is False
 
 
-# def test_sampling_all_args():
-#     n = 10
-#     poisson_problem.discretise_domain(n, 'grid', locations=['D'])
+def test_variables_correct_order_sampling():
+    n = 10
+    poisson_problem = Poisson()
+    poisson_problem.discretise_domain(n,
+                                      'grid',
+                                      locations=['D'],
+                                      variables=['x'])
+    poisson_problem.discretise_domain(n,
+                                      'grid',
+                                      locations=['D'],
+                                      variables=['y'])
+    assert poisson_problem.input_pts['D'].labels == sorted(
+        poisson_problem.input_variables)
 
-# def test_sampling_all_kwargs():
-#     n = 10
-#     poisson_problem.discretise_domain(n=n, mode='latin', 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,
+                                      'grid',
+                                      locations=['D'],
+                                      variables=['y'])
+    poisson_problem.discretise_domain(n,
+                                      'grid',
+                                      locations=['D'],
+                                      variables=['x'])
+    assert poisson_problem.input_pts['D'].labels == sorted(
+        poisson_problem.input_variables)
 
-# def test_sampling_dict():
-#     n = 10
-#     poisson_problem.discretise_domain(
-#         {'variables': ['x', 'y'], 'mode': 'grid', 'n': n}, locations=['D'])
 
-# def test_sampling_mixed_args_kwargs():
-#     n = 10
-#     with pytest.raises(ValueError):
-#         poisson_problem.discretise_domain(n, mode='latin', locations=['D'])
+def test_add_points():
+    poisson_problem = Poisson()
+    poisson_problem.discretise_domain(0,
+                                      'random',
+                                      locations=['D'],
+                                      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'))
+
+
+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
diff --git a/tests/test_solvers/test_supervised_solver.py b/tests/test_solvers/test_supervised_solver.py
index 912480bb..8ceadcd9 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()