Remove deprecation warning, add on_after_load_sample option

	modified:   data.py
	modified:   simulator.py

Author: cweniger

swyft/lightning/data.py

@@ -48,6 +48,7 @@ def __init__(
         batch_size: int = 32,
         num_workers: int = 0,
         shuffle: bool = False,
+        on_after_load_sample: Optional[callable] = None,
     ):
         super().__init__()
         self.data = data
@@ -60,6 +61,7 @@ def __init__(
         self.batch_size = batch_size
         self.num_workers = num_workers
         self.shuffle = shuffle
+        self.on_after_load_sample = on_after_load_sample
 
     @staticmethod
     def _get_lengths(fractions, N):
@@ -72,15 +74,15 @@ def _get_lengths(fractions, N):
 
     def setup(self, stage: str):
         if isinstance(self.data, Samples):
-            dataset = self.data.get_dataset()
+            dataset = self.data.get_dataset(on_after_load_sample = self.on_after_load_sample)
             splits = torch.utils.data.random_split(dataset, self.lengths)
             self.dataset_train, self.dataset_val, self.dataset_test = splits
         elif isinstance(self.data, swyft.ZarrStore):
             idxr1 = (0, self.lengths[1])
             idxr2 = (self.lengths[1], self.lengths[1] + self.lengths[2])
             idxr3 = (self.lengths[1] + self.lengths[2], len(self.data))
             self.dataset_train = self.data.get_dataset(
-                idx_range=idxr1, on_after_load_sample=None
+                idx_range=idxr1, on_after_load_sample=self.on_after_load_sample
             )
             self.dataset_val = self.data.get_dataset(
                 idx_range=idxr2, on_after_load_sample=None

swyft/lightning/simulator.py

@@ -76,15 +76,14 @@ def get_dataloader(
         repeat=None,
         num_workers=0,
     ):
-        """(Deprecated) Generator function to directly generate a dataloader object.
+        """Generator function to directly generate a dataloader object.
 
         Args:
             batch_size: batch_size for dataloader
             shuffle: shuffle for dataloader
             on_after_load_sample: see `get_dataset`
             repeat: If not None, Wrap dataset in RepeatDatasetWrapper
         """
-        print("WARNING: Deprecated")
         dataset = self.get_dataset(on_after_load_sample=on_after_load_sample)
         if repeat is not None:
             dataset = swyft.lightning.data.RepeatDatasetWrapper(dataset, repeat=repeat)