Allows for adding labels in dataset config

src/graphnet/data/dataset/dataset.py

@@ -85,6 +85,18 @@ def parse_graph_definition(cfg: dict) -> GraphDefinition:
     return graph_definition
 
 
+def parse_labels(cfg: dict) -> Dict[str, Label]:
+    """Construct Label from DatasetConfig."""
+    assert cfg["labels"] is not None
+
+    labels = {}
+    for key in cfg["labels"].keys():
+        labels[key] = load_module(cfg["labels"][key]["class_name"])(
+            **cfg["labels"][key]["arguments"]
+        )
+    return labels
+
+
 class Dataset(
     Logger,
     Configurable,
@@ -147,6 +159,8 @@ def from_config(  # type: ignore[override]
         cfg = source.dict()
         if cfg["graph_definition"] is not None:
             cfg["graph_definition"] = parse_graph_definition(cfg)
+        if cfg["labels"] is not None:
+            cfg["labels"] = parse_labels(cfg)
         return source._dataset_class(**cfg)
 
     @classmethod
@@ -230,6 +244,7 @@ def __init__(
         loss_weight_default_value: Optional[float] = None,
         seed: Optional[int] = None,
         use_cache: bool = True,
+        labels: Optional[Dict[str, Any]] = None,
     ):
         """Construct Dataset.
 
@@ -277,6 +292,7 @@ def __init__(
                 events ~ event_no % 5 > 0"`).
             graph_definition: Method that defines the graph representation.
             use_cache: Whether or not to save indices and selections to a temporary cache for faster initializing.
+            labels: Dictionary of labels to be added to the dataset.
         """
         # Base class constructor
         super().__init__(name=__name__, class_name=self.__class__.__name__)
@@ -301,6 +317,7 @@ def __init__(
         self._truth_table = truth_table
         self._loss_weight_default_value = loss_weight_default_value
         self._graph_definition = deepcopy(graph_definition)
+        self._labels = labels
 
         if node_truth is not None:
             assert isinstance(node_truth_table, str)
@@ -351,6 +368,9 @@ def __init__(
             use_cache=use_cache,
         )
 
+        if self._labels is not None:
+            for key in self._labels.keys():
+                self.add_label(self._labels[key])
         # Implementation-specific initialisation.
         self._init()
 

src/graphnet/utilities/config/dataset_config.py

@@ -56,6 +56,7 @@ class DatasetConfig(BaseConfig):
     seed: Optional[int] = None
     graph_definition: Any = None
     use_cache: bool = True
+    labels: Optional[Dict[str, Any]] = None
 
     def __init__(self, **data: Any) -> None:
         """Construct `DataConfig`.