Update solvers (#434)

* Enable DDP training with batch_size=None and add validity check for split sizes
* Refactoring SolverInterfaces (#435)
* Solver update + weighting
* Updating PINN for 0.2
* Modify GAROM + tests
* Adding more versatile loggers
* Disable compilation when running on Windows
* Fix tests

---------
Co-authored-by: Dario Coscia <[email protected]>
Co-authored-by: giovanni <[email protected]>
Co-authored-by: FilippoOlivo <[email protected]>

Author: 3 people

pina/callbacks/processing_callbacks.py

@@ -1,7 +1,5 @@
 """PINA Callbacks Implementations"""
 
-from lightning.pytorch.core.module import LightningModule
-from lightning.pytorch.trainer.trainer import Trainer
 import torch
 import copy
 
@@ -16,66 +14,64 @@ class MetricTracker(Callback):
 
     def __init__(self, metrics_to_track=None):
         """
-        PINA Implementation of a Lightning Callback for Metric Tracking.
+        Lightning Callback for Metric Tracking.
 
-        This class provides functionality to track relevant metrics during
-        the training process.
+        Tracks specific metrics during the training process.
 
-        :ivar _collection: A list to store collected metrics after each
-        training epoch.
+        :ivar _collection: A list to store collected metrics after each epoch.
 
-        :param trainer: The trainer object managing the training process.
-        :type trainer: pytorch_lightning.Trainer
-
-        :return: A dictionary containing aggregated metric values.
-        :rtype: dict
-
-        Example:
-            >>> tracker = MetricTracker()
-            >>> # ... Perform training ...
-            >>> metrics = tracker.metrics
+        :param metrics_to_track: List of metrics to track. Defaults to train/val loss.
+        :type metrics_to_track: list, optional
         """
         super().__init__()
         self._collection = []
-        if metrics_to_track is not None:
-            metrics_to_track = ['train_loss_epoch', 'train_loss_step', 'val_loss']
-        self.metrics_to_track = metrics_to_track
+        # Default to tracking 'train_loss' and 'val_loss' if not specified
+        self.metrics_to_track = metrics_to_track or ['train_loss', 'val_loss']
 
     def on_train_epoch_end(self, trainer, pl_module):
         """
         Collect and track metrics at the end of each training epoch.
 
         :param trainer: The trainer object managing the training process.
         :type trainer: pytorch_lightning.Trainer
-        :param pl_module: Placeholder argument.
+        :param pl_module: The model being trained (not used here).
         """
-        super().on_train_epoch_end(trainer, pl_module)
+        # Track metrics after the first epoch onwards
         if trainer.current_epoch > 0:
-            self._collection.append(
-                copy.deepcopy(trainer.logged_metrics)
-            )  # track them
+            # Append only the tracked metrics to avoid unnecessary data
+            tracked_metrics = {
+                k: v for k, v in trainer.logged_metrics.items() 
+                if k in self.metrics_to_track
+            }
+            self._collection.append(copy.deepcopy(tracked_metrics))
 
     @property
     def metrics(self):
         """
-        Aggregate collected metrics during training.
+        Aggregate collected metrics over all epochs.
 
         :return: A dictionary containing aggregated metric values.
         :rtype: dict
         """
-        common_keys = set.intersection(*map(set, self._collection))
-        v = {
+        if not self._collection:
+            return {}
+
+        # Get intersection of keys across all collected dictionaries
+        common_keys = set(self._collection[0]).intersection(*self._collection[1:])
+        
+        # Stack the metric values for common keys and return
+        return {
             k: torch.stack([dic[k] for dic in self._collection])
-            for k in common_keys
+            for k in common_keys if k in self.metrics_to_track
         }
-        return v
+
 
 
 class PINAProgressBar(TQDMProgressBar):
 
     BAR_FORMAT = "{l_bar}{bar}| {n_fmt}/{total_fmt} [{elapsed}<{remaining}, {rate_noinv_fmt}{postfix}]"
 
-    def __init__(self, metrics="val_loss", **kwargs):
+    def __init__(self, metrics="val", **kwargs):
         """
         PINA Implementation of a Lightning Callback for enriching the progress
         bar.
@@ -131,14 +127,6 @@ def get_metrics(self, trainer, model):
             pbar_metrics = {
                 key: pbar_metrics[key] for key in self._sorted_metrics
             }
-        duplicates = list(standard_metrics.keys() & pbar_metrics.keys())
-        if duplicates:
-            rank_zero_warn(
-                f"The progress bar already tracks a metric with the name(s) '{', '.join(duplicates)}' and"
-                f" `self.log('{duplicates[0]}', ..., prog_bar=True)` will overwrite this value. "
-                " If this is undesired, change the name or override `get_metrics()` in the progress bar callback.",
-            )
-
         return {**standard_metrics, **pbar_metrics}
 
     def on_fit_start(self, trainer, pl_module):
@@ -154,7 +142,7 @@ def on_fit_start(self, trainer, pl_module):
         for key in self._sorted_metrics:
             if (
                 key not in trainer.solver.problem.conditions.keys()
-                and key != "mean"
+                and key != "train" and key != "val"
             ):
                 raise KeyError(f"Key '{key}' is not present in the dictionary")
         # add the loss pedix

pina/collector.py

@@ -1,5 +1,4 @@
-from . import LabelTensor
-from .utils import check_consistency, merge_tensors
+from .utils import check_consistency
 
 
 class Collector:
@@ -8,11 +7,6 @@ def __init__(self, problem):
         # creating a hook between collector and problem
         self.problem = problem
 
-        # this variable is used to store the data in the form:
-        # {'[condition_name]' :
-        #           {'input_points' : Tensor,
-        #            '[equation/output_points/conditional_variables]': Tensor}
-        # }
         # those variables are used for the dataloading
         self._data_collections = {name: {} for name in self.problem.conditions}
         self.conditions_name = {

pina/data/data_module.py

@@ -1,6 +1,5 @@
 import logging
 from lightning.pytorch import LightningDataModule
-import math
 import torch
 from ..label_tensor import LabelTensor
 from torch.utils.data import DataLoader, BatchSampler, SequentialSampler, \
@@ -10,8 +9,38 @@
 from ..collector import Collector
 
 class DummyDataloader:
-    def __init__(self, dataset, device):
-        self.dataset = dataset.get_all_data()
+    """"
+    Dummy dataloader used when batch size is None. It callects all the data
+    in self.dataset and returns it when it is called a single batch.
+    """
+
+    def __init__(self, dataset):
+        """
+        param dataset: The dataset object to be processed.
+        :notes:
+            - **Distributed Environment**:
+                - Divides the dataset across processes using the
+                    rank and world size.
+                - Fetches only the portion of data corresponding to
+                    the current process.
+            - **Non-Distributed Environment**:
+                - Fetches the entire dataset.
+        """
+        if (torch.distributed.is_available() and
+                torch.distributed.is_initialized()):
+            rank = torch.distributed.get_rank()
+            world_size = torch.distributed.get_world_size()
+            if len(dataset) < world_size:
+                raise RuntimeError(
+                    "Dimension of the dataset smaller than world size."
+                    " Increase the size of the partition or use a single GPU")
+            idx, i = [], rank
+            while i < len(dataset):
+                idx.append(i)
+                i += world_size
+            self.dataset = dataset.fetch_from_idx_list(idx)
+        else:
+            self.dataset = dataset.get_all_data()
 
     def __iter__(self):
         return self
@@ -50,7 +79,7 @@ def _collate_standard_dataloader(self, batch):
             for arg in condition_args:
                 data_list = [batch[idx][condition_name][arg] for idx in range(
                     min(len(batch),
-                    self.max_conditions_lengths[condition_name]))]
+                        self.max_conditions_lengths[condition_name]))]
                 if isinstance(data_list[0], LabelTensor):
                     single_cond_dict[arg] = LabelTensor.stack(data_list)
                 elif isinstance(data_list[0], torch.Tensor):
@@ -61,7 +90,6 @@ def _collate_standard_dataloader(self, batch):
             batch_dict[condition_name] = single_cond_dict
         return batch_dict
 
-
     def __call__(self, batch):
         return self.callable_function(batch)
 
@@ -99,6 +127,7 @@ def __init__(self,
                  ):
         """
         Initialize the object, creating dataset based on input problem
+        :param problem: Problem where data are defined
         :param train_size: number/percentage of elements in train split
         :param test_size: number/percentage of elements in test split
         :param val_size: number/percentage of elements in evaluation split
@@ -112,6 +141,9 @@ def __init__(self,
         self.shuffle = shuffle
         self.repeat = repeat
 
+        # Check if the splits are correct
+        self._check_slit_sizes(train_size, test_size, val_size, predict_size)
+
         # Begin Data splitting
         splits_dict = {}
         if train_size > 0:
@@ -179,23 +211,28 @@ def _split_condition(condition_dict, splits_dict):
         len_condition = len(condition_dict['input_points'])
 
         lengths = [
-            int(math.floor(len_condition * length)) for length in
+            int(len_condition * length) for length in
             splits_dict.values()
         ]
 
         remainder = len_condition - sum(lengths)
         for i in range(remainder):
             lengths[i % len(lengths)] += 1
-        splits_dict = {k: v for k, v in zip(splits_dict.keys(), lengths)
+
+        splits_dict = {k: max(1, v) for k, v in zip(splits_dict.keys(), lengths)
                        }
         to_return_dict = {}
         offset = 0
+
         for stage, stage_len in splits_dict.items():
             to_return_dict[stage] = {k: v[offset:offset + stage_len]
                                      for k, v in condition_dict.items() if
                                      k != 'equation'
                                      # Equations are NEVER dataloaded
                                      }
+            if offset + stage_len > len_condition:
+                offset = len_condition - 1
+                continue
             offset += stage_len
         return to_return_dict
 
@@ -234,6 +271,26 @@ def _apply_shuffle(condition_dict, len_data):
                 dataset_dict[key].update({condition_name: data})
         return dataset_dict
 
+
+    def _create_dataloader(self, split, dataset):
+        shuffle = self.shuffle if split == 'train' else False
+        # Use custom batching (good if batch size is large)
+        if self.batch_size is not None:
+            sampler = PinaSampler(dataset, self.batch_size,
+                                  shuffle, self.automatic_batching)
+            if self.automatic_batching:
+                collate = Collator(self.find_max_conditions_lengths(split))
+
+            else:
+                collate = Collator(None, dataset)
+            return DataLoader(dataset, self.batch_size,
+                              collate_fn=collate, sampler=sampler)
+        dataloader = DummyDataloader(dataset)
+        dataloader.dataset = self._transfer_batch_to_device(
+            dataloader.dataset, self.trainer.strategy.root_device, 0)
+        self.transfer_batch_to_device = self._transfer_batch_to_device_dummy
+        return dataloader
+
     def find_max_conditions_lengths(self, split):
         max_conditions_lengths = {}
         for k, v in self.collector_splits[split].items():
@@ -250,60 +307,28 @@ def val_dataloader(self):
         """
         Create the validation dataloader
         """
-        # Use custom batching (good if batch size is large)
-        if self.batch_size is not None:
-            sampler = PinaSampler(self.val_dataset, self.batch_size,
-                                  self.shuffle, self.automatic_batching)
-            if self.automatic_batching:
-                collate = Collator(self.find_max_conditions_lengths('val'))
-            else:
-                collate = Collator(None, self.val_dataset)
-            return DataLoader(self.val_dataset, self.batch_size,
-                              collate_fn=collate, sampler=sampler)
-        dataloader = DummyDataloader(self.val_dataset,
-                                     self.trainer.strategy.root_device)
-        dataloader.dataset = self._transfer_batch_to_device(dataloader.dataset,
-                                                            self.trainer.strategy.root_device,
-                                                            0)
-        self.transfer_batch_to_device = self._transfer_batch_to_device_dummy
-        return dataloader
+        return self._create_dataloader('val', self.val_dataset)
 
     def train_dataloader(self):
         """
         Create the training dataloader
         """
-        # Use custom batching (good if batch size is large)
-        if self.batch_size is not None:
-            sampler = PinaSampler(self.train_dataset, self.batch_size,
-                                  self.shuffle, self.automatic_batching)
-            if self.automatic_batching:
-                collate = Collator(self.find_max_conditions_lengths('train'))
-
-            else:
-                collate = Collator(None, self.train_dataset)
-            return DataLoader(self.train_dataset, self.batch_size,
-                            collate_fn=collate, sampler=sampler)
-        dataloader = DummyDataloader(self.train_dataset,
-                                     self.trainer.strategy.root_device)
-        dataloader.dataset = self._transfer_batch_to_device(dataloader.dataset,
-                                            self.trainer.strategy.root_device,
-                                            0)
-        self.transfer_batch_to_device = self._transfer_batch_to_device_dummy
-        return dataloader
+        return self._create_dataloader('train', self.train_dataset)
 
     def test_dataloader(self):
         """
         Create the testing dataloader
         """
-        raise NotImplementedError("Test dataloader not implemented")
+        return self._create_dataloader('test', self.test_dataset)
 
     def predict_dataloader(self):
         """
         Create the prediction dataloader
         """
         raise NotImplementedError("Predict dataloader not implemented")
 
-    def _transfer_batch_to_device_dummy(self, batch, device, dataloader_idx):
+    @staticmethod
+    def _transfer_batch_to_device_dummy(batch, device, dataloader_idx):
         return batch
 
     def _transfer_batch_to_device(self, batch, device, dataloader_idx):
@@ -312,10 +337,34 @@ def _transfer_batch_to_device(self, batch, device, dataloader_idx):
         training loop and is used to transfer the batch to the device.
         """
         batch = [
-            (k, super(LightningDataModule, self).transfer_batch_to_device(v,
-                                                                device,
-                                                                dataloader_idx))
+            (k,
+             super(LightningDataModule, self).transfer_batch_to_device(
+                 v, device, dataloader_idx))
             for k, v in batch.items()
         ]
 
         return batch
+
+    @staticmethod
+    def _check_slit_sizes(train_size, test_size, val_size, predict_size):
+        """
+        Check if the splits are correct
+        """
+        if train_size < 0 or test_size < 0 or val_size < 0 or predict_size < 0:
+            raise ValueError("The splits must be positive")
+        if abs(train_size + test_size + val_size + predict_size - 1) > 1e-6:
+            raise ValueError("The sum of the splits must be 1")
+
+    @property
+    def input_points(self):
+        """
+        # TODO
+        """
+        to_return = {}
+        if hasattr(self, "train_dataset") and self.train_dataset is not None:
+            to_return["train"] = self.train_dataset.input_points
+        if hasattr(self, "val_dataset") and self.val_dataset is not None:
+            to_return["val"] = self.val_dataset.input_points
+        if hasattr(self, "test_dataset") and self.test_dataset is not None:
+            to_return = self.test_dataset.input_points
+        return to_return