Merge pull request #481 from asogaard/concatdataset-to-ensembledataset

Replace ConcatDataset with EnsembleDataset

src/graphnet/data/dataset.py

@@ -14,10 +14,8 @@
     Iterable,
 )
 
-from tqdm import tqdm
 import numpy as np
 import torch
-from torch.utils.data import ConcatDataset
 from torch_geometric.data import Data
 
 from graphnet.constants import GRAPHNET_ROOT_DIR
@@ -46,9 +44,9 @@ def from_config(  # type: ignore[override]
         source: Union[DatasetConfig, str],
     ) -> Union[
         "Dataset",
-        ConcatDataset,
+        "EnsembleDataset",
         Dict[str, "Dataset"],
-        Dict[str, ConcatDataset],
+        Dict[str, "EnsembleDataset"],
     ]:
         """Construct `Dataset` instance from `source` configuration."""
         if isinstance(source, str):
@@ -75,9 +73,9 @@ def from_config(  # type: ignore[override]
     def concatenate(
         cls,
         datasets: List["Dataset"],
-    ) -> ConcatDataset:
+    ) -> "EnsembleDataset":
         """Concatenate multiple `Dataset`s into one instance."""
-        return ConcatDataset(datasets)
+        return EnsembleDataset(datasets)
 
     @classmethod
     def _construct_datasets_from_dict(
@@ -90,7 +88,7 @@ def _construct_datasets_from_dict(
         for key, selection in selections.items():
             config.selection = selection
             dataset = Dataset.from_config(config)
-            assert isinstance(dataset, (Dataset, ConcatDataset))
+            assert isinstance(dataset, (Dataset, EnsembleDataset))
             datasets[key] = dataset
 
         # Reset `selections`.

src/graphnet/utilities/config/dataset_config.py

@@ -95,12 +95,12 @@ def __init__(self, **data: Any) -> None:
                 (...)
             }
             >>> dataset.config.dump("dataset.yml")
-            >>> datasets: Dict[str, ConcatDataset] = Dataset.from_config(
+            >>> datasets: Dict[str, EnsembleDataset] = Dataset.from_config(
                 "dataset.yml"
             )
             >>> datasets
             {
-                "train": ConcatDataset(...),
+                "train": EnsembleDataset(...),
                 (...)
             }