Fixes

Author: FilippoOlivo

pina/data/data_module.py

@@ -334,7 +334,7 @@ def __init__(
         # collector = Collector(problem)
         # collector.store_fixed_data()
         # collector.store_sample_domains()
-        problem.aggregate_data()
+        problem.collect_data()
 
         # Check if the splits are correct
         self._check_slit_sizes(train_size, test_size, val_size)
@@ -363,7 +363,9 @@ def __init__(
             # raises NotImplementedError
             self.val_dataloader = super().val_dataloader
 
-        self.data_splits = self._create_splits(problem.data, splits_dict)
+        self.data_splits = self._create_splits(
+            problem.collected_data, splits_dict
+        )
         self.transfer_batch_to_device = self._transfer_batch_to_device
 
     def setup(self, stage=None):

pina/problem/abstract_problem.py

@@ -38,7 +38,30 @@ def __init__(self):
                     self.domains[cond_name] = cond.domain
                     cond.domain = cond_name
 
-        self.data = None
+        self._collect_data = {}
+
+    @property
+    def collected_data(self):
+        """
+        Return the collected data from the problem's conditions.
+
+        :return: The collected data.
+        :rtype: dict
+        """
+        if not self._collect_data:
+            raise RuntimeError(
+                "You have to call collect_data() before accessing the data."
+            )
+        return self._collect_data
+
+    @collected_data.setter
+    def collected_data(self, data):
+        """
+        Set the collected data from the problem's conditions.
+
+        :param dict data: The collected data.
+        """
+        self._collect_data = data
 
     #  back compatibility 0.1
     @property
@@ -281,11 +304,11 @@ def add_points(self, new_points_dict):
                 [self.discretised_domains[k], v]
             )
 
-    def aggregate_data(self):
+    def collect_data(self):
         """
         Aggregate data from the problem's conditions into a single dictionary.
         """
-        self.data = {}
+        data = {}
         if not self.are_all_domains_discretised:
             raise RuntimeError(
                 "All domains must be discretised before aggregating data."
@@ -295,11 +318,12 @@ def aggregate_data(self):
             if hasattr(condition, "domain"):
                 samples = self.discretised_domains[condition.domain]
 
-                self.data[condition_name] = {
+                data[condition_name] = {
                     "input": samples,
                     "equation": condition.equation,
                 }
             else:
                 keys = condition.__slots__
                 values = [getattr(condition, name) for name in keys]
-                self.data[condition_name] = dict(zip(keys, values))
+                data[condition_name] = dict(zip(keys, values))
+        self.collected_data = data

tests/test_problem.py

@@ -98,21 +98,23 @@ def test_aggregate_data():
         target=LabelTensor(torch.tensor([[0.0]]), labels=["u"]),
     )
     poisson_problem.discretise_domain(0, "random", domains="all")
-    poisson_problem.aggregate_data()
-    assert isinstance(poisson_problem.data, dict)
+    poisson_problem.collect_data()
+    assert isinstance(poisson_problem.collected_data, dict)
     for name, conditions in poisson_problem.conditions.items():
-        assert name in poisson_problem.data.keys()
+        assert name in poisson_problem.collected_data.keys()
         if isinstance(conditions, InputTargetCondition):
-            assert "input" in poisson_problem.data[name].keys()
-            assert "target" in poisson_problem.data[name].keys()
+            assert "input" in poisson_problem.collected_data[name].keys()
+            assert "target" in poisson_problem.collected_data[name].keys()
         elif isinstance(conditions, DomainEquationCondition):
-            assert "input" in poisson_problem.data[name].keys()
-            assert "target" not in poisson_problem.data[name].keys()
-            assert "equation" in poisson_problem.data[name].keys()
+            assert "input" in poisson_problem.collected_data[name].keys()
+            assert "target" not in poisson_problem.collected_data[name].keys()
+            assert "equation" in poisson_problem.collected_data[name].keys()
 
 
 def test_wrong_aggregate_data():
     poisson_problem = Poisson()
     poisson_problem.discretise_domain(0, "random", domains=["D"])
     with pytest.raises(RuntimeError):
-        poisson_problem.aggregate_data()
+        poisson_problem.collected_data()
+    with pytest.raises(RuntimeError):
+        poisson_problem.collect_data()