Conditions refactoring (#758)

Author: FilippoOlivo

pina/_src/condition/condition_base.py

@@ -34,7 +34,10 @@ def __init__(self, **kwargs):
         """
         super().__init__()
         self.data = self.store_data(**kwargs)
+<<<<<<< HEAD
         self.has_custom_dataloader_fn = False
+=======
+>>>>>>> 4ca4993a (Conditions refactoring (#758))
 
     @property
     def problem(self):
@@ -87,8 +90,12 @@ def automatic_batching_collate_fn(cls, batch):
         if not batch:
             return {}
         instance_class = batch[0].__class__
+<<<<<<< HEAD
         batch = instance_class.create_batch(batch)
         return batch
+=======
+        return instance_class.create_batch(batch)
+>>>>>>> 4ca4993a (Conditions refactoring (#758))
 
     @staticmethod
     def collate_fn(batch, condition):
@@ -106,11 +113,15 @@ def collate_fn(batch, condition):
         return data
 
     def create_dataloader(
+<<<<<<< HEAD
         self,
         dataset,
         batch_size,
         automatic_batching,
         **kwargs,
+=======
+        self, dataset, batch_size, shuffle, automatic_batching
+>>>>>>> 4ca4993a (Conditions refactoring (#758))
     ):
         """
         Create a DataLoader for the condition.
@@ -121,14 +132,23 @@ def create_dataloader(
         :rtype: torch.utils.data.DataLoader
         """
         if batch_size == len(dataset):
+<<<<<<< HEAD
             return DummyDataloader(dataset)
         return DataLoader(
             dataset=dataset,
+=======
+            pass  # will be updated in the near future
+        return DataLoader(
+            dataset=dataset,
+            batch_size=batch_size,
+            shuffle=shuffle,
+>>>>>>> 4ca4993a (Conditions refactoring (#758))
             collate_fn=(
                 partial(self.collate_fn, condition=self)
                 if not automatic_batching
                 else self.automatic_batching_collate_fn
             ),
+<<<<<<< HEAD
             batch_size=batch_size,
             **kwargs,
         )
@@ -146,3 +166,6 @@ def switch_dataloader_fn(self, create_dataloader_fn):
         # the new function
         self.has_custom_dataloader_fn = True
         self.create_dataloader = create_dataloader_fn
+=======
+        )
+>>>>>>> 4ca4993a (Conditions refactoring (#758))

pina/_src/condition/data_manager.py

@@ -120,6 +120,7 @@ def create_batch(items):
                     else torch.stack
                 )
                 batch_data[k] = batch_fn(vals)
+                batch_data[k] = batch_fn(vals, dim=0)
             else:
                 batch_data[k] = sample
         return batch_data

pina/_src/condition/domain_equation_condition.py

@@ -31,7 +31,11 @@ class DomainEquationCondition(ConditionBase):
     # Available slots
     __fields__ = ["domain", "equation"]
 
+<<<<<<< HEAD
     _avail_domain_cls = (DomainInterface, str)
+=======
+    _avail_domain_cls = DomainInterface
+>>>>>>> 4ca4993a (Conditions refactoring (#758))
     _avail_equation_cls = EquationInterface
 
     def __new__(cls, domain, equation):