fix trainer + add tests

Author: GiovanniCanali

pina/_src/callback/refinement/base_refinement.py

@@ -99,7 +99,7 @@ def on_train_start(self, trainer, solver):
         # Initialize dataset and compute initial population size
         self._dataset = trainer.datamodule.train_datasets
         self._initial_population_size = {
-            cond: self.dataset[cond].length
+            cond: self.dataset[cond].dataset_length
             for cond in self._condition_to_update
         }
 

pina/_src/condition/time_series_condition.py

@@ -200,7 +200,7 @@ def evaluate(self, batch, solver):
             raise ValueError(
                 "The provided input tensor must have at least 4 dimensions:"
                 " [trajectories, windows, time_steps, *features]."
-                f" Got shape {batch["input"].shape}."
+                f" Got shape {batch['input'].shape}."
             )
 
         # Copy the kwargs to avoid modifying the original settings

pina/_src/core/trainer.py

@@ -10,12 +10,6 @@ class _Aggregator:
     iteration of multiple training conditions within a single training loop.
     """
 
-    _AVAIL_BATCHING_MODES = {
-        "common_batch_size",
-        "proportional",
-        "separate_conditions",
-    }
-
     def __init__(self, dataloaders, batching_mode):
         """
         Initialization of the :class:`_Aggregator` class.
@@ -27,17 +21,9 @@ def __init__(self, dataloaders, batching_mode):
             uniform batch sizes across conditions, ``"proportional"`` for batch
             sizes proportional to dataset sizes, and ``"separate_conditions"``
             for iterating through each condition separately.
-        :raises ValueError: If an invalid batching mode is provided.
         :raises NotImplementedError: If the selected batching mode is not yet
             implemented.
         """
-        # Check consistency
-        if batching_mode not in self._AVAIL_BATCHING_MODES:
-            raise ValueError(
-                f"Invalid batching mode '{batching_mode}'. "
-                f"Available options are: {self._AVAIL_BATCHING_MODES}"
-            )
-
         # Raise not implemented error for separate_conditions mode
         if batching_mode == "separate_conditions":
             raise NotImplementedError(

pina/_src/data/aggregator.py

@@ -12,13 +12,6 @@ class _Creator:
     behavior to specific training requirements
     """
 
-    # Available batching modes
-    _AVAIL_BATCHING_MODES = {
-        "common_batch_size",
-        "proportional",
-        "separate_conditions",
-    }
-
     def __init__(
         self,
         batching_mode,
@@ -53,15 +46,7 @@ def __init__(
         :param dict[str, BaseCondition] conditions: The mapping between
             condition names and condition objects responsible for data loader
             creation.
-        :raises ValueError: If an invalid batching mode is provided.
         """
-        # Check consistency
-        if batching_mode not in self._AVAIL_BATCHING_MODES:
-            raise ValueError(
-                f"Invalid batching mode '{batching_mode}'. "
-                f"Available options are: {self._AVAIL_BATCHING_MODES}"
-            )
-
         # Initialize attributes
         self.batching_mode = batching_mode
         self.batch_size = batch_size

pina/_src/data/creator.py

@@ -99,10 +99,13 @@ def __init__(
         # Move domain discretisation into conditions subsets
         self.problem.move_discretisation_into_conditions()
 
-        # Verify which splits are zero
-        self._has_train = train_size > 0
-        self._has_val = val_size > 0
-        self._has_test = test_size > 0
+        # If no splits are defined, use the default dataloaders
+        if train_size == 0:
+            self.train_dataloader = super().train_dataloader
+        if val_size == 0:
+            self.val_dataloader = super().val_dataloader
+        if test_size == 0:
+            self.test_dataloader = super().test_dataloader
 
         # Otherwise, create the condition splits and initialize the creator
         self._create_condition_splits(train_size, test_size)
@@ -244,14 +247,6 @@ def train_dataloader(self):
             dataloaders.
         :rtype: _Aggregator
         """
-        # If no training split is defined, return the default dataloader
-        if not self._has_train:
-            return super().train_dataloader()
-
-        # If the training dataloaders have not been created yet, call setup
-        if not hasattr(self, "train_datasets"):
-            self.setup("fit")
-
         return _Aggregator(
             self.creator(self.train_datasets),
             batching_mode=self.batching_mode,
@@ -265,14 +260,6 @@ def val_dataloader(self):
             dataloaders.
         :rtype: _Aggregator
         """
-        # If no validation split is defined, return the default dataloader
-        if not self._has_val:
-            return super().val_dataloader()
-
-        # If the validation dataloaders have not been created yet, call setup
-        if not hasattr(self, "val_datasets"):
-            self.setup("fit")
-
         return _Aggregator(
             self.creator(self.val_datasets), batching_mode=self.batching_mode
         )
@@ -285,14 +272,6 @@ def test_dataloader(self):
             dataloaders.
         :rtype: _Aggregator
         """
-        # If no test split is defined, return the default dataloader
-        if not self._has_test:
-            return super().test_dataloader()
-
-        # If the test dataloaders have not been created yet, call setup
-        if not hasattr(self, "test_datasets"):
-            self.setup("test")
-
         return _Aggregator(
             self.creator(self.test_datasets),
             batching_mode=self.batching_mode,

pina/_src/data/data_module.py

@@ -296,14 +296,3 @@ def are_all_domains_discretised(self):
         :rtype: bool
         """
         return all(d in self.discretised_domains for d in self.domains)
-
-
-# Back-compatibility with version 0.2, to be removed soon
-class AbstractProblem(BaseProblem):
-    def __init__(self, *args, **kwargs):
-        warnings.warn(
-            "AbstractProblem is deprecated, use BaseProblem instead.",
-            DeprecationWarning,
-            stacklevel=2,
-        )
-        super().__init__(*args, **kwargs)

pina/_src/problem/base_problem.py

@@ -1,7 +1,6 @@
 """Module for the Problems."""
 
 __all__ = [
-    "AbstractProblem",  # back-compatibility with version 0.2, to be removed soon
     "ProblemInterface",
     "BaseProblem",
     "SpatialProblem",
@@ -18,4 +17,19 @@
 from pina._src.problem.inverse_problem import InverseProblem
 
 # Back-compatibility with version 0.2, to be removed soon
-from pina._src.problem.base_problem import AbstractProblem
+import warnings
+
+_DEPRECATED_IMPORTS = {"AbstractProblem": "BaseProblem"}
+
+
+def __getattr__(name):
+    if name in _DEPRECATED_IMPORTS:
+
+        warnings.warn(
+            f"Importing '{name}' from 'pina.problem' is deprecated; use "
+            f"pina.problem.{_DEPRECATED_IMPORTS[name]} instead.",
+            DeprecationWarning,
+            stacklevel=2,
+        )
+
+        return globals()[_DEPRECATED_IMPORTS[name]]

pina/problem/__init__.py

@@ -69,10 +69,6 @@ def test_constructor(batching_mode):
     # Initialize the aggregator
     _Aggregator(dataloaders, batching_mode=batching_mode)
 
-    # Should fail if an invalid batching mode is provided
-    with pytest.raises(ValueError):
-        _Aggregator(dataloaders, batching_mode="invalid_mode")
-
     # Should raise NotImplementedError for separate_conditions mode
     with pytest.raises(NotImplementedError):
         _Aggregator(dataloaders, batching_mode="separate_conditions")

tests/test_data/test_aggregator.py

@@ -51,7 +51,10 @@ def create_dataloader(
 }
 
 
-@pytest.mark.parametrize("batching_mode", _Creator._AVAIL_BATCHING_MODES)
+@pytest.mark.parametrize(
+    "batching_mode",
+    ["common_batch_size", "separate_conditions", "proportional"],
+)
 def test_constructor(batching_mode):
 
     _Creator(
@@ -64,18 +67,6 @@ def test_constructor(batching_mode):
         conditions=dataloaders,
     )
 
-    # Should fail if an invalid batching mode is provided
-    with pytest.raises(ValueError):
-        _Creator(
-            batching_mode="invalid_mode",
-            batch_size=4,
-            shuffle=False,
-            automatic_batching=True,
-            num_workers=0,
-            pin_memory=False,
-            conditions=dataloaders,
-        )
-
 
 @pytest.mark.parametrize(
     "batching_mode, batch_size, expected_batch_sizes, expected_max_len",