getting rid of _create_data

Author: sevmag

src/graphnet/models/data_representation/data_representation.py

@@ -170,7 +170,7 @@ def forward(  # type: ignore
         input_features = self._detector(input_features, input_feature_names)
 
         # Create data & get new final data feature names
-        data = self._create_data(input_features=input_features)
+        data = Data(x=input_features)
 
         # Attach number of pulses as static attribute.
         data.n_pulses = torch.tensor(len(input_features), dtype=torch.int32)
@@ -413,20 +413,6 @@ def _set_output_feature_names(
         """Set the final data output feature names."""
         raise NotImplementedError
 
-    @abstractmethod
-    def _create_data(self, input_features: torch.Tensor) -> Data:
-        """Create data from input features.
-
-        Enforce the dtype of the feature tensor.
-        E.g.: `data.x = data.x.type(self.dtype)`
-        if the training data is stored in `data.x`.
-
-        Should return:
-            - data: torch_geometric.data.Data object representing the event.
-            - data_feature_names: List of feature names in the data object.
-        """
-        raise NotImplementedError
-
     def _label_repeater(self, label: torch.Tensor, data: Data) -> torch.Tensor:
         """Handle the label repetition.
 

src/graphnet/models/data_representation/graphs/graph_definition.py

@@ -5,7 +5,7 @@
 and can be passed to dataloaders during training and deployment.
 """
 
-from typing import List, Optional, Dict, Union, Tuple, Any, Callable
+from typing import List, Optional, Dict, Union, Any, Callable
 import torch
 from numpy.random import Generator
 import numpy as np
@@ -120,23 +120,6 @@ def _set_output_feature_names(
         self._node_definition.set_output_feature_names(input_feature_names)
         return self._node_definition._output_feature_names
 
-    def _create_data(
-        self, input_features: torch.Tensor
-    ) -> Tuple[Data, List[str]]:
-        # Create graph & get new node feature names
-        data = self._node_definition(input_features)
-        if self._sort_by is not None:
-            data.x = data.x[data.x[:, self._sort_by].sort()[1]]
-
-        # Enforce dtype
-        data.x = data.x.type(self.dtype)
-
-        # Assign edges
-        if self._edge_definition is not None:
-            data = self._edge_definition(data)
-
-        return data
-
     def forward(  # type: ignore
         self,
         input_features: np.ndarray,
@@ -178,6 +161,18 @@ def forward(  # type: ignore
             data_path=data_path,
         )
 
+        # Create graph & get new node feature names
+        data = self._node_definition(data.x)
+        if self._sort_by is not None:
+            data.x = data.x[data.x[:, self._sort_by].sort()[1]]
+
+        # Enforce dtype
+        data.x = data.x.type(self.dtype)
+
+        # Assign edges
+        if self._edge_definition is not None:
+            data = self._edge_definition(data)
+
         if self._add_static_features:
             data = self._add_features_individually(
                 data,
@@ -186,17 +181,6 @@ def forward(  # type: ignore
 
         return data
 
-    def _forward_end(
-        self,
-        data: Data,
-        data_feature_names: List[str],
-    ) -> Data:
-        """Add processing steps at the end of the forward pass."""
-        # Add original features as attributes
-        if self._add_static_features:
-            data = self._add_features_individually(data, data_feature_names)
-        return data
-
     def _add_features_individually(
         self,
         data: Data,