Feature/rand split return recording

dn3/data/dataset.py

@@ -1,3 +1,4 @@
+from typing import Optional
 import mne
 import torch
 import copy
@@ -15,6 +16,7 @@
 from pathlib import Path
 from torch.utils.data import Dataset as TorchDataset
 from torch.utils.data import ConcatDataset, DataLoader
+from torch.utils.data.dataset import Subset as TorchSubset
 
 
 class DN3ataset(TorchDataset):
@@ -180,8 +182,6 @@ def to_numpy(self, batch_size=64, batch_transforms: list = None, num_workers=4,
                 loaded = [np.concatenate([loaded[i], batch[i]], axis=0) for i in range(len(batch))]
 
         return loaded
-
-
 class _Recording(DN3ataset, ABC):
     """
     Abstract base class for any supported recording
@@ -433,6 +433,26 @@ def get_targets(self):
         return np.apply_along_axis(lambda x: self.epoch_codes_to_class_labels[x[0]], 1,
                                    self.epochs.events[list(self._skip_map.values()), -1, np.newaxis]).squeeze()
 
+class DN3ataSubSet(DN3ataset):
+    """
+    Wrap a torch subset of a DN3ataset.
+    """
+    def __init__(self, dn3ata: DN3ataset, subset: TorchSubset):
+        DN3ataset.__init__(self)
+        self.dataset = subset.dataset
+        self.indices = subset.indices
+        if not hasattr(dn3ata, 'get_targets'):
+            raise ValueError("dn3ata must have a get_targets method")
+        self.targets = dn3ata.get_targets()[subset.indices]
+
+    def __getitem__(self, idx):
+        return TorchSubset.__getitem__(self, idx)
+
+    def __len__(self):
+        return TorchSubset.__len__(self)
+
+    def get_targets(self):
+        return self.targets
 
 class Thinker(DN3ataset, ConcatDataset):
     """
@@ -608,7 +628,7 @@ def split(self, training_sess_ids=None, validation_sess_ids=None, testing_sess_i
             if len(use_sessions) > 0:
                 print("Warning: sessions specified do not span all sessions. Skipping {} sessions.".format(
                     len(use_sessions)))
-                return training, validating, testing
+                return self._dn3_or_none(training), self._dn3_or_none(validating), self._dn3_or_none(testing)
 
         # Split up the rest if there is anything left
         if len(use_sessions) > 0:
@@ -622,8 +642,14 @@ def split(self, training_sess_ids=None, validation_sess_ids=None, testing_sess_i
                     validating, remainder = rand_split(remainder, frac=validation_frac)
 
         training = remainder if training is None else training
-
-        return training, validating, testing
+
+        return self._dn3_or_none(training), self._dn3_or_none(validating), self._dn3_or_none(testing)
+
+    def _dn3_or_none(self, subset: Optional[DN3ataset]) -> Optional[DN3ataset]:
+        if subset is None or type(subset) is DN3ataset:
+            return subset
+
+        return DN3ataSubSet(self, subset)
 
     def preprocess(self, preprocessor: Preprocessor, apply_transform=True, sessions=None, **kwargs):
         """

dn3/utils.py

@@ -53,6 +53,7 @@ def rand_split(dataset, frac=0.75):
     if frac >= 1:
         return dataset
     samples = len(dataset)
+
     return random_split(dataset, lengths=[round(x) for x in [samples*frac, samples*(1-frac)]])