Fix bugs in DataLoader and tutorial5 notebook

pina/__init__.py

@@ -1,10 +1,10 @@
 __all__ = [
     "Trainer", "LabelTensor", "Plotter", "Condition", "SamplePointDataset",
-    "PinaDataModule", "PinaDataLoader", 'TorchOptimizer', 'Graph'
+    "PinaDataModule", "PinaDataLoader", 'TorchOptimizer', 'Graph', 'LabelParameter'
 ]
 
 from .meta import *
-from .label_tensor import LabelTensor
+from .label_tensor import LabelTensor, LabelParameter
 from .solvers.solver import SolverInterface
 from .trainer import Trainer
 from .plotter import Plotter

pina/condition/data_condition.py

@@ -15,15 +15,15 @@ class DataConditionInterface(ConditionInterface):
     """
 
     __slots__ = ["input_points", "conditional_variables"]
+    condition_type = ['unsupervised']
 
     def __init__(self, input_points, conditional_variables=None):
         """
-        TODO
+        TODO : add docstring
         """
         super().__init__()
         self.input_points = input_points
         self.conditional_variables = conditional_variables
-        self._condition_type = 'unsupervised'
 
     def __setattr__(self, key, value):
         if (key == 'input_points') or (key == 'conditional_variables'):

pina/condition/domain_equation_condition.py

@@ -13,15 +13,14 @@ class DomainEquationCondition(ConditionInterface):
     """
 
     __slots__ = ["domain", "equation"]
-
+    condition_type = ['physics']
     def __init__(self, domain, equation):
         """
-        TODO
+        TODO : add docstring
         """
         super().__init__()
         self.domain = domain
         self.equation = equation
-        self._condition_type = 'physics'
 
     def __setattr__(self, key, value):
         if key == 'domain':

pina/condition/input_equation_condition.py

@@ -14,15 +14,14 @@ class InputPointsEquationCondition(ConditionInterface):
     """
 
     __slots__ = ["input_points", "equation"]
-
+    condition_type = ['physics']
     def __init__(self, input_points, equation):
         """
-        TODO
+        TODO : add docstring
         """
         super().__init__()
         self.input_points = input_points
         self.equation = equation
-        self._condition_type = 'physics'
 
     def __setattr__(self, key, value):
         if key == 'input_points':

pina/condition/input_output_condition.py

@@ -13,15 +13,14 @@ class InputOutputPointsCondition(ConditionInterface):
     """
 
     __slots__ = ["input_points", "output_points"]
-
+    condition_type = ['supervised']
     def __init__(self, input_points, output_points):
         """
-        TODO
+        TODO : add docstring
         """
         super().__init__()
         self.input_points = input_points
         self.output_points = output_points
-        self._condition_type = ['supervised', 'physics']
 
     def __setattr__(self, key, value):
         if (key == 'input_points') or (key == 'output_points'):

pina/data/base_dataset.py

@@ -1,12 +1,12 @@
 """
 Basic data module implementation
 """
+import math
 import torch
 import logging
-
 from torch.utils.data import Dataset
-
 from ..label_tensor import LabelTensor
+from .pina_subset import PinaSubset
 
 
 class BaseDataset(Dataset):
@@ -41,7 +41,7 @@ def __init__(self, problem=None, device=torch.device('cpu')):
         super().__init__()
         self.empty = True
         self.problem = problem
-        self.device = device
+        self.device = torch.device('cpu')
         self.condition_indices = None
         for slot in self.__slots__:
             setattr(self, slot, [])
@@ -53,7 +53,7 @@ def __init__(self, problem=None, device=torch.device('cpu')):
 
     def _init_from_problem(self, collector_dict):
         """
-        TODO
+        TODO : Add docstring
         """
         for name, data in collector_dict.items():
             keys = list(data.keys())
@@ -109,14 +109,14 @@ def initialize(self):
         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)
+                torch.tensor(
+                    [self.conditions_idx[i]] * self.num_el_per_condition[i],
+                    dtype=torch.uint8)
                 for i in range(len(self.num_el_per_condition))
             ],
-                                               dim=0)
+                dim=0)
             for slot in self.__slots__:
                 current_attribute = getattr(self, slot)
                 if all(isinstance(a, LabelTensor) for a in current_attribute):
@@ -152,6 +152,50 @@ def apply_shuffle(self, indices):
             if slot != 'equation':
                 attribute = getattr(self, slot)
                 if isinstance(attribute, (LabelTensor, torch.Tensor)):
-                    setattr(self, 'slot', attribute[[indices]])
+                    setattr(self, slot, attribute[[indices]].detach())
                 if isinstance(attribute, list):
-                    setattr(self, 'slot', [attribute[i] for i in indices])
+                    setattr(self, slot, [attribute[i] for i in indices])
+        self.condition_indices = self.condition_indices[indices]
+
+    def eval_splitting_lengths(self, lengths):
+        if sum(lengths) - 1 < 1e-3:
+            len_dataset = len(self)
+            lengths = [
+                int(math.floor(len_dataset * length)) for length in lengths
+            ]
+            remainder = len(self) - 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")
+        return lengths
+
+    def dataset_split(self, 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
+        """
+
+        lengths = self.eval_splitting_lengths(lengths)
+
+        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))
+            self.apply_shuffle(indices)
+
+        offsets = [
+            sum(lengths[:i]) if i > 0 else 0 for i in range(len(lengths))
+        ]
+        return [
+            PinaSubset(self, slice(offset, offset + length))
+            for offset, length in zip(offsets, lengths)
+        ]

pina/data/data_module.py

@@ -2,15 +2,12 @@
 This module provide basic data management functionalities
 """
 
-import math
-import torch
 import logging
-from pytorch_lightning import LightningDataModule
+from lightning.pytorch 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):
@@ -23,8 +20,8 @@ def __init__(self,
                  problem,
                  device,
                  train_size=.7,
-                 test_size=.1,
-                 val_size=.2,
+                 test_size=.2,
+                 val_size=.1,
                  predict_size=0.,
                  batch_size=None,
                  shuffle=True,
@@ -61,28 +58,31 @@ def __init__(self,
         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)
+        else:
+            self.train_dataloader = super().train_dataloader
+
         if test_size > 0:
             self.split_length.append(test_size)
             self.split_names.append('test')
-            self.loader_functions['test_dataloader'] = lambda: PinaDataLoader(
-                self.splits['test'], self.batch_size, self.condition_names)
+        else:
+            self.test_dataloader = super().test_dataloader
+
         if val_size > 0:
             self.split_length.append(val_size)
             self.split_names.append('val')
-            self.loader_functions['val_dataloader'] = lambda: PinaDataLoader(
-                self.splits['val'], self.batch_size, self.condition_names)
+        else:
+            self.val_dataloader = super().val_dataloader
+
         if predict_size > 0:
             self.split_length.append(predict_size)
             self.split_names.append('predict')
-            self.loader_functions['predict_dataloader'] = lambda: PinaDataLoader(
-                self.splits['predict'], self.batch_size, self.condition_names)
+        else:
+            self.predict_dataloader = super().predict_dataloader
+
         self.splits = {k: {} for k in self.split_names}
         self.shuffle = shuffle
-
-        for k, v in self.loader_functions.items():
-            setattr(self, k, v)
+        self.has_setup_fit = False
+        self.has_setup_test = False
 
     def prepare_data(self):
         if self.datasets is None:
@@ -98,57 +98,21 @@ def setup(self, stage=None):
         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)
+                    splits = dataset.dataset_split(
+                        self.split_length,
+                        shuffle=self.shuffle)
                     for i in range(len(self.split_length)):
                         self.splits[self.split_names[i]][
                             dataset.data_type] = splits[i]
+            self.has_setup_fit = True
         elif stage == 'test':
-            raise NotImplementedError("Testing pipeline not implemented yet")
+            if self.has_setup_fit is False:
+                raise NotImplementedError(
+                    "You must call setup with stage='fit' "
+                    "first")
         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 
@@ -177,3 +141,31 @@ def _create_datasets(self):
                 dataset.initialize()
                 datasets.append(dataset)
         self.datasets = datasets
+
+    def val_dataloader(self):
+        """
+        Create the validation dataloader
+        """
+        return PinaDataLoader(self.splits['val'], self.batch_size,
+                              self.condition_names, device=self.device)
+
+    def train_dataloader(self):
+        """
+        Create the training dataloader
+        """
+        return PinaDataLoader(self.splits['train'], self.batch_size,
+                              self.condition_names, device=self.device)
+
+    def test_dataloader(self):
+        """
+        Create the testing dataloader
+        """
+        return PinaDataLoader(self.splits['test'], self.batch_size,
+                              self.condition_names, device=self.device)
+
+    def predict_dataloader(self):
+        """
+        Create the prediction dataloader
+        """
+        return PinaDataLoader(self.splits['predict'], self.batch_size,
+                              self.condition_names, device=self.device)