Merge pull request #788 from sevmag/restructure-Graph_Definition

Restructure graph definition

Author: sevmag

src/graphnet/data/curated_datamodule.py

@@ -11,9 +11,14 @@
 from glob import glob
 
 from .datamodule import GraphNeTDataModule
-from graphnet.models.graphs import GraphDefinition
+from graphnet.models.data_representation import (
+    GraphDefinition,
+    DataRepresentation,
+)
 from graphnet.data.dataset import ParquetDataset, SQLiteDataset
 
+from graphnet.utilities.logging import Logger
+
 
 class CuratedDataset(GraphNeTDataModule):
     """Generic base class for curated datasets.
@@ -26,8 +31,9 @@ class CuratedDataset(GraphNeTDataModule):
 
     def __init__(
         self,
-        graph_definition: GraphDefinition,
         download_dir: str,
+        data_representation: Optional[DataRepresentation] = None,
+        graph_definition: Optional[GraphDefinition] = None,
         truth: Optional[List[str]] = None,
         features: Optional[List[str]] = None,
         backend: str = "parquet",
@@ -38,8 +44,10 @@ def __init__(
         """Construct CuratedDataset.
 
         Args:
-            graph_definition: Method that defines the data representation.
             download_dir: Directory to download dataset to.
+            data_representation: Method that defines the data representation.
+            graph_definition: Method that defines the graph representation.
+                NOTE: DEPRECATED Use `data_representation` instead.
             truth (Optional): List of event-level truth to include. Will
                             include all available information if not given.
             features (Optional): List of input features from pulsemap to use.
@@ -54,9 +62,17 @@ def __init__(
             test_dataloader_kwargs (Optional): Arguments for the test
                                     DataLoader. Default None.
         """
+        if (data_representation is None) & (graph_definition is not None):
+            data_representation = graph_definition
+        elif (data_representation is None) & (graph_definition is None):
+            # Code stops
+            raise TypeError(
+                "__init__() missing 1 required keyword argument:"
+                "'data_representation'"
+            )
+        self._data_representation = data_representation
         # From user
         self._download_dir = download_dir
-        self._graph_definition = graph_definition
         self._backend = backend.lower()
 
         # Checks
@@ -85,6 +101,15 @@ def __init__(
             test_selection=test_selec,
         )
 
+        if graph_definition is not None:
+            # Code continues after warning
+            self.warning(
+                "DeprecationWarning: Argument `graph_definition` will be"
+                " deprecated in GraphNeT 2.0. Please use `data_representation`"
+                " instead."
+                ""
+            )
+
     @abstractmethod
     def prepare_data(self) -> None:
         """Download and prepare data."""
@@ -249,6 +274,28 @@ def dataset_dir(self) -> str:
         )
         return dataset_dir
 
+    # DEPRECATION: REMOVE AT 2.0 LAUNCH
+    # See https://github.com/graphnet-team/graphnet/issues/647
+    @property
+    def _graph_definition(self) -> DataRepresentation:
+        """Return the graph definition."""
+        # needed for the call in _prepare_args
+        # call before Logger init
+        if hasattr(self, "_logger"):
+            self.warning(
+                "DeprecationWarning: `_graph_definition` will be deprecated in"
+                " GraphNeT 2.0. Please use `_data_representation` instead."
+            )
+        else:
+            Logger(log_folder=None).warning_once(
+                (
+                    "`graphnet.models.graphs` will be depricated soon. "
+                    "All functionality has been moved to "
+                    "`graphnet.models.data_representation`."
+                )
+            )
+        return self._data_representation  # type: ignore
+
 
 class ERDAHostedDataset(CuratedDataset):
     """A base class for dataset/datamodule hosted at ERDA.

src/graphnet/data/dataset/dataset.py

@@ -31,8 +31,10 @@
 )
 from graphnet.exceptions.exceptions import ColumnMissingException
 from graphnet.utilities.logging import Logger
-from graphnet.models.graphs import GraphDefinition
-
+from graphnet.models.data_representation import (
+    GraphDefinition,
+    DataRepresentation,
+)
 from graphnet.utilities.config.parsing import (
     get_all_grapnet_classes,
 )
@@ -58,11 +60,11 @@ def load_module(class_name: str) -> Type:
     return namespace_classes[class_name]
 
 
-def parse_graph_definition(cfg: dict) -> GraphDefinition:
-    """Construct GraphDefinition from DatasetConfig."""
-    assert cfg["graph_definition"] is not None
+def parse_data_representation(data_rep_cfg: dict) -> DataRepresentation:
+    """Construct DataRepresentation from DatasetConfig."""
+    assert data_rep_cfg is not None
 
-    args = cfg["graph_definition"]["arguments"]
+    args = data_rep_cfg["arguments"]
     classes = {}
     for arg in args.keys():
         if isinstance(args[arg], dict):
@@ -75,10 +77,8 @@ def parse_graph_definition(cfg: dict) -> GraphDefinition:
 
     new_cfg = deepcopy(args)
     new_cfg.update(classes)
-    graph_definition = load_module(cfg["graph_definition"]["class_name"])(
-        **new_cfg
-    )
-    return graph_definition
+    data_representation = load_module(data_rep_cfg["class_name"])(**new_cfg)
+    return data_representation
 
 
 def parse_labels(cfg: dict) -> Dict[str, Label]:
@@ -122,9 +122,18 @@ def from_config(  # type: ignore[override]
             "`DatasetConfig`"
         )
 
-        assert (
-            "graph_definition" in source.dict().keys()
-        ), "`DatasetConfig` incompatible with current GraphNeT version."
+        if "data_representation" not in source.dict().keys():
+            if "graph_definition" in source.dict().keys():
+                Logger(log_folder=None).warning_once(
+                    "DeprecationWarning: Field `graph_definition` will be"
+                    " deprecated in GraphNeT 2.0. Please use "
+                    "`data_representation` instead."
+                )
+            else:
+                raise TypeError(
+                    "`DatasetConfig` incompatible with "
+                    "current GraphNeT version."
+                )
 
         # Parse set of `selection``.
         if isinstance(source.selection, dict):
@@ -137,8 +146,24 @@ def from_config(  # type: ignore[override]
             return cls._construct_dataset_from_list_of_strings(source)
 
         cfg = source.dict()
-        if cfg["graph_definition"] is not None:
-            cfg["graph_definition"] = parse_graph_definition(cfg)
+
+        if (
+            "data_representation" in cfg
+            and cfg["data_representation"] is not None
+        ):
+            cfg["data_representation"] = parse_data_representation(
+                cfg["data_representation"]
+            )
+        elif "graph_definition" in cfg and cfg["graph_definition"] is not None:
+            Logger(log_folder=None).warning_once(
+                "DeprecationWarning: Field `graph_definition` will be"
+                " deprecated in GraphNeT 2.0. Please use "
+                "`data_representation` instead."
+            )
+            cfg["graph_definition"] = parse_data_representation(
+                cfg["graph_definition"]
+            )
+
         if cfg["labels"] is not None:
             cfg["labels"] = parse_labels(cfg)
 
@@ -216,11 +241,12 @@ def _resolve_graphnet_paths(
     def __init__(
         self,
         path: Union[str, List[str]],
-        graph_definition: GraphDefinition,
         pulsemaps: Union[str, List[str]],
         features: List[str],
         truth: List[str],
         *,
+        graph_definition: Optional[GraphDefinition] = None,
+        data_representation: Optional[DataRepresentation] = None,
         node_truth: Optional[List[str]] = None,
         index_column: str = "event_no",
         truth_table: str = "truth",
@@ -278,8 +304,13 @@ def __init__(
                 subset of events when resolving a string-based selection (e.g.,
                 `"10000 random events ~ event_no % 5 > 0"` or `"20% random
                 events ~ event_no % 5 > 0"`).
-            graph_definition: Method that defines the graph representation.
+            data_representation: Method that defines the data representation.
             labels: Dictionary of labels to be added to the dataset.
+
+            graph_definition: Method that defines the graph representation.
+                NOTE: DEPRECATED Use `data_representation` instead.
+                # DEPRECATION: REMOVE AT 2.0 LAUNCH
+                # See https://github.com/graphnet-team/graphnet/issues/647
         """
         # Base class constructor
         super().__init__(name=__name__, class_name=self.__class__.__name__)
@@ -303,9 +334,26 @@ def __init__(
         self._index_column = index_column
         self._truth_table = truth_table
         self._loss_weight_default_value = loss_weight_default_value
-        self._graph_definition = deepcopy(graph_definition)
+
+        if data_representation is None:
+            if graph_definition is not None:
+                data_representation = graph_definition
+                # Code continues after warning
+                self.warning(
+                    "DeprecationWarning: Argument `graph_definition` "
+                    "will be deprecated in GraphNeT 2.0. "
+                    "Please use `data_representation` instead."
+                )
+            else:
+                # Code stops
+                raise TypeError(
+                    "__init__() missing 1 required keyword argument:"
+                    "'data_representation'"
+                )
+
+        self._data_representation = deepcopy(data_representation)
         self._labels = labels
-        self._string_column = graph_definition._detector.string_index_name
+        self._string_column = data_representation._detector.string_index_name
 
         if node_truth is not None:
             assert isinstance(node_truth_table, str)
@@ -390,6 +438,17 @@ def truth_table(self) -> str:
         """Name of the table containing event-level truth information."""
         return self._truth_table
 
+    # DEPRECATION PROPERTY: REMOVE AT 2.0 LAUNCH
+    # See https://github.com/graphnet-team/graphnet/issues/647
+    @property
+    def _graph_definition(self) -> DataRepresentation:
+        """Return the graph definition."""
+        self.warning(
+            "DeprecationWarning: `_graph_definition` will be deprecated in"
+            " GraphNeT 2.0. Please use `_data_representation` instead."
+        )
+        return self._data_representation
+
     # Abstract method(s)
     @abstractmethod
     def _init(self) -> None:
@@ -647,8 +706,8 @@ def _create_graph(
 
         assert isinstance(features, np.ndarray)
         # Construct graph data object
-        assert self._graph_definition is not None
-        graph = self._graph_definition(
+        assert self._data_representation is not None
+        graph = self._data_representation(
             input_features=node_features,
             input_feature_names=self._features,
             truth_dicts=truth_dicts,

src/graphnet/data/dataset/parquet/parquet_dataset.py

@@ -18,7 +18,10 @@
 from bisect import bisect_right
 from collections import OrderedDict
 
-from graphnet.models.graphs import GraphDefinition
+from graphnet.models.data_representation import (
+    GraphDefinition,
+    DataRepresentation,
+)
 from graphnet.data.dataset import Dataset
 from graphnet.exceptions.exceptions import ColumnMissingException
 
@@ -29,11 +32,12 @@ class ParquetDataset(Dataset):
     def __init__(
         self,
         path: str,
-        graph_definition: GraphDefinition,
         pulsemaps: Union[str, List[str]],
         features: List[str],
         truth: List[str],
         *,
+        data_representation: Optional[DataRepresentation] = None,
+        graph_definition: Optional[GraphDefinition] = None,
         node_truth: Optional[List[str]] = None,
         index_column: str = "event_no",
         truth_table: str = "truth",
@@ -92,7 +96,9 @@ def __init__(
                 subset of events when resolving a string-based selection (e.g.,
                 `"10000 random events ~ event_no % 5 > 0"` or `"20% random
                 events ~ event_no % 5 > 0"`).
+            data_representation: Method that defines the data representation.
             graph_definition: Method that defines the graph representation.
+                NOTE: DEPRECATED Use `data_representation` instead.
             cache_size: Number of files to cache in memory.
                         Must be at least 1. Defaults to 1.
             labels: Dictionary of labels to be added to the dataset.
@@ -116,6 +122,7 @@ def __init__(
             loss_weight_default_value=loss_weight_default_value,
             seed=seed,
             graph_definition=graph_definition,
+            data_representation=data_representation,
             labels=labels,
         )
 

src/graphnet/models/data_representation/__init__.py

@@ -0,0 +1,20 @@
+"""Modules for constructing data.
+
+´DataRepresentation´ defines the basic structure for representing data.
+´GraphDefinition´ defines graphs with different nodes and their features, as
+well as the edges between them.
+"""
+
+from .data_representation import DataRepresentation
+from .graphs import (
+    GraphDefinition,
+    KNNGraph,
+    EdgelessGraph,
+    KNNGraphRRWP,
+    KNNGraphRWSE,
+    NodeDefinition,
+    NodesAsPulses,
+    PercentileClusters,
+    NodeAsDOMTimeSeries,
+    IceMixNodes,
+)